#!/usr/bin/env python3
"""
09_figures.py  ->  ../output/figures/*.png (300 dpi) + *.pdf (vector)

Figure A  construct map        - where every element + fragment + fusion overhang sits (PLACEMENT)
Figure B  codon-optimization   - CAI / GC / size per ORF vs CHO-K1 targets (EVIDENCE)
Figure C  overhang fidelity    - 5x5 ligation-specificity matrix, Potapov/Pryor data (EVIDENCE)
"""
import json, os
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch, Rectangle
import numpy as np
import openpyxl
from report_style import PALETTE

HERE = os.path.dirname(os.path.abspath(__file__))
FIG = os.path.abspath(os.path.join(HERE, "..", "output", "figures"))
os.makedirs(FIG, exist_ok=True)
NAVY = "#1B2A4A"
plt.rcParams.update({"font.family": "DejaVu Sans", "font.size": 9, "svg.fonttype": "none",
                     "axes.linewidth": 0.8, "pdf.fonttype": 42})

opt = json.load(open(os.path.join(HERE, "optimized_cds.json")))
fus = json.load(open(os.path.join(HERE, "fusion_sites.json")))
frag = json.load(open(os.path.join(HERE, "fragments.json")))
from Bio import SeqIO
gb = SeqIO.read(os.path.abspath(os.path.join(HERE, "..", "output", "Trastuzumab_IgG1_GoldenGate_CHO_construct.gb")), "genbank")
L = len(gb.seq)

def save(fig, name):
    fig.savefig(os.path.join(FIG, name + ".png"), dpi=300, bbox_inches="tight")
    fig.savefig(os.path.join(FIG, name + ".pdf"), bbox_inches="tight")
    plt.close(fig)

# =================================================================== Figure A
# category -> colour (cohesive palette anchored on the report navy)
CAT = {"CMV enhancer/promoter": ("promoter", "#1B2A4A"), "SV40 promoter": ("promoter", "#3B5B92"),
       "AmpR promoter": ("promoter", "#7E57C2"), "chimeric intron": ("intron", "#8FA3C8"),
       "EMCV IRES": ("IRES", "#0E7C7B"), "WPRE": ("regulatory", "#5BA89A"),
       "bGH poly(A)": ("poly(A)", "#B0843B"), "SV40 poly(A)": ("poly(A)", "#CDA256"),
       "pUC origin": ("ori", "#888888"),
       "heavy chain": ("antibody", "#C0392B"), "light chain": ("antibody", "#E07B54"),
       "puromycin": ("selection", "#2E8B57"), "AmpR (bla)": ("selection", "#9B59B6")}
def cat_of(lab):
    for k, v in CAT.items():
        if lab.startswith(k) or k.split()[0] in lab:
            return v
    return ("other", "#bbbbbb")

elements, junctions = [], []
for f in gb.features:
    lab = f.qualifiers.get("label", ["?"])[0]
    s, e = int(f.location.start), int(f.location.end)
    if f.type == "misc_feature" and lab.startswith("fusion site"):
        junctions.append((s, lab.split("(")[1].rstrip(")"), lab.split()[2]))  # (pos, overhang, Jname)
    elif "Kozak" in lab:
        continue                                                              # 6-bp, too small to label
    elif f.type in ("promoter", "intron", "polyA_signal", "rep_origin", "CDS", "misc_feature"):  # incl. EMCV IRES, WPRE
        short = (lab.split(":")[0].replace(" enhancer/promoter", " prom").replace(" poly(A)", " pA")
                 .replace("puromycin N-acetyltransferase (pac, CHO-opt)", "pac (PuroR)")
                 .replace("chimeric intron", "intron").replace("pUC origin", "pUC ori")
                 .replace(" (promoter + bla CDS)", "").replace("AmpR (bla)", "AmpR (bla)"))
        if short.startswith("heavy chain"): short = "heavy chain (IgG1)"
        if short.startswith("light chain"): short = "light chain (kappa)"
        elements.append((s, e, short, cat_of(lab)[1], cat_of(lab)[0]))

# fragment spans (inner + right overhang), tiling the plasmid
ORDER = frag["fragment_order"]
spans, x = [], 0
for n in ORDER:
    il = frag["fragments"][n]["inner_len"] + 4
    spans.append((x, x + il, n, frag["fragments"][n]["inner_len"])); x += il

