from ctypes import *
from .. import ndarray
from ..stream import *
import numpy as np
import os


def _load_nccl_lib():
    """Load libary in build/lib."""
    curr_path = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
    lib_path = os.path.join(curr_path, '../../../build/lib/')
    path_to_so_file = os.path.join(lib_path, "lib_nccl_runtime.so")
    lib = CDLL(path_to_so_file, RTLD_LOCAL)
    return lib


lib_nccl = _load_nccl_lib()


class NCCL_Communicator():

    def __init__(self, devs, devs_number):
        self.comms = (c_int64 * devs_number)(0)
        self.streams = (c_int64 * devs_number)(0)
        self.stream_handles = []
        self.devs = (c_int * devs_number)(*devs)
        self.devs_number = c_int(devs_number)
        self.send_buff = None
        self.recv_buff = None

    def _create_streams(self):
        for i in range(self.devs_number.value):
            self.stream_handles.append(create_stream_handle(ndarray.gpu(i)))
            lib_nccl.update_stream(i, self.streams, c_int64(
                self.stream_handles[-1].handle.contents.handle))

    def _destroy_streams(self):
        self.stream_handles = []
        lib_nccl.free_streams(self.streams, self.devs, self.devs_number)

    def _init_NCCL(self):
        lib_nccl.init_NCCL(self.comms, self.devs, self.devs_number)

    def _destroy_NCCL_comms(self):
        lib_nccl.finish_NCCL(self.comms, self.devs_number)

    def _stream_sync(self):
        lib_nccl.Synchronize_streams(self.streams, self.devs, self.devs_number)

    def _allreduce(self, send_buff, recv_buff, size):
        lib_nccl.NCCL_AllReduce(
            send_buff, recv_buff, size, self.comms, self.streams, self.devs_number)

    def get_send_buff(self, array):
        self.send_buff = (c_void_p * self.devs_number.value)(*array)

    def get_recv_buff(self, array):
        self.recv_buff = (c_void_p * self.devs_number.value)(*array)

    def all_reduce(self, send_array, recv_array, size):
        self.get_send_buff(send_array)
        self.get_recv_buff(recv_array)
        self._allreduce(self.send_buff, self.recv_buff, size)

    def show_property(self):
        print("self.comms = ", self.comms)
        print("self.streams = ", self.streams)
        print("self.devs = ")
        lib_nccl.for_each(self.devs, self.devs_number)
        print("self.devs_number = ", self.devs_number.value)

    def All_Reduce_Ndarray(self, gradient_list, allreduced_gradient_list):
        gradient_buff = []
        allreduced_gradient_buff = []
        for i in range(self.devs_number.value):
            gradient_buff.append(gradient_list[i].handle.contents.data)
            allreduced_gradient_buff.append(
                allreduced_gradient_list[i].handle.contents.data)
        length = 1
        for i in range(gradient_list[0].handle.contents.ndim):
            length = length * gradient_list[0].handle.contents.shape[i]

        self.all_reduce(gradient_buff, allreduced_gradient_buff, length)


def nccl_communicator(devs, devs_number):
    return NCCL_Communicator(devs, devs_number)