ws_expand_rvec Subroutine

  subroutine ws_expand_rvec(ws_distance, use_ws_distance, num_wann, nrpts, irvec, ndegen, &
                            irvec_full, nrpts_full, ir_map, ir_origin, error, comm)
    !================================================!
    !! Build the fully expanded list of lattice vectors, i.e. the set of all
    !! R + T that occur in the Wigner-Seitz mapping computed by ws_translate_dist,
    !! together with the index map ir_map(ideg, i, j, ir) that sends a degenerate
    !! image of the pair (i, j) at the folded vector irvec(:, ir) to its position
    !! in that list.
    !!
    !! The expanded list is ordered lexicographically, so that it does not depend
    !! on the order in which the vectors are discovered.
    !!
    !! If use_ws_distance is false there is nothing to expand: irvec_full is
    !! irvec and ir_map is unused.
    !================================================!

    use w90_types, only: ws_distance_type

    implicit none

    type(ws_distance_type), intent(in) :: ws_distance
    type(w90_error_type), allocatable, intent(out) :: error
    type(w90_comm_type), intent(in) :: comm

    logical, intent(in) :: use_ws_distance
    integer, intent(in) :: num_wann
    integer, intent(in) :: nrpts
    integer, intent(in) :: irvec(3, nrpts)
    integer, intent(in) :: ndegen(nrpts)

    integer, allocatable, intent(out) :: irvec_full(:, :)
    integer, intent(out) :: nrpts_full
    integer, allocatable, intent(out) :: ir_map(:, :, :, :)
    integer, intent(out) :: ir_origin
    !! index of R = 0 in the expanded list

    ! local variables
    integer :: i, j, ideg, ir, i1, i2, i3, ierr, max_ndeg, rpt_origin
    integer :: ivdum(3), ivmin(3), ivmax(3)
    integer, allocatable :: index_box(:, :, :)

    rpt_origin = 0
    do ir = 1, nrpts
      if (all(irvec(:, ir) == 0)) rpt_origin = ir
    end do
    if (rpt_origin == 0) then
      call set_error_fatal(error, 'R=0 is not in the list of lattice vectors.', comm)
      return
    end if
    if (ndegen(rpt_origin) /= 1) then
      call set_error_fatal(error, 'ndegen for R=0 is not 1.', comm)
      return
    end if

    if (.not. use_ws_distance) then
      nrpts_full = nrpts
      ir_origin = rpt_origin

      allocate (irvec_full(3, nrpts_full), stat=ierr)
      if (ierr /= 0) then
        call set_error_alloc(error, 'Error in allocating irvec_full in ws_expand_rvec', comm)
        return
      end if
      ! ws_apply_ndegen does not consult ir_map when there is nothing to expand,
      ! so allocate it only to have something to pass
      allocate (ir_map(1, 1, 1, 1), stat=ierr)
      if (ierr /= 0) then
        call set_error_alloc(error, 'Error in allocating ir_map in ws_expand_rvec', comm)
        return
      end if
      ir_map = -1

      irvec_full = irvec
      return
    end if

    ! Check degeneracy factor ws_distance%ndeg for a Wannier function with itself,
    ! i.e. R = 0 and i = j, is 1.
    do ir = 1, nrpts
      do i = 1, num_wann
        do ideg = 1, ws_distance%ndeg(i, i, ir)
          if (all(ws_distance%irdist(:, ideg, i, i, ir) == 0)) then
            if (ws_distance%ndeg(i, i, ir) /= 1) then
              call set_error_fatal(error, 'ws_distance%ndeg for R=0 and i=j is not 1.', comm)
              return
            end if
          end if
        end do
      end do
    end do

    max_ndeg = maxval(ws_distance%ndeg)

    ! Mark every vector that occurs in irdist on an integer box spanning them all,
    ! then walk the box in lexicographic order to number the vectors found.
    ! Unused slots of irdist are zero, which is a vector of the list anyway.
    ! R_wz_sc bounds the box by +-2*(ws_search_size + 1)*mp_grid, and in practice
    ! it is a small multiple of mp_grid: a few MB of integers at worst.
    do i = 1, 3
      ivmin(i) = minval(ws_distance%irdist(i, :, :, :, :))
      ivmax(i) = maxval(ws_distance%irdist(i, :, :, :, :))
    end do

    allocate (index_box(ivmin(1):ivmax(1), ivmin(2):ivmax(2), ivmin(3):ivmax(3)), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating index_box in ws_expand_rvec', comm)
      return
    end if
    index_box = 0

    nrpts_full = 0
    do ir = 1, nrpts
      do j = 1, num_wann
        do i = 1, num_wann
          do ideg = 1, ws_distance%ndeg(i, j, ir)
            ivdum = ws_distance%irdist(:, ideg, i, j, ir)
            if (index_box(ivdum(1), ivdum(2), ivdum(3)) == 0) then
              index_box(ivdum(1), ivdum(2), ivdum(3)) = 1
              nrpts_full = nrpts_full + 1
            end if
          end do
        end do
      end do
    end do

    allocate (irvec_full(3, nrpts_full), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating irvec_full in ws_expand_rvec', comm)
      return
    end if
    allocate (ir_map(max_ndeg, num_wann, num_wann, nrpts), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating ir_map in ws_expand_rvec', comm)
      return
    end if

    ! each marked slot is visited once, so a slot still holding the mark 1 has
    ! not been numbered yet
    ir = 0
    do i1 = ivmin(1), ivmax(1)
      do i2 = ivmin(2), ivmax(2)
        do i3 = ivmin(3), ivmax(3)
          if (index_box(i1, i2, i3) == 1) then
            ir = ir + 1
            index_box(i1, i2, i3) = ir
            irvec_full(:, ir) = (/i1, i2, i3/)
          end if
        end do
      end do
    end do
    ir_origin = index_box(0, 0, 0)

    ir_map = -1
    do ir = 1, nrpts
      do j = 1, num_wann
        do i = 1, num_wann
          do ideg = 1, ws_distance%ndeg(i, j, ir)
            ivdum = ws_distance%irdist(:, ideg, i, j, ir)
            ir_map(ideg, i, j, ir) = index_box(ivdum(1), ivdum(2), ivdum(3))
          end do
        end do
      end do
    end do

    deallocate (index_box, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating index_box in ws_expand_rvec', comm)
      return
    end if
    !================================================!
  end subroutine ws_expand_rvec