"""F3-conf-04: multi-horizon forecasts of Canadian quarterly CPI inflation vs seasonal naive.
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
import pandas as pd
from scipy import stats

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DATA = os.path.join(ROOT, "data"); OUT = os.path.join(ROOT, "out")
HORIZONS = [1, 2, 4]; PMAX = 8; HAC_LAGS = 3; MODEL = "ar_seas_aic"

q = pd.read_csv(os.path.join(DATA, "cpi_quarterly.csv"), parse_dates=["quarter"])
m = pd.read_csv(os.path.join(DATA, "cpi_monthly.csv"))
n_q_raw, n_m_raw = len(q), len(m)
q = q.dropna(subset=["infl_annualised"]).reset_index(drop=True)   # drops 1995Q1 only
y_all = q.set_index("quarter")["infl_annualised"]
dates = list(y_all.index)
origins = pd.date_range("2015-10-01", "2025-07-01", freq="QS-JAN")
assert len(origins) == 40

def design(y, qtr, p, start):
    """Rows t=start..T-1: [Q1..Q4 dummies, y_{t-1..t-p}]; target y_t."""
    T = len(y)
    X = [np.eye(4)[qtr[start:T]]]
    for j in range(1, p + 1):
        X.append(y[start - j:T - j, None])
    return np.hstack(X), y[start:T]

def fit_select(y, qtr):
    """AIC lag selection p in 0..PMAX on common sample (first PMAX obs held out), then OLS refit on max sample."""
    best = None
    for p in range(PMAX + 1):
        X, Y = design(y, qtr, p, PMAX)
        b, *_ = np.linalg.lstsq(X, Y, rcond=None)
        ssr = np.sum((Y - X @ b) ** 2); n = len(Y)
        aic = n * np.log(ssr / n) + 2 * X.shape[1]
        if best is None or aic < best[0]:
            best = (aic, p)
    p = best[1]
    X, Y = design(y, qtr, p, p)
    b, *_ = np.linalg.lstsq(X, Y, rcond=None)
    return p, b, len(Y)

def iterate(y, qtr_last, p, b, H):
    hist = list(y); out = []
    for h in range(1, H + 1):
        qq = (qtr_last + h) % 4
        f = b[qq] + sum(b[4 + j - 1] * hist[-j] for j in range(1, p + 1))
        hist.append(f); out.append(f)
    return out

rows = []; lag_choice = {}; n_last = None
for t in origins:
    idx = dates.index(t)
    hist = y_all.iloc[: idx + 1]                     # quarters <= t only
    assert hist.index.max() == t
    y = hist.values; qtr = ((hist.index.month - 1) // 3).values
    p, b, n_est = fit_select(y, qtr)
    lag_choice[str(t.date())] = p; n_last = n_est
    fc = iterate(y, qtr[-1], p, b, max(HORIZONS))
    for h in HORIZONS:
        tgt = t + pd.DateOffset(months=3 * h)
        actual = float(y_all.loc[tgt])
        sn_date = tgt - pd.DateOffset(months=12)
        assert sn_date <= t
        sn = float(y_all.loc[sn_date])
        rows.append(dict(origin_date=str(t.date()), target_date=str(tgt.date()), horizon=h, forecast=sn, actual=actual, model="snaive"))
        rows.append(dict(origin_date=str(t.date()), target_date=str(tgt.date()), horizon=h, forecast=float(fc[h - 1]), actual=actual, model=MODEL))
fc_df = pd.DataFrame(rows).sort_values(["model", "horizon", "origin_date"]).reset_index(drop=True)
fc_df.to_csv(os.path.join(OUT, "forecasts.csv"), index=False, float_format="%.6f")
fc_df = pd.read_csv(os.path.join(OUT, "forecasts.csv"))   # metrics computed from the written file

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

res = {"task_id": "F3-conf-04", "model_name": MODEL,
       "model_description": ("Iterated (recursive) AR(p) for annualised quarterly inflation with four quarter-of-year intercepts, "
                             "p in 0..8 selected by AIC on a common sample at each origin, estimated by OLS on an expanding window of all quarters <= origin (from 1995Q2)."),
       "n_origins": len(origins), "first_origin": str(origins[0].date()), "last_origin": str(origins[-1].date())}
err = {}
for mod, key in [("snaive", "snaive"), (MODEL, "model")]:
    for h in HORIZONS:
        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 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
n = len(d); dbar = float(np.mean(d))
se = float(np.sqrt(hac_var_mean(d, HAC_LAGS, "bartlett")))
dm = dbar / se
res.update(loss_diff_mean_h4=dbar, loss_diff_se_h4=se, hac_lags_h4=HAC_LAGS, hac_kernel_h4="bartlett",
           dm_stat_h4=dm, dm_pvalue_h4=float(2 * stats.norm.sf(abs(dm))))
# sensitivity (supplementary keys)
sens = {}
for L, k in [(4, "bartlett"), (6, "bartlett"), (3, "rectangular")]:
    v = hac_var_mean(d, L, k)
    if v > 0:
        s_ = np.sqrt(v); sens[f"{k}_L{L}"] = {"se": float(s_), "dm": float(dbar / s_), "p_normal": float(2 * stats.norm.sf(abs(dbar / s_)))}
    else:
        sens[f"{k}_L{L}"] = {"se": None, "dm": None, "p_normal": None, "note": "non-positive long-run variance"}
hln = np.sqrt((n + 1 - 2 * 4 + 4 * 3 / n) / n)   # Harvey-Leybourne-Newbold small-sample factor, h=4
sens["hln_bartlett_L3_t39"] = {"dm": float(dm * hln), "p_t": float(2 * stats.t.sf(abs(dm * hln), n - 1))}
res["sensitivity_h4"] = sens
res["n_loss_diff_positive_h4"] = int(np.sum(d > 0))
res["selected_lags_by_origin"] = lag_choice
res.update(feasibility="complete", feasibility_reason=None,
           data_issues=[{"issue": "infl_annualised missing for 1995Q1 (no preceding quarter)", "handling": "row dropped; estimation starts 1995Q2"},
                        {"issue": "extreme value 12.15 in 2022Q2 (inflation surge)", "handling": "genuine observation, kept"}],
           n_input_rows={"cpi_quarterly.csv": n_q_raw, "cpi_monthly.csv": n_m_raw},
           n_analysed=int(n_last), random_seed=None)
with open(os.path.join(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 != "selected_lags_by_origin"}, indent=2))
print("selected lags:", sorted(set(lag_choice.values())))