"""F3-conf-04: multi-horizon forecasts of Canadian quarterly NSA CPI inflation.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Model: SAR-AIC = OLS AR(p) with 4 quarterly intercepts (seasonal dummies), p in {0..4}
chosen by AIC at each origin on a common sample, expanding window, iterated multi-step forecasts.
No randomness is used.
"""
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__)))
DATA, OUT = os.path.join(ROOT, "data"), os.path.join(ROOT, "out")
MODEL = "sar_aic"
PMAX, HS, HAC_LAG, KERNEL = 4, (1, 2, 4), 4, "bartlett"

q_raw = pd.read_csv(os.path.join(DATA, "cpi_quarterly.csv"), parse_dates=["quarter"])
m_raw = pd.read_csv(os.path.join(DATA, "cpi_monthly.csv"))
n_q, n_m = len(q_raw), len(m_raw)

excl = []
q = q_raw.dropna(subset=["infl_annualised"]).sort_values("quarter").reset_index(drop=True)
excl.append(dict(step=1, rule="drop rows with missing infl_annualised (1995Q1, no previous quarter)",
                 n_before=n_q, n_after=len(q), n_removed=n_q - len(q),
                 rationale="target undefined for the first quarter"))
pd.DataFrame(excl).to_csv(os.path.join(OUT, "exclusions.csv"), index=False)

y = q["infl_annualised"].to_numpy()
dates = q["quarter"]
qtr = dates.dt.quarter.to_numpy()
idx = {d: i for i, d in enumerate(dates)}
origins = pd.date_range("2015-10-01", "2025-07-01", freq="QS-JAN")
assert len(origins) == 40


def fit_sar(yw, qw):
    """Fit AR(p)+seasonal dummies for p=0..PMAX on common sample; return best by AIC."""
    T = len(yw)
    best = None
    for p in range(PMAX + 1):
        rows = np.arange(PMAX, T)
        D = (qw[rows][:, None] == np.arange(1, 5)[None, :]).astype(float)
        L = np.column_stack([yw[rows - k] for k in range(1, p + 1)]) if p else np.empty((len(rows), 0))
        X = np.hstack([D, L]); Y = yw[rows]
        b, *_ = np.linalg.lstsq(X, Y, rcond=None)
        rss = np.sum((Y - X @ b) ** 2); 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


def iterate(yw, qw, p, b, hmax):
    hist = list(yw); qs = list(qw); out = []
    for h in range(1, hmax + 1):
        qn = qs[-1] % 4 + 1
        f = b[qn - 1] + sum(b[4 + k - 1] * hist[-k] for k in range(1, p + 1))
        out.append(f); hist.append(f); qs.append(qn)
    return out


rows, chosen_p, n_last = [], [], None
for t in origins:
    i = idx[t]
    yw, qw = y[: i + 1], qtr[: i + 1]           # information set: quarters <= t
    assert dates.iloc[i] == t and dates.iloc[: i + 1].max() <= t
    aic, p, b, n = fit_sar(yw, qw)
    chosen_p.append(p); n_last = n
    fc = iterate(yw, qw, p, b, max(HS))
    for h in HS:
        tgt = i + h
        assert tgt - 4 <= i
        common = dict(origin_date=t.strftime("%Y-%m-%d"), target_date=dates.iloc[tgt].strftime("%Y-%m-%d"),
                      horizon=h, actual=y[tgt])
        rows.append({**common, "forecast": y[tgt - 4], "model": "snaive"})
        rows.append({**common, "forecast": fc[h - 1], "model": MODEL})

F = pd.DataFrame(rows)[["origin_date", "target_date", "horizon", "forecast", "actual", "model"]]
F = F.sort_values(["model", "horizon", "origin_date"]).reset_index(drop=True)
F.to_csv(os.path.join(OUT, "forecasts.csv"), index=False, float_format="%.6f")
F["e"] = F["actual"] - F["forecast"]


