text: import os
  import pickle
  import numpy as np
  import pandas as pd
  from sklearn.model_selection import train_test_split
  
  import torch
  from torch.utils.data import Dataset
  
  from pysnn.file_io import Events, read_2d_spikes
  
  
  #########################################################
  # Utility functions
  #########################################################
  def _train_test_split_classification(datasets, labels=None, train_size=0.8, seed=42):
      if not isinstance(datasets, (list, tuple)):
          datasets = [datasets]
          labels = [labels]
      train = []
      test = []
      lbls = ["sample", "label"]
  
      # Split each data subset into desired train test
      for idx, dset in enumerate(datasets):
          if labels[idx] is not None:
              x_trn, x_tst, y_trn, y_tst = train_test_split(
                  dset, labels[idx], train_size=train_size, random_state=seed
              )
              trn = pd.DataFrame({lbls[0]: x_trn, lbls[1]: y_trn})
              tst = pd.DataFrame({lbls[0]: x_tst, lbls[1]: y_tst})
          else:
              trn, tst = train_test_split(dset, train_size=train_size, random_state=seed)
              if not isinstance(trn, pd.DataFrame):
                  trn = pd.DataFrame(trn, columns=lbls)
                  tst = pd.DataFrame(tst, columns=lbls)
              else:
                  trn.columns = lbls
                  tst.columns = lbls
                  trn.reset_index(inplace=True, drop=True)
                  tst.reset_index(inplace=True, drop=True)
  
          train.append(trn)
          test.append(tst)
  
      train = pd.concat(train)
      test = pd.concat(test)
  
      return train, test
  
  
  def _list_dir_content(root_dir):
      content = {}
      for root, dirs, files in os.walk(root_dir):
          for label, name in enumerate(dirs):
              subdir = os.path.join(root, name)
              dir_content = os.listdir(subdir)
              content[name] = [os.path.join(subdir, im) for im in dir_content]
      return content
  
  
  def _concat_dir_content(content):
      ims = []
      labels = []
      names = []
      for idx, (name, data) in enumerate(content.items()):
          if not isinstance(data, (list, tuple)):
              data = [data]
          ims += data
          labels += [idx for _ in range(len(data))]
          names += [name for _ in range(len(data))]
      df = pd.DataFrame({"sample": ims, "label": labels})
      return df, names
  
  
  def train_test(root_dir, train_size=0.8, seed=42):
      r"""Split dataset into train and test sets.
      
      Takes in a directory where it looks for sub-directories. Content of each directory is split into train and test subsets.
  
      :param root_dir: Directory containing data.
      :param train_size: Percentage of the data to be assigned to training set, 1 - train_size is assigned as test set.
      :param seed: Seed for random number generator.
      """
      content = _list_dir_content(root_dir)
      data, _ = _concat_dir_content(content)
      train, test = _train_test_split_classification(
          data, train_size=train_size, seed=seed
      )
      return train, test
  
  
  class NeuromorphicDataset(Dataset):
      r"""Class that wraps around several neuromorphic datasets.
  
      The class adheres to regular PyTorch dataset conventions.
  
      :param data: Data for dataset formatted in a pd.DataFrame.
      :param sampling_time: Duration of interval between samples.
      :param sample_length: Total duration of a single sample.
      :param height: Number of pixels in height direction.
      :param width: Number of pixels in width direction.
      :param im_transform: Image transforms, same convention as for PyTorch datasets.
      :param lbl_transform: Lable transforms, same convention as for PyTorch datasets.
      """
  
      def __init__(
          self,
          data,
          sampling_time,
          sample_length,
          height,
          width,
          im_transform=None,
          lbl_transform=None,
      ):
          self.data = data
          self.im_transform = im_transform
          self.lbl_transform = lbl_transform
  
          self.sampling_time = sampling_time
          self.sample_length = sample_length
          self.height = height
          self.width = width
          self.n_time_bins = int(sample_length / sampling_time)
          self.im_template = torch.zeros((2, height, width, self.n_time_bins))
  
      def __len__(self):
          return len(self.data)
  
      def __getitem__(self, idx):
          # Image sample
          im_name = self.data.iloc[idx, 0]
          input_spikes = read_2d_spikes(im_name)
          sample = input_spikes.to_spike_tensor(
              self.im_template, sampling_time=self.sampling_time
          )
  
          # Label
          label = self.data.iloc[idx, 1]
  
          # Apply transforms
          if self.im_transform:
              sample = self.im_transform(sample)
          if self.lbl_transform:
              label = self.lbl_transform(label)
  
          return (sample, label)
  
  
  #########################################################
  # Neurmorphic Caltech 101
  #########################################################
  def ncaltech_train_test(
      root,
      sampling_time=1,
      sample_length=300,
      height=200,
      width=300,
      im_transform=None,
      lbl_transform=None,
  ):
      r"""Neurmorphic version of the Caltech-101 dataset, obtained from:
  
          https://www.garrickorchard.com/datasets/n-caltech101
      
      'Converting Static Image Datasets to Spiking Neuromorphic Datasets Using Saccades' by G. Orchard et al.
  
      :return: :class:`NeuromorphicDataset` for both and training and test data.
      """
      train, test = train_test(root)
      train_dataset = NeuromorphicDataset(
          train,
          sampling_time,
          sample_length,
          height,
          width,
          im_transform=im_transform,
          lbl_transform=lbl_transform,
      )
      test_dataset = NeuromorphicDataset(
          test,
          sampling_time,
          sample_length,
          height,
          width,
          im_transform=im_transform,
          lbl_transform=lbl_transform,
      )
      return train_dataset, test_dataset
  
  
  #########################################################
  # Neuromorphic MNIST
  #########################################################
  def nmnist_train_test(
      root,
      sampling_time=1,
      sample_length=300,
      height=34,
      width=34,
      im_transform=None,
      lbl_transform=None,
  ):
      r"""Neurmorphic version of the MNIST dataset, obtained from:
  
          https://www.garrickorchard.com/datasets/n-mnist
      
      'Converting Static Image Datasets to Spiking Neuromorphic Datasets Using Saccades' by G. Orchard et al.
  
      :return: :class:`NeuromorphicDataset` for both and training and test data.
      """
      train_content = _list_dir_content(os.path.join(root, "Train"))
      train, _ = _concat_dir_content(train_content)
      train_dataset = NeuromorphicDataset(
          train,
          sampling_time,
          sample_length,
          height,
          width,
          im_transform=im_transform,
          lbl_transform=lbl_transform,
      )
  
      test_content = _list_dir_content(os.path.join(root, "Test"))
      test, _ = _concat_dir_content(test_content)
      test_dataset = NeuromorphicDataset(
          test,
          sampling_time,
          sample_length,
          height,
          width,
          im_transform=im_transform,
          lbl_transform=lbl_transform,
      )
  
      return train_dataset, test_dataset
  
  
  #########################################################
  # Neuromorphic MNIST
  #########################################################
  def ncars_train_test(
      root,
      sampling_time=1,
      sample_length=100,
      height=100,
      width=120,
      im_transform=None,
      lbl_transform=None,
  ):
      r"""Neurmorphic dataset containing images of cars or background, obtained from:
  
          https://www.prophesee.ai/dataset-n-cars/
      
      This is a two class problem.
      
      :return: :class:`NeuromorphicDataset` for both and training and test data.
      """
      train_content = _list_dir_content(os.path.join(root, "train"))
      train, _ = _concat_dir_content(train_content)
      train_dataset = NeuromorphicDataset(
          train,
          sampling_time,
          sample_length,
          height,
          width,
          im_transform=im_transform,
          lbl_transform=lbl_transform,
      )
  
      test_content = _list_dir_content(os.path.join(root, "test"))
      test, _ = _concat_dir_content(test_content)
      test_dataset = NeuromorphicDataset(
          test,
          sampling_time,
          sample_length,
          height,
          width,
          im_transform=im_transform,
          lbl_transform=lbl_transform,
      )
  
      return train_dataset, test_dataset
  
  
  ############################
  # Boolean
  ############################
  class Boolean(Dataset):
      r"""Dataset for generating event-based, boolean data samples. Can be used to construct AND, OR, XOR datasets.
      
      :param data_encoder: ``Encoder`` to convert scalar data values to spiketrains.
      :param data_transform: Image transforms, same convention as for PyTorch datasets.
      :param lbl_transform: Label transforms, same convention as for PyTorch datasets.
      :repeats: Number of times to repeat the 4 samples within a single iteration of the dataset, total is 4 * repeats.
      """
  
      def __init__(
          self, data_encoder=None, data_transform=None, lbl_transform=None, repeats=1
      ):
          self.data_encoder = data_encoder
          self.data_transform = data_transform
          self.lbl_transform = lbl_transform
          self.data = torch.tensor([[0, 0], [1, 0], [0, 1], [1, 1]], dtype=torch.float)
          self.data = torch.repeat_interleave(self.data, int(repeats), dim=1)
  
      def __len__(self):
          return len(self.data)
  
      def __getitem__(self, idx):
          sample = self.data[idx].unsqueeze(0)
          label = self.labels[idx]
  
          # Sample transforms
          if self.data_transform:
              sample = self.data_transform(sample)
          if self.data_encoder:
              sample = self.data_encoder(sample)
  
          # Label transforms
          if self.lbl_transform:
              label = self.lbl_transform(label)
  
          return sample, label
  
  
  class XOR(Boolean):
      r"""XOR dataset, inherits directly from :class:`Boolean`."""
  
      def __init__(
          self, data_encoder=None, data_transform=None, lbl_transform=None, repeats=1
      ):
          super(XOR, self).__init__(data_encoder, data_transform, lbl_transform, repeats)
          self.labels = torch.tensor([[0], [1], [1], [0]])
  
  
  class AND(Boolean):
      r"""AND dataset, inherits directly from :class:`Boolean`."""
  
      def __init__(
          self, data_encoder=None, data_transform=None, lbl_transform=None, repeats=1
      ):
          super(AND, self).__init__(data_encoder, data_transform, lbl_transform, repeats)
          self.labels = torch.tensor([[0], [0], [0], [1]])
  
  
  class OR(Boolean):
      r"""OR dataset, inherits directly from :class:`Boolean`."""
  
      def __init__(
          self, data_encoder=None, data_transform=None, lbl_transform=None, repeats=1
      ):
          super(OR, self).__init__(data_encoder, data_transform, lbl_transform, repeats)
          self.labels = torch.tensor([[0], [1], [1], [1]])
  
  
  ############################
  # transforms
  ############################
  class DiscretizeFloat:
      r"""Discretize float to nearest integer values, returns a torch.Float."""
  
      def __call__(self, tensor):
          tensor = tensor.int()
          return tensor.float()
  
  
  class BooleanNoise:
      r"""Add noise to the integer boolean values, e.g. 1 -> 0.8 and 0 -> 0.2
      
      :param low_thresh: Upper boundary for the negative values.
      :param high_thresh: Lower boundary for the positive values.
      """
  
      def __init__(self, low_thresh, high_thresh):
          self.low_distr = torch.distributions.Uniform(0.0, low_thresh)
          self.high_distr = torch.distributions.Uniform(high_thresh, 1.0)
  
      def __call__(self, x):
          zero_idx = x == 0
          one_idx = x == 1
          if zero_idx.any():
              x[zero_idx] = self.low_distr.sample(zero_idx.shape)
          if one_idx.any():
              x[one_idx] = self.high_distr.sample(one_idx.shape)
          return x
  
  
  class Intensity:
      r"""Multiplication with a fixed scalar value."""
  
      def __init__(self, intensity):
          self.intensity = intensity
  
      def __call__(self, x):
          return x * self.intensity
  
  
  if __name__ == "__main__":
      train = XOR()
      print(train[0][0].shape)
      print(train[0][1].shape)
      print(len(train))
  
      # ncalt_train, ncalt_test = ncaltech_train_test(root_dir + "ncaltech101")
      # print("caltech")
      # print(ncalt_train[0][0].shape)
      # print(ncalt_train[0][1])
      # print(len(ncalt_train))
  
      # nmnist_train, nmnist_test = nmnist_train_test(root_dir + "nmnist")
      # print("\nnmnist")
      # print(nmnist_train[0][0].shape)
      # print(nmnist_train[0][1])
      # print(len(nmnist_train))
  
      # ncars_train, ncars_test = ncars_train_test(root_dir + "n-cars")
      # print("\nncars")
      # print(ncars_train[0][0].shape)
      # print(ncars_train[0][1])
      # print(len(ncars_train))
  