# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
"""DreamWaQ PPO — aligned with Manaro-Alpha/DreamWaQ.

CENet (VAE):  history(225) → [128,64] → latent(16) + vel_est(3) = code(19)
Decoder:       code(19) → [64,128] → next_obs(45)
Actor:         code(19) + obs(45) = 64 → [512,256,128] → action(12)
Critic:        privileged_obs(235) → [512,256,128] → value(1)

VAE Loss: reconstruction_MSE + velocity_MSE + beta * KL
"""

import flax.linen as nn
import jax
import jax.numpy as jnp
import numpy as np
import optax
import pickle, os

from skrl.agents.jax.ppo import PPO as BasePPO
from skrl.memories.jax import RandomMemory
from skrl.models.jax import DeterministicMixin, GaussianMixin, Model
from skrl.trainers.jax import SequentialTrainer
from skrl.utils import set_seed

from motrix_envs import registry as env_registry
from motrix_envs.locomotion.go1.dreamwaq import DreamWaQTask
from motrix_rl import registry, utils
from motrix_rl.skrl import get_log_dir

# Ensure RL config registered
import motrix_rl.tasks.go1_dreamwaq  # noqa: F401


# ═══════════════════════════════════════════════════════════════════════
# CENet (VAE) — exact Manaro-Alpha architecture
# ═══════════════════════════════════════════════════════════════════════

class CENet(nn.Module):
    """VAE for terrain imagination + velocity estimation.

    Encoder: 225 → Linear(128) → ELU → Linear(64) → ELU
    4 heads on 64-dim bottleneck:
      - latent_mu, latent_logvar: 64 → 16 each
      - vel_mu, vel_logvar: 64 → 3 each
    Decoder: (latent+vel=19) → Linear(64) → ELU → Linear(128) → ELU → Linear(45)
    """
    num_obs: int = 45
    num_history: int = 5
    num_latent: int = 16   # latent z dim
    cenet_out_dim: int = 19  # code = vel(3) + latent(16)

    def setup(self):
        # Orthogonal weight init matching upstream torch.nn.init.orthogonal_
        kernel_init = nn.initializers.orthogonal(jnp.sqrt(2))

        # Encoder (matches Manaro-Alpha: 225 → 128 → 64)
        self.enc_fc1 = nn.Dense(128, kernel_init=kernel_init)
        self.enc_fc2 = nn.Dense(64, kernel_init=kernel_init)

        # 4 output heads
        self.latent_mu = nn.Dense(self.num_latent, kernel_init=kernel_init)
        self.latent_logvar = nn.Dense(self.num_latent, kernel_init=kernel_init)
        self.vel_mu = nn.Dense(3, kernel_init=kernel_init)
        self.vel_logvar = nn.Dense(3, kernel_init=kernel_init)

        # Decoder (matches Manaro-Alpha: 19 → 64 → 128 → 45)
        self.dec_fc1 = nn.Dense(64, kernel_init=kernel_init)
        self.dec_fc2 = nn.Dense(128, kernel_init=kernel_init)
        self.dec_out = nn.Dense(self.num_obs, kernel_init=kernel_init)

    def encode(self, obs_history):
        """obs_history: (B, T, 45) or (B, T*45) → (latent_mu, latent_logvar, vel_mu, vel_logvar)."""
        x = obs_history.reshape(obs_history.shape[0], -1)
        x = nn.elu(self.enc_fc1(x))
        x = nn.elu(self.enc_fc2(x))
        return (self.latent_mu(x), self.latent_logvar(x),
                self.vel_mu(x), self.vel_logvar(x))

    def reparameterize(self, mu, logvar, rng):
        std = jnp.exp(0.5 * logvar)
        return mu + std * jax.random.normal(rng, mu.shape)

    def decode(self, z, vel):
        x = jnp.concatenate([z, vel], axis=-1)
        x = nn.elu(self.dec_fc1(x))
        x = nn.elu(self.dec_fc2(x))
        return self.dec_out(x)

    def __call__(self, obs_history, rng):
        lm, ll, vm, vl = self.encode(obs_history)
        rng_z, rng_v = jax.random.split(rng)
        z = self.reparameterize(lm, ll, rng_z)
        vel = self.reparameterize(vm, vl, rng_v)
        return {
            'z': z, 'vel_est': vel,
            'latent_mu': lm, 'latent_logvar': ll,
            'vel_mu': vm, 'vel_logvar': vl,
            'obs_pred': self.decode(z, vel),
        }

    def sample(self, obs_history, rng):
        """Sample (z, vel) for rollout (adds noise)."""
        lm, ll, vm, vl = self.encode(obs_history)
        rng_z, rng_v = jax.random.split(rng)
        z = self.reparameterize(lm, ll, rng_z)
        vel = self.reparameterize(vm, vl, rng_v)
        return z, vel

    def inference(self, obs_history):
        """Mean (z, vel) for deployment (no noise)."""
        lm, ll, vm, vl = self.encode(obs_history)
        return lm, vm

    def loss_fn(self, obs_history, current_obs, base_vel, rng, kl_weight=1.0):
        """VAE loss: reconstructs current obs (last history frame), NOT next obs.

        Upstream: decode_target = obs_batch (the same timestep's observation,
        not a future prediction). This is an autoencoder over current observation.
        """
        lm, ll, vm, vl = self.encode(obs_history)
        rng_z, rng_v = jax.random.split(rng)
        z = self.reparameterize(lm, ll, rng_z)
        vel = self.reparameterize(vm, vl, rng_v)
        pred = self.decode(z, vel)

        recons = jnp.mean(jnp.square(pred - current_obs), axis=-1)
        vel_loss = jnp.mean(jnp.square(vel - base_vel), axis=-1)
        kld = -0.5 * jnp.sum(1 + ll - jnp.square(lm) - jnp.exp(ll), axis=-1)

        return {
            'loss': jnp.mean(recons + vel_loss + kl_weight * kld),
            'recons_loss': jnp.mean(recons),
            'vel_loss': jnp.mean(vel_loss),
            'kld_loss': jnp.mean(kld),
        }


