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

#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/funcs/blas/blas.h"
#include "paddle/phi/kernels/funcs/detail/activation_functions.h"
#include "paddle/phi/kernels/funcs/lstm_compute.h"
#include "paddle/phi/kernels/funcs/lstm_utils.h"

namespace phi {

template <typename T, typename Context>
void LSTMKernel(const Context& dev_ctx,
                const DenseTensor& input,
                const optional<DenseTensor>& h0,
                const optional<DenseTensor>& c0,
                const DenseTensor& weight,
                const DenseTensor& bias,
                bool use_peepholes,
                bool is_reverse,
                bool is_test,
                const std::string& gate_activation,
                const std::string& cell_activation,
                const std::string& candidate_activation,
                DenseTensor* hidden,
                DenseTensor* cell,
                DenseTensor* batch_gate,
                DenseTensor* batch_cell_pre_act) {
  auto* hidden_t0 = h0.get_ptr();
  auto* cell_t0 = c0.get_ptr();

  DenseTensor* batch_gate_new = nullptr;
  DenseTensor batch_gate_temp;
  if (is_test) {
    batch_gate_new = &batch_gate_temp;
    batch_gate_new->Resize(input.dims());
  } else {
    batch_gate_new = batch_gate;
  }

  dev_ctx.template Alloc<T>(batch_gate_new);
  dev_ctx.template Alloc<T>(hidden);
  dev_ctx.template Alloc<T>(cell);

  funcs::DenseTensor2BatchFunctor<Context, T> to_batch;
  to_batch(dev_ctx, input, batch_gate_new, true, is_reverse);

  auto in_dims = input.dims();
  int64_t frame_size = in_dims[1] / 4;
  DDim dims({in_dims[0], frame_size});

  if (bias.initialized()) {
    DenseTensor b = bias;
    b.Resize({bias.numel(), 1});
    DenseTensor gate_bias = b.Slice(0, 4 * frame_size);
    funcs::RowwiseAdd<Context, T> add_bias;
    add_bias(dev_ctx, *batch_gate_new, gate_bias, batch_gate_new);
  }

  funcs::LstmMetaValue<T> lstm_value;
  if (bias.initialized() && use_peepholes) {
    T* bias_data = const_cast<T*>(bias.data<T>());
    // the code style in LstmMetaValue will be updated later.

    lstm_value.check_ig = bias_data + 4 * frame_size;
    lstm_value.check_fg = lstm_value.check_ig + frame_size;
    lstm_value.check_og = lstm_value.check_fg + frame_size;
  } else {
    lstm_value.check_ig = nullptr;
    lstm_value.check_fg = nullptr;
    lstm_value.check_og = nullptr;
  }
  lstm_value.prev_state_value = nullptr;
  DenseTensor ordered_c0;

  Vector<size_t> order(batch_gate_new->lod()[2]);

  if (cell_t0) {
    // Since the batch computing for LSTM reorders the input sequence
    // according to their length. The initialized cell state also needs
    // to reorder.
    ReorderInitState<Context, T>(dev_ctx, *cell_t0, order, &ordered_c0, true);
    lstm_value.prev_state_value = ordered_c0.data<T>();
  }

  // Use the local variable as here.
  DenseTensor batch_hidden, batch_cell, batch_cell_pre_act_temp;
  DenseTensor* batch_cell_pre_act_p;
  if (is_test) {
    batch_cell_pre_act_p = &batch_cell_pre_act_temp;
  } else {
    batch_cell_pre_act_p = batch_cell_pre_act;
  }
  batch_hidden.Resize(dims);
  batch_cell.Resize(dims);
  dev_ctx.template Alloc<T>(&batch_hidden);
  dev_ctx.template Alloc<T>(&batch_cell);
  batch_cell_pre_act_p->Resize(dims);
  dev_ctx.template Alloc<T>(batch_cell_pre_act_p);

  auto batch_starts = batch_gate_new->lod()[0];
  size_t num_batch = batch_starts.size() - 1;
  auto gate_act = funcs::detail::GetActivationType(gate_activation);
  auto cell_act = funcs::detail::GetActivationType(cell_activation);
  auto cand_act = funcs::detail::GetActivationType(candidate_activation);

  auto blas = funcs::GetBlas<Context, T>(dev_ctx);
  for (size_t n = 0; n < num_batch; n++) {
    int bstart = static_cast<int>(batch_starts[n]);
    int bend = static_cast<int>(batch_starts[n + 1]);

    DenseTensor gate_t = batch_gate_new->Slice(bstart, bend);
    DenseTensor out_t = batch_hidden.Slice(bstart, bend);
    DenseTensor cell_t = batch_cell.Slice(bstart, bend);
    DenseTensor cell_pre_act_t = batch_cell_pre_act_p->Slice(bstart, bend);

    int cur_batch_size = bend - bstart;

    if (n > 0) {
      int pre_h_start = static_cast<int>(batch_starts[n - 1]);
      int pre_h_end = pre_h_start + cur_batch_size;
      auto pre_hidden_t = batch_hidden.Slice(pre_h_start, pre_h_end);
      blas.MatMul(pre_hidden_t,
                  false,
                  weight,
                  false,
                  static_cast<T>(1.0),
                  &gate_t,
                  static_cast<T>(1.0));
    } else if (hidden_t0 != nullptr) {
      // If n == 0 and there is no initialized hidden state, that is to say
      // the H0 is zeros, the calculation W_h * H0 will be skipped.
      // If n == 0 and there is initialized hidden state, calculate W_h * H0.

      // Since the batch computing for LSTM reorders the input sequence
      // according to their length. The initialized hidden state also needs
      // to reorder.
      DenseTensor ordered_h0;
      ReorderInitState<Context, T>(
          dev_ctx, *hidden_t0, order, &ordered_h0, true);
      blas.MatMul(ordered_h0,
                  false,
                  weight,
                  false,
                  static_cast<T>(1.0),
                  &gate_t,
                  static_cast<T>(1.0));
    }

    lstm_value.gate_value = gate_t.data<T>();
    lstm_value.output_value = out_t.data<T>();
    lstm_value.state_value = cell_t.data<T>();
    lstm_value.state_active_value = cell_pre_act_t.data<T>();
    T cell_clip = 0.0;
    funcs::LstmUnitFunctor<Context, T>::compute(dev_ctx,
                                                lstm_value,
                                                frame_size,
                                                cur_batch_size,
                                                cell_clip,
                                                gate_act,
                                                cell_act,
                                                cand_act);
    lstm_value.prev_state_value = lstm_value.state_value;
  }

  funcs::Batch2DenseTensorFunctor<Context, T> to_seq;
  batch_hidden.set_lod(batch_gate_new->lod());
  // restore the output hidden in DenseTensor from the batch hidden
  to_seq(dev_ctx, batch_hidden, hidden);

  batch_cell.set_lod(batch_gate_new->lod());
  // restore the output cell state in DenseTensor from the batch cell
  to_seq(dev_ctx, batch_cell, cell);
}

template <typename T, typename Context>
void LSTMGradKernel(const Context& dev_ctx,
                    const DenseTensor& input_in,
                    const optional<DenseTensor>& h0_in,
                    const optional<DenseTensor>& c0_in,
                    const DenseTensor& weight_in,
                    const DenseTensor& bias_in,
                    const DenseTensor& hidden_in,
                    const DenseTensor& cell_in,
                    const DenseTensor& batch_gate_in,
                    const DenseTensor& batch_cell_pre_act_in,
                    const DenseTensor& hidden_grad,
                    bool use_peepholes,
                    bool is_reverse,
                    bool is_test,
                    const std::string& gate_activation,
                    const std::string& cell_activation,
                    const std::string& candidate_activation,
                    DenseTensor* input_grad,
                    DenseTensor* h0_grad,
                    DenseTensor* c0_grad,
                    DenseTensor* weight_grad,
                    DenseTensor* bias_grad) {
  auto* input = &input_in;
  auto* weight = &weight_in;
  auto* bias = &bias_in;

  auto* hidden_out = &hidden_in;
  auto* cell_out = &cell_in;

  auto* batch_gate = &batch_gate_in;
  auto* batch_cell_pre_act = &batch_cell_pre_act_in;

  auto* hidden_g = &hidden_grad;

  auto* in_g = input_grad;
  auto* weight_g = weight_grad;
  auto* bias_g = bias_grad;

  auto* h0 = h0_in.get_ptr();
  auto* c0 = c0_in.get_ptr();

  auto* h0_g = h0_grad;
  auto* c0_g = c0_grad;

  funcs::SetConstant<Context, T> zero;
  if (weight_g) {
    dev_ctx.template Alloc<T>(weight_g);
    zero(dev_ctx, weight_g, static_cast<T>(0.0));
  }

  // ordered_h0/c0 is the reordered hidden/cell initialization.
  // ordered_h0_g/c0_g is the reordered gradient of hidden/cell
  // initialization.
  DenseTensor ordered_h0, ordered_c0, ordered_h0_g, ordered_c0_g;
  Vector<size_t> order(batch_gate->lod()[2]);

  if (c0) {
    ReorderInitState<Context, T>(dev_ctx, *c0, order, &ordered_c0, true);
  }
  if (c0 && c0_g) {
    ordered_c0_g.Resize(c0_g->dims());
    dev_ctx.template Alloc<T>(&ordered_c0_g);
  }

  auto in_dims = input->dims();
  auto out_dims = hidden_g->dims();
  int64_t frame_size = in_dims[1] / 4;
  PADDLE_ENFORCE_EQ(frame_size,
                    out_dims[1],
                    common::errors::InvalidArgument(
                        "The second dimension of Input(hidden_grad) should be "
                        "%d, but received %d in LSTM@GRAD operator.",
                        frame_size,
                        out_dims[1]));

  funcs::LstmMetaValue<T> lstm_value;
  if (bias && use_peepholes) {
    T* bias_data = const_cast<T*>(bias->data<T>());
    lstm_value.check_ig = bias_data + 4 * frame_size;
    lstm_value.check_fg = lstm_value.check_ig + frame_size;
    lstm_value.check_og = lstm_value.check_fg + frame_size;
  } else {
    lstm_value.check_ig = nullptr;
    lstm_value.check_fg = nullptr;
    lstm_value.check_og = nullptr;
  }

  funcs::LstmMetaGrad<T> lstm_grad;

  if (bias && bias_g) {
    dev_ctx.template Alloc<T>(bias_g);
    zero(dev_ctx, bias_g, static_cast<T>(0.0));
  }
  if (bias && bias_g && use_peepholes) {
    T* bias_g_data = bias_g->data<T>();
    lstm_grad.check_ig_grad = bias_g_data + 4 * frame_size;
    lstm_grad.check_fg_grad = lstm_grad.check_ig_grad + frame_size;
    lstm_grad.check_og_grad = lstm_grad.check_fg_grad + frame_size;
  } else {
    lstm_grad.check_ig_grad = nullptr;
    lstm_grad.check_fg_grad = nullptr;
    lstm_grad.check_og_grad = nullptr;
  }

  funcs::DenseTensor2BatchFunctor<Context, T> to_batch;

  auto ToBatch = [&batch_gate, &to_batch](const Context& dev_ctx,
                                          const DenseTensor& src,
                                          const DDim& dims,
                                          DenseTensor& dst) {
    dst.Resize(dims);
    dev_ctx.template Alloc<T>(&dst);
    dst.set_lod(batch_gate->lod());
    to_batch(dev_ctx, src, &dst, false);
  };

  DenseTensor batch_hidden, batch_hidden_g, batch_cell;
  ToBatch(dev_ctx, *hidden_out, out_dims, batch_hidden);
  ToBatch(dev_ctx, *hidden_g, out_dims, batch_hidden_g);
  ToBatch(dev_ctx, *cell_out, out_dims, batch_cell);

  DenseTensor batch_cell_g, batch_gate_g;
  batch_cell_g.Resize(out_dims);
  dev_ctx.template Alloc<T>(&batch_cell_g);
  // TODO(qingqing) support the case output cell has gradient.
  // to_batch(dev_ctx, *cell_g, batch_cell_g, false);
  zero(dev_ctx, &batch_cell_g, static_cast<T>(0.0));
  batch_gate_g.Resize(batch_gate->dims());
  dev_ctx.template Alloc<T>(&batch_gate_g);
  batch_gate_g.set_lod(batch_gate->lod());

  auto gate_act = funcs::detail::GetActivationType(gate_activation);
  auto cell_act = funcs::detail::GetActivationType(cell_activation);
  auto cand_act = funcs::detail::GetActivationType(candidate_activation);

  auto batch_starts = batch_gate->lod()[0];
  size_t num_batch = batch_starts.size() - 1;
  auto blas = funcs::GetBlas<Context, T>(dev_ctx);
  for (int n = static_cast<int>(num_batch) - 1; n >= 0; n--) {
    int bstart = static_cast<int>(batch_starts[n]);
    int bend = static_cast<int>(batch_starts[n + 1]);

    DenseTensor gate = batch_gate->Slice(bstart, bend);
    DenseTensor cell = batch_cell.Slice(bstart, bend);
    DenseTensor cell_pre_act = batch_cell_pre_act->Slice(bstart, bend);
    lstm_value.gate_value = gate.data<T>();
    lstm_value.state_value = cell.data<T>();
    lstm_value.state_active_value = cell_pre_act.data<T>();

    DenseTensor out_g = batch_hidden_g.Slice(bstart, bend);
    DenseTensor gate_g = batch_gate_g.Slice(bstart, bend);
    DenseTensor cell_g = batch_cell_g.Slice(bstart, bend);
    lstm_grad.state_grad = cell_g.data<T>();
    lstm_grad.gate_grad = gate_g.data<T>();
    lstm_grad.output_grad = out_g.data<T>();

    if (n > 0) {
      int bstart_pre = static_cast<int>(batch_starts[n - 1]);
      DenseTensor cell_pre = batch_cell.Slice(bstart_pre, bstart);
      DenseTensor cell_pre_g = batch_cell_g.Slice(bstart_pre, bstart);
      lstm_value.prev_state_value = cell_pre.data<T>();
      lstm_grad.prev_state_grad = cell_pre_g.data<T>();
    } else {
      lstm_value.prev_state_value = c0 ? ordered_c0.data<T>() : nullptr;
      lstm_grad.prev_state_grad = c0_g ? ordered_c0_g.data<T>() : nullptr;
    }

    // lstm_value.output_value not used in bp, set to nullptr
    // lstm_grad.state_active_grad not used in bp, set to nullptr
    lstm_value.output_value = nullptr;
    lstm_grad.state_active_grad = nullptr;
    int cur_batch_size = bend - bstart;
    T cell_clip = 0.0;
    funcs::LstmUnitGradFunctor<Context, T>::compute(dev_ctx,
                                                    lstm_value,
                                                    lstm_grad,
                                                    frame_size,
                                                    cur_batch_size,
                                                    cell_clip,
                                                    gate_act,
                                                    cell_act,
                                                    cand_act);

    if (n > 0) {
      int pre_h_start = static_cast<int>(batch_starts[n - 1]);
      int pre_h_end = pre_h_start + cur_batch_size;
      auto pre_hidden_g = batch_hidden_g.Slice(pre_h_start, pre_h_end);
      blas.MatMul(gate_g,
                  false,
                  *weight,
                  true,
                  static_cast<T>(1.0),
                  &pre_hidden_g,
                  static_cast<T>(1.0));
      if (weight_g) {
        /* backward weight */
        auto pre_hidden = batch_hidden.Slice(pre_h_start, pre_h_end);
        blas.MatMul(pre_hidden,
                    true,
                    gate_g,
                    false,
                    static_cast<T>(1.0),
                    weight_g,
                    static_cast<T>(1.0));
      }
    } else {
      if (h0 && weight_g) {
        ReorderInitState<Context, T>(dev_ctx, *h0, order, &ordered_h0, true);
        blas.MatMul(ordered_h0,
                    true,
                    gate_g,
                    false,
                    static_cast<T>(1.0),
                    weight_g,
                    static_cast<T>(1.0));
      }
      if (h0 && h0_g) {
        ordered_h0_g.Resize(h0_g->dims());
        dev_ctx.template Alloc<T>(&ordered_h0_g);
        blas.MatMul(gate_g,
                    false,
                    *weight,
                    true,
                    static_cast<T>(1.0),
                    &ordered_h0_g,
                    static_cast<T>(0.0));
      }
    }
  }

  funcs::Batch2DenseTensorFunctor<Context, T> to_seq;
  if (in_g) {
    /* backward data */
    dev_ctx.template Alloc<T>(in_g);
    to_seq(dev_ctx, batch_gate_g, in_g);
  }
  if (bias && bias_g) {
    /* backward bias */
    DenseTensor b_g = *bias_g;
    b_g.Resize({bias_g->numel(), 1});
    DenseTensor gate_bias_g = b_g.Slice(0, 4 * frame_size);
    funcs::ColwiseSum<Context, T> col_sum;
    col_sum(dev_ctx, batch_gate_g, &gate_bias_g);
  }

  if (h0 && h0_g) {
    ReorderInitState<Context, T>(dev_ctx, ordered_h0_g, order, h0_g, false);
  }
  if (c0 && c0_g) {
    ReorderInitState<Context, T>(dev_ctx, ordered_c0_g, order, c0_g, false);
  }
}
}  // namespace phi