# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import functools
import os
import pickle
import random
import warnings
from collections import OrderedDict

import numpy as np
import torch
import torch.distributed as dist

from scepter.modules.utils.model import StdMsg

__all__ = [
    'gather_data', 'we', 'broadcast', 'barrier', 'reduce_scatter', 'reduce',
    'all_reduce', 'send', 'recv', 'isend', 'irecv', 'scatter',
    'shared_random_seed'
]

try:
    from onnxruntime.transformers.benchmark_helper import set_random_seed
except Exception:

    def set_random_seed(seed):
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        np.random.seed(seed)
        random.seed(seed)


def get_dist_info():
    if dist.is_available() and dist.is_initialized():
        return dist.get_rank(), dist.get_world_size()
    else:
        return 0, 1


def gather_data(data):
    """ Gather tensors and other picklable objects to rank 0.
    Will recursively walk through inner list and dict values.

    Args:
        data (any): Anything.

    Returns:
        A object has same structure with input `data`.
    """
    if not we.is_distributed:
        return data
    if isinstance(data, torch.Tensor):
        return gather_gpu_tensors(data)
    elif isinstance(data, dict):
        # Keep in order, dict type DO NOT guarantee a fixed key order
        keys = sorted(list(data.keys()))
        ret = OrderedDict()
        for key in keys:
            ret[key] = gather_data(data[key])
        return ret
    elif isinstance(data, list):
        return gather_list(data)
    else:
        return gather_picklable(data)


def gather_list(data):
    """ Gather list of picklable objects to a new list on rank 0.
    Will NOT recursively walk through.

    Args:
        data (list): List of picklable things.

    Returns:
        A new flat list.
    """
    rank, _ = get_dist_info()
    list_of_list = gather_picklable(data)
    if rank == 0:
        return sum(list_of_list, [])


def gather_picklable(data):
    """ Gather picklable object to a list on rank 0.
    Will NOT recursively walk through.

    Args:
        data (picklable): Picklable data.

    Returns:
        A list contains data collected.
    """
    from packaging import version
    from torch.version import __version__
    if version.parse(__version__) < version.parse('1.8.0'):
        return _gather_picklable_custom(data)
    else:
        rank, world_size = we.rank, we.world_size
        obj_list = [None for _ in range(world_size)]
        dist.all_gather_object(obj_list, data)
        if rank == 0:
            return obj_list


def _gather_picklable_custom(data):
    """ Custom implementation function to gather picklable object to a list on rank 0.
    If torch version is lower than 1.8.0, use this.

    Args:
        data (picklable): Picklable data.

    Returns:
        A list contains data collected.
    """
    import pickle
    byte_tensor = torch.tensor(bytearray(pickle.dumps(data)),
                               dtype=torch.uint8,
                               device='cuda')
    rank, world_size = we.rank, we.world_size
    shape_tensor = torch.tensor(byte_tensor.shape, device='cuda')
    shape_list = [shape_tensor.clone() for _ in range(world_size)]
    dist.all_gather(shape_list, shape_tensor)
    shape_max = torch.tensor(shape_list).max()

    tensor_send = torch.zeros(shape_max,
                              dtype=byte_tensor.dtype,
                              device='cuda')
    tensor_send[0:shape_tensor[0]] = byte_tensor
    tensor_list = [torch.zeros_like(tensor_send) for _ in range(world_size)]
    dist.all_gather(tensor_list, tensor_send)

    if rank == 0:
        data_out = []
        for tensor_recv, shape_recv in zip(tensor_list, shape_list):
            new_data = pickle.loads(
                tensor_recv[:shape_recv[0]].cpu().numpy().tobytes())
            data_out.append(new_data)
        return data_out


def gather_gpu_tensors(tensor, all_recv=False, is_cat=True):
    """
    Args:
        tensor (torch.Tensor):
        all_recv: Gather tensor to rank 0 and concat it.

    Returns:
        A new tensor.
    """
    assert dist.get_backend() == 'nccl'

    device = tensor.device
    if device.type == 'cpu':
        tensor = tensor.to(we.device_id)

    rank, world_size = we.rank, we.world_size

    shape_tensor = torch.tensor(tensor.shape[0], device='cuda')
    shape_list = [shape_tensor.clone() for _ in range(world_size)]
    dist.all_gather(shape_list, shape_tensor)
    shape_max = torch.tensor(shape_list).max()

    tensor_send = torch.zeros((shape_max, *tensor.shape[1:]),
                              dtype=tensor.dtype,
                              device='cuda')
    tensor_send[0:tensor.shape[0]] = tensor
    tensor_list = [torch.zeros_like(tensor_send) for _ in range(world_size)]
    dist.all_gather(tensor_list, tensor_send)
    if not all_recv:
        if rank == 0:
            if not is_cat:
                return tensor_list, shape_list
            tensors_out = []
            for tensor_recv, shape_recv in zip(tensor_list, shape_list):
                tensors_out.append(tensor_recv[0:shape_recv])
            tensor_out = torch.cat(tensors_out).contiguous()
            if device.type == 'cpu':
                tensor_out = tensor_out.cpu()
            del tensor_list, shape_list
            return tensor_out
        else:
            del tensor_list, shape_list
    else:
        if not is_cat:
            return tensor_list, shape_list
        tensors_out = []
        for tensor_recv, shape_recv in zip(tensor_list, shape_list):
            tensors_out.append(tensor_recv[0:shape_recv])
        tensor_out = torch.cat(tensors_out).contiguous()
        if device.type == 'cpu':
            tensor_out = tensor_out.cpu()
        del tensor_list, shape_list
        return tensor_out


