import numpy

from cupy import cuda
from cupy._creation.basic import _new_like_order_and_strides
from cupy.core import internal


def _update_shape(a, shape):
    if shape is None and a is not None:
        shape = a.shape
    elif isinstance(shape, int):
        shape = (shape,)
    else:
        shape = tuple(shape)
    return shape


def empty_pinned(shape, dtype=float, order='C'):
    """Returns an array without initializing the elements.

    This is a convenience function which is just :func:`numpy.empty`,
    except that the underlying memory is allocated from CuPy's pinned
    memory pool.

    Args:
        shape (int or tuple of ints): Dimensionalities of the array.
        dtype: Data type specifier.
        order ({'C', 'F'}): Row-major (C-style) or column-major
            (Fortran-style) order.

    Returns:
        numpy.ndarray: A new array with elements not initialized.

    .. seealso:: :func:`numpy.empty`

    """
    shape = _update_shape(None, shape)
    nbytes = internal.prod(shape) * numpy.dtype(dtype).itemsize
    mem = cuda.alloc_pinned_memory(nbytes)
    out = numpy.ndarray(shape, dtype=dtype, buffer=mem, order=order)
    return out


def empty_like_pinned(a, dtype=None, order='K', subok=None, shape=None):
    """Returns a new array with same shape and dtype of a given array.

    This is a convenience function which is just :func:`numpy.empty_like`,
    except that the underlying memory is allocated from CuPy's pinned
    memory pool.

    This function currently does not support ``subok`` option.

    Args:
        a (numpy.ndarray or cupy.ndarray): Base array.
        dtype: Data type specifier. The data type of ``a`` is used by default.
        order ({'C', 'F', 'A', or 'K'}): Overrides the memory layout of the
            result. ``'C'`` means C-order, ``'F'`` means F-order, ``'A'`` means
            ``'F'`` if ``a`` is Fortran contiguous, ``'C'`` otherwise.
            ``'K'`` means match the layout of ``a`` as closely as possible.
        subok: Not supported yet, must be None.
        shape (int or tuple of ints): Overrides the shape of the result. If
            ``order='K'`` and the number of dimensions is unchanged, will try
            to keep order, otherwise, ``order='C'`` is implied.

    Returns:
        numpy.ndarray: A new array with same shape and dtype of ``a`` with
        elements not initialized.

    .. seealso:: :func:`numpy.empty_like`

    """
    # We're kinda duplicating the code here because order='K' needs special
    # treatment: strides need to be computed
    if subok is not None:
        raise TypeError('subok is not supported yet')
    if dtype is None:
        dtype = a.dtype
    shape = _update_shape(a, shape)
    order, strides, _ = _new_like_order_and_strides(
        a, dtype, order, shape, get_memptr=False)
    nbytes = internal.prod(shape) * numpy.dtype(dtype).itemsize
    mem = cuda.alloc_pinned_memory(nbytes)
    out = numpy.ndarray(shape, dtype=dtype, buffer=mem,
                        strides=strides, order=order)
    return out


def zeros_pinned(shape, dtype=float, order='C'):
    """Returns a new array of given shape and dtype, filled with zeros.

    This is a convenience function which is just :func:`numpy.zeros`,
    except that the underlying memory is allocated from CuPy's pinned
    memory pool.

    Args:
        shape (int or tuple of ints): Dimensionalities of the array.
        dtype: Data type specifier.
        order ({'C', 'F'}): Row-major (C-style) or column-major
            (Fortran-style) order.

    Returns:
        numpy.ndarray: An array filled with zeros.

    .. seealso:: :func:`numpy.zeros`

    """
    out = empty_pinned(shape, dtype, order)
    numpy.copyto(out, 0, casting='unsafe')
    return out


def zeros_like_pinned(a, dtype=None, order='K', subok=None, shape=None):
    """Returns an array of zeros with same shape and dtype as a given array.

    This is a convenience function which is just :func:`numpy.zeros_like`,
    except that the underlying memory is allocated from CuPy's pinned
    memory pool.

    This function currently does not support ``subok`` option.

    Args:
        a (numpy.ndarray or cupy.ndarray): Base array.
        dtype: Data type specifier. The dtype of ``a`` is used by default.
        order ({'C', 'F', 'A', or 'K'}): Overrides the memory layout of the
            result. ``'C'`` means C-order, ``'F'`` means F-order, ``'A'`` means
            ``'F'`` if ``a`` is Fortran contiguous, ``'C'`` otherwise.
            ``'K'`` means match the layout of ``a`` as closely as possible.
        subok: Not supported yet, must be None.
        shape (int or tuple of ints): Overrides the shape of the result. If
            ``order='K'`` and the number of dimensions is unchanged, will try
            to keep order, otherwise, ``order='C'`` is implied.

    Returns:
        numpy.ndarray: An array filled with zeros.

    .. seealso:: :func:`numpy.zeros_like`

    """
    out = empty_like_pinned(a, dtype, order, subok, shape)
    numpy.copyto(out, 0, casting='unsafe')
    return out

@@ -401,49 +401,74 @@
 ns a
-n array without initializing the elements
+ new, uninitialized NumPy array with the given shape%0A    and dtype
 .%0A%0A 
@@ -581,52 +581,25 @@
  is 
-allocated from CuPy's pinned%0A    memory pool
+pinned/pagelocked
 .%0A%0A 
@@ -1259,24 +1259,45 @@
 eturns a new
+, uninitialized NumPy
  array with 
@@ -1288,32 +1288,36 @@
 umPy array with 
+the 
 same shape and d
@@ -1324,13 +1324,28 @@
 type
+%0A    as those
  of 
-a
+the
  giv
@@ -1475,52 +1475,25 @@
  is 
-allocated from CuPy's pinned%0A    memory pool
+pinned/pagelocked
 .%0A%0A 
@@ -3132,25 +3132,55 @@
 ns a new
- array of
+, zero-initialized NumPy array with the
  given s
@@ -3179,24 +3179,28 @@
  given shape
+%0A   
  and dtype, 
@@ -3201,27 +3201,8 @@
 type
-, filled with zeros
 .%0A%0A 
@@ -3315,52 +3315,25 @@
  is 
-allocated from CuPy's pinned%0A    memory pool
+pinned/pagelocked
 .%0A%0A 
@@ -3854,29 +3854,51 @@
 ns a
-n array of zeros with
+ new, zero-initialized NumPy array with the
  sam
@@ -3918,13 +3918,28 @@
 type
- as a
+%0A    as those of the
  giv
@@ -4069,52 +4069,25 @@
  is 
-allocated from CuPy's pinned%0A    memory pool
+pinned/pagelocked
 .%0A%0A 

update docstrings