"""F2-conf-03: Monte Carlo coverage of OLS CIs under heteroskedasticity (fixed design)."""
import json, os
import numpy as np, pandas as pd
from scipy import stats

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
df = pd.read_csv(os.path.join(ROOT, "data", "design.csv"))
issues = []
if df[["x", "sigma"]].isna().any().any():
    issues.append({"issue": "missing values", "handling": "none expected; would abort"})
if (df["sigma"] <= 0).any():
    issues.append({"issue": "non-positive sigma", "handling": "flagged"})
if len(df) != 30 or not np.array_equal(df["i"].values, np.arange(1, 31)):
    issues.append({"issue": "row count/index not 1..30", "handling": "flagged"})

ALPHA, BETA, R, SEED = 1.0, 0.5, 2000, 20261030
x = df["x"].to_numpy(float); sig = df["sigma"].to_numpy(float); n = len(x)
X = np.column_stack([np.ones(n), x])
XtXi = np.linalg.inv(X.T @ X)
A = XtXi @ X.T                      # (2, n): beta_hat = A @ y
tcrit = stats.t.ppf(0.975, n - 2)

# (c) exact sd: Var(beta_hat) = sum_i a_i^2 sigma_i^2, a_i = (x_i - xbar)/Sxx
sd_exact = float(np.sqrt(np.sum(A[1] ** 2 * sig ** 2)))

rng = np.random.default_rng(SEED)
eps = rng.standard_normal((R, n))
Y = ALPHA + BETA * x + sig * eps     # (R, n)
B = Y @ A.T                          # (R, 2)
resid = Y - B @ X.T
s2 = (resid ** 2).sum(1) / (n - 2)
se_conv = np.sqrt(s2 * XtXi[1, 1])
hc0 = (resid ** 2) @ (A[1] ** 2)     # sandwich variance for slope
se_hc1 = np.sqrt(hc0 * n / (n - 2))
bh = B[:, 1]
cov_conv = float(np.mean(np.abs(bh - BETA) <= tcrit * se_conv))
cov_hc1 = float(np.mean(np.abs(bh - BETA) <= tcrit * se_hc1))
mcse = float(np.sqrt(0.95 * 0.05 / R))
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(bh, ddof=1)),
    "mean_se_conventional": float(se_conv.mean()),
    "conventional_ci_undercovers": bool(cov_conv < 0.95 - 3 * mcse),
    "n_replications": R,
    "n_obs_per_replication": n,
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"design.csv": int(len(df))},
    "random_seed": SEED,
    # auxiliary
    "mcse_binomial_at_095": mcse,
    "mcse_coverage_conventional": float(np.sqrt(cov_conv * (1 - cov_conv) / R)),
    "mcse_coverage_hc1": float(np.sqrt(cov_hc1 * (1 - cov_hc1) / R)),
    "mean_se_hc1": float(se_hc1.mean()),
    # sqrt of exact E[s^2]*(X'X)^-1_22, E[s^2] = tr(M Sigma)/(n-2); upper bound for E[se_conv] (Jensen)
    "sqrt_exact_E_var_conv": float(np.sqrt(np.trace((np.eye(n) - X @ A) @ np.diag(sig**2)) / (n - 2) * XtXi[1, 1])),
    "t_crit": float(tcrit),
}
with open(os.path.join(ROOT, "out", "results.json"), "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))