import os
import shutil
import subprocess
from A3PI_AcdtoolTemplate import AcdtoolTemplate

def run_meshconvert(workflow, step_arg):
    
    #Load all necessary acdtool data from workflow input
    A3PI_mode = workflow.config.get('RUN_PARAMETERS', 'A3PI_mode')
    path = workflow.config.get('PATHS', 'acdtool_path')
    file = workflow.config.get('CUBIT_ELEMENT_'+step_arg, 'output_file')
    
    if A3PI_mode == 'single':
        subprocess.call('srun -n 1 ' + path + ' meshconvert ' + file, 
                        shell=True, cwd=workflow.folder)
    elif A3PI_mode == 'worker':
        from libensemble.executors.executor import Executor
        exctr = Executor.executor
        os.chdir(workflow.folder)
        app_arg = ' meshconvert ' + file
        task = exctr.submit(app_name='acdtool', num_procs=1,
                            app_args=app_arg, wait_on_run=True)
        task.wait()
        if task.success is False:
            task.kill()
        os.chdir('..')

def run_rfpost(workflow, step_arg):
    
    #Load all necessary acdtool data from workflow input
    A3PI_mode = workflow.config.get('RUN_PARAMETERS', 'A3PI_mode')
    cores = workflow.config.get('RUN_PARAMETERS', 'acdtool_cores')
    path = workflow.config.get('PATHS', 'acdtool_path')
    file = workflow.config.get('RFPOST_MODE_'+step_arg, 'file')
    
    rfpost_mode = 'RFPOST_MODE_'+step_arg
    rfpost_temp = make_rfpost_template(workflow, rfpost_mode)
    rfpost_temp.write_file(workflow.folder + '/' + file)
    
    if A3PI_mode == 'single':
        subprocess.call('srun -n ' + cores + ' --cpu-bind=cores ' + path 
                        + ' postprocess rf ' + file,
                        shell=True, cwd=workflow.folder)
    elif A3PI_mode == 'worker':
        from libensemble.executors.executor import Executor
        exctr = Executor.executor
        os.chdir(workflow.folder)
        app_arg = ' postprocess rf ' + file
        task = exctr.submit(app_name='acdtool', num_procs=eval(cores),
                            app_args=app_arg, wait_on_run=True,
                            extra_args='--cpu-bind=cores')
        task.wait()
        if task.success is False:
            task.kill()
        os.chdir('..')
    
    #Find file ID for given mode and rename acdtool outputs to data_.h5 files
    mode_fid = int(workflow.config.get('OMEGA3P_MODE_'+step_arg, 'mode_fid'))
    shutil.move(workflow.folder + '/E_Real.h5', 
                workflow.folder + '/data' + str(mode_fid) + '.h5')
    shutil.move(workflow.folder + '/E_Imag.h5', 
                workflow.folder + '/data' + str(mode_fid+1) + '.h5')
    shutil.move(workflow.folder + '/B_Real.h5', 
                workflow.folder + '/data' + str(mode_fid+2) + '.h5')
    shutil.move(workflow.folder + '/B_Imag.h5', 
                workflow.folder + '/data' + str(mode_fid+3) + '.h5')
    
def make_rfpost_template(workflow, rf_mode):
    
    rfpost_temp = AcdtoolTemplate()
    
    #Set omega3p output filename
    results = workflow.config.get(rf_mode, 'RFField_ResultDir')
    rfpost_temp.set_result_file(results)
    
    #Set frequency scan ID and mode ID
    if workflow.config.has_option(rf_mode, 'RFField_FreqScanID'):
        freq_scan = workflow.config.get(rf_mode, 'RFField_FreqScanID')
    else:
        freq_scan = '0'
    if workflow.config.has_option(rf_mode, 'RFField_ModeID'):
        mode_id = workflow.config.get(rf_mode, 'RFField_ModeID')
    else:
        mode_id = '0'
    rfpost_temp.set_mode_info(freq_scan, mode_id)
    
    #Set mesh symmetry options
    if workflow.config.has_option(rf_mode, 'RFField_xsymmetry'):
        xsym = workflow.config.get(rf_mode, 'RFField_xsymmetry')
    else:
        xsym = 'none'
    if workflow.config.has_option(rf_mode, 'RFField_ysymmetry'):
        ysym = workflow.config.get(rf_mode, 'RFField_ysymmetry')
    else:
        ysym = 'none'
    rfpost_temp.set_symmetry(xsym, ysym)
    
    if workflow.config.has_option(rf_mode, 'RFField_gradient'):
        gradient = workflow.config.get(rf_mode, 'RFField_gradient')
        cavityBeta = workflow.config.get(rf_mode, 'RFField_cavityBeta')
        revPflow = workflow.config.get(rf_mode, 'RFField_reversePowerFlow')
        x0 = workflow.config.get(rf_mode, 'RFField_x0')
        y0 = workflow.config.get(rf_mode, 'RFField_y0')
        gz1 = workflow.config.get(rf_mode, 'RFField_gz1')
        gz2 = workflow.config.get(rf_mode, 'RFField_gz2')
        npoint = workflow.config.get(rf_mode, 'RFField_npoint')
        fmnx = workflow.config.get(rf_mode, 'RFField_fmnx')
        fmny = workflow.config.get(rf_mode, 'RFField_fmny')
        fmnz = workflow.config.get(rf_mode, 'RFField_fmnz')
        rfpost_temp.set_gradient(x0, y0, gz1, gz2, npoint, fmnx, fmny, fmnz, 
                         gradient, cavityBeta, revPflow)
    
    if workflow.config.has_option(rf_mode, 'OpenPMD_IMPACT_ionoff'):
        ionoff = workflow.config.get(rf_mode, 'OpenPMD_IMPACT_ionoff')
        nx = workflow.config.get(rf_mode, 'OpenPMD_IMPACT_nx')
        ny = workflow.config.get(rf_mode, 'OpenPMD_IMPACT_ny')
        nz = workflow.config.get(rf_mode, 'OpenPMD_IMPACT_nz')
        x1 = workflow.config.get(rf_mode, 'OpenPMD_IMPACT_x1')
        x2 = workflow.config.get(rf_mode, 'OpenPMD_IMPACT_x2')
        y1 = workflow.config.get(rf_mode, 'OpenPMD_IMPACT_y1')
        y2 = workflow.config.get(rf_mode, 'OpenPMD_IMPACT_y2')
        z1 = workflow.config.get(rf_mode, 'OpenPMD_IMPACT_z1')
        z2 = workflow.config.get(rf_mode, 'OpenPMD_IMPACT_z2')
        rfpost_temp.set_openpmd(nx, ny, nz, x1, x2, y1, y2, z1, z2, ionoff)
    
    return rfpost_temp