"""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

FIRST, LAST = pd.Period("1960Q1", "Q"), pd.Period("2019Q4", "Q")
issues, excl = [], []

u_raw = pd.read_csv("data/UNRATE.csv")
g_raw = pd.read_csv("data/GDPC1.csv")
n_rows = {"UNRATE.csv": int(len(u_raw)), "GDPC1.csv": int(len(g_raw))}

# --- unemployment: monthly -> quarterly average
u = u_raw.copy()
u["date"] = pd.to_datetime(u["observation_date"])
u["UNRATE"] = pd.to_numeric(u["UNRATE"], errors="coerce")
miss = u.loc[u["UNRATE"].isna(), "observation_date"].tolist()
if miss:
    issues.append({"issue": f"UNRATE missing value(s) at {', '.join(miss)}",
                   "handling": "outside the 1959Q4-2019Q4 span needed; dropped by window filter, no effect on estimates"})
u["q"] = u["date"].dt.to_period("Q")
# quarterly average only for quarters with all 3 monthly values present
cnt_all = u.groupby("q")["UNRATE"].count()
uq_all = u.groupby("q")["UNRATE"].mean()[cnt_all == 3]

# --- GDP
g = g_raw.copy()
g["q"] = pd.to_datetime(g["observation_date"]).dt.to_period("Q")
g["GDPC1"] = pd.to_numeric(g["GDPC1"], errors="coerce")
assert g["GDPC1"].notna().all() and not g["q"].duplicated().any()
gq = g.set_index("q")["GDPC1"]

# Exclusion ledger in quarterly units, starting from the 319 GDPC1 quarters
panel = pd.DataFrame({"y": gq})
n0 = len(panel)
panel["u"] = uq_all.reindex(panel.index)
panel = panel.dropna()
excl.append(("merge_unrate", "keep GDPC1 quarters with a complete (3-month, non-missing) UNRATE quarterly average",
             n0, len(panel), n0 - len(panel),
             "UNRATE (945 monthly rows) averaged to quarters; 1947 has no UNRATE, 2025Q4 lacks Oct-2025 value"))
n0 = len(panel)
panel = panel.loc[FIRST - 1:LAST]
excl.append(("window", "keep 1959Q4-2019Q4 (1959Q4 only for differencing)", n0, len(panel), n0 - len(panel),
             "information rule: fixed sample window, no extension"))
assert len(panel) == 241 and (cnt_all.loc[FIRST - 1:LAST] == 3).all()
df = panel.sort_index()
df["du"] = df["u"].diff()
df["g"] = 100 * np.log(df["y"] / df["y"].shift(1))
n0 = len(df)
df = df.loc[FIRST:LAST].dropna()
excl.append(("differencing", "drop 1959Q4 (levels used only to form 1960Q1 differences)", n0, len(df), n0 - len(df),
             "first difference unavailable for base quarter"))
T = len(df)
assert T == 240

X = sm.add_constant(df["g"].values)
y = df["du"].values
ols = sm.OLS(y, X).fit()
# Newey-West (1994) plug-in rule of thumb: floor(4*(T/100)^(2/9))
L = int(math.floor(4 * (T / 100) ** (2 / 9)))
hac = sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": L, "kernel": "bartlett", "use_correction": False})

# manual check of HAC SE
e = ols.resid; Xe = X * e[:, None]
S = Xe.T @ Xe
for j in range(1, L + 1):
    w = 1 - j / (L + 1); G = Xe[j:].T @ Xe[:-j]; S += w * (G + G.T)
XtXi = np.linalg.inv(X.T @ X)
se_manual = float(np.sqrt((XtXi @ S @ XtXi)[1, 1]))
assert abs(se_manual - hac.bse[1]) < 1e-10

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)))

sens = {}
for l in [0, 1, 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])

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(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_stat_hac": float(beta / se_hac), "p_value_hac": p_hac,
    "residual_acf_lag1": float(pd.Series(e).autocorr(1)),
    "se_hac_by_lags": sens,
    "beta_ci_iid_low": float(beta - z * ols.bse[1]), "beta_ci_iid_high": float(beta + z * ols.bse[1]),
}
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
pd.DataFrame(excl, columns=["step", "rule", "n_before", "n_after", "n_removed", "rationale"]).to_csv("out/exclusions.csv", index=False)

fig, ax = plt.subplots(figsize=(6.5, 4.5))
ax.scatter(df["g"], df["du"], s=12, alpha=0.6)
xs = np.linspace(df["g"].min(), df["g"].max(), 100)
ax.plot(xs, alpha + beta * xs, color="red", label=f"OLS: du = {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("Change in unemployment rate du_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=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))