"""Real-time forecasts of US real GDP growth (F3-conf-02).
Reads only data/; writes out/results.json, out/forecasts.csv, out/exclusions.csv. Deterministic (no randomness)."""
import glob, json, os
import numpy as np, pandas as pd
from scipy import stats

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
D = os.path.join(ROOT, "data"); O = os.path.join(ROOT, "out")
MODEL = "ar1_1985_excovid"
EST_START = pd.Timestamp("1985-01-01")          # first quarter allowed in the estimation sample (as lag)
COVID = {pd.Timestamp("2020-04-01"), pd.Timestamp("2020-07-01")}  # pandemic outlier quarters

def read_series(path):
    df = pd.read_csv(path)
    df.columns = ["date", "value"]
    df["date"] = pd.to_datetime(df["date"])
    return df.sort_values("date").reset_index(drop=True)

def fit_ar1(df, ledger=None):
    """OLS AR(1) with intercept on y_t ~ y_{t-1}, t-1 >= EST_START, dropping pairs touching COVID quarters."""
    y = df.set_index("date")["value"]
    pairs = pd.DataFrame({"y": y, "ylag": y.shift(1), "lagdate": pd.Series(y.index, index=y.index).shift(1)}).dropna()
    n0 = len(df)
    p1 = pairs[pairs["lagdate"] >= EST_START]
    p2 = p1[~(p1.index.isin(list(COVID)) | p1["lagdate"].isin(list(COVID)))]
    if ledger is not None:
        ledger += [("1", "regression pairs (y_t, y_{t-1}) with y_{t-1} dated >= 1985-01-01 (pre-1985 quarters and the first, lag-less row dropped)", n0, len(p1), n0 - len(p1),
                    "post-Great-Moderation estimation sample; avoids high-volatility early postwar regime"),
                   ("2", "drop pairs in which y_t or y_{t-1} is 2020Q2 or 2020Q3", len(p1), len(p2), len(p1) - len(p2),
                    "pandemic collapse/rebound (about -28/+35 % SAAR) would dominate OLS and bias the AR coefficient")]
    X = np.column_stack([np.ones(len(p2)), p2["ylag"].values])
    beta, *_ = np.linalg.lstsq(X, p2["y"].values, rcond=None)
    return beta, len(p2)

def main():
    fr = pd.read_csv(os.path.join(D, "first_release.csv"), parse_dates=["quarter", "first_release_vintage_date"])
    frmap = dict(zip(fr["quarter"], fr["first_release_value"]))
    cur = read_series(os.path.join(D, "current_vintage.csv"))
    curmap = dict(zip(cur["date"], cur["value"]))
    files = sorted(glob.glob(os.path.join(D, "vintages", "*.csv")))
    n_input_rows = {"current_vintage.csv": int(len(cur)), "first_release.csv": int(len(fr))}
    rows, issues, ledger, coefs = [], [], [], {}
    for i, f in enumerate(files):
        origin = os.path.basename(f).split("_vintage_")[1].replace(".csv", "")
        v = read_series(f)
        n_input_rows[os.path.relpath(f, D)] = int(len(v))
        last_date, last_val = v["date"].iloc[-1], float(v["value"].iloc[-1])
        assert v["value"].notna().all() and not v["date"].duplicated().any()
        is_last = i == len(files) - 1
        beta, n_ar = fit_ar1(v, ledger if is_last else None)
        coefs[origin] = [round(float(b), 6) for b in beta]
        if is_last:
            n_analysed = n_ar
        fc_ar, prev = {}, last_val
        for h in (1, 2):
            prev = beta[0] + beta[1] * prev
            fc_ar[h] = prev
        for h in (1, 2):
            tgt = last_date + pd.DateOffset(months=3 * h)
            act = frmap.get(tgt, np.nan)
            if np.isnan(act):
                issues.append({"issue": f"no first release for {tgt.date()}", "handling": "target excluded"})
            for m, fc, ntr in (("naive", last_val, len(v)), (MODEL, fc_ar[h], n_ar)):
                rows.append(dict(origin_date=origin, target_date=str(tgt.date()), horizon=h, forecast=round(float(fc), 10),
                                 actual=act, model=m, vintage_date=origin, train_last_date=str(last_date.date()),
                                 train_last_value=last_val, n_train=int(ntr)))
    fc = pd.DataFrame(rows)
    fc.to_csv(os.path.join(O, "forecasts.csv"), index=False)
    pd.DataFrame(ledger, columns=["step", "rule", "n_before", "n_after", "n_removed", "rationale"]).to_csv(
        os.path.join(O, "exclusions.csv"), index=False)

    res = {"task_id": "F3-conf-02", "model_name": MODEL,
           "model_description": ("AR(1) with intercept estimated by OLS separately at each origin on that origin's vintage, "
                                 "sample of pairs (y_t, y_{t-1}) with y_{t-1} from 1985Q1 to the vintage's last quarter, excluding pairs "
                                 "involving 2020Q2 or 2020Q3; h=2 forecast by iterating the h=1 forecast."),
           "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):
            s = fc[(fc.model == m) & (fc.horizon == h)].dropna(subset=["actual"])
            e = s["forecast"] - s["actual"]
            res[f"rmse_{tag}_h{h}"] = float(np.sqrt(np.mean(e ** 2)))
            res[f"mae_{tag}_h{h}"] = float(np.mean(np.abs(e)))
    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}"]
    for h in (1, 2):  # Diebold-Mariano (squared loss) with Harvey-Leybourne-Newbold small-sample correction
        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).values
        n = len(d); dc = d - d.mean()
        lrv = np.mean(dc ** 2) + 2 * sum(np.sum(dc[k:] * dc[:-k]) / n for k in range(1, h))
        dm = d.mean() / np.sqrt(lrv / n)
        hln = dm * np.sqrt((n + 1 - 2 * h + h * (h - 1) / n) / n)
        res[f"dm_hln_stat_h{h}"] = float(hln)
        res[f"dm_hln_pvalue_two_sided_h{h}"] = float(2 * stats.t.sf(abs(hln), df=n - 1))
    t1 = fc[(fc.model == "naive") & (fc.horizon == 1)]["target_date"]
    revs = [abs(curmap[pd.Timestamp(t)] - frmap[pd.Timestamp(t)]) for t in t1]
    res["mean_abs_revision_h1_targets"] = float(np.mean(revs))
    issues.append({"issue": "2025Q3 advance estimate delayed by Oct-Nov 2025 federal shutdown; 8th origin is 2025-12-31, not 2025-10-31",
                   "handling": "used the provided 2025-12-31 vintage as the 8th origin; targets remain 2025Q4 (h=1) and 2026Q1 (h=2)"})
    issues.append({"issue": "pandemic outliers 2020Q2/2020Q3 in every vintage",
                   "handling": "excluded from AR(1) estimation pairs (see exclusions.csv); not relevant to the naive forecast"})
    res.update({"feasibility": "complete" if fc["actual"].notna().all() else "partial",
                "feasibility_reason": None if fc["actual"].notna().all() else "missing first releases for some targets",
                "data_issues": issues, "n_input_rows": n_input_rows, "n_analysed": int(n_analysed), "random_seed": None,
                "ar1_coefficients_by_origin": coefs})
    with open(os.path.join(O, "results.json"), "w") as fh:
        json.dump(res, fh, indent=2)
    print(json.dumps({k: v for k, v in res.items() if k.startswith(("rmse", "mae", "rel", "mean"))}, indent=2))

if __name__ == "__main__":
    main()