"""Tests for causal acceleration, friction, and replayable H2 ablations."""

from __future__ import annotations

import sys
import unittest
from pathlib import Path

import numpy as np


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

from core.estimation_signals import (  # noqa: E402
    AccelerationEstimatorConfig,
    AccelerationEstimateStatus,
    CalibratedResidualWrenchEstimator,
    CausalAccelerationEstimator,
    JointFrictionCalibration,
    ResidualAblation,
    ResidualTorqueModel,
    WrenchEstimatorCalibration,
)
from core.wrench_solver import UndampedSVDSolver  # noqa: E402


class EstimationSignalsTest(unittest.TestCase):
    def test_acceleration_estimator_is_causal_and_rejects_bad_dt(self) -> None:
        config = AccelerationEstimatorConfig(
            cutoff_hz=10.0,
            minimum_dt_s=0.001,
            maximum_dt_s=0.1,
        )
        estimator = CausalAccelerationEstimator(2, config)
        first = estimator.update(np.array([0.0, 0.0]), 1.0)
        self.assertIs(first.status, AccelerationEstimateStatus.INITIALIZING)

        second = estimator.update(np.array([0.1, -0.2]), 1.01)
        self.assertIs(second.status, AccelerationEstimateStatus.VALID)
        np.testing.assert_allclose(second.raw_acceleration, [10.0, -20.0])
        expected_alpha = 1.0 - np.exp(-2.0 * np.pi * 10.0 * 0.01)
        np.testing.assert_allclose(
            second.acceleration,
            expected_alpha * np.array([10.0, -20.0]),
        )

        invalid = estimator.update(np.array([99.0, 99.0]), 1.0101)
        self.assertIs(invalid.status, AccelerationEstimateStatus.INVALID_DT)
        # Invalid data do not advance the differentiator state.
        recovered = estimator.update(np.array([0.2, -0.4]), 1.02)
        self.assertIs(recovered.status, AccelerationEstimateStatus.VALID)
        np.testing.assert_allclose(recovered.raw_acceleration, [10.0, -20.0])

    def test_friction_fit_recovers_synthetic_coefficients(self) -> None:
        velocity = np.linspace(-1.2, 1.2, 101)
        velocity_samples = np.column_stack((velocity, 0.7 * velocity))
        true = JointFrictionCalibration(
            coulomb_nm=np.array([0.4, 0.2]),
            viscous_nm_per_rad_s=np.array([0.08, 0.12]),
            smoothing_velocity_rad_s=0.03,
        )
        torque_samples = np.vstack(
            [true.torque(sample) for sample in velocity_samples]
        )
        fitted = JointFrictionCalibration.fit(
            velocity_samples,
            torque_samples,
            smoothing_velocity_rad_s=0.03,
        )
        np.testing.assert_allclose(fitted.coulomb_nm, true.coulomb_nm, atol=1e-13)
        np.testing.assert_allclose(
            fitted.viscous_nm_per_rad_s,
            true.viscous_nm_per_rad_s,
            atol=1e-13,
        )

    def test_nominal_and_ablation_residuals_are_explicit(self) -> None:
        friction = JointFrictionCalibration(
            coulomb_nm=np.array([0.3, 0.2]),
            viscous_nm_per_rad_s=np.array([0.1, 0.05]),
        )
        calibration = WrenchEstimatorCalibration(
            joint_bias_nm=np.array([0.12, -0.08]),
            friction=friction,
            characteristic_length_m=0.4,
            damping=0.02,
            calibration_id="fixture-v1",
        )
        model = ResidualTorqueModel(calibration)
        velocity = np.array([0.5, -0.4])
        contact = np.array([1.2, -0.7])
        rigid = np.array([3.0, 4.0])
        measured = (
            rigid
            + contact
            + calibration.joint_bias_nm
            + friction.torque(velocity)
        )

        nominal = model.compute(measured, rigid, velocity)
        no_bias = model.compute(
            measured, rigid, velocity, ablation=ResidualAblation.no_bias()
        )
        no_friction = model.compute(
            measured,
            rigid,
            velocity,
            ablation=ResidualAblation.no_friction(),
        )
        np.testing.assert_allclose(nominal.residual_nm, contact, atol=1e-15)
        np.testing.assert_allclose(
            no_bias.residual_nm,
            contact + calibration.joint_bias_nm,
            atol=1e-15,
        )
        np.testing.assert_allclose(
            no_friction.residual_nm,
            contact + friction.torque(velocity),
            atol=1e-15,
        )
        self.assertFalse(calibration.joint_bias_nm.flags.writeable)

    def test_calibrated_estimator_accepts_matched_undamped_solver(self) -> None:
        joint_count = 7
        calibration = WrenchEstimatorCalibration(
            joint_bias_nm=np.zeros(joint_count),
            friction=JointFrictionCalibration.zeros(joint_count),
            characteristic_length_m=0.5,
            damping=0.03,
            relative_rank_tolerance=1e-8,
            calibration_id="synthetic-v1",
        )
        solver = UndampedSVDSolver(
            0.5, relative_rank_tolerance=1e-8
        )
        estimator = CalibratedResidualWrenchEstimator(calibration, solver)
        rng = np.random.default_rng(12)
        jacobian = rng.normal(size=(6, joint_count))
        wrench = np.array([2.0, -1.0, 3.0, 0.2, -0.3, 0.4])
        residual = jacobian.T @ wrench
        estimate = estimator.estimate(
            jacobian,
            measured_torque_nm=residual,
            rigid_body_torque_nm=np.zeros(joint_count),
            joint_velocity_rad_s=np.zeros(joint_count),
        )
        np.testing.assert_allclose(
            estimate.solve.wrench, wrench, rtol=1e-13, atol=1e-13
        )

        mismatched = UndampedSVDSolver(
            0.6, relative_rank_tolerance=1e-8
        )
        with self.assertRaisesRegex(ValueError, "characteristic length"):
            CalibratedResidualWrenchEstimator(calibration, mismatched)


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