"""The twin and the oracle against the numbers they were built from (configs/stringman.yaml, docs/MACHINE.md)."""
import sys
from pathlib import Path

import mujoco
import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from cablepilot.model.oracle import OracleWorldModel  # noqa: E402
from cablepilot.sim.kin import line_speeds, static_tensions, worst_hang_tensions  # noqa: E402
from cablepilot.sim.plant import Plant  # noqa: E402
from cablepilot.sim.scene import Indices, build_model, load_config, mounts, settle_spools, total_mass  # noqa: E402
from cablepilot.sim.state import set_from_state, to_state  # noqa: E402

CFG = load_config()


def _hang(world=False, timestep=None, p=(0.0, 0.0, 1.2), seconds=3.0):
    m = build_model(CFG, world=world, timestep=timestep, start=p)
    d = mujoco.MjData(m)
    idx = Indices(m)
    settle_spools(m, d, idx, CFG)
    for _ in range(int(seconds / m.opt.timestep)):
        mujoco.mj_step(m, d)
    return m, d, idx


def test_hangs_still_with_static_tensions():
    m, d, idx = _hang()
    x = to_state(m, d, idx)
    assert np.allclose(x[:3], [0, 0, 1.2], atol=0.003)
    T = -d.actuator_force[idx.spool]
    assert np.allclose(T, static_tensions(CFG, x[:3], total_mass(m)), atol=0.2)   # symmetric pose: 4 equal shares
    assert T.max() < CFG["lines"]["max_safe_tension_N"]


def test_swing_period_close_to_firmware():
    """fw definitions.pole_length_carbon270 = 0.3089 m -> 1.115 s; the twin must be within 5 %."""
    m, d, idx = _hang()
    d.qvel[idx.gv + 3] = 0.3
    ys = []
    for _ in range(int(6 / m.opt.timestep)):
        mujoco.mj_step(m, d)
        ys.append(to_state(m, d, idx)[7])
    ys = np.array(ys) - np.mean(ys)
    up = np.where(np.diff(np.sign(ys)) > 0)[0]
    period = np.diff(up).mean() * m.opt.timestep
    assert abs(period / (2 * np.pi * np.sqrt(0.3089 / 9.81)) - 1) < 0.05


def test_timestep_convergence():
    """Plant step 0.004 s vs the firmware twin's 0.002 s: same hang within 1 mm."""
    m1, d1, i1 = _hang(timestep=0.004, p=(0.8, -0.5, 1.0))
    m2, d2, i2 = _hang(timestep=0.002, p=(0.8, -0.5, 1.0))
    assert np.linalg.norm(to_state(m1, d1, i1)[:3] - to_state(m2, d2, i2)[:3]) < 0.001


def test_state_roundtrip():
    m = build_model(CFG, world=False)
    d = mujoco.MjData(m)
    idx = Indices(m)
    x = np.array([0.3, -0.2, 1.1, 0.05, -0.02, 0.01, 0.04, -0.03, 0.1, 0.05, 0.7, 0.2, 0.4, 0.1])
    set_from_state(m, d, idx, x, CFG)
    mujoco.mj_forward(m, d)
    assert np.allclose(to_state(m, d, idx), x, atol=1e-6)


def test_kinematics_helpers():
    M = mounts(CFG)
    p = np.array([0.4, 0.3, 1.0])
    T = static_tensions(CFG, p, 1.0)
    u = (M - p) / np.linalg.norm(M - p, axis=1, keepdims=True)
    assert np.allclose(u.T @ T, [0, 0, 9.81], atol=1e-6)
    v = np.array([0.0, 0.0, 0.1])                    # going up: every line shortens
    assert (line_speeds(M, p, v) < 0).all()
    Tw = worst_hang_tensions(torch.tensor(M, dtype=torch.float32), torch.tensor(p, dtype=torch.float32), 1.0).numpy()
    assert (Tw >= T - 1e-4).all()


def test_oracle_matches_plant_for_a_smooth_move():
    dt, Tn = 0.2, 5
    pl = Plant(CFG)
    pl.reset([0.2, 0.1, 1.1])
    x = pl.state()
    acts = np.zeros((Tn, 5))
    acts[:, 0], acts[:, 2] = 0.08, -0.05
    pred = OracleWorldModel(CFG, dt, timestep=0.01).rollout(x[None], acts[:1], acts[None])[0].numpy()
    for k in range(Tn):
        pl.command(pl.state()[:3], acts[k])
        pl.step(dt)
    assert np.linalg.norm(pred[-1, :3] - pl.state()[:3]) < 0.01      # within 1 cm after 1 s


# ----------------------------------------------------------------------------------------------- the SO-101 hand (M3)
SO = load_config(hand="so101")


def test_so101_mass_from_its_parts():
    """Every part's mass counts (MuJoCo drops geom masses in a body with an <inertial>: they ride on welded bodies)."""
    m = build_model(SO, world=False)
    h, b = SO["so101"], SO["body"]
    parts = (b["gantry_mass_kg"] + b["pole_mass_kg"] + h["roll_servo_mass_kg"] + h["gripper_mass_kg"] + h["jaw_mass_kg"]
             + h["camera_mass_kg"])
    assert abs(total_mass(m) - parts) < 1e-6


def test_so101_hangs_still_and_roundtrips():
    m = build_model(SO, world=False)
    d = mujoco.MjData(m)
    idx = Indices(m)
    settle_spools(m, d, idx, SO)
    for _ in range(int(3 / m.opt.timestep)):
        mujoco.mj_step(m, d)
    assert np.allclose(to_state(m, d, idx)[:3], [0, 0, 1.2], atol=0.003)
    x = np.array([0.3, -0.2, 1.1, 0.05, -0.02, 0.01, 0.04, -0.03, 0.1, 0.05, 0.7, 0.2, 0.4, 0.1])
    set_from_state(m, d, idx, x, SO)
    mujoco.mj_forward(m, d)
    assert np.allclose(to_state(m, d, idx), x, atol=1e-6)


def test_so101_jaw_convention():
    """closure 1 = the pads meet (SO-101 -10 deg), closure 0 = jaw_open_deg; the plant reports the SO-101's scale."""
    from cablepilot.sim.plant import Plant
    pl = Plant(SO)
    m, d, idx = pl.model, pl.data, pl.idx
    assert abs(np.degrees(idx.travel) - (SO["so101"]["jaw_open_deg"] - SO["so101"]["jaw_closed_deg"])) < 1e-3   # written with 6 decimals
    fl, fr = m.geom("finger_l_pad").id, m.geom("finger_r_pad").id
    for close, gap_ok in ((1.0, lambda g: abs(g) < 0.001), (0.0, lambda g: g > 0.05)):
        pl.reset([0.0, 0.0, 1.0], close=close)
        assert gap_ok(mujoco.mj_geomDistance(m, d, fl, fr, 0.2, None))
        want = SO["so101"]["jaw_closed_deg"] if close == 1.0 else SO["so101"]["jaw_open_deg"]
        assert abs(pl.jaw_angle_deg() - want) < 0.5