dee/core/prime_editor.py

"""Prime editing — design the pegRNA, with the geometry derived rather than guessed.

Every transversion and every small indel used to hit
``prime_editing_unavailable`` and stop. That was honest but it was also most of
the disease-causing variation in the genome: base editors make exactly two
changes (A>G, C>T), so a compiler without prime editing refuses the majority of
what a therapeutic programme actually brings it. This module is that pass.

WHY THIS FILE IS MOSTLY COORDINATE ARITHMETIC
---------------------------------------------
A pegRNA is four things stitched together, and three of them are positions:

    5'-[ spacer 20 nt ]-[ scaffold ]-[ RT template ]-[ PBS ]-3'

The spacer picks the nick. The PBS has to be the reverse complement of the
bases immediately 5' of that nick **on the strand Cas9 nicked**. The RT
template has to encode the corrected sequence continuing 3' from the nick, on
that same strand, read backwards. Get any one of those the wrong way round and
you get a well-formed RNA that installs nothing — or worse, installs something
else. There is no assay in this module that would catch it, and no scientist
reads a 60-mer and notices.

So the derivation is written out below, and then it is CHECKED: every pegRNA is
built, applied back to the patient sequence in silico, and the result compared
to the reference base for base. :func:`_build` refuses if reconstruction does
not reproduce the reference exactly. Same discipline as
``compiler.restores_wildtype`` and ``edits.py``'s refuse-on-mismatch, applied to
a construct with three coordinate systems in it.

THE DERIVATION
--------------
``find_guides`` reports ``position`` as the 1-based leftmost FORWARD index of
the 20-nt spacer footprint, on both strands. Let ``start = position - 1``.

Cas9 nicks 3 bp 5' of the PAM, i.e. between spacer positions 17 and 18. The
base that keeps the free 3'-OH is spacer position 17, so its forward index is
``spacer_pos_to_offset(position, strand, 20, 17)`` — the same mapping the
base-editing pass already uses, verified over 160 spacer positions across both
strands before anything was built on it. Concretely:

    '+' guide   nick base at start + 16   synthesis runs to INCREASING index
    '-' guide   nick base at start +  3   synthesis runs to DECREASING index

The nicked strand's 3' end is extended, so the RT template lies DOWNSTREAM of
the nick along that strand and the PBS lies upstream:

    '+'   PBS covers forward [nick-k+1, nick]      RTT covers [nick+1, ...]
    '-'   PBS covers forward [nick, nick+k-1]      RTT covers [..., nick-1]

Written 5'->3', the pegRNA's 3' extension is RT template FIRST, then PBS —
because reverse transcriptase reads the template 3'->5' while writing DNA
5'->3'. Which yields, and this is the identity worth remembering:

    extension == reverse_complement( corrected sequence spanning
                                     PBS region + newly written region,
                                     read on the nicked strand )

For a '-' guide that reverse complement cancels the strand flip and the RT
template comes out as plain forward-strand corrected sequence. That asymmetry
is real, not a bug, and it is the single most likely thing to be implemented
backwards.

WHAT IS DERIVED vs WHAT IS BORROWED
-----------------------------------
The geometry above is arithmetic — it follows from the enzymes and can be
re-derived by anyone. Everything that is an *empirical* design preference
(how long a PBS to start with, how far a PE3 nick should sit) is a published
heuristic, and every one of them is in :data:`HEURISTICS` with its source
attached. The distinction matters for the same reason it matters in
``compiler_validation``: a number that came from a paper and a number that
came from the genetic code are not the same kind of claim, and a design record
that blurs them is not auditable.

Pure logic: no network, no GPU, no model. Deterministic.
"""

from __future__ import annotations

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

from dee.core.compiler import Diagnostic, spacer_pos_to_offset

__all__ = [
    "PegRNA", "NickingGuide", "ExtensionVariant", "GeometryError",
    "design_pegrnas", "nick_index", "choose_pbs_len", "HEURISTICS",
    "SCAFFOLD_PLACEHOLDER", "pegrna_to_dict",
]

_COMPLEMENT = {"A": "T", "T": "A", "C": "G", "G": "C", "N": "N"}

SPACER_LEN = 20
NICK_SPACER_POS = 17        # Cas9 nicks 3 bp 5' of the PAM

# The pegRNA is emitted as spacer + <scaffold> + extension rather than as one
# orderable string, and that is deliberate. The scaffold is a construct
# decision, not a design output: plain sgRNA scaffold, an optimised variant, or
# an epegRNA with a 3' structured motif are different orders with different
# efficiencies. Typing a 76-mer from memory into a sequence somebody synthesises
# is exactly the class of mistake this module exists to prevent, so the slot is
# named and left to the person who knows which vector they are cloning into.
SCAFFOLD_PLACEHOLDER = "<sgRNA scaffold — supply the one your construct uses>"

