"""
Create fake 3-feature fingerprint data by generating random vectors,
separating into two clusters, then using the clusters as class labels.
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d
from sklearn.cluster import KMeans
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression

import sys

def add_noise(label, prob):
    def flip(bit):
        return 0 if bit == 1 else 1
    return flip(label) if np.random.random() < prob else label

def gen_data():
    feat1 = np.random.random(1000)
    feat2 = np.random.random(1000)
    feat3 = np.random.random(1000)
    df = pd.DataFrame({"feature1": feat1, "feature2": feat2, "feature3": feat3})

    kmeans = KMeans(n_clusters=2, random_state=0).fit(df)

    labels = kmeans.predict(df)
    labels_with_noise = [add_noise(label, .10) for label in labels]

    df["label"] = labels_with_noise
    return df

def main(args):
    # One of the consequences of Python's stupid whitespace-based syntax is
    # that there's no do-while loop, so you have to use hacks like this.
    score = 1.0

    # Don't want the data to be too hard or too easy
    while score > .85 or score < .80:
        df = gen_data()
        X = df.drop("label", axis=1)
        y = df["label"]

        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.20)
        model = LogisticRegression(solver="liblinear").fit(X_train, y_train)
        score = model.score(X_test, y_test)
        print(score)
    df.index.name = "id"
    df.to_csv("fake-fingerprints.csv")

    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")

    ones = df[df["label"] == 1]
    zeroes = df[df["label"] == 0]

    ax.scatter(ones["feature1"], ones["feature2"],
               ones["feature3"], marker="+",
               color="green")
    ax.scatter(zeroes["feature1"], zeroes["feature2"], zeroes["feature3"],
               marker="_", color="red")

    ax.set_xlabel('feature1')
    ax.set_ylabel('feature2')
    ax.set_zlabel('feature3')
    ax.set_title("Fake Fingerprints Scatterplot")
    #plt.show()
    fig.savefig("fake-fingerprints-scatter.png")

if __name__=="__main__":
    main(sys.argv)