"""F2-conf-03: Monte Carlo coverage of OLS 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 and out/exclusions.csv. Seed fixed."""
import json, hashlib, os
import numpy as np, pandas as pd
from scipy import stats

SEED = 20261030
ALPHA, BETA, N, R, LEVEL = 1.0, 0.5, 30, 2000, 0.95
os.makedirs("out", exist_ok=True)

path = "data/design.csv"
raw = pd.read_csv(path)
n_raw = len(raw)
issues = []
# audit checks (no filtering expected)
d = raw.copy()
excl = []
n0 = len(d)
d = d.dropna(subset=["x", "sigma"])
excl.append(("1", "drop rows with missing x or sigma", n0, len(d), n0 - len(d), "model requires x and sigma"))
n0 = len(d)
d = d[d["sigma"] > 0]
excl.append(("2", "drop rows with sigma <= 0", n0, len(d), n0 - len(d), "sigma is a standard deviation"))
n0 = len(d)
d = d.drop_duplicates(subset=["i"])
excl.append(("3", "drop duplicate observation index i", n0, len(d), n0 - len(d), "each i appears once in design"))
pd.DataFrame(excl, columns=["step", "rule", "n_before", "n_after", "n_removed", "rationale"]).to_csv("out/exclusions.csv", index=False)
for s in excl:
    if s[4] > 0:
        issues.append({"issue": f"{s[4]} rows removed by rule '{s[1]}'", "handling": "excluded"})
if len(d) != N:
    issues.append({"issue": f"design has {len(d)} usable rows, expected {N}", "handling": "reported; feasibility partial"})

d = d.sort_values("i")
x = d["x"].to_numpy(float); sig = d["sigma"].to_numpy(float); n = len(x)
X = np.column_stack([np.ones(n), x])
XtX_inv = np.linalg.inv(X.T @ X)
H = XtX_inv @ X.T                      # 2 x n, beta_hat = H y
w = H[1]                               # slope weights
sd_exact = float(np.sqrt(np.sum(w**2 * sig**2)))  # exact: Var(b) = sum w_i^2 sigma_i^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 @ H.T                                           # R x 2
resid = Y - B @ X.T
s2 = np.sum(resid**2, axis=1) / (n - 2)
se_conv = np.sqrt(s2 * XtX_inv[1, 1])
hc0 = np.sum((w[None, :] ** 2) * resid**2, axis=1)    # sandwich [ (X'X)^-1 X' diag(e^2) X (X'X)^-1 ]_{11}
se_hc1 = np.sqrt(hc0 * n / (n - 2))
tcrit = stats.t.ppf(0.5 + LEVEL / 2, n - 2)
b = B[:, 1]
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)
# exact expectation of conventional variance (for checking mean_se_conventional): E[s2]=tr(M Sigma)/(n-2)
M = np.eye(n) - X @ H
E_s2 = float(np.trace(M @ np.diag(sig**2)) / (n - 2))

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(b, ddof=1)),
    "mean_se_conventional": float(np.mean(se_conv)),
    "conventional_ci_undercovers": undercovers,
    "n_replications": R,
    "n_obs_per_replication": int(n),
    "feasibility": "complete" if n == N else "partial",
    "feasibility_reason": None if n == N else "design row count differs from 30",
    "data_issues": issues,
    "n_input_rows": {"design.csv": int(n_raw)},
    "random_seed": SEED,
    # auxiliary
    "t_critical_28": float(tcrit),
    "mcse_coverage_conventional": mcse_conv,
    "mcse_coverage_hc1": mcse_hc1,
    "mcse_at_nominal_0.95": mcse_nominal,
    "undercover_threshold": float(LEVEL - 3 * mcse_nominal),
    "undercovers_using_estimated_p_se": bool(cov_conv < LEVEL - 3 * mcse_conv),
    "mean_se_hc1": float(np.mean(se_hc1)),
    "se_conv_population_sqrtEs2": float(np.sqrt(E_s2 * XtX_inv[1, 1])),
    "mean_beta_mc": float(np.mean(b)),
    "design_sha256": hashlib.sha256(open(path, "rb").read()).hexdigest(),
}
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))