#!/usr/bin/env python3
"""Unit tests for final applied-port haptic energy supervision.

These tests use no URDF, simulator, or Pinocchio model.  They exercise the
energy supervisor and the renderer's nominal shaping order directly.
"""

from __future__ import annotations

import os
import sys
import unittest

import numpy as np

CODE_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if CODE_DIR not in sys.path:
    sys.path.insert(0, CODE_DIR)

from core.haptic_render import (  # noqa: E402
    AppliedPortDiagnostics,
    AppliedPortEnergySupervisor,
    HapticRenderer,
    TankParams,
)


def make_supervisor(E0: float = 3.0,
                    E_min: float = 1.0,
                    E_max: float = 5.0) -> AppliedPortEnergySupervisor:
    return AppliedPortEnergySupervisor(
        TankParams(
            E_min=E_min,
            E_max=E_max,
            alpha_floor=0.0,
            alpha_ceil=1.0,
        ),
        E0=E0,
    )


def make_renderer_state(size: int,
                        *,
                        E0: float = 3.0,
                        feedback_strength: float = 1.0,
                        torque_limit=None,
                        torque_rate_limit=None) -> HapticRenderer:
    """Construct the pure torque pipeline without a Pinocchio dependency."""
    renderer = object.__new__(HapticRenderer)
    renderer.feedback_strength = float(feedback_strength)
    renderer.tank = make_supervisor(E0=E0)
    renderer.tau_alpha = 1.0
    renderer.torque_limit = torque_limit
    renderer.torque_rate_limit = torque_rate_limit
    renderer._tau_fb_state = np.zeros(size, dtype=float)
    renderer._tau_applied_prev = np.zeros(size, dtype=float)
    return renderer


class FakeChestJacobian:
    def __init__(self, jacobian: np.ndarray):
        self.jacobian = np.asarray(jacobian, dtype=float)

    def chest_jacobian(self, _q_m, _qd_m):
        return self.jacobian.copy()


class AppliedPortEnergySupervisorTest(unittest.TestCase):
    def test_low_level_apply_returns_complete_diagnostics(self):
        supervisor = make_supervisor(E0=3.0)
        candidate = np.array([2.0, -1.0])
        tau_app, diagnostics = supervisor.apply(
            candidate,
            np.array([0.5, 1.0]),
            dt=0.25,
        )

        self.assertIsInstance(diagnostics, AppliedPortDiagnostics)
        np.testing.assert_allclose(diagnostics.tau_candidate, candidate)
        np.testing.assert_allclose(diagnostics.tau_applied, tau_app)
        self.assertAlmostEqual(diagnostics.rho, 1.0)
        self.assertAlmostEqual(diagnostics.E_before, 3.0)
        self.assertAlmostEqual(diagnostics.E_preclip, 3.0)
        self.assertAlmostEqual(diagnostics.candidate_power, 0.0)
        self.assertAlmostEqual(diagnostics.power, 0.0)
        self.assertAlmostEqual(diagnostics.E_after, 3.0)
        self.assertFalse(diagnostics.fail_safe_active)
        self.assertIs(supervisor.last_diagnostics, diagnostics)

    def test_positive_power_discharges_once(self):
        supervisor = make_supervisor(E0=3.0)
        alpha, tau_app = supervisor.project_and_account(
            np.array([2.0, 0.0]),
            np.array([1.0, 0.0]),
            dt=0.25,
        )

        self.assertAlmostEqual(alpha, 1.0)
        np.testing.assert_allclose(tau_app, [2.0, 0.0])
        self.assertAlmostEqual(supervisor.last_power, 2.0)
        self.assertAlmostEqual(supervisor.E, 2.5)

    def test_negative_power_charges_once(self):
        supervisor = make_supervisor(E0=3.0)
        alpha, tau_app = supervisor.project_and_account(
            np.array([-2.0, 1.0]),
            np.array([1.0, 0.0]),
            dt=0.25,
        )

        self.assertAlmostEqual(alpha, 1.0)
        np.testing.assert_allclose(tau_app, [-2.0, 1.0])
        self.assertAlmostEqual(supervisor.last_power, -2.0)
        self.assertAlmostEqual(supervisor.E, 3.5)

    def test_zero_power_does_not_change_energy_or_torque(self):
        supervisor = make_supervisor(E0=3.0)
        candidate = np.array([2.0, -1.0])
        alpha, tau_app = supervisor.project_and_account(
            candidate,
            np.zeros(2),
            dt=0.25,
        )

        self.assertAlmostEqual(alpha, 1.0)
        np.testing.assert_allclose(tau_app, candidate)
        self.assertAlmostEqual(supervisor.last_power, 0.0)
        self.assertAlmostEqual(supervisor.E, 3.0)

    def test_energy_projection_hits_E_min_exactly(self):
        supervisor = make_supervisor(E0=1.25)
        alpha, tau_app = supervisor.project_and_account(
            np.array([2.0]),
            np.array([1.0]),
            dt=0.5,
        )

        self.assertAlmostEqual(alpha, 0.25)
        np.testing.assert_allclose(tau_app, [0.5])
        self.assertAlmostEqual(supervisor.last_power, 0.5)
        self.assertAlmostEqual(supervisor.E, 1.0)

        alpha_at_min, tau_at_min = supervisor.project_and_account(
            np.array([4.0]),
            np.array([1.0]),
            dt=0.1,
        )
        self.assertAlmostEqual(alpha_at_min, 0.0)
        np.testing.assert_allclose(tau_at_min, [0.0])
        self.assertAlmostEqual(supervisor.E, 1.0)

    def test_nonpositive_dt_uses_zero_fail_safe_without_accounting(self):
        supervisor = make_supervisor(E0=3.0)
        alpha, tau_app = supervisor.project_and_account(
            np.array([2.0]),
            np.array([1.0]),
            dt=0.0,
        )

        self.assertAlmostEqual(alpha, 0.0)
        np.testing.assert_allclose(tau_app, [0.0])
        self.assertTrue(supervisor.fail_safe_active)
        self.assertAlmostEqual(supervisor.E, 3.0)


