"""Okun's law, US quarterly 1960Q1-2019Q4, OLS + Newey-West HAC. Run from workspace root."""
import json, os
import numpy as np, pandas as pd, statsmodels.api as sm
import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt

OUT = "out"; os.makedirs(OUT, exist_ok=True)
un = pd.read_csv("data/UNRATE.csv"); gd = pd.read_csv("data/GDPC1.csv")
n_rows = {"UNRATE.csv": int(len(un)), "GDPC1.csv": int(len(gd))}
issues = []

un["date"] = pd.to_datetime(un.observation_date); gd["date"] = pd.to_datetime(gd.observation_date)
for name, d, col in [("UNRATE", un, "UNRATE"), ("GDPC1", gd, "GDPC1")]:
    assert d.date.duplicated().sum() == 0, f"duplicate dates in {name}"
    miss = d.loc[d[col].isna(), "observation_date"].tolist()
    if miss:
        issues.append({"issue": f"{name} has missing value(s) at {miss}",
                       "handling": "Outside the 1959Q4-2019Q4 window needed for the regression; not used. "
                                   "Quarterly averages inside the window are required to have all 3 months (checked)."})

un["q"] = un.date.dt.to_period("Q"); gd["q"] = gd.date.dt.to_period("Q")
uq = un.groupby("q")["UNRATE"].agg(["mean", "count"])
lo, hi = pd.Period("1959Q4", "Q"), pd.Period("2019Q4", "Q")
uq = uq.loc[lo:hi]; assert (uq["count"] == 3).all() and len(uq) == 241
g = gd.set_index("q")["GDPC1"].loc[lo:hi]; assert len(g) == 241 and g.notna().all()

df = pd.DataFrame({"u": uq["mean"], "y": g})
df["du"] = df.u.diff(); df["g"] = 100 * np.log(df.y).diff()
df = df.loc[pd.Period("1960Q1", "Q"):hi].dropna()
assert len(df) == 240
T = len(df)

X = sm.add_constant(df["g"].values); yv = df["du"].values
ols = sm.OLS(yv, X).fit()
# Bandwidth: Newey-West (1994)/Stock-Watson rule of thumb floor(4*(T/100)^(2/9))
L = int(np.floor(4 * (T / 100) ** (2 / 9)))
hac = sm.OLS(yv, X).fit(cov_type="HAC", cov_kwds={"maxlags": L, "kernel": "bartlett", "use_correction": False})
b, a = float(ols.params[1]), float(ols.params[0])
se_hac = float(hac.bse[1]); z = 1.959963984540054
ci = (b - z * se_hac, b + z * se_hac)
p_hac = float(hac.pvalues[1])

# sensitivity of HAC SE to bandwidth (reported, not part of headline)
sens = {}
for l in [0, 1, 2, 4, 6, 8, 12, int(np.floor(0.75 * T ** (1 / 3)))]:
    r = sm.OLS(yv, X).fit(cov_type="HAC", cov_kwds={"maxlags": l, "use_correction": False})
    sens[str(l)] = float(r.bse[1])
resid_acf = [float(np.corrcoef(ols.resid[k:], ols.resid[:-k])[0, 1]) for k in range(1, 5)]

res = {
    "task_id": "F2-conf-04",
    "beta_hat": b, "alpha_hat": a, "r_squared": float(ols.rsquared),
    "se_iid": float(ols.bse[1]), "se_hac": se_hac, "hac_lags": L,
    "beta_ci_low": float(ci[0]), "beta_ci_high": float(ci[1]),
    "okun_negative_5pct": bool(b < 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,
}
with open(f"{OUT}/results.json", "w") as f: json.dump(res, f, indent=2)
with open(f"{OUT}/supplementary.json", "w") as f:
    json.dump({"p_value_hac_normal": p_hac, "t_hac": b / se_hac, "se_alpha_hac": float(hac.bse[0]),
               "hac_se_by_lags": sens, "resid_autocorr_lags1to4": resid_acf}, f, indent=2)

plt.figure(figsize=(6, 4.5))
plt.scatter(df["g"], df["du"], s=12, alpha=0.6)
xs = np.linspace(df["g"].min(), df["g"].max(), 50)
plt.plot(xs, a + b * xs, "r-", label=f"OLS: Δu = {a:.3f} {b:+.3f} g")
plt.xlabel("Real GDP growth g_t (% q/q, not annualised)"); plt.ylabel("Δ unemployment rate (pp)")
plt.title("Okun's law, US 1960Q1–2019Q4 (T=240)"); plt.legend(); plt.grid(alpha=.3); plt.tight_layout()
plt.savefig(f"{OUT}/figure1.png", dpi=150)
print(json.dumps(res, indent=2))