/*
 GAMG geometric-algebric multigrid PC - Mark Adams 2011
 */

#include <../src/ksp/pc/impls/gamg/gamg.h> /*I "petscpc.h" I*/
#include <petscblaslapack.h>
#include <petscdm.h>
#include <petsc/private/kspimpl.h>

typedef struct {
  PetscInt   nsmooths;
  PetscInt   aggressive_coarsening_levels; // number of aggressive coarsening levels (square or MISk)
  MatCoarsen crs;
} PC_GAMG_AGG;

/*@
  PCGAMGSetNSmooths - Set number of smoothing steps (1 is typical) used for multigrid on all the levels

  Logically Collective

  Input Parameters:
+ pc - the preconditioner context
- n  - the number of smooths

  Options Database Key:
. -pc_gamg_agg_nsmooths <nsmooth, default=1> - number of smoothing steps to use with smooth aggregation

  Level: intermediate

.seealso: `PCMG`, `PCGAMG`
@*/
PetscErrorCode PCGAMGSetNSmooths(PC pc, PetscInt n)
{
  PetscFunctionBegin;
  PetscValidHeaderSpecific(pc, PC_CLASSID, 1);
  PetscValidLogicalCollectiveInt(pc, n, 2);
  PetscTryMethod(pc, "PCGAMGSetNSmooths_C", (PC, PetscInt), (pc, n));
  PetscFunctionReturn(PETSC_SUCCESS);
}

