MODULE domain_mod
   USE cellset_mod, ONLY: t_cellset
   USE domain_param
   IMPLICIT NONE
   SAVE

   TYPE t_domain
      INTEGER :: face
      INTEGER :: iim
      INTEGER :: jjm
      INTEGER :: ii_begin
      INTEGER :: jj_begin
      INTEGER :: ii_end
      INTEGER :: jj_end
      INTEGER :: ii_nb
      INTEGER :: jj_nb
      INTEGER :: ii_begin_glo
      INTEGER :: jj_begin_glo
      INTEGER :: ii_end_glo
      INTEGER :: jj_end_glo
      INTEGER, POINTER  :: assign_domain(:, :)
      INTEGER, POINTER  :: assign_cell_glo(:, :)
      INTEGER, POINTER  :: assign_i(:, :)
      INTEGER, POINTER  :: assign_j(:, :)
      INTEGER, POINTER  :: edge_assign_domain(:, :, :)
      INTEGER, POINTER  :: edge_assign_i(:, :, :)
      INTEGER, POINTER  :: edge_assign_j(:, :, :)
      INTEGER, POINTER  :: edge_assign_pos(:, :, :)
      INTEGER, POINTER  :: edge_assign_sign(:, :, :)
      INTEGER, POINTER  :: vertex_assign_domain(:, :, :)
      INTEGER, POINTER  :: vertex_assign_i(:, :, :)
      INTEGER, POINTER  :: vertex_assign_j(:, :, :)
      INTEGER, POINTER  :: vertex_assign_pos(:, :, :)
      REAL, POINTER     :: xyz(:, :, :)
      REAL, POINTER     :: neighbour(:, :, :, :)
      INTEGER, POINTER  :: delta(:, :)
      REAL, POINTER     :: vertex(:, :, :, :)
      LOGICAL, POINTER  :: own(:, :)
      LOGICAL, POINTER  :: outside(:, :)
      INTEGER, POINTER  :: ne(:, :, :)

      INTEGER :: t_right
      INTEGER :: t_rup
      INTEGER :: t_lup
      INTEGER :: t_left
      INTEGER :: t_ldown
      INTEGER :: t_rdown

      INTEGER :: u_right
      INTEGER :: u_rup
      INTEGER :: u_lup
      INTEGER :: u_left
      INTEGER :: u_ldown
      INTEGER :: u_rdown

      INTEGER :: z_rup
      INTEGER :: z_up
      INTEGER :: z_lup
      INTEGER :: z_ldown
      INTEGER :: z_down
      INTEGER :: z_rdown

      INTEGER :: t_pos(6)
      INTEGER :: u_pos(6)
      INTEGER :: z_pos(6)

      TYPE(t_cellset), POINTER :: primal_own, dual_own, edge_own, &
                                  primal_all, dual_all, edge_all ! halo included, used for debug by writefield
   END TYPE t_domain

   TYPE t_mesh
      INTEGER :: ndomain, primal_num, edge_num, dual_num
      TYPE(t_domain), ALLOCATABLE :: domain(:)
      TYPE(t_cellset), ALLOCATABLE :: primal_own(:), dual_own(:), edge_own(:), &
                                      primal_all(:), dual_all(:), edge_all(:)
   END type t_mesh

   INTEGER :: ndomain
   INTEGER :: ndomain_glo

   TYPE(t_mesh), TARGET :: mesh_loc, mesh_glo
   TYPE(t_domain), POINTER :: domain(:), domain_glo(:)

   LOGICAL :: swap_needed = .TRUE. ! .FALSE. if a thread always computes on the same domain
   !$OMP THREADPRIVATE(swap_needed)

   INTEGER, ALLOCATABLE :: domglo_rank(:)  ! size : ndomain_glo : mpi rank assigned to a domain
   INTEGER, ALLOCATABLE :: domglo_loc_ind(:) ! size : ndomain_glo : corresponding local indice for a global domain indice
   INTEGER, ALLOCATABLE :: domloc_glo_ind(:) ! size : ndomain : corresponding global indice for a local domain indice

   LOGICAL, ALLOCATABLE :: assigned_domain(:) ! size : ndomain : is an omp task is assigned to this domain
!$OMP THREADPRIVATE(assigned_domain)

