"""F3-conf-04: multi-horizon forecasts of Canadian quarterly CPI inflation.
Run from workspace root:  ~/wp29/venv/bin/python out/analysis.py
Reads only data/, writes out/results.json and out/forecasts.csv. Deterministic (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__)))
D = lambda *p: os.path.join(ROOT, *p)

q = pd.read_csv(D("data", "cpi_quarterly.csv"), parse_dates=["quarter"])
m = pd.read_csv(D("data", "cpi_monthly.csv"))
n_q_rows, n_m_rows = len(q), len(m)
q = q.dropna(subset=["infl_annualised"]).reset_index(drop=True)   # exclusion step 1 (1995Q1)
y = pd.Series(q["infl_annualised"].values, index=pd.PeriodIndex(q["quarter"], freq="Q"))

ORIGINS = pd.period_range("2015Q4", "2025Q3", freq="Q")
HORIZONS = [1, 2, 4]
MODEL = "direct_sdum_ybar4"
HAC_LAGS, HAC_KERNEL = 4, "bartlett"


def snaive(t, h):
    src = t + h - 4
    assert src <= t
    return float(y[src])


def model_forecast(t, h):
    """Direct OLS: y_{s+h} = sum_q g_q*1[quarter(s+h)=q] + b*mean(y_{s-3..s}) + e,
    estimated on all pairs with s+h <= t (only data dated <= t)."""
    ys = y[y.index <= t]                       # information set at origin t
    ybar4 = ys.rolling(4).mean()
    rows, tgt = [], []
    for s in ys.index:
        if s + h > t or np.isnan(ybar4[s]):
            continue
        rows.append(_x(s + h, ybar4[s])); tgt.append(ys[s + h])
    X, Y = np.array(rows), np.array(tgt)
    beta, *_ = np.linalg.lstsq(X, Y, rcond=None)
    fc = float(_x(t + h, ybar4[t]) @ beta)
    return fc, len(Y)


def _x(target_period, ybar):
    d = np.zeros(4); d[target_period.quarter - 1] = 1.0
    return np.r_[d, ybar]


recs, n_last = [], {}
for t in ORIGINS:
    for h in HORIZONS:
        act = float(y[t + h])
        recs.append(dict(origin_date=t.start_time.date().isoformat(), target_date=(t + h).start_time.date().isoformat(),
                         horizon=h, forecast=snaive(t, h), actual=act, model="snaive"))
        fc, n = model_forecast(t, h)
        if t == ORIGINS[-1]:
            n_last[h] = n
        recs.append(dict(origin_date=t.start_time.date().isoformat(), target_date=(t + h).start_time.date().isoformat(),
                         horizon=h, forecast=fc, actual=act, model=MODEL))
fc = pd.DataFrame(recs)
fc["error"] = fc["actual"] - fc["forecast"]
fc.drop(columns="error").to_csv(D("out", "forecasts.csv"), index=False, float_format="%.6f")


def hac_var_mean(d, L, kernel):
    d = np.asarray(d) - np.mean(d); T = len(d)
    g = lambda k: np.sum(d[k:] * d[:T - k]) / T
    w = (lambda k: 1 - k / (L + 1)) if kernel == "bartlett" else (lambda k: 1.0)
    lrv = g(0) + 2 * sum(w(k) * g(k) for k in range(1, L + 1))
    return lrv / T


res = dict(task_id="F3-conf-04", model_name=MODEL,
           model_description=("Direct (one regression per horizon) OLS of y_{t+h} on four target-quarter seasonal dummies "
                              "and the mean of the last four quarterly inflation rates at the origin (y_t..y_{t-3}), "
                              "no intercept, re-estimated at each origin on an expanding window of all pairs with "
                              "target quarter <= origin (sample starts 1995Q2); no model selection."),
           n_origins=len(ORIGINS), first_origin=ORIGINS[0].start_time.date().isoformat(),
           last_origin=ORIGINS[-1].start_time.date().isoformat())
err = {}
for h in HORIZONS:
    for mod, key in [("snaive", "snaive"), (MODEL, "model")]:
        e = fc[(fc.horizon == h) & (fc.model == mod)].sort_values("origin_date")["error"].values
        assert len(e) == 40
        err[(key, h)] = 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 HORIZONS:
    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
dbar = float(np.mean(d)); se = float(np.sqrt(hac_var_mean(d, HAC_LAGS, HAC_KERNEL)))
dm = dbar / se
res.update(loss_diff_mean_h4=dbar, loss_diff_se_h4=se, hac_lags_h4=HAC_LAGS, hac_kernel_h4=HAC_KERNEL,
           dm_stat_h4=dm, dm_pvalue_h4=float(2 * stats.norm.sf(abs(dm))))
# sensitivity analyses (not the primary test)
T = len(d); h = 4
hln = np.sqrt((T + 1 - 2 * h + h * (h - 1) / T) / T)
sens = {"bartlett_lag3": None, "rectangular_lag3": None}
for name, L, k in [("bartlett_lag3", 3, "bartlett"), ("rectangular_lag3", 3, "rectangular"), ("bartlett_lag6", 6, "bartlett")]:
    v = hac_var_mean(d, L, k)
    if v > 0:
        s = dbar / np.sqrt(v)
        sens[name] = dict(se=float(np.sqrt(v)), dm_stat=float(s), p_normal=float(2 * stats.norm.sf(abs(s))))
    else:
        sens[name] = dict(se=None, note="non-positive long-run variance estimate")
s_hln = dm * hln
sens["primary_with_HLN_correction_t39"] = dict(dm_stat=float(s_hln), p_t=float(2 * stats.t.sf(abs(s_hln), T - 1)))
res["sensitivity_dm_h4"] = sens
res.update(feasibility="complete", feasibility_reason=None,
           data_issues=[{"issue": "infl_annualised missing for 1995Q1 (no prior quarter)", "handling": "row dropped (exclusion step 1)"},
                        {"issue": "infl_annualised differs from 400*dln(cpi_avg) by up to ~3e-6 due to 6-dp rounding of cpi_avg",
                         "handling": "file values used as given"}],
           n_input_rows={"cpi_quarterly.csv": n_q_rows, "cpi_monthly.csv": n_m_rows},
           n_analysed=int(n_last[1]), n_analysed_by_horizon_last_origin={f"h{k}": int(v) for k, v in n_last.items()},
           random_seed=None)
with open(D("out", "results.json"), "w") as f:
    json.dump(res, f, indent=2)
for k in ["rmse_snaive_h1", "rmse_model_h1", "rmse_snaive_h2", "rmse_model_h2", "rmse_snaive_h4", "rmse_model_h4",
          "loss_diff_mean_h4", "loss_diff_se_h4", "dm_stat_h4", "dm_pvalue_h4"]:
    print(f"{k}: {res[k]:.4f}")