from __future__ import annotations

from dataclasses import dataclass
import math

from .core import HydrologyError


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


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


@dataclass(frozen=True)
class ReservoirFloodInput:
    catchment_area_km2: float
    rainfall_decay_exponent: float
    peak_flow_coefficient: float
    design_24h_rainfall_mm: float
    net_rainfall_mm: float


@dataclass(frozen=True)
class ReservoirFloodResult:
    rainfall_intensity_mm_h: float
    peak_flow_m3s: float
    flood_volume_m3: float
    formula: str


def calculate_reservoir_flood(data: ReservoirFloodInput) -> ReservoirFloodResult:
    area = _positive(data.catchment_area_km2, "汇水面积F")
    exponent = _positive(data.rainfall_decay_exponent, "暴雨递减指数n")
    coefficient = _positive(data.peak_flow_coefficient, "洪峰流量参数K")
    rain24 = _positive(data.design_24h_rainfall_mm, "设计频率24h降雨量H24P")
    net_rain = _non_negative(data.net_rainfall_mm, "一次净雨量h")
    intensity = 24.0 ** (exponent - 1.0) * rain24
    peak = coefficient * intensity * area ** (2.0 / 3.0)
    volume = 1000.0 * net_rain * area
    return ReservoirFloodResult(intensity, peak, volume, "Q=KSF^(2/3)；S=24^(n-1)H24P；W=1000hF")


@dataclass(frozen=True)
class MountainFloodInput:
    catchment_area_km2: float
    rainfall_intensity_mm_h: float
    rainfall_decay_exponent: float
    runoff_modulus: float
    peak_flow_coefficient: float


@dataclass(frozen=True)
class MountainFloodResult:
    no_rainfall_data_peak_m3s: float
    rainfall_small_area_peak_m3s: float
    rainfall_large_area_peak_m3s: float
    formula: str


def calculate_mountain_flood_peak(data: MountainFloodInput) -> MountainFloodResult:
    area = _positive(data.catchment_area_km2, "汇水面积F")
    intensity = _positive(data.rainfall_intensity_mm_h, "暴雨力S")
    exponent = _positive(data.rainfall_decay_exponent, "暴雨递减指数n")
    runoff_modulus = _positive(data.runoff_modulus, "径流模数K")
    coefficient = _positive(data.peak_flow_coefficient, "洪峰流量参数C")
    return MountainFloodResult(
        runoff_modulus * area**exponent,
        coefficient * intensity * area,
        coefficient * intensity * area ** (2.0 / 3.0),
        "无降雨资料Q=KF^n；有降雨资料F<3km2时Q=CSF；F≥3km2时Q=CSF^(2/3)",
    )


@dataclass(frozen=True)
class ReservoirSedimentationInput:
    catchment_area_km2: float
    erosion_modulus_t_km2: float
    service_life_years: float
    dry_unit_weight_t_m3: float


@dataclass(frozen=True)
class ReservoirSedimentationResult:
    annual_sediment_t: float
    sediment_capacity_m3: float
    formula: str


def calculate_reservoir_sedimentation(data: ReservoirSedimentationInput) -> ReservoirSedimentationResult:
    area = _positive(data.catchment_area_km2, "流域面积F")
    modulus = _positive(data.erosion_modulus_t_km2, "侵蚀模数M")
    years = _positive(data.service_life_years, "水库使用年限T")
    dry_weight = _positive(data.dry_unit_weight_t_m3, "淤沙体干容重γ")
    annual = modulus * area
    capacity = annual * years / dry_weight
    return ReservoirSedimentationResult(annual, capacity, "W=MF；V=WT/γ")


@dataclass(frozen=True)
class MonthlyReservoirOperationRow:
    month: str
    inflow: float
    demand: float
    surplus: float
    deficit: float
    storage: float
    abandoned: float
    note: str


@dataclass(frozen=True)
class ReservoirAnnualRegulationResult:
    capacity_monthly_flow: float
    rows: list[MonthlyReservoirOperationRow]
    total_surplus: float
    total_deficit: float
    total_abandoned: float
    formula: str


