#!/usr/bin/env python3
"""
Do national-inventory prioritization choices matter more for industrialized
countries? A paired secondary analysis of a historical GHG emissions dataset.

Self-contained analysis script.
  - Downloads the four pinned source CSV files.
  - Verifies MD5 checksums against the published values (records observed).
  - Runs the pre-registered confirmatory tests and robustness checks.
  - Writes all figures (figures/), tables (tables/ + datapackage.json),
    and results.json.
Re-running this script reproduces every reported number exactly.

Pre-registered question
------------------------
The dataset provides, for each national area, two parallel estimates of the
same emissions time series built with different source-prioritization rules:
one that prioritizes country-reported inventory data ("CR") and one that
prioritizes third-party datasets ("TP"). For the national total GHG estimate
(aggregate Kyoto-GHG basket, AR4 GWP-100, national total excl. LULUCF,
year 2019), does the choice of prioritization rule change the estimate MORE
for Annex I (industrialized) countries than for the rest?

  H1 (directional): the per-country absolute symmetric relative difference
      |2*(TP-CR)/(TP+CR)| is stochastically larger for Annex I countries.
  H0: the two groups have the same distribution of |rel diff|.

Primary test  : two-sided Mann-Whitney U on |rel diff|, alpha = 0.05.
Effect size   : Cliff's delta with a 10,000-resample bootstrap 95% CI.
Secondary test: Fisher's exact on (Annex I vs not) x (scenarios exactly
      equal vs not) -- a proxy for whether country-reported data exist at all.

Robustness    : (a) repeat primary test for 1990/2000/2010/2019;
                (b) recompute 2019 from the no-rounding source variant;
                (c) repeat 2019 test for CO2 and CH4 separately.
"""
from __future__ import annotations
import hashlib
import json
import os
import sys
import urllib.request
from pathlib import Path

import numpy as np
import pandas as pd
from scipy import stats
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

# --------------------------------------------------------------------------
# 0. Determinism
# --------------------------------------------------------------------------
SEED = 42
RNG = np.random.default_rng(SEED)
N_BOOT = 10_000

# --------------------------------------------------------------------------
# 1. Pinned inputs
# --------------------------------------------------------------------------
BASE = "https://zenodo.org/api/records/5175154/files/{name}/content"
FILES = {
    "main": (
        "Guetschow-et-al-2021-PRIMAP-hist_v2.3_28_Jul_2021.csv",
        "7b29dfb49ca97dcf594bc639767f1e2a",
    ),
    "no_extrap": (
        "Guetschow-et-al-2021-PRIMAP-hist_v2.3_no_extrap_28_Jul_2021.csv",
        "127989b66294d50873ef22a50fe60d14",
    ),
    "no_extrap_no_rounding": (
        "Guetschow-et-al-2021-PRIMAP-hist_v2.3_no_extrap_no_rounding_28_Jul_2021.csv",
        "88fda69069c09e505ca9d5e6cdfc18fe",
    ),
    "no_rounding": (
        "Guetschow-et-al-2021-PRIMAP-hist_v2.3_no_rounding_28_Jul_2021.csv",
        "99aa3f88ca69be675dacafd5bb32e766",
    ),
}

HERE = Path(__file__).resolve().parent
DATA = HERE / "data"
FIGS = HERE / "figures"
TABS = HERE / "tables"
for d in (DATA, FIGS, TABS):
    d.mkdir(parents=True, exist_ok=True)

