import os
import shutil
import pandas as pd

def load_excel(file_path):
    """載入 Excel 文件"""
    return pd.read_excel(file_path, sheet_name='Sheet1')

def ensure_directories_exist(target_dir, modalities):
    """確保目標目錄存在"""
    for modality in modalities:
        os.makedirs(os.path.join(target_dir, modality), exist_ok=True)

def find_file_recursive(filename, search_path):
    """在給定目錄中遞迴搜索文件。"""
    for root, dirs, files in os.walk(search_path):
        if filename in files:
            return os.path.join(root, filename)
    return None

def extract_patient_and_date(file_path):
    """從文件路徑中提取病人檔名和日期。"""
    parts = file_path.split(os.sep)
    if len(parts) >= 4:
        return parts[-4], parts[-3]  # 假設病人檔名在倒數第四層，日期在倒數第三層
    return None, None

def is_swi_file(file_name):
    """檢查檔案名稱是否包含 'swi'"""
    return 'swi' in file_name.lower()

def copy_files(df, source_dir, target_dir):
    """遍歷數據框並將對應的 .warp.nii.gz 文件複製到相應的模態資料夾，若無 .warp.nii.gz 則複製原始的 .nii.gz 文件，保留病人檔名和日期"""
    processed_patients = set()
    
    for index, row in df.iterrows():
        file_name = row['name']
        modality = row['modality_nn']
        
        # 檢查是否為 swi 檔案
        if is_swi_file(file_name):
            print(f'跳過 swi 檔案: {file_name}')
            continue
        
        # 修改來源檔名：將 .nii.gz 替換為 .warp.nii.gz
        warp_file_name = file_name.replace('.nii.gz', '.warp.nii.gz')
        
        # 在源目錄中遞迴搜索 .warp.nii.gz 文件
        warp_file_path = find_file_recursive(warp_file_name, source_dir)
        
        if warp_file_path:
            # 如果找到 .warp.nii.gz 文件，複製該文件
            patient_id, date = extract_patient_and_date(warp_file_path)
            if patient_id and date:
                if patient_id not in processed_patients:
                    processed_patients.add(patient_id)
                
                # 複製到個別的模態資料夾
                target_dir_path = os.path.join(target_dir, modality, patient_id, date)
                os.makedirs(target_dir_path, exist_ok=True)
                target_file_path = os.path.join(target_dir_path, warp_file_name)
                shutil.copy2(warp_file_path, target_file_path)
                print(f'已複製 {warp_file_name} 到 {target_dir_path} 資料夾。')

                # 如果是 CT 或 CTC，也複製到 CT+CTC 資料夾
                if modality in ['CT', 'CTC']:
                    combined_target_dir_path = os.path.join(target_dir, 'CT+CTC', patient_id, date)
                    os.makedirs(combined_target_dir_path, exist_ok=True)
                    combined_target_file_path = os.path.join(combined_target_dir_path, warp_file_name)
                    shutil.copy2(warp_file_path, combined_target_file_path)
                    print(f'已複製 {warp_file_name} 到 {combined_target_dir_path} 資料夾。')
        else:
            # 如果找不到 .warp.nii.gz 文件，則複製原始的 .nii.gz 文件
            original_file_path = find_file_recursive(file_name, source_dir)
            if original_file_path:
                patient_id, date = extract_patient_and_date(original_file_path)
                if patient_id and date:
                    if patient_id not in processed_patients:
                        processed_patients.add(patient_id)

                    # 複製到個別的模態資料夾
                    target_dir_path = os.path.join(target_dir, modality, patient_id, date)
                    os.makedirs(target_dir_path, exist_ok=True)
                    target_file_path = os.path.join(target_dir_path, file_name)
                    shutil.copy2(original_file_path, target_file_path)
                    print(f'已複製 {file_name} 到 {target_dir_path} 資料夾。')

                    # 如果是 CT 或 CTC，也複製到 CT+CTC 資料夾
                    if modality in ['CT', 'CTC']:
                        combined_target_dir_path = os.path.join(target_dir, 'CT+CTC', patient_id, date)
                        os.makedirs(combined_target_dir_path, exist_ok=True)
                        combined_target_file_path = os.path.join(combined_target_dir_path, file_name)
                        shutil.copy2(original_file_path, combined_target_file_path)
                        print(f'已複製 {file_name} 到 {combined_target_dir_path} 資料夾。')
            else:
                print(f'{file_name} 和 {warp_file_name} 均在源目錄中不存在。')

def main():
    file_path = '/mnt/1218/onlylian/modality_by_nn.xlsx'
    source_dir = '/mnt/1218/onlylian/N6_24Q3'
    target_dir = '/mnt/1218/onlylian/sorted_modality_files'
    
    modalities = ['CT', 'CTC', 'T1', 'T1C', 'T2', 'FLAIR', 'CT+CTC']  # 加入 CT+CTC
    
    df = load_excel(file_path)
    ensure_directories_exist(target_dir, modalities)
    copy_files(df, source_dir, target_dir)
    
    print('文件分類和複製完成。')

if __name__ == "__main__":
    main()