from shaded.linalg_utils import null_space, projection, residue
import numpy as np
import pytest


@pytest.mark.parametrize('basis',
                         [[[1, 0, 0]],
                          [[1, 0, 0], [0, 1, 0]],
                          (2, 3),
                          (100, 200)])
def test_null_space(basis):
    if isinstance(basis, tuple):
        basis = np.random.random(basis)
    basis = np.array(basis)
    normal = null_space(basis)
    prod = np.dot(normal, basis.T)
    zero_as_prod = np.zeros(prod.shape)
    assert np.allclose(prod, zero_as_prod), "The null space is not normal to the original space"


@pytest.mark.parametrize('basis, vector',
                          [([[1, 0, 0], [0, 1, 0], [0, 0, 1]], [1, 2, 3]),
                           ([[1, 0, 0], [0, 1, 0], [0, 0, 1]], [[0, 1, 2], [3, 4, 5]])])
def test_projection(basis, vector):
    basis = np.array(basis)
    vector = np.array(vector)
    vect_proj = projection(basis, vector)
    assert np.allclose(vector, vect_proj), "The projection differs from the original vector, something is not right"


@pytest.mark.parametrize('X, span',
                         [(np.random.random((100, 3)), np.array([[1, 0, 0], [0, 1, 0]])),
                          (np.random.random((100, 4)), np.array([[1, 0, 0, 0]]))])
def test_proj_plus_res(X, span):
    assert np.allclose(projection(span.T, X) + residue(span.T, X), X)


@pytest.mark.parametrize('span_vectors, vectors',
                         [([[1, 0], [0, 1]], [[2, 3]]),
                          ([[1, 0, 0], [0, 1, 0], [0, 0, 1]], [[2, 3, 4]])])
def test_residue(span_vectors, vectors):
    span_vectors = np.array(span_vectors)
    vectors = np.array(vectors)
    assert np.allclose(residue(span_vectors, vectors), np.zeros(vectors.shape))