"""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 data/ only; writes out/results.json, out/forecasts.csv, out/exclusions.csv. Deterministic (no randomness)."""
import json
import os
import numpy as np
import pandas as pd
from scipy import stats

DATA = "data/ip_predictors.csv"
OUT = "out"
LAMBDA = 50.0
TRAIN_START = pd.Timestamp("1992-03-01")
MEAN_START = pd.Timestamp("1992-04-01")
ORIGINS = pd.date_range("2021-09-01", "2026-08-01", freq="MS")
FEATURES = ["ip_l0", "ip_l1", "ip_l2", "oil_chg", "fedfunds", "permit_chg", "neworder_chg", "tcu"]

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

# Feature vector z_s (all values dated <= s) and the h=1 target y_{s+1}
feat = pd.DataFrame(index=df.index)
feat["ip_l0"] = df["ip_growth"]
feat["ip_l1"] = df["ip_growth"].shift(1)
feat["ip_l2"] = df["ip_growth"].shift(2)
for c in ["oil_chg", "fedfunds", "permit_chg", "neworder_chg", "tcu"]:
    feat[c] = df[c]
target_next = df["ip_growth"].shift(-1)  # y_{s+1}, used only as a training label for s+1 <= t


def ridge_forecast(t):
    """Forecast of ip_growth(t+1) using only data dated <= t."""
    train_idx = feat.index[(feat.index >= TRAIN_START) & (feat.index + pd.DateOffset(months=1) <= t)]
    X = feat.loc[train_idx, FEATURES].to_numpy(float)
    y = target_next.loc[train_idx].to_numpy(float)
    # sanity: every label date is <= t
    assert (train_idx + pd.DateOffset(months=1)).max() <= t
    assert np.isfinite(X).all() and np.isfinite(y).all()
    mu = X.mean(axis=0)
    sd = X.std(axis=0, ddof=1)
    Z = (X - mu) / sd
    ybar = y.mean()
    b = np.linalg.solve(Z.T @ Z + LAMBDA * np.eye(Z.shape[1]), Z.T @ (y - ybar))
    zt = (feat.loc[t, FEATURES].to_numpy(float) - mu) / sd
    return float(ybar + zt @ b), len(train_idx), b


rows = []
n_train = {}
coefs = {}
for t in ORIGINS:
    tgt = t + pd.DateOffset(months=1)
    actual = float(df.loc[tgt, "ip_growth"])
    naive = float(df.loc[t, "ip_growth"])
    mean_fc = float(df.loc[MEAN_START:t, "ip_growth"].mean())
    ridge_fc, ntr, b = ridge_forecast(t)
    n_train[t] = ntr
    coefs[t] = b
    for m, f in [("naive", naive), ("mean", mean_fc), ("ridge", ridge_fc)]:
        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)
os.makedirs(OUT, exist_ok=True)
fc.to_csv(os.path.join(OUT, "forecasts.csv"), index=False, float_format="%.10f")

err = {m: (g["forecast"] - g["actual"]).to_numpy() for m, g in fc.groupby("model", sort=False)}
rmse = {m: float(np.sqrt(np.mean(e ** 2))) for m, e in err.items()}
mae = {m: float(np.mean(np.abs(e))) for m, e in err.items()}


def dm_test(e_model, e_base):
    """Diebold-Mariano (h=1, squared loss) with Harvey-Leybourne-Newbold small-sample correction.
    d = e_model^2 - e_base^2; negative mean favours the model."""
    d = e_model ** 2 - e_base ** 2
    n = len(d)
    stat = d.mean() / np.sqrt(d.var(ddof=0) / n)
    stat *= np.sqrt((n + 1 - 2 + 0) / n)  # HLN factor for h=1: sqrt((n+1-2h+h(h-1)/n)/n)
    p = 2 * stats.t.sf(abs(stat), df=n - 1)
    return float(stat), float(p)


dm_naive = dm_test(err["ridge"], err["naive"])
dm_mean = dm_test(err["ridge"], err["mean"])

