"""F3-conf-04: multi-horizon forecasts of Canadian quarterly CPI inflation (NSA, annualised).
Model: AR(p) + quarter-of-year dummies, OLS, iterated multi-step, p in 0..4 chosen by AIC
on each expanding window (data dated <= origin only). Baseline: seasonal naive.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
"""
import json
import numpy as np
import pandas as pd
from scipy import stats

DATA = "data"
OUT = "out"
MODEL = "ar_sdum_aic"
PMAX = 4
H = [1, 2, 4]

q = pd.read_csv(f"{DATA}/cpi_quarterly.csv")
m = pd.read_csv(f"{DATA}/cpi_monthly.csv")
n_rows = {"cpi_quarterly.csv": int(len(q)), "cpi_monthly.csv": int(len(m))}
q["quarter"] = pd.to_datetime(q["quarter"])
q = q.sort_values("quarter").reset_index(drop=True)

# ---- data checks ----
issues = []
assert not q["quarter"].duplicated().any()
assert (q["quarter"].dt.to_period("Q").diff().iloc[1:].apply(lambda x: x.n) == 1).all()
m["q"] = pd.to_datetime(m["month"]).dt.to_period("Q")
agg = m.groupby("q")["cpi"].agg(["mean", "count"])
assert (agg["count"] == 3).all() and len(agg) == len(q)
assert np.max(np.abs(agg["mean"].values - q["cpi_avg"].values)) < 1e-5
recomp = 400 * np.log(q["cpi_avg"]).diff()
assert np.nanmax(np.abs(recomp - q["infl_annualised"])) < 1e-4
issues.append({"issue": "infl_annualised is missing for 1995Q1 (first difference undefined)",
               "handling": "1995Q1 excluded from all estimation samples; consistency of cpi_avg with the monthly file and of infl_annualised with 400*dln(cpi_avg) verified (max abs diff < 1e-5)"})
issues.append({"issue": "Evaluation window contains the 2020 COVID shock and the 2021-2023 inflation surge (outliers relative to the pre-2015 estimation sample)",
               "handling": "No winsorising or exclusion; errors computed on the data as given, noted as a limitation"})

y_all = q["infl_annualised"].values
dates = q["quarter"].values
qtr = q["quarter"].dt.quarter.values


def fit_ar_sdum(y, qq, p, start):
    """OLS of y_t on 4 quarter dummies and p lags, t = start..T-1 (common sample)."""
    T = len(y)
    rows = range(start, T)
    X = []
    for t in rows:
        d = [1.0 if qq[t] == k else 0.0 for k in (1, 2, 3, 4)]
        X.append(d + [y[t - j] for j in range(1, p + 1)])
    X = np.array(X)
    Y = y[start:]
    beta, *_ = np.linalg.lstsq(X, Y, rcond=None)
    resid = Y - X @ beta
    n = len(Y)
    k = X.shape[1]
    aic = n * np.log(resid @ resid / n) + 2 * k
    return beta, aic, n


def forecast_iter(y, qq, beta, p, hmax):
    path = list(y)
    qpath = list(qq)
    out = []
    for h in range(1, hmax + 1):
        qn = (qpath[-1] % 4) + 1
        d = [1.0 if qn == k else 0.0 for k in (1, 2, 3, 4)]
        x = np.array(d + [path[-j] for j in range(1, p + 1)])
        f = float(x @ beta)
        out.append(f)
        path.append(f)
        qpath.append(qn)
    return out


origins = pd.date_range("2015-10-01", "2025-07-01", freq="QS")
assert len(origins) == 40
rows = []
chosen_p = []
n_last = None
for o in origins:
    io = int(np.where(dates == np.datetime64(o))[0][0])
    y = y_all[1:io + 1]          # drop 1995Q1 (NaN); data dated <= origin only
    qq = qtr[1:io + 1]
    start = PMAX                 # common sample across p for AIC comparability
    best = None
    for p in range(0, PMAX + 1):
        beta, aic, n = fit_ar_sdum(y, qq, p, start)
        if best is None or aic < best[1]:
            best = (p, aic, beta)
    p, _, beta = best
    beta, _, n = fit_ar_sdum(y, qq, p, p)  # re-estimate on the longest sample for chosen p
    n_last = n
    chosen_p.append(p)
    fc = forecast_iter(y, qq, beta, p, max(H))
    for h in H:
        it = io + h
        act = float(y_all[it])
        tdate = pd.Timestamp(dates[it])
        rows.append(dict(origin_date=o.date().isoformat(), target_date=tdate.date().isoformat(),
                         horizon=h, forecast=fc[h - 1], actual=act, model=MODEL))
        rows.append(dict(origin_date=o.date().isoformat(), target_date=tdate.date().isoformat(),
                         horizon=h, forecast=float(y_all[it - 4]), actual=act, model="snaive"))

