from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from PythonModels.components.cylinder import Cylinder
from PythonModels.components.orifice import Orifice
from PythonModels.components.pipe import Pipe
from PythonModels.components.tank import Tank
from PythonModels.components.tee import Tee
from PythonModels.core.medium import IdealGasMedium
from PythonModels.core.network import SimulationNetwork
from PythonModels.core.solver import SolveIVPConfig, integrate_ode
from PythonModels.systems.testmodel_closure import (
    BranchClosureComponents,
    InitializationDiagnostics,
    TestModelClosure,
    TestModelClosureComponents,
    TestModelSnapshot,
)


@dataclass(frozen=True)
class CylinderConfig:
    volume: float = 0.01
    p0: float = 35e6
    T0: float = 300.0


@dataclass(frozen=True)
class OrificeConfig:
    K: float = 1e-5


@dataclass(frozen=True)
class TankConfig:
    volume: float = 0.1
    p0: float = 1e5
    T0: float = 300.0


@dataclass(frozen=True)
class PipeConfig:
    length: float = 5.0
    diameter: float = 0.02
    lambda_darcy: float = 0.02
    p0: float = 1e5
    T0: float = 300.0


@dataclass(frozen=True)
class BranchConfig:
    orifice: OrificeConfig = field(default_factory=OrificeConfig)
    pipe: PipeConfig = field(default_factory=PipeConfig)


@dataclass(frozen=True)
class TestModelConfig:
    cylinder: CylinderConfig = field(default_factory=CylinderConfig)
    upper_branch: BranchConfig = field(default_factory=BranchConfig)
    lower_branch: BranchConfig = field(default_factory=BranchConfig)
    tank: TankConfig = field(default_factory=TankConfig)


