/* Copyright (c) 2021 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 "Eigen/Core"
#include "Eigen/LU"
#include "paddle/phi/backends/all_context.h"
#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/core/enforce.h"

namespace phi {
namespace funcs {

// for TransposeNormal
static std::vector<int> getNewAxis(const int b_rank) {
  std::vector<int> axis_1 = {0};
  std::vector<int> axis_2 = {1, 0};
  std::vector<int> axis_3 = {0, 2, 1};
  std::vector<int> axis_4 = {0, 1, 3, 2};
  std::vector<int> axis_5 = {0, 1, 2, 4, 3};
  std::vector<int> axis_6 = {0, 1, 2, 3, 5, 4};
  std::vector<int> axis_7 = {0, 1, 2, 3, 4, 6, 5};
  std::vector<int> axis_8 = {0, 1, 2, 3, 4, 5, 7, 6};
  std::vector<int> axis_9 = {0, 1, 2, 3, 4, 5, 6, 8, 7};
  switch (b_rank) {
    case 1:
      return axis_1;
      break;
    case 2:
      return axis_2;
      break;
    case 3:
      return axis_3;
      break;
    case 4:
      return axis_4;
      break;
    case 5:
      return axis_5;
      break;
    case 6:
      return axis_6;
      break;
    case 7:
      return axis_7;
      break;
    case 8:
      return axis_8;
      break;
    default:
      return axis_9;
  }
}

// for Resize
static std::vector<int64_t> getNewDimsVec(const DDim& b_dims) {
  std::vector<int64_t> b_dims_vec = vectorize(b_dims);
  int size = b_dims_vec.size();
  if (size >= 2) {
    // swap the last 2 elements in b_dims_vec
    int64_t temp = b_dims_vec[size - 1];
    b_dims_vec[size - 1] = b_dims_vec[size - 2];
    b_dims_vec[size - 2] = temp;
    return b_dims_vec;
  }
  PADDLE_ENFORCE_NE(
      b_dims_vec.empty(),
      true,
      common::errors::PreconditionNotMet(
          "The size of tensor b must not be %d after getting new dims", 0));
  // if b_dims_vec.size() == 1, just return original vec
  return b_dims_vec;
}

template <typename Context, typename T>
void compute_solve_eigen(const Context& dev_ctx,
                         const DenseTensor& a,
                         const DenseTensor& b,
                         DenseTensor* out) {
  using Matrix =
      Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;
  using EigenMatrixMap = Eigen::Map<Matrix>;
  using ConstEigenMatrixMap = Eigen::Map<const Matrix>;
  // prepare for a
  const auto& a_mat_dims = a.dims();
  const int a_rank = a_mat_dims.size();
  int n = a_mat_dims[a_rank - 1];
  int64_t a_batch_size =
      a_rank > 2 ? a.numel() / (static_cast<int64_t>(n) * n) : 1;

  // prepare for b
  const auto& b_mat_dims = b.dims();
  const int b_rank = b_mat_dims.size();
  int b_h = n;
  int b_w = b_mat_dims[b_rank - 1];
  int64_t b_batch_size =
      b_rank > 2 ? b.numel() / (static_cast<int64_t>(b_h) * b_w) : 1;

  const T* a_ptr = a.data<T>();
  const T* b_ptr = b.data<T>();
  out->Resize(b_mat_dims);  // make sure the out dims is right

  T* out_ptr = dev_ctx.template Alloc<T>(out);
  if (a_batch_size == b_batch_size) {
    for (int64_t i = 0; i < a_batch_size; ++i) {
      ConstEigenMatrixMap a_mat(a_ptr + static_cast<int64_t>(i) * n * n, n, n);
      ConstEigenMatrixMap b_mat(
          b_ptr + static_cast<int64_t>(i) * b_h * b_w, b_h, b_w);
      EigenMatrixMap out_mat(
          out_ptr + static_cast<int64_t>(i) * b_h * b_w, b_h, b_w);
      Eigen::PartialPivLU<Matrix> lu;
      lu.compute(a_mat);
      const T min_abs_pivot = lu.matrixLU().diagonal().cwiseAbs().minCoeff();
      PADDLE_ENFORCE_GT(
          min_abs_pivot,
          static_cast<T>(0),
          common::errors::InvalidArgument("Input is not invertible."));
      out_mat.noalias() = lu.solve(b_mat);
    }
  } else {
    PADDLE_ENFORCE_EQ(a_batch_size,
                      b_batch_size,
                      common::errors::InvalidArgument(
                          "All input tensors must have the same rank."));
  }
}

// only used for complex input
template <typename T>
void SolveLinearSystem(T* matrix_data,
                       T* rhs_data,
                       T* out_data,
                       int order,
                       int rhs_cols,
                       int batch) {
  using Treal = typename Eigen::NumTraits<T>::Real;

  // cast paddle::complex into std::complex
  std::complex<Treal>* matrix_data_ =
      reinterpret_cast<std::complex<Treal>*>(matrix_data);
  std::complex<Treal>* rhs_data_ =
      reinterpret_cast<std::complex<Treal>*>(rhs_data);
  std::complex<Treal>* out_data_ =
      reinterpret_cast<std::complex<Treal>*>(out_data);

  using Matrix = Eigen::Matrix<std::complex<Treal>,
                               Eigen::Dynamic,
                               Eigen::Dynamic,
                               Eigen::RowMajor>;
  using InputMatrixMap = Eigen::Map<Matrix>;
  using OutputMatrixMap = Eigen::Map<Matrix>;

  for (int i = 0; i < batch; ++i) {
    auto input_matrix =
        InputMatrixMap(matrix_data_ + i * order * order, order, order);
    auto input_rhs =
        InputMatrixMap(rhs_data_ + i * order * rhs_cols, order, rhs_cols);
    auto output =
        OutputMatrixMap(out_data_ + i * order * rhs_cols, order, rhs_cols);

    Eigen::PartialPivLU<Matrix> lu_decomposition(order);
    lu_decomposition.compute(input_matrix);

    const Treal min_abs_piv =
        lu_decomposition.matrixLU().diagonal().cwiseAbs().minCoeff();
    PADDLE_ENFORCE_GT(min_abs_piv,
                      Treal(0),
                      common::errors::InvalidArgument(
                          "Something's wrong with SolveLinearSystem. "));

    output = lu_decomposition.solve(input_rhs);
  }
}

template <typename Context, typename T>
class MatrixSolveFunctor {
 public:
  void operator()(const Context& dev_ctx,
                  const DenseTensor& a,
                  const DenseTensor& b,
                  DenseTensor* out);
};

}  // namespace funcs
}  // namespace phi