fc = pd.DataFrame(rows)
fc.to_csv(f"{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) with four quarter-of-year intercept dummies, "
                             "estimated by OLS on an expanding window (1995Q2 to origin), lag order p in 0..4 "
                             "re-selected by AIC at every origin on a common sample, then re-estimated on the full window."),
       "n_origins": 40, "first_origin": origins[0].date().isoformat(),
       "last_origin": origins[-1].date().isoformat()}
for h in H:
    for mod, tag in (("snaive", "snaive"), (MODEL, "model")):
        e = fc[(fc.horizon == h) & (fc.model == mod)]["e"].values
        assert len(e) == 40
        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):
    d = np.asarray(d) - np.mean(d)
    n = len(d)
    lrv = d @ d / n
    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]) / n
    return np.sqrt(lrv / n) if lrv > 0 else np.nan


def dm_d(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 = dm_d(4)
L, K = 4, "bartlett"   # Bartlett with L=4: non-zero weights on autocovariances 1..4 (>= h-1 = 3)
se = hac_se(d4, L, K)
res["loss_diff_mean_h4"] = float(np.mean(d4))
res["loss_diff_se_h4"] = float(se)
res["hac_lags_h4"] = L
res["hac_kernel_h4"] = K
res["dm_stat_h4"] = float(np.mean(d4) / se)
# two-sided p-value from Student t with n-1 df (Harvey-Leybourne-Newbold style small-sample reference)
res["dm_pvalue_h4"] = float(2 * stats.t.sf(abs(res["dm_stat_h4"]), df=len(d4) - 1))

# sensitivity (not required keys)
sens = {}
for (LL, KK) in [(3, "bartlett"), (4, "bartlett"), (6, "bartlett"), (8, "bartlett"), (3, "rectangular")]:
    s = hac_se(d4, LL, KK)
    st = np.mean(d4) / s
    sens[f"{KK}_L{LL}"] = {"se": float(s), "stat": float(st),
                           "p_t": float(2 * stats.t.sf(abs(st), len(d4) - 1)),
                           "p_normal": float(2 * stats.norm.sf(abs(st)))}
# HLN correction factor for h=4 with rectangular L=3
n = len(d4); h = 4
hln = np.sqrt((n + 1 - 2 * h + h * (h - 1) / n) / n)
sens["hln_rectangular_L3"] = {"stat": float(sens["rectangular_L3"]["stat"] * hln),
                              "p_t": float(2 * stats.t.sf(abs(sens["rectangular_L3"]["stat"] * hln), n - 1))}
for hh in (1, 2):
    dd = dm_d(hh)
    s = hac_se(dd, max(hh, 1), "bartlett")
    sens[f"dm_h{hh}_bartlett_L{max(hh,1)}"] = {"mean": float(np.mean(dd)), "se": float(s),
                                              "stat": float(np.mean(dd) / s),
                                              "p_t": float(2 * stats.t.sf(abs(np.mean(dd) / s), n - 1))}
# excluding 2020-2023 targets at h=4
mask = ~fc[(fc.horizon == 4) & (fc.model == "snaive")].sort_values("origin_date")["target_date"].str[:4].isin(["2020", "2021", "2022", "2023"]).values
sens["h4_excl_targets_2020_2023"] = {"n": int(mask.sum()), "loss_diff_mean": float(np.mean(d4[mask]))}

res["feasibility"] = "complete"
res["feasibility_reason"] = None
res["data_issues"] = issues
res["n_input_rows"] = n_rows
res["n_analysed"] = int(n_last)
res["random_seed"] = None
res["supplementary"] = {"chosen_p_by_origin": dict(zip([o.date().isoformat() for o in origins], chosen_p)),
                        "dm_sensitivity": sens,
                        "dm_pvalue_reference": "Student t with n-1 = 39 df"}
with open(f"{OUT}/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))