"""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, csv
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

d = pd.read_csv("data/design.csv")
n_raw = len(d)
issues = []
# audit checks (no filtering needed if all pass)
assert list(d.columns) == ["i", "x", "sigma"]
assert not d.isna().any().any(), "missing values in design"
assert n_raw == N, f"expected {N} rows, got {n_raw}"
assert (d["sigma"] > 0).all()
assert d["i"].tolist() == list(range(1, N + 1))
assert d["x"].std() > 0
x = d["x"].to_numpy(float); sig = d["sigma"].to_numpy(float)

os.makedirs("out", exist_ok=True)
with open("out/exclusions.csv", "w", newline="") as f:
    w = csv.writer(f)
    w.writerow(["step", "rule", "n_before", "n_after", "n_removed", "rationale"])
    w.writerow(["1", "drop rows with missing/invalid x or sigma", n_raw, n_raw, 0,
                "audit found no missing, duplicate or non-positive values; nothing removed"])

X = np.column_stack([np.ones(N), x])
XtX_inv = np.linalg.inv(X.T @ X)
xc = x - x.mean(); Sxx = xc @ xc
wts = xc / Sxx                     # beta_hat - beta = sum w_i sigma_i eps_i
sd_beta_exact = float(np.sqrt(np.sum(wts**2 * sig**2)))
tcrit = stats.t.ppf(0.5 + LEVEL / 2, N - 2)

rng = np.random.default_rng(SEED)
eps = rng.standard_normal((R, N))           # one row per replication
Y = ALPHA + BETA * x[None, :] + sig[None, :] * eps
A = X @ XtX_inv                             # (N,2): beta_hat = Y @ A
B = Y @ A
resid = Y - B @ X.T
s2 = (resid**2).sum(1) / (N - 2)
se_conv = np.sqrt(s2 * XtX_inv[1, 1])
# HC0 slope variance: sum_i a_i1^2 e_i^2  (= [(X'X)^-1 X' diag(e^2) X (X'X)^-1]_{11})
hc0 = (resid**2) @ (A[:, 1]**2)
se_hc1 = np.sqrt(hc0 * N / (N - 2))
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)

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(se_conv.mean()),
    "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": int(n_raw)},
    "random_seed": SEED,
    # supplementary keys
    "t_critical_28": float(tcrit),
    "mcse_coverage_at_nominal": mcse_nominal,
    "undercoverage_threshold": float(LEVEL - 3 * mcse_nominal),
    "mcse_coverage_conventional": float(np.sqrt(cov_conv * (1 - cov_conv) / R)),
    "mcse_coverage_hc1": float(np.sqrt(cov_hc1 * (1 - cov_hc1) / R)),
    "sd_beta_homoskedastic_formula": float(np.sqrt(np.mean(sig**2) / Sxx)),
    "mean_se_hc1": float(se_hc1.mean()),
    "mean_beta_hat": float(bhat.mean()),
    "mcse_mean_beta_hat": float(np.std(bhat, ddof=1) / np.sqrt(R)),
    "mcse_mean_se_conventional": float(np.std(se_conv, ddof=1) / np.sqrt(R)),
}
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))