"""F6-conf-03 Okun's law replication. Run from workspace root:
   ~/wp29/venv/bin/python out/analysis.py
Reads data/ only; 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)
SEED = None  # no randomness used
START, END = pd.Timestamp("1960-01-01"), pd.Timestamp("2019-10-01")
SPLIT = pd.Timestamp("1984-01-01")

gdp = pd.read_csv("data/GDPC1.csv")
un = pd.read_csv("data/UNRATE.csv")
n_input = {"GDPC1.csv": int(len(gdp)), "UNRATE.csv": int(len(un))}
gdp["date"] = pd.to_datetime(gdp["observation_date"]); un["date"] = pd.to_datetime(un["observation_date"])
gdp["GDPC1"] = pd.to_numeric(gdp["GDPC1"], errors="coerce")
un["UNRATE"] = pd.to_numeric(un["UNRATE"], errors="coerce")

data_issues = []
miss_un = un.loc[un["UNRATE"].isna(), "observation_date"].tolist()
if miss_un:
    data_issues.append({"issue": f"UNRATE missing (empty value) for month(s) {miss_un}",
                        "handling": "Outside 1959Q4-2019Q4 window; affected quarter set to missing and in any case excluded by the window restriction, no imputation."})
miss_g = gdp.loc[gdp["GDPC1"].isna(), "observation_date"].tolist()
if miss_g:
    data_issues.append({"issue": f"GDPC1 missing for {miss_g}", "handling": "dropped if outside window; fatal if inside"})
data_issues.append({"issue": "Files extend beyond 2019Q4 (GDPC1 to 2026Q3, UNRATE to 2026-09), including COVID-19 quarters",
                    "handling": "Excluded by the information rule: no quarter after 2019Q4 enters any estimate."})
data_issues.append({"issue": "GDPC1 is the current (2026) vintage, revised data, not real-time",
                    "handling": "Used as provided; noted as limitation (estimates reflect revised GDP)."})

# Quarterly unemployment: mean of 3 months, require complete quarter
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 quarter -> missing
g = gdp.set_index(gdp["date"].dt.to_period("Q"))["GDPC1"]
df = pd.DataFrame({"u": uq["mean"], "y": g}).sort_index()
df["du"] = df["u"].diff()
df["g"] = 100 * np.log(df["y"]).diff()
df["date"] = df.index.to_timestamp()

excl = []
n0 = len(df)
step1 = df[(df["date"] >= START) & (df["date"] <= END)]
excl.append(["1", "keep quarters 1960Q1-2019Q4 (union of GDPC1/UNRATE quarterly index)", n0, len(step1), n0 - len(step1),
             "information rule: regression sample; pre-1960 used only as lag for 1960Q1 difference; post-2019 excluded"])
step2 = step1.dropna(subset=["du", "g"])
excl.append(["2", "drop quarters with missing du or g", len(step1), len(step2), len(step1) - len(step2),
             "no missing values expected inside the window"])
pd.DataFrame(excl, columns=["step", "rule", "n_before", "n_after", "n_removed", "rationale"]).to_csv(f"{OUT}/exclusions.csv", index=False)
d = step2.copy()
assert len(d) == 240, len(d)
assert d["date"].max() <= END

def nw_lags(n):  # Newey-West (1994) plug-in rule of thumb
    return int(np.floor(4 * (n / 100) ** (2 / 9)))

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

full_ols, full_hac, L = fit(d)
s1 = d[d["date"] < SPLIT]; s2 = d[d["date"] >= SPLIT]
o1, h1, L1 = fit(s1); o2, h2, L2 = fit(s2)

# baseline / model metrics
baseline_rmse = float(np.sqrt(np.mean((d["du"] - d["du"].mean()) ** 2)))
model_rmse = float(np.sqrt(np.mean(full_ols.resid ** 2)))
beta = float(full_ols.params[1]); alpha = 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

# Chow test
n, k = len(d), 2
ssr_p, ssr_1, ssr_2 = full_ols.ssr, o1.ssr, o2.ssr
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))


# Supplementary robustness for the stability question (not the required Chow test):
# HAC Wald tests in a pooled regression with subperiod-2 dummy and interaction.
D = (d["date"] >= SPLIT).astype(float).values
Xi = np.column_stack([np.ones(n), d["g"].values, D, D * d["g"].values])
hi_fit = sm.OLS(d["du"].values, Xi).fit(cov_type="HAC", cov_kwds={"maxlags": L, "kernel": "bartlett", "use_correction": True})
w_joint = hi_fit.wald_test(np.array([[0, 0, 1, 0], [0, 0, 0, 1]]), use_f=True, scalar=True)
w_slope = hi_fit.wald_test(np.array([[0, 0, 0, 1]]), use_f=True, scalar=True)
var_ratio = (ssr_1 / (len(s1) - 2)) / (ssr_2 / (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",
    "sample_start": d["date"].min().strftime("%Y-%m-%d"),
    "sample_end": d["date"].max().strftime("%Y-%m-%d"),
    "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,
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": data_issues, "n_input_rows": n_input,
    "n_analysed": int(n), "random_seed": SEED,
    # supplementary (not required)
    "ols_classical_se": float(full_ols.bse[1]),
    "r_squared": float(full_ols.rsquared),
    "se_sub1_hac": float(h1.bse[1]), "se_sub2_hac": float(h2.bse[1]),
    "hac_lags_sub1": L1, "hac_lags_sub2": L2,
    "intercept_sub1": float(o1.params[0]), "intercept_sub2": float(o2.params[0]),
    "ssr_pooled": float(ssr_p), "ssr_sub1": float(ssr_1), "ssr_sub2": float(ssr_2),
    "supp_hac_wald_joint_F": float(w_joint.statistic), "supp_hac_wald_joint_p": float(w_joint.pvalue),
    "supp_hac_wald_slope_F": float(w_slope.statistic), "supp_hac_wald_slope_p": float(w_slope.pvalue),
    "supp_resid_var_ratio_sub1_over_sub2": float(var_ratio),
    "supp_var_ratio_F_p_two_sided": float(2 * min(stats.f.sf(var_ratio, len(s1) - 2, len(s2) - 2), stats.f.cdf(var_ratio, len(s1) - 2, len(s2) - 2))),
    "sensitivity_effect_se_by_lag": {str(l): float(sm.OLS(d["du"].values, sm.add_constant(d["g"].values)).fit(
        cov_type="HAC", cov_kwds={"maxlags": l, "kernel": "bartlett", "use_correction": True}).bse[1]) for l in (0, 2, 4, 8, 12)},
}
with open(f"{OUT}/results.json", "w") as f:
    json.dump(res, f, indent=2)

tab = pd.DataFrame([
    ["full_1960Q1-2019Q4", n, alpha, beta, se],
    ["sub1_1960Q1-1983Q4", len(s1), o1.params[0], o1.params[1], h1.bse[1]],
    ["sub2_1984Q1-2019Q4", len(s2), o2.params[0], o2.params[1], h2.bse[1]],
], columns=["period", "n", "intercept", "slope", "slope_se"])
tab.to_csv(f"{OUT}/table1.csv", index=False)

fig, ax = plt.subplots(figsize=(7, 5))
ax.scatter(s1["g"], s1["du"], s=14, alpha=.6, label="1960Q1–1983Q4")
ax.scatter(s2["g"], s2["du"], s=14, alpha=.6, marker="^", label="1984Q1–2019Q4")
xs = np.linspace(d["g"].min(), d["g"].max(), 50)
ax.plot(xs, o1.params[0] + o1.params[1] * xs, lw=2, label=f"fit 1: β={o1.params[1]:.3f}")
ax.plot(xs, o2.params[0] + o2.params[1] * xs, lw=2, ls="--", label=f"fit 2: β={o2.params[1]:.3f}")
ax.axhline(0, c="grey", lw=.5); ax.axvline(0, c="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(f"{OUT}/figure1.png", dpi=120)

print(json.dumps({k: res[k] for k in ["effect_estimate", "effect_se", "hac_lags", "effect_ci_low", "effect_ci_high",
                                      "beta_sub1", "beta_sub2", "chow_f_stat", "chow_p_value", "stability_conclusion",
                                      "conclusion", "baseline_metric", "model_metric"]}, indent=1))