"""How much does the polar field k years AFTER a cycle's maximum say about the NEXT cycle's amplitude?
Polar field: MWO polar faculae calibrated to WSO (annual N in Aug-Sep, S in Feb-Mar; Munoz-Jaramillo et al. 2012) before 1976,
WSO 20 nHz-filtered Nf, Sf (10-day, ../data/wso_polar.csv) after. Strength = (|N| + |S|)/2 in the new polarity
(WSO: signed with the polarity expected at the next minimum, clipped at 0; MWO counts are unsigned).
For each lag k, fit R_{n+1} = a + b*P_n(t_max_n + k) over cycles 14..24 (targets 15..25), report LOO RMSE and the C26 value
where data exist (WSO valid to 2026-01)."""
import csv, os, numpy as np
HERE = os.path.dirname(os.path.abspath(__file__)); D = os.path.join(HERE, '..', 'data')
cyc = {int(r['cycle']): r for r in csv.DictReader(open(os.path.join(D, 'cycles.csv')))}
tmax = {n: float(cyc[n]['max_decyear']) for n in cyc}; amp = {n: float(cyc[n]['max_sn']) for n in cyc}
fac = np.array([[float(x) for x in l.strip().split(',')] for l in open(os.path.join(D, 'mwo_faculae', 'e_PField_MWO_WSO.csv')) if l[0].isdigit()])
wso = []
for r in csv.DictReader(open(os.path.join(D, 'wso_polar.csv'))):
    if r['Nf'] != 'XXX' and r['Sf'] != 'XXX':
        y, m, d = map(int, r['date'].split('-')); wso.append((y + (m - 1) / 12 + (d - 1) / 365.25, float(r['Nf']) / 100, float(r['Sf']) / 100))
wso = np.array(wso)
# sign of the north polar field at the minimum that STARTS cycle n+1: +1 at 1976 (start C21), alternating
def north_sign_at_start(c): return 1 if (c - 21) % 2 == 0 else -1
def polar(n, t):
    """polar strength at time t for the field that will seed cycle n+1"""
    if t >= 1976.5:
        if t > wso[-1, 0] + 0.05: return np.nan
        i = np.argmin(abs(wso[:, 0] - t)); s = north_sign_at_start(n + 1)
        return max(0.0, (s * wso[i, 1] - s * wso[i, 2]) / 2)
    okn = ~np.isnan(fac[:, 1]); oks = ~np.isnan(fac[:, 7])
    if not okn.any() or t > 2007.5: return np.nan
    iN = np.argmin(abs(fac[okn, 0] - t)); iS = np.argmin(abs(fac[oks, 6] - t))
    if abs(fac[okn, 0][iN] - t) > 0.75 or abs(fac[oks, 6][iS] - t) > 0.75: return np.nan
    return (abs(fac[okn, 1][iN]) + abs(fac[oks, 7][iS])) / 2
print(' k   n   r     b     LOO-RMSE  (climatology sd 58)   C26 x   C26 pred')
out = []
for k in [0.5, 1.0, 1.25, 1.5, 2, 3, 4, 5, 6]:
    X, Y = [], []
    for n in range(14, 25):
        x = polar(n, tmax[n] + k)
        if not np.isnan(x): X.append(x); Y.append(amp[n + 1])
    X, Y = np.array(X), np.array(Y)
    if len(X) < 5: continue
    res = []
    for i in range(len(X)):
        m = np.arange(len(X)) != i; b, a = np.polyfit(X[m], Y[m], 1); res.append(Y[i] - (a + b * X[i]))
    b, a = np.polyfit(X, Y, 1); r = np.corrcoef(X, Y)[0, 1]
    x26 = polar(25, tmax[25] + k)
    p26 = a + b * x26 if not np.isnan(x26) else np.nan
    rm = np.sqrt(np.mean(np.square(res)))
    out.append(dict(k=k, n=len(X), r=round(r, 2), loo_rmse=round(rm, 1), x26=round(x26, 3) if not np.isnan(x26) else '', pred26=round(p26, 1) if not np.isnan(p26) else ''))
    print(f'{k:4.2f} {len(X):3d} {r:5.2f} {b:6.1f} {rm:8.1f}   {"":20s} {x26:7.3f} {p26:8.1f}')
with open(os.path.join(HERE, 'early_polar.csv'), 'w', newline='') as f:
    w = csv.DictWriter(f, fieldnames=list(out[0])); w.writeheader(); w.writerows(out)
print('WSO last valid sample:', wso[-1])