dee/core/edits.py

"""Applying an edit to a construct — the half of the loop that was missing.

Every other tool in this engine READS. It scores substitutions, folds a
protein, maps a plasmid, ranks guides. None of them change anything, which
means the agent could describe an edit in beautiful detail and never make one:
the scientist still had to open the editor and type it themselves.

This module is the other half. It is deliberately small and deliberately
paranoid, because it is the first thing in the codebase that alters a
construct, and a silent off-by-one here is not a bug report — it is a
scientist ordering the wrong DNA.

Three rules it will not bend:

  1. REFUSE ON MISMATCH. Every edit states what it expects to find. If the
     construct does not have that residue or that base at that position, the
     edit does not happen and the caller is told exactly what is there
     instead. The same rule base_editor.codon_consequence already follows:
     an alignment that has slipped must never be papered over.
  2. NEVER GUESS THE FRAME. Residue numbering only means something if the
     sequence is a coding sequence. If it is not, a residue-level edit is
     refused rather than measured from position 0 and hoped for.
  3. SAY WHICH CODON. "R175H" does not name a codon — His is CAT or CAC, and
     which one you install is a real decision with a real effect on
     expression. The host's preferred codon is used and the choice is
     REPORTED, never left implicit.

Pure: no network, no model, no torch. Everything here is checkable by hand.
"""
from __future__ import annotations

import re
from typing import Any, Dict, Optional

from dee.core.base_editor import CODON_TABLE

# "R175H" — wild-type residue, 1-based position, replacement residue.
_PROTEIN_EDIT = re.compile(r"^([A-Za-z])(\d+)([A-Za-z*])$")
# "A1523G" on DNA reads identically to a protein edit, so the two are told
# apart by the `level` argument rather than by guessing from the letters.
_DNA_EDIT = re.compile(r"^([ACGTacgt])(\d+)([ACGTacgt])$")

_AA_LETTERS = set("ACDEFGHIKLMNPQRSTVWY*")

# Refuse rather than churn: an "edit" this large is a replacement, and it
# should go through the editor where the user can see the whole thing.
MAX_EDITS_PER_CALL = 20


def _clean_dna(seq: str) -> str:
    return re.sub(r"[^ACGTNacgtn]", "", str(seq or "")).upper()


def translate(dna: str) -> str:
    """Codons → one-letter amino acids. Partial trailing codon ignored."""
    dna = _clean_dna(dna)
    return "".join(CODON_TABLE.get(dna[i:i + 3], "X")
                   for i in range(0, len(dna) - len(dna) % 3, 3))


def is_coding(dna: str) -> bool:
    """Is this plausibly a CDS — a whole number of codons, starting ATG and
    ending on a stop, with no internal stop?

    Strict on purpose. A residue number quoted against a sequence that is not
    a reading frame is a number about nothing, and the cost of being wrong is
    an edit at the wrong codon.
    """
    dna = _clean_dna(dna)
    if len(dna) < 6 or len(dna) % 3:
        return False
    if not dna.startswith("ATG"):
        return False
    aa = translate(dna)
    return aa.endswith("*") and "*" not in aa[:-1]


def _preferred_codon(aa: str, host: str) -> str:
    """The host's most-used codon for this residue.

    Reuses the same tables the library encoder uses (dee.core.codon), so a
    codon Turing installs and a codon the Directed Evolution pipeline installs
    are the same codon — not two different opinions about the same host.
    """
    from dee.core import codon as _codon
    table = _codon._resolve_table(host or "e_coli")
    return _codon._best_codon(aa.upper(), table)


def parse_edit(spec: str, level: str = "protein") -> Optional[Dict[str, Any]]:
    """'R175H' → {'from': 'R', 'pos': 175, 'to': 'H'}. None if malformed.

    `level` decides which alphabet applies; it is never inferred, because
    'A123G' is a valid edit in both and guessing would silently edit the
    wrong thing on exactly the sequences where it matters most.
    """
    spec = str(spec or "").strip()
    m = (_PROTEIN_EDIT if level == "protein" else _DNA_EDIT).match(spec)
    if not m:
        return None
    src, pos, dst = m.group(1).upper(), int(m.group(2)), m.group(3).upper()
    if pos < 1:
        return None
    if level == "protein" and (src not in _AA_LETTERS or dst not in _AA_LETTERS):
        return None
    return {"from": src, "pos": pos, "to": dst}


