"""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 is used."""
import json, os
import numpy as np
import pandas as pd
from scipy import stats

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_1985"
EST_START = pd.Timestamp("1985-01-01")
OUTLIERS = np.array([pd.Timestamp("2020-04-01"), pd.Timestamp("2020-07-01")], dtype="datetime64[ns]")


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


def ar1_fit(s):
    """s: DataFrame(date,value) from the origin vintage only. Returns (c, phi, n_pairs_used, ledger)."""
    led = []
    w = s[s.date >= EST_START].reset_index(drop=True)
    led.append(("window", f"keep quarters >= {EST_START.date()}", len(s), len(w)))
    y, x = w.value.values[1:], w.value.values[:-1]
    dy, dx = w.date.values[1:], w.date.values[:-1]
    keep = ~(np.isin(dy, OUTLIERS) | np.isin(dx, OUTLIERS))
    led.append(("lag_pairs", "form (y_t, y_t-1) pairs; first window quarter used only as a lag", len(w), len(y)))
    led.append(("pandemic", "drop pairs with y_t or y_t-1 in 2020Q2-2020Q3", len(y), int(keep.sum())))
    X = np.column_stack([np.ones(keep.sum()), x[keep]])
    beta, *_ = np.linalg.lstsq(X, y[keep], rcond=None)
    return float(beta[0]), float(beta[1]), int(keep.sum()), led


RATIONALE = {"window": "post-1984 (Great Moderation) sample, chosen a priori",
             "lag_pairs": "AR(1) needs one lag",
             "pandemic": "extreme COVID quarters would dominate OLS"}


def main():
    n_rows = {}
    fr = pd.read_csv(os.path.join(DATA, "first_release.csv"), parse_dates=["quarter", "first_release_vintage_date"])
    n_rows["first_release.csv"] = len(fr)
    cur = read_series(os.path.join(DATA, "current_vintage.csv"))  # revision statistic ONLY
    n_rows["current_vintage.csv"] = len(cur)
    frmap = dict(zip(fr.quarter, fr.first_release_value))
    frvint = dict(zip(fr.quarter, fr.first_release_vintage_date))

    rows, excl = [], []
    for o in ORIGINS:
        rel = f"vintages/A191RL1Q225SBEA_vintage_{o}.csv"
        s = read_series(os.path.join(DATA, rel))  # the ONLY data used at this origin
        n_rows[rel] = len(s)
        assert s.value.notna().all() and not s.date.duplicated().any()
        last_d, last_v = s.date.iloc[-1], float(s.value.iloc[-1])
        c, phi, n_used, led = ar1_fit(s)
        for step, rule, nb, na in led:
            excl.append(dict(step=f"{o}:{step}", rule=rule, n_before=nb, n_after=na,
                             n_removed=nb - na, rationale=RATIONALE[step]))
        f = last_v
        for h in (1, 2):
            tgt = last_d + pd.DateOffset(months=3 * h)
            assert tgt in frmap, f"no first release for {tgt}"
            assert frvint[tgt] > pd.Timestamp(o), "target first release not after origin"
            f = c + phi * f  # iterated AR(1) forecast
            for model, fc, n in (("naive", last_v, len(s)), (MODEL, f, n_used)):
                rows.append(dict(origin_date=o, target_date=tgt.strftime("%Y-%m-%d"), horizon=h,
                                 forecast=round(fc, 6), actual=float(frmap[tgt]), model=model,
                                 vintage_date=o, train_last_date=last_d.strftime("%Y-%m-%d"),
                                 train_last_value=last_v, n_train=n))
        print(f"{o}: last {last_d.date()}={last_v}, AR(1) c={c:.3f} phi={phi:.3f} n_pairs={n_used}")

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

    res = dict(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, "
                                  "sample 1985Q1 to last published quarter, excluding regression pairs involving "
                                  "2020Q2-2020Q3; h=2 forecast iterated from the h=1 forecast."),
               n_origins=len(ORIGINS), first_origin=ORIGINS[0], last_origin=ORIGINS[-1])
    for h in (1, 2):
        for m, tag in (("naive", "naive"), (MODEL, "model")):
            e = fc[(fc.horizon == h) & (fc.model == m)].eval("forecast - actual").values
            assert len(e) == 8
            res[f"rmse_{tag}_h{h}"] = round(float(np.sqrt(np.mean(e ** 2))), 6)
            res[f"mae_{tag}_h{h}"] = round(float(np.mean(np.abs(e))), 6)
    for h in (1, 2):
        res[f"rel_improvement_rmse_h{h}"] = round(1 - res[f"rmse_model_h{h}"] / res[f"rmse_naive_h{h}"], 6)
    for h in (1, 2):  # descriptive loss-differential summary (n=8: indicative only)
        en = fc[(fc.horizon == h) & (fc.model == "naive")].eval("forecast - actual").values
        em = fc[(fc.horizon == h) & (fc.model == MODEL)].eval("forecast - actual").values
        d = en ** 2 - em ** 2
        res[f"n_origins_model_better_sqerr_h{h}"] = int((d > 0).sum())
        res[f"dm_tstat_sqerr_h{h}"] = round(float(d.mean() / (d.std(ddof=1) / np.sqrt(len(d)))), 6)
        res[f"dm_pvalue_two_sided_h{h}"] = round(float(2 * stats.t.sf(abs(res[f"dm_tstat_sqerr_h{h}"]), len(d) - 1)), 6)
    t1 = pd.to_datetime(fc[(fc.horizon == 1) & (fc.model == "naive")].target_date)
    cmap = dict(zip(cur.date, cur.value))
    res["mean_abs_revision_h1_targets"] = round(float(np.mean([abs(cmap[t] - frmap[t]) for t in t1])), 6)
    res["feasibility"] = "complete"
    res["feasibility_reason"] = None
    res["data_issues"] = [
        {"issue": "Pandemic outliers 2020Q2 (about -28) and 2020Q3 (about +35) present in every vintage.",
         "handling": "Regression pairs involving these quarters excluded from AR(1) estimation; naive unaffected."},
        {"issue": "Origins unevenly spaced: 2025Q3 advance estimate delayed by the 2025 federal shutdown, so origin 8 is 2025-12-31.",
         "handling": "Used as specified; target 2025Q4 is the h=2 target of origin 7 and the h=1 target of origin 8."},
        {"issue": "current_vintage.csv contains revisions unknown at the origins.",
         "handling": "Used only for mean_abs_revision_h1_targets, never for estimation or forecast input."},
    ]
    res["n_input_rows"] = n_rows
    res["n_analysed"] = n_rows[f"vintages/A191RL1Q225SBEA_vintage_{ORIGINS[-1]}.csv"]
    res["n_model_estimation_pairs_last_origin"] = int(fc[fc.model == MODEL].n_train.iloc[-1])
    res["random_seed"] = None
    with open(os.path.join(OUT, "results.json"), "w") as fh:
        json.dump(res, fh, indent=2)
    print(json.dumps({k: v for k, v in res.items() if k.startswith(("rmse", "mae", "rel", "mean"))}, indent=1))


if __name__ == "__main__":
    main()