# Published design preferences. Not geometry: these are empirical, they are
# starting points rather than optima, and each is attributed so a reviewer can
# check it instead of taking this file's word for it.
HEURISTICS: Dict[str, Dict[str, str]] = {
    "pbs_length": {
        "rule": "Start at 12-13 nt for a PBS footprint of ~40-60% GC. Test "
                "longer (14-15 nt) where GC is low and shorter (8-11 nt) "
                "where GC is high, then screen the range.",
        "source": "Anzalone et al. 2019 (PMID 31634902); Addgene prime-editing "
                  "design tips",
    },
    "rtt_length": {
        "rule": "Start with an RT template of ~10-16 nt. For edits close to "
                "the nick, screen a short (9-12), medium (13-16) and long "
                "(17-20) template rather than committing to one.",
        "source": "Anzalone et al. 2019 (PMID 31634902); pegFinder — Chow "
                  "et al., Nat Biomed Eng 2020",
    },
    "no_c_at_extension_start": {
        "rule": "The first base of the pegRNA's 3' extension should not be C. "
                "A C there is thought to pair with G81 of the scaffold and "
                "disturb the fold Cas9 binds; separately, efficiency suffers "
                "when the last templated nucleotide written is G.",
        "source": "Addgene prime-editing design tips, summarising Anzalone "
                  "et al. 2019",
    },
    "pe3_nick_window": {
        "rule": "A PE3 nick on the non-edited strand works best roughly 40-90 "
                "nt from the pegRNA nick and 3' of the edit. Design tools "
                "commonly widen the search to 40-150 nt.",
        "source": "Anzalone et al. 2019 (PMID 31634902); pegFinder default "
                  "search band",
    },
    "pe3b": {
        "rule": "A PE3b nicking guide matches only the EDITED strand, so the "
                "second nick cannot happen until the edit is installed. That "
                "avoids two simultaneous nicks and lowers indel byproducts.",
        "source": "Anzalone et al. 2019 (PMID 31634902)",
    },
}

# Bounds used for screening and for warnings. Named so nothing downstream
# re-invents them, and deliberately expressed as ranges to screen rather than
# single "correct" values, because that is what the sources actually say.
PBS_MIN, PBS_MAX = 8, 17
PBS_DEFAULT = 13
HOMOLOGY_DEFAULT = 10           # nt of RT template beyond the edit
HOMOLOGY_MIN_SAFE = 5
MAX_NICK_TO_EDIT = 30           # beyond this, route it but say it is a stretch
PE3_BAND = (40, 90)
PE3_SEARCH_BAND = (40, 150)


class GeometryError(Exception):
    """A pegRNA cannot be built at this position. Carries the reason."""


def _revcomp(seq: str) -> str:
    return "".join(_COMPLEMENT.get(b, "N") for b in reversed(seq.upper()))


def _gc_pct(seq: str) -> float:
    if not seq:
        return 0.0
    gc = sum(1 for b in seq.upper() if b in "GC")
    return round(100.0 * gc / len(seq), 1)


def nick_index(position: int, strand: str) -> int:
    """0-based forward index of the base bearing the free 3'-OH after nicking.

    Thin, named, and reusing the base-editing pass's own mapping on purpose:
    two independent implementations of a spacer-to-genome transform is how the
    two drift apart and one of them silently designs against the wrong strand.
    """
    return spacer_pos_to_offset(position, strand, SPACER_LEN, NICK_SPACER_POS)


def choose_pbs_len(footprint: str, *, default: int = PBS_DEFAULT) -> int:
    """PBS length from the GC content of the region it will anneal to.

    A starting point, not an optimum — see HEURISTICS["pbs_length"]. The
    screen matrix on every pegRNA exists because the honest answer is "test
    several", and a tool that emits one number implies it knows which.
    """
    gc = _gc_pct(footprint)
    if gc < 40.0:
        return min(default + 1, PBS_MAX)
    if gc > 60.0:
        return max(default - 2, PBS_MIN)
    return default


@dataclass
class ExtensionVariant:
    """One (PBS length, RT template length) pair to screen at the bench."""
    pbs_len: int
    rtt_len: int
    extension: str
    homology: int
    first_base_is_c: bool


@dataclass
class NickingGuide:
    """A second nick — PE3, or PE3b when it only matches the edited strand."""
    spacer: str
    pam: str
    strand: str              # opposite the strand the pegRNA nicked
    position: int            # 1-based forward start of the 20-nt footprint
    nick_offset: int         # 0-based forward index bearing the free 3'-OH
    signed_offset: int       # nt from the pegRNA nick, + = 3' along synthesis
    on_target_score: float
    pe3b: bool = False
    in_optimal_band: bool = False
    coords: str = "patient"  # "patient" | "edited" — PE3b is found on edited
    note: str = ""


