#!/usr/bin/env python3
"""
05_design_donor.py -- Build the HDR donor for C-terminal mCherry tagging of TP53.

Architecture (gene-sense, 5'->3'):
  [5' homology arm ~800 bp | ...p53 CDS ending at last sense codon D393]
  [GS-linker (GGGGS)x3 | mCherry CDS, NO stop]
  [3' homology arm ~800 bp | native TGA stop + 3'UTR]
Reading frame: ...D393 - (GGGGS)3 - mCherry - TGA(native). Fusion reuses endogenous
stop + 3'UTR (preserves native post-transcriptional regulation).

PAM re-cut protection:
  * PRIMARY guide g3: its 20-nt protospacer STRADDLES the insertion point, so the insert
    splits the target -- the integrated allele and the donor are NOT a Cas9 substrate.
    No edit needed. We PROVE this by searching the edited allele for the g3 target.
  * BACKUP guide g2: its target lies entirely inside the 5' arm (p53 coding), so it WOULD
    be re-cut. We introduce SILENT (synonymous) mutation(s) that break its NGG PAM, making
    one universal donor safe for either guide. We prove the edit is silent and abolishes g2.

Outputs: scripts/donor.json, output/sequences staged later.
"""
import json, os, math, re
from Bio.Seq import Seq

HERE = os.path.dirname(os.path.abspath(__file__))
locus = json.load(open(os.path.join(HERE, "locus.json")))
gd = json.load(open(os.path.join(HERE, "guide_candidates.json")))["guides"]
mch = json.load(open(os.path.join(HERE, "mcherry_opt.json")))
top = "".join(l.strip() for l in open(os.path.join(HERE, "tp53_genesense.fa")) if not l.startswith(">")).upper()

ins = locus["insertion_point_genesense_idx"]
ARM = 800
FLANK = 500   # extra genomic context beyond each arm for primer anchoring
G = {g["id"]: g for g in gd}
g3, g2 = G["TP53ct_g3"], G["TP53ct_g2"]

SYN = {}  # codon -> list of synonymous codons
from Bio.Data.CodonTable import standard_dna_table as T
aa2cod = {}
for cod, aa in T.forward_table.items():
    aa2cod.setdefault(aa, []).append(cod)
def synonyms(codon):
    aa = str(Seq(codon).translate())
    return [c for c in aa2cod.get(aa, []) if c != codon]

# ---- verify the genomic CDS tail is intron-free & in-frame (codon math safe near stop) ----
tail_nt = top[ins-60:ins]
tail_aa = str(Seq(tail_nt).translate())
assert tail_aa == locus["protein_cterm_12aa"][-20:] or tail_aa.endswith(locus["protein_cterm_12aa"]), \
    f"genomic CDS tail not in-frame/intron-free: {tail_aa} vs {locus['protein_cterm_12aa']}"
print(f"[frame] genomic CDS tail translates {tail_aa} (matches p53 C-term) -> codon math near stop is valid")

def codon_span(x):
    """0-based [start,end) of the CDS codon containing gene-sense position x (x<ins, exon 11)."""
    k = math.ceil((ins - x)/3)
    s = ins - 3*k
    return s, s+3

insert = mch["linker_nt"] + mch["mcherry_nt"]
assert len(insert) % 3 == 0

# ---------- base arms ----------
arm5 = list(top[ins-ARM:ins])         # mutable; ends at last sense codon (D393)
arm3 = top[ins:ins+ARM]               # starts with native stop TGA
assert "".join(arm5[-3:]) == locus["last_sense_codon"]
assert arm3[:3] == locus["stop_codon"]

def target23_present(seq, spacer, pam_strandinfo):
    """Is an intact NGG target for `spacer` present in seq (either strand)? Returns count of exact 23-mer matches w/ NGG."""
    s = seq.upper(); rc = str(Seq(s).reverse_complement())
    n = 0
    for strand_seq in (s, rc):
        i = strand_seq.find(spacer)
        while i != -1:
            pam = strand_seq[i+20:i+23]
            if len(pam) == 3 and pam[1:] == "GG":
                n += 1
            i = strand_seq.find(spacer, i+1)
    return n

# ---------- PRIMARY g3: prove the insertion disrupts it ----------
edited_pre = top[ins-ARM:ins] + insert + arm3        # edited allele w/o backup-blocking edits
g3_in_wt   = target23_present(top[ins-60:ins+60], g3["spacer"], None)
g3_in_edit = target23_present(edited_pre, g3["spacer"], None)
print(f"[g3] intact NGG target near junction: WT={g3_in_wt}  edited-allele={g3_in_edit}  "
      f"-> {'DISRUPTED by insertion (no PAM-block needed)' if g3_in_edit==0 else 'STILL PRESENT (!)'}")
