"""
Template for implementing QLearner  (c) 2015 Tucker Balch
"""

import numpy as np
import random as rand

class QLearner(object):

    def __init__(self, \
        num_states=100, \
        num_actions = 4, \
        alpha = 0.2, \
        gamma = 0.9, \
        rar = 0.5, \
        radr = 0.99, \
        dyna = 10, \
        verbose = False):

        self.verbose = verbose
        self.num_actions = num_actions
        self.num_states=num_states
        self.actions=range(0,num_actions)
        self.Q = np.random.uniform(low=-1,high=1,size=(num_states,num_actions))
        self.Tc = 0.00001*np.ones(shape=(self.num_states, self.num_actions, self.num_states))
        self.T = np.zeros(shape=(self.num_states, self.num_actions, self.num_states))
        self.R = np.zeros(shape=(self.num_states, self.num_actions))
        self.gamma=gamma
        self.alpha=alpha
        self.rar=rar
        self.radr=radr
        self.dyna=dyna
        self.s = 0
        self.a = 0
        

    def querysetstate(self, s):
        """
        @summary: Update the state without updating the Q-table
        @param s: The new state
        @returns: The selected action
        """
        self.s = s
        prand = np.random.random()
        if prand < self.rar:
        	action = np.random.randint(0, self.num_actions-1)
        else:
        	action=self.Q[s, :].argmax()
        self.a=action
        self.rar = self.rar * self.radr
        	
        if self.verbose: print "s =", s,"a =",action
        return action

    def query(self,s_prime,r):
        """
        @summary: Update the Q table and return an action
        @param s_prime: The new state
        @param r: The ne state
        @returns: The selected action
        """    	
        '''update q'''
        maxqnew=self.Q[s_prime].max()
        value=r+self.gamma*maxqnew
        update=self.alpha*(value-self.Q[self.s,self.a])
        self.Q[self.s,self.a]+=update
        '''throw dice'''
        prand = np.random.random()
        if prand < self.rar:
        	action = np.random.randint(0, self.num_actions-1)
        else:
        	action=self.Q[s_prime, :].argmax()
        self.rar = self.rar * self.radr
        '''dyna part'''
        self.Tc[self.s, self.a, s_prime]+=1
        	
        self.T[self.s, self.a, s_prime]=self.Tc[self.s, self.a, s_prime]/np.sum(self.Tc[self.s, self.a,:], axis=-1)
        self.R[self.s, self.a]=(1-self.alpha)*self.R[self.s, self.a]+self.alpha*r
        for i in range(0,self.dyna):
        	rand_s=np.random.randint(0, self.num_states-1)
        	rand_a=np.random.randint(0, self.num_actions-1)
        	#draw s' from T[s,a,:]
        	rand=np.random.uniform()
        	prob=0
        	rand_s_prime=0
        	for j in range(0, self.num_states):
        		prob+=self.T[rand_s, rand_a, j]
        		if prob>=rand:
        			rand_s_prime=j
        			break
        	rand_r=self.R[rand_s, rand_a]
        	'''update Q inside dyna loop'''
        	maxqnew=self.Q[rand_s_prime].max()
        	value=rand_r+self.gamma*maxqnew
        	update=self.alpha*(value-self.Q[rand_s,rand_a])
        	self.Q[rand_s,rand_a]+=update
        	rand_a=self.Q[rand_s_prime, :].argmax()
        	rand_s=rand_s_prime
        
        #self.verbose=False
        if self.verbose: print "s =", s_prime,"a =",action,"r =",r
        self.s=s_prime
        self.a=action
        return action

if __name__=="__main__":
    print "Remember Q from Star Trek? Well, this isn't him"
