#!/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 glob
import shutil
import slmr

DEBUG = True
#  DEBUG = False


def main():
    parser = argparse.ArgumentParser(
            description='''
========================================================================================
 sLMR: a scripting system for LMR (Feng Zhu, fengzhu@usc.edu)
----------------------------------------------------------------------------------------

 Usage example of running LMR with Slurm:

    slrm job -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__)
            )

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

    parser_job = subparsers.add_parser('job', help='SLURM job submitting')
    parser_pp = subparsers.add_parser('pp', help='post-processing')
    pp_subparsers = parser_pp.add_subparsers(help='post-processing tasks')
    pp_subparsers.dest = 'task'

    parser_pp_pick_files = pp_subparsers.add_parser('pick_files', help='pick files')

    #===================================================================================
    # SLURM job submitting
    #===================================================================================
    parser_job.add_argument(
            '-c',
            '--config',
            required=True, help='path of the config YAML file')

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

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

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

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

    parser_job.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_job.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_job.add_argument(
            '-nn',
            '--nodename',
            help='node name; if not set, will be chosen by Slurm')

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

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

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


    #===================================================================================
    # post-processing
    #===================================================================================
    parser_pp_pick_files.add_argument(
            '-f',
            '--filename',
            required=True,
            help='pick files with the specified filename')

    parser_pp_pick_files.add_argument(
            '-d',
            '--directories', nargs='*',
            required=True,
            help='pick files from specified directories')

    parser_pp_pick_files.add_argument(
            '-s',
            '--saveto',
            required=True,
            help='the destination directory to store the found files')

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

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

    if args.mode == 'job':
        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.nthreads is not None:
            cfg_dict['core']['nthreads'] = int(args.nthreads)

        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['wrapper']['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)

    elif args.mode == 'pp':
        if args.task == 'pick_files':
            pick_files(args.filename, args.directories, args.saveto)

#===================================================================================
# SLURM job submitting
#===================================================================================
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.nthreads is not None:
        sbatch.write("""
#SBATCH --ntasks="""+args.nthreads+""" """)
    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')


#===================================================================================
# post-processing
#===================================================================================
def pick_files(filename, directories, saveto):
    if not os.path.exists(saveto):
        os.makedirs(saveto)

    nfile = 0
    for directory in directories:
        path = os.path.join(directory, './*/', filename)
        res = glob.glob(path)
        for f in res:
            nfile += 1
            save_filename = '{}_{:04d}'.format(filename, nfile)
            save_path = os.path.join(saveto, save_filename)
            shutil.copy(f, save_path)
            print('sLMR >>> {} -> {}'.format(f, save_path))

    print('sLMR >>> Found {} files named as: {} '.format(nfile, filename))
    print('sLMR >>> Saving them to: {} >>> DONE'.format(saveto))


if __name__ == '__main__':
    main()
