from __future__ import annotations

import unittest

from app.simulation.components.amesim.flow.pipes import AmesimPnl0002
from app.simulation.core.medium import IdealGasMedium
from app.simulation.reporting.amesim_results import AmesimResults
from app.simulation.reporting.pnl0002_replay import (
    Pnl0002ReplayPaths,
    replay_pnl0002_amesim_states,
)


class Pnl0002ReplayTests(unittest.TestCase):
    def setUp(self) -> None:
        self.pipe = AmesimPnl0002(
            "pneumatic_83",
            IdealGasMedium(),
            diam=0.02,
            le=2.0,
            rr=0.00225,
        )
        self.paths = Pnl0002ReplayPaths(
            center_pressure="center_p",
            center_temperature="center_T",
            port_1_pressure="port_1_p",
            port_1_temperature="port_1_T",
            port_1_mass_flow="port_1_dm",
            port_2_pressure="port_2_p",
            port_2_temperature="port_2_T",
            port_2_mass_flow="port_2_dm",
            reynolds="re",
            friction_factor="ff",
        )

    def _matching_results(self) -> AmesimResults:
        center_pressure = 300000.0
        center_temperature = 320.0
        port_1_pressure = 500000.0
        port_1_temperature = 330.0
        port_2_pressure = 100000.0
        port_2_temperature = 310.0
        flow_1 = self.pipe.mass_flow(
            port_1_pressure,
            center_pressure,
            port_1_temperature,
        )
        flow_2 = self.pipe.mass_flow(
            port_2_pressure,
            center_pressure,
            center_temperature,
        )
        reynolds_1 = self.pipe.reynolds_number(flow_1, port_1_temperature)
        reynolds_2 = self.pipe.reynolds_number(flow_2, center_temperature)
        friction = 0.5 * (
            self.pipe.friction_factor(reynolds_1)
            + self.pipe.friction_factor(reynolds_2)
        )
        return AmesimResults(
            times=(0.0,),
            variables=(),
            saved_variable_indices=(),
            series_by_data_path={
                "center_p": (center_pressure,),
                "center_T": (center_temperature,),
                "port_1_p": (port_1_pressure,),
                "port_1_T": (port_1_temperature,),
                "port_1_dm": (flow_1 / -1.0e-3,),
                "port_2_p": (port_2_pressure,),
                "port_2_T": (port_2_temperature,),
                "port_2_dm": (flow_2 / -1.0e-3,),
                "re": (0.5 * (reynolds_1 + reynolds_2),),
                "ff": (friction,),
            },
            final_values_by_data_path={},
        )

    def test_matching_saved_states_replay_without_integration(self) -> None:
        initial_state = self.pipe.get_state_vector()

        report = replay_pnl0002_amesim_states(
            self.pipe,
            self._matching_results(),
            self.paths,
        )

        self.assertEqual(report["mode"], "amesim-state-replay-no-integration")
        self.assertEqual(report["pointCount"], 1)
        self.assertEqual(self.pipe.get_state_vector(), initial_state)
        row = report["rows"][0]
        self.assertAlmostEqual(row["port1MassFlowError"], 0.0, places=15)
        self.assertAlmostEqual(row["port2MassFlowError"], 0.0, places=15)
        self.assertAlmostEqual(row["reynoldsError"], 0.0, places=12)
        self.assertAlmostEqual(row["frictionFactorError"], 0.0, places=15)

    def test_missing_saved_variable_has_actionable_error(self) -> None:
        results = self._matching_results()
        results.series_by_data_path.pop("ff")

        with self.assertRaisesRegex(
            ValueError,
            "AMESim replay variable is not saved: ff",
        ):
            replay_pnl0002_amesim_states(self.pipe, results, self.paths)


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