import torch


def get_error_instance_ids_by_class(model, batch_ids, batched_evaluation_data, batched_evaluation_target,
                                    device="cpu"):
    """
    Return the instance ids corresponding to errors of the model by class.

    Args:
        model: The model being evaluated.
        batched_evaluation_data: A single batch of input data to be passed to our model.
        batched_evaluation_target: A single batch of the corresponding output targets.
        device: A string with values either "cpu" or "cuda" to indicate the
            device that Pytorch is performing training on. By default this
            value is "cpu". But in case your models reside on the GPU, make sure
            to set this to "cuda". This makes sure that the input and target
            tensors are transferred to the GPU during training.
    Returns:
        A dictionary of key / value pairs, where the key is the output class
        and the value is the list of instance ids corresponding to misclassification
        errors of the model within that class.
    """
    with torch.no_grad():
        if device != "cpu":
            batched_evaluation_data = batched_evaluation_data.to(device)
            batched_evaluation_target = batched_evaluation_target.to(device)
        _, _, model_output_logsoftmax = model(batched_evaluation_data)
        class_error_instance_ids = {}
        model_diff = (torch.argmax(model_output_logsoftmax, 1) - batched_evaluation_target).float()
        model_errors = model_diff.nonzero().view(-1).tolist()
        if len(model_errors) > 0:
            error_instance_indices = torch.tensor(batch_ids).index_select(
                0, torch.tensor(model_errors)).tolist()
            target_error_classes = batched_evaluation_target[model_errors].view(-1).tolist()
            for class_label in set(batched_evaluation_target.view(-1).tolist()):
                class_error_instance_ids[class_label] = []
            for error_instance_index, error_class in zip(error_instance_indices, target_error_classes):
                class_error_instance_ids[error_class].append(error_instance_index)

        return class_error_instance_ids