from calendar import monthrange
from collections import deque
import numpy as np
import pandas as pd

TERMINAL = {-888, -999}


def _period_days(dt):
    """Return the number of days in the dekadal period starting at *dt*.

    Day 1 or 11 → 10 days; day 21 → remaining days in the month.
    """
    day = dt.day
    if day in (1, 11):
        return 10
    if day == 21:
        return monthrange(dt.year, dt.month)[1] - 20
    raise ValueError(f"Unexpected start day {day} in {dt}")


def _topo_sort(routes, basin_ids):
    """Topological sort of basin_ids so upstream basins come first."""
    # Build upstream -> downstream edges, filtered to basins in df
    id_set = set(basin_ids)
    # routes keys are strings, values are ints
    downstream = {}
    upstream_of = {bid: [] for bid in id_set}
    in_degree = {bid: 0 for bid in id_set}

    for src_str, val in routes.items():
        src = int(src_str)
        if src not in id_set:
            continue
        dst = val[0]
        if dst in TERMINAL or dst not in id_set:
            downstream[src] = None
        else:
            downstream[src] = dst
            upstream_of[dst].append(src)
            in_degree[dst] += 1

    # Kahn's algorithm
    queue = deque(bid for bid, deg in in_degree.items() if deg == 0)
    order = []
    while queue:
        bid = queue.popleft()
        order.append(bid)
        ds = downstream.get(bid)
        if ds is not None:
            in_degree[ds] -= 1
            if in_degree[ds] == 0:
                queue.append(ds)

    if len(order) != len(id_set):
        missing = id_set - set(order)
        raise ValueError(f"Cycle detected or disconnected basins: {missing}")

    return order, upstream_of


