#!/usr/bin/env python3
"""
03_screen_sites.py  ->  parts_clean.json

Screen every part (regulatory + backbone + coding) for internal BsaI/BsmBI/SapI
on both strands and domesticate the two sites found:
  - chimeric_intron : 1 BsaI motif in the intron BODY (away from 5' donor/3'
                      acceptor/branch point) -> silent single-nt transition.
  - AmpR (bla CDS)  : 1 BsaI motif inside the beta-lactamase ORF -> SYNONYMOUS
                      codon change (E. coli function preserved).
Coding ORFs (HC/LC/pac) were already made site-free during codon optimization.
"""
import json, os
from Bio.Seq import Seq
from Bio.Restriction import BsaI, BsmBI, SapI
from Bio.Data.CodonTable import standard_dna_table as SCT

HERE = os.path.dirname(os.path.abspath(__file__))
MOTIFS = {"BsaI": ["GGTCTC", "GAGACC"], "BsmBI": ["CGTCTC", "GAGACG"], "SapI": ["GCTCTTC", "GAAGAGC"]}
ALL = [m for v in MOTIFS.values() for m in v]

def sites(s):
    s = Seq(s); return {"BsaI": len(BsaI.search(s)), "BsmBI": len(BsmBI.search(s)), "SapI": len(SapI.search(s))}
def nsites(s):    # total Type IIS occurrences (both strands, via Biopython)
    d = sites(s); return d["BsaI"] + d["BsmBI"] + d["SapI"]
def clean(s): return nsites(s) == 0
def find_motif(s):
    hits = [(s.find(m), m) for m in ALL if s.find(m) >= 0]
    return min(hits) if hits else (-1, None)

SYN = {}                                   # codon -> list of synonymous codons
for aa in set(SCT.forward_table.values()):
    cods = [c for c, a in SCT.forward_table.items() if a == aa]
    for c in cods: SYN[c] = cods

def domesticate_noncoding(seq):
    """Break each motif with the single substitution that most reduces site count.
    Tries motif positions centre-first (away from element edges); each accepted
    edit strictly lowers the total count, guaranteeing termination."""
    log = []; s = seq
    while not clean(s):
        i, m = find_motif(s); cur = nsites(s); chosen = None
        for off in sorted(range(len(m)), key=lambda k: abs(k - len(m) // 2)):
            pos = i + off
            for nt in "ACGT":
                if nt == s[pos]: continue
                trial = s[:pos] + nt + s[pos+1:]
                if nsites(trial) < cur:
                    chosen = (pos, s[pos], nt, trial); break
            if chosen: break
        if not chosen: raise RuntimeError("could not domesticate non-coding motif")
        pos, a, b, s = chosen; log.append((pos, a, b))
    return s, log

def domesticate_cds(seq, cds_start):
    """Synonymous substitution(s) that remove motifs inside an embedded CDS."""
    log = []; s = seq; prot0 = str(Seq(seq[cds_start:]).translate())
    while not clean(s):
        i, m = find_motif(s); cur = nsites(s); chosen = None
        for cdx in range((i - cds_start) // 3, (i + len(m) - 1 - cds_start) // 3 + 1):
            cstart = cds_start + cdx * 3
            cod = s[cstart:cstart+3]
            for alt in SYN.get(cod, []):
                if alt == cod: continue
                trial = s[:cstart] + alt + s[cstart+3:]
                if nsites(trial) < cur and str(Seq(trial[cds_start:]).translate()) == prot0:
                    chosen = (cstart, cod, alt, trial); break
            if chosen: break
        if not chosen: raise RuntimeError("could not domesticate CDS motif synonymously")
        cstart, a, b, s = chosen; log.append((cstart, a, b))
    return s, log

parts = json.load(open(os.path.join(HERE, "parts.json")))
opt = json.load(open(os.path.join(HERE, "optimized_cds.json")))

clean_parts = {"dna_parts": {}, "coding_cds": {}, "domestication_log": {}}

for name, d in parts["dna_parts"].items():
    seq = d["seq"]
    if clean(seq):
        clean_parts["dna_parts"][name] = {**d, "seq": seq, "domesticated": False}
        continue
    if name == "AmpR":
        # bla CDS starts after the 105-bp AmpR promoter; verify start codon
        cds_start = 105
        assert seq[cds_start:cds_start+3] == "ATG", f"AmpR CDS frame check: {seq[cds_start:cds_start+3]}"
        new, log = domesticate_cds(seq, cds_start)
        # protein preserved
        assert str(Seq(new[cds_start:]).translate()) == str(Seq(seq[cds_start:]).translate())
    else:
        new, log = domesticate_noncoding(seq)
        assert len(new) == len(seq)
    clean_parts["dna_parts"][name] = {**d, "seq": new, "domesticated": True}
    clean_parts["domestication_log"][name] = [{"pos": p, "from": a, "to": b} for (p, a, b) in log]
    print(f"{name}: domesticated {len(log)} change(s) {log} -> sites now {sites(new)}  (len {len(new)}, unchanged={len(new)==len(seq)})")

# coding ORFs (already site-free)
for k in ("heavy_chain", "light_chain", "pac"):
    cds = opt[k]["cds"]
    assert clean(cds), f"{k} unexpectedly has a site"
    clean_parts["coding_cds"][k] = {"cds": cds, "protein": opt[k]["protein"], "CAI": opt[k]["CAI"], "GC_pct": opt[k]["GC_pct"]}

# global re-verify
allok = all(clean(v["seq"]) for v in clean_parts["dna_parts"].values()) and \
        all(clean(v["cds"]) for v in clean_parts["coding_cds"].values())
clean_parts["all_parts_site_free"] = allok
json.dump(clean_parts, open(os.path.join(HERE, "parts_clean.json"), "w"), indent=2)
print("\nALL parts site-free:", allok, "-> wrote parts_clean.json")
