#Preprocessing file

import os

import numpy as np

import nibabel as nib

import time

import math

import random

import glob

import collections

import sys

import random

from random import randint

import time

import datetime

import augmentations as AUG

#------------------------------------------------------------------------------

#------------------------------------------------------------------------------

#Basic Utils

#------------------------------------------------------------------------------

#------------------------------------------------------------------------------

def getTime():

    return datetime.datetime.now().strftime("%Y-%m-%d-%H-%M")

def printDimensions(img_path = 'pre/FLAIR.nii.gz', segm_path = 'wmh.nii.gz', data_folder = '../Data/MS2017a/'):

    scan_folders = glob.glob(data_folder + 'scans/*')

    for sf in scan_folders:

        file_num = os.path.basename(sf)

        img = nib.load(os.path.join(sf, img_path))

        print(file_num, img.shape)

def extractMeanDataStats(size = [200, 200, 100], 

                        postfix = '_200x200x100orig', 

                        main_folder_path = '../../Data/MS2017b/', 

                        ):

    scan_folders = glob.glob(main_folder_path + 'scans/*')

    img_path = 'pre/FLAIR' + postfix + '.nii.gz'

    segm_path = 'wmh' + postfix + '.nii.gz'

    shape_ = [len(scan_folders), size[0], size[1], size[2]]

    arr = np.zeros(shape_)

    for i, sf in enumerate(scan_folders):

        arr[i, :,:,:] =  numpyFromScan(os.path.join(sf,img_path)).squeeze()

    arr /= len(scan_folders)

    means = np.mean(arr)

    stds = np.std(arr, axis = 0)

    np.save(main_folder_path + 'extra_data/std' + postfix, stds)

    np.save(main_folder_path + 'extra_data/mean' + postfix, means)

def getExperimentInfo(experiment_str):

    exp_arr = experiment_str.split('_')

    isPriv = bool(int(exp_arr[1]))

    withASPP = bool(int(exp_arr[2]))

    dilations_str = exp_arr[0]

    dilation_arr = [int(i) for i in dilations_str.split('x')]

    return dilation_arr, isPriv, withASPP

#------------------------------------------------------------------------------

#------------------------------------------------------------------------------

##GENERAL PREPROCESSING

#------------------------------------------------------------------------------

#------------------------------------------------------------------------------

#convert .nii.gz file to scan

def numpyFromScan(path, get_affine = False, makebin = False):

    img = nib.load(path)

    img_np = img.get_data()

    #reshape to size1 x size2  - > size1 x size2 x 1

    img_np = np.expand_dims(img_np, axis=len(img_np.shape))

    #img_np = img_np.reshape([img_np.shape[0], img_np.shape[1], 1])

    if makebin:

        img_np[img_np == 2] = 0

    if get_affine:

        return img_np, img.get_affine()

    return img_np

def saveScan(img_np, affine, save_path, header = None):

    if header:

        nft_img = nib.Nifti1Image(img_np, affine, header = header)

    else:

        nft_img = nib.Nifti1Image(img_np, affine)

    nib.save(nft_img, save_path)

#get list of validation/train sets

def splitTrainVal(train_fraction, data_folder = '../Data/MS2017a/'):

    scan_folders = glob.glob(data_folder + 'scans/*')

    num_scans = len(scan_folders)

    indices = np.random.permutation(num_scans)

    train_indices = indices[0:int(num_scans*train_fraction)]

    val_indices = indices[int(num_scans*train_fraction):]

    train_scan_folders = [scan_folders[i] for i in train_indices]

    val_scan_folders = [scan_folders[i] for i in val_indices]

    return train_scan_folders, val_scan_folders

#call this once to split training data

