<reponame>radiome-flow/radiome<filename>radiome/core/cli.py
import argparse
import logging
import os
import shutil
import subprocess
import sys
import tempfile
from datetime import datetime
from pathlib import Path

import psutil
import yaml

from radiome.core import __version__, __author__, __email__
from radiome.core import context
from radiome.core.execution import pipeline
from radiome.core.utils.s3 import S3Resource


def parse_args(args):
    parser = argparse.ArgumentParser(description='Radiome Pipeline Runner')
    parser.add_argument('bids_dir', help='The directory with the input dataset '
                                         'formatted according to the BIDS standard. '
                                         'Use the format s3://bucket/path/to/bidsdir'
                                         'to read data directly from an S3 bucket.'
                                         ' This may require AWS S3 credentials specified via the'
                                         ' --aws_input_creds or --aws_input_creds_profile option.')
    parser.add_argument('outputs_dir', help='The directory where the output files should be stored.'
                                            'Use the format s3://bucket/path/to/bidsdir'
                                            'to write data directly to an S3 bucket.'
                                            ' This may require AWS S3 credentials specified via the'
                                            ' --aws_input_creds or --aws_input_creds_profile option.')
    parser.add_argument('--config_file', help='The path for the pipeline config file.', required=True)
    parser.add_argument('--working_dir', help='The local directory where temporary file and logs reside. If not set,'
                                              'the output dir will be used. If the output dir is an S3 bucket,'
                                              'you must provide a local path or a temporary directory will be created.')
    parser.add_argument('--participant_label', help='The label of the participant'
                                                    ' that should be analyzed. The label '
                                                    'corresponds to sub-<participant_label> from the BIDS spec '
                                                    '(so it does not include "sub-"). If this parameter is not '
                                                    'provided all participants should be analyzed. Multiple '
                                                    'participants can be specified with a space separated list.',
                        nargs="*", type=str)
    parser.add_argument('--aws_input_creds_path', help='The Path for credentials for reading from S3.'
                                                       ' If not provided and s3 paths are specified in the data config'
                                                       ' we will try to access the bucket anonymously'
                                                       ' use the string "env" to indicate that input credentials should'
                                                       ' read from the environment. (E.g. when using AWS iam roles).',
                        default=None)
    parser.add_argument('--aws_input_creds_profile', help='The profile name for credentials for writing to S3.'
                                                          ' If not provided and s3 paths are specified in the output '
                                                          ' directory we will try to access the bucket anonymously'
                                                          ' use the string "env" to indicate that output credentials '
                                                          ' should read from the environment. '
                                                          '(E.g. when using AWS iam roles).',
                        default=None)
    parser.add_argument('--aws_output_creds_path', help='The Path for credentials for writing to S3.'
                                                        ' If not provided and s3 paths are specified in the output'
                                                        ' directory we will try to access the bucket anonymously'
                                                        ' use the string "env" to indicate that output credentials'
                                                        ' should read from the environment. '
                                                        '(E.g. when using AWS iam roles).',
                        default=None)
    parser.add_argument('--aws_output_creds_profile', help='The profile name for credentials for writing to S3.'
                                                           ' If not provided and s3 paths are specified in the output'
                                                           ' directory we will try to access the bucket anonymously'
                                                           ' use the string "env" to indicate that output credentials'
                                                           ' should read from the environment. '
                                                           '(E.g. when using AWS iam roles).',
                        default=None)
    parser.add_argument('--n_cpus', type=int,
                        help='Number of execution '
                             ' resources available for the pipeline')
    parser.add_argument('--mem_mb', type=int,
                        help='Amount of RAM available to the pipeline in megabytes.'
                             ' Included for compatibility with BIDS-Apps standard, but mem_gb is preferred')
    parser.add_argument('--mem_gb', type=float,
                        help='Amount of RAM available to the pipeline in gigabytes.'
                             ' if this is specified along with mem_mb, this flag will take precedence.')
    parser.add_argument('--save_working_dir', action='store_true',
                        help='Save the contents of the working directory.')
    parser.add_argument('--disable_file_logging', action='store_true',
                        help='Disable file logging, this is useful for clusters that have disabled file locking.')
    parser.add_argument('--diagnostics', action='store_true',
                        help='Enable diagnostics dashboard of execution engine.')
    parser.add_argument('--enable_bids_validator',
                        help='skips bids validation',
                        action='store_true')
    parser.add_argument('--bids_validator_config', help='JSON file specifying configuration of '
                                                        'bids-validator: See https://github.com/INCF/bids-validator '
                                                        'for more info')
    parser.add_argument('-v', '--version', action='version',
                        version=f'Radiome version: {__version__}, email: {__email__}, author: {__author__}')
    return parser.parse_args(args)


