"""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.
No randomness is used (random_seed = null)."""
import json, os
import numpy as np
import 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)
S0, S1 = pd.Period("1960Q1", "Q"), pd.Period("2019Q4", "Q")
SPLIT = pd.Period("1984Q1", "Q")

gdp_raw = pd.read_csv("data/GDPC1.csv")
un_raw = pd.read_csv("data/UNRATE.csv")
n_input_rows = {"GDPC1.csv": int(len(gdp_raw)), "UNRATE.csv": int(len(un_raw))}
data_issues = []

# --- unemployment: quarterly average of monthly values (require all 3 months) ---
un = un_raw.copy()
un["date"] = pd.to_datetime(un["observation_date"])
un["UNRATE"] = pd.to_numeric(un["UNRATE"], errors="coerce")
nan_months = un.loc[un["UNRATE"].isna(), "observation_date"].tolist()
if nan_months:
    data_issues.append({"issue": f"UNRATE missing value(s) for month(s) {nan_months}",
                        "handling": "Quarter containing a missing month is not averaged (set missing); "
                                    "all such months lie outside the 1959Q4-2019Q4 window, so no effect on estimates."})
un["q"] = un["date"].dt.to_period("Q")
uq = un.groupby("q")["UNRATE"].agg(["mean", "count"])
uq.loc[uq["count"] < 3, "mean"] = np.nan  # incomplete quarters
u = uq["mean"].rename("u")

gdp = gdp_raw.copy()
gdp["q"] = pd.to_datetime(gdp["observation_date"]).dt.to_period("Q")
assert not gdp["q"].duplicated().any()
g_level = gdp.set_index("q")["GDPC1"].astype(float)

# --- merge and exclusion ledger (unit: quarters of GDPC1) ---
ledger = []
def log(step, rule, before, after, why):
    ledger.append(dict(step=step, rule=rule, n_before=before, n_after=after,
                       n_removed=before - after, rationale=why))

df = pd.DataFrame({"gdp": g_level})
n0 = len(df)
df = df.join(u, how="left")
n1 = int(df["u"].notna().sum())
df_ok = df[df["u"].notna()].copy()
log(1, "drop GDPC1 quarters without a complete 3-month UNRATE quarterly average", n0, n1,
    "UNRATE starts 1948-01 (1947 quarters have no unemployment data); 2025Q4 has a missing month")
# differences require consecutive quarters
full = df.reindex(pd.period_range(df.index.min(), df.index.max(), freq="Q"))
full["du"] = full["u"].diff()
full["g"] = 100 * np.log(full["gdp"]).diff()
df_d = full.loc[df_ok.index].dropna(subset=["du", "g"])
log(2, "drop quarters whose first difference cannot be formed (previous quarter missing)", n1, len(df_d),
    "first available quarter (1948Q1) and the quarter after the incomplete 2025Q4")
d = df_d.loc[(df_d.index >= S0) & (df_d.index <= S1)].copy()
log(3, "restrict to regression window 1960Q1-2019Q4", len(df_d), len(d),
    "TASK.md information rule: quarters outside 1960Q1-2019Q4 must not enter any estimate")
pd.DataFrame(ledger).to_csv(f"{OUT}/exclusions.csv", index=False)

expected = pd.period_range(S0, S1, freq="Q")
assert len(d) == 240 and (d.index == expected).all(), "sample is not the 240 consecutive quarters"
# 1959Q4 levels used for 1960Q1 differences: check present
assert pd.notna(full.loc[pd.Period("1959Q4", "Q"), "u"]) and pd.notna(full.loc[pd.Period("1959Q4", "Q"), "gdp"])

# --- baseline ---
du, g = d["du"].values, d["g"].values
n = len(d)
baseline_rmse = float(np.sqrt(np.mean((du - du.mean()) ** 2)))

# --- full-sample OLS with Newey-West HAC ---
HAC_LAGS = int(np.floor(4 * (n / 100) ** (2 / 9)))  # Newey-West (1994) rule of thumb -> 4 for n=240
X = sm.add_constant(g)
ols = sm.OLS(du, X).fit()
hac = sm.OLS(du, X).fit(cov_type="HAC", cov_kwds={"maxlags": HAC_LAGS, "use_correction": True})
alpha, beta = float(ols.params[0]), float(ols.params[1])
se_b = float(hac.bse[1])
z = stats.norm.ppf(0.975)
ci = (beta - z * se_b, beta + z * se_b)
model_rmse = float(np.sqrt(np.mean(ols.resid ** 2)))
# lag-length sensitivity (reported, not used for the headline)
hac_sens = {}
for L in [0, 2, 4, 8, 12]:
    r = sm.OLS(du, X).fit(cov_type="HAC", cov_kwds={"maxlags": L, "use_correction": True})
    hac_sens[str(L)] = float(r.bse[1])
ols_se_classical = float(ols.bse[1])

