"""Sea-floor detail for refined areas:
turbidity canyons and their fans, sediment ponds, fine volcanic fields on the sunken plateaus, vents (data) and the
refined sea-floor fields. Only sea cells inside the area change, and never above SEA_TOP_M; land and the halo keep
their heights. Deterministic: the same cells and world values give the same floor. Generated terrain (`idea`)."""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np
from scipy.spatial import cKDTree

import h3par
import worldgen_path  # noqa: F401  (mapgen on sys.path)
from mapgen import plateaus as PL, seabed as SB
from mapgen.graph import accumulate, components, priority_flood, receiver_levels, smooth_km, steepest_receivers
from mapgen.noise import name_seed
from mapgen.sphere import east_north, gc_dist_km, great_circle_point, latlon_to_xyz
from mapgen.zones import smootherstep

SEA_TOP_M = -5.0               # the pass never lifts a sea cell above this (no new land)
WALL_M = 1.0e7                 # land in the sea-floor routing
CANYON = {"k": 0.1, "m": 0.5, "n": 1.0, "min_area_km2": 500.0, "min_slope": 2.0, "cap_m": 1500.0}   # slope: m/km
FAN = {"flat_slope": 1.0, "max_m": 100.0, "km": 15.0}
POND = {"min_depth_m": 50.0, "fill": 0.3}
VENT_TYPES = ["black smoker", "white smoker", "diffuse", "cold seep"]
VENT_MINERALS = ["copper-iron sulfides", "zinc sulfides and barite", "iron-manganese oxides", "methane carbonates"]
VENT = {"rate": 0.03, "seep_rate": 0.002, "min_depth_m": 200.0, "ridge_km": 60.0, "reach_km": 3.0}
TEMP_C = (np.array([300.0, 100.0, 10.0, 0.0]), np.array([100.0, 200.0, 90.0, 5.0]))   # per type: from, span


def canyons(g, z, sea, halo, P=CANYON, F=FAN):
    """(heights, canyon depth, fan thickness). Turbidity flows follow the steepest descent over the sea floor (land is
    a wall); where their contributing sea area and slope are high they cut (stream-power form, like rivers, capped
    at cap_m); the material they carry settles as a fan where they first reach flat floor (slope < flat_slope),
    spread over ≈ km and at most max_m thick. Canyons only lower, fans only raise; the halo stays."""
    z = np.asarray(z, dtype=np.float64)
    sea, halo = np.asarray(sea, bool), np.asarray(halo, bool)
    active = sea & ~halo
    recv, slope, _ = steepest_receivers(g, np.where(sea, z, WALL_M))
    levels = receiver_levels(recv)
    area = accumulate(recv, levels, np.where(sea, g.area_km2, 0.0))
    cut = np.where(active & (area >= P["min_area_km2"]) & (slope >= P["min_slope"]),
                   np.minimum(P["cap_m"], P["k"] * area ** P["m"] * slope ** P["n"]), 0.0)
    vol = accumulate(recv, levels, cut * g.area_km2 / 1000.0)                  # km³ carried down
    flat = active & (slope < F["flat_slope"])
    ar = np.arange(g.n)
    src = active & ~flat & (recv != ar) & (vol > 0)
    entry = np.zeros(g.n, bool)
    entry[recv[src]] = True
    entry &= flat
    dep = np.where(entry, vol * 1000.0 / g.area_km2, 0.0)                       # m over the entry cell
    fan = np.where(active & (cut == 0), np.clip(smooth_km(g, dep, F["km"]), 0.0, F["max_m"]), 0.0)
    zc = z - cut
    zn = np.where(fan > 0, np.maximum(zc, np.minimum(zc + fan, SEA_TOP_M)), zc)
    return zn, cut, zn - zc