def generateTrainValFile(train_fraction, main_folder = '../Data/MS2017a/', postfix = ''):

    train_folders, val_folders = splitTrainVal(0.8, data_folder=main_folder)

    img_path = '/pre/FLAIR' + postfix + '.nii.gz'

    train_folder_names = [train_folders[i].split(main_folder)[1] + img_path for i in range(len(train_folders))]

    val_folder_names = [val_folders[i].split(main_folder)[1]  + img_path for i in range(len(val_folders))]

    f_train = open(main_folder + 'train' + postfix + '.txt', 'w+')

    f_val = open(main_folder + 'val' + postfix + '.txt', 'w+')

    for fn in train_folder_names:

        f_train.write(fn + '\n')

    for fn in val_folder_names:

        f_val.write(fn + '\n')

    f_train.close()

    f_val.close()

def read_file(path_to_file, pretext = ''):

    with open(path_to_file) as f:

        img_list = []

        for line in f:

            img_list.append(pretext + line[:-1])

    return img_list

def generateTestFile(folder):

    pass

#------------------------------------------------------------------------------

#------------------------------------------------------------------------------

##END GENERAL PREPROCESSING

#------------------------------------------------------------------------------

#------------------------------------------------------------------------------

#func1()

#simpleSplitTrainVal(0.8)

#generateTrainValFile(0.8)

#------------------------------------------------------------------------------

#------------------------------------------------------------------------------

##SLICES PREPROCESSING

#------------------------------------------------------------------------------

#------------------------------------------------------------------------------

#generate a slices folder containing all slices

def generateImgSlicesFolder(data_folder = '../Data/MS2017a/scans/'):

    scan_folders = glob.glob(data_folder + '*')

    for sf in scan_folders:

        slice_dir_path = os.path.join(sf, 'slices/')

        if not os.path.exists(slice_dir_path):

            print('Creating directory at:' , slice_dir_path)

            os.makedirs(slice_dir_path)

        img = nib.load(os.path.join(sf, 'pre/FLAIR.nii.gz'))

        img_np = img.get_data()

        img_affine = img.affine

        print(sf)

        print('The img shape', img_np.shape[2])

        for i in range(img_np.shape[2]):

            slice_img_np = img_np[:,:,i]

            nft_img = nib.Nifti1Image(slice_img_np, img_affine)

            nib.save(nft_img, slice_dir_path + 'FLAIR_' + str(i) + '.nii.gz')

            if os.path.basename(sf) == '0':

                slice_img = nib.load(slice_dir_path + 'FLAIR_' + str(i) + '.nii.gz').get_data() / 5

                print('DID I GET HERE?')

                print('Writing to', str(i) + '.jpg')

def generateGTSlicesFolder(data_folder = '../Data/MS2017a/scans/'):

    scan_folders = glob.glob(data_folder + '*')

    for sf in scan_folders:

        slice_dir_path = os.path.join(sf, 'gt_slices/')

        if not os.path.exists(slice_dir_path):

            print('Creating directory at:' , slice_dir_path)

            os.makedirs(slice_dir_path)

        img = nib.load(os.path.join(sf, 'wmh.nii.gz'))

        img_np = img.get_data()

        img_affine = img.affine

        print(sf)

        print('The img shape', img_np.shape[2])

        for i in range(img_np.shape[2]):

            slice_img_np = img_np[:,:,i]

            nft_img = nib.Nifti1Image(slice_img_np, img_affine)

            nib.save(nft_img, slice_dir_path + 'wmh_' + str(i) + '.nii.gz')

            if os.path.basename(sf) == '0':

                slice_img = nib.load(slice_dir_path + 'wmh_' + str(i) + '.nii.gz').get_data() * 256

                #cv2.imwrite('temp/' + str(i) + '.jpg', slice_img)

def splitTrainVal_Slices(train_fraction, data_folder = '../Data/MS2017a/scans/'):

    scan_folders = glob.glob(data_folder + '/*/slices/*')

    num_scans = len(scan_folders)

    indices = np.random.permutation(num_scans)

    train_indices = indices[0:int(num_scans*train_fraction)]

    val_indices = indices[int(num_scans*train_fraction):]

    train_scan_folders = [scan_folders[i] for i in train_indices]

    val_scan_folders = [scan_folders[i] for i in val_indices]

    return train_scan_folders, val_scan_folders

