tests/test_crispr.py

"""Regression net for the LIVE CRISPR guide-design engine (dee/core/crispr.py).

This engine is used by real paying customers, so these tests lock down the
properties that matter most:

  * guide coordinates round-trip on BOTH strands (a wrong position →
    customer orders a guide against the wrong locus),
  * CFD off-target scores match the published Doench 2016 matrix,
  * on-target / KO / indel / base-editor outputs stay in sane ranges,
  * Cas12a PAM rules are enforced,
  * malformed / hostile input fails safe (no crash, no garbage guides).

Pure stdlib + the module under test. No network. No I/O.
"""
import time

import pytest

from dee.core import crispr as C


def _rc(s: str) -> str:
    return s.translate(str.maketrans("ACGT", "TGCA"))[::-1]


# A sequence engineered to carry guides on both strands for both enzymes.
SEQ_BOTH = (
    "ATGGCCTACGACGCATCTTCCCGACAACGTGGACCGGTACCGGAGGTTTAA"
    "CCGGTACCGGATCGATCGGGCCCAGGTACCGGATTACCGGTACCGGATCC"
    "GGGTACCGGATCGGGCCCAGGTTTAACCGGTACCGGATCGATCGGGAA"
)


# ─────────────────────────── helpers ───────────────────────────

def test_revcomp_basic():
    assert C._revcomp("ACGT") == "ACGT"
    assert C._revcomp("AAAA") == "TTTT"
    assert C._revcomp("ATGC") == "GCAT"


def test_normalize_strips_fasta_whitespace_and_junk():
    out = C._normalize_input(">my header\nacgt acgt\nNNNN\n", min_len=4)
    assert out == "ACGTACGTNNNN"  # N kept (valid base char), case-folded


def test_normalize_rejects_too_short():
    with pytest.raises(ValueError):
        C._normalize_input("ACGT", min_len=23)


def test_gc_pct():
    assert C._gc_pct("GGCC") == 100.0
    assert C._gc_pct("ATAT") == 0.0
    assert C._gc_pct("ATGC") == 50.0
    assert C._gc_pct("") == 0.0


# ─────────────────────────── CFD matrix ───────────────────────────
# These are the off-target numbers customers rely on. Pin them.

def test_cfd_perfect_match_with_canonical_pam_is_one():
    # A perfect spacer:spacer match under the optimal TGG PAM = 1.0.
    assert C._cfd_pair("A" * 20, "A" * 20, "TGG") == 1.0


def test_cfd_perfect_match_applies_pam_multiplier():
    # Same perfect match, weaker NGG variant: AGG multiplier = 0.259.
    assert C._cfd_pair("A" * 20, "A" * 20, "AGG") == pytest.approx(0.259)


def test_cfd_pam_proximal_mismatch_zeros_score():
    # Mismatch at position 20 (PAM-proximal) with rA:dC is fully
    # disallowed in Doench 2016 (penalty 0.0) → whole score 0.
    b = "A" * 19 + "C"
    assert C._cfd_pair("A" * 20, b, "TGG") == 0.0


def test_cfd_distal_mismatch_is_tolerated():
    # Mismatch at position 1 (5' / distal) rA:dC retains 0.857.
    b = "C" + "A" * 19
    assert C._cfd_pair("A" * 20, b, "TGG") == pytest.approx(0.857)


def test_cfd_non_ngg_pam_is_zero():
    assert C._cfd_pair("A" * 20, "A" * 20, "AAA") == 0.0


def test_cfd_wrong_length_is_zero():
    assert C._cfd_pair("A" * 19, "A" * 19, "TGG") == 0.0


# ─────────────────────────── on-target ───────────────────────────

def test_pam_score_ordering():
    # Doench: TGG > CGG > AGG > GGG.
    assert (C._spcas9_pam_score("TGG") > C._spcas9_pam_score("CGG")
            > C._spcas9_pam_score("AGG") > C._spcas9_pam_score("GGG"))
    assert C._spcas9_pam_score("TAA") == 0.0  # not NGG


