// Copyright 2018 Xiaomi, Inc.  All rights reserved.
//
// 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.

#ifndef MACE_KERNELS_STRIDED_SLICE_H_
#define MACE_KERNELS_STRIDED_SLICE_H_

#include <algorithm>
#include <memory>
#include <vector>

#include "mace/core/future.h"
#include "mace/core/tensor.h"
#include "mace/public/mace.h"

namespace mace {
namespace kernels {

template <DeviceType D, typename T>
struct StridedSliceFunctor {
  StridedSliceFunctor(int begin_mask,
                      int end_mask,
                      int ellipsis_mask,
                      int new_axis_mask,
                      int shrink_axis_mask)
      : begin_mask_(begin_mask),
        end_mask_(end_mask),
        ellipsis_mask_(ellipsis_mask),
        new_axis_mask_(new_axis_mask),
        shrink_axis_mask_(shrink_axis_mask) {}

  MaceStatus operator()(const Tensor *input,
                        const Tensor *begin_indices,
                        const Tensor *end_indices,
                        const Tensor *strides,
                        Tensor *output,
                        StatsFuture *future) {
    MACE_UNUSED(future);
    MACE_CHECK(ellipsis_mask_ == 0 && new_axis_mask_ == 0,
               "ellipsis_mask and new_axis_mask are not supported yet.");

    Tensor::MappingGuard input_guard(input);
    Tensor::MappingGuard begin_indices_guard(begin_indices);
    Tensor::MappingGuard end_indices_guard(end_indices);
    Tensor::MappingGuard strides_guard(strides);
    const T *input_data = input->data<T>();
    const int32_t *begin_indices_data = begin_indices->data<int32_t>();
    const int32_t *end_indices_data = end_indices->data<int32_t>();
    const int32_t *strides_data = strides->data<int32_t>();

    std::vector<index_t> output_shape;
    std::vector<index_t> real_begin_indices(input->dim_size(), 0);
    std::vector<index_t> real_end_indices(input->dim_size(), 0);
    for (index_t d = 0; d < input->dim_size(); ++d) {
      index_t dim_len = input->dim(d);
      if (begin_mask_ & (1 << d)) {
        real_begin_indices[d] = strides_data[d] > 0 ? 0 : dim_len - 1;
      } else {
        real_begin_indices[d] = (begin_indices_data[d] + dim_len) % dim_len;
      }
      if (end_mask_ & (1 << d)) {
        real_end_indices[d] = strides_data[d] > 0 ? dim_len : -1;
      } else {
        real_end_indices[d] =
            end_indices_data[d] < -dim_len
                ? -1
                : (end_indices_data[d] < 0
                       ? (end_indices_data[d] + dim_len)
                       : std::min(static_cast<index_t>(end_indices_data[d]),
                                  dim_len));
      }

      int32_t out_dim_len = std::max(
          0.f, std::ceil((real_end_indices[d] - real_begin_indices[d]) /
                         static_cast<float>(strides_data[d])));
      if (!(shrink_axis_mask_ & (1 << d))) {
        output_shape.push_back(out_dim_len);
      } else {
        MACE_CHECK(out_dim_len == 1,
                   "cannot shrink axis that has len > 1, dim(", d, "): [",
                   real_begin_indices[d], ", ", real_end_indices[d], "]");
      }
    }

    std::vector<index_t> dim_stride(input->dim_size(), 1);
    for (index_t d = input->dim_size() - 2; d >= 0; --d) {
      dim_stride[d] = dim_stride[d + 1] * input->dim(d + 1);
    }

    MACE_RETURN_IF_ERROR(output->Resize(output_shape));
    Tensor::MappingGuard output_guard(output);
    T *output_data = output->mutable_data<T>();

    bool slice_by_first_axis = true;
    if (strides_data[0] != 1) {
      slice_by_first_axis = false;
    } else {
      for (index_t d = 1; d < input->dim_size(); ++d) {
        if (strides_data[d] != 1 || real_begin_indices[d] != 0 ||
            real_end_indices[d] != input->dim(d)) {
          slice_by_first_axis = false;
          break;
        }
      }
    }

    if (slice_by_first_axis) {
      memcpy(output_data, input_data + real_begin_indices[0] * dim_stride[0],
             sizeof(T) * (real_end_indices[0] - real_begin_indices[0]) *
                 dim_stride[0]);
    } else {
      if (input->dim_size() == 1) {
        for (index_t i = real_begin_indices[0];
             strides_data[0] > 0 ? i < real_end_indices[0]
                                 : i > real_end_indices[0];
             i += strides_data[0]) {
          *output_data++ = input_data[i];
        }

      } else if (input->dim_size() == 2) {
        for (index_t i = real_begin_indices[0];
             strides_data[0] > 0 ? i < real_end_indices[0]
                                 : i > real_end_indices[0];
             i += strides_data[0]) {
          for (index_t j = real_begin_indices[1];
               strides_data[1] > 0 ? j < real_end_indices[1]
                                   : j > real_end_indices[1];
               j += strides_data[1]) {
            *output_data++ = input_data[i * dim_stride[0] + j];
          }
        }
      } else {
        MACE_NOT_IMPLEMENTED;
      }
    }

    return MACE_SUCCESS;
  }

  int begin_mask_;
  int end_mask_;
  int ellipsis_mask_;
  int new_axis_mask_;
  int shrink_axis_mask_;
};

}  // namespace kernels
}  // namespace mace

#endif  // MACE_KERNELS_STRIDED_SLICE_H_
