"""Dataclasses that capture the thermal-hydraulic state of the plant."""

from __future__ import annotations

from dataclasses import dataclass, field, asdict

from . import constants

from .generator import GeneratorState


def clamp(value: float, min_value: float, max_value: float) -> float:
    return max(min_value, min(max_value, value))


@dataclass
class CoreState:
    fuel_temperature: float  # Kelvin
    neutron_flux: float  # neutrons/cm^2-s equivalent
    reactivity_margin: float  # delta rho
    power_output_mw: float  # MW thermal
    burnup: float  # fraction of fuel consumed
    clad_temperature: float | None = None  # Kelvin
    xenon_inventory: float = 0.0
    iodine_inventory: float = 0.0
    delayed_precursors: list[float] = field(default_factory=list)
    fission_product_inventory: dict[str, float] = field(default_factory=dict)
    emitted_particles: dict[str, float] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if self.clad_temperature is None:
            self.clad_temperature = self.fuel_temperature

    def update_burnup(self, dt: float) -> None:
        produced_energy_mwh = self.power_output_mw * (dt / 3600.0)
        self.burnup = clamp(self.burnup + produced_energy_mwh * 1e-5, 0.0, 0.99)

    def add_products(self, products: dict[str, float]) -> None:
        for element, amount in products.items():
            self.fission_product_inventory[element] = self.fission_product_inventory.get(element, 0.0) + amount

    def add_emitted_particles(self, particles: dict[str, float]) -> None:
        for name, amount in particles.items():
            self.emitted_particles[name] = self.emitted_particles.get(name, 0.0) + amount


@dataclass
class CoolantLoopState:
    temperature_in: float  # K
    temperature_out: float  # K
    pressure: float  # MPa
    mass_flow_rate: float  # kg/s
    steam_quality: float  # fraction of vapor
    inventory_kg: float = 0.0  # bulk mass of coolant
    level: float = 1.0  # fraction full relative to nominal volume
    energy_j: float = 0.0  # stored thermal/latent energy for the loop

    def average_temperature(self) -> float:
        return 0.5 * (self.temperature_in + self.temperature_out)


@dataclass
class TurbineState:
    steam_enthalpy: float  # kJ/kg
    shaft_power_mw: float
    electrical_output_mw: float
    condenser_temperature: float
    load_demand_mw: float = 0.0
    load_supplied_mw: float = 0.0
    status: str = "OFF"


@dataclass
class PumpState:
    active: bool
    flow_rate: float
    pressure: float
    status: str = "OFF"


@dataclass
class PlantState:
    core: CoreState
    primary_loop: CoolantLoopState
    secondary_loop: CoolantLoopState
    turbines: list[TurbineState]
    primary_pumps: list[PumpState] = field(default_factory=list)
    secondary_pumps: list[PumpState] = field(default_factory=list)
    generators: list[GeneratorState] = field(default_factory=list)
    aux_draws: dict[str, float] = field(default_factory=dict)
    heat_exchanger_efficiency: float = 0.0
    primary_to_secondary_delta_t: float = 0.0
    time_elapsed: float = field(default=0.0)

    def snapshot(self) -> dict[str, float]:
        return {
            "time_elapsed": self.time_elapsed,
            "core_temp": self.core.fuel_temperature,
            "core_power": self.core.power_output_mw,
            "neutron_flux": self.core.neutron_flux,
            "primary_outlet_temp": self.primary_loop.temperature_out,
            "secondary_pressure": self.secondary_loop.pressure,
            "turbine_electric": self.total_electrical_output(),
            "products": self.core.fission_product_inventory,
            "particles": self.core.emitted_particles,
            "primary_pumps": [pump.active for pump in self.primary_pumps],
            "secondary_pumps": [pump.active for pump in self.secondary_pumps],
            "generators": [gen.running or gen.starting for gen in self.generators],
        }

    def total_electrical_output(self) -> float:
        return sum(t.electrical_output_mw for t in self.turbines)

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_dict(cls, data: dict) -> "PlantState":
        core_blob = dict(data["core"])
        inventory = core_blob.pop("fission_product_inventory", {})
        particles = core_blob.pop("emitted_particles", {})
        # Backwards/forwards compatibility for optional core fields.
        core_blob.pop("dnb_margin", None)
        core_blob.pop("subcooling_margin", None)
        core_blob.setdefault("clad_temperature", core_blob.get("fuel_temperature", 295.0))
        turbines_blob = data.get("turbines")
        if turbines_blob is None:
            # Compatibility with previous single-turbine snapshots.
            old_turbine = data.get("turbine")
            turbines_blob = [old_turbine] if old_turbine else []
        turbines = [TurbineState(**t) for t in turbines_blob]
        prim_pumps_blob = data.get("primary_pumps", [])
        sec_pumps_blob = data.get("secondary_pumps", [])
        generators_blob = data.get("generators", [])
        generators = [GeneratorState(**g) for g in generators_blob]
        aux_draws = data.get("aux_draws", {})
        hx_eff = data.get("heat_exchanger_efficiency", 0.0)
        delta_t = data.get("primary_to_secondary_delta_t", 0.0)
        return cls(
            core=CoreState(**core_blob, fission_product_inventory=inventory, emitted_particles=particles),
            primary_loop=CoolantLoopState(**_with_energy(data["primary_loop"])),
            secondary_loop=CoolantLoopState(**_with_energy(data["secondary_loop"])),
            turbines=turbines,
            primary_pumps=[PumpState(**p) for p in prim_pumps_blob],
            secondary_pumps=[PumpState(**p) for p in sec_pumps_blob],
            generators=generators,
            aux_draws=aux_draws,
            heat_exchanger_efficiency=hx_eff,
            primary_to_secondary_delta_t=delta_t,
            time_elapsed=data.get("time_elapsed", 0.0),
        )


def _with_energy(loop_blob: dict) -> dict:
    """Backwards compatibility: derive energy if missing."""
    if "energy_j" in loop_blob:
        return loop_blob
    energy = 0.5 * (loop_blob.get("temperature_in", 295.0) + loop_blob.get("temperature_out", 295.0))
    energy *= loop_blob.get("inventory_kg", 0.0) * constants.COOLANT_HEAT_CAPACITY
    out = dict(loop_blob)
    out["energy_j"] = energy
    return out