# (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).
import json, numpy as np
import pandas as pd
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import TwoSlopeNorm

df = pd.read_json("skew_grid.json")
d = df[df.ratio == 1.5]
cvs = sorted(d.cv.unique()); ns = sorted(d.n.unique())
def grid(col):
    M = np.zeros((len(cvs), len(ns)))
    for i, c in enumerate(cvs):
        for j, n in enumerate(ns):
            M[i, j] = d[(d.cv == c) & (d.n == n)][col].values[0]
    return M
LR = grid("diff_lr"); LM = grid("diff_lm")

fig, axes = plt.subplots(1, 2, figsize=(11, 4.6))
X = np.arange(len(ns)); Y = np.arange(len(cvs))

# Panel A: log-t minus raw-t
im0 = axes[0].imshow(LR, origin="lower", aspect="auto", cmap="YlOrRd", vmin=0, vmax=max(0.2, LR.max()))
axes[0].set_title("Transform beats the NAIVE test\n(log-scale t  −  raw-scale t)", fontsize=11)
fig.colorbar(im0, ax=axes[0], label="power gained")
cs = axes[0].contour(LR, levels=[0.05, 0.10, 0.15], colors="k", linewidths=0.6, alpha=0.5)
axes[0].clabel(cs, fmt=lambda v: f"+{v:.2f}", fontsize=7)

# Panel B: log-t minus MW (diverging around 0)
vmax = max(abs(LM.min()), abs(LM.max()), 0.03)
im1 = axes[1].imshow(LM, origin="lower", aspect="auto", cmap="RdBu_r",
                     norm=TwoSlopeNorm(vcenter=0, vmin=-vmax, vmax=vmax))
axes[1].set_title("Transform vs the NON-PARAMETRIC test\n(log-scale t  −  Mann–Whitney)", fontsize=11)
fig.colorbar(im1, ax=axes[1], label="power gained")

for ax in axes:
    ax.set_xticks(X); ax.set_xticklabels(ns)
    ax.set_yticks(Y); ax.set_yticklabels([f"{c:.2g}" for c in cvs])
    ax.set_xlabel("sample size per group, n")
    ax.set_ylabel("coefficient of variation, CV")
fig.suptitle("Where the log-transformed t-test improves — true median ratio 1.5, log-normal data", y=1.02, fontsize=12)
fig.tight_layout()
fig.savefig("fig9_skew_locus.png", dpi=150, bbox_inches="tight", metadata={"Author": "William J. Dwyer (ORCID 0009-0004-0855-7222)", "Title": "fig9_skew_locus", "Subject": "Generated by fig_skew_review.py for the Comparable-Strata Designer; MIT License; cite Zenodo DOI 10.5281/zenodo.20709963"})
print("saved; LR range", round(LR.min(),3), round(LR.max(),3), " LM range", round(LM.min(),3), round(LM.max(),3))
