"""F3-conf-04: seasonal AR (iterated) vs seasonal-naive forecasts of Canadian quarterly CPI inflation.
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, os
import numpy as np, pandas as pd
from scipy import stats

OUT = "out"
os.makedirs(OUT, exist_ok=True)
q = pd.read_csv("data/cpi_quarterly.csv", parse_dates=["quarter"])
m = pd.read_csv("data/cpi_monthly.csv")
n_q_raw, n_m_raw = len(q), len(m)
q = q.sort_values("quarter").reset_index(drop=True)

data_issues = []
# consistency checks against monthly file
mm = m.copy(); mm["q"] = pd.to_datetime(mm["month"]).dt.to_period("Q")
agg = mm.groupby("q")["cpi"].agg(["mean", "size"])
assert (agg["size"] == 3).all() and len(agg) == n_q_raw
assert np.abs(agg["mean"].values - q["cpi_avg"].values).max() < 1e-5
recomp = 400 * np.diff(np.log(q["cpi_avg"].values))
assert np.abs(recomp - q["infl_annualised"].values[1:]).max() < 1e-4
assert (q["quarter"].diff().dropna().dt.days.between(89, 92)).all()
data_issues.append({"issue": "infl_annualised is empty for 1995Q1 (no previous quarter)",
                    "handling": "row excluded from all estimation samples (exclusion ledger step 1)"})

excl = []
df = q.dropna(subset=["infl_annualised"]).reset_index(drop=True)
excl.append(("1", "drop rows with missing infl_annualised", n_q_raw, len(df),
             "first difference undefined for 1995Q1"))
y_all = df["infl_annualised"].values
dates = df["quarter"]
qtr = dates.dt.quarter.values  # 1..4
pos = {d: i for i, d in enumerate(dates)}

PMAX = 4
H = [1, 2, 4]
origins = pd.date_range("2015-10-01", "2025-07-01", freq="QS")
assert len(origins) == 40


def design(y, qq, p, start):
    """rows t = start..len(y)-1; regressors: 4 seasonal dummies + p lags."""
    T = len(y)
    X = [np.eye(4)[qq[start:T] - 1]]
    for i in range(1, p + 1):
        X.append(y[start - i:T - i][:, None])
    return np.hstack(X), y[start:T]


def fit_select(y, qq):
    best = None
    for p in range(PMAX + 1):
        X, Y = design(y, qq, p, PMAX)  # common sample for AIC comparison
        b, *_ = np.linalg.lstsq(X, Y, rcond=None)
        rss = float(((Y - X @ b) ** 2).sum()); n = len(Y)
        aic = n * np.log(rss / n) + 2 * X.shape[1]
        if best is None or aic < best[0]:
            best = (aic, p, b, n)
    return best[1], best[2], best[3]


rows = []; sel = []
for t in origins:
    i = pos[t]
    y = y_all[: i + 1]; qq = qtr[: i + 1]          # information set: quarters <= t
    p, b, n_est = fit_select(y, qq)
    sel.append(p)
    hist = list(y); fc = {}
    for h in range(1, max(H) + 1):                  # iterated multi-step
        qt = ((qtr[i] - 1 + h) % 4) + 1
        f = b[qt - 1] + sum(b[4 + k] * hist[-1 - k] for k in range(p))
        hist.append(f); fc[h] = f
    for h in H:
        tgt = t + pd.DateOffset(months=3 * h)
        j = pos[tgt]; assert j == i + h
        actual = y_all[j]
        sn = y_all[j - 4]; assert j - 4 <= i
        rows.append((t.date().isoformat(), tgt.date().isoformat(), h, sn, actual, "snaive"))
        rows.append((t.date().isoformat(), tgt.date().isoformat(), h, fc[h], actual, "sar_aic"))
    last_n = n_est; last_p = p

fc_df = pd.DataFrame(rows, columns=["origin_date", "target_date", "horizon", "forecast", "actual", "model"])
fc_df.to_csv(f"{OUT}/forecasts.csv", index=False, float_format="%.6f")

n_last_avail = pos[origins[-1]] + 1
excl.append(("2", "last estimation sample: quarters after last origin 2025-07-01 not used",
             len(df), n_last_avail, "information-availability rule (quarters <= origin only)"))
excl.append(("3", f"drop first {PMAX} quarters to form lags (common AIC sample, PMAX={PMAX})",
             n_last_avail, last_n, "lagged regressors undefined at start of sample"))
pd.DataFrame([dict(step=s, rule=r, n_before=a, n_after=b_, n_removed=a - b_, rationale=ra)
              for s, r, a, b_, ra in excl]).to_csv(f"{OUT}/exclusions.csv", index=False)

res = {"task_id": "F3-conf-04", "model_name": "sar_aic",
       "model_description": ("Iterated AR(p) with four quarterly seasonal intercepts for annualised quarterly inflation, "
                             "p in {0..4} chosen by AIC at each origin on a common sample (first 4 obs dropped), "
                             "OLS on an expanding window from 1995Q2 to the origin, multi-step forecasts by iteration."),
       "n_origins": len(origins), "first_origin": origins[0].date().isoformat(),
       "last_origin": origins[-1].date().isoformat()}
err = {}
for mod, key in [("snaive", "snaive"), ("sar_aic", "model")]:
    for h in H:
        s = fc_df[(fc_df.model == mod) & (fc_df.horizon == h)].sort_values("origin_date")
        e = (s.actual - s.forecast).values; err[(key, h)] = e
        res[f"rmse_{key}_h{h}"] = float(np.sqrt(np.mean(e ** 2)))
        res[f"mae_{key}_h{h}"] = float(np.mean(np.abs(e)))
for h in H:
    res[f"rel_improvement_rmse_h{h}"] = 1 - res[f"rmse_model_h{h}"] / res[f"rmse_snaive_h{h}"]

# DM test at h=4: Bartlett (Newey-West) long-run variance, lag 3 (= h-1), no small-sample correction
d = err[("snaive", 4)] ** 2 - err[("model", 4)] ** 2
T = len(d); L = 3
dc = d - d.mean()
lrv = dc @ dc / T
for k in range(1, L + 1):
    lrv += 2 * (1 - k / (L + 1)) * (dc[k:] @ dc[:-k]) / T
se = float(np.sqrt(lrv / T))
dm = float(d.mean() / se)
res.update({"loss_diff_mean_h4": float(d.mean()), "loss_diff_se_h4": se, "hac_lags_h4": L,
            "hac_kernel_h4": "bartlett", "dm_stat_h4": dm,
            "dm_pvalue_h4": float(2 * stats.norm.sf(abs(dm))),
            "feasibility": "complete", "feasibility_reason": None, "data_issues": data_issues,
            "n_input_rows": {"cpi_quarterly.csv": n_q_raw, "cpi_monthly.csv": n_m_raw},
            "n_analysed": int(last_n), "random_seed": None})
# supplementary (not required): sensitivity of the DM SE to lag/kernel, and lag orders selected
sens = {}
for LL in [3, 4, 6]:
    v = dc @ dc / T
    for k in range(1, LL + 1):
        v += 2 * (1 - k / (LL + 1)) * (dc[k:] @ dc[:-k]) / T
    sens[f"bartlett_L{LL}"] = float(np.sqrt(v / T))
v = dc @ dc / T + 2 * sum((dc[k:] @ dc[:-k]) / T for k in range(1, 4))
sens["rectangular_L3"] = float(np.sqrt(v / T)) if v > 0 else None
res["supp_dm_se_sensitivity_h4"] = sens
res["supp_hln_pvalue_h4"] = float(2 * stats.t.sf(abs(dm) * np.sqrt((T + 1 - 2 * 4 + 4 * 3 / T) / T), T - 1))
res["supp_selected_p_counts"] = {str(k): int(v) for k, v in zip(*np.unique(sel, return_counts=True))}
res["supp_selected_p_last_origin"] = int(last_p)
with open(f"{OUT}/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps({k: v for k, v in res.items() if k != "data_issues"}, indent=1))