"""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
d = pd.read_csv("data/ip_predictors.csv", parse_dates=["month"])
n_rows = len(d)
d = d.sort_values("month").set_index("month")
assert d.index.is_unique
months = d.index
issues = []
# data checks
gaps = pd.Series(months).diff().dt.days.dropna()
if not gaps.between(28, 31).all():
    issues.append({"issue": "non-consecutive months", "handling": "none needed check failed"})
na = d.isna().sum()
issues.append({"issue": f"missing values: neworder_chg {int(na['neworder_chg'])} (1992-01/02); other columns {int(na.drop('neworder_chg').sum())}",
               "handling": "outside every training window (training starts 1992-03, lags reach back to 1992-01 only for ip_growth); no imputation"})
if not np.isnan(d["ip_growth"].iloc[0]):
    issues.append({"issue": "README says the first month of each change series is empty, but ip_growth (and other change columns) for 1992-01 are populated",
                   "handling": "values kept as given; ip_growth 1992-01 is used only as the second lag of the 1992-03 training row, as the task specification implies"})
covid = d.loc["2020-03-01":"2020-07-01", "ip_growth"]
issues.append({"issue": f"extreme COVID-19 observations in 2020 (ip_growth range {covid.min():.2f} to {covid.max():.2f})",
               "handling": "kept in all training windows as specified (no outlier treatment); they influence standardisation statistics and coefficients"})

ip = d["ip_growth"]
feat = pd.DataFrame({
    "ip0": ip, "ip1": ip.shift(1), "ip2": ip.shift(2),
    "oil": d["oil_chg"], "ff": d["fedfunds"], "perm": d["permit_chg"],
    "nord": d["neworder_chg"], "tcu": d["tcu"]}, index=months)
target_next = ip.shift(-1)  # y_{s+1} indexed by 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(ip.loc[tgt])
    # naive
    f_naive = float(ip.loc[t])
    # expanding mean 1992-04..t
    f_mean = float(ip.loc["1992-04-01":t].mean())
    # ridge: training s in [1992-03, t-1]
    s_idx = months[(months >= "1992-03-01") & (months <= t - pd.offsets.MonthBegin(1))]
    X = feat.loc[s_idx].to_numpy(float)
    y = target_next.loc[s_idx].to_numpy(float)
    assert np.isfinite(X).all() and np.isfinite(y).all()
    assert (s_idx + pd.offsets.MonthBegin(1)).max() <= t  # no info after t
    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 = (feat.loc[t].to_numpy(float) - mu) / sd
    f_ridge = float(ybar + zt @ b)
    n_train_last = len(s_idx)
    for m, f in (("naive", f_naive), ("mean", f_mean), ("ridge", f_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)
os.makedirs("out", exist_ok=True)
fc.to_csv("out/forecasts.csv", index=False)

def err(m):
    e = fc.loc[fc.model == m, "forecast"] - fc.loc[fc.model == m, "actual"]
    return float(np.sqrt((e ** 2).mean())), float(e.abs().mean())
r = {k: err(k) for k in ("naive", "mean", "ridge")}
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": 8,
    "rmse_naive": r["naive"][0], "mae_naive": r["naive"][1],
    "rmse_mean": r["mean"][0], "mae_mean": r["mean"][1],
    "rmse_ridge": r["ridge"][0], "mae_ridge": r["ridge"][1],
    "rel_improvement_vs_naive": 1 - r["ridge"][0] / r["naive"][0],
    "rel_improvement_vs_mean": 1 - r["ridge"][0] / r["mean"][0],
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": issues, "n_input_rows": {"ip_predictors.csv": n_rows},
    "n_analysed": n_train_last, "random_seed": 12345,  # used only by the bootstrap in out/uncertainty.json
}
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))

# ---- Uncertainty diagnostics (written to out/uncertainty.json; not part of results.json) ----
from scipy import stats
def dm(base, model="ridge"):
    e1 = (fc.loc[fc.model == model, "forecast"].values - fc.loc[fc.model == model, "actual"].values)
    e0 = (fc.loc[fc.model == base, "forecast"].values - fc.loc[fc.model == base, "actual"].values)
    dl = e1 ** 2 - e0 ** 2
    n = len(dl)
    stat = dl.mean() / np.sqrt(dl.var(ddof=0) / n)
    stat_hln = stat * np.sqrt((n + 1 - 2 + 0) / n)  # Harvey-Leybourne-Newbold correction, h=1
    p = 2 * stats.t.sf(abs(stat_hln), df=n - 1)
    return {"mean_loss_diff_sq": float(dl.mean()), "dm_stat_hln": float(stat_hln), "p_two_sided": float(p)}
rng = np.random.default_rng(12345)
def boot_rel(base, B=5000, block=6):
    e1 = (fc.loc[fc.model == "ridge", "forecast"].values - fc.loc[fc.model == "ridge", "actual"].values)
    e0 = (fc.loc[fc.model == base, "forecast"].values - fc.loc[fc.model == base, "actual"].values)
    n = len(e1); out = []
    for _ in range(B):
        starts = rng.integers(0, n - block + 1, size=int(np.ceil(n / block)))
        idx = np.concatenate([np.arange(s, s + block) for s in starts])[:n]
        out.append(1 - np.sqrt((e1[idx] ** 2).mean()) / np.sqrt((e0[idx] ** 2).mean()))
    return [float(np.quantile(out, 0.025)), float(np.quantile(out, 0.975))]
unc = {"dm_ridge_vs_naive": dm("naive"), "dm_ridge_vs_mean": dm("mean"),
       "boot95_rel_impr_vs_naive": boot_rel("naive"), "boot95_rel_impr_vs_mean": boot_rel("mean"),
       "bootstrap": "moving block, block=6, B=5000, seed=12345"}
with open("out/uncertainty.json", "w") as f:
    json.dump(unc, f, indent=2)
print(json.dumps(unc, indent=2))