// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
//
// 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 <string>
#include "paddle/phi/backends/cpu/cpu_info.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/funcs/blas/blas.h"
#include "paddle/phi/kernels/funcs/cpu_vec.h"
#include "paddle/phi/kernels/funcs/fc_functor.h"
#include "paddle/utils/optional.h"

namespace phi {

// y[i] = (x[i] + bias[0]) > 0 ? (x[i] + bias[0]) : 0;
template <typename T>
inline void bias_relu(const int n, const T* x, const T* bias, T* y) {
  if (bias) {
    funcs::vec_add_bias<T, backends::cpu::avx>(n, *bias, x, y);
    funcs::vec_relu<T, backends::cpu::avx>(n, y, y);
  } else {
    funcs::vec_relu<T, backends::cpu::avx>(n, x, y);
  }
}

template <typename T>
inline void vec_softmax(const int n, const T* x, T* y) {
  T scalar = x[0];
  // max
  for (int i = 1; i < n; ++i) {
    scalar = scalar < x[i] ? x[i] : scalar;
  }
  funcs::vec_add_bias<T, backends::cpu::avx>(n, -scalar, x, y);  // sub
  funcs::vec_exp<T>(n, y, y);                                    // exp
  // sum
  scalar = T(0);
  for (int i = 0; i < n; ++i) {
    scalar += y[i];
  }
  funcs::vec_scal<T>(n, static_cast<T>(1) / scalar, y);  // scale
}

template <typename T, typename Context>
void AttentionLSTMKernel(const Context& dev_ctx,
                         const DenseTensor& x_in,
                         const DenseTensor& c0_in,
                         const optional<DenseTensor>& h0_in,
                         const DenseTensor& attention_weight_in,
                         const optional<DenseTensor>& attention_bias_in,
                         const optional<DenseTensor>& attention_scalar_in,
                         const optional<DenseTensor>& attention_scalar_bias_in,
                         const DenseTensor& lstm_weight_in,
                         const DenseTensor& lstm_bias_in,
                         const std::string& gate_activation,
                         const std::string& cell_activation,
                         const std::string& candidate_activation,
                         DenseTensor* hidden,
                         DenseTensor* cell,
                         DenseTensor* attentioned_x,
                         DenseTensor* attention_fc_out,
                         DenseTensor* lstm_x,
                         DenseTensor* lstm_out) {
  auto* x = &x_in;
  auto* h0 = h0_in.get_ptr();
  auto* c0 = &c0_in;
  auto* atten_w = &attention_weight_in;
  auto* atten_b = attention_bias_in.get_ptr();
  auto* atten_scalar = attention_scalar_in.get_ptr();
  auto* atten_scalar_bias = attention_scalar_bias_in.get_ptr();
  auto* lstm_w = &lstm_weight_in;
  auto* lstm_b = &lstm_bias_in;

  auto* hidden_out = hidden;
  auto* cell_out = cell;
  auto* atted_x = attentioned_x;
  auto* fc_out = attention_fc_out;

  // some shape should be reshape here since infershape can not get lod info
  auto x_lod = x->lod();
  const int N = static_cast<int>(x_lod[0].size() - 1);  // batch size
  auto x_dims = x->dims();                              // T x M
  auto w_dims = lstm_w->dims();                         // (D+M) x 4D
  const int total_T = static_cast<int>(x_dims[0]);
  const int M = static_cast<int>(x_dims[1]);      // x frame size
  const int D = static_cast<int>(w_dims[1] / 4);  // gate frame size
  const int D2 = static_cast<int>(D * 2);
  const int D3 = static_cast<int>(D * 3);
  const int D4 = static_cast<int>(w_dims[1]);
  int max_seq_len = static_cast<int>(x_lod[0][1]);
  for (int i = 1; i < N; ++i) {
    int len = static_cast<int>(x_lod[0][i + 1] - x_lod[0][i]);
    max_seq_len = max_seq_len < len ? len : max_seq_len;
  }
  PADDLE_ENFORCE_EQ(
      x_lod.size(),
      1UL,
      common::errors::InvalidArgument("Input(X)'s lod size must be 1."));
  PADDLE_ENFORCE_EQ(
      c0->dims()[0],
      N,
      common::errors::InvalidArgument("C0 dims should be %d x %d.", N, D));
  fc_out->Resize({max_seq_len, 1});

  std::function<void(const int, const T*, T*)> act_gate, act_cell, act_cand;
  auto& act_gate_str = gate_activation;
  auto& act_cell_str = cell_activation;
  auto& act_cand_str = candidate_activation;
  if (backends::cpu::MayIUse(backends::cpu::avx)) {
    funcs::VecActivations<T, backends::cpu::avx> act_functor;
    act_gate = act_functor(act_gate_str);
    act_cell = act_functor(act_cell_str);
    act_cand = act_functor(act_cand_str);
  } else {
    funcs::VecActivations<T, backends::cpu::isa_any> act_functor;
    act_gate = act_functor(act_gate_str);
    act_cell = act_functor(act_cell_str);
    act_cand = act_functor(act_cand_str);
  }

  const T* x_data = x->data<T>();
  const T* h0_data = h0 ? h0->data<T>() : NULL;
  const T* c0_data = c0->data<T>();
  const T* lstm_w_data = lstm_w->data<T>();
  const T* lstm_b_data = lstm_b->data<T>();
  const T* atten_w_data = atten_w->data<T>();
  const T* atten_b_data = atten_b ? atten_b->data<T>() : NULL;
  const T* atten_scalar_data = atten_scalar ? atten_scalar->data<T>() : NULL;
  const T* atten_scalar_bias_data =
      atten_scalar_bias ? atten_scalar_bias->data<T>() : NULL;

  T* hidden_out_data = dev_ctx.template Alloc<T>(hidden_out);
  T* cell_out_data = dev_ctx.template Alloc<T>(cell_out);
  T* atted_x_data = dev_ctx.template Alloc<T>(atted_x);
  T* fc_out_data = dev_ctx.template Alloc<T>(fc_out);
  T* lstm_x_data = dev_ctx.template Alloc<T>(lstm_x);
  T* lstm_out_data = dev_ctx.template Alloc<T>(lstm_out);

  // x(TxM) * fc (Mx1) part of atten_wgt(M+D)x1
  auto blas = funcs::GetBlas<CPUContext, T>(dev_ctx);

  funcs::FCFunctor<Context, T> fc;
  fc(dev_ctx, total_T, 1, M, x_data, atten_w_data, atted_x_data, atten_b_data);

  const T* cur_atten_x_data = atted_x_data;
  const T* cur_x_data = x_data;
  const T* prev_cell_data = NULL;
  const T* prev_hidden_data = NULL;
  T* cur_cell_out_data = cell_out_data;
  T* cur_hidden_out_data = hidden_out_data;
  for (int i = 0; i < N; ++i) {
    int seq_len = static_cast<int>(x_lod[0][i + 1] - x_lod[0][i]);
    prev_cell_data = c0_data + i * D;
    prev_hidden_data = h0_data ? h0_data + i * D : NULL;
    for (int step = 0; step < seq_len; ++step) {
      /// 1. compute attention vector
      // 1a. prev_cell(1xD) * fc(D) rest part of atten_wgt
      T prev_cell_bias = blas.DOT(D, prev_cell_data, atten_w_data + M);
      // 1b. add cell bias and relu
      bias_relu<T>(seq_len, cur_atten_x_data, &prev_cell_bias, fc_out_data);
      // 1c. fc scalar
      if (atten_scalar_data) {
        blas.SCAL(seq_len, *atten_scalar_data, fc_out_data);
        bias_relu<T>(seq_len, fc_out_data, atten_scalar_bias_data, fc_out_data);
      }
      // 1d. softmax
      vec_softmax<T>(seq_len, fc_out_data, fc_out_data);
      // mul x(seq_len*M) and sum pool
      fc(dev_ctx, 1, M, seq_len, fc_out_data, cur_x_data, lstm_x_data);

      /// 2. compute LSTM step
      // lstm weight : concat[forget , input , output , tilde]
      // shape : (D + M) x (4 * D)
      // fc inputX(1xM) * weightX(M*(4D))  => 1 x 4D
      blas.MatMul(1, D4, M, lstm_x_data, lstm_w_data + D * D4, lstm_out_data);
      if (prev_hidden_data) {
        blas.GEMM(CblasNoTrans,
                  CblasNoTrans,
                  1,
                  D4,
                  D,
                  static_cast<T>(1),
                  prev_hidden_data,
                  D,
                  lstm_w_data,
                  D4,
                  static_cast<T>(1),
                  lstm_out_data,
                  D4);
      }
      // since input is 1xM, so can use add bias
      blas.VADD(D4, lstm_b_data, lstm_out_data, lstm_out_data);

      // gate act: sigmoid
      act_gate(D3, lstm_out_data, lstm_out_data);
      // candidate act: tanh
      act_cand(D, lstm_out_data + D3, lstm_out_data + D3);

      // a = forget * prev_cell
      blas.VMUL(D, lstm_out_data, prev_cell_data, lstm_out_data);

      // b = input * tilde
      blas.VMUL(D, lstm_out_data + D, lstm_out_data + D3, lstm_out_data + D);

      // cell_out = a + b
      blas.VADD(D, lstm_out_data, lstm_out_data + D, cur_cell_out_data);

      // state act tanh(cell_out) * output_gate
      act_cell(D, cur_cell_out_data, lstm_out_data);
      blas.VMUL(D, lstm_out_data, lstm_out_data + D2, cur_hidden_out_data);

      prev_hidden_data = cur_hidden_out_data;
      prev_cell_data = cur_cell_out_data;
      cur_cell_out_data = cur_cell_out_data + D;
      cur_hidden_out_data = cur_hidden_out_data + D;
    }
    cur_x_data = cur_x_data + seq_len * M;
    cur_atten_x_data = cur_atten_x_data + seq_len;
  }
}
}  // namespace phi

PD_REGISTER_KERNEL(
    attention_lstm, CPU, ALL_LAYOUT, phi::AttentionLSTMKernel, float, double) {}