text: ################################################################################
  # Copyright (c) 2021 ContinualAI.                                              #
  # Copyrights licensed under the MIT License.                                   #
  # See the accompanying LICENSE file for terms.                                 #
  #                                                                              #
  # Date: 15-06-2021                                                             #
  # Author(s): Lorenzo Pellegrini                                                #
  # E-mail: contact@continualai.org                                              #
  # Website: avalanche.continualai.org                                           #
  ################################################################################
  
  from collections import defaultdict
  from typing import Sequence, Iterable, Dict, Optional, Iterator
  
  from avalanche.benchmarks.utils import AvalancheDataset
  
  
  class LazyDatasetSequence(Sequence[AvalancheDataset]):
      """
      A lazily initialized sequence of datasets.
  
      This class provides a way to lazily generate and store the datasets
      linked to each experience. This class uses a generator to get the sequence
      of datasets but it can also be used with a more classic statically
      initialized Sequence (like a list).
  
      This class will also keep track of the targets and task labels field of the
      generated datasets.
      """
  
      def __init__(
              self,
              experience_generator: Iterable[AvalancheDataset],
              stream_length: int):
          self._exp_source: Optional[Iterable[AvalancheDataset]] = \
              experience_generator
          """
          The source of the experiences stream, as an Iterable.
          
          Can be a simple Sequence or a Generator.
          
          This field is kept for reference and debugging. The actual generator
          is kept in the `_exp_generator` field, which stores the iterator.
          
          This field is None when if all the experiences have been loaded.
          """
  
          self._next_exp_id: int = 0
          """
          The ID of the next experience that will be generated.
          """
  
          self._loaded_experiences: Dict[int, AvalancheDataset] = dict()
          """
          The sequence of experiences obtained from the generator.
          """
  
          self._stream_length: int = stream_length
          """
          The length of the stream.
          """
          try:
              self._exp_generator: Optional[Iterator[AvalancheDataset]] = iter(
                  self._exp_source
              )
          except TypeError as e:
              if callable(self._exp_source):
                  # https://stackoverflow.com/a/17092033
                  raise ValueError(
                      'The provided generator is not iterable. When using a '
                      'generator function based on "yield", remember to pass the'
                      ' result of that function, not the '
                      'function itself!') from None
              raise e
          """
          The experience generator, as an Iterator.
          
          This field is None when if all the experiences have been loaded.
          """
  
          self.targets_field_sequence: Dict[int, Optional[Sequence]] = \
              defaultdict(lambda: None)
          """
          A dictionary mapping each experience to its `targets` field.
          
          This dictionary contains the targets field of datasets generated up to
          now, including the ones of dropped experiences.
          """
  
          self.task_labels_field_sequence: Dict[int, Optional[Sequence[int]]] = \
              defaultdict(lambda: None)
          """
          A dictionary mapping each experience to its `targets_task_labels` field.
  
          This dictionary contains the task labels of datasets generated up to
          now, including the ones of dropped experiences.
          """
  
      def __len__(self) -> int:
          """
          Gets the length of the stream (number of experiences).
  
          :return: The length of the stream.
          """
          return self._stream_length
  
      def __getitem__(self, exp_idx: int) -> AvalancheDataset:
          """
          Gets the dataset associated to an experience.
  
          :param exp_idx: The ID of the experience.
          :return: The dataset associated to the experience.
          """
          exp_idx = int(exp_idx)  # Handle single element tensors
          self.load_all_experiences(exp_idx)
          if exp_idx not in self._loaded_experiences:
              raise RuntimeError(f'Experience {exp_idx} has been dropped')
  
          return self._loaded_experiences[exp_idx]
  
      def get_experience_if_loaded(self, exp_idx: int) -> \
              Optional[AvalancheDataset]:
          """
          Gets the dataset associated to an experience.
  
          Differently from `__getitem__`, this will return None if the experience
          has not been (lazily) loaded yet.
  
          :param exp_idx: The ID of the experience.
          :return: The dataset associated to the experience or None if the
              experience has not been loaded yet or if it has been dropped.
          """
          exp_idx = int(exp_idx)  # Handle single element tensors
          if exp_idx >= len(self):
              raise IndexError(f'The stream doesn\'t contain {exp_idx+1}'
                               f'experiences')
  
          return self._loaded_experiences.get(exp_idx, None)
  
      def drop_previous_experiences(self, to_exp: int) -> None:
          """
          Drop the reference to experiences up to a certain experience ID
          (inclusive).
  
          This means that experiences with ID [0, from_exp] will be released.
          Beware that the associated object will be valid until all the references
          to it are dropped.
  
          :param to_exp: The ID of the last exp to drop (inclusive). If None,
              the whole stream will be loaded. Can be a negative number, in
              which case this method doesn't have any effect. Can be greater
              or equal to the stream length, in which case all currently loaded
              experiences will be dropped.
          :return: None
          """
  
          to_exp = int(to_exp)  # Handle single element tensors
          if to_exp < 0:
              return
  
          to_exp = min(to_exp, len(self)-1)
  
          for exp_id in range(0, to_exp+1):
              if exp_id in self._loaded_experiences:
                  del self._loaded_experiences[exp_id]
  
      def load_all_experiences(self, to_exp: int = None) -> None:
          """
          Load all experiences up to a certain experience ID (inclusive).
  
          Beware that this won't re-load any already dropped experience.
  
          :param to_exp: The ID of the last exp to load (inclusive). If None,
              the whole stream will be loaded.
          :return: None
          """
          if to_exp is None:
              to_exp = len(self) - 1
          else:
              to_exp = int(to_exp)  # Handle single element tensors
  
          if to_exp >= len(self):
              raise IndexError(f'The stream doesn\'t contain {to_exp+1}'
                               f'experiences')
  
          if self._next_exp_id > to_exp:
              # Nothing to do
              return
  
          for exp_id in range(self._next_exp_id, to_exp+1):
              try:
                  generated_exp: AvalancheDataset = next(self._exp_generator)
              except StopIteration:
                  raise RuntimeError(
                      f'Unexpected end of stream. The generator was supposed to '
                      f'generate {len(self)} experiences, but an error occurred '
                      f'while generating experience {exp_id}.')
  
              if not isinstance(generated_exp, AvalancheDataset):
                  raise ValueError(
                      'All experience datasets must be subclasses of'
                      ' AvalancheDataset')
  
              self._loaded_experiences[exp_id] = generated_exp
              self.targets_field_sequence[exp_id] = generated_exp.targets
              self.task_labels_field_sequence[exp_id] = \
                  generated_exp.targets_task_labels
              self._next_exp_id += 1
  
          if self._next_exp_id == len(self):
              # Release all references to the generator
              self._exp_generator = None
              self._exp_source = None
  
  
  __all__ = [
      'LazyDatasetSequence'
  ]
  