from typing import Union

import numpy as np
import torch
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel

from tartrl.algorithms.base_algorithm import BaseAlgorithm
from tartrl.utils.util import check
from tartrl.modules.utils.util import get_gard_norm, huber_loss, mse_loss

from tartrl.modules.networks.utils.distributed_utils import reduce_tensor


class PPOAlgorithm(BaseAlgorithm):
    def __init__(
            self,
            args,
            init_module,
            device: Union[str, torch.device] = "cpu",
    ) -> None:
        self.use_joint_action_loss = args.use_joint_action_loss
        super(PPOAlgorithm, self).__init__(args, init_module, device)

    def ppo_update(self, sample, turn_on=True):
        for optimizer in self.algo_module.optimizers.values():
            optimizer.zero_grad()

        share_obs_batch, obs_batch, rnn_states_batch, rnn_states_critic_batch, actions_batch, \
            value_preds_batch, return_batch, masks_batch, active_masks_batch, old_action_log_probs_batch, \
            adv_targ, available_actions_batch = sample

        old_action_log_probs_batch = check(old_action_log_probs_batch).to(**self.tpdv)
        adv_targ = check(adv_targ).to(**self.tpdv)
        value_preds_batch = check(value_preds_batch).to(**self.tpdv)
        return_batch = check(return_batch).to(**self.tpdv)
        active_masks_batch = check(active_masks_batch).to(**self.tpdv)

        if self.use_amp:
            with torch.cuda.amp.autocast():
                loss_list, value_loss, policy_loss, dist_entropy, ratio = self.prepare_loss(share_obs_batch,
                                                                                            obs_batch,
                                                                                            rnn_states_batch,
                                                                                            rnn_states_critic_batch,
                                                                                            actions_batch,
                                                                                            masks_batch,
                                                                                            available_actions_batch,
                                                                                            old_action_log_probs_batch,
                                                                                            adv_targ, value_preds_batch,
                                                                                            return_batch,
                                                                                            active_masks_batch, turn_on)
            for loss in loss_list:
                self.algo_module.scaler.scale(loss).backward()
        else:
            loss_list, value_loss, policy_loss, dist_entropy, ratio = self.prepare_loss(share_obs_batch,
                                                                                        obs_batch,
                                                                                        rnn_states_batch,
                                                                                        rnn_states_critic_batch,
                                                                                        actions_batch,
                                                                                        masks_batch,
                                                                                        available_actions_batch,
                                                                                        old_action_log_probs_batch,
                                                                                        adv_targ, value_preds_batch,
                                                                                        return_batch,
                                                                                        active_masks_batch, turn_on)
            for loss in loss_list:
                loss.backward()

        if "transformer" in self.algo_module.models:
            if self._use_max_grad_norm:
                grad_norm = nn.utils.clip_grad_norm_(self.algo_module.models['transformer'].parameters(),
                                                     self.max_grad_norm)
            else:
                grad_norm = get_gard_norm(self.algo_module.models['transformer'].parameters())
            critic_grad_norm = grad_norm
            actor_grad_norm = grad_norm

        else:
            if self._use_max_grad_norm:
                actor_grad_norm = nn.utils.clip_grad_norm_(self.algo_module.models['policy'].parameters(),
                                                           self.max_grad_norm)
            else:
                actor_grad_norm = get_gard_norm(self.algo_module.models['policy'].parameters())

            if self._use_max_grad_norm:
                critic_grad_norm = nn.utils.clip_grad_norm_(self.algo_module.models['critic'].parameters(),
                                                            self.max_grad_norm)
            else:
                critic_grad_norm = get_gard_norm(self.algo_module.models['critic'].parameters())

        if self.use_amp:
            for optimizer in self.algo_module.optimizers.values():
                self.algo_module.scaler.unscale_(optimizer)

            for optimizer in self.algo_module.optimizers.values():
                self.algo_module.scaler.step(optimizer)

            self.algo_module.scaler.update()
        else:
            for optimizer in self.algo_module.optimizers.values():
                optimizer.step()

        if self.world_size > 1:
            torch.cuda.synchronize()

        return value_loss, critic_grad_norm, policy_loss, dist_entropy, actor_grad_norm, ratio

    def cal_value_loss(
            self,
            value_normalizer,
            values,
            value_preds_batch,
            return_batch,
            active_masks_batch
    ):
        value_pred_clipped = value_preds_batch + (values - value_preds_batch).clamp(-self.clip_param, self.clip_param)

        if self._use_popart or self._use_valuenorm:
            value_normalizer.update(return_batch)
            error_clipped = value_normalizer.normalize(return_batch) - value_pred_clipped
            error_original = value_normalizer.normalize(return_batch) - values
        else:
            error_clipped = return_batch - value_pred_clipped
            error_original = return_batch - values

        if self._use_huber_loss:
            value_loss_clipped = huber_loss(error_clipped, self.huber_delta)
            value_loss_original = huber_loss(error_original, self.huber_delta)
        else:
            value_loss_clipped = mse_loss(error_clipped)
            value_loss_original = mse_loss(error_original)

        if self._use_clipped_value_loss:
            value_loss = torch.max(value_loss_original, value_loss_clipped)
        else:
            value_loss = value_loss_original

        if self._use_value_active_masks:
            value_loss = (value_loss * active_masks_batch).sum() / active_masks_batch.sum()
        else:
            value_loss = value_loss.mean()

        return value_loss

    def to_single_np(
            self,
            input
    ):
        reshape_input = input.reshape(-1, self._num_agents, *input.shape[1:])
        return reshape_input[:, 0, ...]

    def prepare_loss(
            self,
            share_obs_batch,
            obs_batch,
            rnn_states_batch,
            rnn_states_critic_batch,
            actions_batch,
            masks_batch,
            available_actions_batch,
            old_action_log_probs_batch,
            adv_targ,
            value_preds_batch,
            return_batch,
            active_masks_batch,
            turn_on
    ):
        loss_list = []

        if self.use_joint_action_loss:
            share_obs_batch = self.to_single_np(share_obs_batch)
            rnn_states_critic_batch = self.to_single_np(rnn_states_critic_batch)
            critic_masks_batch = self.to_single_np(masks_batch)
            value_preds_batch = self.to_single_np(value_preds_batch)
            return_batch = self.to_single_np(return_batch)
            adv_targ = adv_targ.reshape(-1, self._num_agents, 1)
            adv_targ = adv_targ[:, 0, :]

        else:
            critic_masks_batch = masks_batch

        values, action_log_probs, dist_entropy, policy_values = \
            self.algo_module.evaluate_actions(
                share_obs_batch,
                obs_batch,
                rnn_states_batch,
                rnn_states_critic_batch,
                actions_batch,
                masks_batch,
                available_actions_batch,
                active_masks_batch,
                critic_masks_batch=critic_masks_batch
            )

        if self.use_joint_action_loss:
            action_log_probs_copy = action_log_probs.reshape(-1, self._num_agents, action_log_probs.shape[-1]).sum(
                dim=(1, -1), keepdim=True).reshape(-1, 1)
            old_action_log_probs_batch_copy = old_action_log_probs_batch.reshape(-1, self._num_agents,
                                                                                 old_action_log_probs_batch.shape[
                                                                                     -1]).sum(dim=(1, -1),
                                                                                              keepdim=True).reshape(-1,
                                                                                                                    1)

            active_masks_batch = active_masks_batch.reshape(-1, self._num_agents, 1)
            active_masks_batch = active_masks_batch[:, 0, :]

            ratio = torch.exp(action_log_probs_copy - old_action_log_probs_batch_copy)
        else:
            ratio = torch.exp(action_log_probs - old_action_log_probs_batch)

        if self.dual_clip_ppo:
            ratio = torch.min(ratio, self.dual_clip_coeff)

        surr1 = ratio * adv_targ
        surr2 = torch.clamp(ratio, 1.0 - self.clip_param, 1.0 + self.clip_param) * adv_targ

        surr_final = torch.min(surr1, surr2)

        if self._use_policy_active_masks:
            policy_action_loss = (-torch.sum(surr_final, dim=-1,
                                             keepdim=True) * active_masks_batch).sum() / active_masks_batch.sum()
        else:
            policy_action_loss = -torch.sum(surr_final, dim=-1, keepdim=True).mean()

        if self._use_policy_vhead:
            if isinstance(self.algo_module.models["actor"], DistributedDataParallel):
                policy_value_normalizer = self.algo_module.models["actor"].module.value_normalizer
            else:
                policy_value_normalizer = self.algo_module.models["actor"].value_normalizer
            policy_value_loss = self.cal_value_loss(policy_value_normalizer, policy_values, value_preds_batch,
                                                    return_batch, active_masks_batch)
            policy_loss = policy_action_loss + policy_value_loss * self.policy_value_loss_coef
        else:
            policy_loss = policy_action_loss

        # critic update
        if isinstance(self.algo_module.models["critic"], DistributedDataParallel):
            value_normalizer = self.algo_module.models["critic"].module.value_normalizer
        else:
            value_normalizer = self.algo_module.get_critic_value_normalizer()
        value_loss = self.cal_value_loss(value_normalizer, values, value_preds_batch, return_batch,
                                         active_masks_batch)

        if "transformer" in self.algo_module.models:
            loss = policy_loss - dist_entropy * self.entropy_coef + value_loss * self.value_loss_coef
            loss_list.append(loss)
        else:
            if turn_on:
                final_p_loss = policy_loss - dist_entropy * self.entropy_coef

                loss_list.append(final_p_loss)
            final_v_loss = value_loss * self.value_loss_coef
            loss_list.append(final_v_loss)
        return loss_list, value_loss, policy_loss, dist_entropy, ratio



    def train(self, buffer, turn_on=True):
        if self._use_popart or self._use_valuenorm:
            if isinstance(self.algo_module.models["critic"], DistributedDataParallel):
                value_normalizer = self.algo_module.models["critic"].module.value_normalizer
            else:
                value_normalizer = self.algo_module.get_critic_value_normalizer()
            advantages = buffer.returns[:-1] - value_normalizer.denormalize(buffer.value_preds[:-1])
        else:
            advantages = buffer.returns[:-1] - buffer.value_preds[:-1]

        if self._use_adv_normalize:
            advantages = (advantages - advantages.mean()) / (
                    advantages.std() + 1e-5)

        advantages_copy = advantages.copy()
        advantages_copy[buffer.active_masks[:-1] == 0.0] = np.nan
        mean_advantages = np.nanmean(advantages_copy)
        std_advantages = np.nanstd(advantages_copy)
        advantages = (advantages - mean_advantages) / (std_advantages + 1e-5)

        train_info = {}

        train_info['value_loss'] = 0
        train_info['policy_loss'] = 0
        if self._use_new_network:
            train_info['id_loss'] = 0
            train_info['id_entropy'] = 0

        train_info['dist_entropy'] = 0
        train_info['actor_grad_norm'] = 0
        train_info['critic_grad_norm'] = 0
        train_info['ratio'] = 0
        if self.world_size > 1:
            train_info['reduced_value_loss'] = 0
            train_info['reduced_policy_loss'] = 0
            if self._use_new_network:
                train_info['reduced_id_loss'] = 0
                train_info['reduced_id_entropy'] = 0

        for _ in range(self.ppo_epoch):
            if 'transformer' in self.algo_module.models:
                data_generator = buffer.feed_forward_generator_transformer(advantages, self.num_mini_batch)
            elif self._use_recurrent_policy:
                if self.use_joint_action_loss:
                    data_generator = buffer.recurrent_generator_v3(advantages, self.num_mini_batch,
                                                                   self.data_chunk_length)
                else:
                    data_generator = buffer.recurrent_generator(advantages, self.num_mini_batch, self.data_chunk_length)
            elif self._use_naive_recurrent:
                data_generator = buffer.naive_recurrent_generator(advantages, self.num_mini_batch)
            else:
                data_generator = buffer.feed_forward_generator(advantages, self.num_mini_batch)

            for sample in data_generator:
                value_loss, critic_grad_norm, policy_loss, dist_entropy, actor_grad_norm, ratio \
                    = self.ppo_update(sample, turn_on)

                if self.world_size > 1:
                    train_info['reduced_value_loss'] += reduce_tensor(value_loss.data, self.world_size)
                    train_info['reduced_policy_loss'] += reduce_tensor(policy_loss.data, self.world_size)

                train_info['value_loss'] += value_loss.item()
                train_info['policy_loss'] += policy_loss.item()

                train_info['dist_entropy'] += dist_entropy.item()
                train_info['actor_grad_norm'] += actor_grad_norm
                train_info['critic_grad_norm'] += critic_grad_norm
                train_info['ratio'] += ratio.mean().item()

        num_updates = self.ppo_epoch * self.num_mini_batch

        for k in train_info.keys():
            train_info[k] /= num_updates

        for optimizer in self.algo_module.optimizers.values():
            if hasattr(optimizer, 'sync_lookahead'):
                optimizer.sync_lookahead()

        return train_info