fig, ax = plt.subplots(figsize=(11, 4.6))
BACK_Y, BACK_H = 0.0, 0.46
# backbone elements
for i, (s, e, lab, col, cat) in enumerate(elements):
    ax.add_patch(FancyBboxPatch((s, BACK_Y), e - s, BACK_H, boxstyle="round,pad=0,rounding_size=12",
                                fc=col, ec="white", lw=0.8, mutation_aspect=0.02, zorder=3))
    # staggered labels above (3 levels) with leader lines
    lvl = [0.95, 1.45, 1.95][i % 3]
    xc = (s + e) / 2
    ax.plot([xc, xc], [BACK_Y + BACK_H, lvl - 0.06], color="#999", lw=0.5, zorder=2)
    ax.text(xc, lvl, lab, ha="center", va="bottom", fontsize=7.8, color=NAVY, fontweight="bold")
# junction markers (overhangs) below the backbone
for pos, oh, jn in junctions:
    ax.plot([pos, pos], [BACK_Y - 0.18, BACK_Y], color=NAVY, lw=1.0, zorder=4)
    ax.text(pos, BACK_Y - 0.24, f"{jn}\n{oh}", ha="center", va="top", fontsize=6.8,
            color=NAVY, family="monospace")
# fragment track
FRAG_Y = -1.35
fcols = ["#1B2A4A", "#0E7C7B", "#B0843B", "#2E8B57", "#7E57C2"]
for (s, e, n, il), c in zip(spans, fcols):
    ax.add_patch(Rectangle((s, FRAG_Y), e - s, 0.34, fc=c, ec="white", lw=1.2, alpha=0.92, zorder=3))
    ax.text((s + e) / 2, FRAG_Y + 0.17, f"{n.split('_')[0]}\n{il} bp", ha="center", va="center",
            fontsize=7.2, color="white", fontweight="bold")
ax.text(0, FRAG_Y - 0.18, "Golden Gate fragments (BsaI):", ha="left", va="top", fontsize=7.5, color=NAVY, style="italic")

# scale bar
ax.plot([0, L], [-1.95, -1.95], color="#444", lw=0.8)
for tick in range(0, L + 1, 1000):
    ax.plot([tick, tick], [-1.95, -2.02], color="#444", lw=0.8)
    ax.text(tick, -2.12, f"{tick}", ha="center", va="top", fontsize=6.5, color="#444")
ax.text(L, -2.12, "bp", ha="left", va="top", fontsize=6.5, color="#444")

# legend (category swatches)
seen = {}
for _, _, _, col, cat in elements:
    seen.setdefault(cat, col)
seen["fragment"] = "#cccccc"
from matplotlib.patches import Patch
handles = [Patch(fc=col, ec="white", label=cat) for cat, col in seen.items() if cat != "fragment"]
ax.legend(handles=handles, loc="upper center", bbox_to_anchor=(0.5, 1.16), ncol=len(handles),
          fontsize=6.8, frameon=False, handlelength=1.1, columnspacing=1.1)
ax.set_xlim(-120, L + 160); ax.set_ylim(-2.35, 2.5)
ax.axis("off")
ax.set_title(f"Single-vector trastuzumab IgG1 expression construct  -  {L} bp, one-pot BsaI Golden Gate (5 fragments)",
             fontsize=10.5, color=NAVY, fontweight="bold", pad=16)
save(fig, "figure_A_construct_map")

