import os
import pandas as pd
import ast
import SimpleITK as sitk

root_dir = "Data/MAMA-MIA"
data_dir= os.path.join(root_dir, "images")
output_dir = os.path.join(root_dir, "4D_images")
segmentation_dir = os.path.join(root_dir, "segmentations", "expert")
clinical_data = pd.read_excel(os.path.join(root_dir, "clinical_and_imaging_info.xlsx"))

data_dict = {}
for sub in os.listdir(data_dir):
    if os.path.isdir(os.path.join(data_dir, sub)):
        data_dict[sub] = {"images": [], "timepoints": []}
        for img in os.listdir(os.path.join(data_dir, sub)):
            if img.endswith(".nii.gz"):
                data_dict[sub]["images"].append(os.path.join(data_dir, sub, img))

        data_dict[sub]["images"].sort()

        acq_times = clinical_data[clinical_data["patient_id"] == sub]["acquisition_times"].values[0]
        if isinstance(acq_times, str):
            data_dict[sub]["timepoints"] = ast.literal_eval(acq_times)
        else:
            data_dict[sub]["timepoints"] = acq_times
        data_dict[sub]["segmentation"] = os.path.join(segmentation_dir, f"{sub}.nii.gz")

for sub in data_dict.keys():
    if not os.path.exists(os.path.join(output_dir, sub, f"{sub}.nii.gz")):
        image_list = []
        for img in data_dict[sub]["images"]:
            sitk_img = sitk.ReadImage(img)
            image_list.append(sitk_img)
        sitk_4D_image = sitk.JoinSeries(image_list)
        os.makedirs(os.path.join(output_dir, sub), exist_ok=True)
        sitk.WriteImage(sitk_4D_image, os.path.join(output_dir, sub, f"{sub}.nii.gz"))