from __future__ import absolute_import as _abs
import ctypes as _ctypes
import numpy as np
from nnvm import symbol
from nnvm._base import c_str, check_call, _LIB, c_array, nn_uint

SessionHandle = _ctypes.c_void_p
nn_float = _ctypes.c_float

def _get_numpy(cptr, dtype, shape):
    if dtype != 0:
        raise ValueError("only float32 is supported so far")
    size = 1
    for s in shape:
        size *= s
    if size != 0 and shape:
        dbuffer = (nn_float * size).from_address(_ctypes.addressof(cptr.contents))
        return np.frombuffer(dbuffer, dtype=np.float32).reshape(shape).copy()
    else:
        return None

class Session(object):
    def __init__(self, config='cpu'):
        handle = SessionHandle()
        check_call(_LIB.NNSessionCreate(_ctypes.byref(handle), c_str(config)))
        self.handle = handle

    def __del__(self):
        check_call(_LIB.NNSessionClose(self.handle))

    def run(self, fetch, feed_dict=None):
        if isinstance(fetch, list):
            fetch = symbol.Group(fetch)
        feed_dict = feed_dict if feed_dict else {}
        feed_placeholders = []
        feed_dptr = []
        feed_dtype = []
        feed_shape_csr_ptr = [0]
        feed_shape_data = []
        src_list = []

        for k, v in feed_dict.items():
            assert isinstance(k, symbol.Symbol)
            assert isinstance(v, np.ndarray)
            feed_placeholders.append(k.handle)
            # only convert to float32 for now
            source_array = np.ascontiguousarray(v, dtype=np.float32)
            # leep src_list alive for the period
            src_list.append(source_array)
            feed_dptr.append(source_array.ctypes.data_as(_ctypes.c_void_p))
            feed_dtype.append(0)
            feed_shape_data.extend(source_array.shape)
            feed_shape_csr_ptr.append(len(feed_shape_data))
        out_size = nn_uint()
        out_dptr = _ctypes.POINTER(_ctypes.POINTER(nn_float))()
        out_dtype = _ctypes.POINTER(nn_uint)()
        out_shape_ndim = _ctypes.POINTER(nn_uint)()
        out_shape_data = _ctypes.POINTER(_ctypes.POINTER(nn_uint))()

        check_call(_LIB.NNSessionRun(
            self.handle, fetch.handle, nn_uint(len(src_list)),
            c_array(_ctypes.c_void_p, feed_placeholders),
            c_array(_ctypes.c_void_p, feed_dptr),
            c_array(nn_uint, feed_dtype),
            c_array(nn_uint, feed_shape_csr_ptr),
            c_array(nn_uint, feed_shape_data),
            _ctypes.byref(out_size),
            _ctypes.byref(out_dptr),
            _ctypes.byref(out_dtype),
            _ctypes.byref(out_shape_ndim),
            _ctypes.byref(out_shape_data)))
        ret = []
        for i in range(out_size.value):
            shape = tuple(out_shape_data[i][:out_shape_ndim[i]])
            ret.append(_get_numpy(out_dptr[i], out_dtype[i], shape))

        return ret[0] if len(ret) == 1 else ret