/****************************************************************************
 * Copyright (c) 2025, ArborX authors                                       *
 * All rights reserved.                                                     *
 *                                                                          *
 * This file is part of the ArborX library. ArborX is                       *
 * distributed under a BSD 3-clause license. For the licensing terms see    *
 * the LICENSE file in the top-level directory.                             *
 *                                                                          *
 * SPDX-License-Identifier: BSD-3-Clause                                    *
 ****************************************************************************/

#ifndef ARBORX_INTERP_MOVING_LEAST_SQUARES_COEFFICIENTS_HPP
#define ARBORX_INTERP_MOVING_LEAST_SQUARES_COEFFICIENTS_HPP

#include <ArborX_GeometryTraits.hpp>
#include <ArborX_Point.hpp>
#include <detail/ArborX_AccessTraits.hpp>
#include <detail/ArborX_InterpDetailsCompactRadialBasisFunction.hpp>
#include <detail/ArborX_InterpDetailsPolynomialBasis.hpp>
#include <kokkos_ext/ArborX_KokkosExtAccessibilityTraits.hpp>
#include <misc/ArborX_SymmetricSVD.hpp>

#include <Kokkos_Core.hpp>
#include <Kokkos_Profiling_ScopedRegion.hpp>
#include <Kokkos_Swap.hpp>

namespace ArborX::Interpolation::Details
{

template <typename SourcePoints, typename TargetAccess, typename Coefficients,
          typename ExecutionSpace, typename CRBFunc = CRBF::Wendland<0>,
          typename PolynomialDegree = PolynomialDegree<2>>
class MovingLeastSquaresCoefficientsKernel
{
private:
  using ScratchMemorySpace = typename ExecutionSpace::scratch_memory_space;

  using SourcePoint = typename SourcePoints::non_const_value_type;
  using TargetPoint = typename TargetAccess::value_type;

  using CoefficientsType = typename Coefficients::non_const_value_type;

  static constexpr int dimension = GeometryTraits::dimension_v<SourcePoint>;
  static constexpr int degree = PolynomialDegree::value;
  static constexpr int poly_size = polynomialBasisSize<dimension, degree>();

  template <typename T>
  using ScratchView =
      Kokkos::View<T, ScratchMemorySpace, Kokkos::MemoryUnmanaged>;

  using LocalSourcePoints = Kokkos::Subview<SourcePoints, int, Kokkos::ALL_t>;
  using LocalPhi = ScratchView<CoefficientsType *>;
  using LocalVandermonde = ScratchView<CoefficientsType *[poly_size]>;
  using LocalMoment = ScratchView<CoefficientsType[poly_size][poly_size]>;
  using LocalSVDDiag = ScratchView<CoefficientsType[poly_size]>;
  using LocalSVDUnit = ScratchView<CoefficientsType[poly_size][poly_size]>;
  using LocalCoefficients = Kokkos::Subview<Coefficients, int, Kokkos::ALL_t>;

public:
  MovingLeastSquaresCoefficientsKernel(ExecutionSpace const &,
                                       TargetAccess const &target_access,
                                       SourcePoints const &source_points,
                                       Coefficients const &coefficients)
      : _target_access(target_access)
      , _source_points(source_points)
      , _coefficients(coefficients)
      , _num_targets(target_access.size())
      , _num_neighbors(source_points.extent_int(1))
  {}

