import numpy as np

from ...random import cp_tensor, check_random_state
from ..robust_decomposition import robust_pca
from ... import backend as T

def test_RPCA():
    """Test for RPCA"""
    tol = 1e-5

    sample = np.array([[1, 2, 3, 4],
                       [2, 4, 6, 8]])
    clean = T.tensor(np.vstack([sample[None, ...]]*100))
    noise_probability = 0.05
    rng = check_random_state(12345)
    noise = T.tensor(rng.choice([0., 100., -100.], size=clean.shape, replace=True,
                                p=[1 - noise_probability, noise_probability/2, noise_probability/2]))
    tensor = clean + noise
    clean_pred, noise_pred = robust_pca(tensor, mask=None, reg_E=0.4, mu_max=10e12, 
                                        learning_rate=1.2,
                                        n_iter_max=200, tol=tol, verbose=True)
    # check recovery 
    T.assert_array_almost_equal(tensor, clean_pred+noise_pred, decimal=tol)
    # check low rank recovery
    T.assert_array_almost_equal(clean, clean_pred, decimal=1)
    # Check for sparsity of the gross error
    assert T.sum(noise_pred > 0.01) == T.sum(noise > 0.01)
    # check sparse gross error recovery
    T.assert_array_almost_equal(noise, noise_pred, decimal=1)

    ############################
    # Test with missing values #
    ############################
    corrupted_clean = T.to_numpy(clean)
    # Add some corruption (missing values, replaced by ones)
    mask = rng.choice([0, 1], clean.shape, replace=True, p=[0.05, 0.95])
    corrupted_clean[mask == 0] = 1
    corrupted_clean = T.tensor(corrupted_clean)
    tensor = T.tensor(corrupted_clean) + noise
    mask = T.tensor(mask)
    # Decompose the tensor
    clean_pred, noise_pred = robust_pca(tensor, mask=mask, reg_E=0.4, mu_max=10e12, 
                                        learning_rate=1.2,
                                        n_iter_max=200, tol=tol, verbose=True)
    # check recovery 
    T.assert_array_almost_equal(tensor, clean_pred+noise_pred, decimal=tol)
    # check low rank recovery
    T.assert_array_almost_equal(corrupted_clean*mask, clean_pred*mask, decimal=1)
    # check sparse gross error recovery
    T.assert_array_almost_equal(noise*mask, noise_pred*mask, decimal=1)

    # Check for recovery of the corrupted/missing part
    mask = 1 - mask
    error = T.norm((clean*mask - clean_pred*mask), 2)/T.norm(clean*mask, 2)
    T.assert_(error <= 10e-3)