"""Era events: pure functions of positions and the event's config, so the era
builder (world cells) and the refinement (fine cells) apply the same event."""
from __future__ import annotations

import numpy as np
from scipy.spatial import cKDTree

from .fields import pressure
from .graph import components
from .noise import fbm, name_seed
from .pipeline import StageError
from .sphere import gc_dist_km, latlon_to_xyz
from .zones import smootherstep


def disintegrate(xyz, z, land, ev: dict, radius_km: float):
    """(new heights, z_ref): inside radius_km of the centre min(old, z_ref − depth · sqrt(1 − (d/r)²)); z_ref = mean
    ground height of the land in the ring 0.9–1.0 r before the event, None when no land lies in that ring (the bowl
    then hangs from 0 m). Material vanishes: no ejecta, no rim."""
    z = np.asarray(z, dtype=np.float64)
    d = gc_dist_km(np.asarray(xyz, dtype=np.float64), latlon_to_xyz(*ev["center"]), radius_km)
    r = float(ev["radius_km"])
    ring = np.asarray(land, bool) & (d >= 0.9 * r) & (d <= r)
    z_ref = float(np.mean(z[ring])) if ring.any() else None   # None: no land on its rim (the bowl hangs from 0 m)
    bowl = (0.0 if z_ref is None else z_ref) - ev["depth_m"] * np.sqrt(np.clip(1.0 - (d / r) ** 2, 0.0, 1.0))
    return np.where(d < r, np.minimum(z, bowl), z), z_ref


def volcano(xyz, z, ev: dict, radius_km: float):
    """Heights after a volcano (only raises): cone, shield or caldera; peak_m above the old ground (peak_mode "above")
    or as an absolute height ("absolute")."""
    z = np.asarray(z, dtype=np.float64)
    d = gc_dist_km(np.asarray(xyz, dtype=np.float64), latlon_to_xyz(*ev["center"]), radius_km)
    x = np.clip(d / ev["radius_km"], 0.0, 1.0)
    prof = {"cone": (1.0 - x) ** 1.5, "shield": (1.0 - x * x) ** 2,
            "caldera": np.maximum((1.0 - x) ** 1.5 - 0.6 * np.exp(-(x / 0.15) ** 2), 0.0)}[ev.get("shape", "cone")]
    peak = float(ev["peak_m"])
    new = z + peak * prof if ev.get("peak_mode", "above") == "above" else z + np.maximum(peak - z, 0.0) * prof
    return np.maximum(z, new)


def landmass(g, land, seed_latlon, event_name: str):
    """The connected land (grid cells) containing the seed point."""
    lab = components(g, land)
    i = g.cell_index(*seed_latlon)
    if lab[i] < 0:
        land = np.flatnonzero(np.asarray(land, bool))
        hint = ""
        if len(land):
            j = land[np.argmax(g.xyz[land] @ g.xyz[i])]
            hint = f"; the nearest land is at [{g.lat[j]:.1f}, {g.lon[j]:.1f}]"
        raise StageError(f"event {event_name}: its seed {list(seed_latlon)} lies in the sea in this era{hint} "
                         f"(move the seed in config/eras.toml)")
    return lab == lab[i]


def zone_weight(xyz, ev: dict, radius_km: float, seed: int, land_xyz=None):
    """0–1 per point: 1 inside, 0 outside; a linear ramp edge_km wide at the edge (sharp), or the smooth profile.
    circle: centre + radius_km. landmass: within reach of the landmass cells land_xyz (unit vectors), the reach
    varying between reach_km by low-frequency noise."""
    xyz = np.asarray(xyz, dtype=np.float64)
    if ev["shape"] == "circle":
        d = gc_dist_km(xyz, latlon_to_xyz(*ev["center"]), radius_km)
        reach = np.full(len(xyz), float(ev["radius_km"]))
    else:
        chord, _ = cKDTree(np.asarray(land_xyz, dtype=np.float64)).query(xyz)
        d = 2.0 * np.arcsin(np.clip(chord / 2.0, 0.0, 1.0)) * radius_km
        lo, hi = ev["reach_km"]
        reach = lo + (hi - lo) * np.clip(0.5 + 1.5 * fbm(xyz, name_seed(seed, ev["name"]), 3, 1.5), 0.0, 1.0)
    if "edge_km" in ev:
        return np.clip((reach - d) / ev["edge_km"], 0.0, 1.0)
    return 1.0 - smootherstep(d / reach)


def apply_zone(fields: dict, w, ev: dict, z_surface_m, scale_height_m: float) -> dict:
    """Fields inside a zone event: absolute values blended by w (gravity_g, o2_fraction, fire_reactivity; pressure_bar
    is the sea-level value, falling with height as elsewhere); po2_bar follows. Never height. Cells with w = 0 keep
    their values exactly."""
    w = np.asarray(w, dtype=np.float64)
    decay = pressure(1.0, z_surface_m, scale_height_m)
    new = {}
    for k, v in ev["fields"].items():
        base = np.asarray(fields[k], dtype=np.float64)
        new[k] = (base / decay * (1.0 - w) + v * w) * decay if k == "pressure_bar" else base * (1.0 - w) + v * w
    o2 = new.get("o2_fraction", np.asarray(fields["o2_fraction"], dtype=np.float64))
    new["po2_bar"] = o2 * new.get("pressure_bar", np.asarray(fields["pressure_bar"], dtype=np.float64))
    return {k: np.where(w > 0, v, fields[k]).astype(np.asarray(fields[k]).dtype) for k, v in new.items()}