import json
import os
import shutil
import glob
import re
import pandas as pd

# 創建 dataset.json 內容
dataset_json = {
    "channel_names": {
        "0": "CT+CTC",
        "1": "T1C"
    },
    "labels": {
        "background": "0",
        "Brainstem": "1",
        "Right_Eye": "2",
        "Left_Eye": "3",
        "Optic_Chiasm": "4",
        "Right_Optic_Nerve": "5",
        "Left_Optic_Nerve": "6",
        "TV": "7"
    },
    "numTraining": 0,
    "file_ending": ".nii.gz"
}

# 定義目錄
source_image1_dir = "/home/onlylian/resampled_CT+CTC_TV_image"
source_image2_dir = "/home/onlylian/resampled_T1C_TV_image"
source_label_dir = "/home/onlylian/resampled_tv_label"
target_base_dir = "/home/onlylian/nnUNet/nnUNet_raw/Dataset022_OAR_TV"
imagesTr_dir = os.path.join(target_base_dir, "imagesTr")
labelsTr_dir = os.path.join(target_base_dir, "labelsTr")

# 創建目錄
os.makedirs(imagesTr_dir, exist_ok=True)
os.makedirs(labelsTr_dir, exist_ok=True)

def extract_patient_id(filename):
    match = re.search(r'(CT\+CTC|T1C|label)_([^.]+)\.nii\.gz', os.path.basename(filename))
    return match.group(2) if match else None

# 獲取並排序原始圖像文件
source_image1_files = sorted(glob.glob(os.path.join(source_image1_dir, "*.nii.gz")))
source_image2_files = sorted(glob.glob(os.path.join(source_image2_dir, "*.nii.gz")))

processed_files = 0
patient_mapping = []

# 創建病歷號到文件名的映射
label_files = glob.glob(os.path.join(source_label_dir, "*.nii.gz"))
label_map = {extract_patient_id(f): f for f in label_files}

print(f"找到 {len(label_map)} 個標籤文件")
print(f"找到 {len(source_image1_files)} 個 CT+CTC 圖像文件")
print(f"找到 {len(source_image2_files)} 個 T1C 圖像文件")

# 處理所有資料
for index, (source_image1_file, source_image2_file) in enumerate(zip(source_image1_files, source_image2_files), start=1):
    if index % 10 == 0:  # 每處理10個文件顯示一次進度
        print(f"\n正在處理檔案 {index}/{len(source_image1_files)}:")
    
    patient_id1 = extract_patient_id(source_image1_file)
    patient_id2 = extract_patient_id(source_image2_file)
    
    if patient_id1 != patient_id2 or patient_id1 is None or patient_id2 is None:
        print(f"警告: 病歷號不匹配或無法提取，跳過這對文件: {source_image1_file}, {source_image2_file}")
        continue

    # 使用新的文件名格式
    new_filename = f"OAR_TV_{index:03d}"
    
    # 複製並重命名圖像文件
    target_image1_file = os.path.join(imagesTr_dir, f"{new_filename}_0000.nii.gz")
    target_image2_file = os.path.join(imagesTr_dir, f"{new_filename}_0001.nii.gz")
    
    shutil.copyfile(source_image1_file, target_image1_file)
    shutil.copyfile(source_image2_file, target_image2_file)

    # 處理標籤文件
    if patient_id1 in label_map:
        source_label_file = label_map[patient_id1]
        target_label_file = os.path.join(labelsTr_dir, f"{new_filename}.nii.gz")
        shutil.copyfile(source_label_file, target_label_file)
        processed_files += 1
        
        # 添加到映射列表
        patient_mapping.append({
            "編號": f"{index:03d}",
            "病歷號": patient_id1
        })

        # 輸出處理細節
        print(f"CT+CTC: {source_image1_file}")
        print(f"T1C: {source_image2_file}")
        print(f"複製圖像檔案: {source_image1_file} 到 {target_image1_file}")
        print(f"複製圖像檔案: {source_image2_file} 到 {target_image2_file}")
        print(f"複製標籤檔案: {source_label_file} 到 {target_label_file}")
        print()  # 空行分隔每筆記錄
    else:
        print(f"警告: 未找到 {patient_id1} 的標籤檔案。")

print("\n所有資料處理完畢。")

# 更新 dataset.json
dataset_json["numTraining"] = processed_files

# 寫入 dataset.json 檔案
dataset_json_path = os.path.join(target_base_dir, "dataset.json")
with open(dataset_json_path, 'w') as f:
    json.dump(dataset_json, f, indent=4)
print(f"dataset.json 已寫入到 {dataset_json_path}")

print(f"總共處理了 {processed_files} 個完整的資料集（圖像對加標籤）")

# 創建 Excel 文件
df = pd.DataFrame(patient_mapping)
excel_path = os.path.join(target_base_dir, "patient_mapping.xlsx")
df.to_excel(excel_path, index=False)
print(f"病人映射表已保存到 {excel_path}")

# 輸出處理摘要
print("\n處理摘要:")
print(f"總共找到的 CT+CTC 圖像: {len(source_image1_files)}")
print(f"總共找到的 T1C 圖像: {len(source_image2_files)}")
print(f"總共找到的標籤文件: {len(label_map)}")
print(f"成功處理的完整資料集: {processed_files}")
print(f"跳過的文件對數量: {len(source_image1_files) - processed_files}")