"""Stage-2 expansion predictors (item 2): subseasonal teleconnection + persistence features.

These extend the Stage-1 forecast-only feature block with predictors that carry real subseasonal
signal, all LEAKAGE-SAFE by construction:

  init-date teleconnection scalars (broadcast across cells of an init):
    enso_nino34   observed NINO3.4 SST anomaly for the init month (../HW nino34.txt). ENSO state
                  at forecast time; an observed initial condition, not a target-year forecast.
    mjo_pc1,      OMI MJO principal components and amplitude at the init date (../HW omi.txt).
    mjo_pc2,      PCs are used directly (not a phase angle) to avoid any convention ambiguity.
    mjo_amp

  per (init, cell) persistence, from observations available AT init and the fold-safe threshold:
    f_ante_ehdfreq  fraction of the ANTECEDENT_DAYS before init that were observed EHDs
                    (obs Tx >= min(Tx90p,35)); antecedent heat persistence.
    f_ante_sevmean  mean observed severity (obs Tx - threshold) over that antecedent window.

Both persistence features use the SAME fold-fit threshold `thr_by_doy` that Stage-1 uses for
labels, so they inherit the leakage rule (the threshold never sees the held-out year) and use only
observations dated strictly before the init (a legitimate initial condition). The MJO/ENSO indices
are observed states at init; the reforecast is control-only, so ensemble-spread predictors are
unavailable in training (documented — see STAGE2_PLAN §3 / MEMORY).

The design contract: this module returns feature columns in the SAME (init, cell) row order that
`features.assemble_year` produces (init slowest, cell fastest), so callers can hstack them onto X.
"""
from __future__ import annotations
import os
import numpy as np
import pandas as pd

EXT_FEATURE_NAMES = ["enso_nino34", "mjo_pc1", "mjo_pc2", "mjo_amp",
                     "f_ante_ehdfreq", "f_ante_sevmean"]
ANTECEDENT_DAYS = 15

# gridded candidate predictors (block-aggregated by scripts/20a_build_ext_grid.py)
GRID_LOCAL_NAMES = ["sm20", "sm100", "st100", "skt", "cape", "tcw", "tcc", "d2m"]
GRID_LS_NAMES = ["z500_reg", "z500_sub", "z500_grad", "t850_reg"]
GRID_FEATURE_NAMES = GRID_LOCAL_NAMES + GRID_LS_NAMES


# ----------------------------------------------------------------- gridded predictor cache
def load_extgrid(outputs_dir, year):
    """Load the per-year gridded candidate-predictor cache written by stage 20a."""
    return np.load(os.path.join(outputs_dir, "ext_predictors", f"extgrid_{year}.npz"),
                   allow_pickle=True)


def ext_grid_block(extgrid, block_index, ncell):
    """Gridded ext columns for one block, in assemble_year row order (init slow, cell fast).

    Local surface fields are per (init, cell); large-scale indices are per (init) and broadcast
    across the block's cells. Returns X_grid[ninit*ncell, n_local+n_ls]; column order =
    GRID_LOCAL_NAMES + GRID_LS_NAMES.
    """
    local = np.asarray(extgrid["local"])[:, block_index, :, :]     # (ninit, ncell, nloc)
    ninit, nc, nloc = local.shape
    assert nc == ncell, f"extgrid ncell {nc} != {ncell}"
    Xloc = local.reshape(ninit * ncell, nloc)
    large = np.asarray(extgrid["large"])[:, block_index, :]        # (ninit, nls)
    Xls = np.repeat(large, ncell, axis=0)
    return np.column_stack([Xloc, Xls]).astype(np.float32)


# ----------------------------------------------------------------- teleconnection index parsers
def load_nino34(path):
    """Return dict (year, month) -> NINO3.4 anomaly. File cols: YR MON N12 A N3 A N4 A N34 A."""
    out = {}
    with open(path) as f:
        for line in f:
            p = line.split()
            if len(p) < 10 or not p[0].isdigit():
                continue
            out[(int(p[0]), int(p[1]))] = float(p[9])   # NINO3.4 ANOM
    return out


