"""Stage `plates`: grow plates from seeds, assign Euler motions, classify boundaries."""
from __future__ import annotations

import numpy as np
from scipy.spatial import cKDTree

from .config import params
from .graph import nearest_source
from .grid import edge_blocks
from .noise import fbm
from .pipeline import StageError
from .sphere import motion_to_omega, velocity

NONE, CONV, DIV, TRANS = 0, 1, 2, 3
DEFAULTS = {"land_cost": 0.25, "noise": 0.35, "noise_freq": 3.0, "warp_km": 1000.0, "warp_freq": 1.5,
            "min_rate_m_yr": 0.005}


def edge_convergence(g, vel):
    """Per directed edge s→d: closing speed (m/yr, >0 converging) and tangential slip speed."""
    t, src, dst = g.edge_tangents, g.src, g.dst
    along, tang = np.empty(len(dst)), np.empty(len(dst))
    for s in edge_blocks(len(dst)):                      # per edge: the same values, small temporaries
        rel = vel[dst[s]] - vel[src[s]]
        along[s] = np.sum(rel * t[s], axis=1)
        tang[s] = np.linalg.norm(rel - along[s][:, None] * t[s], axis=1)
    return -along, tang


def classify_boundaries(g, plate, vel, min_rate):
    conv, tang = edge_convergence(g, vel)
    s, d = g.src, g.dst
    b = plate[s] != plate[d]
    cnt = np.bincount(s[b], minlength=g.n)
    c_mean = np.bincount(s[b], weights=conv[b], minlength=g.n) / np.maximum(cnt, 1)
    t_mean = np.bincount(s[b], weights=tang[b], minlength=g.n) / np.maximum(cnt, 1)
    other = np.full(g.n, -1, np.int16)
    other[s[b]] = plate[d[b]]
    typ = np.zeros(g.n, np.int8)
    isb = cnt > 0
    typ[isb] = TRANS
    typ[isb & (c_mean > min_rate)] = CONV
    typ[isb & (c_mean < -min_rate)] = DIV
    rate = np.where(typ == CONV, c_mean, np.where(typ == DIV, -c_mean, np.where(isb, t_mean, 0.0)))
    return typ, rate.astype(np.float32), other


def _warp_labels(g, plate, seeds, P, seed, continental_kind):
    """Domain warp: each cell takes the label found at a noise-displaced point → meandering boundaries.
    Only same-kind swaps (cont↔cont, ocean↔ocean): ocean–continent margins stay where growth put them."""
    if P["warp_km"] <= 0:
        return plate
    amp = P["warp_km"] / g.radius_km
    disp = np.stack([fbm(g.xyz, seed + 101 + k, 3, P["warp_freq"]) for k in range(3)], axis=1)
    disp /= max(float(disp.std()), 1e-12)              # warp_km = RMS displacement per component
    q = g.xyz + amp * disp
    q /= np.linalg.norm(q, axis=1, keepdims=True)
    _, j = cKDTree(g.xyz).query(q)
    out = plate[j].astype(np.int16)
    keep = continental_kind[out] != continental_kind[plate]
    out[keep] = plate[keep]
    out[seeds] = np.arange(len(seeds), dtype=np.int16)
    return out


def run(ctx) -> dict:
    g = ctx.grid
    P = params(ctx.cfg, "plates", DEFAULTS)
    land, land_hint = ctx.need("sk_land", "m_land_hint")
    plates = ctx.tect["plate"]
    seeds = np.array([g.cell_index(*p["seed"]) for p in plates], dtype=np.int64)
    if len(np.unique(seeds)) < len(seeds):
        raise StageError("two plate seeds fall in the same cell; move one in tectonics.toml")
    hint = np.clip(land + 0.5 * land_hint, 0.0, 1.0)
    nz = fbm(g.xyz, ctx.seed + 11, octaves=4, freq=P["noise_freq"])
    mult = np.where(hint[g.dst] > 0.5, P["land_cost"], 1.0) * (1.0 + P["noise"] * nz[g.dst])
    _, src = nearest_source(g, seeds, g.edge_km * np.maximum(mult, 0.05))
    order = np.argsort(seeds)
    plate = order[np.searchsorted(seeds[order], src)].astype(np.int16)
    kinds = np.array([p["kind"] == "continental" for p in plates])
    plate = _warp_labels(g, plate, seeds, P, ctx.seed, kinds)
    omegas = np.array([motion_to_omega(*p["seed"], *p["motion"], g.radius_km) for p in plates])
    vel = velocity(g.xyz, omegas[plate], g.radius_km)
    btype, brate, bother = classify_boundaries(g, plate, vel, P["min_rate_m_yr"])
    return {"plate": plate, "vel": vel, "plate_continental": kinds,
            "bnd_type": btype, "bnd_rate": brate, "bnd_other": bother}