import numpy as np
from ... import backend as T
from ...base import fold, unfold
from .._kronecker import kronecker
from ..n_mode_product import mode_dot, multi_mode_dot

def test_mode_dot():
    """Test for mode_dot (n_mode_product)"""
    X = T.tensor([[[1, 13],
                   [4, 16],
                   [7, 19],
                   [10, 22]],

                  [[2, 14],
                   [5, 17],
                   [8, 20],
                   [11, 23]],

                  [[3, 15],
                   [6, 18],
                   [9, 21],
                   [12, 24]]])
    # tensor times matrix
    U = T.tensor([[1, 3, 5],
                  [2, 4, 6]])
    true_res = T.tensor([[[22, 130],
                          [49, 157],
                          [76, 184],
                          [103, 211]],
                         [[28, 172],
                          [64, 208],
                          [100, 244],
                          [136, 280]]])
    res = mode_dot(X, U, 0)
    T.assert_array_equal(true_res, res)

    #######################
    # tensor times vector #
    #######################
    U = T.tensor([1, 2, 3, 4])
    true_res = T.tensor([[70, 190],
                         [80, 200],
                         [90, 210]])
    res = mode_dot(X, U, 1)
    T.assert_array_equal(true_res, res)

    # Using equivalence with unfolded expression
    X = T.tensor(np.random.random((2, 4, 5)))
    U = T.tensor(np.random.random((3, 4)))
    true_res = fold(T.dot(U, unfold(X, 1)), 1, (2, 3, 5))
    res = mode_dot(X, U, 1)
    assert (res.shape == (2, 3, 5))
    T.assert_array_almost_equal(true_res, res)

    #########################################
    # Test for errors that should be raised #
    #########################################
    with T.assert_raises(ValueError):
        mode_dot(X, U, 0)
    # Same test for the vector case
    with T.assert_raises(ValueError):
        mode_dot(X, U[:, 0], 0)
    # Cannot take mode product of tensor with tensor
    with T.assert_raises(ValueError):
        mode_dot(X, X, 0)

    # Test using the equivalence with unfolded expression
    X = T.tensor(np.random.random((2, 4, 5)))
    U = T.tensor(np.random.random((3, 4)))
    res = unfold(mode_dot(X, U, 1), 1)
    T.assert_array_almost_equal(T.dot(U, unfold(X, 1)), res)


def test_multi_mode_dot():
    """Test for multi_mode_dot

    Notes
    -----
    First a numerical test (ie compute by hand and check)
    Then use that the following expressions are equivalent:

    * X x_1 U_1 x ... x_n U_n
    * U_1 x unfold(X, 1) x kronecker(U_2, ..., U_n).T
    * U_1 x unfold(X x_2 U_2 x ... x_n U_n)
    """
    X = T.tensor([[1, 2],
                  [0, -1]])
    U = [T.tensor([2, 1]),
         T.tensor([-1, 1])]
    true_res = T.tensor([1])
    res = multi_mode_dot(X, U, [0, 1])
    T.assert_array_equal(true_res, res)

    X = T.tensor(np.arange(12).reshape((3, 4)))
    U = T.tensor(np.random.random((3, 5)))
    res_1 = multi_mode_dot(X, [U], modes=[0], transpose=True)
    res_2 = T.dot(U.T, X)
    T.assert_array_almost_equal(res_1, res_2)

    dims = [2, 3, 4, 5]
    X = T.tensor(np.random.randn(*dims))
    factors = [T.tensor(np.random.rand(dims[i], X.shape[i])) for i in range(X.ndim)]
    true_res = T.dot(T.dot(factors[0], unfold(X, 0)), kronecker(factors[1:]).T)
    n_mode_res = multi_mode_dot(X, factors)
    T.assert_array_almost_equal(true_res, unfold(n_mode_res, 0))
    for i in range(X.ndim):
        indices = [j for j in range(X.ndim) if j != i]
        sub_factors = [factors[j] for j in indices]
        true_res = T.dot(T.dot(factors[i], unfold(X, i)), kronecker(sub_factors).T)
        res = unfold(n_mode_res, i)
        temp = multi_mode_dot(X, sub_factors, modes=indices)
        res2 = T.dot(factors[i], unfold(temp, i))
        T.assert_equal(true_res.shape, res.shape, err_msg='shape should be {}, is {}'.format(true_res.shape, res.shape))
        T.assert_array_almost_equal(true_res, res)
        T.assert_array_almost_equal(true_res, res2)

    # Test skipping a factor
    dims = [2, 3, 4, 5]
    X = T.tensor(np.random.randn(*dims))
    factors = [T.tensor(np.random.rand(dims[i], X.shape[i])) for i in range(X.ndim)]
    res_1 = multi_mode_dot(X, factors, skip=1)
    res_2 = multi_mode_dot(X, [factors[0]] + factors[2:], modes=[0, 2, 3])
    T.assert_array_equal(res_1, res_2)