dee/core/offtarget.py

"""Genome-wide off-target search for CRISPR guides.

Lazy-downloads a small genome (currently only E. coli K-12 MG1655 is
fully implemented on the free HF Space tier — mammalian genomes need
pre-built indexes hosted externally, Phase 2B-2 work) and scores every
guide against every plausible off-target with the CFD matrix from
Doench 2016.

Architecture:

    1. FASTA download — lazy, on first request that names the organism.
       The 4.6 Mb E. coli genome takes ~3 s from NCBI's E-utils. Cached
       to /tmp/turingdna_genomes/ for the container's lifetime. Cold-
       start re-download cost is ~3 s.

    2. Kmer index build — extract every 23-mer matching {spacer}{NGG}
       on both strands. For E. coli that's ~240k sites. We organize
       them by their 8-nt PAM-proximal seed (positions 13-20 of the
       spacer + the 3-nt PAM) so a guide query only scores a small
       candidate list (typically <200 entries) instead of brute-forcing
       all 240k. Build cost ~5 s on E. coli; in-memory size ~20 MB.

    3. Query — given a guide spacer, find the bucket whose seed matches
       (or differs by ≤1 nt, since 1 seed mismatch is the empirical
       tolerance for cleavage), CFD-score each candidate, return ranked
       hits above CFD threshold.

Thread-safe: build is guarded by a single lock so concurrent first
requests don't duplicate the download/build work.

Privacy: the user's guide spacer never leaves the HF Space. Only the
public genome FASTA is fetched (from NCBI, anonymous). The user can
inspect the cached FASTA in /tmp.

This module is intentionally pure-Python with no new dependencies.
Bowtie2 / BWA would be ~10× faster but add 50+ MB of native binaries
to the Docker image. For E. coli (5 s queries) the speed gain isn't
worth the deployment cost. Mammalian Phase 2B-2 may revisit.
"""

from __future__ import annotations

import gzip
import logging
import os
import re
import threading
import time
import urllib.error
import urllib.request
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

logger = logging.getLogger("dee.offtarget")

# Local import — CFD matrix + PAM penalties live in crispr.py to keep
# the Doench 2016 numbers in one place. Lazy import to avoid a
# circular reference (crispr.py imports from this module too via
# find_guides extension).
def _cfd_score(spacer_a: str, spacer_b: str, pam_b: str) -> float:
    from dee.core.crispr import _cfd_pair as _f
    return _f(spacer_a, spacer_b, pam_b)


# ─── Genome sources ────────────────────────────────────────────────
# NCBI E-utils endpoint for retrieving full FASTAs. Robust + fast for
# small organisms; for human/mouse this would be 3 GB+ and impractical
# to download per container start, hence the "Coming soon" status.
GENOME_SOURCES: Dict[str, Dict[str, object]] = {
    "ecoli": {
        "name":      "E. coli K-12 MG1655",
        "accession": "NC_000913.3",
        # E-utils efetch — chosen over the FTP URL because it's reliably
        # CORS-permissive and stable across NCBI's server moves.
        "url":       "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi"
                     "?db=nuccore&id=556503834&rettype=fasta&retmode=text",
        "is_gzip":   False,
        "size_mb":   4.6,
        "scope":     "full genome",
        "ready":     True,
    },
    "human": {
        # CDS-only (exome). Catches the off-targets users actually care
        # about — those inside coding regions. The full GRCh38 (3 GB) is
        # impractical to download + index inside a free HF container; the
        # ~30 MB Ensembl CDS bundle is the pragmatic tradeoff and matches
        # what CRISPR users typically screen against in clinical contexts.
        "name":      "Homo sapiens (GRCh38, CDS only)",
        "accession": "GRCh38",
        "url":       "https://ftp.ensembl.org/pub/release-112/fasta/homo_sapiens/cds/"
                     "Homo_sapiens.GRCh38.cds.all.fa.gz",
        "is_gzip":   True,
        "size_mb":   30,
        "scope":     "exome (CDS only)",
        "ready":     True,
    },
    "mouse": {
        "name":      "Mus musculus (GRCm39, CDS only)",
        "accession": "GRCm39",
        "url":       "https://ftp.ensembl.org/pub/release-112/fasta/mus_musculus/cds/"
                     "Mus_musculus.GRCm39.cds.all.fa.gz",
        "is_gzip":   True,
        "size_mb":   25,
        "scope":     "exome (CDS only)",
        "ready":     True,
    },
}

