"""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.
No randomness is used (random_seed = null)."""
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 = pd.read_csv("data/GDPC1.csv", na_values=["."])
un = pd.read_csv("data/UNRATE.csv", na_values=["."])
n_input_rows = {"GDPC1.csv": int(len(gdp)), "UNRATE.csv": int(len(un))}
gdp["observation_date"] = pd.to_datetime(gdp["observation_date"])
un["observation_date"] = pd.to_datetime(un["observation_date"])

data_issues = []
miss = un[un["UNRATE"].isna()]["observation_date"].dt.strftime("%Y-%m-%d").tolist()
if miss:
    data_issues.append({"issue": f"UNRATE missing value(s) at {miss} (blank in raw FRED CSV)",
                        "handling": "outside the 1959Q4-2019Q4 window; dropped with all post-2019Q4 observations, no effect on estimates"})
data_issues.append({"issue": "Both series extend beyond 2019Q4 (to 2026) and are the 2026 vintage (revised, not real-time)",
                    "handling": "all observations after 2019Q4 excluded before computation; vintage noted as limitation"})

# Quarterly unemployment = mean of 3 months; require all 3 months present
un["q"] = un["observation_date"].dt.to_period("Q")
uq = un.groupby("q")["UNRATE"].agg(["mean", "count"])
gdp["q"] = gdp["observation_date"].dt.to_period("Q")
df = gdp.set_index("q")[["GDPC1"]].join(uq, how="left")

excl = []
def log(step, rule, before, after, why):
    excl.append(dict(step=step, rule=rule, n_before=before, n_after=after, n_removed=before - after, rationale=why))

n0 = len(df)
win = df.loc[pd.Period("1959Q4"):pd.Period("2019Q4")].copy()
log(1, "keep quarters 1959Q4-2019Q4 (GDPC1 quarters as unit)", n0, len(win),
    "task window; 1959Q4 needed only for lagged levels; post-2019Q4 forbidden")
assert (win["count"] == 3).all() and win["GDPC1"].notna().all(), "incomplete data inside window"
win["u"] = win["mean"]
win["du"] = win["u"].diff()
win["g"] = 100 * np.log(win["GDPC1"]).diff()
reg = win.dropna(subset=["du", "g"]).copy()
log(2, "drop 1959Q4 (no lagged level, differences undefined)", len(win), len(reg),
    "first-difference construction")
pd.DataFrame(excl).to_csv(f"{OUT}/exclusions.csv", index=False)
assert len(reg) == 240 and str(reg.index[0]) == "1960Q1" and str(reg.index[-1]) == "2019Q4"

y, X = reg["du"].values, sm.add_constant(reg["g"].values)
n = len(y)
# Baseline: mean-only prediction
baseline_rmse = float(np.sqrt(np.mean((y - y.mean()) ** 2)))
# Full sample OLS with Newey-West HAC (Bartlett) SE
HAC_LAGS = int(np.floor(4 * (n / 100) ** (2 / 9)))  # Newey-West (1994) rule of thumb -> 4
ols = sm.OLS(y, X).fit()
hac = sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": HAC_LAGS, "use_correction": True})
alpha, beta = float(ols.params[0]), float(ols.params[1])
se = float(hac.bse[1])
z = stats.norm.ppf(0.975)
ci = (beta - z * se, beta + z * se)
model_rmse = float(np.sqrt(np.mean(ols.resid ** 2)))

# Subperiods
s1 = reg.index <= pd.Period("1983Q4")
s2 = ~s1
fits = {}
for name, m in [("sub1", s1), ("sub2", s2)]:
    Xs = sm.add_constant(reg["g"].values[m])
    f = sm.OLS(y[m], Xs).fit()
    fh = sm.OLS(y[m], Xs).fit(cov_type="HAC", cov_kwds={"maxlags": HAC_LAGS, "use_correction": True})
    fits[name] = (f, fh, int(m.sum()))
ssr_p = ols.ssr
ssr_u = fits["sub1"][0].ssr + fits["sub2"][0].ssr
k = 2
chow_f = float(((ssr_p - ssr_u) / k) / (ssr_u / (n - 2 * k)))
chow_p = float(stats.f.sf(chow_f, k, n - 2 * k))

