"""Numbers for /guides/how-to-interpret-confidence-intervals/ (constructed examples).
Python 3.13, NumPy, SciPy 1.18.1. Every interval is a two-sided Welch CI computed from summary
statistics unless stated. Run: python example.py  -> example_output.json and ci_coverage.svg
"""
import json
import math
import numpy as np
from scipy import stats

out = {}


def welch_ci(m1, s1, n1, m2, s2, n2, level=0.95):
    """Difference m2 - m1 with a two-sided Welch CI and p-value, from summary statistics."""
    v1, v2 = s1**2 / n1, s2**2 / n2
    se = math.sqrt(v1 + v2)
    df = (v1 + v2) ** 2 / (v1**2 / (n1 - 1) + v2**2 / (n2 - 1))
    d = m2 - m1
    q = stats.t.ppf(1 - (1 - level) / 2, df)
    t = d / se
    p = 2 * stats.t.sf(abs(t), df)
    return {"diff": d, "se": se, "df": df, "t": t, "p": p, "lo": d - q * se, "hi": d + q * se}


# 1. Three results: large but uncertain (log2 scale), small but precise, imprecise.
a = welch_ci(0.0, 0.5, 6, 1.0, 0.5, 6)          # log2 protein level, 6 mice per group
a.update({"ratio": 2 ** a["diff"], "ratio_lo": 2 ** a["lo"], "ratio_hi": 2 ** a["hi"]})
b = welch_ci(25.0, 1.0, 180, 25.3, 1.0, 180)     # body weight in g, 180 mice per group
c = welch_ci(10.0, 2.0, 5, 11.5, 2.0, 5)         # 5 mice per group
out["three_results"] = {"A_log2_ratio": a, "B_weight": b, "C_imprecise": c}

# 2. Overlapping group CIs vs the test of the difference (independent groups, SD 2, n 10 each).
rows = []
n, s, m1 = 10, 2.0, 10.0
q = stats.t.ppf(0.975, n - 1)
moe = q * s / math.sqrt(n)                       # margin of error of each group's 95% CI
for gap in [2.0, 2.6, 2.86, 3.5]:
    r = welch_ci(m1, s, n, m1 + gap, s, n)
    lo1, hi1 = m1 - moe, m1 + moe
    lo2, hi2 = m1 + gap - moe, m1 + gap + moe
    overlap = hi1 - lo2
    rows.append({"gap": gap, "ci_group1": [lo1, hi1], "ci_group2": [lo2, hi2], "overlap": overlap,
                 "overlap_over_moe": overlap / moe, "p_diff": r["p"], "ci_diff": [r["lo"], r["hi"]]})
out["overlap"] = {"moe": moe, "rows": rows}

# 3. Width at 90%, 95% and 99% for the same data (result C).
out["levels"] = {str(L): welch_ci(10.0, 2.0, 5, 11.5, 2.0, 5, L) for L in (0.90, 0.95, 0.99)}

# 4. Coverage: what "95%" means, and what happens when the model is wrong.
rng = np.random.default_rng(20261110)        # seed fixed before running, not searched


def run(n, runs, draw_x, draw_y, true_diff, keep=0):
    hit, first = 0, []
    for _ in range(runs):
        x, y = draw_x(n), draw_y(n)
        r = welch_ci(x.mean(), x.std(ddof=1), n, y.mean(), y.std(ddof=1), n)
        ok = r["lo"] <= true_diff <= r["hi"]
        hit += ok
        if len(first) < keep:
            first.append([r["lo"], r["hi"], bool(ok)])
    return hit / runs, first


cov = {}
cov["normal_n6"], first50 = run(6, 10000, lambda k: rng.normal(0, 1, k), lambda k: rng.normal(1, 1, k), 1.0, keep=50)
cov["normal_n3"], _ = run(3, 10000, lambda k: rng.normal(0, 1, k), lambda k: rng.normal(1, 1, k), 1.0)
# Model wrong: strongly skewed values, raw-scale analysis; the target is the true mean difference.
true_skew = math.exp(1.0**2 / 2) - math.exp(0.3**2 / 2)
cov["skewed_n6_raw"], _ = run(6, 10000, lambda k: rng.lognormal(0, 0.3, k), lambda k: rng.lognormal(0, 1.0, k), true_skew)
out["coverage"] = cov
out["first50_missed"] = sum(1 for f in first50 if not f[2])

# Figure: the first 50 intervals at n = 6 (true difference 1.0).
W, H, pad = 640, 560, 40
xs = [v for f in first50 for v in f[:2]]
lo, hi = min(xs + [1.0]) - 0.2, max(xs + [1.0]) + 0.2


def X(v):
    return pad + (v - lo) / (hi - lo) * (W - 2 * pad)


parts = [f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {W} {H}" role="img" '
         'aria-label="Fifty 95% confidence intervals from repeated simulated experiments; most cross the true difference of 1.0 and a few, drawn in red, miss it.">',
         f'<rect width="{W}" height="{H}" fill="#ffffff"/>',
         f'<line x1="{X(1.0):.1f}" y1="20" x2="{X(1.0):.1f}" y2="{H - 40}" stroke="#222" stroke-dasharray="4 3"/>',
         f'<text x="{X(1.0):.1f}" y="14" font-family="sans-serif" font-size="12" text-anchor="middle">true difference = 1.0</text>']
for i, (l, h, ok) in enumerate(first50):
    y = 30 + i * 9.6
    col = "#555555" if ok else "#c0392b"
    parts.append(f'<line x1="{X(l):.1f}" y1="{y:.1f}" x2="{X(h):.1f}" y2="{y:.1f}" stroke="{col}" stroke-width="{2 if ok else 3}"/>')
for v in np.arange(math.ceil(lo), math.floor(hi) + 1):
    parts.append(f'<text x="{X(v):.1f}" y="{H - 22}" font-family="sans-serif" font-size="12" text-anchor="middle">{v:g}</text>')
parts.append(f'<text x="{W / 2}" y="{H - 4}" font-family="sans-serif" font-size="12" text-anchor="middle">difference between groups (red: the interval misses the true value)</text>')
parts.append('</svg>')
with open("ci_coverage.svg", "w") as fh:
    fh.write("\n".join(parts) + "\n")


def rnd(o):
    if isinstance(o, float):
        return round(o, 6)
    if isinstance(o, dict):
        return {k: rnd(v) for k, v in o.items()}
    if isinstance(o, list):
        return [rnd(v) for v in o]
    return o


with open("example_output.json", "w") as fh:
    json.dump(rnd(out), fh, indent=1)
print(json.dumps(rnd(out), indent=1))