"""Numbers and figure for /guides/sd-vs-sem-error-bars/.

Constructed data (written by hand, no random generation): body-weight gain (g) in 5 control and 5 treated
mice, one value per mouse. Python 3.13, NumPy, SciPy 1.18.1, matplotlib 3.11.2.
Run: python example.py   (writes example_output.json and sd-vs-sem.svg next to this file)
"""
import json
from pathlib import Path

import matplotlib
matplotlib.use("svg")
import matplotlib.pyplot as plt
import numpy as np
from scipy import stats

HERE = Path(__file__).parent
control = np.array([3.1, 4.6, 2.4, 5.2, 3.9])
treated = np.array([5.8, 7.4, 4.9, 8.6, 6.3])


def summary(x):
    n = len(x)
    m = x.mean()
    sd = x.std(ddof=1)
    sem = sd / np.sqrt(n)
    tq = stats.t.ppf(0.975, n - 1)
    return {"n": n, "mean": m, "sd": sd, "sem": sem, "t_multiplier": tq,
            "ci_low": m - tq * sem, "ci_high": m + tq * sem}


out = {"control": summary(control), "treated": summary(treated)}
w = stats.ttest_ind(treated, control, equal_var=False)
out["welch"] = {"t": w.statistic, "df": w.df, "p": w.pvalue,
                "diff": treated.mean() - control.mean(),
                "ci": list(w.confidence_interval(0.95))}

# 1. What a mean +/- 1 SEM bar covers, and the multiplier a 95% CI needs, by number of animals
cov = []
for n in [3, 4, 5, 6, 8, 10, 20, 30]:
    cov.append({"n": n, "sem_bar_coverage": 2 * stats.t.cdf(1, n - 1) - 1,
                "t_multiplier_95": stats.t.ppf(0.975, n - 1)})
out["sem_bar_coverage"] = cov


# 2. Two independent groups, equal n and equal SD: p-value (Student's t) when bars touch, gap or overlap
def p_from_diff(n, se, diff):
    t = diff / (se * np.sqrt(2))
    return 2 * stats.t.sf(abs(t), 2 * n - 2)


gaps = []
for n in [3, 5, 10]:
    se = 1.0
    tq = stats.t.ppf(0.975, n - 1)
    gaps.append({
        "n": n,
        "sem_bars_touch": p_from_diff(n, se, 2 * se),
        "gap_of_one_sem": p_from_diff(n, se, 3 * se),
        "ci_bars_touch": p_from_diff(n, se, 2 * tq * se),
        "ci_overlap_half_arm": p_from_diff(n, se, 1.5 * tq * se),
        "ci_overlap_one_arm": p_from_diff(n, se, 1.0 * tq * se),
    })
out["overlap_p"] = gaps

# 3. Simulation: how often each bar contains the true mean (normal data, true mean 0, SD 1)
rng = np.random.default_rng(20261110)  # seed fixed before running; not searched
REPS, BOOT = 20_000, 1_000
sim = []
for n in [3, 5, 10]:
    x = rng.standard_normal((REPS, n))
    m = x.mean(axis=1)
    se = x.std(axis=1, ddof=1) / np.sqrt(n)
    tq = stats.t.ppf(0.975, n - 1)
    hit_sem = np.mean(np.abs(m) <= se)
    hit_2sem = np.mean(np.abs(m) <= 2 * se)
    hit_t = np.mean(np.abs(m) <= tq * se)
    hits = 0
    for start in range(0, REPS, 1_000):  # percentile bootstrap of the mean, the interval seaborn draws by default
        xs = x[start:start + 1_000]
        idx = rng.integers(0, n, size=(xs.shape[0], BOOT, n))
        bm = np.take_along_axis(np.broadcast_to(xs[:, None, :], idx.shape), idx, axis=2).mean(axis=2)
        lo, hi = np.percentile(bm, [2.5, 97.5], axis=1)
        hits += int(np.sum((lo <= 0) & (0 <= hi)))
    sim.append({"n": n, "mean_pm_1sem": hit_sem, "mean_pm_2sem": hit_2sem,
                "t_95ci": hit_t, "bootstrap_95ci": hits / REPS})
out["coverage_simulation"] = {"reps": REPS, "bootstrap_resamples": BOOT, "seed": 20261110, "rows": sim}

# 4. Recovering SD from a published SEM
out["sd_from_sem_example"] = {"sem": 0.52, "n": 5, "sd": 0.52 * np.sqrt(5)}


def r(v):
    if isinstance(v, dict):
        return {k: r(x) for k, x in v.items()}
    if isinstance(v, list):
        return [r(x) for x in v]
    if isinstance(v, (float, np.floating)):
        return float(f"{float(v):.6g}")
    if isinstance(v, np.integer):
        return int(v)
    return v


(HERE / "example_output.json").write_text(json.dumps(r(out), indent=2) + "\n")
print(json.dumps(r(out), indent=2))

# Figure: the same ten mice shown four ways
plt.rcParams["svg.hashsalt"] = "sd-vs-sem"
plt.rcParams["font.size"] = 10
fig, axes = plt.subplots(1, 4, figsize=(10, 3.4), sharey=True, facecolor="white")
offsets = np.array([-0.08, 0.04, -0.03, 0.08, 0.0])
panels = [("Points and mean", None), ("Mean ± SD", "sd"), ("Mean ± SEM", "sem"), ("Mean ± 95% CI", "ci")]
for ax, (title, kind) in zip(axes, panels):
    ax.set_facecolor("white")
    for i, (x, name) in enumerate([(control, "control"), (treated, "treated")]):
        s = out[name]
        ax.scatter(i + offsets, x, s=22, color="#555555", zorder=3)
        if kind is None:
            ax.hlines(s["mean"], i - 0.2, i + 0.2, color="black", lw=2)
            continue
        if kind == "sd":
            lo, hi = s["mean"] - s["sd"], s["mean"] + s["sd"]
        elif kind == "sem":
            lo, hi = s["mean"] - s["sem"], s["mean"] + s["sem"]
        else:
            lo, hi = s["ci_low"], s["ci_high"]
        ax.errorbar(i + 0.25, s["mean"], yerr=[[s["mean"] - lo], [hi - s["mean"]]], fmt="o",
                    color="black", capsize=4, lw=1.5)
    ax.set_xticks([0, 1], ["Control", "Treated"])
    ax.set_xlim(-0.5, 1.6)
    ax.set_title(title)
axes[0].set_ylabel("Body-weight gain (g)")
fig.text(0.5, 0.01, "Same 5 + 5 mice in every panel", ha="center", fontsize=9)
fig.tight_layout(rect=(0, 0.04, 1, 1))
fig.savefig(HERE / "sd-vs-sem.svg", metadata={"Date": None}, facecolor="white")