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

#pragma once

#include <memory>
#include "paddle/phi/backends/all_context.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/funcs/activation_functor.h"
#include "paddle/phi/kernels/funcs/blas/blas.h"
#include "paddle/phi/kernels/funcs/eigen/common.h"
#include "paddle/utils/optional.h"
namespace phi {

enum GRUActivationType { identity = 0, sigmoid = 1, tanh = 2, relu = 3 };

template <typename T, typename Device, typename X, typename Y>
void ActCompute(const int act_type, const Device& d, X x, Y y, Place place) {
  if (act_type == identity) {
    y.device(d) = x;
  } else if (act_type == sigmoid) {
    funcs::SigmoidFunctor<T>()(d, x, y);
  } else if (act_type == tanh) {
    funcs::TanhFunctor<T>()(d, x, y);
  } else if (act_type == relu) {
    if (place == CPUPlace())
      funcs::ReluCPUFunctor<T>()(d, x, y);
    else
      funcs::ReluCUDAFunctor<T>()(d, x, y);
  } else {
    PADDLE_THROW(common::errors::Unimplemented(
        "Unsupported activation type, only supports identity, sigmoid, tanh "
        "and relu."));
  }
}

#define ACT_COMPUTE ActCompute<T>

template <typename T, typename Context>
void GRUUnitKernel(const Context& dev_ctx,
                   const DenseTensor& input,
                   const DenseTensor& hidden_prev,
                   const DenseTensor& weight,
                   const optional<DenseTensor>& bias,
                   int activation,
                   int gate_activation,
                   bool origin_mode,
                   DenseTensor* gate,
                   DenseTensor* reset_hidden_prev,
                   DenseTensor* hidden) {
  auto* input_p = &input;
  auto* hidden_prev_p = &hidden_prev;

  dev_ctx.template Alloc<T>(gate);
  dev_ctx.template Alloc<T>(reset_hidden_prev);
  dev_ctx.template Alloc<T>(hidden);

  int64_t batch_size = input_p->dims()[0];
  int64_t frame_size = hidden_prev_p->dims()[1];

  auto x = EigenMatrix<T>::From(input);
  auto h_p = EigenMatrix<T>::From(hidden_prev);
  auto g = EigenMatrix<T>::From(*gate);
  auto r_h_p = EigenMatrix<T>::From(*reset_hidden_prev);
  auto h = EigenMatrix<T>::From(*hidden);
  auto& place = *dev_ctx.eigen_device();

  // calculate unactivated gate outputs
  if (bias) {
    auto b = EigenMatrix<T>::From(bias.get());
    g.device(place) =
        x + b.reshape(Eigen::array<int64_t, 2>({{1, frame_size * 3}}))
                .broadcast(Eigen::array<int64_t, 2>({{batch_size, 1}}));
  } else {
    g.device(place) = x;
  }
  const T* hidden_prev_data = hidden_prev.data<T>();
  const T* weight_data = weight.data<T>();
  T* gate_data = gate->data<T>();
  T* reset_hidden_prev_data = reset_hidden_prev->data<T>();
  auto blas = funcs::GetBlas<Context, T>(dev_ctx);
  blas.GEMM(false,
            false,
            batch_size,
            2 * frame_size,
            frame_size,
            1,
            hidden_prev_data,
            frame_size,
            weight_data,
            frame_size * 2,
            1,
            gate_data,
            frame_size * 3);

  // calculate activated gate
  Eigen::array<int64_t, 2> extents{{batch_size, frame_size}};
  Eigen::array<int64_t, 2> u_offsets{{0, 0}};
  ACT_COMPUTE(gate_activation,
              place,
              g.slice(u_offsets, extents),
              g.slice(u_offsets, extents),
              dev_ctx.GetPlace());
  auto u = g.slice(u_offsets, extents);  // update gate
  Eigen::array<int64_t, 2> r_offsets{{0, frame_size}};
  ACT_COMPUTE(gate_activation,
              place,
              g.slice(r_offsets, extents),
              g.slice(r_offsets, extents),
              dev_ctx.GetPlace());
  auto r = g.slice(r_offsets, extents);  // reset gate
  r_h_p.device(place) = r * h_p;         // reset previous hidden state
  blas.GEMM(false,
            false,
            batch_size,
            frame_size,
            frame_size,
            1,
            reset_hidden_prev_data,
            frame_size,
            weight_data + frame_size * frame_size * 2,
            frame_size,
            1,
            gate_data + frame_size * 2,
            frame_size * 3);

  Eigen::array<int64_t, 2> c_offsets{{0, frame_size * 2}};
  ACT_COMPUTE(activation,
              place,
              g.slice(c_offsets, extents),
              g.slice(c_offsets, extents),
              dev_ctx.GetPlace());
  auto c = g.slice(c_offsets, extents);  // output candidate

  // calculate final output
  if (origin_mode) {
    h.device(place) = c + u * (h_p - c);  // (1 - u) * c + u * h_p
  } else {
    h.device(place) = u * (c - h_p) + h_p;  // u * c + (1 - u) * h_p
  }
}

template <typename T,
          typename Device,
          typename X,
          typename Y,
          typename DX,
          typename DY>
void ActGradCompute(
    const int act_type, const Device& d, X x, Y y, DX dx, DY dy) {
  // x is dummy and won't be used even in Relu(use y instead)
  if (act_type == identity)
    dx.device(d) = dy;
  else if (act_type == sigmoid)
    funcs::SigmoidGradFunctor<T>()(d, x, y, dy, dx);
  else if (act_type == tanh)
    funcs::TanhGradFunctor<T>()(d, x, y, dy, dx);
  else if (act_type == relu)
    funcs::ReluGradFunctor<T>()(d, x, y, dy, dx);
  else
    PADDLE_THROW(common::errors::Unimplemented(
        "Unsupported activation type, only supports identity, sigmoid, tanh "
        "and relu."));
}

#define ACT_GRAD_COMPUTE ActGradCompute<T>

template <typename T, typename Context>
void GRUUnitGradKernel(const Context& dev_ctx,
                       const DenseTensor& input,
                       const DenseTensor& hidden_prev,
                       const DenseTensor& weight,
                       const optional<DenseTensor>& bias,
                       const DenseTensor& gate,
                       const DenseTensor& reset_hidden_prev,
                       const DenseTensor& hidden_grad,
                       int activation,
                       int gate_activation,
                       bool origin_mode,
                       DenseTensor* input_grad,
                       DenseTensor* hidden_prev_grad,
                       DenseTensor* weight_grad,
                       DenseTensor* bias_grad) {
  DenseTensor gate_grad;
  DenseTensor reset_hidden_prev_grad;

  const T* hidden_prev_data = hidden_prev.data<T>();
  const T* weight_data = weight.data<T>();
  gate_grad.Resize(input.dims());
  T* gate_grad_data = dev_ctx.template Alloc<T>(&gate_grad);
  const T* reset_hidden_prev_data = reset_hidden_prev.data<T>();
  reset_hidden_prev_grad.Resize(reset_hidden_prev.dims());
  T* reset_hidden_prev_grad_data =
      dev_ctx.template Alloc<T>(&reset_hidden_prev_grad);

  auto h_p = EigenMatrix<T>::From(hidden_prev);
  auto g = EigenMatrix<T>::From(gate);
  auto d_h = EigenMatrix<T>::From(hidden_grad);
  auto d_g = EigenMatrix<T>::From(gate_grad);
  auto d_r_h_p = EigenMatrix<T>::From(reset_hidden_prev_grad);
  auto& place = *dev_ctx.eigen_device();

  int64_t batch_size = input.dims()[0];
  int64_t frame_size = hidden_prev.dims()[1];

  Eigen::array<int64_t, 2> extents{{batch_size, frame_size}};
  Eigen::array<int64_t, 2> u_offsets{{0, 0}};
  auto u = g.slice(u_offsets, extents);  // update gate
  Eigen::array<int64_t, 2> r_offsets{{0, frame_size}};
  auto r = g.slice(r_offsets, extents);  // reset gate
  Eigen::array<int64_t, 2> c_offsets{{0, frame_size * 2}};
  auto c = g.slice(c_offsets, extents);  // output candidate

  // backward for unactivated update gate
  if (origin_mode) {
    ACT_GRAD_COMPUTE(gate_activation,
                     place,
                     u,
                     u,
                     d_g.slice(u_offsets, extents),
                     d_h * (h_p - c));
    // backward for unactivated output candidate
    ACT_GRAD_COMPUTE(
        activation, place, c, c, d_g.slice(c_offsets, extents), d_h * (1 - u));
  } else {
    ACT_GRAD_COMPUTE(gate_activation,
                     place,
                     u,
                     u,
                     d_g.slice(u_offsets, extents),
                     d_h * (c - h_p));
    // backward for unactivated output candidate
    ACT_GRAD_COMPUTE(
        activation, place, c, c, d_g.slice(c_offsets, extents), d_h * u);
  }
  // backward for reset_hidden_prev
  auto blas = funcs::GetBlas<Context, T>(dev_ctx);
  blas.GEMM(false,
            true,
            batch_size,
            frame_size,
            frame_size,
            1,
            gate_grad_data + frame_size * 2,
            frame_size * 3,
            weight_data + frame_size * frame_size * 2,
            frame_size,
            0,
            reset_hidden_prev_grad_data,
            frame_size);
  // backward for unactivated reset gate
  ACT_GRAD_COMPUTE(gate_activation,
                   place,
                   r,
                   r,
                   d_g.slice(r_offsets, extents),
                   d_r_h_p * h_p);
  // backward for weight
  if (weight_grad) {
    T* weight_grad_data = dev_ctx.template Alloc<T>(weight_grad);
    // backward for state_weight
    blas.GEMM(true,
              false,
              frame_size,
              frame_size,
              batch_size,
              1,
              reset_hidden_prev_data,
              frame_size,
              gate_grad_data + frame_size * 2,
              frame_size * 3,
              0,
              weight_grad_data + frame_size * frame_size * 2,
              frame_size);

    // backward for update_gate_weight and reset_gate_weight
    blas.GEMM(true,
              false,
              frame_size,
              frame_size * 2,
              batch_size,
              1,
              hidden_prev_data,
              frame_size,
              gate_grad_data,
              frame_size * 3,
              0,
              weight_grad_data,
              frame_size * 2);
  }
  // backward for hidden_prev
  if (hidden_prev_grad) {
    T* hidden_prev_grad_data = dev_ctx.template Alloc<T>(hidden_prev_grad);
    auto d_h_p = EigenMatrix<T>::From(*hidden_prev_grad);
    if (origin_mode) {
      d_h_p.device(place) = d_r_h_p * r + d_h * u;
    } else {
      d_h_p.device(place) = d_r_h_p * r + d_h * (1 - u);
    }
    blas.GEMM(false,
              true,
              batch_size,
              frame_size,
              frame_size * 2,
              1,
              gate_grad_data,
              frame_size * 3,
              weight_data,
              frame_size * 2,
              1,
              hidden_prev_grad_data,
              frame_size);
  }
  // backward for input
  if (input_grad) {
    dev_ctx.template Alloc<T>(input_grad);
    auto d_x = EigenMatrix<T>::From(*input_grad);
    d_x.device(place) = d_g;
  }
  // backward for bias
  if (bias_grad) {
    dev_ctx.template Alloc<T>(bias_grad);
    auto d_b = EigenVector<T>::Flatten(*bias_grad);
    d_b.device(place) = d_g.sum(Eigen::array<int64_t, 1>({{0}}));
  }
}
}  // namespace phi