"""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, os
import numpy as np
import pandas as pd
from scipy import stats

DATA = "data/UNRATE.csv"
OUT = "out"
P = 2                      # AR order on first differences
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)
s = pd.Series(raw["UNRATE"].values, index=pd.to_datetime(raw["observation_date"]), dtype=float)
assert s.index.is_monotonic_increasing and not s.index.duplicated().any()
assert len(pd.date_range(s.index[0], s.index[-1], freq="MS")) == len(s)


def fit_ar_diff(y):
    """OLS of d_t on const, d_{t-1..t-P}; rows with any NaN dropped. y: level series up to origin."""
    d = y.diff()
    X = pd.concat([d.shift(k) for k in range(1, P + 1)], axis=1)
    df = pd.concat([d, X], axis=1).dropna()
    Y = df.iloc[:, 0].values
    Xm = np.column_stack([np.ones(len(df)), df.iloc[:, 1:].values])
    beta, *_ = np.linalg.lstsq(Xm, Y, rcond=None)
    return beta, len(df)


def forecast_model(y, origin, h):
    """Forecast level at origin+h using only y (already truncated to <= origin)."""
    beta, nreg = fit_ar_diff(y)
    last_obs = y.last_valid_index()
    # state: interior missing values interpolated from data <= origin (state only, not training)
    ys = y.loc[:last_obs].interpolate(method="linear", limit_area="inside")
    steps = h + (len(pd.date_range(last_obs, origin, freq="MS")) - 1)  # extra steps if origin missing
    level = ys.iloc[-1]
    dl = list(ys.diff().iloc[-P:].values[::-1])  # most recent first
    for _ in range(steps):
        dn = beta[0] + sum(beta[k + 1] * dl[k] for k in range(P))
        level += dn
        dl = [dn] + dl[:-1]
    return level, nreg


rows = []
for t in ORIGINS:
    y = s.loc[:t]                       # information set: observations dated <= t
    naive = y.dropna().iloc[-1]
    for h in HORIZONS:
        tgt = t + pd.DateOffset(months=h)
        actual = s.get(tgt, np.nan)
        fm, nreg = forecast_model(y, t, h)
        for m, f in (("naive", naive), (MODEL_NAME, fm)):
            rows.append(dict(origin_date=t.strftime("%Y-%m-%d"), target_date=tgt.strftime("%Y-%m-%d"),
                             horizon=h, forecast=round(float(f), 10),
                             actual=(None if pd.isna(actual) else float(actual)), model=m))
fc = pd.DataFrame(rows)
os.makedirs(OUT, exist_ok=True)
fc.to_csv(os.path.join(OUT, "forecasts.csv"), index=False)

res = {"task_id": "F3-dev-01", "model_name": MODEL_NAME,
       "model_description": ("AR(2) with intercept on first differences of UNRATE, estimated by OLS on all available "
                             "observations from 1948-01 up to each origin (expanding window; rows touching missing "
                             "values dropped), iterated h-step forecasts of differences cumulated onto the last "
                             "observed level."),
       "n_origins": int(len(ORIGINS)), "first_origin": ORIGINS[0].strftime("%Y-%m-%d"),
       "last_origin": ORIGINS[-1].strftime("%Y-%m-%d")}
extra = {}
for h in HORIZONS:
    sub = fc[(fc.horizon == h) & fc.actual.notna()]
    en = (sub[sub.model == "naive"].set_index("origin_date").pipe(lambda d: d.actual - d.forecast))
    em = (sub[sub.model == MODEL_NAME].set_index("origin_date").pipe(lambda d: d.actual - d.forecast))
    assert (en.index == em.index).all()
    n = len(en)
    res[f"n_evaluated_h{h}"] = int(n)
    for name, e in (("naive", en), ("model", em)):
        res[f"rmse_{name}_h{h}"] = float(np.sqrt(np.mean(e ** 2)))
        res[f"mae_{name}_h{h}"] = float(np.mean(np.abs(e)))
    res[f"rel_improvement_rmse_h{h}"] = 1 - res[f"rmse_model_h{h}"] / res[f"rmse_naive_h{h}"]
    # Diebold-Mariano (squared-error loss), rectangular HAC with h-1 lags, Harvey-Leybourne-Newbold correction.
    # Note: with a missing pair the series has a gap; treated as consecutive (approximation, documented).
    dlt = (em.values ** 2 - en.values ** 2)
    dbar = dlt.mean(); dc = dlt - 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
    dm = dbar / np.sqrt(var)
    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), df=n - 1))
    extra[f"mean_error_naive_h{h}"] = float(en.mean())
    extra[f"mean_error_model_h{h}"] = float(em.mean())

# exclusion ledger (observation level, last training window)
n_nonmiss = int(s.notna().sum())
last_window = s.loc[:ORIGINS[-1]]
n_analysed = int(last_window.notna().sum())
_, nreg_last = fit_ar_diff(last_window)
ex = pd.DataFrame([
    dict(step=1, rule="drop rows with missing UNRATE value", n_before=n_raw, n_after=n_nonmiss,
         n_removed=n_raw - n_nonmiss, rationale="empty value field (2025-10-01) cannot be used for estimation"),
    dict(step=2, rule="drop observations dated after last origin 2026-06-01", n_before=n_nonmiss,
         n_after=n_analysed, n_removed=n_nonmiss - n_analysed,
         rationale="information rule: last training window uses only data <= last origin (later months are evaluation targets only)"),
])
ex.to_csv(os.path.join(OUT, "exclusions.csv"), index=False)

missing_dates = [d.strftime("%Y-%m-%d") for d in s.index[s.isna()]]
res.update({
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": [
        {"issue": f"missing observation(s) at {missing_dates}",
         "handling": ("target pairs with missing actual kept in ledger with empty actual and excluded from error "
                      "statistics; origin 2025-10 uses latest earlier observation (naive) / iterates one extra step "
                      "(model); for later origins the missing level is linearly interpolated from data <= origin for "
                      "the AR state only; regression rows touching it are dropped")},
        {"issue": "COVID-19 outliers (2020-04 = 14.8%, large swings 2020-04..2020-08) in the estimation sample",
         "handling": "kept (no outlier treatment, decided a priori); they influence OLS coefficients - see report"},
        {"issue": "final-vintage (revised) data, not real-time vintages",
         "handling": "results describe pseudo-out-of-sample accuracy on revised data; no real-time claim"},
    ],
    "n_input_rows": {"UNRATE.csv": int(n_raw)},
    "n_analysed": n_analysed,
    "n_regression_rows_last_window": int(nreg_last),
    "random_seed": None,
})
res.update(extra)
beta_last, _ = fit_ar_diff(last_window)
res["ar_coefficients_last_window"] = {"const": float(beta_last[0]), "phi1": float(beta_last[1]), "phi2": float(beta_last[2])}
with open(os.path.join(OUT, "results.json"), "w") as fh:
    json.dump(res, fh, indent=2)
for h in HORIZONS:
    print(f"h={h}: n={res[f'n_evaluated_h{h}']} RMSE naive={res[f'rmse_naive_h{h}']:.4f} model={res[f'rmse_model_h{h}']:.4f} "
          f"rel_impr={res[f'rel_improvement_rmse_h{h}']:.4f} DM={res[f'dm_stat_h{h}']:.3f} p={res[f'dm_pvalue_two_sided_h{h}']:.3f}")