#include <petscdmplex.h>
#include <iostream>


PetscInt SetupDOFs(DM dm);
PetscInt CreateGlobalStiffnessMatrix(DM dm, Mat *K);

PetscErrorCode ierr;

int main(int argc, char **argv) {
    static char help[] = "";
    ierr = PetscInitialize(&argc, &argv, NULL, help); CHKERRQ(ierr);

    DM dm;
    PetscInt dim = 3;
    PetscBool interpolate = PETSC_TRUE;
    PetscInt cells[]= {2,2,1};
    ierr = DMPlexCreateHexBoxMesh(PETSC_COMM_WORLD, dim, cells, DM_BOUNDARY_NONE,DM_BOUNDARY_NONE,DM_BOUNDARY_NONE, &dm); CHKERRQ(ierr);

    DM dist;
    ierr = DMPlexDistribute(dm, 0, nullptr, &dist); CHKERRQ(ierr);
    if (dist){
        ierr = DMDestroy(&dm); CHKERRQ(ierr);
        dm = dist;
    }

    ierr = SetupDOFs(dm); CHKERRQ(ierr);

    Mat mat;
    ierr = CreateGlobalStiffnessMatrix(dm, &mat); CHKERRQ(ierr);

    PetscSynchronizedFlush(PETSC_COMM_WORLD,PETSC_STDOUT);
    PetscFinalize();

    return 0;
}


// Assign dofs (3 for each node)
PetscInt SetupDOFs(DM dm) {
    PetscSection s;

    ierr = DMGetDefaultSection(dm,&s); CHKERRQ(ierr);
    ierr = PetscSectionSetNumFields(s, 1); CHKERRQ(ierr);
    PetscInt pStart, pEnd, cStart, cEnd, fStart, fEnd, vStart, vEnd, eStart, eEnd;
    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);

    // set dofs for each point in hasse diagram
    for (PetscInt c = cStart; c < cEnd; ++c) {
        ierr = PetscSectionSetDof(s, c, 0); CHKERRQ(ierr);
    }
    for (PetscInt f = fStart; f < fEnd; ++f) {
        ierr = PetscSectionSetDof(s, f, 0); CHKERRQ(ierr);
    }
    for (PetscInt e = eStart; e < eEnd; ++e) {
        ierr = PetscSectionSetDof(s, e, 0); CHKERRQ(ierr);
    }
    for (PetscInt  v = vStart; v < vEnd; ++v) {
        ierr = PetscSectionSetDof(s, v, 3); CHKERRQ(ierr);
    }
    ierr = PetscSectionSetUp(s); CHKERRQ(ierr);
    ierr = DMSetDefaultSection(dm, s); CHKERRQ(ierr);

    return 0;
}

PetscInt CreateGlobalStiffnessMatrix(DM dm, Mat *K){
    PetscInt cStart, cEnd, nStart, nEnd;
    PetscSection localSection;

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

    ierr = DMCreateMatrix(dm, K);  CHKERRQ(ierr);

    // print matrix debug information
    PetscViewer  viewer;
    ierr = PetscViewerASCIIGetStdout(PETSC_COMM_WORLD, &viewer); CHKERRQ(ierr);
    ierr = PetscViewerSetType(viewer, PETSCVIEWERASCII); CHKERRQ(ierr);
    ierr = PetscViewerSetFormat(viewer, PETSC_VIEWER_ASCII_INFO_DETAIL); CHKERRQ(ierr);
    ierr = MatView(*K,viewer);  CHKERRQ(ierr);

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

    // Edof vector
    PetscInt edof[24];

    // element matrix
    PetscScalar ke[24*24];

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

    // Element loop
    for (PetscInt e = cStart; e < cEnd; e++ ){
        PetscBool useCone = PETSC_TRUE;
        PetscInt clojureCount;
        PetscInt *clojureIndex = nullptr;

        ierr = DMPlexGetTransitiveClosure(dm, e, useCone, &clojureCount, &clojureIndex);  CHKERRQ(ierr);
        int edofIndex = 0;
        for (int v = 0;v<clojureCount;v++) {
            PetscInt cPoint = clojureIndex[v*2];
            PetscInt cOrientation = clojureIndex[v*2+1];
            bool isVertex = cPoint >= nStart && cPoint < nEnd; // is vertex if it exist in the vertex range
            if (isVertex) {
                PetscInt offset;
                ierr = PetscSectionGetOffset(localSection, cPoint, &offset);  CHKERRQ(ierr);
                PetscInt dofSize;
                ierr = PetscSectionGetDof(localSection, cPoint, &dofSize);  CHKERRQ(ierr);
                for (PetscInt 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);

    return 0;
}