"""
convert dcm 2 nii
1. read dicom series
2. clarify and sort series
3. skip scout series
4. convert series to 3d
5. resample and save as nii

## usage ##
---------------------------------
-- parent(root)
  |
  -- Folder1(subdir)
  |    |
  |    -- DICOMS
       /  -- 4d-dicom file
  -- Folder2
      |
      -- DICOMS
---------------------------------
input: root
output: subdir, include several nii files determined by series

dependencies:
python 3.9.2 (used for log)
SimpleITK
numpy

logging:
support output to console and file

"""

import SimpleITK as sitk
import numpy as np
import logging
from pathlib import Path


def set_logger(mylogger, file):
    mylogger.setLevel(logging.DEBUG)
    # create file handler which logs even debug messages
    fh = logging.FileHandler(file)
    fh.setLevel(logging.DEBUG)
    # create console handler with a higher log level
    ch = logging.StreamHandler()
    ch.setLevel(logging.DEBUG)
    # create formatter and add it to the handlers
    formatter = logging.Formatter('%(asctime)s [%(name)s] [%(levelname)s] %(message)s')
    ch.setFormatter(formatter)
    fh.setFormatter(formatter)
    # add the handlers to logger
    mylogger.addHandler(ch)
    mylogger.addHandler(fh)


# resample image using new spacing
def resample_image(itk_image, out_spacing):
    original_spacing = itk_image.GetSpacing()
    original_size = itk_image.GetSize()

    # image size according to new spacing
    out_size = [
        int(np.round(original_size[0] * original_spacing[0] / out_spacing[0])),
        int(np.round(original_size[1] * original_spacing[1] / out_spacing[1])),
        int(np.round(original_size[2] * original_spacing[2] / out_spacing[2]))
    ]

    resample = sitk.ResampleImageFilter()
    resample.SetOutputSpacing(out_spacing)
    resample.SetSize(out_size)
    resample.SetOutputDirection(itk_image.GetDirection())
    resample.SetOutputOrigin(itk_image.GetOrigin())
    resample.SetTransform(sitk.Transform())
    resample.SetDefaultPixelValue(itk_image.GetPixelIDValue())

    resample.SetInterpolator(sitk.sitkBSpline)
    # print(resample)
    return resample.Execute(itk_image)


# read and save nii file
# def resample_nii(input_path, out_file):
#     original_img = sitk.ReadImage(input_path)
#     resample_nii_image(original_img, out_file)


def resample_nii_image(nii_img, out_file):
    pix_spacing = nii_img.GetSpacing()
    logger.info('Original image Spacing：')
    logger.info(pix_spacing)
    logger.info('Original image Size：')
    logger.info(nii_img.GetSize())

    new_pix_spacing = [pix_spacing[0], pix_spacing[0], pix_spacing[0]]
    resample_img = resample_image(nii_img, new_pix_spacing)
    logger.info('Resampled image Spacing：')
    logger.info(resample_img.GetSpacing())
    logger.info('Resampled image Size：')
    logger.info(resample_img.GetSize())

    sitk.WriteImage(resample_img, out_file)


def convert_dcm2nii_resample(directory_path):
    reader = sitk.ImageSeriesReader()
    reader.MetaDataDictionaryArrayUpdateOn()
    reader.LoadPrivateTagsOn()
    series_ids = sitk.ImageSeriesReader.GetGDCMSeriesIDs(directory_path)
    arr_index = 0
    for series_ID in series_ids:
        arr_index += 1
        logger.info("series id: " + series_ID)
        dicom_names = reader.GetGDCMSeriesFileNames(directory_path, series_ID)
        # if 3 == arr_index:
        #     dicom_names = dicom_names[1409:2152]
        # else:
        #     print("skip series number: ", str(arr_index))
        #     continue
        reader.SetFileNames(dicom_names)
        image3d = reader.Execute()

        series_number = reader.GetMetaData(0, '0020|0011')
        # series_description = reader.GetMetaData(0, '0008|103e')
        # if series_number == '3 ':
        #     dicom_names = dicom_names[1409:2152]
        #     reader.SetFileNames(dicom_names)
        #     image3d = reader.Execute()
        # else:
        #     print("skip series number: ", str(series_number))
        #     continue
        # determine if it is scout protocol
        slice_number = 0
        # slice_keys = reader.GetMetaDataKeys(slice_number)
        # reader.HasMetaDataKey(slice_number, '0008|0008')
        # image type
        image_type = reader.GetMetaData(slice_number, '0008|0008')

        scout_flag = "LOCALIZER"
        if scout_flag in image_type:
            logger.info("skip scout")
            continue

        # saved path
        save_path = Path(directory_path).parent.as_posix()
        # save_path = save_path + "//" + str(series_number)+ "_" + str(series_description) + ".nii.gz"
        save_path = save_path + "//" + str(series_number) + "_" + ".nii.gz"
        # save_path = re.sub(r"[\n\t\s]*", "", save_path)
        # save_path1 = save_path.strip()
        # resample
        resample_nii_image(image3d, save_path)
        logger.info("saved file: ")
        logger.info(save_path)


def get_paths(input_path, out_paths):
    path = Path(input_path)
    for subdir in path.iterdir():
        if subdir.is_dir():
            get_paths(subdir, out_paths)
        else:
            if "nii.gz" in subdir.name:
                continue
            else:
                out_paths.append(subdir.parent)
                break


def main():
    path = Path(r"C:\Data\mydata\SID221153")
    logger.info("--------------------task begin-----------------------")
    input_dir = []
    get_paths(path, input_dir)
    logger.info("folder size: " + str(len(input_dir)))
    logger.info(input_dir)

    for convert_dir in input_dir:
        logger.info("---begin dir---")
        logger.info(convert_dir)
        convert_dcm2nii_resample(Path(convert_dir).as_posix())
        logger.info("---finish dir---")
    logger.info("--------------------task finish-----------------------")


logger = logging.getLogger('dcm2nii')
set_logger(logger, "log/convert_test.log")

if __name__ == '__main__':
    main()