# I: improt
from sklearn.datasets import load_svmlight_file
from sklearn.preprocessing import OneHotEncoder
from sklearn.preprocessing import Normalizer
import numpy as np
import matplotlib.pyplot as plt
from SS_ProxHD import SS_ProxHD
from SPIDER import SpiderMED_nonsmooth
from SS_ProxHD_EIM import SS_ProxHD_EIM

from loss import logistic_loss_nonconvex_nonsmooth, logistic_loss_nonconvex_gradient, Robust_linear_regression_loss_nonsmooth, Robust_linear_regression_gradient
# Compare SSPHD-EIM at different b

# II: load data
dataset_ = ['a9a', 'ijcnn1', 'w8a', 'covtype.libsvm.binary', 'gisette_scale', 'rcv1_train.binary']
# dataset_ = ['a9a']

save_file = True
plot_fig = False
#alg = "logistic"
alg = "linear"
for dataset in dataset_:
    experiment_name = 'test_3_linear_nonsmooth' + '_' + dataset
    opt = {}
    if dataset == 'a9a':
        X, Y = load_svmlight_file('data/a9a')
        n = X.shape[0]
        d = X.shape[1]
        # enc = OneHotEncoder()
        # Y = enc.fit_transform(Y.reshape(-1, 1))
        Y = [0 if e == -1 else e for e in Y]
        # c = Y.shape[1]
        max_iterations = 150
        opt['l1_weight'] = 0.01
        opt["batch_size"] = 256
        opt["tol_grad"] = -1
    elif dataset == 'w8a':
        X, Y = load_svmlight_file('data/w8a')
        n = X.shape[0]
        d = X.shape[1]
        # enc = OneHotEncoder()
        # Y = enc.fit_transform(Y.reshape(-1, 1))
        Y = [0 if e == -1 else e for e in Y]
        # c = Y.shape[1]
        max_iterations = 150
        opt['l1_weight'] = 0.01
        opt["batch_size"] = 256
        opt["tol_grad"] = -1
    elif dataset == 'ijcnn1':
        X, Y = load_svmlight_file('data/ijcnn1')
        n = X.shape[0]
        d = X.shape[1]
        # enc = OneHotEncoder()
        # Y = enc.fit_transform(Y.reshape(-1, 1))
        Y = [0 if e == -1 else e for e in Y]
        # c = Y.shape[1]
        max_iterations = 2000
        opt['l1_weight'] = 0.001
        opt["batch_size"] = 512
        opt["tol_grad"] = -1
    elif dataset == 'covtype.libsvm.binary':
        X, Y = load_svmlight_file('data/covtype.libsvm.binary')
        n = X.shape[0]
        d = X.shape[1]
        # enc = OneHotEncoder()
        # Y = enc.fit_transform(Y.reshape(-1, 1))
        Y = [0 if e == 2 else e for e in Y]
        # c = Y.shape[1]
        max_iterations = 1000
        opt['l1_weight'] = 0.00001
        opt["batch_size"] = 1024
        opt["tol_grad"] = -1
    elif dataset == 'gisette_scale':
        X, Y = load_svmlight_file('data/gisette_scale')
        n = X.shape[0]
        d = X.shape[1]
        # enc = OneHotEncoder()
        # Y = enc.fit_transform(Y.reshape(-1, 1))
        Y = [0 if e == -1 else e for e in Y]
        # c = Y.shape[1]
        max_iterations = 10
        opt['l1_weight'] = 0.1
        opt["batch_size"] = 256
        opt["tol_grad"] = -1
    elif dataset == 'rcv1_train.binary':
        X, Y = load_svmlight_file('data/rcv1_train.binary')
        n = X.shape[0]
        d = X.shape[1]
        # enc = OneHotEncoder()
        # Y = enc.fit_transform(Y.reshape(-1, 1))
        Y = [0 if e == -1 else e for e in Y]
        # c = Y.shape[1]
        max_iterations = 10
        opt['l1_weight'] = 0.01
        opt["batch_size"] = 512
        opt["tol_grad"] = -1
    '''
    min_max_scaler = MinMaxScaler()
    X = min_max_scaler.fit_transform(X.toarray())
    '''
    normalizer = Normalizer().fit(X.toarray())
    X = normalizer.transform(X.toarray())
    # X = X.toarray()
    # Y = Y.toarray()
    Y = np.array(Y)
    w = np.ones(d) * 2
    # w = np.ones(d) * 2
    # w = np.random.normal(0, 1, d)

    # III: loss function
    # loss_function = logistic_loss_nonconvex_nonsmooth
    # gradient = logistic_loss_nonconvex_gradient
    loss_function = Robust_linear_regression_loss_nonsmooth
    gradient = Robust_linear_regression_gradient

    # IV: algortihms
    opt["alpha"] = 0.1 # nonconvex regularizer
    opt["tau"] = 1e-10  # small term for log function in case of log(0)
    opt["eta"] = 0.01  # step_size
    opt["nonsmooth"] = True
    opt['momentum'] = 10


    flag_SpiderMED = False


    # save Data
    loss = []
    grads = []
    samples = []
    times = []
    list_params = []

    EID_gma = 1
    if (dataset in ["a9a", "rcv1_train.binary"]):
        b = [70]
        K = [18]
    elif (dataset in ["w8a", "ijcnn1"]):
        b = [110]
        K = [14]
    elif (dataset in ["covtype.libsvm.binary"]):
        b = [130]
        K = [22]
    elif (dataset in ["gisette_scale"]):
        b = [30]
        K = [20]
    for i in range(len(b)):
        print('\n', i, '\n')
        opt["batch_size"] = b[i]
        opt['max_iterations'] = K[i]
        (_loss, _grad, _sample, _time) = SS_ProxHD_EIM(w, X, Y, loss_function, gradient, beta=0.001, gamma=EID_gma, **opt)
        list_params.append(r'$\gamma$='+str(EID_gma)+r' b='+str(opt["batch_size"]))
        loss.append(_loss)
        grads.append(_grad)
        _sample = [e / n for e in _sample]
        samples.append(_sample)
        times.append(_time)

    EID_gma = 1.2
    if (dataset in ["a9a", "rcv1_train.binary"]):
        b = [10, 13, 16]
        K = [10, 14, 13]
    elif (dataset in ["w8a", "ijcnn1"]):
        b = [12, 16, 20]
        K = [16, 14, 10]
    elif (dataset in ["covtype.libsvm.binary"]):
        b = [18, 22, 26]
        K = [16, 15, 14]
    elif (dataset in ["gisette_scale"]):
        b = [2, 4, 6]
        K = [20, 16, 14]
    for i in range(len(b)):
        print('\n', i, '\n')
        opt["batch_size"] = b[i]
        opt['max_iterations'] = K[i]
        (_loss, _grad, _sample, _time) = SS_ProxHD_EIM(w, X, Y, loss_function, gradient, beta=0.001, gamma=EID_gma, **opt)
        list_params.append(r'$\gamma$='+str(EID_gma)+r' b='+str(opt["batch_size"]))
        loss.append(_loss)
        grads.append(_grad)
        _sample = [e / n for e in _sample]
        samples.append(_sample)
        times.append(_time)

    # save data in dat
    filename = "dat/Experiment_3" + experiment_name + ".dat"
    thefile = open(filename, 'w')
    for i in range(len(loss)):
     thefile.write("%s\n" % loss[i])
     thefile.write("%s\n" % grads[i])
     thefile.write("%s\n" % samples[i])
     thefile.write("%s\n" % times[i])
    thefile.close()


    # plot

    colors = ['#1B2631', '#C0392B', '#9B59B6', 'Maroon', '#1E8449', '#0343DF', '#E67E22', '#95A5A6', '#FF97F2',
              '#34495E']
    # markers = ['s', '8', '>', '<', 'P', '*', 'd', 'X']
    markers = [None, None, None, None, None, None]
    linestyles = ['-.', '-.', '-.', '-.', '-.', '-.']
    etalinestyles = ['--', '--', '--', '--', '--', '--']

    # gradient/epoch
    fig = plt.figure()
    plt.yscale('log')

    for i in range(len(list_params)):
        plt.plot(samples[i], grads[i], linestyles[i % 10], marker=markers[i % 10], color=colors[i % 10], linewidth=2,
                 markersize=12)
    plt.legend(list_params, fontsize=13, loc=1)
    plt.title(dataset)
    plt.ylabel("grad (log)", fontsize=18)
    plt.xlabel("number of epochs", fontsize=18)
    fig.savefig(f'result/test_3/{experiment_name}_Grad_eta{opt["eta"]}_L{opt["l1_weight"]}_epoch.pdf', format='pdf', dpi=1000)
    fig.savefig(f'result/test_3/{experiment_name}_Grad_eta{opt["eta"]}_L{opt["l1_weight"]}_epoch.png', bbox_inches='tight', dpi=1000)
    #plt.show()

    # gradient/s
    fig = plt.figure()
    plt.yscale('log')

    for i in range(len(list_params)):
        plt.plot(times[i], grads[i], linestyles[i % 10], marker=markers[i % 10], color=colors[i % 10], linewidth=2,
                 markersize=12)
    plt.legend(list_params, fontsize=13, loc=1)
    plt.title(dataset)
    plt.ylabel("grad (log)", fontsize=18)
    plt.xlabel("time(s)", fontsize=18)
    fig.savefig(f'result/test_3/{experiment_name}_Grad_eta{opt["eta"]}_L{opt["l1_weight"]}_s.pdf', format='pdf', dpi=1000)
    fig.savefig(f'result/test_3/{experiment_name}_Grad_eta{opt["eta"]}_L{opt["l1_weight"]}_s.png', bbox_inches='tight', dpi=1000)
    #plt.show()

    # plot loss/epoch
    fig = plt.figure()
    plt.yscale('log')

    # min_loss = loss[0][0]
    # small_term = 0
    # for i in range(len(list_params)):
    #     for e in loss[i]:
    #         if e < min_loss:
    #             min_loss = e
    # min_loss = min_loss - small_term
    min_loss = 0
    for i in range(len(list_params)):
        temp = [e - min_loss for e in loss[i]]
        loss[i] = temp
    for i in range(len(list_params)):
        plt.plot(samples[i], loss[i], linestyles[i % 10], marker=markers[i % 10], color=colors[i % 10], linewidth=2,
                 markersize=12)
    plt.legend(list_params, fontsize=13, loc=1)
    plt.title(dataset)
    plt.ylabel("$f-f^*$(log)", fontsize=18)
    plt.xlabel("number of epochs", fontsize=18)
    fig.savefig(f'result/test_3/{experiment_name}_Loss_eta{opt["eta"]}_L{opt["l1_weight"]}_epoch.pdf', format='pdf', dpi=1000)
    fig.savefig(f'result/test_3/{experiment_name}_Loss_eta{opt["eta"]}_L{opt["l1_weight"]}_epoch.png', bbox_inches='tight', dpi=1000)
    #plt.show()

    # plot loss/s
    fig = plt.figure()
    plt.yscale('log')

    # min_loss = loss[0][0]
    # small_term = 0
    # for i in range(len(list_params)):
    #     for e in loss[i]:
    #         if e < min_loss:
    #             min_loss = e
    # min_loss = min_loss - small_term
    min_loss = 0
    for i in range(len(list_params)):
        temp = [e - min_loss for e in loss[i]]
        loss[i] = temp
    for i in range(len(list_params)):
        plt.plot(times[i], loss[i], linestyles[i % 10], marker=markers[i % 10], color=colors[i % 10], linewidth=2,
                 markersize=12)
    plt.legend(list_params, fontsize=13, loc=1)
    plt.title(dataset)
    plt.ylabel("$f-f^*$(log)", fontsize=18)
    plt.xlabel("time(s)", fontsize=18)
    fig.savefig(f'result/test_3/{experiment_name}_Loss_eta{opt["eta"]}_L{opt["l1_weight"]}_s.pdf', format='pdf', dpi=1000)
    fig.savefig(f'result/test_3/{experiment_name}_Loss_eta{opt["eta"]}_L{opt["l1_weight"]}_s.png', bbox_inches='tight', dpi=1000)
    #plt.show()