sequence_core.py

"""
sequence_core.py
================
Core nucleic acid transformation algorithms supporting both DNA and RNA sequences.
Includes pedagogical loop-based approaches and high-performance vectorized lookup tables.
"""

from typing import Dict, Optional

# Standard DNA Watson-Crick complement mapping
DNA_COMPLEMENT_MAP: Dict[str, str] = {
    'A': 'T', 'a': 't',
    'T': 'A', 't': 'a',
    'C': 'G', 'c': 'g',
    'G': 'C', 'g': 'c',
    'N': 'N', 'n': 'n',
    'X': 'X', 'x': 'x',
}
DNA_TRANS_TABLE = str.maketrans("ATCGNatcgn", "TAGCNtagcn")

# RNA Watson-Crick complement mapping (RNA -> RNA)
RNA_COMPLEMENT_MAP: Dict[str, str] = {
    'A': 'U', 'a': 'u',
    'U': 'A', 'u': 'a',
    'C': 'G', 'c': 'g',
    'G': 'C', 'g': 'c',
    'N': 'N', 'n': 'n',
    'X': 'X', 'x': 'x',
}
RNA_TRANS_TABLE = str.maketrans("AUCGNaucgn", "UAGCNuagcn")

# Template DNA to mRNA transcription mapping (Template 3'->5' to mRNA 5'->3')
# A -> U, T -> A, C -> G, G -> C
TEMPLATE_DNA_TO_RNA_TABLE = str.maketrans("ATCGNatcgn", "UAGCNuagcn")

# Coding DNA (Sense) to mRNA mapping (T -> U)
CODING_DNA_TO_RNA_TABLE = str.maketrans("Tt", "Uu")
RNA_TO_CODING_DNA_TABLE = str.maketrans("Uu", "Tt")


def reverse_sequence(seq: str) -> str:
    """
    Reverse a nucleotide sequence using Python slicing.
    
    Parameters
    ----------
    seq : str
        Input nucleotide sequence (5' to 3' or 3' to 5').
        
    Returns
    -------
    str
        Reversed sequence string.
        
    Examples
    --------
    >>> reverse_sequence("ATGCCG")
    'GCCGTA'
    """
    return seq[::-1]


def complement_dna(seq: str, method: str = "table") -> str:
    """
    Generate the complementary DNA strand.
    
    Watson-Crick rules:
    - Adenine (A) <-> Thymine (T)
    - Cytosine (C) <-> Guanine (G)
    
    Parameters
    ----------
    seq : str
        DNA nucleotide sequence.
    method : str, default "table"
        'table' for high-speed str.translate, or 'loop' for educational step-by-step logic.
        
    Returns
    -------
    str
        Complementary DNA sequence.
    """
    if method == "table":
        # Fast translation table lookup
        return seq.translate(DNA_TRANS_TABLE)
    elif method == "loop":
        comp = []
        for base in seq:
            upper_base = base.upper()
            if upper_base == "A":
                comp.append("T" if base.isupper() else "t")
            elif upper_base == "T":
                comp.append("A" if base.isupper() else "a")
            elif upper_base == "C":
                comp.append("G" if base.isupper() else "g")
            elif upper_base == "G":
                comp.append("C" if base.isupper() else "c")
            elif upper_base in ("N", "X"):
                comp.append(base)
            else:
                comp.append("?")
        return "".join(comp)
    else:
        raise ValueError(f"Unknown method '{method}'. Choose 'table' or 'loop'.")


def reverse_complement_dna(seq: str, method: str = "table") -> str:
    """
    Compute the reverse-complement of a DNA sequence.
    
    This simulates standard primer binding or sequencing read alignment on the opposite strand:
    Step 1: Invert sequence order (Reverse, 3' <- 5').
    Step 2: Substitute Watson-Crick complements.
    
    Parameters
    ----------
    seq : str
        Input DNA sequence (5' -> 3').
    method : str, default "table"
        'table' or 'loop'.
        
    Returns
    -------
    str
        Reverse complement sequence (5' -> 3' of opposite strand).
    """
    if method == "table":
        return seq[::-1].translate(DNA_TRANS_TABLE)
    elif method == "loop":
        # Can complement then reverse, or reverse and complement
        return complement_dna(seq[::-1], method="loop")
    else:
        raise ValueError(f"Unknown method '{method}'. Choose 'table' or 'loop'.")


