/**
 * @file Atoms/Ghosts.cpp
 * @author Miguel Molinos (@migmolper)
 * @brief
 * @version 0.1
 * @date 2024-08-22
 *
 * @copyright Copyright (c) 2024
 *
 */

#include "Macros.hpp"
#include "petscdmswarm.h"

extern PetscMPIInt size_MPI;
extern PetscMPIInt rank_MPI;

extern PetscInt ndiv_mesh_X;
extern PetscInt ndiv_mesh_Y;
extern PetscInt ndiv_mesh_Z;

extern double box_x_min;
extern double box_x_max;
extern double box_y_min;
extern double box_y_max;
extern double box_z_min;
extern double box_z_max;

static PetscErrorCode GetElementCoords(const PetscScalar _coords[],
                                       const PetscInt e2n[],
                                       PetscScalar el_coords[]);

static bool In_Out_Buffer(const Eigen::Vector3d mean_q,
                          const PetscScalar el_coords[]);

/********************************************************************************/

PetscErrorCode create_ghost_atoms(DMD* Simulation, double buffer_width) {

  unsigned int dim = NumberDimensions;

  PetscFunctionBegin;

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    Check we have the correct migration algorithm
      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  DMSwarmMigrateType atoms_migrate_type;
  DMSwarmGetMigrateType(Simulation->atomistic_data, &atoms_migrate_type);

  if (atoms_migrate_type != DMSWARM_MIGRATE_BASIC) {

    PetscCall(
        PetscError(PETSC_COMM_WORLD, __LINE__, "create_ghost_atoms", __FILE__,
                   PETSC_ERR_RETURN, PETSC_ERROR_INITIAL,
                   "Set DMSWARM_MIGRATE_BASIC using DMSwarmSetMigrateType"));

    PetscFunctionReturn(PETSC_ERR_RETURN);
  }

  //! Get global size
  PetscInt n_global_size = 0;
  PetscCall(DMSwarmGetSize(Simulation->atomistic_data, &n_global_size));

  //! Get local size
  PetscInt n_local_size = 0;
  PetscCall(DMSwarmGetLocalSize(Simulation->atomistic_data, &n_local_size));

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    Get the global (x,y,z) indices of the lower left corner and size of the
    local region, excluding ghost points.
    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  Vec coords;
  PetscInt nel, npe;
  const PetscScalar* _coords;
  const PetscInt* element_list;
  PetscScalar el_coords[24];
  double X_ll = box_x_max, Y_ll = box_y_max, Z_ll = box_z_max;
  double X_ur = box_x_min, Y_ur = box_y_min, Z_ur = box_z_min;

  //! @brief Get mesh information
  PetscCall(DMGetCoordinatesLocal(Simulation->background_mesh, &coords));
  PetscCall(VecGetArrayRead(coords, &_coords));
  PetscCall(
      DMDAGetElements(Simulation->background_mesh, &nel, &npe, &element_list));

  for (PetscInt eidx = 0; eidx < nel; eidx++) {

    /* get coords for the element */
    const PetscInt* element = &element_list[npe * eidx];
    PetscCall(GetElementCoords(_coords, element, el_coords));

    X_ll = DMIN(X_ll, el_coords[0 * dim + 0]);
    Y_ll = DMIN(Y_ll, el_coords[0 * dim + 1]);
    Z_ll = DMIN(Z_ll, el_coords[0 * dim + 2]);

    X_ur = DMAX(X_ur, el_coords[7 * dim + 0]);
    Y_ur = DMAX(Y_ur, el_coords[7 * dim + 1]);
    Z_ur = DMAX(Z_ur, el_coords[7 * dim + 2]);
  }

  PetscCall(VecRestoreArrayRead(coords, &_coords));

  PetscCall(DMDARestoreElements(Simulation->background_mesh, &nel, &npe,
                                &element_list));

  MPI_Barrier(MPI_COMM_WORLD);

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    Get the list of processes neighboring this one.
    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  const PetscMPIInt* neig_ranks_ptr;
  PetscCall(DMDAGetNeighbors(Simulation->background_mesh, &neig_ranks_ptr));

  Eigen::Map<const List1D> ranks_list_1D(neig_ranks_ptr, 27);

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    Mark ghost particles in the buffer region of the domain.
    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */

  // Direction (-1 -1 -1)

  int num_ghost_m1_m1_m1 = 0;

  int rank_m1_m1_m1 = ranks_list_1D(0);
  if (rank_m1_m1_m1 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ll;
    buffer_coords[1] = Y_ll;
    buffer_coords[2] = Z_ll;

    buffer_coords[3] = X_ll + buffer_width;
    buffer_coords[4] = Y_ll + buffer_width;
    buffer_coords[5] = Z_ll + buffer_width;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_m1_m1_m1;
        num_ghost_m1_m1_m1++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_m1_m1_m1
              << " from rank: " << rank_MPI << " to: " << rank_m1_m1_m1
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (0 -1 -1)

  int num_ghost_0_m1_m1 = 0;

  int rank_0_m1_m1 = ranks_list_1D(1);
  if (rank_0_m1_m1 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ll;
    buffer_coords[1] = Y_ll;
    buffer_coords[2] = Z_ll;

    buffer_coords[3] = X_ur;
    buffer_coords[4] = Y_ll + buffer_width;
    buffer_coords[5] = Z_ll + buffer_width;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_0_m1_m1;
        num_ghost_0_m1_m1++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_0_m1_m1
              << " from rank: " << rank_MPI << " to: " << rank_0_m1_m1
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (1 -1 -1)

  int num_ghost_p1_m1_m1 = 0;

  int rank_p1_m1_m1 = ranks_list_1D(2);
  if (rank_p1_m1_m1 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ur - buffer_width;
    buffer_coords[1] = Y_ll;
    buffer_coords[2] = Z_ll;

    buffer_coords[3] = X_ur;
    buffer_coords[4] = Y_ll + buffer_width;
    buffer_coords[5] = Z_ll + buffer_width;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_p1_m1_m1;
        num_ghost_p1_m1_m1++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_p1_m1_m1
              << " from rank: " << rank_MPI << " to: " << rank_p1_m1_m1
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (-1 0 -1)

  int num_ghost_m1_0_m1 = 0;

  int rank_m1_0_m1 = ranks_list_1D(3);
  if (rank_m1_0_m1 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ll;
    buffer_coords[1] = Y_ll;
    buffer_coords[2] = Z_ll;

    buffer_coords[3] = X_ll + buffer_width;
    buffer_coords[4] = Y_ur;
    buffer_coords[5] = Z_ll + buffer_width;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_m1_0_m1;
        num_ghost_m1_0_m1++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_m1_0_m1
              << " from rank: " << rank_MPI << " to: " << rank_m1_0_m1
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (0 0 -1)
  int num_ghost_0_0_m1 = 0;
  int rank_0_0_m1 = ranks_list_1D(4);
  if (rank_0_0_m1 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ll;
    buffer_coords[1] = Y_ll;
    buffer_coords[2] = Z_ll;

    buffer_coords[3] = X_ur;
    buffer_coords[4] = Y_ur;
    buffer_coords[5] = Z_ll + buffer_width;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_0_0_m1;
        num_ghost_0_0_m1++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_0_0_m1
              << " from rank: " << rank_MPI << " to: " << rank_0_0_m1
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (1 0 -1)
  int num_ghost_p1_0_m1 = 0;
  int rank_p1_0_m1 = ranks_list_1D(5);
  if (rank_p1_0_m1 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ur - buffer_width;
    buffer_coords[1] = Y_ll;
    buffer_coords[2] = Z_ll;

    buffer_coords[3] = X_ur;
    buffer_coords[4] = Y_ur;
    buffer_coords[5] = Z_ll + buffer_width;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_p1_0_m1;
        num_ghost_p1_0_m1++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_p1_0_m1
              << " from rank: " << rank_MPI << " to: " << rank_p1_0_m1
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (-1 1 -1)
  int num_ghost_m1_p1_m1 = 0;
  int rank_m1_p1_m1 = ranks_list_1D(6);
  if (rank_m1_p1_m1 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ll;
    buffer_coords[1] = Y_ur - buffer_width;
    buffer_coords[2] = Z_ll;

    buffer_coords[3] = X_ll + buffer_width;
    buffer_coords[4] = Y_ur;
    buffer_coords[5] = Z_ll + buffer_width;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_m1_p1_m1;
        num_ghost_m1_p1_m1++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_m1_p1_m1
              << " from rank: " << rank_MPI << " to: " << rank_m1_p1_m1
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (0 1 -1)
  int num_ghost_0_p1_m1 = 0;
  int rank_0_p1_m1 = ranks_list_1D(7);
  if (rank_0_p1_m1 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ll;
    buffer_coords[1] = Y_ur - buffer_width;
    buffer_coords[2] = Z_ll;

    buffer_coords[3] = X_ur;
    buffer_coords[4] = Y_ur;
    buffer_coords[5] = Z_ll + buffer_width;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_0_p1_m1;
        num_ghost_0_p1_m1++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_0_p1_m1
              << " from rank: " << rank_MPI << " to: " << rank_0_p1_m1
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (1 1 -1)
  int num_ghost_p1_p1_m1 = 0;
  int rank_p1_p1_m1 = ranks_list_1D(8);
  if (rank_p1_p1_m1 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ur - buffer_width;
    buffer_coords[1] = Y_ur - buffer_width;
    buffer_coords[2] = Z_ll;

    buffer_coords[3] = X_ur;
    buffer_coords[4] = Y_ur;
    buffer_coords[5] = Z_ll + buffer_width;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_p1_p1_m1;
        num_ghost_p1_p1_m1++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_p1_p1_m1
              << " from rank: " << rank_MPI << " to: " << rank_p1_p1_m1
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (-1 -1 0)
  int num_ghost_m1_m1_0 = 0;
  int rank_m1_m1_0 = ranks_list_1D(9);
  if (rank_m1_m1_0 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ll;
    buffer_coords[1] = Y_ll;
    buffer_coords[2] = Z_ll;

    buffer_coords[3] = X_ll + buffer_width;
    buffer_coords[4] = Y_ll + buffer_width;
    buffer_coords[5] = Z_ur;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_m1_m1_0;
        num_ghost_m1_m1_0++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_m1_m1_0
              << " from rank: " << rank_MPI << " to: " << rank_m1_m1_0
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (0 -1 0)
  int num_ghost_0_m1_0 = 0;
  int rank_0_m1_0 = ranks_list_1D(10);
  if (rank_0_m1_0 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ll;
    buffer_coords[1] = Y_ll;
    buffer_coords[2] = Z_ll;

    buffer_coords[3] = X_ur;
    buffer_coords[4] = Y_ll + buffer_width;
    buffer_coords[5] = Z_ur;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_0_m1_0;
        num_ghost_0_m1_0++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_0_m1_0
              << " from rank: " << rank_MPI << " to: " << rank_0_m1_0
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (1 -1 0)
  int num_ghost_p1_m1_0 = 0;
  int rank_p1_m1_0 = ranks_list_1D(11);
  if (rank_p1_m1_0 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ur - buffer_width;
    buffer_coords[1] = Y_ll;
    buffer_coords[2] = Z_ll;

    buffer_coords[3] = X_ur;
    buffer_coords[4] = Y_ll + buffer_width;
    buffer_coords[5] = Z_ur;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_p1_m1_0;
        num_ghost_p1_m1_0++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_p1_m1_0
              << " from rank: " << rank_MPI << " to: " << rank_p1_m1_0
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (-1 0 0)
  int num_ghost_m1_0_0 = 0;
  int rank_m1_0_0 = ranks_list_1D(12);
  if (rank_m1_0_0 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ll;
    buffer_coords[1] = Y_ll;
    buffer_coords[2] = Z_ll;

    buffer_coords[3] = X_ll + buffer_width;
    buffer_coords[4] = Y_ur;
    buffer_coords[5] = Z_ur;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_m1_0_0;
        num_ghost_m1_0_0++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_m1_0_0
              << " from rank: " << rank_MPI << " to: " << rank_m1_0_0
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (0 0 0) <- current rank (we skip it)

  // Direction (1 0 0)
  int num_ghost_p1_0_0 = 0;
  int rank_p1_0_0 = ranks_list_1D(14);
  if (rank_p1_0_0 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ur - buffer_width;
    buffer_coords[1] = Y_ll;
    buffer_coords[2] = Z_ll;

    buffer_coords[3] = X_ur;
    buffer_coords[4] = Y_ur;
    buffer_coords[5] = Z_ur;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_p1_0_0;
        num_ghost_p1_0_0++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_p1_0_0
              << " from rank: " << rank_MPI << " to: " << rank_p1_0_0
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (-1 1 0)
  int num_ghost_m1_p1_0 = 0;
  int rank_m1_p1_0 = ranks_list_1D(15);
  if (rank_m1_p1_0 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ll;
    buffer_coords[1] = Y_ur - buffer_width;
    buffer_coords[2] = Z_ll;

    buffer_coords[3] = X_ll + buffer_width;
    buffer_coords[4] = Y_ur;
    buffer_coords[5] = Z_ur;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_m1_p1_0;
        num_ghost_m1_p1_0++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_m1_p1_0
              << " from rank: " << rank_MPI << " to: " << rank_m1_p1_0
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (0 1 0)
  int num_ghost_0_p1_0 = 0;
  int rank_0_p1_0 = ranks_list_1D(16);
  if (rank_0_p1_0 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ll;
    buffer_coords[1] = Y_ur - buffer_width;
    buffer_coords[2] = Z_ll;

    buffer_coords[3] = X_ur;
    buffer_coords[4] = Y_ur;
    buffer_coords[5] = Z_ur;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_0_p1_0;
        num_ghost_0_p1_0++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_0_p1_0
              << " from rank: " << rank_MPI << " to: " << rank_0_p1_0
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (1 1 0)
  int num_ghost_p1_p1_0 = 0;
  int rank_p1_p1_0 = ranks_list_1D(17);
  if (rank_p1_p1_0 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ur - buffer_width;
    buffer_coords[1] = Y_ur - buffer_width;
    buffer_coords[2] = Z_ll;

    buffer_coords[3] = X_ur;
    buffer_coords[4] = Y_ur;
    buffer_coords[5] = Z_ur;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_p1_p1_0;
        num_ghost_p1_p1_0++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_p1_p1_0
              << " from rank: " << rank_MPI << " to: " << rank_p1_p1_0
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (-1 -1 1)
  int num_ghost_m1_m1_p1 = 0;
  int rank_m1_m1_p1 = ranks_list_1D(18);
  if (rank_m1_m1_p1 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ll;
    buffer_coords[1] = Y_ll;
    buffer_coords[2] = Z_ur - buffer_width;

    buffer_coords[3] = X_ll + buffer_width;
    buffer_coords[4] = Y_ll + buffer_width;
    buffer_coords[5] = Z_ur;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_m1_m1_p1;
        num_ghost_m1_m1_p1++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_m1_m1_p1
              << " from rank: " << rank_MPI << " to: " << rank_m1_m1_p1
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (0 -1 1)
  int num_ghost_0_m1_p1 = 0;
  int rank_0_m1_p1 = ranks_list_1D(19);
  if (rank_0_m1_p1 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ll;
    buffer_coords[1] = Y_ll;
    buffer_coords[2] = Z_ur - buffer_width;

    buffer_coords[3] = X_ur;
    buffer_coords[4] = Y_ll + buffer_width;
    buffer_coords[5] = Z_ur;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_0_m1_p1;
        num_ghost_0_m1_p1++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_0_m1_p1
              << " from rank: " << rank_MPI << " to: " << rank_0_m1_p1
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (1 -1 1)
  int num_ghost_p1_m1_p1 = 0;
  int rank_p1_m1_p1 = ranks_list_1D(20);
  if (rank_p1_m1_p1 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ur - buffer_width;
    buffer_coords[1] = Y_ll;
    buffer_coords[2] = Z_ur - buffer_width;

    buffer_coords[3] = X_ur;
    buffer_coords[4] = Y_ll + buffer_width;
    buffer_coords[5] = Z_ur;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_p1_m1_p1;
        num_ghost_p1_m1_p1++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_p1_m1_p1
              << " from rank: " << rank_MPI << " to: " << rank_p1_m1_p1
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (-1 0 1)
  int num_ghost_m1_0_p1 = 0;
  int rank_m1_0_p1 = ranks_list_1D(21);
  if (rank_m1_0_p1 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ll;
    buffer_coords[1] = Y_ll;
    buffer_coords[2] = Z_ur - buffer_width;

    buffer_coords[3] = X_ll + buffer_width;
    buffer_coords[4] = Y_ur;
    buffer_coords[5] = Z_ur;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_m1_0_p1;
        num_ghost_m1_0_p1++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_m1_0_p1
              << " from rank: " << rank_MPI << " to: " << rank_m1_0_p1
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (0 0 1)
  int num_ghost_0_0_p1 = 0;
  int rank_0_0_p1 = ranks_list_1D(22);
  if (rank_0_0_p1 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ll;
    buffer_coords[1] = Y_ll;
    buffer_coords[2] = Z_ur - buffer_width;

    buffer_coords[3] = X_ur;
    buffer_coords[4] = Y_ur;
    buffer_coords[5] = Z_ur;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_0_0_p1;
        num_ghost_0_0_p1++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_0_0_p1
              << " from rank: " << rank_MPI << " to: " << rank_0_0_p1
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (1 0 1)
  int num_ghost_p1_0_p1 = 0;
  int rank_p1_0_p1 = ranks_list_1D(23);
  if (rank_p1_0_p1 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ur - buffer_width;
    buffer_coords[1] = Y_ll;
    buffer_coords[2] = Z_ur - buffer_width;

    buffer_coords[3] = X_ur;
    buffer_coords[4] = Y_ur;
    buffer_coords[5] = Z_ur;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_p1_0_p1;
        num_ghost_p1_0_p1++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_p1_0_p1
              << " from rank: " << rank_MPI << " to: " << rank_p1_0_p1
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (-1 1 1)
  int num_ghost_m1_p1_p1 = 0;
  int rank_m1_p1_p1 = ranks_list_1D(24);
  if (rank_m1_p1_p1 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ll;
    buffer_coords[1] = Y_ur - buffer_width;
    buffer_coords[2] = Z_ur - buffer_width;

    buffer_coords[3] = X_ll + buffer_width;
    buffer_coords[4] = Y_ur;
    buffer_coords[5] = Z_ur;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_m1_p1_p1;
        num_ghost_m1_p1_p1++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_m1_p1_p1
              << " from rank: " << rank_MPI << " to: " << rank_m1_p1_p1
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (0 1 1)
  int num_ghost_0_p1_p1 = 0;
  int rank_0_p1_p1 = ranks_list_1D(25);
  if (rank_0_p1_p1 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ll;
    buffer_coords[1] = Y_ur - buffer_width;
    buffer_coords[2] = Z_ur - buffer_width;

    buffer_coords[3] = X_ur;
    buffer_coords[4] = Y_ur;
    buffer_coords[5] = Z_ur;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_0_p1_p1;
        num_ghost_0_p1_p1++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_0_p1_p1
              << " from rank: " << rank_MPI << " to: " << rank_0_p1_p1
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  // Direction (1 1 1)
  int num_ghost_p1_p1_p1 = 0;
  int rank_p1_p1_p1 = ranks_list_1D(26);
  if (rank_p1_p1_p1 >= 0) {

    PetscScalar buffer_coords[6];
    buffer_coords[0] = X_ur - buffer_width;
    buffer_coords[1] = Y_ur - buffer_width;
    buffer_coords[2] = Z_ur - buffer_width;

    buffer_coords[3] = X_ur;
    buffer_coords[4] = Y_ur;
    buffer_coords[5] = Z_ur;

    PetscInt* swarm_rank_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank,
                              NULL, NULL, (void**)&swarm_rank_ptr));

    PetscScalar* mean_q_ptr;
    PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmPICField_coor,
                              NULL, NULL, (void**)&mean_q_ptr));
    Eigen::Map<MatrixType> mean_q(mean_q_ptr, n_local_size, dim);

    for (int site_i = 0; site_i < n_local_size; site_i++) {

      Eigen::Vector3d mean_q_i = mean_q.row(site_i);

      if (In_Out_Buffer(mean_q_i, buffer_coords)) {  //!
        swarm_rank_ptr[site_i] = rank_p1_p1_p1;
        num_ghost_p1_p1_p1++;
      }
    }

    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data,
                                  DMSwarmPICField_coor, NULL, NULL,
                                  (void**)&mean_q_ptr));
    PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                  NULL, NULL, (void**)&swarm_rank_ptr));

    std::cout << "Number of ghost atoms: " << num_ghost_p1_p1_p1
              << " from rank: " << rank_MPI << " to: " << rank_p1_p1_p1
              << std::endl;
  }

  PetscCall(DMSwarmMigrate(Simulation->atomistic_data, PETSC_FALSE));

  PetscFunctionReturn(PETSC_SUCCESS);
}

/********************************************************************************/

PetscErrorCode destroy_ghost_atoms(DMD* Simulation) {

  PetscFunctionBeginUser;

  //! Get global size
  PetscInt n_global_size = 0;
  PetscCall(DMSwarmGetSize(Simulation->atomistic_data, &n_global_size));

  //! Get local size
  PetscInt n_local_size = 0;
  PetscCall(DMSwarmGetLocalSize(Simulation->atomistic_data, &n_local_size));

  //!
  PetscInt num_remove_point = 0;

  PetscInt* swarm_rank_ptr;
  PetscCall(DMSwarmGetField(Simulation->atomistic_data, DMSwarmField_rank, NULL,
                            NULL, (void**)&swarm_rank_ptr));

  PetscInt* rank_ptr;
  PetscCall(DMSwarmGetField(Simulation->atomistic_data, "MPI-rank", NULL, NULL,
                            (void**)&rank_ptr));

  for (int site_i = 0; site_i < n_local_size; site_i++) {

    if ((swarm_rank_ptr[site_i] != rank_ptr[site_i]) &&
        (swarm_rank_ptr[site_i] == rank_MPI)) {
      num_remove_point += 1;
    }

    if ((swarm_rank_ptr[site_i] != rank_ptr[site_i]) &&
        (rank_ptr[site_i] != rank_MPI)) {
      swarm_rank_ptr[site_i] = rank_ptr[site_i];
    }
  }

  PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, DMSwarmField_rank,
                                NULL, NULL, (void**)&swarm_rank_ptr));
  PetscCall(DMSwarmRestoreField(Simulation->atomistic_data, "MPI-rank", NULL,
                                NULL, (void**)&rank_ptr));

  for (int site_i = 0; site_i < num_remove_point; site_i++) {
    PetscCall(DMSwarmRemovePoint(Simulation->atomistic_data));
  }

  std::cout << "Number of removed ghost atoms: " << num_remove_point
            << " at rank: " << rank_MPI << std::endl;

  PetscFunctionReturn(PETSC_SUCCESS);
}

/********************************************************************************/

static PetscErrorCode GetElementCoords(const PetscScalar _coords[],
                                       const PetscInt e2n[],
                                       PetscScalar el_coords[]) {
  PetscInt dim = NumberDimensions;
  PetscInt n_nodes = 8;

  PetscFunctionBeginUser;
  /* get coords for the element */
  for (PetscInt i = 0; i < n_nodes; i++) {
    for (PetscInt d = 0; d < dim; d++)
      el_coords[dim * i + d] = _coords[dim * e2n[i] + d];
  }
  PetscFunctionReturn(PETSC_SUCCESS);
}

/********************************************************************************/

static bool In_Out_Buffer(const Eigen::Vector3d mean_q,
                          const PetscScalar el_coords[]) {

  unsigned int dim = NumberDimensions;

  //! Lower left coordinates of the brick
  PetscReal xI_lw, yI_lw, zI_lw;
  xI_lw = el_coords[0];
  yI_lw = el_coords[1];
  zI_lw = el_coords[2];

  //! Upper right coordinates of the brick
  PetscReal xI_up, yI_up, zI_up;
  xI_up = el_coords[3];
  yI_up = el_coords[4];
  zI_up = el_coords[5];

  if ((mean_q(0) > xI_lw) &&  //!
      (mean_q(0) < xI_up) &&   //!
      (mean_q(1) > yI_lw) &&  //!
      (mean_q(1) < yI_up) &&   //!
      (mean_q(2) > zI_lw) &&  //!
      (mean_q(2) < zI_up)) {
    return true;
  } else {
    return false;
  }
}

/********************************************************************************/