################################################################################
# Copyright (c) 2021 ContinualAI.                                              #
# Copyrights licensed under the MIT License.                                   #
# See the accompanying LICENSE file for terms.                                 #
#                                                                              #
# Date: 28-05-2020                                                             #
# Author(s): Lorenzo Pellegrini                                                #
# E-mail: contact@continualai.org                                              #
# Website: avalanche.continualai.org                                           #
################################################################################

from typing import Sequence, Any
import torch
from avalanche.benchmarks.utils.dataset_definitions import \
    ISupportedClassificationDataset


def _exp_structure_from_assignment(
        dataset: ISupportedClassificationDataset[Any],
        assignment: Sequence[Sequence[int]],
        n_classes: int):
    n_experiences = len(assignment)
    exp_structure = [[0 for _ in range(n_classes)]
                     for _ in range(n_experiences)]

    for exp_id in range(n_experiences):
        exp_targets = [int(dataset.targets[pattern_idx])
                       for pattern_idx in assignment[exp_id]]
        cls_ids, cls_counts = torch.unique(torch.as_tensor(
            exp_targets), return_counts=True)

        for unique_idx in range(len(cls_ids)):
            exp_structure[exp_id][int(cls_ids[unique_idx])] += \
                int(cls_counts[unique_idx])

    return exp_structure


__all__ = [
    '_exp_structure_from_assignment'
]