def generateTrainValFile_Slices(train_fraction, main_folder = '../Data/MS2017a/'):

    train_folders, val_folders = splitTrainVal_Slices(0.8)

    train_folder_names = [train_folders[i].split(main_folder)[1] for i in range(len(train_folders))]

    val_folder_names = [val_folders[i].split(main_folder)[1] for i in range(len(val_folders))]

    f_train = open(main_folder + 'train_slices.txt', 'w+')

    f_val = open(main_folder + 'val_slices.txt', 'w+')

    for fn in train_folder_names:

        f_train.write(fn + '\n')

    for fn in val_folder_names:

        f_val.write(fn + '\n')

    f_train.close()

    f_val.close()

#Use this to save the images quickly (for testing purposes)

def quickSave(img, wmh, gif, n):

    nft_img = nib.Nifti1Image(img.squeeze(), np.eye(4))

    nib.save(nft_img, n + '_img.nii.gz')

    nft_img = nib.Nifti1Image(wmh.squeeze(), np.eye(4))

    nib.save(nft_img, n + '_wmh.nii.gz')

    if gif is not None:

        nft_img = nib.Nifti1Image(gif.squeeze(), np.eye(4))

        nib.save(nft_img, n + '_gif.nii.gz')

#------------------------------------------------------------------------------

#END OF SLICES PREPROCESSING

#------------------------------------------------------------------------------

#------------------------------------------------------------------------------

#3D PREPROCESSING

#------------------------------------------------------------------------------

#go through every 3D object from training set and every patch of size NxNxN

#save resulting 3D volumes in one of the two folders based on what the center pixel of the image is

def extractCenterPixelPatches(N = 33, main_folder = '../Data/MS2017b/', postfix = ''):

    if N % 2 != 1:

        print('N must be odd')

        sys.exit()

    img_path = 'pre/FLAIR' + postfix + '.nii.gz'

    segm_path = 'wmh' + postfix + '.nii.gz'

    folders = ['lesion', 'other']

    patch_folder_path = os.path.join(main_folder, 'centerPixelPatches' + postfix + '_' + str(N))

    if not os.path.exists(patch_folder_path):

        for f in folders:

            os.makedirs(os.path.join(main_folder, patch_folder_path, f))

    scan_folders = glob.glob(main_folder + 'scans/*')

    counter = 0

    f_lesion_txt = open(os.path.join(patch_folder_path, 'lesion', 'center_locs.txt'), 'w+')

    f_other_txt = open(os.path.join(patch_folder_path, 'other', 'center_locs.txt'), 'w+')

    #This is only for training data

    img_list = read_file(main_folder + 'train' + postfix + '.txt', pretext = main_folder)

    print('Gathering training images from ' + main_folder + 'train' + postfix + '.txt')

    #remove pre/FLAIR_s.nii.gz from path. Only want up to folder name

    scan_folders = [img_list[i][:-len(img_path)] for i in range(len(img_list))]

    num_lesion = 0

    num_other = 0

    num_background = 0

    for sf in scan_folders:

        folder_num = sf.split('/')[-2]

        #read the FLAIR img

        img = nib.load(os.path.join(sf, img_path))

        img_affine = img.affine

        img_np = img.get_data()

        #read the wmh img

        wmh = nib.load(os.path.join(sf, segm_path))

        wmh_affine = wmh.affine

        wmh_np = wmh.get_data()

        #reshape to size1 x size2 -> size1 x size2 x 1

        img_np = img_np.reshape([img_np.shape[0], img_np.shape[1], img_np.shape[2], 1])

        wmh_np = wmh_np.reshape([wmh_np.shape[0], wmh_np.shape[1], wmh_np.shape[2], 1])

        #loop through every size

        for x in range(img_np.shape[0] - N + 1):

            for y in range(img_np.shape[1] - N + 1):

                for z in range(img_np.shape[2] - N + 1):

                    wmh_patch = wmh_np[x:x+N, y:y+N, z:z+N]

                    M = (N + 1) / 2

                    center_pixel = wmh_patch[M,M,M]

                    #folder_num | x | y | z

                    location_name = str(folder_num) + '|' + str(x) + '|' + str(y) + '|' + str(z)

                    if center_pixel == 1:

                        num_lesion += 1

                        f_lesion_txt.write(location_name + '\n')

                    elif center_pixel == 2:

                        num_other += 1

                        f_other_txt.write(location_name + '\n')

        counter += 1

        print(str(counter) + ' / ' + str(len(scan_folders)))

    f_lesion_txt.close()

    f_other_txt.close()

    print('Num background: ', num_background)

    print('Num lesion', num_lesion)

    print('Num other', num_other)

    print('Done!')

        #TEMPORARY

        #if sf.split('/')

