/*
 * Copyright (C) 2018 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 "annotator/model-executor.h"

#include "annotator/quantization.h"
#include "utils/base/logging.h"

namespace libtextclassifier3 {

TensorView<float> ModelExecutor::ComputeLogits(
    const TensorView<float>& features, tflite::Interpreter* interpreter) const {
  if (!interpreter) {
    return TensorView<float>::Invalid();
  }
  interpreter->ResizeInputTensor(kInputIndexFeatures, features.shape());
  if (interpreter->AllocateTensors() != kTfLiteOk) {
    TC3_VLOG(1) << "Allocation failed.";
    return TensorView<float>::Invalid();
  }

  SetInput<float>(kInputIndexFeatures, features, interpreter);

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

  return OutputView<float>(kOutputIndexLogits, interpreter);
}

std::unique_ptr<TFLiteEmbeddingExecutor> TFLiteEmbeddingExecutor::FromBuffer(
    const flatbuffers::Vector<uint8_t>* model_spec_buffer, int embedding_size,
    int quantization_bits,
    const Model_::EmbeddingPruningMask* embedding_pruning_mask) {
  std::unique_ptr<TfLiteModelExecutor> executor =
      TfLiteModelExecutor::FromBuffer(model_spec_buffer);
  if (!executor) {
    TC3_LOG(ERROR) << "Could not load TFLite model for embeddings.";
    return nullptr;
  }

  std::unique_ptr<tflite::Interpreter> interpreter =
      executor->CreateInterpreter();
  if (!interpreter) {
    TC3_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)) {
    TC3_LOG(ERROR) << "Mismatch in quantization parameters.";
    return nullptr;
  }

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

TFLiteEmbeddingExecutor::TFLiteEmbeddingExecutor(
    std::unique_ptr<TfLiteModelExecutor> executor, 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,
    const Model_::EmbeddingPruningMask* embedding_pruning_mask)
    : executor_(std::move(executor)),
      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)) {
  if ((embedding_pruning_mask != nullptr) &&
      (embedding_pruning_mask->enabled())) {
    for (int i = 0; i < embedding_pruning_mask->pruning_mask()->size(); i++) {
      pruning_mask_.push_back((*(embedding_pruning_mask->pruning_mask()))[i]);
    }
    ComputePrefixCounts();
    full_num_buckets_ = embedding_pruning_mask->full_num_buckets();
    pruned_row_bucket_id_ = embedding_pruning_mask->pruned_row_bucket_id();
  } else {
    full_num_buckets_ = num_buckets;
  }
}

void TFLiteEmbeddingExecutor::ComputePrefixCounts() {
  // Pre-compute the prefix sums.
  // For each i in {0, 1,...,pruning_mask_.size()-1}, we compute number of 1s
  // in binary representations of the uint64 values in pruning_mask_ before
  // index i. We set pruned_row_bucket_id_ to the total number of 1s
  // in binary representations of all values in pruning_mask_.
  int count = 0;
  for (const uint64 mask : pruning_mask_) {
    prefix_counts_.push_back(count);
    count += __builtin_popcountll(mask);
  }
}

int TFLiteEmbeddingExecutor::PruneBucketId(int bucket_id) const {
  // Implements auxiliary data structure for computing the pruned index of a
  // given bucket_id.
  // If bucket_id is present in pruning_mask_, we compute floor(bucket_id/64),
  // look it up in the auxiliary array prefix_counts_, and add to it the number
  // of 1s before before bucket_id % 64 in the 64-bit sequence
  // pruning_mask_[floor(bucket_id/64)].
  // If bucket_id is absent from pruning_mask_, we return pruned_row_bucket_id_.
  const int bucket_id_major = bucket_id >> 6;
  const int bucket_id_minor = bucket_id & 63;
  uint64_t one = 1;
  if (!(pruning_mask_[bucket_id_major] & (one << bucket_id_minor)))
    return pruned_row_bucket_id_;
  const uint64 zero = 0;
  uint64 minor_mask;
  if (bucket_id_minor == 0)
    minor_mask = zero;
  else
    minor_mask = ((~zero) >> (64 - bucket_id_minor));
  return prefix_counts_[bucket_id_major] +
         __builtin_popcountll(pruning_mask_[bucket_id_major] & minor_mask);
}

bool TFLiteEmbeddingExecutor::AddEmbedding(
    const TensorView<int>& sparse_features, float* dest, int dest_size) const {
  if (dest_size != output_embedding_size_) {
    TC3_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];
    int full_num_buckets;
    if (!pruning_mask_.empty()) {
      full_num_buckets = full_num_buckets_;
    } else {
      full_num_buckets = num_buckets_;
    }
    if (bucket_id >= full_num_buckets) {
      return false;
    }
    int final_bucket_id;
    if (!pruning_mask_.empty()) {
      final_bucket_id = PruneBucketId(bucket_id);
    } else {
      final_bucket_id = bucket_id;
    }
    if (!DequantizeAdd(scales_->data.f, embeddings_->data.uint8,
                       bytes_per_embedding_, num_sparse_features,
                       quantization_bits_, final_bucket_id, dest, dest_size)) {
      return false;
    }
  }
  return true;
}

}  // namespace libtextclassifier3