#!/usr/bin/env python3
"""Does Foster & Rahmstorf's finding that warming accelerated after about 2014 survive a null
simulation of their whole procedure? Reads data/parameters.json and the two Zenodo files
data/external.json points at; writes results/R1.json (their tests on their data),
results/R2.json (the full-pipeline null), results/null_draws.csv, and results/qtable.json."""
import functools, io, json, sys, time, zipfile
from pathlib import Path
import numpy as np, pandas as pd
from scipy import signal, stats
from statsmodels.tsa.arima.model import ARIMA

ROOT = Path(sys.argv[1] if len(sys.argv) > 1 else ".")
sys.path.insert(0, str(ROOT / "code"))
import fr, tmax  # noqa: E402

P = json.loads((ROOT / "data/parameters.json").read_text())
OUT = ROOT / "results"
OUT.mkdir(exist_ok=True)
SS = np.random.SeedSequence(P["seed"])
STREAMS = dict(zip(["qtable", "observed", "null"], SS.spawn(3)))
T0 = time.time()


def log(*a):
    print(f"[{time.time() - T0:7.1f}s]", *a, flush=True)


def r6(x):
    if isinstance(x, (np.integer,)):
        return int(x)
    if isinstance(x, (float, np.floating)):
        return float(f"{float(x):.6g}")
    if isinstance(x, dict):
        return {k: r6(v) for k, v in x.items()}
    if isinstance(x, (list, tuple)):
        return [r6(v) for v in x]
    return x


# ---------------------------------------------------------------- data
with zipfile.ZipFile(ROOT / "data/foster-rahmstorf/AdjustData.zip") as z:
    obs = pd.read_csv(io.BytesIO(z.read("ObservedData.csv")))
pub = pd.read_csv(ROOT / "data/foster-rahmstorf/GMST_adjusted.csv", comment="#")
volc_all = np.nan_to_num(obs.volc.to_numpy(float), nan=0.0)
Y0, Y1 = P["test_years"]
N_YEARS = Y1 - Y0 + 1
TEST = tmax.TmaxTest(N_YEARS)


def series(name):
    ok = np.isfinite(obs[name].to_numpy(float))
    t = obs.t.to_numpy(float)[ok]
    return t, obs[name].to_numpy(float)[ok], fr.Adjuster(t, volc_all[ok], obs.NINO34.to_numpy(float)[ok],
                                                          obs.ssn.to_numpy(float)[ok], fit=tuple(P["fit_window"]),
                                                          nback=P["nback"])


def quad_test(t, x):
    """The quadratic test with Foster & Rahmstorf (2011)'s ARMA(1,1) correction: the OLS standard error of
    the quadratic coefficient times sqrt(nu), nu = 1 + 2 r1 / (1 - r2 / r1), from the residuals' lag-1 and
    lag-2 autocorrelations. Returns the acceleration (degC per decade^2) and z."""
    tc = t - t.mean()
    X = np.column_stack([np.ones_like(tc), tc, tc * tc])
    b = np.linalg.lstsq(X, x, rcond=None)[0]
    e = x - X @ b
    n = len(x)
    se = np.sqrt(e @ e / (n - 3) * np.linalg.inv(X.T @ X)[2, 2])
    ec = e - e.mean()
    r1, r2 = (ec[1:] @ ec[:-1]) / (ec @ ec), (ec[2:] @ ec[:-2]) / (ec @ ec)
    nu = 1 + 2 * r1 / (1 - r2 / r1) if r1 > 0 and r2 / r1 < 1 else 1.0
    nu = max(nu, 1.0)
    return 200 * b[2], b[2] / (se * np.sqrt(nu))


def tests_on(t, z):
    """Both of Foster & Rahmstorf's tests on one monthly series: annual means for years Y0..Y1 for the
    changepoint test, monthly and annual values for the quadratic test."""
    years, ann = fr.annual(t, z, Y0, Y1)
    m = (t >= Y0) & (t < Y1 + 1)
    tm, k, _ = TEST.stats(ann[None])
    phi, _, _ = tmax.ar1_null_fit(ann)
    _, zq_m = quad_test(t[m], z[m])
    _, zq_a = quad_test(years + 0.5, ann)
    return dict(tmax=float(tm[0]), change_year=int(Y0 - 1 + k[0]), phi=phi, zq_monthly=float(zq_m), zq_annual=float(zq_a))


