# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
"""CENet Actor 模型——DreamWaQ 的 VAE 编码器 + MLP Actor。

架构：
  CENetVAE:  obs_history(225) → encoder[128,64] → latent(16) + vel_est(3) = code(19)
             code(19) → decoder[64,128] → 重建 obs(45)
  CENetActorModel (继承 MLPModel):
              code(19) + obs(45) = 64 → MLP[512,256,128] → action(12)

用法：
  class_name = "motrix_rl.rslrl.torch.models.cenet_actor:CENetActorModel"
  obs_groups = {"actor": ["policy", "obs_history"], "critic": ["privileged_obs"]}
"""

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

from rsl_rl.models.mlp_model import MLPModel
from rsl_rl.modules import EmpiricalNormalization, HiddenState


# ═══ CENet VAE ═══

class CENetVAE(nn.Module):
    """CENet VAE 编码器-解码器。

    obs_history(225) → encoder → latent(16) + vel_est(3) = code(19)
    code(19) → decoder → 重建 obs(45)
    """

    def __init__(self, cenet_in_dim: int = 225, cenet_out_dim: int = 19,
                 activation: str = "elu"):
        super().__init__()
        act = _get_activation(activation)
        self.cenet_in_dim = cenet_in_dim
        self.cenet_out_dim = cenet_out_dim

        # 编码器：225 → 128 → 64
        self.encoder = nn.Sequential(
            nn.Linear(cenet_in_dim, 128), act,
            nn.Linear(128, 64), act,
        )
        # 潜变量头：64 → 16 (mean + logvar)
        self.encode_mean_latent = nn.Linear(64, cenet_out_dim - 3)
        self.encode_logvar_latent = nn.Linear(64, cenet_out_dim - 3)
        # 速度估计头：64 → 3 (mean + logvar)
        self.encode_mean_vel = nn.Linear(64, 3)
        self.encode_logvar_vel = nn.Linear(64, 3)
        # 解码器：19 → 64 → 128 → 45
        self.decoder = nn.Sequential(
            nn.Linear(cenet_out_dim, 64), act,
            nn.Linear(64, 128), act,
            nn.Linear(128, 45),
        )

    def reparameterise(self, mean: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
        """重参数化技巧：从 N(mean, exp(logvar/2)) 采样。"""
        std = torch.exp(logvar * 0.5)
        eps = torch.randn_like(std)
        return mean + std * eps

    def forward(self, obs_history: torch.Tensor):
        """前向传播。

        Returns:
            code:       (N, 19)  潜变量 [vel_sample(3) + latent_sample(16)]
            code_vel:   (N, 3)   速度估计采样
            decode:     (N, 45)  重建观测
            mean_vel:   (N, 3)   速度估计均值
            logvar_vel: (N, 3)   速度估计对数方差
            mean_latent:   (N, 16)  潜变量均值
            logvar_latent: (N, 16)  潜变量对数方差
        """
        h = self.encoder(obs_history)
        mean_latent = self.encode_mean_latent(h)
        logvar_latent = self.encode_logvar_latent(h)
        mean_vel = self.encode_mean_vel(h)
        logvar_vel = self.encode_logvar_vel(h)
        code_latent = self.reparameterise(mean_latent, logvar_latent)
        code_vel = self.reparameterise(mean_vel, logvar_vel)
        code = torch.cat((code_vel, code_latent), dim=-1)
        decode = self.decoder(code)
        return code, code_vel, decode, mean_vel, logvar_vel, mean_latent, logvar_latent

    def deterministic_code(self, obs_history: torch.Tensor) -> torch.Tensor:
        """推理模式：使用均值而非采样，产生确定性 code(19)。"""
        h = self.encoder(obs_history)
        mean_latent = self.encode_mean_latent(h)
        mean_vel = self.encode_mean_vel(h)
        return torch.cat((mean_vel, mean_latent), dim=-1)


# ═══ CENet Actor Model ═══

class CENetActorModel(MLPModel):
    """CENet Actor：继承 MLPModel，在 get_latent() 中注入 VAE code。

    class_name = "motrix_rl.rslrl.torch.models.cenet_actor:CENetActorModel"
    """

    def __init__(self, obs: TensorDict, obs_groups: dict[str, list[str]],
                 obs_set: str, output_dim: int,
                 cenet_in_dim: int = 225, cenet_out_dim: int = 19,
                 activation: str = "elu", **kwargs):
        # 必须在 super().__init__ 之前设置，因为 _get_latent_dim() 会被父类构造函数调用
        self._history_dim = cenet_in_dim
        self._code_dim = cenet_out_dim
        # 禁用观测归一化（VAE 输出已是归一化后的 code，维度也不匹配）
        kwargs["obs_normalization"] = False

        super().__init__(obs, obs_groups, obs_set, output_dim, **kwargs)

        # 在 nn.Module.__init__ 之后创建 VAE 子模块
        self.vae = CENetVAE(cenet_in_dim, cenet_out_dim, activation)
        self._last_cenet_output = None

    def _update_distribution(self, obs: torch.Tensor) -> None:
        """覆盖父类 — 强制 std > 0 再创建 Normal 分布（防止 NaN）。"""
        # 先 clamp std，再调父类创建分布
        if self.stochastic and not self.state_dependent_std:
            with torch.no_grad():
                if self.noise_std_type == "scalar":
                    self.std.nan_to_num_(nan=0.5, posinf=1.0, neginf=1.0)
                    self.std.clamp_(min=1e-6)
                elif self.noise_std_type == "log":
                    self.log_std.nan_to_num_(nan=0.0, posinf=5.0, neginf=-5.0)
                    self.log_std.clamp_(min=-20.0, max=10.0)
        super()._update_distribution(obs)

    def _get_latent_dim(self) -> int:
        """Actor 实际输入：code(19) + policy(45) = 64。"""
        return self.obs_dim - self._history_dim + self._code_dim

    def get_latent(self, obs: TensorDict, masks: torch.Tensor | None = None,
                   hidden_state: HiddenState = None) -> torch.Tensor:
        """提取观测 → VAE 编码 → 拼接 code + policy → 返回 latent(64)。"""
        policy_obs = obs["policy"]           # (N, 45)
        obs_history = obs["obs_history"]     # (N, 225)

        out = self.vae(obs_history)
        self._last_cenet_output = out
        code, code_vel, decode, mean_vel, logvar_vel, mean_latent, logvar_latent = out

        # 防止 VAE NaN 传播到下游
        if torch.isnan(code).any():
            code = torch.nan_to_num(code, nan=0.0)
        if torch.isnan(policy_obs).any():
            policy_obs = torch.nan_to_num(policy_obs, nan=0.0)

        latent = torch.cat([code, policy_obs], dim=-1)  # (N, 64)
        return latent

    def update_normalization(self, obs: TensorDict) -> None:
        """CENetActor 使用 obs_normalization=False，此方法为空。"""
        pass


# ═══ 工具函数 ═══

def _get_activation(act_name: str) -> nn.Module:
    """解析激活函数名称。"""
    _map = {
        "elu": nn.ELU, "selu": nn.SELU, "relu": nn.ReLU,
        "lrelu": nn.LeakyReLU, "tanh": nn.Tanh, "sigmoid": nn.Sigmoid,
    }
    if act_name in _map:
        return _map[act_name]()

    raise ValueError(f"未知激活函数: {act_name}。可选: {list(_map.keys())}")
