import sys
import numpy as np
import nibabel as nib
import os
from nibabel.freesurfer.mghformat import MGHHeader
from scipy.ndimage import affine_transform
from numpy.linalg import inv
import glob


import argparse
import torch

#import nnunet.inference.predict
from nnunet.inference.predict import predict_from_folder
from nnunet.paths import default_plans_identifier, network_training_output_dir, default_cascade_trainer, default_trainer
from batchgenerators.utilities.file_and_folder_operations import join, isdir
from nnunet.utilities.task_name_id_conversion import convert_id_to_task_name


def correct_header(input, output):
    #args = parser.parse_args()
    #input = args.input_folder
    #output =
    img1 = nib.load(input)
    corr_affine = img1.get_qform()
    img1.set_sform(corr_affine)
    img1.update_header()
    img1.set_data_dtype(img1.get_data_dtype())
    nib.save(img1, output)

def correct_num_col(input, output):
    img1 = nib.load(input)
    #my_mr = os.path.join(path_in, case_name)
    #mr = nib.load(my_mr)
    aff = img1.get_qform()
    mr_mat = img1.get_fdata()
    mr_shape = mr_mat.shape
    if len(mr_shape) > 3:
        l0 = mr_mat.shape[0]
        l1 = mr_mat.shape[1]
        l2 = mr_mat.shape[2]
        l3 = mr_mat.shape[3]
        new_mr = mr_mat.reshape(l0,l1,l2)
        new_mr = nib.Nifti1Image(new_mr, affine = aff)
        new_mr.set_data_dtype(img1.get_data_dtype())
        nib.save(new_mr, output)
#   nib.save(new_mr, output)
#else:
#   new_mr = nib.Nifti1Image(mr_mat, affine = aff)
#   new_mr.to_filename(output)
#   nib.save(new_mr, output)

