# -*- coding: utf-8 -*-
"""Mass-estimation models for Brayton cycle components."""

from dataclasses import dataclass
import math


PCHE_STATE_UNIT_CONVENTION = {
    "P": "kPa",
    "T": "K",
    "h": "kJ/kg",
    "mass_flow_rate": "kg/s",
    "density": "kg/m3",
    "viscosity": "Pa*s",
    "thermal_conductivity": "W/(m*K)",
    "specific_heat": "J/(kg*K)",
    "length": "m",
    "area": "m2",
    "pressure_drop": "Pa",
    "mass": "kg",
}


@dataclass(frozen=True)
class PCHEGeometry:
    """Default PCHE channel geometry from Yuan et al."""

    channel_diameter_m: float = 0.002
    channel_pitch_m: float = 0.0024
    plate_thickness_m: float = 0.0015

    def __post_init__(self):
        _require_positive("channel_diameter_m", self.channel_diameter_m)
        _require_positive("channel_pitch_m", self.channel_pitch_m)
        _require_positive("plate_thickness_m", self.plate_thickness_m)

    @property
    def flow_area_m2(self):
        return pche_channel_flow_area(self)

    @property
    def wetted_perimeter_m(self):
        return pche_channel_wetted_perimeter(self)

    @property
    def equivalent_diameter_m(self):
        return pche_equivalent_diameter(self)


@dataclass(frozen=True)
class PCHEMaterial:
    """Default Inconel 617 material data from Yuan et al."""

    density_kg_m3: float = 8360.0
    thermal_conductivity_w_m_k: float = 21.0

    def __post_init__(self):
        _require_positive("density_kg_m3", self.density_kg_m3)
        _require_positive(
            "thermal_conductivity_w_m_k",
            self.thermal_conductivity_w_m_k,
        )


@dataclass(frozen=True)
class PCHEDesignOptions:
    """Numerical and design limits for the Yuan-style PCHE calculation."""

    num_segments: int = 50
    allowable_pressure_drop_ratio: float = 0.01
    outlet_temperature_tolerance_k: float = 1.0e-3
    max_iterations: int = 100
    reynolds_transition: float = 2300.0

    def __post_init__(self):
        if self.num_segments < 1:
            raise ValueError("num_segments must be at least 1")
        if self.max_iterations < 1:
            raise ValueError("max_iterations must be at least 1")
        _require_positive(
            "allowable_pressure_drop_ratio",
            self.allowable_pressure_drop_ratio,
        )
        _require_positive(
            "outlet_temperature_tolerance_k",
            self.outlet_temperature_tolerance_k,
        )
        _require_positive("reynolds_transition", self.reynolds_transition)


def _require_positive(name, value):
    _require_finite(name, value)
    if value <= 0:
        raise ValueError(f"{name} must be positive")


def _require_finite(name, value):
    if not math.isfinite(value):
        raise ValueError(f"{name} must be finite")


def validate_pche_state(state, state_name="state", require_transport=False):
    """Validate the state fields expected by the PCHE mass model."""
    required = ("P", "T", "h")
    missing = [key for key in required if key not in state]
    if missing:
        raise ValueError(f"{state_name} missing required keys: {', '.join(missing)}")

    _require_positive(f"{state_name}.P", state["P"])
    _require_positive(f"{state_name}.T", state["T"])
    _require_finite(f"{state_name}.h", state["h"])

    if require_transport:
        transport_keys = (
            "rho",
            "cp_mass",
            "mu",
            "thermal_conductivity",
        )
        missing = [key for key in transport_keys if key not in state]
        if missing:
            raise ValueError(
                f"{state_name} missing transport keys: {', '.join(missing)}"
            )
        for key in transport_keys:
            _require_positive(f"{state_name}.{key}", state[key])


def pche_channel_flow_area(geometry=None):
    """Return the semicircular PCHE channel flow area in m2."""
    geometry = geometry or PCHEGeometry()
    diameter = geometry.channel_diameter_m
    return math.pi * diameter**2 / 8.0


def pche_channel_wetted_perimeter(geometry=None):
    """Return the Yuan PCHE channel wetted perimeter in m."""
    geometry = geometry or PCHEGeometry()
    diameter = geometry.channel_diameter_m
    return math.pi * diameter / 2.0 + diameter


