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

ROOT = os.getcwd(); OUT = os.path.join(ROOT, "out"); os.makedirs(OUT, exist_ok=True)
FIRST, LAST = pd.Period("1960Q1", "Q"), pd.Period("2019Q4", "Q")
EXPECTED_SHA = {"UNRATE.csv": "621ff6fadf50910cb7a21ee20ba610c247cb043b29b4e05c4c1a72fadca2a06c",
                "GDPC1.csv": "1c6c6e9f59e383e20c849dd299175c7a307526c09abdc18f18da9754fc5b3343"}
issues, excl = [], []

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

for f, h in EXPECTED_SHA.items():
    if sha(os.path.join("data", f)) != h:
        issues.append({"issue": f"{f} SHA-256 differs from data/README.md", "handling": "used as given; flagged"})

u = pd.read_csv("data/UNRATE.csv", dtype={"UNRATE": str}, keep_default_na=False)
g = pd.read_csv("data/GDPC1.csv", dtype={"GDPC1": str}, keep_default_na=False)
n_rows = {"UNRATE.csv": int(len(u)), "GDPC1.csv": int(len(g))}
u["date"] = pd.to_datetime(u.observation_date); g["date"] = pd.to_datetime(g.observation_date)
u["UNRATE"] = pd.to_numeric(u.UNRATE.replace("", np.nan)); g["GDPC1"] = pd.to_numeric(g.GDPC1.replace("", np.nan))
assert u.date.is_unique and g.date.is_unique

# Step 1: drop months with missing UNRATE (monthly rows)
miss = u[u.UNRATE.isna()]
if len(miss):
    issues.append({"issue": "UNRATE missing (empty field) for month(s) " + ", ".join(miss.observation_date) +
                   "; outside the 1959Q4-2019Q4 window used",
                   "handling": "row dropped; affects no quarter used in the analysis"})
n0 = len(u); u = u.dropna(subset=["UNRATE"])
# (monthly row drop is documented in data_issues; the exclusion ledger below is in quarterly units
#  so that it chains from the quarterly panel to n_analysed)

# Quarterly aggregation
u["q"] = u.date.dt.to_period("Q")
uq = u.groupby("q").UNRATE.agg(["mean", "count"])
g["q"] = g.date.dt.to_period("Q"); gq = g.set_index("q").GDPC1
need = pd.period_range(FIRST - 1, LAST, freq="Q")
assert (uq["count"].reindex(need) == 3).all(), "incomplete quarter of monthly UNRATE in window"
assert gq.reindex(need).notna().all(), "missing GDPC1 in window"

# Step 2: build quarterly panel over the union of quarters, then restrict to the window
allq = pd.period_range(min(uq.index.min(), gq.index.min()), max(uq.index.max(), gq.index.max()), freq="Q")
full = pd.DataFrame({"u": uq["mean"], "Y": gq}).reindex(allq)
full["du"] = full.u.diff(); full["g"] = 100 * np.log(full.Y / full.Y.shift(1))
sample = full.loc[FIRST:LAST]
excl.append(("1", "restrict to regression window 1960Q1-2019Q4 (1959Q4 levels used only for differences)",
             len(full), len(sample), "unit: quarters (union of both series 1947Q1-2026Q3); window fixed by TASK.md; quarters 1947Q1-1959Q4 and 2020Q1-2026Q3 out of window"))
n_b = len(sample); sample = sample.dropna(subset=["du", "g"])
excl.append(("2", "drop window quarters with missing du or g", n_b, len(sample), "none expected"))
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)

n = len(sample); assert n == 240
X = sm.add_constant(sample.g.values); y = sample.du.values
ols = sm.OLS(y, X).fit()
# Bandwidth: Newey-West (1994) rule of thumb for the Bartlett kernel, floor(4*(T/100)^(2/9))
L = int(np.floor(4 * (n / 100) ** (2 / 9)))
def hac_se(lag):
    r = sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": lag, "kernel": "bartlett", "use_correction": False})
    return float(r.bse[1])
b, a = float(ols.params[1]), float(ols.params[0])
se_iid, se_hac = float(ols.bse[1]), hac_se(L)
z = float(stats.norm.ppf(0.975))
lo, hi = b - z * se_hac, b + z * se_hac
p_hac = float(2 * stats.norm.sf(abs(b / se_hac)))
sens = {str(l): hac_se(l) for l in [0, 1, 2, 4, 6, 8, 12]}
res_acf = acf(ols.resid, nlags=4, fft=False)[1:]

results = {
    "task_id": "F2-conf-04", "beta_hat": b, "alpha_hat": a, "r_squared": float(ols.rsquared),
    "se_iid": se_iid, "se_hac": se_hac, "hac_lags": L,
    "beta_ci_low": float(lo), "beta_ci_high": float(hi),
    "okun_negative_5pct": bool(b < 0 and p_hac < 0.05),
    "mean_growth": float(sample.g.mean()),
    "first_quarter": str(sample.index[0]), "last_quarter": str(sample.index[-1]),
    "n_analysed": int(n), "feasibility": "complete", "feasibility_reason": None,
    "data_issues": issues, "n_input_rows": n_rows, "random_seed": None,
    # supplementary keys (not required by TASK.md)
    "t_stat_hac": float(b / se_hac), "p_value_hac": p_hac, "z_crit_95": z,
    "baseline_ci_iid_low": float(b - z * se_iid), "baseline_ci_iid_high": float(b + z * se_iid),
    "se_hac_by_lag": sens, "resid_acf_lags1to4": [float(x) for x in res_acf],
    "mean_du": float(sample.du.mean()),
}
with open(os.path.join(OUT, "results.json"), "w") as fh:
    json.dump(results, fh, indent=2)

fig, ax = plt.subplots(figsize=(6, 4.5))
ax.scatter(sample.g, sample.du, s=12, alpha=0.7)
xs = np.linspace(sample.g.min(), sample.g.max(), 50)
ax.plot(xs, a + b * xs, "r-", label=f"OLS: Δu = {a:.3f} + ({b:.3f}) g")
ax.axhline(0, color="grey", lw=0.5); ax.axvline(0, color="grey", lw=0.5)
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=240)"); 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"]}))