_GENOME_CACHE_DIR = os.environ.get("TURINGDNA_GENOME_CACHE", "/tmp/turingdna_genomes")
_KMER_CACHE: Dict[str, "KmerIndex"] = {}
_BUILD_LOCK = threading.Lock()
# Per-organism build state — distinct from _BUILD_LOCK because we want
# concurrent requests for DIFFERENT organisms to proceed in parallel,
# and we want to detect "already building" without acquiring the lock.
_BUILDING: Dict[str, bool] = {}
_BUILDING_LOCK = threading.Lock()
# Max seconds to BLOCK a user request waiting for an index. If the
# build takes longer than this, the request returns without genome
# off-target data + a "building" status flag. The build keeps running
# in the original thread that triggered it, so subsequent requests
# eventually find a populated cache.
_INDEX_WAIT_BUDGET_S = 12.0

# Seed = the 8 nt immediately 5' of the PAM (positions 13-20 of the
# spacer). A perfect seed match plus PAM is required for SpCas9 to bind
# stably; allowing 1 seed mismatch covers most empirically-observed
# off-targets while keeping the candidate set small.
_SEED_LEN  = 8
_MAX_SEED_MISMATCHES = 1
_MAX_TOTAL_MISMATCHES = 4
_CFD_KEEP_THRESHOLD   = 0.05   # off-targets below this aren't worth showing


@dataclass
class OffTargetHit:
    chrom:           str    # FASTA contig / chromosome identifier
    position_1:      int    # 1-based start position on the chrom
    strand:          str    # '+' or '-'
    target_spacer:   str    # 20-nt target spacer (genome-side)
    target_pam:      str    # 3-nt PAM (genome-side)
    cfd:             float  # CFD score, [0, 1]
    n_mismatches:    int    # total mismatches in the 20-nt spacer


@dataclass
class KmerIndex:
    """In-memory off-target index for one organism.

    Structure: seed (last SEED_LEN nt of spacer) → list of full kmer
    matches and their genomic locations. The seed is the most
    selective region for Cas9 binding, so seed-bucketing prunes the
    search space dramatically.
    """
    organism:  str
    n_chroms:  int
    n_sites:   int
    # seed → list of (full_spacer, pam, chrom, position_1, strand)
    by_seed:   Dict[str, List[Tuple[str, str, str, int, str]]]


# ─── Public API ────────────────────────────────────────────────────

def is_organism_ready(organism: str) -> bool:
    """True if the organism has a real index pipeline (vs. UI placeholder)."""
    cfg = GENOME_SOURCES.get(organism)
    return bool(cfg and cfg.get("ready"))


def index_status(organism: str) -> str:
    """Returns 'ready' if the kmer index is cached in memory,
    'building' if a build is currently in flight (or just kicked off),
    'unavailable' if the organism isn't recognized,
    'n/a' if no organism was requested."""
    if not organism:
        return "n/a"
    if not is_organism_ready(organism):
        return "unavailable"
    if organism in _KMER_CACHE:
        return "ready"
    return "building"


def kick_off_build(organism: str) -> None:
    """Start a background build for `organism` if one isn't already in
    progress. Returns immediately. Idempotent. Used by prewarm hooks +
    by find_genomic_offtargets when it times out — the next request
    benefits from the build that this one started."""
    if not is_organism_ready(organism) or organism in _KMER_CACHE:
        return
    with _BUILDING_LOCK:
        if _BUILDING.get(organism):
            return
        _BUILDING[organism] = True

    def _worker():
        try:
            _get_kmer_index(organism)  # actual download + build
        finally:
            with _BUILDING_LOCK:
                _BUILDING[organism] = False

    t = threading.Thread(target=_worker, name=f"kmer-build-{organism}", daemon=True)
    t.start()


