# Copyright (c) 2019-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#

import math
import itertools
import torch
import torch.nn as nn
import torch.nn.functional as F
import model_utils
from copy import deepcopy
from entmax import entmax_bisect


N_MAX_POSITIONS = 512  # maximum input sequence length


def Embedding(num_embeddings, embedding_dim, padding_idx=None):
    m = nn.Embedding(num_embeddings, embedding_dim, padding_idx=padding_idx)
    nn.init.normal_(m.weight, mean=0, std=embedding_dim ** -0.5)
    if padding_idx is not None:
        nn.init.constant_(m.weight[padding_idx], 0)
    return m


def Linear(in_features, out_features, bias=True):
    m = nn.Linear(in_features, out_features, bias)
    # nn.init.normal_(m.weight, mean=0, std=1)
    # nn.init.xavier_uniform_(m.weight)
    # nn.init.constant_(m.bias, 0.)
    return m


class PredLayer(nn.Module):
    """
    Prediction layer (cross_entropy or adaptive_softmax).
    """
    def __init__(self, emb_dim, n_words):
        super().__init__()
        self.n_words = n_words
        self.proj = Linear(emb_dim, n_words, bias=True)

    def forward(self, x, y, get_scores=False):
        """
        Compute the loss, and optionally the scores.
        """
        scores = self.proj(x).view(-1, self.n_words)
        loss = F.cross_entropy(scores, y, reduction='mean')

        return scores, loss

    def get_scores(self, x):
        """
        Compute scores.
        """
        return self.proj(x)


class MultiHeadAttention(nn.Module):

    NEW_ID = itertools.count()

    def __init__(self, n_heads, dim, dropout):
        super().__init__()
        self.layer_id = next(MultiHeadAttention.NEW_ID)
        self.dim = dim
        self.n_heads = n_heads
        self.dropout = dropout
        assert self.dim % self.n_heads == 0

        self.q_lin = Linear(dim, dim)
        self.k_lin = Linear(dim, dim)
        self.v_lin = Linear(dim, dim)
        self.out_lin = Linear(dim, dim)

    def forward(self, input, mask, kv=None, cache=None):
        """
        Self-attention (if kv is None) or attention over source sentence (provided by kv).
        """
        # Input is (bs, qlen, dim)
        # Mask is (bs, klen) (non-causal) or (bs, klen, klen)
        bs, qlen, dim = input.size()
        if kv is None:
            klen = qlen if cache is None else cache['slen'] + qlen
        else:
            klen = kv.size(1)
        assert dim == self.dim, 'Dimensions do not match: %s input vs %s configured' % (dim, self.dim)
        n_heads = self.n_heads
        dim_per_head = dim // n_heads
        mask_reshape = (bs, 1, qlen, klen) if mask.dim() == 3 else (bs, 1, 1, klen)

        def shape(x):
            """  projection """
            return x.view(bs, -1, self.n_heads, dim_per_head).transpose(1, 2)

        def unshape(x):
            """  compute context """
            return x.transpose(1, 2).contiguous().view(bs, -1, self.n_heads * dim_per_head)

        q = shape(self.q_lin(input))                                          # (bs, n_heads, qlen, dim_per_head)
        if kv is None:
            k = shape(self.k_lin(input))                                      # (bs, n_heads, qlen, dim_per_head)
            v = shape(self.v_lin(input))                                      # (bs, n_heads, qlen, dim_per_head)
        elif cache is None or self.layer_id not in cache:
            k = v = kv
            k = shape(self.k_lin(k))                                          # (bs, n_heads, qlen, dim_per_head)
            v = shape(self.v_lin(v))                                          # (bs, n_heads, qlen, dim_per_head)

        if cache is not None:
            if self.layer_id in cache:
                if kv is None:
                    k_, v_ = cache[self.layer_id]
                    k = torch.cat([k_, k], dim=2)                             # (bs, n_heads, klen, dim_per_head)
                    v = torch.cat([v_, v], dim=2)                             # (bs, n_heads, klen, dim_per_head)
                else:
                    k, v = cache[self.layer_id]
            cache[self.layer_id] = (k, v)

        q = q / math.sqrt(dim_per_head)                                       # (bs, n_heads, qlen, dim_per_head)
        scores = torch.matmul(q, k.transpose(2, 3))                           # (bs, n_heads, qlen, klen)
        mask = (mask == 0).view(mask_reshape).expand_as(scores)               # (bs, n_heads, qlen, klen)
        scores.masked_fill_(mask, -float('inf'))                              # (bs, n_heads, qlen, klen)

        weights = F.softmax(scores.float(), dim=-1).type_as(scores)           # (bs, n_heads, qlen, klen)
        drop_weights = F.dropout(weights, p=self.dropout, training=self.training)  # (bs, n_heads, qlen, klen)
        context = torch.matmul(drop_weights, v)                                    # (bs, n_heads, qlen, dim_per_head)
        context = unshape(context)                                            # (bs, qlen, dim)

        return self.out_lin(context), weights


