import warnings
warnings.filterwarnings("ignore", category=UserWarning)

import os, nibabel
import sys, getopt
import PIL
from PIL import Image
import imageio
import scipy.misc
import numpy as np
import glob
from torch.utils import data
import torch
import random
from .augmentations import augment_spatial, Compose, RandomRotate, PaddingCenterCrop
from skimage import transform

from scipy.ndimage.interpolation import map_coordinates
from scipy.ndimage.filters import gaussian_filter

def augment_gamma(data_sample, gamma_range=(0.5, 2), invert_image=False, epsilon=1e-7, per_channel=False,
                  retain_stats=False):
    if invert_image:
        data_sample = - data_sample
    if not per_channel:
        if retain_stats:
            mn = data_sample.mean()
            sd = data_sample.std()
        if np.random.random() < 0.5 and gamma_range[0] < 1:
            gamma = np.random.uniform(gamma_range[0], 1)
        else:
            gamma = np.random.uniform(max(gamma_range[0], 1), gamma_range[1])
        minm = data_sample.min()
        rnge = data_sample.max() - minm
        data_sample = np.power(((data_sample - minm) / float(rnge + epsilon)), gamma) * rnge + minm
        if retain_stats:
            data_sample = data_sample - data_sample.mean() + mn
            data_sample = data_sample / (data_sample.std() + 1e-8) * sd
    else:
        for c in range(data_sample.shape[0]):
            if retain_stats:
                mn = data_sample[c].mean()
                sd = data_sample[c].std()
            if np.random.random() < 0.5 and gamma_range[0] < 1:
                gamma = np.random.uniform(gamma_range[0], 1)
            else:
                gamma = np.random.uniform(max(gamma_range[0], 1), gamma_range[1])
            minm = data_sample[c].min()
            rnge = data_sample[c].max() - minm
            data_sample[c] = np.power(((data_sample[c] - minm) / float(rnge + epsilon)), gamma) * float(rnge + epsilon) + minm
            if retain_stats:
                data_sample[c] = data_sample[c] - data_sample[c].mean() + mn
                data_sample[c] = data_sample[c] / (data_sample[c].std() + 1e-8) * sd
    if invert_image:
        data_sample = - data_sample
    return data_sample

class AC17Data(data.Dataset):

    def __init__(self,
                 root,
                 split='train',
                 augmentations=None,
                 img_norm=True,
                 k=5,
                 k_split=1,
                 target_size=(352, 352)
                 ):
        self.target_size = target_size
        self.ROOT_PATH = root
        self.split = split
        self.k = 5
        self.split_len = int(200/self.k)
        self.k_split = int(k_split)
        self.img_norm = img_norm
        self.augmentations = augmentations
        self.TRAIN_IMG_PATH = os.path.join(root, 'train', 'img')
        self.TRAIN_SEG_PATH = os.path.join(root, 'train', 'seg')
        self.list = self.read_files()

    def read_files(self):
        d = []
        txt_file = os.path.join('./data/data_series.txt')
        split_range = list(range((self.k_split-1)*self.split_len, self.k_split*self.split_len))
        with open(txt_file, 'r') as f:
            for i, line in enumerate(f):
                l = line.split(' ')
                key, val = int(l[0]), int(l[1])
                # append as train set
                if self.split == 'train' and i not in split_range:
                    d.append([key, val])

                #append as val set
                elif self.split == 'val' and i in split_range:
                    d.append([key, val])
                    if i > split_range[-1]: #do not need to iterate through rest of file
                        return d
        return d

    def __len__(self):
        if self.split == "train":
            return 200 - self.split_len
        else:
            return self.split_len

    def __getitem__(self, i): # i is index
        filename = "patient%03d_frame%02d" % (self.list[i][0], self.list[i][1])
        full_img_path = os.path.join(self.TRAIN_IMG_PATH, filename)
        full_seg_path = os.path.join(self.TRAIN_SEG_PATH, filename)
        img = nibabel.load(full_img_path+".nii")
        seg = nibabel.load(full_seg_path+"_gt.nii").get_data()
        pix_dim = img.header.structarr['pixdim'][1]
        img = np.array(img.get_data())
        seg = np.array(seg)

        pre_shape = img.shape[:-1]
        ratio = float(pix_dim/1.25)
        scale_vector = [ratio, ratio, 1]

        img = transform.rescale(img,
                                scale_vector,
                                order=1,
                                preserve_range=True,
                                multichannel=False,
                                mode='constant')
        seg = transform.rescale(seg,
                                scale_vector,
                                order=0,
                                preserve_range=True,
                                multichannel=False,
                                mode='constant')

        if self.augmentations is not None:
            img = img.transpose(2, 0, 1)
            seg = seg.transpose(2, 0, 1)
            img_c = np.zeros((img.shape[0], self.target_size[0], self.target_size[1]))
            seg_c = np.zeros((seg.shape[0], self.target_size[0], self.target_size[1]))

            for z in range(img.shape[0]):
                #mu = img[z].mean()
                #sigma = img[z].std()
                if img[z].min() > 0:
                    img[z] -= img[z].min()

                img_tmp, seg_tmp = self.augmentations(img[z].astype(np.uint32), seg[z].astype(np.uint8))
                img_tmp = augment_gamma(img_tmp)

                if self.img_norm:
                    #img_tmp = (img_tmp - img_tmp.min())/(img_tmp.max() - img_tmp.min() + 1e-10)
                    #img_tmp = (img_tmp - 0.5)*2.0
                    mu = img_tmp.mean()
                    sigma = img_tmp.std()
                    img_tmp = (img_tmp - mu) / (sigma+1e-10)
                img_c[z] = img_tmp
                seg_c[z] = seg_tmp

            img = img_c.transpose(1,2,0)
            seg = seg_c.transpose(1,2,0)
        
        img = torch.from_numpy(img).float()
        seg = torch.from_numpy(seg).long()

        diagnosis = -1
        p_id = self.list[i][0]
        if 1 <= p_id and p_id <= 20:
            diagnosis = 0 #DCM
        elif 21 <= p_id and p_id <= 40:
            diagnosis = 1 #HCM
        elif 41 <= p_id and p_id <= 60:
            diagnosis = 2 #MINF
        elif 61 <= p_id and p_id <= 80:
            diagnosis = 3 #NOR
        elif 81 <= p_id and p_id <= 100:
            diagnosis = 4 #RV

        data_dict = {
            "name": filename,
            "image": img,
            "mask": seg,
            "class": diagnosis
        }

        return data_dict

    def _transform(self, img, mask):
        if img.ndim == 3:
            img = np.expand_dims(img, axis=0)
            img = np.concatenate((img, img, img), axis=0)
        img = torch.from_numpy(img).float()
        mask = torch.from_numpy(mask).long()
        return img, mask


