"""F3-conf-01: pseudo-real-time forecasts of Canadian monthly GDP YoY growth.
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
import numpy as np, pandas as pd
from scipy import stats

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

raw = pd.read_csv(DATA)
n_raw = len(raw)
d = raw.copy()
d["month"] = pd.to_datetime(d["month"])
d = d.set_index("month").sort_index()
assert d.index.is_unique and (d.index.to_series().diff().dropna().dt.days <= 31).all()
y = d["gdp_yoy"]; e = d["emp_yoy"]
M = pd.DateOffset(months=1)
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])

def design(t, h):
    """Training rows for the direct model at origin t, horizon h.
    Row indexed by pseudo-origin s: target y(s+h), regressors y(s-2), e(s-1).
    The information set is truncated first: gdp <= t-2, emp <= t-1."""
    gdp_last = t - 2 * M; emp_last = t - 1 * M
    yv = y[y.index <= gdp_last].dropna(); ev = e[e.index <= emp_last].dropna()
    rows = pd.DataFrame({"s": d.index})
    steps = [("0", "candidate pseudo-origins s = every month in file", len(rows), len(rows), "all monthly rows")]
    rows["tgt"] = rows.s + pd.DateOffset(months=h); rows["ylag"] = rows.s - 2 * M; rows["elag"] = rows.s - M
    nb = len(rows); rows = rows[rows.tgt.isin(yv.index)]
    steps.append(("1", f"target y(s+{h}) observed at origin (<= {gdp_last.date()})", nb, len(rows), "training targets must be published at origin"))
    nb = len(rows); rows = rows[rows.ylag.isin(yv.index) & rows.elag.isin(ev.index)]
    steps.append(("2", "lagged regressors y(s-2), e(s-1) exist", nb, len(rows), "start of sample"))
    nb = len(rows); rows = rows[~(in_covid(rows.tgt) | in_covid(rows.ylag) | in_covid(rows.elag))]
    steps.append(("3", "drop rows with target or regressor month in 2020-03..2021-12", nb, len(rows), "pandemic outliers would dominate OLS (pre-specified)"))
    Y = yv.loc[rows.tgt].values
    X = np.column_stack([np.ones(len(rows)), yv.loc[rows.ylag].values, ev.loc[rows.elag].values])
    x_new = np.array([1.0, yv.loc[gdp_last], ev.loc[emp_last]])
    return X, Y, x_new, rows, steps

recs = []; last_info = {}
for t in origins:
    for h in H:
        tgt = t + pd.DateOffset(months=h)
        actual = float(y.loc[tgt])
        gdp_last = t - 2 * M
        recs.append(dict(origin_date=t.date().isoformat(), target_date=tgt.date().isoformat(), horizon=h,
                         forecast=float(y.loc[gdp_last]), actual=actual, model="naive",
                         last_target_date_used=gdp_last.date().isoformat(), last_predictor_date_used=""))
        X, Y, xn, rows, steps = design(t, h)
        beta, *_ = np.linalg.lstsq(X, Y, rcond=None)
        last_t = max(rows.tgt.max(), rows.ylag.max(), gdp_last)
        last_e = max(rows.elag.max(), t - M)
        assert last_t <= t - 2 * M and last_e <= t - M
        recs.append(dict(origin_date=t.date().isoformat(), target_date=tgt.date().isoformat(), horizon=h,
                         forecast=float(xn @ beta), actual=actual, model=MODEL,
                         last_target_date_used=last_t.date().isoformat(), last_predictor_date_used=last_e.date().isoformat()))
        last_info[h] = (len(Y), steps, beta)

fc = pd.DataFrame(recs)
fc["forecast"] = fc.forecast.round(6)
fc.to_csv(f"{OUT}/forecasts.csv", index=False)

def err(m, h):
    s = fc[(fc.model == m) & (fc.horizon == h)]
    u = s.forecast - s.actual
    return float(np.sqrt((u ** 2).mean())), float(u.abs().mean())

res = {"task_id": "F3-conf-01", "model_name": MODEL,
       "model_description": ("Direct OLS per horizon h re-estimated at each origin t on an expanding window: "
           "gdp_yoy(s+h) = a + b*gdp_yoy(s-2) + c*emp_yoy(s-1), using all pseudo-origins s with s+h <= t-2 "
           "(1998 onward, excluding rows whose target or regressor month lies in 2020-03..2021-12); "
           "the forecast plugs in gdp_yoy(t-2) and emp_yoy(t-1)."),
       "n_origins": len(origins), "first_origin": origins[0].date().isoformat(), "last_origin": origins[-1].date().isoformat(),
       "last_target_date_available": (origins[-1] - 2 * M).date().isoformat(),
       "last_predictor_date_available": (origins[-1] - M).date().isoformat()}
for m, k in [("naive", "naive"), (MODEL, "model")]:
    for h in H:
        r, a = err(m, h); res[f"rmse_{k}_h{h}"] = round(r, 6); res[f"mae_{k}_h{h}"] = round(a, 6)
for h in H:
    res[f"rel_improvement_rmse_h{h}"] = round(1 - res[f"rmse_model_h{h}"] / res[f"rmse_naive_h{h}"], 6)
res["feasibility"] = "complete"; res["feasibility_reason"] = None
res["data_issues"] = [
    {"issue": "gdp_yoy and emp_yoy are populated for 1998-01..1998-12 although data/README.md says they are empty for the first 12 months; the 1997 levels needed to verify them are not in the file.",
     "handling": "Used as given (values from 1999-01 on were verified against the levels to within 5e-5 rounding error)."},
    {"issue": "GDP columns missing for 2026-08 and 2026-09 (not yet published).",
     "handling": "Not needed: latest target is 2026-07 (origin 2026-04, h=3); no rows dropped."},
    {"issue": "Pandemic outliers (GDP YoY between -16.2% and +18.5% in 2020-03..2021-12).",
     "handling": "Pre-specified exclusion of estimation rows whose target or regressor month is in 2020-03..2021-12; evaluation targets (2022-06 onward) are unaffected."},
    {"issue": "Data are the current vintage, not first releases.",
     "handling": "Exercise simulates publication timing only; no real-time validity claimed."}]
res["n_input_rows"] = {"canada_gdp_employment.csv": n_raw}
res["n_analysed"] = last_info[1][0]
res["n_analysed_h3"] = last_info[3][0]
res["last_origin_coefficients"] = {f"h{h}": {"const": round(float(b[0]), 6), "gdp_yoy_lag2": round(float(b[1]), 6), "emp_yoy_lag1": round(float(b[2]), 6)} for h, (_, _, b) in last_info.items()}
res["random_seed"] = None
# Diebold-Mariano (squared loss, rectangular HAC with h-1 lags, Harvey-Leybourne-Newbold adjustment, t(n-1))
dm = {}
for h in H:
    a = fc[(fc.model == "naive") & (fc.horizon == h)].sort_values("origin_date")
    b = fc[(fc.model == MODEL) & (fc.horizon == h)].sort_values("origin_date")
    dl = (a.forecast.values - a.actual.values) ** 2 - (b.forecast.values - b.actual.values) ** 2
    n = len(dl); dc = dl - dl.mean()
    v = dc @ dc / n
    for k in range(1, h):
        v += 2 * (dc[k:] @ dc[:-k]) / n
    st = float(dl.mean() / np.sqrt(v / n) * np.sqrt((n + 1 - 2 * h + h * (h - 1) / n) / n))
    dm[f"h{h}"] = {"dm_stat": round(st, 4), "p_value_two_sided": round(float(2 * stats.t.sf(abs(st), n - 1)), 4)}
res["dm_test_naive_vs_model"] = dm
with open(f"{OUT}/results.json", "w") as f:
    json.dump(res, f, indent=2)

ex = []
for h in H:
    for st in last_info[h][1]:
        ex.append(dict(step=f"h{h}_last_origin_{st[0]}", rule=st[1], n_before=st[2], n_after=st[3], n_removed=st[2] - st[3], rationale=st[4]))
pd.DataFrame(ex).to_csv(f"{OUT}/exclusions.csv", index=False)
print(json.dumps({k: res[k] for k in res if k.startswith(("rmse", "mae", "rel", "n_analysed"))}, indent=1))