import os
import subprocess
from A3PI_CubitTemplate import CubitTemplate

def run_cubit(workflow, step_arg):
    
    #Load all necessary Cubit data from workflow input
    A3PI_mode = workflow.config.get('RUN_PARAMETERS', 'A3PI_mode')
    path = workflow.config.get('PATHS', 'cubit_path')
    file = workflow.config.get('CUBIT_ELEMENT_'+step_arg, 'file')
    keys = workflow.config.get('CUBIT_ELEMENT_'+step_arg, 'keys')
    values = workflow.config.get('CUBIT_ELEMENT_'+step_arg, 'values')
    output_file = workflow.config.get('CUBIT_ELEMENT_'+step_arg, 'output_file')
    
    if not isinstance(keys,list):
        keys = [keys]
        values = [values]
    
    #Update values in Cubit journal and write file to workflow folder
    cubit_temp = CubitTemplate(file)
    cubit_temp.set_value(dict(zip(keys, values)))
    cubit_temp.set_export(output_file)
    cubit_temp.write_file(workflow.folder + '/' + file)
    
    if A3PI_mode == 'single':
        subprocess.call('srun -n 1 ' + path + ' -nog ' + file, 
                        shell=True, cwd=workflow.folder)
    elif A3PI_mode == 'optimize':
        from libensemble.executors.executor import Executor
        exctr = Executor.executor
        os.chdir(workflow.folder)
        app_arg = ' -nog ' + file
        task = exctr.submit(app_name='cubit', num_procs=1,
                            app_args=app_arg, wait_on_run=True)
        task.wait(cubit_wait)
        if task.success is False:
                task.kill()
        os.chdir('..')