# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
"""CTS PPO v2 — Teacher/Student encoders with distillation loss.

Matches go2_rl_gym CTS training flow:
  Phase 1: Standard PPO update with teacher-computed latent
  Phase 2: Distillation: MSE(teacher_latent, student_latent) on student envs

Architecture:
  Teacher: privileged_obs(76) → MLP(512,256) → L2Norm → latent(32)
  Student: history_obs(225) → MLP(512,256,256) → L2Norm → latent(32)
  Actor:   student_obs(45) + latent(32) = 77 → [512,256,128] → action(12)
  Critic:  student_obs(45) + latent(32) = 77 → [512,256,128] → value(1)

The wrapper computes latent online (teacher for 75% envs, student for 25%)
and augments observations to 77-dim for standard PPO. After each PPO update,
a distillation step trains the student encoder to match the teacher's latent.
"""

import flax.linen as nn
import jax
import jax.numpy as jnp
import numpy as np
import optax

from skrl.agents.jax.ppo import PPO as BasePPO
from skrl.memories.jax import RandomMemory
from skrl.trainers.jax import SequentialTrainer
from skrl.utils import set_seed

from motrix_envs import registry as env_registry
from motrix_rl import registry, utils
from motrix_rl.skrl import get_log_dir


# ═══════════════════════════════════════════════════════════════════════════
# Encoder Modules (Flax)
# ═══════════════════════════════════════════════════════════════════════════

class TeacherEncoder(nn.Module):
    latent_dim: int = 32
    hidden_dims: tuple = (512, 256)

    @nn.compact
    def __call__(self, x):
        for d in self.hidden_dims:
            x = nn.elu(nn.Dense(d)(x))
        return nn.Dense(self.latent_dim)(x)


class StudentEncoder(nn.Module):
    latent_dim: int = 32
    hidden_dims: tuple = (512, 256, 256)

    @nn.compact
    def __call__(self, x):
        for d in self.hidden_dims:
            x = nn.elu(nn.Dense(d)(x))
        return nn.Dense(self.latent_dim)(x)


# ═══════════════════════════════════════════════════════════════════════════
# CTS Env Wrapper
# ═══════════════════════════════════════════════════════════════════════════

class CTSEnvWrapper:
    """Wraps CTS env. Computes latent online, returns augmented obs (77-dim)."""

    def __init__(self, env, teacher_ratio=0.75, latent_dim=32,
                 teacher_enc=None, teacher_params=None,
                 student_enc=None, student_params=None):
        self.__env = env
        self._teacher_ratio = teacher_ratio
        self._latent_dim = latent_dim
        self._num_envs = env.num_envs
        self._t_enc = teacher_enc
        self._t_params = teacher_params
        self._s_enc = student_enc
        self._s_params = student_params
        from gymnasium import spaces
        self._obs_space = spaces.Box(-np.inf, np.inf,
                                     (45 + latent_dim,), dtype=np.float32)

    @property
    def num_envs(self): return self._num_envs
    @property
    def num_teacher(self): return int(self._num_envs * self._teacher_ratio)
    @property
    def teacher_mask(self):
        m = np.zeros(self._num_envs, dtype=np.bool)
        m[:self.num_teacher] = True
        return m
    @property
    def student_mask(self): return ~self.teacher_mask
    @property
    def privileged_obs(self):
        return self.__env._state.info.get("privileged_obs",
                np.zeros((self._num_envs, 76), dtype=np.float32))
    @property
    def history_obs(self): return self.__env.history_obs
    @property
    def observation_space(self): return self._obs_space
    @property
    def action_space(self): return self.__env.action_space
    @property
    def device(self): return None
    @property
    def num_agents(self): return 1
    @property
    def state_preprocessor(self): return None
    @property
    def value_preprocessor(self): return None
    @property
    def world_info_idx(self): return None
    @property
    def random_timestep(self): return False
    @property
    def running(self): return True
    @property
    def _env(self): return self.__env

    def _augment(self, states):
        n = states.shape[0]
        latent = np.zeros((n, self._latent_dim), dtype=np.float32)
        t_mask = self.teacher_mask
        if t_mask.any() and self._t_params is not None:
            priv = jnp.array(self.privileged_obs[t_mask])
            latent[t_mask] = np.array(self._t_enc.apply(self._t_params, priv))
        s_mask = self.student_mask
        if s_mask.any() and self._s_params is not None:
            hist = jnp.array(self.history_obs[s_mask])
            latent[s_mask] = np.array(self._s_enc.apply(self._s_params, hist))
        return np.concatenate([states, latent], axis=-1)

    def reset(self):
        state = self.__env.init_state()
        return self._augment(state.obs), state.info

    def step(self, actions):
        state = self.__env.step(np.array(actions))
        return (self._augment(state.obs), state.reward.reshape(-1, 1),
                state.terminated.reshape(-1, 1),
                state.truncated.reshape(-1, 1), state.info)

    def render(self, *a, **kw): pass
    def close(self): pass


