"""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
from scipy import stats

DATA = "data"; OUT = "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_covid_excl"
# Equations whose dependent variable or lag falls in 2020Q2/2020Q3 are dropped
COVID = pd.to_datetime(["2020-04-01", "2020-07-01"])

n_rows, issues = {}, []

def read_series(rel):
    d = pd.read_csv(os.path.join(DATA, rel))
    n_rows[rel] = len(d)
    d.columns = ["date", "value"]
    d["date"] = pd.to_datetime(d["date"])
    d["value"] = pd.to_numeric(d["value"], errors="coerce")
    exp = pd.date_range(d.date.iloc[0], d.date.iloc[-1], freq="QS")
    if d.value.isna().any() or d.date.duplicated().any() or len(exp) != len(d):
        issues.append({"issue": f"{rel}: missing/duplicate/gapped quarters",
                       "handling": "rows with missing values dropped"})
        d = d.dropna()
    return d.set_index("date")["value"]

fr = pd.read_csv(os.path.join(DATA, "first_release.csv"))
n_rows["first_release.csv"] = len(fr)
fr["quarter"] = pd.to_datetime(fr["quarter"])
first = fr.set_index("quarter")["first_release_value"]
cur = read_series("current_vintage.csv")  # used ONLY for the revision statistic

def fit_ar1(y):
    """OLS AR(1) with intercept; drop equations touching COVID quarters."""
    yt, yl = y.values[1:], y.values[:-1]
    dt, dl = y.index[1:], y.index[:-1]
    keep = ~(dt.isin(COVID) | dl.isin(COVID))
    X = np.column_stack([np.ones(keep.sum()), yl[keep]])
    b, *_ = np.linalg.lstsq(X, yt[keep], rcond=None)
    return b, int(keep.sum())

rows, fit_info = [], []
for o in ORIGINS:
    rel = f"vintages/A191RL1Q225SBEA_vintage_{o}.csv"
    y = read_series(rel)
    last_d, last_v = y.index[-1], float(y.iloc[-1])
    (c, phi), neq = fit_ar1(y)
    fit_info.append({"origin": o, "const": c, "phi": phi, "n_eq": neq})
    f = last_v
    for h in (1, 2):
        tgt = last_d + pd.DateOffset(months=3 * h)
        f = c + phi * f
        act = float(first[tgt]) if tgt in first.index else np.nan
        base = dict(origin_date=o, target_date=tgt.strftime("%Y-%m-%d"), horizon=h,
                    actual=act, vintage_date=o, train_last_date=last_d.strftime("%Y-%m-%d"),
                    train_last_value=last_v, n_train=len(y))
        rows.append({**base, "forecast": last_v, "model": "naive"})
        rows.append({**base, "forecast": round(float(f), 6), "model": MODEL})
    n_last = len(y)

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)
fc["err"] = fc["actual"] - fc["forecast"]

res = {}
for h in (1, 2):
    for m, lab in (("naive", "naive"), (MODEL, "model")):
        e = fc[(fc.horizon == h) & (fc.model == m)]["err"]
        res[f"rmse_{lab}_h{h}"] = float(np.sqrt((e ** 2).mean()))
        res[f"mae_{lab}_h{h}"] = float(e.abs().mean())
    res[f"rel_improvement_rmse_h{h}"] = 1 - res[f"rmse_model_h{h}"] / res[f"rmse_naive_h{h}"]

t1 = pd.to_datetime(fc[(fc.horizon == 1) & (fc.model == "naive")]["target_date"])
rev = [abs(float(cur[t]) - float(first[t])) for t in t1]
mean_rev = float(np.mean(rev))

# Diebold-Mariano with Harvey-Leybourne-Newbold small-sample correction (squared loss)
dm = {}
for h in (1, 2):
    en = fc[(fc.horizon == h) & (fc.model == "naive")]["err"].values
    em = fc[(fc.horizon == h) & (fc.model == MODEL)]["err"].values
    d = en ** 2 - em ** 2; T = len(d); dbar = d.mean()
    g = [np.mean((d[k:] - dbar) * (d[:T - k] - dbar)) for k in range(h)]
    v = (g[0] + 2 * sum(g[1:])) / T
    if v <= 0:
        v = g[0] / T
    stat = dbar / np.sqrt(v) * np.sqrt((T + 1 - 2 * h + h * (h - 1) / T) / T)
    dm[f"h{h}"] = {"stat": float(stat), "p_two_sided": float(2 * stats.t.sf(abs(stat), T - 1)),
                   "mean_loss_diff_naive_minus_model": float(dbar)}

# Descriptive robustness: per-origin wins and leave-one-origin-out relative improvement
robust = {}
for h in (1, 2):
    a = fc[(fc.horizon == h)].pivot(index="origin_date", columns="model", values="err")
    wins = int((a[MODEL].abs() < a["naive"].abs()).sum())
    loo = {o: float(1 - np.sqrt((a[MODEL].drop(o) ** 2).mean()) / np.sqrt((a["naive"].drop(o) ** 2).mean()))
           for o in a.index}
    robust[f"h{h}"] = {"n_origins_model_abs_error_smaller": wins, "loo_rel_improvement_rmse": loo}

issues.append({"issue": "2020Q2 (-28) and 2020Q3 (+35) COVID outliers in every vintage would dominate OLS AR estimates",
               "handling": "AR(1) equations whose dependent value or lag is 2020Q2 or 2020Q3 are excluded from estimation (equivalent to impulse dummies for 2020Q2-Q4); naive unaffected"})
issues.append({"issue": "2025Q3 advance estimate delayed by the Oct-Nov 2025 federal shutdown; 8th origin is 2025-12-31, and the h=1 target 2025Q4 first release is dated 2026-02-28",
               "handling": "used as provided; origin spacing is irregular for the last origin"})
issues.append({"issue": "Past values differ across vintages (e.g. 2023Q4: 3.3 in 2024-01-31 vintage, 3.4 in 2024-04-30, 3.2 in 2024-10-31) and annual revisions alter whole history",
               "handling": "each origin uses only its own vintage file; actuals are first releases"})

out = {"task_id": "F3-conf-02", "model_name": MODEL,
       "model_description": ("AR(1) with intercept estimated by OLS on the full origin vintage (1947Q2 to last published quarter), "
                             "excluding equations whose dependent value or lag is 2020Q2 or 2020Q3, re-estimated at each origin; "
                             "h=2 forecast iterated."),
       "n_origins": len(ORIGINS), "first_origin": ORIGINS[0], "last_origin": ORIGINS[-1],
       **{k: res[k] for k in ["rmse_naive_h1", "mae_naive_h1", "rmse_naive_h2", "mae_naive_h2",
                              "rmse_model_h1", "mae_model_h1", "rmse_model_h2", "mae_model_h2",
                              "rel_improvement_rmse_h1", "rel_improvement_rmse_h2"]},
       "mean_abs_revision_h1_targets": mean_rev,
       "feasibility": "complete", "feasibility_reason": None,
       "data_issues": issues, "n_input_rows": n_rows, "n_analysed": n_last, "random_seed": None,
       "diebold_mariano_hln": dm, "robustness": robust, "ar_fits": fit_info,
       "revisions_h1_targets": {t.strftime("%Y-%m-%d"): {"first": float(first[t]), "current": float(cur[t])} for t in t1}}
with open(os.path.join(OUT, "results.json"), "w") as fh:
    json.dump(out, fh, indent=2)
print(json.dumps(out, indent=2))
print(fc.to_string())