# Copyright (c) 2025, Istituto Italiano di Tecnologia
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause


from __future__ import annotations

from typing import Optional, Tuple, Dict, Any

import torch
import torch.nn as nn
import torch.optim as optim

# External modules providing the actor-critic model, storage utilities, and AMP components.
from rsl_rl.modules import ActorCritic
from rsl_rl.storage import RolloutStorage

from amp_rsl_rl.storage import ReplayBuffer
from amp_rsl_rl.networks import Discriminator
from amp_rsl_rl.utils import AMPLoader


class AMP_PPO:
    """
    AMP_PPO implements Adversarial Motion Prior (AMP) combined with Proximal Policy Optimization (PPO).

    This class integrates reinforcement learning via PPO with adversarial imitation learning using AMP.
    It improves policy training by leveraging both interactions with the environment and expert motion data.
    The discriminator is used to force the policy's behavior to align with expert demonstrations.

    Parameters
    ----------
    actor_critic : ActorCritic
        The policy network implementing both the actor (for action selection) and the critic (for value estimation).
    discriminator : Discriminator
        The adversarial discriminator model that distinguishes between expert transitions and those generated by the policy.
    amp_data : AMPLoader
        Data loader that provides batches of expert motion data.
    amp_normalizer : Optional[Any]
        Optional normalizer to preprocess AMP observations before they are fed into the discriminator.
    num_learning_epochs : int, default=1
        Number of passes over the collected data per update iteration.
    num_mini_batches : int, default=1
        Number of mini-batches to divide each epoch's data into.
    clip_param : float, default=0.2
        PPO clipping parameter that bounds the policy update step.
    gamma : float, default=0.998
        Discount factor for future rewards.
    lam : float, default=0.95
        Lambda parameter for Generalized Advantage Estimation (GAE).
    value_loss_coef : float, default=1.0
        Coefficient for the value function loss term in the PPO loss.
    entropy_coef : float, default=0.0
        Coefficient for the entropy regularization term (encouraging exploration).
    learning_rate : float, default=1e-3
        Initial learning rate for the optimizer.
    max_grad_norm : float, default=1.0
        Maximum gradient norm for clipping gradients during backpropagation.
    use_clipped_value_loss : bool, default=True
        Flag indicating whether to use a clipped value loss, as in the original PPO implementation.
    schedule : str, default="fixed"
        Learning rate schedule mode ("fixed" or "adaptive" based on KL divergence).
    desired_kl : float, default=0.01
        Target KL divergence for the adaptive learning rate schedule.
    amp_replay_buffer_size : int, default=100000
        Maximum number of policy transitions stored in the replay buffer for AMP training.
    device : str, default="cpu"
        The device (CPU or GPU) on which the models will be computed.
    """

    actor_critic: ActorCritic

    def __init__(
        self,
        actor_critic: ActorCritic,
        discriminator: Discriminator,
        amp_data: AMPLoader,
        amp_normalizer: Optional[Any],
        num_learning_epochs: int = 1,
        num_mini_batches: int = 1,
        clip_param: float = 0.2,
        gamma: float = 0.998,
        lam: float = 0.95,
        value_loss_coef: float = 1.0,
        entropy_coef: float = 0.0,
        learning_rate: float = 1e-3,
        max_grad_norm: float = 1.0,
        use_clipped_value_loss: bool = True,
        schedule: str = "fixed",
        desired_kl: float = 0.01,
        amp_replay_buffer_size: int = 100000,
        device: str = "cpu",
    ) -> None:
        # Set device and learning hyperparameters
        self.device: str = device
        self.desired_kl: float = desired_kl
        self.schedule: str = schedule
        self.learning_rate: float = learning_rate

        # Set up the discriminator and move it to the appropriate device.
        self.discriminator: Discriminator = discriminator.to(self.device)
        self.amp_transition: RolloutStorage.Transition = RolloutStorage.Transition()

        # Determine observation dimension used in the replay buffer.
        # The discriminator expects concatenated observations, so the replay buffer uses half the dimension.
        obs_dim: int = self.discriminator.input_dim // 2
        self.amp_storage: ReplayBuffer = ReplayBuffer(
            obs_dim=obs_dim, buffer_size=amp_replay_buffer_size, device=device
        )
        self.amp_data: AMPLoader = amp_data
        self.amp_normalizer: Optional[Any] = amp_normalizer

        # Set up the actor-critic (policy) and move it to the device.
        self.actor_critic = actor_critic
        self.actor_critic.to(self.device)
        self.storage: Optional[RolloutStorage] = (
            None  # Will be initialized later once environment parameters are known
        )

        # Create optimizer for both the actor-critic and the discriminator.
        # Note: Weight decay is set differently for discriminator trunk and head.
        params = [
            {"params": self.actor_critic.parameters(), "name": "actor_critic"},
            {
                "params": self.discriminator.trunk.parameters(),
                "weight_decay": 10e-4,
                "name": "amp_trunk",
            },
            {
                "params": self.discriminator.linear.parameters(),
                "weight_decay": 10e-2,
                "name": "amp_head",
            },
        ]
        self.optimizer: optim.Adam = optim.Adam(params, lr=learning_rate)
        self.transition: RolloutStorage.Transition = RolloutStorage.Transition()

        # PPO-specific parameters
        self.clip_param: float = clip_param
        self.num_learning_epochs: int = num_learning_epochs
        self.num_mini_batches: int = num_mini_batches
        self.value_loss_coef: float = value_loss_coef
        self.entropy_coef: float = entropy_coef
        self.gamma: float = gamma
        self.lam: float = lam
        self.max_grad_norm: float = max_grad_norm
        self.use_clipped_value_loss: bool = use_clipped_value_loss

    def init_storage(
        self,
        num_envs: int,
        num_transitions_per_env: int,
        actor_obs_shape: Tuple[int, ...],
        critic_obs_shape: Tuple[int, ...],
        action_shape: Tuple[int, ...],
    ) -> None:
        """
        Initializes the storage for collected transitions during interactions with the environment.

        Parameters
        ----------
        num_envs : int
            Number of parallel environments.
        num_transitions_per_env : int
            Number of transitions to store per environment.
        actor_obs_shape : tuple
            Shape of the observations for the actor.
        critic_obs_shape : tuple
            Shape of the observations for the critic.
        action_shape : tuple
            Shape of the actions taken by the policy.
        """
        self.storage = RolloutStorage(
            training_type="rl",
            num_envs=num_envs,
            num_transitions_per_env=num_transitions_per_env,
            obs_shape=actor_obs_shape,
            privileged_obs_shape=critic_obs_shape,
            actions_shape=action_shape,
            rnd_state_shape=None,
            device=self.device,
        )

    def test_mode(self) -> None:
        """
        Sets the actor-critic model to evaluation mode.
        """
        self.actor_critic.test()

    def train_mode(self) -> None:
        """
        Sets the actor-critic model to training mode.
        """
        self.actor_critic.train()

    def act(self, obs: torch.Tensor, critic_obs: torch.Tensor) -> torch.Tensor:
        """
        Selects an action based on the current observation and critic observation.
        It also records the necessary data (actions, log probabilities, values, etc.) in a transition.

        Parameters
        ----------
        obs : torch.Tensor
            Observation used by the actor network.
        critic_obs : torch.Tensor
            Observation used by the critic network for value estimation.

        Returns
        -------
        torch.Tensor
            The selected actions.
        """
        # If using a recurrent network, retrieve the hidden states.
        if self.actor_critic.is_recurrent:
            self.transition.hidden_states = self.actor_critic.get_hidden_states()
        # Compute actions and related statistics, ensuring we detach tensors to avoid gradient issues.
        self.transition.actions = self.actor_critic.act(obs).detach()
        self.transition.values = self.actor_critic.evaluate(critic_obs).detach()
        self.transition.actions_log_prob = self.actor_critic.get_actions_log_prob(
            self.transition.actions
        ).detach()
        self.transition.action_mean = self.actor_critic.action_mean.detach()
        self.transition.action_sigma = self.actor_critic.action_std.detach()
        # Record the observations before taking an environment step.
        self.transition.observations = obs
        self.transition.privileged_observations = critic_obs
        return self.transition.actions

    def act_amp(self, amp_obs: torch.Tensor) -> None:
        """
        Records the AMP observation (from the policy) into the AMP transition storage.

        Parameters
        ----------
        amp_obs : torch.Tensor
            The AMP observation from the policy.
        """
        self.amp_transition.observations = amp_obs

    def process_env_step(
        self, rewards: torch.Tensor, dones: torch.Tensor, infos: Dict[str, Any]
    ) -> None:
        """
        Processes the outcome of an environment step by recording rewards and handling bootstrapping
        for timeouts. Also resets the actor-critic's hidden state if necessary.

        Parameters
        ----------
        rewards : torch.Tensor
            Rewards received from the environment.
        dones : torch.Tensor
            Tensor indicating if an episode is finished.
        infos : dict
            Additional information from the environment, which may include 'time_outs'.
        """
        # Assign rewards to the current transition.
        self.transition.rewards = rewards.clone()
        self.transition.dones = dones

        # Handle bootstrapping on time-outs to adjust rewards.
        if "time_outs" in infos:
            self.transition.rewards += self.gamma * torch.squeeze(
                self.transition.values
                * infos["time_outs"].unsqueeze(1).to(self.device),
                1,
            )

        # Add the transition to storage and clear the temporary transition.
        self.storage.add_transitions(self.transition)
        self.transition.clear()
        # Reset the actor-critic states for environments that are done.
        self.actor_critic.reset(dones)

    def process_amp_step(self, amp_obs: torch.Tensor) -> None:
        """
        Processes an AMP step by inserting the current AMP observation and the new observation (from expert data)
        into the AMP replay buffer. Clears the temporary AMP transition afterwards.

        Parameters
        ----------
        amp_obs : torch.Tensor
            The new AMP observation (from expert data or policy update).
        """
        self.amp_storage.insert(self.amp_transition.observations, amp_obs)
        self.amp_transition.clear()

    def compute_returns(self, last_critic_obs: torch.Tensor) -> None:
        """
        Computes the discounted returns and advantages based on the critic's evaluation of the last observation.

        Parameters
        ----------
        last_critic_obs : torch.Tensor
            The critic observation after the last environment step.
        """
        last_values = self.actor_critic.evaluate(last_critic_obs).detach()
        self.storage.compute_returns(last_values, self.gamma, self.lam)

    def discriminator_policy_loss(
        self, discriminator_output: torch.Tensor
    ) -> torch.Tensor:
        """
        Computes the loss for the discriminator when classifying policy-generated transitions.
        Uses binary cross-entropy loss where the target label for policy transitions is 0.

        Parameters
        ----------
        discriminator_output : torch.Tensor
            The raw logits output from the discriminator for policy data.

        Returns
        -------
        torch.Tensor
            The computed policy loss.
        """
        loss_fn = torch.nn.BCEWithLogitsLoss()
        expected = torch.zeros_like(discriminator_output).to(self.device)
        return loss_fn(discriminator_output, expected)

    def discriminator_expert_loss(
        self, discriminator_output: torch.Tensor
    ) -> torch.Tensor:
        """
        Computes the loss for the discriminator when classifying expert transitions.
        Uses binary cross-entropy loss where the target label for expert transitions is 1.

        Parameters
        ----------
        discriminator_output : torch.Tensor
            The raw logits output from the discriminator for expert data.

        Returns
        -------
        torch.Tensor
            The computed expert loss.
        """
        loss_fn = torch.nn.BCEWithLogitsLoss()
        expected = torch.ones_like(discriminator_output).to(self.device)
        return loss_fn(discriminator_output, expected)

    def update(self) -> Tuple[float, float, float, float, float, float, float, float]:
        """
        Performs a single update step for both the actor-critic (PPO) and the AMP discriminator.
        It iterates over mini-batches of data, computes surrogate, value, AMP and gradient penalty losses,
        performs adaptive learning rate scheduling (if enabled), and updates model parameters.

        Returns
        -------
        tuple
            A tuple containing mean losses and statistics:
            (mean_value_loss, mean_surrogate_loss, mean_amp_loss, mean_grad_pen_loss,
             mean_policy_pred, mean_expert_pred, mean_accuracy_policy, mean_accuracy_expert)
        """
        # Initialize mean loss and accuracy statistics.
        mean_value_loss: float = 0.0
        mean_surrogate_loss: float = 0.0
        mean_amp_loss: float = 0.0
        mean_grad_pen_loss: float = 0.0
        mean_policy_pred: float = 0.0
        mean_expert_pred: float = 0.0
        mean_accuracy_policy: float = 0.0
        mean_accuracy_expert: float = 0.0
        mean_accuracy_policy_elem: float = 0.0
        mean_accuracy_expert_elem: float = 0.0

        # Create data generators for mini-batch sampling.
        if self.actor_critic.is_recurrent:
            generator = self.storage.reccurent_mini_batch_generator(
                self.num_mini_batches, self.num_learning_epochs
            )
        else:
            generator = self.storage.mini_batch_generator(
                self.num_mini_batches, self.num_learning_epochs
            )

        # Generator for policy-generated AMP transitions.
        amp_policy_generator = self.amp_storage.feed_forward_generator(
            num_mini_batch=self.num_learning_epochs * self.num_mini_batches,
            mini_batch_size=self.storage.num_envs
            * self.storage.num_transitions_per_env
            // self.num_mini_batches,
            allow_replacement=True,
        )

        # Generator for expert AMP data.
        amp_expert_generator = self.amp_data.feed_forward_generator(
            self.num_learning_epochs * self.num_mini_batches,
            self.storage.num_envs
            * self.storage.num_transitions_per_env
            // self.num_mini_batches,
        )

        # Loop over mini-batches from the environment transitions and AMP data.
        for sample, sample_amp_policy, sample_amp_expert in zip(
            generator, amp_policy_generator, amp_expert_generator
        ):
            # Unpack the mini-batch sample from the environment.
            (
                obs_batch,
                critic_obs_batch,
                actions_batch,
                target_values_batch,
                advantages_batch,
                returns_batch,
                old_actions_log_prob_batch,
                old_mu_batch,
                old_sigma_batch,
                hid_states_batch,
                masks_batch,
                rnd_state_batch,
            ) = sample

            # Forward pass through the actor to get current policy outputs.
            self.actor_critic.act(
                obs_batch, masks=masks_batch, hidden_states=hid_states_batch[0]
            )
            actions_log_prob_batch = self.actor_critic.get_actions_log_prob(
                actions_batch
            )
            value_batch = self.actor_critic.evaluate(
                critic_obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1]
            )
            mu_batch = self.actor_critic.action_mean
            sigma_batch = self.actor_critic.action_std
            entropy_batch = self.actor_critic.entropy

            # Adaptive learning rate adjustment based on KL divergence if schedule is "adaptive".
            if self.desired_kl is not None and self.schedule == "adaptive":
                with torch.inference_mode():
                    kl = torch.sum(
                        torch.log(sigma_batch / old_sigma_batch + 1.0e-5)
                        + (
                            torch.square(old_sigma_batch)
                            + torch.square(old_mu_batch - mu_batch)
                        )
                        / (2.0 * torch.square(sigma_batch))
                        - 0.5,
                        axis=-1,
                    )
                    kl_mean = torch.mean(kl)

                    if kl_mean > self.desired_kl * 2.0:
                        self.learning_rate = max(1e-5, self.learning_rate / 1.5)
                    elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0:
                        self.learning_rate = min(1e-2, self.learning_rate * 1.5)

                    for param_group in self.optimizer.param_groups:
                        param_group["lr"] = self.learning_rate

            # Compute the PPO surrogate loss.
            ratio = torch.exp(
                actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch)
            )
            surrogate = -torch.squeeze(advantages_batch) * ratio
            surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp(
                ratio, 1.0 - self.clip_param, 1.0 + self.clip_param
            )
            surrogate_loss = torch.max(surrogate, surrogate_clipped).mean()

            # Compute the value function loss.
            if self.use_clipped_value_loss:
                value_clipped = target_values_batch + (
                    value_batch - target_values_batch
                ).clamp(-self.clip_param, self.clip_param)
                value_losses = (value_batch - returns_batch).pow(2)
                value_losses_clipped = (value_clipped - returns_batch).pow(2)
                value_loss = torch.max(value_losses, value_losses_clipped).mean()
            else:
                value_loss = (returns_batch - value_batch).pow(2).mean()

            # Combine surrogate loss, value loss and entropy regularization to form PPO loss.
            ppo_loss = (
                surrogate_loss
                + self.value_loss_coef * value_loss
                - self.entropy_coef * entropy_batch.mean()
            )

            # Process AMP loss by unpacking policy and expert AMP samples.
            policy_state, policy_next_state = sample_amp_policy
            expert_state, expert_next_state = sample_amp_expert

            # Normalize AMP observations if a normalizer is provided.
            if self.amp_normalizer is not None:
                with torch.no_grad():
                    policy_state = self.amp_normalizer.normalize(policy_state)
                    policy_next_state = self.amp_normalizer.normalize(policy_next_state)
                    expert_state = self.amp_normalizer.normalize(expert_state)
                    expert_next_state = self.amp_normalizer.normalize(expert_next_state)

            # Pass concatenated state transitions to the discriminator.
            policy_d = self.discriminator(
                torch.cat([policy_state, policy_next_state], dim=-1)
            )
            expert_d = self.discriminator(
                torch.cat([expert_state, expert_next_state], dim=-1)
            )

            # Compute discriminator losses for expert and policy data.
            expert_loss = self.discriminator_expert_loss(expert_d)
            policy_loss = self.discriminator_policy_loss(policy_d)

            # AMP loss is the average of expert and policy losses.
            amp_loss = 0.5 * (expert_loss + policy_loss)

            # Compute gradient penalty to stabilize discriminator training.
            grad_pen_loss = self.discriminator.compute_grad_pen(
                *sample_amp_expert, lambda_=10
            )

            # The final loss combines the PPO loss with AMP losses.
            loss = ppo_loss + (amp_loss + grad_pen_loss)

            # Backpropagation and optimizer step.
            self.optimizer.zero_grad()
            loss.backward()
            nn.utils.clip_grad_norm_(self.actor_critic.parameters(), self.max_grad_norm)
            self.optimizer.step()

            # Update the normalizer with current policy and expert AMP observations.
            if self.amp_normalizer is not None:
                self.amp_normalizer.update(policy_state)
                self.amp_normalizer.update(expert_state)

            # Compute probabilities from the discriminator logits.
            policy_d_prob = torch.sigmoid(policy_d)
            expert_d_prob = torch.sigmoid(expert_d)

            # Update running statistics.
            mean_value_loss += value_loss.item()
            mean_surrogate_loss += surrogate_loss.item()
            mean_amp_loss += amp_loss.item()
            mean_grad_pen_loss += grad_pen_loss.item()
            mean_policy_pred += policy_d_prob.mean().item()
            mean_expert_pred += expert_d_prob.mean().item()

            # Calculate the accuracy of the discriminator.
            mean_accuracy_policy += torch.sum(
                torch.round(policy_d_prob) == torch.zeros_like(policy_d_prob)
            ).item()
            mean_accuracy_expert += torch.sum(
                torch.round(expert_d_prob) == torch.ones_like(expert_d_prob)
            ).item()

            # Record the total number of elements processed.
            mean_accuracy_expert_elem += expert_d_prob.numel()
            mean_accuracy_policy_elem += policy_d_prob.numel()

        # Average the statistics over all mini-batches.
        num_updates = self.num_learning_epochs * self.num_mini_batches
        mean_value_loss /= num_updates
        mean_surrogate_loss /= num_updates
        mean_amp_loss /= num_updates
        mean_grad_pen_loss /= num_updates
        mean_policy_pred /= num_updates
        mean_expert_pred /= num_updates
        mean_accuracy_policy /= mean_accuracy_policy_elem
        mean_accuracy_expert /= mean_accuracy_expert_elem

        # Clear the storage for the next update cycle.
        self.storage.clear()

        return (
            mean_value_loss,
            mean_surrogate_loss,
            mean_amp_loss,
            mean_grad_pen_loss,
            mean_policy_pred,
            mean_expert_pred,
            mean_accuracy_policy,
            mean_accuracy_expert,
        )
