import os

import re

import argparse

import numpy as np

import random

import monai

import time

# from networks import build_net

import logging

import os

import sys

import tempfile

from glob import glob

from ignite.metrics import Accuracy

import nibabel as nib

import torch

import argparse

from monai.data import CacheDataset, DataLoader, Dataset

import SimpleITK as sitk

from monai.inferers import sliding_window_inference

from monai.metrics import DiceMetric

from monai.data import NiftiSaver, create_test_image_3d, list_data_collate

from collections import OrderedDict

from monai.handlers import (MeanDice, StatsHandler, ValidationHandler, CheckpointSaver, LrScheduleHandler, CheckpointLoader,

                         SegmentationSaver, TensorBoardImageHandler, TensorBoardStatsHandler)

from monai.inferers import SimpleInferer, SlidingWindowInferer

from monai.utils import set_determinism

import re

from monai.data import create_test_image_3d, list_data_collate

from monai.inferers import sliding_window_inference

from monai.transforms import (Activationsd,MeanEnsembled, GaussianSmoothd, CropForegroundd, ThresholdIntensityd, Activations,AsDiscrete, LoadImaged, AsChannelFirstd, VoteEnsembled, AsDiscreted, Compose, AddChanneld, Transpose, ConcatItemsd,

                              ScaleIntensityd, Resized,ToTensord, RandSpatialCropd, Rand3DElasticd, RandAffined, RandGaussianSmoothd, SpatialPadd,

    Spacingd, Orientationd, RandShiftIntensityd, BorderPadd, RandGaussianNoised, RandAdjustContrastd,NormalizeIntensityd,RandFlipd, KeepLargestConnectedComponent)

from monai.engines import (

    EnsembleEvaluator,

    SupervisedEvaluator,

    SupervisedTrainer

)

from skimage.measure import label

def getLargestCC(segmentation):

    labels = label(segmentation)

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

    list_seg=list(zip(unique, counts))[1:] # the 0 label is by default background so take the rest

    largest=max(list_seg, key=lambda x:x[1])[0]

    labels_max=(labels == largest).astype(int)

    return labels_max

def Padding(image, reference):

    size_new = reference.GetSize()

    output_size = tuple(size_new)

    resampler = sitk.ResampleImageFilter()

    resampler.SetOutputSpacing(reference.GetSpacing())

    resampler.SetSize(output_size)

    # resample on label

    resampler.SetInterpolator(sitk.sitkNearestNeighbor)

    resampler.SetOutputOrigin(reference.GetOrigin())

    resampler.SetOutputDirection(reference.GetDirection())

    image = resampler.Execute(image)

    return image

def resize(img, new_size, interpolator):

    # img = sitk.ReadImage(img)

    dimension = img.GetDimension()

    # Physical image size corresponds to the largest physical size in the training set, or any other arbitrary size.

    reference_physical_size = np.zeros(dimension)

    reference_physical_size[:] = [(sz - 1) * spc if sz * spc > mx else mx for sz, spc, mx in

                                  zip(img.GetSize(), img.GetSpacing(), reference_physical_size)]

    # Create the reference image with a zero origin, identity direction cosine matrix and dimension

    reference_origin = np.zeros(dimension)

    reference_direction = np.identity(dimension).flatten()

    reference_size = new_size

    reference_spacing = [phys_sz / (sz - 1) for sz, phys_sz in zip(reference_size, reference_physical_size)]

    reference_image = sitk.Image(reference_size, img.GetPixelIDValue())

    reference_image.SetOrigin(reference_origin)

    reference_image.SetSpacing(reference_spacing)

    reference_image.SetDirection(reference_direction)

    # Always use the TransformContinuousIndexToPhysicalPoint to compute an indexed point's physical coordinates as

    # this takes into account size, spacing and direction cosines. For the vast majority of images the direction

    # cosines are the identity matrix, but when this isn't the case simply multiplying the central index by the

    # spacing will not yield the correct coordinates resulting in a long debugging session.

    reference_center = np.array(

        reference_image.TransformContinuousIndexToPhysicalPoint(np.array(reference_image.GetSize()) / 2.0))

    # Transform which maps from the reference_image to the current img with the translation mapping the image

    # origins to each other.

    transform = sitk.AffineTransform(dimension)

    transform.SetMatrix(img.GetDirection())

    transform.SetTranslation(np.array(img.GetOrigin()) - reference_origin)

    # Modify the transformation to align the centers of the original and reference image instead of their origins.

    centering_transform = sitk.TranslationTransform(dimension)

    img_center = np.array(img.TransformContinuousIndexToPhysicalPoint(np.array(img.GetSize()) / 2.0))

    centering_transform.SetOffset(np.array(transform.GetInverse().TransformPoint(img_center) - reference_center))

    # centered_transform = sitk.Transform(transform)

    # centered_transform.AddTransform(centering_transform)

    centered_transform = sitk.CompositeTransform([transform, centering_transform])

    # Using the linear interpolator as these are intensity images, if there is a need to resample a ground truth

    # segmentation then the segmentation image should be resampled using the NearestNeighbor interpolator so that

    # no new labels are introduced.

    return sitk.Resample(img, reference_image, centered_transform, interpolator, 0.0)