  template <typename TeamMember>
  KOKKOS_FUNCTION void operator()(TeamMember const &member) const
  {
    auto const &scratch = member.thread_scratch(0);

    int target = member.league_rank() * member.team_size() + member.team_rank();
    if (target >= _num_targets)
      return;

    auto target_point = _target_access(target);
    auto source_points = Kokkos::subview(_source_points, target, Kokkos::ALL);
    LocalPhi phi(scratch, _num_neighbors);
    LocalVandermonde vandermonde(scratch, _num_neighbors);
    LocalMoment moment(scratch);
    LocalSVDDiag svd_diag(scratch);
    LocalSVDUnit svd_unit(scratch);
    auto coefficients = Kokkos::subview(_coefficients, target, Kokkos::ALL);

    // The goal is to compute the following line vector for each target point:
    // p(x).[P^T.PHI.P]^-1.P^T.PHI
    // Where:
    // - p(x) is the polynomial basis of point x (line vector).
    // - P is the multidimensional Vandermonde matrix built from the source
    //   points, i.e., each line is the polynomial basis of a source point.
    // - PHI is the diagonal weight matrix / CRBF evaluated at each source
    // point.

    // We first change the origin of the evaluation to be at the target point
    // within about unit distance which is equivalent to recentering and scaling
    // the monomial basis. This lets us use p(0) which is [1 0 ... 0] and leads
    // to a more stable algorithm, see, e.g., (3.8) in
    // Mirzaei, Davoud. "Analysis of moving least squares approximation
    // revisited." Journal of Computational and Applied Mathematics 282 (2015):
    // 237-250.
    sourceNormalization(target_point, source_points);

    // This computes PHI given the source points.
    phiComputation(source_points, phi);

    // This builds the Vandermonde (P) matrix
    vandermondeComputation(source_points, vandermonde);

    // We then create what is called the moment matrix, which is P^T.PHI.P. By
    // construction, it is symmetric.
    momentComputation(phi, vandermonde, moment);

    // We need the inverse of P^T.PHI.P, and because it is symmetric, we can use
    // the symmetric SVD algorithm to get it.
    ::ArborX::Details::symmetricSVDKernel(moment, svd_diag, svd_unit);

    // Check if matrix is singular
    auto [is_singular, tolerance] = isSingularMatrix(svd_diag);
    if (is_singular)
    {
      constexpr int n = svd_diag.static_extent(0);

#ifndef NDEBUG
      int n_zero_values = 0;
      for (int i = 0; i < n; ++i)
        if (Kokkos::abs(svd_diag(i)) < tolerance)
          ++n_zero_values;
      KOKKOS_ASSERT(n_zero_values > 0);
#endif
      ::ArborX::Details::symmetricMatrixFromSVD(svd_diag, svd_unit, moment);

      // Bring matrix in row-echelon form to find rows that form a full rank
      // matrix

      // svd_unit can be overwritten now and we store a copy of the moment
      // matrix in it to compute the row echelon form
      for (int i = 0; i < n; ++i)
        for (int j = 0; j < n; ++j)
          svd_unit(i, j) = moment(i, j);

      Kokkos::Array<int, n> permutation;
      for (int i = 0; i < n; ++i)
        permutation[i] = i;

      rowEchelonForm(svd_unit, permutation, tolerance);

#ifndef NDEBUG
      int n_zero_values_in_echelon_form = 0;
      for (int i = 0; i < n; ++i)
      {
        CoefficientsType max = 0;
        for (int j = i; j < n; ++j)
          if (Kokkos::abs(svd_unit(i, j)) > max)
            max = Kokkos::abs(svd_unit(i, j));
        if (max < tolerance)
          ++n_zero_values_in_echelon_form;
      }
      KOKKOS_ASSERT(n_zero_values_in_echelon_form == n_zero_values);
#endif
      for (int i = 0; i < n; ++i)
      {
        CoefficientsType max = 0;
        for (int j = i; j < n; ++j)
          if (Kokkos::abs(svd_unit(i, j)) > max)
            max = Kokkos::abs(svd_unit(i, j));
        if (max < tolerance)
        {
          int row_index = permutation[i];
          // Eliminate corresponding column from Vandermonde matrix and
          // eliminate row and column from moment matrix storing 1 one the
          // diagonal
          for (int j = 0; j < n; ++j)
            moment(row_index, j) = moment(j, row_index) = 0;
          for (unsigned int j = 0; j < vandermonde.static_extent(0); ++j)
            vandermonde(j, row_index) = 0;
          moment(row_index, row_index) = 1.;
        }
      }
      // Compute a SVD for the new matrix
      ::ArborX::Details::symmetricSVDKernel(moment, svd_diag, svd_unit);
#ifndef NDEBUG
      for (int i = 0; i < n; ++i)
        if (Kokkos::abs(svd_unit(i, i)) < tolerance)
          Kokkos::abort("Didn't eliminate all zero eigenvalues!");
#endif
    }

    // Store [P^T.PHI.P]^-1 in moment
    ::ArborX::Details::symmetricPseudoInverseSVDKernel(svd_diag, svd_unit,
                                                       moment);

    // Finally, the result is produced by computing
    // p(0).[P^T.PHI.P]^-1.P^T.PHI
    coefficientsComputation(phi, vandermonde, moment, coefficients);
  }

