"""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, os
import numpy as np, pandas as pd
from scipy import stats

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
D = os.path.join(ROOT, "data"); O = os.path.join(ROOT, "out")
HORIZONS = [1, 2, 4]; PMAX = 4; HAC_LAGS = 4; KERNEL = "bartlett"
MODEL = "sar_dummy_aic"

q_raw = pd.read_csv(os.path.join(D, "cpi_quarterly.csv"), parse_dates=["quarter"])
m_raw = pd.read_csv(os.path.join(D, "cpi_monthly.csv"), parse_dates=["month"])
q = q_raw.dropna(subset=["infl_annualised"]).reset_index(drop=True)   # exclusion step 1
y_all = q["infl_annualised"].to_numpy(); dates = q["quarter"]; qtr = dates.dt.quarter.to_numpy()

# data-issue check: quarterly series reproducible from monthly file
m = m_raw.copy(); m["q"] = m["month"].dt.to_period("Q")
agg = m.groupby("q")["cpi"].agg(["mean", "count"]); agg = agg[agg["count"] == 3]
rec = 400 * np.log(agg["mean"]).diff().dropna().to_numpy()
max_dev = float(np.max(np.abs(rec - y_all)))
data_issues = [{"issue": "infl_annualised missing for 1995Q1 (first difference undefined)",
                "handling": "row excluded from modelling (exclusions.csv step 1)"},
               {"issue": "large genuine shocks: 2008Q4 (-6.0), 2020Q2 COVID, 2022Q2 (+12.1)",
                "handling": "retained; no outlier treatment, since any such treatment would require post-origin judgement"}]
assert max_dev < 1e-5, "quarterly file not reproducible from monthly file"

def design(y, s, p):
    """rows t=p..n-1: 4 quarter dummies + p own lags."""
    n = len(y); X = []
    for t in range(p, n):
        X.append([1.0 * (s[t] == k) for k in (1, 2, 3, 4)] + [y[t - i] for i in range(1, p + 1)])
    return np.array(X), y[p:]

def fit_forecast(y, s, hmax):
    """Expanding-window fit on data <= origin; AIC selects p in 0..PMAX on common sample; iterated forecasts."""
    best = None
    for p in range(PMAX + 1):
        X, Y = design(y, s, p); X, Y = X[PMAX - p:], Y[PMAX - p:]     # common sample
        b, *_ = np.linalg.lstsq(X, Y, rcond=None); r = Y - X @ b
        aic = len(Y) * np.log(r @ r / len(Y)) + 2 * X.shape[1]
        if best is None or aic < best[0]: best = (aic, p)
    p = best[1]
    X, Y = design(y, s, p); b, *_ = np.linalg.lstsq(X, Y, rcond=None)
    hist = list(y); fc = []; s_last = s[-1]
    for h in range(1, hmax + 1):
        sq = (s_last - 1 + h) % 4 + 1
        x = [1.0 * (sq == k) for k in (1, 2, 3, 4)] + [hist[-i] for i in range(1, p + 1)]
        f = float(np.dot(x, b)); fc.append(f); hist.append(f)
    return fc, p, len(Y)

origins = pd.date_range("2015-10-01", "2025-07-01", freq="QS")
rows = []; p_used = []; n_last = None
for o in origins:
    i0 = int(np.where(dates == o)[0][0])
    y, s = y_all[: i0 + 1], qtr[: i0 + 1]            # information set: quarters <= origin
    fc, p, nobs = fit_forecast(y, s, max(HORIZONS)); p_used.append(p); n_last = nobs
    for h in HORIZONS:
        tgt = i0 + h; actual = float(y_all[tgt])
        rows.append(dict(origin_date=o.date().isoformat(), target_date=dates[tgt].date().isoformat(), horizon=h,
                         forecast=float(y_all[tgt - 4]), actual=actual, model="snaive"))
        assert tgt - 4 <= i0
        rows.append(dict(origin_date=o.date().isoformat(), target_date=dates[tgt].date().isoformat(), horizon=h,
                         forecast=fc[h - 1], actual=actual, model=MODEL))
F = pd.DataFrame(rows); F.to_csv(os.path.join(O, "forecasts.csv"), index=False, float_format="%.6f")

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

res = {"task_id": "F3-conf-04", "model_name": MODEL,
       "model_description": ("Seasonal autoregression: OLS of quarterly annualised inflation on four quarter dummies and p own lags, "
                             "p in {0,...,4} chosen by AIC on a common sample at each origin, re-estimated by OLS on the expanding window "
                             "from 1995Q2 to the origin, iterated multi-step forecasts."),
       "n_origins": len(origins), "first_origin": origins[0].date().isoformat(), "last_origin": origins[-1].date().isoformat()}
E = {}
for h in HORIZONS:
    for mod, key in (("snaive", "snaive"), (MODEL, "model")):
        g = F[(F.horizon == h) & (F.model == mod)].sort_values("origin_date")
        e = (g.actual - g.forecast).to_numpy(); E[(h, key)] = 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 HORIZONS:
    res[f"rel_improvement_rmse_h{h}"] = 1 - res[f"rmse_model_h{h}"] / res[f"rmse_snaive_h{h}"]
d = E[(4, "snaive")] ** 2 - E[(4, "model")] ** 2; T = len(d)
se, _ = hac_se(d, HAC_LAGS, KERNEL)
res.update({"loss_diff_mean_h4": float(d.mean()), "loss_diff_se_h4": se, "hac_lags_h4": HAC_LAGS, "hac_kernel_h4": KERNEL,
            "dm_stat_h4": float(d.mean() / se), "dm_pvalue_h4": float(2 * stats.norm.sf(abs(d.mean() / se)))})
# sensitivity of the h=4 test (supplementary keys)
sens = {}
for L, k in ((3, "bartlett"), (3, "rectangular"), (6, "bartlett")):
    s_, lrv = hac_se(d, L, k)
    sens[f"{k}_L{L}"] = {"se": s_, "lrv": lrv, "dm_stat": (float(d.mean() / s_) if s_ else None),
                         "pvalue_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)          # Harvey-Leybourne-Newbold correction
dm_hln = res["dm_stat_h4"] * hln
sens["hln_corrected_main"] = {"dm_stat": float(dm_hln), "pvalue_t_T-1": float(2 * stats.t.sf(abs(dm_hln), T - 1))}
res["sensitivity_h4"] = sens
res["ar_lag_order_selected"] = {"min": int(min(p_used)), "max": int(max(p_used)),
                                "counts": {str(p): int(p_used.count(p)) for p in sorted(set(p_used))}}
res.update({"feasibility": "complete", "feasibility_reason": None, "data_issues": data_issues,
            "n_input_rows": {"cpi_quarterly.csv": int(len(q_raw)), "cpi_monthly.csv": int(len(m_raw))},
            "n_analysed": int(n_last), "random_seed": None})
with open(os.path.join(O, "results.json"), "w") as f: json.dump(res, f, indent=2)
print(json.dumps({k: v for k, v in res.items() if k not in ("data_issues", "model_description")}, indent=1))