"""F3-dev-01: rolling-origin (expanding window) forecasts of 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. Deterministic (no randomness)."""
import json
import os
import numpy as np
import pandas as pd

DATA = "data/UNRATE.csv"
OUT = "out"
P = 2                      # AR lag order on first differences (fixed a priori)
HORIZONS = [1, 3]
MODEL_NAME = "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")
s = raw.set_index("date")["UNRATE"].astype(float)
assert s.index.is_monotonic_increasing and not s.index.duplicated().any()
assert (s.index == pd.date_range(s.index[0], s.index[-1], freq="MS")).all()
missing_dates = [d.strftime("%Y-%m-%d") for d in s.index[s.isna()]]


def fit_ar_diff(y):
    """OLS of dy_t on const + dy_{t-1..t-P}, complete cases only. y: level series (may contain NaN)."""
    dy = y.diff()
    X = pd.concat([dy.shift(k) for k in range(1, P + 1)], axis=1)
    df = pd.concat([dy, X], axis=1).dropna()
    Y = df.iloc[:, 0].to_numpy()
    Xm = np.column_stack([np.ones(len(df)), df.iloc[:, 1:].to_numpy()])
    beta, *_ = np.linalg.lstsq(Xm, Y, rcond=None)
    return beta, len(df)


def forecast_ar_diff(y, origin, hmax):
    """Forecast levels at origin+1..origin+hmax using only y dated <= origin.
    Missing levels inside the window (and the origin month if missing) are filled sequentially
    with the model's own one-step forecast from earlier data (no information after the origin)."""
    y = y.loc[:origin]
    beta, nobs = fit_ar_diff(y)
    lv = y.copy()
    # fill internal missing values (needed only if they are among the last P+1 levels used)
    for i in np.where(lv.isna().to_numpy())[0]:
        if i >= len(lv) - (P + 1) - hmax:   # only recent gaps matter for the recursion
            d = lv.diff().iloc[i - P:i].to_numpy()[::-1]   # dy_{i-1},...,dy_{i-P}
            lv.iloc[i] = lv.iloc[i - 1] + beta[0] + beta[1:] @ d
    assert not lv.iloc[-(P + 1):].isna().any()
    hist_d = list(lv.diff().iloc[-P:].to_numpy()[::-1])   # most recent first
    level = lv.iloc[-1]
    out = []
    for _ in range(hmax):
        d = beta[0] + beta[1:] @ np.array(hist_d[:P])
        level = level + d
        hist_d.insert(0, d)
        out.append(level)
    return out, beta, nobs


rows = []
last_window_n = None
for t in ORIGINS:
    window = s.loc[:t]
    naive = window.dropna().iloc[-1]
    fc, beta, nobs = forecast_ar_diff(s, t, max(HORIZONS))
    last_window_n = int(window.notna().sum())
    for h in HORIZONS:
        target = t + pd.DateOffset(months=h)
        actual = s.get(target, np.nan)
        for model, f in (("naive", naive), (MODEL_NAME, fc[h - 1])):
            rows.append(dict(origin_date=t.strftime("%Y-%m-%d"), target_date=target.strftime("%Y-%m-%d"),
                             horizon=h, forecast=round(float(f), 10),
                             actual=(None if pd.isna(actual) else float(actual)), model=model))

fc_df = pd.DataFrame(rows)
assert len(fc_df) == len(ORIGINS) * len(HORIZONS) * 2
assert (pd.to_datetime(fc_df.target_date) > pd.to_datetime(fc_df.origin_date)).all()
fc_df.to_csv(os.path.join(OUT, "forecasts.csv"), index=False, float_format="%.10g")

res = {"task_id": "F3-dev-01", "model_name": MODEL_NAME,
       "model_description": ("AR(2) with intercept on first differences of UNRATE, levels forecast by iterating the "
                             "difference equation; estimated by OLS on complete cases over an expanding window from "
                             "1948 (first usable dependent obs 1948-04) to the origin; lag order fixed a priori; a "
                             "missing in-window level is replaced for forecasting by the model's own forecast from "
                             "earlier data."),
       "n_origins": len(ORIGINS), "first_origin": ORIGINS[0].strftime("%Y-%m-%d"),
       "last_origin": ORIGINS[-1].strftime("%Y-%m-%d")}
