"""Stage `climate`: seasonal insolation → temperature, 3-cell winds + monsoons, moisture transport → rain."""
from __future__ import annotations

import numpy as np
from scipy import sparse

from . import ocean as OC
from .config import params
from .graph import bicgstab_jacobi, distance_to, gradient, nearest_source, pmap, smooth_km
from .grid import rowdot_at
from .sphere import east_north, latlon_to_xyz, rotate_about, tangent_dir

S0 = 1361.0
SEA_FREEZE_C = -1.8          # the exported SST never goes below sea water's freezing point (ice-covered sea)
SEASONS = ("jun", "dec", "eq")
DEFAULTS = {
    "t_a": -55.2, "t_b": 0.3206, "t_c": -2.974e-4, "land_seasonal": 0.45, "land_summer": 0.9, "continentality_summer_max": 1.5, "ocean_seasonal": 0.15,
    "continentality_km": 1500.0, "continentality_max": 1.8, "lapse_c_per_km": 6.5,
    "current_c": 4.0, "current_reach_km": 600.0, "current_leak_km": 150.0, "heat_transport_km": 300.0,
    "hadley_edge_deg": 20.0, "ferrel_edge_deg": 55.0, "itcz_shift_deg": 8.0,
    "trade_u": -6.0, "trade_v": 2.0, "westerly_u": 8.0, "westerly_v": 1.0, "polar_u": -4.0, "polar_v": 1.0,
    "monsoon_k": 1.0, "monsoon_length_km": 1500.0, "monsoon_speed_scale": 6000.0,
    "coriolis_min_deg": 20.0, "coriolis_span_deg": 50.0,
    "base_rate": 0.25, "conv_rate": 2.0, "front_rate": 0.8, "front_lat_deg": 40.0, "front_width_deg": 10.0, "itcz_width_deg": 8.0, "oro_rate": 20.0, "subsidence": 0.2,
    "min_rate": 0.05, "cc_per_c": 0.07, "recycle": 0.83, "eddy_k_m2s": 2.2e6,
    "eddy_wind_ms": 8.0,
    "global_mean_mm": 1000.0,
    "lock": False, "lock_at": [0.0, 0.0], "lock_day_c": 120.0, "lock_night_c": -200.0, "lock_wind_ms": 10.0,
    "lock_melt_c": 0.0, "lock_melt_width_c": 25.0,
}


def t_of_q(q, P):
    """Radiative-equilibrium-like surface temperature (°C) from insolation (W/m²); quadratic fit to Earth-like
    zonal means for a 20° tilt: equator 27, 45° ≈ 15, 60° ≈ 3, pole ≈ −22 (concave: damps polar-day summers)."""
    return P["t_a"] + P["t_b"] * q + P["t_c"] * q * q


def declinations(tilt):
    return {"jun": tilt, "dec": -tilt, "eq": 0.0}


def insolation(lat, decl):
    """Daily-mean top-of-atmosphere insolation (W/m²)."""
    phi = np.radians(np.clip(lat, -89.9999, 89.9999))
    d = np.radians(decl)
    h0 = np.arccos(np.clip(-np.tan(phi) * np.tan(d), -1.0, 1.0))
    return S0 / np.pi * (h0 * np.sin(phi) * np.sin(d) + np.cos(phi) * np.cos(d) * np.sin(h0))


def zonal_mean(g, f, bin_deg=2.0):
    b = np.floor((g.lat + 90.0) / bin_deg).astype(np.int64)
    s = np.bincount(b, weights=f * g.area_km2)
    w = np.bincount(b, weights=g.area_km2)
    return (s / np.maximum(w, 1e-12))[b]


def current_anomaly(g, land, P):
    """Subtropical gyres: cold water off west coasts, warm off east coasts; leaks onto coastal land."""
    if not land.any():
        return np.zeros(g.n)
    d, src = nearest_source(g, np.flatnonzero(land))
    e, _ = east_north(g.xyz)
    east_comp = np.sum(tangent_dir(g.xyz, g.xyz[np.maximum(src, 0)]) * e, axis=1)
    band = np.sin(np.radians((np.clip(np.abs(g.lat), 10.0, 50.0) - 10.0) * 4.5))
    a = -P["current_c"] * np.sign(east_comp) * band * np.exp(-d / P["current_reach_km"])
    a[land] = 0.0
    return smooth_km(g, a, P["current_leak_km"])


