"""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."""
import json, os, glob
import numpy as np
import 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"]
SAMPLE_START = pd.Timestamp("1985-01-01")
EXCLUDED_Q = [pd.Timestamp("2020-04-01"), pd.Timestamp("2020-07-01")]
MODEL = "ar1_post1985"
DESC = ("AR(1) with intercept estimated by OLS on the origin vintage, sample 1985Q1 to the last quarter "
        "of that vintage, with 2020Q2 and 2020Q3 set to missing (no regression pair uses them); "
        "h=2 forecast obtained by iterating the AR(1) recursion.")

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

def qshift(ts, k):
    return ts + pd.DateOffset(months=3 * k)

n_input_rows = {}
vint = {}
for o in ORIGINS:
    rel = f"vintages/A191RL1Q225SBEA_vintage_{o}.csv"
    s = read_series(os.path.join(DATA, rel))
    vint[o] = s
    n_input_rows[rel] = int(len(s))
cur = read_series(os.path.join(DATA, "current_vintage.csv"))  # used ONLY for revision statistic
n_input_rows["current_vintage.csv"] = int(len(cur))
fr = pd.read_csv(os.path.join(DATA, "first_release.csv"), parse_dates=["quarter", "first_release_vintage_date"])
n_input_rows["first_release.csv"] = int(len(fr))
fr_val = fr.set_index("quarter")["first_release_value"].astype(float)
fr_vd = fr.set_index("quarter")["first_release_vintage_date"]

data_issues = []
for o, s in vint.items():
    assert s.index.is_monotonic_increasing and not s.index.duplicated().any() and s.notna().all()
    assert (s.index.to_series().diff().dropna().dt.days <= 92).all()
data_issues.append({"issue": "Pandemic outliers 2020Q2/2020Q3 (about -28% and +35% SAAR) in every vintage.",
                    "handling": "Set to missing in the AR(1) estimation sample (documented in out/exclusions.csv); naive forecasts unaffected."})
data_issues.append({"issue": "Historical values differ across vintages (annual revisions, e.g. 2024Q4 2.3 -> 2.4 -> 1.9 across vintages).",
                    "handling": "Each origin uses only its own vintage file; current vintage used only for the revision statistic."})
data_issues.append({"issue": "Eighth origin is 2025-12-31 rather than 2025-10-31 (2025Q3 advance estimate delayed by the federal shutdown); its h=1 target 2025Q4 first release dated 2026-02-28.",
                    "handling": "Used as given in TASK.md; origins are not equally spaced in calendar time but each still forecasts 1 and 2 quarters past its last published quarter."})

def fit_ar1(s):
    y = s[s.index >= SAMPLE_START].copy()
    n_after_start = len(y)
    y.loc[y.index.isin(EXCLUDED_Q)] = np.nan
    lag = y.shift(1)
    # require the lag to be the immediately preceding quarter (index is contiguous quarterly)
    m = y.notna() & lag.notna()
    X = np.column_stack([np.ones(m.sum()), lag[m].values])
    beta, *_ = np.linalg.lstsq(X, y[m].values, rcond=None)
    return beta, int(y.notna().sum()), n_after_start, int(m.sum())

rows, excl = [], []
fits = {}
for o in ORIGINS:
    s = vint[o]
    T, yT = s.index[-1], float(s.iloc[-1])
    beta, n_used, n_after_start, n_pairs = fit_ar1(s)
    fits[o] = {"c": float(beta[0]), "phi": float(beta[1]), "n_pairs": n_pairs}
    excl.append({"step": f"{o}:sample_start", "rule": "drop quarters before 1985-01-01 (model only)",
                 "n_before": len(s), "n_after": n_after_start, "n_removed": len(s) - n_after_start,
                 "rationale": "post-Great-Moderation estimation sample for AR(1)"})
    excl.append({"step": f"{o}:pandemic", "rule": "set 2020-04-01 and 2020-07-01 to missing (model only)",
                 "n_before": n_after_start, "n_after": n_used, "n_removed": n_after_start - n_used,
                 "rationale": "extreme pandemic swings would dominate OLS AR coefficient"})
    f_ar = {}
    prev = yT
    for h in (1, 2):
        prev = beta[0] + beta[1] * prev
        f_ar[h] = float(prev)
    for h in (1, 2):
        tgt = qshift(T, h)
        assert tgt in fr_val.index, f"missing first release for {tgt}"
        assert fr_vd[tgt] >= pd.Timestamp(o), "actual published before origin"
        for model, fc, ntr in (("naive", yT, len(s)), (MODEL, f_ar[h], n_used)):
            rows.append({"origin_date": o, "target_date": tgt.strftime("%Y-%m-%d"), "horizon": h,
                         "forecast": round(float(fc), 10), "actual": float(fr_val[tgt]), "model": model,
                         "vintage_date": o, "train_last_date": T.strftime("%Y-%m-%d"),
                         "train_last_value": yT, "n_train": int(ntr)})

