#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.
#

from collections.abc import Iterable
from typing import Any, List, Optional

import numpy as np
import pandas
import torch
from torch.utils.data import Dataset, DistributedSampler

from raydp.spark.context import save_to_ray
from raydp.spark.resource_manager.exchanger import SharedDataset
from raydp.spark.utils import BLOCK_SIZE_BIT, divide_blocks


class _Dataset(Dataset):
    def __init__(self,
                 feature_columns: List[str] = None,
                 feature_shapes: Optional[List[Any]] = None,
                 feature_types: Optional[List[torch.dtype]] = None,
                 label_column: str = None,
                 label_type: Optional[torch.dtype] = None):
        """
        :param feature_columns: the feature columns in df
        :param feature_shapes: the each feature shape that need to return when loading this
               dataset. If it is not None, it's size must match the size of feature_columns.
               If it is None, we guess all are scalar value and return all as a tensor when
               loading this dataset.
        :param feature_types: the feature types. All will be casted into torch.float by default
        :param label_column: the label column in df
        :param label_type: the label type. It will be casted into torch.float by default.
        """
        super(_Dataset, self).__init__()
        self._feature_columns = feature_columns
        self._feature_shapes = feature_shapes
        self._feature_types = feature_types
        self._label_column = label_column
        self._label_type = label_type

        self._feature_tensor = None
        self._label_tensor = None

    def _check_and_convert(self):
        # convert to list for convenience
        if not isinstance(self._feature_columns, List):
            self._feature_columns = [self._feature_columns]

        if self._feature_shapes:
            if not isinstance(self._feature_shapes, list):
                self._feature_shapes = [self._feature_shapes]

            assert len(self._feature_columns) == len(self._feature_shapes), \
                "The feature_shapes size must match the feature_columns"
            for i in range(len(self._feature_shapes)):
                if not isinstance(self._feature_shapes[i], Iterable):
                    self._feature_shapes[i] = [self._feature_shapes[i]]

        if self._feature_types:
            if not isinstance(self._feature_types, list):
                self._feature_types = [self._feature_types]

            assert len(self._feature_columns) == len(self._feature_types), \
                "The feature_types size must match the feature_columns"
            for i in range(len(self._feature_types)):
                assert all(isinstance(dtype, torch.dtype) for dtype in self._feature_types), \
                    "All value in feature_types should be torch.dtype instance"

        if not self._feature_shapes and self._feature_types:
            assert all(dtype == self._feature_types[0] for dtype in self._feature_types), \
                "All dtypes should be same when feature_shapes doesn't provide"

        if not self._feature_types:
            self._feature_types = [torch.float] * len(self._feature_columns)

        if not self._label_type:
            self._label_type = torch.float

    def _convert_to_tensor(self, df):
        if self._feature_shapes:
            tensors = []
            for col, shape, dtype in zip(self._feature_columns, self._feature_shapes,
                                         self._feature_types):
                column = df[col].values
                if column.dtype == np.object:
                    if isinstance(column[0], np.ndarray):
                        column = np.stack(column)
                    elif isinstance(column[0], (list, tuple)):
                        column = list(column)
                    else:
                        raise Exception(
                            f"Column {col}'s type: {type(column[0])} is not supported. It must "
                            "be numpy built in type or numpy object of (ndarray, list, tuple)")

                t = torch.as_tensor(column, dtype=dtype)
                if shape != [0]:
                    t = t.view(*(-1, *shape))
                tensors.append(t)
            self._feature_tensor = tensors
        else:
            feature_columns = (self._feature_columns if
                               len(self._feature_columns) > 1 else self._feature_columns[0])
            feature_df = df[feature_columns].values
            t = torch.as_tensor(feature_df, dtype=self._feature_types[0])
            self._feature_tensor = [t]

        label_df = df[self._label_column].values
        self._label_tensor = torch.as_tensor(label_df, dtype=self._label_type)

    def _get_next(self, index):
        label = self._label_tensor[index]
        features = [tensor[index] for tensor in self._feature_tensor]
        return (*features, label)


