import random

import os

from PIL import Image

import torch

from torch.utils.data import Dataset, DataLoader

import torchvision.transforms as transforms

import pandas as pd

import re

import numpy as np

import utils

# borrowed from http://pytorch.org/tutorials/advanced/neural_style_tutorial.html

# and http://pytorch.org/tutorials/beginner/data_loading_tutorial.html

# define a training image loader that specifies transforms on images. See documentation for more details.

def get_tfms_3d(split, params):

    if split == "train":

        def tfms(x):

            # x = random_3d_crop(x, params.n_crop_vox)

            x = normalize(x, params)

            x = random_crop(x, params.n_crop_vox)

            # batchgenerators transforms expect bath dim and channel dim

            # add these and squeeze off later

            x = np.expand_dims(np.expand_dims(x, 0), 0)

            # x = transforms3d.spatial_transforms.augment_mirroring(x)[0]

            x = np.squeeze(x)

            return x

    else:

        def tfms(x):

            x = normalize(x, params)

            x = unpad(x, int(params.n_crop_vox/2))

            return x

    return tfms

def get_tfms(split = "train", size = 51):

    if split == "train":

        tfms = transforms.Compose([

            # transforms.CenterCrop(70),

            transforms.RandomCrop(size),

            # transforms.Resize((size, size)),

            transforms.RandomHorizontalFlip(),

            transforms.RandomVerticalFlip(),

            # transforms.RandomRotation(90),

            # transforms.Resize((224, 224)),

            # transforms.RandomResizedCrop(size, scale = (.9, 1)),

            # transforms.RandomRotation(12),

            # transforms.Resize((224, 224)),  # resize the image to 64x64 (remove if images are already 64x64),

            # transforms.RandomHorizontalFlip(),  # randomly flip image horizontally

            # transforms.RandomAffine(10, translate=(.1, .1), scale=(.1, .1), shear=.1, resample=False, fillcolor=0),

            transforms.ToTensor()

            # normalize_2d

            ])  # transform it into a torch tensor

    else:

        tfms = transforms.Compose([

            transforms.CenterCrop(size),

            # transforms.Resize((size, size)),

            transforms.ToTensor()

            # normalize_2d

            ])

    return tfms

def normalize_2d(x):

    return x / 255

def normalize(x, params=None):

    if params is None:

        MIN_BOUND = -1000.; MAX_BOUND = 600.0; PIXEL_MEAN = .25

    else:

        MIN_BOUND = params.hu_min; MAX_BOUND = params.hu_max; PIXEL_MEAN = params.pix_mean

    x = (x - MIN_BOUND) / (MAX_BOUND - MIN_BOUND)

    x[x > (1 - PIXEL_MEAN)] = 1.

    x[x < (0 - PIXEL_MEAN)] - 0.

    return x

def random_crop(x, num_vox=3):

    starts = np.random.choice(range(num_vox), replace=True, size=(x.ndim,))

    ends = x.shape - (num_vox - starts)

    for i in range(x.ndim):

        x = x.take(indices=range(starts[i],ends[i]), axis=i)

    return x

def unpad(x, n=2):

    """

    Skim off n-entries in 3 dimensions

    """

    assert type(x) is np.ndarray

    if n>0:

        x = x[n:-n,n:-n,n:-n]

    return x

