"""
From https://stackoverflow.com/questions/64356953/batch-wise-beam-search-in-pytorch
"""

import torch
import math

def backtrace_indices(indices,backpointers):
    """
    Use backpointers to reconstruct each timestep of each hypothesis in the beam

    indices: tensor, batch_size, beam_size, seq_len
    backpointers: list, len = seq_len, each tensor has length of beam_size 
    """
    beam_size = indices.shape[1]

    # start at end of sequence
    hypos = [indices[:,:,-1]]
    for i in reversed(range(len(backpointers))):
        back = backpointers[i].flatten()
        prev_step = indices[:,:,i][:,back]
        hypos = [prev_step] + hypos
    return torch.stack(hypos,dim=-1)


def beam_search_decoder(post, k):
    """Beam Search Decoder

    Parameters:

        post(Tensor) – the posterior of network.
        k(int) – beam size of decoder.

    Outputs:

        indices(Tensor) – a beam of index sequence.
        log_prob(Tensor) – a beam of log likelihood of sequence.

    Shape:

        post: (batch_size, seq_length, vocab_size).
        indices: (batch_size, beam_size, seq_length).
        log_prob: (batch_size, beam_size).

    Examples:

        >>> post = torch.softmax(torch.randn([32, 20, 1000]), -1)
        >>> indices, log_prob = beam_search_decoder(post, 3)

    """
    batch_size, seq_length,vocab_size = post.shape
    log_post = post#post.log()
    import pdb;pdb.set_trace()
    log_prob, indices = log_post[:, 0, :].topk(k, sorted=True)
    indices = indices.unsqueeze(-1)
    backpointers = []
    for i in range(1, seq_length):
        log_prob = log_prob.unsqueeze(-1) + log_post[:, i, :].unsqueeze(1).repeat(1, k, 1)
        log_prob, index = log_prob.view(batch_size, -1).topk(k, sorted=True)
        back = index/k # points to the idx of the hypothesis at the prev timestep that led to the best result in at this timestep
        backpointers.append(back)
        index = index%post.shape[-1]
        indices = torch.cat([indices, index.unsqueeze(-1)], dim=-1)
    
    return indices, log_prob, backpointers
                                                        
