"""F2-conf-04: Okun's law, US quarterly 1960Q1-2019Q4, OLS with Newey-West HAC SE.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py"""
import json, os, math
import numpy as np, pandas as pd
import statsmodels.api as sm
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

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

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

data_issues = []
un = un_raw.copy()
un["UNRATE"] = pd.to_numeric(un["UNRATE"], errors="coerce")
un["q"] = pd.to_datetime(un["observation_date"]).dt.to_period("Q")
miss = un[un["UNRATE"].isna()]
for _, r in miss.iterrows():
    inwin = LEVEL_START <= r["q"] <= LAST
    data_issues.append({"issue": f"UNRATE missing for {r['observation_date']}",
                        "handling": ("inside window: quarter would be dropped" if inwin else
                                     "outside the 1959Q4-2019Q4 window; no effect on the analysis")})
gdp = gdp_raw.copy()
gdp["GDPC1"] = pd.to_numeric(gdp["GDPC1"], errors="coerce")
gdp["q"] = pd.to_datetime(gdp["observation_date"]).dt.to_period("Q")
for _, r in gdp[gdp["GDPC1"].isna()].iterrows():
    data_issues.append({"issue": f"GDPC1 missing for {r['observation_date']}", "handling": "quarter dropped"})

# quarterly average of the three monthly observations; require all 3 present
uq = un.groupby("q")["UNRATE"].agg(["mean", "count"])
uq = uq[uq["count"] == 3]["mean"].rename("u")
df = gdp.set_index("q")[["GDPC1"]].join(uq, how="left").sort_index()
excl = []
n0 = len(df)  # n_raw = GDPC1 quarters
df = df.dropna(subset=["u", "GDPC1"])
excl.append(("0", "drop quarters lacking GDPC1 or a complete 3-month UNRATE average", n0, len(df),
             "1947Q1-Q4 precede UNRATE start; 2025Q4 lacks Oct-2025 UNRATE; all outside window"))
n0 = len(df)
df = df[(df.index >= LEVEL_START) & (df.index <= LAST)]
excl.append(("1", "keep quarters 1959Q4-2019Q4 (levels)", n0, len(df),
             "task window; 1959Q4 kept only for differencing"))
assert len(df) == 241 and (df.index[1:].astype(int) - df.index[:-1].astype(int) == 1).all()
df["du"] = df["u"].diff()
df["g"] = 100 * np.log(df["GDPC1"] / df["GDPC1"].shift(1))
n1 = len(df)
df = df.dropna(subset=["du", "g"])
excl.append(("2", "drop 1959Q4 (no prior level for differences)", n1, len(df),
             "first-difference loses one observation"))
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(os.path.join(OUT, "exclusions.csv"), index=False)

T = len(df)
X = sm.add_constant(df["g"].values)
y = df["du"].values
ols = sm.OLS(y, X).fit()
# Newey-West (1994) plug-in style rule of thumb: 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 = 1.959963984540054
lo, hi = beta - z * se_hac, beta + z * se_hac
t_hac = 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, "use_correction": False})
    sens[str(L)] = round(float(r.bse[1]), 10)
resid = ols.resid
rho = [float(np.corrcoef(resid[k:], resid[:-k])[0, 1]) for k in range(1, 5)]

results = {
    "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": lo, "beta_ci_high": hi,
    "okun_negative_5pct": bool(beta < 0 and abs(t_hac) > z),
    "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": data_issues, "n_input_rows": n_input_rows, "random_seed": None,
    # supplementary keys
    "t_stat_hac": float(t_hac),
    "p_value_hac_two_sided": float(2 * (1 - __import__("scipy.stats").stats.norm.cdf(abs(t_hac)))),
    "se_hac_by_lags": sens,
    "resid_autocorr_lag1_4": rho,
    "mean_du": float(df["du"].mean()),
}
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(results, f, indent=2)

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(), 50)
ax.plot(xs, alpha + beta * xs, color="red", label=f"OLS: Δu = {alpha:.3f} {beta:+.3f}·g")
ax.axhline(0, lw=0.5, color="k"); ax.axvline(0, lw=0.5, color="k")
ax.set_xlabel("g_t: real GDP growth, % q/q (not annualised)")
ax.set_ylabel("Δu_t: change in unemployment rate, pp")
ax.set_title("Okun's law, US 1960Q1–2019Q4 (n=%d)" % T); ax.legend()
fig.tight_layout(); fig.savefig(os.path.join(OUT, "figure1.png"), dpi=120)
print(json.dumps({k: results[k] for k in ["beta_hat", "se_iid", "se_hac", "hac_lags", "beta_ci_low",
                                          "beta_ci_high", "okun_negative_5pct", "n_analysed"]}, indent=1))