"""The RAPID record of the Atlantic overturning at 26.5°N, as used on factchecked.science/amoc-heat-loss/.

Usage:
  python3 -I rapid_trends.py MOC_V2024_1A MOC_V2024_2 [results.json]

  MOC_V2024_1A  moc_transports_200404_2024327.nc from RAPID release v2024.1a, the one Bell et al. (2026) use:
                BODC, doi:10.5285/48d0bf43-0598-ceb2-e063-7086abc062f1 (12-hourly overturning, 2 April 2004 to
                27 March 2024)
  MOC_V2024_2   RAPID_transports_v2024-2.nc from the newest release, v2024-2: BODC,
                doi:10.5285/5c39bf4e-5140-5546-e063-7086abc027bc (monthly, April 2004 to November 2024)

Both come inside the zip that BODC's "Download" button prepares on each DOI's page; take the two files out of the
zips. The data are under the UK Open Government Licence v3.0. RAPID's acknowledgement: "Data from the RAPID AMOC
observing project is funded by the Natural Environment Research Council, U.S. National Science Foundation (NSF)
with support from NOAA. They are freely available from https://rapid.ac.uk/."

What it prints (and saves, if a .json name is given):
- v2024.1a: the trend of calendar-year means 2005-2023 with the preprint's uncertainty recipe, which reproduces the
  "-0.89 +/- 1.5 Sv per decade" of its Fig. A1, both from all 12-hourly values and from monthly means;
- v2024.1a and v2024-2: straight-line trends of the monthly means over several periods, with the change over each
  period in Sv and as a share of the fitted start (New Scientist's figure is about 10% since 2004);
- v2024-2: the year means, the 2024 mean (January to November) and the trend of the heat transport.

Needs Python 3 with NumPy, SciPy and h5py in the interpreter's own site-packages (python3 -I ignores packages
installed with pip install --user). Reads the files as data only. Statistics: trendstats.py (same folder).
"""
import hashlib
import json
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import h5py  # noqa: E402
import numpy as np  # noqa: E402

from trendstats import decimal_year, fit, summary  # noqa: E402

sys.stdout.reconfigure(encoding="utf-8")
if len(sys.argv) < 3:
    sys.exit(__doc__)
path_a, path_b = sys.argv[1], sys.argv[2]
out_json = sys.argv[3] if len(sys.argv) > 3 else None


def sha256(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 22), b""):
            h.update(chunk)
    return h.hexdigest()


def text(v):
    return v.decode() if isinstance(v, bytes) else str(v)


res = {"inputs": {}, "v2024_1a": {}, "v2024_2": {}}
for p in (path_a, path_b):
    res["inputs"][os.path.basename(p)] = {"bytes": os.path.getsize(p), "sha256": sha256(p)}
    print(f"{os.path.basename(p)}: {os.path.getsize(p):,} bytes, SHA-256 {res['inputs'][os.path.basename(p)]['sha256']}")

# ---- v2024.1a: 12-hourly overturning (moc_mar_hc10), days since 2004-04-01 --------------------------------------
with h5py.File(path_a, "r") as f:
    days = f["time"][:]
    amoc = f["moc_mar_hc10"][:].astype(float)
    version_a = text(f.attrs.get("version", ""))
amoc[amoc <= -9999] = np.nan
when = np.datetime64("2004-04-01T00:00:00") + np.round(days * 86400).astype("timedelta64[s]")
year = when.astype("datetime64[Y]").astype(int) + 1970
month = when.astype("datetime64[M]")
print(f"\nRAPID {version_a}: {len(amoc)} values, {str(when[0])[:10]} to {str(when[-1])[:10]},"
      f" {int(np.isnan(amoc).sum())} missing")
res["v2024_1a"]["version"] = version_a

years = np.arange(2005, 2024)
ann_12h = np.array([np.nanmean(amoc[year == y]) for y in years])
months = np.unique(month)
mon_mean = np.array([np.nanmean(amoc[month == m]) for m in months])
mon_year = months.astype("datetime64[Y]").astype(int) + 1970
ann_mon = np.array([mon_mean[mon_year == y].mean() for y in years])