def temperatures(g, z, land, P, tilt, cur=None):
    decl = declinations(tilt)
    Q = {s: insolation(g.lat, d) for s, d in decl.items()}
    q_ann = (Q["jun"] + Q["dec"] + 2 * Q["eq"]) / 4
    t_ann = t_of_q(q_ann, P)
    dist_ocean = distance_to(g, ~land) if (~land).any() else np.full(g.n, 1e4)
    cont = np.clip(1.0 + dist_ocean / P["continentality_km"], 1.0, P["continentality_max"])
    cur = current_anomaly(g, land, P) if cur is None else cur
    lapse = -P["lapse_c_per_km"] * np.maximum(z, 0.0) / 1000.0
    def season(s):
        raw = t_of_q(Q[s], P)
        land_resp = np.where(raw > t_ann, P["land_summer"] * np.minimum(cont, P["continentality_summer_max"]),
                             P["land_seasonal"] * cont)       # land heats faster in summer (dry, low heat capacity)
        resp = np.where(land, land_resp, P["ocean_seasonal"])
        return smooth_km(g, t_ann + resp * (raw - t_ann) + cur, P["heat_transport_km"]) + lapse
    return dict(zip(SEASONS, pmap(season, SEASONS))), dist_ocean


def biotemperature(tmean, trange, n=12):
    ph = np.linspace(0.0, 2 * np.pi, n, endpoint=False)
    t = np.asarray(tmean)[:, None] + (np.asarray(trange)[:, None] / 2) * np.sin(ph)[None, :]
    return np.clip(t, 0.0, 30.0).mean(axis=1)


def band_winds(g, itcz_lat, P):
    """3-cell surface winds (m/s) relative to the thermal equator."""
    phi = g.lat - itcz_lat
    a = np.abs(phi)
    sgn = np.where(phi >= 0, 1.0, -1.0)
    s1 = 0.5 * (1 + np.tanh((a - P["hadley_edge_deg"]) / 3.0))
    s2 = 0.5 * (1 + np.tanh((a - P["ferrel_edge_deg"]) / 4.0))
    u = (1 - s1) * P["trade_u"] + (s1 - s2) * P["westerly_u"] + s2 * P["polar_u"]
    v = sgn * (-(1 - s1) * P["trade_v"] + (s1 - s2) * P["westerly_v"] - s2 * P["polar_v"])
    e, n = east_north(g.xyz)
    return u[:, None] * e + v[:, None] * n


def monsoon_winds(g, T_s, land, P):
    """Thermal lows over hot land / highs over cold land, flow deflected by Coriolis."""
    anom = np.where(land, T_s - zonal_mean(g, T_s), 0.0)
    press = -P["monsoon_k"] * smooth_km(g, anom, P["monsoon_length_km"])
    flow = -gradient(g, press)
    theta = np.radians(P["coriolis_min_deg"] + P["coriolis_span_deg"] * np.abs(np.sin(np.radians(g.lat))))
    return rotate_about(g.xyz, flow, -np.sign(g.lat) * theta) * P["monsoon_speed_scale"]


def _rate(per_1000km, spacing_km):
    return 1.0 - np.exp(-np.maximum(per_1000km, 0.0) * spacing_km / 1000.0)


def eddy_mixing(g, P):
    """kappa · (nbr − I): the per-step eddy mixing with the neighbours — the same for every season (build once)."""
    kappa = P["eddy_k_m2s"] / (P["eddy_wind_ms"] * g.spacing_km * 1000.0)
    nbr = sparse.csr_matrix((1.0 / g.counts[g.src], (g.src, g.dst)), shape=(g.n, g.n))
    return kappa * (nbr - sparse.identity(g.n, format="csr"))