def transcribe_template_to_rna(template_dna: str, method: str = "table") -> str:
    """
    Transcribe a DNA template strand (3' to 5') into an mRNA transcript (5' to 3').
    
    Complementary Base Pairing:
    - Template A -> mRNA U
    - Template T -> mRNA A
    - Template C -> mRNA G
    - Template G -> mRNA C
    
    Parameters
    ----------
    template_dna : str
        DNA template sequence.
    method : str, default "table"
        'table' or 'loop'.
        
    Returns
    -------
    str
        mRNA transcript.
    """
    if method == "table":
        return template_dna.translate(TEMPLATE_DNA_TO_RNA_TABLE)
    elif method == "loop":
        rna = []
        for base in template_dna:
            upper_base = base.upper()
            if upper_base == "A":
                rna.append("U" if base.isupper() else "u")
            elif upper_base == "T":
                rna.append("A" if base.isupper() else "a")
            elif upper_base == "C":
                rna.append("G" if base.isupper() else "g")
            elif upper_base == "G":
                rna.append("C" if base.isupper() else "c")
            elif upper_base in ("N", "X"):
                rna.append(base)
            else:
                rna.append("?")
        return "".join(rna)
    else:
        raise ValueError(f"Unknown method '{method}'. Choose 'table' or 'loop'.")


def transcribe_coding_to_rna(coding_dna: str) -> str:
    """
    Transcribe coding DNA strand (sense strand, 5' to 3') to mRNA.
    Simply replaces Thymine (T) with Uracil (U).
    """
    return coding_dna.translate(CODING_DNA_TO_RNA_TABLE)


def reverse_transcribe_to_dna(rna_seq: str) -> str:
    """
    Reverse transcribe RNA (e.g. viral genomic RNA) to cDNA coding sequence (U -> T).
    """
    return rna_seq.translate(RNA_TO_CODING_DNA_TABLE)


def complement_rna(rna_seq: str, method: str = "table") -> str:
    """
    Generate the complementary RNA strand for an RNA molecule.
    
    Rules:
    - A <-> U
    - C <-> G
    """
    if method == "table":
        return rna_seq.translate(RNA_TRANS_TABLE)
    elif method == "loop":
        comp = []
        for base in rna_seq:
            upper_base = base.upper()
            if upper_base == "A":
                comp.append("U" if base.isupper() else "u")
            elif upper_base == "U":
                comp.append("A" if base.isupper() else "a")
            elif upper_base == "C":
                comp.append("G" if base.isupper() else "g")
            elif upper_base == "G":
                comp.append("C" if base.isupper() else "c")
            elif upper_base in ("N", "X"):
                comp.append(base)
            else:
                comp.append("?")
        return "".join(comp)
    else:
        raise ValueError(f"Unknown method '{method}'. Choose 'table' or 'loop'.")


def reverse_complement_rna(rna_seq: str, method: str = "table") -> str:
    """
    Compute the reverse complement of an RNA sequence.
    """
    if method == "table":
        return rna_seq[::-1].translate(RNA_TRANS_TABLE)
    elif method == "loop":
        return complement_rna(rna_seq[::-1], method="loop")
    else:
        raise ValueError(f"Unknown method '{method}'. Choose 'table' or 'loop'.")


def is_valid_sequence(seq: str, seq_type: str = "DNA", allow_ambiguous: bool = True) -> bool:
    """
    Validate whether a sequence only contains valid nucleotide characters.
    """
    canonical_dna = set("ATCGatcg")
    canonical_rna = set("AUCGaucg")
    ambiguous = set("NXnx-")

    valid_chars = canonical_dna if seq_type.upper() == "DNA" else canonical_rna
    if allow_ambiguous:
        valid_chars = valid_chars.union(ambiguous)

    return all(char in valid_chars for char in seq)