"""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
import pandas as pd

LAM = 50.0
SEED = None  # no randomness used
FIRST, LAST = pd.Timestamp("2021-09-01"), pd.Timestamp("2026-08-01")
TRAIN_START = pd.Timestamp("1992-03-01")
MEAN_START = pd.Timestamp("1992-04-01")
FEATS = ["ip_growth", "ip_l1", "ip_l2", "oil_chg", "fedfunds", "permit_chg", "neworder_chg", "tcu"]

os.makedirs("out", exist_ok=True)
raw = pd.read_csv("data/ip_predictors.csv", parse_dates=["month"])
n_raw = len(raw)
d = raw.sort_values("month").set_index("month")
assert d.index.is_unique
d["ip_l1"] = d["ip_growth"].shift(1)   # lags use only past values
d["ip_l2"] = d["ip_growth"].shift(2)
d["y_next"] = d["ip_growth"].shift(-1)  # target of row s = ip_growth_{s+1}

origins = pd.date_range(FIRST, LAST, freq="MS")
rows, n_train_last, excl = [], None, []
for t in origins:
    tgt = t + pd.DateOffset(months=1)
    actual = float(d.loc[tgt, "ip_growth"])
    avail = d.loc[:t]  # information set: dated <= t
    # baselines
    naive = float(avail.loc[t, "ip_growth"])
    mean = float(avail.loc[MEAN_START:t, "ip_growth"].mean())
    # ridge training rows: TRAIN_START <= s, s+1 <= t
    tr = avail.loc[TRAIN_START:t - pd.DateOffset(months=1)]
    Xr = tr[FEATS].to_numpy(float)
    # y_{s+1} with s+1 <= t, taken from the available block
    yr = avail["ip_growth"].shift(-1).loc[tr.index].to_numpy(float)
    assert np.isfinite(Xr).all() and np.isfinite(yr).all(), t
    mu, sd = Xr.mean(0), Xr.std(0, ddof=1)
    ybar = yr.mean()
    Z = (Xr - mu) / sd
    b = np.linalg.solve(Z.T @ Z + LAM * np.eye(Z.shape[1]), Z.T @ (yr - ybar))
    zt = avail.loc[t, FEATS].to_numpy(float)
    assert np.isfinite(zt).all()
    ridge = float(ybar + ((zt - mu) / sd) @ b)
    for m, f in (("naive", naive), ("mean", mean), ("ridge", ridge)):
        rows.append(dict(origin_date=t.strftime("%Y-%m-%d"), target_date=tgt.strftime("%Y-%m-%d"),
                         horizon=1, forecast=f, actual=actual, model=m))
    n_train_last = len(tr)

fc = pd.DataFrame(rows)
fc.to_csv("out/forecasts.csv", index=False, float_format="%.10f")

def metrics(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())

res = {"task_id": "F3-conf-03", "n_origins": len(origins),
       "first_origin": FIRST.strftime("%Y-%m-%d"), "last_origin": LAST.strftime("%Y-%m-%d"),
       "lambda": LAM, "normalisation_window": "expanding", "n_features": len(FEATS)}
for m in ("naive", "mean", "ridge"):
    res[f"rmse_{m}"], res[f"mae_{m}"] = metrics(m)
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["baseline_best_rmse"] = min(res["rmse_naive"], res["rmse_mean"])

# Uncertainty: Diebold-Mariano tests (squared-error loss, h=1, normal approx, Harvey et al. small-sample correction)
from scipy import stats
def dm(base):
    e_r = (fc.loc[fc.model == "ridge", "forecast"].to_numpy() - fc.loc[fc.model == "ridge", "actual"].to_numpy())
    e_b = (fc.loc[fc.model == base, "forecast"].to_numpy() - fc.loc[fc.model == base, "actual"].to_numpy())
    dd = e_b ** 2 - e_r ** 2  # >0 favours ridge
    n = len(dd)
    stat = dd.mean() / np.sqrt(dd.var(ddof=0) / n)
    stat *= np.sqrt((n + 1 - 2 + 0) / n)  # HLN correction for h=1
    p = 2 * stats.t.sf(abs(stat), df=n - 1)
    return float(stat), float(p)
res["dm_stat_ridge_vs_naive"], res["dm_pvalue_ridge_vs_naive"] = dm("naive")
res["dm_stat_ridge_vs_mean"], res["dm_pvalue_ridge_vs_mean"] = dm("mean")
res.update({"feasibility": "complete", "feasibility_reason": None,
            "data_issues": [
                {"issue": "neworder_chg missing for 1992-01 and 1992-02 (series starts 1992-02).",
                 "handling": "Not used: earliest neworder_chg needed is 1992-03 (first training row)."},
                {"issue": "README states the first month (1992-01) of every change series is empty, but ip_growth, oil_chg and permit_chg are populated for 1992-01.",
                 "handling": "Values used as given; only ip_growth 1992-01 enters (lag-2 feature of training row 1992-03)."},
                {"issue": "Extreme COVID-19 observations (e.g. ip_growth 2020-04 = -14.1, oil_chg 2020-03..05) lie in every training window.",
                 "handling": "Kept, as the task prescribes the training window; no outlier rule applied."},
                {"issue": "Current-vintage data; publication lags ignored.",
                 "handling": "Errors computed on revised data as instructed; no real-time validity claimed."}],
            "n_input_rows": {"ip_predictors.csv": n_raw},
            "n_analysed": int(n_train_last), "random_seed": SEED})
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)

# exclusion ledger: raw rows -> training rows at last origin
n_after_start = int((d.index >= TRAIN_START).sum())
n_final = n_train_last
pd.DataFrame([
    dict(step=1, rule="drop months before 1992-03-01 (first training row; lag/neworder availability)",
         n_before=n_raw, n_after=n_after_start, n_removed=n_raw - n_after_start,
         rationale="TASK.md ridge spec: training rows start at s=1992-03-01"),
    dict(step=2, rule="drop months s with s+1 > last origin 2026-08-01 (target not available at origin)",
         n_before=n_after_start, n_after=n_final, n_removed=n_after_start - n_final,
         rationale="information rule: training target y_{s+1} must be dated <= t"),
]).to_csv("out/exclusions.csv", index=False)
print(json.dumps({k: res[k] for k in res if k.startswith(("rmse", "mae", "rel_", "n_", "dm_"))}, indent=1))