# Copyright 1999-2021 Alibaba Group Holding Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pandas as pd

from ... import opcodes as OperandDef
from ...core import OutputType, recursive_tile
from ...serialization.serializables import ListField, BoolField
from ...tensor.base.sort import _validate_sort_psrs_kinds
from ...utils import calc_nsplits
from ..utils import (
    parse_index,
    validate_axis,
    build_concatenated_rows_frame,
    standardize_range_index,
)
from ..operands import DATAFRAME_TYPE
from .core import DataFrameSortOperand
from .psrs import DataFramePSRSOperandMixin, execute_sort_index


class DataFrameSortIndex(DataFrameSortOperand, DataFramePSRSOperandMixin):
    _op_type_ = OperandDef.SORT_INDEX

    level = ListField("level", default=None)
    sort_remaining = BoolField("sort_remaining", default=None)

    @classmethod
    def _tile(cls, op):
        df = op.inputs[0]

        if op.axis == 0:
            if df.chunk_shape[op.axis] == 1:
                if op.output_types[0] == OutputType.dataframe:
                    df = build_concatenated_rows_frame(df)
                    out_chunks = []
                    for chunk in df.chunks:
                        chunk_op = op.copy().reset_key()
                        out_chunks.append(
                            chunk_op.new_chunk(
                                [chunk],
                                shape=chunk.shape,
                                index=chunk.index,
                                index_value=chunk.index_value,
                                columns_value=chunk.columns_value,
                                dtypes=chunk.dtypes,
                            )
                        )
                    new_op = op.copy()
                    kws = op.outputs[0].params.copy()
                    kws["nsplits"] = df.nsplits
                    kws["chunks"] = out_chunks
                    return new_op.new_dataframes(op.inputs, **kws)
                else:
                    out_chunks = []
                    for chunk in df.chunks:
                        chunk_op = op.copy().reset_key()
                        out_chunks.append(
                            chunk_op.new_chunk(
                                [chunk],
                                shape=chunk.shape,
                                index=chunk.index,
                                index_value=chunk.index_value,
                                name=chunk.name,
                                dtype=chunk.dtype,
                            )
                        )
                    new_op = op.copy()
                    kws = op.outputs[0].params.copy()
                    kws["nsplits"] = df.nsplits
                    kws["chunks"] = out_chunks
                    return new_op.new_seriess(op.inputs, **kws)
            else:
                if op.output_types[0] == OutputType.dataframe:
                    df = build_concatenated_rows_frame(df)
                if op.na_position != "last":  # pragma: no cover
                    raise NotImplementedError("Only support puts NaNs at the end.")
                # use parallel sorting by regular sampling
                return (yield from cls._tile_psrs(op, df))
        else:
            assert op.axis == 1

            sorted_columns = list(
                df.columns_value.to_pandas().sort_values(ascending=op.ascending)
            )
            r = [(yield from recursive_tile(df[sorted_columns]))]
            if op.ignore_index:
                chunks = r[0].chunks
                yield chunks
                out = op.outputs[0]
                chunks = standardize_range_index(chunks, axis=0)
                new_op = op.copy()
                return new_op.new_dataframes(
                    op.inputs,
                    shape=out.shape,
                    chunks=chunks,
                    nsplits=calc_nsplits({c.index: c.shape for c in chunks}),
                    index_value=out.index_value,
                    columns_value=out.columns_value,
                    dtypes=out.dtypes,
                )
            return r

    @classmethod
    def execute(cls, ctx, op: "DataFrameSortIndex"):
        in_data = ctx[op.inputs[0].key]
        result = execute_sort_index(in_data, op)
        if op.nrows is not None:
            result = result.head(op.nrows)
        ctx[op.outputs[0].key] = result

    def _call_dataframe(self, df):
        if self.ignore_index:
            index_value = parse_index(pd.RangeIndex(df.shape[0]))
        else:
            index_value = df.index_value
        if self.axis == 0:
            return self.new_dataframe(
                [df],
                shape=df.shape,
                dtypes=df.dtypes,
                index_value=index_value,
                columns_value=df.columns_value,
            )
        else:
            dtypes = df.dtypes.sort_index(ascending=self.ascending)
            columns_value = parse_index(dtypes.index, store_data=True)
            return self.new_dataframe(
                [df],
                shape=df.shape,
                dtypes=dtypes,
                index_value=index_value,
                columns_value=columns_value,
            )

    def _call_series(self, series):
        if self.axis != 0:  # pragma: no cover
            raise TypeError(f"Invalid axis: {self.axis}")
        if self.ignore_index:
            index_value = parse_index(pd.RangeIndex(series.shape[0]))
        else:
            index_value = series.index_value

        return self.new_series(
            [series],
            shape=series.shape,
            dtype=series.dtype,
            index_value=index_value,
            name=series.name,
        )

    def __call__(self, a):
        if isinstance(a, DATAFRAME_TYPE):
            self.output_types = [OutputType.dataframe]
            return self._call_dataframe(a)
        else:
            self.output_types = [OutputType.series]
            return self._call_series(a)


