"""Rolling-origin (expanding window) forecasts of US UNRATE: naive vs AR(2) on first differences.
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. No randomness."""
import json, os
import numpy as np, pandas as pd
from scipy import stats

DATA = "data/UNRATE.csv"; OUT = "out"
P = 2                      # AR order on first differences (fixed a priori)
HORIZONS = [1, 3]
MODEL = "AR2_diff"
origins = pd.date_range("2023-07-01", "2026-06-01", freq="MS")

raw = pd.read_csv(DATA, dtype={"observation_date": str})
n_raw = len(raw)
raw["date"] = pd.to_datetime(raw["observation_date"], format="%Y-%m-%d")
y = raw.set_index("date")["UNRATE"].astype(float)
assert y.index.is_monotonic_increasing and not y.index.duplicated().any()
assert (y.index == pd.date_range(y.index[0], y.index[-1], freq="MS")).all()
missing_dates = [d.strftime("%Y-%m-%d") for d in y.index[y.isna()]]

def ar_forecast(s, steps):
    """s: series with obs dated <= origin (may contain NaN). Returns forecasts 1..steps months
    ahead of the last *observed* month."""
    s = s.loc[:s.last_valid_index()]
    d = s.diff()
    X = pd.concat({f"l{i}": d.shift(i) for i in range(1, P + 1)}, axis=1)
    reg = pd.concat([d.rename("y"), X], axis=1).dropna()        # drop rows touching missing values
    A = np.column_stack([np.ones(len(reg)), reg[[f"l{i}" for i in range(1, P + 1)]].values])
    beta, *_ = np.linalg.lstsq(A, reg["y"].values, rcond=None)
    # forecasting state: interior missing values linearly interpolated using only obs <= origin
    si = s.interpolate(method="linear", limit_area="inside")
    lags = list(si.diff().values[-P:][::-1])               # [d_T, d_{T-1}]
    level, out = si.values[-1], []
    for _ in range(steps):
        dhat = beta[0] + np.dot(beta[1:], lags[:P])
        level += dhat; out.append(level); lags = [dhat] + lags[:P - 1]
    return np.array(out), len(reg)

rows = []
for t in origins:
    s = y.loc[:t]                                # information rule: only obs dated <= t
    last_obs = s.last_valid_index()
    gap = (t.year - last_obs.year) * 12 + (t.month - last_obs.month)   # months origin is ahead of last obs
    fc, nreg = ar_forecast(s, gap + max(HORIZONS))
    for h in HORIZONS:
        tgt = t + pd.DateOffset(months=h)
        act = y.get(tgt, np.nan)
        rows.append(dict(origin_date=t.strftime("%Y-%m-%d"), target_date=tgt.strftime("%Y-%m-%d"),
                         horizon=h, forecast=float(s.loc[last_obs]), actual=act, model="naive"))
        rows.append(dict(origin_date=t.strftime("%Y-%m-%d"), target_date=tgt.strftime("%Y-%m-%d"),
                         horizon=h, forecast=float(fc[gap + h - 1]), actual=act, model=MODEL))
fc_df = pd.DataFrame(rows)
fc_df.to_csv(f"{OUT}/forecasts.csv", index=False, float_format="%.6f")

res = {"task_id": "F3-dev-01", "model_name": MODEL,
       "model_description": ("AR(2) with intercept on monthly first differences of UNRATE, estimated by OLS on an "
                             "expanding window starting 1948-01 (regression rows touching missing values dropped), "
                             "re-estimated at every origin; level forecasts by iterating the difference recursion."),
       "n_origins": len(origins), "first_origin": origins[0].strftime("%Y-%m-%d"),
       "last_origin": origins[-1].strftime("%Y-%m-%d")}
dm = {}
for h in HORIZONS:
    sub = fc_df[fc_df.horizon == h]
    ev = sub.dropna(subset=["actual"])
    e = {m: (ev[ev.model == m].set_index("origin_date")["actual"] - ev[ev.model == m].set_index("origin_date")["forecast"])
         for m in ["naive", MODEL]}
    n = len(e["naive"]); assert n == len(e[MODEL])
    res[f"n_evaluated_h{h}"] = int(n)
    for m, tag in [("naive", "naive"), (MODEL, "model")]:
        res[f"rmse_{tag}_h{h}"] = float(np.sqrt(np.mean(e[m] ** 2)))
        res[f"mae_{tag}_h{h}"] = float(np.mean(np.abs(e[m])))
    # Diebold-Mariano (squared loss), HAC variance with h-1 Bartlett-free (rectangular) lags, HLN small-sample correction
    dl = (e[MODEL] ** 2 - e["naive"] ** 2).values
    dbar = dl.mean(); dc = dl - dbar
    gam = [np.sum(dc[k:] * dc[:n - k]) / n for k in range(h)]
    var = (gam[0] + 2 * sum(gam[1:])) / n
    if var <= 0: var = gam[0] / n
    hln = np.sqrt((n + 1 - 2 * h + h * (h - 1) / n) / n)
    stat = hln * dbar / np.sqrt(var)
    dm[h] = (float(stat), float(2 * stats.t.sf(abs(stat), df=n - 1)), float(dbar))
for h in HORIZONS:
    res[f"rel_improvement_rmse_h{h}"] = 1 - res[f"rmse_model_h{h}"] / res[f"rmse_naive_h{h}"]
for h in HORIZONS:
    res[f"dm_stat_h{h}"], res[f"dm_pvalue_h{h}"], res[f"mean_loss_diff_h{h}"] = dm[h]
res.update({
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": [
        {"issue": f"Missing UNRATE value(s) at {', '.join(missing_dates)} (empty field).",
         "handling": "Origin 2025-10-01: naive uses 2025-09 value; AR forecasts iterated from last observed month (2025-09). "
                     "Pairs with missing target (origin 2025-09 h1, origin 2025-07 h3) kept in ledger with empty actual and "
                     "excluded from error statistics. In later training windows, regression rows touching the missing month are "
                     "dropped; the forecasting state uses linear interpolation between observations dated <= origin."},
        {"issue": "Final (revised) vintage, not real-time data.",
         "handling": "Errors measure forecastability of the revised series; no real-time validity claimed."},
        {"issue": "COVID-19 spike 2020-03..2020-05 (4.4 -> 14.8 -> 13.2) inside every training window.",
         "handling": "Kept as genuine data; no outlier treatment (specification fixed a priori, not tuned on test errors)."}],
    "n_input_rows": {"UNRATE.csv": int(n_raw)},
})
last_win = y.loc[:origins[-1]]
res["n_analysed"] = int(last_win.notna().sum())
res["n_regression_rows_last_window"] = int(ar_forecast(last_win, 1)[1])
res["random_seed"] = None
with open(f"{OUT}/results.json", "w") as f:
    json.dump(res, f, indent=2)

# exclusion ledger (training data for the last window)
n1 = int(y.notna().sum()); n2 = res["n_analysed"]
pd.DataFrame([
    dict(step=1, rule="drop rows with missing UNRATE", n_before=n_raw, n_after=n1, n_removed=n_raw - n1,
         rationale="value not available (2025-10-01)"),
    dict(step=2, rule="drop observations dated after last origin 2026-06-01", n_before=n1, n_after=n2, n_removed=n1 - n2,
         rationale="information rule: training only on obs <= origin (these are used only as evaluation actuals)"),
]).to_csv(f"{OUT}/exclusions.csv", index=False)
print(json.dumps({k: v for k, v in res.items() if k not in ("data_issues", "model_description")}, indent=1))