// @HEADER
// *****************************************************************************
//        MueLu: A package for multigrid based preconditioning
//
// Copyright 2012 NTESS and the MueLu contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER

#ifndef MUELU_ZOLTANINTERFACE_DECL_HPP
#define MUELU_ZOLTANINTERFACE_DECL_HPP

#include "MueLu_ConfigDefs.hpp"

#if defined(HAVE_MUELU_ZOLTAN) && defined(HAVE_MPI)

#include <zoltan_cpp.h>

#include <Xpetra_Matrix_fwd.hpp>
#include <Xpetra_VectorFactory_fwd.hpp>

#include "MueLu_SingleLevelFactoryBase.hpp"
#include "MueLu_ZoltanInterface_fwd.hpp"

#include "MueLu_Level_fwd.hpp"
#include "MueLu_FactoryBase_fwd.hpp"

namespace MueLu {

/*!
  @class ZoltanInterface
  @brief Interface to Zoltan library.
  @ingroup Rebalancing

  This interface provides access to partitioning methods in Zoltan.
  Currently, it supports the RCB algorithm only.

  ## Input/output of ZoltanInterface ##

  ### User parameters of ZoltanInterface ###
  Parameter | type | default | master.xml | validated | requested | description
  ----------|------|---------|:----------:|:---------:|:---------:|------------
  | A                                      | Factory | null  |   | * | * | Generating factory of the matrix A used during the prolongator smoothing process |
  | Coordinates                            | Factory | null  |   | * | * | Factory generating coordinates vector used for rebalancing (RCB algorithm)
  | number of partitions                   | GO      | - |  |  |  | Short-cut parameter set by RepartitionFactory. Avoid repartitioning algorithms if only one partition is necessary (see details below)

  The * in the @c master.xml column denotes that the parameter is defined in the @c master.xml file.<br>
  The * in the @c validated column means that the parameter is declared in the list of valid input parameters (see ZoltanInterface::GetValidParameters).<br>
  The * in the @c requested column states that the data is requested as input with all dependencies (see ZoltanInterface::DeclareInput).

  ### Variables provided by ZoltanInterface ###

  After ZoltanInterface::Build the following data is available (if requested)

  Parameter | generated by | description
  ----------|--------------|------------
  | Partition       | ZoltanInterface   | GOVector based on the Row map of A (DOF-based) containing the process id the DOF should be living in after rebalancing/repartitioning

  The "Partition" vector is used as input for the RepartitionFactory class.
  If Re-partitioning/rebalancing is necessary it uses the "Partition" variable to create the corresponding Xpetra::Import object which then is used
  by the RebalanceFactory classes (e.g., RebalanceAcFactory, RebalanceTransferFactory,...) to rebalance the coarse level operators.

  The RepartitionHeuristicFactory calculates how many partitions are to be built when performing rebalancing.
  It stores the result in the "number of partitions" variable on the current level (type = GO).
  If it is "number of partitions=1" we skip the Zoltan call and just create an dummy "Partition" vector containing zeros only.
  If no repartitioning is necessary (i.e., just keep the current partitioning) we return "Partition = Teuchos::null".
  If "number of partitions" > 1, the algorithm tries to find the requested number of partitions.

  @note: The ZoltanInterface class does not have all the features availabe in the Zoltan2Interface class
*/

// FIXME: this class should not be templated
template <class Scalar,
          class LocalOrdinal  = DefaultLocalOrdinal,
          class GlobalOrdinal = DefaultGlobalOrdinal,
          class Node          = DefaultNode>
class ZoltanInterface : public SingleLevelFactoryBase {
#undef MUELU_ZOLTANINTERFACE_SHORT
#include "MueLu_UseShortNames.hpp"

 public:
  //! @name Constructors/Destructors
  //@{

  //! Constructor
  ZoltanInterface() {}

  //! Destructor
  virtual ~ZoltanInterface() {}
  //@}

  RCP<const ParameterList> GetValidParameterList() const;

  //! @name Input
  //@{
  void DeclareInput(Level &level) const;
  //@}

  //! @name Build methods.
  //@{
  void Build(Level &level) const;

  //@}

  //! @name Query methods (really functions) required by Zoltan.
  //@{

  /*! Callback function that returns the local number of objects. Required by Zoltan.

  In this case, the number of objects is the number of local rows.

  @param data (in) void pointer to an Xpetra::Matrix.
  @param ierr (out) error code.
  */
  static int GetLocalNumberOfRows(void *data, int *ierr);

  /*! Callback function that returns the local number of nonzeros in the matrix. Required by Zoltan.

  FIXME: Note that this will not work properly for non-point matrices.

  @param data (in) void pointer to an Xpetra::Matrix
  @param weights (out) array whose <tt>i</tt><sup>th</sup> entry is the number of nonzeros in local row \c i.
  @param ierr (out) error code
  */
  static void GetLocalNumberOfNonzeros(void *data, int NumGidEntries, int NumLidEntries, ZOLTAN_ID_PTR gids,
                                       ZOLTAN_ID_PTR lids, int wgtDim, float *weights, int *ierr);

  /*! Callback function that returns the problem dimension. Required by Zoltan.

  @param data (in) void pointer to integer dimension
  @param ierr (out) error code
  */
  static int GetProblemDimension(void *data, int *ierr);

  /*! Callback function that returns the problem dimension. Required by Zoltan.

  @param data (in) void pointer to Xpetra::MultiVector.
  @param coordinates (out) array of double coordinates, arranged like so: [x1 y1 z1 x2 y2 z2 ...].
  @param ierr (out) error code

  TODO -- should I return a view of the coordinates instead of copying them?
  */
  static void GetProblemGeometry(void *data, int numGIDEntries, int numLIDEntries, int numObjectIDs,
                                 ZOLTAN_ID_PTR gids, ZOLTAN_ID_PTR lids, int dim, double *coordinates, int *ierr);

  //@}

 private:
  static ArrayRCP<typename Teuchos::ScalarTraits<Scalar>::magnitudeType> coalesceCoordinates(ArrayRCP<typename Teuchos::ScalarTraits<Scalar>::magnitudeType> coord, LocalOrdinal blksize);

};  // class ZoltanInterface

}  // namespace MueLu

#define MUELU_ZOLTANINTERFACE_SHORT
#endif  // if defined(HAVE_MUELU_ZOLTAN) && defined(HAVE_MPI)

#endif  // MUELU_ZOLTANINTERFACE_DECL_HPP
