"""Standalone S10 MJCF/FK checks that do not require mjlab or torch.

Run with the system Python:
  python3 tests/test_s10_model.py

The test validates:
1. The S10 model in ``s10_supermimic/robot/model/s10.xml`` loads.
2. Joint names match the training/deployment contract.
3. The default pose puts the wheel bottom at ~0.434 m above the ground.
4. The analytic wheel-bottom formula in ``core/rolling_rewards.py`` matches
   MuJoCo forward kinematics for random joint angles and arbitrary base tilt.
"""

from __future__ import annotations

import math
from pathlib import Path

import mujoco
import numpy as np


ROOT = Path(__file__).resolve().parents[1]
MODEL = ROOT / "s10_supermimic" / "robot" / "model" / "s10.xml"
LEGS = ("fl", "fr", "hl", "hr")
LEG_JOINT_NAMES = tuple(
    f"{leg}_{joint}_joint" for leg in LEGS for joint in ("hipx", "hipy", "knee")
)
WHEEL_JOINT_NAMES = tuple(f"{leg}_wheel_joint" for leg in LEGS)
DEFAULT_Q = {
    "fl": (0.1, -0.35, 0.65),
    "fr": (-0.1, -0.35, 0.65),
    "hl": (0.1, 0.35, -0.65),
    "hr": (-0.1, 0.35, -0.65),
}
# Geometry from S10.xml.
X0 = {"fl": 0.2277, "fr": 0.2277, "hl": -0.2277, "hr": -0.2277}
Y0 = {"fl": 0.0574, "fr": -0.0574, "hl": 0.0574, "hr": -0.0574}
SIDE = {"fl": 1.0, "fr": -1.0, "hl": 1.0, "hr": -1.0}
HIPY_Y = 0.0779
WHEEL_Y = 0.045891
LINK = 0.18
WHEEL_RADIUS = 0.081
WHEEL_HALF_WIDTH = 0.0185


def analytic_wheel_bottom(q, leg, base_z=0.5, up=(0.0, 0.0, 1.0)) -> float:
    """NumPy translation of the torch implementation in rolling_rewards.py."""
    hipx, hipy, knee = (float(v) for v in q)
    side = SIDE[leg]
    x = LINK * math.sin(knee)
    z = -LINK - LINK * math.cos(knee)
    x3 = math.cos(hipy) * x - math.sin(hipy) * z
    z3 = math.sin(hipy) * x + math.cos(hipy) * z
    y3 = side * (HIPY_Y + WHEEL_Y)
    y4 = math.cos(hipx) * y3 + math.sin(hipx) * z3
    z4 = -math.sin(hipx) * y3 + math.cos(hipx) * z3
    point = np.array([X0[leg] + x3, Y0[leg] + y4, z4])
    up = np.asarray(up, dtype=float)
    up /= np.linalg.norm(up)
    axle_z = -up[1] * math.cos(hipx) + up[2] * math.sin(hipx)
    radius_z = WHEEL_RADIUS * math.sqrt(max(1.0 - axle_z * axle_z, 0.0))
    radius_z += WHEEL_HALF_WIDTH * abs(axle_z)
    return base_z + float(point @ up) - radius_z


def wheel_geom_id(model: mujoco.MjModel, leg: str) -> int:
    body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, f"{leg}_wheel")
    for geom_id in range(model.ngeom):
        if model.geom_bodyid[geom_id] != body_id:
            continue
        if model.geom_type[geom_id] == mujoco.mjtGeom.mjGEOM_CYLINDER:
            return geom_id
    raise AssertionError(f"No cylinder collision geom for {leg}_wheel")


def measured_wheel_bottom(model, data, geom_id: int) -> float:
    axis = data.geom_xmat[geom_id].reshape(3, 3)[:, 2]
    axle_z = abs(float(axis[2]))
    radius_z = WHEEL_RADIUS * math.sqrt(max(1.0 - axle_z * axle_z, 0.0))
    radius_z += WHEEL_HALF_WIDTH * axle_z
    return float(data.geom_xpos[geom_id][2] - radius_z)


def set_pose(model, data, q_map, base_z=0.5, quat=(1.0, 0.0, 0.0, 0.0)):
    data.qpos[:] = 0.0
    data.qpos[2] = base_z
    data.qpos[3:7] = quat
    for joint_name, value in q_map.items():
        joint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, joint_name)
        data.qpos[model.jnt_qposadr[joint_id]] = value
    mujoco.mj_forward(model, data)


def main() -> None:
    assert MODEL.is_file(), MODEL
    model = mujoco.MjModel.from_xml_path(str(MODEL))
    data = mujoco.MjData(model)

    joint_names = [model.joint(i).name for i in range(model.njnt)]
    expected = ("floating_base",) + LEG_JOINT_NAMES + WHEEL_JOINT_NAMES
    assert set(joint_names) == set(expected), joint_names
    assert model.nq == 23 and model.nv == 22 and model.njnt == 17
    assert model.nu == 0, "Robot XML must not contain actuators; mjlab adds them."

    default_q = {
        f"{leg}_{joint}_joint": value
        for leg in LEGS
        for joint, value in zip(("hipx", "hipy", "knee"), DEFAULT_Q[leg])
    }
    set_pose(model, data, default_q)
    for leg in LEGS:
        geom_id = wheel_geom_id(model, leg)
        bottom = measured_wheel_bottom(model, data, geom_id)
        pred = analytic_wheel_bottom(DEFAULT_Q[leg], leg)
        assert abs(bottom - pred) < 1e-9, (leg, bottom, pred)
        assert abs((0.5 - bottom) - 0.434145) < 1e-4, (leg, bottom)

    rng = np.random.default_rng(0)
    max_error = 0.0
    for _ in range(200):
        quat = rng.normal(size=4)
        quat /= np.linalg.norm(quat)
        q_map = {}
        q_by_leg = {}
        for leg in LEGS:
            q = np.array([
                rng.uniform(-0.6109, 0.6109),
                rng.uniform(-2.5307, 2.5307),
                rng.uniform(-2.7227, 2.7227),
            ])
            q_by_leg[leg] = q
            for joint, value in zip(("hipx", "hipy", "knee"), q):
                q_map[f"{leg}_{joint}_joint"] = value
        set_pose(model, data, q_map, base_z=0.5, quat=quat)

        rotation = data.xmat[
            mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "base_link")
        ].reshape(3, 3)
        up = rotation.T @ np.array([0.0, 0.0, 1.0])
        for leg in LEGS:
            geom_id = wheel_geom_id(model, leg)
            measured = measured_wheel_bottom(model, data, geom_id)
            pred = analytic_wheel_bottom(q_by_leg[leg], leg, base_z=0.5, up=up)
            max_error = max(max_error, abs(measured - pred))
    assert max_error < 1e-9, max_error
    print(f"S10 model/FK checks passed (max error {max_error:.2e} m)")


if __name__ == "__main__":
    main()