cool/core/base_editor.py

"""Base editor and prime editor suitability analysis for plants.

Implements:
1. CBE (Cytosine Base Editor) — C→T conversion, window 4-8 from PAM-distal end
2. ABE (Adenine Base Editor) — A→G conversion, window 4-7
3. Prime Editing pegRNA designer — PBS + RT template scoring
4. Cas12a (TTTV) guide efficiency — nucleotide position preference model
5. Nucleosome occupancy approximation — TA/AA periodicity for CRISPR accessibility

Sources:
- CBE/ABE windows: Komor 2016, Gaudelli 2017; plant validation IJMS 2025
- Cas12a plant efficiency: PMC 2025 Random Forest features
- Nucleosome occupancy: Nature Methods 2025 (TDAC-seq); TF/AT periodicity model
- Prime editing in plants: MDPI Genes 2025 (TwinPE 44.2% efficiency)
"""
import math
import re

try:
    import RNA as _RNA
    _VIENNA_OK = True
except ImportError:
    _VIENNA_OK = False

from .codons import CODON_TO_AMINO_ACID
from Bio.SeqUtils import MeltingTemp as mt
from Bio.Seq import Seq


# ── 1. CBE Suitability Scanner ────────────────────────────────────────────────

# CBE deaminase window: positions 4-8 counting from the 5' end of the spacer
# (PAM-distal = position 1). Cytosines in this window are converted to T.
# Context preferences (highest to lowest efficiency in plants):
#   TC context: high (TC·target is optimal for BE3 and evoBE4)
#   CC context: medium
#   AC/GC context: lower
# Source: Komor 2016 Nature; Jin 2019 Plant Biotechnology Journal.

CBE_WINDOW = range(3, 8)   # 0-indexed positions 3-7 = spacer positions 4-8

_CBE_CONTEXT_SCORE = {
    'TC': 1.0, 'CC': 0.75, 'AC': 0.55, 'GC': 0.45,
}


def cbe_sites(dna: str, pam_type: str = "NGG") -> list:
    """Find all cytosines editable by CBE within their PAM-proximal spacer windows.

    Returns list of dicts: {position, spacer, C_positions_in_window, efficiency_score,
    pam, context_scores}
    """
    dna = dna.upper()
    pam_seqs = {"NGG": ["AGG", "TGG", "CGG", "GGG"],
                "TTTV": ["TTTA", "TTTC", "TTTG"]}
    pams = pam_seqs.get(pam_type, pam_seqs["NGG"])
    pam_len = len(pams[0])
    sites = []

    for i in range(len(dna) - 20 - pam_len + 1):
        pam = dna[i + 20:i + 20 + pam_len]
        if pam not in pams:
            continue
        spacer = dna[i:i + 20]
        c_positions = [p for p in CBE_WINDOW if p < len(spacer) and spacer[p] == 'C']
        if not c_positions:
            continue

        eff_scores = []
        for p in c_positions:
            context = spacer[p - 1:p + 1] if p > 0 else 'NC'
            eff_scores.append(_CBE_CONTEXT_SCORE.get(context, 0.35))

        sites.append({
            "type": "CBE",
            "pam_type": pam_type,
            "spacer_start": i,
            "spacer": spacer,
            "pam": pam,
            "editable_positions": c_positions,      # 0-indexed in spacer
            "efficiency_score": round(max(eff_scores), 3),
            "context_scores": eff_scores,
        })

    return sorted(sites, key=lambda x: x["efficiency_score"], reverse=True)


# ── 2. ABE Suitability Scanner ────────────────────────────────────────────────

# ABE deaminase window: positions 4-7 (0-indexed 3-6).
# Adenines in this window are converted A→G (creates G on the non-edited strand).
# Context preferences: TA and AA contexts best (Gaudelli 2017, Nature).

ABE_WINDOW = range(3, 7)   # 0-indexed positions 3-6 = spacer positions 4-7

_ABE_CONTEXT_SCORE = {
    'TA': 1.0, 'AA': 0.90, 'GA': 0.65, 'CA': 0.55,
}


def abe_sites(dna: str, pam_type: str = "NGG") -> list:
    """Find all adenines editable by ABE within their PAM-proximal spacer windows."""
    dna = dna.upper()
    pams_map = {"NGG": ["AGG", "TGG", "CGG", "GGG"],
                "TTTV": ["TTTA", "TTTC", "TTTG"]}
    pams = pams_map.get(pam_type, pams_map["NGG"])
    pam_len = len(pams[0])
    sites = []

    for i in range(len(dna) - 20 - pam_len + 1):
        pam = dna[i + 20:i + 20 + pam_len]
        if pam not in pams:
            continue
        spacer = dna[i:i + 20]
        a_positions = [p for p in ABE_WINDOW if p < len(spacer) and spacer[p] == 'A']
        if not a_positions:
            continue

        eff_scores = []
        for p in a_positions:
            context = spacer[p - 1:p + 1] if p > 0 else 'NA'
            eff_scores.append(_ABE_CONTEXT_SCORE.get(context, 0.40))

        sites.append({
            "type": "ABE",
            "pam_type": pam_type,
            "spacer_start": i,
            "spacer": spacer,
            "pam": pam,
            "editable_positions": a_positions,
            "efficiency_score": round(max(eff_scores), 3),
            "context_scores": eff_scores,
        })

    return sorted(sites, key=lambda x: x["efficiency_score"], reverse=True)


