/*
     Mechanism for register PETSc matrix types
*/
#include <petsc/private/matimpl.h> /*I "petscmat.h" I*/

PetscBool MatRegisterAllCalled = PETSC_FALSE;

/*
   Contains the list of registered Mat routines
*/
PetscFunctionList MatList = NULL;

/* MatGetRootType_Private - Gets the root type of the input matrix's type (e.g., MATAIJ for MATSEQAIJ)

   Not Collective

   Input Parameter:
.  mat      - the input matrix, could be sequential or MPI

   Output Parameter:
.  rootType  - the root matrix type

   Level: developer

.seealso: `MatGetType()`, `MatSetType()`, `MatType`, `Mat`
*/
PetscErrorCode MatGetRootType_Private(Mat mat, MatType *rootType)
{
  PetscBool   found = PETSC_FALSE;
  MatRootName names = MatRootNameList;
  MatType     inType;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(mat, MAT_CLASSID, 1);
  PetscCall(MatGetType(mat, &inType));
  while (names) {
    PetscCall(PetscStrcmp(inType, names->mname, &found));
    if (!found) PetscCall(PetscStrcmp(inType, names->sname, &found));
    if (found) {
      found     = PETSC_TRUE;
      *rootType = names->rname;
      break;
    }
    names = names->next;
  }
  if (!found) *rootType = inType;
  PetscFunctionReturn(PETSC_SUCCESS);
}