@dataclass
class PegRNA:
    rank: int
    strand: str              # '+' = the pegRNA nicks the forward strand
    position: int            # 1-based forward start of the spacer footprint
    spacer: str
    pam: str
    nick_offset: int
    nick_to_edit: int        # nt of template written before the first change
    pbs: str
    pbs_len: int
    pbs_gc: float
    pbs_tm: Optional[float]
    rtt: str
    rtt_len: int
    extension: str           # 5'->3', RT template then PBS
    flap: str                # the DNA actually synthesised, 5'->3'
    flap_forward: str        # the same new DNA read on the FORWARD strand
    flap_start: int          # reference index where flap_forward begins
    nick_ref: int            # reference index the nick falls immediately after
    homology: int            # nt of template beyond the edit
    on_target_score: float
    composite_score: float
    reconstructed: bool      # the in-silico apply reproduced the reference
    re_engages_edited_allele: bool
    warnings: List[Diagnostic] = field(default_factory=list)
    nicks: List[NickingGuide] = field(default_factory=list)
    screen: List[ExtensionVariant] = field(default_factory=list)

    @property
    def parts(self) -> List[str]:
        """What you actually order, in order, with the scaffold left to you."""
        return [self.spacer, SCAFFOLD_PLACEHOLDER, self.extension]


@dataclass
class _Geom:
    flap: str
    rtt: str
    pbs: str
    extension: str
    nick_to_edit: int
    homology: int
    # The same new DNA expressed on the FORWARD strand, with the reference
    # index it starts at. Computed here rather than in the UI because it is
    # geometry, and because for an indel the reference and patient coordinate
    # systems differ by the indel's length — a renderer that assumed they
    # matched would draw the new strand off by that much and look right.
    flap_forward: str = ""
    flap_start: int = 0
    nick_ref: int = 0        # reference-coordinate boundary of the nick


def _corrected_nicked_strand(window: str, lo: int, hi: int, strand: str) -> str:
    """Reference sequence over forward span [lo, hi), read 5'->3' on the strand
    the pegRNA nicked. For a '-' guide that is the reverse complement, which is
    the whole reason this is a named function and not an inline slice."""
    seg = window[lo:hi]
    return seg if strand == "+" else _revcomp(seg)


def _check_nick_against_pam(patient: str, nick: int, strand: str,
                            pam: str) -> None:
    """The nick must sit 3 bp 5' of the PAM on the engaged strand.

    This exists because the reconstruction check below CANNOT catch a wrong
    nick. Everything in `_build` is derived from `nick`, so shifting it by one
    shifts the PBS and the RT template together and the reconstruction still
    reproduces the reference perfectly — it just describes a nick Cas9 will
    never make. Verified by deliberately shifting `nick_index` by +/-1: the
    reconstruction passed both times and five pegRNAs came out clean.

    So the nick is pinned to the PAM instead, which is the physical fact that
    fixes it, and is stated here in terms that do not reuse `nick_index`'s
    arithmetic.
    """
    if strand == "+":
        observed = patient[nick + 4:nick + 7]
    else:
        observed = _revcomp(patient[nick - 6:nick - 3])
    if observed != pam.upper():
        raise GeometryError(
            f"the nick at {nick} is not 3 bp 5' of the PAM on the {strand} "
            f"strand (found {observed!r} where {pam!r} was reported). Refusing "
            "to build a pegRNA around a nick position Cas9 would not make.")


