# from utils import generate_rationality_dirichlet, generate_params, generate_theta, compute_params, ERM, DRO, gradient_descent
from utils import *
import sys
from math import inf
import os
import pandas as pd
from sklearn.cluster import KMeans
import pickle
import random
from copy import deepcopy
random.seed(0)

dset = pd.read_csv('./data/NYC_82341_59_5_1.data')
dset = dset.to_numpy()
full_data = []
U = []
for i in range(dset.shape[0]):
    temp = dset[i][0].split()
    if(i == 0):
        full_data.append([int(j) for j in temp])
    elif(i == 1):
        U.append([float(j) for j in temp])
        # full_data.append([0.0 for j in temp])
        full_data.append([float(j) for j in temp])
    else:
        full_data.append([float(temp[j]) for j in range(len(temp))])
        # full_data.append([float(temp[j]) - U[0][j] for j in range(len(temp))])
full_data = np.array(full_data)
# print(np.max(np.exp(full_data[2:,])))
# print(full_data[1:,0])
# print(np.max(full_data[2:,0]))
# print(np.min(full_data[2:,0]))
# print(U[0][:10])
# for i in range(10):
#     print(full_data[4,i] - U[])
# print(full_data[:,1])

trans_data = full_data.T
# Filter data
# count = 0
# upper = 1.0
# lower = -1.0
for upper in [2.0]:
    for lower in [0.0]:
        if(upper == lower):
            continue
        count = 0
        filtered_data = []
        for i in range(len(trans_data)):
            u = trans_data[i][1]
            v_max = np.max(trans_data[i,2:])
            if (v_max - u >= lower and v_max - u <= upper):
                filtered_data.append(trans_data[i])
                count += 1
        print(count)
        filtered_data = np.array(filtered_data)
        # Shuffle full_data to form train-test split
        N_train = int(0.8 * count)
        random.shuffle(filtered_data)
        # print(filtered_data[0])
        full_data = filtered_data.T
        test_data = full_data[:,N_train:]
        full_data = full_data[:,:N_train]

        features = full_data[1:].T
        weights = np.array(full_data[0])
        test_features = test_data[1:].T 
        test_weights = np.array(test_data[0])
        # print(np.max(features))
        # print(np.min(features))
        N_act = len(weights)
        print(N_act)
        # scale = 1.0
        # features = np.exp(-features/scale)
        # test_features = np.exp(-test_features/scale)
        features = features 
        test_features = test_features 

        # # log = {}
        # # all_losses = []
        # # log_file = './kmeans/exp{}_scale.npy'.format(scale)
        # # for N in range(50,10001,50):
        N = 50
        kmeans_loc = './kmeans/filtered_kmeans_{}clusters_{}upper_{}lower_{}train'.format(N, upper, lower, N_train)
        # kmeans_loc = './kmeans/another_kmeans_{}clusters'.format(N)
        if(os.path.isfile(kmeans_loc)):
            print("Using Saved K-Means")
            kmeans = pickle.load(open(kmeans_loc, 'rb'))
        else:
            print("Running new K-Means")
            kmeans = KMeans(n_clusters=N).fit(features,sample_weight=weights)
            pickle.dump(kmeans, open(kmeans_loc, 'wb'))
            print("Saved K-Means")

        #     # print ("N : ",N," Inertia : ",kmeans.inertia_)
        # #     all_losses.append(kmeans.inertia_)
        # # log['losses'] = all_losses
        # # np.save(log_file, log)
        # # print(all_losses)

        # # # # print(kmeans.inertia_)
        C = kmeans.cluster_centers_

        # To compute fractions


        # print(C[0])
        # L = []
        # for i in range(N):
        #     L.append(np.max(C[i][1:]) - C[i][0])
        # L.sort()
        # print(L)
        # Update new counts
        Y = kmeans.predict(features)
        s = np.zeros(N+1)
        for i in range(len(Y)):
            # s[Y[i]] += weights[Y[i]]
            s[Y[i]] += weights[i]

        print(s)
        # print(C)
        for i in range(len(C)):
            # C[i,:] = C[i,:] - np.max(C[i,:])
            C[i,:] = C[i,:] - C[i,0]

        for i in range(len(test_features)):
            # test_features[i,:] = test_features[i,:] - np.max(test_features[i,:])
            test_features[i,:] = test_features[i,:] - test_features[i,0]

        # print(C[0])
        # print(C[1])
        # print(C[2])
        # print(C[3])

        m = 59
        # L = []
        # for i in range(1,m+1):
        #     L.append(np.mean(C[:,i]))
        # L.sort()
        # print(L)

        # # Prepare Constants
        U = np.exp(C[:,0])
        V = np.exp(C[:,1:])
        U_test = np.exp(test_features[:,0])
        V_test = np.exp(test_features[:,1:])
        print("Range U")
        print(np.max(U))
        print(np.min(U))
        print("Range V")
        print(np.max(V))
        print(np.min(V))

        tL, tU = upper_lower_t(V, U)
        print("Range t")
        print(np.max(tU))
        print(np.min(tL))
        print(tL)


        X = np.ones(m)

        temp = FLP_values(X, V, U)
        avg_d, std_d = weighted_avg_and_std(temp, s[:-1])
        print(avg_d)
        print(std_d)

        print(weights)
        print(sum(s))
        print(N_train)
        for xi in [1e2, 1e3, 1e4]:
            for M in range(10, 15, 2):
        # for xi in [1e3]:
        #     for M in [15]:
                # print(sum(test_weights))
                print(xi, " ", M)
                # print(M)
                log = {}
                # xi = 1e3
                # xi = 0.01 * N_act * (N_act - 1)
                X_DRO = FLP_DRO(m, N, sum(s), M, tL, tU, V, 0, V, U, s, xi)
                X_ERM = FLP_DRO(m, N, sum(s), M, tL, tU, V, 0, V, U, s, 0)
                # X_ERM = FLP_ERM(m, N, N_train, M, tL, tU, V, 0, V, U, s[:-1])
                X_DRO = [round(i) for i in X_DRO]
                X_ERM = [round(i) for i in X_ERM]
                print(sum(X_DRO))
                print(sum(X_ERM))
                print(X_DRO)
                print(X_ERM)
                log_file = './logs/filtered{}N_{}upper_{}lower_{}train_FLP_weights_M{}_xi{}'.format(N, upper, lower, N_train, M,xi)
                log['weights'] = [X_DRO, X_ERM]
                print("Train")
                values_DRO = np.array(FLP_values(X_DRO, V, U))
                values_ERM = np.array(FLP_values(X_ERM, V, U))
                avg_d, std_d = weighted_avg_and_std(values_DRO, s[:-1])
                avg_e, std_e = weighted_avg_and_std(values_ERM, s[:-1])
                print(avg_d)
                print(avg_e)
                print(std_d)
                print(std_e)
                print("Test")
                values_DRO = np.array(FLP_values(X_DRO, V_test, U_test))
                values_ERM = np.array(FLP_values(X_ERM, V_test, U_test))
                arg_d = np.argsort(values_DRO)
                arg_e = np.argsort(values_ERM)
                l = int(0.05 * (count - N_train))
                avg_d, std_d = weighted_avg_and_std(values_DRO[arg_d[:l]], test_weights[arg_d[:l]])
                avg_e, std_e = weighted_avg_and_std(values_ERM[arg_e[:l]], test_weights[arg_e[:l]])
                print(avg_d)
                print(avg_e)
                avg_d, std_d = weighted_avg_and_std(values_DRO, test_weights)
                avg_e, std_e = weighted_avg_and_std(values_ERM, test_weights)
                print(avg_d)
                print(avg_e)
                print(std_d)
                print(std_e)
                np.save(log_file, log)




