import six
import datetime
import numpy as np
import pandas as pd

from pandas.core.groupby import SeriesGroupBy

import ibis
import ibis.expr.operations as ops

from ibis.pandas.dispatch import execute_node, pre_execute
from ibis.pandas.core import (
    numeric_types, integer_types, timestamp_types, timedelta_types)


@execute_node.register(ops.Strftime, pd.Timestamp, six.string_types)
def execute_strftime_timestamp_str(op, data, format_string, **kwargs):
    return data.strftime(format_string)


@execute_node.register(ops.Strftime, pd.Series, six.string_types)
def execute_strftime_series_str(op, data, format_string, **kwargs):
    return data.dt.strftime(format_string)


@execute_node.register(ops.ExtractTemporalField, pd.Timestamp)
def execute_extract_timestamp_field_timestamp(op, data, **kwargs):
    field_name = type(op).__name__.lower().replace('extract', '')
    return getattr(data, field_name)


@execute_node.register(ops.ExtractMillisecond, pd.Timestamp)
def execute_extract_millisecond_timestamp(op, data, **kwargs):
    return int(data.microsecond // 1000.0)


@execute_node.register(ops.ExtractTemporalField, pd.Series)
def execute_extract_timestamp_field_series(op, data, **kwargs):
    field_name = type(op).__name__.lower().replace('extract', '')
    return getattr(data.dt, field_name).astype(np.int32)


@execute_node.register(
    ops.BetweenTime,
    pd.Series,
    (pd.Series, str, datetime.time),
    (pd.Series, str, datetime.time),
)
def execute_between_time(op, data, lower, upper, **kwargs):
    indexer = pd.DatetimeIndex(data).indexer_between_time(
        lower, upper)
    result = np.zeros(len(data), dtype=np.bool_)
    result[indexer] = True
    return result


@execute_node.register(ops.Date, pd.Series)
def execute_timestamp_date(op, data, **kwargs):
    return data.dt.floor('d')


@execute_node.register((ops.TimestampTruncate, ops.DateTruncate), pd.Series)
def execute_timestamp_truncate(op, data, **kwargs):
    dtype = 'datetime64[{}]'.format(op.unit)
    array = data.values.astype(dtype)
    return pd.Series(array, name=data.name)


@execute_node.register(ops.IntervalFromInteger, pd.Series)
def execute_interval_from_integer_series(op, data, **kwargs):
    resolution = '{}s'.format(op.resolution)

    def convert_to_offset(n):
        return pd.offsets.DateOffset(**{resolution: n})

    return data.apply(convert_to_offset)


@execute_node.register(ops.TimestampAdd, timestamp_types, timedelta_types)
def execute_timestamp_add_datetime_timedelta(op, left, right, **kwargs):
    return pd.Timestamp(left) + pd.Timedelta(right)


@execute_node.register(ops.TimestampAdd, timestamp_types, pd.Series)
def execute_timestamp_add_datetime_series(op, left, right, **kwargs):
    return pd.Timestamp(left) + right


@execute_node.register(ops.IntervalAdd, timedelta_types, timedelta_types)
def execute_interval_add_delta_delta(op, left, right, **kwargs):
    return op.op(pd.Timedelta(left), pd.Timedelta(right))


@execute_node.register(ops.IntervalAdd, timedelta_types, pd.Series)
@execute_node.register(
    ops.IntervalMultiply, timedelta_types, numeric_types + (pd.Series,)
)
def execute_interval_add_multiply_delta_series(op, left, right, **kwargs):
    return op.op(pd.Timedelta(left), right)


@execute_node.register(
    (ops.TimestampAdd, ops.IntervalAdd), pd.Series, timedelta_types)
def execute_timestamp_interval_add_series_delta(op, left, right, **kwargs):
    return left + pd.Timedelta(right)


@execute_node.register(
    (ops.TimestampAdd, ops.IntervalAdd), pd.Series, pd.Series)
def execute_timestamp_interval_add_series_series(op, left, right, **kwargs):
    return left + right


@execute_node.register(ops.TimestampSub, timestamp_types, timedelta_types)
def execute_timestamp_sub_datetime_timedelta(op, left, right, **kwargs):
    return pd.Timestamp(left) - pd.Timedelta(right)


@execute_node.register(
    (ops.TimestampDiff, ops.TimestampSub), timestamp_types, pd.Series)
def execute_timestamp_diff_sub_datetime_series(op, left, right, **kwargs):
    return pd.Timestamp(left) - right


@execute_node.register(ops.TimestampSub, pd.Series, timedelta_types)
def execute_timestamp_sub_series_timedelta(op, left, right, **kwargs):
    return left - pd.Timedelta(right)


@execute_node.register(
    (ops.TimestampDiff, ops.TimestampSub), pd.Series, pd.Series)
def execute_timestamp_diff_sub_series_series(op, left, right, **kwargs):
    return left - right


@execute_node.register(ops.TimestampDiff, timestamp_types, timestamp_types)
def execute_timestamp_diff_datetime_datetime(op, left, right, **kwargs):
    return pd.Timestamp(left) - pd.Timestamp(right)


@execute_node.register(ops.TimestampDiff, pd.Series, timestamp_types)
def execute_timestamp_diff_series_datetime(op, left, right, **kwargs):
    return left - pd.Timestamp(right)


@execute_node.register(
    ops.IntervalMultiply, pd.Series, numeric_types + (pd.Series,)
)
@execute_node.register(
    ops.IntervalFloorDivide,
    (pd.Timedelta, pd.Series),
    numeric_types + (pd.Series,)
)
def execute_interval_multiply_fdiv_series_numeric(op, left, right, **kwargs):
    return op.op(left, right)


@execute_node.register(ops.TimestampFromUNIX, (pd.Series,) + integer_types)
def execute_timestamp_from_unix(op, data, **kwargs):
    return pd.to_datetime(data, unit=op.unit)


@pre_execute.register(ops.TimestampNow)
@pre_execute.register(ops.TimestampNow, ibis.client.Client)
def pre_execute_timestamp_now(op, *args, **kwargs):
    return {op: pd.Timestamp('now')}


@execute_node.register(ops.DayOfWeekIndex, six.string_types + (datetime.date,))
def execute_day_of_week_index_any(op, value, **kwargs):
    return pd.Timestamp(value).dayofweek


@execute_node.register(ops.DayOfWeekIndex, pd.Series)
def execute_day_of_week_index_series(op, data, **kwargs):
    return data.dt.dayofweek.astype(np.int16)


@execute_node.register(ops.DayOfWeekIndex, SeriesGroupBy)
def execute_day_of_week_index_series_group_by(op, data, **kwargs):
    groupings = data.grouper.groupings
    return data.obj.dt.dayofweek.astype(np.int16).groupby(groupings)


def day_name(obj):
    """Backwards compatible name of day getting function.

    Parameters
    ----------
    obj : Union[Series, pd.Timestamp]

    Returns
    -------
    str
        The name of the day corresponding to `obj`
    """
    try:
        return obj.day_name()
    except AttributeError:
        return obj.weekday_name


@execute_node.register(ops.DayOfWeekName, six.string_types + (datetime.date,))
def execute_day_of_week_name_any(op, value, **kwargs):
    return day_name(pd.Timestamp(value))


@execute_node.register(ops.DayOfWeekName, pd.Series)
def execute_day_of_week_name_series(op, data, **kwargs):
    return day_name(data.dt)


@execute_node.register(ops.DayOfWeekName, SeriesGroupBy)
def execute_day_of_week_name_series_group_by(op, data, **kwargs):
    return day_name(data.obj.dt).groupby(data.grouper.groupings)
