import torch


def atomic(species, coordinates, model, AEVC, device=None):
    """
    Compute atomic contributions.

    Parameters
    ----------
    species: np.array
        Atomic species (mapped to indices)
    coordinates: np.array
        Atomic coordinates
    model: torch.nn.Module
        Trained model
    AEVC: torchani.AEVComputer
        AEV computer
    device:
        Computation device

    Returns
    -------
    np.array
        Atomic contributions
    """

    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Move data to device and add batch dimension
    species = species.to(device).unsqueeze(0)
    coordinates = (
        coordinates.clone().detach().requires_grad_(True).to(device).unsqueeze(0)
    )

    aevs = AEVC.forward((species, coordinates)).aevs

    atomic_contributions = model._forward_atomic(species, aevs)

    assert atomic_contributions.shape == species.shape

    return atomic_contributions