from race import *
from utils import *
from storm import *
from baselines import *
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import sys
import scipy.stats
import pickle

#todo: train, val test split
#    : median + varience in the center function
def experimentSTORM(datasetName, dataset, Reps, P, method, Itr, alpha, ridge_lambda, init, REPS, avgcnt, hash_type):
    print ("STORM --------------------------------------------------------------------------------")
    print ("Method: ", method)

    Lt = []
    Ls = []
    sparseMem = []
    FinalRes = []
    for i in range(0,avgcnt):
        Results1 = stormRegression(datasetName, dataset, Reps, P, method, Itr, alpha, ridge_lambda, init, REPS, hash_type)
        FinalRes.append(np.asarray(Results1[1:]))
    FinalRes = np.asarray(FinalRes)
    FinalRes = np.swapaxes(FinalRes,0,2)
    print ("Results1: ", FinalRes[:,:,:])

    AllrepRes = []
    for i in range(len(Reps)):
        print ("min train loss: ", center(FinalRes[i,0,:]))#, " at theta: ", Results1[0][4][a])
        print ("val loss at min train loss: ", center(FinalRes[i,1,:]))#, " at theta: ", Results1[0][4][a])
        AllrepRes.append([center(FinalRes[i,0,:]), center(FinalRes[i,1,:]), FinalRes[i,2,0]])
    return AllrepRes

        # losses_train1 = np.array(Results1[0][5])
        # losses_test1 = np.array(Results1[0][6])
        #
        # # pickle.dump(Results, open("Results.p", "wb") )
        # losses_train1 = np.array(Results1[0][5])
        # losses_test1 = np.array(Results1[0][6])
        # a = np.argmin(losses_train1)
        # Lt.append(losses_train1[a])
        # Ls.append(losses_test1[a])
        # sparseMem.append(Results1[0][7])

    # print ("min train loss: ", center(Lt))#, " at theta: ", Results1[0][4][a])
    # # print ("min val loss for 1edit jump: ", losses_test1[b])#, " at theta: ", Results1[0][4][b])
    # print ("val loss at min train loss: ", center(Ls))#, " at theta: ", Results1[0][4][a])
    # print ("STROM theta is: ", Results1[0][4][a], mseLoss(Results1[0][4][a], dataset))
    # plotTheta(Results1[0][4][a], dataset[int(dataset.shape[0]*0.8):, :])
    # return [center(Lt), center(Ls), np.median(np.array(sparseMem))]

def experimentMean(datasetName, dataset):
    print("just mean--------------------------------------------------------------")
    [trainloss_M, testloss_M] = Mean(dataset)
    # print ( round(trainloss_M,3), round(testloss_M,3))
    return [center(trainloss_M), center(testloss_M)]

def experimentClarkson(datasetName, dataset, rho, avgcnt):
    print("Clarkson2009 regression coefficients-------------------------------------------------------------")
    Lt = []
    Ls = []
    for i in range(0, avgcnt):
        NI = dataset.shape[0]
        tr = int(NI*(0.8)) # 80% points in training
        ts = NI-tr

        dataset1 = dataset[:int(tr),:]
        d = dataset1.shape[1] -1
        m = math.ceil(rho*dataset1.shape[0])

        S = Clarkson2009(m,d)
        dataset1 = dataset1 / np.sqrt(d+2)
        theta_Ck = S.fit(dataset1)
        # print("Clarkson2009 sketch size (in terms of # floating point values)")
        # print(S.size())

        Lt.append(mseLoss(np.append(theta_Ck, -1), dataset1))
        Ls.append(mseLoss(np.append(theta_Ck, -1), dataset[tr:,:]))

    # print (center(Lt), center(Ls))
    return [center(Lt), center(Ls)]

def experimentLeastSquare(datasetName, dataset, avgcnt, ridge_lambda):
    print ("Least Square --------------------------------------------------------------------------------")
    # trainloss_ls_avg = 0
    # testloss_ls_avg = 0
    Lt = []
    Ls = []
    avgcnt = 3 # we don't need many exp for this
    for i in range(0, avgcnt):
        NI = dataset.shape[0]
        tr = math.ceil(NI*(0.8)) # 80% points in training
        ts = NI-tr

        X_train = dataset[:tr,:-1]
        y_train = dataset[:tr,-1]

        # create linear regression object
        if ridge_lambda>0:
            reg = Ridge(alpha=ridge_lambda)
        else:
            reg = linear_model.LinearRegression()
        reg.fit(X_train, y_train)
        theta = reg.coef_
        Lt.append(mseLoss(np.append(theta, -1), dataset[:tr,:]))
        Ls.append(mseLoss(np.append(theta, -1), dataset[tr:,:]))

    # print (theta)
    plotTheta(theta, dataset[int(dataset.shape[0]*0.8):, :])
    # print (center(Lt), center(Ls))
    return [center(Lt), center(Ls)]


def experimentRandomSampling(datasetName, dataset, rho, ridge_lambda, avgcnt):
    print ("Random sampling --------------------------------------------------------------------------------")
    Lt = []
    Ls = []
    for i in range(0, avgcnt):
        NI = dataset.shape[0]
        tr = math.ceil(NI*(0.8)) # 80% points in training
        ts = NI-tr
        Ns = math.ceil(tr*rho)

        dataset1 = dataset[:int(tr),:]
        randomSampler = Sampler(method = "RS")
        theta = randomSampler.fit(dataset1, Ns, ridge_lambda)
        Lt.append(mseLoss(np.append(theta, -1), dataset[:tr,:]))
        Ls.append(mseLoss(np.append(theta, -1), dataset[tr:,:]))

    plotTheta(theta, dataset[int(dataset.shape[0]*0.8):, :])
    # print (center(Lt), center(Ls))
    return [center(Lt), center(Ls)]


def experimentLeverageSampling(datasetName, dataset, rho, avgcnt):
    print ("leverage Sampling --------------------------------------------------------------------------------")
    Lt = []
    Ls = []
    for i in range(0, avgcnt):
        NI = dataset.shape[0]
        tr = math.ceil(NI*(0.8)) # 80% points in training
        ts = NI-tr
        Ns = math.ceil(tr*rho)

        dataset1 = dataset[:int(tr),:]
        randomSampler = Sampler(method = "LS")
        theta = randomSampler.fit(dataset1, Ns)
        Lt.append(mseLoss(np.append(theta, -1), dataset[:tr,:]))
        Ls.append(mseLoss(np.append(theta, -1), dataset[tr:,:]))

    plotTheta(theta, dataset[int(dataset.shape[0]*0.8):, :])
    # print (center(Lt), center(Ls))
    return [center(Lt), center(Ls)]
