const char help[] = "Test DMPlexCoordinatesToReference().\n\n";

#include <petscdm.h>
#include <petscdmplex.h>

static PetscErrorCode testIdentity(DM dm, PetscBool dmIsSimplicial, PetscInt cell, PetscRandom randCtx, PetscInt numPoints, PetscReal tol)
{
  PetscInt       i, j, dimC, dimR;
  PetscReal      *preimage, *mapped, *inverted;

  PetscFunctionBegin;
  PetscCall(DMGetDimension(dm,&dimR));
  PetscCall(DMGetCoordinateDim(dm,&dimC));

  PetscCall(DMGetWorkArray(dm,dimR * numPoints,MPIU_REAL,&preimage));
  PetscCall(DMGetWorkArray(dm,dimC * numPoints,MPIU_REAL,&mapped));
  PetscCall(DMGetWorkArray(dm,dimR * numPoints,MPIU_REAL,&inverted));

  for (i = 0; i < dimR * numPoints; i++) {
    PetscCall(PetscRandomGetValueReal(randCtx, &preimage[i]));
  }
  if (dmIsSimplicial && dimR > 1) {
    for (i = 0; i < numPoints; i++) {
      for (j = 0; j < dimR; j++) {
        PetscReal x = preimage[i * dimR + j];
        PetscReal y = preimage[i * dimR + ((j + 1) % dimR)];

        preimage[i * dimR + ((j + 1) % dimR)] = -1. + (y + 1.) * 0.5 * (1. - x);
      }
    }
  }

  PetscCall(DMPlexReferenceToCoordinates(dm,cell,numPoints,preimage,mapped));
  PetscCall(DMPlexCoordinatesToReference(dm,cell,numPoints,mapped,inverted));

  for (i = 0; i < numPoints; i++) {
    PetscReal max = 0.;

    for (j = 0; j < dimR; j++) {
      max = PetscMax(max,PetscAbsReal(preimage[i * dimR + j] - inverted[i * dimR + j]));
    }
    if (max > tol) {
      PetscCall(PetscPrintf(PETSC_COMM_SELF,"Bad inversion for cell %" PetscInt_FMT " with error %g (tol %g): (",cell,(double)max,(double)tol));
      for (j = 0; j < dimR; j++) {
        PetscCall(PetscPrintf(PETSC_COMM_SELF,"%+f",(double) preimage[i * dimR + j]));
        if (j < dimR - 1) {
          PetscCall(PetscPrintf(PETSC_COMM_SELF,","));
        }
      }
      PetscCall(PetscPrintf(PETSC_COMM_SELF,") --> ("));
      for (j = 0; j < dimC; j++) {
        PetscCall(PetscPrintf(PETSC_COMM_SELF,"%+f",(double) mapped[i * dimC + j]));
        if (j < dimC - 1) {
          PetscCall(PetscPrintf(PETSC_COMM_SELF,","));
        }
      }
      PetscCall(PetscPrintf(PETSC_COMM_SELF,") --> ("));
      for (j = 0; j < dimR; j++) {
        PetscCall(PetscPrintf(PETSC_COMM_SELF,"%+f",(double) inverted[i * dimR + j]));
        if (j < dimR - 1) {
          PetscCall(PetscPrintf(PETSC_COMM_SELF,","));
        }
      }
      PetscCall(PetscPrintf(PETSC_COMM_SELF,")\n"));
    } else {
      char   strBuf[BUFSIZ] = {'\0'};
      size_t offset = 0, count;

      PetscCall(PetscSNPrintfCount(strBuf + offset,BUFSIZ-offset,"Good inversion for cell %" PetscInt_FMT ": (",&count,cell));
      offset += count;
      for (j = 0; j < dimR; j++) {
        PetscCall(PetscSNPrintfCount(strBuf + offset,BUFSIZ-offset,"%+f",&count,(double) preimage[i * dimR + j]));
        offset += count;
        if (j < dimR - 1) {
          PetscCall(PetscSNPrintfCount(strBuf + offset,BUFSIZ-offset,",",&count));
          offset += count;
        }
      }
      PetscCall(PetscSNPrintfCount(strBuf + offset,BUFSIZ-offset,") --> (",&count));
      offset += count;
      for (j = 0; j < dimC; j++) {
        PetscCall(PetscSNPrintfCount(strBuf + offset,BUFSIZ-offset,"%+f",&count,(double) mapped[i * dimC + j]));
        offset += count;
        if (j < dimC - 1) {
          PetscCall(PetscSNPrintfCount(strBuf + offset,BUFSIZ-offset,",",&count));
          offset += count;
        }
      }
      PetscCall(PetscSNPrintfCount(strBuf + offset,BUFSIZ-offset,") --> (",&count));
      offset += count;
      for (j = 0; j < dimR; j++) {
        PetscCall(PetscSNPrintfCount(strBuf + offset,BUFSIZ-offset,"%+f",&count,(double) inverted[i * dimR + j]));
        offset += count;
        if (j < dimR - 1) {
          PetscCall(PetscSNPrintfCount(strBuf + offset,BUFSIZ-offset,",",&count));
          offset += count;
        }
      }
      PetscCall(PetscSNPrintfCount(strBuf + offset,BUFSIZ-offset,")\n",&count));
      PetscCall(PetscInfo(dm,"%s",strBuf));
    }
  }

  PetscCall(DMRestoreWorkArray(dm,dimR * numPoints,MPIU_REAL,&inverted));
  PetscCall(DMRestoreWorkArray(dm,dimC * numPoints,MPIU_REAL,&mapped));
  PetscCall(DMRestoreWorkArray(dm,dimR * numPoints,MPIU_REAL,&preimage));
  PetscFunctionReturn(0);
}

