import pandas as pd

from .encoding import Encoding
from .autovizwidget import AutoVizWidget


def infer_vegalite_type(data):
    """
    From an array-like input, infer the correct vega typecode
    ('O', 'N', 'Q', or 'T')
    Parameters
    ----------
    data: Numpy array or Pandas Series
    """

    typ = pd.lib.infer_dtype(data)

    if typ in ['floating', 'mixed-integer-float', 'integer',
               'mixed-integer', 'complex']:
        typecode = 'Q'
    elif typ in ['string', 'bytes', 'categorical', 'boolean', 'mixed', 'unicode']:
        typecode = 'N'
    elif typ in ['datetime', 'datetime64', 'timedelta',
                 'timedelta64', 'date', 'time', 'period']:
        typecode = 'T'
    else:
        typecode = 'N'

    return typecode


def _validate_custom_order(order):
    assert len(order) == 4
    list_to_check = list(order)
    list_to_check.sort()
    assert list_to_check == ['N', 'O', 'Q', 'T']


def _classify_data_by_type(data, order, skip=None):
    """Get O, N, Q, or T vegalite type for all columns in data except if in skip."""
    if skip is None:
        skip = []

    d = dict()
    for typ in order:
        d[typ] = []

    for column_name in data:
        if column_name not in skip:
            typ = infer_vegalite_type(data[column_name])
            d[typ].append(column_name)

    return d


def select_x(data, order=None):
    """
    Helper function that does a best effort of selecting an automatic x axis.
    Returns None if it cannot find x axis.
    """
    if data is None:
        return None

    if len(data) < 1:
        return None

    if order is None:
        order = ['T', 'O', 'N', 'Q']
    else:
        _validate_custom_order(order)

    d = _classify_data_by_type(data, order)

    chosen_x = None
    for typ in order:
        if len(d[typ]) >= 1:
            chosen_x = d[typ][0]
            break

    return chosen_x


def select_y(data, x_name, order=None, aggregator=None):
    """
    Helper function that does a best effort of selecting an automatic y axis.
    It won't set the same axis that x is set to again.
    Returns None if it cannot find y axis.
    """
    if data is None:
        return None

    if len(data) < 2:
        return None

    if x_name is None:
        return None

    if order is None:
        order = ['Q', 'O', 'N', 'T']
    else:
        _validate_custom_order(order)

    d = _classify_data_by_type(data, order, [x_name])

    # Choose the first column found on the following order: Q, O, N, T
    chosen_y = None
    for typ in order:
        if len(d[typ]) >= 1:
            chosen_y = d[typ][0]
            break

    return chosen_y


def display_dataframe(df):
    selected_x = select_x(df)
    selected_y = select_y(df, selected_x)
    encoding = Encoding(chart_type=Encoding.chart_type_table, x=selected_x, y=selected_y,
                        y_aggregation=Encoding.y_agg_max)
    return AutoVizWidget(df, encoding)