def _build(window: str, offset: int, wt_allele: str, patient_allele: str,
           patient: str, nick: int, strand: str, pbs_len: int,
           homology: int, pam: str) -> _Geom:
    """Build one pegRNA extension and prove it installs the reference.

    `window` is the WILD-TYPE reference and is therefore also the INTENDED
    sequence — a therapeutic correction runs patient -> reference, so the
    target of the edit is the window itself. That makes the reconstruction
    check a comparison against something the caller supplied rather than
    against something this function computed, which is the only version of the
    check worth having.

    THREE checks, because each catches something the others cannot:
      1. the nick is 3 bp from the PAM        — catches a wrong nick
      2. reconstruction reproduces `window`   — catches a wrong RT template
      3. revcomp(extension) reads as the corrected nicked strand across
         PBS + flap                           — catches a wrong or
                                                wrong-orientation PBS, which
                                                reconstruction never touches

    Check 3 was added after a deliberate mutation — complementing the PBS on
    the '-' branch, the classic prime-editing mistake — produced five pegRNAs
    that passed reconstruction, because reconstruction only ever looks at the
    flap. A pegRNA whose PBS cannot anneal primes nothing.

    Raises GeometryError, with a reason, whenever a pegRNA cannot be built
    here. Refusing per-position is normal and expected: most spacers near a
    variant put the nick on the wrong side of it.
    """
    wlen, plen = len(wt_allele), len(patient_allele)
    lo, hi = offset, offset + plen          # changed span, PATIENT coordinates
    delta = wlen - plen

    if homology < 1:
        raise GeometryError("the RT template needs at least one base of "
                            "homology beyond the edit to anneal")
    if not (PBS_MIN <= pbs_len <= PBS_MAX):
        raise GeometryError(f"PBS length {pbs_len} is outside {PBS_MIN}-{PBS_MAX}")

    _check_nick_against_pam(patient, nick, strand, pam)

    if strand == "+":
        # Synthesis runs to increasing index, so the lesion must sit strictly
        # 3' of the nick. A '+' guide whose nick is past the variant cannot
        # reach it, however good the spacer is.
        if lo < nick + 1:
            raise GeometryError("the lesion is 5' of this guide's nick; a "
                                "pegRNA can only write downstream of the nick")
        end = offset + wlen + homology       # edited-coordinate end of the flap
        if end > len(window):
            raise GeometryError("not enough reference 3' of the edit for the "
                                "homology arm")
        if nick - pbs_len + 1 < 0:
            raise GeometryError("not enough sequence 5' of the nick for the PBS")
        flap = window[nick + 1:end]
        rtt = _revcomp(flap)
        pbs = _revcomp(patient[nick - pbs_len + 1:nick + 1])
        rebuilt = patient[:nick + 1] + flap + patient[hi + homology:]
        nick_to_edit = lo - nick
        span = (nick - pbs_len + 1, end)     # PBS + flap, forward coordinates
        flap_forward, flap_start, nick_ref = flap, nick + 1, nick
    else:
        # Synthesis runs to decreasing index: the lesion must sit 5' of the
        # nick in forward coordinates, and the RT template reads out as plain
        # forward-strand sequence because the two reverse complements cancel.
        if hi > nick:
            raise GeometryError("the lesion is 3' of this guide's nick on the "
                                "forward strand; an antisense pegRNA writes "
                                "toward decreasing coordinates")
        e_nick = nick + delta                # the nick in edited coordinates
        if offset - homology < 0:
            raise GeometryError("not enough reference 5' of the edit for the "
                                "homology arm")
        if nick + pbs_len > len(patient):
            raise GeometryError("not enough sequence 3' of the nick for the PBS")
        flap_fwd = window[offset - homology:e_nick]
        flap = _revcomp(flap_fwd)
        rtt = flap_fwd
        pbs = patient[nick:nick + pbs_len]
        rebuilt = window[:offset - homology] + flap_fwd + patient[nick:]
        nick_to_edit = nick - hi + 1
        span = (offset - homology, e_nick + pbs_len)
        flap_forward, flap_start, nick_ref = flap_fwd, offset - homology, e_nick - 1

    if not flap:
        raise GeometryError("the RT template came out empty")

    # CHECK 2 — the RT template installs the intended sequence. Everything
    # above is arithmetic across three coordinate systems; this is the line
    # that knows whether the arithmetic was right.
    if rebuilt != window:
        raise GeometryError(
            "in-silico reconstruction did not reproduce the reference "
            f"({len(rebuilt)} nt vs {len(window)} nt) — refusing to emit a "
            "pegRNA whose flap does not install the intended sequence")

    # CHECK 3 — the PBS. Reconstruction above never reads the PBS, so it
    # cannot tell an annealing primer-binding site from a scrambled one. This
    # asserts the identity from the module docstring instead: the whole 3'
    # extension, reverse-complemented, must read as the corrected nicked
    # strand running continuously across the PBS footprint and the new flap.
    # A complemented, mis-sized or wrongly-ordered PBS breaks it.
    extension = rtt + pbs
    if span[0] < 0 or span[1] > len(window):
        raise GeometryError("the PBS + flap span runs off the reference window")
    if _revcomp(extension) != _corrected_nicked_strand(window, span[0],
                                                       span[1], strand):
        raise GeometryError(
            "the 3' extension does not read as the corrected nicked strand "
            "across the PBS and the new flap — the primer-binding site cannot "
            "anneal where this pegRNA nicks")

    return _Geom(flap=flap, rtt=rtt, pbs=pbs, extension=extension,
                 nick_to_edit=nick_to_edit, homology=homology,
                 flap_forward=flap_forward, flap_start=flap_start,
                 nick_ref=nick_ref)


