#!/usr/bin/env python3
"""Our re-analysis of Karakhanyan and Molodykh (2018), "Spatial distribution of temperature during
geomagnetic disturbances", Solar-Terrestrial Physics 4(4), 59-62 (https://doi.org/10.12737/stp-44201808;
Russian edition: Solnechno-Zemnaya Fizika 4(4), 76-81, https://doi.org/10.12737/szf-44201808).

The code behind section 4 and the other numbers marked "Our analysis" on the Fact-Checked Science page
https://factchecked.science/pc-index-temperature/

The paper builds an "index of efficiency of the solar activity effect",
    A = k * a * log|PC| * exp(-(lat - 60)^2 / 20^2),
where a is the anomaly of absolute humidity in the 925-700 hPa layer and PC the polar-cap geomagnetic index,
and correlates it with 925-700 hPa temperature anomalies at 16 areas during four disturbances (2002, 2005).
This script rebuilds the index from the same public data and asks whether its geomagnetic part adds
anything: humidity alone, the PC series of other disturbances, and weather from weeks away from the storms.

Run it in an empty folder:

    python3 pc_index_rerun.py
    python3 pc_index_rerun.py --cache <download folder> --out <results.json>

It needs Python 3 with NumPy and netCDF4. The first run downloads about 29 MB into ./pc_index_data:
12 subsets of NCEP/NCAR Reanalysis 1 daily means from NOAA PSL's THREDDS server (air temperature and
specific humidity at 925, 850 and 700 hPa, 20-90 N, 2002-01-01..06-30 and 2005-04-01..11-30), the
reanalysis's model terrain, and NASA's OMNI2 hourly files for 2002 and 2005. Later runs reuse them and take
under a minute. Results go to ./pc-index-numbers.json, with checksums of the data read; two summary lines are
printed.

Data (all public):
    NCEP/NCAR Reanalysis 1 (Kalnay et al. 1996, https://doi.org/10.1175/1520-0477(1996)077<0437:TNYRP>2.0.CO;2),
        "data provided by the NOAA PSL, Boulder, Colorado, USA, from their website at https://psl.noaa.gov"
        (the acknowledgement PSL asks for; free to use).
    OMNI2 hourly data set (Papitashvili and King 2020, https://doi.org/10.48322/1shr-ht18), NASA SPDF, CC0 1.0:
        word 52 is the PC(N) index, word 39 Kp, word 41 Dst (from WDC Kyoto).
PSL rebuilds each subset on request, so the netCDF files differ byte for byte between downloads (their
history attribute holds the time); the script records a SHA-256 of the data values instead, in
"data_checksums_sha256" in the results, which should match ours. The OMNI files are static.

Our choices where the paper is silent (marked * on the page):
    T = mean of the 925, 850 and 700 hPa temperatures*; a = water-vapour density from specific humidity,
    pressure and temperature at each level, averaged over the three*; anomalies from the mean of the 15 days
    before day 0 (days -15..-1)*; correlation over days -1..6, the span of the paper's Figs. 3-4*; each area
    the mean of the 3 x 3 grid points of its 5 x 5 degree box*; k = 1 (a positive constant can't change r);
    the latitude factor multiplies each area's box-mean humidity*, so it cancels from each area's correlation
    but not from curves averaged over areas at different latitudes; PC = daily mean of hourly PC(N)*;
    day 0 = 8 September 2005 (stated) and the day after the flare for the other three events*.
    Two forms of the geomagnetic factor: "printed", log10|PC| as the paper prints it, which turns negative
    when |PC| < 1; and "positive", log10(1 + |PC|), our guess at what the paper's Figs. 3-4 show (the plotted
    index keeps the humidity anomaly's sign on days when the plotted PC is below 1). A third form, with the
    daily mean of |hourly PC| inside either factor ("r_T_A_pc_abs_hourly"), is reported area by area as a
    sensitivity check.

Tests (rules fixed before the first run; added after it, without changing any rule: the "positive" variant,
the averaged curves and the hemispheric ranks; and, after a fresh review on 2026-10-28, the averaged curves
without the latitude weight, the share of picks at 0.9 or more and the counts for the third form):
    T0  r(T, A) at the paper's 16 areas, beside the paper's values.
    T1  r(T, a): humidity alone, no geomagnetic data. Also r(T, log PC) and r(T, PC).
    T2  wrong storm: the same weather and areas with the PC series of every other disturbed 8-day window of
        2002 and 2005 (daily PC > 1.5 on one of days 0-4, not overlapping the event's days -15..6).
    T3a no storm, fixed areas: the real PC series with weather from day 0 shifted 20-60 days either way
        (1-day steps, within the downloaded period, not overlapping any event's days -1..6).
    T3b no storm, areas chosen the paper's way: in each window, among 5 x 5 degree boxes centred at 45-80 N,
        the two warm-on-day-0 and the two cold-on-day-0 boxes with the largest temperature range over days
        -1..6 (at least 10 degrees of latitude or 20 of longitude apart), scored against the storm's index.
    T4  the hemisphere: area-weighted share of 20-90 N where r > 0.5, storm windows against shifted ones.
Extras: the events in OMNI's record, the days with daily PC > 1.5 in 2002 and 2005, the sample size the
paper's "+-" values imply, standard-atmosphere heights, the Clausius-Clapeyron rate, and the model terrain.

What this is and isn't: a rebuild, not the authors' code. The paper doesn't give k, the log base, how it
treats |PC| < 1, which PC version it used, the onset days of three events, or how it averaged the areas and
the layer, so our correlations are not the paper's: with the printed formula the median at the 16 areas is
0.36 against the paper's 0.90, with the positive factor 0.68. What holds in both forms: humidity alone
tracks temperature at least as closely as the index; the real storm's PC series does no better than about a
third of other storms' series (in the median area 31% of them match or beat it with the printed factor, 33%
with the positive one); and the areas chosen in the storm windows, and the hemispheric shares, look like those
of windows weeks away. With the positive factor the index also scores about as well with other storms' series
or with weather from weeks away; with the printed one it scores poorly with all of them.
"""
import argparse, csv, datetime, gzip, hashlib, io, json, math, os, sys, urllib.request

