from __future__ import annotations

import math
from dataclasses import dataclass

from .core import HydrologyError


def _positive(value: float, label: str) -> float:
    value = float(value)
    if not math.isfinite(value) or value <= 0:
        raise HydrologyError(f"{label}必须大于0")
    return value


def _non_negative(value: float, label: str) -> float:
    value = float(value)
    if not math.isfinite(value) or value < 0:
        raise HydrologyError(f"{label}不能小于0")
    return value


@dataclass(frozen=True)
class HydropowerOutputInput:
    design_flow: float
    gross_head: float
    head_loss: float
    turbine_efficiency_percent: float
    output_coefficient: float


@dataclass(frozen=True)
class HydropowerOutputResult:
    working_head: float
    station_output_kw: float
    turbine_input_kw: float
    turbine_output_kw: float
    formula: str


def calculate_hydropower_output(data: HydropowerOutputInput) -> HydropowerOutputResult:
    flow = _positive(data.design_flow, "引水流量Q")
    gross_head = _positive(data.gross_head, "总水头")
    head_loss = _non_negative(data.head_loss, "水头损失")
    if gross_head <= head_loss:
        raise HydrologyError("总水头必须大于水头损失")
    efficiency = _positive(data.turbine_efficiency_percent, "水轮机效率") / 100.0
    coefficient = _positive(data.output_coefficient, "出力系数K")
    working_head = gross_head - head_loss
    turbine_input = 9.81 * flow * working_head
    turbine_output = turbine_input * efficiency
    station_output = coefficient * flow * working_head
    return HydropowerOutputResult(
        working_head,
        station_output,
        turbine_input,
        turbine_output,
        "N=KQH；Nω=9.81QH；N=9.81QHη",
    )


@dataclass(frozen=True)
class InstalledCapacityInput:
    guarantee_flow: float
    working_head: float
    output_coefficient: float
    ratio_coefficient: float
    average_annual_generation_10k_kwh: float
    annual_utilization_hours: float


@dataclass(frozen=True)
class InstalledCapacityResult:
    guarantee_output_kw: float
    capacity_by_ratio_kw: float
    capacity_by_utilization_hours_kw: float
    formula: str


def calculate_installed_capacity(data: InstalledCapacityInput) -> InstalledCapacityResult:
    flow = _positive(data.guarantee_flow, "设计保证率流量Q")
    head = _positive(data.working_head, "工作水头H")
    coefficient = _positive(data.output_coefficient, "出力系数k")
    ratio = _positive(data.ratio_coefficient, "经验系数a")
    generation = _positive(data.average_annual_generation_10k_kwh, "多年平均发电量E年")
    utilization_hours = _positive(data.annual_utilization_hours, "装机容量年利用小时数T装")
    guarantee_output = coefficient * flow * head
    capacity_by_ratio = ratio * guarantee_output
    capacity_by_hours = generation * 10000.0 / utilization_hours
    return InstalledCapacityResult(
        guarantee_output,
        capacity_by_ratio,
        capacity_by_hours,
        "Np=kQH；Ny=aNp；Ny=E年/T装",
    )


@dataclass(frozen=True)
class CandidateEconomicInput:
    capacities_kw: list[float]
    investments_10k_yuan: list[float]
    annual_generations_10k_kwh: list[float]
    electricity_price_yuan_kwh: float
    operation_cost_rate_percent: float
    design_load_kw: float


@dataclass(frozen=True)
class CandidateEconomicRow:
    capacity_kw: float
    investment_10k_yuan: float
    annual_generation_10k_kwh: float
    gross_revenue_10k_yuan: float
    operation_cost_10k_yuan: float
    net_revenue_10k_yuan: float
    payback_years: float
    meets_load: bool


@dataclass(frozen=True)
class CandidateEconomicResult:
    rows: list[CandidateEconomicRow]
    selected_capacity_kw: float | None
    formula: str


