#include <petsc/private/dmimpl.h> /*I "petscdm.h" I*/

/*@
   DMGetLocalVector - Gets a PETSc vector that may be used with the DM local routines. This vector has spaces for the ghost values.

   Not Collective

   Input Parameter:
.  dm - the dm

   Output Parameter:
.  g - the local vector

   Level: beginner

   Note:
   The vector values are NOT initialized and may have garbage in them, so you may need
   to zero them.

   The output parameter, g, is a regular PETSc vector that should be returned with
   DMRestoreLocalVector() DO NOT call VecDestroy() on it.

   This is intended to be used for vectors you need for a short time, like within a single function call.
   For vectors that you intend to keep around (for example in a C struct) or pass around large parts of your
   code you should use DMCreateLocalVector().

   VecStride*() operations can be useful when using DM with dof > 1

.seealso: `DMCreateGlobalVector()`, `VecDuplicate()`, `VecDuplicateVecs()`,
          `DMDACreate1d()`, `DMDACreate2d()`, `DMDACreate3d()`, `DMGlobalToLocalBegin()`,
          `DMGlobalToLocalEnd()`, `DMLocalToGlobalBegin()`, `DMCreateLocalVector()`, `DMRestoreLocalVector()`,
          `VecStrideMax()`, `VecStrideMin()`, `VecStrideNorm()`
@*/
PetscErrorCode DMGetLocalVector(DM dm, Vec *g)
{
  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  PetscValidPointer(g, 2);
  for (PetscInt i = 0; i < DM_MAX_WORK_VECTORS; i++) {
    if (dm->localin[i]) {
      DM vdm;

      *g             = dm->localin[i];
      dm->localin[i] = NULL;

      PetscCall(VecGetDM(*g, &vdm));
      PetscCheck(!vdm, PetscObjectComm((PetscObject)vdm), PETSC_ERR_LIB, "Invalid vector");
      PetscCall(VecSetDM(*g, dm));
      goto alldone;
    }
  }
  PetscCall(DMCreateLocalVector(dm, g));

alldone:
  for (PetscInt i = 0; i < DM_MAX_WORK_VECTORS; i++) {
    if (!dm->localout[i]) {
      dm->localout[i] = *g;
      break;
    }
  }
  PetscFunctionReturn(0);
}

/*@
   DMRestoreLocalVector - Returns a PETSc vector that was
     obtained from `DMGetLocalVector()`. Do not use with vector obtained via
     `DMCreateLocalVector()`.

   Not Collective

   Input Parameters:
+  dm - the dm
-  g - the local vector

   Level: beginner

.seealso: `DM`, `DMCreateGlobalVector()`, `VecDuplicate()`, `VecDuplicateVecs()`,
          `DMDACreate1d()`, `DMDACreate2d()`, `DMDACreate3d()`, `DMGlobalToLocalBegin()`,
          `DMGlobalToLocalEnd()`, `DMLocalToGlobalBegin()`, `DMCreateLocalVector()`, `DMGetLocalVector()`
@*/
PetscErrorCode DMRestoreLocalVector(DM dm, Vec *g)
{
  PetscInt i, j;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  PetscValidPointer(g, 2);
  for (j = 0; j < DM_MAX_WORK_VECTORS; j++) {
    if (*g == dm->localout[j]) {
      DM vdm;

      PetscCall(VecGetDM(*g, &vdm));
      PetscCheck(vdm == dm, PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_WRONGSTATE, "Invalid vector");
      PetscCall(VecSetDM(*g, NULL));
      dm->localout[j] = NULL;
      for (i = 0; i < DM_MAX_WORK_VECTORS; i++) {
        if (!dm->localin[i]) {
          dm->localin[i] = *g;
          goto alldone;
        }
      }
    }
  }
  PetscCall(VecDestroy(g));
alldone:
  *g = NULL;
  PetscFunctionReturn(0);
}

