// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include "bahdanau_attention.h"
#include "core/providers/cpu/rnn/rnn_helpers.h"

#include <stdexcept>
#include <memory.h>

using onnxruntime::rnn::detail::Allocate;

namespace onnxruntime {
namespace contrib {

template <typename T>
BahdanauAttention<T>::BahdanauAttention(AllocatorPtr allocator, const logging::Logger& logger,
                                        int batch_size, int max_memory_step, int memory_depth,
                                        int query_depth, int attn_depth, bool normalize)
    : allocator_(allocator), logger_(logger), batch_size_(batch_size), max_memory_steps_(max_memory_step), memory_depth_(memory_depth), query_depth_(query_depth), attn_depth_(attn_depth), normalize_(normalize) {
  values_ = Allocate(allocator_, batch_size_ * max_memory_steps_ * memory_depth_, values_ptr_, true);
  keys_ = Allocate(allocator_, batch_size_ * max_memory_steps_ * attn_depth_, keys_ptr_, true);
  processed_query_ = Allocate(allocator_, batch_size_ * attn_depth_, processed_query_ptr_, true);
  mem_seq_lengths_ = Allocate(allocator_, batch_size_, mem_seq_lengths_ptr_, true);

  ORT_ENFORCE(!normalize_, "not support normalize yet.");
}

template <typename T>
void BahdanauAttention<T>::SetWeights(
    const gsl::span<const T>& attn_weights,
    const gsl::span<const T>& query_layer_weights,
    const gsl::span<const T>& memory_layer_weights) {
  attention_v_ = attn_weights;                   //[attn_depth_]
  query_layer_weights_ = query_layer_weights;    //[query_depth_, attn_depth_]
  memory_layer_weights_ = memory_layer_weights;  //[memory_depth_, attn_depth_]
}

template <typename T>
const gsl::span<const T> BahdanauAttention<T>::Values() const {
  return values_;
}

template <typename T>
const gsl::span<const T> BahdanauAttention<T>::Keys() const {
  return keys_;
}

template <typename T>
int BahdanauAttention<T>::GetMaxMemorySteps() const {
  return max_memory_steps_;
}

template <typename T>
bool BahdanauAttention<T>::NeedPrevAlignment() const {
  return false;
}

template <typename T>
void BahdanauAttention<T>::PrepareMemory(
    const gsl::span<const T>& memory,
    const gsl::span<const int>& memory_sequence_lengths) {
  std::copy(memory.cbegin(), memory.cend(), values_.begin());
  if (memory_sequence_lengths.empty()) {
    std::fill(mem_seq_lengths_.begin(), mem_seq_lengths_.end(), max_memory_steps_);
  } else {
    std::copy(memory_sequence_lengths.cbegin(), memory_sequence_lengths.cend(), mem_seq_lengths_.begin());
  }

  for (int b = 0; b < batch_size_; b++) {
    int mem_steps = mem_seq_lengths_[b];
    ORT_ENFORCE(mem_steps <= max_memory_steps_ && mem_steps > 0,
                "Real memory steps ", mem_steps, " is not in (0, ", max_memory_steps_, "]");
  }

  math::GemmEx<T, CPUMathUtil>(CblasNoTrans, CblasNoTrans,
                               batch_size_ * max_memory_steps_, attn_depth_, memory_depth_, T{1.0},
                               memory.data(), memory_depth_,
                               memory_layer_weights_.data(), attn_depth_, T{0.0},
                               keys_.data(), attn_depth_, &CPUMathUtil::Instance());
}

template <typename T>
static void SoftmaxInplace(const gsl::span<T>& alignments) {
  T* x = alignments.data();
  size_t len = alignments.size();

  double sum = 0.0;

  for (size_t i = 0; i < len; i++) {
    T e = exp(x[i]);
    sum += e;
    x[i] = e;
  }

  if (sum == 0.0) {
    for (size_t i = 0; i < len; i++) {
      x[i] = static_cast<T>(1.0 / len);
    }
  } else {
    for (size_t i = 0; i < len; i++) {
      x[i] = static_cast<T>(x[i] / sum);
    }
  }
}

/**
  * Args:
  *     queries: Tensor, shape `[batch_size_, query_depth_]` to compare to keys.
  *     keys_: Processed memory, shape `[batch_size_, max_memory_step_, attn_depth_]`.
  */
template <typename T>
void BahdanauAttention<T>::Compute(
    const gsl::span<const T>& queries,
    const gsl::span<const T>&,  // Not used by bahdanau attention
    const gsl::span<T>& output,
    const gsl::span<T>& aligns) const {
  //process query in dense query layer without bias
  math::GemmEx<T, CPUMathUtil>(CblasNoTrans, CblasNoTrans,
                               batch_size_, attn_depth_, query_depth_, T{1.0},
                               queries.data(), query_depth_,
                               query_layer_weights_.data(), attn_depth_, T{0.0},
                               processed_query_.data(), attn_depth_, &CPUMathUtil::Instance());

  std::fill(aligns.begin(), aligns.end(), T{});

  // return math_ops.reduce_sum(v * math_ops.tanh(keys + processed_query), [2])
  for (int b = 0; b < batch_size_; b++) {
    T* alignments = aligns.data() + b * max_memory_steps_;
    const T* keys = keys_.data() + b * max_memory_steps_ * attn_depth_;
    const T* query = processed_query_.data() + b * attn_depth_;

    int mem_steps = mem_seq_lengths_[b];
    for (int step = 0; step < mem_steps; step++) {
      const T* keys_on_step = keys + step * attn_depth_;
      T* dest = alignments + step;

      // reduce_sum(v * tanh(keys[step] + query)) on last dimension
      *dest = T(0.0);
      for (int i = 0; i < attn_depth_; i++) {
        *dest += (attention_v_[i] * tanh(keys_on_step[i] + query[i]));
      }
    }

    SoftmaxInplace(gsl::span<T>{alignments, mem_steps});

    // Calculate the context
    auto outspan = output.subspan(b * memory_depth_);
    auto values = values_.subspan(b * max_memory_steps_ * memory_depth_);
    math::GemmEx<T, CPUMathUtil>(CblasNoTrans, CblasNoTrans,
                                 1, memory_depth_, max_memory_steps_, T{1.0},
                                 alignments, max_memory_steps_,
                                 values.data(), memory_depth_, T{0.0},
                                 outspan.data(), memory_depth_, &CPUMathUtil::Instance());
  }
}

template class BahdanauAttention<float>;

}  // namespace contrib
}  // namespace onnxruntime
