"""This module provides BeamSampling class for generating completions using beam search sampling."""

from math import exp, log2
from random import random
from typing import Any, Callable, Iterable, Tuple

from sign_language_translator.models.language_models.abstract_language_model import (
    LanguageModel,
)
from sign_language_translator.utils import sample_one_index


class BeamSampling:
    """BeamSampling class for generating completions using beam search sampling.

    Args:
        model (LanguageModel): The language model used for generating completions.
        beam_width (int, optional): The beam width for beam search. Defaults to 3.
        start_of_sequence_token (str, optional): The start of sequence token. Defaults to "[".
        end_of_sequence_token (str, optional): The end of sequence token. Defaults to "]".
        max_length (int, optional): The maximum length of the generated completions. Defaults to 33.
        scoring_function (Callable[[Iterable, float], float], optional):
            The scoring function used to score the completions. It should accept the generated
            sequence and its overall log probability as arguments. Defaults to a linear function.

        return_log_of_probability (bool, optional):
            A flag indicating whether to return the probability of the completions
            or log2 of it. Defaults to True.
    """

    def __init__(
        self,
        model: LanguageModel,
        beam_width: int = 3,
        start_of_sequence_token="[",
        end_of_sequence_token="]",
        max_length: int = 37,
        scoring_function: Callable[[Iterable, float], float] = (
            lambda seq, log_prob: 10.0 + log_prob / len(seq)  # type: ignore
        ),
        return_log_of_probability: bool = True,
    ) -> None:
        self.model = model
        self.start_of_sequence_token = start_of_sequence_token
        self.end_of_sequence_token = end_of_sequence_token
        self.beam_width = beam_width
        self.max_length = max_length
        self.scoring_function = scoring_function
        self.return_log_of_probability = return_log_of_probability

    def __call__(self, context: Iterable | None = None) -> Iterable:
        return self.complete(initial_context=context)

    def complete(
        self,
        initial_context: Iterable | None = None,
        append_func: Callable[[Any, Any], Any] = lambda context, token: (
            context + [token]
            if isinstance(context, list)
            else context + (token,)
            if isinstance(context, tuple)
            else context + token
        ),
    ) -> Tuple[Iterable, float]:
        """Generate completions based on the given initial context.

        Args:
            initial_context (Iterable | None, optional): The initial context for completion generation. Defaults to None.
            append_func (Callable[[Any, Any], Any], optional): a function that can append the generated next token to provided context.
                Defaults to a lambda function that can append to list, tuple & str.

        Returns:
            Tuple[Iterable, float]: One generated completion and its score.
        """

        if initial_context is None:
            initial_context = [self.start_of_sequence_token]

        branches = [(initial_context, 0.0)]

        for _ in range(self.max_length):
            n_branches = round(self.beam_width + random() * 0.8 - 0.4)

            # Expand
            new_branches = []
            for context_, score_ in branches:
                if (  # sequence completed
                    context_[-1] == self.end_of_sequence_token  # type: ignore
                    or len(context_) >= self.max_length  # type: ignore
                ):
                    # just as it was. no change.
                    new_branches.append((context_, score_))
                    continue

                # append next tokens
                for _ in range(n_branches):
                    next_token, prob = self.model.next(context_)
                    if next_token == self.model.unknown_token:
                        score = score_
                        next_context = context_
                    else:
                        next_context = append_func(context_, next_token)
                        score = score_ + log2(prob)

                    # no repeats; for diversity.
                    if (next_context, score) not in new_branches:
                        new_branches.append((next_context, score))

            # Prune
            new_branches = sorted(
                new_branches,
                key=lambda item: self.scoring_function(*item),
                reverse=True,
            )[:n_branches]

            if branches == new_branches:
                # no branch has grown further: stop
                break

            branches = new_branches

        # softmax: turn each branch's scores into a probability distribution
        weights = [self.scoring_function(seq, log_prob) for seq, log_prob in branches]
        weights = [exp(w) for w in weights]
        weights = [w / sum(weights) for w in weights]

        # select one branch
        selected_completion, score = branches[sample_one_index(weights)]

        # reformat score
        if not self.return_log_of_probability:
            score = 2 ** score # math.exp2(score)

        return selected_completion, score
