// Copyright (c) 2025 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 "paddle/phi/infermeta/binary.h"
#include "paddle/phi/kernels/funcs/blas/blas.h"
#include "paddle/phi/kernels/funcs/math_function.h"
#include "paddle/phi/kernels/funcs/matrix_solve.h"
#include "paddle/phi/kernels/impl/lu_kernel_impl.h"
#include "paddle/phi/kernels/lu_solve_grad_kernel.h"
#include "paddle/phi/kernels/lu_solve_kernel.h"
#include "paddle/phi/kernels/lu_unpack_kernel.h"
#include "paddle/phi/kernels/triangular_solve_kernel.h"

namespace phi {

template <typename T, typename Context>
DenseTensor GetMH(const Context& dev_ctx, const DenseTensor x) {
  DenseTensor x_mH;
  Tensor_Conj<Context, T>(dev_ctx, x, &x_mH);
  return Transpose2DTo6D<Context, T>(dev_ctx, x_mH);
}

template <typename T, typename Context>
void LuSolveGradKernel(const Context& dev_ctx,
                       const DenseTensor& b,
                       const DenseTensor& lu,
                       const DenseTensor& pivots,
                       const DenseTensor& out,
                       const DenseTensor& out_grad,
                       const std::string& trans,
                       DenseTensor* b_grad,
                       DenseTensor* lu_grad) {
  if (b_grad != nullptr) {
    dev_ctx.template Alloc<T>(b_grad);
    std::string trans_t = (trans == "N") ? "T" : "N";
    LuSolveKernel<T, Context>(dev_ctx, out_grad, lu, pivots, trans_t, b_grad);
  }

  if (lu_grad != nullptr) {
    dev_ctx.template Alloc<T>(lu_grad);

    DenseTensor p, l, u, l_mH, u_mH;
    MetaTensor meta_p(&p);
    MetaTensor meta_l(&l);
    MetaTensor meta_u(&u);
    bool unpack_pivots = (trans == "N") ? false : true;
    LUUnpackInferMeta(
        lu, pivots, true, unpack_pivots, &meta_p, &meta_l, &meta_u);
    LUUnpackKernel<T, Context>(
        dev_ctx, lu, pivots, true, unpack_pivots, &p, &l, &u);
    l_mH = GetMH<T, Context>(dev_ctx, l);
    u_mH = GetMH<T, Context>(dev_ctx, u);
    if (trans == "N") {
      // gR = U^{-H}op_2(-gX)op_2(X)^Ha
      DenseTensor gR, psi_tmp, out_mH;
      out_mH = GetMH<T, Context>(dev_ctx, out);

      auto blas = funcs::GetBlas<Context, T>(dev_ctx);
      auto out_grad_dims = out_grad.dims();
      auto mat_dim_l = funcs::CreateMatrixDescriptor(out_grad_dims, 0, false);
      auto out_mH_dims = out_mH.dims();
      auto mat_dim_g = funcs::CreateMatrixDescriptor(out_mH_dims, 0, false);
      psi_tmp.Resize(lu.dims());
      dev_ctx.template Alloc<T>(&psi_tmp);
      blas.MatMul(out_grad,
                  mat_dim_l,
                  out_mH,
                  mat_dim_g,
                  static_cast<T>(-1),
                  &psi_tmp,
                  static_cast<T>(0));
      TriangularSolveKernel<T, Context>(
          dev_ctx, u_mH, psi_tmp, false, false, false, &gR);

      // gL = (L^{-H} gR U^H).tril(-1)
      DenseTensor mul_tmp, gL;
      auto gr_dims = gR.dims();
      auto mat_dim_r = funcs::CreateMatrixDescriptor(gr_dims, 0, false);
      auto gu_dims = u_mH.dims();
      auto mat_dim_u = funcs::CreateMatrixDescriptor(gu_dims, 0, false);
      mul_tmp.Resize(gr_dims);
      dev_ctx.template Alloc<T>(&mul_tmp);
      blas.MatMul(gR,
                  mat_dim_r,
                  u_mH,
                  mat_dim_u,
                  static_cast<T>(1),
                  &mul_tmp,
                  static_cast<T>(0));
      TriangularSolveKernel<T, Context>(
          dev_ctx, l_mH, mul_tmp, true, false, true, &gL);

      auto phil_rank = gL.dims().size();
      auto phir_rank = gR.dims().size();
      funcs::ForRange<Context> l_for_range(dev_ctx, gL.numel());
      funcs::TrilTriuCompute<T> tril_computer(gL.data<T>(),
                                              -1,
                                              true,
                                              gL.dims()[phil_rank - 2],
                                              gL.dims()[phil_rank - 1],
                                              gL.data<T>());
      l_for_range(tril_computer);

      funcs::ForRange<Context> r_for_range(dev_ctx, gR.numel());
      funcs::TrilTriuCompute<T> triu_computer(gR.data<T>(),
                                              0,
                                              false,
                                              gR.dims()[phir_rank - 2],
                                              gR.dims()[phir_rank - 1],
                                              gR.data<T>());
      r_for_range(triu_computer);
      Tensor_Add<Context, T>(dev_ctx, gL, gR, lu_grad);
    } else {
      DenseTensor gR, p_mT, tem_out, out_grad_mH, tem_out1, tem_out2, tem_out3,
          gU;
      p_mT = Transpose2DTo6D<Context, T>(dev_ctx, p);
      auto PmTdims = p_mT.dims();
      auto Outdims = out.dims();
      auto mat_dim_p = funcs::CreateMatrixDescriptor(PmTdims, 0, false);
      auto mat_dim_o = funcs::CreateMatrixDescriptor(Outdims, 0, false);
      tem_out.Resize(Outdims);
      dev_ctx.template Alloc<T>(&tem_out);
      auto blas = funcs::GetBlas<Context, T>(dev_ctx);
      // gR = -P^T op_3(X)op_1(op_2(gX))P
      blas.MatMul(p_mT,
                  mat_dim_p,
                  out,
                  mat_dim_o,
                  static_cast<T>(-1),
                  &tem_out,
                  static_cast<T>(0));
      out_grad_mH = GetMH<T, Context>(dev_ctx, out_grad);
      auto TemOutdims = tem_out.dims();
      auto OutGradmHdims = out_grad_mH.dims();
      auto mat_dim_tem_out =
          funcs::CreateMatrixDescriptor(TemOutdims, 0, false);
      auto mat_dim_out_grad_mH =
          funcs::CreateMatrixDescriptor(OutGradmHdims, 0, false);
      tem_out1.Resize(lu.dims());
      dev_ctx.template Alloc<T>(&tem_out1);
      blas.MatMul(tem_out,
                  mat_dim_tem_out,
                  out_grad_mH,
                  mat_dim_out_grad_mH,
                  static_cast<T>(1),
                  &tem_out1,
                  static_cast<T>(0));
      auto TemOutdims1 = tem_out1.dims();
      auto pdims = p.dims();
      auto mat_dim_tem_out1 =
          funcs::CreateMatrixDescriptor(TemOutdims1, 0, false);
      auto mat_dim_p1 = funcs::CreateMatrixDescriptor(pdims, 0, false);
      tem_out2.Resize(TemOutdims1);
      dev_ctx.template Alloc<T>(&tem_out2);
      blas.MatMul(tem_out1,
                  mat_dim_tem_out1,
                  p,
                  mat_dim_p1,
                  static_cast<T>(1),
                  &tem_out2,
                  static_cast<T>(0));
      // gR = gR L^{-H}
      TriangularSolveKernel<T, Context>(
          dev_ctx, l_mH, tem_out2, true, true, true, &gR);
      // gU = (L^H gR U^{-H}).triu()
      auto LmHdims = l_mH.dims();
      auto gRdims = gR.dims();
      auto mat_dim_l_mh = funcs::CreateMatrixDescriptor(LmHdims, 0, false);
      auto mat_dim_gr = funcs::CreateMatrixDescriptor(gRdims, 0, false);
      tem_out3.Resize(LmHdims);
      dev_ctx.template Alloc<T>(&tem_out3);
      blas.MatMul(l_mH,
                  mat_dim_l_mh,
                  gR,
                  mat_dim_gr,
                  static_cast<T>(1),
                  &tem_out3,
                  static_cast<T>(0));
      TriangularSolveKernel<T, Context>(
          dev_ctx, u_mH, tem_out3, false, true, false, &gU);

      auto phiu_rank = gU.dims().size();
      auto phir_rank = gR.dims().size();
      funcs::ForRange<Context> l_for_range(dev_ctx, gR.numel());
      funcs::TrilTriuCompute<T> tril_computer(gR.data<T>(),
                                              -1,
                                              true,
                                              gR.dims()[phir_rank - 2],
                                              gR.dims()[phir_rank - 1],
                                              gR.data<T>());
      l_for_range(tril_computer);

      funcs::ForRange<Context> r_for_range(dev_ctx, gU.numel());
      funcs::TrilTriuCompute<T> triu_computer(gU.data<T>(),
                                              0,
                                              false,
                                              gU.dims()[phiu_rank - 2],
                                              gU.dims()[phiu_rank - 1],
                                              gU.data<T>());
      r_for_range(triu_computer);
      Tensor_Add<Context, T>(dev_ctx, gR, gU, lu_grad);
    }
  }
}

}  // namespace phi