import torch


def proportion_weighting(
    spikes: torch.Tensor,
    assignments: torch.Tensor,
    proportions: torch.Tensor,
    n_labels: int,
) -> torch.Tensor:
    # language=rst
    """
    Classify data with the label with highest average spiking activity over all neurons,
    weighted by class-wise proportion.

    :param spikes: Binary tensor of shape ``(n_samples, time, n_neurons)`` of a single
        layer's spiking activity.
    :param assignments: A vector of shape ``(n_neurons,)`` of neuron label assignments.
    :param proportions: A matrix of shape ``(n_neurons, n_labels)`` giving the per-class
        proportions of neuron spiking activity.
    :param n_labels: The number of target labels in the data.
    :return: Predictions tensor of shape ``(n_samples,)`` resulting from the "proportion
        weighting" classification scheme.
    """
    n_samples = spikes.size(0)

    # Sum over time dimension (spike ordering doesn't matter).
    spikes = spikes.sum(1)

    rates = torch.zeros((n_samples, n_labels), device=spikes.device)
    for i in range(n_labels):
        # Count the number of neurons with this label assignment.
        n_assigns = torch.sum(assignments == i).float()

        if n_assigns > 0:
            # Get indices of samples with this label.
            indices = torch.nonzero(assignments == i).view(-1)

            # Compute layer-wise firing rate for this label.
            rates[:, i] += (
                torch.sum((proportions[:, i] * spikes)[:, indices], 1) / n_assigns
            )

    # Predictions are arg-max of layer-wise firing rates.
    predictions = torch.sort(rates, dim=1, descending=True)[1][:, 0]

    return predictions