"""Okun's law (difference form), US 1960Q1-2019Q4. Run from workspace root:
   ~/wp29/venv/bin/python out/analysis.py
Reads only data/; writes out/results.json, out/table1.csv, out/figure1.png, out/exclusions.csv."""
import json, os
import numpy as np, 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)
gdp_raw = pd.read_csv("data/GDPC1.csv")
un_raw = pd.read_csv("data/UNRATE.csv")
n_input_rows = {"GDPC1.csv": int(len(gdp_raw)), "UNRATE.csv": int(len(un_raw))}
data_issues = []

# --- unemployment: monthly -> quarterly average (require all 3 months) ---
un = un_raw.copy()
un["date"] = pd.to_datetime(un["observation_date"])
un["UNRATE"] = pd.to_numeric(un["UNRATE"], errors="coerce")
miss = un.loc[un["UNRATE"].isna(), "observation_date"].tolist()
if miss:
    data_issues.append({"issue": f"UNRATE missing (blank) values at {miss}",
                        "handling": "All outside the 1959Q4-2019Q4 window; quarters with <3 valid months set to NaN; no effect on estimates."})
un["q"] = un["date"].dt.to_period("Q")
gq = un.groupby("q")["UNRATE"]
u_q = gq.mean().where(gq.count() == 3)

# --- GDP ---
gdp = gdp_raw.copy()
gdp["q"] = pd.to_datetime(gdp["observation_date"]).dt.to_period("Q")
gdp = gdp.set_index("q")["GDPC1"].astype(float)
data_issues.append({"issue": "Files extend beyond 2019Q4 (GDPC1 to 2026Q3, UNRATE to 2026-09); GDPC1 is the current vintage (revised data), not real-time.",
                    "handling": "All quarters after 2019Q4 dropped before any computation of estimates; revised-vintage use noted as limitation."})

df = pd.DataFrame({"gdp": gdp, "u": u_q}).sort_index()
df = df[df.index.isin(gdp.index)]  # GDP quarters define the panel
excl = []
n0 = len(df)
q_start_lag, q_start, q_end = pd.Period("1959Q4"), pd.Period("1960Q1"), pd.Period("2019Q4")
d1 = df[df.index >= q_start_lag]
excl.append(["1", "drop quarters before 1959Q4", n0, len(d1), n0 - len(d1), "outside question (1959Q4 kept only as lag level)"])
d2 = d1[d1.index <= q_end]
excl.append(["2", "drop quarters after 2019Q4", len(d1), len(d2), len(d1) - len(d2), "information rule: post-2019Q4 must not enter any estimate"])
d2 = d2.copy()
d2["du"] = d2["u"].diff()
d2["g"] = 100 * np.log(d2["gdp"]).diff()
d3 = d2[d2.index >= q_start].dropna(subset=["du", "g"])
excl.append(["3", "drop 1959Q4 (lag-only quarter; no difference)", len(d2), len(d3), len(d2) - len(d3), "first differences need previous quarter"])
pd.DataFrame(excl, columns=["step", "rule", "n_before", "n_after", "n_removed", "rationale"]).to_csv(f"{OUT}/exclusions.csv", index=False)
assert len(d3) == 240 and d3[["du", "g"]].notna().all().all()
d = d3

# --- baseline and full-sample OLS ---
y, x = d["du"].values, d["g"].values
n = len(y)
baseline_rmse = float(np.sqrt(np.mean((y - y.mean()) ** 2)))
X = sm.add_constant(x)
ols = sm.OLS(y, X).fit()
HAC_LAGS = int(np.floor(4 * (n / 100) ** (2 / 9)))  # Newey-West (1994) rule of thumb -> 4 for n=240
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 = float(hac.bse[1])
z = stats.norm.ppf(0.975)
ci_lo, ci_hi = beta - z * se, beta + z * se
model_rmse = float(np.sqrt(np.mean(ols.resid ** 2)))

# --- subperiods and Chow test ---
m1 = d.index <= pd.Period("1983Q4"); m2 = ~m1
fits = {}
for name, m in [("sub1", m1), ("sub2", m2)]:
    Xi = sm.add_constant(x[m]); fi = sm.OLS(y[m], Xi).fit()
    fh = sm.OLS(y[m], Xi).fit(cov_type="HAC", cov_kwds={"maxlags": HAC_LAGS, "kernel": "bartlett", "use_correction": False})
    fits[name] = (int(m.sum()), fi, fh)