# ── 3. Prime Editing pegRNA Designer ─────────────────────────────────────────
# Prime editing requires:
#   (a) A spacer (20 nt, same design as Cas9 gRNA)
#   (b) Primer Binding Site (PBS): reverse complement of 3'-flap, 8-15 nt,
#       Tm 55-65°C
#   (c) RT template: encodes the desired edit + downstream sequence, 10-30 nt
#
# TwinPE in plants (MDPI 2025) achieved 44.2% efficiency with PBS Tm 58-63°C,
# RT template 20-25 nt, minimal secondary structure in the extension.

PBS_TM_OPTIMAL = (58, 63)   # °C
RT_TEMPLATE_OPTIMAL = (20, 25)  # nt


def design_pegrna(spacer: str, edit_site_offset: int, edit_seq: str,
                  flap_3prime: str) -> dict:
    """Design a pegRNA for prime editing.

    Parameters
    ----------
    spacer          : 20-nt protospacer (PAM-distal to PAM-proximal).
    edit_site_offset: number of nt from the nick site to the edit position.
    edit_seq        : the sequence change to encode in the RT template.
    flap_3prime     : 3'-flap sequence on the non-template strand (used to
                      compute PBS = reverse complement of the 3' end of this).

    Returns
    -------
    dict with PBS sequence, Tm, RT template, structure score, and overall score.
    """
    # PBS: reverse complement of last 13 nt of the 3'-flap (adjustable 8-15 nt)
    for pbs_len in range(13, 7, -1):
        pbs = str(Seq(flap_3prime[-pbs_len:]).reverse_complement())
        pbs_tm = float(mt.Tm_NN(Seq(pbs)))
        if PBS_TM_OPTIMAL[0] <= pbs_tm <= PBS_TM_OPTIMAL[1]:
            break
    else:
        pbs_tm = float(mt.Tm_NN(Seq(pbs)))

    # RT template: edit + downstream context
    rt = edit_seq + flap_3prime[:edit_site_offset]
    rt_len = len(rt)

    # Structure of the 3' extension (PBS + RT): should be low MFE
    extension = pbs + rt
    struct_score = 1.0
    if _VIENNA_OK and len(extension) >= 8:
        _ss, mfe = _RNA.fold(extension.replace('T', 'U'))
        struct_score = max(0.0, min(1.0, 1.0 + mfe / 20.0))

    # Overall pegRNA score
    tm_score = 1.0 if PBS_TM_OPTIMAL[0] <= pbs_tm <= PBS_TM_OPTIMAL[1] else \
        max(0.0, 1.0 - abs(pbs_tm - 60.0) / 10.0)
    len_score = 1.0 if RT_TEMPLATE_OPTIMAL[0] <= rt_len <= RT_TEMPLATE_OPTIMAL[1] else \
        max(0.0, 1.0 - abs(rt_len - 22.0) / 15.0)
    overall = round((tm_score * 0.4 + len_score * 0.3 + struct_score * 0.3), 3)

    return {
        "pbs": pbs,
        "pbs_tm": round(pbs_tm, 1),
        "rt_template": rt,
        "rt_length": rt_len,
        "extension_structure_score": round(struct_score, 3),
        "overall_score": overall,
    }


# ── 4. Cas12a Guide Efficiency Scorer ─────────────────────────────────────────
# Random Forest features from PMC 2025 (+15% over existing algorithms):
# - PAM: TTTV (TTTA > TTTC > TTTG, in order of efficiency)
# - Position 1 (PAM-proximal): A preferred (+0.15)
# - Position 1-5: A/T preferred, G/C penalised at 1-3
# - Seed region (13-20, PAM-proximal): positions 14-20 most critical
# - Homopolymer T-runs penalised (premature U6 termination)
# - GC content 40-60% optimal

_PAM_EFFICIENCY = {"TTTA": 1.00, "TTTC": 0.90, "TTTG": 0.80}

