"""Rolling-origin 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 and out/forecasts.csv. Deterministic (no randomness)."""
import json, os
import numpy as np, pandas as pd

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUT = os.path.join(ROOT, "out")
raw = pd.read_csv(os.path.join(ROOT, "data", "UNRATE.csv"), dtype={"observation_date": str})
n_input = len(raw)
s = pd.Series(raw["UNRATE"].values, index=pd.PeriodIndex(raw["observation_date"], freq="M"))
missing = [str(p.to_timestamp().date()) for p in s.index[s.isna()]]

P = 2  # AR order on first differences, fixed a priori
MODEL = "AR2_diff"
origins = pd.period_range("2023-07", "2026-06", freq="M")
H = [1, 3]

def fit_ar_diff(y):
    """OLS of dy_t on const, dy_{t-1..t-P}. y: pandas Series (no NaN, contiguous)."""
    dy = np.diff(y.values)
    X = np.column_stack([np.ones(len(dy) - P)] + [dy[P - k:len(dy) - k] for k in range(1, P + 1)])
    beta, *_ = np.linalg.lstsq(X, dy[P:], rcond=None)
    return beta, dy

def forecast_ar_diff(beta, y_last, dy_hist, steps):
    hist = list(dy_hist[-P:]); lvl = y_last; path = []
    for _ in range(steps):
        d = beta[0] + sum(beta[k] * hist[-k] for k in range(1, P + 1))
        hist.append(d); lvl = lvl + d; path.append(lvl)
    return path

rows = []; last_n = None
for t in origins:
    w = s[s.index <= t]                       # information set: dated <= t
    obs = w.dropna()
    last_obs_p = obs.index[-1]                # latest observed month <= t
    naive = float(obs.iloc[-1])
    # model input: trim trailing missing, interpolate interior gaps using only data <= t
    wm = w.loc[:last_obs_p].interpolate(method="linear")
    beta, dy = fit_ar_diff(wm)
    maxsteps = max(H) + (t - last_obs_p).n
    path = forecast_ar_diff(beta, float(wm.iloc[-1]), dy, maxsteps)
    last_n = int(obs.shape[0])
    for h in H:
        tgt = t + h
        act = s.get(tgt, np.nan)
        steps = (tgt - last_obs_p).n
        for m, f in (("naive", naive), (MODEL, path[steps - 1])):
            rows.append({"origin_date": str(t.to_timestamp().date()), "target_date": str(tgt.to_timestamp().date()),
                         "horizon": h, "forecast": round(float(f), 10),
                         "actual": (None if pd.isna(act) else float(act)), "model": m})

fc = pd.DataFrame(rows)
fc.to_csv(os.path.join(OUT, "forecasts.csv"), index=False, float_format="%.10g")

res = {"task_id": "F3-dev-01", "model_name": MODEL,
       "model_description": ("AR(2) on first differences of UNRATE with intercept, estimated by OLS on the expanding window "
                             "from 1948-01 (first usable difference regression row 1948-04) to the origin; levels forecast by "
                             "iterating the difference recursion and cumulating; lag order fixed a priori, no selection."),
       "n_origins": len(origins), "first_origin": str(origins[0].to_timestamp().date()),
       "last_origin": str(origins[-1].to_timestamp().date())}
for h in H:
    ev = fc[(fc.horizon == h) & fc.actual.notna()]
    res[f"n_evaluated_h{h}"] = int((ev.model == "naive").sum())
    for m, key in (("naive", "naive"), (MODEL, "model")):
        e = ev[ev.model == m]; err = e.actual - e.forecast
        res[f"rmse_{key}_h{h}"] = float(np.sqrt(np.mean(err ** 2)))
        res[f"mae_{key}_h{h}"] = float(np.mean(np.abs(err)))
from scipy import stats
for h in H:
    ev = fc[(fc.horizon == h) & fc.actual.notna()]
    en = (ev[ev.model == "naive"].set_index("origin_date").eval("actual - forecast"))
    em = (ev[ev.model == MODEL].set_index("origin_date").eval("actual - forecast")).loc[en.index]
    d = (em ** 2 - en ** 2).values; n = len(d); dc = d - d.mean()
    lrv = np.mean(dc ** 2) + 2 * sum((1 - k / h) * np.mean(dc[k:] * dc[:-k]) for k in range(1, h))  # Bartlett, h-1 lags
    dm = d.mean() / np.sqrt(lrv / n)
    dm *= np.sqrt((n + 1 - 2 * h + h * (h - 1) / n) / n)  # Harvey-Leybourne-Newbold correction
    res[f"dm_stat_h{h}"] = float(dm)
    res[f"dm_pvalue_h{h}"] = float(2 * stats.t.sf(abs(dm), df=n - 1))
for h in H:
    res[f"rel_improvement_rmse_h{h}"] = 1 - res[f"rmse_model_h{h}"] / res[f"rmse_naive_h{h}"]
res.update({"feasibility": "complete", "feasibility_reason": None,
    "data_issues": [{"issue": f"Missing UNRATE value at {d} (empty field)",
                     "handling": ("Pairs targeting this month kept in ledger with empty actual and excluded from error stats; "
                                  "naive at origin 2025-10 uses 2025-09; model forecasts from 2025-09 at that origin and "
                                  "linearly interpolates the gap (using only data <= origin) at later origins.")} for d in missing]
                   + [{"issue": "Final-vintage (revised) data, not real-time vintages",
                       "handling": "Errors computed on data as given; no real-time validity claimed."}],
    "n_input_rows": {"UNRATE.csv": int(n_input)}, "n_analysed": last_n, "random_seed": None})
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(res, f, indent=2)
for k in ("n_evaluated_h1", "n_evaluated_h3", "rmse_naive_h1", "rmse_model_h1", "rmse_naive_h3", "rmse_model_h3",
          "rel_improvement_rmse_h1", "rel_improvement_rmse_h3", "dm_stat_h1", "dm_pvalue_h1", "dm_stat_h3", "dm_pvalue_h3"):
    print(k, res[k])