"""Tree-based codon simulation under the configured null model."""

from __future__ import annotations

from dataclasses import dataclass
from functools import lru_cache

import numpy as np
from scipy.linalg import expm

from babappa.io.fasta import CodonAlignment
from babappa.io.newick import PhyloTree, TreeNode
from babappa.null.model import CodonNullModel, codon_rate_matrix


@dataclass
class SimulatorContext:
    q: np.ndarray
    codons: list[str]
    pi: np.ndarray
    codon_to_idx: dict[str, int]


@lru_cache(maxsize=512)
def _transition_matrix_cached(q_bytes: bytes, shape0: int, shape1: int, t: float) -> np.ndarray:
    q = np.frombuffer(q_bytes, dtype=float).reshape((shape0, shape1))
    return expm(q * t)


def _transition_matrix(q: np.ndarray, branch_length: float, rate_multiplier: float = 1.0) -> np.ndarray:
    t = float(branch_length) * float(rate_multiplier)
    if t == 0:
        return np.eye(q.shape[0], dtype=float)
    return _transition_matrix_cached(q.tobytes(), q.shape[0], q.shape[1], t)


def _sample_state(prob: np.ndarray, rng: np.random.Generator) -> int:
    return int(rng.choice(len(prob), p=prob))


def _gamma_rate(alpha: float, rng: np.random.Generator) -> float:
    # Mean-one gamma with shape alpha and scale 1/alpha.
    if alpha <= 0:
        return 1.0
    return float(rng.gamma(shape=alpha, scale=1.0 / alpha))


def _simulate_site_recursive(
    node: TreeNode,
    parent_state: int,
    ctx: SimulatorContext,
    rng: np.random.Generator,
    rate_multiplier: float,
    out: dict[str, int],
) -> None:
    for child in node.children:
        p = _transition_matrix(ctx.q, child.length, rate_multiplier=rate_multiplier)
        child_state = _sample_state(p[parent_state], rng)
        if child.is_leaf():
            if child.name is None:
                raise ValueError("Leaf without taxon name")
            out[child.name] = child_state
        else:
            _simulate_site_recursive(
                child,
                child_state,
                ctx=ctx,
                rng=rng,
                rate_multiplier=rate_multiplier,
                out=out,
            )


def build_simulator_context(model: CodonNullModel) -> SimulatorContext:
    q, codons, pi = codon_rate_matrix(model)
    return SimulatorContext(q=q, codons=codons, pi=pi, codon_to_idx={c: i for i, c in enumerate(codons)})


def simulate_alignment(
    tree: PhyloTree,
    model: CodonNullModel,
    n_codons: int,
    gene_name: str,
    seed: int,
) -> CodonAlignment:
    rng = np.random.default_rng(seed)
    ctx = build_simulator_context(model)
    taxa = tree.leaf_names()
    seq_chunks: dict[str, list[str]] = {t: [] for t in taxa}

    for _ in range(n_codons):
        root_state = _sample_state(ctx.pi, rng)
        rate_multiplier = _gamma_rate(model.gamma_alpha, rng) if model.rate_variation else 1.0
        state_by_taxon: dict[str, int] = {}

        if tree.root.is_leaf():
            if tree.root.name is None:
                raise ValueError("Root leaf has no name")
            state_by_taxon[tree.root.name] = root_state
        else:
            _simulate_site_recursive(
                tree.root,
                root_state,
                ctx=ctx,
                rng=rng,
                rate_multiplier=rate_multiplier,
                out=state_by_taxon,
            )

        for taxon in taxa:
            state = state_by_taxon.get(taxon, root_state)
            seq_chunks[taxon].append(ctx.codons[state])

    seqs = {t: "".join(codons) for t, codons in seq_chunks.items()}
    return CodonAlignment(gene=gene_name, sequences=seqs)


def simulate_replicates(
    tree: PhyloTree,
    model: CodonNullModel,
    n_codons: int,
    n_reps: int,
    base_seed: int,
    gene_prefix: str,
) -> list[CodonAlignment]:
    out: list[CodonAlignment] = []
    rng = np.random.default_rng(base_seed)
    seeds = [int(x) for x in rng.integers(1, 2**31 - 1, size=n_reps)]
    for i, s in enumerate(seeds):
        out.append(simulate_alignment(tree, model, n_codons=n_codons, gene_name=f"{gene_prefix}_{i+1}", seed=s))
    return out
