import os

import numpy as np

import torch

import pandas as pd

from torch.utils.data import Dataset, DataLoader

from torchvision import transforms

from skimage import transform

from utils.Config import opt

import matplotlib.pylab as plt

from sklearn.model_selection import train_test_split

import utils.array_tool as at

import matplotlib.patches as patches

from data.data_utils import read_image

DSB_BBOX_LABEL_NAMES = ('p')  # Pneumonia

def inverse_normalize(img):

    if opt.caffe_pretrain:

        img = img + (np.array([122.7717, 115.9465, 102.9801]).reshape(3, 1, 1))

        return img[::-1, :, :].clip(min=0, max=255)

    # approximate un-normalize for visualize

    return (img * 0.225 + 0.45).clip(min=0, max=1) * 255

"""Transforms:

Data augmentation

"""

class Transform(object):

    def __init__(self, img_size):

        self.img_size = img_size

    def __call__(self, in_data):

        img_id, img, mask = in_data['img_id'], in_data['image'], in_data['mask']

        _, H, W = img.shape

        img, mask = preprocess(img, mask, self.img_size)

        return {'img_id': img_id, 'image': img.copy(), 'mask': mask.copy()}

def preprocess(img, mask, img_size):

    C, H, W = img.shape

    img = img / 255.

    img = transform.resize(img, (C, img_size, img_size), mode='reflect')

    mask = mask.astype(np.float32)

    mask = transform.resize(mask, (1, img_size, img_size), mode='reflect')

    # both the longer and shorter should be less than

    # max_size and min_size

    if opt.caffe_pretrain:

        normalize = caffe_normalize

    else:

        normalize = pytorch_normalze

    img = normalize(img)

    return img, mask

def pytorch_normalze(img):

    """

    https://discuss.pytorch.org/t/how-to-preprocess-input-for-pre-trained-networks/683

    https://github.com/pytorch/vision/issues/223

    return appr -1~1 RGB

    """

    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],

                                std=[0.229, 0.224, 0.225])

    img = normalize(torch.from_numpy(img))

    return img.numpy()

def caffe_normalize(img):

    """

    return appr -125-125 BGR

    """

    img = img[[2, 1, 0], :, :]  # RGB-BGR

    img = img * 255

    mean = np.array([122.7717, 115.9465, 102.9801]).reshape(3, 1, 1)

    img = (img - mean).astype(np.float32, copy=True)

    return img

class RSNADataset(Dataset):

    def __init__(self, root_dir, img_id, mask_id, transform=True):

        """

        Args:

        :param root_dir (string): Directory with all the images

        :param img_id (list): lists of image id

        :param train: if equals true, then read training set, so the output is image, mask and imgId

                      if equals false, then read testing set, so the output is image and imgId

        :param transform (callable, optional): Optional transform to be applied on a sample

        """

        self.root_dir = root_dir

        self.img_id = img_id

        self.mask_id = mask_id

        self.transform = transform

        self.tsf = Transform(opt.img_size)

    def __len__(self):

        return len(self.img_id)

    def __getitem__(self, idx):

        img_path = os.path.join(self.root_dir, 'images', self.img_id[idx].split('.')[0], 'image.png')

        mask_path = os.path.join(self.root_dir, 'masks', self.mask_id[idx])

        image = read_image(img_path, np.float32, False)

        mask = read_image(mask_path, np.uint8, False)

        sample = {'img_id': self.img_id[idx], 'image':image.copy(), 'mask':mask.copy()}

        if self.transform:

            sample = self.tsf(sample)

        return sample

class RSNADatasetTest(Dataset):

    def __init__(self, root_dir, transform=True):

        """

        Args:

        :param root_dir (string): Directory with all the images

        :param img_id (list): lists of image id

        :param train: if equals true, then read training set, so the output is image, mask and imgId

                      if equals false, then read testing set, so the output is image and imgId

        :param transform (callable, optional): Optional transform to be applied on a sample

        """

        self.root_dir = root_dir

        self.img_id = os.listdir(root_dir)

        self.transform = transform

        self.tsf = Transform(opt.img_size)

    def __len__(self):

        return len(self.img_id)

    def __getitem__(self, idx):

        img_path = os.path.join(self.root_dir, self.img_id[idx], 'image.png')

        image = read_image(img_path, np.float32, False)

        C, H, W = image.shape

        image = image / 255.

        image = transform.resize(image, (C, opt.img_size, opt.img_size), mode='reflect')

        if opt.caffe_pretrain:

            normalize = caffe_normalize

        else:

            normalize = pytorch_normalze

        image = normalize(image)

        sample = {'img_id': self.img_id[idx], 'image': image.copy()}

        return sample

