/******************************************************************************
 * Copyright (c) 2024 Lawrence Livermore National Security, LLC and other
 * HYPRE Project Developers. See the top-level COPYRIGHT file for details.
 *
 * SPDX-License-Identifier: MIT
 ******************************************************************************/

#ifndef MGR_HEADER
#define MGR_HEADER

#include "internal/amg.h"
#include "internal/compatibility.h"
#include "internal/containers.h"
#include "internal/fsai.h"
#include "internal/ilu.h"
#include "internal/precon_reuse.h"
#include "internal/utils.h"
#if HYPRE_CHECK_MIN_VERSION(30100, 55)
#include "internal/schwarz.h"
#endif

struct NestedKrylov_args_struct;
struct Stats_struct;
struct YAMLnode_struct;
typedef struct MGR_args_struct MGR_args;

enum
{
   MAX_MGR_LEVELS = 32,

   /* hypre's level-0 custom F-solver callback path is keyed through AMG's
    * F-relaxation type. We use it for wrapped Schwarz F-relaxation because it
    * keeps hypre's extracted A_FF matrix alive through setup and solve. */
   MGR_FRLX_TYPE_CUSTOM_SOLVER_CB = 2,
   MGR_FRLX_TYPE_NESTED_MGR       = 202,

   /* Disable hypre's built-in global smoother slot when a user smoother is
    * installed via HYPRE_MGRSetGlobalSmootherAtLevel. */
   MGR_GRLX_TYPE_USER_SMOOTHER = -1,
   MGR_SOLVER_TYPE_SCHWARZ     = 34,
};

typedef struct MGRComponentReuse_args_struct
{
   int              present;
   int              warned_runtime_unsupported;
   int              warned_policy_unsupported;
   int              warned_type_unsupported;
   PreconReuse_args args;
} MGRComponentReuse_args;

/*--------------------------------------------------------------------------
 * Coarsest level solver arguments struct
 *--------------------------------------------------------------------------*/

typedef struct MGRcls_args_struct
{
   HYPRE_Int              type;
   MGRComponentReuse_args reuse;

   int                              use_krylov;
   struct NestedKrylov_args_struct *krylov;

   /* Only one coarsest solver is active at a time; store its args in a union.
    * Note: anonymous union is a GNU C extension. */
   union
   {
      AMG_args  amg;
      ILU_args  ilu;
      FSAI_args fsai;
#if HYPRE_CHECK_MIN_VERSION(30100, 55)
      Schwarz_args schwarz;
#endif
   };
} MGRcls_args;

/*--------------------------------------------------------------------------
 * F-Relaxation arguments struct
 *--------------------------------------------------------------------------*/

typedef struct MGRfrlx_args_struct
{
   HYPRE_Int              type;
   HYPRE_Int              num_sweeps;
   MGRComponentReuse_args reuse;

   int                              use_krylov;
   struct NestedKrylov_args_struct *krylov;
   MGR_args                        *mgr;

   /* Only one fine-relaxation solver is active at a time. */
   union
   {
      AMG_args  amg;
      ILU_args  ilu;
      FSAI_args fsai;
#if HYPRE_CHECK_MIN_VERSION(30100, 55)
      Schwarz_args schwarz;
#endif
   };
} MGRfrlx_args;

/*--------------------------------------------------------------------------
 * Global-Relaxation arguments struct
 *--------------------------------------------------------------------------*/

typedef struct MGRgrlx_args_struct
{
   HYPRE_Int              type;
   HYPRE_Int              num_sweeps;
   MGRComponentReuse_args reuse;

   int                              use_krylov;
   struct NestedKrylov_args_struct *krylov;

   /* Only one global-relaxation solver is active at a time. */
   union
   {
      AMG_args  amg;
      ILU_args  ilu;
      FSAI_args fsai;
#if HYPRE_CHECK_MIN_VERSION(30100, 55)
      Schwarz_args schwarz;
#endif
   };
} MGRgrlx_args;

/*--------------------------------------------------------------------------
 * MGR level arguments struct
 *--------------------------------------------------------------------------*/

