# (c) 2026 William J. Dwyer. Free and open-source software under the MIT License (see LICENSE).
# If you use or build upon this software, please retain this notice and cite the accompanying
# methodology paper together with the archived software (Zenodo DOI: 10.5281/zenodo.20709963).
"""Rescue-zone map: power gained by the screen-guided log-scale t-test over the
naive raw-scale t-test, across the (sample size n) x (coefficient of variation)
plane, for a true doubling of the median."""
import numpy as np
from scipy import stats
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap

rng = np.random.default_rng(99)
REPS = 2500
shift = np.log(2.0)   # true median ratio 2

def ln_params(cv):
    s2 = np.log(cv**2 + 1.0); return -s2/2.0, np.sqrt(s2)

n_grid = np.array([3, 4, 5, 6, 8, 10, 15, 20, 30, 50, 75, 100])
cv_grid = np.round(np.arange(0.1, 3.01, 0.1), 2)

gain = np.zeros((len(cv_grid), len(n_grid)))
for j, n in enumerate(n_grid):
    for i, cv in enumerate(cv_grid):
        me, se = ln_params(cv)
        g1 = rng.lognormal(me, se, (REPS, n))
        g2 = rng.lognormal(me + shift, se, (REPS, n))
        praw = np.mean(stats.ttest_ind(g1, g2, axis=1, equal_var=False).pvalue < .05)
        plog = np.mean(stats.ttest_ind(np.log(g1), np.log(g2), axis=1, equal_var=False).pvalue < .05)
        gain[i, j] = plog - praw

# ---- plot ----
fig, ax = plt.subplots(figsize=(8.4, 5.2))
cmap = LinearSegmentedColormap.from_list("rescue", ["#f7f7f7", "#fee08b", "#fc8d59", "#b2182b"])
X, Y = np.meshgrid(n_grid, cv_grid)
cf = ax.contourf(X, Y, gain, levels=np.linspace(0, 0.20, 11), cmap=cmap, extend="max")
cl = ax.contour(X, Y, gain, levels=[0.05, 0.10, 0.15], colors="black", linewidths=0.7, alpha=0.5)
ax.clabel(cl, fmt=lambda v: f"+{v:.2f}", fontsize=8)
cb = fig.colorbar(cf, ax=ax, pad=0.02)
cb.set_label("power gained by the log-scale t-test\n(vs. raw-scale t-test; true median ratio 2)")

ax.set_xscale("log")
ax.set_xticks([3, 5, 10, 20, 50, 100]); ax.set_xticklabels([3, 5, 10, 20, 50, 100])
ax.set_xlabel("sample size per group, n")
ax.set_ylabel("coefficient of variation, CV = SD / mean")
ax.set_xlim(3, 100); ax.set_ylim(0.1, 3.0)

# zone boundaries
for cv0, lab in [(0.3, "CV = 0.3"), (0.5, "CV = 0.5")]:
    ax.axhline(cv0, color="grey", ls="--", lw=1)
# resolution floor: Mann-Whitney cannot reach p<0.05 for n<=3 (two-sided)
ax.axvspan(3, 3.7, color="#3b3b6d", alpha=0.18)
ax.text(3.15, 2.7, "rank test\ncannot reach\np < 0.05", rotation=90, va="top", ha="left", fontsize=7.5, color="#3b3b6d")

# zone labels (semi-transparent backgrounds keep them legible over the colored field and contours)
_box = dict(boxstyle="round,pad=0.28", fc="white", ec="none", alpha=0.72)
_box2 = dict(boxstyle="round,pad=0.2", fc="white", ec="none", alpha=0.6)
ax.text(11, 0.18, "effectively normal — use the raw-scale t-test", fontsize=9, style="italic",
        color="#333", ha="center", va="center", bbox=_box2)
ax.text(6.2, 0.40, "borderline", fontsize=9, style="italic", color="#333", ha="center", va="center", bbox=_box2)
ax.text(30, 2.66, "RESCUE ZONE", fontsize=15, fontweight="bold", color="#7a0a1e",
        ha="center", va="center", bbox=_box, zorder=5)
ax.text(30, 2.37, "screen → log-scale t-test recovers power", fontsize=9.5, color="#7a0a1e",
        ha="center", va="center", bbox=_box, zorder=5)

# mark the worked examples from the paper
for (n, cv, txt) in [(20, 1.0, "CV 1.0, n 20\n(+0.15)"), (30, 2.0, "CV 2.0, n 30\n(+0.24)")]:
    ax.plot(n, cv, "o", color="black", ms=6, mfc="white", mew=1.5, zorder=6)
    ax.annotate(txt, (n, cv), textcoords="offset points", xytext=(8, 6), fontsize=8,
                bbox=_box2, zorder=6)

ax.set_title("The rescue zone: where a screen-justified log-scale t-test most helps", fontsize=12)
fig.tight_layout()
fig.savefig("fig8_rescue_zone.png", dpi=150, bbox_inches="tight", metadata={"Author": "William J. Dwyer (ORCID 0009-0004-0855-7222)", "Title": "fig8_rescue_zone", "Subject": "Generated by fig_rescue_zone.py for the Comparable-Strata Designer; MIT License; cite Zenodo DOI 10.5281/zenodo.20709963"})
print("saved; gain range", round(gain.min(), 3), round(gain.max(), 3))