def resample_sitk_image(sitk_image, spacing=None, interpolator=None, fill_value=0):

    # https://github.com/SimpleITK/SlicerSimpleFilters/blob/master/SimpleFilters/SimpleFilters.py

    _SITK_INTERPOLATOR_DICT = {

        'nearest': sitk.sitkNearestNeighbor,

        'linear': sitk.sitkLinear,

        'gaussian': sitk.sitkGaussian,

        'label_gaussian': sitk.sitkLabelGaussian,

        'bspline': sitk.sitkBSpline,

        'hamming_sinc': sitk.sitkHammingWindowedSinc,

        'cosine_windowed_sinc': sitk.sitkCosineWindowedSinc,

        'welch_windowed_sinc': sitk.sitkWelchWindowedSinc,

        'lanczos_windowed_sinc': sitk.sitkLanczosWindowedSinc

    }

    if isinstance(sitk_image, str):

        sitk_image = sitk.ReadImage(sitk_image)

    num_dim = sitk_image.GetDimension()

    if not interpolator:

        interpolator = 'linear'

        pixelid = sitk_image.GetPixelIDValue()

        if pixelid not in [1, 2, 4]:

            raise NotImplementedError(

                'Set `interpolator` manually, '

                'can only infer for 8-bit unsigned or 16, 32-bit signed integers')

        if pixelid == 1:  # 8-bit unsigned int

            interpolator = 'nearest'

    orig_pixelid = sitk_image.GetPixelIDValue()

    orig_origin = sitk_image.GetOrigin()

    orig_direction = sitk_image.GetDirection()

    orig_spacing = np.array(sitk_image.GetSpacing())

    orig_size = np.array(sitk_image.GetSize(), dtype=np.int)

    if not spacing:

        min_spacing = orig_spacing.min()

        new_spacing = [min_spacing] * num_dim

    else:

        new_spacing = [float(s) for s in spacing]

    assert interpolator in _SITK_INTERPOLATOR_DICT.keys(), \

        '`interpolator` should be one of {}'.format(_SITK_INTERPOLATOR_DICT.keys())

    sitk_interpolator = _SITK_INTERPOLATOR_DICT[interpolator]

    new_size = orig_size * (orig_spacing / new_spacing)

    new_size = np.ceil(new_size).astype(np.int)  # Image dimensions are in integers

    new_size = [int(s) for s in new_size]  # SimpleITK expects lists, not ndarrays

    resample_filter = sitk.ResampleImageFilter()

    resample_filter.SetOutputSpacing(new_spacing)

    resample_filter.SetSize(new_size)

    resample_filter.SetOutputDirection(orig_direction)

    resample_filter.SetOutputOrigin(orig_origin)

    resample_filter.SetTransform(sitk.Transform())

    resample_filter.SetDefaultPixelValue(orig_pixelid)

    resample_filter.SetInterpolator(sitk_interpolator)

    resample_filter.SetDefaultPixelValue(fill_value)

    resampled_sitk_image = resample_filter.Execute(sitk_image)

    return resampled_sitk_image

