// 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/core/op_lite.h"
#include <list>
#include <set>
#include <utility>
#include <vector>
#include "lite/core/op_registry.h"
#include "lite/utils/string.h"

namespace paddle {
namespace lite {

bool OpLite::InferShape() {
  auto UseCache = [&, this]() -> bool {
    if (last_input_shapes_.empty()) {
      return false;
    }
    if (last_input_shapes_.size() == input_tensor_ptrs_cache_.size()) {
      for (size_t i = 0; i < input_tensor_ptrs_cache_.size(); i++) {
        if (last_input_shapes_[i] != input_tensor_ptrs_cache_[i]->dims() ||
            last_input_lods_[i] != input_tensor_ptrs_cache_[i]->lod()) {
          return false;
        }
      }
      return true;
    }
    return false;
  };
  if (InferShapeWithCache() && UseCache()) {
    for (size_t i = 0; i < output_tensor_ptrs_cache_.size(); i++) {
      output_tensor_ptrs_cache_[i]->Resize(last_output_shapes_[i]);
      output_tensor_ptrs_cache_[i]->set_lod(last_output_lods_[i]);
    }
  } else {
    this->InferShapeImpl();
    if (InferShapeWithCache()) {
      last_output_shapes_.clear();
      last_output_lods_.clear();
      for (size_t i = 0; i < output_tensor_ptrs_cache_.size(); i++) {
        last_output_shapes_.push_back(output_tensor_ptrs_cache_[i]->dims());
        last_output_lods_.push_back(output_tensor_ptrs_cache_[i]->lod());
      }
      last_input_shapes_.clear();
      last_input_lods_.clear();
      for (size_t i = 0; i < input_tensor_ptrs_cache_.size(); i++) {
        last_input_shapes_.push_back(input_tensor_ptrs_cache_[i]->dims());
        last_input_lods_.push_back(input_tensor_ptrs_cache_[i]->lod());
      }
    }
  }
  return true;
}

std::vector<std::unique_ptr<KernelBase>> OpLite::CreateKernels(
    const std::vector<Place> &places, const std::string &kernel_type) {
  std::vector<std::unique_ptr<KernelBase>> kernels;
  CHECK(!op_type_.empty()) << "op_type_ should be set first";

  auto pick_kernel = [&](const Place &place) {
    auto ks = KernelRegistry::Global().Create(
        op_type_, place.target, place.precision, place.layout);
    VLOG(5) << "pick kernel for " << op_type_ << " " << place.DebugString()
            << " get " << ks.size() << " kernels";
    for (auto &&it : ks) {
      AttachKernel(it.get());
      kernels.emplace_back(std::move(it));
    }
  };

  if (!kernel_type.empty()) {
    Place place;
    std::string op_type, alias;
    KernelBase::ParseKernelType(kernel_type, &op_type, &alias, &place);
    pick_kernel(place);
    CHECK(!kernels.empty()) << "no kernel for kernel type " << kernel_type;
    return kernels;
  }

  std::set<Place> expanded_places(places.begin(), places.end());
  for (auto &place : places) {
    // Pick kernels those support any Precision and any DataLayout, For example:
    // kARM,kFloat,kNCHW -> kARM,kFloat,kAny; kARM,kAny,kNCHW; kARM,kAny,kAny
    expanded_places.insert(
        Place(place.target, place.precision, DATALAYOUT(kAny)));
    expanded_places.insert(Place(place.target, PRECISION(kAny), place.layout));
    expanded_places.insert(
        Place(place.target, PRECISION(kAny), DATALAYOUT(kAny)));
  }

  std::set<TargetType> targets;
  for (auto place : expanded_places) {
    pick_kernel(place);
    targets.insert(place.target);
  }

  VLOG(5) << "op " << op_type_ << " get " << kernels.size() << " kernels";
  return kernels;
}

bool OpLite::Run() {
  CHECK(kernel_);
  SyncInputEvents();

  kernel_->Launch();

  RecordOutputEvents();
  return true;
}

bool OpLite::Attach(const cpp::OpDesc &opdesc, lite::Scope *scope) {
  // valid_places_.clear();
  CHECK(scope != nullptr);
  // CHECK(!op_info_.get());
  scope_ = scope;
  op_info_.reset(
      new OpInfo(opdesc));  // Force clean the out-of-date infomation.
  return AttachImpl(*op_info(), scope);
}

#ifdef LITE_ON_FLATBUFFERS_DESC_VIEW
bool OpLite::Attach(const cpp::OpDescWrite &opdesc, lite::Scope *scope) {
  CHECK(scope != nullptr);
  scope_ = scope;
  return AttachImpl(opdesc, scope);
}
#endif

const Tensor *OpLite::GetTensor(lite::Scope *scope,
                                const std::string &name) const {
  auto *var = scope->FindVar(name);
  CHECK(var) << "no variable called " << name << " found";
  return &var->Get<lite::Tensor>();
}

Tensor *OpLite::GetMutableTensor(lite::Scope *scope,
                                 const std::string &name) const {
  auto *var = scope->FindVar(name);
  CHECK(var) << "no variable called " << name << " found";
  return var->GetMutable<lite::Tensor>();
}

void OpLite::AttachInput(const cpp::OpDesc &op_desc,
                         lite::Scope *scope,
                         const std::string &input_name,
                         bool is_dispensable,
                         lite::Tensor **input_var) {
  bool is_have_input =
      op_desc.HasInput(input_name) && op_desc.Input(input_name).size() > 0;
  CHECK(is_dispensable || is_have_input);
  if (is_have_input) {
    std::string input_var_name = op_desc.Input(input_name).front();
    *input_var = scope->FindVar(input_var_name)->GetMutable<lite::Tensor>();
  }
}

void OpLite::AttachOutput(const cpp::OpDesc &op_desc,
                          lite::Scope *scope,
                          const std::string &output_name,
                          bool is_dispensable,
                          lite::Tensor **output_var) {
  bool is_have_output =
      op_desc.HasOutput(output_name) && op_desc.Output(output_name).size() > 0;
  CHECK(is_dispensable || is_have_output);
  if (is_have_output) {
    std::string output_var_name = op_desc.Output(output_name).front();
    *output_var = scope->FindVar(output_var_name)->GetMutable<lite::Tensor>();
  }
}

bool OpInfo::GetInputArgname(const std::string &value_name,
                             std::string *out) const {
  for (auto &item : inputs()) {
    auto it = std::find(item.second.begin(), item.second.end(), value_name);
    if (it != item.second.end()) {
      *out = item.first;
      return true;
    }
  }
  return false;
}

bool OpInfo::GetOutputArgname(const std::string &value_name,
                              std::string *out) const {
  for (auto &item : outputs()) {
    auto it = std::find(item.second.begin(), item.second.end(), value_name);
    if (it != item.second.end()) {
      *out = item.first;
      return true;
    }
  }
  return false;
}

bool OpInfo::GetInputIndex(const std::string &input_name, int *out) const {
  for (auto &item : inputs()) {
    auto it = std::find(item.second.begin(), item.second.end(), input_name);
    if (it != item.second.end()) {
      *out = it - item.second.begin();
      return true;
    }
  }
  return false;
}

bool OpInfo::GetOutputIndex(const std::string &output_name, int *out) const {
  for (auto &item : outputs()) {
    auto it = std::find(item.second.begin(), item.second.end(), output_name);
    if (it != item.second.end()) {
      *out = it - item.second.begin();
      return true;
    }
  }
  return false;
}

bool OpInfo::HasInputScale(const std::string &name, bool is_scale_name) const {
  bool res = false;
  if (is_scale_name) {
    res = HasAttr(name);
  } else {
    std::string argname;
    int index;
    if (GetInputArgname(name, &argname) && GetInputIndex(name, &index)) {
      res = HasAttr(argname + to_string(index) + "_scale");
    }
  }
  return res;
}

bool OpInfo::HasOutputScale(const std::string &name, bool is_scale_name) const {
  bool res = false;
  if (is_scale_name) {
    res = HasAttr(name);
  } else {
    std::string argname;
    int index;
    if (GetOutputArgname(name, &argname) && GetOutputIndex(name, &index)) {
      res = HasAttr(argname + to_string(index) + "_scale");
    }
  }
  return res;
}

void OpInfo::SetInputScale(const std::string &name,
                           const std::vector<float> &scale_value,
                           bool is_scale_name) {
  std::string scale_name;
  if (is_scale_name) {
    scale_name = name;
  } else {
    std::string argname;
    int index;
    CHECK(GetInputArgname(name, &argname));
    CHECK(GetInputIndex(name, &index));
    CHECK(scale_value.size() > 0)
        << "Error in SetInputScale: the scales should not be empty";
    scale_name = argname + to_string(index) + "_scale";
  }
  SetAttr<std::vector<float>>(scale_name, scale_value);
}

void OpInfo::SetOutputScale(const std::string &name,
                            const std::vector<float> &scale_value,
                            bool is_scale_name) {
  std::string scale_name;
  if (is_scale_name) {
    scale_name = name;
  } else {
    std::string argname;
    int index;
    CHECK(GetOutputArgname(name, &argname));
    CHECK(GetOutputIndex(name, &index));
    CHECK(scale_value.size() > 0)
        << "Error in SetOutputScale: the scales should not be empty";
    scale_name = argname + to_string(index) + "_scale";
  }
  SetAttr<std::vector<float>>(scale_name, scale_value);
}

std::vector<float> OpInfo::GetInputScale(const std::string &name,
                                         bool is_scale_name) const {
  std::string scale_name;
  if (is_scale_name) {
    scale_name = name;
  } else {
    std::string argname;
    int index;
    CHECK(GetInputArgname(name, &argname));
    CHECK(GetInputIndex(name, &index));
    scale_name = argname + to_string(index) + "_scale";
  }
  return GetAttr<std::vector<float>>(scale_name);
}

std::vector<float> OpInfo::GetOutputScale(const std::string &name,
                                          bool is_scale_name) const {
  std::string scale_name;
  if (is_scale_name) {
    scale_name = name;
  } else {
    std::string argname;
    int index;
    CHECK(GetOutputArgname(name, &argname));
    CHECK(GetOutputIndex(name, &index));
    scale_name = argname + to_string(index) + "_scale";
  }
  return GetAttr<std::vector<float>>(scale_name);
}

}  // namespace lite
}  // namespace paddle