# Supplementary: HAC test of slope-only difference (interaction model) and lag sensitivity
d2 = s2.astype(float)
Xi = np.column_stack([np.ones(n), reg["g"].values, d2, d2 * reg["g"].values])
inter = sm.OLS(y, Xi).fit(cov_type="HAC", cov_kwds={"maxlags": HAC_LAGS, "use_correction": True})
slope_diff_p = float(inter.pvalues[3]); int_diff_p = float(inter.pvalues[2])
R = np.array([[0, 0, 1, 0], [0, 0, 0, 1]], dtype=float)
wald = inter.wald_test(R, use_f=True, scalar=True)
hac_joint_F, hac_joint_p = float(wald.statistic), float(wald.pvalue)
lag_sens = {str(L): float(sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": L, "use_correction": True}).bse[1])
            for L in (0, 2, 8, 12)}
resid_var = {"sub1": float(fits["sub1"][0].ssr / (fits["sub1"][2] - 2)), "sub2": float(fits["sub2"][0].ssr / (fits["sub2"][2] - 2))}

conclusion = "positive_effect" if ci[0] > 0 else ("negative_effect" if ci[1] < 0 else "no_detectable_effect")
results = {
    "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": reg.index[0].start_time.strftime("%Y-%m-%d"),
    "sample_end": reg.index[-1].start_time.strftime("%Y-%m-%d"),
    "n_sub1": fits["sub1"][2],
    "n_sub2": fits["sub2"][2],
    "baseline_metric": baseline_rmse,
    "model_metric": model_rmse,
    "effect_estimate": beta,
    "intercept": alpha,
    "effect_se": se,
    "hac_lags": HAC_LAGS,
    "effect_ci_low": float(ci[0]),
    "effect_ci_high": float(ci[1]),
    "beta_sub1": float(fits["sub1"][0].params[1]),
    "beta_sub2": float(fits["sub2"][0].params[1]),
    "chow_f_stat": chow_f,
    "chow_p_value": chow_p,
    "stability_conclusion": "stable" if chow_p >= 0.05 else "unstable",
    "conclusion": conclusion,
    # supplementary (not required)
    "ols_classical_se_full": float(ols.bse[1]),
    "r_squared_full": float(ols.rsquared),
    "intercept_sub1": float(fits["sub1"][0].params[0]),
    "intercept_sub2": float(fits["sub2"][0].params[0]),
    "beta_sub1_hac_se": float(fits["sub1"][1].bse[1]),
    "beta_sub2_hac_se": float(fits["sub2"][1].bse[1]),
    "implied_zero_du_growth_full": float(-alpha / beta),
    "hac_slope_diff_p_value": slope_diff_p,
    "hac_intercept_diff_p_value": int_diff_p,
    "hac_joint_break_f_stat": hac_joint_F,
    "hac_joint_break_p_value": hac_joint_p,
    "hac_se_lag_sensitivity": lag_sens,
    "resid_variance_sub1": resid_var["sub1"],
    "resid_variance_sub2": resid_var["sub2"],
}
with open(f"{OUT}/results.json", "w") as fh_:
    json.dump(results, fh_, indent=2)

rows = [("full_1960Q1_2019Q4", n, alpha, beta, se)]
for lab, key in [("sub1_1960Q1_1983Q4", "sub1"), ("sub2_1984Q1_2019Q4", "sub2")]:
    f, fh, nn = fits[key]
    rows.append((lab, nn, float(f.params[0]), float(f.params[1]), float(fh.bse[1])))
pd.DataFrame(rows, columns=["period", "n", "intercept", "slope", "slope_se"]).to_csv(f"{OUT}/table1.csv", index=False)

fig, ax = plt.subplots(figsize=(7, 5))
g = reg["g"].values
ax.scatter(g[s1], y[s1], s=14, alpha=0.6, label="1960Q1-1983Q4", color="tab:blue")
ax.scatter(g[s2], y[s2], s=14, alpha=0.6, label="1984Q1-2019Q4", color="tab:orange")
xx = np.linspace(g.min(), g.max(), 50)
for key, c in [("sub1", "tab:blue"), ("sub2", "tab:orange")]:
    a, b = fits[key][0].params
    ax.plot(xx, a + b * xx, color=c, lw=2, label=f"fit {key}: du = {a:.2f} {b:+.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(json.dumps({k: results[k] for k in ["effect_estimate", "effect_se", "effect_ci_low", "effect_ci_high",
      "beta_sub1", "beta_sub2", "chow_f_stat", "chow_p_value", "baseline_metric", "model_metric"]}, indent=1))