"""Stage `elevation`: tectonic relief + Lightening provinces + hotspots + hints; sea level solved."""
from __future__ import annotations

import numpy as np

from .config import params
from .crust import PRE_OROGEN, SCAR, center_dist, lip_profile, name_seed
from .fields import gravity_mod
from .graph import OCEAN_MIN_KM2, distance_to, nearest_source, ocean_mask, smooth_km
from .noise import fbm, ridged
from . import plateaus as PL
from .pipeline import StageError
from .plates import edge_convergence
from .sphere import east_north, gc_dist_km, great_circle_point, latlon_to_xyz

OVER, SUB, COLLISION = 1, 2, 3
DEFAULTS = {
    "continental_base_m": 400.0, "continental_noise_m": 350.0,
    "margin_km": 500.0, "shelf_m": -200.0, "slope_km": 250.0,
    "coast_noise_m": 900.0, "coast_band_km": 600.0, "coast_noise_freq": 6.0,
    "ridge_depth_m": 2500.0, "age_depth_coeff": 350.0, "abyss_m": 6500.0,
    "rate_full_m_yr": 0.05, "min_rate_m_yr": 0.005,
    "trench_depth_m": 4000.0, "trench_width_km": 70.0,
    "arc_cont_m": 5500.0, "arc_ocean_m": 3500.0, "arc_offset_km": 200.0, "arc_width_km": 110.0,
    "collision_peak_m": 9500.0, "collision_width_km": 260.0, "plateau_m": 4500.0, "plateau_km": 800.0,
    "rift_depth_m": 1200.0, "rift_shoulder_m": 800.0,
    "pre_orogen_m": 1200.0, "lip_m": 1500.0, "lip_step_m": 300.0,
    "apron_m": 500.0, "scar_drop_m": 1500.0,
    "hotspot_m": 5500.0, "hotspot_spacing_km": 150.0, "hotspot_radius_km": 70.0,
    "hint_m": 2500.0, "hint_km": 80.0, "land_hint_m": 800.0, "spire_m": 2500.0, "detail_m": 250.0,
    "max_land_m": 12000.0, "min_ocean_m": -11000.0,
}


def solve_sea_level(z, area, land_fraction):
    order = np.argsort(-z)
    cum = np.cumsum(area[order]) / area.sum()
    k = min(int(np.searchsorted(cum, land_fraction)), len(z) - 1)
    return z - z[order[k]]


def solve_sea_level_connected(g, z, land_fraction, min_sea_km2=OCEAN_MIN_KM2, iters=30):
    """Shift z so land = everything outside the connected ocean covers land_fraction (interior pits stay land)."""
    area, tot = g.area_km2, g.area_km2.sum()
    if abs(area[~ocean_mask(g, z, min_sea_km2)].sum() / tot - land_fraction) <= 0.5 * area.min() / tot:
        return z                                   # already there: a flat sea floor would pull the bisection onto it
    s0 = float(np.asarray(z)[np.argsort(-z)][min(int(np.searchsorted(np.cumsum(area[np.argsort(-z)]) / tot,
                                                                           land_fraction)), len(z) - 1)])
    lo, hi = s0 - 3000.0, s0 + 3000.0
    for _ in range(iters):
        mid = 0.5 * (lo + hi)
        if area[~ocean_mask(g, z - mid, min_sea_km2)].sum() / tot > land_fraction:
            lo = mid
        else:
            hi = mid
    return z - 0.5 * (lo + hi)


def _smoothstep(a, b, x):
    t = np.clip((x - a) / (b - a), 0.0, 1.0)
    return t * t * (3 - 2 * t)


def roles(g, plate, vel, continental, ocean_age, plate_continental, min_rate):
    conv, _ = edge_convergence(g, vel)
    s, d = g.src, g.dst
    e = (plate[s] != plate[d]) & (conv > min_rate)
    cs, cd = continental[s[e]], continental[d[e]]
    ks, kd = plate_continental[plate[s[e]]], plate_continental[plate[d[e]]]
    over_s = np.where(cs & ~cd, True, np.where(~cs & cd, False,
                      np.where(ks & ~kd, True, np.where(~ks & kd, False, ocean_age[s[e]] < ocean_age[d[e]]))))
    role_e = np.where(cs & cd, COLLISION, np.where(over_s, OVER, SUB)).astype(np.int8)
    order = np.argsort(conv[e])
    role = np.zeros(g.n, np.int8)
    role[s[e][order]] = role_e[order]
    return role


