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

DATA = "data/UNRATE.csv"; OUT = "out"
P = 2                                   # AR lags on first differences (fixed a priori, not selected)
COVID = (pd.Timestamp("2020-04-01"), pd.Timestamp("2020-09-01"))  # differences dated here are excluded from estimation
HORIZONS = [1, 3]
MODEL = "AR2_diff"

raw = pd.read_csv(DATA, dtype={"UNRATE": "float64"})
n_rows = len(raw)
raw["observation_date"] = pd.to_datetime(raw["observation_date"])
s_full = raw.set_index("observation_date")["UNRATE"].asfreq("MS")
origins = pd.date_range("2023-07-01", "2026-06-01", freq="MS")

def design(u, exclude_covid=True):
    """Regression rows y=du_t on [1, du_{t-1},...,du_{t-P}], using observed (non-interpolated) data only."""
    du = u.diff()
    X = pd.concat({f"l{k}": du.shift(k) for k in range(1, P + 1)}, axis=1)
    df = pd.concat([du.rename("y"), X], axis=1).dropna()
    if exclude_covid:
        bad = pd.Series(False, index=df.index)
        for k in range(0, P + 1):          # drop row if y or any lag is dated inside the COVID window
            d = df.index - pd.DateOffset(months=k)
            bad |= (d >= COVID[0]) & (d <= COVID[1])
        df = df[~bad]
    return df

def ar_forecast(u_obs, origin, hmax, exclude_covid=True):
    """Fit on data <= origin; return dict h->level forecast for h=1..hmax relative to origin."""
    u = u_obs[u_obs.index <= origin]
    df = design(u, exclude_covid)
    X = np.column_stack([np.ones(len(df)), df[[f"l{k}" for k in range(1, P + 1)]].values])
    beta, *_ = np.linalg.lstsq(X, df["y"].values, rcond=None)
    # state: interpolate interior gaps (uses only data <= origin); trailing gaps are forecast
    last_obs = u.last_valid_index()
    ui = u.loc[:last_obs].interpolate(limit_area="inside")
    gap = (origin.year - last_obs.year) * 12 + origin.month - last_obs.month
    du = list(ui.diff().values[-P:])
    level = ui.iloc[-1]; path = {}
    for step in range(1, gap + hmax + 1):
        d = beta[0] + sum(beta[k] * du[-k] for k in range(1, P + 1))
        du.append(d); level = level + d
        if step > gap: path[step - gap] = level
    return path, len(df), beta

