"""F3-conf-02: real-time forecasts of US real GDP growth (A191RL1Q225SBEA).
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Reads only data/; writes out/results.json, out/forecasts.csv, out/exclusions.csv.
No randomness is used (random_seed = null)."""
import json, os, glob
import numpy as np
import pandas as pd

DATA = "data"; OUT = "out"
MODEL = "ar1"
EST_START = pd.Timestamp("1985-01-01")
COVID_DEP = pd.to_datetime(["2020-04-01", "2020-07-01", "2020-10-01"])  # dependent quarters dropped
DESC = ("AR(1) with intercept, y_t = c + phi*y_{t-1} + e_t, estimated by OLS on the origin vintage "
        "from 1985Q1 (first lag) to the vintage's last quarter, dropping observations whose dependent "
        "quarter is 2020Q2-2020Q4 (COVID outliers as dependent or lag); h=2 by iterating the 1-step forecast.")

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

def add_q(ts, h):
    return ts + pd.DateOffset(months=3 * h)

def ar1_fit(s):
    """s: vintage DataFrame. Returns (c, phi, n_obs, exclusion steps)."""
    steps = []
    n0 = len(s)
    w = s[s.date >= EST_START].reset_index(drop=True)
    steps.append(("window", "keep quarters >= 1985-01-01", n0, len(w), "Great-Moderation-era estimation window"))
    y = w.value.values[1:]; x = w.value.values[:-1]; dts = w.date.values[1:]
    steps.append(("lag", "first quarter of window used only as lag", len(w), len(y), "AR(1) loses one observation"))
    keep = ~np.isin(dts, COVID_DEP.values)
    steps.append(("covid", "drop dependent quarters 2020Q2-2020Q4", len(y), int(keep.sum()),
                  "COVID outliers (-28/+35 % SAAR) would dominate OLS"))
    y, x = y[keep], x[keep]
    X = np.column_stack([np.ones_like(x), x])
    beta, *_ = np.linalg.lstsq(X, y, rcond=None)
    return float(beta[0]), float(beta[1]), int(len(y)), steps

def main():
    vfiles = sorted(glob.glob(os.path.join(DATA, "vintages", "A191RL1Q225SBEA_vintage_*.csv")))
    fr = pd.read_csv(os.path.join(DATA, "first_release.csv"), parse_dates=["quarter", "first_release_vintage_date"])
    cur = read_series(os.path.join(DATA, "current_vintage.csv"))  # used ONLY for revision statistic
    n_input_rows = {"current_vintage.csv": len(cur), "first_release.csv": len(fr)}
    issues = []
    rows = []; excl = None; h1_targets = []; coefs = {}
    for f in vfiles:
        vdate = pd.Timestamp(os.path.basename(f).split("_vintage_")[1][:-4])
        s = read_series(f)
        n_input_rows[os.path.relpath(f, DATA)] = len(s)
        assert s.value.notna().all() and not s.date.duplicated().any()
        last_d, last_v = s.date.iloc[-1], float(s.value.iloc[-1])
        # sanity: last vintage value equals first release of that quarter
        frl = fr.loc[fr.quarter == last_d, "first_release_value"]
        if len(frl) and abs(float(frl.iloc[0]) - last_v) > 1e-9:
            issues.append({"issue": f"vintage {vdate.date()} last value differs from first release",
                           "handling": "vintage value used for forecasting, first release as actual"})
        c, phi, n_ar, steps = ar1_fit(s)
        excl = steps  # keep last origin's ledger
        coefs[vdate.strftime('%Y-%m-%d')] = {'c': round(c, 6), 'phi': round(phi, 6)}
        f1 = c + phi * last_v; f2 = c + phi * f1
        for h, fm in [(1, f1), (2, f2)]:
            tgt = add_q(last_d, h)
            a = fr.loc[fr.quarter == tgt]
            if len(a) == 0:
                actual = np.nan
                issues.append({"issue": f"no first release for {tgt.date()}", "handling": "excluded from errors"})
            else:
                actual = float(a.first_release_value.iloc[0])
                assert a.first_release_vintage_date.iloc[0] > vdate  # actual published after origin
            if h == 1: h1_targets.append(tgt)
            base = dict(origin_date=vdate.strftime("%Y-%m-%d"), target_date=tgt.strftime("%Y-%m-%d"), horizon=h,
                        actual=actual, vintage_date=vdate.strftime("%Y-%m-%d"),
                        train_last_date=last_d.strftime("%Y-%m-%d"), train_last_value=last_v)
            rows.append({**base, "forecast": last_v, "model": "naive", "n_train": len(s)})
            rows.append({**base, "forecast": round(fm, 6), "model": MODEL, "n_train": n_ar})
    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)

    def err(model, h):
        d = fc[(fc.model == model) & (fc.horizon == h)].dropna(subset=["actual"])
        e = d.forecast - d.actual
        return float(np.sqrt(np.mean(e ** 2))), float(np.mean(np.abs(e)))
    res = {"task_id": "F3-conf-02", "model_name": MODEL, "model_description": DESC,
           "n_origins": len(vfiles),
           "first_origin": os.path.basename(vfiles[0]).split("_vintage_")[1][:-4],
           "last_origin": os.path.basename(vfiles[-1]).split("_vintage_")[1][:-4]}
    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)
    revs = []
    for t in h1_targets:
        cv = float(cur.loc[cur.date == t, "value"].iloc[0])
        fv = float(fr.loc[fr.quarter == t, "first_release_value"].iloc[0])
        revs.append(abs(cv - fv))
    res["mean_abs_revision_h1_targets"] = round(float(np.mean(revs)), 6)
    res["feasibility"] = "complete"; res["feasibility_reason"] = None
    issues.append({"issue": "COVID quarters 2020Q2-2020Q3 are extreme outliers in every vintage",
                   "handling": "AR(1) estimation drops observations with dependent quarter 2020Q2-2020Q4; naive unaffected"})
    issues.append({"issue": "8th origin is 2025-12-31 (2025Q3 advance estimate delayed by the 2025 shutdown); origins are not equally spaced",
                   "handling": "used as given in TASK.md; horizons defined relative to last quarter in each vintage"})
    res["data_issues"] = issues
    res["n_input_rows"] = n_input_rows
    res["n_analysed"] = int(fc[(fc.origin_date == res["last_origin"]) & (fc.model == MODEL)].n_train.iloc[0])
    res["random_seed"] = None
    res["ar1_coefficients_by_origin"] = coefs
    for h in (1, 2):
        a = fc[(fc.horizon == h)].pivot(index="origin_date", columns="model", values="forecast")
        act = fc[(fc.horizon == h) & (fc.model == "naive")].set_index("origin_date").actual
        res[f"n_origins_model_abs_error_lower_h{h}"] = int(((a[MODEL] - act).abs() < (a["naive"] - act).abs()).sum())
    with open(os.path.join(OUT, "results.json"), "w") as fh:
        json.dump(res, fh, indent=2)
    ex = pd.DataFrame(excl, columns=["step", "rule", "n_before", "n_after", "rationale"])
    ex["n_removed"] = ex.n_before - ex.n_after
    ex[["step", "rule", "n_before", "n_after", "n_removed", "rationale"]].to_csv(
        os.path.join(OUT, "exclusions.csv"), index=False)
    print(json.dumps({k: v for k, v in res.items() if k.startswith(("rmse", "mae", "rel", "mean"))}, indent=1))

if __name__ == "__main__":
    main()