def broadcast(tensor, src, group=None, **kwargs):
    if we.is_distributed:
        return dist.broadcast(tensor, src, group, **kwargs)


def barrier():
    if we.is_distributed:
        dist.barrier()


@functools.lru_cache()
def get_global_gloo_group():
    backend = dist.get_backend()
    assert backend in ['gloo', 'nccl']
    if backend == 'nccl':
        return dist.new_group(backend='gloo')
    else:
        return dist.group.WORLD


def reduce_scatter(output,
                   input_list,
                   op=dist.ReduceOp.SUM,
                   group=None,
                   **kwargs):
    if we.is_distributed:
        return dist.reduce_scatter(output, input_list, op, group, **kwargs)


def all_reduce(tensor, op=dist.ReduceOp.SUM, group=None, **kwargs):
    if we.is_distributed:
        return dist.all_reduce(tensor, op, group, **kwargs)


def reduce(tensor, dst, op=dist.ReduceOp.SUM, group=None, **kwargs):
    if we.is_distributed:
        return dist.reduce(tensor, dst, op, group, **kwargs)


def _serialize_to_tensor(data):
    buffer = pickle.dumps(data)
    storage = torch.ByteStorage.from_buffer(buffer)
    tensor = torch.ByteTensor(storage)
    return tensor


def _unserialize_from_tensor(recv_data):
    buffer = recv_data.cpu().numpy().tobytes()
    return pickle.loads(buffer)


def send(tensor, dst, group=None, **kwargs):
    if we.is_distributed:
        assert tensor.is_contiguous(
        ), 'ops.send requires the tensor to be contiguous()'
        return dist.send(tensor, dst, group, **kwargs)


def recv(tensor, src=None, group=None, **kwargs):
    if we.is_distributed:
        assert tensor.is_contiguous(
        ), 'ops.recv requires the tensor to be contiguous()'
        return dist.recv(tensor, src, group, **kwargs)


def isend(tensor, dst, group=None, **kwargs):
    if we.is_distributed:
        assert tensor.is_contiguous(
        ), 'ops.isend requires the tensor to be contiguous()'
        return dist.isend(tensor, dst, group, **kwargs)


def irecv(tensor, src=None, group=None, **kwargs):
    if we.is_distributed:
        assert tensor.is_contiguous(
        ), 'ops.irecv requires the tensor to be contiguous()'
        return dist.irecv(tensor, src, group, **kwargs)


def scatter(data, scatter_list=None, src=0, group=None, **kwargs):
    r"""NOTE: only supports CPU tensor communication.
    """
    world_size = we.world_size
    if world_size == 1:
        data.copy_(scatter_list[0])
    if group is None:
        group = get_global_gloo_group()
    return dist.scatter(data, scatter_list, src, group, **kwargs)


def shared_random_seed():
    seed = np.random.randint(2**31)
    all_seeds, _ = gather_gpu_tensors(seed, all_recv=True, is_cat=False)
    return all_seeds[0]


global we


def mp_worker(gpu, ngpus_per_node, cfg, fn, pmi_rank, world_size, work_env):
    rank = pmi_rank * ngpus_per_node + gpu
    work_env.device_id = gpu % ngpus_per_node
    work_env.rank = rank
    dist.init_process_group(backend='nccl', world_size=world_size, rank=rank)
    torch.backends.cudnn.deterministic = cfg.ENV.get('CUDNN_DETERMINISTIC',
                                                     True)
    torch.backends.cudnn.benchmark = cfg.ENV.get('CUDNN_BENCHMARK', False)
    torch.cuda.set_device(work_env.device_id)
    if work_env.logger is not None:
        work_env.logger.info(
            f'Now running in the distributed environment with world size {work_env.world_size}!'
        )
        work_env.logger.info(f'PMI rank {pmi_rank}!')
        work_env.logger.info(f'Nums of gpu {ngpus_per_node}!')
        work_env.logger.info(
            f'Current rank {work_env.rank} current devices num {ngpus_per_node} '
            f'current machine rank {pmi_rank} and all world size {world_size}')

    we.set_env(work_env.get_env())
    fn(cfg)


