"""Deterministic simulation replay — roll out recorded actions in MuJoCo.

Given a parameter vector and a recorded action sequence, builds a MuJoCo
model with overridden physics parameters, replays the actions, and returns
the simulated trajectory for comparison with the real recording.

This module is the inner loop of the CMA-ES optimizer: it is called once
per candidate parameter vector per generation.

Motor parameters are **locked** from the unified sysid result.
The optimizer only fits
pendulum/arm inertial parameters, pendulum joint dynamics, and
``ctrl_limit``.  The asymmetric motor model (bias, deadzone, gear
compensation, Coulomb + Stribeck friction, viscous damping) is applied
via ``ActuatorConfig.transform_ctrl()`` and ``compute_motor_force()`` —
the same code the training runners use, so sim == sysid by construction.
"""

from __future__ import annotations

import dataclasses
import os
import tempfile
import xml.etree.ElementTree as ET
from pathlib import Path

import mujoco
import numpy as np
import yaml

from src.core.robot import ActuatorConfig, JointConfig, RobotConfig
from src.runners.mujoco import ActuatorLimits, load_mujoco_model
from src.sysid._urdf import patch_link_inertials


# ── Locked motor parameters (from the unified sysid) ────────────────
# These are FIXED and not optimised.  They come from the unified
# 28-param sysid run (assets/rotary_cartpole/sysid_result.json,
# cost 0.925) — Stribeck friction + action bias + ~96 ms motor lag.

LOCKED_MOTOR_PARAMS: dict[str, float] = {
    "actuator_gear_pos": 0.846499,
    "actuator_gear_neg": 1.183733,
    "actuator_filter_tau": 0.096263,
    "motor_damping_pos": 0.013165,
    "motor_damping_neg": 0.015452,
    "motor_armature": 0.001676,
    "motor_frictionloss_pos": 0.014244,
    "motor_frictionloss_neg": 0.001005,
    "stribeck_friction_boost": 0.068594,
    "stribeck_vel": 5.279594,
    "motor_deadzone_pos": 0.181097,
    "motor_deadzone_neg": 0.202072,
    "action_bias": 0.056566,
}


# ── Tunable parameter specification ──────────────────────────────────


@dataclasses.dataclass
class ParamSpec:
    """Specification for a single tunable parameter."""

    name: str
    default: float
    lower: float
    upper: float
    log_scale: bool = False  # optimise in log-space (masses, inertias)


# Pendulum sysid parameters — motor params are LOCKED (not here).
# Order matters: the optimizer maps a flat vector to these specs.
# Defaults are from the URDF exported by Fusion 360.
ROTARY_CARTPOLE_PARAMS: list[ParamSpec] = [
    # ── Arm link (URDF) ──────────────────────────────────────────
    ParamSpec("arm_mass", 0.02110, 0.005, 0.08, log_scale=True),
    ParamSpec("arm_com_x", -0.00710, -0.03, 0.03),
    ParamSpec("arm_com_y", 0.00085, -0.02, 0.02),
    ParamSpec("arm_com_z", 0.00795, -0.02, 0.02),
    # ── Pendulum link (URDF) ─────────────────────────────────────
    ParamSpec("pendulum_mass", 0.03937, 0.010, 0.10, log_scale=True),
    ParamSpec("pendulum_com_x", 0.06025, 0.01, 0.15),
    ParamSpec("pendulum_com_y", -0.07602, -0.20, 0.0),
    ParamSpec("pendulum_com_z", -0.00346, -0.05, 0.05),
    ParamSpec("pendulum_ixx", 6.20e-05, 1e-07, 1e-03, log_scale=True),
    ParamSpec("pendulum_iyy", 3.70e-05, 1e-07, 1e-03, log_scale=True),
    ParamSpec("pendulum_izz", 7.83e-05, 1e-07, 1e-03, log_scale=True),
    ParamSpec("pendulum_ixy", -6.93e-06, -1e-03, 1e-03),
    # ── Pendulum joint dynamics ──────────────────────────────────
    ParamSpec("pendulum_damping", 0.0001, 1e-6, 0.05, log_scale=True),
    ParamSpec("pendulum_frictionloss", 0.0001, 1e-6, 0.05, log_scale=True),
    # ── Hardware realism (control pipeline) ────────────────────
    ParamSpec("ctrl_limit", 0.588, 0.45, 0.70),  # MAX_MOTOR_SPEED / 255
]