def precipitation(g, wind, T, z, land, itcz_lat, P, mixing=None):
    """Steady-state moisture transport along the wind on the cell graph; returns rain (relative units).
    mixing: eddy_mixing(g, P), when several seasons share it."""
    t = g.edge_tangents
    out = np.maximum(rowdot_at(wind, g.src, t), 0.0)
    tot = np.bincount(g.src, weights=out, minlength=g.n)
    frac = np.where(tot[g.src] > 0, out / np.maximum(tot[g.src], 1e-12), 0.0)
    Tm = sparse.csr_matrix((frac, (g.dst, g.src)), shape=(g.n, g.n))
    stay = (tot <= 0).astype(np.float64)
    upslope = np.maximum(np.sum(wind * gradient(g, np.maximum(z, 0.0) / 1000.0), axis=1), 0.0)
    conv = P["conv_rate"] * np.exp(-((g.lat - itcz_lat) / P["itcz_width_deg"]) ** 2) * np.clip((T - 10.0) / 20.0, 0, 1)
    subs = P["subsidence"] * np.exp(-((np.abs(g.lat - itcz_lat) - P["hadley_edge_deg"]) / 6.0) ** 2)
    front = P["front_rate"] * np.exp(-((np.abs(g.lat - itcz_lat) - P["front_lat_deg"]) / P["front_width_deg"]) ** 2)
    per = np.maximum(P["base_rate"] + conv + front + P["oro_rate"] * upslope - subs, P["min_rate"])
    r = _rate(per, g.spacing_km)
    evap = np.where(land, 0.0, np.exp(P["cc_per_c"] * (np.clip(T, -2.0, 35.0) - 25.0)))
    # per advection step (one cell, time h/U): rain out, move downwind, eddy-mix with neighbours
    mix = eddy_mixing(g, P) if mixing is None else mixing
    eye = sparse.identity(g.n, format="csr")
    keep = sparse.diags(1.0 - r)
    recyc = sparse.diags(np.where(land, P["recycle"] * r, 0.0))   # land evapotranspiration returns rain
    system = (eye - (Tm @ keep + sparse.diags(stay) @ keep) - mix - recyc).tocsr()
    W, info = bicgstab_jacobi(system, evap, evap / np.maximum(r, 1e-6), 1.0 / system.diagonal(), 1e-7, 5000)
    if info != 0:
        raise ValueError(f"precipitation: moisture solve did not converge (info={info})")
    return r * np.maximum(W, 0.0)


def run_locked(ctx, g, z, land, P) -> dict:
    """One face always to the sun: temperature by sun angle, no seasons; surface wind from night to the sun point."""
    sub = latlon_to_xyz(*P["lock_at"])
    mu = np.maximum(g.xyz @ sub, 0.0)
    t = P["lock_night_c"] + (P["lock_day_c"] - P["lock_night_c"]) * mu ** 0.25
    lapse = -P["lapse_c_per_km"] * np.maximum(z, 0.0) / 1000.0
    T = smooth_km(g, t, P["heat_transport_km"]) + lapse
    dist_ocean = distance_to(g, ~land) if (~land).any() else np.full(g.n, 1e4)
    wind = P["lock_wind_ms"] * tangent_dir(g.xyz, np.broadcast_to(sub, g.xyz.shape))
    rain = np.exp(-((T - P["lock_melt_c"]) / P["lock_melt_width_c"]) ** 2)   # meltwater and frost in the twilight ring
    k = P["global_mean_mm"] / max(np.sum(rain * g.area_km2) / g.area_km2.sum(), 1e-12)
    out = {}
    for s in SEASONS:
        out[f"wind_{s}"] = wind.astype(np.float32)
        out[f"P_{s}"] = rain * k
        out[f"T_{s}"] = T
    out["P_ann"] = rain * k
    out["T_mean"] = T
    out["T_range"] = np.zeros(g.n)
    out["T_min"] = T
    out["biotemp"] = biotemperature(T, out["T_range"])
    out["PET"] = 58.93 * out["biotemp"]
    out["dist_ocean_km"] = dist_ocean
    return out


