'''
Copyright (c) Microsoft Corporation. All rights reserved.
Licensed under the MIT License.
'''

'''
This code does the following:
(1) converts PET DICOM images in units Bq/ml to decay-corrected SUV and saved as 3D NIFTI files
(2) converts CT DICOM images to NIFTI
(3) converts DICOM RTSTRUCT images to NIFTI (using rt-utils)
'''
#%%
import SimpleITK as sitk
from pydicom import dcmread, FileDataset
from rt_utils import RTStructBuilder, RTStruct
import numpy as np
import dateutil
import pandas as pd
import os
import time

#%%
'''
Script to convert PET and CT dicom series to niftii files. Works under 
the assumption that the rescale slope and intercept in the PET dicom 
series map image intensities to Bq/mL. Saved PET files will have image
intensities of SUVbw, and saved CT files will have HU units.

'''
def bqml_to_suv(dcm_file: FileDataset) -> float:
    
    # Calculates the SUV conversion factor from Bq/mL to SUVbw using 
    # the dicom header information in one of the images from a dicom series
    
    nuclide_dose = dcm_file[0x054, 0x0016][0][0x0018, 0x1074].value  # Total injected dose (Bq)
    weight = dcm_file[0x0010, 0x1030].value  # Patient weight (Kg)
    half_life = float(dcm_file[0x054, 0x0016][0][0x0018, 0x1075].value)  # Radionuclide half life (s)

    parse = lambda x: dateutil.parser.parse(x)

    series_time = str(dcm_file[0x0008, 0x00031].value)  # Series start time (hh:mm:ss)
    series_date = str(dcm_file[0x0008, 0x00021].value)  # Series start date (yyy:mm:dd)
    series_datetime_str = series_date + ' ' + series_time
    series_dt = parse(series_datetime_str)

    nuclide_time = str(dcm_file[0x054, 0x0016][0][0x0018, 0x1072].value)  # Radionuclide time of injection (hh:mm:ss)
    nuclide_datetime_str = series_date + ' ' + nuclide_time
    nuclide_dt = parse(nuclide_datetime_str)

    delta_time = (series_dt - nuclide_dt).total_seconds()
    decay_correction = 2 ** (-1 * delta_time/half_life)
    suv_factor = (weight * 1000) / (decay_correction * nuclide_dose)

    return(suv_factor)

def get_filtered_roi_list(rois):
    filtered_rois = []
    for roi in rois:
        if roi.endswith('PETEdge'):
            filtered_rois.append(roi)
        else:
            pass
    return filtered_rois


def load_merge_masks(rtstruct: RTStruct) -> np.ndarray:
    '''
    Load and merge masks from a dicom RTStruct. All of the
    masks in the RTStruct will be merged. Add an extra line
    of code if you want to filter for/out certain masks.
    '''
    rois = rtstruct.get_roi_names()
    rois = get_filtered_roi_list(rois)
    masks = []
    for roi in rois:
        print(roi)
        mask_3d = rtstruct.get_roi_mask_by_name(roi).astype(int)
        masks.append(mask_3d)

    final_mask = sum(masks)  # sums element-wise
    final_mask = np.where(final_mask>=1, 1, 0)
    # Reorient the mask to line up with the reference image
    final_mask = np.moveaxis(final_mask, [0, 1, 2], [1, 2, 0])

    return final_mask

############################################################################################
########  Update the three variables below with the locations of your choice  ##############
############################################################################################
save_dir_ct = '' # path to directory where your new CT files in NIFTI format will be written
save_dir_pt = '' # path to directory where your new PET files in NIFTI format will be written
save_dir_gt = '' # path to directory where your new GT files in NIFTI format will be written
############################################################################################
############################################################################################
############################################################################################

cases = pd.read_csv('dicom_ctpt_to_nifti_conversion_file.csv')
cases = list(cases.itertuples(index=False, name=None)) 
structs = pd.read_csv('dicom_rtstruct_to_nifti_conversion_file.csv')
structs = list(structs.itertuples(index=False, name=None))
# Execution
start = time.time()

for case in cases:
    patient_id, ct_folder, pet_folder, convert = case
    if convert=='N':
        continue
    print(f'Converting patient Id: {patient_id}')

    # Convert CT series
    ct_reader = sitk.ImageSeriesReader()
    ct_series_names = ct_reader.GetGDCMSeriesFileNames(ct_folder)
    ct_reader.SetFileNames(ct_series_names)
    ct = ct_reader.Execute()
    sitk.WriteImage(ct, os.path.join(save_dir_ct, f"{patient_id}_0000.nii.gz"), imageIO='NiftiImageIO')
    print('Saved nifti CT')

    # Convert PET series
    pet_reader = sitk.ImageSeriesReader()
    pet_series_names = pet_reader.GetGDCMSeriesFileNames(pet_folder)
    pet_reader.SetFileNames(pet_series_names)
    pet = pet_reader.Execute()

    pet_img = dcmread(pet_series_names[0])  # read one of the images for header info
    suv_factor = bqml_to_suv(pet_img)
    pet = sitk.Multiply(pet, suv_factor)
    sitk.WriteImage(pet, os.path.join(save_dir_pt, f"{patient_id}_0001.nii.gz"), imageIO='NiftiImageIO')
    print('Saved nifti PET')

# Execution
for struct in structs:
    patient_id, struct_folder, ref_folder, convert = struct
    if convert=='N':
        continue

    # print('Converting RTStruct for patient {}'.format(num))
    # Get all the paths in order
    struct_file = os.listdir(struct_folder)[0]
    struct_path = os.path.join(struct_folder, struct_file)

    # Create the mask
    rtstruct = RTStructBuilder.create_from(dicom_series_path= ref_folder, rt_struct_path=struct_path)
    final_mask = load_merge_masks(rtstruct)

    # Load original DICOM image for reference
    reader = sitk.ImageSeriesReader()
    dicom_names = reader.GetGDCMSeriesFileNames(ref_folder)
    reader.SetFileNames(dicom_names)
    ref_img = reader.Execute()

    # Properly reference and convert the mask to an image object
    mask_img = sitk.GetImageFromArray(final_mask)
    mask_img.CopyInformation(ref_img)
    sitk.WriteImage(mask_img, os.path.join(save_dir_gt, f"{patient_id}.nii.gz"), imageIO="NiftiImageIO")

    print('Patient {} mask saved'.format(patient_id))