"""F2-conf-02: OLS slope of y on x with Newey-West (Bartlett) HAC inference.
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, os
import numpy as np, pandas as pd
from scipy import stats
import statsmodels.api as sm

DATA = "data/series.csv"
os.makedirs("out", exist_ok=True)

raw = pd.read_csv(DATA, dtype={"month": str})
n_raw = len(raw)
issues = []
excl = []

# Audit: missing values, duplicates, gaps
d = raw.copy()
nb = len(d); d = d.dropna(subset=["x", "y"])
excl.append(("1", "drop rows with missing x or y", nb, len(d), nb - len(d), "OLS requires complete cases"))
nb = len(d); d = d.drop_duplicates(subset=["month"])
excl.append(("2", "drop duplicated months", nb, len(d), nb - len(d), "one observation per month"))
per = pd.PeriodIndex(d["month"], freq="M")
gaps = int(((per[1:] - per[:-1]).map(lambda z: z.n) != 1).sum())
if gaps:
    issues.append({"issue": f"{gaps} gaps/non-consecutive months", "handling": "kept; HAC lags computed on row order"})
if n_raw - len(d):
    issues.append({"issue": f"{n_raw-len(d)} rows removed (missing/duplicate)", "handling": "excluded, see out/exclusions.csv"})
d = d.sort_values("month").reset_index(drop=True)

y = d["y"].to_numpy(float); x = d["x"].to_numpy(float); T = len(y)
X = np.column_stack([np.ones(T), x])
XtX_inv = np.linalg.inv(X.T @ X)
b = XtX_inv @ X.T @ y
u = y - X @ b
k = X.shape[1]

# Naive iid SE (memo's measure)
s2 = u @ u / (T - k)
se_iid = math.sqrt(s2 * XtX_inv[1, 1])

# First-order residual autocorrelation (OLS of u_t on u_{t-1}, no intercept; residuals have mean 0)
rho_resid = float(np.sum(u[1:] * u[:-1]) / np.sum(u[:-1] ** 2))

def nw_se(L):
    """Newey-West Bartlett-kernel HAC SE of the slope, no small-sample df correction."""
    g = X * u[:, None]
    S = g.T @ g
    for j in range(1, L + 1):
        w = 1 - j / (L + 1)
        G = g[j:].T @ g[:-j]
        S += w * (G + G.T)
    V = XtX_inv @ S @ XtX_inv
    return math.sqrt(V[1, 1])

# 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 * (alpha(1) * T)^(1/3), alpha(1) = 4 r^2 / ((1-r)^2 (1+r)^2)
v = (x - x.mean()) * u
r_v = float(np.sum(v[1:] * v[:-1]) / np.sum(v[:-1] ** 2))
alpha1 = 4 * r_v ** 2 / ((1 - r_v) ** 2 * (1 + r_v) ** 2)
andrews_bw = 1.1447 * (alpha1 * T) ** (1 / 3)
hac_lags = int(math.floor(andrews_bw))
nw_rule = int(math.floor(4 * (T / 100) ** (2 / 9)))  # Newey-West (1994) rule of thumb, for sensitivity

se_hac = nw_se(hac_lags)
t_hac = float(b[1] / se_hac)
p_hac = float(2 * stats.norm.sf(abs(t_hac)))
z = stats.norm.ppf(0.975)
ci_low, ci_high = float(b[1] - z * se_hac), float(b[1] + z * se_hac)

# Cross-check against statsmodels (use_correction=False -> no df correction)
sm_fit = sm.OLS(y, X).fit(cov_type="HAC", cov_kwds={"maxlags": hac_lags, "use_correction": False})
assert abs(sm_fit.bse[1] - se_hac) < 1e-10, (sm_fit.bse[1], se_hac)

sens = {}
for L in sorted(set([0, nw_rule, hac_lags, 12, 24])):
    s = nw_se(L); t = b[1] / s
    sens[f"lags_{L}"] = {"se": s, "t": float(t), "p_normal": float(2 * stats.norm.sf(abs(t))),
                         "p_t_df": float(2 * stats.t.sf(abs(t), T - k))}
p_iid = float(2 * stats.t.sf(abs(b[1] / se_iid), T - k))

memo_claim_supported = bool(p_hac < 0.05)

res = {
    "task_id": "F2-conf-02",
    "beta_hat": float(b[1]), "alpha_hat": float(b[0]),
    "se_iid": se_iid, "t_iid": float(b[1] / se_iid), "p_iid": p_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_low, "beta_ci_high": ci_high,
    "memo_claim_supported": memo_claim_supported,
    "n_analysed": int(T),
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"series.csv": int(n_raw)},
    "random_seed": None,
    "bandwidth_rule": "Andrews (1991) AR(1) plug-in, Bartlett kernel, on slope score (x-xbar)*u; floor",
    "andrews_bandwidth_raw": float(andrews_bw), "score_rho": r_v, "nw_rule_of_thumb_lags": nw_rule,
    "hac_sensitivity": sens,
}
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
pd.DataFrame(excl, columns=["step", "rule", "n_before", "n_after", "n_removed", "rationale"]).to_csv("out/exclusions.csv", index=False)
print(f"beta={b[1]:.4f} se_iid={se_iid:.4f} se_hac({hac_lags})={se_hac:.4f} t={t_hac:.3f} p={p_hac:.4f} supported={memo_claim_supported}")