#!/usr/bin/env python3
#
#  Michael Kuiper <kuiper.mike@gmail.com>
#

"""
\033[34;1m
mdwf (md workflow) \033[0m
mdwf is a python program used to help setup, run and monitor multiple molecular 
dynamics jobs run on a cluster based on a customisable directory structure. 

\033[22mUsage: mdwf (options) [<dir>] 
       mdwf -h |--help 
\033[34;1m
Options: \033[0m
  -h --help         Show this screen. 
  -m --monitor      Monitor Job status on cluster. 
  -i --initialize   Initialize jobs directories.
  -p --populate     Populate or update job directories with job scripts. 
  -c --checkjob     Checks job input before running simulations.
  -b --benchmark    Setup benchmarking runs for job optimization. 
  -s --start        Start all jobs.
  -g --gather       Collect all OutputFiles.
  -e --extend       Extend or shorten the runs of simulations. 
  --stop_jobs       Stops all running jobs immediately. 
  --pause           Pause all jobs jobs, (allows run to finish). 
  --restart         Restart all production jobs. 
  --clear_stops     Clear pausejob and stop flags. 
  --recover         Recover and cleanup stopped jobs. 
  --erase_all_data  Erases all job data. Use with caution!
  --clone           Clone directory structure without data. 
"""

import sys
import os 
from mdwf_lib import mdwf_functions as mdwf
from mdwf_lib import docopt

# ansi color variables for formatting purposes: 
DEFAULT = '\033[0m'
RED     = '\033[31m'
GREEN   = '\033[32m'
BLUE    = '\033[0;34m'

def main():
    args = docopt(__doc__, version=0.1)  # create argument dictionary
    arg_list = list(args.keys())   
    action = '--help'                    # default action 
    for i in arg_list: 
        if args[i] == True:
            print("{}mdwf action:{}{}".format(GREEN,DEFAULT,i))
            action = i 
    if action == '--help':
        print(__doc__)
    else:
        process(action)

def process(action):
    if action == '--monitor': 
        print("{}monitoring job status:{}".format( BLUE, DEFAULT ))
        mdwf.monitor_jobs()

    if action == '--initialize': 
        print("{}initializing job directories:{}".format( BLUE, DEFAULT ))
        mdwf.initialize_job_directories()

    if action == '--populate': 
        print("{}populating job directories:{}".format( BLUE, DEFAULT ))
        mdwf.populate_job_directories()

    if action == '--checkjob':
        print("{}checking job input:{}".format( BLUE, DEFAULT )) 
        mdwf.check_job()

    if action == '--benchmark':
        print("{}benchmarking job:{}".format( BLUE, DEFAULT )) 
        mdwf.benchmark()

    if action == '--start':
        print("{}starting jobs:{}".format( BLUE, DEFAULT )) 
        mdwf.start_all_jobs()

    if action == '--restart':
        print("{}restarting production jobs:{}".format( BLUE, DEFAULT )) 
        mdwf.restart_all_production_jobs()

    if action == '--gather':
        print("{}gathering data files: (list under /Analysis) {}".format( BLUE, DEFAULT )) 
        mdwf.gather_jobs()

    if action == '--extend':
        print("{}extend the runs: {}".format( BLUE, DEFAULT )) 
        try:
            a = int(eval(raw_input("Enter number of runs extentions: ")))
        except ValueError:
            print("Not an integer value...")
        mdwf.extend_jobs(a)

    if action == '--recover':
        print("{}recovering and restoring crashed jobs:{}".format( BLUE, DEFAULT )) 
        mdwf.recover_all_jobs()

    if action == '--clear_stops':
        print("{}clearing pausejob flags:{}".format( BLUE, DEFAULT )) 
        mdwf.clear_jobs()

    if action == '--stop_jobs':
        print("{}stopping all jobs:{}".format( BLUE, DEFAULT )) 
        mdwf.stop_jobs()

    if action == '--pause':
        print("{}pausing all jobs:{}".format( BLUE, DEFAULT )) 
        mdwf.pause_jobs()

    if action == '--erase_all_data':
        print("{}erasing data:{}".format( BLUE, DEFAULT ))
        mdwf.erase_all_data()

    if action == '--clone':
        print("{}cloning job_directory: (not implimented yet){}".format( BLUE, DEFAULT )) 
        mdwf.clone()

if __name__ == '__main__':

    main()
