import numpy as np

from ... import backend as T
from ..proximal import (
    svd_thresholding,
    smoothness_prox,
    soft_thresholding,
    hals_nnls,
    fista,
    active_set_nnls,
    procrustes,
    hard_thresholding,
    soft_sparsity_prox,
    simplex_prox,
    normalized_sparsity_prox,
    monotonicity_prox,
    unimodality_prox,
    l2_prox,
    l2_square_prox,
    admm,
)
from ...testing import assert_, assert_array_equal, assert_array_almost_equal
from tensorly import tensor_to_vec, truncated_svd
import pytest

# Author: Jean Kossaifi
skip_tensorflow = pytest.mark.skipif(
    (T.get_backend() == "tensorflow"),
    reason=f"Indexing with list not supported in TensorFlow",
)


def test_smoothness():
    """Test for smoothness operator"""
    tensor = T.tensor([[1, 2, 1.5], [4, -6, -0.5], [0.2, 1.2, -3.4]])
    copy_tensor = T.copy(tensor)
    parameter = 0.2
    tol = 0.5
    res = smoothness_prox(copy_tensor, parameter)
    true_res = T.tensor([[1.2, 0.83, 0.98], [3.11, -4.12, 0.57], [0.58, 0.26, -2.51]])
    error = T.norm(true_res - res, 2) / T.norm(true_res, 2)
    assert_(error < tol)
    # Check that we did not change the original tensor
    assert_array_equal(copy_tensor, tensor)


def test_hard_thresholding():
    """Test for hard_thresholding operator"""
    tensor = T.tensor([[1, 2, 1.5], [4, -6, -0.5], [0.2, 1.2, -3.4]])
    copy_tensor = T.copy(tensor)
    threshold = 2
    res = hard_thresholding(tensor, threshold)
    true_res = T.tensor([[0, 0, 0], [4, -6, 0], [0, 0, 0]])
    assert_array_almost_equal(true_res, res)
    # Check that we did not change the original tensor
    assert_array_equal(copy_tensor, tensor)


def test_soft_sparsity():
    """Test for soft_sparsity operator"""
    tensor = T.tensor([[0.5, 1.3, 4.5], [0.8, 0.3, 2]])
    threshold = 2
    res = soft_sparsity_prox(tensor, threshold)
    true_res = T.tensor([[0.85, 1.5, 2.0], [1.15, 0.5, 0.0]])
    assert_array_almost_equal(true_res, res)


def test_simplex():
    """Test for simplex operator"""
    # 1d Tensor
    tensor = T.tensor([0.4, 0.5, 0.6])
    res = simplex_prox(tensor, 1)
    true_res = T.tensor([0.23, 0.33, 0.43])
    assert_array_almost_equal(true_res, res, decimal=2)

    # 2d Tensor
    tensor = T.tensor([[0.4, 1.5, 1], [0.5, 2, 3], [0.6, 0.3, 2.9]])
    res = simplex_prox(tensor, 1)
    true_res = T.tensor([[0.23, 0.25, 0], [0.33, 0.75, 0.55], [0.43, 0, 0.45]])
    assert_array_almost_equal(true_res, res, decimal=2)


def test_normalized_sparsity():
    """Test for normalized_sparsity operator"""
    tensor = T.tensor([2, 3, 4])
    res = normalized_sparsity_prox(tensor, 2)
    true_res = T.tensor([0, 0.6, 0.8])
    assert_array_almost_equal(true_res, res)


def test_monotonicity():
    """Test for monotonicity operator"""
    tensor = T.tensor(np.random.rand(20, 10))
    # Monotone increasing
    tensor_monoton = monotonicity_prox(tensor)
    assert_(np.all(np.diff(tensor_monoton, axis=0) >= 0))
    # Monotone decreasing
    tensor_monoton = monotonicity_prox(tensor, decreasing=True)
    assert_(np.all(np.diff(tensor_monoton, axis=0) <= 0))


def test_unimodality():
    """Test for unimdality operator"""
    tensor = T.tensor(np.random.rand(10, 10))
    tensor_unimodal = unimodality_prox(tensor)
    for i in range(10):
        max_location = T.argmax(tensor_unimodal[:, i])
        if max_location == 0:
            assert_(np.all(np.diff(tensor_unimodal[:, i], axis=0) <= 0))
        elif max_location == (T.shape(tensor)[0] - 1):
            assert_(np.all(np.diff(tensor_unimodal[:, i], axis=0) >= 0))
        else:
            assert_(
                np.all(np.diff(tensor_unimodal[: int(max_location), i], axis=0) >= 0)
            )
            assert_(
                np.all(np.diff(tensor_unimodal[int(max_location) :, i], axis=0) <= 0)
            )


def test_l2_prox():
    """Test for l2 prox operator"""
    tensor = T.tensor([2, 4, 4])
    res = l2_prox(tensor, 3)
    true_res = T.tensor([1, 2, 2])
    assert_array_almost_equal(true_res, res)


def test_squared_l2_prox():
    """Test for squared l2 prox operator"""
    tensor = T.tensor([3, 6, 9])
    res = l2_square_prox(tensor, 0.5)
    true_res = T.tensor([1.5, 3, 4.5])
    assert_array_almost_equal(true_res, res)


