"""Ocean circulation for stage `climate`: wind-driven surface currents, sea-surface temperature, Ekman upwelling
and a productivity index.

Currents: the Stommel model on the sphere, r∇²ψ + βψ_x = k·curl(τ/ρH), solved for the stream function ψ over
every cell; land is the same fluid with `land_friction`× the friction (Brinkman penalisation), so coasts block the
flow and islands need no special treatment. u = k × ∇ψ. Western boundary currents come out ≈ r/β wide.
"""
from __future__ import annotations

import numpy as np
from scipy import sparse
from scipy.sparse import linalg as splinalg

from .graph import bicgstab_jacobi, distance_to, gradient, pmap, smooth_km
from .grid import rowdot_at
from .sphere import east_north

DEFAULTS = {
    "enabled": True, "friction_days": 5.0, "layer_m": 150.0, "stress_k": 1.0, "land_friction": 1000.0,
    "rho_air": 1.2, "drag": 1.3e-3, "direct_max_cells": 500000,
    "relax_days": 300.0, "kappa_m2s": 1000.0,
    "ekman_min_lat": 3.0, "upwell_coast_km": 100.0,
    "prod_upwell": 0.6, "prod_shelf": 0.4, "prod_mix": 0.3, "upwell_ref_m_yr": 100.0,
    "prod_front": 0.4, "front_min": 0.5, "front_ref": 1.0,
}
RHO_W = 1025.0
YEAR_S = 3.15576e7


def omega(day_hours):
    return 2.0 * np.pi / (float(day_hours) * 3600.0)


def divergence(g, V):
    """Graph divergence of a tangent field (V per km): (2/k_i) Σ_j V_j·t_ij / d_ij (exact for linear V on a hex
    lattice; the V_i terms cancel)."""
    w = rowdot_at(V, g.dst, g.edge_tangents) / g.edge_km
    return np.bincount(g.src, weights=w, minlength=g.n) * (2.0 / g.counts)


def _edge_operator(g, w):
    """Sparse operator f ↦ Σ_j w_ij (f_j − f_i)."""
    A = sparse.csr_matrix((w, (g.src, g.dst)), shape=(g.n, g.n))
    return A - sparse.diags(np.asarray(A.sum(axis=1)).ravel())


def wind_stress(wind, P):
    """Bulk formula τ = k·ρ_air·C_d·|w|·w (N/m²)."""
    wind = np.asarray(wind, dtype=np.float64)
    return P["stress_k"] * P["rho_air"] * P["drag"] * np.linalg.norm(wind, axis=1, keepdims=True) * wind


def streamfunction(g, ocean, wind, P, day_hours):
    """ψ (m²/s) of the wind-driven surface flow; one direct sparse solve over every cell."""
    Om = omega(day_hours)
    R = g.radius_km * 1e3
    beta = 2.0 * Om * np.cos(np.radians(g.lat)) / R                    # 1/(m·s)
    beta30 = 2.0 * Om * np.cos(np.radians(30.0)) / R
    r = max(1.0 / (P["friction_days"] * 86400.0), beta30 * g.spacing_km * 1e3)   # boundary layer ≥ one cell
    F = np.where(ocean[:, None], wind_stress(wind, P), 0.0) / (RHO_W * P["layer_m"])
    curl = divergence(g, np.cross(F, g.xyz)) * 1e-3                    # k·curl F = ∇·(F×k), 1/s²
    fr = np.where(ocean, 1.0, P["land_friction"])
    fe = 2.0 / (1.0 / fr[g.src] + 1.0 / fr[g.dst])                    # harmonic mean: coasts count as land
    L = _edge_operator(g, 4.0 / (g.counts[g.src] * g.edge_km ** 2) * fe)          # ∇·(fr∇), per km²
    e, _ = east_north(g.xyz)
    Dx = _edge_operator(g, (2.0 / g.counts[g.src]) * rowdot_at(e, g.src, g.edge_tangents) / g.edge_km)
    A = (L + sparse.diags(beta / r * 1e3) @ Dx).tolil()                # β/r·1e3: per km
    b = 1e6 * curl / r
    # gauge: ψ is defined up to a constant; pin it deep inland, where land friction soaks up the solve's
    # compatibility residual (pinned at sea it would leave a point vortex there)
    k = int(np.argmax(distance_to(g, ocean))) if (~ocean).any() else 0
    A[k, :] = 0.0
    A[k, k] = 1.0
    b[k] = 0.0
    return splinalg.spsolve(A.tocsc(), b)


def velocity(g, psi):
    """u = k × ∇ψ (m/s), tangent 3-vectors."""
    return np.cross(g.xyz, gradient(g, psi) * 1e-3)


def _coarse_currents(g, ocean, wind, P, day_hours):
    """Grids too big for a direct solve: solve on the parent H3 resolution, carry u down, smooth."""
    from .grid import build_grid, cell_parents
    cg = build_grid(g.res - 1, g.radius_km)
    parent = np.searchsorted(cg.ids, cell_parents(g.ids, g.res - 1))
    cnt = np.maximum(np.bincount(parent, minlength=cg.n), 1)
    c_ocean = np.bincount(parent, weights=ocean.astype(float), minlength=cg.n) / cnt > 0.5
    c_wind = np.stack([np.bincount(parent, weights=wind[:, k], minlength=cg.n) / cnt for k in range(3)], axis=1)
    cu = currents(cg, c_ocean, c_wind, P, day_hours)
    u = np.stack(pmap(lambda k: smooth_km(g, cu[parent, k], cg.spacing_km), range(3)), axis=1)
    return u - np.sum(u * g.xyz, axis=1, keepdims=True) * g.xyz       # back into the tangent plane