  auto makePolicy(ExecutionSpace const &space) const
  {
    Kokkos::TeamPolicy dummy_policy(space, 1, Kokkos::AUTO);
    dummy_policy.set_scratch_size(0, Kokkos::PerThread(perTargetMem()));
    int team_size =
        dummy_policy.team_size_recommended(*this, Kokkos::ParallelForTag{});
    if (team_size != 0)
    {
      int league_size = (_num_targets + team_size - 1) / team_size;
      return Kokkos::TeamPolicy(space, league_size, team_size)
          .set_scratch_size(0, Kokkos::PerThread(perTargetMem()));
    }
    return Kokkos::TeamPolicy(space, _num_targets, 1, 1)
        .set_scratch_size(0, Kokkos::PerTeam(perTargetMem()));
  }

private:
  std::size_t perTargetMem() const
  {
    std::size_t val = 0;
    val += LocalPhi::shmem_size(_num_neighbors);
    val += LocalVandermonde::shmem_size(_num_neighbors);
    val += LocalMoment::shmem_size();
    val += LocalSVDDiag::shmem_size();
    val += LocalSVDUnit::shmem_size();
    return val;
  }

  // Recenters the source points so that the target is at the origin within
  // about unit distance
  KOKKOS_FUNCTION void
  sourceNormalization(TargetPoint const &target_point,
                      LocalSourcePoints &source_points) const
  {
    CoefficientsType radius = Kokkos::Experimental::epsilon_v<CoefficientsType>;
    for (int neighbor = 0; neighbor < _num_neighbors; neighbor++)
    {
      CoefficientsType const norm =
          distance(source_points(neighbor), target_point);
      radius = Kokkos::max(radius, norm);
    }
    // The one at the limit would be 0 due to how CRBFs work
    radius *= CoefficientsType(1.1);

    for (int neighbor = 0; neighbor < _num_neighbors; neighbor++)
      for (int k = 0; k < dimension; k++)
        source_points(neighbor)[k] =
            (source_points(neighbor)[k] - target_point[k]) / radius;
  }

  // Computes the weight matrix
  KOKKOS_FUNCTION void phiComputation(LocalSourcePoints const &source_points,
                                      LocalPhi &phi) const
  {
    for (int neighbor = 0; neighbor < _num_neighbors; neighbor++)
      phi(neighbor) = CRBF::evaluate<CRBFunc>(source_points(neighbor));
  }

  // Computes the vandermonde matrix
  KOKKOS_FUNCTION void
  vandermondeComputation(LocalSourcePoints const &source_points,
                         LocalVandermonde &vandermonde) const
  {
    for (int neighbor = 0; neighbor < _num_neighbors; neighbor++)
    {
      auto basis = evaluatePolynomialBasis<degree>(source_points(neighbor));
      for (int k = 0; k < poly_size; k++)
        vandermonde(neighbor, k) = basis[k];
    }
  }

  // Computes the moment matrix
  KOKKOS_FUNCTION void momentComputation(LocalPhi const &phi,
                                         LocalVandermonde const &vandermonde,
                                         LocalMoment &moment) const
  {
    for (int i = 0; i < poly_size; i++)
      for (int j = 0; j < poly_size; j++)
      {
        moment(i, j) = 0;
        for (int neighbor = 0; neighbor < _num_neighbors; neighbor++)
          moment(i, j) += vandermonde(neighbor, i) * vandermonde(neighbor, j) *
                          phi(neighbor);
      }
  }

  KOKKOS_FUNCTION Kokkos::pair<bool, CoefficientsType>
  isSingularMatrix(LocalSVDDiag const &svd_diag) const
  {
    CoefficientsType max_eigen = 0;
    for (int i = 0; i < poly_size; i++)
      max_eigen = Kokkos::max(Kokkos::abs(svd_diag(i)), max_eigen);

    constexpr auto epsilon = Kokkos::Experimental::epsilon_v<CoefficientsType>;
    auto tolerance = poly_size * max_eigen * epsilon;
    for (int i = 0; i < poly_size; i++)
      if (Kokkos::abs(svd_diag(i)) <= tolerance)
        return {true, tolerance};
    return {false, tolerance};
  }

  KOKKOS_FUNCTION void
  rowEchelonForm(LocalMoment &matrix,
                 Kokkos::Array<int, poly_size> &permutation,
                 CoefficientsType tolerance) const
  {
    int pivot_row = 0;
    int pivot_col = 0;

    while (pivot_row < poly_size && pivot_col < poly_size)
    {
      // Find the k-th pivot:
      int i_max = 0;
      CoefficientsType max_entry = 0;
      for (int i = pivot_row; i < poly_size; ++i)
        if (Kokkos::abs(matrix(i, pivot_col)) > max_entry)
        {
          max_entry = Kokkos::abs(matrix(i, pivot_col));
          i_max = i;
        }

      // No pivot in this column, pass to next column
      if (Kokkos::abs(matrix(i_max, pivot_col)) < tolerance)
      {
        ++pivot_col;
        continue;
      }
      // swap rows(pivot_row, i_max)
      for (int i = 0; i < poly_size; ++i)
        Kokkos::kokkos_swap(matrix(pivot_row, i), matrix(i_max, i));
      Kokkos::kokkos_swap(permutation[i_max], permutation[pivot_row]);
      // Do for all rows below pivot:
      auto const pivot = matrix(pivot_row, pivot_col);
      for (int i = pivot_row + 1; i < poly_size; ++i)
      {
        CoefficientsType f = matrix(i, pivot_col) / pivot;
        // Fill with zeros the lower part of pivot column:
        matrix(i, pivot_col) = 0;
        // Do for all remaining elements in current row:
        for (int j = pivot_col + 1; j < poly_size; ++j)
          matrix(i, j) -= matrix(pivot_row, j) * f;
      }
      // Increase pivot row and column
      ++pivot_row;
      ++pivot_col;
    }
  }

