#!/usr/bin/env python3
"""
04_codon_optimize_mcherry.py -- Fetch authoritative mCherry, codon-optimize for human,
forbid BsaI/BsmBI sites. Also define the codon-optimized flexible GS linker.

mCherry reference: GenBank AY678264 (Shaner et al. 2004, Nat Biotechnol) -- the original
mCherry. We translate the fetched CDS and verify it matches the canonical 236-aa mCherry
before using it (no hand-typed protein sequence trusted blindly).

For an ENDOGENOUS C-terminal fusion we reuse the native TP53 stop codon, so the mCherry
ORF here carries NO stop codon (it reads through into the genomic TGA in the 3' arm).

Outputs: scripts/mcherry_opt.json
"""
import json, os, time, sys
from Bio import Entrez, SeqIO
from Bio.Seq import Seq
from Bio.Restriction import BsaI, BsmBI
from dnachisel import (DnaOptimizationProblem, CodonOptimize, EnforceTranslation,
                       AvoidPattern, EnforceGCContent, EnzymeSitePattern,
                       AvoidHairpins, reverse_translate)

Entrez.email = "bioinformatics@labterna.com"
HERE = os.path.dirname(os.path.abspath(__file__))

# Canonical mCherry protein (FPbase / Shaner 2004), used ONLY to validate the fetched CDS.
MCHERRY_CANON = ("MVSKGEEDNMAIIKEFMRFKVHMEGSVNGHEFEIEGEGEGRPYEGTQTAKLKVTKGGPLPFAWDILSPQF"
                 "MYGSKAYVKHPADIPDYLKLSFPEGFKWERVMNFEDGGVVTVTQDSSLQDGEFIYKVKLRGTNFPSDGPV"
                 "MQKKTMGWEASSERMYPEDGALKGEIKQRLKLKDGGHYDAEVKTTYKAKKPVQLPGAYNVNIKLDITSHN"
                 "EDYTIVEQYERAEGRHSTGGMDELYK")

def fetch_mcherry_protein():
    """Fetch AY678264, return its CDS protein (stop stripped). Fallback to canonical."""
    try:
        h = Entrez.efetch(db="nucleotide", id="AY678264", rettype="gb", retmode="text")
        rec = SeqIO.read(h, "genbank"); h.close()
        cds = next(f for f in rec.features if f.type == "CDS")
        prot = str(cds.qualifiers.get("translation", [""])[0])
        if not prot:
            prot = str(Seq(str(cds.extract(rec.seq))).translate(to_stop=True))
        return prot, "GenBank AY678264 (Shaner et al. 2004)"
    except Exception as e:
        sys.stderr.write(f"[warn] AY678264 fetch failed ({e}); using FPbase canonical mCherry\n")
        return MCHERRY_CANON, "FPbase canonical mCherry (network fetch unavailable)"

def has_sites(seq):
    s = Seq(seq)
    return {"BsaI": len(BsaI.search(s)), "BsmBI": len(BsmBI.search(s))}

