# coding=utf-8
# Copyright 2020 The HuggingFace Inc. team
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import warnings
from typing import List, Optional, Union, TYPE_CHECKING

import torch
import torch.distributed as dist
from torch import nn

from transformers.generation.configuration_utils import GenerationConfig

from .sbs_helpers.gumbel import gumbel, gumbel_with_maximum
from .sbs_helpers.logits_process import LogitsProcessorList, LogitsProcessor
from transformers.generation.stopping_criteria import StoppingCriteriaList, validate_stopping_criteria
from .beam_constrained_decoder import ConstrainedDecoder
from .contrastive_decoder import ContrastiveDecoder
from .utils import GenerationStrategy, BeamSearchDecoderOnlyOutput, BeamSearchEncoderDecoderOutput, \
    BeamSearchOutput
from .beam_utils import BeamSearchScorer

if TYPE_CHECKING:
    from transformers.modeling_utils import PreTrainedModel, GenerationMixin
    from transformers.generation.streamers import BaseStreamer


class OldSBSLogitProcessor(LogitsProcessor):

    def __init__(self, num_beams, batch_size=2):
        self.num_beams = num_beams
        self.batch_size = batch_size

    def __call__(self, input_ids, scores, **kwargs):
        """

        :param input_ids: shape (batch_size * num_beams, seq_len)
        :param scores: shape (batch_size * num_beams, vocab_size)
        :param kwargs.beam_log_probs: shape (batch_size * num_beams,)
        :param kwargs.past_processed_scores: tuple of size seq_len with tensors (batch_size * num_beams, vocab_size)
        :return: shape (batch_size, vocab_size)
        """

        assert kwargs.get('beam_log_probs', None) is not None, "beam_log_probs must be provided"
        assert kwargs.get('past_scores', None) is not None, "past_scores is None. Add output_" \
                                                                      "scores=True to the generate function arguments"
        assert kwargs.get('beam_indices', None) is not None, "beam_indices must be provided"
        beam_log_probs = kwargs['beam_log_probs']  # shape (batch_size * num_beams,)
        beam_log_probs = beam_log_probs.view(-1,1)  # shape (batch_size * num_beams,)
        past_scores = kwargs['past_scores']  # tuple: seq_len tensors of dim (batch_size * num_beams, vocab_size)
        beam_indices = kwargs['beam_indices']  # shape (batch_size * num_beams,)

        device = input_ids.device

        scores = beam_log_probs.view(-1, 1) + scores  # shape (batch_size * num_beams, vocab_size)
        # scores = scores.clamp(min=-1e9)

        if len(past_scores) == 0:  # first token
            last_gumbels = gumbel(size=(self.batch_size,)).to(device)  # shape (batch_size, )
            # expand (batch_size, ) to (batch_size * num_beams, )
            last_gumbels = last_gumbels.repeat_interleave(self.num_beams, dim=0)
        else:
            last_token_scores = past_scores[-1]  # shape (batch_size * num_beams, vocab_size)
            last_tokens = input_ids[:, -1]  # shape (batch_size * num_beams, )
            last_beam_indices = torch.tensor(tuple(tup[-1] for tup in beam_indices), device=device)
            last_gumbels = last_token_scores[last_beam_indices, last_tokens]
        new_gumbels, _ = gumbel_with_maximum(scores, last_gumbels)  # shape (batch_size * num_beams, ) todo check this view
        return new_gumbels


