"""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
Reads only data/design.csv; writes out/results.json deterministically."""
import json, os
import numpy as np
import pandas as pd
from scipy import stats
import statsmodels.api as sm

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

d = pd.read_csv("data/design.csv")
n_rows = int(len(d))
data_issues = []
assert n_rows == 30 and d[["x", "sigma"]].notna().all().all() and (d.sigma > 0).all()
x = d["x"].to_numpy(float)
sig = d["sigma"].to_numpy(float)
n = len(x)
df = n - 2
tcrit = stats.t.ppf(0.5 + LEVEL / 2, df)

# (c) exact SD of betahat: betahat - beta = sum w_i sigma_i eps_i, w_i = (x_i - xbar)/Sxx
xc = x - x.mean()
Sxx = np.sum(xc ** 2)
w = xc / Sxx
sd_exact = float(np.sqrt(np.sum(w ** 2 * sig ** 2)))

# Monte Carlo (only eps redrawn; x, sigma fixed)
rng = np.random.default_rng(SEED)
eps = rng.standard_normal((R, n))
Y = ALPHA + BETA * x[None, :] + sig[None, :] * eps
X = np.column_stack([np.ones(n), x])
XtX_inv = np.linalg.inv(X.T @ X)
P = XtX_inv @ X.T                   # 2 x n
B = Y @ P.T                         # R x 2 coefficients
resid = Y - B @ X.T                 # R x n
s2 = np.sum(resid ** 2, axis=1) / df
se_conv = np.sqrt(s2 * XtX_inv[1, 1])
a = P[1]                            # slope row of (X'X)^-1 X'
var_hc0 = np.sum((a[None, :] ** 2) * resid ** 2, axis=1)
se_hc1 = np.sqrt(var_hc0 * n / df)
bhat = B[:, 1]

cov_conv = float(np.mean(np.abs(bhat - BETA) <= tcrit * se_conv))
cov_hc1 = float(np.mean(np.abs(bhat - BETA) <= tcrit * se_hc1))
mcse_nominal = float(np.sqrt(LEVEL * (1 - LEVEL) / R))
undercovers = bool(cov_conv < LEVEL - 3 * mcse_nominal)

# Cross-check against statsmodels on replication 0
m0 = sm.OLS(Y[0], X).fit()
assert np.isclose(m0.bse[1], se_conv[0])
assert np.isclose(m0.get_robustcov_results("HC1").bse[1], se_hc1[0])
assert np.isclose(np.sqrt(np.sum(a ** 2 * sig ** 2)), sd_exact)

H = np.diag(X @ XtX_inv @ X.T)
results = {
    "task_id": "F2-conf-03",
    "coverage_conventional": cov_conv,
    "coverage_hc1": cov_hc1,
    "sd_beta_exact": sd_exact,
    "sd_beta_mc": float(np.std(bhat, 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": data_issues,
    "n_input_rows": {"design.csv": n_rows},
    "random_seed": SEED,
    # supplementary diagnostics
    "t_critical_28": float(tcrit),
    "mcse_coverage_conventional": float(np.sqrt(cov_conv * (1 - cov_conv) / R)),
    "mcse_coverage_hc1": float(np.sqrt(cov_hc1 * (1 - cov_hc1) / R)),
    "mcse_at_nominal_095": mcse_nominal,
    "undercoverage_threshold": float(LEVEL - 3 * mcse_nominal),
    "mean_se_hc1": float(np.mean(se_hc1)),
    "mean_beta_hat": float(np.mean(bhat)),
    "sd_se_conventional_mc": float(np.std(se_conv, ddof=1)),
    "sqrt_expected_se_conv_sq_exact": float(np.sqrt(XtX_inv[1, 1] * np.sum(sig ** 2 * (1 - H)) / df)),
}
os.makedirs("out", exist_ok=True)
with open("out/results.json", "w") as f:
    json.dump(results, f, indent=2)
print(json.dumps(results, indent=2))