import numpy as np
import numpy.linalg as la
import mdp_solver
import simple_chain
import util
import birl

##test script for running gradient descent for bayesian inverse reinforcement learning
##domain is a simple markov chain (see simple_chain.py) 
##TODO I haven't incorporated a prior so this really is more of a maximum likelihood rather than bayesian irl algorithm

reward = [1,0,0,0,0,1] #true expert reward
terminals = [0] #no terminals, you can change this if you want
gamma = 0.9 #discount factor for mdp
chain = simple_chain.SimpleChain(reward, terminals, gamma) #create markov chain
print "expert reward"
print chain.reward
pi_star, V_star = mdp_solver.policy_iteration(chain) #solve for expert policy
print "expert policy"
print pi_star
print "expert value function"
print V_star
Q_star = mdp_solver.calc_qvals(chain, pi_star, V_star, gamma)
print "expert Q-values"
print Q_star

#give optimal action in each (non-terminal) state as demonstration
#we can test giving demonstrations in some but not all states, or even noisy demonstrations to see what happens if we want
demo = [(state, np.argmax(Q_star[state,:])) for state in range(chain.num_states) if state not in terminals]
print "demonstration", demo


####### gradient descent starting from random guess at expert's reward
reward_guess = np.reshape([np.random.randint(-10,10) for _ in range(chain.num_states)],(chain.num_states,1))

#create new mdp with reward_guess as reward
mdp = simple_chain.SimpleChain(reward_guess, terminals, gamma) #create markov chain
mdp_temp = mdp #temporal mdp used in BTLS

num_steps = 200
step_size = 1 # BTLS Initialization with step size = 1
alpha = 0.01
beta = 0.5 # BTLS variables, alpha is related to the constraint finishing each iteration, beta determines how fast step size will decrease
print "----- gradient descent ------"
for step in range(num_steps):
    flag = 0; # Variable that checks whether the BTLS iteration should finish or not
    #calculate optimal policy for current estimate of reward
    pi, V = mdp_solver.policy_iteration(mdp)
    #print "new policy"
    #print pi_star
    #calculate Q values for current estimate of reward
    Q = mdp_solver.calc_qvals(mdp, pi, V, gamma)
    #print "new Qvals"
    #print log-likelihood
    print "log-likelihood posterior", birl.demo_log_likelihood(demo, Q)

    #calculate gradient of posterior wrt reward
    grad = birl.calc_reward_gradient(demo, mdp, mdp.R, eta=1.0)
    while flag == 0:
        R_temp = mdp_temp.R + step_size * grad
        mdp_temp.set_reward(R_temp)
        pi_temp, V_temp = mdp_solver.policy_iteration(mdp_temp)
        Q_temp = mdp_solver.calc_qvals(mdp_temp, pi_temp, V_temp, gamma)
        func = birl.demo_log_likelihood(demo, Q_temp)
        approx = birl.demo_log_likelihood(demo,Q) + alpha * step_size * pow(la.norm(grad,2),2)
        if func >= approx:
            step_size = beta * step_size
    #update reward
    R_new = mdp.R + step_size * grad
    #print "new reward"
    #print R_new
    #update mdp with new reward 
    mdp.set_reward(R_new)

print "recovered reward"
print R_new
pi, V = mdp_solver.policy_iteration(mdp)
print "resulting optimal policy"
print pi
print "policy difference"
print np.linalg.norm(pi_star - pi)
