from sklearn import datasets
from sklearn.model_selection import train_test_split
import numpy as np
import matplotlib.pyplot as plt
from diffprivlib.models import GaussianNB


def main():
    dataset = datasets.load_iris()
    x_train, x_test, y_train, y_test = train_test_split(dataset.data, dataset.target, test_size=0.2)

    clf = GaussianNB()
    clf.fit(x_train, y_train)

    clf.predict(x_test)

    print("Test accuracy: %f" % clf.score(x_test, y_test))

    epsilons = np.logspace(-2, 2, 50)
    bounds = ([4.3, 2.0, 1.1, 0.1], [7.9, 4.4, 6.9, 2.5])
    accuracy = list()

    for epsilon in epsilons:
        clf = GaussianNB(bounds=bounds, epsilon=epsilon)
        clf.fit(x_train, y_train)

        accuracy.append(clf.score(x_test, y_test))

    plt.semilogx(epsilons, accuracy)
    plt.title("Differentially private Naive Bayes accuracy")
    plt.xlabel("epsilon")
    plt.ylabel("Accuracy")
    plt.show()


if __name__ == '__main__':
    main()