#during training we will sample uniformly between the two folders (uniformly select folder and uniformly select training sample)

def extractPatchBatch(batch_size, patch_size, img_list, 

                        onlyLesions = False, center_pixel = False, 

                        main_folder_path = '../../Data/MS2017b/', 

                        postfix = '', with_priv = False):

    img_b = np.zeros([batch_size, 1, patch_size, patch_size, patch_size])

    label_b = np.zeros([batch_size, 1, patch_size, patch_size, patch_size]) 

    gif_b = None

    if with_priv:

        gif_b = np.zeros([batch_size, 1, patch_size, patch_size, patch_size])   

    for i in range(batch_size):

        if center_pixel:

            center_pixel_folder_path = main_folder_path + 'centerPixelPatches' + postfix + '_' + str(patch_size)

            locs_lesion = open(os.path.join(center_pixel_folder_path, 'lesion', 'center_locs.txt')).readlines()

            locs_other =  open(os.path.join(center_pixel_folder_path, 'other', 'center_locs.txt')).readlines()

            img_patch, gt_patch, gif_patch = getCenterPixelPatch(patch_size, img_list, locs_lesion, locs_other, 

                                                            onlyLesions, main_folder_path, postfix, with_priv)

        else:

            img_patch, gt_patch, gif_patch = getRandomPatch(patch_size, img_list, onlyLesions, main_folder_path, postfix, with_priv)

        img_b[i, :,:,:,:] = img_patch

        label_b[i, :,:,:,:] = gt_patch

        if with_priv:

            gif_b[i, :,:,:,:] = gif_patch

    return img_b, label_b, gif_b    

##################################################################################################################################

##################################################################################################################################

##################################Patch functions####################################################################

def getRandomPatch(patch_size, img_list, onlyLesions, main_folder_path, postfix, with_priv = False):

    img_str = img_list[randint(0, len(img_list)- 1)].rstrip()

    gt_str = img_str.replace('slices', 'gt_slices').replace('FLAIR', 'wmh').replace('/pre','')

    img_np = numpyFromScan(os.path.join(main_folder_path, img_str))

    gt_np = numpyFromScan(os.path.join(main_folder_path, gt_str), makebin = onlyLesions)

    img_np = img_np.transpose(3,0,1,2)

    gt_np = gt_np.transpose(3,0,1,2)

    img_dims = img_np.shape

    x = randint(0, img_dims[1]-patch_size-1)

    y = randint(0, img_dims[2]-patch_size-1)

    z = randint(0, img_dims[3]-patch_size-1)

    img_np_patch = img_np[:, x:x+patch_size, y:y+patch_size, z:z+patch_size]

    gt_np_patch = gt_np[:, x:x+patch_size, y:y+patch_size, z:z+patch_size]

    if with_priv:

        gif_str = img_str.replace('scans', 'gifs').replace('FLAIR','parcellation').replace('/pre','')

        gif_np = numpyFromScan(os.path.join(main_folder_path, gif_str))

        gif_np = gif_np.transpose(3,0,1,2)

        gif_np_patch = gif_np[:, x:x+patch_size, y:y+patch_size, z:z+patch_size]

        return img_np_patch, gt_np_patch, gif_np_patch

    #draw 3 numbers between patch_size 

    return img_np_patch, gt_np_patch, None

