"""`mapgen.py check`: physical sanity checks (fail) + environment summary (info)."""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np

from . import plateaus as PL, zones as ZN
from .climate import DEFAULTS as CLIMATE_DEFAULTS
from .config import params
from .environment import GROUND_NAMES, LANDFORM_NAMES, REGIONS
from .fields import gravity_mod
from .graph import gradient
from .ice import ICE_NAMES
from .sphere import gc_dist_km, latlon_to_xyz


def check_land_fraction(z, area, target, tol=0.02):
    f = area[z > 0].sum() / area.sum()
    return None if abs(f - target) <= tol else f"land fraction {f:.3f} not within {tol} of target {target}"


def check_hypsometry(z, area):
    edges = np.arange(-11000, 12500, 500)
    h, _ = np.histogram(z, bins=edges, weights=area)
    c = (edges[:-1] + edges[1:]) / 2
    ocean = (c >= -7000) & (c <= -1000)
    landb = (c >= -500) & (c <= 3000)
    io, il = np.flatnonzero(ocean)[np.argmax(h[ocean])], np.flatnonzero(landb)[np.argmax(h[landb])]
    valley = h[io:il + 1].min()
    if h[io] > 1.5 * valley and h[il] > 1.5 * valley:
        return None
    return "hypsometry not bimodal (expected separate ocean-floor and continent peaks)"


def check_max_elevation(z, gmod, cap=12000.0):
    lim = cap * np.clip(1.0 / gmod, 1.0, 3.0)
    bad = z > lim + 1.0
    return None if not bad.any() else f"{int(bad.sum())} cells above {cap:.0f} m (outside low-g zones)"


def check_drainage(recv, z, endorheic, terminal_ok=None):
    roots = recv == np.arange(len(recv))
    bad = roots & (z > 0) & ~endorheic
    if bad.any():
        return f"{int(bad.sum())} land cells are sinks outside endorheic basins"
    if terminal_ok is not None:
        dry = roots & endorheic & ~terminal_ok
        if dry.any():
            return f"{int(dry.sum())} endorheic terminals are neither lake nor salt flat"
    return None


def check_no_uphill(zf, recv, z, endorheic, lake=None):
    ar = np.arange(len(recv))
    sel = (z > 0) & ~endorheic & (recv != ar)
    if lake is not None:
        sel &= ~lake                                  # flat lake surfaces: routing direction is a convention
    bad = sel & (zf[recv] >= zf)
    return None if not bad.any() else f"{int(bad.sum())} cells route uphill on the filled surface"


def _band_means(g, f, edges):
    a = np.abs(g.lat)
    return [np.sum(f[(a >= lo) & (a < hi)] * g.area_km2[(a >= lo) & (a < hi)]) /
            max(g.area_km2[(a >= lo) & (a < hi)].sum(), 1e-9) for lo, hi in zip(edges, edges[1:])]


def check_temperature(g, t_mean):
    m = _band_means(g, t_mean, list(range(0, 91, 10)))
    ok = all(b <= a + 0.5 for a, b in zip(m, m[1:])) and m[0] > m[-1] + 10
    return None if ok else f"zonal-mean temperature does not fall poleward: {np.round(m, 1).tolist()}"


def check_rain_bands(g, p_ann, hadley_edge):
    """Dry belt ≈ Hadley edge −2..+8°, storm track ≈ edge +12..+25° (scaled from Earth's 30° edge)."""
    h = hadley_edge
    eq = _band_means(g, p_ann, [0, 10])[0]
    sub = _band_means(g, p_ann, [h - 2, h + 8])[0]
    mid = _band_means(g, p_ann, [h + 12, h + 25])[0]
    if eq > sub and mid > sub:
        return None
    return (f"no subtropical dry belt: P(0-10)={eq:.0f}, P({h - 2:.0f}-{h + 8:.0f})={sub:.0f}, "
            f"P({h + 12:.0f}-{h + 25:.0f})={mid:.0f} mm/yr")


def check_rain_shadow(g, data):
    z = np.maximum(np.asarray(data["elevation_eroded_m"], dtype=np.float64), 0.0) / 1000.0
    grad = gradient(g, z)
    land = z > 0
    ww, lw = [], []
    for s in ("jun", "dec"):
        up = np.sum(np.asarray(data[f"wind_{s}"], dtype=np.float64) * grad, axis=1)
        p = np.asarray(data[f"P_{s}"])
        ww.append(p[land & (up > 0.02)])
        lw.append(p[land & (up < -0.02)])
    w, l = np.concatenate(ww), np.concatenate(lw)
    if len(w) == 0 or len(l) == 0:
        return None
    return None if w.mean() > l.mean() else f"windward slopes ({w.mean():.0f}) not wetter than lee ({l.mean():.0f})"