def lrv(d, L, kernel):
    d = d - d.mean(); T = len(d)
    g = lambda j: np.sum(d[j:] * d[: T - j]) / T
    w = (lambda j: 1 - j / (L + 1)) if kernel == "bartlett" else (lambda j: 1.0)
    return g(0) + 2 * sum(w(j) * g(j) for j in range(1, L + 1))


res = dict(task_id="F3-conf-04", model_name=MODEL,
           model_description=("OLS regression of quarterly annualised inflation on four quarterly seasonal intercepts "
                              "plus p own lags, p in {0,...,4} selected by AIC at each origin on a common sample "
                              "(first 4 obs held out for lags), expanding window from 1995Q2 to the origin, "
                              "iterated (recursive) multi-step forecasts for h=1,2,4."),
           n_origins=len(origins), first_origin=origins[0].strftime("%Y-%m-%d"),
           last_origin=origins[-1].strftime("%Y-%m-%d"))
for h in HS:
    for mname, key in (("snaive", "snaive"), (MODEL, "model")):
        e = F[(F.model == mname) & (F.horizon == h)]["e"].to_numpy()
        assert len(e) == 40
        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 HS:
    res[f"rel_improvement_rmse_h{h}"] = 1 - res[f"rmse_model_h{h}"] / res[f"rmse_snaive_h{h}"]

e_s = F[(F.model == "snaive") & (F.horizon == 4)].sort_values("origin_date")["e"].to_numpy()
e_m = F[(F.model == MODEL) & (F.horizon == 4)].sort_values("origin_date")["e"].to_numpy()
d = e_s ** 2 - e_m ** 2; T = len(d)
se = float(np.sqrt(lrv(d, HAC_LAG, KERNEL) / T))
dm = float(d.mean() / se)
res.update(loss_diff_mean_h4=float(d.mean()), loss_diff_se_h4=se, hac_lags_h4=HAC_LAG, hac_kernel_h4=KERNEL,
           dm_stat_h4=dm, dm_pvalue_h4=float(2 * stats.norm.sf(abs(dm))))
# sensitivity (extra keys)
v_rect = lrv(d, 3, "rectangular")
res["sens_rect_lag3_lrv_positive"] = bool(v_rect > 0)
res["sens_rect_lag3_dm_stat_h4"] = float(d.mean() / np.sqrt(v_rect / T)) if v_rect > 0 else None
h = 4
hln = np.sqrt((T + 1 - 2 * h + h * (h - 1) / T) / T)
se_r3 = np.sqrt(v_rect / T) if v_rect > 0 else se
res["sens_hln_dm_stat_h4"] = float(hln * d.mean() / se_r3)
res["sens_hln_dm_pvalue_h4_t39"] = float(2 * stats.t.sf(abs(res["sens_hln_dm_stat_h4"]), T - 1))
for L in (3, 6, 8):
    res[f"sens_bartlett_lag{L}_dm_stat_h4"] = float(d.mean() / np.sqrt(lrv(d, L, "bartlett") / T))
res["selected_p_counts"] = {str(k): int(v) for k, v in zip(*np.unique(chosen_p, return_counts=True))}
res.update(feasibility="complete", feasibility_reason=None,
           data_issues=[{"issue": "infl_annualised missing for 1995Q1 (no previous quarter)",
                         "handling": "row dropped (exclusions.csv step 1)"},
                        {"issue": "recomputing 400*dln(cpi_avg) differs from infl_annualised by up to ~3e-6 due to 6-decimal rounding of cpi_avg",
                         "handling": "used infl_annualised as given (immaterial)"}],
           n_input_rows={"cpi_quarterly.csv": n_q, "cpi_monthly.csv": n_m},
           n_analysed=int(n_last), random_seed=None)
with open(os.path.join(OUT, "results.json"), "w") as fh:
    json.dump(res, fh, indent=2)
print(json.dumps(res, indent=2))