import numpy as np
from rl.agent.dqn import DQN
from rl.util import logger
from keras.models import Sequential


class DoubleDQN(DQN):

    '''
    The base class of double DQNs
    '''

    def build_model(self):
        super(DoubleDQN, self).build_model()

        model2 = Sequential.from_config(self.model.get_config())
        logger.info("Model 2 summary")
        model2.summary()
        self.model2 = model2

        self.model2.compile(
            loss='mean_squared_error', optimizer=self.optimizer)
        logger.info("Models built and compiled")
        return self.model, self.model2

    def train(self, sys_vars):
        '''
        Training is for the Q function (NN) only
        otherwise (e.g. policy) see self.update()
        step 1,2,3,4 of algo.
        '''
        loss_total = 0
        for _epoch in range(self.n_epoch):
            minibatch = self.memory.rand_minibatch(self.batch_size)
            # note the computed values below are batched in array
            Q_states = self.model.predict(minibatch['states'])

            # Select max using model 2
            Q_next_states_select = self.model2.predict(
                minibatch['next_states'])
            Q_next_states_max_ind = np.argmax(Q_next_states_select, axis=1)
            # if more than one max, pick 1st
            if (Q_next_states_max_ind.shape[0] > 1):
                Q_next_states_max_ind = Q_next_states_max_ind[0]
            # Evaluate max using model 1

            Q_next_states = self.model.predict(minibatch['next_states'])
            Q_next_states_max = Q_next_states[:, Q_next_states_max_ind]

            # make future reward 0 if exp is terminal
            Q_targets_a = minibatch['rewards'] + self.gamma * \
                (1 - minibatch['terminals']) * Q_next_states_max
            # set batch Q_targets of a as above, the rest as is
            # minibatch['actions'] is one-hot encoded
            Q_targets = minibatch['actions'] * Q_targets_a[:, np.newaxis] + \
                (1 - minibatch['actions']) * Q_states

            loss = self.model.train_on_batch(minibatch['states'], Q_targets)
            loss_total += loss

            # Switch model 1 and model 2
            temp = self.model
            self.model = self.model2
            self.model2 = temp
        avg_loss = loss_total / self.n_epoch
        sys_vars['loss'].append(avg_loss)
        return avg_loss

@@ -607,183 +607,758 @@
 def 
-train(self, sys_vars):%0A        '''%0A        Training is f
+compute_Q_states(self, minibatch):%0A        Q_states = self.model.predict(minibatch%5B'states'%5D)%0A        # Different from (single) dqn: Select max using model 2%0A        Q_next_states_select = self.model2.predict(%0A            minibatch%5B'next_states'%5D)%0A        Q_next_states_max_ind = np.argmax(Q_next_states_select, axis=1)%0A        # if m
 or
+e
  th
-e Q function (NN) only%0A        otherwise (e.g. policy) see self.update()%0A        step 1,2,3,4 of algo.%0A        '''
+an one max, pick 1st%0A        if (Q_next_states_max_ind.shape%5B0%5D %3E 1):%0A            Q_next_states_max_ind = Q_next_states_max_ind%5B0%5D%0A%0A        # same as dqn again, but use Q_next_states_max_ind above%0A        Q_next_states = self.model.predict(minibatch%5B'next_states'%5D)%0A        Q_next_states_max = Q_next_states%5B:, Q_next_states_max_ind%5D%0A        return (Q_states, Q_next_states_max)%0A%0A    def train(self, sys_vars):
 %0A   
@@ -1416,24 +1416,50 @@
 f.n_epoch):%0A
+            # same as dqn%0A
             
@@ -1518,73 +1518,8 @@
 ze)%0A
-            # note the computed values below are batched in array
 %0A   
@@ -1523,32 +1523,33 @@
             
+(
 Q_states
  = self.mode
@@ -1540,669 +1540,63 @@
 ates
- = self.model.predict(minibatch%5B'states'%5D)%0A%0A            # Select max using model 2%0A            Q_next_states_select = self.model2.predict(%0A                minibatch%5B'next_states'%5D)%0A            Q_next_states_max_ind = np.argmax(Q_next_states_select, axis=1)%0A            # if more than one max, pick 1st%0A            if (Q_next_states_max_ind.shape%5B0%5D %3E 1):%0A                Q_next_states_max_ind = Q_next_states_max_ind%5B0%5D%0A            # Evaluate max using model 1%0A%0A            Q_next_states = self.model.predict(minibatch%5B'next_states'%5D)%0A            Q_next_states_max = Q_next_states%5B:, Q_next_states_max_ind%5D%0A%0A            # make future reward 0 if exp is terminal
+, Q_next_states_max) = self.compute_Q_states(minibatch)
 %0A   
@@ -1617,180 +1617,41 @@
 gets
-_a
  = 
-minibatch%5B'rewards'%5D + self.gamma * %5C%0A                (1 - minibatch%5B'terminals'%5D) * Q_next_states_max%0A            # set batch Q_targets of a as above, the rest as is%0A
+self.compute_Q_targets(%0A      
     
@@ -1655,20 +1655,16 @@
         
-   #
  minibat
@@ -1669,170 +1669,38 @@
 atch
-%5B'actions'%5D is one-hot encoded%0A            Q_targets = minibatch%5B'actions'%5D * Q_targets_a%5B:, np.newaxis%5D + %5C%0A                (1 - minibatch%5B'actions'%5D) * Q
+, Q_states, Q_next
 _states
+_max)
 %0A%0A  

refactor double_dqn compute Q method