"""F3-conf-03: rolling-origin ridge vs naive / expanding-mean forecasts of US IP growth (h=1).
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Reads only data/; writes out/results.json, out/forecasts.csv, out/exclusions.csv. No randomness."""
import json, hashlib
from pathlib import Path
import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parent.parent
DATA = ROOT / "data" / "ip_predictors.csv"
OUT = ROOT / "out"
LAM = 50.0
FIRST_ORIGIN, LAST_ORIGIN = pd.Timestamp("2021-09-01"), pd.Timestamp("2026-08-01")
TRAIN_START = pd.Timestamp("1992-03-01")
MEAN_START = pd.Timestamp("1992-04-01")

raw = pd.read_csv(DATA, parse_dates=["month"])
n_raw = len(raw)
d = raw.sort_values("month").set_index("month")
assert not d.index.duplicated().any()
assert (d.index == pd.date_range(d.index[0], d.index[-1], freq="MS")).all(), "non-contiguous months"

# feature matrix z_s
F = pd.DataFrame({
    "ip0": d["ip_growth"], "ip1": d["ip_growth"].shift(1), "ip2": d["ip_growth"].shift(2),
    "oil": d["oil_chg"], "ff": d["fedfunds"], "permit": d["permit_chg"],
    "neworder": d["neworder_chg"], "tcu": d["tcu"]})
y_next = d["ip_growth"].shift(-1)  # target for row s is ip_{s+1}
months = d.index

origins = pd.date_range(FIRST_ORIGIN, LAST_ORIGIN, freq="MS")
rows = []
n_train_last = None
for t in origins:
    tgt = t + pd.offsets.MonthBegin(1)
    actual = float(d.loc[tgt, "ip_growth"])
    # naive
    rows.append((t, tgt, d.loc[t, "ip_growth"], actual, "naive"))
    # expanding mean 1992-04..t
    rows.append((t, tgt, d.loc[MEAN_START:t, "ip_growth"].mean(), actual, "mean"))
    # ridge: training rows s in [1992-03, t-1] (s+1 <= t)
    s_last = t - pd.offsets.MonthBegin(1)
    Xtr = F.loc[TRAIN_START:s_last].to_numpy(float)
    ytr = y_next.loc[TRAIN_START:s_last].to_numpy(float)
    assert np.isfinite(Xtr).all() and np.isfinite(ytr).all()
    assert y_next.loc[TRAIN_START:s_last].index.max() + pd.offsets.MonthBegin(1) <= t
    mu = Xtr.mean(0); sd = Xtr.std(0, ddof=1); ybar = ytr.mean()
    Z = (Xtr - mu) / sd
    b = np.linalg.solve(Z.T @ Z + LAM * np.eye(Z.shape[1]), Z.T @ (ytr - ybar))
    zt = (F.loc[t].to_numpy(float) - mu) / sd
    rows.append((t, tgt, float(ybar + zt @ b), actual, "ridge"))
    n_train_last = len(ytr)

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, float_format="%.8f")

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]}
err = {}
for m in ["naive", "mean", "ridge"]:
    e = (fc.loc[fc.model == m, "actual"] - fc.loc[fc.model == m, "forecast"]).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["feasibility"] = "complete"
res["feasibility_reason"] = None
res["data_issues"] = [
    {"issue": "neworder_chg missing for 1992-01 and 1992-02 (NEWORDER starts 1992-02).",
     "handling": "Outside every training window (training starts 1992-03 as specified); no imputation needed."},
    {"issue": "data/README.md says the first month (1992-01) of each change series is empty, but ip_growth, oil_chg and permit_chg have values in 1992-01.",
     "handling": "Values used as given; only ip_growth 1992-01 enters (as lag ip_{s-2} for s=1992-03), as implied by the task's feature definition."},
    {"issue": "Current-vintage (revised) data with publication lags ignored; training windows include the 2020 COVID swing (ip_growth minimum -14.14 in 2020-04) while evaluation-period actuals lie in [-1.38, 1.33].",
     "handling": "No adjustment or outlier treatment (not part of the specified model); results are not real-time and are interpreted accordingly."},
]
res["n_input_rows"] = {"ip_predictors.csv": n_raw}
res["n_analysed"] = n_train_last
res["random_seed"] = None
# supplementary (not required): Diebold-Mariano-type tests on squared-error loss, h=1, plain variance
from scipy import stats
supp = {}
for base in ["naive", "mean"]:
    dl = err[base] ** 2 - err["ridge"] ** 2
    dm = dl.mean() / np.sqrt(dl.var(ddof=1) / len(dl))
    supp[f"dm_stat_ridge_vs_{base}"] = float(dm)
    supp[f"dm_pvalue_ridge_vs_{base}"] = float(2 * (1 - stats.t.cdf(abs(dm), df=len(dl) - 1)))
res["supplementary"] = supp
(OUT / "results.json").write_text(json.dumps(res, indent=2))

# exclusion ledger: raw rows -> rows used as ridge training rows at last origin
last = origins[-1]
steps = [("raw", "all rows read", n_raw, n_raw)]
n1 = int((months >= TRAIN_START).sum())
steps.append(("1", "drop months before 1992-03 (no lags / neworder_chg missing)", n_raw, n1))
n2 = int(((months >= TRAIN_START) & (months <= last - pd.offsets.MonthBegin(1))).sum())
steps.append(("2", "drop months s with s+1 > last origin 2026-08 (target not available at last origin)", n1, n2))
ex = pd.DataFrame([(s, r, a, b, a - b) for s, r, a, b in steps],
                  columns=["step", "rule", "n_before", "n_after", "n_removed"])
ex["rationale"] = ["input", "task-specified training start; features need ip lags and neworder_chg",
                   "information rule: s+1 <= t at the last origin"]
ex.to_csv(OUT / "exclusions.csv", index=False)
assert n2 == n_train_last
print(json.dumps({k: v for k, v in res.items() if k.startswith(("rmse", "mae", "rel"))}, indent=2))