import numpy as np
from nets import PolicyNet, ValueNet
import torch.nn.functional as F
from rl_utils import *


class MomentumNPG:
    """A momentum-based NPG agent.
    Attributes:
        actors: policy networks for all agents (including copies of other agents).
        critic: value network for the i-th agent.
        gamma: discount factor.
        lmbda: lambda for GAE.
        kl_constraint: KL constraint for NPG.
        min_isw: minimum importance weight.
        beta: beta for momentum-based variance reduction
    """
    def __init__(self, num_agents, state_dim, action_dim, lmbda, kl_constraint,
                 critic_lr, gamma, device, min_isw, beta):
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.num_agents = num_agents
        self.actors = []
        for _ in range(self.num_agents):
            self.actors.append(PolicyNet(self.state_dim, self.action_dim).to(device))
        self.critic = ValueNet(self.state_dim).to(device)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=critic_lr)
        self.gamma = gamma
        self.lmbda = lmbda
        self.kl_constraint = kl_constraint
        self.device = device
        self.min_isw = min_isw
        self.beta = beta

    def take_actions(self, state):
        """Giving a state, take an action using policy network."""
        state = torch.tensor([state], dtype=torch.float).to(self.device)
        actions = []
        for actor in self.actors:
            probs = actor(state)
            action_dist = torch.distributions.Categorical(probs)
            action = action_dist.sample()
            actions.append(action)
        return actions

    def calc_log_probs(self, transition_dict):
        """Giving a transition dictionary, calculate log propobilites for a poliy."""
        states = torch.tensor(transition_dict['states'], dtype=torch.float).to(self.device)
        actions_list = torch.tensor(transition_dict['actions']).T.to(self.device)
        old_log_probs_list = []
        log_probs_list = []
        for idx, actor in enumerate(self.actors):
            # old_log_probs do not require gradients.
            old_log_probs = torch.log(actor(states).gather(1, actions_list[idx].unsqueeze(1))).detach()
            old_log_probs_list.append(old_log_probs)
            # log_probs require gradients.
            log_probs = torch.log(actor(states).gather(1, actions_list[idx].unsqueeze(1)))
            log_probs_list.append(log_probs)
        return old_log_probs_list, log_probs_list

    def compute_grads(self, advantage, old_log_probs_list, log_probs_list):
        """Calculate gradients for the advantage function."""
        obj_grad_list = []
        for actor, old_log_probs, log_probs in zip(self.actors, old_log_probs_list, log_probs_list):
            ratio = torch.exp(log_probs - old_log_probs)
            surrogate_obj = torch.mean(ratio * advantage)
            grads = torch.autograd.grad(surrogate_obj, actor.parameters())
            obj_grad = torch.cat([grad.view(-1) for grad in grads]).detach()
            obj_grad_list.append(obj_grad)
        return obj_grad_list

    def compute_IS_weight(self, actions_list, states, phis, min_isw):
        """Calculate importance weight."""
        probs_list = []
        old_probs_list = []
        for idx, (actor, phi) in enumerate(zip(self.actors, phis)):
            probs = actor(states).gather(1, actions_list[idx].unsqueeze(1)).detach()
            probs_list.append(probs)
            old_policy_probs = phi(states).gather(1, actions_list[idx].unsqueeze(1)).detach()
            old_probs_list.append(old_policy_probs)

        probs_list_t = torch.stack(probs_list).squeeze()
        old_policy_probs_t = torch.stack(old_probs_list).squeeze()
        prod_probs = torch.prod(probs_list_t, dim=0)
        prod_old_probs = torch.prod(old_policy_probs_t, dim=0)
        weights = prod_old_probs / (prod_probs + 1e-8)
        weight = torch.prod(weights)
        weight = np.max((min_isw, weight))
        weight_list = [weight for _ in self.actors]
        return weight_list

    def compute_grad_traj_prev_weights(self, states, actions_list, phis, advantage):
        """Calculate gradients of old policy using the trajectory generated by new policy"""
        obj_grad_list = []
        for idx, phi in enumerate(phis):
            old_policy_log_probs = torch.log(phi(states).gather(1, actions_list[idx].unsqueeze(1))).detach()
            log_probs = torch.log(phi(states).gather(1, actions_list[idx].unsqueeze(1)))
            ratio = torch.exp(log_probs - old_policy_log_probs)
            old_policy_surrogate_obj = torch.mean(ratio * advantage)
            grads = torch.autograd.grad(old_policy_surrogate_obj, phi.parameters())
            obj_grad = torch.cat([grad.view(-1) for grad in grads]).detach()
            obj_grad_list.append(obj_grad)
        return obj_grad_list

    def compute_v_list(self, transition_dict, advantage, prev_v_list, phis, beta):
        """Generate gradient estimator based on momentum-based variance reduction."""
        old_log_probs_list, log_probs_list = self.calc_log_probs(transition_dict)
        states = torch.tensor(transition_dict['states'], dtype=torch.float).to(self.device)
        actions_list = torch.tensor(transition_dict['actions']).T.to(self.device)

        isw_list = self.compute_IS_weight(actions_list, states, phis, self.min_isw)
        prev_g_list = self.compute_grad_traj_prev_weights(states, actions_list, phis, advantage)
        grad_list = self.compute_grads(advantage, old_log_probs_list, log_probs_list)

        grad_v_list = []
        for grad, prev_v, prev_g, isw in zip(grad_list, prev_v_list, prev_g_list, isw_list):
            grad_v = beta * grad + (1 - beta) * (prev_v + grad - isw * prev_g)
            grad_v_list.append(grad_v)
        return grad_v_list

    def update_value(self, transition_dict):
        """Update value network and calculate advantage functions."""
        states = torch.tensor(transition_dict['states'],
                              dtype=torch.float).to(self.device)
        rewards = torch.tensor(transition_dict['rewards'],
                               dtype=torch.float).view(-1, 1).to(self.device)
        next_states = torch.tensor(transition_dict['next_states'],
                                   dtype=torch.float).to(self.device)
        dones = torch.tensor(transition_dict['dones'],
                             dtype=torch.float).view(-1, 1).to(self.device)
        td_target = rewards + self.gamma * self.critic(next_states) * (1 - dones)
        td_delta = td_target - self.critic(states)
        critic_loss = torch.mean(F.mse_loss(self.critic(states), td_target.detach()))
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()

        advantage = compute_advantage(self.gamma, self.lmbda, td_delta.cpu()).to(self.device)
        return advantage

    def hessian_matrix_vector_product(self, states, old_action_dists, vector, idx):
        """Calculate product between Hessian of KL and gradient vector."""
        new_action_dists = torch.distributions.Categorical(self.actors[idx](states))
        # Derivative of KL divergence.
        kl = torch.mean(torch.distributions.kl.kl_divergence(old_action_dists, new_action_dists))
        kl_grad = torch.autograd.grad(kl, self.actors[idx].parameters(), create_graph=True)
        kl_grad_vector = torch.cat([grad.view(-1) for grad in kl_grad])
        # Dot product between gradient of KL and gradient vector.
        kl_grad_vector_product = torch.dot(kl_grad_vector, vector)
        # Derivative of kl_grad_vector_product.
        grad2 = torch.autograd.grad(kl_grad_vector_product, self.actors[idx].parameters())
        grad2_vector = torch.cat([grad.view(-1) for grad in grad2])
        return grad2_vector

    def conjugate_gradient(self, grad, states, old_action_dists, idx):
        """Conjugate gradient method for the calculation of $H^t_j \times y^{t+1}_j$.
        See: https://github.com/thu-ml/tianshou/blob/master/tianshou/policy/modelfree/npg.py
        """
        x = torch.zeros_like(grad)
        r = grad.clone()
        p = grad.clone()
        rdotr = torch.dot(r, r)
        for i in range(10):
            Hp = self.hessian_matrix_vector_product(states, old_action_dists, p, idx)
            alpha = rdotr / torch.dot(p, Hp)
            x += alpha * p
            r -= alpha * Hp
            new_rdotr = torch.dot(r, r)
            if new_rdotr < 1e-10:
                break
            beta = new_rdotr / rdotr
            p = r + beta * p
            rdotr = new_rdotr
        return x

    def compute_precondition_with_y(self, states_list, y_list):
        """Calculate precondition for NPG."""
        states = torch.tensor(states_list, dtype=torch.float).to(self.device)
        vec_grad_list = []  # policy list for all agents (including copies of other agents)
        for idx, y in enumerate(y_list):
            old_action_dists = torch.distributions.Categorical(self.actors[idx](states).detach())
            # Calculation of Hy using conjugate gradient method.
            descent_direction = self.conjugate_gradient(y, states, old_action_dists, idx)
            # Calculate NPG update coefficient.
            Hd = self.hessian_matrix_vector_product(states, old_action_dists, descent_direction, idx)
            max_coef = torch.sqrt(2 * self.kl_constraint / (torch.dot(descent_direction, Hd) + 1e-8))
            vec_grad = max_coef * descent_direction
            vec_grad_list.append(vec_grad)
        return vec_grad_list


