from __future__ import annotations

from collections.abc import Mapping
from types import SimpleNamespace
import unittest

import numpy as np
from scipy.integrate._ivp.common import num_jac

from app.main import ReactFlowProjectPayload, compile_reactflow_network
from app.simulation.components.experimental.storage.cylinder import Cylinder
from app.simulation.components.experimental.storage.tank import Tank
from app.simulation.core.base import DynamicComponent
from app.simulation.core.ports import PortDefinition
from app.simulation.core.medium import IdealGasMedium
from app.simulation.solvers.mechanical import MechanicalConstraintGroup
from app.simulation.systems.generic import GenericFluidSystem
from app.simulation.systems.network import Endpoint
from tests.test_amesim_pnrp17_xml import pnrp17_coupled_project
from tests.test_generic_system_xml_simulation import component_node, physical_edge
from tests.test_system_xml_protocol import physical_port


class _DynamicVolumeSource(DynamicComponent):
    PORTS = (PortDefinition.pneumatic("port"),)
    state_size = 1

    def __init__(self, name: str) -> None:
        super().__init__(name)
        self.state = 0.25
        self.port = self.register_declared_port("port")

    def get_state_vector(self) -> list[float]:
        return [self.state]

    def set_state_vector(self, values: list[float]) -> None:
        self.state = float(values[0])

    def refresh_thermodynamic_ports(self):
        return None

    def state_derivative_from_ports(
        self,
        connected_h: Mapping[str, float],
    ) -> list[float]:
        return [0.0]

    def pneumatic_volume_outputs(self) -> Mapping[str, tuple[float, float]]:
        return {"port": (self.state, 0.0)}


class _TrustedDynamicVolumeSource(_DynamicVolumeSource):
    pass


_TrustedDynamicVolumeSource.__module__ = "app.simulation.components.synthetic"


class _UntrustedDynamicVolumeSource(_DynamicVolumeSource):
    pass


def cross_domain_storage_project() -> ReactFlowProjectPayload:
    """PNL storage -> variable chamber -> pneumatic piston -> two masses."""

    base = pnrp17_coupled_project()
    nodes = [node.model_dump() for node in base.nodes]
    nodes.append(
        component_node(
            "line_storage",
            "amesim_pnl0001",
            [
                physical_port("port_1", "bidirectional", "left"),
                physical_port("port_2", "bidirectional", "right"),
            ],
            {
                "diam": 0.01,
                "le": 1.0,
                "rr": 1.0e-5,
                "k": 1.35,
                "kth": 0.0,
                "extemp": 300.0,
                "gi": 0.0,
                "mode": 2.0,
                "p0": 200000.0,
                "T0": 300.0,
            },
        )
    )
    edges = [
        edge.model_dump()
        for edge in base.edges
        if edge.id != "edge-boundary-1"
    ]
    edges.extend(
        (
            physical_edge(
                "edge-chamber-line",
                "chamber_1",
                "port_1",
                "line_storage",
                "port_1",
            ),
            physical_edge(
                "edge-line-boundary",
                "line_storage",
                "port_2",
                "boundary_1",
                "port_1",
            ),
        )
    )
    return ReactFlowProjectPayload(
        projectSchemaVersion=base.projectSchemaVersion,
        name="cross-domain-jacobian-sparsity",
        nodes=nodes,
        edges=edges,
        simulation=base.simulation.model_dump(),
    )


def state_slices(system: GenericFluidSystem) -> dict[str, slice]:
    result: dict[str, slice] = {}
    cursor = 0
    for entry in system.mechanical_state_reducer.state_entries:
        if isinstance(entry, MechanicalConstraintGroup):
            entry_size = 2
            names = tuple(component.name for component in entry.components)
        else:
            entry_size = entry.state_size
            names = (entry.name,)
        state_slice = slice(cursor, cursor + entry_size)
        for name in names:
            result[name] = state_slice
        cursor += entry_size
    return result


