/* 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/phi/core/tensor_utils.h"
#include "paddle/phi/kernels/expand_as_kernel.h"
#include "paddle/phi/kernels/funcs/matrix_solve.h"
#include "paddle/phi/kernels/funcs/reduce_functor.h"
#include "paddle/phi/kernels/squeeze_kernel.h"
#include "paddle/phi/kernels/unsqueeze_kernel.h"

namespace phi {

// check the input other is vector_case or not
static inline bool is_vector_rhs(const DenseTensor& input,
                                 const DenseTensor& other) {
  auto x_dim = input.dims();
  auto y_dim = other.dims();
  auto x_dim_size = x_dim.size();
  auto y_dim_size = y_dim.size();
  std::vector<int64_t> x_dims_vec = vectorize(x_dim);
  std::vector<int64_t> y_dims_vec = vectorize(y_dim);

  std::vector<int64_t>::const_iterator f = x_dims_vec.begin();
  std::vector<int64_t>::const_iterator l = x_dims_vec.end() - 1;
  std::vector<int64_t> x_dims_vec_cut(f, l);  // input.shape[:-1]

  std::vector<int64_t> expected_batched_rhs_shape(x_dims_vec_cut);
  bool vector_case =
      y_dim_size == 1 || (x_dim_size - 1 == y_dim_size &&
                          y_dims_vec == (expected_batched_rhs_shape));

  return vector_case;
}

// Prepared for the broadcast operation
static std::vector<int64_t> get_broadcast_batch_portion(
    std::vector<int64_t> x, std::vector<int64_t> y) {
  size_t size_x = x.size();
  size_t size_y = y.size();
  size_t size = std::max(size_x, size_y);
  std::vector<int64_t> batchPortion(size);
  ptrdiff_t i = (ptrdiff_t)size - 1;
  for (; i >= 0; --i) {
    ptrdiff_t offset = size - i - 1;
    ptrdiff_t dim_x = size_x - offset - 1;
    ptrdiff_t dim_y = size_y - offset - 1;
    int64_t x_size = (dim_x >= 0) ? x[dim_x] : 1;
    int64_t y_size = (dim_y >= 0) ? y[dim_y] : 1;
    PADDLE_ENFORCE_EQ(
        (x_size == y_size || x_size == 1 || y_size == 1),
        true,
        common::errors::PreconditionNotMet(
            "The size of tensor x (%d) must match the size of tensor y "
            "(%d) at non-singleton dimension %d.",
            x_size,
            y_size,
            i));

    batchPortion[i] = x_size != 1 ? x_size : y_size;
  }
  return batchPortion;
}

// broadcast the batch dimensions of tensor x and tensor y.
static inline std::tuple<std::vector<int64_t>, std::vector<int64_t>>
get_broadcast_dims(const DenseTensor& x, const DenseTensor& y) {
  std::vector<int64_t> x_dims_vec = vectorize(x.dims());
  std::vector<int64_t> y_dims_vec = vectorize(y.dims());
  std::vector<int64_t>::const_iterator f1 = x_dims_vec.begin();
  std::vector<int64_t>::const_iterator l1 = x_dims_vec.end() - 2;
  std::vector<int64_t> x_dims_vec_cut(f1, l1);

  std::vector<int64_t>::const_iterator f2 = y_dims_vec.begin();
  std::vector<int64_t>::const_iterator l2 = y_dims_vec.end() - 2;
  std::vector<int64_t> y_dims_vec_cut(f2, l2);

  std::vector<int64_t> expand_batch_portion =
      get_broadcast_batch_portion(x_dims_vec_cut, y_dims_vec_cut);
  std::vector<int64_t> x_expand_size({expand_batch_portion});
  x_expand_size.insert(x_expand_size.end(),
                       {x_dims_vec[static_cast<int>(x_dims_vec.size()) - 2],
                        x_dims_vec[static_cast<int>(x_dims_vec.size()) - 1]});
  std::vector<int64_t> y_expand_size({expand_batch_portion});
  y_expand_size.insert(y_expand_size.end(),
                       {y_dims_vec[static_cast<int>(y_dims_vec.size()) - 2],
                        y_dims_vec[static_cast<int>(y_dims_vec.size()) - 1]});

  return std::make_tuple(x_expand_size, y_expand_size);
}

template <typename Context, typename T>
static void linalg_solve(const Context& dev_ctx,
                         const DenseTensor& x,
                         const DenseTensor& y,
                         DenseTensor* out) {
  dev_ctx.template Alloc<T>(out);
  funcs::MatrixSolveFunctor<Context, T> mat_solve;

  // input y can be vector or matrix
  // but need to be unsqueezed if y is a vector
  bool is_vector = false;
  is_vector = is_vector_rhs(x, y);

  DenseTensor tmp_y;
  if (is_vector) {
    dev_ctx.Alloc(&tmp_y, y.dtype());

    Unsqueeze<T, Context>(dev_ctx, y, {-1}, &tmp_y, nullptr);
  } else {
    tmp_y.Resize(y.dims());
    dev_ctx.Alloc(&tmp_y, y.dtype());

    Copy(dev_ctx, y, dev_ctx.GetPlace(), false, &tmp_y);
  }

  DenseTensor tmp_x;
  tmp_x.Resize(x.dims());
  dev_ctx.Alloc(&tmp_x, x.dtype());
  Copy(dev_ctx, x, dev_ctx.GetPlace(), false, &tmp_x);

  std::vector<int64_t> x_broadcast_dims;
  std::vector<int64_t> y_broadcast_dims;
  std::tie(x_broadcast_dims, y_broadcast_dims) =
      get_broadcast_dims(tmp_x, tmp_y);

  DenseTensor tmp_x_bc;

  ExpandAsKernel<T, Context>(
      dev_ctx, tmp_x, nullptr, x_broadcast_dims, &tmp_x_bc);

  DenseTensor tmp_y_bc;
  ExpandAsKernel<T, Context>(
      dev_ctx, tmp_y, nullptr, y_broadcast_dims, &tmp_y_bc);

  auto x_dim = x.dims();
  auto y_dim = y.dims();
  auto x_dim_size = x_dim.size();
  auto y_dim_size = y_dim.size();

  if (is_vector) {                 // vector case
    out->Resize(tmp_y_bc.dims());  // out.unsqueeze(-1)
    mat_solve(dev_ctx, tmp_x_bc, tmp_y_bc, out);

    DenseTensor out_tmp;
    out_tmp.Resize(out->dims());
    out_tmp = *out;

    Squeeze<T, Context>(dev_ctx, out_tmp, {-1}, out);
  } else {
    PADDLE_ENFORCE_EQ(
        x_dim[x_dim_size - 1],
        y_dim[y_dim_size - 2],
        common::errors::InvalidArgument(
            "Matrix X1 with dimension greater than 2 and any matrix Y1,"
            "the matrix X1's width must be equal with matrix Y1's "
            "height. But received X's shape = [%s], X1's shape = [%s], X1's "
            "width = %s; Y's shape = [%s], Y1's shape = [%s], Y1's height = "
            "%s.",
            x_dim,
            x_dim,
            x_dim[x_dim_size - 1],
            y_dim,
            y_dim,
            y_dim[y_dim_size - 2]));
    mat_solve(dev_ctx, tmp_x_bc, tmp_y_bc, out);
  }
}

template <typename T, typename Context>
void SolveKernel(const Context& dev_ctx,
                 const DenseTensor& x,
                 const DenseTensor& y,
                 DenseTensor* out) {
  if (x.numel() == 0 || y.numel() == 0) {
    auto x_dims = x.dims();
    auto y_dims = y.dims();
    std::vector<int64_t> out_dims;
    if (y_dims.size() == 1) {
      out_dims =
          std::vector<int64_t>(x_dims.Get(), x_dims.Get() + x_dims.size() - 2);
      out_dims.push_back(y_dims[y_dims.size() - 1]);
    } else {
      // broadcast
      std::vector<int> x_shape(x_dims.Get(), x_dims.Get() + x_dims.size() - 2);
      std::vector<int> y_shape(y_dims.Get(), y_dims.Get() + y_dims.size() - 2);
      auto x_it = x_shape.rbegin();
      auto y_it = y_shape.rbegin();
      while (x_it != x_shape.rend() || y_it != y_shape.rend()) {
        int x_dim = (x_it != x_shape.rend()) ? *x_it : 1;
        int y_dim = (y_it != y_shape.rend()) ? *y_it : 1;
        if (x_dim == 0 || y_dim == 0) {
          out_dims.push_back(0);
        } else {
          out_dims.push_back(std::max(x_dim, y_dim));
        }
        if (x_it != x_shape.rend()) ++x_it;
        if (y_it != y_shape.rend()) ++y_it;
      }
      std::reverse(out_dims.begin(), out_dims.end());
      out_dims.insert(out_dims.end(),
                      y_dims.Get() + y_dims.size() - 2,
                      y_dims.Get() + y_dims.size());
    }
    out->Resize(out_dims);
    dev_ctx.template Alloc<T>(out);
    return;
  }
  linalg_solve<Context, T>(dev_ctx, x, y, out);
}

}  // namespace phi