"""F2-conf-02: OLS slope of y on x with HAC (Newey-West, Bartlett) SE.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Reads only data/series.csv; writes out/results.json, out/exclusions.csv, out/sensitivity.csv.
No randomness is used (random_seed = null)."""
import json, math
import numpy as np, pandas as pd
import statsmodels.api as sm
from scipy import stats

d = pd.read_csv("data/series.csv")
n_raw = len(d)
issues = []
excl = []
# --- audit / exclusions
n0 = len(d); d = d.dropna(subset=["month", "x", "y"])
excl.append(("1", "drop rows with missing month/x/y", n0, len(d), n0 - len(d), "regression requires complete cases"))
n0 = len(d); d = d.drop_duplicates(subset=["month"])
excl.append(("2", "drop duplicated months", n0, len(d), n0 - len(d), "one observation per month"))
per = pd.PeriodIndex(d["month"], freq="M")
d = d.iloc[np.argsort(per.asi8)].reset_index(drop=True)
per = pd.PeriodIndex(d["month"], freq="M")
gaps = int((np.diff(per.asi8) != 1).sum())
if gaps:
    issues.append({"issue": f"{gaps} gaps in monthly sequence", "handling": "kept; HAC lags computed on row order"})
pd.DataFrame(excl, columns=["step", "rule", "n_before", "n_after", "n_removed", "rationale"]).to_csv("out/exclusions.csv", index=False)

X = sm.add_constant(d["x"].to_numpy()); y = d["y"].to_numpy(); T = len(y)

# --- baseline: conventional OLS (the memo's measure)
ols = sm.OLS(y, X).fit()
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))

# --- bandwidth: Andrews (1991) AR(1) plug-in for the Bartlett kernel, applied to the
# slope score v_t = (x_t - xbar) u_t (the series whose long-run variance enters the slope SE).
v = (d["x"].to_numpy() - d["x"].mean()) * u
rho_v = float(np.sum(v[1:] * v[:-1]) / np.sum(v[:-1] ** 2))
a1 = 4 * rho_v**2 / ((1 - rho_v) ** 2 * (1 + rho_v) ** 2)
andrews_bw = 1.1447 * (a1 * T) ** (1 / 3)
hac_lags = int(math.ceil(andrews_bw))
nw94_rule = int(math.floor(4 * (T / 100) ** (2 / 9)))  # Newey-West (1994) rule of thumb, for comparison

def hac(L):
    r = sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": L, "kernel": "bartlett", "use_correction": False})
    return float(r.bse[1])

se_hac = hac(hac_lags)
t_hac = beta_hat / se_hac
p_hac = float(2 * stats.norm.sf(abs(t_hac)))
z = stats.norm.ppf(0.975)
lo, hi = beta_hat - z * se_hac, beta_hat + z * se_hac

# sensitivity over bandwidths (reported, not used for the verdict)
rows = []
for L in sorted(set([0, 1, 2, nw94_rule, 6, 8, 10, 12, hac_lags, 16, 20, 24, 30, 36])):
    s = hac(L); t = beta_hat / s
    rows.append({"lags": L, "se": s, "t": t, "p_normal": 2 * stats.norm.sf(abs(t))})
sens = pd.DataFrame(rows); sens.to_csv("out/sensitivity.csv", index=False)

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(2 * stats.t.sf(abs(beta_hat / se_iid), T - 2)),
    "rho_resid": rho_resid, "rho_score": rho_v,
    "se_hac": se_hac, "hac_lags": hac_lags,
    "andrews_bandwidth_raw": andrews_bw, "nw1994_rule_lags": nw94_rule,
    "se_hac_nw1994_rule": hac(nw94_rule),
    "t_hac": t_hac, "p_hac": p_hac,
    "beta_ci_low": float(lo), "beta_ci_high": float(hi),
    "memo_claim_supported": bool(p_hac < 0.05),
    "sensitivity_max_p": float(sens.loc[sens.lags >= nw94_rule, "p_normal"].max()),
    "sensitivity_min_p": float(sens.loc[sens.lags >= nw94_rule, "p_normal"].min()),
    "n_analysed": int(T),
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"series.csv": int(n_raw)},
    "random_seed": None,
}
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2)); print(sens.to_string(index=False))