"""F2-conf-02: OLS slope of y on x with Newey-West HAC inference. Run from workspace root."""
import json, math, os
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)
d = d.dropna(subset=["x", "y"])
months = pd.PeriodIndex(d["month"], freq="M")
assert (np.diff(months.asi8) == 1).all(), "month gaps"
n = len(d)
X = sm.add_constant(d["x"].to_numpy()); y = d["y"].to_numpy()

# P2 baseline: conventional iid OLS (memo's measure)
ols = sm.OLS(y, X).fit()
alpha_hat, beta_hat = map(float, ols.params)
se_iid = float(ols.bse[1]); t_iid = beta_hat / se_iid
p_iid = float(2 * stats.t.sf(abs(t_iid), n - 2))
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:  S_T = 1.1447 * (a1 * T)^(1/3),
# a1 = 4 rho^2 / ((1-rho)^2 (1+rho)^2). Rounded up to an integer.
v = (d["x"].to_numpy() - d["x"].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)
andrews_bw = 1.1447 * (a1 * 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

def hac(L):
    r = sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": L, "kernel": "bartlett", "use_correction": True})
    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)
ci = (beta_hat - z * se_hac, beta_hat + z * se_hac)

sens = {}
for L in sorted({nw_rule, 6, 8, 12, 18, 24, hac_lags}):
    s = hac(L); sens[str(L)] = {"se_hac": s, "t_hac": beta_hat / s, "p_hac": float(2 * stats.norm.sf(abs(beta_hat / s)))}

res = {
    "task_id": "F2-conf-02",
    "beta_hat": beta_hat, "alpha_hat": alpha_hat, "se_iid": se_iid,
    "baseline_t_iid": t_iid, "baseline_p_iid": p_iid,
    "rho_resid": rho_resid, "rho_score": rho_v, "andrews_bandwidth_raw": andrews_bw,
    "nw_rule_of_thumb_lags": nw_rule,
    "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],
    "se_ratio_hac_to_iid": se_hac / se_iid,
    "memo_claim_supported": bool(p_hac < 0.05),
    "hac_sensitivity": sens,
    "n_analysed": int(n), "feasibility": "complete", "feasibility_reason": None,
    "data_issues": [], "n_input_rows": {"series.csv": int(n_raw)}, "random_seed": None,
}
os.makedirs("out", exist_ok=True)
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))