class AC17_2DLoad():
    def __init__(self, dataset, split='train', deform=True):
        super(AC17_2DLoad, self).__init__()
        self.data = []
        self.split = split
        self.deform = deform

        for i in range(dataset.__len__()):
            d = dataset[i]
            for x in range(d["image"].shape[-1]):
                entry = {}
                entry["image"] = d["image"].permute(2,0,1)[x]
                entry["mask"] = d["mask"].permute(2,0,1)[x]
                entry["name"] = d["name"] + "_z"+str(x)
                entry["class"] = d["class"]
                self.data.append(entry)
    
    def __getitem__(self, i):
        #return self.data[i]
        img = self.data[i]["image"]
        seg = self.data[i]["mask"]
       
        if self.split == 'train': #50% chance of deformation
            img = img.double().numpy()
            seg = seg.double().numpy()
            
            if random.uniform(0, 1.0) <= 0.5 and self.deform==True:
                if len(img.shape) == 2:
                    img = np.expand_dims(img, axis=2)
                seg = np.expand_dims(seg, axis=2)
                stacked = np.concatenate((img, seg), axis=2)
                red = self.random_elastic_deformation(stacked, alpha=500, sigma=20).transpose(2,0,1)
                img, seg = red[0], red[1]
            
            if img.ndim == 2:
                img = np.expand_dims(img, axis=0)
                img = np.concatenate((img, img, img), axis=0)
            # End Random Elastic Deformation
            d = {"image":torch.from_numpy(img).float(),
                 "mask":torch.from_numpy(seg),
                 "name":self.data[i]["name"],
                 "class":self.data[i]["class"]}
            return d

        elif self.split == 'val' or self.deform == False:
            img = img.unsqueeze(0)
            img = torch.cat([img, img, img], 0)
            
            d = {"image":img.float(),
                 "mask":seg,
                 "name":self.data[i]["name"],
                 "class":self.data[i]["class"]}
            return d


    def __len__(self):
        return len(self.data)

    def random_elastic_deformation(self, image, alpha, sigma, mode='nearest',
                                   random_state=None):
        """Elastic deformation of images as described in [Simard2003]_.
    ..  [Simard2003] Simard, Steinkraus and Platt, "Best Practices for
        Convolutional Neural Networks applied to Visual Document Analysis", in
        Proc. of the International Conference on Document Analysis and
        Recognition, 2003.
        """
        assert len(image.shape) == 3

        if random_state is None:
            random_state = np.random.RandomState(None)

        height, width, channels = image.shape

        dx = gaussian_filter(2*random_state.rand(height, width) - 1,
                         sigma, mode="constant", cval=0) * alpha
        dy = gaussian_filter(2*random_state.rand(height, width) - 1,
                         sigma, mode="constant", cval=0) * alpha

        x, y = np.meshgrid(np.arange(height), np.arange(width), indexing='ij')
        indices = (np.repeat(np.ravel(x+dx), channels),
                np.repeat(np.ravel(y+dy), channels),
                np.tile(np.arange(channels), height*width))

        values = map_coordinates(image, indices, order=1, mode=mode)

        return values.reshape((height, width, channels))

if __name__ == '__main__':

    DATA_DIR = ""
    augs = Compose([PaddingCenterCrop(352)])
    dataset = AC17Data(DATA_DIR, augmentations=augs)
    ac17 = AC17_2DLoad(dataset)
    dloader = torch.utils.data.DataLoader(ac17,batch_size=10)
    for idx, batch in enumerate(dloader):
        img, mask = batch['image'], batch['mask']
        print(mask.shape, img.shape, mask.max(), mask.min(), img.max(), img.min())