class RayDataset(_Dataset):
    """
    Store Spark DataFrame or koalas.DataFrame into ray object store and wrap into a torch
    Dataset which could be used by torch DataLoader.
    """
    def __init__(self,
                 df: Any = None,
                 feature_columns: List[str] = None,
                 feature_shapes: Optional[List[Any]] = None,
                 feature_types: Optional[List[torch.dtype]] = None,
                 label_column: str = None,
                 label_type: Optional[torch.dtype] = None):
        """
        :param df: Spark DataFrame or Koalas.DataFrame
        """
        super(RayDataset, self).__init__(feature_columns, feature_shapes,
                                         feature_types, label_column, label_type)
        self._unresolved_shared_dataset: SharedDataset = None
        self._resolved_shared_dataset: SharedDataset = None
        self._previous_block_index = -1

        self._check_and_convert()

        if df is not None:
            self._unresolved_shared_dataset = save_to_ray(df)

    def _resolve_with_indices(self,
                              indices: List[int],
                              plasma_store_socket_name: Optional[str]):
        resolved_shared_dataset = self._unresolved_shared_dataset.subset(indices)
        resolved_shared_dataset.set_plasma_store_socket_name(plasma_store_socket_name)
        resolved_shared_dataset.resolve()
        self._resolved_shared_dataset = resolved_shared_dataset

    def __getitem__(self, index):
        block_index = index >> BLOCK_SIZE_BIT
        block_inner_index = (block_index << BLOCK_SIZE_BIT) ^ index
        if block_index != self._previous_block_index:
            self._previous_block_index = block_index
            df = self._resolved_shared_dataset[block_index]
            self._convert_to_tensor(df)
        return self._get_next(block_inner_index)

    def __len__(self):
        """Get the total size"""
        return self._unresolved_shared_dataset.total_size()

    def block_sizes(self) -> List[int]:
        """Get the block sizes"""
        return self._unresolved_shared_dataset.partition_sizes()

    @classmethod
    def _custom_deserialize(cls,
                            data_set: SharedDataset,
                            feature_columns: List[str],
                            feature_shapes: List[Any],
                            feature_types: List[torch.dtype],
                            label_column: str,
                            label_type: torch.dtype):
        instance = cls(
            None, feature_columns, feature_shapes, feature_types, label_column, label_type)
        instance._unresolved_shared_dataset = data_set
        return instance

    def __reduce__(self):
        return (RayDataset._custom_deserialize,
                (self._unresolved_shared_dataset, self._feature_columns, self._feature_shapes,
                 self._feature_types, self._label_column, self._label_type))


class BlockSetSampler(DistributedSampler):
    """
    A distributed sampler for BlockSet.

    We will shuffle the blocks order and then shuffle the block inner if shuffle is set to True.
    """
    def __init__(self, dataset, num_replicas=None, rank=None, shuffle=True, init_lazy=True):
        assert isinstance(dataset, RayDataset)
        self._args = (dataset, num_replicas, rank, shuffle)
        self._inited = False

        self._block_indices = None
        self._selected_indices = None

        if not init_lazy:
            self._init_lazy()

    def _init_lazy(self):
        """
        This is a workaround because of ray sgd call initialize the data creator before of
        setup distributed components.
        """
        if not self._inited:
            super(BlockSetSampler, self).__init__(*self._args)
            self._split_blocks()
            self._inited = True

    def _split_blocks(self):
        block_indexes, packed_selected_indexes = divide_blocks(
            self.dataset.block_sizes(), self.num_replicas, self.rank, self.shuffle)
        self._block_indices = block_indexes
        self._selected_indices = packed_selected_indexes

    def resolve(self, plasma_store_socket_name: Optional[str] = None):
        """Manually trigger the underlying object transfer."""
        self._init_lazy()
        self.dataset._resolve_with_indices(self._block_indices,
                                           plasma_store_socket_name)

    @property
    def block_indices(self):
        return self._block_indices

    def __iter__(self):
        self.resolve()
        # deterministically shuffle based on epoch
        np.random.seed(self.epoch)
        block_indices = list(range(len(self._block_indices)))
        if self.shuffle:
            np.random.shuffle(block_indices)

        indices = []
        for index in block_indices:
            tmp = self._selected_indices[index]
            tmp = np.copy(tmp)
            if self.shuffle:
                np.random.shuffle(tmp)
            indices += tmp.tolist()

        return iter(indices)

    def __len__(self):
        # if we use `if sampler` to determine whether the sampler is None,
        # it will call this method. This can be happened when the BlockSetSampler
        # used in the evaluation in ray TorchTrainer.
        self._init_lazy()
        return self.num_samples


class PandasDataset(_Dataset):
    """
    A pandas dataset which support feature columns with different shapes.
    """
    def __init__(self,
                 df: pandas.DataFrame = None,
                 feature_columns: List[str] = None,
                 feature_shapes: Optional[List[Any]] = None,
                 feature_types: Optional[List[torch.dtype]] = None,
                 label_column: str = None,
                 label_type: Optional[torch.dtype] = None):
        """
        :param df: pandas DataFrame
        """
        super(PandasDataset, self).__init__(feature_columns, feature_shapes,
                                            feature_types, label_column, label_type)
        self._check_and_convert()

        self._size = len(df)
        self._convert_to_tensor(df)

    def __getitem__(self, index):
        return self._get_next(index)

    def __len__(self):
        return self._size
