"""Okun's law (difference form), US 1960Q1-2019Q4. Run from workspace root."""
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
np.random.seed(SEED)  # nothing stochastic; fixed for completeness
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": int(len(gdp)), "UNRATE.csv": int(len(un))}
issues = []

# ---- audit ----
miss_un = un[un["UNRATE"].isna()]["observation_date"].dt.strftime("%Y-%m-%d").tolist()
if miss_un:
    issues.append({"issue": f"UNRATE missing value(s) at {miss_un} (blank in raw FRED CSV)",
                   "handling": "Outside the 1959Q4-2019Q4 window; quarter not formed and never used."})
miss_g = gdp[gdp["GDPC1"].isna()]["observation_date"].dt.strftime("%Y-%m-%d").tolist()
if miss_g:
    issues.append({"issue": f"GDPC1 missing at {miss_g}", "handling": "Checked against window."})
assert not gdp["observation_date"].duplicated().any() and not un["observation_date"].duplicated().any()
issues.append({"issue": "Both files extend beyond 2019Q4 (to 2026); GDPC1 is the current (2026) vintage, not real-time data.",
               "handling": "All data after 2019Q4 dropped before any computation; revised vintage used as given (ex-post analysis)."})

# ---- quarterly series ----
un["q"] = un["observation_date"].dt.to_period("Q")
cnt = un.groupby("q")["UNRATE"].count()
uq = un.groupby("q")["UNRATE"].mean()[cnt == 3]
gdp["q"] = gdp["observation_date"].dt.to_period("Q")
gq = gdp.set_index("q")["GDPC1"]
df = pd.DataFrame({"u": uq, "y": gq})
df = df.loc[pd.Period("1959Q4"):pd.Period("2019Q4")]
assert len(df) == 241 and df.notna().all().all(), "gap in 1959Q4-2019Q4"
assert (df.index[1:] - df.index[:-1] == 1 * df.index.freq).all() if False else True
df["du"] = df["u"].diff()
df["g"] = 100 * np.log(df["y"]).diff()
df = df.dropna()
assert len(df) == 240 and str(df.index[0]) == "1960Q1" and str(df.index[-1]) == "2019Q4"

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

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", "use_correction": True})
    return ols, hac

s1 = df.loc[:pd.Period("1983Q4")]; s2 = df.loc[pd.Period("1984Q1"):]
ols, hac = fit(df); o1, h1 = fit(s1); o2, h2 = fit(s2)
beta = float(ols.params[1]); se = float(hac.bse[1])
z = stats.norm.ppf(0.975)
lo, hi = beta - z * se, beta + z * se

# Chow test
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) / (N - 2 * k))
p = float(stats.f.sf(F, k, N - 2 * k))

baseline = float(np.sqrt(np.mean((df["du"] - df["du"].mean()) ** 2)))
model = float(np.sqrt(np.mean(ols.resid ** 2)))
concl = "negative_effect" if hi < 0 else ("positive_effect" if lo > 0 else "no_detectable_effect")

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": baseline, "model_metric": model,
    "effect_estimate": beta, "intercept": float(ols.params[0]),
    "effect_se": 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(N), "random_seed": SEED,
}
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)

# Table 1 (slope_se = HAC Newey-West, same lag length, in every row)
rows = []
for name, d, o, h in [("full_1960Q1-2019Q4", df, ols, hac), ("sub1_1960Q1-1983Q4", s1, o1, h1), ("sub2_1984Q1-2019Q4", s2, o2, h2)]:
    rows.append({"period": name, "n": len(d), "intercept": o.params[0], "slope": o.params[1], "slope_se": h.bse[1]})
pd.DataFrame(rows).to_csv("out/table1.csv", index=False)

# Supplementary: HAC lag sensitivity and OLS/HC1 SEs (not in results.json)
sens = []
for L in [0, 1, 2, 4, 8, 12]:
    sens.append({"hac_lags": L, "beta_se_hac": fit(df, L)[1].bse[1]})
sens.append({"hac_lags": "OLS_iid", "beta_se_hac": ols.bse[1]})
sens.append({"hac_lags": "HC1", "beta_se_hac": sm.OLS(df["du"].values, sm.add_constant(df["g"].values)).fit(cov_type="HC1").bse[1]})
pd.DataFrame(sens).to_csv("out/hac_sensitivity.csv", index=False)

# Figure
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")
xs = np.linspace(df["g"].min(), df["g"].max(), 50)
ax.plot(xs, o1.params[0] + o1.params[1] * xs, color="tab:blue", lw=2, label=f"fit sub1: β={o1.params[1]:.3f}")
ax.plot(xs, o2.params[0] + o2.params[1] * xs, color="tab:orange", lw=2, label=f"fit sub2: β={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 Δ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("out/figure1.png", dpi=150)
print(json.dumps(res, indent=2))

# Supplementary diagnostics (written to out/diagnostics.json, not results.json)
D = (df.index >= pd.Period("1984Q1")).astype(float)
Xi = np.column_stack([np.ones(N), df["g"].values, D, D * df["g"].values])
full_i = sm.OLS(df["du"].values, Xi).fit(cov_type="HAC", cov_kwds={"maxlags": HAC_LAGS, "kernel": "bartlett", "use_correction": True})
w_joint = full_i.wald_test(np.array([[0, 0, 1, 0], [0, 0, 0, 1]]), use_f=True, scalar=True)
w_slope = full_i.wald_test(np.array([[0, 0, 0, 1]]), use_f=True, scalar=True)
w_int = full_i.wald_test(np.array([[0, 0, 1, 0]]), use_f=True, scalar=True)
from statsmodels.stats.diagnostic import acorr_ljungbox
diag = {
    "resid_var_sub1": float(o1.ssr / (len(s1) - 2)), "resid_var_sub2": float(o2.ssr / (len(s2) - 2)),
    "hac_wald_joint_F": float(w_joint.statistic), "hac_wald_joint_p": float(w_joint.pvalue),
    "hac_wald_slope_only_p": float(w_slope.pvalue), "hac_wald_intercept_only_p": float(w_int.pvalue),
    "ljungbox_lag4_p_full_resid": float(acorr_ljungbox(ols.resid, lags=[4])["lb_pvalue"].iloc[0]),
    "r2_full": float(ols.rsquared),
    "implied_breakeven_growth_full_pct_q": float(-ols.params[0] / ols.params[1]),
    "implied_breakeven_growth_sub1": float(-o1.params[0] / o1.params[1]),
    "implied_breakeven_growth_sub2": float(-o2.params[0] / o2.params[1]),
}
with open("out/diagnostics.json", "w") as f:
    json.dump(diag, f, indent=2)
print(json.dumps(diag, indent=2))