"""GPU-batched MuJoCo simulation using MJX (JAX backend).

MJX runs all environments in parallel on GPU via JAX, providing
~10-100x speedup over the CPU MuJoCoRunner for large env counts (1024+).

Requirements:
    pip install 'jax[cuda12]'    # NVIDIA GPU
    pip install jax               # CPU fallback
"""

import dataclasses
import os

import structlog
import torch

# JAX (MJX physics) shares the GPU with PyTorch (policy + optimizer). By
# default JAX preallocates ~75% of GPU memory on init, starving torch and
# causing OOM at the first PPO update. Disable preallocation so both libraries
# grow on demand — essential on small vGPU slices (e.g. a 6 GB HAMI slice).
os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")

try:
    import jax
    import jax.numpy as jnp
    import mujoco
    from mujoco import mjx
except ImportError as e:
    raise ImportError(
        "MJX runner requires JAX and MuJoCo MJX.  Install with:\n"
        "  pip install 'jax[cuda12]'   # GPU\n"
        "  pip install jax              # CPU\n"
    ) from e

import numpy as np

from src.core.env import BaseEnv
from src.core.runner import BaseRunner, BaseRunnerConfig
from src.runners.mujoco import (
    ActuatorLimits,
    load_mujoco_model,
)

log = structlog.get_logger()


@dataclasses.dataclass
class MJXRunnerConfig(BaseRunnerConfig):
    num_envs: int = 1024
    device: str = "cuda"
    dt: float = 0.002
    substeps: int = 10


