"""Okun's law, U.S. 1960Q1-2019Q4, OLS with Newey-West HAC SE. Reads only data/, writes out/."""
import json, hashlib, os
import numpy as np, pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy import stats

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
D, O = os.path.join(ROOT, "data"), os.path.join(ROOT, "out")
FIRST, LAST = pd.Period("1960Q1", "Q"), pd.Period("2019Q4", "Q")

def sha(p):
    return hashlib.sha256(open(p, "rb").read()).hexdigest()

un = pd.read_csv(os.path.join(D, "UNRATE.csv"))
gd = pd.read_csv(os.path.join(D, "GDPC1.csv"))
n_rows = {"UNRATE.csv": int(len(un)), "GDPC1.csv": int(len(gd))}
data_issues = []
miss = un[un["UNRATE"].isna()]
for d in miss["observation_date"]:
    data_issues.append({"issue": f"UNRATE missing (empty) value at {d}",
                        "handling": "Quarter containing it has <3 monthly values and is dropped; it lies outside the 1959Q4-2019Q4 levels window, so results are unaffected."})

# quarterly average unemployment, only quarters with all 3 months
un["q"] = pd.to_datetime(un["observation_date"]).dt.to_period("Q")
g = un.groupby("q")["UNRATE"]
uq = g.mean()[g.count() == 3].rename("u")
n_un_q = int(un["q"].nunique())
gd["q"] = pd.to_datetime(gd["observation_date"]).dt.to_period("Q")
gq = gd.set_index("q")["GDPC1"].astype(float)

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

df = pd.concat([gq, uq], axis=1, join="inner").sort_index()
log(1, "keep GDPC1 quarters that have a complete (3-month) UNRATE quarterly average", len(gq), len(df),
    "UNRATE starts 1948Q1 (GDPC1 starts 1947Q1); 2025Q4 lacks October 2025 UNRATE")
# differences only between consecutive calendar quarters
idx = df.index
consec = pd.Series([False] + [idx[i].ordinal - idx[i - 1].ordinal == 1 for i in range(1, len(idx))], index=idx)
df["du"] = df["u"].diff()
df["g"] = 100 * np.log(df["GDPC1"] / df["GDPC1"].shift(1))
df.loc[~consec, ["du", "g"]] = np.nan
nb = len(df); df = df.dropna(subset=["du", "g"])
log(2, "drop quarters whose previous calendar quarter is unavailable (first differences undefined)", nb, len(df),
    "first available quarter and quarters following a dropped quarter")
nb = len(df); df = df[(df.index >= FIRST) & (df.index <= LAST)]
log(3, "restrict to regression window 1960Q1-2019Q4", nb, len(df), "task-defined sample; no extension of window")
assert len(df) == 240 and df.index[0] == FIRST and df.index[-1] == LAST
pd.DataFrame(ledger).to_csv(os.path.join(O, "exclusions.csv"), index=False)

# OLS
y = df["du"].to_numpy(); x = df["g"].to_numpy(); T = len(y)
X = np.column_stack([np.ones(T), x])
XtXi = np.linalg.inv(X.T @ X)
b = XtXi @ X.T @ y
e = y - X @ b
k = X.shape[1]
r2 = 1 - (e @ e) / np.sum((y - y.mean()) ** 2)
se_iid = float(np.sqrt((e @ e) / (T - k) * XtXi[1, 1]))

def nw_se(L):
    """Newey-West (1987) Bartlett-kernel HAC, no small-sample df correction."""
    u = X * e[:, None]
    S = u.T @ u
    for l in range(1, L + 1):
        w = 1 - l / (L + 1)
        G = u[l:].T @ u[:-l]
        S += w * (G + G.T)
    V = XtXi @ S @ XtXi
    return float(np.sqrt(V[1, 1]))

# Bandwidth: Newey-West (1994) / Stock-Watson rule of thumb floor(4*(T/100)^(2/9))
L = int(np.floor(4 * (T / 100) ** (2 / 9)))
se_hac = nw_se(L)
z = stats.norm.ppf(0.975)
lo, hi = float(b[1] - z * se_hac), float(b[1] + z * se_hac)
t_hac = float(b[1] / se_hac)
p_hac = float(2 * stats.norm.sf(abs(t_hac)))
neg = bool(b[1] < 0 and p_hac < 0.05)

# cross-check with statsmodels (no df correction)
import statsmodels.api as sm
m = sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": L, "use_correction": False})
assert abs(m.bse[1] - se_hac) < 1e-10 and abs(m.params[1] - b[1]) < 1e-12
m0 = sm.OLS(y, X).fit()
assert abs(m0.bse[1] - se_iid) < 1e-10

sens = {str(l): nw_se(l) for l in [0, 1, 2, 4, 6, 8, 12]}
rho1 = float(np.corrcoef(e[1:], e[:-1])[0, 1])

res = {
    "task_id": "F2-conf-04",
    "beta_hat": float(b[1]), "alpha_hat": float(b[0]), "r_squared": float(r2),
    "se_iid": se_iid, "se_hac": se_hac, "hac_lags": L,
    "beta_ci_low": lo, "beta_ci_high": hi,
    "okun_negative_5pct": neg,
    "mean_growth": float(x.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_rows, "random_seed": None,
    # supplementary
    "baseline_beta_ci_low_iid": float(b[1] - z * se_iid), "baseline_beta_ci_high_iid": float(b[1] + z * se_iid),
    "t_stat_hac": t_hac, "p_value_hac": p_hac, "critical_value": float(z),
    "se_hac_by_lags": sens, "resid_lag1_autocorr": rho1,
    "mean_du": float(y.mean()),
    "input_sha256": {f: sha(os.path.join(D, f)) for f in ["UNRATE.csv", "GDPC1.csv"]},
}
with open(os.path.join(O, "results.json"), "w") as f:
    json.dump(res, f, indent=2)

fig, ax = plt.subplots(figsize=(6.5, 4.5))
ax.scatter(x, y, s=12, alpha=0.6, label="quarters 1960Q1-2019Q4")
xs = np.linspace(x.min(), x.max(), 100)
ax.plot(xs, b[0] + b[1] * xs, "r-", label=f"OLS: Δu = {b[0]:.3f} {b[1]:+.3f} g")
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.axhline(0, color="grey", lw=0.5); ax.axvline(0, color="grey", lw=0.5)
ax.legend(fontsize=8); fig.tight_layout()
fig.savefig(os.path.join(O, "figure1.png"), dpi=120)
print(f"beta={b[1]:.4f} se_iid={se_iid:.4f} se_hac(L={L})={se_hac:.4f} CI=[{lo:.4f},{hi:.4f}] neg5%={neg}")