"""Okun's law (difference form), US 1960Q1-2019Q4. Run from workspace root:
   ~/wp29/venv/bin/python out/analysis.py"""
import json, math
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

SEED = 0  # no randomness used; fixed for completeness
np.random.seed(SEED)
OUT = "out"
issues = []

gdp = pd.read_csv("data/GDPC1.csv", parse_dates=["observation_date"])
ur = pd.read_csv("data/UNRATE.csv", parse_dates=["observation_date"])
n_rows = {"GDPC1.csv": int(len(gdp)), "UNRATE.csv": int(len(ur))}

# ---- audit ----
for name, d, col in [("GDPC1.csv", gdp, "GDPC1"), ("UNRATE.csv", ur, "UNRATE")]:
    if d["observation_date"].duplicated().any():
        raise SystemExit(f"duplicate dates in {name}")
    miss = d.loc[d[col].isna(), "observation_date"].dt.strftime("%Y-%m-%d").tolist()
    if miss:
        inside = [m for m in miss if "1959-10-01" <= m <= "2019-12-01"]
        issues.append({"issue": f"{name}: missing value(s) at {miss}",
                       "handling": "outside the 1959Q4-2019Q4 window needed; not used" if not inside
                       else "inside window -- quarter dropped"})
# monthly -> quarterly average, require 3 non-missing months
ur["q"] = ur["observation_date"].dt.to_period("Q")
uq = ur.groupby("q")["UNRATE"].agg(["mean", "count"])
uq = uq[uq["count"] == 3]["mean"].rename("u")
gdp["q"] = gdp["observation_date"].dt.to_period("Q")
gq = gdp.set_index("q")["GDPC1"]
df = pd.concat([uq, gq], axis=1).sort_index()
df["du"] = df["u"].diff()
df["g"] = 100 * np.log(df["GDPC1"]).diff()
# ensure consecutive quarters for differencing
df = df.loc[pd.Period("1960Q1"):pd.Period("2019Q4")]
assert len(df) == 240 and df[["du", "g"]].notna().all().all(), "sample incomplete"
assert (np.diff(df.index.asi8) == 1).all()
issues.append({"issue": "Post-2019 observations (incl. COVID-19 2020 outliers) present in files",
               "handling": "excluded per task; no estimate uses quarters after 2019Q4"})
issues.append({"issue": "GDPC1 is the current (2026) vintage, revised long after the fact; not real-time data",
               "handling": "used as is (task asks for a historical relationship); noted as limitation"})

T = len(df)
HAC_LAGS = int(math.floor(4 * (T / 100) ** (2 / 9)))  # Newey-West (1994) rule of thumb

def fit(d, lags=HAC_LAGS):
    X = sm.add_constant(d["g"].values)
    ols = sm.OLS(d["du"].values, X).fit()
    hac = sm.OLS(d["du"].values, X).fit(cov_type="HAC", cov_kwds={"maxlags": lags, "kernel": "bartlett"})
    return ols, hac

ols, hac = fit(df)
s1 = df.loc[:pd.Period("1983Q4")]; s2 = df.loc[pd.Period("1984Q1"):]
o1, h1 = fit(s1); o2, h2 = fit(s2)

beta, alpha = ols.params[1], ols.params[0]
se = hac.bse[1]
z = stats.norm.ppf(0.975)
ssr_p, ssr_1, ssr_2 = ols.ssr, o1.ssr, o2.ssr
k = 2
F = ((ssr_p - (ssr_1 + ssr_2)) / k) / ((ssr_1 + ssr_2) / (T - 2 * k))
p = float(stats.f.sf(F, k, T - 2 * k))
lo, hi = beta - z * se, beta + z * se
concl = "positive_effect" if lo > 0 else ("negative_effect" if hi < 0 else "no_detectable_effect")