CONTAINS

   SUBROUTINE allocate_mesh(ndom, mesh)
      INTEGER, INTENT(IN)       :: ndom
      TYPE(t_mesh), TARGET, INTENT(OUT) :: mesh
      INTEGER :: ind
      mesh%ndomain = ndom
      ALLOCATE (mesh%domain(ndom))
      ALLOCATE (mesh%primal_own(ndom))
      ALLOCATE (mesh%primal_all(ndom))
      ALLOCATE (mesh%dual_own(ndom))
      ALLOCATE (mesh%dual_all(ndom))
      ALLOCATE (mesh%edge_own(ndom))
      ALLOCATE (mesh%edge_all(ndom))
      DO ind = 1, ndom
         mesh%domain(ind)%primal_own => mesh%primal_own(ind)
         mesh%domain(ind)%primal_all => mesh%primal_all(ind)
         mesh%domain(ind)%dual_own => mesh%dual_own(ind)
         mesh%domain(ind)%dual_all => mesh%dual_all(ind)
         mesh%domain(ind)%edge_own => mesh%edge_own(ind)
         mesh%domain(ind)%edge_all => mesh%edge_all(ind)
      END DO
   END SUBROUTINE allocate_mesh

   SUBROUTINE create_domain
      USE grid_param
      USE metric_var_mod, ONLY: nb_face
      USE mpipara
      USE ioipsl
      INTEGER :: ind, nf, ni, nj, i, j
      INTEGER :: quotient, rest
      INTEGER :: halo_i, halo_j

      TYPE(t_domain), POINTER :: d

      ndomain_glo = nsplit_i*nsplit_j*nb_face
      CALL allocate_mesh(ndomain_glo, mesh_glo)
      domain_glo => mesh_glo%domain
      ALLOCATE (domglo_rank(ndomain_glo))
      ALLOCATE (domglo_loc_ind(ndomain_glo))

      halo_i = 0
      CALL getin("halo_i", halo_i)
      halo_j = 1
      CALL getin("halo_j", halo_j)

      ind = 0
      DO nf = 1, nb_face
         DO nj = 1, nsplit_j
            DO ni = 1, nsplit_i
               ind = ind + 1
               d => domain_glo(ind)
               quotient = (iim_glo/nsplit_i)
               rest = MOD(iim_glo, nsplit_i)
               IF (ni - 1 < rest) THEN
                  d%ii_nb = quotient + 1
                  d%ii_begin_glo = (quotient + 1)*(ni - 1) + 1
               ELSE
                  d%ii_nb = quotient
                  d%ii_begin_glo = (quotient + 1)*rest + (ni - 1 - rest)*quotient + 1
               END IF
               d%ii_end_glo = d%ii_begin_glo + d%ii_nb - 1

               IF (ni /= 1) THEN
                  d%ii_nb = d%ii_nb + 1
                  d%ii_begin_glo = d%ii_begin_glo - 1
               END IF

               quotient = (jjm_glo/nsplit_j)
               rest = MOD(jjm_glo, nsplit_j)
               IF (nj - 1 < rest) THEN
                  d%jj_nb = quotient + 1
                  d%jj_begin_glo = (quotient + 1)*(nj - 1) + 1
               ELSE
                  d%jj_nb = quotient
                  d%jj_begin_glo = (quotient + 1)*rest + (nj - 1 - rest)*quotient + 1
               END IF
               d%jj_end_glo = d%jj_begin_glo + d%jj_nb - 1

               IF (nj /= 1) THEN
                  d%jj_nb = d%jj_nb + 1
                  d%jj_begin_glo = d%jj_begin_glo - 1
               END IF

               d%iim = d%ii_nb + 2*halo + halo_i*2
               d%jjm = d%jj_nb + 2*halo + halo_j*2
               d%ii_begin = halo + 1 + halo_i
               d%jj_begin = halo + 1 + halo_j
               d%ii_end = d%ii_begin + d%ii_nb - 1
               d%jj_end = d%jj_begin + d%jj_nb - 1
               d%face = nf
               ALLOCATE (d%assign_domain(d%iim, d%jjm))
               ALLOCATE (d%assign_cell_glo(d%iim, d%jjm))
               ALLOCATE (d%assign_i(d%iim, d%jjm))
               ALLOCATE (d%assign_j(d%iim, d%jjm))
               ALLOCATE (d%edge_assign_domain(0:5, d%iim, d%jjm))
               ALLOCATE (d%edge_assign_i(0:5, d%iim, d%jjm))
               ALLOCATE (d%edge_assign_j(0:5, d%iim, d%jjm))
               ALLOCATE (d%edge_assign_pos(0:5, d%iim, d%jjm))
               ALLOCATE (d%edge_assign_sign(0:5, d%iim, d%jjm))
               ALLOCATE (d%vertex_assign_domain(0:5, d%iim, d%jjm))
               ALLOCATE (d%vertex_assign_i(0:5, d%iim, d%jjm))
               ALLOCATE (d%vertex_assign_j(0:5, d%iim, d%jjm))
               ALLOCATE (d%vertex_assign_pos(0:5, d%iim, d%jjm))
               ALLOCATE (d%delta(d%iim, d%jjm))
               ALLOCATE (d%xyz(3, d%iim, d%jjm))
               ALLOCATE (d%vertex(3, 0:5, d%iim, d%jjm))
               ALLOCATE (d%neighbour(3, 0:5, d%iim, d%jjm))
               ALLOCATE (d%own(d%iim, d%jjm))
               ALLOCATE (d%outside(d%iim, d%jjm))
               d%outside(:, :) = .TRUE.
               ALLOCATE (d%ne(0:5, d%iim, d%jjm))

               IF (is_mpi_root) PRINT *, "Domain ", ind, " : size ", d%ii_nb, "x", d%jj_nb

            END DO
         END DO
      END DO

   END SUBROUTINE create_domain

   SUBROUTINE copy_domain(d1, d2)
      TYPE(t_domain), TARGET, INTENT(IN) :: d1
      TYPE(t_domain), INTENT(OUT) :: d2

      d2%iim = d1%iim
      d2%jjm = d1%jjm
      d2%ii_begin = d1%ii_begin
      d2%jj_begin = d1%jj_begin
      d2%ii_end = d1%ii_end
      d2%jj_end = d1%jj_end
      d2%ii_nb = d1%ii_nb
      d2%jj_nb = d1%jj_nb
      d2%ii_begin_glo = d1%ii_begin_glo
      d2%jj_begin_glo = d1%jj_begin_glo
      d2%ii_end_glo = d1%ii_end_glo
      d2%jj_end_glo = d1%jj_end_glo
      d2%assign_domain => d1%assign_domain
      d2%assign_cell_glo => d1%assign_cell_glo
      d2%assign_i => d1%assign_i
      d2%assign_j => d1%assign_j
      d2%edge_assign_domain => d1%edge_assign_domain
      d2%edge_assign_i => d1%edge_assign_i
      d2%edge_assign_j => d1%edge_assign_j
      d2%edge_assign_pos => d1%edge_assign_pos
      d2%edge_assign_sign => d1%edge_assign_sign
      d2%vertex_assign_domain => d1%vertex_assign_domain
      d2%vertex_assign_i => d1%vertex_assign_i
      d2%vertex_assign_j => d1%vertex_assign_j
      d2%vertex_assign_pos => d1%vertex_assign_pos
      d2%xyz => d1%xyz
      d2%neighbour => d1%neighbour
      d2%delta => d1%delta
      d2%vertex => d1%vertex
      d2%own => d1%own
      d2%outside => d1%outside
      d2%ne => d1%ne

      d2%t_right = d1%t_right
      d2%t_rup = d1%t_rup
      d2%t_lup = d1%t_lup
      d2%t_left = d1%t_left
      d2%t_ldown = d1%t_ldown
      d2%t_rdown = d1%t_rdown

      d2%u_right = d1%u_right
      d2%u_rup = d1%u_rup
      d2%u_lup = d1%u_lup
      d2%u_left = d1%u_left
      d2%u_ldown = d1%u_ldown
      d2%u_rdown = d1%u_rdown

      d2%z_rup = d1%z_rup
      d2%z_up = d1%z_up
      d2%z_lup = d1%z_lup
      d2%z_ldown = d1%z_ldown
      d2%z_down = d1%z_down
      d2%z_rdown = d1%z_rdown

      d2%t_pos = d1%t_pos
      d2%u_pos = d1%u_pos
      d2%z_pos = d1%z_pos

   END SUBROUTINE copy_domain

   SUBROUTINE assign_cell
      USE metric
      USE mpipara, ONLY : is_mpi_master
      INTEGER :: ind_d, ind, ind2, e, v
      INTEGER :: nf, nf2
      INTEGER :: i, j, k, ii, jj
      TYPE(t_domain), POINTER :: d
      INTEGER :: delta
      INTEGER :: nbp, nn, M
      REAL :: rn
      INTEGER, ALLOCATABLE :: nb_own(:)
      INTEGER :: mean_own, side

      ALLOCATE (nb_own(0:ndomain_glo))
      nb_own(:) = 0

      cell_glo(:)%assign_domain = -1
      edge_glo(:)%assign_domain = -1
      vertices_glo(:)%assign_domain = -1

      DO ind_d = 1, ndomain_glo
         d => domain_glo(ind_d)
         nf = d%face
         DO j = d%jj_begin, d%jj_end
            DO i = d%ii_begin, d%ii_end
               ii = d%ii_begin_glo - d%ii_begin + i
               jj = d%jj_begin_glo - d%jj_begin + j
               ind = vertex_glo(ii, jj, nf)%ind
               cell_glo(ind)%assign_domain = 0
            END DO
         END DO
      END DO

      DO ind_d = 1, ndomain_glo
         d => domain_glo(ind_d)
         nf = d%face
         DO j = d%jj_begin - 1, d%jj_end + 1
            DO i = d%ii_begin - 1, d%ii_end + 1
               ii = d%ii_begin_glo - d%ii_begin + i
               jj = d%jj_begin_glo - d%jj_begin + j
               ind = vertex_glo(ii, jj, nf)%ind
               IF (cell_glo(ind)%assign_domain == 0) d%outside(i, j) = .FALSE.
            END DO
         END DO
      END DO

