"""Okun's law, US 1960Q1-2019Q4 (task F6-conf-03). Run from workspace root:
   ~/wp29/venv/bin/python out/analysis.py
Reads data/ only; 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 = pd.read_csv("data/GDPC1.csv", parse_dates=["observation_date"])
un = pd.read_csv("data/UNRATE.csv", parse_dates=["observation_date"])
n_input_rows = {"GDPC1.csv": int(len(gdp)), "UNRATE.csv": int(len(un))}
data_issues = []
miss = un[un["UNRATE"].isna()]
if len(miss):
    data_issues.append({"issue": f"UNRATE missing for {', '.join(miss.observation_date.dt.strftime('%Y-%m-%d'))}",
                        "handling": "Outside the 1959Q4-2019Q4 window; removed by the 2019Q4 cutoff, no imputation needed."})
data_issues.append({"issue": "Both files extend beyond 2019Q4 (to 2026), outside the question.",
                    "handling": "All rows after 2019-12-31 dropped before any computation."})

# --- quarterly unemployment: simple average of the 3 months; require complete quarters in window
un["q"] = un.observation_date.dt.to_period("Q")
uq = un.groupby("q")["UNRATE"].agg(["mean", "count"])
gdp["q"] = gdp.observation_date.dt.to_period("Q")
df = gdp.set_index("q")[["GDPC1"]].join(uq, how="left")

LAG0, START, END, SPLIT = pd.Period("1959Q4"), pd.Period("1960Q1"), pd.Period("2019Q4"), pd.Period("1984Q1")
excl = []
def step(d, mask, name, rationale):
    nb = len(d); d2 = d[mask]
    excl.append(dict(step=len(excl) + 1, rule=name, n_before=nb, n_after=len(d2), n_removed=nb - len(d2), rationale=rationale))
    return d2
df = step(df, df.index >= LAG0, "drop quarters before 1959Q4", "Not needed; 1959Q4 kept only as lag level for 1960Q1 differences")
df = step(df, df.index <= END, "drop quarters after 2019Q4", "Information rule: post-2019Q4 quarters outside the question")
assert (df["count"] == 3).all() and df["mean"].notna().all() and df["GDPC1"].notna().all(), "incomplete data in window"
df = df.assign(u=df["mean"], du=df["mean"].diff(), g=100 * np.log(df["GDPC1"]).diff())
df = step(df, df.index >= START, "drop 1959Q4 (lag-only row)", "First differences undefined for the 1959Q4 row in-sample")
pd.DataFrame(excl).to_csv(f"{OUT}/exclusions.csv", index=False)
assert len(df) == 240 and df[["du", "g"]].notna().all().all()

def nw_lags(n):  # Newey-West (1994) rule of thumb
    return int(np.floor(4 * (n / 100) ** (2 / 9)))

def fit(d):
    X = sm.add_constant(d["g"].values)
    L = nw_lags(len(d))
    r = sm.OLS(d["du"].values, X).fit(cov_type="HAC", cov_kwds={"maxlags": L, "kernel": "bartlett", "use_correction": True})
    return r, L

full, L = fit(df)
s1, s2 = df[df.index < SPLIT], df[df.index >= SPLIT]
r1, L1 = fit(s1); r2, L2 = fit(s2)

n = len(df); k = 2
ssr_p, ssr_1, ssr_2 = full.ssr, r1.ssr, r2.ssr
F = ((ssr_p - (ssr_1 + ssr_2)) / k) / ((ssr_1 + ssr_2) / (n - 2 * k))
p = float(stats.f.sf(F, k, n - 2 * k))

# supplementary robustness: HAC Wald test of break in (alpha, beta) and in beta alone
D = (df.index >= SPLIT).astype(float)
Xb = np.column_stack([np.ones(n), df.g.values, D, D * df.g.values])
rb = sm.OLS(df.du.values, Xb).fit(cov_type="HAC", cov_kwds={"maxlags": L, "kernel": "bartlett", "use_correction": True})
w_joint = rb.wald_test(np.array([[0, 0, 1, 0], [0, 0, 0, 1]]), scalar=True, use_f=False)
w_slope = rb.wald_test(np.array([[0, 0, 0, 1]]), scalar=True, use_f=False)

beta, se = float(full.params[1]), float(full.bse[1])
z = stats.norm.ppf(0.975)
lo, hi = beta - z * se, beta + z * se
concl = "positive_effect" if lo > 0 else ("negative_effect" if hi < 0 else "no_detectable_effect")
baseline = float(np.sqrt(np.mean((df.du - df.du.mean()) ** 2)))
model = float(np.sqrt(np.mean(full.resid ** 2)))

tab = pd.DataFrame([
    dict(period="full 1960Q1-2019Q4", n=n, intercept=full.params[0], slope=full.params[1], slope_se=full.bse[1]),
    dict(period="sub1 1960Q1-1983Q4", n=len(s1), intercept=r1.params[0], slope=r1.params[1], slope_se=r1.bse[1]),
    dict(period="sub2 1984Q1-2019Q4", n=len(s2), intercept=r2.params[0], slope=r2.params[1], slope_se=r2.bse[1]),
])
tab.to_csv(f"{OUT}/table1.csv", index=False, float_format="%.6f")

fig, ax = plt.subplots(figsize=(7, 5))
ax.scatter(s1.g, s1.du, s=14, alpha=.6, label="1960Q1-1983Q4", color="tab:blue")
ax.scatter(s2.g, s2.du, s=14, alpha=.6, label="1984Q1-2019Q4", color="tab:orange", marker="s")
for d, r, c in [(s1, r1, "tab:blue"), (s2, r2, "tab:orange")]:
    xs = np.linspace(d.g.min(), d.g.max(), 50)
    ax.plot(xs, r.params[0] + r.params[1] * xs, color=c, lw=2, label=f"fit: du = {r.params[0]:.3f} + {r.params[1]:.3f} g")
ax.axhline(0, color="grey", lw=.5); ax.axvline(0, color="grey", lw=.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); plt.close(fig)

res = dict(
    task_id="F6-conf-03", feasibility="complete", feasibility_reason=None, data_issues=data_issues,
    n_input_rows=n_input_rows, n_analysed=int(full.nobs), random_seed=None,
    sample_start=str(df.index[0].start_time.date()), sample_end=str(df.index[-1].start_time.date()),
    n_sub1=int(len(s1)), n_sub2=int(len(s2)),
    baseline_metric=baseline, model_metric=model,
    effect_estimate=beta, intercept=float(full.params[0]), effect_se=se, hac_lags=int(L),
    effect_ci_low=float(lo), effect_ci_high=float(hi),
    beta_sub1=float(r1.params[1]), beta_sub2=float(r2.params[1]),
    chow_f_stat=float(F), chow_p_value=p,
    stability_conclusion="stable" if p >= 0.05 else "unstable", conclusion=concl,
    # supplementary (not required)
    ols_se_beta_iid=float(sm.OLS(df.du.values, sm.add_constant(df.g.values)).fit().bse[1]),
    hac_lags_sub1=int(L1), hac_lags_sub2=int(L2), r_squared=float(full.rsquared),
    ssr_pooled=float(ssr_p), ssr_sub1=float(ssr_1), ssr_sub2=float(ssr_2),
    intercept_sub1=float(r1.params[0]), intercept_sub2=float(r2.params[0]),
    hac_wald_break_joint_chi2=float(w_joint.statistic), hac_wald_break_joint_p=float(w_joint.pvalue),
    hac_wald_break_slope_chi2=float(w_slope.statistic), hac_wald_break_slope_p=float(w_slope.pvalue),
    sigma_resid_sub1=float(np.sqrt(ssr_1 / len(s1))), sigma_resid_sub2=float(np.sqrt(ssr_2 / len(s2))),
    se_beta_sub1=float(r1.bse[1]), se_beta_sub2=float(r2.bse[1]),
)
with open(f"{OUT}/results.json", "w") as f:
    json.dump(res, f, indent=2)
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))