def lpi_conform(input, output):
    img = nib.load(input)
    h1 = MGHHeader.from_header(img)
    
    h1.set_data_shape([256, 256, 256])
    h1.set_zooms([1, 1, 1])
    
    h1['Mdc'] = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
    h1['Pxyz_c'] = img.affine.dot(np.hstack((np.array(img.shape[:3]) / 2.0, [1])))[:3]
    
    ras2ras=np.array([[1.0, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
    vox2vox=inv(h1.get_affine())@ ras2ras @ img.affine
    
    new_data = affine_transform(img.get_fdata(), inv(vox2vox), output_shape=h1.get_data_shape(), order=1)
    new_img = nib.MGHImage(new_data, h1.get_affine(), h1)
    
    #new_img.set_data_dtype(np.uint8)
    new_img.set_data_dtype(img.get_data_dtype())
    nib.save(new_img, output)


def arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--input_folder', help="Must contain all modalities for each patient in the correct"
                        " order (same as training). Files must be named "
                        "CASENAME_XXXX.nii.gz where XXXX is the modality "
                        "identifier (0000, 0001, etc)", required=True)
    parser.add_argument('-o', '--output_folder', required=True, help="folder for saving predictions")
                        #parser.add_argument('-t', '--task_name', help='task name or task ID, required.',
                        #                    default=default_plans_identifier, required=True)
                        
    parser.add_argument('-tr', '--trainer_class_name',
                                            help='Name of the nnUNetTrainer used for 2D U-Net, full resolution 3D U-Net and low resolution '
                                            'U-Net. The default is %s. If you are running inference with the cascade and the folder '
                                            'pointed to by --lowres_segmentations does not contain the segmentation maps generated by '
                                            'the low resolution U-Net then the low resolution segmentation maps will be automatically '
                                            'generated. For this case, make sure to set the trainer class here that matches your '
                                            '--cascade_trainer_class_name (this part can be ignored if defaults are used).'
                                            % default_trainer,
                                            required=False,
                                            default=default_trainer)
    parser.add_argument('-ctr', '--cascade_trainer_class_name',
                                            help="Trainer class name used for predicting the 3D full resolution U-Net part of the cascade."
                                            "Default is %s" % default_cascade_trainer, required=False,
                                            default=default_cascade_trainer)
                        
                        #parser.add_argument('-m', '--model', help="2d, 3d_lowres, 3d_fullres or 3d_cascade_fullres. Default: 3d_fullres",
                        #                    default="3d_fullres", required=False)
                        
    parser.add_argument('-p', '--plans_identifier', help='do not touch this unless you know what you are doing',
                                            default=default_plans_identifier, required=False)
                        
    parser.add_argument('-f', '--folds', nargs='+', default='None',
                                            help="folds to use for prediction. Default is None which means that folds will be detected "
                                            "automatically in the model output folder")
                        
    parser.add_argument('-z', '--save_npz', required=False, action='store_true',
                                            help="use this if you want to ensemble these predictions with those of other models. Softmax "
                                            "probabilities will be saved as compressed numpy arrays in output_folder and can be "
                                            "merged between output_folders with nnUNet_ensemble_predictions")
                        
    parser.add_argument('-l', '--lowres_segmentations', required=False, default='None',
                                            help="if model is the highres stage of the cascade then you can use this folder to provide "
                                            "predictions from the low resolution 3D U-Net. If this is left at default, the "
                                            "predictions will be generated automatically (provided that the 3D low resolution U-Net "
                                            "network weights are present")
                        
    parser.add_argument("--part_id", type=int, required=False, default=0, help="Used to parallelize the prediction of "
                                            "the folder over several GPUs. If you "
                                            "want to use n GPUs to predict this "
                                            "folder you need to run this command "
                                            "n times with --part_id=0, ... n-1 and "
                                            "--num_parts=n (each with a different "
                                            "GPU (for example via "
                                            "CUDA_VISIBLE_DEVICES=X)")
                        
    parser.add_argument("--num_parts", type=int, required=False, default=1,
                                            help="Used to parallelize the prediction of "
                                            "the folder over several GPUs. If you "
                                            "want to use n GPUs to predict this "
                                            "folder you need to run this command "
                                            "n times with --part_id=0, ... n-1 and "
                                            "--num_parts=n (each with a different "
                                            "GPU (via "
                                            "CUDA_VISIBLE_DEVICES=X)")
                        
    parser.add_argument("--num_threads_preprocessing", required=False, default=6, type=int, help=
                                            "Determines many background processes will be used for data preprocessing. Reduce this if you "
                                            "run into out of memory (RAM) problems. Default: 6")
                        
    parser.add_argument("--num_threads_nifti_save", required=False, default=2, type=int, help=
                                            "Determines many background processes will be used for segmentation export. Reduce this if you "
                                            "run into out of memory (RAM) problems. Default: 2")
                        
    parser.add_argument("--disable_tta", required=False, default=False, action="store_true",
                                            help="set this flag to disable test time data augmentation via mirroring. Speeds up inference "
                                            "by roughly factor 4 (2D) or 8 (3D)")
                        
    parser.add_argument("--overwrite_existing", required=False, default=False, action="store_true",
                                            help="Set this flag if the target folder contains predictions that you would like to overwrite")
                        
    parser.add_argument("--mode", type=str, default="normal", required=False, help="Hands off!")
    parser.add_argument("--all_in_gpu", type=str, default="None", required=False, help="can be None, False or True. "
                                            "Do not touch.")
    parser.add_argument("--step_size", type=float, default=0.5, required=False, help="don't touch")
                        # parser.add_argument("--interp_order", required=False, default=3, type=int,
                        #                     help="order of interpolation for segmentations, has no effect if mode=fastest. Do not touch this.")
                        # parser.add_argument("--interp_order_z", required=False, default=0, type=int,
                        #                     help="order of interpolation along z is z is done differently. Do not touch this.")
                        # parser.add_argument("--force_separate_z", required=False, default="None", type=str,
                        #                     help="force_separate_z resampling. Can be None, True or False, has no effect if mode=fastest. "
                        #                          "Do not touch this.")
    parser.add_argument('-chk',
                                            help='checkpoint name, default: model_final_checkpoint',
                                            required=False,
                                            default='model_final_checkpoint')
    parser.add_argument('--disable_mixed_precision', default=False, action='store_true', required=False,
                                            help='Predictions are done with mixed precision by default. This improves speed and reduces '
                                            'the required vram. If you want to disable mixed precision you can set this flag. Note '
                                            'that yhis is not recommended (mixed precision is ~2x faster!)')
                        
    return parser


def inference(parser):
    args = parser.parse_args()
    directory = 'preprocessed'
    path = os.path.join(args.input_folder, directory)
    input_folder = path
    #input_folder = args.input_folder
    output_folder = args.output_folder
    part_id = args.part_id
    num_parts = args.num_parts
    folds = args.folds
    save_npz = args.save_npz
    lowres_segmentations = args.lowres_segmentations
    num_threads_preprocessing = args.num_threads_preprocessing
    num_threads_nifti_save = args.num_threads_nifti_save
    disable_tta = args.disable_tta
    step_size = args.step_size
    # interp_order = args.interp_order
    # interp_order_z = args.interp_order_z
    # force_separate_z = args.force_separate_z
    overwrite_existing = args.overwrite_existing
    mode = args.mode
    all_in_gpu = args.all_in_gpu
    #model = args.model
    model = '3d_fullres'
    trainer_class_name = args.trainer_class_name
    cascade_trainer_class_name = args.cascade_trainer_class_name
    
    #task_name = args.task_name
    task_name = 'Task054_Mri'
    
    if not task_name.startswith("Task"):
        task_id = int(task_name)
        task_name = convert_id_to_task_name(task_id)

    assert model in ["2d", "3d_lowres", "3d_fullres", "3d_cascade_fullres"], "-m must be 2d, 3d_lowres, 3d_fullres or " \
        "3d_cascade_fullres"
    

    if lowres_segmentations == "None":
        lowres_segmentations = None
    
    if isinstance(folds, list):
        if folds[0] == 'all' and len(folds) == 1:
            pass
        else:
            folds = [int(i) for i in folds]
    elif folds == "None":
        folds = None
    else:
        raise ValueError("Unexpected value for argument folds")
    
    assert all_in_gpu in ['None', 'False', 'True']
    if all_in_gpu == "None":
        all_in_gpu = None
    elif all_in_gpu == "True":
        all_in_gpu = True
    elif all_in_gpu == "False":
        all_in_gpu = False


    trainer = trainer_class_name
        
    model_folder_name = join(network_training_output_dir, model, task_name, trainer + "__" +
                                 args.plans_identifier)
    print("using model stored in ", model_folder_name)
    assert isdir(model_folder_name), "model output folder not found. Expected: %s" % model_folder_name

    predict_from_folder(model_folder_name, input_folder, output_folder, folds, save_npz, num_threads_preprocessing,
                                                         num_threads_nifti_save, lowres_segmentations, part_id, num_parts, not disable_tta,
                                                         overwrite_existing=overwrite_existing, mode=mode, overwrite_all_in_gpu=all_in_gpu,
                                                         mixed_precision=not args.disable_mixed_precision,
                                                         step_size=step_size, checkpoint_name=args.chk)




def preprocessing(parser):
    args = parser.parse_args()
    input = args.input_folder
    #parser = arguments()
    #args = parser.parse_args()
    #input_folder = args.input_folder
    input_folder = input
    os.chdir(input)
    directory = 'preprocessed'
    path = os.path.join(input_folder, directory)
    if not os.path.exists(path):
        os.mkdir(path)
    for file in glob.glob('*_0000.nii.gz'):
        input_img = os.path.join(input_folder, file)
        output_img = os.path.join(path, file)
        print('Pre-processing: ', file)
        correct_header(input_img, output_img)
        correct_num_col(output_img, output_img)
        lpi_conform(output_img, output_img)


def main():
    parser = arguments()
    preprocessing(parser)
    
    parser = arguments()
    print('Segmenting')
    inference(parser)


if __name__ == "__main__":
    main()