# Extended set: full params + motor armature (compensates for the
# motor-only sysid having captured arm/pendulum loading in armature).
EXTENDED_PARAMS: list[ParamSpec] = ROTARY_CARTPOLE_PARAMS + [
    ParamSpec("motor_armature", 0.00277, 0.0005, 0.02, log_scale=True),
]


# Reduced set: only the 6 most impactful pendulum parameters.
# Good for a fast first pass — converges in ~50 generations.
REDUCED_PARAMS: list[ParamSpec] = [
    ParamSpec("pendulum_mass", 0.03937, 0.010, 0.10, log_scale=True),
    ParamSpec("pendulum_com_x", 0.06025, 0.01, 0.15),
    ParamSpec("pendulum_com_y", -0.07602, -0.20, 0.0),
    ParamSpec("pendulum_ixx", 6.20e-05, 1e-07, 1e-03, log_scale=True),
    ParamSpec("pendulum_damping", 0.0001, 1e-6, 0.05, log_scale=True),
    ParamSpec("pendulum_frictionloss", 0.0001, 1e-6, 0.05, log_scale=True),
]


PARAM_SETS: dict[str, list[ParamSpec]] = {
    "full": ROTARY_CARTPOLE_PARAMS,
    "extended": EXTENDED_PARAMS,
    "reduced": REDUCED_PARAMS,
}


def params_to_dict(
    values: np.ndarray, specs: list[ParamSpec] | None = None
) -> dict[str, float]:
    """Convert a flat parameter vector to a named dict."""
    if specs is None:
        specs = ROTARY_CARTPOLE_PARAMS
    return {s.name: float(values[i]) for i, s in enumerate(specs)}


def defaults_vector(specs: list[ParamSpec] | None = None) -> np.ndarray:
    """Return the default parameter vector (in natural space)."""
    if specs is None:
        specs = ROTARY_CARTPOLE_PARAMS
    return np.array([s.default for s in specs], dtype=np.float64)


def bounds_arrays(
    specs: list[ParamSpec] | None = None,
) -> tuple[np.ndarray, np.ndarray]:
    """Return (lower, upper) bound arrays."""
    if specs is None:
        specs = ROTARY_CARTPOLE_PARAMS
    lo = np.array([s.lower for s in specs], dtype=np.float64)
    hi = np.array([s.upper for s in specs], dtype=np.float64)
    return lo, hi


# ── MuJoCo model building with parameter overrides ──────────────────


