"""Stage `seabed`: sea-floor data for every ocean cell — vent potential,
sea-floor type, minerals, bottom temperature, sediment. Land cells get 0 / "—" (the fields describe the sea floor)."""
from __future__ import annotations

import numpy as np
from scipy.spatial import cKDTree

from . import plateaus as PL
from .config import params
from .elevation import hotspot_track
from .graph import distance_to
from .noise import fbm
from .sphere import latlon_to_xyz

(SB_NONE, SB_TERRIGENOUS, SB_CARBONATE, SB_CLAY, SB_BASALT, SB_VOLCANIC, SB_CONTINENTAL, SB_VENTS,
 SB_TRENCH) = range(9)
SEABED_NAMES = ["—", "terrigenous sediment", "carbonate ooze", "pelagic clay", "young basalt", "volcanic",
                "continental / plateau", "vent field", "trench"]
MI_NONE, MI_SULFIDES, MI_NODULES, MI_COBALT, MI_PHOSPHORITE = range(5)
MINERAL_NAMES = ["none", "polymetallic sulfides", "manganese nodules", "cobalt crusts", "phosphorite"]
DEFAULTS = {"ridge_km": 100.0, "back_arc_km": 300.0, "back_arc_width_km": 150.0, "back_arc": 0.6,
            "hotspot_km": 150.0, "hotspot_spacing_km": 150.0, "terrigenous_km": 300.0, "carbonate_depth_m": 4500.0,
            "carbonate_t_c": 10.0, "young_myr": 10.0, "trench_km": 80.0, "trench_depth_m": 5000.0,
            "vent_field": 0.8, "sulfides": 0.7, "volcanic": 0.3, "thermocline_m": 800.0}


def _cells_near(tree, radius_km, p, reach_km):
    return np.asarray(tree.query_ball_point(p, 2.0 * np.sin(min(np.pi, reach_km / radius_km) / 2.0)), dtype=np.int64)


def _bump(v, g, tree, p, r_km, amp):
    i = _cells_near(tree, g.radius_km, p, 3.0 * r_km)
    if len(i):
        d = g.radius_km * np.arccos(np.clip(g.xyz[i] @ p, -1.0, 1.0))
        v[i] = np.maximum(v[i], amp * np.exp(-(d / r_km) ** 2))


def volcanic_potential(g, tect: dict, vel, seed: int, P: dict):
    """0–1 per cell from hotspot chains (strongest at the active end) and plateau volcanic fields."""
    tree, v = cKDTree(g.xyz), np.zeros(g.n)
    for h in tect.get("hotspot", []):
        for p, s in hotspot_track(g, h, vel, P["hotspot_spacing_km"]):
            _bump(v, g, tree, p, P["hotspot_km"], 1.0 - s / h["length_km"])
    for pl in tect.get("plateau", []):
        f, vent = PL.features(pl, seed, g.radius_km), float(pl.get("vent", 1.0))
        _bump(v, g, tree, latlon_to_xyz(*f["hotspot"]), 200.0, vent)
        for c in f["cones"]:
            _bump(v, g, tree, latlon_to_xyz(c["lat"], c["lon"]), c["radius_km"] + 30.0, 0.8 * vent)
        for c in f["calderas"]:
            _bump(v, g, tree, latlon_to_xyz(c["lat"], c["lon"]), c["radius_km"] + 30.0, vent)
    return v


def run(ctx) -> dict:
    g = ctx.grid
    P = params(ctx.cfg, "seabed", DEFAULTS)
    z, ocean, cont, age, d_div, d_over, d_sub, t_mean, vel = ctx.need(
        "elevation_eroded_m", "ocean", "continental", "ocean_age_myr", "d_div_km", "d_over_km", "d_sub_km", "T_mean",
        "vel")
    z, ocean, cont = np.asarray(z, dtype=np.float64), np.asarray(ocean, bool), np.asarray(cont, bool)
    age, t_mean = np.asarray(age, dtype=np.float64), np.asarray(t_mean, dtype=np.float64)
    plateau_id = np.asarray(ctx.data.get("plateau_id", np.full(g.n, -1)))
    depth = np.maximum(-z, 0.0)
    volc = volcanic_potential(g, ctx.tect, np.asarray(vel), ctx.seed, P)
    ridge = np.exp(-(np.asarray(d_div) / P["ridge_km"]) ** 2)
    back_arc = P["back_arc"] * np.exp(-((np.asarray(d_over) - P["back_arc_km"]) / P["back_arc_width_km"]) ** 2)
    cluster = np.clip(0.5 + 1.5 * fbm(g.xyz, ctx.seed + 301, 4, 30.0), 0.0, 1.0)    # vents come in fields
    vent = np.where(ocean, np.maximum.reduce([ridge, back_arc, volc]) * (0.5 + 0.5 * cluster), 0.0)
    d_land = distance_to(g, ~ocean) if (~ocean).any() else np.full(g.n, np.inf)

    t = np.full(g.n, SB_CLAY, np.int8)
    t[(depth < P["carbonate_depth_m"]) & (t_mean > P["carbonate_t_c"])] = SB_CARBONATE
    t[cont] = SB_CONTINENTAL
    t[d_land < P["terrigenous_km"]] = SB_TERRIGENOUS
    t[~cont & (age < P["young_myr"])] = SB_BASALT
    t[volc > P["volcanic"]] = SB_VOLCANIC
    t[(np.asarray(d_sub) < P["trench_km"]) & (depth > P["trench_depth_m"])] = SB_TRENCH
    t[vent > P["vent_field"]] = SB_VENTS
    t[~ocean] = SB_NONE

    m = np.zeros(g.n, np.int8)
    m[(t == SB_CLAY) & (age > 30.0) & (d_land > 1000.0) & (depth > 4000.0)] = MI_NODULES
    m[((t == SB_VOLCANIC) | (plateau_id >= 0)) & (depth >= 800.0) & (depth <= 2500.0)] = MI_COBALT
    m[cont & (depth < 500.0)] = MI_PHOSPHORITE
    m[vent > P["sulfides"]] = MI_SULFIDES
    m[~ocean] = MI_NONE

    deep = 1.0 + 3.0 * np.cos(np.radians(g.lat)) ** 2                          # ≈ 1 °C polar … 4 °C tropical
    w = np.clip((P["thermocline_m"] - depth) / P["thermocline_m"], 0.0, 1.0)   # shallow floors: towards the surface
    bottom = np.where(ocean, deep + w * (np.maximum(t_mean, -1.8) - deep), 0.0)
    terr = 2500.0 * np.exp(-d_land / 250.0)                                     # aprons off the continents
    sed = np.where(cont, 300.0 + terr, np.minimum(5.0 * age, 800.0) + terr)     # pelagic rain ≈ 5 m per Myr
    sed = np.where(ocean, sed, 0.0)
    return {"vent_potential": vent.astype(np.float32), "seabed_type": t, "seabed_mineral": m,
            "bottom_temp_c": bottom.astype(np.float32), "sediment_m": sed.astype(np.float32)}