"""Stage `hydrology`: depression filling, lakes vs endorheic basins, discharge, rivers, Strahler order."""
from __future__ import annotations

import numpy as np
from scipy import sparse
from scipy.sparse import csgraph

from .config import params
from .crust import RIFT
from .graph import accumulate, components, distance_to, priority_flood, receiver_levels, steepest_receivers
from .noise import fbm
from .pipeline import StageError

DEFAULTS = {"fill_eps_m": 0.01, "min_depth_m": 1.0, "river_min_km3_yr": 2.0,
            "salt_flat_max_p_mm": 300.0, "salt_flat_fraction": 0.2, "dry_lake_fraction": 0.05,
            # lake beds: deepest point = k × area^exp m (Earth-like: ≈95 m at
            # 1,000 km², ≈230 m at 20,000 km², ≈610 m at 500,000 km²) × 0.5–2 (noise at the lake), × rift_x in
            # rifts, × arid_x for dry terminal lakes; never shallower than lake_min_m (refinement keeps ≥ 15 m)
            "lake_depth_k": 12.0, "lake_depth_exp": 0.3, "lake_min_m": 25.0, "lake_max_m": 1800.0,
            "lake_rift_x": 2.5, "lake_arid_x": 0.4, "lake_arid_p_mm": 400.0, "lake_shore": 0.25}


def strahler(recv, levels, river):
    n = len(recv)
    order = np.zeros(n, np.int8)
    mx = np.zeros(n, np.int8)
    cnt = np.zeros(n, np.int16)
    for lv in reversed(levels):
        cells = lv[river[lv]]
        if len(cells) == 0:
            continue
        order[cells] = np.where(cnt[cells] >= 2, mx[cells] + 1, np.maximum(mx[cells], 1))
        nonroot = cells[recv[cells] != cells]
        r, o = recv[nonroot], order[nonroot]
        np.maximum.at(mx, r, o)
        np.add.at(cnt, r, (o == mx[r]).astype(np.int16))
    return order


def _groups(lab):
    idx = np.argsort(lab, kind="stable")
    ls = lab[idx]
    start = np.searchsorted(ls, 0)
    idx, ls = idx[start:], ls[start:]
    cuts = np.flatnonzero(np.diff(ls)) + 1
    return np.split(idx, cuts) if len(idx) else []


def _leaves(recv, lab, x, limit=100000):
    """Does the flow from depression cell x leave its depression for good (not back in over shallow ground)?"""
    own, y = lab[x], recv[x]
    for _ in range(limit):
        if lab[y] == own:
            return False
        if lab[y] >= 0 or recv[y] == y:
            return True
        y = recv[y]
    return True


def _leaves_all(recv, lab, levels, limit=100000):
    """_leaves for every cell at once: walking down from recv[x], the first cell that is in a depression or a root
    (stop) and how many steps away it is; x leaves unless that cell is in x's own depression (or the walk is longer
    than limit steps)."""
    n = len(recv)
    stop, d = np.arange(n), np.zeros(n, np.int64)
    for lv in levels[1:]:                                    # roots first: a cell's receiver is done before it
        free = lv[lab[lv] < 0]
        stop[free] = stop[recv[free]]
        d[free] = d[recv[free]] + 1
    y = recv
    return (d[y] >= limit) | (lab[stop[y]] != lab)


def _route_inside(g, dep, lab, recv, targets):
    """Within each depression, point every cell along the shortest intra-depression path to its target cell."""
    e = dep[g.src] & dep[g.dst] & (lab[g.src] == lab[g.dst])
    m = sparse.csr_matrix((g.edge_km[e], (g.src[e], g.dst[e])), shape=(g.n, g.n))
    _, pred, _ = csgraph.dijkstra(m, directed=False, indices=targets, min_only=True, return_predecessors=True)
    out = recv.copy()
    inside = dep & (pred >= 0)
    out[inside] = pred[inside]
    return out