import numpy as np
import netCDF4 as nc

NCSS = ('https://psl.noaa.gov/thredds/ncss/grid/Datasets/ncep.reanalysis/Dailies/pressure/{var}.{year}.nc'
        '?var={var}&north=90&south=20&west=0&east=357.5&vertCoord={lev}'
        '&time_start={t0}T00:00:00Z&time_end={t1}T00:00:00Z&accept=netcdf4')
HGT = ('https://psl.noaa.gov/thredds/ncss/grid/Datasets/ncep.reanalysis/surface/hgt.sfc.nc'
       '?var=hgt&north=90&south=20&west=0&east=357.5&accept=netcdf4')
OMNI = 'https://spdf.gsfc.nasa.gov/pub/data/omni/low_res_omni/omni2_{year}.dat'
OMNI_SHA256 = {2002: 'ad0316866abee6f17741e8940d02f8199570ce8a107e2cb45c2eff01972a560a',
               2005: '9b99d08300d8e2a2cbeaeab073b2e0d83e31b5c925c0f465860430602ad07102'}
PERIODS = {2002: ('2002-01-01', '2002-06-30'), 2005: ('2005-04-01', '2005-11-30')}
LEVELS = (925, 850, 700)

# The paper's table (EN p. 60; RU p. 78): flare date, then (lat, lon east, +1 warm / -1 cold, the paper's r, its +-)
EVENTS = [
    ('2002-03-22', [(60, -130, +1, 0.90, 0.18), (65, -40, +1, 0.94, 0.14), (80, -135, -1, 0.93, 0.15), (50, 15, -1, 0.90, 0.18)]),
    ('2002-04-16', [(50, -160, +1, 0.68, 0.26), (65, 165, -1, 0.96, 0.10)]),
    ('2005-06-16', [(65, -100, +1, 0.91, 0.17), (50, 15, +1, 0.88, 0.18), (55, -120, -1, 0.64, 0.31), (65, 60, -1, 0.85, 0.20)]),
    ('2005-09-07', [(70, 180, +1, 0.90, 0.18), (80, 105, +1, 0.51, 0.35), (55, -50, +1, 0.78, 0.26),
                    (75, -90, -1, 0.97, 0.10), (65, 140, -1, 0.56, 0.34), (70, 45, -1, 0.93, 0.15)]),
]
# The R values in the text for September 2005's averaged curves (EN pp. 61-62)
TEXT_R = [('warm areas, index (Fig. 3)', 0.94, 0.14), ('warm areas, PC and aa', 0.80, 0.24),
          ('cold areas, index (Fig. 4)', 0.92, 0.16), ('cold areas, PC and aa', -0.68, 0.30)]
