import numpy as np
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import RBF, ConstantKernel as C

import numpy as np
import pandas as pd
from sklearn.gaussian_process.kernels import (RBF, Matern, RationalQuadratic,
                                              ExpSineSquared, DotProduct,
                                              ConstantKernel)
from matplotlib import pyplot as plt
import joblib


class Simulation:
    def __init__(self,allRats, ratNum=0,seed=None):
        ##initialize GP


        cRat=allRats[ratNum]
        x = np.atleast_2d(np.linspace(0, 4, 100)).T
        y_pred=cRat.sample_y(x,n_samples=2,random_state=seed)
        y_pred=np.array(y_pred)
        y_avg=np.mean(y_pred,axis=1)


        alphas=[.55,.09,.05,.05,.2,.35]
        kernel = 1.0 * ExpSineSquared(length_scale=1.0, periodicity=3.0,
                                       length_scale_bounds=(0.1, 10.0),
                                       periodicity_bounds=(1.0, 10.0))
        gp = GaussianProcessRegressor(kernel=kernel, n_restarts_optimizer=9, alpha=alphas[ratNum],random_state=3)
        gp.fit(x, y_avg)


        self.gp=gp
        self.x_groundTruth=x
        self.y_groundTruth=y_avg

        ###Best Params
        x = np.atleast_2d(np.linspace(0, 4, 500)).T
        y_pred,sigma=gp.predict(x, return_std=True)
        obj=np.multiply(x.reshape(500,),y_pred)
        i=np.argmax(obj)
        self.bestP=x[i]
        self.ratNum = ratNum
        #print("BESTP",self.bestP)

        """plt.figure()
        plt.plot(x, np.multiply(x.reshape(500,), y_pred))
        plt.title('Objective')

        plt.figure()
        plt.plot(x, y_pred, 'b-', label='Prediction')

        plt.fill(np.concatenate([x, x[::-1]]),
                 np.concatenate([y_pred - 1.9600 * sigma,
                                 (y_pred + 1.9600 * sigma)[::-1]]),
                 alpha=.5, fc='b', ec='None', label='95% confidence interval', )
        plt.show()
        """










    def get_RatModel(self):

        return self.gp


    def simulateRat(self,action):

        out=self.gp.predict(action)

        """x = np.atleast_2d(np.linspace(0, 4, 500)).T
        y_pred, sigma = self.gp.predict(x, return_std=True)
        plt.figure()
        plt.plot(x, np.multiply(x.reshape(500, ), y_pred))
        plt.title('Objective')
        plt.scatter(action,np.multiply(action.reshape(-1, ), out))
        plt.show()
        """

        return out

    def predictRat(self,action):
        out=self.gp.predict(action)

        return out



    def updateState(self,state):
        self.currentState=state

