import numpy as np
from rl.policy.base_policy import Policy
from rl.util import log_self


class BoltzmannPolicy(Policy):

    '''
    The Boltzmann policy, where prob dist for selection
    p = exp(Q/tau) / sum(Q[a]/tau)
    '''

    def __init__(self, env_spec,
                 init_tau=5., final_tau=0.5, exploration_anneal_episodes=20,
                 **kwargs):  # absorb generic param without breaking
        super(BoltzmannPolicy, self).__init__(env_spec)
        self.init_tau = init_tau
        self.final_tau = final_tau
        self.tau = self.init_tau
        self.exploration_anneal_episodes = exploration_anneal_episodes
        self.clip_val = 100000
        log_self(self)

    def select_action(self, state):
        agent = self.agent
        state = np.reshape(state, (1, state.shape[0]))
        Q_state = agent.model.predict(state)[0]  # extract from batch predict
        assert Q_state.ndim == 1
        Q_state = Q_state.astype('float32')  # fix precision nan issue
        Q_state = Q_state - np.amax(Q_state)  # prevent overflow
        exp_values = np.clip(
            np.exp(Q_state / self.tau), -self.clip_val, self.clip_val)
        assert not np.isnan(exp_values).any()
        probs = np.array(exp_values / np.sum(exp_values))
        probs /= probs.sum()  # renormalize to prevent floating pt error
        action = np.random.choice(agent.env_spec['actions'], p=probs)
        return action

    def update(self, sys_vars):
        '''strategy to update tau in agent'''
        epi = sys_vars['epi']
        rise = self.final_tau - self.init_tau
        slope = rise / float(self.exploration_anneal_episodes)
        self.tau = max(slope * epi + self.init_tau, self.final_tau)
        return self.tau


class DoubleDQNBoltzmannPolicy(BoltzmannPolicy):

    '''
    Same as the Boltzmann policy but for a Double DQN agent
    '''

    def __init__(self, env_spec,
                 init_tau=5., final_tau=0.5, exploration_anneal_episodes=20,
                 **kwargs):  # absorb generic param without breaking
        super(DoubleDQNBoltzmannPolicy, self).__init__(env_spec)

    def select_action(self, state):
        agent = self.agent
        state = np.reshape(state, (1, state.shape[0]))
        Q_state1 = agent.model.predict(state)[0]  # extract from batch predict
        Q_state2 = agent.model_2.predict(state)[0]  # extract from batch predict
        Q_state = Q_state1 + Q_state2
        assert Q_state.ndim == 1
        Q_state = Q_state.astype('float32')  # fix precision nan issue
        Q_state = Q_state - np.amax(Q_state)  # prevent overflow
        exp_values = np.clip(
            np.exp(Q_state / self.tau), -self.clip_val, self.clip_val)
        assert not np.isnan(exp_values).any()
        probs = np.array(exp_values / np.sum(exp_values))
        probs /= probs.sum()  # renormalize to prevent floating pt error
        action = np.random.choice(agent.env_spec['actions'], p=probs)
        return action

@@ -2094,24 +2094,109 @@
 _init__(
-env_spec
+%0A            env_spec, init_tau, final_tau,%0A            exploration_anneal_episodes, **kwargs
 )%0A%0A    d
@@ -2318,50 +2318,8 @@
     
-  Q_state1 = agent.model.predict(state)%5B0%5D
   # 
@@ -2360,17 +2360,17 @@
  Q_state
-2
+1
  = agent
@@ -2379,10 +2379,8 @@
 odel
-_2
 .pre
@@ -2397,38 +2397,59 @@
 )%5B0%5D
+%0A
   
-# extract from batch predict
+      Q_state2 = agent.model_2.predict(state)%5B0%5D
 %0A   

fix param missing from DoubleDQNBoltzmannPolicy