from __future__ import annotations

from dataclasses import dataclass
import math
from typing import Sequence

from .core import HydrologyError


@dataclass(frozen=True)
class ReservoirRoutingInput:
    reservoir_name: str
    inflow_hydrograph: Sequence[tuple[float, float]]
    elevation_storage_curve: Sequence[tuple[float, float]]
    elevation_discharge_curve: Sequence[tuple[float, float]]
    initial_elevation_m: float
    flood_limit_level_m: float | None = None
    check_flood_level_m: float | None = None


@dataclass(frozen=True)
class ReservoirRoutingStep:
    index: int
    time_h: float
    interval_h: float
    inflow_m3s: float
    average_inflow_m3s: float | None
    elevation_m: float
    storage_m3: float
    outflow_m3s: float
    average_outflow_m3s: float | None
    interval_inflow_volume_m3: float | None
    interval_outflow_volume_m3: float | None
    storage_change_m3: float | None
    balance_error_m3: float | None


@dataclass(frozen=True)
class ReservoirRoutingResult:
    reservoir_name: str
    steps: tuple[ReservoirRoutingStep, ...]
    peak_inflow_m3s: float
    peak_inflow_time_h: float
    peak_outflow_m3s: float
    peak_outflow_time_h: float
    peak_delay_h: float
    peak_reduction_percent: float
    max_elevation_m: float
    max_elevation_time_h: float
    initial_storage_m3: float
    final_storage_m3: float
    total_inflow_volume_m3: float
    total_outflow_volume_m3: float
    net_storage_change_m3: float
    mass_balance_error_m3: float
    mass_balance_relative_percent: float
    warnings: tuple[str, ...]


def _validated_curve(
    values: Sequence[tuple[float, float]],
    label: str,
    *,
    require_strict_y: bool,
) -> tuple[tuple[float, float], ...]:
    points = tuple((float(x), float(y)) for x, y in values)
    if len(points) < 2:
        raise HydrologyError(f"{label}至少需要2个数据点")
    for index, (x, y) in enumerate(points):
        if not math.isfinite(x) or not math.isfinite(y):
            raise HydrologyError(f"{label}第{index + 1}行包含无效数字")
        if y < 0:
            raise HydrologyError(f"{label}第{index + 1}行的数值不能为负数")
        if index == 0:
            continue
        previous_x, previous_y = points[index - 1]
        if x <= previous_x:
            raise HydrologyError(f"{label}第一列必须严格递增，不能重复或倒序")
        if require_strict_y and y <= previous_y:
            raise HydrologyError(f"{label}库容必须随水位严格增加")
        if not require_strict_y and y < previous_y:
            raise HydrologyError(f"{label}泄量不能随水位减小")
    return points


def _validated_inflow(values: Sequence[tuple[float, float]]) -> tuple[tuple[float, float], ...]:
    points = tuple((float(time_h), float(flow)) for time_h, flow in values)
    if len(points) < 2:
        raise HydrologyError("入库洪水过程线至少需要2个时刻")
    for index, (time_h, flow) in enumerate(points):
        if not math.isfinite(time_h) or not math.isfinite(flow):
            raise HydrologyError(f"入库洪水过程线第{index + 1}行包含无效数字")
        if flow < 0:
            raise HydrologyError(f"入库洪水过程线第{index + 1}行流量不能为负数")
        if index and time_h <= points[index - 1][0]:
            raise HydrologyError("入库洪水过程线时间必须严格递增，不能重复或倒序")
    return points


def _interpolate(points: Sequence[tuple[float, float]], x: float, label: str) -> float:
    lower = points[0][0]
    upper = points[-1][0]
    tolerance = max(1.0, abs(lower), abs(upper)) * 1e-12
    if x < lower - tolerance or x > upper + tolerance:
        raise HydrologyError(f"{label}超出曲线范围：{x:.3f}不在{lower:.3f}至{upper:.3f}之间")
    if x <= lower:
        return points[0][1]
    if x >= upper:
        return points[-1][1]
    for (x1, y1), (x2, y2) in zip(points, points[1:]):
        if x <= x2:
            ratio = (x - x1) / (x2 - x1)
            return y1 + ratio * (y2 - y1)
    return points[-1][1]


