#!/usr/bin/env python3
"""Join the IHFC Global Heat Flow Database (R2024 v.2026.03) to the Seton et al. (2020)
seafloor age grid, per PREREGISTRATION.md, and write sites.json (the per-site table the page
and verify.mjs read) and summary.json (the pre-registered bin statistics).

Inputs live in raw/ (not committed; fetch.sh downloads them).
"""
import json, math, os, random, sys
import numpy as np
import netCDF4

HERE = os.path.dirname(os.path.abspath(__file__))
RAW = os.path.join(HERE, "raw")
GHF = os.path.join(RAW, "IHFC_2024_GHFDB_v.2026.03.txt")
AGE = os.path.join(RAW, "age.2020.1.GTS2012.6m.nc")

MYR = 365.25 * 86400 * 1e6
PARAMS = {
    "primary": {"k": 3.3, "kappa": 1.0e-6, "dT": 1300.0},
}
BINS = [(0, 5), (5, 10), (10, 20), (20, 40), (40, 60), (60, 80), (80, 100), (100, 140), (140, 200)]


def C(p):
    """q (mW/m2) at t = 1 Myr for a half-space: k dT / sqrt(pi kappa), in mW m^-2 Myr^0.5."""
    return p["k"] * p["dT"] / math.sqrt(math.pi * p["kappa"]) / math.sqrt(MYR) * 1000


def kelvin_age(q, c):
    return (c / q) ** 2


def read_ghf():
    rows = []
    with open(GHF, encoding="latin-1") as f:
        lines = [l.rstrip("\r\n") for l in f]
    hdr_i = next(i for i, l in enumerate(lines) if l.startswith("q\tq_uncertainty"))
    hdr = lines[hdr_i].split("\t")
    col = {h: i for i, h in enumerate(hdr) if h}
    seen = set()
    for l in lines[hdr_i + 1:]:
        c = l.split("\t")
        if len(c) < len(col):
            c += [""] * (len(col) - len(c))
        pid = c[col["ID_parent"]]
        if not pid or pid in seen:
            continue
        seen.add(pid)
        try:
            q = float(c[col["q"]]); lat = float(c[col["lat_NS"]]); lon = float(c[col["long_EW"]])
        except ValueError:
            continue
        rows.append({
            "id": pid, "q": q, "lat": lat, "lon": lon,
            "env": c[col["environment"]].strip(),
            "qual": c[col["Quality_Score_Parent"]].strip(),
            "ref": c[col["publication_reference"]].strip() or c[col["data_reference"]].strip(),
            "year": c[col["Year"]].strip(),
        })
    return rows


def land_gradients():
    """Every child row on land (onshore (continental)) that reports a mean temperature gradient
    (K/km): T_grad_mean, or T_grad_mean_cor where only the corrected value is given."""
    with open(GHF, encoding="latin-1") as f:
        lines = [l.rstrip("\r\n") for l in f]
    hdr_i = next(i for i, l in enumerate(lines) if l.startswith("q\tq_uncertainty"))
    hdr = lines[hdr_i].split("\t"); col = {h: i for i, h in enumerate(hdr) if h}
    out = []
    for l in lines[hdr_i + 1:]:
        c = l.split("\t")
        if len(c) < len(col):
            c += [""] * (len(col) - len(c))
        if not c[col["environment"]].strip().lower().startswith("[onshore (continental)"):
            continue
        v = c[col["T_grad_mean"]] or c[col["T_grad_mean_cor"]]
        try:
            v = float(v)
        except ValueError:
            continue
        if v > 0:
            out.append(v)
    return out