/*@
   DMGetGlobalVector - Gets a PETSc vector that may be used with the `DM` global routines.

   Collective on dm

   Input Parameter:
.  dm - the dm

   Output Parameter:
.  g - the global vector

   Level: beginner

   Note:
   The vector values are NOT initialized and may have garbage in them, so you may need
   to zero them.

   The output parameter, g, is a regular PETSc vector that should be returned with
   `DMRestoreGlobalVector()` DO NOT call `VecDestroy()` on it.

   This is intended to be used for vectors you need for a short time, like within a single function call.
   For vectors that you intend to keep around (for example in a C struct) or pass around large parts of your
   code you should use `DMCreateGlobalVector()`.

   VecStride*() operations can be useful when using `DM` with dof > 1

.seealso: `DM`, `DMCreateGlobalVector()`, `VecDuplicate()`, `VecDuplicateVecs()`,
          `DMDACreate1d()`, `DMDACreate2d()`, `DMDACreate3d()`, `DMGlobalToLocalBegin()`,
          `DMGlobalToLocalEnd()`, `DMLocalToGlobalBegin()`, `DMCreateLocalVector()`, `DMRestoreLocalVector()`
          `VecStrideMax()`, `VecStrideMin()`, `VecStrideNorm()`
@*/
PetscErrorCode DMGetGlobalVector(DM dm, Vec *g)
{
  PetscInt i;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  PetscValidPointer(g, 2);
  for (i = 0; i < DM_MAX_WORK_VECTORS; i++) {
    if (dm->globalin[i]) {
      DM vdm;

      *g              = dm->globalin[i];
      dm->globalin[i] = NULL;

      PetscCall(VecGetDM(*g, &vdm));
      PetscCheck(!vdm, PetscObjectComm((PetscObject)vdm), PETSC_ERR_LIB, "Invalid vector");
      PetscCall(VecSetDM(*g, dm));
      goto alldone;
    }
  }
  PetscCall(DMCreateGlobalVector(dm, g));

alldone:
  for (i = 0; i < DM_MAX_WORK_VECTORS; i++) {
    if (!dm->globalout[i]) {
      dm->globalout[i] = *g;
      break;
    }
  }
  PetscFunctionReturn(0);
}

/*@
   DMClearGlobalVectors - Destroys all the global vectors that have been stashed in this DM

   Collective on dm

   Input Parameter:
.  dm - the dm

   Level: developer

.seealso: `DMCreateGlobalVector()`, `VecDuplicate()`, `VecDuplicateVecs()`,
          `DMDACreate1d()`, `DMDACreate2d()`, `DMDACreate3d()`, `DMGlobalToLocalBegin()`,
          `DMGlobalToLocalEnd()`, `DMLocalToGlobalBegin()`, `DMCreateLocalVector()`, `DMRestoreLocalVector()`
          `VecStrideMax()`, `VecStrideMin()`, `VecStrideNorm()`
@*/
PetscErrorCode DMClearGlobalVectors(DM dm)
{
  PetscInt i;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  for (i = 0; i < DM_MAX_WORK_VECTORS; i++) {
    Vec g;

    PetscCheck(!dm->globalout[i], PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_WRONGSTATE, "Clearing DM of global vectors that has a global vector obtained with DMGetGlobalVector()");
    g               = dm->globalin[i];
    dm->globalin[i] = NULL;
    if (g) {
      DM vdm;

      PetscCall(VecGetDM(g, &vdm));
      PetscCheck(!vdm, PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_WRONGSTATE, "Clearing global vector that has a DM attached");
    }
    PetscCall(VecDestroy(&g));
  }
  PetscFunctionReturn(0);
}

/*@
   DMClearLocalVectors - Destroys all the local vectors that have been stashed in this DM

   Collective on dm

   Input Parameter:
.  dm - the dm

   Level: developer

.seealso: `DMCreateLocalVector()`, `VecDuplicate()`, `VecDuplicateVecs()`,
          `DMDACreate1d()`, `DMDACreate2d()`, `DMDACreate3d()`, `DMLocalToLocalBegin()`,
          `DMLocalToLocalEnd()`, `DMLocalToLocalBegin()`, `DMCreateLocalVector()`, `DMRestoreLocalVector()`
          `VecStrideMax()`, `VecStrideMin()`, `VecStrideNorm()`
@*/
PetscErrorCode DMClearLocalVectors(DM dm)
{
  PetscInt i;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  for (i = 0; i < DM_MAX_WORK_VECTORS; i++) {
    Vec g;

    PetscCheck(!dm->localout[i], PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_WRONGSTATE, "Clearing DM of local vectors that has a local vector obtained with DMGetLocalVector()");
    g              = dm->localin[i];
    dm->localin[i] = NULL;
    if (g) {
      DM vdm;

      PetscCall(VecGetDM(g, &vdm));
      PetscCheck(!vdm, PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_WRONGSTATE, "Clearing local vector that has a DM attached");
    }
    PetscCall(VecDestroy(&g));
  }
  PetscFunctionReturn(0);
}

