"""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

DATA, OUT = "data", "out"
os.makedirs(OUT, exist_ok=True)
FIRST_Q, LAST_Q, SPLIT_Q = pd.Period("1960Q1"), pd.Period("2019Q4"), pd.Period("1984Q1")

gdp_raw = pd.read_csv(f"{DATA}/GDPC1.csv", parse_dates=["observation_date"])
ur_raw = pd.read_csv(f"{DATA}/UNRATE.csv", parse_dates=["observation_date"])
n_input_rows = {"GDPC1.csv": int(len(gdp_raw)), "UNRATE.csv": int(len(ur_raw))}

data_issues = []
ur = ur_raw.copy()
ur["UNRATE"] = pd.to_numeric(ur["UNRATE"], errors="coerce")
ur["q"] = ur["observation_date"].dt.to_period("Q")
miss = ur[ur["UNRATE"].isna()]
if len(miss):
    inside = bool(((miss["q"] >= FIRST_Q - 1) & (miss["q"] <= LAST_Q)).any())
    data_issues.append({
        "issue": f"UNRATE has {len(miss)} missing monthly value(s): "
                 + ", ".join(d.strftime('%Y-%m') for d in miss["observation_date"]),
        "handling": ("inside the estimation window: affected quarter dropped (3 valid months required)" if inside else
                     "outside the 1959Q4-2019Q4 window needed for estimation, so no effect; quarters with "
                     "missing months are never averaged (3 valid months required); no imputation")})
data_issues.append({"issue": "Both files extend beyond 2019Q4 (to 2026); post-2019 quarters are outside the question",
                    "handling": "truncated to 1959Q4-2019Q4 before any computation (1959Q4 used only as lag level)"})
data_issues.append({"issue": "GDPC1 is the 2026 FRED vintage (chained 2017 dollars), not real-time data",
                    "handling": "used as provided; results are conditional on this vintage (noted as limitation)"})

qagg = ur.groupby("q")["UNRATE"].agg(["mean", "count"])
qagg = qagg[qagg["count"] == 3]          # complete quarters only
u = qagg["mean"].rename("u")
gdp = gdp_raw.assign(q=gdp_raw["observation_date"].dt.to_period("Q")).set_index("q")["GDPC1"].astype(float)

df = pd.concat([u, gdp], axis=1, join="inner").sort_index()
df = df[(df.index >= FIRST_Q - 1) & (df.index <= LAST_Q)]
assert len(df) == 241 and (np.diff(df.index.asi8) == 1).all(), "gap in quarterly series"
df["du"] = df["u"].diff()
df["g"] = 100 * np.log(df["GDPC1"]).diff()

n0 = len(df)
d = df.dropna()
pd.DataFrame([("1", "drop 1959Q4 (lag level only; difference undefined)", n0, len(d), n0 - len(d),
               "1959Q4 level used only to form 1960Q1 differences; post-2019Q4 rows were never in the window")],
             columns=["step", "rule", "n_before", "n_after", "n_removed", "rationale"]).to_csv(
    f"{OUT}/exclusions.csv", index=False)
assert len(d) == 240 and d.index[0] == FIRST_Q and d.index[-1] == LAST_Q

def ols(sub):
    return sm.OLS(sub["du"].values, sm.add_constant(sub["g"].values)).fit()

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

def hac_se(fit, lags):
    return float(fit.get_robustcov_results(cov_type="HAC", maxlags=lags, kernel="bartlett",
                                           use_correction=False).bse[1])

n = len(d)
hac_lags = nw_lags(n)
full = ols(d)
beta, alpha = float(full.params[1]), float(full.params[0])
se = hac_se(full, hac_lags)
z = stats.norm.ppf(0.975)
ci_lo, ci_hi = beta - z * se, beta + z * se

baseline_rmse = float(np.sqrt(np.mean((d["du"] - d["du"].mean()) ** 2)))
model_rmse = float(np.sqrt(np.mean(full.resid ** 2)))

s1, s2 = d[d.index < SPLIT_Q], d[d.index >= SPLIT_Q]
f1, f2 = ols(s1), ols(s2)
k = 2
ssr_u = f1.ssr + f2.ssr
chow_f = float(((full.ssr - ssr_u) / k) / (ssr_u / (n - 2 * k)))
chow_p = float(stats.f.sf(chow_f, k, n - 2 * k))

table = pd.DataFrame([
    {"period": "full_1960Q1-2019Q4", "n": n, "intercept": alpha, "slope": beta, "slope_se": se},
    {"period": "sub1_1960Q1-1983Q4", "n": len(s1), "intercept": float(f1.params[0]),
     "slope": float(f1.params[1]), "slope_se": hac_se(f1, nw_lags(len(s1)))},
    {"period": "sub2_1984Q1-2019Q4", "n": len(s2), "intercept": float(f2.params[0]),
     "slope": float(f2.params[1]), "slope_se": hac_se(f2, nw_lags(len(s2)))},
])
table.to_csv(f"{OUT}/table1.csv", index=False)

# robustness (supplementary): HAC Wald tests in the interaction model du = a + b g + c D + d D*g
D = (d.index >= SPLIT_Q).astype(float)
Xi = np.column_stack([np.ones(n), d["g"].values, D, D * d["g"].values])
fi = sm.OLS(d["du"].values, Xi).fit().get_robustcov_results(cov_type="HAC", maxlags=hac_lags,
                                                             kernel="bartlett", use_correction=False)
w_joint = fi.wald_test(np.array([[0, 0, 1, 0], [0, 0, 0, 1]]), use_f=False, scalar=True)
w_slope = fi.wald_test(np.array([[0, 0, 0, 1]]), use_f=False, scalar=True)
w_int = fi.wald_test(np.array([[0, 0, 1, 0]]), use_f=False, scalar=True)

sens = {str(L): hac_se(full, L) for L in (0, 2, 4, 8, 12)}
concl = "positive_effect" if ci_lo > 0 else ("negative_effect" if ci_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": d.index[0].start_time.strftime("%Y-%m-%d"),
    "sample_end": d.index[-1].start_time.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": hac_lags,
    "effect_ci_low": float(ci_lo), "effect_ci_high": float(ci_hi),
    "beta_sub1": float(f1.params[1]), "beta_sub2": float(f2.params[1]),
    "chow_f_stat": chow_f, "chow_p_value": chow_p,
    "stability_conclusion": "stable" if chow_p >= 0.05 else "unstable",
    "conclusion": concl,
    # supplementary keys
    "ols_classical_se": float(full.bse[1]),
    "r_squared": float(full.rsquared),
    "hac_se_by_lag": sens,
    "intercept_sub1": float(f1.params[0]), "intercept_sub2": float(f2.params[0]),
    "beta_sub1_hac_se": float(table.loc[1, "slope_se"]), "beta_sub2_hac_se": float(table.loc[2, "slope_se"]),
    "hac_lags_sub1": nw_lags(len(s1)), "hac_lags_sub2": nw_lags(len(s2)),
    "chow_df": [k, int(n - 2 * k)],
    "hac_wald_joint_p": float(w_joint.pvalue), "hac_wald_slope_diff_p": float(w_slope.pvalue),
    "hac_wald_intercept_diff_p": float(w_int.pvalue),
    "okun_growth_sub1": float(-f1.params[0] / f1.params[1]), "okun_growth_sub2": float(-f2.params[0] / f2.params[1]),
    "n_window_quarters": int(n0),
    "mean_du": float(d["du"].mean()), "mean_g": float(d["g"].mean()),
}
with open(f"{OUT}/results.json", "w") as fh:
    json.dump(res, fh, indent=2)

fig, ax = plt.subplots(figsize=(7, 5))
ax.scatter(s1["g"], s1["du"], s=14, alpha=0.7, label="1960Q1–1983Q4")
ax.scatter(s2["g"], s2["du"], s=14, alpha=0.7, marker="^", label="1984Q1–2019Q4")
for sub, fit, ls in ((s1, f1, "-"), (s2, f2, "--")):
    xs = np.linspace(sub["g"].min(), sub["g"].max(), 50)
    ax.plot(xs, fit.params[0] + fit.params[1] * xs, ls, color="k",
            label=f"OLS {sub.index[0]}–{sub.index[-1]}: β={fit.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 1960Q1–2019Q4")
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))