"""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
import pandas as pd

DATA = "data/canada_gdp_employment.csv"
COVID = (pd.Timestamp("2020-03-01"), pd.Timestamp("2021-12-01"))  # training exclusion window
MODEL = "adl_emp_direct"

raw = pd.read_csv(DATA, parse_dates=["month"])
n_raw = len(raw)
d = raw.set_index("month").sort_index()
months = d.index
pos = {m: i for i, m in enumerate(months)}

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

def in_covid(m):
    return COVID[0] <= m <= COVID[1]

def training_rows(t, h):
    """Direct regression rows indexed by base month m: target y[m+k], regressors y[m], e[m+1], k=h+2.
    Only data available at end of month t: gdp months <= t-2, emp months <= t-1."""
    k = h + 2
    gcut, ecut = shift(t, -2), shift(t, -1)
    log = []
    m_all = [m for m in months if not np.isnan(d.loc[m, "gdp_yoy"])]
    rows = list(m_all)
    log.append(("A", "base month has non-missing gdp_yoy in file", n_raw, len(rows)))
    rows = [m for m in rows if m <= gcut]
    log.append(("B", f"base month gdp available at origin (<= t-2 = {gcut.date()})", log[-1][3], len(rows)))
    rows = [m for m in rows if shift(m, k) <= gcut and shift(m, 1) <= ecut and shift(m, k) in pos]
    log.append(("C", f"training target month m+{k} observed at origin (<= {gcut.date()})", log[-1][3], len(rows)))
    rows = [m for m in rows if not (in_covid(m) or in_covid(shift(m, 1)) or in_covid(shift(m, k)))]
    log.append(("D", "drop rows whose regressor or target month lies in 2020-03..2021-12 (pandemic)", log[-1][3], len(rows)))
    return rows, log

origins = pd.date_range("2022-05-01", "2026-04-01", freq="MS")
assert len(origins) == 48
recs = []
last_log = {}
for t in origins:
    gcut, ecut = shift(t, -2), shift(t, -1)
    y_now = d.loc[gcut, "gdp_yoy"]; e_now = d.loc[ecut, "emp_yoy"]
    assert not np.isnan(y_now) and not np.isnan(e_now)
    for h in (1, 3):
        k = h + 2
        target = shift(t, h)
        actual = float(d.loc[target, "gdp_yoy"])
        recs.append(dict(origin_date=t.date().isoformat(), target_date=target.date().isoformat(), horizon=h,
                         forecast=float(y_now), actual=actual, model="naive",
                         last_target_date_used=gcut.date().isoformat(), last_predictor_date_used=""))
        rows, log = training_rows(t, h)
        Y = np.array([d.loc[shift(m, k), "gdp_yoy"] for m in rows])
        X = np.column_stack([np.ones(len(rows)), [d.loc[m, "gdp_yoy"] for m in rows],
                             [d.loc[shift(m, 1), "emp_yoy"] for m in rows]])
        beta, *_ = np.linalg.lstsq(X, Y, rcond=None)
        fc = float(beta @ np.array([1.0, y_now, e_now]))
        last_tgt = max(max(shift(m, k) for m in rows), gcut)
        last_emp = max(max(shift(m, 1) for m in rows), ecut)
        assert last_tgt <= gcut and last_emp <= ecut
        recs.append(dict(origin_date=t.date().isoformat(), target_date=target.date().isoformat(), horizon=h,
                         forecast=fc, actual=actual, model=MODEL,
                         last_target_date_used=last_tgt.date().isoformat(),
                         last_predictor_date_used=last_emp.date().isoformat()))
        if t == origins[-1]:
            last_log[h] = (log, len(rows))

fc = pd.DataFrame(recs)
fc.to_csv("out/forecasts.csv", index=False, float_format="%.6f")

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

res = {"task_id": "F3-conf-01", "model_name": MODEL,
       "model_description": ("Direct OLS regression per horizon h of gdp_yoy[m+h+2] on a constant, gdp_yoy[m] and emp_yoy[m+1], "
                             "re-estimated at each origin t on an expanding sample from 1998 using only pairs observed at t "
                             "(gdp <= t-2, emp <= t-1), excluding pairs touching 2020-03..2021-12; forecast plugs in gdp_yoy[t-2] and emp_yoy[t-1]."),
       "n_origins": int(len(origins)), "first_origin": origins[0].date().isoformat(),
       "last_origin": origins[-1].date().isoformat(),
       "last_target_date_available": shift(origins[-1], -2).date().isoformat(),
       "last_predictor_date_available": shift(origins[-1], -1).date().isoformat()}
for model, tag in (("naive", "naive"), (MODEL, "model")):
    for h in (1, 3):
        r, a = err(model, h)
        res[f"rmse_{tag}_h{h}"] = r; res[f"mae_{tag}_h{h}"] = 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 empty for 2026-08 and 2026-09 (not yet published)",
                 "handling": "never needed: latest GDP month used is 2026-02 (t-2 at last origin); latest target is 2026-07."},
                {"issue": "README says gdp_yoy/emp_yoy are empty for the first 12 months, but the file has values for 1998-01..1998-12 that cannot be recomputed from the levels in the file",
                 "handling": "used as given (values for 1999 onward match 100*(L_m/L_{m-12}-1) to <5e-5); noted as a documentation inconsistency."},
                {"issue": "pandemic outliers 2020-03..2021-12 (gdp_yoy from -16.2 to +18.5 percent) would dominate OLS estimates",
                 "handling": "training pairs whose regressor or target month falls in 2020-03..2021-12 are excluded (rule fixed a priori; window ends before the first origin's information set)."},
                {"issue": "current-vintage data, not historical first releases",
                 "handling": "exercise simulates publication timing only; no real-time validity claimed."}],
            "n_input_rows": {"canada_gdp_employment.csv": int(n_raw)},
            "n_analysed": int(last_log[1][1]), "random_seed": None})
# Diebold-Mariano test (squared loss), Newey-West variance with h+1 lags (overlap of k=h+2-step errors)
from scipy import stats
for h in (1, 3):
    a = fc[(fc.model == "naive") & (fc.horizon == h)].reset_index(drop=True)
    b = fc[(fc.model == MODEL) & (fc.horizon == h)].reset_index(drop=True)
    dl = ((a.forecast - a.actual) ** 2 - (b.forecast - b.actual) ** 2).to_numpy()
    n = len(dl); dc = dl - dl.mean(); L = h + 1
    v = dc @ dc / n + 2 * sum((1 - j / (L + 1)) * (dc[j:] @ dc[:-j]) / n for j in range(1, L + 1))
    stat = dl.mean() / np.sqrt(v / n)
    res[f"dm_stat_h{h}"] = float(stat)
    res[f"dm_pvalue_two_sided_h{h}"] = float(2 * (1 - stats.t.cdf(abs(stat), n - 1)))
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)

# exclusion ledger: last origin, h=1 estimation sample (n_analysed); h=3 shown for reference
ex = []
for h in (1,):
    for step, rule, nb, na in last_log[h][0]:
        ex.append(dict(step=f"h{h}_{step}", rule=f"[origin 2026-04-01, h={h}] {rule}", n_before=nb, n_after=na,
                       n_removed=nb - na, rationale="information-availability rule" if step in "BC" else
                       ("missing value (unpublished GDP)" if step == "A" else "pandemic outliers excluded from estimation")))
pd.DataFrame(ex).to_csv("out/exclusions.csv", index=False)
print(json.dumps({k: v for k, v in res.items() if k.startswith(("rmse", "mae", "rel", "n_"))}, indent=1))