"""F2-conf-04: Okun's law, U.S. quarterly 1960Q1-2019Q4, OLS + Newey-West HAC SE.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Reads only data/; writes out/results.json, out/exclusions.csv, out/figure1.png. No randomness."""
import json, math
from pathlib import Path
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

ROOT = Path(__file__).resolve().parent.parent
DATA, OUT = ROOT / "data", ROOT / "out"
OUT.mkdir(exist_ok=True)
FIRST, LAST = pd.Period("1960Q1", "Q"), pd.Period("2019Q4", "Q")
LEVEL_START = FIRST - 1  # 1959Q4

un = pd.read_csv(DATA / "UNRATE.csv")
gd = pd.read_csv(DATA / "GDPC1.csv")
n_input_rows = {"UNRATE.csv": int(len(un)), "GDPC1.csv": int(len(gd))}
data_issues = []
excl = []

# --- unemployment: monthly -> quarterly average (require all 3 months) ---
un["date"] = pd.to_datetime(un["observation_date"])
assert not un["date"].duplicated().any()
miss = un[un["UNRATE"].isna()]
for d in miss["observation_date"]:
    p = pd.Period(d, "M").asfreq("Q")
    inside = LEVEL_START <= p <= LAST
    data_issues.append({"issue": f"UNRATE missing value at {d} (empty field)",
                        "handling": ("inside the needed 1959Q4-2019Q4 range; affected quarter would be dropped"
                                     if inside else
                                     "outside the 1959Q4-2019Q4 range needed for the regression; dropped, no effect on results")})
# monthly UNRATE is an input to the quarterly panel, not the unit of analysis; the exclusion
# ledger follows the quarterly panel (unit of analysis). Monthly rows: window filter then dropna.
un["q"] = un["date"].dt.to_period("Q")
un = un[(un["q"] >= LEVEL_START) & (un["q"] <= LAST)].dropna(subset=["UNRATE"])
qu = un.groupby("q")["UNRATE"].agg(["mean", "count"])
assert (qu["count"] == 3).all(), "incomplete quarter in window"
u = qu["mean"]

# --- GDP ---
gd["q"] = pd.to_datetime(gd["observation_date"]).dt.to_period("Q")
assert not gd["q"].duplicated().any() and gd["GDPC1"].notna().all() and (gd["GDPC1"] > 0).all()
n0 = len(gd); gd = gd[(gd["q"] >= LEVEL_START) & (gd["q"] <= LAST)]
excl.append(("G1", "keep GDPC1 quarters in 1959Q4-2019Q4", n0, len(gd), "window rule: outside 1960Q1-2019Q4 only 1959Q4 levels are used"))
y = gd.set_index("q")["GDPC1"]

n0 = len(y)
df = pd.concat({"u": u, "gdp": y}, axis=1, join="inner").sort_index()
excl.append(("M1", "inner-join with quarters having 3 non-missing UNRATE months", n0, len(df), "u_t requires complete quarter (none missing in window)"))
assert df.index.is_monotonic_increasing and len(df) == 241 and df.notna().all().all()
assert (df.index[1:] - df.index[:-1] == 1 * df.index.freq).all() if df.index.freq else True
df["du"] = df["u"].diff()
df["g"] = 100 * np.log(df["gdp"] / df["gdp"].shift(1))
n0 = len(df); reg = df.loc[FIRST:LAST].dropna()
excl.append(("Q1", "drop 1959Q4 (levels-only quarter used for first differences)", n0, len(reg), "differencing"))
assert len(reg) == 240

pd.DataFrame(excl, columns=["step", "rule", "n_before", "n_after", "rationale"]).assign(
    n_removed=lambda d: d.n_before - d.n_after)[["step", "rule", "n_before", "n_after", "n_removed", "rationale"]
].to_csv(OUT / "exclusions.csv", index=False)

# --- OLS ---
T = len(reg)
X = sm.add_constant(reg["g"].values)
Y = reg["du"].values
ols = sm.OLS(Y, X).fit()
# Newey-West (1994) plug-in rule of thumb for Bartlett kernel: floor(4 (T/100)^(2/9))
hac_lags = int(math.floor(4 * (T / 100) ** (2 / 9)))
hac = sm.OLS(Y, X).fit(cov_type="HAC", cov_kwds={"maxlags": hac_lags, "kernel": "bartlett", "use_correction": False})
beta, alpha = float(ols.params[1]), float(ols.params[0])
se_iid, se_hac = float(ols.bse[1]), float(hac.bse[1])
z = stats.norm.ppf(0.975)
ci = (beta - z * se_hac, beta + z * se_hac)
t_hac = beta / se_hac
p_hac = 2 * stats.norm.sf(abs(t_hac))

# sensitivity over bandwidths (not required keys; reported as extras)
sens = {}
for L in [0, 1, 2, 4, 8, 12, int(math.floor(0.75 * T ** (1 / 3)))]:
    r = sm.OLS(Y, X).fit(cov_type="HAC", cov_kwds={"maxlags": L, "use_correction": False})
    sens[str(L)] = float(r.bse[1])
hac_dfc = float(sm.OLS(Y, X).fit(cov_type="HAC", cov_kwds={"maxlags": hac_lags, "use_correction": True}).bse[1])
resid_ac1 = float(np.corrcoef(ols.resid[1:], ols.resid[:-1])[0, 1])

fq = lambda p: f"{p.year}Q{p.quarter}"
res = {
    "task_id": "F2-conf-04",
    "beta_hat": beta, "alpha_hat": alpha, "r_squared": float(ols.rsquared),
    "se_iid": se_iid, "se_hac": se_hac, "hac_lags": hac_lags,
    "beta_ci_low": float(ci[0]), "beta_ci_high": float(ci[1]),
    "okun_negative_5pct": bool(beta < 0 and p_hac < 0.05),
    "mean_growth": float(reg["g"].mean()),
    "first_quarter": fq(reg.index[0]), "last_quarter": fq(reg.index[-1]),
    "n_analysed": int(T),
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": data_issues, "n_input_rows": n_input_rows, "random_seed": None,
    # extras
    "baseline_ci_iid_low": float(beta - z * se_iid), "baseline_ci_iid_high": float(beta + z * se_iid),
    "t_stat_hac": float(t_hac), "p_value_hac": float(p_hac), "z_crit": float(z),
    "se_hac_small_sample_corrected": hac_dfc,
    "se_hac_by_lags": sens, "resid_lag1_autocorr": resid_ac1,
    "mean_du": float(reg["du"].mean()),
    "hac_lag_rule": "floor(4*(T/100)^(2/9)) (Newey-West 1994 rule of thumb, Bartlett kernel)",
}
(OUT / "results.json").write_text(json.dumps(res, indent=2) + "\n")

fig, ax = plt.subplots(figsize=(6.5, 4.5))
ax.scatter(reg["g"], reg["du"], s=12, alpha=0.7)
gg = np.linspace(reg["g"].min(), reg["g"].max(), 50)
ax.plot(gg, alpha + beta * gg, color="C3", 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("Change in unemployment rate Δ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)

print(f"beta={beta:.4f} alpha={alpha:.4f} R2={ols.rsquared:.3f} se_iid={se_iid:.4f} "
      f"se_hac(L={hac_lags})={se_hac:.4f} CI=[{ci[0]:.4f},{ci[1]:.4f}] neg5%={res['okun_negative_5pct']}")