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)
                self.run_case(testname, inputs)

        def run_case(self, testname, inputs):
            inputs_ref = [theano.shared(inp) for inp in inputs]
            inputs_tst = [theano.shared(inp) for inp in inputs]

            try:
                node_ref = safe_make_node(self.op, *inputs_ref)
                node_tst = safe_make_node(self.op, *inputs_tst)
            except Exception as exc:
                err_msg = (
                    "Test %s::%s: Error occurred while making " "a node with inputs %s"
                ) % (self.gpu_op, testname, inputs)
                exc.args += (err_msg,)
                raise

            try:
                f_ref = inplace_func([], node_ref.outputs, mode=mode_nogpu)
                f_tst = inplace_func([], node_tst.outputs, mode=mode_gpu)
            except Exception as exc:
                err_msg = (
                    "Test %s::%s: Error occurred while trying to " "make a Function"
                ) % (self.gpu_op, testname)
                exc.args += (err_msg,)
                raise

            self.assertFunctionContains1(f_tst, self.gpu_op)

            ref_e = None
            try:
                expecteds = f_ref()
            except Exception as exc:
                ref_e = exc

            try:
                variables = f_tst()
            except Exception as exc:
                if ref_e is None:
                    err_msg = (
                        "Test %s::%s: exception when calling the " "Function"
                    ) % (self.gpu_op, testname)
                    exc.args += (err_msg,)
                    raise
                else:
                    # if we raised an exception of the same type we're good.
                    if isinstance(exc, type(ref_e)):
                        return
                    else:
                        err_msg = (
                            "Test %s::%s: exception raised during test "
                            "call was not the same as the reference "
                            "call (got: %s, expected %s)"
                            % (self.gpu_op, testname, type(exc), type(ref_e))
                        )
                        exc.args += (err_msg,)
                        raise

            for i, (variable, expected) in enumerate(zip(variables, expecteds)):
                condition = (
                    variable.dtype != expected.dtype
                    or variable.shape != expected.shape
                    or not TensorType.values_eq_approx(variable, expected)
                )
                assert not condition, (
                    "Test %s::%s: Output %s gave the wrong "
                    "value. With inputs %s, expected %s "
                    "(dtype %s), got %s (dtype %s)."
                    % (
                        self.op,
                        testname,
                        i,
                        inputs,
                        expected,
                        expected.dtype,
                        variable,
                        variable.dtype,
                    )
                )

            for description, check in self.checks.items():
                assert check(inputs, variables), (
                    "Test %s::%s: Failed check: %s " "(inputs were %s, ouputs were %s)"
                ) % (self.op, testname, description, inputs, variables)

    Checker.__name__ = name
    if hasattr(Checker, "__qualname__"):
        Checker.__qualname__ = name
    return Checker


def test_transfer_cpu_gpu():
    a = tt.fmatrix("a")
    g = GpuArrayType(dtype="float32", broadcastable=(False, False))("g")

    av = np.asarray(rng.rand(5, 4), dtype="float32")
    gv = gpuarray.array(av, context=get_context(test_ctx_name))

    f = theano.function([a], GpuFromHost(test_ctx_name)(a))
    fv = f(av)
    assert GpuArrayType.values_eq(fv, gv)

    f = theano.function([g], host_from_gpu(g))
    fv = f(gv)
    assert 
np.all(fv == av)