def numericalSort(value):

    numbers = re.compile(r'(\d+)')

    parts = numbers.split(value)

    parts[1::2] = map(int, parts[1::2])

    return parts

def lstFiles(Path):

    images_list = []  # create an empty list, the raw image data files is stored here

    for dirName, subdirList, fileList in os.walk(Path):

        for filename in fileList:

            if ".nii.gz" in filename.lower():

                images_list.append(os.path.join(dirName, filename))

            elif ".nii" in filename.lower():

                images_list.append(os.path.join(dirName, filename))

            elif ".mhd" in filename.lower():

                images_list.append(os.path.join(dirName, filename))

    images_list = sorted(images_list, key=numericalSort)

    return images_list

def new_state_dict(file_name):

    state_dict = torch.load(file_name)

    new_state_dict = OrderedDict()

    for k, v in state_dict.items():

        if k[:6] == 'module':

            name = k[7:]

            new_state_dict[name] = v

        else:

            new_state_dict[k] = v

    return new_state_dict

def new_state_dict_cpu(file_name):

    state_dict = torch.load(file_name, map_location='cpu')

    new_state_dict_cpu = OrderedDict()

    for k, v in state_dict.items():

        if k[:6] == 'module':

            name = k[7:]

            new_state_dict_cpu[name] = v

        else:

            new_state_dict_cpu[k] = v

    return new_state_dict_cpu

def from_numpy_to_itk(image_np, image_itk):

    # read image file

    reader = sitk.ImageFileReader()

    reader.SetFileName(image_itk)

    image_itk = reader.Execute()

    image_np = np.transpose(image_np, (2, 1, 0))

    image = sitk.GetImageFromArray(image_np)

    image.SetDirection(image_itk.GetDirection())

    image.SetSpacing(image_itk.GetSpacing())

    image.SetOrigin(image_itk.GetOrigin())

    return image

# function to keep track of the cropped area and coordinates

def statistics_crop(image, resolution):

    files = [{"image": image}]

    reader = sitk.ImageFileReader()

    reader.SetFileName(image)

    image_itk = reader.Execute()

    original_resolution = image_itk.GetSpacing()

    # original size

    transforms = Compose([

        LoadImaged(keys=['image']),

        AddChanneld(keys=['image']),

        ToTensord(keys=['image'])])

    data = monai.data.Dataset(data=files, transform=transforms)

    loader = DataLoader(data, batch_size=1, num_workers=0, pin_memory=torch.cuda.is_available())

    loader = monai.utils.misc.first(loader)

    im, = (loader['image'][0])

    vol = im.numpy()

    original_shape = vol.shape

    # cropped foreground size

    transforms = Compose([

        LoadImaged(keys=['image']),

        AddChanneld(keys=['image']),

        CropForegroundd(keys=['image'], source_key='image', start_coord_key='foreground_start_coord',

                        end_coord_key='foreground_end_coord', ),  # crop CropForeground

        ToTensord(keys=['image', 'foreground_start_coord', 'foreground_end_coord'])])

    data = monai.data.Dataset(data=files, transform=transforms)

    loader = DataLoader(data, batch_size=1, num_workers=0, pin_memory=torch.cuda.is_available())

    loader = monai.utils.misc.first(loader)

    im, coord1, coord2 = (loader['image'][0], loader['foreground_start_coord'][0], loader['foreground_end_coord'][0])

    vol = im[0].numpy()

    coord1 = coord1.numpy()

    coord2 = coord2.numpy()

    crop_shape = vol.shape

    if resolution is not None:

        transforms = Compose([

            LoadImaged(keys=['image']),

            AddChanneld(keys=['image']),

            CropForegroundd(keys=['image'], source_key='image'),  # crop CropForeground

            Spacingd(keys=['image'], pixdim=resolution, mode=('bilinear')),  # resolution

            ToTensord(keys=['image'])])

        data = monai.data.Dataset(data=files, transform=transforms)

        loader = DataLoader(data, batch_size=1, num_workers=0, pin_memory=torch.cuda.is_available())

        loader = monai.utils.misc.first(loader)

        im, = (loader['image'][0])

        vol = im.numpy()

        resampled_size = vol.shape

    else:

        resampled_size = original_shape

    return original_shape, crop_shape, coord1, coord2, resampled_size, original_resolution