class GenericJacobianSparsityTests(unittest.TestCase):
    def setUp(self) -> None:
        self.system = GenericFluidSystem(
            compile_reactflow_network(cross_domain_storage_project())
        )

    def test_external_volume_connects_mechanical_and_nearby_storage_states(self) -> None:
        slices = state_slices(self.system)
        pattern = self.system.jacobian_sparsity().toarray().astype(bool)
        line_states = range(
            slices["line_storage"].start,
            slices["line_storage"].stop,
        )
        mechanical_states = [
            state_index
            for name in ("piston_mass", "cylinder_mass")
            for state_index in range(slices[name].start, slices[name].stop)
        ]

        self.assertTrue(
            pattern[np.ix_(tuple(line_states), tuple(mechanical_states))].all()
        )
        self.assertTrue(
            pattern[np.ix_(tuple(mechanical_states), tuple(line_states))].all()
        )

    def _apply_dynamic_volume_dependency_probe(
        self,
        source: _DynamicVolumeSource,
    ) -> list[set[int]]:
        medium = IdealGasMedium()
        receiver = Cylinder("receiver", medium, V=0.1, p0=200_000.0)
        remote = Tank("remote", medium, V=0.1, p0=100_000.0)
        system = GenericFluidSystem.__new__(GenericFluidSystem)
        system.pneumatic_volume_resolver = SimpleNamespace(
            _output_components=(source,),
            _connected_endpoint={
                Endpoint(source.name, "port"): SimpleNamespace(
                    connected_endpoint=Endpoint("receiver", "port_b")
                )
            },
        )
        dependencies = [
            {0, 1},
            {0, 1},
            {1, 2},
        ]
        system._add_pneumatic_volume_state_dependencies(
            dependencies,
            (source, receiver, remote),
            {"source": 0, "receiver": 1, "remote": 2},
        )
        return dependencies

    def test_dynamic_volume_source_ode_state_drives_remote_pneumatic_state(
        self,
    ) -> None:
        dependencies = self._apply_dynamic_volume_dependency_probe(
            _TrustedDynamicVolumeSource("source")
        )

        self.assertIn(0, dependencies[2])
        self.assertIn(2, dependencies[0])

    def test_untrusted_volume_source_disables_ode_jacobian_sparsity(self) -> None:
        dependencies = self._apply_dynamic_volume_dependency_probe(
            _UntrustedDynamicVolumeSource("source")
        )

        self.assertEqual(dependencies, [{0, 1, 2}] * 3)

    def test_dense_numerical_jacobian_has_no_significant_entry_outside_pattern(
        self,
    ) -> None:
        state = np.asarray(self.system.consistent_initial_state_vector(0.0))
        slices = state_slices(self.system)
        # Move one piston face away from the zero-volume reference so the
        # chamber/line flow has a measurable local volume derivative.  This is
        # an operating-point probe only; the state remains well inside the
        # chamber's positive total-volume domain.
        state[slices["piston_mass"].stop - 1] = 1.0e-3

        def evaluate_one(values: np.ndarray) -> np.ndarray:
            return np.asarray(
                self.system.rhs(0.0, [float(value) for value in values])
            )

        def evaluate(_time: float, values: np.ndarray) -> np.ndarray:
            if values.ndim == 1:
                return evaluate_one(values)
            return np.column_stack(
                [evaluate_one(values[:, index]) for index in range(values.shape[1])]
            )

        derivative = evaluate(0.0, state)
        absolute_tolerance = np.asarray(
            self.system.mechanical_state_reducer.absolute_tolerances(1.0e-8)
        )
        dense_jacobian, _factor = num_jac(
            evaluate,
            0.0,
            state,
            derivative,
            absolute_tolerance / 1.0e-6,
            None,
            None,
        )
        pattern = self.system.jacobian_sparsity().toarray().astype(bool)
        missed = np.abs(np.asarray(dense_jacobian)) * ~pattern
        column_scale = np.maximum(
            1.0,
            np.max(np.abs(np.asarray(dense_jacobian)), axis=0),
        )

        self.assertFalse(
            np.any(missed > 1.0e-6 * column_scale[np.newaxis, :]),
            f"maximum omitted derivative was {float(np.max(missed))}",
        )
        line_rows = range(
            slices["line_storage"].start,
            slices["line_storage"].stop,
        )
        piston_columns = range(
            slices["piston_mass"].start,
            slices["piston_mass"].stop,
        )
        self.assertGreater(
            float(
                np.max(
                    np.abs(
                        np.asarray(dense_jacobian)[
                            np.ix_(tuple(line_rows), tuple(piston_columns))
                        ]
                    )
                )
            ),
            1.0,
        )


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