# =================================================================== Figure B
fig, axes = plt.subplots(1, 3, figsize=(11, 3.4))
orfs = [("Heavy chain", "heavy_chain"), ("Light chain", "light_chain"), ("pac (PuroR)", "pac")]
labels = [o[0] for o in orfs]
# B1 CAI
cai = [opt[k]["CAI"] for _, k in orfs]
b = axes[0].bar(labels, cai, color=["#C0392B", "#E07B54", "#2E8B57"], width=0.6, zorder=3)
axes[0].axhline(0.85, ls="--", lw=1.2, color=NAVY); axes[0].text(2.4, 0.857, "target 0.85", ha="right", fontsize=7.5, color=NAVY)
axes[0].set_ylim(0.7, 1.0); axes[0].set_ylabel("CAI (CHO-K1)"); axes[0].set_title("Codon adaptation index", fontsize=9.5, color=NAVY)
for r, v in zip(b, cai): axes[0].text(r.get_x()+r.get_width()/2, v+0.004, f"{v:.3f}", ha="center", fontsize=7.5, fontweight="bold")
# B2 GC%
gcv = [opt[k]["GC_pct"] for _, k in orfs]
b = axes[1].bar(labels, gcv, color=["#C0392B", "#E07B54", "#2E8B57"], width=0.6, zorder=3)
axes[1].axhspan(40, 60, color="#2E8B57", alpha=0.10); axes[1].axhline(60, ls=":", color="#888", lw=1)
axes[1].set_ylim(0, 80); axes[1].set_ylabel("GC %"); axes[1].set_title("GC content", fontsize=9.5, color=NAVY)
for r, v in zip(b, gcv): axes[1].text(r.get_x()+r.get_width()/2, v+1.2, f"{v:.1f}", ha="center", fontsize=7.5, fontweight="bold")
# B3 fragment sizes
fl = [frag["fragments"][n]["synth_len"] for n in ORDER]
names = [n.split("_")[0] for n in ORDER]
b = axes[2].bar(names, fl, color=["#1B2A4A", "#0E7C7B", "#B0843B", "#2E8B57", "#7E57C2"], width=0.62, zorder=3)
axes[2].set_ylabel("Synth length (bp)"); axes[2].set_title("Golden Gate fragments", fontsize=9.5, color=NAVY)
for r, v in zip(b, fl): axes[2].text(r.get_x()+r.get_width()/2, v+25, f"{v}", ha="center", fontsize=7.2, fontweight="bold")
axes[2].set_ylim(0, max(fl)*1.18)
for a in axes:
    a.spines[["top", "right"]].set_visible(False); a.tick_params(labelsize=8)
    for lab in a.get_xticklabels(): lab.set_rotation(12)
fig.suptitle("Codon optimization for CHO-K1 and fragment QC  (all ORFs: 0 internal BsaI / BsmBI / SapI)",
             fontsize=10.5, color=NAVY, fontweight="bold", y=1.03)
fig.tight_layout()
save(fig, "figure_B_codon_optimization_qc")

# =================================================================== Figure C
ws = openpyxl.load_workbook(os.path.join(HERE, "fidelity_data", "pone.0238592.s001.xlsx"),
                            read_only=True, data_only=True)["S1 Table. BsaI-HFv2"]
rows = list(ws.iter_rows(values_only=True)); cols = [c for c in rows[0][1:] if c]
M = {}
for r in rows[1:]:
    if r[0] is None: continue
    for cl, v in zip(cols, r[1:]): M[(r[0], cl)] = v or 0
comp = {"A": "T", "T": "A", "G": "C", "C": "G"}
rc = lambda x: "".join(comp[b] for b in reversed(x))
lig = lambda a, b: M.get((a, b), 0)
OHs = fus["overhangs"]
n = len(OHs)
mat = np.zeros((n, n))
for i in range(n):
    for j in range(n):
        if i == j:
            mat[i, j] = lig(OHs[i], rc(OHs[i]))                  # on-target
        else:
            mat[i, j] = lig(OHs[i], rc(OHs[j])) + lig(OHs[j], rc(OHs[i]))  # cross-talk

fig, ax = plt.subplots(figsize=(6.0, 5.2))
im = ax.imshow(mat, cmap="GnBu", aspect="equal")
ax.set_xticks(range(n)); ax.set_yticks(range(n))
ax.set_xticklabels(OHs, family="monospace", fontsize=9); ax.set_yticklabels(OHs, family="monospace", fontsize=9)
ax.set_xlabel("fusion site (overhang)", color=NAVY); ax.set_ylabel("fusion site (overhang)", color=NAVY)
vmax = mat.max()
for i in range(n):
    for j in range(n):
        v = int(mat[i, j])
        ax.text(j, i, v, ha="center", va="center", fontsize=8.5,
                color="white" if v > vmax*0.55 else NAVY,
                fontweight="bold" if i == j else "normal")
ax.set_title("Fusion-site ligation specificity\n(Potapov/Pryor BsaI-HFv2; diagonal = on-target, off-diagonal = cross-talk)",
             fontsize=9.5, color=NAVY, pad=10)
cb = fig.colorbar(im, fraction=0.046, pad=0.04); cb.set_label("T4 ligase ligation count", fontsize=8)
ax.text(0.5, -0.16, f"Predicted set fidelity {fus['predicted_set_fidelity_pct']:.1f}%  -  all off-diagonal cross-talk = 0",
        transform=ax.transAxes, ha="center", fontsize=8.5, color=NAVY, fontweight="bold")
fig.tight_layout()
save(fig, "figure_C_overhang_fidelity")

print("wrote 3 figures (PNG 300dpi + PDF) to", FIG)
for f in sorted(os.listdir(FIG)): print("  ", f)