def ponds(g, z, sea, halo, sediment_m, P=POND):
    """(heights, fill thickness): closed sea-floor hollows deeper than min_depth_m below their spill level fill
    with sediment to a flat floor at min(spill level, lowest point + fill × the mean sediment there). Never above
    the spill level, never lower."""
    z = np.asarray(z, dtype=np.float64)
    sea, halo = np.asarray(sea, bool), np.asarray(halo, bool)
    out = z.copy()
    if not (sea & ~halo).any():
        return out, np.zeros(g.n)
    sinks = sea & halo
    if not sinks.any():                                       # sea closed inside the area: its deepest cell drains
        sinks = np.zeros(g.n, bool)
        s = np.flatnonzero(sea)
        sinks[s[np.argmin(z[s])]] = True
    zf = priority_flood(g, np.where(sea, z, WALL_M), sinks, eps=0.0)
    hollow = sea & ~halo & (zf - z > 0.5)
    if not hollow.any():
        return out, np.zeros(g.n)
    lab = components(g, hollow)
    idx = np.flatnonzero(hollow)
    L = lab[idx]
    n = int(L.max()) + 1
    low = np.full(n, np.inf)
    np.minimum.at(low, L, z[idx])
    spill = np.full(n, -np.inf)
    np.maximum.at(spill, L, zf[idx])
    cnt = np.bincount(L, minlength=n)
    sup = np.bincount(L, weights=np.asarray(sediment_m, dtype=np.float64)[idx], minlength=n) / np.maximum(cnt, 1)
    level = np.minimum(spill, low + P["fill"] * sup)
    deep = (spill - low) >= P["min_depth_m"]
    lv = np.where(deep[L], np.minimum(level[L], SEA_TOP_M), -np.inf)
    out[idx] = np.maximum(z[idx], lv)
    return out, out - z


def _u(ids, salt: int) -> np.ndarray:
    """Uniform [0, 1) per H3 cell id (splitmix64): a cell draws the same number in every area and every era."""
    x = np.asarray(ids, dtype=np.uint64) + np.uint64((int(salt) * 0x9E3779B97F4A7C15) & 0xFFFFFFFFFFFFFFFF)
    x = x + np.uint64(0x9E3779B97F4A7C15)
    x = (x ^ (x >> np.uint64(30))) * np.uint64(0xBF58476D1CE4E5B9)
    x = (x ^ (x >> np.uint64(27))) * np.uint64(0x94D049BB133111EB)
    x = x ^ (x >> np.uint64(31))
    return (x >> np.uint64(11)).astype(np.float64) / float(1 << 53)


def _around(rng, lat, lon, reach_km, radius_km):
    c = latlon_to_xyz(lat, lon)
    e, n = east_north(c[None])
    az = rng.uniform(0.0, 2.0 * np.pi)
    return great_circle_point(c, np.cos(az) * n[0] + np.sin(az) * e[0], rng.uniform(0.0, reach_km), radius_km)


def small_features(p: dict, feat: dict, seed: int, radius_km: float) -> list:
    """A plateau's fine volcanic features, deterministic per name and independent of any area: 2–6 cones (2–15 km
    wide, 0.2–2 km tall) around each world-scale cone, 1–2 pit calderas in each world caldera, and a fissure ridge
    from the hotspot to each caldera (the hotspot's feeding track)."""
    rng = np.random.default_rng(name_seed(seed, p["name"]) + 11)
    vent = float(p.get("vent", 1.0))
    out = []
    for c in feat["cones"]:
        for _ in range(int(rng.integers(2, 7))):
            out.append({"kind": "cone", "at": _around(rng, c["lat"], c["lon"], c["radius_km"], radius_km),
                        "radius_km": float(rng.uniform(1.0, 7.5)), "height_m": float(rng.uniform(200.0, 2000.0)) * vent})
    for c in feat["calderas"]:
        for _ in range(int(rng.integers(1, 3))):
            out.append({"kind": "pit", "at": _around(rng, c["lat"], c["lon"], 0.5 * c["radius_km"], radius_km),
                        "radius_km": float(rng.uniform(2.0, 6.0)), "height_m": 150.0 * vent, "floor_m": -250.0 * vent})
    hot = latlon_to_xyz(*feat["hotspot"])
    for c in feat["calderas"]:
        out.append({"kind": "fissure", "at": hot, "to": latlon_to_xyz(c["lat"], c["lon"]), "radius_km": 1.5,
                    "height_m": float(rng.uniform(100.0, 400.0)) * vent})
    return out


def _segment_km(xyz, a, b, radius_km, step_km=1.0):
    n = max(2, int(gc_dist_km(a, b, radius_km) / step_km) + 1)
    t = np.linspace(0.0, 1.0, n)[:, None]
    pts = a * (1.0 - t) + b * t
    pts /= np.linalg.norm(pts, axis=1, keepdims=True)
    chord, _ = cKDTree(pts).query(xyz, workers=max(1, h3par.workers()))
    return 2.0 * np.arcsin(np.clip(chord / 2.0, 0.0, 1.0)) * radius_km


