# pylint: disable=no-self-use,invalid-name

import numpy
from numpy.testing import assert_almost_equal, assert_array_almost_equal
import keras.backend as K

from deep_qa.tensors.backend import l1_normalize
from deep_qa.tensors.masked_operations import masked_batch_dot, masked_softmax


class TestMaskedOperations:
    def test_masked_batch_dot_masks_properly(self):
        embedding_dim = 3
        a_length = 4
        b_length = 5
        batch_size = 2

        tensor_a = numpy.random.rand(batch_size, a_length, embedding_dim)
        tensor_b = numpy.random.rand(batch_size, b_length, embedding_dim)
        mask_a = numpy.ones((batch_size, a_length))
        mask_a[1, 3] = 0
        mask_b = numpy.ones((batch_size, b_length))
        mask_b[1, 2] = 0
        result = K.eval(masked_batch_dot(K.variable(tensor_a),
                                         K.variable(tensor_b),
                                         K.variable(mask_a),
                                         K.variable(mask_b)))
        assert numpy.all(result[0, :, :] != numpy.zeros((a_length, b_length)))
        assert numpy.any(result[1, 0, :] != numpy.zeros((b_length)))
        assert numpy.any(result[1, 1, :] != numpy.zeros((b_length)))
        assert numpy.any(result[1, 2, :] != numpy.zeros((b_length)))
        assert numpy.all(result[1, 3, :] == numpy.zeros((b_length)))
        assert numpy.any(result[1, :, 0] != numpy.zeros((a_length)))
        assert numpy.any(result[1, :, 1] != numpy.zeros((a_length)))
        assert numpy.all(result[1, :, 2] == numpy.zeros((a_length)))
        assert numpy.any(result[1, :, 3] != numpy.zeros((a_length)))
        assert numpy.any(result[1, :, 4] != numpy.zeros((a_length)))

        result = K.eval(masked_batch_dot(K.variable(tensor_a),
                                         K.variable(tensor_b),
                                         None,
                                         None))
        assert numpy.all(result[0, :, :] != numpy.zeros((a_length, b_length)))
        assert numpy.all(result[1, :, :] != numpy.zeros((a_length, b_length)))

        result = K.eval(masked_batch_dot(K.variable(tensor_a),
                                         K.variable(tensor_b),
                                         K.variable(mask_a),
                                         None))
        assert numpy.all(result[0, :, :] != numpy.zeros((a_length, b_length)))
        assert numpy.any(result[1, 0, :] != numpy.zeros((b_length)))
        assert numpy.any(result[1, 1, :] != numpy.zeros((b_length)))
        assert numpy.any(result[1, 2, :] != numpy.zeros((b_length)))
        assert numpy.all(result[1, 3, :] == numpy.zeros((b_length)))
        assert numpy.any(result[1, :, 0] != numpy.zeros((a_length)))
        assert numpy.any(result[1, :, 1] != numpy.zeros((a_length)))
        assert numpy.any(result[1, :, 2] != numpy.zeros((a_length)))
        assert numpy.any(result[1, :, 3] != numpy.zeros((a_length)))
        assert numpy.any(result[1, :, 4] != numpy.zeros((a_length)))

        result = K.eval(masked_batch_dot(K.variable(tensor_a),
                                         K.variable(tensor_b),
                                         None,
                                         K.variable(mask_b)))
        assert numpy.all(result[0, :, :] != numpy.zeros((a_length, b_length)))
        assert numpy.any(result[1, 0, :] != numpy.zeros((b_length)))
        assert numpy.any(result[1, 1, :] != numpy.zeros((b_length)))
        assert numpy.any(result[1, 2, :] != numpy.zeros((b_length)))
        assert numpy.any(result[1, 3, :] != numpy.zeros((b_length)))
        assert numpy.any(result[1, :, 0] != numpy.zeros((a_length)))
        assert numpy.any(result[1, :, 1] != numpy.zeros((a_length)))
        assert numpy.all(result[1, :, 2] == numpy.zeros((a_length)))
        assert numpy.any(result[1, :, 3] != numpy.zeros((a_length)))
        assert numpy.any(result[1, :, 4] != numpy.zeros((a_length)))

    def test_masked_batch_dot_handles_uneven_tensors(self):
        # We're going to test masked_batch_dot with tensors of shape (batch_size, a_length,
        # embedding_dim) and (batch_size, embedding_dim).  The result should have shape
        # (batch_size, a_length).
        embedding_dim = 3
        a_length = 5
        batch_size = 2

        tensor_a = numpy.random.rand(batch_size, a_length, embedding_dim)
        tensor_b = numpy.random.rand(batch_size, embedding_dim)
        mask_a = numpy.ones((batch_size, a_length))
        mask_a[0, 3] = 0
        mask_b = numpy.ones((batch_size,))
        mask_b[1] = 0
        result = K.eval(masked_batch_dot(K.variable(tensor_a),
                                         K.variable(tensor_b),
                                         K.variable(mask_a),
                                         K.variable(mask_b)))
        assert result[0, 0] != 0
        assert result[0, 1] != 0
        assert result[0, 2] != 0