def find_genomic_offtargets(
    guide_spacer: str,
    organism: str,
    max_results: int = 20,
) -> List[OffTargetHit]:
    """Score the given guide against every plausible off-target in the
    organism's genome. Returns hits ranked by CFD descending, capped at
    max_results. Returns an EMPTY LIST in these cases (caller should
    consult index_status() to distinguish them):
      - guide isn't 20 nt (Cas12a, malformed input)
      - organism isn't supported / placeholder
      - kmer index isn't cached AND a build is in progress — we won't
        block the HTTP request waiting for it. Instead we kick off a
        background build (idempotent if already running) and return.
        The next request once the build is done will find the cache."""
    if len(guide_spacer) != 20:
        return []
    if not is_organism_ready(organism):
        return []

    # Fast path: already cached. Use it.
    index = _KMER_CACHE.get(organism)
    if index is None:
        # Cold path: kick off the background build (no-op if already
        # running) and return empty. Frontend renders a "still
        # building, refresh in 2 min" banner via the API status field.
        kick_off_build(organism)
        return []

    guide_seed = guide_spacer[-_SEED_LEN:]
    # Search the perfect-seed bucket plus all 1-mismatch seed buckets.
    # 1 mismatch × 4 bases × 8 positions = 24 alternative seeds.
    candidate_seeds: List[str] = [guide_seed]
    for i in range(_SEED_LEN):
        for b in "ACGT":
            if b == guide_seed[i]:
                continue
            candidate_seeds.append(guide_seed[:i] + b + guide_seed[i + 1:])

    hits: List[OffTargetHit] = []
    seen: set = set()  # dedupe identical (chrom, pos, strand) hits
    for seed in candidate_seeds:
        bucket = index.by_seed.get(seed)
        if not bucket:
            continue
        for (target_spacer, target_pam, chrom, pos, strand) in bucket:
            key = (chrom, pos, strand)
            if key in seen:
                continue
            # Quick total-mismatch filter before CFD computation.
            n_mm = sum(1 for i in range(20) if guide_spacer[i] != target_spacer[i])
            if n_mm > _MAX_TOTAL_MISMATCHES:
                continue
            # CFD score (matrix lives in crispr.py).
            cfd = _cfd_score(guide_spacer, target_spacer, target_pam)
            if cfd < _CFD_KEEP_THRESHOLD:
                continue
            seen.add(key)
            hits.append(OffTargetHit(
                chrom=chrom,
                position_1=pos,
                strand=strand,
                target_spacer=target_spacer,
                target_pam=target_pam,
                cfd=cfd,
                n_mismatches=n_mm,
            ))

    hits.sort(key=lambda h: -h.cfd)
    return hits[:max_results]


# ─── Lazy index construction ───────────────────────────────────────

def _get_kmer_index(organism: str) -> Optional[KmerIndex]:
    """Return the cached kmer index, building it if this is the first
    request after a cold start. Thread-safe."""
    cached = _KMER_CACHE.get(organism)
    if cached is not None:
        return cached
    with _BUILD_LOCK:
        cached = _KMER_CACHE.get(organism)  # double-check after lock
        if cached is not None:
            return cached
        try:
            t0 = time.time()
            fasta_path = _ensure_genome_downloaded(organism)
            index = _build_kmer_index(organism, fasta_path)
            elapsed = time.time() - t0
            logger.info(
                "Built kmer index for %s: %d sites across %d chroms in %.1f s",
                organism, index.n_sites, index.n_chroms, elapsed,
            )
            _KMER_CACHE[organism] = index
            return index
        except Exception as exc:  # noqa: BLE001
            logger.exception("Failed to build kmer index for %s: %s", organism, exc)
            return None


def _ensure_genome_downloaded(organism: str) -> str:
    """Download the genome FASTA if not already cached on disk. Returns
    the path to the cached (uncompressed) FASTA. Handles both plain-text
    and gzip-compressed sources."""
    cfg = GENOME_SOURCES[organism]
    if not cfg.get("url"):
        raise RuntimeError(f"No download URL configured for {organism}")
    os.makedirs(_GENOME_CACHE_DIR, exist_ok=True)
    path = os.path.join(_GENOME_CACHE_DIR, f"{organism}.fasta")
    if os.path.exists(path) and os.path.getsize(path) > 1000:
        return path  # already cached, looks healthy

    # For larger organisms, log a clear "this will take a minute" message
    # so the server logs make first-request latency obvious.
    size_mb = cfg.get("size_mb", 0)
    logger.info(
        "Downloading %s genome (%s, ~%.0f MB) from %s",
        organism, cfg.get("scope", "?"), size_mb, cfg["url"],
    )
    req = urllib.request.Request(
        str(cfg["url"]),
        headers={"User-Agent": "TuringDNA/1.0 (https://turingdna.com)"},
    )
    # Longer timeout for the larger CDS bundles — Ensembl FTP is usually
    # fast but the 30 MB human file occasionally takes 30-45 s.
    with urllib.request.urlopen(req, timeout=300) as resp:
        data = resp.read()
    if cfg.get("is_gzip"):
        # Decompress before writing so the kmer-index reader doesn't need
        # to special-case gzip. ~30 MB → ~80 MB uncompressed for human;
        # well within /tmp's free-tier 50 GB.
        logger.info("Decompressing %s genome (%.1f MB compressed)",
                    organism, len(data) / 1e6)
        data = gzip.decompress(data)
    with open(path, "wb") as f:
        f.write(data)
    logger.info("Cached %s genome at %s (%.1f MB)",
                organism, path, os.path.getsize(path) / 1e6)
    return path


