"""Straight-line trends and their 95% intervals, as used on factchecked.science/amoc-heat-loss/.

Imported by rapid_trends.py and deepc_region.py; keep it in the same folder. Needs NumPy and SciPy.

fit(t, y, kmax) fits y = a t + b by ordinary least squares and returns a dict with
- slope_per_yr, intercept, fit_start, fit_end, mean, n and r1 (the lag-1 autocorrelation of the residuals);
- ci95_ols_per_yr: the white-noise 95% half-width, sigma_U * t(0.975, N - 2);
- ci95_ar1_per_yr: with N_eff = N (1 - r1) / (1 + r1), capped at N (the form of the preprint's eq. 6);
- ci95_bell_per_yr: the preprint's own recipe (Bell et al. 2026, eqs. 2-5): N_eff(K) = N / (1 + 2 sum_{k<=K} (1 - k/N) rho_k),
  with rho_k = sum_{i<=N-k} e_i e_{i+k} / sum e_i^2, minimised over 0 <= K <= kmax, N_eff(0) = N;
  sigma_eff = sigma_U sqrt((N - 2) / (N_eff - 2)); half-width = sigma_eff * t(0.975, N_eff - 2).
  The preprint does not say what to do when the denominator is zero or negative; we skip those K
  (treating N_eff as infinite) and report every N_eff(K) so the choice is visible.
"""
import math

import numpy as np
from scipy import stats


def fit(t, y, kmax=9):
    t = np.asarray(t, float)
    y = np.asarray(y, float)
    ok = np.isfinite(y) & np.isfinite(t)
    t, y = t[ok], y[ok]
    n = len(y)
    a, b = np.polyfit(t, y, 1)
    e = y - (a * t + b)
    s2 = (e ** 2).sum() / (n - 2)
    sig_u = math.sqrt(s2 / ((t - t.mean()) ** 2).sum())
    denom = (e ** 2).sum()
    rho = [(e[:-k] * e[k:]).sum() / denom for k in range(1, min(n, kmax + 2))]
    neff_k = [float(n)]
    for k_max in range(1, kmax + 1):
        d = 1 + 2 * sum((1 - k / n) * rho[k - 1] for k in range(1, k_max + 1))
        neff_k.append(n / d if d > 0 else float("inf"))
    neff_bell = min(neff_k)
    r1 = rho[0]
    neff_ar1 = min(float(n), n * (1 - r1) / (1 + r1))

    def half(neff):
        if not math.isfinite(neff) or neff <= 2.5:
            return float("nan")
        return sig_u * math.sqrt((n - 2) / (neff - 2)) * stats.t.ppf(0.975, neff - 2)

    return {
        "n": n, "t0": float(t[0]), "t1": float(t[-1]),
        "slope_per_yr": float(a), "intercept": float(b),
        "fit_start": float(a * t[0] + b), "fit_end": float(a * t[-1] + b), "mean": float(y.mean()),
        "r1": float(r1), "neff_k": [round(x, 2) if math.isfinite(x) else None for x in neff_k],
        "neff_bell": float(neff_bell), "neff_ar1": float(neff_ar1),
        "ci95_bell_per_yr": half(neff_bell), "ci95_ar1_per_yr": half(neff_ar1),
        "ci95_ols_per_yr": float(sig_u * stats.t.ppf(0.975, n - 2)),
    }


def decimal_year(dt64):
    """numpy datetime64 values -> decimal years."""
    dt64 = np.asarray(dt64).astype("datetime64[s]")
    y = dt64.astype("datetime64[Y]").astype(int) + 1970
    start = np.array([np.datetime64(f"{yy}-01-01T00:00:00") for yy in y])
    end = np.array([np.datetime64(f"{yy + 1}-01-01T00:00:00") for yy in y])
    return y + (dt64 - start).astype(float) / (end - start).astype(float)


def summary(r, scale=10.0, digits=3):
    """A compact, rounded record of a fit for the results file (slopes and half-widths per decade)."""
    span = r["t1"] - r["t0"]
    change = r["slope_per_yr"] * span
    return {
        "n": r["n"],
        "slope_per_decade": round(scale * r["slope_per_yr"], digits),
        "ci95_preprint_method": round(scale * r["ci95_bell_per_yr"], digits) if math.isfinite(r["ci95_bell_per_yr"]) else None,
        "ci95_ar1": round(scale * r["ci95_ar1_per_yr"], digits) if math.isfinite(r["ci95_ar1_per_yr"]) else None,
        "ci95_white_noise": round(scale * r["ci95_ols_per_yr"], digits),
        "neff_preprint_method": round(r["neff_bell"], 1), "neff_ar1": round(r["neff_ar1"], 1), "r1": round(r["r1"], 3),
        "fit_start": round(r["fit_start"], 3), "fit_end": round(r["fit_end"], 3),
        "years_spanned": round(span, 2), "change": round(change, 3),
        "change_pct_of_fit_start": round(100 * change / r["fit_start"], 1),
        "mean": round(r["mean"], 3),
    }