# Copyright (c) 2018, NVIDIA CORPORATION.

from __future__ import print_function, division

import pandas as pd
import numpy as np
import pickle
from copy import deepcopy, copy
from numba.cuda.cudadrv.devicearray import DeviceNDArray

from librmm_cffi import librmm as rmm

from cudf.dataframe import columnops
from cudf.utils import cudautils, utils, ioutils
from cudf.dataframe.buffer import Buffer
from cudf.dataframe.numerical import NumericalColumn
from cudf.dataframe.column import Column
from cudf.dataframe.datetime import DatetimeColumn
from cudf.dataframe.categorical import CategoricalColumn
from cudf.dataframe.string import StringColumn
from cudf.comm.serialize import register_distributed_serializer


class Index(object):
    """The root interface for all Series indexes.
    """
    is_monotonic = None
    is_monotonic_increasing = None
    is_monotonic_decreasing = None

    def serialize(self, serialize):
        """Serialize into pickle format suitable for file storage or network
        transmission.

        Parameters
        ---
        serialize:  A function provided by register_distributed_serializer
        middleware.
        """
        header = {}
        header['payload'], frames = serialize(pickle.dumps(self))
        header['frame_count'] = len(frames)
        return header, frames

    @classmethod
    def deserialize(cls, deserialize, header, frames):
        """Convert from pickle format into Index

        Parameters
        ---
        deserialize:  A function provided by register_distributed_serializer
        middleware.
        header: The data header produced by the serialize function.
        frames: The serialized data
        """
        payload = deserialize(header['payload'],
                              frames[:header['frame_count']])
        return pickle.loads(payload)

    def take(self, indices):
        """Gather only the specific subset of indices

        Parameters
        ---
        indices: An array-like that maps to values contained in this Index.
        """
        assert indices.dtype.kind in 'iu'
        if indices.size == 0:
            # Empty indices
            return RangeIndex(indices.size)
        else:
            # Gather
            index = cudautils.gather(data=self.gpu_values, index=indices)
            col = self.as_column().replace(data=Buffer(index))
            return as_index(col)

    def argsort(self, ascending=True):
        return self.as_column().argsort(ascending=ascending)

    @property
    def values(self):
        return np.asarray([i for i in self.as_column()])

    def to_pandas(self):
        return pd.Index(self.as_column().to_pandas(), name=self.name)

    def to_arrow(self):
        return self.as_column().to_arrow()

    @ioutils.doc_to_dlpack()
    def to_dlpack(self):
        """{docstring}"""
        import cudf.io.dlpack as dlpack
        return dlpack.to_dlpack(self)

    @property
    def gpu_values(self):
        return self.as_column().to_gpu_array()

    def min(self):
        return self.as_column().min()

    def max(self):
        return self.as_column().max()

    def sum(self):
        return self.as_column().sum()

    def find_segments(self):
        """Return the beginning index for segments

        Returns
        -------
        result : NumericalColumn
        """
        segments, _ = self._find_segments()
        return segments

    def _find_segments(self):
        seg, markers = cudautils.find_segments(self.gpu_values)
        return NumericalColumn(data=Buffer(seg), dtype=seg.dtype), markers

    @classmethod
    def _concat(cls, objs):
        data = Column._concat([o.as_column() for o in objs])
        return as_index(data)

    def _apply_op(self, fn, other=None):
        from cudf.dataframe.series import Series
        idx_series = Series(self)
        op = getattr(idx_series, fn)
        if other is not None:
            return as_index(op(other))
        else:
            return as_index(op())

    def __add__(self, other):
        return self._apply_op('__add__', other)

    def __radd__(self, other):
        return self._apply_op('__radd__', other)

    def __sub__(self, other):
        return self._apply_op('__sub__', other)

    def __rsub__(self, other):
        return self._apply_op('__rsub__', other)

    def __mul__(self, other):
        return self._apply_op('__mul__', other)

    def __rmul__(self, other):
        return self._apply_op('__rmul__', other)

    def __mod__(self, other):
        return self._apply_op('__mod__', other)

    def __rmod__(self, other):
        return self._apply_op('__rmod__', other)

    def __pow__(self, other):
        return self._apply_op('__pow__', other)

    def __floordiv__(self, other):
        return self._apply_op('__floordiv__', other)

    def __rfloordiv__(self, other):
        return self._apply_op('__rfloordiv__', other)

    def __truediv__(self, other):
        return self._apply_op('__truediv__', other)

    def __rtruediv__(self, other):
        return self._apply_op('__rtruediv__', other)

    __div__ = __truediv__

    def __eq__(self, other):
        return self._apply_op('__eq__', other)

    def __ne__(self, other):
        return self._apply_op('__ne__', other)

    def __lt__(self, other):
        return self._apply_op('__lt__', other)

    def __le__(self, other):
        return self._apply_op('__le__', other)

    def __gt__(self, other):
        return self._apply_op('__gt__', other)

    def __ge__(self, other):
        return self._apply_op('__ge__', other)

    def equals(self, other):
        if len(self) != len(other):
            return False
        return (self == other)._values.all()

    def join(self, other, method, how='left', return_indexers=False):
        column_join_res = self.as_column().join(
            other.as_column(), how=how, return_indexers=return_indexers,
            method=method)
        if return_indexers:
            joined_col, indexers = column_join_res
            joined_index = as_index(joined_col)
            return joined_index, indexers
        else:
            return column_join_res

    def rename(self, name):
        """
        Alter Index name.

        Defaults to returning new index.

        Parameters
        ----------
        name : label
            Name(s) to set.

        Returns
        -------
        Index

        Difference from pandas:
          * Not supporting: inplace
        """
        out = self.copy(deep=False)
        out.name = name

        return out.copy(deep=True)