/* MatGetMPIMatType_Private - Gets the MPI type corresponding to the input matrix's type (e.g., MATMPIAIJ for MATSEQAIJ)

   Not Collective

   Input Parameter:
.  mat      - the input matrix, could be sequential or MPI

   Output Parameter:
.  MPIType  - the parallel (MPI) matrix type

   Level: developer

.seealso: `MatGetType()`, `MatSetType()`, `MatType`, `Mat`
*/
PetscErrorCode MatGetMPIMatType_Private(Mat mat, MatType *MPIType)
{
  PetscBool   found = PETSC_FALSE;
  MatRootName names = MatRootNameList;
  MatType     inType;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(mat, MAT_CLASSID, 1);
  PetscCall(MatGetType(mat, &inType));
  while (names) {
    PetscCall(PetscStrcmp(inType, names->sname, &found));
    if (!found) PetscCall(PetscStrcmp(inType, names->mname, &found));
    if (!found) PetscCall(PetscStrcmp(inType, names->rname, &found));
    if (found) {
      found    = PETSC_TRUE;
      *MPIType = names->mname;
      break;
    }
    names = names->next;
  }
  PetscCheck(found, PETSC_COMM_SELF, PETSC_ERR_SUP, "No corresponding parallel (MPI) type for this matrix");
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@
  MatSetType - Builds matrix object for a particular matrix type

  Collective

  Input Parameters:
+ mat    - the matrix object
- matype - matrix type

  Options Database Key:
. -mat_type  <method> - Sets the type; see `MatType`

  Level: intermediate

  Note:
  See `MatType` for possible values

.seealso: [](ch_matrices), `Mat`, `PCSetType()`, `VecSetType()`, `MatCreate()`, `MatType`
@*/
PetscErrorCode MatSetType(Mat mat, MatType matype)
{
  PetscBool   sametype, found, subclass = PETSC_FALSE, matMPI = PETSC_FALSE, requestSeq = PETSC_FALSE;
  MatRootName names = MatRootNameList;
  PetscErrorCode (*r)(Mat);

  PetscFunctionBegin;
  PetscValidHeaderSpecific(mat, MAT_CLASSID, 1);

  /* make this special case fast */
  PetscCall(PetscObjectTypeCompare((PetscObject)mat, matype, &sametype));
  if (sametype) PetscFunctionReturn(PETSC_SUCCESS);

  /* suppose with one MPI process, one created an MPIAIJ (mpiaij) matrix with MatCreateMPIAIJWithArrays(), and later tried
     to change its type via '-mat_type aijcusparse'. Even there is only one MPI rank, we need to adapt matype to
     'mpiaijcusparse' so that it will be treated as a subclass of MPIAIJ and proper MatCovert() will be called.
  */
  if (((PetscObject)mat)->type_name) PetscCall(PetscStrbeginswith(((PetscObject)mat)->type_name, "mpi", &matMPI)); /* mat claims itself is an 'mpi' matrix */
  if (matype) PetscCall(PetscStrbeginswith(matype, "seq", &requestSeq));                                           /* user is requesting a 'seq' matrix */
  PetscCheck(!(matMPI && requestSeq), PetscObjectComm((PetscObject)mat), PETSC_ERR_SUP, "Changing an MPI matrix (%s) to a sequential one (%s) is not allowed. Please remove the 'seq' prefix from your matrix type.", ((PetscObject)mat)->type_name, matype);

  while (names) {
    PetscCall(PetscStrcmp(matype, names->rname, &found));
    if (found) {
      PetscMPIInt size;
      PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)mat), &size));
      if (size == 1 && !matMPI) matype = names->sname; /* try to align the requested type (matype) with the existing type per seq/mpi */
      else matype = names->mname;
      break;
    }
    names = names->next;
  }

  PetscCall(PetscObjectTypeCompare((PetscObject)mat, matype, &sametype));
  if (sametype) PetscFunctionReturn(PETSC_SUCCESS);

  PetscCall(PetscFunctionListFind(MatList, matype, &r));
  PetscCheck(r, PetscObjectComm((PetscObject)mat), PETSC_ERR_ARG_UNKNOWN_TYPE, "Unknown Mat type given: %s", matype);

  if (mat->assembled && ((PetscObject)mat)->type_name) PetscCall(PetscStrbeginswith(matype, ((PetscObject)mat)->type_name, &subclass));
  if (subclass) { /* mat is a subclass of the requested 'matype'? */
    PetscCall(MatConvert(mat, matype, MAT_INPLACE_MATRIX, &mat));
    PetscFunctionReturn(PETSC_SUCCESS);
  }
  if (names && mat->assembled) {
    PetscCall(PetscStrbeginswith(names->rname, "sell", &sametype));
    if (sametype) {                                                  /* mattype is MATSELL or its subclass */
      PetscCall(MatConvert(mat, MATSELL, MAT_INPLACE_MATRIX, &mat)); /* convert to matsell first */
      PetscCall(MatConvert(mat, matype, MAT_INPLACE_MATRIX, &mat));
      PetscFunctionReturn(PETSC_SUCCESS);
    }
  }
  PetscTryTypeMethod(mat, destroy);
  mat->ops->destroy = NULL;

  /* should these null spaces be removed? */
  PetscCall(MatNullSpaceDestroy(&mat->nullsp));
  PetscCall(MatNullSpaceDestroy(&mat->nearnullsp));

  PetscCall(PetscMemzero(mat->ops, sizeof(struct _MatOps)));
  mat->preallocated  = PETSC_FALSE;
  mat->assembled     = PETSC_FALSE;
  mat->was_assembled = PETSC_FALSE;

  /*
   Increment, rather than reset these: the object is logically the same, so its logging and
   state is inherited.  Furthermore, resetting makes it possible for the same state to be
   obtained with a different structure, confusing the PC.
  */
  mat->nonzerostate++;
  PetscCall(PetscObjectStateIncrease((PetscObject)mat));

  /* create the new data structure */
  PetscCall((*r)(mat));
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@
  MatGetType - Gets the matrix type as a string from the matrix object.

  Not Collective

  Input Parameter:
. mat - the matrix

  Output Parameter:
. type - name of matrix type

  Level: intermediate

