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 RiverDischargeEstimateInput:
    average_depth: float
    water_width: float
    water_surface_slope: float
    roughness: float


@dataclass(frozen=True)
class RiverDischargeEstimateResult:
    velocity: float
    flow: float
    formula: str


def calculate_river_discharge_estimate(data: RiverDischargeEstimateInput) -> RiverDischargeEstimateResult:
    depth = _positive(data.average_depth, "河道平均水深H")
    width = _positive(data.water_width, "河道水面宽B")
    slope = _positive(data.water_surface_slope, "水面坡降i")
    roughness = _positive(data.roughness, "河床糙率n")
    velocity = depth ** (2.0 / 3.0) * math.sqrt(slope) / roughness
    flow = width * depth * velocity
    return RiverDischargeEstimateResult(velocity, flow, "v=H^(2/3)*sqrt(i)/n；Q=B*H*v")


@dataclass(frozen=True)
class TrainingLineWidthInput:
    flow: float
    roughness: float
    training_depth: float
    water_surface_slope: float


@dataclass(frozen=True)
class TrainingLineWidthResult:
    width: float
    formula: str


def calculate_training_line_width(data: TrainingLineWidthInput) -> TrainingLineWidthResult:
    flow = _positive(data.flow, "整治流量Q")
    roughness = _positive(data.roughness, "河床糙率n")
    depth = _positive(data.training_depth, "整治水深H")
    slope = _positive(data.water_surface_slope, "水面平均坡降J")
    width = flow * roughness / (depth ** (5.0 / 3.0) * math.sqrt(slope))
    return TrainingLineWidthResult(width, "B=Q*n/(H^(5/3)*sqrt(J))")


@dataclass(frozen=True)
class ManningLeveeCrestInput:
    flow: float
    bottom_width: float
    side_slope: float
    roughness: float
    bed_slope: float
    wind_speed: float
    fetch: float
    safety_freeboard: float


@dataclass(frozen=True)
class ManningLeveeCrestResult:
    water_depth: float
    wave_runup: float
    crest_freeboard: float
    levee_height: float
    check_flow: float
    formula: str


def _trapezoid_manning_flow(depth: float, bottom_width: float, side_slope: float, roughness: float, slope: float) -> float:
    area = (bottom_width + side_slope * depth) * depth
    wetted = bottom_width + 2.0 * depth * math.sqrt(1.0 + side_slope * side_slope)
    radius = area / wetted
    chezy = radius ** (1.0 / 6.0) / roughness
    return area * chezy * math.sqrt(radius * slope)


def _solve_trapezoid_depth(flow: float, bottom_width: float, side_slope: float, roughness: float, slope: float) -> float:
    low = 1e-9
    high = 1.0
    while _trapezoid_manning_flow(high, bottom_width, side_slope, roughness, slope) < flow:
        high *= 2.0
        if high > 10000:
            raise HydrologyError("未能反算设计水深，请检查流量、底宽、糙率和坡降")
    for _ in range(90):
        mid = (low + high) / 2.0
        if _trapezoid_manning_flow(mid, bottom_width, side_slope, roughness, slope) < flow:
            low = mid
        else:
            high = mid
    return (low + high) / 2.0


def calculate_manning_levee_crest(data: ManningLeveeCrestInput) -> ManningLeveeCrestResult:
    flow = _positive(data.flow, "设计流量Q")
    bottom_width = _positive(data.bottom_width, "河底宽b")
    side_slope = _non_negative(data.side_slope, "边坡系数m")
    roughness = _positive(data.roughness, "河床糙率n")
    slope = _positive(data.bed_slope, "河床比降i")
    wind_speed = _non_negative(data.wind_speed, "风速v")
    fetch = _non_negative(data.fetch, "吹程D")
    safety = _non_negative(data.safety_freeboard, "安全超高")
    depth = _solve_trapezoid_depth(flow, bottom_width, side_slope, roughness, slope)
    wave_runup = 0.0208 * wind_speed ** 1.25 * fetch ** (1.0 / 3.0) if wind_speed and fetch else 0.0
    freeboard = wave_runup + safety
    return ManningLeveeCrestResult(
        water_depth=depth,
        wave_runup=wave_runup,
        crest_freeboard=freeboard,
        levee_height=depth + freeboard,
        check_flow=_trapezoid_manning_flow(depth, bottom_width, side_slope, roughness, slope),
        formula="Q=A*C*sqrt(R*i)；e=0.0208*v^(5/4)*D^(1/3)；Δh=e+安全超高；H=h+Δh",
    )


