/*
 * 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 "common/embedding-network.h"

#include <math.h>

#include "common/simple-adder.h"
#include "util/base/integral_types.h"
#include "util/base/logging.h"

namespace libtextclassifier {
namespace nlp_core {

namespace {

// Returns true if and only if matrix does not use any quantization.
bool CheckNoQuantization(const EmbeddingNetworkParams::Matrix &matrix) {
  if (matrix.quant_type != QuantizationType::NONE) {
    TC_LOG(ERROR) << "Unsupported quantization";
    TC_DCHECK(false);  // Crash in debug mode.
    return false;
  }
  return true;
}

// Initializes a Matrix object with the parameters from the MatrixParams
// source_matrix.  source_matrix should not use quantization.
//
// Returns true on success, false on error.
bool InitNonQuantizedMatrix(const EmbeddingNetworkParams::Matrix &source_matrix,
                            EmbeddingNetwork::Matrix *mat) {
  mat->resize(source_matrix.rows);

  // Before we access the weights as floats, we need to check that they are
  // really floats, i.e., no quantization is used.
  if (!CheckNoQuantization(source_matrix)) return false;
  const float *weights =
      reinterpret_cast<const float *>(source_matrix.elements);
  for (int r = 0; r < source_matrix.rows; ++r) {
    (*mat)[r] = EmbeddingNetwork::VectorWrapper(weights, source_matrix.cols);
    weights += source_matrix.cols;
  }
  return true;
}

// Initializes a VectorWrapper object with the parameters from the MatrixParams
// source_matrix.  source_matrix should have exactly one column and should not
// use quantization.
//
// Returns true on success, false on error.
bool InitNonQuantizedVector(const EmbeddingNetworkParams::Matrix &source_matrix,
                            EmbeddingNetwork::VectorWrapper *vector) {
  if (source_matrix.cols != 1) {
    TC_LOG(ERROR) << "wrong #cols " << source_matrix.cols;
    return false;
  }
  if (!CheckNoQuantization(source_matrix)) {
    TC_LOG(ERROR) << "unsupported quantization";
    return false;
  }
  // Before we access the weights as floats, we need to check that they are
  // really floats, i.e., no quantization is used.
  if (!CheckNoQuantization(source_matrix)) return false;
  const float *weights =
      reinterpret_cast<const float *>(source_matrix.elements);
  *vector = EmbeddingNetwork::VectorWrapper(weights, source_matrix.rows);
  return true;
}

// Computes y = weights * Relu(x) + b where Relu is optionally applied.
template <typename ScaleAdderClass>
bool SparseReluProductPlusBias(bool apply_relu,
                               const EmbeddingNetwork::Matrix &weights,
                               const EmbeddingNetwork::VectorWrapper &b,
                               const VectorSpan<float> &x,
                               EmbeddingNetwork::Vector *y) {
  // Check that dimensions match.
  if ((x.size() != weights.size()) || weights.empty()) {
    TC_LOG(ERROR) << x.size() << " != " << weights.size();
    return false;
  }
  if (weights[0].size() != b.size()) {
    TC_LOG(ERROR) << weights[0].size() << " != " << b.size();
    return false;
  }

  y->assign(b.data(), b.data() + b.size());
  ScaleAdderClass adder(y->data(), y->size());

  const int x_size = x.size();
  for (int i = 0; i < x_size; ++i) {
    const float &scale = x[i];
    if (apply_relu) {
      if (scale > 0) {
        adder.LazyScaleAdd(weights[i].data(), scale);
      }
    } else {
      adder.LazyScaleAdd(weights[i].data(), scale);
    }
  }
  return true;
}
}  // namespace

bool EmbeddingNetwork::ConcatEmbeddings(
    const std::vector<FeatureVector> &feature_vectors, Vector *concat) const {
  concat->resize(concat_layer_size_);

  // Invariant 1: feature_vectors contains exactly one element for each
  // embedding space.  That element is itself a FeatureVector, which may be
  // empty, but it should be there.
  if (feature_vectors.size() != embedding_matrices_.size()) {
    TC_LOG(ERROR) << feature_vectors.size()
                  << " != " << embedding_matrices_.size();
    return false;
  }

  // "es_index" stands for "embedding space index".
  for (int es_index = 0; es_index < feature_vectors.size(); ++es_index) {
    // Access is safe by es_index loop bounds and Invariant 1.
    EmbeddingMatrix *const embedding_matrix =
        embedding_matrices_[es_index].get();
    if (embedding_matrix == nullptr) {
      // Should not happen, hence our terse log error message.
      TC_LOG(ERROR) << es_index;
      return false;
    }

    // Access is safe due to es_index loop bounds.
    const FeatureVector &feature_vector = feature_vectors[es_index];

    // Access is safe by es_index loop bounds, Invariant 1, and Invariant 2.
    const int concat_offset = concat_offset_[es_index];

    if (!GetEmbeddingInternal(feature_vector, embedding_matrix, concat_offset,
                              concat->data(), concat->size())) {
      TC_LOG(ERROR) << es_index;
      return false;
    }
  }
  return true;
}

bool EmbeddingNetwork::GetEmbedding(const FeatureVector &feature_vector,
                                    int es_index, float *embedding) const {
  EmbeddingMatrix *const embedding_matrix = embedding_matrices_[es_index].get();
  if (embedding_matrix == nullptr) {
    // Should not happen, hence our terse log error message.
    TC_LOG(ERROR) << es_index;
    return false;
  }
  return GetEmbeddingInternal(feature_vector, embedding_matrix, 0, embedding,
                              embedding_matrices_[es_index]->dim());
}

bool EmbeddingNetwork::GetEmbeddingInternal(
    const FeatureVector &feature_vector,
    EmbeddingMatrix *const embedding_matrix, const int concat_offset,
    float *concat, int concat_size) const {
  const int embedding_dim = embedding_matrix->dim();
  const bool is_quantized =
      embedding_matrix->quant_type() != QuantizationType::NONE;
  const int num_features = feature_vector.size();
  for (int fi = 0; fi < num_features; ++fi) {
    // Both accesses below are safe due to loop bounds for fi.
    const FeatureType *feature_type = feature_vector.type(fi);
    const FeatureValue feature_value = feature_vector.value(fi);
    const int feature_offset =
        concat_offset + feature_type->base() * embedding_dim;

    // Code below updates max(0, embedding_dim) elements from concat, starting
    // with index feature_offset.  Check below ensures these updates are safe.
    if ((feature_offset < 0) ||
        (feature_offset + embedding_dim > concat_size)) {
      TC_LOG(ERROR) << fi << ": " << feature_offset << " " << embedding_dim
                    << " " << concat_size;
      return false;
    }

    // Pointer to float / uint8 weights for relevant embedding.
    const void *embedding_data;

    // Multiplier for each embedding weight.
    float multiplier;

    if (feature_type->is_continuous()) {
      // Continuous features (encoded as FloatFeatureValue).
      FloatFeatureValue float_feature_value(feature_value);
      const int id = float_feature_value.id;
      embedding_matrix->get_embedding(id, &embedding_data, &multiplier);
      multiplier *= float_feature_value.weight;
    } else {
      // Discrete features: every present feature has implicit value 1.0.
      // Hence, after we grab the multiplier below, we don't multiply it by
      // any weight.
      embedding_matrix->get_embedding(feature_value, &embedding_data,
                                      &multiplier);
    }

    // Weighted embeddings will be added starting from this address.
    float *concat_ptr = concat + feature_offset;

    if (is_quantized) {
      const uint8 *quant_weights =
          reinterpret_cast<const uint8 *>(embedding_data);
      for (int i = 0; i < embedding_dim; ++i, ++quant_weights, ++concat_ptr) {
        // 128 is bias for UINT8 quantization, only one we currently support.
        *concat_ptr += (static_cast<int>(*quant_weights) - 128) * multiplier;
      }
    } else {
      const float *weights = reinterpret_cast<const float *>(embedding_data);
      for (int i = 0; i < embedding_dim; ++i, ++weights, ++concat_ptr) {
        *concat_ptr += *weights * multiplier;
      }
    }
  }
  return true;
}

bool EmbeddingNetwork::ComputeLogits(const VectorSpan<float> &input,
                                     Vector *scores) const {
  return EmbeddingNetwork::ComputeLogitsInternal(input, scores);
}

bool EmbeddingNetwork::ComputeLogits(const Vector &input,
                                     Vector *scores) const {
  return EmbeddingNetwork::ComputeLogitsInternal(input, scores);
}

bool EmbeddingNetwork::ComputeLogitsInternal(const VectorSpan<float> &input,
                                             Vector *scores) const {
  return FinishComputeFinalScoresInternal<SimpleAdder>(input, scores);
}

template <typename ScaleAdderClass>
bool EmbeddingNetwork::FinishComputeFinalScoresInternal(
    const VectorSpan<float> &input, Vector *scores) const {
  // This vector serves as an alternating storage for activations of the
  // different layers. We can't use just one vector here because all of the
  // activations of  the previous layer are needed for computation of
  // activations of the next one.
  std::vector<Vector> h_storage(2);

  // Compute pre-logits activations.
  VectorSpan<float> h_in(input);
  Vector *h_out;
  for (int i = 0; i < hidden_weights_.size(); ++i) {
    const bool apply_relu = i > 0;
    h_out = &(h_storage[i % 2]);
    h_out->resize(hidden_bias_[i].size());
    if (!SparseReluProductPlusBias<ScaleAdderClass>(
            apply_relu, hidden_weights_[i], hidden_bias_[i], h_in, h_out)) {
      return false;
    }
    h_in = VectorSpan<float>(*h_out);
  }

  // Compute logit scores.
  if (!SparseReluProductPlusBias<ScaleAdderClass>(
          true, softmax_weights_, softmax_bias_, h_in, scores)) {
    return false;
  }

  return true;
}

bool EmbeddingNetwork::ComputeFinalScores(
    const std::vector<FeatureVector> &features, Vector *scores) const {
  return ComputeFinalScores(features, {}, scores);
}

bool EmbeddingNetwork::ComputeFinalScores(
    const std::vector<FeatureVector> &features,
    const std::vector<float> extra_inputs, Vector *scores) const {
  // If we haven't successfully initialized, return without doing anything.
  if (!is_valid()) return false;

  Vector concat;
  if (!ConcatEmbeddings(features, &concat)) return false;

  if (!extra_inputs.empty()) {
    concat.reserve(concat.size() + extra_inputs.size());
    for (int i = 0; i < extra_inputs.size(); i++) {
      concat.push_back(extra_inputs[i]);
    }
  }

  scores->resize(softmax_bias_.size());
  return ComputeLogits(concat, scores);
}

EmbeddingNetwork::EmbeddingNetwork(const EmbeddingNetworkParams *model) {
  // We'll set valid_ to true only if construction is successful.  If we detect
  // an error along the way, we log an informative message and return early, but
  // we do not crash.
  valid_ = false;

  // Fill embedding_matrices_, concat_offset_, and concat_layer_size_.
  const int num_embedding_spaces = model->GetNumEmbeddingSpaces();
  int offset_sum = 0;
  for (int i = 0; i < num_embedding_spaces; ++i) {
    concat_offset_.push_back(offset_sum);
    const EmbeddingNetworkParams::Matrix matrix = model->GetEmbeddingMatrix(i);
    if (matrix.quant_type != QuantizationType::UINT8) {
      TC_LOG(ERROR) << "Unsupported quantization for embedding #" << i << ": "
                    << static_cast<int>(matrix.quant_type);
      return;
    }

    // There is no way to accomodate an empty embedding matrix.  E.g., there is
    // no way for get_embedding to return something that can be read safely.
    // Hence, we catch that error here and return early.
    if (matrix.rows == 0) {
      TC_LOG(ERROR) << "Empty embedding matrix #" << i;
      return;
    }
    embedding_matrices_.emplace_back(new EmbeddingMatrix(matrix));
    const int embedding_dim = embedding_matrices_.back()->dim();
    offset_sum += embedding_dim * model->GetNumFeaturesInEmbeddingSpace(i);
  }
  concat_layer_size_ = offset_sum;

  // Invariant 2 (trivial by the code above).
  TC_DCHECK_EQ(concat_offset_.size(), embedding_matrices_.size());

  const int num_hidden_layers = model->GetNumHiddenLayers();
  if (num_hidden_layers < 1) {
    TC_LOG(ERROR) << "Wrong number of hidden layers: " << num_hidden_layers;
    return;
  }
  hidden_weights_.resize(num_hidden_layers);
  hidden_bias_.resize(num_hidden_layers);

  for (int i = 0; i < num_hidden_layers; ++i) {
    const EmbeddingNetworkParams::Matrix matrix =
        model->GetHiddenLayerMatrix(i);
    const EmbeddingNetworkParams::Matrix bias = model->GetHiddenLayerBias(i);
    if (!InitNonQuantizedMatrix(matrix, &hidden_weights_[i]) ||
        !InitNonQuantizedVector(bias, &hidden_bias_[i])) {
      TC_LOG(ERROR) << "Bad hidden layer #" << i;
      return;
    }
  }

  if (!model->HasSoftmaxLayer()) {
    TC_LOG(ERROR) << "Missing softmax layer";
    return;
  }
  const EmbeddingNetworkParams::Matrix softmax = model->GetSoftmaxMatrix();
  const EmbeddingNetworkParams::Matrix softmax_bias = model->GetSoftmaxBias();
  if (!InitNonQuantizedMatrix(softmax, &softmax_weights_) ||
      !InitNonQuantizedVector(softmax_bias, &softmax_bias_)) {
    TC_LOG(ERROR) << "Bad softmax layer";
    return;
  }

  // Everything looks good.
  valid_ = true;
}

int EmbeddingNetwork::EmbeddingSize(int es_index) const {
  return embedding_matrices_[es_index]->dim();
}

}  // namespace nlp_core
}  // namespace libtextclassifier