typedef struct MGRlvl_args_struct
{
   StackIntArray f_dofs;

   HYPRE_Int prolongation_type;
   HYPRE_Int restriction_type;
   HYPRE_Int coarse_level_type;

   MGRfrlx_args f_relaxation;
   MGRgrlx_args g_relaxation;
} MGRlvl_args;

/*--------------------------------------------------------------------------
 * MGR preconditioner arguments struct
 *--------------------------------------------------------------------------*/

struct MGR_args_struct
{
   IntArray      *dofmap;
   HYPRE_IJVector vec_nn;
   uint64_t       rbm_input_generation;
   /* Per-level application-provided coarse (Schur) operators (coarse_level_type: user) */
   const HYPRE_IJMatrix *coarse_schur;

   HYPRE_Int  non_c_to_f;
   HYPRE_Int  pmax;
   HYPRE_Int  max_iter;
   HYPRE_Int  num_levels;
   HYPRE_Int  relax_type; /* TODO: we shouldn't need this */
   HYPRE_Int  print_level;
   HYPRE_Int  cycle;            /* 1=V-cycle (default), 2=W-cycle */
   HYPRE_Int  cycle_smooth_pos; /* 1=pre, 2=post, 3=pre+post */
   HYPRE_Int  nonglk_max_elmts;
   HYPRE_Real tolerance;
   HYPRE_Real coarse_th;

   MGRlvl_args level[MAX_MGR_LEVELS - 1];
   MGRcls_args coarsest_level;

   HYPRE_Solver csolver;
   HYPRE_Int    csolver_type;
   HYPRE_Int    keep_csolver;
   HYPRE_Solver frelax[MAX_MGR_LEVELS - 1];
   HYPRE_Solver grelax[MAX_MGR_LEVELS - 1];
   HYPRE_Int    keep_frelax[MAX_MGR_LEVELS - 1];
   HYPRE_Int    keep_grelax[MAX_MGR_LEVELS - 1];
   HYPRE_Int    num_active_levels;
   HYPRE_Int    active_level_map[MAX_MGR_LEVELS - 1];
   HYPRE_Int   *point_marker_data;
};

/*--------------------------------------------------------------------------
 * Public prototypes
 *--------------------------------------------------------------------------*/

void hypredrv_MGRSetDefaultArgs(MGR_args *);
void hypredrv_MGRSetArgs(void *, const struct YAMLnode_struct *);
void hypredrv_MGRSetDofmap(MGR_args *, IntArray *);
void hypredrv_MGRSetDofLabels(const DofLabelMap *);
void hypredrv_MGRSetNearNullSpace(MGR_args *, HYPRE_IJVector);
void hypredrv_MGRSetCoarseSchur(MGR_args *, const HYPRE_IJMatrix *);
void hypredrv_MGRCreate(MGR_args *, HYPRE_Solver *, const struct Stats_struct *, int);
int  hypredrv_MGRComponentReuseSetupMode(MGR_args *, const struct Stats_struct *, int);
int  hypredrv_MGRComponentReuseShouldKeepOuter(const MGR_args *, const IntArray *,
                                               const struct Stats_struct *, int);
void hypredrv_MGRRefreshComponentsForSetup(MGR_args *, HYPRE_Solver, const IntArray *,
                                           const struct Stats_struct *, int);
void hypredrv_MGRSelectCachedSolversToKeep(MGR_args *, const IntArray *,
                                           const struct Stats_struct *, int);
void hypredrv_MGRCountCachedSolvers(const MGR_args *, int *, int *, int *);
void hypredrv_MGRCountKeepFlags(const MGR_args *, int *, int *, int *);
void hypredrv_MGRDestroyCachedSolvers(MGR_args *, int);
void hypredrv_MGRForgetCachedSolvers(MGR_args *);
void hypredrv_MGRDestroyNestedSolverArgs(MGR_args *);
int  hypredrv_MGRNestedFRelaxWrapperIsLive(HYPRE_Solver);
HYPRE_Solver hypredrv_MGRNestedFRelaxWrapperGetInner(HYPRE_Solver);
HYPRE_Solver hypredrv_MGRNestedFRelaxWrapperDetachInner(HYPRE_Solver);
void         hypredrv_MGRNestedFRelaxWrapperFree(HYPRE_Solver *);

#endif /* MGR_HEADER */
