import random as _random
from dataclasses import dataclass as _dataclass
from typing import List as _List
from typing import Optional as _Optional
from typing import Tuple as _Tuple
from typing import Union as _Union

import torch as _torch
from torch import nn as _nn

from tamm.generation.token_generators import TokenGenerator as _TokenGenerator
from tamm.generation.token_generators.base import (
    TokenGeneratorOutput as _TokenGeneratorOutput,
)
from tamm.generation.token_generators.next_token_logic import (
    NextTokenLogic as _NextTokenLogic,
)
from tamm.generation.token_generators.utils import (
    KVCacheType,
    create_random_number_generator,
    find_eos_index,
    model_forward_step,
    pad_after_termination,
    post_process_token_ids,
    setup_kv_cache,
)
from tamm.layers.transformer.kv_cache import BaseKVCache as _BaseKVCache


@_dataclass
class SpeculativeDecodingTokenGeneratorOutput(_TokenGeneratorOutput):
    """
    Output generated from :py:class:`SpeculativeDecodingTokenGenerator`, consisting of token_ids,
    number of suggested tokens by the draft generator and number of accepted draft tokens for each
    sequence in the input batch.

    Attributes:
        token_ids (:py:class:`_torch.Tensor`): Tokens produced by the speculative decoding generator.
          Does not include ``input_ids``.
        suggested_token_counts (:py:class:`_torch.Tensor`): Total number of tokens suggested by the draft model.
          The tensor has shape ``(batch_size, N)``, where ``N`` is the number of times draft generator
          is called for speculative decoding.
        accepted_token_counts (:py:class:`_torch.Tensor`): Total number of tokens which were generated by
         the draft model and accepted as final tokens. The tensor has shape ``(batch_size, N)``,
         where ``N`` is the number of times draft generator is called for speculative decoding.
    """

    token_ids: _torch.Tensor
    suggested_token_counts: _torch.Tensor
    accepted_token_counts: _torch.Tensor