def pche_equivalent_diameter(geometry=None):
    """Return the Yuan PCHE equivalent diameter in m."""
    area = pche_channel_flow_area(geometry)
    perimeter = pche_channel_wetted_perimeter(geometry)
    return 4.0 * area / perimeter


def flow_velocity(
    mass_flow_rate_kg_s,
    density_kg_m3,
    flow_area_m2,
    parallel_channels=1,
):
    """Return average channel velocity in m/s."""
    _require_positive("mass_flow_rate_kg_s", mass_flow_rate_kg_s)
    _require_positive("density_kg_m3", density_kg_m3)
    _require_positive("flow_area_m2", flow_area_m2)
    if parallel_channels < 1:
        raise ValueError("parallel_channels must be at least 1")

    channel_mass_flow = mass_flow_rate_kg_s / parallel_channels
    return channel_mass_flow / (density_kg_m3 * flow_area_m2)


def reynolds_number(
    density_kg_m3,
    velocity_m_s,
    hydraulic_diameter_m,
    viscosity_pa_s,
):
    """Return Reynolds number."""
    _require_positive("density_kg_m3", density_kg_m3)
    _require_positive("velocity_m_s", velocity_m_s)
    _require_positive("hydraulic_diameter_m", hydraulic_diameter_m)
    _require_positive("viscosity_pa_s", viscosity_pa_s)
    return density_kg_m3 * velocity_m_s * hydraulic_diameter_m / viscosity_pa_s


def prandtl_number(
    specific_heat_j_kg_k,
    viscosity_pa_s,
    thermal_conductivity_w_m_k,
):
    """Return Prandtl number."""
    _require_positive("specific_heat_j_kg_k", specific_heat_j_kg_k)
    _require_positive("viscosity_pa_s", viscosity_pa_s)
    _require_positive(
        "thermal_conductivity_w_m_k",
        thermal_conductivity_w_m_k,
    )
    return specific_heat_j_kg_k * viscosity_pa_s / thermal_conductivity_w_m_k


def gnielinski_f_star(reynolds):
    """Return the turbulent f* term used in Yuan et al.'s Nu correlation."""
    _require_positive("reynolds", reynolds)
    denominator = 1.82 * math.log10(reynolds) - 1.64
    if denominator == 0:
        raise ValueError("reynolds gives a zero denominator in f* correlation")
    return 1.0 / denominator**2


def nusselt_number(reynolds, prandtl, transition_re=2300.0):
    """Return Nusselt number using Yuan et al.'s laminar/turbulent formulas."""
    _require_positive("reynolds", reynolds)
    _require_positive("prandtl", prandtl)
    _require_positive("transition_re", transition_re)

    if reynolds < transition_re:
        return 4.36

    f_star = gnielinski_f_star(reynolds)
    numerator = (f_star / 8.0) * (reynolds - 1000.0) * prandtl
    denominator = 1.0 + 12.7 * math.sqrt(f_star / 8.0) * (
        prandtl ** (2.0 / 3.0) - 1.0
    )
    if denominator == 0:
        raise ValueError("Nusselt correlation denominator is zero")
    return numerator / denominator


def heat_transfer_coefficient(
    nusselt,
    thermal_conductivity_w_m_k,
    hydraulic_diameter_m,
):
    """Return convective heat-transfer coefficient in W/(m2*K)."""
    _require_positive("nusselt", nusselt)
    _require_positive(
        "thermal_conductivity_w_m_k",
        thermal_conductivity_w_m_k,
    )
    _require_positive("hydraulic_diameter_m", hydraulic_diameter_m)
    return nusselt * thermal_conductivity_w_m_k / hydraulic_diameter_m


def overall_heat_transfer_coefficient(
    hot_h_w_m2_k,
    cold_h_w_m2_k,
    wall_thickness_m,
    wall_thermal_conductivity_w_m_k,
):
    """Return total heat-transfer coefficient K in W/(m2*K)."""
    _require_positive("hot_h_w_m2_k", hot_h_w_m2_k)
    _require_positive("cold_h_w_m2_k", cold_h_w_m2_k)
    _require_positive("wall_thickness_m", wall_thickness_m)
    _require_positive(
        "wall_thermal_conductivity_w_m_k",
        wall_thermal_conductivity_w_m_k,
    )

    resistance = (
        1.0 / hot_h_w_m2_k
        + 1.0 / cold_h_w_m2_k
        + wall_thickness_m / wall_thermal_conductivity_w_m_k
    )
    return 1.0 / resistance


