from torch import Tensor, FloatTensor, LongTensor

from torch_geometric.data import Data

from torch.utils.data import Dataset

from torch import is_tensor

from numpy import load, array, int64

from os import listdir

from Extracting_Planes import Extract_And_Convert

from time import time

class PairData(Data):

  r"""

    PyTorch Geometric data class used for the ImageCHD dataset.

  """

  def __inc__(self, key, value, *args, **kwargs):

      if key == 'edge_index':

          return self.x.size(0)

      return super().__inc__(key, value, *args, **kwargs)

  def __init__(self,

               x: Tensor | None = None,

               edge_index: Tensor | None = None,

               edge_attr: Tensor | None = None,

               y: Tensor | int | float | None = None,

               pos: Tensor | None = None,

               time: Tensor | None = None,

               **kwargs):

     r"""

        Arguments for proper ImageCHD processing:

            x (Tensor): Source coronary-CT image as graph.

            y (Tensor): Ground truth segmentation as graph.

            edge_index (Tensor): Adjacency matrix.

            adj_count (int): Source image width.

     """

     super().__init__(x, edge_index, edge_attr, y, pos, time, **kwargs)

class CHD_Dataset(Dataset):

  r"""

    PyTorch dataset class used for the ImageCHD dataset.

  """

  def __init__(self, metadata, directory, adjacency):

    r"""

    Arguments:

        metadata (DataFrame): Pandas DataFrame containing dataset information.

        directory (string):   Path to the directory of the dataset.

        adjacency (Dictionary): Dictionary of adjacency matrices.

    """

    self.metadata = metadata

    self.directory = directory

    self.image_dir = directory + 'IMAGES/'

    self.label_dir = directory + 'LABELS/'

    self.adjacency = adjacency

  def __len__(self):

    return len(self.metadata)

  def __getitem__(self, idx):

    if is_tensor(idx):

      idx = idx.tolist()

    image, label = Extract_And_Convert(path_to_image = self.image_dir \

                                        + str(self.metadata['index'][idx]) + '.nii.gz',

                                       path_to_label = self.label_dir \

                                        + str(self.metadata['index'][idx]) + '.nii.gz',

                                       plane_type = self.metadata['Type'][idx],

                                       plane_index = self.metadata['Indice'][idx])

    # start = time()

    adj_matrix = self.adjacency[str(self.metadata['Adjacency_count'][idx])]

    # print('Adj_matrix assign time: ', time() - start)

    # start = time()

    sample = PairData(x = FloatTensor(image),

                      edge_index = LongTensor(adj_matrix),

                      y = LongTensor(label),

                      adj_count = self.metadata['Adjacency_count'][idx])

    # print('Sample preparation time: ', time() - start)

    return sample

  def get(self, idx):

    return self.__getitem__(idx)

def __Load_Adjacency__(path):

  files = listdir(path)

  adjacency = {}

  for f in files:

    adjacency[f.split('_')[2].split('.')[0]] = array(load(path + f), dtype = int64)

  return adjacency