static PetscErrorCode identityEmbedding(PetscInt dim, PetscReal time, const PetscReal *x, PetscInt Nf, PetscScalar *u, void *ctx)
{
  PetscInt i;

  for (i = 0; i < dim; i++) {u[i] = x[i];};
  return 0;
}

int main(int argc, char **argv)
{
  PetscRandom    randCtx;
  PetscInt       dim, dimC, isSimplex, isFE, numTests = 10;
  PetscReal      perturb = 0.1, tol = 10. * PETSC_SMALL;

  PetscFunctionBeginUser;
  PetscCall(PetscInitialize(&argc,&argv,NULL,help));
  PetscCall(PetscRandomCreate(PETSC_COMM_SELF,&randCtx));
  PetscCall(PetscRandomSetInterval(randCtx,-1.,1.));
  PetscCall(PetscRandomSetFromOptions(randCtx));
  PetscOptionsBegin(PETSC_COMM_WORLD,NULL,"ex21",NULL);
  PetscCall(PetscOptionsReal("-vertex_perturbation","scale of random vertex distortion",NULL,perturb,&perturb,NULL));
  PetscCall(PetscOptionsBoundedInt("-num_test_points","number of points to test",NULL,numTests,&numTests,NULL,0));
  PetscOptionsEnd();
  for (dim = 1; dim <= 3; dim++) {
    for (dimC = dim; dimC <= PetscMin(3,dim + 1); dimC++) {
      for (isSimplex = 0; isSimplex < 2; isSimplex++) {
        for (isFE = 0; isFE < 2; isFE++) {
          DM           dm;
          Vec          coords;
          PetscScalar *coordArray;
          PetscReal    noise;
          PetscInt     i, n, order = 1;

          PetscCall(DMPlexCreateReferenceCell(PETSC_COMM_SELF, DMPolytopeTypeSimpleShape(dim, isSimplex ? PETSC_TRUE : PETSC_FALSE), &dm));
          if (isFE) {
            DM             dmCoord;
            PetscSpace     sp;
            PetscFE        fe;
            Vec            localCoords;
            PetscErrorCode (*funcs[1])(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *) = {identityEmbedding};
            void           *ctxs[1] = {NULL};

            PetscCall(PetscFECreateDefault(PetscObjectComm((PetscObject) dm),dim,dim,isSimplex ? PETSC_TRUE : PETSC_FALSE,isSimplex ? NULL : "tensor_" ,PETSC_DEFAULT,&fe));
            PetscCall(PetscFEGetBasisSpace(fe,&sp));
            PetscCall(PetscSpaceGetDegree(sp,&order,NULL));
            PetscCall(DMSetField(dm,0,NULL,(PetscObject)fe));
            PetscCall(DMCreateDS(dm));
            PetscCall(DMCreateLocalVector(dm,&localCoords));
            PetscCall(DMProjectFunctionLocal(dm,0,funcs,ctxs,INSERT_VALUES,localCoords));
            PetscCall(VecSetDM(localCoords,NULL)); /* This is necessary to prevent a reference loop */
            PetscCall(DMClone(dm,&dmCoord));
            PetscCall(DMSetField(dmCoord,0,NULL,(PetscObject)fe));
            PetscCall(PetscFEDestroy(&fe));
            PetscCall(DMCreateDS(dmCoord));
            PetscCall(DMSetCoordinateDM(dm,dmCoord));
            PetscCall(DMDestroy(&dmCoord));
            PetscCall(DMSetCoordinatesLocal(dm,localCoords));
            PetscCall(VecDestroy(&localCoords));
          }
          PetscCall(PetscInfo(dm,"Testing %s%" PetscInt_FMT " %" PetscInt_FMT "D mesh embedded in %" PetscInt_FMT "D\n",isSimplex ? "P" : "Q" , order, dim, dimC));
          PetscCall(DMGetCoordinatesLocal(dm,&coords));
          PetscCall(VecGetLocalSize(coords,&n));
          if (dimC > dim) { /* reembed in higher dimension */
            PetscSection sec, newSec;
            PetscInt     pStart, pEnd, p, i, newN;
            Vec          newVec;
            DM           coordDM;
            PetscScalar  *newCoordArray;

            PetscCall(DMGetCoordinateSection(dm,&sec));
            PetscCall(PetscSectionCreate(PetscObjectComm((PetscObject)sec),&newSec));
            PetscCall(PetscSectionSetNumFields(newSec,1));
            PetscCall(PetscSectionGetChart(sec,&pStart,&pEnd));
            PetscCall(PetscSectionSetChart(newSec,pStart,pEnd));
            for (p = pStart; p < pEnd; p++) {
              PetscInt nDof;

              PetscCall(PetscSectionGetDof(sec,p,&nDof));
              PetscCheck(nDof % dim == 0,PETSC_COMM_SELF,PETSC_ERR_PLIB,"Coordinate section point %" PetscInt_FMT " has %" PetscInt_FMT " dofs != 0 mod %" PetscInt_FMT,p,nDof,dim);
              PetscCall(PetscSectionSetDof(newSec,p,(nDof/dim)*dimC));
              PetscCall(PetscSectionSetFieldDof(newSec,p,0,(nDof/dim)*dimC));
            }
            PetscCall(PetscSectionSetUp(newSec));
            PetscCall(PetscSectionGetStorageSize(newSec,&newN));
            PetscCall(VecCreateSeq(PETSC_COMM_SELF,newN,&newVec));
            PetscCall(VecGetArray(newVec,&newCoordArray));
            PetscCall(VecGetArray(coords,&coordArray));
            for (i = 0; i < n / dim; i++) {
              PetscInt j;

              for (j = 0; j < dim; j++) {
                newCoordArray[i * dimC + j] = coordArray[i * dim + j];
              }
              for (; j < dimC; j++) {
                newCoordArray[i * dimC + j] = 0.;
              }
            }
            PetscCall(VecRestoreArray(coords,&coordArray));
            PetscCall(VecRestoreArray(newVec,&newCoordArray));
            PetscCall(DMSetCoordinateDim(dm,dimC));
            PetscCall(DMSetCoordinateSection(dm,dimC,newSec));
            PetscCall(DMSetCoordinatesLocal(dm,newVec));
            PetscCall(VecDestroy(&newVec));
            PetscCall(PetscSectionDestroy(&newSec));
            PetscCall(DMGetCoordinatesLocal(dm,&coords));
            PetscCall(DMGetCoordinateDM(dm,&coordDM));
            if (isFE) {
              PetscFE fe;

              PetscCall(PetscFECreateDefault(PetscObjectComm((PetscObject) dm),dim,dimC,isSimplex ? PETSC_TRUE : PETSC_FALSE,isSimplex ? NULL : "tensor_",PETSC_DEFAULT,&fe));
              PetscCall(DMSetField(coordDM,0,NULL,(PetscObject)fe));
              PetscCall(PetscFEDestroy(&fe));
              PetscCall(DMCreateDS(coordDM));
            }
            PetscCall(DMSetCoordinateDim(coordDM,dimC));
            PetscCall(VecGetLocalSize(coords,&n));
          }
          PetscCall(VecGetArray(coords,&coordArray));
          for (i = 0; i < n; i++) {
            PetscCall(PetscRandomGetValueReal(randCtx,&noise));
            coordArray[i] += noise * perturb;
          }
          PetscCall(VecRestoreArray(coords,&coordArray));
          PetscCall(DMSetCoordinatesLocal(dm, coords));
          PetscCall(testIdentity(dm, isSimplex ? PETSC_TRUE : PETSC_FALSE, 0, randCtx, numTests, tol));
          PetscCall(DMDestroy(&dm));
        }
      }
    }
  }
  PetscCall(PetscRandomDestroy(&randCtx));
  PetscCall(PetscFinalize());
  return 0;
}

/*TEST

  test:
    suffix: 0
    args: -petscspace_degree 2 -tensor_petscspace_degree 2

TEST*/