import pickle
import numpy as np
import random
from sklearn.utils import shuffle
from sklearn import preprocessing
import matplotlib.pyplot as plt
import sys

epsilon = [0.5]
G = 1
C = 4
L = 1

experiments = 5

def sigmoid(N, data, weight):
    x = np.dot(data, weight)
    result = np.zeros(N, )
    for i in range(N):
        if x[i] >= 0:
            result[i] = 1. / (1 + np.exp(-x[i]))
        else:
            result[i] = np.exp(x[i]) / (1 + np.exp(x[i]))
    return result

def sigmoid2(data, weight):
    x = np.dot(data, weight)
    if x >= 0:
        return 1. / (1 + np.exp(-x))
    else:
        return np.exp(x) / (1 + np.exp(x))

def classifier(x, weights):
    prob = sigmoid2(x, weights)
    if prob > 0.5:
        return 1.0
    else:
        return 0.0

def cost(N, prob, labels):
    a = labels * np.log(prob + 1e-10)
    b = (1 - labels) * np.log(1 - prob + 1e-10)
    c = a + b
    return -1 / N * np.sum(c)


def deminsion(argv):
    if argv == 'CC':
        X, y = pickle.load(open('CreditCard_data.p', 'rb'))

        train_data_num = 800

        eta = 0.1
        T = 100

        seed = 1

    elif argv == 'Bank':
        X, y = pickle.load(open('bank_data.p', 'rb'))

        train_data_num = 25000

        eta = 0.1
        T = 100

        seed = 0

    elif argv == 'Adult':
        X, y = pickle.load(open('adult_data.p', 'rb'))

        for id in range(y.shape[0]):
            if y[id] == -1:
                y[id] = 0

        train_data_num = 30162

        eta = 0.05
        T = 200

        seed = 0

    n, d = X.shape[0], X.shape[1]

    X, y = shuffle(X, y, random_state=0)
    X = preprocessing.scale(X)

    random.seed(seed)
    np.random.seed(seed)

    # Normalized Gradient Perturbation
    print('Normalized Gradient Perturbation......')

    acc_nor_mean, acc_nor_var = [], []

    for c in range(0, 9):

        md = np.zeros((n, c * d))
        newX = np.concatenate((X, md), axis=1)

        train_X = newX[0:train_data_num]
        train_y = y[0:train_data_num]

        test_X = newX[train_data_num:]
        test_y = y[train_data_num:]
        test_N = test_X.shape[0]

        trainn, newd = train_X.shape[0], train_X.shape[1]
        delta = 1 / trainn
        w_init = np.ones(newd)

        print('Dimension: ', newd, 'Processing...')

        for e in epsilon:
            std = C * G * (T ** 0.5) * (np.log(1 / delta) ** 0.5) / (trainn * e)
            acc_nor = []
            for experiment in range(experiments):
                w = w_init
                for t in range(T):
                    h = sigmoid(trainn, train_X, w)
                    error = h - train_y
                    error = error / trainn
                    noise = np.random.normal(loc=0.0, scale=std, size=newd)
                    gradient = np.matmul(train_X.transpose(), error) + L * w
                    gradient_l2 = np.linalg.norm(gradient, ord=2)
                    if gradient_l2 < 1:
                        gradient = gradient / gradient_l2

                    w = w - eta * (gradient + noise)

                error_item = 0.0
                for j in range(test_N):
                    predict = classifier(test_X[j], w)
                    if int(predict) != int(test_y[j]):
                        error_item += 1

                acc_nor.append(1 - error_item / test_N)
            acc_nor_mean.append(np.mean(acc_nor))
            acc_nor_var.append(np.var(acc_nor))
    print('Mean: ', acc_nor_mean)
    print('Var:', acc_nor_var)

    p = [i * d for i in range(1, 10)]
    color = ['red']
    line = ['-']
    marker = ['o']
    plt.plot(p, acc_nor_mean, color=color[0], linestyle=line[0], lw=3, markersize=10, marker=marker[0], label='m-NGP')
    plt.xlabel('Dimension $p$')
    plt.ylabel('Accuracy')
    plt.legend(['m-NGP'])
    plt.xticks(p)
    plt.show()
    print('################################')


if __name__ == "__main__":
    deminsion(sys.argv[1])
