"""F3-conf-04: multi-horizon forecasts of Canadian quarterly CPI inflation (NSA, annualised).
Model: AR(p) with quarterly seasonal dummies, OLS, iterated multi-step, p in 0..4 chosen by AIC
on each expanding window (common sample). Baseline: seasonal naive. DM test at h=4 with Bartlett HAC."""
import json, numpy as np, pandas as pd
from scipy import stats
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
OUT = ROOT / "out"
q = pd.read_csv(ROOT / "data/cpi_quarterly.csv", parse_dates=["quarter"])
m = pd.read_csv(ROOT / "data/cpi_monthly.csv", parse_dates=["month"])
issues = []
# --- data checks
m = m.sort_values("month").reset_index(drop=True)
if m["cpi"].isna().any(): issues.append({"issue": "missing monthly CPI", "handling": "none needed for quarterly input"})
q = q.sort_values("quarter").reset_index(drop=True)
exp_q = pd.date_range(q.quarter.iloc[0], q.quarter.iloc[-1], freq="QS")
assert len(exp_q) == len(q) and (exp_q == q.quarter).all(), "quarterly gaps"
mq = m.assign(quarter=m.month.dt.to_period("Q").dt.start_time).groupby("quarter")["cpi"].agg(["mean", "size"])
mq = mq[mq["size"] == 3]["mean"]
chk = q.set_index("quarter")["cpi_avg"] - mq.reindex(q.quarter).values
infl_chk = 400 * np.log(q.cpi_avg).diff()
maxdev_avg = float(np.nanmax(np.abs(chk.values))); maxdev_inf = float(np.nanmax(np.abs(infl_chk - q.infl_annualised)))
issues.append({"issue": "infl_annualised is empty for 1995Q1 (no previous quarter)",
               "handling": "1995Q1 dropped from the inflation series; estimation starts 1995Q2"})
if maxdev_avg > 1e-5 or maxdev_inf > 1e-4:
    issues.append({"issue": f"quarterly file deviates from monthly recomputation (max cpi_avg dev {maxdev_avg:.2e}, infl dev {maxdev_inf:.2e})",
                   "handling": "quarterly file used as given (it is the designated input)"})

y = q.set_index("quarter")["infl_annualised"].dropna()
dates = list(y.index)
H = [1, 2, 4]; PMAX = 4; MODEL = "ar_seas_aic"
origins = pd.date_range("2015-10-01", "2025-07-01", freq="QS")
assert len(origins) == 40

def design(yv, qtr, p, start):
    # rows t=start..len-1 ; regressors: 4 seasonal dummies + p lags
    X = [np.eye(4)[qtr[start:]]]
    for l in range(1, p + 1):
        X.append(yv[start - l:len(yv) - l][:, None])
    return np.hstack(X), yv[start:]

