#!/usr/bin/env python3

# ==============================================================================
#  Author: Feng Zhu (fengzhu@usc.edu)
#  Date: 2018-07-03 15:20:58
# ==============================================================================

import argparse
import yaml
import os
import subprocess
import slmr

#  DEBUG = True
DEBUG = False


def main():
    parser = argparse.ArgumentParser(
            description='''
========================================================================================
 sLMR: run LMR with Slurm (Feng Zhu, fengzhu@usc.edu)
----------------------------------------------------------------------------------------

 Usage example:

    slrm -c config.yml -n 4 -nn hungus -rp 0 2000 -em slmr@gmail.com -x test_ccsm4

    # -c config.yml: using "config.yml" as a configuration template
    # -n 4 -nn hungus: run LRM with 4 threads on the node "hungus"
    # -rp 0 2000: reconstruction period to be 0-2000
    # -em slmr@gmail.com: notification will be sent to slmr@gmail.com
    # -x test_ccsm4: the experiment is named as "test_ccsm4"

========================================================================================
            ''', formatter_class=argparse.RawTextHelpFormatter)

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

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

    # for LMR
    parser.add_argument(
            '-x',
            '--exp',
            help='name of the experiment')

    parser.add_argument(
            '-ne',
            '--nens',
            help='number of ensemble members; if not set, will use the setting in the YAML file')

    parser.add_argument(
            '-n',
            '--nthread',
            help='number of threads; if not set, will use the setting in the YAML file')

    parser.add_argument(
            '-s',
            '--seed',
            help='random seed number; if not set, will use the setting in the YAML file')

    parser.add_argument(
            '-rp',
            '--recon_period', nargs='*',
            help=''' time span for reconstruction in the form of start end,
            e.g., "-rp 0 2000" (if more than two numbers are given, only the first two are valid);
            if not set, will use the setting in the YAML file
            ''')

    parser.add_argument(
            '-ir',
            '--iter_range', nargs='*',
            help=''' Monte-Carlo iterations
            e.g., "-ir 0 10" (if more than two numbers are given, only the first two are valid);
            if not set, will use the setting in the YAML file
            ''')

    # for Slurm
    parser.add_argument(
            '-nn',
            '--nodename',
            help='node name; if not set, will be chosen by Slurm')

    parser.add_argument(
            '-em',
            '--email',
            help='email address; if not set, notification will be diabled')

    parser.add_argument(
            '-jn',
            '--jobname',
            help='name of the Slurm job')

    parser.add_argument(
            '-hr',
            '--hour',
            help='running hours; if not set, use 72 hours')

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

    cfg_path = os.path.abspath(args.config)
    cfg_dict = yaml.load(open(cfg_path))
    work_dir = os.path.dirname(cfg_path)

    # update cfg with args
    if args.nthread is not None:
        cfg_dict['core']['nthread'] = int(args.nthread)

    if args.exp is not None:
        cfg_dict['core']['nexp'] = args.exp

    if args.nens is not None:
        cfg_dict['core']['nens'] = int(args.nens)

    if args.seed is not None:
        cfg_dict['core']['seed'] = int(args.seed)

    if args.recon_period is not None:
        cfg_dict['core']['recon_period'] = (int(args.recon_period[0]), int(args.recon_period[1]))

    if args.iter_range is not None:
        cfg_dict['core']['iter_range'] = (int(args.iter_range[0]), int(args.iter_range[1]))


    yml_path = make_yaml(work_dir, cfg_dict)

    job_path = make_sbatch(work_dir, cfg_dict, args, yml_path)

    if not DEBUG:
        run_job(job_path)


def make_yaml(work_dir, cfg_dict):
    ''' output a new YAML file
    '''
    nexp = cfg_dict['core']['nexp']
    yml_path = os.path.join(work_dir, 'config_{}.yml'.format(nexp))

    with open(yml_path, 'w') as f:
        yaml.dump(cfg_dict, f, default_flow_style=False)

    print('sLMR >>> Configuration created! >>>', yml_path)
    return yml_path


def make_sbatch(work_dir, cfg_dict, args, yml_path):
    ''' output a shell script for sbatch
    '''
    nexp = cfg_dict['core']['nexp']

    job_path = os.path.join(work_dir, 'run_{}.sh'.format(nexp))

    # create a new file
    with open(job_path, 'w') as f:
        f.write('')

    sbatch = open(job_path, 'a')
    sbatch.write("""#!/bin/bash
#SBATCH --nodes=1
#SBATCH --mem=30000 """)

    if args.hour is not None:
        sbatch.write("""
#SBATCH --time="""+args.hour+""":00:00 """)
    else:
        sbatch.write("""
#SBATCH --time=72:00:00 """)

    if args.jobname is not None:
        sbatch.write("""
#SBATCH --job-name="""+args.jobname+""" """)
    else:
        sbatch.write("""
#SBATCH --job-name=LMR """)

    if args.nthread is not None:
        sbatch.write("""
#SBATCH --ntasks="""+args.nthread+""" """)
    else:
        sbatch.write("""
#SBATCH --ntasks=1""")

    if args.email is not None:
        sbatch.write("""
#SBATCH --mail-user="""+args.email+"""
#SBATCH --mail-type=ALL """)

    if args.nodename is not None:
        sbatch.write("""
#SBATCH --nodelist="""+args.nodename+""" """)

    sbatch.write("""

python -u LMR_wrapper.py """+yml_path+""" &> """+nexp+""".log """)

    sbatch.close()

    print('sLMR >>> Slurm job created! >>>', job_path)

    return job_path


def run_job(job_path):
    cmd = 'sbatch '+job_path
    subprocess.call(cmd, shell=True)
    print('sLMR >>> Job submitted! >>> DONE')

if __name__ == '__main__':
    main()
