"""F3-conf-01: pseudo-real-time forecasts of Canadian monthly GDP YoY growth (h=1,3).
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, hashlib
from pathlib import Path
import numpy as np, pandas as pd

DATA = Path("data/canada_gdp_employment.csv"); OUT = Path("out")
MODEL = "ardl_emp_direct"
COVID = (pd.Timestamp("2020-03-01"), pd.Timestamp("2021-12-01"))  # pre-specified exclusion window
H = (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
idx = d.index
def shift(m, k): return m + pd.DateOffset(months=k)
def ymd(m): return m.strftime("%Y-%m-%d")

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

def design(h):
    """Direct-regression rows indexed by target month m: y(m) on [1, y(m-h-2), e(m-h-1), e(m-h-2)].
    At origin t=m-h these regressors are gdp at t-2 and emp at t-1, t-2."""
    rows = pd.DataFrame(index=idx)
    rows["y"] = d["gdp_yoy"]
    rows["y_lag"] = d["gdp_yoy"].shift(h + 2)
    rows["e_lag1"] = d["emp_yoy"].shift(h + 1)
    rows["e_lag2"] = d["emp_yoy"].shift(h + 2)
    # months involved for COVID exclusion: target month m and regressor months m-h-2 .. m-h-1
    m = rows.index
    first_month = m - pd.DateOffset(months=h + 2)
    covid = np.array([(a <= COVID[1]) and (b >= COVID[0]) for a, b in zip(first_month, m)])
    rows["covid"] = covid
    return rows

ledger, excl = [], []
last_n = {}
for h in H:
    X = design(h)
    for t in origins:
        tgt = shift(t, h)
        actual = float(d.loc[tgt, "gdp_yoy"])
        g_av, e_av = shift(t, -2), shift(t, -1)        # information boundaries at origin t
        # naive
        naive = float(d.loc[g_av, "gdp_yoy"])
        ledger.append(dict(origin_date=ymd(t), target_date=ymd(tgt), horizon=h, forecast=naive,
                           actual=actual, model="naive", last_target_date_used=ymd(g_av),
                           last_predictor_date_used=""))
        # model: training targets must be observed at origin -> target month <= t-2
        s0 = X[X.index <= g_av]
        s1 = s0.dropna(subset=["y", "y_lag", "e_lag1", "e_lag2"])
        s2 = s1[~s1["covid"]]
        # guard: regressor months of all training rows lie before the target month <= t-2
        A = np.column_stack([np.ones(len(s2)), s2[["y_lag", "e_lag1", "e_lag2"]].values])
        beta, *_ = np.linalg.lstsq(A, s2["y"].values, rcond=None)
        xnew = np.array([1.0, d.loc[g_av, "gdp_yoy"], d.loc[e_av, "emp_yoy"], d.loc[shift(t, -2), "emp_yoy"]])
        assert not np.isnan(xnew).any()
        fc = float(xnew @ beta)
        ledger.append(dict(origin_date=ymd(t), target_date=ymd(tgt), horizon=h, forecast=fc,
                           actual=actual, model=MODEL, last_target_date_used=ymd(g_av),
                           last_predictor_date_used=ymd(e_av)))
        if t == origins[-1]:
            last_n[h] = len(s2)
            excl += [
                dict(step=f"h{h}_1", rule=f"origin {ymd(t)}, h={h}: keep rows with target month <= t-2 ({ymd(g_av)}) (information rule)",
                     n_before=n_raw, n_after=len(s0), n_removed=n_raw - len(s0),
                     rationale="targets after t-2 not yet published at origin (includes 2 rows with empty GDP)"),
                dict(step=f"h{h}_2", rule="drop rows whose lagged regressors fall before 1998-01 (not in file)",
                     n_before=len(s0), n_after=len(s1), n_removed=len(s0) - len(s1),
                     rationale="lagged regressors unavailable"),
                dict(step=f"h{h}_3", rule="drop rows whose target or regressor months fall in 2020-03..2021-12",
                     n_before=len(s1), n_after=len(s2), n_removed=len(s1) - len(s2),
                     rationale="pre-specified COVID outlier exclusion; extreme swings would dominate OLS"),
            ]

fc = pd.DataFrame(ledger)
fc["forecast"] = fc["forecast"].round(6)
fc = fc.sort_values(["model", "horizon", "origin_date"]).reset_index(drop=True)
fc.to_csv(OUT / "forecasts.csv", index=False)
pd.DataFrame(excl).to_csv(OUT / "exclusions.csv", index=False)

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

res = dict(task_id="F3-conf-01", model_name=MODEL,
    model_description=("Direct h-step OLS, separately per horizon: gdp_yoy(t+h) on a constant, gdp_yoy(t-2), "
        "emp_yoy(t-1) and emp_yoy(t-2), re-estimated at each origin on an expanding window of all target months "
        "from 1998 up to t-2 (training rows with target or regressor months in 2020-03..2021-12 excluded)."),
    n_origins=len(origins), first_origin=ymd(origins[0]), last_origin=ymd(origins[-1]),
    last_target_date_available=ymd(shift(origins[-1], -2)),
    last_predictor_date_available=ymd(shift(origins[-1], -1)))
for h in H:
    rn, mn = err("naive", h); rm, mm = err(MODEL, h)
    res.update({f"rmse_naive_h{h}": round(rn, 6), f"mae_naive_h{h}": round(mn, 6),
                f"rmse_model_h{h}": round(rm, 6), f"mae_model_h{h}": round(mm, 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.update(feasibility="complete", feasibility_reason=None,
    data_issues=[
        {"issue": "gdp_level/gdp_yoy empty for 2026-08 and 2026-09 (not yet published)",
         "handling": "never needed: latest target is 2026-07 and latest GDP used is 2026-02"},
        {"issue": "README says gdp_yoy/emp_yoy are empty for the first 12 months, but the file has values from 1998-01 (computed from pre-1998 levels not in the file)",
         "handling": "values used as given; recomputation from levels matches exactly wherever possible (from 1999-01)"},
        {"issue": "COVID-19 swings 2020-03..2021-12 (gdp_yoy from -16.2 to +18.5 %)",
         "handling": "training rows touching this window excluded from model estimation (pre-specified); evaluation period 2022-06..2026-07 unaffected"},
        {"issue": "data are current vintage, 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(last_n[1]),
    n_analysed_h3=int(last_n[3]),
    random_seed=None,
    input_sha256=hashlib.sha256(DATA.read_bytes()).hexdigest())
# --- Uncertainty: Diebold-Mariano (squared loss), Newey-West variance with h+1 lags
# (errors overlap: the naive information gap is h+2 months), Harvey-Leybourne-Newbold small-sample correction.
from scipy import stats
for h in H:
    a = fc[(fc.model == "naive") & (fc.horizon == h)].set_index("origin_date")
    b = fc[(fc.model == MODEL) & (fc.horizon == h)].set_index("origin_date")
    dl = ((b.forecast - b.actual) ** 2 - (a.forecast - a.actual) ** 2).values
    T, L = len(dl), h + 1
    u = dl - dl.mean()
    lrv = u @ u / T + 2 * sum((1 - k / (L + 1)) * (u[k:] @ u[:-k]) / T for k in range(1, L + 1))
    dm = dl.mean() / np.sqrt(lrv / T)
    hln = dm * np.sqrt((T + 1 - 2 * h + h * (h - 1) / T) / T)
    res[f"dm_stat_h{h}"] = round(float(hln), 4)
    res[f"dm_pvalue_two_sided_h{h}"] = round(float(2 * stats.t.cdf(-abs(hln), df=T - 1)), 4)
res["dm_note"] = "negative stat = model has lower squared loss than naive; HLN-corrected, NW lags h+1, t(T-1) reference"
(OUT / "results.json").write_text(json.dumps(res, indent=2))
print(json.dumps({k: v for k, v in res.items() if k.startswith(("rmse", "mae", "rel", "n_"))}, indent=1))