from __future__ import annotations

import math
from collections import Counter, defaultdict
from dataclasses import dataclass, field
from typing import Iterable, Sequence

import numpy as np

NEG_INF = float("-inf")


def logadd(*values: float) -> float:
    finite = [value for value in values if value != NEG_INF]
    if not finite:
        return NEG_INF
    maximum = max(finite)
    return maximum + math.log(sum(math.exp(value - maximum) for value in finite))


def ctc_collapse(path: Sequence[int], blank_id: int = 0) -> tuple[int, ...]:
    output = []
    previous = None
    for token in path:
        token = int(token)
        if token != blank_id and token != previous:
            output.append(token)
        previous = token
    return tuple(output)


@dataclass
class CharNGramLM:
    order: int
    vocab_size: int
    smoothing: float = 0.1
    _context_counts: dict[tuple[int, ...], Counter] = field(default_factory=dict, init=False)
    _context_totals: Counter = field(default_factory=Counter, init=False)

    def __post_init__(self) -> None:
        if self.order < 1:
            raise ValueError("order must be at least 1")
        if self.vocab_size < 1:
            raise ValueError("vocab_size must be positive")
        if self.smoothing <= 0:
            raise ValueError("smoothing must be positive")

    def fit(self, sequences: Iterable[Sequence[int]]) -> "CharNGramLM":
        context_counts: dict[tuple[int, ...], Counter] = defaultdict(Counter)
        context_totals = Counter()
        history_size = self.order - 1
        for sequence in sequences:
            history: tuple[int, ...] = ()
            for raw_token in sequence:
                token = int(raw_token)
                context = history[-history_size:] if history_size else ()
                context_counts[context][token] += 1
                context_totals[context] += 1
                history = (*history, token)
        self._context_counts = dict(context_counts)
        self._context_totals = context_totals
        return self

    def log_prob(self, history: Sequence[int], token: int) -> float:
        history_size = self.order - 1
        context = tuple(history[-history_size:]) if history_size else ()
        # Back off to shorter contexts when an n-gram history was unseen.
        while context and context not in self._context_totals:
            context = context[1:]
        counts = self._context_counts.get(context, Counter())
        total = self._context_totals.get(context, 0)
        numerator = counts.get(int(token), 0) + self.smoothing
        denominator = total + self.smoothing * self.vocab_size
        return math.log(numerator / denominator)


def prefix_beam_search(
    log_probs: np.ndarray,
    *,
    beam_width: int = 10,
    blank_id: int = 0,
    token_topk: int | None = None,
    lm: CharNGramLM | None = None,
    lm_weight: float = 0.0,
    token_bonus: float = 0.0,
) -> tuple[int, ...]:
    """Decode one [time, classes] CTC log-probability matrix.

    Acoustic paths are summed exactly within the retained prefix beam. The optional
    character LM and token bonus affect beam ranking but never alter CTC path sums.
    """
    values = np.asarray(log_probs)
    if values.ndim != 2:
        raise ValueError("log_probs must have shape [time, classes]")
    if beam_width < 1:
        raise ValueError("beam_width must be positive")
    if not 0 <= blank_id < values.shape[1]:
        raise ValueError("blank_id is outside the class dimension")
    if token_topk is not None and token_topk < 1:
        raise ValueError("token_topk must be positive when provided")

    beams: dict[tuple[int, ...], tuple[float, float]] = {(): (0.0, NEG_INF)}
    lm_scores: dict[tuple[int, ...], float] = {(): 0.0}

    def rank(prefix: tuple[int, ...], scores: tuple[float, float]) -> float:
        return logadd(*scores) + lm_weight * lm_scores[prefix] + token_bonus * len(prefix)

    for frame in values:
        if token_topk is None or token_topk >= len(frame):
            tokens = list(range(len(frame)))
        else:
            top = np.argpartition(frame, -token_topk)[-token_topk:].tolist()
            tokens = top if blank_id in top else [blank_id, *top]

        next_beams: dict[tuple[int, ...], tuple[float, float]] = {}

        def update(prefix: tuple[int, ...], blank_score: float = NEG_INF,
                   nonblank_score: float = NEG_INF) -> None:
            old_blank, old_nonblank = next_beams.get(prefix, (NEG_INF, NEG_INF))
            next_beams[prefix] = (
                logadd(old_blank, blank_score),
                logadd(old_nonblank, nonblank_score),
            )

        for prefix, (p_blank, p_nonblank) in beams.items():
            total = logadd(p_blank, p_nonblank)
            blank_logp = float(frame[blank_id])
            update(prefix, blank_score=total + blank_logp)

            for token in tokens:
                token = int(token)
                if token == blank_id:
                    continue
                token_logp = float(frame[token])
                if prefix and token == prefix[-1]:
                    # Repeating without a blank collapses back to the same prefix.
                    update(prefix, nonblank_score=p_nonblank + token_logp)
                    # Repeating after a blank creates a second character.
                    extended = (*prefix, token)
                    update(extended, nonblank_score=p_blank + token_logp)
                else:
                    extended = (*prefix, token)
                    update(extended, nonblank_score=total + token_logp)

                if extended not in lm_scores:
                    prior = lm_scores[prefix]
                    lm_scores[extended] = prior + (lm.log_prob(prefix, token) if lm is not None else 0.0)

        ranked = sorted(next_beams.items(), key=lambda item: rank(item[0], item[1]), reverse=True)
        beams = dict(ranked[:beam_width])

    return max(beams.items(), key=lambda item: rank(item[0], item[1]))[0]