fc = pd.DataFrame(rows)
fc["error"] = fc["actual"] - fc["forecast"]
fc.drop(columns="error").to_csv(os.path.join(OUT, "forecasts.csv"), index=False)
# Main ledger: single chain for the last origin (whose model training set defines n_analysed);
# the same two rules applied at every origin are listed in exclusions_by_origin.csv.
pd.DataFrame([e for e in excl if e["step"].startswith(ORIGINS[-1])]).to_csv(os.path.join(OUT, "exclusions.csv"), index=False)
pd.DataFrame(excl).to_csv(os.path.join(OUT, "exclusions_by_origin.csv"), index=False)

def metrics(model, h):
    e = fc[(fc.model == model) & (fc.horizon == h)]["error"].values
    assert len(e) == len(ORIGINS)
    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": DESC,
       "n_origins": len(ORIGINS), "first_origin": ORIGINS[0], "last_origin": ORIGINS[-1]}
for model, tag in (("naive", "naive"), (MODEL, "model")):
    for h in (1, 2):
        r, a = metrics(model, h)
        res[f"rmse_{tag}_h{h}"] = round(r, 10)
        res[f"mae_{tag}_h{h}"] = round(a, 10)
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}"], 10)
h1t = pd.to_datetime(fc[(fc.model == "naive") & (fc.horizon == 1)]["target_date"])
rev = [abs(float(cur[t]) - float(fr_val[t])) for t in h1t]
res["mean_abs_revision_h1_targets"] = round(float(np.mean(rev)), 10)
res["feasibility"] = "complete"
res["feasibility_reason"] = None
res["data_issues"] = data_issues
res["n_input_rows"] = n_input_rows
res["n_analysed"] = int(fc[(fc.origin_date == ORIGINS[-1]) & (fc.model == MODEL)]["n_train"].iloc[0])
res["random_seed"] = None
# supplementary diagnostics (not required keys)
res["supp_ar1_coefficients_by_origin"] = fits
res["supp_abs_revision_h1_targets"] = {t.strftime("%Y-%m-%d"): round(r, 10) for t, r in zip(h1t, rev)}
for model, tag in (("naive", "naive"), (MODEL, "model")):
    for h in (1, 2):
        res[f"supp_mean_error_{tag}_h{h}"] = round(float(fc[(fc.model == model) & (fc.horizon == h)]["error"].mean()), 10)
# Diebold-Mariano-type comparison (squared-error loss), simple t-stat on loss differential, small-sample caveat
from scipy import stats
for h in (1, 2):
    en = fc[(fc.model == "naive") & (fc.horizon == h)]["error"].values
    em = fc[(fc.model == MODEL) & (fc.horizon == h)]["error"].values
    d = en ** 2 - em ** 2
    n = len(d)
    # HAC variance with h-1 lags (rectangular), fallback to plain variance if non-positive
    g0 = np.mean((d - d.mean()) ** 2)
    v = g0 + sum(2 * np.mean((d[k:] - d.mean()) * (d[:-k] - d.mean())) for k in range(1, h))
    if v <= 0:
        v = g0
    dm = float(d.mean() / np.sqrt(v / n))
    hln = dm * np.sqrt((n + 1 - 2 * h + h * (h - 1) / n) / n)
    res[f"supp_dm_hln_stat_h{h}"] = round(float(hln), 10)
    res[f"supp_dm_hln_pvalue_two_sided_h{h}"] = round(float(2 * stats.t.sf(abs(hln), df=n - 1)), 10)

with open(os.path.join(OUT, "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_abs", "supp_dm"))}, indent=2))