from __future__ import print_function, division

import os

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

from tensorflow.python.client import device_lib

print(device_lib.list_local_devices())

from Vnet.model_vnet3d_multilabel import Vnet3dModuleMultiLabel

from Vnet.util import getLargestConnectedCompont

from dataprocess.utils import calcu_dice, file_name_path

from dataprocess.data3dprepare import normalize

import numpy as np

import SimpleITK as sitk

def inference():

    """

    Vnet network segmentation brats fine segmatation

    :return:

    """

    channel = 4

    numclass = 4

    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"

    out_mask_name = "_outseg.nii.gz"

    # step1 init vnet model

    depth_z = 48

    Vnet3d = Vnet3dModuleMultiLabel(240, 240, depth_z, channels=channel, numclass=numclass,

                                    costname=("categorical_dice",), inference=True,

                                    model_path="log\segmeation2mm\weighted_categorical_crossentropy\model\Vnet3d.pd-10000")

    brats_path = "D:\Data\\brats18\\test"

    # step2 get all test image path

    dice_values0 = []

    dice_values1 = []

    dice_values2 = []

    dice_values3 = []

    path_list = file_name_path(brats_path)

    # step3 get test image(4 model) and mask

    for subsetindex in range(len(path_list)):

        # step4 load test image(4 model) and mask as ndarray

        brats_subset_path = brats_path + "/" + str(path_list[subsetindex]) + "/"

        flair_image = brats_subset_path + str(path_list[subsetindex]) + flair_name

        t1_image = brats_subset_path + str(path_list[subsetindex]) + t1_name

        t1ce_image = brats_subset_path + str(path_list[subsetindex]) + t1ce_name

        t2_image = brats_subset_path + str(path_list[subsetindex]) + t2_name

        mask_image = brats_subset_path + str(path_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)

        label = sitk.GetArrayFromImage(mask)

        # step5 mormazalation test image(4 model) and merage to 4 channels ndarray

        flair_array = normalize(flair_array)

        t1_array = normalize(t1_array)

        t1ce_array = normalize(t1ce_array)

        t2_array = normalize(t2_array)

        imagez, height, width = np.shape(flair_array)[0], np.shape(flair_array)[1], np.shape(flair_array)[2]

        fourmodelimagearray = np.zeros((imagez, height, width, channel), np.float)

        fourmodelimagearray[:, :, :, 0] = flair_array

        fourmodelimagearray[:, :, :, 1] = t1_array

        fourmodelimagearray[:, :, :, 2] = t1ce_array

        fourmodelimagearray[:, :, :, 3] = t2_array

        ys_pd_array = np.zeros((imagez, height, width), np.uint8)

        # step6 predict test image(4 model)

        last_depth = 0

        for depth in range(0, imagez // depth_z, 1):

            patch_xs = fourmodelimagearray[depth * depth_z:(depth + 1) * depth_z, :, :, :]

            pathc_pd = Vnet3d.prediction(patch_xs)

            ys_pd_array[depth * depth_z:(depth + 1) * depth_z, :, :] = pathc_pd

            last_depth = depth

        if imagez != depth_z * last_depth:

            patch_xs = fourmodelimagearray[(imagez - depth_z):imagez, :, :, :]

            pathc_pd = Vnet3d.prediction(patch_xs)

            ys_pd_array[(imagez - depth_z):imagez, :, :] = pathc_pd

        ys_pd_array = np.clip(ys_pd_array, 0, 255).astype('uint8')

        all_ys_pd_array = ys_pd_array.copy()

        all_ys_pd_array[ys_pd_array != 0] = 1

        outmask = getLargestConnectedCompont(sitk.GetImageFromArray(all_ys_pd_array))

        ys_pd_array[outmask == 0] = 0

        # step7 calcu test mask and predict mask dice value

        batch_ys = label.copy()

        batch_ys[label == 4] = 3

        dice_value0 = 0

        dice_value1 = 0

        dice_value2 = 0

        dice_value3 = 0

        for num_class in range(4):

            ys_pd_array_tmp = ys_pd_array.copy()

            batch_ys_tmp = batch_ys.copy()

            ys_pd_array_tmp[ys_pd_array == num_class] = 1

            batch_ys_tmp[label == num_class] = 1

            if num_class == 0:

                dice_value0 = calcu_dice(ys_pd_array_tmp, batch_ys_tmp, 1)

            if num_class == 1:

                dice_value1 = calcu_dice(ys_pd_array_tmp, batch_ys_tmp, 1)

            if num_class == 2:

                dice_value2 = calcu_dice(ys_pd_array_tmp, batch_ys_tmp, 1)

            if num_class == 3:

                dice_value3 = calcu_dice(ys_pd_array_tmp, batch_ys_tmp, 1)

        print("index,dice:", (subsetindex, dice_value0, dice_value1, dice_value2, dice_value3))

        dice_values0.append(dice_value0)

        dice_values1.append(dice_value1)

        dice_values2.append(dice_value2)

        dice_values3.append(dice_value3)

        # step8 out put predict mask

        ys_pd_array = ys_pd_array.astype('float')

        outputmask = np.zeros((imagez, height, width), np.uint8)

        outputmask[ys_pd_array == 1] = 1

        outputmask[ys_pd_array == 2] = 2

        outputmask[ys_pd_array == 3] = 4

        ys_pd_itk = sitk.GetImageFromArray(outputmask)

        ys_pd_itk.SetSpacing(mask.GetSpacing())

        ys_pd_itk.SetOrigin(mask.GetOrigin())

        ys_pd_itk.SetDirection(mask.GetDirection())

        out_mask_image = brats_subset_path + str(path_list[subsetindex]) + out_mask_name

        sitk.WriteImage(ys_pd_itk, out_mask_image)

    average0 = sum(dice_values0) / len(dice_values0)

    average1 = sum(dice_values1) / len(dice_values1)

    average2 = sum(dice_values2) / len(dice_values2)

    average3 = sum(dice_values3) / len(dice_values3)

    print("average dice:", (average0, average1, average2, average3))

inference()