static char help[] = "Spectral element access patterns with Plex\n\n";

#include <petscdmplex.h>

typedef struct {
  PetscInt  Nf; /* Number of fields */
  PetscInt *Nc; /* Number of components per field */
  PetscInt *k;  /* Spectral order per field */
} AppCtx;

static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options) {
  PetscInt  len;
  PetscBool flg;

  PetscFunctionBeginUser;
  options->Nf = 0;
  options->Nc = NULL;
  options->k  = NULL;

  PetscOptionsBegin(comm, "", "SEM Problem Options", "DMPLEX");
  PetscCall(PetscOptionsBoundedInt("-num_fields", "The number of fields", "ex6.c", options->Nf, &options->Nf, NULL, 0));
  if (options->Nf) {
    len = options->Nf;
    PetscCall(PetscMalloc1(len, &options->Nc));
    PetscCall(PetscOptionsIntArray("-num_components", "The number of components per field", "ex6.c", options->Nc, &len, &flg));
    PetscCheck(!flg || !(len != options->Nf), PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Length of components array is %" PetscInt_FMT " should be %" PetscInt_FMT, len, options->Nf);
    len = options->Nf;
    PetscCall(PetscMalloc1(len, &options->k));
    PetscCall(PetscOptionsIntArray("-order", "The spectral order per field", "ex6.c", options->k, &len, &flg));
    PetscCheck(!flg || !(len != options->Nf), PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Length of order array is %" PetscInt_FMT " should be %" PetscInt_FMT, len, options->Nf);
  }
  PetscOptionsEnd();
  PetscFunctionReturn(0);
}

static PetscErrorCode LoadData2D(DM dm, PetscInt Ni, PetscInt Nj, PetscInt clSize, Vec u, AppCtx *user) {
  PetscInt     i, j, f, c;
  PetscScalar *closure;

  PetscFunctionBeginUser;
  PetscCall(PetscMalloc1(clSize, &closure));
  for (j = 0; j < Nj; ++j) {
    for (i = 0; i < Ni; ++i) {
      PetscInt ki, kj, o = 0;
      PetscCall(PetscArrayzero(closure, clSize));

      for (f = 0; f < user->Nf; ++f) {
        PetscInt ioff = i * user->k[f], joff = j * user->k[f];

        for (kj = 0; kj <= user->k[f]; ++kj) {
          for (ki = 0; ki <= user->k[f]; ++ki) {
            for (c = 0; c < user->Nc[f]; ++c) closure[o++] = ((kj + joff) * (Ni * user->k[f] + 1) + ki + ioff) * user->Nc[f] + c;
          }
        }
      }
      PetscCall(DMPlexVecSetClosure(dm, NULL, u, j * Ni + i, closure, INSERT_VALUES));
    }
  }
  PetscCall(PetscFree(closure));
  PetscFunctionReturn(0);
}

static PetscErrorCode LoadData3D(DM dm, PetscInt Ni, PetscInt Nj, PetscInt Nk, PetscInt clSize, Vec u, AppCtx *user) {
  PetscInt     i, j, k, f, c;
  PetscScalar *closure;

  PetscFunctionBeginUser;
  PetscCall(PetscMalloc1(clSize, &closure));
  for (k = 0; k < Nk; ++k) {
    for (j = 0; j < Nj; ++j) {
      for (i = 0; i < Ni; ++i) {
        PetscInt ki, kj, kk, o = 0;
        PetscCall(PetscArrayzero(closure, clSize));

        for (f = 0; f < user->Nf; ++f) {
          PetscInt ioff = i * user->k[f], joff = j * user->k[f], koff = k * user->k[f];

          for (kk = 0; kk <= user->k[f]; ++kk) {
            for (kj = 0; kj <= user->k[f]; ++kj) {
              for (ki = 0; ki <= user->k[f]; ++ki) {
                for (c = 0; c < user->Nc[f]; ++c) closure[o++] = (((kk + koff) * (Nj * user->k[f] + 1) + kj + joff) * (Ni * user->k[f] + 1) + ki + ioff) * user->Nc[f] + c;
              }
            }
          }
        }
        PetscCall(DMPlexVecSetClosure(dm, NULL, u, (k * Nj + j) * Ni + i, closure, INSERT_VALUES));
      }
    }
  }
  PetscCall(PetscFree(closure));
  PetscFunctionReturn(0);
}

static PetscErrorCode CheckPoint(DM dm, Vec u, PetscInt point, AppCtx *user) {
  PetscSection       s;
  PetscScalar       *a;
  const PetscScalar *array;
  PetscInt           dof, d;

  PetscFunctionBeginUser;
  PetscCall(DMGetLocalSection(dm, &s));
  PetscCall(VecGetArrayRead(u, &array));
  PetscCall(DMPlexPointLocalRead(dm, point, array, &a));
  PetscCall(PetscSectionGetDof(s, point, &dof));
  PetscCall(PetscPrintf(PETSC_COMM_SELF, "Point %" PetscInt_FMT ": ", point));
  for (d = 0; d < dof; ++d) {
    if (d > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
    PetscCall(PetscPrintf(PETSC_COMM_SELF, "%2.0f", (double)PetscRealPart(a[d])));
  }
  PetscCall(PetscPrintf(PETSC_COMM_SELF, "\n"));
  PetscCall(VecRestoreArrayRead(u, &array));
  PetscFunctionReturn(0);
}

static PetscErrorCode ReadData2D(DM dm, Vec u, AppCtx *user) {
  PetscInt cStart, cEnd, cell;

  PetscFunctionBeginUser;
  PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
  for (cell = cStart; cell < cEnd; ++cell) {
    PetscScalar *closure = NULL;
    PetscInt     closureSize, ki, kj, f, c, foff = 0;

    PetscCall(DMPlexVecGetClosure(dm, NULL, u, cell, &closureSize, &closure));
    PetscCall(PetscPrintf(PETSC_COMM_SELF, "Cell %" PetscInt_FMT "\n", cell));
    for (f = 0; f < user->Nf; ++f) {
      PetscCall(PetscPrintf(PETSC_COMM_SELF, "  Field %" PetscInt_FMT "\n", f));
      for (kj = user->k[f]; kj >= 0; --kj) {
        for (ki = 0; ki <= user->k[f]; ++ki) {
          if (ki > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, "  "));
          for (c = 0; c < user->Nc[f]; ++c) {
            if (c > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ","));
            PetscCall(PetscPrintf(PETSC_COMM_SELF, "%2.0f", (double)PetscRealPart(closure[(kj * (user->k[f] + 1) + ki) * user->Nc[f] + c + foff])));
          }
        }
        PetscCall(PetscPrintf(PETSC_COMM_SELF, "\n"));
      }
      PetscCall(PetscPrintf(PETSC_COMM_SELF, "\n\n"));
      foff += PetscSqr(user->k[f] + 1);
    }
    PetscCall(DMPlexVecRestoreClosure(dm, NULL, u, cell, &closureSize, &closure));
    PetscCall(PetscPrintf(PETSC_COMM_SELF, "\n\n"));
  }
  PetscFunctionReturn(0);
}

static PetscErrorCode ReadData3D(DM dm, Vec u, AppCtx *user) {
  PetscInt cStart, cEnd, cell;

  PetscFunctionBeginUser;
  PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
  for (cell = cStart; cell < cEnd; ++cell) {
    PetscScalar *closure = NULL;
    PetscInt     closureSize, ki, kj, kk, f, c, foff = 0;

    PetscCall(DMPlexVecGetClosure(dm, NULL, u, cell, &closureSize, &closure));
    PetscCall(PetscPrintf(PETSC_COMM_SELF, "Cell %" PetscInt_FMT "\n", cell));
    for (f = 0; f < user->Nf; ++f) {
      PetscCall(PetscPrintf(PETSC_COMM_SELF, "  Field %" PetscInt_FMT "\n", f));
      for (kk = user->k[f]; kk >= 0; --kk) {
        for (kj = user->k[f]; kj >= 0; --kj) {
          for (ki = 0; ki <= user->k[f]; ++ki) {
            if (ki > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, "  "));
            for (c = 0; c < user->Nc[f]; ++c) {
              if (c > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ","));
              PetscCall(PetscPrintf(PETSC_COMM_SELF, "%2.0f", (double)PetscRealPart(closure[((kk * (user->k[f] + 1) + kj) * (user->k[f] + 1) + ki) * user->Nc[f] + c + foff])));
            }
          }
          PetscCall(PetscPrintf(PETSC_COMM_SELF, "\n"));
        }
        PetscCall(PetscPrintf(PETSC_COMM_SELF, "\n"));
      }
      PetscCall(PetscPrintf(PETSC_COMM_SELF, "\n\n"));
      foff += PetscSqr(user->k[f] + 1);
    }
    PetscCall(DMPlexVecRestoreClosure(dm, NULL, u, cell, &closureSize, &closure));
    PetscCall(PetscPrintf(PETSC_COMM_SELF, "\n\n"));
  }
  PetscFunctionReturn(0);
}

