# Copyright 1999-2020 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

import numpy as np

from ...config import options
from ...serialize import ValueType, StringField, TupleField
from ..utils import normalize_shape, decide_chunk_sizes
from ..core import TensorOrder
from ..operands import TensorOperand, TensorOperandMixin


class TensorDataSource(TensorOperand, TensorOperandMixin):
    """
    Tensor data source base class, provide universal tile logic,
    subclass can overwrite tile method.
    """

    __slots__ = ()

    def to_chunk_op(self, *args):
        chunk_shape = args[0]
        chunk_op = self.copy().reset_key()
        chunk_op.extra_params = {'size': chunk_shape}  # to make op key different
        return chunk_op

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

        chunk_size = tensor.extra_params.raw_chunk_size or options.chunk_size
        chunk_size = decide_chunk_sizes(tensor.shape, chunk_size, tensor.dtype.itemsize)
        chunk_size_idxes = (range(len(size)) for size in chunk_size)

        out_chunks = []
        for chunk_shape, chunk_idx in zip(itertools.product(*chunk_size),
                                          itertools.product(*chunk_size_idxes)):
            chunk_op = op.to_chunk_op(chunk_shape, chunk_idx, chunk_size)
            out_chunk = chunk_op.new_chunk(None, shape=chunk_shape, index=chunk_idx,
                                           order=tensor.order)
            out_chunks.append(out_chunk)

        new_op = op.copy()
        return new_op.new_tensors(op.inputs, tensor.shape, chunks=out_chunks, nsplits=chunk_size,
                                  order=tensor.order, **tensor.extra_params)


class TensorNoInput(TensorDataSource):
    """
    Tensor operand with no inputs.
    """

    def check_inputs(self, inputs):
        # no inputs
        if inputs and len(inputs) > 0:
            raise ValueError("Tensor data source has no inputs")

    def _new_chunks(self, inputs, kws=None, **kw):
        shape = kw.get('shape', None)
        self.extra_params['shape'] = shape  # set shape to make the operand key different
        return super()._new_chunks(inputs, kws=kws, **kw)

    def _new_tileables(self, inputs, kws=None, **kw):
        shape = kw.get('shape', None)
        self.extra_params['shape'] = shape  # set shape to make the operand key different
        return super()._new_tileables(inputs, kws=kws, **kw)

    def __call__(self, shape, chunk_size=None, order=None):
        shape = normalize_shape(shape)
        order = TensorOrder.C_ORDER if order is None else order
        return self.new_tensor(None, shape, raw_chunk_size=chunk_size, order=order)


class TensorHasInput(TensorDataSource):
    """
    Tensor operand with a single input.
    """

    @property
    def input(self):
        return self._input

    def check_inputs(self, inputs):
        # no inputs
        if len(inputs) != 1:
            raise ValueError("Tensor can only have 1 input")

    def _set_inputs(self, inputs):
        super()._set_inputs(inputs)
        self._input = self._inputs[0]

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

        out_chunks = []
        for c in op.input.chunks:
            out_chunk = op.copy().reset_key().new_chunk([c], shape=c.shape,
                                                        index=c.index, order=output.order)
            out_chunks.append(out_chunk)

        new_op = op.copy()
        return new_op.new_tensors(op.inputs, output.shape, order=output.order,
                                  chunks=out_chunks, nsplits=op.input.nsplits)

    def __call__(self, a, order=None):
        order = a.order if order is None else order
        return self.new_tensor([a], a.shape, order=order)


class TensorLike(TensorHasInput):
    def _set_inputs(self, inputs):
        super()._set_inputs(inputs)
        if self.dtype is None:
            self._dtype = self.input.dtype
        if self.gpu is None:
            self._gpu = self.input.op.gpu

        # FIXME: remove when cupy supports other dtypes
        if self._gpu and self._dtype not in (np.float32, np.float64):
            raise NotImplementedError('Sparse tensor on GPU only supports float32 and float64')


class TensorFromHDF5Like(TensorNoInput):
    _filename = StringField('filename')
    _group = StringField('group')
    _dataset = StringField('dataset')
    _axis_offsets = TupleField('axis_offsets', ValueType.int64)

    def __init__(self, filename=None, group=None, dataset=None,
                 dtype=None, **kw):
        super().__init__(_filename=filename, _group=group,
                         _dataset=dataset, _dtype=dtype, **kw)

    @property
    def filename(self):
        return self._filename

    @property
    def group(self):
        return self._group

    @property
    def dataset(self):
        return self._dataset

    @property
    def axis_offsets(self):
        return self._axis_offsets

    @property
    def path(self):
        return self.get_path(self.group, self.dataset)

    def to_chunk_op(self, *args):
        _, chunk_index, nsplits = args
        chunk_op = super().to_chunk_op(*args)
        cum_offsets = [[0] + np.cumsum(ns).tolist() for ns in nsplits]
        axis_offsets = []
        for axis, idx in enumerate(chunk_index):
            axis_offsets.append(cum_offsets[axis][idx])
        chunk_op._axis_offsets = tuple(axis_offsets)
        return chunk_op

    @staticmethod
    def get_path(group, dataset):
        paths = []
        if group:
            paths.append(group)
        paths.append(dataset)
        return '/'.join(paths)
