#!/usr/bin/env python3

# ==============================================================================
#  Author: Feng Zhu (fzhu@nuist.edu.cn)
#  Date: 2022-05-22
# ==============================================================================

import argparse
import subprocess
import os
import cfr

def main():
    parser = argparse.ArgumentParser(
            description='''
========================================================================================
 cfr: a scripting system for CFR (Feng Zhu, fzhu@nuist.edu.cn)
----------------------------------------------------------------------------------------
 Usage example:
    cfr reconjob -c config.yml -vb True
    # -c config.yml: run the reconstruction job according to config.yml
    # -vb True: output the verbose runtime information
    # -s 0 9: set seeds as integers from 0 to 9
========================================================================================
            ''', formatter_class=argparse.RawTextHelpFormatter)

    parser.add_argument(
            '-v',
            '--version',
            action='version',
            version='%(prog)s version: {}'.format(cfr.__version__)
            )

    subparsers = parser.add_subparsers(help='running mode')
    subparsers.dest = 'mode'

    parser_reconjob = subparsers.add_parser('reconjob', help='reconstruction job')
    parser_reconjob.add_argument(
        '-c',
        '--config',
        required=True, help='path of the config YAML file')

    parser_reconjob.add_argument(
        '-vb',
        '--verbose',
        action=argparse.BooleanOptionalAction,
        help='output the verbose runtime information')

    parser_reconjob.add_argument(
        '-s',
        '--seeds',
        nargs='*',
        default=None,
        help='the start and end of the random seeds for reconstruction')

    parser_reconjob.add_argument(
        '-b',
        '--bsub',
        action=argparse.BooleanOptionalAction,
        help='submit a BSUB job')

    parser_reconjob.add_argument(
        '-r',
        '--run',
        action=argparse.BooleanOptionalAction,
        help='prepare the job without running Monte-Carlo')

    #===================================================================================
    # Executing
    #===================================================================================

    # parse the input command line
    args = parser.parse_args()

    if args.mode == 'reconjob':

        if args.bsub:
            cfg_path = os.path.abspath(args.config)
            work_dir = os.path.dirname(cfg_path)
            bsub_path = make_bsub(work_dir, cfg_path, args.seeds, args.run, args.verbose)
            run_bsub(bsub_path)
        else:
            if args.seeds is None:
                seeds = None
            elif len(args.seeds) == 1:
                seeds = int(args.seeds[0])
            elif len(args.seeds) == 2:
                seeds = list(range(int(args.seeds[0]), int(args.seeds[-1])+1))
            else:
                raise ValueError('Wrong number of seeds')

            cfr.run_da_cfg(args.config, seeds=seeds, run_mc=args.run, verbose=args.verbose)

def make_bsub(work_dir, cfg_path, seeds, run, verbose):
    bsub_path = os.path.join(work_dir, 'cfr.bsub')
    with open(bsub_path, 'w') as f:
        f.write('''#!/bin/bash
#BSUB -J CFR
#BSUB -n 1
#BSUB -o %J.out
        ''')

        spell = f'cfr reconjob -c {cfg_path}'
        if seeds is not None:
            if len(seeds) == 1:
                spell += f' -s {int(seeds[0])}'
            elif len(seeds) == 2:
                spell += f' -s {int(seeds[0])} {int(seeds[-1])}'
        if verbose: spell += ' -vb'
        if run: spell += ' -r'
        f.write(f'''
{spell}
        ''')
    
    return bsub_path


def run_bsub(bsub_path):
    cmd = f'bsub < {bsub_path}'
    subprocess.call(cmd, shell=True)

if __name__ == '__main__':
    main()