"""Okun's law, US 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
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy import stats

D, O = "data", "out"
os.makedirs(O, exist_ok=True)
FIRST, LAST = pd.Period("1960Q1", "Q"), pd.Period("2019Q4", "Q")

un = pd.read_csv(f"{D}/UNRATE.csv")
gd = pd.read_csv(f"{D}/GDPC1.csv")
n_rows = {"UNRATE.csv": int(len(un)), "GDPC1.csv": int(len(gd))}
issues = []
excl = []

# --- audit ---
for name, df, col in [("UNRATE", un, "UNRATE"), ("GDPC1", gd, "GDPC1")]:
    df["date"] = pd.to_datetime(df["observation_date"])
    df[col] = pd.to_numeric(df[col], errors="coerce")
    assert not df["date"].duplicated().any(), f"duplicate dates in {name}"
    miss = df.loc[df[col].isna(), "observation_date"].tolist()
    if miss:
        inside = [m for m in miss if FIRST - 1 <= pd.Period(m, "M").asfreq("Q") <= LAST]
        issues.append({"issue": f"{name}: missing (empty) value(s) at {miss}",
                       "handling": "outside the 1959Q4-2019Q4 levels window; ignored, no effect on estimates" if not inside
                       else "INSIDE window - affected quarters cannot be formed"})

# --- quarterly unemployment: mean of 3 months, require all 3 non-missing ---
un["q"] = un["date"].dt.to_period("Q")
uq = un.groupby("q")["UNRATE"].agg(["mean", "count"])  # count = non-missing months
gd["q"] = gd["date"].dt.to_period("Q")
gq = gd.set_index("q")["GDPC1"]

qs = pd.period_range(FIRST - 1, LAST, freq="Q")  # 1959Q4..2019Q4 levels
lev = pd.DataFrame(index=qs)
lev["u"] = uq["mean"].reindex(qs)
lev["nmon"] = uq["count"].reindex(qs)
lev["gdp"] = gq.reindex(qs)
assert (lev["nmon"] == 3).all(), "incomplete months in window"
assert lev["gdp"].notna().all() and (lev["gdp"] > 0).all()
lev["du"] = lev["u"].diff()
lev["g"] = 100 * np.log(lev["gdp"] / lev["gdp"].shift(1))

q_all = pd.period_range(min(uq.index.min(), gq.index.min()), max(uq.index.max(), gq.index.max()), freq="Q")
excl.append(("1", "restrict to 1959Q4-2019Q4 (calendar quarters spanned by either file)", len(q_all), len(qs),
             len(q_all) - len(qs), "task window plus one lag quarter; window not extended"))
dat = lev.loc[FIRST:LAST, ["du", "g"]]
excl.append(("2", "drop 1959Q4 (lag-only level quarter)", len(qs), len(dat), len(qs) - len(dat),
             "1959Q4 only supplies levels for 1960Q1 differences"))
assert dat.notna().all().all() and len(dat) == 240
pd.DataFrame(excl, columns=["step", "rule", "n_before", "n_after", "n_removed", "rationale"]).to_csv(f"{O}/exclusions.csv", index=False)

y, x = dat["du"].values, dat["g"].values
X = sm.add_constant(x)
T = len(y)

# --- baseline: OLS with iid SE ---
ols = sm.OLS(y, X).fit()
# --- HAC: Newey-West (1994) rule-of-thumb bandwidth floor(4*(T/100)^(2/9)) ---
L = int(np.floor(4 * (T / 100) ** (2 / 9)))
hac = sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": L, "kernel": "bartlett", "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 over bandwidths (reported, not used for decision)
sens = {}
for l in [0, 2, 4, 6, 8, 12]:
    r = sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": l, "kernel": "bartlett", "use_correction": False})
    sens[str(l)] = float(r.bse[1])
resid_ac1 = float(np.corrcoef(ols.resid[1:], ols.resid[:-1])[0, 1])

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(np.mean(x)),
    "first_quarter": str(dat.index[0]), "last_quarter": str(dat.index[-1]),
    "n_analysed": int(T),
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": issues, "n_input_rows": n_rows, "random_seed": None,
    "baseline_ci_iid_low": float(beta - z * ols.bse[1]), "baseline_ci_iid_high": float(beta + z * ols.bse[1]),
    "t_hac": float(beta / se_hac), "p_value_hac": p_hac,
    "resid_lag1_autocorr": resid_ac1,
    "se_hac_sensitivity_by_lags": sens,
    "mean_du": float(np.mean(y)),
}
with open(f"{O}/results.json", "w") as f:
    json.dump(res, f, indent=2)

fig, ax = plt.subplots(figsize=(6.5, 4.5))
ax.scatter(x, y, s=12, alpha=0.6)
xx = np.linspace(x.min(), x.max(), 100)
ax.plot(xx, alpha + beta * xx, "r-", label=f"OLS: du = {alpha:.3f} + ({beta:.3f}) g")
ax.axhline(0, c="grey", lw=0.5); ax.axvline(0, c="grey", lw=0.5)
ax.set_xlabel("g_t: real GDP growth, % q/q (not annualised)")
ax.set_ylabel("change in quarterly unemployment rate, pp")
ax.set_title("Okun's law, US 1960Q1-2019Q4 (n=240)"); ax.legend()
fig.tight_layout(); fig.savefig(f"{O}/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"]}))