"""
This script enables to launch predictions with SynthSeg from the terminal.

If you use this code, please cite one of the SynthSeg papers:
https://github.com/BBillot/SynthSeg/blob/master/bibtex.bib

Copyright 2020 Benjamin Billot

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in
compliance with the License. You may obtain a copy of the License at
https://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software distributed under the License is
distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied. See the License for the specific language governing permissions and limitations under the
License.
"""

# python imports
import os
import sys
from argparse import ArgumentParser

# add main folder to python path and import ./SynthSeg/predict_synthseg.py
synthseg_home = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(sys.argv[0]))))
sys.path.append(synthseg_home)
model_dir = os.path.join(synthseg_home, 'models')
labels_dir = os.path.join(synthseg_home, 'data/labels_classes_priors')
from SynthSeg.predict_synthseg import predict


# parse arguments
parser = ArgumentParser(description="SynthSeg", epilog='\n')

# input/outputs
parser.add_argument("--i", help="Image(s) to segment. Can be a path to an image or to a folder.")
parser.add_argument("--o", help="Segmentation output(s). Must be a folder if --i designates a folder.")
parser.add_argument("--parc", action="store_true", help="(optional) Whether to perform cortex parcellation.")
parser.add_argument("--robust", action="store_true", help="(optional) Whether to use robust predictions (slower).")
parser.add_argument("--fast", action="store_true", help="(optional) Bypass some postprocessing for faster predictions.")
parser.add_argument("--ct", action="store_true", help="(optional) Clip intensities to [0,80] for CT scans.")
parser.add_argument("--vol", help="(optional) Path to output CSV file with volumes (mm3) for all regions and subjects.")
parser.add_argument("--qc", help="(optional) Path to output CSV file with qc scores for all subjects.")
parser.add_argument("--post", help="(optional) Posteriors output(s). Must be a folder if --i designates a folder.")
parser.add_argument("--resample", help="(optional) Resampled image(s). Must be a folder if --i designates a folder.")
parser.add_argument("--crop", nargs='+', type=int, help="(optional) Size of 3D patches to analyse. Default is 192.")
parser.add_argument("--threads", type=int, default=1, help="(optional) Number of cores to be used. Default is 1.")
parser.add_argument("--cpu", action="store_true", help="(optional) Enforce running with CPU rather than GPU.")
parser.add_argument("--v1", action="store_true", help="(optional) Use SynthSeg 1.0 (updated 25/06/22).")

# check for no arguments
if len(sys.argv) < 2:
    parser.print_help()
    sys.exit(1)

# parse commandline
args = vars(parser.parse_args())

# print SynthSeg version and checks boolean params for SynthSeg-robust
if args['robust']:
    args['fast'] = True
    assert not args['v1'], 'The flag --v1 cannot be used with --robust since SynthSeg-robust only came out with 2.0.'
    version = 'SynthSeg-robust 2.0'
else:
    version = 'SynthSeg 1.0' if args['v1'] else 'SynthSeg 2.0'
    if args['fast']:
        version += ' (fast)'
print('\n' + version + '\n')

# enforce CPU processing if necessary
if args['cpu']:
    print('using CPU, hiding all CUDA_VISIBLE_DEVICES')
    os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

# limit the number of threads to be used if running on CPU
import tensorflow as tf
if args['threads'] == 1:
    print('using 1 thread')
else:
    print('using %s threads' % args['threads'])
tf.config.threading.set_inter_op_parallelism_threads(args['threads'])
tf.config.threading.set_intra_op_parallelism_threads(args['threads'])

# path models
if args['robust']:
    args['path_model_segmentation'] = os.path.join(model_dir, 'synthseg_robust_2.0.h5')
else:
    args['path_model_segmentation'] = os.path.join(model_dir, 'synthseg_2.0.h5')
args['path_model_parcellation'] = os.path.join(model_dir, 'synthseg_parc_2.0.h5')
args['path_model_qc'] = os.path.join(model_dir, 'synthseg_qc_2.0.h5')

# path labels
args['labels_segmentation'] = os.path.join(labels_dir, 'synthseg_segmentation_labels_2.0.npy')
args['labels_denoiser'] = os.path.join(labels_dir, 'synthseg_denoiser_labels_2.0.npy')
args['labels_parcellation'] = os.path.join(labels_dir, 'synthseg_parcellation_labels.npy')
args['labels_qc'] = os.path.join(labels_dir, 'synthseg_qc_labels_2.0.npy')
args['names_segmentation_labels'] = os.path.join(labels_dir, 'synthseg_segmentation_names_2.0.npy')
args['names_parcellation_labels'] = os.path.join(labels_dir, 'synthseg_parcellation_names.npy')
args['names_qc_labels'] = os.path.join(labels_dir, 'synthseg_qc_names_2.0.npy')
args['topology_classes'] = os.path.join(labels_dir, 'synthseg_topological_classes_2.0.npy')
args['n_neutral_labels'] = 19

# use previous model if needed
if args['v1']:
    args['path_model_segmentation'] = os.path.join(model_dir, 'synthseg_1.0.h5')
    args['labels_segmentation'] = args['labels_segmentation'].replace('_2.0.npy', '.npy')
    args['labels_qc'] = args['labels_qc'].replace('_2.0.npy', '.npy')
    args['names_segmentation_labels'] = args['names_segmentation_labels'].replace('_2.0.npy', '.npy')
    args['names_qc_labels'] = args['names_qc_labels'].replace('_2.0.npy', '.npy')
    args['topology_classes'] = args['topology_classes'].replace('_2.0.npy', '.npy')
    args['n_neutral_labels'] = 18

# run prediction
predict(path_images=args['i'],
        path_segmentations=args['o'],
        path_model_segmentation=args['path_model_segmentation'],
        labels_segmentation=args['labels_segmentation'],
        robust=args['robust'],
        fast=args['fast'],
        v1=args['v1'],
        do_parcellation=args['parc'],
        n_neutral_labels=args['n_neutral_labels'],
        names_segmentation=args['names_segmentation_labels'],
        labels_denoiser=args['labels_denoiser'],
        path_posteriors=args['post'],
        path_resampled=args['resample'],
        path_volumes=args['vol'],
        path_model_parcellation=args['path_model_parcellation'],
        labels_parcellation=args['labels_parcellation'],
        names_parcellation=args['names_parcellation_labels'],
        path_model_qc=args['path_model_qc'],
        labels_qc=args['labels_qc'],
        path_qc_scores=args['qc'],
        names_qc=args['names_qc_labels'],
        cropping=args['crop'],
        topology_classes=args['topology_classes'],
        ct=args['ct'])