"""Steam generator and turbine performance models."""

from __future__ import annotations

from dataclasses import dataclass
import logging

from . import constants
from .state import CoolantLoopState, TurbineState

LOGGER = logging.getLogger(__name__)


@dataclass
class SteamGenerator:
    drum_volume_m3: float = 200.0

    def steam_enthalpy(self, loop: CoolantLoopState) -> float:
        base = 2_700.0  # kJ/kg saturated steam
        quality_adjustment = 500.0 * loop.steam_quality
        return base + quality_adjustment


@dataclass
class Turbine:
    generator_efficiency: float = constants.GENERATOR_EFFICIENCY
    mechanical_efficiency: float = constants.STEAM_TURBINE_EFFICIENCY
    rated_output_mw: float = 400.0  # cap per unit electrical output
    spool_time: float = constants.TURBINE_SPOOL_TIME
    throttle: float = 1.0  # 0-1 valve position

    def step(
        self,
        loop: CoolantLoopState,
        state: TurbineState,
        steam_power_mw: float = 0.0,
        dt: float = 1.0,
    ) -> None:
        effective_mass_flow = loop.mass_flow_rate * max(0.0, loop.steam_quality)
        if steam_power_mw <= 0.0 and (loop.steam_quality <= 0.01 or effective_mass_flow <= 10.0):
            # No steam available; turbine should idle.
            state.shaft_power_mw = 0.0
            state.electrical_output_mw = 0.0
            state.load_demand_mw = 0.0
            state.load_supplied_mw = 0.0
            state.steam_enthalpy = 0.0
            state.condenser_temperature = max(305.0, loop.temperature_in - 20.0)
            return

        throttle = min(constants.TURBINE_THROTTLE_MAX, max(constants.TURBINE_THROTTLE_MIN, self.throttle))
        throttle_eff = 1.0 - constants.TURBINE_THROTTLE_EFFICIENCY_DROP * (constants.TURBINE_THROTTLE_MAX - throttle)

        enthalpy = 2_700.0 + loop.steam_quality * 600.0
        mass_flow = effective_mass_flow * 0.6 * throttle
        computed_power = (enthalpy * mass_flow / 1_000.0) / 1_000.0
        available_power = steam_power_mw if steam_power_mw > 0 else computed_power
        backpressure_loss = 1.0 - _backpressure_penalty(loop)
        shaft_power_mw = available_power * self.mechanical_efficiency * throttle_eff * backpressure_loss
        electrical = shaft_power_mw * self.generator_efficiency
        if electrical > self.rated_output_mw:
            electrical = self.rated_output_mw
            shaft_power_mw = electrical / max(1e-6, self.generator_efficiency)
        condenser_temp = max(305.0, loop.temperature_in - 20.0)
        state.steam_enthalpy = enthalpy
        state.shaft_power_mw = _ramp(state.shaft_power_mw, shaft_power_mw, dt, self.spool_time)
        state.electrical_output_mw = _ramp(state.electrical_output_mw, electrical, dt, self.spool_time)
        state.condenser_temperature = condenser_temp
        LOGGER.debug(
            "Turbine output: shaft=%.1fMW electrical=%.1fMW condenser=%.1fK",
            shaft_power_mw,
            electrical,
            condenser_temp,
        )


def _ramp(current: float, target: float, dt: float, time_constant: float) -> float:
    if time_constant <= 0.0:
        return target
    alpha = min(1.0, max(0.0, dt / max(1e-6, time_constant)))
    return current + (target - current) * alpha


def _backpressure_penalty(loop: CoolantLoopState) -> float:
    base = constants.CONDENSER_BASE_PRESSURE_MPA
    max_p = constants.CONDENSER_MAX_PRESSURE_MPA
    pressure = max(base, min(max_p, loop.pressure))
    if pressure <= base:
        return 0.0
    frac = (pressure - base) / max(1e-6, max_p - base)
    return min(constants.CONDENSER_BACKPRESSURE_PENALTY, frac * constants.CONDENSER_BACKPRESSURE_PENALTY)