"""F3-conf-03: rolling-origin ridge vs naive / expanding-mean forecasts of IP growth (h=1).
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py"""
import json, os
import numpy as np, pandas as pd

LAM = 50.0
FEATS = ["ip0", "ip1", "ip2", "oil_chg", "fedfunds", "permit_chg", "neworder_chg", "tcu"]
os.makedirs("out", exist_ok=True)
raw = pd.read_csv("data/ip_predictors.csv")
n_rows = len(raw)
d = raw.copy()
d["month"] = pd.to_datetime(d["month"])
d = d.set_index("month").sort_index()

issues = []
# checks
assert d.index.is_unique and d.index.is_monotonic_increasing
full = pd.date_range(d.index[0], d.index[-1], freq="MS")
assert len(full) == len(d), "gaps in monthly index"
na = d.isna().sum()
issues.append({"issue": "Missing values: " + ", ".join(f"{k}={int(v)}" for k, v in na[na>0].items()) + " (neworder_chg for 1992-01/02).",
               "handling": "Not used: earliest ridge training row is 1992-03, whose features (incl. lags to 1992-01 of ip_growth) are complete."})
if pd.notna(d["ip_growth"].iloc[0]) or pd.notna(d["oil_chg"].iloc[0]):
    issues.append({"issue": "README says the first month (1992-01) of each change series is empty, but ip_growth, oil_chg and permit_chg have values for 1992-01 (presumably computed from 1991-12 levels).",
                   "handling": "Values used as given; ip_growth 1992-01 enters only as the second lag of the 1992-03 training row; the expanding-mean baseline starts at 1992-04 as specified."})
mn = d["ip_growth"].idxmin()
issues.append({"issue": f"Extreme COVID-19 observations in ip_growth (min {d['ip_growth'].min():.3f} at {mn:%Y-%m}, max {d['ip_growth'].max():.3f} at {d['ip_growth'].idxmax():%Y-%m}) lie in every training window.",
               "handling": "Kept as given (task specifies the estimator exactly); noted as a limitation."})

X = pd.DataFrame(index=d.index)
X["ip0"] = d["ip_growth"]; X["ip1"] = d["ip_growth"].shift(1); X["ip2"] = d["ip_growth"].shift(2)
for c in ["oil_chg", "fedfunds", "permit_chg", "neworder_chg", "tcu"]:
    X[c] = d[c]
y_next = d["ip_growth"].shift(-1)  # y_{s+1} aligned to s

origins = pd.date_range("2021-09-01", "2026-08-01", freq="MS")
rows = []
n_train_last = None
for t in origins:
    tgt = t + pd.offsets.MonthBegin(1)
    actual = float(d.loc[tgt, "ip_growth"])
    # baselines
    naive = float(d.loc[t, "ip_growth"])
    mean_ = float(d.loc[pd.Timestamp("1992-04-01"):t, "ip_growth"].mean())
    # ridge training rows: 1992-03 <= s <= t-1 (so s+1 <= t)
    s_end = t - pd.offsets.MonthBegin(1)
    Ztr = X.loc[pd.Timestamp("1992-03-01"):s_end, FEATS].to_numpy(float)
    ytr = y_next.loc[pd.Timestamp("1992-03-01"):s_end].to_numpy(float)
    assert not np.isnan(Ztr).any() and not np.isnan(ytr).any()
    mu = Ztr.mean(0); sd = Ztr.std(0, ddof=1); ybar = ytr.mean()
    Z = (Ztr - mu) / sd
    b = np.linalg.solve(Z.T @ Z + LAM * np.eye(Z.shape[1]), Z.T @ (ytr - ybar))
    zt = X.loc[t, FEATS].to_numpy(float)
    ridge = float(ybar + ((zt - mu) / sd) @ b)
    n_train_last = len(ytr)
    for m, f in [("naive", naive), ("mean", mean_), ("ridge", ridge)]:
        rows.append({"origin_date": t.strftime("%Y-%m-%d"), "target_date": tgt.strftime("%Y-%m-%d"),
                     "horizon": 1, "forecast": f, "actual": actual, "model": m})
fc = pd.DataFrame(rows)
fc.to_csv("out/forecasts.csv", index=False)

met = {}
for m in ["naive", "mean", "ridge"]:
    e = fc[fc.model == m]; err = e.forecast - e.actual
    met[f"rmse_{m}"] = float(np.sqrt((err**2).mean())); met[f"mae_{m}"] = float(err.abs().mean())

res = {"task_id": "F3-conf-03", "n_origins": len(origins),
       "first_origin": origins[0].strftime("%Y-%m-%d"), "last_origin": origins[-1].strftime("%Y-%m-%d"),
       "lambda": LAM, "normalisation_window": "expanding", "n_features": len(FEATS), **met,
       "rel_improvement_vs_naive": 1 - met["rmse_ridge"] / met["rmse_naive"],
       "rel_improvement_vs_mean": 1 - met["rmse_ridge"] / met["rmse_mean"],
       "feasibility": "complete", "feasibility_reason": None, "data_issues": issues,
       "n_input_rows": {"ip_predictors.csv": n_rows}, "n_analysed": n_train_last, "random_seed": 20261030}  # used only for the supplementary bootstrap (out/uncertainty.json)
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))

# Supplementary uncertainty: Diebold-Mariano (squared loss, h=1, Newey-West lag 0..3 shown with
# Harvey-Leybourne-Newbold small-sample correction) and a moving-block bootstrap of RMSE ratios.
from scipy import stats
unc = {}
e = {m: (fc[fc.model == m].forecast - fc[fc.model == m].actual).to_numpy() for m in ["naive", "mean", "ridge"]}
T = len(origins)
for base in ["naive", "mean"]:
    dl = e["ridge"]**2 - e[base]**2
    dbar = dl.mean()
    out = {}
    for L in [0, 3]:
        g = [np.mean((dl[k:] - dbar) * (dl[:T-k] - dbar)) for k in range(L + 1)]
        v = g[0] + 2 * sum((1 - k / (L + 1)) * g[k] for k in range(1, L + 1))
        dm = dbar / np.sqrt(v / T) * np.sqrt((T + 1 - 2 + 0) / T)  # HLN factor for h=1
        out[f"nw_lag{L}"] = {"dm_stat": float(dm), "p_two_sided": float(2 * stats.t.sf(abs(dm), T - 1))}
    rng = np.random.default_rng(20261030)
    B, bl = 5000, 6
    ratios = []
    for _ in range(B):
        starts = rng.integers(0, T - bl + 1, size=int(np.ceil(T / bl)))
        idx = np.concatenate([np.arange(s0, s0 + bl) for s0 in starts])[:T]
        ratios.append(1 - np.sqrt(np.mean(e["ridge"][idx]**2)) / np.sqrt(np.mean(e[base][idx]**2)))
    out["mean_loss_diff_ridge_minus_base"] = float(dbar)
    out["block_bootstrap_rel_improvement_90ci"] = [float(np.percentile(ratios, 5)), float(np.percentile(ratios, 95))]
    out["bootstrap"] = {"B": B, "block_length": bl, "seed": 20261030}
    unc[f"ridge_vs_{base}"] = out
with open("out/uncertainty.json", "w") as f:
    json.dump(unc, f, indent=2)
print(json.dumps(unc, indent=2))