#!/usr/bin/env python

import os
import sys
sys.path.append("/mnt/93E8-0534/AIDAN/HDBET/")
from HD_BET.run import run_hd_bet
from HD_BET.utils import maybe_mkdir_p, subfiles
import HD_BET

def hd_bet(input_file_or_dir,output_file_or_dir,mode,device,tta,pp=1,save_mask=0,overwrite_existing=1):

    if output_file_or_dir is None:
        output_file_or_dir = os.path.join(os.path.dirname(input_file_or_dir),
                                          os.path.basename(input_file_or_dir).split(".")[0] + "_bet")

    
    params_file = os.path.join(HD_BET.__path__[0], "model_final.py")
    config_file = os.path.join(HD_BET.__path__[0], "config.py")

    assert os.path.abspath(input_file_or_dir) != os.path.abspath(output_file_or_dir), "output must be different from input"

    if device == 'cpu':
        pass
    else:
        device = int(device)

    if os.path.isdir(input_file_or_dir):
        maybe_mkdir_p(output_file_or_dir)
        input_files = subfiles(input_file_or_dir, suffix='_0000.nii.gz', join=False)

        if len(input_files) == 0:
            raise RuntimeError("input is a folder but no nifti files (.nii.gz) were found in here")

        output_files = [os.path.join(output_file_or_dir, i) for i in input_files]
        input_files = [os.path.join(input_file_or_dir, i) for i in input_files]
    else:
        if not output_file_or_dir.endswith('.nii.gz'):
            output_file_or_dir += '.nii.gz'
            assert os.path.abspath(input_file_or_dir) != os.path.abspath(output_file_or_dir), "output must be different from input"
        
        output_files = [output_file_or_dir]
        input_files = [input_file_or_dir]

    if tta == 0:
        tta = False
    elif tta == 1:
        tta = True
    else:
        raise ValueError("Unknown value for tta: %s. Expected: 0 or 1" % str(tta))

    if overwrite_existing == 0:
        overwrite_existing = False
    elif overwrite_existing == 1:
        overwrite_existing = True
    else:
        raise ValueError("Unknown value for overwrite_existing: %s. Expected: 0 or 1" % str(overwrite_existing))

    if pp == 0:
        pp = False
    elif pp == 1:
        pp = True
    else:
        raise ValueError("Unknown value for pp: %s. Expected: 0 or 1" % str(pp))

    if save_mask == 0:
        save_mask = False
    elif save_mask == 1:
        save_mask = True
    else:
        raise ValueError("Unknown value for pp: %s. Expected: 0 or 1" % str(pp))

    run_hd_bet(input_files, output_files, mode, config_file, device, pp, tta, save_mask, overwrite_existing)


if __name__ == "__main__":
    print("\n########################")
    print("If you are using hd-bet, please cite the following paper:")
    print("Isensee F, Schell M, Tursunova I, Brugnara G, Bonekamp D, Neuberger U, Wick A, Schlemmer HP, Heiland S, Wick W,"
           "Bendszus M, Maier-Hein KH, Kickingereder P. Automated brain extraction of multi-sequence MRI using artificial"
           "neural networks. arXiv preprint arXiv:1901.11341, 2019.")
    print("########################\n")

    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--input', help='input. Can be either a single file name or an input folder. If file: must be '
                                       'nifti (.nii.gz) and can only be 3D. No support for 4d images, use fslsplit to '
                                       'split 4d sequences into 3d images. If folder: all files ending with .nii.gz '
                                       'within that folder will be brain extracted.', required=True, type=str)
    parser.add_argument('-o', '--output', help='output. Can be either a filename or a folder. If it does not exist, the folder'
                                     ' will be created', required=False, type=str)
    parser.add_argument('-mode', type=str, default='accurate', help='can be either \'fast\' or \'accurate\'. Fast will '
                                                                'use only one set of parameters whereas accurate will '
                                                                'use the five sets of parameters that resulted from '
                                                                'our cross-validation as an ensemble. Default: '
                                                                    'accurate',
                        required=False)
    parser.add_argument('-device', default='0', type=str, help='used to set on which device the prediction will run. '
                                                               'Must be either int or str. Use int for GPU id or '
                                                               '\'cpu\' to run on CPU. When using CPU you should '
                                                               'consider disabling tta. Default for -device is: 0',
                        required=False)
    parser.add_argument('-tta', default=1, required=False, type=int, help='whether to use test time data augmentation '
                                                                          '(mirroring). 1= True, 0=False. Disable this '
                                                                          'if you are using CPU to speed things up! '
                                                                          'Default: 1')
    parser.add_argument('-pp', default=1, type=int, required=False, help='set to 0 to disabe postprocessing (remove all'
                                                                         ' but the largest connected component in '
                                                                         'the prediction. Default: 1')
    parser.add_argument('-s', '--save_mask', default=1, type=int, required=False, help='if set to 0 the segmentation '
                                                                                       'mask will not be '
                                                                                       'saved')
    parser.add_argument('--overwrite_existing', default=1, type=int, required=False, help="set this to 0 if you don't "
                                                                                          "want to overwrite existing "
                                                                                          "predictions")

    args = parser.parse_args()

    hd_bet(args.input,args.output,args.mode,args.device,args.tta,args.pp,args.save_mask,args.overwrite_existing)