"""Worked example for the ProteinIQ guide "Primer dimers" (/guides/primer-dimers).

Recomputes every primer, dimer and Tm value quoted in the guide.

Requirements: Python 3.10+, primer3-py 2.2.0 (Primer3 2.6.1 thermodynamics), CPU only.
Template: human ACTB mRNA, NCBI RefSeq NM_001101.5 (NM_001101.5.fasta in this folder).

Run:  python primer_dimer_example.py
Writes primer_dimer_example.json next to this script.
"""

import json
import math
from pathlib import Path

import primer3

HERE = Path(__file__).resolve().parent
TEMPLATE = "".join(
    line.strip()
    for line in (HERE / "NM_001101.5.fasta").read_text().splitlines()
    if not line.startswith(">")
)

# Primer3 default PCR conditions (the same defaults as /app/primer3 and /app/oligo-analyzer)
CONDITIONS = dict(mv_conc=50.0, dv_conc=1.5, dntp_conc=0.6, dna_conc=50.0)

# Pair chosen from ACTB to illustrate a 3'-3' overlap: it passes length, GC, Tm and GC-clamp rules
DIMER_PAIR = ("CGCAAGTACTCCGTGTGGATC", "CTCCTGCTTGCTGATCCACA")

# SantaLucia (1998) unified nearest-neighbor dG37 in 1 M NaCl, kcal/mol (Table 1, unified column)
NN_DG37 = {
    "AA/TT": -1.00, "AT/TA": -0.88, "TA/AT": -0.58, "CA/GT": -1.45, "GT/CA": -1.44,
    "CT/GA": -1.28, "GA/CT": -1.30, "CG/GC": -2.17, "GC/CG": -2.24, "GG/CC": -1.84,
}
INIT_GC, INIT_AT, SYMMETRY = 0.98, 1.03, 0.43
COMP = str.maketrans("ACGT", "TGCA")


def revcomp(seq):
    return seq.translate(COMP)[::-1]


def nn_key(step):
    """Return the table key for a 5'->3' dinucleotide step on the top strand."""
    key = f"{step}/{step.translate(COMP)}"
    if key in NN_DG37:
        return key
    # The same stack read from the other strand, e.g. TG/AC == CA/GT
    other = revcomp(step)
    return f"{other}/{other.translate(COMP)}"


def santalucia_dg37(top, na_molar=1.0):
    """dG37 of a perfect duplex of `top` with its complement (no dangling ends)."""
    stacks = [nn_key(top[i:i + 2]) for i in range(len(top) - 1)]
    dg = sum(NN_DG37[k] for k in stacks)
    for end in (top[0], top[-1]):
        dg += INIT_GC if end in "GC" else INIT_AT
    if top == revcomp(top):
        dg += SYMMETRY
    # Salt correction (SantaLucia 1998, eq. 7): N = phosphates in the duplex / 2
    n_half_phosphates = len(top) - 1
    dg_salt = dg - 0.114 * n_half_phosphates * math.log(na_molar)
    return {"stacks": stacks, "dg_1M": round(dg, 2), "dg_salt": round(dg_salt, 2)}


def three_prime_overlap(left, right):
    """Longest k for which the last k bases of `left` pair with the last k bases of `right`."""
    best = 0
    for k in range(1, min(len(left), len(right)) + 1):
        if left[-k:] == revcomp(right[-k:]):
            best = k
    return best


def structure(result):
    return {
        "found": bool(result.structure_found),
        "dg_kcal": round(result.dg / 1000, 2),
        "tm_c": round(result.tm, 1),
        "ascii": result.ascii_structure_lines,
    }