def calculate_candidate_economics(data: CandidateEconomicInput) -> CandidateEconomicResult:
    if not (len(data.capacities_kw) == len(data.investments_10k_yuan) == len(data.annual_generations_10k_kwh)):
        raise HydrologyError("备选装机容量、投资和多年平均发电量数量必须一致")
    price = _positive(data.electricity_price_yuan_kwh, "单位电价")
    op_rate = _non_negative(data.operation_cost_rate_percent, "年运行费率") / 100.0
    design_load = _positive(data.design_load_kw, "设计水平年负荷")
    rows: list[CandidateEconomicRow] = []
    for capacity, investment, generation in zip(data.capacities_kw, data.investments_10k_yuan, data.annual_generations_10k_kwh):
        capacity = _positive(capacity, "备选装机容量")
        investment = _positive(investment, "建站估算总投资")
        generation = _positive(generation, "多年平均发电量")
        gross = generation * price
        operation = investment * op_rate
        net = gross - operation
        payback = math.inf if net <= 0 else (investment + operation) / net
        rows.append(
            CandidateEconomicRow(
                capacity,
                investment,
                generation,
                gross,
                operation,
                net,
                payback,
                capacity >= design_load,
            )
        )
    feasible = [row for row in rows if row.meets_load and row.net_revenue_10k_yuan > 0]
    selected = max(feasible, key=lambda row: (row.net_revenue_10k_yuan, -row.payback_years)).capacity_kw if feasible else None
    return CandidateEconomicResult(rows, selected, "k1=m1E年；k2=m2M；k=k1-k2；还本年限=(M+k2)/k")


@dataclass(frozen=True)
class PressurePipeHydropowerInput:
    design_flow: float
    max_design_flow: float
    design_velocity: float
    design_head: float
    max_head: float
    coefficient_k: float
    weld_coefficient: float
    allowable_stress_kg_cm2: float


@dataclass(frozen=True)
class PressurePipeHydropowerResult:
    diameter_by_beyer_m: float
    diameter_by_velocity_m: float
    selected_diameter_m: float
    strength_thickness_mm: float
    minimum_thickness_mm: float
    calculated_thickness_mm: float
    design_thickness_mm: float
    formula: str


def calculate_pressure_pipe_hydropower(data: PressurePipeHydropowerInput) -> PressurePipeHydropowerResult:
    flow = _positive(data.design_flow, "管道设计流量Q")
    max_flow = _positive(data.max_design_flow, "管道最大设计流量Qmax")
    velocity = _positive(data.design_velocity, "管道经济流速v")
    head = _positive(data.design_head, "设计水头H")
    max_head = _positive(data.max_head, "最大设计水头")
    coefficient = _positive(data.coefficient_k, "系数k")
    weld = _positive(data.weld_coefficient, "钢管焊缝系数φ")
    stress = _positive(data.allowable_stress_kg_cm2, "钢板允许应力[σ]")
    diameter_beyer = (coefficient * max_flow**3 / head) ** (1.0 / 7.0)
    diameter_velocity = math.sqrt(4.0 * flow / (math.pi * velocity))
    selected = diameter_velocity
    strength_thickness = 100.0 * selected * max_head / (2.0 * weld * stress)
    minimum_thickness = selected * 1000.0 / 130.0
    calculated = max(strength_thickness, minimum_thickness)
    design = max(calculated + 2.0, 6.0)
    return PressurePipeHydropowerResult(
        diameter_beyer,
        diameter_velocity,
        selected,
        strength_thickness,
        minimum_thickness,
        calculated,
        design,
        "D=(kQmax^3/H)^(1/7)；D=sqrt(4Q/(πv))；δ=100DH/(2φ[σ])；δ≥D/130；δ设=max(δ计+2,6)",
    )


@dataclass(frozen=True)
class PipeSubmergenceInput:
    pipe_diameter: float
    flow: float
    upstream_level: float
    downstream_level: float
    axis_angle_deg: float
    is_round_pipe: bool = True
    square_side_length: float | None = None


@dataclass(frozen=True)
class PipeSubmergenceResult:
    unit_width_flow: float
    water_level_difference: float
    minimum_submergence_depth: float
    formula: str


def calculate_pipe_minimum_submergence(data: PipeSubmergenceInput) -> PipeSubmergenceResult:
    diameter = _positive(data.pipe_diameter, "管道内径D")
    flow = _positive(data.flow, "管中流量Q")
    upstream = float(data.upstream_level)
    downstream = float(data.downstream_level)
    angle = _positive(data.axis_angle_deg, "管道轴线与水平线夹角δ")
    width = diameter if data.is_round_pipe else _positive(data.square_side_length or 0.0, "方管边长b")
    unit_flow = flow / width
    z = abs(upstream - downstream)
    h_min = 1.414 * diameter * math.sin(math.radians(angle)) * (z / (unit_flow ** (2.0 / 3.0))) ** 0.31
    return PipeSubmergenceResult(unit_flow, z, h_min, "Hmin=1.414D*sinδ*(Z/q^(2/3))^0.31")


@dataclass(frozen=True)
class VentPipeInput:
    air_flow: float
    flow_coefficient: float
    elastic_modulus_mpa: float
    safety_factor: float
    wall_thickness: float
    pipe_diameter: float