ssr_p = ols.ssr; ssr_u = fits["sub1"][1].ssr + fits["sub2"][1].ssr; k = 2
F = ((ssr_p - ssr_u) / k) / (ssr_u / (n - 2 * k))
p = float(stats.f.sf(F, k, n - 2 * k))

# supplementary: HAC Wald test of slope-only break (dummy-interacted regression)
D = m2.astype(float)
Xi_full = np.column_stack([np.ones(n), x, D, D * x])
fit_int = sm.OLS(y, Xi_full).fit(cov_type="HAC", cov_kwds={"maxlags": HAC_LAGS, "kernel": "bartlett", "use_correction": False})
slope_break_p = float(fit_int.pvalues[3])

conclusion = "negative_effect" if ci_hi < 0 else ("positive_effect" if ci_lo > 0 else "no_detectable_effect")
res = {
    "task_id": "F6-conf-03", "feasibility": "complete", "feasibility_reason": None,
    "data_issues": data_issues, "n_input_rows": n_input_rows, "n_analysed": int(n), "random_seed": None,
    "sample_start": d.index[0].start_time.strftime("%Y-%m-%d"),
    "sample_end": d.index[-1].start_time.strftime("%Y-%m-%d"),
    "n_sub1": fits["sub1"][0], "n_sub2": fits["sub2"][0],
    "baseline_metric": baseline_rmse, "model_metric": model_rmse,
    "effect_estimate": beta, "intercept": alpha, "effect_se": se, "hac_lags": HAC_LAGS,
    "effect_ci_low": float(ci_lo), "effect_ci_high": float(ci_hi),
    "beta_sub1": float(fits["sub1"][1].params[1]), "beta_sub2": float(fits["sub2"][1].params[1]),
    "chow_f_stat": float(F), "chow_p_value": p,
    "stability_conclusion": "stable" if p >= 0.05 else "unstable",
    "conclusion": conclusion,
    # supplementary (not required by TASK.md)
    "ols_classical_se_full": float(ols.bse[1]),
    "hac_se_sensitivity": {str(L): float(sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": L, "kernel": "bartlett", "use_correction": False}).bse[1]) for L in (0, 2, 8, 12)},
    "slope_break_hac_pvalue": slope_break_p,
    "intercept_sub1": float(fits["sub1"][1].params[0]), "intercept_sub2": float(fits["sub2"][1].params[0]), "r_squared_full": float(ols.rsquared),
}
with open(f"{OUT}/results.json", "w") as fh_:
    json.dump(res, fh_, indent=2)

# --- table: slope_se is HAC (Newey-West, same lag) in every row ---
rows = [["full_1960Q1_2019Q4", n, alpha, beta, se]]
for name, lab in [("sub1", "sub1_1960Q1_1983Q4"), ("sub2", "sub2_1984Q1_2019Q4")]:
    ni, fi, fh = fits[name]
    rows.append([lab, ni, float(fi.params[0]), float(fi.params[1]), float(fh.bse[1])])
pd.DataFrame(rows, columns=["period", "n", "intercept", "slope", "slope_se"]).to_csv(f"{OUT}/table1.csv", index=False)

# --- figure ---
fig, ax = plt.subplots(figsize=(7, 5))
grid = np.linspace(x.min(), x.max(), 50)
for (name, m, c, lab) in [("sub1", m1, "tab:blue", "1960Q1-1983Q4"), ("sub2", m2, "tab:red", "1984Q1-2019Q4")]:
    fi = fits[name][1]
    ax.scatter(x[m], y[m], s=12, alpha=0.6, color=c, label=f"{lab} (n={m.sum()})")
    ax.plot(grid, fi.params[0] + fi.params[1] * grid, color=c, lw=2,
            label=f"fit {lab}: slope={fi.params[1]:.3f}")
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(f"Okun's law, US 1960Q1-2019Q4 (Chow F={F:.2f}, p={p:.3f})")
ax.legend(fontsize=8); fig.tight_layout(); fig.savefig(f"{OUT}/figure1.png", dpi=120)
print(json.dumps({k: res[k] for k in ["effect_estimate", "effect_se", "effect_ci_low", "effect_ci_high",
                                      "beta_sub1", "beta_sub2", "chow_f_stat", "chow_p_value",
                                      "baseline_metric", "model_metric"]}, indent=1))