static char help[] = "";

#include <petscdmplex.h>

#undef __FUNCT__
#define __FUNCT__ "SetupDOFs"
/* Assign dofs (3 for each node) */
PetscErrorCode SetupDOFs(DM dm) {
  PetscSection   s;
  PetscInt       pStart, pEnd, cStart, cEnd, c, fStart, fEnd, f, vStart, vEnd, v, eStart, eEnd, e;
  PetscErrorCode ierr;

  PetscFunctionBeginUser;
  ierr = PetscSectionCreate(PetscObjectComm((PetscObject) dm), &s);CHKERRQ(ierr);
  ierr = PetscSectionSetNumFields(s, 1);CHKERRQ(ierr);
  ierr = DMPlexGetChart(dm, &pStart, &pEnd);CHKERRQ(ierr);
  ierr = DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd);CHKERRQ(ierr);
  ierr = DMPlexGetHeightStratum(dm, 1, &fStart, &fEnd);CHKERRQ(ierr);
  ierr = DMPlexGetHeightStratum(dm, 2, &eStart, &eEnd);CHKERRQ(ierr);
  ierr = DMPlexGetDepthStratum(dm, 0, &vStart, &vEnd);CHKERRQ(ierr);
  ierr = PetscSectionSetChart(s, pStart, pEnd);CHKERRQ(ierr);

#if 0
  for (c = cStart; c < cEnd; ++c) {
    ierr = PetscSectionSetDof(s, c, 0);CHKERRQ(ierr);
  }
  for (f = fStart; f < fEnd; ++f) {
    ierr = PetscSectionSetDof(s, f, 0);CHKERRQ(ierr);
  }
  for (e = eStart; e < eEnd; ++e) {
    ierr = PetscSectionSetDof(s, e, 0);CHKERRQ(ierr);
  }
#endif
  for (v = vStart; v < vEnd; ++v) {
    ierr = PetscSectionSetDof(s, v, 3);CHKERRQ(ierr);
    ierr = PetscSectionSetFieldDof(s, v, 0, 3);CHKERRQ(ierr);
  }
  ierr = PetscSectionSetUp(s);CHKERRQ(ierr);
  ierr = DMSetDefaultSection(dm, s);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "CreateGlobalStiffnessMatrix"
PetscInt CreateGlobalStiffnessMatrix(DM dm, Mat *K)
{
  PetscSection   localSection;
  PetscInt       cStart, cEnd, nStart, nEnd, e;
  PetscInt       edof[24];  // edof vector
  PetscScalar    ke[24*24]; // element matrix
  PetscErrorCode ierr;

  PetscFunctionBeginUser;
  ierr = DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd);CHKERRQ(ierr);
  ierr = DMPlexGetDepthStratum(dm, 0, &nStart, &nEnd);CHKERRQ(ierr);

  ierr = DMCreateMatrix(dm, K);CHKERRQ(ierr);
  ierr = MatViewFromOptions(*K, NULL, "-mat_view");CHKERRQ(ierr);

  ierr = DMGetDefaultSection(dm, &localSection);CHKERRQ(ierr);

  for (e = 0; e < 24*24; e++) {
    ke[e] = 1.0; // placeholder value - replace with local stiffness matrix
  }

  for (e = cStart; e < cEnd; e++) {
    PetscBool useCone = PETSC_TRUE;
    PetscInt  clojureCount;
    PetscInt *clojureIndex = NULL;
    PetscInt  edofIndex    = 0, v;

    ierr = DMPlexGetTransitiveClosure(dm, e, useCone, &clojureCount, &clojureIndex);CHKERRQ(ierr);
    for (v = 0; v < clojureCount; ++v) {
      PetscInt  cPoint       = clojureIndex[v*2];
      PetscInt  cOrientation = clojureIndex[v*2+1];
      PetscBool isVertex     = cPoint >= nStart && cPoint < nEnd ? PETSC_TRUE : PETSC_FALSE; // is vertex if it exist in the vertex range

      if (isVertex) {
        PetscInt offset, dofSize, i;

        ierr = PetscSectionGetOffset(localSection, cPoint, &offset);CHKERRQ(ierr);
        ierr = PetscSectionGetDof(localSection, cPoint, &dofSize);CHKERRQ(ierr);
        for (i = 0; i < dofSize; ++i){
          edof[edofIndex] = offset + i; // assumes that dofSize always equals 3
          edofIndex++;
        }
      }
    }
    ierr = DMPlexRestoreTransitiveClosure(dm, e, useCone, &clojureCount, &clojureIndex);CHKERRQ(ierr);
    ierr = MatSetValuesLocal(*K, 24, edof, 24, edof, ke, ADD_VALUES);CHKERRQ(ierr);
  }
  ierr = MatAssemblyBegin(*K, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);
  ierr = MatAssemblyEnd(*K, MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "main"
int main(int argc, char **argv) {
  DM             dm, dist;
  Mat            mat;
  PetscInt       dim         = 3;
  PetscInt       cells[]     = {2, 2, 1};
  PetscBool      interpolate = PETSC_TRUE;
  PetscErrorCode ierr;

  ierr = PetscInitialize(&argc, &argv, NULL, help); CHKERRQ(ierr);
  ierr = DMPlexCreateHexBoxMesh(PETSC_COMM_WORLD, dim, cells, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, &dm); CHKERRQ(ierr);
  ierr = DMPlexDistribute(dm, 0, NULL, &dist); CHKERRQ(ierr);
  if (dist) {
    ierr = DMDestroy(&dm);CHKERRQ(ierr);
    dm   = dist;
  }
  ierr = SetupDOFs(dm);CHKERRQ(ierr);
  ierr = CreateGlobalStiffnessMatrix(dm, &mat);CHKERRQ(ierr);
  PetscFinalize();
  return 0;
}