def build_net_CT(patch_size,resolution):

    from monai.networks.layers import Norm

    sizes, spacings = patch_size, resolution

    strides, kernels = [], []

    while True:

        spacing_ratio = [sp / min(spacings) for sp in spacings]

        stride = [2 if ratio <= 2 and size >= 8 else 1 for (ratio, size) in zip(spacing_ratio, sizes)]

        kernel = [3 if ratio <= 2 else 1 for ratio in spacing_ratio]

        if all(s == 1 for s in stride):

            break

        sizes = [i / j for i, j in zip(sizes, stride)]

        spacings = [i * j for i, j in zip(spacings, stride)]

        kernels.append(kernel)

        strides.append(stride)

    strides.insert(0, len(spacings) * [1])

    kernels.append(len(spacings) * [3])

    # # create Unet

    nn_Unet = monai.networks.nets.DynUNet(

        spatial_dims=3,

        in_channels=1,

        out_channels=1,

        kernel_size=kernels,

        strides=strides,

        upsample_kernel_size=strides[1:],

        res_block=True,

    )

    return nn_Unet

def crop_window(prostate_contour):

    # Cut data, restricted to the prostate contours + a pitch per direction per dimension.

    """

    nrrd has the following format, assuming to watch the patient from the front:

    (x, y, z)

    x: left to right (ascending)

    y: front to back (ascending)

    z: bottom to top (ascending)

    """

    pitch = 5

    pattern = np.where(prostate_contour == 1)

    minx = np.min(pattern[0]) - pitch

    maxx = np.max(pattern[0]) + pitch

    miny = np.min(pattern[1]) - pitch

    maxy = np.max(pattern[1]) + pitch

    minz = np.min(pattern[2]) - pitch

    maxz = np.max(pattern[2]) + pitch

    if (maxx - minx) % 2 != 0:

        maxx += 1

    if (maxy - miny) % 2 != 0:

        maxy += 1

    if (maxz - minz) % 2 != 0:

        maxz += 1

    """

    Choose all tensors to have size of 64x64x64

    """

    limit = 32

    while maxx - minx < limit:

        maxx += 1

        minx -= 1

    while maxy - miny < limit:

        maxy += 1

        miny -= 1

    while maxz - minz < limit:

        maxz += 1

        minz -= 1

    return minx, maxx, miny, maxy, minz, maxz

def uniform_img_dimensions(image, label, nearest):

    image_array = sitk.GetArrayFromImage(image)

    image_array = np.transpose(image_array, axes=(2, 1, 0))  # reshape array from itk z,y,x  to  x,y,z

    image_shape = image_array.shape

    if nearest is True:

        label = resample_sitk_image(label, spacing=image.GetSpacing(), interpolator='nearest')

        res = resize(label,image_shape,sitk.sitkNearestNeighbor)

        res = (np.rint(sitk.GetArrayFromImage(res)))

        res = sitk.GetImageFromArray(res.astype('uint8'))

        # print(res.GetSize())

    else:

        label = resample_sitk_image(label, spacing=image.GetSpacing(), interpolator='linear')

        res = resize(label, image_shape, sitk.sitkLinear)

        res = (np.rint(sitk.GetArrayFromImage(res)))

        res = sitk.GetImageFromArray(res.astype('float'))

    res.SetDirection(image.GetDirection())

    res.SetOrigin(image.GetOrigin())

    res.SetSpacing(image.GetSpacing())

    return image, res