def volcanic(g, sea, halo, plateau_id, plateaus: list, seed: int):
    """(height to add, nearness 0–1 to a feature) from the fine volcanic features of the plateaus under the area's
    sea cells (vent = 0: none)."""
    add, near = np.zeros(g.n), np.zeros(g.n)
    for k, p in enumerate(plateaus):
        on = np.flatnonzero((np.asarray(plateau_id) == k) & sea & ~halo)
        if not len(on) or float(p.get("vent", 1.0)) <= 0:
            continue
        xyz = g.xyz[on]
        for f in small_features(p, PL.features(p, seed, g.radius_km), seed, g.radius_km):
            r = f["radius_km"]
            if f["kind"] == "fissure":
                d = _segment_km(xyz, f["at"], f["to"], g.radius_km)
                h = f["height_m"] * np.exp(-(d / r) ** 2)
            else:
                d = gc_dist_km(xyz, f["at"], g.radius_km)
                if f["kind"] == "cone":
                    h = f["height_m"] * np.clip(1.0 - d / r, 0.0, 1.0) ** 1.5
                else:
                    h = f["height_m"] * np.exp(-((d - r) / (0.25 * r)) ** 2) + f["floor_m"] * (1.0 - smootherstep(d / r))
            add[on] += h
            near[on] = np.maximum(near[on], np.exp(-(d / (r + 3.0)) ** 2))
    return add, near


def clamp_plateaus(z, sea, halo, plateau_id, plateaus: list):
    """Hidden plateaus keep every sea cell at least HIDDEN_MAX_M (−550 m) deep."""
    z = np.asarray(z, dtype=np.float64).copy()
    for k, p in enumerate(plateaus):
        if not p.get("islands", False):
            m = (np.asarray(plateau_id) == k) & sea & ~halo
            z[m] = np.minimum(z[m], PL.HIDDEN_MAX_M)
    return z


def vents(g, z, sea, halo, potential, near, d_div_km, sediment_m, bottom_c, seed: int, V=VENT):
    """(per-cell vent strength 0–1, vents). A sea cell deeper than min_depth_m holds a vent with probability
    rate × potential × (0.5 + nearness to a volcanic feature); on a ridge axis (d_div_km < ridge_km) vents are black
    smokers, near a feature black (60 %) or white smokers, elsewhere diffuse; cold seeps sit in thick sediment where
    the potential is low. Temperature and flow class (0–2) per vent; the strength spreads ≈ reach_km around each."""
    depth = -np.asarray(z, dtype=np.float64)
    cand = sea & ~halo & (depth > V["min_depth_m"])
    pot = np.clip(np.asarray(potential, dtype=np.float64), 0.0, 1.0)
    near = np.asarray(near, dtype=np.float64)
    u = [_u(g.ids, seed * 8 + s) for s in range(5)]
    hot = cand & (u[0] < V["rate"] * pot * (0.5 + near))
    seep = cand & ~hot & (pot < 0.2) & (np.asarray(sediment_m) > 1000.0) & (u[1] < V["seep_rate"])
    i = np.flatnonzero(hot | seep)
    ridge = np.asarray(d_div_km)[i] < V["ridge_km"]
    close = near[i] > 0.3
    t = np.where(seep[i], 3, np.where(ridge | (close & (u[2][i] < 0.6)), 0, np.where(close, 1, 2))).astype(np.int8)
    lo, span = TEMP_C
    temp = np.where(t == 3, np.asarray(bottom_c, dtype=np.float64)[i], 0.0) + lo[t] + span[t] * u[3][i]
    flow = np.minimum((3.0 * u[4][i] * np.sqrt(pot[i])).astype(np.int8), 2).astype(np.int8)
    strength = np.zeros(g.n)
    if len(i):
        reach = V["reach_km"] / g.radius_km
        d, k = cKDTree(g.xyz[i]).query(g.xyz, distance_upper_bound=3.0 * reach, workers=max(1, h3par.workers()))
        ok = np.isfinite(d) & sea
        w = 0.5 + 0.25 * flow
        strength[ok] = np.clip(w[k[ok]] * np.exp(-(d[ok] / reach) ** 2), 0.0, 1.0)
    out = {"cell": g.ids[i].astype(np.uint64), "lat": g.lat[i].astype(np.float64), "lon": g.lon[i].astype(np.float64),
           "type": t, "temp_c": temp.astype(np.float32), "flow": flow, "mineral": t.copy()}
    return strength, out


