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

DATA, OUT = "data", "out"
MODEL = "ar1_post1985_excovid"
SAMPLE_START = pd.Timestamp("1985-01-01")            # first dependent-variable quarter
COVID_EXCL = pd.to_datetime(["2020-04-01", "2020-07-01", "2020-10-01"])  # rows whose y_t or y_{t-1} is 2020Q2/Q3

def read_series(path):
    d = pd.read_csv(path, parse_dates=["observation_date"])
    return pd.Series(d.iloc[:, 1].astype(float).values, index=d["observation_date"])

def qshift(ts, k):
    return (pd.Period(ts, freq="Q") + k).start_time

def ar1_fit(s, ledger=None):
    """OLS AR(1) with intercept on the vintage series s."""
    n0 = len(s)
    df = pd.DataFrame({"y": s, "ylag": s.shift(1)})
    df1 = df[(df.index >= SAMPLE_START) & df["ylag"].notna()]
    df2 = df1[~df1.index.isin(COVID_EXCL)]
    if ledger is not None:
        ledger.append(("1", f"keep regression rows with dependent quarter >= {SAMPLE_START.date()} (lag needs prior quarter)",
                       n0, len(df1), n0 - len(df1), "post-1984 (Great Moderation) sample; pre-1985 volatility regime not representative of 2024-26"))
        ledger.append(("2", "drop rows t in {2020Q2,2020Q3,2020Q4} (y_t or y_{t-1} is a COVID quarter)",
                       len(df1), len(df2), len(df1) - len(df2), "pandemic outliers (~-28, ~+35 pp) would dominate OLS"))
    X = np.column_stack([np.ones(len(df2)), df2["ylag"].values])
    beta, *_ = np.linalg.lstsq(X, df2["y"].values, rcond=None)
    return beta, len(df2)

def main():
    vfiles = sorted(glob.glob(os.path.join(DATA, "vintages", "A191RL1Q225SBEA_vintage_*.csv")))
    cur = read_series(os.path.join(DATA, "current_vintage.csv"))
    fr = pd.read_csv(os.path.join(DATA, "first_release.csv"),
                     parse_dates=["quarter", "first_release_vintage_date"]).set_index("quarter")
    n_input_rows = {"current_vintage.csv": int(len(cur)), "first_release.csv": int(len(fr))}
    rows, h1_targets, coefs, ledger = [], [], {}, []
    for i, f in enumerate(vfiles):
        vdate = os.path.basename(f).split("_")[-1].replace(".csv", "")
        s = read_series(f)
        n_input_rows[os.path.relpath(f, DATA)] = int(len(s))
        assert s.index.is_monotonic_increasing and not s.index.duplicated().any() and s.notna().all()
        last_d, last_v = s.index[-1], float(s.iloc[-1])
        is_last = i == len(vfiles) - 1
        beta, n_reg = ar1_fit(s, ledger if is_last else None)
        coefs[vdate] = [float(beta[0]), float(beta[1]), n_reg]
        f1 = beta[0] + beta[1] * last_v
        f2 = beta[0] + beta[1] * f1
        for h, fm in ((1, f1), (2, f2)):
            tgt = qshift(last_d, h)
            actual = float(fr.loc[tgt, "first_release_value"])
            assert fr.loc[tgt, "first_release_vintage_date"] > pd.Timestamp(vdate)  # chronology
            base = dict(origin_date=vdate, target_date=tgt.strftime("%Y-%m-%d"), horizon=h, actual=actual,
                        vintage_date=vdate, train_last_date=last_d.strftime("%Y-%m-%d"), train_last_value=last_v)
            rows.append({**base, "forecast": last_v, "model": "naive", "n_train": int(len(s))})
            rows.append({**base, "forecast": float(fm), "model": MODEL, "n_train": int(n_reg)})
            if h == 1:
                h1_targets.append(tgt)
        if is_last:
            n_analysed = n_reg
    fc = pd.DataFrame(rows)[["origin_date", "target_date", "horizon", "forecast", "actual", "model",
                             "vintage_date", "train_last_date", "train_last_value", "n_train"]]
    fc = fc.sort_values(["origin_date", "horizon", "model"]).reset_index(drop=True)
    fc.to_csv(os.path.join(OUT, "forecasts.csv"), index=False, float_format="%.10g")
    pd.DataFrame(ledger, columns=["step", "rule", "n_before", "n_after", "n_removed", "rationale"]).to_csv(
        os.path.join(OUT, "exclusions.csv"), index=False)

    def err(model, h):
        d = fc[(fc.model == model) & (fc.horizon == h)]
        e = d["forecast"] - d["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 at each origin on that origin's vintage only, "
                                 "dependent quarters 1985Q1 to the last vintage quarter, excluding regression rows t=2020Q2-2020Q4 "
                                 "(COVID outliers in y_t or y_{t-1}); h=2 forecast iterated from the h=1 forecast."),
           "n_origins": len(vfiles),
           "first_origin": fc.origin_date.min(), "last_origin": fc.origin_date.max()}
    for h in (1, 2):
        res[f"rmse_naive_h{h}"], res[f"mae_naive_h{h}"] = err("naive", h)
        res[f"rmse_model_h{h}"], res[f"mae_model_h{h}"] = err(MODEL, h)
    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}"]
    for h in (1, 2):
        a = fc[(fc.model == "naive") & (fc.horizon == h)].set_index("origin_date")
        b = fc[(fc.model == MODEL) & (fc.horizon == h)].set_index("origin_date")
        res[f"n_origins_model_abs_error_smaller_h{h}"] = int(((b.forecast - b.actual).abs() < (a.forecast - a.actual).abs()).sum())
    rev = [abs(float(cur.loc[t]) - float(fr.loc[t, "first_release_value"])) for t in h1_targets]
    res["mean_abs_revision_h1_targets"] = float(np.mean(rev))
    res.update({
        "feasibility": "complete", "feasibility_reason": None,
        "data_issues": [
            {"issue": "COVID-19 outliers 2020Q2 (~-28) and 2020Q3 (~+35) in every vintage",
             "handling": "AR regression rows t=2020Q2-2020Q4 excluded (dependent or lag is a COVID quarter); naive unaffected"},
            {"issue": "Irregular origin spacing: 8th origin 2025-12-31 instead of 2025-10-31 (2025Q3 advance estimate delayed by federal shutdown)",
             "handling": "Used as provided; horizons defined relative to last quarter in each vintage, so targets are 2025Q4 and 2026Q1"},
            {"issue": "Values of earlier quarters differ across vintages (revisions)",
             "handling": "Each origin uses only its own vintage; current vintage used only for revision statistic"}],
        "n_input_rows": n_input_rows, "n_analysed": int(n_analysed), "random_seed": None,
        "ar1_coefficients_by_origin": {k: {"intercept": v[0], "phi": v[1], "n_reg": v[2]} for k, v in coefs.items()},
        "abs_revision_h1_targets": {t.strftime("%Y-%m-%d"): r for t, r in zip(h1_targets, rev)},
    })
    with open(os.path.join(OUT, "results.json"), "w") as fh:
        json.dump(res, fh, indent=2)
    print(json.dumps({k: res[k] for k in res if k.startswith(("rmse", "mae", "rel", "mean"))}, indent=2))

if __name__ == "__main__":
    main()