_FASTA_HEADER_RE = re.compile(r"^>(\S+)")


def _read_fasta(path: str) -> List[Tuple[str, str]]:
    """Returns list of (chrom_id, sequence). All sequences uppercase,
    non-ACGT chars dropped (Ns + IUPAC ambiguity codes can't be used
    as off-target candidates anyway)."""
    chroms: List[Tuple[str, str]] = []
    cur_id: Optional[str] = None
    cur_parts: List[str] = []
    with open(path, "r") as f:
        for line in f:
            line = line.rstrip("\n")
            if line.startswith(">"):
                if cur_id is not None:
                    chroms.append((cur_id, "".join(cur_parts).upper()))
                m = _FASTA_HEADER_RE.match(line)
                cur_id = m.group(1) if m else line[1:].strip()
                cur_parts = []
            else:
                cur_parts.append(line.strip())
        if cur_id is not None:
            chroms.append((cur_id, "".join(cur_parts).upper()))
    # Drop ambiguity codes — keep only canonical bases.
    cleaned: List[Tuple[str, str]] = []
    for cid, seq in chroms:
        cleaned.append((cid, re.sub(r"[^ACGT]", "N", seq)))
    return cleaned


_NGG = re.compile(r"(?=([ACGT]GG))")


def _revcomp(seq: str) -> str:
    return seq.translate(str.maketrans("ACGTacgt", "TGCAtgca"))[::-1]


def _build_kmer_index(organism: str, fasta_path: str) -> KmerIndex:
    """Scan both strands of every chromosome for NGG sites with at
    least 20 nt of upstream context, organize them by 8-nt PAM-proximal
    seed for fast guide lookup."""
    chroms = _read_fasta(fasta_path)
    by_seed: Dict[str, List[Tuple[str, str, str, int, str]]] = {}
    n_sites = 0

    for (chrom, seq) in chroms:
        # Forward strand. NGG at position p means 20-nt spacer at [p-20, p).
        for m in _NGG.finditer(seq):
            p = m.start()
            if p < 20:
                continue
            spacer = seq[p - 20:p]
            if "N" in spacer:
                continue
            pam = m.group(1)
            seed = spacer[-_SEED_LEN:]
            by_seed.setdefault(seed, []).append((spacer, pam, chrom, p - 20 + 1, "+"))
            n_sites += 1

        # Reverse strand: complement the search to keep coordinates on
        # the forward strand for display. The 20-nt spacer on the - strand
        # corresponds to a forward-strand region; we store its
        # forward-strand position.
        rc = _revcomp(seq)
        n = len(seq)
        for m in _NGG.finditer(rc):
            p = m.start()
            if p < 20:
                continue
            spacer_rc = rc[p - 20:p]
            if "N" in spacer_rc:
                continue
            pam_rc = m.group(1)
            seed = spacer_rc[-_SEED_LEN:]
            # The forward-strand coordinate: the spacer on rc occupies
            # rc[p-20:p], which complements forward[n-p : n-(p-20)] =
            # forward[n-p : n-p+20]. Forward-strand 1-based start =
            # (n - p) + 1.
            fwd_pos = n - p + 1
            by_seed.setdefault(seed, []).append((spacer_rc, pam_rc, chrom, fwd_pos, "-"))
            n_sites += 1

    return KmerIndex(
        organism=organism,
        n_chroms=len(chroms),
        n_sites=n_sites,
        by_seed=by_seed,
    )


# ─── Convenience: prewarm cache ────────────────────────────────────

def prewarm(organism: str = "ecoli") -> bool:
    """Trigger the lazy index build now. Used in tests / startup hooks
    if you want to absorb the first-request latency outside of a user
    request."""
    return _get_kmer_index(organism) is not None