"""Stage `crust`: continental vs oceanic crust, age classes, oceanic age."""
from __future__ import annotations

import numpy as np

from .config import params
from . import plateaus as PL
from .graph import distance_to, smooth_km
from .noise import fbm, name_seed, ridged
from .plates import CONV, DIV
from .sphere import azimuth_deg, gc_dist_km, latlon_to_xyz

OCEANIC, CRATON, PRE_OROGEN, POST_OROGEN, RIFT, LIP, SCAR, VOLCANO = range(8)
AGE_NAMES = ["oceanic", "craton (3g era)", "pre-Lightening orogen", "post-Lightening orogen",
             "rift", "Lightening basalt province", "collapse scar", "overshoot volcano"]
DEFAULTS = {"continental_threshold": 0.25, "shelf_km": 450.0, "edge_noise": 0.5, "edge_noise_freq": 4.0, "edge_noise_gain": 0.65, "orogen_width_km": 700.0, "rift_width_km": 200.0,
            "pre_orogen_threshold": 0.72, "half_spreading_km_myr": 30.0, "max_ocean_age_myr": 200.0,
            "scar_inner": 0.4, "scar_outer": 1.6, "scar_halfwidth_deg": 25.0}


def center_dist(g, center):
    return gc_dist_km(g.xyz, latlon_to_xyz(*center), g.radius_km)


def lip_profile(g, lip, seed):
    r = lip["radius_km"] * (1.0 + 0.25 * fbm(g.xyz, name_seed(seed, lip["name"]), 3, 6.0))
    x = center_dist(g, lip["center"]) / r
    return np.clip((1.0 - x) / 0.25, 0.0, 1.0)


def _scar_azimuths(v, seed):
    if "scar_azimuths" in v:
        return list(v["scar_azimuths"])
    rng = np.random.default_rng(name_seed(seed, v["name"]))
    return list(rng.uniform(0, 360, size=2))


def patch_field(g, patch: dict, seed: int):
    """[[land_patch]]: land hint (strength) over a disc whose radius varies ± 25 % · edge_noise (fractal edge)."""
    d = center_dist(g, patch["center"])
    r = float(patch["radius_km"])
    out = np.zeros(g.n)
    near = d < 1.3 * r
    if near.any():
        w = np.clip(2.0 * fbm(g.xyz[near], name_seed(seed, patch["name"]), 6, 12.0), -1.0, 1.0)
        out[near] = patch.get("strength", 1.0) * (d[near] < r * (1.0 + 0.25 * patch.get("edge_noise", 0.0) * w))
    return out


def run(ctx) -> dict:
    g = ctx.grid
    P = params(ctx.cfg, "crust", DEFAULTS)
    land, land_hint, btype = ctx.need("sk_land", "m_land_hint", "bnd_type")
    rough = (fbm(g.xyz, ctx.seed + 23, 7, P["edge_noise_freq"], gain=P["edge_noise_gain"])
             if P["edge_noise"] > 0 else None)

    def continental_of(raw):
        field = smooth_km(g, raw, P["shelf_km"])
        if rough is not None:     # fractal margins: bays, peninsulas, offshore continental fragments
            f = np.clip(field, 0.0, 1.0)
            near = 4.0 * f * (1.0 - f)                         # strongest at the margin; none deep inland / far offshore
            field = field + P["edge_noise"] * near * rough / max(float(rough.std()), 1e-12)
        else:
            field = np.maximum(field, (raw > 0.5).astype(float))
        cont = field > P["continental_threshold"]
        for m in ctx.tect.get("microcontinent", []):
            cont |= center_dist(g, m["center"]) < m["radius_km"]
        return cont

    raw = np.clip(land + 0.5 * land_hint, 0.0, 1.0)
    continental_base = continental_of(raw)                     # without land patches: its sea level is the world's
    patches = sum((patch_field(g, p, ctx.seed) for p in ctx.tect.get("land_patch", [])), np.zeros(g.n))
    continental = continental_of(np.clip(raw + patches, 0.0, 1.0)) if patches.any() else continental_base.copy()
    plateau_id = PL.cell_ids(g.xyz, ctx.tect.get("plateau", []), ctx.seed, g.radius_km)
    continental |= plateau_id >= 0                             # sunken plateaus: continental crust that never rose
    d_conv = distance_to(g, btype == CONV)
    d_div = distance_to(g, btype == DIV)
    age = np.full(g.n, CRATON, np.int8)
    age[continental & (ridged(g.xyz, ctx.seed + 21, 4, 4.0) > P["pre_orogen_threshold"])] = PRE_OROGEN
    age[continental & (d_div < P["rift_width_km"])] = RIFT
    age[continental & (d_conv < P["orogen_width_km"])] = POST_OROGEN
    age[~continental] = OCEANIC
    for lip in ctx.tect.get("lip", []):
        age[lip_profile(g, lip, ctx.seed) > 0] = LIP
    for v in ctx.tect.get("volcano", []):
        d = center_dist(g, v["center"])
        r = v["radius_km"]
        age[d < r] = VOLCANO
        az = azimuth_deg(latlon_to_xyz(*v["center"]), g.xyz)
        for a in _scar_azimuths(v, ctx.seed):
            diff = np.abs((az - a + 180.0) % 360.0 - 180.0)
            age[(diff < P["scar_halfwidth_deg"]) & (d > P["scar_inner"] * r) & (d < P["scar_outer"] * r)] = SCAR
    oceanic_age = np.minimum(d_div / P["half_spreading_km_myr"], P["max_ocean_age_myr"])
    ocean_age = np.where(continental & (plateau_id < 0), 0.0, oceanic_age)   # plateaus: the floor around them
    return {"continental": continental, "age_class": age, "ocean_age_myr": ocean_age.astype(np.float32),
            "d_conv_km": d_conv, "d_div_km": d_div, "continental_base": continental_base, "plateau_id": plateau_id,
            "land_patch": patches.astype(np.float32)}