import os
import glob
import SimpleITK as sitk
from difflib import get_close_matches

def process_images(input_base_path, output_base_path, output_prefix):
    # 定義組織標籤對應關係
    labels = {
        "background": "0",
        "Struct_Brain_Stem": "1",
        "Struct_Right_Eye": "2",
        "Struct_Left_Eye": "3",
        "Struct_Optic_Chiasm": "4",
        "Struct_Right_Optic_Nerve": "5",
        "Struct_Left_Optic_Nerve": "6",
        "TV": "7"  # Tumor
    }

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

    # 定義重新取樣函數
    def resample_image(image, reference_image):
        resample = sitk.ResampleImageFilter()
        resample.SetReferenceImage(reference_image)
        resample.SetInterpolator(sitk.sitkNearestNeighbor)  # 使用最近鄰插值避免標籤變化
        resample.SetDefaultPixelValue(0)
        resampled_image = resample.Execute(image)
        return resampled_image

    # 遍歷所有病人的資料夾
    for patient_id in os.listdir(input_base_path):
        patient_folder_path = os.path.join(input_base_path, patient_id)
        if not os.path.isdir(patient_folder_path):
            continue  # 跳過非資料夾的項目

        # 確定病人資料夾下的日期資料夾
        date_folder = os.listdir(patient_folder_path)[0]
        ORGAN_folder_path = os.path.join(patient_folder_path, date_folder)

        # 創建一個空的影像來合併所有標籤
        combined_image = None

        # 列出病人資料夾內的所有檔案
        all_files = os.listdir(ORGAN_folder_path)

        # 遍歷每個標籤並合併
        for organ_name, label in labels.items():
            if organ_name == "TV":
                # 動態搜尋 TV 檔案
                tv_files = glob.glob(os.path.join(ORGAN_folder_path, "TV-*.nii.gz"))
                if len(tv_files) == 0:
                    print(f"TV file not found for patient {patient_id}, skipping TV...")
                    continue
                organ_path = tv_files[0]  # 取第一個找到的檔案
            else:
                # 使用模糊比對尋找最接近的檔案名
                matches = get_close_matches(organ_name, all_files, n=1, cutoff=0.6)
                if len(matches) == 0:
                    print(f"No close match found for {organ_name} for patient {patient_id}, skipping...")
                    continue
                organ_filename = matches[0]
                organ_path = os.path.join(ORGAN_folder_path, organ_filename)

            if not os.path.exists(organ_path):
                print(f"File {organ_path} not found for patient {patient_id}, skipping...")
                continue

            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)
            else:
                # 如果器官影像的尺寸和合併影像不同，進行重新取樣
                if organ_image.GetSize() != combined_image.GetSize():
                    print(f"Resampling organ {organ_name} for patient {patient_id} due to size mismatch...")
                    organ_image = resample_image(organ_image, combined_image)

            # 將器官影像添加到合併影像中
            organ_array = sitk.GetArrayFromImage(organ_image)
            combined_array = sitk.GetArrayFromImage(combined_image)

            # 更新合併影像的陣列，根據標籤值來設定非背景區域
            combined_array[organ_array > 0] = int(label)

            # 將更新的陣列轉回 SimpleITK 影像
            combined_image = sitk.GetImageFromArray(combined_array)
            combined_image.CopyInformation(organ_image)

        # 保存合併後的影像
        output_filename = f'{output_prefix}_label_{patient_id}.nii.gz'
        output_path = os.path.join(output_base_path, output_filename)
        sitk.WriteImage(combined_image, output_path)
        print(f'Saved combined image for patient {patient_id} to {output_path}')

if __name__ == "__main__":
    # 處理 T1C 影像
    T1C_input_path = '/home/onlylian/sorted_OAR_Target_files/T1C'
    T1C_output_path = '/home/onlylian/T1C_OAR_Target_label_folder'
    process_images(T1C_input_path, T1C_output_path, "T1C")

    # 處理 CT+CTC 影像
    CT_CTC_input_path = '/home/onlylian/sorted_OAR_Target_files/T1C'  # 注意：這裡使用相同的輸入路徑，如果需要不同的路徑請修改
    CT_CTC_output_path = '/home/onlylian/CT+CTC_OAR_Target_label_folder'
    process_images(CT_CTC_input_path, CT_CTC_output_path, "CT+CTC")