"""Okun's law (difference form), US 1960Q1-2019Q4. Run from workspace root."""
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

np.random.seed(0)
OUT = "out"
gdp = pd.read_csv("data/GDPC1.csv", parse_dates=["observation_date"])
un = pd.read_csv("data/UNRATE.csv", parse_dates=["observation_date"])
n_rows = {"GDPC1.csv": len(gdp), "UNRATE.csv": len(un)}
issues = []

# ---- audit
for name, df, col, freq in [("GDPC1.csv", gdp, "GDPC1", "QS"), ("UNRATE.csv", un, "UNRATE", "MS")]:
    d = df["observation_date"]
    if d.duplicated().any():
        issues.append({"issue": f"{name}: duplicated dates", "handling": "would abort"}); raise SystemExit
    full = pd.date_range(d.min(), d.max(), freq=freq)
    if len(full) != len(d):
        issues.append({"issue": f"{name}: {len(full)-len(d)} missing dates", "handling": "checked against sample window"})
    miss = df[df[col].isna()]
    for dt in miss["observation_date"]:
        inside = pd.Timestamp("1959-10-01") <= dt <= pd.Timestamp("2019-12-31")
        issues.append({"issue": f"{name}: missing value at {dt.date()}",
                       "handling": ("inside sample window -- quarter dropped" if inside else
                                    "outside the 1959Q4-2019Q4 window used; no effect on estimates")})
    if (df[col].dropna() <= 0).any():
        issues.append({"issue": f"{name}: non-positive values", "handling": "inspected"})
issues.append({"issue": "Both series extend beyond 2019Q4 (to 2026) and are current vintages (revised data, GDP in chained 2017 dollars), not real-time",
               "handling": "observations after 2019Q4 excluded before any computation; revision/vintage effects acknowledged as limitation"})

# ---- build quarterly data (only through 2019Q4; 1959Q4 needed for first difference)
un = un[(un.observation_date >= "1959-10-01") & (un.observation_date <= "2019-12-01")].copy()
un["q"] = un.observation_date.dt.to_period("Q")
cnt = un.groupby("q")["UNRATE"].count()
assert (cnt == 3).all(), "incomplete quarter in window"
uq = un.groupby("q")["UNRATE"].mean()
gdp = gdp[(gdp.observation_date >= "1959-10-01") & (gdp.observation_date <= "2019-10-01")].copy()
gdp["q"] = gdp.observation_date.dt.to_period("Q")
gq = gdp.set_index("q")["GDPC1"]
df = pd.DataFrame({"u": uq, "y": gq})
df["du"] = df.u.diff()
df["g"] = 100 * np.log(df.y).diff()
df = df.loc[pd.Period("1960Q1"):pd.Period("2019Q4")].dropna()
assert len(df) == 240

def ols(d, hac=None):
    X = sm.add_constant(d["g"].values)
    m = sm.OLS(d["du"].values, X)
    return m.fit(cov_type="HAC", cov_kwds={"maxlags": hac, "kernel": "bartlett"}) if hac is not None else m.fit()

T = len(df)
L = int(np.floor(4 * (T / 100) ** (2 / 9)))   # Newey-West (1994) rule of thumb -> 4
full = ols(df, L)
s1 = df.loc[:pd.Period("1983Q4")]; s2 = df.loc[pd.Period("1984Q1"):]
r1, r2 = ols(s1, L), ols(s2, L)
# HAC lag sensitivity
sens = {lag: float(ols(df, lag).bse[1]) for lag in [0, 2, 4, 8, 12]}
print("HAC SE sensitivity by lag:", sens)
print("OLS (iid) SE:", float(ols(df).bse[1]))

# Chow
ssr_p = float(ols(df).ssr); ssr_u = float(ols(s1).ssr + ols(s2).ssr); k = 2
F = ((ssr_p - ssr_u) / k) / (ssr_u / (T - 2 * k))
p = float(stats.f.sf(F, k, T - 2 * k))