def currents(g, ocean, wind, P, day_hours):
    """Surface current (n,3) m/s; 0 on land."""
    ocean = np.asarray(ocean, bool)
    wind = np.asarray(wind, dtype=np.float64)
    if not ocean.any():
        return np.zeros((g.n, 3))
    if g.n > P["direct_max_cells"]:
        u = _coarse_currents(g, ocean, wind, P, day_hours)
    else:
        u = velocity(g, streamfunction(g, ocean, wind, P, day_hours))
    return np.where(ocean[:, None], u, 0.0)


def _masked_edges(g, mask, w):
    return _edge_operator(g, np.where(mask[g.src] & mask[g.dst], w, 0.0))


def sst(g, ocean, u, T_eq, P):
    """Annual-mean sea-surface temperature (°C): steady u·∇T − κ∇²T = λ(T_eq − T) over the ocean (upwind
    advection, no flux into land); land keeps T_eq."""
    ocean = np.asarray(ocean, bool)
    T_eq = np.asarray(T_eq, dtype=np.float64)
    lam = 1.0 / (P["relax_days"] * 86400.0)
    up = (4.0 / g.counts[g.src]) * np.maximum(-rowdot_at(u, g.src, g.edge_tangents), 0.0) / (g.edge_km * 1e3)
    adv = -_masked_edges(g, ocean, up)                                 # Σ_upstream c_ij (T_i − T_j), 1/s
    dif = _masked_edges(g, ocean, 4.0 / (g.counts[g.src] * g.edge_km ** 2)) * (P["kappa_m2s"] * 1e-6)
    A = (adv - dif) / lam + sparse.identity(g.n)
    A = (sparse.diags(ocean.astype(float)) @ A + sparse.diags((~ocean).astype(float))).tocsr()
    T, info = bicgstab_jacobi(A, T_eq, T_eq.copy(), 1.0 / A.diagonal(), 1e-9, 5000)
    if info != 0:
        raise ValueError(f"sst: solver did not converge (info={info})")
    return T


def upwelling(g, ocean, wind, P, day_hours):
    """Ekman pumping (m/yr, + up): w = ∇·M, M = τ×k/(ρf), |f| floored at `ekman_min_lat`. Transport pointing off a
    coast leaves the coast cell (land carries none), so coastal upwelling needs no separate rule. Smoothed over
    `upwell_coast_km`; 0 on land."""
    ocean = np.asarray(ocean, bool)
    Om = omega(day_hours)
    fmin = 2.0 * Om * np.sin(np.radians(P["ekman_min_lat"]))
    f = np.where(g.lat >= 0, 1.0, -1.0) * np.maximum(np.abs(2.0 * Om * np.sin(np.radians(g.lat))), fmin)
    tau = np.where(ocean[:, None], wind_stress(wind, P), 0.0)
    M = np.cross(tau, g.xyz) / (RHO_W * f[:, None])                    # m²/s
    w = np.where(ocean, divergence(g, M) * 1e-3 * YEAR_S, 0.0)
    return np.where(ocean, smooth_km(g, w, P["upwell_coast_km"]), 0.0)


def sst_front(g, ocean, sst_c):
    """|∇SST| over the sea (°C per 100 km); edges to land count as flat, so coasts are no front."""
    sst_c = np.asarray(sst_c, dtype=np.float64)
    both = ocean[g.src] & ocean[g.dst]
    w = np.where(both, sst_c[g.dst] - sst_c[g.src], 0.0) / g.edge_km
    grad = np.stack([np.bincount(g.src, weights=w * g.edge_tangents[:, c], minlength=g.n) for c in range(3)], axis=1)
    return np.where(ocean, np.linalg.norm(grad, axis=1) * (2.0 / g.counts) * 100.0, 0.0)


def productivity(g, ocean, w, z, t_range, sst_c, P):
    """0–1 sea productivity: upwelling (saturating), shallow shelf, winter mixing and SST fronts (where warm and cold
    currents meet, e.g. a Brazil–Malvinas confluence; gradients under `front_min` °C/100 km add nothing), dimmed
    toward the poles."""
    ocean = np.asarray(ocean, bool)
    depth = np.maximum(-np.asarray(z, dtype=np.float64), 0.0)
    x = np.maximum(w, 0.0) / P["upwell_ref_m_yr"]
    shelf = np.clip((1000.0 - depth) / 800.0, 0.0, 1.0)
    mix = np.clip(np.asarray(t_range) / 20.0, 0.0, 1.0) * np.clip((20.0 - np.asarray(sst_c)) / 20.0, 0.0, 1.0)
    light = 0.3 + 0.7 * np.cos(np.radians(g.lat))
    xf = np.maximum(sst_front(g, ocean, sst_c) - P["front_min"], 0.0) / P["front_ref"]
    p = (P["prod_upwell"] * x / (1.0 + x) + P["prod_shelf"] * shelf + P["prod_mix"] * mix
         + P["prod_front"] * xf / (1.0 + xf)) * light
    return np.where(ocean, np.clip(p, 0.0, 1.0), 0.0)