from copy import deepcopy

import numpy as np
import pytest

import theano
import theano.tensor as tt
from theano import Mode, function
from theano.gof import Apply, generic
from theano.gof.op import PureOp
from theano.ifelse import ifelse


class IfElseIfElseIf(PureOp):
    def __init__(self, inplace=False):
        # check destroyhandler and others to ensure that a view_map with
        self.inplace = inplace
        # multiple inputs can work
        assert not self.inplace

    def make_node(self, c1, t1, c2, t2, c3, t3, f3):
        assert t1.type == f3.type
        assert t2.type == t3.type
        assert t3.type == f3.type
        return Apply(self, [c1, t1, c2, t2, c3, t3, f3], [t1.type()])

    def make_thunk(self, node, storage_map, compute_map, no_recycling, impl):

        input_computed = [compute_map[v] for v in node.inputs]
        output_computed = [compute_map[v] for v in node.outputs]
        input_registers = [storage_map[v] for v in node.inputs]
        output_registers = [storage_map[v] for v in node.outputs]

        outtype = node.outputs[0].type

        def thunk():
            if not input_computed[0][0]:
                return [0]
            else:
                truthval = input_registers[0][0]
                if truthval:
                    if not input_computed[1][0]:
                        return [1]
                    else:
                        output_computed[0][0] = 1
                        output_registers[0][0] = outtype.filter(
                            deepcopy(input_registers[1][0])
                        )
                        return []
                else:
                    if not input_computed[2][0]:
                        return [2]
                    else:
                        truthval = input_registers[2][0]
                        if truthval:
                            if not input_computed[3][0]:
                                return [3]
                            else:
                                output_computed[0][0] = 1
                                output_registers[0][0] = outtype.filter(
                                    deepcopy(input_registers[3][0])
                                )
                                return []
                        else:
                            if not input_computed[4][0]:
                                return [4]
                            else:
                                truthval = input_registers[4][0]
                                if truthval:
                                    if not input_computed[5][0]:
                                        return [5]
                                    else:
                                        output_computed[0][0] = 1
                                        output_registers[0][0] = outtype.filter(
                                            deepcopy(input_registers[5][0])
                                        )
                                        return []
                                else:
                                    if not input_computed[6][0]:
                                        return [6]
                                    else:
                                        output_computed[0][0] = 1
                                        output_registers[0][0] = outtype.filter(
                                            deepcopy(input_registers[6][0])
                                        )
                                        return []

        thunk.lazy = True
        return thunk


class NotImplementedOpException(Exception):
    pass


class NotImplementedOp(PureOp):
    def make_node(self, x):
        return Apply(self, [x], [x.type()])

    def make_thunk(self, node, storage_map, compute_map, no_recycling, impl):
        def thunk():
            raise NotImplementedOpException()

        thunk.lazy = False
        return thunk


def test_ifelse():
    a = tt.scalar()
    b = generic()
    c = generic()

    notimpl = NotImplementedOp()

    lazys = [True]
    # We need lazy to end up being True for this test.
    if theano.config.vm.lazy in [True, None]:
        lazys = [True, None]

    cloops = [True, False]

    if theano.config.cxx == "":
        cloops = [False]

    for cloop in cloops:
        for lazy in lazys:
            linker = theano.gof.vm.VM_Linker(use_cloop=cloop, lazy=lazy)
            f = function(
                [a, b, c],
                ifelse(a, notimpl(b), c),
                mode=Mode(linker=linker, optimizer="fast_run"),
            )

            with pytest.raises(NotImplementedOpException):
                f(1, "a", "b")

            assert f(0, "a", "b") == "b"


def test_nested():
    notimpl = NotImplementedOp()
    ifelseifelseif = IfElseIfElseIf()

    x1 = tt.scalar("x1")
    x2 = tt.scalar("x2")
    c1 = tt.scalar("c1")
    c2 = tt.scalar("c2")
    t1 = ifelse(c1, x1, notimpl(x2))
    t1.name = "t1"
    t2 = t1 * 10
    t2.name = "t2"
    t3 = ifelse(c2, t2, x1 + t1)
    t3.name = "t3"
    t4 = ifelseifelseif(tt.eq(x1, x2), x1, tt.eq(x1, 5), x2, c2, t3, t3 + 0.5)
    t4.name = "t4"

    linker = theano.gof.vm.VM_Linker(lazy=False)
    f = function([c1, c2, x1, x2], t4, mode=Mode(linker=linker, optimizer="fast_run"))
    with pytest.raises(NotImplementedOpException):
        f(1, 0, np.array(10, dtype=x1.dtype), 0)

    linker = theano.gof.vm.VM_Linker(lazy=True)
    f = function([c1, c2, x1, x2], t4, mode=Mode(linker=linker, optimizer="fast_run"))
    assert f(1, 0, np.array(10, dtype=x1.dtype), 0) == 20.5
