"""F3-conf-01: GDP y/y growth forecasts under publication lags.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Naive: y(t-2). Model 'adl_direct_emp': direct OLS per horizon
  y(tau+h) = a + b*y(tau-2) + c*e(tau-1), training rows tau with tau+h <= t-2,
  excluding rows touching 2020-03..2021-12 (COVID); forecast plugs y(t-2), e(t-1).
No randomness is used.
"""
import json, os
import numpy as np, pandas as pd

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DATA = os.path.join(ROOT, "data", "canada_gdp_employment.csv")
OUT = os.path.join(ROOT, "out")
MODEL = "adl_direct_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)
d = d.set_index("month")
months = d.index
origins = pd.date_range("2022-05-01", "2026-04-01", freq="MS")
H = [1, 3]
M = lambda m, k: m + pd.DateOffset(months=k)
fmt = lambda x: x.strftime("%Y-%m-%d")
in_covid = lambda m: (m >= COVID[0]) & (m <= COVID[1])


def training_rows(t, h, return_ledger=False, exclude_covid=True):
    """Training rows available at origin t for horizon h (vectorised)."""
    gdp_last = M(t, -2)   # latest GDP month observed at origin t
    emp_last = M(t, -1)   # latest employment month observed at origin t
    y = d["gdp_yoy"].where(months <= gdp_last)   # information truncation
    e = d["emp_yoy"].where(months <= emp_last)
    df = pd.DataFrame({"y_tgt": y.shift(-h), "y_lag": y.shift(2), "e_lag": e.shift(1)}, index=months)
    ledger = []
    n0 = len(df)
    df1 = df.dropna()
    ledger.append(("drop_unavailable", f"target month tau+{h} <= {fmt(gdp_last)}, y(tau-2) and e(tau-1) observed",
                   n0, len(df1), "information-availability rule / series edges"))
    tau = df1.index
    bad = in_covid(tau) | in_covid(tau + pd.DateOffset(months=h)) | in_covid(tau - pd.DateOffset(months=2)) | in_covid(tau - pd.DateOffset(months=1))
    df2 = df1[~np.asarray(bad)] if exclude_covid else df1
    ledger.append(("exclude_covid", "drop rows whose target/predictor months fall in 2020-03..2021-12",
                   len(df1), len(df2), "COVID outliers dominate OLS; decided with pre-2022 info"))
    # sanity: every used month respects the rule
    assert (df2.index + pd.DateOffset(months=h)).max() <= gdp_last
    return (df2, ledger) if return_ledger else df2


rows = []
n_analysed = None
for t in origins:
    y_t2 = d.at[M(t, -2), "gdp_yoy"]
    e_t1 = d.at[M(t, -1), "emp_yoy"]
    for h in H:
        tgt = M(t, h)
        actual = d.at[tgt, "gdp_yoy"]
        rows.append(dict(origin_date=fmt(t), target_date=fmt(tgt), horizon=h, forecast=y_t2, actual=actual,
                         model="naive", last_target_date_used=fmt(M(t, -2)), last_predictor_date_used=""))
        tr = training_rows(t, h)
        X = np.column_stack([np.ones(len(tr)), tr["y_lag"], tr["e_lag"]])
        beta, *_ = np.linalg.lstsq(X, tr["y_tgt"].values, rcond=None)
        fc = float(beta @ np.array([1.0, y_t2, e_t1]))
        rows.append(dict(origin_date=fmt(t), target_date=fmt(tgt), horizon=h, forecast=fc, actual=actual,
                         model=MODEL, last_target_date_used=fmt(M(t, -2)), last_predictor_date_used=fmt(M(t, -1))))

fc = pd.DataFrame(rows)
assert len(fc) == 192 and fc["actual"].notna().all()
fc.to_csv(os.path.join(OUT, "forecasts.csv"), index=False, float_format="%.6f")

# exclusion ledger for the last estimation sample (last origin, h=1); n_raw = all file rows
t_last = origins[-1]
tr_last, led = training_rows(t_last, 1, return_ledger=True)
n_analysed = len(tr_last)
assert led[0][2] == n_raw
ex = pd.DataFrame([dict(step=s, rule=r, n_before=a, n_after=b, n_removed=a - b, rationale=w) for s, r, a, b, w in led])
ex.to_csv(os.path.join(OUT, "exclusions.csv"), index=False)
n_last_h3 = len(training_rows(t_last, 3))