DAYS = list(range(-1, 7))
BASE = list(range(-15, 0))
UA = {'User-Agent': 'Mozilla/5.0 (factchecked.science re-run)'}


def fetch(url, path):
    if os.path.exists(path) and os.path.getsize(path) > 0:
        return
    print('downloading', url, file=sys.stderr)
    req = urllib.request.Request(url, headers=UA)
    with urllib.request.urlopen(req, timeout=300) as r:
        data = r.read()
    if data[:2] == b'\x1f\x8b':
        data = gzip.decompress(data)
    tmp = path + '.part'
    with open(tmp, 'wb') as f:
        f.write(data)
    os.replace(tmp, path)


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


def ncep_name(var, lev, year):
    t0, t1 = PERIODS[year]
    return f'{var}.{lev}hPa.{t0}_{t1}.N20-90.nc'


def load_year(cache, year, checks):
    T, R = [], []
    for lev in LEVELS:
        arrays = {}
        for var in ('air', 'shum'):
            path = os.path.join(cache, ncep_name(var, lev, year))
            fetch(NCSS.format(var=var, year=year, lev=lev, t0=PERIODS[year][0], t1=PERIODS[year][1]), path)
            d = nc.Dataset(path)
            x = np.asarray(d.variables[var][:, 0, :, :], dtype='f8')
            lat = np.asarray(d.variables['lat'][:], dtype='f8')
            lon = np.asarray(d.variables['lon'][:], dtype='f8')
            tt = d.variables['time']
            dates = [datetime.date(v.year, v.month, v.day) for v in nc.num2date(tt[:], tt.units)]
            checks[os.path.basename(path)] = hashlib.sha256(np.asarray(d.variables[var][:], dtype='<f4').tobytes()).hexdigest()
            d.close()
            arrays[var] = (x, lat, lon, dates)
        ta, lat, lon, dates = arrays['air']
        q = arrays['shum'][0]
        assert np.allclose(lat, arrays['shum'][1]) and np.allclose(lon, arrays['shum'][2]) and dates == arrays['shum'][3]
        p = lev * 100.0
        e = q * p / (0.622 + 0.378 * q)          # vapour pressure, Pa
        rho = e / (461.5 * ta) * 1000.0           # water-vapour density, g m-3
        T.append(ta - 273.15)
        R.append(rho)
    return dates, lat, lon, np.mean(T, axis=0), np.mean(R, axis=0)


def load_omni(cache, checks):
    pc, pcabs, kp, dst = {}, {}, {}, {}
    for year in (2002, 2005):
        path = os.path.join(cache, f'omni2_{year}.dat')
        fetch(OMNI.format(year=year), path)
        digest = sha256_file(path)
        checks[os.path.basename(path)] = digest
        if digest != OMNI_SHA256[year]:
            print(f'note: omni2_{year}.dat differs from the file we used (SHA-256 {digest})', file=sys.stderr)
        acc = {}
        with open(path, encoding='ascii') as f:
            for line in f:
                w = line.split()
                if len(w) < 55:
                    continue
                day = datetime.date(int(w[0]), 1, 1) + datetime.timedelta(days=int(w[1]) - 1)
                a = acc.setdefault(day, {'pc': [], 'kp': [], 'dst': []})
                if float(w[51]) < 999:
                    a['pc'].append(float(w[51]))
                if int(w[38]) < 99:
                    a['kp'].append(int(w[38]) / 10)
                if int(w[40]) < 99999:
                    a['dst'].append(int(w[40]))
        for day, a in acc.items():
            if a['pc']:
                pc[day] = sum(a['pc']) / len(a['pc'])
                pcabs[day] = sum(abs(v) for v in a['pc']) / len(a['pc'])
            if a['kp']:
                kp[day] = sum(a['kp']) / len(a['kp'])
            if a['dst']:
                dst[day] = min(a['dst'])
    return pc, pcabs, kp, dst


