import os
import argparse
import subprocess

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 flo images", required=True)
    parser.add_argument("-o", "--output", help="Output folder for resampled images", required=True)
    args = parser.parse_args()

    # Get the script's location
    script_location = os.path.dirname(os.path.realpath(__file__))

    # Construct paths to reference image and affine transformation file
    ref_image_path = os.path.join(script_location, "..", "output", "C57BL6_mouse.iso_registerd.nii.gz")
    affine_trafo_path = os.path.join(script_location, "..", "output", "affineTrafo.txt")

    # 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 = [os.path.join(args.input, image) for image in os.listdir(args.input) if image.endswith(".nii") or image.endswith(".gz")]

    # Loop through flo images and perform resampling
    for flo_image in flo_images:
        output_filename = os.path.basename(flo_image).replace(".nii", "_registered.nii")
        output_path = os.path.join(args.output, output_filename)

        # Run the resample command
        command = [
            "reg_resample",
            "-ref", ref_image_path,
            "-flo", flo_image,
            "-res", output_path,
            "-trans", affine_trafo_path,
            "-inter", "0"
        ]

        subprocess.run(command)

if __name__ == "__main__":
    main()