"""F3-conf-01: direct forecasts of Canadian monthly GDP yoy growth under publication lags.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py"""
import json, numpy as np, pandas as pd

DATA = "data/canada_gdp_employment.csv"
d = pd.read_csv(DATA, parse_dates=["month"])
n_rows = len(d)
d = d.set_index("month").asfreq("MS")
assert len(d) == n_rows  # no gaps
MODEL = "ardl_emp_direct"
COVID = (pd.Timestamp("2020-03-01"), pd.Timestamp("2021-12-01"))
HS = [1, 3]
origins = pd.date_range("2022-05-01", "2026-04-01", freq="MS")
assert len(origins) == 48

def in_covid(m):
    return (m >= COVID[0]) & (m <= COVID[1])

rows, n_last = [], {}
for t in origins:
    g_last = t - pd.DateOffset(months=2)   # latest GDP month published at end of t
    e_last = t - pd.DateOffset(months=1)   # latest LFS month published at end of t
    y = d["gdp_yoy"].where(d.index <= g_last)      # mask unavailable info
    e = d["emp_yoy"].where(d.index <= e_last)
    for h in HS:
        target = t + pd.DateOffset(months=h)
        actual = d["gdp_yoy"].get(target, np.nan)
        # naive
        rows.append(dict(origin_date=t, target_date=target, horizon=h, forecast=y[g_last],
                         actual=actual, model="naive", last_target_date_used=g_last,
                         last_predictor_date_used=pd.NaT))
        # direct regression: y_s = a + b*y_{s-h-2} + c*e_{s-h-1}, all with s <= g_last
        tr = pd.DataFrame({"y": y, "ylag": y.shift(h + 2), "elag": e.shift(h + 1)})
        idx = tr.index
        bad = in_covid(idx) | in_covid(idx - pd.DateOffset(months=h + 2)) | in_covid(idx - pd.DateOffset(months=h + 1))
        tr = tr[~bad].dropna()
        assert tr.index.max() <= g_last
        X = np.column_stack([np.ones(len(tr)), tr["ylag"], tr["elag"]])
        beta, *_ = np.linalg.lstsq(X, tr["y"].values, rcond=None)
        assert not np.isnan(e[e_last]) and not np.isnan(y[g_last])
        fc = beta @ np.array([1.0, y[g_last], e[e_last]])
        rows.append(dict(origin_date=t, target_date=target, horizon=h, forecast=fc,
                         actual=actual, model=MODEL, last_target_date_used=g_last,
                         last_predictor_date_used=e_last))
        if t == origins[-1]:
            n_last[h] = len(tr)

f = pd.DataFrame(rows)
for c in ["origin_date", "target_date", "last_target_date_used", "last_predictor_date_used"]:
    f[c] = f[c].dt.strftime("%Y-%m-%d")
f = f.fillna({"last_predictor_date_used": ""})
f["forecast"] = f["forecast"].round(6)
f.to_csv("out/forecasts.csv", index=False)

res = {"task_id": "F3-conf-01", "model_name": MODEL,
       "model_description": ("Direct OLS per horizon h: gdp_yoy(s) = a + b*gdp_yoy(s-h-2) + c*emp_yoy(s-h-1), "
                             "expanding window from 1998 using only targets observed at the origin (s <= t-2), "
                             "dropping training rows whose target or regressor month lies in 2020-03..2021-12 (COVID); "
                             "forecast plugs in gdp_yoy(t-2) and emp_yoy(t-1)."),
       "n_origins": len(origins), "first_origin": "2022-05-01", "last_origin": "2026-04-01",
       "last_target_date_available": (origins[-1] - pd.DateOffset(months=2)).strftime("%Y-%m-%d"),
       "last_predictor_date_available": (origins[-1] - pd.DateOffset(months=1)).strftime("%Y-%m-%d")}
stats = {}
for m, tag in [("naive", "naive"), (MODEL, "model")]:
    for h in HS:
        s = f[(f.model == m) & (f.horizon == h)]
        err = s.forecast - s.actual
        assert err.notna().sum() == 48
        res[f"rmse_{tag}_h{h}"] = float(np.sqrt(np.mean(err ** 2)))
        res[f"mae_{tag}_h{h}"] = float(np.mean(np.abs(err)))
for h in HS:
    res[f"rel_improvement_rmse_h{h}"] = 1 - res[f"rmse_model_h{h}"] / res[f"rmse_naive_h{h}"]
# Diebold-Mariano (squared loss, HAC with h-1 lags, Harvey small-sample correction) - supplementary
from scipy import stats as st
dm = {}
for h in HS:
    a = f[(f.model == "naive") & (f.horizon == h)].reset_index(drop=True)
    b = f[(f.model == MODEL) & (f.horizon == h)].reset_index(drop=True)
    dl = ((b.forecast - b.actual) ** 2 - (a.forecast - a.actual) ** 2).values
    n = len(dl); dbar = dl.mean(); u = dl - dbar
    lr = u @ u / n + 2 * sum((1 - k / h) * (u[k:] @ u[:-k]) / n for k in range(1, h))
    stat = dbar / np.sqrt(lr / n) * np.sqrt((n + 1 - 2 * h + h * (h - 1) / n) / n)
    dm[f"h{h}"] = {"dm_stat": float(stat), "p_two_sided": float(2 * st.t.sf(abs(stat), n - 1))}
res.update({"feasibility": "complete", "feasibility_reason": None,
            "data_issues": [
                {"issue": "COVID-19 outliers in 2020-2021 (gdp_yoy from -16.2% to +18.5%, emp_yoy similar) would dominate OLS fits.",
                 "handling": "Training rows with target or regressor month in 2020-03..2021-12 excluded from model estimation; evaluation sample (targets 2022-06..2026-07) unaffected; naive unaffected."},
                {"issue": "Data are the current (final) vintage, not first releases; yoy columns match levels to rounding (max abs diff ~5e-5).",
                 "handling": "Used as given; results simulate publication timing only, no real-time validity claimed."},
                {"issue": "GDP columns empty for 2026-08 and 2026-09.",
                 "handling": "Not needed: latest target month is 2026-07; GDP never used beyond t-2."}],
            "n_input_rows": {"canada_gdp_employment.csv": n_rows},
            "n_analysed": int(n_last[1]),
            "n_analysed_note": f"training rows at last origin: h=1 {n_last[1]}, h=3 {n_last[3]} (n_analysed reports h=1)", "random_seed": None,
            "supplementary_dm_test": dm})
json.dump(res, open("out/results.json", "w"), indent=2)
print(json.dumps(res, indent=2))