static PetscErrorCode PCGAMGSetNSmooths_AGG(PC pc, PetscInt n)
{
  PC_MG       *mg          = (PC_MG *)pc->data;
  PC_GAMG     *pc_gamg     = (PC_GAMG *)mg->innerctx;
  PC_GAMG_AGG *pc_gamg_agg = (PC_GAMG_AGG *)pc_gamg->subctx;

  PetscFunctionBegin;
  pc_gamg_agg->nsmooths = n;
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@
  PCGAMGSetAggressiveLevels -  Use aggressive coarsening on first n levels

  Logically Collective

  Input Parameters:
+ pc - the preconditioner context
- n  - 0, 1 or more

  Options Database Key:
. -pc_gamg_aggressive_coarsening <n,default = 1> - Number of levels to square the graph on before aggregating it

  Level: intermediate

.seealso: `PCGAMG`, `PCGAMGSetThreshold()`
@*/
PetscErrorCode PCGAMGSetAggressiveLevels(PC pc, PetscInt n)
{
  PetscFunctionBegin;
  PetscValidHeaderSpecific(pc, PC_CLASSID, 1);
  PetscValidLogicalCollectiveInt(pc, n, 2);
  PetscTryMethod(pc, "PCGAMGSetAggressiveLevels_C", (PC, PetscInt), (pc, n));
  PetscFunctionReturn(PETSC_SUCCESS);
}

static PetscErrorCode PCGAMGSetAggressiveLevels_AGG(PC pc, PetscInt n)
{
  PC_MG       *mg          = (PC_MG *)pc->data;
  PC_GAMG     *pc_gamg     = (PC_GAMG *)mg->innerctx;
  PC_GAMG_AGG *pc_gamg_agg = (PC_GAMG_AGG *)pc_gamg->subctx;

  PetscFunctionBegin;
  pc_gamg_agg->aggressive_coarsening_levels = n;
  PetscFunctionReturn(PETSC_SUCCESS);
}

static PetscErrorCode PCSetFromOptions_GAMG_AGG(PC pc, PetscOptionItems *PetscOptionsObject)
{
  PC_MG       *mg          = (PC_MG *)pc->data;
  PC_GAMG     *pc_gamg     = (PC_GAMG *)mg->innerctx;
  PC_GAMG_AGG *pc_gamg_agg = (PC_GAMG_AGG *)pc_gamg->subctx;

  PetscFunctionBegin;
  PetscOptionsHeadBegin(PetscOptionsObject, "GAMG-AGG options");
  {
    PetscBool flg;
    PetscCall(PetscOptionsInt("-pc_gamg_agg_nsmooths", "smoothing steps for smoothed aggregation, usually 1", "PCGAMGSetNSmooths", pc_gamg_agg->nsmooths, &pc_gamg_agg->nsmooths, NULL));
    PetscCall(
      PetscOptionsInt("-pc_gamg_square_graph", "Number of aggressive coarsening (MIS-2) levels from finest (alias for -pc_gamg_aggressive_coarsening, deprecated)", "PCGAMGSetAggressiveLevels", pc_gamg_agg->aggressive_coarsening_levels, &pc_gamg_agg->aggressive_coarsening_levels, &flg));
    if (!flg) {
      PetscCall(PetscOptionsInt("-pc_gamg_aggressive_coarsening", "Number of aggressive coarsening (MIS-2) levels from finest", "PCGAMGSetAggressiveLevels", pc_gamg_agg->aggressive_coarsening_levels, &pc_gamg_agg->aggressive_coarsening_levels, NULL));
    } else {
      PetscCall(PetscOptionsInt("-pc_gamg_aggressive_coarsening", "Number of aggressive coarsening (MIS-2) levels from finest", "PCGAMGSetAggressiveLevels", pc_gamg_agg->aggressive_coarsening_levels, &pc_gamg_agg->aggressive_coarsening_levels, &flg));
      if (flg) PetscCall(PetscInfo(pc, "Warning: both -pc_gamg_square_graph and -pc_gamg_aggressive_coarsening are used. -pc_gamg_square_graph is deprecated, Number of aggressive levels is %d\n", (int)pc_gamg_agg->aggressive_coarsening_levels));
    }
  }
  PetscOptionsHeadEnd();
  PetscFunctionReturn(PETSC_SUCCESS);
}

static PetscErrorCode PCDestroy_GAMG_AGG(PC pc)
{
  PC_MG   *mg      = (PC_MG *)pc->data;
  PC_GAMG *pc_gamg = (PC_GAMG *)mg->innerctx;

  PetscFunctionBegin;
  PetscCall(PetscFree(pc_gamg->subctx));
  PetscCall(PetscObjectComposeFunction((PetscObject)pc, "PCGAMGSetNSmooths_C", NULL));
  PetscCall(PetscObjectComposeFunction((PetscObject)pc, "PCGAMGSetAggressiveLevels_C", NULL));
  PetscCall(PetscObjectComposeFunction((PetscObject)pc, "PCSetCoordinates_C", NULL));
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*
   PCSetCoordinates_AGG

   Collective

   Input Parameter:
   . pc - the preconditioner context
   . ndm - dimension of data (used for dof/vertex for Stokes)
   . a_nloc - number of vertices local
   . coords - [a_nloc][ndm] - interleaved coordinate data: {x_0, y_0, z_0, x_1, y_1, ...}
*/

static PetscErrorCode PCSetCoordinates_AGG(PC pc, PetscInt ndm, PetscInt a_nloc, PetscReal *coords)
{
  PC_MG   *mg      = (PC_MG *)pc->data;
  PC_GAMG *pc_gamg = (PC_GAMG *)mg->innerctx;
  PetscInt arrsz, kk, ii, jj, nloc, ndatarows, ndf;
  Mat      mat = pc->pmat;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(pc, PC_CLASSID, 1);
  PetscValidHeaderSpecific(mat, MAT_CLASSID, 1);
  nloc = a_nloc;

  /* SA: null space vectors */
  PetscCall(MatGetBlockSize(mat, &ndf));               /* this does not work for Stokes */
  if (coords && ndf == 1) pc_gamg->data_cell_cols = 1; /* scalar w/ coords and SA (not needed) */
  else if (coords) {
    PetscCheck(ndm <= ndf, PETSC_COMM_SELF, PETSC_ERR_PLIB, "degrees of motion %" PetscInt_FMT " > block size %" PetscInt_FMT, ndm, ndf);
    pc_gamg->data_cell_cols = (ndm == 2 ? 3 : 6); /* displacement elasticity */
    if (ndm != ndf) PetscCheck(pc_gamg->data_cell_cols == ndf, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Don't know how to create null space for ndm=%" PetscInt_FMT ", ndf=%" PetscInt_FMT ".  Use MatSetNearNullSpace().", ndm, ndf);
  } else pc_gamg->data_cell_cols = ndf; /* no data, force SA with constant null space vectors */
  pc_gamg->data_cell_rows = ndatarows = ndf;
  PetscCheck(pc_gamg->data_cell_cols > 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "pc_gamg->data_cell_cols %" PetscInt_FMT " <= 0", pc_gamg->data_cell_cols);
  arrsz = nloc * pc_gamg->data_cell_rows * pc_gamg->data_cell_cols;

  if (!pc_gamg->data || (pc_gamg->data_sz != arrsz)) {
    PetscCall(PetscFree(pc_gamg->data));
    PetscCall(PetscMalloc1(arrsz + 1, &pc_gamg->data));
  }
  /* copy data in - column oriented */
  for (kk = 0; kk < nloc; kk++) {
    const PetscInt M    = nloc * pc_gamg->data_cell_rows; /* stride into data */
    PetscReal     *data = &pc_gamg->data[kk * ndatarows]; /* start of cell */
    if (pc_gamg->data_cell_cols == 1) *data = 1.0;
    else {
      /* translational modes */
      for (ii = 0; ii < ndatarows; ii++) {
        for (jj = 0; jj < ndatarows; jj++) {
          if (ii == jj) data[ii * M + jj] = 1.0;
          else data[ii * M + jj] = 0.0;
        }
      }

      /* rotational modes */
      if (coords) {
        if (ndm == 2) {
          data += 2 * M;
          data[0] = -coords[2 * kk + 1];
          data[1] = coords[2 * kk];
        } else {
          data += 3 * M;
          data[0]         = 0.0;
          data[M + 0]     = coords[3 * kk + 2];
          data[2 * M + 0] = -coords[3 * kk + 1];
          data[1]         = -coords[3 * kk + 2];
          data[M + 1]     = 0.0;
          data[2 * M + 1] = coords[3 * kk];
          data[2]         = coords[3 * kk + 1];
          data[M + 2]     = -coords[3 * kk];
          data[2 * M + 2] = 0.0;
        }
      }
    }
  }
  pc_gamg->data_sz = arrsz;
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*
   PCSetData_AGG - called if data is not set with PCSetCoordinates.
      Looks in Mat for near null space.
      Does not work for Stokes

  Input Parameter:
   . pc -
   . a_A - matrix to get (near) null space out of.
*/
static PetscErrorCode PCSetData_AGG(PC pc, Mat a_A)
{
  PC_MG       *mg      = (PC_MG *)pc->data;
  PC_GAMG     *pc_gamg = (PC_GAMG *)mg->innerctx;
  MatNullSpace mnull;

  PetscFunctionBegin;
  PetscCall(MatGetNearNullSpace(a_A, &mnull));
  if (!mnull) {
    DM dm;
    PetscCall(PCGetDM(pc, &dm));
    if (!dm) PetscCall(MatGetDM(a_A, &dm));
    if (dm) {
      PetscObject deformation;
      PetscInt    Nf;

      PetscCall(DMGetNumFields(dm, &Nf));
      if (Nf) {
        PetscCall(DMGetField(dm, 0, NULL, &deformation));
        PetscCall(PetscObjectQuery((PetscObject)deformation, "nearnullspace", (PetscObject *)&mnull));
        if (!mnull) PetscCall(PetscObjectQuery((PetscObject)deformation, "nullspace", (PetscObject *)&mnull));
      }
    }
  }

  if (!mnull) {
    PetscInt bs, NN, MM;
    PetscCall(MatGetBlockSize(a_A, &bs));
    PetscCall(MatGetLocalSize(a_A, &MM, &NN));
    PetscCheck(MM % bs == 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "MM %" PetscInt_FMT " must be divisible by bs %" PetscInt_FMT, MM, bs);
    PetscCall(PCSetCoordinates_AGG(pc, bs, MM / bs, NULL));
  } else {
    PetscReal         *nullvec;
    PetscBool          has_const;
    PetscInt           i, j, mlocal, nvec, bs;
    const Vec         *vecs;
    const PetscScalar *v;

    PetscCall(MatGetLocalSize(a_A, &mlocal, NULL));
    PetscCall(MatNullSpaceGetVecs(mnull, &has_const, &nvec, &vecs));
    for (i = 0; i < nvec; i++) {
      PetscCall(VecGetLocalSize(vecs[i], &j));
      PetscCheck(j == mlocal, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Attached null space vector size %" PetscInt_FMT " != matrix size %" PetscInt_FMT, j, mlocal);
    }
    pc_gamg->data_sz = (nvec + !!has_const) * mlocal;
    PetscCall(PetscMalloc1((nvec + !!has_const) * mlocal, &nullvec));
    if (has_const)
      for (i = 0; i < mlocal; i++) nullvec[i] = 1.0;
    for (i = 0; i < nvec; i++) {
      PetscCall(VecGetArrayRead(vecs[i], &v));
      for (j = 0; j < mlocal; j++) nullvec[(i + !!has_const) * mlocal + j] = PetscRealPart(v[j]);
      PetscCall(VecRestoreArrayRead(vecs[i], &v));
    }
    pc_gamg->data           = nullvec;
    pc_gamg->data_cell_cols = (nvec + !!has_const);
    PetscCall(MatGetBlockSize(a_A, &bs));
    pc_gamg->data_cell_rows = bs;
  }
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*
  formProl0 - collect null space data for each aggregate, do QR, put R in coarse grid data and Q in P_0

  Input Parameter:
   . agg_llists - list of arrays with aggregates -- list from selected vertices of aggregate unselected vertices
   . bs - row block size
   . nSAvec - column bs of new P
   . my0crs - global index of start of locals
   . data_stride - bs*(nloc nodes + ghost nodes) [data_stride][nSAvec]
   . data_in[data_stride*nSAvec] - local data on fine grid
   . flid_fgid[data_stride/bs] - make local to global IDs, includes ghosts in 'locals_llist'

  Output Parameter:
   . a_data_out - in with fine grid data (w/ghosts), out with coarse grid data
   . a_Prol - prolongation operator
*/
static PetscErrorCode formProl0(PetscCoarsenData *agg_llists, PetscInt bs, PetscInt nSAvec, PetscInt my0crs, PetscInt data_stride, PetscReal data_in[], const PetscInt flid_fgid[], PetscReal **a_data_out, Mat a_Prol)
{
  PetscInt        Istart, my0, Iend, nloc, clid, flid = 0, aggID, kk, jj, ii, mm, nSelected, minsz, nghosts, out_data_stride;
  MPI_Comm        comm;
  PetscReal      *out_data;
  PetscCDIntNd   *pos;
  PCGAMGHashTable fgid_flid;

  PetscFunctionBegin;
  PetscCall(PetscObjectGetComm((PetscObject)a_Prol, &comm));
  PetscCall(MatGetOwnershipRange(a_Prol, &Istart, &Iend));
  nloc = (Iend - Istart) / bs;
  my0  = Istart / bs;
  PetscCheck((Iend - Istart) % bs == 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Iend %" PetscInt_FMT " - Istart %" PetscInt_FMT " must be divisible by bs %" PetscInt_FMT, Iend, Istart, bs);
  Iend /= bs;
  nghosts = data_stride / bs - nloc;

  PetscCall(PCGAMGHashTableCreate(2 * nghosts + 1, &fgid_flid));
  for (kk = 0; kk < nghosts; kk++) PetscCall(PCGAMGHashTableAdd(&fgid_flid, flid_fgid[nloc + kk], nloc + kk));

  /* count selected -- same as number of cols of P */
  for (nSelected = mm = 0; mm < nloc; mm++) {
    PetscBool ise;
    PetscCall(PetscCDEmptyAt(agg_llists, mm, &ise));
    if (!ise) nSelected++;
  }
  PetscCall(MatGetOwnershipRangeColumn(a_Prol, &ii, &jj));
  PetscCheck((ii / nSAvec) == my0crs, PETSC_COMM_SELF, PETSC_ERR_PLIB, "ii %" PetscInt_FMT " /nSAvec %" PetscInt_FMT "  != my0crs %" PetscInt_FMT, ii, nSAvec, my0crs);
  PetscCheck(nSelected == (jj - ii) / nSAvec, PETSC_COMM_SELF, PETSC_ERR_PLIB, "nSelected %" PetscInt_FMT " != (jj %" PetscInt_FMT " - ii %" PetscInt_FMT ")/nSAvec %" PetscInt_FMT, nSelected, jj, ii, nSAvec);

  /* aloc space for coarse point data (output) */
  out_data_stride = nSelected * nSAvec;

  PetscCall(PetscMalloc1(out_data_stride * nSAvec, &out_data));
  for (ii = 0; ii < out_data_stride * nSAvec; ii++) out_data[ii] = PETSC_MAX_REAL;
  *a_data_out = out_data; /* output - stride nSelected*nSAvec */

  /* find points and set prolongation */
  minsz = 100;
  for (mm = clid = 0; mm < nloc; mm++) {
    PetscCall(PetscCDSizeAt(agg_llists, mm, &jj));
    if (jj > 0) {
      const PetscInt lid = mm, cgid = my0crs + clid;
      PetscInt       cids[100]; /* max bs */
      PetscBLASInt   asz = jj, M = asz * bs, N = nSAvec, INFO;
      PetscBLASInt   Mdata = M + ((N - M > 0) ? N - M : 0), LDA = Mdata, LWORK = N * bs;
      PetscScalar   *qqc, *qqr, *TAU, *WORK;
      PetscInt      *fids;
      PetscReal     *data;

      /* count agg */
      if (asz < minsz) minsz = asz;

      /* get block */
      PetscCall(PetscMalloc5(Mdata * N, &qqc, M * N, &qqr, N, &TAU, LWORK, &WORK, M, &fids));

      aggID = 0;
      PetscCall(PetscCDGetHeadPos(agg_llists, lid, &pos));
      while (pos) {
        PetscInt gid1;
        PetscCall(PetscCDIntNdGetID(pos, &gid1));
        PetscCall(PetscCDGetNextPos(agg_llists, lid, &pos));

        if (gid1 >= my0 && gid1 < Iend) flid = gid1 - my0;
        else {
          PetscCall(PCGAMGHashTableFind(&fgid_flid, gid1, &flid));
          PetscCheck(flid >= 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Cannot find gid1 in table");
        }
        /* copy in B_i matrix - column oriented */
        data = &data_in[flid * bs];
        for (ii = 0; ii < bs; ii++) {
          for (jj = 0; jj < N; jj++) {
            PetscReal d                       = data[jj * data_stride + ii];
            qqc[jj * Mdata + aggID * bs + ii] = d;
          }
        }
        /* set fine IDs */
        for (kk = 0; kk < bs; kk++) fids[aggID * bs + kk] = flid_fgid[flid] * bs + kk;
        aggID++;
      }

      /* pad with zeros */
      for (ii = asz * bs; ii < Mdata; ii++) {
        for (jj = 0; jj < N; jj++, kk++) qqc[jj * Mdata + ii] = .0;
      }

      /* QR */
      PetscCall(PetscFPTrapPush(PETSC_FP_TRAP_OFF));
      PetscCallBLAS("LAPACKgeqrf", LAPACKgeqrf_(&Mdata, &N, qqc, &LDA, TAU, WORK, &LWORK, &INFO));
      PetscCall(PetscFPTrapPop());
      PetscCheck(INFO == 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "xGEQRF error");
      /* get R - column oriented - output B_{i+1} */
      {
        PetscReal *data = &out_data[clid * nSAvec];
        for (jj = 0; jj < nSAvec; jj++) {
          for (ii = 0; ii < nSAvec; ii++) {
            PetscCheck(data[jj * out_data_stride + ii] == PETSC_MAX_REAL, PETSC_COMM_SELF, PETSC_ERR_PLIB, "data[jj*out_data_stride + ii] != %e", (double)PETSC_MAX_REAL);
            if (ii <= jj) data[jj * out_data_stride + ii] = PetscRealPart(qqc[jj * Mdata + ii]);
            else data[jj * out_data_stride + ii] = 0.;
          }
        }
      }

      /* get Q - row oriented */
      PetscCallBLAS("LAPACKorgqr", LAPACKorgqr_(&Mdata, &N, &N, qqc, &LDA, TAU, WORK, &LWORK, &INFO));
      PetscCheck(INFO == 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "xORGQR error arg %" PetscBLASInt_FMT, -INFO);

      for (ii = 0; ii < M; ii++) {
        for (jj = 0; jj < N; jj++) qqr[N * ii + jj] = qqc[jj * Mdata + ii];
      }

      /* add diagonal block of P0 */
      for (kk = 0; kk < N; kk++) { cids[kk] = N * cgid + kk; /* global col IDs in P0 */ }
      PetscCall(MatSetValues(a_Prol, M, fids, N, cids, qqr, INSERT_VALUES));
      PetscCall(PetscFree5(qqc, qqr, TAU, WORK, fids));
      clid++;
    } /* coarse agg */
  }   /* for all fine nodes */
  PetscCall(MatAssemblyBegin(a_Prol, MAT_FINAL_ASSEMBLY));
  PetscCall(MatAssemblyEnd(a_Prol, MAT_FINAL_ASSEMBLY));
  PetscCall(PCGAMGHashTableDestroy(&fgid_flid));
  PetscFunctionReturn(PETSC_SUCCESS);
}

static PetscErrorCode PCView_GAMG_AGG(PC pc, PetscViewer viewer)
{
  PC_MG       *mg          = (PC_MG *)pc->data;
  PC_GAMG     *pc_gamg     = (PC_GAMG *)mg->innerctx;
  PC_GAMG_AGG *pc_gamg_agg = (PC_GAMG_AGG *)pc_gamg->subctx;

  PetscFunctionBegin;
  PetscCall(PetscViewerASCIIPrintf(viewer, "      AGG specific options\n"));
  PetscCall(PetscViewerASCIIPrintf(viewer, "        Number of levels to square graph %d\n", (int)pc_gamg_agg->aggressive_coarsening_levels));
  PetscCall(PetscViewerASCIIPrintf(viewer, "        Number smoothing steps %d\n", (int)pc_gamg_agg->nsmooths));
  PetscFunctionReturn(PETSC_SUCCESS);
}

static PetscErrorCode PCGAMGCreateGraph_AGG(PC pc, Mat Amat, Mat *a_Gmat)
{
  PC_MG          *mg          = (PC_MG *)pc->data;
  PC_GAMG        *pc_gamg     = (PC_GAMG *)mg->innerctx;
  PC_GAMG_AGG    *pc_gamg_agg = (PC_GAMG_AGG *)pc_gamg->subctx;
  const PetscReal vfilter     = pc_gamg->threshold[pc_gamg->current_level];
  PetscBool       ishem;
  const char     *prefix;

  PetscFunctionBegin;
  PetscCall(PetscLogEventBegin(petsc_gamg_setup_events[GAMG_GRAPH], 0, 0, 0, 0));
  /* Note: depending on the algorithm that will be used for computing the coarse grid points this should pass PETSC_TRUE or PETSC_FALSE as the first argument */

  /* MATCOARSENHEM requires numerical weights for edges so ensure they are computed */
  PetscCall(MatCoarsenCreate(PetscObjectComm((PetscObject)pc), &pc_gamg_agg->crs));
  PetscCall(PetscObjectGetOptionsPrefix((PetscObject)pc, &prefix));
  PetscCall(PetscObjectSetOptionsPrefix((PetscObject)pc_gamg_agg->crs, prefix));
  PetscCall(MatCoarsenSetFromOptions(pc_gamg_agg->crs));
  PetscCall(PetscObjectTypeCompare((PetscObject)pc_gamg_agg->crs, MATCOARSENHEM, &ishem));

  PetscCall(MatCreateGraph(Amat, PETSC_TRUE, (vfilter >= 0 || ishem) ? PETSC_TRUE : PETSC_FALSE, vfilter, a_Gmat));
  PetscCall(PetscLogEventEnd(petsc_gamg_setup_events[GAMG_GRAPH], 0, 0, 0, 0));
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*
   PCGAMGCoarsen_AGG - supports squaring the graph (deprecated) and new graph for
     communication of QR data used with HEM and MISk coarsening

  Input Parameter:
   . a_pc - this

  Input/Output Parameter:
   . a_Gmat1 - graph to coarsen (in), graph off processor edges for QR gather scatter (out)

  Output Parameter:
   . agg_lists - list of aggregates

*/
static PetscErrorCode PCGAMGCoarsen_AGG(PC a_pc, Mat *a_Gmat1, PetscCoarsenData **agg_lists)
{
  PC_MG       *mg          = (PC_MG *)a_pc->data;
  PC_GAMG     *pc_gamg     = (PC_GAMG *)mg->innerctx;
  PC_GAMG_AGG *pc_gamg_agg = (PC_GAMG_AGG *)pc_gamg->subctx;
  Mat          mat, Gmat1 = *a_Gmat1; /* aggressive graph */
  IS           perm;
  PetscInt     Istart, Iend, Ii, nloc, bs, nn;
  PetscInt    *permute, *degree;
  PetscBool   *bIndexSet;
  PetscReal    hashfact;
  PetscInt     iSwapIndex;
  PetscRandom  random;

  PetscFunctionBegin;
  PetscCall(PetscLogEventBegin(petsc_gamg_setup_events[GAMG_COARSEN], 0, 0, 0, 0));
  PetscCall(MatGetLocalSize(Gmat1, &nn, NULL));
  PetscCall(MatGetBlockSize(Gmat1, &bs));
  PetscCheck(bs == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "bs %" PetscInt_FMT " must be 1", bs);
  nloc = nn / bs;

  /* get MIS aggs - randomize */
  PetscCall(PetscMalloc2(nloc, &permute, nloc, &degree));
  PetscCall(PetscCalloc1(nloc, &bIndexSet));
  for (Ii = 0; Ii < nloc; Ii++) permute[Ii] = Ii;
  PetscCall(PetscRandomCreate(PETSC_COMM_SELF, &random));
  PetscCall(MatGetOwnershipRange(Gmat1, &Istart, &Iend));
  for (Ii = 0; Ii < nloc; Ii++) {
    PetscInt nc;
    PetscCall(MatGetRow(Gmat1, Istart + Ii, &nc, NULL, NULL));
    degree[Ii] = nc;
    PetscCall(MatRestoreRow(Gmat1, Istart + Ii, &nc, NULL, NULL));
  }
  for (Ii = 0; Ii < nloc; Ii++) {
    PetscCall(PetscRandomGetValueReal(random, &hashfact));
    iSwapIndex = (PetscInt)(hashfact * nloc) % nloc;
    if (!bIndexSet[iSwapIndex] && iSwapIndex != Ii) {
      PetscInt iTemp        = permute[iSwapIndex];
      permute[iSwapIndex]   = permute[Ii];
      permute[Ii]           = iTemp;
      iTemp                 = degree[iSwapIndex];
      degree[iSwapIndex]    = degree[Ii];
      degree[Ii]            = iTemp;
      bIndexSet[iSwapIndex] = PETSC_TRUE;
    }
  }
  // create minimum degree ordering
  PetscCall(PetscSortIntWithArray(nloc, degree, permute));

  PetscCall(PetscFree(bIndexSet));
  PetscCall(PetscRandomDestroy(&random));
  PetscCall(ISCreateGeneral(PETSC_COMM_SELF, nloc, permute, PETSC_USE_POINTER, &perm));
  PetscCall(PetscLogEventBegin(petsc_gamg_setup_events[GAMG_MIS], 0, 0, 0, 0));
  PetscCall(MatCoarsenSetGreedyOrdering(pc_gamg_agg->crs, perm));
  PetscCall(MatCoarsenSetAdjacency(pc_gamg_agg->crs, Gmat1));
  PetscCall(MatCoarsenSetStrictAggs(pc_gamg_agg->crs, PETSC_TRUE));
  if (pc_gamg->current_level < pc_gamg_agg->aggressive_coarsening_levels) PetscCall(MatCoarsenMISKSetDistance(pc_gamg_agg->crs, 2)); // hardwire to MIS-2
  else PetscCall(MatCoarsenMISKSetDistance(pc_gamg_agg->crs, 1));                                                                    // MIS
  PetscCall(MatCoarsenApply(pc_gamg_agg->crs));
  PetscCall(MatCoarsenViewFromOptions(pc_gamg_agg->crs, NULL, "-mat_coarsen_view"));
  PetscCall(MatCoarsenGetData(pc_gamg_agg->crs, agg_lists)); /* output */
  PetscCall(MatCoarsenDestroy(&pc_gamg_agg->crs));

  PetscCall(ISDestroy(&perm));
  PetscCall(PetscFree2(permute, degree));
  PetscCall(PetscLogEventEnd(petsc_gamg_setup_events[GAMG_MIS], 0, 0, 0, 0));

  {
    PetscCoarsenData *llist = *agg_lists;
    /* see if we have a matrix that takes precedence (returned from MatCoarsenApply) */
    PetscCall(PetscCDGetMat(llist, &mat));
    if (mat) {
      PetscCall(MatDestroy(&Gmat1));
      *a_Gmat1 = mat; /* output */
    }
  }
  PetscCall(PetscLogEventEnd(petsc_gamg_setup_events[GAMG_COARSEN], 0, 0, 0, 0));
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*
 PCGAMGProlongator_AGG

 Input Parameter:
 . pc - this
 . Amat - matrix on this fine level
 . Graph - used to get ghost data for nodes in
 . agg_lists - list of aggregates
 Output Parameter:
 . a_P_out - prolongation operator to the next level
 */
static PetscErrorCode PCGAMGProlongator_AGG(PC pc, Mat Amat, Mat Gmat, PetscCoarsenData *agg_lists, Mat *a_P_out)
{
  PC_MG         *mg      = (PC_MG *)pc->data;
  PC_GAMG       *pc_gamg = (PC_GAMG *)mg->innerctx;
  const PetscInt col_bs  = pc_gamg->data_cell_cols;
  PetscInt       Istart, Iend, nloc, ii, jj, kk, my0, nLocalSelected, bs;
  Mat            Prol;
  PetscMPIInt    size;
  MPI_Comm       comm;
  PetscReal     *data_w_ghost;
  PetscInt       myCrs0, nbnodes = 0, *flid_fgid;
  MatType        mtype;

  PetscFunctionBegin;
  PetscCall(PetscObjectGetComm((PetscObject)Amat, &comm));
  PetscCheck(col_bs >= 1, comm, PETSC_ERR_PLIB, "Column bs cannot be less than 1");
  PetscCall(PetscLogEventBegin(petsc_gamg_setup_events[GAMG_PROL], 0, 0, 0, 0));
  PetscCallMPI(MPI_Comm_size(comm, &size));
  PetscCall(MatGetOwnershipRange(Amat, &Istart, &Iend));
  PetscCall(MatGetBlockSize(Amat, &bs));
  nloc = (Iend - Istart) / bs;
  my0  = Istart / bs;
  PetscCheck((Iend - Istart) % bs == 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "(Iend %" PetscInt_FMT " - Istart %" PetscInt_FMT ") not divisible by bs %" PetscInt_FMT, Iend, Istart, bs);

  /* get 'nLocalSelected' */
  for (ii = 0, nLocalSelected = 0; ii < nloc; ii++) {
    PetscBool ise;
    /* filter out singletons 0 or 1? */
    PetscCall(PetscCDEmptyAt(agg_lists, ii, &ise));
    if (!ise) nLocalSelected++;
  }

  /* create prolongator, create P matrix */
  PetscCall(MatGetType(Amat, &mtype));
  PetscCall(MatCreate(comm, &Prol));
  PetscCall(MatSetSizes(Prol, nloc * bs, nLocalSelected * col_bs, PETSC_DETERMINE, PETSC_DETERMINE));
  PetscCall(MatSetBlockSizes(Prol, bs, col_bs));
  PetscCall(MatSetType(Prol, mtype));
#if PetscDefined(HAVE_DEVICE)
  PetscBool flg;
  PetscCall(MatBoundToCPU(Amat, &flg));
  PetscCall(MatBindToCPU(Prol, flg));
  if (flg) PetscCall(MatSetBindingPropagates(Prol, PETSC_TRUE));
#endif
  PetscCall(MatSeqAIJSetPreallocation(Prol, col_bs, NULL));
  PetscCall(MatMPIAIJSetPreallocation(Prol, col_bs, NULL, col_bs, NULL));

  /* can get all points "removed" */
  PetscCall(MatGetSize(Prol, &kk, &ii));
  if (!ii) {
    PetscCall(PetscInfo(pc, "%s: No selected points on coarse grid\n", ((PetscObject)pc)->prefix));
    PetscCall(MatDestroy(&Prol));
    *a_P_out = NULL; /* out */
    PetscCall(PetscLogEventEnd(petsc_gamg_setup_events[GAMG_PROL], 0, 0, 0, 0));
    PetscFunctionReturn(PETSC_SUCCESS);
  }
  PetscCall(PetscInfo(pc, "%s: New grid %" PetscInt_FMT " nodes\n", ((PetscObject)pc)->prefix, ii / col_bs));
  PetscCall(MatGetOwnershipRangeColumn(Prol, &myCrs0, &kk));

  PetscCheck((kk - myCrs0) % col_bs == 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "(kk %" PetscInt_FMT " -myCrs0 %" PetscInt_FMT ") not divisible by col_bs %" PetscInt_FMT, kk, myCrs0, col_bs);
  myCrs0 = myCrs0 / col_bs;
  PetscCheck((kk / col_bs - myCrs0) == nLocalSelected, PETSC_COMM_SELF, PETSC_ERR_PLIB, "(kk %" PetscInt_FMT "/col_bs %" PetscInt_FMT " - myCrs0 %" PetscInt_FMT ") != nLocalSelected %" PetscInt_FMT ")", kk, col_bs, myCrs0, nLocalSelected);

  /* create global vector of data in 'data_w_ghost' */
  PetscCall(PetscLogEventBegin(petsc_gamg_setup_events[GAMG_PROLA], 0, 0, 0, 0));
  if (size > 1) { /* get ghost null space data */
    PetscReal *tmp_gdata, *tmp_ldata, *tp2;
    PetscCall(PetscMalloc1(nloc, &tmp_ldata));
    for (jj = 0; jj < col_bs; jj++) {
      for (kk = 0; kk < bs; kk++) {
        PetscInt         ii, stride;
        const PetscReal *tp = pc_gamg->data + jj * bs * nloc + kk;
        for (ii = 0; ii < nloc; ii++, tp += bs) tmp_ldata[ii] = *tp;

        PetscCall(PCGAMGGetDataWithGhosts(Gmat, 1, tmp_ldata, &stride, &tmp_gdata));

        if (!jj && !kk) { /* now I know how many total nodes - allocate TODO: move below and do in one 'col_bs' call */
          PetscCall(PetscMalloc1(stride * bs * col_bs, &data_w_ghost));
          nbnodes = bs * stride;
        }
        tp2 = data_w_ghost + jj * bs * stride + kk;
        for (ii = 0; ii < stride; ii++, tp2 += bs) *tp2 = tmp_gdata[ii];
        PetscCall(PetscFree(tmp_gdata));
      }
    }
    PetscCall(PetscFree(tmp_ldata));
  } else {
    nbnodes      = bs * nloc;
    data_w_ghost = (PetscReal *)pc_gamg->data;
  }

  /* get 'flid_fgid' TODO - move up to get 'stride' and do get null space data above in one step (jj loop) */
  if (size > 1) {
    PetscReal *fid_glid_loc, *fiddata;
    PetscInt   stride;

    PetscCall(PetscMalloc1(nloc, &fid_glid_loc));
    for (kk = 0; kk < nloc; kk++) fid_glid_loc[kk] = (PetscReal)(my0 + kk);
    PetscCall(PCGAMGGetDataWithGhosts(Gmat, 1, fid_glid_loc, &stride, &fiddata));
    PetscCall(PetscMalloc1(stride, &flid_fgid)); /* copy real data to in */
    for (kk = 0; kk < stride; kk++) flid_fgid[kk] = (PetscInt)fiddata[kk];
    PetscCall(PetscFree(fiddata));

    PetscCheck(stride == nbnodes / bs, PETSC_COMM_SELF, PETSC_ERR_PLIB, "stride %" PetscInt_FMT " != nbnodes %" PetscInt_FMT "/bs %" PetscInt_FMT, stride, nbnodes, bs);
    PetscCall(PetscFree(fid_glid_loc));
  } else {
    PetscCall(PetscMalloc1(nloc, &flid_fgid));
    for (kk = 0; kk < nloc; kk++) flid_fgid[kk] = my0 + kk;
  }
  PetscCall(PetscLogEventEnd(petsc_gamg_setup_events[GAMG_PROLA], 0, 0, 0, 0));
  /* get P0 */
  PetscCall(PetscLogEventBegin(petsc_gamg_setup_events[GAMG_PROLB], 0, 0, 0, 0));
  {
    PetscReal *data_out = NULL;
    PetscCall(formProl0(agg_lists, bs, col_bs, myCrs0, nbnodes, data_w_ghost, flid_fgid, &data_out, Prol));
    PetscCall(PetscFree(pc_gamg->data));

    pc_gamg->data           = data_out;
    pc_gamg->data_cell_rows = col_bs;
    pc_gamg->data_sz        = col_bs * col_bs * nLocalSelected;
  }
  PetscCall(PetscLogEventEnd(petsc_gamg_setup_events[GAMG_PROLB], 0, 0, 0, 0));
  if (size > 1) PetscCall(PetscFree(data_w_ghost));
  PetscCall(PetscFree(flid_fgid));

  *a_P_out = Prol; /* out */

  PetscCall(PetscLogEventEnd(petsc_gamg_setup_events[GAMG_PROL], 0, 0, 0, 0));
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*
   PCGAMGOptProlongator_AGG

  Input Parameter:
   . pc - this
   . Amat - matrix on this fine level
 In/Output Parameter:
   . a_P - prolongation operator to the next level
*/
static PetscErrorCode PCGAMGOptProlongator_AGG(PC pc, Mat Amat, Mat *a_P)
{
  PC_MG       *mg          = (PC_MG *)pc->data;
  PC_GAMG     *pc_gamg     = (PC_GAMG *)mg->innerctx;
  PC_GAMG_AGG *pc_gamg_agg = (PC_GAMG_AGG *)pc_gamg->subctx;
  PetscInt     jj;
  Mat          Prol = *a_P;
  MPI_Comm     comm;
  KSP          eksp;
  Vec          bb, xx;
  PC           epc;
  PetscReal    alpha, emax, emin;

  PetscFunctionBegin;
  PetscCall(PetscObjectGetComm((PetscObject)Amat, &comm));
  PetscCall(PetscLogEventBegin(petsc_gamg_setup_events[GAMG_OPT], 0, 0, 0, 0));

  /* compute maximum singular value of operator to be used in smoother */
  if (0 < pc_gamg_agg->nsmooths) {
    /* get eigen estimates */
    if (pc_gamg->emax > 0) {
      emin = pc_gamg->emin;
      emax = pc_gamg->emax;
    } else {
      const char *prefix;

      PetscCall(MatCreateVecs(Amat, &bb, NULL));
      PetscCall(MatCreateVecs(Amat, &xx, NULL));
      PetscCall(KSPSetNoisy_Private(bb));

      PetscCall(KSPCreate(comm, &eksp));
      PetscCall(KSPSetNestLevel(eksp, pc->kspnestlevel));
      PetscCall(PCGetOptionsPrefix(pc, &prefix));
      PetscCall(KSPSetOptionsPrefix(eksp, prefix));
      PetscCall(KSPAppendOptionsPrefix(eksp, "pc_gamg_esteig_"));
      {
        PetscBool isset, sflg;
        PetscCall(MatIsSPDKnown(Amat, &isset, &sflg));
        if (isset && sflg) PetscCall(KSPSetType(eksp, KSPCG));
      }
      PetscCall(KSPSetErrorIfNotConverged(eksp, pc->erroriffailure));
      PetscCall(KSPSetNormType(eksp, KSP_NORM_NONE));

      PetscCall(KSPSetInitialGuessNonzero(eksp, PETSC_FALSE));
      PetscCall(KSPSetOperators(eksp, Amat, Amat));

      PetscCall(KSPGetPC(eksp, &epc));
      PetscCall(PCSetType(epc, PCJACOBI)); /* smoother in smoothed agg. */

      PetscCall(KSPSetTolerances(eksp, PETSC_DEFAULT, PETSC_DEFAULT, PETSC_DEFAULT, 10)); // 10 is safer, but 5 is often fine, can override with -pc_gamg_esteig_ksp_max_it -mg_levels_ksp_chebyshev_esteig 0,0.25,0,1.2

      PetscCall(KSPSetFromOptions(eksp));
      PetscCall(KSPSetComputeSingularValues(eksp, PETSC_TRUE));
      PetscCall(KSPSolve(eksp, bb, xx));
      PetscCall(KSPCheckSolve(eksp, pc, xx));

      PetscCall(KSPComputeExtremeSingularValues(eksp, &emax, &emin));
      PetscCall(PetscInfo(pc, "%s: Smooth P0: max eigen=%e min=%e PC=%s\n", ((PetscObject)pc)->prefix, (double)emax, (double)emin, PCJACOBI));
      PetscCall(VecDestroy(&xx));
      PetscCall(VecDestroy(&bb));
      PetscCall(KSPDestroy(&eksp));
    }
    if (pc_gamg->use_sa_esteig) {
      mg->min_eigen_DinvA[pc_gamg->current_level] = emin;
      mg->max_eigen_DinvA[pc_gamg->current_level] = emax;
      PetscCall(PetscInfo(pc, "%s: Smooth P0: level %" PetscInt_FMT ", cache spectra %g %g\n", ((PetscObject)pc)->prefix, pc_gamg->current_level, (double)emin, (double)emax));
    } else {
      mg->min_eigen_DinvA[pc_gamg->current_level] = 0;
      mg->max_eigen_DinvA[pc_gamg->current_level] = 0;
    }
  } else {
    mg->min_eigen_DinvA[pc_gamg->current_level] = 0;
    mg->max_eigen_DinvA[pc_gamg->current_level] = 0;
  }

  /* smooth P0 */
  for (jj = 0; jj < pc_gamg_agg->nsmooths; jj++) {
    Mat tMat;
    Vec diag;

    PetscCall(PetscLogEventBegin(petsc_gamg_setup_events[GAMG_OPTSM], 0, 0, 0, 0));

    /* smooth P1 := (I - omega/lam D^{-1}A)P0 */
    PetscCall(PetscLogEventBegin(petsc_gamg_setup_matmat_events[pc_gamg->current_level][2], 0, 0, 0, 0));
    PetscCall(MatMatMult(Amat, Prol, MAT_INITIAL_MATRIX, PETSC_DEFAULT, &tMat));
    PetscCall(PetscLogEventEnd(petsc_gamg_setup_matmat_events[pc_gamg->current_level][2], 0, 0, 0, 0));
    PetscCall(MatProductClear(tMat));
    PetscCall(MatCreateVecs(Amat, &diag, NULL));
    PetscCall(MatGetDiagonal(Amat, diag)); /* effectively PCJACOBI */
    PetscCall(VecReciprocal(diag));
    PetscCall(MatDiagonalScale(tMat, diag, NULL));
    PetscCall(VecDestroy(&diag));

    /* TODO: Set a PCFailedReason and exit the building of the AMG preconditioner */
    PetscCheck(emax != 0.0, PetscObjectComm((PetscObject)pc), PETSC_ERR_PLIB, "Computed maximum singular value as zero");
    /* TODO: Document the 1.4 and don't hardwire it in this routine */
    alpha = -1.4 / emax;

    PetscCall(MatAYPX(tMat, alpha, Prol, SUBSET_NONZERO_PATTERN));
    PetscCall(MatDestroy(&Prol));
    Prol = tMat;
    PetscCall(PetscLogEventEnd(petsc_gamg_setup_events[GAMG_OPTSM], 0, 0, 0, 0));
  }
  PetscCall(PetscLogEventEnd(petsc_gamg_setup_events[GAMG_OPT], 0, 0, 0, 0));
  *a_P = Prol;
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*
   PCCreateGAMG_AGG

  Input Parameter:
   . pc -
*/
PetscErrorCode PCCreateGAMG_AGG(PC pc)
{
  PC_MG       *mg      = (PC_MG *)pc->data;
  PC_GAMG     *pc_gamg = (PC_GAMG *)mg->innerctx;
  PC_GAMG_AGG *pc_gamg_agg;

  PetscFunctionBegin;
  /* create sub context for SA */
  PetscCall(PetscNew(&pc_gamg_agg));
  pc_gamg->subctx = pc_gamg_agg;

  pc_gamg->ops->setfromoptions = PCSetFromOptions_GAMG_AGG;
  pc_gamg->ops->destroy        = PCDestroy_GAMG_AGG;
  /* reset does not do anything; setup not virtual */

  /* set internal function pointers */
  pc_gamg->ops->creategraph       = PCGAMGCreateGraph_AGG;
  pc_gamg->ops->coarsen           = PCGAMGCoarsen_AGG;
  pc_gamg->ops->prolongator       = PCGAMGProlongator_AGG;
  pc_gamg->ops->optprolongator    = PCGAMGOptProlongator_AGG;
  pc_gamg->ops->createdefaultdata = PCSetData_AGG;
  pc_gamg->ops->view              = PCView_GAMG_AGG;

  pc_gamg_agg->aggressive_coarsening_levels = 1;
  pc_gamg_agg->nsmooths                     = 1;

  PetscCall(PetscObjectComposeFunction((PetscObject)pc, "PCGAMGSetNSmooths_C", PCGAMGSetNSmooths_AGG));
  PetscCall(PetscObjectComposeFunction((PetscObject)pc, "PCGAMGSetAggressiveLevels_C", PCGAMGSetAggressiveLevels_AGG));
  PetscCall(PetscObjectComposeFunction((PetscObject)pc, "PCSetCoordinates_C", PCSetCoordinates_AGG));
  PetscFunctionReturn(PETSC_SUCCESS);
}