# 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 itertools
from numbers import Integral

import pandas as pd
import numpy as np

from ... import opcodes as OperandDef
from ...config import options
from ...core import ENTITY_TYPE, OutputType, recursive_tile
from ...serialization.serializables import AnyField, Int32Field, BoolField
from ...tensor.core import TENSOR_TYPE, TENSOR_CHUNK_TYPE
from ...tensor.datasource import tensor as astensor
from ...utils import has_unknown_shape
from ..align import align_dataframe_series, align_dataframe_dataframe
from ..core import (
    SERIES_TYPE,
    SERIES_CHUNK_TYPE,
    DATAFRAME_TYPE,
    DATAFRAME_CHUNK_TYPE,
    is_chunk_meta_lazy,
)
from ..merge import DataFrameConcat
from ..operands import DataFrameOperand, DataFrameOperandMixin
from ..utils import parse_index, in_range_index
from .utils import calc_columns_index


class SeriesIndex(DataFrameOperand, DataFrameOperandMixin):
    _op_module_ = "series"
    _op_type_ = OperandDef.INDEX

    _labels = AnyField("labels")

    _combine_size = Int32Field("combine_size")
    _is_intermediate = BoolField("is_intermediate")

    def __init__(
        self,
        labels=None,
        combine_size=None,
        is_intermediate=None,
        output_types=None,
        **kw,
    ):
        super().__init__(
            _labels=labels,
            _combine_size=combine_size,
            _is_intermediate=is_intermediate,
            _output_types=output_types,
            **kw,
        )

    @property
    def labels(self):
        return self._labels

    @property
    def combine_size(self):
        return self._combine_size

    @property
    def is_intermediate(self):
        return self._is_intermediate

    def __call__(self, series, name=None):
        return self.new_tileable([series], dtype=series.dtype, name=name)

    def _new_tileables(self, inputs, kws=None, **kw):
        # Override this method to automatically decide the output type,
        # when `labels` is a list, we will set `output_types` as series,
        # otherwise it will be a scalar.
        output_types = getattr(self, "_output_types", None)
        shape = kw.pop("shape", None)
        is_scalar = not isinstance(self._labels, list)
        if not output_types:
            output_types = [OutputType.scalar] if is_scalar else [OutputType.series]
            self.output_types = output_types
        if shape is None:
            shape = () if is_scalar else ((len(self._labels)),)
            kw["shape"] = shape
        if not is_scalar:
            index_value = kw.pop("index_value", None) or parse_index(
                pd.Index(self._labels)
            )
            kw["index_value"] = index_value
        return super()._new_tileables(inputs, kws=kws, **kw)

    def _new_chunks(self, inputs, kws=None, **kw):
        # Override this method to automatically decide the output type,
        # when `labels` is a list, we will set `output_types` as series,
        # otherwise it will be a scalar.
        output_types = getattr(self, "_output_types", None)
        is_scalar = not isinstance(self._labels, list)
        if not output_types:
            output_types = [OutputType.scalar] if is_scalar else [OutputType.series]
            self.output_types = output_types
        if kw.get("shape", None) is None:
            shape = () if is_scalar else ((len(self._labels)),)
            kw["shape"] = shape
        if not is_scalar:
            index_value = kw.pop("index_value", None) or parse_index(
                pd.Index(self._labels)
            )
            kw["index_value"] = index_value
        else:
            # tensor chunk cannot accept index_value
            kw.pop("index_value", None)
        return super()._new_chunks(inputs, kws=kws, **kw)

    @classmethod
    def _calc_chunk_index(cls, label, chunk_indexes):
        for i, index in enumerate(chunk_indexes):
            if isinstance(index, pd.RangeIndex) and in_range_index(label, index):
                return i
            elif label in index:
                return i
        raise TypeError(f"label {label} doesn't exist")

    @classmethod
    def _tile_one_chunk(cls, op):
        in_series = op.inputs[0]
        out_series = op.outputs[0]

        index_op = SeriesIndex(labels=op.labels)
        kw = {"name": out_series.name} if hasattr(out_series, "name") else {}
        index_chunk = index_op.new_chunk(in_series.chunks, dtype=out_series.dtype, **kw)
        new_op = op.copy()
        nsplits = ((len(op.labels),),) if isinstance(op.labels, list) else ()
        return new_op.new_tileables(
            op.inputs, chunks=[index_chunk], nsplits=nsplits, dtype=out_series.dtype
        )

    @classmethod
    def _tree_getitem(cls, op):
        """
        DataFrame doesn't store the index value except RangeIndex or specify `store=True` in `parse_index`,
        So we build a tree structure to avoid too much dependence for getitem node.
        """
        out_series = op.outputs[0]
        combine_size = options.combine_size
        chunks = op.inputs[0].chunks
        while len(chunks) > combine_size:
            new_chunks = []
            for i in range(0, len(chunks), combine_size):
                chks = chunks[i : i + combine_size]
                if len(chks) == 1:
                    chk = chks[0]
                else:
                    concat_op = DataFrameConcat(output_types=[OutputType.series])
                    chk = concat_op.new_chunk(chks, dtype=chks[0].dtype)
                chk_op = SeriesIndex(labels=op.labels, is_intermediate=True)
                kw = {"name": out_series.name} if hasattr(out_series, "name") else {}
                chk = chk_op.new_chunk(
                    [chk],
                    shape=(np.nan,),
                    dtype=chk.dtype,
                    index_value=parse_index(pd.RangeIndex(-1)),
                    **kw,
                )
                new_chunks.append(chk)
            chunks = new_chunks

        concat_op = DataFrameConcat(output_types=[OutputType.series])
        kw = {"name": out_series.name} if hasattr(out_series, "name") else {}
        kw["index"] = (0,)
        chk = concat_op.new_chunk(chunks, dtype=chunks[0].dtype, **kw)
        index_op = SeriesIndex(labels=op.labels)
        chunk = index_op.new_chunk([chk], dtype=chk.dtype, **kw)
        new_op = op.copy()
        nsplits = ((len(op.labels),),) if isinstance(op.labels, list) else ()
        kw = out_series.params
        kw["nsplits"] = nsplits
        kw["chunks"] = [chunk]
        return new_op.new_tileables(op.inputs, kws=[kw])

    @classmethod
    def tile(cls, op):
        in_series = op.inputs[0]
        out_series = op.outputs[0]

        if len(in_series.chunks) == 1:
            return cls._tile_one_chunk(op)
        if not in_series.index_value.has_value():
            return cls._tree_getitem(op)

        chunk_indexes = [c.index_value.to_pandas() for c in in_series.chunks]
        if not isinstance(op.labels, list):
            selected_chunk = in_series.chunks[
                cls._calc_chunk_index(op.labels, chunk_indexes)
            ]
            index_op = op.copy().reset_key()
            out_chunk = index_op.new_chunk(
                [selected_chunk], shape=(), dtype=selected_chunk.dtype
            )
            new_op = op.copy()
            return new_op.new_scalars(
                op.inputs, dtype=out_series.dtype, chunks=[out_chunk]
            )
        else:
            # When input series's index is RangeIndex(5), chunk_size is 3, and labels is [4, 2, 3, 4],
            # Combine the labels in the same chunk, so the splits will be [[4], [2], [3, 4]],
            # the corresponding chunk index is [1, 0, 1].
            selected_index = [
                cls._calc_chunk_index(label, chunk_indexes) for label in op.labels
            ]
            condition = np.where(np.diff(selected_index))[0] + 1
            column_splits = np.split(op.labels, condition)
            column_indexes = np.split(selected_index, condition)

            out_chunks = []
            nsplits = []
            for i, (labels, idx) in enumerate(zip(column_splits, column_indexes)):
                index_op = SeriesIndex(labels=list(labels))
                c = in_series.chunks[idx[0]]
                nsplits.append(len(labels))
                index_value = parse_index(
                    pd.Index([], dtype=c.index_value.to_pandas().dtype), c, labels
                )
                out_chunks.append(
                    index_op.new_chunk(
                        [c],
                        shape=(len(labels),),
                        dtype=c.dtype,
                        index_value=index_value,
                        name=c.name,
                        index=(i,),
                    )
                )
            new_op = op.copy()
            return new_op.new_seriess(
                op.inputs,
                shape=out_series.shape,
                dtype=out_series.dtype,
                index_value=out_series.index_value,
                nsplits=(tuple(nsplits),),
                chunks=out_chunks,
                name=out_series.name,
            )

    @classmethod
    def execute(cls, ctx, op):
        series = ctx[op.inputs[0].key]
        labels = op.labels
        if op.is_intermediate:
            # for intermediate result, it is always a series even if labels is a scalar.
            labels = labels if isinstance(labels, list) else [labels]
            labels = [label for label in set(labels) if label in series]
        ctx[op.outputs[0].key] = series[labels]