# Position-specific nucleotide preferences for Cas12a (1-based, 1=PAM-proximal)
# Derived from pooled-screen data in PMC 2025 (Table S4 approximate values)
_CAS12A_POS_WEIGHTS = {
    1:  {'A': +0.15, 'T': +0.05, 'C': -0.10, 'G': -0.10},
    2:  {'A': +0.10, 'T': +0.08, 'C': -0.08, 'G': -0.10},
    3:  {'T': +0.12, 'A': +0.06, 'C': -0.08, 'G': -0.10},
    4:  {'T': +0.10, 'A': +0.08, 'C': -0.05, 'G': -0.08},
    5:  {'A': +0.08, 'T': +0.08, 'C': -0.04, 'G': -0.05},
    # Seed region (13-20, counting from 5' of spacer = positions 13-20):
    13: {'G': +0.08, 'C': +0.05, 'A': -0.04, 'T': -0.05},
    14: {'G': +0.10, 'C': +0.08, 'A': -0.06, 'T': -0.08},
    15: {'G': +0.12, 'C': +0.10, 'A': -0.08, 'T': -0.10},
    16: {'G': +0.12, 'C': +0.10, 'A': -0.08, 'T': -0.10},
    17: {'G': +0.15, 'C': +0.12, 'A': -0.10, 'T': -0.12},
    18: {'G': +0.15, 'C': +0.12, 'A': -0.10, 'T': -0.12},
    19: {'G': +0.18, 'C': +0.14, 'A': -0.12, 'T': -0.14},
    20: {'G': +0.20, 'C': +0.16, 'A': -0.14, 'T': -0.16},
}


def cas12a_guide_score(spacer: str, pam: str = "TTTA") -> float:
    """Score a 23-nt Cas12a guide (spacer only, 20nt) for plant efficiency.

    Returns [0, 1]; 0.8+ = high efficiency predicted.
    """
    spacer = spacer.upper()[:20]
    if len(spacer) < 20:
        return 0.0

    # PAM efficiency base
    score = _PAM_EFFICIENCY.get(pam.upper(), 0.70)

    # Position-specific preferences (1-indexed from 5' of spacer)
    for i, base in enumerate(spacer, start=1):
        score += _CAS12A_POS_WEIGHTS.get(i, {}).get(base, 0.0)

    # GC content penalty
    gc = sum(1 for b in spacer if b in 'GC') / 20
    if not (0.40 <= gc <= 0.60):
        score -= 0.15 * abs(gc - 0.50)

    # Poly-T run penalty
    if 'TTTT' in spacer:
        score -= 0.25

    return round(max(0.0, min(1.0, score)), 4)


def design_cas12a_guides(dna: str, guide_num: int = 3) -> list:
    """Find and score Cas12a (TTTV) guides in a DNA sequence."""
    dna = dna.upper()
    pam_seqs = {"TTTA": 0, "TTTC": 1, "TTTG": 2}
    guides = []
    for i in range(len(dna) - 24 + 1):
        pam = dna[i:i + 4]
        if pam not in pam_seqs:
            continue
        spacer = dna[i + 4:i + 24]
        if len(spacer) < 20:
            continue
        eff = cas12a_guide_score(spacer, pam)
        guides.append({"guide": spacer, "pam": pam, "position": i,
                       "efficiency": eff, "type": "Cas12a"})
    guides.sort(key=lambda g: g["efficiency"], reverse=True)
    return guides[:guide_num]


# ── 5. Nucleosome Occupancy Approximation ─────────────────────────────────────
# Nucleosome DNA wrapping is characterised by:
# - AA/TT dinucleotides every ~10 bp (in phase with the minor groove)
# - TA dinucleotides at specific positions within the nucleosome
# - GC-rich sequences resist nucleosome wrapping (more open chromatin)
# Source: Nature Methods 2025 (TDAC-seq single-nucleotide resolution).
# Accessible chromatin = better Cas9/Cas12a activity (order-of-magnitude effect).

def nucleosome_occupancy_score(dna: str) -> float:
    """Estimate nucleosome occupancy in [0, 1]; 0=open/accessible, 1=occluded.

    Based on AA/TT 10-bp periodicity and GC content.
    High score means the guide RNA target is likely wrapped in a nucleosome
    and will have REDUCED Cas9/Cas12a activity.
    """
    dna = dna.upper()
    n = len(dna)
    if n < 10:
        return 0.5

    # 1. AA/TT 10-bp periodicity — nucleosome signal
    # Count AA and TT at offsets that are multiples of 10 bp.
    aa_tt_periodic = 0
    total_dinucs = 0
    for start in range(0, n - 10, 10):
        window = dna[start:start + 10]
        for i in range(len(window) - 1):
            di = window[i:i + 2]
            if di in ('AA', 'TT', 'TA'):
                aa_tt_periodic += 1
            total_dinucs += 1
    # Fraction of scored dinucleotides that are AA/TT/TA (each 10-bp window
    # contributes 9 dinucleotides). Dividing by the actual dinucleotide count
    # keeps the score in [0, 1]; the previous /3 over-counted ~3× and saturated.
    periodicity_score = aa_tt_periodic / total_dinucs if total_dinucs else 0

    # 2. GC content — inverse of occupancy
    gc = sum(1 for b in dna if b in 'GC') / n
    gc_accessibility = gc  # more GC = more accessible

    # Nucleosome occupancy = high periodicity + low GC
    occupancy = 0.5 * periodicity_score + 0.5 * (1.0 - gc_accessibility)
    return round(max(0.0, min(1.0, occupancy)), 4)


def crispr_chromatin_score(dna: str) -> float:
    """CRISPR accessibility score accounting for nucleosome occupancy.

    Returns [0, 1]; 1.0 = open chromatin, predicted highest Cas activity.
    """
    return round(1.0 - nucleosome_occupancy_score(dna), 4)