#!/usr/bin/env python3

# ==============================================================================
#  Author: Feng Zhu (fengzhu@ucar.edu)
#  Date: 2023-06-17
# ==============================================================================

import argparse
import subprocess
import os
import cfr

def main():
    parser = argparse.ArgumentParser(
            description='''
========================================================================================
 cfr: a scripting system for CFR
----------------------------------------------------------------------------------------
 Usage example for DA:
    cfr da -c config.yml -vb -s 1 2 -r
    # -c config.yml: run the reconstruction job according to config.yml
    # -vb: output the verbose runtime information
    # -s 1 2: set seeds as integers from 1 to 2
    # -r: run the Monte-Carlo iterations for PDA

 Usage example for GraphEM:
    cfr graphem -c config.yml -vb
    # -c config.yml: run the reconstruction job according to config.yml
    # -vb: output the verbose runtime information
========================================================================================
            ''', 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'

    # DA
    parser_da = subparsers.add_parser('da', help='run a DA-based reconstruction')

    parser_da.add_argument(
        '-c',
        '--config',
        required=True, help='path of the config YAML file')

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

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

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

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

    # GraphEM
    parser_graphem = subparsers.add_parser('graphem', help='run a GraphEM-based reconstruction')

    parser_graphem.add_argument(
        '-c',
        '--config',
        required=True, help='path of the config YAML file')

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

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

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

    job = cfr.ReconJob()
    if args.mode == 'da':
        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')

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

    elif args.mode == 'graphem':
        job.run_graphem_cfg(args.config, verbose=args.verbose)

        # 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')

        #     job = cfr.ReconJob()
        #     job.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()