import argparse
import subprocess
import shlex
import os
import glob
import csv
import nibabel as nib
import numpy as np
from tqdm import tqdm

def main():
    parser = argparse.ArgumentParser(description="Image Registration and Resampling")
    
    parser.add_argument("-i", "--input", required=True, help="Path to main folder of the nii Converted Data that are already processed with AIDAmri, e.g proc_data")
    parser.add_argument("-m", "--masks", required=True, help="Path to the folder containing the drawn masks. They should be in the MTPL space!")
    
    args = parser.parse_args()
    
    missing_matrix_paths = []

    # Parsing for all BiasBet.nii.gz files in T2w folders after a given input path.
    T2BetPath = os.path.join(args.input,"**","anat","*BiasBet.nii.gz")
    T2BetFiles = glob.glob(T2BetPath,recursive=True)
    
    # Parsing for all given masks.
    FloatingMasksPath = os.path.join(args.masks, "*.nii*")
    FloatingMasksFiles = glob.glob(FloatingMasksPath)
    
    total_iterations = len(T2BetFiles) * len(FloatingMasksFiles)
    progress_bar = tqdm(total=total_iterations, desc='Registration Progress (Mask to Anatomical T2w)', unit='files')

    for ff in T2BetFiles:
        for mm in FloatingMasksFiles:
            MaskName = os.path.basename(mm).replace(".nii","").replace(".gz","")
            NewMaskName = MaskName + "_anat" + ".nii.gz"
            NewMaskFolder = os.path.join(os.path.dirname(ff),"RegisteredTractMasks")
            NewMaskAddress = os.path.join(NewMaskFolder,NewMaskName)
            if not os.path.exists(NewMaskFolder):
                os.makedirs(NewMaskFolder)

            Affine_matrix_files = glob.glob(os.path.join(os.path.dirname(ff),"*MatrixAff.txt"))
            if not Affine_matrix_files:
                missing_matrix_paths.append(os.path.dirname(ff))
                continue
            Affine_matrix = Affine_matrix_files[0]

            # Execute the command
            command = f"reg_resample -ref {ff} -flo {mm} -res {NewMaskAddress} -trans {Affine_matrix} -inter 0"        
            command_args = shlex.split(command)

            try:
                result = subprocess.run(command_args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                print(f"Output:\n{result.stdout.decode()}")
                print(f"Errors:\n{result.stderr.decode()}")
            except Exception as e:
                print(f"An error occurred: {e}")
                
            progress_bar.update(1)
    
    registered_masks_path = os.path.join(args.input,"**","RegisteredTractMasks", "*.nii.gz")
    registered_masks_files = glob.glob(registered_masks_path, recursive=True)

    total_iterations = len(registered_masks_files)
    progress_bar = tqdm(total=total_iterations, desc='Flipping Masks', unit='files')

    for mask_file in registered_masks_files:
        # Load the mask file
        mask_nifti = nib.load(mask_file)
        mask_data = mask_nifti.get_fdata()

        # Flip the mask along the Y and Z axes
        #flipped_mask = np.flip(np.flip(mask_data, axis=0), axis=1)
        flipped_mask = np.flip(mask_data, axis=0)
        

        # Determine the directory to save the adjusted mask
        mask_dir = os.path.dirname(os.path.dirname(mask_file))
        adjusted_masks_dir = os.path.join(mask_dir, "RegisteredTractMasks")



        # Define the new file name and save the flipped mask
        if "anat" in mask_file:
            adjusted_mask_file = os.path.join(adjusted_masks_dir, os.path.basename(mask_file).replace(".nii.gz",".nii.gz"))
            nib.save(nib.Nifti1Image(flipped_mask, mask_nifti.affine), adjusted_mask_file)

        progress_bar.update(1)

    progress_bar.close()
    print("Flipping of RegisteredTractMasks and saving to RegisteredTractMasks_adjusted: DONE\n")
        
    progress_bar.close()
    print("Registered all masks to t2w space of individual mice: DONE\n")
    
    # Parsing for all SmoothMicoBet.nii.gz files in DTI folders after a given input path.
    DTIBetPath = os.path.join(args.input,"**","dwi*","*SmoothMicoBet.nii.gz")
    DTIBetFiles = glob.glob(DTIBetPath,recursive=True)
    
    total_iterations = len(DTIBetFiles) * len(FloatingMasksFiles)
    progress_bar = tqdm(total=total_iterations, desc='Registration Progress (to diffusion dwi) and adjusting for fa, ad, md maps', unit='files')

    for ff in DTIBetFiles:
        temp = os.path.dirname(os.path.dirname(ff))
        t2MasksPath = os.path.join(temp,"anat","RegisteredTractMasks","*nii.gz")
        FloatingMasksFiles = glob.glob(t2MasksPath)
        for mm in FloatingMasksFiles:
            MaskName = os.path.basename(mm).replace("anat","dwi")
            NewMaskName = MaskName
            NewMaskFolder = os.path.join(os.path.dirname(ff),"RegisteredTractMasks")
            NewMaskAddress = os.path.join(NewMaskFolder,NewMaskName)
            if not os.path.exists(NewMaskFolder):
                os.makedirs(NewMaskFolder)

            Affine_matrix_files = glob.glob(os.path.join(os.path.dirname(ff),"*MatrixAff.txt"))
            if not Affine_matrix_files:
                missing_matrix_paths.append(os.path.dirname(ff))
                continue
            Affine_matrix = Affine_matrix_files[0]

            command = f"reg_resample -ref {ff} -flo {mm} -res {NewMaskAddress} -trans {Affine_matrix} -inter 0"        
            command_args = shlex.split(command)

            try:
                result = subprocess.run(command_args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                print(f"Output:\n{result.stdout.decode()}")
                print(f"Errors:\n{result.stderr.decode()}")
            except Exception as e:
                print(f"An error occurred: {e}")
                
            progress_bar.update(1)
    
    progress_bar.close()
    print("Registered all masks to t2w space of individual mice: DONE\n")

    # Adjust by flipping for the fa0 and ad files
    # Locate all files in the "RegisteredTractMasks" folder
    registered_masks_path = os.path.join(args.input,"**","RegisteredTractMasks", "*.nii.gz")
    registered_masks_files = glob.glob(registered_masks_path, recursive=True)

    total_iterations = len(registered_masks_files)
    progress_bar = tqdm(total=total_iterations, desc='Flipping Masks', unit='files')

    for mask_file in registered_masks_files:
        # Load the mask file
        mask_nifti = nib.load(mask_file)
        mask_data = mask_nifti.get_fdata()

        # Flip the mask along the Y and Z axes
        #flipped_mask = np.flip(np.flip(mask_data, axis=1), axis=2)
        flipped_mask = np.flip(mask_data, axis=2)
        

        # Determine the directory to save the adjusted mask
        mask_dir = os.path.dirname(os.path.dirname(mask_file))
        adjusted_masks_dir = os.path.join(mask_dir, "DSI_studio", "RegisteredTractMasks_adjusted")

        # Create the directory if it doesn't exist
        if not os.path.exists(adjusted_masks_dir) and not "anat" in mask_file:
            os.makedirs(adjusted_masks_dir)

        # Define the new file name and save the flipped mask
        if not "anat" in mask_file:
            adjusted_mask_file = os.path.join(adjusted_masks_dir, os.path.basename(mask_file).replace(".nii.gz","_flipped.nii.gz"))
            nib.save(nib.Nifti1Image(flipped_mask, mask_nifti.affine), adjusted_mask_file)

        progress_bar.update(1)

    progress_bar.close()
    print("Flipping of RegisteredTractMasks and saving to RegisteredTractMasks_adjusted: DONE\n")

    # Save the missing matrix paths to a CSV file
    missing_matrix_csv_path = os.path.join(args.input, "missing_matrix_paths.csv")
    with open(missing_matrix_csv_path, mode='w', newline='') as file:
        writer = csv.writer(file)
        if not missing_matrix_paths:
            writer.writerow(["Registration of all masks/maps was executed. All folders contained an Affine transformation matrix"])
        else:
            for path in missing_matrix_paths:
                writer.writerow([path])
    
            print("Saved missing matrix paths to CSV: DONE\n")

        
if __name__ == "__main__":
    main()