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

H = [1, 2, 4]; PMAX = 4; MODEL = "ar_seasdum"
q = pd.read_csv("data/cpi_quarterly.csv", parse_dates=["quarter"])
m = pd.read_csv("data/cpi_monthly.csv")
n_rows = {"cpi_quarterly.csv": int(len(q)), "cpi_monthly.csv": int(len(m))}
s = q.dropna(subset=["infl_annualised"]).reset_index(drop=True)
y = s["infl_annualised"].to_numpy(); dates = s["quarter"]; qtr = dates.dt.quarter.to_numpy()
origins = pd.date_range("2015-10-01", "2025-07-01", freq="QS")
assert len(origins) == 40
pos = {d: i for i, d in enumerate(dates)}

def design(yv, qv, p, start):
    T = len(yv); rows = range(start, T)
    X = [[1.0 * (qv[t] == k) for k in (1, 2, 3, 4)] + [yv[t - j] for j in range(1, p + 1)] for t in rows]
    return np.array(X), yv[start:]

def fit_forecast(i_end):
    yv, qv = y[: i_end + 1], qtr[: i_end + 1]          # only quarters <= origin
    best = None
    for p in range(PMAX + 1):                           # AIC on common sample (start = PMAX)
        X, Y = design(yv, qv, 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, b, n)
    _, p, b, n = best
    hist = list(yv); qs = list(qv); fc = {}
    for h in range(1, max(H) + 1):                      # iterated multi-step
        qn = (qs[-1] % 4) + 1
        x = [1.0 * (qn == k) for k in (1, 2, 3, 4)] + [hist[-j] for j in range(1, p + 1)]
        f = float(np.dot(x, b)); hist.append(f); qs.append(qn); fc[h] = f
    return fc, p, n

rows = []; plist = []; n_last = None
for o in origins:
    i = pos[o]; fc, p, n = fit_forecast(i); plist.append(p); n_last = n
    for h in H:
        tgt = i + h; act = float(y[tgt])
        assert dates[tgt] == o + pd.DateOffset(months=3 * h)
        assert tgt - 4 <= i
        rows.append(dict(origin_date=o.date().isoformat(), target_date=dates[tgt].date().isoformat(), horizon=h,
                         forecast=float(y[tgt - 4]), actual=act, model="snaive"))
        rows.append(dict(origin_date=o.date().isoformat(), target_date=dates[tgt].date().isoformat(), horizon=h,
                         forecast=fc[h], actual=act, model=MODEL))
F = pd.DataFrame(rows); F["err"] = F["actual"] - F["forecast"]
os.makedirs("out", exist_ok=True)
F.drop(columns="err").to_csv("out/forecasts.csv", index=False, float_format="%.6f")

R = {"task_id": "F3-conf-04", "model_name": MODEL,
     "model_description": ("OLS AR(p) on annualised quarterly inflation with four quarterly seasonal intercepts, "
                           "p in {0..4} chosen by AIC at each origin on a common sample, expanding window from 1995Q2 "
                           "to the origin, iterated multi-step forecasts."),
     "n_origins": 40, "first_origin": origins[0].date().isoformat(), "last_origin": origins[-1].date().isoformat()}
for h in H:
    for mod, key in (("snaive", "snaive"), (MODEL, "model")):
        e = F[(F.horizon == h) & (F.model == mod)]["err"].to_numpy()
        R[f"rmse_{key}_h{h}"] = float(np.sqrt(np.mean(e ** 2))); R[f"mae_{key}_h{h}"] = float(np.mean(np.abs(e)))
for h in H:
    R[f"rel_improvement_rmse_h{h}"] = 1 - R[f"rmse_model_h{h}"] / R[f"rmse_snaive_h{h}"]

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

e_s = F[(F.horizon == 4) & (F.model == "snaive")].sort_values("origin_date")["err"].to_numpy()
e_m = F[(F.horizon == 4) & (F.model == MODEL)].sort_values("origin_date")["err"].to_numpy()
d = e_s ** 2 - e_m ** 2; T = len(d); L = 3
dbar = float(d.mean()); se = float(np.sqrt(hac_var_mean(d, L)))
dm = dbar / se
R.update({"loss_diff_mean_h4": dbar, "loss_diff_se_h4": se, "hac_lags_h4": L, "hac_kernel_h4": "bartlett",
          "dm_stat_h4": dm, "dm_pvalue_h4": float(2 * stats.norm.sf(abs(dm)))})
# sensitivity (not primary)
sens = {}
for LL in (4, 6, 8):
    sens[f"bartlett_L{LL}_se"] = float(np.sqrt(hac_var_mean(d, LL)))
    sens[f"bartlett_L{LL}_pvalue"] = float(2 * stats.norm.sf(abs(dbar / np.sqrt(hac_var_mean(d, LL)))))
vr = hac_var_mean(d, 3, "rectangular"); sens["rectangular_L3_se"] = float(np.sqrt(vr)) if vr > 0 else None
hln = np.sqrt((T + 1 - 2 * 4 + 4 * 3 / T) / T)          # Harvey-Leybourne-Newbold, h=4
sens["hln_stat_bartlett_L3"] = float(dm * hln)
sens["hln_pvalue_t39"] = float(2 * stats.t.sf(abs(dm * hln), T - 1))
sens["iid_se_L0"] = float(np.sqrt(hac_var_mean(d, 0)))
sens["d_autocorr_lag1_to_4"] = [float(np.corrcoef(d[j:], d[:-j])[0, 1]) for j in range(1, 5)]
sens["n_origins_model_better_sqerr_h4"] = int((d > 0).sum())
R["sensitivity_h4"] = sens
R["selected_ar_order_counts"] = {str(k): int(v) for k, v in pd.Series(plist).value_counts().sort_index().items()}
R.update({"feasibility": "complete", "feasibility_reason": None,
          "data_issues": [{"issue": "infl_annualised missing for 1995Q1 (no prior quarter)", "handling": "dropped; series starts 1995Q2"},
                          {"issue": "infl_annualised reproduced from cpi_avg only to ~3e-6 due to 6-dp rounding of cpi_avg", "handling": "used infl_annualised as given"},
                          {"issue": "COVID-2020 and 2021-2023 inflation-surge outliers within evaluation window", "handling": "kept (genuine data); noted as threat to DM stationarity assumption"}],
          "n_input_rows": n_rows, "n_analysed": int(n_last), "random_seed": None})
with open("out/results.json", "w") as f: json.dump(R, f, indent=2)
print(json.dumps({k: R[k] for k in R if k.startswith(("rmse", "mae", "rel", "dm", "loss"))}, indent=1))