"""F3-dev-01: rolling-origin (expanding window) forecasts of UNRATE, naive vs AR(2) on first differences.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
"""
import json, os
import numpy as np
import pandas as pd

DATA = "data/UNRATE.csv"
OUT = "out"
MODEL = "AR2_diff"
P = 2
HORIZONS = [1, 3]
ORIGINS = pd.date_range("2023-07-01", "2026-06-01", freq="MS")

raw = pd.read_csv(DATA, dtype={"observation_date": str, "UNRATE": str}, keep_default_na=False)
n_rows = len(raw)
df = pd.DataFrame({"date": pd.to_datetime(raw["observation_date"], format="%Y-%m-%d"),
                   "y": pd.to_numeric(raw["UNRATE"].replace("", np.nan), errors="coerce")})
issues = []
miss = df.loc[df.y.isna(), "date"].dt.strftime("%Y-%m-%d").tolist()
assert df.date.is_unique and df.date.is_monotonic_increasing
full = pd.date_range(df.date.min(), df.date.max(), freq="MS")
assert len(full) == len(df), "calendar gaps present"
s = df.set_index("date")["y"]
if miss:
    issues.append({"issue": f"Missing value(s) (empty field) at {', '.join(miss)}",
                   "handling": "Excluded from estimation (any regression row whose dependent or lagged difference touches the missing month is dropped); "
                               "naive uses latest earlier observation; when missing is the latest month at an origin the AR model iterates forward from the last observed month "
                               "(extra steps); when it is interior to the window it is linearly interpolated from neighbours dated <= origin only to build lag inputs; "
                               "(origin,horizon) pairs whose target is missing are kept in the ledger with empty actual and excluded from error statistics."})
covid = s["2020-03-01":"2020-06-01"]
issues.append({"issue": "COVID-19 outlier: unemployment jumped from %.1f (2020-03) to %.1f (2020-04), extreme first differences in training sample" % (s["2020-03-01"], s["2020-04-01"]),
               "handling": "Retained in the primary specification (pre-specified, no outlier treatment); a sensitivity variant excluding 2020-03..2020-12 differences from estimation is reported in report.md only."})


def fit_ar_diff(y, p=P, exclude=None):
    """OLS AR(p) with intercept on dy, using rows with no missing values. y: level Series up to origin."""
    d = y.diff()
    X = pd.concat({f"l{k}": d.shift(k) for k in range(1, p + 1)}, axis=1)
    data = pd.concat([d.rename("dy"), X], axis=1)
    if exclude is not None:
        data = data[~((data.index >= exclude[0]) & (data.index <= exclude[1]))]
    data = data.dropna()
    A = np.column_stack([np.ones(len(data)), data[[f"l{k}" for k in range(1, p + 1)]].values])
    beta, *_ = np.linalg.lstsq(A, data["dy"].values, rcond=None)
    return beta, len(data)


def forecast_ar_diff(y, t, horizons, exclude=None):
    y = y[:t]
    beta, nobs = fit_ar_diff(y, exclude=exclude)
    last_obs = y.last_valid_index()
    yy = y[:last_obs].interpolate(method="linear", limit_area="inside")  # interior gaps only, data <= t
    lev = list(yy.values)
    dif = list(np.diff(yy.values))
    steps_extra = (t.year - last_obs.year) * 12 + (t.month - last_obs.month)
    out = {}
    H = max(horizons) + steps_extra
    for step in range(1, H + 1):
        dnew = beta[0] + sum(beta[k] * dif[-k] for k in range(1, P + 1))
        dif.append(dnew)
        lev.append(lev[-1] + dnew)
        hh = step - steps_extra
        if hh in horizons:
            out[hh] = lev[-1]
    return out, nobs


def run(exclude=None):
    rows, nobs_last = [], None
    for t in ORIGINS:
        yt = s[:t]
        naive = yt.dropna().iloc[-1]
        fc, nobs = forecast_ar_diff(s, t, HORIZONS, exclude)
        nobs_last = nobs
        for h in HORIZONS:
            tgt = t + pd.DateOffset(months=h)
            act = s.get(tgt, np.nan)
            for m, f in (("naive", naive), (MODEL, fc[h])):
                rows.append({"origin_date": t.strftime("%Y-%m-%d"), "target_date": tgt.strftime("%Y-%m-%d"),
                             "horizon": h, "forecast": round(float(f), 6),
                             "actual": act, "model": m})
    return pd.DataFrame(rows), nobs_last


def stats(led):
    r = {}
    for h in HORIZONS:
        for m, tag in (("naive", "naive"), (MODEL, "model")):
            sub = led[(led.horizon == h) & (led.model == m) & led.actual.notna()]
            e = sub.actual - sub.forecast
            r[f"rmse_{tag}_h{h}"] = float(np.sqrt(np.mean(e ** 2)))
            r[f"mae_{tag}_h{h}"] = float(np.mean(np.abs(e)))
            r[f"n_evaluated_h{h}"] = int(len(sub))
        r[f"rel_improvement_rmse_h{h}"] = 1 - r[f"rmse_model_h{h}"] / r[f"rmse_naive_h{h}"]
    return r


