"""F3-conf-02: real-time forecasts of US real GDP growth from ALFRED vintages.
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, glob
import numpy as np, pandas as pd

DATA, OUT = "data", "out"
ORIGINS = ["2024-01-31", "2024-04-30", "2024-07-31", "2024-10-31",
           "2025-01-31", "2025-04-30", "2025-07-31", "2025-12-31"]
MODEL = "ar1_post1985_excovid"
SAMPLE_START = pd.Timestamp("1985-01-01")
COVID = [pd.Timestamp(x) for x in ("2020-04-01", "2020-07-01")]

def vpath(d): return f"vintages/A191RL1Q225SBEA_vintage_{d}.csv"

def read_series(rel):
    df = pd.read_csv(os.path.join(DATA, rel))
    df.columns = ["date", "value"]
    df["date"] = pd.to_datetime(df["date"])
    return df.sort_values("date").reset_index(drop=True)

n_input_rows = {}
cur = read_series("current_vintage.csv"); n_input_rows["current_vintage.csv"] = len(cur)
fr = pd.read_csv(os.path.join(DATA, "first_release.csv")); n_input_rows["first_release.csv"] = len(fr)
fr["quarter"] = pd.to_datetime(fr["quarter"]); fr["first_release_vintage_date"] = pd.to_datetime(fr["first_release_vintage_date"])
fr_map = fr.set_index("quarter")
cur_map = cur.set_index("date")["value"]

def ar1_design(v):
    """Regression obs t with y_t, y_{t-1}; keep t >= 1985Q1, drop if t or t-1 is a pandemic outlier quarter."""
    y = v["value"].to_numpy(); d = v["date"].to_numpy()
    df = pd.DataFrame({"date": d[1:], "lagdate": d[:-1], "y": y[1:], "ylag": y[:-1]})
    steps = [("raw_regression_pairs", len(df))]
    df = df[df["date"] >= SAMPLE_START]; steps.append(("t>=1985Q1", len(df)))
    df = df[~df["date"].isin(COVID) & ~df["lagdate"].isin(COVID)]; steps.append(("drop_pandemic_2020Q2-Q4", len(df)))
    return df, steps

rows, excl, data_issues = [], [], []
for o in ORIGINS:
    rel = vpath(o); v = read_series(rel); n_input_rows[rel] = len(v)
    assert v["value"].notna().all() and not v["date"].duplicated().any()
    last_d, last_y = v["date"].iloc[-1], float(v["value"].iloc[-1])
    # model estimation (vintage o only)
    df, steps = ar1_design(v)
    X = np.column_stack([np.ones(len(df)), df["ylag"].to_numpy()])
    beta, *_ = np.linalg.lstsq(X, df["y"].to_numpy(), rcond=None)
    c, phi = float(beta[0]), float(beta[1])
    f_model, prev = {}, last_y
    for h in (1, 2):
        prev = c + phi * prev; f_model[h] = prev
    if o == ORIGINS[-1]:
        excl.append(dict(step=0, rule="rows in last-origin vintage file", n_before=len(v), n_after=len(v), n_removed=0,
                         rationale="raw input (vintage 2025-12-31)"))
        excl.append(dict(step=1, rule="form AR(1) pairs (first quarter has no lag)", n_before=len(v), n_after=steps[0][1],
                         n_removed=len(v) - steps[0][1], rationale="1947Q2 has no lagged value in file"))
        excl.append(dict(step=2, rule="keep t >= 1985Q1", n_before=steps[0][1], n_after=steps[1][1],
                         n_removed=steps[0][1] - steps[1][1], rationale="post-Great-Moderation estimation sample (a priori)"))
        excl.append(dict(step=3, rule="drop t in {2020Q2,2020Q3,2020Q4} (pandemic outlier as y_t or y_{t-1})", n_before=steps[1][1],
                         n_after=steps[2][1], n_removed=steps[1][1] - steps[2][1], rationale="extreme pandemic outliers (-28%, +35%) would dominate OLS"))
        n_analysed = steps[2][1]
    for h in (1, 2):
        tgt = last_d + pd.DateOffset(months=3 * h)
        if tgt not in fr_map.index:
            raise RuntimeError(f"no first release for {tgt}")
        actual = float(fr_map.loc[tgt, "first_release_value"])
        assert fr_map.loc[tgt, "first_release_vintage_date"] > pd.Timestamp(o), "actual published before origin"
        for m, f, n in (("naive", last_y, len(v)), (MODEL, f_model[h], len(df))):
            rows.append(dict(origin_date=o, target_date=tgt.strftime("%Y-%m-%d"), horizon=h, forecast=round(f, 10),
                             actual=actual, model=m, vintage_date=o, train_last_date=last_d.strftime("%Y-%m-%d"),
                             train_last_value=last_y, n_train=n, ar_const=c if m != "naive" else np.nan,
                             ar_phi=phi if m != "naive" else np.nan))

fc = pd.DataFrame(rows)
os.makedirs(OUT, exist_ok=True)
fc.drop(columns=["ar_const", "ar_phi"]).to_csv(os.path.join(OUT, "forecasts.csv"), index=False)
pd.DataFrame(excl).to_csv(os.path.join(OUT, "exclusions.csv"), index=False)

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

res = {"task_id": "F3-conf-02", "model_name": MODEL,
       "model_description": ("AR(1) with intercept estimated by OLS separately on each origin's vintage, regression sample "
                             "t = 1985Q1 to the last quarter in the vintage excluding t in 2020Q2-2020Q4 (pandemic outliers as "
                             "y_t or y_{t-1}); h=2 forecast by iterating the h=1 forecast."),
       "n_origins": len(ORIGINS), "first_origin": ORIGINS[0], "last_origin": ORIGINS[-1]}
for m, k in (("naive", "naive"), (MODEL, "model")):
    for h in (1, 2):
        r, a = err(m, h); res[f"rmse_{k}_h{h}"] = r; res[f"mae_{k}_h{h}"] = a
for h in (1, 2):
    res[f"rel_improvement_rmse_h{h}"] = 1 - res[f"rmse_model_h{h}"] / res[f"rmse_naive_h{h}"]
t1 = fc[(fc.model == "naive") & (fc.horizon == 1)]
revs = [abs(float(cur_map.loc[pd.Timestamp(t)]) - a) for t, a in zip(t1.target_date, t1.actual)]
# Diebold-Mariano (squared loss) with Harvey-Leybourne-Newbold small-sample correction; supplementary
from scipy import stats
for h in (1, 2):
    en = fc[(fc.model == "naive") & (fc.horizon == h)].sort_values("origin_date")
    em = fc[(fc.model == MODEL) & (fc.horizon == h)].sort_values("origin_date")
    dl = (en["forecast"].to_numpy() - en["actual"].to_numpy()) ** 2 - (em["forecast"].to_numpy() - em["actual"].to_numpy()) ** 2
    T = len(dl); dbar = dl.mean(); g = [np.mean((dl[k:] - dbar) * (dl[:T - k] - dbar)) for k in range(h)]
    var = (g[0] + 2 * sum(g[1:])) / T
    dm = dbar / np.sqrt(var) if var > 0 else np.nan
    hln = dm * np.sqrt((T + 1 - 2 * h + h * (h - 1) / T) / T)
    res[f"supp_dm_hln_stat_h{h}"] = float(hln)
    res[f"supp_dm_hln_pvalue_two_sided_h{h}"] = float(2 * stats.t.sf(abs(hln), df=T - 1))
    res[f"supp_n_model_wins_sqloss_h{h}"] = int(np.sum(dl > 0))
res["mean_abs_revision_h1_targets"] = float(np.mean(revs))
res["feasibility"] = "complete"; res["feasibility_reason"] = None
res["data_issues"] = [
    {"issue": "Pandemic outlier quarters 2020Q2 (about -28%) and 2020Q3 (about +35%) present in every vintage",
     "handling": "AR(1) regression observations with 2020Q2 or 2020Q3 as dependent or lagged value are excluded (decided before evaluation); naive unaffected"},
    {"issue": "Eighth origin is 2025-12-31 instead of 2025-10-31 because the 2025Q3 advance estimate was delayed by the federal shutdown",
     "handling": "Used as given; targets defined relative to the last quarter in that vintage (2025Q4, 2026Q1)"},
    {"issue": "current_vintage.csv contains revisions unknown at the origins",
     "handling": "Used only for the revision statistic mean_abs_revision_h1_targets, never for estimation or forecast input"}]
res["n_input_rows"] = n_input_rows
res["n_analysed"] = int(n_analysed)
res["random_seed"] = None
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(res, f, indent=2)
# per-origin diagnostics (printed only; consistent with written results)
print(fc[["origin_date", "target_date", "horizon", "model", "forecast", "actual", "n_train", "ar_const", "ar_phi"]].to_string(index=False))
print("revisions h1 targets:", [round(x, 2) for x in revs])
print(json.dumps({k: v for k, v in res.items() if k.startswith(("rmse", "mae", "rel", "mean"))}, indent=1))