import os

import re

import numpy as np

from PIL import Image as PILImage

import torch

import torch.nn.functional as F

from torch.utils import data as data

from torchvision import transforms as transforms

import matplotlib.pyplot as plt

from skimage.measure import label as method_label

from skimage.measure import regionprops

# dataset for GGO and C segmentation

class CovidCTData(data.Dataset):

    def __init__(self, **kwargs):

        self.mask_type = kwargs['mask_type']

        self.ignore_ = kwargs['ignore_small']

        # ignore small areas?

        if self.ignore_:

           self.area_th = 100

        else:

           self.area_th = 1

        self.stage = kwargs['stage']

        # this returns the path to imgs dir

        self.data = kwargs['data']

        # this returns the path to

        self.gt = kwargs['gt']

        # IMPORTANT: the order of images and masks must be the same

        self.sorted_data = sorted(os.listdir(self.data))

        self.sorted_gt = sorted(os.listdir(self.gt))

        self.fname = None

        self.img_fname = None

    # this method normalizes the image and converts it to Pytorch tensor

    # Here we use pytorch transforms functionality, and Compose them together,

    def transform_img(self, img):

        # Faster R-CNN does the normalization

        t_ = transforms.Compose([

            # transforms.ToPILImage(),

            # transforms.Resize(img_size),

            transforms.ToTensor(),

        ])

        img = t_(img)

        return img

    # inputs: box coords (min_row, min_col, max_row, max_col)

    # array HxW from whic to extract a single object's mask

    # each isolated mask should have a different label, lab>0

    # masks are binary uint8 type

    def extract_single_mask(self, mask, lab):

        _mask = np.zeros(mask.shape, dtype=np.uint8)

        area = mask == lab

        _mask[area] = 1

        return _mask

    def load_img(self, idx):

        im = PILImage.open(os.path.join(self.data, self.sorted_data[idx]))

        self.img_fname = os.path.join(self.data, self.sorted_data[idx])

        im = self.transform_img(im)

        return im

    def load_labels_covid_ctscan_data(self, idx):

        list_of_bboxes = []

        labels = []

        list_of_masks = []

        # load bbox

        self.fname = os.path.join(self.gt, self.sorted_gt[idx])

        # extract bboxes from the mask

        mask = np.array(PILImage.open(self.fname))

        # only GGO: merge C and background

        # or merge GGO and C into a single mask

        # or keep separate masks

        if self.mask_type == "ggo":

           mask[mask==3] = 0

        elif self.mask_type == "merge":

           mask[mask==3] = 2

        # array  (NUM_CLASS_IN_IMNG, H,W) without bgr+lungs class (merge Class 0 and 1)

        # THIS IS IMPORTANT! CAN TRIGGER CUDA ERROR

        mask_classes = mask == np.unique(mask)[:, None, None][2:]

        # extract bounding boxes and masks for each object

        for _idx, m in enumerate(mask_classes):

            lab_mask = method_label(m)

            regions = regionprops(lab_mask)

            for _i, r in enumerate(regions):

                # get rid of really small ones:

                if r.area > self.area_th:

                    box_coords = (r.bbox[1], r.bbox[0], r.bbox[3], r.bbox[2])

                    list_of_bboxes.append(box_coords)

                    labels.append(_idx + 1)

                    # create a mask for one object, append to the list of masks

                    mask_obj = self.extract_single_mask(lab_mask, r.label)

                    list_of_masks.append(mask_obj)

        # create labels for Mask R-CNN

        # DO NOT CHANGE THESE DATATYPES!

        lab = {}

        list_of_bboxes = torch.as_tensor(list_of_bboxes, dtype=torch.float)

        labels = torch.tensor(labels, dtype=torch.int64)

        masks = torch.tensor(list_of_masks, dtype=torch.uint8)

        lab['labels'] = labels

        lab['boxes'] = list_of_bboxes

        lab['masks'] = masks

        lab['fname'] = self.fname

        lab['img_name'] = self.img_fname

        return lab

    # 'magic' method: size of the dataset

    def __len__(self):

        return len(os.listdir(self.data))

    # return one datapoint

    def __getitem__(self, idx):

        X = self.load_img(idx)

        y = self.load_labels_covid_ctscan_data(idx)

        return X, y