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

DATA = "data/canada_gdp_employment.csv"
MODEL = "ardl_direct_emp"
HORIZONS = [1, 3]
# Pre-specified (not tuned on the evaluation sample): drop training rows touching
# the COVID collapse/rebound months, whose y/y values (-16% .. +18.5%) dominate OLS.
COVID_START, COVID_END = pd.Timestamp("2020-03-01"), pd.Timestamp("2021-06-01")

raw = pd.read_csv(DATA)
n_rows = len(raw)
d = raw.copy()
d["month"] = pd.to_datetime(d["month"])
d = d.set_index("month").sort_index()
assert d.index.is_unique
assert (d.index.to_series().diff().dropna().dt.days.between(28, 31)).all()

# data checks
issues = []
recomp_g = 100 * (d.gdp_level / d.gdp_level.shift(12) - 1)
recomp_e = 100 * (d.emp_level / d.emp_level.shift(12) - 1)
mg = (recomp_g - d.gdp_yoy).abs().max()
me = (recomp_e - d.emp_yoy).abs().max()
issues.append({"issue": f"yoy columns are rounded to 4 decimals; recomputation from levels (where 12-month lags exist in file) differs by at most {mg:.2e} (GDP) / {me:.2e} (employment) pp",
               "handling": "used the provided gdp_yoy/emp_yoy columns as given"})
issues.append({"issue": "gdp_level/gdp_yoy missing for 2026-08 and 2026-09 (not yet published)",
               "handling": "not needed: last target month used is 2026-07 (origin 2026-04, h=3); availability rule applied anyway"})
issues.append({"issue": "COVID-19 outliers: gdp_yoy ranges from -16.2% (2020-04) to +18.5% (2021-04), emp_yoy similarly",
               "handling": "training rows whose target or predictor month lies in 2020-03..2021-06 are excluded from model estimation (pre-specified); naive forecasts unaffected; no evaluation target falls in that window"})
issues.append({"issue": "file is the current (revised) vintage, not real-time first releases",
               "handling": "publication timing simulated only; no real-time validity claimed"})

y = d["gdp_yoy"]
e = d["emp_yoy"]
months = d.index

def shift(m, k):
    return m + pd.DateOffset(months=k)

def in_covid(m):
    return COVID_START <= m <= COVID_END

origins = pd.date_range("2022-05-01", "2026-04-01", freq="MS")
assert len(origins) == 48

rows = []
n_last = {}
for t in origins:
    g_last = shift(t, -2)  # latest GDP month available
    e_last = shift(t, -1)  # latest employment month available
    y_av = y[(y.index <= g_last)].dropna()
    e_av = e[(e.index <= e_last)].dropna()
    assert y_av.index.max() == g_last and e_av.index.max() == e_last
    for h in HORIZONS:
        target = shift(t, h)
        actual = float(y.loc[target])
        # naive
        rows.append(dict(origin_date=t, target_date=target, horizon=h,
                         forecast=float(y_av.loc[g_last]), actual=actual, model="naive",
                         last_target_date_used=g_last, last_predictor_date_used=None))
        # direct regression: y_{s+h} = a + b*y_{s-2} + c*e_{s-1}, s = pseudo-origin
        X, Y = [], []
        for s in months:
            ty, py, pe = shift(s, h), shift(s, -2), shift(s, -1)
            if ty > g_last:  # training target must be observed at origin t
                continue
            if py not in y_av.index or pe not in e_av.index or ty not in y_av.index:
                continue
            if any(in_covid(m) for m in (ty, py, pe)):
                continue
            X.append([1.0, y_av.loc[py], e_av.loc[pe]])
            Y.append(y_av.loc[ty])
        X, Y = np.array(X), np.array(Y)
        beta, *_ = np.linalg.lstsq(X, Y, rcond=None)
        xf = np.array([1.0, y_av.loc[g_last], e_av.loc[e_last]])
        n_last[h] = len(Y)
        rows.append(dict(origin_date=t, target_date=target, horizon=h,
                         forecast=float(xf @ beta), actual=actual, model=MODEL,
                         last_target_date_used=g_last, last_predictor_date_used=e_last))

