#!/usr/bin/env python3
"""
06_design_primers.py -- Genotyping & Sanger primers for the TP53-mCherry knock-in.

Three primer sets on the knock-in allele template (ki_context = 500bp genomic | 5'arm 800 |
insert 753 | 3'arm 800 | 500bp genomic):
  1. 5' junction  : F OUTSIDE 5' arm (genomic upstream) + R in insert  -> proves 5' integration
  2. 3' junction  : F in insert + R OUTSIDE 3' arm (genomic downstream) -> proves 3' integration
  3. internal/across-insert (also WT-vs-KI zygosity): F & R in the arms flanking the insert ->
     WT band (small) vs KI band (WT+753); spans the whole insert for Sanger confirmation.

Junction primers anchored OUTSIDE the arms cannot prime on the donor plasmid (only ~800bp arms),
so a band appears only on correctly integrated genomic DNA -- not from episomal/random donor.

Specificity: isPcr against local hg38.2bit for every genomic-binding primer.
Outputs: scripts/primers.json
"""
import json, os, subprocess, primer3
from Bio.Seq import Seq

HERE = os.path.dirname(os.path.abspath(__file__))
TWOBIT = os.path.expanduser("~/genomes/hg38.2bit")
ISPCR = os.path.expanduser("~/genomes/isPcr")
donor = json.load(open(os.path.join(HERE, "donor.json")))
ki = donor["ki_context"]; m = donor["ki_map"]
j5, j3 = m["j5_pos"], m["j3_pos"]            # junction coords in ki_context
LC = m["left_ctx_len"]; arm5 = m["arm5_len"]; arm3 = m["arm3_len"]
insert_len = m["insert_len"]

GLOBAL = dict(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_MIN_GC=40.0, PRIMER_MAX_GC=60.0, PRIMER_MAX_POLY_X=4,
              PRIMER_SALT_MONOVALENT=50.0, PRIMER_NUM_RETURN=5)

def design(name, ok_region, size_range):
    seq_args = dict(SEQUENCE_ID=name, SEQUENCE_TEMPLATE=ki,
                    SEQUENCE_PRIMER_PAIR_OK_REGION_LIST=[ok_region])
    g = dict(GLOBAL); g["PRIMER_PRODUCT_SIZE_RANGE"] = [size_range]
    res = primer3.bindings.design_primers(seq_args, g)
    if res.get("PRIMER_PAIR_NUM_RETURNED", 0) == 0:
        raise RuntimeError(f"{name}: primer3 returned 0 pairs ({res.get('PRIMER_ERROR','')}) "
                           f"left={res.get('PRIMER_LEFT_EXPLAIN')} right={res.get('PRIMER_RIGHT_EXPLAIN')}")
    f_seq = res["PRIMER_LEFT_0_SEQUENCE"]; r_seq = res["PRIMER_RIGHT_0_SEQUENCE"]
    f_pos = res["PRIMER_LEFT_0"]; r_pos = res["PRIMER_RIGHT_0"]    # (start,len); right start is 3'->5' top coord
    return dict(name=name, fwd=f_seq, rev=r_seq,
                fwd_tm=round(res["PRIMER_LEFT_0_TM"],1), rev_tm=round(res["PRIMER_RIGHT_0_TM"],1),
                fwd_gc=round(res["PRIMER_LEFT_0_GC_PERCENT"],1), rev_gc=round(res["PRIMER_RIGHT_0_GC_PERCENT"],1),
                fwd_start=f_pos[0], fwd_len=f_pos[1], rev_end=r_pos[0], rev_len=r_pos[1],
                product_size_ki=res["PRIMER_PAIR_0_PRODUCT_SIZE"])

# 1. 5' junction: F in left_ctx (outside 5' arm), R in insert near its 5' end
p5 = design("TP53mCh_5pJxn", [120, LC-130, j5, 480], [700, 1400])
# 2. 3' junction: F in insert near 3' end, R in right_ctx (outside 3' arm)
p3 = design("TP53mCh_3pJxn", [j3-500, 480, LC+arm5+insert_len+arm3, 480], [700, 1400])
# 3. internal/across-insert + WT/KI zygosity: F in 5'arm, R in 3'arm flanking insert
pint = design("TP53mCh_insertSeq", [j5-260, 210, j3+40, 260], [800, 1350])
pint["product_size_wt"] = pint["product_size_ki"] - insert_len

