'''

Data loaders and manipulation

BY: Yuval

'''

import os

from skimage.io import imread

import numpy as np

import pandas as pd

import torch

from skimage import io, transform

from torch.utils.data import Dataset, DataLoader

from torchvision import transforms, utils

from .myio import load_one_image

from itertools import product

from scipy.ndimage import zoom

from tqdm import tqdm_notebook

from multiprocessing import Pool

def narrow(arr,axis,start,end):

    # A numpy implementation to torch.narrow

    if axis<0:

        axis = arr.ndim-axis

    return arr[(slice(None),)*axis+(slice(start,end,1),)+(slice(None),)*(arr.ndim-axis-1)]

def _shift(arr, shift, axis, fill_value=None):

    ''' Shits an image. Fill the empty space

        Arges:

            arr  : numpy array with the image to shift

            shift: shift step

            axis : axis to shift

            fill_value: a value to fill empty spaces - default - None => Use min value in array

        Return:

            new shifted array

        Update: Yuval 12/10/2019

    '''

    if shift == 0:

        return arr

    if fill_value is None:

        fill_value = arr.min()

    if axis < 0:

        axis += arr.ndim

    dim_size = arr.shape[axis]

    after_start = dim_size - shift

    slice_shape=list(arr.shape)

    slice_shape[axis]=abs(shift)

    if shift < 0:

        after_start = -shift

        shift = dim_size - abs(shift)

        before = np.ones(slice_shape)*fill_value

        after = narrow(arr,axis, after_start, shift)

    else:

        before = narrow(arr,axis, 0, dim_size - shift)

        after = np.ones(slice_shape)*fill_value

    return np.concatenate([after, before], axis)

class MyTransform():

    '''

    My implementation for image transformation class

    Args:

     flip        : do a right/left mirroring.  default - False

     mean_change : mean change.                default - 0

     std_change  : std change.                 default - 0

     crop=None   : crop image to size. If tuple - x,y. the crop position will chabge randomly.  

                   default -None, keep image. the crop position will

     seed        : seed to use for random.     default - None

     zoom=0.0    : Zoom images, by             default - 0 (stay the same)

     rotate=0    : Rotation Angle, Deg         default - 0 

     shift=0     : Shift image                 default - 0

     out_size    : Output image size, int (x=y) or tuple 

                                               default - None => keep input                                           

   Methods:

       random: random transform using the init parameters

       Args:

           imags - numpy array with one or more images, for multiple images the first dim should be the channle

       Returns 

           numpy array with randomly transformed images of size out_size

   Updated: Yuval 12/10/19

   '''

    def __init__(self,

                 flip=False,

                 mean_change=0,

                 std_change=0,

                 crop=None,

                 seed=None,

                 zoom=0.0,

                 rotate=0,

                 shift=0,

                 out_size=None,

                 normal=True,

                 anti_aliasing=True

                ):

        self.do_flip = flip

        self.rotate_angle=rotate

        self.mean_change = mean_change

        self.std_change = std_change

        np.random.seed(seed)

        self.zoom_factor=zoom

        self.anti_aliasing=anti_aliasing

        if isinstance (crop,tuple):

            self.cropx=crop[0]

            self.cropy=crop[1]

        else:

            self.cropx=crop

            self.cropy=crop

        if isinstance (shift,tuple):

            self.shiftx=shift[0]

            self.shifty=shift[1]

        else:

            self.shiftx=shift

            self.shifty=shift

        if isinstance (out_size,tuple):

            self.out_sizex=out_size[0]

            self.out_sizey=out_size[1]

        else:

            self.out_sizex=out_size

            self.out_sizey=out_size

        if normal:

            self.randf = lambda n: torch.randn(n).numpy() 

        else:        

            self.randf = lambda n: 2.0*torch.rand(n).numpy()-1.0

    def random(self,imgs):

        sqz=False

        imgs=imgs.copy()

        if len(imgs.shape)==2:

            imgs=np.expand_dims(imgs, axis=0)

            sqz=True

        cropx,cropy = imgs.shape[1:3] if self.cropx is None else (self.cropx,self.cropy)

        out_sizex,out_sizey = imgs.shape[1:3] if self.out_sizex is None else (self.out_sizex,self.out_sizey)

        imgs=imgs.transpose(1,2,0)

        if (self.std_change>0) or (self.mean_change>0):