class TestModelSystem:
    """Runnable first-pass Python system for the current Testmodel topology.

    This version keeps the component split from the Modelica model while keeping
    the downstream tee-tank pressure coupling in the ODE framework. The original
    Modelica system is a tighter DAE because both pipe outlets discharge into an
    ideal lossless junction directly connected to the tank. Here the branch
    outlet flows are solved from a pressure-consistent energy balance so the
    outlet is no longer driven by an arbitrary conductance parameter.
    """

    def __init__(
        self,
        medium: IdealGasMedium | None = None,
        config: TestModelConfig | None = None,
    ) -> None:
        self.medium = medium or IdealGasMedium()
        self.config = config or TestModelConfig()

        self.mycylinder = Cylinder(
            name="mycylinder",
            medium=self.medium,
            V=self.config.cylinder.volume,
            p0=self.config.cylinder.p0,
            T0=self.config.cylinder.T0,
        )
        self.mytee = Tee(name="mytee")
        self.myorifice = Orifice(name="myorifice", K=self.config.upper_branch.orifice.K)
        self.mypipe = Pipe(
            name="mypipe",
            medium=self.medium,
            L=self.config.upper_branch.pipe.length,
            D=self.config.upper_branch.pipe.diameter,
            lambda_darcy=self.config.upper_branch.pipe.lambda_darcy,
            p0=self.config.upper_branch.pipe.p0,
            T0=self.config.upper_branch.pipe.T0,
        )
        self.myorifice1 = Orifice(name="myorifice1", K=self.config.lower_branch.orifice.K)
        self.mypipe1 = Pipe(
            name="mypipe1",
            medium=self.medium,
            L=self.config.lower_branch.pipe.length,
            D=self.config.lower_branch.pipe.diameter,
            lambda_darcy=self.config.lower_branch.pipe.lambda_darcy,
            p0=self.config.lower_branch.pipe.p0,
            T0=self.config.lower_branch.pipe.T0,
        )
        self.mytee1 = Tee(name="mytee1")
        self.mytank = Tank(
            name="mytank",
            medium=self.medium,
            V=self.config.tank.volume,
            p0=self.config.tank.p0,
            T0=self.config.tank.T0,
        )

        self.network = SimulationNetwork(name="Testmodel")
        for component in (
            self.mycylinder,
            self.mytee,
            self.myorifice,
            self.mypipe,
            self.myorifice1,
            self.mypipe1,
            self.mytee1,
            self.mytank,
        ):
            self.network.add_component(component)

        self.network.connect("mycylinder", "port_b", "mytee", "port_in")
        self.network.connect("mytee", "port_out1", "myorifice", "port_a")
        self.network.connect("myorifice", "port_b", "mypipe", "port_a")
        self.network.connect("mypipe", "port_b", "mytee1", "port_out2")
        self.network.connect("mytee", "port_out2", "myorifice1", "port_a")
        self.network.connect("myorifice1", "port_b", "mypipe1", "port_a")
        self.network.connect("mypipe1", "port_b", "mytee1", "port_out1")
        self.network.connect("mytee1", "port_in", "mytank", "port_a")

        self.closure = TestModelClosure(
            medium=self.medium,
            components=TestModelClosureComponents(
                cylinder=self.mycylinder,
                upstream_tee=self.mytee,
                upper_branch=BranchClosureComponents(
                    name="upper_branch",
                    orifice=self.myorifice,
                    pipe=self.mypipe,
                ),
                lower_branch=BranchClosureComponents(
                    name="lower_branch",
                    orifice=self.myorifice1,
                    pipe=self.mypipe1,
                ),
                downstream_tee=self.mytee1,
                tank=self.mytank,
            ),
            initial_state_vector=self.initial_state_vector,
            apply_state_vector=self.apply_state_vector,
        )

    def initial_state_vector(self) -> list[float]:
        return self.network.initial_state_vector()

    def apply_state_vector(self, values: list[float]) -> None:
        self.network.apply_state_vector(values)

    def consistent_initial_state_vector(self) -> list[float]:
        return self.closure.consistent_initial_state_vector()

    @property
    def last_solve_diagnostics(self):
        return self.closure.last_solve_diagnostics

    def initialize_consistent_state(
        self,
        max_iterations: int = 12,
        state_tolerance: float = 1e-9,
        flow_tolerance: float = 1e-9,
        enthalpy_tolerance: float = 1e-6,
        pressure_tolerance: float = 1e-6,
        strict_internal_solvers: bool = False,
    ) -> InitializationDiagnostics:
        return self.closure.initialize_consistent_state(
            max_iterations=max_iterations,
            state_tolerance=state_tolerance,
            flow_tolerance=flow_tolerance,
            enthalpy_tolerance=enthalpy_tolerance,
            pressure_tolerance=pressure_tolerance,
            strict_internal_solvers=strict_internal_solvers,
        )

    def project_downstream_pressure_constraints(self, *, strict: bool = False) -> None:
        self.closure.project_downstream_pressure_constraints(strict=strict)

    def snapshot(
        self,
        state_vector: list[float] | None = None,
        *,
        strict: bool = False,
    ) -> TestModelSnapshot:
        return self.closure.snapshot(state_vector, strict=strict)

    def rhs(self, _t: float, state_vector: list[float]) -> list[float]:
        return self.closure.rhs(state_vector)

    @staticmethod
    def _legacy_branch_series_key_map() -> tuple[tuple[str, str, str], tuple[str, str, str]]:
        return (
            ("upper_branch", "branch_upper.in", "branch_upper.out"),
            ("lower_branch", "branch_lower.in", "branch_lower.out"),
        )

    @classmethod
    def _legacy_branch_series_keys_by_name(cls) -> dict[str, tuple[str, str]]:
        return {
            branch_name: (inlet_key, outlet_key)
            for branch_name, inlet_key, outlet_key in cls._legacy_branch_series_key_map()
        }

    @staticmethod
    def _generic_branch_series_keys(branch_name: str) -> tuple[str, str, str]:
        return (
            f"branch.{branch_name}.p",
            f"branch.{branch_name}.in",
            f"branch.{branch_name}.out",
        )

    @staticmethod
    def _legacy_branch_pressure_keys_by_name() -> dict[str, str]:
        return {
            "upper_branch": "mypipe.p",
            "lower_branch": "mypipe1.p",
        }

    @classmethod
    def _append_legacy_branch_series_aliases(
        cls,
        series: dict[str, list[float]],
    ) -> dict[str, list[float]]:
        legacy_branch_series_keys = cls._legacy_branch_series_keys_by_name()
        legacy_branch_pressure_keys = cls._legacy_branch_pressure_keys_by_name()
        for branch_name, (legacy_inlet_key, legacy_outlet_key) in legacy_branch_series_keys.items():
            pressure_key, generic_inlet_key, generic_outlet_key = cls._generic_branch_series_keys(
                branch_name
            )
            series[legacy_branch_pressure_keys[branch_name]] = list(series[pressure_key])
            series[legacy_inlet_key] = list(series[generic_inlet_key])
            series[legacy_outlet_key] = list(series[generic_outlet_key])
        return series

    def simulate(
        self,
        config: SolveIVPConfig | None = None,
        t_eval: list[float] | None = None,
    ) -> Any:
        return integrate_ode(
            rhs=self.rhs,
            initial_state=self.consistent_initial_state_vector(),
            config=config or SolveIVPConfig(),
            t_eval=t_eval,
        )

    def evaluate_solution(self, solution: Any) -> dict[str, list[float]]:
        series = {
            "time": [],
            "mycylinder.p": [],
            "mycylinder.T": [],
            "mytank.p": [],
            "mytank.T": [],
        }
        for branch_name, _, _ in self._legacy_branch_series_key_map():
            pressure_key, inlet_key, outlet_key = self._generic_branch_series_keys(branch_name)
            series[pressure_key] = []
            series[inlet_key] = []
            series[outlet_key] = []

        for index, time_value in enumerate(solution.t):
            state_vector = [row[index] for row in solution.y]
            snapshot = self.snapshot(state_vector)
            series["time"].append(float(time_value))
            series["mycylinder.p"].append(snapshot.cylinder.p)
            series["mycylinder.T"].append(snapshot.cylinder.T)
            series["mytank.p"].append(snapshot.tank.p)
            series["mytank.T"].append(snapshot.tank.T)
            for branch in snapshot.branches:
                pressure_key, generic_inlet_key, generic_outlet_key = self._generic_branch_series_keys(
                    branch.name
                )
                series[pressure_key].append(branch.pipe.p)
                series[generic_inlet_key].append(branch.inlet_flow)
                series[generic_outlet_key].append(branch.outlet_flow)

        return self._append_legacy_branch_series_aliases(series)


def build_testmodel() -> SimulationNetwork:
    """Compatibility helper for callers that only need the topology."""

    return TestModelSystem().network