import numpy as np
import pytest

import theano
import theano.tensor as tt

# Don't import test classes otherwise they get tested as part of the file
from tests import unittest_tools as utt
from tests.gpuarray.config import mode_with_gpu, mode_without_gpu, test_ctx_name
from tests.tensor.test_basic import (
    TestAlloc,
    TestComparison,
    TestJoinAndSplit,
    TestReshape,
)
from tests.tensor.utils import rand, safe_make_node
from theano.gpuarray.basic_ops import (
    GpuAlloc,
    GpuAllocEmpty,
    GpuContiguous,
    GpuEye,
    GpuFromHost,
    GpuJoin,
    GpuReshape,
    GpuSplit,
    GpuToGpu,
    GpuTri,
    HostFromGpu,
    gpu_contiguous,
    gpu_join,
    host_from_gpu,
)
from theano.gpuarray.elemwise import GpuDimShuffle, GpuElemwise
from theano.gpuarray.subtensor import GpuSubtensor
from theano.gpuarray.type import GpuArrayType, get_context, gpuarray_shared_constructor
from theano.tensor import TensorType
from theano.tensor.basic import alloc


pygpu = pytest.importorskip("pygpu")
gpuarray = pygpu.gpuarray

utt.seed_rng()
rng = np.random.RandomState(seed=utt.fetch_seed())


def inplace_func(
    inputs,
    outputs,
    mode=None,
    allow_input_downcast=False,
    on_unused_input="raise",
    name=None,
):
    if mode is None:
        mode = mode_with_gpu
    return theano.function(
        inputs,
        outputs,
        mode=mode,
        allow_input_downcast=allow_input_downcast,
        accept_inplace=True,
        on_unused_input=on_unused_input,
        name=name,
    )


def fake_shared(value, name=None, strict=False, allow_downcast=None, **kwargs):
    from theano.tensor.sharedvar import scalar_constructor, tensor_constructor

    for c in (gpuarray_shared_constructor, tensor_constructor, scalar_constructor):
        try:
            return c(
                value, name=name, strict=strict, allow_downcast=allow_downcast, **kwargs
            )
        except TypeError:
            continue


def rand_gpuarray(*shape, **kwargs):
    r = rng.rand(*shape) * 2 - 1
    dtype = kwargs.pop("dtype", theano.config.floatX)
    cls = kwargs.pop("cls", None)
    if len(kwargs) != 0:
        raise TypeError("Unexpected argument %s", list(kwargs.keys())[0])
    return gpuarray.array(r, dtype=dtype, cls=cls, context=get_context(test_ctx_name))


def makeTester(
    name,
    op,
    gpu_op,
    cases,
    checks=None,
    mode_gpu=mode_with_gpu,
    mode_nogpu=mode_without_gpu,
    skip=False,
    eps=1e-10,
):
    if checks is None:
        checks = {}

    _op = op
    _gpu_op = gpu_op
    _cases = cases
    _skip = skip
    _checks = checks

    class Checker(utt.OptimizationTestMixin):
        op = staticmethod(_op)
        gpu_op = staticmethod(_gpu_op)
        cases = _cases
        skip = _skip
        checks = _checks

        def setup_method(self):
            eval(self.__class__.__module__ + "." + self.__class__.__name__)

        def test_all(self):
            if skip:
                pytest.skip(skip)

            for testname, inputs in cases.items():
                for _ in range(len(inputs)):
                    if type(inputs[_]) is float:
                        inputs[_] = 
np.asarray(inputs[_], dtype=theano.config.floatX)