text: """
  Defines DataLoader classes to use.
  """
  import logging
  from collections import defaultdict
  from itertools import groupby
  
  import torch
  import torch.nn as nn
  import numpy as np
  from torch._six import int_classes as _int_classes
  from torch.utils.data import DataLoader
  from torch.utils.data.sampler import WeightedRandomSampler, RandomSampler, Sampler
  import torch.nn.functional as F
  
  from pet_ct.util.util import flex_stack, get_batch_size, soft_to_hard
  
  def pad_inputs(inputs, max_length):
      """
      """
      inputs_length = inputs.size()[0]
      pad_length = max_length - inputs_length
      pad_dims = (0, 0, 0, 0, 0, 0, 0, pad_length) # pads first dim evenly
      padded_inputs = F.pad(inputs, pad=pad_dims)
      return padded_inputs
  
  def pad_scans(scan_inputs, value=0):
      """
      """
      max_length = max(scan_inputs, key=lambda x: x.shape[0]).shape[0]
      padded_inputs = []
      for inputs in scan_inputs:
          inputs_length = inputs.size()[0]
          pad_length = max_length - inputs_length
          pad_dims = (0, 0, 0, 0, 0, 0, 0, pad_length) # pads first dim evenly
          padded = F.pad(inputs, pad=pad_dims, value=value)
          padded_inputs.append(padded)
  
      return padded_inputs
  
  
  def pad_targets(targets, value=-1):
      max_length = max(targets, key=lambda x: x.shape[-1]).shape[-1]
      padded_targets = []
      for tgt in targets:
          tgt_length = tgt.shape[-1]
          pad_length = max_length - tgt_length
          pad_dims = (0, pad_length)
          padded_targets.append(F.pad(tgt, pad=pad_dims, value=value))
      return padded_targets
  
  
  def exam_collate(batch_list):
      """ Collate function fro a multi-task dataset.
      args:
          exam_list (list) list of exams
      """
      max_length = max(batch_list, key=lambda x: x[0].size())[0].size()[0]
      all_inputs = []
      all_targets = []
      all_info = []
  
      for inputs, targets, info in batch_list:
          all_inputs.append(pad_inputs(inputs, max_length))
          all_info.append(info)
          all_targets.append(targets)
  
      # stack targets and inputs
      all_targets = flex_stack(all_targets, dim=0)
      all_inputs = flex_stack(all_inputs, dim=0)
      return all_inputs, all_targets, all_info
  
  
  def mt_exam_collate(batch_list):
      """ Collate function for a multi-task dataset.
      args:
          exam_list (list) list of exams
      """
      max_exam = max(batch_list, key=lambda x: x[0].size()[0])
      max_length = max_exam[0].size()[0]
  
      all_inputs = []
      all_targets = defaultdict(list)
      all_info = []
  
      for inputs, targets, info in batch_list:
          all_inputs.append(pad_inputs(inputs, max_length))
          all_info.append(info)
          for task, target in targets.items():
              all_targets[task].append(target)
  
      # stack targets and inputs
      all_targets = {task: flex_stack(targets, dim=0) for task, targets in all_targets.items()}
      all_inputs = flex_stack(all_inputs, dim=0)
      return all_inputs, all_targets, all_info
  
  
  def mt_mi_exam_collate(batch_list):
      """ Collate function for a multi-task multi-input dataset
      """
      all_inputs = defaultdict(list)
      all_targets = defaultdict(list)
      all_info = []
  
      for inputs, targets, info in batch_list:
          all_info.append(info)
  
          for name, inpt in inputs.items():
              all_inputs[name].append(inpt)
  
          for task, target in targets.items():
              all_targets[task].append(target)
  
      if "scan" in all_inputs:
          all_inputs["scan"] = pad_scans(all_inputs["scan"], value=-1)
  
      if "mlm" in all_targets:
          all_targets["mlm"] = pad_targets(all_targets["mlm"], value=-1)
      elif "fdg_abnorm" in all_targets:
          all_targets["fdg_abnorm"] = pad_targets(all_targets["fdg_abnorm"], value=-1)
  
      # stack targets and inputs
      all_targets = {task: flex_stack(targets, dim=0) for task, targets in all_targets.items()}
      all_inputs = {name: flex_stack(inputs, dim=0) for name, inputs in all_inputs.items()}
      return all_inputs, all_targets, all_info
  
  
  class ExamBatchSampler(Sampler):
  
      def __init__(self, batch_size, num_slices, sampler=None,
                   weights=None, num_samples=None, replacement=None, shuffle=None,
                   drop_last=False):
          """
          Creates batches of exams with same number of slices.
  
          TODO: Implement `drop_last`.
          """
          if not isinstance(num_samples, _int_classes) or isinstance(num_samples, bool) or \
                  num_samples <= 0:
              raise ValueError("num_samples should be a positive integer "
                               "value, but got num_samples={}".format(num_samples))
          if not isinstance(replacement, bool):
              raise ValueError("replacement should be a boolean value, but got "
                               "replacement={}".format(replacement))
          if sampler is not None and shuffle:
              raise ValueError("sampler option is mutually exclusive with shuffle")
          if not isinstance(batch_size, _int_classes) or isinstance(batch_size, bool) or \
                  batch_size <= 0:
              raise ValueError("batch_size should be a positive integer value, "
                               "but got batch_size={}".format(batch_size))
  
          self.batch_size = batch_size
          self.num_samples = num_samples
          self.num_slices = num_slices
  
          self.sampler = sampler
          if self.sampler is not None:
              self.weights = weights if weights is not None else torch.ones(len(num_slices))
          else:
              self.weights = torch.ones(len(num_slices))
  
          self.replacement = replacement
          self.shuffle = shuffle
  
          self.drop_last = drop_last
  
      def __iter__(self):
          """
          """
          if self.sampler is not None:
              samples = torch.multinomial(self.weights, self.num_samples,
                                          replacement=self.replacement)
          else:
              if self.shuffle:
                  samples = torch.multinomial(self.weights, self.num_samples,
                                              replacement=False)
              else:
                  samples = torch.tensor(range(self.num_samples))
          samples = sorted(samples, key=lambda idx: self.num_slices[idx])
  
          curr_iter = 0
          batches = []
          while curr_iter < self.num_samples:
              batch = [samples[curr_iter]]
              batch_slices = self.num_slices[samples[curr_iter]]
  
              offset = self.batch_size
              for i in range(1, self.batch_size):
                  if curr_iter + i < self.num_samples and \
                      batch_slices == self.num_slices[samples[curr_iter + i]]:
                      batch.append(samples[curr_iter + i])
                  else:
                      offset = i
                      break
              batches.append(batch)
              curr_iter = curr_iter + offset
  
          batch_idxs = torch.randperm(len(batches)).tolist()
          for batch_idx in batch_idxs:
              yield batches[batch_idx]
  
      def __len__(self):
          """
          This is approximate because we cannot know number of batches a priori.
          """
          if self.drop_last:
              return self.num_samples // self.batch_size
          else:
              return (self.num_samples + self.batch_size - 1) // self.batch_size
  
  
  
  class ExamDataLoader(DataLoader):
      def __init__(self,
                   dataset,
                   batch_size=1,
                   shuffle=False,
                   num_workers=6,
                   sampler=None,
                   num_samples=1000,
                   replacement=False,
                   class_probs=None,
                   pin_memory=False):
          """
          """
          # get example weights so examples are sampled according to class_probs
          if sampler in {"WeightedRandomSampler", "RandomSampler"}:
              if sampler == "WeightedRandomSampler":
                  classes = dataset.get_targets()
                  counts = torch.bincount(classes)
                  weights = torch.zeros_like(classes, dtype=torch.float)
                  for example_idx, class_idx in enumerate(classes):
                      class_prob = class_probs[class_idx] / float(counts[class_idx])
                      weights[example_idx] = class_prob
                  sampler = WeightedRandomSampler(weights=weights, num_samples=num_samples,
                                                  replacement=replacement)
              elif sampler == "RandomSampler":
                  weights = None
                  sampler = RandomSampler(data_source=dataset, num_samples=num_samples,
                                          replacement=True)
          elif sampler is not None:
              raise ValueError(f"Sampler {sampler} not supported.")
          else:
              num_samples = len(dataset)
              weights = None
  
          if batch_size > 1:
              num_slices = dataset.get_num_slices()
              batch_sampler = ExamBatchSampler(batch_size, num_slices,
                                               sampler=sampler,
                                               weights=weights,
                                               num_samples=num_samples,
                                               replacement=replacement,
                                               shuffle=shuffle,
                                               drop_last=False)
  
              super().__init__(dataset=dataset, num_workers=num_workers,
                               batch_sampler=batch_sampler, pin_memory=pin_memory,
                               collate_fn=mt_exam_collate)
          else:
              super().__init__(dataset=dataset, batch_size=batch_size, shuffle=shuffle,
                               num_workers=num_workers, sampler=sampler, pin_memory=pin_memory,
                               collate_fn=mt_exam_collate)
  
  
  class MTExamDataLoader(DataLoader):
      def __init__(self,
                   dataset,
                   batch_size=1,
                   shuffle=False,
                   num_workers=6,
                   sampler=None,
                   num_samples=1000,
                   replacement=False,
                   weight_task=None,
                   class_probs=None,
                   pin_memory=False):
          """
          """
          if sampler in {"WeightedRandomSampler", "RandomSampler"}:
              # get example weights so examples are sampled according to class_probs
              if sampler == "WeightedRandomSampler":
                  classes = []
                  for target in dataset.get_targets(tasks=[weight_task], hard=True):
                      target_class = target[weight_task]
                      classes.append(target_class)
                  classes = torch.stack(classes)
                  if classes.shape[-1] > 1:
                      classes = soft_to_hard(classes, break_ties="random").long()
                  classes = torch.LongTensor(classes)
                  counts = torch.bincount(classes)
                  weights = torch.zeros_like(classes, dtype=torch.float)
                  for example_idx, class_idx in enumerate(classes):
                      class_prob = class_probs[class_idx] / float(counts[class_idx])
                      weights[example_idx] = class_prob
                  sampler = WeightedRandomSampler(weights=weights, num_samples=num_samples,
                                                  replacement=replacement)
              elif sampler == "RandomSampler":
                  weights = None
                  sampler = RandomSampler(data_source=dataset, num_samples=num_samples, 
                                          replacement=True)
          elif sampler is not None:
              raise ValueError(f"Sampler {sampler} not supported.")
          else:
              num_samples = len(dataset)
              weights = None
  
          if batch_size > 1:
              num_slices = dataset.get_num_slices()
              batch_sampler = ExamBatchSampler(batch_size, num_slices,
                                               sampler=sampler,
                                               weights=weights,
                                               num_samples=num_samples,
                                               replacement=replacement,
                                               shuffle=shuffle,
                                               drop_last=False)
  
              super().__init__(dataset=dataset, num_workers=num_workers,
                               batch_sampler=batch_sampler, pin_memory=pin_memory,
                               collate_fn=mt_exam_collate)
          else:
              super().__init__(dataset=dataset, batch_size=batch_size, shuffle=shuffle,
                               num_workers=num_workers, sampler=sampler, pin_memory=pin_memory,
                               collate_fn=mt_exam_collate)
  
  class MTAdaptiveDataloader(DataLoader):
  
      def __init__(self, dataset,
                   batch_size=1,
                   shuffle=False,
                   num_workers=6,
                   num_samples=1000,
                   weight_task=None,
                   pin_memory=False):
          """
          """
          self.num_samples = num_samples
          sampler = RandomSampler(data_source=dataset, num_samples=num_samples,
                                  replacement=True)
          super().__init__(dataset=dataset, batch_size=batch_size, shuffle=shuffle,
                           num_workers=num_workers, sampler=sampler)
  
      def update(self, metrics, metric="roc_auc"):
          """
          """
          for task in self.dataset.tasks:
              value = metrics.get_metric(metric=metric, task=task)
  
  
          weights = torch.zeros(len(dataset))
          for example_idx, targets in enumerate(self.dataset.get_targets(hard=True)):
              weights[example_idx] = 0
  
          sampler = WeightedRandomSampler(weights=weights, num_samples=self.num_samples, replacement=True)
          self.batch_sampler = BatchSampler(sampler, self.batch_size, self.drop_last)
  
  
  class MTMIExamDataLoader(DataLoader):
      def __init__(self,
                   dataset,
                   batch_size=1,
                   shuffle=False,
                   num_workers=6,
                   sampler=None,
                   num_samples=1000,
                   replacement=False,
                   weight_task=None,
                   class_probs=None,
                   pin_memory=False):
          """
          """
          # get example weights so examples are sampled according to class_probs
          if sampler == "WeightedRandomSampler":
              self.num_samples = int(round(num_samples / batch_size))
              classes = torch.LongTensor([target[weight_task]
                                          for target in dataset.get_targets([weight_task])])
  
              counts = torch.bincount(classes)
              weights = torch.zeros_like(classes, dtype=torch.float)
              for example_idx, class_idx in enumerate(classes):
                  class_prob = class_probs[class_idx] / float(counts[class_idx])
                  weights[example_idx] = class_prob
              sampler = WeightedRandomSampler(weights=weights, num_samples=num_samples,
                                              replacement=replacement)
          elif sampler == "RandomSampler":
              self.num_samples = int(round(num_samples / batch_size))
              sampler = RandomSampler(data_source=dataset, num_samples=num_samples,
                                      replacement=True)
          elif sampler is not None:
              raise ValueError(f"Sampler {sampler} not supported.")
  
          super().__init__(dataset=dataset, batch_size=batch_size, shuffle=shuffle,
                           num_workers=num_workers, sampler=sampler, pin_memory=pin_memory,
                           collate_fn=mt_mi_exam_collate)
  
      def __len__(self):
          """
          """
          if hasattr(self, "num_samples"):
              return self.num_samples
          else:
              return super().__len__()
  
  
  class DynamicRandomSampler(Sampler):
  
      def __init__(self, weights, num_samples, replacement=False):
          if not isinstance(num_samples, int) or isinstance(num_samples, bool) or \
                  num_samples <= 0:
              raise ValueError("num_samples should be a positive integer "
                               "value, but got num_samples={}".format(num_samples))
          if not isinstance(replacement, bool):
              raise ValueError("replacement should be a boolean value, but got "
                               "replacement={}".format(replacement))
          self.weights = torch.as_tensor(weights, dtype=torch.double)
          self.num_samples = num_samples
          self.replacement = replacement
          if not replacement:
              self.mask = torch.ones_like(self.weights, dtype=torch.double)
  
      def __iter__(self):
          for _ in range(self.num_samples):
              weights = self.weights if self.replacement else self.weights * self.mask
              idx = int(torch.multinomial(weights, 1, self.replacement))
              if self.replacement:
                  self.mask[idx] = 0.0
              yield idx
  
      def __len__(self):
          return self.num_samples
  
  
  class MTDynamicDataLoader(DataLoader):
  
      def __init__(self, dataset,
                   batch_size=1,
                   num_workers=6,
                   num_samples=1000,
                   update_period=10,
                   replacement=False,
                   priority_metric="roc_auc",
                   priority_scale=100,
                   pin_memory=False):
          # initialize uniform weights
          self.exams = list(dataset.get_targets(hard=True))
  
          weights = torch.ones(len(self.exams))
          self.sampler = DynamicRandomSampler(weights, num_samples, replacement=replacement)
  
          self.priority_metric = priority_metric
          self.priority_scale = priority_scale
          self.task_to_priority = defaultdict(lambda: 1)
          self.task_to_samples = defaultdict(list)
          self.sample_idxs = []
          self.tasks = [{"name": task,
                         "cardinality": len(self.exams[0][task])}
                        for task in dataset.tasks]
          self._init_demands()
  
          # support for periodic updates
          self.update_counter = 0
          self.update_period = update_period
          self.queued_idxs = []
  
          super().__init__(dataset=dataset, batch_size=batch_size,
                           num_workers=num_workers, sampler=self.sampler, pin_memory=pin_memory,
                           collate_fn=mt_exam_collate)
  
      def update_epoch(self, metrics):
  
          # reset
          self.task_to_samples = defaultdict(list)
          self.sample_idxs = []
          self._init_demands()
          self.update_counter = 0
          self.queued_idxs = []
  
          # update priorities
          task_to_metric = {}
          total = 0
          for task in self.tasks:
              metric = metrics.get_metric(self.priority_metric, task["name"])
              task_to_metric[task["name"]] = metric
              total += metric
  
          self.task_to_priority = {
              task: self.priority_scale * (1 - metric/total) for task, metric in task_to_metric.items()
          }
          logging.info(task_to_metric)
          logging.info(self.task_to_priority)
  
      def update_batch(self, idxs):
          """
          """
          self.queued_idxs.extend(idxs)
          self.update_counter += len(idxs)
  
          if self.update_counter % self.update_period == 0:
              self.update_demands()
              self.update_weights()
  
  
      def _init_demands(self):
          """
          """
          self.task_to_demands = {}
          for task in self.tasks:
              self.task_to_demands[task["name"]] = np.full((1, task["cardinality"]), fill_value=1/task["cardinality"])
              self.task_to_samples[task["name"]] = []
  
      def get_loss_weights(self, targets):
          """
          """
          task_to_weight = {}
          for task, tgt in targets.items():
              tgt = int(torch.argmax(tgt).cpu().numpy())
  
              task_to_weight[task] = self.task_to_demands[task][-1, tgt]
          return task_to_weight
  
      def update_demands(self):
          """
          """
          for idx in self.queued_idxs:
              self.sample_idxs.append(idx)
              exam = self.exams[idx]
              for task, tgt in exam.items():
                  self.task_to_samples[task].append(int(torch.argmax(tgt)))
  
              for task in self.tasks:
                  samples = np.array(self.task_to_samples[task["name"]])
                  total = np.sum(self.task_to_demands[task["name"]][np.arange(samples.size), samples])
                  new_demands = []
                  for tgt in range(task["cardinality"]):
                      tgt_total = np.sum(self.task_to_demands[task["name"]]
                                         [np.where(samples == tgt), tgt])
                      new_demand = (1 - (tgt_total + 1) / (total + task["cardinality"]))
                      new_demand *= self.task_to_priority[task["name"]]
                      new_demands.append(new_demand)
                  self.task_to_demands[task["name"]] = np.append(self.task_to_demands[task["name"]],
                                                                 np.array(new_demands)[None, :], axis=0)
          self.queued_idxs = []
  
  
      def update_weights(self):
          weights = []
          for exam in self.exams:
              weight = self._get_weight(exam)
              weights.append(weight)
          self.sampler.weights = torch.nn.functional.softmax(torch.tensor(weights), dim=0).double()
  
      def _get_weight(self, exam):
          """
          """
          weight = 0
          for task in self.tasks:
              tgt = torch.argmax(exam[task["name"]])
              weight += self.task_to_demands[task["name"]][-1, tgt]
          return weight
  