"""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

OUT = "out"; os.makedirs(OUT, exist_ok=True)
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")]:
    assert not d.observation_date.duplicated().any(), name + " duplicate dates"
    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-31"]
        issues.append({"issue": f"{name}: missing value(s) at {miss}",
                       "handling": "outside 1959Q4-2019Q4 estimation window; no effect on estimates" if not inside
                       else "inside window: quarter dropped"})
    if (d[col].dropna() <= 0).any():
        issues.append({"issue": f"{name}: non-positive values", "handling": "flagged"})
# monthly continuity
months = pd.date_range(ur.observation_date.min(), ur.observation_date.max(), freq="MS")
assert len(months) == len(ur), "UNRATE month gaps"
qs = pd.date_range(gdp.observation_date.min(), gdp.observation_date.max(), freq="QS")
assert len(qs) == len(gdp), "GDPC1 quarter gaps"

# --- construct quarterly series ---
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")
g = gdp.set_index(gdp.observation_date.dt.to_period("Q"))["GDPC1"]
df = pd.concat([uq, g], axis=1).sort_index()
df["du"] = df["u"].diff()
df["g"] = 100 * np.log(df["GDPC1"]).diff()
df = df.loc[pd.Period("1960Q1"):pd.Period("2019Q4")]
assert len(df) == 240 and df[["du", "g"]].notna().all().all()
issues.append({"issue": "Observations after 2019Q4 present in both files (incl. 2020 COVID outliers); GDPC1 is the 2026 vintage (revised data, chained 2017 $)",
               "handling": "post-2019Q4 data excluded per task; current vintage used as provided (log growth rates invariant to base year)"})
# outlier screen on in-sample data
z = (df[["du", "g"]] - df[["du", "g"]].mean()) / df[["du", "g"]].std()
big = df.index[(z.abs() > 4).any(axis=1)].astype(str).tolist()
if big:
    issues.append({"issue": f"Large (|z|>4) in-sample observations: {big}",
                   "handling": "retained (genuine recession/recovery quarters); influence discussed in report"})

def ols(d, hac=None):
    X = sm.add_constant(d["g"])
    m = sm.OLS(d["du"], 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
full = ols(df, L); full_ols = ols(df)
s1 = df.loc[:pd.Period("1983Q4")]; s2 = df.loc[pd.Period("1984Q1"):]
r1, r2 = ols(s1, L), ols(s2, L)
ssr_p = full_ols.ssr; ssr_u = 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))

beta = float(full.params["g"]); se = float(full.bse["g"])
z975 = stats.norm.ppf(0.975)
lo, hi = beta - z975 * se, beta + z975 * se
concl = "positive_effect" if lo > 0 else ("negative_effect" if hi < 0 else "no_detectable_effect")

# sensitivity (reported in report, not in results keys)
sens = {f"lags_{l}": float(ols(df, l).bse["g"]) for l in (0, 2, 4, 8, 12)}
sens["ols_iid_se"] = float(full_ols.bse["g"])
sens["beta_excl_z4"] = float(ols(df.drop(pd.PeriodIndex(big, freq="Q"))).params["g"]) if big else beta
# HAC-robust Wald test of break (interaction model)
d2 = df.copy(); d2["D"] = (d2.index >= pd.Period("1984Q1")).astype(float); d2["Dg"] = d2["D"] * d2["g"]
mw = sm.OLS(d2["du"], sm.add_constant(d2[["g", "D", "Dg"]])).fit(cov_type="HAC", cov_kwds={"maxlags": L})
wt = mw.wald_test("D = 0, Dg = 0", use_f=True, scalar=True)
sens["hac_wald_break_F"] = float(wt.statistic); sens["hac_wald_break_p"] = float(wt.pvalue)
sens["hac_slope_diff"] = float(mw.params["Dg"]); sens["hac_slope_diff_p"] = float(mw.pvalues["Dg"])
sens["hac_intercept_diff"] = float(mw.params["D"]); sens["hac_intercept_diff_p"] = float(mw.pvalues["D"])
sens["breakeven_sub1"] = float(-r1.params["const"] / r1.params["g"]); sens["breakeven_sub2"] = float(-r2.params["const"] / r2.params["g"])
sens["sub1_hac_se"] = float(r1.bse["g"]); sens["sub2_hac_se"] = float(r2.bse["g"])
sens["sub1_alpha"] = float(r1.params["const"]); sens["sub2_alpha"] = float(r2.params["const"])
sens["r2_full"] = float(full_ols.rsquared)
sens["breakeven_growth_full"] = float(-full.params["const"] / beta)
json.dump(sens, open(f"{OUT}/sensitivity.json", "w"), indent=2)

res = {
    "task_id": "F6-conf-03",
    "sample_start": df.index[0].start_time.strftime("%Y-%m-%d"),
    "sample_end": df.index[-1].start_time.strftime("%Y-%m-%d"),
    "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_ols.resid ** 2))),
    "effect_estimate": beta, "intercept": float(full.params["const"]),
    "effect_se": se, "hac_lags": L, "effect_ci_low": float(lo), "effect_ci_high": float(hi),
    "beta_sub1": float(r1.params["g"]), "beta_sub2": float(r2.params["g"]),
    "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": None,
}
json.dump(res, open(f"{OUT}/results.json", "w"), indent=2)

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

fig, ax = plt.subplots(figsize=(7, 5))
for d, r, c, lab in [(s1, r1, "tab:blue", "1960Q1–1983Q4"), (s2, r2, "tab:red", "1984Q1–2019Q4")]:
    ax.scatter(d.g, d.du, s=12, alpha=0.6, color=c, label=f"{lab} (n={len(d)})")
    xs = np.linspace(d.g.min(), d.g.max(), 50)
    ax.plot(xs, r.params["const"] + r.params["g"] * xs, color=c,
            label=f"fit {lab}: Δu = {r.params['const']:.2f} {r.params['g']:+.3f}·g")
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("Δ unemployment rate Δu_t (pp)")
ax.set_title("Okun's law, US 1960Q1–2019Q4"); ax.legend(fontsize=8)
fig.tight_layout(); fig.savefig(f"{OUT}/figure1.png", dpi=150)
print(json.dumps(res, indent=2)); print(json.dumps(sens, indent=2))