def _on_plate(g, plate, role_mask, other=None):
    d, src = nearest_source(g, np.flatnonzero(role_mask))
    ok = src >= 0
    same = ok & (plate == plate[np.maximum(src, 0)])
    if other is not None:
        same |= ok & (plate == other[np.maximum(src, 0)])
    return np.where(same, d, np.inf), np.maximum(src, 0)


def hotspot_track(g, h: dict, vel, spacing_km: float):
    """(unit vector, km from the active end) along a hotspot chain; the chain follows the plate's motion."""
    c = latlon_to_xyz(*h["center"])
    v = vel[g.cell_index(*h["center"])]
    sp = np.linalg.norm(v)
    t = v / sp if sp > 0 else east_north(c[None])[0][0]
    return [(great_circle_point(c, t, k * spacing_km, g.radius_km), k * spacing_km)
            for k in range(int(h["length_km"] // spacing_km) + 1)]


def _relief(ctx, g, P, cont):
    """Heights before the sea-level solve for the continental mask `cont` → (z, role, d_over, d_sub, d_coll)."""
    R, seed = g.radius_km, ctx.seed
    (plate, vel, pk, brate, bother, age, ocean_age, d_div, land_hint, sk_mtn, mtn_hint) = ctx.need(
        "plate", "vel", "plate_continental", "bnd_rate", "bnd_other", "age_class", "ocean_age_myr", "d_div_km",
        "m_land_hint", "sk_mountains", "m_mountain_hint")
    f = lambda r: np.clip(r / P["rate_full_m_yr"], 0.3, 1.0)

    # passive-margin profile: interior plateau → coastal ramp → shelf → slope → abyssal floor
    d_in = distance_to(g, ~cont) if (~cont).any() else np.full(g.n, np.inf)
    d_out = distance_to(g, cont) if cont.any() else np.full(g.n, np.inf)
    ramp = _smoothstep(0.0, P["margin_km"], d_in)
    z_cont = P["shelf_m"] + (P["continental_base_m"] - P["shelf_m"]) * ramp
    z_cont = z_cont + P["continental_noise_m"] * fbm(g.xyz, seed + 31, 5, 3.0)
    z_ocean = -np.minimum(P["ridge_depth_m"] + P["age_depth_coeff"] * np.sqrt(ocean_age), P["abyss_m"])
    z_ocean = P["shelf_m"] + (z_ocean - P["shelf_m"]) * _smoothstep(0.0, P["slope_km"], d_out)
    z = np.where(cont, z_cont, z_ocean)
    # multi-scale coastal noise: sea level cuts it → bays, headlands, drowned valleys, offshore islands
    d_edge = np.where(cont, d_in, d_out)
    z += P["coast_noise_m"] * np.exp(-d_edge / P["coast_band_km"]) * fbm(g.xyz, seed + 91, 7, P["coast_noise_freq"])

    role = roles(g, plate, vel, cont, ocean_age, pk, P["min_rate_m_yr"])
    d_over, s_over = _on_plate(g, plate, role == OVER)
    d_sub, s_sub = _on_plate(g, plate, role == SUB)
    d_coll, s_coll = _on_plate(g, plate, role == COLLISION, other=bother)

    z += -P["trench_depth_m"] * f(brate[s_sub]) * np.exp(-(d_sub / P["trench_width_km"]) ** 2)
    arc_h = np.where(cont, P["arc_cont_m"], P["arc_ocean_m"])
    z += arc_h * f(brate[s_over]) * np.exp(-((d_over - P["arc_offset_km"]) / P["arc_width_km"]) ** 2)
    w = P["collision_width_km"]
    z += P["collision_peak_m"] * f(brate[s_coll]) * np.exp(-(d_coll / w) ** 2)
    plateau = _smoothstep(0.5 * w, 1.5 * w, d_coll) * (1 - _smoothstep(0.7 * P["plateau_km"], P["plateau_km"], d_coll))
    z += np.where(cont, P["plateau_m"] * f(brate[s_coll]) * plateau, 0.0)
    rift = -P["rift_depth_m"] * np.exp(-(d_div / 60.0) ** 2) + P["rift_shoulder_m"] * np.exp(-((d_div - 120.0) / 60.0) ** 2)
    z += np.where(cont, rift, 0.0)

    z += np.where(age == PRE_OROGEN, P["pre_orogen_m"] * ridged(g.xyz, seed + 21, 4, 4.0), 0.0)
    for lip in ctx.tect.get("lip", []):
        prof = lip_profile(g, lip, seed)
        z += np.floor(P["lip_m"] * prof / P["lip_step_m"]) * P["lip_step_m"]
    for v in ctx.tect.get("volcano", []):
        d = center_dist(g, v["center"])
        r, H = v["radius_km"], v.get("height_m", 7000.0)
        z += H * np.clip(1 - d / r, 0, 1) ** 1.5 - 0.35 * H * np.exp(-(d / (0.12 * r)) ** 2)
        z += P["apron_m"] * np.exp(-((d - r) / (0.3 * r)) ** 2) * (0.5 + 0.5 * fbm(g.xyz, name_seed(seed, v["name"]), 3, 20.0))
    z -= np.where(age == SCAR, P["scar_drop_m"], 0.0)

    for h in ctx.tect.get("hotspot", []):
        peaks = np.zeros(g.n)
        for p, s in hotspot_track(g, h, vel, P["hotspot_spacing_km"]):
            dk = gc_dist_km(g.xyz, p, R)
            hk = P["hotspot_m"] * (1 - s / h["length_km"])
            peaks = np.maximum(peaks, hk * np.clip(1 - dk / P["hotspot_radius_km"], 0, 1) ** 1.2)
        z += peaks

    hint = smooth_km(g, np.clip(sk_mtn + mtn_hint, 0, 1), P["hint_km"])
    z += np.where(cont, P["hint_m"] * hint, 0.0)
    z += P["land_hint_m"] * land_hint
    z += P["detail_m"] * fbm(g.xyz, seed + 41, 6, 12.0)
    return z, role, d_over, d_sub, d_coll


def run(ctx) -> dict:
    g = ctx.grid
    P = params(ctx.cfg, "elevation", DEFAULTS)
    continental, sk_land, land_hint, m_grav, m_lock = ctx.need("continental", "sk_land", "m_land_hint",
                                                               "m_gravity_zones", "m_lock")
    plateau_id = np.asarray(ctx.data.get("plateau_id", np.full(g.n, -1)))
    cont = np.asarray(continental) & (plateau_id < 0)          # plateaus: ocean until their surfaces are set (§3)
    base = np.asarray(ctx.data.get("continental_base", continental)) & (plateau_id < 0)
    z, role, d_over, d_sub, d_coll = _relief(ctx, g, P, cont)

    target = ctx.cfg["build"]["land_fraction"]
    cont_area = g.area_km2[base].sum() / g.area_km2.sum()
    if target > cont_area + 0.005:
        raise StageError(f"land_fraction {target} exceeds continental crust area {cont_area:.3f}: sea level would "
                         f"lift ocean ridges into land; lower [build] land_fraction or raise [crust] shelf_km")
    min_sea = ctx.cfg.get("erosion", {}).get("min_sea_km2", OCEAN_MIN_KM2)
    sea_ref = np.zeros(g.n, bool)                              # sea in the world without land patches
    if np.array_equal(base, cont):
        z = solve_sea_level_connected(g, z, target, min_sea)
    else:   # land patches (the eastern continent made whole) must not move every other coast: the sea level of the
        z_raw = _relief(ctx, g, P, base)[0]                  # world without them, applied to the world with them
        z_ref = solve_sea_level_connected(g, z_raw, target, min_sea)
        z = z - float(np.mean(z_raw - z_ref))
        sea_ref = ocean_mask(g, z_ref, min_sea)
        patch = smooth_km(g, np.asarray(ctx.data.get("land_patch", np.zeros(g.n)), dtype=np.float64), P["hint_km"])
        z += P["land_hint_m"] * patch                        # a patch is a land hint in height too (bare crust
                                                             # sits below the world's sea level: land)
    mult = np.clip(1.0 / gravity_mod(ctx.cfg, m_grav), 1.0, 3.0)
    spires = P["spire_m"] * (mult - 1) * ridged(g.xyz, ctx.seed + 71, 4, 40.0)
    z = np.where(z > 0, z * mult + spires, z)
    target_land = (sk_land + 0.5 * land_hint) > 0.5
    lock = m_lock > 0.5
    z = np.where(lock & target_land, np.maximum(z, 50.0), np.where(lock & ~target_land, np.minimum(z, -50.0), z))
    z = PL.apply(g, z, ctx.tect.get("plateau", []), plateau_id, ctx.seed)     # sunken plateaus (after the solve)
    z = np.clip(z, P["min_ocean_m"], P["max_land_m"] * mult)
    added = ~ocean_mask(g, z, min_sea) & (sea_ref | (plateau_id >= 0))       # land the sea-level solve never saw
    return {"elevation_m": z.astype(np.float32), "role": role, "d_over_km": d_over, "d_sub_km": d_sub,
            "d_coll_km": d_coll, "land_added": added}