#!/usr/bin/env python3
"""
07_figures.py -- Publication-grade figures for the TP53-mCherry knock-in report.
  A: knock-in strategy -- (top) SpCas9 guides at the stop codon; (bottom) p53 domain map
     showing the C-terminal mCherry fusion after the regulatory CTD.
  B: specificity evidence -- two-tool off-target cross-validation (pseudogene-safety).
  C: HDR donor map + genotyping/Sanger primer strategy.
Colours come from the report PALETTE so figures match the PDF exactly.
Saves PNG (300 dpi) + vector PDF for each to output/figures/.
"""
import json, os
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle, FancyBboxPatch, Patch
from report_style import PALETTE

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "..", "output", "figures")
os.makedirs(OUT, exist_ok=True)
def C(k): c = PALETTE[k]; return (c.red, c.green, c.blue)
NAVY, GREEN, AMBER, RED = C("navy"), C("green"), C("amber"), C("red")
LGRAY, MGRAY, CAP, DARK = C("light_gray"), C("mid_gray"), C("caption"), C("dark_gray")
BLUE = (0.18, 0.40, 0.62); CHERRY = (0.78, 0.10, 0.22); TEAL = (0.10, 0.52, 0.50)

locus = json.load(open(os.path.join(HERE, "locus.json")))
gd = {g["id"]: g for g in json.load(open(os.path.join(HERE, "guide_candidates.json")))["guides"]}
cris = json.load(open(os.path.join(HERE, "crispor_summary.json")))["guides"]
offt = json.load(open(os.path.join(HERE, "offtarget_results.json")))
ff = json.load(open(os.path.join(HERE, "flashfry_summary.json")))["guides"]
donor = json.load(open(os.path.join(HERE, "donor.json")))
prim = json.load(open(os.path.join(HERE, "primers.json")))["sets"]
mch = json.load(open(os.path.join(HERE, "mcherry_opt.json")))

plt.rcParams.update({"font.family": "DejaVu Sans", "font.size": 9,
                     "axes.linewidth": 0.8, "savefig.dpi": 300})

def save(fig, name):
    fig.savefig(os.path.join(OUT, name + ".png"), dpi=300, bbox_inches="tight", pad_inches=0.10)
    fig.savefig(os.path.join(OUT, name + ".pdf"), bbox_inches="tight", pad_inches=0.10)
    plt.close(fig); print("  wrote", name)

