"""F3-conf-04: AR(p)+seasonal dummies vs seasonal naive, Canadian quarterly CPI inflation.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py"""
import json
import numpy as np, pandas as pd
from scipy import stats

D = "data/"; O = "out/"
q = pd.read_csv(D + "cpi_quarterly.csv", parse_dates=["quarter"])
m = pd.read_csv(D + "cpi_monthly.csv", parse_dates=["month"])
n_rows = {"cpi_quarterly.csv": int(len(q)), "cpi_monthly.csv": int(len(m))}

# ---- data checks ----
issues = []
assert q["quarter"].is_monotonic_increasing
gaps = (q["quarter"].dt.to_period("Q").astype(int).diff().dropna() != 1).sum()
assert gaps == 0, "non-contiguous quarters"
# recompute quarterly series from monthly file to verify derivation
m["q"] = m["month"].dt.to_period("Q")
mq = m.groupby("q")["cpi"].agg(["mean", "size"])
mq = mq[mq["size"] == 3]
rec = 400 * np.log(mq["mean"]).diff()
chk = pd.Series(rec.values, index=mq.index.to_timestamp()).reindex(q["quarter"])
maxdiff = float(np.nanmax(np.abs(chk.values - q["infl_annualised"].values)))
assert maxdiff < 1e-4, maxdiff
issues.append({"issue": "infl_annualised is empty for 1995Q1 (no previous quarter)",
               "handling": "first quarter dropped; series starts 1995Q2 (126 obs)"})
