"""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; deterministic (no randomness)."""
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 = []
# audit checks
assert d[["month", "x", "y"]].notna().all().all(), "missing values"
per = pd.PeriodIndex(d["month"], freq="M")
assert per.is_monotonic_increasing and per.is_unique
gaps = int(((per[1:] - per[:-1]).map(lambda z: z.n) != 1).sum())
assert gaps == 0

X = sm.add_constant(d["x"].to_numpy())
y = d["y"].to_numpy()
T = len(y)
ols = sm.OLS(y, X).fit()
alpha_hat, beta_hat = map(float, ols.params)
se_iid = float(ols.bse[1])
u = ols.resid
rho = 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. Chosen over the fixed NW rule-of-thumb
# floor(4(T/100)^(2/9)) because errors (rho~0.7) and x are both persistent,
# which makes the score persistent and the rule of thumb too short.
v = (d["x"].to_numpy() - d["x"].mean()) * u
rv = float(np.sum(v[1:] * v[:-1]) / np.sum(v ** 2))
a1 = 4 * rv ** 2 / ((1 - rv) ** 2 * (1 + rv) ** 2)
S_andrews = 1.1447 * (a1 * T) ** (1 / 3)
hac_lags = int(math.ceil(S_andrews))
nw_rule = int(math.floor(4 * (T / 100) ** (2 / 9)))

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

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

sens = {}
for L in sorted({0, nw_rule, 8, 12, hac_lags, 24, 36}):
    s = hac(L)
    sens[str(L)] = {"se": s, "t": beta_hat / s, "p": 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,
    "t_iid": beta_hat / se_iid,
    "p_iid": float(2 * stats.t.sf(abs(beta_hat / se_iid), T - 2)),
    "rho_resid": rho,
    "rho_score": rv,
    "andrews_bandwidth_raw": S_andrews,
    "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": lo,
    "beta_ci_high": hi,
    "hac_sensitivity": sens,
    "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,
}
with open("out/results.json", "w") as fh:
    json.dump(res, fh, indent=2)
print(json.dumps({k: res[k] for k in ["beta_hat", "se_iid", "rho_resid", "hac_lags", "se_hac", "t_hac", "p_hac", "memo_claim_supported"]}, indent=2))