# ---------------------------------------------------------------- critical values of T_max under AR(1)
lo, hi, step = P["qtable_phi"]
PHI_GRID = np.round(np.arange(lo, hi + step / 2, step), 10)
rng_q = np.random.default_rng(STREAMS["qtable"])
QT = []
for ph in PHI_GRID:
    vals = np.concatenate([TEST.stats(tmax.simulate_ar1(rng_q, ph, N_YEARS, P["batch"]))[0]
                           for _ in range(P["qtable_sims"] // P["batch"])])
    QT.append(np.sort(vals))
QT = np.array(QT)
log("critical-value table done", QT.shape)


def nominal_p(T, phi):
    """The AR(1) Monte Carlo p-value of T_max at the fitted phi, interpolated between table rows."""
    T, phi = np.atleast_1d(T), np.clip(np.atleast_1d(phi), PHI_GRID[0], PHI_GRID[-1])
    i = np.clip(np.searchsorted(PHI_GRID, phi) - 1, 0, len(PHI_GRID) - 2)
    w = (phi - PHI_GRID[i]) / (PHI_GRID[i + 1] - PHI_GRID[i])
    n = QT.shape[1]
    sf = lambda row, x: (1 + n - np.searchsorted(QT[row], x, side="left")) / (n + 1)
    return (1 - w) * np.array([sf(r, x) for r, x in zip(i, T)]) + w * np.array([sf(r + 1, x) for r, x in zip(i, T)])


json.dump(r6({"phi": PHI_GRID.tolist(), "sims": int(QT.shape[1]),
              "q95": np.quantile(QT, 0.95, axis=1).tolist(), "q99": np.quantile(QT, 0.99, axis=1).tolist()}),
          open(OUT / "qtable.json", "w"), indent=1)

# ---------------------------------------------------------------- their tests on their data
rng_o = np.random.default_rng(STREAMS["observed"])
R1, WORLD = {}, {}
for name in P["datasets"]:
    t, x, A = series(name)
    adj, k, coef = A.adjust(x)
    pj = pub.set_index(np.round(pub.time, 4))[name + ".adj"].reindex(np.round(t, 4)).to_numpy()
    dev = (pj - adj) - np.nanmean(pj - adj)
    res = tests_on(t, adj)
    raw = tests_on(t, x)
    sims = np.concatenate([TEST.stats(tmax.simulate_ar1(rng_o, res["phi"], N_YEARS, P["batch"]))[0]
                           for _ in range(P["observed_sims"] // P["batch"])])
    R1[name] = {
        "lags_months": {"volcanic": int(k[0]), "nino34": int(k[1]), "sunspots": int(k[2])},
        "coef": {"volcanic": coef[1], "nino34": coef[2], "sunspots": coef[3]},
        "max_abs_dev_from_published_adjusted": float(np.nanmax(np.abs(dev))),
        "adjusted": {**res, "q95": float(np.quantile(sims, 0.95)), "q99": float(np.quantile(sims, 0.99)),
                     "p_tmax": float((1 + (sims >= res["tmax"]).sum()) / (1 + len(sims))),
                     "p_quad_monthly": float(2 * stats.norm.sf(abs(res["zq_monthly"]))),
                     "p_quad_annual": float(2 * stats.norm.sf(abs(res["zq_annual"])))},
        "unadjusted": {**raw, "p_tmax": float(nominal_p(raw["tmax"], raw["phi"])[0]),
                       "p_quad_monthly": float(2 * stats.norm.sf(abs(raw["zq_monthly"]))),
                       "p_quad_annual": float(2 * stats.norm.sf(abs(raw["zq_annual"])))},
    }
    WORLD[name] = (t, x, A, adj, coef, k)
    log(name, "observed", R1[name]["adjusted"])
R1["all_adjusted_p_tmax_below_0_01"] = all(R1[n]["adjusted"]["p_tmax"] < 0.01 for n in P["datasets"])
R1["all_unadjusted_p_tmax_at_least_0_05"] = all(R1[n]["unadjusted"]["p_tmax"] >= 0.05 for n in P["datasets"])
json.dump(r6(R1), open(OUT / "R1.json", "w"), indent=1, sort_keys=True)


# ---------------------------------------------------------------- the null world and its simulations
def null_world(t, adj, coef, k, A):
    """No acceleration since Y0: a straight line fitted to the adjusted series from Y0 on, and before Y0
    its 20-year lowess, shifted to meet the line; plus the exogenous terms as fitted to the real data."""
    post = t >= Y0
    m = post & (t < Y1 + 1)
    a, b = np.polyfit(t[m] - Y0, adj[m], 1)[::-1]
    smooth = fr.lowess_r(t, adj, 20 / (t.max() - t.min()))
    g0 = np.where(post, a + b * (t - Y0), 0.0)
    first = np.argmax(post)
    g0[:first] = smooth[:first] + (g0[first] - smooth[first])
    exo = sum(coef[j + 1] * A.L[j][k[j]] for j in range(3))
    return g0, exo, smooth, b


def fit_noise(resid):
    f = ARIMA(resid, order=(1, 0, 1), trend="n").fit()
    ar, ma, s2 = f.params[0], f.params[1], f.params[2]
    return {"ar": float(ar), "ma": float(ma), "sd": float(np.sqrt(s2))}


@functools.lru_cache(maxsize=None)
def frac_weights(d, n):
    """Coefficients of (1 - B)^-d up to lag n - 1 (negate d for (1 - B)^d)."""
    w = np.ones(n)
    for j in range(1, n):
        w[j] = w[j - 1] * (j - 1 + d) / j
    return w


def fit_long_memory(resid):
    """ARFIMA(1, d, 0): d from the Geweke-Porter-Hudak log-periodogram regression over the lowest
    floor(sqrt(n)) Fourier frequencies, clipped to [0, 0.45]; then an AR(1) fitted by least squares to
    the series fractionally differenced by d."""
    x = resid - resid.mean()
    n = len(x)
    m = int(np.floor(np.sqrt(n)))
    lam = 2 * np.pi * np.arange(1, m + 1) / n
    I = np.abs(np.fft.fft(x)[1:m + 1]) ** 2 / (2 * np.pi * n)
    slope = np.polyfit(np.log(4 * np.sin(lam / 2) ** 2), np.log(I), 1)[0]
    d = float(np.clip(-slope, 0.0, 0.45))
    u = np.convolve(x, frac_weights(-d, n))[:n]
    ar = float(np.dot(u[1:], u[:-1]) / np.dot(u[:-1], u[:-1]))
    sd = float(np.std(u[1:] - ar * u[:-1]))
    return {"d": d, "ar": ar, "sd": sd}


def simulate_noise(rng, model, prm, n, burn):
    z = prm["sd"] * rng.standard_normal(n + burn)
    if model == "long_memory":
        u = signal.lfilter([1.0], [1.0, -prm["ar"]], z)
        e = signal.fftconvolve(u, frac_weights(prm["d"], n + burn))[:n + burn]
    else:
        e = signal.lfilter([1.0, prm["ma"]], [1.0, -prm["ar"]], z)
    return e[burn:]


rng_n = np.random.default_rng(STREAMS["null"])
R2, draws = {"noise_models": {}}, []
for model in P["noise_models"]:
    R2["noise_models"][model] = {}
    for name in P["datasets"]:
        t, x, A, adj, coef, k = WORLD[name]
        g0, exo, smooth, slope = null_world(t, adj, coef, k, A)
        if model == "lowess":
            m = (t >= P["fit_window"][0]) & (t < P["fit_window"][1])
            resid = (adj - smooth)[m]
        else:
            m = (t >= Y0) & (t < Y1 + 1)
            resid = (adj - g0)[m]
        prm = fit_long_memory(resid) if model == "long_memory" else fit_noise(resid)
        out = []
        for i in range(P["null_sims"]):
            xi = g0 + exo + simulate_noise(rng_n, model, prm, len(t), P["burn_in_months"])
            zi, _, _ = A.adjust(xi)
            out.append(tests_on(t, zi))
        d = pd.DataFrame(out)
        d["p_nominal"] = nominal_p(d.tmax.to_numpy(), d.phi.to_numpy())
        o = R1[name]["adjusted"]
        S = len(d)
        R2["noise_models"][model][name] = {
            "noise": {**prm, "months": int(len(resid))},
            "null_trend_per_decade": 10 * slope,
            "size_tmax_005": float((d.p_nominal < 0.05).mean()),
            "size_tmax_001": float((d.p_nominal < 0.01).mean()),
            "tmax_q95": float(d.tmax.quantile(0.95)), "tmax_q99": float(d.tmax.quantile(0.99)),
            "p_tmax": float((1 + (d.tmax >= o["tmax"]).sum()) / (1 + S)),
            "size_quad_monthly_005": float((2 * stats.norm.sf(d.zq_monthly.abs()) < 0.05).mean()),
            "size_quad_annual_005": float((2 * stats.norm.sf(d.zq_annual.abs()) < 0.05).mean()),
            "p_quad_monthly": float((1 + (d.zq_monthly.abs() >= abs(o["zq_monthly"])).sum()) / (1 + S)),
            "p_quad_annual": float((1 + (d.zq_annual.abs() >= abs(o["zq_annual"])).sum()) / (1 + S)),
            "phi_hat_median": float(d.phi.median()),
            "phi_hat_outside_table": int(((d.phi < PHI_GRID[0]) | (d.phi > PHI_GRID[-1])).sum()),
            "sims": S,
        }
        draws.append(d.assign(noise_model=model, dataset=name))
        log(model, name, {k_: R2["noise_models"][model][name][k_] for k_ in ("size_tmax_005", "p_tmax", "p_quad_monthly", "p_quad_annual")})


def verdict(ps, strong, weak):
    if all(p < strong for p in ps):
        return "robust"
    if sum(p >= weak for p in ps) >= 3:
        return "not_robust"
    return "mixed"


prim = R2["noise_models"][P["noise_models"][0]]
R2["decision"] = {
    "tmax": verdict([prim[n]["p_tmax"] for n in P["datasets"]], 0.01, 0.05),
    "quad_monthly": verdict([prim[n]["p_quad_monthly"] for n in P["datasets"]], 0.02, 0.05),
    "quad_annual": verdict([prim[n]["p_quad_annual"] for n in P["datasets"]], 0.02, 0.05),
}
R2["primary_noise_model"] = P["noise_models"][0]
json.dump(r6(R2), open(OUT / "R2.json", "w"), indent=1, sort_keys=True)
pd.concat(draws).to_csv(OUT / "null_draws.csv", index=False, float_format="%.6g")
log("done")