pvrun

#!/usr/bin/env python3

import sys
import os

from paravision.utils import read_files

import importlib
import subprocess
import argparse
from addict import Dict
from pathlib import Path

from rich import print, print_json

def parse_outer_args():
    """
    parser for the main script. 
    """

    ap = argparse.ArgumentParser()

    ap.add_argument('cmd', help='command')
    ap.add_argument("-np", "--nproc", type=int, default=1, help="Screenshot the given object")
    ap.add_argument("-lsm", "--load-scientific-colormaps", action='store_true', help="Load all the ScientificColourMaps7 presets into ParaView before starting.")

    args, unknown =  ap.parse_known_args()

    args = Dict(vars(args))

    if args['load_scientific_colormaps']: 
        load_scientific_colormaps()

    return args, unknown

def run_module(module:str, args:list, nproc:int=1): 
    command = importlib.util.find_spec(module).origin
    run_command(command, args, nproc)

def run_command(commandstr:str, args:list, nproc:int=1):
    ## Needed because subprocess.run couldn't find modules
    myenv = os.environ
    pythonpath = os.pathsep.join(sys.path)
    myenv.update({'PYTHONPATH': pythonpath})

    if nproc > 1:
        command_pre = [ 'mpiexec', '-np', str(nproc), 'pvbatch' ]
        pvbatch_path = subprocess.run(['which', 'pvbatch'], capture_output=True, env=myenv)
        print(f"Using pvbatch from path: {pvbatch_path.stdout.decode('ascii').strip()}")
        mpiexec_path = subprocess.run(['which', 'mpiexec'], capture_output=True, env=myenv)
        print(f"Using mpiexec from path: {mpiexec_path.stdout.decode('ascii').strip()}")
    else:
        command_pre = [ 'pvpython' ]
        pvpython_path = subprocess.run(['which', 'pvpython'], capture_output=True, env=myenv)
        print(f"Using pvpython from path: {pvpython_path.stdout.decode('ascii').strip()}")

    command = commandstr.split()
    command.extend(args)

    command = command_pre + command

    print(command)
    subprocess.run(command)


def main():

    args, subargs = parse_outer_args()
    print_json(data=args)
    print(subargs)

    # argv_sub = sys.argv[1:]
    argv_sub = subargs

    plugin_map = {
            'chromatogram'          : 'paravision.chromatogram',
            'shell_chromatograms'   : 'paravision.shell_chromatograms',
            'grm2d'                 : 'paravision.grm2d',
            'radial_shell_integrate': 'paravision.radial_shell_integrate',
            'column_snapshot'       : 'paravision.column_snapshot',
            'column_snapshot_fast'  : 'paravision.column_snapshot_fast',
            'volume_integral'       : 'paravision.volume_integral',
            'bead_loading'          : 'paravision.bead_loading',
            'mass_flux'             : 'paravision.mass_flux',
            'animate'               : 'paravision.animate',
            'screenshot'            : 'paravision.screenshot',
            'save_projection'       : 'paravision.save_projection',
            'infogeneric'           : 'paravision.infogeneric',
            'radial_porosity'       : 'paravision.radial_porosity',
            'radial_shell_flux'     : 'paravision.radial_shell_flux',
            'hotspots'              : 'paravision.hotspots',
            'screenshot_with_edges' : 'paravision.screenshot_with_edges',
            'plot_over_line'        : 'paravision.plot_over_line',
            'chromoo'               : 'paravision.chromoo',
            'mesh_info'             : 'paravision.mesh_info'
            }

    if args.cmd.replace('-','_') in plugin_map: 
        run_module(plugin_map[args.cmd.replace('-','_')], argv_sub, args.nproc)
    elif Path(args.cmd).exists(): 
        run_command(args.cmd, argv_sub, args.nproc)
    elif args.cmd == 'pipeline': 
        ## NOTE: Pipeline operations below [EXPERIMENTAL]
        ## NOTE: Doesn't work in parallel!!
        ## The idea is to provide the sequence of operations on the commandline
        ##  and execute it here
        # FIXME: I've changed the argument parsing a couple of times without updating this block
        from paravision.project import project
        from paravision.screenshot import screenshot
        supported_operations = {
            'project': project,
            'screenshot': screenshot
        }

        reader = read_files(args['FILES'], filetype=args['filetype'])
        object = reader
        args.pipeline = args.pipeline or []

        for operation in args.pipeline:
            object = supported_operations[operation](object, args)
    else: 
        raise ValueError(f'{args.cmd} is an invalid command. Please ensure that the order of commandline args is correct.')

if __name__ == "__main__":
    main()