"""RX1 fully-dynamic robot environment (nu=40 joints + 2 weld bits).

Action vector layout (action_dim = nu + 2):
  action[0:nu]  — position targets for all 40 actuators (rad)
  action[nu]    — weld_red_cube  active bit (>0.5 → attach)
  action[nu+1]  — weld_blue_cube active bit (>0.5 → attach)

Proprio vector (nq + nv dims):
  qpos[0:nu_joints] + qvel[0:nu_joints]  (excludes free-joint cube DOFs)
"""
from __future__ import annotations

import re
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, Optional

import mujoco
import numpy as np

_ROOT    = Path(__file__).resolve().parents[1]
_XML     = _ROOT / "rx1_mujoco" / "rx1_scene.xml"
_MESHDIR = (_ROOT / "rx1_mujoco" / "meshes_mjcf").as_posix() + "/"


def _load_model() -> mujoco.MjModel:
    xml = _XML.read_text(encoding="utf-8")
    xml = re.sub(r'meshdir="[^"]*"', f'meshdir="{_MESHDIR}"', xml)
    return mujoco.MjModel.from_xml_string(xml)


from env.humanoid_env import RobotObs  # noqa: E402
from models.rx1_ik import discover_cube_welds  # noqa: E402


class RX1Env:
    """
    MuJoCo RX1 environment — all 40 joints fully dynamic.

    The action dimension is discovered from the loaded model (nu + 2 weld bits)
    so it stays correct after any rebuild_scene changes.
    """

    CAMERA_W = 224
    CAMERA_H = 224

    def __init__(self, config: dict):
        self.config = config

        self.model = _load_model()
        self.data  = mujoco.MjData(self.model)

        # Discover dims from the model. Every manipulable cube contributes one
        # weld bit, discovered dynamically so the action layout follows the scene
        # (adding a cube in build_scene.py needs no code change here).
        self._nu         = self.model.nu          # 40 joints
        self._cube_welds = discover_cube_welds(self.model)   # [(body, eq_id)]
        self.N_ACTUATORS = self._nu + len(self._cube_welds)  # joints + 1 bit/cube
        self.proprio_dim = self._nu * 2           # qpos + qvel for robot joints
        # action index of each cube's weld bit (nu + i) — matches the policy's
        # ordering since both derive it from discover_cube_welds().
        self._weld_bit = {body: self._nu + i
                          for i, (body, _eq) in enumerate(self._cube_welds)}

        vis_cfg       = config.get("vision", {})
        self.CAMERA_W = vis_cfg.get("width",  224)
        self.CAMERA_H = vis_cfg.get("height", 224)
        self._primary_cam = vis_cfg.get("camera", "ego")
        mv = vis_cfg.get("multi_view", {})
        self._multi_view_cameras: Dict[str, str] = {
            "overview": mv.get("overview", "overview"),
        }

        # Camera-based perception: when enabled, the ego (head) camera's RGB+depth
        # is rendered at a low rate and objects are localised by colour + depth
        # deprojection — genuine vision, not privileged sim state. Rendering runs
        # inside step() (the mj_step thread) so the GL context never conflicts
        # with the passive viewer; the result is cached for perceive_objects().
        self._vision_enabled = bool(vis_cfg.get("enabled", False))
        self._perceive_every = int(vis_cfg.get("perceive_every", 10))  # ctrl steps
        self._ego_rgb_r  = None
        self._ego_depth_r = None
        self._vision_objects: Dict[str, dict] = {}   # colour -> {"pos", "px"}
        self._perceive_n = 0

        sim_cfg = config.get("simulation", {})
        self.model.opt.timestep = sim_cfg.get("physics_timestep", 0.002)
        self._n_substeps = sim_cfg.get("n_substeps", 4)

        # Slew-rate limit on commanded joint targets. The PyBullet reference
        # (humanoid_torso/rx1_sim.py) caps joint speed at maxVelocity=3.0 rad/s;
        # we enforce the same by limiting how far each ctrl target may move per
        # control step. Without this a target that jumps (e.g. rest→pregrasp)
        # makes the position servo slam the joint and the robot flails.
        act_cfg = config.get("actuator", {})
        max_joint_speed = act_cfg.get("max_joint_speed", 3.0)   # rad/s
        ctrl_dt = self.model.opt.timestep * self._n_substeps
        self._max_dctrl = max_joint_speed * ctrl_dt             # rad per step

        self._renderer: Optional[mujoco.Renderer] = None

    # ------------------------------------------------------------------ public

    def reset(self, render_vision: bool = False) -> RobotObs:
        mujoco.mj_resetData(self.model, self.data)
        for _body, eq_id in self._cube_welds:
            self.data.eq_active[eq_id] = 0
        mujoco.mj_forward(self.model, self.data)
        # Seed the servo targets with the spawn configuration so the position
        # actuators hold the robot still until the brain issues commands
        # (prevents an initial lurch toward the all-zero ctrl pose).
        self.data.ctrl[:self._nu] = self.data.qpos[:self._nu]
        return self._make_obs(render_vision=render_vision)

    def step(self, action: np.ndarray,
             render_vision: bool = False,
             render_multi_view: bool = False) -> RobotObs:
        nu   = self._nu
        clip = self.config.get("actuator", {}).get("control_clip", 3.14)

        # Joint position targets, clamped to the configured control range and
        # the per-actuator limits baked into the model.
        ctrl = np.clip(action[:nu], -clip, clip)
        lo = self.model.actuator_ctrlrange[:, 0]
        hi = self.model.actuator_ctrlrange[:, 1]
        ctrl = np.clip(ctrl, lo, hi)

        # Slew-rate limit: move each target at most _max_dctrl toward the goal
        # this step, so joint speed stays bounded (≈ reference maxVelocity).
        prev = self.data.ctrl[:nu]
        delta = np.clip(ctrl - prev, -self._max_dctrl, self._max_dctrl)
        self.data.ctrl[:nu] = prev + delta

        # Weld bits — one per cube, in discover_cube_welds() order. Capture the
        # current relative pose on the rising edge so the grasped object holds
        # exactly where it sits in the palm (ported from
        # rx1_brain/control.py:set_grasp_weld) instead of snapping to the wrist.
        for i, (_body, eq_id) in enumerate(self._cube_welds):
            bit = nu + i
            if len(action) > bit:
                self._set_weld(eq_id, action[bit] > 0.5)

        for _ in range(self._n_substeps):
            mujoco.mj_step(self.model, self.data)

        # Camera perception on the mj_step thread (GL-safe), throttled.
        if self._vision_enabled:
            self._perceive_n += 1
            if self._perceive_n % self._perceive_every == 0:
                try:
                    self._update_vision()
                except Exception as e:        # never let a render error kill control
                    if self._perceive_n % (self._perceive_every * 20) == 0:
                        import logging
                        logging.getLogger("rx1_env").debug("vision render failed: %s", e)

        return self._make_obs(render_vision=render_vision,
                              render_multi_view=render_multi_view)

    def _set_weld(self, eq_id: int, active: bool):
        """Toggle a hand↔object weld, capturing relpose on activation."""
        was_active = bool(self.data.eq_active[eq_id])
        if active and not was_active:
            b1 = self.model.eq_obj1id[eq_id]
            b2 = self.model.eq_obj2id[eq_id]
            p1, r1 = self.data.xpos[b1], self.data.xmat[b1].reshape(3, 3)
            p2, r2 = self.data.xpos[b2], self.data.xmat[b2].reshape(3, 3)
            relpos = r1.T @ (p2 - p1)
            quat   = np.zeros(4)
            mujoco.mju_mat2Quat(quat, (r1.T @ r2).reshape(9))
            self.model.eq_data[eq_id, 0:3]  = 0.0
            self.model.eq_data[eq_id, 3:6]  = relpos
            self.model.eq_data[eq_id, 6:10] = quat
        self.data.eq_active[eq_id] = 1 if active else 0

    # ------------------------------------------------------------------ perception

    def perceive_objects(self) -> Dict[str, dict]:
        """
        Perceive manipulable objects. Each object's colour, identity and weld are
        known from the model; its POSITION comes from the ego camera (colour +
        depth deprojection) when vision is enabled and the object is in view —
        otherwise it falls back to the body's sim state. The 'source' field says
        which. Either way the target is not a hardcoded constant.
        """
        out: Dict[str, dict] = {}
        for j in range(self.model.njnt):
            if self.model.jnt_type[j] != mujoco.mjtJoint.mjJNT_FREE:
                continue
            bid  = self.model.jnt_bodyid[j]
            name = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_BODY, bid)
            qadr = self.model.jnt_qposadr[j]
            gt_pos = self.data.qpos[qadr:qadr + 3].copy()
            color, rgba = "object", None
            for g in range(self.model.ngeom):
                if self.model.geom_bodyid[g] == bid:
                    rgba = np.array(self.model.geom_rgba[g])
                    matid = int(self.model.geom_matid[g])
                    if matid >= 0:
                        rgba = np.array(self.model.mat_rgba[matid])
                    color = self._color_name(rgba)
                    break
            eq_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_EQUALITY,
                                      f"weld_{name}")
            # vision wins when this colour was detected in the ego camera
            seen = self._vision_objects.get(color)
            if seen is not None:
                pos, source = seen["pos"].copy(), "vision"
            else:
                pos, source = gt_pos, "state"
            out[name] = {"pos": pos, "color": color, "eq_id": int(eq_id),
                         "weld_bit": self._weld_bit.get(name),
                         "source": source, "rgba": None if rgba is None else np.array(rgba)}
        bid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "box")
        if bid >= 0:
            out["box"] = {"pos": self.data.xpos[bid].copy(), "color": "box",
                          "eq_id": -1, "weld_bit": None, "source": "state", "rgba": None}
        return out

    # ── camera-based object localisation ──────────────────────────────────────

    def _update_vision(self):
        """Render the ego camera and localise each coloured object by colour +
        depth deprojection. Updates self._vision_objects (colour -> pos)."""
        if self._ego_rgb_r is None:
            self._ego_rgb_r   = mujoco.Renderer(self.model, self.CAMERA_H, self.CAMERA_W)
            self._ego_depth_r = mujoco.Renderer(self.model, self.CAMERA_H, self.CAMERA_W)
            self._ego_depth_r.enable_depth_rendering()
        self._ego_rgb_r.update_scene(self.data, camera="ego")
        self._ego_depth_r.update_scene(self.data, camera="ego")
        rgb   = self._ego_rgb_r.render().astype(np.int32)
        depth = self._ego_depth_r.render()

        H, W = depth.shape
        r, g, b = rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2]
        near = (depth > 0.05) & (depth < 1.0)        # table region; rejects sky/floor
        masks = {
            "red":  near & (r > 120) & (g < 90) & (b < 90),
            "blue": near & (b > 140) & ((b - g) > 60) & (r >= g - 10),
            "green": near & (g > 120) & (r < 90) & (b < 90),
            "yellow": near & (r > 140) & (g > 140) & (b < 90),
        }
        cam = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_CAMERA, "ego")
        fovy = np.deg2rad(self.model.cam_fovy[cam])
        fy = (H / 2.0) / np.tan(fovy / 2.0); fx = fy
        cx, cy = W / 2.0, H / 2.0
        cpos = self.data.cam_xpos[cam]
        cmat = self.data.cam_xmat[cam].reshape(3, 3)

        found = {}
        for color, mask in masks.items():
            ys, xs = np.where(mask)
            if len(xs) < 8:                          # not confidently in view
                continue
            u, v = int(np.median(xs)), int(np.median(ys))
            dd = float(depth[v, u])
            p_cam = np.array([(u - cx) / fx * dd, -(v - cy) / fy * dd, -dd])
            found[color] = {"pos": cpos + cmat @ p_cam, "px": int(mask.sum())}
        self._vision_objects = found

    @staticmethod
    def _color_name(rgba) -> str:
        r, g, b = float(rgba[0]), float(rgba[1]), float(rgba[2])
        if r > 0.6 and g < 0.5 and b < 0.5:  return "red"
        if b > 0.6 and r < 0.6:              return "blue"
        if g > 0.6 and r < 0.6 and b < 0.6:  return "green"
        if r > 0.6 and g > 0.6 and b < 0.5:  return "yellow"
        return "object"

    # ------------------------------------------------------------------ obs

    def _make_obs(self, render_vision: bool,
                  render_multi_view: bool = False) -> RobotObs:
        vision = self._get_vision()      if render_vision    else None
        mv     = self._get_multi_view()  if render_multi_view else None
        return RobotObs(
            vision=vision,
            touch=np.zeros(0, dtype=np.float32),
            imu=np.zeros(6, dtype=np.float32),
            proprio=self._get_proprio(),
            sim_time=float(self.data.time),
            multi_view=mv,
        )

    def _get_vision(self) -> np.ndarray:
        self._ensure_renderer()
        self._renderer.update_scene(self.data, camera=self._primary_cam)
        return self._renderer.render().copy()

    def _get_multi_view(self) -> Dict[str, np.ndarray]:
        self._ensure_renderer()
        views: Dict[str, np.ndarray] = {}
        for role, cam in self._multi_view_cameras.items():
            self._renderer.update_scene(self.data, camera=cam)
            views[role] = self._renderer.render().copy()
        return views

    def _get_proprio(self) -> np.ndarray:
        nu = self._nu
        return np.concatenate([
            self.data.qpos[:nu].copy(),
            self.data.qvel[:nu].copy(),
        ]).astype(np.float32)

    def _ensure_renderer(self):
        if self._renderer is None:
            self._renderer = mujoco.Renderer(self.model, self.CAMERA_H, self.CAMERA_W)

    @property
    def dt(self) -> float:
        return self.model.opt.timestep * self._n_substeps

    def close(self):
        for r in (self._renderer, self._ego_rgb_r, self._ego_depth_r):
            try:
                if r is not None:
                    r.close()
            except Exception:
                pass
        self._renderer = self._ego_rgb_r = self._ego_depth_r = None