import torch
from MetaEnvironment import Environment
from BayesianOpt.MetaBayesian import BayesianEnvironment
from DDPG.model import *
import MILPs.MILPModel_2layer as MILP
import MILPs.infoGainMILP_ChanceConstraints as MILPSafety
from BayesianOpt import BayesianAcquisitionFunction
from matplotlib import pyplot as pl
import MetaLearnBO.run_experiments as MBO
import time
from scipy.stats import norm
from sklearn.ensemble import RandomForestRegressor
import joblib

def calcDiff(prevA,prevO,a,gp):
    o=gp.predict(a)
    allDiffA=np.zeros((50,8))
    allDiffO = np.zeros((50, 8))
    for i in range(prevA.shape[0]):
        allDiffA[:,i]=(prevA[i]-a).reshape(-1,)
        allDiffO[:,i]=(prevO[i]-o).reshape(-1,)
    return np.hstack((allDiffA,allDiffO))

def calcU(gp,a):
    _,sigma=gp.predict(a,return_std=True)

    return np.array([sigma])

def findBestA(lal,actions,out,gp,step):
    numSamples=50
    possA=np.atleast_2d(np.linspace(0, 4, numSamples)).T
    uncertainty=calcU(gp,possA)
    diff=calcDiff(actions,out,possA,gp)
    print(possA.shape,diff.shape,uncertainty.shape,np.array([step]).repeat(numSamples).reshape(-1,1).shape)
    allFeatures= np.hstack((possA, diff, uncertainty.reshape(numSamples,1), np.array([step]).repeat(numSamples).reshape(-1,1)))
    r=lal.predict(allFeatures)


    i=np.argmax(np.array(r))
    return possA[i]



e = Environment(optimizeModel=False)
eBay = BayesianEnvironment(optimizeModel=False)
"""
numIters=4000
numSteps=8
##############RANDOM####################
DAErrorRandom2=np.zeros((numIters,numSteps+1))
ParamErrorRandom2=np.zeros((numIters,numSteps+1))

featNum=19
allFeatures=np.zeros((numIters*numSteps,featNum))
allError=np.zeros(numIters*numSteps)
c=0
for j in range(0,numIters):
    print("Iter",j)
    np.random.seed(j)
    ratN = np.random.randint(0, 5)
    print('rat',ratN)
    r_prev,error,DAError = eBay.reset(ratNum=ratN,seed=j)
    prevError=error

    for i in range(0,numSteps):
        np.random.seed(None)
        a=np.random.uniform(low=0.0, high=4.0, size=(1, 1)).reshape(1, 1)
        a = np.asarray(a).reshape(-1, 1)
        uncertainty=calcU(eBay.ratEst.rat,a)
        diff=calcDiff(eBay.ratEst.actions,eBay.ratEst.out,a,eBay.ratEst.rat)

        r1,error,bestParamError,DAError=eBay.step(a)
        allError[c]=r1
        allFeatures[c,:]=np.hstack((a,diff,uncertainty,np.array([i]).reshape(1,1)))

        prevError=error
        c+=1


parameters = {'est': 2000, 'depth': 40, 'feat': 6 }
lalModel1 = RandomForestRegressor(n_estimators = parameters['est'], max_depth = parameters['depth'],
                                 max_features=parameters['feat'], oob_score=True, n_jobs=8)

lalModel1.fit(allFeatures, np.ravel(allError))
print('Oob score = ', lalModel1.oob_score_)
joblib.dump(lalModel1, 'lalModel2.sav')
"""
lalModel1=joblib.load('lalModel2.sav')
numIters=10
numSteps=8

for j in range(0,numIters):
    np.random.seed(j)
    print("Iter",j)
    ratN = np.random.randint(0, 5)
    print("rat",ratN)
    r_prev,error,DAError = eBay.reset(ratNum=ratN,seed=j)
    prevError=error

    for i in range(0,numSteps):
        a =findBestA(lalModel1,eBay.ratEst.actions,eBay.ratEst.out,eBay.ratEst.rat,i)
        print("action",a)
        r1,error,bestParamError,DAError=eBay.step(a.reshape(1,1))