class LIDCDataset(Dataset):

    """

    A standard PyTorch definition of Dataset which defines the functions __len__ and __getitem__.

    """

    def __init__(self, data_dir, transform, df, setting, params):

        """

        Store the filenames of the jpgs to use. Specifies transforms to apply on images.

        Args:

            data_dir: (string) directory containing the dataset

            transform: (torchvision.transforms) transformation to apply on image

        """

        self.setting = setting

        self.params = params

        self.data_dir = data_dir

        self.transform = transform

        self.df = df

        self.mode3d = setting.mode3d

        self.covar_mode = setting.covar_mode

        self.fase = setting.fase

        # print(df.head())

        # print(df.dtypes)

        assert ("name" in df.columns)

        self.name_col  = df.columns.get_loc("name")

        self.label_col = df.columns.get_loc(setting.outcome[0])

        self.data_cols = list(set(range(len(self.df.columns))) - 

                                set([self.name_col, self.label_col]))

        # split of data, which contains covariate data that is not name or label

        if self.covar_mode:

            self.data = self.df.loc[:,"t"].values

        # if len(self.data_cols) > 0:

        #     self.data = self.df.iloc[:,self.data_cols]

        df['x_true'] = df.x

        # calculate a transformation of x to assess robustness of method to different measurements

        df['x'] = df.x + params.size_offset

        if params.size_measurement == 'area':

            pass

        elif params.size_measurement == 'diameter':

            df['x'] = df.x.values ** (1/2)

        elif params.size_measurement == 'volume':

            df['x'] = df.x.values ** (3/2)

        else:

            raise ValueError(f'dont know how to measure size in {params.size_measurement}, pick area, diameter or volume')

        # renormalize x to make sure that whatever measurement is used, the MSE is comparable

        df['x'] = (df.x - df.x.mean()) / df.x.std()

    def __len__(self):

        # return size of dataset

        return self.df.shape[0]

    def __getitem__(self, idx):

        """

        Fetch index idx image and labels from dataset. Perform transforms on image.

        Args:

            idx: (int) index in [0, 1, ..., size_of_dataset-1]

        Returns:

            image: (Tensor) transformed image

            label: (int) corresponding label of image

        """

        # image = Image.open(self.fpath_dict[self.idx_to_id[idx]]).convert("RGB")  # PIL image

        if self.mode3d:

            image = np.load(os.path.join(self.data_dir, self.df.iloc[idx, self.name_col]))

            image = image.astype(np.float32)

            image = self.transform(image)

            image = torch.from_numpy(image).unsqueeze(0)

        else:

            img_name = os.path.join(self.data_dir, 

                                    self.df.iloc[idx, self.name_col])

            image = Image.open(img_name).convert("L") # use rgb for resnet compatibility; L for grayscale

            image = self.transform(image)

        label = torch.from_numpy(np.array(self.df.iloc[idx, self.label_col], dtype = np.float32))

        sample = {"image": image, 'label': label}

        for variable in ["x", "y", "z", "t", 'x_true']:

            if variable in self.df.columns:

                sample[variable] = self.df[variable].values[idx].astype(np.float32)

        if self.setting.fase == "feature":

            sample[self.setting.outcome[0]] = self.df[self.setting.outcome[0]].values[idx].astype(np.float32)

        return sample

def fetch_dataloader(args, params, setting, types = ["train"], df = None):

    """

    Fetches the DataLoader object for each type in types from data_dir.

    Args:

        types: (list) has one or more of 'train', 'val', 'test' depending on which data is required

        data_dir: (string) directory containing the dataset

        df: pandas dataframe containing at least name, label and split

        params: (Params) hyperparameters

    Returns:

        data: (dict) contains the DataLoader object for each type in types

    """

    if setting.gen_model == "":

        if setting.mode3d:

            data_dir = "data"

        else:

            data_dir = "slices"

    else:

        data_dir = os.path.join(setting.home, "data")

    if df is None:

        df = pd.read_csv(os.path.join(data_dir, "labels.csv"))

    dataloaders = {}

    if not setting.mode3d:

        pass

        # print(df.name.tolist()[:5])

        # df["name"] = df.apply(lambda x: os.path.join(x["split"], x["name"]), axis=1)

        # print(df.name.tolist()[:5])

    # make sure the dataframe has no index

    df_cols = df.columns

    df = df.reset_index()

    df = df[df_cols]

    try:

        assert setting.outcome[0] in df.columns

    except:

        print(f"outcome {setting.outcome[0]} not in df.columns:")

        print("\n".join(df.columns))

        raise

    if "split" in df.columns:

        splits = [x for x in types if x in df.split.unique().tolist()]

    else:

        df["split"] = types[0]

        splits = types

    df_grp = df.groupby("split")

    # for split in ['train', 'val', 'test']:

    for split, df_split in df_grp:

        df_split = df_split.drop("split", axis = 1)

        if split in types:

            # path = os.path.join(data_dir, split)

            path = data_dir

            if setting.mode3d:

                tfms = get_tfms_3d(split, params)

                # tfms = []

            else:

                tfms = get_tfms(split, params.size)

            # use the train_transformer if training data, else use eval_transformer without random flip

            if split == 'train':

                dl = DataLoader(LIDCDataset(path, tfms, df_split, setting, params), 

                                        shuffle=True,

                                        num_workers=params.num_workers,

                                        batch_size=params.batch_size, 

                                        pin_memory=params.cuda)

                                        # batch_size = batch_size,

                                        # num_workers=2,

                                        # pin_memory=True)

            else:

                # dl = DataLoader(SEGMENTATIONDataset(path, eval_transformer, df[df.split.isin([split])]), 

                dl = DataLoader(LIDCDataset(path, tfms, df_split, setting, params), 

                                batch_size=params.batch_size,

                                num_workers=params.num_workers,

                                shuffle=False,

                                pin_memory=params.cuda)

                                # batch_size = batch_size,

                                # num_workers=2,

                                # pin_memory=True)

            dataloaders[split] = dl

    return dataloaders