class Workenv(object):
    def __init__(self):
        self.initialized = False
        self.is_distributed = False
        self.sync_bn = False
        self.rank = 0
        self.world_size = 1
        self.device_id = 0
        self.device_count = 1
        self.seed = 2023
        self.debug = False
        self.use_pl = False
        self.launcher = 'spawn' if torch.cuda.device_count() > 1 else None
        self.data_online = False
        self.share_storage = False

    def init_env(self, config, fn, logger=None):
        # if use pytorch_lightning: then direct use pytorch_lightning.
        config.ENV = config.get('ENV', {})
        self.seed = config.ENV.get('SEED', 2023)
        self.debug = os.environ.get('ES_DEBUG', None) == 'true'
        set_random_seed(self.seed)
        if logger is not None:
            logger.info(f'And running with seed {self.seed}!')
        if config.ENV.get('USE_PL', False):
            self.use_pl = config.ENV.USE_PL
            fn(config)
            return
        if hasattr(config, 'args') and hasattr(config.args, 'launcher'):
            self.launcher = config.args.launcher
        if logger is None:
            self.logger = StdMsg(name='env')
        else:
            self.logger = logger

        self.data_online = os.environ.get('DATA_ONLINE', None) == 'true'
        self.share_storage = os.environ.get('SHARE_STORAGE', None) == 'true'

        if not torch.cuda.is_available():
            self.device_id = 'cpu'
            fn(config)
            return

        if (os.environ.get('WORLD_SIZE') is None or int(os.environ.get('WORLD_SIZE')) == 1) \
                and torch.cuda.device_count() == 1 and not self.launcher == 'dist':
            self.device_id = 0
            fn(config)
            return

        if self.launcher == 'torchrun':
            try:
                torch.multiprocessing.set_start_method('spawn')
            except Exception as e:
                warnings.warn(f'{e}')
            # checking train mode is distributed or not
            if not os.environ.get('WORLD_SIZE') is None:
                if self.logger is not None:
                    self.logger.info(
                        f"Now running in the distributed environment with {os.environ.get('WORLD_SIZE')}!"
                    )
                self.is_distributed = True
            if not self.initialized:
                if self.is_distributed:
                    self.backend = config.ENV.get('BACKEND', 'nccl')
                    self.sync_bn = config.ENV.get('SYNC_BN', False)
                    dist.init_process_group(backend=self.backend)
                    # dist.barrier()
                self.initialized = True
                if dist.is_initialized():
                    self.rank, self.world_size = dist.get_rank(
                    ), dist.get_world_size()
                    if self.logger is not None:
                        self.logger.info(f'And running in rank {self.rank}!')
                        self.logger.info(
                            f"And cuda visible devices {os.environ.get('CUDA_VISIBLE_DEVICES')}!"
                        )
                else:
                    self.rank, self.world_size = 0, 1
                local_devices = os.environ.get(
                    'LOCAL_WORLD_SIZE') or torch.cuda.device_count()
                local_devices = int(local_devices)
                self.device_count = local_devices
                self.device_id = self.rank % local_devices
            self.logger.info(f"We's attributes: \n"
                             f' launcher {self.launcher} \n'
                             f' rank {self.rank} \n'
                             f' world size {self.world_size} \n'
                             f' device_id {self.device_id}')
            torch.cuda.set_device(self.device_id)
            torch.backends.cudnn.deterministic = config.ENV.get(
                'CUDNN_DETERMINISTIC', True)
            torch.backends.cudnn.benchmark = config.ENV.get(
                'CUDNN_BENCHMARK', False)
            fn(config)
        else:
            import torch.multiprocessing as mp
            if 'MASTER_ADDR' not in os.environ:
                os.environ['MASTER_ADDR'] = 'localhost'
            if 'MASTER_PORT' not in os.environ:
                os.environ['MASTER_PORT'] = '14567'
            pmi_rank = int(os.environ.get('RANK', 0))
            pmi_world_size = int(os.environ.get('WORLD_SIZE', 1))
            ngpus_per_node = os.environ.get(
                'LOCAL_WORLD_SIZE') or torch.cuda.device_count()
            ngpus_per_node = int(ngpus_per_node)
            self.device_count = ngpus_per_node
            world_size = ngpus_per_node * pmi_world_size
            self.world_size = world_size
            if self.world_size > 1:
                self.is_distributed = True
                self.initialized = True
            if self.is_distributed:
                self.backend = config.ENV.get('BACKEND', 'nccl')
                self.sync_bn = config.ENV.get('SYNC_BN', False)
            mp.spawn(mp_worker,
                     nprocs=ngpus_per_node,
                     args=(ngpus_per_node, config, fn, pmi_rank, world_size,
                           self))

    def get_env(self):
        return self.__dict__

    def set_env(self, we_env):
        for k, v in we_env.items():
            setattr(self, k, v)
        set_random_seed(self.seed)

    def __str__(self):
        environ_str = f'Now running in the distributed environment with world size {self.world_size}\n!'
        environ_str += f'Current pod have {self.device_count} devices!\n'
        environ_str += f'Current task executes on device {self.device_id}!\n'
        environ_str += f"Current task's global rank is {self.rank} \n"
        environ_str += f"Current task's data online is set {self.data_online}"
        environ_str += f"Current task's share storage is set {self.share_storage}"
        environ_str += f"Current task's global seed is set {self.seed}"
        return environ_str


we = Workenv()