#!/usr/bin/env python3
"""
04_fusion_sites.py  ->  fusion_sites.json

Select 5 mutually-orthogonal BsaI fusion-site overhangs (one per fragment
junction of the 5-fragment assembly) from the Potapov/Pryor ligation-fidelity
data (BsaI-HFv2 x T4 ligase, 37 C; Pryor 2020 S1 Table, extending Potapov 2018).

Verifies: non-palindromic, no reverse-complement collision, pairwise Hamming
distance, and predicted set fidelity. The fidelity estimator is CALIBRATED by
reproducing the ~81 % value Pryor 2020 reports for the common 11-overhang plant
(MoClo) set before it is trusted for selection.
"""
import json, os, itertools, openpyxl
HERE = os.path.dirname(os.path.abspath(__file__))
XLSX = os.path.join(HERE, "fidelity_data", "pone.0238592.s001.xlsx")

# ---- load 256x256 ligation-frequency matrix ---------------------------------
ws = openpyxl.load_workbook(XLSX, read_only=True, data_only=True)["S1 Table. BsaI-HFv2"]
rows = list(ws.iter_rows(values_only=True))
cols = [c for c in rows[0][1:] if c]
M = {}
for r in rows[1:]:
    rl = r[0]
    if rl is None: continue
    for cl, v in zip(cols, r[1:]):
        M[(rl, cl)] = v or 0
OH = cols[:]                                  # all 256 overhangs
COMP = {"A": "T", "T": "A", "G": "C", "C": "G"}
def rc(x): return "".join(COMP[b] for b in reversed(x))
def is_pal(x): return x == rc(x)
def lig(a, b): return M.get((a, b), 0)

def fidelity(S):
    """NEB-style predicted fraction of correct ligations for fusion-site set S.
    Product over the N fusion sites (top-strand overhang of each); for each site
    the off-target competitors are every overhang species actually present in the
    reaction, i.e. both strands (O and rc(O)) of all sites in the pool."""
    pool = []
    for o in S:
        pool.append(o)
        if not is_pal(o): pool.append(rc(o))
    f = 1.0
    for i in S:                              # one term per junction (not per strand)
        total = sum(lig(i, j) for j in pool)
        if total == 0: return 0.0
        f *= lig(i, rc(i)) / total
    return f

def hamming(a, b): return sum(x != y for x, y in zip(a, b))

# ---- CALIBRATION: reproduce Pryor 2020 plant/MoClo set fidelity (~0.81) -----
PLANT = ["GGAG", "TGAC", "TCCC", "TACT", "CCAT", "AATG", "AGCC", "TTCG", "GCTT", "GGTA", "CGCT"]
plant_fid = fidelity(PLANT)
print(f"Calibration: plant/MoClo 11-set predicted fidelity = {plant_fid*100:.1f}%  (Pryor 2020 reports ~81%)")
assert 0.70 <= plant_fid <= 0.87, "fidelity estimator not calibrated to published value"

# ---- candidate pool: non-palindromic, strong on-target, GC 25-75% -----------
def gc(x): return (x.count("G") + x.count("C"))
cand = [o for o in OH if not is_pal(o) and lig(o, rc(o)) >= 300 and 1 <= gc(o) <= 3]
# keep one of each {o, rc(o)} pair to avoid trivially proposing complements
seen, pool = set(), []
for o in sorted(cand, key=lambda x: -lig(x, rc(x))):
    if o in seen or rc(o) in seen: continue
    seen.add(o); pool.append(o)
print(f"Candidate overhangs after palindrome/strength/GC filter: {len(pool)}")

# ---- greedy build to 5, then exhaustive local improvement -------------------
N = 5
def valid(S):
    # distinct, no member is RC of another, pairwise Hamming >= 2
    for a, b in itertools.combinations(S, 2):
        if a == b or a == rc(b) or hamming(a, b) < 2: return False
    return True

# greedy seed
best = [max(pool, key=lambda o: lig(o, rc(o)))]
while len(best) < N:
    cands = [(fidelity(best + [o]), o) for o in pool if o not in best and valid(best + [o])]
    best.append(max(cands)[1])
# local search: try replacing each member with any candidate to raise fidelity
improved = True
while improved:
    improved = False
    cur = fidelity(best)
    for idx in range(N):
        for o in pool:
            if o in best: continue
            trial = best[:idx] + [o] + best[idx+1:]
            if valid(trial) and fidelity(trial) > cur:
                best, cur, improved = trial, fidelity(trial), True
best = sorted(best, key=lambda o: -lig(o, rc(o)))
fid = fidelity(best)

# ---- worst cross-talk pair in final set -------------------------------------
pool_species = [s for o in best for s in ([o] if is_pal(o) else [o, rc(o)])]
worst = max(((lig(i, j), i, j) for i in pool_species for j in pool_species if j != rc(i)),
            default=(0, "", ""))
on_min = min(lig(o, rc(o)) for o in best)

junctions = ["J1_intron|HC", "J2_HC|IRES", "J3_polyA|SV40p", "J4_SV40pA|ori", "J5_AmpR|CMV"]
out = {
  "source": "Pryor et al. 2020 PLoS ONE (S1 Table, BsaI-HFv2 x T4 DNA ligase, 37C); fidelity model from Potapov et al. 2018",
  "calibration_plant_set_fidelity_pct": round(plant_fid*100, 1),
  "overhangs": best,
  "assignment": dict(zip(junctions, best)),
  "predicted_set_fidelity_pct": round(fid*100, 2),
  "min_on_target_count": on_min,
  "worst_offtarget": {"count": worst[0], "pair": f"{worst[1]}/{worst[2]}"},
  "pairwise": [{"pair": f"{a}/{b}", "hamming": hamming(a, b), "a_is_rc_b": a == rc(b)}
               for a, b in itertools.combinations(best, 2)],
  "all_non_palindromic": all(not is_pal(o) for o in best),
  "no_rc_collision": all(rc(a) not in best for a in best),
  "gc_per_overhang": {o: gc(o) for o in best},
}
json.dump(out, open(os.path.join(HERE, "fusion_sites.json"), "w"), indent=2)

print("\nSelected fusion-site overhangs (5'->3', top strand):")
for j, o in out["assignment"].items():
    print(f"  {j:18} {o}  (on-target={lig(o,rc(o))}, GC={gc(o)}/4)")
print(f"\nPredicted set fidelity: {fid*100:.2f}%   min on-target={on_min}   "
      f"worst off-target={worst[0]} ({worst[1]}/{worst[2]})")
print("non-palindromic:", out["all_non_palindromic"], " no RC collision:", out["no_rc_collision"],
      " min pairwise Hamming:", min(hamming(a, b) for a, b in itertools.combinations(best, 2)))