res = {
    "task_id": "F6-conf-03",
    "sample_start": "1960-01-01", "sample_end": "2019-10-01",
    "n_sub1": int(len(s1)), "n_sub2": int(len(s2)),
    "baseline_metric": float(np.sqrt(np.mean((df.du - df.du.mean()) ** 2))),
    "model_metric": float(np.sqrt(np.mean(ols.resid ** 2))),
    "effect_estimate": float(beta), "intercept": float(alpha),
    "effect_se": float(se), "hac_lags": HAC_LAGS,
    "effect_ci_low": float(lo), "effect_ci_high": float(hi),
    "beta_sub1": float(o1.params[1]), "beta_sub2": float(o2.params[1]),
    "chow_f_stat": float(F), "chow_p_value": p,
    "stability_conclusion": "stable" if p >= 0.05 else "unstable",
    "conclusion": concl,
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": issues, "n_input_rows": n_rows,
    "n_analysed": int(T), "random_seed": SEED,
}
with open(f"{OUT}/results.json", "w") as f:
    json.dump(res, f, indent=2)

# table: slope_se is HAC (Bartlett, same lag length) for every row
pd.DataFrame([
    {"period": "full_1960Q1_2019Q4", "n": T, "intercept": alpha, "slope": beta, "slope_se": se},
    {"period": "sub1_1960Q1_1983Q4", "n": len(s1), "intercept": o1.params[0], "slope": o1.params[1], "slope_se": h1.bse[1]},
    {"period": "sub2_1984Q1_2019Q4", "n": len(s2), "intercept": o2.params[0], "slope": o2.params[1], "slope_se": h2.bse[1]},
]).to_csv(f"{OUT}/table1.csv", index=False)

# supplementary diagnostics (not in results.json)
rows = []
for L in [0, 2, 4, 8, 12]:
    rows.append({"hac_lags": L, "beta_se": fit(df, L)[1].bse[1]})
pd.DataFrame(rows).to_csv(f"{OUT}/hac_lag_sensitivity.csv", index=False)
# Wald test of beta equality with HAC (interaction model), robust to heteroskedasticity
D = (df.index >= pd.Period("1984Q1")).astype(float)
Xi = np.column_stack([np.ones(T), df.g, D, D * df.g])
wi = sm.OLS(df.du.values, Xi).fit(cov_type="HAC", cov_kwds={"maxlags": HAC_LAGS})
wt_joint = wi.wald_test(np.array([[0, 0, 1, 0], [0, 0, 0, 1]]), scalar=True)
wt_slope = wi.wald_test(np.array([[0, 0, 0, 1]]), scalar=True)
lb = sm.stats.acorr_ljungbox(ols.resid, lags=[4, 8])
bp = sm.stats.het_breuschpagan(ols.resid, sm.add_constant(df.g.values))
with open(f"{OUT}/diagnostics.json", "w") as f:
    json.dump({"hac_wald_joint_stat": float(wt_joint.statistic), "hac_wald_joint_p": float(wt_joint.pvalue),
               "hac_wald_slope_only_stat": float(wt_slope.statistic), "hac_wald_slope_only_p": float(wt_slope.pvalue),
               "ljungbox_p_lag4": float(lb.lb_pvalue.iloc[0]), "ljungbox_p_lag8": float(lb.lb_pvalue.iloc[1]),
               "breusch_pagan_p": float(bp[1]),
               "r2_full": float(ols.rsquared), "alpha_sub1": float(o1.params[0]), "alpha_sub2": float(o2.params[0]),
               "implied_zero_du_growth_full": float(-alpha / beta), "implied_zero_du_growth_sub1": float(-o1.params[0] / o1.params[1]), "implied_zero_du_growth_sub2": float(-o2.params[0] / o2.params[1])}, f, indent=2)

# figure
fig, ax = plt.subplots(figsize=(7, 5))
ax.scatter(s1.g, s1.du, s=14, alpha=.6, label="1960Q1–1983Q4", color="C0")
ax.scatter(s2.g, s2.du, s=14, alpha=.6, label="1984Q1–2019Q4", color="C1", marker="^")
xs = np.linspace(df.g.min(), df.g.max(), 50)
ax.plot(xs, o1.params[0] + o1.params[1] * xs, color="C0", lw=2, label=f"fit 1: β={o1.params[1]:.3f}")
ax.plot(xs, o2.params[0] + o2.params[1] * xs, color="C1", lw=2, label=f"fit 2: β={o2.params[1]:.3f}")
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("Δ unemployment rate (pp)")
ax.set_title(f"Okun's law, US 1960–2019 (Chow F={F:.2f}, p={p:.3f})")
ax.legend(fontsize=8); fig.tight_layout(); fig.savefig(f"{OUT}/figure1.png", dpi=150)
print(json.dumps(res, indent=2))