from __future__ import print_function, division
import numpy as np
import SimpleITK as sitk
from dataprocess.utils import file_name_path

flair_name = "_flair.nii.gz"
t1_name = "_t1.nii.gz"
t1ce_name = "_t1ce.nii.gz"
t2_name = "_t2.nii.gz"
mask_name = "_seg.nii.gz"


def subimage_generator(image, mask, patch_block_size, numberxy, numberz):
    """
    generate the sub images and masks with patch_block_size
    :param image:
    :param patch_block_size:
    :param stride:
    :return:
    """
    width = np.shape(image)[1]
    height = np.shape(image)[2]
    imagez = np.shape(image)[0]
    block_width = np.array(patch_block_size)[1]
    block_height = np.array(patch_block_size)[2]
    blockz = np.array(patch_block_size)[0]
    stridewidth = (width - block_width) // numberxy
    strideheight = (height - block_height) // numberxy
    stridez = (imagez - blockz) // numberz
    # step 1:if stridez is bigger 1,return  numberxy * numberxy * numberz samples
    if stridez >= 1 and stridewidth >= 1 and strideheight >= 1:
        step_width = width - (stridewidth * numberxy + block_width)
        step_width = step_width // 2
        step_height = height - (strideheight * numberxy + block_height)
        step_height = step_height // 2
        step_z = imagez - (stridez * numberz + blockz)
        step_z = step_z // 2
        hr_samples_list = []
        hr_mask_samples_list = []
        for z in range(step_z, numberz * (stridez + 1) + step_z, numberz):
            for x in range(step_width, numberxy * (stridewidth + 1) + step_width, numberxy):
                for y in range(step_height, numberxy * (strideheight + 1) + step_height, numberxy):
                    if np.max(mask[z:z + blockz, x:x + block_width, y:y + block_height]) != 0:
                        hr_samples_list.append(image[z:z + blockz, x:x + block_width, y:y + block_height])
                        hr_mask_samples_list.append(mask[z:z + blockz, x:x + block_width, y:y + block_height])
        hr_samples = np.array(hr_samples_list).reshape((len(hr_samples_list), blockz, block_width, block_height))
        hr_mask_samples = np.array(hr_mask_samples_list).reshape(
            (len(hr_mask_samples_list), blockz, block_width, block_height))
        return hr_samples, hr_mask_samples
    # step 2:other sutitation,return one samples
    else:
        nb_sub_images = 1 * 1 * 1
        hr_samples = np.zeros(shape=(nb_sub_images, blockz, block_width, block_height), dtype=np.float)
        hr_mask_samples = np.zeros(shape=(nb_sub_images, blockz, block_width, block_height), dtype=np.float)
        rangz = lambda imagez, blockz: imagez if imagez < blockz else blockz
        rangwidth = lambda width, block_width: width if width < block_width else block_width
        rangheight = lambda height, block_height: height if width < block_height else block_height
        hr_samples[0, 0:blockz, 0:block_width, 0:block_height] = image[0:rangz, 0:rangwidth, 0:rangheight]
        hr_mask_samples[0, 0:blockz, 0:block_width, 0:block_height] = mask[0:rangz, 0:rangwidth, 0:rangheight]
        return hr_samples, hr_mask_samples


def make_patch(image, mask, patch_block_size, numberxy, numberz):
    """
    make number patch
    :param image:[depth,512,512]
    :param patch_block: such as[64,128,128]
    :return:[samples,64,128,128]
    expand the dimension z range the subimage:[startpostion-blockz//2:endpostion+blockz//2,:,:]
    """
    image_subsample, mask_subsample = subimage_generator(image=image, mask=mask, patch_block_size=patch_block_size,
                                                         numberxy=numberxy, numberz=numberz)
    return image_subsample, mask_subsample


def gen_image_mask(flairimg, t1img, t1ceimg, t2img, segimg, index, shape, numberxy, numberz, trainImage, trainMask,
                   part):
    """
    :param flairimg:
    :param t1img:
    :param t1ceimg:
    :param t2img:
    :param segimg:
    :param index:
    :param shape:
    :param numberxy:
    :param numberz:
    :param trainImage:
    :param trainMask:
    :param part:
    :return:
    """
    # step 2 get subimages (numberxy*numberxy*numberz,64, 128, 128)
    sub_flairimages,sub_maskimages = make_patch(flairimg,segimg, patch_block_size=shape, numberxy=numberxy, numberz=numberz)
    sub_t1images,_ = make_patch(t1img,segimg, patch_block_size=shape, numberxy=numberxy, numberz=numberz)
    sub_t1ceimages,_ = make_patch(t1ceimg,segimg, patch_block_size=shape, numberxy=numberxy, numberz=numberz)
    sub_t2images,_ = make_patch(t2img,segimg, patch_block_size=shape, numberxy=numberxy, numberz=numberz)
    # step 3 only save subimages (numberxy*numberxy*numberz,64, 128, 128)
    samples, imagez, height, width = np.shape(sub_flairimages)[0], np.shape(sub_flairimages)[1], \
                                     np.shape(sub_flairimages)[2], np.shape(sub_flairimages)[3]
    for j in range(samples):
        sub_masks = sub_maskimages.astype(np.float)
        sub_masks = np.clip(sub_masks, 0, 255).astype('uint8')
        if np.max(sub_masks[j, :, :, :]) != 0:
            """
            merage 4 model image into 4 channel (imagez,width,height,channel)
            """
            fourmodelimagearray = np.zeros((imagez, height, width, 4), np.float)
            filepath1 = trainImage + "\\" + str(part) + "_" + str(index) + "_" + str(j) + ".npy"
            filepath = trainMask + "\\" + str(part) + "_" + str(index) + "_" + str(j) + ".npy"
            flairimage = sub_flairimages[j, :, :, :]
            flairimage = flairimage.astype(np.float)
            fourmodelimagearray[:, :, :, 0] = flairimage
            t1image = sub_t1images[j, :, :, :]
            t1image = t1image.astype(np.float)
            fourmodelimagearray[:, :, :, 1] = t1image
            t1ceimage = sub_t1ceimages[j, :, :, :]
            t1ceimage = t1ceimage.astype(np.float)
            fourmodelimagearray[:, :, :, 2] = t1ceimage
            t2image = sub_t2images[j, :, :, :]
            t2image = t2image.astype(np.float)
            fourmodelimagearray[:, :, :, 3] = t2image
            np.save(filepath1, fourmodelimagearray)
            np.save(filepath, sub_masks[j, :, :, :])