@dataclass(frozen=True)
class RevetmentThicknessInput:
    layer_count: float
    coefficient: float
    block_weight_t: float
    block_unit_weight_kn_m3: float


@dataclass(frozen=True)
class RevetmentThicknessResult:
    thickness: float
    formula: str


def calculate_revetment_thickness(data: RevetmentThicknessInput) -> RevetmentThicknessResult:
    layer_count = _positive(data.layer_count, "护面层数n")
    coefficient = _positive(data.coefficient, "系数c")
    block_weight = _positive(data.block_weight_t, "块体个体重量Q")
    unit_weight = _positive(data.block_unit_weight_kn_m3, "块体重度rb")
    thickness = layer_count * coefficient * (block_weight / (0.1 * unit_weight)) ** (1.0 / 3.0)
    return RevetmentThicknessResult(thickness, "t=n*c*(Q/(0.1*rb))^(1/3)")


@dataclass(frozen=True)
class RevetmentBlockSizeInput:
    velocity: float
    water_density_t_m3: float
    stone_density_t_m3: float
    stability_coefficient: float
    block_length: float
    block_width: float
    block_height: float
    gravity: float = 9.81


@dataclass(frozen=True)
class RevetmentBlockSizeResult:
    stable_spherical_diameter: float
    actual_spherical_diameter: float
    length: float
    width: float
    height: float
    stable: bool
    formula: str


def calculate_revetment_block_size(data: RevetmentBlockSizeInput) -> RevetmentBlockSizeResult:
    velocity = _positive(data.velocity, "河道水流流速v")
    water_density = _positive(data.water_density_t_m3, "水的密度γ")
    stone_density = _positive(data.stone_density_t_m3, "石块密度γs")
    coefficient = _positive(data.stability_coefficient, "石块运动稳定系数c")
    length = _positive(data.block_length, "石块长a")
    width = _positive(data.block_width, "石块宽b")
    height = _positive(data.block_height, "石块高h")
    gravity = _positive(data.gravity, "重力加速度g")
    if stone_density <= water_density:
        raise HydrologyError("石块密度必须大于水的密度")
    stable_d = velocity * velocity / (2.0 * coefficient * coefficient * gravity * ((stone_density - water_density) / water_density))
    actual_d = 1.24 * (length * width * height) ** (1.0 / 3.0)
    stable = actual_d + 0.005 >= stable_d
    return RevetmentBlockSizeResult(
        stable_spherical_diameter=stable_d,
        actual_spherical_diameter=actual_d,
        length=length,
        width=width,
        height=height,
        stable=stable,
        formula="d=v^2/[2*c^2*g*((γs-γ)/γ)]；d=1.24*(a*b*h)^(1/3)",
    )


@dataclass(frozen=True)
class ObliqueBankScourInput:
    angle_deg: float
    side_slope: float
    soil_particle_diameter: float
    flow: float
    section_area: float
    contracted_area: float
    gravity: float = 9.81


@dataclass(frozen=True)
class ObliqueBankScourResult:
    local_velocity: float
    scour_depth: float
    formula: str


def calculate_oblique_bank_scour(data: ObliqueBankScourInput) -> ObliqueBankScourResult:
    angle = _positive(data.angle_deg, "水流与岸坡交角α")
    if angle >= 90:
        raise HydrologyError("水流与岸坡交角α应小于90°")
    side_slope = _non_negative(data.side_slope, "岸坡边坡系数m")
    particle = _non_negative(data.soil_particle_diameter, "坡脚土壤粒径d")
    flow = _positive(data.flow, "河道设计流量Q")
    area = _positive(data.section_area, "原河道过水断面面积ω")
    contracted = _non_negative(data.contracted_area, "缩窄部分过水断面面积ωp")
    gravity = _positive(data.gravity, "重力加速度g")
    if area <= contracted:
        raise HydrologyError("原河道过水断面面积必须大于缩窄部分面积")
    local_velocity = flow / (area - contracted)
    scour = 23.0 * math.tan(math.radians(angle) / 2.0) * local_velocity * local_velocity / (
        math.sqrt(1.0 + side_slope * side_slope) * gravity
    ) - 30.0 * particle
    return ObliqueBankScourResult(local_velocity, max(scour, 0.0), "hp=23*tan(α/2)*vj^2/(sqrt(1+m^2)*g)-30d；vj=Q/(ω-ωp)")


