# Copyright (c) Microsoft Corporation. 
# Licensed under the MIT license.

import torch
import numpy as np
class Beam(object):
    def __init__(self, size,sos,eos,prob):
        self.size = size
        topk=prob.shape[-1]
        self.tt = torch.cuda

        # The score for each translation on the beam.
        self.scores = self.tt.FloatTensor(size*topk).zero_()
        self.scores=self.scores.view(size,topk)+torch.log(prob)[None,:]

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

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

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


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

    def getCurrentState(self):
        "Get the outputs for the current timestep."
        batch = self.tt.LongTensor(self.nextYs[-1]).view(-1, 1)
        return batch

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

    def advance(self, wordLk):
        """
        Given prob over words for every last beam `wordLk` and attention
        `attnOut`: Compute and update the beam search.

        Parameters:

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

        Returns: True if beam search is complete.
        """
        numWords = wordLk.size(2)

        # Sum the previous scores.
        if len(self.prevKs) > 0:
            beamLk = wordLk + self.scores.unsqueeze(2).expand_as(wordLk)
            # Don't let EOS have children.
            for i in range(self.nextYs[-1].size(0)):
                if self.nextYs[-1][i] == self._eos:
                    beamLk[i] = -1e20
        else:
            beamLk = wordLk[:1]+ self.scores.unsqueeze(2).expand_as(wordLk)[:1]
        flatBeamLk = torch.exp(beamLk).sum(1).view(-1)
        bestScores, bestScoresId = flatBeamLk.topk(self.size, 0, True, True)

        # bestScoresId is flattened beam x word array, so calculate which
        # word and beam each score came from
        prevK = bestScoresId // numWords
        self.prevKs.append(prevK)
        self.nextYs.append((bestScoresId - prevK * numWords))


        for i in range(self.nextYs[-1].size(0)):
            self.scores[i,:]=beamLk[self.prevKs[-1][i],:,self.nextYs[-1][i]]
            if self.nextYs[-1][i] == self._eos:
                s = torch.exp(self.scores[i]).sum()
                self.finished.append((s, len(self.nextYs) - 1, i))

        # End condition is when top-of-beam is EOS and no global score.
        if self.nextYs[-1][0] == self._eos:
            self.eosTop = True

    def done(self):
        return self.eosTop and len(self.finished) >=self.size

    def getFinal(self):
        if len(self.finished) == 0:
            self.finished.append((torch.exp(self.scores[0]).sum(), len(self.nextYs) - 1, 0))
        self.finished.sort(key=lambda a: -a[0])
        if len(self.finished) != self.size:
            unfinished=[]
            for i in range(self.nextYs[-1].size(0)):
                if self.nextYs[-1][i] != self._eos:
                    s = torch.exp(self.scores[i]).sum()
                    unfinished.append((s, len(self.nextYs) - 1, i)) 
            unfinished.sort(key=lambda a: -a[0])
            self.finished+=unfinished[:self.size-len(self.finished)]
        return self.finished[:self.size]

    def getHyp(self, beam_res):
        """
        Walk back to construct the full hypothesis.
        """
        hyps=[]
        for _,timestep, k in beam_res:
            hyp = []
            for j in range(len(self.prevKs[:timestep]) - 1, -1, -1):
                hyp.append(self.nextYs[j+1][k])
                k = self.prevKs[j][k]
            hyps.append(hyp[::-1])
        return hyps
    
    def buildTargetTokens(self, preds):
        sentence=[]
        for pred in preds:
            tokens = []
            for tok in pred:
                if tok==self._eos :
                    break
                tokens.append(tok)
            sentence.append(tokens)
        return sentence    

# Copyright (c) Microsoft Corporation. 
# Licensed under the MIT license.