def _shares(values, area, names, mask):
    tot = area[mask].sum()
    out = []
    for i, nm in enumerate(names):
        s = area[mask & (values == i)].sum() / max(tot, 1e-9)
        if s >= 0.005:
            out.append(f"{nm} {100 * s:.1f}%")
    return ", ".join(out)



def check_plateaus(g, data, plateaus, seed) -> list:
    """Spec §8: every plateau's top within top_m (± relief), hidden ones ≤ −500 m, island ones with land."""
    if not plateaus:
        return []
    z = np.asarray(data["elevation_eroded_m"], dtype=np.float64)
    ocean, ids = np.asarray(data["ocean"]), np.asarray(data["plateau_id"])
    out = []
    for k, p in enumerate(plateaus):
        idx = np.flatnonzero(ids == k)
        if len(idx) == 0:
            out.append(f"plateau {p['name']}: no cells at this resolution")
            continue
        _, b = PL.semi_axes(p)
        core = idx[(1.0 - PL.rho(g.xyz[idx], p, seed, g.radius_km)) * b > PL.MARGIN_KM]
        lo, hi = p["top_m"]
        if len(core):
            med = float(np.median(-z[core]))
            if not lo - PL.RELIEF_M <= med <= hi + PL.RELIEF_M:
                out.append(f"plateau {p['name']}: median top depth {med:.0f} m outside {lo:.0f}–{hi:.0f} m")
        if p.get("islands", False):
            isl = [g.cell_index(c["lat"], c["lon"]) for c in PL.features(p, seed, g.radius_km)["cones"] if c["island"]]
            if ocean[idx].all() and all(ocean[i] for i in isl):
                out.append(f"plateau {p['name']}: no island breaks the surface")
        elif z[idx].max() > -500.0:
            out.append(f"plateau {p['name']}: a hidden plateau reaches {z[idx].max():.0f} m (limit −500 m)")
    return out


def check_zones(g, data, max_o2_points=0.23, max_g=0.014, per_km=30.0) -> list:
    """Spec §5, the gradual rule: steepest change per day's walk of the O₂ fraction (points) and of gravity."""
    out = []
    for key, scale, lim, unit in (("o2_fraction", 100.0, max_o2_points, "O₂ points"), ("gravity_g", 1.0, max_g, "g")):
        if key not in data:
            continue
        f = np.asarray(data[key], dtype=np.float64) * scale
        steep = float(np.max(np.abs(f[g.dst] - f[g.src]) / g.edge_km)) * per_km
        if steep > lim + 1e-9:
            out.append(f"zones: {key} changes {steep:.3f} {unit} per {per_km:.0f} km (limit {lim})")
    return out


def edited_areas(xyz, tect, radius_km, plateau_km=300.0, patch_km=1500.0):
    """Cells the revision changes by design: plateaus + 300 km, land patches + 1,500 km, low-gravity zones
    (their relief)."""
    xyz = np.asarray(xyz, dtype=np.float64)
    near = lambda c, r: gc_dist_km(xyz, latlon_to_xyz(*c), radius_km) < r
    m = np.zeros(len(xyz), bool)
    for p in tect.get("plateau", []):
        m |= near(p["center"], PL.semi_axes(p)[0] * (1 + PL.EDGE_WARP) / (1 - PL.EDGE_WARP) + plateau_km)
    for p in tect.get("land_patch", []):
        m |= near(p["center"], p["radius_km"] * 1.3 + patch_km)
    for z in tect.get("zone", []):
        if z["field"] == "gravity":
            m |= near(z["center"], z["radius_km"] / (1 - ZN.WARP))
    return m


O2_FIELDS = ("m_o2_zones", "o2_fraction", "po2_bar")        # what the base's O₂ zones change (no relief)


def o2_areas(xyz, tect, radius_km):
    """Cells the base's O₂ zones reach: compare skips only the O₂ fields there."""
    xyz = np.asarray(xyz, dtype=np.float64)
    m = np.zeros(len(xyz), bool)
    for z in tect.get("zone", []):
        if z["field"] == "o2":
            m |= gc_dist_km(xyz, latlon_to_xyz(*z["center"]), radius_km) < z["radius_km"] / (1 - ZN.WARP)
    return m


