"""Intrinsics of TVM-Python Hybrid Script for Python emulation runtime"""

import numpy

class _range(object):
    """Base class of the loop ranges in hybrid script"""
    def __init__(self, a, b=None):
        if b is None:
            self.low = 0
            self.ext = a
        else:
            self.low = a
            self.ext = b

    def __iter__(self):
        i = 0
        while i < self.ext:
            yield i + self.low
            i += 1


class bind(_range): #pylint: disable=invalid-name
    def __init__(self, tag, ext):
        super(bind, self).__init__(ext)
        self.tag = tag


unroll = vectorize = parallel = _range #pylint: disable=invalid-name


def allocate(shape, dtype='float32', scope='global'): #pylint: disable=unused-argument
    """Allocate a buffer with given shape

    Parameters
    ----------
    shape: Tuple
        The shape of the tensor to be allocated
    dtype: string
        The data type of the tensor
    scope: string
        The storage scope of the tensor

    Returns
    -------
    tensor: numpy.array
        The tensor allocated
    """
    return numpy.zeros(shape).astype(dtype)

output_tensor = allocate #pylint: disable=invalid-name

def popcount(x):
    """
    Count ones in the binary representation of number x

    Parameters
    ----------
    x: Integer
        The number to be counted

    Returns
    -------
    cnt: Integer
        The number of ones in the binary representation of number x
    """
    cnt = 0
    while x:
        x -= x & -x
        cnt += 1
    return cnt


def sigmoid(x):
    """
    Sigmoid function of x, aka 1/(1+exp(-x)).

    Parameters
    ----------
    x: a real number

    Returns
    -------
    res: a real number
        The result of sigmoid function
    """
    return 1 / (1 + numpy.exp(-x))


HYBRID_GLOBALS = {
    'unroll'       : unroll,
    'vectorize'    : vectorize,
    'parallel'     : parallel,
    'allocate'     : allocate,
    'output_tensor': output_tensor,
    'bind'         : bind,
    'sqrt'         : numpy.sqrt,
    'log'          : numpy.log,
    'tanh'         : numpy.tanh,
    'power'        : numpy.power,
    'exp'          : numpy.exp,
    'sigmoid'      : sigmoid,
    'popcount'     : popcount
}
