#!/usr/bin/env python3
"""
05_assemble_construct.py

Assemble the final circular plasmid in silico from the site-free parts + the
selected fusion overhangs, build the 5 synthesizable Golden Gate fragments
(BsaI-flanked), and VERIFY:
  (1) the assembled product carries ZERO internal BsaI/BsmBI/SapI sites,
  (2) each fragment carries exactly one inward GGTCTC and one inward GAGACC
      (=2 BsaI) and zero BsmBI/SapI,
  (3) BsaI digestion of every fragment exposes the designed 4-nt overhang,
  (4) shared overhangs chain the fragments unambiguously, and re-ligating them
      reproduces the assembled plasmid exactly.

Outputs:
  ../output/Trastuzumab_IgG1_GoldenGate_CHO_construct.gb  (annotated, circular)
  fragments.json, assembled.fasta  (intermediates for downstream steps)
"""
import json, os
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio.SeqFeature import SeqFeature, FeatureLocation
from Bio import SeqIO
from Bio.Restriction import BsaI, BsmBI, SapI

HERE = os.path.dirname(os.path.abspath(__file__))
OUT  = os.path.abspath(os.path.join(HERE, "..", "output"))
parts = json.load(open(os.path.join(HERE, "parts_clean.json")))
fus   = json.load(open(os.path.join(HERE, "fusion_sites.json")))
P = {k: v["seq"] for k, v in parts["dna_parts"].items()}
C = parts["coding_cds"]

def rc(s): return str(Seq(s).reverse_complement())
def nsites(s):
    s = Seq(s); return {"BsaI": len(BsaI.search(s)), "BsmBI": len(BsmBI.search(s)), "SapI": len(SapI.search(s))}

KOZAK, STOP = "GCCACC", "TAA"
OH = fus["assignment"]
J1, J2, J3, J4, J5 = OH["J1_intron|HC"], OH["J2_HC|IRES"], OH["J3_polyA|SV40p"], OH["J4_SV40pA|ori"], OH["J5_AmpR|CMV"]

# ---- INNER bodies (the part of each fragment that survives into the plasmid) --
HC_cds, LC_cds, PAC_cds = C["heavy_chain"]["cds"], C["light_chain"]["cds"], C["pac"]["cds"]
assert HC_cds[:3] == "ATG" and LC_cds[:3] == "ATG" and PAC_cds[:3] == "ATG"
assert P["EMCV_IRES"].endswith("ATG")           # IRES terminal ATG = LC start codon

INNER = {
 "F1_promoter":  P["CMV_enh_prom"] + P["chimeric_intron"],
 "F2_heavy":     KOZAK + HC_cds + STOP,
 "F3_IRES_light":P["EMCV_IRES"] + LC_cds[3:] + STOP + P["WPRE"] + P["bGH_polyA"],   # IRES ATG serves as LC Met
 "F4_puro":      P["SV40_promoter"] + KOZAK + PAC_cds + STOP + P["SV40_polyA"],
 "F5_backbone":  P["pUC_ori"] + P["AmpR"],
}
LEFT  = {"F1_promoter": J5, "F2_heavy": J1, "F3_IRES_light": J2, "F4_puro": J3, "F5_backbone": J4}
RIGHT = {"F1_promoter": J1, "F2_heavy": J2, "F3_IRES_light": J3, "F4_puro": J4, "F5_backbone": J5}
ORDER = ["F1_promoter", "F2_heavy", "F3_IRES_light", "F4_puro", "F5_backbone"]

