import numpy as np
from copy import deepcopy
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from replay_buffers import ReplayBuffer, PrioritisedReplayBuffer
from tqdm import tqdm

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("device = ", device)


class OUProcess():
    """Implementation of Ornstein–Uhlenbeck process for noise generation."""

    def __init__(self, size, seed, mu=0, theta=0.1, sigma=0.1):
        """OU process initialiser.

        Args:
            size (int): dimension of the noise tensor required.
            seed (int): seed for Gaussian noise generation.
            mu (int): mean value for of noise generated.
            theta (int): drift scaling parameter.
            sigma (int): noise scaling parameter.
        """
        self.state = np.zeros(size)
        self.mu = mu
        self.size = size
        self.theta = theta
        self.sigma = sigma
        np.random.seed(seed)

    def sample(self):
        """Returns noise generated by OU process."""
        noise = np.random.normal(0.0, 1.0, self.size)
        self.state += self.theta * (self.mu - self.state) + self.sigma * noise
        return self.state


class CriticNetwork(nn.Module):
    """A critic network.

    A class that creates a neural network to use as a Q function
    approximation and as critic in the DDPG algorithm.
    """

    def __init__(self, state_size, action_size, hidden_layers):
        """Critic network initialisation.

        Args:
            state_size (int): dimension of state space.
            action_size (int): dimension of action space.
            hidden_layers (list[int]): list of dimensions for the hidden layers required.
        """
        super().__init__()
        self.input = nn.Linear(state_size + action_size, hidden_layers[0])
        self.hidden_layers = nn.ModuleList()
        for layer_in, layer_out in zip(hidden_layers[:-1], hidden_layers[1:]):
            self.hidden_layers.append(nn.Linear(layer_in, layer_out))
        self.output = nn.Linear(hidden_layers[-1], 1)

    def forward(self, state, action):
        """Forward propagation step for critic network.

        When given a state-action pair an approximation of its Q value
        is returned.
        """
        state_action = torch.cat((state, action), 2)
        x = self.input(state_action)
        x = F.relu(x)

        for layer in self.hidden_layers:
            x = layer(x)
            x = F.relu(x)

        value = self.output(x)

        return value


class ActorNetwork(nn.Module):
    """An actor network.

    A class that creates a neural network that will learn to approximate an
    optimal policy and is used as the actor network in the DDPG algorithm
    """

    def __init__(self, state_size, action_size, hidden_layers):
        """Critic network initialisation.

        Args:
            state_size (int): dimension of state space.
            action_size (int): dimension of action space.
            hidden_layers (list[int]): list of dimensions for the hidden layers required.
        """
        super().__init__()
        self.input = nn.Linear(state_size, hidden_layers[0])
        self.hidden_layers = nn.ModuleList()
        for layer_in, layer_out in zip(hidden_layers[:-1], hidden_layers[1:]):
            self.hidden_layers.append(nn.Linear(layer_in, layer_out))
        self.output = nn.Linear(hidden_layers[-1], action_size)

    def forward(self, state):
        """ Forward propagation step of the actor network.

        When passed a state this function approximates the optimal action to take in that state.
        """
        state = state.to(device)

        x = self.input(state)
        x = F.relu(x)

        for layer in self.hidden_layers:
            x = layer(x)
            x = F.relu(x)
        output = self.output(x)

        return torch.tanh(output)


