"""F3-conf-04: multi-horizon forecasts of Canadian quarterly CPI inflation.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Reads only data/, writes out/results.json and out/forecasts.csv. Deterministic (no randomness)."""
import json, hashlib, os
import numpy as np, pandas as pd
from scipy import stats

DATA = "data"; OUT = "out"
HORIZONS = [1, 2, 4]; PMAX = 4
MODEL = "sar_aic"

def sha(p): return hashlib.sha256(open(p, "rb").read()).hexdigest()

q_raw = pd.read_csv(f"{DATA}/cpi_quarterly.csv", parse_dates=["quarter"])
m_raw = pd.read_csv(f"{DATA}/cpi_monthly.csv", parse_dates=["month"])
n_rows = {"cpi_quarterly.csv": int(len(q_raw)), "cpi_monthly.csv": int(len(m_raw))}
q = q_raw.dropna(subset=["infl_annualised"]).reset_index(drop=True)   # drops 1995Q1 only
assert (q.quarter.diff().dropna().dt.days.between(89, 93)).all(), "gap in quarters"
y_all = q.infl_annualised.to_numpy(); dates = q.quarter
qtr_all = dates.dt.quarter.to_numpy()

def design(y, qtr, p, start):
    """rows t=start..T-1: const, Q2,Q3,Q4 dummies, y_{t-1..t-p}"""
    T = len(y); rows = range(start, T)
    X = [[1.0] + [float(qtr[t] == k) for k in (2, 3, 4)] + [y[t - j] for j in range(1, p + 1)] for t in rows]
    return np.array(X), y[start:]

def fit(y, qtr, p, start):
    X, Y = design(y, qtr, p, start)
    b, *_ = np.linalg.lstsq(X, Y, rcond=None)
    rss = float(((Y - X @ b) ** 2).sum()); return b, rss, len(Y)

def forecast_origin(y, qtr):
    """y, qtr: data up to and including the origin only. AIC choice of p on common sample (start=PMAX),
    refit chosen p on its full available sample, iterate forecasts to h=4."""
    aic = {}
    for p in range(PMAX + 1):
        _, rss, n = fit(y, qtr, p, PMAX)
        aic[p] = n * np.log(rss / n) + 2 * (4 + p)
    p = min(aic, key=aic.get)
    b, _, n = fit(y, qtr, p, p)
    hist = list(y); qn = int(qtr[-1]); fc = {}
    for h in range(1, max(HORIZONS) + 1):
        qn = qn % 4 + 1
        x = [1.0] + [float(qn == k) for k in (2, 3, 4)] + [hist[-j] for j in range(1, p + 1)]
        f = float(np.dot(b, x)); hist.append(f); fc[h] = f
    return fc, p, n

origins = pd.date_range("2015-10-01", "2025-07-01", freq="QS")
assert len(origins) == 40
idx = {d: i for i, d in enumerate(dates)}
rows, chosen_p = [], []
for o in origins:
    i = idx[o]
    y, qt = y_all[: i + 1], qtr_all[: i + 1]          # information set: quarters <= origin
    assert dates.iloc[i] == o and dates.iloc[: i + 1].max() <= o
    fc, p, n_last = forecast_origin(y, qt); chosen_p.append(p)
    for h in HORIZONS:
        tgt = i + h
        sn = y_all[tgt - 4]; assert tgt - 4 <= i
        for mname, f in (("snaive", sn), (MODEL, fc[h])):
            rows.append(dict(origin_date=o.strftime("%Y-%m-%d"), target_date=dates.iloc[tgt].strftime("%Y-%m-%d"),
                             horizon=h, forecast=round(float(f), 6), actual=float(y_all[tgt]), model=mname))
fdf = pd.DataFrame(rows)
fdf.to_csv(f"{OUT}/forecasts.csv", index=False)

def err(m, h):
    d = fdf[(fdf.model == m) & (fdf.horizon == h)].sort_values("origin_date")
    return (d.actual - d.forecast).to_numpy()

def lrv(d, L, kernel):
    u = d - d.mean(); T = len(u); s = u @ u / T
    for j in range(1, L + 1):
        w = 1.0 if kernel == "rectangular" else 1 - j / (L + 1)
        s += 2 * w * (u[j:] @ u[:-j]) / T
    return s

res = {"task_id": "F3-conf-04", "model_name": MODEL,
       "model_description": ("Iterated (recursive) seasonal AR(p): OLS of infl on a constant, three quarterly dummies and p own lags, "
                             "p in 0..4 re-selected by AIC at each origin on a common sample, expanding window from 1995Q2 "
                             "through the origin quarter; multi-step forecasts by iterating the fitted equation."),
       "n_origins": len(origins), "first_origin": origins[0].strftime("%Y-%m-%d"), "last_origin": origins[-1].strftime("%Y-%m-%d")}
for h in HORIZONS:
    for m, tag in (("snaive", "snaive"), (MODEL, "model")):
        e = err(m, h)
        res[f"rmse_{tag}_h{h}"] = float(np.sqrt(np.mean(e ** 2))); res[f"mae_{tag}_h{h}"] = float(np.mean(np.abs(e)))
for h in HORIZONS:
    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(MODEL, 4) ** 2; T = len(d)
L, kern = 3, "rectangular"
v = lrv(d, L, kern)
if v <= 0:                       # pre-specified fallback
    kern = "bartlett"; v = lrv(d, L, kern)
se = float(np.sqrt(v / 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=kern,
           dm_stat_h4=dm, dm_pvalue_h4=float(2 * stats.norm.sf(abs(dm))))
# sensitivity (not required keys)
sens = {}
for k, l in (("bartlett", 3), ("bartlett", 4), ("bartlett", 6), ("rectangular", 3)):
    vv = lrv(d, l, k); s = float(np.sqrt(vv / T)) if vv > 0 else None
    sens[f"{k}_L{l}"] = {"se": s, "dm": (float(d.mean() / s) if s else None),
                         "p_normal": (float(2 * stats.norm.sf(abs(d.mean() / s))) if s else None)}
h = 4; hln = np.sqrt((T + 1 - 2 * h + h * (h - 1) / T) / T)
sens["HLN_rectangular_L3_t39"] = {"dm": dm * hln, "p_t": float(2 * stats.t.sf(abs(dm * hln), T - 1))}
res["sensitivity_h4"] = sens
res["chosen_p_by_origin"] = chosen_p
res.update(feasibility="complete", feasibility_reason=None,
           data_issues=[{"issue": "infl_annualised missing for 1995Q1 (no prior quarter)", "handling": "row dropped (exclusions.csv step 1)"},
                        {"issue": "files contain quarters after the last origin (to 2026Q3)", "handling": "estimation truncated at each origin; later quarters used only as realised targets"}],
           n_input_rows=n_rows, n_analysed=int(n_last), random_seed=None,
           input_sha256={"cpi_quarterly.csv": sha(f"{DATA}/cpi_quarterly.csv"), "cpi_monthly.csv": sha(f"{DATA}/cpi_monthly.csv")})
json.dump(res, open(f"{OUT}/results.json", "w"), indent=2)
print(json.dumps({k: v for k, v in res.items() if k.startswith(("rmse", "mae", "rel", "loss", "dm", "hac"))}, indent=1))