dis_main Subroutine

  subroutine dis_main(dis_control, dis_spheres, dis_manifold, kmesh_info, kpt_latt, sitesym, &
                      print_output, m_matrix_orig_local, u_matrix, u_matrix_opt, eigval, &
                      real_lattice, omega_invariant, num_bands, num_kpts, num_wann, gamma_only, &
                      lsitesymmetry, stdout, timer, dist_k, error, comm)
    !================================================!
    !
    !! Main disentanglement routine
    !
    !================================================!
    use w90_comms, only: comms_bcast, w90_comm_type, mpirank
    use w90_constants, only: dp, cmplx_0, cmplx_1
    use w90_error
    use w90_io, only: io_stopwatch_start, io_stopwatch_stop
    use w90_sitesym, only: sitesym_replace_d_matrix_band, sitesym_symmetrize_u_matrix, &
                           sitesym_symmetrize_zmatrix, sitesym_dis_extract_symmetry
    use w90_types, only: dis_manifold_type, kmesh_info_type, print_output_type, timer_list_type
    use w90_utility, only: utility_recip_lattice_base
    use w90_wannier90_types, only: dis_control_type, dis_spheres_type, sitesym_type

    ! arguments
    integer, intent(in) :: num_bands, num_kpts, num_wann
    integer, intent(in) :: stdout
    integer, intent(in) :: dist_k(:)

    logical, intent(in) :: lsitesymmetry
    logical, intent(in) :: gamma_only

    real(kind=dp), pointer, intent(in) :: eigval(:, :) ! (num_bands, num_kpts)
    real(kind=dp), intent(in) :: kpt_latt(:, :)
    real(kind=dp), intent(inout) :: omega_invariant
    real(kind=dp), intent(in) :: real_lattice(3, 3)

    complex(kind=dp), intent(inout) :: u_matrix(:, :, :) ! (num_wann, num_wann, num_kpts)
    complex(kind=dp), intent(inout) :: u_matrix_opt(:, :, :) ! (num_bands, num_wann, num_kpts)
    complex(kind=dp), intent(inout) :: m_matrix_orig_local(:, :, :, :) ! this is the only "m matrix" here now

    type(dis_control_type), intent(inout)  :: dis_control
    type(dis_manifold_type), intent(inout) :: dis_manifold
    type(dis_spheres_type), intent(in) :: dis_spheres
    type(kmesh_info_type), intent(in) :: kmesh_info
    type(print_output_type), intent(in) :: print_output
    type(sitesym_type), intent(inout) :: sitesym
    type(w90_comm_type), intent(in) :: comm
    type(timer_list_type), intent(inout) :: timer
    type(w90_error_type), allocatable, intent(out) :: error

    ! internal variables
    real(kind=dp) :: recip_lattice(3, 3), volume
    integer :: nkp, nkp2, nn, j, ierr, nkp_global
    logical :: linner                         !! Is there a frozen window
    logical :: lfrozen(num_bands, num_kpts)   !! true if the i-th band inside outer window is frozen
    integer :: ndimfroz(num_kpts)             !! number of frozen bands at nkp-th k point
    integer :: indxfroz(num_bands, num_kpts)  !! number of bands inside outer window at nkp-th k point
    integer :: indxnfroz(num_bands, num_kpts) !! outer-window band index for the i-th non-frozen state
    complex(kind=dp), allocatable :: a_matrix(:, :, :) ! (num_bands, num_wann, num_kpts)
    !! (equals 1 if it is the bottom of outer window)

    real(kind=dp), allocatable :: eigval_opt(:, :)  !! At input it contains a large set of eigenvalues. At
    !! it is slimmed down to contain only those inside the energy window.

    complex(kind=dp), allocatable :: cwb(:, :), cww(:, :)

    ! pllel setup
    integer :: nkrank, ikg, ikl, my_node_id
    integer, allocatable :: global_k(:)
    logical :: on_root = .false.

    my_node_id = mpirank(comm)
    on_root = (my_node_id == 0)
    nkrank = count(dist_k == my_node_id) ! this routine must proceed also in the case of zero k-points this rank, to ensure collective communications are matched

    allocate (a_matrix(num_bands, num_wann, num_kpts), stat=ierr) ! a_matrix is local to disentangle()
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating a_matrix in dis_main', comm)
      return
    end if
    a_matrix = u_matrix_opt ! initial projections are passed to this routine via u_matrix_opt

    allocate (global_k(nkrank), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating global_k in dis_main', comm)
      return
    end if
    global_k = huge(1); ikl = 1
    do ikg = 1, num_kpts
      if (dist_k(ikg) == my_node_id) then
        global_k(ikl) = ikg
        ikl = ikl + 1
      end if
    end do

    if (print_output%timing_level > 0) call io_stopwatch_start('dis: main', timer)

    call utility_recip_lattice_base(real_lattice, recip_lattice, volume)

    if (print_output%iprint > 0) write (stdout, '(/1x,a)') &
      '*------------------------------- DISENTANGLE --------------------------------*'

    ! Allocate arrays
    allocate (eigval_opt(num_bands, num_kpts), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating eigval_opt in dis_main', comm)
      return
    end if
    eigval_opt(1:num_bands, 1:num_kpts) = eigval(1:num_bands, 1:num_kpts)

    ! Set up energy windows
    if (dis_manifold%frozen_proj) then
      call dis_windows_proj(dis_manifold, eigval_opt, a_matrix, m_matrix_orig_local, &
                            indxfroz, indxnfroz, ndimfroz, dis_manifold%nfirstwin, &
                            print_output%iprint, kmesh_info%nnlist, kmesh_info%nntot, num_bands, &
                            num_kpts, num_wann, print_output%timing_level, lfrozen, linner, &
                            on_root, stdout, dist_k, global_k, my_node_id, timer, error, comm)
      if (allocated(error)) return
    else
      call dis_windows(dis_spheres, dis_manifold, eigval_opt, kpt_latt, recip_lattice, indxfroz, &
                       indxnfroz, ndimfroz, print_output%iprint, num_bands, num_kpts, num_wann, &
                       print_output%timing_level, lfrozen, linner, on_root, stdout, timer, error, comm)
      if (allocated(error)) return
    end if

    ! Construct the unitarized projection
    call dis_project(a_matrix, u_matrix_opt, dis_manifold%ndimwin, dis_manifold%nfirstwin, &
                     num_bands, num_kpts, num_wann, print_output%timing_level, on_root, &
                     print_output%iprint, timer, error, stdout, comm)
    if (allocated(error)) return

    ! If there is an inner window, need to modify projection procedure
    ! (Sec. III.G SMV)
    if (linner) then
      if (lsitesymmetry) then
        call set_error_fatal(error, 'in symmetry-adapted mode, frozen window not implemented yet', &
                             comm)
        return
      end if
      if (print_output%iprint > 0) write (stdout, '(3x,a)') 'Using an inner window (linner = T)'
      call dis_proj_froz(u_matrix_opt, indxfroz, ndimfroz, dis_manifold%ndimwin, &
                         print_output%iprint, num_bands, num_kpts, num_wann, &
                         print_output%timing_level, lfrozen, on_root, timer, error, stdout, comm)
      if (allocated(error)) return
    else
      if (print_output%iprint > 0) write (stdout, '(3x,a)') 'No inner window (linner = F)'
    end if

    ! Debug
    call internal_check_orthonorm(u_matrix_opt, dis_manifold%ndimwin, num_kpts, num_wann, &
                                  print_output%timing_level, on_root, timer, error, stdout, comm)
    if (allocated(error)) return

    ! For frozen_proj, these are done inside dis_windows_proj()
    if (.not. dis_manifold%frozen_proj) then
      ! Slim down the original Mmn(k,b)

      call internal_slim_m(m_matrix_orig_local, dis_manifold%ndimwin, dis_manifold%nfirstwin, &
                           kmesh_info%nnlist, kmesh_info%nntot, num_bands, print_output%timing_level, &
                           timer, dist_k, global_k, error, comm)
      if (allocated(error)) return

      dis_manifold%lwindow = .false.
      do nkp = 1, num_kpts
        do j = dis_manifold%nfirstwin(nkp), dis_manifold%nfirstwin(nkp) + dis_manifold%ndimwin(nkp) - 1
          dis_manifold%lwindow(j, nkp) = .true.
        end do
      end do
    end if

    if (lsitesymmetry) then
      call sitesym_symmetrize_u_matrix(sitesym, u_matrix_opt, num_bands, num_bands, num_kpts, &
                                       num_wann, stdout, error, comm, dis_manifold%lwindow)
      if (allocated(error)) return
    end if

    !RS: calculate initial U_{opt}(Rk) from U_{opt}(k)
    ! Extract the optimally-connected num_wann-dimensional subspaces

    if (gamma_only) then
      call dis_extract_gamma(dis_control, kmesh_info, print_output, dis_manifold, &
                             m_matrix_orig_local, u_matrix_opt, eigval_opt, omega_invariant, &
                             indxnfroz, ndimfroz, num_bands, num_kpts, num_wann, timer, error, &
                             stdout, comm)
      if (allocated(error)) return
    else
      call dis_extract(dis_control, kmesh_info, sitesym, print_output, dis_manifold, &
                       m_matrix_orig_local, u_matrix_opt, eigval_opt, omega_invariant, indxnfroz, &
                       ndimfroz, my_node_id, num_bands, num_kpts, num_wann, lsitesymmetry, timer, &
                       nkrank, global_k, error, stdout, comm)
      if (allocated(error)) return
    end if

    ! Allocate workspace
    allocate (cwb(num_wann, num_bands), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating cwb in dis_main', comm)
      return
    end if
    allocate (cww(num_wann, num_wann), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating cww in dis_main', comm)
      return
    end if

    ! Find the num_wann x num_wann overlap matrices between
    ! the basis states of the optimal subspaces
    do nkp = 1, nkrank
      nkp_global = global_k(nkp)
      do nn = 1, kmesh_info%nntot
        nkp2 = kmesh_info%nnlist(nkp_global, nn)
        call zgemm('C', 'N', num_wann, dis_manifold%ndimwin(nkp2), dis_manifold%ndimwin(nkp_global), &
                   cmplx_1, u_matrix_opt(:, :, nkp_global), num_bands, &
                   m_matrix_orig_local(:, :, nn, nkp), num_bands, cmplx_0, cwb, num_wann)
        call zgemm('N', 'N', num_wann, num_wann, dis_manifold%ndimwin(nkp2), cmplx_1, cwb, &
                   num_wann, u_matrix_opt(:, :, nkp2), num_bands, cmplx_0, cww, num_wann)
        m_matrix_orig_local(1:num_wann, 1:num_wann, nn, nkp) = cww(:, :)
      end do
    end do

    ! Find the initial u_matrix
    if (lsitesymmetry) call sitesym_replace_d_matrix_band(sitesym, num_wann, error, comm)

    if (gamma_only) then
      call internal_find_u_gamma(a_matrix, u_matrix, u_matrix_opt, dis_manifold%ndimwin, num_wann, &
                                 print_output%timing_level, stdout, timer, error, comm)
      if (allocated(error)) return
    else
      call internal_find_u(sitesym, a_matrix, u_matrix, u_matrix_opt, dis_manifold%ndimwin, &
                           num_bands, num_kpts, num_wann, print_output%timing_level, &
                           lsitesymmetry, on_root, stdout, timer, error, comm)
      if (allocated(error)) return
    end if

    !zero the unused elements of u_matrix_opt (just in case...)
    do nkp = 1, num_kpts
      do j = 1, num_wann
        if (dis_manifold%ndimwin(nkp) < num_bands) &
          u_matrix_opt(dis_manifold%ndimwin(nkp) + 1:, j, nkp) = cmplx_0
      end do
    end do

    ! Deallocate workspace
    deallocate (cww, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating cww in dis_main', comm)
      return
    end if
    deallocate (cwb, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating cwb in dis_main', comm)
      return
    end if
    deallocate (global_k, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating global_k in dis_main', comm)
      return
    end if
    deallocate (a_matrix, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating a_matrix in dis_main', comm)
      return
    end if

    if (print_output%timing_level > 0 .and. on_root) call io_stopwatch_stop('dis: main', timer)

    return
    !================================================!
  end subroutine dis_main