# pylint: disable=too-many-locals
class SpeculativeDecodingTokenGenerator(_TokenGenerator):
    """Token generator that uses speculative decoding strategy. Based on the paper:
    `Fast Inference from Transformers via Speculative Decoding <https://arxiv.org/pdf/2211.17192>`_

    Args:
        model (:obj:`nn.Module`): The target model instance from which tokens need to be sampled.
        next_token_logic (:obj:`_NextTokenLogic`): Logic specifying how target model's output logits should be
         used to generate the next token.
        draft_generator (:obj:`_TokenGenerator`): Token generator for producing draft tokens to be used for
          speculative decoding. draft_generator should provide an option called output_probabilities. When this
          option is set to ``True``, the draft generator should return probabilities used for generating the tokens,
          along with the generated tokens.
        max_draft_tokens (:obj:`int`): The maximum number of draft tokens generated in each round of speculative
          decoding. This argument is called ``gamma`` in the paper.
        cache_type (:obj:`KVCacheType`, optional): Type of KV cache to use. Defaults to ``KVCacheType.Speculative``.
          Only supported types are ``KVCacheType.Speculative`` and ``None``.
        cache_dtype (:obj:`torch.dtype`, optional): Data type for KV cache. Defaults to ``None``.
    """

    def __init__(
        self,
        model: _nn.Module,
        next_token_logic: _NextTokenLogic,
        draft_generator: _TokenGenerator,
        max_draft_tokens: int,
        cache_type: _Optional[KVCacheType] = KVCacheType.SPECULATIVE,
        cache_dtype: _Optional[_torch.dtype] = None,
    ):
        super().__init__(
            model=model,
            eos_id=draft_generator.eos_ids,
            pad_id=draft_generator.pad_id,
        )
        self._verify_draft_generator_compatibility(draft_generator)
        self._draft_generator = draft_generator
        self._next_token_logic = next_token_logic
        self._cache_type = cache_type
        self._cache_dtype = cache_dtype
        self._first_sd_call = True
        self.max_draft_tokens = max_draft_tokens

    def _verify_draft_generator_compatibility(self, draft_generator: _TokenGenerator):
        assert hasattr(draft_generator, "output_probabilities"), (
            "draft_generator should provide an option called output_probabilities to return probabilities used for "
            f"generation, to be compatible with {self.__class__.__name__}."
        )

    @_torch.inference_mode()
    def generate(
        self,
        input_ids: _torch.Tensor,
        max_new_tokens: _Union[int, _List[int]],
        seed: _Optional[int] = None,
    ) -> SpeculativeDecodingTokenGeneratorOutput:
        token_ids = input_ids.clone()

        bs, prompt_length = token_ids.shape
        device = token_ids.device

        prng = _random.Random(seed)

        if isinstance(max_new_tokens, list):
            tokens_until_max = _torch.tensor(max_new_tokens, device=device)
        else:
            assert isinstance(max_new_tokens, int)
            tokens_until_max = _torch.full((bs,), max_new_tokens, device=device)

        max_max_new_tokens = max_new_tokens  # max across all batch sequences
        if isinstance(max_new_tokens, list):
            max_max_new_tokens = max(max_new_tokens)

        decode_terminator = [x.to(device) for x in self.eos_ids]

        all_suggested = _torch.zeros((bs, 0), device=device)
        all_accepted = _torch.zeros((bs, 0), device=device)

        should_terminate = _torch.zeros((bs,), device=device, dtype=_torch.bool)

        target_cache = setup_kv_cache(
            model=self.model,
            batch_size=bs,
            cache_len=prompt_length + max_max_new_tokens,
            input_ids=input_ids,
            pad_id=self.pad_id,
            cache_type=self._cache_type,
            cache_dtype=self._cache_dtype,
        )

        self._first_sd_call = True

        while not should_terminate.all().item():
            # we set num_speculative_tokens based on number of expected tokens to reach max_new_tokens
            # we use max_remaining_tokens - 1 because last token comes from target model

            max_remaining_tokens = tokens_until_max.amax().item()
            num_speculative_tokens = min(
                self.max_draft_tokens, max_remaining_tokens - 1
            )

            # using no persistent draft cache, so we always pass in full input for next round
            output = self._speculative_decode_once(
                token_ids,
                num_speculative_tokens,
                target_cache,
                random_number_generator=prng,
            )

            token_ids = _torch.concat([token_ids, output.token_ids], dim=-1)

            non_pad_tokens = (output.token_ids != self.pad_id).sum(dim=-1)
            tokens_until_max -= non_pad_tokens

            # ignore counts if a sequence has terminated
            output.suggested_token_counts[should_terminate] = 0
            output.accepted_token_counts[should_terminate] = 0

            all_suggested = _torch.concat(
                [all_suggested, output.suggested_token_counts[:, None]],
                dim=-1,
            )
            all_accepted = _torch.concat(
                [all_accepted, output.accepted_token_counts[:, None]],
                dim=-1,
            )

            should_terminate_round = self._should_terminate_decoding(
                tokens_until_max,
                decode_terminator,
                output.token_ids,
            )

            should_terminate |= should_terminate_round

            if should_terminate.all().item():
                break

        token_ids = post_process_token_ids(
            token_ids[:, prompt_length:], max_new_tokens, self.pad_id
        )
        return SpeculativeDecodingTokenGeneratorOutput(
            token_ids=pad_after_termination(token_ids, self.pad_id, decode_terminator),
            suggested_token_counts=all_suggested.to(_torch.int32),
            accepted_token_counts=all_accepted.to(_torch.int32),
        )

    def _speculative_decode_once(
        self,
        token_ids: _torch.Tensor,
        num_speculative_tokens: int,
        target_cache: _Optional[_BaseKVCache],
        random_number_generator: _random.Random,
    ) -> SpeculativeDecodingTokenGeneratorOutput:
        """
        Generates num_speculative_tokens speculative tokens from draft generator and 1 token from
        the target model. Uses rejection sampling to accept n <= num_speculative_tokens speculative tokens.

        Args:
            token_ids (:obj:`_torch.Tensor`): Input token ids for the draft generator. Since we do not use
              a persistent cache for graft generation, we pass all the tokens generated so far to the draft
              generator.
            num_speculative_tokens (:obj:`int`): Number of tokens to generate from draft generator.
            target_cache: (:obj:`_BaseKVCache` or obj:`None`): KV cache for the target model.
            random_number_generator (:obj:`_random.Random`): A pseudo random number generator used to
              generating seeds used for sampling form draft generator and target model's probability distribution.
        """
        device = token_ids.device
        batch_size = token_ids.shape[0]

        target_input_ids = token_ids
        if target_cache and not self._first_sd_call:
            # New sample generated from target in last round needs to be fed to target
            target_input_ids = target_input_ids[:, -1:]

        if num_speculative_tokens == 0:
            target_logits = self._compute_target_logits(
                target_input_ids,
                target_prompt_len=token_ids.shape[-1],
                draft_token_count=0,
                target_cache=target_cache,
                with_temporary_cache=False,
            )
            final_token_ids = self._next_token_logic(
                target_logits[:, -1],
                rng=create_random_number_generator(
                    device=device, seed=random_number_generator.getrandbits(32)
                ),
                return_probabilities=False,
            ).token_ids

            return SpeculativeDecodingTokenGeneratorOutput(
                token_ids=final_token_ids,
                suggested_token_counts=_torch.zeros((batch_size,), device=device),
                accepted_token_counts=_torch.zeros((batch_size,), device=device),
            )

        # using no persistent cache, so always pass in full input to draft model;
        # at output, we get newly generated tokens and their logits
        self._draft_generator.output_probabilities = True  # type: ignore[attr-defined]

        draft_output = self._draft_generator.generate(
            token_ids,
            num_speculative_tokens,
            seed=random_number_generator.getrandbits(32),
        )
        draft_tokens = draft_output.token_ids
        draft_probs = draft_output.probabilities  # type: ignore[attr-defined]

        draft_token_count = draft_tokens.shape[1]
        # draft_token_count can be less than num_speculative_tokens if draft has early terminated in
        # all elements in the batch
        early_termination_amount = (draft_tokens == self.pad_id).sum(dim=-1)

        # append draft token ids to target token ids
        target_input_ids = _torch.cat([target_input_ids, draft_tokens], dim=-1)

        # on first target call, we treat the last target_input_id as if it was generated during the "last" target call.
        # because from second target call onwards, target_input_id always have an extra token generated
        # during the last call prepended, which is not present in the target cache but is fed to the target model
        target_prompt_len = token_ids.shape[-1] - 1

        target_logits = self._compute_target_logits(
            target_input_ids,
            target_prompt_len=target_prompt_len,
            draft_token_count=draft_token_count,
            target_cache=target_cache,
            with_temporary_cache=True,
        )

        (
            target_token,
            num_accepted_tokens,
        ) = self._verify_from_target(
            draft_tokens,
            draft_probs,
            target_logits,
            draft_token_count,
            random_number_generator,
        )  # evolve initial seed deterministically on every round

        max_accepted_tokens = _torch.max(num_accepted_tokens).item()

        # merge temporary cache into main cache, +1 because we want to merge cache values corresponding
        # to the token generated by target model from previous round
        if target_cache:
            target_cache.merge(
                merge_length=num_accepted_tokens + 1,
                max_merge_length=max_accepted_tokens + 1,
                previous_length=target_prompt_len,
            )

        accepted = draft_tokens.clone()
        # align sequences across batch dimension
        accepted[
            _torch.arange(draft_token_count, device=device)
            >= num_accepted_tokens[:, None]
        ] = self.pad_id
        accepted = accepted[:, :max_accepted_tokens]
        final_token_ids = _torch.concat([accepted, target_token], dim=-1)

        suggested_count = draft_token_count - early_termination_amount
        accepted_count = _torch.min(num_accepted_tokens, suggested_count)

        self._first_sd_call = False

        return SpeculativeDecodingTokenGeneratorOutput(
            token_ids=final_token_ids,
            suggested_token_counts=suggested_count,
            accepted_token_counts=accepted_count,
        )

    def _compute_target_logits(
        self,
        target_input_ids: _torch.Tensor,
        target_prompt_len: int,
        draft_token_count: int,
        target_cache: _Optional[_BaseKVCache],
        with_temporary_cache: bool,
    ) -> _torch.Tensor:
        if self._first_sd_call and target_cache:
            # pre-fill cache
            model_forward_step(
                model=self.model,
                input_ids=target_input_ids[:, :target_prompt_len],
                next_tokens=None,
                pad_id=self.pad_id,
                kv_cache=target_cache,
                prompt_len=target_prompt_len,
                token_idx=None,
            )
            # remove all pre-filled tokens
            target_input_ids = target_input_ids[:, target_prompt_len:]
        if target_cache and with_temporary_cache:
            target_cache.create_temporary_cache_view(
                speculation_length=draft_token_count + 1
            )
        output = model_forward_step(
            model=self.model,
            input_ids=target_input_ids,
            next_tokens=None,
            pad_id=self.pad_id,
            kv_cache=target_cache,
            prompt_len=target_prompt_len,
            token_idx=None,
        )
        # draft_token_count tokens from draft and one token from target
        # output shape: batch_size x draft_token_count + 1 x vocab_size
        return output.logits[
            :, -(draft_token_count + 1) :, : self.model.config.vocab_size
        ]

    def _verify_from_target(
        self,
        draft_tokens: _torch.Tensor,
        draft_probs: _torch.Tensor,
        target_logits: _torch.Tensor,
        draft_token_count: int,
        random_number_generator: _random.Random,
    ) -> _Tuple[_torch.Tensor, _torch.Tensor]:
        device = draft_tokens.device

        target_probs = self._next_token_logic.compute_probabilities(
            target_logits
        )  # shape: batch_size x draft_token_count + 1 x vocab_size

        # sample next token assuming all draft tokens are accepted
        # we do this to make selection of acceptance tokens without if/else loops
        # this token is always rejected, since we resample from new_probs later
        target_sample_seed = random_number_generator.getrandbits(32)
        target_token = self._sample_from_target_probs(
            target_probs[:, -1], seed=target_sample_seed
        )  # shape: batch_size x 1

        all_draft_tokens = _torch.concat(
            [draft_tokens[:, -draft_token_count:].clone(), target_token],
            dim=1,
        )  # shape: batch_size x draft_token_count + 1

        # add inf to the end of draft probs, this is to align draft prob shape
        # with the target_probs shape
        draft_probs = _torch.concat(
            [
                draft_probs,
                _torch.inf * _torch.ones_like(draft_probs[..., -1, :]).unsqueeze(1),
            ],
            dim=-2,
        )

        # get the acceptance probability for every logit
        # since draft_probs is inf at the last index, acceptance_prob is 0
        acceptance_prob = _torch.clamp(target_probs / draft_probs, max=1.0)

        # find the acceptance probability for the suggested logits
        acceptance_prob = _torch.gather(
            acceptance_prob, -1, all_draft_tokens[..., None]
        ).squeeze(
            -1
        )  # shape: batch_size x draft_token_count + 1

        # roll the die
        roll_seed = random_number_generator.getrandbits(32)
        roll = _torch.rand(
            acceptance_prob.shape,
            generator=create_random_number_generator(device=device, seed=roll_seed),
            device=acceptance_prob.device,
        )

        # since acceptance_prob is 0 at the last index, at least
        # one token (the one from target) would be rejected
        rejected = roll > acceptance_prob

        num_accepted_tokens = rejected.to(_torch.int32).argmax(dim=-1)

        # resample from adjusted target distribution
        # new_probs will be zero at the last index (corresponding to target token)
        new_probs = _torch.clamp(target_probs - draft_probs, min=0)
        new_probs /= _torch.sum(new_probs, keepdim=True, dim=-1)
        # when all draft tokens are accepted, since we set new_probs to be same as
        # target_probs at the last index, we end up drawing from the target_probs, which
        # is what we want
        new_probs[:, -1, :] = target_probs[:, -1, :]

        new_probs = new_probs[
            _torch.arange(new_probs.shape[0], device=device), num_accepted_tokens, :
        ]

        target_token = self._sample_from_target_probs(
            new_probs, seed=target_sample_seed
        )

        return target_token, num_accepted_tokens

    @staticmethod
    def _should_terminate_decoding(
        tokens_until_max: _torch.Tensor,
        decode_terminator: _List[_torch.Tensor],
        token_ids: _torch.Tensor,
    ) -> _torch.Tensor:
        max_gen_termination = tokens_until_max <= 0

        termination_indices = find_eos_index(token_ids, decode_terminator)
        eos_reached = termination_indices < _torch.iinfo(_torch.int32).max

        should_terminate = _torch.logical_or(eos_reached, max_gen_termination)

        return should_terminate

    @staticmethod
    def _sample_from_target_probs(probs: _torch.Tensor, seed: int) -> _torch.Tensor:
        rng = create_random_number_generator(device=probs.device, seed=seed)
        return _torch.multinomial(probs, num_samples=1, generator=rng)
