import torch
from typing import Dict
from typing import Any


def compute_feature_loss(feature_logits: torch.Tensor, feature_targets: torch.Tensor,
                         feature_dimensions: Dict[Any, int]):
    """Computes losses on the given features.

    Parameters
    ----------
    feature_logits : torch.Tensor
        2-d tensor of feature logits
    feature_targets : torch.Tensor
        1-d integer tensor of true feature labels
    feature_dimensions : Dict[str, int]
        A list of integers representing the dimension of each feature.

    Returns
    -------
    losses : torch.Tensor
        An array of losses on each feature
    accuracies : torch.Tensor
        An array of accuracies on each feature
    labels : torch.Tensor
        The provided array of targets
    predictions : torch.Tensor
        An array of predicted labels according to arg-max predictions.
    """
    current_offset = 0

    losses = feature_logits.new_empty(len(feature_dimensions))
    accuracies = feature_logits.new_empty(len(feature_dimensions))

    predictions = feature_logits.new_empty((feature_logits.shape[0], len(feature_dimensions)), dtype=torch.int64)

    for i, feature_dim in enumerate(feature_dimensions.values()):
        logits = torch.narrow(feature_logits, 1, current_offset, feature_dim)
        prediction = torch.argmax(logits, dim=-1)

        accuracies[i] = (prediction == feature_targets[:, i]).float().mean()
        losses[i] = torch.nn.functional.cross_entropy(logits, feature_targets[:, i], reduction='sum')
        predictions[:, i] = prediction

        current_offset += feature_dim

    labels = feature_targets

    return losses, accuracies, labels, predictions