"""F2-conf-02: OLS slope of y on x with Newey-West HAC inference.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Reads only data/series.csv (memo numbers are checked against, not used). No randomness."""
import json, math
import numpy as np, pandas as pd, statsmodels.api as sm
from scipy import stats

d = pd.read_csv("data/series.csv")
n_raw = len(d)
issues = []
# --- data checks (no exclusions needed if all pass) ---
per = pd.PeriodIndex(d["month"], freq="M")
gaps = int(((per[1:].astype("int64") - per[:-1].astype("int64")) != 1).sum())
n_dup = int(d["month"].duplicated().sum())
n_na = int(d[["x", "y"]].isna().any(axis=1).sum())
if gaps: issues.append({"issue": f"{gaps} non-consecutive month steps", "handling": "kept; HAC assumes equal spacing"})
if n_dup: issues.append({"issue": f"{n_dup} duplicate months", "handling": "reported"})
if n_na:
    issues.append({"issue": f"{n_na} rows with missing x/y", "handling": "dropped"})
    d = d.dropna(subset=["x", "y"])
d = d.sort_values("month").reset_index(drop=True)
T = len(d)

X = sm.add_constant(d["x"].values)
y = d["y"].values
ols = sm.OLS(y, X).fit()  # conventional iid SE (memo baseline)
alpha_hat, beta_hat = map(float, ols.params)
se_iid = float(ols.bse[1])
u = ols.resid
rho_resid = float(np.sum(u[1:] * u[:-1]) / np.sum(u * u))  # standard first-order autocorrelation coefficient

# --- bandwidth: Andrews (1991) AR(1) plug-in for Bartlett kernel, applied to the slope score v_t=(x_t-xbar)u_t ---
v = (d["x"].values - d["x"].values.mean()) * u
rho_v = float(np.sum(v[1:] * v[:-1]) / np.sum(v * v))
a1 = 4 * rho_v**2 / ((1 - rho_v) ** 2 * (1 + rho_v) ** 2)
S_andrews = 1.1447 * (a1 * T) ** (1 / 3)
nw_rule = int(math.floor(4 * (T / 100) ** (2 / 9)))  # Newey-West (1994) rule of thumb, reported as sensitivity
hac_lags = max(int(math.floor(S_andrews)), nw_rule)

def hac(L):
    r = sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": L, "kernel": "bartlett", "use_correction": True})
    se = float(r.bse[1]); t = beta_hat / se
    return se, t, float(2 * stats.norm.sf(abs(t)))

se_hac, t_hac, p_hac = hac(hac_lags)
z = stats.norm.ppf(0.975)
ci_lo, ci_hi = beta_hat - z * se_hac, beta_hat + z * se_hac
sens = {str(L): dict(zip(["se", "t", "p"], hac(L))) for L in sorted({nw_rule, 4, 8, 12, 16, 24, hac_lags})}

res = {
    "task_id": "F2-conf-02",
    "beta_hat": beta_hat, "alpha_hat": alpha_hat, "se_iid": se_iid,
    "t_iid": beta_hat / se_iid, "p_iid": float(ols.pvalues[1]),
    "rho_resid": rho_resid, "se_hac": se_hac, "hac_lags": hac_lags,
    "t_hac": t_hac, "p_hac": p_hac, "beta_ci_low": ci_lo, "beta_ci_high": ci_hi,
    "memo_claim_supported": bool(p_hac < 0.05),
    "n_analysed": int(T), "feasibility": "complete", "feasibility_reason": None,
    "data_issues": issues, "n_input_rows": {"series.csv": int(n_raw)}, "random_seed": None,
    "bandwidth_details": {"rho_score": rho_v, "andrews_alpha1": a1, "andrews_S": S_andrews,
                          "nw_rule_of_thumb_lags": nw_rule, "rule": "max(floor(Andrews AR(1) plug-in), floor(4(T/100)^(2/9)))"},
    "hac_sensitivity": sens,
    "se_ratio_hac_over_iid": se_hac / se_iid,
    "memo_reproduced": {"alpha": round(alpha_hat, 4) == 2.0925, "beta": round(beta_hat, 4) == 0.1017,
                         "se": round(se_iid, 4) == 0.0448, "t": round(beta_hat / se_iid, 2) == 2.27,
                         "p": round(float(ols.pvalues[1]), 4) == 0.0238},
    "any_sensitivity_p_lt_005": bool(min(s["p"] for s in sens.values()) < 0.05),
}
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(f"beta={beta_hat:.4f} se_iid={se_iid:.4f} se_hac={se_hac:.4f} (L={hac_lags}) t_hac={t_hac:.3f} p_hac={p_hac:.4f} "
      f"CI=[{ci_lo:.4f},{ci_hi:.4f}] memo_supported={res['memo_claim_supported']}")