"""F2-conf-02: re-estimate OLS slope of y on x with a 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 and out/exclusions.csv. No randomness."""
import json, math
import numpy as np, pandas as pd, statsmodels.api as sm
from scipy import stats

df = pd.read_csv("data/series.csv")
n_raw = len(df)
issues = []
excl = []
# --- audit / filters (none expected to remove rows) ---
n0 = len(df); df = df.dropna(subset=["x", "y"]); excl.append(("1", "drop rows with missing x or y", n0, len(df), n0 - len(df), "OLS requires complete cases"))
n0 = len(df); df = df.drop_duplicates(subset=["month"]); excl.append(("2", "drop duplicate months", n0, len(df), n0 - len(df), "one observation per month"))
per = pd.PeriodIndex(df["month"], freq="M")
gaps = int((np.diff(per.asi8) != 1).sum())
if gaps:
    issues.append({"issue": f"{gaps} non-consecutive month steps", "handling": "kept; HAC lags interpreted as observation lags"})
pd.DataFrame(excl, columns=["step", "rule", "n_before", "n_after", "n_removed", "rationale"]).to_csv("out/exclusions.csv", index=False)

y = df["y"].to_numpy(float); x = df["x"].to_numpy(float); T = len(y)
X = sm.add_constant(x)
ols = sm.OLS(y, X).fit()
alpha_hat, beta_hat = float(ols.params[0]), float(ols.params[1])
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 = (x - x.mean()) * u
rv = float(np.sum(v[1:] * v[:-1]) / np.sum(v[:-1] ** 2))
a1 = 4 * rv**2 / ((1 - rv) ** 2 * (1 + rv) ** 2)
andrews_bw = 1.1447 * (a1 * T) ** (1 / 3)
hac_lags = int(math.floor(andrews_bw))
nw_rule_lags = int(math.floor(4 * (T / 100) ** (2 / 9)))

def hac_fit(L):
    return sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": L, "kernel": "bartlett", "use_correction": False})

hac = hac_fit(hac_lags)
se_hac = float(hac.bse[1])

# manual cross-check of the sandwich (Bartlett weights 1 - j/(L+1))
Z = X * u[:, None]
S = Z.T @ Z
for j in range(1, hac_lags + 1):
    G = Z[j:].T @ Z[:-j]
    S += (1 - j / (hac_lags + 1)) * (G + G.T)
XtXi = np.linalg.inv(X.T @ X)
se_manual = float(np.sqrt((XtXi @ S @ XtXi)[1, 1]))
assert abs(se_manual - se_hac) / se_hac < 1e-8, (se_manual, se_hac)

t_hac = beta_hat / se_hac
p_hac = float(2 * stats.norm.sf(abs(t_hac)))
z = float(stats.norm.ppf(0.975))
ci_low, ci_high = beta_hat - z * se_hac, beta_hat + z * se_hac

sens = {}
for L in sorted(set([0, nw_rule_lags, hac_lags, 12, 16, 24])):
    f = hac_fit(L); se = float(f.bse[1]); t = beta_hat / se
    sens[str(L)] = {"se": se, "t": t, "p": float(2 * stats.norm.sf(abs(t)))}

res = {
    "task_id": "F2-conf-02",
    "beta_hat": beta_hat, "alpha_hat": alpha_hat,
    "se_iid": 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": float(t_hac), "p_hac": p_hac,
    "beta_ci_low": float(ci_low), "beta_ci_high": float(ci_high),
    "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,
    "baseline_se_iid_t_p": [se_iid, float(ols.tvalues[1]), float(ols.pvalues[1])],
    "hac_bandwidth_details": {"method": "Andrews (1991) AR(1) plug-in, Bartlett kernel, on score (x-xbar)*u",
                              "rho_score": rv, "andrews_bw_continuous": andrews_bw,
                              "nw_rule_of_thumb_lags": nw_rule_lags, "small_sample_correction": False,
                              "reference_distribution": "standard normal"},
    "hac_lag_sensitivity": sens,
    "hac_lag_sensitivity_min_p_excl_0": min(v_["p"] for k, v_ in sens.items() if k != "0"),
}
with open("out/results.json", "w") as fh:
    json.dump(res, fh, indent=2)
print(f"beta={beta_hat:.4f} se_iid={se_iid:.4f} se_hac={se_hac:.4f} (L={hac_lags}) t={t_hac:.3f} p={p_hac:.4f} "
      f"CI=[{ci_low:.4f},{ci_high:.4f}] memo_supported={p_hac < 0.05}")