from __future__ import annotations

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_amesim_pnvo001_signal_xml import signal_edge, signal_port
from tests.test_generic_system_xml_simulation import component_node, physical_edge
from tests.test_system_xml_protocol import physical_port


MECMAS21_DEFAULTS = {
    "mass": 2.0,
    "fstick": 0.0,
    "fcoul": 0.0,
    "rvisc": 0.0,
    "wind": 0.0,
    "dvel": 1.0e-6,
    "restdvel": 1.0e-6,
    "restcoeff": 0.65,
    "astrib": 1.0e-3,
    "xmin": -1.0,
    "Kbmin": 1.0e9,
    "Dbmin": 1.0e4,
    "Pdmin": 1.0e-4,
    "xmax": 1.0,
    "Kbmax": 1.0e9,
    "Dbmax": 1.0e4,
    "Pdmax": 1.0e-4,
    "theta": 0.0,
    "useFriction": 1.0,
    "stoptype": 4.0,
    "discContactOption": 1.0,
    "strib": 1.0,
    "frictionType": 1.0,
    "v0": 0.0,
    "x0": 0.0,
}


def mechanical_port(name: str, side: str) -> dict[str, str]:
    return physical_port(name, "bidirectional", side, domain="mechanical")


def zero_force_mass_project() -> ReactFlowProjectPayload:
    return ReactFlowProjectPayload(
        name="amesim-mechanical-zero-force-smoke",
        nodes=[
            component_node("zero_left", "amesim_f000", [mechanical_port("port_1", "right")]),
            component_node(
                "mass_1",
                "amesim_mecmas21",
                [mechanical_port("port_1", "left"), mechanical_port("port_2", "right")],
                MECMAS21_DEFAULTS,
            ),
            component_node("zero_right", "amesim_f000", [mechanical_port("port_1", "left")]),
        ],
        edges=[
            physical_edge("edge-1", "zero_left", "port_1", "mass_1", "port_1"),
            physical_edge("edge-2", "mass_1", "port_2", "zero_right", "port_1"),
        ],
        simulation={"t_start": 0.0, "t_stop": 0.02, "step": 0.01, "max_step": 0.01, "method": "BDF"},
    )


def signal_force_mass_project() -> ReactFlowProjectPayload:
    parameters = dict(MECMAS21_DEFAULTS)
    parameters["mass"] = 2.0
    return ReactFlowProjectPayload(
        name="amesim-mechanical-signal-force-smoke",
        nodes=[
            component_node(
                "force_signal",
                "amesim_ud00",
                [signal_port("out", "output", "right")],
                {
                    "tstart": 0.0,
                    "start1": 10.0,
                    "end1": 10.0,
                    "t1": 1.0,
                    "start2": 10.0,
                    "end2": 10.0,
                    "t2": 0.0,
                    "start3": 10.0,
                    "end3": 10.0,
                    "t3": 0.0,
                    "start4": 10.0,
                    "end4": 10.0,
                    "t4": 0.0,
                    "start5": 10.0,
                    "end5": 10.0,
                    "t5": 0.0,
                    "start6": 10.0,
                    "end6": 10.0,
                    "t6": 0.0,
                    "start7": 10.0,
                    "end7": 10.0,
                    "t7": 0.0,
                    "start8": 10.0,
                    "end8": 10.0,
                    "t8": 0.0,
                    "nstages": 1.0,
                    "iscyclic": 0.0,
                },
            ),
            component_node(
                "force_1",
                "amesim_forc",
                [signal_port("res", "input", "left"), mechanical_port("port_2", "right")],
            ),
            component_node(
                "mass_1",
                "amesim_mecmas21",
                [mechanical_port("port_1", "left"), mechanical_port("port_2", "right")],
                parameters,
            ),
            component_node("zero_right", "amesim_f000", [mechanical_port("port_1", "left")]),
        ],
        edges=[
            signal_edge("signal-1", "force_signal", "out", "force_1", "res"),
            physical_edge("edge-1", "force_1", "port_2", "mass_1", "port_1"),
            physical_edge("edge-2", "mass_1", "port_2", "zero_right", "port_1"),
        ],
        simulation={"t_start": 0.0, "t_stop": 0.02, "step": 0.01, "max_step": 0.005, "method": "BDF"},
    )


class AmesimMechanicalXmlTests(unittest.TestCase):
    def test_zero_force_mechanical_project_compiles_and_simulates(self) -> None:
        xml = build_reactflow_system_xml(zero_force_mass_project())
        report = validate_system_xml_document(xml)

        self.assertTrue(report.valid, report.as_dict())
        network = compile_reactflow_network(zero_force_mass_project())
        self.assertTrue(network.pressure_flow_structure_dict()["isSquare"])
        result = run_system_xml_simulation(xml)

        self.assertTrue(result["success"], result["message"])
        self.assertEqual(result["series"]["time"], [0.0, 0.01, 0.02])
        self.assertEqual(result["series"]["mass_1.v"], [0.0, 0.0, 0.0])
        self.assertEqual(result["series"]["mass_1.x"], [0.0, 0.0, 0.0])

    def test_signal_force_mechanical_project_compiles_and_simulates(self) -> None:
        xml = build_reactflow_system_xml(signal_force_mass_project())
        report = validate_system_xml_document(xml)

        self.assertTrue(report.valid, report.as_dict())
        network = compile_reactflow_network(signal_force_mass_project())
        self.assertTrue(network.pressure_flow_structure_dict()["isSquare"])
        result = run_system_xml_simulation(xml)

        self.assertTrue(result["success"], result["message"])
        self.assertEqual(result["series"]["force_signal.out.signal"], [10.0, 10.0, 10.0])
        self.assertEqual(result["series"]["force_1.res.signal"], [10.0, 10.0, 10.0])
        self.assertAlmostEqual(result["series"]["mass_1.a"][0], 5.0)
        self.assertGreater(result["series"]["mass_1.v"][-1], 0.0)
        self.assertGreater(result["series"]["mass_1.x"][-1], 0.0)


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