"""River valleys for the deep zoom: a graded water level per river cell, a valley style from
landform, ground, rock and climate, and the valley cut into procedural heights. Procedural (`idea`), deterministic,
continuous across tiles. Lives outside mapgen/ so changing it never invalidates the build cache."""
from __future__ import annotations

import math
import threading

import numpy as np
from scipy.spatial import cKDTree

from mapgen.noise import value_noise

THETA = 0.45            # graded slope ∝ Q^-θ (discharge stands in for drainage area)
KS = 2.0                # m/km at Q = 1 km³/yr on a craton in average rock
KNICK_M_KM = 15.0       # steepest drop below lakes and steep reaches: a cataract, not a cliff
MAX_REACH_KM = 40.0
FLOOR_MAX_KM = 15.0     # the widest valley floor (half-width beyond the channel)
FADE_M = 2000.0         # a valley narrower than 2 px is lifted by up to this much: it fades in, never pops
CHUNK_KM = 100.0        # spread-out queries (long profiles) are answered in compact chunks
AGE_KS = {"craton (3g era)": 1.0, "pre-Lightening orogen": 1.5, "post-Lightening orogen": 3.0, "rift": 1.5,
          "Lightening basalt province": 2.0, "collapse scar": 1.5, "overshoot volcano": 3.0, "oceanic": 1.0}
ROCK = {"granite/gneiss": 1.6, "metamorphic": 1.6, "basalt": 1.8, "andesite": 1.4, "limestone": 1.3,
        "sandstone/shale": 0.7, "oceanic basalt": 1.5}                  # hardness: graded steepness and caps
CAP_M = {"plain": 150, "hills": 400, "mountains": 1200, "plateau": 1500, "rift valley": 600, "escarpment": 800,
         "volcanic arc": 900, "volcanic massif": 900, "basalt plateau": 1200, "dunes": 60, "badlands": 300, "ocean": 0}
STYLE = {"plain": (15, 40, 1.0), "hills": (5, 150, 0.5), "mountains": (1.5, 600, 0.0), "plateau": (1.2, 1500, 0.0),
         "rift valley": (5, 300, 0.3), "escarpment": (1.5, 800, 0.0), "volcanic arc": (1.5, 600, 0.0),
         "volcanic massif": (1.5, 600, 0.0), "basalt plateau": (1.2, 1500, 0.0), "dunes": (8, 60, 0.5),
         "badlands": (3, 400, 0.2), "ocean": (15, 40, 1.0)}          # landform → (floor × half-width, wall m/km, meander)
GROUND = {"floodplain": (30, 20, 1.5), "delta": (30, 20, 1.5), "wetland": (20, 20, 1.2), "bog": (20, 20, 1.2)}
DETAIL_M = {"ocean": 300, "plain": 50, "hills": 240, "mountains": 900, "plateau": 120, "rift valley": 300,
            "escarpment": 500, "volcanic arc": 700, "volcanic massif": 700, "basalt plateau": 120, "dunes": 60,
            "badlands": 240}   # ≈ 2 × the procedural detail's amplitude (tiles.DEEP): how far above the cell it reaches
WALL_ROCK = {"granite/gneiss": 1.5, "metamorphic": 1.5, "basalt": 1.5, "oceanic basalt": 1.5, "limestone": 1.4,
             "andesite": 1.2, "sandstone/shale": 0.6}
STATE = ("cap", "level", "seg_a", "seg_b", "a_xyz", "b_xyz", "level_a", "level_b", "seg_len", "half_w", "floor",
         "wall", "meander", "amp", "width", "reach")      # what a RiverNet computes (serve cache: saved, not recomputed)


def half_width_km(q_km3_yr):
    """Channel half-width: 4·√Q(m³/s) metres (hydraulic geometry)."""
    return 0.004 * np.sqrt(np.asarray(q_km3_yr, dtype=np.float64) * 31.7)


DRAW_MIN_PX = 0.02                        # streams narrower than this share of a pixel are not drawn (refined
ALWAYS_HW = float(half_width_km(2.0))     # 0.2 km³/yr streams appear from ≈ 0.5 km/px); world rivers (≥ 2 km³/yr)
                                          # are drawn at every zoom