# ═══════════════════════════════════════════════════════════════════════════
# CTS Agent — adds distillation step after PPO update
# ═══════════════════════════════════════════════════════════════════════════

class CTSAgent(BasePPO):
    """PPO agent with post-update distillation on student encoder."""

    def __init__(self, models, memory, cfg, observation_space, action_space, device,
                 teacher_enc=None, student_enc=None, cts_wrapper=None,
                 teacher_params=None, student_params=None,
                 distill_coef=0.1, student_lr=1e-3):
        super().__init__(models=models, memory=memory, cfg=cfg,
                         observation_space=observation_space,
                         action_space=action_space, device=device)
        self._t_enc = teacher_enc
        self._s_enc = student_enc
        self._t_params = teacher_params
        self._s_params = student_params
        self._wrapper = cts_wrapper
        self._distill_coef = distill_coef
        self._latent_dim = getattr(teacher_enc, 'latent_dim', 32) if teacher_enc else 32

        # Separate optimizer for student encoder
        if student_params is not None:
            self._s_optimizer = optax.adam(student_lr)
            self._s_opt_state = self._s_optimizer.init(student_params)
        else:
            self._s_optimizer = None

    def update(self):
        """PPO update + distillation phase."""
        result = super().update()

        # Phase 2: Distillation — MSE(teacher_latent, student_latent) on student envs
        if self._s_params is not None and self._s_optimizer is not None:
            self._distillation_step()

        return result

    def _distillation_step(self):
        """Train student encoder to match teacher latent.

        go2_rl_gym Phase 2: only on student env samples.
        Loss = || teacher_enc(privileged_obs) - student_enc(history_obs) ||²
        """
        wrapper = self._wrapper
        if wrapper is None:
            return

        s_mask = wrapper.student_mask
        if not s_mask.any():
            return

        # Get current student env observations
        priv = jnp.array(wrapper.privileged_obs[s_mask])   # [S, 76]
        hist = jnp.array(wrapper.history_obs[s_mask])       # [S, 225]

        # L2 normalize both before computing MSE (match go2_rl_gym)
        def l2_norm(x):
            return x / (jnp.linalg.norm(x, axis=-1, keepdims=True) + 1e-8)

        def distill_loss(s_params):
            t_latent = l2_norm(jax.lax.stop_gradient(
                self._t_enc.apply(self._t_params, priv)))
            s_latent = l2_norm(self._s_enc.apply(s_params, hist))
            return jnp.mean(jnp.square(t_latent - s_latent))

        loss, grads = jax.value_and_grad(distill_loss)(self._s_params)
        updates, self._s_opt_state = self._s_optimizer.update(
            grads, self._s_opt_state, self._s_params)
        self._s_params = optax.apply_updates(self._s_params, updates)

        # Track
        self.tracking_data["Loss / Distillation"].append(loss)


# ═══════════════════════════════════════════════════════════════════════════
# Trainer
# ═══════════════════════════════════════════════════════════════════════════

