"""F2-conf-04: Okun's law, US quarterly 1960Q1-2019Q4, OLS + Newey-West HAC SE.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py"""
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"
un = pd.read_csv("data/UNRATE.csv")
gd = pd.read_csv("data/GDPC1.csv")
n_rows = {"UNRATE.csv": int(len(un)), "GDPC1.csv": int(len(gd))}
issues = []

un["date"] = pd.to_datetime(un["observation_date"])
gd["date"] = pd.to_datetime(gd["observation_date"])
un["UNRATE"] = pd.to_numeric(un["UNRATE"], errors="coerce")
gd["GDPC1"] = pd.to_numeric(gd["GDPC1"], errors="coerce")

miss_u = un.loc[un["UNRATE"].isna(), "observation_date"].tolist()
if miss_u:
    issues.append({"issue": f"UNRATE missing value(s) at {miss_u} (empty field in the FRED file)",
                   "handling": "Outside the 1959Q4-2019Q4 window needed for estimation; no effect. "
                               "Quarterly averages require all 3 months present, so affected quarters would be NaN."})
miss_g = gd.loc[gd["GDPC1"].isna(), "observation_date"].tolist()
if miss_g:
    issues.append({"issue": f"GDPC1 missing value(s) at {miss_g}", "handling": "Quarter set to NaN."})

un["q"] = un["date"].dt.to_period("Q")
uq = un.groupby("q")["UNRATE"].agg(["mean", "count"])
uq.loc[uq["count"] < 3, "mean"] = np.nan  # require all 3 months
gd["q"] = gd["date"].dt.to_period("Q")
assert not gd["q"].duplicated().any()
df = pd.DataFrame({"u": uq["mean"]}).join(gd.set_index("q")["GDPC1"], how="outer").sort_index()
# ensure consecutive quarters
full = pd.period_range(df.index.min(), df.index.max(), freq="Q")
df = df.reindex(full)
df["du"] = df["u"].diff()
df["g"] = 100 * np.log(df["GDPC1"]).diff()
s = df.loc[pd.Period("1960Q1"):pd.Period("2019Q4")]
assert len(s) == 240 and s[["du", "g"]].notna().all().all()

y = s["du"].values
X = sm.add_constant(s["g"].values)
T = len(y)
ols = sm.OLS(y, X).fit()
# Newey-West (1994) rule-of-thumb bandwidth for the Bartlett kernel: floor(4*(T/100)^(2/9))
L = int(np.floor(4 * (T / 100) ** (2 / 9)))
hac = sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": L, "kernel": "bartlett", "use_correction": True})
beta, alpha = float(ols.params[1]), float(ols.params[0])
se_iid = float(ols.bse[1]); se_hac = float(hac.bse[1])
z = stats.norm.ppf(0.975)
ci = (beta - z * se_hac, beta + z * se_hac)
p_hac = float(2 * stats.norm.sf(abs(beta / se_hac)))

# sensitivity of HAC SE to bandwidth (reported in report, not a required key)
sens = {}
for lag in [0, 1, 2, 4, 6, 8, 12, 16]:
    r = sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": lag, "kernel": "bartlett", "use_correction": True})
    sens[lag] = float(r.bse[1])
acf_res = [float(np.corrcoef(ols.resid[k:], ols.resid[:-k])[0, 1]) for k in range(1, 5)]

res = {
    "task_id": "F2-conf-04",
    "beta_hat": beta, "alpha_hat": alpha, "r_squared": float(ols.rsquared),
    "se_iid": se_iid, "se_hac": se_hac, "hac_lags": L,
    "beta_ci_low": float(ci[0]), "beta_ci_high": float(ci[1]),
    "okun_negative_5pct": bool(beta < 0 and p_hac < 0.05),
    "mean_growth": float(s["g"].mean()),
    "first_quarter": str(s.index[0]), "last_quarter": str(s.index[-1]),
    "n_analysed": int(T),
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": issues, "n_input_rows": n_rows, "random_seed": None,
}
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(res, f, indent=2)
with open(os.path.join(OUT, "diagnostics.json"), "w") as f:
    json.dump({"p_value_hac_normal": p_hac, "t_hac": beta / se_hac,
               "hac_se_by_lag": sens, "resid_acf_1to4": acf_res,
               "dw": float(sm.stats.durbin_watson(ols.resid))}, f, indent=2)

fig, ax = plt.subplots(figsize=(6.5, 4.5))
ax.scatter(s["g"], s["du"], s=12, alpha=0.6)
gx = np.linspace(s["g"].min(), s["g"].max(), 50)
ax.plot(gx, alpha + beta * gx, "r-", label=f"OLS: du = {alpha:.3f} {beta:+.3f} g")
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 (n=240)")
ax.axhline(0, color="grey", lw=0.5); ax.axvline(0, color="grey", lw=0.5)
ax.legend(); fig.tight_layout(); fig.savefig(os.path.join(OUT, "figure1.png"), dpi=130)
print(json.dumps(res, indent=2))