// Copyright (c) 2019 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 "lite/kernels/arm/merge_lod_tensor_compute.h"
#include <string>
#include <utility>
#include <vector>
#include "lite/backends/arm/math/funcs.h"
#include "lite/core/op_registry.h"
#include "lite/core/tensor.h"
#include "lite/core/type_system.h"

namespace paddle {
namespace lite {
namespace kernels {
namespace arm {

struct CopyRange {
  size_t begin;
  size_t end;
};

void MergeLodTensorCompute::Run() {
  auto &param = Param<operators::MergeLodTensorParam>();
  const lite::Tensor *x = param.x;
  const lite::Tensor *mask = param.mask;
  const lite::Tensor *in_true = param.in_true;
  const lite::Tensor *in_false = param.in_false;
  lite::Tensor *out = param.out;
  int level = param.level;

  CHECK(in_true->IsInitialized() || in_false->IsInitialized());

  auto &in_true_dim = in_true->dims();
  auto &in_false_dim = in_false->dims();

  // only merge the first dim
  int64_t batch_size = 0;
  std::vector<int64_t> out_shape;
  if (in_true->IsInitialized()) {
    batch_size += in_true->dims()[0];
  }
  if (in_false->IsInitialized()) {
    batch_size += in_false->dims()[0];
  }
  out_shape.push_back(batch_size);
  if (in_true->IsInitialized()) {
    for (int i = 1; i < in_true_dim.size(); i++) {
      out_shape.push_back(in_true_dim[i]);
    }
  } else {
    for (int i = 1; i < in_false_dim.size(); i++) {
      out_shape.push_back(in_false_dim[i]);
    }
  }
  out->Resize(out_shape);

  size_t base_num = static_cast<size_t>(out->numel() / batch_size);
  auto *out_data = out->mutable_data<float>();
  auto *out_lod = out->mutable_lod();
  out_lod->clear();
  auto &mask_dim = mask->dims();
  auto *mask_data = mask->data<bool>();

  size_t out_offset = 0;
  size_t in_true_idx = 0;
  size_t in_false_idx = 0;
  for (size_t i = 0; i < static_cast<size_t>(mask_dim[0]); i++) {
    const Tensor *input = nullptr;
    size_t *in_idx = nullptr;
    if (static_cast<int>(mask_data[i]) == 0) {
      input = in_false;
      in_idx = &in_false_idx;
    } else {
      input = in_true;
      in_idx = &in_true_idx;
    }
    auto lod_and_offset = lite::arm::math::GetSubLoDAndAbsoluteOffset(
        input->lod(), *in_idx, (*in_idx) + 1, 0);
    auto &lod_length = lod_and_offset.first;

    lite::arm::math::AppendLoD(out_lod, lod_length);

    size_t start_offset = lod_and_offset.second.first;
    size_t end_offset = lod_and_offset.second.second;

    CHECK(end_offset >= start_offset);
    size_t len = end_offset - start_offset;
    if (len == 0) {
      continue;
    }
    auto *in_src = input->data<float>() + base_num * start_offset;
    auto *out_dest = out_data + base_num * out_offset;
    size_t copy_num = base_num * len * sizeof(float);
    memcpy(out_dest, in_src, copy_num);
    out_offset += len;
    (*in_idx) += 1;
  }

  for (size_t i = 0; i < level; i++) {
    out_lod->insert(out_lod->begin(), x->lod()[i]);
  }

  return;
}

}  // namespace arm
}  // namespace kernels
}  // namespace lite
}  // namespace paddle

REGISTER_LITE_KERNEL(merge_lod_tensor,
                     kARM,
                     kFloat,
                     kNCHW,
                     paddle::lite::kernels::arm::MergeLodTensorCompute,
                     def)
    .BindInput("X", {LiteType::GetTensorTy(TARGET(kARM))})
    .BindInput("Mask", {LiteType::GetTensorTy(TARGET(kARM), PRECISION(kBool))})
    .BindInput("InTrue", {LiteType::GetTensorTy(TARGET(kARM))})
    .BindInput("InFalse", {LiteType::GetTensorTy(TARGET(kARM))})
    .BindOutput("Out", {LiteType::GetTensorTy(TARGET(kARM))})
    .Finalize();