def main():
    rows = read_ghf()
    n_parents = len(rows)
    d = netCDF4.Dataset(AGE)
    lons = np.array(d["lon"][:]); lats = np.array(d["lat"][:]); z = d["z"][:]
    z = np.ma.filled(z.astype("float64"), np.nan)
    dlon = lons[1] - lons[0]; dlat = lats[1] - lats[0]

    def age_at(lat, lon):
        # nearest node of the 0.1 degree grid (nodes at -90 + 0.1 i, -180 + 0.1 j). Positions
        # given to two decimals often sit exactly halfway between nodes, so the tie rule matters
        # and is fixed here as "round half up", written as the same IEEE arithmetic the verifier
        # uses (Python's round() is half-to-even, JavaScript's Math.round is half-up).
        lon = ((lon + 180) % 360) - 180
        i = math.floor((lat + 90) * 10 + 0.5 + 1e-9); j = math.floor((lon + 180) * 10 + 0.5 + 1e-9)
        if not (0 <= i < len(lats) and 0 <= j < len(lons)):
            return float("nan")
        return float(z[i, j])

    counts = {"parents": n_parents}
    off = [r for r in rows if r["env"].lower().startswith("[offshore")]
    counts["offshore"] = len(off)
    for r in off:
        r["age"] = age_at(r["lat"], r["lon"])
    ocean = [r for r in off if not math.isnan(r["age"])]
    counts["offshore_on_grid"] = len(ocean)
    sites = [r for r in ocean if r["q"] > 0]
    counts["selected"] = len(sites)

    c = C(PARAMS["primary"])
    rng = random.Random(1862)

    def boot_median(xs, n=2000):
        xs = sorted(xs)
        if not xs:
            return None
        meds = []
        for _ in range(n):
            s = sorted(rng.choice(xs) for _ in xs)
            m = len(s)
            meds.append(s[m // 2] if m % 2 else 0.5 * (s[m // 2 - 1] + s[m // 2]))
        meds.sort()
        return [meds[int(0.025 * n)], meds[int(0.975 * n) - 1]]

    def median(xs):
        s = sorted(xs); m = len(s)
        return s[m // 2] if m % 2 else 0.5 * (s[m // 2 - 1] + s[m // 2])

    def bin_table(ss, cc):
        out = []
        for lo, hi in BINS:
            b = [r for r in ss if lo <= r["age"] < hi]
            tk = [kelvin_age(r["q"], cc) for r in b]
            R = [kelvin_age(r["q"], cc) / r["age"] for r in b if r["age"] > 0]
            out.append({
                "lo": lo, "hi": hi, "n": len(b),
                "median_q": median([r["q"] for r in b]) if b else None,
                "median_tK": median(tk) if tk else None,
                "median_tK_ci": boot_median(tk),
                "median_R": median(R) if R else None,
                "median_R_ci": boot_median(R),
                "median_age": median([r["age"] for r in b]) if b else None,
            })
        return out

    def spearman(x, y):
        def rank(v):
            o = sorted(range(len(v)), key=lambda i: v[i]); r = [0.0] * len(v); i = 0
            while i < len(o):
                j = i
                while j + 1 < len(o) and v[o[j + 1]] == v[o[i]]:
                    j += 1
                for k in range(i, j + 1):
                    r[o[k]] = (i + j) / 2 + 1
                i = j + 1
            return r
        rx, ry = rank(x), rank(y); n = len(x)
        mx, my = sum(rx) / n, sum(ry) / n
        sxy = sum((a - mx) * (b - my) for a, b in zip(rx, ry))
        sx = math.sqrt(sum((a - mx) ** 2 for a in rx)); sy = math.sqrt(sum((b - my) ** 2 for b in ry))
        return sxy / (sx * sy)

    primary = bin_table(sites, c)
    rho = spearman([r["q"] ** -2 for r in sites], [r["age"] for r in sites])
    gdh1 = bin_table(sites, 510.0)
    assessed = [r for r in sites if r["qual"] and not r["qual"].startswith("Ux")]
    assessed_tab = bin_table(assessed, c)

    def binof(lo, hi, tab):
        return next(b for b in tab if b["lo"] == lo and b["hi"] == hi)

    young = [kelvin_age(r["q"], c) / r["age"] for r in sites if 0 < r["age"] < 20]
    P1 = median(young) > 1
    b100, b140 = binof(100, 140, primary), binof(140, 200, primary)
    P2 = (b100["median_R"] < 1 and b140["median_R"] < 1 and
          50 <= b100["median_tK"] <= 150 and 50 <= b140["median_tK"] <= 150)
    mid = [b for b in primary if b["lo"] >= 20 and b["hi"] <= 80 and b["n"] >= 50]
    P3 = any(0.67 <= b["median_R"] <= 1.5 for b in mid)
    P4 = rho > 0

    # --- secondary, not pre-registered: one vote per 1-degree cell (many surveys put dozens of
    # probes within a few km, and the bootstrap above treats them as independent)
    cells = {}
    for r in sites:
        key = (math.floor(r["lat"]), math.floor(r["lon"]))
        cells.setdefault(key, []).append(r)
    cell_sites = [{"q": median([x["q"] for x in v]), "age": median([x["age"] for x in v])} for v in cells.values()]
    declustered = bin_table(cell_sites, c)
    # --- context, not pre-registered: the same clock on land
    land = [r for r in rows if r["env"].lower().startswith("[onshore (continental)") and r["q"] > 0]
    land_tk = sorted(kelvin_age(r["q"], c) for r in land)
    grads = sorted(land_gradients())

    summary = {
        "counts": counts, "C_primary": c, "params": PARAMS,
        "bins_primary": primary, "bins_gdh1": gdh1,
        "bins_assessed_only": assessed_tab, "n_assessed": len(assessed),
        "median_R_under_20": median(young), "n_under_20": len(young),
        "spearman_qm2_age": rho,
        "n_cells": len(cell_sites), "bins_declustered_1deg": declustered,
        "land": {"n": len(land), "median_q": median([r["q"] for r in land]), "median_tK": median(land_tk),
                 "tK_q25": land_tk[len(land_tk)//4], "tK_q75": land_tk[3*len(land_tk)//4],
                 "tK_p90": land_tk[int(0.9 * len(land_tk))], "n_over_1000": sum(t > 1000 for t in land_tk),
                 "gradient_n": len(grads), "gradient_median_K_per_km": median(grads),
                 "gradient_q25": grads[len(grads)//4], "gradient_q75": grads[3*len(grads)//4]},
        "predictions": {"P1": P1, "P2": P2, "P3": P3, "P4": P4},
        "P2_detail": {"tK_100_140": b100["median_tK"], "tK_140_200": b140["median_tK"],
                      "R_100_140": b100["median_R"], "R_140_200": b140["median_R"]},
    }
    PUB = os.path.join(HERE, "..", "..", "public", "strata", "kelvin-age-of-the-earth", "data")
    os.makedirs(PUB, exist_ok=True)
    summary["sources"] = {
        "heat_flow": "IHFC Global Heat Flow Database, Release 2024, v.2026.03, doi:10.5880/fidgeo.2024.014, CC BY 4.0",
        "seafloor_age": "Seton et al. (2020) age.2020.1.GTS2012.6m.nc, doi:10.5281/zenodo.6782543 (v1.1), CC BY 4.0",
    }
    with open(os.path.join(PUB, "summary.json"), "w") as f:
        json.dump(summary, f, indent=1)
    # per-site table: the database's parent id (R24-P000001 -> 1), lat, lon, q, the grid's age at
    # the nearest node. q and age are written exactly as read, so the verifier can re-read both.
    def num(pid):
        return int(pid.split("-P")[1])
    sites_out = [[num(r["id"]), r["lat"], r["lon"], r["q"], round(r["age"], 4)] for r in sites]
    with open(os.path.join(PUB, "sites.json"), "w") as f:
        json.dump({"columns": ["parent_id (R24-P######)", "lat", "lon", "q_mWm2", "age_Myr"], "rows": sites_out}, f, separators=(",", ":"))
    print(json.dumps({k: summary[k] for k in ["counts", "C_primary", "median_R_under_20", "spearman_qm2_age", "predictions", "P2_detail"]}, indent=1))
    print("declustered", len(cell_sites)); [print(b["lo"], b["hi"], b["n"], round(b["median_tK"],1) if b["median_tK"] else None, round(b["median_R"],2) if b["median_R"] else None) for b in declustered]
    print("land", summary["land"])
    print("gdh1"); [print(b["lo"], b["hi"], round(b["median_R"],2)) for b in gdh1]
    print("assessed", len(assessed)); [print(b["lo"], b["hi"], b["n"], round(b["median_R"],2) if b["median_R"] else None) for b in assessed_tab]
    for b in primary:
        print(b["lo"], b["hi"], b["n"], b["median_q"], None if b["median_tK"] is None else round(b["median_tK"], 1),
              None if b["median_R"] is None else round(b["median_R"], 2), b["median_R_ci"])


if __name__ == "__main__":
    main()