/*@
   DMRestoreGlobalVector - Returns a PETSc vector that
     obtained from DMGetGlobalVector(). Do not use with vector obtained via
     DMCreateGlobalVector().

   Not Collective

   Input Parameters:
+  dm - the dm
-  g - the global vector

   Level: beginner

.seealso: `DMCreateGlobalVector()`, `VecDuplicate()`, `VecDuplicateVecs()`,
          `DMDACreate1d()`, `DMDACreate2d()`, `DMDACreate3d()`, `DMGlobalToGlobalBegin()`,
          `DMGlobalToGlobalEnd()`, `DMGlobalToGlobal()`, `DMCreateLocalVector()`, `DMGetGlobalVector()`
@*/
PetscErrorCode DMRestoreGlobalVector(DM dm, Vec *g)
{
  PetscInt i, j;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  PetscValidPointer(g, 2);
  PetscCall(VecSetErrorIfLocked(*g, 2));
  for (j = 0; j < DM_MAX_WORK_VECTORS; j++) {
    if (*g == dm->globalout[j]) {
      DM vdm;

      PetscCall(VecGetDM(*g, &vdm));
      PetscCheck(vdm == dm, PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_WRONGSTATE, "Invalid vector");
      PetscCall(VecSetDM(*g, NULL));
      dm->globalout[j] = NULL;
      for (i = 0; i < DM_MAX_WORK_VECTORS; i++) {
        if (!dm->globalin[i]) {
          dm->globalin[i] = *g;
          goto alldone;
        }
      }
    }
  }
  PetscCall(VecDestroy(g));
alldone:
  *g = NULL;
  PetscFunctionReturn(0);
}

/*@C
   DMHasNamedGlobalVector - check for a named, persistent global vector

   Not Collective

   Input Parameters:
+  dm - DM to hold named vectors
-  name - unique name for Vec

   Output Parameter:
.  exists - true if the vector was previously created

   Level: developer

   Note: If a Vec with the given name does not exist, it is created.

.seealso: `DMGetNamedGlobalVector()`, `DMRestoreNamedLocalVector()`
@*/
PetscErrorCode DMHasNamedGlobalVector(DM dm, const char *name, PetscBool *exists)
{
  DMNamedVecLink link;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  PetscValidCharPointer(name, 2);
  PetscValidBoolPointer(exists, 3);
  *exists = PETSC_FALSE;
  for (link = dm->namedglobal; link; link = link->next) {
    PetscBool match;
    PetscCall(PetscStrcmp(name, link->name, &match));
    if (match) {
      *exists = PETSC_TRUE;
      break;
    }
  }
  PetscFunctionReturn(0);
}

/*@C
   DMGetNamedGlobalVector - get access to a named, persistent global vector

   Collective on dm

   Input Parameters:
+  dm - DM to hold named vectors
-  name - unique name for Vec

   Output Parameter:
.  X - named Vec

   Level: developer

   Note: If a Vec with the given name does not exist, it is created.

.seealso: `DMRestoreNamedGlobalVector()`
@*/
PetscErrorCode DMGetNamedGlobalVector(DM dm, const char *name, Vec *X)
{
  DMNamedVecLink link;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  PetscValidCharPointer(name, 2);
  PetscValidPointer(X, 3);
  for (link = dm->namedglobal; link; link = link->next) {
    PetscBool match;

    PetscCall(PetscStrcmp(name, link->name, &match));
    if (match) {
      DM vdm;

      PetscCheck(link->status == DMVEC_STATUS_IN, PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_WRONGSTATE, "Vec name '%s' already checked out", name);
      PetscCall(VecGetDM(link->X, &vdm));
      PetscCheck(!vdm, PetscObjectComm((PetscObject)vdm), PETSC_ERR_LIB, "Invalid vector");
      PetscCall(VecSetDM(link->X, dm));
      goto found;
    }
  }

  /* Create the Vec */
  PetscCall(PetscNew(&link));
  PetscCall(PetscStrallocpy(name, &link->name));
  PetscCall(DMCreateGlobalVector(dm, &link->X));
  link->next      = dm->namedglobal;
  dm->namedglobal = link;

found:
  *X           = link->X;
  link->status = DMVEC_STATUS_OUT;
  PetscFunctionReturn(0);
}

/*@C
   DMRestoreNamedGlobalVector - restore access to a named, persistent global vector

   Collective on dm

   Input Parameters:
+  dm - DM on which the vector was gotten
.  name - name under which the vector was gotten
-  X - Vec to restore

   Level: developer

.seealso: `DMGetNamedGlobalVector()`
@*/
PetscErrorCode DMRestoreNamedGlobalVector(DM dm, const char *name, Vec *X)
{
  DMNamedVecLink link;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  PetscValidCharPointer(name, 2);
  PetscValidPointer(X, 3);
  PetscValidHeaderSpecific(*X, VEC_CLASSID, 3);
  for (link = dm->namedglobal; link; link = link->next) {
    PetscBool match;

    PetscCall(PetscStrcmp(name, link->name, &match));
    if (match) {
      DM vdm;

      PetscCall(VecGetDM(*X, &vdm));
      PetscCheck(link->status == DMVEC_STATUS_OUT, PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_WRONGSTATE, "Vec name '%s' was not checked out", name);
      PetscCheck(link->X == *X, PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_INCOMP, "Attempt to restore Vec name '%s', but Vec does not match the cache", name);
      PetscCheck(vdm == dm, PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_WRONGSTATE, "Invalid vector");

      link->status = DMVEC_STATUS_IN;
      PetscCall(VecSetDM(link->X, NULL));
      *X = NULL;
      PetscFunctionReturn(0);
    }
  }
  SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_INCOMP, "Could not find Vec name '%s' to restore", name);
}