# --------------------------------------------------------------------------
# 2. Fixed group definitions (frozen at pre-registration)
# --------------------------------------------------------------------------
# Annex I Parties to the UNFCCC, as ISO3 codes. 40 members.
# Note: Cyprus and Malta are NOT Annex I; the EU is not a national ISO3 area.
ANNEX1 = {
    "AUS", "AUT", "BLR", "BEL", "BGR", "CAN", "HRV", "CZE", "DNK", "EST",
    "FIN", "FRA", "DEU", "GRC", "HUN", "ISL", "IRL", "ITA", "JPN", "LVA",
    "LIE", "LTU", "LUX", "MCO", "NLD", "NZL", "NOR", "POL", "PRT", "ROU",
    "RUS", "SVK", "SVN", "ESP", "SWE", "CHE", "TUR", "UKR", "GBR", "USA",
}
# Non-country / supranational / region aggregates present in the 'area (ISO3)'
# column -- excluded so the comparison is country-to-country.
NON_COUNTRY = {
    "EARTH", "ANNEXI", "NONANNEXI", "AOSIS", "BASIC", "LDC", "UMBRELLA",
    "EU27BX", "ANT", "ATA",
}

ENTITY_TOTAL = "KYOTOGHG (AR4GWP100)"   # aggregate Kyoto basket, AR4 GWP-100
CAT_TOTAL = "M.0.EL"                     # national total excluding LULUCF
SCEN_CR = "HISTCR"                        # country-reported prioritized
SCEN_TP = "HISTTP"                        # third-party prioritized

# --------------------------------------------------------------------------
# 3. Download + verify
# --------------------------------------------------------------------------
def md5(path: Path) -> str:
    h = hashlib.md5()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def fetch(name: str, expected: str) -> tuple[Path, str, bool]:
    path = DATA / name
    if not (path.exists() and md5(path) == expected):
        url = BASE.format(name=name)
        for attempt in range(3):
            urllib.request.urlretrieve(url, path)
            if md5(path) == expected:
                break
    observed = md5(path)
    return path, observed, observed == expected


# --------------------------------------------------------------------------
# 4. Statistics helpers
# --------------------------------------------------------------------------
def sym_rel_diff(tp: np.ndarray, cr: np.ndarray) -> np.ndarray:
    """Symmetric relative difference 2*(TP-CR)/(TP+CR)."""
    denom = (tp + cr) / 2.0
    out = np.where(denom != 0, (tp - cr) / denom, 0.0)
    return out


def cliffs_delta(a: np.ndarray, b: np.ndarray) -> float:
    """Cliff's delta = P(a>b) - P(a<b), computed via rank sums (exact)."""
    a = np.asarray(a, float)
    b = np.asarray(b, float)
    n_a, n_b = len(a), len(b)
    # rank-based O(n log n): U statistic relation
    combined = np.concatenate([a, b])
    ranks = stats.rankdata(combined)
    r_a = ranks[:n_a].sum()
    u_a = r_a - n_a * (n_a + 1) / 2.0        # Mann-Whitney U for group a
    delta = (2.0 * u_a) / (n_a * n_b) - 1.0
    return float(delta)


def bootstrap_delta_ci(a: np.ndarray, b: np.ndarray, rng, n_boot=N_BOOT,
                       alpha=0.05) -> tuple[float, float]:
    a = np.asarray(a, float)
    b = np.asarray(b, float)
    n_a, n_b = len(a), len(b)
    deltas = np.empty(n_boot)
    for i in range(n_boot):
        sa = a[rng.integers(0, n_a, n_a)]
        sb = b[rng.integers(0, n_b, n_b)]
        deltas[i] = cliffs_delta(sa, sb)
    lo, hi = np.percentile(deltas, [100 * alpha / 2, 100 * (1 - alpha / 2)])
    return float(lo), float(hi)


def national_totals(df: pd.DataFrame, entity: str, year: str) -> pd.DataFrame:
    """Return per-country CR/TP totals for one entity/year (national total)."""
    sub = df[(df["entity"] == entity) &
             (df["category (IPCC2006_PRIMAP)"] == CAT_TOTAL)]
    wide = sub.pivot_table(index="area (ISO3)",
                           columns="scenario (PRIMAP-hist)",
                           values=year)
    wide = wide.dropna(subset=[SCEN_CR, SCEN_TP])
    wide = wide[~wide.index.isin(NON_COUNTRY)]
    wide = wide.copy()
    wide["is_annex1"] = wide.index.isin(ANNEX1)
    wide["rel_diff"] = sym_rel_diff(wide[SCEN_TP].values, wide[SCEN_CR].values)
    wide["abs_rel_diff"] = np.abs(wide["rel_diff"])
    wide["exact_equal"] = wide[SCEN_CR].values == wide[SCEN_TP].values
    return wide


