""" Shared functions to load/generate data """
import numpy as np
import string
import random
import math
import itertools
from impyute.dataset.corrupt import Corruptor
from impyute.util import BadInputError

def randu(bound=(0, 10), shape=(5, 5), missingness="mcar", thr=0.2, dtype="int"):
    """ Return randomly generated dataset of numbers with uniformly
    distributed values between bound[0] and bound[1]

    Parameters
    ----------
    bound:tuple (start,stop)
        Determines the range of values in the matrix. Index 0 for start
        value and index 1 for stop value. Start is inclusive, stop is
        exclusive.
    shape:tuple(optional)
        Size of the randomly generated data
    missingness: ('mcar', 'mar', 'mnar')
        Type of missingness you want in your dataset
    thr: float between [0,1]
        Percentage of missing data in generated data
    dtype: ('int','float')
        Type of data

    Returns
    -------
    numpy.ndarray
    """
    if dtype == "int":
        data = np.random.randint(bound[0], bound[1], size=shape).astype(float)
    elif dtype == "float":
        data = np.random.uniform(bound[0], bound[1], size=shape)
    corruptor = Corruptor(data, thr=thr)
    raw_data = getattr(corruptor, missingness)()
    return raw_data


def randn(theta=(0, 1), shape=(5, 5), missingness="mcar", thr=0.2, dtype="float"):
    """ Return randomly generated dataset of numbers with normally
    distributed values with given and sigma.

    Parameters
    ----------
    theta: tuple (mu, sigma)
        Determines the range of values in the matrix
    shape:tuple(optional)
        Size of the randomly generated data
    missingness: ('mcar', 'mar', 'mnar')
        Type of missingness you want in your dataset
    thr: float between [0,1]
        Percentage of missing data in generated data
    dtype: ('int','float')
        Type of data

    Returns
    -------
    numpy.ndarray
    """
    mean, sigma = theta
    data = np.random.normal(mean, sigma, size=shape)
    if dtype == "int":
        data = np.round(data)
    elif dtype == "float":
        pass
    corruptor = Corruptor(data, thr=thr)
    raw_data = getattr(corruptor, missingness)()
    return raw_data

def randc(nlevels=5, shape=(5, 5), missingness="mcar", thr=0.2):
    """ Return randomly generated dataset with uniformly distributed categorical data (alphabetic character)

    Parameters
    ----------
    :param nlevels: int
        Specify the number of different categories in the dataset
    :param shape: tuple(optional)
        Size of the randomly generated data
    :param missingness: string in ('mcar', 'mar', 'mnar')
        Type of missingness you want in your dataset
    :param thr: float between [0,1]
        Percentage of missing data in generated data
    :return:
    """
    if shape[0]*shape[1] < nlevels:
        raise 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



def test_data(mask=np.zeros((3, 3), dtype=bool)):
    """ Returns a dataset to use with tests (INTERNAL USE - FOR UNIT TESTING)

    mask: True/False array, same size as dataset
        Use True where missing values should occur and False everywhere else
    th: float between[0,1]
        Percentage of missing data in generated dataset
    """
    shape = np.shape(mask)
    data = np.reshape(np.arange(np.product(shape)), shape).astype("float")
    data[mask] = np.nan
    return data


def mnist(missingness="mcar", thr=0.2):
    """ Loads corrupted MNIST

    Parameters
    ----------
    missingness: ('mcar', 'mar', 'mnar')
        Type of missigness you want in your dataset
    th: float between [0,1]
        Percentage of missing data in generated data

    Returns
    -------
    numpy.ndarray
    """
    from sklearn.datasets import fetch_mldata
    dataset = fetch_mldata('MNIST original')
    corruptor = Corruptor(dataset.data, thr=thr)
    data = getattr(corruptor, missingness)()
    return {"X": data, "Y": dataset.target}
@@ -2420,23 +2420,16 @@
 ---%0A    
-:param 
 nlevels:
@@ -2433,16 +2433,16 @@
 ls: int%0A
+
         
@@ -2503,23 +2503,16 @@
 set%0A    
-:param 
 shape: t
@@ -2574,23 +2574,16 @@
 ata%0A    
-:param 
 missingn
@@ -2682,15 +2682,8 @@
     
-:param 
 thr:
@@ -2760,20 +2760,50 @@
 ata%0A
+%0A
     
-:r
+R
 eturn
-:
+s%0A    -------%0A    numpy.ndarray
 %0A   

update the docstring