"""F3-dev-01: rolling-origin forecasts of US UNRATE, naive vs ARIMA(2,1,0).
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Reads only data/UNRATE.csv; writes out/results.json, out/forecasts.csv, out/exclusions.csv.
"""
import json, os, warnings
import numpy as np
import pandas as pd
from scipy import stats
from statsmodels.tsa.statespace.sarimax import SARIMAX

warnings.filterwarnings("ignore")
OUT = "out"
os.makedirs(OUT, exist_ok=True)
MODEL = "ARIMA210_nc"
HORIZONS = [1, 3]

raw = pd.read_csv("data/UNRATE.csv", dtype={"observation_date": str})
n_raw = len(raw)
s = pd.Series(pd.to_numeric(raw["UNRATE"], errors="coerce").values,
              index=pd.DatetimeIndex(pd.to_datetime(raw["observation_date"]), freq="MS"))
assert s.index.is_monotonic_increasing and not s.index.duplicated().any()
missing_dates = [d.strftime("%Y-%m-%d") for d in s.index[s.isna()]]

origins = pd.date_range("2023-07-01", "2026-06-01", freq="MS")
rows = []
for t in origins:
    train = s.loc[:t]                       # only observations dated <= t
    # naive: latest available observation dated <= t
    naive = float(train.dropna().iloc[-1])
    # model: ARIMA(2,1,0), no constant, Gaussian MLE via Kalman filter (handles NaN)
    mod = SARIMAX(train, order=(2, 1, 0), trend="n")
    res = mod.fit(disp=False, method="lbfgs", maxiter=500)
    fc = res.forecast(steps=max(HORIZONS))
    for h in HORIZONS:
        tgt = t + pd.DateOffset(months=h)
        act = s.get(tgt, np.nan)
        for name, f in (("naive", naive), (MODEL, float(fc.iloc[h - 1]))):
            rows.append({"origin_date": t.strftime("%Y-%m-%d"), "target_date": tgt.strftime("%Y-%m-%d"),
                         "horizon": h, "forecast": round(f, 10),
                         "actual": (None if pd.isna(act) else float(act)), "model": name})
fc_df = pd.DataFrame(rows)
fc_df.to_csv(f"{OUT}/forecasts.csv", index=False, na_rep="")

def err(model, h):
    d = fc_df[(fc_df.model == model) & (fc_df.horizon == h) & fc_df.actual.notna()].sort_values("origin_date")
    return (d.actual - d.forecast).to_numpy(dtype=float)

def dm_test(e1, e2, h):
    """Diebold-Mariano (squared loss), HAC variance with h-1 rectangular lags, HLN small-sample correction."""
    d = e1 ** 2 - e2 ** 2
    n = len(d); dbar = d.mean(); dc = d - dbar
    g = [np.sum(dc[k:] * dc[:n - k]) / n for k in range(h)]
    v = (g[0] + 2 * sum(g[1:])) / n
    if v <= 0:
        return None, None
    dm = dbar / np.sqrt(v)
    hln = np.sqrt((n + 1 - 2 * h + h * (h - 1) / n) / n)
    stat = dm * hln
    p = 2 * stats.t.sf(abs(stat), df=n - 1)
    return float(stat), float(p)

R = {"task_id": "F3-dev-01", "model_name": MODEL,
     "model_description": ("ARIMA(2,1,0) without constant, i.e. AR(2) on first differences of UNRATE, "
                           "estimated by Gaussian maximum likelihood (Kalman filter, statsmodels SARIMAX, missing months skipped) "
                           "on an expanding window starting 1948-01-01 and ending at the origin, re-estimated at each origin; "
                           "multi-step forecasts are iterated."),
     "n_origins": int(len(origins)), "first_origin": origins[0].strftime("%Y-%m-%d"),
     "last_origin": origins[-1].strftime("%Y-%m-%d")}
for h in HORIZONS:
    en, em = err("naive", h), err(MODEL, h)
    assert len(en) == len(em)
    R[f"n_evaluated_h{h}"] = int(len(en))
for h in HORIZONS:
    en = err("naive", h)
    R[f"rmse_naive_h{h}"] = float(np.sqrt(np.mean(en ** 2)))
    R[f"mae_naive_h{h}"] = float(np.mean(np.abs(en)))
for h in HORIZONS:
    em = err(MODEL, h)
    R[f"rmse_model_h{h}"] = float(np.sqrt(np.mean(em ** 2)))
    R[f"mae_model_h{h}"] = float(np.mean(np.abs(em)))
for h in HORIZONS:
    R[f"rel_improvement_rmse_h{h}"] = 1.0 - R[f"rmse_model_h{h}"] / R[f"rmse_naive_h{h}"]
for h in HORIZONS:
    st, p = dm_test(err("naive", h), err(MODEL, h), h)
    R[f"dm_stat_h{h}"] = st      # >0 means model has lower squared error
    R[f"dm_pvalue_h{h}"] = p
n_missing_targets = {h: int(((fc_df.model == "naive") & (fc_df.horizon == h) & fc_df.actual.isna()).sum()) for h in HORIZONS}
R["feasibility"] = "complete"
R["feasibility_reason"] = None
R["data_issues"] = [
    {"issue": f"Missing observation(s) in UNRATE.csv: {missing_dates}",
     "handling": ("Naive at origin 2025-10-01 uses the 2025-09-01 value; the ARIMA likelihood/filter skips the missing month; "
                  f"(origin,horizon) pairs targeting it are kept in forecasts.csv with empty actual and excluded from error statistics "
                  f"(h1: {n_missing_targets[1]}, h3: {n_missing_targets[3]}).")},
    {"issue": "Data are the final (revised) vintage retrieved 2026-10-30, not real-time vintages.",
     "handling": "Evaluation is pseudo-out-of-sample on revised data; no real-time validity is claimed."},
    {"issue": "COVID-19 spike (2020-04 = 14.8%) lies inside every training window.",
     "handling": "Retained (genuine observation); noted as a limitation for Gaussian MLE of the AR coefficients."},
]
R["n_input_rows"] = {"UNRATE.csv": int(n_raw)}
last_train = s.loc[:origins[-1]]
R["n_analysed"] = int(last_train.notna().sum())
R["random_seed"] = None
with open(f"{OUT}/results.json", "w") as f:
    json.dump(R, f, indent=2)

# exclusion ledger: raw rows -> observations in the last training window
n1 = int(s.notna().sum())
excl = pd.DataFrame([
    {"step": 1, "rule": "drop rows with missing UNRATE value", "n_before": n_raw, "n_after": n1,
     "n_removed": n_raw - n1, "rationale": "empty value field (2025-10-01); no data to use"},
    {"step": 2, "rule": "drop observations dated after last origin (2026-06-01) from training",
     "n_before": n1, "n_after": R["n_analysed"], "n_removed": n1 - R["n_analysed"],
     "rationale": "information rule: training uses only obs <= origin; these serve only as evaluation actuals"},
])
excl.to_csv(f"{OUT}/exclusions.csv", index=False)
print(json.dumps({k: R[k] for k in R if k.startswith(("rmse", "mae", "rel", "n_eval", "dm"))}, indent=1))