"""Robot config loading + motor model unit tests."""

import math
from pathlib import Path

import pytest
import torch

from src.core.robot import ActuatorConfig, load_robot_config

ROBOT_DIR = Path(__file__).resolve().parent.parent / "assets" / "rotary_cartpole"


# ── Loading ──────────────────────────────────────────────────────────


def test_canonical_robot_yaml_loads_full_motor_model():
    robot = load_robot_config(ROBOT_DIR)
    act = robot.actuators[0]

    assert act.has_motor_model
    # Tuned (unified sysid) values must survive the round-trip.
    assert act.filter_tau == pytest.approx(0.096263)
    assert act.stribeck_friction_boost == pytest.approx(0.068594)
    assert act.stribeck_vel == pytest.approx(5.279594)
    assert act.action_bias == pytest.approx(0.056566)
    assert act.gear == pytest.approx((0.846499, 1.183733))


def test_unknown_actuator_keys_are_ignored_not_fatal(tmp_path):
    (tmp_path / "dummy.urdf").write_text("<robot name='x'/>")
    (tmp_path / "robot.yaml").write_text(
        "urdf: dummy.urdf\n"
        "actuators:\n"
        "  - joint: j\n"
        "    type: motor\n"
        "    gear: [1.0, 1.0]\n"
        "    some_future_field: 42\n"
    )
    robot = load_robot_config(tmp_path)  # must not raise
    assert robot.actuators[0].joint == "j"


# ── transform_ctrl: clip → bias → deadzone → gear ────────────────────


@pytest.fixture
def act() -> ActuatorConfig:
    return ActuatorConfig(
        joint="m",
        gear=(0.8, 1.2),
        ctrl_range=(-0.6, 0.6),
        deadzone=(0.15, 0.20),
        frictionloss=(0.014, 0.001),
        damping=(0.013, 0.015),
        stribeck_friction_boost=0.07,
        stribeck_vel=5.0,
        action_bias=0.05,
    )


def test_transform_ctrl_clips_to_ctrl_range(act):
    # 1.0 clips to 0.6, then +bias=0.65, gear comp 0.8/1.0 → 0.52
    out = act.transform_ctrl(1.0)
    assert out == pytest.approx((0.6 + 0.05) * 0.8 / 1.0)


def test_transform_ctrl_deadzone_zeroes_small_commands(act):
    # 0.05 + bias 0.05 = 0.10 < dz_pos 0.15 → 0
    assert act.transform_ctrl(0.05) == 0.0
    # -0.15 + bias 0.05 = -0.10 > -dz_neg -0.20 → 0
    assert act.transform_ctrl(-0.15) == 0.0


def test_transform_ctrl_gear_compensation_is_asymmetric(act):
    pos = act.transform_ctrl(0.5)   # (0.55) * 0.8
    neg = act.transform_ctrl(-0.5)  # (-0.45) * 1.2
    assert pos == pytest.approx(0.55 * 0.8)
    assert neg == pytest.approx(-0.45 * 1.2)


def test_transform_action_matches_transform_ctrl_elementwise(act):
    vals = torch.linspace(-1.2, 1.2, 49)
    batched = act.transform_action(vals.clone())
    scalar = torch.tensor([act.transform_ctrl(float(v)) for v in vals])
    assert torch.allclose(batched, scalar, atol=1e-6)


# ── compute_motor_force: Coulomb + Stribeck + damping ────────────────


def test_friction_opposes_motion(act):
    assert act.compute_motor_force(vel=2.0, ctrl=0.0) < 0
    assert act.compute_motor_force(vel=-2.0, ctrl=0.0) > 0
    assert act.compute_motor_force(vel=0.0, ctrl=0.0) == 0.0


def test_stribeck_boost_decays_with_speed(act):
    """Friction torque magnitude (minus damping) is higher near standstill."""
    no_strb = ActuatorConfig(
        joint="m", gear=act.gear, frictionloss=act.frictionloss,
        damping=(0.0, 0.0),
    )
    with_strb = ActuatorConfig(
        joint="m", gear=act.gear, frictionloss=act.frictionloss,
        damping=(0.0, 0.0),
        stribeck_friction_boost=0.07, stribeck_vel=5.0,
    )
    v_slow, v_fast = 0.1, 50.0
    extra_slow = abs(with_strb.compute_motor_force(v_slow, 0.0)) - abs(
        no_strb.compute_motor_force(v_slow, 0.0))
    extra_fast = abs(with_strb.compute_motor_force(v_fast, 0.0)) - abs(
        no_strb.compute_motor_force(v_fast, 0.0))

    assert extra_slow == pytest.approx(
        0.07 * math.exp(-((v_slow / 5.0) ** 2)), abs=1e-9)
    assert extra_fast < extra_slow
    assert extra_fast == pytest.approx(0.0, abs=1e-9)


def test_friction_scale_dr_multiplies_friction(act):
    base = act.compute_motor_force(1.0, 0.0, friction_scale=1.0, damping_scale=0.0)
    # damping_scale=0 isolates the friction term
    doubled = act.compute_motor_force(1.0, 0.0, friction_scale=2.0, damping_scale=0.0)
    assert doubled == pytest.approx(2.0 * base)