def primary_test(wide: pd.DataFrame, rng) -> dict:
    """Mann-Whitney U + Cliff's delta on |rel diff|, Annex I vs non-Annex I."""
    a = wide.loc[wide["is_annex1"], "abs_rel_diff"].values   # Annex I
    b = wide.loc[~wide["is_annex1"], "abs_rel_diff"].values  # non-Annex I
    u, p = stats.mannwhitneyu(a, b, alternative="two-sided")
    delta = cliffs_delta(a, b)
    lo, hi = bootstrap_delta_ci(a, b, rng)
    return {
        "n_annex1": int(len(a)),
        "n_non_annex1": int(len(b)),
        "median_abs_rel_diff_annex1": float(np.median(a)),
        "median_abs_rel_diff_non_annex1": float(np.median(b)),
        "mann_whitney_u": float(u),
        "p_value": float(p),
        "cliffs_delta": float(delta),
        "cliffs_delta_ci95": [lo, hi],
    }


def secondary_test(wide: pd.DataFrame) -> dict:
    """Fisher's exact: Annex I vs not  x  exactly-equal vs not."""
    a1 = wide["is_annex1"].values
    eq = wide["exact_equal"].values
    # 2x2: rows = Annex I / non; cols = not-equal / equal
    table = np.array([
        [int(((a1) & (~eq)).sum()), int(((a1) & (eq)).sum())],
        [int(((~a1) & (~eq)).sum()), int(((~a1) & (eq)).sum())],
    ])
    odds, p = stats.fisher_exact(table, alternative="two-sided")
    return {
        "contingency": {
            "annex1_notequal": int(table[0, 0]),
            "annex1_equal": int(table[0, 1]),
            "nonannex1_notequal": int(table[1, 0]),
            "nonannex1_equal": int(table[1, 1]),
        },
        "frac_equal_annex1": float(table[0, 1] / table[0].sum()),
        "frac_equal_non_annex1": float(table[1, 1] / table[1].sum()),
        "odds_ratio": float(odds),
        "p_value": float(p),
    }


# --------------------------------------------------------------------------
# 5. Figures (colourblind-safe Okabe-Ito palette; no tool names)
# --------------------------------------------------------------------------
CB = {"annex1": "#D55E00", "non": "#0072B2", "grey": "#999999"}
plt.rcParams.update({
    "figure.dpi": 150, "savefig.dpi": 150,
    "font.size": 11, "axes.titlesize": 12, "axes.labelsize": 11,
    "axes.spines.top": False, "axes.spines.right": False,
    "xtick.direction": "out", "ytick.direction": "out",
})


def fig1_distribution(wide: pd.DataFrame, path: Path):
    a = wide.loc[wide["is_annex1"], "abs_rel_diff"].values
    b = wide.loc[~wide["is_annex1"], "abs_rel_diff"].values
    eps = 1e-4  # floor so exact-equal (0) points remain visible on log axis
    fig, ax = plt.subplots(figsize=(7, 4.5))
    for i, (vals, lab, col) in enumerate([
            (b, f"Non-Annex I (n={len(b)})", CB["non"]),
            (a, f"Annex I (n={len(a)})", CB["annex1"])]):
        y = np.full(len(vals), i) + RNG.uniform(-0.12, 0.12, len(vals))
        ax.scatter(np.maximum(vals, eps), y, s=18, alpha=0.55,
                   color=col, edgecolor="none", zorder=3)
        med = np.median(vals)
        ax.plot([max(med, eps)] * 2, [i - 0.28, i + 0.28],
                color="black", lw=2.2, zorder=4)
    ax.set_xscale("log")
    ax.set_yticks([0, 1])
    ax.set_yticklabels(["Non-Annex I", "Annex I"])
    ax.set_xlabel("Absolute symmetric relative difference between\n"
                  "prioritization scenarios, 2019 national total GHG "
                  "(log scale; floored at 1e-4)")
    ax.set_title("Prioritization choice shifts industrialized-country totals more")
    ax.axvline(0.01, color=CB["grey"], ls=":", lw=1)
    ax.text(0.01, 1.55, "1%", color=CB["grey"], ha="center", fontsize=9)
    fig.tight_layout()
    fig.savefig(path)
    plt.close(fig)


