import os

import torch

import numpy as np

import math

import random

import cv2 as cv

import nibabel as nib

import torch

from torch.utils import data

import torchvision.transforms as transforms

import matplotlib.pyplot as plt

import pandas as pd

from data_loader.preprocess import readVol,to_uint8,IR_to_uint8,histeq,preprocessed,get_stacked,rotate,calc_crop_region,calc_max_region_list,crop,get_edge

class MR18loader_CV(data.Dataset):

    def __init__(self,root='../../data/',val_num=5,is_val=False,

                 is_transform=False,is_flip=False,is_rotate=False,is_crop=False,is_histeq=False,forest=5):

        self.root=root

        self.val_num=val_num

        self.is_val=is_val

        self.is_transform=is_transform

        self.is_flip=is_flip

        self.is_rotate=is_rotate

        self.is_crop=is_crop

        self.is_histeq=is_histeq

        self.forest=forest

        self.n_classes=11

        # Back: Background

        # GM:   Cortical GM(red), Basal ganglia(green)

        # WM:   WM(yellow), WM lesions(blue)

        # CSF:  CSF(pink), Ventricles(light blue)

        # Back: Cerebellum(white), Brainstem(dark red)

        self.color=np.asarray([[0,0,0],[0,0,255],[0,255,0],[0,255,255],[255,0,0],\

                [255,0,255],[255,255,0],[255,255,255],[0,0,128],[0,128,0],[128,0,0]]).astype(np.uint8)

        # Back , CSF , GM , WM

        self.label_test=[0,2,2,3,3,1,1,0,0]

        # nii paths

        self.T1path=[self.root+'training/'+name+'/pre/reg_T1.nii.gz' for name in ['1','4','5','7','14','070','148']]

        self.IRpath=[self.root+'training/'+name+'/pre/IR.nii.gz' for name in ['1','4','5','7','14','070','148']]

        self.T2path=[self.root+'training/'+name+'/pre/FLAIR.nii.gz' for name in ['1','4','5','7','14','070','148']]

        self.lblpath=[self.root+'training/'+name+'/segm.nii.gz' for name in ['1','4','5','7','14','070','148']]

        # val path

        self.val_T1path=self.T1path[self.val_num-1]

        self.val_IRpath=self.IRpath[self.val_num-1]

        self.val_T2path=self.T2path[self.val_num-1]

        self.val_lblpath=self.lblpath[self.val_num-1]

        # train path

        self.train_T1path=[temp for temp in self.T1path if temp not in [self.val_T1path]]

        self.train_IRpath=[temp for temp in self.IRpath if temp not in [self.val_IRpath]]

        self.train_T2path=[temp for temp in self.T2path if temp not in [self.val_T2path]]

        self.train_lblpath=[temp for temp in self.lblpath if temp not in [self.val_lblpath]]

        if self.is_val==False:

            print('training data')

            T1_nii=[to_uint8(readVol(path)) for path in self.train_T1path]

            IR_nii=[IR_to_uint8(readVol(path)) for path in self.train_IRpath]

            T2_nii=[to_uint8(readVol(path)) for path in self.train_T2path]

            lbl_nii=[readVol(path) for path in self.train_lblpath]

            if self.is_flip:

                vol_num=len(T1_nii)

                for nums in range(vol_num):

                    T1_nii.append(np.array([cv.flip(slice_,1) for slice_ in T1_nii[nums]]))

                    IR_nii.append(np.array([cv.flip(slice_,1) for slice_ in IR_nii[nums]]))

                    T2_nii.append(np.array([cv.flip(slice_,1) for slice_ in T2_nii[nums]]))

                    lbl_nii.append(np.array([cv.flip(slice_,1) for slice_ in lbl_nii[nums]]))

            if self.is_histeq:

                print('hist equalizing......')

                T1_nii=[histeq(vol) for vol in T1_nii]

                IR_nii=[vol for vol in IR_nii]

                T2_nii=[vol for vol in T2_nii]

            print('get stacking......')

            T1_stack_lists=[get_stacked(vol,self.forest) for vol in T1_nii]

            IR_stack_lists=[get_stacked(vol,self.forest) for vol in IR_nii]

            T2_stack_lists=[get_stacked(vol,self.forest) for vol in T2_nii]

            lbl_stack_lists=[get_stacked(vol,self.forest) for vol in lbl_nii]

            if self.is_rotate:

                print('rotating......')

                angle_list=[5,-5,10,-10,15,-15]

                sample_num=len(T1_stack_lists)

                for angle in angle_list:

                    for sample_index in range(sample_num):

                        T1_stack_lists.append(rotate(T1_stack_lists[sample_index],angle,interp=cv.INTER_CUBIC).copy())

                        IR_stack_lists.append(rotate(IR_stack_lists[sample_index],angle,interp=cv.INTER_CUBIC).copy())

                        T2_stack_lists.append(rotate(T2_stack_lists[sample_index],angle,interp=cv.INTER_CUBIC).copy())

                        lbl_stack_lists.append(rotate(lbl_stack_lists[sample_index],angle,interp=cv.INTER_NEAREST).copy())

            if self.is_crop:

                print('cropping......')

                region_lists=[calc_max_region_list(calc_crop_region(T1_stack_list,50,5),self.forest) for T1_stack_list in T1_stack_lists]

                self.region_lists=region_lists

                T1_stack_lists=[crop(stack_list,region_lists[list_index]) for list_index,stack_list in enumerate(T1_stack_lists)]

                IR_stack_lists=[crop(stack_list,region_lists[list_index]) for list_index,stack_list in enumerate(IR_stack_lists)]

                T2_stack_lists=[crop(stack_list,region_lists[list_index]) for list_index,stack_list in enumerate(T2_stack_lists)]

                lbl_stack_lists=[crop(stack_list,region_lists[list_index]) for list_index,stack_list in enumerate(lbl_stack_lists)]

            '''

            print('len=',len(T1_stack_lists))

            T1_path_list=[]

            IR_path_list=[]

            T2_path_list=[]

            lbl_path_list=[]

            range_list=[]

            name=['1','4','5','7','14','070','148']

            f_n=['n','f']

            ang=['0','5','-5','10','-10','15','-15']

            save_path='../../../../data/'

            for sam_i,sample in enumerate(T1_stack_lists):

                for img_j,img in enumerate(sample):

                    T1_path_list.append('imgs/'+'T1/'+'{}_{}_{}_{}.png'.format(name[sam_i%7],f_n[(int(sam_i/7))%2],ang[int(sam_i/14)],img_j))

                    path=save_path+'imgs/'+'T1/'+'{}_{}_{}_{}.png'.format(name[sam_i%7],f_n[(int(sam_i/7))%2],ang[int(sam_i/14)],img_j)

                    cv.imwrite(path,img)

            for sam_i,sample in enumerate(IR_stack_lists):

                for img_j,img in enumerate(sample):

                    IR_path_list.append('imgs/'+'IR/'+'{}_{}_{}_{}.png'.format(name[sam_i%7],f_n[(int(sam_i/7))%2],ang[int(sam_i/14)],img_j))

                    path=save_path+'imgs/'+'IR/'+'{}_{}_{}_{}.png'.format(name[sam_i%7],f_n[(int(sam_i/7))%2],ang[int(sam_i/14)],img_j)

                    cv.imwrite(path,img)

            for sam_i,sample in enumerate(T2_stack_lists):

                for img_j,img in enumerate(sample):

                    T2_path_list.append('imgs/'+'T2/'+'{}_{}_{}_{}.png'.format(name[sam_i%7],f_n[(int(sam_i/7))%2],ang[int(sam_i/14)],img_j))

                    path=save_path+'imgs/'+'T2/'+'{}_{}_{}_{}.png'.format(name[sam_i%7],f_n[(int(sam_i/7))%2],ang[int(sam_i/14)],img_j)

                    cv.imwrite(path,img)

            for sam_i,sample in enumerate(lbl_stack_lists):

                for img_j,img in enumerate(sample):

                    lbl_path_list.append('lbls/'+'{}_{}_{}_{}.png'.format(name[sam_i%7],f_n[(int(sam_i/7))%2],ang[int(sam_i/14)],img_j))

                    path=save_path+'lbls/'+'{}_{}_{}_{}.png'.format(name[sam_i%7],f_n[(int(sam_i/7))%2],ang[int(sam_i/14)],img_j)

                    print(img.shape)

                    cv.imwrite(path,img)

            for sam_i,sample in enumerate(region_lists):

                for img_j,img in enumerate(sample):

                    range_list.append(img)

            range_array=np.array(range_list)

            y_min_list=range_array[:,0]

            y_max_list=range_array[:,1]

            x_min_list=range_array[:,2]

            x_max_list=range_array[:,3]

            df=pd.DataFrame({   'T1':T1_path_list,'IR':IR_path_list,'T2':T2_path_list,'lbl':lbl_path_list,

                                'y_min':y_min_list,'y_max':y_max_list,'x_min':x_min_list,'x_max':x_max_list})

            print(df)

            df.to_csv("index.csv")

            '''

            # get means

            T1mean,IRmean,T2mean=0.0,0.0,0.0

            for samples in T1_stack_lists:

                for stacks in samples:

                    T1mean=T1mean+np.mean(stacks)

            T1mean=T1mean/(len(T1_stack_lists)*len(T1_stack_lists[0]))

            print('T1 mean = ',T1mean)

            self.T1mean=T1mean

            for samples in IR_stack_lists:

                for stacks in samples:

                    IRmean=IRmean+np.mean(stacks)

            IRmean=IRmean/(len(IR_stack_lists)*len(IR_stack_lists[0]))

            print('IR mean = ',IRmean)

            self.IRmean=IRmean

            for samples in T2_stack_lists:

                for stacks in samples:

                    T2mean=T2mean+np.mean(stacks)

            T2mean=T2mean/(len(T2_stack_lists)*len(T2_stack_lists[0]))

            print('T2 mean = ',T2mean)

            self.T2mean=T2mean

            # get edegs

            print('getting edges')

            edge_stack_lists=[]

            for samples in lbl_stack_lists:

                edge_stack_lists.append(get_edge(samples))

            # transform

            if self.is_transform:

                print('transforming')

                for sample_index in range(len(T1_stack_lists)):

                    for stack_index in range(len(T1_stack_lists[0])):

                        T1_stack_lists[sample_index][stack_index],  \

                        IR_stack_lists[sample_index][stack_index],  \

                        T2_stack_lists[sample_index][stack_index],  \

                        lbl_stack_lists[sample_index][stack_index], \

                        edge_stack_lists[sample_index][stack_index]=\

                        self.transform(                             \

                        T1_stack_lists[sample_index][stack_index],  \

                        IR_stack_lists[sample_index][stack_index],  \

                        T2_stack_lists[sample_index][stack_index],  \

                        lbl_stack_lists[sample_index][stack_index], \

                        edge_stack_lists[sample_index][stack_index])

        else:

            print('validating data')

            T1_nii=to_uint8(readVol(self.val_T1path))

            IR_nii=IR_to_uint8(readVol(self.val_IRpath))

            T2_nii=to_uint8(readVol(self.val_T2path))

            lbl_nii=readVol(self.val_lblpath)

            if self.is_histeq:

                print('hist equalizing......')

                T1_nii=histeq(T1_nii)

                IR_nii=IR_nii

                T1_nii=T1_nii

            print('get stacking......')

            T1_stack_lists=get_stacked(T1_nii,self.forest)

            IR_stack_lists=get_stacked(IR_nii,self.forest)

            T2_stack_lists=get_stacked(T2_nii,self.forest)

            lbl_stack_lists=get_stacked(lbl_nii,self.forest)

            if self.is_crop:

                print('cropping......')

                region_lists=calc_max_region_list(calc_crop_region(T1_stack_lists,50,5),self.forest)

                self.region_lists=region_lists

                T1_stack_lists=crop(T1_stack_lists,region_lists)

                IR_stack_lists=crop(IR_stack_lists,region_lists)

                T2_stack_lists=crop(T2_stack_lists,region_lists)

                lbl_stack_lists=crop(lbl_stack_lists,region_lists)

            # get means

            T1mean,IRmean,T2mean=0.0,0.0,0.0

            for stacks in T1_stack_lists:

                T1mean=T1mean+np.mean(stacks)

            T1mean=T1mean/(len(T1_stack_lists))

            print('T1 mean = ',T1mean)

            self.T1mean=T1mean

            for stacks in IR_stack_lists:

                IRmean=IRmean+np.mean(stacks)

            IRmean=IRmean/(len(IR_stack_lists))

            print('IR mean = ',IRmean)

            self.IRmean=IRmean

            for stacks in T2_stack_lists:

                T2mean=T2mean+np.mean(stacks)

            T2mean=T2mean/(len(T2_stack_lists))

            print('T2 mean = ',T2mean)

            self.T2mean=T2mean

            # get edges

            print('getting edges')

            edge_stack_lists=get_edge(lbl_stack_lists)

            # transform

            if self.is_transform:

                print('transforming')

                for stack_index in range(len(T1_stack_lists)):

                    T1_stack_lists[stack_index],  \

                    IR_stack_lists[stack_index],  \

                    T2_stack_lists[stack_index],  \

                    lbl_stack_lists[stack_index], \

                    edge_stack_lists[stack_index]=\

                    self.transform(               \

                    T1_stack_lists[stack_index],  \

                    IR_stack_lists[stack_index],  \

                    T2_stack_lists[stack_index],  \

                    lbl_stack_lists[stack_index], \

                    edge_stack_lists[stack_index])

        # data ready

        self.T1_stack_lists=T1_stack_lists

        self.IR_stack_lists=IR_stack_lists

        self.T2_stack_lists=T2_stack_lists

        self.lbl_stack_lists=lbl_stack_lists

        self.edge_stack_lists=edge_stack_lists

    def __len__(self):

        return (self.is_val)and(48)or(48*6*7*2)

    def __getitem__(self,index):

        # get train or validation data

        if self.is_val==False:

            set_index=range(len(self.T1_stack_lists))

            img_index=range(len(self.T1_stack_lists[0]))

            return  \

                self.region_lists[set_index[int(index/48)]][img_index[int(index%48)]],  \

                self.T1_stack_lists[set_index[int(index/48)]][img_index[int(index%48)]],\

                self.IR_stack_lists[set_index[int(index/48)]][img_index[int(index%48)]],\

                self.T2_stack_lists[set_index[int(index/48)]][img_index[int(index%48)]],\

                self.lbl_stack_lists[set_index[int(index/48)]][img_index[int(index%48)]]

                #self.edge_stack_lists[set_index[int(index/48)]][img_index[int(index%48)]]

        else:

            img_index=range(len(self.T1_stack_lists))

            return  \

                self.region_lists[img_index[int(index)]],   \

                self.T1_stack_lists[img_index[int(index)]], \

                self.IR_stack_lists[img_index[int(index)]], \

                self.T2_stack_lists[img_index[int(index)]], \

                self.lbl_stack_lists[img_index[int(index)]]

                #self.edge_stack_lists[img_index[int(index)]]

    def transform(self,imgT1,imgIR,imgT2,lbl,edge):

        imgT1=torch.from_numpy((imgT1.transpose(2,0,1).astype(np.float)-self.T1mean)/255.0).float()

        imgIR=torch.from_numpy((imgIR.transpose(2,0,1).astype(np.float)-self.IRmean)/255.0).float()

        imgT2=torch.from_numpy((imgT2.transpose(2,0,1).astype(np.float)-self.T2mean)/255.0).float()

        lbl=torch.from_numpy(lbl.transpose(2,0,1)).long()

        edge=torch.from_numpy(edge.transpose(2,0,1)/255).float()

        return imgT1,imgIR,imgT2,lbl,edge

    def decode_segmap(self,label_mask):

        r,g,b=label_mask.copy(),label_mask.copy(),label_mask.copy()

        for ll in range(0,self.n_classes):

            r[label_mask==ll]=self.color[ll,2]

            g[label_mask==ll]=self.color[ll,1]

            b[label_mask==ll]=self.color[ll,0]

        rgb=np.zeros((label_mask.shape[0],label_mask.shape[1],3))

        rgb[:,:,0],rgb[:,:,1],rgb[:,:,2]=r,g,b

        return rgb

    def lbl_totest(self,pred):

        pred_test=np.zeros((pred.shape[0],pred.shape[1]),np.uint8)

        for ll in range(9):

            pred_test[pred==ll]=self.label_test[ll]

        return pred_test