static PetscErrorCode SetSymmetries(DM dm, PetscSection s, AppCtx *user) {
  PetscInt dim, f, o, i, j, k, c, d;
  DMLabel  depthLabel;

  PetscFunctionBegin;
  PetscCall(DMGetDimension(dm, &dim));
  PetscCall(DMGetLabel(dm, "depth", &depthLabel));
  for (f = 0; f < user->Nf; f++) {
    PetscSectionSym sym;

    if (user->k[f] < 3) continue; /* No symmetries needed for order < 3, because no cell, facet, edge or vertex has more than one node */
    PetscCall(PetscSectionSymCreateLabel(PetscObjectComm((PetscObject)s), depthLabel, &sym));

    for (d = 0; d <= dim; d++) {
      if (d == 1) {
        PetscInt   numDof  = user->k[f] - 1;
        PetscInt   numComp = user->Nc[f];
        PetscInt   minOrnt = -1;
        PetscInt   maxOrnt = 1;
        PetscInt **perms;

        PetscCall(PetscCalloc1(maxOrnt - minOrnt, &perms));
        for (o = minOrnt; o < maxOrnt; o++) {
          PetscInt *perm;

          if (!o) { /* identity */
            perms[o - minOrnt] = NULL;
          } else {
            PetscCall(PetscMalloc1(numDof * numComp, &perm));
            for (i = numDof - 1, k = 0; i >= 0; i--) {
              for (j = 0; j < numComp; j++, k++) perm[k] = i * numComp + j;
            }
            perms[o - minOrnt] = perm;
          }
        }
        PetscCall(PetscSectionSymLabelSetStratum(sym, d, numDof * numComp, minOrnt, maxOrnt, PETSC_OWN_POINTER, (const PetscInt **)perms, NULL));
      } else if (d == 2) {
        PetscInt   perEdge = user->k[f] - 1;
        PetscInt   numDof  = perEdge * perEdge;
        PetscInt   numComp = user->Nc[f];
        PetscInt   minOrnt = -4;
        PetscInt   maxOrnt = 4;
        PetscInt **perms;

        PetscCall(PetscCalloc1(maxOrnt - minOrnt, &perms));
        for (o = minOrnt; o < maxOrnt; o++) {
          PetscInt *perm;

          if (!o) continue; /* identity */
          PetscCall(PetscMalloc1(numDof * numComp, &perm));
          /* We want to perm[k] to list which *localArray* position the *sectionArray* position k should go to for the given orientation*/
          switch (o) {
          case 0: break; /* identity */
          case -2:       /* flip along (-1,-1)--( 1, 1), which swaps edges 0 and 3 and edges 1 and 2.  This swaps the i and j variables */
            for (i = 0, k = 0; i < perEdge; i++) {
              for (j = 0; j < perEdge; j++, k++) {
                for (c = 0; c < numComp; c++) perm[k * numComp + c] = (perEdge * j + i) * numComp + c;
              }
            }
            break;
          case -1: /* flip along (-1, 0)--( 1, 0), which swaps edges 0 and 2.  This reverses the i variable */
            for (i = 0, k = 0; i < perEdge; i++) {
              for (j = 0; j < perEdge; j++, k++) {
                for (c = 0; c < numComp; c++) perm[k * numComp + c] = (perEdge * (perEdge - 1 - i) + j) * numComp + c;
              }
            }
            break;
          case -4: /* flip along ( 1,-1)--(-1, 1), which swaps edges 0 and 1 and edges 2 and 3.  This swaps the i and j variables and reverse both */
            for (i = 0, k = 0; i < perEdge; i++) {
              for (j = 0; j < perEdge; j++, k++) {
                for (c = 0; c < numComp; c++) perm[k * numComp + c] = (perEdge * (perEdge - 1 - j) + (perEdge - 1 - i)) * numComp + c;
              }
            }
            break;
          case -3: /* flip along ( 0,-1)--( 0, 1), which swaps edges 3 and 1.  This reverses the j variable */
            for (i = 0, k = 0; i < perEdge; i++) {
              for (j = 0; j < perEdge; j++, k++) {
                for (c = 0; c < numComp; c++) perm[k * numComp + c] = (perEdge * i + (perEdge - 1 - j)) * numComp + c;
              }
            }
            break;
          case 1: /* rotate section edge 1 to local edge 0.  This swaps the i and j variables and then reverses the j variable */
            for (i = 0, k = 0; i < perEdge; i++) {
              for (j = 0; j < perEdge; j++, k++) {
                for (c = 0; c < numComp; c++) perm[k * numComp + c] = (perEdge * (perEdge - 1 - j) + i) * numComp + c;
              }
            }
            break;
          case 2: /* rotate section edge 2 to local edge 0.  This reverse both i and j variables */
            for (i = 0, k = 0; i < perEdge; i++) {
              for (j = 0; j < perEdge; j++, k++) {
                for (c = 0; c < numComp; c++) perm[k * numComp + c] = (perEdge * (perEdge - 1 - i) + (perEdge - 1 - j)) * numComp + c;
              }
            }
            break;
          case 3: /* rotate section edge 3 to local edge 0.  This swaps the i and j variables and then reverses the i variable */
            for (i = 0, k = 0; i < perEdge; i++) {
              for (j = 0; j < perEdge; j++, k++) {
                for (c = 0; c < numComp; c++) perm[k * numComp + c] = (perEdge * j + (perEdge - 1 - i)) * numComp + c;
              }
            }
            break;
          default: break;
          }
          perms[o - minOrnt] = perm;
        }
        PetscCall(PetscSectionSymLabelSetStratum(sym, d, numDof * numComp, minOrnt, maxOrnt, PETSC_OWN_POINTER, (const PetscInt **)perms, NULL));
      }
    }
    PetscCall(PetscSectionSetFieldSym(s, f, sym));
    PetscCall(PetscSectionSymDestroy(&sym));
  }
  PetscCall(PetscSectionViewFromOptions(s, NULL, "-section_with_sym_view"));
  PetscFunctionReturn(0);
}

