"""F3-conf-04: multi-horizon forecasts of Canadian quarterly CPI inflation.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Reads data/ only; writes out/results.json, out/forecasts.csv, out/exclusions.csv,
out/sensitivity_dm.csv. Deterministic (no randomness)."""
import json, os
import numpy as np, pandas as pd
from scipy import stats

OUT = "out"
HORIZONS = [1, 2, 4]
PMAX = 4
MODEL = "dsar_bic"

q = pd.read_csv("data/cpi_quarterly.csv", parse_dates=["quarter"])
m = pd.read_csv("data/cpi_monthly.csv")
n_q, n_m = len(q), len(m)
q = q.sort_values("quarter").reset_index(drop=True)

excl = []
n0 = len(q)
y_all = q.dropna(subset=["infl_annualised"]).reset_index(drop=True)
excl.append(["1", "drop quarters with missing infl_annualised", n0, len(y_all), n0 - len(y_all),
             "1995Q1 has no previous quarter, so its inflation is undefined"])

origins = pd.date_range("2015-10-01", "2025-07-01", freq="QS")
assert len(origins) == 40
dates = y_all["quarter"].reset_index(drop=True)
y = y_all["infl_annualised"].to_numpy()
qtr = dates.dt.quarter.to_numpy()
pos = {d: i for i, d in enumerate(dates)}


def design(yv, qv, idx, h, p):
    """Rows for direct regression of y[s+h] on target-quarter dummies and y[s],...,y[s-p+1]."""
    X = []
    for s in idx:
        d = [1.0 if qv[s + h] == k else 0.0 for k in (1, 2, 3, 4)] if s + h < len(qv) else None
        X.append((d, [yv[s - j] for j in range(p)]))
    return X


def fit_forecast(t_idx, h):
    """Direct model using only observations with index <= t_idx (quarters <= origin)."""
    yv = y[: t_idx + 1]
    qv = qtr[: t_idx + 1]
    # target quarter of the forecast (deterministic calendar; known at origin)
    tq = ((qv[t_idx] - 1 + h) % 4) + 1
    # common estimation sample for all p: s from PMAX-1 to t_idx-h
    S = np.arange(PMAX - 1, t_idx - h + 1)
    Y = yv[S + h]
    D = np.array([[1.0 if qv[s + h] == k else 0.0 for k in (1, 2, 3, 4)] for s in S])
    best = None
    for p in range(PMAX + 1):
        L = np.column_stack([yv[S - j] for j in range(p)]) if p else np.empty((len(S), 0))
        X = np.hstack([D, L])
        beta, *_ = np.linalg.lstsq(X, Y, rcond=None)
        rss = float(np.sum((Y - X @ beta) ** 2))
        n, k = X.shape
        bic = n * np.log(rss / n) + k * np.log(n)
        if best is None or bic < best[0]:
            best = (bic, p, beta, n)
    _, p, beta, n = best
    x = np.r_[[1.0 if tq == k else 0.0 for k in (1, 2, 3, 4)], [yv[t_idx - j] for j in range(p)]]
    return float(x @ beta), p, n


rows, psel = [], []
for t in origins:
    ti = pos[t]
    for h in HORIZONS:
        tgt = dates[ti + h]
        assert tgt == t + pd.DateOffset(months=3 * h)
        actual = float(y[ti + h])
        sn = float(y[ti + h - 4])  # quarter t+h-4 <= t
        assert dates[ti + h - 4] <= t
        f, p, n = fit_forecast(ti, h)
        psel.append({"origin": t, "h": h, "p": p, "n_reg": n})
        for mod, fc in (("snaive", sn), (MODEL, f)):
            rows.append({"origin_date": t.strftime("%Y-%m-%d"), "target_date": tgt.strftime("%Y-%m-%d"),
                         "horizon": h, "forecast": fc, "actual": actual, "model": mod})
fc = pd.DataFrame(rows)
fc.to_csv(os.path.join(OUT, "forecasts.csv"), index=False, float_format="%.6f")
psel = pd.DataFrame(psel)

last_ti = pos[origins[-1]]
n_last = last_ti + 1
excl.append(["2", "last estimation sample: quarters <= last origin 2025-07-01", len(y_all), n_last,
             len(y_all) - n_last, "information rule: forecasts at origin t use only quarters <= t; "
             "later quarters serve only as targets"])
pd.DataFrame(excl, columns=["step", "rule", "n_before", "n_after", "n_removed", "rationale"]).to_csv(
    os.path.join(OUT, "exclusions.csv"), index=False)

