#!/usr/bin/env python3
"""PA1: thông tin NGOÀI dự báo EC — tiền sử quan trắc và chỉ số quy mô lớn.

Every feature the project has used so far comes out of a single ECMWF run, so the
model has never been told what the atmosphere or the land surface was doing when
the forecast was issued.  Two sources are added here, both of which a real-time
system has in hand at initialisation:

  hist   observed Tmax history at the grid point, its region and the country,
         taken STRICTLY from days before the initialisation date;
  idx    the state of the two large-scale modes that set the odds of a Vietnamese
         heat spell weeks ahead: the MJO (NOAA/PSL OMI, daily) and ENSO
         (CPC ERSST5 Nino3.4 monthly anomaly).

ON THE PROJECT'S "OBSERVATIONS ARE TARGETS ONLY" RULE.  That rule exists to stop
the concurrent observation leaking into its own forecast, and it is not relaxed
here: `history_end_offset` is -1 day, so a sample initialised on date D may only
see observed days up to D-1.  The target lives at D+lt with lt >= 0, so no
observation at or after the forecast valid time is ever visible.  Every window is
built by slicing the continuous observed series at indices strictly below the
index of D, which makes the guarantee structural rather than a matter of care.

The Tmax climatology the anomalies are taken against is fitted on the SIXTEEN
TRAINING YEARS only, matching the rule adopted for every other fitted quantity.

Output: outputs/AI_Inp/external.npz
    hist   (ninit, npt, nh)   float32
    idx    (ninit, nx)        float32
    hist_names, idx_names
"""

from __future__ import annotations

import argparse
import datetime as dt
import os
import sys

import numpy as np
from netCDF4 import Dataset

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import ehd_common as C
import build_cube as B
from calc_climatology import circular_window_sum

EXT_DIR = os.path.join(C.OUT, "AI_Inp")
EXT_FILE = os.path.join(EXT_DIR, "external.npz")
EXTERNAL_DATA = os.path.join(C.ROOT, "data", "external")
OMI_FILE = os.path.join(EXTERNAL_DATA, "omi.txt")
NINO_FILE = os.path.join(EXTERNAL_DATA, "nino34.txt")

OBS_YEAR0 = 2002                   # first year of OBS/Tx_EC_0.25_2002-2023.nc
WINDOWS = (5, 10, 30)
HISTORY_END_OFFSET = -1            # last observed day a sample may see, relative to init


def obs_row(date):
    """Row of `date` in the continuous 365-day observed series."""
    return (date.year - OBS_YEAR0) * 365 + (C.doy_noleap(date) - 1)


def load_obs_tx(grid):
    """Observed Tmax as (nday, npt), and its day-of-year index."""
    with Dataset(C.OBS_TX_CONT) as ds:
        tx = np.ma.filled(ds["tasmax"][:], np.nan).astype("f4")
    tx = grid.to_points(tx)                                   # (nday, npt)
    nday = tx.shape[0]
    doy = np.arange(nday) % 365 + 1
    year = OBS_YEAR0 + np.arange(nday) // 365
    return tx, doy, year


def tx_climatology(tx, doy, year, npt):
    """31-day window mean of observed Tmax per (doy, point), TRAINING YEARS ONLY."""
    keep = np.isin(year, C.TRAIN_YEARS)
    s0 = np.zeros((365, npt)); s1 = np.zeros((365, npt))
    ok = np.isfinite(tx[keep])
    v = np.where(ok, tx[keep], 0.0).astype("f8")
    d = doy[keep] - 1
    np.add.at(s0, d, ok.astype("f8"))
    np.add.at(s1, d, v)
    n = circular_window_sum(s0)
    a = circular_window_sum(s1)
    with np.errstate(invalid="ignore", divide="ignore"):
        return np.where(n > 0, a / np.maximum(n, 1), np.nan).astype("f4")


