# Py_add_dcm_DVH.py
# 这个代码是在安装了pydicom的环境运行的，
# 用来导出把dcm_folder里面的患者目录下面的All_ROI_DVH_data.txt
# 添加到RD.dcm的ds.DVHSequence里面

import os
import zipfile
import pydicom
from pydicom.dataset import Dataset
from pydicom.sequence import Sequence

# 设置DICOM文件夹路径
dcm_folder = r'X:\dcm_temp'

# 解压ZIP文件的函数
def unzip_files(dcm_folder):
    for root, _, files in os.walk(dcm_folder):
        for file in files:
            if file.endswith('.zip'):
                zip_file_path = os.path.join(root, file)
                extract_folder = os.path.splitext(zip_file_path)[0]
                if not os.path.exists(extract_folder):
                    os.makedirs(extract_folder)
                with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
                    zip_ref.extractall(extract_folder)
                print(f'Unzipped {file} to {extract_folder}')

# 处理DICOM和DVH的函数
def process_dvh_and_dcm_files(dcm_folder):
    for patient_dir in os.listdir(dcm_folder):
        patient_path = os.path.join(dcm_folder, patient_dir)
        if os.path.isdir(patient_path):
            # 查找DICOM文件
            rtstruct_file_path = None
            dicom_file_path = None
            for file in os.listdir(patient_path):
                file_path = os.path.join(patient_path, file)
                if file_path.endswith('.dcm'):
                    dcm_data = pydicom.dcmread(file_path, stop_before_pixels=True)
                    if dcm_data.Modality == 'RTSTRUCT':
                        rtstruct_file_path = file_path
                    elif dcm_data.Modality == 'RTDOSE':
                        dicom_file_path = file_path

            # 查找DVH数据
            txt_file_path = os.path.join(patient_path, "All_ROI_DVH_data.txt")

            if txt_file_path and rtstruct_file_path and dicom_file_path:
                # 读取RTSTRUCT文件以获取ROI的ROINumber
                rtstruct = pydicom.dcmread(rtstruct_file_path)
                roi_number_map = {roi_seq.ROIName: roi_seq.ROINumber for roi_seq in rtstruct.StructureSetROISequence}

                # 读取TXT文件中的DVH数据
                with open(txt_file_path, 'r') as file:
                    lines = file.readlines()

                # 读取现有的RTDOSE DICOM文件
                dicom_data = pydicom.dcmread(dicom_file_path)

                # 检查并初始化DVHSequence
                if not hasattr(dicom_data, 'DVHSequence'):
                    dicom_data.DVHSequence = Sequence()

                # 处理每个ROI的数据
                i = 0
                while i < len(lines):
                    if lines[i].startswith("ROI Name:"):
                        roi_name = lines[i].split(":")[1].strip()
                        dvh_type = lines[i+1].split(":")[1].strip()
                        dose_units = lines[i+2].split(":")[1].strip()
                        dose_type = lines[i+3].split(":")[1].strip()
                        dvh_dose_scaling = float(lines[i+4].split(":")[1].strip())
                        dvh_volume_units = lines[i+5].split(":")[1].strip()
                        dvh_number_of_bins = int(lines[i+6].split(":")[1].strip())
                        dvh_minimum_dose = float(lines[i+7].split(":")[1].strip())
                        dvh_maximum_dose = float(lines[i+8].split(":")[1].strip())
                        dvh_mean_dose = float(lines[i+9].split(":")[1].strip())
                        dvh_data = []
                        i += 11
                        while not lines[i].startswith("End of ROI"):
                            dose_value, volume_value = map(float, lines[i].strip().split(":"))
                            dvh_data.append((dose_value, volume_value))
                            i += 1

                        # 创建一个新的DVH序列项
                        dvh_item = Dataset()
                        dvh_item.DVHType = dvh_type
                        dvh_item.DoseUnits = dose_units
                        dvh_item.DoseType = dose_type
                        dvh_item.DVHDoseScaling = dvh_dose_scaling
                        dvh_item.DVHVolumeUnits = dvh_volume_units
                        dvh_item.DVHNumberOfBins = dvh_number_of_bins
                        dvh_item.DVHMinimumDose = dvh_minimum_dose
                        dvh_item.DVHMaximumDose = dvh_maximum_dose
                        dvh_item.DVHMeanDose = dvh_mean_dose
                        dvh_item.DVHData = '\\'.join([f'{dose}\\{volume}' for dose, volume in dvh_data])

                        # 获取Referenced ROI Number
                        referenced_roi_number = roi_number_map.get(roi_name)
                        if referenced_roi_number:
                            referenced_roi_item = Dataset()
                            referenced_roi_item.ReferencedROINumber = referenced_roi_number
                            referenced_roi_item.DVHROIContributionType = 'INCLUDED'
                            dvh_item.DVHReferencedROISequence = Sequence([referenced_roi_item])

                            # 将DVH序列项添加到DICOM文件
                            dicom_data.DVHSequence.append(dvh_item)

                    i += 1

                # 覆盖保存修改后的RTDOSE DICOM文件
                dicom_data.save_as(dicom_file_path)

                print(f'Modified DICOM file saved to: {dicom_file_path}')

# 执行解压和处理功能
if __name__ == '__main__':
    unzip_files(dcm_folder)
    process_dvh_and_dcm_files(dcm_folder)