#include "petsc.h"
#include <iostream>

static void
fn_integral(PetscInt,
            PetscInt,
            PetscInt,
            const PetscInt u_off[],
            const PetscInt[],
            const PetscScalar u[],
            const PetscScalar u_t[],
            const PetscScalar u_x[],
            const PetscInt a_off[],
            const PetscInt[],
            const PetscScalar[],
            const PetscScalar[],
            const PetscScalar[],
            PetscReal,
            const PetscReal[],
            PetscInt,
            const PetscScalar[],
            PetscScalar val[])
{
    val[0] = u[0];
}

int
main(int argc, char * argv[])
{
    PetscInitialize(&argc, &argv, nullptr, nullptr);

    MPI_Comm comm = PETSC_COMM_WORLD;
    PetscInt dim = 2;

    PetscReal coords[] = {
        -1, -1,
        0, -1,
        1, -1,
        -1, 0,
        0, 0,
        1, 0,
        -1, 1,
        0, 1,
        1, 1
    };

    PetscInt cells[] = {
        0, 1, 4, 3,
        1, 2, 5, 4,
        3, 4, 7, 6,
        4, 5, 8, 7
    };

    DM dm;
    DMPlexCreateFromCellListPetsc(comm, dim, 4, 9, 4, PETSC_FALSE, cells, dim, coords, &dm);
    DMSetFromOptions(dm);
    DMSetUp(dm);

    PetscBool simplex;
    DMPlexIsSimplex(dm, &simplex);
    PetscFE fe;
    PetscFECreateLagrange(comm, dim, 1, simplex, 1, PETSC_DETERMINE, &fe);
    PetscFESetName(fe, "u");
    DMSetField(dm, 0, nullptr, (PetscObject) fe);
    PetscFEDestroy(&fe);

    DMCreateDS(dm);

    Vec vec;
    DMCreateGlobalVector(dm, &vec);

    PetscSection section;
    DMGetGlobalSection(dm, &section);

    PetscScalar vals[] = { -3., 0, 1., -2., -1, -2, 1., 0, -3. };
    for (int i = 0; i < 9; i++) {
        PetscInt offset;
        PetscSectionGetFieldOffset(section, 4 + i, 0, &offset);
        VecSetValue(vec, offset, vals[i], INSERT_VALUES);
    }

    PetscDS ds;
    DMGetDS(dm, &ds);
    PetscDSSetObjective(ds, 0, fn_integral);

    PetscScalar integral;
    DMPlexComputeIntegralFEM(dm, vec, &integral, nullptr);

    std::cerr << "int = " << integral << std::endl;

    PetscFinalize();

    return 0;
}