class HapticRendererPipelineTest(unittest.TestCase):
    def test_external_supervisor_candidate_requires_explicit_commit(self):
        renderer = make_renderer_state(
            1,
            torque_limit=2.0,
            torque_rate_limit=10.0,
        )
        candidate, valid = renderer.shape_mapped_reaction_candidate(
            np.array([-10.0]),
            qd_m=np.zeros(1),
            dt=0.1,
        )
        self.assertTrue(valid)
        np.testing.assert_allclose(candidate, [-1.0])
        np.testing.assert_allclose(renderer._tau_applied_prev, [0.0])
        renderer.commit_applied(np.array([0.25]))
        np.testing.assert_allclose(renderer._tau_applied_prev, [0.25])

    def test_generalized_reaction_keeps_environment_on_device_sign(self):
        renderer = make_renderer_state(2)
        reaction = np.array([-2.0, 1.0])

        tau_app, rho = renderer.render_mapped_reaction(
            reaction,
            qd_m=np.zeros(2),
            dt=0.01,
        )

        np.testing.assert_allclose(tau_app, reaction)
        self.assertAlmostEqual(rho, 1.0)

    def test_direct_baseline_is_exact_J_transpose_F(self):
        renderer = object.__new__(HapticRenderer)
        jacobian = np.arange(12, dtype=float).reshape(6, 2) / 10.0
        renderer.CJ_master = FakeChestJacobian(jacobian)
        wrench = np.array([1.0, -2.0, 0.5, 3.0, -1.0, 2.0])

        tau_direct = renderer.direct_from_CF(
            q_m=np.zeros(2),
            qd_m=np.zeros(2),
            CF_int_slave_C=wrench,
        )

        np.testing.assert_allclose(tau_direct, jacobian.T @ wrench)

    def test_rate_limit_precedes_energy_projection(self):
        renderer = make_renderer_state(
            2,
            torque_rate_limit=np.array([2.0, 1.0]),
        )

        tau_step_1, alpha_1 = renderer._shape_and_supervise(
            tau_direct=np.array([-10.0, 10.0]),
            qd_m=np.zeros(2),
            dt=0.1,
        )
        tau_step_2, alpha_2 = renderer._shape_and_supervise(
            tau_direct=np.array([-10.0, 10.0]),
            qd_m=np.zeros(2),
            dt=0.1,
        )

        np.testing.assert_allclose(tau_step_1, [0.2, -0.1])
        np.testing.assert_allclose(tau_step_2, [0.4, -0.2])
        self.assertAlmostEqual(alpha_1, 1.0)
        self.assertAlmostEqual(alpha_2, 1.0)
        self.assertTrue(renderer.last_rate_limit_active)
        self.assertFalse(renderer.last_torque_saturation_active)
        np.testing.assert_allclose(
            renderer.tank.last_tau_candidate, tau_step_2
        )

    def test_saturation_precedes_energy_projection(self):
        renderer = make_renderer_state(
            2,
            torque_limit=np.array([3.0, 4.0]),
        )

        tau_app, alpha = renderer._shape_and_supervise(
            tau_direct=np.array([-10.0, 10.0]),
            qd_m=np.zeros(2),
            dt=0.1,
        )

        self.assertAlmostEqual(alpha, 1.0)
        np.testing.assert_allclose(tau_app, [3.0, -4.0])
        self.assertFalse(renderer.last_rate_limit_active)
        self.assertTrue(renderer.last_torque_saturation_active)
        np.testing.assert_allclose(
            renderer.tank.last_tau_candidate, [3.0, -4.0]
        )

    def test_projection_is_last_and_uses_saturated_candidate(self):
        renderer = make_renderer_state(
            1,
            E0=1.1,
            torque_limit=2.0,
        )

        tau_app, alpha = renderer._shape_and_supervise(
            tau_direct=np.array([-10.0]),
            qd_m=np.array([1.0]),
            dt=1.0,
        )

        # Direct -> nominal +10 -> saturation +2 -> energy projection +0.1.
        self.assertAlmostEqual(alpha, 0.05)
        np.testing.assert_allclose(renderer.tank.last_tau_candidate, [2.0])
        np.testing.assert_allclose(renderer.tank.last_tau_applied, [0.1])
        np.testing.assert_allclose(tau_app, renderer.tank.last_tau_applied)
        np.testing.assert_allclose(renderer._tau_applied_prev, tau_app)
        self.assertAlmostEqual(renderer.tank.last_power, float(tau_app[0]))
        self.assertAlmostEqual(renderer.tank.E, 1.0)

    def test_stepwise_offline_energy_reconstruction_matches_exactly(self):
        renderer = make_renderer_state(
            2,
            E0=2.5,
            feedback_strength=1.5,
            torque_limit=np.array([2.0, 1.5]),
            torque_rate_limit=np.array([4.0, 3.0]),
        )
        dt = 0.1
        direct_sequence = [
            np.array([-2.0, 1.0]),
            np.array([-4.0, 2.0]),
            np.array([1.0, -2.0]),
            np.array([0.0, 0.0]),
            np.array([-5.0, -5.0]),
            np.array([2.0, 2.0]),
        ]
        velocity_sequence = [
            np.array([1.0, 0.5]),
            np.array([0.7, -0.4]),
            np.array([1.0, 1.0]),
            np.zeros(2),
            np.array([-0.6, -0.2]),
            np.array([0.5, -0.8]),
        ]

        energy_offline = 2.5
        for tau_direct, qd_m in zip(direct_sequence, velocity_sequence):
            tau_app, _ = renderer._shape_and_supervise(
                tau_direct=tau_direct,
                qd_m=qd_m,
                dt=dt,
            )
            applied_power = float(np.dot(tau_app, qd_m))
            energy_offline = float(np.clip(
                energy_offline - applied_power * dt,
                renderer.tank.tp.E_min,
                renderer.tank.tp.E_max,
            ))

            np.testing.assert_allclose(
                renderer.tank.last_tau_applied, tau_app
            )
            self.assertAlmostEqual(renderer.tank.last_power, applied_power)
            self.assertAlmostEqual(renderer.tank.E, energy_offline, places=14)


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