"""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. No randomness."""
import json, os, glob
import numpy as np, pandas as pd

D, O = "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_1985_excovid"
SAMPLE_START = pd.Timestamp("1985-04-01")            # first dependent-variable quarter
COVID_EXCL = pd.to_datetime(["2020-04-01", "2020-07-01", "2020-10-01"])  # dependent quarters dropped

def read_series(path):
    d = pd.read_csv(path)
    d.columns = ["date", "value"]
    d["date"] = pd.to_datetime(d["date"])
    return d

n_input_rows, issues = {}, []
fr = pd.read_csv(os.path.join(D, "first_release.csv"), parse_dates=["quarter", "first_release_vintage_date"])
n_input_rows["first_release.csv"] = len(fr)
cur = read_series(os.path.join(D, "current_vintage.csv"))
n_input_rows["current_vintage.csv"] = len(cur)
frmap = dict(zip(fr["quarter"], fr["first_release_value"]))
frdate = dict(zip(fr["quarter"], fr["first_release_vintage_date"]))

def ar1_fit(s):
    """s: vintage DataFrame. OLS of y_t on (1, y_{t-1}) over the restricted sample."""
    y = s.set_index("date")["value"]
    df = pd.DataFrame({"y": y, "ylag": y.shift(1)}).dropna()
    df = df[(df.index >= SAMPLE_START) & (~df.index.isin(COVID_EXCL))]
    X = np.column_stack([np.ones(len(df)), df["ylag"].values])
    beta, *_ = np.linalg.lstsq(X, df["y"].values, rcond=None)
    return beta, len(df), df

rows, excl_rows = [], []
for o in ORIGINS:
    rel = f"vintages/A191RL1Q225SBEA_vintage_{o}.csv"
    s = read_series(os.path.join(D, rel))
    n_input_rows[rel] = len(s)
    assert s["value"].notna().all() and not s["date"].duplicated().any()
    last_date, last_val = s["date"].iloc[-1], float(s["value"].iloc[-1])
    beta, n_ar, df_ar = ar1_fit(s)
    if o == ORIGINS[-1]:
        n_raw = len(s)
        n1 = int((s["date"] >= SAMPLE_START).sum())
        n2 = n_ar
        excl_rows += [
            dict(step=1, rule="drop dependent quarters before 1985-04-01", n_before=n_raw, n_after=n1, n_removed=n_raw - n1,
                 rationale="AR estimated on post-1985 (Great Moderation) sample, fixed a priori; 1947Q2 also lacks a lag"),
            dict(step=2, rule="drop dependent quarters 2020Q2-2020Q4", n_before=n1, n_after=n2, n_removed=n1 - n2,
                 rationale="COVID outliers (-28%, +35%) would dominate OLS as y_t or y_{t-1}"),
        ]
        n_analysed = n2
    path_model, prev = [], last_val
    for h in (1, 2):
        prev = beta[0] + beta[1] * prev
        path_model.append(prev)
    for h in (1, 2):
        tgt = last_date + pd.DateOffset(months=3 * h)
        assert tgt in frmap, f"no first release for {tgt}"
        assert frdate[tgt] > pd.Timestamp(o), "actual must be published after origin"
        for model, fc, ntr in (("naive", last_val, len(s)), (MODEL, path_model[h - 1], n_ar)):
            rows.append(dict(origin_date=o, target_date=tgt.strftime("%Y-%m-%d"), horizon=h,
                             forecast=round(float(fc), 10), actual=float(frmap[tgt]), model=model,
                             vintage_date=o, train_last_date=last_date.strftime("%Y-%m-%d"),
                             train_last_value=last_val, n_train=int(ntr)))
    if abs(last_val - frmap.get(last_date, np.nan)) > 1e-9:
        issues.append({"issue": f"vintage {o} last value differs from first_release for {last_date.date()}",
                       "handling": "vintage value used as input (information rule); first release used only as actual"})

fc = pd.DataFrame(rows)
fc.to_csv(os.path.join(O, "forecasts.csv"), index=False)
pd.DataFrame(excl_rows).to_csv(os.path.join(O, "exclusions.csv"), index=False)

def err(model, h):
    e = fc[(fc.model == model) & (fc.horizon == h)]
    e = e["forecast"] - e["actual"]
    assert len(e) == 8
    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 at each origin on that origin's vintage only, dependent sample "
                             "1985Q2 to the last vintage quarter excluding dependent quarters 2020Q2-2020Q4 (COVID outliers); "
                             "h=2 forecast by iterating the AR(1)."),
       "n_origins": len(ORIGINS), "first_origin": ORIGINS[0], "last_origin": ORIGINS[-1]}
for h in (1, 2):
    rn, mn = err("naive", h); rm, mm = err(MODEL, h)
    res[f"rmse_naive_h{h}"], res[f"mae_naive_h{h}"] = rn, mn
    res[f"rmse_model_h{h}"], res[f"mae_model_h{h}"] = rm, mm
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}"]
tg = pd.to_datetime(fc[(fc.horizon == 1) & (fc.model == "naive")]["target_date"])
curmap = dict(zip(cur["date"], cur["value"]))
res["mean_abs_revision_h1_targets"] = float(np.mean([abs(curmap[t] - frmap[t]) for t in tg]))
res["feasibility"] = "complete"; res["feasibility_reason"] = None
issues.append({"issue": "COVID-19 outliers 2020Q2 (about -28%) and 2020Q3 (about +35%) in all vintages",
               "handling": "dependent quarters 2020Q2-2020Q4 excluded from AR(1) estimation (decided before evaluation)"})
issues.append({"issue": "8th origin is 2025-12-31 (2025Q3 advance estimate delayed by the 2025 federal shutdown); 2025Q4 first release dated 2026-02-28",
               "handling": "used as provided; origin dates and first-release dates checked to respect chronology"})
res["data_issues"] = issues
res["n_input_rows"] = n_input_rows
res["n_analysed"] = int(n_analysed)
res["random_seed"] = None
# supplementary uncertainty summaries (n=8 per horizon: descriptive only)
from scipy import stats
for h in (1, 2):
    a = fc[(fc.model == "naive") & (fc.horizon == h)].reset_index(drop=True)
    b = fc[(fc.model == MODEL) & (fc.horizon == h)].reset_index(drop=True)
    d = (a.forecast - a.actual) ** 2 - (b.forecast - b.actual) ** 2   # >0: model better
    n = len(d); dc = d - d.mean()
    lrv = np.sum(dc ** 2) / n + sum(2 * (1 - k / h) * np.sum(dc[k:].values * dc[:-k].values) / n for k in range(1, h))
    t = d.mean() / np.sqrt(lrv / n)
    res[f"supp_n_origins_model_lower_sqerr_h{h}"] = int((d > 0).sum())
    res[f"supp_dm_stat_h{h}"] = float(t)
    res[f"supp_dm_pvalue_t7_h{h}"] = float(2 * stats.t.sf(abs(t), df=n - 1))
res["ar1_coefficients_last_origin"] = {"intercept": float(beta[0]), "phi": float(beta[1])}
with open(os.path.join(O, "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", "mean"))}, indent=2))