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

DATA = "data/ip_predictors.csv"
LAMBDA = 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")
FEATS = ["ip_growth", "ip_lag1", "ip_lag2", "oil_chg", "fedfunds", "permit_chg", "neworder_chg", "tcu"]

raw = pd.read_csv(DATA)
n_raw = len(raw)
sha = hashlib.sha256(open(DATA, "rb").read()).hexdigest()
d = raw.copy()
d["month"] = pd.to_datetime(d["month"])
d = d.sort_values("month").set_index("month")
assert d.index.is_unique and (d.index.to_series().diff().dropna().dt.days.between(28, 31)).all()
# Lags use only past rows (shift forward in time); next-month target kept in a separate column
d["ip_lag1"] = d["ip_growth"].shift(1)
d["ip_lag2"] = d["ip_growth"].shift(2)
d["y_next"] = d["ip_growth"].shift(-1)   # target of row s = ip_growth of s+1 (only used when s+1 <= t)

origins = pd.date_range(FIRST_ORIGIN, LAST_ORIGIN, freq="MS")
rows, n_train_last = [], None
for t in origins:
    tgt = t + pd.DateOffset(months=1)
    actual = float(d.loc[tgt, "ip_growth"])
    avail = d.loc[:t]                                   # information set: rows dated <= t
    # --- baselines ---
    naive = float(avail.loc[t, "ip_growth"])
    mean_f = float(avail.loc[MEAN_START:t, "ip_growth"].mean())
    # --- ridge ---
    tr = avail.loc[TRAIN_START:t - pd.DateOffset(months=1)]   # s+1 <= t
    X = tr[FEATS].to_numpy(float); y = tr["y_next"].to_numpy(float)
    assert np.isfinite(X).all() and np.isfinite(y).all()
    assert tr.index.max() + pd.DateOffset(months=1) <= 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 + LAMBDA * np.eye(Z.shape[1]), Z.T @ (y - ybar))
    zt = (avail.loc[t, FEATS].to_numpy(float) - mu) / sd
    ridge = float(ybar + zt @ b)
    n_train_last = len(tr)
    for m, f in (("naive", naive), ("mean", mean_f), ("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))

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

def err(m):
    e = fc.loc[fc.model == m, "actual"] - fc.loc[fc.model == m, "forecast"]
    return float(np.sqrt(np.mean(e ** 2))), float(np.mean(np.abs(e)))

def dm(base):
    """Diebold-Mariano test (squared-error loss, h=1, Harvey et al. small-sample correction).
    d = e_base^2 - e_ridge^2; positive mean => ridge more accurate."""
    from scipy import stats
    e_b = (fc.loc[fc.model == base, "actual"] - fc.loc[fc.model == base, "forecast"]).to_numpy()
    e_r = (fc.loc[fc.model == "ridge", "actual"] - fc.loc[fc.model == "ridge", "forecast"]).to_numpy()
    dd = e_b ** 2 - e_r ** 2; T = len(dd)
    stat = dd.mean() / np.sqrt(dd.var(ddof=0) / T) * np.sqrt((T - 1) / T)   # HLN factor for h=1
    return float(stat), float(2 * stats.t.sf(abs(stat), df=T - 1))

rmse_n, mae_n = err("naive"); rmse_m, mae_m = err("mean"); rmse_r, mae_r = err("ridge")

# exclusion ledger (row-level, for the training set at the last origin)
last = origins[-1]
ex = []
n0 = n_raw
n1 = int((d.index >= TRAIN_START).sum());          ex.append(("1", "month >= 1992-03-01 (training start per spec)", n0, n1))
n2 = int(((d.index >= TRAIN_START) & (d.index <= last - pd.DateOffset(months=1))).sum())
ex.append(("2", "s+1 <= last origin 2026-08-01 (target observed at origin)", n1, n2))
pd.DataFrame([dict(step=s, rule=r, n_before=a, n_after=b_, n_removed=a - b_,
                   rationale=("rows 1992-01/02 lack full lag/neworder features; spec starts at 1992-03" if s == "1"
                              else "rows dated after 2026-07 have targets after the last origin (information rule)"))
              for s, r, a, b_ in ex]).to_csv("out/exclusions.csv", index=False)
assert n2 == n_train_last

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": LAMBDA,
    "normalisation_window": "expanding",
    "n_features": len(FEATS),
    "baseline_note": "naive and mean are the task-required baselines",
    "rmse_naive": rmse_n, "mae_naive": mae_n,
    "rmse_mean": rmse_m, "mae_mean": mae_m,
    "rmse_ridge": rmse_r, "mae_ridge": mae_r,
    "rel_improvement_vs_naive": 1 - rmse_r / rmse_n,
    "rel_improvement_vs_mean": 1 - rmse_r / rmse_m,
    "dm_stat_ridge_vs_naive": dm("naive")[0], "dm_pvalue_ridge_vs_naive": dm("naive")[1],
    "dm_stat_ridge_vs_mean": dm("mean")[0], "dm_pvalue_ridge_vs_mean": dm("mean")[1],
    "dm_note": "Diebold-Mariano, squared-error loss, h=1, HLN correction, t(59) reference; positive stat => ridge more accurate",
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": [
        {"issue": "data/README.md states the first month (1992-01) of every change series is empty, but ip_growth, oil_chg and permit_chg have values at 1992-01; only neworder_chg is missing (1992-01, 1992-02).",
         "handling": "Values used as given. Only ip_growth 1992-01 enters the analysis (lag-2 feature of the s=1992-03 training row, as specified); missing neworder_chg rows precede the training window."},
        {"issue": "Extreme COVID-19 observations (ip_growth 2020-04 = -14.14, 2020-06 = +6.35; oil_chg 2020-03..05) lie inside every training window.",
         "handling": "Kept; the task fully specifies the model without outlier treatment. Noted as a limitation."},
        {"issue": "Current-vintage (revised) data; publication lags ignored.",
         "handling": "Errors computed on the data as given; no real-time validity claimed."},
    ],
    "n_input_rows": {"ip_predictors.csv": n_raw},
    "input_sha256": {"ip_predictors.csv": sha},
    "n_analysed": n_train_last,
    "random_seed": None,
}
json.dump(res, open("out/results.json", "w"), indent=2)
print(json.dumps({k: res[k] for k in res if k.startswith(("rmse", "mae", "rel", "n_"))}, indent=2))