# import os
# import pandas as pd

# # 要扫描的根目录
# root_dir = '/mnt/1218/onlylian/sorted_modality_files'

# # 为 DataFrame 准备数据
# data = []

# # 遍历目录
# for subdir, _, files in os.walk(root_dir):
#     for file in files:
#         if file.endswith('.nii.gz'):
#             parts = subdir.split('/')
#             try:
#                 if len(parts) >= 8:  # 标准结构
#                     modalities = parts[4]  # 提取模态类型
#                     patient_id = parts[5]  # 提取患者信息
#                     date = parts[6]  # 提取日期
#                     name = file  # 提取文件名
#                 elif len(parts) == 7:  # 当模态类型与具体模态合并时
#                     modalities = parts[4]  # 提取模态类型
#                     patient_id = parts[5]  # 提取患者信息
#                     date = parts[6]  # 提取日期
#                     name = file  # 提取文件名
#                 elif len(parts) == 6:  # 当只有日期和模态类型时
#                     modalities = parts[4]  # 提取模态类型
#                     patient_id = parts[5]  # 提取患者信息
#                     date = ''  # 没有日期信息
#                     name = file  # 提取文件名
#                 else:
#                     print(f"Path {subdir} does not have enough parts to extract data")
#                     continue
#                 data.append([modalities, patient_id, date, name])
#                 print(f"Added: {modalities}/{patient_id}/{date}/{name}")  # 调试信息
#             except IndexError as e:
#                 print(f"Error processing path {subdir}: {e}")

# # 检查数据是否已经添加
# if not data:
#     print("No data found. Please check the directory structure and files.")

# # 创建 DataFrame
# df = pd.DataFrame(data, columns=['Modalities', 'patient_id', 'date', 'name'])
# # 输出到 Excel
# output_path = '/mnt/1218/onlylian/patient_modalities_files.xlsx'
# df.to_excel(output_path, index=False)

# print(f"Excel 文件已生成在: {output_path}")

import os
import pandas as pd
from datetime import datetime

# 设置根目录
root_dir = '/mnt/1218/onlylian/sorted_modality_files'

# 为 DataFrame 准备数据
data = []

# 添加计数器和集合来跟踪统计信息
total_files = 0
unique_patients = set()
unique_modalities = set()
processed_dates = set()

def process_file_path(file_path):
    """
    处理文件路径，提取相关信息
    示例路径: /mnt/1218/onlylian/sorted_modality_files/CT/2U5BG75W/20240802/1.1_CyberKnife_head(MAR)_20240802155850_6.nii.gz
    """
    try:
        # 分割路径
        parts = file_path.split('/')
        
        # 确保文件是 .nii.gz 格式
        if not parts[-1].endswith('.nii.gz'):
            return None
            
        # 从路径中提取所需信息
        modalities = parts[-4]  # CT
        patient_id = parts[-3]  # 2U5BG75W
        date = parts[-2]        # 20240802
        name = parts[-1]        # filename.nii.gz
        
        return [modalities, patient_id, date, name]
        
    except IndexError as e:
        print(f"Error processing path {file_path}: {e}")
        return None

# 获取当前时间作为运行标记
current_time = datetime.now().strftime("%Y%m%d_%H%M%S")

print("开始扫描文件...")

# 遍历目录
for subdir, _, files in os.walk(root_dir):
    for file in files:
        if file.endswith('.nii.gz'):
            full_path = os.path.join(subdir, file)
            file_info = process_file_path(full_path)
            
            if file_info:
                data.append(file_info)
                total_files += 1
                unique_patients.add(file_info[1])  # 添加病人ID到集合
                unique_modalities.add(file_info[0])  # 添加模态类型到集合
                processed_dates.add(file_info[2])    # 添加日期到集合
                
                # 每处理100个文件显示一次进度
                if total_files % 100 == 0:
                    print(f"已处理 {total_files} 个文件...")

# 检查数据是否已经添加
if not data:
    print("未找到数据。请检查目录结构和文件。")
else:
    # 创建 DataFrame
    df = pd.DataFrame(data, columns=['Modalities', 'patient_id', 'date', 'name'])
    
    # 生成带时间戳的输出文件名，以避免覆盖现有文件
    output_path = f'/mnt/1218/onlylian/patient_modalities_files_{current_time}.xlsx'
    
    # 输出到 Excel
    df.to_excel(output_path, index=False)
    
    # 打印统计信息
    print("\n处理完成！统计信息：")
    print(f"总文件数: {total_files}")
    print(f"病人数量: {len(unique_patients)}")
    print(f"模态类型数: {len(unique_modalities)}")
    print(f"检查日期数: {len(processed_dates)}")
    print("\n各模态类型及其文件数量：")
    modality_counts = df['Modalities'].value_counts()
    for modality, count in modality_counts.items():
        print(f"{modality}: {count} 个文件")
    
    print(f"\nExcel 文件已生成在: {output_path}")
    
    # 输出病人ID分布统计
    patient_file_counts = df['patient_id'].value_counts()
    print("\n病人文件数量分布：")
    print(f"最少文件数: {patient_file_counts.min()}")
    print(f"最多文件数: {patient_file_counts.max()}")
    print(f"平均每个病人文件数: {patient_file_counts.mean():.2f}")

# 错误处理和日志
try:
    # 同时保存一份日志文件
    log_path = f'/mnt/1218/onlylian/processing_log_{current_time}.txt'
    with open(log_path, 'w') as log_file:
        log_file.write(f"处理时间: {current_time}\n")
        log_file.write(f"总文件数: {total_files}\n")
        log_file.write(f"病人数量: {len(unique_patients)}\n")
        log_file.write(f"模态类型数: {len(unique_modalities)}\n")
        log_file.write(f"检查日期数: {len(processed_dates)}\n")
        log_file.write("\n各模态类型及其文件数量：\n")
        for modality, count in modality_counts.items():
            log_file.write(f"{modality}: {count} 个文件\n")
        
    print(f"\n处理日志已保存在: {log_path}")
except Exception as e:
    print(f"保存日志文件时出错: {e}")