def test_soft_thresholding():
    """Test for shrinkage"""
    # small test
    tensor = T.tensor(
        [
            [[1, 2, 3], [4.3, -1.2, 3]],
            [[0.5, -5, -1.3], [1.2, 3.7, -9]],
            [[-2, 0, 1.0], [0.5, -0.5, 1.1]],
        ]
    )
    threshold = 1.1
    copy_tensor = T.copy(tensor)
    res = soft_thresholding(tensor, threshold)
    true_res = T.tensor(
        [
            [[0, 0.9, 1.9], [3.2, -0.1, 1.9]],
            [[0, -3.9, -0.2], [0.1, 2.6, -7.9]],
            [[-0.9, 0, 0], [0, -0, 0]],
        ]
    )
    # account for floating point errors: np array have a precision of around 2e-15
    # check np.finfo(np.float64).eps
    assert_array_almost_equal(true_res, res)
    # Check that we did not change the original tensor
    assert_array_equal(copy_tensor, tensor)

    # Another test
    tensor = T.tensor([[1, 2, 1.5], [4, -6, -0.5], [0.2, 1.02, -3.4]])
    copy_tensor = T.copy(tensor)
    threshold = 1.1
    true_res = T.tensor([[0, 0.9, 0.4], [2.9, -4.9, 0], [0, 0, -2.3]])
    res = soft_thresholding(tensor, threshold)
    assert_array_almost_equal(true_res, res)
    assert_array_equal(copy_tensor, tensor)

    # Test with missing values
    tensor = T.tensor([[1, 2, 1.5], [4, -6, -0.5], [0.2, 1.02, -3.4]])
    copy_tensor = T.copy(tensor)
    mask = T.tensor([[0, 1, 1], [1, 0, 1], [1, 1, 0]])
    threshold = 1.1 * mask
    true_res = T.tensor([[1, 0.9, 0.4], [2.9, -6, 0], [0, 0, -3.4]])
    res = soft_thresholding(tensor, threshold)
    assert_array_almost_equal(true_res, res)
    assert_array_equal(copy_tensor, tensor)


def test_svd_thresholding():
    """Test for singular_value_thresholding operator"""
    U = T.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
    singular_values = T.tensor([0.4, 2.1, -2])
    tensor = T.dot(U, T.reshape(singular_values, (-1, 1)) * T.transpose(U))
    shrinked_singular_values = T.tensor([0, 1.6, -1.5])
    true_res = T.dot(U, T.reshape(shrinked_singular_values, (-1, 1)) * T.transpose(U))
    res = svd_thresholding(tensor, 0.5)
    assert_array_almost_equal(true_res, res)


def test_procrustes():
    """Test for procrustes operator"""
    U = T.tensor(np.random.rand(20, 10))
    S, _, V = truncated_svd(U, n_eigenvecs=min(U.shape))
    true_res = T.dot(S, V)
    res = procrustes(U)
    assert_array_almost_equal(true_res, res)


def test_hals_nnls():
    """Test for hals_nnls operator"""
    a = T.tensor(np.random.rand(20, 10))
    true_res = T.tensor(np.random.rand(10, 1))
    b = T.dot(a, true_res)
    atb = T.dot(T.transpose(a), b)
    ata = T.dot(T.transpose(a), a)
    x_hals = hals_nnls(atb, ata)[0]
    assert_array_almost_equal(true_res, x_hals, decimal=2)


def test_fista():
    """Test for fista operator"""
    a = T.tensor(np.random.rand(20, 10))
    true_res = T.tensor(np.random.rand(10, 1))
    b = T.dot(a, true_res)
    atb = T.dot(T.transpose(a), b)
    ata = T.dot(T.transpose(a), a)
    x_fista = fista(atb, ata, tol=10e-16, n_iter_max=5000)
    assert_array_almost_equal(true_res, x_fista, decimal=2)


def test_admm():
    """Test for admm operator"""
    a = T.tensor(np.random.rand(20, 10))
    true_res = T.tensor(np.random.rand(10, 10))
    b = T.dot(a, true_res)
    atb = T.dot(T.transpose(a), b)
    ata = T.dot(T.transpose(a), a)
    dual = T.zeros(T.shape(atb))
    x_init = T.zeros(T.shape(atb))
    x_admm, _, _ = admm(T.transpose(atb), T.transpose(ata), x=x_init, dual_var=dual)
    assert_array_almost_equal(true_res, T.transpose(x_admm), decimal=2)


@skip_tensorflow
def test_active_set_nnls():
    """Test for active_set_nnls operator"""
    a = T.tensor(np.random.rand(20, 10))
    true_res = T.tensor(np.random.rand(10, 1))
    b = T.dot(a, true_res)
    atb = T.dot(T.transpose(a), b)
    ata = T.dot(T.transpose(a), a)
    x_as = active_set_nnls(tensor_to_vec(atb), ata)
    x_as = T.reshape(x_as, T.shape(atb))
    assert_array_almost_equal(true_res, x_as, decimal=2)
