from __future__ import annotations

from dataclasses import dataclass
from math import isfinite, sqrt

from app.simulation.core.ports import PortState, VariableRole
from app.simulation.systems.network import SimulationNetwork


class AlgebraicSolveError(RuntimeError):
    def __init__(self, message: str, diagnostics: "AlgebraicSolveDiagnostics") -> None:
        super().__init__(message)
        self.diagnostics = diagnostics


@dataclass(frozen=True)
class AlgebraicUnknown:
    component: str
    port: str
    variable: str
    role: VariableRole
    state: PortState

    @property
    def id(self) -> str:
        return f"{self.component}.{self.port}.{self.variable}"

    def read(self) -> float:
        return float(getattr(self.state, self.variable))

    def write(self, value: float) -> None:
        setattr(self.state, self.variable, float(value))


@dataclass(frozen=True)
class AlgebraicSolveDiagnostics:
    success: bool
    message: str
    evaluations: int
    pressure_scale: float
    flow_scale: float
    max_scaled_residual: float
    max_raw_residual: float

    def as_dict(self) -> dict[str, object]:
        return {
            "success": self.success,
            "message": self.message,
            "evaluations": self.evaluations,
            "pressureScale": self.pressure_scale,
            "flowScale": self.flow_scale,
            "maxScaledResidual": self.max_scaled_residual,
            "maxRawResidual": self.max_raw_residual,
        }