# ═══════════════════════════════════════════════════════════════════════
# Environment Wrapper
# ═══════════════════════════════════════════════════════════════════════

class DreamWaQWrapper:
    """Wraps DreamWaQTask, computes CENet code online.

    Augmented observation layout (254-dim):
      [code_vel(3) | code_latent(16) | obs(45) | base_vel(3) | heights(187)]

    Actor uses: [0:64]  = code(19) + obs(45)
    Critic uses: [19:254] = obs(45) + base_vel(3) + heights(187)
    """

    def __init__(self, env, cenet, cenet_params, rng=None, vae_buf_size=24):
        self._env = env
        self._cenet = cenet
        self._cenet_params = cenet_params
        self._num_envs = env.num_envs
        self._rng = rng if rng is not None else jax.random.PRNGKey(0)
        self._code_dim = 19   # vel(3) + latent(16)
        self._aug_dim = self._code_dim + 45 + 3 + 187  # 254

        from gymnasium import spaces
        self._obs_space = spaces.Box(-np.inf, np.inf, (self._aug_dim,), dtype=np.float32)

        # VAE training buffer
        self._vae_buf = []
        self._vae_buf_size = vae_buf_size

    @property
    def num_envs(self): return self._num_envs
    @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 privileged_obs(self):
        return self._env._state.info.get("privileged_obs",
                np.zeros((self._num_envs, 235), dtype=np.float32))

    @property
    def obs_history(self):
        return self._env._state.info.get("obs_history",
                np.zeros((self._num_envs, 5, 45), dtype=np.float32))

    @property
    def base_vel(self):
        return self._env._state.info.get("base_vel",
                np.zeros((self._num_envs, 3), dtype=np.float32))

    @property
    def vae_training_data(self):
        return self._vae_buf

    def _augment(self, obs):
        """Compute CENet code and build augmented obs."""
        hist = jnp.array(self.obs_history)
        self._rng, rng_key = jax.random.split(self._rng)
        z, vel = self._cenet.apply(self._cenet_params, hist, rng_key, method=self._cenet.sample)
        code = np.concatenate([np.array(vel), np.array(z)], axis=-1)  # (N, 19)

        priv = self.privileged_obs
        heights = priv[:, 45+3:]  # (N, 187)
        base_vel_norm = self.base_vel * np.array([2.0, 2.0, 1.0], dtype=np.float32)

        return np.concatenate([code, obs, base_vel_norm, heights], axis=-1)

    def reset(self):
        state = self._env.init_state()
        self._vae_buf = []
        return self._augment(state.obs), state.info

    def step(self, actions):
        old_hist = self.obs_history.copy()
        state = self._env.step(np.array(actions, dtype=np.float32))

        # Store for VAE training: history → current_obs (last frame of history)
        # Upstream: decode target = obs_batch (current observation, not next)
        self._vae_buf.append((old_hist, old_hist[:, -1, :].copy(), self.base_vel.copy()))
        if len(self._vae_buf) > self._vae_buf_size:
            self._vae_buf.pop(0)

        return (self._augment(state.obs),
                state.reward.reshape(-1, 1).astype(np.float32),
                state.terminated.reshape(-1, 1).astype(np.float32),
                state.truncated.reshape(-1, 1).astype(np.float32),
                state.info)

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