def _solve_next_elevation(
    target: float,
    interval_seconds: float,
    storage_curve: Sequence[tuple[float, float]],
    discharge_curve: Sequence[tuple[float, float]],
    minimum_elevation: float,
    maximum_elevation: float,
) -> float:
    def storage_indication(elevation: float) -> float:
        storage = _interpolate(storage_curve, elevation, "水位库容曲线")
        outflow = _interpolate(discharge_curve, elevation, "水位泄量曲线")
        return 2.0 * storage / interval_seconds + outflow

    lower_value = storage_indication(minimum_elevation)
    upper_value = storage_indication(maximum_elevation)
    tolerance = max(1.0, abs(target), abs(lower_value), abs(upper_value)) * 1e-13
    if target < lower_value - tolerance:
        raise HydrologyError(
            "调洪水位低于曲线下限，请补充更低水位的库容、泄量关系，或检查初始水位和入流数据"
        )
    if target > upper_value + tolerance:
        raise HydrologyError(
            "调洪水位超过曲线上限，请补充更高水位的库容、泄量关系，不能直接外推"
        )
    if abs(target - lower_value) <= tolerance:
        return minimum_elevation
    if abs(target - upper_value) <= tolerance:
        return maximum_elevation

    low = minimum_elevation
    high = maximum_elevation
    for _ in range(150):
        middle = (low + high) / 2.0
        value = storage_indication(middle)
        if abs(value - target) <= tolerance:
            return middle
        if value < target:
            low = middle
        else:
            high = middle
    return (low + high) / 2.0


