# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
"""DreamWaQ PPO——标准 PPO + CENet VAE 自编码器损失。

在 PPO.update() 的每个 mini-batch 中，额外计算：
  - 速度估计损失：MSE(code_vel, 真实 base_vel)（来自 privileged_obs）
  - 观测重建损失：MSE(decode, 当前 policy obs)
  - KL 散度损失：beta * KL(N(mean, var) || N(0, 1))

class_name = "motrix_rl.rslrl.torch.train.dreamwaq_ppo:DreamWaQPPO"
"""

import torch
import torch.nn as nn
from tensordict import TensorDict

from rsl_rl.algorithms import PPO
from rsl_rl.storage import RolloutStorage
from rsl_rl.utils import resolve_callable, resolve_obs_groups


class DreamWaQPPO(PPO):
    """PPO + Beta-VAE 自编码器损失——DreamWaQ CENet 训练。"""

    # VAE 损失权重（与上游 beta=1.0 一致）
    vae_beta: float = 1.0

    @staticmethod
    def construct_algorithm(obs: TensorDict, env, cfg: dict, device: str) -> "DreamWaQPPO":
        """构造 DreamWaQ PPO 算法——创建 CENetActor + MLP Critic。

        与父类 PPO.construct_algorithm 的区别：
          - actor 使用 CENetActorModel（含 VAE）
          - obs_groups 中 actor=["policy", "obs_history"], critic=["privileged_obs"]
        """
        # 提取 DreamWaQ 特有参数
        vae_beta = cfg.pop("vae_beta", 1.0)

        # 解析 actor / critic 类
        actor_class = resolve_callable(cfg["actor"].pop("class_name"))
        critic_class = resolve_callable(cfg["critic"].pop("class_name"))

        # 解析观测分组
        obs_groups = resolve_obs_groups(obs, cfg["obs_groups"], ["actor", "critic"])

        # 创建 actor（CENetActorModel）
        actor = actor_class(obs, obs_groups, "actor", env.num_actions, **cfg["actor"]).to(device)

        # 创建 critic（标准 MLPModel，输入 privileged_obs）
        critic = critic_class(obs, obs_groups, "critic", 1, **cfg["critic"]).to(device)

        # 初始化 rollout 存储
        storage = RolloutStorage(
            "rl", env.num_envs, cfg["num_steps_per_env"], obs, [env.num_actions], device
        )

        # 提取算法参数（移除 DreamWaQPPO 特有 key，剩余传给父类 PPO.__init__）
        algo_cfg = dict(cfg["algorithm"])
        for dw_key in ("class_name", "vae_beta", "cenet_in_dim", "cenet_out_dim"):
            algo_cfg.pop(dw_key, None)

        # 创建 DreamWaQPPO 实例
        alg = DreamWaQPPO(actor, critic, storage, device=device, **algo_cfg)
        alg.vae_beta = vae_beta
        return alg

    def update(self) -> dict[str, float]:
        """标准 PPO update + VAE 自编码器损失。"""
        mean_value_loss = 0.0
        mean_surrogate_loss = 0.0
        mean_autoenc_loss = 0.0

        if self.actor.is_recurrent or self.critic.is_recurrent:
            generator = self.storage.recurrent_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)

        for (
            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,
        ) in generator:
            # NaN 检测 — 数据有 NaN 就跳过这个 batch
            if torch.isnan(obs_batch["policy"]).any() or torch.isnan(obs_batch["obs_history"]).any():
                continue
            # ── 标准 PPO 前向（各调用一次，复用结果）──
            self.actor(obs_batch, masks=masks_batch, stochastic_output=True)
            value_batch = self.critic(obs_batch, masks=masks_batch)

            # ── 动作对数概率 ──
            actions_log_prob_batch = self.actor.get_output_log_prob(actions_batch)

            # ── 价值损失 ──
            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 = torch.square(value_batch - returns_batch)
                value_losses_clipped = torch.square(value_clipped - returns_batch)
                value_loss = torch.max(value_losses, value_losses_clipped).mean()
            else:
                value_loss = torch.square(returns_batch - value_batch).mean()

            # ── 代理损失 ──
            surrogate_loss = self._compute_surrogate_loss(
                actions_batch, actions_log_prob_batch,
                old_actions_log_prob_batch, advantages_batch)

            # ── KL 自适应 schedule ──
            if self.desired_kl is not None and self.schedule == 'adaptive':
                with torch.inference_mode():
                    kl = torch.sum(
                        torch.log(self.actor.output_std / old_sigma_batch + 1e-5)
                        + (old_sigma_batch.pow(2) + (old_mu_batch - self.actor.output_mean).pow(2))
                        / (2.0 * self.actor.output_std.pow(2))
                        - 0.5, dim=-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

            # ── 熵 ──
            entropy_batch = self.actor.output_entropy
            entropy_loss = entropy_batch.mean()

            # ── VAE 自编码器损失 ──
            autoenc_loss = self._compute_vae_loss(obs_batch)

            # ── 总损失 ──
            loss = (
                surrogate_loss
                + self.value_loss_coef * value_loss
                - self.entropy_coef * entropy_loss
                + autoenc_loss
            )

            # ── 梯度更新 ──
            self.optimizer.zero_grad()
            loss.backward()
            if self.max_grad_norm is not None:
                nn.utils.clip_grad_norm_(self.actor.parameters(), self.max_grad_norm)
                nn.utils.clip_grad_norm_(self.critic.parameters(), self.max_grad_norm)
                nn.utils.clip_grad_value_(self.actor.parameters(), 10.0)
            self.optimizer.step()
            # 每次更新后强制 std > 0，防止数值异常导致 NaN
            if hasattr(self.actor, 'std') and self.actor.stochastic:
                with torch.no_grad():
                    self.actor.std.nan_to_num_(nan=0.5, posinf=1.0, neginf=1.0)
                    self.actor.std.clamp_(min=1e-6)

            # ── 累计日志 ──
            mean_value_loss += value_loss.item()
            mean_surrogate_loss += surrogate_loss.item()
            mean_autoenc_loss += autoenc_loss.item()

        num_updates = self.num_learning_epochs * self.num_mini_batches
        mean_value_loss /= num_updates
        mean_surrogate_loss /= num_updates
        mean_autoenc_loss /= num_updates

        self.storage.clear()

        return {
            "value_loss": mean_value_loss,
            "surrogate_loss": mean_surrogate_loss,
            "autoenc_loss": mean_autoenc_loss,
        }

    def _compute_surrogate_loss(
        self, actions_batch, actions_log_prob_batch,
        old_actions_log_prob_batch, advantages_batch
    ) -> torch.Tensor:
        """计算 PPO 代理损失（从父类 PPO.update() 中提取）。"""
        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)
        return torch.max(surrogate, surrogate_clipped).mean()

    def _compute_vae_loss(self, obs_batch: TensorDict) -> torch.Tensor:
        """计算 CENet VAE 损失。

        obs_batch 包含：
          - "policy": 当前观测 (N, 45) = 重建目标
          - "obs_history": 观测历史 (N, 225) = VAE 编码器输入
          - "privileged_obs": 特权观测 (N, 247), 其中 [45:48] 是 base_vel

        损失组成：
          1. 速度估计损失：MSE(code_vel, base_vel_gt)
          2. 观测重建损失：MSE(decoded_obs, policy_obs)
          3. KL 散度：beta * KL(q(z|history) || N(0,1))
        """
        # 从 actor 获取最近一次 CENet 前向输出
        cenet_out = getattr(self.actor, "_last_cenet_output", None)
        if cenet_out is None:
            return torch.tensor(0.0, device=obs_batch.device)

        code, code_vel, decode, mean_vel, logvar_vel, mean_latent, logvar_latent = cenet_out

        # 速度估计目标：privileged_obs 中的 base_vel（索引 45:48）
        vel_target = obs_batch["privileged_obs"][:, 45:48]
        # 观测重建目标：当前 policy obs
        obs_target = obs_batch["policy"]

        mse = nn.functional.mse_loss
        estimation_loss = mse(code_vel, vel_target)
        reconstruction_loss = mse(decode, obs_target)
        # KL 散度：上游用 raw sum（非 mean），然后除以 num_mini_batches
        logvar_latent = torch.clamp(logvar_latent, -20.0, 10.0)
        kl_loss = -0.5 * torch.sum(
            1 + logvar_latent - mean_latent.pow(2) - logvar_latent.exp()
        )

        # 与上游完全一致：除以 num_mini_batches
        autoenc_loss = (
            estimation_loss + reconstruction_loss + self.vae_beta * kl_loss
        ) / self.num_mini_batches
        if torch.isnan(autoenc_loss) or torch.isinf(autoenc_loss):
            return torch.tensor(0.0, device=obs_batch.device)
        return autoenc_loss