! assign edge and vertex in the inner domain
      nbp = 0
      DO ind_d = 1, ndomain_glo
         d => domain_glo(ind_d)
         nf = d%face
         DO j = d%jj_begin, d%jj_end
            DO i = d%ii_begin, d%ii_end
               ii = d%ii_begin_glo - d%ii_begin + i
               jj = d%jj_begin_glo - d%jj_begin + j
               ind = vertex_glo(ii, jj, nf)%ind
               delta = vertex_glo(ii, jj, nf)%delta
               IF (cell_glo(ind)%assign_face == nf) THEN
                  nb_own(cell_glo(ind)%assign_domain) = nb_own(cell_glo(ind)%assign_domain) - 1
                  nb_own(ind_d) = nb_own(ind_d) + 1
                  cell_glo(ind)%assign_domain = ind_d
                  cell_glo(ind)%assign_i = i
                  cell_glo(ind)%assign_j = j
               END IF

               IF (i + 1 <= d%ii_end) CALL assign_edge(ind_d, ind, i, j, delta, 0)
               IF (j + 1 <= d%jj_end) CALL assign_edge(ind_d, ind, i, j, delta, 1)
               IF (i - 1 >= d%ii_begin .AND. j + 1 <= d%jj_end) CALL assign_edge(ind_d, ind, i, j, delta, 2)
               IF (i - 1 >= d%ii_begin) CALL assign_edge(ind_d, ind, i, j, delta, 3)
               IF (j - 1 >= d%jj_begin) CALL assign_edge(ind_d, ind, i, j, delta, 4)
               IF (i + 1 <= d%ii_end .AND. j - 1 >= d%jj_begin) CALL assign_edge(ind_d, ind, i, j, delta, 5)

               IF (i + 1 <= d%ii_end .AND. j + 1 <= d%jj_end) CALL assign_vertex(ind_d, ind, i, j, delta, 0)
               IF (i - 1 >= d%ii_begin .AND. j + 1 <= d%jj_end) CALL assign_vertex(ind_d, ind, i, j, delta, 1)
               IF (i - 1 >= d%ii_begin .AND. j + 1 <= d%jj_end) CALL assign_vertex(ind_d, ind, i, j, delta, 2)
               IF (i - 1 >= d%ii_begin .AND. j - 1 >= d%jj_begin) CALL assign_vertex(ind_d, ind, i, j, delta, 3)
               IF (i + 1 <= d%ii_end .AND. j - 1 >= d%jj_begin) CALL assign_vertex(ind_d, ind, i, j, delta, 4)
               IF (i + 1 <= d%ii_end .AND. j - 1 >= d%jj_begin) CALL assign_vertex(ind_d, ind, i, j, delta, 5)

               IF (j == d%jj_begin .OR. j == d%jj_begin .OR. i == d%ii_begin .OR. i == d%ii_end) nbp = nbp + 1

            END DO
         END DO
      END DO