def _still_ocean(g, t_mean):
    """No circulation (ocean disabled, locked world): zero currents/upwelling/productivity, SST = T_mean."""
    z = np.zeros(g.n, np.float32)
    return {"current": np.zeros((g.n, 3), np.float32), "current_speed": z, "sst": np.asarray(t_mean, np.float32),
            "upwelling": z.copy(), "productivity": z.copy()}


def run(ctx) -> dict:
    g = ctx.grid
    P = params(ctx.cfg, "climate", DEFAULTS)
    tilt = float(ctx.cfg["planet"]["tilt_deg"])
    (z,) = ctx.need("elevation_eroded_m")
    z = z.astype(np.float64)
    water = ctx.data.get("open_water", ctx.data.get("ocean"))            # big inland basins are water to the air
    land = ~np.asarray(water) if water is not None else z > 0
    O = params(ctx.cfg, "ocean", OC.DEFAULTS)
    if P["lock"]:
        out = run_locked(ctx, g, z, land, P)
        out.update(_still_ocean(g, out["T_mean"]))
        return out
    sea = ~land
    coupled = bool(O["enabled"]) and bool(sea.any())
    T, dist_ocean = temperatures(g, z, land, P, tilt, cur=np.zeros(g.n) if coupled else None)
    decl = declinations(tilt)
    def season_wind(s):
        itcz = P["itcz_shift_deg"] * decl[s] / max(tilt, 1e-9)
        wind = band_winds(g, itcz, P)
        return wind + monsoon_winds(g, T[s], land, P) if s != "eq" else wind
    winds = dict(zip(SEASONS, pmap(season_wind, SEASONS)))
    if coupled:                     # one pass: winds from current-free temperatures, then currents carry heat
        day = float(ctx.cfg["planet"]["day_hours"])
        w_ann = (winds["jun"] + winds["dec"] + 2 * winds["eq"]) / 4
        u = OC.currents(g, sea, w_ann, O, day)
        T_eq = (T["jun"] + T["dec"] + 2 * T["eq"]) / 4
        T_s = OC.sst(g, sea, u, T_eq, O)
        cur = smooth_km(g, np.where(sea, T_s - T_eq, 0.0), P["current_leak_km"])
        T, _ = temperatures(g, z, land, P, tilt, cur=cur)
    out = {}
    mixing = eddy_mixing(g, P)
    rain = dict(zip(SEASONS, pmap(lambda s: precipitation(g, winds[s], T[s], z, land,
                                                           P["itcz_shift_deg"] * decl[s] / max(tilt, 1e-9), P, mixing),
                                  SEASONS)))
    del mixing
    for s in SEASONS:
        out[f"wind_{s}"] = winds[s].astype(np.float32)
    ann = (rain["jun"] + rain["dec"] + 2 * rain["eq"]) / 4
    k = P["global_mean_mm"] / max(np.sum(ann * g.area_km2) / g.area_km2.sum(), 1e-12)
    for s in SEASONS:
        out[f"P_{s}"] = rain[s] * k
        out[f"T_{s}"] = T[s]
    out["P_ann"] = ann * k
    out["T_mean"] = (T["jun"] + T["dec"] + 2 * T["eq"]) / 4
    out["T_range"] = np.abs(T["jun"] - T["dec"])
    out["T_min"] = np.minimum(T["jun"], T["dec"])
    out["biotemp"] = biotemperature(out["T_mean"], out["T_range"])
    out["PET"] = 58.93 * out["biotemp"]
    out["dist_ocean_km"] = dist_ocean
    if coupled:
        sst_c = np.where(sea, np.maximum(T_s, SEA_FREEZE_C), out["T_mean"])
        w_up = OC.upwelling(g, sea, w_ann, O, day)
        out.update({"current": u.astype(np.float32),
                    "current_speed": np.linalg.norm(u, axis=1).astype(np.float32),
                    "sst": sst_c.astype(np.float32),
                    "upwelling": w_up.astype(np.float32),
                    "productivity": OC.productivity(g, sea, w_up, z, out["T_range"], sst_c, O).astype(np.float32)})
    else:
        out.update(_still_ocean(g, out["T_mean"]))
    return out