import torch


def _calculate_aggregation_loss_known(logits_aggregation, aggregate_mask,
                                aggregation_function_id, config):
    """Calculates aggregation loss when its type is known during training.
    In the weakly supervised setting, the only known information is that for
    cell selection examples, "no aggregation" should be predicted. For other
    examples (those that require aggregation), no loss is accumulated.
    In the setting where aggregation type is always known, standard cross entropy
    loss is accumulated for all examples.
    Args:
        logits_aggregation: torch.FloatTensor[batch_size, num_aggregation_labels]
        aggregate_mask: torch.FloatTensor[batch_size]
        aggregation_function_id: torch.LongTensor[batch_size]
    Returns:
        aggregation_loss_known: torch.FloatTensor[batch_size, num_aggregation_labels]
    """
    if config.use_answer_as_supervision:
        # Prepare "no aggregation" targets for cell selection examples.
        target_aggregation = torch.zeros_like(aggregate_mask, dtype=torch.long)
    else:
        # Use aggregation supervision as the target.
        target_aggregation = aggregation_function_id

    batch_size = aggregate_mask.size()[0]

    one_hot_labels = torch.nn.functional.one_hot(target_aggregation,
                                                 num_classes=config.num_aggregation_labels).type(torch.float32)
    log_probs = torch.nn.functional.log_softmax(logits_aggregation, dim=-1)

    # torch.FloatTensor[batch_size]
    per_example_aggregation_intermediate = -torch.sum(
        one_hot_labels * log_probs, dim=-1)
    if config.use_answer_as_supervision:
        # Accumulate loss only for examples requiring cell selection
        # (no aggregation).
        return per_example_aggregation_intermediate * (1 - aggregate_mask)
    else:
        return per_example_aggregation_intermediate