def calculate_reservoir_annual_regulation(
    months: list[str],
    inflows: list[float],
    demands: list[float],
    effective_capacity_m3: float,
    seconds_per_month: float = 30.0 * 24.0 * 3600.0,
) -> ReservoirAnnualRegulationResult:
    if not (len(months) == len(inflows) == len(demands)):
        raise HydrologyError("月份、来水流量和用水流量数量必须一致")
    capacity = _positive(effective_capacity_m3, "水库有效库容") / _positive(seconds_per_month, "月换算秒数")
    storage = 0.0
    rows: list[MonthlyReservoirOperationRow] = []
    total_surplus = total_deficit = total_abandoned = 0.0
    for month, inflow, demand in zip(months, inflows, demands):
        inflow = _non_negative(inflow, "来水流量")
        demand = _non_negative(demand, "用水流量")
        surplus = max(inflow - demand, 0.0)
        deficit = max(demand - inflow, 0.0)
        abandoned = 0.0
        note = ""
        if surplus > 0:
            storage += surplus
            if storage > capacity:
                abandoned = storage - capacity
                storage = capacity
                note = "库满弃水期"
            else:
                note = "空库蓄水期" if storage <= surplus + 1e-9 else "蓄水期"
        elif deficit > 0:
            storage -= deficit
            if storage < 0:
                deficit += abs(storage)
                storage = 0.0
                note = "缺水期"
            else:
                note = "供水期"
        rows.append(MonthlyReservoirOperationRow(month, inflow, demand, surplus, max(demand - inflow, 0.0), storage, abandoned, note))
        total_surplus += surplus
        total_deficit += max(demand - inflow, 0.0)
        total_abandoned += abandoned
    return ReservoirAnnualRegulationResult(
        capacity,
        rows,
        total_surplus,
        total_deficit,
        total_abandoned,
        "余水=来水-用水；亏水=用水-来水；蓄水量=上月蓄水+本月来水-本月用水-弃水",
    )


@dataclass(frozen=True)
class EarthDamFreeboardInput:
    water_depth_m: float
    wind_speed_m_s: float
    fetch_m: float
    wind_water_axis_angle_deg: float
    wave_period_s: float
    upstream_slope: float
    slope_permeability_coefficient: float
    empirical_coefficient: float
    safety_margin_m: float
    friction_coefficient: float = 3.6e-6
    gravity_m_s2: float = 9.81


@dataclass(frozen=True)
class EarthDamFreeboardResult:
    wave_height_m: float
    wave_length_m: float
    wave_runup_m: float
    wind_setup_m: float
    freeboard_m: float
    crest_elevation_m: float
    formula: str


def calculate_earth_dam_freeboard(data: EarthDamFreeboardInput) -> EarthDamFreeboardResult:
    depth = _positive(data.water_depth_m, "坝前水深")
    wind = _non_negative(data.wind_speed_m_s, "风速")
    fetch = _non_negative(data.fetch_m, "风区长度")
    period = _positive(data.wave_period_s, "平均波周期")
    slope = _positive(data.upstream_slope, "上游坝坡系数m")
    kd = _positive(data.slope_permeability_coefficient, "坝坡渗透性及糙率系数kΔ")
    kw = _positive(data.empirical_coefficient, "经验系数kW")
    safety = _non_negative(data.safety_margin_m, "安全加高A")
    friction = _positive(data.friction_coefficient, "综合摩阻系数k")
    gravity = _positive(data.gravity_m_s2, "重力加速度g")
    angle = math.radians(float(data.wind_water_axis_angle_deg))
    wave_height = (period / 4.438) ** 2
    wave_length = gravity * period**2 / (2.0 * math.pi)
    if slope <= 1.25:
        r0 = 1.24 if slope <= 0 else 1.45 if slope <= 0.5 else 2.2 if slope <= 1.0 else 2.5
        runup = kd * kw * r0 * wave_height
    else:
        runup = kd * kw / math.sqrt(1.0 + slope**2) * math.sqrt(wave_height * wave_length)
    wind_setup = friction * wind**2 * fetch * math.cos(angle) / (2.0 * gravity * depth)
    freeboard = runup + wind_setup + safety
    return EarthDamFreeboardResult(
        wave_height,
        wave_length,
        runup,
        wind_setup,
        freeboard,
        depth + freeboard,
        "Δh=Rm+e+A；Rm=kΔkW/sqrt(1+m^2)*sqrt(hmLm)或kΔkWR0hm；e=kV²Dcosβ/(2gHm)",
    )


@dataclass(frozen=True)
class EarthDamProfileInput:
    dam_body_height_m: float
    upstream_slope: float
    downstream_slope: float
    curtain_constant_m: float