# --- subperiods and Chow test ---
m1 = d.index < SPLIT
def fit(mask):
    y, x = du[mask], sm.add_constant(g[mask])
    r = sm.OLS(y, x).fit()
    rh = sm.OLS(y, x).fit(cov_type="HAC", cov_kwds={"maxlags": HAC_LAGS, "use_correction": True})
    return r, rh
r1, r1h = fit(m1)
r2, r2h = fit(~m1)
ssr_p = float(ols.ssr); ssr_u = float(r1.ssr + r2.ssr)
k = 2
df2 = n - 2 * k
F = ((ssr_p - ssr_u) / k) / (ssr_u / df2)
p = float(stats.f.sf(F, k, df2))

# robustness (supplementary): interacted model, HAC Wald tests of equality
D = (~m1).astype(float)
Xi = np.column_stack([np.ones(n), g, D, D * g])
ri = sm.OLS(du, Xi).fit(cov_type="HAC", cov_kwds={"maxlags": HAC_LAGS, "use_correction": True})
w_joint = ri.wald_test(np.array([[0, 0, 1, 0], [0, 0, 0, 1]]), use_f=False, scalar=True)
w_slope = ri.wald_test(np.array([[0, 0, 0, 1]]), use_f=False, scalar=True)
w_int = ri.wald_test(np.array([[0, 0, 1, 0]]), use_f=False, scalar=True)
resid_var_sub1 = float(r1.ssr / (m1.sum() - 2)); resid_var_sub2 = float(r2.ssr / ((~m1).sum() - 2))

if ci[0] > 0: concl = "positive_effect"
elif ci[1] < 0: concl = "negative_effect"
else: concl = "no_detectable_effect"

data_issues.append({"issue": "GDPC1 is the current (2026) vintage and extends to 2026Q3; UNRATE extends to 2026-09",
                    "handling": "Observations after 2019Q4 dropped before estimation (exclusion step 3); "
                                "revised (not real-time) GDP is used, consistent with an ex-post descriptive estimand."})

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[0].start_time.date()),
    "sample_end": str(d.index[-1].start_time.date()),
    "n_sub1": int(m1.sum()),
    "n_sub2": int((~m1).sum()),
    "baseline_metric": baseline_rmse,
    "model_metric": model_rmse,
    "effect_estimate": beta,
    "intercept": alpha,
    "effect_se": se_b,
    "hac_lags": HAC_LAGS,
    "effect_ci_low": float(ci[0]),
    "effect_ci_high": float(ci[1]),
    "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,
    # supplementary (not required by TASK.md)
    "intercept_sub1": float(r1.params[0]),
    "intercept_sub2": float(r2.params[0]),
    "beta_sub1_hac_se": float(r1h.bse[1]),
    "beta_sub2_hac_se": float(r2h.bse[1]),
    "effect_se_ols_classical": ols_se_classical,
    "effect_se_hac_by_lag": hac_sens,
    "r_squared": float(ols.rsquared),
    "chow_df": [k, int(df2)],
    "ssr_pooled": ssr_p,
    "ssr_unrestricted": ssr_u,
    "hac_wald_joint_equality_p": float(w_joint.pvalue),
    "hac_wald_slope_equality_p": float(w_slope.pvalue),
    "hac_wald_intercept_equality_p": float(w_int.pvalue),
    "resid_var_sub1": resid_var_sub1,
    "resid_var_sub2": resid_var_sub2,
    "growth_for_stable_u_sub1": float(-r1.params[0] / r1.params[1]),
    "growth_for_stable_u_sub2": float(-r2.params[0] / r2.params[1]),
}
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_b),
    dict(period="sub1 1960Q1-1983Q4", n=int(m1.sum()), intercept=float(r1.params[0]), slope=float(r1.params[1]), slope_se=float(r1h.bse[1])),
    dict(period="sub2 1984Q1-2019Q4", n=int((~m1).sum()), intercept=float(r2.params[0]), slope=float(r2.params[1]), slope_se=float(r2h.bse[1])),
]).to_csv(f"{OUT}/table1.csv", index=False)

fig, ax = plt.subplots(figsize=(7, 5))
ax.scatter(g[m1], du[m1], s=14, alpha=0.6, label="1960Q1-1983Q4")
ax.scatter(g[~m1], du[~m1], s=14, alpha=0.6, marker="^", label="1984Q1-2019Q4")
for r, mask, lab in [(r1, m1, "fit 1960-1983"), (r2, ~m1, "fit 1984-2019")]:
    xs = np.linspace(g[mask].min(), g[mask].max(), 50)
    ax.plot(xs, r.params[0] + r.params[1] * xs,
            label=f"{lab}: du = {r.params[0]:.3f} {r.params[1]:+.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("Change in unemployment rate du_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=120)

print(f"n={n} beta={beta:.4f} HAC se={se_b:.4f} (L={HAC_LAGS}) CI=[{ci[0]:.4f},{ci[1]:.4f}] "
      f"beta1={r1.params[1]:.4f} beta2={r2.params[1]:.4f} Chow F={F:.3f} p={p:.4g}")