"""
This tutorial uses a simple implementation of Deep Deterministic Policy Gradient (DDPG), 
derived from keras examples, refactored to work with this tutorial.

This implementation uses tensorflow, and keras wrapers for the neural networks to do 
the heavy lifting of the implementation for us. 
For more examles visit: https://keras.io/examples/rl/ddpg_pendulum/
"""

import gym
import tensorflow as tf
from tensorflow.keras import layers
import numpy as np
import matplotlib.pyplot as plt

from exploration import OUActionNoise
from rpm import Buffer, update_target


class DDPG:
    def __init__(self,
                 problem_name,
                 num_states,
                 num_actions,
                 lower_bound,
                 upper_bound,
                 total_episodes=100):

        self.problem_name = problem_name
        self.num_states = num_states
        self.num_actions = num_actions
        self.lower_bound = lower_bound
        self.upper_bound = upper_bound

        self.std_dev = 0.2
        self.ou_noise = OUActionNoise(mean=np.zeros(1),
                                      std_deviation=float(self.std_dev) *
                                      np.ones(1))

        self.actor_model = self.get_actor()
        self.critic_model = self.get_critic()

        self.target_actor = self.get_actor()
        self.target_critic = self.get_critic()

        # Making the weights equal initially
        self.target_actor.set_weights(self.actor_model.get_weights())
        self.target_critic.set_weights(self.critic_model.get_weights())

        # Learning rate for actor-critic models
        self.critic_lr = 0.002
        self.actor_lr = 0.001

        self.critic_optimizer = tf.keras.optimizers.Adam(self.critic_lr)
        self.actor_optimizer = tf.keras.optimizers.Adam(self.actor_lr)

        self.total_episodes = total_episodes
        # Discount factor for future rewards
        self.gamma = 0.99
        # Used to update target networks
        self.tau = 0.005

        self.buffer = Buffer(50000,
                             64,
                             gamma=self.gamma,
                             num_states=self.num_states,
                             num_actions=self.num_actions)

    def get_actor(self):
        # Initialize weights between -3e-3 and 3-e3
        last_init = tf.random_uniform_initializer(minval=-0.003, maxval=0.003)

        inputs = layers.Input(shape=(self.num_states, ))
        out = layers.Dense(256, activation="relu")(inputs)
        out = layers.Dense(256, activation="relu")(out)
        outputs = layers.Dense(self.num_actions,
                               activation="tanh",
                               kernel_initializer=last_init)(out)

        # Our upper bound is 2.0 for Pendulum.
        outputs = outputs * self.upper_bound
        model = tf.keras.Model(inputs, outputs)
        return model

    def get_critic(self):
        # State as input
        state_input = layers.Input(shape=(self.num_states))
        state_out = layers.Dense(16, activation="relu")(state_input)
        state_out = layers.Dense(32, activation="relu")(state_out)

        # Action as input
        action_input = layers.Input(shape=(self.num_actions))
        action_out = layers.Dense(32, activation="relu")(action_input)

        # Both are passed through seperate layer before concatenating
        concat = layers.Concatenate()([state_out, action_out])

        out = layers.Dense(256, activation="relu")(concat)
        out = layers.Dense(256, activation="relu")(out)
        outputs = layers.Dense(self.num_actions)(out)

        # Outputs single value for give state-action
        model = tf.keras.Model([state_input, action_input], outputs)

        return model

    def policy(self, state, noise_object):
        sampled_actions = tf.squeeze(self.actor_model(state))
        noise = noise_object()
        # Adding noise to action
        sampled_actions = sampled_actions.numpy() + noise

        # We make sure action is within bounds
        legal_action = np.clip(sampled_actions, self.lower_bound,
                               self.upper_bound)

        return [np.squeeze(legal_action)]

    def save_parameters(self):

        # Save the weights
        actor_model.save_weights("pendulum_actor.h5")
        critic_model.save_weights("pendulum_critic.h5")

        target_actor.save_weights("pendulum_target_actor.h5")
        target_critic.save_weights("pendulum_target_critic.h5")

    def load_parameters(self):

        # Save the weights
        actor_model.load_parameters("pendulum_actor.h5")
        critic_model.load_parameters("pendulum_critic.h5")

        target_actor.load_parameters("pendulum_target_actor.h5")
        target_critic.load_parameters("pendulum_target_critic.h5")


if __name__ == '__main__':

    # problem = "Pendulum-v0"
    # env = gym.make(problem)

    # num_states = env.observation_space.shape[0]
    # print("Size of State Space ->  {}".format(num_states))
    # num_actions = env.action_space.shape[0]
    # print("Size of Action Space ->  {}".format(num_actions))

    # upper_bound = env.action_space.high[0]
    # lower_bound = env.action_space.low[0]

    # print("Max Value of Action ->  {}".format(upper_bound))
    # print("Min Value of Action ->  {}".format(lower_bound))

    problem = "Pendulum-v0"
    env = gym.make(problem)

    num_states = env.observation_space.shape[0]
    print("Size of State Space ->  {}".format(num_states))
    num_actions = env.action_space.shape[0]
    print("Size of Action Space ->  {}".format(num_actions))

    upper_bound = env.action_space.high[0]
    lower_bound = env.action_space.low[0]

    print("Max Value of Action ->  {}".format(upper_bound))
    print("Min Value of Action ->  {}".format(lower_bound))

    ddpg_agent = DDPG(problem_name=problem,
                      num_states=num_states,
                      num_actions=num_actions,
                      lower_bound=lower_bound,
                      upper_bound=upper_bound,
                      total_episodes=100)

    # To store reward history of each episode
    ep_reward_list = []
    # To store average reward history of last few episodes
    avg_reward_list = []

    # Takes about 4 min to train
    for ep in range(ddpg_agent.total_episodes):

        prev_state = env.reset()
        episodic_reward = 0
        render_episodes = 10
        render = not (ep % render_episodes)

        while True:
            # Uncomment this to see the Actor in action
            # But not in a python notebook.
            # env.render()

            tf_prev_state = tf.expand_dims(tf.convert_to_tensor(prev_state), 0)
            action = ddpg_agent.policy(tf_prev_state, ddpg_agent.ou_noise)
            # Recieve state and reward from environment.
            state, reward, done, info = env.step(action)
            if render: env.render()

            ddpg_agent.buffer.record((prev_state, action, reward, state))
            episodic_reward += reward

            ddpg_agent.buffer.learn(ddpg_agent.target_actor,
                                    ddpg_agent.target_critic,
                                    ddpg_agent.actor_model,
                                    ddpg_agent.critic_model,
                                    ddpg_agent.actor_optimizer,
                                    ddpg_agent.critic_optimizer)

            update_target(ddpg_agent.target_actor.variables,
                          ddpg_agent.actor_model.variables, ddpg_agent.tau)
            update_target(ddpg_agent.target_critic.variables,
                          ddpg_agent.critic_model.variables, ddpg_agent.tau)

            # End this episode when `done` is True
            if done:
                break

            prev_state = state

        ep_reward_list.append(episodic_reward)

        # Mean of last 40 episodes
        avg_reward = np.mean(ep_reward_list[-40:])
        print("Episode * {} * Avg Reward is ==> {}".format(ep, avg_reward))
        avg_reward_list.append(avg_reward)

    # Plotting graph
    # Episodes versus Avg. Rewards
    plt.plot(avg_reward_list)
    plt.xlabel("Episode")
    plt.ylabel("Avg. Epsiodic Reward")
    plt.show()