!--------------------------------------------------------------------------------------------------!
!   CP2K: A general program to perform molecular dynamics simulations                              !
!   Copyright (C) 2000 - 2018  CP2K developers group                                               !
!--------------------------------------------------------------------------------------------------!

! **************************************************************************************************
!> \brief Rountines to calculate gradients of RI-GPW-MP2 energy using pw
!> \par History
!>      10.2013 created [Mauro Del Ben]
! **************************************************************************************************
MODULE mp2_ri_grad
   USE atomic_kind_types,               ONLY: atomic_kind_type,&
                                              get_atomic_kind_set
   USE basis_set_types,                 ONLY: gto_basis_set_type
   USE cell_types,                      ONLY: cell_type,&
                                              pbc
   USE cp_blacs_env,                    ONLY: cp_blacs_env_type
   USE cp_control_types,                ONLY: dft_control_type
   USE cp_dbcsr_operations,             ONLY: copy_dbcsr_to_fm
   USE cp_eri_mme_interface,            ONLY: cp_eri_mme_param
   USE cp_fm_struct,                    ONLY: cp_fm_struct_create,&
                                              cp_fm_struct_release,&
                                              cp_fm_struct_type
   USE cp_fm_types,                     ONLY: cp_fm_create,&
                                              cp_fm_get_info,&
                                              cp_fm_indxg2l,&
                                              cp_fm_indxg2p,&
                                              cp_fm_indxl2g,&
                                              cp_fm_release,&
                                              cp_fm_set_all,&
                                              cp_fm_type
   USE cp_gemm_interface,               ONLY: cp_gemm
   USE cp_para_env,                     ONLY: cp_para_env_create,&
                                              cp_para_env_release
   USE cp_para_types,                   ONLY: cp_para_env_type
   USE dbcsr_api,                       ONLY: &
        dbcsr_add, dbcsr_copy, dbcsr_copy_into_existing, dbcsr_create, &
        dbcsr_deallocate_matrix_set, dbcsr_multiply, dbcsr_p_type, dbcsr_release, dbcsr_set, &
        dbcsr_transposed, dbcsr_type, dbcsr_type_no_symmetry, dbcsr_type_symmetric
   USE gaussian_gridlevels,             ONLY: gaussian_gridlevel
   USE input_constants,                 ONLY: do_eri_gpw,&
                                              do_eri_mme
   USE kinds,                           ONLY: dp
   USE mathconstants,                   ONLY: fourpi
   USE message_passing,                 ONLY: &
        mp_alltoall, mp_comm_split_direct, mp_irecv, mp_isend, mp_request_null, mp_sendrecv, &
        mp_sum, mp_wait, mp_waitall
   USE mp2_eri,                         ONLY: mp2_eri_2c_integrate,&
                                              mp2_eri_3c_integrate,&
                                              mp2_eri_deallocate_forces,&
                                              mp2_eri_force
   USE mp2_types,                       ONLY: integ_mat_buffer_type,&
                                              integ_mat_buffer_type_2D,&
                                              mp2_type
   USE orbital_pointers,                ONLY: ncoset
   USE particle_types,                  ONLY: particle_type
   USE pw_env_types,                    ONLY: pw_env_get,&
                                              pw_env_type
   USE pw_methods,                      ONLY: pw_copy,&
                                              pw_derive,&
                                              pw_integral_ab,&
                                              pw_scale,&
                                              pw_transfer
   USE pw_poisson_methods,              ONLY: pw_poisson_solve
   USE pw_poisson_types,                ONLY: pw_poisson_type
   USE pw_pool_types,                   ONLY: pw_pool_create_pw,&
                                              pw_pool_give_back_pw,&
                                              pw_pool_type
   USE pw_types,                        ONLY: COMPLEXDATA1D,&
                                              REALDATA3D,&
                                              REALSPACE,&
                                              RECIPROCALSPACE,&
                                              pw_p_type
   USE qs_collocate_density,            ONLY: calculate_rho_elec,&
                                              calculate_wavefunction
   USE qs_environment_types,            ONLY: get_qs_env,&
                                              qs_environment_type
   USE qs_force_types,                  ONLY: qs_force_type
   USE qs_integrate_potential,          ONLY: integrate_pgf_product_rspace,&
                                              integrate_v_rspace
   USE qs_kind_types,                   ONLY: get_qs_kind,&
                                              qs_kind_type
   USE qs_ks_types,                     ONLY: qs_ks_env_type
   USE qs_neighbor_list_types,          ONLY: neighbor_list_set_p_type
   USE realspace_grid_types,            ONLY: realspace_grid_desc_p_type,&
                                              realspace_grid_p_type,&
                                              rs_grid_release,&
                                              rs_grid_retain
   USE rs_pw_interface,                 ONLY: potential_pw2rs
   USE task_list_types,                 ONLY: task_list_type
   USE util,                            ONLY: get_limit
   USE virial_types,                    ONLY: virial_type
#include "./base/base_uses.f90"

   IMPLICIT NONE

   PRIVATE

   CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'mp2_ri_grad'

   PUBLIC :: calc_ri_mp2_nonsep

CONTAINS