int main(int argc, char **argv) {
  DM           dm;
  PetscSection s;
  Vec          u;
  AppCtx       user;
  PetscInt     dim, size = 0, f;

  PetscFunctionBeginUser;
  PetscCall(PetscInitialize(&argc, &argv, NULL, help));
  PetscCall(ProcessOptions(PETSC_COMM_WORLD, &user));
  PetscCall(DMCreate(PETSC_COMM_WORLD, &dm));
  PetscCall(DMSetType(dm, DMPLEX));
  PetscCall(DMSetFromOptions(dm));
  PetscCall(DMViewFromOptions(dm, NULL, "-dm_view"));
  PetscCall(DMGetDimension(dm, &dim));
  /* Create a section for SEM order k */
  {
    PetscInt *numDof, d;

    PetscCall(PetscMalloc1(user.Nf * (dim + 1), &numDof));
    for (f = 0; f < user.Nf; ++f) {
      for (d = 0; d <= dim; ++d) numDof[f * (dim + 1) + d] = PetscPowInt(user.k[f] - 1, d) * user.Nc[f];
      size += PetscPowInt(user.k[f] + 1, d) * user.Nc[f];
    }
    PetscCall(DMSetNumFields(dm, user.Nf));
    PetscCall(DMPlexCreateSection(dm, NULL, user.Nc, numDof, 0, NULL, NULL, NULL, NULL, &s));
    PetscCall(SetSymmetries(dm, s, &user));
    PetscCall(PetscFree(numDof));
  }
  PetscCall(DMSetLocalSection(dm, s));
  /* Create spectral ordering and load in data */
  PetscCall(DMPlexSetClosurePermutationTensor(dm, PETSC_DETERMINE, NULL));
  PetscCall(DMGetLocalVector(dm, &u));
  switch (dim) {
  case 2: PetscCall(LoadData2D(dm, 2, 2, size, u, &user)); break;
  case 3: PetscCall(LoadData3D(dm, 2, 2, 2, size, u, &user)); break;
  }
  /* Remove ordering and check some values */
  PetscCall(PetscSectionSetClosurePermutation(s, (PetscObject)dm, dim, NULL));
  switch (dim) {
  case 2:
    PetscCall(CheckPoint(dm, u, 0, &user));
    PetscCall(CheckPoint(dm, u, 13, &user));
    PetscCall(CheckPoint(dm, u, 15, &user));
    PetscCall(CheckPoint(dm, u, 19, &user));
    break;
  case 3:
    PetscCall(CheckPoint(dm, u, 0, &user));
    PetscCall(CheckPoint(dm, u, 13, &user));
    PetscCall(CheckPoint(dm, u, 15, &user));
    PetscCall(CheckPoint(dm, u, 19, &user));
    break;
  }
  /* Recreate spectral ordering and read out data */
  PetscCall(DMPlexSetClosurePermutationTensor(dm, PETSC_DETERMINE, s));
  switch (dim) {
  case 2: PetscCall(ReadData2D(dm, u, &user)); break;
  case 3: PetscCall(ReadData3D(dm, u, &user)); break;
  }
  PetscCall(DMRestoreLocalVector(dm, &u));
  PetscCall(PetscSectionDestroy(&s));
  PetscCall(DMDestroy(&dm));
  PetscCall(PetscFree(user.Nc));
  PetscCall(PetscFree(user.k));
  PetscCall(PetscFinalize());
  return 0;
}