#            for i,ix in enumerate(self.channels):

#                imgs[i]=imgs[i]*np.random.normal(loc=1,scale=self.std_change)+np.random.normal(loc=0,scale=self.mean_change)

            imgs=self.change_mean_std(imgs,self.randf(1)[0]*self.mean_change,1+self.randf(1)[0]*self.std_change)

        if self.do_flip:

            if (torch.randint(low=0,high=2,size=(1,))[0]>0):

                imgs = self.flip(imgs)

        if self.rotate_angle>0:

            angle=int(torch.randint(-self.rotate_angle,self.rotate_angle,(1,))[0])

            imgs=self.rotate(imgs,angle)       

        if self.shiftx>0:

            imgs=self.img_shift(imgs,np.random.randint(-self.shiftx,self.shiftx),np.random.randint(-self.shifty,self.shifty))

        if self.zoom_factor!=0:

            if isinstance(self.zoom_factor,tuple):

                factor_x=1+self.randf(1)[0]*self.zoom_factor[0]

                factor_y=(1+self.randf(1)[0]*self.zoom_factor[1])*factor_x

                factor=(factor_x,factor_y)

            else:

                factor=1+np.random.randn(1)[0]*self.zoom_factor

            imgs=self.zoom(imgs,factor)

        x0=max(imgs.shape[1]//2-cropx//2,0)

        y0=max(imgs.shape[0]//2-cropy//2,0)

        imgs=self.crop(imgs,x0,y0,cropx,cropy)

        if (imgs.shape[0]!=out_sizey) or (imgs.shape[1]!=out_sizex):

            imgs=self.resize(imgs,out_sizex, out_sizey)

        imgs=imgs.transpose(2,0,1)        

        if sqz:

            imgs=imgs.squeeze(0)

        return imgs

    def flip(self,img,axis=1):

        return np.flip(img,axis=axis)

    def img_shift(self,img,x,y):

        return _shift(_shift(img,x,1),y,0)

    def crop(self,img,x,y,width,hight):

        if width>img.shape[1]:

            img=np.concatenate([np.ones((img.shape[0],(width-img.shape[1])//2+1,img.shape[-1]))*img.min(),

                                img,

                                np.ones((img.shape[0],(width-img.shape[1])//2+1,img.shape[-1]))*img.min()],1)

        if hight>img.shape[0]:

            img=np.concatenate([np.ones(((hight-img.shape[0])//2+1,img.shape[1],img.shape[-1]))*img.min(),

                                img,

                                np.ones(((hight-img.shape[0])//2+1,img.shape[1],img.shape[-1]))*img.min()],0)

        return img[x:x+width,y:y+hight,:]

    def change_mean_std(self,img,mean,std):

        if (isinstance(mean,list)):

            for i,(m,s) in zip(mean,std):

                img[...,i] = img[...,i]*s+m

        else:

            img = img*std + mean

        return img

    def resize(self,img,width,hight):

        return transform.resize(img,(hight,width),anti_aliasing=self.anti_aliasing)

    def zoom(self,img,factor):

#        timg=transform.rescale(img,1.0+factor,multichannel=True,mode='constant',cval=float(img.min()))

        return transform.rescale(img,factor,multichannel=True,mode='constant',cval=float(img.min()))

    def rotate(self,img,angle,resize=True):

        return transform.rotate(img, angle, resize=resize, center=None, order=1, 

                                mode='constant', cval=img.min(), clip=True, preserve_range=False)   

class sampler():

    '''

    sampler class for RSNA 2019. sample the images according to the tagets vector

    Args:

        arr:            numpy array with the target vectors

        norm_ratio:     float - the ratio of sampling for all zero target vector

        sampled ratios: a numpy vector,len: arr.shape[-1], sampling ratio by target value.

        unique_col:     numpy vector length arr.shape[0], 

                        with values which will be uniqued (no 2 samples would have the same value in this column

                        default: None (don't use)

    Methods:

        __call__:

        Args:

            index_arr: index vector, sample only from this index. default: None

    Update: Yuval 12/10/19

    '''

    def __init__(self,arr,norm_ratio,sampled_ratios,unique_col=None):

        self.arr=arr

        self.norm_ratio = norm_ratio

        self.sampled_ratios = sampled_ratios

        self.unique_col=unique_col

    def do_unique(self,indxes):

        if self.unique_col is not None:

            u,ind = np.unique(self.unique_col[indxes],return_index=True)

            return indxes[ind]

        else:

            return indxes

    def __call__(self,index_arr=None):

        if index_arr is None:

            index_arr = Ellipsis

        sampled = []

        indxes=np.argwhere(~self.arr[index_arr].any(axis=1)>0).squeeze()

        np.random.shuffle(indxes)

        indxes = self.do_unique(indxes)    

        sampled.append(indxes[:int(indxes.shape[0]*(self.norm_ratio-np.floor(self.norm_ratio)))])

        for i in range(int(self.norm_ratio)):

            indxes=np.argwhere(~self.arr[index_arr].any(axis=1)>0).squeeze()

            np.random.shuffle(indxes)  

            sampled.append(self.do_unique(indxes))

        for i,s in enumerate(self.sampled_ratios):

            s_=s

            if s_>1:

                s_=np.floor(s_)

                for j in range(int(s_)):

                    indxes=np.argwhere(self.arr[index_arr][...,i]>0).squeeze()

                    np.random.shuffle(indxes) 

                    sampled.append(self.do_unique(indxes))

                s_=s-s_

            if s_>0:

                    indxes=np.argwhere(self.arr[index_arr][...,i]>0).squeeze()

                    np.random.shuffle(indxes)

                    indxes = self.do_unique(indxes)

                    sampled.append(indxes[:int(indxes.shape[0]*s_)])

        return np.concatenate(sampled)

class simple_sampler():

    '''

    Simple sampler, will shuffle and sample a part of an array (dim - 0)

    Args:

        arr   :numpy array

        ratio :float, the sampling ratio, if>1 the same as =1

    Methods:

        __call__:

            Return numpy vector, type long, with sampled indexes

    Update: Yuval 12/10/19

    '''

    def __init__(self,arr,ratio):

        self.arr=arr

        self.ratio = ratio

    def __call__(self):

        indxes=np.arange(self.arr.shape[0])

        np.random.shuffle(indxes)

        return indxes[:int(self.arr.shape[0]*self.ratio)]

class Mixup():

    '''

    Method for mixup augmentation - TODO doc

    '''

    def __init__(self,alpha=0.4,device='gpu'):

        self.alpha=alpha

        self.device=device

    def __call__(self,images,targets):

        lambd = np.random.beta(self.alpha, self.alpha, targets.size(0))

        lambd = np.abs(lambd-0.5)+0.5 #np.concatenate([lambd[:,None], 1-lambd[:,None]], 1).max(1)

        shuffle = torch.randperm(targets.size(0)).to(self.device)

        lambd=torch.tensor(lambd,dtype=torch.float).to(self.device)

        out_images = (lambd*images.transpose(0,-1)+(1-lambd)*images[shuffle].transpose(0,-1)).transpose(0,-1)

        out_targets = torch.cat([targets.unsqueeze(-1),

                                 targets[shuffle].unsqueeze(-1),

                                 lambd.expand_as(targets.transpose(0,-1)).transpose(0,-1).unsqueeze(-1)],-1)

#        out_targets = (lambd*targets.transpose(0,-1)+(1-lambd)*targets[shuffle].transpose(0,-1)).transpose(0,-1)

        return out_images, out_targets

class ImageDataset(Dataset):

    '''

        RSNA 2019 Image (DICOM) dataset to use in Pytorch dataloader.

        Base class: Dataset

        Args:

            df              : Data frame with the image ids

            base_path       : File path for the images 

            transform=None  : Transfor method. to perform after the images are loaded. default: None - no transform

            out_shape=None  : Expected output shape - used only for sanity check.      default: None - no check

            window_eq=False : Do window equaliztion: (for backward competability, don't use it anymore use WSO)

                              False - No equalization

                              True  - Do equalization with window = [(40,80),(80,200),(600,2800)]

                              tuple/list shaped as above 

            equalize         : Equalize - return (image-mean)/std

            rescale=False    : Use DICOM parameters for rescale, done automaticaly if windows_eq!=False

       Update:Yuval 12/10/19       

    '''

    def __init__(self, df, base_path, transform=None,out_shape=None,window_eq=False,equalize=True,rescale=False):

        super(ImageDataset, self).__init__()

        self.df = df

        self.pids = df.PatientID.values

        self.transform = transform

        self.base_path = base_path

        self.out_shape=out_shape

        self.window_eq=window_eq

        self.equalize = equalize

        self.rescale=rescale

    def __len__(self):

        return self.pids.shape[0] 

    def __getitem__(self, idx):

        sample=load_one_image(self.pids[idx],equalize=self.equalize,base_path=self.base_path,file_path='',

                              window_eq=self.window_eq,rescale=self.rescale)

        sample = torch.tensor(sample,dtype=torch.float) \

            if self.transform is None else torch.tensor(self.transform(sample),dtype=torch.float)

        if len(sample.shape)==2:

            sample = sample.unsqueeze(0)

        if self.out_shape is not None:

            if sample.shape != self.out_shape:

                print ("Error in idx {}".format(idx))

                print (sample.shape,sample)

                sample = torch.randn(self.out_shape)*1e-5

        return sample

class FeatursDataset(Dataset):

    '''

        RSNA 2019 features dataset to use in Pytorch dataloader.

        Base class: Dataset

        Args:

            df              : Data frame

            features        : pytorch tensor with features.

                              Shape:

                                  option1 - (df.shape[0],num_of_features) - normal mode

                                  option2 - Not Implemented here yet

                                            (df.shape[0],N,num_of_features) - TTA mode, 

                                            will select random feature vector from same raw. 

            num_neighbors   : int, Number of neighbor to return, output will be shape (1+2*num_neighbors,features.shape[-1])

            ref_column      : string, The name of the column in df with the series id

            order_column    : string, The name of the column in df with the data which will determine the neighbors

            target_columns  : list of strings/None, names of column in df with the target data, 

                              default None - no target data will be returned

        Methods:

            __calls__:

                return: sample - tensor size (1+2*num_neighbors,features.shape[-1])

                        if target column is defined, return tuple with the 2nd valiable: 

                            targets - tensor size (1+2*num_neighbors,len(target_columns)), dtype=torch.float

        Update:Yuval 12/10/19       

    '''

    def __init__(self, df, features,num_neighbors, ref_column,order_column,target_columns=None):

        """

        Args:

            Todo

        """

        super(FeatursDataset, self).__init__()

        self.df = df.sort_values([ref_column,order_column])

        self.num_neighbors = num_neighbors

        self.ref_column = ref_column

        self.target_columns=target_columns

        self.target_tensor=None if target_columns is None else torch.tensor(df[self.target_columns].values,dtype=torch.float)

        self.features=features

        self.ref_arr=np.zeros((self.df.shape[0],1+2*self.num_neighbors),dtype=np.long)

        for i in range(-self.num_neighbors,self.num_neighbors+1):

            self.ref_arr[:,i+self.num_neighbors]=np.where(self.df[ref_column]==self.df[ref_column].shift(i),

                                                          np.roll(self.df.index.values,i),

                                                          self.df.index.values)

        self.ref_arr=torch.tensor(self.ref_arr[np.argsort(self.ref_arr[:,self.num_neighbors])])

    def __len__(self):

        return self.ref_arr.shape[0] 

    def __getitem__(self, idx):

        sample=self.features[self.ref_arr[idx]]

        return sample if self.target_tensor is None else (sample, self.target_tensor[idx])

class FeatursDatasetCor(Dataset):

    """Not Used, like FeatursDataset but determine neighbors according to feature distance"""

    def __init__(self, df, features,num_neighbors, ref_column,target_columns=None):

        """

        Args:

            Todo

        """

        super(FeatursDatasetCor, self).__init__()

        self.num_neighbors = num_neighbors

        self.ref_column = ref_column

        self.target_columns=target_columns

        self.target_tensor=None if target_columns is None else torch.tensor(df[self.target_columns].values,dtype=torch.float)

        self.features=features

        self.ref_arr=np.zeros((df.shape[0],1+2*self.num_neighbors),dtype=np.long)

        unq,si=np.unique(df[self.ref_column].values,return_inverse=True)

        for i in tqdm_notebook(range(unq.shape[0]), leave=False):

            sinx = np.where(si==i)[0]

            r=np.corrcoef(self.features[sinx].numpy())

            self.ref_arr[sinx]=sinx[np.argsort(-r)][:,:1+2*self.num_neighbors]

    def __len__(self):

        return self.ref_arr.shape[0] 

    def __getitem__(self, idx):

        sample=self.features[self.ref_arr[idx]]

        return sample if self.target_tensor is None else (sample, self.target_tensor[idx])

class FullHeadImageDataset(Dataset):

    '''

        RSNA 2019  full head dataset to use in Pytorch dataloader.

        return all the slices from a scan in the right order.

        Base class: Dataset

        Args:

            df              : Data frame

            base_path       : File path for the images 

            SeriesIDs       : numpy array with scan series ids, each call the mathod will return one full series

            transform       : Transfor method. to perform after the images are loaded. 

                              The same transformation is done for all images in a series

                              default: None - no transform

            out_shape       : Expected output shape - used only for sanity check.

                              default: None - no check

            window_eq       : Do window equaliztion: (for backward competability, don't use it anymore use WSO)

                              False - No equalization [default]

                              True  - Do equalization with window = [(40,80),(80,200),(600,2800)]

                              tuple/list shaped as above 

            equalize        : Equalize - return (image-mean)/std [default - False]

            rescale         : Use DICOM parameters for rescale, done automaticaly if windows_eq!=False

                               default - True

            ref_column      : string, The name of the column in df with the series id

            order_column    : string, The name of the column in df with the data which will determine the neighbors

            target_columns  : list of strings/None, names of column in df with the target data, 

                              default None - no target data will be returned

        Methods:

            __calls__:

                return: sample - tensor size (# of images in series,image shape)

                        if target column is defined, return tuple with the 2nd valiable: 

                            targets - tensor size (# of images in series,len(target_columns)), dtype=torch.float

        Update:Yuval 12/10/19       

    '''

    def __init__(self, df,

                 base_path,

                 SeriesIDs,

                 ref_column,

                 order_column,

                 transform=None,

                 window_eq=False,

                 equalize=False,

                 rescale=True, 

                 target_columns=None,

                 full_transform=True):

        super(FullHeadImageDataset, self).__init__()

        self.df = df

        self.SeriesIDs=SeriesIDs

        self.ref_column=ref_column

        self.order_column=order_column

        self.target_columns=target_columns

        self.pids = df.PatientID.values

        self.transform = transform

        self.base_path = base_path

        self.window_eq=window_eq

        self.equalize = equalize

        self.rescale=rescale

        self.full_transform=full_transform

        self.ref_arr=df[ref_column].values

        self.order_arr=df[order_column].values

        self.target_tensor=None if target_columns is None else torch.tensor(df[self.target_columns].values,dtype=torch.float)

    def __len__(self):

        return self.SeriesIDs.shape[0]

    def __getitem__(self, idx):

        head_idx=np.where(self.ref_arr==self.SeriesIDs[idx])[0]

        sorted_head_idx=head_idx[np.argsort(self.order_arr[head_idx])]

        samples=[]

        for i in sorted_head_idx:

            sample=load_one_image(self.pids[i],equalize=self.equalize,base_path=self.base_path,file_path='',

                                  window_eq=self.window_eq,rescale=self.rescale)[None]

            if (not self.full_transform) and (self.transform is not None):

                sample = self.transform(sample)

            samples.append(sample)

        headimages=np.concatenate(samples,0)

        headimages = torch.tensor(headimages,dtype=torch.float) \

                if ((self.transform is None) or (not self.full_transform)) else torch.tensor(self.transform(headimages),dtype=torch.float)

        headimages=headimages[:,None]  # lat's make a batch out of it.

        if self.target_tensor is not None:

            targets=self.target_tensor[sorted_head_idx]

        return headimages if self.target_tensor is None else (headimages, targets)

class FullHeadDataset(Dataset):

    '''

        RSNA 2019 full head scan features dataset to use in Pytorch dataloader.

        Base class: Dataset

        Args:

            df              : Data frame

            SeriesIDs       : numpy array with scan series ids, each call the mathod will return one full series

            features        : pytorch tensor with features.

                              Shape:

                                  option1 - (df.shape[0],num_of_features) - normal mode

                                  option2 - (df.shape[0],N,num_of_features) - TTA mode, 

                                            will select random feature vector from same raw. 

            ref_column      : string, The name of the column in df with the series id

            order_column    : string, The name of the column in df with the data which will determine the neighbors

            target_columns  : list of strings/None, names of column in df with the target data, 

                              default None - no target data will be returned

        Methods:

            __calls__:

                return: sample - tensor size (# of images in series,features.shape[-1])

                        if target column is defined, return tuple with the 2nd valiable: 

                            targets - tensor size (# of images in series,len(target_columns)), dtype=torch.float

        Update:Yuval 12/10/19       

    '''

    def __init__(self, df, SeriesIDs,features, ref_column,order_column,target_columns=None,max_len=60,multi=1):

        """

        Args:

            Todo

        """

        super(FullHeadDataset, self).__init__()

        self.ref_column = ref_column

        self.target_columns=target_columns

        self.target_tensor=None if target_columns is None else torch.tensor(df[self.target_columns].values,dtype=torch.float)

        self.features=features

        self.ref_arr=df[ref_column].values

        self.order_arr=df[order_column].values

        self.max_len=max_len

        self.SeriesIDs=SeriesIDs

        self.multi=multi

    def __len__(self):

        return self.SeriesIDs.shape[0] 

    def __getitem__(self, idx):

        sample = torch.zeros((self.max_len,self.features.shape[-1]*self.multi),dtype=torch.float)

        head_idx=np.where(self.ref_arr==self.SeriesIDs[idx])[0]

        sorted_head_idx=head_idx[np.argsort(self.order_arr[head_idx])]

        if self.features.dim()==3:         

            if self.multi>1:

                tta_idx=torch.zeros((head_idx.shape[0],self.multi),dtype=torch.long)

                for i in range(head_idx.shape[0]):

                    tta_idx[i]=torch.randperm(self.features.shape[1],dtype=torch.long)[:self.multi]

#                tta_idx2=(tta_idx+torch.LongTensor(head_idx.shape[0]).random_(1, self.features.shape[1]))%self.features.shape[1]

                sample[:head_idx.shape[0]]=torch.cat([self.features[sorted_head_idx,tta_idx[:,i]] for i in range(self.multi)],-1)

            else:

                tta_idx=torch.LongTensor(head_idx.shape[0]).random_(0, self.features.shape[1])

                sample[:head_idx.shape[0]]=self.features[sorted_head_idx,tta_idx]

        else:

            sample[:head_idx.shape[0]]=self.features[sorted_head_idx]

        if self.target_tensor is not None:

            targets = -1*torch.ones((self.max_len,self.target_tensor.shape[-1]),dtype=torch.float)

            targets[:head_idx.shape[0]]=self.target_tensor[sorted_head_idx]

        return sample if self.target_tensor is None else (sample, targets)

class DatasetCat(Dataset):

    '''

    Concatenate datasets for Pytorch dataloader

    The normal pytorch implementation does it only for raws. this is a "column" implementation

    Arges:

        datasets: list of datasets, of the same length

    Updated: Yuval 12/10/2019

    '''

    def __init__(self,datasets):

        '''

        Args: datasets - an iterable containing the datasets

        '''

        super(DatasetCat, self).__init__()

        self.datasets=datasets

        assert len(self.datasets)>0

        for dataset in datasets:

            assert len(self.datasets[0])==len(dataset),"Datasets length should be equal"

    def __len__(self):

        return len(self.datasets[0])

    def __getitem__(self, idx):

        outputs = tuple(dataset.__getitem__(idx) for i in self.datasets for dataset in (i if isinstance(i, tuple) else (i,)))

        return tuple(output for i in outputs for output in (i if isinstance(i, tuple) else (i,)))