/*
 * Copyright (c) 2018 Samsung Electronics Co., Ltd. All Rights Reserved
 * 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 "FullyConnectedLayer.h"

#include "tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h"
#include "tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h"
#include "kernel/cpu/OperationUtils.h"

#include <mutex>

namespace neurun
{
namespace kernel
{
namespace cpu
{

FullyConnectedLayer::FullyConnectedLayer()
    : _inputData(nullptr), _weightsData(nullptr), _biasData(nullptr), _outputData(nullptr),
      _inputShape(), _weightsShape(), _biasShape(), _outputShape(),
      _activation(ANEURALNETWORKS_FUSED_NONE), _inputType(OperandType::SCALAR_FLOAT32)
{
  // DO NOTHING
}

// executionMutex is used to protect concurrent access of non-threadsafe resources
// like gemmlowp::GemmContext.
// std::mutex is safe for pthreads on Android.
static std::mutex executionMutex;
bool FullyConnectedLayer::fullyConnectedFloat32()
{
  int total_input_size = 1;
  for (int i = 0; i < _inputShape.dimensions.size(); i++)
  {
    total_input_size *= _inputShape.dimensions[i];
  }

  int input_size = _weightsShape.dimensions[1];
  const int batch_size = total_input_size / input_size;
  const int num_units = _weightsShape.dimensions[0];

  TfLiteFusedActivation act = convertFusedActivation(_activation);

  ::tflite::tensor_utils::VectorBatchVectorAssign(reinterpret_cast<const float *>(_biasData),
                                                  num_units, batch_size,
                                                  reinterpret_cast<float *>(_outputData));

  // Compute output += weight * input
  ::tflite::tensor_utils::MatrixBatchVectorMultiplyAccumulate(
      reinterpret_cast<const float *>(_weightsData), num_units, input_size,
      reinterpret_cast<const float *>(_inputData), batch_size,
      reinterpret_cast<float *>(_outputData), /*result_stride=*/1);

  // Apply activation function
  ::tflite::tensor_utils::ApplyActivationToVector(reinterpret_cast<float *>(_outputData),
                                                  batch_size * num_units, act,
                                                  reinterpret_cast<float *>(_outputData));

  return true;
}

bool FullyConnectedLayer::fullyConnectedQuant8()
{
  throw std::runtime_error{"FullyConnectedLayer : Not tested for TENSOR_QUANT8_ASYMM"};
}

void FullyConnectedLayer::configure(uint8_t *inputData, const Shape inputShape,
                                    uint8_t *weightsData, const Shape weightsShape,
                                    uint8_t *biasData, const Shape biasShape, FuseCode activation,
                                    uint8_t *outputData, const Shape outputShape)
{
  _inputData = inputData;
  _inputShape = inputShape;
  _inputType = inputShape.type;
  _weightsData = weightsData;
  _weightsShape = weightsShape;
  _biasData = biasData;
  _biasShape = biasShape;
  _activation = activation;
  _outputData = outputData;
  _outputShape = outputShape;
}

void FullyConnectedLayer::run()
{
  if (_inputType == OperandType::TENSOR_FLOAT32)
  {
    fullyConnectedFloat32();
  }
  else if (_inputType == OperandType::TENSOR_QUANT8_ASYMM)
  {
    throw std::runtime_error{"FullyConnectedLayer : Not tested for TENSOR_QUANT8_ASYMM"};
    // fullyConnectedQuant8();
  }
}

} // namespace cpu
} // namespace kernel
} // namespace neurun