# I: improt
from sklearn.datasets import load_svmlight_file
from sklearn.preprocessing import OneHotEncoder
from sklearn.preprocessing import MinMaxScaler, Normalizer
import numpy as np
import matplotlib.pyplot as plt
import json_func
from SS_ProxHD import SS_ProxHD_return_eta

from loss import logistic_loss_nonconvex_nonsmooth, logistic_loss_nonconvex_gradient, Robust_linear_regression_loss_nonsmooth, Robust_linear_regression_gradient
# Compare SSPHD at different initial learning rates eta0

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

plot_fig = True
save_fig = True
test_loss = True
small_term = 1e-2
# alg = "logistic"
alg = "linear"
for dataset in dataset_:
    experiment_name = 'test_2_different_hyperlearning_rates_logistic_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 = 500
        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 = 500
        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 = 2800
        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 = 300
        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 = 15
        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 = 35
        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
    if alg == "logistic":
        loss_function = logistic_loss_nonconvex_nonsmooth
        gradient = logistic_loss_nonconvex_gradient
    elif alg == "linear":
        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["inner_loop_m"] = 30
    opt["eta"] = 0.05  # step_size
    opt["nonsmooth"] = True
    opt['momentum'] = 10

    opt['max_iterations'] = max_iterations

    flag_SS_ProxHD = True
    eta_list = [0.1, 0.05, 0.01, 1e-4, 1e-10, 1e-14]

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

    file_path = "result/test_2/json/"
    if flag_SS_ProxHD == True:
        for eta in eta_list:
            opt["eta"] = eta
            list_params.append(r'eta_0=' + str(eta))
            file_name = f'{alg}_{dataset}_{list_params[-1]}_eta1e-14_m{opt["inner_loop_m"]}_b{opt["batch_size"]}_L{opt["l1_weight"]}.json'
            #print(file_name)
            json_data = json_func.read_json(file_path, file_name)
            _loss = json_data['loss'][:max_iterations + 1]
            _grad = json_data['grads'][:max_iterations + 1]
            _sample = json_data['samples'][:max_iterations + 1]
            _time = json_data['times'][:max_iterations + 1]
            _eta = json_data['eta0'][:max_iterations + 1]
            list_params[-1] = r'$eta_0$=' + str(eta)
            loss.append(_loss)
            grads.append(_grad)
            # _sample = [e / n for e in _sample]
            samples.append(_sample)
            times.append(_time)
            eta0.append(_eta)

    # plot

    colors = ['red', 'orange', 'gold', 'fuchsia', 'pink', 'slateblue', 'royalblue']
    # markers = ['s', '8', '>', '<', 'P', '*', 'd', 'X']
    markers = [None, None, None, None, None, None, None]
    linestyles = ['-.', '-.', '-.', '-.', '-.', '-.', '-.']
    etalinestyles = ['--', '--', '--', '--', '--', '--', '--']
    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='best')
    plt.title(dataset)
    plt.ylabel("grad (log)", fontsize=18)
    plt.xlabel("number of epochs", fontsize=18)
    if plot_fig == True:
        plt.show()
    if save_fig == True:
        fig.savefig(f'result/test_2/{experiment_name}_Grad_m{opt["inner_loop_m"]}_b{opt["batch_size"]}_L{opt["l1_weight"]}.pdf', format='pdf', dpi=1000)
        fig.savefig(f'result/test_2/{experiment_name}_Grad_m{opt["inner_loop_m"]}_b{opt["batch_size"]}_L{opt["l1_weight"]}.png', bbox_inches='tight', dpi=1000)


    # plot loss
    fig = plt.figure()
    plt.yscale('log')
    if test_loss == False:
        min_loss = loss[0][0]
        small_term = 1e-25
        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
    else:
        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='best')
    plt.title(dataset)
    plt.ylabel("$f-f^*$(log)", fontsize=18)
    plt.xlabel("number of epochs", fontsize=18)
    if plot_fig == True:
        plt.show()
    if save_fig == True:
        fig.savefig(f'result/test_2/{experiment_name}_Loss_m{opt["inner_loop_m"]}_b{opt["batch_size"]}_L{opt["l1_weight"]}.pdf', format='pdf', dpi=1000)
        fig.savefig(f'result/test_2/{experiment_name}_Loss_m{opt["inner_loop_m"]}_b{opt["batch_size"]}_L{opt["l1_weight"]}.png', bbox_inches='tight', dpi=1000)


    # plot eta
    fig = plt.figure()

    for i in range(len(list_params)):
        plt.plot(samples[i], eta0[i], etalinestyles[i % 10], marker=markers[i % 10], color=colors[i % 10], linewidth=1,
                     markersize=8)
    plt.legend(list_params, fontsize=13, loc='best')
    plt.title(dataset)
    plt.ylabel("eta", fontsize=18)
    plt.xlabel("number of epochs", fontsize=18)
    if plot_fig == True:
        plt.show()

    if save_fig == True:
        fig.savefig(f'result/test_2/{experiment_name}_ETA_m{opt["inner_loop_m"]}_b{opt["batch_size"]}_L{opt["l1_weight"]}.pdf', format='pdf', dpi=1000)
        fig.savefig(f'result/test_2/{experiment_name}_ETA_m{opt["inner_loop_m"]}_b{opt["batch_size"]}_L{opt["l1_weight"]}.png', bbox_inches='tight', dpi=1000)