def sort_index(
    a,
    axis=0,
    level=None,
    ascending=True,
    inplace=False,
    kind="quicksort",
    na_position="last",
    sort_remaining=True,
    ignore_index: bool = False,
    parallel_kind="PSRS",
    psrs_kinds=None,
):
    """
    Sort object by labels (along an axis).

    Parameters
    ----------
    a : Input DataFrame or Series.
    axis : {0 or 'index', 1 or 'columns'}, default 0
        The axis along which to sort.  The value 0 identifies the rows,
        and 1 identifies the columns.
    level : int or level name or list of ints or list of level names
        If not None, sort on values in specified index level(s).
    ascending : bool, default True
        Sort ascending vs. descending.
    inplace : bool, default False
        If True, perform operation in-place.
    kind : {'quicksort', 'mergesort', 'heapsort'}, default 'quicksort'
        Choice of sorting algorithm. See also ndarray.np.sort for more
        information.  `mergesort` is the only stable algorithm. For
        DataFrames, this option is only applied when sorting on a single
        column or label.
    na_position : {'first', 'last'}, default 'last'
        Puts NaNs at the beginning if `first`; `last` puts NaNs at the end.
        Not implemented for MultiIndex.
    sort_remaining : bool, default True
        If True and sorting by level and index is multilevel, sort by other
        levels too (in order) after sorting by specified level.
    ignore_index : bool, default False
        If True, the resulting axis will be labeled 0, 1, …, n - 1.
    parallel_kind: {'PSRS'}, optional.
        Parallel sorting algorithm, for the details, refer to:
        http://csweb.cs.wfu.edu/bigiron/LittleFE-PSRS/build/html/PSRSalgorithm.html
    psrs_kinds: Sorting algorithms during PSRS algorithm.

    Returns
    -------
    sorted_obj : DataFrame or None
        DataFrame with sorted index if inplace=False, None otherwise.
    """
    if na_position not in ["last", "first"]:  # pragma: no cover
        raise TypeError(f"Invalid na_position: {na_position}")
    psrs_kinds = _validate_sort_psrs_kinds(psrs_kinds)
    axis = validate_axis(axis, a)
    level = level if isinstance(level, (list, tuple)) else [level]
    op = DataFrameSortIndex(
        level=level,
        axis=axis,
        ascending=ascending,
        inplace=inplace,
        kind=kind,
        na_position=na_position,
        sort_remaining=sort_remaining,
        ignore_index=ignore_index,
        parallel_kind=parallel_kind,
        psrs_kinds=psrs_kinds,
        gpu=a.op.is_gpu(),
    )
    sorted_a = op(a)
    if inplace:
        a.data = sorted_a.data
    else:
        return sorted_a