def _sweep(recv, levels, water, outlets, cap):
    """Accumulate flow upstream→downstream; at each spilling-lake outlet remove up to `cap` (lake evaporation)."""
    acc = np.asarray(water, dtype=np.float64).copy()
    loss = np.zeros(len(acc))
    is_out = np.zeros(len(acc), bool)
    is_out[outlets] = True
    cap_cell = np.zeros(len(acc))
    cap_cell[outlets] = cap
    buf = np.zeros(len(acc))                    # per level only the receivers change (graph.accumulate): same floats
    if len(levels) > 1:
        acc += 0.0
    for lv in reversed(levels[1:]):
        push = acc[lv].copy()
        o = is_out[lv]
        if o.any():
            cells = lv[o]
            lost = np.minimum(acc[cells], cap_cell[cells])
            loss[cells] = lost
            push[o] = acc[cells] - lost
        r = recv[lv]
        buf[r] = 0.0
        np.add.at(buf, r, push)
        acc[r] = acc[r] + buf[r]
    return acc, loss


def lake_levels(z, zf, lab, lake, lake_id, endo):
    """Each lake cell's water level (NaN elsewhere): a spilling lake stands at its spill height; a terminal lake at the
    lowest ground of its basin it does not cover (the next cell to flood). Carving the beds leaves both unchanged."""
    level = np.full(len(z), np.nan)
    m = int(lab.max()) + 1 if len(lab) and lab.max() >= 0 else 0
    dry = np.full(m, np.inf)                                    # per depression: its lowest uncovered ground
    sel = (lab >= 0) & ~lake
    np.minimum.at(dry, lab[sel], z[sel])
    for c in _groups(np.where(lake, lake_id, -1)):
        k = lab[c[0]]
        level[c] = dry[k] if endo[c[0]] and k >= 0 and np.isfinite(dry[k]) else zf[c].min()
    return level


def lake_beds(g, z, level, lake, lake_id, endo, rift, p_ann, seed, P, only=None):
    """Lake beds carved below their level (heights unchanged elsewhere): each lake's deepest point from its area
    (P lake_*), noise keyed by where the lake lies (not its id), the depth rising from lake_shore × that at the shore
    to all of it at the cell farthest from shore; never above the ground (min), never shallower than lake_min_m.
    only: carve just the lakes touching these cells (eras: where events changed the ground)."""
    z = np.asarray(z, dtype=np.float64).copy()
    if not lake.any():
        return z
    shore = distance_to(g, ~lake) if (~lake).any() else np.full(g.n, g.spacing_km)
    for c in _groups(np.where(lake, lake_id, -1)):
        if only is not None and not only[c].any():
            continue
        area = g.area_km2[c].sum()
        ctr = g.xyz[c].mean(0)
        ctr = ctr / max(np.linalg.norm(ctr), 1e-12)
        d = P["lake_depth_k"] * area ** P["lake_depth_exp"] * 2.0 ** float(fbm(ctr[None], seed + 4421, 3, 6.0)[0] * 1.4)
        if rift[c].mean() > 0.3:
            d *= P["lake_rift_x"]
        if endo[c[0]] and p_ann[c].mean() < P["lake_arid_p_mm"]:
            d *= P["lake_arid_x"]
        d = float(np.clip(d, P["lake_min_m"], P["lake_max_m"]))
        t = shore[c] / max(shore[c].max(), 1e-9)
        prof = np.maximum(d * (P["lake_shore"] + (1 - P["lake_shore"]) * t ** 0.6), P["lake_min_m"])
        z[c] = np.minimum(z[c], level[c] - prof)
    return z


