"""Okun's law (difference form), US 1960Q1-2019Q4. 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, 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

OUT = "out"; os.makedirs(OUT, exist_ok=True)
SEED = None  # no randomness used
LEVEL_START, END = pd.Timestamp("1959-10-01"), pd.Timestamp("2019-10-01")
S1_END, S2_START = pd.Timestamp("1983-10-01"), pd.Timestamp("1984-01-01")

gdp = pd.read_csv("data/GDPC1.csv", parse_dates=["observation_date"])
un = pd.read_csv("data/UNRATE.csv", parse_dates=["observation_date"])
n_input = {"GDPC1.csv": int(len(gdp)), "UNRATE.csv": int(len(un))}
excl = []
def led(step, rule, nb, na, why):
    excl.append(dict(step=step, rule=rule, n_before=nb, n_after=na, n_removed=nb - na, rationale=why))

# --- unemployment: monthly -> quarterly average (complete quarters only)
nb = len(un); un_nonmiss = un.dropna(subset=["UNRATE"])
led("UNRATE_1", "drop months with blank UNRATE", nb, len(un_nonmiss), "2025-10 value blank in FRED file; outside sample")
nb = len(un_nonmiss)
un_w = un_nonmiss[(un_nonmiss.observation_date >= LEVEL_START) & (un_nonmiss.observation_date < END + pd.DateOffset(months=3))]
led("UNRATE_2", "keep months 1959-10..2019-12", nb, len(un_w), "levels needed for 1960Q1-2019Q4 differences; post-2019 excluded by rule")
q = un_w.observation_date.dt.to_period("Q").dt.start_time
grp = un_w.groupby(q)["UNRATE"]
cnt = grp.count()
assert (cnt == 3).all(), "incomplete quarter in sample window"
u = grp.mean().rename("u")
led("UNRATE_3", "average 3 months -> quarters", len(un_w), len(u), "u_t = mean of 3 monthly values (rows are now quarters)")

# --- GDP
nb = len(gdp)
gdp_w = gdp[(gdp.observation_date >= LEVEL_START) & (gdp.observation_date <= END)].set_index("observation_date")["GDPC1"]
led("GDPC1_1", "keep quarters 1959Q4..2019Q4", nb, len(gdp_w), "levels needed for 1960Q1-2019Q4 growth; post-2019 excluded by rule")
assert gdp_w.notna().all()

df = pd.concat([u, gdp_w], axis=1, join="inner")
assert len(df) == 241 and (df.index.to_series().diff().dropna().dt.days.between(89, 93)).all()
df["du"] = df["u"].diff()
df["g"] = 100 * np.log(df["GDPC1"]).diff()
nb = len(df); df = df.dropna()
led("MERGED_1", "first difference (drop 1959Q4)", nb, len(df), "1959Q4 only provides lagged levels")
assert df.index.min() == pd.Timestamp("1960-01-01") and df.index.max() == END and len(df) == 240
pd.DataFrame(excl).to_csv(f"{OUT}/exclusions.csv", index=False)

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

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

full, full_hac, L = fit(df)
d1 = df[df.index <= S1_END]; d2 = df[df.index >= S2_START]
f1, h1, L1 = fit(d1); f2, h2, L2 = fit(d2)
assert len(d1) + len(d2) == len(df)

beta = float(full.params["g"]); alpha = float(full.params["const"])
se = float(full_hac.bse[1])
z = stats.norm.ppf(0.975)
lo, hi = beta - z * se, beta + z * se

# Chow test (classical, SSR-based)
n, k = len(df), 2
ssr_p, ssr_1, ssr_2 = full.ssr, f1.ssr, f2.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: dummy-interaction model with HAC SEs (robust to the heteroskedasticity the classical Chow ignores)
D = (df.index >= S2_START).astype(float)
Xi = pd.DataFrame({"const": 1.0, "g": df.g, "D": D, "Dg": D * df.g}, index=df.index)
inter = sm.OLS(df.du, Xi).fit().get_robustcov_results(cov_type="HAC", maxlags=L, kernel="bartlett", use_correction=True)
wald_slope_p = float(inter.pvalues[3]); wald_int_p = float(inter.pvalues[2])
wald_joint = inter.wald_test(np.array([[0, 0, 1, 0], [0, 0, 0, 1]]), scalar=True, use_f=True)
wald_joint_p = float(wald_joint.pvalue)
resvar1 = float(f1.ssr / (len(d1) - 2)); resvar2 = float(f2.ssr / (len(d2) - 2))

baseline = float(np.sqrt(np.mean((df.du - df.du.mean()) ** 2)))
model = float(np.sqrt(np.mean(full.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.min().strftime("%Y-%m-%d"),
    "sample_end": df.index.max().strftime("%Y-%m-%d"),
    "n_sub1": int(len(d1)), "n_sub2": int(len(d2)),
    "baseline_metric": baseline, "model_metric": model,
    "effect_estimate": beta, "intercept": alpha, "effect_se": se, "hac_lags": L,
    "effect_ci_low": float(lo), "effect_ci_high": float(hi),
    "beta_sub1": float(f1.params["g"]), "beta_sub2": float(f2.params["g"]),
    "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": [
        {"issue": "UNRATE 2025-10-01 value is blank in the FRED file", "handling": "dropped; outside the 1959Q4-2019Q4 window, no effect on estimates"},
        {"issue": "Files extend to 2026 (beyond question period)", "handling": "truncated to 1959Q4-2019Q4 levels before differencing; no post-2019Q4 data enters any estimate"},
        {"issue": "Data are the 2026 vintage (revised, chained 2017 dollars), not real-time", "handling": "used as provided; noted as limitation"},
    ],
    "n_input_rows": n_input, "n_analysed": int(full.nobs), "random_seed": SEED,
    # supplementary
    "r_squared_full": float(full.rsquared),
    "effect_se_ols_classical": float(full.bse["g"]),
    "chow_df": [k, n - 2 * k],
    "hac_lags_sub1": L1, "hac_lags_sub2": L2,
    "slope_se_sub1_hac": float(h1.bse[1]), "slope_se_sub2_hac": float(h2.bse[1]),
    "intercept_sub1": float(f1.params["const"]), "intercept_sub2": float(f2.params["const"]),
    "effect_se_hac_lags8": float(full.get_robustcov_results(cov_type="HAC", maxlags=8, use_correction=True).bse[1]),
    "hac_wald_slope_diff_p": wald_slope_p, "hac_wald_intercept_diff_p": wald_int_p,
    "hac_wald_joint_p": wald_joint_p, "resid_var_sub1": resvar1, "resid_var_sub2": resvar2,
    "mean_du": float(df.du.mean()), "mean_g": float(df.g.mean()),
}
with open(f"{OUT}/results.json", "w") as fh:
    json.dump(res, fh, indent=2)

pd.DataFrame([
    dict(period="full_1960Q1-2019Q4", n=len(df), intercept=alpha, slope=beta, slope_se=se),
    dict(period="sub1_1960Q1-1983Q4", n=len(d1), intercept=f1.params["const"], slope=f1.params["g"], slope_se=h1.bse[1]),
    dict(period="sub2_1984Q1-2019Q4", n=len(d2), intercept=f2.params["const"], slope=f2.params["g"], slope_se=h2.bse[1]),
]).to_csv(f"{OUT}/table1.csv", index=False)

fig, ax = plt.subplots(figsize=(7, 5))
ax.scatter(d1.g, d1.du, s=14, alpha=0.6, label=f"1960Q1-1983Q4 (n={len(d1)})", color="tab:blue")
ax.scatter(d2.g, d2.du, s=14, alpha=0.6, label=f"1984Q1-2019Q4 (n={len(d2)})", color="tab:orange", marker="s")
for d, f, c, nm in [(d1, f1, "tab:blue", "sub1"), (d2, f2, "tab:orange", "sub2")]:
    xs = np.linspace(d.g.min(), d.g.max(), 50)
    ax.plot(xs, f.params["const"] + f.params["g"] * xs, color=c, lw=2,
            label=f"{nm} fit: du = {f.params['const']:.3f} + ({f.params['g']:.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(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=150)

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",
                                       "baseline_metric", "model_metric", "conclusion", "stability_conclusion"]}, indent=1))