y = q.set_index("quarter")["infl_annualised"].dropna()
dates = list(y.index)
yv = y.values
qtr = np.array([d.quarter for d in 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(yv, qtr, p, start):
    """rows i=start..end: y_i on const, Q2-Q4 dummies, y_{i-1..i-p}"""
    idx = np.arange(start, len(yv))
    X = [np.ones(len(idx))] + [(qtr[idx] == s).astype(float) for s in (2, 3, 4)]
    X += [yv[idx - k] for k in range(1, p + 1)]
    return np.column_stack(X), yv[idx]

def fit_forecast(ys, qs, hmax):
    # AIC lag selection p in 0..PMAX on common sample (first PMAX obs held out)
    best = None
    for p in range(PMAX + 1):
        X, Y = design(ys, qs, p, PMAX)
        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)
    p = best[1]
    X, Y = design(ys, qs, p, p)          # re-estimate on all usable obs
    b, *_ = np.linalg.lstsq(X, Y, rcond=None)
    hist = list(ys); qcur = qs[-1]; fc = []
    for h in range(1, hmax + 1):          # iterated forecasts
        qn = (qcur + h - 1) % 4 + 1
        x = [1.0] + [float(qn == s) for s in (2, 3, 4)] + [hist[-k] for k in range(1, p + 1)]
        f = float(np.dot(b, x)); fc.append(f); hist.append(f)
    return fc, p, len(Y)

rows = []; lags_used = []
for t in origins:
    i = dates.index(t)
    ys, qs = yv[: i + 1], qtr[: i + 1]      # only quarters <= t
    fc, p, nobs = fit_forecast(ys, qs, max(H))
    lags_used.append(p)
    for h in H:
        tgt = dates[i + h]
        rows.append(dict(origin_date=t.date(), target_date=tgt.date(), horizon=h,
                         forecast=fc[h - 1], actual=yv[i + h], model="ar_seas_aic"))
        rows.append(dict(origin_date=t.date(), target_date=tgt.date(), horizon=h,
                         forecast=yv[i + h - 4], actual=yv[i + h], model="snaive"))
fcst = pd.DataFrame(rows).sort_values(["model", "horizon", "origin_date"])
fcst.to_csv(O + "forecasts.csv", index=False)
n_analysed = nobs  # obs in last (2025Q3 origin) estimation regression

def err(mod, h):
    s = fcst[(fcst.model == mod) & (fcst.horizon == h)].sort_values("origin_date")
    return (s.actual - s.forecast).values

def hac_se(d, L, kernel="bartlett"):
    d = d - d.mean(); n = len(d)
    lrv = d @ d / n
    for k in range(1, L + 1):
        w = 1 - k / (L + 1) if kernel == "bartlett" else 1.0
        lrv += 2 * w * (d[k:] @ d[:-k]) / n
    return np.sqrt(lrv / n), lrv

res = {"task_id": "F3-conf-04", "model_name": "ar_seas_aic",
       "model_description": ("OLS autoregression of quarterly annualised inflation on a constant, "
           "three quarterly seasonal dummies and p own lags (p in 0..4 chosen by AIC at each origin on a "
           "common sample), estimated on an expanding window from 1995Q2 to the origin, with iterated "
           "multi-step forecasts (h=1,2,4) and known future seasonal dummies."),
       "n_origins": 40, "first_origin": str(origins[0].date()), "last_origin": str(origins[-1].date())}
for h in H:
    es, em = err("snaive", h), err("ar_seas_aic", h)
    res[f"rmse_snaive_h{h}"] = float(np.sqrt(np.mean(es ** 2)))
    res[f"mae_snaive_h{h}"] = float(np.mean(np.abs(es)))
    res[f"rmse_model_h{h}"] = float(np.sqrt(np.mean(em ** 2)))
    res[f"mae_model_h{h}"] = float(np.mean(np.abs(em)))
for h in H:
    res[f"rel_improvement_rmse_h{h}"] = 1 - res[f"rmse_model_h{h}"] / res[f"rmse_snaive_h{h}"]

d = err("snaive", 4) ** 2 - err("ar_seas_aic", 4) ** 2
L = 3
se, _ = hac_se(d, L, "bartlett")
dm = d.mean() / se
res.update({"loss_diff_mean_h4": float(d.mean()), "loss_diff_se_h4": float(se),
            "hac_lags_h4": L, "hac_kernel_h4": "bartlett", "dm_stat_h4": float(dm),
            "dm_pvalue_h4": float(2 * stats.norm.sf(abs(dm))),
            "feasibility": "complete", "feasibility_reason": None, "data_issues": issues,
            "n_input_rows": n_rows, "n_analysed": int(n_analysed), "random_seed": None})

# sensitivity (not required keys; stored separately)
sens = {"selected_lags_by_origin": lags_used,
        "derivation_check_max_abs_diff": maxdiff}
n = len(d)
for kern, lag in [("bartlett", 4), ("bartlett", 6), ("rectangular", 3)]:
    s_, lrv = hac_se(d, lag, kern)
    sens[f"{kern}_L{lag}"] = {"se": float(s_) if lrv > 0 else None,
        "dm": float(d.mean() / s_) if lrv > 0 else None,
        "p_norm": float(2 * stats.norm.sf(abs(d.mean() / s_))) if lrv > 0 else None}
hln = np.sqrt((n + 1 - 2 * 4 + 4 * 3 / n) / n)
sens["hln_bartlett_L3"] = {"dm": float(dm * hln), "p_t39": float(2 * stats.t.sf(abs(dm * hln), n - 1))}
for h in (1, 2):
    dh = err("snaive", h) ** 2 - err("ar_seas_aic", h) ** 2
    s_, _ = hac_se(dh, max(h - 1, 1))
    sens[f"dm_h{h}_bartlett_L{max(h-1,1)}"] = {"mean": float(dh.mean()), "se": float(s_),
        "dm": float(dh.mean() / s_), "p_norm": float(2 * stats.norm.sf(abs(dh.mean() / s_)))}
# excluding 2021-2023 inflation surge targets
fm = fcst[fcst.horizon == 4].copy()
fm["td"] = pd.to_datetime(fm.target_date)
keep = ~fm.td.between("2021-01-01", "2023-12-31")
for mod in ("snaive", "ar_seas_aic"):
    s = fm[keep & (fm.model == mod)]
    sens[f"rmse_h4_excl_2021_2023_{mod}"] = float(np.sqrt(((s.actual - s.forecast) ** 2).mean()))
sens["n_h4_excl_2021_2023"] = int((keep & (fm.model == "snaive")).sum())

json.dump(res, open(O + "results.json", "w"), indent=2)
json.dump(sens, open(O + "sensitivity.json", "w"), indent=2)
print(json.dumps(res, indent=2)); print(json.dumps(sens, indent=2))