class RangeIndex(Index):
    """An iterable integer index defined by a starting value and ending value.
    Can be sliced and indexed arbitrarily without allocating memory for the
    complete structure.

    Properties
    ---
    _start: The first value
    _stop: The last value
    name: Name of the index
    """

    def __init__(self, start, stop=None, name=None):
        """RangeIndex(size), RangeIndex(start, stop)

        Parameters
        ----------
        start, stop: int
        name: string
        """
        if isinstance(start, range):
            therange = start
            start = therange.start
            stop = therange.stop
        if stop is None:
            start, stop = 0, start
        self._start = int(start)
        self._stop = int(stop)
        self.name = name

    def copy(self, deep=True):
        if(deep):
            result = deepcopy(self)
        else:
            result = copy(self)
        result.name = self.name
        return result

    def __repr__(self):
        return "{}(start={}, stop={})".format(self.__class__.__name__,
                                              self._start, self._stop)

    def __len__(self):
        return max(0, self._stop - self._start)

    def __getitem__(self, index):
        if isinstance(index, slice):
            start, stop, step, sln = utils.standard_python_slice(len(self),
                                                                 index)
            start += self._start
            stop += self._start
            if sln == 0:
                return RangeIndex(0)
            else:
                return index_from_range(start, stop, step)

        elif isinstance(index, int):
            index = utils.normalize_index(index, len(self))
            index += self._start
            return index
        elif isinstance(index, (list, np.ndarray)):
            index = np.array(index)
            index = rmm.to_device(index)

        if isinstance(index, (DeviceNDArray)):
            return self.take(index)
        else:
            raise ValueError(index)

    def __eq__(self, other):
        return super(RangeIndex, self).__eq__(other)

    def equals(self, other):
        if isinstance(other, RangeIndex):
            return (self._start == other._start and self._stop == other._stop)
        else:
            return (self == other)._values.all()

    @property
    def dtype(self):
        return np.dtype(np.int64)

    @property
    def _values(self):
        return self.as_column()

    @property
    def is_contiguous(self):
        return True

    @property
    def size(self):
        return max(0, self._stop - self._start)

    def find_label_range(self, first, last):
        # clip first to range
        if first is None or first < self._start:
            begin = self._start
        elif first < self._stop:
            begin = first
        else:
            begin = self._stop
        # clip last to range
        if last is None:
            end = self._stop
        elif last < self._start:
            end = begin
        elif last < self._stop:
            end = last + 1
        else:
            end = self._stop
        # shift to index
        return begin - self._start, end - self._start

    def as_column(self):
        if len(self) > 0:
            vals = cudautils.arange(self._start, self._stop, dtype=self.dtype)
        else:
            vals = rmm.device_array(0, dtype=self.dtype)
        return NumericalColumn(data=Buffer(vals), dtype=vals.dtype)

    def to_gpu_array(self):
        return self.as_column().to_gpu_array()

    def to_pandas(self):
        return pd.RangeIndex(start=self._start, stop=self._stop,
                             dtype=self.dtype, name=self.name)


def index_from_range(start, stop=None, step=None):
    vals = cudautils.arange(start, stop, step, dtype=np.int64)
    return as_index(vals)