assert g3_in_edit == 0, "g3 target unexpectedly intact in edited allele"

# ---------- BACKUP g2: silent PAM-block in 5' arm ----------
# g2 is +strand; its PAM (NGG) sits in gene-sense at [pam_j, pam_j+3); the GG are at pam_j+1,pam_j+2.
pam_j = g2["pam_j"]
assert top[pam_j+1] == "G" and top[pam_j+2] == "G", "g2 PAM is not NGG as expected"
print(f"\n[g2] +strand PAM NGG at gene-sense {pam_j}-{pam_j+2} (bases {top[pam_j:pam_j+3]}); target in 5' arm")
g2_in_wtarm = target23_present(top[ins-ARM:ins], g2["spacer"], None)
print(f"[g2] intact NGG target in unedited 5' arm: {g2_in_wtarm} (will be re-cut unless blocked)")

# try to break the PAM by a synonymous substitution at pam_j+2 (3rd PAM base) or pam_j+1
def apply_pam_block():
    for pam_pos in (pam_j+2, pam_j+1):                 # prefer breaking the most PAM-proximal G last
        cs, ce = codon_span(pam_pos)
        if cs < ins-ARM:   # codon must be inside the arm
            continue
        orig_codon = "".join(arm5[cs-(ins-ARM):ce-(ins-ARM)])
        rel = pam_pos - cs                              # which base of the codon is the PAM G
        for alt in synonyms(orig_codon):
            if alt[rel] != "G":                         # breaks the G at the PAM position
                # apply
                for k in range(3):
                    arm5[cs-(ins-ARM)+k] = alt[k]
                return dict(strategy="PAM-disrupt", codon_pos_genesense=cs,
                            aa=str(Seq(orig_codon).translate()), orig_codon=orig_codon,
                            new_codon=alt, pam_base_changed=pam_pos)
    return None

block = apply_pam_block()
arm5s = "".join(arm5)
# verify silent: the CDS portion of arm5 (from first full codon after arm start, to ins) translates unchanged
cds_start_in_arm = (ins - ARM) + ((ARM) % 3)          # align to codon boundary (ins is boundary)
# simpler: compare last 60 codons region translation
def cds_tail_translation(seq_arm):
    s = seq_arm[-60:]                                   # last 60 nt = last 20 codons, all exon 11
    return str(Seq(s).translate())
orig_tail = cds_tail_translation(top[ins-ARM:ins])
new_tail  = cds_tail_translation(arm5s)
assert block is not None, "could not find a silent PAM-block for g2"
assert orig_tail == new_tail, f"PAM-block is NOT silent!\n{orig_tail}\n{new_tail}"
g2_after = target23_present(arm5s, g2["spacer"], None)
print(f"[g2] silent block: codon {block['aa']} {block['orig_codon']}->{block['new_codon']} "
      f"at gene-sense {block['codon_pos_genesense']} (changed base {block['pam_base_changed']})")
print(f"[g2] translation of CDS tail unchanged: {new_tail}")
print(f"[g2] intact NGG target after block: {g2_after} -> {'BLOCKED' if g2_after==0 else 'STILL CUTTABLE (!)'}")
assert g2_after == 0
# ensure g3 still disrupted and unaffected by the g2 edit (g2 codon < g3 target)
assert block["pam_base_changed"] < g3["pam_j"], "g2 block overlaps g3 target region!"

# ---------- assemble donor ----------
donor = arm5s + insert + arm3
# BsaI/BsmBI scan: insert must be clean (client requirement); report arm/junction sites too
def _find(seq, motif):
    """1-based positions of `motif` on BOTH strands (overlapping). Counts RECOGNITION
    sites regardless of whether the downstream cut falls within a linear fragment -- a
    site at the very 3' end is still a site once the donor is cloned into a plasmid.
    (Bio.Restriction's linear .search() silently drops such terminal sites.)"""
    s = seq.upper(); rc = str(Seq(motif).reverse_complement())
    pos = [m.start()+1 for m in re.finditer("(?=("+motif+"))", s)]
    pos += [m.start()+1 for m in re.finditer("(?=("+rc+"))", s)]
    return sorted(set(pos))
def sites(seq):
    # BsaI = GGTCTC, BsmBI = CGTCTC (both non-palindromic -> search both strands)
    return {"BsaI": _find(seq, "GGTCTC"), "BsmBI": _find(seq, "CGTCTC")}
