"""F2-conf-04: Okun's law, U.S. quarterly 1960Q1-2019Q4, OLS with Newey-West HAC SE.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Reads only data/. Deterministic (no randomness; random_seed = null)."""
import json, os
import numpy as np, pandas as pd
from scipy import stats
import statsmodels.api as sm
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

OUT = "out"
os.makedirs(OUT, exist_ok=True)
FIRST_Q, LAST_Q = pd.Period("1960Q1", "Q"), pd.Period("2019Q4", "Q")
LEVEL_START = FIRST_Q - 1  # 1959Q4 levels needed for first differences

un_raw = pd.read_csv("data/UNRATE.csv", keep_default_na=False, dtype=str)
gdp_raw = pd.read_csv("data/GDPC1.csv", keep_default_na=False, dtype=str)
n_input_rows = {"UNRATE.csv": int(len(un_raw)), "GDPC1.csv": int(len(gdp_raw))}

data_issues, excl = [], []
def log_excl(step, rule, nb, na, why):
    excl.append(dict(step=step, rule=rule, n_before=nb, n_after=na, n_removed=nb - na, rationale=why))

# ---- unemployment: monthly -> quarterly average
un = un_raw.copy()
un["date"] = pd.to_datetime(un["observation_date"])
un["v"] = pd.to_numeric(un["UNRATE"].replace("", np.nan), errors="coerce")
miss = un[un["v"].isna()]
for d in miss["date"]:
    q = d.to_period("Q")
    inwin = LEVEL_START <= q <= LAST_Q
    data_issues.append({"issue": f"UNRATE missing (empty field) for {d:%Y-%m}",
                        "handling": ("inside required window -> quarter dropped" if inwin else
                                     "outside 1959Q4-2019Q4 level window; not used, no effect on estimates")})
un["q"] = un["date"].dt.to_period("Q")
assert not un["date"].duplicated().any()
g = un.groupby("q")["v"]
uq = pd.DataFrame({"u": g.mean(), "n_months": g.count()})
# ---- GDP
gd = gdp_raw.copy()
gd["q"] = pd.to_datetime(gd["observation_date"]).dt.to_period("Q")
gd["Y"] = pd.to_numeric(gd["GDPC1"].replace("", np.nan), errors="coerce")
assert not gd["q"].duplicated().any() and gd["Y"].notna().all()
gq = gd.set_index("q")[["Y"]]

# ---- quarterly panel restricted to level window 1959Q4-2019Q4
panel = uq.join(gq, how="outer").sort_index()
n_q_all = len(panel)
lev = panel.loc[LEVEL_START:LAST_Q].copy()
log_excl("1", "keep quarters 1959Q4-2019Q4 (levels window; information rule)", n_q_all, len(lev),
         "task fixes regression window 1960Q1-2019Q4; 1959Q4 levels needed only for first differences")
assert (lev["n_months"] == 3).all() and lev["Y"].notna().all(), "incomplete quarter in window"
assert len(pd.period_range(LEVEL_START, LAST_Q, freq="Q")) == len(lev)
lev["du"] = lev["u"].diff()
lev["g"] = 100 * np.log(lev["Y"] / lev["Y"].shift(1))
reg = lev.loc[FIRST_Q:LAST_Q].dropna(subset=["du", "g"])
log_excl("2", "drop 1959Q4 (no prior level for differencing; levels-only quarter)", len(lev), len(reg),
         "1959Q4 is used only as the base level for 1960Q1 differences")
n = len(reg)
assert n == 240

# ---- OLS
y, x = reg["du"].values, reg["g"].values
X = sm.add_constant(x)
ols = sm.OLS(y, X).fit()
alpha_hat, beta_hat = map(float, ols.params)
se_iid = float(ols.bse[1])

# Newey-West bandwidth: classic rule floor(4*(T/100)^(2/9)) (Newey & West 1994);
# cross-check with floor(0.75*T^(1/3)) (Stock & Watson textbook rule). Both = 4 for T=240.
L_nw94 = int(np.floor(4 * (n / 100) ** (2 / 9)))
L_sw = int(np.floor(0.75 * n ** (1 / 3)))
hac_lags = L_nw94

