import torch


def get_context_vector(encoded_sents, target, future=False, use_cuda=False):
    """
    Get create the context vector for the sentence given at index target for 
    state classification. Do this by max pooling the sentences before the target sentence
    Args:
        encoded_sents (Tensor, [num_sents, batch, encoder dim])
        target (int) : Index of the target sentence (starts at 0)
        future (bool) : If true, use the vectors from the future instead of the past 
    Ret:
        context vector (Tensor [batch, encoder dim])
    """

    if target == 0 and not future:
        #return encoded_target[0, :, :]
        return torch.zeros(encoded_sents[0,:,:].shape).cuda() if use_cuda else torch.zeros(encoded_sents[0,:,:].shape)
    elif target == encoded_sents.shape[0]-1 and future:
        #return encoded_target[encoded_sents.shape[0]-1, :, :]
        return torch.zeros(encoded_sents[0,:,:].shape).cuda() if use_cuda else torch.zeros(encoded_sents[0,:,:].shape)


    if not future:
        sents = encoded_sents[:target, :, :] #[sents, batch, encoder dim]
    else:
        sents = encoded_sents[target+1:, :, :] #[sents, batch, encoder dim]

#    maxpool, _ = torch.max(sents, dim=0) #[batch, encoder dim]
    maxpool = torch.mean(sents, dim=0) #[batch, encoder dim]

    return maxpool