res = {"task_id": "F3-conf-01", "model_name": MODEL,
       "model_description": ("Direct (per-horizon) OLS regression of gdp_yoy(t+h) on a constant, gdp_yoy(t-2) and emp_yoy(t-1), "
                             "re-estimated at each origin on an expanding window of training rows (from 1998-03) whose target month is <= t-2, "
                             "excluding rows touching 2020-03..2021-12."),
       "n_origins": int(len(origins)), "first_origin": fmt(origins[0]), "last_origin": fmt(t_last),
       "last_target_date_available": fmt(M(t_last, -2)), "last_predictor_date_available": fmt(M(t_last, -1))}
for h in H:
    for m, lab in [("naive", "naive"), (MODEL, "model")]:
        s = fc[(fc.horizon == h) & (fc.model == m)]
        err = s["forecast"] - s["actual"]
        res[f"rmse_{lab}_h{h}"] = float(np.sqrt(np.mean(err ** 2)))
        res[f"mae_{lab}_h{h}"] = float(np.mean(np.abs(err)))
for h in H:
    res[f"rel_improvement_rmse_h{h}"] = 1 - res[f"rmse_model_h{h}"] / res[f"rmse_naive_h{h}"]

# Uncertainty: Diebold-Mariano test (squared loss), HAC variance with h-1 lags (Newey-West), HLN small-sample correction
from scipy import stats
dm = {}
for h in H:
    en = fc[(fc.horizon == h) & (fc.model == "naive")].sort_values("origin_date")
    em = fc[(fc.horizon == h) & (fc.model == MODEL)].sort_values("origin_date")
    dl = (em["forecast"].values - em["actual"].values) ** 2 - (en["forecast"].values - en["actual"].values) ** 2
    T = len(dl); dbar = dl.mean(); u = dl - dbar
    L = h + 1  # Newey-West lags h+1 since effective GDP horizon is h+2 months
    v = np.mean(u * u) + 2 * sum((1 - k / (L + 1)) * np.mean(u[k:] * u[:-k]) for k in range(1, L + 1))
    stat = dbar / np.sqrt(v / T)
    hk = h + 2
    stat *= np.sqrt((T + 1 - 2 * hk + hk * (hk - 1) / T) / T)
    p = 2 * stats.t.sf(abs(stat), df=T - 1)
    dm[f"h{h}"] = {"mean_loss_diff_model_minus_naive": float(dbar), "dm_stat_hln": float(stat), "p_value_two_sided": float(p), "nw_lags": L}
res["dm_test"] = dm
res["n_training_last_origin_h3"] = int(n_last_h3)
# Sensitivity: same model without the COVID exclusion (not the headline model)
sens = {}
for h in H:
    errs = []
    for t in origins:
        tr = training_rows(t, h, exclude_covid=False)
        X = np.column_stack([np.ones(len(tr)), tr["y_lag"], tr["e_lag"]])
        b, *_ = np.linalg.lstsq(X, tr["y_tgt"].values, rcond=None)
        f_ = float(b @ np.array([1.0, d.at[M(t, -2), "gdp_yoy"], d.at[M(t, -1), "emp_yoy"]]))
        errs.append(f_ - d.at[M(t, h), "gdp_yoy"])
    errs = np.array(errs)
    sens[f"rmse_h{h}"] = float(np.sqrt(np.mean(errs ** 2)))
    sens[f"rel_improvement_rmse_h{h}"] = 1 - sens[f"rmse_h{h}"] / res[f"rmse_naive_h{h}"]
res["sensitivity_no_covid_exclusion"] = sens
res.update({"feasibility": "complete", "feasibility_reason": None,
            "data_issues": [
                {"issue": "README states gdp_yoy is empty for the first 12 months, but the file contains yoy values for 1998 (both series).",
                 "handling": "Used as given; yoy recomputed from levels agrees to rounding (<5e-5) where possible."},
                {"issue": "gdp columns missing for 2026-08 and 2026-09 (not yet published).",
                 "handling": "Not needed: latest target (2026-07) and latest GDP input (2026-02) are present."},
                {"issue": "COVID-19 outliers (gdp_yoy from -16.2 to +18.5 in 2020-2021).",
                 "handling": "Model training rows touching 2020-03..2021-12 excluded; naive unaffected; see out/exclusions.csv."},
                {"issue": "Current-vintage data (revised), not 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(n_analysed), "random_seed": None})
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps({k: v for k, v in res.items() if k.startswith(("rmse", "mae", "rel", "n_"))}, indent=2))
print(json.dumps(dm, indent=2))