"""Okun's law, U.S. 1960Q1-2019Q4, OLS with Newey-West HAC SE.  Run from workspace root."""
import json, os
import numpy as np, pandas as pd
import statsmodels.api as sm
from scipy import stats
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

OUT = "out"
FIRST, LAST, BASE = pd.Period("1960Q1"), pd.Period("2019Q4"), pd.Period("1959Q4")
issues = []

un = pd.read_csv("data/UNRATE.csv", dtype={"observation_date": str})
gd = pd.read_csv("data/GDPC1.csv", dtype={"observation_date": str})
n_rows = {"UNRATE.csv": int(len(un)), "GDPC1.csv": int(len(gd))}
un["UNRATE"] = pd.to_numeric(un["UNRATE"], errors="coerce")
gd["GDPC1"] = pd.to_numeric(gd["GDPC1"], errors="coerce")
for d in (un, gd):
    d["date"] = pd.to_datetime(d["observation_date"])
    assert not d["date"].duplicated().any()
miss = un.loc[un["UNRATE"].isna(), "observation_date"].tolist()
if miss:
    issues.append({"issue": f"UNRATE missing value(s) at {', '.join(miss)}",
                   "handling": "outside the 1959Q4-2019Q4 window; not used. Quarterly averages require all 3 months (asserted for every quarter used)."})

un["q"] = un["date"].dt.to_period("Q")
uq = un.groupby("q")["UNRATE"].agg(["mean", "count"])
gd["q"] = gd["date"].dt.to_period("Q")
gq = gd.set_index("q")["GDPC1"]

n0 = len(gq)
panel = pd.DataFrame({"gdp": gq}).join(uq, how="left")
panel = panel.loc[(panel.index >= BASE) & (panel.index <= LAST)]
assert (panel["count"] == 3).all() and panel["gdp"].notna().all() and (panel["gdp"] > 0).all()
n1 = len(panel)
panel["du"] = panel["mean"].diff()
panel["g"] = 100 * np.log(panel["gdp"] / panel["gdp"].shift(1))
df = panel.loc[panel.index >= FIRST].copy()
n2 = len(df)
assert df[["du", "g"]].notna().all().all()
assert n2 == 240 and df.index[0] == FIRST and df.index[-1] == LAST

pd.DataFrame([
    dict(step=1, rule="keep quarters 1959Q4-2019Q4 (GDPC1 quarters with complete 3-month UNRATE average)",
         n_before=n0, n_after=n1, n_removed=n0 - n1, rationale="window fixed by TASK.md; 1959Q4 kept only as difference base"),
    dict(step=2, rule="drop 1959Q4 (no lagged level for first difference)", n_before=n1, n_after=n2,
         n_removed=n1 - n2, rationale="first differences need t-1 levels; regression sample 1960Q1-2019Q4"),
]).to_csv(os.path.join(OUT, "exclusions.csv"), index=False)

y, X = df["du"].values, sm.add_constant(df["g"].values)
T = len(y)
ols = sm.OLS(y, X).fit()
# Bandwidth: Newey-West (1994) plug-in-free rule floor(4*(T/100)^(2/9)) -> 4 for T=240
L = int(np.floor(4 * (T / 100) ** (2 / 9)))
def hac(lags):
    # classic Newey-West: Bartlett weights 1 - j/(L+1), no small-sample df correction
    return sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": lags, "use_correction": False})
h = hac(L)
# manual cross-check of the HAC SE
e = ols.resid; Xe = X * e[:, None]; S = Xe.T @ Xe
for j in range(1, L + 1):
    G = Xe[j:].T @ Xe[:-j]; S += (1 - j / (L + 1)) * (G + G.T)
XXi = np.linalg.inv(X.T @ X); se_manual = np.sqrt((XXi @ S @ XXi)[1, 1])
beta, se_h = float(ols.params[1]), float(h.bse[1])
assert abs(se_manual - se_h) < 1e-10
z = stats.norm.ppf(0.975)
lo, hi = beta - z * se_h, beta + z * se_h
p_hac = float(2 * stats.norm.sf(abs(beta / se_h)))
resid_ac = [float(pd.Series(e).autocorr(k)) for k in (1, 2, 3, 4)]

res = {
    "task_id": "F2-conf-04",
    "beta_hat": beta, "alpha_hat": float(ols.params[0]), "r_squared": float(ols.rsquared),
    "se_iid": float(ols.bse[1]), "se_hac": se_h, "hac_lags": L,
    "beta_ci_low": float(lo), "beta_ci_high": float(hi),
    "okun_negative_5pct": bool(beta < 0 and p_hac < 0.05),
    "mean_growth": float(df["g"].mean()),
    "first_quarter": str(df.index[0]), "last_quarter": str(df.index[-1]),
    "n_analysed": int(T), "feasibility": "complete", "feasibility_reason": None,
    "data_issues": issues, "n_input_rows": n_rows, "random_seed": None,
    # supplementary
    "t_hac": float(beta / se_h), "p_value_hac": p_hac,
    "baseline_ci_iid_low": float(beta - z * ols.bse[1]), "baseline_ci_iid_high": float(beta + z * ols.bse[1]),
    "resid_autocorr_lag1_4": resid_ac,
    "sensitivity_se_hac_by_lag": {str(k): float(hac(k).bse[1]) for k in (0, 2, 4, 5, 8, 12)},
    "mean_du": float(df["du"].mean()),
}
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(res, f, indent=2)

fig, ax = plt.subplots(figsize=(6.5, 4.5))
ax.scatter(df["g"], df["du"], s=12, alpha=0.6)
gg = np.linspace(df["g"].min(), df["g"].max(), 50)
ax.plot(gg, res["alpha_hat"] + beta * gg, color="C3",
        label=f"OLS: Δu = {res['alpha_hat']:.3f} + ({beta:.3f}) g")
ax.axhline(0, lw=0.5, color="k"); ax.axvline(0, lw=0.5, color="k")
ax.set_xlabel("Real GDP growth g_t (% q/q, not annualised)")
ax.set_ylabel("Δ unemployment rate Δu_t (pp)")
ax.set_title("Okun's law, U.S. 1960Q1–2019Q4 (n=240)")
ax.legend(); fig.tight_layout(); fig.savefig(os.path.join(OUT, "figure1.png"), dpi=120)
print(json.dumps({k: res[k] for k in ("beta_hat", "se_iid", "se_hac", "hac_lags", "beta_ci_low", "beta_ci_high", "okun_negative_5pct")}, indent=1))