from libensemble.resources import mpi_resources
from libensemble.executors.executor import jassert
import argparse
import logging

logger = logging.getLogger(__name__)
# To change logging level for just this module
# logger.setLevel(logging.DEBUG)


class MPIRunner:

    @staticmethod
    def get_runner(mpi_runner_type, runner_name=None):

        mpi_runners = {
            'mpich': MPICH_MPIRunner,
            'openmpi': OPENMPI_MPIRunner,
            'aprun': APRUN_MPIRunner,
            'srun': SRUN_MPIRunner,
            'jsrun': JSRUN_MPIRunner,
            'custom': MPIRunner
        }
        mpi_runner = mpi_runners[mpi_runner_type]
        if runner_name is not None:
            runner = mpi_runner(runner_name)
        else:
            runner = mpi_runner()
        return runner

    def __init__(self, run_command='mpiexec'):
        self.run_command = run_command
        self.mpi_command = [self.run_command, '{extra_args}']
        self.subgroup_launch = False
        self.mfile_support = False
        self.arg_nprocs = ('--LIBE_NPROCS_ARG_EMPTY',)
        self.arg_nnodes = ('--LIBE_NNODES_ARG_EMPTY',)
        self.arg_ppn = ('--LIBE_PPN_ARG_EMPTY',)

    def _get_parser(self, p_args, nprocs, nnodes, ppn):
        parser = argparse.ArgumentParser(description='Parse extra_args', allow_abbrev=False)
        parser.add_argument(*nprocs, type=int, dest='num_procs', default=None)
        parser.add_argument(*nnodes, type=int, dest='num_nodes', default=None)
        parser.add_argument(*ppn, type=int, dest='procs_per_node', default=None)
        args, _ = parser.parse_known_args(p_args)
        return args

    def _parse_extra_args(self, num_procs, num_nodes, procs_per_node,
                          hyperthreads, extra_args):

        splt_extra_args = extra_args.split()
        p_args = self._get_parser(splt_extra_args, self.arg_nprocs, self.arg_nnodes, self.arg_ppn)

        # Only fill from extra_args if not set by portable options
        if num_procs is None:
            num_procs = p_args.num_procs
        if num_nodes is None:
            num_nodes = p_args.num_nodes
        if procs_per_node is None:
            procs_per_node = p_args.procs_per_node

        extra_args = ' '.join(splt_extra_args)
        return num_procs, num_nodes, procs_per_node, p_args

    def _rm_replicated_args(self, num_procs, num_nodes, procs_per_node, p_args):
        if p_args.num_procs is not None:
            num_procs = None
        if p_args.num_nodes is not None:
            num_nodes = None
        if p_args.procs_per_node is not None:
            procs_per_node = None
        return num_procs, num_nodes, procs_per_node

    def express_spec(self, task, num_procs, num_nodes,
                     procs_per_node, machinefile,
                     hyperthreads, extra_args,
                     resources, workerID):

        hostlist = None
        machinefile = None
        # Always use host lists (unless uneven mapping)
        hostlist = mpi_resources.get_hostlist(resources, num_nodes)
        return hostlist, machinefile

    def get_mpi_specs(self, task, num_procs, num_nodes,
                      procs_per_node, machinefile,
                      hyperthreads, extra_args,
                      resources, workerID):
        "Form the mpi_specs dictionary."

        # Return auto_resource variables inc. extra_args additions
        if extra_args:
            num_procs, num_nodes, procs_per_node, p_args = \
                self._parse_extra_args(num_procs, num_nodes, procs_per_node,
                                       hyperthreads, extra_args=extra_args)

        hostlist = None
        if machinefile and not self.mfile_support:
            logger.warning('User machinefile ignored - not supported by {}'.format(self.run_command))
            machinefile = None

        if machinefile is None and resources is not None:
            num_procs, num_nodes, procs_per_node = \
                mpi_resources.get_resources(resources, num_procs, num_nodes,
                                            procs_per_node, hyperthreads)
            hostlist, machinefile = \
                self.express_spec(task, num_procs, num_nodes,
                                  procs_per_node, machinefile,
                                  hyperthreads, extra_args,
                                  resources, workerID)
        else:
            num_procs, num_nodes, procs_per_node = \
                mpi_resources.task_partition(num_procs, num_nodes,
                                             procs_per_node, machinefile)

        # Remove portable variable if in extra_args
        if extra_args:
            num_procs, num_nodes, procs_per_node = \
                self._rm_replicated_args(num_procs, num_nodes,
                                         procs_per_node, p_args)

        return {'num_procs': num_procs,
                'num_nodes': num_nodes,
                'procs_per_node': procs_per_node,
                'extra_args': extra_args,
                'machinefile': machinefile,
                'hostlist': hostlist}