def fig2_exact_agreement(sec: dict, path: Path):
    fig, ax = plt.subplots(figsize=(5.2, 4.5))
    groups = ["Non-Annex I", "Annex I"]
    fracs = [sec["frac_equal_non_annex1"], sec["frac_equal_annex1"]]
    counts = [(sec["contingency"]["nonannex1_equal"],
               sec["contingency"]["nonannex1_equal"] + sec["contingency"]["nonannex1_notequal"]),
              (sec["contingency"]["annex1_equal"],
               sec["contingency"]["annex1_equal"] + sec["contingency"]["annex1_notequal"])]
    cols = [CB["non"], CB["annex1"]]
    bars = ax.bar(groups, fracs, color=cols, width=0.6, zorder=3)
    for bar, (k, n) in zip(bars, counts):
        ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.02,
                f"{k}/{n}", ha="center", va="bottom", fontsize=10)
    ax.set_ylim(0, 1)
    ax.set_ylabel("Fraction of countries where the two\nscenarios are exactly identical (2019)")
    ax.set_title("Where country-reported data are absent,\nboth scenarios coincide")
    fig.tight_layout()
    fig.savefig(path)
    plt.close(fig)


def fig3_temporal(years_res: dict, path: Path):
    years = sorted(int(y) for y in years_res)
    deltas = [years_res[str(y)]["cliffs_delta"] for y in years]
    los = [years_res[str(y)]["cliffs_delta_ci95"][0] for y in years]
    his = [years_res[str(y)]["cliffs_delta_ci95"][1] for y in years]
    fig, ax = plt.subplots(figsize=(6.5, 4.3))
    yerr = np.array([np.array(deltas) - np.array(los),
                     np.array(his) - np.array(deltas)])
    ax.errorbar(years, deltas, yerr=yerr, fmt="o-", color=CB["annex1"],
                capsize=4, lw=2, markersize=7, zorder=3)
    ax.axhline(0, color=CB["grey"], ls="--", lw=1)
    ax.set_ylim(-0.1, 1.0)
    ax.set_xlabel("Year")
    ax.set_ylabel("Cliff's delta (Annex I vs non-Annex I)\nof |scenario difference|")
    ax.set_title("Effect persists across three decades")
    ax.set_xticks(years)
    fig.tight_layout()
    fig.savefig(path)
    plt.close(fig)


def fig4_by_gas(gas_res: dict, path: Path):
    gases = list(gas_res)
    deltas = [gas_res[g]["cliffs_delta"] for g in gases]
    los = [gas_res[g]["cliffs_delta_ci95"][0] for g in gases]
    his = [gas_res[g]["cliffs_delta_ci95"][1] for g in gases]
    fig, ax = plt.subplots(figsize=(6, 4.3))
    ypos = np.arange(len(gases))
    xerr = np.array([np.array(deltas) - np.array(los),
                     np.array(his) - np.array(deltas)])
    ax.errorbar(deltas, ypos, xerr=xerr, fmt="o", color=CB["annex1"],
                capsize=4, markersize=8, lw=2, zorder=3)
    ax.axvline(0, color=CB["grey"], ls="--", lw=1)
    ax.set_yticks(ypos)
    ax.set_yticklabels(gases)
    ax.set_xlim(-0.2, 1.0)
    ax.set_xlabel("Cliff's delta (Annex I vs non-Annex I) of |scenario difference|, 2019")
    ax.set_title("Same direction for individual gases")
    fig.tight_layout()
    fig.savefig(path)
    plt.close(fig)