class CanonicalRats:
    def __init__(self):
        #file = 'C:/Users/mschr/CORELABProjects/RNS-Meta/Data/Data_Memory.csv'
        """file = 'E:\ResearchGatech\RNS-Meta\Data\Data_Memory.csv'
        data = pd.read_csv(
            file, sep=',')



        X = data['amplitude']
        X = np.array(X)
        y = data['discriminant_score_b']
        y = np.array(y)

        X1 = X[0:25].reshape(-1, 1)  # .45
        y1 = y[0:25].reshape(-1)
        kernel1 = 1.0 * ExpSineSquared(length_scale=1.0, periodicity=3.0,
                                      length_scale_bounds=(0.1, 10.0),
                                      periodicity_bounds=(1.0, 10.0))
        self.rat1 = GaussianProcessRegressor(kernel=kernel1, n_restarts_optimizer=9, alpha=.45,random_state=2)
        self.rat1.fit(X1, y1)
        joblib.dump(self.rat1, 'C:/Users/mschr/CORELABProjects/NeuripsFinalCode/RNS-Meta/Data/rat1.sav')
        self.rat1=joblib.load('C:/Users/mschr/CORELABProjects/NeuripsFinalCode/RNS-Meta/Data/rat1.sav')

        X2 = X[26:50].reshape(-1, 1)  # .09
        y2 = y[26:50].reshape(-1)
        kernel2 = 1.0 * ExpSineSquared(length_scale=1.0, periodicity=3.0,
                                      length_scale_bounds=(0.1, 10.0),
                                      periodicity_bounds=(1.0, 10.0))
        self.rat2 = GaussianProcessRegressor(kernel=kernel2, n_restarts_optimizer=9, alpha=.09,random_state=3)
        self.rat2.fit(X2, y2)
        joblib.dump(self.rat2, 'C:/Users/mschr/CORELABProjects/NeuripsFinalCode/RNS-Meta/Data/rat2.sav')

        X3 = X[51:74].reshape(-1, 1)  # .05
        y3 = y[51:74].reshape(-1)
        kernel3 = 1.0 * ExpSineSquared(length_scale=1.0, periodicity=3.0,
                                      length_scale_bounds=(0.1, 10.0),
                                      periodicity_bounds=(1.0, 10.0))
        self.rat3 = GaussianProcessRegressor(kernel=kernel3, n_restarts_optimizer=9, alpha=.05,random_state=3)
        self.rat3.fit(X3, y3)
        joblib.dump(self.rat3, 'C:/Users/mschr/CORELABProjects/NeuripsFinalCode/RNS-Meta/Data/rat3.sav')

        X4 = X[75:121].reshape(-1, 1)  # .05
        y4 = y[75:121].reshape(-1)
        kernel4 = 1.0 * ExpSineSquared(length_scale=1.0, periodicity=3.0,
                                      length_scale_bounds=(0.1, 10.0),
                                      periodicity_bounds=(1.0, 10.0))
        self.rat4 = GaussianProcessRegressor(kernel=kernel4, n_restarts_optimizer=9, alpha=.05,random_state=3)
        self.rat4.fit(X4, y4)
        joblib.dump(self.rat4, 'C:/Users/mschr/CORELABProjects/NeuripsFinalCode/RNS-Meta/Data/rat4.sav')

        X5 = X[122:148].reshape(-1, 1)  # .2
        y5 = y[122:148].reshape(-1)
        kernel5 = 1.0 * ExpSineSquared(length_scale=1.0, periodicity=3.0,
                                      length_scale_bounds=(0.1, 10.0),
                                      periodicity_bounds=(1.0, 10.0))
        self.rat5 = GaussianProcessRegressor(kernel=kernel5, n_restarts_optimizer=9, alpha=.2,random_state=3)
        self.rat5.fit(X5, y5)
        joblib.dump(self.rat5, 'C:/Users/mschr/CORELABProjects/NeuripsFinalCode/RNS-Meta/Data/rat5.sav')

        X6 = X[150:174].reshape(-1, 1)  # .35
        y6 = y[150:174].reshape(-1)
        kernel6 = 1.0 * ExpSineSquared(length_scale=1.0, periodicity=3.0,
                                      length_scale_bounds=(0.1, 10.0),
                                      periodicity_bounds=(1.0, 10.0))
        self.rat6 = GaussianProcessRegressor(kernel=kernel6, n_restarts_optimizer=9, alpha=.35,random_state=6)
        self.rat6.fit(X6, y6)
        joblib.dump(self.rat6, 'C:/Users/mschr/CORELABProjects/NeuripsFinalCode/RNS-Meta/Data/rat6.sav')
        """

        self.rat1 = joblib.load('C:/Users/mschr/CORELABProjects/NeuripsFinalCode/RNS-Meta/Data/rat1.sav')
        self.rat2 = joblib.load('C:/Users/mschr/CORELABProjects/NeuripsFinalCode/RNS-Meta/Data/rat2.sav')
        self.rat3 = joblib.load('C:/Users/mschr/CORELABProjects/NeuripsFinalCode/RNS-Meta/Data/rat3.sav')
        self.rat4 = joblib.load('C:/Users/mschr/CORELABProjects/NeuripsFinalCode/RNS-Meta/Data/rat4.sav')
        self.rat5 = joblib.load('C:/Users/mschr/CORELABProjects/NeuripsFinalCode/RNS-Meta/Data/rat5.sav')
        self.rat6 = joblib.load('C:/Users/mschr/CORELABProjects/NeuripsFinalCode/RNS-Meta/Data/rat6.sav')


    def getCanonicalRat(self):
        return [self.rat1,self.rat2,self.rat3,self.rat4,self.rat5,self.rat6]