def _build_model(
    robot_path: Path,
    params: dict[str, float],
    motor_params: dict[str, float] | None = None,
) -> tuple[mujoco.MjModel, ActuatorConfig]:
    """Build a MuJoCo model with sysid overrides.

    Returns (model, actuator) — use ``actuator.transform_ctrl()`` and
    ``actuator.compute_motor_force()`` in the rollout loop.
    """
    if motor_params is None:
        motor_params = LOCKED_MOTOR_PARAMS

    robot_path = Path(robot_path).resolve()

    # ── Patch URDF inertial parameters to a temp file ────────────
    robot_yaml = yaml.safe_load((robot_path / "robot.yaml").read_text())
    urdf_path = robot_path / robot_yaml["urdf"]

    tree = ET.parse(urdf_path)
    patch_link_inertials(tree.getroot(), params)

    fd, tmp_urdf = tempfile.mkstemp(
        suffix=".urdf", prefix="_sysid_", dir=str(robot_path),
    )
    os.close(fd)
    tmp_urdf_path = Path(tmp_urdf)
    tree.write(str(tmp_urdf_path), xml_declaration=True, encoding="unicode")

    # ── Build RobotConfig with full motor sysid values ───────────
    gear_pos = motor_params.get("actuator_gear_pos", 0.424182)
    gear_neg = motor_params.get("actuator_gear_neg", 0.425031)
    motor_armature = params.get(
        "motor_armature",
        motor_params.get("motor_armature", 0.00277342),
    )
    pend_damping = params.get("pendulum_damping", 0.0001)
    pend_frictionloss = params.get("pendulum_frictionloss", 0.0001)

    act_cfg = robot_yaml["actuators"][0]
    ctrl_lo, ctrl_hi = act_cfg.get("ctrl_range", [-1.0, 1.0])

    # The fitted ctrl_limit overrides the YAML ctrl_range so the rollout
    # saturates at exactly the identified PWM bound.
    if "ctrl_limit" in params:
        ctrl_lo, ctrl_hi = -params["ctrl_limit"], params["ctrl_limit"]

    actuator = ActuatorConfig(
        joint=act_cfg["joint"],
        type="motor",
        gear=(gear_pos, gear_neg),
        ctrl_range=(ctrl_lo, ctrl_hi),
        deadzone=(
            motor_params.get("motor_deadzone_pos", 0.181),
            motor_params.get("motor_deadzone_neg", 0.202),
        ),
        damping=(
            motor_params.get("motor_damping_pos", 0.013),
            motor_params.get("motor_damping_neg", 0.015),
        ),
        frictionloss=(
            motor_params.get("motor_frictionloss_pos", 0.014),
            motor_params.get("motor_frictionloss_neg", 0.001),
        ),
        filter_tau=motor_params.get("actuator_filter_tau", 0.096),
        viscous_quadratic=motor_params.get("viscous_quadratic", 0.0),
        back_emf_gain=motor_params.get("back_emf_gain", 0.0),
        stribeck_friction_boost=motor_params.get("stribeck_friction_boost", 0.0),
        stribeck_vel=motor_params.get("stribeck_vel", 2.0),
        action_bias=motor_params.get("action_bias", 0.0),
    )

    robot = RobotConfig(
        urdf_path=tmp_urdf_path,
        actuators=[actuator],
        joints={
            "motor_joint": JointConfig(
                damping=0.0,
                armature=motor_armature,
                frictionloss=0.0,
            ),
            "pendulum_joint": JointConfig(
                damping=pend_damping,
                frictionloss=pend_frictionloss,
            ),
        },
    )

    try:
        model = load_mujoco_model(robot)
    finally:
        tmp_urdf_path.unlink(missing_ok=True)

    return model, actuator





# ── Simulation rollout ───────────────────────────────────────────────


def rollout(
    robot_path: str | Path,
    params: dict[str, float],
    actions: np.ndarray,
    sim_dt: float = 0.002,
    substeps: int = 10,
    motor_params: dict[str, float] | None = None,
) -> dict[str, np.ndarray]:
    """Replay recorded actions in MuJoCo with overridden parameters.

    Parameters
    ----------
    robot_path : asset directory
    params : named parameter overrides (pendulum/arm only)
    actions : (N,) normalised actions [-1, 1] from the recording
    sim_dt : MuJoCo physics timestep
    substeps : physics substeps per control step
    motor_params : locked motor params (default: LOCKED_MOTOR_PARAMS)

    Returns
    -------
    dict with keys: motor_angle, motor_vel, pendulum_angle, pendulum_vel
    Each is an (N,) numpy array of simulated values.
    """
    if motor_params is None:
        motor_params = LOCKED_MOTOR_PARAMS

    robot_path = Path(robot_path).resolve()
    model, actuator = _build_model(robot_path, params, motor_params)
    model.opt.timestep = sim_dt
    data = mujoco.MjData(model)
    mujoco.mj_resetData(model, data)

    n = len(actions)

    sim_motor_angle = np.zeros(n, dtype=np.float64)
    sim_motor_vel = np.zeros(n, dtype=np.float64)
    sim_pend_angle = np.zeros(n, dtype=np.float64)
    sim_pend_vel = np.zeros(n, dtype=np.float64)

    limits = ActuatorLimits(model)

    for i in range(n):
        # transform_ctrl clips to the (fitted) ctrl_range internally.
        ctrl = actuator.transform_ctrl(float(actions[i]))
        data.ctrl[0] = ctrl

        for _ in range(substeps):
            limits.enforce(model, data)
            data.qfrc_applied[0] = actuator.compute_motor_force(data.qvel[0], ctrl)
            mujoco.mj_step(model, data)

        sim_motor_angle[i] = data.qpos[0]
        sim_pend_angle[i] = data.qpos[1]
        sim_motor_vel[i] = data.qvel[0]
        sim_pend_vel[i] = data.qvel[1]

    return {
        "motor_angle": sim_motor_angle,
        "motor_vel": sim_motor_vel,
        "pendulum_angle": sim_pend_angle,
        "pendulum_vel": sim_pend_vel,
    }


