import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
fig = plt.figure(figsize=(11, 6.2))
gs = fig.add_gridspec(1, 2, width_ratios=[1.15, 1.0], wspace=0.28)
concord, discord, empty = "#2c6fbb", "#c0504d", "#e8e8e8"
# Panel A: 2x2 concordance (allele class x clinical direction)
axA = fig.add_subplot(gs[0, 0])
M = np.array([[(lof["phenotype_class"]=="Hypopigmentation").sum(), (lof["phenotype_class"]=="Hyperpigmentation").sum()],
[(dom["phenotype_class"]=="Hypopigmentation").sum(), (dom["phenotype_class"]=="Hyperpigmentation").sum()]])
cellcol = [[concord, discord], [empty, empty]]
for i in range(2):
for j in range(2):
n = M[i, j]
axA.add_patch(plt.Rectangle((j, 1-i), 1, 1, facecolor=cellcol[i][j] if n else empty, edgecolor="white", lw=2))
axA.text(j+0.5, 1-i+0.5, str(n), ha="center", va="center", fontsize=20, fontweight="bold",
color="white" if (n and cellcol[i][j]!=empty) else "#555")
axA.set_xlim(0,2); axA.set_ylim(0,2)
axA.set_xticks([0.5,1.5]); axA.set_xticklabels(["Hypopigmentation","Hyperpigmentation"])
axA.set_yticks([1.5,0.5]); axA.set_yticklabels(["Loss-of-function\n(recessive / X-linked)","Dominant /\ncongenital / acquired"])
axA.set_xlabel("Clinical pigmentation direction (OMIM)")
axA.tick_params(length=0)
for s in axA.spines.values(): s.set_visible(False)
axA.set_title("Positive-regulator prediction (LoF -> hypopigmentation)\nholds only under a loss-of-function allele", loc="left", fontsize=9)
axA.text(1.02, -0.14, f"LoF: {lof['correct'].sum()}/{len(lof)} concordant | base rate P(hypo|LoF)={base_hypo:.2f} | permutation p<1e-5",
transform=axA.transAxes, ha="right", va="top", fontsize=7.5, color="#7a7a7a")
# Panel B: per-gene, colored by concordance, sized by n direction sources
axB = fig.add_subplot(gs[0, 1])
allg = pd.concat([lof, dom])
allg["concord"] = allg["correct"]
allg = allg.sort_values(["is_LoF","correct","gene"], ascending=[False, True, True]).reset_index(drop=True)
y = np.arange(len(allg))
xx = allg["bajpai_effect"].fillna(0.3)
cols_pt = [concord if c else discord for c in allg["concord"]]
axB.hlines(y, 0, xx, color="#cfcfcf", lw=1.1, zorder=1)
axB.scatter(xx, y, c=cols_pt, s=34 + 10*allg["n_direction_sources"], zorder=2, edgecolor="white", lw=0.5)
axB.set_yticks(y); axB.set_yticklabels([f"$\\it{{{g}}}$" for g in allg["gene"]], fontsize=6.0)
axB.set_ylim(-0.7, len(allg)-0.3)
axB.set_xlabel("Bajpai CRISPR effect on melanin\n(higher = knockdown reduces melanin more)")
axB.set_title("Per-gene concordance; point size = # independent\ndirection sources", loc="left", fontsize=9)
leg = [Line2D([0],[0],marker="o",color="w",markerfacecolor=concord,markersize=7,label="Concordant with the law"),
Line2D([0],[0],marker="o",color="w",markerfacecolor=discord,markersize=7,label="Discordant (dominant / non-LoF allele)")]
axB.legend(handles=leg, loc="lower right", fontsize=6.5, frameon=False)
fig.suptitle("A functional melanin screen predicts Mendelian pigmentation-disorder direction - conditioned on allele mechanism",
x=0.5, y=1.005, fontsize=11, fontweight="bold")
fig.text(0.5, -0.02,
"Genes shared between the Bajpai 2023 CRISPR melanin screen / NB6 GRN / NB7 signed network (function) and D'Arcy 2023 OMIM disorders (clinic). "
"The loss-of-function prediction holds for every recessive/X-linked gene and is not asserted for dominant alleles - each discordant dominant "
"acts by a documented non-loss-of-function mechanism (TYR melanoma allele; CDKN2A/KIT proliferative; PSENEN Notch; BNC2 developmental).",
ha="center", va="top", fontsize=6.6, color="#333", wrap=True)
FIGDIR = ROOT / "notebooks" / "figures"
FIGDIR.mkdir(parents=True, exist_ok=True)
fig.savefig(FIGDIR / "nb10_direction_law.png", dpi=200, bbox_inches="tight")
print("saved", FIGDIR / "nb10_direction_law.png")