import unittest
import scmkl
import numpy as np
import anndata as ad
from test_calculate_z import create_test_adata
from test_optimize_alpha import read_h5ad


class TestGetOptimalD(unittest.TestCase):
    """
    This unittest class is used to evaluate whether scmkl.get_optimal_d 
    is working properly for binary and multiclass setups.
    """
    def test_calculate_d(self):
        """
        This funcation ensures that the calculate_d function is working 
        correctly.
        """
        msg = "D was incorrectly calculated."
        self.assertEqual(scmkl.calculate_d(10000), 222, msg)
        self.assertEqual(scmkl.calculate_d(1000), 100, msg)


    def test_get_median_size(self):
        """
        This function ensure that the median run size for multiclass 
        runs is working correctly.
        """
        x = np.arange(60000).reshape(60000,1)
        x = np.concatenate([x, x, x], axis=1)
        vars = {f'gene_{i}' for i in range(x.shape[1])}
        labs = [
            'class_1', 'class_2', 'class_2', 
            'class_1', 'class_1', 'class_3'
            ]
        labs *= 10000

        group_dict = {
            'Group1': ['gene_0', 'gene_2'],
            'Group2': ['gene_0', 'gene_1']
        }

        adata = ad.AnnData(X=x)
        adata.var_names = vars
        adata.obs['labels'] = labs

        print(scmkl.get_median_size(adata, 1.5))


if '__main__' == __name__:
    unittest.main()