# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
"""Adapter: MotrixLab DreamWaQ NpEnv -> upstream rsl_rl-1.0.2 VecEnv interface.

Upstream OnPolicyRunner expects env.step() to return a 7-tuple:
    (obs, privileged_obs, prev_privileged_obs, obs_hist, rewards, dones, extras)
and env.reset() to return (obs, privileged_obs, prev_privileged_obs, obs_hist).

MotrixLab's DreamWaQTask (numpy, CPU physics) already provides:
    state.obs                  -> obs           (N, 45)
    info["privileged_obs"]     -> privileged    (N, 235) = obs(45)+base_vel(3)+heights(187)
    info["obs_history"]        -> history       (N, 5, 45)  (flattened to 225)
    state.terminated/truncated -> dones / time_outs

This wrapper does numpy<->torch transfer and maintains prev_privileged_obs.
"""
import numpy as np
import torch


class DwaqVecEnv:
    """Wraps a MotrixLab DreamWaQ NpEnv for the upstream DreamWaQ rsl_rl runner."""

    def __init__(self, env, device, num_obs=45, num_privileged_obs=235,
                 num_obs_hist=5, num_actions=12, clip_actions=100.0):
        self._env = env
        self.device = device
        self.num_envs = env.num_envs
        self.num_obs = num_obs
        self.num_privileged_obs = num_privileged_obs
        self.num_obs_hist = num_obs_hist
        self.num_actions = num_actions
        self._clip_actions = clip_actions
        self.max_episode_length = int(getattr(env.cfg, "max_episode_steps", None) or 1000)

        n = self.num_envs
        self.obs_buf = torch.zeros(n, num_obs, device=device)
        self.obs_hist_buf = torch.zeros(n, num_obs_hist * num_obs, device=device)
        self.privileged_obs_buf = torch.zeros(n, num_privileged_obs, device=device)
        self.prev_privileged_obs_buf = torch.zeros(n, num_privileged_obs, device=device)
        self.rew_buf = torch.zeros(n, device=device)
        self.reset_buf = torch.ones(n, device=device, dtype=torch.long)
        self.episode_length_buf = torch.zeros(n, device=device, dtype=torch.long)
        self.extras = {}

        self._env.init_state()
        self._sync_from_state(self._env._state)

    # ── helpers ──
    def _t(self, arr):
        return torch.from_numpy(np.ascontiguousarray(arr, dtype=np.float32)).to(self.device)

    def _sync_from_state(self, state):
        self.obs_buf = self._t(state.obs)
        info = state.info
        priv = info.get("privileged_obs")
        if priv is None:
            priv = np.zeros((self.num_envs, self.num_privileged_obs), dtype=np.float32)
        self.privileged_obs_buf = self._t(priv)
        hist = info.get("obs_history")  # (N, 5, 45)
        if hist is None:
            hist = np.zeros((self.num_envs, self.num_obs_hist, self.num_obs), dtype=np.float32)
        self.obs_hist_buf = self._t(hist).reshape(self.num_envs, -1)
        self.episode_length_buf = torch.from_numpy(
            np.asarray(info.get("steps", np.zeros(self.num_envs)), dtype=np.int64)).to(self.device)

    # ── VecEnv interface ──
    def step(self, actions):
        actions = torch.clip(actions, -self._clip_actions, self._clip_actions)
        actions_np = actions.detach().cpu().numpy().astype(np.float32)
        # prev_privileged = privileged BEFORE this step (matches upstream)
        self.prev_privileged_obs_buf = self.privileged_obs_buf.clone()
        state = self._env.step(actions_np)
        self._sync_from_state(state)

        done = np.logical_or(state.terminated, state.truncated)
        self.rew_buf = self._t(state.reward).reshape(-1)
        self.reset_buf = torch.from_numpy(done.astype(np.int64)).to(self.device)
        self.extras = {
            "time_outs": torch.from_numpy(np.asarray(state.truncated, dtype=np.float32)).to(self.device),
        }
        # pass per-term reward logging through if present
        terms = state.info.get("reward_terms")
        if terms:
            self.extras["episode"] = {f"rew_{k}": float(v) for k, v in terms.items()}
        # Log terrain_level so tensorboard tracks curriculum progress.
        tl = state.info.get("terrain_level")
        if tl is not None:
            self.extras["episode"] = self.extras.get("episode", {})
            self.extras["episode"]["terrain_level"] = float(np.asarray(tl).mean())

        # rewards/dones are (N,) like upstream (storage views them to (N,1) internally)
        return (self.obs_buf, self.privileged_obs_buf, self.prev_privileged_obs_buf,
                self.obs_hist_buf, self.rew_buf, self.reset_buf, self.extras)

    def reset(self, env_ids=None):
        self._env.init_state()
        self._sync_from_state(self._env._state)
        self.prev_privileged_obs_buf = torch.zeros_like(self.privileged_obs_buf)
        return (self.obs_buf, self.privileged_obs_buf,
                self.prev_privileged_obs_buf, self.obs_hist_buf)

    def get_observations(self):
        return self.obs_buf, self.obs_hist_buf

    def get_privileged_observations(self):
        return self.privileged_obs_buf, self.prev_privileged_obs_buf
