"""F3-conf-01: Canadian monthly GDP y/y growth forecasts under publication lags.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Reads only data/; writes out/results.json, out/forecasts.csv, out/exclusions.csv. Deterministic (no randomness)."""
import json, os
import numpy as np, pandas as pd
from scipy import stats

DATA = "data/canada_gdp_employment.csv"
OUT = "out"
MODEL = "ardl_emp_direct"
COVID = (pd.Timestamp("2020-03-01"), pd.Timestamp("2021-12-01"))  # pre-specified training exclusion window
HORIZONS = (1, 3)

raw = pd.read_csv(DATA, parse_dates=["month"])
n_raw = len(raw)
d = raw.set_index("month").sort_index()
assert d.index.is_unique
origins = pd.date_range("2022-05-01", "2026-04-01", freq="MS")
assert len(origins) == 48
M = lambda ts, k: ts + pd.DateOffset(months=k)

def design(gdp, emp, h, exclude_covid):
    """Direct regression rows: target y(m) on [1, y(m-h-2), emp(m-h-1)], using only the masked series."""
    rows = []
    for m in gdp.dropna().index:
        ml_y, ml_e = M(m, -(h + 2)), M(m, -(h + 1))
        if ml_y not in gdp.index or ml_e not in emp.index or pd.isna(gdp.get(ml_y)) or pd.isna(emp.get(ml_e)):
            continue
        if exclude_covid and any(COVID[0] <= x <= COVID[1] for x in (m, ml_y, ml_e)):
            continue
        rows.append((m, gdp[m], gdp[ml_y], emp[ml_e]))
    return pd.DataFrame(rows, columns=["m", "y", "ylag", "elag"]).set_index("m")

def fit_forecast(t, h, exclude_covid=True):
    gdp = d["gdp_yoy"][d.index <= M(t, -2)]          # GDP available: months <= t-2
    emp = d["emp_yoy"][d.index <= M(t, -1)]          # LFS available: months <= t-1
    X = design(gdp, emp, h, exclude_covid)
    A = np.column_stack([np.ones(len(X)), X["ylag"], X["elag"]])
    beta, *_ = np.linalg.lstsq(A, X["y"].values, rcond=None)
    x_new = np.array([1.0, gdp[M(t, -2)], emp[M(t, -1)]])
    return float(x_new @ beta), X, gdp.index.max(), emp.index.max()

recs, sens, n_last = [], [], {}
for t in origins:
    for h in HORIZONS:
        tgt = M(t, h)
        actual = float(d.loc[tgt, "gdp_yoy"])
        assert not np.isnan(actual)
        last_g = M(t, -2)
        recs.append(dict(origin_date=t, target_date=tgt, horizon=h, forecast=float(d.loc[last_g, "gdp_yoy"]),
                         actual=actual, model="naive", last_target_date_used=last_g, last_predictor_date_used=pd.NaT))
        f, X, lg, le = fit_forecast(t, h, True)
        assert lg <= M(t, -2) and le <= M(t, -1) and X.index.max() <= M(t, -2)
        recs.append(dict(origin_date=t, target_date=tgt, horizon=h, forecast=f, actual=actual, model=MODEL,
                         last_target_date_used=lg, last_predictor_date_used=le))
        fs, *_ = fit_forecast(t, h, False)
        sens.append(dict(h=h, e=fs - actual))
        if t == origins[-1]:
            n_last[h] = len(X)
fc = pd.DataFrame(recs)
fc["error"] = fc["forecast"] - fc["actual"]

def rmse(e): return float(np.sqrt(np.mean(np.square(e))))
def mae(e): return float(np.mean(np.abs(e)))

def dm_test(e1, e2, h):
    """Diebold-Mariano (squared loss), HAC variance with h-1 Bartlett-free rectangular lags, Harvey et al. correction."""
    dl = e1 ** 2 - e2 ** 2
    n = len(dl); dbar = dl.mean(); dc = dl - dbar
    g = [np.sum(dc[k:] * dc[:n - k]) / n for k in range(h)]
    v = (g[0] + 2 * sum(g[1:])) / n
    if v <= 0:
        v = g[0] / n
    stat = dbar / np.sqrt(v) * np.sqrt((n + 1 - 2 * h + h * (h - 1) / n) / n)
    return float(stat), float(2 * stats.t.sf(abs(stat), df=n - 1))