res = {"task_id": "F3-conf-04", "model_name": MODEL,
       "model_description": ("Direct (one regression per horizon) OLS model: infl(s+h) on four target-quarter "
                             "seasonal dummies plus p lags infl(s),...,infl(s-p+1), p in 0..4 chosen by BIC on a "
                             "common sample, re-estimated and re-selected at each origin on an expanding window "
                             "from 1995Q2 using only quarters <= origin."),
       "n_origins": len(origins), "first_origin": origins[0].strftime("%Y-%m-%d"),
       "last_origin": origins[-1].strftime("%Y-%m-%d")}
err = {}
for h in HORIZONS:
    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).to_numpy()
        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}"]


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


d = err[("snaive", 4)] ** 2 - err[("model", 4)] ** 2
T = len(d)
LAGS, KERNEL = 3, "bartlett"
dbar = float(d.mean())
se = float(np.sqrt(lrv(d, LAGS, KERNEL) / T))
dm = dbar / se
res.update({"loss_diff_mean_h4": dbar, "loss_diff_se_h4": se, "hac_lags_h4": LAGS, "hac_kernel_h4": KERNEL,
            "dm_stat_h4": dm, "dm_pvalue_h4": float(2 * stats.norm.sf(abs(dm)))})

# sensitivity of the h=4 test (reported, not primary)
sens = []
for kern, L in (("bartlett", 3), ("bartlett", 4), ("bartlett", 6), ("rectangular", 3)):
    v = lrv(d, L, kern)
    if v > 0:
        s_ = np.sqrt(v / T); st = dbar / s_
        hln = st * np.sqrt((T + 1 - 2 * 4 + 4 * 3 / T) / T) if kern == "rectangular" else np.nan
        sens.append({"kernel": kern, "lags": L, "se": s_, "dm_stat": st, "p_normal": 2 * stats.norm.sf(abs(st)),
                     "hln_stat": hln, "p_hln_t39": 2 * stats.t.sf(abs(hln), T - 1) if kern == "rectangular" else np.nan})
    else:
        sens.append({"kernel": kern, "lags": L, "se": np.nan, "dm_stat": np.nan, "p_normal": np.nan,
                     "hln_stat": np.nan, "p_hln_t39": np.nan})
pd.DataFrame(sens).to_csv(os.path.join(OUT, "sensitivity_dm.csv"), index=False, float_format="%.6f")

res.update({"feasibility": "complete", "feasibility_reason": None,
            "data_issues": [{"issue": "infl_annualised missing for 1995Q1 (no prior quarter)",
                             "handling": "row excluded from all estimation samples"},
                            {"issue": "large genuine outliers (2008Q4, 2020Q2, 2022Q1-Q2) in the target",
                             "handling": "kept as observed; no winsorising; noted in report"}],
            "n_input_rows": {"cpi_quarterly.csv": n_q, "cpi_monthly.csv": n_m},
            "n_analysed": int(n_last), "random_seed": None,
            "supplementary": {
                "n_analysed_definition": "quarters with observed inflation <= last origin (1995Q2-2025Q3); regression obs in last fit are fewer (lags/horizon)",
                "n_reg_last_origin": {f"h{h}": int(psel[(psel.origin == origins[-1]) & (psel.h == h)].n_reg.iloc[0]) for h in HORIZONS},
                "selected_p_counts": {f"h{h}": {str(k): int(v) for k, v in psel[psel.h == h].p.value_counts().sort_index().items()} for h in HORIZONS},
                "mean_loss_diff_h1": float(np.mean(err[("snaive", 1)] ** 2 - err[("model", 1)] ** 2)),
                "mean_loss_diff_h2": float(np.mean(err[("snaive", 2)] ** 2 - err[("model", 2)] ** 2)),
                "dm_sensitivity_h4": [{k: (None if (isinstance(v, float) and np.isnan(v)) else (float(v) if not isinstance(v, str) and not isinstance(v, int) else v)) for k, v in r.items()} for r in sens],
                "loss_diff_mean_h4_excl_2021_2022_targets": float(np.mean(d[~pd.to_datetime(fc[(fc.horizon == 4) & (fc.model == 'snaive')].sort_values('origin_date').target_date).dt.year.isin([2021, 2022]).to_numpy()])),
            }})
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(res, f, indent=2)
for k in [k for k in res if k.startswith(("rmse", "mae", "rel_", "loss", "dm_"))]:
    print(f"{k}: {res[k]:.4f}")