class OldStochasticBeamSearchDecoder(GenerationStrategy):
    """
    Implements beam search decoding method.
    """

    def __init__(self, config: GenerationConfig = None, check_config: bool = True):
        super().__init__(config, check_config=check_config)
        self.scorer = None

    @classmethod
    def param_check(cls, config: GenerationConfig):
        if config.num_beams <= 1:
            raise ValueError(
                f"{cls.__name__} cannot be used with num_beams <= 1, but num_beams is {config.num_beams}."
            )
        if config.num_beam_groups != 1:
            raise ValueError(
                f"{cls.__name__} cannot be used with num_beam_groups != 1, "
                f"but num_beam_groups is {config.num_beam_groups}."
            )
        if config.do_sample:
            raise ValueError(
                f"{cls.__name__} cannot be used with sampling."
            )
        if ConstrainedDecoder.param_test(config):
            raise ValueError(
                f"{cls.__name__} cannot be used with constraints."
            )
        if ContrastiveDecoder.param_test(config):
            raise ValueError(
                f"{cls.__name__} cannot be used with contrastive penalties."
            )

    def __call__(self,
                 model: Union["PreTrainedModel", "GenerationMixin"],
                 input_ids: torch.LongTensor,
                 assistant_model: Optional["PreTrainedModel"] = None,
                 top_k: Optional[int] = None,
                 top_p: Optional[float] = None,
                 typical_p: Optional[float] = None,
                 penalty_alpha: Optional[float] = None,
                 do_sample: bool = False,
                 logits_processor: Optional[LogitsProcessorList] = None,
                 logits_warper: Optional[LogitsProcessorList] = None,
                 max_length: Optional[int] = None,
                 stopping_criteria: Optional[StoppingCriteriaList] = None,
                 pad_token_id: Optional[int] = None,
                 eos_token_id: Optional[Union[int, List[int]]] = None,
                 output_attentions: Optional[bool] = None,
                 output_hidden_states: Optional[bool] = None,
                 output_scores: Optional[bool] = True,
                 return_dict_in_generate: Optional[bool] = True,
                 synced_gpus: bool = False,
                 streamer: Optional["BaseStreamer"] = None,
                 sequential: Optional[bool] = None,
                 beam_scores_type: Optional[torch.dtype] = torch.float,
                 **model_kwargs,
                 ) -> Union[BeamSearchOutput, torch.LongTensor]:
        r"""
        Generates sequences of token ids for models with a language modeling head using **beam search decoding** and
        can be used for text-decoder, text-to-text, speech-to-text, and vision-to-text models.

        <Tip warning={true}>

        In most cases, you do not need to call [`~generation.GenerationMixin.beam_search`] directly. Use generate()
        instead. For an overview of generation strategies and code examples, check the [following
        guide](../generation_strategies).

        </Tip>

        Parameters:
            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
                The sequence used as a prompt for the generation.
            beam_scorer (`transformers.generation.BeamScorer`):
                An derived instance of [`BeamScorer`] that defines how beam hypotheses are constructed, stored and
                sorted during generation. For more information, the documentation of [`BeamScorer`] should be read.
            logits_processor (`LogitsProcessorList`, *optional*):
                An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`]
                used to modify the prediction scores of the language modeling head applied at each generation step.
            stopping_criteria (`StoppingCriteriaList`, *optional*):
                An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`]
                used to tell if the generation loop should stop.
            max_length (`int`, *optional*, defaults to 20):
                **DEPRECATED**. Use `logits_processor` or `stopping_criteria` directly to cap the number of generated
                tokens. The maximum length of the sequence to be generated.
            pad_token_id (`int`, *optional*):
                The id of the *padding* token.
            eos_token_id (`Union[int, List[int]]`, *optional*):
                The id of the *end-of-sequence* token. Optionally, use a list to set multiple *end-of-sequence* tokens.
            output_attentions (`bool`, *optional*, defaults to `False`):
                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
                returned tensors for more details.
            output_hidden_states (`bool`, *optional*, defaults to `False`):
                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
                for more details.
            output_scores (`bool`, *optional*, defaults to `False`):
                Whether or not to return the prediction scores. See `scores` under returned tensors for more details.
            return_dict_in_generate (`bool`, *optional*, defaults to `False`):
                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
            synced_gpus (`bool`, *optional*, defaults to `False`):
                Whether to continue running the while loop until max_length (needed for ZeRO stage 3)
            model_kwargs:
                Additional model specific kwargs will be forwarded to the `forward` function of the model. If model is
                an encoder-decoder model the kwargs should include `encoder_outputs`.

        Return:
            [`generation.BeamSearchDecoderOnlyOutput`], [`~generation.BeamSearchEncoderDecoderOutput`] or
            `torch.LongTensor`: A `torch.LongTensor` containing the generated tokens (default behaviour) or a
            [`~generation.BeamSearchDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and
            `return_dict_in_generate=True` or a [`~generation.BeamSearchEncoderDecoderOutput`] if
            `model.config.is_encoder_decoder=True`.


        Examples:

        ```python
        >>> from transformers.generation import BeamSearchScorer        >>> from transformers import (
        ...     AutoTokenizer,
        ...     AutoModelForSeq2SeqLM,
        ...     LogitsProcessorList,
        ...     MinLengthLogitsProcessor,
        ...     )
        >>> import torch

        >>> tokenizer = AutoTokenizer.from_pretrained("t5-base")
        >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")

        >>> encoder_input_str = "translate English to German: How old are you?"
        >>> encoder_input_ids = tokenizer(encoder_input_str, return_tensors="pt").input_ids


        >>> # lets run beam search using 3 beams
        >>> num_beams = 3
        >>> # define decoder start token ids
        >>> input_ids = torch.ones((num_beams, 1), device=model.device, dtype=torch.long)
        >>> input_ids = input_ids * model.config.decoder_start_token_id

        >>> # add encoder_outputs to model keyword arguments
        >>> model_kwargs = {
        ...     "encoder_outputs": model.get_encoder()(
        ...         encoder_input_ids.repeat_interleave(num_beams, dim=0), return_dict=True
        ...     )
        ... }

        >>> # instantiate beam scorer
        >>> beam_scorer = BeamSearchScorer(
        ...     batch_size=1,
        ...     num_beams=num_beams,
        ...     device=model.device,
        ... )

        >>> # instantiate logits processors
        >>> logits_processor = LogitsProcessorList(
        ...     [
        ...         MinLengthLogitsProcessor(5, eos_token_id=model.config.eos_token_id),
        ...     ]
        ... )

        >>> outputs = model.beam_search(input_ids, beam_scorer, logits_processor=logits_processor, **model_kwargs)

        >>> tokenizer.batch_decode(outputs, skip_special_tokens=True)
        ['Wie alt bist du?']
        ```"""
        if self.config.num_beams <= 1:
            raise ValueError(
                f"{self.__class__.__name__} cannot be used with `num_beams == 1`. Please use `num_beams > 1`."
            )
        if logits_processor is not None:
            logits_processor = LogitsProcessorList(logits_processor)
            if any(isinstance(p, OldSBSLogitProcessor) for p in logits_processor):
                pass
            else:
                logits_processor.append(OldSBSLogitProcessor(num_beams=self.config.num_beams, batch_size=input_ids.shape[0]))
        else:
            logits_processor = LogitsProcessorList([OldSBSLogitProcessor(num_beams=self.config.num_beams,
                                                                         batch_size=input_ids.shape[0])])

        if self.check_config:
            if self.config.num_return_sequences > self.config.num_beams:
                raise ValueError("`num_return_sequences` has to be smaller or equal to `num_beams`.")

            if stopping_criteria.max_length is None:
                raise ValueError("`max_length` needs to be a stopping_criteria for now.")

        batch_size = input_ids.shape[0]
        # 11. prepare beam search scorer
        if self.scorer is None:
            self.scorer = BeamSearchScorer(
                batch_size=batch_size,
                num_beams=self.config.num_beams,
                device=input_ids.device,
                length_penalty=self.config.length_penalty,
                do_early_stopping=self.config.early_stopping,
                num_beam_hyps_to_keep=self.config.num_return_sequences,
                max_length=self.config.max_length,
            )
        # 12. interleave input_ids with `num_beams` additional sequences per batch
        input_ids, model_kwargs = model._expand_inputs_for_generation(
            input_ids=input_ids,
            expand_size=self.config.num_beams,
            is_encoder_decoder=model.config.is_encoder_decoder,
            **model_kwargs,
        )
        # 13. run beam search

        # init values
        logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
        stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
        if max_length is not None:
            warnings.warn(
                "`max_length` is deprecated in this function, use"
                " `stopping_criteria=StoppingCriteriaList(MaxLengthCriteria(max_length=max_length))` instead.",
                UserWarning,
            )
            stopping_criteria = validate_stopping_criteria(stopping_criteria, max_length)
        if len(stopping_criteria) == 0:
            warnings.warn("You don't have defined any stopping_criteria, this will likely loop forever", UserWarning)
        pad_token_id = pad_token_id if pad_token_id is not None else self.config.pad_token_id
        eos_token_id = eos_token_id if eos_token_id is not None else self.config.eos_token_id
        if isinstance(eos_token_id, int):
            eos_token_id = [eos_token_id]
        output_scores = output_scores if output_scores is not None else self.config.output_scores
        output_attentions = (
            output_attentions if output_attentions is not None else self.config.output_attentions
        )
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict_in_generate = (
            return_dict_in_generate
            if return_dict_in_generate is not None
            else self.config.return_dict_in_generate
        )

        batch_size = len(self.scorer._beam_hyps)
        num_beams = self.scorer.num_beams

        batch_beam_size, cur_len = input_ids.shape

        if num_beams * batch_size != batch_beam_size:
            raise ValueError(
                f"Batch dimension of `input_ids` should be {num_beams * batch_size}, but is {batch_beam_size}."
            )

        # init attention / hidden states / scores tuples
        scores = () if (return_dict_in_generate and output_scores) else None
        beam_indices = (
            tuple(() for _ in range(batch_beam_size)) if (return_dict_in_generate and output_scores) else None
        )
        decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
        cross_attentions = () if (return_dict_in_generate and output_attentions) else None
        decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None

        # if model is an encoder-decoder, retrieve encoder attention weights and hidden states
        if return_dict_in_generate and model.config.is_encoder_decoder:
            encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None
            encoder_hidden_states = (
                model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None
            )

        # initialise score of first beam with 0 and the rest with -1e9. This makes sure that only tokens
        # of the first beam are considered to avoid sampling the exact same tokens across all beams.
        beam_scores = torch.zeros((batch_size, num_beams), dtype=beam_scores_type, device=input_ids.device)
        beam_scores[:, 1:] = -float("inf")
        beam_scores = beam_scores.view((batch_size * num_beams,))

        this_peer_finished = False  # used by synced_gpus only
        while True:
            if synced_gpus:
                # Under synced_gpus the `forward` call must continue until all gpus complete their sequence.
                # The following logic allows an early break if all peers finished generating their sequence
                this_peer_finished_flag = torch.tensor(0.0 if this_peer_finished else 1.0).to(input_ids.device)
                # send 0.0 if we finished, 1.0 otherwise
                dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM)
                # did all peers finish? the reduced sum will be 0.0 then
                if this_peer_finished_flag.item() == 0.0:
                    break

            model_inputs = model.prepare_inputs_for_generation(input_ids, **model_kwargs)

            outputs = model(
                **model_inputs,
                return_dict=True,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
            )

            if synced_gpus and this_peer_finished:
                cur_len = cur_len + 1
                continue  # don't waste resources running the code we don't need

            next_token_logits = outputs.logits[:, -1, :]
            next_token_scores = nn.functional.log_softmax(
                next_token_logits, dim=-1
            )  # (batch_size * num_beams, vocab_size)

            model_log_probs = next_token_scores.clone()
            next_token_scores_processed = logits_processor(input_ids,
                                                           next_token_scores,
                                                           beam_log_probs=beam_scores,
                                                           past_scores=scores,
                                                           beam_indices=beam_indices, )
            next_token_scores = next_token_scores_processed
            # Store scores, attentions and hidden_states when required
            if return_dict_in_generate:
                if output_scores:
                    scores += (next_token_scores_processed,)
                if output_attentions:
                    decoder_attentions += (
                        (outputs.decoder_attentions,) if model.config.is_encoder_decoder else (outputs.attentions,)
                    )
                    if model.config.is_encoder_decoder:
                        cross_attentions += (outputs.cross_attentions,)

                if output_hidden_states:
                    decoder_hidden_states += (
                        (outputs.decoder_hidden_states,)
                        if model.config.is_encoder_decoder
                        else (outputs.hidden_states,)
                    )

            # reshape for beam search
            vocab_size = next_token_scores.shape[-1]
            next_token_scores = next_token_scores.view(batch_size, num_beams * vocab_size)

            # Sample 2 next tokens for each beam (so we have some spare tokens and match output of beam search)
            next_token_scores, next_tokens = torch.topk(
                next_token_scores, 2 * num_beams, dim=1, largest=True, sorted=True
            )

            model_log_probs = model_log_probs.view((batch_size, num_beams * vocab_size)).gather(dim=1,
                                                                                                index=next_tokens)

            next_indices = torch.div(next_tokens, vocab_size, rounding_mode="floor")
            next_tokens = next_tokens % vocab_size

            # stateless
            beam_outputs = self.scorer.process(
                input_ids,
                next_token_scores,
                next_tokens,
                next_indices,
                pad_token_id=pad_token_id,
                eos_token_id=eos_token_id,
                beam_indices=beam_indices,
                beam_scores=beam_scores,
                next_true_log_probs=model_log_probs,
            )

            beam_scores = beam_outputs["next_beam_scores"]
            beam_next_tokens = beam_outputs["next_beam_tokens"]
            beam_idx = beam_outputs["next_beam_indices"]
            beam_gumbels = beam_outputs["next_beam_gumbels"]

            input_ids = torch.cat([input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1)

            model_kwargs = model._update_model_kwargs_for_generation(
                outputs, model_kwargs, is_encoder_decoder=model.config.is_encoder_decoder
            )
            if model_kwargs["past_key_values"] is not None:
                model_kwargs["past_key_values"] = model._reorder_cache(model_kwargs["past_key_values"], beam_idx)

            if return_dict_in_generate and output_scores:
                beam_indices = tuple((beam_indices[beam_idx[i]] + (beam_idx[i],) for i in range(len(beam_indices))))

            # increase cur_len
            cur_len = cur_len + 1

            if self.scorer.is_done or stopping_criteria(input_ids, scores):
                if not synced_gpus:
                    break
                else:
                    this_peer_finished = True

        sequence_outputs = self.scorer.finalize(
            input_ids,
            beam_scores,
            next_tokens,
            next_indices,
            pad_token_id=pad_token_id,
            eos_token_id=eos_token_id,
            max_length=stopping_criteria.max_length,
            beam_indices=beam_indices,
            beam_gumbels=beam_gumbels,
        )

        if return_dict_in_generate:
            if not output_scores:
                sequence_outputs["sequence_scores"] = None

            if model.config.is_encoder_decoder:
                return BeamSearchEncoderDecoderOutput(
                    sequences=sequence_outputs["sequences"],
                    sequences_scores=sequence_outputs["sequence_scores"],
                    scores=scores,
                    beam_indices=sequence_outputs["beam_indices"],
                    encoder_attentions=encoder_attentions,
                    encoder_hidden_states=encoder_hidden_states,
                    decoder_attentions=decoder_attentions,
                    cross_attentions=cross_attentions,
                    decoder_hidden_states=decoder_hidden_states,
                    beam_gumbels=sequence_outputs["beam_gumbels"],
                )
            else:
                return BeamSearchDecoderOnlyOutput(
                    sequences=sequence_outputs["sequences"],
                    sequences_scores=sequence_outputs["sequence_scores"],
                    scores=scores,
                    beam_indices=sequence_outputs["beam_indices"],
                    attentions=decoder_attentions,
                    hidden_states=decoder_hidden_states,
                    beam_gumbels=sequence_outputs["beam_gumbels"],
                )
        else:
            return sequence_outputs["sequences"]