def main():
    prot, src = fetch_mcherry_protein()
    print(f"[mCherry] source: {src}; length={len(prot)} aa")
    # validate against canonical (allow it to BE the canonical)
    if prot != MCHERRY_CANON:
        # report differences but proceed only if highly similar (engineered variants differ slightly)
        diffs = [(i+1, a, b) for i,(a,b) in enumerate(zip(prot, MCHERRY_CANON)) if a!=b]
        print(f"[mCherry] NOTE: fetched protein differs from canonical at {len(diffs)} positions: {diffs[:10]}")
        if len(prot) != len(MCHERRY_CANON) or len(diffs) > 3:
            print("[mCherry] using canonical mCherry to be safe")
            prot = MCHERRY_CANON; src += " (validated vs canonical; used canonical)"
    else:
        print("[mCherry] fetched protein is byte-identical to canonical mCherry 236-aa. OK")

    target_len = len(prot)*3
    # initial reverse translation (human) as a starting sequence
    seq0 = reverse_translate(prot, table="Standard", randomize_codons=False)
    problem = DnaOptimizationProblem(
        sequence=seq0,
        constraints=[
            EnforceTranslation(translation=prot),
            AvoidPattern(EnzymeSitePattern("BsaI")),    # GGTCTC (both strands)
            AvoidPattern(EnzymeSitePattern("BsmBI")),   # CGTCTC (both strands)
            EnforceGCContent(mini=0.35, maxi=0.70, window=50),
            AvoidPattern("9xA"), AvoidPattern("9xT"),
            AvoidPattern("6xG"), AvoidPattern("6xC"),
        ],
        objectives=[CodonOptimize(species="h_sapiens", method="use_best_codon")],
        logger=None,
    )
    problem.resolve_constraints()
    problem.optimize()
    mcherry_nt = str(problem.sequence)
    # verify translation identical and no stop
    tr = str(Seq(mcherry_nt).translate())
    assert tr == prot, f"translation mismatch:\n{tr}\n{prot}"
    sites = has_sites(mcherry_nt)
    assert sites["BsaI"] == 0 and sites["BsmBI"] == 0, f"restriction sites remain: {sites}"
    gc = round(100*(mcherry_nt.count("G")+mcherry_nt.count("C"))/len(mcherry_nt), 1)
    print(f"[mCherry-opt] {len(mcherry_nt)} nt, GC={gc}%, BsaI/BsmBI sites={sites}, translation OK (no stop)")

    # ---- flexible GS linker: (GGGGS)x3 = 15 aa ----
    # Codon-DIVERSIFIED by hand (not use_best_codon) so the 3 repeats do NOT form a perfect
    # 45-bp tandem repeat -- tandem repeats are flagged by synthesis vendors and can recombine.
    # Three GGGGS repeats, each with distinct Gly/Ser codons:
    linker_aa = "GGGGS"*3
    linker_nt = ("GGAGGCGGTGGGTCT"   # G  G  G  G  S
                 "GGCGGTGGAGGCAGC"   # G  G  G  G  S
                 "GGTGGAGGCGGATCA")  # G  G  G  G  S
    assert str(Seq(linker_nt).translate()) == linker_aa, "linker translation mismatch"
    lsites = has_sites(linker_nt)
    assert lsites["BsaI"]==0 and lsites["BsmBI"]==0, f"linker has sites {lsites}"
    import re
    assert not re.search(r"(.)\1{4,}", linker_nt), "linker has a >=5 homopolymer run"
    # no exact tandem repeat of a full 15-mer repeat unit
    assert len({linker_nt[0:15], linker_nt[15:30], linker_nt[30:45]}) == 3, "linker repeats identical"
    print(f"[linker] {linker_aa} -> {linker_nt} ({len(linker_nt)} nt) sites={lsites} (codon-diversified)")

    # combined insert coding (linker + mCherry, NO stop -- reuses native TGA)
    insert_nt = linker_nt + mcherry_nt
    isites = has_sites(insert_nt)
    print(f"[insert linker+mCherry] {len(insert_nt)} nt, BsaI/BsmBI={isites}, "
          f"in-frame={len(insert_nt)%3==0}, translates={str(Seq(insert_nt).translate())[:8]}...{str(Seq(insert_nt).translate())[-8:]}")
    assert isites["BsaI"]==0 and isites["BsmBI"]==0 and len(insert_nt)%3==0

    out = dict(mcherry_protein=prot, mcherry_protein_len=len(prot), mcherry_source=src,
               mcherry_nt=mcherry_nt, mcherry_gc=gc, mcherry_sites=sites,
               linker_aa=linker_aa, linker_nt=linker_nt,
               insert_coding_nt=insert_nt, insert_sites=isites,
               note="mCherry CDS carries NO stop codon; fusion reuses native TP53 TGA in 3' arm")
    json.dump(out, open(os.path.join(HERE, "mcherry_opt.json"), "w"), indent=2)
    print("[ok] wrote mcherry_opt.json")

if __name__ == "__main__":
    main()