def _footprint(position: int, strand: str) -> Tuple[int, int]:
    """Forward span, inclusive, of spacer + PAM for a 20-nt Cas9 guide."""
    start = position - 1
    if strand == "+":
        return start, start + SPACER_LEN + 2       # PAM sits 3' of the spacer
    return start - 3, start + SPACER_LEN - 1       # PAM sits at lower indices


def _overlaps_lesion(position: int, strand: str, lo: int, hi: int) -> bool:
    """Does spacer+PAM cover any base the edit changes?

    An empty patient span (the correction INSERTS bases) is treated as the
    single junction it sits at, because an insertion between two covered bases
    still destroys the match.
    """
    a, b = _footprint(position, strand)
    if hi <= lo:                                   # insertion point
        return a <= lo <= b + 1
    return not (hi - 1 < a or lo > b)


def _tm(seq: str) -> Optional[float]:
    try:
        from dee.core import primers as _p
        return _p.tm_c(seq)
    except Exception:                              # noqa: BLE001
        return None


def _screen_matrix(window: str, offset: int, wt_allele: str,
                   patient_allele: str, patient: str, nick: int, strand: str,
                   pam: str, pbs_chosen: int, homology_chosen: int
                   ) -> List[ExtensionVariant]:
    """The lengths a bench scientist should actually order.

    Both sources say to screen rather than to trust one length, so the tool
    emits the screen. Only combinations that BUILD and pass reconstruction get
    in — an invalid geometry is silently absent rather than listed as
    something to try.
    """
    pbs_set = sorted({max(PBS_MIN, pbs_chosen - 3), pbs_chosen,
                      min(PBS_MAX, pbs_chosen + 2)})
    hom_set = sorted({max(1, HOMOLOGY_MIN_SAFE), homology_chosen,
                      homology_chosen + 6})
    out: List[ExtensionVariant] = []

    def rows_for(h: int) -> List[ExtensionVariant]:
        rows = []
        for k in pbs_set:
            try:
                g = _build(window, offset, wt_allele, patient_allele, patient,
                           nick, strand, k, h, pam)
            except GeometryError:
                continue
            rows.append(ExtensionVariant(
                pbs_len=k, rtt_len=len(g.rtt), extension=g.extension,
                homology=h, first_base_is_c=g.extension[:1].upper() == "C"))
        return rows

    for h in hom_set:
        out.extend(rows_for(h))

    # The first base of the extension is the 5' end of the RT template, so it
    # depends ONLY on the template length — varying the PBS can never change
    # it. A screen where every row starts with C therefore offers no way to act
    # on the no-C warning, which makes the warning noise. Walk the template
    # length outward by one base at a time until a row avoids it.
    if out and all(v.first_base_is_c for v in out):
        for step in (1, -1, 2, -2, 3, -3, 4):
            h = homology_chosen + step
            if h < 1 or h in hom_set:
                continue
            rows = rows_for(h)
            if rows and not rows[0].first_base_is_c:
                out.extend(rows)
                break
    return out


