import numpy as np
import tensorflow as tf
try:
  from tensorflow.math import floormod as tf_floormod
except ImportError: # for older tf-1.x versions
  from tensorflow import floormod as tf_floormod

from onnx_tf.handlers.backend_handler import BackendHandler
from onnx_tf.handlers.handler import onnx_op
from onnx_tf.handlers.handler import tf_func


@onnx_op("Slice")
@tf_func(tf.strided_slice)
class Slice(BackendHandler):

  @classmethod
  def version_1(cls, node, **kwargs):
    tensor_dict = kwargs["tensor_dict"]
    x = tensor_dict[node.inputs[0]]

    full_sizes = x.get_shape().as_list()
    full_begin = [0] * len(full_sizes)

    starts = node.attrs.get("starts")
    ends = node.attrs.get("ends")
    slice_len = len(starts)
    axes = node.attrs.get("axes", list(range(slice_len)))

    for i in range(slice_len):
      starts[i] = full_sizes[
          axes[i]] + starts[i] if starts[i] < 0 else starts[i]
      ends[i] = full_sizes[axes[i]] + ends[i] if ends[i] < 0 else ends[i]
      if full_sizes[axes[i]] is not None:
        ends[i] = np.min([full_sizes[axes[i]], ends[i]])
        starts[i] = np.min([full_sizes[axes[i]], starts[i]])
      full_begin[axes[i]] = starts[i]
      full_sizes[axes[i]] = ends[i] - starts[i]

    return [
        cls.make_tensor_from_onnx_node(
            node,
            tf_func=tf.slice,
            inputs=[
                tensor_dict[node.inputs[0]],
                tf.constant(full_begin),
                tf.constant(full_sizes)
            ],
            **kwargs)
    ]

  @classmethod
  def version_10(cls, node, **kwargs):
    tensor_dict = kwargs["tensor_dict"]
    input_tensor = tensor_dict[node.inputs[0]]
    starts = tensor_dict[node.inputs[1]]
    ends = tensor_dict[node.inputs[2]]

    # first of all, get the input tensor shape
    input_tensor_shape = tf.shape(input_tensor, out_type=ends.dtype)

    axes = tensor_dict[node.inputs[3]] if len(
        node.inputs) >= 4 else tf.range(tf.shape(starts, out_type=ends.dtype)[0])

    # process negative axes
    input_rank = tf.cast(tf.rank(input_tensor), axes.dtype)
    axes = tf_floormod(tf.add(axes, input_rank), input_rank)

    # expand a dimension of 1 at the end
    sparse_indices = tf.expand_dims(axes, -1)

    # build the indexed dimension sizes as sparse_shape
    sparse_shape = tf.gather_nd(
        params=input_tensor_shape, indices=sparse_indices)
    sparse_shape = tf.cast(sparse_shape, ends.dtype)

    # take care of starts, ends that are larger than the dim size.
    starts_min = tf.minimum(starts, sparse_shape)
    ends_min = tf.minimum(ends, sparse_shape)

    # take care of starts, ends that are negative
    is_starts_negative = tf.less(starts_min, tf.zeros_like(starts_min))
    starts_final = tf.where(is_starts_negative, starts_min + sparse_shape,
                            starts_min)
    is_ends_negative = tf.less(ends_min, tf.zeros_like(ends_min))
    ends_final = tf.where(is_ends_negative, ends_min + sparse_shape, ends_min)

    # need to densify everything for the inputs to slice
    # the output shape is the input_tensor rank
    output_shape = tf.reshape(tf.rank(input_tensor), [1])
    output_shape = tf.cast(output_shape, ends.dtype)

    # create dense tensor, pad 0 as default begins
    dense_begins = tf.sparse_to_dense(sparse_indices, output_shape,
                                      starts_final)
    # create dense tensor, pad -1 for next step
    dense_ends = tf.sparse_to_dense(
        sparse_indices,
        output_shape,
        ends_final,
        default_value=tf.constant(-1, dtype=dense_begins.dtype))
    # replace -1 with respective dimension sizes
    dense_ends = tf.where(
        tf.equal(dense_ends, tf.constant(-1, dtype=dense_begins.dtype)),
        input_tensor_shape, dense_ends)

    # create dense tensor for steps if not already so
    if len(node.inputs) >= 5:
      dense_steps = tf.sparse_to_dense(
          sparse_indices,
          output_shape,
          tensor_dict[node.inputs[4]],
          default_value=tf.constant(1, dtype=tensor_dict[node.inputs[4]].dtype))
    else:
      dense_steps = tf.ones(input_tensor_shape.shape, ends.dtype)

    return [
        cls.make_tensor_from_onnx_node(
            node,
            inputs=[
                tensor_dict[node.inputs[0]], dense_begins, dense_ends,
                dense_steps
            ],
            **kwargs)
    ]

  @classmethod
  def version_11(cls, node, **kwargs):
    return cls.version_10(node, **kwargs)