def calc_runoff(df, routes, f_values=None, output="mixed_flux"):
    """Calculate routed runoff using a linear reservoir model.

    Parameters
    ----------
    df : DataFrame
        Must have columns: SDG6_SB_ID, start_datetime, AETI (mm/day), PCP (mm/day).
    routes : dict
        Mapping basin_id (str) -> [downstream_id, area_km2] or downstream_id (int).
        -888/-999 = terminal basin.
    f_values : dict, optional
        Response factor per basin_id (int). Defaults to 0.1 for all.
    output : {"volume", "volume_flux", "depth", "depth_flux", "mixed", "mixed_flux"}, optional
        Controls output units (default "volume").
        - "volume": PCP, AETI, Q, ΔS in m³.
        - "volume_flux": PCP, AETI, Q, ΔS in m³/s.
        - "depth": PCP, AETI, Q, ΔS in mm.
        - "depth_flux": PCP, AETI, Q, ΔS in mm/day.
        - "mixed": PCP, AETI in mm; Q, ΔS in m³.
        - "mixed_flux": PCP, AETI in mm/day; Q, ΔS in m³/s.

    Returns
    -------
    DataFrame
        Input df with added columns depending on *output* mode.
    """
    df = df.sort_values(["SDG6_SB_ID", "start_datetime"]).reset_index(drop=True)

    if df.empty:
        for col in ("Q", "ΔS"):
            df[col] = np.float64()
        return df

    basin_ids = df["SDG6_SB_ID"].unique()
    default_f = 0.1

    # Extract areas from routes (km² -> m²)
    areas = {}
    for src_str, val in routes.items():
        if isinstance(val, list) and len(val) >= 2:
            areas[int(src_str)] = val[1] * 1e6  # km² to m²

    if f_values is None:
        f_values = {}

    # Compute period length (days) for each row from start_datetime
    df["_period_days"] = df["start_datetime"].apply(_period_days)

    # Initialize result columns
    df["QIN"] = 0.0
    df["Storage"] = 0.0
    df["Q"] = 0.0

    order, upstream_of = _topo_sort(routes, basin_ids)

    for bid in order:
        mask = df["SDG6_SB_ID"] == bid
        idx = df.index[mask]
        f = f_values.get(bid, default_f)
        area = areas.get(bid, 1e6)
        period_days = df.loc[idx, "_period_days"].values

        # mm/day to m³ per period: mm/day × days × m² / 1000
        df.loc[idx, "PCP"] = df.loc[idx, "PCP"].values * period_days * area / 1e3
        df.loc[idx, "AETI"] = df.loc[idx, "AETI"].values * period_days * area / 1e3

        # Qin: sum of upstream Qout per time step (m³ per period)
        ups = upstream_of[bid]
        if ups:
            up_mask = df["SDG6_SB_ID"].isin(ups)
            qin_series = df.loc[up_mask].groupby("start_datetime")["Q"].sum()
            qin = df.loc[idx, "start_datetime"].map(qin_series).fillna(0.0).values
        else:
            qin = 0.0

        # Net water input for this period (m³): inflow + precipitation - evapotranspiration
        balance = qin + df.loc[idx, "PCP"].values - df.loc[idx, "AETI"].values

        # Sequential storage/outflow computation using a linear reservoir model.
        # Storage is the cumulative water volume (m³) held in the basin, updated
        # each step by adding the new balance and subtracting the previous outflow.
        n = len(idx)
        storage = np.zeros(n, dtype="float64")
        qout = np.zeros(n, dtype="float64")

        for t in range(n):
            if t == 0:
                storage[t] = balance[t]
            else:
                storage[t] = storage[t - 1] + balance[t] - qout[t - 1]
            if storage[t] < 0.0:
                storage[t] = 0.0
            qout[t] = storage[t] * f  # outflow is a fraction of current storage

        df.loc[idx, "QIN"] = qin
        df.loc[idx, "Storage"] = storage
        df.loc[idx, "Q"] = qout

    if output == "volume":
        # ΔS from water balance (m³)
        df["ΔS"] = df["QIN"].values + df["PCP"].values - df["AETI"].values - df["Q"].values
        df.drop(columns=["Storage"], inplace=True)
    elif output == "volume_flux":
        period_s = df["_period_days"].values * 86400
        # ΔS from water balance (m³ → m³/s)
        delta_s = df["QIN"].values + df["PCP"].values - df["AETI"].values - df["Q"].values
        df["PCP"] = df["PCP"].values / period_s
        df["AETI"] = df["AETI"].values / period_s
        df["Q"] = df["Q"].values / period_s
        df["ΔS"] = delta_s / period_s
        df.drop(columns=["Storage"], inplace=True)
    elif output == "depth":
        # m³ → mm: value_m3 / area_m2 * 1000
        area_per_row = df["SDG6_SB_ID"].map(areas).values
        # ΔS from water balance (m³ → mm)
        delta_s = df["QIN"].values + df["PCP"].values - df["AETI"].values - df["Q"].values
        for col in ("PCP", "AETI", "Q"):
            df[col] = df[col].values / area_per_row * 1e3
        df["ΔS"] = delta_s / area_per_row * 1e3
        df.drop(columns=["Storage"], inplace=True)
    elif output == "depth_flux":
        # m³ → mm/day: value_m3 / area_m2 * 1000 / period_days
        area_per_row = df["SDG6_SB_ID"].map(areas).values
        period_days = df["_period_days"].values
        # ΔS from water balance (m³ → mm/day)
        delta_s = df["QIN"].values + df["PCP"].values - df["AETI"].values - df["Q"].values
        for col in ("PCP", "AETI", "Q"):
            df[col] = df[col].values / area_per_row * 1e3 / period_days
        df["ΔS"] = delta_s / area_per_row * 1e3 / period_days
        df.drop(columns=["Storage"], inplace=True)
    elif output == "mixed":
        # PCP, AETI: m³ → mm; Q, ΔS remain in m³
        area_per_row = df["SDG6_SB_ID"].map(areas).values
        # ΔS from water balance (m³)
        delta_s = df["QIN"].values + df["PCP"].values - df["AETI"].values - df["Q"].values
        for col in ("PCP", "AETI"):
            df[col] = df[col].values / area_per_row * 1e3
        df["ΔS"] = delta_s
        df.drop(columns=["Storage"], inplace=True)
    elif output == "mixed_flux":
        # PCP, AETI: m³ → mm/day; Q: m³ → m³/s; ΔS: m³ → m³/s
        area_per_row = df["SDG6_SB_ID"].map(areas).values
        period_days = df["_period_days"].values
        period_s = period_days * 86400
        # ΔS from water balance (m³ → m³/s)
        delta_s = df["QIN"].values + df["PCP"].values - df["AETI"].values - df["Q"].values
        for col in ("PCP", "AETI"):
            df[col] = df[col].values / area_per_row * 1e3 / period_days
        df["Q"] = df["Q"].values / period_s
        df["ΔS"] = delta_s / period_s
        df.drop(columns=["Storage"], inplace=True)

    df.drop(columns=["_period_days", "QIN"], inplace=True)
    return df


if __name__ == "__main__":

    import json
    from pathlib import Path
    from holoviz_utils.hydrology import calc_runoff

    _routes = json.loads(
        Path(
            "/Users/hmcoerver/Library/Mobile Documents/com~apple~CloudDocs/Repositories/holoviz-utils/projects/wapor_discharge/fao_sb_ids.json"
        ).read_text()
    )

    df_in = pd.read_pickle(
        "/Users/hmcoerver/Library/Mobile Documents/com~apple~CloudDocs/Repositories/holoviz-utils/in_progress/runoff_routing/test_data.pickle"
    )

    df = calc_runoff(df_in, _routes)