if __name__ == "__main__":
    led, nobs_last = run()
    st = stats(led)
    os.makedirs(OUT, exist_ok=True)
    led.to_csv(os.path.join(OUT, "forecasts.csv"), index=False, float_format="%.6g")
    res = {
        "task_id": "F3-dev-01",
        "model_name": MODEL,
        "model_description": "AR(2) with intercept on monthly first differences of UNRATE, estimated by OLS on an expanding window starting with the first available difference (1948-04 as dependent observation) up to the origin, iterated recursively and cumulated to levels for h-step forecasts.",
        "n_origins": len(ORIGINS),
        "first_origin": ORIGINS[0].strftime("%Y-%m-%d"),
        "last_origin": ORIGINS[-1].strftime("%Y-%m-%d"),
    }
    for k in ["n_evaluated_h1", "n_evaluated_h3", "rmse_naive_h1", "mae_naive_h1", "rmse_naive_h3", "mae_naive_h3",
              "rmse_model_h1", "mae_model_h1", "rmse_model_h3", "mae_model_h3",
              "rel_improvement_rmse_h1", "rel_improvement_rmse_h3"]:
        res[k] = st[k]
    res.update({"feasibility": "complete", "feasibility_reason": None, "data_issues": issues,
                "n_input_rows": {"UNRATE.csv": n_rows}, "n_analysed": int(nobs_last), "random_seed": None})
    with open(os.path.join(OUT, "results.json"), "w") as f:
        json.dump(res, f, indent=2)
    print(json.dumps(res, indent=2))

    # ---------------- supplementary (report only) ----------------
    from scipy import stats as sps
    sup = {}
    beta_last, _ = fit_ar_diff(s[:ORIGINS[-1]])
    beta_first, _ = fit_ar_diff(s[:ORIGINS[0]])
    sup["coef_first_origin"] = [float(b) for b in beta_first]
    sup["coef_last_origin"] = [float(b) for b in beta_last]

    def dm(led, h):
        a = led[(led.horizon == h) & (led.model == "naive") & led.actual.notna()].reset_index(drop=True)
        b = led[(led.horizon == h) & (led.model == MODEL) & led.actual.notna()].reset_index(drop=True)
        d = (a.actual - b.forecast) ** 2 - (a.actual - a.forecast) ** 2  # model loss - naive loss
        n = len(d); dbar = d.mean(); dc = d - dbar
        lrv = np.sum(dc ** 2) / n
        for k in range(1, h):  # rectangular HAC with h-1 lags
            lrv += 2 * np.sum(dc[k:].values * dc[:-k].values) / n
        stat = dbar / np.sqrt(lrv / n)
        hln = np.sqrt((n + 1 - 2 * h + h * (h - 1) / n) / n)  # Harvey-Leybourne-Newbold correction
        stat *= hln
        p = 2 * sps.t.sf(abs(stat), df=n - 1)
        return {"mean_loss_diff_model_minus_naive": float(dbar), "dm_hln_stat": float(stat), "p_value_two_sided": float(p), "n": int(n)}
    sup["dm_h1"] = dm(led, 1); sup["dm_h3"] = dm(led, 3)

    # bootstrap-free descriptive: share of origins where model abs error < naive abs error
    for h in HORIZONS:
        a = led[(led.horizon == h) & (led.model == "naive") & led.actual.notna()].reset_index(drop=True)
        b = led[(led.horizon == h) & (led.model == MODEL) & led.actual.notna()].reset_index(drop=True)
        sup[f"share_model_better_abs_h{h}"] = float(np.mean(np.abs(b.actual - b.forecast) < np.abs(a.actual - a.forecast)))
        sup[f"share_ties_abs_h{h}"] = float(np.mean(np.isclose(np.abs(b.actual - b.forecast), np.abs(a.actual - a.forecast))))

    led2, _ = run(exclude=(pd.Timestamp("2020-03-01"), pd.Timestamp("2020-12-01")))
    st2 = stats(led2)
    sup["sensitivity_exclude_covid_2020"] = {k: st2[k] for k in st2 if k.startswith(("rmse_model", "mae_model", "rel_"))}
    sup["sensitivity_exclude_covid_coef_last_origin"] = [float(b) for b in fit_ar_diff(s[:ORIGINS[-1]], exclude=(pd.Timestamp("2020-03-01"), pd.Timestamp("2020-12-01")))[0]]
    with open(os.path.join(OUT, "supplementary.json"), "w") as f:
        json.dump(sup, f, indent=2)
    print(json.dumps(sup, indent=2))