def fit_forecast(ys, hmax):
    yv = ys.values; qtr = np.array([d.month // 3 for d in ys.index])  # 0..3
    best = None
    for p in range(PMAX + 1):
        X, Y = design(yv, qtr, p, PMAX)  # common sample for AIC
        b, *_ = np.linalg.lstsq(X, Y, rcond=None)
        n = len(Y); s2 = np.sum((Y - X @ b) ** 2) / n
        aic = n * np.log(s2) + 2 * X.shape[1]
        if best is None or aic < best[0]: best = (aic, p)
    p = best[1]
    X, Y = design(yv, qtr, p, p)  # re-estimate on max sample for chosen p
    b, *_ = np.linalg.lstsq(X, Y, rcond=None)
    hist = list(yv); lastq = qtr[-1]; fc = []
    for h in range(1, hmax + 1):
        qq = (lastq + h) % 4
        f = b[qq] + sum(b[4 + l - 1] * hist[-l] for l in range(1, p + 1))
        hist.append(f); fc.append(f)
    return fc, p, len(Y)

rows = []; plist = []
for o in origins:
    ys = y[y.index <= o]
    fc, p, nobs = fit_forecast(ys, max(H)); plist.append(p)
    for h in H:
        tgt = o + pd.DateOffset(months=3 * h)
        act = float(y.loc[tgt])
        sn = float(y.loc[tgt - pd.DateOffset(months=12)])
        assert tgt - pd.DateOffset(months=12) <= o
        rows.append(dict(origin_date=o.date(), target_date=tgt.date(), horizon=h, forecast=sn, actual=act, model="snaive"))
        rows.append(dict(origin_date=o.date(), target_date=tgt.date(), horizon=h, forecast=fc[h - 1], actual=act, model=MODEL))
n_last = nobs
fc = pd.DataFrame(rows); fc.to_csv(OUT / "forecasts.csv", index=False, float_format="%.6f")
fc["e"] = fc.actual - fc.forecast

res = {"task_id": "F3-conf-04", "model_name": MODEL,
       "model_description": ("Iterated multi-step AR(p) for annualised NSA quarterly inflation with four quarterly seasonal intercepts, "
                             "p in {0,...,4} chosen by AIC on a common sample at each origin, estimated by OLS on an expanding window from 1995Q2 to the origin."),
       "n_origins": 40, "first_origin": str(origins[0].date()), "last_origin": str(origins[-1].date())}
for h in H:
    for mod, tag in [("snaive", "snaive"), (MODEL, "model")]:
        e = fc[(fc.horizon == h) & (fc.model == mod)].e.values
        res[f"rmse_{tag}_h{h}"] = float(np.sqrt(np.mean(e ** 2)))
        res[f"mae_{tag}_h{h}"] = float(np.mean(np.abs(e)))
for h in H:
    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"):
    d = np.asarray(d); 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

def dloss(h):
    es = fc[(fc.horizon == h) & (fc.model == "snaive")].sort_values("origin_date").e.values
    em = fc[(fc.horizon == h) & (fc.model == MODEL)].sort_values("origin_date").e.values
    return es ** 2 - em ** 2
d4 = dloss(4); L = 4
se, _ = hac_se(d4, L)
dm = d4.mean() / se
res.update({"loss_diff_mean_h4": float(d4.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)))})
# sensitivity (reported in report, not required keys)
sens = {}
for k, LL in [("bartlett", 3), ("bartlett", 4), ("bartlett", 6), ("rectangular", 3)]:
    s, lrv = hac_se(d4, LL, k)
    t = d4.mean() / s if lrv > 0 else float("nan")
    n = len(d4); hh = 4
    hln = np.sqrt((n + 1 - 2 * hh + hh * (hh - 1) / n) / n)
    sens[f"{k}_L{LL}"] = {"se": float(s), "stat": float(t), "p_normal": float(2 * stats.norm.sf(abs(t))),
                          "p_t39": float(2 * stats.t.sf(abs(t), n - 1)),
                          "hln_stat": float(t * hln), "hln_p_t39": float(2 * stats.t.sf(abs(t * hln), n - 1))}
for h in [1, 2]:
    d = dloss(h); s, _ = hac_se(d, max(h - 1, 1))
    sens[f"h{h}_bartlett_L{max(h-1,1)}"] = {"mean": float(d.mean()), "se": float(s), "stat": float(d.mean() / s),
                                          "p_normal": float(2 * stats.norm.sf(abs(d.mean() / s)))}
res.update({"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(n_last), "random_seed": None})
json.dump(res, open(OUT / "results.json", "w"), indent=2)
json.dump({"sensitivity_h4_and_h12": sens, "ar_order_by_origin": dict(zip([str(o.date()) for o in origins], plist)),
           "check_max_dev_cpi_avg": maxdev_avg, "check_max_dev_infl": maxdev_inf},
          open(OUT / "diagnostics.json", "w"), indent=2)
print(json.dumps(res, indent=2))