from ..bclr_one import BayesCC
import pytest
import numpy as np

X = np.load("bclr_one.npy")

prior_meanBAD1 = [0]
prior_meanBAD2 = [[3,4], [1,2]]
prior_covBAD1 = np.array([[3, 1], [2,2]])
prior_covBAD2 = np.array([[1, 1], [1,0.999]])
prior_covBAD3 = np.array([[1,0,0], [0,1,0], [0,0,1]])
prior_covBAD4 = np.array([[1,0], [1,0], [0,1]])
prior_mean = [0,0]
prior_cov = np.array([[1,0], [0,1]])

@pytest.mark.parametrize('prior_meanBAD', [prior_meanBAD1, prior_meanBAD2])
def test_bad_mean(prior_meanBAD):
    with pytest.raises(ValueError) as e_bad_mean:
        BayesCC(X, prior_mean = prior_meanBAD, prior_cov = prior_cov, n_iter=1000)
    
    assert e_bad_mean.type is ValueError

@pytest.mark.parametrize('prior_covBAD', [prior_covBAD1, prior_covBAD2, prior_covBAD3, prior_covBAD4])
def test_bad_cov(prior_covBAD):
    with pytest.raises(ValueError) as e_bad_cov:
        BayesCC(X, prior_mean = prior_mean, prior_cov = prior_covBAD, n_iter=1000)
    
    assert e_bad_cov.type is ValueError

bclr_good = BayesCC(X, prior_mean = prior_mean, prior_cov = prior_cov, n_iter=1000)
bclr_good.fit()

def test_cp():
    bclr_good.transform()
    assert int(bclr_good.post_k_mode) == 200
    
def test_plot():
    plot_res = bclr_good.plot_k()
    assert plot_res is None
