/*
 * 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.
 */

#ifndef LIBTEXTCLASSIFIER_COMMON_EMBEDDING_NETWORK_H_
#define LIBTEXTCLASSIFIER_COMMON_EMBEDDING_NETWORK_H_

#include <memory>
#include <vector>

#include "common/embedding-network-params.h"
#include "common/feature-extractor.h"
#include "common/vector-span.h"
#include "util/base/integral_types.h"
#include "util/base/logging.h"
#include "util/base/macros.h"

namespace libtextclassifier {
namespace nlp_core {

// Classifier using a hand-coded feed-forward neural network.
//
// No gradient computation, just inference.
//
// Classification works as follows:
//
// Discrete features -> Embeddings -> Concatenation -> Hidden+ -> Softmax
//
// In words: given some discrete features, this class extracts the embeddings
// for these features, concatenates them, passes them through one or two hidden
// layers (each layer uses Relu) and next through a softmax layer that computes
// an unnormalized score for each possible class.  Note: there is always a
// softmax layer.
class EmbeddingNetwork {
 public:
  // Class used to represent an embedding matrix.  Each row is the embedding on
  // a vocabulary element.  Number of columns = number of embedding dimensions.
  class EmbeddingMatrix {
   public:
    explicit EmbeddingMatrix(const EmbeddingNetworkParams::Matrix source_matrix)
        : rows_(source_matrix.rows),
          cols_(source_matrix.cols),
          quant_type_(source_matrix.quant_type),
          data_(source_matrix.elements),
          row_size_in_bytes_(GetRowSizeInBytes(cols_, quant_type_)),
          quant_scales_(source_matrix.quant_scales) {}

    // Returns vocabulary size; one embedding for each vocabulary element.
    int size() const { return rows_; }

    // Returns number of weights in embedding of each vocabulary element.
    int dim() const { return cols_; }

    // Returns quantization type for this embedding matrix.
    QuantizationType quant_type() const { return quant_type_; }

    // Gets embedding for k-th vocabulary element: on return, sets *data to
    // point to the embedding weights and *scale to the quantization scale (1.0
    // if no quantization).
    void get_embedding(int k, const void **data, float *scale) const {
      if ((k < 0) || (k >= size())) {
        TC_LOG(ERROR) << "Index outside [0, " << size() << "): " << k;

        // In debug mode, crash.  In prod, pretend that k is 0.
        TC_DCHECK(false);
        k = 0;
      }
      *data = reinterpret_cast<const char *>(data_) + k * row_size_in_bytes_;
      if (quant_type_ == QuantizationType::NONE) {
        *scale = 1.0;
      } else {
        *scale = Float16To32(quant_scales_[k]);
      }
    }

   private:
    static int GetRowSizeInBytes(int cols, QuantizationType quant_type) {
      switch (quant_type) {
        case QuantizationType::NONE:
          return cols * sizeof(float);
        case QuantizationType::UINT8:
          return cols * sizeof(uint8);
        default:
          TC_LOG(ERROR) << "Unknown quant type: "
                        << static_cast<int>(quant_type);
          return 0;
      }
    }

    // Vocabulary size.
    const int rows_;

    // Number of elements in each embedding.
    const int cols_;

    const QuantizationType quant_type_;

    // Pointer to the embedding weights, in row-major order.  This is a pointer
    // to an array of floats / uint8, depending on the quantization type.
    // Not owned.
    const void *const data_;

    // Number of bytes for one row.  Used to jump to next row in data_.
    const int row_size_in_bytes_;

    // Pointer to quantization scales.  nullptr if no quantization.  Otherwise,
    // quant_scales_[i] is scale for embedding of i-th vocabulary element.
    const float16 *const quant_scales_;

    TC_DISALLOW_COPY_AND_ASSIGN(EmbeddingMatrix);
  };

  // An immutable vector that doesn't own the memory that stores the underlying
  // floats.  Can be used e.g., as a wrapper around model weights stored in the
  // static memory.
  class VectorWrapper {
   public:
    VectorWrapper() : VectorWrapper(nullptr, 0) {}