for h in HORIZONS:
    sub = fc_df[(fc_df.horizon == h) & fc_df.actual.notna()]
    n_ev = sub[sub.model == "naive"].shape[0]
    res[f"n_evaluated_h{h}"] = int(n_ev)
for h in HORIZONS:
    for model, key in (("naive", "naive"), (MODEL_NAME, "model")):
        sub = fc_df[(fc_df.horizon == h) & (fc_df.model == model) & fc_df.actual.notna()]
        e = sub.actual.to_numpy() - sub.forecast.to_numpy()
        res[f"rmse_{key}_h{h}"] = float(np.sqrt(np.mean(e ** 2)))
        res[f"mae_{key}_h{h}"] = float(np.mean(np.abs(e)))
for h in HORIZONS:
    res[f"rel_improvement_rmse_h{h}"] = 1.0 - res[f"rmse_model_h{h}"] / res[f"rmse_naive_h{h}"]

excluded_pairs = fc_df[fc_df.actual.isna() & (fc_df.model == "naive")][["origin_date", "horizon"]]
res.update({
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": [{"issue": f"Missing UNRATE value at {d} (empty field).",
                     "handling": ("Kept as missing. Naive forecast at origin 2025-10-01 uses 2025-09-01. AR estimation "
                                  "uses complete cases only; for forecasting, the missing level is replaced by the "
                                  "model's own forecast from data before it. (origin, horizon) pairs whose target is "
                                  "this month are kept in forecasts.csv with empty actual and excluded from error "
                                  "statistics: " + "; ".join(f"{r.origin_date} h={r.horizon}"
                                                             for r in excluded_pairs.itertuples()) + ".")}
                    for d in missing_dates] + [
        {"issue": "Data are the final (revised) vintage, not real-time vintages.",
         "handling": "Errors are pseudo-out-of-sample on revised data; no real-time validity is claimed."}],
    "n_input_rows": {"UNRATE.csv": int(n_raw)},
    "n_analysed": last_window_n,
    "random_seed": None,
})
# diagnostic (not a required key): coefficients at last origin
res["diag_ar_coef_last_origin"] = [float(b) for b in beta]
res["diag_ar_nobs_last_origin"] = int(nobs)


# Diebold-Mariano test (squared-error loss), HAC long-run variance with rectangular kernel, h-1 lags,
# Harvey-Leybourne-Newbold small-sample correction; t(n-1) reference. Diagnostic only.
from scipy import stats
for h in HORIZONS:
    a = fc_df[(fc_df.horizon == h) & fc_df.actual.notna()]
    en = (a[a.model == "naive"].actual - a[a.model == "naive"].forecast).to_numpy()
    em = (a[a.model == MODEL_NAME].actual - a[a.model == MODEL_NAME].forecast).to_numpy()
    d = em ** 2 - en ** 2          # positive => model worse
    n = len(d); dbar = d.mean(); dc = d - dbar
    lrv = np.sum(dc ** 2) / n + 2 * sum(np.sum(dc[k:] * dc[:-k]) / n for k in range(1, h))
    dm = dbar / np.sqrt(lrv / n)
    hln = dm * np.sqrt((n + 1 - 2 * h + h * (h - 1) / n) / n)
    res[f"diag_dm_stat_h{h}"] = float(hln)
    res[f"diag_dm_pvalue_two_sided_h{h}"] = float(2 * stats.t.sf(abs(hln), df=n - 1))

with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(res, f, indent=2)

# exclusion ledger (observation-level)
n1 = int(s.notna().sum())
n2 = int(s.loc[:ORIGINS[-1]].notna().sum())
pd.DataFrame([
    dict(step=1, rule="drop rows with missing UNRATE", n_before=n_raw, n_after=n1, n_removed=n_raw - n1,
         rationale="empty value field (" + ", ".join(missing_dates) + ")"),
    dict(step=2, rule="drop observations dated after last origin 2026-06-01 from training", n_before=n1, n_after=n2,
         n_removed=n1 - n2, rationale="information rule: training uses only data dated <= origin (used only as targets)"),
]).to_csv(os.path.join(OUT, "exclusions.csv"), index=False)

for h in HORIZONS:
    print(f"h={h}: n={res[f'n_evaluated_h{h}']} RMSE naive={res[f'rmse_naive_h{h}']:.4f} "
          f"model={res[f'rmse_model_h{h}']:.4f} rel.impr={res[f'rel_improvement_rmse_h{h}']:.4f}")