class DDPGAgent():
    """ An agent that learns using the DDPG algorithm.

    An agent that is trained using the actor-critic based DDPG method that
    uses a neural network with an associated target network for both actor
    and critic that are softly updated. Additionally, a replay buffer is used
    for experience selection during training and an Ornstein–Uhlenbeck process
    to generate exploratory noise which decays linearly during training.
    """

    def __init__(
        self,
        buffer_size,
        seed,
        state_size,
        action_size,
        actor_network,
        critic_network,
        gamma,
        tau,
        l2_decay,
        learning_rate_actor,
        learning_rate_critic,
        noise_decay,
        min_noise,
        prioritised_replay_buffer=False,
        alpha=None,
        beta=None,
        beta_increment_size=None,
        base_priority=None,
        max_priority=None,
        episode_scores=None,
        step_number=0,
    ):
        """DDPGAgent initialisation function.

        Args:
            buffer_size (int): maximum size of the replay buffer.
            seed (int): random seed used for batch selection.

            state_size (int): dimension of the observed state space.
            action_size (int): dimension of action space.
            hidden_layers_actor (list[int]): list of dimensions for the hidden layers required for the actor.
            hidden_layers_critic (list[int]): list of dimensions for the hidden layers required for the critic.

            gamma (int): discount factor for future expected returns.
            tau (int): soft update factor used to define how much to shift
                       target network parameters towards current network parameter.
            l2_decay (int): a regularisation factor for the critic network.

            noise_decay (int): linear decay rate for noise.
            noise_min (int): minimum value for the noise during training.

            learning_rate_actor (int): learning rate for the actors gradient decent optimisation.
            learning_rate_critic (int): learning rate for the critics gradient decent optimisation.

            priority_replay_buffer (bool): set true to use priority replay buffer.
            alpha (float): priority scaling hyperparameter.
            beta_zero (float): importance sampling scaling hyperparameter.
            beta_increment_size (float): beta annealing rate.
            base_priority (float): base priority to ensure non-zero sampling probability.
            max_priority (float): initial maximum priority.

            episode_scores (list[int]): rewards gained in previous training episodes.
                                        (this is primarily used to reloading saved agents)
            step_number (int): number of steps the agent has taken.
                               (this is primarily used to reloading saved agents)
        """
        self.buffer_size = buffer_size
        self.seed = seed

        self.prioritised_replay_buffer = prioritised_replay_buffer
        self.alpha = alpha
        self.beta = beta
        self.beta_increment_size = beta_increment_size
        self.base_priority = base_priority
        self.max_priority = max_priority

        if prioritised_replay_buffer:
            print("using PER")

            self.replay_buffer = PrioritisedReplayBuffer(
                buffer_size,
                alpha,
                beta,
                beta_increment_size,
                base_priority,
                max_priority,
                seed,
            )
        else:
            self.replay_buffer = ReplayBuffer(buffer_size, seed)

        self.state_size = state_size
        self.action_size = action_size

        self.actor_net = actor_network.to(device)
        self.actor_target = deepcopy(actor_network).to(device)
        self.actor_optimizer = optim.Adam(self.actor_net.parameters(), lr=learning_rate_actor)

        self.critic_net = critic_network.to(device)
        self.critic_target = deepcopy(critic_network).to(device)
        self.critic_optimizer = optim.Adam(self.critic_net.parameters(), lr=learning_rate_critic, weight_decay=l2_decay)

        self.noise = OUProcess(self.action_size, self.seed)

        self.gamma = gamma
        self.tau = tau
        self.l2_decay = l2_decay
        self.noise_decay = noise_decay
        self.min_noise = min_noise

        self.learning_rate_actor = learning_rate_actor
        self.learning_rate_critic = learning_rate_critic

        if episode_scores is None:
            self.episode_scores = []
        else:
            self.episode_scores = episode_scores
        self.step_number = step_number

    def step(self, state, action, reward, next_state, done, batch_size):
        """Function that add experience to buffer and learns after each environment step."""

        self.replay_buffer.add(state, action, reward, next_state, done)
        self.learn(batch_size)
        self.step_number += 1

    def act(self, state, add_noise=True):
        """Uses the actor network to choose an action given current state."""
        state = torch.FloatTensor(state).to(device)
        self.actor_net.eval()
        with torch.no_grad():
            action = self.actor_net.forward(state).cpu().data.numpy()
        self.actor_net.train()

        if add_noise:
            action += self.noise.sample() * max((self.noise_decay**self.step_number), self.min_noise)
            return np.clip(action, -1., 1.)
        return action

    def learn(self, batch_size):
        """Samples experiences from the replay buffer and uses them to learn."""

        if len(self.replay_buffer) > 10 * batch_size:

            experience = self.replay_buffer.sample(batch_size)

            states = torch.FloatTensor(experience[0]).to(device)
            actions = torch.FloatTensor(experience[1]).to(device)
            rewards = torch.FloatTensor(experience[2]).unsqueeze(1).to(device)
            next_states = torch.FloatTensor(experience[3]).to(device)
            done_tensor = torch.FloatTensor(experience[4]).unsqueeze(1).to(device)

            """Update of critic network."""
            next_actions = self.actor_target.forward(next_states)

            Q_target_next = self.critic_target.forward(next_states, next_actions)
            Q_target = rewards + self.gamma * Q_target_next * (1 - done_tensor)

            Q_expected = self.critic_net(states, actions)

            if self.prioritised_replay_buffer:
                idx_list = experience[5]
                weights = torch.FloatTensor(experience[6]).unsqueeze(1).to(device)
                td_error = (Q_target - Q_expected)
                priority_list = torch.abs(td_error.squeeze().detach()).cpu().numpy()
                self.replay_buffer.update(idx_list, priority_list)
                critic_loss = torch.mean(weights * td_error**2)
            else:
                critic_loss = F.mse_loss(Q_expected, Q_target)

            self.critic_optimizer.zero_grad()
            critic_loss.backward()
            torch.nn.utils.clip_grad_norm_(self.critic_net.parameters(), 0.5)
            self.critic_optimizer.step()

            """Update of actor network."""
            actor_loss = -self.critic_net(states, self.actor_net(states)).mean()

            self.actor_optimizer.zero_grad()
            actor_loss.backward()
            torch.nn.utils.clip_grad_norm_(self.actor_net.parameters(), 1)

            self.actor_optimizer.step()

            self.soft_update_targets()

    def soft_update_targets(self):
        """Soft updates the actor and critic network parameters."""
        for target_param, critic_param in zip(self.critic_target.parameters(), self.critic_net.parameters()):
            target_param.data = self.tau * critic_param.data + (1 - self.tau) * target_param.data

        for target_param, actor_param in zip(self.actor_target.parameters(), self.actor_net.parameters()):
            target_param.data = self.tau * actor_param.data + (1 - self.tau) * target_param.data

    def save_agent(self, name, path=""):
        params = (
            self.buffer_size,
            self.seed,
            self.state_size,
            self.action_size,
            self.actor_net,
            self.critic_net,
            self.gamma,
            self.tau,
            self.l2_decay,
            self.learning_rate_actor,
            self.learning_rate_critic,
            self.noise_decay,
            self.min_noise,
            self.prioritised_replay_buffer,
            self.alpha,
            self.beta,
            self.beta_increment_size,
            self.base_priority,
            self.max_priority,
            self.episode_scores,
            self.step_number,
        )
        checkpoint = {"params": params,
                      "actor_optimizer": self.actor_optimizer.state_dict(),
                      "critic_optimizer": self.critic_optimizer.state_dict()
                      }
        path += name
        torch.save(checkpoint, path)