@dataclass(frozen=True)
class VentPipeResult:
    allowable_negative_pressure_mpa: float
    vent_area_m2: float
    vent_diameter_m: float
    formula: str


def calculate_vent_pipe(data: VentPipeInput) -> VentPipeResult:
    air_flow = _positive(data.air_flow, "通气量Q0")
    coefficient = _positive(data.flow_coefficient, "通气孔流量系数C")
    elastic = _positive(data.elastic_modulus_mpa, "管壁材料弹性模数E")
    safety = _positive(data.safety_factor, "管道安全系数K")
    thickness = _positive(data.wall_thickness, "管壁厚度t")
    diameter = _positive(data.pipe_diameter, "管道内径D0")
    negative_pressure = 2.0 * elastic / safety * (thickness / diameter) ** 3
    area = air_flow / (1265.0 * coefficient * math.sqrt(negative_pressure))
    vent_diameter = math.sqrt(4.0 * area / math.pi)
    return VentPipeResult(negative_pressure, area, vent_diameter, "F=Q0/(1265C√ΔP)；ΔP=2E/K*(t/D0)^3")


@dataclass(frozen=True)
class WaterHammerInput:
    pipe_diameter: float
    wall_thickness: float
    pipe_material_elastic_mpa: float
    water_bulk_modulus_mpa: float
    pipe_length: float
    closing_time: float
    initial_velocity: float
    final_velocity: float
    steady_head: float
    water_density_t_m3: float = 1.0
    gravity: float = 9.81


@dataclass(frozen=True)
class WaterHammerResult:
    wave_velocity: float
    reflection_time: float
    hammer_type: str
    direct_head: float
    sigma: float
    relative_head: float
    indirect_head: float
    maximum_head: float
    maximum_pressure_kn_m2: float
    allowable_head: float
    formula: str


def _allowed_hammer_percent(head: float) -> float:
    if head < 40.0:
        return 70.0
    if head < 100.0:
        return 50.0
    return 25.0


def calculate_water_hammer(data: WaterHammerInput) -> WaterHammerResult:
    diameter = _positive(data.pipe_diameter, "管道内径D")
    thickness = _positive(data.wall_thickness, "管壁厚度δ")
    pipe_elastic = _positive(data.pipe_material_elastic_mpa, "管壁材料纵向弹性模数E")
    water_bulk = _positive(data.water_bulk_modulus_mpa, "水的体积弹性模数K")
    length = _positive(data.pipe_length, "管道长度L")
    closing_time = _positive(data.closing_time, "阀门或导叶启闭时间Ts")
    initial_velocity = _non_negative(data.initial_velocity, "启闭前流速V0")
    final_velocity = _non_negative(data.final_velocity, "启闭后流速V")
    head = _positive(data.steady_head, "恒定水头H0")
    gravity = _positive(data.gravity, "重力加速度g")
    density = _positive(data.water_density_t_m3, "水密度ρ")
    wave_velocity = 1435.0 / math.sqrt(1.0 + diameter * water_bulk / (thickness * pipe_elastic))
    reflection_time = 2.0 * length / wave_velocity
    hammer_type = "直接水锤" if reflection_time >= closing_time else "间接水锤"
    direct_head = abs(wave_velocity * (final_velocity - initial_velocity) / gravity)
    sigma = initial_velocity * length / (gravity * head * closing_time)
    relative_head = sigma / 2.0 * (math.sqrt(sigma * sigma + 4.0) + sigma)
    indirect_head = relative_head * head
    maximum_head = head + indirect_head
    maximum_pressure = maximum_head * density * gravity
    allowable_head = head * _allowed_hammer_percent(head) / 100.0
    return WaterHammerResult(
        wave_velocity,
        reflection_time,
        hammer_type,
        direct_head,
        sigma,
        relative_head,
        indirect_head,
        maximum_head,
        maximum_pressure,
        allowable_head,
        "a=1435/sqrt(1+DK/(δE))；TR=2L/a；直接ΔH=-a(V-V0)/g；间接ζ=σ/2(√(σ²+4)+σ)",
    )


@dataclass(frozen=True)
class PipeHeadLossInput:
    velocity: float
    diameter: float
    length: float
    roughness_mm: float
    local_loss_coefficient: float
    kinematic_viscosity: float = 1.519e-6
    gravity: float = 9.81


@dataclass(frozen=True)
class PipeHeadLossResult:
    reynolds_number: float
    friction_factor: float
    friction_loss: float
    local_loss: float
    total_loss: float
    formula: str


