/******************************************************************************
 * 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 AMG_HEADER
#define AMG_HEADER

#include <stdint.h>
#include "HYPRE_parcsr_ls.h"
#include "internal/cheby.h"
#include "internal/containers.h"
#include "internal/fsai.h"
#include "internal/ilu.h"

struct YAMLnode_struct;

/*--------------------------------------------------------------------------
 * AMG complex smoother arguments struct
 *--------------------------------------------------------------------------*/

typedef struct AMGsmt_args_struct
{
   HYPRE_Int type;
   HYPRE_Int num_levels;
   HYPRE_Int num_sweeps;
   FSAI_args fsai;
   ILU_args  ilu;
} AMGsmt_args;

/*--------------------------------------------------------------------------
 * AMG relaxation arguments struct
 *--------------------------------------------------------------------------*/

typedef struct AMGrlx_args_struct
{
   HYPRE_Int  type;
   HYPRE_Int  down_type;
   HYPRE_Int  up_type;
   HYPRE_Int  coarse_type;
   HYPRE_Int  down_sweeps;
   HYPRE_Int  up_sweeps;
   HYPRE_Int  coarse_sweeps;
   HYPRE_Int  num_sweeps;
   HYPRE_Int  order;
   HYPRE_Int  points;
   HYPRE_Real weight;
   HYPRE_Real outer_weight;
   Cheby_args chebyshev;
} AMGrlx_args;

/*--------------------------------------------------------------------------
 * AMG aggressive coarsening arguments struct
 *--------------------------------------------------------------------------*/

typedef struct AMGagg_args_struct
{
   HYPRE_Int  num_levels;
   HYPRE_Int  num_paths;
   HYPRE_Int  prolongation_type;
   HYPRE_Int  max_nnz_row;
   HYPRE_Int  P12_max_elements;
   HYPRE_Real P12_trunc_factor;
   HYPRE_Real trunc_factor;
} AMGagg_args;

/*--------------------------------------------------------------------------
 * AMG coarsening arguments struct
 *--------------------------------------------------------------------------*/

typedef struct AMGcsn_args_struct
{
   HYPRE_Int  type;
   HYPRE_Int  rap2;
   HYPRE_Int  mod_rap2;
   HYPRE_Int  keep_transpose;
   HYPRE_Int  sabs;
   HYPRE_Int  num_functions;
   HYPRE_Int  filter_functions;
   HYPRE_Int  nodal;
   HYPRE_Int  nodal_type;         /* HYPRE_BoomerAMGSetNodal mode for GM/LN interp */
   HYPRE_Int  interp_vec_variant; /* -1 = auto, 1 = GM1, 2 = GM2, 3+ = LN interp */
   HYPRE_Int  smooth_interp_vecs; /* -1 = auto; smooth the interp (rigid-body) vectors */
   HYPRE_Int  interp_vec_qmax; /* -1 = auto; max nnz per row added by GM/LN expansion */
   HYPRE_Real interp_vec_abs_q_trunc; /* drop GM/LN Q entries below this magnitude */
   HYPRE_Int  seq_amg_th;
   HYPRE_Int  min_coarse_size;
   HYPRE_Int  max_coarse_size;
   HYPRE_Int  max_levels;
   HYPRE_Real max_row_sum;
   HYPRE_Real strong_th;
} AMGcsn_args;

/*--------------------------------------------------------------------------
 * AMG interpolation arguments struct
 *--------------------------------------------------------------------------*/

typedef struct AMGint_args_struct
{
   HYPRE_Int  prolongation_type;
   HYPRE_Int  restriction_type;
   HYPRE_Int  max_nnz_row;
   HYPRE_Real trunc_factor;
   HYPRE_Real restrict_strong_th;
   HYPRE_Real restrict_filter_th;
} AMGint_args;

/*--------------------------------------------------------------------------
 * AMG preconditioner arguments struct
 *--------------------------------------------------------------------------*/

typedef struct AMG_args_struct
{
   AMGint_args interpolation;
   AMGagg_args aggressive;
   AMGcsn_args coarsening;
   AMGrlx_args relaxation;
   AMGsmt_args smoother;

   HYPRE_Int  max_iter;
   HYPRE_Int  print_level;
   HYPRE_Real tolerance;

   /* GM/LN interpolation variant implied by how the near-null space was built
      (full system vs. projected onto an MGR subblock). Used when the YAML knob
      coarsening.interp_vec_variant is left on "auto". */
   HYPRE_Int       interp_vec_variant;
   HYPRE_Int       num_rbms;
   HYPRE_ParVector rbms[3];
   uint64_t        rbm_source_generation;
   size_t          rbm_source_labels_size;
   uint64_t        rbm_source_labels_hash;
} AMG_args;

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

void hypredrv_AMGSetDefaultArgs(AMG_args *);
void hypredrv_AMGSetArgs(void *, const struct YAMLnode_struct *);
void hypredrv_AMGCreate(const AMG_args *, HYPRE_Solver *);
void hypredrv_AMGSetRBMs(AMG_args *, HYPRE_IJVector);
void hypredrv_AMGSetProjectedRBMs(AMG_args *, HYPRE_IJVector, const IntArray *, uint64_t,
                                  const int *, size_t);
void hypredrv_AMGDestroyRBMs(AMG_args *);
void hypredrv_AMGSetDofFunc(const AMG_args *, const IntArray *, HYPRE_Solver,
                            HYPRE_IJMatrix);

#endif /* AMG_HEADER */