    // Constructs a vector wrapper around the size consecutive floats that start
    // at address data.  Note: the underlying data should be alive for at least
    // the lifetime of this VectorWrapper object.  That's trivially true if data
    // points to statically allocated data :)
    VectorWrapper(const float *data, int size) : data_(data), size_(size) {}

    int size() const { return size_; }

    const float *data() const { return data_; }

   private:
    const float *data_;  // Not owned.
    int size_;

    // Doesn't own anything, so it can be copied and assigned at will :)
  };

  typedef std::vector<VectorWrapper> Matrix;
  typedef std::vector<float> Vector;

  // Constructs an embedding network using the parameters from model.
  //
  // Note: model should stay alive for at least the lifetime of this
  // EmbeddingNetwork object.
  explicit EmbeddingNetwork(const EmbeddingNetworkParams *model);

  virtual ~EmbeddingNetwork() {}

  // Returns true if this EmbeddingNetwork object has been correctly constructed
  // and is ready to use.  Idea: in case of errors, mark this EmbeddingNetwork
  // object as invalid, but do not crash.
  bool is_valid() const { return valid_; }

  // Runs forward computation to fill scores with unnormalized output unit
  // scores. This is useful for making predictions.
  //
  // Returns true on success, false on error (e.g., if !is_valid()).
  bool ComputeFinalScores(const std::vector<FeatureVector> &features,
                          Vector *scores) const;

  // Same as above, but allows specification of extra neural network inputs that
  // will be appended to the embedding vector build from features.
  bool ComputeFinalScores(const std::vector<FeatureVector> &features,
                          const std::vector<float> extra_inputs,
                          Vector *scores) const;

  // Constructs the concatenated input embedding vector in place in output
  // vector concat.  Returns true on success, false on error.
  bool ConcatEmbeddings(const std::vector<FeatureVector> &features,
                        Vector *concat) const;

  // Sums embeddings for all features from |feature_vector| and adds result
  // to values from the array pointed-to by |output|.  Embeddings for continuous
  // features are weighted by the feature weight.
  //
  // NOTE: output should point to an array of EmbeddingSize(es_index) floats.
  bool GetEmbedding(const FeatureVector &feature_vector, int es_index,
                    float *embedding) const;

  // Runs the feed-forward neural network for |input| and computes logits for
  // softmax layer.
  bool ComputeLogits(const Vector &input, Vector *scores) const;

  // Same as above but uses a view of the feature vector.
  bool ComputeLogits(const VectorSpan<float> &input, Vector *scores) const;

  // Returns the size (the number of columns) of the embedding space es_index.
  int EmbeddingSize(int es_index) const;

 protected:
  // Builds an embedding for given feature vector, and places it from
  // concat_offset to the concat vector.
  bool GetEmbeddingInternal(const FeatureVector &feature_vector,
                            EmbeddingMatrix *embedding_matrix,
                            int concat_offset, float *concat,
                            int embedding_size) const;

  // Templated function that computes the logit scores given the concatenated
  // input embeddings.
  bool ComputeLogitsInternal(const VectorSpan<float> &concat,
                             Vector *scores) const;

  // Computes the softmax scores (prior to normalization) from the concatenated
  // representation.  Returns true on success, false on error.
  template <typename ScaleAdderClass>
  bool FinishComputeFinalScoresInternal(const VectorSpan<float> &concat,
                                        Vector *scores) const;

  // Set to true on successful construction, false otherwise.
  bool valid_ = false;

  // Network parameters.

  // One weight matrix for each embedding space.
  std::vector<std::unique_ptr<EmbeddingMatrix>> embedding_matrices_;

  // concat_offset_[i] is the input layer offset for i-th embedding space.
  std::vector<int> concat_offset_;

  // Size of the input ("concatenation") layer.
  int concat_layer_size_;

  // One weight matrix and one vector of bias weights for each hiden layer.
  std::vector<Matrix> hidden_weights_;
  std::vector<VectorWrapper> hidden_bias_;

  // Weight matrix and bias vector for the softmax layer.
  Matrix softmax_weights_;
  VectorWrapper softmax_bias_;
};

}  // namespace nlp_core
}  // namespace libtextclassifier

#endif  // LIBTEXTCLASSIFIER_COMMON_EMBEDDING_NETWORK_H_