"""F3-dev-01: rolling-origin (expanding window) forecasts of UNRATE.
Model: AR(p) on monthly first differences, intercept, OLS, p chosen by BIC in 0..6
at every origin, estimation sample from 1948, with impulse dummies for the
2020-04 and 2020-05 COVID changes. Deterministic; no randomness.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
"""
import json, os
import numpy as np, pandas as pd
from scipy import stats

DATA = "data/UNRATE.csv"; OUT = "out"
MODEL = "AR_BIC_diff"
PMAX = 6
DUMMY_DATES = [pd.Timestamp("2020-04-01"), pd.Timestamp("2020-05-01")]
HORIZONS = [1, 3]

raw = pd.read_csv(DATA)
n_rows = len(raw)
df = raw.copy()
df["observation_date"] = pd.to_datetime(df["observation_date"])
df["UNRATE"] = pd.to_numeric(df["UNRATE"], errors="coerce")
df = df.set_index("observation_date").asfreq("MS")  # full monthly grid
y_all = df["UNRATE"]

origins = pd.date_range("2023-07-01", "2026-06-01", freq="MS")


def fit_predict(y, origin, hmax, use_dummies=True):
    """y: level series on monthly grid, already truncated to dates <= origin."""
    dy = y.diff()  # NaN where either month missing -> those rows dropped in estimation
    best = None
    for p in range(PMAX + 1):
        X = [np.ones(len(dy))]
        for k in range(1, p + 1):
            X.append(dy.shift(k).values)
        cols = len(X)
        if use_dummies:
            for d in DUMMY_DATES:
                if d <= origin:
                    X.append((dy.index == d).astype(float))
        X = np.column_stack(X); Y = dy.values
        ok = ~np.isnan(Y) & ~np.isnan(X).any(1)
        # common estimation sample across p for comparable BIC: start after PMAX lags
        ok &= np.arange(len(Y)) > PMAX
        Xo, Yo = X[ok], Y[ok]
        b, *_ = np.linalg.lstsq(Xo, Yo, rcond=None)
        e = Yo - Xo @ b; n = len(Yo); kpar = Xo.shape[1]
        bic = n * np.log(e @ e / n) + kpar * np.log(n)
        if best is None or bic < best[0]:
            best = (bic, p, b[:cols], n)
    _, p, b, nreg = best
    # state for forecasting: interior missing levels linearly interpolated (uses only data <= origin);
    yy = y.copy()
    last_obs = yy.last_valid_index()
    yy.loc[:last_obs] = yy.loc[:last_obs].interpolate()
    dyy = list(yy.loc[:last_obs].diff().values)
    level = yy.loc[last_obs]
    steps = (origin.year - last_obs.year) * 12 + origin.month - last_obs.month + hmax
    fc = []
    for _ in range(steps):
        d = b[0] + sum(b[k] * dyy[-k] for k in range(1, p + 1))
        dyy.append(d); level = level + d; fc.append(level)
    # fc[i] is forecast for last_obs + i+1 months; return by horizon relative to origin
    off = steps - hmax
    return {h: fc[off + h - 1] for h in range(1, hmax + 1)}, p, nreg


def run(use_dummies=True):
    rows, ps = [], []
    for t in origins:
        y = y_all.loc[:t]
        naive = y.loc[y.last_valid_index()]
        fc, p, nreg = fit_predict(y, t, max(HORIZONS), use_dummies)
        ps.append(p)
        for h in HORIZONS:
            tgt = t + pd.DateOffset(months=h)
            act = y_all.get(tgt, np.nan)
            for m, f in (("naive", naive), (MODEL, fc[h])):
                rows.append(dict(origin_date=t.strftime("%Y-%m-%d"), target_date=tgt.strftime("%Y-%m-%d"),
                                 horizon=h, forecast=round(float(f), 6),
                                 actual=None if pd.isna(act) else float(act), model=m))
    return pd.DataFrame(rows), ps, nreg


led, ps, nreg_last = run(True)
led.to_csv(os.path.join(OUT, "forecasts.csv"), index=False, float_format="%.6f")
led_nd, _, _ = run(False)  # sensitivity only (no COVID dummies)


def errs(L, m, h):
    s = L[(L.model == m) & (L.horizon == h) & L.actual.notna()]
    e = s.actual - s.forecast
    return e.values, s.origin_date.values