def apply_edits(sequence: str, specs, *, level: str = "protein",
                host: str = "e_coli") -> Dict[str, Any]:
    """Apply one or more edits to a construct.

    Returns {ok, sequence, applied[], length, ...} or {ok: False, error}.
    Nothing is applied unless EVERY edit validates — a half-applied set is a
    construct nobody asked for, and the agent would have no way to tell which
    half it got.
    """
    dna = _clean_dna(sequence)
    if not dna:
        return {"ok": False, "error": "No sequence to edit."}
    if isinstance(specs, str):
        specs = [specs]
    specs = [s for s in (specs or []) if str(s or "").strip()]
    if not specs:
        return {"ok": False, "error": "No edit given."}
    if len(specs) > MAX_EDITS_PER_CALL:
        return {"ok": False, "error": (
            f"{len(specs)} edits in one call — cap is {MAX_EDITS_PER_CALL}. "
            f"A change this large is a replacement, not an edit; build it in "
            f"the editor where the whole construct is visible.")}

    if level not in ("protein", "dna"):
        return {"ok": False, "error": "level must be 'protein' or 'dna'."}

    parsed = []
    for spec in specs:
        p = parse_edit(spec, level)
        if not p:
            return {"ok": False, "error": (
                f"Couldn't read “{spec}” as a{'n amino-acid' if level == 'protein' else ' DNA'} "
                f"edit. Expected e.g. "
                f"{'R175H (wild-type residue, position, replacement)' if level == 'protein' else 'A1523G'}.")}
        parsed.append((spec, p))

    seen = {}
    for spec, p in parsed:
        if p["pos"] in seen:
            return {"ok": False, "error": (
                f"Two edits at position {p['pos']} ({seen[p['pos']]} and {spec}). "
                f"Pick one.")}
        seen[p["pos"]] = spec

    # ── protein-level ────────────────────────────────────────────────────
    if level == "protein":
        if not is_coding(dna):
            return {"ok": False, "kind": "not_coding", "error": (
                "Residue numbering needs a coding sequence — this one is not a "
                "clean reading frame (whole codons, ATG start, single stop at "
                "the end). Give the CDS, or use level='dna' with base "
                "positions.")}
        aa = translate(dna)
        chars = list(dna)
        applied = []
        for spec, p in parsed:
            idx = p["pos"] - 1
            if idx >= len(aa) - 1:      # -1: never edit the stop by residue
                return {"ok": False, "error": (
                    f"{spec}: this protein is {len(aa) - 1} residues long "
                    f"(plus the stop), so there is no residue {p['pos']}.")}
            actual = aa[idx]
            if actual != p["from"]:
                return {"ok": False, "kind": "mismatch", "error": (
                    f"{spec}: residue {p['pos']} is {actual}, not {p['from']}. "
                    f"Not editing — the numbering does not line up, and "
                    f"guessing which one you meant is how the wrong residue "
                    f"gets mutated.")}
            codon_before = dna[idx * 3:idx * 3 + 3]
            codon_after = _preferred_codon(p["to"], host)
            chars[idx * 3:idx * 3 + 3] = list(codon_after)
            applied.append({
                "edit": f"{p['from']}{p['pos']}{p['to']}",
                "level": "protein",
                "residue": p["pos"],
                "codon_before": codon_before,
                "codon_after": codon_after,
                # The whole reason this field exists: His is CAT or CAC and
                # the choice is ours, so it is stated rather than buried.
                "codon_note": (f"{p['to']} installed as {codon_after} — the "
                               f"most-used {p['to']} codon in {host}."),
                "dna_from": (idx * 3) + 1,        # 1-based, inclusive
                "dna_to": (idx * 3) + 3,
            })
        new_dna = "".join(chars)
        return {
            "ok": True, "sequence": new_dna, "length": len(new_dna),
            "applied": applied, "level": "protein", "host": host,
            "length_changed": len(new_dna) - len(dna),
            "protein_before": aa, "protein_after": translate(new_dna),
        }

    # ── base-level ───────────────────────────────────────────────────────
    chars = list(dna)
    applied = []
    for spec, p in parsed:
        idx = p["pos"] - 1
        if idx >= len(dna):
            return {"ok": False, "error": (
                f"{spec}: this construct is {len(dna):,} bp, so there is no "
                f"base {p['pos']}.")}
        actual = dna[idx]
        if actual != p["from"]:
            return {"ok": False, "kind": "mismatch", "error": (
                f"{spec}: base {p['pos']} is {actual}, not {p['from']}. Not "
                f"editing — check the coordinate system, since an off-by-one "
                f"here silently mutates the neighbour.")}
        chars[idx] = p["to"]
        applied.append({"edit": f"{p['from']}{p['pos']}{p['to']}",
                        "level": "dna", "dna_from": p["pos"], "dna_to": p["pos"]})
    new_dna = "".join(chars)
    return {
        "ok": True, "sequence": new_dna, "length": len(new_dna),
        "applied": applied, "level": "dna", "host": host,
        "length_changed": 0,
    }