# ============================================================ Figure A
def fig_A():
    fig = plt.figure(figsize=(9.2, 7.4))
    gs = fig.add_gridspec(2, 1, height_ratios=[1.18, 0.92], hspace=0.30)

    # ---------- top: guides at the stop codon (x = bp relative to insertion point) ----------
    ax = fig.add_subplot(gs[0]); ax.set_xlim(-42, 18); ax.set_ylim(-2.7, 3.0); ax.axis("off")
    ax.set_title("A.  SpCas9 guides at the TP53 stop codon (C-terminal insertion site)",
                 loc="left", fontsize=11.5, fontweight="bold", color=NAVY, pad=10)
    ax.plot([-42, 18], [0, 0], color="black", lw=1.4, zorder=1)               # backbone
    ax.add_patch(Rectangle((-42, -0.30), 42, 0.60, color=BLUE, alpha=0.20, zorder=0))   # CDS
    ax.add_patch(Rectangle((0, -0.30), 3, 0.60, color=AMBER, alpha=0.90, zorder=2))     # stop
    ax.add_patch(Rectangle((3, -0.30), 15, 0.60, color=TEAL, alpha=0.16, zorder=0))     # 3'UTR
    ax.text(-38, 0.46, "p53 coding sequence (...Asp$_{393}$)", ha="left", va="bottom", fontsize=8.4, color=BLUE)
    ax.text(11, 0.46, "3' UTR", ha="center", va="bottom", fontsize=8.4, color=TEAL)
    ax.annotate("native stop (TGA)", xy=(2.7, 0.30), xytext=(5.0, 1.05), fontsize=7.6,
                color=AMBER, fontweight="bold", ha="left", va="center",
                arrowprops=dict(arrowstyle="-", color=AMBER, lw=1.0))
    # insertion arrow
    ax.annotate("", xy=(0, 0.32), xytext=(0, 1.75), arrowprops=dict(arrowstyle="-|>", color=CHERRY, lw=2.4))
    ax.text(0, 2.0, "mCherry insertion\n(in-frame, before native stop)", ha="center", va="bottom",
            fontsize=8.6, color=CHERRY, fontweight="bold")

    def guide(yc, x0, x1, pam0, pam1, cut, name, scores, col, up):
        ax.add_patch(Rectangle((x0, yc-0.16), x1-x0, 0.32, color=col, alpha=0.85, zorder=3))   # protospacer
        ax.add_patch(Rectangle((pam0, yc-0.16), pam1-pam0, 0.32, color="black", alpha=0.62, zorder=4))  # PAM
        ax.text((pam0+pam1)/2, yc, "PAM", color="white", fontsize=5.4, ha="center", va="center", zorder=5)
        ax.plot([cut, cut], [yc-0.42, yc+0.42], color=RED, lw=1.8, zorder=6)                   # cut tick
        ty = yc+0.34 if up else yc-0.34
        ax.text((x0+x1)/2, ty, name, ha="center", va="bottom" if up else "top",
                fontsize=8.2, color=col, fontweight="bold")
        ax.text((x0+x1)/2, ty+(0.34 if up else -0.34), scores, ha="center", va="bottom" if up else "top",
                fontsize=7.4, color=DARK)
    # g2 (+strand) upper lane: protospacer rel [-35,-15], PAM [-15,-12], cut -18
    guide(1.15, -35, -15, -15, -12, gd["TP53ct_g2"]["dist_to_insert"],
          "g2  backup  (+ strand)", "MIT 55 / CFD 73  |  silent PAM-block", BLUE, up=True)
    # g3 (-strand) lower lane: protospacer rel [-9,+11], PAM [-12,-9], cut -6
    guide(-1.15, -9, 11, -12, -9, gd["TP53ct_g3"]["dist_to_insert"],
          "g3  PRIMARY  (- strand)", "MIT 70 / CFD 87  |  Doench'16 49", CHERRY, up=False)
    ax.text(-41, -2.4, "g3 protospacer straddles the insertion point -> split by the insert "
            "(self-disabled, no PAM-block needed)", fontsize=7.6, color=CHERRY, ha="left", va="center")
    ax.plot([-42, -32], [2.7, 2.7], color="black", lw=1.3); ax.text(-37, 2.5, "10 bp", ha="center", va="top", fontsize=7)
    ax.text(17.5, 2.7, "red | = predicted Cas9 cut", fontsize=6.8, color=RED, ha="right", va="top")

    # ---------- bottom: p53 domain map + C-terminal fusion ----------
    ax2 = fig.add_subplot(gs[1]); ax2.axis("off")
    ax2.set_title("B.  C-terminal mCherry fusion relative to p53 domain structure",
                  loc="left", fontsize=11.5, fontweight="bold", color=NAVY, pad=8)
    ax2.set_xlim(-12, 393+60+236+12); ax2.set_ylim(-2.0, 1.7)
    domains = [("TAD", 1, 63, (0.45,0.55,0.70)), ("PRD", 64, 92, (0.62,0.68,0.78)),
               ("DNA-binding domain", 94, 292, BLUE), ("OD", 323, 356, TEAL),
               ("CTD", 363, 393, AMBER)]
    for lab,s,e,col in domains:
        ax2.add_patch(Rectangle((s,-0.28),e-s,0.56,color=col,ec="black",lw=0.6))
        ax2.text((s+e)/2, 0, lab, ha="center", va="center", fontsize=7.2,
                 color="white" if lab in ("DNA-binding domain","OD","CTD") else "black",
                 fontweight="bold")
    # linker + mCherry after 393
    lk0=393; lk1=393+60; mc1=lk1+236
    ax2.add_patch(Rectangle((lk0,-0.18),lk1-lk0,0.36,color=TEAL,ec="black",lw=0.6))
    ax2.add_patch(Rectangle((lk1,-0.28),mc1-lk1,0.56,fc=CHERRY,ec="black",lw=0.7))
    ax2.text((lk1+mc1)/2,0,"mCherry",ha="center",va="center",fontsize=8.5,color="white",fontweight="bold")
    ax2.annotate("(GGGGS)x3\nlinker", xy=((lk0+lk1)/2,0.18), xytext=((lk0+lk1)/2,1.15),
                 ha="center", fontsize=7, color=TEAL, fontweight="bold",
                 arrowprops=dict(arrowstyle="-", color=TEAL, lw=0.9))
    ax2.annotate("fused after residue 393\n(downstream of the CTD)", xy=(393,-0.28), xytext=(300,-1.4),
                 ha="center", fontsize=7.3, color=AMBER, fontweight="bold",
                 arrowprops=dict(arrowstyle="-|>", color=AMBER, lw=1.2))
    ax2.text((mc1)/2, 1.5, "CTD (363-393) is the regulatory domain (MDM2 ubiquitination, PTMs) -> "
             "validate tagged-p53 function", ha="center", fontsize=7.4, color=CAP)
    # residue ruler
    for r in [1,100,200,300,393]:
        ax2.plot([r,r],[-0.28,-0.40],color="black",lw=0.7); ax2.text(r,-0.52,str(r),ha="center",fontsize=6.2,color=CAP)
    save(fig, "figure_A_knockin_strategy")

