/*
 * 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 "EmbeddingLookup.h"

#include "NeuralNetworksWrapper.h"
#include "gmock/gmock-matchers.h"
#include "gtest/gtest.h"

using ::testing::FloatNear;
using ::testing::Matcher;

namespace android {
namespace nn {
namespace wrapper {

namespace {

std::vector<Matcher<float>> ArrayFloatNear(const std::vector<float>& values,
                                           float max_abs_error=1.e-6) {
  std::vector<Matcher<float>> matchers;
  matchers.reserve(values.size());
  for (const float& v : values) {
    matchers.emplace_back(FloatNear(v, max_abs_error));
  }
  return matchers;
}

}  // namespace

using ::testing::ElementsAreArray;

#define FOR_ALL_INPUT_AND_WEIGHT_TENSORS(ACTION) \
  ACTION(Value, float)                           \
  ACTION(Lookup, int)

// For all output and intermediate states
#define FOR_ALL_OUTPUT_TENSORS(ACTION) \
  ACTION(Output, float)

class EmbeddingLookupOpModel {
 public:
  EmbeddingLookupOpModel(std::initializer_list<uint32_t> index_shape,
                         std::initializer_list<uint32_t> weight_shape) {
    auto it = weight_shape.begin();
    rows_ = *it++;
    columns_ = *it++;
    features_ = *it;

    std::vector<uint32_t> inputs;

    OperandType LookupTy(Type::TENSOR_INT32, index_shape);
    inputs.push_back(model_.addOperand(&LookupTy));

    OperandType ValueTy(Type::TENSOR_FLOAT32, weight_shape);
    inputs.push_back(model_.addOperand(&ValueTy));

    std::vector<uint32_t> outputs;

    OperandType OutputOpndTy(Type::TENSOR_FLOAT32, weight_shape);
    outputs.push_back(model_.addOperand(&OutputOpndTy));

    auto multiAll = [](const std::vector<uint32_t> &dims) -> uint32_t {
        uint32_t sz = 1;
        for (uint32_t d : dims) { sz *= d; }
        return sz;
    };

    Value_.insert(Value_.end(), multiAll(weight_shape), 0.f);
    Output_.insert(Output_.end(), multiAll(weight_shape), 0.f);

    model_.addOperation(ANEURALNETWORKS_EMBEDDING_LOOKUP, inputs, outputs);
    model_.identifyInputsAndOutputs(inputs, outputs);

    model_.finish();
  }

  void Invoke() {
    ASSERT_TRUE(model_.isValid());

    Compilation compilation(&model_);
    compilation.finish();
    Execution execution(&compilation);

#define SetInputOrWeight(X, T)                                               \
  ASSERT_EQ(execution.setInput(EmbeddingLookup::k##X##Tensor, X##_.data(),   \
                               sizeof(T) * X##_.size()),                     \
            Result::NO_ERROR);

    FOR_ALL_INPUT_AND_WEIGHT_TENSORS(SetInputOrWeight);

#undef SetInputOrWeight

#define SetOutput(X, T)                                                       \
  ASSERT_EQ(execution.setOutput(EmbeddingLookup::k##X##Tensor, X##_.data(),   \
                                sizeof(T) * X##_.size()),                     \
            Result::NO_ERROR);

    FOR_ALL_OUTPUT_TENSORS(SetOutput);

#undef SetOutput

    ASSERT_EQ(execution.compute(), Result::NO_ERROR);
  }

#define DefineSetter(X, T)                       \
  void Set##X(const std::vector<T>& f) {         \
    X##_.insert(X##_.end(), f.begin(), f.end()); \
  }

  FOR_ALL_INPUT_AND_WEIGHT_TENSORS(DefineSetter);

#undef DefineSetter

  void Set3DWeightMatrix(const std::function<float(int, int, int)>& function) {
    for (uint32_t i = 0; i < rows_; i++) {
      for (uint32_t j = 0; j < columns_; j++) {
        for (uint32_t k = 0; k < features_; k++) {
          Value_[(i * columns_ + j) * features_ + k] = function(i, j, k);
        }
      }
    }
  }

  const std::vector<float> &GetOutput() const { return Output_; }

 private:
  Model model_;
  uint32_t rows_;
  uint32_t columns_;
  uint32_t features_;

#define DefineTensor(X, T) std::vector<T> X##_;

  FOR_ALL_INPUT_AND_WEIGHT_TENSORS(DefineTensor);
  FOR_ALL_OUTPUT_TENSORS(DefineTensor);

#undef DefineTensor
};

// TODO: write more tests that exercise the details of the op, such as
// lookup errors and variable input shapes.
TEST(EmbeddingLookupOpTest, SimpleTest) {
  EmbeddingLookupOpModel m({3}, {3, 2, 4});
  m.SetLookup({1, 0, 2});
  m.Set3DWeightMatrix(
      [](int i, int j, int k) { return i + j / 10.0f + k / 100.0f; });

  m.Invoke();

  EXPECT_THAT(m.GetOutput(),
              ElementsAreArray(ArrayFloatNear({
                  1.00, 1.01, 1.02, 1.03, 1.10, 1.11, 1.12, 1.13,  // Row 1
                  0.00, 0.01, 0.02, 0.03, 0.10, 0.11, 0.12, 0.13,  // Row 0
                  2.00, 2.01, 2.02, 2.03, 2.10, 2.11, 2.12, 2.13,  // Row 2
              })));
}

}  // namespace wrapper
}  // namespace nn
}  // namespace android