# ---- assembled circular plasmid: INNER_i + RIGHT_OH_i around the circle -------
segments = []   # (label, type, seq) in plasmid order, including overhang scars
def add(label, typ, seq): segments.append((label, typ, seq))
add("CMV enhancer/promoter", "promoter", P["CMV_enh_prom"])
add("chimeric intron", "intron", P["chimeric_intron"])
add(f"fusion site J1 ({J1})", "misc_feature", J1)
add("Kozak", "misc_feature", KOZAK)
add("heavy chain: Ig-kappa leader + trastuzumab VH + human IgG1 (CHO-opt)", "CDS_HC", HC_cds + STOP)
add(f"fusion site J2 ({J2})", "misc_feature", J2)
add("EMCV IRES", "misc_feature", P["EMCV_IRES"])
add("light chain: Ig-kappa leader + trastuzumab VL + kappa C (CHO-opt)", "CDS_LC", LC_cds[3:] + STOP)
add("WPRE", "misc_feature", P["WPRE"])
add("bGH poly(A)", "polyA_signal", P["bGH_polyA"])
add(f"fusion site J3 ({J3})", "misc_feature", J3)
add("SV40 promoter", "promoter", P["SV40_promoter"])
add("Kozak", "misc_feature", KOZAK)
add("puromycin N-acetyltransferase (pac, CHO-opt)", "CDS_PAC", PAC_cds + STOP)
add("SV40 poly(A)", "polyA_signal", P["SV40_polyA"])
add(f"fusion site J4 ({J4})", "misc_feature", J4)
add("pUC origin", "rep_origin", P["pUC_ori"])
add("AmpR (promoter + bla CDS)", "AmpR", P["AmpR"])
add(f"fusion site J5 ({J5})", "misc_feature", J5)

assembled = "".join(s for _, _, s in segments)
sites_asm = nsites(assembled)
print(f"Assembled plasmid: {len(assembled)} bp")
print(f"  internal Type IIS sites: {sites_asm}  ->", "CLEAN" if sites_asm == {'BsaI':0,'BsmBI':0,'SapI':0} else "*** NOT CLEAN ***")

# ---- annotate GenBank --------------------------------------------------------
rec = SeqRecord(Seq(assembled), id="pTRAST_GG", name="pTRAST_GG",
                description="Single-vector trastuzumab IgG1 mAb expression vector for CHO-K1 (BsaI Golden Gate)")
rec.annotations.update({"molecule_type": "DNA", "topology": "circular", "data_file_division": "SYN",
                        "organism": "synthetic construct", "date": "29-JUN-2026"})
TYPEMAP = {"CDS_HC": "CDS", "CDS_LC": "CDS", "CDS_PAC": "CDS", "AmpR": "CDS"}
pos = 0; cds_check = {}
for label, typ, seq in segments:
    start, end = pos, pos + len(seq); pos = end
    if typ in ("CDS_HC", "CDS_LC", "CDS_PAC"):
        cstart = start - 3 if typ == "CDS_LC" else start   # LC initiator ATG is the IRES terminal ATG
        prot = str(Seq(assembled[cstart:end]).translate(to_stop=True))
        q = {"label": label, "translation": prot, "codon_start": 1}
        rec.features.append(SeqFeature(FeatureLocation(cstart, end), type="CDS", qualifiers=q))
        cds_check[typ] = prot
    elif typ == "AmpR":
        rec.features.append(SeqFeature(FeatureLocation(start, start+105), type="promoter",
                                       qualifiers={"label": "AmpR promoter"}))
        rec.features.append(SeqFeature(FeatureLocation(start+105, end), type="CDS",
                                       qualifiers={"label": "AmpR (bla)",
                                                   "translation": str(Seq(seq[105:]).translate(to_stop=True))}))
    else:
        rec.features.append(SeqFeature(FeatureLocation(start, end), type=TYPEMAP.get(typ, typ),
                                       qualifiers={"label": label}))
# fine-grained sub-features on the antibody chains (signal peptide / V / C); use
# types figure A ignores so the construct map is unaffected.
def add_sub(cds_label, blocks):
    def lab(f):
        v = f.qualifiers.get("label", "")
        return v[0] if isinstance(v, list) else v
    cds = next(f for f in rec.features if f.type == "CDS" and lab(f) == cds_label)
    p = int(cds.location.start)
    for name, ln, typ in blocks:
        rec.features.append(SeqFeature(FeatureLocation(p, p + ln), type=typ, qualifiers={"label": name}))
        p += ln
add_sub("heavy chain: Ig-kappa leader + trastuzumab VH + human IgG1 (CHO-opt)",
        [("Ig-kappa signal peptide", 60, "sig_peptide"), ("trastuzumab VH", 360, "region"),
         ("human IgG1 constant CH1-hinge-CH2-CH3", 990, "region")])
add_sub("light chain: Ig-kappa leader + trastuzumab VL + kappa C (CHO-opt)",
        [("Ig-kappa signal peptide", 60, "sig_peptide"), ("trastuzumab VL", 321, "region"),
         ("human kappa constant", 321, "region")])