# ═══════════════════════════════════════════════════════════════════════
# DreamWaQ Agent — PPO + CENet training
# ═══════════════════════════════════════════════════════════════════════

class DreamWaQAgent(BasePPO):
    """PPO agent with asymmetric critic + CENet (VAE) training."""

    def __init__(self, models, memory, cfg, observation_space, action_space, device,
                 cenet=None, cenet_params=None, wrapper=None,
                 kl_weight=1.0, vae_lr=1e-3, save_dir=None):
        super().__init__(models=models, memory=memory, cfg=cfg,
                         observation_space=observation_space,
                         action_space=action_space, device=device)
        self._cenet = cenet
        self._cenet_params = cenet_params
        self._wrapper = wrapper
        self._kl_weight = kl_weight
        self._save_dir = save_dir
        self._vae_count = 0
        self._rng = jax.random.PRNGKey(42)

        if cenet_params is not None:
            self._vae_opt = optax.adam(vae_lr)
            self._vae_opt_state = self._vae_opt.init(cenet_params)
        else:
            self._vae_opt = None

    # Override SKRL's internal _update (NOT 'update')
    def _update(self, timestep: int, timesteps: int):
        super()._update(timestep, timesteps)

        # Track detailed reward terms
        wrapper = self._wrapper
        if wrapper is not None:
            try:
                terms = wrapper._env._state.info.get("reward_terms", {})
                for name, val in terms.items():
                    self.tracking_data[f"Reward / {name}"].append(float(val))
            except Exception:
                pass

        # CENet training after each PPO update
        if self._cenet_params is not None and self._vae_opt is not None:
            self._vae_step()

    def _vae_step(self):
        """Train CENet: reconstruct current obs (last history frame) + velocity supervision + KL."""
        wrapper = self._wrapper
        if wrapper is None or len(wrapper.vae_training_data) == 0:
            return

        hist_np, current_obs_np, base_vel_np = wrapper.vae_training_data[-1]
        hist = jnp.array(hist_np)
        current_obs = jnp.array(current_obs_np)
        base_vel = jnp.array(base_vel_np)

        n = wrapper.num_envs
        valid = jnp.linalg.norm(hist.reshape(n, -1), axis=-1) > 1e-6
        if not valid.any():
            return

        self._rng, vae_rng = jax.random.split(self._rng)

        def loss_fn(p):
            return self._cenet.apply(
                p, hist, current_obs, base_vel, vae_rng, self._kl_weight,
                method=self._cenet.loss_fn)['loss']

        loss, grads = jax.value_and_grad(loss_fn)(self._cenet_params)
        updates, self._vae_opt_state = self._vae_opt.update(
            grads, self._vae_opt_state, self._cenet_params)
        self._cenet_params = optax.apply_updates(self._cenet_params, updates)

        # Sync wrapper
        wrapper._cenet_params = self._cenet_params

        self._vae_count += 1
        if self._save_dir and self._vae_count % 200 == 0:
            pickle.dump(self._cenet_params,
                        open(f"{self._save_dir}/vae_{self._vae_count}.pkl", "wb"))

        self.tracking_data["Loss / VAE"].append(loss)


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

