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

DATA, OUT = "data", "out"
EST_START = pd.Timestamp("1985-01-01")             # first dependent quarter in AR sample
COVID_DEP = [pd.Timestamp(d) for d in ("2020-04-01", "2020-07-01", "2020-10-01")]
MODEL = "ar1_ols"

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

def read_series(path):
    d = pd.read_csv(path)
    s = pd.Series(d.iloc[:, 1].astype(float).values, index=pd.to_datetime(d.iloc[:, 0]))
    return s, len(d)

n_input_rows = {}
cur, n_input_rows["current_vintage.csv"] = read_series(f"{DATA}/current_vintage.csv")
fr = pd.read_csv(f"{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 = fr.set_index("quarter")

vfiles = sorted(glob.glob(f"{DATA}/vintages/A191RL1Q225SBEA_vintage_*.csv"))
data_issues = [
    {"issue": "COVID-19 outliers 2020Q2 (about -28%) and 2020Q3 (about +35%) in every vintage would dominate OLS AR estimates",
     "handling": "AR(1) estimation sample excludes dependent quarters 2020Q2-2020Q4, so 2020Q2/2020Q3 enter neither as regressand nor regressor; naive forecast unaffected"},
    {"issue": "Eighth origin is 2025-12-31 (2025Q3 advance estimate delayed by the 2025 federal shutdown) and 2025Q4 first release dated 2026-02-28",
     "handling": "Used as given; at that origin the h=1/h=2 targets are 2025Q4/2026Q1 as defined by the task"},
    {"issue": "Historical quarters are revised between vintages (e.g. 2025Q1: -0.3, -0.5, -0.6 in successive vintages)",
     "handling": "At each origin only that origin's vintage file is used; current vintage used only for the revision statistic"},
]

rows, excl = [], None
for vf in vfiles:
    origin = os.path.basename(vf).split("_vintage_")[1][:-4]
    y, n = read_series(vf)
    n_input_rows["vintages/" + os.path.basename(vf)] = n
    assert y.index.is_monotonic_increasing and not y.index.duplicated().any()
    T = y.index[-1]; yT = float(y.iloc[-1])
    # sanity: last value equals first release of that quarter (if listed)
    if T in fr.index:
        assert abs(fr.loc[T, "first_release_value"] - yT) < 1e-9, (origin, T)
    # AR(1) OLS estimation sample (dependent quarters t, regressor y_{t-1}) from this vintage only
    ylag = y.shift(1)
    dep = pd.DataFrame({"y": y, "ylag": ylag}).dropna()
    s1 = dep[dep.index >= EST_START]
    s2 = s1[~s1.index.isin(COVID_DEP)]
    X = np.column_stack([np.ones(len(s2)), s2["ylag"].values])
    beta, *_ = np.linalg.lstsq(X, s2["y"].values, rcond=None)
    c, phi = float(beta[0]), float(beta[1])
    f_ar = {1: c + phi * yT}; f_ar[2] = c + phi * f_ar[1]
    for h in (1, 2):
        tgt = qshift(T, h)
        assert pd.Timestamp(origin) < fr.loc[tgt, "first_release_vintage_date"], "look-ahead"
        act = float(fr.loc[tgt, "first_release_value"])
        base = dict(origin_date=origin, target_date=tgt.strftime("%Y-%m-%d"), horizon=h, actual=act,
                    vintage_date=origin, train_last_date=T.strftime("%Y-%m-%d"), train_last_value=yT)
        rows.append({**base, "forecast": yT, "model": "naive", "n_train": n})
        rows.append({**base, "forecast": round(f_ar[h], 10), "model": MODEL, "n_train": len(s2)})
    last = dict(origin=origin, n=n, n_dep=len(dep), n1=len(s1), n2=len(s2), c=c, phi=phi)

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(f"{OUT}/forecasts.csv", index=False)

# exclusion ledger for the last origin's model training data
L = last
excl = pd.DataFrame([
    ["1_lag", "drop first quarter (no lagged value for AR(1))", L["n"], L["n_dep"], L["n"] - L["n_dep"], "AR(1) needs y_{t-1}"],
    ["2_window", "keep dependent quarters >= 1985-01-01", L["n_dep"], L["n1"], L["n_dep"] - L["n1"], "post-1984 (Great Moderation) dynamics, fixed a priori"],
    ["3_covid", "drop dependent quarters 2020Q2-2020Q4", L["n1"], L["n2"], L["n1"] - L["n2"], "COVID outliers would dominate OLS"],
], columns=["step", "rule", "n_before", "n_after", "n_removed", "rationale"])
excl.to_csv(f"{OUT}/exclusions.csv", index=False)

def err(model, h):
    d = fc[(fc.model == model) & (fc.horizon == h)]
    e = d["forecast"] - d["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 on each origin's vintage only, dependent quarters 1985Q1 to the "
                             "vintage's last quarter excluding 2020Q2-2020Q4, re-estimated at every origin; h=2 forecast by iteration."),
       "n_origins": len(vfiles), "first_origin": fc.origin_date.min(), "last_origin": fc.origin_date.max()}
for m, tag in (("naive", "naive"), (MODEL, "model")):
    for h in (1, 2):
        r, a = err(m, h)
        res[f"rmse_{tag}_h{h}"] = r; res[f"mae_{tag}_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}"]
# small-sample comparison diagnostics: squared-error loss differential (naive - model)
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")
    ea, eb = (a.forecast - a.actual), (b.forecast - b.actual)
    d = (ea ** 2 - eb ** 2).values
    n_ = len(d); dbar = d.mean(); g = [np.mean((d[k:] - dbar) * (d[:n_ - k] - dbar)) for k in range(h)]
    lrv = g[0] + 2 * sum(g[1:])                     # rectangular HAC with h-1 lags (Diebold-Mariano)
    res[f"dm_stat_h{h}"] = float(dbar / np.sqrt(lrv / n_))
    hln = np.sqrt((n_ + 1 - 2 * h + h * (h - 1) / n_) / n_)  # Harvey-Leybourne-Newbold small-sample correction
    res[f"dm_hln_pvalue_two_sided_h{h}"] = float(2 * stats.t.sf(abs(hln * res[f"dm_stat_h{h}"]), n_ - 1))
    res[f"n_model_smaller_abs_error_h{h}"] = int((eb.abs() < ea.abs()).sum())
t1 = pd.to_datetime(fc[(fc.horizon == 1) & (fc.model == "naive")]["target_date"])
rev = [abs(cur.loc[t] - fr.loc[t, "first_release_value"]) for t in t1]
res["mean_abs_revision_h1_targets"] = float(np.mean(rev))
res.update({"feasibility": "complete", "feasibility_reason": None, "data_issues": data_issues,
            "n_input_rows": n_input_rows, "n_analysed": int(L["n2"]), "random_seed": None,
            "last_origin_ar1_intercept": L["c"], "last_origin_ar1_phi": L["phi"],
            "last_origin_ar1_implied_mean": L["c"] / (1 - L["phi"])})
with open(f"{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 not in ("data_issues", "n_input_rows")}, indent=2))