# verify annotated CDS translations match the intended proteins
prot = json.load(open(os.path.join(HERE, "parts.json")))["proteins"]
assert cds_check["CDS_HC"] == prot["heavy_chain"]["aa"], "HC CDS annotation translation mismatch"
assert cds_check["CDS_LC"] == prot["light_chain"]["aa"], "LC CDS annotation translation mismatch"
assert cds_check["CDS_PAC"] == prot["pac"]["aa"], "pac CDS annotation translation mismatch"
print("  annotated CDS translations match HC/LC/pac proteins: True")

os.makedirs(OUT, exist_ok=True)
GB = os.path.join(OUT, "Trastuzumab_IgG1_GoldenGate_CHO_construct.gb")
SeqIO.write(rec, GB, "genbank")

# ---- build synthesizable BsaI-flanked fragments ------------------------------
PAD, SP = "CAGTCA", "A"                 # 6-bp end pad (>=6 bp lets BsaI cut near ends); 1-nt BsaI spacer
BSAI_F, BSAI_R = "GGTCTC", "GAGACC"     # GGTCTC(1/5); GAGACC = rc(GGTCTC)
frags = {}
for name in ORDER:
    body = LEFT[name] + INNER[name] + RIGHT[name]                    # overhang..INNER..overhang (survives)
    synth = PAD + BSAI_F + SP + LEFT[name] + INNER[name] + RIGHT[name] + SP + BSAI_R + PAD
    frags[name] = {"left_oh": LEFT[name], "right_oh": RIGHT[name],
                   "inner_len": len(INNER[name]), "survives": body, "synth": synth, "synth_len": len(synth)}

# ---- VERIFY fragments --------------------------------------------------------
ok = (sites_asm == {"BsaI": 0, "BsmBI": 0, "SapI": 0})
for name in ORDER:
    f = frags[name]; s = f["synth"]; sit = nsites(s)
    # BsaI exactly 2 (one forward GGTCTC, one reverse GAGACC), no BsmBI/SapI
    nF, nR = s.count(BSAI_F), s.count(BSAI_R)
    # overhang exposed by left (forward) site = 4 nt after GGTCTC + 1 spacer
    iF = s.index(BSAI_F); left_exposed = s[iF+6+1: iF+6+1+4]
    # overhang exposed by right (reverse) site = 4 nt immediately 5' of (spacer + GAGACC)
    iR = s.rindex(BSAI_R); right_exposed = s[iR-1-4: iR-1]
    good = (sit["BsaI"] == 2 and sit["BsmBI"] == 0 and sit["SapI"] == 0 and nF == 1 and nR == 1
            and left_exposed == f["left_oh"] and right_exposed == f["right_oh"])
    ok = ok and good
    print(f"  {name:14} synth={f['synth_len']:>4}bp inner={f['inner_len']:>4}  "
          f"BsaI={sit['BsaI']} BsmBI={sit['BsmBI']} SapI={sit['SapI']}  "
          f"OH L/R={left_exposed}/{right_exposed} (want {f['left_oh']}/{f['right_oh']})  {'OK' if good else 'FAIL'}")

# ---- VERIFY in-silico re-ligation reproduces the plasmid ---------------------
# chain fragments by shared overhangs starting at F1, concatenating INNER+RIGHT_OH
chain = ""
for name in ORDER:
    assert frags[name]["right_oh"] == LEFT[ORDER[(ORDER.index(name)+1) % 5]], f"overhang chain break at {name}"
    chain += INNER[name] + RIGHT[name]
religated_ok = (chain == assembled)
print(f"\nRe-ligation reproduces assembled plasmid exactly: {religated_ok}")
ok = ok and religated_ok

json.dump({"assembled_len": len(assembled), "assembled_sites": sites_asm,
           "fragments": {k: {kk: vv for kk, vv in v.items() if kk != 'survives'} for k, v in frags.items()},
           "fragment_order": ORDER, "overhangs": OH, "all_checks_pass": ok},
          open(os.path.join(HERE, "fragments.json"), "w"), indent=2)
open(os.path.join(HERE, "assembled.fasta"), "w").write(f">pTRAST_GG_assembled {len(assembled)}bp circular\n{assembled}\n")
print("\nALL assembly checks pass:", ok)
print("wrote", GB)
