"""Real-time forecasts of US real GDP growth (F3-conf-02).
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"
EST_START = pd.Timestamp("1985-01-01")
COVID = [pd.Timestamp("2020-04-01"), pd.Timestamp("2020-07-01")]
MODEL = "ar1_1985_excovid"

n_rows, issues = {}, []

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")
    nmiss = int(df["value"].isna().sum())
    if nmiss:
        issues.append({"issue": f"{nmiss} 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").reset_index(drop=True)
    exp = pd.date_range(df["date"].iloc[0], df["date"].iloc[-1], freq="QS")
    if len(exp) != len(df):
        issues.append({"issue": f"quarterly gaps in {path}", "handling": "lag pairs require consecutive quarters"})
    return df

files = sorted(glob.glob(os.path.join(DATA, "vintages", "*.csv")))
vint = {os.path.basename(f).split("_vintage_")[1][:-4]: read_series(f) for f in files}
cur = read_series(os.path.join(DATA, "current_vintage.csv"))
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"])
frmap = dict(zip(fr["quarter"], fr["first_release_value"]))

def fit_ar1(df):
    s = df.set_index("date")["value"]
    prev = s.shift(1); prev_date = pd.Series(s.index, index=s.index).shift(1)
    d = pd.DataFrame({"y": s, "x": prev, "pd": prev_date})
    # consecutive quarter requirement
    d = d[d["pd"] == (d.index - pd.offsets.QuarterBegin(1, startingMonth=1))]
    d = d[d.index >= EST_START]
    d = d[~d.index.isin(COVID) & ~d["pd"].isin(COVID)]
    d = d.dropna()
    X = np.column_stack([np.ones(len(d)), d["x"].values])
    beta, *_ = np.linalg.lstsq(X, d["y"].values, rcond=None)
    return beta, len(d)

rows = []
for origin, df in sorted(vint.items()):
    last_date, last_val = df["date"].iloc[-1], float(df["value"].iloc[-1])
    (c, phi), n_reg = fit_ar1(df)
    f = last_val
    for h in (1, 2):
        tgt = last_date + pd.DateOffset(months=3 * h)
        f = c + phi * f
        actual = frmap.get(tgt, np.nan)
        base = dict(origin_date=origin, target_date=tgt.strftime("%Y-%m-%d"), horizon=h,
                    actual=actual, vintage_date=origin,
                    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(df)})
        rows.append({**base, "forecast": round(float(f), 6), "model": MODEL, "n_train": n_reg,
                     "ar_const": c, "ar_phi": phi})
fc = pd.DataFrame(rows)
coefs = fc[fc.model == MODEL].groupby("origin_date")[["ar_const", "ar_phi"]].first()
fc = fc.drop(columns=["ar_const", "ar_phi"])
cols = ["origin_date", "target_date", "horizon", "forecast", "actual", "model",
        "vintage_date", "train_last_date", "train_last_value", "n_train"]
fc = fc[cols]
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"})

res = {"task_id": "F3-conf-02", "model_name": MODEL,
       "model_description": ("AR(1) with intercept estimated by OLS at each origin on the origin's vintage, "
                             "sample 1985Q1 to last published quarter, excluding regression pairs involving "
                             "2020Q2-2020Q3 (COVID outliers); h=2 forecast iterated."),
       "n_origins": len(vint), "first_origin": min(vint), "last_origin": max(vint)}
for m, tag in (("naive", "naive"), (MODEL, "model")):
    for h in (1, 2):
        e = fc[(fc.model == m) & (fc.horizon == h)].eval("forecast-actual").dropna()
        res[f"rmse_{tag}_h{h}"] = round(float(np.sqrt((e**2).mean())), 4)
        res[f"mae_{tag}_h{h}"] = round(float(e.abs().mean()), 4)
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}"], 4)
curmap = dict(zip(cur["date"], cur["value"]))
t1 = pd.to_datetime(fc[(fc.model == "naive") & (fc.horizon == 1)]["target_date"])
revs = [abs(curmap[t] - frmap[t]) for t in t1]
res["mean_abs_revision_h1_targets"] = round(float(np.mean(revs)), 4)
issues.append({"issue": "No 2025-10-31 vintage: 2025Q3 advance estimate delayed by federal shutdown; 8th origin is 2025-12-31 (gap of 5 months between last two origins)",
               "handling": "used origins as provided; horizons defined relative to last quarter in each file"})
issues.append({"issue": "Extreme COVID observations 2020Q2 (~-28%) and 2020Q3 (~+35%) dominate OLS AR estimates",
               "handling": "regression pairs with dependent or lagged value in 2020Q2-2020Q3 excluded from AR estimation (pre-specified, not tuned on outcomes)"})
res.update({"feasibility": "complete", "feasibility_reason": None, "data_issues": issues,
            "n_input_rows": n_rows,
            "n_analysed": int(fc[(fc.model == MODEL) & (fc.origin_date == max(vint))]["n_train"].iloc[0]),
            "random_seed": None})
with open(os.path.join(OUT, "results.json"), "w") as fh:
    json.dump(res, fh, indent=2)
# diagnostics (not a required deliverable)
coefs.to_csv(os.path.join(OUT, "ar_coefficients.csv"))
print(json.dumps(res, indent=2)); print(fc.to_string()); print(coefs)
print("h1 target revisions:", dict(zip([t.strftime('%Y-%m-%d') for t in t1], [round(r,2) for r in revs])))

# ---- supplementary: uncertainty and sensitivity (written to out/supplementary.json) ----
from scipy import stats
sup = {}
for h in (1, 2):
    a = fc[(fc.model == "naive") & (fc.horizon == h)].reset_index(drop=True)
    b = fc[(fc.model == MODEL) & (fc.horizon == h)].reset_index(drop=True)
    d = (a.forecast - a.actual) ** 2 - (b.forecast - b.actual) ** 2
    n = len(d); dbar = d.mean()
    # HAC variance with h-1 lags (rectangular), Harvey-Leybourne-Newbold correction
    g = [((d - dbar) * (d - dbar).shift(k)).sum() / n for k in range(h)]
    v = (g[0] + 2 * sum(g[1:])) / n
    dm = dbar / np.sqrt(v) if v > 0 else np.nan
    k = np.sqrt((n + 1 - 2 * h + h * (h - 1) / n) / n)
    stat = dm * k
    sup[f"h{h}"] = {"mean_sq_loss_diff_naive_minus_model": round(float(dbar), 4),
                    "dm_hln_stat": round(float(stat), 4),
                    "dm_pvalue_two_sided_t(n-1)": round(float(2 * stats.t.sf(abs(stat), n - 1)), 4),
                    "n_origins_model_abs_error_smaller": int(((b.forecast - b.actual).abs() < (a.forecast - a.actual).abs()).sum()),
                    "n": n}
# sensitivity: AR(1) without excluding COVID quarters
COVID_SAVE = COVID[:]; COVID.clear()
sens = {1: [], 2: []}
for origin, df in sorted(vint.items()):
    (c, phi), _ = fit_ar1(df); f = float(df["value"].iloc[-1]); ld = df["date"].iloc[-1]
    for h in (1, 2):
        f = c + phi * f
        sens[h].append(f - frmap[ld + pd.DateOffset(months=3 * h)])
COVID.extend(COVID_SAVE)
for h in (1, 2):
    e = np.array(sens[h])
    sup[f"sensitivity_ar1_incl_covid_h{h}"] = {"rmse": round(float(np.sqrt((e**2).mean())), 4),
                                               "mae": round(float(np.abs(e).mean()), 4)}
sup["h1_target_abs_revisions"] = dict(zip([t.strftime('%Y-%m-%d') for t in t1], [round(float(r), 2) for r in revs]))
with open(os.path.join(OUT, "supplementary.json"), "w") as fh:
    json.dump(sup, fh, indent=2)
print(json.dumps(sup, indent=2))