@dataclass(frozen=True)
class ParallelBankScourInput:
    pre_scour_depth: float
    average_velocity: float
    allowable_velocity: float
    exponent: float
    scour_coefficient: float


@dataclass(frozen=True)
class ParallelBankScourResult:
    post_scour_depth: float
    scour_depth: float
    formula: str


def calculate_parallel_bank_scour(data: ParallelBankScourInput) -> ParallelBankScourResult:
    depth = _positive(data.pre_scour_depth, "冲刷前水深h")
    velocity = _positive(data.average_velocity, "河槽平均流速vp")
    allowable = _positive(data.allowable_velocity, "河床允许不冲流速vH")
    exponent = _positive(data.exponent, "系数n")
    coefficient = _positive(data.scour_coefficient, "冲刷系数p")
    post_depth = coefficient * depth + ((velocity / allowable) ** exponent - 1.0)
    return ParallelBankScourResult(post_depth, max(post_depth - depth, 0.0), "hB=p*h+[(vp/vH)^n-1]")


@dataclass(frozen=True)
class BendScourInput:
    flow: float
    water_width: float
    straight_average_depth: float
    bed_particle_diameter: float
    water_density_t_m3: float
    sediment_density_t_m3: float
    gravity: float = 9.81


@dataclass(frozen=True)
class BendScourResult:
    lambda_coefficient: float
    max_depth_from_water_surface: float
    formula: str


def calculate_bend_scour(data: BendScourInput) -> BendScourResult:
    flow = _positive(data.flow, "河道流量Q")
    width = _positive(data.water_width, "弯道水面宽B")
    straight_depth = _positive(data.straight_average_depth, "直段平均水深H0")
    particle = _positive(data.bed_particle_diameter, "河床砂平均粒径d")
    water_density = _positive(data.water_density_t_m3, "水的重率γ")
    sediment_density = _positive(data.sediment_density_t_m3, "河床砂重率γs")
    gravity = _positive(data.gravity, "重力加速度g")
    if sediment_density <= water_density:
        raise HydrologyError("河床砂重率必须大于水的重率")
    lambda_coefficient = 0.64 * math.exp(3.61 * particle / straight_depth)
    denominator = width * particle ** (1.0 / 3.0) * math.sqrt(gravity * ((sediment_density - water_density) / water_density))
    hmax = (lambda_coefficient * flow / denominator) ** (6.0 / 7.0)
    return BendScourResult(lambda_coefficient, hmax, "Hmax=[λQ/(B*d^(1/3)*sqrt(g*((γs-γ)/γ)))]^(6/7)；λ=0.64e^(3.61d/H0)")


@dataclass(frozen=True)
class LockDamScourInput:
    unit_discharge: float
    bed_particle_diameter: float
    water_depth: float


@dataclass(frozen=True)
class LockDamScourResult:
    post_scour_depth: float
    formula: str


def calculate_lock_dam_scour(data: LockDamScourInput) -> LockDamScourResult:
    q = _positive(data.unit_discharge, "单宽流量q")
    particle = _positive(data.bed_particle_diameter, "河床砂平均粒径d")
    depth = _positive(data.water_depth, "河床水深h")
    post_depth = 0.332 * q / (math.sqrt(particle) * (depth / particle) ** (1.0 / 6.0))
    return LockDamScourResult(post_depth, "hp=0.332*q/[sqrt(d)*(h/d)^(1/6)]")


@dataclass(frozen=True)
class ShortGroyneDesignInput:
    flow: float
    roughness: float
    average_depth: float
    water_surface_slope: float
    length_factor: float = 0.3
    spacing_factor: float = 2.0


@dataclass(frozen=True)
class ShortGroyneDesignResult:
    main_channel_width: float
    groyne_length: float
    groyne_spacing: float
    formula: str


def calculate_short_groyne_design(data: ShortGroyneDesignInput) -> ShortGroyneDesignResult:
    width = calculate_training_line_width(
        TrainingLineWidthInput(
            flow=data.flow,
            roughness=data.roughness,
            training_depth=data.average_depth,
            water_surface_slope=data.water_surface_slope,
        )
    ).width
    length_factor = _positive(data.length_factor, "丁坝长度系数")
    spacing_factor = _positive(data.spacing_factor, "丁坝间距系数")
    length = length_factor * width
    spacing = spacing_factor * length
    return ShortGroyneDesignResult(width, length, spacing, "B=Qn/(H^(5/3)*sqrt(J))；L=(0.2~0.3)B；S=(1~2)L")
