from __future__ import annotations

import math
import unittest

from app.main import (
    ReactFlowProjectPayload,
    build_reactflow_system_xml,
    compile_reactflow_network,
    run_system_xml_simulation,
)
from app.system_xml import validate_system_xml_document
from tests.test_generic_system_xml_simulation import component_node, physical_edge
from tests.test_system_xml_protocol import physical_port


def _signal_port(name: str, role: str, side: str) -> dict[str, str | None]:
    return {
        "name": name,
        "kind": "signal",
        "domain": "signal",
        "nominalRole": role,
        "positiveFlowDirection": None,
        "side": side,
    }


def _signal_edge(
    edge_id: str,
    source: str,
    source_port: str,
    target: str,
    target_port: str,
) -> dict[str, str]:
    return {
        "id": edge_id,
        "source": source,
        "sourceHandle": source_port,
        "target": target,
        "targetHandle": target_port,
    }


def _plug(component_id: str) -> dict[str, object]:
    return component_node(
        component_id,
        "amesim_pnpl01",
        [physical_port("port_1", "bidirectional", "left")],
        {},
    )


def helium_step_equalization_project() -> ReactFlowProjectPayload:
    """Programmatic form of the user XML that failed at the 3.04 s sample."""

    return ReactFlowProjectPayload(
        projectSchemaVersion=1,
        name="amesim-helium-step-equalization-regression",
        nodes=[
            component_node(
                "amesim_helium_medium_1",
                "amesim_helium_medium",
                [],
                {"gi": 1.0, "property_model": 0.0},
            ),
            component_node(
                "amesim_pnch023_1",
                "amesim_pnch023",
                [
                    physical_port("port_1", "bidirectional", "left"),
                    physical_port("port_2", "bidirectional", "right"),
                ],
                {
                    "gi": 1.0,
                    "cvol": 0.057,
                    "kth": 0.0,
                    "sth": 0.1,
                    "extemp": 293.15,
                    "p0": 15_300_000.0,
                    "T0": 293.15,
                },
            ),
            component_node(
                "amesim_pnch012_1",
                "amesim_pnch012",
                [
                    physical_port("port_1", "bidirectional", "left"),
                    physical_port("port_2", "bidirectional", "right"),
                    physical_port("port_3", "bidirectional", "left"),
                    physical_port("port_4", "bidirectional", "right"),
                ],
                {
                    "gi": 1.0,
                    "cvol0": 0.015,
                    "kth": 1500.0,
                    "sth": 0.7,
                    "extemp": 293.15,
                    "p0": 100_000.0,
                    "T0": 293.15,
                    "vol1": 0.0,
                    "vol2": 0.0,
                    "vol3": 0.0,
                    "vol4": 0.0,
                    "dvol1": 0.0,
                    "dvol2": 0.0,
                    "dvol3": 0.0,
                    "dvol4": 0.0,
                },
            ),
            component_node(
                "amesim_pnvo001_1",
                "amesim_pnvo001",
                [
                    _signal_port("res", "input", "left"),
                    physical_port("port_2", "bidirectional", "left"),
                    physical_port("port_3", "bidirectional", "right"),
                ],
                {
                    "gi": 1.0,
                    "cq": 0.45,
                    "area0": 7.85e-5,
                    "Cv": 0.5,
                    "Kv": 0.4,
                    "flowset": 1.0,
                    "opening0": 1.0,
                },
            ),
            component_node(
                "amesim_step0_1",
                "amesim_step0",
                [_signal_port("out", "output", "right")],
                {"initial": 0.0, "final": 1.0, "time": 0.04},
            ),
            _plug("amesim_pnpl01_1"),
            _plug("amesim_pnpl01_2"),
            _plug("amesim_pnpl01_3"),
            _plug("amesim_pnpl01_4"),
        ],
        edges=[
            physical_edge(
                "edge-high-plug",
                "amesim_pnpl01_1",
                "port_1",
                "amesim_pnch023_1",
                "port_1",
            ),
            physical_edge(
                "edge-high-valve",
                "amesim_pnch023_1",
                "port_2",
                "amesim_pnvo001_1",
                "port_2",
            ),
            physical_edge(
                "edge-valve-low",
                "amesim_pnvo001_1",
                "port_3",
                "amesim_pnch012_1",
                "port_3",
            ),
            physical_edge(
                "edge-low-port-4",
                "amesim_pnch012_1",
                "port_4",
                "amesim_pnpl01_4",
                "port_1",
            ),
            physical_edge(
                "edge-low-port-2",
                "amesim_pnch012_1",
                "port_2",
                "amesim_pnpl01_3",
                "port_1",
            ),
            physical_edge(
                "edge-low-port-1",
                "amesim_pnpl01_2",
                "port_1",
                "amesim_pnch012_1",
                "port_1",
            ),
            _signal_edge(
                "edge-step-valve",
                "amesim_step0_1",
                "out",
                "amesim_pnvo001_1",
                "res",
            ),
        ],
        simulation={
            "t_start": 0.0,
            "t_stop": 10.0,
            "step": 0.02,
            "max_step": 0.002,
            "method": "BDF",
        },
    )


class AmesimHeliumStepLongRunTests(unittest.TestCase):
    def test_helium_step_equalization_completes_beyond_3_04_seconds(self) -> None:
        project = helium_step_equalization_project()
        network = compile_reactflow_network(project)
        self.assertEqual(
            network.components["amesim_pnch023_1"].medium.name,
            "AMESimHeliumPengRobinson",
        )
        self.assertEqual(
            network.components["amesim_pnch012_1"].medium.name,
            "AMESimHeliumPengRobinson",
        )

        xml = build_reactflow_system_xml(project)
        report = validate_system_xml_document(xml)
        self.assertTrue(report.valid, report.as_dict())

        result = run_system_xml_simulation(xml)

        self.assertTrue(result["success"], result["message"])
        self.assertEqual(result["status"], "completed")
        self.assertAlmostEqual(result["simulatedUntil"], 10.0)

        series = result["series"]
        self.assertIn(3.04, series["time"])
        self.assertTrue(
            all(
                math.isfinite(value)
                for values in series.values()
                for value in values
            )
        )

        high_pressure = series["amesim_pnch023_1.p"]
        low_pressure = series["amesim_pnch012_1.p"]
        initial_pressure_gap = abs(high_pressure[0] - low_pressure[0])
        final_pressure_gap = abs(high_pressure[-1] - low_pressure[-1])
        self.assertLess(high_pressure[-1], high_pressure[0])
        self.assertGreater(low_pressure[-1], low_pressure[0])
        self.assertLess(final_pressure_gap, initial_pressure_gap)
        self.assertLess(final_pressure_gap, 10.0)
        self.assertLess(
            abs(series["amesim_pnvo001_1.port_2.m_flow"][-1]),
            1.0e-4,
        )


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