# SPDX-License-Identifier: Apache-2.0

import os
from typing import Optional, List
import torch
import torch.jit
import torch.nn.functional as F

from vllm.model_executor.layers.spec_decode_base_sampler import (
    SpecDecodeDeterministicBaseSampler)
from vllm.logger import init_logger

logger = init_logger(__name__)


class TypicalAcceptanceSampler(SpecDecodeDeterministicBaseSampler):
    """Apply typical acceptance sampling as described in section 3.3.1 in 
        "MEDUSA: Simple LLM Inference Acceleration Framework with 
        Multiple Decoding Heads"
        https://arxiv.org/pdf/2401.10774
    """

    def __init__(
        self,
        posterior_threshold: float,
        posterior_alpha: float,
        strict_mode: bool = False,
    ):
        """Create a Typical Acceptance Sampler.

        Args:
            strict_mode: Whether or not to perform shape/device/dtype checks
            during sampling. This catches correctness issues but adds
            nontrivial latency.
            posterior_threshold : A threshold value that sets a lower bound 
            on the posterior probability of a token in target model for it
            to be accepted.
            posterior_alpha : A scaling factor for the entropy-based
            threshold in typical acceptance sampling.
        """
        self._posterior_threshold = posterior_threshold
        self._posterior_alpha = posterior_alpha
        super().__init__(strict_mode=strict_mode)

        self.tree_decoding = (os.environ.get('VLLM_TREE_DECODING') == '1')

    def forward(
        self,
        target_with_bonus_probs: torch.Tensor,
        bonus_token_ids: torch.Tensor,
        draft_probs: torch.Tensor,
        draft_token_ids: torch.Tensor,
        cart_candidates: Optional[torch.Tensor] = None,
        best_candidates: Optional[List] = None,
        accept_lengths: Optional[List] = None,
        first_step_flags: Optional[List] = None,
    ) -> torch.Tensor:
        """Sample token ids using typical acceptance sampling. This accepts 
        or rejects tokens proposed by the draft model using the probability
        of each token according to the draft and target models.

        In the worst case where all draft tokens are rejected, it is guaranteed
        one token will be emitted.

        In the case where all draft tokens are accepted, the bonus token will be
        accepted.

        Args:
            target_probs: The probability distribution over token ids given
                context according to the target model.
            shape = [batch_size, num_speculative_tokens, vocab_size]

            bonus_token_ids: The "bonus" token ids that are accepted iff all
                speculative tokens in a sequence are accepted.
            shape = [batch_size, num_bonus_tokens]

            draft_probs: This parameter is unused by the acceptance sampler.

            draft_token_ids: The token ids that were sampled from the draft
                probabilities.
            shape = [batch_size, num_speculative_tokens]

            cart_candidates: tree-style cartesian candidates
            best_candidates: pending to write best candidates index 
            accept_lengths: pending to write accept lengths 
            first_step_flags: whether this is the first decoding step

        Returns:
            output_token_ids: The token ids sampled via rejection sampling,
                or -1 if unable to sample a token because the previous token
                was rejected.
            shape = [batch_size, num_speculative_tokens + num_bonus_tokens]
        """
        # Only perform shape/dtype/device checking in strict mode, as it adds
        # overhead.
        if self._strict_mode:
            self._raise_if_incorrect_input(target_with_bonus_probs,
                                           draft_token_ids, bonus_token_ids)
            
        if not self.tree_decoding:
            target_probs = target_with_bonus_probs[:, :-1]
            accepted = self._evaluate_accepted_tokens(target_probs,
                                                    draft_token_ids)
            recovered_token_ids = self._get_recovered_token_ids(target_probs)
            output_token_ids = self._create_output(accepted, recovered_token_ids,
                                                draft_token_ids,
                                                bonus_token_ids)
        else:
            assert cart_candidates is not None
            target_probs = target_with_bonus_probs
            output_token_ids = self._evaluate_accepted_tokens_tree_style(target_probs,
                                                        draft_token_ids,
                                                        cart_candidates,
                                                        best_candidates,
                                                        accept_lengths,
                                                        first_step_flags)
        return output_token_ids
    
    def _evaluate_accepted_tokens_tree_style(self, target_probs, draft_token_ids, 
                                             cart_candidates, output_best_candidates,
                                             accept_lengths, first_step_flags):
        r"""
        Evaluates and returns a mask of accepted tokens based on the
        posterior probabilities.

        Parameters:
        ----------
        target_probs : torch.Tensor
            A tensor of shape (batch_size, k, vocab_size) representing 
            the probabilities of each token in the vocabulary for each
            position in the proposed sequence. This is the distribution
            generated by the target model.
        draft_token_ids : torch.Tensor
            A tensor of shape (batch_size, k) representing the proposed
            token ids.
        cart_candidates : torch.Tensor
            A tensor of shape (batch_size, retrieve_size, tree_depth)
            representing the cart candidates of tree proposals.

        A draft token_id x_{n+k} is accepted if it satisfies the
        following condition
    
        .. math::
            p_{\text{original}}(x_{n+k} | x_1, x_2, \dots, x_{n+k-1}) > 
            \min \left( \epsilon, \delta * \exp \left(
                -H(p_{\text{original}}(
                    \cdot | x_1, x_2, \ldots, x_{n+k-1})) \right) \right)
        
        where :math:`p_{\text{original}}` corresponds to target_probs 
        and :math:`\epsilon` and :math:`\delta` correspond to hyperparameters
        specified using self._posterior_threshold and self._posterior_alpha

        This method computes the posterior probabilities for the given
        draft token ids based on the provided target probabilities. It
        calculates the entropy of the posterior distribution and determines
        a dynamic threshold for each token position using the provided
        posterior_threshold and posterior_alpha values. The method then
        returns a boolean mask indicating which tokens can be accepted.

        Returns:
        -------
        torch.Tensor
            A boolean tensor of shape (batch_size, k) where each element
            indicates whether the corresponding draft token has been accepted
            or rejected. True indicates acceptance and false indicates
            rejection.
            
        """
        target_probs = target_probs[:, :, :-1]
        device = target_probs.device
        batch_size = cart_candidates.shape[0]
        candidates_prob = torch.gather(
            target_probs, dim=-1, index=cart_candidates[:, :, 1:].unsqueeze(-1)
        ).squeeze(-1) # [batch_size, retrieve_size, max_depth]
        posterior_entropy = -torch.sum(
            target_probs * torch.log(target_probs + 1e-5), dim=-1
        )  # torch.sum(torch.log(*)) is faster than torch.prod  [batch_size, retrieve_size, max_depth]
        threshold = torch.minimum(
            torch.ones_like(posterior_entropy) * self._posterior_threshold,
            torch.exp(-posterior_entropy) * self._posterior_alpha,
        )
        posterior_mask = candidates_prob > threshold # [batch_size, retrieve_size, max_depth]
        candidates_accept_length = (torch.cumprod(posterior_mask, dim=2)).sum(dim=-1) # [batch_size, retrieve_size]

        # Choose the best candidate based on the evaluated posterior probabilities
        accept_length, _ = candidates_accept_length.max(dim=-1) # [batch_size]
        if torch.any(accept_length > 0):
            valid_index = (candidates_accept_length == accept_length.unsqueeze(-1)).unsqueeze(-1) # [batch_size, retrieve_size, 1]

            candidates_prob = candidates_prob * valid_index # [batch_size, retrieve_size, max_depth]
            valid_index = torch.arange(candidates_prob.shape[-1], device=device).unsqueeze(0).unsqueeze(0).repeat(
                batch_size, candidates_prob.shape[1], 1) # [batch_size, retrieve_size, max_depth]
            valid_index = (valid_index < accept_length.unsqueeze(1).unsqueeze(2).repeat(1, candidates_prob.shape[1], 1)) # [batch_size, retrieve_size, 1]
            candidates_prob = candidates_prob*valid_index # [batch_size, retrieve_size, max_depth]

            # add 1e-3 to avoid zero value
            likelihood = torch.sum(torch.log(candidates_prob + 1e-3), dim=-1) # [batch_size, retrieve_size]

            best_candidate = torch.argmax(likelihood, dim=-1) # [batch_size]
        else:
            # Choose the best candidate
            best_candidate = torch.zeros((batch_size), dtype=torch.long, device=device) # [batch_size]

        k = draft_token_ids.shape[-1]
        output_token_id_list = []

        accept_length_list = accept_length.cpu().tolist()
        #logger.info("accept_length:%s", accept_length_list)
        for i in range(batch_size):
            output_best_candidates.append(best_candidate[i])
            accept_lengths.append(accept_length_list[i])

            if not first_step_flags[i]:
                select_indices = cart_candidates[i, best_candidate[i], : accept_length[i] + 1]
                select_indices = F.pad(select_indices, (0, k - 1 - accept_length[i]), 'constant', -1)
            else:
                select_indices = cart_candidates[i, best_candidate[i], 1 : accept_length[i] + 1]
                select_indices = F.pad(select_indices, (0, k - accept_length[i]), 'constant', -1)
            output_token_id_list.append(select_indices)

        return torch.stack(output_token_id_list, dim=0)


    def _evaluate_accepted_tokens(self, target_probs, draft_token_ids):
        r"""
        Evaluates and returns a mask of accepted tokens based on the
        posterior probabilities.

        Parameters:
        ----------
        target_probs : torch.Tensor
            A tensor of shape (batch_size, k, vocab_size) representing 
            the probabilities of each token in the vocabulary for each
            position in the proposed sequence. This is the distribution
            generated by the target model.
        draft_token_ids : torch.Tensor
            A tensor of shape (batch_size, k) representing the proposed
            token ids.

        A draft token_id x_{n+k} is accepted if it satisfies the
        following condition
    
        .. math::
            p_{\text{original}}(x_{n+k} | x_1, x_2, \dots, x_{n+k-1}) > 
            \min \left( \epsilon, \delta * \exp \left(
                -H(p_{\text{original}}(
                    \cdot | x_1, x_2, \ldots, x_{n+k-1})) \right) \right)
        
        where :math:`p_{\text{original}}` corresponds to target_probs 
        and :math:`\epsilon` and :math:`\delta` correspond to hyperparameters
        specified using self._posterior_threshold and self._posterior_alpha

        This method computes the posterior probabilities for the given
        draft token ids based on the provided target probabilities. It
        calculates the entropy of the posterior distribution and determines
        a dynamic threshold for each token position using the provided
        posterior_threshold and posterior_alpha values. The method then
        returns a boolean mask indicating which tokens can be accepted.

        Returns:
        -------
        torch.Tensor
            A boolean tensor of shape (batch_size, k) where each element
            indicates whether the corresponding draft token has been accepted
            or rejected. True indicates acceptance and false indicates
            rejection.
            
        """
        device = target_probs.device
        candidates_prob = torch.gather(
            target_probs, dim=-1,
            index=draft_token_ids.unsqueeze(-1)).squeeze(-1)
        # A small constant added to prevent computing the logarithm of zero,
        # which can lead to undefined values.
        epsilon = 1e-5
        posterior_entropy = -torch.sum(
            target_probs * torch.log(target_probs + epsilon), dim=-1)
        threshold = torch.minimum(
            torch.ones_like(posterior_entropy, device=device) *
            self._posterior_threshold,
            torch.exp(-posterior_entropy) * self._posterior_alpha,
        )
        accepted_mask = candidates_prob > threshold
        return accepted_mask

    def _get_recovered_token_ids(self, target_probs):
        """
        The recovered token ids will fill the first unmatched token
        by the target token.

        Parameters
        ----------
        target_probs : torch.Tensor
            A tensor of shape (batch_size, k, vocab_size) containing 
            the target probability distribution

        Returns
        -------
        torch.Tensor
            A tensor of shape (batch_size, k) with the recovered token
            ids which are selected from target probs.
        """
        max_indices = torch.argmax(target_probs, dim=-1)

        return max_indices
