#!/usr/bin/env python

from radical.entk import AppManager
from repex.Sync import SynchronousExchange
import os, sys, time, pprint, json, pickle
import radical.utils as ru
import radical.analytics as ra
import radical.entk as re

# ------------------------------------------------------------------------------
# Set default verbosity


os.environ['RADICAL_SAGA_VERBOSE']         = 'INFO'
os.environ['RP_ENABLE_OLD_DEFINES']        = 'True'
os.environ['RADICAL_PROFILE']              = 'True'
os.environ['RADICAL_ENTK_PROFILE']         = 'True'
os.environ['RADICAL_ENTK_VERBOSE']         = 'INFO'
os.environ['RP_ENABLE_OLD_DEFINES']        = 'True'
os.environ['SAGA_PTY_SSH_TIMEOUT']         = '2000'
os.environ['RADICAL_VERBOSE']              = 'INFO'
os.environ['RADICAL_PILOT_PROFILE']        = 'True'
os.environ['RADICAL_REPEX_SYNCEX_PROFILE'] = 'True'
os.environ['RADICAL_REPEX_RUN_PROFILE']    = 'True'
os.environ['RADICAL_PILOT_DBURL'] = "mongodb://smush:key1209@ds147361.mlab.com:47361/db_repex_4"


simconfig = sys.argv[1]
resconfig = sys.argv[2]

with open(simconfig) as simconf:
    simdata = json.load(simconf)

with open(resconfig) as resconf:
    resdata = json.load(resconf)

replicas       = simdata["replicas"]
replica_cores  = simdata["replica_cores"]
cycles         = simdata["cycles"]    #0 cycles = no exchange
exchange_method = simdata["exchange_method"]  #'exchangeMethods/TempEx.py' #/path/to/your/exchange/method
md_executable  = resdata["md_executable"]  #'/u/sciteam/mushnoor/amber/amber14/bin/sander.MPI' #/path/to/your/MD/Executable
python_path    = resdata["py_executable"]
timesteps      = simdata["timesteps"] #100 #Number of timesteps between exchanges
basename       = simdata["basename"]     #'ace-ala'
min_temp       = simdata["min_temp"]
max_temp       = simdata["max_temp"]
pilot_cores    = resdata["cpus"]  
pre_exec       = resdata["pre_exec"]

if __name__ == '__main__':
    
    #if resdata["resource"] is "local.localhost":


    try:
        res_dict ={
                    "resource"      : str(resdata["resource"]),
                    "walltime"      : int(resdata["walltime"]),
                    "cpus"          : pilot_cores,
                    "gpus_per_node" : int(resdata["gpus_per_node"]),
                    "access_schema" : str(resdata["access_schema"]),
                    "queue"         : str(resdata["queue"]),
                    "project"       : str(resdata["project"]),  
                  }

    except:

        print resdata["resource"]
        res_dict ={
                    "resource"      : str(resdata["resource"]),
                    "walltime"      : int(resdata["walltime"]),
                    "cpus"          : pilot_cores,
 
                  }

    uid1 = ru.generate_id('radical.repex.run')
    logger = ru.get_logger('radical.repex.run')
    prof = ru.Profiler(name=uid1)
    prof.prof('Create_Workflow_0', uid=uid1)

      
    synchronous_exchange=SynchronousExchange()
    

    appman                  = AppManager(autoterminate=False, port=32769)  # Create Application Manager
    appman.resource_desc = res_dict # Assign resource manager to the Application Manager      
   
    exchange                = synchronous_exchange.init_cycle(replicas, replica_cores, python_path, md_executable, exchange_method, min_temp, max_temp, timesteps, basename, pre_exec)
    
    appman.workflow = set([exchange]) # Assign the workflow as a set of Pipelines to the Application Manager 

    prof.prof('Run_Cycle_0', uid=uid1)

    appman.run() # Run the Application Manager 

    prof.prof('End_Cycle_0', uid=uid1)
 
    
    for cycle in range (cycles):

        prof.prof('Create_Workflow_{0}'.format(cycle+1), uid=uid1)

                          
        exchange_gen            = synchronous_exchange.general_cycle(replicas, replica_cores, cycle, python_path, md_executable, exchange_method, pre_exec)
        
        appman.workflow = set([exchange_gen]) # Assign the workflow as a set of Pipelines to the Application Manager       

        prof.prof('Run_Cycle_{0}'.format(cycle+1), uid=uid1)

                            

        appman.run() # Run the Application Manager

        prof.prof('End_Cycle_{0}'.format(cycle+1), uid=uid1)

    appman.resource_terminate()

    mdtasks  = synchronous_exchange.mdtasklist
    extasks  = synchronous_exchange.extasklist

    pwd = os.getcwd()
    session = ra.Session(sid   = './%s'%appman.sid,
                         stype = 'radical.entk',
                         src   = pwd)


    mdtask_uid_map = dict()
    for task in mdtasks:
        mdtask_uid_map[task.name] = task.uid
        #print task.name

    extask_uid_map = dict()
    for task in extasks:
        extask_uid_map[task.name] = task.uid

        
    def get_mdtask_uids(task_names_list):
        return [mdtask_uid_map[task.name] for task in task_names_list]
    def get_extask_uids(task_names_list):
        return [extask_uid_map[task.name] for task in task_names_list]

    #Write MD and EX task lists to files

    with open('MDLIST_%s'%appman.sid, 'wb') as mdlist:
        pickle.dump(get_mdtask_uids(mdtasks), mdlist)

    with open('EXLIST_%s'%appman.sid, 'wb') as exlist:
        pickle.dump(get_extask_uids(extasks), exlist)
                    
 
    with open('MDLIST_%s'%appman.sid, 'rb') as mdlist:
        md_task_list=pickle.load(mdlist)

    with open('EXLIST_%s'%appman.sid, 'rb') as exlist:
        ex_task_list=pickle.load(exlist)
                

    md                      = session.filter(etype='task', inplace=False,  uid=md_task_list)
    md_scheduling_durations = md.duration([re.states.SCHEDULED, re.states.SUBMITTED])
    md_dequeuing_durations = md.duration([re.states.COMPLETED, re.states.DONE])
    md_durations            = md.duration([re.states.SUBMITTED, re.states.COMPLETED])

    ex                      = session.filter(etype='task', inplace=False,  uid=ex_task_list)
    ex_scheduling_durations = ex.duration([re.states.SCHEDULED, re.states.SUBMITTED])
    ex_dequeuing_durations = ex.duration([re.states.COMPLETED, re.states.DONE])
    ex_durations            = ex.duration([re.states.SUBMITTED, re.states.COMPLETED])

    total                   = session.filter(etype='task', inplace=False)
    total_durations         = total.duration([re.states.SCHEDULED, re.states.DONE])

    print "Total MD duration is ", md_durations
    print "MD Scheduling duration is ", md_scheduling_durations
    print "MD Dequeuing  duration is ", md_dequeuing_durations
    print "Total EX duration is ", ex_durations
    print "EX Scheduling duration is ", ex_scheduling_durations
    print "EX Dequeuing  duration is ", ex_dequeuing_durations
    print "total duration is " , total_durations

    