.seealso: [](ch_matrices), `Mat`, `MatType`, `MatSetType()`
@*/
PetscErrorCode MatGetType(Mat mat, MatType *type)
{
  PetscFunctionBegin;
  PetscValidHeaderSpecific(mat, MAT_CLASSID, 1);
  PetscAssertPointer(type, 2);
  *type = ((PetscObject)mat)->type_name;
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@
  MatGetVecType - Gets the vector type the matrix will return with `MatCreateVecs()`

  Not Collective

  Input Parameter:
. mat - the matrix

  Output Parameter:
. vtype - name of vector type

  Level: intermediate

.seealso: [](ch_matrices), `Mat`, `MatType`, `MatSetVecType()`, `VecType`
@*/
PetscErrorCode MatGetVecType(Mat mat, VecType *vtype)
{
  PetscFunctionBegin;
  PetscValidHeaderSpecific(mat, MAT_CLASSID, 1);
  PetscAssertPointer(vtype, 2);
  *vtype = mat->defaultvectype;
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@
  MatSetVecType - Set the vector type the matrix will return with `MatCreateVecs()`

  Collective

  Input Parameters:
+ mat   - the matrix object
- vtype - vector type

  Level: advanced

  Note:
  This is rarely needed in practice since each matrix object internally sets the proper vector type.

.seealso: [](ch_matrices), `Mat`, `VecType`, `VecSetType()`, `MatGetVecType()`
@*/
PetscErrorCode MatSetVecType(Mat mat, VecType vtype)
{
  PetscFunctionBegin;
  PetscValidHeaderSpecific(mat, MAT_CLASSID, 1);
  PetscCall(PetscFree(mat->defaultvectype));
  PetscCall(PetscStrallocpy(vtype, &mat->defaultvectype));
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@C
  MatRegister - Adds a new matrix type implementation that is usable as a `Mat` in PETSc

  Not Collective, No Fortran Support

  Input Parameters:
+ sname    - name of a new user-defined matrix type
- function - routine to create method context

  Level: advanced

  Note:
  `MatRegister()` may be called multiple times to add several user-defined solvers.

  A simpler alternative to using `MatRegister()` for an application specific matrix format is to use `MatCreateShell()`, which
  generates a `Mat` of `MatType` `MATSHELL`. One can then use `MatShellSetContext()` and `MatShellSetOperation()` to provide
  the data structures and routines customized for their matrix format.

  Example Usage:
.vb
   MatRegister("my_mat", MyMatCreate);
.ve

  Then, your solver can be chosen with the procedural interface via `MatSetType(Mat, "my_mat")` or at runtime via the option
  `-mat_type my_mat`

.seealso: [](ch_matrices), `Mat`, `MatType`, `MatSetType()`, `MatRegisterAll()`
@*/
PetscErrorCode MatRegister(const char sname[], PetscErrorCode (*function)(Mat))
{
  PetscFunctionBegin;
  PetscCall(MatInitializePackage());
  PetscCall(PetscFunctionListAdd(&MatList, sname, function));
  PetscFunctionReturn(PETSC_SUCCESS);
}

MatRootName MatRootNameList = NULL;

/*@
  MatRegisterRootName - Registers a name that can be used for either a sequential or its corresponding parallel matrix type.

  Input Parameters:
+ rname - the rootname, for example, `MATAIJ`
. sname - the name of the sequential matrix type, for example, `MATSEQAIJ`
- mname - the name of the parallel matrix type, for example, `MATMPIAIJ`

  Level: developer

  Notes:
  `MatSetType()` and `-mat_type name` will automatically use the sequential or parallel version
  based on the size of the MPI communicator associated with the matrix.

  The matrix rootname should not be confused with the base type of the function
  `PetscObjectBaseTypeCompare()`

  Developer Notes:
  PETSc vectors have a similar rootname that indicates PETSc should automatically select the appropriate `VecType` based on the
  size of the communicator but it is implemented by simply having additional `VecCreate_RootName()` registerer routines that dispatch to the
  appropriate creation routine. Why have two different ways of implementing the same functionality for different types of objects? It is
  confusing.

.seealso: [](ch_matrices), `Mat`, `MatType`, `PetscObjectBaseTypeCompare()`
@*/
PetscErrorCode MatRegisterRootName(const char rname[], const char sname[], const char mname[])
{
  MatRootName names;

  PetscFunctionBegin;
  PetscCall(PetscNew(&names));
  PetscCall(PetscStrallocpy(rname, &names->rname));
  PetscCall(PetscStrallocpy(sname, &names->sname));
  PetscCall(PetscStrallocpy(mname, &names->mname));
  if (!MatRootNameList) {
    MatRootNameList = names;
  } else {
    MatRootName next = MatRootNameList;
    while (next->next) next = next->next;
    next->next = names;
  }
  PetscFunctionReturn(PETSC_SUCCESS);
}