#XXX not implemented for onlyLesions = True

def getCenterPixelPatch(patch_size, img_list, locs_lesion, locs_other, 

                            onlyLesions, main_folder_path, postfix, with_priv = False):

    b = random.uniform(0.5, 3.5)

    #segm class = 1

    if b < 1.5:

        loc_str = locs_lesion[randint(0, len(locs_lesion) - 1)].rstrip()

    #segm class = 2

    elif b > 1.5 and b < 2.5 and (not onlyLesions):

        loc_str = locs_other[randint(0,len(locs_other) - 1)].rstrip()

    #segm class = 3

    else:

        loc_str = getBackgroundLoc(patch_size, img_list, onlyLesions, main_folder_path)

    #extract patch given folder number, location of top left edge and patch size

    #---------------------------------------------------------------------------

    folder_num_str, x, y, z = parseLocStr(loc_str)

    img_type_path = 'pre/FLAIR' + postfix + '.nii.gz'

    gt_type_path = 'wmh' + postfix + '.nii.gz'

    #read the file

    img_np = numpyFromScan(os.path.join(main_folder_path, 'scans', folder_num_str, img_type_path))

    gt_np = numpyFromScan(os.path.join(main_folder_path, 'scans', folder_num_str, gt_type_path), makebin = onlyLesions)

    #extract the patch

    patch_img_np = img_np[x:x+patch_size, y:y+patch_size, z:z+patch_size, :]

    patch_gt_np = gt_np[x:x+patch_size, y:y+patch_size, z:z+patch_size, :]

    #reshape to 1 x dim1 x dim2 x dim3

    patch_img_np = patch_img_np.transpose((3,0,1,2))

    patch_gt_np = patch_gt_np.transpose((3,0,1,2))

    if with_priv:

        gif_type_path = 'parcellation' + postfix + '.nii.gz'

        gif_np = numpyFromScan(os.path.join(main_folder_path, 'gifs', folder_num_str, gif_type_path))

        patch_gif_np = gif_np[x:x+patch_size, y:y+patch_size, z:z+patch_size, :]

        patch_gif_np = patch_gif_np.transpose((3,0,1,2))

        return patch_img_np, patch_gt_np, patch_gif_np

    return patch_img_np, patch_gt_np, None

def getBackgroundLoc(patch_size, img_list, onlyLesions, main_folder_path):

    num_generated = 0

    found_background = False

    #choose a random 3D image

    img_str = img_list[randint(0, len(img_list)- 1)].rstrip()

    curr_wmh_str = img_str.replace('slices', 'gt_slices').replace('FLAIR', 'wmh').replace('/pre','')

    wmh_np = numpyFromScan(os.path.join(main_folder_path, curr_wmh_str), makebin = onlyLesions)

    img_dims = wmh_np.shape

    folder_num = curr_wmh_str.split('/')[1]

    #print('THE FOLDER NUM', folder_num)

    while not found_background:

        x = randint(0, img_dims[0]-patch_size-1)

        y = randint(0, img_dims[1]-patch_size-1)

        z = randint(0, img_dims[2]-patch_size-1)

        #Load and check center pixel

        if wmh_np[x + ((patch_size - 1)/2), y + ((patch_size-1)/2), z + ((patch_size-1)/2)] == 0:

            found_background = True

            loc_str = str(folder_num) + '|' + str(x) + '|' + str(y) + '|' + str(z)

            return loc_str

        num_generated += 1

    #print('Num generated until a background batch was found: ', num_generated)

    return loc_str

def parseLocStr(loc_str):

    s = loc_str.split('|')

    return s[0], int(s[1]), int(s[2]), int(s[3])