def latw(phi):
    return math.exp(-((phi - 60.0) ** 2) / 20.0 ** 2)


def corr(x, y):
    x = np.asarray(x, float); y = np.asarray(y, float)
    if np.std(x) == 0 or np.std(y) == 0:
        return float('nan')
    return float(np.corrcoef(x, y)[0, 1])


def factor(p, variant):
    p = np.abs(np.asarray(p, float))
    if variant == 'positive':
        return np.log10(1.0 + p)
    return np.log10(np.maximum(p, 1e-3))


def shift(d, k):
    return d + datetime.timedelta(days=k)


def overlaps(d0, others, span=(-1, 6), other_span=(-1, 6)):
    a0, a1 = shift(d0, span[0]), shift(d0, span[1])
    return any(a0 <= shift(o, other_span[1]) and shift(o, other_span[0]) <= a1 for o in others)


def main():
    ap = argparse.ArgumentParser(description=__doc__.split('\n\n')[0])
    ap.add_argument('--cache', default='pc_index_data', help='folder for the downloads (default ./pc_index_data)')
    ap.add_argument('--out', default='pc-index-numbers.json', help='results file (default ./pc-index-numbers.json)')
    args = ap.parse_args()
    os.makedirs(args.cache, exist_ok=True)
    checks = {}

    DATA = {y: load_year(args.cache, y, checks) for y in PERIODS}
    PC, PCABS, KP, DST = load_omni(args.cache, checks)
    path = os.path.join(args.cache, 'hgt.sfc.N20-90.nc')
    fetch(HGT, path)
    dh = nc.Dataset(path)
    H = np.asarray(dh.variables['hgt'][0], dtype='f8'); hlat = np.asarray(dh.variables['lat'][:]); hlon = np.asarray(dh.variables['lon'][:])
    dh.close()

    EV_D0 = {fl: shift(datetime.date.fromisoformat(fl), 1) for fl, _ in EVENTS}

    def pc_series(d0, table=PC):
        return np.array([table[shift(d0, k)] for k in DAYS])

    def window_ok(d0, y):
        dates = DATA[y][0]
        return shift(d0, -15) in dates and shift(d0, 6) in dates

    def box_idx(lat, lon, lat0, lon0):
        lon0 = lon0 % 360
        li = [i for i, v in enumerate(lat) if abs(v - lat0) <= 2.5 + 1e-6]
        lj = [j for j, v in enumerate(lon) if min(abs(v - lon0), 360 - abs(v - lon0)) <= 2.5 + 1e-6]
        return li, lj

    def box_series(y, lat0, lon0, d0):
        dates, lat, lon, T, R = DATA[y]
        li, lj = box_idx(lat, lon, lat0, lon0)
        i0 = dates.index(d0)
        tb = T[:, li][:, :, lj].mean(axis=(1, 2)); rb = R[:, li][:, :, lj].mean(axis=(1, 2))
        b = [i0 + k for k in BASE]; w = [i0 + k for k in DAYS]
        return tb[w] - tb[b].mean(), rb[w] - rb[b].mean()

    def box_weighted_humidity(y, lat0, lon0, d0):
        """The humidity anomaly weighted by latitude at each grid point before the box mean (a check on whether
        the weighting cancels from an area's correlation when it varies across the box)."""
        dates, lat, lon, T, R = DATA[y]
        li, lj = box_idx(lat, lon, lat0, lon0)
        i0 = dates.index(d0)
        sub = R[:, li][:, :, lj]
        b = [i0 + k for k in BASE]; w = [i0 + k for k in DAYS]
        an = sub[w] - sub[b].mean(axis=0)
        wts = np.array([latw(lat[i]) for i in li])[None, :, None]
        return (an * wts).mean(axis=(1, 2))

    def field_anoms(y, d0):
        dates, lat, lon, T, R = DATA[y]
        i0 = dates.index(d0)
        b = [i0 + k for k in BASE]; w = [i0 + k for k in DAYS]
        return T[w] - T[b].mean(axis=0), R[w] - R[b].mean(axis=0), lat, lon

    def box_mean_field(x):
        xp = np.concatenate([x[..., -1:], x, x[..., :1]], axis=-1)
        s = np.zeros_like(x)
        for di in (-1, 0, 1):
            for dj in (-1, 0, 1):
                s += np.roll(np.roll(xp, -dj, axis=-1)[..., 1:-1], -di, axis=-2)
        s /= 9.0
        s[..., 0, :] = np.nan; s[..., -1, :] = np.nan
        return s

    def corr_map(a, b):
        am = a - a.mean(axis=0); bm = b - b.mean(axis=0)
        with np.errstate(invalid='ignore', divide='ignore'):
            return (am * bm).sum(axis=0) / np.sqrt((am ** 2).sum(axis=0) * (bm ** 2).sum(axis=0))

    def select_and_score(y, w0, lf):
        ta, ra, lat, lon = field_anoms(y, w0)
        tb = box_mean_field(ta); rb = box_mean_field(ra)
        A = rb * lf[:, None, None] * np.array([latw(v) for v in lat])[None, :, None]
        rr = corr_map(tb, A); rng = tb.max(axis=0) - tb.min(axis=0)
        picks = []
        for sign in (+1, -1):
            cand = sorted(((rng[i, j], i, j) for i, la in enumerate(lat) if 45 <= la <= 80
                           for j in range(len(lon)) if not np.isnan(tb[1, i, j]) and sign * tb[1, i, j] > 0), reverse=True)
            chosen = []
            for _, i, j in cand:
                if all(not (abs(lat[i] - lat[ii]) < 10 and min(abs(lon[j] - lon[jj]), 360 - abs(lon[j] - lon[jj])) < 20)
                       for ii, jj in chosen):
                    chosen.append((i, j))
                if len(chosen) == 2:
                    break
            for i, j in chosen:
                picks.append({'type': 'warm' if sign > 0 else 'cold', 'lat': float(lat[i]), 'lon': float(lon[j]),
                              'r': float(rr[i, j])})
        return picks

    def hemis(y, w0, lf):
        ta, ra, lat, lon = field_anoms(y, w0)
        A = ra * lf[:, None, None] * np.array([latw(v) for v in lat])[None, :, None]
        wts = np.cos(np.radians(lat))[:, None] * np.ones((1, len(lon)))
        def share(m):
            ok = ~np.isnan(m)
            return float((wts * (m > 0.5) * ok).sum() / (wts * ok).sum())
        return share(corr_map(ta, A)), share(corr_map(ta, ra))

    # placebo PC windows (T2): any 8-day window of 2002 and 2005 with daily PC > 1.5 on one of days 0..4
    pc_windows = []
    for d0 in sorted(PC):
        try:
            s = pc_series(d0)
        except KeyError:
            continue
        if max(s[1:6]) > 1.5:
            pc_windows.append(d0)

    shifts_for = {}
    for fl, _ in EVENTS:
        d0 = EV_D0[fl]; y = d0.year
        shifts_for[fl] = [shift(d0, s) for s in list(range(-60, -19)) + list(range(20, 61))
                          if window_ok(shift(d0, s), y) and not overlaps(shift(d0, s), list(EV_D0.values()))]

    res = {'variants': {}}
    for variant in ('printed', 'positive'):
        V = {'areas': [], 'T3b': {}, 'T4': [], 'sept2005_averaged': {}}
        all_real_picks, all_plc_picks = [], []
        for fl, rows in EVENTS:
            d0 = EV_D0[fl]; y = d0.year
            pcs = pc_series(d0); lf = factor(pcs, variant)
            plc = [w for w in pc_windows if not overlaps(w, [d0], span=(-1, 6), other_span=(-15, 6))]
            for la, lo, sg, rp, pm in rows:
                ta, ra = box_series(y, la, lo, d0)
                rA = corr(ta, ra * lf * latw(la))
                p2 = [v for v in (corr(ta, ra * factor(pc_series(w), variant) * latw(la)) for w in plc) if not math.isnan(v)]
                p3 = []
                for w0 in shifts_for[fl]:
                    tb, rb = box_series(y, la, lo, w0)
                    v = corr(tb, rb * lf * latw(la))
                    if not math.isnan(v):
                        p3.append(v)
                V['areas'].append({
                    'flare': fl, 'day0': d0.isoformat(), 'lat': la, 'lon': lo, 'type': 'warm' if sg > 0 else 'cold',
                    'paper_r': rp, 'r_T_A': rA,
                    'r_T_A_gridpoint_weight': corr(ta, box_weighted_humidity(y, la, lo, d0) * lf),
                    'r_T_A_pc_abs_hourly': corr(ta, ra * factor(pc_series(d0, PCABS), variant) * latw(la)),
                    'r_T_a': corr(ta, ra), 'r_T_PC': corr(ta, pcs), 'r_T_logPC': corr(ta, lf),
                    'T_day0': float(ta[1]),
                    'T2_n': len(p2), 'T2_median': float(np.median(p2)), 'T2_share_ge_real': float(np.mean([v >= rA for v in p2])),
                    'T2_share_ge_0.5': float(np.mean([v >= 0.5 for v in p2])),
                    'T3a_n': len(p3), 'T3a_median': float(np.median(p3)), 'T3a_share_ge_real': float(np.mean([v >= rA for v in p3])),
                    'T_anom': [round(float(v), 3) for v in ta], 'a_anom': [round(float(v), 4) for v in ra]})
            real = select_and_score(y, d0, lf)
            for p in real:
                p['flare'] = fl
            all_real_picks += real
            shares = []
            for w0 in shifts_for[fl]:
                for p in select_and_score(y, w0, lf):
                    if not math.isnan(p['r']):
                        all_plc_picks.append(p['r'])
                shares.append(hemis(y, w0, lf))
            sA, sa = hemis(y, d0, lf)
            V['T4'].append({'flare': fl, 'n_windows': len(shares), 'share_r_T_A_gt_0.5': sA, 'share_r_T_a_gt_0.5': sa,
                            'placebo_median_share_r_T_A': float(np.median([v[0] for v in shares])),
                            'placebo_median_share_r_T_a': float(np.median([v[1] for v in shares])),
                            'placebo_share_of_windows_ge_real': float(np.mean([v[0] >= sA for v in shares]))})
        rr = [p['r'] for p in all_real_picks]
        V['T3b'] = {'real_picks': all_real_picks, 'real_median': float(np.median(rr)), 'placebo_n': len(all_plc_picks),
                    'placebo_median': float(np.median(all_plc_picks)),
                    'placebo_share_ge_0.5': float(np.mean([v >= 0.5 for v in all_plc_picks])),
                    'placebo_share_ge_0.8': float(np.mean([v >= 0.8 for v in all_plc_picks])),
                    'placebo_share_ge_0.9': float(np.mean([v >= 0.9 for v in all_plc_picks])),
                    'placebo_quantiles_10_25_50_75_90': [float(np.quantile(all_plc_picks, q)) for q in (0.1, 0.25, 0.5, 0.75, 0.9)],
                    'placebo_histogram_bins_-1_to_1_step_0.1': [int(v) for v in np.histogram(all_plc_picks, bins=np.linspace(-1, 1, 21))[0]]}
        d9 = EV_D0['2005-09-07']; pc9 = pc_series(d9); l9 = factor(pc9, variant)
        for typ in ('warm', 'cold'):
            ar = [a for a in V['areas'] if a['flare'] == '2005-09-07' and a['type'] == typ]
            Tm = np.mean([a['T_anom'] for a in ar], axis=0)
            Am = np.mean([np.array(a['a_anom']) * l9 * latw(a['lat']) for a in ar], axis=0)
            Am_flat = np.mean([np.array(a['a_anom']) * l9 for a in ar], axis=0)   # the same, without the latitude weight
            V['sept2005_averaged'][typ] = {'r_T_A': corr(Tm, Am), 'r_T_A_no_latitude_weight': corr(Tm, Am_flat),
                                           'r_T_a': corr(Tm, np.mean([a['a_anom'] for a in ar], axis=0)),
                                           'r_T_PC': corr(Tm, pc9),
                                           'latitude_weights': [round(latw(a['lat']), 3) for a in ar]}
        A16 = V['areas']
        med = lambda k: float(np.median([a[k] for a in A16 if not math.isnan(a[k])]))
        V['summary'] = {k: med(k) for k in ('paper_r', 'r_T_A', 'r_T_A_pc_abs_hourly', 'r_T_a', 'r_T_PC', 'T2_median',
                                             'T2_share_ge_real', 'T3a_median', 'T3a_share_ge_real')}
        V['summary'].update({
            'n_humidity_ge_index': sum(a['r_T_a'] >= a['r_T_A'] for a in A16),
            'n_humidity_ge_index_pc_abs_hourly': sum(a['r_T_a'] >= a['r_T_A_pc_abs_hourly'] for a in A16),
            'max_abs_change_with_gridpoint_weight': max(abs(a['r_T_A_gridpoint_weight'] - a['r_T_A']) for a in A16),
            'n_humidity_gt_0.7': sum(a['r_T_a'] > 0.7 for a in A16),
            'n_index_gt_0.7': sum(a['r_T_A'] > 0.7 for a in A16),
            'n_T2_half_or_more_ge_0.5': sum(a['T2_share_ge_0.5'] >= 0.5 for a in A16),
            'warm_r_T_PC_positive': sum(a['r_T_PC'] > 0 for a in A16 if a['type'] == 'warm'),
            'cold_r_T_PC_negative': sum(a['r_T_PC'] < 0 for a in A16 if a['type'] == 'cold'),
            'n_day0_sign_as_paper': sum((a['T_day0'] > 0) == (a['type'] == 'warm') for a in A16)})
        res['variants'][variant] = V

    # extras ---------------------------------------------------------------------------------------------
    ev = []
    for fl, _ in EVENTS:
        d0 = EV_D0[fl]
        days = [shift(d0, k) for k in DAYS]; base = [shift(d0, k) for k in BASE]
        ev.append({'flare': fl, 'day0': d0.isoformat(), 'pc_days_-1_6': [round(PC[d], 3) for d in days],
                   'max_daily_pc': max(PC[d] for d in days), 'days_pc_gt_1.5': sum(PC[d] > 1.5 for d in days),
                   'days_abs_pc_lt_1': sum(abs(PC[d]) < 1 for d in days), 'min_dst': min(DST[d] for d in days),
                   'max_daily_kp': max(KP[d] for d in days), 'baseline_min_dst': min(DST[d] for d in base),
                   'baseline_min_dst_date': min(base, key=lambda d: DST[d]).isoformat()})
    res['events'] = ev
    cand = {}
    for year in (2002, 2005):
        days = sorted(d for d in PC if d.year == year and PC[d] > 1.5)
        eps = []
        for d in days:
            if eps and (d - eps[-1][-1]).days <= 3:
                eps[-1].append(d)
            else:
                eps.append([d])
        cand[str(year)] = {'days_pc_gt_1.5': len(days), 'episodes_gaps_up_to_3_days_merged': len(eps)}
    res['candidate_days'] = cand

    def p_two_sided(r, n):
        from math import lgamma, exp, log, sqrt
        df = n - 2; t = abs(r) * sqrt(df / (1 - r * r)); x = df / (df + t * t); a, b = df / 2.0, 0.5
        def cf(a, b, x):
            qab, qap, qam = a + b, a + 1, a - 1
            c, d = 1.0, 1 - qab * x / qap
            d = 1 / (d if abs(d) > 1e-300 else 1e-300); h = d
            for m in range(1, 200):
                m2 = 2 * m
                aa = m * (b - m) * x / ((qam + m2) * (a + m2)); d = 1 + aa * d; d = 1 / (d if abs(d) > 1e-300 else 1e-300)
                c = 1 + aa / c; c = c if abs(c) > 1e-300 else 1e-300; h *= d * c
                aa = -(a + m) * (qab + m) * x / ((a + m2) * (qap + m2)); d = 1 + aa * d; d = 1 / (d if abs(d) > 1e-300 else 1e-300)
                c = 1 + aa / c; c = c if abs(c) > 1e-300 else 1e-300; de = d * c; h *= de
                if abs(de - 1) < 3e-14:
                    break
            return h
        bt = exp(lgamma(a + b) - lgamma(a) - lgamma(b) + a * log(x) + b * log(1 - x))
        return bt * cf(a, b, x) / a if x < (a + 1) / (a + b + 2) else 1 - bt * cf(b, a, 1 - x) / b

    table = [(r, pm) for _, rows in EVENTS for (_, _, _, r, pm) in rows] + [(r, pm) for _, r, pm in TEXT_R]
    lo, hi = 0.0, 0.999
    for _ in range(60):
        mid = (lo + hi) / 2
        lo, hi = (mid, hi) if p_two_sided(mid, 8) > 0.05 else (lo, mid)
    res['plus_minus'] = {'implied_n': [round(2 + (1 - r * r) / (pm * pm), 2) for r, pm in table],
                         'nominal_p_n8_table': [round(p_two_sided(r, 8), 4) for r, _ in table[:16]],
                         'critical_r_p05_n8': round(hi, 3),
                         'table_values_above_critical': sum(abs(r) > hi for r, _ in table[:16])}
    T0, L, p0, g0, Rd = 288.15, 0.0065, 1013.25, 9.80665, 287.053
    res['standard_atmosphere_m'] = {str(p): round(T0 / L * (1 - (p / p0) ** (Rd * L / g0))) for p in LEVELS}
    res['clausius_clapeyron_percent_per_K'] = {str(tc): round(100 * 2.501e6 / (461.5 * (273.15 + tc) ** 2), 1) for tc in (0, 10)}
    terr = []
    for fl, rows in EVENTS:
        for la, lo, *_ in rows:
            li, lj = box_idx(hlat, hlon, la, lo)
            sub = H[np.ix_(li, lj)]
            terr.append({'lat': la, 'lon': lo, 'box_mean_m': round(float(sub.mean())), 'box_max_m': round(float(sub.max()))})
    res['terrain'] = terr
    res['data_checksums_sha256'] = checks

    with open(args.out, 'w', encoding='utf-8', newline='\n') as f:
        json.dump(res, f, indent=1, ensure_ascii=False)
        f.write('\n')
    for variant in ('printed', 'positive'):
        s = res['variants'][variant]['summary']; t = res['variants'][variant]['T3b']
        print(f'{variant:9s} median r at the 16 areas: index {s["r_T_A"]:.2f}, humidity alone {s["r_T_a"]:.2f}, '
              f'wrong storms {s["T2_median"]:.2f} (the paper: {s["paper_r"]:.2f}); areas chosen the paper\'s way '
              f'away from the storms: median {t["placebo_median"]:.2f}, r >= 0.8 in {100 * t["placebo_share_ge_0.8"]:.0f}% '
              f'of {t["placebo_n"]}; storm windows: median {t["real_median"]:.2f}')
    print('results written to', args.out)


if __name__ == '__main__':
    main()