def route_reservoir_flood(data: ReservoirRoutingInput) -> ReservoirRoutingResult:
    """Route a reservoir flood with level-pool Modified Puls storage indication."""

    inflow = _validated_inflow(data.inflow_hydrograph)
    storage_curve = _validated_curve(
        data.elevation_storage_curve,
        "水位库容曲线",
        require_strict_y=True,
    )
    discharge_curve = _validated_curve(
        data.elevation_discharge_curve,
        "水位泄量曲线",
        require_strict_y=False,
    )
    initial_elevation = float(data.initial_elevation_m)
    if not math.isfinite(initial_elevation):
        raise HydrologyError("初始水位必须是有效数字")

    minimum_elevation = max(storage_curve[0][0], discharge_curve[0][0])
    maximum_elevation = min(storage_curve[-1][0], discharge_curve[-1][0])
    if minimum_elevation >= maximum_elevation:
        raise HydrologyError("水位库容曲线与水位泄量曲线没有共同的有效水位范围")
    if not minimum_elevation <= initial_elevation <= maximum_elevation:
        raise HydrologyError(
            f"初始水位必须位于两条曲线共同范围{minimum_elevation:.3f}至{maximum_elevation:.3f}m内"
        )

    for value, label in (
        (data.flood_limit_level_m, "汛限水位"),
        (data.check_flood_level_m, "校核洪水位"),
    ):
        if value is not None and not math.isfinite(float(value)):
            raise HydrologyError(f"{label}必须是有效数字")

    initial_storage = _interpolate(storage_curve, initial_elevation, "水位库容曲线")
    initial_outflow = _interpolate(discharge_curve, initial_elevation, "水位泄量曲线")
    steps: list[ReservoirRoutingStep] = [
        ReservoirRoutingStep(
            index=0,
            time_h=inflow[0][0],
            interval_h=0.0,
            inflow_m3s=inflow[0][1],
            average_inflow_m3s=None,
            elevation_m=initial_elevation,
            storage_m3=initial_storage,
            outflow_m3s=initial_outflow,
            average_outflow_m3s=None,
            interval_inflow_volume_m3=None,
            interval_outflow_volume_m3=None,
            storage_change_m3=None,
            balance_error_m3=None,
        )
    ]

    total_inflow = 0.0
    total_outflow = 0.0
    for index in range(1, len(inflow)):
        previous = steps[-1]
        previous_inflow = inflow[index - 1][1]
        time_h, current_inflow = inflow[index]
        interval_h = time_h - inflow[index - 1][0]
        interval_seconds = interval_h * 3600.0
        target = (
            previous_inflow
            + current_inflow
            + 2.0 * previous.storage_m3 / interval_seconds
            - previous.outflow_m3s
        )
        elevation = _solve_next_elevation(
            target,
            interval_seconds,
            storage_curve,
            discharge_curve,
            minimum_elevation,
            maximum_elevation,
        )
        storage = _interpolate(storage_curve, elevation, "水位库容曲线")
        outflow = _interpolate(discharge_curve, elevation, "水位泄量曲线")
        average_inflow = (previous_inflow + current_inflow) / 2.0
        average_outflow = (previous.outflow_m3s + outflow) / 2.0
        inflow_volume = average_inflow * interval_seconds
        outflow_volume = average_outflow * interval_seconds
        storage_change = storage - previous.storage_m3
        balance_error = storage_change - (inflow_volume - outflow_volume)
        total_inflow += inflow_volume
        total_outflow += outflow_volume
        steps.append(
            ReservoirRoutingStep(
                index=index,
                time_h=time_h,
                interval_h=interval_h,
                inflow_m3s=current_inflow,
                average_inflow_m3s=average_inflow,
                elevation_m=elevation,
                storage_m3=storage,
                outflow_m3s=outflow,
                average_outflow_m3s=average_outflow,
                interval_inflow_volume_m3=inflow_volume,
                interval_outflow_volume_m3=outflow_volume,
                storage_change_m3=storage_change,
                balance_error_m3=balance_error,
            )
        )

    peak_inflow_step = max(steps, key=lambda step: step.inflow_m3s)
    peak_outflow_step = max(steps, key=lambda step: step.outflow_m3s)
    max_elevation_step = max(steps, key=lambda step: step.elevation_m)
    net_storage_change = steps[-1].storage_m3 - steps[0].storage_m3
    mass_balance_error = total_inflow - total_outflow - net_storage_change
    balance_base = max(abs(total_inflow), 1.0)
    relative_error = mass_balance_error / balance_base * 100.0
    if peak_inflow_step.inflow_m3s > 0:
        reduction = (
            peak_inflow_step.inflow_m3s - peak_outflow_step.outflow_m3s
        ) / peak_inflow_step.inflow_m3s * 100.0
    else:
        reduction = 0.0

    warnings: list[str] = []
    if data.flood_limit_level_m is not None and max_elevation_step.elevation_m > float(data.flood_limit_level_m):
        warnings.append(
            f"最高水位{max_elevation_step.elevation_m:.3f}m超过汛限水位{float(data.flood_limit_level_m):.3f}m。"
        )
    if data.check_flood_level_m is not None and max_elevation_step.elevation_m > float(data.check_flood_level_m):
        warnings.append(
            f"最高水位{max_elevation_step.elevation_m:.3f}m超过校核洪水位{float(data.check_flood_level_m):.3f}m，当前方案不满足该控制条件。"
        )
    curve_margin = maximum_elevation - max_elevation_step.elevation_m
    curve_span = maximum_elevation - minimum_elevation
    if curve_margin <= max(0.02, curve_span * 0.02):
        warnings.append("最高水位已接近水位关系曲线上限，建议补充更高水位的数据后复核。")
    if steps[-1].inflow_m3s > max(peak_inflow_step.inflow_m3s * 0.05, 0.01):
        warnings.append("入库过程线末值仍较大，洪水过程可能尚未结束，建议继续补充退水段。")
    if steps[-1].outflow_m3s > max(peak_outflow_step.outflow_m3s * 0.05, 0.01) and steps[-1].storage_m3 > initial_storage * 1.001:
        warnings.append("计算末时刻库容尚未回落至初始状态，建议延长过程线直至退水基本结束。")
    if abs(relative_error) > 1e-6:
        warnings.append(f"总水量平衡相对误差为{relative_error:.6f}%，请检查输入与计算时段。")

    return ReservoirRoutingResult(
        reservoir_name=data.reservoir_name.strip() or "未命名水库",
        steps=tuple(steps),
        peak_inflow_m3s=peak_inflow_step.inflow_m3s,
        peak_inflow_time_h=peak_inflow_step.time_h,
        peak_outflow_m3s=peak_outflow_step.outflow_m3s,
        peak_outflow_time_h=peak_outflow_step.time_h,
        peak_delay_h=peak_outflow_step.time_h - peak_inflow_step.time_h,
        peak_reduction_percent=reduction,
        max_elevation_m=max_elevation_step.elevation_m,
        max_elevation_time_h=max_elevation_step.time_h,
        initial_storage_m3=steps[0].storage_m3,
        final_storage_m3=steps[-1].storage_m3,
        total_inflow_volume_m3=total_inflow,
        total_outflow_volume_m3=total_outflow,
        net_storage_change_m3=net_storage_change,
        mass_balance_error_m3=mass_balance_error,
        mass_balance_relative_percent=relative_error,
        warnings=tuple(warnings),
    )
