"""F3-conf-01: direct forecasts of Canadian GDP yoy growth under publication lags.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py"""
import json, numpy as np, pandas as pd

DATA = "data/canada_gdp_employment.csv"
d = pd.read_csv(DATA, parse_dates=["month"])
n_rows = len(d)
d = d.set_index("month").sort_index()
y, e = d["gdp_yoy"], d["emp_yoy"]
months = list(d.index)
pos = {m: i for i, m in enumerate(months)}

# data checks
issues = []
assert d.index.is_unique
yy = (d.gdp_level / d.gdp_level.shift(12) * 100 - 100).round(4)
ee = (d.emp_level / d.emp_level.shift(12) * 100 - 100).round(4)
assert (yy - y).abs().max() < 1e-3 and (ee - e).abs().max() < 1e-3
issues.append({"issue": "gdp_yoy/emp_yoy are filled for 1998-01..1998-12 although README says they are empty for the first 12 months (levels for 1997 are not in the file).",
               "handling": "Values used as given; for 1999-01 onward yoy columns were verified to match levels to 4 decimals."})
issues.append({"issue": "GDP columns empty for 2026-08 and 2026-09 (not yet published).",
               "handling": "Expected; not needed since the last target is 2026-07."})
issues.append({"issue": "COVID-19 shock: yoy GDP growth ranges from -16.2% to +18.5% in 2020-2021, extreme outliers for OLS.",
               "handling": "Pre-specified (not tuned on test errors): training pairs in which the target or any regressor month lies in 2020-03..2021-12 are dropped from the model's estimation sample; the forecast at each origin still uses the latest available values."})

COVID = (pd.Timestamp("2020-03-01"), pd.Timestamp("2021-12-01"))
import sys
EXCL = '--no-covid-exclusion' not in sys.argv
def covid(m): return EXCL and COVID[0] <= m <= COVID[1]

MODEL = "ardl_direct_emp"
origins = pd.date_range("2022-05-01", "2026-04-01", freq="MS")
assert len(origins) == 48
rows, n_last_by_h = [], {}
for t in origins:
    it = pos[t]
    last_g, last_e = months[it - 2], months[it - 1]   # information set at end of month t
    for h in (1, 3):
        tgt = months[it + h]
        actual = float(y[tgt])
        rows.append(dict(origin_date=t.date(), target_date=tgt.date(), horizon=h,
                         forecast=float(y[last_g]), actual=actual, model="naive",
                         last_target_date_used=last_g.date(), last_predictor_date_used=""))
        # direct regression: y[s+h] = a + b*y[s-2] + c*e[s-1], pseudo-origins s with s+h <= t-2
        X, Y = [], []
        for s in range(2, it - 2 - h + 1):
            ms = [months[s + h], months[s - 2], months[s - 1]]
            if any(covid(m) for m in ms): continue
            if months[s + h] > last_g or months[s - 1] > last_e: raise RuntimeError
            X.append([1.0, y.iloc[s - 2], e.iloc[s - 1]]); Y.append(y.iloc[s + h])
        X, Y = np.array(X), np.array(Y)
        beta = np.linalg.lstsq(X, Y, rcond=None)[0]
        fc = float(beta @ np.array([1.0, y[last_g], e[last_e]]))
        rows.append(dict(origin_date=t.date(), target_date=tgt.date(), horizon=h,
                         forecast=fc, actual=actual, model=MODEL,
                         last_target_date_used=last_g.date(), last_predictor_date_used=last_e.date()))
        if t == origins[-1]:
            n_last_by_h[h] = len(Y)
f = pd.DataFrame(rows).sort_values(["model", "horizon", "origin_date"])
if EXCL: f.to_csv("out/forecasts.csv", index=False)

res = {"task_id": "F3-conf-01", "model_name": MODEL,
       "model_description": ("Direct h-step OLS 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 from 1998 using only pairs whose target is published "
                             "at the origin, excluding pairs touching 2020-03..2021-12."),
       "n_origins": 48, "first_origin": "2022-05-01", "last_origin": "2026-04-01",
       "last_target_date_available": str(months[pos[origins[-1]] - 2].date()),
       "last_predictor_date_available": str(months[pos[origins[-1]] - 1].date())}
for h in (1, 3):
    for m, k in (("naive", "naive"), (MODEL, "model")):
        err = f[(f.model == m) & (f.horizon == h)].eval("forecast-actual")
        assert len(err) == 48
        res[f"rmse_{k}_h{h}"] = float(np.sqrt((err ** 2).mean()))
        res[f"mae_{k}_h{h}"] = float(err.abs().mean())
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}"]
# supplementary: Diebold-Mariano (HAC, lag h-1, Harvey small-sample correction), squared loss
from scipy import stats
dm = {}
for h in (1, 3):
    a = f[(f.model == "naive") & (f.horizon == h)].eval("forecast-actual").values
    b = f[(f.model == MODEL) & (f.horizon == h)].eval("forecast-actual").values
    dl = b ** 2 - a ** 2; T = len(dl); dc = dl - dl.mean()
    lr = (dc @ dc) / T + 2 * sum((1 - k / h) * (dc[k:] @ dc[:-k]) / T for k in range(1, h)) if h > 1 else (dc @ dc) / T
    stat = dl.mean() / np.sqrt(lr / T) * np.sqrt((T + 1 - 2 * h + h * (h - 1) / T) / T)
    dm[f"h{h}"] = {"dm_stat": float(stat), "p_two_sided": float(2 * stats.t.sf(abs(stat), T - 1))}
res.update({"feasibility": "complete", "feasibility_reason": None, "data_issues": issues,
            "n_input_rows": {"canada_gdp_employment.csv": n_rows}, "n_analysed": n_last_by_h[1], "n_analysed_by_horizon": {f"h{k}": v for k, v in n_last_by_h.items()},
            "random_seed": None, "supplementary_dm_test": dm})
if EXCL: json.dump(res, open("out/results.json", "w"), indent=2)
print(json.dumps(res, indent=2))