def fields(pv: dict, sea, halo, cut, fan, fill, volc, strength) -> dict:
    """The refined sea-floor fields (0 / none on land): the world cell's values; sea over a world land cell (a refined
    coast) starts as terrigenous sediment; canyons scour (their fans and pond fill add sediment) and read as
    terrigenous; fine volcanic features read as volcanic floor; vent fields over the vents (sulfides, no sediment,
    warmer water)."""
    sea = np.asarray(sea, bool)
    t = np.asarray(pv["seabed_type"]).astype(np.int8).copy()
    m = np.asarray(pv["seabed_mineral"]).astype(np.int8).copy()
    sed = np.asarray(pv["sediment_m"], dtype=np.float64).copy()
    bt = np.asarray(pv["bottom_temp_c"], dtype=np.float64).copy()
    orphan = sea & (t == SB.SB_NONE)
    t[orphan] = SB.SB_TERRIGENOUS
    sed[orphan] = 300.0
    bt[orphan] = np.maximum(np.asarray(pv["T_mean"], dtype=np.float64)[orphan], -1.8)
    cut, fan, fill = (np.asarray(a, dtype=np.float64) for a in (cut, fan, fill))
    volc, strength = np.asarray(volc, dtype=np.float64), np.asarray(strength, dtype=np.float64)
    sed = np.where(cut > 0, 0.3 * sed, sed) + fan + fill
    t[(cut > 50.0) | (fan > 1.0)] = SB.SB_TERRIGENOUS
    t[np.abs(volc) > 50.0] = SB.SB_VOLCANIC
    vf = strength > 0.5
    t[vf] = SB.SB_VENTS
    m[vf] = SB.MI_SULFIDES
    sed[vf] = 0.0
    bt = bt + 20.0 * strength
    land = ~sea
    t[land], m[land] = SB.SB_NONE, SB.MI_NONE
    z32 = lambda a: np.where(land, 0.0, a).astype(np.float32)
    return {"seabed_type": t, "seabed_mineral": m, "bottom_temp_c": z32(bt), "sediment_m": z32(sed),
            "vent": z32(strength), "canyon_m": z32(cut), "fan_m": z32(fan)}


PV_KEYS = ("vent_potential", "seabed_type", "seabed_mineral", "bottom_temp_c", "sediment_m", "d_div_km", "T_mean")


def run(g, z, sea, halo, parent, world_a, plateaus: list, seed: int):
    """The whole pass: (heights, refined sea-floor fields incl. plateau_id, vents)."""
    sea, halo = np.asarray(sea, bool), np.asarray(halo, bool)
    pv = {k: np.asarray(world_a[k])[parent] for k in PV_KEYS}
    plateau_id = (PL.cell_ids(g.xyz, plateaus, seed, g.radius_km) if plateaus
                  else np.full(g.n, -1, np.int16))
    z0 = np.asarray(z, dtype=np.float64)
    zc, cut, fan = canyons(g, z0, sea, halo)
    zp, fill = ponds(g, zc, sea, halo, pv["sediment_m"])
    add, near = volcanic(g, sea, halo, plateau_id, plateaus, seed)
    zv = np.where(sea & ~halo, np.minimum(zp + add, np.maximum(zp, SEA_TOP_M)), zp)
    zn = clamp_plateaus(zv, sea, halo, plateau_id, plateaus)
    strength, vt = vents(g, zn, sea, halo, pv["vent_potential"], near, pv["d_div_km"], pv["sediment_m"],
                         pv["bottom_temp_c"], seed)
    f = fields(pv, sea, halo, cut, fan, fill, zv - zp, strength)
    f["plateau_id"] = plateau_id.astype(np.int16)
    return zn, f, vt


VENT_RGB = [(230, 40, 30), (245, 245, 245), (250, 160, 40), (40, 220, 230)]   # by VENT_TYPES


