#!/usr/bin/env python3
"""Regression tests for the rigid-body closed-loop simulation."""

from dataclasses import replace
from pathlib import Path
import sys
import unittest

import numpy as np


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

from simulate_closed_loop import (  # noqa: E402
    DEFAULT_CONFIG_PATH,
    SCENARIOS,
    SimulationConfig,
    Wall,
    build_mapper,
    load_models,
    load_simulation_config,
    make_wall,
    master_endpoint_configurations,
    simulate_scenario,
)
from core.energy_audit import audit_haptic_energy  # noqa: E402


class ClosedLoopSimulationTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.models = load_models(add_simulated_tcp=True)

    def test_yaml_configuration_is_complete(self) -> None:
        config = load_simulation_config(DEFAULT_CONFIG_PATH)
        self.assertEqual(config.slave_contact_frame, "R_EE_SIM")
        self.assertEqual(config.dt, 0.002)
        self.assertEqual(config.mapping_hz, 50.0)
        self.assertLess(config.energy_min, config.energy_max)

    def test_wall_is_unilateral_dissipative_and_force_limited(self) -> None:
        wall = Wall(
            point=np.zeros(3),
            normal=np.array([1.0, 0.0, 0.0]),
            stiffness=100.0,
            damping=10.0,
            force_limit=5.0,
        )
        free_wrench, free_penetration = wall.wrench(
            np.array([-0.1, 0.0, 0.0]),
            np.array([1.0, 0.0, 0.0]),
        )
        np.testing.assert_allclose(free_wrench, 0.0)
        self.assertEqual(free_penetration, 0.0)

        contact_wrench, penetration = wall.wrench(
            np.array([0.1, 0.0, 0.0]),
            np.array([2.0, 0.0, 0.0]),
        )
        self.assertAlmostEqual(penetration, 0.1)
        np.testing.assert_allclose(contact_wrench[:3], [-5.0, 0.0, 0.0])
        self.assertLessEqual(float(contact_wrench[:3] @ np.array([2.0, 0.0, 0.0])), 0.0)

    def test_generalized_mapping_preserves_virtual_work(self) -> None:
        mapper = build_mapper(self.models)
        q_master, _ = master_endpoint_configurations()
        _, differential, debug = mapper.retarget_with_differential(q_master)
        self.assertTrue(debug["differential_valid"])

        qd_master = np.array([0.2, -0.1, 0.15, -0.3, 0.05, 0.08, -0.04])
        tau_slave = np.array([-1.0, 0.4, 0.8, -0.5, 0.2, -0.3, 0.1])
        tau_master = differential.T @ tau_slave
        self.assertAlmostEqual(
            float(tau_master @ qd_master),
            float(tau_slave @ (differential @ qd_master)),
            places=12,
        )

    def test_short_rigid_body_run_is_finite_and_respects_energy_floor(self) -> None:
        config = replace(
            SimulationConfig(),
            duration=0.8,
            feedback_delay_s=0.04,
            contact_probe_fraction=0.0,
        )
        mapper = build_mapper(self.models)
        wall, _, q_slave_start = make_wall(config, self.models, mapper)
        result = simulate_scenario(
            SCENARIOS[0],
            config,
            self.models,
            wall,
            q_slave_start,
        )

        self.assertTrue(result.metrics["completed"])
        self.assertTrue(result.metrics["finite_state"])
        self.assertEqual(result.metrics["master_joint_limit_events"], 0)
        self.assertEqual(result.metrics["slave_joint_limit_events"], 0)
        self.assertEqual(result.metrics["differential_fallback_count"], 0)
        self.assertGreater(result.metrics["wall_contact_duration_s"], 0.0)
        self.assertGreaterEqual(
            result.metrics["tank_energy_min_J"],
            config.energy_min - 1e-12,
        )

    def test_matched_source_map_and_independent_energy_audit(self) -> None:
        config = replace(
            SimulationConfig(),
            duration=0.6,
            feedback_delay_s=0.04,
            contact_probe_fraction=0.0,
        )
        mapper = build_mapper(self.models)
        wall, _, q_slave_start = make_wall(config, self.models, mapper)
        matched = next(
            scenario
            for scenario in SCENARIOS
            if scenario.key == "matched_wrench_energy"
        )
        result = simulate_scenario(
            matched, config, self.models, wall, q_slave_start
        )
        active = result.logs["return_packet_active"] > 0.5
        self.assertTrue(np.any(active))
        self.assertTrue(np.all(result.logs["source_map_id"][active] >= 0))
        self.assertTrue(np.all(result.logs["map_id"][active] >= 0))

        audit = audit_haptic_energy(
            tau_candidate=result.logs["tau_master_candidate"],
            tau_projected=result.logs["tau_master_applied"],
            tau_accepted=result.logs["tau_master_accepted"],
            qd_master=result.logs["qd_master"],
            dt=config.dt,
            energy_initial=config.energy_initial,
            energy_min=config.energy_min,
            energy_max=config.energy_max,
            logged_preclip=result.logs["energy_preclip"],
        )
        self.assertLessEqual(audit.max_floor_deficit, 1e-12)
        self.assertLessEqual(audit.preclip_log_max_error, 1e-12)

    def test_time_domain_popc_is_a_distinct_executable_baseline(self) -> None:
        config = replace(
            SimulationConfig(),
            duration=0.6,
            feedback_delay_s=0.02,
            contact_probe_fraction=0.0,
        )
        mapper = build_mapper(self.models)
        wall, _, q_slave_start = make_wall(config, self.models, mapper)
        popc = next(
            scenario for scenario in SCENARIOS if scenario.key == "proposed_popc"
        )
        result = simulate_scenario(
            popc, config, self.models, wall, q_slave_start
        )
        self.assertTrue(result.metrics["completed"])
        self.assertIsNone(result.metrics["rho_min"])
        self.assertTrue(np.all(np.isfinite(result.logs["popc_damping_gain"])))


if __name__ == "__main__":
    unittest.main()