"""Okun's law, U.S. 1960Q1-2019Q4, OLS with Newey-West HAC SE. Run from workspace root."""
import json, math
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

issues = []
u = pd.read_csv("data/UNRATE.csv", parse_dates=["observation_date"])
g = pd.read_csv("data/GDPC1.csv", parse_dates=["observation_date"])
n_rows = {"UNRATE.csv": int(len(u)), "GDPC1.csv": int(len(g))}

for name, df, col in [("UNRATE.csv", u, "UNRATE"), ("GDPC1.csv", g, "GDPC1")]:
    if df["observation_date"].duplicated().any():
        raise SystemExit(f"duplicate dates in {name}")
    miss = df[df[col].isna()]
    if len(miss):
        dates = ", ".join(d.strftime("%Y-%m") for d in miss["observation_date"])
        issues.append({"issue": f"{name}: {len(miss)} missing value(s) at {dates}",
                       "handling": "Outside the 1959Q4-2019Q4 window needed; dropped, no effect on the analysis."})

u = u.dropna()
u["q"] = u["observation_date"].dt.to_period("Q")
cnt = u.groupby("q")["UNRATE"].count()
uq = u.groupby("q")["UNRATE"].mean()
g["q"] = g["observation_date"].dt.to_period("Q")
gq = g.dropna().set_index("q")["GDPC1"]

win = pd.period_range("1959Q4", "2019Q4", freq="Q")
assert (cnt.reindex(win) == 3).all(), "incomplete months in window"
assert gq.reindex(win).notna().all(), "missing GDP in window"
uw, gw = uq.reindex(win), gq.reindex(win)
du = uw.diff().iloc[1:]
gr = (100 * np.log(gw)).diff().iloc[1:]
assert len(du) == 240

X = sm.add_constant(gr.values)
y = du.values
T = len(y)
L = int(math.floor(4 * (T / 100) ** (2 / 9)))   # Newey-West (1994) plug-in rule of thumb -> 4
ols = sm.OLS(y, X).fit()
hac = sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": L, "use_correction": False})
beta, alpha = float(ols.params[1]), float(ols.params[0])
se_hac = float(hac.bse[1])
z = stats.norm.ppf(0.975)
lo, hi = beta - z * se_hac, beta + z * se_hac
p_hac = float(2 * stats.norm.sf(abs(beta / se_hac)))

# sensitivity (printed only)
sens = {}
for lag in [0, 1, 2, 4, 6, 8, 12, 16]:
    r = sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": lag, "use_correction": False})
    sens[lag] = float(r.bse[1])
resid_ac = [float(pd.Series(ols.resid).autocorr(k)) for k in range(1, 5)]

res = {
    "task_id": "F2-conf-04",
    "beta_hat": beta, "alpha_hat": alpha, "r_squared": float(ols.rsquared),
    "se_iid": float(ols.bse[1]), "se_hac": se_hac, "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(gr.mean()),
    "first_quarter": str(du.index[0]), "last_quarter": str(du.index[-1]),
    "n_analysed": int(T), "feasibility": "complete", "feasibility_reason": None,
    "data_issues": issues, "n_input_rows": n_rows, "random_seed": None,
}
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)

print(json.dumps(res, indent=2))
print("p_hac", p_hac, "t_hac", beta / se_hac)
print("HAC SE by lag:", sens)
print("resid autocorr 1-4:", resid_ac)
r_corr = sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": L, "use_correction": True})
print("HAC SE with df correction:", float(r_corr.bse[1]))

fig, ax = plt.subplots(figsize=(6, 4.5))
ax.scatter(gr.values, y, s=12, alpha=0.7)
xs = np.linspace(gr.min(), gr.max(), 50)
ax.plot(xs, alpha + beta * xs, color="red",
        label=f"OLS: Δu = {alpha:.3f} {beta:+.3f}·g")
ax.axhline(0, color="grey", lw=0.5); ax.axvline(0, color="grey", lw=0.5)
ax.set_xlabel("Real GDP growth g_t (%, q/q, not annualised)")
ax.set_ylabel("Δu_t (pp)")
ax.set_title("Okun's law, U.S. 1960Q1–2019Q4 (n=240)")
ax.legend(); fig.tight_layout(); fig.savefig("out/figure1.png", dpi=150)