"""F2-conf-03: Monte Carlo coverage of OLS slope CIs under heteroskedasticity (fixed design).
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Reads only data/design.csv; writes out/results.json deterministically."""
import json
from pathlib import Path
import numpy as np
import pandas as pd
from scipy import stats

ALPHA, BETA, N, R, LEVEL = 1.0, 0.5, 30, 2000, 0.95
SEED = 20261030  # fixed seed for numpy.random.default_rng (PCG64)

d = pd.read_csv("data/design.csv")
n_rows = int(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 = float(xc @ xc)
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)))
sd_beta_conv_formula = float(np.sqrt(np.mean(sig**2) / Sxx))  # naive homoskedastic SD (diagnostic only)

tcrit = float(stats.t.ppf(0.5 + LEVEL / 2, N - 2))

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)
B = Y @ X @ XtX_inv                            # (R,2) OLS coefficients (XtX_inv symmetric)
bhat = B[:, 1]
E = Y - B @ X.T                                # residuals
s2 = np.sum(E**2, axis=1) / (N - 2)
se_conv = np.sqrt(s2 * XtX_inv[1, 1])
# HC0 sandwich for slope: (X'X)^-1 X' diag(e^2) X (X'X)^-1, element [1,1]; HC1 = HC0 * n/(n-2)
A = X @ XtX_inv                                # row i: a_i' ; slope weight = A[:,1] (= w)
hc0 = (E**2) @ (A[:, 1] ** 2)
se_hc1 = np.sqrt(hc0 * N / (N - 2))

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))
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)

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(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": issues,
    "n_input_rows": {"design.csv": n_rows},
    "random_seed": SEED,
    # supplementary (not required) diagnostics
    "t_critical_28": tcrit,
    "mcse_binomial_nominal": mcse_nominal,
    "mcse_coverage_conventional": mcse_conv,
    "mcse_coverage_hc1": mcse_hc1,
    "undercoverage_threshold": LEVEL - 3 * mcse_nominal,
    "mean_se_hc1": float(np.mean(se_hc1)),
    "sd_beta_homoskedastic_formula": sd_beta_conv_formula,
    "mean_beta_hat_mc": float(np.mean(bhat)),
}
Path("out").mkdir(exist_ok=True)
Path("out/results.json").write_text(json.dumps(res, indent=2))
for k, v in res.items():
    print(f"{k}: {v}")