/*TEST

  # Spectral ordering 2D 0-5
  testset:
    args: -dm_plex_simplex 0 -dm_plex_box_faces 2,2

    test:
      suffix: 0
      args: -num_fields 1 -num_components 1 -order 2
    test:
      suffix: 1
      args: -num_fields 1 -num_components 1 -order 3
    test:
      suffix: 2
      args: -num_fields 1 -num_components 1 -order 5
    test:
      suffix: 3
      args: -num_fields 1 -num_components 2 -order 2
    test:
      suffix: 4
      args: -num_fields 2 -num_components 1,1 -order 2,2
    test:
      suffix: 5
      args: -num_fields 2 -num_components 1,2 -order 2,3

  # Spectral ordering 3D 6-11
  testset:
    args: -dm_plex_dim 3 -dm_plex_simplex 0 -dm_plex_box_faces 2,2,2

    test:
      suffix: 6
      args: -num_fields 1 -num_components 1 -order 2
    test:
      suffix: 7
      args: -num_fields 1 -num_components 1 -order 3
    test:
      suffix: 8
      args: -num_fields 1 -num_components 1 -order 5
    test:
      suffix: 9
      args: -num_fields 1 -num_components 2 -order 2
    test:
      suffix: 10
      args: -num_fields 2 -num_components 1,1 -order 2,2
    test:
      suffix: 11
      args: -num_fields 2 -num_components 1,2 -order 2,3

TEST*/