def load_agent(file_path, show_parameters=True):
    """ Reloads agents saved using DQNAgent.save_agent function.

        Parameters:
            file_path (str): file path of the saved agent being reloaded.

    """
    if torch.cuda.is_available():
        map_location = lambda storage, loc: storage.cuda()
    else:
        map_location = 'cpu'
    checkpoint = torch.load(file_path, map_location=map_location)
    agent = DDPGAgent(*checkpoint["params"])

    agent.actor_optimizer.load_state_dict(checkpoint["actor_optimizer"])
    agent.critic_optimizer.load_state_dict(checkpoint["critic_optimizer"])

    if show_parameters:
        print("Agent loaded with parameters: \n", agent.__dict__)

    return agent


def train_DDPG_agent(agent, env, episodes, batch_size, name, path="", print_every=100):
    """ Used to train DDPGAgent in the a Unity environment.

    Returns list of scores for each episode and additionally saves a copy of
    the final agent and of the agent that scored the highest average score.

    Args:
        agent (DDPGAgent): a DDPG agent to act in the environment.
        env (Unity environment): a Unity environment for the agent.
        episodes (int): number of episode to train the agent for.
        batch_size (int): size of experience batch used during learning.
        name (str): name used to save the best preforming and final agent.
        path (str): file path for the saved agent.
        print_every (int): how often to print the average score of the agent
    """

    brain_name = env.brain_names[0]
    name_best = name.split(".")[0] + "_best." + name.split(".")[1]
    best_average = 0

    for i in tqdm(range(episodes)):
        env_info = env.reset(train_mode=True)[brain_name]
        num_agents = len(env_info.agents)
        states = env_info.vector_observations
        scores = np.zeros(num_agents)
        dones = False
        while not np.any(dones):

            actions = agent.act(states)
            env_info = env.step(actions)[brain_name]
            next_states = env_info.vector_observations
            rewards = env_info.rewards
            dones = env_info.local_done
            agent.step(states, actions, rewards, next_states, dones, batch_size)
            scores += env_info.rewards
            states = next_states

        agent.episode_scores.append(scores)

        if i % print_every == print_every - 1:

            avg_score = np.mean(agent.episode_scores[-print_every:])
            print("Episode: {}  Average Score: {}".format(i+1, avg_score))
            if avg_score > best_average:
                best_average = avg_score
                agent.save_agent(name_best, path)
                print("Best agent so far has been saved.")

    agent.save_agent(name, path)
    print("Final agent has been saved.")
    return agent.episode_scores


def test_agent(agent, env, episodes=1, print_every=1, quick_view=False):
    """Tests agent using noiseless actor network and returns a list of episode rewards.

    Args:
        agent (DDPGAgent): a DDPG agent to act in the environment.
        env (Unity environment): a Unity environment for the agent.
        episodes (int): number of episode to test the agent over.
        print_every (int): how often to print the average score of the agent.
                           (setting to zero will stop the print)
        quick_view (bool): set to False to run environment in slow mode and observer agents actions.
    """
    brain_name = env.brain_names[0]
    episode_scores = []

    for i in range(episodes):
        env_info = env.reset(train_mode=quick_view)[brain_name]
        num_agents = len(env_info.agents)
        states = env_info.vector_observations
        scores = np.zeros(num_agents)
        dones = False

        while not np.any(dones):
            actions = agent.act(states, False)
            env_info = env.step(actions)[brain_name]
            next_states = env_info.vector_observations
            dones = env_info.local_done
            scores += env_info.rewards
            states = next_states
        episode_scores.append(scores)
        if i == 0:
            pass
        elif i+1 % print_every == 0:

            avg_score = np.mean(episode_scores[-print_every:])
            print("Episode: {}  Average Score: {}".format(i+1, avg_score))

    print("Average reward over {} tests was {} with a SD of {}".format(episodes, np.mean(episode_scores), np.std(episode_scores)))
    return episode_scores


def check_solved(averaged_training_scores):
    """Checks if agent as solved environment.

    Args:
        averaged_training_scores(list/panda.Series): a list of training scores averaged over 100 consecutive episodes.
    """
    for i, avg_score in enumerate(averaged_training_scores):
        if avg_score >= 30:
            print("Environment solved after {} training episodes.".format(i+1))
            return
    print("Environment not solved.")
