from __future__ import print_function, division

import os

import torch

import pandas as pd

from skimage import io, transform

import numpy as np

from torch.utils.data import Dataset, DataLoader

from torchvision import transforms, utils

from PIL import Image, ImageOps

from random import random, randint

# Ignore warnings

import warnings

import pdb

warnings.filterwarnings("ignore")

def make_dataset(root,mode):

  """   Takes in the root directory and mode(train or val or test) as inputs

  then joins the path with the folder of the specified mode.Applies normalize

  function to each img of the folder and returns a list of tuples containing

  image and its corresponding mask.

  Returns

  -------

  tuple : Normalized image and its annotation.

  """

  assert mode in ['train', 'val', 'test']

  items = []

  if mode == 'train':

      train_img_path = os.path.join(root, 'Train/train_image')

      train_mask_path = os.path.join(root, 'Train/train_mask')

      images = os.listdir(train_img_path)

      labels = os.listdir(train_mask_path)

      images.sort()

      labels.sort()

      for it_im, it_gt in zip(images, labels):

          item = (os.path.join(train_img_path, it_im), os.path.join(train_mask_path, it_gt))

          items.append(item)

  elif mode == 'val':

      val_img_path = os.path.join(root, 'Val/val_img')

      val_mask_path = os.path.join(root, 'Val/val_mask')

      images = os.listdir(val_img_path)

      labels = os.listdir(val_mask_path)

      images.sort()

      labels.sort()

      for it_im, it_gt in zip(images, labels):

          item = (os.path.join(val_img_path, it_im), os.path.join(val_mask_path, it_gt))

          items.append(item)

  else:

      test_img_path = os.path.join(root, 'Test/test_img')

     # test_mask_path = os.path.join(root, 'Test/test_mask')

      images = os.listdir(test_img_path)

      #labels = os.listdir(test_mask_path)

      images.sort()

      #labels.sort()

      for it_im in images:

          item = os.path.join(test_img_path, it_im)

          items.append(item)

  return items

class MedicalImageDataset(Dataset):

    """ GI dataset."""

    def __init__(self, mode, root_dir, transform=None):

        """

        Args:

            csv_file (string): Path to the csv file with annotations.

            root_dir (string): Directory with all the images.

            transform (callable, optional): Optional transform to be applied

                on a sample.

        """

        self.root_dir = root_dir

        self.mode=mode

        self.transform = transform

        self.imgs = make_dataset(root_dir, mode)

    def __len__(self):

        return len(self.imgs)

    def __getitem__(self,index):

      if self.mode== 'test':

         img_path=self.imgs[index]

         img =np.array( Image.open(img_path))

         img_shape=np.array(img).shape

         if self.transform:

            augmented = self.transform(image=img)

            image=augmented["image"]

         return image

      else:

        img_path, mask_path = self.imgs[index]

        # print("{} and {}".format(img_path,mask_path))

        img = np.array(Image.open(img_path))  # .convert('RGB')

        # mask = Image.open(mask_path)  # .convert('RGB')

        # img = Image.open(img_path).convert('L')

        mask = np.array(Image.open(mask_path).convert('L'))

        # print('{} and {}'.format(img_path,mask_path))

        if self.transform:

            augmented = self.transform(image=img,mask=mask)

            image=augmented["image"]

            mask=augmented["mask"]

        return [image, mask]