insert_sites = sites(insert)
donor_sites = sites(donor)
print(f"\n[donor] length={len(donor)} bp = arm5({len(arm5s)}) + insert({len(insert)}) + arm3({len(arm3)})")
print(f"[donor] insert BsaI/BsmBI: {insert_sites}  (client requirement: zero) ")
print(f"[donor] whole-donor BsaI sites (1-based cut pos): {donor_sites['BsaI']}")
print(f"[donor] whole-donor BsmBI sites: {donor_sites['BsmBI']}")
assert insert_sites["BsaI"]==[] and insert_sites["BsmBI"]==[], "insert has forbidden sites!"

# verify final fusion ORF: p53 tail + linker + mCherry + native stop
fusion_orf = top[ins-60:ins] + insert + arm3[:3]      # last 20 p53 codons + insert + TGA
fusion_aa = str(Seq(fusion_orf).translate())
print(f"\n[fusion] ...{fusion_aa[:20]} ... {fusion_aa[-20:]}")
assert fusion_aa.endswith("MDELYK*"), f"fusion does not end mCherry+stop: {fusion_aa[-10:]}"
assert fusion_aa.count("*") == 1, "unexpected internal stop in fusion"

# ---------- knock-in & WT context for primer design (genomic flanks beyond arms) ----------
left_ctx  = top[ins-ARM-FLANK : ins-ARM]              # genomic, OUTSIDE 5' arm (upstream)
right_ctx = top[ins+ARM : ins+ARM+FLANK]              # genomic, OUTSIDE 3' arm (downstream)
wt_context = top[ins-ARM-FLANK : ins+ARM+FLANK]
ki_context = left_ctx + donor + right_ctx             # full knock-in allele w/ flanks
# coordinate map within ki_context
ki_map = dict(left_ctx_len=len(left_ctx), arm5_len=len(arm5s), insert_len=len(insert),
              arm3_len=len(arm3), right_ctx_len=len(right_ctx),
              j5_pos=len(left_ctx)+len(arm5s),                       # 5' junction (arm5|insert)
              j3_pos=len(left_ctx)+len(arm5s)+len(insert),           # 3' junction (insert|arm3)
              insert_start=len(left_ctx)+len(arm5s),
              insert_end=len(left_ctx)+len(arm5s)+len(insert))

def gs_to_chr1(i): return locus["window_chr_0based_halfopen"][0] + (len(top)-1 - i) + 1

out = dict(
    arm_len=ARM, flank_for_primers=FLANK,
    insertion_point_genesense_idx=ins,
    insertion_chr17=f"{locus['chrom']}:{locus['native_stop_chr17_1based'][0]} (before native stop, minus strand)",
    linker_aa=mch["linker_aa"], linker_nt=mch["linker_nt"],
    mcherry_nt=mch["mcherry_nt"], mcherry_len=len(mch["mcherry_nt"]),
    insert_nt=insert, insert_len=len(insert),
    arm5_nt=arm5s, arm3_nt=arm3,
    arm5_chr17=f"{locus['chrom']}:{gs_to_chr1(ins-1)}-{gs_to_chr1(ins-ARM)} (minus strand, 5' arm)",
    arm3_chr17=f"{locus['chrom']}:{gs_to_chr1(ins+ARM-1)}-{gs_to_chr1(ins)} (minus strand, 3' arm)",
    donor_nt=donor, donor_len=len(donor),
    donor_bsaI=donor_sites["BsaI"], donor_bsmBI=donor_sites["BsmBI"],
    insert_bsaI=insert_sites["BsaI"], insert_bsmBI=insert_sites["BsmBI"],
    fusion_cterm_aa=fusion_aa[-30:],
    primary_guide=dict(id=g3["id"], spacer=g3["spacer"], pam=g3["pam"], strand=g3["strand"],
                       recut_protection="intrinsic: insertion splits the protospacer (no edit needed)",
                       intact_target_in_edited_allele=g3_in_edit),
    backup_guide=dict(id=g2["id"], spacer=g2["spacer"], pam=g2["pam"], strand=g2["strand"],
                      recut_protection="engineered silent PAM-block in 5' arm", pam_block=block,
                      intact_target_after_block=g2_after),
    wt_context=wt_context, ki_context=ki_context, ki_map=ki_map,
)
json.dump(out, open(os.path.join(HERE, "donor.json"), "w"), indent=2)
print(f"\n[ok] donor.json written. Donor {len(donor)} bp; KI context {len(ki_context)} bp (with {FLANK} bp flanks).")
print(f"     5' arm: {out['arm5_chr17']}")
print(f"     3' arm: {out['arm3_chr17']}")
