// Copyright (c) 2022 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/common/enforce.h"
#include "paddle/phi/common/memory_utils.h"
#include "paddle/phi/core/enforce.h"
#include "paddle/utils/optional.h"

#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/kernels/activation_kernel.h"
#include "paddle/phi/kernels/elementwise_subtract_kernel.h"
#include "paddle/phi/kernels/matmul_kernel.h"
#include "paddle/phi/kernels/reduce_sum_kernel.h"

#if defined(PADDLE_WITH_CUDA)
#include "paddle/phi/backends/dynload/cusolver.h"
#endif

#if defined(PADDLE_WITH_HIP)
#include "paddle/phi/backends/dynload/rocsolver.h"
#endif

#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
#include "paddle/phi/backends/gpu/gpu_context.h"
#endif

namespace phi {

inline int GetBatchCount(const DDim& dims) {
  int64_t count = 1;
  int num_dims = dims.size();
  for (int i = 0; i < num_dims - 2; ++i) {
    count *= dims[i];
  }
  PADDLE_ENFORCE_LE_INT_MAX(count, "lstsq batch count");
  return static_cast<int>(count);
}

inline int GetMatrixStride(const DDim& dims) {
  int num_dims = dims.size();
  int64_t stride = dims[num_dims - 1] * dims[num_dims - 2];
  PADDLE_ENFORCE_LE_INT_MAX(stride, "lstsq matrix stride");
  return static_cast<int>(stride);
}

inline bool IsComplexDtype(const DataType& type) {
  return (type == DataType::COMPLEX64 || type == DataType::COMPLEX128);
}

template <typename Context, typename T>
inline void GetResidualsTensor(const Context& dev_ctx,
                               const DenseTensor& x,
                               const DenseTensor& y,
                               const std::string& driver,
                               DenseTensor* solution,
                               DenseTensor* residuals,
                               DenseTensor* rank) {
  auto x_dims = x.dims();
  int dim_size = x_dims.size();
  int64_t m = x_dims[dim_size - 2];
  int64_t n = x_dims[dim_size - 1];

  // Note(zrr1999): Although m and n are declared as int64_t, the rank tensor
  // stores int values (see rank->data<int>() usage below), so effectively these
  // dimensions are limited to int range in the current implementation.
  if (m > n && driver != "gelsy") {
    bool compute_residuals = true;
    if ((driver == "gelss" || driver == "gelsd") && rank->numel() != 0) {
      if (dim_size == 2) {
        compute_residuals = rank->data<int>()[0] == n;
      } else {
        compute_residuals = std::all_of(rank->data<int>(),
                                        rank->data<int>() + rank->numel(),
                                        [n](int r) { return r == n; });
      }
    }
    if (compute_residuals) {
      DenseTensor matmul_tensor =
          Matmul<T>(dev_ctx, x, *solution, false, false);
      DenseTensor sub_tensor = Subtract<T>(dev_ctx, matmul_tensor, y);
      DenseTensor pow_tensor;
      pow_tensor.Resize(sub_tensor.dims());
      dev_ctx.template Alloc<T>(&pow_tensor);
      PowKernel<T>(dev_ctx, sub_tensor, Scalar(2), &pow_tensor);

      auto sum_tensor = Sum<T>(
          dev_ctx, pow_tensor, IntArray({-2}), pow_tensor.dtype(), false);
      Copy<Context>(dev_ctx, sum_tensor, dev_ctx.GetPlace(), true, residuals);
      return;
    }
  }

  IntArray empty_shape({0});
  DenseTensor empty_tensor = Empty<T, Context>(dev_ctx, empty_shape);
  Copy<Context>(dev_ctx, empty_tensor, dev_ctx.GetPlace(), true, residuals);
}

#ifdef PADDLE_WITH_HIP
template <typename Context, typename T>
inline void BatchedOrmqr(const Context& dev_ctx,
                         bool left,
                         bool transpose,
                         int batch_size,
                         int m,
                         int n,
                         int k,
                         T* a,
                         int a_stride,
                         T* tau,
                         int tau_stride,
                         T* other,
                         int other_stride);

#define FUNC_WITH_TYPES(m) m(float, s) m(double, d)
#define ORMQR_BATCH_INSTANCE(T, C)                                        \
  template <>                                                             \
  inline void BatchedOrmqr<GPUContext, T>(const GPUContext& dev_ctx,      \
                                          bool left,                      \
                                          bool transpose,                 \
                                          int batch_size,                 \
                                          int m,                          \
                                          int n,                          \
                                          int k,                          \
                                          T* a,                           \
                                          int a_stride,                   \
                                          T* tau,                         \
                                          int tau_stride,                 \
                                          T* other,                       \
                                          int other_stride) {             \
    auto side = left ? rocblas_side_left : rocblas_side_right;            \
    auto trans =                                                          \
        transpose ? rocblas_operation_transpose : rocblas_operation_none; \
    int lda = std::max<int>(1, left ? m : n);                             \
    int ldc = std::max<int>(1, m);                                        \
    auto handle = dev_ctx.cusolver_dn_handle();                           \
    for (int i = 0; i < batch_size; ++i) {                                \
      T* a_working_ptr = &a[i * a_stride];                                \
      T* tau_working_ptr = &tau[i * tau_stride];                          \
      T* other_working_ptr = &other[i * other_stride];                    \
      PADDLE_ENFORCE_GPU_SUCCESS(                                         \
          dynload::rocsolver_##C##ormqr(handle,                           \
                                        side,                             \
                                        trans,                            \
                                        m,                                \
                                        n,                                \
                                        k,                                \
                                        a_working_ptr,                    \
                                        lda,                              \
                                        tau_working_ptr,                  \
                                        other_working_ptr,                \
                                        ldc));                            \
    }                                                                     \
  }
