import os
import argparse
import numpy as np
import nibabel as nib
import glob

def main():
    parser = argparse.ArgumentParser(description="Resample multiple flo images using a reference image and an affine transformation file.")
    parser.add_argument("-i", "--input", help="Input folder containing viraltracing images", required=True)
    parser.add_argument("-o", "--output", help="Output folder for flipped images", required=True)
    args = parser.parse_args()

    # Create the output folder if it doesn't exist
    if not os.path.exists(args.output):
        os.makedirs(args.output)

    # List all flo images in the input folder
    flo_images = glob.glob(os.path.join(args.input, "**", "*.nii.gz"), recursive=True)

    # Loop through flo images and perform resampling
    for flo_image in flo_images:
        output_filename = os.path.basename(flo_image).replace(".nii.gz", "_flipped.nii.gz")
        output_path = os.path.join(args.output, output_filename)
        mask_output_path = os.path.join(args.output, output_filename.replace("_flipped.nii.gz", "_mask.nii.gz"))
        
        # Load the image using nibabel
        img = nib.load(flo_image)
        data = img.get_fdata()
        
        # Perform the dimension manipulation to achieve the desired shape
        data_swapped = np.swapaxes(data, 0, 2)  # Swap axes as per the desired shape (264, 160, 228) -> (228, 160, 264)
        data_flipped = np.flip(data_swapped, axis=2)  # Flip along the third axis
        
        # Save the modified data back to a NIfTI file
        img_flipped = nib.Nifti1Image(data_flipped, img.affine)
        nib.save(img_flipped, output_path)
        
        # Create mask data
        mask_data = np.where(data_flipped != 0, 1, 0)
        mask_img = nib.Nifti1Image(mask_data.astype(np.uint8), img.affine)
        nib.save(mask_img, mask_output_path)

if __name__ == "__main__":
    main()