def test_gc_score_shape():
    assert C._gc_score(50.0) == 1.0
    assert C._gc_score(10.0) == 0.0
    assert C._gc_score(95.0) == 0.0


@pytest.mark.parametrize("spacer,pam", [
    ("ACGTACGTACGTACGTACGT", "TGG"),
    ("GGGGGGGGGGGGGGGGGGGG", "AGG"),
    ("ATATATATATATATATATAT", "CGG"),
])
def test_on_target_score_bounded(spacer, pam):
    s = C._score_spcas9_on_target(spacer, pam)
    assert 0.0 <= s <= 1.0


# ─────────────────────────── KO efficacy ───────────────────────────

def test_ko_efficacy_bounds_and_monotonicity():
    early, _ = C._score_ko_efficacy(10, 1000)
    late, _ = C._score_ko_efficacy(950, 1000)
    assert 0.0 <= late <= early <= 1.0
    # base frameshift fraction (Allen 2018) at the 5' end ≈ 0.67.
    assert C._score_ko_efficacy(0, 1000)[0] == pytest.approx(0.67, abs=0.01)


def test_ko_reasoning_has_no_semicolon():
    # Georgian/German Excel locales treat ';' as the CSV delimiter.
    for cut in (10, 400, 800, 990):
        _, why = C._score_ko_efficacy(cut, 1000)
        assert ";" not in why


# ─────────────────────────── indel prediction ───────────────────────────

def test_indel_distribution_normalizes_to_one():
    spacer = "ACGTACGTACGTACGTACGT"
    ctx = "GGGG" + spacer + "TGG" + "AAA"
    predicted, label, fs_pct, dom_pct = C._predict_indels(spacer, "TGG", ctx, "cas9")
    assert predicted, "expected at least one predicted indel outcome"
    assert sum(freq for _, freq in predicted) == pytest.approx(1.0, abs=1e-6)
    assert 0.0 <= fs_pct <= 100.0
    assert 0.0 <= dom_pct <= 100.0
    assert "%" in label


def test_indel_cas12a_is_placeholder_not_fabricated():
    predicted, label, fs_pct, _ = C._predict_indels("A" * 23, "TTTA", "A" * 30, "cas12a")
    assert predicted == []
    assert "Cas12a" in label


# ─────────────────────────── base-editor scan ───────────────────────────

def test_base_editor_scan_windows():
    # spacer pos:  1234 5678 ...
    spacer = "AAAC" + "A" * 16  # C at position 4; A at 5,6,7 (and 8)
    cbe, abe, summary = C._scan_base_editor(spacer, "cas9")
    assert cbe == "4"            # CBE window 4-8 → only the C at pos 4
    assert abe == "5,6,7"        # ABE window 4-7 → A at 5,6,7 (pos4 is C)
    assert "CBE" in summary and "ABE" in summary


def test_base_editor_scan_skips_cas12a_and_short():
    assert C._scan_base_editor("A" * 23, "cas12a") == ("", "", "")
    assert C._scan_base_editor("AAA", "cas9") == ("", "", "")


# ─────────────────────────── find_guides integration ───────────────────────────

def _assert_guides_roundtrip(seq, guides):
    """Every guide must reconstruct from (position, strand, spacer)."""
    for g in guides:
        p = g.position - 1                      # 0-based forward footprint
        glen = len(g.spacer)
        footprint = seq[p:p + glen]
        if g.strand == "+":
            assert footprint == g.spacer, f"+ guide {g.rank} pos {g.position}"
        else:
            assert footprint == _rc(g.spacer), f"- guide {g.rank} pos {g.position}"


def test_cas9_guides_roundtrip_both_strands():
    guides = C.find_guides(SEQ_BOTH, enzyme="cas9", max_results=200)
    assert guides
    assert {g.strand for g in guides} == {"+", "-"}, "expected guides on both strands"
    _assert_guides_roundtrip(SEQ_BOTH, guides)


