"""Okun's law, US 1960Q1-2019Q4 (task F6-conf-03). Run from workspace root:
   ~/wp29/venv/bin/python out/analysis.py
Reads only data/, writes out/results.json, out/table1.csv, out/figure1.png, out/exclusions.csv."""
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)
gdp_raw = pd.read_csv("data/GDPC1.csv", parse_dates=["observation_date"])
un_raw = pd.read_csv("data/UNRATE.csv", parse_dates=["observation_date"])
n_input_rows = {"GDPC1.csv": int(len(gdp_raw)), "UNRATE.csv": int(len(un_raw))}
data_issues = []

# --- unemployment: quarterly mean of 3 monthly values (complete quarters only)
un = un_raw.copy()
miss = un[un["UNRATE"].isna()]
if len(miss):
    data_issues.append({"issue": f"UNRATE missing for month(s) {', '.join(miss.observation_date.dt.strftime('%Y-%m').tolist())}",
                        "handling": "Outside the 1959Q4-2019Q4 window needed for estimation; quarter would be incomplete and is not used. No imputation."})
un["q"] = un["observation_date"].dt.to_period("Q")
uq = un.groupby("q")["UNRATE"].agg(["mean", "count"])
uq.loc[un.groupby("q")["UNRATE"].apply(lambda s: s.notna().sum()) < 3, "mean"] = np.nan
gdp = gdp_raw.copy(); gdp["q"] = gdp["observation_date"].dt.to_period("Q")
gq = gdp.set_index("q")["GDPC1"]
df = pd.DataFrame({"u": uq["mean"], "gdp": gq}).sort_index()
df["du"] = df["u"].diff()
df["g"] = 100 * np.log(df["gdp"]).diff()
# require consecutive quarters for differences
idx = df.index
consec = pd.Series([True] + [(idx[i] - idx[i-1]).n == 1 for i in range(1, len(idx))], index=idx)
df.loc[~consec, ["du", "g"]] = np.nan

data_issues.append({"issue": "GDPC1 is the current (2026) vintage in chained 2017 dollars, not real-time data; historical growth rates reflect later revisions.",
                    "handling": "Used as provided (the task specifies this series); noted as a limitation."})

# --- exclusion ledger (unit: quarterly observations of the merged panel)
excl = []
n0 = len(df)
q_start, q_end = pd.Period("1960Q1", "Q"), pd.Period("2019Q4", "Q")
d1 = df[df.index >= q_start]; excl.append(("1", "drop quarters before 1960Q1", n0, len(d1), "outside question window (1959Q4 levels used only for differencing)"))
d2 = d1[d1.index <= q_end]; excl.append(("2", "drop quarters after 2019Q4", len(d1), len(d2), "information rule: post-2019Q4 must not enter any estimate"))
d3 = d2.dropna(subset=["du", "g"]); excl.append(("3", "drop quarters with missing du or g", len(d2), len(d3), "incomplete data"))
pd.DataFrame([dict(step=s, rule=r, n_before=a, n_after=b, n_removed=a-b, rationale=t) for s, r, a, b, t in excl]).to_csv(f"{OUT}/exclusions.csv", index=False)
d = d3.copy()
assert len(d) == 240 and d.index.min() == q_start and d.index.max() == q_end

# --- baseline: mean-only
baseline_rmse = float(np.sqrt(np.mean((d.du - d.du.mean())**2)))

# --- HAC lag: Newey-West (1994) rule of thumb floor(4*(T/100)^(2/9))
def nw_lag(T): return int(np.floor(4 * (T / 100) ** (2 / 9)))
L = nw_lag(len(d))

def fit(sub):
    X = sm.add_constant(sub["g"].values); y = sub["du"].values
    ols = sm.OLS(y, X).fit()
    hac = sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": L, "kernel": "bartlett", "use_correction": False})
    return ols, hac

full_ols, full_hac = fit(d)
s1 = d[d.index <= pd.Period("1983Q4", "Q")]; s2 = d[d.index >= pd.Period("1984Q1", "Q")]
o1, h1 = fit(s1); o2, h2 = fit(s2)
beta, alpha = float(full_ols.params[1]), float(full_ols.params[0])
se = float(full_hac.bse[1])
z = stats.norm.ppf(0.975)
lo, hi = beta - z * se, beta + z * se
model_rmse = float(np.sqrt(np.mean(full_ols.resid**2)))

# --- Chow test (classical SSR F, df 2, n-4)
k = 2; n = len(d)
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))

