import dgl
import numpy as np
import backend as F
import unittest, pytest
import os

@unittest.skipIf(os.name == 'nt', reason='Cython only works on linux')
def test_cython():
    import dgl._ffi._cy3.core

@pytest.mark.parametrize('arg', [1, 2.3])
def test_callback(arg):
    def cb(x):
        return x + 1
    ret = dgl._api_internal._TestPythonCallback(cb, arg)
    assert ret == arg + 1

@pytest.mark.parametrize('dtype', [F.float32, F.float64, F.int32, F.int64])
def _test_callback_array(dtype):
    def cb(x):
        return F.to_dgl_nd(F.from_dgl_nd(x) + 1)
    arg = F.copy_to(F.tensor([1, 2, 3], dtype=dtype), F.ctx())
    ret = F.from_dgl_nd(dgl._api_internal._TestPythonCallback(cb, F.to_dgl_nd(arg)))
    assert np.allclose(F.asnumpy(ret), F.asnumpy(arg) + 1)

@pytest.mark.parametrize('arg', [1, 2.3])
def test_callback_thread(arg):
    def cb(x):
        return x + 1
    ret = dgl._api_internal._TestPythonCallbackThread(cb, arg)
    assert ret == arg + 1

@pytest.mark.parametrize('dtype', [F.float32, F.float64, F.int32, F.int64])
def _test_callback_array_thread(dtype):
    def cb(x):
        return F.to_dgl_nd(F.from_dgl_nd(x) + 1)
    arg = F.copy_to(F.tensor([1, 2, 3], dtype=dtype), F.ctx())
    ret = F.from_dgl_nd(dgl._api_internal._TestPythonCallbackThread(cb, F.to_dgl_nd(arg)))
    assert np.allclose(F.asnumpy(ret), F.asnumpy(arg) + 1)