"""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
"""
import json
import numpy as np
import pandas as pd

DATA = "data/UNRATE.csv"
MODEL = "AR2_diff"
P = 2
HORIZONS = [1, 3]

raw = pd.read_csv(DATA, dtype={"observation_date": str, "UNRATE": str}, keep_default_na=False)
n_rows = len(raw)
raw["date"] = pd.to_datetime(raw["observation_date"], format="%Y-%m-%d")
raw["y"] = pd.to_numeric(raw["UNRATE"].replace("", np.nan))
s = raw.set_index("date")["y"]
assert s.index.is_monotonic_increasing and not s.index.duplicated().any()
full_idx = pd.date_range(s.index[0], s.index[-1], freq="MS")
assert len(full_idx) == len(s), "non-contiguous months"
missing = [d.strftime("%Y-%m-%d") for d in s.index[s.isna()]]

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


def ar_forecast(train, steps):
    """train: series of levels (no trailing NaN, interior NaN interpolated). AR(P)+const on diffs, OLS."""
    d = train.diff().dropna().values
    Y = d[P:]
    X = np.column_stack([np.ones(len(Y))] + [d[P - k:len(d) - k] for k in range(1, P + 1)])
    beta, *_ = np.linalg.lstsq(X, Y, rcond=None)
    hist = list(d[-P:])
    level = train.values[-1]
    path = []
    for _ in range(steps):
        dn = beta[0] + sum(beta[k] * hist[-k] for k in range(1, P + 1))
        hist.append(dn)
        level += dn
        path.append(level)
    return np.array(path), len(train)


rows = []
n_last = None
for t in origins:
    avail = s[s.index <= t]
    last_obs_date = avail.last_valid_index()
    train = avail.loc[:last_obs_date]  # drop trailing missing
    # interior gaps: linear interpolation using only data <= last_obs_date (<= t)
    train = train.interpolate(method="linear", limit_area="inside")
    naive = train.iloc[-1]
    gap = (t.year - last_obs_date.year) * 12 + (t.month - last_obs_date.month)
    path, n_used = ar_forecast(train, gap + max(HORIZONS))
    n_last = n_used
    for h in HORIZONS:
        target = t + pd.DateOffset(months=h)
        actual = s.get(target, np.nan)
        for model, fc in (("naive", naive), (MODEL, path[gap + h - 1])):
            rows.append({"origin_date": t.strftime("%Y-%m-%d"), "target_date": target.strftime("%Y-%m-%d"),
                         "horizon": h, "forecast": round(float(fc), 6),
                         "actual": None if pd.isna(actual) else float(actual), "model": model})

fc = pd.DataFrame(rows)
fc.to_csv("out/forecasts.csv", index=False)

res = {"task_id": "F3-dev-01", "model_name": MODEL,
       "model_description": ("AR(2) with intercept on first differences of the unemployment rate, estimated by OLS on an "
                             "expanding window starting 1948-03 (first usable difference lag) up to each origin, interior "
                             "missing months linearly interpolated within the window, iterated multi-step forecasts "
                             "cumulated back to levels."),
       "n_origins": len(origins), "first_origin": origins[0].strftime("%Y-%m-%d"),
       "last_origin": origins[-1].strftime("%Y-%m-%d")}
for h in HORIZONS:
    ev = fc[(fc.horizon == h) & fc.actual.notna()]
    res[f"n_evaluated_h{h}"] = int((ev.model == "naive").sum())
for h in HORIZONS:
    for m, tag in (("naive", "naive"), (MODEL, "model")):
        ev = fc[(fc.horizon == h) & (fc.model == m) & fc.actual.notna()]
        e = ev.forecast - ev.actual
        res[f"rmse_{tag}_h{h}"] = float(np.sqrt(np.mean(e ** 2)))
        res[f"mae_{tag}_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}"]

# Diebold-Mariano style check (squared loss, HAC variance with h-1 lags), reported in report only
dm = {}
for h in HORIZONS:
    a = fc[(fc.horizon == h) & (fc.model == "naive") & fc.actual.notna()].set_index("origin_date")
    b = fc[(fc.horizon == h) & (fc.model == MODEL) & fc.actual.notna()].set_index("origin_date")
    dl = ((b.forecast - b.actual) ** 2 - (a.forecast - a.actual) ** 2).values
    n = len(dl); dbar = dl.mean(); dc = dl - dbar
    v = dc @ dc / n
    for k in range(1, h):
        v += 2 * (1 - k / h) * (dc[k:] @ dc[:-k]) / n
    stat = dbar / np.sqrt(v / n)
    from scipy import stats
    dm[f"h{h}"] = {"mean_loss_diff_model_minus_naive": float(dbar), "dm_stat": float(stat),
                   "p_two_sided_t": float(2 * stats.t.sf(abs(stat), n - 1)), "n": n}

res.update({"feasibility": "complete", "feasibility_reason": None,
            "data_issues": [{"issue": f"Missing observation(s) (empty value) at {', '.join(missing)}",
                             "handling": ("Targets at missing months kept in ledger with empty actual and excluded from "
                                          "error statistics (affects origin 2025-09 h=1 and origin 2025-07 h=3); at origin "
                                          "2025-10 the naive forecast uses 2025-09 and the AR forecast iterates from 2025-09; "
                                          "for later origins the interior gap is linearly interpolated inside the training "
                                          "window (uses only data <= origin) for AR estimation.")}],
            "n_input_rows": {"UNRATE.csv": n_rows}, "n_analysed": int(n_last), "random_seed": None})
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
with open("out/dm_tests.json", "w") as f:
    json.dump(dm, f, indent=2)
print(json.dumps(res, indent=2)); print(json.dumps(dm, indent=2))