class DataFrameIndex(DataFrameOperand, DataFrameOperandMixin):
    _op_type_ = OperandDef.INDEX

    col_names = AnyField("col_names", default=None)

    # for bool index
    mask = AnyField("mask", default=None)
    identical_index = BoolField("identical_index")

    def __init__(self, output_types=None, **kw):
        output_types = output_types or [OutputType.series]
        super().__init__(_output_types=output_types, **kw)

    def _set_inputs(self, inputs):
        super()._set_inputs(inputs)
        if isinstance(self.col_names, ENTITY_TYPE):
            self.col_names = self._inputs[0]
        if isinstance(self.mask, ENTITY_TYPE):
            self.mask = self._inputs[-1]

    def __call__(self, df):
        if self.col_names is not None:
            # if col_names is a list, return a DataFrame, else return a Series
            col_names = self.col_names
            if not isinstance(col_names, list):
                col_names = [col_names]
                is_list = False
            else:
                is_list = True

            dtypes_list = df._get_dtypes_by_columns(col_names)
            if is_list or len(dtypes_list) > 1:
                if len(col_names) != len(dtypes_list):
                    col_names = df._get_columns_by_columns(col_names)
                columns = parse_index(pd.Index(col_names), store_data=True)
                return self.new_dataframe(
                    [df],
                    shape=(df.shape[0], len(col_names)),
                    dtypes=pd.Series(dtypes_list, index=col_names, dtype=np.dtype("O")),
                    index_value=df.index_value,
                    columns_value=columns,
                )
            else:
                dtype = dtypes_list[0]
                return self.new_series(
                    [df],
                    shape=(df.shape[0],),
                    dtype=dtype,
                    index_value=df.index_value,
                    name=self.col_names,
                )
        else:
            if isinstance(self.mask, (SERIES_TYPE, DATAFRAME_TYPE, TENSOR_TYPE)):
                index_value = parse_index(
                    pd.Index(
                        [],
                        dtype=df.index_value.to_pandas().dtype,
                        name=df.index_value.name,
                    ),
                    df,
                    self.mask,
                )
                return self.new_dataframe(
                    [df, self.mask],
                    shape=(np.nan, df.shape[1]),
                    dtypes=df.dtypes,
                    index_value=index_value,
                    columns_value=df.columns_value,
                )
            else:
                index_value = parse_index(
                    pd.Index(
                        [],
                        dtype=df.index_value.to_pandas().dtype,
                        name=df.index_value.name,
                    ),
                    df,
                    self.mask,
                )
                return self.new_dataframe(
                    [df],
                    shape=(np.nan, df.shape[1]),
                    dtypes=df.dtypes,
                    index_value=index_value,
                    columns_value=df.columns_value,
                )

    @classmethod
    def tile(cls, op):
        if op.col_names is not None:
            return cls.tile_with_columns(op)
        else:
            return (yield from cls.tile_with_mask(op))

    @classmethod
    def tile_with_mask(cls, op: "DataFrameIndex"):
        in_df = op.inputs[0]
        out_df = op.outputs[0]

        out_chunks = []

        if isinstance(op.mask, (SERIES_TYPE, DATAFRAME_TYPE, TENSOR_TYPE)):
            mask = op.inputs[1]

            if hasattr(mask, "index_value") and mask.ndim == 1 and op.identical_index:
                if has_unknown_shape(in_df, mask):
                    yield
                nsplits = ((np.nan,) * in_df.chunk_shape[0], in_df.nsplits[1])
                out_shape = in_df.chunk_shape
                df_chunks = in_df.chunks
                aligned_mask = yield from recursive_tile(
                    mask.rechunk(in_df.nsplits[: mask.ndim])
                )
                mask_chunks = aligned_mask.chunks
            elif isinstance(mask, SERIES_TYPE):
                nsplits, out_shape, df_chunks, mask_chunks = align_dataframe_series(
                    in_df, mask, axis="index"
                )
            elif isinstance(mask, DATAFRAME_TYPE):
                nsplits, out_shapes, df_chunks, mask_chunks = align_dataframe_dataframe(
                    in_df, mask
                )
                out_shape = out_shapes[0]
            else:
                # tensor
                nsplits = in_df.nsplits
                mask = yield from recursive_tile(mask.rechunk(nsplits[: mask.ndim]))
                out_shape = in_df.chunk_shape
                df_chunks = in_df.chunks
                mask_chunks = mask.chunks
            out_chunk_indexes = itertools.product(*(range(s) for s in out_shape))

            out_chunks = []
            for i, idx, df_chunk in zip(
                itertools.count(), out_chunk_indexes, df_chunks
            ):
                if op.mask.ndim == 1:
                    mask_chunk = mask_chunks[df_chunk.index[0]]
                else:
                    mask_chunk = mask_chunks[i]
                out_chunk = (
                    op.copy()
                    .reset_key()
                    .new_chunk(
                        [df_chunk, mask_chunk],
                        index=idx,
                        shape=(np.nan, df_chunk.shape[1]),
                    )
                )
                out_chunk._set_tileable_meta(
                    tileable_key=out_df.key,
                    nsplits=nsplits,
                    index_value=out_df.index_value,
                    columns_value=out_df.columns_value,
                    dtypes=out_df.dtypes,
                )
                out_chunks.append(out_chunk)
        else:
            if has_unknown_shape(in_df):
                yield
            nsplits_acc = np.cumsum((0,) + in_df.nsplits[0])
            for idx in range(in_df.chunk_shape[0]):
                for idxj in range(in_df.chunk_shape[1]):
                    in_chunk = in_df.cix[idx, idxj]
                    chunk_op = op.copy().reset_key()
                    chunk_op.mask = op.mask.iloc[
                        nsplits_acc[idx] : nsplits_acc[idx + 1]
                    ]
                    out_chunk = chunk_op.new_chunk(
                        [in_chunk],
                        index=in_chunk.index,
                        shape=(np.nan, in_chunk.shape[1]),
                        dtypes=in_chunk.dtypes,
                        index_value=in_df.index_value,
                        columns_value=in_chunk.columns_value,
                    )
                    out_chunks.append(out_chunk)

        nsplits_on_columns = tuple(c.shape[1] for c in out_chunks if c.index[0] == 0)
        row_chunk_num = len([c.shape[0] for c in out_chunks if c.index[1] == 0])
        nsplits = ((np.nan,) * row_chunk_num, nsplits_on_columns)

        new_op = op.copy()
        return new_op.new_dataframes(
            op.inputs,
            shape=out_df.shape,
            dtypes=out_df.dtypes,
            index_value=out_df.index_value,
            columns_value=out_df.columns_value,
            chunks=out_chunks,
            nsplits=nsplits,
        )

    @classmethod
    def tile_with_columns(cls, op):
        in_df = op.inputs[0]
        out_df = op.outputs[0]
        col_names = op.col_names
        chunk_meta_lazy = is_chunk_meta_lazy(in_df.chunks[0])
        if out_df.ndim < 2:
            # Series
            column_index = calc_columns_index(col_names, in_df)[0]
            out_chunks = []
            dtype = in_df.dtypes[col_names]
            out_nsplits = (in_df.nsplits[0],)
            for i in range(in_df.chunk_shape[0]):
                c = in_df.cix[(i, column_index)]
                chunk_op = DataFrameIndex(col_names=col_names)
                if chunk_meta_lazy:
                    out_chunk = chunk_op.new_chunk(
                        [c],
                        shape=(c.shape[0],),
                        index=(i,),
                        dtype=dtype,
                        name=col_names,
                    )
                    out_chunk._set_tileable_meta(
                        tileable_key=out_df.key,
                        nsplits=out_nsplits,
                        index_value=out_df.index_value,
                    )
                else:
                    out_chunk = chunk_op.new_chunk(
                        [c],
                        shape=(c.shape[0],),
                        index=(i,),
                        dtype=dtype,
                        index_value=c.index_value,
                        name=col_names,
                    )
                out_chunks.append(out_chunk)
            new_op = op.copy()
            params = out_df.params.copy()
            params["chunks"] = out_chunks
            params["nsplits"] = out_nsplits
            return new_op.new_seriess(op.inputs, kws=[params])
        else:
            # combine columns into one chunk and keep the columns order at the same time.
            # When chunk columns are ['c1', 'c2', 'c3'], ['c4', 'c5'],
            # selected columns are ['c2', 'c3', 'c4', 'c2'], `column_splits` will be
            # [(['c2', 'c3'], 0), ('c4', 1), ('c2', 0)].
            if not isinstance(col_names, _list_like_types):
                col_names = [col_names]
            selected_index = [calc_columns_index(col, in_df) for col in col_names]
            selected_index = list(itertools.chain.from_iterable(selected_index))
            condition = np.where(np.diff(selected_index))[0] + 1
            column_splits = np.split(col_names, condition)
            column_indexes = np.split(selected_index, condition)

            out_chunks = [[] for _ in range(in_df.chunk_shape[0])]
            nsplits = [in_df.nsplits[0], []]
            column_nsplits = nsplits[1]
            for i, (columns, column_idx) in enumerate(
                zip(column_splits, column_indexes)
            ):
                try:
                    dtypes = in_df.dtypes[columns]
                except ValueError:  # pragma: no cover
                    dtypes = in_df.dtypes[list(columns)]
                column_nsplits.append(len(dtypes))
                for j in range(in_df.chunk_shape[0]):
                    c = in_df.cix[(j, column_idx[0])]
                    index_op = DataFrameIndex(
                        col_names=list(columns), output_types=[OutputType.dataframe]
                    )
                    if chunk_meta_lazy:
                        out_chunk = index_op.new_chunk(
                            [c], shape=(c.shape[0], len(dtypes)), index=(j, i)
                        )
                        out_chunk._set_tileable_meta(
                            tileable_key=out_df.key,
                            nsplits=nsplits,
                            index_value=out_df.index_value,
                            columns_value=out_df.columns_value,
                            dtypes=out_df.dtypes,
                        )
                    else:
                        out_chunk = index_op.new_chunk(
                            [c],
                            shape=(c.shape[0], len(dtypes)),
                            index=(j, i),
                            dtypes=dtypes,
                            index_value=c.index_value,
                            columns_value=parse_index(
                                pd.Index(dtypes.index), store_data=True
                            ),
                        )
                    out_chunks[j].append(out_chunk)
            out_chunks = [item for cl in out_chunks for item in cl]
            new_op = op.copy()
            params = out_df.params.copy()
            params["chunks"] = out_chunks
            params["nsplits"] = nsplits
            return new_op.new_dataframes(op.inputs, kws=[params])

    @classmethod
    def execute(cls, ctx, op: "DataFrameIndex"):
        if op.mask is None:
            df = ctx[op.inputs[0].key]
            ctx[op.outputs[0].key] = df[op.col_names]
        else:
            df = ctx[op.inputs[0].key]
            if isinstance(
                op.mask, (SERIES_CHUNK_TYPE, DATAFRAME_CHUNK_TYPE, TENSOR_CHUNK_TYPE)
            ):
                mask = ctx[op.inputs[1].key]
            else:
                mask = op.mask
            if hasattr(mask, "reindex_like") and not op.identical_index:
                mask = mask.reindex_like(df).fillna(False)
            ctx[op.outputs[0].key] = df[mask]

    @classmethod
    def estimate_size(cls, ctx: dict, op: "DataFrameIndex"):
        super().estimate_size(ctx, op)
        result_size = ctx[op.outputs[0].key][0]
        ctx[op.outputs[0].key] = (result_size, result_size)