! dynamic load balancing for owned cells. Get ownership of a cell if local cells number < mean cell number
! use random shot on frontier

      mean_own = sum(nb_own(1:ndomain_glo))/ndomain_glo
      IF (is_mpi_master) PRINT *, "Before balance : Mean own ", sum(nb_own(1:ndomain_glo))/ndomain_glo, &
         "Min own", minval(nb_own(1:ndomain_glo)), "Max own", maxval(nb_own(1:ndomain_glo))
! set random seed to get reproductibility when using a different number of process

      CALL RANDOM_SEED(size=M)
      CALL RANDOM_SEED(put=(/(i, i=1, M)/))

      DO nn = 1, 1000*nbp

         CALL random_number(rn)
         ind_d = INT(ndomain_glo*rn) + 1
         IF (nb_own(ind_d) >= mean_own) CYCLE

         d => domain_glo(ind_d)
         nf = d%face

         CALL random_number(rn)
         side = INT(4*rn)
         IF (side == 0) THEN
            CALL random_number(rn)
            i = INT((d%ii_end - d%ii_begin + 1)*rn) + d%ii_begin
            j = d%jj_begin
         ELSE IF (side == 1) THEN
            CALL random_number(rn)
            i = INT((d%ii_end - d%ii_begin + 1)*rn) + d%ii_begin
            j = d%jj_end
         ELSE IF (side == 2) THEN
            CALL random_number(rn)
            j = INT((d%jj_end - d%jj_begin + 1)*rn) + d%jj_begin
            i = d%ii_begin
         ELSE IF (side == 3) THEN
            CALL random_number(rn)
            j = INT((d%jj_end - d%jj_begin + 1)*rn) + d%jj_begin
            i = d%ii_end
         END IF

         ii = d%ii_begin_glo - d%ii_begin + i
         jj = d%jj_begin_glo - d%jj_begin + j
         ind = vertex_glo(ii, jj, nf)%ind
         delta = vertex_glo(ii, jj, nf)%delta
         IF (cell_glo(ind)%assign_face == nf) THEN
            nb_own(cell_glo(ind)%assign_domain) = nb_own(cell_glo(ind)%assign_domain) - 1
            nb_own(ind_d) = nb_own(ind_d) + 1
            cell_glo(ind)%assign_domain = ind_d
            cell_glo(ind)%assign_i = i
            cell_glo(ind)%assign_j = j
         END IF

      END DO

      IF (is_mpi_master) PRINT *, "After balance :Mean own ", sum(nb_own(1:ndomain_glo))/ndomain_glo, &
         "Min own", minval(nb_own(1:ndomain_glo)), "Max own", maxval(nb_own(1:ndomain_glo))

! assign missing cells and edge, vertex for limited area config

! first assign cells and vertex on the first neighbourough
      DO ind_d = 1, ndomain_glo
         d => domain_glo(ind_d)
         nf = d%face
         DO j = d%jj_begin, d%jj_end
            DO i = d%ii_begin, d%ii_end
               ii = d%ii_begin_glo - d%ii_begin + i
               jj = d%jj_begin_glo - d%jj_begin + j
               ind = vertex_glo(ii, jj, nf)%ind
               delta = vertex_glo(ii, jj, nf)%delta
               IF (cell_glo(ind)%assign_face == nf) THEN
                  IF (cell_glo(ind)%assign_domain == -1) THEN
                     cell_glo(ind)%assign_domain = ind_d
                     cell_glo(ind)%assign_i = i
                     cell_glo(ind)%assign_j = j
                  END IF
               END IF

               CALL assign_edge_if_not(ind_d, ind, i, j, delta, 0)
               CALL assign_edge_if_not(ind_d, ind, i, j, delta, 1)
               CALL assign_edge_if_not(ind_d, ind, i, j, delta, 2)
               CALL assign_edge_if_not(ind_d, ind, i, j, delta, 3)
               CALL assign_edge_if_not(ind_d, ind, i, j, delta, 4)
               CALL assign_edge_if_not(ind_d, ind, i, j, delta, 5)

               CALL assign_vertex_if_not(ind_d, ind, i, j, delta, 0)
               CALL assign_vertex_if_not(ind_d, ind, i, j, delta, 1)
               CALL assign_vertex_if_not(ind_d, ind, i, j, delta, 2)
               CALL assign_vertex_if_not(ind_d, ind, i, j, delta, 3)
               CALL assign_vertex_if_not(ind_d, ind, i, j, delta, 4)
               CALL assign_vertex_if_not(ind_d, ind, i, j, delta, 5)

            END DO
         END DO
      END DO

