import numpy as np
from QLearningEnv import boxEnvironment, robotAgent
import keras
from keras.models import Sequential
from keras.layers import Dense
from keras.regularizers import l2

class learningAgent():
    def __init__(self,tau=1, memoryMax= 1000000, alpha=0.5, gamma=0.7, updateFeq = 500, badMemory= 3000):
        self.alpha = alpha
        self.gamma = gamma 
        self.tau = tau
        self.a = np.array([-2,0,2])
        self.updateFeq = updateFeq
        self.memoryMax = memoryMax
        self.memoryCounter = -1
        self.badMemory = badMemory
        self.initPhase = True

    def setSensors(self, number=5):
        self.distanceSensorValue= np.zeros(max(number,5))
        self.dirtSensor = np.array([0])
        if self.memoryCounter<0:
            self.MemoryStates = np.inf*np.ones( (self.memoryMax,number+2) )
            self.MemoryActions = np.inf*np.ones(self.memoryMax)
            self.FeedbackDirt = np.inf*np.ones(self.memoryMax)
            self.FeedbackCol = np.zeros(self.memoryMax)
            self.Rewards = np.zeros(self.memoryMax)

    def _dumpSensorsData(self):
        if self.memoryCounter<1:
            dDiff = 0 
        else:
            dDiff = self.distanceSensorValue[2]/10 - self.MemoryStates[self.memoryCounter-1,2]
            dDiff = dDiff*10
        temp = np.hstack( (self.distanceSensorValue/10,self.dirtSensor*10**1,dDiff) )
        self.MemoryStates[self.memoryCounter,:] = temp

    def _dumpActionData(self,w):
        self.MemoryActions[self.memoryCounter] = w

    def collisionFeedback(self, yesOrNo): 
        self.FeedbackCol[self.memoryCounter] = int(yesOrNo)
        if np.min(self.MemoryStates[self.memoryCounter,:]) > 0.005:
            self.FeedbackCol[self.memoryCounter] = False

    def dirtCleaned(self,dirt): 
        if self.memoryCounter<0: return
        self.FeedbackDirt[self.memoryCounter] = dirt
        
    def calcReward(self):
        self.Rewards[self.memoryCounter] = 0.05
        self.Rewards[self.memoryCounter] += 10**4*self.FeedbackDirt[self.memoryCounter]
        self.Rewards[self.memoryCounter] -= self.FeedbackCol[self.memoryCounter]*10    
        self.Rewards[self.memoryCounter] -= 10*np.exp(- 5*self.distanceSensorValue[2]) 
        temp1 = 10*np.exp(- 5*self.distanceSensorValue[4])
        temp2 = 10*np.exp(- 5*self.distanceSensorValue[0])
        if self.MemoryActions[self.memoryCounter] > 0.001 : 
            self.Rewards[self.memoryCounter] -= temp1 
        elif self.MemoryActions[self.memoryCounter] < -0.001: 
            self.Rewards[self.memoryCounter] -= temp2
        else:
            self.Rewards[self.memoryCounter] -= (temp1 + temp2)/2
            
    def _initQ(self):
        XTrain = np.hstack( (self.MemoryStates[0:self.memoryCounter,:],self.MemoryActions[0:self.memoryCounter,None]))
        self.Q = Sequential()
        self.Q.add(Dense(256, input_dim=XTrain.shape[1], activation='relu',kernel_regularizer=l2(0.0001)))
        self.Q.add(Dense(128, activation='relu',kernel_regularizer=l2(0.0001)))
        self.Q.add(Dense(64, activation='relu',kernel_regularizer=l2(0.0001)))
        self.Q.add(Dense(1, activation='linear'))
        clist = [keras.callbacks.EarlyStopping(monitor='loss',patience=5,verbose=0,mode='auto')]
        self.Q.compile(loss='mean_squared_error', optimizer='adam')
        self.Q.fit(XTrain,self.alpha*self.Rewards[0:XTrain.shape[0]], epochs=50, callbacks=clist, batch_size=5, verbose=False)
            
    def _updateQ(self):
        learnSet  = np.arange(0,self.memoryCounter-1) 
        index1 = np.flatnonzero(self.Rewards[learnSet]<-0.9)
        if index1.shape[0]>self.badMemory:
            index = np.random.choice(index1.shape[0], self.badMemory, replace=False)
            index1 = index1[index]
        samplesleft = min(index1.shape[0]*5,self.memoryCounter-index1.shape[0]-1)
        index2 = np.random.choice(self.memoryCounter-1, samplesleft, replace=False)
        learnSet = np.hstack( (index1[None,:],index2[None,:]) )
        learnSet = np.reshape(learnSet, learnSet.shape[1])
        X = np.hstack( (self.MemoryStates[learnSet,:],self.MemoryActions[learnSet,None]))

        Qalt = np.squeeze(self.Q.predict(X))
        QChooose = np.zeros( (Qalt.shape[0],len(self.a)) )
        for i in range(len(self.a)): #*\label{code:qlearning:agent:0}
            a = self.a[i]*np.ones_like(self.MemoryActions[learnSet,None])
            Xtemp = np.hstack( (self.MemoryStates[learnSet+1,:],a))
            QChooose[:,i] = np.squeeze(self.Q.predict(Xtemp)) #*\label{code:qlearning:agent:1}
        Qmax = np.max(QChooose,axis=1) #*\label{code:qlearning:agent:2}
        r = self.Rewards[learnSet]
        QNeu = (1-self.alpha)*Qalt+self.alpha*(r + self.gamma*Qmax) #*\label{code:qlearning:agent:3}
        self.Q.fit(X,QNeu, epochs=3, batch_size=5, verbose=False) #*\label{code:qlearning:agent:4}
        
    def _softmax(self,Qa):
        return np.exp(Qa/self.tau) / np.sum(np.exp(Qa/self.tau)) 

    def chooseAction(self,Qa):
        toChoose = np.arange(0,len(Qa))
        pW = self._softmax(Qa)
        if np.any(np.isnan(pW)) or np.any(np.isinf(pW)):
            choosenA = np.random.randint(0,len(Qa))
        else:
            choosenA = np.random.choice(toChoose,replace=False, p=pW)
        return(choosenA)
    
    def changeAngularVelocity(self):
        self.calcReward() #*\label{code:qlearning:agent:112}

        self.memoryCounter += 1 
        if self.memoryCounter%self.updateFeq == 0 and not self.initPhase:
            self._updateQ()        
        if self.memoryCounter > min(500,self.memoryMax-1) and self.initPhase:
            self.initPhase = False
            self._initQ()
            self._updateQ()

        self._dumpSensorsData() #*\label{code:qlearning:agent:13}        
        if self.initPhase: #*\label{code:qlearning:agent:5}
            w = 4*(np.random.rand(1) - 0.5)
            if self.distanceSensorValue[1] < 0.75 or self.distanceSensorValue[3] < 0.75:
                w = 1.0 if self.distanceSensorValue[1] < self.distanceSensorValue[3] else -1.0
            if self.distanceSensorValue[0] < 0.5 or self.distanceSensorValue[4] < 0.5:
                w = 1.5 if self.distanceSensorValue[0] < self.distanceSensorValue[4] else -1.5
            if self.distanceSensorValue[2] < 0.5 : w = -2 #*\label{code:qlearning:agent:6}
        else:
            Qa = np.zeros_like(self.a)
            j = self.MemoryStates[self.memoryCounter,:].shape[0]
            x = np.zeros( (1, j+1) )
            x[0,0:j] = self.MemoryStates[self.memoryCounter,:] #*\label{code:qlearning:agent:7}
            for i in range(len(self.a)):
                x[0,self.MemoryStates[self.memoryCounter,:].shape[0]] = self.a[i] #*\label{code:qlearning:agent:8}
                Qa[i] = self.Q.predict(x)[0] #*\label{code:qlearning:agent:9}
            ca = self.chooseAction(Qa) #*\label{code:qlearning:agent:10}
            w  = self.a[ca] #*\label{code:qlearning:agent:11}        
        self._dumpActionData(w) #*\label{code:qlearning:agent:12}
        return(w)

    def CoopData(self,X,Y,nX,nY):
        pass

#    def addOldMemory(self,States,Actions,FeedbackDirt,FeedbackCol,counter):
#        self.setSensors(number=States.shape[1]-1)
#        if self.memoryCounter < 0 : self.memoryCounter = 0
#        self.MemoryStates[self.memoryCounter:self.memoryCounter+counter,:] = np.copy(States[0:counter,:])
#        self.MemoryActions[self.memoryCounter:self.memoryCounter+counter] = np.copy(Actions[0:counter])
#        self.FeedbackDirt[self.memoryCounter:self.memoryCounter+counter] = np.copy(FeedbackDirt[0:counter])
#        self.FeedbackCol[self.memoryCounter:self.memoryCounter+counter] = np.copy(FeedbackCol[0:counter])
#        self.memoryCounter += counter-1
#        self.calcReward(recomputeAll=True)
        
if __name__ == '__main__':
    np.random.seed(42)
    jack = learningAgent(tau=0.3, gamma=0.8, alpha=0.4,badMemory= 10000) 
    jack.updateFeq = 500
    for i in range(7):
        print('new world:',i)
        myNewBox = boxEnvironment()
        myNewBox.addRobot(jack)
        myNewBox.startSim(300)
        myNewBox.evaluation()
    myNewBox.plotScene()