def _colebrook_friction_factor(reynolds: float, diameter: float, roughness_m: float) -> float:
    if reynolds <= 2000.0:
        return 64.0 / reynolds
    if reynolds < 4000.0:
        # Explicit transition approximation between laminar and turbulent limits.
        turbulent = _colebrook_friction_factor(4000.0, diameter, roughness_m)
        return .032 + (turbulent - .032) * (reynolds - 2000.0) / 2000.0
    factor = 0.02
    for _ in range(30):
        right = -2.0 * math.log10(2.51 / (reynolds * math.sqrt(factor)) + roughness_m / (3.7 * diameter))
        new_factor = 1.0 / (right * right)
        if abs(new_factor - factor) < 1e-9:
            return new_factor
        factor = new_factor
    return factor


def calculate_pipe_head_loss(data: PipeHeadLossInput) -> PipeHeadLossResult:
    velocity = _positive(data.velocity, "管中流速v")
    diameter = _positive(data.diameter, "管道直径d")
    length = _positive(data.length, "管道长度L")
    roughness_m = _non_negative(data.roughness_mm, "管壁粗糙度Δ") / 1000.0
    zeta = _non_negative(data.local_loss_coefficient, "局部水头损失系数Σζ")
    viscosity = _positive(data.kinematic_viscosity, "水的运动黏滞系数γ")
    gravity = _positive(data.gravity, "重力加速度g")
    reynolds = velocity * diameter / viscosity
    factor = _colebrook_friction_factor(reynolds, diameter, roughness_m)
    velocity_head = velocity * velocity / (2.0 * gravity)
    friction = factor * length * velocity_head / diameter
    local = zeta * velocity_head
    return PipeHeadLossResult(reynolds, factor, friction, local, friction + local, "h=λLv²/(2gd)+Σζv²/(2g)；λ为Darcy系数：Re≤2000采用64/Re，Re≥4000迭代Colebrook，过渡区线性插值需复核；内部计算不提前舍入。")


@dataclass(frozen=True)
class ExternalPressureInput:
    wall_thickness_m: float
    inner_diameter_m: float
    steel_yield_strength_mpa: float
    steel_elastic_modulus_mpa: float


@dataclass(frozen=True)
class ExternalPressureResult:
    exposed_pipe_kpa: float
    buried_pipe_mpa: float
    formula: str


def calculate_external_pressure(data: ExternalPressureInput) -> ExternalPressureResult:
    thickness = _positive(data.wall_thickness_m, "管壁计算厚度δ")
    diameter = _positive(data.inner_diameter_m, "钢管内径D0")
    yield_strength = _positive(data.steel_yield_strength_mpa, "钢材屈服强度σs")
    elastic = _positive(data.steel_elastic_modulus_mpa, "钢材弹模E")
    exposed_mpa = 2.0 * elastic * (thickness / diameter) ** 3
    thickness_cm = thickness * 100.0
    diameter_cm = diameter * 100.0
    buried_mpa = 3440.0 * (thickness_cm / diameter_cm) ** 1.7 * yield_strength**0.25
    return ExternalPressureResult(exposed_mpa * 1000.0, buried_mpa, "明管pcr=2E(δ/D0)^3；埋管pcr=3440(δ/r1)^1.7σs^0.25")


CONDUCTOR_ALLOWABLE_CURRENT_A: dict[str, dict[int, float]] = {
    "LGJ-10": {80: 86},
    "LGJ-16": {70: 105, 80: 108},
    "LGJ-25": {70: 130, 80: 138},
    "LGJ-35": {70: 175, 80: 183},
    "LGJ-50": {70: 210, 80: 215},
    "LGJ-70": {70: 265, 80: 260},
    "LGJ-95": {70: 330, 80: 352},
    "LGJ-120": {70: 380, 80: 401},
    "LGJ-150": {70: 445, 80: 452},
    "LGJ-185": {70: 510, 80: 531},
    "LGJ-240": {70: 610, 80: 613},
    "LGJ-300": {70: 690, 80: 755},
    "LGJ-400": {70: 835, 80: 840},
}


@dataclass(frozen=True)
class TransmissionConductorInput:
    max_load_kw: float
    rated_voltage_kv: float
    economic_current_density_a_mm2: float
    power_factor: float
    allowable_temperature_c: int = 80


@dataclass(frozen=True)
class TransmissionConductorResult:
    economic_section_mm2: float
    recommended_conductor: str | None
    sustained_current_a: float
    selected_allowable_current_a: float | None
    is_selected_sufficient: bool | None
    formula: str


