from __future__ import division
import torch

"""
 Class for managing the internals of the beam search process.
 Takes care of beams, back pointers, and scores.
 Adapted from openNMT-py (https://github.com/OpenNMT/OpenNMT-py).
"""


class Beam(object):
    def __init__(self, size, go, eos, n_best=1, cuda=False, vocab=None, global_scorer=None):
        self.size = size
        self.tt = torch.cuda if cuda else torch

        # The score for each translation on the beam.
        self.scores = self.tt.FloatTensor(size).zero_()
        self.all_scores = []

        # The backpointers at each time-step.
        self.prev_ks = []

        # The outputs at each time-step.
        self.next_ys = [self.tt.LongTensor(size).fill_(0)]
        self.next_ys[0][0] = go
        self.vocab = vocab

        # Has EOS topped the beam yet.
        self._eos = eos
        self.eos_top = False

        # The attentions (matrix) for each time.
        self.attn = []

        # Time and k pair for finished.
        self.finished = []
        self.n_best = n_best

        # Information for global scoring.
        self.global_scorer = global_scorer
        self.global_state = {}

    def get_current_state(self):
        """Get the outputs for the current timestep."""
        return self.next_ys[-1]

    def get_current_origin(self):
        """Get the backpointers for the current timestep."""
        return self.prev_ks[-1]

    def advance(self, word_lk, attn_out, output_fixed_size=False):
        """
        Given prob over words for every last beam `word_lk` and attention
        `attn_out`: Compute and update the beam search.

        Parameters:

        * `word_lk`- probs of advancing from the last step (K x words)
        * `attn_out`- attention at the last step

        Returns: True if beam search is complete.
        """
        num_words = word_lk.size(1)

        if output_fixed_size:  # Block the EOS symbol since the size is fixed
            word_lk = word_lk.clone()
            word_lk[:, self._eos] = float('-inf')

        # Sum the previous scores.
        if len(self.prev_ks) > 0:
            beam_lk = word_lk + self.scores.unsqueeze(1).expand_as(word_lk)
            # Don't let EOS have children.
            for i in range(self.next_ys[-1].size(0)):
                if self.next_ys[-1].cpu()[i] == self._eos:
                    beam_lk[i] = -1e20
        else:
            beam_lk = word_lk[0]
        flat_beam_lk = beam_lk.view(-1)
        best_scores, best_scores_id = flat_beam_lk.topk(self.size, 0, True, True)

        self.all_scores.append(self.scores)
        self.scores = best_scores

        # best_scores_id is flattened beam x word array, so calculate which word and beam each score came from
        prev_k = best_scores_id / num_words
        self.prev_ks.append(prev_k)
        self.next_ys.append((best_scores_id - prev_k * num_words))
        self.attn.append(attn_out.index_select(0, prev_k))

        if self.global_scorer is not None:
            self.global_scorer.updateGlobalState(self)

        for i in range(self.next_ys[-1].size(0)):
            if self.next_ys[-1][i] == self._eos:
                s = self.scores[i]
                if self.global_scorer is not None:
                    global_scores = self.global_scorer.score(self, self.scores)
                    s = global_scores[i]
                self.finished.append((s, len(self.next_ys) - 1, i))

        # End condition is when top-of-beam is EOS.
        if self.next_ys[-1][0] == self._eos:
            self.eos_top = True

    def done(self):
        return self.eos_top and len(self.finished) >= self.n_best

    def sort_finished(self, minimum=None):
        if minimum is not None and minimum > 1:
            # Add the current contents of the beam to the finished set, even if they don't end with eos
            for i in range(len(self.scores)):
                s = self.scores[i]
                if self.globalScorer is not None:
                    global_scores = self.globalScorer.score(self, self.scores)
                    s = global_scores[i]
                self.finished.append((s, len(self.nextYs) - 1, i))
        else:
            # Always add the best performer in the beam to the finished set, even if it doesn't end with eos
            s = self.scores[0]
            self.finished.append((s, len(self.next_ys) - 1, 0))

        self.finished.sort(key=lambda a: -a[0])
        scores = [sc for sc, _, _ in self.finished]
        ks = [(t, k) for _, t, k in self.finished]
        return scores, ks

    def get_hyp(self, timestep, k):
        """
        Walk back to construct the full hypothesis.
        """
        hyp, attn = [], []
        for j in range(len(self.prev_ks[:timestep]) - 1, -1, -1):
            hyp.append(self.next_ys[j + 1][k])
            attn.append(self.attn[j][k])
            k = self.prev_ks[j][k]
        return hyp[::-1], torch.stack(attn[::-1])