def compare(old: dict, new: dict, exclude, extra=None) -> list:
    """Per field, over the cells outside `exclude` (and outside extra[field] for that field): share changed, largest
    change, 99th percentile."""
    keep0 = ~np.asarray(exclude, bool)
    lines = [f"cells compared: {int(keep0.sum())} of {len(keep0)} (outside the edited areas)"]
    for k in sorted(set(old) & set(new)):
        a, b = np.asarray(old[k]), np.asarray(new[k])
        if a.shape != b.shape or a.shape[:1] != keep0.shape or a.dtype.kind not in "biuf" or b.dtype.kind not in "biuf":
            continue
        keep = keep0 & ~np.asarray(extra[k], bool) if extra and k in extra else keep0
        a, b = a.astype(np.float64)[keep], b.astype(np.float64)[keep]
        with np.errstate(invalid="ignore"):
            diff = np.where((a == b) | (np.isnan(a) & np.isnan(b)), 0.0, np.abs(b - a))
        if diff.ndim > 1:
            diff = diff.reshape(len(diff), -1).max(axis=1)
        if not len(diff):
            lines.append(f"{k}: no cells to compare")
            continue
        lines.append(f"{k}: {np.mean(diff > 0):.1%} changed, max {diff.max():.4g}, p99 {np.percentile(diff, 99):.4g}")
    return lines


def compare_report(root: Path, res: int, old_path: Path) -> int:
    from . import config as C
    _, tect = C.load(root)
    out = root / "out" / f"r{res}"
    with np.load(old_path) as z:
        old = {k: z[k] for k in z.files}
    with np.load(out / "cells.npz") as z:
        new = {k: z[k] for k in z.files}
    for name, arrays in ((old_path, old), (out / "cells.npz", new)):
        if "g_ids" not in arrays:
            print(f"error: {name} has no g_ids: not the cells.npz of a world build")
            return 2
    if not np.array_equal(old["g_ids"], new["g_ids"]):
        print("error: the two worlds have different cells (another resolution?)")
        return 2
    radius = float(json.loads((out / "cells_meta.json").read_text())["radius_km"])
    o2 = o2_areas(new["g_xyz"], tect, radius)
    for line in compare(old, new, edited_areas(new["g_xyz"], tect, radius), {k: o2 for k in O2_FIELDS}):
        print(line)
    return 0

def info_lines(ctx):
    g, d = ctx.grid, ctx.data
    z = np.asarray(d["elevation_eroded_m"])
    land = ~np.asarray(d["ocean"]) if "ocean" in d else z > 0
    a = g.area_km2
    sk = np.asarray(d["sk_land"]) > 0.5
    iou = (land & sk).sum() / max((land | sk).sum(), 1)
    sites = PL.site_report(g, d, ctx.tect["plateau"]) if ctx.tect.get("plateau") and "plateau_id" in d else []
    return [f"cells {g.n} (H3 res {ctx.res}); land {100 * a[land].sum() / a.sum():.1f}% "
            f"({a[land].sum() / 1e6:.0f}M km², Earth: 29%, 149M km²)",
            f"sketch land overlap (IoU): {iou:.2f}",
            "regions: " + _shares(np.asarray(d["hold_region"]), a, REGIONS, land),
            "landforms: " + _shares(np.asarray(d["landform"]), a, LANDFORM_NAMES, land),
            "ground: " + _shares(np.asarray(d["ground"]), a, GROUND_NAMES, land),
            "ice: " + _shares(np.asarray(d["ice"]), a, ICE_NAMES, np.ones(g.n, bool))] + [f"site: {s}" for s in sites]


def run_checks(ctx):
    g, d, cfg = ctx.grid, ctx.data, ctx.cfg
    z = np.asarray(d["elevation_eroded_m"], dtype=np.float64)
    if "ocean" in d:   # interior lows below sea level are land: sign-correct z for the land/sea checks
        z = np.where(np.asarray(d["ocean"]), np.minimum(z, -0.001), np.maximum(z, 0.001))
    endo = np.asarray(d["endorheic"])
    results = [
        check_land_fraction(z, g.area_km2, cfg["build"]["land_fraction"]),
        check_hypsometry(z, g.area_km2),
        check_max_elevation(z, gravity_mod(cfg, d["m_gravity_zones"])),
        check_drainage(np.asarray(d["recv"]), z, endo, np.asarray(d["lake"]) | np.asarray(d["salt_flat"])),
        check_no_uphill(np.asarray(d["z_filled_m"]), np.asarray(d["recv"]), z, endo, np.asarray(d["lake"])),
        check_temperature(g, np.asarray(d["T_mean"])),
        check_rain_bands(g, np.asarray(d["P_ann"]), params(cfg, "climate", CLIMATE_DEFAULTS)["hadley_edge_deg"]),
        check_rain_shadow(g, d),
    ]
    results += check_plateaus(g, d, ctx.tect.get("plateau", []), ctx.seed) + check_zones(g, d)
    return [r for r in results if r], info_lines(ctx)


def report(ctx) -> int:
    fails, infos = run_checks(ctx)
    for line in infos:
        print("info:", line)
    for f in fails:
        print("FAIL:", f)
    print("check:", "OK" if not fails else f"{len(fails)} failure(s)")
    return 1 if fails else 0