def build_context(args) -> context.Context:
    mapping = {}
    print("Building the pipeline....")
    # Load config.
    if not os.path.isfile(args.config_file):
        raise FileNotFoundError(f"Can't find config file {args.config_file}!")
    with open(args.config_file, 'r') as f:
        mapping['pipeline_config'] = yaml.safe_load(f)
    print(f'Config file: {args.config_file}.')

    # Check the output directory.
    if args.outputs_dir.lower().startswith("s3://"):
        mapping['working_dir'] = args.working_dir or tempfile.mkdtemp(prefix='rdm')
        mapping['outputs_dir'] = S3Resource(args.outputs_dir, mapping['working_dir'], args.aws_output_creds_path,
                                            args.aws_output_creds_profile)
    else:
        mapping['working_dir'] = args.working_dir or os.path.join(args.outputs_dir, 'scratch')
        mapping['outputs_dir'] = os.path.abspath(args.outputs_dir)
        Path(mapping['outputs_dir']).mkdir(parents=True, exist_ok=True)

    mapping['working_dir'] = os.path.abspath(mapping['working_dir'])
    Path(mapping['working_dir']).mkdir(parents=True, exist_ok=True)
    print(f'Working directory: {mapping["working_dir"]}.')
    print(f'Output directory: {mapping["outputs_dir"]}.')

    # Check the input dataset.
    if args.bids_dir.lower().startswith("s3://"):
        mapping['inputs_dir'] = S3Resource(args.bids_dir, mapping['working_dir'], args.aws_input_creds_path,
                                           args.aws_input_creds_profile)
    else:
        if not os.path.exists(args.bids_dir):
            raise FileNotFoundError(f"Can't find {args.bids_dir}!")
        mapping['inputs_dir'] = os.path.abspath(args.bids_dir)
    print(f'Input directory: {mapping["inputs_dir"]}.')

    # Participant label
    mapping['participant_label'] = args.participant_label
    if mapping['participant_label'] is not None:
        print(f'Participants to process: {[f"sub-{label}" for label in args.participant_label]}')

    # Set up the logging
    log_format = '%(asctime)s-15s %(name)s %(levelname)s: %(message)s'
    if args.disable_file_logging:
        logging.basicConfig(stream=sys.stdout, level=logging.INFO, format=log_format)
    else:
        log_path = f'{mapping["working_dir"]}/{datetime.now().strftime("radiome_%Y_%m_%d_%H_%M.log")}'
        logging.basicConfig(
            level=logging.INFO,
            format=log_format,
            handlers=[logging.FileHandler(log_path)]
        )
        print(f'Logging at {log_path}')

    # Set up maximum memory allowed.
    if args.mem_gb:
        mapping['memory'] = int(args.mem_gb * 1024)
    elif args.mem_mb:
        mapping['memory'] = args.mem_mb
    else:
        mapping['memory'] = 6 * 1024
    print(f'Max Memory Available: {mapping["memory"]}MB')

    # Set up max core allowed
    mapping['n_cpus'] = min(args.n_cpus or 1, psutil.cpu_count())
    print(f'Max Cores Available: {mapping["n_cpus"]}')

    # BIDS Validation
    if args.enable_bids_validator and not args.bids_dir.lower().startswith('s3://'):
        if not shutil.which('bids-validator'):
            raise OSError('BIDS Validator is not correctly set up in your system!'
                          'Please refer to https://github.com/bids-standard/bids-validator'
                          'Command line version section for more information.')
        commands = ['bids-validator', f'--config {args.bids_validator_config}',
                    mapping['inputs_dir']] if args.bids_validator_config else ['bids-validator', mapping['inputs_dir']]
        print('Validating inputs again BIDS standard.......')
        completed_process = subprocess.run(commands, capture_output=True, universal_newlines=True)
        if completed_process.returncode != 0:
            raise ValueError(
                f'BIDS Validation failed. The error information is: {completed_process.stdout.splitlines()}')
        else:
            print('BIDS Validation passed. Continue')

    # flags
    mapping['save_working_dir'] = bool(args.save_working_dir)
    mapping['diagnostics'] = bool(args.diagnostics)

    return context.Context(**mapping)


def print_info() -> None:
    print('####  Running Radiome')
    print(f'Version: {__version__}')
    print(f"Start time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")


def main(args=None):
    if args is None:
        args = sys.argv[1:]
    params = parse_args(args)
    try:
        print_info()
        ctx = build_context(params)
        pipeline.build(ctx)
    except Exception as e:
        print(f'{type(e).__name__}:{e}', file=sys.stderr)
        return 1
    else:
        return 0


if __name__ == "__main__":
    sys.exit(main())
