/*
 * Copyright (C) 2017 The Android Open Source Project
 *
 * 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 "model-executor.h"

#include "quantization.h"
#include "util/base/logging.h"

namespace libtextclassifier2 {
namespace internal {
bool FromModelSpec(const tflite::Model* model_spec,
                   std::unique_ptr<const tflite::FlatBufferModel>* model) {
  *model = tflite::FlatBufferModel::BuildFromModel(model_spec);
  if (!(*model) || !(*model)->initialized()) {
    TC_LOG(ERROR) << "Could not build TFLite model from a model spec. ";
    return false;
  }
  return true;
}
}  // namespace internal

std::unique_ptr<tflite::Interpreter> ModelExecutor::CreateInterpreter() const {
  std::unique_ptr<tflite::Interpreter> interpreter;
  tflite::InterpreterBuilder(*model_, builtins_)(&interpreter);
  return interpreter;
}

std::unique_ptr<TFLiteEmbeddingExecutor> TFLiteEmbeddingExecutor::Instance(
    const flatbuffers::Vector<uint8_t>* model_spec_buffer, int embedding_size,
    int quantization_bits) {
  const tflite::Model* model_spec =
      flatbuffers::GetRoot<tflite::Model>(model_spec_buffer->data());
  flatbuffers::Verifier verifier(model_spec_buffer->data(),
                                 model_spec_buffer->Length());
  std::unique_ptr<const tflite::FlatBufferModel> model;
  if (!model_spec->Verify(verifier) ||
      !internal::FromModelSpec(model_spec, &model)) {
    TC_LOG(ERROR) << "Could not load TFLite model.";
    return nullptr;
  }

  std::unique_ptr<tflite::Interpreter> interpreter;
  tflite::ops::builtin::BuiltinOpResolver builtins;
  tflite::InterpreterBuilder(*model, builtins)(&interpreter);
  if (!interpreter) {
    TC_LOG(ERROR) << "Could not build TFLite interpreter for embeddings.";
    return nullptr;
  }

  if (interpreter->tensors_size() != 2) {
    return nullptr;
  }
  const TfLiteTensor* embeddings = interpreter->tensor(0);
  if (embeddings->dims->size != 2) {
    return nullptr;
  }
  int num_buckets = embeddings->dims->data[0];
  const TfLiteTensor* scales = interpreter->tensor(1);
  if (scales->dims->size != 2 || scales->dims->data[0] != num_buckets ||
      scales->dims->data[1] != 1) {
    return nullptr;
  }
  int bytes_per_embedding = embeddings->dims->data[1];
  if (!CheckQuantizationParams(bytes_per_embedding, quantization_bits,
                               embedding_size)) {
    TC_LOG(ERROR) << "Mismatch in quantization parameters.";
    return nullptr;
  }

  return std::unique_ptr<TFLiteEmbeddingExecutor>(new TFLiteEmbeddingExecutor(
      std::move(model), quantization_bits, num_buckets, bytes_per_embedding,
      embedding_size, scales, embeddings, std::move(interpreter)));
}

TFLiteEmbeddingExecutor::TFLiteEmbeddingExecutor(
    std::unique_ptr<const tflite::FlatBufferModel> model, int quantization_bits,
    int num_buckets, int bytes_per_embedding, int output_embedding_size,
    const TfLiteTensor* scales, const TfLiteTensor* embeddings,
    std::unique_ptr<tflite::Interpreter> interpreter)
    : model_(std::move(model)),
      quantization_bits_(quantization_bits),
      num_buckets_(num_buckets),
      bytes_per_embedding_(bytes_per_embedding),
      output_embedding_size_(output_embedding_size),
      scales_(scales),
      embeddings_(embeddings),
      interpreter_(std::move(interpreter)) {}

bool TFLiteEmbeddingExecutor::AddEmbedding(
    const TensorView<int>& sparse_features, float* dest, int dest_size) const {
  if (dest_size != output_embedding_size_) {
    TC_LOG(ERROR) << "Mismatching dest_size and output_embedding_size: "
                  << dest_size << " " << output_embedding_size_;
    return false;
  }
  const int num_sparse_features = sparse_features.size();
  for (int i = 0; i < num_sparse_features; ++i) {
    const int bucket_id = sparse_features.data()[i];
    if (bucket_id >= num_buckets_) {
      return false;
    }

    if (!DequantizeAdd(scales_->data.f, embeddings_->data.uint8,
                       bytes_per_embedding_, num_sparse_features,
                       quantization_bits_, bucket_id, dest, dest_size)) {
      return false;
    }
  }
  return true;
}

TensorView<float> ComputeLogitsHelper(const int input_index_features,
                                      const int output_index_logits,
                                      const TensorView<float>& features,
                                      tflite::Interpreter* interpreter) {
  if (!interpreter) {
    return TensorView<float>::Invalid();
  }
  interpreter->ResizeInputTensor(input_index_features, features.shape());
  if (interpreter->AllocateTensors() != kTfLiteOk) {
    TC_VLOG(1) << "Allocation failed.";
    return TensorView<float>::Invalid();
  }

  TfLiteTensor* features_tensor =
      interpreter->tensor(interpreter->inputs()[input_index_features]);
  int size = 1;
  for (int i = 0; i < features_tensor->dims->size; ++i) {
    size *= features_tensor->dims->data[i];
  }
  features.copy_to(features_tensor->data.f, size);

  if (interpreter->Invoke() != kTfLiteOk) {
    TC_VLOG(1) << "Interpreter failed.";
    return TensorView<float>::Invalid();
  }

  TfLiteTensor* logits_tensor =
      interpreter->tensor(interpreter->outputs()[output_index_logits]);

  std::vector<int> output_shape(logits_tensor->dims->size);
  for (int i = 0; i < logits_tensor->dims->size; ++i) {
    output_shape[i] = logits_tensor->dims->data[i];
  }

  return TensorView<float>(logits_tensor->data.f, output_shape);
}

}  // namespace libtextclassifier2