"""F6-conf-03: Okun's law (difference form), US 1960Q1-2019Q4.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Reads only data/GDPC1.csv and data/UNRATE.csv. Deterministic (no randomness used)."""
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)
START, END, BREAK = pd.Period("1960Q1"), pd.Period("2019Q4"), pd.Period("1984Q1")
excl = []

gdp = pd.read_csv("data/GDPC1.csv", na_values=[".", ""], parse_dates=["observation_date"])
ur = pd.read_csv("data/UNRATE.csv", na_values=[".", ""], parse_dates=["observation_date"])
n_input = {"GDPC1.csv": int(len(gdp)), "UNRATE.csv": int(len(ur))}
issues = []
miss = ur[ur.UNRATE.isna()].observation_date.dt.strftime("%Y-%m-%d").tolist()
if miss:
    issues.append({"issue": f"UNRATE missing value(s) at {miss}",
                   "handling": "outside the 1959Q4-2019Q4 window; that quarter is incomplete and is excluded by the date truncation; no effect on estimates"})
if gdp.GDPC1.isna().any():
    issues.append({"issue": "GDPC1 missing values", "handling": "outside window; dropped by truncation"})
issues.append({"issue": "Both files extend beyond 2019Q4 (to 2026)",
               "handling": "quarters after 2019Q4 removed before any estimation (information rule)"})
issues.append({"issue": "GDPC1 is the 2026 vintage (revised data, chained 2017 dollars), not real-time data",
               "handling": "used as is; the question is historical; noted as a limitation"})

# Quarterly unemployment = mean of the 3 monthly values; keep only complete quarters
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([g, uq], axis=1, join="inner").sort_index()
nm = len(df)
df = df[df.index <= END]                      # information rule: nothing after 2019Q4
excl.append(dict(step=1, rule="drop quarters after 2019Q4 from merged quarterly panel (GDPC1 x complete-quarter UNRATE averages)",
                 n_before=nm, n_after=len(df), n_removed=nm - len(df),
                 rationale="information rule: quarters after 2019Q4 must not enter any estimate"))
ords = np.array([p.ordinal for p in df.index])
assert (np.diff(ords) == 1).all(), "non-consecutive quarters"
df["du"] = df.u.diff()
df["g"] = 100 * np.log(df.GDPC1).diff()

n0 = len(df)
d = df[(df.index >= START)].dropna()
excl.append(dict(step=2, rule="keep merged quarters 1960Q1-2019Q4 with non-missing du and g",
                 n_before=n0, n_after=len(d), n_removed=n0 - len(d),
                 rationale="regression sample per TASK; 1947Q1-1959Q4 quarters (GDP/UNRATE overlap begins 1948Q1) only supply lagged levels"))
assert len(d) == 240 and d.du.notna().all() and d.g.notna().all()
pd.DataFrame(excl).to_csv(f"{OUT}/exclusions.csv", index=False)

n = len(d)
HAC_LAGS = int(np.floor(4 * (n / 100) ** (2 / 9)))  # Newey-West (1994) rule of thumb -> 4 for n=240

def fit(sub):
    X = sm.add_constant(sub.g.values); y = sub.du.values
    ols = sm.OLS(y, X).fit()
    rob = ols.get_robustcov_results(cov_type="HAC", maxlags=HAC_LAGS, kernel="bartlett", use_correction=False)
    return ols, rob

full_ols, full_hac = fit(d)
s1 = d[d.index < BREAK]; s2 = d[d.index >= BREAK]
o1, h1 = fit(s1); o2, h2 = fit(s2)

alpha, beta = full_ols.params
se = float(full_hac.bse[1]); z = stats.norm.ppf(0.975)
lo, hi = beta - z * se, beta + z * se
baseline = float(np.sqrt(np.mean((d.du - d.du.mean()) ** 2)))
model_rmse = float(np.sqrt(np.mean(full_ols.resid ** 2)))

k = 2
ssr_p, ssr1, ssr2 = full_ols.ssr, o1.ssr, o2.ssr
F = ((ssr_p - (ssr1 + ssr2)) / k) / ((ssr1 + ssr2) / (n - 2 * k))
p = float(stats.f.sf(F, k, n - 2 * k))

sens = {}
for L in [0, 2, 4, 6, 8, 12]:
    r = full_ols.get_robustcov_results(cov_type="HAC", maxlags=L, kernel="bartlett", use_correction=False)
    sens[str(L)] = float(r.bse[1])

concl = "positive_effect" if lo > 0 else ("negative_effect" if hi < 0 else "no_detectable_effect")

res = {
    "task_id": "F6-conf-03",
    "sample_start": str(d.index[0].start_time.date()), "sample_end": str(d.index[-1].start_time.date()),
    "n_sub1": int(len(s1)), "n_sub2": int(len(s2)),
    "baseline_metric": baseline, "model_metric": model_rmse,
    "effect_estimate": float(beta), "intercept": float(alpha),
    "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_input,
    "n_analysed": int(n), "random_seed": None,
    # supplementary keys
    "r_squared": float(full_ols.rsquared), "ols_classical_se_beta": float(full_ols.bse[1]),
    "effect_se_lag_sensitivity": sens,
    "intercept_sub1": float(o1.params[0]), "intercept_sub2": float(o2.params[0]),
    "beta_sub1_hac_se": float(h1.bse[1]), "beta_sub2_hac_se": float(h2.bse[1]),
    "chow_df": [k, int(n - 2 * k)], "mean_growth": float(d.g.mean()), "mean_du": float(d.du.mean()),
    "resid_rmse_sub1": float(np.sqrt(np.mean(o1.resid ** 2))), "resid_rmse_sub2": float(np.sqrt(np.mean(o2.resid ** 2))),
    "n_merged_quarters": int(nm),
    "implied_breakeven_growth_full": float(-alpha / beta),
}
with open(f"{OUT}/results.json", "w") as f:
    json.dump(res, f, indent=2)

pd.DataFrame([
    dict(period="full_1960Q1_2019Q4", n=n, intercept=alpha, slope=beta, slope_se=se),
    dict(period="sub1_1960Q1_1983Q4", n=len(s1), intercept=o1.params[0], slope=o1.params[1], slope_se=h1.bse[1]),
    dict(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)

fig, ax = plt.subplots(figsize=(7, 5))
for s, o, c, lab in [(s1, o1, "tab:blue", "1960Q1–1983Q4"), (s2, o2, "tab:orange", "1984Q1–2019Q4")]:
    ax.scatter(s.g, s.du, s=14, alpha=.6, color=c, label=lab)
    xs = np.linspace(s.g.min(), s.g.max(), 50)
    ax.plot(xs, o.params[0] + o.params[1] * xs, color=c, lw=2,
            label=f"OLS {lab}: Δu = {o.params[0]:.2f} {o.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 (%, quarterly, not annualised)")
ax.set_ylabel("Change in 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({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", "stability_conclusion", "conclusion"]}, indent=1))