class DreamWaQTrainer:
    def __init__(self, env_name="go1-dreamwaq-walk", cfg_override=None):
        self._env_name = env_name
        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

        # Init CENet
        rng = jax.random.PRNGKey(cfg.runner.seed)
        rng, crng = jax.random.split(rng)
        cenet = CENet(num_obs=45, num_history=env_cfg.num_history,
                       num_latent=env_cfg.num_latent,
                       cenet_out_dim=env_cfg.cenet_out_dim)
        cenet_params = cenet.init(crng, jnp.ones((1, env_cfg.num_history, 45)), jax.random.PRNGKey(0))

        # Wrapper
        wrapper = DreamWaQWrapper(env, cenet, cenet_params,
                                   rng=jax.random.PRNGKey(cfg.runner.seed + 1))
        env.init_state()

        print(f"[DreamWaQ] {cfg.num_envs} envs, CENet[128,64], latent=16+3=19, "
              f"aug_dim={wrapper.observation_space.shape[0]}")

        # Build models
        from motrix_rl.skrl.jax.train.ppo import Trainer as StdTrainer, _add_runtime_config

        policy_cfg = cfg.runner.models.policy
        value_cfg = cfg.runner.models.value

        class Policy(GaussianMixin, Model):
            def __init__(self, observation_space, action_space, device=None, **kwargs):
                Model.__init__(self, observation_space, action_space, device, **kwargs)
                GaussianMixin.__init__(self, policy_cfg.clip_actions,
                                       policy_cfg.clip_log_std,
                                       policy_cfg.min_log_std, policy_cfg.max_log_std,
                                       policy_cfg.reduction)

            @nn.compact
            def __call__(self, inputs, role):
                kernel_init = nn.initializers.orthogonal(jnp.sqrt(2))
                x = inputs["states"]
                # Actor: code(19) + obs(45) = first 64 dims
                x = x[:, :64]
                for d in policy_cfg.hiddens:
                    x = nn.elu(nn.Dense(d, kernel_init=kernel_init)(x))
                x = nn.Dense(self.num_actions, kernel_init=kernel_init)(x)
                log_std = self.param("log_std",
                    lambda _: jnp.full(self.num_actions, float(policy_cfg.initial_log_std)))
                return x, log_std, {}

        class Value(DeterministicMixin, Model):
            def __init__(self, observation_space, action_space, device=None, **kwargs):
                Model.__init__(self, observation_space, action_space, device, **kwargs)
                DeterministicMixin.__init__(self, value_cfg.clip_actions)

            @nn.compact
            def __call__(self, inputs, role):
                kernel_init = nn.initializers.orthogonal(jnp.sqrt(2))
                x = inputs["states"]
                # Critic: obs(45) + base_vel(3) + heights(187) = 235
                # Layout: [code(19) | obs(45) | base_vel(3) | heights(187)]
                x_c = jnp.concatenate([x[:, 19:64], x[:, 64:254]], axis=-1)
                for d in value_cfg.hiddens:
                    x_c = nn.elu(nn.Dense(d, kernel_init=kernel_init)(x_c))
                return nn.Dense(1, kernel_init=kernel_init)(x_c), {}

        models = {
            "policy": Policy(wrapper.observation_space, wrapper.action_space, None),
            "value": Value(wrapper.observation_space, wrapper.action_space, None),
        }
        for role, m in models.items():
            m.init_state_dict(role)

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

        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 = DreamWaQAgent(
            models=models, memory=memory, cfg=ppo_cfg,
            observation_space=wrapper.observation_space,
            action_space=wrapper.action_space, device=None,
            cenet=cenet, cenet_params=cenet_params, wrapper=wrapper,
            kl_weight=1.0, vae_lr=1e-3, save_dir=log_dir,
        )

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

        # Save final VAE params
        pickle.dump(agent._cenet_params, open(f"{log_dir}/cenet_params.pkl", "wb"))
        print(f"[DreamWaQ] CENet weights saved → {log_dir}")

    def play(self, policy_path=None):
        """Render env with trained policy (uses wrapper + CENet)."""
        from motrix_rl.skrl.jax.wrap_np import SkrlNpWrapper as RenderWrapper
        cfg = self._rlcfg
        raw_env = env_registry.make(self._env_name, num_envs=min(64, cfg.play_num_envs))
        env = RenderWrapper(raw_env)  # adds render() method
        env_cfg = raw_env._cfg

        import msgpack
        rng = jax.random.PRNGKey(cfg.runner.seed)
        rng, crng = jax.random.split(rng)
        cenet = CENet()
        log_dir = get_log_dir(self._env_name, rllib="skrl", agent_name="DreamWaQ_PPO")

        # Load CENet params
        if policy_path and os.path.exists(policy_path):
            # Load full agent checkpoint and extract
            with open(policy_path, 'rb') as f:
                ckpt = pickle.load(f)
            # Load policy params
            policy_raw = msgpack.unpackb(ckpt['policy'])
            policy_params = {}
            for name, val in policy_raw['params'].items():
                if isinstance(val, dict):
                    policy_params[name] = {
                        k: _decode_msgpack_array(v) for k, v in val.items()
                    }
                else:
                    policy_params[name] = _decode_msgpack_array(val)

            # Try loading VAE from base skrl dir
            log_dir = get_log_dir(self._env_name, rllib="skrl", agent_name="DreamWaQ_PPO")
            vae_path = os.path.join(log_dir, "cenet_params.pkl")
            if not os.path.exists(vae_path):
                # Find latest vae_N.pkl
                vae_files = sorted([f for f in os.listdir(log_dir) if f.startswith("vae_")])
                if vae_files:
                    vae_path = os.path.join(log_dir, vae_files[-1])
            if os.path.exists(vae_path):
                with open(vae_path, 'rb') as f:
                    cenet_params = pickle.load(f)
                print(f"[Play] VAE loaded: {vae_path}")
            else:
                print("[Play] WARNING: VAE params not found, using random init")
                cenet_params = cenet.init(crng, jnp.ones((1, 5, 45)), jax.random.PRNGKey(0))
        else:
            print("[Play] No policy, random actions")
            policy_params = None
            cenet_params = cenet.init(crng, jnp.ones((1, 5, 45)), jax.random.PRNGKey(0))

        wrapper = DreamWaQWrapper(raw_env, cenet, cenet_params,
                                   rng=jax.random.PRNGKey(cfg.runner.seed + 1))
        raw_env.init_state()

        import time
        n = raw_env._num_envs
        print(f"[Play] {n} envs, press Ctrl+C to stop")
        try:
            while True:
                if policy_params is not None:
                    aug_obs = wrapper._augment(raw_env._state.obs)
                    actions = np.array(_policy_forward(jnp.array(aug_obs), policy_params))
                else:
                    actions = np.zeros((n, 12), dtype=np.float32)

                wrapper.step(actions)
                env.render()
                time.sleep(0.01)
        except KeyboardInterrupt:
            env.close()
            print("[Play] Stopped")