class MJXRunner(BaseRunner[MJXRunnerConfig]):
    """GPU-batched MuJoCo runner using MJX (JAX).

    Physics runs entirely on GPU via JAX; observations flow to
    PyTorch through zero-copy DLPack transfers.
    """

    def __init__(self, env: BaseEnv, config: MJXRunnerConfig):
        super().__init__(env, config)

    @property
    def num_envs(self) -> int:
        return self.config.num_envs

    @property
    def device(self) -> torch.device:
        return torch.device(self.config.device)

    # ── Initialization ───────────────────────────────────────────────

    def _sim_initialize(self, config: MJXRunnerConfig) -> None:
        # Step 1: Load CPU model (reuses URDF → MJCF → actuator injection)
        self._mj_model = load_mujoco_model(self.env.robot)
        self._mj_model.opt.timestep = config.dt
        self._nq = self._mj_model.nq
        self._nv = self._mj_model.nv
        self._nu = self._mj_model.nu

        # Step 2: Put model on GPU
        self._mjx_model = mjx.put_model(self._mj_model)

        # Step 3: Default reset state on GPU
        default_data = mujoco.MjData(self._mj_model)
        mujoco.mj_forward(self._mj_model, default_data)
        self._default_mjx_data = mjx.put_data(self._mj_model, default_data)

        # Env-defined initial-state distribution (shared with the CPU
        # runner) — baked into the JIT reset as constants.
        qpos_lo, qpos_hi, qvel_lo, qvel_hi = self.env.initial_state_ranges(
            self._nq, self._nv,
        )
        self._init_qpos_lo = jnp.asarray(qpos_lo)
        self._init_qpos_hi = jnp.asarray(qpos_hi)
        self._init_qvel_lo = jnp.asarray(qvel_lo)
        self._init_qvel_hi = jnp.asarray(qvel_hi)

        # Step 4: Initialise all environments with randomized states
        self._rng = jax.random.PRNGKey(42)
        self._batch_data = self._make_batched_data(config.num_envs)

        # Step 4b: Build motor model info (ctrl_idx, qvel_idx, ActuatorConfig)
        self._motor_info: list[tuple[int, int]] = []
        self._motor_acts: list = []
        for ctrl_idx, act in enumerate(self.env.robot.actuators):
            if act.has_motor_model:
                jnt_id = mujoco.mj_name2id(
                    self._mj_model, mujoco.mjtObj.mjOBJ_JOINT, act.joint,
                )
                qvel_idx = self._mj_model.jnt_dofadr[jnt_id]
                self._motor_info.append((ctrl_idx, qvel_idx))
                self._motor_acts.append(act)

        # Step 5: JIT-compile the hot-path functions
        self._compile_jit_fns(config.substeps)

        # Keep one CPU MjData for offscreen rendering
        self._render_data = mujoco.MjData(self._mj_model)

        # Per-env DR scale arrays (synced from torch on every reset).
        # Initialised to 1.0 here because _setup_domain_rand runs after this.
        self._mjx_fr = jnp.ones(config.num_envs)
        self._mjx_dp = jnp.ones(config.num_envs)
        self._mjx_tq = jnp.ones(config.num_envs)

        log.info(
            "mjx_runner_ready",
            num_envs=config.num_envs,
            substeps=config.substeps,
            jax_devices=str(jax.devices()),
        )

    def _make_batched_data(self, n: int):
        """Create *n* environments with env-defined initial randomization."""
        self._rng, k1, k2 = jax.random.split(self._rng, 3)
        pq = jax.random.uniform(
            k1, (n, self._nq),
            minval=self._init_qpos_lo, maxval=self._init_qpos_hi,
        )
        pv = jax.random.uniform(
            k2, (n, self._nv),
            minval=self._init_qvel_lo, maxval=self._init_qvel_hi,
        )

        default = self._default_mjx_data
        model = self._mjx_model

        def _init_one(pq_i, pv_i):
            d = default.replace(
                qpos=default.qpos + pq_i,
                qvel=default.qvel + pv_i,
            )
            return mjx.forward(model, d)

        return jax.vmap(_init_one)(pq, pv)

    def _compile_jit_fns(self, substeps: int) -> None:
        """Pre-compile the two hot-path functions so the first call is fast."""
        model = self._mjx_model
        default = self._default_mjx_data

        lim = ActuatorLimits(self._mj_model)
        act_jnt_ids = jnp.array(lim.jnt_ids)
        act_limited = jnp.array(lim.limited)
        act_lo = jnp.array(lim.lo)
        act_hi = jnp.array(lim.hi)
        act_gs = jnp.array(lim.gear_sign)

        # ── Motor model params (JAX arrays for JIT) ─────────────────
        # Must stay in lock-step with ActuatorConfig.transform_ctrl() /
        # compute_motor_force() in src/core/robot.py — sysid fits against
        # the CPU implementation.
        _has_motor = len(self._motor_info) > 0
        if _has_motor:
            acts = self._motor_acts
            _ctrl_ids = jnp.array([c for c, _ in self._motor_info])
            _qvel_ids = jnp.array([q for _, q in self._motor_info])
            _ctrl_lo = jnp.array([a.ctrl_range[0] for a in acts])
            _ctrl_hi = jnp.array([a.ctrl_range[1] for a in acts])
            _bias = jnp.array([a.action_bias for a in acts])
            _dz_pos = jnp.array([a.deadzone[0] for a in acts])
            _dz_neg = jnp.array([a.deadzone[1] for a in acts])
            _gear_pos = jnp.array([a.gear[0] for a in acts])
            _gear_neg = jnp.array([a.gear[1] for a in acts])
            _gear_avg = jnp.array([a.gear_avg for a in acts])
            _fl_pos = jnp.array([a.frictionloss[0] for a in acts])
            _fl_neg = jnp.array([a.frictionloss[1] for a in acts])
            _strb_boost = jnp.array([a.stribeck_friction_boost for a in acts])
            _strb_vel = jnp.array([a.stribeck_vel for a in acts])
            _damp_pos = jnp.array([a.damping[0] for a in acts])
            _damp_neg = jnp.array([a.damping[1] for a in acts])
            _visc_quad = jnp.array([a.viscous_quadratic for a in acts])
            _back_emf = jnp.array([a.back_emf_gain for a in acts])

        # ── Batched step (N substeps per call) ──────────────────────
        # fr/dp/tq_scale are per-env (num_envs,) domain-randomization
        # multipliers (1.0 = off).  Passed as args (not closure constants) so
        # resampling them every episode does NOT trigger JIT recompilation.
        @jax.jit
        def step_fn(data, fr_scale, dp_scale, tq_scale):
            fr = fr_scale[:, None]   # broadcast over motor actuators
            dp = dp_scale[:, None]
            tq = tq_scale[:, None]

            # Software limit switch: clamp ctrl once before substeps.
            pos = data.qpos[:, act_jnt_ids]
            ctrl = data.ctrl
            at_hi = act_limited & (pos >= act_hi) & (act_gs * ctrl > 0)
            at_lo = act_limited & (pos <= act_lo) & (act_gs * ctrl < 0)
            ctrl = jnp.where(at_hi | at_lo, 0.0, ctrl)

            if _has_motor:
                # Clip → bias → deadzone → asymmetric gear compensation
                # (same order as ActuatorConfig.transform_ctrl).
                mc = ctrl[:, _ctrl_ids]
                mc = jnp.clip(mc, _ctrl_lo, _ctrl_hi)
                mc = mc + _bias
                mc = jnp.where((mc >= 0) & (mc < _dz_pos), 0.0, mc)
                mc = jnp.where((mc < 0) & (mc > -_dz_neg), 0.0, mc)
                gear_dir = jnp.where(mc >= 0, _gear_pos, _gear_neg)
                mc = mc * gear_dir / _gear_avg
                mc = mc * tq                       # torque_scale (DR)
                ctrl = ctrl.at[:, _ctrl_ids].set(mc)

            data = data.replace(ctrl=ctrl)

            def body(_, d):
                if _has_motor:
                    vel = d.qvel[:, _qvel_ids]
                    mc = d.ctrl[:, _ctrl_ids]

                    # Coulomb + Stribeck friction (direction-dependent) × DR
                    fl = jnp.where(vel > 0, _fl_pos, _fl_neg)
                    fl = fl + _strb_boost * jnp.exp(
                        -((jnp.abs(vel) / _strb_vel) ** 2)
                    )
                    fl = fl * fr
                    torque = -jnp.where(
                        jnp.abs(vel) > 1e-6, jnp.sign(vel) * fl, 0.0,
                    )
                    # Viscous damping (direction-dependent) × DR scale
                    damp = jnp.where(vel > 0, _damp_pos, _damp_neg) * dp
                    torque = torque - damp * vel
                    # Quadratic velocity drag
                    torque = torque - _visc_quad * vel * jnp.abs(vel)
                    # Back-EMF torque reduction
                    bemf = _back_emf * vel * jnp.sign(mc)
                    torque = torque - jnp.where(
                        jnp.abs(mc) > 1e-6, bemf, 0.0,
                    )
                    torque = jnp.clip(torque, -10.0, 10.0)
                    d = d.replace(
                        qfrc_applied=d.qfrc_applied.at[:, _qvel_ids].set(torque),
                    )

                return jax.vmap(mjx.step, in_axes=(None, 0))(model, d)

            return jax.lax.fori_loop(0, substeps, body, data)

        self._jit_step = step_fn

        # ── Selective reset ─────────────────────────────────────────
        init_qpos_lo = self._init_qpos_lo
        init_qpos_hi = self._init_qpos_hi
        init_qvel_lo = self._init_qvel_lo
        init_qvel_hi = self._init_qvel_hi

        @jax.jit
        def reset_fn(data, mask, rng):
            rng, k1, k2 = jax.random.split(rng, 3)
            ne = data.qpos.shape[0]

            pq = jax.random.uniform(
                k1, (ne, default.qpos.shape[0]),
                minval=init_qpos_lo, maxval=init_qpos_hi,
            )
            pv = jax.random.uniform(
                k2, (ne, default.qvel.shape[0]),
                minval=init_qvel_lo, maxval=init_qvel_hi,
            )

            m = mask[:, None]  # (num_envs, 1) broadcast helper

            new_qpos = jnp.where(m, default.qpos + pq, data.qpos)
            new_qvel = jnp.where(m, default.qvel + pv, data.qvel)
            new_ctrl = jnp.where(m, 0.0, data.ctrl)

            new_data = data.replace(qpos=new_qpos, qvel=new_qvel, ctrl=new_ctrl)
            new_data = jax.vmap(mjx.forward, in_axes=(None, 0))(model, new_data)

            return new_data, rng

        self._jit_reset = reset_fn

    # ── Step / Reset ─────────────────────────────────────────────────

    def _sim_step(self, actions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        # PyTorch → JAX  (zero-copy on GPU via DLPack)
        actions_jax = jnp.from_dlpack(actions.detach().contiguous())

        # Set ctrl & run N substeps for all environments (with per-env DR scales)
        self._batch_data = self._batch_data.replace(ctrl=actions_jax)
        self._batch_data = self._jit_step(
            self._batch_data, self._mjx_fr, self._mjx_dp, self._mjx_tq,
        )

        # JAX → PyTorch  (zero-copy on GPU via DLPack, cast to float32)
        qpos = torch.from_dlpack(self._batch_data.qpos.astype(jnp.float32))
        qvel = torch.from_dlpack(self._batch_data.qvel.astype(jnp.float32))
        return qpos, qvel

    def _sim_reset(self, env_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        # Build boolean mask (fixed shape → no JIT recompilation)
        mask = torch.zeros(
            self.config.num_envs, dtype=torch.bool, device=self.device,
        )
        mask[env_ids] = True
        mask_jax = jnp.from_dlpack(mask)

        self._batch_data, self._rng = self._jit_reset(
            self._batch_data, mask_jax, self._rng,
        )

        # Sync per-env DR scales (torch → JAX) for the step fn.  BaseRunner
        # resamples self._dr_scales just before this call, so re-deriving the
        # full arrays here keeps the JAX copies current for every env.
        self._mjx_fr = jnp.from_dlpack(self._dr_scales["friction_scale"].contiguous())
        self._mjx_dp = jnp.from_dlpack(self._dr_scales["damping_scale"].contiguous())
        self._mjx_tq = jnp.from_dlpack(self._dr_scales["torque_scale"].contiguous())

        # Return the FULL batch (BaseRunner indexes the reset envs in torch)
        # — avoids a GPU→CPU sync + JAX gather on every step with a done env.
        qpos = torch.from_dlpack(self._batch_data.qpos.astype(jnp.float32))
        qvel = torch.from_dlpack(self._batch_data.qvel.astype(jnp.float32))
        return qpos, qvel

    # ── Rendering ────────────────────────────────────────────────────

    def _render_frame(self, env_idx: int = 0) -> np.ndarray:
        """Offscreen render — copies one env's state from GPU to CPU."""
        self._render_data.qpos[:] = np.asarray(self._batch_data.qpos[env_idx])
        self._render_data.qvel[:] = np.asarray(self._batch_data.qvel[env_idx])
        self._render_data.ctrl[:] = np.asarray(self._batch_data.ctrl[env_idx])
        mujoco.mj_forward(self._mj_model, self._render_data)

        if not hasattr(self, "_offscreen_renderer"):
            self._offscreen_renderer = mujoco.Renderer(
                self._mj_model, width=640, height=480,
            )
        self._offscreen_renderer.update_scene(self._render_data)
        return self._offscreen_renderer.render().copy()