! assign edges and vertex on the second neighbourough
      DO ind_d = 1, ndomain_glo
         d => domain_glo(ind_d)
         nf = d%face
         DO j = d%jj_begin - 1, d%jj_end + 1
            DO i = d%ii_begin - 1, d%ii_end + 1
               ii = d%ii_begin_glo - d%ii_begin + i
               jj = d%jj_begin_glo - d%jj_begin + j
               ind = vertex_glo(ii, jj, nf)%ind
               delta = vertex_glo(ii, jj, nf)%delta
               IF (cell_glo(ind)%assign_face == nf) THEN
                  IF (cell_glo(ind)%assign_domain == -1) THEN
                     cell_glo(ind)%assign_domain = ind_d
                     cell_glo(ind)%assign_i = i
                     cell_glo(ind)%assign_j = j
                  END IF
               END IF

               IF (i + 1 <= d%ii_end + 1) CALL assign_edge_if_not(ind_d, ind, i, j, delta, 0)
               IF (j + 1 <= d%jj_end + 1) CALL assign_edge_if_not(ind_d, ind, i, j, delta, 1)
               IF (i - 1 >= d%ii_begin - 1 .AND. j + 1 <= d%jj_end + 1) CALL assign_edge_if_not(ind_d, ind, i, j, delta, 2)
               IF (i - 1 >= d%ii_begin - 1) CALL assign_edge_if_not(ind_d, ind, i, j, delta, 3)
               IF (j - 1 >= d%jj_begin - 1) CALL assign_edge_if_not(ind_d, ind, i, j, delta, 4)
               IF (i + 1 <= d%ii_end + 1 .AND. j - 1 >= d%jj_begin - 1) CALL assign_edge_if_not(ind_d, ind, i, j, delta, 5)

               IF (i + 1 <= d%ii_end + 1 .AND. j + 1 <= d%jj_end + 1) CALL assign_vertex_if_not(ind_d, ind, i, j, delta, 0)
               IF (i - 1 >= d%ii_begin - 1 .AND. j + 1 <= d%jj_end + 1) CALL assign_vertex_if_not(ind_d, ind, i, j, delta, 1)
               IF (i - 1 >= d%ii_begin - 1 .AND. j + 1 <= d%jj_end + 1) CALL assign_vertex_if_not(ind_d, ind, i, j, delta, 2)
               IF (i - 1 >= d%ii_begin - 1 .AND. j - 1 >= d%jj_begin - 1) CALL assign_vertex_if_not(ind_d, ind, i, j, delta, 3)
               IF (i + 1 <= d%ii_end + 1 .AND. j - 1 >= d%jj_begin - 1) CALL assign_vertex_if_not(ind_d, ind, i, j, delta, 4)
               IF (i + 1 <= d%ii_end + 1 .AND. j - 1 >= d%jj_begin - 1) CALL assign_vertex_if_not(ind_d, ind, i, j, delta, 5)

            END DO
         END DO
      END DO

