#!/usr/bin/env python
#################################################################################
# Copyright 2018 ROBOTIS CO., LTD.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#################################################################################

# Authors: Gilbert #

import rospy
import os
import json
import numpy as np
import random
import time
import sys
sys.path.append(os.path.dirname(os.path.abspath(os.path.dirname(__file__))))
from collections import deque
from std_msgs.msg import Float32MultiArray
from environment_stage_1 import Env
from tensorflow.keras.models import Sequential, load_model,Model
from tensorflow.keras.optimizers import RMSprop
from tensorflow.keras.layers import Dense, Dropout, Activation , Input, Flatten, Conv2D, MaxPooling2D,Concatenate
from IPython.display import SVG,display
from keras.utils.vis_utils import model_to_dot
from tensorflow.keras.utils import plot_model
from tensorflow.keras import backend as K
from tensorflow.python.client import device_lib
import tensorflow as tf

EPISODES = 3000

class ReinforceAgent():
    def __init__(self, state_size,state_size_2, action_size):
        self.pub_result = rospy.Publisher('result', Float32MultiArray, queue_size=5)
        self.dirPath = os.path.dirname(os.path.realpath(__file__))
        self.dirPath = '/home/upm/RL_ws(18-8-2020)/src/hsr_navigation/save_model/stage_1_'
        self.result = Float32MultiArray()
        config = tf.ConfigProto()
        config.gpu_options.allow_growth=True
        sess = tf.Session(config=config)
       
        self.load_model = True
        self.load_episode = 240
        self.state_size = state_size
        self.state_size_2 = state_size_2
        self.action_size = action_size
        self.episode_step = 6000
        self.target_update = 2000
        self.discount_factor = 0.99
        self.learning_rate = 0.00025
        self.epsilon = 1.0
        self.epsilon_decay = 0.99
        self.epsilon_min = 0.05
        #defult was 64
        self.batch_size = 5
        self.train_start = 5
        self.memory = deque(maxlen=1000000)

        #rospy.loginfo(device_lib.list_local_devices())
        
        self.model = self.buildModel()
        self.target_model = self.buildModel()

        self.updateTargetModel()

        if self.load_model:
            self.model.set_weights(load_model(self.dirPath+str(self.load_episode)+".h5").get_weights())

            with open(self.dirPath+str(self.load_episode)+'.json') as outfile:
                param = json.load(outfile)
                self.epsilon = param.get('epsilon')

    def buildModel(self):
        input = Input(shape=(self.state_size,))
        input_2 = Input(shape=(self.state_size_2))

        dropout = 0.2
        out = Conv2D(128, kernel_size=(5,5), padding="same", activation="relu")(input_2)
        out = MaxPooling2D()(out)
        out = Conv2D(128, kernel_size=(3,3), padding="same", activation="relu")(out)
        out = MaxPooling2D()(out)
        out = Flatten()(out)
        out = Dense(24, activation="relu")(out)
        out = Dense(24, activation="relu")(out)
        out = Dense(6, activation="linear")(out)
        concat_A = Concatenate()([out, input])
        out = Dense(64, activation="relu", kernel_initializer='lecun_uniform')(concat_A)
        out = Dense(64, activation="relu", kernel_initializer='lecun_uniform')(out)
        out = Dropout(dropout)(out)
        out = Dense(self.action_size, activation="relu", kernel_initializer='lecun_uniform')(out)
        out = Activation('linear')(out)
        model = Model(inputs=[input,input_2], outputs=out)
        model.compile(optimizer=RMSprop(lr=self.learning_rate, rho=0.9, epsilon=1e-06), loss="mse")
        model.summary()
        #SVG(model_to_dot(model).create(prog='dot',format='svg'))
        #plot_model(model, show_shapes=True, to_file='/home/upm/Desktop/network.png')
       
        return model

    def getQvalue(self, reward, next_target, done):
        if done:
            return reward
        else:
            return reward + self.discount_factor * np.amax(next_target)

    def updateTargetModel(self):
        self.target_model.set_weights(self.model.get_weights())

    def getAction(self, state,camera_state):
        if np.random.rand() <= self.epsilon:
            self.q_value = np.zeros(self.action_size)
            return random.randrange(self.action_size)
        else:
            with tf.device("GPU:0"):
                q_value = self.model.predict([state.reshape(1, len(state)),np.expand_dims(camera_state, axis=0)])
            self.q_value = q_value
            
            return np.argmax(q_value[0])

    def appendMemory(self, state, action, reward, next_state, done,state_camera,next_state_camera):
        self.memory.append((state, action, reward, next_state, done,state_camera,next_state_camera))

    def trainModel(self, target=False):
        with tf.device("GPU:0"):
            mini_batch = random.sample(self.memory, self.batch_size)
            X_batch = np.empty((0, self.state_size), dtype=np.float64)
            Y_batch = np.empty((0, self.action_size), dtype=np.float64)
            #RGB
            #X_batch_camera = np.empty(((0, 480, 640, 3)), dtype=np.float64)
            #Gray
            X_batch_camera = np.empty(((0, 300, 400,1)), dtype=np.float64)
            for i in range(self.batch_size):
                states = mini_batch[i][0]
                actions = mini_batch[i][1]
                rewards = mini_batch[i][2]
                next_states = mini_batch[i][3]
                dones = mini_batch[i][4]
                state_camera = np.array(mini_batch[i][5])
                next_state_camera = mini_batch[i][6]


                q_value = self.model.predict([states.reshape(1, len(states)),np.expand_dims(state_camera, axis=0)])
                self.q_value = q_value

                if target:
                    next_target = self.target_model.predict([next_states.reshape(1, len(next_states)),np.expand_dims(next_state_camera, axis=0)])
                else:
                    next_target = self.model.predict([next_states.reshape(1, len(next_states)),np.expand_dims(next_state_camera, axis=0)])

                next_q_value = self.getQvalue(rewards, next_target, dones)

                X_batch = np.append(X_batch, np.array([states.copy()]), axis=0)
                Y_sample = q_value.copy()

                Y_sample[0][actions] = next_q_value
                Y_batch = np.append(Y_batch, np.array([Y_sample[0]]), axis=0)

                X_batch_camera =  np.append(X_batch_camera, np.array([state_camera.copy()]), axis=0)
                if dones:
                    X_batch = np.append(X_batch, np.array([next_states.copy()]), axis=0)
                    X_batch_camera =  np.append(X_batch_camera, np.array([next_state_camera.copy()]), axis=0)
                    Y_batch = np.append(Y_batch, np.array([[rewards] * self.action_size]), axis=0)
    #todo
            self.model.fit([X_batch,X_batch_camera], Y_batch, batch_size=self.batch_size, epochs=1, verbose=0)