def valley_style(landform: str, ground: str, rock: str, rain_mm: float):
    """(floor × half-width, wall m/km, meander factor) for a river cell."""
    floor, wall, meander = GROUND.get(ground) or STYLE.get(landform, (5, 150, 0.5))
    k = WALL_ROCK.get(rock, 1.0)
    dry = 1.4 if rain_mm < 500 else 0.8 if rain_mm > 1500 else 1.0
    return floor / k, wall * k * dry, meander


def valley_surface(d_km, level, half_w, floor, wall, rough=0.0):
    """Height (m) of a valley at distance d from its centreline: the water level in the channel, a floor 1–3 m above
    it, then walls rising at `wall` m/km. `rough` (the terrain's own procedural detail, m) roughens the walls,
    fading in over the first 0.5 km above the floor, so they look eroded rather than planar."""
    d = np.asarray(d_km, dtype=np.float64)
    fl = np.clip((d - half_w) / np.maximum(floor, 1e-9), 0.0, 1.0)
    up = np.maximum(d - half_w - floor, 0.0)
    wall_m = np.maximum(0.0, wall * up + np.clip(up / 0.5, 0.0, 1.0) * rough)   # rough walls never dip below the floor
    return np.where(d <= half_w, level, level + 1.0 + 2.0 * fl + wall_m)


def _rownorm(p):
    """np.linalg.norm(p, axis=1) for (n, 3), without its per-call overhead: the same squares summed in the same order."""
    return np.sqrt(p[:, 0] * p[:, 0] + p[:, 1] * p[:, 1] + p[:, 2] * p[:, 2])


def _vnorm(v):
    """np.linalg.norm of one 3-vector (it is sqrt(v·v))."""
    return math.sqrt(v.dot(v))


def densify(a, b, t0, t1, step):
    """Unit vectors along a→b (a short great-circle segment) at the lattice t = k·step/|b−a|, k integer, t in [t0, t1]."""
    d = b - a
    L = _vnorm(d)
    if L == 0:
        return a[None, :]
    k = np.arange(np.ceil(t0 * L / step), np.floor(t1 * L / step) + 1)
    t = np.append(k * step / L, [] if t1 < 1 else [1.0])   # the segment's end joins the next one
    if len(t) == 0:
        return np.zeros((0, 3))
    p = a[None, :] * (1 - t[:, None]) + b[None, :] * t[:, None]
    return p / _rownorm(p)[:, None]


