import re
import sys
import subprocess
import os
import psutil
from .launcher_base import Launcher

class DistributedTrainingLauncher(Launcher):
    '''
    Launcher for distributed traning with MPI launcher
    '''
    def add_params(self, parser):
        group = parser.add_argument_group('Distributed Training Arguments With oneCCL backend')
        group.add_argument(
            '--nnodes',
            type=int,
            default=0,
            help='Number of machines/devices to use for distributed training',
        )
        group.add_argument(
            '--nprocs-per-node',
            '--nprocs_per_node',
            type=int,
            default=0,
            help='Number of processes run on each machine/device',
        )
        # ccl control
        group.add_argument(
            '--ccl-worker-count',
            '--ccl_worker_count',
            default=4,
            type=int,
            help='Number of cores per rank for ccl communication',
        )
        group.add_argument(
            '--logical-cores-for-ccl',
            '--logical_cores_for_ccl',
            action='store_true',
            default=False,
            help='Use logical cores for the ccl worker.',
        )
        # mpi control
        group.add_argument(
            '--master-addr',
            '--master_addr',
            default='127.0.0.1',
            type=str,
            help='Address of master node (rank 0), should be either '
            'IP address or hostname of node 0. For '
            'single node multi-proc training, the '
            '--master-addr can simply be 127.0.0.1',
        )
        group.add_argument(
            '--master-port',
            '--master_port',
            default=29500,
            type=int,
            help='Port on master node (rank 0) for communication during distributed training',
        )
        group.add_argument(
            '--hostfile',
            default='hostfile',
            type=str,
            help='Set the hostfile for multi-node multi-proc '
            'training. The hostfile includes a node address list '
            'containing either IP addresses or hostnames of computation nodes.',
        )
        group.add_argument(
            '--extra-mpi-params',
            '--extra_mpi_params',
            default='',
            type=str,
            help='Extra parameters for mpiexec.hydra '
            'except for -np -ppn -hostfile and -genv I_MPI_PIN_DOMAIN',
        )

    def get_pin_domain_affinity(self, cpu_pools, ccl_worker_count, logical_cores_for_ccl=False):
        '''
        I_MPI_PIN_DOMAIN specify the cores used for every MPI process.
        1) use physical core for oneccl
           The first ccl_worker_count cores of every rank for ccl communication
           and the other cores will be used to do computation.
           For example: on CascadeLake 8280 CPU, 2 ranks on one node. ccl_worker_count=4
           CCL_WORKER_COUNT=4
           CCL_WORKER_AFFINITY="0,1,2,3,28,29,30,31"
           I_MPI_PIN_DOMAIN=[0xffffff0,0xffffff00000000]
        2) use logical core oneccl
           The first ccl_worker_count logical cores which is correponding to the
           first ccl_worker_count physical cores are used as the ccl cores.
           For example: on CascadeLake 8280 CPU, 2 ranks on one node. ccl_worker_count=4
           CCL_WORKER_COUNT=4
           CCL_WORKER_AFFINITY="56,57,58,59,84,85,86,87"
           I_MPI_PIN_DOMAIN=[0xfffffff,0xfffffff0000000]
        '''
        domain_binaries = []
        affinity = []
        for pool in cpu_pools:
            if logical_cores_for_ccl and len([c for c in pool if not c.is_physical_core]) < ccl_worker_count:
                self.verbose('warning', 'Argument --logical-cores-for-ccl is set but no enough logical cores are available. Disable this argument.')
                logical_cores_for_ccl = False
                break
        for pool in cpu_pools:
            domain_binary = 0
            cores = []
            if logical_cores_for_ccl:
                affinity.extend([str(c.cpu) for c in pool if not c.is_physical_core][:ccl_worker_count])
                cores = [str(c.cpu) for c in pool if c.is_physical_core]
            else:
                physical_cores = [str(c.cpu) for c in pool if c.is_physical_core]
                assert ccl_worker_count < len(physical_cores), f'ccl_worker_count ({ccl_worker_count}) cannot exceed number of available cores ({len(physical_cores)}).'
                affinity.extend(physical_cores[:ccl_worker_count])
                cores = physical_cores[ccl_worker_count:]
            for c in cores:
                domain_binary |= (1 << int(c))
            domain_binaries.append(hex(domain_binary))
        return {'pin_domain': f'[{",".join(domain_binaries)}]', 'affinity': ','.join(affinity)}

    def launch(self, args):
        '''
        Set ENVs and launch MPI process for distributed training.
        '''
        assert not(args.logical_cores_for_ccl and args.use_logical_cores), 'Can\'t use --logical-cores-for-ccl and --use-logical-cores at the same time.'
        if args.nnodes > 1:
            assert os.path.exists(args.hostfile), 'A hostfile is required when you perform multi-node distributed training. Please create the hostfile which includes ip addresses of nodes that you will use for the distributed computation workload.'
            ipv4_addr_pattern = r'^(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$'
            ip_list = []
            with open(args.hostfile) as f:
                for line in f:
                    line = line.strip().strip('\n')
                    # is_valid = re.match(ipv4_addr_pattern, line)
                    # if not is_valid:
                    #     logger.error('{} is not valid IPV4 address'.format(line))
                    #     exit(-1)
                    # else:
                    #     ip_list.append(line)
                    ip_list.append(line)
            assert len(ip_list) >= args.nnodes, f'Number of IP addresses in the hostfile ({len(ip_list)}) should be >= the nnodes argument ({args.nnodes}).'
            master_check = False
            dic = psutil.net_if_addrs()
            for adapter in dic:
                snicList = dic[adapter]
                for snic in snicList:
                    if snic.address == ip_list[0]:
                        master_check = True
            assert master_check, 'MASTER_ADDR is incorrect. Please make sure the first line ({ip_list[0]}) of the hostfile is the ip address of the current node.'

            self.verbose('info', 'Begin to validate SSH connections')
            args.master_addr = ip_list[0]
            for ip in ip_list[1:]:
                completed_process = subprocess.run(f'ssh -o PasswordAuthentication=no {ip} ":"', shell=True)
                assert completed_process.returncode == 0, f'Passwordless SSH login to {ip} failed, please make sure you have a SSH public key setup correctly.'
                self.verbose('info', f'Connection from the master node {args.master_addr} to the slave node {ip} succeeded.')

        nodes_list = self.parse_list_argument(args.nodes_list)
        args.nprocs_per_node = len(set([c.node for c in self.cpuinfo.pool_all])) if len(nodes_list) == 0 else len(nodes_list)
        ncores_per_instance = args.ncores_per_instance
        if ncores_per_instance > 0:
            if not args.logical_cores_for_ccl or len([c for c in self.cpuinfo.pool_all if not c.is_physical_core]) < args.nprocs_per_node * args.ccl_worker_count:
                ncores_per_instance += args.ccl_worker_count
            ncores_per_instance = len([c for c in self.cpuinfo.pool_all if c.core < ncores_per_instance])
        self.cpuinfo.gen_pools_ondemand(
                ninstances=args.nprocs_per_node,
                ncores_per_instance=ncores_per_instance,
                use_logical_cores=True,
                use_e_cores=args.use_e_cores,
                nodes_list=nodes_list
                )

        self.set_memory_allocator(args.memory_allocator, False, ['jemalloc'])
        self.set_omp_runtime(args.omp_runtime, True)
        omp_num_threads = len([c for c in self.cpuinfo.pools_ondemand[0] if c.is_physical_core])
        if not args.logical_cores_for_ccl:
            omp_num_threads -= args.ccl_worker_count
        self.add_env('OMP_NUM_THREADS', str(omp_num_threads))

        # set distributed related environmental variables
        self.add_env('MASTER_ADDR', args.master_addr)
        self.add_env('MASTER_PORT', str(args.master_port))
        pin_domain_affinity = self.get_pin_domain_affinity(
                self.cpuinfo.pools_ondemand,
                args.ccl_worker_count,
                args.logical_cores_for_ccl,
                )
        self.add_env('I_MPI_PIN_DOMAIN', pin_domain_affinity['pin_domain'])
        self.add_env('CCL_WORKER_COUNT', str(args.ccl_worker_count))
        self.add_env('CCL_WORKER_AFFINITY', pin_domain_affinity['affinity'])

        ld_preload_marker = 'LD_PRELOAD_UNSET'
        ld_preload_bk = os.environ['LD_PRELOAD'] if 'LD_PRELOAD' in os.environ else ld_preload_marker
        if len(self.ld_preload) > 0:
            os.environ['LD_PRELOAD'] = ':'.join(self.ld_preload)
            self.verbose('info', f'LD_PRELOAD={os.environ["LD_PRELOAD"]}')
        else:
            if 'LD_PRELOAD' in os.environ:
                del os.environ['LD_PRELOAD']
        for k,v in self.environ_set.items():
            self.verbose('info', f'env: {k}={v}')

        cmd = ['mpiexec.hydra']
        genvs = [f'-genv {k}={v}' for k,v in self.environ_set.items()]
        mpi_config = f"-l -np {args.nnodes * args.nprocs_per_node} -ppn {args.nprocs_per_node} {' '.join(genvs)} "
        mpi_config += args.extra_mpi_params
        if args.nnodes > 1:
            mpi_config += f' -hostfile {args.hostfile}'
        cmd.extend(mpi_config.split())
        with_python = not args.no_python
        if with_python:
            cmd.append(sys.executable)
            cmd.append('-u')
        if args.module:
            cmd.append('-m')
        cmd.append(args.program)
        cmd.extend(args.program_args)
        log_name = f'{args.log_file_prefix}.log'
        log_name = os.path.join(args.log_dir, log_name)
        cmd_s = ' '.join(cmd)
        if args.log_dir:
            cmd_s = f'{cmd_s} 2>&1 | tee {log_name}'
        self.verbose('info', f'cmd: {cmd_s}')
        process = subprocess.Popen(cmd_s, env=os.environ, shell=True)
        process.wait()
        if ld_preload_bk == ld_preload_marker:
            if 'LD_PRELOAD' in os.environ:
                del os.environ['LD_PRELOAD']
        else:
            os.environ['LD_PRELOAD'] = ld_preload_bk
        if args.log_dir:
            log_fns = []
            for i in range(args.nnodes * args.nprocs_per_node):
                log_name_rank = f'{args.log_file_prefix}_rank_{i}.log'
                log_name_rank = os.path.join(args.log_dir, log_name_rank)
                fn = open(log_name_rank, 'w')
                log_fns.append(fn)
            with open(log_name) as fp:
                for line in fp:
                    m = re.match('\[(\d+)\] (.*)', line.strip())
                    if m:
                        log_fns[int(m.group(1))].write(f'{m.group(2)}\n')
                    else:
                        self.verbose('warning', f'Failed to detect rank id from log file {log_name} at line "{line.strip()}".')
            for fn in log_fns:
                fn.close()

if __name__ == '__main__':
    pass
