"""F3-conf-04: multi-horizon forecasts of Canadian quarterly CPI inflation.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Model: direct h-step OLS regression of y_{s+h} on 4 quarter-of-target dummies and
p lags y_s..y_{s-p+1}, p in {0..4} chosen by BIC at each origin and horizon on a
common sample, expanding window using only quarters <= origin. No randomness."""
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__)))
DATA, OUT = os.path.join(ROOT, "data"), os.path.join(ROOT, "out")
H = [1, 2, 4]; PMAX = 4; MODEL = "dirSAR_bic"

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_rows = {"cpi_quarterly.csv": int(len(q)), "cpi_monthly.csv": int(len(m))}
q = q.dropna(subset=["infl_annualised"]).reset_index(drop=True)  # drops 1995Q1 only
dates = q["quarter"]; y = q["infl_annualised"].to_numpy(); qtr = dates.dt.quarter.to_numpy()
origins = pd.date_range("2015-10-01", "2025-07-01", freq="QS")
assert len(origins) == 40

def design(yy, qq, h, p, start):
    """Rows s (index of origin-in-sample) with s>=start and s+h<len(yy)."""
    S = np.arange(start, len(yy) - h)
    X = [(qq[S + h] == k).astype(float) for k in (1, 2, 3, 4)]
    X += [yy[S - i] for i in range(p)]
    return np.column_stack(X), yy[S + h]

def fit_forecast(yy, qq, h):
    """yy, qq: data up to and including origin (last element = origin)."""
    start = PMAX - 1  # common sample for all p
    best = None
    for p in range(PMAX + 1):
        X, Y = design(yy, qq, h, p, start)
        b, *_ = np.linalg.lstsq(X, Y, rcond=None)
        n = len(Y); rss = np.sum((Y - X @ b) ** 2)
        bic = n * np.log(rss / n) + X.shape[1] * np.log(n)
        if best is None or bic < best[0]:
            best = (bic, p)
    p = best[1]
    X, Y = design(yy, qq, h, p, max(p - 1, 0))  # re-estimate on all usable rows
    b, *_ = np.linalg.lstsq(X, Y, rcond=None)
    tq = (qq[-1] - 1 + h) % 4 + 1
    x = [float(tq == k) for k in (1, 2, 3, 4)] + [yy[-1 - i] for i in range(p)]
    return float(np.dot(b, x)), p, len(Y)

rows, chosen_p, nobs_last = [], {h: [] for h in H}, {}
for o in origins:
    io = int(np.where(dates == o)[0][0])
    yy, qq = y[: io + 1], qtr[: io + 1]          # information set: quarters <= origin
    assert dates.iloc[io] == o and dates.iloc[: io + 1].max() <= o
    for h in H:
        it = io + h; tdate = dates.iloc[it]; act = y[it]
        sn = y[it - 4]; assert it - 4 <= io
        f, p, n = fit_forecast(yy, qq, h)
        chosen_p[h].append(p)
        if o == origins[-1]: nobs_last[h] = n
        for mod, val in (("snaive", sn), (MODEL, f)):
            rows.append(dict(origin_date=o.strftime("%Y-%m-%d"), target_date=tdate.strftime("%Y-%m-%d"),
                             horizon=h, forecast=round(val, 6), actual=round(act, 6), model=mod))
fc = pd.DataFrame(rows)
fc.to_csv(os.path.join(OUT, "forecasts.csv"), index=False)

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

res = {"task_id": "F3-conf-04", "model_name": MODEL,
       "model_description": ("Direct h-step OLS regression of inflation at t+h on four target-quarter dummies and p lags "
                             "of inflation (y_t..y_{t-p+1}), p in 0..4 selected by BIC per origin and horizon, "
                             "expanding window from 1995Q2 to the origin."),
       "n_origins": len(origins), "first_origin": origins[0].strftime("%Y-%m-%d"),
       "last_origin": origins[-1].strftime("%Y-%m-%d")}
err = {}
for mod, tag in (("snaive", "snaive"), (MODEL, "model")):
    for h in H:
        s = fc[(fc.model == mod) & (fc.horizon == h)]
        e = (s.actual - s.forecast).to_numpy(); err[(tag, h)] = e
        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}"]
d = err[("snaive", 4)] ** 2 - err[("model", 4)] ** 2
L = 4; se = hac_se(d, L, "bartlett"); n = len(d)
res.update({"loss_diff_mean_h4": float(np.mean(d)), "loss_diff_se_h4": se, "hac_lags_h4": L,
            "hac_kernel_h4": "bartlett", "dm_stat_h4": float(np.mean(d) / se)})
res["dm_pvalue_h4"] = float(2 * stats.norm.sf(abs(res["dm_stat_h4"])))
# sensitivity (not required): other HAC choices and small-sample t reference
sens = {}
for k, LL in (("bartlett", 3), ("bartlett", 8), ("rectangular", 3)):
    s2 = hac_se(d, LL, k)
    sens[f"{k}_L{LL}"] = None if s2 is None else {"se": s2, "stat": float(np.mean(d) / s2),
        "p_normal": float(2 * stats.norm.sf(abs(np.mean(d) / s2))),
        "p_t39": float(2 * stats.t.sf(abs(np.mean(d) / s2), n - 1))}
sens["main_p_t39"] = float(2 * stats.t.sf(abs(res["dm_stat_h4"]), n - 1))
res["sensitivity_dm_h4"] = sens
res["dm_tests_h1_h2_bartlett_L_h"] = {}
for h in (1, 2):
    dh = err[("snaive", h)] ** 2 - err[("model", h)] ** 2; sh = hac_se(dh, h)
    res["dm_tests_h1_h2_bartlett_L_h"][f"h{h}"] = {"mean": float(np.mean(dh)), "se": sh,
        "stat": float(np.mean(dh) / sh), "p_normal": float(2 * stats.norm.sf(abs(np.mean(dh) / sh)))}
res.update({"feasibility": "complete", "feasibility_reason": None,
            "data_issues": [{"issue": "infl_annualised missing for 1995Q1 (no previous quarter)",
                             "handling": "row dropped; inflation sample starts 1995Q2"},
                            {"issue": "infl_annualised recomputed from rounded cpi_avg differs by up to ~3e-6",
                             "handling": "file values used as given (negligible)"}],
            "n_input_rows": n_rows,
            "n_analysed": int(np.sum(dates <= origins[-1])),
            "n_regression_rows_last_origin": {f"h{h}": nobs_last[h] for h in H},
            "selected_p_counts": {f"h{h}": {str(k): chosen_p[h].count(k) for k in range(PMAX + 1)} for h in H},
            "baseline_note": "seasonal-naive baseline = keys *_snaive_*",
            "random_seed": None})
with open(os.path.join(OUT, "results.json"), "w") as fh:
    json.dump(res, fh, indent=2)
print(json.dumps({k: v for k, v in res.items() if k.startswith(("rmse", "mae", "rel", "loss", "dm_stat", "dm_p"))}, indent=1))