import torch
import numpy as np
import torch.nn as nn
import itertools
from torch.nn import functional as F


class Beam(object):
    def __init__(self, beam_size):
        self.beam_size = beam_size

        self.candidates = []
        self.scores = []

    def step(self, prob, prev_beam, f_done):
        pre_score = prob.new_tensor(prev_beam.scores)

        score = prob + pre_score.unsqueeze(-1).expand_as(prob)
        if score.numel() < self.beam_size:
            nbest_score, nbest_ix = score.view(-1).topk(score.numel(), largest=False)
        else:
            nbest_score, nbest_ix = score.view(-1).topk(self.beam_size, largest=False)

        beam_ix = nbest_ix / prob.size(1)
        token_ix = nbest_ix - beam_ix * prob.size(1)

        done_list, remain_list = [], []
        prev_candidates = prev_beam.candidates
        for b_score, b_ix, t_ix in itertools.zip_longest(nbest_score.tolist(), beam_ix.tolist(), token_ix.tolist()):
            candidate = prev_candidates[b_ix] + [t_ix]

            if f_done(candidate):
                done_list.append([candidate, b_score])
            else:
                remain_list.append(b_ix)
                self.candidates.append(candidate)
                self.scores.append(b_score)
        return done_list, remain_list