def _decode_msgpack_array(ext) -> np.ndarray | None:
    """Decode flax-serialized msgpack ExtType."""
    import msgpack
    if not hasattr(ext, "code"): return None
    parts = msgpack.unpackb(ext.data, raw=False)
    if not isinstance(parts, list) or len(parts) < 3: return None
    shape = []
    def _flatten(s):
        if isinstance(s, list):
            for x in s: _flatten(x)
        elif isinstance(s, int): shape.append(s)
    _flatten(parts[0])
    return np.frombuffer(parts[2], dtype=np.dtype(parts[1])).reshape(shape)


def _policy_forward(x, params):
    """Forward pass through policy (first 64 dims only, apply elu)."""
    import jax.numpy as jnp
    import jax
    w0, b0 = jnp.array(params['Dense_0']['kernel']), jnp.array(params['Dense_0']['bias'])
    w1, b1 = jnp.array(params['Dense_1']['kernel']), jnp.array(params['Dense_1']['bias'])
    w2, b2 = jnp.array(params['Dense_2']['kernel']), jnp.array(params['Dense_2']['bias'])
    w3, b3 = jnp.array(params['Dense_3']['kernel']), jnp.array(params['Dense_3']['bias'])
    x = x[:, :64]  # code(19) + obs(45)
    x = jax.nn.elu(x @ w0 + b0)
    x = jax.nn.elu(x @ w1 + b1)
    x = jax.nn.elu(x @ w2 + b2)
    return x @ w3 + b3
