#include <petsc.h>
#include <metis.h>

int main(int argc, char **argv)
{
    idx_t *addresses, *connectivities;
    // Number of vertices:
    int n = 1000000;

    PetscCall(PetscInitialize(&argc, &argv, NULL, NULL));
    PetscCall(PetscOptionsGetInt(NULL, NULL, "-n", &n, NULL));
    // Addresses should for example be [0 1 3 5 7 8] for 5 vertices
    PetscCall(PetscMalloc1(n+1, &addresses));
    addresses[0] = 0; addresses[1] = 1;
    for (int i = 2; i < n; ++i) addresses[i] = addresses[i-1] + 2;
    addresses[n] = addresses[n-1]+1;
   
#if 1
    // Connected vertices should for example be [1 0 2 1 3 2 4 3] for 5 vertices
    PetscCall(PetscMalloc1(2*(n-1), &connectivities));
    connectivities[0] = 1;
    for (int i = 1; i < n-1; i++)
    {
        connectivities[2*i-1] = i-1;
        connectivities[2*i+0] = i+1;
    }
    connectivities[2*(n-1)-1] = connectivities[2*(n-1)-2]-1;
   
    idx_t nvtxs = n;
    idx_t nEdges = n-1;
   
    idx_t metisOptions[METIS_NOPTIONS];
    METIS_SetDefaultOptions(metisOptions);
    metisOptions[METIS_OPTION_NUMBERING] = 0; // c numbering (start at 0)

    idx_t ncon = 1;
    idx_t nParts = 2;
    idx_t objval;
    idx_t part[nvtxs];
   
    int ret = METIS_PartGraphRecursive ( &nvtxs, &ncon, addresses, connectivities, NULL, NULL, NULL, &nParts, NULL, NULL, metisOptions, &objval, part );
    PetscCall(PetscFree(connectivities));
#endif
    PetscCall(PetscFree(addresses));
    PetscCall(PetscFinalize());
    return 0;
}