import SimpleITK as sitk
import os
import numpy as np

def standardize_origin(image, target_origin):
    """將影像的原點標準化為目標原點"""
    return sitk.Resample(
        image,
        image.GetSize(),
        sitk.Transform(),
        sitk.sitkLinear,
        target_origin,
        image.GetSpacing(),
        image.GetDirection(),
        0,
        image.GetPixelID()
    )

def resample_image(image, reference_image, is_label=False):
    resampler = sitk.ResampleImageFilter()
    resampler.SetReferenceImage(reference_image)
    resampler.SetTransform(sitk.Transform())
    resampler.SetOutputSpacing(reference_image.GetSpacing())
    resampler.SetSize(reference_image.GetSize())
    resampler.SetOutputDirection(reference_image.GetDirection())
    resampler.SetOutputOrigin(reference_image.GetOrigin())
    
    if is_label:
        resampler.SetInterpolator(sitk.sitkNearestNeighbor)
        resampler.SetDefaultPixelValue(0)
    else:
        resampler.SetInterpolator(sitk.sitkLinear)
        resampler.SetDefaultPixelValue(image.GetPixelIDValue())
    
    return resampler.Execute(image)

def process_patient(t1c_image_path, ctctc_image_path, label_path, output_t1c_dir, output_ctctc_dir, output_label_dir, patient_id):
    print(f"\nProcessing patient: {patient_id}")
    
    # 讀取影像
    t1c_image = sitk.ReadImage(t1c_image_path)
    ctctc_image = sitk.ReadImage(ctctc_image_path)
    label = sitk.ReadImage(label_path)

    print(f"  Original T1C image size: {t1c_image.GetSize()}, origin: {t1c_image.GetOrigin()}")
    print(f"  Original CT+CTC image size: {ctctc_image.GetSize()}, origin: {ctctc_image.GetOrigin()}")
    print(f"  Original label size: {label.GetSize()}, origin: {label.GetOrigin()}")

    # 使用 T1C 影像作為參考
    reference_image = t1c_image

    # 重採樣 CT+CTC 影像以匹配 T1C 影像
    resampled_ctctc = resample_image(ctctc_image, reference_image)
    
    # 重採樣 label 以匹配 T1C 影像
    resampled_label = resample_image(label, reference_image, is_label=True)

    print(f"  Resampled CT+CTC image size: {resampled_ctctc.GetSize()}, origin: {resampled_ctctc.GetOrigin()}")
    print(f"  Resampled label size: {resampled_label.GetSize()}, origin: {resampled_label.GetOrigin()}")

    # 保存重採樣後的影像和標籤
    t1c_output_path = os.path.join(output_t1c_dir, f"T1C_{patient_id}.nii.gz")
    ctctc_output_path = os.path.join(output_ctctc_dir, f"CT+CTC_{patient_id}.nii.gz")
    label_output_path = os.path.join(output_label_dir, f"label_{patient_id}.nii.gz")

    sitk.WriteImage(t1c_image, t1c_output_path)
    sitk.WriteImage(resampled_ctctc, ctctc_output_path)
    sitk.WriteImage(resampled_label, label_output_path)

    print(f"  T1C image saved to: {t1c_output_path}")
    print(f"  Resampled CT+CTC image saved to: {ctctc_output_path}")
    print(f"  Resampled label saved to: {label_output_path}")

def process_directory(t1c_dir, ctctc_dir, label_dir, output_t1c_dir, output_ctctc_dir, output_label_dir):
    print("Starting directory processing")
    os.makedirs(output_t1c_dir, exist_ok=True)
    os.makedirs(output_ctctc_dir, exist_ok=True)
    os.makedirs(output_label_dir, exist_ok=True)

    # 獲取所有病人 ID
    patient_ids = set()
    for filename in os.listdir(t1c_dir):
        if filename.startswith("T1C_") and filename.endswith(".nii.gz"):
            patient_id = filename[4:-7]  # 去掉 "T1C_" 和 ".nii.gz"
            patient_ids.add(patient_id)

    for patient_id in patient_ids:
        t1c_path = os.path.join(t1c_dir, f"T1C_{patient_id}.nii.gz")
        ctctc_path = os.path.join(ctctc_dir, f"CT+CTC_{patient_id}.nii.gz")
        label_path = os.path.join(label_dir, f"T1C_label_{patient_id}.nii.gz")  # 假設 label 文件名與 T1C 相同
        
        if os.path.exists(t1c_path) and os.path.exists(ctctc_path) and os.path.exists(label_path):
            process_patient(t1c_path, ctctc_path, label_path, output_t1c_dir, output_ctctc_dir, output_label_dir, patient_id)
        else:
            print(f"  WARNING: Missing files for patient {patient_id}")

# 主程序
t1c_dir = '/home/onlylian/T1C_image_folder'
ctctc_dir = '/home/onlylian/CT+CTC_image_folder'
label_dir = '/home/onlylian/T1C_label_folder'  # 使用 T1C 的 label
output_t1c_dir = '/home/onlylian/resampled_T1C_image'
output_ctctc_dir = '/home/onlylian/resampled_CT+CTC_image'
output_label_dir = '/home/onlylian/resampled_label'

print("Starting image preprocessing")
process_directory(t1c_dir, ctctc_dir, label_dir, output_t1c_dir, output_ctctc_dir, output_label_dir)

print("All processing completed.")