wann_main Subroutine

  subroutine wann_main(ham_logical, kmesh_info, kpt_latt, wann_control, omega, sitesym, &
                       print_output, wannier_data, ws_region, w90_calculation, ham_k, ham_r, &
                       m_matrix_loc, u_matrix, real_lattice, wannier_centres_translated, irvec, &
                       mp_grid, ndegen, nrpts, num_kpts, num_proj, num_wann, optimisation, &
                       rpt_origin, bands_plot_mode, transport_mode, lsitesymmetry, stdout, &
                       timer, dist_k, error, comm)
    !================================================!
    !
    !! Calculate the Unitary Rotations to give Maximally Localised Wannier Functions
    !
    !================================================
    use w90_constants, only: dp, cmplx_1, cmplx_0, twopi, cmplx_i
    use w90_io, only: io_wallclocktime, io_stopwatch_start, io_stopwatch_stop
    use w90_wannier90_types, only: wann_control_type, w90_calculation_type, wann_omega_type, &
                                   sitesym_type, ham_logical_type
    use w90_types, only: kmesh_info_type, print_output_type, wannier_data_type, ws_region_type, &
                         timer_list_type
    use w90_utility, only: utility_frac_to_cart, utility_zgemm
    use w90_sitesym, only: sitesym_symmetrize_gradient
    use w90_comms, only: mpisize, mpirank, comms_allreduce, w90_comm_type
    use w90_hamiltonian, only: hamiltonian_setup

    implicit none

    ! arguments
    type(ham_logical_type), intent(inout)    :: ham_logical
    type(kmesh_info_type), intent(in)        :: kmesh_info
    type(ws_region_type), intent(in)         :: ws_region
    type(print_output_type), intent(in)      :: print_output
    type(wann_control_type), intent(inout)   :: wann_control
    type(wann_omega_type), intent(inout)     :: omega
    type(sitesym_type), intent(in)           :: sitesym
    type(w90_calculation_type), intent(in)   :: w90_calculation
    type(w90_comm_type), intent(in)           :: comm
    type(wannier_data_type), intent(inout)   :: wannier_data
    type(timer_list_type), intent(inout)     :: timer
    type(w90_error_type), allocatable, intent(out) :: error

    integer, intent(in) :: mp_grid(3)
    integer, intent(in) :: num_kpts
    integer, intent(in) :: num_proj
    integer, intent(in) :: num_wann
    integer, intent(in) :: optimisation
    integer, intent(inout), allocatable :: irvec(:, :)
    integer, intent(inout), allocatable :: ndegen(:)
    integer, intent(inout) :: nrpts
    integer, intent(inout) :: rpt_origin
    integer, intent(in) :: stdout
    integer, intent(in) :: dist_k(:)

    real(kind=dp), intent(in) :: kpt_latt(:, :)
    real(kind=dp), intent(inout), allocatable :: wannier_centres_translated(:, :)
    real(kind=dp), intent(in) :: real_lattice(3, 3)

    complex(kind=dp), intent(inout), allocatable :: ham_k(:, :, :)
    complex(kind=dp), intent(inout), allocatable :: ham_r(:, :, :)
    !complex(kind=dp), intent(inout) :: m_matrix(:, :, :, :)
    complex(kind=dp), intent(inout) :: m_matrix_loc(:, :, :, :)
    complex(kind=dp), intent(inout) :: u_matrix(:, :, :)

    logical, intent(in) :: lsitesymmetry

    character(len=*), intent(in) :: bands_plot_mode
    character(len=*), intent(in) :: transport_mode

    ! local variables
    type(localisation_vars_type) :: old_spread
    type(localisation_vars_type) :: wann_spread
    type(localisation_vars_type) :: trial_spread

    ! Data to avoid large allocation within iteration loop
    real(kind=dp), allocatable :: rnkb_loc(:, :, :)
    real(kind=dp), allocatable :: ln_tmp(:, :, :)
    real(kind=dp), allocatable :: ln_tmp_loc(:, :, :)
    real(kind=dp), allocatable :: sheet(:, :, :)
    real(kind=dp), allocatable :: rave(:, :), r2ave(:), rave2(:)
    ! guiding centres
    real(kind=dp), allocatable :: rguide(:, :)

    complex(kind=dp), allocatable :: u_matrix_loc(:, :, :)
    complex(kind=dp), allocatable :: cdq_loc(:, :, :) ! the only large array sent from process to process in the main loop
    complex(kind=dp), allocatable :: cdodq_loc(:, :, :)
    complex(kind=dp), allocatable :: csheet(:, :, :)
    complex(kind=dp), allocatable :: cdodq(:, :, :)
    complex(kind=dp), allocatable :: cdodq_r(:, :, :)
    complex(kind=dp), allocatable :: k_to_r(:, :)
    complex(kind=dp), allocatable :: cdodq_precond(:, :, :)
    complex(kind=dp), allocatable :: cdodq_precond_loc(:, :, :)

    !local arrays not passed into subroutines
    complex(kind=dp), allocatable  :: cwschur1(:), cwschur2(:)
    complex(kind=dp), allocatable  :: cwschur3(:), cwschur4(:)
    complex(kind=dp), allocatable  :: cdq(:, :, :)!,cdqkeep(:,:,:)
    ! cdqkeep is replaced by cdqkeep_loc
    complex(kind=dp), allocatable  :: cdqkeep_loc(:, :, :)
    complex(kind=dp), allocatable  :: cz(:, :)
    complex(kind=dp), allocatable  :: cmtmp(:, :), tmp_cdq(:, :)
    ! complex(kind=dp), allocatable  :: m0(:,:,:,:),u0(:,:,:)
    ! m0 and u0 are replaced by m0_loc and u0_loc
    complex(kind=dp), allocatable  :: m0_loc(:, :, :, :), u0_loc(:, :, :)
    complex(kind=dp), allocatable  :: cwork(:)

    real(kind=dp), allocatable  :: evals(:)
    real(kind=dp), allocatable  :: rwork(:)
    real(kind=dp), allocatable :: history(:)
    real(kind=dp), allocatable :: rnr0n2(:)

    logical :: first_pass
    !! Used to trigger the calculation of the invarient spread we only need to do this on entering wann_main (_gamma)
    real(kind=dp) :: lambda_loc

    integer, allocatable :: global_k(:)
    complex(kind=dp) :: rdotk
    integer :: conv_count, noise_count, page_unit
    integer :: i, n, iter, ind, ierr, iw, ncg, nkp, nkp_loc
    integer :: irguide
    integer :: irpt, loop_kpt
    integer :: nkrank
    logical :: lconverged, lrandom, lfirst
    logical :: lprint, ldump, lquad
    real(kind=dp) :: doda0
    real(kind=dp) :: falphamin, alphamin
    real(kind=dp) :: gcfac, gcnorm1, gcnorm0
    real(kind=dp) :: save_spread

    ! pllel setup
    logical :: on_root = .false.
    integer :: num_nodes, my_node_id

    num_nodes = mpisize(comm)
    my_node_id = mpirank(comm)
    if (my_node_id == 0) on_root = .true.

    if (print_output%timing_level > 0 .and. print_output%iprint > 0) then
      call io_stopwatch_start('wann: main', timer)
    end if

    first_pass = .true.

    nkrank = count(dist_k == my_node_id) ! number k this rank, for dimensioning
    ! there is no need to round up to 1, but less than zero is nonsense
    if (nkrank < 0) then
      call set_error_fatal(error, 'kpt decomposition nonsensical in wann_main', comm)
      return
    end if
    allocate (global_k(nkrank), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating local kpoint distribution in wann_main', comm)
      return
    end if
    loop_kpt = 1
    do i = 1, num_kpts
      if (dist_k(i) == my_node_id) then
        global_k(loop_kpt) = i
        loop_kpt = loop_kpt + 1
      end if
    end do

    ! Allocate stuff
    allocate (history(wann_control%conv_window), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error allocating history in wann_main', comm)
      return
    end if
    allocate (ln_tmp(num_wann, kmesh_info%nntot, num_kpts), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating ln_tmp in wann_main', comm)
      return
    end if
    if (wann_control%constrain%selective_loc) then
      allocate (rnr0n2(wann_control%constrain%slwf_num), stat=ierr)
      if (ierr /= 0) then
        call set_error_alloc(error, 'Error in allocating rnr0n2 in wann_main', comm)
        return
      end if
    end if

    ! sub vars passed into other subs
    allocate (csheet(num_wann, kmesh_info%nntot, num_kpts), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating csheet in wann_main', comm)
      return
    end if
    allocate (cdodq(num_wann, num_wann, num_kpts), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating cdodq in wann_main', comm)
      return
    end if
    allocate (sheet(num_wann, kmesh_info%nntot, num_kpts), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating sheet in wann_main', comm)
      return
    end if
    allocate (rave(3, num_wann), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating rave in wann_main', comm)
      return
    end if
    allocate (r2ave(num_wann), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating r2ave in wann_main', comm)
      return
    end if
    allocate (rave2(num_wann), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating rave2 in wann_main', comm)
      return
    end if
    allocate (rguide(3, num_wann), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating rguide in wann_main', comm)
      return
    end if

    if (wann_control%precond) then
      call hamiltonian_setup(ham_logical, print_output, ws_region, w90_calculation, ham_k, ham_r, &
                             real_lattice, wannier_centres_translated, irvec, mp_grid, ndegen, &
                             num_kpts, num_wann, nrpts, rpt_origin, bands_plot_mode, stdout, &
                             timer, error, transport_mode, comm)
      if (allocated(error)) return

      allocate (cdodq_r(num_wann, num_wann, nrpts), stat=ierr)
      if (ierr /= 0) then
        call set_error_alloc(error, 'Error in allocating cdodq_r in wann_main', comm)
        return
      end if
      allocate (cdodq_precond(num_wann, num_wann, num_kpts), stat=ierr)
      if (ierr /= 0) then
        call set_error_alloc(error, 'Error in allocating cdodq_precond in wann_main', comm)
        return
      end if

      ! this method of computing the preconditioning is much more efficient, but requires more RAM
      if (optimisation >= 3) then
        allocate (k_to_r(num_kpts, nrpts), stat=ierr)
        if (ierr /= 0) then
          call set_error_alloc(error, 'Error in allocating k_to_r in wann_main', comm)
          return
        end if

        do irpt = 1, nrpts
          do loop_kpt = 1, num_kpts
            rdotk = twopi*dot_product(kpt_latt(:, loop_kpt), real(irvec(:, irpt), dp))
            k_to_r(loop_kpt, irpt) = exp(-cmplx_i*rdotk)
          end do
        end do
      end if
    end if

    csheet = cmplx_1; cdodq = cmplx_0
    sheet = 0.0_dp; rave = 0.0_dp; r2ave = 0.0_dp; rave2 = 0.0_dp; rguide = 0.0_dp

    ! sub vars not passed into other subs
    allocate (cwschur1(num_wann), cwschur2(10*num_wann), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating cwshur1 in wann_main', comm)
      return
    end if
    allocate (cwschur3(num_wann), cwschur4(num_wann), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating cwshur3 in wann_main', comm)
      return
    end if
    allocate (cdq(num_wann, num_wann, num_kpts), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating cdq in wann_main', comm)
      return
    end if
    allocate (rnkb_loc(num_wann, kmesh_info%nntot, nkrank), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating rnkb_loc in wann_main', comm)
      return
    end if
    if (wann_control%use_ss_functional) then
      allocate (ln_tmp_loc(num_wann, kmesh_info%nntot, 1), stat=ierr)
    else
      allocate (ln_tmp_loc(num_wann, kmesh_info%nntot, nkrank), stat=ierr)
    end if
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating ln_tmp_loc in wann_main', comm)
      return
    end if
    allocate (u_matrix_loc(num_wann, num_wann, nkrank), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating u_matrix_loc in wann_main', comm)
      return
    end if
    if (wann_control%precond) then
      allocate (cdodq_precond_loc(num_wann, num_wann, nkrank), stat=ierr)
      if (ierr /= 0) then
        call set_error_alloc(error, 'Error in allocating cdodq_precond_loc in wann_main', comm)
        return
      end if
    end if

    ! initialize local u matrix with global one
    do nkp_loc = 1, nkrank
      nkp = global_k(nkp_loc)
      u_matrix_loc(:, :, nkp_loc) = u_matrix(:, :, nkp)
    end do

    allocate (cdq_loc(num_wann, num_wann, nkrank), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating cdq_loc in wann_main', comm)
      return
    end if
    allocate (cdodq_loc(num_wann, num_wann, nkrank), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating cdodq_loc in wann_main', comm)
      return
    end if
    allocate (cdqkeep_loc(num_wann, num_wann, nkrank), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating cdqkeep_loc in wann_main', comm)
      return
    end if
    if (optimisation > 0) then
      allocate (m0_loc(num_wann, num_wann, kmesh_info%nntot, nkrank), stat=ierr)
    end if
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating m0_loc in wann_main', comm)
      return
    end if
    allocate (u0_loc(num_wann, num_wann, nkrank), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating u0_loc in wann_main', comm)
      return
    end if
    allocate (cz(num_wann, num_wann), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating cz in wann_main', comm)
      return
    end if
    allocate (cmtmp(num_wann, num_wann), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating cmtmp in wann_main', comm)
      return
    end if
    allocate (tmp_cdq(num_wann, num_wann), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating tmp_cdq in wann_main', comm)
      return
    end if
    allocate (evals(num_wann), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating evals in wann_main', comm)
      return
    end if
    allocate (cwork(4*num_wann), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating cwork in wann_main', comm)
      return
    end if
    allocate (rwork(3*num_wann - 2), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating rwork in wann_main', comm)
      return
    end if

    cwschur1 = cmplx_0; cwschur2 = cmplx_0; cwschur3 = cmplx_0; cwschur4 = cmplx_0
    cdq = cmplx_0; cz = cmplx_0; cmtmp = cmplx_0; cdqkeep_loc = cmplx_0; cdq_loc = cmplx_0
    gcnorm1 = 0.0_dp; gcnorm0 = 0.0_dp

    ! initialise rguide to projection centres (Cartesians in units of Ang)
    if (wann_control%guiding_centres%enable) then
      do n = 1, num_wann
        call utility_frac_to_cart(wann_control%guiding_centres%centres(:, n), rguide(:, n), &
                                  real_lattice)
      end do
    end if

    if (print_output%iprint > 0) then
      write (stdout, *)
      write (stdout, '(1x,a)') '*------------------------------- WANNIERISE ---------------------------------*'
      write (stdout, '(1x,a)') '+--------------------------------------------------------------------+<-- CONV'
      if (trim(print_output%length_unit) == 'Ang') then
        write (stdout, '(1x,a)') '| Iter  Delta Spread     RMS Gradient      Spread (Ang^2)      Time  |<-- CONV'
      else
        write (stdout, '(1x,a)') '| Iter  Delta Spread     RMS Gradient      Spread (Bohr^2)     Time  |<-- CONV'
      end if
      write (stdout, '(1x,a)') '+--------------------------------------------------------------------+<-- CONV'
      write (stdout, *)
    end if

    irguide = 0
    if (wann_control%guiding_centres%enable .and. (wann_control%guiding_centres%num_no_guide_iter .le. 0)) then
      call wann_phases(csheet, sheet, rguide, irguide, num_wann, kmesh_info, num_kpts, &
                       wann_control%use_ss_functional, m_matrix_loc, print_output%timing_level, &
                       print_output%iprint, timer, nkrank, global_k, error, comm)
      if (allocated(error)) return

      irguide = 1
    end if

    ! constrained centres part
    lambda_loc = 0.0_dp
    if (wann_control%constrain%selective_loc .and. wann_control%constrain%constrain) then
      lambda_loc = wann_control%constrain%lambda
    end if

    ! calculate initial centers and spread
    call wann_omega(csheet, sheet, rave, r2ave, rave2, wann_spread, num_wann, kmesh_info, &
                    num_kpts, print_output, wann_control%use_ss_functional, wann_control%constrain, &
                    omega%invariant, ln_tmp_loc, m_matrix_loc, lambda_loc, first_pass, timer, &
                    nkrank, global_k, error, comm)
    if (allocated(error)) return

    ! public variables
    if (.not. wann_control%constrain%selective_loc) then
      omega%total = wann_spread%om_tot
      omega%invariant = wann_spread%om_i
      omega%tilde = wann_spread%om_d + wann_spread%om_od
    else
      omega%total = wann_spread%om_tot
      ! omega_invariant = wann_spread%om_iod
      ! omega_tilde = wann_spread%om_d + wann_spread%om_nu
    end if

    ! public arrays of Wannier centres and spreads
    wannier_data%centres = rave
    wannier_data%spreads = r2ave - rave2

    if (wann_control%lfixstep) lquad = .false.

    ncg = 0
    iter = 0
    old_spread%om_tot = 0.0_dp

    ! print initial state
    if (print_output%iprint > 0) then
      write (stdout, '(1x,a78)') repeat('-', 78)
      write (stdout, '(1x,a)') 'Initial State'
      do iw = 1, num_wann
        write (stdout, 1000) iw, (rave(ind, iw)*print_output%lenconfac, ind=1, 3), &
          (r2ave(iw) - rave2(iw))*print_output%lenconfac**2
      end do
      write (stdout, 1001) (sum(rave(ind, :))*print_output%lenconfac, ind=1, 3), &
        (sum(r2ave) - sum(rave2))*print_output%lenconfac**2
      write (stdout, *)
      if (wann_control%constrain%selective_loc .and. wann_control%constrain%constrain) then
        write (stdout, '(1x,i6,2x,E12.3,2x,F15.10,2x,F18.10,3x,F8.2,2x,a)') iter, &
          (wann_spread%om_tot - old_spread%om_tot)*print_output%lenconfac**2, &
          sqrt(abs(gcnorm1))*print_output%lenconfac, &
          wann_spread%om_tot*print_output%lenconfac**2, io_wallclocktime(), '<-- CONV'
        write (stdout, '(7x,a,F15.7,a,F15.7,a,F15.7,a,F15.7,a)') &
          'O_D=', wann_spread%om_d*print_output%lenconfac**2, &
          ' O_IOD=', (wann_spread%om_iod + wann_spread%om_nu)*print_output%lenconfac**2, &
          ' O_TOT=', wann_spread%om_tot*print_output%lenconfac**2, ' <-- SPRD'
        write (stdout, '(1x,a78)') repeat('-', 78)
      elseif (wann_control%constrain%selective_loc .and. .not. wann_control%constrain%constrain) then
        write (stdout, '(1x,i6,2x,E12.3,2x,F15.10,2x,F18.10,3x,F8.2,2x,a)') iter, &
          (wann_spread%om_tot - old_spread%om_tot)*print_output%lenconfac**2, &
          sqrt(abs(gcnorm1))*print_output%lenconfac, &
          wann_spread%om_tot*print_output%lenconfac**2, io_wallclocktime(), '<-- CONV'
        write (stdout, '(7x,a,F15.7,a,F15.7,a,F15.7,a)') &
          'O_D=', wann_spread%om_d*print_output%lenconfac**2, &
          ' O_IOD=', wann_spread%om_iod*print_output%lenconfac**2, &
          ' O_TOT=', wann_spread%om_tot*print_output%lenconfac**2, ' <-- SPRD'
        write (stdout, '(1x,a78)') repeat('-', 78)
      else
        write (stdout, '(1x,i6,2x,E12.3,2x,F15.10,2x,F18.10,3x,F8.2,2x,a)') iter, &
          (wann_spread%om_tot - old_spread%om_tot)*print_output%lenconfac**2, &
          sqrt(abs(gcnorm1))*print_output%lenconfac, &
          wann_spread%om_tot*print_output%lenconfac**2, io_wallclocktime(), '<-- CONV'
        write (stdout, '(8x,a,F15.7,a,F15.7,a,F15.7,a)') &
          'O_D=', wann_spread%om_d*print_output%lenconfac**2, ' O_OD=', &
          wann_spread%om_od*print_output%lenconfac**2, &
          ' O_TOT=', wann_spread%om_tot*print_output%lenconfac**2, ' <-- SPRD'
        write (stdout, '(1x,a78)') repeat('-', 78)
      end if
    end if

    lconverged = .false.
    lfirst = .true.
    lrandom = .false.
    conv_count = 0
    noise_count = 0

    if (.not. wann_control%lfixstep .and. optimisation <= 0) then
      open (newunit=page_unit, status='scratch', form='unformatted')
    end if

    ! main iteration loop
    do iter = 1, wann_control%num_iter

      lprint = .false.
      if ((mod(iter, wann_control%num_print_cycles) .eq. 0) .or. (iter .eq. 1) &
          .or. (iter .eq. wann_control%num_iter)) lprint = .true.

      ldump = .false.
      if ((wann_control%num_dump_cycles .gt. 0) .and. &
          (mod(iter, wann_control%num_dump_cycles) .eq. 0)) ldump = .true.

      if (lprint .and. print_output%iprint > 0) write (stdout, '(1x,a,i6)') 'Cycle: ', iter

      if (wann_control%guiding_centres%enable .and. &
          (iter .gt. wann_control%guiding_centres%num_no_guide_iter) &
          .and. (mod(iter, wann_control%guiding_centres%num_guide_cycles) .eq. 0)) then
        call wann_phases(csheet, sheet, rguide, irguide, num_wann, kmesh_info, num_kpts, &
                         wann_control%use_ss_functional, m_matrix_loc, print_output%timing_level, &
                         print_output%iprint, timer, nkrank, global_k, error, comm)
        if (allocated(error)) return

        irguide = 1
      end if

      ! calculate gradient of omega
      if (lsitesymmetry .or. wann_control%precond) then
        call wann_domega(csheet, sheet, rave, num_wann, kmesh_info, num_kpts, &
                         wann_control%constrain, wann_control%use_ss_functional, lsitesymmetry, &
                         ln_tmp_loc, m_matrix_loc, rnkb_loc, cdodq_loc, lambda_loc, &
                         print_output%timing_level, sitesym, timer, nkrank, global_k, error, comm, &
                         print_output%iprint, cdodq)
        if (allocated(error)) return

      else
        call wann_domega(csheet, sheet, rave, num_wann, kmesh_info, num_kpts, &
                         wann_control%constrain, wann_control%use_ss_functional, lsitesymmetry, &
                         ln_tmp_loc, m_matrix_loc, rnkb_loc, cdodq_loc, lambda_loc, &
                         print_output%timing_level, sitesym, timer, nkrank, global_k, error, comm, &
                         print_output%iprint)
        if (allocated(error)) return

      end if

      if (lprint .and. print_output%iprint > 2) &
        write (stdout, *) ' LINE --> Iteration                     :', iter

      ! calculate search direction (cdq)
      if (wann_control%precond) then
        call precond_search_direction(cdodq, cdodq_r, cdodq_precond, cdodq_precond_loc, k_to_r, &
                                      wann_spread, num_wann, num_kpts, kpt_latt, real_lattice, &
                                      nrpts, irvec, ndegen, optimisation, timer)
      end if
      call internal_search_direction(cdodq_precond_loc, cdqkeep_loc, iter, lprint, lrandom, &
                                     noise_count, ncg, gcfac, gcnorm0, gcnorm1, doda0, &
                                     wann_control, num_wann, kmesh_info%wbtot, cdq_loc, cdodq_loc, &
                                     stdout, timer, error, comm)
      if (allocated(error)) return

      if (lsitesymmetry) then
        ! symmetrize_graident requires all k (IBZ and related FBZ points for each op)
        cdq(:, :, :) = 0.0_dp
        do nkp_loc = 1, nkrank
          nkp = global_k(nkp_loc)
          cdq(:, :, nkp) = cdq_loc(:, :, nkp_loc)
        end do
        call comms_allreduce(cdq(1, 1, 1), num_wann*num_wann*num_kpts, 'SUM', error, comm)

        ! called in parallel; alternatively broadcast
        call sitesym_symmetrize_gradient(sitesym, cdq, 2, num_kpts, num_wann, error, comm)
        do nkp_loc = 1, nkrank
          cdq_loc(:, :, nkp_loc) = cdq(:, :, global_k(nkp_loc))
        end do
      end if

      ! save search direction
      cdqkeep_loc(:, :, :) = cdq_loc(:, :, :)

      ! check whether we're doing fixed step lengths
      if (wann_control%lfixstep) then

        alphamin = wann_control%fixed_step

        ! or a parabolic line search
      else

        ! take trial step
        cdq_loc(:, :, :) = cdqkeep_loc(:, :, :)*(wann_control%trial_step/(4.0_dp*kmesh_info%wbtot))

        ! store original U and M before rotating
        u0_loc = u_matrix_loc

        if (optimisation <= 0) then
          write (page_unit) m_matrix_loc
          rewind (page_unit)
        else
          m0_loc = m_matrix_loc
        end if

        ! update U and M
        call internal_new_u_and_m(cdq, cmtmp, tmp_cdq, cwork, rwork, evals, cwschur1, cwschur2, &
                                  cwschur3, cwschur4, cz, num_wann, num_kpts, kmesh_info, &
                                  lsitesymmetry, cdq_loc, u_matrix_loc, m_matrix_loc, &
                                  print_output%timing_level, stdout, sitesym, timer, nkrank, &
                                  global_k, error, comm)
        if (allocated(error)) return

        ! calculate spread at trial step
        call wann_omega(csheet, sheet, rave, r2ave, rave2, trial_spread, num_wann, kmesh_info, &
                        num_kpts, print_output, wann_control%use_ss_functional, wann_control%constrain, &
                        omega%invariant, ln_tmp_loc, m_matrix_loc, lambda_loc, first_pass, timer, &
                        nkrank, global_k, error, comm)
        if (allocated(error)) return

        ! Calculate optimal step (alphamin)
        call internal_optimal_step(wann_spread, trial_spread, doda0, alphamin, falphamin, lquad, &
                                   lprint, wann_control%trial_step, stdout, timer)
      end if

      ! print line search information
      if (lprint .and. print_output%iprint > 2) then
        write (stdout, *) ' LINE --> Spread at initial point       :', &
          wann_spread%om_tot*print_output%lenconfac**2
        if (.not. wann_control%lfixstep) &
          write (stdout, *) ' LINE --> Spread at trial step          :', &
          trial_spread%om_tot*print_output%lenconfac**2
        write (stdout, *) ' LINE --> Slope along search direction  :', &
          doda0*print_output%lenconfac**2
        write (stdout, *) ' LINE --> ||SD gradient||^2             :', &
          gcnorm1*print_output%lenconfac**2
        if (.not. wann_control%lfixstep) then
          write (stdout, *) ' LINE --> Trial step length             :', wann_control%trial_step
          if (lquad) then
            write (stdout, *) ' LINE --> Optimal parabolic step length :', alphamin
            write (stdout, *) ' LINE --> Spread at predicted minimum   :', &
              falphamin*print_output%lenconfac**2
          end if
        else
          write (stdout, *) ' LINE --> Fixed step length             :', wann_control%fixed_step
        end if
        write (stdout, *) ' LINE --> CG coefficient                :', gcfac
      end if

      ! if taking a fixed step or if parabolic line search was successful
      if (wann_control%lfixstep .or. lquad) then

        ! take optimal step
        cdq_loc(:, :, :) = cdqkeep_loc(:, :, :)*(alphamin/(4.0_dp*kmesh_info%wbtot))

        ! if doing a line search then restore original U and M before rotating
        if (.not. wann_control%lfixstep) then
          u_matrix_loc = u0_loc
          if (optimisation <= 0) then
            read (page_unit) m_matrix_loc
            rewind (page_unit)
          else
            m_matrix_loc = m0_loc
          end if
        end if

        ! update U and M
        call internal_new_u_and_m(cdq, cmtmp, tmp_cdq, cwork, rwork, evals, cwschur1, cwschur2, &
                                  cwschur3, cwschur4, cz, num_wann, num_kpts, kmesh_info, &
                                  lsitesymmetry, cdq_loc, u_matrix_loc, m_matrix_loc, &
                                  print_output%timing_level, stdout, sitesym, timer, nkrank, &
                                  global_k, error, comm)
        if (allocated(error)) return

        call wann_spread_copy(wann_spread, old_spread)

        ! calculate the new centers and spread
        call wann_omega(csheet, sheet, rave, r2ave, rave2, wann_spread, num_wann, kmesh_info, &
                        num_kpts, print_output, wann_control%use_ss_functional, wann_control%constrain, &
                        omega%invariant, ln_tmp_loc, m_matrix_loc, lambda_loc, first_pass, timer, &
                        nkrank, global_k, error, comm)
        if (allocated(error)) return

        ! parabolic line search was unsuccessful, use trial step already taken
      else

        call wann_spread_copy(wann_spread, old_spread)
        call wann_spread_copy(trial_spread, wann_spread)

      end if

      ! print the new centers and spreads
      if (lprint .and. print_output%iprint > 0) then
        do iw = 1, num_wann
          write (stdout, 1000) iw, (rave(ind, iw)*print_output%lenconfac, ind=1, 3), &
            (r2ave(iw) - rave2(iw))*print_output%lenconfac**2
        end do
        write (stdout, 1001) (sum(rave(ind, :))*print_output%lenconfac, ind=1, 3), &
          (sum(r2ave) - sum(rave2))*print_output%lenconfac**2
        write (stdout, *)
        if (wann_control%constrain%selective_loc .and. wann_control%constrain%constrain) then
          write (stdout, '(1x,i6,2x,E12.3,2x,F15.10,2x,F18.10,3x,F8.2,2x,a)') &
            iter, (wann_spread%om_tot - old_spread%om_tot)*print_output%lenconfac**2, &
            sqrt(abs(gcnorm1))*print_output%lenconfac, &
            wann_spread%om_tot*print_output%lenconfac**2, io_wallclocktime(), '<-- CONV'
          write (stdout, '(7x,a,F15.7,a,F15.7,a,F15.7,a)') &
            'O_IOD=', (wann_spread%om_iod + wann_spread%om_nu)*print_output%lenconfac**2, &
            ' O_D=', wann_spread%om_d*print_output%lenconfac**2, &
            ' O_TOT=', wann_spread%om_tot*print_output%lenconfac**2, ' <-- SPRD'
          write (stdout, '(a,E15.7,a,E15.7,a,E15.7,a)') &
            'Delta: O_IOD=', ((wann_spread%om_iod + wann_spread%om_nu) - &
                              (old_spread%om_iod + wann_spread%om_nu))*print_output%lenconfac**2, &
            ' O_D=', (wann_spread%om_d - old_spread%om_d)*print_output%lenconfac**2, &
            ' O_TOT=', (wann_spread%om_tot - old_spread%om_tot)*print_output%lenconfac**2, ' <-- DLTA'
          write (stdout, '(1x,a78)') repeat('-', 78)
        elseif (wann_control%constrain%selective_loc .and. .not. wann_control%constrain%constrain) then
          write (stdout, '(1x,i6,2x,E12.3,2x,F15.10,2x,F18.10,3x,F8.2,2x,a)') &
            iter, (wann_spread%om_tot - old_spread%om_tot)*print_output%lenconfac**2, &
            sqrt(abs(gcnorm1))*print_output%lenconfac, &
            wann_spread%om_tot*print_output%lenconfac**2, io_wallclocktime(), '<-- CONV'
          write (stdout, '(7x,a,F15.7,a,F15.7,a,F15.7,a)') &
            'O_IOD=', wann_spread%om_iod*print_output%lenconfac**2, &
            ' O_D=', wann_spread%om_d*print_output%lenconfac**2, &
            ' O_TOT=', wann_spread%om_tot*print_output%lenconfac**2, ' <-- SPRD'
          write (stdout, '(a,E15.7,a,E15.7,a,E15.7,a)') &
            'Delta: O_IOD=', (wann_spread%om_iod - old_spread%om_iod)*print_output%lenconfac**2, &
            ' O_D=', (wann_spread%om_d - old_spread%om_d)*print_output%lenconfac**2, &
            ' O_TOT=', (wann_spread%om_tot - old_spread%om_tot)*print_output%lenconfac**2, ' <-- DLTA'
          write (stdout, '(1x,a78)') repeat('-', 78)
        else
          write (stdout, '(1x,i6,2x,E12.3,2x,F15.10,2x,F18.10,3x,F8.2,2x,a)') &
            iter, (wann_spread%om_tot - old_spread%om_tot)*print_output%lenconfac**2, &
            sqrt(abs(gcnorm1))*print_output%lenconfac, &
            wann_spread%om_tot*print_output%lenconfac**2, io_wallclocktime(), '<-- CONV'
          write (stdout, '(8x,a,F15.7,a,F15.7,a,F15.7,a)') &
            'O_D=', wann_spread%om_d*print_output%lenconfac**2, &
            ' O_OD=', wann_spread%om_od*print_output%lenconfac**2, &
            ' O_TOT=', wann_spread%om_tot*print_output%lenconfac**2, ' <-- SPRD'
          write (stdout, '(1x,a,E15.7,a,E15.7,a,E15.7,a)') &
            'Delta: O_D=', (wann_spread%om_d - old_spread%om_d)*print_output%lenconfac**2, &
            ' O_OD=', (wann_spread%om_od - old_spread%om_od)*print_output%lenconfac**2, &
            ' O_TOT=', (wann_spread%om_tot - old_spread%om_tot)*print_output%lenconfac**2, ' <-- DLTA'
          write (stdout, '(1x,a78)') repeat('-', 78)
        end if
      end if

      ! Public array of Wannier centres and spreads
      wannier_data%centres = rave
      wannier_data%spreads = r2ave - rave2

      ! Public variables
      if (.not. wann_control%constrain%selective_loc) then
        omega%total = wann_spread%om_tot
        omega%tilde = wann_spread%om_d + wann_spread%om_od
      else
        omega%total = wann_spread%om_tot
        !omega_tilde = wann_spread%om_d + wann_spread%om_nu
      end if

!JJ      if (ldump) then
!JJ        ! Before calling w90_wannier90_readwrite_write_chkpt, I need to gather on the root node
!JJ        ! the u_matrix from the u_matrix_loc. No need to broadcast it since
!JJ        ! it's printed by the root node only
!JJ        u_matrix(:, :, :) = 0.0_dp
!JJ        m_matrix(:, :, :, :) = 0.0_dp
!JJ        do nkp_loc = 1, nkrank
!JJ          nkp = displs(my_node_id) + nkp_loc
!JJ          u_matrix(:, :, nkp) = u_matrix_loc(:, :, nkp_loc)
!JJ          m_matrix(:, :, :, nkp) = m_matrix_loc(:, :, :, nkp_loc)
!JJ        enddo
!JJ
!JJ        wwk = num_wann*num_wann*num_kpts
!JJ        call comms_reduce(u_matrix(1, 1, 1), wwk, 'SUM', error, comm)
!JJ        if (allocated(error)) return
!JJ        call comms_reduce(m_matrix(1, 1, 1, 1), wwk*kmesh_info%nntot, 'SUM', error, comm)
!JJ        if (allocated(error)) return
!JJ
!JJ        if (on_root) then
!JJ          call w90_wannier90_readwrite_write_chkpt('postdis', exclude_bands, wannier_data, &
!JJ                                                   kmesh_info, kpt_latt, num_kpts, dis_manifold, &
!JJ                                                   num_bands, num_wann, u_matrix, u_matrix_opt, &
!JJ                                                   m_matrix, mp_grid, real_lattice, &
!JJ                                                   omega%invariant, have_disentangled, stdout, &
!JJ                                                   seedname)
!JJ        endif
!JJ      endif

      if (wann_control%conv_window .gt. 1) then
        call internal_test_convergence(old_spread, wann_spread, history, save_spread, iter, &
                                       conv_count, noise_count, lconverged, lrandom, lfirst, &
                                       wann_control, error, comm)
        if (allocated(error)) return
      end if

      if (lconverged) then
        if (print_output%iprint > 0) then
          write (stdout, '(/13x,a,es10.3,a,i2,a)') '<<<     Delta <', wann_control%conv_tol, &
            '  over ', wann_control%conv_window, ' iterations     >>>'
          write (stdout, '(13x,a/)') '<<< Wannierisation convergence criteria satisfied >>>'
        end if
        exit
      end if

    end do
    ! end of the minimization loop

    ! copy from local u matrix back to full matrix & reduce
    u_matrix(:, :, :) = 0.0_dp
    do nkp_loc = 1, nkrank
      nkp = global_k(nkp_loc)
      u_matrix(:, :, nkp) = u_matrix_loc(:, :, nkp_loc)
    end do
    call comms_allreduce(u_matrix(1, 1, 1), num_wann*num_wann*num_kpts, 'SUM', error, comm)
    if (allocated(error)) return

    ! Evaluate the penalty functional
    if (wann_control%constrain%selective_loc .and. wann_control%constrain%constrain) then
      rnr0n2 = 0.0_dp
      do iw = 1, wann_control%constrain%slwf_num
        rnr0n2(iw) = (wannier_data%centres(1, iw) - wann_control%constrain%centres(iw, 1))**2 &
                     + (wannier_data%centres(2, iw) - wann_control%constrain%centres(iw, 2))**2 &
                     + (wannier_data%centres(3, iw) - wann_control%constrain%centres(iw, 3))**2
      end do
    end if

    if (print_output%iprint > 0) then
      write (stdout, '(1x,a)') 'Final State'
      do iw = 1, num_wann
        write (stdout, 1000) iw, (rave(ind, iw)*print_output%lenconfac, ind=1, 3), &
          (r2ave(iw) - rave2(iw))*print_output%lenconfac**2
      end do
      write (stdout, 1001) (sum(rave(ind, :))*print_output%lenconfac, ind=1, 3), &
        (sum(r2ave) - sum(rave2))*print_output%lenconfac**2
      write (stdout, *)
      if (wann_control%constrain%selective_loc .and. wann_control%constrain%constrain) then
        write (stdout, '(3x,a21,a,f15.9)') '     Spreads ('//trim(print_output%length_unit)//'^2)', &
          '       Omega IOD_C   = ', (wann_spread%om_iod + wann_spread%om_nu)*print_output%lenconfac**2
        write (stdout, '(3x,a,f15.9)') '     ================       Omega D       = ', &
          wann_spread%om_d*print_output%lenconfac**2
        write (stdout, '(3x,a,f15.9)') '                            Omega Rest    = ', &
          (sum(r2ave) - sum(rave2) + wann_spread%om_tot)*print_output%lenconfac**2
        write (stdout, '(3x,a,f15.9)') '                            Penalty func  = ', &
          sum(rnr0n2(:))
        write (stdout, '(3x,a21,a,f15.9)') 'Final Spread ('//trim(print_output%length_unit)//'^2)', &
          '       Omega Total_C = ', wann_spread%om_tot*print_output%lenconfac**2
        write (stdout, '(1x,a78)') repeat('-', 78)
      else if (wann_control%constrain%selective_loc .and. .not. wann_control%constrain%constrain) then
        write (stdout, '(3x,a21,a,f15.9)') '     Spreads ('//trim(print_output%length_unit)//'^2)', &
          '       Omega IOD    = ', wann_spread%om_iod*print_output%lenconfac**2
        write (stdout, '(3x,a,f15.9)') '     ================       Omega D      = ', &
          wann_spread%om_d*print_output%lenconfac**2
        write (stdout, '(3x,a,f15.9)') '                            Omega Rest   = ', &
          (sum(r2ave) - sum(rave2) + wann_spread%om_tot)*print_output%lenconfac**2
        write (stdout, '(3x,a21,a,f15.9)') 'Final Spread ('//trim(print_output%length_unit)//'^2)', &
          '       Omega Total  = ', wann_spread%om_tot*print_output%lenconfac**2
        write (stdout, '(1x,a78)') repeat('-', 78)
      else
        write (stdout, '(3x,a21,a,f15.9)') '     Spreads ('//trim(print_output%length_unit)//'^2)', &
          '       Omega I      = ', wann_spread%om_i*print_output%lenconfac**2
        write (stdout, '(3x,a,f15.9)') '     ================       Omega D      = ', &
          wann_spread%om_d*print_output%lenconfac**2
        write (stdout, '(3x,a,f15.9)') '                            Omega OD     = ', &
          wann_spread%om_od*print_output%lenconfac**2
        write (stdout, '(3x,a21,a,f15.9)') 'Final Spread ('//trim(print_output%length_unit)//'^2)', &
          '       Omega Total  = ', wann_spread%om_tot*print_output%lenconfac**2
        write (stdout, '(1x,a78)') repeat('-', 78)
      end if
    end if

    if (wann_control%guiding_centres%enable) then
      call wann_phases(csheet, sheet, rguide, irguide, num_wann, kmesh_info, num_kpts, &
                       wann_control%use_ss_functional, m_matrix_loc, print_output%timing_level, &
                       print_output%iprint, timer, nkrank, global_k, error, comm)
      if (allocated(error)) return
    end if

    ! check unitarity of u
    call wann_check_unitarity(num_kpts, num_wann, u_matrix, print_output%timing_level, &
                              print_output%iprint, stdout, timer, error, comm)
    if (allocated(error)) return

    ! deallocate sub vars not passed into other subs
    deallocate (rwork, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating rwork in wann_main', comm)
      return
    end if
    deallocate (cwork, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating cwork in wann_main', comm)
      return
    end if
    deallocate (evals, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating evals in wann_main', comm)
      return
    end if
    deallocate (tmp_cdq, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating tmp_cdq in wann_main', comm)
      return
    end if
    deallocate (cmtmp, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating cmtmp in wann_main', comm)
      return
    end if
    deallocate (cz, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating cz in wann_main', comm)
      return
    end if
    deallocate (cdq, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating cdq in wann_main', comm)
      return
    end if
    deallocate (ln_tmp_loc, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating ln_tmp_loc in wann_main', comm)
      return
    end if
    deallocate (rnkb_loc, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating rnkb_loc in wann_main', comm)
      return
    end if
    deallocate (u_matrix_loc, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating u_matrix_loc in wann_main', comm)
      return
    end if
    deallocate (cdq_loc, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating cdq_loc in wann_main', comm)
      return
    end if
    deallocate (cdodq_loc, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating cdodq_loc in wann_main', comm)
      return
    end if
    deallocate (cdqkeep_loc, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating cdqkeep_loc in wann_main', comm)
      return
    end if
    deallocate (cwschur3, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating cwschur3 in wann_main', comm)
      return
    end if
    deallocate (cwschur1, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating cwschur1 in wann_main', comm)
      return
    end if
    if (wann_control%precond) then
      if (optimisation >= 3) then
        deallocate (k_to_r, stat=ierr)
        if (ierr /= 0) then
          call set_error_dealloc(error, 'Error in deallocating k_to_r in wann_main', comm)
          return
        end if
      end if
      deallocate (cdodq_r, stat=ierr)
      if (ierr /= 0) then
        call set_error_dealloc(error, 'Error in deallocating cdodq_r in wann_main', comm)
        return
      end if
      deallocate (cdodq_precond, stat=ierr)
      if (ierr /= 0) then
        call set_error_dealloc(error, 'Error in deallocating cdodq_precond in wann_main', comm)
        return
      end if
      deallocate (cdodq_precond_loc, stat=ierr)
      if (ierr /= 0) then
        call set_error_dealloc(error, 'Error in deallocating cdodq_precond_loc in wann_main', comm)
        return
      end if
    end if

    ! deallocate sub vars passed into other subs
    deallocate (rguide, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating rguide in wann_main', comm)
      return
    end if
    deallocate (rave2, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating rave2 in wann_main', comm)
      return
    end if
    deallocate (rave, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating rave in wann_main', comm)
      return
    end if
    deallocate (sheet, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating sheet in wann_main', comm)
      return
    end if
    deallocate (cdodq, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating cdodq in wann_main', comm)
      return
    end if
    deallocate (csheet, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating csheet in wann_main', comm)
      return
    end if
    if (wann_control%constrain%selective_loc) then
      deallocate (rnr0n2, stat=ierr)
      if (ierr /= 0) then
        call set_error_dealloc(error, 'Error in deallocating rnr0n2 in wann_main', comm)
        return
      end if
    end if
    deallocate (ln_tmp, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating ln_tmp in wann_main', comm)
      return
    end if
    deallocate (u0_loc, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error in deallocating u0_loc in wann_main', comm)
      return
    end if
    if (optimisation > 0) then
      deallocate (m0_loc, stat=ierr)
      if (ierr /= 0) then
        call set_error_dealloc(error, 'Error in deallocating m0_loc in wann_main', comm)
        return
      end if
    end if
    deallocate (history, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error deallocating history in wann_main', comm)
      return
    end if

    if (print_output%timing_level > 0 .and. print_output%iprint > 0) then
      call io_stopwatch_stop('wann: main', timer)
    end if

    if (.not. wann_control%lfixstep .and. optimisation <= 0) close (page_unit) !close scratch file

    return

1000 format(2x, 'WF centre and spread', i5, 2x, '(', f10.6, ',', f10.6, ',', f10.6, ' )', f15.8)
1001 format(2x, 'Sum of centres and spreads', 1x, '(', f10.6, ',', f10.6, ',', f10.6, ' )', f15.8)

  contains

    !================================================!
    subroutine internal_random_noise(conv_noise_amp, num_wann, nkrank, cdq_loc)
      !================================================!
      !
      !! Add some random noise to the search direction
      !! to help escape from local minima
      !
      !================================================!
      use w90_constants, only: cmplx_0
      use w90_comms, only: w90_comm_type

      implicit none

      ! arguments
      integer, intent(in) :: num_wann
      integer, intent(in) :: nkrank
      real(kind=dp), intent(in) :: conv_noise_amp
      complex(kind=dp), intent(inout) :: cdq_loc(:, :, :)

      ! local
      integer :: ikp, iw, jw, ierr
      real(kind=dp), allocatable :: noise_real(:, :), noise_imag(:, :)
      complex(kind=dp), allocatable :: cnoise(:, :)

      ! Allocate
      allocate (noise_real(num_wann, num_wann), stat=ierr)
      if (ierr /= 0) then
        call set_error_alloc(error, 'Error allocating noise_real in wann_main: random_noise', comm)
        return
      end if
      allocate (noise_imag(num_wann, num_wann), stat=ierr)
      if (ierr /= 0) then
        call set_error_alloc(error, 'Error allocating noise_imag in wann_main: random_noise', comm)
        return
      end if
      allocate (cnoise(num_wann, num_wann), stat=ierr)
      if (ierr /= 0) then
        call set_error_alloc(error, 'Error allocating cnoise in wann_main: random_noise', comm)
        return
      end if

      ! Initialise
      cnoise = cmplx_0; noise_real = 0.0_dp; noise_imag = 0.0_dp

      ! cdq is a num_wann x num_wann x num_kpts anti-hermitian array
      ! to which we add a random anti-hermitian matrix

      do ikp = 1, nkrank
        do iw = 1, num_wann
          call random_seed()
          call random_number(noise_real(:, iw))
          call random_seed()
          call random_number(noise_imag(:, iw))
        end do
        do jw = 1, num_wann
          do iw = 1, jw
            if (iw .eq. jw) then
              cnoise(iw, jw) = cmplx(0.0_dp, noise_imag(iw, jw), dp)
            else
              cnoise(iw, jw) = cmplx(noise_real(iw, jw), noise_imag(iw, jw), dp)
            end if
            cnoise(jw, iw) = -conjg(cnoise(iw, jw))
          end do
        end do
        ! Add noise to search direction
        cdq_loc(:, :, ikp) = cdq_loc(:, :, ikp) + conv_noise_amp*cnoise(:, :)
      end do

      ! Deallocate
      deallocate (cnoise, stat=ierr)
      if (ierr /= 0) then
        call set_error_dealloc(error, 'Error deallocating cnoise in wann_main: random_noise', comm)
        return
      end if
      deallocate (noise_imag, stat=ierr)
      if (ierr /= 0) then
        call set_error_dealloc(error, 'Error deallocating noise_imag in wann_main: random_noise', comm)
        return
      end if
      deallocate (noise_real, stat=ierr)
      if (ierr /= 0) then
        call set_error_dealloc(error, 'Error deallocating noise_real in wann_main: random_noise', comm)
        return
      end if

      return

    end subroutine internal_random_noise

    !================================================!
    subroutine precond_search_direction(cdodq, cdodq_r, cdodq_precond, cdodq_precond_loc, k_to_r, &
                                        wann_spread, num_wann, num_kpts, kpt_latt, real_lattice, &
                                        nrpts, irvec, ndegen, optimisation, timer)
      !================================================!
      !
      !! Calculate the conjugate gradients search
      !! direction using the Fletcher-Reeves formula:
      !!
      !!     cg_coeff = [g(i).g(i)]/[g(i-1).g(i-1)]
      !
      !================================================!

      use w90_constants, only: cmplx_0, cmplx_1, cmplx_i, twopi
      use w90_io, only: io_stopwatch_start, io_stopwatch_stop
      use w90_types, only: timer_list_type

      implicit none

      ! arguments
      type(localisation_vars_type), intent(in) :: wann_spread
      type(timer_list_type), intent(inout) :: timer

      complex(kind=dp), intent(in) :: cdodq(:, :, :)
      complex(kind=dp), intent(inout) :: cdodq_r(:, :, :)
      complex(kind=dp), intent(inout) :: cdodq_precond(:, :, :)
      complex(kind=dp), intent(inout) :: cdodq_precond_loc(:, :, :)
      ! k_to_r depends on optimisation flag
      complex(kind=dp), allocatable, intent(in) :: k_to_r(:, :)

      real(kind=dp), intent(in) :: kpt_latt(:, :)
      real(kind=dp), intent(in) :: real_lattice(3, 3)

      integer, intent(in) :: num_wann, num_kpts
      integer, intent(in) :: nrpts
      integer, intent(in) :: irvec(:, :)
      integer, intent(in) :: ndegen(:)
      integer, intent(in) :: optimisation

      ! local
      complex(kind=dp) :: fac, rdotk
      real(kind=dp) :: rvec_cart(3)
      real(kind=dp) :: alpha_precond
      integer :: irpt, loop_kpt

      if (print_output%timing_level > 1 .and. print_output%iprint > 0) then
        call io_stopwatch_start('wann: main: search_direction', timer)
      end if

      ! gcnorm1 = Tr[gradient . gradient] -- NB gradient is anti-Hermitian
      ! gcnorm1 = real(zdotc(num_kpts*num_wann*num_wann,cdodq,1,cdodq,1),dp)

      cdodq_r(:, :, :) = 0 ! intermediary gradient in R space
      cdodq_precond(:, :, :) = 0
      cdodq_precond_loc(:, :, :) = 0

      ! convert to real space in cdodq_r
      ! Two algorithms: either double loop or GEMM. GEMM is much more efficient but requires more RAM
      ! Ideally, we should implement FFT-based filtering here
      if (optimisation >= 3) then
        call zgemm('N', 'N', num_wann*num_wann, nrpts, num_kpts, cmplx_1, cdodq, &
                   num_wann*num_wann, k_to_r, num_kpts, cmplx_0, cdodq_r, num_wann*num_wann)
        cdodq_r = cdodq_r/real(num_kpts, dp)
      else
        do irpt = 1, nrpts
          do loop_kpt = 1, num_kpts
            rdotk = twopi*dot_product(kpt_latt(:, loop_kpt), real(irvec(:, irpt), dp))
            fac = exp(-cmplx_i*rdotk)/real(num_kpts, dp)
            cdodq_r(:, :, irpt) = cdodq_r(:, :, irpt) + fac*cdodq(:, :, loop_kpt)
          end do
        end do
      end if

      ! filter cdodq_r in real space by 1/(1+R^2/alpha)

      ! this alpha coefficient is more or less arbitrary, and could
      ! be tweaked further: the point is to have something that has
      ! the right units, and is not too small (or the filtering is
      ! too severe) or too high (or the filtering does nothing).
      !
      ! the descent direction produced has a different magnitude
      ! than the one without preconditionner, so the values of
      ! trial_step are not consistent
      alpha_precond = 10*wann_spread%om_tot/num_wann
      do irpt = 1, nrpts
        rvec_cart = matmul(real_lattice(:, :), real(irvec(:, irpt), dp))
        cdodq_r(:, :, irpt) = cdodq_r(:, :, irpt)*1/(1 + dot_product(rvec_cart, rvec_cart)/ &
                                                     alpha_precond)
      end do

      ! go back to k space
      if (optimisation >= 3) then
        do irpt = 1, nrpts
          cdodq_r(:, :, irpt) = cdodq_r(:, :, irpt)/real(ndegen(irpt), dp)
        end do
        call zgemm('N', 'C', num_wann*num_wann, num_kpts, nrpts, cmplx_1, cdodq_r, &
                   num_wann*num_wann, k_to_r, num_kpts, cmplx_0, cdodq_precond, num_wann*num_wann)
      else
        do irpt = 1, nrpts
          do loop_kpt = 1, num_kpts
            rdotk = twopi*dot_product(kpt_latt(:, loop_kpt), real(irvec(:, irpt), dp))
            fac = exp(cmplx_i*rdotk)/real(ndegen(irpt), dp)
            cdodq_precond(:, :, loop_kpt) = cdodq_precond(:, :, loop_kpt) + &
                                            fac*cdodq_r(:, :, irpt)
          end do
        end do
      end if
      do nkp_loc = 1, nkrank
        nkp = global_k(nkp_loc)
        cdodq_precond_loc(:, :, nkp_loc) = cdodq_precond(:, :, nkp)
      end do

    end subroutine precond_search_direction

    !================================================!
    subroutine internal_search_direction(cdodq_precond_loc, cdqkeep_loc, iter, lprint, lrandom, &
                                         noise_count, ncg, gcfac, gcnorm0, gcnorm1, doda0, &
                                         wann_control, num_wann, wbtot, cdq_loc, cdodq_loc, &
                                         stdout, timer, error, comm)
      !================================================!
      !
      !! Calculate the conjugate gradients search
      !! direction using the Fletcher-Reeves formula:
      !!
      !!     cg_coeff = [g(i).g(i)]/[g(i-1).g(i-1)]
      !
      !================================================!
      use w90_io, only: io_stopwatch_start, io_stopwatch_stop
      use w90_comms, only: comms_allreduce, w90_comm_type
      use w90_wannier90_types, only: wann_control_type
      use w90_types, only: timer_list_type

      implicit none

      ! arguments
      type(wann_control_type), intent(in) :: wann_control
      type(timer_list_type), intent(inout) :: timer
      type(w90_error_type), allocatable, intent(out) :: error
      type(w90_comm_type), intent(in) :: comm

      integer, intent(in) :: iter
      integer, intent(in) :: noise_count
      integer, intent(in) :: num_wann
      integer, intent(inout) :: ncg
      integer, intent(in) :: stdout

      complex(kind=dp), allocatable, intent(inout) :: cdodq_precond_loc(:, :, :)
      complex(kind=dp), intent(in) :: cdodq_loc(:, :, :)
      complex(kind=dp), intent(inout) :: cdqkeep_loc(:, :, :)
      complex(kind=dp), intent(inout) :: cdq_loc(:, :, :)

      real(kind=dp), intent(inout) :: gcnorm0, gcnorm1
      real(kind=dp), intent(in) :: wbtot
      real(kind=dp), intent(out) :: doda0
      real(kind=dp), intent(out) :: gcfac

      logical, intent(in) :: lprint
      logical, intent(inout) :: lrandom

      ! local
      integer :: m
      complex(kind=dp) :: zres

      m = count(dist_k == mpirank(comm))*num_wann*num_wann ! for dimensioning

      if ((.not. wann_control%precond) .and. print_output%timing_level > 1 .and. print_output%iprint > 0) then
        call io_stopwatch_start('wann: main: search_direction', timer)
      end if

      ! gcnorm1 = Tr[gradient . gradient] -- NB gradient is anti-Hermitian
      gcnorm1 = 0
      if (m > 0) then
        if (wann_control%precond) then
          ! compute (zdotc) cdodq_precond_loc.cdodq_loc^c
          call zgemv('c', m, 1, cmplx_1, cdodq_precond_loc, m, cdodq_loc, 1, cmplx_0, zres, 1)
          gcnorm1 = real(zres, dp)
        else
          ! compute (zdotc) cdodq_loc.cdodq_loc^c
          call zgemv('c', m, 1, cmplx_1, cdodq_loc, m, cdodq_loc, 1, cmplx_0, zres, 1)
          gcnorm1 = real(zres, dp)
        end if
      end if
      call comms_allreduce(gcnorm1, 1, 'SUM', error, comm)
      if (allocated(error)) return

      ! calculate cg_coefficient
      if ((iter .eq. 1) .or. (ncg .ge. wann_control%num_cg_steps)) then
        gcfac = 0.0_dp                 ! Steepest descents
        ncg = 0
      else
        if (gcnorm0 .gt. epsilon(1.0_dp)) then
          gcfac = gcnorm1/gcnorm0     ! Fletcher-Reeves CG coefficient
          ! prevent CG coefficient from getting too large
          if (gcfac .gt. 3.0_dp) then
            if (lprint .and. print_output%iprint > 2) &
              write (stdout, *) ' LINE --> CG coeff too large. Resetting :', gcfac
            gcfac = 0.0_dp
            ncg = 0
          else
            ncg = ncg + 1
          end if
        else
          gcfac = 0.0_dp
          ncg = 0
        end if
      end if

      ! save for next iteration
      gcnorm0 = gcnorm1

      ! calculate search direction
      if (wann_control%precond) then
        cdq_loc(:, :, :) = cdodq_precond_loc(:, :, :) + cdqkeep_loc(:, :, :)*gcfac !! JRY not MPI
      else
        cdq_loc(:, :, :) = cdodq_loc(:, :, :) + cdqkeep_loc(:, :, :)*gcfac
      end if

      ! add some random noise to search direction, if required
      if (lrandom) then
        if (print_output%iprint > 0) write (stdout, '(a,i3,a,i3,a)') &
          ' [ Adding random noise to search direction. Time ', noise_count, ' / ', &
          wann_control%conv_noise_num, ' ]'
        call internal_random_noise(wann_control%conv_noise_amp, num_wann, nkrank, cdq_loc)
      end if

      ! calculate gradient along search direction - Tr[gradient . search direction]
      ! NB gradient is anti-hermitian
      ! compute (zdotc) cdodq_loc.cdq_loc^c
      zres = 0
      if (m > 0) call zgemv('c', m, 1, cmplx_1, cdodq_loc, m, cdq_loc, 1, cmplx_0, zres, 1)
      doda0 = -real(zres, dp)

      call comms_allreduce(doda0, 1, 'SUM', error, comm)
      if (allocated(error)) return

      doda0 = doda0/(4.0_dp*wbtot)

      ! check search direction is not uphill
      if (doda0 .gt. 0.0_dp) then
        ! if doing a CG step then reset CG
        if (ncg .gt. 0) then
          if (lprint .and. print_output%iprint > 2 .and. print_output%iprint > 0) &
            write (stdout, *) ' LINE --> Search direction uphill: resetting CG'
          cdq_loc(:, :, :) = cdodq_loc(:, :, :)
          if (lrandom) then
            call internal_random_noise(wann_control%conv_noise_amp, num_wann, nkrank, cdq_loc)
          end if
          ncg = 0
          gcfac = 0.0_dp

          ! re-calculate gradient along search direction
          ! compute (zdotc) cdodq_loc.cdq_loc^c
          zres = 0
          if (m > 0) call zgemv('c', m, 1, cmplx_1, cdodq_loc, m, cdq_loc, 1, cmplx_0, zres, 1)
          doda0 = -real(zres, dp)

          call comms_allreduce(doda0, 1, 'SUM', error, comm)
          if (allocated(error)) return
          doda0 = doda0/(4.0_dp*wbtot)

          ! if search direction still uphill then reverse search direction
          if (doda0 .gt. 0.0_dp) then
            if (lprint .and. print_output%iprint > 2 .and. print_output%iprint > 0) &
              write (stdout, *) ' LINE --> Search direction still uphill: reversing'
            cdq_loc(:, :, :) = -cdq_loc(:, :, :)
            doda0 = -doda0
          end if
          ! if doing a SD step then reverse search direction
        else
          if (lprint .and. print_output%iprint > 2 .and. print_output%iprint > 0) &
            write (stdout, *) ' LINE --> Search direction uphill: reversing'
          cdq_loc(:, :, :) = -cdq_loc(:, :, :)
          doda0 = -doda0
        end if
      end if

      if (print_output%timing_level > 1 .and. print_output%iprint > 0) then
        call io_stopwatch_stop('wann: main: search_direction', timer)
      end if

      lrandom = .false.
      return
    end subroutine internal_search_direction

    !================================================!
    subroutine internal_optimal_step(wann_spread, trial_spread, doda0, alphamin, falphamin, lquad, &
                                     lprint, trial_step, stdout, timer)
      !================================================!
      !
      !! Calculate the optimal step length based on a
      !! parabolic line search
      !
      !================================================!
      use w90_io, only: io_stopwatch_start, io_stopwatch_stop
      use w90_comms, only: w90_comm_type
      use w90_types, only: timer_list_type

      implicit none

      type(localisation_vars_type), intent(in) :: trial_spread
      type(localisation_vars_type), intent(in) :: wann_spread
      type(timer_list_type), intent(inout) :: timer
      integer, intent(in) :: stdout
      real(kind=dp), intent(in) :: doda0
      real(kind=dp), intent(in) :: trial_step
      real(kind=dp), intent(out) :: alphamin, falphamin
      logical, intent(out) :: lquad
      logical, intent(in) :: lprint

      ! local variables
      real(kind=dp) :: fac, shift, eqa, eqb

      if (print_output%timing_level > 1 .and. print_output%iprint > 0) then
        call io_stopwatch_start('wann: main: optimal_step', timer)
      end if

      fac = trial_spread%om_tot - wann_spread%om_tot
      if (abs(fac) .gt. tiny(1.0_dp)) then
        fac = 1.0_dp/fac
        shift = 1.0_dp
      else
        fac = 1.0e6_dp
        shift = fac*trial_spread%om_tot - fac*wann_spread%om_tot
      end if
      eqb = fac*doda0
      eqa = shift - eqb*trial_step
      if (abs(eqa/(fac*wann_spread%om_tot)) .gt. epsilon(1.0_dp)) then
        lquad = .true.
        alphamin = -0.5_dp*eqb/eqa*(trial_step**2)
        falphamin = wann_spread%om_tot &
                    - 0.25_dp*eqb*eqb/(fac*eqa)*(trial_step**2)
      else
        if (lprint .and. print_output%iprint > 2) write (stdout, *) &
          ' LINE --> Parabolic line search unstable: using trial step'
        lquad = .false.
        alphamin = trial_step
        falphamin = trial_spread%om_tot
      end if

      if (doda0*alphamin .gt. 0.0_dp) then
        if (lprint .and. print_output%iprint > 2) write (stdout, *) &
          ' LINE --> Line search unstable : using trial step'
        lquad = .false.
        alphamin = trial_step
        falphamin = trial_spread%om_tot
      end if

      if (print_output%timing_level > 1 .and. print_output%iprint > 0) then
        call io_stopwatch_stop('wann: main: optimal_step', timer)
      end if

      return

    end subroutine internal_optimal_step

    !================================================!
    subroutine internal_new_u_and_m(cdq, cmtmp, tmp_cdq, cwork, rwork, evals, cwschur1, cwschur2, &
                                    cwschur3, cwschur4, cz, num_wann, num_kpts, kmesh_info, &
                                    lsitesymmetry, cdq_loc, u_matrix_loc, m_matrix_loc, &
                                    timing_level, stdout, sitesym, timer, nkrank, global_k, error, &
                                    comm)
      !================================================!
      !
      !! Update U and M matrices after a trial step
      !
      !================================================!

      use w90_constants, only: cmplx_i
      use w90_sitesym, only: sitesym_symmetrize_rotation
      use w90_wannier90_types, only: sitesym_type
      use w90_io, only: io_stopwatch_start, io_stopwatch_stop
      use w90_comms, only: comms_allreduce, w90_comm_type
      use w90_utility, only: utility_zgemm
      use w90_types, only: kmesh_info_type, timer_list_type

      implicit none

      type(kmesh_info_type), intent(in) :: kmesh_info
      type(sitesym_type), intent(in) :: sitesym
      type(timer_list_type), intent(inout) :: timer
      type(w90_error_type), allocatable, intent(out) :: error
      type(w90_comm_type), intent(in) :: comm

      complex(kind=dp), intent(inout) :: cdq(:, :, :)
      complex(kind=dp), intent(inout) :: cdq_loc(:, :, :)
      complex(kind=dp), intent(inout) :: cmtmp(:, :), tmp_cdq(:, :) ! really just local?
      complex(kind=dp), intent(inout) :: cwork(:)
      complex(kind=dp), intent(inout) :: cwschur1(:), cwschur2(:)
      complex(kind=dp), intent(inout) :: cwschur3(:), cwschur4(:)
      complex(kind=dp), intent(inout) :: cz(:, :)
      complex(kind=dp), intent(inout) :: m_matrix_loc(:, :, :, :)
      complex(kind=dp), intent(inout) :: u_matrix_loc(:, :, :)

      integer, intent(in) :: nkrank
      integer, intent(in) :: global_k(:)
      integer, intent(in) :: timing_level
      integer, intent(in) :: num_wann, num_kpts
      integer, intent(in) :: stdout
      logical, intent(in) :: lsitesymmetry

      real(kind=dp), intent(inout) :: evals(:)
      real(kind=dp), intent(inout) :: rwork(:)

      ! local vars
      integer :: i, nkp, nn, nkp2, nsdim, nkp_loc, info
      logical :: ltmp
      integer :: my_node_id

      my_node_id = mpirank(comm)

      if (timing_level > 1 .and. print_output%iprint > 0) call io_stopwatch_start('wann: main: u_and_m', timer)

      do nkp_loc = 1, nkrank
        nkp = global_k(nkp_loc)
        if (lsitesymmetry) then
          if (sitesym%ir2ik(sitesym%ik2ir(nkp)) .ne. nkp) cycle
        end if

        ! cdq(nkp) is anti-Hermitian; tmp_cdq = i*cdq  is Hermitian
        tmp_cdq(:, :) = cmplx_i*cdq_loc(:, :, nkp_loc)
        ! Hermitian matrix eigen-solver
        call zheev('V', 'U', num_wann, tmp_cdq, num_wann, evals, cwork, 4*num_wann, rwork, info)
        if (info .ne. 0) then
          if (print_output%iprint > 0) write (stdout, *) &
            'wann_main: ZHEEV in internal_new_u_and_m failed, info= ', info
          if (print_output%iprint > 0) write (stdout, *) '           trying Schur decomposition instead'
          tmp_cdq(:, :) = cdq_loc(:, :, nkp_loc)
          call zgees('V', 'N', ltmp, num_wann, tmp_cdq, num_wann, nsdim, &
                     cwschur1, cz, num_wann, cwschur2, 10*num_wann, cwschur3, &
                     cwschur4, info)
          if (info .ne. 0) then
            if (print_output%iprint > 0) write (stdout, *) 'wann_main: SCHUR failed, info= ', info
            call set_error_fatal(error, 'wann_main: problem computing schur form 1', comm)
            return
          end if
          do i = 1, num_wann
            tmp_cdq(:, i) = cz(:, i)*exp(cwschur1(i))
          end do
          ! cmtmp   = tmp_cdq . cz^{dagger}
          call utility_zgemm(cmtmp, tmp_cdq, 'N', cz, 'C', num_wann)
          cdq_loc(:, :, nkp_loc) = cmtmp(:, :)
        else
          do i = 1, num_wann
            cmtmp(:, i) = tmp_cdq(:, i)*exp(-cmplx_i*evals(i))
          end do
          ! cdq(nkp)   = cmtmp . tmp_cdq^{dagger}
          call utility_zgemm(cdq_loc(:, :, nkp_loc), cmtmp, 'N', tmp_cdq, 'C', num_wann)
        end if
      end do

      ! each process communicates its result to other processes
      cdq(:, :, :) = 0.0_dp
      do nkp_loc = 1, nkrank
        nkp = global_k(nkp_loc)
        cdq(:, :, nkp) = cdq_loc(:, :, nkp_loc)
      end do
      call comms_allreduce(cdq(1, 1, 1), num_wann*num_wann*num_kpts, 'SUM', error, comm)
      if (allocated(error)) return

!!$      do nkp = 1, num_kpts
!!$         tmp_cdq(:,:) = cdq(:,:,nkp)
!!$         call zgees ('V', 'N', ltmp, num_wann, tmp_cdq, num_wann, nsdim, &
!!$              cwschur1, cz, num_wann, cwschur2, 10 * num_wann, cwschur3, &
!!$              cwschur4, info)
!!$         if (info.ne.0) then
!!$            write(stdout,*) 'SCHUR: ', info
!!$            call io_error('wann_main: problem computing schur form 1')
!!$         endif
!!$         do i=1,num_wann
!!$            tmp_cdq(:,i) = cz(:,i) * exp(cwschur1(i))
!!$         enddo
!!$         ! cmtmp   = tmp_cdq . cz^{dagger}
!!$         call utility_zgemm(cmtmp,tmp_cdq,'N',cz,'C',num_wann)
!!$         cdq(:,:,nkp)=cmtmp(:,:)
!!$      enddo

      if (lsitesymmetry) then
        call sitesym_symmetrize_rotation(sitesym, cdq, num_kpts, num_wann, error, comm)
        if (allocated(error)) return
        do nkp_loc = 1, nkrank
          nkp = global_k(nkp_loc)
          cdq_loc(:, :, nkp_loc) = cdq(:, :, nkp)
        end do
      end if

      ! the orbitals are rotated
      do nkp_loc = 1, nkrank
        ! cmtmp = U(k) . cdq(k)
        call utility_zgemm(cmtmp, u_matrix_loc(:, :, nkp_loc), 'N', cdq_loc(:, :, nkp_loc), 'N', &
                           num_wann)
        u_matrix_loc(:, :, nkp_loc) = cmtmp(:, :)
      end do

      ! and the M_ij are updated
      do nkp_loc = 1, nkrank
        nkp = global_k(nkp_loc)
        do nn = 1, kmesh_info%nntot
          nkp2 = kmesh_info%nnlist(nkp, nn)
          ! tmp_cdq = cdq^{dagger} . M

          ! note: m_matrix_loc is dimensioned larger than block copied here
          ! the striding used here likely incurs some overhead; ideally we should avoid it
          ! Jerome Jackson Jun 24
          call utility_zgemm(tmp_cdq, cdq(:, :, nkp), 'C', m_matrix_loc(1:num_wann, 1:num_wann, nn, nkp_loc), 'N', &
                             num_wann)
          ! cmtmp = tmp_cdq . cdq
          call utility_zgemm(cmtmp, tmp_cdq, 'N', cdq(:, :, nkp2), 'N', num_wann)
          ! note striding
          m_matrix_loc(1:num_wann, 1:num_wann, nn, nkp_loc) = cmtmp(:, :)
        end do
      end do

      if (timing_level > 1) call io_stopwatch_stop('wann: main: u_and_m', timer)
    end subroutine internal_new_u_and_m

  end subroutine wann_main