# ---- annotate roles / where each primer binds ----
def zone(pos):
    if pos < LC: return "genomic upstream (OUTSIDE 5' arm)"
    if pos < LC+arm5: return "5' homology arm"
    if pos < LC+arm5+insert_len: return "insert (linker/mCherry)"
    if pos < LC+arm5+insert_len+arm3: return "3' homology arm"
    return "genomic downstream (OUTSIDE 3' arm)"
for p in (p5, p3, pint):
    p["fwd_zone"] = zone(p["fwd_start"]); p["rev_zone"] = zone(p["rev_end"])

# sanity: junction anchors are truly outside the arms
assert p5["fwd_start"]+p5["fwd_len"] <= LC, "5' junction F not fully outside 5' arm"
assert p3["rev_end"] >= LC+arm5+insert_len+arm3, "3' junction R not fully outside 3' arm"
print("[5' junction]", p5["fwd"], "/", p5["rev"], "KI product", p5["product_size_ki"], "bp")
print("   F:", p5["fwd_zone"], "| R:", p5["rev_zone"])
print("[3' junction]", p3["fwd"], "/", p3["rev"], "KI product", p3["product_size_ki"], "bp")
print("   F:", p3["fwd_zone"], "| R:", p3["rev_zone"])
print("[internal/zygosity]", pint["fwd"], "/", pint["rev"],
      "-> WT", pint["product_size_wt"], "bp / KI", pint["product_size_ki"], "bp")

# ---- isPcr specificity for all GENOMIC-binding primers ----
# insert primers (mCherry) are absent from hg38 by design, so junction PCRs cannot prime on
# unmodified genome. We validate the 4 genomic primers via genomic-only pairings.
gp = {"out5F": p5["fwd"], "out3R": p3["rev"], "intF": pint["fwd"], "intR": pint["rev"]}
ispcr_in = os.path.join(HERE, "_ispcr_primers.txt")
with open(ispcr_in, "w") as fh:
    fh.write(f"TP53_internal_WT\t{gp['intF']}\t{gp['intR']}\n")          # expect ~WT band at TP53
    fh.write(f"TP53_wholespan_WT\t{gp['out5F']}\t{gp['out3R']}\n")        # both outside primers
    fh.write(f"TP53_5p_genomic\t{gp['out5F']}\t{gp['intR']}\n")          # 5' outside F specificity
    fh.write(f"TP53_3p_genomic\t{gp['intF']}\t{gp['out3R']}\n")          # 3' outside R specificity
ispcr_out = os.path.join(HERE, "_ispcr_out.fa")
subprocess.run([ISPCR, TWOBIT, ispcr_in, ispcr_out, "-out=fa", "-maxSize=5000"],
               check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
products = {}
for line in open(ispcr_out):
    if line.startswith(">"):
        # >chrN:start-end name ... product size
        h = line[1:].split()
        loc = h[0]; pname = h[1] if len(h) > 1 else "?"
        products.setdefault(pname, []).append(loc)
print("\n[isPcr vs hg38] genomic-primer specificity:")
for pname in ["TP53_internal_WT","TP53_wholespan_WT","TP53_5p_genomic","TP53_3p_genomic"]:
    locs = products.get(pname, [])
    chrom_ok = all(l.startswith("chr17:") for l in locs)
    print(f"   {pname:20} {len(locs)} product(s): {locs}  {'OK (unique chr17)' if len(locs)==1 and chrom_ok else 'REVIEW' }")
os.remove(ispcr_in); os.remove(ispcr_out)

primers = {"sets": {"5_prime_junction": p5, "3_prime_junction": p3, "internal_across_insert": pint},
           "isPcr_products": products,
           "notes": "Junction-PCR insert primers are mCherry (absent from hg38) -> no product on "
                    "unmodified genome; a band appears only on correct integration. Internal pair gives "
                    f"WT={pint['product_size_wt']}bp vs KI={pint['product_size_ki']}bp for zygosity."}
json.dump(primers, open(os.path.join(HERE, "primers.json"), "w"), indent=2)
print("\n[ok] primers.json written")
