"""F2-conf-02: OLS slope of y on x with Newey-West HAC inference.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Reads only data/series.csv; writes out/results.json. No randomness used."""
import json, math
from pathlib import Path
import numpy as np
import pandas as pd
from scipy import stats
import statsmodels.api as sm

ROOT = Path(__file__).resolve().parent.parent
OUT = ROOT / "out"

df = pd.read_csv(ROOT / "data" / "series.csv", dtype={"month": str})
n_raw = len(df)
issues = []
# data checks
per = pd.PeriodIndex(df["month"], freq="M")
assert per.is_monotonic_increasing and not per.duplicated().any()
assert ((per[1:] - per[:-1]).map(lambda d: d.n) == 1).all(), "gaps in months"
d = df.dropna(subset=["x", "y"])
assert len(d) == n_raw
T = 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])
t_iid = float(ols.tvalues[1]); p_iid = float(ols.pvalues[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 = 1.1447 * (a1 * T)^(1/3),
# a1 = 4 rho_v^2 / ((1-rho_v)^2 (1+rho_v)^2). Rounded up to an integer lag count.
v = (d["x"].to_numpy() - d["x"].mean()) * u
rho_v = float(np.sum(v[1:] * v[:-1]) / np.sum(v[:-1] ** 2))
a1 = 4 * rho_v ** 2 / ((1 - rho_v) ** 2 * (1 + rho_v) ** 2)
andrews_bw = 1.1447 * (a1 * T) ** (1 / 3)
nw_rule = int(math.floor(4 * (T / 100) ** (2 / 9)))  # Newey-West (1994) rule of thumb
hac_lags = int(max(math.ceil(andrews_bw), nw_rule))

def hac(L):
    r = sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": L, "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 = float(stats.norm.ppf(0.975))
ci_lo, ci_hi = beta_hat - z * se_hac, beta_hat + z * se_hac

# sensitivity of the verdict to the bandwidth (reported, not used for the primary key)
sens = {}
for L in sorted(set([nw_rule, hac_lags, 8, 12, 18, 24, int(math.floor(0.75 * T ** (1 / 3)))])):
    s = hac(L); t = beta_hat / s
    sens[str(L)] = {"se_hac": s, "t_hac": t, "p_hac": 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,
    "rho_resid": rho_resid,
    "se_hac": se_hac,
    "hac_lags": hac_lags,
    "t_hac": t_hac,
    "p_hac": p_hac,
    "beta_ci_low": ci_lo,
    "beta_ci_high": ci_hi,
    "memo_claim_supported": bool(p_hac < 0.05),
    "n_analysed": T,
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"series.csv": n_raw},
    "random_seed": None,
    # supplementary keys
    "baseline_t_iid": t_iid,
    "baseline_p_iid": p_iid,
    "rho_score": rho_v,
    "andrews_bandwidth_raw": andrews_bw,
    "nw1994_rule_lags": nw_rule,
    "se_ratio_hac_over_iid": se_hac / se_iid,
    "hac_sensitivity_by_lag": sens,
}
OUT.mkdir(exist_ok=True)
(OUT / "results.json").write_text(json.dumps(res, indent=2))
print(f"beta={beta_hat:.4f} alpha={alpha_hat:.4f} se_iid={se_iid:.4f} (t={t_iid:.2f}, p={p_iid:.4f})")
print(f"rho_resid={rho_resid:.3f} rho_score={rho_v:.3f} lags={hac_lags} se_hac={se_hac:.4f} t={t_hac:.2f} p={p_hac:.4f}")
print(f"CI=[{ci_lo:.4f},{ci_hi:.4f}] memo_claim_supported={res['memo_claim_supported']}")
for L, s in sens.items(): print(L, {k: round(v, 4) for k, v in s.items()})