@dataclass(frozen=True)
class EarthDamProfileResult:
    crest_width_m: float
    bottom_width_m: float
    curtain_depth_m: float
    formula: str


def calculate_earth_dam_profile(data: EarthDamProfileInput) -> EarthDamProfileResult:
    height = _positive(data.dam_body_height_m, "水库坝体高度H")
    upstream = _positive(data.upstream_slope, "上游坝坡坡率m上游")
    downstream = _positive(data.downstream_slope, "下游坝坡坡率m下游")
    curtain_constant = _non_negative(data.curtain_constant_m, "帷幕常数C")
    crest_width = math.sqrt(height)
    bottom_width = height * upstream + crest_width + height * downstream
    curtain_depth = height / 3.0 + curtain_constant
    return EarthDamProfileResult(crest_width, bottom_width, curtain_depth, "B=sqrt(H)；b=Hm上游+B+Hm下游；S=H/3+C")


@dataclass(frozen=True)
class DamSlopeProtectionInput:
    slope_coefficient: float
    slope_change_coefficient: float
    block_density_t_m3: float
    water_density_t_m3: float
    wind_speed_m_s: float
    fetch_m: float
    wave_period_s: float
    gravity_m_s2: float = 9.81


@dataclass(frozen=True)
class DamSlopeProtectionResult:
    wave_height_m: float
    wave_length_m: float
    wave_steepness_ratio: float
    equivalent_stone_diameter_m: float
    protection_thickness_m: float
    formula: str


def calculate_dam_slope_protection(data: DamSlopeProtectionInput) -> DamSlopeProtectionResult:
    slope = _positive(data.slope_coefficient, "坝坡系数m")
    kt = _positive(data.slope_change_coefficient, "随坡率变化系数Kt")
    block_density = _positive(data.block_density_t_m3, "块石密度ρk")
    water_density = _positive(data.water_density_t_m3, "水的密度ρw")
    if block_density <= water_density:
        raise HydrologyError("块石密度必须大于水的密度")
    wind = _positive(data.wind_speed_m_s, "最大风速W")
    fetch = _positive(data.fetch_m, "风区长度D")
    period = _positive(data.wave_period_s, "平均波周期Tm")
    gravity = _positive(data.gravity_m_s2, "重力加速度g")
    wave_height = wind**2 / gravity * 0.0076 * wind ** (-1.0 / 12.0) * (gravity * fetch / wind**2) ** (1.0 / 3.0)
    wave_length = gravity * period**2 / (2.0 * math.pi)
    ratio = wave_length / wave_height
    stone_diameter = (
        1.018
        * kt
        * (water_density / (block_density - water_density))
        * (math.sqrt(slope**2 + 1.0) / (slope * (slope + 2.0)))
        * wave_height
    )
    coefficient = 1.67 if ratio <= 15.0 else 1.82
    thickness = coefficient * stone_diameter / kt
    return DamSlopeProtectionResult(
        wave_height,
        wave_length,
        ratio,
        stone_diameter,
        thickness,
        "d=1.018Kt*ρw/(ρk-ρw)*sqrt(m²+1)/(m(m+2))*hp；Lm/hp≤15时t=1.67d/Kt，否则t=1.82d/Kt",
    )


@dataclass(frozen=True)
class ReservoirCapacityLayer:
    height_m: float
    upper_area_m2: float
    lower_area_m2: float


@dataclass(frozen=True)
class ReservoirCapacityInput:
    layers: list[ReservoirCapacityLayer]
    cone_height_m: float
    cone_top_area_m2: float


@dataclass(frozen=True)
class ReservoirCapacityResult:
    layer_volumes_m3: list[float]
    cone_volume_m3: float
    total_capacity_m3: float
    formula: str


def calculate_reservoir_capacity(data: ReservoirCapacityInput) -> ReservoirCapacityResult:
    if not data.layers:
        raise HydrologyError("至少需要输入一层柱体库容")
    volumes: list[float] = []
    for layer in data.layers:
        height = _non_negative(layer.height_m, "分层高度bi")
        upper = _non_negative(layer.upper_area_m2, "上面水面面积A1i")
        lower = _non_negative(layer.lower_area_m2, "下面水面面积A2i")
        volumes.append(0.5 * height * (upper + lower))
    cone_height = _non_negative(data.cone_height_m, "底部锥体高度b")
    cone_area = _non_negative(data.cone_top_area_m2, "锥体上部水面面积Az")
    cone_volume = cone_height * cone_area / 3.0
    total = sum(volumes) + cone_volume
    return ReservoirCapacityResult(volumes, cone_volume, total, "Vi=1/2*bi(A1i+A2i)；Vz=1/3*bAz；V=ΣVi+Vz")