def build_history(grid, init_dates):
    """Observed-history features for every (init, point)."""
    tx, doy, year = load_obs_tx(grid)
    clim = tx_climatology(tx, doy, year, grid.npt)
    anom = tx - clim[doy - 1]                                  # (nday, npt)

    with Dataset(os.path.join(C.OUT, "Cli", "Qp_tx_obs.nc")) as ds:
        thr = grid.to_points(np.ma.filled(ds["Tx_thres_obs"][:], np.nan).astype("f4"))
    ehd = (tx >= thr[doy - 1]).astype("f4")
    ehd[~np.isfinite(tx)] = np.nan

    reg = grid.reg
    names = (["h_tx%02d" % w for w in WINDOWS]
             + ["h_ehd30", "h_tx30_reg", "h_ehd30_reg", "h_tx30_nat", "h_tx30_grad"])
    out = np.full((len(init_dates), grid.npt, len(names)), np.nan, dtype="f4")

    for i, d0 in enumerate(init_dates):
        end = obs_row(d0) + HISTORY_END_OFFSET + 1        # exclusive; strictly before d0
        for j, w in enumerate(WINDOWS):
            lo = max(end - w, 0)
            out[i, :, j] = np.nanmean(anom[lo:end], axis=0) if end > lo else np.nan
        lo30 = max(end - 30, 0)
        e30 = np.nanmean(ehd[lo30:end], axis=0) if end > lo30 else np.full(grid.npt, np.nan)
        a30 = out[i, :, WINDOWS.index(30)]
        out[i, :, len(WINDOWS)] = e30
        for r in range(1, 8):                              # region means of the 30-day state
            m = reg == r
            if m.any():
                out[i, m, len(WINDOWS) + 1] = np.nanmean(a30[m])
                out[i, m, len(WINDOWS) + 2] = np.nanmean(e30[m])
        nat = np.nanmean(a30)
        out[i, :, len(WINDOWS) + 3] = nat
        # how far this point departs from the national state: local vs large scale
        out[i, :, len(WINDOWS) + 4] = a30 - nat
    return out, names


def load_omi():
    a = np.loadtxt(OMI_FILE)
    key = (a[:, 0].astype(int) * 10000 + a[:, 1].astype(int) * 100
           + a[:, 2].astype(int))
    return dict(zip(key.tolist(), a[:, 3:6].tolist()))


def load_nino():
    rows = [l.split() for l in open(NINO_FILE).read().splitlines()[1:] if l.strip()]
    a = np.array([[float(x) for x in r] for r in rows])
    return {int(y) * 100 + int(m): v for y, m, v in zip(a[:, 0], a[:, 1], a[:, 9])}


def build_indices(init_dates):
    """MJO and ENSO state at each initialisation date."""
    omi, nino = load_omi(), load_nino()
    names = ["x_omi1", "x_omi2", "x_omi_amp", "x_omi1_m15", "x_omi2_m15",
             "x_nino34", "x_nino34_3m"]
    out = np.full((len(init_dates), len(names)), np.nan, dtype="f4")
    for i, d0 in enumerate(init_dates):
        last = d0 + dt.timedelta(days=HISTORY_END_OFFSET)
        k = last.year * 10000 + last.month * 100 + last.day
        if k in omi:
            out[i, 0:3] = omi[k]
        w = [omi[kk] for kk in
             ((last - dt.timedelta(days=t)).year * 10000
              + (last - dt.timedelta(days=t)).month * 100
              + (last - dt.timedelta(days=t)).day for t in range(15))
             if kk in omi]
        if w:
            out[i, 3:5] = np.mean(w, axis=0)[:2]
        mk = last.year * 100 + last.month
        if mk in nino:
            out[i, 5] = nino[mk]
        three = []
        y, m = last.year, last.month
        for _ in range(3):
            if y * 100 + m in nino:
                three.append(nino[y * 100 + m])
            m -= 1
            if m == 0:
                y, m = y - 1, 12
        if three:
            out[i, 6] = float(np.mean(three))
    return out, names


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default=EXT_FILE)
    args = ap.parse_args()

    grid = C.Grid()
    meta = B.load_meta()
    init_dates = [dt.date.fromisoformat(str(s)) for s in meta["init_dates"]]
    print("[ext] %d initialisations, %d land points" % (len(init_dates), grid.npt),
          flush=True)
    print("[ext] history may see observed days up to init%+d"
          % HISTORY_END_OFFSET, flush=True)

    idx, idx_names = build_indices(init_dates)
    for j, n in enumerate(idx_names):
        v = idx[:, j]
        print("   %-12s thieu %4d/%d, mean %7.3f, sd %6.3f"
              % (n, int(np.isnan(v).sum()), v.size, np.nanmean(v), np.nanstd(v)),
              flush=True)

    hist, hist_names = build_history(grid, init_dates)
    for j, n in enumerate(hist_names):
        v = hist[:, :, j]
        print("   %-12s thieu %6d/%d, mean %7.3f, sd %6.3f"
              % (n, int(np.isnan(v).sum()), v.size, np.nanmean(v), np.nanstd(v)),
              flush=True)

    os.makedirs(EXT_DIR, exist_ok=True)
    np.savez(args.out, hist=hist, idx=idx,
             hist_names=np.array(hist_names), idx_names=np.array(idx_names),
             history_end_offset=np.array([HISTORY_END_OFFSET]))
    print("[done] -> %s" % args.out, flush=True)


if __name__ == "__main__":
    main()