def friction_factor(reynolds, transition_re=2300.0):
    """Return Darcy friction factor from Yuan et al.'s piecewise relation."""
    _require_positive("reynolds", reynolds)
    _require_positive("transition_re", transition_re)
    if reynolds < transition_re:
        return 64.0 / reynolds
    return 0.3164 / reynolds**0.25


def pressure_drop_pa(
    friction_factor_value,
    length_m,
    hydraulic_diameter_m,
    density_kg_m3,
    velocity_m_s,
):
    """Return channel pressure drop in Pa."""
    _require_positive("friction_factor_value", friction_factor_value)
    _require_positive("length_m", length_m)
    _require_positive("hydraulic_diameter_m", hydraulic_diameter_m)
    _require_positive("density_kg_m3", density_kg_m3)
    _require_positive("velocity_m_s", velocity_m_s)

    return (
        friction_factor_value
        * length_m
        / hydraulic_diameter_m
        * density_kg_m3
        * velocity_m_s**2
        / 2.0
    )


def calculate_turbine_mass(component, Pe=None, A=0.5, **kwargs):
    """Estimate turbine/TAC mass with the empirical TAC formula.

    Formula from the provided reference:
        M = A * sqrt(0.5*pi*(30.522*ln(Pe) - 5.7178))

    If Pe is not supplied, the current turbine specific work ``Wt`` is used
    directly as Pe without unit conversion.
    """
    if not 0.4 <= A <= 0.8:
        raise ValueError("A should be within the recommended range 0.4-0.8")

    if Pe is None:
        Pe = abs(component.variables["Wt"])

    if Pe <= 0:
        raise ValueError("Pe must be positive")

    fit_term = 30.522 * math.log(Pe) - 5.7178
    if fit_term <= 0:
        raise ValueError(
            "Pe is outside the valid logarithmic domain for this "
            "empirical mass formula"
        )

    mass = A * math.sqrt(0.5 * math.pi * fit_term)
    return {
        "mass": mass,
        "A": A,
        "Pe": Pe,
        "fit_term": fit_term,
        "formula": "A*sqrt(0.5*pi*(30.522*ln(Pe)-5.7178))",
    }


def calculate_heater_mass(
    component,
    P_heat_mwt=None,
    mass_flow_rate=None,
    shielding_mass_ton=2.8,
    include_shielding=True,
    **kwargs,
):
    """Estimate reactor and shielding mass for the heater module.

    Reactor empirical formula from the provided reference:
        M_reactor = 0.2195 * P_heat + 0.09836

    P_heat is the reactor thermal power in MWt, and masses are in tons.
    The shielding mass is added as a constant 2.8 ton by default.

    If P_heat_mwt is not supplied, it is estimated from Q_in and
    mass_flow_rate, assuming Q_in is kJ/kg and mass_flow_rate is kg/s:
        P_heat_mwt = abs(Q_in) * mass_flow_rate / 1000
    """
    if P_heat_mwt is None:
        if mass_flow_rate is None:
            raise ValueError(
                "Heater mass calculation requires P_heat_mwt, or "
                "mass_flow_rate to estimate P_heat_mwt from Q_in."
            )
        P_heat_mwt = abs(component.variables["Q_in"]) * mass_flow_rate / 1000

    if P_heat_mwt <= 0:
        raise ValueError("P_heat_mwt must be positive")
    if shielding_mass_ton < 0:
        raise ValueError("shielding_mass_ton must be non-negative")

    reactor_mass_ton = 0.2195 * P_heat_mwt + 0.09836
    shielding_mass = shielding_mass_ton if include_shielding else 0.0
    total_mass_ton = reactor_mass_ton + shielding_mass

    return {
        "mass": total_mass_ton,
        "reactor_mass_ton": reactor_mass_ton,
        "shielding_mass_ton": shielding_mass,
        "P_heat_mwt": P_heat_mwt,
        "include_shielding": include_shielding,
        "formula": "0.2195*P_heat+0.09836+shielding_mass",
    }