def uniform_img_dimensions_internal(image, label, nearest):

    name_label = label

    image = sitk.ReadImage(image)

    label = sitk.ReadImage(label)

    image_array = sitk.GetArrayFromImage(image)

    image_array = np.transpose(image_array, axes=(2, 1, 0))  # reshape array from itk z,y,x  to  x,y,z

    image_shape = image_array.shape

    if nearest is True:

        label = resample_sitk_image(label, spacing=image.GetSpacing(), interpolator='nearest')

        res = resize(label,image_shape,sitk.sitkNearestNeighbor)

        res = (np.rint(sitk.GetArrayFromImage(res)))

        res = sitk.GetImageFromArray(res.astype('uint8'))

        # print(res.GetSize())

    else:

        label = resample_sitk_image(label, spacing=image.GetSpacing(), interpolator='linear')

        res = resize(label, image_shape, sitk.sitkLinear)

        res = (np.rint(sitk.GetArrayFromImage(res)))

        res = sitk.GetImageFromArray(res.astype('float'))

    res.SetDirection(image.GetDirection())

    res.SetOrigin(image.GetOrigin())

    res.SetSpacing(image.GetSpacing())

    sitk.WriteImage(res, name_label)

def normalize_PET(image_itk, value):

    # read image file

    image_np = sitk.GetArrayFromImage(image_itk)

    image_np = image_np/value

    image = sitk.GetImageFromArray(image_np)

    image.SetDirection(image_itk.GetDirection())

    image.SetSpacing(image_itk.GetSpacing())

    image.SetOrigin(image_itk.GetOrigin())

    return image

def processing_itk(label_CT, image_PET, label_PET, gluteus, new_resolution, patch_size):

    gluteus = sitk.ReadImage(gluteus)

    label_CT = sitk.ReadImage(label_CT)

    image_PET = sitk.ReadImage(image_PET)

    if label_PET is not None:

        label_PET = sitk.ReadImage(label_PET)

    if new_resolution is not None:

        image_PET = resample_sitk_image(image_PET, spacing=new_resolution, interpolator='linear')

    label_CT = Padding(label_CT, image_PET)

    gluteus = Padding(gluteus, image_PET)

    image_PET, label_CT = uniform_img_dimensions(image_PET, label_CT, True)

    image_PET, gluteus = uniform_img_dimensions(image_PET, gluteus, True)

    # new part for Pet tumor_background normalization

    gluteos_ROI_array = sitk.GetArrayFromImage(gluteus)

    gluteos_ROI_index = np.where(gluteos_ROI_array == 1)

    PET_array = sitk.GetArrayFromImage(image_PET)

    avg = np.mean(PET_array[gluteos_ROI_index])

    image_PET = normalize_PET(image_PET, avg)

    # end normalization

    if label_PET is not None:

        label_PET = Padding(label_PET, image_PET)

        image_PET, label_PET = uniform_img_dimensions(image_PET, label_PET, True)

    label_CT_array = sitk.GetArrayFromImage(label_CT)

    minx, maxx, miny, maxy, minz, maxz = crop_window(label_CT_array)

    roiFilter = sitk.RegionOfInterestImageFilter()

    roiFilter.SetSize(patch_size)

    roiFilter.SetIndex([int(minz), int(miny), int(minx)])

    label_CT = roiFilter.Execute(label_CT)

    image_PET = roiFilter.Execute(image_PET)

    if label_PET is not None:

        label_PET = roiFilter.Execute(label_PET)

    else:

        label_PET = None

    sitk.WriteImage(label_CT, 'mask_crop.nii')

    sitk.WriteImage(image_PET, 'result.nii')

    if label_PET is not None:

        sitk.WriteImage(label_PET, 'label_crop.nii')

def gaussian2(image):

    resacleFilter = sitk.RescaleIntensityImageFilter()

    resacleFilter.SetOutputMaximum(255)

    resacleFilter.SetOutputMinimum(0)

    image = resacleFilter.Execute(image)  # set intensity 0-255

    gaussianFilter = sitk.SmoothingRecursiveGaussianImageFilter()

    gaussianFilter.SetSigma(3)

    image = gaussianFilter.Execute(image)

    resacleFilter = sitk.RescaleIntensityImageFilter()

    resacleFilter.SetOutputMaximum(1)

    resacleFilter.SetOutputMinimum(0)

    image = resacleFilter.Execute(image)  # set intensity 0-255

    thresholdFilter = sitk.BinaryThresholdImageFilter()

    thresholdFilter.SetLowerThreshold(0.5)

    thresholdFilter.SetUpperThreshold(2)

    thresholdFilter.SetInsideValue(1)

    thresholdFilter.SetOutsideValue(0)

    image = thresholdFilter.Execute(image)

    return image