"""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/; writes out/results.json, out/exclusions.csv, out/figure1.png. No randomness."""
import json, math, os
import numpy as np
import pandas as pd
import statsmodels.api as sm
from scipy import stats
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

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

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

def add_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["date"] = pd.to_datetime(un["observation_date"])
un["v"] = pd.to_numeric(un["UNRATE"].replace("", np.nan), errors="coerce")
miss = un[un["v"].isna()]
if len(miss):
    inwin = [d for d in miss["date"] if LEVEL_START <= d.to_period("Q") <= LAST]
    data_issues.append({
        "issue": f"UNRATE missing (empty) for {', '.join(miss['observation_date'])}",
        "handling": ("outside the 1959Q4-2019Q4 window; does not affect the analysis" if not inwin
                     else "INSIDE window: affected quarter averaged over available months")})
assert not un["date"].duplicated().any() and not gd["observation_date"].duplicated().any()
un["q"] = un["date"].dt.to_period("Q")
m = un[(un["q"] >= LEVEL_START) & (un["q"] <= LAST)]
assert m["v"].notna().all()
cnt = m.groupby("q")["v"].count()
assert (cnt == 3).all(), "incomplete quarter in window"
u = m.groupby("q")["v"].mean()

# --- GDP
gd["q"] = pd.to_datetime(gd["observation_date"]).dt.to_period("Q")
gd["v"] = pd.to_numeric(gd["GDPC1"], errors="coerce")
assert gd["v"].notna().all()
y = gd.set_index("q")["v"]
y = y[(y.index >= LEVEL_START) & (y.index <= LAST)]

df = pd.DataFrame({"u": u, "y": y})
assert df.notna().all().all() and len(df) == 241
df["du"] = df["u"].diff()
df["g"] = 100 * np.log(df["y"] / df["y"].shift(1))

# exclusion ledger (in units of quarters, starting from the full GDP quarterly range)
nq_all = len(pd.period_range(gd["q"].min(), gd["q"].max(), freq="Q"))
add_excl(1, "keep quarters 1959Q4-2019Q4 (levels window)", nq_all, len(df),
         "task window 1960Q1-2019Q4 plus 1959Q4 levels for first differences; window not extended")
d = df.dropna().loc[FIRST:LAST]
add_excl(2, "drop 1959Q4 (no first difference; levels only)", len(df), len(d),
         "1959Q4 used only as lag level for 1960Q1 differences")
n = len(d)
assert n == 240

# --- OLS: baseline iid SE and Newey-West HAC SE
X = sm.add_constant(d["g"].values)
Y = d["du"].values
ols = sm.OLS(Y, X).fit()
T = n
hac_lags = int(math.floor(4 * (T / 100) ** (2 / 9)))  # Newey-West (1994) plug-in rule -> 4
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)
lo, hi = beta - z * se_hac, beta + z * se_hac
p_hac = float(2 * stats.norm.sf(abs(beta / se_hac)))

# sensitivity of HAC SE to bandwidth (reported, not used for headline)
sens = {}
for L in [0, 1, 2, 4, 6, 8, 12]:
    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(pd.Series(ols.resid).autocorr(1))

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(lo), "beta_ci_high": float(hi),
    "okun_negative_5pct": bool(beta < 0 and p_hac < 0.05),
    "mean_growth": float(d["g"].mean()),
    "first_quarter": str(d.index[0]), "last_quarter": str(d.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_hac": float(beta / se_hac), "p_value_hac": p_hac,
    "se_hac_by_lags": sens, "se_hac_lag4_dfcorrected": hac_dfc,
    "residual_autocorr_lag1": resid_ac1,
    "mean_du": float(d["du"].mean()),
    "baseline_ci_iid_low": float(beta - z * se_iid), "baseline_ci_iid_high": float(beta + z * se_iid),
}
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(res, f, indent=2)
pd.DataFrame(excl).to_csv(os.path.join(OUT, "exclusions.csv"), index=False)

fig, ax = plt.subplots(figsize=(6.5, 4.5))
ax.scatter(d["g"], d["du"], s=12, alpha=0.6)
gx = np.linspace(d["g"].min(), d["g"].max(), 100)
ax.plot(gx, alpha + beta * gx, 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("Δ unemployment rate Δu_t (pp)")
ax.set_title("Okun's law, U.S. 1960Q1–2019Q4 (n=240)")
ax.legend(); fig.tight_layout()
fig.savefig(os.path.join(OUT, "figure1.png"), dpi=120)
print(json.dumps({k: res[k] for k in ["beta_hat", "alpha_hat", "se_iid", "se_hac", "hac_lags",
                                     "beta_ci_low", "beta_ci_high", "okun_negative_5pct", "n_analysed"]}, indent=1))