def test_ranks_are_dense_and_sorted_by_composite():
    guides = C.find_guides(SEQ_BOTH, enzyme="cas9", max_results=200)
    assert [g.rank for g in guides] == list(range(1, len(guides) + 1))
    comps = [g.composite_score for g in guides]
    assert comps == sorted(comps, reverse=True)


def test_guide_fields_in_range():
    for g in C.find_guides(SEQ_BOTH, enzyme="cas9", max_results=50):
        assert g.enzyme == "cas9"
        assert 0.0 <= g.composite_score <= 1.0
        assert 0.0 <= g.on_target_score <= 1.0
        assert 0.0 <= g.cfd_max_offtarget <= 1.0
        assert 0.0 <= g.ko_efficacy <= 1.0
        assert 0.0 <= g.gc_pct <= 100.0
        assert len(g.spacer) == 20
        assert g.pam[1:] == "GG"


def test_max_results_and_min_score_respected():
    few = C.find_guides(SEQ_BOTH, enzyme="cas9", max_results=3)
    assert len(few) <= 3
    filtered = C.find_guides(SEQ_BOTH, enzyme="cas9", min_score=0.5, max_results=200)
    assert all(g.composite_score >= 0.5 for g in filtered)


def test_cas12a_pam_rules():
    # TTTA + 23-nt spacer on the forward strand; a TTTT must NOT be used.
    seq = "GGGGGG" + "TTTA" + "ACGTACGTACGTACGTACGTACG" + "GGGGGGTTTTGGGGGG"
    guides = C.find_guides(seq, enzyme="cas12a", max_results=50)
    assert guides
    for g in guides:
        assert g.enzyme == "cas12a"
        assert len(g.spacer) == 23
        assert g.pam.startswith("TTT") and g.pam[3] in "ACG"   # TTTV, never TTTT
    _assert_guides_roundtrip(seq, guides)


def test_cloning_oligos_populated_when_vector_chosen():
    guides = C.find_guides(SEQ_BOTH, enzyme="cas9", max_results=5, vector_id="px459_v2")
    for g in guides:
        assert g.sense_oligo == "CACCG" + g.spacer
        assert g.antisense_oligo == "AAAC" + _rc(g.spacer) + "C"
        assert g.cloning_addgene_id == "62988"


# ─────────────────────────── hostile / edge input ───────────────────────────

def test_html_input_fails_safe():
    with pytest.raises(ValueError):
        C.find_guides("<script>alert('xss')</script>", enzyme="cas9")


def test_protein_input_fails_safe():
    with pytest.raises(ValueError):
        C.find_guides("MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQ", enzyme="cas9")


def test_n_runs_do_not_crash_and_are_excluded():
    seq = "ACGT" * 5 + "N" * 25 + "GGG" + "ACGT" * 10 + "TGG" + "ACGT" * 5
    guides = C.find_guides(seq, enzyme="cas9", max_results=50)
    for g in guides:
        assert "N" not in g.spacer and "N" not in g.pam


def test_moderate_input_completes_quickly():
    # ~2 kb. Locks in that a realistic gene-sized paste returns fast.
    # NOTE (follow-up): self-off-target is O(k^2) in the number of PAM
    # sites; a multi-hundred-kb paste can be slow. Tracked separately.
    import random
    random.seed(7)
    seq = "".join(random.choice("ACGT") for _ in range(2000))
    t0 = time.time()
    guides = C.find_guides(seq, enzyme="cas9", max_results=50)
    assert guides
    assert time.time() - t0 < 10.0


# ─────────────────────────── CSV/XLSX flattening ───────────────────────────

def test_csv_rows_header_and_width_consistent():
    guides = C.find_guides(SEQ_BOTH, enzyme="cas9", max_results=5, vector_id="px459_v2")
    rows = C.guides_to_csv_rows(guides)
    header = rows[0]
    assert rows[1:], "expected at least one data row"
    assert all(len(r) == len(header) for r in rows), "ragged CSV/XLSX rows"
    # No semicolons in any cell (Georgian/German Excel `;`-delimiter safety).
    for r in rows:
        for cell in r:
            assert ";" not in cell