def get_train_loader(root_dir, batch_size=16, shuffle=False, num_workers=4, pin_memory=False):

    """Utility function for loading and returning training and validation Dataloader

    :param root_dir: the root directory of data set

    :param batch_size: batch size of training and validation set

    :param split: if split data set to training set and validation set

    :param shuffle: if shuffle the image in training and validation set

    :param num_workers: number of workers loading the data, when using CUDA, set to 1

    :param val_ratio: ratio of validation set size

    :param pin_memory: store data in CPU pin buffer rather than memory. when using CUDA, set to True

    :return:

        - train_loader: Dataloader for training

    """

    img_ids = os.listdir(root_dir)

    img_ids.sort()

    transformed_dataset = RSNADataset(root_dir=root_dir, img_id=img_ids, transform=True)

    dataloader = DataLoader(transformed_dataset, batch_size=batch_size,

                            shuffle=shuffle, num_workers=num_workers, pin_memory=pin_memory)

    return dataloader

def get_train_val_loader(root_dir, batch_size=16, val_ratio=0.2, shuffle=False, num_workers=4, pin_memory=False):

    """Utility function for loading and returning training and validation Dataloader

    :param root_dir: the root directory of data set

    :param batch_size: batch size of training and validation set

    :param split: if split data set to training set and validation set

    :param shuffle: if shuffle the image in training and validation set

    :param num_workers: number of workers loading the data, when using CUDA, set to 1

    :param val_ratio: ratio of validation set size

    :param pin_memory: store data in CPU pin buffer rather than memory. when using CUDA, set to True

    :return:

        - train_loader: Dataloader for training

        - valid_loader: Dataloader for validation

    """

    df = pd.read_csv(os.path.join(opt.root_dir, 'train.csv'))

    img_id, mask_id = list(df['image']), list(df['label'])

    train_img_id, val_img_id, train_mask_id, val_mask_id = train_test_split(img_id, mask_id, random_state=42, test_size=val_ratio, shuffle=False)

    train_dataset = RSNADataset(root_dir=root_dir, img_id=train_img_id, mask_id=train_mask_id, transform=True)

    train_loader = DataLoader(train_dataset, batch_size=batch_size,

                            shuffle=shuffle, num_workers=num_workers, pin_memory=pin_memory)

    val_dataset = RSNADataset(root_dir=root_dir, img_id=val_img_id, mask_id=val_mask_id, transform=True)

    val_loader = DataLoader(val_dataset, batch_size=batch_size,

                             shuffle=shuffle, num_workers=num_workers, pin_memory=pin_memory)

    return train_loader, val_loader

def get_test_loader(batch_size=16, shuffle=False, num_workers=4, pin_memory=False):

    """Utility function for loading and returning training and validation Dataloader

    :param root_dir: the root directory of data set

    :param batch_size: batch size of training and validation set

    :param shuffle: if shuffle the image in training and validation set

    :param num_workers: number of workers loading the data, when using CUDA, set to 1

    :param pin_memory: store data in CPU pin buffer rather than memory. when using CUDA, set to True

    :return:

        - testloader: Dataloader of all the test set

    """

    transformed_dataset = RSNADatasetTest(root_dir=opt.test_root)

    testloader = DataLoader(transformed_dataset, batch_size=batch_size,

                            shuffle=shuffle, num_workers=num_workers, pin_memory=pin_memory)

    return testloader

def show_batch_train(sample_batched):

    """

    Visualize one training image and its corresponding bbox

    """

    img_id, image, mask = sample_batched['img_id'], sample_batched['image'], sample_batched['mask']

    image, mask = np.squeeze(at.tonumpy(image)), np.squeeze(at.tonumpy(mask))

    image = inverse_normalize(image)

    combined = np.multiply(image, mask)

    ax1 = plt.subplot(121)

    ax1.imshow(image / 255.)

    ax1.set_title(img_id[0])

    ax2 = plt.subplot(122)

    ax2.imshow(combined / 255.)

    ax2.set_title(img_id[0])

    plt.show()

def show_batch_test(sample_batch):

    img_id, image = sample_batch['img_id'], sample_batch['image']

    image = inverse_normalize(np.squeeze(at.tonumpy(image[0])))

    plt.figure()

    plt.imshow(image/255)

    plt.show()

if __name__ == '__main__':

    # Load training & validation set

    # train_loader, val_loader = get_train_val_loader(opt.root_dir, batch_size=1, val_ratio=0.2,

    #                                                 shuffle=False, num_workers=opt.num_workers,

    #                                                 pin_memory=opt.pin_memory)

    # for i_batch, sample in enumerate(val_loader):

    #     print(sample['img_id'], ', ', sample['image'].shape, ', ', sample['mask'].shape)

    #     show_batch_train(sample)

    test_loader = get_test_loader(batch_size=1, shuffle=False,

                                                num_workers=opt.num_workers,

                                                pin_memory=opt.pin_memory)

    for i_batch, sample in enumerate(test_loader):

        print(sample['img_id'], ', ', sample['image'].shape)

        show_batch_test(sample)