for label, vals in (("annual_2005_2023_from_12h", ann_12h), ("annual_2005_2023_from_monthly", ann_mon)):
    r = fit(years + 0.5, vals, kmax=9)
    res["v2024_1a"][label] = summary(r)
    s = res["v2024_1a"][label]
    print(f"  {label}: {s['slope_per_decade']:+.3f} Sv/decade, 95% +/-{s['ci95_preprint_method']:.2f} by the preprint's"
          f" recipe (N_eff {s['neff_preprint_method']}), +/-{s['ci95_ar1']:.2f} AR(1)")
res["v2024_1a"]["annual_means_2005_2023_from_12h"] = {str(y): round(float(v), 3) for y, v in zip(years, ann_12h)}
res["v2024_1a"]["mean_2005_2023"] = round(float(ann_12h.mean()), 2)

mon_t = decimal_year(months.astype("datetime64[s]") + np.timedelta64(15, "D"))
for label, sel in (("monthly_2004-04_2023-12", months <= np.datetime64("2023-12")),
                   ("monthly_2004-04_2024-02", months <= np.datetime64("2024-02")),
                   ("monthly_2004-04_2023-02", months <= np.datetime64("2023-02"))):
    r = fit(mon_t[sel], mon_mean[sel], kmax=24)
    res["v2024_1a"][label] = summary(r)
    s = res["v2024_1a"][label]
    print(f"  {label}: {s['slope_per_decade']:+.3f} Sv/decade, 95% +/-{s['ci95_ar1']:.2f} AR(1),"
          f" +/-{s['ci95_preprint_method']:.2f} preprint's recipe; change {s['change']:+.2f} Sv,"
          f" {s['change_pct_of_fit_start']:+.1f}% of the fitted start")

# ---- v2024-2: monthly overturning in depth space and heat transport, days since 1950-01-01 ------------------------
with h5py.File(path_b, "r") as f:
    days = f["time_1m"][:]
    amoc2 = f["amoc_depth_1m"][:].astype(float)
    mht = f["heat_trans_1m"][:].astype(float)
    version_b = text(f.attrs.get("Dataset_version", ""))
when2 = np.datetime64("1950-01-01T00:00:00") + np.round(days * 86400).astype("timedelta64[s]")
months2 = when2.astype("datetime64[M]")
year2 = when2.astype("datetime64[Y]").astype(int) + 1970
t2 = decimal_year(when2)
print(f"\nRAPID {version_b}: {len(amoc2)} months, {str(months2[0])} to {str(months2[-1])}")
res["v2024_2"]["version"] = version_b
for label, sel in (("monthly_2004-04_2023-12", months2 <= np.datetime64("2023-12")),
                   ("monthly_2004-04_2024-11", months2 <= np.datetime64("2024-11"))):
    r = fit(t2[sel], amoc2[sel], kmax=24)
    res["v2024_2"][label] = summary(r)
    s = res["v2024_2"][label]
    print(f"  {label}: {s['slope_per_decade']:+.3f} Sv/decade, 95% +/-{s['ci95_ar1']:.2f} AR(1),"
          f" +/-{s['ci95_preprint_method']:.2f} preprint's recipe; change {s['change']:+.2f} Sv,"
          f" {s['change_pct_of_fit_start']:+.1f}% of the fitted start")
ann2 = {int(y): float(amoc2[year2 == y].mean()) for y in range(2005, 2025)}
res["v2024_2"]["annual_means_2005_2024"] = {str(y): round(v, 3) for y, v in ann2.items()}
res["v2024_2"]["months_in_2024"] = int((year2 == 2024).sum())
res["v2024_2"]["mean_2005_2023"] = round(float(np.mean([ann2[y] for y in range(2005, 2024)])), 2)
print(f"  mean 2005-2023 {res['v2024_2']['mean_2005_2023']} Sv; 2024 ({res['v2024_2']['months_in_2024']} months)"
      f" {ann2[2024]:.2f} Sv")
r = fit(t2, mht, kmax=24)
res["v2024_2"]["heat_transport_monthly_2004-04_2024-11"] = summary(r, digits=4)
s = res["v2024_2"]["heat_transport_monthly_2004-04_2024-11"]
print(f"  heat transport: mean {s['mean']:.3f} PW; {s['slope_per_decade']:+.4f} PW/decade,"
      f" 95% +/-{s['ci95_ar1']:.3f} AR(1)")

if out_json:
    with open(out_json, "w", encoding="utf-8", newline="\n") as fh:
        json.dump(res, fh, indent=1, ensure_ascii=False)
        fh.write("\n")
    print("\nwrote", out_json)