if __name__ == '__main__':
    rospy.init_node('hsrb_dqn_stage_1')
    pub_result = rospy.Publisher('result', Float32MultiArray, queue_size=5)
    pub_get_action = rospy.Publisher('get_action', Float32MultiArray, queue_size=5)
    result = Float32MultiArray()
    get_action = Float32MultiArray()

    #RGB
    #state_size_2 = (480, 640, 3)
    
    #Gray
    state_size_2 = (300, 400,1)

    state_size = 182
    action_size = 5

    env = Env(action_size)

    agent = ReinforceAgent(state_size,state_size_2, action_size)
    scores, episodes = [], []
    global_step = 0
    start_time = time.time()

    for e in range(agent.load_episode + 1, EPISODES):
        done = False
        state,camera_state = env.reset()
        
        score = 0
        for t in range(agent.episode_step):
            action = agent.getAction(state,camera_state)
            if e < 200:
                next_state, reward, done,next_camera_state = env.step(action, False)
            else:
                next_state, reward, done,next_camera_state = env.step(action, True)
            agent.appendMemory(state, action, reward, next_state, done,camera_state,next_camera_state)

            if len(agent.memory) >= agent.train_start:
                if global_step <= agent.target_update:
                    agent.trainModel()
                else:
                    agent.trainModel(True)

            score += reward
            state = next_state
            camera_state = next_camera_state
            get_action.data = [action, score, reward]
            pub_get_action.publish(get_action)

            if e % 10 == 0:
                agent.model.save(agent.dirPath + str(e) + '.h5')
                with open(agent.dirPath + str(e) + '.json', 'w') as outfile:
                    json.dump(param_dictionary, outfile)

            if t >= 500:
                rospy.loginfo("Time out!!")
                done = True

            if done:
                result.data = [score, np.max(agent.q_value)]
                pub_result.publish(result)
                agent.updateTargetModel()
                scores.append(score)
                episodes.append(e)
                m, s = divmod(int(time.time() - start_time), 60)
                h, m = divmod(m, 60)

                rospy.loginfo('Ep: %d score: %.2f memory: %d epsilon: %.2f time: %d:%02d:%02d',
                              e, score, len(agent.memory), agent.epsilon, h, m, s)
                param_keys = ['epsilon']
                param_values = [agent.epsilon]
                param_dictionary = dict(zip(param_keys, param_values))
                break

            global_step += 1
            if global_step % agent.target_update == 0:
                rospy.loginfo("UPDATE TARGET NETWORK")

        if agent.epsilon > agent.epsilon_min:
            agent.epsilon *= agent.epsilon_decay