# ============================================================ Figure B
def fig_B():
    fig, (axL, axR) = plt.subplots(1, 2, figsize=(9.4, 4.3), gridspec_kw={"width_ratios":[1.25,1]})
    order = ["TP53ct_g3","TP53ct_g2","TP53ct_g4","TP53ct_g1","TP53ct_g5","TP53ct_g6","TP53ct_g7"]
    labs = ["g3","g2","g4","g1","g5","g6","g7"]
    mit = [cris[g]["mit_spec"] for g in order]; cfd = [cris[g]["cfd_spec"] for g in order]
    x = np.arange(len(order)); w = 0.38
    b1 = axL.bar(x-w/2, mit, w, label="MIT specificity", color=NAVY)
    b2 = axL.bar(x+w/2, cfd, w, label="CFD specificity", color=TEAL)
    b1[0].set_color(CHERRY); b2[0].set_color(CHERRY); b2[0].set_alpha(0.55)
    axL.set_xticks(x); axL.set_xticklabels(labs, fontsize=8.5)
    axL.set_ylabel("CRISPOR specificity score (0-100)", fontsize=8.5)
    axL.set_ylim(0, 108)
    axL.set_title("A.  Guide specificity (CRISPOR, hg38)", loc="left", fontsize=10.5, fontweight="bold", color=NAVY)
    # explicit proxy handles so the legend swatches match the majority (navy/teal) bars,
    # not the crimson g3-highlight (which is labelled "primary" on the bar itself)
    axL.legend(handles=[Patch(color=NAVY, label="MIT specificity"), Patch(color=TEAL, label="CFD specificity")],
               fontsize=7.5, frameon=False, loc="upper center", ncol=2, bbox_to_anchor=(0.55, 1.02))
    axL.spines[["top","right"]].set_visible(False)
    axL.text(0, mit[0]+3, "primary", ha="center", fontsize=6.8, color=CHERRY, fontweight="bold")
    axL.text(1, max(mit[1],cfd[1])+3, "backup", ha="center", fontsize=6.8, color=BLUE, fontweight="bold")
    axL.text(6, 100, "higher = more specific", fontsize=6.8, color=CAP, va="top", ha="right")

    g3o = offt["TP53ct_g3"]; co_bins = g3o["bins"]
    def co(mm): return co_bins.get(f"NGG_{mm}mm",0)+co_bins.get(f"NAG_{mm}mm",0)
    cas = [co(2), co(3), co(4)]
    cr = cris["TP53ct_g3"]["ot_by_mm"]; crv = [cr["2mm"], cr["3mm"], cr["4mm"]]
    ffb = ff["TP53ct_g3"]["ot_by_mm"]; ffv = [ffb.get("2",0), ffb.get("3",0), ffb.get("4",0)]
    xm = np.arange(3); w = 0.27
    axR.bar(xm-w, cas, w, label="Cas-OFFinder", color=BLUE)
    axR.bar(xm,   crv, w, label="CRISPOR", color=AMBER)
    axR.bar(xm+w, ffv, w, label="FlashFry", color=TEAL)
    for i,(a,b,f) in enumerate(zip(cas,crv,ffv)):
        axR.text(i-w, a+3, str(a), ha="center", fontsize=6.0, color=BLUE)
        axR.text(i,   b+3, str(b), ha="center", fontsize=6.0, color=AMBER)
        axR.text(i+w, f+3, str(f), ha="center", fontsize=6.0, color=TEAL)
    axR.set_xticks(xm); axR.set_xticklabels(["2 mm","3 mm","4 mm"], fontsize=8.5)
    axR.set_ylabel("g3 off-target sites", fontsize=8.5); axR.set_ylim(0, 185)
    axR.set_title("B.  g3 off-targets: three-tool agreement", loc="left", fontsize=10.5, fontweight="bold", color=NAVY)
    axR.legend(fontsize=7.2, frameon=False, loc="upper left")
    axR.spines[["top","right"]].set_visible(False)
    txt = ("0 off-targets at <=1 mismatch (all 3 tools)\n0 hits in TP53 family / pseudogenes\n"
           "all exonic hits >=4 mm (max CFD 0.17)")
    axR.text(0.97, 0.55, txt, transform=axR.transAxes, ha="right", va="top", fontsize=6.8, color=DARK,
             bbox=dict(boxstyle="round,pad=0.4", fc=C("light_green"), ec=GREEN, lw=1.0))
    save(fig, "figure_B_specificity_evidence")

