from __future__ import annotations

import unittest
from pathlib import Path

from PythonModels.components.amesim_mechanical import (
    AmesimElasticEndstop,
    AmesimMassFrictionEndstops,
    AmesimPistonGeometry,
    circular_area,
    m_to_mm,
    mm_to_m,
)
from PythonModels.reporting.amesim_results import load_test_mql_amesim_results


REPO_ROOT = Path(__file__).resolve().parents[1]
TEST_MQL_AME = REPO_ROOT / "AmesimModels" / "test_mql.ame"


class AmesimMechanicalComponentsTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.amesim_results = load_test_mql_amesim_results(TEST_MQL_AME)

    def test_unit_helpers_and_area(self) -> None:
        self.assertAlmostEqual(mm_to_m(200.0), 0.2)
        self.assertAlmostEqual(m_to_mm(0.37), 370.0)
        self.assertAlmostEqual(circular_area(0.2), 0.031415926535897934)

    def test_piston_geometry_matches_amesim_length_and_volume(self) -> None:
        geometry = AmesimPistonGeometry(piston_diameter_m=0.2, rod_diameter_m=0.001)
        index = -1
        x4 = self.amesim_results.series("x4@pn_brp2_8")[index]
        x5 = self.amesim_results.series("x5@pn_brp2_8")[index]

        self.assertAlmostEqual(
            geometry.chamber_length_mm(x4, x5),
            self.amesim_results.series("length@pn_brp2_8")[index],
        )
        self.assertAlmostEqual(
            geometry.chamber_volume_cm3(x4, x5),
            self.amesim_results.series("vol1@pn_brp2_8")[index],
            delta=1.0e-3,
        )

    def test_piston_geometry_matches_amesim_volume_rate(self) -> None:
        geometry = AmesimPistonGeometry(piston_diameter_m=0.2, rod_diameter_m=0.001)
        index = 100
        v4 = self.amesim_results.series("v4@pn_brp2_8")[index]
        v5 = self.amesim_results.series("v5@pn_brp2_8")[index]

        self.assertAlmostEqual(
            geometry.chamber_volume_rate_l_min(v4, v5),
            self.amesim_results.series("vvol1@pn_brp2_8")[index],
            delta=1.0e-8,
        )

    def test_elastic_endstop_contact_force_matches_amesim_final_sample(self) -> None:
        endstop = AmesimElasticEndstop(
            contact_stiffness_n_per_m=1.0e11,
            contact_damping_n_per_m_per_s=1.0e11,
        )
        index = -1
        gap_mm = self.amesim_results.series("gap@elasticendstop_8")[index]
        penetration_velocity_m_s = self.amesim_results.series("v5@pn_brp2_8")[index]

        self.assertAlmostEqual(
            endstop.contact_force(gap_mm, penetration_velocity_m_s),
            self.amesim_results.series("f1@elasticendstop_8")[index],
            delta=1.0e-3,
        )

    def test_elastic_endstop_returns_zero_before_contact(self) -> None:
        endstop = AmesimElasticEndstop(contact_stiffness_n_per_m=1.0e11)

        self.assertEqual(endstop.static_contact_force(0.1), 0.0)
        self.assertEqual(endstop.penetration_m_from_gap_mm(0.1), 0.0)

    def test_mass_endstop_penetration_and_viscous_friction(self) -> None:
        mass = AmesimMassFrictionEndstops(
            mass_kg=50.0,
            lower_limit_m=-1.0,
            upper_limit_m=0.8,
            lower_stiffness_n_per_m=1.0e9,
            upper_stiffness_n_per_m=1.0e9,
            viscous_friction_n_per_m_per_s=12.0,
        )

        self.assertEqual(mass.lower_penetration_m(0.0), 0.0)
        self.assertEqual(mass.upper_penetration_m(0.0), 0.0)
        self.assertAlmostEqual(mass.lower_static_force_magnitude(-1.001), 1.0e6, delta=1.0e-6)
        self.assertAlmostEqual(mass.upper_static_force_magnitude(0.801), 1.0e6, delta=1.0e-6)
        self.assertAlmostEqual(mass.viscous_friction_force(0.25), -3.0)


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