# Path: impyute/dataset/base.py
# def randc(nlevels=5, shape=(5, 5), missingness="mcar", thr=0.2):
#     """ Return randomly generated dataset with uniformly distributed categorical data (alphabetic character)
# 
#     Parameters
#     ----------
#     nlevels: int
#         Specify the number of different categories in the dataset
#     shape: tuple(optional)
#         Size of the randomly generated data
#     missingness: string in ('mcar', 'mar', 'mnar')
#         Type of missingness you want in your dataset
#     thr: float between [0,1]
#         Percentage of missing data in generated data
# 
#     Returns
#     -------
#     numpy.ndarray
#     """
#     if shape[0]*shape[1] < nlevels:
#         raise error.BadInputError("nlevel exceeds the size of desired dataset. Please decrease the nlevel or increase the shape")
# 
#     length = len(string.ascii_lowercase)
#     n_fold = int(math.floor(math.log(nlevels, length)))
#     cat_pool = list(string.ascii_lowercase)
# 
#     # when nlevel > 26, the alphabetical character is used up, need to generate extra strings as categorical data
#     if n_fold > 0:
#         for i in range(2, n_fold+2):
#             pool_candidate = list(itertools.product(string.ascii_lowercase, repeat=i))
#             cat_pool.extend([''.join(w) for w in pool_candidate])
#             if len(cat_pool) > nlevels:
#                 break
# 
#     cat = random.sample(cat_pool, nlevels)
#     data = np.random.choice(cat, shape, replace=True)
# 
#     # make sure the data frame has nlevel different categories
#     while len(np.unique(data)) != nlevels:
#         data = np.random.choice(cat, shape, replace=True)
# 
#     corruptor = Corruptor(data, thr=thr, dtype=np.str)
#     raw_data = getattr(corruptor, missingness)()
#     return raw_data
# 
# Path: impyute/ops/error.py
# class BadInputError(Exception):
# class BadOutputError(Exception):


import numpy as np
import pytest
from impyute.dataset.base import randc
from impyute.ops import error

def test_raise_error_nlevel_exceed_shape():
    with pytest.raises(error.BadInputError) as e:
        randc(shape=(2, 2))
    expected = "nlevel exceeds the size of desired dataset. Please decrease the nlevel or increase the shape"
    assert str(e.value) == expected

@pytest.mark.parametrize("nlevels, shape", [(5, (5,5)), (9, (3,4)), (100, (20, 20))])
def test_nlevel_categories(nlevels, shape):
    """ideally the returned matrix should have nlevel+1 different categories,
    +1 because the Corrupt class introduce np.nan however, if the missing value
    introduced by Corrupt class happens to replace a group of categories, the
    unique category number would be < nlevel + 1
    """

# Path: impyute/dataset/base.py
# def randc(nlevels=5, shape=(5, 5), missingness="mcar", thr=0.2):
#     """ Return randomly generated dataset with uniformly distributed categorical data (alphabetic character)
# 
#     Parameters
#     ----------
#     nlevels: int
#         Specify the number of different categories in the dataset
#     shape: tuple(optional)
#         Size of the randomly generated data
#     missingness: string in ('mcar', 'mar', 'mnar')
#         Type of missingness you want in your dataset
#     thr: float between [0,1]
#         Percentage of missing data in generated data
# 
#     Returns
#     -------
#     numpy.ndarray
#     """
#     if shape[0]*shape[1] < nlevels:
#         raise error.BadInputError("nlevel exceeds the size of desired dataset. Please decrease the nlevel or increase the shape")
# 
#     length = len(string.ascii_lowercase)
#     n_fold = int(math.floor(math.log(nlevels, length)))
#     cat_pool = list(string.ascii_lowercase)
# 
#     # when nlevel > 26, the alphabetical character is used up, need to generate extra strings as categorical data
#     if n_fold > 0:
#         for i in range(2, n_fold+2):
#             pool_candidate = list(itertools.product(string.ascii_lowercase, repeat=i))
#             cat_pool.extend([''.join(w) for w in pool_candidate])
#             if len(cat_pool) > nlevels:
#                 break
# 
#     cat = random.sample(cat_pool, nlevels)
#     data = np.random.choice(cat, shape, replace=True)
# 
#     # make sure the data frame has nlevel different categories
#     while len(np.unique(data)) != nlevels:
#         data = np.random.choice(cat, shape, replace=True)
# 
#     corruptor = Corruptor(data, thr=thr, dtype=np.str)
#     raw_data = getattr(corruptor, missingness)()
#     return raw_data
# 
# Path: impyute/ops/error.py
# class BadInputError(Exception):
# class BadOutputError(Exception):


# Path: test/dataset/test_randc.py
import numpy as np
import pytest
from impyute.dataset.base import randc
from impyute.ops import error


def test_raise_error_nlevel_exceed_shape():
    with pytest.raises(error.BadInputError) as e:
        randc(shape=(2, 2))
    expected = "nlevel exceeds the size of desired dataset. Please decrease the nlevel or increase the shape"
    assert str(e.value) == expected

@pytest.mark.parametrize("nlevels, shape", [(5, (5,5)), (9, (3,4)), (100, (20, 20))])
def test_nlevel_categories(nlevels, shape):
    """ideally the returned matrix should have nlevel+1 different categories,
    +1 because the Corrupt class introduce np.nan however, if the missing value
    introduced by Corrupt class happens to replace a group of categories, the
    unique category number would be < nlevel + 1
    """

    dataframe = randc(nlevels, shape)