! assign edges and vertex on the third neighbourough
      DO ind_d = 1, ndomain_glo
         d => domain_glo(ind_d)
         nf = d%face
         DO j = d%jj_begin - 1, d%jj_end + 1
            DO i = d%ii_begin - 1, d%ii_end + 1
               ii = d%ii_begin_glo - d%ii_begin + i
               jj = d%jj_begin_glo - d%jj_begin + j
               ind = vertex_glo(ii, jj, nf)%ind
               delta = vertex_glo(ii, jj, nf)%delta
               IF (cell_glo(ind)%assign_face == nf) THEN
                  IF (cell_glo(ind)%assign_domain == -1) THEN
                     cell_glo(ind)%assign_domain = ind_d
                     cell_glo(ind)%assign_i = i
                     cell_glo(ind)%assign_j = j
                  END IF
               END IF

               CALL assign_edge_if_not(ind_d, ind, i, j, delta, 0)
               CALL assign_edge_if_not(ind_d, ind, i, j, delta, 1)
               CALL assign_edge_if_not(ind_d, ind, i, j, delta, 2)
               CALL assign_edge_if_not(ind_d, ind, i, j, delta, 3)
               CALL assign_edge_if_not(ind_d, ind, i, j, delta, 4)
               CALL assign_edge_if_not(ind_d, ind, i, j, delta, 5)

               CALL assign_vertex_if_not(ind_d, ind, i, j, delta, 0)
               CALL assign_vertex_if_not(ind_d, ind, i, j, delta, 1)
               CALL assign_vertex_if_not(ind_d, ind, i, j, delta, 2)
               CALL assign_vertex_if_not(ind_d, ind, i, j, delta, 3)
               CALL assign_vertex_if_not(ind_d, ind, i, j, delta, 4)
               CALL assign_vertex_if_not(ind_d, ind, i, j, delta, 5)

            END DO
         END DO
      END DO

      DO ind_d = 1, ndomain_glo
         d => domain_glo(ind_d)
         nf = d%face
         DO j = d%jj_begin - 1, d%jj_end + 1
            DO i = d%ii_begin - 1, d%ii_end + 1
               ii = d%ii_begin_glo - d%ii_begin + i
               jj = d%jj_begin_glo - d%jj_begin + j
               ind = vertex_glo(ii, jj, nf)%ind
               d%assign_cell_glo(i, j) = ind
               d%assign_domain(i, j) = cell_glo(ind)%assign_domain
               d%assign_i(i, j) = cell_glo(ind)%assign_i
               d%assign_j(i, j) = cell_glo(ind)%assign_j
               delta = vertex_glo(ii, jj, nf)%delta
               d%delta(i, j) = vertex_glo(ii, jj, nf)%delta
               DO k = 0, 5
                  ind2 = vertex_glo(ii, jj, nf)%neighbour(k)
                  IF (ind2 == -1) THEN
                     d%edge_assign_domain(k, i, j) = -1
                  ELSE
                     d%neighbour(:, k, i, j) = cell_glo(ind2)%xyz(:)

                     d%ne(k, i, j) = 1 - 2*MOD(k, 2)

                     e = cell_glo(ind)%edge(MOD(k + delta + 6, 6))
                     d%edge_assign_domain(k, i, j) = edge_glo(e)%assign_domain
                     d%edge_assign_i(k, i, j) = edge_glo(e)%assign_i
                     d%edge_assign_j(k, i, j) = edge_glo(e)%assign_j
                     d%edge_assign_pos(k, i, j) = edge_glo(e)%assign_pos
                     nf2 = domain_glo(edge_glo(e)%assign_domain)%face
                     d%edge_assign_sign(k, i, j) = 1 - 2*MOD(12 + tab_index(nf, nf2, 0), 2)
                     IF (MOD(6 + k + tab_index(nf, nf2, 0), 6) /= edge_glo(e)%assign_pos &
                         .AND. MOD(6 + k + tab_index(nf, nf2, 0), 6) /= MOD(edge_glo(e)%assign_pos + 3, 6)) THEN
                        d%edge_assign_sign(k, i, j) = -d%edge_assign_sign(k, i, j)
                     END IF
                  END IF

                  v = cell_glo(ind)%vertex(MOD(k + delta + 6, 6))
                  IF (v /= -1) THEN
                     d%vertex_assign_domain(k, i, j) = vertices_glo(v)%assign_domain
                     d%vertex_assign_i(k, i, j) = vertices_glo(v)%assign_i
                     d%vertex_assign_j(k, i, j) = vertices_glo(v)%assign_j
                     d%vertex_assign_pos(k, i, j) = vertices_glo(v)%assign_pos
                  ELSE
                     d%vertex_assign_domain(k, i, j) = -1
                  END IF
               END DO
               d%xyz(:, i, j) = cell_glo(ind)%xyz(:)
               IF (d%assign_domain(i, j) == ind_d) THEN
                  d%own(i, j) = .TRUE.
               ELSE
                  d%own(i, j) = .FALSE.
               END IF
            END DO
         END DO
      END DO

      nbp = 0
      DO ind_d = 1, ndomain_glo
         d => domain_glo(ind_d)
         nf = d%face
         DO j = d%jj_begin, d%jj_end
            DO i = d%ii_begin, d%ii_end
               DO k = 0, 5
                  IF (d%edge_assign_domain(k, i, j) == ind_d) nbp = nbp + 1
               END DO
            END DO
         END DO
      END DO

   CONTAINS

      SUBROUTINE assign_edge(ind_d, ind, i, j, delta, k)
         INTEGER :: ind_d, ind, i, j, delta, k
         INTEGER :: e
         e = cell_glo(ind)%edge(MOD(k + delta + 6, 6))
         edge_glo(e)%assign_domain = ind_d
         edge_glo(e)%assign_i = i
         edge_glo(e)%assign_j = j
         edge_glo(e)%assign_pos = k
         edge_glo(e)%assign_delta = delta

      END SUBROUTINE assign_edge

      SUBROUTINE assign_edge_if_not(ind_d, ind, i, j, delta, k)
         INTEGER :: ind_d, ind, i, j, delta, k
         INTEGER :: e
         e = cell_glo(ind)%edge(MOD(k + delta + 6, 6))
         IF (e /= -1) THEN
            IF (edge_glo(e)%assign_domain == -1) THEN
               edge_glo(e)%assign_domain = ind_d
               edge_glo(e)%assign_i = i
               edge_glo(e)%assign_j = j
               edge_glo(e)%assign_pos = k
               edge_glo(e)%assign_delta = delta
            END IF
         END IF
      END SUBROUTINE assign_edge_if_not

      SUBROUTINE assign_vertex(ind_d, ind, i, j, delta, k)
         INTEGER :: ind_d, ind, i, j, delta, k
         INTEGER :: e

         e = cell_glo(ind)%vertex(MOD(k + delta + 6, 6))
         vertices_glo(e)%assign_domain = ind_d
         vertices_glo(e)%assign_i = i
         vertices_glo(e)%assign_j = j
         vertices_glo(e)%assign_pos = k
         vertices_glo(e)%assign_delta = delta

      END SUBROUTINE assign_vertex

      SUBROUTINE assign_vertex_if_not(ind_d, ind, i, j, delta, k)
         INTEGER :: ind_d, ind, i, j, delta, k
         INTEGER :: e

         e = cell_glo(ind)%vertex(MOD(k + delta + 6, 6))
         IF (e /= -1) THEN
            IF (vertices_glo(e)%assign_domain == -1) THEN
               vertices_glo(e)%assign_domain = ind_d
               vertices_glo(e)%assign_i = i
               vertices_glo(e)%assign_j = j
               vertices_glo(e)%assign_pos = k
               vertices_glo(e)%assign_delta = delta
            END IF
         END IF

      END SUBROUTINE assign_vertex_if_not

   END SUBROUTINE assign_cell

   SUBROUTINE compute_boundary
      USE spherical_geom_mod
      INTEGER :: ind_d
      INTEGER :: i, j, k
      TYPE(t_domain), POINTER :: d
      REAL(rstd) :: ng1(3), ng2(3)

      DO ind_d = 1, ndomain_glo
         d => domain_glo(ind_d)
         DO j = d%jj_begin - 1, d%jj_end + 1
            DO i = d%ii_begin - 1, d%ii_end + 1
               DO k = 0, 5
                  ng1 = d%neighbour(:, MOD(k, 6), i, j)
                  ng2 = d%neighbour(:, MOD(k + 1, 6), i, j)
                  IF (sqrt(sum((ng1 - ng2)**2)) < 1e-16) ng2 = d%neighbour(:, MOD(k + 2, 6), i, j)
                  CALL circumcenter(d%xyz(:, i, j), ng1, ng2, d%vertex(:, k, i, j))
               END DO
            END DO
         END DO
      END DO
   END SUBROUTINE compute_boundary

   SUBROUTINE set_neighbour_indice
      USE metric
      INTEGER :: ind_d
      TYPE(t_domain), POINTER :: d

      DO ind_d = 1, ndomain_glo
         d => domain_glo(ind_d)
         d%t_right = 1
         d%t_left = -1
         d%t_rup = d%iim
         d%t_lup = d%iim - 1
         d%t_ldown = -d%iim
         d%t_rdown = -d%iim + 1

         d%u_right = 0
         d%u_lup = d%iim*d%jjm
         d%u_ldown = 2*d%iim*d%jjm

         d%u_rup = d%t_rup + d%u_ldown
         d%u_left = d%t_left + d%u_right
         d%u_rdown = d%t_rdown + d%u_lup

         d%z_up = 0
         d%z_down = d%iim*d%jjm
         d%z_rup = d%t_rup + d%z_down
         d%z_lup = d%t_lup + d%z_down
         d%z_ldown = d%t_ldown + d%z_up
         d%z_rdown = d%t_rdown + d%z_up

         d%t_pos(right) = d%t_right
         d%t_pos(rup) = d%t_rup
         d%t_pos(lup) = d%t_lup
         d%t_pos(left) = d%t_left
         d%t_pos(ldown) = d%t_ldown
         d%t_pos(rdown) = d%t_rdown

         d%u_pos(right) = d%u_right
         d%u_pos(rup) = d%u_rup
         d%u_pos(lup) = d%u_lup
         d%u_pos(left) = d%u_left
         d%u_pos(ldown) = d%u_ldown
         d%u_pos(rdown) = d%u_rdown

         d%z_pos(vrup) = d%z_rup
         d%z_pos(vup) = d%z_up
         d%z_pos(vlup) = d%z_lup
         d%z_pos(vldown) = d%z_ldown
         d%z_pos(vdown) = d%z_down
         d%z_pos(vrdown) = d%z_rdown

      END DO

   END SUBROUTINE set_neighbour_indice

   SUBROUTINE assign_domain
      USE mpipara
      USE grid_param
      USE metric_var_mod, ONLY: nb_face
      INTEGER :: nb_domain(0:mpi_size - 1)
      INTEGER :: rank, ind, ind_glo
      INTEGER :: block_j, jb, i, j, nd_glo, n, nf
      LOGICAL :: exit

      DO rank = 0, mpi_size - 1
         nb_domain(rank) = ndomain_glo/mpi_size
         IF (rank < MOD(ndomain_glo, mpi_size)) nb_domain(rank) = nb_domain(rank) + 1
      END DO

      ndomain = nb_domain(mpi_rank)

      ! ALLOCATE(domain(ndomain))
      CALL allocate_mesh(ndomain, mesh_loc)
      domain => mesh_loc%domain

      ALLOCATE (domloc_glo_ind(ndomain))

      block_j = INT(sqrt(nsplit_i*nsplit_j*nb_face*1./mpi_size))
      exit = .FALSE.
      jb = 1
      i = 1
      j = 1
      ind = 1
      nd_glo = 0
      rank = 0
      DO WHILE (.NOT. exit)

         IF (j == MIN(jb + block_j, nsplit_j*nb_face + 1)) THEN
            j = jb
            i = i + 1
         END IF

         IF (i > nsplit_i) THEN
            i = 1
            jb = jb + block_j
            j = jb
         END IF

         IF (ind > nb_domain(rank)) THEN
            rank = rank + 1
            ind = 1
         END IF
         ind_glo = (j - 1)*nsplit_i + i

         nd_glo = nd_glo + 1
         IF (nd_glo == ndomain_glo) THEN

            exit = .TRUE.
            IF (.NOT. (rank == mpi_size - 1 .AND. ind == nb_domain(rank))) THEN
               PRINT *, "Distribution problem in assign_domain"
               STOP
            END IF

         END IF

         domglo_rank(ind_glo) = rank
         domglo_loc_ind(ind_glo) = ind
         IF (rank == mpi_rank) THEN
            CALL copy_domain(domain_glo(ind_glo), domain(ind))
            domloc_glo_ind(ind) = ind_glo
         END IF

         j = j + 1
         ind = ind + 1

      END DO

      IF (is_mpi_master) THEN

         ind_glo = 0
         PRINT *, ''
         PRINT *, '      MPI PROCESS DISTRIBUTION'
         PRINT *, ''

         WRITE (*, "(' ')", ADVANCE='NO')
         DO n = 1, nsplit_i*7 - 1
            WRITE (*, "('=')", ADVANCE='NO')
         END DO
         PRINT *, ''

         DO nf = 1, nb_face
            DO j = 1, nsplit_j
               IF (j > 1) THEN
                  WRITE (*, "(' ')", ADVANCE='NO')
                  DO n = 1, nsplit_i*7 - 1
                     WRITE (*, "('-')", ADVANCE='NO')
                  END DO
                  PRINT *, ''
               END IF

               WRITE (*, "('|')", ADVANCE='NO')
               DO i = 1, nsplit_i
                  WRITE (*, "(' ','    ',' |')", ADVANCE='NO')
               END DO
               PRINT *, ''

               WRITE (*, "('|')", ADVANCE='NO')
               DO i = 1, nsplit_i
                  ind_glo = ind_glo + 1
                  WRITE (*, "(' ',i4.4  ,' |')", ADVANCE='NO') domglo_rank(ind_glo)
               END DO
               PRINT *, ''

               WRITE (*, "('|')", ADVANCE='NO')
               DO i = 1, nsplit_i
                  WRITE (*, "(' ','    ',' |')", ADVANCE='NO')
               END DO
               PRINT *, ''

            END DO

            WRITE (*, "(' ')", ADVANCE='NO')
            DO n = 1, nsplit_i*7 - 1
               WRITE (*, "('=')", ADVANCE='NO')
            END DO
            PRINT *, ''
         END DO
      END IF

