static char help[] = "";

#include <petscdmplex.h>

const int dof = 1;

#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);
    PetscSynchronizedFlush(PETSC_COMM_WORLD, PETSC_STDOUT);
    ierr = DMPlexGetDepthStratum(dm, 0, &vStart, &vEnd); CHKERRQ(ierr);
    ierr = PetscSectionSetChart(s, pStart, pEnd); CHKERRQ(ierr);

    for (v = vStart; v < vEnd; ++v) {
        ierr = PetscSectionSetDof(s, v, dof); CHKERRQ(ierr);
        ierr = PetscSectionSetFieldDof(s, v, 0, dof); CHKERRQ(ierr);
    }
    ierr = PetscSectionSetUp(s); CHKERRQ(ierr);
    ierr = DMSetDefaultSection(dm, s); CHKERRQ(ierr);
    ierr = PetscSectionDestroy(&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[dof * 8];  // edof vector
    PetscScalar ke[(dof * 8) * (dof * 8)]; // 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 < (8 * dof) * (8 * dof); 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;
                    edofIndex++;
                }
            }
        }

        ierr = DMPlexRestoreTransitiveClosure(dm, e, useCone, &clojureCount, &clojureIndex); CHKERRQ(ierr);
        ierr = MatSetValuesLocal(*K, 8 * dof, edof, 8 * dof, 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[] = {1, 1, 2};
    PetscInt overlap = 0;
    PetscErrorCode ierr;
    PetscViewer view;

    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 = DMPlexSetAdjacencyUseCone(dm, PETSC_FALSE); CHKERRQ(ierr);
    //ierr = DMPlexSetAdjacencyUseClosure(dm, PETSC_TRUE); CHKERRQ(ierr);

    ierr = DMPlexDistribute(dm, overlap, NULL, &dist); CHKERRQ(ierr);
    if (dist) {
        ierr = DMDestroy(&dm); CHKERRQ(ierr);
        dm = dist;
    }
    ierr = SetupDOFs(dm); CHKERRQ(ierr);
    ierr = CreateGlobalStiffnessMatrix(dm, &mat); CHKERRQ(ierr);

    Vec v;
    ierr = DMCreateLocalVector(dm, &v); CHKERRQ(ierr);

    PetscInt vsize;
    ierr = VecGetSize(v, &vsize);  CHKERRQ(ierr);
    PetscSynchronizedPrintf(PETSC_COMM_WORLD, "Loc size: %i\n", vsize);
    PetscSynchronizedFlush(PETSC_COMM_WORLD, PETSC_STDOUT);


    PetscScalar trace;
    ierr = MatGetTrace(mat, &trace);  CHKERRQ(ierr);
    PetscPrintf(PETSC_COMM_WORLD, "Trace of matrix: %f\n", trace);

    ierr = MatDestroy(&mat); CHKERRQ(ierr);
    ierr = VecDestroy(&v); CHKERRQ(ierr);
    ierr = DMDestroy(&dm); CHKERRQ(ierr);
    PetscFinalize();

    return 0;
}