  // Computes the coefficients
  KOKKOS_FUNCTION void coefficientsComputation(
      LocalPhi const &phi, LocalVandermonde const &vandermonde,
      LocalMoment const &moment, LocalCoefficients &coefficients) const
  {
    for (int neighbor = 0; neighbor < _num_neighbors; neighbor++)
    {
      coefficients(neighbor) = 0;
      for (int i = 0; i < poly_size; i++)
        coefficients(neighbor) +=
            moment(0, i) * vandermonde(neighbor, i) * phi(neighbor);
    }
  }

  TargetAccess _target_access;
  SourcePoints _source_points;
  Coefficients _coefficients;
  int _num_targets;
  int _num_neighbors;
};

template <typename CRBFunc, typename PolynomialDegree,
          typename CoefficientsType, typename ExecutionSpace,
          typename SourcePoints, typename TargetAccess>
auto movingLeastSquaresCoefficients(ExecutionSpace const &space,
                                    SourcePoints const &source_points,
                                    TargetAccess const &target_access)
{
  auto guard =
      Kokkos::Profiling::ScopedRegion("ArborX::MovingLeastSquaresCoefficients");

  namespace KokkosExt = ::ArborX::Details::KokkosExt;

  using MemorySpace = typename SourcePoints::memory_space;
  static_assert(
      KokkosExt::is_accessible_from<MemorySpace, ExecutionSpace>::value,
      "Memory space must be accessible from the execution space");

  // SourcePoints is a 2D view of points
  static_assert(Kokkos::is_view_v<SourcePoints> && SourcePoints::rank() == 2,
                "source points must be a 2D view of points");
  static_assert(
      KokkosExt::is_accessible_from<typename SourcePoints::memory_space,
                                    ExecutionSpace>::value,
      "source points must be accessible from the execution space");
  using SourcePoint = typename SourcePoints::non_const_value_type;
  GeometryTraits::check_valid_geometry_traits(SourcePoint{});
  static_assert(GeometryTraits::is_point<SourcePoint>::value,
                "source points elements must be points");
  static_assert(!std::is_const_v<typename SourcePoints::value_type>,
                "source points must be writable");
  constexpr int dimension = GeometryTraits::dimension_v<SourcePoint>;

  // TargetAccess is an access values of points
  static_assert(
      KokkosExt::is_accessible_from<typename TargetAccess::memory_space,
                                    ExecutionSpace>::value,
      "target access must be accessible from the execution space");
  using TargetPoint = typename TargetAccess::value_type;
  GeometryTraits::check_valid_geometry_traits(TargetPoint{});
  static_assert(GeometryTraits::is_point<TargetPoint>::value,
                "target access elements must be points");
  static_assert(dimension == GeometryTraits::dimension_v<TargetPoint>,
                "target and source points must have the same dimension");

  // The number of source groups must be correct
  KOKKOS_ASSERT(std::size_t{target_access.size()} == source_points.extent(0));

  Kokkos::View<CoefficientsType **, MemorySpace> coefficients(
      Kokkos::view_alloc(
          space, Kokkos::WithoutInitializing,
          "ArborX::MovingLeastSquaresCoefficients::coefficients"),
      source_points.extent_int(0), source_points.extent_int(1));

  MovingLeastSquaresCoefficientsKernel<SourcePoints, TargetAccess,
                                       decltype(coefficients), ExecutionSpace,
                                       CRBFunc, PolynomialDegree>
      kernel(space, target_access, source_points, coefficients);

  Kokkos::parallel_for("ArborX::MovingLeastSquaresCoefficients::kernel",
                       kernel.makePolicy(space), kernel);

  return coefficients;
}

} // namespace ArborX::Interpolation::Details

#endif
