"""Canonical default codon neutral model (MG94-like)."""

from __future__ import annotations

from dataclasses import asdict, dataclass

import numpy as np

from babappa.features.codon import GENETIC_CODE, SENSE_CODONS
from babappa.io.fasta import CodonAlignment


@dataclass
class CodonNullModel:
    model_name: str
    codons: list[str]
    codon_frequencies: list[float]
    kappa: float
    omega: float
    codon_frequency_mode: str
    rate_variation: bool
    gamma_alpha: float

    def to_dict(self) -> dict[str, object]:
        return asdict(self)


def _f3x4_codon_frequencies(alignment: CodonAlignment) -> dict[str, float]:
    pos_counts = [dict(A=1.0, C=1.0, G=1.0, T=1.0) for _ in range(3)]  # pseudocounts

    nt_len = alignment.nt_length
    for seq in alignment.sequences.values():
        for i in range(0, nt_len, 3):
            codon = seq[i : i + 3].upper()
            if len(codon) != 3 or any(nt not in {"A", "C", "G", "T"} for nt in codon):
                continue
            if codon not in SENSE_CODONS:
                continue
            for pos in range(3):
                pos_counts[pos][codon[pos]] += 1.0

    pos_freqs: list[dict[str, float]] = []
    for counts in pos_counts:
        total = sum(counts.values())
        pos_freqs.append({k: v / total for k, v in counts.items()})

    probs: dict[str, float] = {}
    for codon in SENSE_CODONS:
        p = pos_freqs[0][codon[0]] * pos_freqs[1][codon[1]] * pos_freqs[2][codon[2]]
        probs[codon] = p

    total = sum(probs.values())
    return {k: v / total for k, v in probs.items()}


def build_codon_frequencies(alignment: CodonAlignment, mode: str = "f3x4") -> dict[str, float]:
    mode_norm = mode.lower()
    if mode_norm == "uniform":
        w = 1.0 / len(SENSE_CODONS)
        return {c: w for c in SENSE_CODONS}
    if mode_norm == "f3x4":
        return _f3x4_codon_frequencies(alignment)
    raise ValueError("Unsupported codon_frequency_mode. Use 'f3x4' or 'uniform'.")


def make_default_null_model(
    alignment: CodonAlignment,
    kappa: float,
    omega: float,
    codon_frequency_mode: str,
    rate_variation: bool,
    gamma_alpha: float,
) -> CodonNullModel:
    freqs = build_codon_frequencies(alignment, mode=codon_frequency_mode)
    return CodonNullModel(
        model_name="mg94_neutral",
        codons=list(SENSE_CODONS),
        codon_frequencies=[float(freqs[c]) for c in SENSE_CODONS],
        kappa=float(kappa),
        omega=float(omega),
        codon_frequency_mode=codon_frequency_mode,
        rate_variation=bool(rate_variation),
        gamma_alpha=float(gamma_alpha),
    )


def is_synonymous(c1: str, c2: str) -> bool:
    return GENETIC_CODE[c1] == GENETIC_CODE[c2]


def is_transition(a: str, b: str) -> bool:
    return (a, b) in {("A", "G"), ("G", "A"), ("C", "T"), ("T", "C")}


def codon_rate_matrix(model: CodonNullModel) -> tuple[np.ndarray, list[str], np.ndarray]:
    codons = model.codons
    pi = np.array(model.codon_frequencies, dtype=float)
    n = len(codons)
    q = np.zeros((n, n), dtype=float)

    for i, ci in enumerate(codons):
        for j, cj in enumerate(codons):
            if i == j:
                continue
            diffs = [k for k in range(3) if ci[k] != cj[k]]
            if len(diffs) != 1:
                continue
            pos = diffs[0]
            rate = model.kappa if is_transition(ci[pos], cj[pos]) else 1.0
            if not is_synonymous(ci, cj):
                rate *= model.omega
            q[i, j] = rate * pi[j]

    np.fill_diagonal(q, -q.sum(axis=1))

    # Scale so expected substitution rate under pi is 1.
    rate = float(-np.sum(pi * np.diag(q)))
    if rate <= 0:
        raise ValueError("Invalid rate matrix scaling (non-positive rate)")
    q /= rate

    return q, codons, pi