/*@C
   DMHasNamedLocalVector - check for a named, persistent local vector

   Not Collective

   Input Parameters:
+  dm - DM to hold named vectors
-  name - unique name for Vec

   Output Parameter:
.  exists - true if the vector was previously created

   Level: developer

   Note: If a Vec with the given name does not exist, it is created.

.seealso: `DMGetNamedGlobalVector()`, `DMRestoreNamedLocalVector()`
@*/
PetscErrorCode DMHasNamedLocalVector(DM dm, const char *name, PetscBool *exists)
{
  DMNamedVecLink link;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  PetscValidCharPointer(name, 2);
  PetscValidBoolPointer(exists, 3);
  *exists = PETSC_FALSE;
  for (link = dm->namedlocal; link; link = link->next) {
    PetscBool match;
    PetscCall(PetscStrcmp(name, link->name, &match));
    if (match) {
      *exists = PETSC_TRUE;
      break;
    }
  }
  PetscFunctionReturn(0);
}

/*@C
   DMGetNamedLocalVector - get access to a named, persistent local vector

   Not Collective

   Input Parameters:
+  dm - DM to hold named vectors
-  name - unique name for Vec

   Output Parameter:
.  X - named Vec

   Level: developer

   Note: If a Vec with the given name does not exist, it is created.

.seealso: `DMGetNamedGlobalVector()`, `DMRestoreNamedLocalVector()`
@*/
PetscErrorCode DMGetNamedLocalVector(DM dm, const char *name, Vec *X)
{
  DMNamedVecLink link;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  PetscValidCharPointer(name, 2);
  PetscValidPointer(X, 3);
  for (link = dm->namedlocal; link; link = link->next) {
    PetscBool match;

    PetscCall(PetscStrcmp(name, link->name, &match));
    if (match) {
      DM vdm;

      PetscCheck(link->status == DMVEC_STATUS_IN, PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_WRONGSTATE, "Vec name '%s' already checked out", name);
      PetscCall(VecGetDM(link->X, &vdm));
      PetscCheck(!vdm, PetscObjectComm((PetscObject)vdm), PETSC_ERR_LIB, "Invalid vector");
      PetscCall(VecSetDM(link->X, dm));
      goto found;
    }
  }

  /* Create the Vec */
  PetscCall(PetscNew(&link));
  PetscCall(PetscStrallocpy(name, &link->name));
  PetscCall(DMCreateLocalVector(dm, &link->X));
  link->next     = dm->namedlocal;
  dm->namedlocal = link;

found:
  *X           = link->X;
  link->status = DMVEC_STATUS_OUT;
  PetscFunctionReturn(0);
}

/*@C
   DMRestoreNamedLocalVector - restore access to a named, persistent local vector

   Not Collective

   Input Parameters:
+  dm - DM on which the vector was gotten
.  name - name under which the vector was gotten
-  X - Vec to restore

   Level: developer

.seealso: `DMRestoreNamedGlobalVector()`, `DMGetNamedLocalVector()`
@*/
PetscErrorCode DMRestoreNamedLocalVector(DM dm, const char *name, Vec *X)
{
  DMNamedVecLink link;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  PetscValidCharPointer(name, 2);
  PetscValidPointer(X, 3);
  PetscValidHeaderSpecific(*X, VEC_CLASSID, 3);
  for (link = dm->namedlocal; link; link = link->next) {
    PetscBool match;

    PetscCall(PetscStrcmp(name, link->name, &match));
    if (match) {
      DM vdm;

      PetscCall(VecGetDM(*X, &vdm));
      PetscCheck(link->status == DMVEC_STATUS_OUT, PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_WRONGSTATE, "Vec name '%s' was not checked out", name);
      PetscCheck(link->X == *X, PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_INCOMP, "Attempt to restore Vec name '%s', but Vec does not match the cache", name);
      PetscCheck(vdm == dm, PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_WRONGSTATE, "Invalid vector");

      link->status = DMVEC_STATUS_IN;
      PetscCall(VecSetDM(link->X, NULL));
      *X = NULL;
      PetscFunctionReturn(0);
    }
  }
  SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_INCOMP, "Could not find Vec name '%s' to restore", name);
}