"""F2-conf-03: Monte Carlo coverage of OLS CIs under heteroskedasticity.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py"""
import json, os
import numpy as np, pandas as pd
from scipy import stats

SEED = 20261030
ALPHA, BETA, R, LEVEL = 1.0, 0.5, 2000, 0.95

d = pd.read_csv("data/design.csv")
issues = []
n_rows = len(d)
assert list(d.columns) == ["i", "x", "sigma"]
if d.isna().any().any(): issues.append({"issue": "missing values", "handling": "none expected; would abort"})
if d["i"].duplicated().any(): issues.append({"issue": "duplicate indices", "handling": "reported"})
if (d["sigma"] <= 0).any(): issues.append({"issue": "non-positive sigma", "handling": "reported"})
x = d["x"].to_numpy(float); sig = d["sigma"].to_numpy(float); n = len(x)
assert n == 30

X = np.column_stack([np.ones(n), x])
XtXi = np.linalg.inv(X.T @ X)
A = XtXi @ X.T                      # 2 x n, beta_hat = A y
# exact variance of slope: sum a_i^2 sigma_i^2
sd_exact = float(np.sqrt(np.sum(A[1] ** 2 * sig ** 2)))

rng = np.random.default_rng(SEED)
eps = rng.standard_normal((R, n))
Y = ALPHA + BETA * x[None, :] + sig[None, :] * eps   # R x n
B = Y @ A.T                                           # R x 2
E = Y - B @ X.T                                       # residuals
s2 = (E ** 2).sum(1) / (n - 2)
se_conv = np.sqrt(s2 * XtXi[1, 1])
hc0 = (E ** 2) @ (A[1] ** 2)                          # sandwich slope element
se_hc1 = np.sqrt(hc0 * n / (n - 2))
tc = stats.t.ppf(0.5 + LEVEL / 2, n - 2)
bh = B[:, 1]
cov_conv = float(np.mean(np.abs(bh - BETA) <= tc * se_conv))
cov_hc1 = float(np.mean(np.abs(bh - BETA) <= tc * se_hc1))
mcse_nom = np.sqrt(LEVEL * (1 - LEVEL) / R)
undercov = bool(cov_conv < LEVEL - 3 * mcse_nom)

res = {
    "task_id": "F2-conf-03",
    "coverage_conventional": cov_conv,
    "coverage_hc1": cov_hc1,
    "sd_beta_exact": sd_exact,
    "sd_beta_mc": float(np.std(bh, ddof=1)),
    "mean_se_conventional": float(se_conv.mean()),
    "conventional_ci_undercovers": undercov,
    "n_replications": R,
    "n_obs_per_replication": n,
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"design.csv": n_rows},
    "random_seed": SEED,
}
os.makedirs("out", exist_ok=True)
with open("out/results.json", "w") as f: json.dump(res, f, indent=2)
# auxiliary diagnostics (printed, not part of results.json)
print(json.dumps(res, indent=2))
print("t crit", tc, "MCSE at 0.95", mcse_nom,
      "MCSE conv", np.sqrt(cov_conv*(1-cov_conv)/R), "MCSE hc1", np.sqrt(cov_hc1*(1-cov_hc1)/R),
      "mean se_hc1", se_hc1.mean(), "naive sd formula", np.sqrt(np.mean(sig**2)*XtXi[1,1]),
      "E[se_conv^2] exact sqrt", np.sqrt(XtXi[1,1]*np.sum(sig**2*np.diag(np.eye(n)-X@XtXi@X.T))/(n-2)))