def run(ctx) -> dict:
    g = ctx.grid
    P = params(ctx.cfg, "hydrology", DEFAULTS)
    z, p_ann, pet = ctx.need("elevation_eroded_m", "P_ann", "PET")
    z = z.astype(np.float64)
    ar = np.arange(g.n)
    ocean = np.asarray(ctx.data["ocean"]) if "ocean" in ctx.data else z <= 0
    aet = p_ann / np.sqrt(1.0 + (p_ann / np.maximum(pet, 1e-6)) ** 2)      # Pike (1964)
    runoff = np.where(ocean, 0.0, np.maximum(p_ann - aet, 0.0))
    water = runoff * g.area_km2 * 1e-6                                       # km³/yr per cell
    zf = priority_flood(g, z, ocean, P["fill_eps_m"])
    depth = zf - z
    dep = ~ocean & (depth > P["min_depth_m"])
    lab = components(g, dep)
    recv1, _, _ = steepest_receivers(g, zf)
    recv1 = np.where(ocean, ar, recv1)
    levels1 = receiver_levels(recv1)
    q1 = accumulate(recv1, levels1, water)
    evap_net = np.maximum(pet - p_ann, 0.0) * g.area_km2 * 1e-6              # full-lake evaporation, km³/yr

    groups = _groups(lab)
    leaves = _leaves_all(recv1, lab, levels1) if len(groups) else None
    outlets = np.zeros(len(groups), np.int64)
    terminals = np.zeros(len(groups), np.int64)
    cap = np.zeros(len(groups))
    for k, c in enumerate(groups):
        ext = c[lab[recv1[c]] != lab[c[0]]]
        ext = ext[leaves[ext]]
        outlets[k] = ext[np.argmax(q1[ext])] if len(ext) else c[np.argmax(q1[c])]
        terminals[k] = c[np.argmin(z[c])]
        cap[k] = evap_net[c].sum()

    # 1) every depression drains to its spill outlet; one upstream-first sweep decides spill vs endorheic
    recv_a = _route_inside(g, dep, lab, recv1, outlets) if len(groups) else recv1
    acc_a, _ = _sweep(recv_a, receiver_levels(recv_a), water, outlets, cap)
    inflow = acc_a[outlets]
    spill = inflow >= cap
    # 2) endorheic depressions drain to their single lowest cell instead
    targets = np.where(spill, outlets, terminals)
    recv = _route_inside(g, dep, lab, recv1, targets) if len(groups) else recv1
    recv[terminals[~spill]] = terminals[~spill]
    try:
        levels = receiver_levels(recv)
    except ValueError as e:
        raise StageError(f"hydrology: {e}") from e
    q, loss = _sweep(recv, levels, water, outlets[spill], cap[spill])

    lake = np.zeros(g.n, bool)
    endo = np.zeros(g.n, bool)
    salt = np.zeros(g.n, bool)
    lake_id = np.full(g.n, -1, np.int32)
    for k, c in enumerate(groups):
        if spill[k]:
            lake[c] = True
            lake_id[c] = k
            continue
        endo[c] = True
        cs = c[np.argsort(z[c])]
        nl = int(np.searchsorted(np.cumsum(evap_net[cs]), inflow[k]))
        if inflow[k] > 0:
            nl = max(nl, 1)                         # the terminal always holds some water
        lake[cs[:nl]] = True
        lake_id[cs[:nl]] = k
        if nl == 0 or (nl < P["dry_lake_fraction"] * len(cs) and p_ann[c].mean() < P["salt_flat_max_p_mm"]):
            a = np.cumsum(g.area_km2[cs])
            ns = max(nl + 1, int(np.searchsorted(a, P["salt_flat_fraction"] * a[-1])))
            salt[cs[nl:ns]] = True
    river = ~ocean & ~lake & (q >= P["river_min_km3_yr"])
    level = lake_levels(z, zf, lab, lake, lake_id, endo)
    rift = np.asarray(ctx.data["age_class"]) == RIFT if "age_class" in ctx.data else np.zeros(g.n, bool)
    zb = lake_beds(g, z, level, lake, lake_id, endo, rift, p_ann, ctx.seed, P, ctx.data.get("lake_carve"))
    cut = np.asarray(ctx.data.get("lake_cut_m", 0.0), dtype=np.float64) + (z - zb)   # sea masks see the uncut ground
    z = zb
    depth = zf - z
    dtype = np.asarray(ctx.data["elevation_eroded_m"]).dtype
    return {"elevation_eroded_m": z.astype(dtype), "lake_level_m": level.astype(np.float32),
            "lake_cut_m": np.broadcast_to(cut, z.shape).astype(np.float32), "z_filled_m": zf, "recv": recv, "discharge_km3_yr": q, "runoff_mm": runoff, "aet_mm": aet,
            "lake": lake, "lake_id": lake_id, "endorheic": endo, "salt_flat": salt, "river": river,
            "strahler": strahler(recv, levels, river), "depression_depth_m": depth, "lake_loss_km3_yr": loss}