def _find_nicking_guides(patient_guides: Sequence[object],
                         edited_guides: Sequence[object],
                         patient: str, window: str, offset: int,
                         wt_allele: str, patient_allele: str,
                         peg_strand: str, peg_nick: int, nick_to_edit: int,
                         *, max_per_kind: int = 3) -> List[NickingGuide]:
    """PE3 and PE3b candidates for one pegRNA.

    PE3 nicks the OTHER strand — the one the pegRNA did not nick — so the
    nicking guide's strand is always the opposite of the pegRNA's. Distance is
    reported SIGNED, measured along the nicked strand's 5'->3' direction, with
    positive meaning 3' of the pegRNA nick. The sign is reported rather than
    silently filtered because "3' of the edit" is the published preference and
    a reader has to be able to see which side a candidate is on.

    Both guide lists are passed in rather than searched here: they are the same
    two lists for every pegRNA at this locus, and re-running the PAM scan once
    per pegRNA made this pass ~6x the cost of the whole design.
    """
    want = "-" if peg_strand == "+" else "+"
    plen, wlen = len(patient_allele), len(wt_allele)
    lo_e, hi_e = offset, offset + wlen
    delta = wlen - plen
    e_nick = peg_nick if peg_strand == "+" else peg_nick + delta

    def signed(nick2: int) -> int:
        return (nick2 - peg_nick) if peg_strand == "+" else (peg_nick - nick2)

    # ── PE3: on the patient sequence, in the published distance band ──────
    pe3: List[NickingGuide] = []
    for g in patient_guides:
        if g.strand != want:
            continue
        n2 = nick_index(g.position, g.strand)
        s = signed(n2)
        lo_band, hi_band = PE3_SEARCH_BAND
        if not (lo_band <= s <= hi_band):
            continue
        in_band = (PE3_BAND[0] <= s <= PE3_BAND[1]) and s > nick_to_edit
        note = ""
        if not in_band:
            if s > PE3_BAND[1]:
                note = (f"{s} nt away — outside the 40-90 nt band that worked "
                        "best in the original report, but inside the range "
                        "design tools search.")
            elif s <= nick_to_edit:
                note = (f"{s} nt away, which is not 3' of the edit "
                        f"({nick_to_edit} nt from the nick).")
        pe3.append(NickingGuide(
            spacer=g.spacer, pam=g.pam, strand=g.strand, position=g.position,
            nick_offset=n2, signed_offset=s,
            on_target_score=round(g.on_target_score, 3),
            in_optimal_band=in_band, coords="patient", note=note))
    pe3.sort(key=lambda x: (not x.in_optimal_band,
                            abs(x.signed_offset - sum(PE3_BAND) // 2),
                            -x.on_target_score))

    # ── PE3b: on the EDITED sequence, spanning the edit, absent from the
    # patient sequence. That last condition is what makes it PE3b rather than
    # a PE3 guide that happens to sit nearby: if the spacer also matches the
    # unedited allele it will nick before the edit lands, which is the failure
    # mode PE3b exists to avoid.
    pe3b: List[NickingGuide] = []
    patient_both = (patient, _revcomp(patient))
    for g in edited_guides:
        if g.strand != want:
            continue
        if not _overlaps_lesion(g.position, g.strand, lo_e, hi_e):
            continue
        if any(g.spacer in s for s in patient_both):
            continue                     # also matches the unedited allele
        n2 = nick_index(g.position, g.strand)
        pe3b.append(NickingGuide(
            spacer=g.spacer, pam=g.pam, strand=g.strand, position=g.position,
            nick_offset=n2, signed_offset=(n2 - e_nick) if peg_strand == "+"
                                          else (e_nick - n2),
            on_target_score=round(g.on_target_score, 3),
            pe3b=True, in_optimal_band=False, coords="edited",
            note="PE3b — this spacer matches only the corrected allele, so the "
                 "second nick cannot occur before the edit is installed. "
                 "Distance is not the selection criterion here."))
    pe3b.sort(key=lambda x: -x.on_target_score)
    return pe3b[:max_per_kind] + pe3[:max_per_kind]


def design_pegrnas(window: str, offset: int, wt_allele: str,
                   patient_allele: str, *,
                   pbs_len: int = 0,
                   homology: int = HOMOLOGY_DEFAULT,
                   max_nick_to_edit: int = MAX_NICK_TO_EDIT,
                   max_results: int = 6,
                   include_pe3: bool = True,
                   ) -> Tuple[List[PegRNA], List[Diagnostic]]:
    """pegRNAs that install the reference allele, ranked, or a reason there are none.

    `window` is the WILD-TYPE reference; `offset` is the 0-based index at which
    the wild-type allele starts within it. The patient sequence is DERIVED here
    rather than accepted, for the same reason the base-editing pass derives it:
    guides have to be found in the sequence that is actually in the cell. A
    pegRNA designed against the reference can have a spacer that does not match
    the patient's own allele, which is a construct that binds nothing.

    Returns (pegrnas, diagnostics). An empty list with a diagnostic is a real
    answer — for many variants no NGG puts a nick on the correct side within
    reach, and that is a fact about PAM availability rather than a threshold
    this module chose.
    """
    from dee.core import crispr as _crispr

    diags: List[Diagnostic] = []
    win = "".join(str(window or "").split()).upper()
    wt = "".join(str(wt_allele or "").split()).upper()
    pt = "".join(str(patient_allele or "").split()).upper()

    if not win:
        diags.append(Diagnostic(
            "error", "no_reference_window",
            "Prime editing needs real sequence — the spacer, the PBS and the "
            "RT template are all read off the locus.",
            "Supply the reference window around the variant."))
        return [], diags
    if not (0 <= offset <= len(win)):
        diags.append(Diagnostic(
            "error", "offset_outside_window",
            f"Offset {offset} is outside the {len(win)}-base window.",
            "Give the 0-based index of the wild-type allele in the window."))
        return [], diags
    if wt and win[offset:offset + len(wt)] != wt:
        diags.append(Diagnostic(
            "error", "reference_mismatch",
            f"The window has {win[offset:offset + len(wt)]!r} at offset "
            f"{offset}, but the wild-type allele was given as {wt!r}.",
            "Refusing rather than templating a correction toward a base the "
            "reference disagrees about. Check the coordinate and the strand."))
        return [], diags

    patient = win[:offset] + pt + win[offset + len(wt):]
    if patient == win:
        diags.append(Diagnostic(
            "error", "no_lesion",
            "The patient sequence derived from these alleles is identical to "
            "the reference.", "Check the alleles are not swapped."))
        return [], diags

    try:
        guides = _crispr.find_guides(patient, mode="knockout", max_results=500)
    except ValueError as exc:
        diags.append(Diagnostic(
            "error", "no_guides_possible", str(exc),
            "Prime editing still needs an NGG PAM. Supply a longer reference "
            "window so more PAMs are in range."))
        return [], diags

    lo = offset
    hi = offset + len(pt)
    out: List[PegRNA] = []
    reasons: Dict[str, int] = {}
    too_far = 0

    for g in guides:
        nick = nick_index(g.position, g.strand)
        footprint = (patient[max(0, nick - PBS_MAX + 1):nick + 1]
                     if g.strand == "+" else patient[nick:nick + PBS_MAX])
        k = pbs_len or choose_pbs_len(footprint)
        try:
            geom = _build(win, offset, wt, pt, patient, nick, g.strand, k,
                          homology, g.pam)
        except GeometryError as exc:
            reasons[str(exc)] = reasons.get(str(exc), 0) + 1
            continue

        if geom.nick_to_edit > max_nick_to_edit:
            too_far += 1
            continue

        warns: List[Diagnostic] = []
        # The one sequence-composition rule with a source behind it.
        if geom.extension[:1] == "C":
            warns.append(Diagnostic(
                "warning", "extension_starts_with_c",
                "The first base of the 3' extension is C. "
                + HEURISTICS["no_c_at_extension_start"]["rule"],
                "Shift the RT template length by one base, or use another "
                "spacer — the screen below includes lengths that avoid it."))
        if geom.homology < HOMOLOGY_MIN_SAFE:
            warns.append(Diagnostic(
                "warning", "short_homology_arm",
                f"Only {geom.homology} nt of RT template extends beyond the "
                "edit. The 3' flap needs homology to the genome to anneal.",
                f"Lengthen the RT template to at least {HOMOLOGY_MIN_SAFE} nt "
                "of post-edit homology."))
        if geom.nick_to_edit > 10:
            warns.append(Diagnostic(
                "note", "edit_far_from_nick",
                f"The edit sits {geom.nick_to_edit} nt from the nick. "
                "Efficiency generally falls as that distance grows, and it "
                "varies by locus and cell type.",
                "Prefer a spacer that nicks closer if one exists; otherwise "
                "screen RT template lengths and expect to optimise."))

        re_engages = not _overlaps_lesion(g.position, g.strand, lo, hi)
        pbs_gc = _gc_pct(geom.pbs)

        out.append(PegRNA(
            rank=0, strand=g.strand, position=g.position, spacer=g.spacer,
            pam=g.pam, nick_offset=nick, nick_to_edit=geom.nick_to_edit,
            pbs=geom.pbs, pbs_len=len(geom.pbs), pbs_gc=pbs_gc,
            pbs_tm=_tm(geom.pbs), rtt=geom.rtt, rtt_len=len(geom.rtt),
            extension=geom.extension, flap=geom.flap,
            flap_forward=geom.flap_forward, flap_start=geom.flap_start,
            nick_ref=geom.nick_ref, homology=geom.homology,
            on_target_score=round(g.on_target_score, 3),
            composite_score=round(g.composite_score, 3),
            reconstructed=True, re_engages_edited_allele=re_engages,
            warnings=warns))

    if not out:
        detail = ""
        if too_far:
            detail = (f" {too_far} guide(s) could template the correction but "
                      f"nick more than {max_nick_to_edit} nt away from it.")
        diags.append(Diagnostic(
            "error", "no_pegrna_reaches_this_lesion",
            "No NGG PAM puts a nick within reach on the correct side of this "
            "lesion." + detail,
            "A pegRNA can only write DOWNSTREAM of its nick, so the nick has "
            "to sit 5' of the edit on the strand it engages — for a sense "
            "guide that means upstream in forward coordinates, for an "
            "antisense guide downstream. Widen the reference window, or use a "
            "Cas9 variant with a different PAM."))
        return [], diags

    # Closer nicks first: the distance from nick to edit is the strongest
    # geometric lever on efficiency. Then prefer designs with no warning,
    # then the spacer's own on-target score.
    out.sort(key=lambda s: (s.nick_to_edit,
                            sum(1 for w in s.warnings if w.level == "warning"),
                            -s.on_target_score))
    kept = out[:max_results]
    for i, s in enumerate(kept, 1):
        s.rank = i

    # The screen matrix and the second-nick search are the expensive parts, so
    # they run only for what is actually being returned. The edited-sequence
    # PAM scan happens once for the whole locus, not once per pegRNA.
    edited_guides: List[object] = []
    if include_pe3:
        try:
            edited_guides = list(_crispr.find_guides(win, mode="knockout",
                                                     max_results=500))
        except ValueError:
            edited_guides = []
    for s in kept:
        s.screen = _screen_matrix(win, offset, wt, pt, patient, s.nick_offset,
                                  s.strand, s.pam, s.pbs_len, s.homology)
        if include_pe3:
            s.nicks = _find_nicking_guides(
                guides, edited_guides, patient, win, offset, wt, pt, s.strand,
                s.nick_offset, s.nick_to_edit)

    if not any(s.nicks for s in kept):
        diags.append(Diagnostic(
            "note", "no_second_nick_in_window",
            "No PE3 or PE3b nicking guide was found. "
            + HEURISTICS["pe3_nick_window"]["rule"],
            "PE2 (pegRNA alone) needs no second nick and produces fewer "
            "indels; PE3 raises efficiency at the cost of byproducts. A "
            "wider reference window is needed to search the full band."))

    if any(not s.re_engages_edited_allele for s in kept):
        diags.append(Diagnostic(
            "note", "edit_disrupts_own_protospacer",
            "At least one pegRNA's own spacer or PAM is changed by the edit it "
            "installs, so the corrected allele is no longer a target for it.",
            "That is usually desirable — it stops the corrected allele being "
            "re-engaged — and it is reported rather than scored because it "
            "trades against how far the edit sits from the nick."))

    diags.append(Diagnostic(
        "note", "pegrna_scaffold_not_supplied",
        "Each pegRNA is emitted as spacer + scaffold + 3' extension. The "
        "scaffold is a construct decision (plain sgRNA, an optimised variant, "
        "or an epegRNA with a 3' motif) and is not supplied here.",
        "Insert the scaffold your vector uses between the spacer and the "
        "extension. Note that the no-C rule above is about the first base of "
        "the extension precisely because of how it pairs with the scaffold."))
    return kept, diags


def pegrna_to_dict(p: PegRNA) -> Dict[str, object]:
    """JSON shape. Verbose on purpose — the record is the deliverable."""
    return {
        "rank": p.rank,
        "strand": p.strand,
        "position": p.position,
        "spacer": p.spacer,
        "pam": p.pam,
        "nick_offset": p.nick_offset,
        "nick_to_edit": p.nick_to_edit,
        "pbs": p.pbs,
        "pbs_len": p.pbs_len,
        "pbs_gc": p.pbs_gc,
        "pbs_tm": p.pbs_tm,
        "rtt": p.rtt,
        "rtt_len": p.rtt_len,
        "extension": p.extension,
        "flap": p.flap,
        "flap_forward": p.flap_forward,
        "flap_start": p.flap_start,
        "nick_ref": p.nick_ref,
        "homology": p.homology,
        "on_target_score": p.on_target_score,
        "composite_score": p.composite_score,
        "reconstructed": p.reconstructed,
        "re_engages_edited_allele": p.re_engages_edited_allele,
        "parts": p.parts,
        "warnings": [{"level": w.level, "code": w.code, "message": w.message,
                      "remedy": w.remedy} for w in p.warnings],
        "nicks": [{"spacer": n.spacer, "pam": n.pam, "strand": n.strand,
                   "position": n.position, "nick_offset": n.nick_offset,
                   "signed_offset": n.signed_offset, "pe3b": n.pe3b,
                   "in_optimal_band": n.in_optimal_band, "coords": n.coords,
                   "on_target_score": n.on_target_score, "note": n.note}
                  for n in p.nicks],
        "screen": [{"pbs_len": v.pbs_len, "rtt_len": v.rtt_len,
                    "extension": v.extension, "homology": v.homology,
                    "first_base_is_c": v.first_base_is_c}
                   for v in p.screen],
    }


def summarise(pegrnas: Sequence[PegRNA]) -> str:
    """One line for the enumerate pass. Counts, not adjectives."""
    if not pegrnas:
        return "no pegRNA reaches this lesion"
    n = len(pegrnas)
    clean = sum(1 for p in pegrnas
                if not any(w.level == "warning" for w in p.warnings))
    pe3 = sum(1 for p in pegrnas if any(not k.pe3b for k in p.nicks))
    pe3b = sum(1 for p in pegrnas if any(k.pe3b for k in p.nicks))
    nearest = min(p.nick_to_edit for p in pegrnas)
    return (f"{n} pegRNA(s); {clean} with no design warning; nearest nick "
            f"{nearest} nt from the edit; {pe3} with a PE3 nicking guide, "
            f"{pe3b} with a PE3b option. Every one was applied in silico and "
            "reproduced the reference exactly.")