# sensitivity of HAC SE to lag length (reported, not used for main CI)
sens = {}
for lag in [0, 2, 4, 8, 12]:
    r = sm.OLS(d.du.values, sm.add_constant(d.g.values)).fit(cov_type="HAC", cov_kwds={"maxlags": lag, "kernel": "bartlett", "use_correction": False})
    sens[str(lag)] = float(r.bse[1])

# supplementary robustness of the stability test: HAC Wald with subperiod-2 dummies
D = (d.index >= pd.Period("1984Q1", "Q")).astype(float)
Xi = np.column_stack([np.ones(n), d.g.values, D, D * d.g.values])
ri = sm.OLS(d.du.values, Xi).fit(cov_type="HAC", cov_kwds={"maxlags": L, "kernel": "bartlett", "use_correction": False})
R = np.array([[0, 0, 1, 0], [0, 0, 0, 1]])
wt = ri.wald_test(R, use_f=False, scalar=True)
hac_wald_chi2, hac_wald_p = float(wt.statistic), float(wt.pvalue)
p_int_diff_hac, p_slope_diff_hac = float(ri.pvalues[2]), float(ri.pvalues[3])
resid_var_ratio = float((o1.ssr / (len(s1) - 2)) / (o2.ssr / (len(s2) - 2)))

conclusion = "positive_effect" if lo > 0 else ("negative_effect" if hi < 0 else "no_detectable_effect")
res = {
    "task_id": "F6-conf-03", "feasibility": "complete", "feasibility_reason": None,
    "data_issues": data_issues, "n_input_rows": n_input_rows, "n_analysed": int(n), "random_seed": None,
    "sample_start": str(d.index.min().start_time.date()), "sample_end": str(d.index.max().start_time.date()),
    "n_sub1": int(len(s1)), "n_sub2": int(len(s2)),
    "baseline_metric": baseline_rmse, "model_metric": model_rmse,
    "effect_estimate": beta, "intercept": alpha, "effect_se": se, "hac_lags": L,
    "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": conclusion,
    # supplementary keys
    "ols_iid_se": float(full_ols.bse[1]), "r_squared": float(full_ols.rsquared),
    "intercept_sub1": float(o1.params[0]), "intercept_sub2": float(o2.params[0]),
    "se_hac_sub1": float(h1.bse[1]), "se_hac_sub2": float(h2.bse[1]),
    "hac_se_by_lag": sens,
    "hac_wald_break_chi2": hac_wald_chi2, "hac_wald_break_p": hac_wald_p,
    "hac_p_intercept_diff": p_int_diff_hac, "hac_p_slope_diff": p_slope_diff_hac,
    "resid_var_ratio_sub1_sub2": resid_var_ratio, "chow_df": [k, n - 2 * k],
}
with open(f"{OUT}/results.json", "w") as f: json.dump(res, f, indent=2)

pd.DataFrame([
    {"period": "full 1960Q1-2019Q4", "n": n, "intercept": alpha, "slope": beta, "slope_se": se},
    {"period": "sub1 1960Q1-1983Q4", "n": len(s1), "intercept": float(o1.params[0]), "slope": float(o1.params[1]), "slope_se": float(h1.bse[1])},
    {"period": "sub2 1984Q1-2019Q4", "n": len(s2), "intercept": float(o2.params[0]), "slope": float(o2.params[1]), "slope_se": float(h2.bse[1])},
]).to_csv(f"{OUT}/table1.csv", index=False)

fig, ax = plt.subplots(figsize=(7, 5))
for sub, o, c, lab in [(s1, o1, "tab:blue", "1960Q1-1983Q4"), (s2, o2, "tab:red", "1984Q1-2019Q4")]:
    ax.scatter(sub.g, sub.du, s=14, alpha=0.6, color=c, label=f"{lab} (n={len(sub)})")
    xs = np.linspace(sub.g.min(), sub.g.max(), 50)
    ax.plot(xs, o.params[0] + o.params[1] * xs, color=c, lw=2, label=f"fit {lab}: slope={o.params[1]:.3f}")
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("Change in unemployment rate Δu_t (pp)")
ax.set_title("Okun's law, US quarterly data"); ax.legend(fontsize=8)
fig.tight_layout(); fig.savefig(f"{OUT}/figure1.png", dpi=120); plt.close(fig)

print(f"n={n} beta={beta:.4f} HAC se={se:.4f} (L={L}) CI=[{lo:.4f},{hi:.4f}] beta1={o1.params[1]:.4f} beta2={o2.params[1]:.4f} Chow F={F:.3f} p={p:.4g}")