# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
import functools
from typing import Callable, Optional

import megengine._internal as mgb

from ..core import set_default_device

_master_ip = None
_master_port = 0
_world_size = 0
_rank = 0
_backend = None


def init_process_group(
    master_ip: str,
    master_port: int,
    world_size: int,
    rank: int,
    dev: int,
    backend: Optional[str] = "nccl",
) -> None:
    """Initialize the distributed process group, and also specify the device used in the current process.

    :param master_ip: IP address of the master node.
    :param master_port: Port available for all processes to communicate.
    :param world_size: Total number of processes participating in the job.
    :param rank: Rank of the current process.
    :param dev: The GPU device id to bind this process to.
    :param backend: Communicator backend, currently support 'nccl' and 'ucx'
    """
    global _master_ip  # pylint: disable=global-statement
    global _master_port  # pylint: disable=global-statement
    global _world_size  # pylint: disable=global-statement
    global _rank  # pylint: disable=global-statement
    global _backend  # pylint: disable=global-statement

    if not isinstance(master_ip, str):
        raise TypeError("Expect type str but got {}".format(type(master_ip)))
    if not isinstance(master_port, int):
        raise TypeError("Expect type int but got {}".format(type(master_port)))
    if not isinstance(world_size, int):
        raise TypeError("Expect type int but got {}".format(type(world_size)))
    if not isinstance(rank, int):
        raise TypeError("Expect type int but got {}".format(type(rank)))
    if not isinstance(backend, str):
        raise TypeError("Expect type str but got {}".format(type(backend)))

    _master_ip = master_ip
    _master_port = master_port
    _world_size = world_size
    _rank = rank
    _backend = backend

    set_default_device(mgb.comp_node("gpu" + str(dev)))

    if rank == 0:
        _master_port = 
mgb.config.create_mm_server("0.0.0.0", master_port)