"""Stage `erosion`: implicit stream-power incision (Braun & Willett 2013) + km-scale hillslope smoothing."""
from __future__ import annotations

import numpy as np

from .config import params
from .crust import CRATON, LIP, OCEANIC, POST_OROGEN, PRE_OROGEN, RIFT, SCAR, VOLCANO
from .elevation import solve_sea_level_connected
from .graph import (accumulate, drop_smooth_cache, ocean_mask, priority_flood, receiver_levels, smooth_km,
                    steepest_receivers)

DEFAULTS = {"min_sea_km2": 5.0e6, "steps": 12, "dt_myr": 5.0, "k": 0.02, "m": 0.5, "hillslope_km": 25.0, "sediment_fill": 0.15}
K_MULT = {OCEANIC: 0.0, CRATON: 1.5, PRE_OROGEN: 1.2, POST_OROGEN: 1.0, RIFT: 1.0, LIP: 0.6, SCAR: 0.9, VOLCANO: 0.8}


def erode(g, z, kmult, P):
    z = np.asarray(z, dtype=np.float64).copy()
    ar = np.arange(g.n)
    for _ in range(int(P["steps"])):
        ocean = ocean_mask(g, z, P["min_sea_km2"])
        zb = np.where(ocean, 0.0, z)                     # base level = sea level, not the seafloor
        zf = priority_flood(g, zb, ocean)
        zb = np.where(ocean, zb, zb + P["sediment_fill"] * (zf - zb))   # sediment infills closed basins
        recv, _, dist = steepest_receivers(g, zf)
        recv = np.where(ocean, ar, recv)
        levels = receiver_levels(recv)
        area = accumulate(recv, levels, g.area_km2)
        F = P["k"] * kmult * P["dt_myr"] * area ** P["m"] / np.where(np.isfinite(dist), dist, 1.0)
        F[ocean] = 0.0
        znew = zb.copy()
        for lv in levels[1:]:
            r = recv[lv]
            znew[lv] = np.minimum(zb[lv], (zb[lv] + F[lv] * znew[r]) / (1.0 + F[lv]))
        znew = smooth_km(g, znew, P["hillslope_km"], keep=True)   # hillslope/sub-grid smoothing, fixed length in km
        z = np.where(ocean, z, znew)
    drop_smooth_cache(g)
    return z


def run(ctx) -> dict:
    g = ctx.grid
    P = params(ctx.cfg, "erosion", DEFAULTS)
    z0, age = ctx.need("elevation_m", "age_class")
    kmult = np.vectorize(K_MULT.get)(age).astype(np.float64)
    z0 = np.asarray(z0, dtype=np.float64)
    added = np.asarray(ctx.data.get("land_added", np.zeros(g.n, bool)), bool)   # land the sea-level solve didn't
    target = ctx.cfg["build"]["land_fraction"] + g.area_km2[added & (z0 > 0)].sum() / g.area_km2.sum()   # see
    z = erode(g, z0, kmult, P)
    z = solve_sea_level_connected(g, z, target, P["min_sea_km2"])
    z = np.maximum(z, -11000.0)
    # the sea is the connected world ocean; a separate basin ≥ min_sea_km2 shaped the coasts above as sea and stays
    # open water for the climate (open_water), but its water is a lake (hydrology), not sea
    return {"elevation_eroded_m": z.astype(np.float32), "ocean": ocean_mask(g, z, np.inf),
            "open_water": ocean_mask(g, z, P["min_sea_km2"])}