def calculate_condenser_mass(
    component,
    Qc_kw=None,
    co2_mass_flow_rate=None,
    coolant_outlet_T=None,
    water_inlet_T=None,
    water_mass_flow_rate=None,
    water_pressure_kpa=101.325,
    refprop_path="C:/Program Files (x86)/REFPROP 10.0+/REFPROP",
    emissivity=0.92,
    surface_temperature=210.15,
    area_density=6.75,
    **kwargs,
):
    """Estimate radiator mass for the condenser module.

    Radiator heat rejection model:
        Qc = phi * sigma * A_rad * (T**4 - T0**4)
        M_rad = kappa * A_rad

    The NaK coolant in the reference is represented here by water. If
    coolant_outlet_T is not supplied, water outlet temperature is evaluated
    with REFPROP from water_inlet_T, water_mass_flow_rate, and Qc_kw.
    """
    if Qc_kw is None:
        if co2_mass_flow_rate is None:
            raise ValueError(
                "Condenser mass calculation requires Qc_kw, or "
                "co2_mass_flow_rate to estimate Qc_kw from Q_out."
            )
        Qc_kw = abs(component.variables["Q_out"]) * co2_mass_flow_rate

    if Qc_kw <= 0:
        raise ValueError("Qc_kw must be positive")
    if emissivity <= 0:
        raise ValueError("emissivity must be positive")
    if area_density <= 0:
        raise ValueError("area_density must be positive")

    if coolant_outlet_T is None:
        coolant_outlet_T = calculate_water_outlet_temperature(
            Qc_kw=Qc_kw,
            water_inlet_T=water_inlet_T,
            water_mass_flow_rate=water_mass_flow_rate,
            water_pressure_kpa=water_pressure_kpa,
            refprop_path=refprop_path,
        )

    temperature_term = coolant_outlet_T**4 - surface_temperature**4
    if temperature_term <= 0:
        raise ValueError(
            "coolant_outlet_T must be higher than surface_temperature for "
            "radiative heat rejection"
        )

    stefan_boltzmann = 5.670374419e-8
    Qc_w = Qc_kw * 1000
    area_m2 = Qc_w / (emissivity * stefan_boltzmann * temperature_term)
    mass_kg = area_density * area_m2

    return {
        "mass": mass_kg,
        "radiator_area_m2": area_m2,
        "Qc_kw": Qc_kw,
        "coolant_outlet_T": coolant_outlet_T,
        "surface_temperature": surface_temperature,
        "emissivity": emissivity,
        "area_density": area_density,
        "mass_unit": "kg",
        "formula": "M_rad=kappa*Qc/(phi*sigma*(T^4-T0^4))",
    }


def calculate_water_outlet_temperature(
    Qc_kw,
    water_inlet_T,
    water_mass_flow_rate,
    water_pressure_kpa,
    refprop_path,
):
    """Calculate water outlet temperature for the condenser mass model."""
    import ctREFPROP.ctREFPROP as ct

    if water_inlet_T is None or water_mass_flow_rate is None:
        raise ValueError(
            "Provide coolant_outlet_T directly, or provide water_inlet_T "
            "and water_mass_flow_rate for REFPROP water calculation."
        )
    if water_mass_flow_rate <= 0:
        raise ValueError("water_mass_flow_rate must be positive")

    water = ct.REFPROPFunctionLibrary(refprop_path)
    water.SETUPdll(1, "WATER.FLD", "HMX.BNC", "DEF")
    water.SETUPdll(2, "SI", "SI", "DEF")
    z = [1.0]
    mw = water.WMOLdll(z)

    inlet = water.TPFLSHdll(water_inlet_T, water_pressure_kpa, z)
    if inlet.ierr > 0:
        raise ValueError(f"REFPROP water inlet calculation error: {inlet.ierr}")

    h_in_mass = inlet.h / mw
    h_out_mass = h_in_mass + Qc_kw / water_mass_flow_rate
    outlet = water.PHFLSHdll(water_pressure_kpa, h_out_mass * mw, z)
    if outlet.ierr > 0:
        raise ValueError(f"REFPROP water outlet calculation error: {outlet.ierr}")

    return outlet.T