import os
import subprocess
from A3PI_Omega3PTemplate import Omega3PTemplate

def run_omega3p(workflow, step_arg):
    
    #Load Omega3P run parameters and path
    A3PI_mode = workflow.config.get('RUN_PARAMETERS', 'A3PI_mode')
    path = workflow.config.get('PATHS', 'omega3p_path')
    omega3p_mode = 'OMEGA3P_MODE_'+step_arg
    omega3p_temp = make_omega3p_template(workflow, omega3p_mode)
    cores = workflow.config.get(omega3p_mode, 'cores')
    
    #Set input filename for Omega3P and write to file
    file = workflow.config.get(omega3p_mode, 'file')
    omega3p_temp.write_file(workflow.folder + '/' + file)
    
    #Run Omega3P in workflow folder with created input file
    if A3PI_mode == 'single':
        if workflow.config.has_option(omega3p_mode, 'cores'):
            cores = workflow.config.get(omega3p_mode, 'cores')
            if workflow.config.has_option(omega3p_mode, 'tasks'):
                tasks = workflow.config.get(omega3p_mode, 'tasks')
            else:
                tasks = cores   #Added for backwards compatibility
                cores = 1
        elif workflow.config.has_option(omega3p_mode, 'tasks'):
            tasks = workflow.config.get(omega3p_mode, 'tasks')
            cores = 1
        else:
            tasks = 1
            cores = 1
    
        subprocess.call('srun -n ' + tasks + ' -c ' + cores + ' --cpu-bind=cores ' 
                        + path + ' ' + file, shell=True, cwd=workflow.folder)
    elif A3PI_mode == 'worker':
        from libensemble.executors.executor import Executor
        exctr = Executor.executor
        os.chdir(workflow.folder)
        task = exctr.submit(app_name='omega3p', num_procs=eval(cores),
                            app_args=file, wait_on_run=True,
                            extra_args='--cpu-bind=cores')
        task.wait()
        if task.success is False:
            task.kill()
            cores = str(int(eval(cores)/2))
            task = exctr.submit(app_name='omega3p', num_procs=eval(cores),
                                app_args=file, wait_on_run=True,
                                extra_args='--cpu-bind=cores')
            task.wait()
            if task.success is False:
                task.kill()
        os.chdir('..')

def make_omega3p_template(workflow, omega3p_mode):
    
    #Create blank Omega3P template
    omega3p_temp = Omega3PTemplate()
    
    #Set model file, if linked to a cubit mesh, use its output with .ncdf suffix
    model_file = workflow.config.get(omega3p_mode, 'ModelInfo_File')
    if model_file.startswith('cubit_element_'):
        model_ind = model_file.strip('cubit_element_')
        model_file = workflow.config.get('CUBIT_ELEMENT_'+model_ind,
                                         'output_file').strip('.gen') + '.ncdf'
    if workflow.config.has_option(omega3p_mode, 'ModelInfo_Tolerant'):
        tolerant = workflow.config.get(omega3p_mode, 'ModelInfo_Tolerant')
    else:
        tolerant = '1'
    omega3p_temp.set_model_file(model_file, tolerant)
    
    #Add boundary condition flags for corresponding reference surfaces
    if workflow.config.has_option(omega3p_mode, 'ModelInfo_BoundaryCondition'):
        bcs = workflow.config.get(omega3p_mode, 'ModelInfo_BoundaryCondition')
        bcs = bcs.split(',')
        for bc in bcs:
            bc_type = bc.strip().split(' ')[0].capitalize()
            bc_refs = bc.strip().split(' ')[1::]
            omega3p_temp.add_boundary_condition(bc_type, bc_refs)
    
    #Add surface material properties for each reference surface
    if workflow.config.has_option(omega3p_mode, 'SurfaceMaterial'):
        sms = workflow.config.get(omega3p_mode, 'SurfaceMaterial').split(',')
        for sm in sms:
            [sm_ref, sm_sigma] = sm.strip().split(' ')
            omega3p_temp.add_surface_material(sm_ref, sm_sigma)
    
    #Add material properties for each reference volume
    if workflow.config.has_option(omega3p_mode, 'Material'):
        mats = workflow.config.get(omega3p_mode, 'Material').split(',')
        for mat in mats:
            mat_ref = mats.strip().split(' ')[0]
            mat_eps = complex(mats.strip().split(' ')[1])
            mat_mu = complex(mats.strip().split(' ')[2])
            omega3p_temp.add_surface_material(mat_ref, mat_eps, mat_mu)
    
    #Set finite element order and curved surfaces options if specified
    fe_order = workflow.config.get(omega3p_mode, 'FiniteElement_Order')
    if workflow.config.has_option(omega3p_mode, 'FiniteElement_CurvedSurfaces'):
        cur = workflow.config.get(omega3p_mode, 'FiniteElement_CurvedSurfaces')
    else:
        cur = 'on'    #Default curved surfaces setting
    omega3p_temp.set_finite_element(fe_order, cur)
    
    #Add PRegions to model for using varying finite element orders
    if workflow.config.has_option(omega3p_mode, 'PRegion'):
        prs = workflow.config.get(omega3p_mode, 'PRegion').split(',')
        for pr in prs:
            pr_ref = pr.strip().split(' ')[0]
            pr_order = pr.strip().split(' ')[0]
            omega3p_temp.add_pregion(pr_ref, pr_order)
    
    #Set eigenvalue parameters and optionally tolerance and max iterations
    num_eigs = workflow.config.get(omega3p_mode, 'EigenSolver_NumEigenvalues')
    freq_shift = workflow.config.get(omega3p_mode, 'EigenSolver_FrequencyShift')
    if workflow.config.has_option(omega3p_mode, 'EigenSolver_Tolerance'):
        tol = workflow.config.get(omega3p_mode, 'EigenSolver_Tolerance')
    else:
        tol = '1e-7'    #Default eigensolver tolerance
    if workflow.config.has_option(omega3p_mode, 'EigenSolver_MaxIterations'):
        miters = workflow.config.get(omega3p_mode, 'EigenSolver_MaxIterations')
    else:
        miters = '200'    #Default maximum iterations
    omega3p_temp.set_eigensolver(num_eigs, freq_shift, tol, miters)
    
    return omega3p_temp