from __future__ import annotations

from math import exp
import unittest

from app.simulation.components.amesim.mechanical.translational import (
    AmesimF000,
    AmesimLstp00a,
    AmesimMecmas21,
)
from app.simulation.core.base import AlgebraicComponent
from app.simulation.core.equations import EquationResidual
from app.simulation.core.medium import IdealGasMedium
from app.simulation.core.ports import PortDefinition
from app.simulation.solvers.algebraic import PressureFlowSolver
from app.simulation.systems.network import SimulationNetwork


class _PressureCoupledMechanicalLoad(AlgebraicComponent):
    PORTS = (
        PortDefinition.mechanical_translational("mechanical"),
        PortDefinition.pneumatic("pneumatic"),
    )

    def __init__(
        self,
        name: str,
        *,
        initial_force: float,
        solved_force: float,
        displacement: float,
    ) -> None:
        super().__init__(name=name)
        self.solved_force = float(solved_force)
        self.mechanical = self.register_declared_port("mechanical")
        self.pneumatic = self.register_declared_port("pneumatic")
        self.mechanical.x = float(displacement)
        self.pneumatic.p = float(initial_force)

    def pressure_flow_equation_residuals(self) -> tuple[EquationResidual, ...]:
        return (
            EquationResidual(
                id=f"{self.name}:velocity_state",
                owner="component",
                owner_id=self.name,
                relation="state",
                variables=(f"{self.name}.mechanical.v",),
                role="effort",
                value=self.mechanical.v,
            ),
            EquationResidual(
                id=f"{self.name}:pressure_force",
                owner="component",
                owner_id=self.name,
                relation="constitutive",
                variables=(
                    f"{self.name}.mechanical.f",
                    f"{self.name}.pneumatic.p",
                ),
                role="flow",
                value=self.mechanical.f + self.pneumatic.p,
            ),
            EquationResidual(
                id=f"{self.name}:pressure_closure",
                owner="component",
                owner_id=self.name,
                relation="constitutive",
                variables=(f"{self.name}.pneumatic.p",),
                role="effort",
                value=self.pneumatic.p - self.solved_force,
            ),
            EquationResidual(
                id=f"{self.name}:zero_mass_flow",
                owner="component",
                owner_id=self.name,
                relation="constitutive",
                variables=(f"{self.name}.pneumatic.m_flow",),
                role="flow",
                value=self.pneumatic.m_flow,
            ),
        )


class _PrescribedMechanicalLoad(AlgebraicComponent):
    PORTS = (PortDefinition.mechanical_translational("port_1"),)

    def __init__(
        self,
        name: str,
        *,
        force: float,
        displacement: float,
        velocity: float,
    ) -> None:
        super().__init__(name=name)
        self.force = float(force)
        self.velocity = float(velocity)
        self.port_1 = self.register_declared_port("port_1")
        self.port_1.x = float(displacement)

    def pressure_flow_equation_residuals(self) -> tuple[EquationResidual, ...]:
        return (
            EquationResidual(
                id=f"{self.name}:velocity_state",
                owner="component",
                owner_id=self.name,
                relation="state",
                variables=(f"{self.name}.port_1.v",),
                role="effort",
                value=self.port_1.v - self.velocity,
            ),
            EquationResidual(
                id=f"{self.name}:force",
                owner="component",
                owner_id=self.name,
                relation="constitutive",
                variables=(f"{self.name}.port_1.f",),
                role="flow",
                value=self.port_1.f - self.force,
            ),
        )


class ContactSolverCausalizationTests(unittest.TestCase):
    def test_nonlinear_binding_tracks_sub_ulp_force_change(self) -> None:
        medium = IdealGasMedium()
        load = _PressureCoupledMechanicalLoad(
            "load",
            initial_force=40.0,
            solved_force=41.0,
            displacement=1.0e9,
        )
        contact = AmesimLstp00a(
            "contact",
            medium,
            gap0=0.0,
            kcont=1.0e11,
            rcont=0.0,
            Pdis=1.0e-7,
            discContactOption=1.0,
        )
        mass = AmesimMecmas21(
            "mass",
            medium,
            mass=1.0,
            useFriction=1.0,
            stoptype=4.0,
            x0=1.0e9,
            v0=0.0,
        )
        zero = AmesimF000("zero")

        network = SimulationNetwork("sub-ulp-dynamic-contact-binding")
        for component in (load, contact, mass, zero):
            network.add_component(component)
        network.connect("load", "mechanical", "contact", "port_1")
        network.connect("contact", "port_2", "mass", "port_1")
        network.connect("mass", "port_2", "zero", "port_1")

        self.assertTrue(network.pressure_flow_structure_dict()["isSquare"])
        diagnostics = PressureFlowSolver(network).solve()

        self.assertTrue(diagnostics.success, diagnostics.message)
        self.assertGreater(diagnostics.evaluations, 0)
        self.assertEqual(diagnostics.jacobian_mode, "dense")
        self.assertFalse(diagnostics.dense_fallback_used)
        self.assertAlmostEqual(load.pneumatic.p, 41.0, delta=1.0e-3)
        self.assertAlmostEqual(load.mechanical.f, -41.0, delta=1.0e-3)
        self.assertAlmostEqual(contact.contact_force, 41.0, delta=1.0e-3)
        self.assertAlmostEqual(contact.penetration, 4.1e-10, delta=1.0e-14)

    def test_negative_contact_binding_uses_nearest_feasible_root(self) -> None:
        medium = IdealGasMedium()
        expected_force = 10.0 - 20.0 * (1.0 - exp(-1.0))
        load = _PrescribedMechanicalLoad(
            "load",
            force=-expected_force,
            displacement=0.0,
            velocity=0.0,
        )
        contact = AmesimLstp00a(
            "contact",
            medium,
            gap0=0.0,
            kcont=100.0,
            rcont=10.0,
            Pdis=0.1,
            discContactOption=1.0,
        )
        mass = AmesimMecmas21(
            "mass",
            medium,
            mass=1.0,
            useFriction=1.0,
            stoptype=4.0,
            x0=-0.08,
            v0=2.0,
        )
        zero = AmesimF000("zero")

        network = SimulationNetwork("negative-contact-force-binding")
        for component in (load, contact, mass, zero):
            network.add_component(component)
        network.connect("load", "port_1", "contact", "port_1")
        network.connect("contact", "port_2", "mass", "port_1")
        network.connect("mass", "port_2", "zero", "port_1")

        self.assertTrue(network.pressure_flow_structure_dict()["isSquare"])
        diagnostics = PressureFlowSolver(network).solve()

        self.assertTrue(diagnostics.success, diagnostics.message)
        self.assertLess(contact.contact_force, 0.0)
        self.assertAlmostEqual(contact.contact_force, expected_force, places=10)
        self.assertAlmostEqual(contact.penetration, 0.1, places=10)


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