import os
import SimpleITK as sitk
import difflib

# 定義資料夾路徑
T1C_folder_path = '/home/onlylian/sorted_OAR_files/T1C'
CTCTC_folder_path = '/home/onlylian/sorted_OAR_files/CT+CTC'
T1C_output_folder_path = '/home/onlylian/T1C_label_folder'
CTCTC_output_folder_path = '/home/onlylian/CT+CTC_label_folder'

# 定義組織標籤對應關係
labels = {
    "background": "0",
    "Brainstem": "1",
    "Right_Eye": "2",
    "Left_Eye": "3",
    "Optic_Chiasm": "4",
    "Right_Optic_Nerve": "5",
    "Left_Optic_Nerve": "6"
}

# 確保輸出資料夾存在
os.makedirs(T1C_output_folder_path, exist_ok=True)
os.makedirs(CTCTC_output_folder_path, exist_ok=True)

def resample_image(image, reference_image):
    resample = sitk.ResampleImageFilter()
    resample.SetReferenceImage(reference_image)
    resample.SetInterpolator(sitk.sitkNearestNeighbor)
    resample.SetDefaultPixelValue(0)
    return resample.Execute(image)

def process_images(input_folder, output_folder, image_type):
    # 獲取所有病人的資料夾路徑
    patient_folders = sorted([os.path.join(dp, d) for dp, dn, filenames in os.walk(input_folder) for d in dn])
    
    file_counter = 1
    
    for patient_folder in patient_folders:
        # 創建空的影像，用來合成所有器官
        combined_image = None
        # 獲取病人文件夾中的所有文件名，並轉換為小寫形式
        patient_files = {f.lower(): f for f in os.listdir(patient_folder)}
        filename = patient_folder.split('/')[-2]
        
        # 遍歷每個病人的所有影像文件
        for organ_name, label in labels.items():
            organ_filename = f'struct_{organ_name}.nii.gz'.lower()
            # 使用 difflib.get_close_matches 查找相似文件名
            possible_matches = difflib.get_close_matches(organ_filename, patient_files.keys(), n=1, cutoff=0.8)
            
            if possible_matches:
                matched_filename = patient_files[possible_matches[0]]
                organ_path = os.path.join(patient_folder, matched_filename)
                organ_image = sitk.ReadImage(organ_path)
                
                if combined_image is None:
                    # 初始化合成影像
                    combined_image = sitk.Image(organ_image.GetSize(), sitk.sitkUInt8)
                    combined_image.CopyInformation(organ_image)
                
                # 重新取樣器官影像以匹配合成影像的大小
                resampled_organ_image = resample_image(organ_image, combined_image)
                
                # 將器官影像添加到合成影像中
                organ_array = sitk.GetArrayFromImage(resampled_organ_image)
                combined_array = sitk.GetArrayFromImage(combined_image)
                combined_array[organ_array > 0] = int(label)
                
                combined_image = sitk.GetImageFromArray(combined_array)
                combined_image.CopyInformation(resampled_organ_image)
        
        # 保存合成影像
        if combined_image:
            output_filename = f'{image_type}_label_{filename}.nii.gz'
            output_path = os.path.join(output_folder, output_filename)
            sitk.WriteImage(combined_image, output_path)
            print(f'Saved combined image to {output_path}')
            file_counter += 1

# 處理 T1C 影像
print("Processing T1C images...")
process_images(T1C_folder_path, T1C_output_folder_path, 'T1C')

# 處理 CT+CTC 影像
print("Processing CT+CTC images...")
process_images(CTCTC_folder_path, CTCTC_output_folder_path, 'CT+CTC')

print("All processing completed.")