import os
import shutil

# 定义所有需要的路径
final_image_folder = '/mnt/1218/onlylian/resampled_T1_image_test'
final_label_folder = '/mnt/1218/onlylian/resampled_label_test'
mapping_file = '/mnt/1218/onlylian/patient_mapping.txt'

# nnUNet路径
base_nnunet_path = '/home/onlylian/nnUNet/nnUNet_raw/Dataset888'
nnunet_image_folder = os.path.join(base_nnunet_path, 'imagesTs')
nnunet_label_folder = os.path.join(base_nnunet_path, 'labelsTs')

# 打印初始状态
print("\nChecking directories:")
print(f"Source image folder exists: {os.path.exists(final_image_folder)}")
print(f"Source label folder exists: {os.path.exists(final_label_folder)}")

# 确保目录存在
try:
    os.makedirs(base_nnunet_path, exist_ok=True)
    os.makedirs(nnunet_image_folder, exist_ok=True)
    os.makedirs(nnunet_label_folder, exist_ok=True)
    print(f"Successfully created/verified nnUNet directories")
except Exception as e:
    print(f"Error creating directories: {e}")

def get_patient_ids(folder_path, prefix='T1_'):
    """從檔案名稱中提取病人ID"""
    if not os.path.exists(folder_path):
        print(f"WARNING: Folder does not exist: {folder_path}")
        return set()
    
    patient_ids = set()
    try:
        for filename in os.listdir(folder_path):
            if filename.endswith('.nii.gz'):
                if filename.startswith(prefix):
                    # 對於T1影像: T1_XXXXX.nii.gz -> XXXXX
                    patient_id = filename[len(prefix):-7]
                    patient_ids.add(patient_id)
                elif prefix == 'T1_' and filename.startswith('label_'):
                    # 對於標籤: label_XXXXX.nii.gz -> XXXXX
                    patient_id = filename[6:-7]
                    patient_ids.add(patient_id)
        print(f"Found {len(patient_ids)} patients in {folder_path}")
    except Exception as e:
        print(f"Error reading directory {folder_path}: {e}")
    return patient_ids

def create_filtered_folder(image_folder, label_folder, output_image_folder, output_label_folder):
    """过滤并重命名文件，同时创建映射文件"""
    print("\nStarting file processing...")
    
    # 获取病人ID
    print("\nGetting patient IDs...")
    image_patients = get_patient_ids(image_folder, 'T1_')
    label_patients = get_patient_ids(label_folder, 'label_')  # 修改為正確的前綴
    
    # 打印一些檔案範例以供驗證
    print("\nSample files in directories:")
    print("Images directory:")
    for f in sorted(os.listdir(image_folder))[:3]:
        print(f"  - {f}")
    print("Labels directory:")
    for f in sorted(os.listdir(label_folder))[:3]:
        print(f"  - {f}")
    
    common_patients = image_patients.intersection(label_patients)
    print(f"\nFound {len(common_patients)} common patients")
    
    mapping_list = []
    processed_files = {"images": 0, "labels": 0}
    
    for index, patient_id in enumerate(sorted(common_patients), 1):
        new_filename = f"T1_{index:03d}"
        mapping_list.append(f"Original ID: {patient_id} -> New ID: {new_filename}")
        
        # 处理图像文件
        image_file = f'T1_{patient_id}.nii.gz'
        src_image_path = os.path.join(image_folder, image_file)
        dst_image_path = os.path.join(output_image_folder, f"{new_filename}_0000.nii.gz")
        
        # 处理标签文件
        label_file = f'label_{patient_id}.nii.gz'  # 修改為正確的命名格式
        src_label_path = os.path.join(label_folder, label_file)
        dst_label_path = os.path.join(output_label_folder, f"{new_filename}.nii.gz")
        
        # 複製檔案並記錄結果
        if os.path.exists(src_image_path):
            try:
                shutil.copy2(src_image_path, dst_image_path)
                processed_files["images"] += 1
                print(f"Copied image: {os.path.basename(dst_image_path)}")
            except Exception as e:
                print(f"Error copying image {src_image_path}: {e}")
        else:
            print(f"Source image not found: {src_image_path}")
        
        if os.path.exists(src_label_path):
            try:
                shutil.copy2(src_label_path, dst_label_path)
                processed_files["labels"] += 1
                print(f"Copied label: {os.path.basename(dst_label_path)}")
            except Exception as e:
                print(f"Error copying label {src_label_path}: {e}")
        else:
            print(f"Source label not found: {src_label_path}")
    
    # 写入映射文件
    try:
        with open(mapping_file, 'w') as f:
            f.write("Patient ID Mapping:\n")
            f.write("=================\n\n")
            for mapping in mapping_list:
                f.write(f"{mapping}\n")
            f.write(f"\nTotal number of patients: {len(mapping_list)}")
        print(f"\nMapping file created: {mapping_file}")
    except Exception as e:
        print(f"Error writing mapping file: {e}")
    
    return len(common_patients), processed_files

print("\nStarting file processing...")
common_count, processed_files = create_filtered_folder(
    final_image_folder, final_label_folder, 
    nnunet_image_folder, nnunet_label_folder
)

print(f"\nProcessing completed:")
print(f"- Total number of common patients: {common_count}")
print(f"- Files processed: {processed_files['images']} images, {processed_files['labels']} labels")
print(f"- Patient ID mapping has been saved to: {mapping_file}")

# 验证最终结果
print("\nFinal verification:")
if os.path.exists(nnunet_label_folder):
    label_files = os.listdir(nnunet_label_folder)
    print(f"Files in labels directory: {len(label_files)}")
    if len(label_files) > 0:
        print("Sample files:")
        for f in sorted(label_files)[:5]:
            print(f"  - {f}")
    else:
        print("WARNING: Labels directory is empty!")
else:
    print("WARNING: Labels directory does not exist!")