dee/core/base_editor.py

"""Base-editing design — a first-class CRISPR mode (Phase 3, M1).

Where the knockout pipeline asks "where will Cas9 cut and how cleanly?",
base editing asks "which single-base change does this guide install, and
what does it do to the protein?". Most competing tools only *flag* that a
guide is base-editable; here we predict the actual edit(s), warn about
bystanders, score editability, and — when a reading frame is known —
report the codon → amino-acid consequence, including premature stops for
DSB-free knockout (CRISPR-STOP / iSTOP).

Two editor families:
  * CBE (cytosine base editors): deaminate C→T on the protospacer
    (non-target) strand within an editor-specific activity window.
  * ABE (adenine base editors): deaminate A→G in a similar window.

Coordinate convention (matches dee/core/crispr.py and _scan_base_editor):
  Spacer positions are 1-indexed from the 5' / PAM-DISTAL end. For SpCas9
  the 20-nt protospacer sits 5'→[1..20]→3' with the NGG PAM at 21-23. The
  canonical editing window is reported in these coordinates.

Honesty: the windows + per-position activity profiles below are the
published literature consensus, NOT a trained model. They rank-order
editable sites well and match each editor's reported window, but exact
per-base efficiencies vary by sequence context and cell type. The UI says
so.

References:
  * Komor AC et al. (2016). Programmable editing of a target base in
    genomic DNA without double-stranded DNA cleavage. Nature 533, 420.
    [BE3, CBE window 4-8]
  * Koblan LW et al. (2018). Improving cytidine and adenine base editors
    ... Nat Biotechnol 36, 843. [BE4max]
  * Thuronyi BW et al. (2019). Continuous evolution of base editors with
    expanded target compatibility. Nat Biotechnol 37, 1070. [evoCDA — wide]
  * Gaudelli NM et al. (2017). Programmable base editing of A•T to G•C ...
    Nature 551, 464. [ABE7.10, window 4-7]
  * Richter MF et al. (2020). Phage-assisted evolution of an adenine base
    editor with improved Cas domain compatibility and activity. Nat
    Biotechnol 38, 883. [ABE8e — wide, processive]
  * Kuscu C et al. (2017). CRISPR-STOP: gene silencing via base editing
    -induced nonsense mutations. Nat Methods 14, 710.
  * Billon P et al. (2017). CRISPR-mediated base editing enables efficient
    disruption of eukaryotic genes through induction of STOP codons. Mol
    Cell 67, 1068. [iSTOP]
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

# ─── Standard genetic code (DNA codons → 1-letter AA; '*' = stop) ─────
CODON_TABLE: Dict[str, str] = {
    "TTT": "F", "TTC": "F", "TTA": "L", "TTG": "L",
    "CTT": "L", "CTC": "L", "CTA": "L", "CTG": "L",
    "ATT": "I", "ATC": "I", "ATA": "I", "ATG": "M",
    "GTT": "V", "GTC": "V", "GTA": "V", "GTG": "V",
    "TCT": "S", "TCC": "S", "TCA": "S", "TCG": "S",
    "CCT": "P", "CCC": "P", "CCA": "P", "CCG": "P",
    "ACT": "T", "ACC": "T", "ACA": "T", "ACG": "T",
    "GCT": "A", "GCC": "A", "GCA": "A", "GCG": "A",
    "TAT": "Y", "TAC": "Y", "TAA": "*", "TAG": "*",
    "CAT": "H", "CAC": "H", "CAA": "Q", "CAG": "Q",
    "AAT": "N", "AAC": "N", "AAA": "K", "AAG": "K",
    "GAT": "D", "GAC": "D", "GAA": "E", "GAG": "E",
    "TGT": "C", "TGC": "C", "TGA": "*", "TGG": "W",
    "CGT": "R", "CGC": "R", "CGA": "R", "CGG": "R",
    "AGT": "S", "AGC": "S", "AGA": "R", "AGG": "R",
    "GGT": "G", "GGC": "G", "GGA": "G", "GGG": "G",
}


@dataclass
class BaseEditor:
    id:           str
    name:         str
    kind:         str            # "CBE" | "ABE"
    target_base:  str            # base edited ON THE PROTOSPACER strand
    result_base:  str            # what it becomes
    window:       Tuple[int, int]  # 1-indexed inclusive spacer positions
    # Per-position relative deamination activity in [0,1]. Positions
    # outside the window are absent (treated as 0 — no editing).
    activity:     Dict[int, float]
    citation:     str
    note:         str = ""


def _profile(window: Tuple[int, int], peak: int, peak_val: float = 1.0,
             edge_val: float = 0.35) -> Dict[int, float]:
    """Triangular activity profile across `window`, maximal at `peak`."""
    lo, hi = window
    prof: Dict[int, float] = {}
    span = max(peak - lo, hi - peak, 1)
    for p in range(lo, hi + 1):
        frac = abs(p - peak) / span
        prof[p] = round(peak_val - (peak_val - edge_val) * frac, 3)
    return prof


# ─── Curated editor library ──────────────────────────────────────────
# Five editors covering the cases real users reach for: a standard +
# an optimized + a wide-window CBE, and a standard + wide ABE.
BASE_EDITORS: Dict[str, BaseEditor] = {
    "be4max": BaseEditor(
        id="be4max", name="BE4max (CBE)", kind="CBE",
        target_base="C", result_base="T", window=(4, 8),
        activity=_profile((4, 8), peak=6),
        citation="Koblan 2018", note="Optimized cytosine base editor — the common default.",
    ),
    "be3": BaseEditor(
        id="be3", name="BE3 (CBE)", kind="CBE",
        target_base="C", result_base="T", window=(4, 8),
        activity=_profile((4, 8), peak=6),
        citation="Komor 2016", note="Original cytosine base editor.",
    ),
    "evocda_be4max": BaseEditor(
        id="evocda_be4max", name="evoCDA1-BE4max (wide CBE)", kind="CBE",
        target_base="C", result_base="T", window=(2, 10),
        activity=_profile((2, 10), peak=6, edge_val=0.45),
        citation="Thuronyi 2019",
        note="Wide-window CBE — more reach, more bystander risk.",
    ),
    "abe8e": BaseEditor(
        id="abe8e", name="ABE8e (ABE)", kind="ABE",
        target_base="A", result_base="G", window=(3, 9),
        activity=_profile((3, 9), peak=6, edge_val=0.45),
        citation="Richter 2020",
        note="Highly processive adenine base editor — wide window.",
    ),
    "abe7.10": BaseEditor(
        id="abe7.10", name="ABE7.10 (ABE)", kind="ABE",
        target_base="A", result_base="G", window=(4, 7),
        activity=_profile((4, 7), peak=6),
        citation="Gaudelli 2017", note="Standard adenine base editor — narrow window.",
    ),
}

DEFAULT_BASE_EDITOR = "be4max"


def list_base_editors(kind: str = "") -> List[BaseEditor]:
    out = list(BASE_EDITORS.values())
    if kind:
        out = [e for e in out if e.kind.upper() == kind.upper()]
    return out


def get_base_editor(editor_id: str) -> Optional[BaseEditor]:
    return BASE_EDITORS.get((editor_id or "").lower())


# ─── Per-guide edit prediction (sequence-only) ───────────────────────

@dataclass
class BaseEditPrediction:
    editor_id:     str
    # Each edit: (spacer_position_1indexed, from_base, to_base, activity)
    edits:         List[Tuple[int, str, str, float]] = field(default_factory=list)
    editability:   float = 0.0      # strongest in-window target activity [0,1]
    has_bystander: bool = False     # >1 target base in the window
    n_targets:     int = 0
    summary:       str = ""         # short label for the table


def predict_base_edits(spacer: str, editor_id: str) -> BaseEditPrediction:
    """Which bases this editor would change in the guide's window.

    Sequence-only: needs no genome context. Reports every target base in
    the activity window, the strongest as the 'editability', and whether
    bystander edits are present.
    """
    ed = get_base_editor(editor_id)
    if ed is None or not spacer:
        return BaseEditPrediction(editor_id=editor_id or "")
    spacer = spacer.upper()
    lo, hi = ed.window
    edits: List[Tuple[int, str, str, float]] = []
    for pos in range(lo, hi + 1):
        idx = pos - 1
        if 0 <= idx < len(spacer) and spacer[idx] == ed.target_base:
            edits.append((pos, ed.target_base, ed.result_base,
                          ed.activity.get(pos, 0.0)))
    editability = max((a for _, _, _, a in edits), default=0.0)
    n = len(edits)
    if n == 0:
        summary = f"no {ed.target_base} in window"
    elif n == 1:
        p, fb, tb, _ = edits[0]
        summary = f"{fb}{p}→{tb}"
    else:
        summary = ", ".join(f"{fb}{p}→{tb}" for p, fb, tb, _ in edits) + " (bystander)"
    return BaseEditPrediction(
        editor_id=ed.id, edits=edits, editability=round(editability, 3),
        has_bystander=n > 1, n_targets=n, summary=summary,
    )


# ─── Codon-level consequence (frame-aware) ───────────────────────────

@dataclass
class CodonConsequence:
    cds_index:    int
    codon_before: str
    codon_after:  str
    aa_before:    str
    aa_after:     str
    kind:         str    # "silent" | "missense" | "nonsense" | "stop_loss"
    creates_stop: bool
    label:        str    # e.g. "Q→* (stop gained)" or "A123A (silent)"


def codon_consequence(cds: str, cds_index: int, ref_base: str,
                      alt_base: str) -> Optional[CodonConsequence]:
    """Translate the effect of changing coding-strand base `cds_index`
    from ref_base→alt_base.

    Returns None (no guess) if the index is out of range, the codon is
    incomplete, or the CDS base doesn't match ref_base (alignment slipped).
    `cds` must be the coding strand 5'→3'; cds_index 0-based.
    """
    if cds_index < 0 or cds_index >= len(cds):
        return None
    if cds[cds_index] != ref_base:
        return None  # alignment mismatch — refuse to fabricate a consequence
    codon_start = (cds_index // 3) * 3
    if codon_start + 3 > len(cds):
        return None
    within = cds_index - codon_start
    codon_before = cds[codon_start:codon_start + 3]
    codon_after = codon_before[:within] + alt_base + codon_before[within + 1:]
    aa_before = CODON_TABLE.get(codon_before, "?")
    aa_after = CODON_TABLE.get(codon_after, "?")
    creates_stop = aa_after == "*" and aa_before != "*"
    if aa_before == aa_after:
        kind = "silent"
    elif creates_stop:
        kind = "nonsense"
    elif aa_before == "*" and aa_after != "*":
        kind = "stop_loss"
    else:
        kind = "missense"
    aa_num = codon_start // 3 + 1  # 1-based residue number
    if kind == "silent":
        label = f"{aa_before}{aa_num}= (silent)"
    elif kind == "nonsense":
        label = f"{aa_before}{aa_num}* (stop gained)"
    elif kind == "stop_loss":
        label = f"*{aa_num}{aa_after} (stop lost)"
    else:
        label = f"{aa_before}{aa_num}{aa_after}"
    return CodonConsequence(
        cds_index=cds_index, codon_before=codon_before, codon_after=codon_after,
        aa_before=aa_before, aa_after=aa_after, kind=kind,
        creates_stop=creates_stop, label=label,
    )