class RiverNet:
    def __init__(self, a: dict, legends: dict, radius_km: float, levels=None, ground=None):
        """levels: given water levels (a refined area); ground: the heights the valleys are cut from (default the
        cells' surface; a refined area passes its valley-shoulder heights, so walls reach up to them)."""
        self.R = float(radius_km)
        riv = np.asarray(a["river"]).astype(bool)
        recv = np.asarray(a["recv"]).astype(np.int64)
        ocean, lake = np.asarray(a["ocean"]).astype(bool), np.asarray(a["lake"]).astype(bool)
        z = np.asarray(a["z_surface_m"], dtype=np.float64)
        zf = np.asarray(a["z_filled_m"], dtype=np.float64)
        q = np.asarray(a["discharge_km3_yr"], dtype=np.float64)
        xyz = np.asarray(a["g_xyz"], dtype=np.float64)
        pick = lambda leg, key, table, default: np.array([table.get(n, default) for n in legends[leg]])[np.asarray(a[key])]
        rock = pick("lithology", "lithology", ROCK, 1.0)
        ks = KS * pick("age_class", "age_class", AGE_KS, 1.0) * rock
        self.cap = pick("landform", "landform", CAP_M, 400) * rock / 1.6
        n = len(riv)
        r = np.where(riv)[0]
        dist = np.zeros(n)
        dist[r] = np.linalg.norm(xyz[r] - xyz[recv[r]], axis=1) * self.R
        if levels is None:
            lev = np.full(n, np.nan)
            lev[lake] = zf[lake]
            lev[ocean] = 0.0
            for i in r[np.argsort(zf[r], kind="stable")]:           # upstream: graded on the receiver's level, capped
                j = recv[i]
                b = lev[j] if np.isfinite(lev[j]) else z[j]
                lev[i] = min(z[i], max(b + ks[i] * q[i] ** -THETA * dist[i], z[i] - self.cap[i]))
            src = np.where((riv | lake) & ~ocean & (recv != np.arange(n)))[0]
            d_src = np.linalg.norm(xyz[src] - xyz[recv[src]], axis=1) * self.R
            for i, di in zip(src[np.argsort(-zf[src], kind="stable")], d_src[np.argsort(-zf[src], kind="stable")]):
                j = recv[i]                                         # downstream: drops limited to a cataract
                if riv[j]:
                    lev[j] = max(lev[j], min(z[j], lev[i] - KNICK_M_KM * di))
            for i in r[np.argsort(-zf[r], kind="stable")]:          # never rising downstream
                j = recv[i]
                if riv[j] and lev[j] > lev[i]:
                    lev[j] = lev[i]
        else:   # given water levels (a refined area: erosion has already cut the valleys; receivers keep theirs)
            lev = np.asarray(levels, dtype=np.float64).copy()
            lev[lake] = zf[lake]
            lev[ocean] = 0.0
        self.level = lev
        out = np.where(lake & riv[recv] & (recv != np.arange(n)))[0]     # lake outlets flow on to their river
        dist[out] = np.linalg.norm(xyz[out] - xyz[recv[out]], axis=1) * self.R
        a_cell = np.concatenate([r, out])
        b_cell = recv[a_cell]
        style_cell = np.concatenate([r, recv[out]])                      # an outlet looks like the river it feeds
        self.seg_a, self.seg_b = a_cell, b_cell
        self.a_xyz, self.b_xyz = xyz[a_cell], xyz[b_cell]
        self.level_a = np.where(lake[a_cell], zf[a_cell], lev[a_cell])
        self.level_b = np.where(ocean[b_cell], 0.0, np.where(lake[b_cell], zf[b_cell], lev[b_cell]))
        self.seg_len = dist[a_cell]
        self.half_w = half_width_km(q[style_cell])
        names = {k: np.asarray(legends[k], dtype=object) for k in ("landform", "ground", "lithology")}
        st = np.array([valley_style(names["landform"][a["landform"][i]], names["ground"][a["ground"][i]],
                                    names["lithology"][a["lithology"][i]], float(a["P_ann"][i])) for i in style_cell]).reshape(-1, 3)
        self.floor = np.minimum(st[:, 0] * self.half_w, FLOOR_MAX_KM)
        self.wall, self.meander = st[:, 1], st[:, 2]
        self.amp = 2.5 * 2 * self.half_w * self.meander                     # meander swing (km)
        zg = z if ground is None else np.asarray(ground, dtype=np.float64)
        depth = np.maximum(np.maximum(zg[a_cell], self.level_a) - self.level_a, 0.0)
        detail = pick("landform", "landform", DETAIL_M, 300)[style_cell]
        self.width = 2 * (self.half_w + self.floor + depth / self.wall)            # as seen at the cell's mean ground
        self.reach = self.half_w + self.floor + np.minimum(MAX_REACH_KM, (depth + detail) / self.wall)   # + detail relief
        self._tree, self._tree_lock = None, threading.Lock()   # built on first use (see tree)
        self.max_extent = float(np.max(self.seg_len / 2 + self.reach + self.amp)) if len(a_cell) else 0.0

    @property
    def tree(self):
        """KD-tree over segment midpoints (unit vectors), built on first use; None without segments."""
        if self._tree is None and len(self.seg_a):
            with self._tree_lock:
                if self._tree is None:
                    mid = np.asarray(self.a_xyz, dtype=np.float64) + np.asarray(self.b_xyz, dtype=np.float64)
                    self._tree = cKDTree(mid / np.linalg.norm(mid, axis=1, keepdims=True))
        return self._tree

    def state(self):
        """(arrays, meta) that from_state turns back into the same net without recomputing it."""
        return {k: getattr(self, k) for k in STATE}, {"R": self.R, "max_extent": self.max_extent}

    @classmethod
    def from_state(cls, arrays: dict, meta: dict) -> "RiverNet":
        net = cls.__new__(cls)
        for k in STATE:
            setattr(net, k, np.asarray(arrays[k]))      # memory maps as plain arrays: same data, cheap indexing
        net.R, net.max_extent = float(meta["R"]), float(meta["max_extent"])
        net._tree, net._tree_lock = None, threading.Lock()
        return net

    # --- geometry ---------------------------------------------------------------------------------------------
    def candidates(self, center, radius_km):
        """Segments whose valley can reach any point within radius_km of a unit vector."""
        if self.tree is None:
            return []
        idx = np.asarray(self.tree.query_ball_point(center, 2 * np.sin(min(np.pi, (radius_km + self.max_extent) / self.R) / 2)),
                         dtype=np.int64)
        if not len(idx):
            return []
        a, ab = self.a_xyz[idx], self.b_xyz[idx] - self.a_xyz[idx]
        tc = np.clip(np.einsum("ij,ij->i", center - a, ab) / np.maximum(np.einsum("ij,ij->i", ab, ab), 1e-30), 0.0, 1.0)
        d = np.linalg.norm(center - (a + tc[:, None] * ab), axis=1) * self.R
        return idx[d <= radius_km + self.reach[idx] + self.amp[idx] + 1.0].tolist()

    def segment_points(self, s, center, within_km, spacing_km):
        """Centreline of segment s near a unit vector: its own global lattice (so any query gets the same points),
        meandered. Returns (unit vectors, t along the segment)."""
        a, b = self.a_xyz[s], self.b_xyz[s]
        ab = b - a
        L = max(_vnorm(ab), 1e-12)
        tc = float(np.dot(center - a, ab) / max(np.dot(ab, ab), 1e-30))
        span = (within_km + self.amp[s]) / self.R / L
        t0, t1 = max(0.0, tc - span), min(1.0, tc + span)
        if t0 > t1:
            return np.zeros((0, 3)), np.zeros(0)
        p = densify(a, b, t0, t1, max(spacing_km, 0.0005) / self.R)
        if len(p) == 0:
            return p, np.zeros(0)
        t = np.clip(np.dot(p - a, ab) / max(np.dot(ab, ab), 1e-30), 0.0, 1.0)
        if self.amp[s] > 0 and self.seg_len[s] > 0:
            side = np.empty_like(p)                    # np.cross(ab, p), the same products and differences
            side[:, 0] = ab[1] * p[:, 2] - ab[2] * p[:, 1]
            side[:, 1] = ab[2] * p[:, 0] - ab[0] * p[:, 2]
            side[:, 2] = ab[0] * p[:, 1] - ab[1] * p[:, 0]
            side /= np.maximum(_rownorm(side), 1e-30)[:, None]
            lam = 11 * 2 * self.half_w[s]
            off = self.amp[s] * value_noise(p * (self.R / lam), 4242) * np.sin(np.pi * t)
            p = p + side * (off / self.R)[:, None]
        return p / _rownorm(p)[:, None], t

    def _chunks(self, xyz):
        """Split query points into compact groups (halving along their widest axis): (indices, centre, radius km)."""
        out, stack = [], [np.arange(len(xyz))]
        while stack:
            ix = stack.pop()
            c = xyz[ix].mean(axis=0)
            c = c / max(np.linalg.norm(c), 1e-12) if np.linalg.norm(c) > 1e-9 else xyz[ix[0]]
            r = float(np.max(np.linalg.norm(xyz[ix] - c, axis=1))) * self.R
            if r <= CHUNK_KM or len(ix) <= 16:
                out.append((ix, c, 2 * self.R * np.arcsin(min(1.0, r / (2 * self.R)))))   # chord → arc
            else:
                p = xyz[ix]
                o = np.argsort(p[:, int(np.argmax(p.max(axis=0) - p.min(axis=0)))], kind="stable")
                stack += [ix[o[: len(ix) // 2]], ix[o[len(ix) // 2:]]]
        return out

    def parts(self, xyz, px_km, spacing_km=None):
        """Per query point, the lowest valley surface over every valley that reaches it, in parts:
        (base = level + floor rise + fade lift, wall rise, roughness share). base is +inf where no valley reaches.
        No lower valley undercuts a reach's banks: each reach's floor is a lower bound, falling away at its wall slope
        beyond it (a reach doubling back below itself leaves a terrace, not dry pits below the water beside it)."""
        n = len(xyz)
        base, wall_up, rough_f = np.full(n, np.inf), np.zeros(n), np.zeros(n)
        best, bank = np.full(n, np.inf), np.full(n, -np.inf)
        spacing = spacing_km or px_km
        for ix, c, r in self._chunks(xyz):
            q = xyz[ix]
            for s in self.candidates(c, r):
                alpha = float(np.clip(self.width[s] / px_km - 1.0, 0.0, 1.0))
                if alpha <= 0:
                    continue                                   # narrower than a pixel: no valley to see
                p, t = self.segment_points(s, c, r + self.reach[s], spacing)
                if len(p) == 0:
                    continue
                d, k = cKDTree(p).query(q, distance_upper_bound=self.reach[s] / self.R)
                hit = k < len(p)
                if not hit.any():
                    continue
                dk = d[hit] * self.R
                lv = self.level_a[s] + (self.level_b[s] - self.level_a[s]) * t[k[hit]]
                fl = np.clip((dk - self.half_w[s]) / max(self.floor[s], 1e-9), 0.0, 1.0)
                up = np.maximum(dk - self.half_w[s] - self.floor[s], 0.0)
                b = lv + 1.0 + 2.0 * fl + (1.0 - alpha) * FADE_M
                w = self.wall[s] * up
                j = ix[hit]
                better = b + w < best[j]
                jb = j[better]
                best[jb], base[jb], wall_up[jb] = (b + w)[better], b[better], w[better]
                rough_f[jb] = np.clip(up / 0.5, 0.0, 1.0)[better]
                np.maximum.at(bank, j, b - w)
        low = bank > best
        base[low], wall_up[low], rough_f[low] = bank[low], 0.0, 0.0
        return base, wall_up, rough_f

    def channel(self, xyz, px_km):
        """Channel water per query point (true width), channel as drawn (≥ 0.6 px), and the water level there."""
        n = len(xyz)
        ch, dr, lev = np.zeros(n, bool), np.zeros(n, bool), np.full(n, np.nan)
        for ix, c, r in self._chunks(xyz):
            pts, lv, hw = [], [], []
            for s in self.candidates(c, r):
                if self.half_w[s] < min(DRAW_MIN_PX * px_km, ALWAYS_HW):
                    continue                                   # a stream far below a pixel wide: not drawn yet
                reach = max(self.half_w[s], 0.6 * px_km)
                a, ab = self.a_xyz[s], self.b_xyz[s] - self.a_xyz[s]
                tc = np.clip(np.dot(c - a, ab) / max(np.dot(ab, ab), 1e-30), 0.0, 1.0)
                if _vnorm(c - (a + tc * ab)) * self.R > r + reach + self.amp[s] + 1.0:
                    continue
                p, t = self.segment_points(s, c, r + reach, px_km / 2)
                keep = _rownorm(p - c) * self.R <= r + reach + px_km   # only points that reach the chunk
                p, t = p[keep], t[keep]
                if len(p):
                    pts.append(p)
                    lv.append(self.level_a[s] + (self.level_b[s] - self.level_a[s]) * t)
                    hw.append(np.full(len(p), self.half_w[s]))
            if not pts:
                continue
            P, LV, HW = np.concatenate(pts), np.concatenate(lv), np.concatenate(hw)
            bound = max(float(HW.max()), 0.6 * px_km) / self.R
            d, k = cKDTree(P).query(xyz[ix], distance_upper_bound=bound)
            hit = k < len(P)
            kk = np.where(hit, k, 0)
            dk = np.where(hit, d * self.R, np.inf)
            ch[ix] = hit & (dk <= HW[kk])
            dr[ix] = hit & (dk <= np.maximum(HW[kk], 0.6 * px_km))
            lev[ix] = np.where(hit, LV[kk], np.nan)
        return ch, dr, lev

    def valleys(self, xyz, px_km, rough=None, spacing_km=None, with_base=False):
        """Valley surface V (m; +inf where none), channel and drawn channel at query points (+ base, see parts)."""
        base, wall_up, f = self.parts(xyz, px_km, spacing_km)
        rg = 0.0 if rough is None else np.asarray(rough, dtype=np.float64)
        V = base + np.maximum(0.0, wall_up + f * rg)
        ch, dr, lev = self.channel(xyz, px_km)
        V = np.where(ch, lev, V)
        return (V, ch, dr, base) if with_base else (V, ch, dr)