"""F3-conf-03: rolling-origin ridge vs naive / expanding mean, h=1. Run from workspace root."""
import json, numpy as np, pandas as pd
DATA = "data/ip_predictors.csv"
LAM = 50.0
d = pd.read_csv(DATA, parse_dates=["month"])
n_rows = len(d)
d = d.sort_values("month").set_index("month")
assert d.index.to_series().diff().dropna().dt.days.between(28, 31).all()  # contiguous monthly
issues = []
na = d.isna().sum()
issues.append({"issue": f"Missing values: { {k: int(v) for k, v in na[na>0].items()} } (neworder_chg empty for 1992-01/02). "
               "README says the first month of each change series is empty, but ip_growth/oil_chg/permit_chg "
               "for 1992-01 are populated in the file.",
               "handling": "No imputation needed: ridge training starts at s=1992-03 (lags reach 1992-01 for ip_growth only), "
                           "expanding mean starts 1992-04 as specified; populated 1992-01 values used as given."})
issues.append({"issue": "Large COVID-19 outliers in 2020 (ip_growth min %.2f, max %.2f)." % (d.ip_growth.min(), d.ip_growth.max()),
               "handling": "Kept as given; no winsorisation (not in specification)."})
ip = d["ip_growth"]
F = pd.DataFrame({"ip0": ip, "ip1": ip.shift(1), "ip2": ip.shift(2), "oil": d.oil_chg, "ff": d.fedfunds,
                  "perm": d.permit_chg, "no": d.neworder_chg, "tcu": d.tcu})
Y = ip.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.DateOffset(months=1)
    actual = float(ip.loc[tgt])
    rows.append((t, tgt, float(ip.loc[t]), actual, "naive"))
    rows.append((t, tgt, float(ip.loc["1992-04-01":t].mean()), actual, "mean"))
    idx = F.index[(F.index >= "1992-03-01") & (F.index + pd.DateOffset(months=1) <= t)]
    X = F.loc[idx].to_numpy(float); y = Y.loc[idx].to_numpy(float)
    assert not np.isnan(X).any() and not np.isnan(y).any()
    mu = X.mean(0); sd = X.std(0, ddof=1); ybar = y.mean()
    Z = (X - mu) / sd
    b = np.linalg.solve(Z.T @ Z + LAM * np.eye(Z.shape[1]), Z.T @ (y - ybar))
    zt = (F.loc[t].to_numpy(float) - mu) / sd
    rows.append((t, tgt, float(ybar + zt @ b), actual, "ridge"))
    n_train_last = len(idx)
fc = pd.DataFrame(rows, columns=["origin_date", "target_date", "forecast", "actual", "model"])
fc.insert(2, "horizon", 1)
fc["origin_date"] = fc.origin_date.dt.strftime("%Y-%m-%d"); fc["target_date"] = fc.target_date.dt.strftime("%Y-%m-%d")
fc.to_csv("out/forecasts.csv", index=False)
err = {}
for m in ["naive", "mean", "ridge"]:
    e = fc[fc.model == m].eval("forecast-actual").to_numpy()
    err[m] = e
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": F.shape[1]}
for m, e in err.items():
    res[f"rmse_{m}"] = float(np.sqrt(np.mean(e**2))); res[f"mae_{m}"] = float(np.mean(np.abs(e)))
res["rel_improvement_vs_naive"] = 1 - res["rmse_ridge"] / res["rmse_naive"]
res["rel_improvement_vs_mean"] = 1 - res["rmse_ridge"] / res["rmse_mean"]
# Supplementary: Diebold-Mariano (squared loss, h=1, HLN small-sample correction)
from scipy import stats
def dm(e1, e2):
    dd = e1**2 - e2**2; n = len(dd); s = dd.mean() / np.sqrt(dd.var(ddof=0) / n)
    s *= np.sqrt((n - 1) / n); return float(s), float(2 * stats.t.sf(abs(s), n - 1))
supp = {"dm_ridge_vs_naive": dm(err["ridge"], err["naive"]), "dm_ridge_vs_mean": dm(err["ridge"], err["mean"])}
# bootstrap CI (iid over origins) for relative improvements
rng = np.random.default_rng(20261030); B = 5000; n = len(origins)
bs = {"naive": [], "mean": []}
for _ in range(B):
    i = rng.integers(0, n, n); r = np.sqrt(np.mean(err["ridge"][i]**2))
    for m in bs: bs[m].append(1 - r / np.sqrt(np.mean(err[m][i]**2)))
supp["boot95_rel_impr_vs_naive"] = [float(np.percentile(bs["naive"], q)) for q in (2.5, 97.5)]
supp["boot95_rel_impr_vs_mean"] = [float(np.percentile(bs["mean"], q)) for q in (2.5, 97.5)]
res.update({"feasibility": "complete", "feasibility_reason": None, "data_issues": issues,
            "n_input_rows": {"ip_predictors.csv": n_rows}, "n_analysed": n_train_last, "random_seed": 20261030,
            "supplementary": supp})
json.dump(res, open("out/results.json", "w"), indent=2)
print(json.dumps(res, indent=2))