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

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

d = pd.read_csv("data/design.csv")
n_rows = len(d)
issues = []
assert n_rows == N and d[["x", "sigma"]].notna().all().all() and (d.sigma > 0).all()
x = d["x"].to_numpy(float); sig = d["sigma"].to_numpy(float)

xc = x - x.mean(); Sxx = np.sum(xc**2)
w = xc / Sxx                       # beta_hat = beta + sum w_i sigma_i eps_i
sd_beta_exact = float(np.sqrt(np.sum(w**2 * sig**2)))
# exact E[s^2] (not E[se]): trace(M Sigma)/(n-2), reported as a diagnostic
X = np.column_stack([np.ones(N), x])
H = X @ np.linalg.solve(X.T @ X, X.T); M = np.eye(N) - H
se_conv_sqrt_Es2 = float(np.sqrt(np.trace(M * sig**2) / (N - 2) / Sxx))

rng = np.random.default_rng(SEED)
eps = rng.standard_normal((R, N))          # only eps redrawn; x, sigma fixed
Y = ALPHA + BETA * x + sig * eps           # R x N
yc = Y - Y.mean(axis=1, keepdims=True)
b = yc @ xc / Sxx
a = Y.mean(axis=1) - b * x.mean()
e = Y - a[:, None] - b[:, None] * x
se_conv = np.sqrt(np.sum(e**2, axis=1) / (N - 2) / Sxx)
hc0 = np.sum((xc**2) * e**2, axis=1) / Sxx**2   # slope element of HC0 sandwich
se_hc1 = np.sqrt(hc0 * N / (N - 2))
tcrit = stats.t.ppf(0.5 + LEVEL / 2, N - 2)
cov_conv = float(np.mean(np.abs(b - BETA) <= tcrit * se_conv))
cov_hc1 = float(np.mean(np.abs(b - BETA) <= tcrit * se_hc1))

mcse_nominal = float(np.sqrt(LEVEL * (1 - LEVEL) / R))
mcse_conv = float(np.sqrt(cov_conv * (1 - cov_conv) / R))
mcse_hc1 = float(np.sqrt(cov_hc1 * (1 - cov_hc1) / R))
undercovers = bool(cov_conv < LEVEL - 3 * mcse_nominal)

# cross-check one replication with statsmodels
try:
    import statsmodels.api as sm
    f = sm.OLS(Y[0], X).fit()
    assert abs(f.bse[1] - se_conv[0]) < 1e-10
    assert abs(f.get_robustcov_results("HC1").bse[1] - se_hc1[0]) < 1e-10
    sm_check = True
except ImportError:
    sm_check = None

res = {
    "task_id": "F2-conf-03",
    "coverage_conventional": cov_conv,
    "coverage_hc1": cov_hc1,
    "sd_beta_exact": sd_beta_exact,
    "sd_beta_mc": float(np.std(b, ddof=1)),
    "mean_se_conventional": float(np.mean(se_conv)),
    "conventional_ci_undercovers": undercovers,
    "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,
    # supplementary
    "t_critical": float(tcrit),
    "mcse_at_nominal": mcse_nominal,
    "mcse_coverage_conventional": mcse_conv,
    "mcse_coverage_hc1": mcse_hc1,
    "undercover_threshold": LEVEL - 3 * mcse_nominal,
    "conventional_ci_undercovers_using_estimated_mcse": bool(cov_conv < LEVEL - 3 * mcse_conv),
    "hc1_ci_undercovers": bool(cov_hc1 < LEVEL - 3 * mcse_nominal),
    "mean_se_hc1": float(np.mean(se_hc1)),
    "mean_beta_hat": float(np.mean(b)),
    "mcse_sd_beta_mc_approx": float(np.std(b, ddof=1) / np.sqrt(2 * (R - 1))),
    "sqrt_expected_conv_variance_exact": se_conv_sqrt_Es2,
    "ratio_mean_se_conv_to_sd_exact": float(np.mean(se_conv) / sd_beta_exact),
    "statsmodels_crosscheck_passed": sm_check,
}
os.makedirs("out", exist_ok=True)
with open("out/results.json", "w") as fh:
    json.dump(res, fh, indent=2)
print(json.dumps(res, indent=2))