def load_omi(path):
    """Return dict (y,m,d) -> (pc1, pc2, amp). File cols: year month day PC1 PC2 amp."""
    out = {}
    with open(path) as f:
        for line in f:
            p = line.split()
            if len(p) < 6 or not p[0].lstrip("-").isdigit():
                continue
            out[(int(p[0]), int(p[1]), int(p[2]))] = (float(p[3]), float(p[4]), float(p[5]))
    return out


def _nearest_omi(omi, y, m, d):
    """OMI at (y,m,d), searching +/- 3 days for gaps."""
    from datetime import date, timedelta
    d0 = date(y, m, d)
    for off in [0, -1, 1, -2, 2, -3, 3]:
        dd = d0 + timedelta(days=off)
        key = (dd.year, dd.month, dd.day)
        if key in omi:
            return omi[key]
    return (0.0, 0.0, 0.0)


# ----------------------------------------------------------------- antecedent persistence
def _ymd(dt):
    return dt.year * 10000 + dt.month * 100 + dt.day


def antecedent_persistence(fc_year, obs_vals, obs_lut, thr_by_doy, land_iy, land_ix,
                           ndays=ANTECEDENT_DAYS):
    """Per (init, cell): antecedent observed EHD frequency and mean severity over the `ndays`
    calendar days before each init. Returns (ehdfreq[ninit,ncell], sevmean[ninit,ncell])."""
    from .features import _doy_noleap
    inits = pd.to_datetime(fc_year["init"].values)
    ninit = len(inits); ncell = land_iy.size
    thr = thr_by_doy.values                                   # (365, lat, lon)
    ehdfreq = np.full((ninit, ncell), np.nan)
    sevmean = np.full((ninit, ncell), np.nan)
    for it, t0 in enumerate(inits):
        days = pd.date_range(end=t0 - pd.Timedelta(days=1), periods=ndays, freq="D")
        doy = _doy_noleap(days)
        rows = np.array([obs_lut.get(_ymd(d), -1) for d in days])
        ok = rows >= 0
        if not ok.any():
            continue
        o = obs_vals[rows[ok]][:, land_iy, land_ix]           # (nok, cell)
        th = thr[doy[ok] - 1][:, land_iy, land_ix]            # (nok, cell)
        with np.errstate(all="ignore"):
            ehd = (o >= th).astype(float)
            sev = o - th
        ehdfreq[it] = np.nanmean(ehd, axis=0)
        sevmean[it] = np.nanmean(sev, axis=0)
    return ehdfreq, sevmean


# ----------------------------------------------------------------- assemble the ext block
def ext_features_year(fc_year, obs_vals, obs_lut, thr_by_doy, land_iy, land_ix,
                      nino34, omi, ndays=ANTECEDENT_DAYS):
    """Return dict with the ext feature columns for ALL (init, cell) rows in assemble_year order,
    plus `names`. Row order: init slowest, cell fastest (np.repeat(init)/np.tile(cell))."""
    inits = pd.to_datetime(fc_year["init"].values)
    ninit = len(inits); ncell = land_iy.size

    enso = np.array([nino34.get((t.year, t.month), 0.0) for t in inits])
    omi_vals = np.array([_nearest_omi(omi, t.year, t.month, t.day) for t in inits])  # (ninit,3)
    ehdfreq, sevmean = antecedent_persistence(fc_year, obs_vals, obs_lut, thr_by_doy,
                                              land_iy, land_ix, ndays)

    enso_c = np.repeat(enso, ncell)
    pc1_c = np.repeat(omi_vals[:, 0], ncell)
    pc2_c = np.repeat(omi_vals[:, 1], ncell)
    amp_c = np.repeat(omi_vals[:, 2], ncell)
    cols = np.column_stack([enso_c, pc1_c, pc2_c, amp_c,
                            ehdfreq.ravel(), sevmean.ravel()])
    return dict(X_ext=cols, names=list(EXT_FEATURE_NAMES))