##################################################################################################################################

##################################################################################################################################

#[518803341, 1496491, 217508]

def classCount(main_folder_path = '../Data/MS2017b/', img_type_path = 'pre/FLAIR_s.nii.gz', gt_type_path = 'wmh_s.nii.gz'):

        scan_folders = glob.glob(main_folder_path + 'scans/*')

        nums = [0, 0 ,0]

        for sf in scan_folders:

            wmh_np = numpyFromScan(os.path.join(sf, gt_type_path))

            unique, counts = np.unique(wmh_np, return_counts= True)

            d = dict(zip(unique, counts))

            for i in range(3):

                try: 

                    nums[i] += d[i]

                except KeyError:

                    pass

        print nums

def extractImgBatch(batch_size, img_list, img_size, onlyLesions = False, main_folder_path = '../Data/MS2017b/', with_priv = False):

    img_b = np.zeros([batch_size, 1, img_size[0], img_size[1], img_size[2]])

    label_b = np.zeros([batch_size, 1, img_size[0], img_size[1], img_size[2]])

    if with_priv:

        gif_b = np.zeros([batch_size, 1, img_size[0], img_size[1], img_size[2]])

    for i in range(batch_size):

        img_str = img_list[randint(0, len(img_list)-1)]

        img_np = numpyFromScan(os.path.join(main_folder_path, img_str))

        img_np = img_np.transpose((3,0,1,2))

        img_b[i, :,:,:,:] = img_np

        wmh_str = img_str.replace('slices', 'gt_slices').replace('FLAIR', 'wmh').replace('/pre','')

        gt_np = numpyFromScan(os.path.join(main_folder_path, wmh_str))

        gt_np = gt_np.transpose((3,0,1,2))

        label_b[i, :,:,:,:] = gt_np

        if with_priv:

            gif_str = img_str.replace('scans','gifs').replace('FLAIR', 'parcellation').replace('/pre','')

            gif_np = numpyFromScan(os.path.join(main_folder_path, gif_str))

            gif_np = gt_np.transpose((3,0,1,2))

            gif_b[i, :,:,:,:] = gif_np

    if with_priv:

        return img_b, label_b, gif_b

    return img_b, label_b, None

#------------------------------------------------------------------------------

#END OF 3D PREPROCESSING

#------------------------------------------------------------------------------

#3D

#generateTrainValFile(0.8, img_path = '/pre/FLAIR_s.nii.gz', main_folder = '../Data/MS2017b/')

#generateTrainValFile(0.8, img_path = '/pre/FLAIR_256x256x256orig.nii.gz', main_folder = '../Data/MS2017b/', postfix='_256x256x256orig')

#generateTrainValFile(0.8, main_folder = '../../Data/MS2017b/', postfix='_200x200x100orig')

#extractCenterPixelPatches(N = 81, main_folder = '../../Data/MS2017b/', postfix = '_200x200x100orig')

#extractCenterPixelPatches(N = 71, main_folder = '../../Data/MS2017b/', postfix = '_200x200x100orig')

#printDimensions(img_path = 'pre/FLAIR_s128.nii.gz', segm_path = 'wmh_s128.nii.gz', data_folder = '../Data/MS2017b/')

#printDimensions(img_path = 'pre/FLAIR.nii.gz', segm_path = 'wmh.nii.gz', data_folder = '../../Data/MS2017b/')

#printDimensions(img_path = 'pre/FLAIR256x256x256orig.nii.gz', segm_path = 'wmh.nii.gz', data_folder = '../../Data/MS2017b/')

#extractCenterPixelPatches()

#extractCenterPixelPatches(N = 91)

#extractCenterPixelPatches(N = 101, main_folder = '../../Data/MS2017b/', postfix = '_256x256x256orig')

#generateTrainValFile(0.8, img_path = '/pre/FLAIR_s128.nii.gz', main_folder = '../Data/MS2017b/', postfix='128x128x64')

#classCount()

#extractMeanDataStats()