"""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/hac_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", dtype={"month": str})
n_raw = len(d)
issues = []
# --- audit / exclusions ---
excl = []
def step(name, rule, mask, why):
    global d
    nb = len(d); d = d[mask].copy(); excl.append(dict(step=name, rule=rule, n_before=nb, n_after=len(d), n_removed=nb-len(d), rationale=why))
step("1", "drop rows with missing month/x/y", d[["month","x","y"]].notna().all(axis=1), "OLS requires complete cases")
step("2", "drop duplicated months", ~d.duplicated("month"), "one observation per month")
per = pd.PeriodIndex(d["month"], freq="M")
gaps = int((np.diff(per.asi8) != 1).sum())
if gaps: issues.append({"issue": f"{gaps} gaps/out-of-order months", "handling": "kept; HAC uses row order"})
pd.DataFrame(excl).to_csv("out/exclusions.csv", index=False)
n = len(d)

y = d["y"].to_numpy(); x = d["x"].to_numpy()
X = sm.add_constant(x)
ols = sm.OLS(y, X).fit()
alpha_hat, beta_hat = ols.params
se_iid = ols.bse[1]
u = ols.resid
rho_resid = float(np.sum(u[1:]*u[:-1]) / np.sum(u**2))

# --- bandwidth: Andrews (1991) AR(1) plug-in for the Bartlett kernel, applied to the
# slope score v_t = (x_t - xbar) * u_t; S_T = 1.1447 * (alpha1 * T)^(1/3),
# alpha1 = 4 rho^2 / ((1-rho)^2 (1+rho)^2). Rounded up to an integer.
v = (x - x.mean()) * u
rho_v = float(np.sum(v[1:]*v[:-1]) / np.sum(v**2))
alpha1 = 4*rho_v**2 / ((1-rho_v)**2 * (1+rho_v)**2)
andrews_bw = 1.1447 * (alpha1 * n) ** (1/3)
hac_lags = int(math.ceil(andrews_bw))
nw_rule = int(math.floor(4 * (n/100) ** (2/9)))   # Newey-West (1994) rule of thumb, for reference

def hac(L):
    r = sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": L, "kernel": "bartlett", "use_correction": True})
    return r
r = hac(hac_lags)
se_hac = float(r.bse[1])
t_hac = float(beta_hat / se_hac)
p_hac = float(2 * stats.norm.sf(abs(t_hac)))
z = stats.norm.ppf(0.975)
ci = (float(beta_hat - z*se_hac), float(beta_hat + z*se_hac))

sens = []
for L in sorted(set([0, 1, 2, 3, nw_rule, 6, 8, hac_lags, 12, 16, 20, 24])):
    rr = hac(L); se = float(rr.bse[1]); t = beta_hat/se
    sens.append(dict(lags=L, se_hac=se, t=t, p_normal=float(2*stats.norm.sf(abs(t)))))
sens = pd.DataFrame(sens); sens.to_csv("out/hac_sensitivity.csv", index=False)

res = {
    "task_id": "F2-conf-02",
    "beta_hat": float(beta_hat), "alpha_hat": float(alpha_hat),
    "se_iid": float(se_iid), "t_iid": float(ols.tvalues[1]), "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[0], "beta_ci_high": ci[1],
    "memo_claim_supported": bool(p_hac < 0.05),
    "n_analysed": int(n),
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"series.csv": int(n_raw)},
    "random_seed": None,
    "bandwidth_details": {"rule": "Andrews (1991) AR(1) plug-in, Bartlett kernel, on slope score (x-xbar)*u, rounded up",
                          "rho_score": rho_v, "andrews_bw_raw": andrews_bw, "nw1994_rule_of_thumb_lags": nw_rule,
                          "small_sample_correction": "statsmodels use_correction=True (n/(n-k))",
                          "p_value_reference": "standard normal"},
    "hac_sensitivity": sens.to_dict(orient="records"),
    "sensitivity_all_p_above_005_for_lags_ge_1": bool((sens.loc[sens.lags >= 1, "p_normal"] > 0.05).all()),
}
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps({k: res[k] for k in ["beta_hat","alpha_hat","se_iid","rho_resid","se_hac","hac_lags","t_hac","p_hac","beta_ci_low","beta_ci_high","memo_claim_supported"]}, indent=1))
print(sens.to_string(index=False))