last = ORIGINS[-1]
n_analysed = int(n_train[last])

# Exclusion ledger (rows of the input file -> training rows at the last origin)
n_after_start = int((feat.index >= TRAIN_START).sum())
excl = pd.DataFrame([
    dict(step=1, rule="month < 1992-03-01 (before training start; ip lags / neworder_chg incomplete)",
         n_before=n_raw, n_after=n_after_start, n_removed=n_raw - n_after_start,
         rationale="Task specifies training rows start at s=1992-03-01; neworder_chg missing for 1992-01/02"),
    dict(step=2, rule="s + 1 > last origin 2026-08-01 (target not available at last origin)",
         n_before=n_after_start, n_after=n_analysed, n_removed=n_after_start - n_analysed,
         rationale="Information rule: label y_{s+1} must be dated <= origin t"),
])
excl.to_csv(os.path.join(OUT, "exclusions.csv"), index=False)
assert n_raw - excl["n_removed"].sum() == n_analysed

results = {
    "task_id": "F3-conf-03",
    "n_origins": int(len(ORIGINS)),
    "first_origin": ORIGINS[0].strftime("%Y-%m-%d"),
    "last_origin": last.strftime("%Y-%m-%d"),
    "lambda": LAMBDA,
    "normalisation_window": "expanding",
    "n_features": len(FEATURES),
    "rmse_naive": rmse["naive"], "mae_naive": mae["naive"],
    "rmse_mean": rmse["mean"], "mae_mean": mae["mean"],
    "rmse_ridge": rmse["ridge"], "mae_ridge": mae["ridge"],
    "rel_improvement_vs_naive": 1 - rmse["ridge"] / rmse["naive"],
    "rel_improvement_vs_mean": 1 - rmse["ridge"] / rmse["mean"],
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": [
        {"issue": "neworder_chg missing for 1992-01 and 1992-02",
         "handling": "Not needed: training rows start at s=1992-03 as specified; no imputation."},
        {"issue": "README states the 1992-01 value of every change series is empty, but ip_growth, oil_chg and permit_chg are populated for 1992-01",
         "handling": "Values used as given (ip_growth 1992-01 enters only as the 2nd lag of the first training row); documentation discrepancy noted."},
        {"issue": "Extreme COVID-era observations (ip_growth 2020-04 = -14.1%, 2020-06; oil_chg 2020-03..05) lie in every training window",
         "handling": "Retained: the model is fully specified by the task; no outlier treatment. Discussed as a limitation."},
        {"issue": "Current-vintage (revised) data and ignored publication lags",
         "handling": "Results interpreted as pseudo-out-of-sample on revised data, not real-time accuracy."},
    ],
    "n_input_rows": {"ip_predictors.csv": int(n_raw)},
    "n_analysed": n_analysed,
    "random_seed": None,
    # supplementary keys
    "n_train_first_origin": int(n_train[ORIGINS[0]]),
    "n_train_last_origin": n_analysed,
    "first_target": (ORIGINS[0] + pd.DateOffset(months=1)).strftime("%Y-%m-%d"),
    "last_target": (last + pd.DateOffset(months=1)).strftime("%Y-%m-%d"),
    "rel_improvement_mae_vs_naive": 1 - mae["ridge"] / mae["naive"],
    "rel_improvement_mae_vs_mean": 1 - mae["ridge"] / mae["mean"],
    "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],
    "ridge_coef_last_origin": {k: float(v) for k, v in zip(FEATURES, coefs[last])},
}
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(results, f, indent=2)
print(json.dumps({k: results[k] for k in ["rmse_naive", "rmse_mean", "rmse_ridge", "mae_naive", "mae_mean",
                                          "mae_ridge", "rel_improvement_vs_naive", "rel_improvement_vs_mean",
                                          "dm_pvalue_ridge_vs_naive", "dm_pvalue_ridge_vs_mean", "n_analysed"]}, indent=1))