! **************************************************************************************************
!> \brief Calcualte the non-separable part of the gradients and update the
!>        Lagrangian
!> \param qs_env ...
!> \param mp2_env ...
!> \param para_env ...
!> \param para_env_sub ...
!> \param dft_control ...
!> \param cell ...
!> \param particle_set ...
!> \param atomic_kind_set ...
!> \param qs_kind_set ...
!> \param mo_coeff ...
!> \param nmo ...
!> \param homo ...
!> \param dimen_RI ...
!> \param Eigenval ...
!> \param my_group_L_start ...
!> \param my_group_L_end ...
!> \param my_group_L_size ...
!> \param sab_orb_sub ...
!> \param rho_r ...
!> \param rho_g ...
!> \param pot_g ...
!> \param mat_munu ...
!> \param pw_env_sub ...
!> \param poisson_env ...
!> \param auxbas_pw_pool ...
!> \param task_list_sub ...
!> \param blacs_env_sub ...
!> \param Eigenval_beta ...
!> \param homo_beta ...
!> \param mo_coeff_beta ...
!> \author Mauro Del Ben
! **************************************************************************************************
   SUBROUTINE calc_ri_mp2_nonsep(qs_env, mp2_env, para_env, para_env_sub, dft_control, cell, particle_set, &
                                 atomic_kind_set, qs_kind_set, mo_coeff, nmo, homo, dimen_RI, Eigenval, &
                                 my_group_L_start, my_group_L_end, my_group_L_size, sab_orb_sub, rho_r, rho_g, pot_g, &
                                 mat_munu, pw_env_sub, poisson_env, auxbas_pw_pool, task_list_sub, &
                                 blacs_env_sub, Eigenval_beta, homo_beta, mo_coeff_beta)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(mp2_type), POINTER                            :: mp2_env
      TYPE(cp_para_env_type), POINTER                    :: para_env, para_env_sub
      TYPE(dft_control_type), POINTER                    :: dft_control
      TYPE(cell_type), POINTER                           :: cell
      TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
      TYPE(atomic_kind_type), DIMENSION(:), POINTER      :: atomic_kind_set
      TYPE(qs_kind_type), DIMENSION(:), POINTER          :: qs_kind_set
      TYPE(cp_fm_type), POINTER                          :: mo_coeff
      INTEGER                                            :: nmo, homo, dimen_RI
      REAL(KIND=dp), DIMENSION(:)                        :: Eigenval
      INTEGER                                            :: my_group_L_start, my_group_L_end, &
                                                            my_group_L_size
      TYPE(neighbor_list_set_p_type), DIMENSION(:), &
         POINTER                                         :: sab_orb_sub
      TYPE(pw_p_type)                                    :: rho_r, rho_g, pot_g
      TYPE(dbcsr_p_type)                                 :: mat_munu
      TYPE(pw_env_type), POINTER                         :: pw_env_sub
      TYPE(pw_poisson_type), POINTER                     :: poisson_env
      TYPE(pw_pool_type), POINTER                        :: auxbas_pw_pool
      TYPE(task_list_type), POINTER                      :: task_list_sub
      TYPE(cp_blacs_env_type), POINTER                   :: blacs_env_sub
      REAL(KIND=dp), DIMENSION(:), OPTIONAL              :: Eigenval_beta
      INTEGER, OPTIONAL                                  :: homo_beta
      TYPE(cp_fm_type), OPTIONAL, POINTER                :: mo_coeff_beta

      CHARACTER(LEN=*), PARAMETER :: routineN = 'calc_ri_mp2_nonsep', &
         routineP = moduleN//':'//routineN

      INTEGER :: alpha, atom_a, beta, dimen, dir, eri_method, handle, handle2, handle3, i, iatom, &
         igrid_level, ikind, iorb, ipgf, iset, itmp(2), L_counter, lb(3), LLL, location(3), &
         my_P_end, my_P_size, my_P_start, na1, na2, natom, ncoa, nseta, offset, sgfa, tp(3), &
         ub(3), virtual, virtual_beta
      INTEGER, ALLOCATABLE, DIMENSION(:)                 :: atom_of_kind, kind_of
      INTEGER, DIMENSION(3)                              :: comp
      INTEGER, DIMENSION(:), POINTER                     :: la_max, la_min, npgfa, nsgfa
      INTEGER, DIMENSION(:, :), POINTER                  :: first_sgfa
      LOGICAL                                            :: alpha_beta, map_it_here, skip_shell, &
                                                            use_virial
      REAL(KIND=dp)                                      :: e_hartree, eps_filter, factor, &
                                                            pair_energy, rab2, total_rho
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:)           :: wf_vector
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :)        :: G_PQ_local, I_ab
      REAL(KIND=dp), DIMENSION(3)                        :: force_a, force_b, ra, rab
      REAL(KIND=dp), DIMENSION(3, 3)                     :: h_stress, my_virial_a, my_virial_b
      REAL(KIND=dp), DIMENSION(:), POINTER               :: set_radius_a
      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: I_tmp2, pab, rpgfa, sphi_a, zeta
      TYPE(cp_eri_mme_param), POINTER                    :: eri_param
      TYPE(cp_fm_struct_type), POINTER                   :: fm_struct_tmp
      TYPE(cp_fm_type), POINTER                          :: L1_mu_i, L1_mu_i_beta, L2_nu_a, &
                                                            L2_nu_a_beta
      TYPE(dbcsr_p_type) :: Lag_mu_i_1, Lag_mu_i_1_beta, Lag_nu_a_2, Lag_nu_a_2_beta, &
         matrix_P_inu, matrix_P_inu_beta, matrix_P_munu, matrix_P_munu_nosym
      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: G_P_ia, G_P_ia_beta, mat_munu_local, &
                                                            matrix_P_munu_local
      TYPE(dbcsr_type), POINTER                          :: mo_coeff_o, mo_coeff_o_beta, mo_coeff_v, &
                                                            mo_coeff_v_beta
      TYPE(gto_basis_set_type), POINTER                  :: basis_set_a
      TYPE(mp2_eri_force), ALLOCATABLE, DIMENSION(:)     :: force_2c, force_3c_aux, force_3c_orb_mu, &
                                                            force_3c_orb_nu
      TYPE(pw_p_type)                                    :: dvg(3), psi_L, psi_L_beta, temp_pw_g
      TYPE(qs_force_type), DIMENSION(:), POINTER         :: force
      TYPE(qs_ks_env_type), POINTER                      :: ks_env
      TYPE(realspace_grid_desc_p_type), DIMENSION(:), &
         POINTER                                         :: rs_descs
      TYPE(realspace_grid_p_type), DIMENSION(:), POINTER :: rs_v
      TYPE(virial_type), POINTER                         :: virial

      CALL timeset(routineN, handle)

      eri_method = mp2_env%eri_method
      eri_param => mp2_env%eri_mme_param

      ! Find out whether we have a closed or open shell
      alpha_beta = .FALSE.
      IF (PRESENT(homo_beta) .AND. PRESENT(Eigenval_beta)) alpha_beta = .TRUE.

      dimen = nmo
      virtual = dimen-homo
      IF (alpha_beta) virtual_beta = dimen-homo_beta
      eps_filter = mp2_env%mp2_gpw%eps_filter
      NULLIFY (mo_coeff_o, mo_coeff_v, G_P_ia, ks_env)
      mo_coeff_o => mp2_env%ri_grad%mo_coeff_o
      mo_coeff_v => mp2_env%ri_grad%mo_coeff_v
      G_P_ia => mp2_env%ri_grad%G_P_ia

      IF (alpha_beta) THEN
         NULLIFY (mo_coeff_o_beta, mo_coeff_v_beta, G_P_ia_beta)
         mo_coeff_o_beta => mp2_env%ri_grad%mo_coeff_o_beta
         mo_coeff_v_beta => mp2_env%ri_grad%mo_coeff_v_beta
         G_P_ia_beta => mp2_env%ri_grad%G_P_ia_beta
      ENDIF

      itmp = get_limit(dimen_RI, para_env_sub%num_pe, para_env_sub%mepos)
      my_P_start = itmp(1)
      my_P_end = itmp(2)
      my_P_size = itmp(2)-itmp(1)+1
      ALLOCATE (G_PQ_local(dimen_RI, my_group_L_size))

      G_PQ_local = 0.0_dp
      IF (.NOT. alpha_beta) THEN
         G_PQ_local(my_P_start:my_P_end, 1:my_group_L_size) = mp2_env%ri_grad%Gamma_PQ
      ELSE
         G_PQ_local(my_P_start:my_P_end, 1:my_group_L_size) = &
            0.50_dp*(mp2_env%ri_grad%Gamma_PQ+mp2_env%ri_grad%Gamma_PQ_beta)
      ENDIF
      DEALLOCATE (mp2_env%ri_grad%Gamma_PQ)
      IF (alpha_beta) THEN
         DEALLOCATE (mp2_env%ri_grad%Gamma_PQ_beta)
      ENDIF
      CALL mp_sum(G_PQ_local, para_env_sub%group)

      ! deallocate here PQ_half, maybe usefull in the future
      ! This is really bad style
      DEALLOCATE (mp2_env%ri_grad%PQ_half)

      ! create matrix holding the back transformation (G_P_inu)
      ALLOCATE (matrix_P_inu%matrix)
      CALL dbcsr_create(matrix_P_inu%matrix, template=mo_coeff_o)
      IF (alpha_beta) THEN
         ALLOCATE (matrix_P_inu_beta%matrix)
         CALL dbcsr_create(matrix_P_inu_beta%matrix, template=mo_coeff_o_beta)
      ENDIF

      ! non symmetric matrix
      ALLOCATE (matrix_P_munu_nosym%matrix)
      CALL dbcsr_create(matrix_P_munu_nosym%matrix, template=mat_munu%matrix, &
                        matrix_type=dbcsr_type_no_symmetry)

      ! create Lagrangian matrices in mixed AO/MO formalism
      ALLOCATE (Lag_mu_i_1%matrix)
      CALL dbcsr_create(Lag_mu_i_1%matrix, template=mo_coeff_o)
      CALL dbcsr_set(Lag_mu_i_1%matrix, 0.0_dp)
      IF (alpha_beta) THEN
         ALLOCATE (Lag_mu_i_1_beta%matrix)
         CALL dbcsr_create(Lag_mu_i_1_beta%matrix, template=mo_coeff_o_beta)
         CALL dbcsr_set(Lag_mu_i_1_beta%matrix, 0.0_dp)
      ENDIF

      ALLOCATE (Lag_nu_a_2%matrix)
      CALL dbcsr_create(Lag_nu_a_2%matrix, template=mo_coeff_v)
      CALL dbcsr_set(Lag_nu_a_2%matrix, 0.0_dp)
      IF (alpha_beta) THEN
         ALLOCATE (Lag_nu_a_2_beta%matrix)
         CALL dbcsr_create(Lag_nu_a_2_beta%matrix, template=mo_coeff_v_beta)
         CALL dbcsr_set(Lag_nu_a_2_beta%matrix, 0.0_dp)
      ENDIF

      ! get forces
      NULLIFY (force, virial)
      CALL get_qs_env(qs_env=qs_env, force=force, virial=virial)

      ! prepare integral derivatives with mme method
      IF (eri_method .EQ. do_eri_mme) THEN
         ALLOCATE (matrix_P_munu_local(my_group_L_size))
         ALLOCATE (mat_munu_local(my_group_L_size))
         L_counter = 0
         DO LLL = my_group_L_start, my_group_L_end
            L_counter = L_counter+1
            ALLOCATE (matrix_P_munu_local(L_counter)%matrix)
            ALLOCATE (mat_munu_local(L_counter)%matrix)
            CALL dbcsr_create(matrix_P_munu_local(L_counter)%matrix, template=mat_munu%matrix, &
                              matrix_type=dbcsr_type_symmetric)
            CALL dbcsr_create(mat_munu_local(L_counter)%matrix, template=mat_munu%matrix, &
                              matrix_type=dbcsr_type_symmetric)
            CALL dbcsr_copy(mat_munu_local(L_counter)%matrix, mat_munu%matrix)
            CALL dbcsr_set(mat_munu_local(L_counter)%matrix, 0.0_dp)

            IF (alpha_beta) THEN
               CALL G_P_transform_MO_to_AO(matrix_P_munu_local(L_counter), matrix_P_munu_nosym, mat_munu, &
                                           G_P_ia(L_counter), matrix_P_inu, &
                                           mo_coeff_v, mo_coeff_o, eps_filter, &
                                           G_P_ia_beta(L_counter), matrix_P_inu_beta, mo_coeff_v_beta, mo_coeff_o_beta)
            ELSE
               CALL G_P_transform_MO_to_AO(matrix_P_munu_local(L_counter), matrix_P_munu_nosym, mat_munu, &
                                           G_P_ia(L_counter), matrix_P_inu, &
                                           mo_coeff_v, mo_coeff_o, eps_filter)
            ENDIF
         ENDDO

         ALLOCATE (I_tmp2(dimen_RI, my_group_L_size))
         I_tmp2(:, :) = 0.0_dp
         CALL mp2_eri_2c_integrate(eri_param, para_env_sub, qs_env, &
                                   basis_type_a="RI_AUX", basis_type_b="RI_AUX", &
                                   hab=I_tmp2, first_b=my_group_L_start, last_b=my_group_L_end, &
                                   eri_method=eri_method, pab=G_PQ_local, force_a=force_2c)

         DEALLOCATE (I_tmp2)
         CALL mp2_eri_3c_integrate(eri_param, para_env_sub, qs_env, &
                                   first_c=my_group_L_start, last_c=my_group_L_end, mat_ab=mat_munu_local, &
                                   basis_type_a="ORB", basis_type_b="ORB", basis_type_c="RI_AUX", &
                                   sab_nl=sab_orb_sub, eri_method=eri_method, &
                                   pabc=matrix_P_munu_local, &
                                   force_a=force_3c_orb_mu, force_b=force_3c_orb_nu, force_c=force_3c_aux)

         L_counter = 0
         DO LLL = my_group_L_start, my_group_L_end
            L_counter = L_counter+1
            ! we recompute matrix_P_inu
            CALL dbcsr_multiply("N", "T", 1.0_dp, mo_coeff_v, G_P_ia(L_counter)%matrix, &
                                0.0_dp, matrix_P_inu%matrix, filter_eps=eps_filter)
            IF (alpha_beta) THEN
               CALL dbcsr_multiply("N", "T", 1.0_dp, mo_coeff_v_beta, G_P_ia_beta(L_counter)%matrix, &
                                   0.0_dp, matrix_P_inu_beta%matrix, filter_eps=eps_filter)
            ENDIF

            IF (alpha_beta) THEN
               CALL update_lagrangian(mat_munu_local(L_counter), matrix_P_inu, Lag_mu_i_1, &
                                      G_P_ia(L_counter), mo_coeff_o, Lag_nu_a_2, &
                                      eps_filter, &
                                      matrix_P_inu_beta, Lag_mu_i_1_beta, G_P_ia_beta(L_counter), mo_coeff_o_beta, Lag_nu_a_2_beta)
            ELSE
               CALL update_lagrangian(mat_munu_local(L_counter), matrix_P_inu, Lag_mu_i_1, &
                                      G_P_ia(L_counter), mo_coeff_o, Lag_nu_a_2, &
                                      eps_filter)

            ENDIF
         ENDDO

         DO ikind = 1, SIZE(force)
            force(ikind)%mp2_non_sep(:, :) = -4.0_dp*force_2c(ikind)%forces(:, :)+ &
                                             force_3c_orb_mu(ikind)%forces(:, :)+ &
                                             force_3c_orb_nu(ikind)%forces(:, :)+ &
                                             force_3c_aux(ikind)%forces(:, :)
         END DO

         CALL mp2_eri_deallocate_forces(force_2c)
         CALL mp2_eri_deallocate_forces(force_3c_aux)
         CALL mp2_eri_deallocate_forces(force_3c_orb_mu)
         CALL mp2_eri_deallocate_forces(force_3c_orb_nu)
         CALL dbcsr_deallocate_matrix_set(matrix_P_munu_local)
         CALL dbcsr_deallocate_matrix_set(mat_munu_local)

      ELSEIF (eri_method == do_eri_gpw) THEN
         ALLOCATE (matrix_P_munu%matrix)
         CALL dbcsr_create(matrix_P_munu%matrix, template=mat_munu%matrix, &
                           matrix_type=dbcsr_type_symmetric)

         CALL get_qs_env(qs_env, ks_env=ks_env)

         natom = SIZE(particle_set)

         ALLOCATE (kind_of(natom))
         ALLOCATE (atom_of_kind(natom))
         CALL get_atomic_kind_set(atomic_kind_set, kind_of=kind_of, atom_of_kind=atom_of_kind)

         ! wave function vector and supporting stuff
         ALLOCATE (wf_vector(dimen_RI))
         NULLIFY (psi_L%pw)
         CALL pw_pool_create_pw(auxbas_pw_pool, psi_L%pw, &
                                use_data=REALDATA3D, &
                                in_space=REALSPACE)
         IF (alpha_beta) THEN
            NULLIFY (psi_L_beta%pw)
            CALL pw_pool_create_pw(auxbas_pw_pool, psi_L_beta%pw, &
                                   use_data=REALDATA3D, &
                                   in_space=REALSPACE)
         ENDIF

         ! check if we want to calculate the virial
         use_virial = virial%pv_availability .AND. (.NOT. virial%pv_numer)

         ! in case virial is required we need auxilliary pw
         ! for calculate the MP2-volume contribution to the virial
         ! (hartree potential derivatives)
         IF (use_virial) THEN
            NULLIFY (temp_pw_g%pw)
            CALL pw_pool_create_pw(auxbas_pw_pool, temp_pw_g%pw, &
                                   use_data=COMPLEXDATA1D, &
                                   in_space=RECIPROCALSPACE)
            DO i = 1, 3
               NULLIFY (dvg(i)%pw)
               CALL pw_pool_create_pw(auxbas_pw_pool, dvg(i)%pw, &
                                      use_data=COMPLEXDATA1D, &
                                      in_space=RECIPROCALSPACE)
            END DO
         END IF

         ! start main loop over auxiliary basis functions
         CALL timeset(routineN//"_loop", handle2)

         L_counter = 0
         DO LLL = my_group_L_start, my_group_L_end
            L_counter = L_counter+1

            IF (alpha_beta) THEN
               CALL G_P_transform_MO_to_AO(matrix_P_munu, matrix_P_munu_nosym, mat_munu, &
                                           G_P_ia(L_counter), &
                                           matrix_P_inu, &
                                           mo_coeff_v, mo_coeff_o, &
                                           eps_filter, &
                                           G_P_ia_beta(L_counter), matrix_P_inu_beta, &
                                           mo_coeff_v_beta, mo_coeff_o_beta)
            ELSE
               CALL G_P_transform_MO_to_AO(matrix_P_munu, matrix_P_munu_nosym, mat_munu, &
                                           G_P_ia(L_counter), &
                                           matrix_P_inu, &
                                           mo_coeff_v, mo_coeff_o, &
                                           eps_filter)
            ENDIF

            ! calculate potential associted to the single aux function
            CALL timeset(routineN//"_wf_pot", handle3)
            wf_vector = 0.0_dp
            wf_vector(LLL) = 1.0_dp
            ! pseudo psi_L
            CALL calculate_wavefunction(mo_coeff, 1, psi_L, rho_g, atomic_kind_set, &
                                        qs_kind_set, cell, dft_control, particle_set, pw_env_sub, &
                                        basis_type="RI_AUX", &
                                        external_vector=wf_vector)

            rho_r%pw%cr3d = psi_L%pw%cr3d
            IF (alpha_beta) THEN
               CALL calculate_wavefunction(mo_coeff_beta, 1, psi_L_beta, rho_g, atomic_kind_set, &
                                           qs_kind_set, cell, dft_control, particle_set, &
                                           pw_env_sub, basis_type='RI_AUX', &
                                           external_vector=wf_vector)
               rho_r%pw%cr3d = 0.50_dp*(rho_r%pw%cr3d+psi_L_beta%pw%cr3d)
            ENDIF
            CALL pw_transfer(rho_r%pw, rho_g%pw)
            CALL pw_poisson_solve(poisson_env, rho_g%pw, pair_energy, pot_g%pw)
            CALL pw_transfer(pot_g%pw, rho_r%pw)
            CALL pw_scale(rho_r%pw, rho_r%pw%pw_grid%dvol)
            CALL timestop(handle3)

            IF (use_virial) THEN
               ! make a copy of the density in G space
               ! calculate the potential derivatives in G space
               CALL timeset(routineN//"_Virial", handle3)
               CALL pw_copy(rho_g%pw, temp_pw_g%pw)
               DO i = 1, 3
                  comp = 0
                  comp(i) = 1
                  CALL pw_copy(pot_g%pw, dvg(i)%pw)
                  CALL pw_derive(dvg(i)%pw, comp)
               END DO
               CALL timestop(handle3)
            END IF

            ! integrate the potential of the single gaussian and update
            ! 2-center forces with Gamma_PQ
            CALL timeset(routineN//"_int_PQ", handle3)
            NULLIFY (rs_v)
            NULLIFY (rs_descs)
            CALL pw_env_get(pw_env_sub, rs_descs=rs_descs, rs_grids=rs_v)
            DO i = 1, SIZE(rs_v)
               CALL rs_grid_retain(rs_v(i)%rs_grid)
            END DO
            CALL potential_pw2rs(rs_v, rho_r, pw_env_sub)

            offset = 0
            DO iatom = 1, natom
               ikind = kind_of(iatom)
               atom_a = atom_of_kind(iatom)
               CALL get_qs_kind(qs_kind=qs_kind_set(ikind), basis_set=basis_set_a, &
                                basis_type="RI_AUX")

               first_sgfa => basis_set_a%first_sgf
               la_max => basis_set_a%lmax
               la_min => basis_set_a%lmin
               npgfa => basis_set_a%npgf
               nseta = basis_set_a%nset
               nsgfa => basis_set_a%nsgf_set
               rpgfa => basis_set_a%pgf_radius
               set_radius_a => basis_set_a%set_radius
               sphi_a => basis_set_a%sphi
               zeta => basis_set_a%zet

               ra(:) = pbc(particle_set(iatom)%r, cell)
               rab = 0.0_dp
               rab2 = 0.0_dp

               force_a(:) = 0.0_dp
               force_b(:) = 0.0_dp
               IF (use_virial) THEN
                  my_virial_a = 0.0_dp
                  my_virial_b = 0.0_dp
               END IF

               DO iset = 1, nseta
                  ncoa = npgfa(iset)*ncoset(la_max(iset))
                  sgfa = first_sgfa(1, iset)

                  ALLOCATE (I_tmp2(ncoa, 1))
                  I_tmp2 = 0.0_dp
                  ALLOCATE (I_ab(nsgfa(iset), 1))
                  I_ab = 0.0_dp
                  ALLOCATE (pab(ncoa, 1))
                  pab = 0.0_dp

                  I_ab(1:nsgfa(iset), 1) = -4.0_dp*G_PQ_local(offset+1:offset+nsgfa(iset), L_counter)

                  CALL dgemm("N", "N", ncoa, 1, nsgfa(iset), &
                             1.0_dp, sphi_a(1, sgfa), SIZE(sphi_a, 1), &
                             I_ab(1, 1), SIZE(I_ab, 1), &
                             0.0_dp, pab(1, 1), SIZE(pab, 1))

                  I_ab = 0.0_dp

                  igrid_level = gaussian_gridlevel(pw_env_sub%gridlevel_info, MINVAL(zeta(:, iset)))
                  map_it_here = .FALSE.
                  IF (.NOT. ALL(rs_v(igrid_level)%rs_grid%desc%perd == 1)) THEN
                     DO dir = 1, 3
                        ! bounds of local grid (i.e. removing the 'wings'), if periodic
                        tp(dir) = FLOOR(DOT_PRODUCT(cell%h_inv(dir, :), ra)*rs_v(igrid_level)%rs_grid%desc%npts(dir))
                        tp(dir) = MODULO(tp(dir), rs_v(igrid_level)%rs_grid%desc%npts(dir))
                        IF (rs_v(igrid_level)%rs_grid%desc%perd(dir) .NE. 1) THEN
                           lb(dir) = rs_v(igrid_level)%rs_grid%lb_local(dir)+rs_v(igrid_level)%rs_grid%desc%border
                           ub(dir) = rs_v(igrid_level)%rs_grid%ub_local(dir)-rs_v(igrid_level)%rs_grid%desc%border
                        ELSE
                           lb(dir) = rs_v(igrid_level)%rs_grid%lb_local(dir)
                           ub(dir) = rs_v(igrid_level)%rs_grid%ub_local(dir)
                        ENDIF
                        ! distributed grid, only map if it is local to the grid
                        location(dir) = tp(dir)+rs_v(igrid_level)%rs_grid%desc%lb(dir)
                     ENDDO
                     IF (lb(1) <= location(1) .AND. location(1) <= ub(1) .AND. &
                         lb(2) <= location(2) .AND. location(2) <= ub(2) .AND. &
                         lb(3) <= location(3) .AND. location(3) <= ub(3)) THEN
                        map_it_here = .TRUE.
                     ENDIF
                  ELSE
                     ! not distributed, just a round-robin distribution over the full set of CPUs
                     IF (MODULO(offset, para_env_sub%num_pe) == para_env_sub%mepos) map_it_here = .TRUE.
                  ENDIF

                  offset = offset+nsgfa(iset)

                  IF (map_it_here) THEN
                     DO ipgf = 1, npgfa(iset)
                        na1 = (ipgf-1)*ncoset(la_max(iset))+1
                        na2 = ipgf*ncoset(la_max(iset))
                        igrid_level = gaussian_gridlevel(pw_env_sub%gridlevel_info, zeta(ipgf, iset))

                        CALL integrate_pgf_product_rspace(la_max=la_max(iset), zeta=zeta(ipgf, iset)/2.0_dp, la_min=la_min(iset), &
                                                          lb_max=0, zetb=zeta(ipgf, iset)/2.0_dp, lb_min=0, &
                                                          ra=ra, rab=rab, rab2=rab2, &
                                                          rsgrid=rs_v(igrid_level)%rs_grid, &
                                                          cell=cell, &
                                                          cube_info=pw_env_sub%cube_info(igrid_level), &
                                                          hab=I_tmp2, &
                                                          pab=pab, &
                                                          o1=na1-1, &
                                                          o2=0, &
                                                          map_consistent=.TRUE., &
                                                          eps_gvg_rspace=dft_control%qs_control%eps_gvg_rspace, &
                                                          calculate_forces=.TRUE., &
                                                          force_a=force_a, force_b=force_b, &
                                                          use_virial=use_virial, my_virial_a=my_virial_a, my_virial_b=my_virial_b)

                     END DO

                     ! CALL dgemm("T","N",nsgfa(iset),1,ncoa,&
                     !             1.0_dp,sphi_a(1,sgfa),SIZE(sphi_a,1),&
                     !             I_tmp2(1,1),SIZE(I_tmp2,1),&
                     !             1.0_dp,I_ab(1,1),SIZE(I_ab,1))
                     ! L_local_col(offset-nsgfa(iset)+1:offset,i_counter)=I_ab(1:nsgfa(iset),1)
                  END IF

                  DEALLOCATE (I_tmp2)
                  DEALLOCATE (I_ab)
                  DEALLOCATE (pab)

               END DO

               force(ikind)%rho_elec(:, atom_a) = &
                  force(ikind)%rho_elec(:, atom_a)+force_a(:)+force_b
               IF (use_virial) THEN
                  virial%pv_virial = virial%pv_virial+my_virial_a+my_virial_b
               END IF
            END DO

            DO i = 1, SIZE(rs_v)
               CALL rs_grid_release(rs_v(i)%rs_grid)
            END DO
            CALL timestop(handle3)
            ! here we are done with the 2 centers

            ! CALL cp_dbcsr_write_sparse_matrix(matrix_P_munu%matrix,4,12,qs_env,para_env_sub,&
            !                                   output_unit=unit_nr)

            ! integrate the potential of the single gaussian and update
            ! 3-center forces
            CALL timeset(routineN//"_int", handle3)
            CALL dbcsr_set(mat_munu%matrix, 0.0_dp)
            CALL integrate_v_rspace(rho_r, hmat=mat_munu, pmat=matrix_P_munu, &
                                    qs_env=qs_env, calculate_forces=.TRUE., compute_tau=.FALSE., gapw=.FALSE., &
                                    pw_env_external=pw_env_sub, &
                                    task_list_external=task_list_sub)
            CALL timestop(handle3)

            IF (alpha_beta) THEN
               CALL update_lagrangian(mat_munu, matrix_P_inu, Lag_mu_i_1, &
                                      G_P_ia(L_counter), mo_coeff_o, Lag_nu_a_2, &
                                      eps_filter, &
                                      matrix_P_inu_beta, Lag_mu_i_1_beta, G_P_ia_beta(L_counter), mo_coeff_o_beta, Lag_nu_a_2_beta)
            ELSE
               CALL update_lagrangian(mat_munu, matrix_P_inu, Lag_mu_i_1, &
                                      G_P_ia(L_counter), mo_coeff_o, Lag_nu_a_2, &
                                      eps_filter)

            ENDIF

            IF (use_virial) THEN
               ! add the volume contribution to the virial due to
               ! the (P|Q) integrals, first we put the full gamme_PQ
               ! pseudo wave-function into grid in order to calculate the
               ! hartree potential derivatives
               CALL timeset(routineN//"_Virial", handle3)
               wf_vector = 0.0_dp
               wf_vector(:) = -2.0_dp*G_PQ_local(:, L_counter)
               CALL calculate_wavefunction(mo_coeff, 1, psi_L, rho_g, atomic_kind_set, &
                                           qs_kind_set, cell, dft_control, particle_set, pw_env_sub, &
                                           basis_type="RI_AUX", &
                                           external_vector=wf_vector)
               ! transfer to reciprocal space and calculate potential
               rho_r%pw%cr3d = psi_L%pw%cr3d
               CALL pw_transfer(rho_r%pw, rho_g%pw)
               CALL pw_poisson_solve(poisson_env, rho_g%pw, pair_energy, pot_g%pw)
               ! update virial with volume term (first calculate hartree like energy (diagonal part of the virial))
               e_hartree = 0.0_dp
               h_stress = 0.0_dp
               e_hartree = pw_integral_ab(temp_pw_g%pw, pot_g%pw)
               DO alpha = 1, 3
                  comp = 0
                  comp(alpha) = 1
                  CALL pw_copy(pot_g%pw, rho_g%pw)
                  CALL pw_derive(rho_g%pw, comp)
                  h_stress(alpha, alpha) = -e_hartree
                  DO beta = alpha, 3
                     h_stress(alpha, beta) = h_stress(alpha, beta) &
                                             -2.0_dp*pw_integral_ab(rho_g%pw, dvg(beta)%pw)/fourpi
                     h_stress(beta, alpha) = h_stress(alpha, beta)
                  END DO
               END DO
               virial%pv_virial = virial%pv_virial+h_stress/REAL(para_env_sub%num_pe, dp)
               CALL timestop(handle3)
            END IF

            ! put the gamma density on grid
            CALL timeset(routineN//"_Gpot", handle3)
            CALL calculate_rho_elec(matrix_p=matrix_P_munu%matrix, &
                                    rho=rho_r, &
                                    rho_gspace=rho_g, &
                                    total_rho=total_rho, &
                                    task_list_external=task_list_sub, &
                                    pw_env_external=pw_env_sub, &
                                    ks_env=ks_env)
            ! calculate associated hartree potential
            ! CALL pw_transfer(rho_r%pw, rho_g%pw)
            !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!1
            CALL pw_poisson_solve(poisson_env, rho_g%pw, pair_energy, pot_g%pw)
            CALL pw_transfer(pot_g%pw, rho_r%pw)
            CALL pw_scale(rho_r%pw, rho_r%pw%pw_grid%dvol)
            CALL timestop(handle3)

            IF (use_virial) THEN
               ! add the volume contribution to the virial coming from
               ! the 3-center integrals (mu nu|P)
               CALL timeset(routineN//"_Virial", handle3)
               e_hartree = 0.0_dp
               h_stress = 0.0_dp
               e_hartree = pw_integral_ab(temp_pw_g%pw, pot_g%pw)
               DO alpha = 1, 3
                  comp = 0
                  comp(alpha) = 1
                  CALL pw_copy(pot_g%pw, rho_g%pw)
                  CALL pw_derive(rho_g%pw, comp)
                  h_stress(alpha, alpha) = -e_hartree
                  DO beta = alpha, 3
                     h_stress(alpha, beta) = h_stress(alpha, beta) &
                                             -2.0_dp*pw_integral_ab(rho_g%pw, dvg(beta)%pw)/fourpi
                     h_stress(beta, alpha) = h_stress(alpha, beta)
                  END DO
               END DO
               virial%pv_virial = virial%pv_virial+h_stress/REAL(para_env_sub%num_pe, dp)
               CALL timestop(handle3)
            END IF

            ! integrate potential with auxiliary basis function derivatives
            NULLIFY (rs_v)
            NULLIFY (rs_descs)
            CALL pw_env_get(pw_env_sub, rs_descs=rs_descs, rs_grids=rs_v)
            DO i = 1, SIZE(rs_v)
               CALL rs_grid_retain(rs_v(i)%rs_grid)
            END DO
            CALL potential_pw2rs(rs_v, rho_r, pw_env_sub)

            offset = 0
            DO iatom = 1, natom
               ikind = kind_of(iatom)
               atom_a = atom_of_kind(iatom)
               CALL get_qs_kind(qs_kind=qs_kind_set(ikind), basis_set=basis_set_a, &
                                basis_type="RI_AUX")

               first_sgfa => basis_set_a%first_sgf
               la_max => basis_set_a%lmax
               la_min => basis_set_a%lmin
               npgfa => basis_set_a%npgf
               nseta = basis_set_a%nset
               nsgfa => basis_set_a%nsgf_set
               rpgfa => basis_set_a%pgf_radius
               set_radius_a => basis_set_a%set_radius
               sphi_a => basis_set_a%sphi
               zeta => basis_set_a%zet

               ra(:) = pbc(particle_set(iatom)%r, cell)
               rab = 0.0_dp
               rab2 = 0.0_dp

               force_a(:) = 0.0_dp
               force_b(:) = 0.0_dp
               IF (use_virial) THEN
                  my_virial_a = 0.0_dp
                  my_virial_b = 0.0_dp
               END IF

               DO iset = 1, nseta
                  ncoa = npgfa(iset)*ncoset(la_max(iset))
                  sgfa = first_sgfa(1, iset)

                  ALLOCATE (I_tmp2(ncoa, 1))
                  I_tmp2 = 0.0_dp
                  ALLOCATE (I_ab(nsgfa(iset), 1))
                  I_ab = 0.0_dp
                  ALLOCATE (pab(ncoa, 1))
                  pab = 0.0_dp

                  skip_shell = .TRUE.
                  DO iorb = 1, nsgfa(iset)
                     IF (iorb+offset == LLL) THEN
                        I_ab(iorb, 1) = 1.0_dp
                        skip_shell = .FALSE.
                     END IF
                  END DO

                  IF (skip_shell) THEN
                     offset = offset+nsgfa(iset)
                     DEALLOCATE (I_tmp2)
                     DEALLOCATE (I_ab)
                     DEALLOCATE (pab)
                     CYCLE
                  END IF

                  CALL dgemm("N", "N", ncoa, 1, nsgfa(iset), &
                             1.0_dp, sphi_a(1, sgfa), SIZE(sphi_a, 1), &
                             I_ab(1, 1), SIZE(I_ab, 1), &
                             0.0_dp, pab(1, 1), SIZE(pab, 1))
                  I_ab = 0.0_dp

                  igrid_level = gaussian_gridlevel(pw_env_sub%gridlevel_info, MINVAL(zeta(:, iset)))
                  map_it_here = .FALSE.
                  IF (.NOT. ALL(rs_v(igrid_level)%rs_grid%desc%perd == 1)) THEN
                     DO dir = 1, 3
                        ! bounds of local grid (i.e. removing the 'wings'), if periodic
                        tp(dir) = FLOOR(DOT_PRODUCT(cell%h_inv(dir, :), ra)*rs_v(igrid_level)%rs_grid%desc%npts(dir))
                        tp(dir) = MODULO(tp(dir), rs_v(igrid_level)%rs_grid%desc%npts(dir))
                        IF (rs_v(igrid_level)%rs_grid%desc%perd(dir) .NE. 1) THEN
                           lb(dir) = rs_v(igrid_level)%rs_grid%lb_local(dir)+rs_v(igrid_level)%rs_grid%desc%border
                           ub(dir) = rs_v(igrid_level)%rs_grid%ub_local(dir)-rs_v(igrid_level)%rs_grid%desc%border
                        ELSE
                           lb(dir) = rs_v(igrid_level)%rs_grid%lb_local(dir)
                           ub(dir) = rs_v(igrid_level)%rs_grid%ub_local(dir)
                        ENDIF
                        ! distributed grid, only map if it is local to the grid
                        location(dir) = tp(dir)+rs_v(igrid_level)%rs_grid%desc%lb(dir)
                     ENDDO
                     IF (lb(1) <= location(1) .AND. location(1) <= ub(1) .AND. &
                         lb(2) <= location(2) .AND. location(2) <= ub(2) .AND. &
                         lb(3) <= location(3) .AND. location(3) <= ub(3)) THEN
                        map_it_here = .TRUE.
                     ENDIF
                  ELSE
                     ! not distributed, just a round-robin distribution over the full set of CPUs
                     IF (MODULO(offset, para_env_sub%num_pe) == para_env_sub%mepos) map_it_here = .TRUE.
                  ENDIF

                  offset = offset+nsgfa(iset)

                  IF (map_it_here) THEN
                     DO ipgf = 1, npgfa(iset)
                        na1 = (ipgf-1)*ncoset(la_max(iset))+1
                        na2 = ipgf*ncoset(la_max(iset))
                        igrid_level = gaussian_gridlevel(pw_env_sub%gridlevel_info, zeta(ipgf, iset))

                        CALL integrate_pgf_product_rspace(la_max=la_max(iset), zeta=zeta(ipgf, iset)/2.0_dp, la_min=la_min(iset), &
                                                          lb_max=0, zetb=zeta(ipgf, iset)/2.0_dp, lb_min=0, &
                                                          ra=ra, rab=rab, rab2=rab2, &
                                                          rsgrid=rs_v(igrid_level)%rs_grid, &
                                                          cell=cell, &
                                                          cube_info=pw_env_sub%cube_info(igrid_level), &
                                                          hab=I_tmp2, &
                                                          pab=pab, &
                                                          o1=na1-1, &
                                                          o2=0, &
                                                          map_consistent=.TRUE., &
                                                          eps_gvg_rspace=dft_control%qs_control%eps_gvg_rspace, &
                                                          calculate_forces=.TRUE., &
                                                          force_a=force_a, force_b=force_b, &
                                                          use_virial=use_virial, my_virial_a=my_virial_a, my_virial_b=my_virial_b)
                     END DO
                     ! CALL dgemm("T","N".0_dp,sphi_a(1,sgfa),SIZE(sphi_a,1),&
                     !             I_tmp2(1,1),SIZE(I_tmp2,1),&
                     !             1.0_dp,I_ab(1,1),SIZE(I_ab,1))
                     ! L_local_col(offset-nsgfa(iset)+1:offset,i_counter)=I_ab(1:nsgfa(iset),1)
                  END IF

                  DEALLOCATE (I_tmp2)
                  DEALLOCATE (I_ab)
                  DEALLOCATE (pab)

               END DO

               force(ikind)%rho_elec(:, atom_a) = &
                  force(ikind)%rho_elec(:, atom_a)+force_a(:)+force_b(:)
               IF (use_virial) THEN
                  virial%pv_virial = virial%pv_virial+my_virial_a+my_virial_b
               END IF
            END DO

            DO i = 1, SIZE(rs_v)
               CALL rs_grid_release(rs_v(i)%rs_grid)
            END DO

         END DO

         CALL timestop(handle2)

         DEALLOCATE (kind_of)
         DEALLOCATE (atom_of_kind)
         DEALLOCATE (wf_vector)
         CALL pw_pool_give_back_pw(auxbas_pw_pool, psi_L%pw)

         IF (alpha_beta) THEN
            CALL pw_pool_give_back_pw(auxbas_pw_pool, psi_L_beta%pw)
         ENDIF

         IF (use_virial) THEN
            CALL pw_pool_give_back_pw(auxbas_pw_pool, temp_pw_g%pw)
            DO i = 1, 3
               CALL pw_pool_give_back_pw(auxbas_pw_pool, dvg(i)%pw)
            END DO
         END IF

         CALL dbcsr_release(matrix_P_munu%matrix)
         DEALLOCATE (matrix_P_munu%matrix)

      ENDIF

      DEALLOCATE (G_PQ_local)

      CALL dbcsr_release(matrix_P_inu%matrix)
      DEALLOCATE (matrix_P_inu%matrix)

      CALL dbcsr_release(matrix_P_munu_nosym%matrix)
      DEALLOCATE (matrix_P_munu_nosym%matrix)

      ! release the full gamma_P_ia structure
      DEALLOCATE (G_P_ia)

      ! Release G_P_ia_beta and and matrix_P_inu_beta
      IF (alpha_beta) THEN
         CALL dbcsr_release(matrix_P_inu_beta%matrix)
         DEALLOCATE (matrix_P_inu_beta%matrix)
         DEALLOCATE (G_P_ia_beta)
      ENDIF

      ! move the forces in the correct place
      IF (eri_method .EQ. do_eri_gpw) THEN
         DO ikind = 1, SIZE(force)
            force(ikind)%mp2_non_sep(:, :) = force(ikind)%rho_elec(:, :)
            force(ikind)%rho_elec(:, :) = 0.0_dp
         END DO
      ENDIF

      ! Now we move from the local matrices to the global ones
      ! defined over all MPI tasks
      ! Start with moving from the DBCSR to FM for the lagrangians
      NULLIFY (L1_mu_i, fm_struct_tmp)
      CALL cp_fm_struct_create(fm_struct_tmp, para_env=para_env_sub, context=blacs_env_sub, &
                               nrow_global=dimen, ncol_global=homo)
      CALL cp_fm_create(L1_mu_i, fm_struct_tmp, name="Lag_mu_i")
      CALL cp_fm_struct_release(fm_struct_tmp)
      CALL cp_fm_set_all(L1_mu_i, 0.0_dp)
      CALL copy_dbcsr_to_fm(matrix=Lag_mu_i_1%matrix, fm=L1_mu_i)

      ! release Lag_mu_i_1
      CALL dbcsr_release(Lag_mu_i_1%matrix)
      DEALLOCATE (Lag_mu_i_1%matrix)

      NULLIFY (L2_nu_a, fm_struct_tmp)
      CALL cp_fm_struct_create(fm_struct_tmp, para_env=para_env_sub, context=blacs_env_sub, &
                               nrow_global=dimen, ncol_global=virtual)
      CALL cp_fm_create(L2_nu_a, fm_struct_tmp, name="Lag_nu_a")
      CALL cp_fm_struct_release(fm_struct_tmp)
      CALL cp_fm_set_all(L2_nu_a, 0.0_dp)
      CALL copy_dbcsr_to_fm(matrix=Lag_nu_a_2%matrix, fm=L2_nu_a)

      ! release Lag_nu_a_2
      CALL dbcsr_release(Lag_nu_a_2%matrix)
      DEALLOCATE (Lag_nu_a_2%matrix)

      ! The same for the beta Lagrangians
      IF (alpha_beta) THEN
         ! Now we move from the local matrices to the global ones
         ! defined over all MPI tasks
         ! Start with moving from the DBCSR to FM for the lagrangians
         NULLIFY (L1_mu_i_beta, fm_struct_tmp)
         CALL cp_fm_struct_create(fm_struct_tmp, para_env=para_env_sub, context=blacs_env_sub, &
                                  nrow_global=dimen, ncol_global=homo_beta)
         CALL cp_fm_create(L1_mu_i_beta, fm_struct_tmp, name="Lag_mu_i_beta")
         CALL cp_fm_struct_release(fm_struct_tmp)
         CALL cp_fm_set_all(L1_mu_i_beta, 0.0_dp)
         CALL copy_dbcsr_to_fm(matrix=Lag_mu_i_1_beta%matrix, fm=L1_mu_i_beta)

         ! release Lag_mu_i_1
         CALL dbcsr_release(Lag_mu_i_1_beta%matrix)
         DEALLOCATE (Lag_mu_i_1_beta%matrix)

         NULLIFY (L2_nu_a_beta, fm_struct_tmp)
         CALL cp_fm_struct_create(fm_struct_tmp, para_env=para_env_sub, context=blacs_env_sub, &
                                  nrow_global=dimen, ncol_global=virtual_beta)
         CALL cp_fm_create(L2_nu_a_beta, fm_struct_tmp, name="Lag_nu_a_beta")
         CALL cp_fm_struct_release(fm_struct_tmp)
         CALL cp_fm_set_all(L2_nu_a_beta, 0.0_dp)
         CALL copy_dbcsr_to_fm(matrix=Lag_nu_a_2_beta%matrix, fm=L2_nu_a_beta)

         ! release Lag_nu_a_2
         CALL dbcsr_release(Lag_nu_a_2_beta%matrix)
         DEALLOCATE (Lag_nu_a_2_beta%matrix)
      ENDIF

      ! Set the factor to multiply P_ij (depends on the open or closed shell)
      factor = 1.0_dp
      IF (alpha_beta) factor = 0.50_dp

      CALL create_W_P(qs_env, mp2_env, mo_coeff, homo, virtual, dimen, para_env, para_env_sub, &
                      Eigenval, L1_mu_i, L2_nu_a, factor, .FALSE.)
      ! Alpha_beta_case
      IF (alpha_beta) THEN
         CALL create_W_P(qs_env, mp2_env, mo_coeff_beta, homo_beta, virtual_beta, dimen, para_env, &
                         para_env_sub, Eigenval_beta, L1_mu_i_beta, L2_nu_a_beta, &
                         factor, .TRUE.)
      ENDIF

      CALL timestop(handle)

   END SUBROUTINE calc_ri_mp2_nonsep

! **************************************************************************************************
!> \brief ...
!> \param G_P_munu ...
!> \param G_P_munu_nosym ...
!> \param mat_munu ...
!> \param G_P_ia ...
!> \param G_P_inu ...
!> \param mo_coeff_v ...
!> \param mo_coeff_o ...
!> \param eps_filter ...
!> \param G_P_ia_beta ...
!> \param G_P_inu_beta ...
!> \param mo_coeff_v_beta ...
!> \param mo_coeff_o_beta ...
! **************************************************************************************************
   SUBROUTINE G_P_transform_MO_to_AO(G_P_munu, G_P_munu_nosym, mat_munu, G_P_ia, G_P_inu, &
                                     mo_coeff_v, mo_coeff_o, eps_filter, &
                                     G_P_ia_beta, G_P_inu_beta, mo_coeff_v_beta, mo_coeff_o_beta)
      TYPE(dbcsr_p_type), INTENT(INOUT)                  :: G_P_munu, G_P_munu_nosym, mat_munu
      TYPE(dbcsr_p_type), INTENT(IN)                     :: G_P_ia
      TYPE(dbcsr_p_type), INTENT(INOUT)                  :: G_P_inu
      TYPE(dbcsr_type), INTENT(IN)                       :: mo_coeff_v, mo_coeff_o
      REAL(KIND=dp), INTENT(IN)                          :: eps_filter
      TYPE(dbcsr_p_type), INTENT(IN), OPTIONAL           :: G_P_ia_beta
      TYPE(dbcsr_p_type), INTENT(INOUT), OPTIONAL        :: G_P_inu_beta
      TYPE(dbcsr_type), INTENT(IN), OPTIONAL             :: mo_coeff_v_beta, mo_coeff_o_beta

      CHARACTER(LEN=*), PARAMETER :: routineN = 'G_P_transform_MO_to_AO', &
         routineP = moduleN//':'//routineN

      INTEGER                                            :: handle
      LOGICAL                                            :: alpha_beta

      alpha_beta = PRESENT(G_P_ia_beta) .AND. PRESENT(G_P_inu_beta) .AND. PRESENT(mo_coeff_v_beta) .AND. PRESENT(mo_coeff_o_beta)

      CALL dbcsr_set(G_P_munu_nosym%matrix, 0.0_dp)
      CALL G_P_transform_alpha_beta(G_P_ia, G_P_inu, G_P_munu_nosym, mo_coeff_v, mo_coeff_o, [1.0_dp, 0.0_dp], eps_filter)

      IF (alpha_beta) THEN
         CALL G_P_transform_alpha_beta(G_P_ia_beta, G_P_inu_beta, G_P_munu_nosym, mo_coeff_v_beta, mo_coeff_o_beta, &
                                       [0.5_dp, 0.5_dp], eps_filter)
      ENDIF

      ! symmetrize
      CALL timeset(routineN//"_symmetrize", handle)
      CALL dbcsr_set(G_P_munu%matrix, 0.0_dp)
      CALL dbcsr_transposed(G_P_munu%matrix, G_P_munu_nosym%matrix)
      CALL dbcsr_add(G_P_munu%matrix, G_P_munu_nosym%matrix, &
                     alpha_scalar=2.0_dp, beta_scalar=2.0_dp)
      ! this is a trick to avoid that integrate_v_rspace starts to cry
      CALL dbcsr_copy_into_existing(mat_munu%matrix, G_P_munu%matrix)
      CALL dbcsr_copy(G_P_munu%matrix, mat_munu%matrix)

      CALL timestop(handle)

   END SUBROUTINE G_P_transform_MO_to_AO

! **************************************************************************************************
!> \brief ...
!> \param G_P_ia ...
!> \param G_P_inu ...
!> \param G_P_munu ...
!> \param mo_coeff_v ...
!> \param mo_coeff_o ...
!> \param fraction_add ...
!> \param eps_filter ...
! **************************************************************************************************
   SUBROUTINE G_P_transform_alpha_beta(G_P_ia, G_P_inu, G_P_munu, mo_coeff_v, mo_coeff_o, fraction_add, eps_filter)
      TYPE(dbcsr_p_type), INTENT(IN)                     :: G_P_ia
      TYPE(dbcsr_p_type), INTENT(INOUT)                  :: G_P_inu, G_P_munu
      TYPE(dbcsr_type), INTENT(IN)                       :: mo_coeff_v, mo_coeff_o
      REAL(KIND=dp), DIMENSION(2), INTENT(IN)            :: fraction_add
      REAL(KIND=dp), INTENT(IN)                          :: eps_filter

      CHARACTER(LEN=*), PARAMETER :: routineN = 'G_P_transform_alpha_beta', &
         routineP = moduleN//':'//routineN

      INTEGER                                            :: handle

      CALL timeset(routineN//"_back_v", handle)
      ! first back-transformation a->nu
      CALL dbcsr_set(G_P_inu%matrix, 0.0_dp)
      CALL dbcsr_multiply("N", "T", 1.0_dp, mo_coeff_v, G_P_ia%matrix, &
                          0.0_dp, G_P_inu%matrix, filter_eps=eps_filter)
      CALL timestop(handle)

      ! second back-transformation i->mu
      CALL timeset(routineN//"_back_o", handle)
      CALL dbcsr_multiply("N", "T", fraction_add(1), G_P_inu%matrix, mo_coeff_o, &
                          fraction_add(2), G_P_munu%matrix, filter_eps=eps_filter)

      CALL timestop(handle)

   END SUBROUTINE G_P_transform_alpha_beta

! **************************************************************************************************
!> \brief ...
!> \param mat_munu ...
!> \param matrix_P_inu ...
!> \param Lag_mu_i_1 ...
!> \param G_P_ia ...
!> \param mo_coeff_o ...
!> \param Lag_nu_a_2 ...
!> \param eps_filter ...
!> \param matrix_P_inu_beta ...
!> \param Lag_mu_i_1_beta ...
!> \param G_P_ia_beta ...
!> \param mo_coeff_o_beta ...
!> \param Lag_nu_a_2_beta ...
! **************************************************************************************************
   SUBROUTINE update_lagrangian(mat_munu, matrix_P_inu, Lag_mu_i_1, &
                                G_P_ia, mo_coeff_o, Lag_nu_a_2, &
                                eps_filter, &
                                matrix_P_inu_beta, Lag_mu_i_1_beta, G_P_ia_beta, mo_coeff_o_beta, Lag_nu_a_2_beta)
      TYPE(dbcsr_p_type), INTENT(IN)                     :: mat_munu
      TYPE(dbcsr_p_type), INTENT(INOUT)                  :: matrix_P_inu, Lag_mu_i_1, G_P_ia
      TYPE(dbcsr_type), POINTER                          :: mo_coeff_o
      TYPE(dbcsr_p_type), INTENT(INOUT)                  :: Lag_nu_a_2
      REAL(KIND=dp), INTENT(IN)                          :: eps_filter
      TYPE(dbcsr_p_type), INTENT(INOUT), OPTIONAL        :: matrix_P_inu_beta, Lag_mu_i_1_beta, &
                                                            G_P_ia_beta
      TYPE(dbcsr_type), OPTIONAL, POINTER                :: mo_coeff_o_beta
      TYPE(dbcsr_p_type), INTENT(INOUT), OPTIONAL        :: Lag_nu_a_2_beta

      CHARACTER(LEN=*), PARAMETER :: routineN = 'update_lagrangian', &
         routineP = moduleN//':'//routineN

      INTEGER                                            :: handle
      LOGICAL                                            :: alpha_beta

      ! update lagrangian
      CALL timeset(routineN//"_Lag", handle)

      alpha_beta = PRESENT(matrix_P_inu_beta) .AND. PRESENT(Lag_mu_i_1_beta) .AND. PRESENT(G_P_ia_beta) &
                   .AND. PRESENT(mo_coeff_o_beta) .AND. PRESENT(Lag_nu_a_2_beta)

      CALL update_lagrangian_alpha_beta(mat_munu, matrix_P_inu, Lag_mu_i_1, &
                                        G_P_ia, mo_coeff_o, Lag_nu_a_2, eps_filter)
      IF (alpha_beta) THEN
         CALL update_lagrangian_alpha_beta(mat_munu, matrix_P_inu_beta, Lag_mu_i_1_beta, &
                                           G_P_ia_beta, mo_coeff_o_beta, Lag_nu_a_2_beta, eps_filter)
      ENDIF

      CALL timestop(handle)

   END SUBROUTINE update_lagrangian

! **************************************************************************************************
!> \brief ...
!> \param mat_munu ...
!> \param matrix_P_inu ...
!> \param Lag_mu_i_1 ...
!> \param G_P_ia ...
!> \param mo_coeff_o ...
!> \param Lag_nu_a_2 ...
!> \param eps_filter ...
! **************************************************************************************************
   SUBROUTINE update_lagrangian_alpha_beta(mat_munu, matrix_P_inu, Lag_mu_i_1, &
                                           G_P_ia, mo_coeff_o, Lag_nu_a_2, &
                                           eps_filter)
      TYPE(dbcsr_p_type), INTENT(IN)                     :: mat_munu
      TYPE(dbcsr_p_type), INTENT(INOUT)                  :: matrix_P_inu, Lag_mu_i_1, G_P_ia
      TYPE(dbcsr_type), POINTER                          :: mo_coeff_o
      TYPE(dbcsr_p_type), INTENT(INOUT)                  :: Lag_nu_a_2
      REAL(KIND=dp), INTENT(IN)                          :: eps_filter

      CHARACTER(LEN=*), PARAMETER :: routineN = 'update_lagrangian_alpha_beta', &
         routineP = moduleN//':'//routineN

      ! first contract mat_munu with the half back transformed Gamma_i_nu
      ! in order to update Lag_mu_i_1
      CALL dbcsr_multiply("N", "N", 1.0_dp, mat_munu%matrix, matrix_P_inu%matrix, &
                          1.0_dp, Lag_mu_i_1%matrix, filter_eps=eps_filter)

      ! transform first index of mat_munu and store the result into matrix_P_inu
      CALL dbcsr_set(matrix_P_inu%matrix, 0.0_dp)
      CALL dbcsr_multiply("N", "N", 1.0_dp, mat_munu%matrix, mo_coeff_o, &
                          0.0_dp, matrix_P_inu%matrix, filter_eps=eps_filter)

      ! contract the transformend matrix_P_inu with the untransformend Gamma_i_a
      ! in order to update Lag_nu_a_2
      CALL dbcsr_multiply("N", "N", -1.0_dp, matrix_P_inu%matrix, G_P_ia%matrix, &
                          1.0_dp, Lag_nu_a_2%matrix, filter_eps=eps_filter)

      ! release the actual gamma_P_ia
      CALL dbcsr_release(G_P_ia%matrix)
      DEALLOCATE (G_P_ia%matrix)

   END SUBROUTINE update_lagrangian_alpha_beta

! **************************************************************************************************
!> \brief ...
!> \param qs_env ...
!> \param mp2_env ...
!> \param mo_coeff ...
!> \param homo ...
!> \param virtual ...
!> \param dimen ...
!> \param para_env ...
!> \param para_env_sub ...
!> \param Eigenval ...
!> \param L1_mu_i ...
!> \param L2_nu_a ...
!> \param factor ...
!> \param alpha_beta ...
! **************************************************************************************************
   SUBROUTINE create_W_P(qs_env, mp2_env, mo_coeff, homo, virtual, dimen, para_env, para_env_sub, &
                         Eigenval, L1_mu_i, L2_nu_a, factor, alpha_beta)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(mp2_type), POINTER                            :: mp2_env
      TYPE(cp_fm_type), POINTER                          :: mo_coeff
      INTEGER                                            :: homo, virtual, dimen
      TYPE(cp_para_env_type), POINTER                    :: para_env, para_env_sub
      REAL(KIND=dp), DIMENSION(:)                        :: Eigenval
      TYPE(cp_fm_type), POINTER                          :: L1_mu_i, L2_nu_a
      REAL(KIND=dp)                                      :: factor
      LOGICAL                                            :: alpha_beta

      CHARACTER(LEN=*), PARAMETER :: routineN = 'create_W_P', routineP = moduleN//':'//routineN

      INTEGER :: color_exchange, comm_exchange, dummy_proc, handle, handle2, handle3, i, i_global, &
         i_local, iiB, iii, iproc, itmp(2), j_global, j_local, jjB, max_col_size, max_row_size, &
         my_B_virtual_end, my_B_virtual_start, mypcol, mypcol_1i, mypcol_2a, myprow, myprow_1i, &
         myprow_2a, ncol_block, ncol_block_1i, ncol_block_2a, ncol_local, ncol_local_1i, &
         ncol_local_2a, npcol, npcol_1i, npcol_2a, nprow, nprow_1i, nprow_2a, nrow_block, &
         nrow_block_1i, nrow_block_2a, nrow_local, nrow_local_1i, nrow_local_2a, number_of_rec, &
         number_of_send, proc_receive, proc_receive_static, proc_send, proc_send_ex
      INTEGER :: proc_send_static, proc_send_sub, proc_shift, rec_col_size, rec_counter, &
         rec_row_size, send_col_size, send_counter, send_pcol, send_prow, send_row_size, &
         size_rec_buffer, size_send_buffer
      INTEGER, ALLOCATABLE, DIMENSION(:) :: iii_vet, map_rec_size, map_send_size, pos_info, &
         pos_info_ex, proc_2_send_pos, proc_map_ex, req_send, sub_proc_map
      INTEGER, ALLOCATABLE, DIMENSION(:, :)              :: grid_2_mepos, mepos_2_grid, &
                                                            mepos_2_grid_1i, mepos_2_grid_2a, &
                                                            sizes, sizes_1i, sizes_2a
      INTEGER, ALLOCATABLE, DIMENSION(:, :, :)           :: col_indeces_info_1i, &
                                                            col_indeces_info_2a, &
                                                            row_indeces_info_1i, &
                                                            row_indeces_info_2a
      INTEGER, DIMENSION(:), POINTER                     :: col_indices, col_indices_1i, &
                                                            col_indices_2a, row_indices, &
                                                            row_indices_1i, row_indices_2a
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :)        :: ab_rec, ab_send, mat_deb, mat_rec, &
                                                            mat_send
      TYPE(cp_blacs_env_type), POINTER                   :: blacs_env
      TYPE(cp_fm_struct_type), POINTER                   :: fm_struct_tmp
      TYPE(cp_fm_type), POINTER                          :: fm_P_ij, L_mu_q
      TYPE(cp_para_env_type), POINTER                    :: para_env_exchange
      TYPE(integ_mat_buffer_type), ALLOCATABLE, &
         DIMENSION(:)                                    :: buffer_rec, buffer_send
      TYPE(integ_mat_buffer_type_2D), ALLOCATABLE, &
         DIMENSION(:)                                    :: buffer_cyclic

      CALL timeset(routineN, handle)

      ! create the globally distributed mixed lagrangian
      NULLIFY (blacs_env)
      CALL get_qs_env(qs_env, blacs_env=blacs_env)

      NULLIFY (L_mu_q, fm_struct_tmp)
      CALL cp_fm_struct_create(fm_struct_tmp, para_env=para_env, context=blacs_env, &
                               nrow_global=dimen, ncol_global=dimen)
      CALL cp_fm_create(L_mu_q, fm_struct_tmp, name="Lag_mu_q")
      CALL cp_fm_struct_release(fm_struct_tmp)
      CALL cp_fm_set_all(L_mu_q, 0.0_dp)

      ! create all information array
      ALLOCATE (pos_info(0:para_env%num_pe-1))
      pos_info = 0
      pos_info(para_env%mepos) = para_env_sub%mepos
      CALL mp_sum(pos_info, para_env%group)

      ALLOCATE (sub_proc_map(-para_env_sub%num_pe:2*para_env_sub%num_pe-1))
      sub_proc_map = 0
      DO i = 0, para_env_sub%num_pe-1
         sub_proc_map(i) = i
         sub_proc_map(-i-1) = para_env_sub%num_pe-i-1
         sub_proc_map(para_env_sub%num_pe+i) = i
      END DO

      ! get matrix information for the global
      CALL cp_fm_get_info(matrix=L_mu_q, &
                          nrow_local=nrow_local, &
                          ncol_local=ncol_local, &
                          row_indices=row_indices, &
                          col_indices=col_indices, &
                          nrow_block=nrow_block, &
                          ncol_block=ncol_block)
      myprow = L_mu_q%matrix_struct%context%mepos(1)
      mypcol = L_mu_q%matrix_struct%context%mepos(2)
      nprow = L_mu_q%matrix_struct%context%num_pe(1)
      npcol = L_mu_q%matrix_struct%context%num_pe(2)

      ALLOCATE (grid_2_mepos(0:nprow-1, 0:npcol-1))
      grid_2_mepos = 0
      grid_2_mepos(myprow, mypcol) = para_env%mepos
      CALL mp_sum(grid_2_mepos, para_env%group)

      ! get matrix information for L1_mu_i
      CALL cp_fm_get_info(matrix=L1_mu_i, &
                          nrow_local=nrow_local_1i, &
                          ncol_local=ncol_local_1i, &
                          row_indices=row_indices_1i, &
                          col_indices=col_indices_1i, &
                          nrow_block=nrow_block_1i, &
                          ncol_block=ncol_block_1i)
      myprow_1i = L1_mu_i%matrix_struct%context%mepos(1)
      mypcol_1i = L1_mu_i%matrix_struct%context%mepos(2)
      nprow_1i = L1_mu_i%matrix_struct%context%num_pe(1)
      npcol_1i = L1_mu_i%matrix_struct%context%num_pe(2)

      ALLOCATE (mepos_2_grid_1i(0:para_env_sub%num_pe-1, 2))
      mepos_2_grid_1i = 0
      mepos_2_grid_1i(para_env_sub%mepos, 1) = myprow_1i
      mepos_2_grid_1i(para_env_sub%mepos, 2) = mypcol_1i
      CALL mp_sum(mepos_2_grid_1i, para_env_sub%group)

      ALLOCATE (sizes_1i(2, 0:para_env_sub%num_pe-1))
      sizes_1i = 0
      sizes_1i(1, para_env_sub%mepos) = nrow_local_1i
      sizes_1i(2, para_env_sub%mepos) = ncol_local_1i
      CALL mp_sum(sizes_1i, para_env_sub%group)

      ! get matrix information for L2_nu_a
      CALL cp_fm_get_info(matrix=L2_nu_a, &
                          nrow_local=nrow_local_2a, &
                          ncol_local=ncol_local_2a, &
                          row_indices=row_indices_2a, &
                          col_indices=col_indices_2a, &
                          nrow_block=nrow_block_2a, &
                          ncol_block=ncol_block_2a)
      myprow_2a = L2_nu_a%matrix_struct%context%mepos(1)
      mypcol_2a = L2_nu_a%matrix_struct%context%mepos(2)
      nprow_2a = L2_nu_a%matrix_struct%context%num_pe(1)
      npcol_2a = L2_nu_a%matrix_struct%context%num_pe(2)

      ALLOCATE (mepos_2_grid_2a(0:para_env_sub%num_pe-1, 2))
      mepos_2_grid_2a = 0
      mepos_2_grid_2a(para_env_sub%mepos, 1) = myprow_2a
      mepos_2_grid_2a(para_env_sub%mepos, 2) = mypcol_2a
      CALL mp_sum(mepos_2_grid_2a, para_env_sub%group)

      ALLOCATE (sizes_2a(2, 0:para_env_sub%num_pe-1))
      sizes_2a = 0
      sizes_2a(1, para_env_sub%mepos) = nrow_local_2a
      sizes_2a(2, para_env_sub%mepos) = ncol_local_2a
      CALL mp_sum(sizes_2a, para_env_sub%group)

      ! Here we perform a ring communication scheme taking into account
      ! for the sub-group distribution of the source matrices.
      ! as a first step we need to redistribute the data within
      ! the subgroup.
      ! In order to do so we have to allocate the structure
      ! that will hold the local data involved in communication, this
      ! structure will be the same for processes in different subgroups
      ! sharing the same position in the subgroup.
      ! -1) create the exchange para_env
      color_exchange = para_env_sub%mepos
      CALL mp_comm_split_direct(para_env%group, comm_exchange, color_exchange)
      NULLIFY (para_env_exchange)
      CALL cp_para_env_create(para_env_exchange, comm_exchange)
      ! crate the proc maps exchange and info
      ALLOCATE (proc_map_ex(-para_env_exchange%num_pe:2*para_env_exchange%num_pe-1))
      DO i = 0, para_env_exchange%num_pe-1
         proc_map_ex(i) = i
         proc_map_ex(-i-1) = para_env_exchange%num_pe-i-1
         proc_map_ex(para_env_exchange%num_pe+i) = i
      END DO
      ALLOCATE (pos_info_ex(0:para_env%num_pe-1))
      pos_info_ex = 0
      pos_info_ex(para_env%mepos) = para_env_exchange%mepos
      CALL mp_sum(pos_info_ex, para_env%group)
      ALLOCATE (sizes(2, 0:para_env_exchange%num_pe-1))
      sizes = 0
      sizes(1, para_env_exchange%mepos) = nrow_local
      sizes(2, para_env_exchange%mepos) = ncol_local
      CALL mp_sum(sizes, para_env_exchange%group)

      ! 0) store some info about indeces of the fm matrices (subgroup)
      CALL timeset(routineN//"_inx", handle2)
      ! matrix L1_mu_i
      max_row_size = MAXVAL(sizes_1i(1, :))
      max_col_size = MAXVAL(sizes_1i(2, :))
      ALLOCATE (row_indeces_info_1i(2, max_row_size, 0:para_env_sub%num_pe-1))
      ALLOCATE (col_indeces_info_1i(2, max_col_size, 0:para_env_sub%num_pe-1))
      row_indeces_info_1i = 0
      col_indeces_info_1i = 0
      dummy_proc = 0
      ! row
      DO iiB = 1, nrow_local_1i
         i_global = row_indices_1i(iiB)
         send_prow = cp_fm_indxg2p(i_global, nrow_block, dummy_proc, &
                                   L_mu_q%matrix_struct%first_p_pos(1), nprow)
         i_local = cp_fm_indxg2l(i_global, nrow_block, dummy_proc, &
                                 L_mu_q%matrix_struct%first_p_pos(1), nprow)
         row_indeces_info_1i(1, iiB, para_env_sub%mepos) = send_prow
         row_indeces_info_1i(2, iiB, para_env_sub%mepos) = i_local
      END DO
      ! col
      DO jjB = 1, ncol_local_1i
         j_global = col_indices_1i(jjB)
         send_pcol = cp_fm_indxg2p(j_global, ncol_block, dummy_proc, &
                                   L_mu_q%matrix_struct%first_p_pos(2), npcol)
         j_local = cp_fm_indxg2l(j_global, ncol_block, dummy_proc, &
                                 L_mu_q%matrix_struct%first_p_pos(2), npcol)
         col_indeces_info_1i(1, jjB, para_env_sub%mepos) = send_pcol
         col_indeces_info_1i(2, jjB, para_env_sub%mepos) = j_local
      END DO
      CALL mp_sum(row_indeces_info_1i, para_env_sub%group)
      CALL mp_sum(col_indeces_info_1i, para_env_sub%group)

      ! matrix L2_nu_a
      max_row_size = MAXVAL(sizes_2a(1, :))
      max_col_size = MAXVAL(sizes_2a(2, :))
      ALLOCATE (row_indeces_info_2a(2, max_row_size, 0:para_env_sub%num_pe-1))
      ALLOCATE (col_indeces_info_2a(2, max_col_size, 0:para_env_sub%num_pe-1))
      row_indeces_info_2a = 0
      col_indeces_info_2a = 0
      ! row
      DO iiB = 1, nrow_local_2a
         i_global = row_indices_2a(iiB)
         send_prow = cp_fm_indxg2p(i_global, nrow_block, dummy_proc, &
                                   L_mu_q%matrix_struct%first_p_pos(1), nprow)
         i_local = cp_fm_indxg2l(i_global, nrow_block, dummy_proc, &
                                 L_mu_q%matrix_struct%first_p_pos(1), nprow)
         row_indeces_info_2a(1, iiB, para_env_sub%mepos) = send_prow
         row_indeces_info_2a(2, iiB, para_env_sub%mepos) = i_local
      END DO
      ! col
      DO jjB = 1, ncol_local_2a
         j_global = col_indices_2a(jjB)+homo
         send_pcol = cp_fm_indxg2p(j_global, ncol_block, dummy_proc, &
                                   L_mu_q%matrix_struct%first_p_pos(2), npcol)
         j_local = cp_fm_indxg2l(j_global, ncol_block, dummy_proc, &
                                 L_mu_q%matrix_struct%first_p_pos(2), npcol)
         col_indeces_info_2a(1, jjB, para_env_sub%mepos) = send_pcol
         col_indeces_info_2a(2, jjB, para_env_sub%mepos) = j_local
      END DO
      CALL mp_sum(row_indeces_info_2a, para_env_sub%group)
      CALL mp_sum(col_indeces_info_2a, para_env_sub%group)
      CALL timestop(handle2)

      ! 1) define the map for sending data in the subgroup starting with L1_mu_i
      CALL timeset(routineN//"_subinfo", handle2)
      ALLOCATE (map_send_size(0:para_env_sub%num_pe-1))
      map_send_size = 0
      DO jjB = 1, ncol_local_1i
         ! j_global=col_indices_1i(jjB)
         ! send_pcol=cp_fm_indxg2p(j_global,ncol_block,dummy_proc,&
         !                         L_mu_q%matrix_struct%first_p_pos(2),npcol)
         send_pcol = col_indeces_info_1i(1, jjB, para_env_sub%mepos)
         DO iiB = 1, nrow_local_1i
            ! i_global=row_indices_1i(iiB)
            ! send_prow=cp_fm_indxg2p(i_global,nrow_block,dummy_proc,&
            !                         L_mu_q%matrix_struct%first_p_pos(1),nprow)
            send_prow = row_indeces_info_1i(1, iiB, para_env_sub%mepos)
            proc_send = grid_2_mepos(send_prow, send_pcol)
            proc_send_sub = pos_info(proc_send)
            map_send_size(proc_send_sub) = map_send_size(proc_send_sub)+1
         END DO
      END DO
      ! and the same for L2_nu_a
      DO jjB = 1, ncol_local_2a
         ! j_global=col_indices_2a(jjB)+homo
         ! send_pcol=cp_fm_indxg2p(j_global,ncol_block,dummy_proc,&
         !                         L_mu_q%matrix_struct%first_p_pos(2),npcol)
         send_pcol = col_indeces_info_2a(1, jjB, para_env_sub%mepos)
         DO iiB = 1, nrow_local_2a
            ! i_global=row_indices_2a(iiB)
            ! send_prow=cp_fm_indxg2p(i_global,nrow_block,dummy_proc,&
            !                         L_mu_q%matrix_struct%first_p_pos(1),nprow)
            send_prow = row_indeces_info_2a(1, iiB, para_env_sub%mepos)
            proc_send = grid_2_mepos(send_prow, send_pcol)
            proc_send_sub = pos_info(proc_send)
            map_send_size(proc_send_sub) = map_send_size(proc_send_sub)+1
         END DO
      END DO
      ! and exchange data in order to create map_rec_size
      ALLOCATE (map_rec_size(0:para_env_sub%num_pe-1))
      map_rec_size = 0
      CALL mp_alltoall(map_send_size, map_rec_size, 1, para_env_sub%group)
      CALL timestop(handle2)

      ! 2) reorder data in sending buffer
      CALL timeset(routineN//"_sub_Bsend", handle2)
      ! count the number of messages (include myself)
      number_of_send = 0
      DO proc_shift = 0, para_env_sub%num_pe-1
         proc_send = sub_proc_map(para_env_sub%mepos+proc_shift)
         IF (map_send_size(proc_send) > 0) THEN
            number_of_send = number_of_send+1
         END IF
      END DO
      ! allocate the structure that will hold the messages to be sent
      ALLOCATE (buffer_send(number_of_send))
      send_counter = 0
      ALLOCATE (proc_2_send_pos(0:para_env_sub%num_pe-1))
      proc_2_send_pos = 0
      DO proc_shift = 0, para_env_sub%num_pe-1
         proc_send = sub_proc_map(para_env_sub%mepos+proc_shift)
         size_send_buffer = map_send_size(proc_send)
         IF (map_send_size(proc_send) > 0) THEN
            send_counter = send_counter+1
            ! allocate the sending buffer (msg)
            ALLOCATE (buffer_send(send_counter)%msg(size_send_buffer))
            buffer_send(send_counter)%msg = 0.0_dp
            buffer_send(send_counter)%proc = proc_send
            proc_2_send_pos(proc_send) = send_counter
         END IF
      END DO
      ! loop over the locally held data and fill the buffer_send
      ! for doing that we need an array that keep track if the
      ! sequential increase of the index for each message
      ALLOCATE (iii_vet(number_of_send))
      iii_vet = 0
      DO jjB = 1, ncol_local_1i
         ! j_global=col_indices_1i(jjB)
         ! send_pcol=cp_fm_indxg2p(j_global,ncol_block,dummy_proc,&
         !                         L_mu_q%matrix_struct%first_p_pos(2),npcol)
         send_pcol = col_indeces_info_1i(1, jjB, para_env_sub%mepos)
         DO iiB = 1, nrow_local_1i
            ! i_global=row_indices_1i(iiB)
            ! send_prow=cp_fm_indxg2p(i_global,nrow_block,dummy_proc,&
            !                         L_mu_q%matrix_struct%first_p_pos(1),nprow)
            send_prow = row_indeces_info_1i(1, iiB, para_env_sub%mepos)
            proc_send = grid_2_mepos(send_prow, send_pcol)
            proc_send_sub = pos_info(proc_send)
            send_counter = proc_2_send_pos(proc_send_sub)
            iii_vet(send_counter) = iii_vet(send_counter)+1
            iii = iii_vet(send_counter)
            buffer_send(send_counter)%msg(iii) = L1_mu_i%local_data(iiB, jjB)
         END DO
      END DO
      ! release the local data of L1_mu_i
      DEALLOCATE (L1_mu_i%local_data)
      ! and the same for L2_nu_a
      DO jjB = 1, ncol_local_2a
         ! j_global=col_indices_2a(jjB)+homo
         ! send_pcol=cp_fm_indxg2p(j_global,ncol_block,dummy_proc,&
         !                         L_mu_q%matrix_struct%first_p_pos(2),npcol)
         send_pcol = col_indeces_info_2a(1, jjB, para_env_sub%mepos)
         DO iiB = 1, nrow_local_2a
            ! i_global=row_indices_2a(iiB)
            ! send_prow=cp_fm_indxg2p(i_global,nrow_block,dummy_proc,&
            !                         L_mu_q%matrix_struct%first_p_pos(1),nprow)
            send_prow = row_indeces_info_2a(1, iiB, para_env_sub%mepos)
            proc_send = grid_2_mepos(send_prow, send_pcol)
            proc_send_sub = pos_info(proc_send)
            send_counter = proc_2_send_pos(proc_send_sub)
            iii_vet(send_counter) = iii_vet(send_counter)+1
            iii = iii_vet(send_counter)
            buffer_send(send_counter)%msg(iii) = L2_nu_a%local_data(iiB, jjB)
         END DO
      END DO
      DEALLOCATE (L2_nu_a%local_data)
      DEALLOCATE (proc_2_send_pos)
      DEALLOCATE (iii_vet)
      CALL timestop(handle2)

      ! 3) create the buffer for receive, post the message with irecv
      !    and send the messages with mp_isend
      CALL timeset(routineN//"_sub_isendrecv", handle2)
      ! count the number of messages to be received
      number_of_rec = 0
      DO proc_shift = 0, para_env_sub%num_pe-1
         proc_receive = sub_proc_map(para_env_sub%mepos-proc_shift)
         IF (map_rec_size(proc_receive) > 0) THEN
            number_of_rec = number_of_rec+1
         END IF
      END DO
      ALLOCATE (buffer_rec(number_of_rec))
      rec_counter = 0
      DO proc_shift = 0, para_env_sub%num_pe-1
         proc_receive = sub_proc_map(para_env_sub%mepos-proc_shift)
         size_rec_buffer = map_rec_size(proc_receive)
         IF (map_rec_size(proc_receive) > 0) THEN
            rec_counter = rec_counter+1
            ! prepare the buffer for receive
            ALLOCATE (buffer_rec(rec_counter)%msg(size_rec_buffer))
            buffer_rec(rec_counter)%msg = 0.0_dp
            buffer_rec(rec_counter)%proc = proc_receive
            ! post the message to be received (not need to send to myself)
            IF (proc_receive /= para_env_sub%mepos) THEN
               CALL mp_irecv(buffer_rec(rec_counter)%msg, proc_receive, para_env_sub%group, &
                             buffer_rec(rec_counter)%msg_req)
            END IF
         END IF
      END DO
      ! send messages
      ALLOCATE (req_send(number_of_send))
      req_send = mp_request_null
      send_counter = 0
      DO proc_shift = 0, para_env_sub%num_pe-1
         proc_send = sub_proc_map(para_env_sub%mepos+proc_shift)
         IF (map_send_size(proc_send) > 0) THEN
            send_counter = send_counter+1
            IF (proc_send == para_env_sub%mepos) THEN
               buffer_rec(send_counter)%msg = buffer_send(send_counter)%msg
            ELSE
               CALL mp_isend(buffer_send(send_counter)%msg, proc_send, para_env_sub%group, &
                             buffer_send(send_counter)%msg_req)
               req_send(send_counter) = buffer_send(send_counter)%msg_req
            END IF
         END IF
      END DO
      DEALLOCATE (map_send_size)
      CALL timestop(handle2)

      ! 4) (if memory is a problem we should move this part after point 5)
      !    Here we create the new buffer for cyclic(ring) communication and
      !    we fill it with the data received from the other member of the
      !    subgroup
      CALL timeset(routineN//"_Bcyclic", handle2)
      ! first allocata new structure
      ALLOCATE (buffer_cyclic(0:para_env_exchange%num_pe-1))
      DO iproc = 0, para_env_exchange%num_pe-1
         rec_row_size = sizes(1, iproc)
         rec_col_size = sizes(2, iproc)
         ALLOCATE (buffer_cyclic(iproc)%msg(rec_row_size, rec_col_size))
         buffer_cyclic(iproc)%msg = 0.0_dp
      END DO
      ! now collect data from other member of the subgroup and fill
      ! buffer_cyclic
      rec_counter = 0
      DO proc_shift = 0, para_env_sub%num_pe-1
         proc_receive = sub_proc_map(para_env_sub%mepos-proc_shift)
         size_rec_buffer = map_rec_size(proc_receive)
         IF (map_rec_size(proc_receive) > 0) THEN
            rec_counter = rec_counter+1

            ! wait for the message
            IF (proc_receive /= para_env_sub%mepos) CALL mp_wait(buffer_rec(rec_counter)%msg_req)

            CALL timeset(routineN//"_fill", handle3)
            iii = 0
            DO jjB = 1, sizes_1i(2, proc_receive)
               ! j_global=cp_fm_indxl2g(jjB,ncol_block_1i,mepos_2_grid_1i(proc_receive,2),&
               !                        L1_mu_i%matrix_struct%first_p_pos(2),npcol_1i)
               ! send_pcol=cp_fm_indxg2p(j_global,ncol_block,dummy_proc,&
               !                         L_mu_q%matrix_struct%first_p_pos(2),npcol)
               ! j_local=cp_fm_indxg2l(j_global,ncol_block,dummy_proc,&
               !                       L_mu_q%matrix_struct%first_p_pos(2),npcol)
               send_pcol = col_indeces_info_1i(1, jjB, proc_receive)
               j_local = col_indeces_info_1i(2, jjB, proc_receive)
               DO iiB = 1, sizes_1i(1, proc_receive)
                  ! i_global=cp_fm_indxl2g(iiB,nrow_block_1i,mepos_2_grid_1i(proc_receive,1),&
                  !                        L1_mu_i%matrix_struct%first_p_pos(1),nprow_1i)
                  ! send_prow=cp_fm_indxg2p(i_global,nrow_block,dummy_proc,&
                  !                         L_mu_q%matrix_struct%first_p_pos(1),nprow)
                  send_prow = row_indeces_info_1i(1, iiB, proc_receive)
                  proc_send = grid_2_mepos(send_prow, send_pcol)
                  proc_send_sub = pos_info(proc_send)
                  IF (proc_send_sub /= para_env_sub%mepos) CYCLE
                  iii = iii+1
                  ! i_local=cp_fm_indxg2l(i_global,nrow_block,dummy_proc,&
                  !                       L_mu_q%matrix_struct%first_p_pos(1),nprow)
                  i_local = row_indeces_info_1i(2, iiB, proc_receive)
                  proc_send_ex = pos_info_ex(proc_send)
                  buffer_cyclic(proc_send_ex)%msg(i_local, j_local) = buffer_rec(rec_counter)%msg(iii)
               END DO
            END DO
            ! and the same for L2_nu_a
            DO jjB = 1, sizes_2a(2, proc_receive)
               ! j_global=cp_fm_indxl2g(jjB,ncol_block_2a,mepos_2_grid_2a(proc_receive,2),&
               !                        L2_nu_a%matrix_struct%first_p_pos(2),npcol_2a)+homo
               ! send_pcol=cp_fm_indxg2p(j_global,ncol_block,dummy_proc,&
               !                         L_mu_q%matrix_struct%first_p_pos(2),npcol)
               ! j_local=cp_fm_indxg2l(j_global,ncol_block,dummy_proc,&
               !                       L_mu_q%matrix_struct%first_p_pos(2),npcol)
               send_pcol = col_indeces_info_2a(1, jjB, proc_receive)
               j_local = col_indeces_info_2a(2, jjB, proc_receive)
               DO iiB = 1, sizes_2a(1, proc_receive)
                  ! i_global=cp_fm_indxl2g(iiB,nrow_block_2a,mepos_2_grid_2a(proc_receive,1),&
                  !                        L2_nu_a%matrix_struct%first_p_pos(1),nprow_2a)
                  ! send_prow=cp_fm_indxg2p(i_global,nrow_block,dummy_proc,&
                  !                         L_mu_q%matrix_struct%first_p_pos(1),nprow)
                  send_prow = row_indeces_info_2a(1, iiB, proc_receive)
                  proc_send = grid_2_mepos(send_prow, send_pcol)
                  proc_send_sub = pos_info(proc_send)
                  IF (proc_send_sub /= para_env_sub%mepos) CYCLE
                  iii = iii+1
                  ! i_local=cp_fm_indxg2l(i_global,nrow_block,dummy_proc,&
                  !                       L_mu_q%matrix_struct%first_p_pos(1),nprow)
                  i_local = row_indeces_info_2a(2, iiB, proc_receive)
                  proc_send_ex = pos_info_ex(proc_send)
                  buffer_cyclic(proc_send_ex)%msg(i_local, j_local) = buffer_rec(rec_counter)%msg(iii)
               END DO
            END DO
            CALL timestop(handle3)

            ! deallocate the received message
            DEALLOCATE (buffer_rec(rec_counter)%msg)
         END IF
      END DO
      DEALLOCATE (row_indeces_info_1i)
      DEALLOCATE (col_indeces_info_1i)
      DEALLOCATE (row_indeces_info_2a)
      DEALLOCATE (col_indeces_info_2a)
      DEALLOCATE (buffer_rec)
      DEALLOCATE (map_rec_size)
      CALL timestop(handle2)

      ! 5)  Wait for all messeges to be sent in the subgroup
      CALL timeset(routineN//"_sub_waitall", handle2)
      CALL mp_waitall(req_send(:))
      DO send_counter = 1, number_of_send
         DEALLOCATE (buffer_send(send_counter)%msg)
      END DO
      DEALLOCATE (buffer_send)
      DEALLOCATE (req_send)
      CALL timestop(handle2)

      ! 6) Start with ring communication
      CALL timeset(routineN//"_ring", handle2)
      proc_send_static = proc_map_ex(para_env_exchange%mepos+1)
      proc_receive_static = proc_map_ex(para_env_exchange%mepos-1)
      max_row_size = MAXVAL(sizes(1, :))
      max_col_size = MAXVAL(sizes(2, :))
      ALLOCATE (mat_send(max_row_size, max_col_size))
      ALLOCATE (mat_rec(max_row_size, max_col_size))
      mat_send = 0.0_dp
      mat_send(1:nrow_local, 1:ncol_local) = buffer_cyclic(para_env_exchange%mepos)%msg(:, :)
      DEALLOCATE (buffer_cyclic(para_env_exchange%mepos)%msg)
      DO proc_shift = 1, para_env_exchange%num_pe-1
         proc_send = proc_map_ex(para_env_exchange%mepos+proc_shift)
         proc_receive = proc_map_ex(para_env_exchange%mepos-proc_shift)

         rec_row_size = sizes(1, proc_receive)
         rec_col_size = sizes(2, proc_receive)

         mat_rec = 0.0_dp
         CALL mp_sendrecv(mat_send, proc_send_static, &
                          mat_rec, proc_receive_static, &
                          para_env_exchange%group)

         mat_send = 0.0_dp
         mat_send(1:rec_row_size, 1:rec_col_size) = mat_rec(1:rec_row_size, 1:rec_col_size)+ &
                                                    buffer_cyclic(proc_receive)%msg(:, :)

         DEALLOCATE (buffer_cyclic(proc_receive)%msg)
      END DO
      ! and finaly
      CALL mp_sendrecv(mat_send, proc_send_static, &
                       mat_rec, proc_receive_static, &
                       para_env_exchange%group)
      L_mu_q%local_data(1:nrow_local, 1:ncol_local) = mat_rec(1:nrow_local, 1:ncol_local)
      DEALLOCATE (buffer_cyclic)
      DEALLOCATE (mat_send)
      DEALLOCATE (mat_rec)
      CALL timestop(handle2)

      DEALLOCATE (proc_map_ex)
      ! release para_env_exchange
      CALL cp_para_env_release(para_env_exchange)

      CALL cp_fm_release(L1_mu_i)
      CALL cp_fm_release(L2_nu_a)
      DEALLOCATE (pos_info_ex)
      DEALLOCATE (grid_2_mepos)
      DEALLOCATE (mepos_2_grid_1i)
      DEALLOCATE (mepos_2_grid_2a)
      DEALLOCATE (sizes)
      DEALLOCATE (sizes_1i)
      DEALLOCATE (sizes_2a)

      ! update the P_ij block of P_MP2 with the
      ! non-singular ij pairs
      CALL timeset(routineN//"_Pij", handle2)
      NULLIFY (fm_P_ij, fm_struct_tmp)
      CALL cp_fm_struct_create(fm_struct_tmp, para_env=para_env, context=blacs_env, &
                               nrow_global=homo, ncol_global=homo)
      CALL cp_fm_create(fm_P_ij, fm_struct_tmp, name="fm_P_ij")
      CALL cp_fm_struct_release(fm_struct_tmp)
      CALL cp_fm_set_all(fm_P_ij, 0.0_dp)

      CALL cp_gemm('T', 'N', homo, homo, dimen, 1.0_dp, &
                   mo_coeff, L_mu_q, 0.0_dp, fm_P_ij, &
                   a_first_col=1, &
                   a_first_row=1, &
                   b_first_col=1, &
                   b_first_row=1, &
                   c_first_col=1, &
                   c_first_row=1)
      ! don't know if it is better to transpose (for communication reasons)
      ! or just recompute the transposed matrix
      CALL cp_gemm('T', 'N', homo, homo, dimen, -2.0_dp, &
                   L_mu_q, mo_coeff, 2.0_dp, fm_P_ij, &
                   a_first_col=1, &
                   a_first_row=1, &
                   b_first_col=1, &
                   b_first_row=1, &
                   c_first_col=1, &
                   c_first_row=1)
      ! we have it, update P_ij local
      CALL cp_fm_get_info(matrix=fm_P_ij, &
                          nrow_local=nrow_local, &
                          ncol_local=ncol_local, &
                          row_indices=row_indices, &
                          col_indices=col_indices)
      DO jjB = 1, ncol_local
         j_global = col_indices(jjB)
         DO iiB = 1, nrow_local
            i_global = row_indices(iiB)
            ! diagonal elements alreasy updated
            IF (j_global == i_global) CYCLE
            ! check if the given element is above the threshold
            IF (ABS(Eigenval(j_global)-Eigenval(i_global)) < mp2_env%ri_mp2%eps_canonical) CYCLE
            IF (.NOT. alpha_beta) THEN
               mp2_env%ri_grad%P_ij(i_global, j_global) = &
                  factor*fm_P_ij%local_data(iiB, jjB)/(Eigenval(j_global)-Eigenval(i_global))
            ELSE
               mp2_env%ri_grad%P_ij_beta(i_global, j_global) = &
                  factor*fm_P_ij%local_data(iiB, jjB)/(Eigenval(j_global)-Eigenval(i_global))
            ENDIF
         END DO
      END DO
      ! release fm_P_ij
      CALL cp_fm_release(fm_P_ij)
      ! mp_sum it (we can avoid mp_sum, but for now let's keep it easy)
      IF (.NOT. alpha_beta) THEN
         CALL mp_sum(mp2_env%ri_grad%P_ij, para_env%group)
      ELSE
         CALL mp_sum(mp2_env%ri_grad%P_ij_beta, para_env%group)
      ENDIF
      CALL timestop(handle2)

      ! Now create and fill the P matrix (MO)
      ! FOR NOW WE ASSUME P_ab AND P_ij ARE REPLICATED OVER EACH MPI RANK
      IF (.NOT. alpha_beta) THEN
         CALL timeset(routineN//"_PMO", handle2)
         NULLIFY (mp2_env%ri_grad%P_mo, fm_struct_tmp)
         CALL cp_fm_struct_create(fm_struct_tmp, para_env=para_env, context=blacs_env, &
                                  nrow_global=dimen, ncol_global=dimen)
         CALL cp_fm_create(mp2_env%ri_grad%P_mo, fm_struct_tmp, name="P_MP2_MO")
         CALL cp_fm_set_all(mp2_env%ri_grad%P_mo, 0.0_dp)
      ELSE
         CALL timeset(routineN//"_PMO", handle2)
         NULLIFY (mp2_env%ri_grad%P_mo_beta, fm_struct_tmp)
         CALL cp_fm_struct_create(fm_struct_tmp, para_env=para_env, context=blacs_env, &
                                  nrow_global=dimen, ncol_global=dimen)
         CALL cp_fm_create(mp2_env%ri_grad%P_mo_beta, fm_struct_tmp, name="P_MP2_MO")
         CALL cp_fm_set_all(mp2_env%ri_grad%P_mo_beta, 0.0_dp)
      ENDIF

      ! start with the (easy) occ-occ block and locally held P_ab elements
      itmp = get_limit(virtual, para_env_sub%num_pe, para_env_sub%mepos)
      my_B_virtual_start = itmp(1)
      my_B_virtual_end = itmp(2)

      IF (.NOT. alpha_beta) THEN
         CALL cp_fm_get_info(matrix=mp2_env%ri_grad%P_mo, &
                             nrow_local=nrow_local, &
                             ncol_local=ncol_local, &
                             row_indices=row_indices, &
                             col_indices=col_indices, &
                             nrow_block=nrow_block, &
                             ncol_block=ncol_block)
         myprow = mp2_env%ri_grad%P_mo%matrix_struct%context%mepos(1)
         mypcol = mp2_env%ri_grad%P_mo%matrix_struct%context%mepos(2)
         nprow = mp2_env%ri_grad%P_mo%matrix_struct%context%num_pe(1)
         npcol = mp2_env%ri_grad%P_mo%matrix_struct%context%num_pe(2)
      ELSE
         CALL cp_fm_get_info(matrix=mp2_env%ri_grad%P_mo_beta, &
                             nrow_local=nrow_local, &
                             ncol_local=ncol_local, &
                             row_indices=row_indices, &
                             col_indices=col_indices, &
                             nrow_block=nrow_block, &
                             ncol_block=ncol_block)
         myprow = mp2_env%ri_grad%P_mo_beta%matrix_struct%context%mepos(1)
         mypcol = mp2_env%ri_grad%P_mo_beta%matrix_struct%context%mepos(2)
         nprow = mp2_env%ri_grad%P_mo_beta%matrix_struct%context%num_pe(1)
         npcol = mp2_env%ri_grad%P_mo_beta%matrix_struct%context%num_pe(2)
      ENDIF

      DO jjB = 1, ncol_local
         j_global = col_indices(jjB)
         DO iiB = 1, nrow_local
            i_global = row_indices(iiB)
            IF (.NOT. alpha_beta) THEN
               IF (i_global <= homo .AND. j_global <= homo) THEN
                  mp2_env%ri_grad%P_mo%local_data(iiB, jjB) = mp2_env%ri_grad%P_ij(i_global, j_global)
               END IF
               IF ((my_B_virtual_start <= i_global-homo .AND. i_global-homo <= my_B_virtual_end) .AND. (j_global > homo)) THEN
                  mp2_env%ri_grad%P_mo%local_data(iiB, jjB) = &
                     mp2_env%ri_grad%P_ab(i_global-homo-my_B_virtual_start+1, j_global-homo)
               END IF
            ELSE
               IF (i_global <= homo .AND. j_global <= homo) THEN
                  mp2_env%ri_grad%P_mo_beta%local_data(iiB, jjB) = mp2_env%ri_grad%P_ij_beta(i_global, j_global)
               END IF
               IF ((my_B_virtual_start <= i_global-homo .AND. i_global-homo <= my_B_virtual_end) .AND. (j_global > homo)) THEN
                  mp2_env%ri_grad%P_mo_beta%local_data(iiB, jjB) = &
                     mp2_env%ri_grad%P_ab_beta(i_global-homo-my_B_virtual_start+1, j_global-homo)
               END IF
            ENDIF
         END DO
      END DO
      ! deallocate the local P_ij
      IF (.NOT. alpha_beta) THEN
         DEALLOCATE (mp2_env%ri_grad%P_ij)
      ELSE
         DEALLOCATE (mp2_env%ri_grad%P_ij_beta)
      ENDIF

      ! send around the sub_group the local data and check if we
      ! have to update our block with external elements
      ALLOCATE (mepos_2_grid(0:para_env_sub%num_pe-1, 2))
      mepos_2_grid = 0
      mepos_2_grid(para_env_sub%mepos, 1) = myprow
      mepos_2_grid(para_env_sub%mepos, 2) = mypcol
      CALL mp_sum(mepos_2_grid, para_env_sub%group)

      ALLOCATE (sizes(2, 0:para_env_sub%num_pe-1))
      sizes = 0
      sizes(1, para_env_sub%mepos) = nrow_local
      sizes(2, para_env_sub%mepos) = ncol_local
      CALL mp_sum(sizes, para_env_sub%group)

      ALLOCATE (ab_rec(nrow_local, ncol_local))
      DO proc_shift = 1, para_env_sub%num_pe-1
         proc_send = sub_proc_map(para_env_sub%mepos+proc_shift)
         proc_receive = sub_proc_map(para_env_sub%mepos-proc_shift)

         send_prow = mepos_2_grid(proc_send, 1)
         send_pcol = mepos_2_grid(proc_send, 2)

         send_row_size = sizes(1, proc_send)
         send_col_size = sizes(2, proc_send)

         ALLOCATE (ab_send(send_row_size, send_col_size))
         ab_send = 0.0_dp

         ! first loop over row since in this way we can cycle
         DO iiB = 1, send_row_size
            i_global = cp_fm_indxl2g(iiB, nrow_block, send_prow, &
                                     mp2_env%ri_grad%P_mo%matrix_struct%first_p_pos(1), nprow)
            IF (i_global <= homo) CYCLE
            i_global = i_global-homo
            IF (.NOT. (my_B_virtual_start <= i_global .AND. i_global <= my_B_virtual_end)) CYCLE
            DO jjB = 1, send_col_size
               IF (.NOT. alpha_beta) THEN
                  j_global = cp_fm_indxl2g(jjB, ncol_block, send_pcol, &
                                           mp2_env%ri_grad%P_mo%matrix_struct%first_p_pos(2), npcol)
               ELSE
                  j_global = cp_fm_indxl2g(jjB, ncol_block, send_pcol, &
                                           mp2_env%ri_grad%P_mo_beta%matrix_struct%first_p_pos(2), npcol)
               ENDIF
               IF (j_global <= homo) CYCLE
               j_global = j_global-homo
               IF (.NOT. alpha_beta) THEN
                  ab_send(iiB, jjB) = mp2_env%ri_grad%P_ab(i_global-my_B_virtual_start+1, j_global)
               ELSE
                  ab_send(iiB, jjB) = mp2_env%ri_grad%P_ab_beta(i_global-my_B_virtual_start+1, j_global)
               ENDIF
            END DO
         END DO

         ab_rec = 0.0_dp
         CALL mp_sendrecv(ab_send, proc_send, &
                          ab_rec, proc_receive, &
                          para_env_sub%group)
         IF (.NOT. alpha_beta) THEN
            mp2_env%ri_grad%P_mo%local_data(1:nrow_local, 1:ncol_local) = &
               mp2_env%ri_grad%P_mo%local_data(1:nrow_local, 1:ncol_local)+ &
               ab_rec(1:nrow_local, 1:ncol_local)
         ELSE
            mp2_env%ri_grad%P_mo_beta%local_data(1:nrow_local, 1:ncol_local) = &
               mp2_env%ri_grad%P_mo_beta%local_data(1:nrow_local, 1:ncol_local)+ &
               ab_rec(1:nrow_local, 1:ncol_local)
         ENDIF

         DEALLOCATE (ab_send)
      END DO
      DEALLOCATE (ab_rec)
      DEALLOCATE (mepos_2_grid)
      DEALLOCATE (sizes)

      ! deallocate the local P_ab'
      IF (.NOT. alpha_beta) THEN
         DEALLOCATE (mp2_env%ri_grad%P_ab)
      ELSE
         DEALLOCATE (mp2_env%ri_grad%P_ab_beta)
      ENDIF
      CALL timestop(handle2)

      ! create also W_MP2_MO
      CALL timeset(routineN//"_WMO", handle2)
      IF (.NOT. alpha_beta) THEN
         NULLIFY (mp2_env%ri_grad%W_mo)
         CALL cp_fm_create(mp2_env%ri_grad%W_mo, fm_struct_tmp, name="W_MP2_MO")
         CALL cp_fm_set_all(mp2_env%ri_grad%W_mo, 0.0_dp)
         CALL cp_fm_struct_release(fm_struct_tmp)

         ! all block
         CALL cp_gemm('T', 'N', dimen, dimen, dimen, 2.0_dp*factor, &
                      L_mu_q, mo_coeff, 0.0_dp, mp2_env%ri_grad%W_mo, &
                      a_first_col=1, &
                      a_first_row=1, &
                      b_first_col=1, &
                      b_first_row=1, &
                      c_first_col=1, &
                      c_first_row=1)

         ! occ-occ block
         CALL cp_gemm('T', 'N', homo, homo, dimen, -2.0_dp*factor, &
                      L_mu_q, mo_coeff, 0.0_dp, mp2_env%ri_grad%W_mo, &
                      a_first_col=1, &
                      a_first_row=1, &
                      b_first_col=1, &
                      b_first_row=1, &
                      c_first_col=1, &
                      c_first_row=1)

         ! occ-virt block
         CALL cp_gemm('T', 'N', homo, virtual, dimen, 2.0_dp*factor, &
                      mo_coeff, L_mu_q, 0.0_dp, mp2_env%ri_grad%W_mo, &
                      a_first_col=1, &
                      a_first_row=1, &
                      b_first_col=homo+1, &
                      b_first_row=1, &
                      c_first_col=homo+1, &
                      c_first_row=1)
      ELSE
         NULLIFY (mp2_env%ri_grad%W_mo_beta)
         CALL cp_fm_create(mp2_env%ri_grad%W_mo_beta, fm_struct_tmp, name="W_MP2_MO")
         CALL cp_fm_set_all(mp2_env%ri_grad%W_mo_beta, 0.0_dp)
         CALL cp_fm_struct_release(fm_struct_tmp)

         ! all block
         CALL cp_gemm('T', 'N', dimen, dimen, dimen, 2.0_dp*factor, &
                      L_mu_q, mo_coeff, 0.0_dp, mp2_env%ri_grad%W_mo_beta, &
                      a_first_col=1, &
                      a_first_row=1, &
                      b_first_col=1, &
                      b_first_row=1, &
                      c_first_col=1, &
                      c_first_row=1)

         ! occ-occ block
         CALL cp_gemm('T', 'N', homo, homo, dimen, -2.0_dp*factor, &
                      L_mu_q, mo_coeff, 0.0_dp, mp2_env%ri_grad%W_mo_beta, &
                      a_first_col=1, &
                      a_first_row=1, &
                      b_first_col=1, &
                      b_first_row=1, &
                      c_first_col=1, &
                      c_first_row=1)

         ! occ-virt block
         CALL cp_gemm('T', 'N', homo, virtual, dimen, 2.0_dp*factor, &
                      mo_coeff, L_mu_q, 0.0_dp, mp2_env%ri_grad%W_mo_beta, &
                      a_first_col=1, &
                      a_first_row=1, &
                      b_first_col=homo+1, &
                      b_first_row=1, &
                      c_first_col=homo+1, &
                      c_first_row=1)
      ENDIF
      CALL timestop(handle2)

      ! IF(para_env%mepos==0) CALL write_array(mp2_env%ri_grad%P_ij(1:homo,1:homo))

      ! Calculate occ-virt block of the lagrangian in MO
      CALL timeset(routineN//"_Ljb", handle2)
      IF (.NOT. alpha_beta) THEN
         NULLIFY (mp2_env%ri_grad%L_jb, fm_struct_tmp)
         CALL cp_fm_struct_create(fm_struct_tmp, para_env=para_env, context=blacs_env, &
                                  nrow_global=homo, ncol_global=virtual)
         CALL cp_fm_create(mp2_env%ri_grad%L_jb, fm_struct_tmp, name="fm_L_jb")
         CALL cp_fm_struct_release(fm_struct_tmp)
         CALL cp_fm_set_all(mp2_env%ri_grad%L_jb, 0.0_dp)

         ! first Virtual
         CALL cp_gemm('T', 'N', homo, virtual, dimen, 2.0_dp*factor, &
                      L_mu_q, mo_coeff, 0.0_dp, mp2_env%ri_grad%L_jb, &
                      a_first_col=1, &
                      a_first_row=1, &
                      b_first_col=homo+1, &
                      b_first_row=1, &
                      c_first_col=1, &
                      c_first_row=1)
         ! then occupied
         CALL cp_gemm('T', 'N', homo, virtual, dimen, 2.0_dp*factor, &
                      mo_coeff, L_mu_q, 1.0_dp, mp2_env%ri_grad%L_jb, &
                      a_first_col=1, &
                      a_first_row=1, &
                      b_first_col=homo+1, &
                      b_first_row=1, &
                      c_first_col=1, &
                      c_first_row=1)
      ELSE
         NULLIFY (mp2_env%ri_grad%L_jb_beta, fm_struct_tmp)
         CALL cp_fm_struct_create(fm_struct_tmp, para_env=para_env, context=blacs_env, &
                                  nrow_global=homo, ncol_global=virtual)
         CALL cp_fm_create(mp2_env%ri_grad%L_jb_beta, fm_struct_tmp, name="fm_L_jb")
         CALL cp_fm_struct_release(fm_struct_tmp)
         CALL cp_fm_set_all(mp2_env%ri_grad%L_jb_beta, 0.0_dp)

         ! first Virtual
         CALL cp_gemm('T', 'N', homo, virtual, dimen, 2.0_dp*factor, &
                      L_mu_q, mo_coeff, 0.0_dp, mp2_env%ri_grad%L_jb_beta, &
                      a_first_col=1, &
                      a_first_row=1, &
                      b_first_col=homo+1, &
                      b_first_row=1, &
                      c_first_col=1, &
                      c_first_row=1)
         ! then occupied
         CALL cp_gemm('T', 'N', homo, virtual, dimen, 2.0_dp*factor, &
                      mo_coeff, L_mu_q, 1.0_dp, mp2_env%ri_grad%L_jb_beta, &
                      a_first_col=1, &
                      a_first_row=1, &
                      b_first_col=homo+1, &
                      b_first_row=1, &
                      c_first_col=1, &
                      c_first_row=1)
      ENDIF
      ! finally release L_mu_q
      CALL cp_fm_release(L_mu_q)
      CALL timestop(handle2)

      ! here we should be done next CPHF

      IF (.FALSE.) THEN
         ALLOCATE (ab_send(homo, virtual))
         ab_send = 0.0_dp
         IF (.NOT. alpha_beta) THEN
            CALL cp_fm_get_info(matrix=mp2_env%ri_grad%L_jb, &
                                nrow_local=nrow_local, &
                                ncol_local=ncol_local, &
                                row_indices=row_indices, &
                                col_indices=col_indices)
         ELSE
            CALL cp_fm_get_info(matrix=mp2_env%ri_grad%L_jb_beta, &
                                nrow_local=nrow_local, &
                                ncol_local=ncol_local, &
                                row_indices=row_indices, &
                                col_indices=col_indices)
         ENDIF
         DO jjB = 1, ncol_local
            j_global = col_indices(jjB)
            DO iiB = 1, nrow_local
               i_global = row_indices(iiB)
               IF (.NOT. alpha_beta) THEN
                  ab_send(i_global, j_global) = mp2_env%ri_grad%L_jb%local_data(iiB, jjB)
               ELSE
                  ab_send(i_global, j_global) = mp2_env%ri_grad%L_jb_beta%local_data(iiB, jjB)
               ENDIF
            END DO
         END DO
         CALL mp_sum(ab_send, para_env%group)
         IF (para_env%mepos == 0) THEN
            WRITE (*, *) 'Ljb'
            CALL write_array(ab_send(1:homo, 1:virtual))
         END IF
         DEALLOCATE (ab_send)
         ALLOCATE (ab_send(homo, homo))
         ALLOCATE (ab_rec(virtual, virtual))
         ALLOCATE (mat_deb(homo, virtual))
         ab_send = 0.0_dp
         ab_rec = 0.0_dp
         mat_deb = 0.0_dp
         IF (.NOT. alpha_beta) THEN
            CALL cp_fm_get_info(matrix=mp2_env%ri_grad%P_mo, &
                                nrow_local=nrow_local, &
                                ncol_local=ncol_local, &
                                row_indices=row_indices, &
                                col_indices=col_indices)
         ELSE
            CALL cp_fm_get_info(matrix=mp2_env%ri_grad%P_mo_beta, &
                                nrow_local=nrow_local, &
                                ncol_local=ncol_local, &
                                row_indices=row_indices, &
                                col_indices=col_indices)
         ENDIF
         IF (.NOT. alpha_beta) THEN
            DO jjB = 1, ncol_local
               j_global = col_indices(jjB)
               DO iiB = 1, nrow_local
                  i_global = row_indices(iiB)
                  IF (i_global <= homo .AND. j_global <= homo) THEN
                     ab_send(i_global, j_global) = mp2_env%ri_grad%P_mo%local_data(iiB, jjB)
                  ELSE
                     IF (i_global > homo .AND. j_global > homo) THEN
                        ab_rec(i_global-homo, j_global-homo) = mp2_env%ri_grad%P_mo%local_data(iiB, jjB)
                     ELSE
                        mat_deb(MIN(i_global, j_global), MAX(i_global, j_global)-homo) = &
                           mat_deb(MIN(i_global, j_global), MAX(i_global, j_global)-homo)+ &
                           mp2_env%ri_grad%P_mo%local_data(iiB, jjB)
                     END IF
                  END IF
               END DO
            END DO
         ELSE
            DO jjB = 1, ncol_local
               j_global = col_indices(jjB)
               DO iiB = 1, nrow_local
                  i_global = row_indices(iiB)
                  IF (i_global <= homo .AND. j_global <= homo) THEN
                     ab_send(i_global, j_global) = mp2_env%ri_grad%P_mo_beta%local_data(iiB, jjB)
                  ELSE
                     IF (i_global > homo .AND. j_global > homo) THEN
                        ab_rec(i_global-homo, j_global-homo) = mp2_env%ri_grad%P_mo_beta%local_data(iiB, jjB)
                     ELSE
                        mat_deb(MIN(i_global, j_global), MAX(i_global, j_global)-homo) = &
                           mat_deb(MIN(i_global, j_global), MAX(i_global, j_global)-homo)+ &
                           mp2_env%ri_grad%P_mo_beta%local_data(iiB, jjB)
                     END IF
                  END IF
               END DO
            END DO
         ENDIF

         CALL mp_sum(ab_send, para_env%group)
         CALL mp_sum(ab_rec, para_env%group)
         CALL mp_sum(mat_deb, para_env%group)
         IF (para_env%mepos == 0) THEN
            WRITE (*, *) 'Pij'
            CALL write_array(ab_send(1:homo, 1:homo))
            WRITE (*, *) 'Pab'
            CALL write_array(ab_rec(1:virtual, 1:virtual))
            WRITE (*, *) 'Pia'
            CALL write_array(mat_deb(1:homo, 1:virtual))
         END IF
         DEALLOCATE (ab_send, ab_rec, mat_deb)
         ALLOCATE (ab_send(dimen, dimen))
         ab_send = 0.0_dp
         IF (.NOT. alpha_beta) THEN
            CALL cp_fm_get_info(matrix=mp2_env%ri_grad%W_mo, &
                                nrow_local=nrow_local, &
                                ncol_local=ncol_local, &
                                row_indices=row_indices, &
                                col_indices=col_indices)
         ELSE
            CALL cp_fm_get_info(matrix=mp2_env%ri_grad%W_mo_beta, &
                                nrow_local=nrow_local, &
                                ncol_local=ncol_local, &
                                row_indices=row_indices, &
                                col_indices=col_indices)
         ENDIF

         DO jjB = 1, ncol_local
            j_global = col_indices(jjB)
            DO iiB = 1, nrow_local
               i_global = row_indices(iiB)
               IF (.NOT. alpha_beta) THEN
                  ab_send(i_global, j_global) = &
                     mp2_env%ri_grad%W_mo%local_data(iiB, jjB)
               ELSE
                  ab_send(i_global, j_global) = &
                     mp2_env%ri_grad%W_mo_beta%local_data(iiB, jjB)
               ENDIF
            END DO
         END DO
         CALL mp_sum(ab_send, para_env%group)
         IF (para_env%mepos == 0) THEN
            WRITE (*, *) 'Wij'
            CALL write_array(ab_send(1:dimen, 1:dimen))
         END IF
         DEALLOCATE (ab_send)
      END IF

      DEALLOCATE (sub_proc_map)
      DEALLOCATE (pos_info)

      CALL timestop(handle)

   END SUBROUTINE create_W_P

! **************************************************************************************************
!> \brief ...
!> \param mat ...
!> \param unitout ...
! **************************************************************************************************
   SUBROUTINE write_array(mat, unitout)
      REAL(KIND=dp), DIMENSION(:, :)                     :: mat
      INTEGER, OPTIONAL                                  :: unitout

      INTEGER                                            :: iii, jjj

      WRITE (*, *)
      DO iii = 1, SIZE(mat, 1)
         WRITE (*, *) iii
         DO jjj = 1, SIZE(mat, 2), 20
            IF (PRESENT(unitout)) THEN
               WRITE (1000+unitout, '(1000F10.5)') mat(iii, jjj:MIN(SIZE(mat, 2), jjj+19))
            ELSE
               WRITE (*, '(1000F10.5)') mat(iii, jjj:MIN(SIZE(mat, 2), jjj+19))
            END IF
         END DO
         WRITE (*, *)
      END DO
      WRITE (*, *)
   END SUBROUTINE

END MODULE mp2_ri_grad