res = {"task_id": "F3-conf-01", "model_name": MODEL,
       "model_description": ("Direct h-step OLS regression per horizon and origin of gdp_yoy(m) on a constant, gdp_yoy(m-h-2) "
                             "and emp_yoy(m-h-1), estimated on an expanding window from 1998 using only training pairs whose target "
                             "month is <= t-2 (and employment <= t-1), excluding pairs touching 2020-03..2021-12 (COVID); "
                             "forecast plugs in gdp_yoy(t-2) and emp_yoy(t-1)."),
       "n_origins": len(origins), "first_origin": str(origins[0].date()), "last_origin": str(origins[-1].date()),
       "last_target_date_available": str(M(origins[-1], -2).date()),
       "last_predictor_date_available": str(M(origins[-1], -1).date())}
sub = {}
for h in HORIZONS:
    en = fc[(fc.horizon == h) & (fc.model == "naive")].sort_values("origin_date")["error"].values
    em = fc[(fc.horizon == h) & (fc.model == MODEL)].sort_values("origin_date")["error"].values
    res[f"rmse_naive_h{h}"], res[f"mae_naive_h{h}"] = rmse(en), mae(en)
    res[f"rmse_model_h{h}"], res[f"mae_model_h{h}"] = rmse(em), mae(em)
    sub[h] = (en, em)
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": [
                {"issue": "README states gdp_yoy is empty for the first 12 months, but gdp_yoy and emp_yoy are populated from 1998-01 (presumably from pre-1998 levels not in the file).",
                 "handling": "Used the provided yoy values as given; they match recomputation from levels to <=5e-5 where verifiable."},
                {"issue": "gdp_level/gdp_yoy missing for 2026-08 and 2026-09 (not yet published).",
                 "handling": "Expected; never needed (last target month 2026-07, last GDP used 2026-02)."},
                {"issue": "Extreme COVID-19 swings 2020-03..2021-06 and base effects through 2021 in both series.",
                 "handling": "Pre-specified exclusion of training pairs touching 2020-03..2021-12 from the model; sensitivity without exclusion reported (sens_* keys). Naive baseline unaffected."},
                {"issue": "Data are the current vintage (revised), not first releases.",
                 "handling": "Exercise simulates publication timing only; no real-time validity claimed."}],
            "n_input_rows": {"canada_gdp_employment.csv": n_raw},
            "n_analysed": int(n_last[1]), "n_analysed_last_origin_h3": int(n_last[3]),
            "random_seed": None})
for h in HORIZONS:
    en, em = sub[h]
    s, p = dm_test(em, en, h)
    res[f"dm_stat_h{h}"], res[f"dm_pvalue_h{h}"] = s, p
    se = np.array([r["e"] for r in sens if r["h"] == h])
    res[f"sens_rmse_model_nocovidexcl_h{h}"] = rmse(se)
    res[f"sens_rel_improvement_rmse_nocovidexcl_h{h}"] = 1 - rmse(se) / res[f"rmse_naive_h{h}"]
res = {k: (round(v, 10) if isinstance(v, float) else v) for k, v in res.items()}

os.makedirs(OUT, exist_ok=True)
out = fc.drop(columns="error").copy()
for c in ["origin_date", "target_date", "last_target_date_used", "last_predictor_date_used"]:
    out[c] = out[c].dt.strftime("%Y-%m-%d")
out["last_predictor_date_used"] = out["last_predictor_date_used"].fillna("")
out["forecast"] = out["forecast"].round(8)
out.to_csv(f"{OUT}/forecasts.csv", index=False)
with open(f"{OUT}/results.json", "w") as fh:
    json.dump(res, fh, indent=2)

# Exclusion ledger for the last estimation sample (last origin, h=1)
t = origins[-1]; h = 1
n1 = int((d.index <= M(t, -2)).sum())
gdp = d["gdp_yoy"][d.index <= M(t, -2)]; emp = d["emp_yoy"][d.index <= M(t, -1)]
n2 = len(design(gdp, emp, h, False)); n3 = n_last[1]
led = pd.DataFrame([
    ("1", "GDP month <= t-2 at last origin 2026-04 (publication lag); drops 2026-03..2026-09", n_raw, n1, "information-availability rule"),
    ("2", f"target month needs gdp_yoy(m-{h+2}) and emp_yoy(m-{h+1}) inside the file", n1, n2, "lags unavailable for first months of file"),
    ("3", "drop pairs with target/regressor month in 2020-03..2021-12", n2, n3, "pre-specified COVID outlier exclusion (model only)")],
    columns=["step", "rule", "n_before", "n_after", "rationale"])
led.insert(4, "n_removed", led.n_before - led.n_after)
led = led[["step", "rule", "n_before", "n_after", "n_removed", "rationale"]]
led.to_csv(f"{OUT}/exclusions.csv", index=False)
assert n_raw - led.n_removed.sum() == res["n_analysed"]
print(json.dumps({k: res[k] for k in res if k.startswith(("rmse", "mae", "rel", "dm", "sens", "n_an"))}, indent=1))