"""F3-conf-01: GDP y/y 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

DATA = "data/canada_gdp_employment.csv"
OUT = "out"
MODEL = "dir_ols_emp"
COVID = (pd.Timestamp("2020-03-01"), pd.Timestamp("2021-12-01"))

d = pd.read_csv(DATA, parse_dates=["month"]).sort_values("month").reset_index(drop=True)
n_raw = len(d)
assert d.month.is_unique
d = d.set_index("month")
M = lambda t, k: t + pd.DateOffset(months=k)
def ds(t): return t.strftime("%Y-%m-%d")
def in_covid(m): return (m >= COVID[0]) & (m <= COVID[1])

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

def build_rows(t, h):
    """Direct regression rows y_s ~ 1 + y_{s-h-2} + x_{s-h-1} using only info at origin t."""
    gdp = d.loc[d.index <= M(t, -2), "gdp_yoy"].dropna()   # GDP known up to t-2
    emp = d.loc[d.index <= M(t, -1), "emp_yoy"].dropna()   # emp known up to t-1
    s = gdp.index
    df = pd.DataFrame({"y": gdp.values,
                       "ylag": gdp.reindex(s - pd.DateOffset(months=h + 2)).values,
                       "xlag": emp.reindex(s - pd.DateOffset(months=h + 1)).values}, index=s)
    n_avail = len(df)
    df = df.dropna()
    n_lagged = len(df)
    bad = in_covid(df.index) | in_covid(df.index - pd.DateOffset(months=h + 2)) | in_covid(df.index - pd.DateOffset(months=h + 1))
    df = df[~bad]
    return df, gdp, emp, n_avail, n_lagged

rows = []
last_info = {}
for t in origins:
    for h in (1, 3):
        tgt = M(t, h)
        actual = float(d.loc[tgt, "gdp_yoy"])
        df, gdp, emp, n_avail, n_lagged = build_rows(t, h)
        y_last, x_last = gdp.index.max(), emp.index.max()
        assert y_last <= M(t, -2) and x_last <= M(t, -1)
        # naive
        rows.append(dict(origin_date=ds(t), target_date=ds(tgt), horizon=h, forecast=float(gdp.loc[M(t, -2)]),
                         actual=actual, model="naive", last_target_date_used=ds(M(t, -2)), last_predictor_date_used=""))
        # model
        X = np.column_stack([np.ones(len(df)), df.ylag.values, df.xlag.values])
        beta, *_ = np.linalg.lstsq(X, df.y.values, rcond=None)
        xnew = np.array([1.0, gdp.loc[M(t, -2)], emp.loc[M(t, -1)]])
        fc = float(xnew @ beta)
        rows.append(dict(origin_date=ds(t), target_date=ds(tgt), horizon=h, forecast=fc, actual=actual, model=MODEL,
                         last_target_date_used=ds(max(df.index.max(), M(t, -2))),
                         last_predictor_date_used=ds(M(t, -1))))
        last_info[h] = dict(n=len(df), n_avail=n_avail, n_lagged=n_lagged, beta=beta.tolist())

fc = pd.DataFrame(rows)
fc.to_csv(os.path.join(OUT, "forecasts.csv"), index=False, float_format="%.6f")

def err(model, h):
    e = fc[(fc.model == model) & (fc.horizon == h)]
    e = e.forecast - e.actual
    assert len(e) == 48
    return float(np.sqrt(np.mean(e ** 2))), float(np.mean(np.abs(e)))

res = {"task_id": "F3-conf-01", "model_name": MODEL,
       "model_description": ("Direct h-step OLS of gdp_yoy(s) on a constant, gdp_yoy(s-h-2) and emp_yoy(s-h-1), "
                             "re-estimated at each origin on an expanding window from 1998 with targets <= t-2, "
                             "excluding rows whose target or regressor month lies in 2020-03..2021-12; "
                             "forecast plugs in gdp_yoy(t-2) and emp_yoy(t-1)."),
       "n_origins": int(len(origins)), "first_origin": ds(origins[0]), "last_origin": ds(origins[-1]),
       "last_target_date_available": ds(M(origins[-1], -2)),
       "last_predictor_date_available": ds(M(origins[-1], -1))}
for m, tag in (("naive", "naive"), (MODEL, "model")):
    for h in (1, 3):
        r, a = err(m, h)
        res[f"rmse_{tag}_h{h}"], res[f"mae_{tag}_h{h}"] = r, a
for h in (1, 3):
    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": "gdp_level/gdp_yoy missing for 2026-08 and 2026-09 (not yet published)",
                 "handling": "Not needed: latest target is 2026-07; rows never used as GDP information."},
                {"issue": "gdp_yoy and emp_yoy for 1998 depend on 1997 levels absent from the file and cannot be re-verified (all later yoy values match levels exactly)",
                 "handling": "Accepted as given."},
                {"issue": "COVID-19 shock 2020-03..2021-12 with yoy values between -16 and +19 percent (extreme leverage)",
                 "handling": "Regression rows touching this window excluded from model estimation (pre-specified); naive unaffected."},
                {"issue": "Current-vintage (revised) data, not historical first releases",
                 "handling": "Only publication timing simulated; no real-time validity claimed."}],
            "n_input_rows": {"canada_gdp_employment.csv": int(n_raw)},
            "n_analysed": int(last_info[1]["n"]), "random_seed": None})
# Uncertainty: Diebold-Mariano (squared loss, HAC variance with h-1 Newey-West lags, Harvey et al. small-sample correction)
from scipy import stats
for h in (1, 3):
    a = fc[(fc.model == "naive") & (fc.horizon == h)].sort_values("origin_date")
    b = fc[(fc.model == MODEL) & (fc.horizon == h)].sort_values("origin_date")
    dd = (a.forecast.values - a.actual.values) ** 2 - (b.forecast.values - b.actual.values) ** 2
    n = len(dd); dm = dd - dd.mean()
    lrv = np.mean(dm ** 2) + 2 * sum((1 - k / h) * np.mean(dm[k:] * dm[:-k]) for k in range(1, h))
    stat = dd.mean() / np.sqrt(lrv / n) * np.sqrt((n + 1 - 2 * h + h * (h - 1) / n) / n)
    res[f"dm_stat_h{h}"] = float(stat)
    res[f"dm_pvalue_two_sided_h{h}"] = float(2 * stats.t.sf(abs(stat), df=n - 1))
# Ablation (diagnostic only): same direct regression without the employment regressor
for h in (1, 3):
    e = []
    for t in origins:
        df, gdp, emp, *_ = build_rows(t, h)
        X = np.column_stack([np.ones(len(df)), df.ylag.values])
        beta, *_ = np.linalg.lstsq(X, df.y.values, rcond=None)
        e.append(beta[0] + beta[1] * gdp.loc[M(t, -2)] - d.loc[M(t, h), "gdp_yoy"])
    res[f"rmse_ablation_no_emp_h{h}"] = float(np.sqrt(np.mean(np.square(e))))
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(res, f, indent=2)

# exclusion ledger for the last estimation sample (origin 2026-04, h=1)
li = last_info[1]
n_info = int((d.index <= M(origins[-1], -2)).sum())
steps = [("info_filter", "keep months <= t-2 (GDP published) at last origin 2026-04-01", n_raw, n_info,
          "publication-lag rule"),
         ("lag_availability", "drop rows lacking gdp_yoy(s-3) or emp_yoy(s-2) (start of sample)", n_info, li["n_lagged"],
          "regressors not in file"),
         ("covid_window", "drop rows with target or regressor month in 2020-03..2021-12", li["n_lagged"], li["n"],
          "pandemic outliers; pre-specified")]
assert li["n_avail"] == n_info
pd.DataFrame([dict(step=s, rule=r, n_before=a, n_after=b, n_removed=a - b, rationale=q) for s, r, a, b, q in steps]
             ).to_csv(os.path.join(OUT, "exclusions.csv"), index=False)
print(json.dumps({k: v for k, v in res.items() if k.startswith(("rmse", "mae", "rel"))}, indent=1))
print("last-origin coefficients [const, gdp lag, emp lag]: h1", np.round(last_info[1]["beta"], 4).tolist(),
      "h3", np.round(last_info[3]["beta"], 4).tolist())