import tensorly as tl

import pytest
import numpy as np
import itertools

from .._tt_cross import tensor_train_cross
from ....tt_tensor import tt_to_tensor
from tensorly.testing import assert_

skip_if_backend = pytest.mark.skipif(
    tl.get_backend() in ("tensorflow", "jax", "cupy"),
    reason=f"Operation not supported in {tl.get_backend()}",
)


@skip_if_backend
def test_tensor_train_cross_1():
    """Test for tensor-train"""

    ## Test 1

    # Create tensor with random elements
    d = 3
    n = 4
    tensor = np.arange(n**d, dtype=float).reshape((n,) * d)
    tensor = tl.tensor(tensor)

    tensor_shape = tensor.shape

    # Find TT decomposition of the tensor
    rank = [1, 3, 3, 1]
    factors = tensor_train_cross(
        tensor, rank, tol=1e-5, n_iter_max=10, random_state=1234
    )
    assert len(factors) == d, "Number of factors should be 4, currently has " + str(
        len(factors)
    )

    # Check that the ranks are correct and that the second mode of each factor
    # has the correct number of elements
    r_prev_iteration = 1
    for k in range(d):
        (r_prev_k, n_k, r_k) = factors[k].shape
        assert tensor_shape[k] == n_k, (
            "Mode 1 of factor "
            + str(k)
            + "needs "
            + str(tensor_shape[k])
            + " dimensions, currently has "
            + str(n_k)
        )
        assert r_prev_k == r_prev_iteration, " Incorrect ranks of factors "
        r_prev_iteration = r_k


@skip_if_backend
def test_tensor_train_cross_2():
    """Test for tensor-train"""
    rng = tl.check_random_state(1234)

    ## Test 2
    # Create tensor with random elements
    tensor = tl.tensor(rng.random_sample([3, 4, 5, 6, 2, 10]))

    # Find TT decomposition of the tensor
    rank = [1, 2, 2, 3, 2, 2, 1]
    factors = tensor_train_cross(tensor, rank, random_state=rng)

    for k in range(6):
        (r_prev, n_k, r_k) = factors[k].shape

        first_error_message = (
            "TT rank " + str(k) + " is greater than the maximum allowed "
        )
        first_error_message += str(r_prev) + " > " + str(rank[k])
        assert r_prev <= rank[k], first_error_message

        first_error_message = (
            "TT rank " + str(k + 1) + " is greater than the maximum allowed "
        )
        first_error_message += str(r_k) + " > " + str(rank[k + 1])
        assert r_k <= rank[k + 1], first_error_message


@skip_if_backend
@pytest.mark.skipif(
    tl.get_backend() in ("mxnet"),
    reason=f"MXNet bug in advanced indexing (Issue 18919).",
)
def test_tensor_train_cross_3():
    """Test for tensor-train"""
    rng = tl.check_random_state(1234)

    ## Test 3
    tol = 10e-5
    tensor = tl.tensor(rng.random_sample([3, 3, 3]))
    factors = tensor_train_cross(tensor, (1, 3, 3, 1), random_state=rng)
    reconstructed_tensor = tt_to_tensor(factors)
    error = tl.norm(reconstructed_tensor - tensor, 2)
    error /= tl.norm(tensor, 2)
    assert_(error < tol, "norm 2 of reconstruction higher than tol")


@skip_if_backend
def test_tensor_train_cross_4():
    """Test for tensor-train"""

    # TEST 4
    # Random tensor is not really compress-able. Test on a tensor as values of a function

    def getEquispaceGrid(n_dim, rng, subdivisions):
        """
        Returns a grid of equally-spaced points in the specified number of dimensions

        n_dim       : The number of dimensions to construct the tensor grid in
        rng         : The maximum dimension coordinate (grid starts at 0)
        subdivisions: Number of subdivisions of the grid to construct
        """

        return np.array(
            [
                np.array(range(subdivisions + 1)) * rng * 1.0 / subdivisions
                for i in range(n_dim)
            ]
        )

    def evaluateGrid(grid, fcn):
        """
        Loops over a grid in specified order and computes the specified function at each
        point in the grid, returning a list of computed values.
        """
        d, n = grid.shape
        values = np.zeros(len(grid[0]) ** len(grid))
        idx = 0
        for permutation in itertools.product(range(len(grid[0])), repeat=len(grid)):
            pt = np.array([grid[i][permutation[i]] for i in range(len(permutation))])
            values[idx] = fcn(pt)
            idx += 1

        return values.reshape((n,) * d)

    def func(X):
        return sum(X) ** 3

    tol = 1e-3
    n = 10
    d = 4
    rng = 1
    grid = getEquispaceGrid(d, rng, n)
    value = evaluateGrid(grid, func)
    value = tl.tensor(value)

    # Find TT decomposition of the tensor
    rank = [1, 4, 4, 4, 1]
    factors = tensor_train_cross(value, rank, tol=tol, random_state=rng)

    approx = tt_to_tensor(factors)
    error = tl.norm(approx - value, 2)
    error /= tl.norm(value, 2)

    print(error)
    assert_(error < 1e-5, "norm 2 of reconstruction higher than tol")