class TransformerFFN(nn.Module):

    def __init__(self, in_dim, dim_hidden, out_dim, dropout, gelu_activation):
        super().__init__()
        self.dropout = dropout
        self.lin1 = Linear(in_dim, dim_hidden)
        self.lin2 = Linear(dim_hidden, out_dim)
        self.act = F.gelu
    #
    # def forward(self, input):
    #     x = self.lin1(input)
    #     x = self.act(x)
    #     x = self.lin2(x)
    #     x = F.dropout(x, p=self.dropout, training=self.training)
    #     return x

    def forward(self, input):
        x = self.lin1(input)
        x = F.dropout(self.act(x), p=self.dropout, training=self.training)
        x = self.lin2(x)
        x = F.dropout(x, p=self.dropout, training=self.training)
        return x


class CacheTransformer(nn.Module):
    def __init__(self, n_words, bos_index, eos_index,
                 pad_index, emb_dim, n_heads, n_layers,
                 dropout, share_embedding, attention_dropout,
                 max_len, use_copy, use_entmax):
        """
        Transformer model decoder.
        """
        super().__init__()

        # encoder / decoder, output layer
        self.is_encoder = False
        self.is_decoder = not self.is_encoder

        self.n_words = n_words
        self.bos_index = bos_index
        self.eos_index = eos_index
        self.pad_index = pad_index
        self.use_copy = use_copy
        self.use_entmax = use_entmax

        # model parameters
        self.dim = emb_dim       # 512 by default
        self.hidden_dim = self.dim * 4  # 2048 by default
        self.n_heads = n_heads   # 8 by default
        self.n_layers = n_layers
        self.dropout = dropout
        self.attention_dropout = attention_dropout
        assert self.dim % self.n_heads == 0, \
            'transformer dim must be a multiple of n_heads'

        # embeddings
        self.position_embeddings = nn.Embedding(max_len, self.dim)
        self.embeddings = nn.Embedding(self.n_words, self.dim)

        # transformer layers
        self.attentions = nn.ModuleList()
        self.layer_norm1 = nn.ModuleList()
        self.ffns = nn.ModuleList()
        self.layer_norm2 = nn.ModuleList()
        if self.is_decoder:
            self.layer_norm15 = nn.ModuleList()
            self.encoder_attn = nn.ModuleList()

        for layer_id in range(self.n_layers):
            self.attentions.append(MultiHeadAttention(
                self.n_heads, self.dim, dropout=self.attention_dropout))
            self.layer_norm1.append(nn.LayerNorm(self.dim, eps=1e-12))
            if self.is_decoder:
                self.layer_norm15.append(nn.LayerNorm(self.dim, eps=1e-12))
                self.encoder_attn.append(
                    MultiHeadAttention(self.n_heads, self.dim,
                                       dropout=self.attention_dropout))

            self.ffns.append(
                TransformerFFN(self.dim, self.hidden_dim,
                               self.dim, dropout=self.dropout,
                               gelu_activation=F.gelu))
            self.layer_norm2.append(nn.LayerNorm(self.dim, eps=1e-12))

        # output layer
        # if self.with_output:
        self.pred_layer = PredLayer(emb_dim, n_words)
        if share_embedding:
            self.pred_layer.proj.weight = self.embeddings.weight
        if use_copy:
            self.copy_generator = \
                model_utils.CopyGenerator(self.dim, use_entmax)

    def _generate_square_subsequent_mask(self, seq):
        mask = (torch.triu(torch.ones(seq, seq)) == 1).transpose(0, 1)
        return mask

    def fwd(self, x, src_enc=None, src_mask=None,
            positions=None, cache=None, src_map=None):
        """
        Inputs:
            `x` LongTensor(slen, bs), containing word indices
            `lengths` LongTensor(bs), containing the length of each sentence
            `causal` Boolean, if True, the attention is only done over previous hidden states
            `positions` LongTensor(slen, bs), containing word positions
            `langs` LongTensor(slen, bs), containing language IDs

            src_enc: bs x seq len x dim
            x: bs x seq len
        """
        bs, slen = x.size()
        if src_enc is not None:
            assert self.is_decoder
            assert src_enc.size(0) == bs

        attn_mask = self._generate_square_subsequent_mask(slen).to(x.device)
        attn_mask = attn_mask.unsqueeze(0).expand(bs, slen, slen)

        # positions
        if positions is None:
            # positions = x.new(slen).long()
            positions = torch.arange(slen).unsqueeze(0).to(x.device).long()
        else:
            assert positions.size() == (slen, bs)
            positions = positions.transpose(0, 1)

        # do not recompute cached elements
        if cache is not None:
            _slen = slen - cache['slen']
            x = x[:, -_slen:]
            positions = positions[:, -_slen:]
            attn_mask = attn_mask[:, -_slen:]

        # embeddings
        tensor = self.embeddings(x.long())
        tensor = tensor + self.position_embeddings(positions).expand_as(tensor)
        tensor = F.dropout(tensor, p=self.dropout, training=self.training)

        # transformer layers
        for i in range(self.n_layers):

            # self attention
            attn, _ = self.attentions[i](tensor, attn_mask, cache=cache)
            attn = F.dropout(attn, p=self.dropout, training=self.training)
            tensor = tensor + attn
            tensor = self.layer_norm1[i](tensor)

            # encoder attention (for decoder only)
            if self.is_decoder and src_enc is not None:
                attn, attn_weight = \
                    self.encoder_attn[i](tensor, src_mask,
                                         kv=src_enc, cache=cache)
                attn = F.dropout(attn, p=self.dropout, training=self.training)
                tensor = tensor + attn
                tensor = self.layer_norm15[i](tensor)

            tensor = tensor + self.ffns[i](tensor)
            tensor = self.layer_norm2[i](tensor)

        # update cache length
        if cache is not None:
            cache['slen'] += tensor.size(1)

        pred_prob = self.pred_layer.get_scores(tensor)

        if self.use_copy:
            assert src_map is not None
            copy_attn = attn_weight[:, 0, :, :].contiguous()
            pred_prob = self.copy_generator(
                hidden=self.copy_generator._bottle(tensor),
                orig_prob=self.copy_generator._bottle(pred_prob),
                attn=self.copy_generator._bottle(copy_attn),
                src_map=src_map)

        return pred_prob, tensor

    def _generate(self, src_enc, src_mask, max_len=200, min_len=0, top_p=None, src_map=None, src_tgt_vocab_map=None):
        """
        Decode a sentence given initial start.
        `x`:
            - LongTensor(bs, slen)
                <EOS> W1 W2 W3 <EOS> <PAD>
                <EOS> W1 W2 W3   W4  <EOS>
        `lengths`:
            - LongTensor(bs) [5, 6]
        `positions`:
            - False, for regular "arange" positions (LM)
            - True, to reset positions from the new generation (MT)
        `langs`:
            - must be None if the model only supports one language
            - lang_id if only one language is involved (LM)
            - (lang_id1, lang_id2) if two languages are involved (MT)
        """

        # input batch
        bs = len(src_mask)
        assert src_enc.size(0) == bs

        # generated sentences
        generated = src_mask.new(bs, max_len)  # upcoming output
        generated.fill_(self.pad_index)       # fill upcoming ouput with <PAD>
        generated[:, 0].fill_(self.bos_index)    # we use <EOS> for <BOS> everywhere

        # current position / max lengths / length of generated sentences / unfinished sentences
        cur_len = 1
        # gen_len = torch.ones(bs).to(src_mask.device).long()
        unfinished_sents = torch.ones(bs).to(src_mask.device).long() #src_len.clone().fill_(1)
        all_scores = torch.zeros(bs).to(src_mask.device)
        # cache compute states
        cache = {'slen': 0}

        while cur_len < max_len:

            # compute word scores
            tensor, _ = self.fwd(
                x=generated[:, :cur_len] if not self.use_copy else generated[:, :cur_len].masked_fill(generated[:, :cur_len].gt(self.n_words - 1), 0),
                src_enc=src_enc,
                src_mask=src_mask,
                cache=cache,
                src_map=src_map
            )
            if self.use_copy:
                tensor = torch.cat(tensor, 1)
                scores, _ = model_utils.collapse_copy_scores(
                    scores=tensor,
                    src_tgt_vocab_map=src_tgt_vocab_map,
                    vocab_size=self.n_words)
                scores[:, self.n_words] = 0.0
            else:
                assert tensor.size() == (bs, 1, self.n_words), (cur_len, max_len, src_enc.size(), tensor.size(), (1, bs, self.n_words))
                scores = tensor[:, -1, :]  # (bs, dim)
            # scores = self.pred_layer.get_scores(tensor)      # (bs, n_words)

            scores[:, 0] = -float('Inf') if not self.use_copy else 0
            scores[:, self.pad_index] = -float('Inf') if not self.use_copy else 0
            scores[:, self.bos_index] = -float('Inf') if not self.use_copy else 0

            if cur_len < min_len:
                scores[:, self.eos_index] = -float('Inf') if not self.use_copy else 0

            # select next words: sample or greedy
            if top_p:
                if self.use_copy:
                    next_words = torch.multinomial(model_utils.top_k_top_p_filtering(scores, top_k=0.0, top_p=top_p if top_p else 0.0, filter_value=0.0, need_softmax=False), 1).squeeze(1)
                    next_scores = (scores + 1e-10).log().gather(1, next_words.unsqueeze(1)).squeeze(1)
                else:
                    next_words = torch.multinomial(F.softmax(model_utils.top_k_top_p_filtering(scores, top_k=0.0, top_p=top_p if top_p else 0.0), dim=1), 1).squeeze(1)
                    next_scores = scores.log_softmax(1).gather(1, next_words.unsqueeze(1)).squeeze(1)
            else:
                if self.use_copy:
                    next_scores, next_words = (scores + 1e-10).log().max(1)
                elif self.use_entmax:
                    next_scores, next_words = (entmax_bisect(scores, 1.2) + 1e-10).log().max(1)
                else:
                    next_scores, next_words = scores.log_softmax(1).max(1)
            assert next_words.size() == (bs,)

            # update generations / lengths / finished sentences / current length
            generated[:, cur_len] = next_words * unfinished_sents + self.pad_index * (1 - unfinished_sents)
            all_scores = all_scores + next_scores * unfinished_sents.float()
            # gen_len.add_(unfinished_sents)
            unfinished_sents.mul_(next_words.ne(self.eos_index).long())
            cur_len = cur_len + 1

            # stop when there is a </s> in each sentence, or if we exceed the maximul length
            if unfinished_sents.max() == 0:
                break

        # add <EOS> to unfinished sentences
        if cur_len == max_len:
            generated[:, -1].masked_fill_(unfinished_sents.bool(), self.eos_index)

        # sanity check
        assert (generated == self.eos_index).sum() == bs

        return generated[:, 1:cur_len].cpu().numpy(), all_scores.cpu().numpy() #, gen_len


    def _generate_beam(self, src_enc, src_mask, beam_size, length_penalty=0.0, early_stopping=False, min_len=0, max_len=200, trigram_blocking=False, return_all=False, src_map=None, src_tgt_vocab_map=None):
        """
        Decode a sentence given initial start.
        `x`:
            - LongTensor(bs, slen)
                <EOS> W1 W2 W3 <EOS> <PAD>
                <EOS> W1 W2 W3   W4  <EOS>
        `lengths`:
            - LongTensor(bs) [5, 6]
        `positions`:
            - False, for regular "arange" positions (LM)
            - True, to reset positions from the new generation (MT)
        `langs`:
            - must be None if the model only supports one language
            - lang_id if only one language is involved (LM)
            - (lang_id1, lang_id2) if two languages are involved (MT)
        """

        # check inputs
        assert src_enc.size(0) == src_mask.size(0)
        assert beam_size >= 1

        # batch size / number of words
        bs = len(src_mask)
        n_words = self.n_words if not self.use_copy else self.n_words + src_tgt_vocab_map.shape[1]

        # expand to beam size the source latent representations / source lengths
        src_enc = src_enc.unsqueeze(1).expand((bs, beam_size) + src_enc.shape[1:]).contiguous().view((bs * beam_size,) + src_enc.shape[1:])
        src_mask = src_mask.unsqueeze(1).expand((bs, beam_size) + src_mask.shape[1:]).contiguous().view((bs * beam_size,) + src_mask.shape[1:])
        if src_tgt_vocab_map is not None:
            src_tgt_vocab_map = src_tgt_vocab_map.unsqueeze(1).expand((bs, beam_size) + src_tgt_vocab_map.shape[1:]).contiguous().view((bs * beam_size,) + src_tgt_vocab_map.shape[1:])
        if src_map is not None:
            src_map = src_map.unsqueeze(1).expand((bs, beam_size) + src_map.shape[1:]).contiguous().view((bs * beam_size,) + src_map.shape[1:])
        # src_len = src_len.unsqueeze(1).expand(bs, beam_size).contiguous().view(-1)

        # generated sentences (batch with beam current hypotheses)
        generated = src_enc.new(bs * beam_size, max_len)  # upcoming output
        generated.fill_(self.pad_index)                   # fill upcoming ouput with <PAD>
        generated[:, 0].fill_(self.bos_index)                # we use <EOS> for <BOS> everywhere

        # generated hypotheses
        generated_hyps = [BeamHypotheses(beam_size, max_len, length_penalty, early_stopping) for _ in range(bs)]
        trigram_set = [set() for _ in range(bs * beam_size)]

        # scores for each sentence in the beam
        beam_scores = src_enc.new(bs, beam_size).fill_(0)
        beam_scores[:, 1:] = -1e9
        beam_scores = beam_scores.view(-1)

        # current position
        cur_len = 1

        # cache compute states
        cache = {'slen': 0}

        # done sentences
        done = [False for _ in range(bs)]

        while cur_len < max_len:

            # compute word scores
            tensor, _ = self.fwd(
                x=generated[:, :cur_len] if not self.use_copy else generated[:, :cur_len].masked_fill(generated[:, :cur_len].gt(self.n_words - 1), 0),
                src_enc=src_enc,
                src_mask=src_mask,
                cache=cache,
                src_map=src_map
            )
            if self.use_copy:
                tensor = torch.cat(tensor, 1)
                scores, _ = model_utils.collapse_copy_scores(
                    scores=tensor,
                    src_tgt_vocab_map=src_tgt_vocab_map,
                    vocab_size=self.n_words)
                scores[:, self.n_words] = 0
            else:
                assert tensor.size() == (bs * beam_size, 1, self.n_words)
                scores = tensor[:, -1, :]               # (bs * beam_size, dim)

            scores[:, 0] = -float('Inf') if not self.use_copy else 0
            scores[:, self.pad_index] = -float('Inf') if not self.use_copy else 0
            scores[:, self.bos_index] = -float('Inf') if not self.use_copy else 0

            if cur_len < min_len:
                scores[:, self.eos_index] = -float('Inf') if not self.use_copy else 0

            if self.use_copy:
                scores = (scores + 1e-10).log()
            elif self.use_entmax:
                scores = torch.log(entmax_bisect(scores, 1.2) + 1e-10)
            else:
                scores = F.log_softmax(scores, dim=-1)       # (bs * beam_size, n_words)

            assert scores.size() == (bs * beam_size, n_words), (scores.shape, (bs * beam_size, n_words))

            # select next words with scores
            _scores = scores + beam_scores[:, None].expand_as(scores)  # (bs * beam_size, n_words)
            _scores = _scores.view(bs, beam_size * n_words)            # (bs, beam_size * n_words)

            next_scores, next_words = torch.sort(_scores, dim=1, descending=True)
            assert next_scores.size() == next_words.size() == (bs, n_words * beam_size)

            # next batch beam content
            # list of (bs * beam_size) tuple(next hypothesis score, next word, current position in the batch)
            next_batch_beam = []

            # for each sentence
            for sent_id in range(bs):

                # if we are done with this sentence
                done[sent_id] = done[sent_id] or generated_hyps[sent_id].is_done(next_scores[sent_id].max().item())
                if done[sent_id]:
                    next_batch_beam.extend([(0, self.pad_index, 0)] * beam_size)  # pad the batch
                    continue

                # next sentence beam content
                next_sent_beam = []
                n_add = 0

                # next words for this sentence
                for idx, value in zip(next_words[sent_id], next_scores[sent_id]):

                    # get beam and word IDs
                    beam_id = idx // n_words
                    word_id = idx % n_words

                    if trigram_blocking and cur_len > 2:
                        trigram = tuple(generated[sent_id * beam_size + beam_id, cur_len-2:cur_len].tolist() + [word_id.item()])
                        if trigram in trigram_set[sent_id * beam_size + beam_id]:
                            continue
                    # end of sentence, or next word
                    if word_id == self.eos_index or cur_len + 1 == max_len:
                        n_add += 1
                        generated_hyps[sent_id].add(generated[sent_id * beam_size + beam_id, :cur_len].clone(), value.item())
                    else:
                        next_sent_beam.append((value, word_id, sent_id * beam_size + beam_id))
                        if trigram_blocking and cur_len > 2:
                            trigram_set[sent_id * beam_size + beam_id].add(trigram)

                    # the beam for next step is full
                    if len(next_sent_beam) == beam_size or (cur_len + 1 == max_len and n_add == beam_size):
                        break

                # update next beam content
                assert len(next_sent_beam) == 0 if cur_len + 1 == max_len else beam_size
                if len(next_sent_beam) == 0:
                    next_sent_beam = [(0, self.pad_index, 0)] * beam_size  # pad the batch
                next_batch_beam.extend(next_sent_beam)
                assert len(next_batch_beam) == beam_size * (sent_id + 1)

            # sanity check / prepare next batch
            assert len(next_batch_beam) == bs * beam_size
            beam_scores = beam_scores.new([x[0] for x in next_batch_beam])
            beam_words = generated.new([x[1] for x in next_batch_beam])
            beam_idx = generated.new([x[2] for x in next_batch_beam]).long()

            # re-order batch and internal states
            trigram_set = [deepcopy(trigram_set[x[2]]) for x in next_batch_beam]
            generated = generated[beam_idx, :]
            generated[:, cur_len] = beam_words
            for k in cache.keys():
                if k != 'slen':
                    cache[k] = (cache[k][0][beam_idx], cache[k][1][beam_idx])

            # update current length
            cur_len = cur_len + 1

            # stop when we are done with each sentence
            if all(done):
                break

        if return_all:
            return generated_hyps

        # select the best hypotheses
        tgt_len = src_enc.new(bs).long()
        best = []
        best_scores = []

        for i, hypotheses in enumerate(generated_hyps):
            best_score, best_hyp = max(hypotheses.hyp, key=lambda x: x[0])
            tgt_len[i] = len(best_hyp) + 1  # +1 for the <EOS> symbol
            best.append(best_hyp)
            best_scores.append(best_score)

        # generate target batch
        decoded = src_enc.new(tgt_len.max().item(), bs).fill_(self.pad_index)
        for i, hypo in enumerate(best):
            decoded[:tgt_len[i] - 1, i] = hypo
            decoded[tgt_len[i] - 1, i] = self.eos_index

        # sanity check
        assert (decoded == self.eos_index).sum() == bs

        return decoded.transpose(0, 1).cpu().numpy(), best_scores, tgt_len.cpu().numpy()


