import os
import sys
import glob
import SimpleITK as sitk



def ReadImageITK(path):
    series_IDs = sitk.ImageSeriesReader.GetGDCMSeriesIDs(path)
    nb_series = len(series_IDs)
    print(nb_series)
    
    series_file_names = sitk.ImageSeriesReader.GetGDCMSeriesFileNames(path, series_IDs[0])
    series_reader = sitk.ImageSeriesReader()
    
    series_reader.SetFileNames(series_file_names)
    
    image3D = series_reader.Execute()
    
    return image3D,sitk.GetArrayFromImage(image3D)


def loadImage(path):
    if os.path.isdir(path):
        img = ReadImageITK(path)
    else:
        img_itk = sitk.ReadImage(path)
        return img_itk,sitk.GetArrayFromImage(img_itk)
    


def WriteMask(mask,origin=None,spacing=None,path=''):
    if origin is None and spacing is None:
        itk_mask = mask
    else:
        itk_mask = sitk.GetImageFromArray(mask)
        itk_mask = sitk.Cast(itk_mask,sitk.sitkUInt8)

        itk_mask.SetSpacing(spacing)
        itk_mask.SetOrigin(origin)
    
    
    print ('Saving img to %s '%path)
    itkWriter = sitk.ImageFileWriter()
    itkWriter.SetUseCompression(True)
    itkWriter.SetFileName(path)
    itkWriter.Execute(itk_mask)