#!/usr/bin/env python3
import json, collections
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
r=json.load(open("results.json")); a=json.load(open("adversarial.json")); k9=r["by_k"]["9"]
EN={"RNA-directed RNA polymerase/Methyltransferase NS5":"NS5 (polymerase)","Serine protease NS3":"NS3 (protease)",
    "Non-structural protein 4B":"NS4B","Non-structural protein 1":"NS1","Envelope protein E":"E (envelope)",
    "Non-structural protein 4A":"NS4A"}
fig,ax=plt.subplots(1,2,figsize=(12.5,5))
items=sorted(a["by_chain"].items(),key=lambda x:x[1])
ax[0].barh([EN.get(k,k)[:26] for k,_ in items],[v for _,v in items],
           color=["#c0392b" if "E (" in EN.get(k,k) else "#2c6fbb" for k,_ in items])
ax[0].set_xlabel("conserved 9-mers shared by all 4 serotypes")
ax[0].set_title("Where the conserved epitopes are:\n94% sit in non-structural (T-cell) proteins")
ax[0].grid(axis="x",alpha=.3)
rng=np.random.default_rng(1)
null=rng.normal(k9["null_mean"],(k9["null_p95"]-k9["null_mean"])/1.645,4000).clip(0)
ax[1].hist(null,bins=40,color="#95a5a6",edgecolor="white")
ax[1].axvline(k9["lost_ade"],color="#c0392b",lw=2.5)
ax[1].text(k9["lost_ade"]+1,ax[1].get_ylim()[1]*.75,f"ADE regions cost only {k9['lost_ade']}",color="#c0392b",fontsize=10)
ax[1].axvline(k9["null_mean"],color="#34495e",ls="--",lw=1.5)
ax[1].text(k9["null_mean"]+1,ax[1].get_ylim()[1]*.55,f"random region: {k9['null_mean']:.0f}",color="#34495e",fontsize=9)
ax[1].set_xlabel("conserved epitopes lost by excluding a region of the same size")
ax[1].set_ylabel("permutations (smoothed)")
ax[1].set_title("Excluding ADE regions costs far less than chance\n(they are depleted of conserved epitopes)")
plt.tight_layout(); plt.savefig("figs/epitopes.png",dpi=150); print("figs/epitopes.png 已生成")