class CTSTrainer:
    def __init__(self, env_name="go1-cts-flat-walk-go2style",
                 sim_backend=None, enable_render=False, cfg_override=None):
        self._env_name = env_name
        self._sim_backend = sim_backend
        rlcfg = registry.default_rl_cfg(env_name, "skrl", backend="jax")
        if cfg_override:
            rlcfg = utils.cfg_override(rlcfg, cfg_override)
        self._rlcfg = rlcfg

    def train(self):
        cfg = self._rlcfg
        set_seed(cfg.runner.seed)

        env = env_registry.make(self._env_name, num_envs=cfg.num_envs)
        env_cfg = env._cfg
        latent_dim = getattr(env_cfg, 'latent_dim', 32)
        priv_dim = getattr(env_cfg, 'num_privileged_obs', 76)
        hist_len = getattr(env_cfg, 'history_length', 5)
        student_dim = getattr(env_cfg, 'num_student_obs', 45)
        teacher_ratio = getattr(env_cfg, 'teacher_env_ratio', 0.75)

        # Init Flax encoders
        rng = jax.random.PRNGKey(cfg.runner.seed)
        rng, t_key, s_key = jax.random.split(rng, 3)
        teacher_enc = TeacherEncoder(latent_dim=latent_dim)
        student_enc = StudentEncoder(latent_dim=latent_dim)
        t_params = teacher_enc.init(t_key, jnp.ones((1, priv_dim)))
        s_params = student_enc.init(s_key, jnp.ones((1, hist_len * student_dim)))

        # CTS wrapper
        wrapper = CTSEnvWrapper(env, teacher_ratio=teacher_ratio, latent_dim=latent_dim,
                                teacher_enc=teacher_enc, teacher_params=t_params,
                                student_enc=student_enc, student_params=s_params)
        env.init_state()

        from gymnasium import spaces
        aug_obs = spaces.Box(-np.inf, np.inf, (student_dim + latent_dim,), dtype=np.float32)
        wrapper._obs_space = aug_obs

        nT = int(cfg.num_envs * teacher_ratio)
        nS = cfg.num_envs - nT
        print(f"[CTS v2] {cfg.num_envs} envs ({nT}T/{nS}S), latent={latent_dim}, "
              f"priv={priv_dim}, hist={hist_len}×{student_dim}, distill_coef=0.1")

        # Standard model creation
        from motrix_rl.skrl.jax.train.ppo import Trainer as StdTrainer, _add_runtime_config
        std = StdTrainer(self._env_name, sim_backend=None, enable_render=False)
        std._rlcfg = cfg
        models = std._make_model(wrapper, cfg)

        ppo_cfg = cfg.runner.agent.to_dict()
        _add_runtime_config(ppo_cfg, wrapper,
                           log_dir=get_log_dir(self._env_name, rllib="skrl", agent_name="CTS_PPO"))

        mem_size = cfg.runner.memory.memory_size
        if mem_size == -1:
            mem_size = ppo_cfg.get("rollouts", 24)
        memory = RandomMemory(memory_size=mem_size, num_envs=cfg.num_envs, device=None)

        agent = CTSAgent(
            models=models, memory=memory, cfg=ppo_cfg,
            observation_space=aug_obs, action_space=env.action_space,
            device=None,
            teacher_enc=teacher_enc, student_enc=student_enc,
            cts_wrapper=wrapper,
            teacher_params=t_params, student_params=s_params,
            distill_coef=0.1, student_lr=1e-3,
        )

        trainer_cfg = {"timesteps": cfg.runner.trainer.timesteps, "headless": True}
        SequentialTrainer(cfg=trainer_cfg, env=wrapper, agents=agent).train()

        # Save encoder params for deployment
        import pickle as _pk
        _log = get_log_dir(self._env_name, rllib="skrl", agent_name="CTS_PPO")
        _pk.dump(t_params, open(f"{_log}/teacher_params.pkl", "wb"))
        _pk.dump(s_params, open(f"{_log}/student_params.pkl", "wb"))
        print(f"[CTS] Encoder weights saved → {_log}")
