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

#include <math.h>
#include <algorithm>
#include <complex>

#include "paddle/common/enforce.h"
#include "paddle/phi/backends/gpu/cuda/cudnn_workspace_helper.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/full_kernel.h"
#include "paddle/phi/kernels/funcs/slice.h"
#include "paddle/phi/kernels/impl/lstsq_kernel_impl.h"
#include "paddle/phi/kernels/impl/qr_kernel_impl.h"
#include "paddle/phi/kernels/impl/tril_triu_kernel_impl.h"
#include "paddle/phi/kernels/lstsq_kernel.h"
#include "paddle/phi/kernels/matmul_kernel.h"
#include "paddle/phi/kernels/transpose_kernel.h"
#include "paddle/phi/kernels/triangular_solve_kernel.h"
namespace phi {

enum class LapackDriverType : int { Gels, Gelsd, Gelsy, Gelss };

template <typename T, typename Context>
void LstsqKernel(const Context& dev_ctx,
                 const DenseTensor& x,
                 const DenseTensor& y,
                 const Scalar& rcond_scalar,
                 const std::string& driver_string,
                 DenseTensor* solution,
                 DenseTensor* residuals,
                 DenseTensor* rank,
                 DenseTensor* singular_values) {
  if (x.numel() == 0 || y.numel() == 0) {
    if (solution) Full<T, Context>(dev_ctx, solution->dims(), 0, solution);
    if (rank) Full<int64_t, Context>(dev_ctx, rank->dims(), 0, rank);
    if (residuals)
      GetResidualsTensor<Context, T>(
          dev_ctx, x, y, driver_string, solution, residuals, rank);
    if (singular_values)
      Full<T, Context>(dev_ctx, singular_values->dims(), 0, singular_values);
    return;
  }

  CUDNN_ENFORCE_TENSOR_SIZE_SUPPORTED(x);
  CUDNN_ENFORCE_TENSOR_SIZE_SUPPORTED(y);

  auto x_dims = x.dims();
  auto y_dims = y.dims();
  int dim_size = x_dims.size();
  const int64_t m_64 = x_dims[dim_size - 2];
  const int64_t n_64 = x_dims[dim_size - 1];
  const int64_t nrhs_64 = y_dims[dim_size - 1];
  PADDLE_ENFORCE_LE_INT_MAX(m_64, "lstsq m");
  PADDLE_ENFORCE_LE_INT_MAX(n_64, "lstsq n");
  PADDLE_ENFORCE_LE_INT_MAX(nrhs_64, "lstsq nrhs");
  const int m = static_cast<int>(m_64);
  const int n = static_cast<int>(n_64);
  const int nrhs = static_cast<int>(nrhs_64);
  int min_mn = std::min(m, n);
  int max_mn = std::max(m, n);
  int k = min_mn;

  int x_stride = GetMatrixStride(x_dims);
  int y_stride = GetMatrixStride(y_dims);
  int tau_stride = min_mn;
  int batch_count = GetBatchCount(x_dims);

  T rcond = rcond_scalar.to<T>();

  DenseTensor new_x;
  new_x.Resize({batch_count, m, n});
  dev_ctx.template Alloc<T>(&new_x);
  Copy<Context>(dev_ctx, x, dev_ctx.GetPlace(), true, &new_x);

  DenseTensor new_y;
  new_y.Resize({batch_count, m, nrhs});
  dev_ctx.template Alloc<T>(&new_y);
  Copy<Context>(dev_ctx, y, dev_ctx.GetPlace(), true, &new_y);

  // Prepare tau
  auto tau_dims_vec = vectorize<int>(x_dims);
  tau_dims_vec.pop_back();
  tau_dims_vec[tau_dims_vec.size() - 1] = min_mn;

  DenseTensor tau;
  tau.Resize(tau_dims_vec);
  auto tau_data = dev_ctx.template Alloc<T>(&tau);

  if (m >= n) {
    DenseTensor tmp_x = TransposeLast2Dim<T>(dev_ctx, new_x);
    DenseTensor tmp_y = TransposeLast2Dim<T>(dev_ctx, new_y);
    auto x_data = tmp_x.data<T>();
    auto y_data = tmp_y.data<T>();

    // step 1, compute QR factorization using geqrf
    BatchedGeqrf<Context, T>(
        dev_ctx, batch_count, m, n, x_data, m, tau_data, x_stride, tau_stride);

    // Step 2, Y <- Q^H Y
    BatchedOrmqr<Context, T>(dev_ctx,
                             true,
                             true,
                             batch_count,
                             m,
                             nrhs,
                             k,
                             x_data,
                             x_stride,
                             tau_data,
                             tau_stride,
                             y_data,
                             y_stride);

    DenseTensor trans_r = TransposeLast2Dim<T>(dev_ctx, tmp_x);
    DenseTensor slice_r =
        funcs::Slice<T>(dev_ctx, trans_r, {-2}, {0}, {min_mn});
    DenseTensor res_r;
    res_r.Resize({batch_count, min_mn, min_mn});
    dev_ctx.template Alloc<T>(&res_r);
    TrilTriuKernel<T>(dev_ctx, slice_r, 0, false, &res_r);

    DenseTensor trans_y = TransposeLast2Dim<T>(dev_ctx, tmp_y);
    DenseTensor slice_y =
        funcs::Slice<T>(dev_ctx, trans_y, {-2}, {0}, {min_mn});

    // Step 3, solve R X = Y
    TriangularSolveKernel<T, Context>(
        dev_ctx, res_r, slice_y, true, false, false, solution);

  } else {
    auto x_data = dev_ctx.template Alloc<T>(&new_x);
    auto y_data = dev_ctx.template Alloc<T>(&new_y);

    // step 1, compute QR factorization using geqrf
    BatchedGeqrf<Context, T>(
        dev_ctx, batch_count, n, m, x_data, n, tau_data, x_stride, tau_stride);

    // Step 2, solve R^H Z = Y
    DenseTensor trans_r = TransposeLast2Dim<T>(dev_ctx, new_x);
    DenseTensor slice_r =
        funcs::Slice<T>(dev_ctx, trans_r, {-2}, {0}, {min_mn});
    DenseTensor res_r;
    res_r.Resize({batch_count, min_mn, min_mn});
    dev_ctx.template Alloc<T>(&res_r);
    TrilTriuKernel<T>(dev_ctx, slice_r, 0, false, &res_r);

    TriangularSolveKernel<T, Context>(
        dev_ctx, res_r, new_y, true, true, false, solution);

    // Step 3, X <- Q Z
    BatchedOrgqr<Context, T>(dev_ctx,
                             batch_count,
                             n,
                             m,
                             min_mn,
                             x_data,
                             n,
                             tau_data,
                             x_stride,
                             tau_stride);

    DenseTensor trans_q = TransposeLast2Dim<T>(dev_ctx, new_x);
    DenseTensor slice_q = funcs::Slice<T>(dev_ctx, trans_q, {-1}, {0}, {m});
    DenseTensor solu_tensor =
        Matmul<T>(dev_ctx, slice_q, *solution, false, false);
    Copy<Context>(dev_ctx, solu_tensor, dev_ctx.GetPlace(), true, solution);
  }
  if (batch_count == 1) solution->Resize({n, nrhs});
  GetResidualsTensor<Context, T>(
      dev_ctx, x, y, driver_string, solution, residuals, rank);
}

}  // namespace phi

PD_REGISTER_KERNEL(lstsq, GPU, ALL_LAYOUT, phi::LstsqKernel, float, double) {}