from __future__ import annotations

import unittest

from PythonModels.core.solver import SolveIVPConfig
from PythonModels.systems.test_mql import TestMqlSystem


class TestMqlPneumaticBranchSimulationTests(unittest.TestCase):
    def test_simulate_pneumatic_branch_returns_ode_solution(self) -> None:
        system = TestMqlSystem()

        solution = system.simulate_pneumatic_branch(
            name="sample_branch",
            upstream_volume_alias="pn_general_chamber",
            orifice_alias="pn_orifice_18",
            downstream_volume_alias="pn_c1_8",
            config=SolveIVPConfig(t_start=0.0, t_stop=1.0e-5, max_step=1.0e-6),
            t_eval=[0.0, 5.0e-6, 1.0e-5],
        )

        self.assertTrue(solution.success)
        self.assertEqual(len(solution.t), 3)
        self.assertEqual(len(solution.y), 4)
        self.assertEqual([len(row) for row in solution.y], [3, 3, 3, 3])

    def test_evaluate_pneumatic_branch_solution_reports_branch_series(self) -> None:
        system = TestMqlSystem()
        solution = system.simulate_pneumatic_branch(
            name="sample_branch",
            upstream_volume_alias="pn_general_chamber",
            orifice_alias="pn_orifice_18",
            downstream_volume_alias="pn_c1_8",
            config=SolveIVPConfig(t_start=0.0, t_stop=1.0e-5, max_step=1.0e-6),
            t_eval=[0.0, 5.0e-6, 1.0e-5],
        )

        series = system.evaluate_pneumatic_branch_solution(
            solution,
            name="sample_branch",
            upstream_volume_alias="pn_general_chamber",
            orifice_alias="pn_orifice_18",
            downstream_volume_alias="pn_c1_8",
        )

        self.assertEqual(series["time"], [0.0, 5.0e-6, 1.0e-5])
        self.assertIn("sample_branch.flow", series)
        self.assertIn("pn_general_chamber.p", series)
        self.assertIn("pn_c1_8.p", series)
        self.assertEqual(len(series["sample_branch.flow"]), 3)
        self.assertGreater(series["sample_branch.flow"][0], 0.0)
        self.assertGreater(series["pn_general_chamber.p"][0], series["pn_general_chamber.p"][-1])
        self.assertLess(series["pn_c1_8.p"][0], series["pn_c1_8.p"][-1])
        self.assertAlmostEqual(
            series["pn_general_chamber.m_flow"][0],
            -series["sample_branch.flow"][0],
        )
        self.assertAlmostEqual(
            series["pn_c1_8.m_flow"][0],
            series["sample_branch.flow"][0],
        )

    def test_pneumatic_branch_spec_drives_closure_and_simulation(self) -> None:
        system = TestMqlSystem()
        spec = system.pneumatic_branch_spec(
            name="sample_branch",
            upstream_volume_alias="pn_general_chamber",
            orifice_alias="pn_orifice_18",
            downstream_volume_alias="pn_c1_8",
            source="manual-test",
        )

        closure = system.pneumatic_branch_closure_from_spec(spec)
        solution = system.simulate_pneumatic_branch_from_spec(
            spec,
            config=SolveIVPConfig(t_start=0.0, t_stop=1.0e-5, max_step=1.0e-6),
            t_eval=[0.0, 1.0e-5],
        )
        series = system.evaluate_pneumatic_branch_solution_from_spec(solution, spec)

        self.assertIs(closure.components.spec, spec)
        self.assertEqual(spec.source, "manual-test")
        self.assertEqual(len(solution.t), 2)
        self.assertIn("sample_branch.flow", series)
        self.assertGreater(series["sample_branch.flow"][0], 0.0)
        self.assertGreater(series["pn_general_chamber.p"][0], series["pn_general_chamber.p"][-1])
        self.assertLess(series["pn_c1_8.p"][0], series["pn_c1_8.p"][-1])

    def test_simulate_pneumatic_branch_rejects_invalid_component_aliases(self) -> None:
        system = TestMqlSystem()

        with self.assertRaises(TypeError):
            system.simulate_pneumatic_branch(
                name="invalid_branch",
                upstream_volume_alias="pn_brp2_8",
                orifice_alias="pn_orifice_18",
                downstream_volume_alias="pn_c1_8",
                config=SolveIVPConfig(t_start=0.0, t_stop=0.0),
            )


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