def windowed_rollout(
    robot_path: str | Path,
    params: dict[str, float],
    recording: dict[str, np.ndarray],
    window_duration: float = 0.5,
    sim_dt: float = 0.002,
    substeps: int = 10,
    motor_params: dict[str, float] | None = None,
) -> dict[str, np.ndarray | float]:
    """Multiple-shooting rollout — split recording into short windows.

    For each window:
      1. Initialize MuJoCo state from the real qpos/qvel at the window start.
      2. Replay the recorded actions within the window.
      3. Record the simulated output.

    Motor dynamics (asymmetric friction, damping, back-EMF, etc.) are
    applied via qfrc_applied using the locked motor sysid parameters.

    Parameters
    ----------
    robot_path : asset directory
    params : named parameter overrides (pendulum/arm only)
    recording : dict with keys time, action, motor_angle, motor_vel,
                pendulum_angle, pendulum_vel (all 1D arrays of length N)
    window_duration : length of each shooting window in seconds
    sim_dt : MuJoCo physics timestep
    substeps : physics substeps per control step
    motor_params : locked motor params (default: LOCKED_MOTOR_PARAMS)

    Returns
    -------
    dict with:
        motor_angle, motor_vel, pendulum_angle, pendulum_vel — (N,) arrays
            (stitched from per-window simulations)
        n_windows — number of windows used
    """
    if motor_params is None:
        motor_params = LOCKED_MOTOR_PARAMS

    robot_path = Path(robot_path).resolve()
    model, actuator = _build_model(robot_path, params, motor_params)
    model.opt.timestep = sim_dt
    data = mujoco.MjData(model)

    times = recording["time"]
    actions = recording["action"]
    real_motor = recording["motor_angle"]
    real_motor_vel = recording["motor_vel"]
    real_pend = recording["pendulum_angle"]
    real_pend_vel = recording["pendulum_vel"]
    n = len(actions)

    sim_motor_angle = np.zeros(n, dtype=np.float64)
    sim_motor_vel = np.zeros(n, dtype=np.float64)
    sim_pend_angle = np.zeros(n, dtype=np.float64)
    sim_pend_vel = np.zeros(n, dtype=np.float64)

    limits = ActuatorLimits(model)

    t0 = times[0]
    t_end = times[-1]
    window_starts: list[int] = []
    current_t = t0
    while current_t < t_end:
        idx = int(np.searchsorted(times, current_t))
        idx = min(idx, n - 1)
        window_starts.append(idx)
        current_t += window_duration

    n_windows = len(window_starts)

    for w, w_start in enumerate(window_starts):
        w_end = window_starts[w + 1] if w + 1 < n_windows else n

        mujoco.mj_resetData(model, data)
        data.qpos[0] = real_motor[w_start]
        data.qpos[1] = real_pend[w_start]
        data.qvel[0] = real_motor_vel[w_start]
        data.qvel[1] = real_pend_vel[w_start]
        data.ctrl[:] = 0.0
        mujoco.mj_forward(model, data)

        for i in range(w_start, w_end):
            # transform_ctrl clips to the (fitted) ctrl_range internally.
            ctrl = actuator.transform_ctrl(float(actions[i]))
            data.ctrl[0] = ctrl

            for _ in range(substeps):
                limits.enforce(model, data)
                data.qfrc_applied[0] = actuator.compute_motor_force(data.qvel[0], ctrl)
                mujoco.mj_step(model, data)

            sim_motor_angle[i] = data.qpos[0]
            sim_pend_angle[i] = data.qpos[1]
            sim_motor_vel[i] = data.qvel[0]
            sim_pend_vel[i] = data.qvel[1]

    return {
        "motor_angle": sim_motor_angle,
        "motor_vel": sim_motor_vel,
        "pendulum_angle": sim_pend_angle,
        "pendulum_vel": sim_pend_vel,
        "n_windows": n_windows,
    }