@dataclass(frozen=True)
class PlaneGateForceInput:
    pressure_on_gate_kN: float
    water_head_m: float
    bearing_type: str
    sliding_support_friction: float = 0.15
    sliding_bearing_f: float = 0.30
    rolling_friction_coefficient: float = 0.10
    wheel_radius_cm: float = 2.0
    rolling_bearing_radius_cm: float = 0.3
    roller_diameter_cm: float = 4.0
    stopwater_friction: float = 0.65
    stopwater_pressure_kN: float = 49.66
    opening_safety_factor: float = 1.20
    closing_safety_factor: float = 1.20
    gate_weight_correction_open: float = 1.05
    gate_weight_correction_close: float = 0.95
    gate_weight_kN: float = 7.76
    ballast_weight_kN: float = 0.0
    downpull_kN: float = 0.0
    water_column_weight_kN: float = 22.06
    uplift_kN: float = 9.26


@dataclass(frozen=True)
class PlaneGateForceResult:
    support_friction_kN: float
    stopwater_friction_kN: float
    opening_force_kN: float
    closing_force_kN: float
    opening_force_t: float
    closing_force_t: float
    selected_capacity_t: float
    formula: str


def calculate_plane_gate_force(data: PlaneGateForceInput) -> PlaneGateForceResult:
    pressure = _non_negative(data.pressure_on_gate_kN, "作用在闸门上的总水压力p")
    bearing_type = data.bearing_type.strip()
    if bearing_type == "滑动支承":
        support = _positive(data.sliding_support_friction, "滑动支承摩擦系数f2") * pressure
    elif bearing_type == "滑动轴承":
        radius = _positive(data.wheel_radius_cm, "滚轮半径R")
        diameter = _positive(data.roller_diameter_cm, "滚柱直径d")
        support = pressure / radius * (_positive(data.sliding_bearing_f, "轴与轴套滑动摩擦系数f1") / diameter + _positive(data.rolling_friction_coefficient, "滚轮滚动摩擦系数f"))
    elif bearing_type == "滚动轴承":
        radius = _positive(data.wheel_radius_cm, "滚轮半径R")
        bearing_radius = _positive(data.rolling_bearing_radius_cm, "滚动轴承平均半径R1")
        diameter = _positive(data.roller_diameter_cm, "滚柱直径d")
        support = pressure * _positive(data.rolling_friction_coefficient, "滚轮滚动摩擦系数f") / radius * (bearing_radius / diameter + 1.0)
    else:
        raise HydrologyError("bearing_type必须为：滑动支承、滑动轴承或滚动轴承")
    stopwater = _non_negative(data.stopwater_pressure_kN, "止水水压力pzs") * _positive(data.stopwater_friction, "止水摩擦系数f3")
    open_force = (
        _positive(data.opening_safety_factor, "摩阻安全系数nT") * (support + stopwater)
        + _non_negative(data.downpull_kN, "下吸力ps")
        + _positive(data.gate_weight_correction_open, "启门力门重修正系数nG2") * _non_negative(data.gate_weight_kN, "门重G")
        + _non_negative(data.ballast_weight_kN, "加重块重Gj")
        + _non_negative(data.water_column_weight_kN, "水柱重Ws")
    )
    close_force = (
        _positive(data.closing_safety_factor, "摩阻安全系数nT") * (support + stopwater)
        - _positive(data.gate_weight_correction_close, "闭门力门重修正系数nG1") * _non_negative(data.gate_weight_kN, "门重G")
        + _non_negative(data.uplift_kN, "上托力pt")
    )
    open_t = open_force / 9.81
    close_t = close_force / 9.81
    return PlaneGateForceResult(
        support,
        stopwater,
        open_force,
        close_force,
        open_t,
        close_t,
        max(open_t, close_t),
        "FQ=nT(Tzd+Tzs)+ps+nG2G+Gj+Ws；FW=nT(Tzd+Tzs)-nG1G+pt",
    )
