from __future__ import annotations

import unittest
from pathlib import Path

from PythonModels.systems.test_mql_computed import (
    run_test_mql_computed_geometry,
    run_test_mql_computed_mechanical_relations,
    run_test_mql_computed_piston_geometry,
)


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


class TestMqlComputedPistonGeometryTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.computed_run = run_test_mql_computed_piston_geometry(TEST_MQL_AME)

    def test_computed_piston_geometry_outputs_schema_valid_series(self) -> None:
        self.assertEqual(self.computed_run.sample_count, 1002)
        self.assertEqual(len(self.computed_run.mechanical_assembly.pistons), 8)
        self.assertEqual(self.computed_run.signal_count, 24)
        self.assertEqual(len(self.computed_run.comparison.metrics), 24)
        self.assertEqual(self.computed_run.output_schema.signal_count, 858)

    def test_computed_piston_geometry_output_paths_are_expected(self) -> None:
        paths = self.computed_run.output.data_paths

        self.assertIn("length@pn_brp2_8", paths)
        self.assertIn("vol1@pn_brp2_8", paths)
        self.assertIn("vvol1@pn_brp2_8", paths)
        self.assertIn("length@pn_brp2_15", paths)
        self.assertEqual(len([path for path in paths if path.startswith("length@")]), 8)
        self.assertEqual(len([path for path in paths if path.startswith("vol1@")]), 8)
        self.assertEqual(len([path for path in paths if path.startswith("vvol1@")]), 8)

    def test_computed_piston_geometry_matches_amesim_with_tight_error_bounds(self) -> None:
        for data_path in self.computed_run.output.data_paths:
            metric = self.computed_run.comparison.metric(data_path)
            with self.subTest(data_path=data_path):
                if data_path.startswith("length@"):
                    self.assertLessEqual(metric.max_abs_error, 3.0e-13)
                elif data_path.startswith("vol1@"):
                    self.assertLessEqual(metric.max_abs_error, 8.0e-12)
                elif data_path.startswith("vvol1@"):
                    self.assertLessEqual(metric.max_abs_error, 3.0e-11)
                else:
                    self.fail(f"Unexpected computed piston geometry path: {data_path}")

    def test_computed_piston_geometry_final_samples_match_expected_values(self) -> None:
        self.assertAlmostEqual(
            self.computed_run.output.series("length@pn_brp2_8")[-1],
            1090.0013536465842,
        )
        self.assertAlmostEqual(
            self.computed_run.output.series("vol1@pn_brp2_8")[-1],
            34242.54636512914,
            delta=1.0e-9,
        )
        self.assertAlmostEqual(
            self.computed_run.output.series("vvol1@pn_brp2_8")[-1],
            8.360141002531034e-07,
            delta=1.0e-11,
        )


class TestMqlComputedGeometryTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.computed_run = run_test_mql_computed_geometry(TEST_MQL_AME)

    def test_computed_geometry_outputs_schema_valid_series(self) -> None:
        self.assertEqual(self.computed_run.sample_count, 1002)
        self.assertEqual(len(self.computed_run.mechanical_assembly.pistons), 8)
        self.assertEqual(len(self.computed_run.pneumatic_assembly.variable_chambers), 8)
        self.assertEqual(self.computed_run.signal_count, 32)
        self.assertEqual(len(self.computed_run.comparison.metrics), 32)
        self.assertEqual(self.computed_run.output_schema.signal_count, 858)

    def test_computed_geometry_includes_variable_chamber_volume_paths(self) -> None:
        paths = self.computed_run.output.data_paths

        self.assertIn("vol@pn_c1_8", paths)
        self.assertIn("vol@pn_c1_15", paths)
        self.assertEqual(len([path for path in paths if path.startswith("vol@pn_c1")]), 8)
        self.assertEqual(len([path for path in paths if path.startswith("length@")]), 8)
        self.assertEqual(len([path for path in paths if path.startswith("vol1@")]), 8)
        self.assertEqual(len([path for path in paths if path.startswith("vvol1@")]), 8)

    def test_computed_variable_chamber_volumes_match_amesim_with_tight_error_bounds(self) -> None:
        for data_path in self.computed_run.output.data_paths:
            metric = self.computed_run.comparison.metric(data_path)
            with self.subTest(data_path=data_path):
                if data_path.startswith("vol@pn_c1"):
                    self.assertLessEqual(metric.max_abs_error, 8.0e-12)

    def test_computed_variable_chamber_volume_final_samples_match_expected_values(self) -> None:
        self.assertAlmostEqual(
            self.computed_run.output.series("vol@pn_c1_8")[-1],
            49242.54636512914,
            delta=1.0e-9,
        )
        self.assertAlmostEqual(
            self.computed_run.output.series("vol@pn_c1_15")[-1],
            49242.546365149814,
            delta=1.0e-9,
        )


class TestMqlComputedMechanicalRelationsTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.computed_run = run_test_mql_computed_mechanical_relations(TEST_MQL_AME)

    def test_computed_mechanical_relations_outputs_schema_valid_series(self) -> None:
        self.assertEqual(self.computed_run.sample_count, 1002)
        self.assertEqual(len(self.computed_run.mechanical_assembly.masses), 10)
        self.assertEqual(self.computed_run.signal_count, 76)
        self.assertEqual(len(self.computed_run.comparison.metrics), 76)
        self.assertEqual(self.computed_run.output_schema.signal_count, 858)

    def test_computed_mechanical_relations_include_expected_paths(self) -> None:
        paths = self.computed_run.output.data_paths

        self.assertIn("x1dup@mass_friction_endstops_10", paths)
        self.assertIn("v1dup@mass_friction_endstops_10", paths)
        self.assertIn("acc1dup@mass_friction_endstops_10", paths)
        self.assertIn("Fmin@mass_friction_endstops_10", paths)
        self.assertIn("Fvisc@mass_friction_endstops_10", paths)
        self.assertIn("Ffric@mass_friction_endstops_10", paths)
        self.assertIn("fzero@zeroforcesource_17", paths)
        self.assertEqual(len([path for path in paths if path.startswith("x1dup@")]), 10)
        self.assertEqual(len([path for path in paths if path.startswith("v1dup@")]), 10)
        self.assertEqual(len([path for path in paths if path.startswith("acc1dup@")]), 10)
        self.assertEqual(len([path for path in paths if path.startswith("Fmin@")]), 10)
        self.assertEqual(len([path for path in paths if path.startswith("Fvisc@")]), 10)
        self.assertEqual(len([path for path in paths if path.startswith("Ffric@")]), 10)
        self.assertEqual(len([path for path in paths if path.startswith("fzero@")]), 16)

    def test_computed_mechanical_relations_match_amesim_exactly(self) -> None:
        for data_path in self.computed_run.output.data_paths:
            metric = self.computed_run.comparison.metric(data_path)
            with self.subTest(data_path=data_path):
                self.assertEqual(metric.max_abs_error, 0.0)

    def test_computed_mechanical_relations_final_samples_match_expected_values(self) -> None:
        self.assertAlmostEqual(
            self.computed_run.output.series("x1dup@mass_friction_endstops_10")[-1],
            -0.3700013536465843,
        )
        self.assertAlmostEqual(
            self.computed_run.output.series("v1dup@mass_friction_endstops_10")[-1],
            -4.435303434246023e-10,
        )
        self.assertAlmostEqual(
            self.computed_run.output.series("acc1dup@mass_friction_endstops_10")[-1],
            3.5662378650158644e-07,
        )
        self.assertEqual(
            self.computed_run.output.series("Fmin@mass_friction_endstops_10")[-1],
            0.0,
        )
        self.assertEqual(
            self.computed_run.output.series("Fvisc@mass_friction_endstops_10")[-1],
            0.0,
        )
        self.assertEqual(
            self.computed_run.output.series("Ffric@mass_friction_endstops_10")[-1],
            0.0,
        )
        self.assertEqual(
            self.computed_run.output.series("fzero@zeroforcesource_17")[-1],
            0.0,
        )


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