from collections import deque
import pandas as pd

TERMINAL = {-888, -999}


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 isinstance(val, list) else val
        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):
    """Calculate routed runoff using a linear reservoir model.

    Parameters
    ----------
    df : DataFrame
        Must have columns: SDG6_SB_ID, start_datetime, AETI (mm), PCP (mm).
    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.

    Returns
    -------
    DataFrame
        Input df with added columns: Qin, Balance, Storage, Qout (all in m³).
    """
    df = df.sort_values(["SDG6_SB_ID", "start_datetime"]).reset_index(drop=True)

    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 = {}

    # Initialize result columns
    df["Qin"] = 0.0
    df["Balance"] = 0.0
    df["Storage"] = 0.0
    df["Qout"] = 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)

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

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

        balance = qin + df.loc[idx, "PCP"].values - df.loc[idx, "AETI"].values

        # Sequential storage/outflow computation
        n = len(idx)
        storage = pd.array([0.0] * n, dtype="float64")
        qout = pd.array([0.0] * 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

        df.loc[idx, "Qin"] = qin
        df.loc[idx, "Balance"] = balance
        df.loc[idx, "Storage"] = storage
        df.loc[idx, "Qout"] = qout

    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)
