"""F3-conf-04: AR + seasonal dummies vs seasonal naive, Canadian quarterly CPI inflation."""
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")
q = pd.read_csv(os.path.join(D, "cpi_quarterly.csv"), parse_dates=["quarter"])
m = pd.read_csv(os.path.join(D, "cpi_monthly.csv"), parse_dates=["month"])
issues = []
# consistency checks
mq = m.set_index("month")["cpi"].resample("QS").agg(["mean", "count"])
mq = mq[mq["count"] == 3]
chk = q.set_index("quarter")["cpi_avg"].reindex(mq.index)
if (chk - mq["mean"]).abs().max() > 1e-5:
    issues.append({"issue": "quarterly cpi_avg differs from monthly averages", "handling": "used quarterly file as given"})
rec = 400 * np.log(q["cpi_avg"]).diff()
if (rec - q["infl_annualised"]).abs().max() > 1e-5:
    issues.append({"issue": "infl_annualised inconsistent with cpi_avg", "handling": "used as given"})
issues.append({"issue": "infl_annualised missing for 1995Q1 (first difference)", "handling": "row dropped from target series; 126 inflation observations used"})
y_full = q.dropna(subset=["infl_annualised"]).set_index("quarter")["infl_annualised"]
assert y_full.index.to_series().diff().dropna().dt.days.between(89, 92).all()

PMAX, HS = 4, [1, 2, 4]
MODEL = "ar_sd_aic"
origins = pd.date_range("2015-10-01", "2025-07-01", freq="QS")
assert len(origins) == 40

def design(y, p, start):
    """rows t=start..n-1: const, 3 seasonal dummies, lags 1..p"""
    n = len(y); qtr = y.index.quarter.values
    X = [np.ones(n - start)]
    for s in (2, 3, 4):
        X.append((qtr[start:] == s).astype(float))
    for l in range(1, p + 1):
        X.append(y.values[start - l:n - l])
    return np.column_stack(X), y.values[start:]

def fit_forecast(y):
    best = None
    for p in range(PMAX + 1):
        X, Y = design(y, p, PMAX)  # common sample for AIC comparability
        b, *_ = np.linalg.lstsq(X, Y, rcond=None)
        rss = ((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
    # iterate forecasts
    hist = list(y.values); qs = list(y.index)
    fc = {}
    for h in range(1, max(HS) + 1):
        d = qs[-1] + pd.DateOffset(months=3 * h)
        x = [1.0] + [float(d.quarter == s) for s in (2, 3, 4)] + [hist[-l] for l in range(1, p + 1)]
        f = float(np.dot(b, x)); hist.append(f); fc[h] = (d, f)
    return fc, p, n

rows, plist = [], []
for t in origins:
    y = y_full[y_full.index <= t]
    fc, p, nobs = fit_forecast(y); plist.append(p)
    for h in HS:
        d, f = fc[h]
        act = y_full.get(d, np.nan)
        sn = y_full[d - pd.DateOffset(months=12)]
        for name, val in (("snaive", sn), (MODEL, f)):
            rows.append({"origin_date": t.strftime("%Y-%m-%d"), "target_date": d.strftime("%Y-%m-%d"),
                         "horizon": h, "forecast": round(float(val), 6), "actual": float(act), "model": name})
fc = pd.DataFrame(rows)
assert fc["actual"].notna().all() and len(fc) == 240
fc.to_csv(os.path.join(O, "forecasts.csv"), index=False)

res = {"task_id": "F3-conf-04", "model_name": MODEL,
       "model_description": ("OLS AR(p) on annualised quarterly inflation with intercept and three quarterly seasonal dummies, "
                             "p in {0..4} selected by AIC at each origin on a common sample, expanding window from 1995Q2 "
                             "(first lag-4-complete obs 1996Q2) to the origin, multi-step forecasts iterated."),
       "n_origins": 40, "first_origin": origins[0].strftime("%Y-%m-%d"), "last_origin": origins[-1].strftime("%Y-%m-%d")}
err = {}
for h in HS:
    for mod, key in (("snaive", "snaive"), (MODEL, "model")):
        s = fc[(fc.horizon == h) & (fc.model == mod)].sort_values("origin_date")
        e = (s.actual - s.forecast).values; err[(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 HS:
    res[f"rel_improvement_rmse_h{h}"] = 1 - res[f"rmse_model_h{h}"] / res[f"rmse_snaive_h{h}"]

def hac_se(d, L, kernel="bartlett"):
    n = len(d); u = d - d.mean()
    lrv = u @ u / n
    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]) / n
    return np.sqrt(lrv / n), lrv

d = err[(4, "snaive")] ** 2 - err[(4, "model")] ** 2
L = 3
se, _ = hac_se(d, L)
dm = d.mean() / se
res.update({"loss_diff_mean_h4": float(d.mean()), "loss_diff_se_h4": float(se), "hac_lags_h4": L,
            "hac_kernel_h4": "bartlett", "dm_stat_h4": float(dm),
            "dm_pvalue_h4": float(2 * stats.norm.sf(abs(dm))),
            "feasibility": "complete", "feasibility_reason": None, "data_issues": issues,
            "n_input_rows": {"cpi_quarterly.csv": int(len(q)), "cpi_monthly.csv": int(len(m))},
            "n_analysed": int(nobs), "random_seed": None})
json.dump(res, open(os.path.join(O, "results.json"), "w"), indent=2)

# sensitivity (printed only, reported in report.md)
n = len(d)
hln = np.sqrt((n + 1 - 2 * 4 + 4 * 3 / n) / n)
print("selected p per origin:", plist)
print(json.dumps({k: v for k, v in res.items() if isinstance(v, (int, float))}, indent=1))
print("HLN-adjusted stat, t(39) p:", dm * hln, 2 * stats.t.sf(abs(dm * hln), n - 1))
for L2, k in ((3, "rectangular"), (4, "bartlett"), (6, "bartlett"), (8, "bartlett")):
    s2, lrv = hac_se(d, L2, k); print(k, L2, "se", s2, "lrv", lrv, "stat", d.mean() / s2 if lrv > 0 else None)
for h in (1, 2):
    dd = err[(h, "snaive")] ** 2 - err[(h, "model")] ** 2
    s2, _ = hac_se(dd, max(h - 1, 1) + 2); print("h", h, "dbar", dd.mean(), "se", s2, "stat", dd.mean() / s2)
# excluding 2021-2023 targets
s = fc[fc.horizon == 4].pivot(index="origin_date", columns="model", values=["forecast", "actual"])
tgt = fc[(fc.horizon == 4) & (fc.model == "snaive")].sort_values("origin_date").target_date.values
mask = ~((tgt >= "2021-01-01") & (tgt <= "2023-10-01"))
print("h4 excl 2021-2023 targets: n", mask.sum(), "rmse sn", np.sqrt(np.mean(err[(4,'snaive')][mask]**2)),
      "rmse model", np.sqrt(np.mean(err[(4,'model')][mask]**2)))