# ============================================================ Figure C
def fig_C():
    km = donor["ki_map"]
    fig = plt.figure(figsize=(9.4, 5.8))
    gs = fig.add_gridspec(2, 1, height_ratios=[1, 1], hspace=0.62)
    Llk = len(mch["linker_nt"]); Lm = len(mch["mcherry_nt"]); dt = donor["donor_len"]
    # ---- top: donor components ----
    ax = fig.add_subplot(gs[0]); ax.axis("off")
    ax.set_title("A.  HDR donor (2,353 bp dsDNA)", loc="left", fontsize=10.5, fontweight="bold", color=NAVY, pad=18)
    segs = [("5' arm 800 bp", 0, 800, BLUE, 0.22), ("(GGGGS)x3", 800, 800+Llk, TEAL, 0.85),
            ("mCherry 708 bp", 800+Llk, 800+Llk+Lm, CHERRY, 0.85), ("3' arm 800 bp", 800+Llk+Lm, dt, AMBER, 0.22)]
    ax.set_xlim(-70, dt+70); ax.set_ylim(-1.9, 2.0)
    for lab,x0,x1,col,al in segs:
        ax.add_patch(Rectangle((x0,-0.32),x1-x0,0.64,color=col,alpha=al,ec="black",lw=0.7))
        if lab != "(GGGGS)x3":
            ax.text((x0+x1)/2,0,lab,ha="center",va="center",fontsize=7.6,
                    color="white" if al>0.5 else "black", fontweight="bold")
    ax.annotate("(GGGGS)x3", xy=(800+Llk/2,0.32), xytext=(800+Llk/2,1.05), ha="center",
                fontsize=6.8, color=TEAL, fontweight="bold", arrowprops=dict(arrowstyle="-",color=TEAL,lw=0.8))
    pb = 800-15
    ax.annotate("silent PAM-block for g2\n(GGG->GGT, Gly)", xy=(pb,0.32), xytext=(pb-70,1.55),
                ha="center", fontsize=7, color=BLUE, fontweight="bold",
                arrowprops=dict(arrowstyle="-|>", color=BLUE, lw=1.2))
    ax.annotate("g3 site split by insert\n(no edit needed)", xy=(800,-0.32), xytext=(800,-1.35),
                ha="center", fontsize=7, color=CHERRY, fontweight="bold",
                arrowprops=dict(arrowstyle="-|>", color=CHERRY, lw=1.2))
    ax.text(dt/2, -1.78, "Insert (linker+mCherry): 0 BsaI / 0 BsmBI sites (re-clonable).   "
            f"Homology arms carry {len(donor['donor_bsaI'])} BsaI + {len(donor['donor_bsmBI'])} BsmBI (native genomic).",
            ha="center", fontsize=7.2, color=CAP)

    # ---- bottom: genotyping primers ----
    ax2 = fig.add_subplot(gs[1]); ax2.axis("off")
    ax2.set_title("B.  Genotyping & Sanger primers (junctions anchored OUTSIDE the arms)",
                  loc="left", fontsize=10.5, fontweight="bold", color=NAVY, pad=8)
    LCl=km["left_ctx_len"]; armv=donor["arm_len"]; ILv=km["insert_len"]; tot=LCl+armv+ILv+armv+LCl
    ax2.set_xlim(-40, tot+40); ax2.set_ylim(-3.4, 1.5)
    zones=[("genomic",0,LCl,MGRAY,0.5),("5' arm",LCl,LCl+armv,BLUE,0.5),
           ("insert",LCl+armv,LCl+armv+ILv,CHERRY,0.85),("3' arm",LCl+armv+ILv,LCl+armv+ILv+armv,AMBER,0.5),
           ("genomic",LCl+armv+ILv+armv,tot,MGRAY,0.5)]
    for lab,x0,x1,col,al in zones:
        ax2.add_patch(Rectangle((x0,0.55),x1-x0,0.40,color=col,alpha=al,ec="black",lw=0.5))
    j5=LCl+armv; j3=LCl+armv+ILv
    ax2.text(j5,1.18,"5' jxn",ha="center",fontsize=6.8,color=NAVY); ax2.text(j3,1.18,"3' jxn",ha="center",fontsize=6.8,color=NAVY)
    ax2.text(LCl/2,1.18,"insert (mCherry)",ha="center",fontsize=6.4,color=CHERRY) if False else None
    def pair(y,F0,F1,R0,R1,label,size):
        ax2.annotate("",xy=(F1,y),xytext=(F0,y),arrowprops=dict(arrowstyle="-|>",color=GREEN,lw=2.2))
        ax2.annotate("",xy=(R0,y),xytext=(R1,y),arrowprops=dict(arrowstyle="-|>",color=RED,lw=2.2))
        ax2.plot([F0,R1],[y-0.16,y-0.16],color=CAP,lw=0.6)
        ax2.text((F0+R1)/2,y-0.55,f"{label}: {size}",ha="center",fontsize=7.2,color=DARK)
    p5=prim["5_prime_junction"]; p3=prim["3_prime_junction"]; pi=prim["internal_across_insert"]
    pair(0.0, p5["fwd_start"], p5["fwd_start"]+p5["fwd_len"], p5["rev_end"]-p5["rev_len"], p5["rev_end"],
         "5' junction PCR", f"{p5['product_size_ki']} bp (KI only)")
    pair(-1.05, p3["fwd_start"], p3["fwd_start"]+p3["fwd_len"], p3["rev_end"]-p3["rev_len"], p3["rev_end"],
         "3' junction PCR", f"{p3['product_size_ki']} bp (KI only)")
    pair(-2.1, pi["fwd_start"], pi["fwd_start"]+pi["fwd_len"], pi["rev_end"]-pi["rev_len"], pi["rev_end"],
         "internal / zygosity", f"WT {pi['product_size_wt']} vs KI {pi['product_size_ki']} bp")
    ax2.text(tot/2,-3.15,"Green = forward, red = reverse.  Junction PCRs give NO product on WT or donor plasmid "
             "-> band only on correct genomic integration.", ha="center", fontsize=6.9, color=CAP)
    save(fig, "figure_C_donor_genotyping")

if __name__ == "__main__":
    print("Building figures ->", OUT)
    fig_A(); fig_B(); fig_C()
    print("[ok] 3 figures (PNG 300dpi + PDF)")
