from __future__ import print_function, division

import SimpleITK as sitk

import numpy as np

import cv2

import os

from glob import glob

def load_itk(filename):

    """

    load mhd files and normalization 0-255

    :param filename:

    :return:

    """

    rescalFilt = sitk.RescaleIntensityImageFilter()

    rescalFilt.SetOutputMaximum(255)

    rescalFilt.SetOutputMinimum(0)

    # Reads the image using SimpleITK

    itkimage = rescalFilt.Execute(sitk.Cast(sitk.ReadImage(filename), sitk.sitkFloat32))

    return itkimage

def removesmallConnectedCompont(sitk_maskimg, rate=0.5):

    cc = sitk.ConnectedComponent(sitk_maskimg)

    stats = sitk.LabelIntensityStatisticsImageFilter()

    stats.Execute(cc, sitk_maskimg)

    maxlabel = 0

    maxsize = 0

    for l in stats.GetLabels():

        size = stats.GetPhysicalSize(l)

        if maxsize < size:

            maxlabel = l

            maxsize = size

    not_remove = []

    for l in stats.GetLabels():

        size = stats.GetPhysicalSize(l)

        if size > maxsize * rate:

            not_remove.append(l)

    labelmaskimage = sitk.GetArrayFromImage(cc)

    outmask = labelmaskimage.copy()

    outmask[labelmaskimage != maxlabel] = 0

    for i in range(len(not_remove)):

        outmask[labelmaskimage == not_remove[i]] = 255

    return outmask

def getLargestConnectedCompont(sitk_maskimg):

    cc = sitk.ConnectedComponent(sitk_maskimg)

    stats = sitk.LabelIntensityStatisticsImageFilter()

    stats.Execute(cc, sitk_maskimg)

    maxlabel = 0

    maxsize = 0

    for l in stats.GetLabels():

        size = stats.GetPhysicalSize(l)

        if maxsize < size:

            maxlabel = l

            maxsize = size

    labelmaskimage = sitk.GetArrayFromImage(cc)

    outmask = labelmaskimage.copy()

    outmask[labelmaskimage == maxlabel] = 255

    outmask[labelmaskimage != maxlabel] = 0

    return outmask

def morphologicaloperation(sitk_maskimg, kernelsize, name='open'):

    if name == 'open':

        morphoimage = sitk.BinaryMorphologicalOpening(sitk_maskimg, [kernelsize, kernelsize, kernelsize])

        labelmaskimage = sitk.GetArrayFromImage(morphoimage)

        outmask = labelmaskimage.copy()

        outmask[labelmaskimage == 1.0] = 255

        return outmask

    if name == 'close':

        morphoimage = sitk.BinaryMorphologicalClosing(sitk_maskimg, [kernelsize, kernelsize, kernelsize])

        labelmaskimage = sitk.GetArrayFromImage(morphoimage)

        outmask = labelmaskimage.copy()

        outmask[labelmaskimage == 1.0] = 255

        return outmask

    if name == 'dilate':

        morphoimage = sitk.BinaryDilate(sitk_maskimg, [kernelsize, kernelsize, kernelsize])

        labelmaskimage = sitk.GetArrayFromImage(morphoimage)

        outmask = labelmaskimage.copy()

        outmask[labelmaskimage == 1.0] = 255

        return outmask

    if name == 'erode':

        morphoimage = sitk.BinaryErode(sitk_maskimg, [kernelsize, kernelsize, kernelsize])

        labelmaskimage = sitk.GetArrayFromImage(morphoimage)

        outmask = labelmaskimage.copy()

        outmask[labelmaskimage == 1.0] = 255

        return outmask

def gettestiamge():

    src = load_itk("D:\Data\LIST\LITS-Challenge-Test-Data\\test-volume-" + str(51) + ".nii")

    srcimg = sitk.GetArrayFromImage(src)

    for i in range(np.shape(srcimg)[0]):

        image = srcimg[i]

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

        cv2.imwrite("D:\Data\LIST\LITS-Challenge-Test-Data\\" + str(51) + "\\" + str(i) + ".bmp", image)

def getmaxsizeimage():

    srcpath = "D:\Data\LIST\LITS-Challenge-Test-Data\\test-volume-" + str(38) + ".nii"

    maskpath = "D:\Data\LIST\\test\PredictMask\\38"

    src = load_itk(srcpath)

    srcimg = sitk.GetArrayFromImage(src)

    maskimage = np.empty(shape=np.shape(srcimg), dtype=np.uint8)

    index = 0

    for _ in os.listdir(maskpath):

        masktmp = cv2.imread(maskpath + "/" + str(index) + ".bmp", cv2.IMREAD_GRAYSCALE)

        maskimage[index, :, :] = masktmp

        index += 1

    sitk_maskimg = sitk.GetImageFromArray(maskimage)

    origin = np.array(src.GetOrigin())

    # Read the spacing along each dimension

    spacing = np.array(src.GetSpacing())

    sitk_maskimg.SetSpacing(spacing)

    sitk_maskimg.SetOrigin(origin)

    maskimage = getLargestConnectedCompont(sitk_maskimg)

    for i in range(np.shape(maskimage)[0]):

        image = maskimage[i]

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

        cv2.imwrite("D:\Data\LIST\\test\PredictMask\\38_1\\" + str(i) + ".bmp", image)

def save_npy2csv(path, name, labelnum=1):

    """

    this is for classify

    :param path:

    :param name:file name

    :param labelnum:classify label

    :param label:

    :return:

    """

    out = open(name, 'w')

    file_list = glob(path + "*.npy")

    out.writelines("index" + "," + "filename" + "\n")

    for index in range(len(file_list)):

        out.writelines(str(labelnum) + "," + file_list[index] + "\n")

# gettestiamge()

# getmaxsizeimage()

# save_npy2csv("G:\Data\LIDC\LUNA16\classsification\\1_aug\\", "nodel_positive.csv", 1)