def normalize(slice, bottom=99, down=1):
    """
    normalize image with mean and std for regionnonzero,and clip the value into range
    :param slice:
    :param bottom:
    :param down:
    :return:
    """
    b = np.percentile(slice, bottom)
    t = np.percentile(slice, down)
    slice = np.clip(slice, t, b)

    image_nonzero = slice[np.nonzero(slice)]
    if np.std(slice) == 0 or np.std(image_nonzero) == 0:
        return slice
    else:
        tmp = (slice - np.mean(image_nonzero)) / np.std(image_nonzero)
        # since the range of intensities is between 0 and 5000 ,
        # the min in the normalized slice corresponds to 0 intensity in unnormalized slice
        # the min is replaced with -9 just to keep track of 0 intensities
        # so that we can discard those intensities afterwards when sampling random patches
        tmp[tmp == tmp.min()] = -9
        return tmp


def prepare3dtraindata(pathhgg_list, bratshgg_path, trainImage, trainMask, shape=(16, 256, 256), numberxy=3,
                       numberz=20, part=1):
    """
    :param pathhgg_list:
    :param bratshgg_path:
    :param trainImage:
    :param trainMask:
    :param shape:
    :param numberxy:
    :param numberz:
    :return:
    """
    """
    load flair_image,t1_image,t1ce_image,t2_image,mask_image
    """
    for subsetindex in range(len(pathhgg_list)):
        brats_subset_path = bratshgg_path + "/" + str(pathhgg_list[subsetindex]) + "/"
        flair_image = brats_subset_path + str(pathhgg_list[subsetindex]) + flair_name
        t1_image = brats_subset_path + str(pathhgg_list[subsetindex]) + t1_name
        t1ce_image = brats_subset_path + str(pathhgg_list[subsetindex]) + t1ce_name
        t2_image = brats_subset_path + str(pathhgg_list[subsetindex]) + t2_name
        mask_image = brats_subset_path + str(pathhgg_list[subsetindex]) + mask_name
        flair_src = sitk.ReadImage(flair_image, sitk.sitkInt16)
        t1_src = sitk.ReadImage(t1_image, sitk.sitkInt16)
        t1ce_src = sitk.ReadImage(t1ce_image, sitk.sitkInt16)
        t2_src = sitk.ReadImage(t2_image, sitk.sitkInt16)
        mask = sitk.ReadImage(mask_image, sitk.sitkUInt8)
        flair_array = sitk.GetArrayFromImage(flair_src)
        t1_array = sitk.GetArrayFromImage(t1_src)
        t1ce_array = sitk.GetArrayFromImage(t1ce_src)
        t2_array = sitk.GetArrayFromImage(t2_src)
        mask_array = sitk.GetArrayFromImage(mask)
        """
        normalize every model image,mask is not normalization
        """
        flair_array = normalize(flair_array)
        t1_array = normalize(t1_array)
        t1ce_array = normalize(t1ce_array)
        t2_array = normalize(t2_array)

        gen_image_mask(flair_array, t1_array, t1ce_array, t2_array, mask_array, subsetindex, shape=shape,
                       numberxy=numberxy, numberz=numberz, trainImage=trainImage, trainMask=trainMask, part=part)


def preparetraindata():
    """
    :return:
    """
    bratshgg_path = "D:\Data\\brats18\HGG"
    bratslgg_path = "D:\Data\\brats18\LGG"
    trainImage = "D:\Data\\brats18\\train\Image"
    trainMask = "D:\Data\\brats18\\train\Mask"
    pathhgg_list = file_name_path(bratshgg_path)
    pathlgg_list = file_name_path(bratslgg_path)

    prepare3dtraindata(pathhgg_list, bratshgg_path, trainImage, trainMask, (64, 128, 128), 3, 15, 1)
    prepare3dtraindata(pathlgg_list, bratslgg_path, trainImage, trainMask, (64, 128, 128), 3, 15, 2)


#preparetraindata()