if __name__=='__main__':

    path='../../../../data/'

    MRloader=MR18loader_CV(root=path,val_num=7,is_val=False,is_transform=True,is_flip=True,is_rotate=True,is_crop=True,is_histeq=True,forest=3)

    loader=data.DataLoader(MRloader, batch_size=1, num_workers=1, shuffle=True)

    for i,(regions,T1s,IRs,T2s,lbls) in enumerate(MRloader):

        print(i)

        #print(T1s.shape)

        #print(regions)

        #print(lbls.min())

        #print(lbls.max())

        #cv.imwrite(str(i)+'.png',T1s[:,:,1])

        #print(region)

        #print(imgT1.shape)

        #print(imgIR.shape)

        #print(imgT2.shape)

        #print(lbl.shape)

        #print('[{},{},{},{}]'.format(imgT1[0,2,40,40],imgIR[0,2,40,40],imgT2[0,2,40,40],lbl[0,2,40,40]))

        #cv.imwrite('T1-'+str(i)+'.png',imgT1[2])

        #cv.imwrite('IR-'+str(i)+'.png',imgIR[2])

        #cv.imwrite('T2-'+str(i)+'.png',imgT2[2])

        #cv.imwrite('lbl-'+str(i)+'.png',MRloader.decode_segmap(lbl[2]))