import numpy as np
from scipy.stats import multivariate_normal
from src.gaussian import GaussianModel


def test_calculate_log_likelihood_single(get_samples_parameterized):
    """
    GIVEN a sample data set
    WHEN the log likelihood is calculated using calculate_log_likelihood
    THEN the calculated log likelihood is close to multivariate_normal model's pdf log value
    :param get_samples_parameterized: Randomly generated sample with 1 point and 2 features
    """
    x = get_samples_parameterized(1, 2)
    x = np.ravel(x)
    likelihood = multivariate_normal.pdf(x, mean=[2.5, 2.5], cov=[[1, 0.5],
                                                                  [0.5, 1]])
    pca = GaussianModel(mean=[2.5, 2.5], cov=[[1, 0.5],
                                              [0.5, 1]])
    assert np.allclose(pca.calculate_log_likelihood(x), np.log(likelihood))


def test_calculate_log_likelihood_multiple(get_samples_parameterized):
    """
    GIVEN a sample data set
    WHEN the log likelihood is calculated using calculate_log_likelihood
    THEN the calculated log likelihood is close to multivariate_normal model's pdf log value
    :param get_samples_parameterized: Randomly generated sample with 10 points and 2 features
    """
    x = get_samples_parameterized(10, 2)
    likelihood = multivariate_normal.pdf(x, mean=[2.5, 2.5], cov=[[1, 0.5],
                                                                  [0.5, 1]])
    pca = GaussianModel(mean=[2.5, 2.5], cov=[[1, 0.5],
                                              [0.5, 1]])
    assert np.allclose(pca.calculate_log_likelihood(x), np.log(likelihood))