fc = pd.DataFrame(rows)
fc["error"] = fc["forecast"] - fc["actual"]
out = fc.drop(columns="error").copy()
for c in ["origin_date", "target_date", "last_target_date_used"]:
    out[c] = out[c].dt.strftime("%Y-%m-%d")
out["last_predictor_date_used"] = out["last_predictor_date_used"].apply(
    lambda v: "" if v is None or pd.isna(v) else pd.Timestamp(v).strftime("%Y-%m-%d"))
out.to_csv("out/forecasts.csv", index=False, float_format="%.6f")
assert len(out) == 192

def metrics(m, h):
    err = fc[(fc.model == m) & (fc.horizon == h)]["error"].values
    return float(np.sqrt(np.mean(err ** 2))), float(np.mean(np.abs(err)))

def dm_test(h):
    """Diebold-Mariano (squared loss). Effective forecast distance from last observed GDP month is h+2,
    so Newey-West HAC uses h+1 lags and the HLN correction uses h+2."""
    from scipy import stats
    a = fc[(fc.model == MODEL) & (fc.horizon == h)].sort_values("origin_date")["error"].values
    b = fc[(fc.model == "naive") & (fc.horizon == h)].sort_values("origin_date")["error"].values
    dl = a ** 2 - b ** 2
    T = len(dl); L = h + 1; he = h + 2
    dc = dl - dl.mean()
    lrv = dc @ dc / T
    for k in range(1, L + 1):
        lrv += 2 * (1 - k / (L + 1)) * (dc[k:] @ dc[:-k]) / T
    stat = dl.mean() / np.sqrt(lrv / T)
    hln = np.sqrt((T + 1 - 2 * he + he * (he - 1) / T) / T)
    stat *= hln
    p = 2 * stats.t.sf(abs(stat), df=T - 1)
    return float(dl.mean()), float(stat), float(p)

res = {"task_id": "F3-conf-01", "model_name": MODEL,
       "model_description": ("Direct h-step OLS of gdp_yoy(t+h) on a constant, gdp_yoy(t-2) and emp_yoy(t-1), "
                             "re-estimated at each origin on an expanding window of all pseudo-origins since 1998 whose "
                             "target month is <= t-2, excluding rows touching 2020-03..2021-06 (COVID)."),
       "n_origins": 48, "first_origin": "2022-05-01", "last_origin": "2026-04-01",
       "last_target_date_available": "2026-02-01", "last_predictor_date_available": "2026-03-01"}
for h in HORIZONS:
    r, a = metrics("naive", h); res[f"rmse_naive_h{h}"] = r; res[f"mae_naive_h{h}"] = a
    r, a = metrics(MODEL, h); res[f"rmse_model_h{h}"] = r; res[f"mae_model_h{h}"] = a
for h in HORIZONS:
    res[f"rel_improvement_rmse_h{h}"] = 1 - res[f"rmse_model_h{h}"] / res[f"rmse_naive_h{h}"]
res.update({"feasibility": "complete", "feasibility_reason": None, "data_issues": issues,
            "n_input_rows": {"canada_gdp_employment.csv": n_rows},
            "n_analysed": int(n_last[1]), "random_seed": None})
# supplementary diagnostics (extra keys)
res["n_analysed_h3_last_origin"] = int(n_last[3])
for h in HORIZONS:
    md, s, p = dm_test(h)
    res[f"dm_mean_lossdiff_h{h}"] = md; res[f"dm_stat_h{h}"] = s; res[f"dm_pvalue_h{h}"] = p
    sub = fc[(fc.horizon == h)]
    res[f"mean_error_naive_h{h}"] = float(sub[sub.model == "naive"].error.mean())
    res[f"mean_error_model_h{h}"] = float(sub[sub.model == MODEL].error.mean())
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps({k: v for k, v in res.items() if k != "data_issues"}, indent=2))