import numpy

try:
    from mock import patch, Mock
except ImportError:
    from unittest.mock import patch, Mock

from kernel_tuner import cuda

def setup_mock(drv):
    context = Mock()
    devprops = { 'MAX_THREADS_PER_BLOCK': 1024,
                'COMPUTE_CAPABILITY_MAJOR': 5,
                'COMPUTE_CAPABILITY_MINOR': 5 }
    context.return_value.get_device.return_value.get_attributes.return_value = devprops
    context.return_value.get_device.return_value.compute_capability.return_value = "55"
    drv.Device.return_value.make_context.return_value = context()
    drv.mem_alloc.return_value = 'mem_alloc'
    return drv

@patch('kernel_tuner.cuda.drv')
def test_ready_argument_list(drv):
    drv = setup_mock(drv)

    size = 5
    a = numpy.int32(75)
    b = numpy.random.randn(size).astype(numpy.float32)
    arguments = [a, b]

    dev = cuda.CudaFunctions(0)
    gpu_args = dev.ready_argument_list(arguments)

    print(drv.mock_calls)
    print(gpu_args)

    drv.mem_alloc.assert_called_once_with(20)
    drv.memcpy_htod.assert_called_once_with('mem_alloc', b)

    assert isinstance(gpu_args[0], numpy.int32)

@patch('kernel_tuner.cuda.DynamicSourceModule')
@patch('kernel_tuner.cuda.drv')
def test_compile(drv, src_mod):

    #setup mocked stuff
    drv = setup_mock(drv)
    dev = cuda.CudaFunctions(0)
    dev.source_mod = src_mod
    dev.source_mod = Mock()
    dev.source_mod.return_value.get_function.return_value = 'func'

    #call compile
    kernel_string = "__global__ void vector_add()"
    func = dev.compile("vector_add", kernel_string)

    #verify behavior
    assert dev.source_mod.call_count == 1
    assert dev.current_module is dev.source_mod.return_value
    assert func == 'func'

    assert kernel_string == list(dev.source_mod.mock_calls[0])[1][0]
    optional_args = list(dev.source_mod.mock_calls[0])[2]
    assert optional_args['code'] == 'sm_55'
    assert optional_args['arch'] == 'compute_55'


def dummy_func(a, b, block=0, grid=0):
    pass

@patch('kernel_tuner.cuda.drv')
def test_benchmark(drv):
    drv = setup_mock(drv)

    drv.Event.return_value.time_since.return_value = 0.1

    dev = cuda.CudaFunctions(0)
    args = [1, 2]
    time = dev.benchmark(dummy_func, args, (1,2), (1,2), False)
    assert time > 0

    assert dev.context.synchronize.call_count == 2*dev.iterations
    assert drv.Event.return_value.record.call_count == 2*dev.iterations
    assert drv.Event.return_value.time_since.call_count == dev.iterations


@patch('kernel_tuner.cuda.drv')
def test_copy_constant_memory_args(drv):
    drv = setup_mock(drv)

    fake_array = numpy.zeros(10).astype(numpy.float32)
    cmem_args = { 'fake_array': fake_array }

    dev = cuda.CudaFunctions(0)
    dev.current_module = Mock()
    dev.current_module.get_global.return_value = ['get_global']

    dev.copy_constant_memory_args(cmem_args)

    drv.memcpy_htod.assert_called_once_with('get_global', fake_array)
    dev.current_module.get_global.assert_called_once_with('fake_array')


@@ -627,16 +627,64 @@
 rn drv%0A%0A
+@patch('kernel_tuner.cuda.DynamicSourceModule')%0A
 @patch('
@@ -735,24 +735,27 @@
 ent_list(drv
+, _
 ):%0A    drv =
@@ -1276,23 +1276,17 @@
 le(drv, 
-src_mod
+_
 ):%0A%0A    
@@ -1367,37 +1367,8 @@
 (0)%0A
-    dev.source_mod = src_mod%0A
     
@@ -2003,16 +2003,64 @@
   pass%0A%0A
+@patch('kernel_tuner.cuda.DynamicSourceModule')%0A
 @patch('
@@ -2101,24 +2101,27 @@
 enchmark(drv
+, _
 ):%0A    drv =
@@ -2545,16 +2545,64 @@
 tions%0A%0A%0A
+@patch('kernel_tuner.cuda.DynamicSourceModule')%0A
 @patch('
@@ -2659,24 +2659,27 @@
 ory_args(drv
+, _
 ):%0A    drv =

fix mocked tests