class GenericIndex(Index):
    """An array of orderable values that represent the indices of another Column

    Attributes
    ---
    _values: A Column object
    name: A string
    """

    def __init__(self, values, name=None):
        from cudf.dataframe.series import Series
        # normalize the input
        if isinstance(values, Series):
            name = values.name
            values = values._column
        elif isinstance(values, columnops.TypedColumnBase):
            values = values
        else:
            values = NumericalColumn(data=Buffer(values), dtype=values.dtype)

        assert isinstance(values, columnops.TypedColumnBase), type(values)
        assert values.null_count == 0

        self._values = values
        self.name = name

    def copy(self, deep=True):
        if (deep):
            result = deepcopy(self)
        else:
            result = copy(self)
        result._values = self._values.copy(deep)
        result.name = self.name
        return result

    def serialize(self, serialize):
        header = {}
        header['payload'], frames = serialize(self._values)
        header['frame_count'] = len(frames)
        return header, frames

    @classmethod
    def deserialize(cls, deserialize, header, frames):
        payload = deserialize(header['payload'],
                              frames[:header['frame_count']])
        return cls(payload)

    def __sizeof__(self):
        return self._values.__sizeof__()

    def __reduce__(self):
        return GenericIndex, tuple([self._values])

    def __len__(self):
        return len(self._values)

    def __repr__(self):
        vals = [self._values[i] for i in range(min(len(self), 10))]
        return "{}({}, dtype={})".format(self.__class__.__name__,
                                         vals, self._values.dtype)

    def __getitem__(self, index):
        res = self._values[index]
        if not isinstance(index, int):
            return as_index(res)
        else:
            return res

    def as_column(self):
        """Convert the index as a Series.
        """
        return self._values

    @property
    def dtype(self):
        return self._values.dtype

    def find_label_range(self, first, last):
        """Find range that starts with *first* and ends with *last*,
        inclusively.

        Returns
        -------
        begin, end : 2-tuple of int
            The starting index and the ending index.
            The *last* value occurs at ``end - 1`` position.
        """
        col = self._values
        begin, end = None, None
        if first is not None:
            begin = col.find_first_value(first)
        if last is not None:
            end = col.find_last_value(last)
            end += 1
        return begin, end


class DatetimeIndex(GenericIndex):
    # TODO this constructor should take a timezone or something to be
    # consistent with pandas
    def __init__(self, values, name=None):
        # we should be more strict on what we accept here but
        # we'd have to go and figure out all the semantics around
        # pandas dtindex creation first which.  For now
        # just make sure we handle np.datetime64 arrays
        # and then just dispatch upstream
        if isinstance(values, np.ndarray) and values.dtype.kind == 'M':
            values = DatetimeColumn.from_numpy(values)
        elif isinstance(values, pd.DatetimeIndex):
            values = DatetimeColumn.from_numpy(values.values)

        self._values = values
        self.name = name

    @property
    def year(self):
        return self.get_dt_field('year')

    @property
    def month(self):
        return self.get_dt_field('month')

    @property
    def day(self):
        return self.get_dt_field('day')

    @property
    def hour(self):
        return self.get_dt_field('hour')

    @property
    def minute(self):
        return self.get_dt_field('minute')

    @property
    def second(self):
        return self.get_dt_field('second')

    def get_dt_field(self, field):
        out_column = self._values.get_dt_field(field)
        # columnops.column_empty_like always returns a Column object
        # but we need a NumericalColumn for GenericIndex..
        # how should this be handled?
        out_column = NumericalColumn(data=out_column.data,
                                     mask=out_column.mask,
                                     null_count=out_column.null_count,
                                     dtype=out_column.dtype)
        return as_index(out_column)


class CategoricalIndex(GenericIndex):
    """An categorical of orderable values that represent the indices of another
    Column

    Attributes
    ---
    _values: A CategoricalColumn object
    name: A string
    """

    def __init__(self, values, name=None):
        if isinstance(values, pd.Series) and \
                pd.api.types.is_categorical_dtype(values.dtype):
            values = CategoricalColumn(
                data=Buffer(values.cat.codes.values),
                categories=values.cat.categories.tolist(),
                ordered=values.cat.ordered
            )
        elif isinstance(values, (pd.Categorical, pd.CategoricalIndex)):
            values = CategoricalColumn(
                data=Buffer(values.codes),
                categories=values.categories.tolist(),
                ordered=values.ordered
            )

        self._values = values
        self.name = name
        self.names = [name]

    @property
    def codes(self):
        return self._values.codes

    @property
    def categories(self):
        return self._values.categories


def as_index(arbitrary, name=None):
    """Create an Index from an arbitrary object

    Currently supported inputs are:

    * ``Column``
    * ``Buffer``
    * ``Series``
    * ``Index``
    * numba device array
    * numpy array
    * pyarrow array
    * pandas.Categorical

    Returns
    -------
    result : subclass of Index
        - CategoricalIndex for Categorical input.
        - DatetimeIndex for Datetime input.
        - GenericIndex for all other inputs.
    """
    # This function should probably be moved to Index.__new__
    if isinstance(arbitrary, Index):
        return arbitrary
    elif isinstance(arbitrary, NumericalColumn):
        return GenericIndex(arbitrary, name=name)
    elif isinstance(arbitrary, DatetimeColumn):
        return DatetimeIndex(arbitrary, name=name)
    elif isinstance(arbitrary, CategoricalColumn):
        return CategoricalIndex(arbitrary, name=name)
    elif isinstance(arbitrary, StringColumn):
        raise NotImplementedError(
            "Strings are not yet supported in the index"
        )
    else:
        if hasattr(arbitrary, 'name') and name is None:
            name = arbitrary.name
        if len(arbitrary) == 0:
            return RangeIndex(0, 0, name=name)
        return as_index(columnops.as_column(arbitrary), name=name)


register_distributed_serializer(RangeIndex)
register_distributed_serializer(GenericIndex)
register_distributed_serializer(DatetimeIndex)
register_distributed_serializer(CategoricalIndex)