class MPICH_MPIRunner(MPIRunner):

    def __init__(self, run_command='mpirun'):
        self.run_command = run_command
        self.subgroup_launch = True
        self.mfile_support = True
        self.arg_nprocs = ('-n', '-np')
        self.arg_nnodes = ('--LIBE_NNODES_ARG_EMPTY',)
        self.arg_ppn = ('--ppn',)
        self.mpi_command = [self.run_command, '--env {env}',
                            '-machinefile {machinefile}',
                            '-hosts {hostlist}', '-np {num_procs}',
                            '--ppn {procs_per_node}', '{extra_args}']


class OPENMPI_MPIRunner(MPIRunner):

    def __init__(self, run_command='mpirun'):
        self.run_command = run_command
        self.subgroup_launch = True
        self.mfile_support = True
        self.arg_nprocs = ('-n', '-np', '-c', '--n')
        self.arg_nnodes = ('--LIBE_NNODES_ARG_EMPTY',)
        self.arg_ppn = ('-npernode',)
        self.mpi_command = [self.run_command, '-x {env}',
                            '-machinefile {machinefile}',
                            '-host {hostlist}', '-np {num_procs}',
                            '-npernode {procs_per_node}', '{extra_args}']

    def express_spec(self, task, num_procs, num_nodes,
                     procs_per_node, machinefile,
                     hyperthreads, extra_args,
                     resources, workerID):

        hostlist = None
        machinefile = None
        # Use machine files for OpenMPI
        # as "-host" requires entry for every rank

        machinefile = "machinefile_autogen"
        if workerID is not None:
            machinefile += "_for_worker_{}".format(workerID)
        machinefile += "_task_{}".format(task.id)
        mfile_created, num_procs, num_nodes, procs_per_node = \
            mpi_resources.create_machinefile(resources, machinefile,
                                             num_procs, num_nodes,
                                             procs_per_node, hyperthreads)
        jassert(mfile_created, "Auto-creation of machinefile failed")

        return hostlist, machinefile


class APRUN_MPIRunner(MPIRunner):

    def __init__(self, run_command='aprun'):
        self.run_command = run_command
        self.subgroup_launch = False
        self.mfile_support = False
        self.arg_nprocs = ('-n',)
        self.arg_nnodes = ('--LIBE_NNODES_ARG_EMPTY',)
        self.arg_ppn = ('-N',)
        self.mpi_command = [self.run_command, '-e {env}',
                            '-L {hostlist}', '-n {num_procs}',
                            '-N {procs_per_node}', '{extra_args}']


class SRUN_MPIRunner(MPIRunner):

    def __init__(self, run_command='srun'):
        self.run_command = run_command
        self.subgroup_launch = False
        self.mfile_support = False
        self.arg_nprocs = ('-n', '--ntasks')
        self.arg_nnodes = ('-N', '--nodes')
        self.arg_ppn = ('--ntasks-per-node',)
        self.mpi_command = [self.run_command, '-w {hostlist}',
                            '--ntasks {num_procs}',
                            '--nodes {num_nodes}',
                            '--ntasks-per-node {procs_per_node}',
                            '{extra_args}']


class JSRUN_MPIRunner(MPIRunner):

    def __init__(self, run_command='jsrun'):
        self.run_command = run_command
        self.subgroup_launch = True
        self.mfile_support = False

        # TODO: Add multiplier to resources checks (for -c/-a)
        self.arg_nprocs = ('--np', '-n')
        self.arg_nnodes = ('--LIBE_NNODES_ARG_EMPTY',)
        self.arg_ppn = ('-r',)
        self.mpi_command = [self.run_command, '-n {num_procs}',
                            '-r {procs_per_node}', '{extra_args}']

    def get_mpi_specs(self, task, num_procs, num_nodes,
                      procs_per_node, machinefile,
                      hyperthreads, extra_args,
                      resources, workerID):

        # Return auto_resource variables inc. extra_args additions
        if extra_args:
            num_procs, num_nodes, procs_per_node, p_args = \
                self._parse_extra_args(num_procs, num_nodes, procs_per_node,
                                       hyperthreads, extra_args=extra_args)

        rm_rpn = True if procs_per_node is None and num_nodes is None else False

        hostlist = None
        if machinefile and not self.mfile_support:
            logger.warning('User machinefile ignored - not supported by {}'.format(self.run_command))
            machinefile = None
        if machinefile is None and resources is not None:
            num_procs, num_nodes, procs_per_node = \
                mpi_resources.get_resources(resources, num_procs, num_nodes,
                                            procs_per_node, hyperthreads)

            # TODO: Create ERF file if mapping worker to resources req.
        else:
            num_procs, num_nodes, procs_per_node = \
                mpi_resources.task_partition(num_procs, num_nodes,
                                             procs_per_node, machinefile)

        # Remove portable variable if in extra_args
        if extra_args:
            num_procs, num_nodes, procs_per_node = \
                self._rm_replicated_args(num_procs, num_nodes,
                                         procs_per_node, p_args)

        if rm_rpn:
            procs_per_node = None

        return {'num_procs': num_procs,
                'num_nodes': num_nodes,
                'procs_per_node': procs_per_node,
                'extra_args': extra_args,
                'machinefile': machinefile,
                'hostlist': hostlist}