def preview(area_dir: Path, path: Path, px: int = 1400) -> Path:
    """A shaded depth image of a refined area (its box, north up) with its vents as dots — black smokers red, white
    smokers white, diffuse orange, cold seeps cyan: a first look at a sea floor before the viewer shows it."""
    from PIL import Image, ImageDraw
    area_dir, path = Path(area_dir), Path(path)
    with np.load(area_dir / "cells.npz") as a:
        inner = ~a["halo"].astype(bool)
        lat, lon = a["g_lat"][inner], a["g_lon"][inner]
        z, sea = a["elevation_eroded_m"][inner].astype(np.float64), a["ocean"][inner].astype(bool)
    c0 = float(np.degrees(np.arctan2(np.sin(np.radians(lon)).mean(), np.cos(np.radians(lon)).mean())))
    x = (lon - c0 + 180.0) % 360.0 - 180.0
    k = float(np.cos(np.radians(lat.mean())))
    W = int(px)
    H = max(2, int(round(W * (lat.max() - lat.min()) / max((x.max() - x.min()) * k, 1e-9))))
    GX, GY = np.meshgrid(np.linspace(x.min(), x.max(), W), np.linspace(lat.max(), lat.min(), H))
    pts = np.column_stack([x * k, lat])
    tree = cKDTree(pts)
    dist, j = tree.query(np.column_stack([GX.ravel() * k, GY.ravel()]))
    spacing = float(np.median(tree.query(pts, k=2)[0][:, 1]))
    Z, S = z[j].reshape(H, W), sea[j].reshape(H, W)
    px_km = (lat.max() - lat.min()) * 111.2 / H
    gy, gx = np.gradient(Z, px_km)
    shade = np.clip(1.0 - 0.015 * (gx - gy), 0.45, 1.35)[..., None]
    t = np.clip(-Z / 6000.0, 0.0, 1.0)[..., None]
    rgb = np.where(S[..., None], (1.0 - t) * np.array([150.0, 210.0, 235.0]) + t * np.array([10.0, 30.0, 80.0]),
                   np.array([125.0, 140.0, 95.0])) * shade
    rgb[(dist > 2.0 * spacing).reshape(H, W)] = 40.0
    im = Image.fromarray(np.clip(rgb, 0, 255).astype(np.uint8))
    if (area_dir / "vents.npz").exists():
        d = ImageDraw.Draw(im)
        with np.load(area_dir / "vents.npz") as v:
            vx = ((v["lon"] - c0 + 180.0) % 360.0 - 180.0 - x.min()) / max(x.max() - x.min(), 1e-9) * (W - 1)
            vy = (lat.max() - v["lat"]) / max(lat.max() - lat.min(), 1e-9) * (H - 1)
            for a_, b_, ty in zip(vx.tolist(), vy.tolist(), v["type"].tolist()):
                d.ellipse([a_ - 3, b_ - 3, a_ + 3, b_ + 3], fill=VENT_RGB[int(ty)], outline=(0, 0, 0))
    path.parent.mkdir(parents=True, exist_ok=True)
    im.save(path, quality=90)
    return path


def preview_region(root: Path, res: int, region_id: str, regions_path: Path | None = None,
                   regions_root: Path | None = None, out_dir: Path | None = None) -> Path:
    """preview() of a region's refined area in every world that has it: <out_dir>/<region>-<era>.jpg."""
    import refine as RF
    root = Path(root)
    regions = RF.load_regions(regions_path or root / "places" / "regions.json")
    if not any(RF.region_key(r) == region_id for r in regions):
        raise SystemExit(f"no region {region_id} in {regions_path or 'places/regions.json'}")
    fine = res + RF.FINER
    h = RF.region_areas(regions, RF.area_cells(regions, fine), fine)[region_id]["area"]
    rroot = Path(regions_root or root / "out" / f"r{res}" / "regions")
    out_dir = Path(out_dir or root / "previews" / f"r{res}" / "seafloor")
    made = []
    for s in sorted(p for p in rroot.glob("*-m*") if p.is_dir()):
        if h and (s / h / "cells.npz").exists():
            era = json.loads((s / "era.json").read_text())["era"] if (s / "era.json").exists() else s.name
            made.append(preview(s / h, out_dir / f"{region_id}-{era}.jpg"))
    if not made:
        raise SystemExit(f"{region_id} is not refined at res {res}: run mapgen.py refine --res {res}")
    return out_dir