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

DATA = "data"; OUT = "out"
MODEL = "ar1_1985_xcovid"
EST_START = pd.Timestamp("1985-01-01")
COVID = {pd.Timestamp("2020-04-01"), pd.Timestamp("2020-07-01")}
issues = []
n_rows = {}

def read_series(rel):
    df = pd.read_csv(os.path.join(DATA, rel))
    n_rows[rel] = int(len(df))
    df.columns = ["date", "value"]
    df["date"] = pd.to_datetime(df["date"])
    df["value"] = pd.to_numeric(df["value"], errors="coerce")
    nmiss = int(df["value"].isna().sum())
    if nmiss:
        issues.append({"issue": f"{rel}: {nmiss} missing/non-numeric values",
                       "handling": "dropped"})
        df = df.dropna()
    s = df.set_index("date")["value"].sort_index()
    exp = pd.date_range(s.index[0], s.index[-1], freq="QS")
    if len(exp) != len(s) or not (exp == s.index).all():
        issues.append({"issue": f"{rel}: irregular/missing quarters", "handling": "noted; AR uses only consecutive pairs"})
    if s.index.duplicated().any():
        issues.append({"issue": f"{rel}: duplicate dates", "handling": "kept last"})
        s = s[~s.index.duplicated(keep="last")]
    return s

fr = pd.read_csv(os.path.join(DATA, "first_release.csv"))
n_rows["first_release.csv"] = int(len(fr))
fr["quarter"] = pd.to_datetime(fr["quarter"])
fr_map = dict(zip(fr["quarter"], fr["first_release_value"]))
cur = read_series("current_vintage.csv")

def fit_ar1(s):
    """OLS AR(1) with intercept on vintage s, from EST_START, dropping pairs
    where y_t or y_{t-1} is a COVID quarter (2020Q2, 2020Q3)."""
    s = s[s.index >= EST_START - pd.offsets.QuarterBegin(startingMonth=1)]
    y = s.values[1:]; x = s.values[:-1]
    dt = s.index[1:]; dl = s.index[:-1]
    consecutive = np.array([(a + pd.offsets.QuarterBegin(startingMonth=1)) == b for a, b in zip(dl, dt)])
    keep = consecutive & np.array([(a not in COVID) and (b not in COVID) for a, b in zip(dl, dt)]) & (dt >= EST_START)
    y, x = y[keep], x[keep]
    X = np.column_stack([np.ones_like(x), x])
    beta, *_ = np.linalg.lstsq(X, y, rcond=None)
    return beta, int(keep.sum())

files = sorted(glob.glob(os.path.join(DATA, "vintages", "A191RL1Q225SBEA_vintage_*.csv")))
rows = []
for f in files:
    rel = os.path.relpath(f, DATA)
    vdate = os.path.basename(f).split("_")[-1].replace(".csv", "")
    s = read_series(rel)
    last_d, last_v = s.index[-1], float(s.iloc[-1])
    if last_d in fr_map and abs(fr_map[last_d] - last_v) > 1e-9:
        issues.append({"issue": f"{rel}: last value {last_v} differs from first release {fr_map[last_d]}",
                       "handling": "vintage value used as training data (per rules)"})
    beta, n_est = fit_ar1(s)
    f_prev = last_v
    for h in (1, 2):
        tgt = last_d + pd.DateOffset(months=3 * h)
        f_prev = beta[0] + beta[1] * f_prev
        actual = fr_map.get(tgt, np.nan)
        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": round(float(f_prev), 6), "model": MODEL, "n_train": n_est})
    print(vdate, "beta", beta.round(3), "n_est", n_est)

fc = pd.DataFrame(rows)[["origin_date", "target_date", "horizon", "forecast", "actual", "model",
                         "vintage_date", "train_last_date", "train_last_value", "n_train"]]
fc.to_csv(os.path.join(OUT, "forecasts.csv"), index=False)
if fc["actual"].isna().any():
    issues.append({"issue": "missing first-release actuals", "handling": "excluded from metrics"})

def err(model, h):
    d = fc[(fc.model == model) & (fc.horizon == h)].dropna(subset=["actual"])
    e = d["forecast"] - d["actual"]
    return float(np.sqrt((e ** 2).mean())), float(e.abs().mean())

res = {"task_id": "F3-conf-02", "model_name": MODEL,
       "model_description": ("AR(1) with intercept estimated by OLS on each origin's vintage from 1985Q1 to the "
                             "last published quarter, excluding regression pairs involving the 2020Q2/2020Q3 "
                             "COVID quarters; h=2 forecast by iterating the AR(1)."),
       "n_origins": len(files),
       "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}"] = round(r, 6); res[f"mae_{tag}_h{h}"] = round(a, 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)
t1 = pd.to_datetime(fc[(fc.horizon == 1) & (fc.model == "naive")]["target_date"])
revs = [abs(cur[t] - fr_map[t]) for t in t1]
res["mean_abs_revision_h1_targets"] = round(float(np.mean(revs)), 6)
# Origin-spacing check (shutdown-delayed 2025Q3 advance estimate)
od = pd.to_datetime(sorted(fc.origin_date.unique()))
gaps = [(b - a).days for a, b in zip(od[:-1], od[1:])]
for a, b, g in zip(od[:-1], od[1:], gaps):
    if g > 100:
        issues.append({"issue": f"irregular origin spacing: {a.date()} -> {b.date()} ({g} days), due to delayed BEA advance estimate (2025 shutdown)",
                       "handling": "kept as given; h is defined relative to last quarter in the vintage, so targets stay consecutive quarters"})
# Small-sample Diebold-Mariano (Harvey-Leybourne-Newbold) test, squared-error loss
from scipy import stats
for h in (1, 2):
    a = fc[(fc.model == "naive") & (fc.horizon == h)].sort_values("origin_date")
    b = fc[(fc.model == MODEL) & (fc.horizon == h)].sort_values("origin_date")
    d = (a.forecast.values - a.actual.values) ** 2 - (b.forecast.values - b.actual.values) ** 2
    T = len(d); dbar = d.mean(); dc = d - dbar
    lrv = (dc @ dc) / T + 2 * sum((dc[k:] @ dc[:-k]) / T for k in range(1, h))
    dm = dbar / np.sqrt(lrv / T)
    hln = dm * np.sqrt((T + 1 - 2 * h + h * (h - 1) / T) / T)
    res[f"dm_hln_stat_h{h}"] = round(float(hln), 4)
    res[f"dm_hln_pvalue_h{h}"] = round(float(2 * stats.t.sf(abs(hln), T - 1)), 4)
res["feasibility"] = "complete"; res["feasibility_reason"] = None
res["data_issues"] = issues
res["n_input_rows"] = n_rows
res["n_analysed"] = int(len(read_series(os.path.relpath(files[-1], DATA))))
res["random_seed"] = None
with open(os.path.join(OUT, "results.json"), "w") as fh:
    json.dump(res, fh, indent=2)
print(json.dumps(res, indent=2))