!    rank=0
!    ind=0
!    DO ind_glo=1,ndomain_glo
!      ind=ind+1
!      domglo_rank(ind_glo)=rank
!      domglo_loc_ind(ind_glo)=ind
!      IF (rank==mpi_rank) THEN
!        CALL copy_domain(domain_glo(ind_glo),domain(ind))
!        domloc_glo_ind(ind)=ind_glo
!      ENDIF
!
!      IF (ind==nb_domain(rank)) THEN
!        rank=rank+1
!        ind=0
!      ENDIF
!    ENDDO

!$OMP PARALLEL
      CALL assign_domain_omp
!$OMP END PARALLEL

   END SUBROUTINE assign_domain

   SUBROUTINE assign_domain_omp
      USE omp_para
      USE mpipara
      INTEGER :: nb_domain
      INTEGER :: rank, ind, i

      ALLOCATE (assigned_domain(ndomain))

      ind = 0
      DO rank = 0, omp_domain_size - 1
         nb_domain = ndomain/omp_domain_size
         IF (rank < MOD(ndomain, omp_domain_size)) nb_domain = nb_domain + 1

         DO i = 1, nb_domain
            ind = ind + 1
            IF (rank == omp_domain_rank) THEN
               assigned_domain(ind) = .TRUE.
               WRITE (*, '(A,I6.6,A,I4.4,A,I4.4,A,I8.8)') "Rank ", mpi_rank, "  task ", rank, "  local domain ", ind, &
                  "  global domain ", domloc_glo_ind(ind)
            ELSE
               assigned_domain(ind) = .FALSE.
            END IF
         END DO

      END DO

      !$acc enter data copyin(assigned_domain(:))

   END SUBROUTINE assign_domain_omp

   SUBROUTINE init_domain_unst(dom)
      USE grid_param, ONLY: primal_num
      TYPE(t_domain) :: dom
      dom%ii_begin = 1
      dom%ii_end = primal_num
      dom%iim = primal_num
      dom%jj_begin = 1
      dom%jj_end = 1
      dom%jjm = 1
      ALLOCATE (dom%assign_domain(primal_num, 1))
      dom%assign_domain(:, :) = 1  ! all cells are assigned to MPI rank 0
   END SUBROUTINE init_domain_unst

   !======================= Define memory layouts for T,U,Z ========================

   SUBROUTINE setup_layout(dom, field_type, layout)
      USE layout_mod, ONLY: t_layout, field_T, field_U, field_Z, order_ij_l_p, order_l_ij_p
      USE grid_param, ONLY: grid_type, grid_unst, grid_ico, primal_num, edge_num, dual_num
      TYPE(t_domain), INTENT(IN)  :: dom(:)
      INTEGER, INTENT(IN)  :: field_type
      TYPE(t_layout), INTENT(OUT) :: layout(:)
      INTEGER :: ind, iim, jjm, ij_size

      SELECT CASE (grid_type)
      CASE (grid_ico)
         DO ind = 1, SIZE(dom)
            iim = dom(ind)%iim
            jjm = dom(ind)%jjm

            IF (field_type == field_T) THEN
               ij_size = iim*jjm
               layout(ind)%oriented = .FALSE.
            ELSE IF (field_type == field_U) THEN
               ij_size = 3*iim*jjm
               layout(ind)%oriented = .TRUE.
            ELSE IF (field_type == field_Z) THEN
               ij_size = 2*iim*jjm
               layout(ind)%oriented = .FALSE.
            END IF

            layout(ind)%ij_size = ij_size
            layout(ind)%index_order = order_ij_l_p
         END DO

      CASE (grid_unst)
         SELECT CASE (field_type)
         CASE (field_T)
            ij_size = primal_num
            layout(:)%oriented = .FALSE.
         CASE (field_U)
            ij_size = edge_num
            layout(:)%oriented = .TRUE.
         CASE (field_Z)
            ij_size = dual_num
            layout(:)%oriented = .FALSE.
         END SELECT
         layout(:)%ij_size = ij_size
         layout(:)%index_order = order_l_ij_p
      END SELECT

   END SUBROUTINE setup_layout

   SUBROUTINE init_layouts
      USE layout_mod
      INTEGER :: id
      ALLOCATE (layout(ndomain, field_type_max))
      ALLOCATE (layout_glo(ndomain_glo, field_type_max))
      DO id = 1, field_type_max
         CALL setup_layout(domain, id, layout(:, id))
         CALL setup_layout(domain_glo, id, layout_glo(:, id))
      END DO
   END SUBROUTINE init_layouts

   !====================================================================

   SUBROUTINE compute_domain
      USE grid_param, ONLY: grid_type, grid_unst, grid_ico
      INTEGER :: ind
      SELECT CASE (grid_type)
      CASE (grid_unst)
         ! FIXME temporary, sequential hack
         ndomain_glo = 1
         CALL allocate_mesh(ndomain_glo, mesh_glo)
         domain_glo => mesh_glo%domain
         ALLOCATE (domglo_rank(ndomain_glo))
         ALLOCATE (domglo_loc_ind(ndomain_glo))
         domglo_rank(:) = 0
         domglo_loc_ind(:) = 1

         ndomain = 1
         CALL allocate_mesh(ndomain, mesh_loc)
         domain => mesh_loc%domain
         ALLOCATE (domloc_glo_ind(ndomain))
         ALLOCATE (assigned_domain(ndomain))
         domloc_glo_ind(:) = 1
         assigned_domain(:) = .TRUE.

         CALL init_domain_unst(domain(1))
         CALL init_domain_unst(domain_glo(1))

         CALL init_layouts

      CASE DEFAULT
         CALL init_domain_param
         CALL create_domain
         CALL assign_cell
         CALL compute_boundary
         CALL set_neighbour_indice
         CALL assign_domain
         CALL init_layouts
      END SELECT

   END SUBROUTINE compute_domain

END MODULE domain_mod
