"""Functions to get some example data
"""
from pyaddons import utils
import random


_SEED = 1984
"""The random seed for data generation (`int`)
"""

_DTYPE_MAP = {'str': str, 'int': int, 'float': float}
"""Mappings between data type names and their actual types (`dict`)
"""


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def create_data(has_header=True, ncols=10, nrows=10, use_dict=False,
                comment=None, comment_lines=10):
    """Create a test data set to work on chunking.

    Parameters
    ----------
    has_header : `bool`, optional, default: `\t`
        Write a header line.
    ncols : `int`, optional, default: `10`
        The number of columns to write.
    nrows : `int`, optional, default: `10`
        The number of rows to write.
    use_dict : `bool`, optional, default: `False`
        Should a `list` of `dict` be returned instead of a `list` of `list`.
    comment : `str`, optional, default: `NoneType`
        Should 10 comment character lines be placed before the start of the
        data. If skiplines > 0 these will be located after them.
    comment_lines : `int`, optional, default: `10`
        If comment is not None, this is the number of comment lines that will
        be added to the data.

    Returns
    -------
    test_data : `list` of (`list` or `dict`)
        The format of the test data will vary depending on the parameters. If
        ``has_header`` is ``True``, then the first element of the list will
        always be a list of strings. If use_dict is False then the remainder of
        the elements will be lists of mixed types depending on the seed. If
        use_dict is True then the remainder of the elements will be dicts with
        the header names as keys and the data as values.
    col_data_types : `list` of `str`
        The data types of the columns, if ``use_dict`` is ``False`` then this
        will be a list of data type names represented as strings, i.e.
        ``'str``. If ``use_dict`` is ``True`` then this will be a dict of
        column names as keys and data type names as values.
    """
    test_data = []

    # Add any comment lines
    if comment is not None:
        for i in range(comment_lines):
            test_data.append([f"{comment} {i}"])

    # Now add a header if we have one
    header = None
    if has_header is True:
        header = [f"column{i}" for i in range(ncols)]
        test_data.append([f"column{i}" for i in range(ncols)])

    data_funcs = [(str, get_string), (float, get_float), (int, get_integer)]
    for i in range(nrows):
        row = []
        col_data_types = []
        for j in range(ncols):
            random.seed(_SEED * j * 2)
            df = data_funcs[random.randint(0, len(data_funcs)-1)]
            row.append(df[1](_SEED * i * j))
            col_data_types.append(df[0].__name__)

        if use_dict is False:
            test_data.append([r for r in row])
        else:
            if has_header is False:
                raise ValueError("need a header for dict data")
            test_data.append(dict([(h, r) for h, r in zip(header, row)]))

    if use_dict is True:
        col_data_types = \
            dict([(h, c) for h, c in zip(header, col_data_types)])

    return test_data, col_data_types


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_string(seed=_SEED):
    """Get a text string for a random integer value between 1-5000 (`str`)
    """
    random.seed(seed)
    return utils.get_text_number(random.randint(1, 5000))


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_float(seed=_SEED):
    """Get a random float value between 0-1 (`float`)
    """
    random.seed(seed)
    return random.random()


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_integer(seed=_SEED):
    """Get a random integer value between 1-5000 (`int`)
    """
    random.seed(seed)
    return random.randint(1, 5000)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def set_row_type(row, dtypes):
    """Convert the data types for a row.

    Parameters
    ----------
    row : `list` or `dict`
        The row to convert.
    dtypes : `list` of `str`
        String representations for the data types in each column.

    Returns
    -------
    converted_row : `list` or `dict`
        The row converted to the dtypes.
    """
    if isinstance(row, list):
        return [_DTYPE_MAP[dt](i) for dt, i in zip(dtypes, row)]
    elif isinstance(row, dict):
        return dict([(k, _DTYPE_MAP[dt](v)) for dt, k, v in
                     zip(dtypes, row.keys(), row.values())])
    else:
        raise TypeError(f"can't handle: {type(row)}")