def nw_se(L, correction=False):
    r = sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": L, "kernel": "bartlett",
                                                  "use_correction": correction})
    return float(r.bse[1])

# manual Bartlett-kernel check (no small-sample df correction)
def nw_manual(L):
    e = ols.resid
    u = X * e[:, None]
    S = u.T @ u
    for j in range(1, L + 1):
        w = 1 - j / (L + 1)
        G = u[j:].T @ u[:-j]
        S += w * (G + G.T)
    XtXi = np.linalg.inv(X.T @ X)
    return float(np.sqrt((XtXi @ S @ XtXi)[1, 1]))

se_hac = nw_se(hac_lags)
assert abs(se_hac - nw_manual(hac_lags)) < 1e-10
z = stats.norm.ppf(0.975)
ci_lo, ci_hi = beta_hat - z * se_hac, beta_hat + z * se_hac
t_hac = beta_hat / se_hac
p_hac = float(2 * stats.norm.sf(abs(t_hac)))
okun_neg = bool(beta_hat < 0 and p_hac < 0.05)

# ---- sensitivity table (written to CSV, not results.json)
resid_ac = [float(pd.Series(ols.resid).autocorr(k)) for k in range(1, 5)]
sens = pd.DataFrame([{"lags": L, "se_hac": nw_se(L), "se_hac_dfcorr": nw_se(L, True)} for L in range(0, 13)])
sens["ci_low"] = beta_hat - z * sens["se_hac"]
sens["ci_high"] = beta_hat + z * sens["se_hac"]
sens["p_value"] = 2 * stats.norm.sf(np.abs(beta_hat / sens["se_hac"]))
sens.to_csv(f"{OUT}/hac_sensitivity.csv", index=False)
pd.DataFrame(excl).to_csv(f"{OUT}/exclusions.csv", index=False)

# ---- figure
fig, ax = plt.subplots(figsize=(6.5, 4.5))
ax.scatter(x, y, s=12, alpha=0.6, label="quarters 1960Q1–2019Q4")
xx = np.linspace(x.min(), x.max(), 100)
ax.plot(xx, alpha_hat + beta_hat * xx, color="C3",
        label=f"OLS: Δu = {alpha_hat:.3f} {beta_hat:+.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=%d)" % n)
ax.legend(fontsize=8)
fig.tight_layout(); fig.savefig(f"{OUT}/figure1.png", dpi=150); plt.close(fig)

results = {
    "task_id": "F2-conf-04",
    "beta_hat": beta_hat, "alpha_hat": alpha_hat, "r_squared": float(ols.rsquared),
    "se_iid": se_iid, "se_hac": se_hac, "hac_lags": hac_lags,
    "beta_ci_low": float(ci_lo), "beta_ci_high": float(ci_hi),
    "okun_negative_5pct": okun_neg,
    "mean_growth": float(np.mean(x)),
    "first_quarter": str(reg.index[0]), "last_quarter": str(reg.index[-1]),
    "n_analysed": int(n),
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": data_issues, "n_input_rows": n_input_rows, "random_seed": None,
    # supplementary keys
    "t_stat_hac": float(t_hac), "p_value_hac": p_hac,
    "baseline_ci_iid_low": float(beta_hat - z * se_iid), "baseline_ci_iid_high": float(beta_hat + z * se_iid),
    "hac_lags_rule_nw1994": L_nw94, "hac_lags_rule_sw": L_sw,
    "resid_autocorr_lag1_4": resid_ac,
    "se_hac_max_over_lags_0_12": float(sens["se_hac"].max()),
    "p_value_hac_max_over_lags_0_12": float(sens["p_value"].max()),
}
with open(f"{OUT}/results.json", "w") as f:
    json.dump(results, f, indent=2)
print(json.dumps({k: results[k] for k in ["beta_hat", "alpha_hat", "r_squared", "se_iid", "se_hac",
                                          "hac_lags", "beta_ci_low", "beta_ci_high", "okun_negative_5pct",
                                          "mean_growth", "n_analysed"]}, indent=1))