"""Codon utilities and canonical bounded embedding."""

from __future__ import annotations

from functools import cache

import numpy as np

NUCLEOTIDES = ("A", "C", "G", "T")
STOP_CODONS = {"TAA", "TAG", "TGA"}

GENETIC_CODE = {
    "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",
    "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",
    "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",
    "TAA": "*",
    "TAG": "*",
    "TGA": "*",
}

SENSE_CODONS = tuple(c for c, aa in GENETIC_CODE.items() if aa != "*")
HYDROPHOBIC_AA = {"A", "V", "I", "L", "M", "F", "Y", "W"}


@cache
def codon_to_embedding(codon: str) -> np.ndarray:
    """Canonical bounded codon embedding u(c) in R^8.

    Coordinates:
    1-4: nucleotide composition fractions (A/C/G/T)
    5: GC fraction
    6: purine fraction (A+G)
    7: transition-like adjacent-pair fraction
    8: amino-acid hydrophobicity indicator in {-1, +1} (0 for ambiguous/gap/stop)
    """
    c = codon.upper()
    if len(c) != 3 or any(base not in {"A", "C", "G", "T"} for base in c):
        return np.zeros(8, dtype=float)

    vec = np.zeros(8, dtype=float)
    for idx, nt in enumerate(NUCLEOTIDES):
        vec[idx] = c.count(nt) / 3.0

    vec[4] = (c.count("G") + c.count("C")) / 3.0
    vec[5] = (c.count("A") + c.count("G")) / 3.0

    transitions = {("A", "G"), ("G", "A"), ("C", "T"), ("T", "C")}
    pairs = [(c[0], c[1]), (c[1], c[2])]
    vec[6] = sum(1 for p in pairs if p in transitions) / 2.0

    aa = GENETIC_CODE.get(c, "X")
    if aa == "*":
        vec[7] = 0.0
    elif aa in HYDROPHOBIC_AA:
        vec[7] = 1.0
    else:
        vec[7] = -1.0

    return vec


def is_valid_sense_codon(codon: str) -> bool:
    c = codon.upper()
    return c in SENSE_CODONS


def is_ambiguous_or_gap(codon: str) -> bool:
    c = codon.upper()
    return len(c) != 3 or any(base not in {"A", "C", "G", "T"} for base in c)