class PressureFlowSolver:
    """Solve the acausal pressure-flow subsystem for a compiled network."""

    def __init__(
        self,
        network: SimulationNetwork,
        *,
        residual_tolerance: float = 1e-7,
        max_evaluations: int = 500,
    ) -> None:
        self.network = network
        self.residual_tolerance = residual_tolerance
        self.max_evaluations = max_evaluations
        self.unknowns = self._build_unknowns()
        self._unknowns_by_id = {unknown.id: unknown for unknown in self.unknowns}
        self.last_diagnostics: AlgebraicSolveDiagnostics | None = None

    def _build_unknowns(self) -> tuple[AlgebraicUnknown, ...]:
        unknowns: list[AlgebraicUnknown] = []
        for component in self.network.components.values():
            for definition in component.port_definitions:
                if definition.kind != "physical":
                    continue
                state = component.get_port(definition.name)
                for variable in definition.variables:
                    if variable.role not in {"effort", "flow"}:
                        continue
                    unknowns.append(
                        AlgebraicUnknown(
                            component=component.name,
                            port=definition.name,
                            variable=variable.name,
                            role=variable.role,
                            state=state,
                        )
                    )
        return tuple(unknowns)

    @staticmethod
    def _port_key(variable: str, expected_variable: str) -> tuple[str, str] | None:
        try:
            component_name, port_name, variable_name = variable.rsplit(".", 2)
        except ValueError:
            return None
        if variable_name != expected_variable:
            return None
        return component_name, port_name

    def _seed_equal_pressures(self) -> None:
        """Lift current state pressures across their complete equality groups.

        Dynamic components refresh their own pressure ports before each closure,
        while connected algebraic ports retain values from the preceding RHS
        evaluation.  Merely filling non-positive pressures therefore leaves a
        stale, and sometimes badly conditioned, nonlinear initial guess.  State
        equations expose the current pressure as ``port.p - target``; use that
        target as the authoritative anchor for every connected/equal port.
        """

        pressure_unknowns = {
            (unknown.component, unknown.port): unknown
            for unknown in self.unknowns
            if unknown.variable == "p"
        }
        if not pressure_unknowns:
            return

        parent = {key: key for key in pressure_unknowns}

        def find(key: tuple[str, str]) -> tuple[str, str]:
            root = key
            while parent[root] != root:
                root = parent[root]
            while parent[key] != key:
                next_key = parent[key]
                parent[key] = root
                key = next_key
            return root

        def union(first: tuple[str, str], second: tuple[str, str]) -> None:
            first_root = find(first)
            second_root = find(second)
            if first_root != second_root:
                parent[second_root] = first_root

        for connection in self.network.connections:
            if connection.kind != "physical":
                continue
            first = connection.endpoint_a.key
            second = connection.endpoint_b.key
            if first in pressure_unknowns and second in pressure_unknowns:
                union(first, second)

        component_equations = {
            component.name: component.pressure_flow_equation_residuals()
            for component in self.network.components.values()
        }
        for equations in component_equations.values():
            for equation in equations:
                if equation.relation != "equal" or equation.role != "effort":
                    continue
                endpoints = [
                    endpoint
                    for variable in equation.variables
                    if (
                        (endpoint := self._port_key(variable, "p"))
                        in pressure_unknowns
                    )
                ]
                for endpoint in endpoints[1:]:
                    union(endpoints[0], endpoint)

        members_by_root: dict[tuple[str, str], list[tuple[str, str]]] = {}
        for endpoint in pressure_unknowns:
            members_by_root.setdefault(find(endpoint), []).append(endpoint)

        anchors_by_root: dict[tuple[str, str], list[float]] = {}
        for equations in component_equations.values():
            for equation in equations:
                if equation.relation != "state" or equation.role != "effort":
                    continue
                endpoints = [
                    endpoint
                    for variable in equation.variables
                    if (
                        (endpoint := self._port_key(variable, "p"))
                        in pressure_unknowns
                    )
                ]
                if len(endpoints) != 1:
                    continue
                endpoint = endpoints[0]
                unknown = pressure_unknowns[endpoint]
                target_pressure = unknown.read() - float(equation.value)
                if not isfinite(target_pressure):
                    continue
                # Keep the state-owned port current even when an invalid model
                # has conflicting storage anchors in one equality group.
                unknown.write(target_pressure)
                anchors_by_root.setdefault(find(endpoint), []).append(target_pressure)

        for root, members in members_by_root.items():
            anchors = anchors_by_root.get(root, [])
            if anchors:
                pressure_scale = max([abs(value) for value in anchors] + [1.0])
                if max(anchors) - min(anchors) > 1.0e-9 * pressure_scale:
                    # A conflicting multi-storage group is structurally invalid;
                    # leave it for the residual solver/preparation diagnostics.
                    continue
                target_pressure = sum(anchors) / len(anchors)
                for endpoint in members:
                    pressure_unknowns[endpoint].write(target_pressure)
                continue

            positive_seed = next(
                (
                    pressure_unknowns[endpoint].read()
                    for endpoint in members
                    if pressure_unknowns[endpoint].read() > 0.0
                ),
                None,
            )
            if positive_seed is None:
                continue
            for endpoint in members:
                unknown = pressure_unknowns[endpoint]
                if unknown.read() <= 0.0:
                    unknown.write(positive_seed)

    def _seed_explicit_mass_flows(self) -> None:
        """Initialize explicit ``m_flow - f(...)`` constitutive relations.

        AMESim orifices and quasi-steady pneumatic lines expose one mass-flow
        unknown with unit coefficient.  Once pressure anchors are current, a
        residual correction places that flow directly on its constitutive
        surface and avoids asking the nonlinear optimizer to discover the
        square-root branch from a stale preceding-step value.
        """

        seeded_ids: set[str] = set()
        for component in self.network.components.values():
            for equation in component.pressure_flow_equation_residuals():
                if equation.relation != "constitutive" or equation.role != "flow":
                    continue
                mass_flow_unknowns = [
                    self._unknowns_by_id[variable]
                    for variable in equation.variables
                    if variable in self._unknowns_by_id
                    and self._unknowns_by_id[variable].variable == "m_flow"
                ]
                if len(mass_flow_unknowns) != 1:
                    continue
                unknown = mass_flow_unknowns[0]
                target_flow = unknown.read() - float(equation.value)
                if not isfinite(target_flow):
                    continue
                unknown.write(target_flow)
                seeded_ids.add(unknown.id)

        # Complete local two-port balances for explicit elements. Connection
        # flow equations remain available to align the adjacent component port.
        for component in self.network.components.values():
            for equation in component.pressure_flow_equation_residuals():
                if equation.relation != "sumToZero" or equation.role != "flow":
                    continue
                mass_flow_unknowns = [
                    self._unknowns_by_id[variable]
                    for variable in equation.variables
                    if variable in self._unknowns_by_id
                    and self._unknowns_by_id[variable].variable == "m_flow"
                ]
                if len(mass_flow_unknowns) != 2:
                    continue
                seeded = [
                    unknown for unknown in mass_flow_unknowns if unknown.id in seeded_ids
                ]
                if len(seeded) != 1:
                    continue
                other = next(
                    unknown for unknown in mass_flow_unknowns if unknown.id not in seeded_ids
                )
                other.write(-seeded[0].read())
                seeded_ids.add(other.id)

        # A physical connector imposes the same sum-to-zero flow rule as a
        # two-port component.  Once an explicit component flow is known, carry
        # that guess to the connected storage/boundary port as well.  For the
        # common volume-orifice-volume topology this makes the seeded state an
        # exact algebraic solution and avoids an unnecessary nonlinear solve on
        # every ODE/Jacobian evaluation.
        for connection in self.network.connections:
            if connection.kind != "physical":
                continue
            endpoint_unknowns = []
            for endpoint in connection.endpoints:
                unknown = self._unknowns_by_id.get(
                    f"{endpoint.component}.{endpoint.port}.m_flow"
                )
                if unknown is not None:
                    endpoint_unknowns.append(unknown)
            if len(endpoint_unknowns) != 2:
                continue
            seeded = [
                unknown for unknown in endpoint_unknowns if unknown.id in seeded_ids
            ]
            if len(seeded) != 1:
                continue
            other = next(
                unknown for unknown in endpoint_unknowns if unknown.id not in seeded_ids
            )
            other.write(-seeded[0].read())
            seeded_ids.add(other.id)

    def _scales(self) -> dict[str, float]:
        pressure_scale = max(
            [
                abs(unknown.read())
                for unknown in self.unknowns
                if unknown.variable == "p" and unknown.read() > 0.0
            ]
            + [1e5]
        )
        estimated_flows = [
            abs(float(getattr(component, "K_eff"))) * sqrt(pressure_scale)
            for component in self.network.components.values()
            if hasattr(component, "K_eff")
        ]
        mass_flow_scale = max(
            estimated_flows
            + [
                abs(unknown.read())
                for unknown in self.unknowns
                if unknown.variable == "m_flow"
            ]
            + [1e-3]
        )
        return {
            "p": pressure_scale,
            "m_flow": mass_flow_scale,
            "x": max(
                [abs(unknown.read()) for unknown in self.unknowns if unknown.variable == "x"]
                + [1.0]
            ),
            "v": max(
                [abs(unknown.read()) for unknown in self.unknowns if unknown.variable == "v"]
                + [1.0]
            ),
            "f": max(
                [abs(unknown.read()) for unknown in self.unknowns if unknown.variable == "f"]
                + [1.0]
            ),
        }

    def solve(self) -> AlgebraicSolveDiagnostics:
        try:
            import numpy as np
            from scipy.optimize import least_squares
        except ImportError as exc:
            raise RuntimeError(
                "Topology-driven simulation requires SciPy; install requirements.txt."
            ) from exc

        self._seed_equal_pressures()
        self._seed_explicit_mass_flows()
        scales = self._scales()
        pressure_scale = scales["p"]
        flow_scale = scales["m_flow"]
        positive_pressures = [
            unknown.read()
            for unknown in self.unknowns
            if unknown.variable == "p" and unknown.read() > 0.0
        ]
        fallback_pressure = (
            sum(positive_pressures) / len(positive_pressures)
            if positive_pressures
            else pressure_scale
        )

        def variable_scale(unknown: AlgebraicUnknown) -> float:
            return scales.get(unknown.variable, max(abs(unknown.read()), 1.0))

        def equation_scale(equation) -> float:
            variable_names = [
                variable.rsplit(".", 1)[-1]
                for variable in equation.variables
            ]
            if equation.role == "flow":
                return scales["f"] if "f" in variable_names else flow_scale
            if equation.role == "effort":
                if "x" in variable_names:
                    return scales["x"]
                if "v" in variable_names:
                    return scales["v"]
                return pressure_scale
            return max([scales.get(name, 1.0) for name in variable_names] + [1.0])

        seeded_equations = self.network.pressure_flow_equation_residuals()
        seeded_scaled = [
            abs(equation.value / equation_scale(equation))
            for equation in seeded_equations
        ]
        seeded_max_scaled_residual = max(seeded_scaled, default=0.0)
        seeded_unknown_values = [
            (unknown, unknown.read()) for unknown in self.unknowns
        ]
        seeded_unknowns_are_feasible = all(
            isfinite(value)
            and (unknown.variable != "p" or value >= 1.0)
            for unknown, value in seeded_unknown_values
        )
        if (
            seeded_unknowns_are_feasible
            and all(isfinite(value) for value in seeded_scaled)
            and seeded_max_scaled_residual <= self.residual_tolerance
        ):
            diagnostics = AlgebraicSolveDiagnostics(
                success=True,
                message="Seeded pressure-flow state satisfies the residual tolerance.",
                evaluations=0,
                pressure_scale=pressure_scale,
                flow_scale=flow_scale,
                max_scaled_residual=seeded_max_scaled_residual,
                max_raw_residual=max(
                    (abs(item.value) for item in seeded_equations),
                    default=0.0,
                ),
            )
            self.last_diagnostics = diagnostics
            return diagnostics

        x0 = np.asarray(
            [
                (
                    unknown.read()
                    if unknown.variable != "p" or unknown.read() > 0.0
                    else fallback_pressure
                )
                / variable_scale(unknown)
                for unknown in self.unknowns
            ],
            dtype=float,
        )
        lower = np.asarray(
            [
                1.0 / pressure_scale if unknown.variable == "p" else -np.inf
                for unknown in self.unknowns
            ]
        )
        upper = np.full(len(self.unknowns), np.inf)

        def assign(values) -> None:
            for unknown, value in zip(self.unknowns, values):
                unknown.write(float(value) * variable_scale(unknown))

        def scaled_residuals(values):
            assign(values)
            equations = self.network.pressure_flow_equation_residuals()
            return np.asarray(
                [
                    equation.value / equation_scale(equation)
                    for equation in equations
                ],
                dtype=float,
            )

        result = least_squares(
            scaled_residuals,
            x0,
            bounds=(lower, upper),
            x_scale="jac",
            ftol=1e-10,
            xtol=1e-10,
            gtol=1e-10,
            max_nfev=self.max_evaluations,
        )
        assign(result.x)
        equations = self.network.pressure_flow_equation_residuals()
        scaled = [
            abs(
                equation.value / equation_scale(equation)
            )
            for equation in equations
        ]
        max_scaled_residual = max(scaled, default=0.0)
        residuals_converged = (
            all(isfinite(value) for value in scaled)
            and max_scaled_residual <= self.residual_tolerance
        )
        optimizer_status_is_acceptable = bool(result.success) or int(result.status) == 0
        success = residuals_converged and optimizer_status_is_acceptable
        diagnostics = AlgebraicSolveDiagnostics(
            success=success,
            message=str(result.message),
            evaluations=int(result.nfev),
            pressure_scale=pressure_scale,
            flow_scale=flow_scale,
            max_scaled_residual=max_scaled_residual,
            max_raw_residual=max((abs(item.value) for item in equations), default=0.0),
        )
        self.last_diagnostics = diagnostics
        if not success:
            raise AlgebraicSolveError(
                "Pressure-flow equations did not converge to the requested tolerance.",
                diagnostics,
            )
        return diagnostics