#!/usr/bin/env python3
"""
02_codon_optimize.py  ->  optimized_cds.json

Codon-optimize the three coding ORFs (heavy chain, light chain, pac) for
Cricetulus griseus CHO-K1 (NCBI taxid 10029, Kazusa usage), targeting CAI > 0.85
while removing all internal BsaI / BsmBI / SapI sites, holding GC in-window, and
avoiding homopolymer runs that hurt gene synthesis.

The two secretion signal peptides (identical 20-aa leader on HC and LC) are
de-duplicated at the DNA level so the construct carries no 60-bp direct repeat.
"""
import json, os, sys, math
import python_codon_tables as pct
from Bio.Seq import Seq
from Bio.Restriction import BsaI, BsmBI, SapI
from dnachisel import (DnaOptimizationProblem, EnforceTranslation, AvoidPattern,
                       EnforceGCContent, CodonOptimize, reverse_translate, EnzymeSitePattern)

HERE = os.path.dirname(os.path.abspath(__file__)); REF = os.path.join(HERE, "refs")
TABLE_JSON = os.path.join(REF, "cho_codon_usage_10029.json")

# ---- CHO codon usage table (cache for offline re-run) ------------------------
if os.path.exists(TABLE_JSON):
    CHO = json.load(open(TABLE_JSON))
else:
    raw = pct.download_codons_table(taxid=10029)           # Kazusa, RNA codons
    CHO = {aa: {cod.replace("U", "T"): f for cod, f in d.items()} for aa, d in raw.items()}
    json.dump(CHO, open(TABLE_JSON, "w"), indent=2)
print("CHO table loaded; Leu CTG fraction =", CHO["L"]["CTG"])

# relative adaptiveness w(codon) = freq/max(freq in synonymous family)
W = {}
for aa, d in CHO.items():
    if aa == "*":            # ignore stop for CAI
        continue
    mx = max(d.values())
    for cod, f in d.items():
        W[cod] = (f / mx) if mx > 0 else 0.0

def cai(dna):
    vals = []
    for i in range(0, len(dna) - len(dna) % 3, 3):
        cod = dna[i:i+3]
        aa = str(Seq(cod).translate())
        if aa in ("M", "W", "*"):     # single-codon AAs and stop excluded
            continue
        w = W.get(cod, 0.0)
        if w > 0:
            vals.append(math.log(w))
    return math.exp(sum(vals) / len(vals)) if vals else 0.0

def n_sites(dna):
    s = Seq(dna)
    return {"BsaI": len(BsaI.search(s)), "BsmBI": len(BsmBI.search(s)), "SapI": len(SapI.search(s))}

def gc(dna):
    return 100.0 * (dna.count("G") + dna.count("C")) / len(dna)

# ---- optimization ------------------------------------------------------------
ENZ = [EnzymeSitePattern(e) for e in ("BsaI", "BsmBI", "SapI")]

def optimize(protein, label):
    seed = reverse_translate(protein)                       # any valid CDS as start point
    constraints = [EnforceTranslation()] \
        + [AvoidPattern(p) for p in ENZ] \
        + [AvoidPattern(h) for h in ("7xA", "7xT", "7xG", "7xC")] \
        + [EnforceGCContent(mini=0.40, maxi=0.70, window=50)]
    prob = DnaOptimizationProblem(
        sequence=seed,
        constraints=constraints,
        objectives=[CodonOptimize(codon_usage_table=CHO, method="use_best_codon")],
        logger=None)
    prob.resolve_constraints()
    prob.optimize()
    dna = prob.sequence
    assert str(Seq(dna).translate()) == protein, f"{label}: translation drift"
    return dna

parts = json.load(open(os.path.join(HERE, "parts.json")))
HC_p = parts["proteins"]["heavy_chain"]["aa"]
LC_p = parts["proteins"]["light_chain"]["aa"]
PAC_p = parts["proteins"]["pac"]["aa"]

HC = optimize(HC_p, "HC")
LC = optimize(LC_p, "LC")
PAC = optimize(PAC_p, "pac")

# ---- de-duplicate the two signal-peptide DNA segments (first 60 bp) ----------
def alt_codon(aa, avoid):
    """highest-CHO-freq synonymous codon for aa that differs from `avoid`."""
    opts = sorted(CHO[aa].items(), key=lambda kv: -kv[1])
    for cod, _ in opts:
        if cod != avoid:
            return cod
    return avoid                                            # M/W: no alternative

swaps = 0
LC_list = [LC[i:i+3] for i in range(0, len(LC), 3)]
HC_sp = [HC[i:i+3] for i in range(0, 60, 3)]
for i in range(20):                                         # 20-aa signal peptide
    aa = str(Seq(LC_list[i]).translate())
    if LC_list[i] == HC_sp[i] and len(CHO[aa]) > 1:
        cand = alt_codon(aa, HC_sp[i])
        trial = "".join(LC_list[:i] + [cand] + LC_list[i+1:])
        if n_sites(trial) == {"BsaI": 0, "BsmBI": 0, "SapI": 0} and "7x" not in trial:
            LC_list[i] = cand; swaps += 1
LC = "".join(LC_list)
assert str(Seq(LC).translate()) == LC_p, "LC translation drift after SP de-dup"
sp_identity = sum(a == b for a, b in zip(HC[:60], LC[:60]))
print(f"SP de-dup: {swaps} codons swapped; HC/LC SP nt identity now {sp_identity}/60")

# ---- verify + report ---------------------------------------------------------
out = {}
ok = True
for label, prot, dna in [("heavy_chain", HC_p, HC), ("light_chain", LC_p, LC), ("pac", PAC_p, PAC)]:
    sites = n_sites(dna); c = cai(dna); g = gc(dna)
    passed = (sites == {"BsaI": 0, "BsmBI": 0, "SapI": 0}) and c > 0.85
    ok = ok and passed
    out[label] = {"protein": prot, "cds": dna, "len_bp": len(dna), "len_aa": len(prot),
                  "CAI": round(c, 4), "GC_pct": round(g, 2), "sites": sites, "pass": passed}
    print(f"{label:12} len={len(dna):4}bp  CAI={c:.3f}  GC={g:.1f}%  sites={sites}  {'PASS' if passed else 'FAIL'}")

out["_meta"] = {"host": "Cricetulus griseus CHO-K1 (taxid 10029, Kazusa)",
                "objective": "use_best_codon (max CAI)",
                "constraints": "no BsaI/BsmBI/SapI; GC 40-70%/50bp window; no >=7 homopolymer",
                "sp_dedup_swaps": swaps, "sp_nt_identity_after": f"{sp_identity}/60"}
json.dump(out, open(os.path.join(HERE, "optimized_cds.json"), "w"), indent=2)
print("\nwrote optimized_cds.json", "(ALL PASS)" if ok else "(CHECK FAILURES)")
sys.exit(0 if ok else 1)
