// Copyright (c) 2022 PaddlePaddle Authors. 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.

#include "fastdeploy/function/split.h"
#include "fastdeploy/utils/utils.h"
#include <cstring>

namespace fastdeploy {
namespace function {

/*
 * All tensors' dimension should be the same and the values of
 * each dimension must be the same, except the axis dimension.
 */
template <typename T> struct SplitFunctor {
 public:
  void operator()(const FDTensor& input,
                  const std::vector<const FDTensor*>& ref_inputs, int axis,
                  std::vector<FDTensor>* outputs) {
    if (input.Numel() == 0) {
      return;
    }

    size_t num = outputs->size();

    int input_rows = 1;
    auto dim_0 = ref_inputs[0]->Shape();
    for (int i = 0; i < axis; ++i) {
      input_rows *= dim_0[i];
    }

    int input_cols = 0;

    std::vector<int64_t> output_cols(outputs->size());
    for (size_t i = 0; i < num; ++i) {
      int t_cols = ref_inputs[i]->Numel() / input_rows;
      input_cols += t_cols;
      output_cols[i] = t_cols;
    }

    // computation
    for (int k = 0; k < input_rows; ++k) {
      const T* src_ptr =
          reinterpret_cast<const T*>(input.Data()) + k * input_cols;
      int col_idx = 0;
      for (size_t j = 0; j < num; ++j) {
        int col_len = output_cols[j];
        auto* out_tensor = &(outputs->at(j));
        if (out_tensor != nullptr) {
          T* dst_ptr = reinterpret_cast<T*>(out_tensor->Data()) + k * col_len;
          std::memcpy(dst_ptr, src_ptr + col_idx, sizeof(T) * col_len);
        }
        col_idx += col_len;
      }
    }
  }
};

inline int GetSplitAxisValue(const FDTensor& x, int axis) {
  int rank = x.Shape().size();
  FDASSERT(axis >= -rank && axis < rank,
           "The axis is expected to be in range of [%d, %d), but got %d", -rank,
           rank, axis);
  if (axis < 0) {
    axis = axis + rank;
  }
  return axis;
}

void CreateSplitOutputs(const FDTensor& x,
                        const std::vector<int>& sections_data,
                        std::vector<FDTensor>* outs, int axis) {
  axis = GetSplitAxisValue(x, axis);
  auto input_axis_dim = x.Shape().at(axis);
  std::vector<int> sections_vec;
  const int unknow_dim_val = -1;
  int unknow_dim_idx = -1;
  int num_of_unknow = 0;
  int sum_of_section = 0;

  for (size_t i = 0; i < sections_data.size(); ++i) {
    sections_vec.push_back(sections_data[i]);
    if (sections_data[i] == unknow_dim_val) {
      num_of_unknow++;
      unknow_dim_idx = i;
    } else {
      sum_of_section += sections_data[i];
    }
  }

  FDASSERT(num_of_unknow <= 1,
           "Only one dimension value of Attr(num_or_sections) "
           "in SplitOp can be -1. "
           "But received Attr(num_or_sections) = [%s].",
           Str(sections_data).c_str());
  if (unknow_dim_idx != -1) {
    // for example, input shape = [4 ,5], axis = 1, sections = [2, 3, -1].
    // input_axis_dim = 5, sum_of_sections = 5.
    // the following check will fail.
    FDASSERT(sum_of_section < input_axis_dim,
             "Sum of Attr(num_or_sections) other than unknown section "
             "must be less than the input's "
             "size "
             "along the split dimension. But received Attr(num_or_sections) "
             "= [%s], input(X)'s shape = [%s], Attr(dim) = %d.",
             Str(sections_data).c_str(), Str(x.Shape()).c_str(), axis);
    sections_vec[unknow_dim_idx] = input_axis_dim - sum_of_section;
  } else {
    FDASSERT(sum_of_section == input_axis_dim,
             "Sum of Attr(num_or_sections) must be equal to the input's "
             "size "
             "along the split dimension. But received Attr(num_or_sections)"
             " = [%s], input(X)'s shape = [%s], Attr(dim) = %d.",
             Str(sections_data).c_str(), Str(x.Shape()).c_str(), axis);
  }
  // fill out dims
  std::vector<std::vector<int64_t>> out_dims(sections_vec.size(), x.Shape());
  for (size_t i = 0; i < sections_vec.size(); ++i) {
    out_dims[i][axis] = sections_vec[i];
  }
  for (size_t i = 0; i < sections_vec.size(); ++i) {
    (*outs)[i].Allocate(out_dims[i], x.Dtype());
  }
}

template <typename T>
void SplitKernel(const FDTensor& x, const std::vector<int>& section,
                 std::vector<FDTensor>* outs, int axis) {
  size_t out_number = section.size();
  outs->resize(out_number);
  CreateSplitOutputs(x, section, outs, axis);

  std::vector<const FDTensor*> shape_refer;
  for (size_t j = 0; j < outs->size(); ++j) {
    shape_refer.emplace_back(&((*outs)[j]));
  }
  SplitFunctor<T> functor;
  functor(x, shape_refer, axis, outs);
}

void Split(const FDTensor& x, const std::vector<int>& num_or_sections,
           std::vector<FDTensor>* out, int axis) {
  FD_VISIT_ALL_TYPES(x.Dtype(), "Split", ([&] {
                       SplitKernel<data_t>(x, num_or_sections, out, axis);
                     }));
}

}  // namespace function
}  // namespace fastdeploy