# Robustness: HAC Wald test of break in (alpha,beta) jointly and in beta only (dummy-interaction regression)
D = (df.index >= pd.Period("1984Q1")).astype(float)
Xi = np.column_stack([np.ones(T), df.g.values, D, D * df.g.values])
mi = sm.OLS(df.du.values, Xi).fit(cov_type="HAC", cov_kwds={"maxlags": L, "kernel": "bartlett"})
w_joint = mi.wald_test(np.array([[0, 0, 1, 0], [0, 0, 0, 1]]), use_f=False, scalar=True)
w_slope = mi.wald_test(np.array([[0, 0, 0, 1]]), use_f=False, scalar=True)
w_int = mi.wald_test(np.array([[0, 0, 1, 0]]), use_f=False, scalar=True)
rob = {"hac_se_by_lag": {str(k_): v for k_, v in sens.items()}, "ols_iid_se_beta": float(ols(df).bse[1]),
       "hac_wald_joint_break_chi2": float(w_joint.statistic), "hac_wald_joint_break_p": float(w_joint.pvalue),
       "hac_wald_slope_break_chi2": float(w_slope.statistic), "hac_wald_slope_break_p": float(w_slope.pvalue),
       "hac_wald_intercept_break_chi2": float(w_int.statistic), "hac_wald_intercept_break_p": float(w_int.pvalue),
       "slope_diff_sub2_minus_sub1": float(mi.params[3]), "slope_diff_hac_se": float(mi.bse[3]),
       "intercept_diff_sub2_minus_sub1": float(mi.params[2]), "intercept_diff_hac_se": float(mi.bse[2]),
       "breakeven_growth_sub1_pct_qq": float(-r1.params[0] / r1.params[1]),
       "breakeven_growth_sub2_pct_qq": float(-r2.params[0] / r2.params[1]),
       "mean_g_sub1": float(s1.g.mean()), "mean_g_sub2": float(s2.g.mean()),
       "durbin_watson_full_resid": float(sm.stats.stattools.durbin_watson(ols(df).resid))}
print("Robustness:", json.dumps(rob, indent=2))

b, se = float(full.params[1]), float(full.bse[1])
z = stats.norm.ppf(0.975)
lo, hi = b - z * se, b + 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(full.resid ** 2))),
    "effect_estimate": b, "intercept": float(full.params[0]),
    "effect_se": se, "hac_lags": L,
    "effect_ci_low": lo, "effect_ci_high": 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,
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": issues, "n_input_rows": n_rows,
    "n_analysed": int(T), "random_seed": 0,
}
os.makedirs(OUT, exist_ok=True)
json.dump(res, open(f"{OUT}/results.json", "w"), indent=2)
json.dump(rob, open(f"{OUT}/robustness.json", "w"), indent=2)

pd.DataFrame([
    {"period": "full_1960Q1-2019Q4", "n": T, "intercept": full.params[0], "slope": full.params[1], "slope_se": full.bse[1]},
    {"period": "sub1_1960Q1-1983Q4", "n": len(s1), "intercept": r1.params[0], "slope": r1.params[1], "slope_se": r1.bse[1]},
    {"period": "sub2_1984Q1-2019Q4", "n": len(s2), "intercept": r2.params[0], "slope": r2.params[1], "slope_se": r2.bse[1]},
]).to_csv(f"{OUT}/table1.csv", index=False)

fig, ax = plt.subplots(figsize=(7, 5))
ax.scatter(s1.g, s1.du, s=12, alpha=.6, label="1960Q1–1983Q4")
ax.scatter(s2.g, s2.du, s=12, alpha=.6, label="1984Q1–2019Q4")
xs = np.linspace(df.g.min(), df.g.max(), 50)
ax.plot(xs, r1.params[0] + r1.params[1] * xs, label=f"fit 1: β={r1.params[1]:.3f}")
ax.plot(xs, r2.params[0] + r2.params[1] * xs, label=f"fit 2: β={r2.params[1]:.3f}")
ax.axhline(0, c="grey", lw=.5); ax.axvline(0, c="grey", lw=.5)
ax.set_xlabel("Real GDP growth g_t (% q/q, not annualised)")
ax.set_ylabel("Δ unemployment rate Δu_t (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))