import torch


def classify_action_with_tsn_bni(input, settings, num_crops=None, n_samples=None, average_scores=True):
    """
    Test routine for tsn bni, for a single action

    :param input: Pytorch variable containing image values
    :param settings: dictionary settings
    :param num_crops: number of crops (should be either 1 or 10) for testing
    :param n_samples: number of samples to evaluate the model on
    :param average_scores: if True, will average the scores obtained from each sample to get a 1 x m array of scores,
     where m is the number of classes. If False with return a n_samples x m array of scores
    :return: the classification scores for the action
    """
    if input.shape[0] != 1:
        # in this case we feed one action per batch since we might be taking loads of crops in
        raise Exception('Please set batch size equal to 1!')

    mode = settings['mode'].lower()
    model = settings['model'].module
    of_stack_size = settings['of_stack_size']
    n_classes = settings['n_classes']

    if n_samples is None:
        n_samples = settings['testing']['n_samples']

    if num_crops is None:
        num_crops = settings['testing']['n_crops']

    if mode == 'rgb':
        length = 3
    elif mode == 'flow':
        length = of_stack_size * 2
    else:
        raise ValueError("Unknown modality " + mode)

    if num_crops == 10:
        input_var = torch.autograd.Variable(input.view(-1, length, input.size(2), input.size(3)), volatile=True)
    elif num_crops == 1:
        input_var = torch.autograd.Variable(input, volatile=True)
    else:
        raise Exception('Cannot deal with num_crops={}'.format(num_crops))

    action_output = model(input_var, do_consensus=False, override_reshape=False).data
    action_output = torch.mean(action_output.view((num_crops, n_samples, n_classes)), 0).view((n_samples,
                                                                                               n_classes))

    if average_scores:
        # averaging the scores
        action_output = torch.autograd.Variable(torch.mean(action_output, 0, keepdim=True).cuda())

    return action_output