#!/usr/bin/env python3
"""
06_design_primers.py  ->  primers.json

Design one junction-spanning colony-PCR primer pair per fragment seam (5 pairs).
Each pair flanks a Golden Gate junction in the assembled plasmid, so a band of
the expected size forms ONLY if the two adjacent fragments are correctly joined
(verifies the full assembly with 5 reactions). Primers are Tm-matched within a
pair, carry a 3' G/C clamp, and amplicons span the seam.
"""
import json, os
import primer3
from Bio.Seq import Seq

HERE = os.path.dirname(os.path.abspath(__file__))
assembled = open(os.path.join(HERE, "assembled.fasta")).read().split("\n", 1)[1].replace("\n", "")
frag = json.load(open(os.path.join(HERE, "fragments.json")))
L = len(assembled)

# overhang (junction scar) start positions, 0-based, in plasmid order
ORDER = frag["fragment_order"]
oh = frag["overhangs"]
# junctions sit at the RIGHT overhang of each fragment; recompute their coordinates
pos, junc = 0, {}
segs = [("F1_promoter", frag["fragments"]["F1_promoter"]["inner_len"]),
        ("F2_heavy",     frag["fragments"]["F2_heavy"]["inner_len"]),
        ("F3_IRES_light",frag["fragments"]["F3_IRES_light"]["inner_len"]),
        ("F4_puro",      frag["fragments"]["F4_puro"]["inner_len"]),
        ("F5_backbone",  frag["fragments"]["F5_backbone"]["inner_len"])]
JNAME = {"F1_promoter":"J1_intron|HC","F2_heavy":"J2_HC|IRES","F3_IRES_light":"J3_polyA|SV40p",
         "F4_puro":"J4_SV40pA|ori","F5_backbone":"J5_AmpR|CMV"}
for name, ilen in segs:
    pos += ilen
    junc[JNAME[name]] = pos          # overhang starts right after this fragment's inner body
    pos += 4                          # 4-nt overhang scar

def circ(center, half):
    return "".join(assembled[(center - half + i) % L] for i in range(2 * half))

GLOBAL = {
    "PRIMER_OPT_SIZE": 20, "PRIMER_MIN_SIZE": 18, "PRIMER_MAX_SIZE": 25,
    "PRIMER_OPT_TM": 60.0, "PRIMER_MIN_TM": 57.0, "PRIMER_MAX_TM": 63.0, "PRIMER_PAIR_MAX_DIFF_TM": 2.0,
    "PRIMER_MIN_GC": 40.0, "PRIMER_MAX_GC": 65.0, "PRIMER_GC_CLAMP": 1,
    "PRIMER_MAX_POLY_X": 4, "PRIMER_PRODUCT_SIZE_RANGE": [[180, 450]], "PRIMER_NUM_RETURN": 5,
}
HALF = 260
design = primer3.bindings.design_primers if hasattr(primer3.bindings, "design_primers") else primer3.bindings.designPrimers

results = {}
for jname, jpos in junc.items():
    win = circ(jpos + 2, HALF)                     # center on middle of 4-nt overhang
    seq_args = {"SEQUENCE_ID": jname, "SEQUENCE_TEMPLATE": win,
                "SEQUENCE_TARGET": [HALF - 10, 20]}   # force pair to flank the seam
    r = design(seq_args, GLOBAL)
    n = r.get("PRIMER_PAIR_NUM_RETURNED", 0)
    if not n:
        results[jname] = {"error": "no primer pair", "explain": r.get("PRIMER_PAIR_EXPLAIN")}
        print(f"{jname}: NO PAIR ({r.get('PRIMER_PAIR_EXPLAIN')})"); continue
    fl, fs = r["PRIMER_LEFT_0"]; rl, rs = r["PRIMER_RIGHT_0"]
    fseq, rseq = r["PRIMER_LEFT_0_SEQUENCE"], r["PRIMER_RIGHT_0_SEQUENCE"]
    # map window-local left start to plasmid coordinate
    f_plasmid = (jpos + 2 - HALF + fl) % L
    results[jname] = {
        "junction_pos_in_plasmid": jpos + 1,
        "forward": {"seq": fseq, "len": len(fseq), "tm": round(r["PRIMER_LEFT_0_TM"], 2),
                    "gc": round(r["PRIMER_LEFT_0_GC_PERCENT"], 1), "plasmid_start": f_plasmid + 1},
        "reverse": {"seq": rseq, "len": len(rseq), "tm": round(r["PRIMER_RIGHT_0_TM"], 2),
                    "gc": round(r["PRIMER_RIGHT_0_GC_PERCENT"], 1)},
        "product_size": r["PRIMER_PAIR_0_PRODUCT_SIZE"],
        "dTm": round(abs(r["PRIMER_LEFT_0_TM"] - r["PRIMER_RIGHT_0_TM"]), 2),
        "spans_seam": fl < HALF - 10 and (fl + (rl - fl)) > HALF + 10,
        "f_3prime": fseq[-1], "r_3prime": rseq[-1],
    }
    d = results[jname]
    print(f"{jname:18} F {fseq:25} Tm{d['forward']['tm']:.1f} | R {rseq:23} Tm{d['reverse']['tm']:.1f} "
          f"| amp={d['product_size']}bp dTm={d['dTm']} 3'={d['f_3prime']}/{d['r_3prime']} span={d['spans_seam']}")

json.dump(results, open(os.path.join(HERE, "primers.json"), "w"), indent=2)
allok = all("error" not in v for v in results.values())
print("\nAll 5 junction pairs designed:", allok, "-> wrote primers.json")