def calculate_transmission_conductor(data: TransmissionConductorInput) -> TransmissionConductorResult:
    load = _positive(data.max_load_kw, "输送电容量P")
    voltage = _positive(data.rated_voltage_kv, "线路额定电压Ue")
    current_density = _positive(data.economic_current_density_a_mm2, "经济电流密度J")
    power_factor = _positive(data.power_factor, "功率因数cosφ")
    if power_factor > 1:
        raise HydrologyError("功率因数cosφ不能大于1")
    section = load / (math.sqrt(3.0) * voltage * current_density * power_factor)
    sustained_current = load / (math.sqrt(3.0) * voltage * power_factor)
    recommended: str | None = None
    allowable: float | None = None
    for name in sorted(CONDUCTOR_ALLOWABLE_CURRENT_A, key=lambda item: int(item.split("-")[1])):
        area = int(name.split("-")[1])
        if area >= section:
            recommended = name
            allowable = CONDUCTOR_ALLOWABLE_CURRENT_A[name].get(data.allowable_temperature_c)
            break
    sufficient = None if allowable is None else allowable >= sustained_current
    return TransmissionConductorResult(section, recommended, sustained_current, allowable, sufficient, "S=P/(√3JUe cosφ)；I=P/(√3Ue cosφ)")


@dataclass(frozen=True)
class TransformerCapacityInput:
    generator_total_active_power_kw: float
    bus_minimum_load_kw: float
    rated_power_factor: float
    standard_capacity_sum_kva: float


@dataclass(frozen=True)
class TransformerCapacityResult:
    required_capacity_kva: float
    is_standard_sufficient: bool
    formula: str


def calculate_transformer_capacity(data: TransformerCapacityInput) -> TransformerCapacityResult:
    total_power = _positive(data.generator_total_active_power_kw, "发电机组总有功功率ΣP")
    min_load = _non_negative(data.bus_minimum_load_kw, "母线侧最小负荷Pmin")
    power_factor = _positive(data.rated_power_factor, "额定功率因数cosφ")
    standard = _positive(data.standard_capacity_sum_kva, "主变压器标准容量总和Sn")
    if min_load >= total_power:
        raise HydrologyError("母线侧最小负荷必须小于发电机组总有功功率")
    required = (total_power - min_load) / power_factor
    return TransformerCapacityResult(required, required <= standard, "S=(ΣP-Pmin)/cosφ≤Sn")


@dataclass(frozen=True)
class LoadForecastInput:
    installed_load_kw: float
    demand_factor: float
    max_load_utilization_hours: float
    base_load_kw: float
    annual_growth_rate_percent: float
    forecast_years: float
    grid_loss_rate_percent: float
    plant_use_rate_percent: float


@dataclass(frozen=True)
class LoadForecastResult:
    max_load_by_demand_kw: float
    annual_energy_by_demand_kwh: float
    load_by_growth_kw: float
    selected_user_max_load_kw: float
    system_supply_max_load_kw: float
    system_generation_max_load_kw: float
    formula: str


def calculate_load_forecast(data: LoadForecastInput) -> LoadForecastResult:
    installed_load = _positive(data.installed_load_kw, "电网总负荷装置功率ΣPz")
    demand_factor = _positive(data.demand_factor, "综合需用系数kzz")
    hours = _positive(data.max_load_utilization_hours, "年最大负荷利用小时数Tm")
    base_load = _positive(data.base_load_kw, "基准年负荷A1")
    growth = _non_negative(data.annual_growth_rate_percent, "年平均增长率i") / 100.0
    years = _positive(data.forecast_years, "设计水平年n")
    grid_loss = _non_negative(data.grid_loss_rate_percent, "网损率") / 100.0
    plant_use = _non_negative(data.plant_use_rate_percent, "厂用电率") / 100.0
    if grid_loss >= 1 or plant_use >= 1:
        raise HydrologyError("网损率和厂用电率必须小于100%")
    demand_load = demand_factor * installed_load
    annual_energy = demand_load * hours
    growth_load = base_load * (1.0 + growth) ** years
    user_load = max(demand_load, growth_load)
    supply_load = user_load / (1.0 - grid_loss)
    generation_load = supply_load / (1.0 - plant_use)
    return LoadForecastResult(
        demand_load,
        annual_energy,
        growth_load,
        user_load,
        supply_load,
        generation_load,
        "Pm=kzzΣPz；E=PmTm；An=A1(1+i)^n；供电最大负荷=用电最大负荷/(1-网损率)；发电最大负荷=供电最大负荷/(1-厂用电率)",
    )
