# This file defines workaround implementation for
# NumPy-compatibility functions that fall back to NumPy/CuPy functions
# for native/cuda devices respecitvely.
# The workaround does not support backprop, and also requires external
# libraries mentioned above.
# Functions defined in this file should be considered to have high priority for
# genuine implementations.
import numpy

import chainerx


try:
    import cupy
except Exception:
    cupy = None


class _DummyContext:
    def __enter__(self):
        pass

    def __exit__(self, type, value, traceback):
        pass


_dummy_context = _DummyContext()


def _to_numpy(array):
    assert isinstance(array, chainerx.ndarray)
    return chainerx.to_numpy(array, copy=False)


def _from_numpy(array):
    assert isinstance(array, numpy.ndarray)
    return chainerx.array(array, copy=False)


def _to_cupy(array):
    assert cupy is not None
    # Convert to cupy.ndarray on the same device as source array
    return cupy.ndarray(
        array.shape,
        array.dtype,
        cupy.cuda.MemoryPointer(
            cupy.cuda.UnownedMemory(
                array.data_ptr + array.offset,
                array.data_size,
                array,
                array.device.index),
            0),
        strides=array.strides)


def _from_cupy(array):
    assert cupy is not None
    assert isinstance(array, cupy.ndarray)
    device = chainerx.get_device('cuda', array.device.id)
    return chainerx._core._fromrawpointer(
        array.data.mem.ptr,
        array.shape,
        array.dtype,
        array.strides,
        device,
        array.data.ptr - array.data.mem.ptr,
        array)


def _from_chainerx(array, check_backprop=True):
    # Converts chainerx.ndarray to numpy/cupy.ndarray.
    # Objects with other types are kept intact.
    # Returns a pair: (xp, cupy device or dummy context, numpy/cupy.ndarray).
    if not isinstance(array, chainerx.ndarray):
        return None, _dummy_context, array
    if check_backprop and array.is_backprop_required():
        raise RuntimeError(
            'ChainerX function fallback using NumPy/CuPy is not '
            'supported for arrays that are connected to a graph.')
    backend_name = array.device.backend.name
    if backend_name == 'native':
        return numpy, _dummy_context, _to_numpy(array)
    if backend_name == 'cuda':
        if cupy is None:
            raise RuntimeError(
                'ChainerX fallback implementation for cuda backend requires '
                'cupy to be installed.')
        array_cupy = _to_cupy(array)
        return cupy, array_cupy.device, array_cupy
    raise RuntimeError(
        'ChainerX fallback implementation only supports native or cuda '
        'backends.')


def _to_chainerx(array):
    # Converts numpy/cupy.ndarray to chainerx.ndarray.
    # Objects with other types are kept intact.
    if isinstance(array, numpy.ndarray):
        return _from_numpy(array)
    elif cupy is not None and isinstance(array, cupy.ndarray):
        return _from_cupy(array)
    return array


def _populate_module_functions():

    def _isfinite(arr):
        xp, dev, arr = _from_chainerx(arr)
        with dev:
            ret = xp.isfinite(arr)
        return _to_chainerx(ret)

    chainerx.isfinite = _isfinite

    def _hstack(arrs):
        assert len(arrs) > 0
        arrs2 = []
        for a in arrs:
            xp, dev, a2 = _from_chainerx(a)
            arrs2.append(a2)
        with dev:
            ret = xp.hstack(arrs2)
        return _to_chainerx(ret)

    chainerx.hstack = _hstack

    def _vstack(arrs):
        assert len(arrs) > 0
        arrs2 = []
        for a in arrs:
            xp, dev, a2 = _from_chainerx(a)
            arrs2.append(a2)
        with dev:
            ret = xp.vstack(arrs2)
        return _to_chainerx(ret)

    chainerx.vstack = _vstack


def _populate_ndarray():
    ndarray = chainerx.ndarray

    # __getitem__ with advanced indexing
    old_getitem = ndarray.__getitem__

    def __getitem__(arr, key):
        try:
            return old_getitem(arr, key)
        except (IndexError, chainerx.DimensionError):
            pass

        is_backprop_required = arr.is_backprop_required()

        xp, dev, arr = _from_chainerx(arr, check_backprop=False)
        _, _, key = _from_chainerx(key, check_backprop=False)

        with dev:
            ret = arr[key]

        # Doing this check after the fallback __getitem__ because the error
        # which caused the fallback might not be due to advanced indexing.
        # In such case the fallback __getitem__ should also raise the error.

        if is_backprop_required:
            raise RuntimeError(
                'ChainerX getitem fallback for advanced indexing is not '
                'supported for arrays that are connected to a graph.')

        return _to_chainerx(ret)

    # __setitem__ with advanced indexing
    def __setitem__(self, key, value):
        if self.is_backprop_required():
            raise RuntimeError(
                'ChainerX setitem fallback for advanced indexing is not '
                'supported for arrays that are connected to a graph.')

        xp, dev, self = _from_chainerx(self)
        _, _, key = _from_chainerx(key)
        _, _, value = _from_chainerx(value)

        with dev:
            self[key] = value

    ndarray.__setitem__ = __setitem__
    ndarray.__getitem__ = __getitem__

    def _min(arr, *args, **kwargs):
        _, dev, arr = _from_chainerx(arr)
        with dev:
            ret = arr.min(*args, **kwargs)
        return _to_chainerx(ret)

    ndarray.min = _min

    def _all(arr, *args, **kwargs):
        _, dev, arr = _from_chainerx(arr)
        with dev:
            ret = arr.all(*args, **kwargs)
        return _to_chainerx(ret)

    ndarray.all = _all

    def _any(arr, *args, **kwargs):
        _, dev, arr = _from_chainerx(arr)
        with dev:
            ret = arr.any(*args, **kwargs)
        return _to_chainerx(ret)

    ndarray.any = _any


def populate():
    _populate_module_functions()
    _populate_ndarray()

@@ -972,318 +972,29 @@
 rn c
-upy.ndarray(%0A        array.shape,%0A        array.dtype,%0A        cupy.cuda.MemoryPointer(%0A            cupy.cuda.UnownedMemory(%0A                array.data_ptr + array.offset,%0A                array.data_size,%0A                array,%0A                array.device.index),%0A            0),%0A        strides=array.strides
+hainerx.to_cupy(array
 )%0A%0A%0A

use chainerx.to_cupy()