# geney/folding/codons.py — Codon table and substitution classification.
"""Standard genetic code and synonymous/non-synonymous classification for SNVs."""

from __future__ import annotations

__all__ = ["classify_substitution", "get_amino_acid", "CODON_TABLE"]

CODON_TABLE = {
    "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",
}


def get_amino_acid(codon: str) -> str:
    """Translate a 3-letter DNA codon to its amino acid (or '*' for stop).

    Returns 'X' for unrecognised codons.
    """
    return CODON_TABLE.get(codon.upper(), "X")


def classify_substitution(cds_seq: str, position_0based: int, alt_base: str) -> str:
    """Classify a single-nucleotide substitution as synonymous or non-synonymous.

    Args:
        cds_seq: CDS DNA sequence (in frame, position 0 is first base of first codon).
        position_0based: 0-based index of the substituted base within the CDS.
        alt_base: Alternative base (A/C/G/T).

    Returns:
        ``'synonymous'`` or ``'non_synonymous'``.

    Raises:
        ValueError: If *position_0based* is out of range or *alt_base* is invalid.
    """
    seq = cds_seq.upper()
    alt = alt_base.upper()
    n = len(seq)

    if position_0based < 0 or position_0based >= n:
        raise ValueError(f"position_0based {position_0based} out of range [0, {n})")
    if alt not in "ACGT" or len(alt) != 1:
        raise ValueError(f"alt_base must be a single character in ACGT, got {alt_base!r}")

    codon_start = (position_0based // 3) * 3
    codon_end = codon_start + 3

    if codon_end > n:
        # Incomplete trailing codon — cannot determine synonymy
        return "non_synonymous"

    ref_codon = seq[codon_start:codon_end]
    pos_in_codon = position_0based % 3
    var_codon = ref_codon[:pos_in_codon] + alt + ref_codon[pos_in_codon + 1:]

    return "synonymous" if get_amino_acid(ref_codon) == get_amino_acid(var_codon) else "non_synonymous"