def run(exclude_covid=True):
    rows = []; n_last = None
    for t in origins:
        u = s_full[s_full.index <= t]
        naive = u.loc[u.last_valid_index()]
        path, n_used, beta = ar_forecast(s_full, t, max(HORIZONS), exclude_covid)
        n_last = n_used
        for h in HORIZONS:
            tgt = t + pd.DateOffset(months=h)
            act = s_full.get(tgt, np.nan)
            for m, f in (("naive", naive), (MODEL, path[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), n_last

def stats(fc):
    out = {}
    for h in HORIZONS:
        out[f"n_evaluated_h{h}"] = int(fc[(fc.horizon == h) & (fc.model == "naive")].actual.notna().sum())
        for m, lab in (("naive", "naive"), (MODEL, "model")):
            e = fc[(fc.horizon == h) & (fc.model == m)].dropna(subset=["actual"])
            err = e.forecast - e.actual
            out[f"rmse_{lab}_h{h}"] = float(np.sqrt(np.mean(err ** 2)))
            out[f"mae_{lab}_h{h}"] = float(np.mean(np.abs(err)))
        out[f"rel_improvement_rmse_h{h}"] = 1 - out[f"rmse_model_h{h}"] / out[f"rmse_naive_h{h}"]
    return out

def dm_test(fc, h):
    """Diebold-Mariano (squared loss) with Newey-West HAC var (h-1 lags) and Harvey-Leybourne-Newbold correction."""
    from scipy import stats as st
    a = fc[(fc.horizon == h) & (fc.model == "naive")].set_index("origin_date")
    b = fc[(fc.horizon == h) & (fc.model == MODEL)].set_index("origin_date")
    ok = a.actual.notna()
    d = ((a.forecast - a.actual) ** 2 - (b.forecast - b.actual) ** 2)[ok].values
    n = len(d); dbar = d.mean(); dc = d - dbar
    v = dc @ dc / n
    for k in range(1, h):
        v += 2 * (1 - k / h) * (dc[k:] @ dc[:-k]) / n
    dmstat = dbar / np.sqrt(v / n)
    hln = dmstat * np.sqrt((n + 1 - 2 * h + h * (h - 1) / n) / n)
    p = 2 * st.t.sf(abs(hln), df=n - 1)
    return float(hln), float(p)

def boot_rel(fc, h, B=2000, block=6, seed=12345):
    """Moving-block bootstrap (over origins) 90% interval for 1 - RMSE_model/RMSE_naive."""
    rng = np.random.default_rng(seed)
    a = fc[(fc.horizon == h) & (fc.model == "naive")].reset_index(drop=True)
    b = fc[(fc.horizon == h) & (fc.model == MODEL)].reset_index(drop=True)
    ok = a.actual.notna().values
    ea = ((a.forecast - a.actual) ** 2).values[ok]; eb = ((b.forecast - b.actual) ** 2).values[ok]
    n = len(ea); reps = []
    for _ in range(B):
        idx = np.concatenate([np.arange(s, s + block) for s in rng.integers(0, n - block + 1, size=int(np.ceil(n / block)))])[:n]
        reps.append(1 - np.sqrt(eb[idx].mean() / ea[idx].mean()))
    return [float(np.quantile(reps, 0.05)), float(np.quantile(reps, 0.95))]

if __name__ == "__main__":
    os.makedirs(OUT, exist_ok=True)
    fc, n_last = run(True)
    fc.to_csv(f"{OUT}/forecasts.csv", index=False, na_rep="")
    res = {"task_id": "F3-dev-01", "model_name": MODEL,
           "model_description": ("AR(2) with intercept on first differences of UNRATE, iterated to levels for h=1,3; "
                                 "lag order fixed a priori; re-estimated by OLS at each origin on an expanding window from 1948-03 "
                                 "(first available difference rows) to the origin, excluding regression rows whose differences fall in "
                                 "2020-04..2020-09 (COVID outliers) and rows involving the missing 2025-10 value."),
           "n_origins": len(origins), "first_origin": origins[0].strftime("%Y-%m-%d"),
           "last_origin": origins[-1].strftime("%Y-%m-%d")}
    res.update(stats(fc))
    res.update({"feasibility": "complete", "feasibility_reason": None,
                "data_issues": [{"issue": "UNRATE for 2025-10-01 is missing (empty field).",
                                 "handling": "Excluded as an evaluation target (origins 2025-09 h=1 and 2025-07 h=3 have empty actual). "
                                             "Naive at origin 2025-10 uses 2025-09. AR model: regression rows involving it are dropped; for the "
                                             "forecast state it is linearly interpolated once 2025-11 is observed (only data <= origin), and at "
                                             "origin 2025-10 the model iterates from 2025-09 (h+1 steps)."},
                                {"issue": "Extreme COVID-19 outliers in 2020 (peak 14.8 in 2020-04; monthly changes up to +10.4 pp).",
                                 "handling": "Regression rows whose differences are dated 2020-04..2020-09 excluded from AR estimation (a fixed rule "
                                             "decided from pre-2023 data); sensitivity without exclusion reported in report.md."}],
                "n_input_rows": {"UNRATE.csv": n_rows}, "n_analysed": int(n_last), "random_seed": 12345})
    json.dump(res, open(f"{OUT}/results.json", "w"), indent=2)
    # supplementary diagnostics (not required keys)
    fc2, _ = run(False)
    supp = {"sensitivity_no_covid_exclusion": stats(fc2),
            "dm_hln": {f"h{h}": dict(zip(["stat", "p_value"], dm_test(fc, h))) for h in HORIZONS},
            "block_bootstrap_90ci_rel_improvement": {f"h{h}": boot_rel(fc, h) for h in HORIZONS},
            "last_origin_coefficients": list(map(float, ar_forecast(s_full, origins[-1], 3)[2]))}
    json.dump(supp, open(f"{OUT}/supplementary.json", "w"), indent=2)
    print(json.dumps(res, indent=2)); print(json.dumps(supp, indent=2))