res = {"task_id": "F3-dev-01", "model_name": MODEL,
       "model_description": ("AR(p) on monthly first differences of UNRATE with intercept, p in 0..6 re-selected by BIC at "
                             "each origin, OLS on an expanding window starting 1948 (first usable change 1948-02, common "
                             "sample after 6 lags) with impulse dummies for the 2020-04 and 2020-05 changes; iterated "
                             "multi-step forecasts cumulated to levels."),
       "n_origins": len(origins), "first_origin": origins[0].strftime("%Y-%m-%d"),
       "last_origin": origins[-1].strftime("%Y-%m-%d")}
extra = {}
for h in HORIZONS:
    en, _ = errs(led, "naive", h); em, _ = errs(led, MODEL, h)
    res[f"n_evaluated_h{h}"] = int(len(en))
    for nm, e in (("naive", en), ("model", em)):
        res[f"rmse_{nm}_h{h}"] = float(np.sqrt(np.mean(e ** 2)))
        res[f"mae_{nm}_h{h}"] = float(np.mean(np.abs(e)))
for h in HORIZONS:
    res[f"rel_improvement_rmse_h{h}"] = 1 - res[f"rmse_model_h{h}"] / res[f"rmse_naive_h{h}"]

# Uncertainty: Diebold-Mariano (squared loss), HAC variance with h-1 Bartlett-free rectangular lags
# (Harvey-Leybourne-Newbold small-sample correction), t(n-1) reference. Note the evaluation set has one gap.
for h in HORIZONS:
    en, _ = errs(led, "naive", h); em, _ = errs(led, MODEL, h)
    d = em ** 2 - en ** 2; n = len(d); dbar = d.mean(); dc = d - dbar
    g = [dc @ dc / n] + [dc[k:] @ dc[:-k] / n for k in range(1, h)]
    var = (g[0] + 2 * sum(g[1:])) / n
    dm = dbar / np.sqrt(var) if var > 0 else np.nan
    hln = dm * np.sqrt((n + 1 - 2 * h + h * (h - 1) / n) / n)
    extra[f"dm_stat_h{h}"] = float(hln)
    extra[f"dm_pvalue_two_sided_h{h}"] = float(2 * stats.t.sf(abs(hln), n - 1))
    e2, _ = errs(led_nd, MODEL, h)
    extra[f"sensitivity_no_covid_dummies_rmse_model_h{h}"] = float(np.sqrt(np.mean(e2 ** 2)))
    extra[f"sensitivity_no_covid_dummies_rel_improvement_h{h}"] = float(1 - np.sqrt(np.mean(e2 ** 2)) / res[f"rmse_naive_h{h}"])
extra["selected_p_counts"] = {str(k): int(v) for k, v in zip(*np.unique(ps, return_counts=True))}
extra["n_regression_rows_last_window"] = int(nreg_last)

n_missing = int(y_all.isna().sum())
res.update({
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": [
        {"issue": f"{n_missing} missing observation(s): " + ", ".join(d.strftime('%Y-%m-%d') for d in y_all.index[y_all.isna()]) + " (empty value field)",
         "handling": "Target pairs with missing actual (origin 2025-09 h1, origin 2025-07 h3) kept in ledger with empty actual and excluded from error statistics; differences touching the missing month dropped from estimation; for forecasting at later origins the interior gap is linearly interpolated from neighbours dated <= origin; at origin 2025-10 the naive forecast uses 2025-09 and the model forecasts 2025-10 as an intermediate step."},
        {"issue": "Extreme COVID-19 jump (2020-04 +10.4 pp, 2020-05 -1.6 pp) inside every training window",
         "handling": "Impulse dummies on the 2020-04 and 2020-05 monthly changes (dates known before all origins); sensitivity without dummies reported in report.md."},
        {"issue": "Final-vintage data (revised seasonal factors), not real-time vintages",
         "handling": "Errors computed on data as given; no real-time validity claimed."}],
    "n_input_rows": {"UNRATE.csv": int(n_rows)},
    "n_analysed": int(y_all.loc[:origins[-1]].notna().sum()),
    "random_seed": None,
})
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(res, f, indent=2)
with open(os.path.join(OUT, "supplementary.json"), "w") as f:
    json.dump(extra, f, indent=2)
print(json.dumps(res, indent=2)); print(json.dumps(extra, indent=2))