"""F2-conf-02: OLS slope of y on x with Newey-West HAC SE. Reads data/ only; writes out/results.json."""
import json, os, numpy as np, pandas as pd
import statsmodels.api as sm
from scipy import stats

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
d = pd.read_csv(os.path.join(ROOT, "data", "series.csv"))
n_raw = len(d)
issues = []
# audit checks (no rows removed unless a problem is found)
per = pd.PeriodIndex(d["month"], freq="M")
assert d[["x", "y"]].notna().all().all(), "missing values"
assert not d["month"].duplicated().any(), "duplicate months"
assert (np.diff(per.asi8) == 1).all(), "gaps in months"
d = d.sort_values("month").reset_index(drop=True)
n = len(d)

X = sm.add_constant(d["x"].to_numpy())
y = d["y"].to_numpy()
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 ** 2))

# Bandwidth: Andrews (1991) AR(1) plug-in for the Bartlett kernel, applied to the
# score of the slope v_t = (x_t - xbar) u_t.  S = 1.1447 (a1 * T)^(1/3),
# a1 = 4 r^2 / ((1-r)^2 (1+r)^2), r = AR(1) coefficient of v_t.  Truncated to integer lags.
v = (d["x"].to_numpy() - d["x"].mean()) * u
r_v = float(np.sum(v[1:] * v[:-1]) / np.sum(v[:-1] ** 2))
a1 = 4 * r_v ** 2 / ((1 - r_v) ** 2 * (1 + r_v) ** 2)
S_andrews = 1.1447 * (a1 * n) ** (1 / 3)
hac_lags = int(np.floor(S_andrews))
nw_rule = int(np.floor(4 * (n / 100) ** (2 / 9)))  # Newey-West (1994) rule-of-thumb, for sensitivity

def hac(L):
    f = sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": L, "kernel": "bartlett", "use_correction": True})
    se = float(f.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_low, ci_high = 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, 8, 12, 16, 24, hac_lags})}
memo_supported = bool(p_hac < 0.05)

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, "rho_score": r_v, "andrews_bandwidth": float(S_andrews),
    "se_hac": se_hac, "hac_lags": hac_lags, "t_hac": t_hac, "p_hac": p_hac,
    "beta_ci_low": float(ci_low), "beta_ci_high": float(ci_high),
    "se_ratio_hac_to_iid": se_hac / se_iid,
    "hac_sensitivity": sens, "nw_rule_of_thumb_lags": nw_rule,
    "sensitivity_all_insignificant": bool(all(s["p"] >= 0.05 for s in sens.values())),
    "memo_claim_supported": memo_supported,
    "n_analysed": int(n), "feasibility": "complete", "feasibility_reason": None,
    "data_issues": issues, "n_input_rows": {"series.csv": int(n_raw)}, "random_seed": None,
}
os.makedirs(os.path.join(ROOT, "out"), exist_ok=True)
with open(os.path.join(ROOT, "out", "results.json"), "w") as fh:
    json.dump(res, fh, indent=2)
pd.DataFrame([{"step": "S0", "rule": "none (no missing, duplicate or gap rows)", "n_before": n_raw,
               "n_after": n, "n_removed": n_raw - n, "rationale": "data audit found no problems"}]
             ).to_csv(os.path.join(ROOT, "out", "exclusions.csv"), index=False)
print(f"beta={beta_hat:.4f} se_iid={se_iid:.4f} se_hac={se_hac:.4f} (L={hac_lags}) t={t_hac:.2f} p={p_hac:.4f} memo_supported={memo_supported}")