# --------------------------------------------------------------------------
# 6. Main
# --------------------------------------------------------------------------
def main():
    results = {
        "meta": {
            "question": ("Does the prioritization scenario (country-reported vs "
                         "third-party) change national total GHG estimates more "
                         "for Annex I than non-Annex I countries?"),
            "seed": SEED, "n_bootstrap": N_BOOT,
            "entity": ENTITY_TOTAL, "category": CAT_TOTAL,
            "primary_year": "2019",
            "annex1_definition_iso3": sorted(ANNEX1),
            "n_annex1_defined": len(ANNEX1),
            "non_country_excluded": sorted(NON_COUNTRY),
        },
        "provenance": {},
    }

    # ---- download + verify all pinned files
    for key, (name, exp) in FILES.items():
        _, obs, ok = fetch(name, exp)
        results["provenance"][name] = {
            "expected_md5": exp, "observed_md5": obs, "verified": bool(ok)}
        if not ok:
            print(f"WARNING checksum mismatch: {name}", file=sys.stderr)

    df_main = pd.read_csv(DATA / FILES["main"][0])

    # ---- CONFIRMATORY -----------------------------------------------------
    wide19 = national_totals(df_main, ENTITY_TOTAL, "2019")
    t1 = primary_test(wide19, np.random.default_rng(SEED))
    t2 = secondary_test(wide19)
    results["confirmatory"] = {"test1_mann_whitney": t1, "test2_fisher_exact": t2}

    fig1_distribution(wide19, FIGS / "fig-1-scenario-divergence-by-group.png")
    fig2_exact_agreement(t2, FIGS / "fig-2-exact-agreement-fraction.png")

    # ---- ROBUSTNESS -------------------------------------------------------
    robustness = {}

    # (a) temporal
    years_res = {}
    for y in ["1990", "2000", "2010", "2019"]:
        wy = national_totals(df_main, ENTITY_TOTAL, y)
        years_res[y] = primary_test(wy, np.random.default_rng(SEED))
    robustness["temporal"] = years_res
    fig3_temporal(years_res, FIGS / "fig-3-effect-over-time.png")

    # (b) rounding sensitivity (recompute 2019 from no_rounding variant)
    df_nr = pd.read_csv(DATA / FILES["no_rounding"][0])
    wide19_nr = national_totals(df_nr, ENTITY_TOTAL, "2019")
    t1_nr = primary_test(wide19_nr, np.random.default_rng(SEED))
    robustness["rounding"] = {
        "main": {"cliffs_delta": t1["cliffs_delta"], "p_value": t1["p_value"],
                 "n_exact_equal_total": int(wide19["exact_equal"].sum())},
        "no_rounding": {"cliffs_delta": t1_nr["cliffs_delta"],
                        "p_value": t1_nr["p_value"],
                        "n_exact_equal_total": int(wide19_nr["exact_equal"].sum())},
        "cliffs_delta_abs_change": abs(t1["cliffs_delta"] - t1_nr["cliffs_delta"]),
    }

    # (c) gas sensitivity
    gas_res = {}
    for gas in ["CO2", "CH4"]:
        wg = national_totals(df_main, gas, "2019")
        gas_res[gas] = primary_test(wg, np.random.default_rng(SEED))
    robustness["by_gas"] = gas_res
    fig4_by_gas(gas_res, FIGS / "fig-4-effect-by-gas.png")

    results["robustness"] = robustness

    # ---- TABLES -----------------------------------------------------------
    # tbl-1: per-country detail, 2019 confirmatory
    tbl1 = wide19.reset_index()[
        ["area (ISO3)", SCEN_CR, SCEN_TP, "rel_diff", "abs_rel_diff",
         "is_annex1", "exact_equal"]].rename(columns={
            "area (ISO3)": "iso3", SCEN_CR: "value_CR_GgCO2e",
            SCEN_TP: "value_TP_GgCO2e"})
    tbl1 = tbl1.sort_values("abs_rel_diff", ascending=False)
    tbl1.to_csv(TABS / "tbl-1-per-country-2019.csv", index=False)

    # tbl-2: confirmatory test summary
    tbl2 = pd.DataFrame([
        {"test": "Mann-Whitney U (|rel diff|, 2019)",
         "statistic": t1["mann_whitney_u"], "p_value": t1["p_value"],
         "effect_size_name": "Cliffs delta", "effect_size": t1["cliffs_delta"],
         "ci95_low": t1["cliffs_delta_ci95"][0], "ci95_high": t1["cliffs_delta_ci95"][1],
         "n_annex1": t1["n_annex1"], "n_non_annex1": t1["n_non_annex1"]},
        {"test": "Fisher exact (exact-agreement x group)",
         "statistic": t2["odds_ratio"], "p_value": t2["p_value"],
         "effect_size_name": "odds_ratio", "effect_size": t2["odds_ratio"],
         "ci95_low": np.nan, "ci95_high": np.nan,
         "n_annex1": t1["n_annex1"], "n_non_annex1": t1["n_non_annex1"]},
    ])
    tbl2.to_csv(TABS / "tbl-2-confirmatory-summary.csv", index=False)

    # tbl-3: temporal robustness
    tbl3 = pd.DataFrame([
        {"year": y, "cliffs_delta": years_res[y]["cliffs_delta"],
         "ci95_low": years_res[y]["cliffs_delta_ci95"][0],
         "ci95_high": years_res[y]["cliffs_delta_ci95"][1],
         "p_value": years_res[y]["p_value"],
         "median_annex1": years_res[y]["median_abs_rel_diff_annex1"],
         "median_non_annex1": years_res[y]["median_abs_rel_diff_non_annex1"],
         "n_annex1": years_res[y]["n_annex1"],
         "n_non_annex1": years_res[y]["n_non_annex1"]}
        for y in ["1990", "2000", "2010", "2019"]])
    tbl3.to_csv(TABS / "tbl-3-temporal-robustness.csv", index=False)

    # tbl-4: gas + rounding robustness
    rows = []
    for gas in ["CO2", "CH4"]:
        rows.append({"check": f"gas={gas}",
                     "cliffs_delta": gas_res[gas]["cliffs_delta"],
                     "ci95_low": gas_res[gas]["cliffs_delta_ci95"][0],
                     "ci95_high": gas_res[gas]["cliffs_delta_ci95"][1],
                     "p_value": gas_res[gas]["p_value"],
                     "n_annex1": gas_res[gas]["n_annex1"],
                     "n_non_annex1": gas_res[gas]["n_non_annex1"]})
    rows.append({"check": "no_rounding variant (2019 KYOTOGHG)",
                 "cliffs_delta": t1_nr["cliffs_delta"],
                 "ci95_low": t1_nr["cliffs_delta_ci95"][0],
                 "ci95_high": t1_nr["cliffs_delta_ci95"][1],
                 "p_value": t1_nr["p_value"],
                 "n_annex1": t1_nr["n_annex1"],
                 "n_non_annex1": t1_nr["n_non_annex1"]})
    pd.DataFrame(rows).to_csv(TABS / "tbl-4-gas-and-rounding-robustness.csv", index=False)

    # ---- results.json -----------------------------------------------------
    with open(HERE / "results.json", "w") as f:
        json.dump(results, f, indent=2, sort_keys=True)

    print("Confirmatory Test 1 (Mann-Whitney U):",
          f"U={t1['mann_whitney_u']:.1f} p={t1['p_value']:.3e} "
          f"delta={t1['cliffs_delta']:.3f} CI{t1['cliffs_delta_ci95']}")
    print("Confirmatory Test 2 (Fisher exact):",
          f"OR={t2['odds_ratio']:.3f} p={t2['p_value']:.3e} "
          f"frac_equal A1={t2['frac_equal_annex1']:.3f} "
          f"nonA1={t2['frac_equal_non_annex1']:.3f}")
    return results


if __name__ == "__main__":
    main()