FUNC_WITH_TYPES(ORMQR_BATCH_INSTANCE);
#endif
#if defined(PADDLE_WITH_CUDA)
template <typename Context, typename T>
inline void BatchedOrmqr(const Context& dev_ctx,
                         bool left,
                         bool transpose,
                         int batch_size,
                         int m,
                         int n,
                         int k,
                         T* a,
                         int a_stride,
                         T* tau,
                         int tau_stride,
                         T* other,
                         int other_stride);

template <>
inline void BatchedOrmqr<GPUContext, float>(const GPUContext& dev_ctx,
                                            bool left,
                                            bool transpose,
                                            int batch_size,
                                            int m,
                                            int n,
                                            int k,
                                            float* a,
                                            int a_stride,
                                            float* tau,
                                            int tau_stride,
                                            float* other,
                                            int other_stride) {
  int lwork = 0;
  auto side = left ? CUBLAS_SIDE_LEFT : CUBLAS_SIDE_RIGHT;
  auto trans = transpose ? CUBLAS_OP_T : CUBLAS_OP_N;
  int lda = std::max<int>(1, left ? m : n);
  int ldc = std::max<int>(1, m);

  auto handle = dev_ctx.cusolver_dn_handle();
  PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnSormqr_bufferSize(
      handle, side, trans, m, n, k, a, lda, tau, other, ldc, &lwork));
  DenseTensor info;
  info.Resize({1});
  int* info_d = dev_ctx.template Alloc<int>(&info);

  for (int i = 0; i < batch_size; ++i) {
    float* a_working_ptr = &a[i * a_stride];
    float* tau_working_ptr = &tau[i * tau_stride];
    float* other_working_ptr = &other[i * other_stride];

    handle = dev_ctx.cusolver_dn_handle();
    DenseTensor workspace;
    workspace.Resize({lwork});
    float* workspace_ptr = dev_ctx.template Alloc<float>(&workspace);

    // compute ormgr
    PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnSormqr(handle,
                                                         side,
                                                         trans,
                                                         m,
                                                         n,
                                                         k,
                                                         a_working_ptr,
                                                         lda,
                                                         tau_working_ptr,
                                                         other_working_ptr,
                                                         ldc,
                                                         workspace_ptr,
                                                         lwork,
                                                         info_d));

    // check the error info
    int info_h;
    memory_utils::Copy(CPUPlace(),
                       &info_h,
                       dev_ctx.GetPlace(),
                       info_d,
                       sizeof(int),
                       dev_ctx.stream());
    PADDLE_ENFORCE_EQ(
        info_h,
        0,
        common::errors::PreconditionNotMet(
            "For batch [%d]: CUSolver info is not zero but [%d]", i, info_h));
  }
}

template <>
inline void BatchedOrmqr<GPUContext, double>(const GPUContext& dev_ctx,
                                             bool left,
                                             bool transpose,
                                             int batch_size,
                                             int m,
                                             int n,
                                             int k,
                                             double* a,
                                             int a_stride,
                                             double* tau,
                                             int tau_stride,
                                             double* other,
                                             int other_stride) {
  int lwork = 0;
  auto side = left ? CUBLAS_SIDE_LEFT : CUBLAS_SIDE_RIGHT;
  auto trans = transpose ? CUBLAS_OP_T : CUBLAS_OP_N;
  int lda = std::max<int>(1, left ? m : n);
  int ldc = std::max<int>(1, m);

  auto handle = dev_ctx.cusolver_dn_handle();
  PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnDormqr_bufferSize(
      handle, side, trans, m, n, k, a, lda, tau, other, ldc, &lwork));
  DenseTensor info;
  info.Resize({1});
  int* info_d = dev_ctx.template Alloc<int>(&info);

  for (int i = 0; i < batch_size; ++i) {
    double* a_working_ptr = &a[i * a_stride];
    double* tau_working_ptr = &tau[i * tau_stride];
    double* other_working_ptr = &other[i * other_stride];

    handle = dev_ctx.cusolver_dn_handle();
    DenseTensor workspace;
    workspace.Resize({lwork});
    double* workspace_ptr = dev_ctx.template Alloc<double>(&workspace);

    // compute ormgr
    PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnDormqr(handle,
                                                         side,
                                                         trans,
                                                         m,
                                                         n,
                                                         k,
                                                         a_working_ptr,
                                                         lda,
                                                         tau_working_ptr,
                                                         other_working_ptr,
                                                         ldc,
                                                         workspace_ptr,
                                                         lwork,
                                                         info_d));

    // check the error info
    int info_h;
    memory_utils::Copy(CPUPlace(),
                       &info_h,
                       dev_ctx.GetPlace(),
                       info_d,
                       sizeof(int),
                       dev_ctx.stream());
    PADDLE_ENFORCE_EQ(
        info_h,
        0,
        common::errors::PreconditionNotMet(
            "For batch [%d]: CUSolver info is not zero but [%d]", i, info_h));
  }
}
#endif

}  // namespace phi