_list_like_types = (list, np.ndarray, SERIES_TYPE, pd.Series, TENSOR_TYPE)


def dataframe_getitem(df, item):
    columns_set = set(df.dtypes.keys())

    if isinstance(item, (np.ndarray, pd.Series)) and item.dtype != np.bool_:
        item = item.tolist()

    if isinstance(item, slice):
        edge = item.start if item.start is not None else item.stop
        if isinstance(edge, Integral):
            return df.iloc[item]
        else:
            return df.loc[item]
    elif isinstance(item, list):
        for col_name in item:
            if col_name not in columns_set:
                raise KeyError(f"{col_name} not in columns")
        op = DataFrameIndex(col_names=item, output_types=[OutputType.dataframe])
    elif isinstance(item, _list_like_types) or hasattr(item, "dtypes"):
        # NB: don't enforce the dtype of `item` to be `bool` since it may be unknown
        if isinstance(item, DATAFRAME_TYPE + SERIES_TYPE):
            identical_index = df.index_value.key == item.index_value.key
        else:
            identical_index = False
        op = DataFrameIndex(
            mask=item,
            identical_index=identical_index,
            output_types=[OutputType.dataframe],
        )
    else:
        if item not in columns_set:
            raise KeyError(f"{item} not in columns {columns_set}")
        op = DataFrameIndex(col_names=item)
    return op(df)


def series_getitem(series, labels, combine_size=None):
    if isinstance(labels, list) or np.isscalar(labels):
        op = SeriesIndex(labels=labels, combine_size=combine_size)
        return op(series, name=series.name)
    elif isinstance(labels, _list_like_types) and astensor(labels).dtype == np.bool_:
        return series.loc[labels]
    elif isinstance(labels, slice):
        edge = labels.start if labels.start is not None else labels.stop
        if isinstance(edge, Integral):
            return series.iloc[labels]
        else:
            return series.loc[labels]
    else:
        raise NotImplementedError(f"type {type(labels)} is not support for getitem")
