import torch
from typing import Dict
from typing import Tuple


def ngram(
    spikes: torch.Tensor,
    ngram_scores: Dict[Tuple[int, ...], torch.Tensor],
    n_labels: int,
    n: int,
) -> torch.Tensor:
    # language=rst
    """
    Predicts between ``n_labels`` using ``ngram_scores``.

    :param spikes: Spikes of shape ``(n_examples, time, n_neurons)``.
    :param ngram_scores: Previously recorded scores to update.
    :param n_labels: The number of target labels in the data.
    :param n: The max size of n-gram to use.
    :return: Predictions per example.
    """
    predictions = []
    for activity in spikes:
        score = torch.zeros(n_labels, device=spikes.device)

        # Aggregate all of the firing neurons' indices
        fire_order = []
        for t in range(activity.size()[0]):
            ordering = torch.nonzero(activity[t].view(-1))
            if ordering.numel() > 0:
                fire_order += ordering[:, 0].tolist()

        # Consider all n-gram sequences.
        for j in range(len(fire_order) - n):
            if tuple(fire_order[j : j + n]) in ngram_scores:
                score += ngram_scores[tuple(fire_order[j : j + n])]

        predictions.append(torch.argmax(score))

    return torch.tensor(predictions, device=spikes.device).long()