"""F3-conf-02: real-time forecasts of US real GDP growth (A191RL1Q225SBEA).
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
"""
import json, glob, os
import numpy as np
import pandas as pd

DATA = "data"; OUT = "out"
MODEL = "ar1_1985_excovid"
EST_START = pd.Timestamp("1985-01-01")
COVID = (pd.Timestamp("2020-01-01"), pd.Timestamp("2021-01-01"))  # quarters excluded from regression pairs
issues = []
n_rows = {}

def read_series(path):
    df = pd.read_csv(path)
    n_rows[os.path.relpath(path, DATA)] = len(df)
    df.columns = ["date", "value"]
    df["date"] = pd.to_datetime(df["date"])
    df["value"] = pd.to_numeric(df["value"], errors="coerce")
    nbad = int(df["value"].isna().sum())
    if nbad:
        issues.append({"issue": f"{nbad} non-numeric/missing values in {path}", "handling": "dropped"})
        df = df.dropna()
    if df["date"].duplicated().any():
        issues.append({"issue": f"duplicate dates in {path}", "handling": "kept last"})
        df = df.drop_duplicates("date", keep="last")
    df = df.sort_values("date").set_index("date")["value"]
    # check regular quarterly spacing
    exp = pd.date_range(df.index[0], df.index[-1], freq="QS")
    if len(exp) != len(df):
        issues.append({"issue": f"gaps in quarterly index of {path}", "handling": "noted; AR pairs use consecutive quarters only"})
    return df

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"])
fr_map = fr.set_index("quarter")["first_release_value"]
cur = read_series(os.path.join(DATA, "current_vintage.csv"))

files = sorted(glob.glob(os.path.join(DATA, "vintages", "A191RL1Q225SBEA_vintage_*.csv")))
rows = []
last_n = None
for f in files:
    vd = os.path.basename(f).split("_vintage_")[1][:-4]
    y = read_series(f)
    last_date, last_val = y.index[-1], float(y.iloc[-1])
    targets = [last_date + pd.DateOffset(months=3 * h) for h in (1, 2)]
    # AR(1) by OLS: y_t = c + phi*y_{t-1} + e, t >= 1985Q1, consecutive quarters, excluding
    # pairs where y_t or y_{t-1} falls in 2020Q1..2021Q1 (pandemic outliers)
    ys = y[y.index >= EST_START - pd.DateOffset(months=3)]
    prev = ys.shift(1); prev_d = pd.Series(ys.index, index=ys.index).shift(1)
    ok = (ys.index >= EST_START) & prev.notna().values
    ok &= ((ys.index - pd.DatetimeIndex(prev_d.fillna(ys.index[0]))).days.values > 80) & \
          ((ys.index - pd.DatetimeIndex(prev_d.fillna(ys.index[0]))).days.values < 100)
    inC = lambda d: (d >= COVID[0]) & (d <= COVID[1])
    ok &= ~inC(ys.index) & ~inC(pd.DatetimeIndex(prev_d.fillna(ys.index[0])))
    Y = ys.values[ok]; X = np.column_stack([np.ones(ok.sum()), prev.values[ok]])
    c, phi = np.linalg.lstsq(X, Y, rcond=None)[0]
    f1 = c + phi * last_val; f2 = c + phi * f1
    for h, t, fm in ((1, targets[0], f1), (2, targets[1], f2)):
        act = float(fr_map[t]) if t in fr_map.index else np.nan
        base = dict(origin_date=vd, target_date=t.strftime("%Y-%m-%d"), horizon=h, actual=act,
                    vintage_date=vd, train_last_date=last_date.strftime("%Y-%m-%d"),
                    train_last_value=last_val)
        rows.append({**base, "forecast": last_val, "model": "naive", "n_train": len(y)})
        rows.append({**base, "forecast": round(float(fm), 6), "model": MODEL, "n_train": int(ok.sum())})
    last_n = len(y)
    print(vd, "c=%.3f phi=%.3f n=%d" % (c, phi, ok.sum()))

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 for some targets", "handling": "excluded from metrics"})

def metrics(m, h):
    d = fc[(fc.model == m) & (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, sample 1985Q1 to the last "
                             "published quarter, excluding regression pairs touching 2020Q1-2021Q1; h=2 forecast by iteration."),
       "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 = metrics(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)
rev = [abs(cur[t] - fr_map[t]) for t in t1]
res["mean_abs_revision_h1_targets"] = round(float(np.mean(rev)), 6)
res["feasibility"] = "complete"; res["feasibility_reason"] = None
issues.append({"issue": "2025Q3 advance estimate delayed by the Oct-Nov 2025 federal shutdown; 8th origin is 2025-12-31 and the 2025Q4 first release is dated 2026-02-28",
               "handling": "used origins as supplied; spacing between last two origins is 5 months, targets still defined relative to last quarter in the vintage"})
issues.append({"issue": "values of earlier quarters differ across vintages (revisions, incl. annual updates)",
               "handling": "each origin uses only its own vintage file; current vintage used only for the revision statistic"})
res["data_issues"] = issues
res["n_input_rows"] = n_rows
res["n_analysed"] = last_n
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))
print(fc.to_string())
print("h1 revisions:", dict(zip(t1.dt.strftime("%Y-%m-%d"), rev)))

# ---- Uncertainty: loss differentials, sign counts, simple DM-type t-stat (HLN small-sample correction) ----
from scipy import stats
unc = {}
for h in (1, 2):
    n_ = fc[(fc.model == "naive") & (fc.horizon == h)].reset_index(drop=True)
    m_ = fc[(fc.model == MODEL) & (fc.horizon == h)].reset_index(drop=True)
    d = (n_.forecast - n_.actual) ** 2 - (m_.forecast - m_.actual) ** 2
    T = len(d); dbar = d.mean()
    g0 = ((d - dbar) ** 2).mean()
    lr = g0 + (2 * ((d - dbar).iloc[1:].values * (d - dbar).iloc[:-1].values).sum() / T if h == 2 else 0)
    lr = lr if lr > 0 else g0
    dm = dbar / np.sqrt(lr / T)
    hln = dm * np.sqrt((T + 1 - 2 * h + h * (h - 1) / T) / T)
    unc[f"h{h}"] = {"mean_sq_loss_diff_naive_minus_model": round(float(dbar), 4),
                    "n_origins_model_better": int((d > 0).sum()), "T": T,
                    "hln_dm_stat": round(float(hln), 3),
                    "p_two_sided_t(T-1)": round(float(2 * stats.t.sf(abs(hln), T - 1)), 3),
                    "loss_diffs": [round(float(x), 3) for x in d]}
with open(os.path.join(OUT, "uncertainty.json"), "w") as fh:
    json.dump(unc, fh, indent=2)
print(json.dumps(unc, indent=2))