import torch
import numpy as np
class Beam(object):
    def __init__(self, size,sos,eos,prob):
        self.size = size
        topk=prob.shape[-1]
        self.tt = torch.cuda

        # The score for each translation on the beam.
        self.scores = self.tt.FloatTensor(size*topk).zero_()
        self.scores=self.scores.view(size,topk)+torch.log(prob)[None,:]

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

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

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


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

    def getCurrentState(self):
        "Get the outputs for the current timestep."
        batch = self.tt.LongTensor(self.nextYs[-1]).view(-1, 1)
        return batch

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

    def advance(self, wordLk):
        """
        Given prob over words for every last beam `wordLk` and attention
        `attnOut`: Compute and update the beam search.

        Parameters:

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

        Returns: True if beam search is complete.
        """
        numWords = wordLk.size(2)

        # Sum the previous scores.
        if len(self.prevKs) > 0:
            beamLk = wordLk + self.scores.unsqueeze(2).expand_as(wordLk)
            # Don't let EOS have children.
            for i in range(self.nextYs[-1].size(0)):
                if self.nextYs[-1][i] == self._eos:
                    beamLk[i] = -1e20
        else:
            beamLk = wordLk[:1]+ self.scores.unsqueeze(2).expand_as(wordLk)[:1]
        flatBeamLk = torch.exp(beamLk).sum(1).view(-1)
        bestScores, bestScoresId = flatBeamLk.topk(self.size, 0, True, True)

        # bestScoresId is flattened beam x word array, so calculate which
        # word and beam each score came from
        prevK = bestScoresId // numWords
        self.prevKs.append(prevK)
        self.nextYs.append((bestScoresId - prevK * numWords))


        for i in range(self.nextYs[-1].size(0)):
            self.scores[i,:]=beamLk[self.prevKs[-1][i],:,self.nextYs[-1][i]]
            if self.nextYs[-1][i] == self._eos:
                s = torch.exp(self.scores[i]).sum()
                self.finished.append((s, len(self.nextYs) - 1, i))

        # End condition is when top-of-beam is EOS and no global score.
        if self.nextYs[-1][0] == self._eos:
            self.eosTop = True

    def done(self):
        return self.eosTop and len(self.finished) >=self.size

    def getFinal(self):
        if len(self.finished) == 0:
            self.finished.append((torch.exp(self.scores[0]).sum(), len(self.nextYs) - 1, 0))
        self.finished.sort(key=lambda a: -a[0])
        if len(self.finished) != self.size:
            unfinished=[]
            for i in range(self.nextYs[-1].size(0)):
                if self.nextYs[-1][i] != self._eos:
                    s = torch.exp(self.scores[i]).sum()
                    unfinished.append((s, len(self.nextYs) - 1, i)) 
            unfinished.sort(key=lambda a: -a[0])
            self.finished+=unfinished[:self.size-len(self.finished)]
        return self.finished[:self.size]

    def getHyp(self, beam_res):
        """
        Walk back to construct the full hypothesis.
        """
        hyps=[]
        for _,timestep, k in beam_res:
            hyp = []
            for j in range(len(self.prevKs[:timestep]) - 1, -1, -1):
                hyp.append(self.nextYs[j+1][k])
                k = self.prevKs[j][k]
            hyps.append(hyp[::-1])
        return hyps
    
    def buildTargetTokens(self, preds):
        sentence=[]
        for pred in preds:
            tokens = []
            for tok in pred:
                if tok==self._eos :
                    break
                tokens.append(tok)
            sentence.append(tokens)
        return sentence    

# Copyright (c) Microsoft Corporation. # Licensed under the MIT license. # The score for each translation on the beam. # The backpointers at each time-step. # The outputs at each time-step. # Has EOS topped the beam yet. # Time and k pair for finished. Given prob over words for every last beam `wordLk` and attention         `attnOut`: Compute and update the beam search.          Parameters:          * `wordLk`- probs of advancing from the last step (K x words)         * `attnOut`- attention at the last step          Returns: True if beam search is complete. # Sum the previous scores. # Don't let EOS have children. # bestScoresId is flattened beam x word array, so calculate which # word and beam each score came from # End condition is when top-of-beam is EOS and no global score. Walk back to construct the full hypothesis.