def describe_pair(left, right):
    check = primer3.design_primers(
        {"SEQUENCE_ID": "ACTB_NM_001101.5", "SEQUENCE_TEMPLATE": TEMPLATE,
         "SEQUENCE_PRIMER": left, "SEQUENCE_PRIMER_REVCOMP": right},
        {"PRIMER_TASK": "check_primers", "PRIMER_PICK_ANYWAY": 1,
         "PRIMER_PRODUCT_SIZE_RANGE": [[50, 1000]]},
    )
    keep = ("TM", "GC_PERCENT", "SELF_ANY_TH", "SELF_END_TH", "HAIRPIN_TH", "END_STABILITY")
    out = {
        "left": left, "right": right,
        "left_template_start_1based": check["PRIMER_LEFT_0"][0] + 1,
        "right_template_end_1based": check["PRIMER_RIGHT_0"][0] + 1,
        "product_bp": check["PRIMER_PAIR_0_PRODUCT_SIZE"],
        "primer3_check": {
            **{f"LEFT_{k}": round(check[f"PRIMER_LEFT_0_{k}"], 2) for k in keep},
            **{f"RIGHT_{k}": round(check[f"PRIMER_RIGHT_0_{k}"], 2) for k in keep},
            "PAIR_COMPL_ANY_TH": round(check["PRIMER_PAIR_0_COMPL_ANY_TH"], 2),
            "PAIR_COMPL_END_TH": round(check["PRIMER_PAIR_0_COMPL_END_TH"], 2),
        },
        "three_prime_overlap_bp": three_prime_overlap(left, right),
    }
    k = out["three_prime_overlap_bp"]
    out["extended_dimer_bp"] = len(left) + len(right) - k if k else None
    for temp in (25.0, 37.0, 60.0):
        tag = f"{int(temp)}C"
        out[f"heterodimer_{tag}"] = structure(
            primer3.calc_heterodimer(left, right, temp_c=temp, output_structure=True, **CONDITIONS))
        out[f"left_end_on_right_{tag}"] = structure(
            primer3.calc_end_stability(left, right, temp_c=temp, **CONDITIONS))
        out[f"right_end_on_left_{tag}"] = structure(
            primer3.calc_end_stability(right, left, temp_c=temp, **CONDITIONS))
    # Structure Tm depends on strand concentration; 500 nM is a common starting primer concentration
    out["heterodimer_tm_c_at_500nM"] = round(primer3.calc_heterodimer(
        left, right, **{**CONDITIONS, "dna_conc": 500.0}).tm, 1)
    for name, seq in (("left", left), ("right", right)):
        out[f"{name}_homodimer_37C"] = structure(
            primer3.calc_homodimer(seq, output_structure=True, **CONDITIONS))
        out[f"{name}_hairpin_37C"] = structure(
            primer3.calc_hairpin(seq, output_structure=True, **CONDITIONS))
    return out


def redesign():
    """Ask Primer3 for new pairs covering the same stretch of ACTB, with default limits."""
    design = primer3.design_primers(
        {"SEQUENCE_ID": "ACTB_NM_001101.5", "SEQUENCE_TEMPLATE": TEMPLATE,
         "SEQUENCE_TARGET": [1110, 30]},
        {"PRIMER_PRODUCT_SIZE_RANGE": [[70, 150]], "PRIMER_NUM_RETURN": 5},
    )
    pairs = []
    for i in range(design["PRIMER_PAIR_NUM_RETURNED"]):
        left, right = design[f"PRIMER_LEFT_{i}_SEQUENCE"], design[f"PRIMER_RIGHT_{i}_SEQUENCE"]
        het = primer3.calc_heterodimer(left, right, **CONDITIONS)
        pairs.append({
            "rank": i, "left": left, "right": right,
            "product_bp": design[f"PRIMER_PAIR_{i}_PRODUCT_SIZE"],
            "three_prime_overlap_bp": three_prime_overlap(left, right),
            "heterodimer_dg_kcal_37C": round(het.dg / 1000, 2),
            "pair_compl_end_th": round(design[f"PRIMER_PAIR_{i}_COMPL_END_TH"], 2),
        })
    return pairs


if __name__ == "__main__":
    candidates = redesign()
    # Take the top-ranked Primer3 pair whose 3' ends do not pair with each other
    chosen = next(p for p in candidates if p["three_prime_overlap_bp"] < 3)
    core = DIMER_PAIR[0][-three_prime_overlap(*DIMER_PAIR):]
    na_equiv = (50.0 + 120 * (1.5 - 0.6) ** 0.5) / 1000  # von Ahsen Mg2+ to Na+ equivalent, M
    result = {
        "primer3_py_version": primer3.__version__,
        "template": "NM_001101.5 (Homo sapiens ACTB mRNA)",
        "conditions_mM_nM": CONDITIONS,
        "dimer_pair": describe_pair(*DIMER_PAIR),
        "nearest_neighbor_core": {
            "top_5to3": core,
            **santalucia_dg37(core),
            "na_equivalent_M": round(na_equiv, 4),
            "dg_at_na_equivalent": santalucia_dg37(core, na_equiv)["dg_salt"],
        },
        "redesign_candidates": candidates,
        "redesigned_pair": describe_pair(chosen["left"], chosen["right"]),
    }
    (HERE / "primer_dimer_example.json").write_text(json.dumps(result, indent=2) + "\n")
    print(json.dumps(result, indent=2))
