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

DATA = Path("data/ip_predictors.csv"); OUT = Path("out"); OUT.mkdir(exist_ok=True)
LAM = 50.0
d = pd.read_csv(DATA, parse_dates=["month"])
n_rows = len(d)
d = d.sort_values("month").set_index("month")
issues = []
assert not d.index.duplicated().any()
assert (d.index.to_series().diff().dropna().dt.days.between(28, 31)).all(), "gaps in months"
# README says first month of each change series is empty; check
first_filled = [c for c in ["ip_growth", "oil_chg", "permit_chg"] if pd.notna(d.iloc[0][c])]
if first_filled:
    issues.append({"issue": f"README states the first month (1992-01) of each change series is empty, but {first_filled} are populated for 1992-01.",
                   "handling": "Values used as given. Only ip_growth 1992-01 enters (as the lag-2 feature of the 1992-03 training row); the task spec explicitly starts training rows at 1992-03, implying this is intended."})
na = d.isna().sum(); na = na[na > 0]
issues.append({"issue": f"Missing values: {na.to_dict()} (neworder_chg 1992-01/02, documented in README).",
               "handling": "Not used: training rows start at 1992-03 per spec, so no imputation needed."})
big = d.loc[d.ip_growth.abs() > 4, "ip_growth"]
issues.append({"issue": "Extreme ip_growth values (|x|>4; 2008-09 hurricanes/financial crisis, 2020 COVID-19) (" + ", ".join(f"{k:%Y-%m}: {v:.2f}" for k, v in big.items()) + ").",
               "handling": "Genuine observations; kept, no winsorising (spec fixes the model). They lie in every training window and influence standardisation and coefficients."})
issues.append({"issue": "Current-vintage (revised) data; publication lags ignored.",
               "handling": "Results interpreted as pseudo-out-of-sample on revised data only, not real-time."})

y = d["ip_growth"]
F = pd.DataFrame({"ip0": y, "ip1": y.shift(1), "ip2": y.shift(2), "oil": d.oil_chg,
                  "ff": d.fedfunds, "perm": d.permit_chg, "no": d.neworder_chg, "tcu": d.tcu})
origins = pd.date_range("2021-09-01", "2026-08-01", freq="MS")
months = d.index
rows = []; n_train_last = None
for t in origins:
    tgt = t + pd.DateOffset(months=1)
    actual = float(y.loc[tgt])
    rows.append((t, tgt, float(y.loc[t]), actual, "naive"))
    rows.append((t, tgt, float(y.loc["1992-04-01":t].mean()), actual, "mean"))
    # training rows s: 1992-03 <= s, s+1 <= t
    s_idx = months[(months >= "1992-03-01") & (months <= t - pd.DateOffset(months=1))]
    X = F.loc[s_idx].to_numpy()
    yt = y.reindex(s_idx + pd.DateOffset(months=1)).to_numpy()
    assert not np.isnan(X).any() and not np.isnan(yt).any()
    assert (s_idx + pd.DateOffset(months=1)).max() <= t  # no future info
    mu = X.mean(0); sd = X.std(0, ddof=1); ybar = yt.mean()
    Z = (X - mu) / sd
    b = np.linalg.solve(Z.T @ Z + LAM * np.eye(Z.shape[1]), Z.T @ (yt - ybar))
    zt = (F.loc[t].to_numpy() - mu) / sd
    rows.append((t, tgt, float(ybar + zt @ b), actual, "ridge"))
    n_train_last = len(s_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)

res = {"task_id": "F3-conf-03", "n_origins": len(origins),
       "first_origin": f"{origins[0]:%Y-%m-%d}", "last_origin": f"{origins[-1]:%Y-%m-%d}",
       "lambda": LAM, "normalisation_window": "expanding", "n_features": F.shape[1]}
err = {}
for m in ["naive", "mean", "ridge"]:
    e = (fc[fc.model == m].forecast - fc[fc.model == m].actual).to_numpy(); err[m] = e
    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"]
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": 0})

# Supplementary uncertainty: Diebold-Mariano (squared loss, h=1, HAC lag 0 -> also report Newey-West lag 4) with HLN correction
from scipy import stats
def dm(e1, e2, L):
    dd = e1**2 - e2**2; T = len(dd); dbar = dd.mean(); u = dd - dbar
    v = u @ u / T + 2 * sum((1 - k/(L+1)) * (u[k:] @ u[:-k]) / T for k in range(1, L+1))
    stat = dbar / np.sqrt(v / T); stat *= np.sqrt((T - 1) / T)  # HLN for h=1
    return float(stat), float(2 * stats.t.sf(abs(stat), T - 1))
supp = {}
for base in ["naive", "mean"]:
    for L in [0, 4]:
        s, p = dm(err[base], err["ridge"], L)
        supp[f"dm_{base}_vs_ridge_lag{L}"] = {"stat": s, "p_value": p}
rng = np.random.default_rng(0)
T = len(err["ridge"]); bl = 6; B = 5000
for base in ["naive", "mean"]:
    vals = []
    for _ in range(B):
        st = rng.integers(0, T - bl + 1, size=int(np.ceil(T / bl)))
        idx = np.concatenate([np.arange(s, s + bl) for s in st])[:T]
        vals.append(1 - np.sqrt(np.mean(err["ridge"][idx]**2)) / np.sqrt(np.mean(err[base][idx]**2)))
    supp[f"block_bootstrap_rel_impr_vs_{base}_95ci"] = [float(np.quantile(vals, .025)), float(np.quantile(vals, .975))]
json.dump(supp, open(OUT / "supplementary_uncertainty.json", "w"), indent=2)
json.dump(res, open(OUT / "results.json", "w"), indent=2)
print(json.dumps(res, indent=2)); print(json.dumps(supp, indent=2))