class BeamHypotheses(object):

    def __init__(self, n_hyp, max_len, length_penalty, early_stopping):
        """
        Initialize n-best list of hypotheses.
        """
        self.max_len = max_len - 1  # ignoring <BOS>
        self.length_penalty = length_penalty
        self.early_stopping = early_stopping
        self.n_hyp = n_hyp
        self.hyp = []
        self.worst_score = 1e9

    def __len__(self):
        """
        Number of hypotheses in the list.
        """
        return len(self.hyp)

    def add(self, hyp, sum_logprobs):
        """
        Add a new hypothesis to the list.
        """
        score = sum_logprobs / len(hyp) ** self.length_penalty
        if len(self) < self.n_hyp or score > self.worst_score:
            self.hyp.append((score, hyp))
            if len(self) > self.n_hyp:
                sorted_scores = sorted([(s, idx) for idx, (s, _) in enumerate(self.hyp)])
                del self.hyp[sorted_scores[0][1]]
                self.worst_score = sorted_scores[1][0]
            else:
                self.worst_score = min(score, self.worst_score)

    def is_done(self, best_sum_logprobs):
        """
        If there are enough hypotheses and that none of the hypotheses being generated
        can become better than the worst one in the heap, then we are done with this sentence.
        """
        if len(self) < self.n_hyp:
            return False
        elif self.early_stopping:
            return True
        else:
            return self.worst_score >= best_sum_logprobs / self.max_len ** self.length_penalty
