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

MODULE optimize_embedding_potential

   USE atomic_kind_types,               ONLY: atomic_kind_type,&
                                              get_atomic_kind,&
                                              get_atomic_kind_set
   USE cell_types,                      ONLY: cell_type
   USE cp_blacs_env,                    ONLY: cp_blacs_env_create,&
                                              cp_blacs_env_release,&
                                              cp_blacs_env_type
   USE cp_control_types,                ONLY: dft_control_type
   USE cp_dbcsr_operations,             ONLY: copy_dbcsr_to_fm
   USE cp_files,                        ONLY: close_file,&
                                              open_file
   USE cp_fm_basic_linalg,              ONLY: cp_fm_column_scale,&
                                              cp_fm_scale,&
                                              cp_fm_scale_and_add,&
                                              cp_fm_trace
   USE cp_fm_diag,                      ONLY: choose_eigv_solver
   USE cp_fm_struct,                    ONLY: cp_fm_struct_create,&
                                              cp_fm_struct_release,&
                                              cp_fm_struct_type
   USE cp_fm_types,                     ONLY: &
        cp_fm_copy_general, cp_fm_create, cp_fm_get_element, cp_fm_get_info, cp_fm_release, &
        cp_fm_set_all, cp_fm_to_fm, cp_fm_to_fm_submat, cp_fm_type, cp_fm_write_unformatted
   USE cp_gemm_interface,               ONLY: cp_gemm
   USE cp_log_handling,                 ONLY: cp_get_default_logger,&
                                              cp_logger_type
   USE cp_output_handling,              ONLY: cp_p_file,&
                                              cp_print_key_finished_output,&
                                              cp_print_key_should_output,&
                                              cp_print_key_unit_nr
   USE cp_para_types,                   ONLY: cp_para_env_type
   USE cp_realspace_grid_cube,          ONLY: cp_pw_to_cube
   USE dbcsr_api,                       ONLY: dbcsr_deallocate_matrix_set,&
                                              dbcsr_p_type
   USE embed_types,                     ONLY: opt_embed_pot_type
   USE force_env_types,                 ONLY: force_env_type
   USE input_constants,                 ONLY: embed_level_shift,&
                                              embed_quasi_newton,&
                                              embed_steep_desc
   USE input_section_types,             ONLY: section_get_ivals,&
                                              section_vals_get_subs_vals,&
                                              section_vals_type,&
                                              section_vals_val_get
   USE kinds,                           ONLY: default_path_length,&
                                              dp
   USE lri_environment_types,           ONLY: lri_kind_type
   USE message_passing,                 ONLY: mp_bcast,&
                                              mp_max,&
                                              mp_sum
   USE particle_list_types,             ONLY: particle_list_type
   USE particle_types,                  ONLY: particle_type
   USE pw_env_types,                    ONLY: pw_env_get,&
                                              pw_env_type
   USE pw_methods,                      ONLY: pw_axpy,&
                                              pw_copy,&
                                              pw_integrate_function,&
                                              pw_scale,&
                                              pw_zero
   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,&
                                              pw_release
   USE qs_collocate_density,            ONLY: calculate_ppl_grid,&
                                              calculate_wavefunction
   USE qs_environment_types,            ONLY: get_qs_env,&
                                              qs_environment_type,&
                                              set_qs_env
   USE qs_integrate_potential_single,   ONLY: integrate_v_rspace_one_center
   USE qs_kind_types,                   ONLY: get_qs_kind,&
                                              qs_kind_type
   USE qs_kinetic,                      ONLY: build_kinetic_matrix
   USE qs_ks_types,                     ONLY: qs_ks_env_type
   USE qs_mo_types,                     ONLY: get_mo_set,&
                                              mo_set_p_type
   USE qs_neighbor_list_types,          ONLY: neighbor_list_set_p_type
   USE qs_subsys_types,                 ONLY: qs_subsys_get,&
                                              qs_subsys_type
   USE realspace_grid_cube,             ONLY: cube_to_pw
#include "./base/base_uses.f90"

   IMPLICIT NONE

   PRIVATE

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

   PUBLIC :: prepare_embed_opt, init_embed_pot, release_opt_embed, calculate_embed_pot_grad, &
             opt_embed_step, print_rho_diff, step_control, max_dens_diff, print_emb_opt_info, &
             conv_check_embed, make_subsys_embed_pot, print_embed_restart, find_aux_dimen, &
             read_embed_pot, understand_spin_states, given_embed_pot

CONTAINS

! **************************************************************************************************
!> \brief Find out whether we need to swap alpha- and beta- spind densities in the second subsystem
!> \brief It's only needed because by default alpha-spins go first in a subsystem.
!> \brief By swapping we impose the constraint:
!> \brief rho_1(alpha) + rho_2(alpha) = rho_total(alpha)
!> \brief rho_1(beta) + rho_2(beta) = rho_total(beta)
!> \param force_env ...
!> \param ref_subsys_number ...
!> \param change_spin ...
!> \param open_shell_embed ...
!> \return ...
!> \author Vladimir Rybkin
! **************************************************************************************************
   SUBROUTINE understand_spin_states(force_env, ref_subsys_number, change_spin, open_shell_embed)
      TYPE(force_env_type), POINTER                      :: force_env
      INTEGER                                            :: ref_subsys_number
      LOGICAL                                            :: change_spin, open_shell_embed

      INTEGER                                            :: i_force_eval, nspins, sub_spin_1, &
                                                            sub_spin_2, total_spin
      INTEGER, DIMENSION(2)                              :: nelectron_spin
      INTEGER, DIMENSION(2, 3)                           :: all_spins
      INTEGER, DIMENSION(3)                              :: all_nspins
      TYPE(dft_control_type), POINTER                    :: dft_control

      change_spin = .FALSE.
      open_shell_embed = .FALSE.
      IF (ref_subsys_number .EQ. 3) THEN
         all_spins = 0
         DO i_force_eval = 1, ref_subsys_number
            CALL get_qs_env(qs_env=force_env%sub_force_env(i_force_eval)%force_env%qs_env, &
                            nelectron_spin=nelectron_spin, dft_control=dft_control)
            all_spins(:, i_force_eval) = nelectron_spin
            nspins = dft_control%nspins
            all_nspins(i_force_eval) = nspins
         ENDDO

         ! Find out whether we need a spin-dependend embedding potential
         IF (.NOT. ((all_nspins(1) .EQ. 1) .AND. (all_nspins(2) .EQ. 1) .AND. (all_nspins(3) .EQ. 1))) &
            open_shell_embed = .TRUE.

         ! If it's open shell, we need to check spin states
         IF (open_shell_embed) THEN

            IF (all_nspins(3) .EQ. 1) THEN
               total_spin = 0
            ELSE
               total_spin = all_spins(1, 3)-all_spins(2, 3)
            ENDIF
            IF (all_nspins(1) .EQ. 1) THEN
               sub_spin_1 = 0
            ELSE
               sub_spin_1 = all_spins(1, 1)-all_spins(2, 1)
            ENDIF
            IF (all_nspins(2) .EQ. 1) THEN
               sub_spin_2 = 0
            ELSE
               sub_spin_2 = all_spins(1, 2)-all_spins(2, 2)
            ENDIF
            IF ((sub_spin_1+sub_spin_2) .EQ. total_spin) THEN
               change_spin = .FALSE.
            ELSE
               IF (ABS(sub_spin_1-sub_spin_2) .EQ. total_spin) THEN
                  change_spin = .TRUE.
               ELSE
                  CPABORT("Spin states of subsystems are not compatible.")
               ENDIF
            ENDIF

         ENDIF ! not open_shell
      ELSE
         CPABORT("Reference subsystem must be the third FORCE_EVAL.")
      ENDIF

   END SUBROUTINE understand_spin_states

! **************************************************************************************************
!> \brief ...
!> \param qs_env ...
!> \param embed_pot ...
!> \param add_const_pot ...
!> \param const_pot ...
!> \param open_shell_embed ...
!> \param spin_embed_pot ...
! **************************************************************************************************
   SUBROUTINE init_embed_pot(qs_env, embed_pot, add_const_pot, const_pot, open_shell_embed, spin_embed_pot)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(pw_p_type), POINTER                           :: embed_pot
      LOGICAL                                            :: add_const_pot
      TYPE(pw_p_type), OPTIONAL, POINTER                 :: const_pot
      LOGICAL                                            :: open_shell_embed
      TYPE(pw_p_type), POINTER                           :: spin_embed_pot

      INTEGER                                            :: nelectrons
      INTEGER, DIMENSION(2)                              :: nelectron_spin
      REAL(KIND=dp)                                      :: factor
      TYPE(pw_env_type), POINTER                         :: pw_env
      TYPE(pw_p_type)                                    :: v_hartree_r_space, vppl
      TYPE(pw_pool_type), POINTER                        :: auxbas_pw_pool

      ! Extract  plane waves environment
      CALL get_qs_env(qs_env=qs_env, pw_env=pw_env, &
                      nelectron_spin=nelectron_spin, &
                      v_hartree_rspace=v_hartree_r_space%pw)

      ! Prepare plane-waves pool
      CALL pw_env_get(pw_env, auxbas_pw_pool=auxbas_pw_pool)

      ! Create embedding potential and set to zero
      NULLIFY (embed_pot)
      ALLOCATE (embed_pot)
      NULLIFY (embed_pot%pw)
      CALL pw_pool_create_pw(auxbas_pw_pool, embed_pot%pw, &
                             use_data=REALDATA3D, &
                             in_space=REALSPACE)
      CALL pw_zero(embed_pot%pw)

      ! Spin embedding potential if asked
      IF (open_shell_embed) THEN
         NULLIFY (spin_embed_pot)
         ALLOCATE (spin_embed_pot)
         NULLIFY (spin_embed_pot%pw)
         CALL pw_pool_create_pw(auxbas_pw_pool, spin_embed_pot%pw, &
                                use_data=REALDATA3D, &
                                in_space=REALSPACE)
         CALL pw_zero(spin_embed_pot%pw)
      ENDIF

      ! Add constant potential if requested
      IF (add_const_pot) THEN
         ! First, pseudopotentials without error function
         NULLIFY (vppl%pw)
         CALL pw_pool_create_pw(auxbas_pw_pool, vppl%pw, &
                                use_data=REALDATA3D, &
                                in_space=REALSPACE)

         CALL calculate_ppl_grid(vppl, qs_env)

         ! Add Fermi-Amaldi potential on the reference density
         NULLIFY (const_pot)
         ALLOCATE (const_pot)
         NULLIFY (const_pot%pw, v_hartree_r_space%pw)

         CALL pw_pool_create_pw(auxbas_pw_pool, const_pot%pw, &
                                use_data=REALDATA3D, &
                                in_space=REALSPACE)

         ! Extract  Hartree potential
         CALL get_qs_env(qs_env=qs_env, pw_env=pw_env, &
                         v_hartree_rspace=v_hartree_r_space%pw)
         CALL pw_copy(v_hartree_r_space%pw, const_pot%pw)

         ! Calculate the number of electrons
         nelectrons = nelectron_spin(1)+nelectron_spin(2)
         factor = (REAL(nelectrons)-1.0_dp)/(REAL(nelectrons))

         ! Scale the Hartree potential to get Fermi-Amaldi
         CALL pw_scale(const_pot%pw, a=factor)

         ! Sum pseudo and Fermi-Amaldi
         CALL pw_axpy(vppl%pw, const_pot%pw)
         CALL pw_copy(const_pot%pw, embed_pot%pw)

         ! Give back plane waves pool for pseudo
         CALL pw_pool_give_back_pw(auxbas_pw_pool, vppl%pw)
      ENDIF

   END SUBROUTINE init_embed_pot

! **************************************************************************************************
!> \brief Creates and allocates objects for optimization of embedding potential
!> \param qs_env ...
!> \param opt_embed ...
!> \param opt_embed_section ...
!> \author Vladimir Rybkin
! **************************************************************************************************
   SUBROUTINE prepare_embed_opt(qs_env, opt_embed, opt_embed_section)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(opt_embed_pot_type)                           :: opt_embed
      TYPE(section_vals_type), POINTER                   :: opt_embed_section

      INTEGER                                            :: diff_size
      TYPE(cp_blacs_env_type), POINTER                   :: blacs_env
      TYPE(cp_fm_struct_type), POINTER                   :: fm_struct
      TYPE(cp_para_env_type), POINTER                    :: para_env

      !TYPE(pw_env_type), POINTER                         :: pw_env
      !TYPE(pw_pool_type), POINTER                        :: auxbas_pw_pool

      ! First, read the input

      CALL read_opt_embed_section(opt_embed, opt_embed_section)

      ! All these are needed for optimization in a finite Guassian basis
      IF (.NOT. opt_embed%grid_opt) THEN
         ! Create blacs environment
         CALL get_qs_env(qs_env=qs_env, &
                         para_env=para_env)
         NULLIFY (blacs_env)
         CALL cp_blacs_env_create(blacs_env=blacs_env, para_env=para_env)

         ! Reveal the dimention of the RI basis
         CALL find_aux_dimen(qs_env, opt_embed%dimen_aux)

         ! Prepare the object for integrals
         CALL make_lri_object(qs_env, opt_embed%lri)

         ! In case if spin embedding potential has to be optimized,
         ! the dimension of variational space is two times larger
         IF (opt_embed%open_shell_embed) THEN
            opt_embed%dimen_var_aux = 2*opt_embed%dimen_aux
         ELSE
            opt_embed%dimen_var_aux = opt_embed%dimen_aux
         ENDIF

         ! Allocate expansion coefficients and gradient
         NULLIFY (opt_embed%embed_pot_grad, opt_embed%embed_pot_coef, opt_embed%step, fm_struct)

         NULLIFY (opt_embed%prev_embed_pot_grad, opt_embed%prev_embed_pot_coef, opt_embed%prev_step)
         CALL cp_fm_struct_create(fm_struct, para_env=para_env, context=blacs_env, &
                                  nrow_global=opt_embed%dimen_var_aux, ncol_global=1)
         CALL cp_fm_create(opt_embed%embed_pot_grad, fm_struct, name="pot_grad")
         CALL cp_fm_create(opt_embed%embed_pot_coef, fm_struct, name="pot_coef")
         CALL cp_fm_create(opt_embed%prev_embed_pot_grad, fm_struct, name="prev_pot_grad")
         CALL cp_fm_create(opt_embed%prev_embed_pot_coef, fm_struct, name="prev_pot_coef")
         CALL cp_fm_create(opt_embed%step, fm_struct, name="step")

         CALL cp_fm_create(opt_embed%prev_step, fm_struct, name="prev_step")

         CALL cp_fm_struct_release(fm_struct)
         CALL cp_fm_set_all(opt_embed%embed_pot_grad, 0.0_dp)
         CALL cp_fm_set_all(opt_embed%prev_embed_pot_grad, 0.0_dp)
         CALL cp_fm_set_all(opt_embed%embed_pot_coef, 0.0_dp)
         CALL cp_fm_set_all(opt_embed%prev_embed_pot_coef, 0.0_dp)
         CALL cp_fm_set_all(opt_embed%step, 0.0_dp)

         CALL cp_fm_set_all(opt_embed%prev_step, 0.0_dp)

         ! Allocate Hessian
         NULLIFY (opt_embed%embed_pot_hess, opt_embed%prev_embed_pot_hess, fm_struct)
         CALL cp_fm_struct_create(fm_struct, para_env=para_env, context=blacs_env, &
                                  nrow_global=opt_embed%dimen_var_aux, ncol_global=opt_embed%dimen_var_aux)
         CALL cp_fm_create(opt_embed%embed_pot_hess, fm_struct, name="pot_Hess")
         CALL cp_fm_create(opt_embed%prev_embed_pot_hess, fm_struct, name="prev_pot_Hess")
         CALL cp_fm_struct_release(fm_struct)

         ! Special structure for the kinetic energy matrix
         NULLIFY (fm_struct, opt_embed%kinetic_mat)
         CALL cp_fm_struct_create(fm_struct, para_env=para_env, context=blacs_env, &
                                  nrow_global=opt_embed%dimen_aux, ncol_global=opt_embed%dimen_aux)
         CALL cp_fm_create(opt_embed%kinetic_mat, fm_struct, name="kinetic_mat")
         CALL cp_fm_struct_release(fm_struct)
         CALL cp_fm_set_all(opt_embed%kinetic_mat, 0.0_dp)

         ! Hessian is set as a unit matrix
         CALL cp_fm_set_all(opt_embed%embed_pot_hess, 0.0_dp, -1.0_dp)
         CALL cp_fm_set_all(opt_embed%prev_embed_pot_hess, 0.0_dp, -1.0_dp)

         ! Release blacs environment
         CALL cp_blacs_env_release(blacs_env)

         !ELSE ! Grid-based optimization
         !my_spins = 1
         !IF (opt_embed%open_shell_embed) my_spins = 2
         !CALL get_qs_env(qs_env=qs_env, pw_env=pw_env)
         !CALL pw_env_get(pw_env, auxbas_pw_pool=auxbas_pw_pool)
         !NULLIFY (opt_embed%prev_grid_grad, opt_embed%prev_grid_pot)
         !ALLOCATE (opt_embed%prev_grid_grad(my_spins))
         !ALLOCATE (opt_embed%prev_grid_pot(my_spins))
         !DO i_spin = 1, my_spins
         !   ALLOCATE (opt_embed%prev_grid_grad(i_spin))
         !   CALL pw_pool_create_pw(auxbas_pw_pool, opt_embed%prev_grid_grad(i_spin)%pw, &
         !                          use_data=REALDATA3D, &
         !                          in_space=REALSPACE)
         !   CALL pw_zero(opt_embed%prev_grid_grad(i_spin)%pw)
         !   ALLOCATE (opt_embed%prev_grid_pot(i_spin))
         !   CALL pw_pool_create_pw(auxbas_pw_pool, opt_embed%prev_grid_pot(i_spin)%pw, &
         !                          use_data=REALDATA3D, &
         !                          in_space=REALSPACE)
         !   CALL pw_zero(opt_embed%prev_grid_pot(i_spin)%pw)
         !ENDDO

      ENDIF

      ! Array to store functional values
      ALLOCATE (opt_embed%w_func(opt_embed%n_iter))
      opt_embed%w_func = 0.0_dp

      ! Allocate max_diff and int_diff
      diff_size = 1
      IF (opt_embed%open_shell_embed) diff_size = 2
      ALLOCATE (opt_embed%max_diff(diff_size))
      ALLOCATE (opt_embed%int_diff(diff_size))

      ! Set allowed energy decrease parameter
      opt_embed%allowed_decrease = 0.0001_dp

      ! Regularization contribution is set to zero
      opt_embed%reg_term = 0.0_dp

      ! Step is accepted in the beginning
      opt_embed%accept_step = .TRUE.
      opt_embed%newton_step = .FALSE.
      opt_embed%last_accepted = 1

      ! Set maximum and minimum trust radii
      opt_embed%max_trad = opt_embed%trust_rad*7.900_dp
      opt_embed%min_trad = opt_embed%trust_rad*0.125*0.065_dp

   END SUBROUTINE prepare_embed_opt

! **************************************************************************************************
!> \brief ...
!> \param opt_embed ...
!> \param opt_embed_section ...
! **************************************************************************************************
   SUBROUTINE read_opt_embed_section(opt_embed, opt_embed_section)
      TYPE(opt_embed_pot_type)                           :: opt_embed
      TYPE(section_vals_type), POINTER                   :: opt_embed_section

      INTEGER                                            :: embed_optimizer

      ! Read keywords
      CALL section_vals_val_get(opt_embed_section, "REG_LAMBDA", &
                                r_val=opt_embed%lambda)

      CALL section_vals_val_get(opt_embed_section, "N_ITER", &
                                i_val=opt_embed%n_iter)

      CALL section_vals_val_get(opt_embed_section, "TRUST_RAD", &
                                r_val=opt_embed%trust_rad)

      CALL section_vals_val_get(opt_embed_section, "DENS_CONV_MAX", &
                                r_val=opt_embed%conv_max)

      CALL section_vals_val_get(opt_embed_section, "DENS_CONV_INT", &
                                r_val=opt_embed%conv_int)

      CALL section_vals_val_get(opt_embed_section, "SPIN_DENS_CONV_MAX", &
                                r_val=opt_embed%conv_max_spin)

      CALL section_vals_val_get(opt_embed_section, "SPIN_DENS_CONV_INT", &
                                r_val=opt_embed%conv_int_spin)

      CALL section_vals_val_get(opt_embed_section, "ADD_CONST_POT", &
                                l_val=opt_embed%add_const_pot)

      CALL section_vals_val_get(opt_embed_section, "READ_EMBED_POT", &
                                l_val=opt_embed%read_embed_pot)

      CALL section_vals_val_get(opt_embed_section, "READ_EMBED_POT_CUBE", &
                                l_val=opt_embed%read_embed_pot_cube)

      CALL section_vals_val_get(opt_embed_section, "GRID_OPT", &
                                l_val=opt_embed%grid_opt)

      CALL section_vals_val_get(opt_embed_section, "OPTIMIZER", i_val=embed_optimizer)
      SELECT CASE (embed_optimizer)
      CASE (embed_steep_desc)
         opt_embed%steep_desc = .TRUE.
      CASE (embed_quasi_newton)
         opt_embed%steep_desc = .FALSE.
         opt_embed%level_shift = .FALSE.
      CASE (embed_level_shift)
         opt_embed%steep_desc = .FALSE.
         opt_embed%level_shift = .TRUE.
      CASE DEFAULT
         opt_embed%steep_desc = .TRUE.
      END SELECT
   END SUBROUTINE read_opt_embed_section

! **************************************************************************************************
!> \brief Find the dimension of the auxiliary basis for the expansion of the embedding potential
!> \param qs_env ...
!> \param dimen_aux ...
! **************************************************************************************************
   SUBROUTINE find_aux_dimen(qs_env, dimen_aux)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      INTEGER                                            :: dimen_aux

      INTEGER                                            :: iatom, ikind, natom, nsgf
      INTEGER, ALLOCATABLE, DIMENSION(:)                 :: kind_of
      TYPE(atomic_kind_type), DIMENSION(:), POINTER      :: atomic_kind_set
      TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
      TYPE(qs_kind_type), DIMENSION(:), POINTER          :: qs_kind_set

      ! First, reveal the dimention of the RI basis
      CALL get_qs_env(qs_env=qs_env, &
                      particle_set=particle_set, &
                      qs_kind_set=qs_kind_set, &
                      atomic_kind_set=atomic_kind_set)

      natom = SIZE(particle_set)
      ALLOCATE (kind_of(natom))

      CALL get_atomic_kind_set(atomic_kind_set, kind_of=kind_of)

      dimen_aux = 0
      DO iatom = 1, natom
         ikind = kind_of(iatom)
         CALL get_qs_kind(qs_kind=qs_kind_set(ikind), nsgf=nsgf, basis_type="RI_AUX")
         dimen_aux = dimen_aux+nsgf
      END DO

      DEALLOCATE (kind_of)

   END SUBROUTINE find_aux_dimen

! **************************************************************************************************
!> \brief Prepare the lri_kind_type object for integrals between density and aux. basis functions
!> \param qs_env ...
!> \param lri ...
! **************************************************************************************************
   SUBROUTINE make_lri_object(qs_env, lri)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(lri_kind_type), DIMENSION(:), POINTER         :: lri

      INTEGER                                            :: ikind, natom, nkind, nsgf
      TYPE(atomic_kind_type), DIMENSION(:), POINTER      :: atomic_kind_set
      TYPE(atomic_kind_type), POINTER                    :: atomic_kind
      TYPE(qs_kind_type), DIMENSION(:), POINTER          :: qs_kind_set

      NULLIFY (atomic_kind, lri)
      CALL get_qs_env(qs_env=qs_env, atomic_kind_set=atomic_kind_set, &
                      qs_kind_set=qs_kind_set)
      nkind = SIZE(atomic_kind_set)

      ALLOCATE (lri(nkind))
      ! Here we need only v_int and acoef (the latter as dummies)
      DO ikind = 1, nkind
         NULLIFY (lri(ikind)%acoef)
         NULLIFY (lri(ikind)%v_int)
         atomic_kind => atomic_kind_set(ikind)
         CALL get_atomic_kind(atomic_kind=atomic_kind, natom=natom)
         CALL get_qs_kind(qs_kind=qs_kind_set(ikind), nsgf=nsgf, basis_type="RI_AUX")
         ALLOCATE (lri(ikind)%acoef(natom, nsgf))
         lri(ikind)%acoef = 0._dp
         ALLOCATE (lri(ikind)%v_int(natom, nsgf))
         lri(ikind)%v_int = 0._dp
      END DO

   END SUBROUTINE make_lri_object

! **************************************************************************************************
!> \brief Read the external embedding potential, not to be optimized
!> \param qs_env ...
! **************************************************************************************************
   SUBROUTINE given_embed_pot(qs_env)
      TYPE(qs_environment_type), POINTER                 :: qs_env

      LOGICAL                                            :: open_shell_embed
      TYPE(dft_control_type), POINTER                    :: dft_control
      TYPE(pw_env_type), POINTER                         :: pw_env
      TYPE(pw_p_type), POINTER                           :: embed_pot, spin_embed_pot
      TYPE(pw_pool_type), POINTER                        :: auxbas_pw_pool_subsys
      TYPE(section_vals_type), POINTER                   :: input, qs_section

      qs_env%given_embed_pot = .TRUE.
      NULLIFY (input, dft_control, embed_pot, spin_embed_pot, embed_pot, spin_embed_pot, &
               qs_section)
      CALL get_qs_env(qs_env=qs_env, &
                      input=input, &
                      dft_control=dft_control, &
                      pw_env=pw_env)
      qs_section => section_vals_get_subs_vals(input, "DFT%QS")
      open_shell_embed = .FALSE.
      IF (dft_control%nspins .EQ. 2) open_shell_embed = .TRUE.

      ! Prepare plane-waves pool
      CALL pw_env_get(pw_env, auxbas_pw_pool=auxbas_pw_pool_subsys)

      ! Create embedding potential
      !CALL get_qs_env(qs_env=qs_env, &
      !                embed_pot=embed_pot)
      ALLOCATE (embed_pot)
      NULLIFY (embed_pot%pw)
      CALL pw_pool_create_pw(auxbas_pw_pool_subsys, embed_pot%pw, &
                             use_data=REALDATA3D, &
                             in_space=REALSPACE)
      IF (open_shell_embed) THEN
         ! Create spin embedding potential
         !  CALL get_qs_env(qs_env=qs_env, &
         !               spin_embed_pot=spin_embed_pot)
         ALLOCATE (spin_embed_pot)
         NULLIFY (spin_embed_pot%pw)
         CALL pw_pool_create_pw(auxbas_pw_pool_subsys, spin_embed_pot%pw, &
                                use_data=REALDATA3D, &
                                in_space=REALSPACE)
      ENDIF
      ! Read the cubes
      CALL read_embed_pot_cube(embed_pot, spin_embed_pot, qs_section, open_shell_embed)

      IF (.NOT. open_shell_embed) THEN
         CALL set_qs_env(qs_env=qs_env, embed_pot=embed_pot)
      ELSE
         CALL set_qs_env(qs_env=qs_env, embed_pot=embed_pot, spin_embed_pot=spin_embed_pot)
      ENDIF

   END SUBROUTINE given_embed_pot

! **************************************************************************************************
!> \brief ...
!> \param qs_env ...
!> \param embed_pot ...
!> \param spin_embed_pot ...
!> \param section ...
!> \param opt_embed ...
! **************************************************************************************************
   SUBROUTINE read_embed_pot(qs_env, embed_pot, spin_embed_pot, section, opt_embed)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(pw_p_type), POINTER                           :: embed_pot, spin_embed_pot
      TYPE(section_vals_type), POINTER                   :: section
      TYPE(opt_embed_pot_type)                           :: opt_embed

      ! Read the potential as a vector in the auxiliary basis
      IF (opt_embed%read_embed_pot) &
         CALL read_embed_pot_vector(qs_env, embed_pot, spin_embed_pot, section, opt_embed%embed_pot_coef, &
                                    opt_embed%open_shell_embed)
      ! Read the potential as a cube (two cubes for open shell)
      IF (opt_embed%read_embed_pot_cube) &
         CALL read_embed_pot_cube(embed_pot, spin_embed_pot, section, opt_embed%open_shell_embed)

   END SUBROUTINE read_embed_pot

! **************************************************************************************************
!> \brief ...
!> \param embed_pot ...
!> \param spin_embed_pot ...
!> \param section ...
!> \param open_shell_embed ...
! **************************************************************************************************
   SUBROUTINE read_embed_pot_cube(embed_pot, spin_embed_pot, section, open_shell_embed)
      TYPE(pw_p_type), POINTER                           :: embed_pot, spin_embed_pot
      TYPE(section_vals_type), POINTER                   :: section
      LOGICAL                                            :: open_shell_embed

      CHARACTER(LEN=default_path_length)                 :: filename
      LOGICAL                                            :: exist
      REAL(KIND=dp)                                      :: scaling_factor

      exist = .FALSE.
      CALL section_vals_val_get(section, "EMBED_CUBE_FILE_NAME", c_val=filename)
      INQUIRE (FILE=filename, exist=exist)
      IF (.NOT. exist) &
         CPABORT("Embedding cube file not found. ")

      scaling_factor = 1.0_dp
      CALL cube_to_pw(embed_pot%pw, filename, scaling_factor)

      ! Spin-dependent part of the potential
      IF (open_shell_embed) THEN
         exist = .FALSE.
         CALL section_vals_val_get(section, "EMBED_SPIN_CUBE_FILE_NAME", c_val=filename)
         INQUIRE (FILE=filename, exist=exist)
         IF (.NOT. exist) &
            CPABORT("Embedding spin cube file not found. ")

         scaling_factor = 1.0_dp
         CALL cube_to_pw(spin_embed_pot%pw, filename, scaling_factor)
      ENDIF

   END SUBROUTINE read_embed_pot_cube

! **************************************************************************************************
!> \brief Read the embedding potential from the binary file
!> \param qs_env ...
!> \param embed_pot ...
!> \param spin_embed_pot ...
!> \param section ...
!> \param embed_pot_coef ...
!> \param open_shell_embed ...
! **************************************************************************************************
   SUBROUTINE read_embed_pot_vector(qs_env, embed_pot, spin_embed_pot, section, embed_pot_coef, &
                                    open_shell_embed)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(pw_p_type), POINTER                           :: embed_pot, spin_embed_pot
      TYPE(section_vals_type), POINTER                   :: section
      TYPE(cp_fm_type), OPTIONAL, POINTER                :: embed_pot_coef
      LOGICAL                                            :: open_shell_embed

      CHARACTER(LEN=default_path_length)                 :: filename
      INTEGER                                            :: dimen_aux, dimen_restart_basis, &
                                                            dimen_var_aux, l_global, LLL, &
                                                            nrow_local, restart_unit
      INTEGER, DIMENSION(:), POINTER                     :: row_indices
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:)           :: coef, coef_read
      TYPE(cp_blacs_env_type), POINTER                   :: blacs_env
      TYPE(cp_fm_struct_type), POINTER                   :: fm_struct
      TYPE(cp_fm_type), POINTER                          :: my_embed_pot_coef
      TYPE(cp_para_env_type), POINTER                    :: para_env

      ! Get the vector dimension
      CALL find_aux_dimen(qs_env, dimen_aux)
      IF (open_shell_embed) THEN
         dimen_var_aux = dimen_aux*2
      ELSE
         dimen_var_aux = dimen_aux
      ENDIF

      ! We need a temporary vector of coefficients
      CALL get_qs_env(qs_env=qs_env, &
                      para_env=para_env)
      NULLIFY (blacs_env)
      NULLIFY (my_embed_pot_coef, fm_struct)
      CALL cp_blacs_env_create(blacs_env=blacs_env, para_env=para_env)
      CALL cp_fm_struct_create(fm_struct, para_env=para_env, context=blacs_env, &
                               nrow_global=dimen_var_aux, ncol_global=1)
      CALL cp_fm_create(my_embed_pot_coef, fm_struct, name="my_pot_coef")
      IF (.NOT. (PRESENT(embed_pot_coef))) THEN
         NULLIFY (embed_pot_coef)
         CALL cp_fm_create(my_embed_pot_coef, fm_struct, name="pot_coef")
      ENDIF

      CALL cp_fm_struct_release(fm_struct)
      CALL cp_fm_set_all(my_embed_pot_coef, 0.0_dp)

      ! Read the coefficients vector
      restart_unit = -1

      ! Allocate the attay to read the coefficients
      ALLOCATE (coef(dimen_var_aux))
      coef = 0.0_dp

      IF (para_env%ionode) THEN

         ! Get the restart file name
         CALL embed_restart_file_name(filename, section)

         CALL open_file(file_name=filename, &
                        file_action="READ", &
                        file_form="UNFORMATTED", &
                        file_status="OLD", &
                        unit_number=restart_unit)

         READ (restart_unit) dimen_restart_basis
         ! Check the dimensions of the bases: the actual and the restart one
         IF (.NOT. (dimen_restart_basis == dimen_aux)) &
            CPABORT("Wrong dimension of the embedding basis in the restart file.")

         ALLOCATE (coef_read(dimen_var_aux))
         coef_read = 0.0_dp

         READ (restart_unit) coef_read
         coef(:) = coef_read(:)
         DEALLOCATE (coef_read)

         ! Close restart file
         CALL close_file(unit_number=restart_unit)

      ENDIF

      ! Broadcast the coefficients on all processes
      CALL mp_bcast(coef, para_env%source, para_env%group)

      ! Copy to fm_type structure
      ! Information about full matrix gradient
      CALL cp_fm_get_info(matrix=my_embed_pot_coef, &
                          nrow_local=nrow_local, &
                          row_indices=row_indices)

      DO LLL = 1, nrow_local
         l_global = row_indices(LLL)
         my_embed_pot_coef%local_data(LLL, 1) = coef(l_global)
      ENDDO

      DEALLOCATE (coef)

      ! Copy to the my_embed_pot_coef to embed_pot_coef
      CALL cp_fm_copy_general(my_embed_pot_coef, embed_pot_coef, para_env)

      ! Build the embedding potential on the grid
      CALL update_embed_pot(embed_pot_coef, dimen_aux, embed_pot, spin_embed_pot, &
                            qs_env, .FALSE., open_shell_embed)

      ! Release my_embed_pot_coef
      CALL cp_fm_release(my_embed_pot_coef)
      IF (.NOT. (PRESENT(embed_pot_coef))) CALL cp_fm_release(embed_pot_coef)

      ! Release blacs environment
      CALL cp_blacs_env_release(blacs_env)

   END SUBROUTINE read_embed_pot_vector

! **************************************************************************************************
!> \brief Find the embedding restart file name
!> \param filename ...
!> \param section ...
! **************************************************************************************************
   SUBROUTINE embed_restart_file_name(filename, section)
      CHARACTER(LEN=default_path_length), INTENT(OUT)    :: filename
      TYPE(section_vals_type), POINTER                   :: section

      LOGICAL                                            :: exist

      exist = .FALSE.
      CALL section_vals_val_get(section, "EMBED_RESTART_FILE_NAME", c_val=filename)
      INQUIRE (FILE=filename, exist=exist)
      IF (.NOT. exist) &
         CPABORT("Embedding restart file not found. ")

   END SUBROUTINE embed_restart_file_name

! **************************************************************************************************
!> \brief Deallocate stuff for optimizing embedding potential
!> \param opt_embed ...
! **************************************************************************************************
   SUBROUTINE release_opt_embed(opt_embed)
      TYPE(opt_embed_pot_type)                           :: opt_embed

      INTEGER                                            :: ikind

      IF (.NOT. opt_embed%grid_opt) THEN
         CALL cp_fm_release(opt_embed%embed_pot_grad)
         CALL cp_fm_release(opt_embed%embed_pot_coef)
         CALL cp_fm_release(opt_embed%step)

         CALL cp_fm_release(opt_embed%prev_step)

         CALL cp_fm_release(opt_embed%embed_pot_hess)
         CALL cp_fm_release(opt_embed%prev_embed_pot_grad)
         CALL cp_fm_release(opt_embed%prev_embed_pot_coef)
         CALL cp_fm_release(opt_embed%prev_embed_pot_hess)
         CALL cp_fm_release(opt_embed%kinetic_mat)
         DEALLOCATE (opt_embed%w_func)
         IF (opt_embed%add_const_pot) THEN
            CALL pw_release(opt_embed%const_pot%pw)
            DEALLOCATE (opt_embed%const_pot)
         ENDIF
         DEALLOCATE (opt_embed%max_diff)
         DEALLOCATE (opt_embed%int_diff)

         DO ikind = 1, SIZE(opt_embed%lri)
            DEALLOCATE (opt_embed%lri(ikind)%v_int)
            DEALLOCATE (opt_embed%lri(ikind)%acoef)
         ENDDO
         DEALLOCATE (opt_embed%lri)
         !ELSE
         !DO ispin = 1, SIZE(opt_embed%prev_grid_pot)
         !   CALL pw_release(opt_embed%prev_grid_pot(ispin)%pw)
         !   CALL pw_release(opt_embed%prev_grid_grad(ispin)%pw)
         !ENDDO
         !DEALLOCATE (opt_embed%prev_grid_pot)
         !DEALLOCATE (opt_embed%prev_grid_grad)
      ENDIF

   END SUBROUTINE release_opt_embed

! **************************************************************************************************
!> \brief Creates a subsystem embedding potential
!> \param qs_env ...
!> \param embed_pot ...
!> \param embed_pot_subsys ...
!> \param spin_embed_pot ...
!> \param spin_embed_pot_subsys ...
!> \param open_shell_embed ...
!> \param change_spin_sign ...
!> \author Vladimir Rybkin
! **************************************************************************************************
   SUBROUTINE make_subsys_embed_pot(qs_env, embed_pot, embed_pot_subsys, &
                                    spin_embed_pot, spin_embed_pot_subsys, open_shell_embed, &
                                    change_spin_sign)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(pw_p_type), POINTER                           :: embed_pot, embed_pot_subsys, &
                                                            spin_embed_pot, spin_embed_pot_subsys
      LOGICAL                                            :: open_shell_embed, change_spin_sign

      TYPE(pw_env_type), POINTER                         :: pw_env
      TYPE(pw_pool_type), POINTER                        :: auxbas_pw_pool_subsys

      ! Extract  plane waves environment
      CALL get_qs_env(qs_env, pw_env=pw_env)

      ! Prepare plane-waves pool
      CALL pw_env_get(pw_env, auxbas_pw_pool=auxbas_pw_pool_subsys)

      ! Create embedding potential and set to zero
      NULLIFY (embed_pot_subsys)
      ALLOCATE (embed_pot_subsys)
      NULLIFY (embed_pot_subsys%pw)
      CALL pw_pool_create_pw(auxbas_pw_pool_subsys, embed_pot_subsys%pw, &
                             use_data=REALDATA3D, &
                             in_space=REALSPACE)

      ! Hard copy the grid
      embed_pot_subsys%pw%cr3d(:, :, :) = embed_pot%pw%cr3d(:, :, :)

      IF (open_shell_embed) THEN
         NULLIFY (spin_embed_pot_subsys)
         ALLOCATE (spin_embed_pot_subsys)
         NULLIFY (spin_embed_pot_subsys%pw)
         CALL pw_pool_create_pw(auxbas_pw_pool_subsys, spin_embed_pot_subsys%pw, &
                                use_data=REALDATA3D, &
                                in_space=REALSPACE)
         ! Hard copy the grid
         IF (change_spin_sign) THEN
            spin_embed_pot_subsys%pw%cr3d(:, :, :) = -spin_embed_pot%pw%cr3d(:, :, :)
         ELSE
            spin_embed_pot_subsys%pw%cr3d(:, :, :) = spin_embed_pot%pw%cr3d(:, :, :)
         ENDIF
      ENDIF

   END SUBROUTINE make_subsys_embed_pot

! **************************************************************************************************
!> \brief Calculates the derivative of the embedding potential wrt to the expansion coefficients
!> \param qs_env ...
!> \param diff_rho_r ...
!> \param diff_rho_spin ...
!> \param opt_embed ...
!> \author Vladimir Rybkin
! **************************************************************************************************

   SUBROUTINE calculate_embed_pot_grad(qs_env, diff_rho_r, diff_rho_spin, opt_embed)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(pw_p_type), POINTER                           :: diff_rho_r, diff_rho_spin
      TYPE(opt_embed_pot_type)                           :: opt_embed

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

      INTEGER                                            :: handle
      TYPE(cp_blacs_env_type), POINTER                   :: blacs_env
      TYPE(cp_fm_struct_type), POINTER                   :: fm_struct
      TYPE(cp_fm_type), POINTER                          :: embed_pot_coeff_spin, &
                                                            embed_pot_coeff_spinless, &
                                                            regular_term, spin_reg, spinless_reg
      TYPE(cp_para_env_type), POINTER                    :: para_env
      TYPE(pw_env_type), POINTER                         :: pw_env
      TYPE(pw_pool_type), POINTER                        :: auxbas_pw_pool

      CALL timeset(routineN, handle)

      ! We destroy the previous gradient and Hessian:
      ! current data are now previous data
      CALL cp_fm_to_fm(opt_embed%embed_pot_grad, opt_embed%prev_embed_pot_grad)
      CALL cp_fm_to_fm(opt_embed%embed_pot_Hess, opt_embed%prev_embed_pot_Hess)

      NULLIFY (pw_env)

      CALL get_qs_env(qs_env=qs_env, pw_env=pw_env, para_env=para_env)

      ! Get plane waves pool
      CALL pw_env_get(pw_env, auxbas_pw_pool=auxbas_pw_pool)

      ! Calculate potential gradient coefficients
      CALL calculate_embed_pot_grad_inner(qs_env, opt_embed%dimen_aux, diff_rho_r, diff_rho_spin, &
                                          opt_embed%embed_pot_grad, &
                                          opt_embed%open_shell_embed, opt_embed%lri)

      ! Add regularization with kinetic matrix
      IF (opt_embed%i_iter .EQ. 1) THEN ! Else it is kept in memory
         CALL compute_kinetic_mat(qs_env, opt_embed%kinetic_mat)
      ENDIF

      NULLIFY (regular_term)
      CALL cp_fm_get_info(matrix=opt_embed%embed_pot_grad, &
                          matrix_struct=fm_struct)
      CALL cp_fm_create(regular_term, fm_struct, name="regular_term")
      CALL cp_fm_set_all(regular_term, 0.0_dp)

      ! In case of open shell embedding we need two terms of dimen_aux=dimen_var_aux/2 for
      ! the spinless and the spin parts
      IF (opt_embed%open_shell_embed) THEN
         ! Prepare auxiliary full matrices
         NULLIFY (embed_pot_coeff_spinless, embed_pot_coeff_spin, spinless_reg, spin_reg, fm_struct, blacs_env)

         !CALL cp_blacs_env_create(blacs_env=blacs_env, para_env=para_env)

         CALL cp_fm_get_info(matrix=opt_embed%embed_pot_coef, context=blacs_env)
         CALL cp_fm_struct_create(fm_struct, para_env=para_env, context=blacs_env, &
                                  nrow_global=opt_embed%dimen_aux, ncol_global=1)
         CALL cp_fm_create(embed_pot_coeff_spinless, fm_struct, name="pot_coeff_spinless")
         CALL cp_fm_create(embed_pot_coeff_spin, fm_struct, name="pot_coeff_spin")
         CALL cp_fm_create(spinless_reg, fm_struct, name="spinless_reg")
         CALL cp_fm_create(spin_reg, fm_struct, name="spin_reg")
         CALL cp_fm_set_all(embed_pot_coeff_spinless, 0.0_dp)
         CALL cp_fm_set_all(embed_pot_coeff_spin, 0.0_dp)
         CALL cp_fm_set_all(spinless_reg, 0.0_dp)
         CALL cp_fm_set_all(spin_reg, 0.0_dp)
         CALL cp_fm_struct_release(fm_struct)

         ! Copy coefficients to the auxiliary structures
         CALL cp_fm_to_fm_submat(msource=opt_embed%embed_pot_coef, &
                                 mtarget=embed_pot_coeff_spinless, &
                                 nrow=opt_embed%dimen_aux, ncol=1, &
                                 s_firstrow=1, s_firstcol=1, &
                                 t_firstrow=1, t_firstcol=1)
         CALL cp_fm_to_fm_submat(msource=opt_embed%embed_pot_coef, &
                                 mtarget=embed_pot_coeff_spin, &
                                 nrow=opt_embed%dimen_aux, ncol=1, &
                                 s_firstrow=opt_embed%dimen_aux+1, s_firstcol=1, &
                                 t_firstrow=1, t_firstcol=1)
         ! Multiply
         CALL cp_gemm(transa="N", transb="N", m=opt_embed%dimen_aux, n=1, &
                      k=opt_embed%dimen_aux, alpha=1.0_dp, &
                      matrix_a=opt_embed%kinetic_mat, matrix_b=embed_pot_coeff_spinless, &
                      beta=0.0_dp, matrix_c=spinless_reg)
         CALL cp_gemm(transa="N", transb="N", m=opt_embed%dimen_aux, n=1, &
                      k=opt_embed%dimen_aux, alpha=1.0_dp, &
                      matrix_a=opt_embed%kinetic_mat, matrix_b=embed_pot_coeff_spin, &
                      beta=0.0_dp, matrix_c=spin_reg)
         ! Copy from the auxiliary structures to the full regularization term
         CALL cp_fm_to_fm_submat(msource=spinless_reg, &
                                 mtarget=regular_term, &
                                 nrow=opt_embed%dimen_aux, ncol=1, &
                                 s_firstrow=1, s_firstcol=1, &
                                 t_firstrow=1, t_firstcol=1)
         CALL cp_fm_to_fm_submat(msource=spin_reg, &
                                 mtarget=regular_term, &
                                 nrow=opt_embed%dimen_aux, ncol=1, &
                                 s_firstrow=1, s_firstcol=1, &
                                 t_firstrow=opt_embed%dimen_aux+1, t_firstcol=1)
         ! Release internally used auxiliary structures
         CALL cp_fm_release(embed_pot_coeff_spinless)
         CALL cp_fm_release(embed_pot_coeff_spin)
         CALL cp_fm_release(spin_reg)
         CALL cp_fm_release(spinless_reg)

      ELSE ! Simply multiply
         CALL cp_gemm(transa="N", transb="N", m=opt_embed%dimen_var_aux, n=1, &
                      k=opt_embed%dimen_var_aux, alpha=1.0_dp, &
                      matrix_a=opt_embed%kinetic_mat, matrix_b=opt_embed%embed_pot_coef, &
                      beta=0.0_dp, matrix_c=regular_term)
      ENDIF

      ! Scale by the regularization parameter and add to the gradient
      CALL cp_fm_scale_and_add(1.0_dp, opt_embed%embed_pot_grad, 4.0_dp*opt_embed%lambda, regular_term)

      ! Calculate the regularization contribution to the energy functional
      CALL cp_fm_trace(opt_embed%embed_pot_coef, regular_term, opt_embed%reg_term)
      opt_embed%reg_term = 2.0_dp*opt_embed%lambda*opt_embed%reg_term

      ! Deallocate regular term
      CALL cp_fm_release(regular_term)

      CALL timestop(handle)

   END SUBROUTINE calculate_embed_pot_grad

! **************************************************************************************************
!> \brief Performs integration for the embedding potential gradient
!> \param qs_env ...
!> \param dimen_aux ...
!> \param rho_r ...
!> \param rho_spin ...
!> \param embed_pot_grad ...
!> \param open_shell_embed ...
!> \param lri ...
!> \author Vladimir Rybkin
! **************************************************************************************************
   SUBROUTINE calculate_embed_pot_grad_inner(qs_env, dimen_aux, rho_r, rho_spin, embed_pot_grad, &
                                             open_shell_embed, lri)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      INTEGER                                            :: dimen_aux
      TYPE(pw_p_type), POINTER                           :: rho_r, rho_spin
      TYPE(cp_fm_type), POINTER                          :: embed_pot_grad
      LOGICAL                                            :: open_shell_embed
      TYPE(lri_kind_type), DIMENSION(:), POINTER         :: lri

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

      INTEGER                                            :: handle, iatom, ikind, l_global, LLL, &
                                                            nrow_local, nsgf, start_pos
      INTEGER, DIMENSION(:), POINTER                     :: row_indices
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:)           :: pot_grad
      TYPE(atomic_kind_type), DIMENSION(:), POINTER      :: atomic_kind_set
      TYPE(cell_type), POINTER                           :: cell
      TYPE(cp_para_env_type), POINTER                    :: para_env
      TYPE(dft_control_type), POINTER                    :: dft_control
      TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
      TYPE(qs_kind_type), DIMENSION(:), POINTER          :: qs_kind_set

! Needed to store integrals

      CALL timeset(routineN, handle)

      CALL get_qs_env(qs_env=qs_env, &
                      particle_set=particle_set, &
                      qs_kind_set=qs_kind_set, &
                      dft_control=dft_control, &
                      cell=cell, &
                      atomic_kind_set=atomic_kind_set, &
                      para_env=para_env)

      ! Create wf_vector and gradient
      IF (open_shell_embed) THEN
         ALLOCATE (pot_grad(dimen_aux*2))
      ELSE
         ALLOCATE (pot_grad(dimen_aux))
      ENDIF

      ! Use lri subroutine
      DO ikind = 1, SIZE(lri)
         lri(ikind)%v_int = 0.0_dp
      END DO

      CALL integrate_v_rspace_one_center(rho_r, qs_env, lri, &
                                         .FALSE., "RI_AUX")
      DO ikind = 1, SIZE(lri)
         CALL mp_sum(lri(ikind)%v_int, para_env%group)
      END DO

      pot_grad = 0.0_dp
      start_pos = 1
      DO ikind = 1, SIZE(lri)
         DO iatom = 1, SIZE(lri(ikind)%v_int, DIM=1)
            nsgf = SIZE(lri(ikind)%v_int(iatom, :))
            pot_grad(start_pos:start_pos+nsgf-1) = lri(ikind)%v_int(iatom, :)
            start_pos = start_pos+nsgf
         ENDDO
      ENDDO

      ! Open-shell embedding
      IF (open_shell_embed) THEN
         DO ikind = 1, SIZE(lri)
            lri(ikind)%v_int = 0.0_dp
         END DO

         CALL integrate_v_rspace_one_center(rho_spin, qs_env, lri, &
                                            .FALSE., "RI_AUX")
         DO ikind = 1, SIZE(lri)
            CALL mp_sum(lri(ikind)%v_int, para_env%group)
         END DO

         start_pos = dimen_aux+1
         DO ikind = 1, SIZE(lri)
            DO iatom = 1, SIZE(lri(ikind)%v_int, DIM=1)
               nsgf = SIZE(lri(ikind)%v_int(iatom, :))
               pot_grad(start_pos:start_pos+nsgf-1) = lri(ikind)%v_int(iatom, :)
               start_pos = start_pos+nsgf
            ENDDO
         ENDDO
      ENDIF

      ! Scale by the cell volume
      pot_grad = pot_grad*rho_r%pw%pw_grid%dvol

      ! Information about full matrix gradient
      CALL cp_fm_get_info(matrix=embed_pot_grad, &
                          nrow_local=nrow_local, &
                          row_indices=row_indices)

      ! Copy the gradient into the full matrix
      DO LLL = 1, nrow_local
         l_global = row_indices(LLL)
         embed_pot_grad%local_data(LLL, 1) = pot_grad(l_global)
      ENDDO

      DEALLOCATE (pot_grad)

      CALL timestop(handle)

   END SUBROUTINE calculate_embed_pot_grad_inner

! **************************************************************************************************
!> \brief Calculates kinetic energy matrix in auxiliary basis in the fm format
!> \param qs_env ...
!> \param kinetic_mat ...
!> \author Vladimir Rybkin
! **************************************************************************************************
   SUBROUTINE compute_kinetic_mat(qs_env, kinetic_mat)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(cp_fm_type), POINTER                          :: kinetic_mat

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

      INTEGER                                            :: handle
      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: matrix_t
      TYPE(neighbor_list_set_p_type), DIMENSION(:), &
         POINTER                                         :: sab_orb
      TYPE(qs_ks_env_type), POINTER                      :: ks_env

      CALL timeset(routineN, handle)

      NULLIFY (ks_env, sab_orb, matrix_t)

      ! First, get the dbcsr structure from the overlap matrix
      CALL get_qs_env(qs_env=qs_env, ks_env=ks_env, sab_orb=sab_orb)

      ! Calculate kinetic matrix
      CALL build_kinetic_matrix(ks_env, matrix_t=matrix_t, &
                                matrix_name="KINETIC ENERGY MATRIX", &
                                basis_type="RI_AUX", &
                                sab_nl=sab_orb, calculate_forces=.FALSE.)

      ! Change to the fm format
      CALL copy_dbcsr_to_fm(matrix_t(1)%matrix, kinetic_mat)

      ! Release memory
      CALL dbcsr_deallocate_matrix_set(matrix_t)

      CALL timestop(handle)

   END SUBROUTINE compute_kinetic_mat

! **************************************************************************************************
!> \brief Takes maximization step in embedding potential optimization
!> \param diff_rho_r ...
!> \param diff_rho_spin ...
!> \param opt_embed ...
!> \param embed_pot ...
!> \param spin_embed_pot ...
!> \param qs_env ...
!> \author Vladimir Rybkin
! **************************************************************************************************
   SUBROUTINE opt_embed_step(diff_rho_r, diff_rho_spin, opt_embed, embed_pot, spin_embed_pot, qs_env)
      TYPE(pw_p_type), POINTER                           :: diff_rho_r, diff_rho_spin
      TYPE(opt_embed_pot_type)                           :: opt_embed
      TYPE(pw_p_type), POINTER                           :: embed_pot, spin_embed_pot
      TYPE(qs_environment_type), POINTER                 :: qs_env

      CHARACTER(LEN=*), PARAMETER :: routineN = 'opt_embed_step', routineP = moduleN//':'//routineN
      REAL(KIND=dp), PARAMETER                           :: thresh = 0.000001_dp

      INTEGER                                            :: handle, l_global, LLL, nrow_local
      INTEGER, DIMENSION(:), POINTER                     :: row_indices
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:)           :: eigenval
      TYPE(cp_fm_struct_type), POINTER                   :: fm_struct
      TYPE(cp_fm_type), POINTER                          :: diag_grad, diag_step, fm_U, fm_U_scale

      CALL timeset(routineN, handle)

      IF (opt_embed%grid_opt) THEN ! Grid based optimization
         CALL grid_based_step(diff_rho_r, diff_rho_spin, opt_embed, embed_pot, spin_embed_pot)

      ELSE ! Finite basis optimization
         ! If the previous step has been rejected, we go back to the previous expansion coefficients
         IF (.NOT. opt_embed%accept_step) &
            CALL cp_fm_scale_and_add(1.0_dp, opt_embed%embed_pot_coef, -1.0_dp, opt_embed%step)

         ! Do a simple steepest descent
         IF (opt_embed%steep_desc) THEN
            IF (opt_embed%i_iter .GT. 2) &
               opt_embed%trust_rad = Barzilai_Borwein(opt_embed%step, opt_embed%prev_step, &
                                                      opt_embed%embed_pot_grad, opt_embed%prev_embed_pot_grad)
            IF (ABS(opt_embed%trust_rad) .GT. opt_embed%max_trad) THEN
               IF (opt_embed%trust_rad .GT. 0.0_dp) THEN
                  opt_embed%trust_rad = opt_embed%max_trad
               ELSE
                  opt_embed%trust_rad = -opt_embed%max_trad
               ENDIF
            ENDIF

            CALL cp_fm_to_fm(opt_embed%step, opt_embed%prev_step)
            CALL cp_fm_scale_and_add(0.0_dp, opt_embed%prev_step, 1.0_dp, opt_embed%step)
            CALL cp_fm_set_all(opt_embed%step, 0.0_dp)
            CALL cp_fm_scale_and_add(1.0_dp, opt_embed%step, opt_embed%trust_rad, opt_embed%embed_pot_grad)
            opt_embed%step_len = opt_embed%trust_rad
         ELSE

            ! First, update the Hessian inverse if needed
            IF (opt_embed%i_iter > 1) THEN
               IF (opt_embed%accept_step) & ! We don't update Hessian if the step has been rejected
                  CALL symm_rank_one_update(opt_embed%embed_pot_grad, opt_embed%prev_embed_pot_grad, &
                                            opt_embed%step, opt_embed%prev_embed_pot_Hess, opt_embed%embed_pot_Hess)
            ENDIF

            ! Add regularization term to the Hessian
            !CALL cp_fm_scale_and_add(1.0_dp, opt_embed%embed_pot_Hess, 4.0_dp*opt_embed%lambda, &
            !                         opt_embed%kinetic_mat)

            ! Else use the first initial Hessian. Now it's just the unit matrix: embed_pot_hess
            ! Second, invert the Hessian
            ALLOCATE (eigenval(opt_embed%dimen_var_aux))
            eigenval = 0.0_dp
            NULLIFY (fm_U, fm_U_scale, diag_grad, diag_step)
            CALL cp_fm_get_info(matrix=opt_embed%embed_pot_hess, &
                                matrix_struct=fm_struct)
            CALL cp_fm_create(fm_U, fm_struct, name="fm_U")
            CALL cp_fm_create(fm_U_scale, fm_struct, name="fm_U")
            CALL cp_fm_set_all(fm_U, 0.0_dp)
            CALL cp_fm_set_all(fm_U_scale, 0.0_dp)
            CALL cp_fm_get_info(matrix=opt_embed%embed_pot_grad, &
                                matrix_struct=fm_struct)
            CALL cp_fm_create(diag_grad, fm_struct, name="diag_grad")
            CALL cp_fm_set_all(diag_grad, 0.0_dp)
            CALL cp_fm_create(diag_step, fm_struct, name="diag_step")
            CALL cp_fm_set_all(diag_step, 0.0_dp)

            ! Store the Hessian as it will be destroyed in diagonalization: use fm_U_scal for it
            CALL cp_fm_to_fm(opt_embed%embed_pot_hess, fm_U_scale)

            ! Diagonalize Hessian
            CALL choose_eigv_solver(opt_embed%embed_pot_hess, fm_U, eigenval)

            ! Copy the Hessian back
            CALL cp_fm_to_fm(fm_U_scale, opt_embed%embed_pot_hess)

            ! Find the step in diagonal representation, begin with gradient
            CALL cp_gemm(transa="T", transb="N", m=opt_embed%dimen_var_aux, n=1, &
                         k=opt_embed%dimen_var_aux, alpha=1.0_dp, &
                         matrix_a=fm_U, matrix_b=opt_embed%embed_pot_grad, beta=0.0_dp, &
                         matrix_c=diag_grad)

            CALL cp_fm_get_info(matrix=opt_embed%embed_pot_coef, &
                                nrow_local=nrow_local, &
                                row_indices=row_indices)

            DO LLL = 1, nrow_local
               l_global = row_indices(LLL)
               IF (ABS(eigenval(l_global)) .GE. thresh) THEN
                  diag_step%local_data(LLL, 1) = &
                     -diag_grad%local_data(LLL, 1)/(eigenval(l_global))
               ELSE
                  diag_step%local_data(LLL, 1) = 0.0_dp
               ENDIF
            ENDDO
            CALL cp_fm_trace(diag_step, diag_step, opt_embed%step_len)

            ! Transform step to a non-diagonal representation
            CALL cp_gemm(transa="N", transb="N", m=opt_embed%dimen_var_aux, n=1, &
                         k=opt_embed%dimen_var_aux, alpha=1.0_dp, &
                         matrix_a=fm_U, matrix_b=diag_step, beta=0.0_dp, &
                         matrix_c=opt_embed%step)

            ! Now use fm_U_scale for scaled eigenvectors
            CALL cp_fm_to_fm(fm_U, fm_U_scale)
            CALL cp_fm_column_scale(fm_U_scale, eigenval)

            CALL cp_fm_release(fm_U_scale)

            ! Scale the step to fit within the trust radius: it it's less already,
            ! then take the Newton step
            CALL cp_fm_trace(opt_embed%step, opt_embed%step, opt_embed%step_len)
            IF (opt_embed%step_len .GT. opt_embed%trust_rad) THEN

               IF (opt_embed%level_shift) THEN
                  ! Find a level shift parameter and apply it
                  CALL level_shift(opt_embed, diag_grad, eigenval, diag_step)
               ELSE ! Just scale
                  CALL cp_fm_trace(diag_step, diag_step, opt_embed%step_len)
                  CALL cp_fm_scale(opt_embed%trust_rad/opt_embed%step_len, diag_step)
               ENDIF
               CALL cp_fm_trace(diag_step, diag_step, opt_embed%step_len)
               ! Transform step to a non-diagonal representation
               CALL cp_gemm(transa="N", transb="N", m=opt_embed%dimen_var_aux, n=1, &
                            k=opt_embed%dimen_var_aux, alpha=1.0_dp, &
                            matrix_a=fm_U, matrix_b=diag_step, beta=0.0_dp, &
                            matrix_c=opt_embed%step)
               CALL cp_fm_trace(opt_embed%step, opt_embed%step, opt_embed%step_len)

               ! Recalculate step in diagonal representation
               opt_embed%newton_step = .FALSE.
            ELSE
               opt_embed%newton_step = .TRUE.
            ENDIF

            ! Release some memory
            DEALLOCATE (eigenval)
            ! Release more memory
            CALL cp_fm_release(diag_grad)
            CALL cp_fm_release(diag_step)
            CALL cp_fm_release(fm_U)

         ENDIF ! grad_descent

         ! Update the coefficients
         CALL cp_fm_scale_and_add(1.0_dp, opt_embed%embed_pot_coef, 1.0_dp, opt_embed%step)

         ! Update the embedding potential
         CALL update_embed_pot(opt_embed%embed_pot_coef, opt_embed%dimen_aux, embed_pot, &
                               spin_embed_pot, qs_env, opt_embed%add_const_pot, &
                               opt_embed%open_shell_embed, opt_embed%const_pot)
      ENDIF ! Grid-based optimization

      CALL timestop(handle)

   END SUBROUTINE opt_embed_step

!
! **************************************************************************************************
!> \brief ...
!> \param diff_rho_r ...
!> \param diff_rho_spin ...
!> \param opt_embed ...
!> \param embed_pot ...
!> \param spin_embed_pot ...
! **************************************************************************************************
   SUBROUTINE grid_based_step(diff_rho_r, diff_rho_spin, opt_embed, embed_pot, spin_embed_pot)
      TYPE(pw_p_type), POINTER                           :: diff_rho_r, diff_rho_spin
      TYPE(opt_embed_pot_type)                           :: opt_embed
      TYPE(pw_p_type), POINTER                           :: embed_pot, spin_embed_pot

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

      INTEGER                                            :: handle

      CALL timeset(routineN, handle)

      opt_embed%step_len = opt_embed%trust_rad
      ! Take the step for spin-free part
      CALL pw_axpy(diff_rho_r%pw, embed_pot%pw, opt_embed%step_len)

      IF (opt_embed%open_shell_embed) THEN
         CALL pw_axpy(diff_rho_spin%pw, spin_embed_pot%pw, opt_embed%step_len)
      ENDIF

      CALL timestop(handle)

   END SUBROUTINE grid_based_step

! **************************************************************************************************
!> \brief ... Adds variable part of to the embedding potential
!> \param embed_pot_coef ...
!> \param dimen_aux ...
!> \param embed_pot ...
!> \param spin_embed_pot ...
!> \param qs_env ...
!> \param add_const_pot ...
!> \param open_shell_embed ...
!> \param const_pot ...
!> \author Vladimir Rybkin
! **************************************************************************************************

   SUBROUTINE update_embed_pot(embed_pot_coef, dimen_aux, embed_pot, spin_embed_pot, &
                               qs_env, add_const_pot, open_shell_embed, const_pot)
      TYPE(cp_fm_type), POINTER                          :: embed_pot_coef
      INTEGER                                            :: dimen_aux
      TYPE(pw_p_type), POINTER                           :: embed_pot, spin_embed_pot
      TYPE(qs_environment_type), POINTER                 :: qs_env
      LOGICAL                                            :: add_const_pot, open_shell_embed
      TYPE(pw_p_type), OPTIONAL, POINTER                 :: const_pot

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

      INTEGER                                            :: handle, l_global, LLL, nrow_local
      INTEGER, DIMENSION(:), POINTER                     :: row_indices
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:)           :: wf_vector
      TYPE(atomic_kind_type), DIMENSION(:), POINTER      :: atomic_kind_set
      TYPE(cell_type), POINTER                           :: cell
      TYPE(cp_blacs_env_type), POINTER                   :: blacs_env
      TYPE(cp_fm_struct_type), POINTER                   :: fm_struct
      TYPE(cp_fm_type), POINTER                          :: embed_pot_coef_spin, &
                                                            embed_pot_coef_spinless, mo_coeff
      TYPE(cp_para_env_type), POINTER                    :: para_env
      TYPE(dft_control_type), POINTER                    :: dft_control
      TYPE(mo_set_p_type), DIMENSION(:), POINTER         :: mos
      TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
      TYPE(pw_env_type), POINTER                         :: pw_env
      TYPE(pw_p_type)                                    :: psi_L, rho_g
      TYPE(pw_pool_type), POINTER                        :: auxbas_pw_pool
      TYPE(qs_kind_type), DIMENSION(:), POINTER          :: qs_kind_set

      CALL timeset(routineN, handle)
      ! Get MO coefficients: we need only the structure, therefore don't care about the spin
      CALL get_qs_env(qs_env=qs_env, &
                      particle_set=particle_set, &
                      qs_kind_set=qs_kind_set, &
                      dft_control=dft_control, &
                      cell=cell, &
                      atomic_kind_set=atomic_kind_set, &
                      pw_env=pw_env, mos=mos, para_env=para_env)
      CALL get_mo_set(mo_set=mos(1)%mo_set, mo_coeff=mo_coeff)

      ! Get plane waves pool
      CALL pw_env_get(pw_env, auxbas_pw_pool=auxbas_pw_pool)

      ! get some of the grids ready
      NULLIFY (rho_g%pw)
      CALL pw_pool_create_pw(auxbas_pw_pool, rho_g%pw, &
                             use_data=COMPLEXDATA1D, &
                             in_space=RECIPROCALSPACE)

      NULLIFY (psi_L%pw)
      CALL pw_pool_create_pw(auxbas_pw_pool, psi_L%pw, &
                             use_data=REALDATA3D, &
                             in_space=REALSPACE)

      ! Create wf_vector and auxiliary wave functions
      ALLOCATE (wf_vector(dimen_aux))
      wf_vector = 0.0_dp

      ! Create auxiliary full matrices for open-shell case
      IF (open_shell_embed) THEN
         NULLIFY (blacs_env)
         CALL cp_fm_get_info(matrix=embed_pot_coef, context=blacs_env)
         CALL cp_fm_struct_create(fm_struct, para_env=para_env, context=blacs_env, &
                                  nrow_global=dimen_aux, ncol_global=1)
         CALL cp_fm_create(embed_pot_coef_spinless, fm_struct, name="pot_coeff_spinless")
         CALL cp_fm_create(embed_pot_coef_spin, fm_struct, name="pot_coeff_spin")
         CALL cp_fm_set_all(embed_pot_coef_spinless, 0.0_dp)
         CALL cp_fm_set_all(embed_pot_coef_spin, 0.0_dp)
         CALL cp_fm_struct_release(fm_struct)

         ! Copy coefficients to the auxiliary structures
         CALL cp_fm_to_fm_submat(embed_pot_coef, &
                                 mtarget=embed_pot_coef_spinless, &
                                 nrow=dimen_aux, ncol=1, &
                                 s_firstrow=1, s_firstcol=1, &
                                 t_firstrow=1, t_firstcol=1)
         CALL cp_fm_to_fm_submat(embed_pot_coef, &
                                 mtarget=embed_pot_coef_spin, &
                                 nrow=dimen_aux, ncol=1, &
                                 s_firstrow=dimen_aux+1, s_firstcol=1, &
                                 t_firstrow=1, t_firstcol=1)

         ! Spinless potential
         CALL cp_fm_get_info(matrix=embed_pot_coef_spinless, &
                             nrow_local=nrow_local, &
                             row_indices=row_indices)

         ! Copy fm_coeff to an array
         DO LLL = 1, nrow_local
            l_global = row_indices(LLL)
            wf_vector(l_global) = embed_pot_coef_spinless%local_data(LLL, 1)
         ENDDO
         CALL mp_sum(wf_vector, para_env%group)

         ! Calculate the variable part of the embedding potential
         CALL calculate_wavefunction(mo_coeff, 1, psi_L, rho_g, atomic_kind_set, &
                                     qs_kind_set, cell, dft_control, particle_set, pw_env, &
                                     basis_type="RI_AUX", &
                                     external_vector=wf_vector)
         ! Update the full embedding potential
         IF (add_const_pot) THEN
            CALL pw_copy(const_pot%pw, embed_pot%pw)
         ELSE
            CALL pw_zero(embed_pot%pw)
         ENDIF

         CALL pw_axpy(psi_L%pw, embed_pot%pw)

         ! Spin-dependent potential
         wf_vector = 0.0_dp
         CALL cp_fm_get_info(matrix=embed_pot_coef_spin, &
                             nrow_local=nrow_local, &
                             row_indices=row_indices)

         ! Copy fm_coeff to an array
         DO LLL = 1, nrow_local
            l_global = row_indices(LLL)
            wf_vector(l_global) = embed_pot_coef_spin%local_data(LLL, 1)
         ENDDO
         CALL mp_sum(wf_vector, para_env%group)

         ! Calculate the variable part of the embedding potential
         CALL calculate_wavefunction(mo_coeff, 1, psi_L, rho_g, atomic_kind_set, &
                                     qs_kind_set, cell, dft_control, particle_set, pw_env, &
                                     basis_type="RI_AUX", &
                                     external_vector=wf_vector)
         ! No constant potential for spin-dependent potential
         CALL pw_zero(spin_embed_pot%pw)
         CALL pw_axpy(psi_L%pw, spin_embed_pot%pw)

      ELSE ! Closed shell

         CALL cp_fm_get_info(matrix=embed_pot_coef, &
                             nrow_local=nrow_local, &
                             row_indices=row_indices)

         ! Copy fm_coeff to an array
         DO LLL = 1, nrow_local
            l_global = row_indices(LLL)
            wf_vector(l_global) = embed_pot_coef%local_data(LLL, 1)
         ENDDO
         CALL mp_sum(wf_vector, para_env%group)

         ! Calculate the variable part of the embedding potential
         CALL calculate_wavefunction(mo_coeff, 1, psi_L, rho_g, atomic_kind_set, &
                                     qs_kind_set, cell, dft_control, particle_set, pw_env)

         CALL calculate_wavefunction(mo_coeff, 1, psi_L, rho_g, atomic_kind_set, &
                                     qs_kind_set, cell, dft_control, particle_set, pw_env, &
                                     basis_type="RI_AUX", &
                                     external_vector=wf_vector)

         ! Update the full embedding potential
         IF (add_const_pot) THEN
            CALL pw_copy(const_pot%pw, embed_pot%pw)
         ELSE
            CALL pw_zero(embed_pot%pw)
         ENDIF

         CALL pw_axpy(psi_L%pw, embed_pot%pw)
      ENDIF ! Open/closed shell

      ! Deallocate memory and release objects
      DEALLOCATE (wf_vector)
      CALL pw_pool_give_back_pw(auxbas_pw_pool, psi_L%pw)
      CALL pw_pool_give_back_pw(auxbas_pw_pool, rho_g%pw)

      IF (open_shell_embed) THEN
         CALL cp_fm_release(embed_pot_coef_spin)
         CALL cp_fm_release(embed_pot_coef_spinless)
      ENDIF

      CALL timestop(handle)

   END SUBROUTINE update_embed_pot

! **************************************************************************************************
!> \brief BFGS update of the inverse Hessian in the full matrix format
!> \param grad ...
!> \param prev_grad ...
!> \param step ...
!> \param prev_inv_Hess ...
!> \param inv_Hess ...
!> \author Vladimir Rybkin
! **************************************************************************************************
   SUBROUTINE inv_Hessian_update(grad, prev_grad, step, prev_inv_Hess, inv_Hess)
      TYPE(cp_fm_type), POINTER                          :: grad, prev_grad, step, prev_inv_Hess, &
                                                            inv_Hess

      INTEGER                                            :: mat_size
      REAL(KIND=dp)                                      :: factor1, s_dot_y, y_dot_B_inv_y
      TYPE(cp_fm_struct_type), POINTER                   :: fm_struct_mat, fm_struct_vec
      TYPE(cp_fm_type), POINTER                          :: B_inv_y, B_inv_y_s, s_s, s_y, s_y_B_inv, &
                                                            y

      ! Recover the dimension
      CALL cp_fm_get_info(matrix=inv_Hess, &
                          nrow_global=mat_size)

      CALL cp_fm_set_all(inv_Hess, 0.0_dp)
      CALL cp_fm_to_fm(prev_inv_Hess, inv_Hess)

      ! Get full matrix structures
      NULLIFY (fm_struct_mat, fm_struct_vec)

      CALL cp_fm_get_info(matrix=prev_inv_Hess, &
                          matrix_struct=fm_struct_mat)
      CALL cp_fm_get_info(matrix=grad, &
                          matrix_struct=fm_struct_vec)

      ! Allocate intermediates
      NULLIFY (B_inv_y, y, s_s, s_y, B_inv_y_s, s_y_B_inv)
      CALL cp_fm_create(B_inv_y, fm_struct_vec, name="B_inv_y")
      CALL cp_fm_create(y, fm_struct_vec, name="y")

      CALL cp_fm_create(s_s, fm_struct_mat, name="s_s")
      CALL cp_fm_create(s_y, fm_struct_mat, name="s_y")
      CALL cp_fm_create(B_inv_y_s, fm_struct_mat, name="B_inv_y_s")
      CALL cp_fm_create(s_y_B_inv, fm_struct_mat, name="s_y_B_inv")

      CALL cp_fm_set_all(B_inv_y, 0.0_dp)
      CALL cp_fm_set_all(s_s, 0.0_dp)
      CALL cp_fm_set_all(s_y, 0.0_dp)
      CALL cp_fm_set_all(B_inv_y_s, 0.0_dp)
      CALL cp_fm_set_all(s_y_B_inv, 0.0_dp)

      ! Calculate intermediates
      ! y the is gradient difference
      CALL cp_fm_get_info(matrix=grad)
      CALL cp_fm_to_fm(grad, y)
      CALL cp_fm_scale_and_add(1.0_dp, y, -1.0_dp, prev_grad)

      ! First term
      CALL cp_gemm(transa="N", transb="N", m=mat_size, n=1, &
                   k=mat_size, alpha=1.0_dp, &
                   matrix_a=prev_inv_Hess, matrix_b=y, beta=0.0_dp, &
                   matrix_c=B_inv_y)

      CALL cp_gemm(transa="N", transb="T", m=mat_size, n=mat_size, &
                   k=1, alpha=1.0_dp, &
                   matrix_a=step, matrix_b=step, beta=0.0_dp, &
                   matrix_c=s_s)

      CALL cp_gemm(transa="N", transb="T", m=mat_size, n=mat_size, &
                   k=1, alpha=1.0_dp, &
                   matrix_a=step, matrix_b=y, beta=0.0_dp, &
                   matrix_c=s_y)

      CALL cp_fm_trace(step, y, s_dot_y)

      CALL cp_fm_trace(y, y, s_dot_y)
      CALL cp_fm_trace(step, step, s_dot_y)

      CALL cp_fm_trace(y, B_inv_y, y_dot_B_inv_y)

      factor1 = (s_dot_y+y_dot_B_inv_y)/(s_dot_y)**2

      CALL cp_fm_scale_and_add(1.0_dp, inv_Hess, factor1, s_s)

      ! Second term
      CALL cp_gemm(transa="N", transb="T", m=mat_size, n=mat_size, &
                   k=1, alpha=1.0_dp, &
                   matrix_a=B_inv_y, matrix_b=step, beta=0.0_dp, &
                   matrix_c=B_inv_y_s)

      CALL cp_gemm(transa="N", transb="N", m=mat_size, n=mat_size, &
                   k=mat_size, alpha=1.0_dp, &
                   matrix_a=s_y, matrix_b=prev_inv_Hess, beta=0.0_dp, &
                   matrix_c=s_y_B_inv)

      CALL cp_fm_scale_and_add(1.0_dp, B_inv_y_s, 1.0_dp, s_y_B_inv)

      ! Assemble the new inverse Hessian
      CALL cp_fm_scale_and_add(1.0_dp, inv_Hess, -s_dot_y, B_inv_y_s)

      ! Deallocate intermediates
      CALL cp_fm_release(y)
      CALL cp_fm_release(B_inv_y)
      CALL cp_fm_release(s_s)
      CALL cp_fm_release(s_y)
      CALL cp_fm_release(B_inv_y_s)
      CALL cp_fm_release(s_y_B_inv)

   END SUBROUTINE inv_Hessian_update

! **************************************************************************************************
!> \brief ...
!> \param grad ...
!> \param prev_grad ...
!> \param step ...
!> \param prev_Hess ...
!> \param Hess ...
! **************************************************************************************************
   SUBROUTINE Hessian_update(grad, prev_grad, step, prev_Hess, Hess)
      TYPE(cp_fm_type), POINTER                          :: grad, prev_grad, step, prev_Hess, Hess

      INTEGER                                            :: mat_size
      REAL(KIND=dp)                                      :: s_b_s, y_t_s
      TYPE(cp_blacs_env_type), POINTER                   :: blacs_env
      TYPE(cp_fm_struct_type), POINTER                   :: fm_struct_mat, fm_struct_vec, &
                                                            fm_struct_vec_t
      TYPE(cp_fm_type), POINTER                          :: B_s, B_s_s_B, s_t_B, y, y_y_t
      TYPE(cp_para_env_type), POINTER                    :: para_env

      ! Recover the dimension
      CALL cp_fm_get_info(matrix=Hess, &
                          nrow_global=mat_size, para_env=para_env)

      CALL cp_fm_set_all(Hess, 0.0_dp)
      CALL cp_fm_to_fm(prev_Hess, Hess)

      ! WARNING: our Hessian must be negative-definite, whereas BFGS makes it positive-definite!
      ! Therefore, we change sign in the beginning and in the end.
      CALL cp_fm_scale(-1.0_dp, Hess)

      ! Create blacs environment
      NULLIFY (blacs_env)
      CALL cp_blacs_env_create(blacs_env=blacs_env, para_env=para_env)

      ! Get full matrix structures
      NULLIFY (fm_struct_mat, fm_struct_vec, fm_struct_vec_t)

      CALL cp_fm_get_info(matrix=prev_Hess, &
                          matrix_struct=fm_struct_mat)
      CALL cp_fm_get_info(matrix=grad, &
                          matrix_struct=fm_struct_vec)
      CALL cp_fm_struct_create(fm_struct_vec_t, para_env=para_env, context=blacs_env, &
                               nrow_global=1, ncol_global=mat_size)

      ! Allocate intermediates
      NULLIFY (y_y_t, B_s, s_t_B, B_s_s_B, y)
      CALL cp_fm_create(B_s, fm_struct_vec, name="B_s")
      CALL cp_fm_create(s_t_B, fm_struct_vec_t, name="s_t_B")
      CALL cp_fm_create(y, fm_struct_vec, name="y")

      CALL cp_fm_create(y_y_t, fm_struct_mat, name="y_y_t")
      CALL cp_fm_create(B_s_s_B, fm_struct_mat, name="B_s_s_B")

      CALL cp_fm_set_all(y_y_t, 0.0_dp)
      CALL cp_fm_set_all(y, 0.0_dp)
      CALL cp_fm_set_all(B_s_s_B, 0.0_dp)
      CALL cp_fm_set_all(B_s, 0.0_dp)
      CALL cp_fm_set_all(s_t_B, 0.0_dp)

      ! Release the structure created only here
      CALL cp_fm_struct_release(fm_struct_vec_t)

      ! Calculate intermediates
      ! y the is gradient difference
      CALL cp_fm_to_fm(grad, y)
      CALL cp_fm_scale_and_add(1.0_dp, y, -1.0_dp, prev_grad)

      ! First term
      CALL cp_gemm(transa="N", transb="T", m=mat_size, n=mat_size, &
                   k=1, alpha=1.0_dp, &
                   matrix_a=y, matrix_b=y, beta=0.0_dp, &
                   matrix_c=y_y_t)

      CALL cp_fm_trace(y, step, y_t_s)

      CALL cp_fm_scale_and_add(1.0_dp, Hess, (1.0_dp/y_t_s), y_y_t)

      ! Second term
      CALL cp_gemm(transa="N", transb="N", m=mat_size, n=1, &
                   k=mat_size, alpha=1.0_dp, &
                   matrix_a=Hess, matrix_b=step, beta=0.0_dp, &
                   matrix_c=B_s)

      CALL cp_fm_trace(B_s, step, s_B_s)

      CALL cp_gemm(transa="T", transb="N", m=1, n=mat_size, &
                   k=mat_size, alpha=1.0_dp, &
                   matrix_a=step, matrix_b=Hess, beta=0.0_dp, &
                   matrix_c=s_t_B)

      CALL cp_gemm(transa="N", transb="N", m=mat_size, n=mat_size, &
                   k=1, alpha=1.0_dp, &
                   matrix_a=B_s, matrix_b=s_t_B, beta=0.0_dp, &
                   matrix_c=B_s_s_B)

      CALL cp_fm_scale_and_add(1.0_dp, Hess, -(1.0_dp/s_B_s), B_s_s_B)

      ! WARNING: our Hessian must be negative-definite, whereas BFGS makes it positive-definite!
      ! Therefore, we change sign in the beginning and in the end.
      CALL cp_fm_scale(-1.0_dp, Hess)

      ! Release blacs environment
      CALL cp_blacs_env_release(blacs_env)

      ! Deallocate intermediates
      CALL cp_fm_release(y_y_t)
      CALL cp_fm_release(B_s_s_B)
      CALL cp_fm_release(B_s)
      CALL cp_fm_release(s_t_B)
      CALL cp_fm_release(y)

   END SUBROUTINE Hessian_update

! **************************************************************************************************
!> \brief ...
!> \param grad ...
!> \param prev_grad ...
!> \param step ...
!> \param prev_Hess ...
!> \param Hess ...
! **************************************************************************************************
   SUBROUTINE symm_rank_one_update(grad, prev_grad, step, prev_Hess, Hess)
      TYPE(cp_fm_type), POINTER                          :: grad, prev_grad, step, prev_Hess, Hess

      INTEGER                                            :: mat_size
      REAL(KIND=dp)                                      :: factor
      TYPE(cp_fm_struct_type), POINTER                   :: fm_struct_mat, fm_struct_vec
      TYPE(cp_fm_type), POINTER                          :: B_x, y, y_B_x_y_B_x

      ! Recover the dimension
      CALL cp_fm_get_info(matrix=Hess, nrow_global=mat_size)

      CALL cp_fm_set_all(Hess, 0.0_dp)
      CALL cp_fm_to_fm(prev_Hess, Hess)

      ! Get full matrix structures
      NULLIFY (fm_struct_mat, fm_struct_vec)

      CALL cp_fm_get_info(matrix=prev_Hess, &
                          matrix_struct=fm_struct_mat)
      CALL cp_fm_get_info(matrix=grad, &
                          matrix_struct=fm_struct_vec)

      ! Allocate intermediates
      NULLIFY (y, B_x, y_B_x_y_B_x)
      CALL cp_fm_create(y, fm_struct_vec, name="y")
      CALL cp_fm_create(B_x, fm_struct_vec, name="B_x")
      CALL cp_fm_create(y_B_x_y_B_x, fm_struct_mat, name="y_B_x_y_B_x")

      CALL cp_fm_set_all(y, 0.0_dp)
      CALL cp_fm_set_all(B_x, 0.0_dp)
      CALL cp_fm_set_all(y_B_x_y_B_x, 0.0_dp)

      ! Calculate intermediates
      ! y the is gradient difference
      CALL cp_fm_to_fm(grad, y)
      CALL cp_fm_scale_and_add(1.0_dp, y, -1.0_dp, prev_grad)

      CALL cp_gemm(transa="N", transb="N", m=mat_size, n=1, &
                   k=mat_size, alpha=1.0_dp, &
                   matrix_a=Hess, matrix_b=step, beta=0.0_dp, &
                   matrix_c=B_x)

      CALL cp_fm_scale_and_add(1.0_dp, y, -1.0_dp, B_x)

      CALL cp_gemm(transa="N", transb="T", m=mat_size, n=mat_size, &
                   k=1, alpha=1.0_dp, &
                   matrix_a=y, matrix_b=y, beta=0.0_dp, &
                   matrix_c=y_B_x_y_B_x)

      ! Scaling factor
      CALL cp_fm_trace(y, step, factor)

      ! Assemble the Hessian
      CALL cp_fm_scale_and_add(1.0_dp, Hess, (1.0_dp/factor), y_B_x_y_B_x)

      ! Deallocate intermediates
      CALL cp_fm_release(y)
      CALL cp_fm_release(B_x)
      CALL cp_fm_release(y_B_x_y_B_x)

   END SUBROUTINE symm_rank_one_update

! **************************************************************************************************
!> \brief Controls the step, changes the trust radius if needed in maximization of the V_emb
!> \param opt_embed ...
!> \author Vladimir Rybkin
! **************************************************************************************************
   SUBROUTINE step_control(opt_embed)
      TYPE(opt_embed_pot_type)                           :: opt_embed

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

      INTEGER                                            :: handle
      REAL(KIND=dp)                                      :: actual_ener_change, ener_ratio, &
                                                            lin_term, pred_ener_change, quad_term
      TYPE(cp_fm_struct_type), POINTER                   :: fm_struct
      TYPE(cp_fm_type), POINTER                          :: H_b

      CALL timeset(routineN, handle)

      NULLIFY (H_b, fm_struct)
      CALL cp_fm_get_info(matrix=opt_embed%embed_pot_grad, &
                          matrix_struct=fm_struct)
      CALL cp_fm_create(H_b, fm_struct, name="H_b")
      CALL cp_fm_set_all(H_b, 0.0_dp)

      ! Calculate the quadratic estimate for the energy
      ! Linear term
      CALL cp_fm_trace(opt_embed%step, opt_embed%embed_pot_grad, lin_term)

      ! Quadratic term
      CALL cp_gemm(transa="N", transb="N", m=opt_embed%dimen_aux, n=1, &
                   k=opt_embed%dimen_aux, alpha=1.0_dp, &
                   matrix_a=opt_embed%embed_pot_Hess, matrix_b=opt_embed%step, &
                   beta=0.0_dp, matrix_c=H_b)
      CALL cp_fm_trace(opt_embed%step, H_b, quad_term)

      pred_ener_change = lin_term+0.5_dp*quad_term

      ! Reveal actual energy change
      actual_ener_change = opt_embed%w_func(opt_embed%i_iter)- &
                           opt_embed%w_func(opt_embed%last_accepted)

      ener_ratio = actual_ener_change/pred_ener_change

      CALL cp_fm_release(H_b)

      IF (actual_ener_change .GT. 0.0_dp) THEN ! If energy increases
         ! We accept step
         opt_embed%accept_step = .TRUE.
         ! If energy change is larger than the predicted one, increase trust radius twice
         ! Else (between 0 and 1) leave as it is, unless Newton step has been taken and if the step is less than max
         IF ((ener_ratio .GT. 1.0_dp) .AND. (.NOT. opt_embed%newton_step) .AND. &
             (opt_embed%trust_rad .LT. opt_embed%max_trad)) &
            opt_embed%trust_rad = 2.0_dp*opt_embed%trust_rad
      ELSE ! Energy decreases
         ! If the decrease is not too large we allow this step to be taken
         ! Otherwise, the step is rejected
         IF (ABS(actual_ener_change) .GE. opt_embed%allowed_decrease) THEN
            opt_embed%accept_step = .FALSE.
         ENDIF
         ! Trust radius is decreased 4 times unless it's smaller than the minimal allowed value
         IF (opt_embed%trust_rad .GE. opt_embed%min_trad) &
            opt_embed%trust_rad = 0.25_dp*opt_embed%trust_rad
      ENDIF

      IF (opt_embed%accept_step) opt_embed%last_accepted = opt_embed%i_iter

      CALL timestop(handle)

   END SUBROUTINE step_control

! **************************************************************************************************
!> \brief ...
!> \param opt_embed ...
!> \param diag_grad ...
!> \param eigenval ...
!> \param diag_step ...
! **************************************************************************************************
   SUBROUTINE level_shift(opt_embed, diag_grad, eigenval, diag_step)
      TYPE(opt_embed_pot_type)                           :: opt_embed
      TYPE(cp_fm_type), POINTER                          :: diag_grad
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:)           :: eigenval
      TYPE(cp_fm_type), POINTER                          :: diag_step

      CHARACTER(LEN=*), PARAMETER :: routineN = 'level_shift', routineP = moduleN//':'//routineN
      INTEGER, PARAMETER                                 :: max_iter = 25
      REAL(KIND=dp), PARAMETER                           :: thresh = 0.00001_dp

      INTEGER                                            :: handle, i_iter, l_global, LLL, &
                                                            min_index, nrow_local
      INTEGER, ALLOCATABLE, DIMENSION(:)                 :: red_eigenval_map
      INTEGER, DIMENSION(:), POINTER                     :: row_indices
      LOGICAL                                            :: converged, do_shift
      REAL(KIND=dp) :: diag_grad_norm, grad_min, hess_min, shift, shift_max, shift_min, step_len, &
         step_minus_trad, step_minus_trad_first, step_minus_trad_max, step_minus_trad_min
      TYPE(cp_para_env_type), POINTER                    :: para_env

      CALL timeset(routineN, handle)

      ! Array properties
      CALL cp_fm_get_info(matrix=opt_embed%embed_pot_coef, &
                          nrow_local=nrow_local, &
                          row_indices=row_indices, &
                          para_env=para_env)

      min_index = MINLOC(ABS(eigenval), dim=1)
      hess_min = eigenval(min_index)
      CALL cp_fm_get_element(diag_grad, min_index, 1, grad_min)

      CALL cp_fm_trace(diag_grad, diag_grad, diag_grad_norm)

      IF (hess_min .LT. 0.0_dp) THEN
         !shift_min = -2.0_dp*(diag_grad_norm/opt_embed%trust_rad - min(hess_min, 0.0_dp))
         !shift_max = max(0.0_dp, -hess_min + 0.5_dp*grad_min/opt_embed%trust_rad)
         !shift_max = MIN(-hess_min+0.5_dp*grad_min/opt_embed%trust_rad, 0.0_dp)
         shift_max = hess_min+0.1
         shift_min = diag_grad_norm/opt_embed%trust_rad
         shift_min = 10.0_dp
         !If (abs(shift_max) .LE. thresh) then
         !   shift_min = -20.0_dp*(diag_grad_norm/opt_embed%trust_rad - min(hess_min, 0.0_dp))
         !Else
         !   shift_min = 20.0_dp*shift_max
         !Endif

         ! The boundary values
         step_minus_trad_max = shifted_step(diag_grad, eigenval, shift_max, opt_embed%trust_rad)
         step_minus_trad_min = shifted_step(diag_grad, eigenval, shift_min, opt_embed%trust_rad)

         ! Find zero by bisection
         converged = .FALSE.
         do_shift = .FALSE.
         IF (ABS(step_minus_trad_max) .LE. thresh) THEN
            shift = shift_max
         ELSE
            IF (ABS(step_minus_trad_min) .LE. thresh) THEN
               shift = shift_min
            ELSE
               DO i_iter = 1, max_iter
                  shift = 0.5_dp*(shift_max+shift_min)
                  step_minus_trad = shifted_step(diag_grad, eigenval, shift, opt_embed%trust_rad)
                  IF (i_iter .EQ. 1) step_minus_trad_first = step_minus_trad
                  IF (step_minus_trad .GT. 0.0_dp) shift_max = shift
                  IF (step_minus_trad .LT. 0.0_dp) shift_min = shift
                  !IF (ABS(shift_max-shift_min) .LT. thresh) converged = .TRUE.
                  IF (ABS(step_minus_trad) .LT. thresh) converged = .TRUE.
                  IF (converged) EXIT
               ENDDO
               IF (ABS(step_minus_trad) .LT. ABS(step_minus_trad_first)) do_shift = .TRUE.
            ENDIF
         ENDIF
         ! Apply level-shifting
         IF (converged .OR. do_shift) THEN
            DO LLL = 1, nrow_local
               l_global = row_indices(LLL)
               IF (ABS(eigenval(l_global)) .GE. thresh) THEN
                  diag_step%local_data(LLL, 1) = &
                     -diag_grad%local_data(LLL, 1)/(eigenval(l_global)-shift)
               ELSE
                  diag_step%local_data(LLL, 1) = 0.0_dp
               ENDIF
            ENDDO
         ENDIF
         IF (.NOT. converged) THEN ! Scale if shift has not been found
            CALL cp_fm_trace(diag_step, diag_step, step_len)
            CALL cp_fm_scale(opt_embed%trust_rad/step_len, diag_step)
         ENDIF

         ! Special case
      ELSE ! Hess min .LT. 0.0_dp
         ! First, find all positive eigenvalues
         ALLOCATE (red_eigenval_map(opt_embed%dimen_var_aux))
         red_eigenval_map = 0
         DO LLL = 1, nrow_local
            l_global = row_indices(LLL)
            IF (eigenval(l_global) .GE. 0.0_dp) THEN
               red_eigenval_map(l_global) = 1
            ENDIF
         ENDDO
         CALL mp_sum(red_eigenval_map, para_env%group)

         ! Set shift as -hess_min and find step on the reduced space of negative-value
         ! eigenvectors
         shift = -hess_min
         DO LLL = 1, nrow_local
            l_global = row_indices(LLL)
            IF (red_eigenval_map(l_global) .EQ. 0) THEN
               IF (ABS(eigenval(l_global)) .GE. thresh) THEN
                  diag_step%local_data(LLL, 1) = &
                     -diag_grad%local_data(LLL, 1)/(eigenval(l_global)-shift)
               ELSE
                  diag_step%local_data(LLL, 1) = 0.0_dp
               ENDIF
            ELSE
               diag_step%local_data(LLL, 1) = 0.0_dp
            ENDIF
         ENDDO

         ! Find the step length of such a step
         CALL cp_fm_trace(diag_step, diag_step, step_len)

      ENDIF

      CALL timestop(handle)

   END SUBROUTINE level_shift

! **************************************************************************************************
!> \brief ...
!> \param diag_grad ...
!> \param eigenval ...
!> \param shift ...
!> \param trust_rad ...
!> \return ...
! **************************************************************************************************
   FUNCTION shifted_step(diag_grad, eigenval, shift, trust_rad) RESULT(step_minus_trad)
      TYPE(cp_fm_type), POINTER                          :: diag_grad
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:)           :: eigenval
      REAL(KIND=dp)                                      :: shift, trust_rad, step_minus_trad

      REAL(KIND=dp), PARAMETER                           :: thresh = 0.000001_dp

      INTEGER                                            :: l_global, LLL, nrow_local
      INTEGER, DIMENSION(:), POINTER                     :: row_indices
      REAL(KIND=dp)                                      :: step, step_1d
      TYPE(cp_para_env_type), POINTER                    :: para_env

      CALL cp_fm_get_info(matrix=diag_grad, &
                          nrow_local=nrow_local, &
                          row_indices=row_indices, &
                          para_env=para_env)

      step = 0.0_dp
      DO LLL = 1, nrow_local
         l_global = row_indices(LLL)
         IF ((ABS(eigenval(l_global)) .GE. thresh) .AND. (ABS(diag_grad%local_data(LLL, 1)) .GE. thresh)) THEN
            step_1d = -diag_grad%local_data(LLL, 1)/(eigenval(l_global)+shift)
            step = step+step_1d**2
         ENDIF
      ENDDO

      CALL mp_sum(step, para_env%group)

      step_minus_trad = SQRT(step)-trust_rad

   END FUNCTION shifted_step

! **************************************************************************************************
!> \brief ...
!> \param step ...
!> \param prev_step ...
!> \param grad ...
!> \param prev_grad ...
!> \return ...
!> \retval length ...
! **************************************************************************************************
   FUNCTION Barzilai_Borwein(step, prev_step, grad, prev_grad) RESULT(length)
      TYPE(cp_fm_type), POINTER                          :: step, prev_step, grad, prev_grad
      REAL(KIND=dp)                                      :: length

      REAL(KIND=dp)                                      :: denominator, numerator
      TYPE(cp_fm_struct_type), POINTER                   :: fm_struct
      TYPE(cp_fm_type), POINTER                          :: grad_diff, step_diff

      ! Get full matrix structures
      NULLIFY (fm_struct)

      CALL cp_fm_get_info(matrix=grad, &
                          matrix_struct=fm_struct)

      ! Allocate intermediates
      NULLIFY (grad_diff, step_diff)
      CALL cp_fm_create(grad_diff, fm_struct, name="grad_diff")
      CALL cp_fm_create(step_diff, fm_struct, name="step_diff")

      ! Calculate intermediates
      CALL cp_fm_to_fm(grad, grad_diff)
      CALL cp_fm_to_fm(step, step_diff)

      CALL cp_fm_scale_and_add(1.0_dp, grad_diff, -1.0_dp, prev_grad)
      CALL cp_fm_scale_and_add(1.0_dp, step_diff, -1.0_dp, prev_step)

      CALL cp_fm_trace(step_diff, grad_diff, numerator)
      CALL cp_fm_trace(grad_diff, grad_diff, denominator)

      ! Release intermediates
      CALL cp_fm_release(grad_diff)
      CALL cp_fm_release(step_diff)

      length = numerator/denominator

   END FUNCTION Barzilai_Borwein

! **************************************************************************************************
!> \brief ...
!> \param diff_rho_r ...
!> \return ...
! **************************************************************************************************
   FUNCTION max_dens_diff(diff_rho_r) RESULT(total_max_diff)
      TYPE(pw_p_type)                                    :: diff_rho_r
      REAL(KIND=dp)                                      :: total_max_diff

      INTEGER                                            :: size_x, size_y, size_z
      REAL(KIND=dp)                                      :: max_diff
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :)     :: grid_3d

      !, i_x, i_y, i_z

      ! Get the sizes
      size_x = SIZE(diff_rho_r%pw%cr3d, 1)
      size_y = SIZE(diff_rho_r%pw%cr3d, 2)
      size_z = SIZE(diff_rho_r%pw%cr3d, 3)

      ! Allocate the density
      ALLOCATE (grid_3d(size_x, size_y, size_z))

      ! Copy density
      grid_3d(:, :, :) = diff_rho_r%pw%cr3d(:, :, :)

      ! Find the maximum absolute value
      max_diff = MAXVAL(ABS(grid_3d))
      total_max_diff = max_diff
      CALL mp_max(total_max_diff, diff_rho_r%pw%pw_grid%para%group)

      ! Deallocate the density
      DEALLOCATE (grid_3d)

   END FUNCTION max_dens_diff

! **************************************************************************************************
!> \brief Prints a cube for the (rho_A + rho_B - rho_ref) to be minimized in embedding
!> \param diff_rho_r ...
!> \param i_iter ...
!> \param qs_env ...
!> \author Vladimir Rybkin
! **************************************************************************************************
   SUBROUTINE print_rho_diff(diff_rho_r, i_iter, qs_env)
      TYPE(pw_p_type), POINTER                           :: diff_rho_r
      INTEGER                                            :: i_iter
      TYPE(qs_environment_type), POINTER                 :: qs_env

      CHARACTER(LEN=default_path_length)                 :: filename, my_pos_cube, title
      INTEGER                                            :: unit_nr
      TYPE(cp_logger_type), POINTER                      :: logger
      TYPE(particle_list_type), POINTER                  :: particles
      TYPE(qs_subsys_type), POINTER                      :: subsys
      TYPE(section_vals_type), POINTER                   :: dft_section, input

      NULLIFY (subsys, input)

      CALL get_qs_env(qs_env=qs_env, &
                      subsys=subsys, &
                      input=input)
      dft_section => section_vals_get_subs_vals(input, "DFT")
      CALL qs_subsys_get(subsys, particles=particles)

      logger => cp_get_default_logger()
      IF (BTEST(cp_print_key_should_output(logger%iter_info, input, &
                                           "DFT%QS%OPT_EMBED%EMBED_DENS_DIFF"), cp_p_file)) THEN
         my_pos_cube = "REWIND"
         WRITE (filename, '(a5,I3.3,a1,I1.1)') "DIFF_", i_iter
         unit_nr = cp_print_key_unit_nr(logger, input, "DFT%QS%OPT_EMBED%EMBED_DENS_DIFF", &
                                        extension=".cube", middle_name=TRIM(filename), file_position=my_pos_cube, &
                                        log_filename=.FALSE.)

         WRITE (title, *) "EMBEDDING DENSITY DIFFERENCE ", " optimization step ", i_iter
         CALL cp_pw_to_cube(diff_rho_r%pw, unit_nr, title, particles=particles, &
                            stride=section_get_ivals(dft_section, "QS%OPT_EMBED%EMBED_DENS_DIFF%STRIDE"))
         CALL cp_print_key_finished_output(unit_nr, logger, input, &
                                           "DFT%QS%OPT_EMBED%EMBED_DENS_DIFF")
      ENDIF

   END SUBROUTINE print_rho_diff

! **************************************************************************************************
!> \brief Print embedding potential as a cube and as a binary (for restarting)
!> \param qs_env ...
!> \param dimen_aux ...
!> \param embed_pot_coef ...
!> \param embed_pot ...
!> \param i_iter ...
!> \param embed_pot_spin ...
!> \param open_shell_embed ...
!> \param grid_opt ...
! **************************************************************************************************
   SUBROUTINE print_embed_restart(qs_env, dimen_aux, embed_pot_coef, embed_pot, i_iter, &
                                  embed_pot_spin, open_shell_embed, grid_opt)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      INTEGER                                            :: dimen_aux
      TYPE(cp_fm_type), POINTER                          :: embed_pot_coef
      TYPE(pw_p_type), POINTER                           :: embed_pot
      INTEGER                                            :: i_iter
      TYPE(pw_p_type), POINTER                           :: embed_pot_spin
      LOGICAL                                            :: open_shell_embed, grid_opt

      CHARACTER(LEN=default_path_length)                 :: filename, my_pos_cube, title
      INTEGER                                            :: unit_nr
      TYPE(cp_logger_type), POINTER                      :: logger
      TYPE(particle_list_type), POINTER                  :: particles
      TYPE(qs_subsys_type), POINTER                      :: subsys
      TYPE(section_vals_type), POINTER                   :: dft_section, input

      NULLIFY (input)
      CALL get_qs_env(qs_env=qs_env, subsys=subsys, &
                      input=input)

      ! First we print an unformatted file
      IF (.NOT. grid_opt) THEN ! Only for finite basis optimization
         logger => cp_get_default_logger()
         IF (BTEST(cp_print_key_should_output(logger%iter_info, input, &
                                              "DFT%QS%OPT_EMBED%EMBED_POT_VECTOR"), cp_p_file)) THEN

            WRITE (filename, '(a10,I3.3)') "embed_pot_", i_iter
            unit_nr = cp_print_key_unit_nr(logger, input, "DFT%QS%OPT_EMBED%EMBED_POT_VECTOR", extension=".wfn", &
                                           middle_name=TRIM(filename), file_form="UNFORMATTED")
            IF (unit_nr > 0) THEN
               WRITE (unit_nr) dimen_aux
            ENDIF
            CALL cp_fm_write_unformatted(embed_pot_coef, unit_nr)
            IF (unit_nr > 0) THEN
               CALL close_file(unit_nr)
            ENDIF
         ENDIF
      ENDIF

      ! Second, cube files
      dft_section => section_vals_get_subs_vals(input, "DFT")
      CALL qs_subsys_get(subsys, particles=particles)

      logger => cp_get_default_logger()
      IF (BTEST(cp_print_key_should_output(logger%iter_info, input, &
                                           "DFT%QS%OPT_EMBED%EMBED_POT_CUBE"), cp_p_file)) THEN
         my_pos_cube = "REWIND"
         WRITE (filename, '(a10,I3.3)') "embed_pot_", i_iter
         unit_nr = cp_print_key_unit_nr(logger, input, "DFT%QS%OPT_EMBED%EMBED_POT_CUBE", &
                                        extension=".cube", middle_name=TRIM(filename), file_position=my_pos_cube, &
                                        log_filename=.FALSE.)

         WRITE (title, *) "EMBEDDING POTENTIAL at optimization step ", i_iter
         CALL cp_pw_to_cube(embed_pot%pw, unit_nr, title, particles=particles)
!, &
!                            stride=section_get_ivals(dft_section, "QS%OPT_EMBED%EMBED_POT_CUBE%STRIDE"))
         CALL cp_print_key_finished_output(unit_nr, logger, input, &
                                           "DFT%QS%OPT_EMBED%EMBED_POT_CUBE")
         IF (open_shell_embed) THEN ! Print spin part of the embedding potential
            my_pos_cube = "REWIND"
            WRITE (filename, '(a15,I3.3)') "spin_embed_pot_", i_iter
            unit_nr = cp_print_key_unit_nr(logger, input, "DFT%QS%OPT_EMBED%EMBED_POT_CUBE", &
                                           extension=".cube", middle_name=TRIM(filename), file_position=my_pos_cube, &
                                           log_filename=.FALSE.)

            WRITE (title, *) "SPIN EMBEDDING POTENTIAL at optimization step ", i_iter
            CALL cp_pw_to_cube(embed_pot_spin%pw, unit_nr, title, particles=particles)
!,  &
!                               stride=section_get_ivals(dft_section, "QS%OPT_EMBED%EMBED_POT_CUBE%STRIDE"))
            CALL cp_print_key_finished_output(unit_nr, logger, input, &
                                              "DFT%QS%OPT_EMBED%EMBED_POT_CUBE")
         ENDIF
      ENDIF

   END SUBROUTINE print_embed_restart

! **************************************************************************************************
!> \brief ...
!> \param output_unit ...
!> \param step_num ...
!> \param opt_embed ...
! **************************************************************************************************
   SUBROUTINE print_emb_opt_info(output_unit, step_num, opt_embed)
      INTEGER                                            :: output_unit, step_num
      TYPE(opt_embed_pot_type)                           :: opt_embed

      IF (output_unit > 0) THEN
         WRITE (UNIT=output_unit, FMT="(/,T2,8('-'),A,I5,1X,12('-'))") &
            "  Optimize embedding potential info at step = ", step_num
         WRITE (UNIT=output_unit, FMT="(T2,A,F20.10)") &
            " Functional value         = ", opt_embed%w_func(step_num)
         IF (step_num .GT. 1) THEN
            WRITE (UNIT=output_unit, FMT="(T2,A,F20.10)") &
               " Real energy change         = ", opt_embed%w_func(step_num)- &
               opt_embed%w_func(step_num-1)

            WRITE (UNIT=output_unit, FMT="(T2,A,F20.10)") &
               " Step size                  = ", opt_embed%step_len

         END IF

         WRITE (UNIT=output_unit, FMT="(T2,A,F20.10)") &
            " Trust radius               = ", opt_embed%trust_rad

         WRITE (UNIT=output_unit, FMT="(T2,51('-'))")
      END IF

   END SUBROUTINE print_emb_opt_info

! **************************************************************************************************
!> \brief ...
!> \param opt_embed ...
!> \param diff_rho_r ...
!> \param diff_rho_spin ...
!> \param output_unit ...
! **************************************************************************************************
   SUBROUTINE conv_check_embed(opt_embed, diff_rho_r, diff_rho_spin, output_unit)
      TYPE(opt_embed_pot_type)                           :: opt_embed
      TYPE(pw_p_type), POINTER                           :: diff_rho_r, diff_rho_spin
      INTEGER                                            :: output_unit

      LOGICAL                                            :: conv_int_diff, conv_max_diff
      REAL(KIND=dp)                                      :: int_diff, int_diff_spin, max_diff, &
                                                            max_diff_spin

      ! Calculate the convergence target values
      opt_embed%max_diff(1) = max_dens_diff(diff_rho_r)
      opt_embed%int_diff(1) = pw_integrate_function(fun=diff_rho_r%pw, oprt='ABS')
      IF (opt_embed%open_shell_embed) THEN
         opt_embed%max_diff(2) = max_dens_diff(diff_rho_spin)
         opt_embed%int_diff(2) = pw_integrate_function(fun=diff_rho_spin%pw, oprt='ABS')
      ENDIF

      ! Find out the convergence
      max_diff = opt_embed%max_diff(1)

      ! Open shell
      IF (opt_embed%open_shell_embed) THEN
         max_diff_spin = opt_embed%max_diff(2)
         IF ((max_diff .LE. opt_embed%conv_max) .AND. (max_diff_spin .LE. opt_embed%conv_max_spin)) THEN
            conv_max_diff = .TRUE.
         ELSE
            conv_max_diff = .FALSE.
         ENDIF
      ELSE
         ! Closed shell
         IF (max_diff .LE. opt_embed%conv_max) THEN
            conv_max_diff = .TRUE.
         ELSE
            conv_max_diff = .FALSE.
         ENDIF
      ENDIF

      int_diff = opt_embed%int_diff(1)
      ! Open shell
      IF (opt_embed%open_shell_embed) THEN
         int_diff_spin = opt_embed%int_diff(2)
         IF ((int_diff .LE. opt_embed%conv_int) .AND. (int_diff_spin .LE. opt_embed%conv_int_spin)) THEN
            conv_int_diff = .TRUE.
         ELSE
            conv_int_diff = .FALSE.
         ENDIF
      ELSE
         ! Closed shell
         IF (int_diff .LE. opt_embed%conv_int) THEN
            conv_int_diff = .TRUE.
         ELSE
            conv_int_diff = .FALSE.
         ENDIF
      ENDIF

      IF ((conv_max_diff) .AND. (conv_int_diff)) THEN
         opt_embed%converged = .TRUE.
      ELSE
         opt_embed%converged = .FALSE.
      ENDIF

      ! Print the information
      IF (output_unit > 0) THEN
         WRITE (UNIT=output_unit, FMT="(/,T2,A)") &
            " Convergence check :"

         ! Maximum value of density
         WRITE (UNIT=output_unit, FMT="(T2,A,F20.10)") &
            " Maximum density difference                = ", max_diff
         WRITE (UNIT=output_unit, FMT="(T2,A,F20.10)") &
            " Convergence limit for max. density diff.  = ", opt_embed%conv_max

         IF (opt_embed%open_shell_embed) THEN

            WRITE (UNIT=output_unit, FMT="(T2,A,F20.10)") &
               " Maximum spin density difference           = ", max_diff_spin
            WRITE (UNIT=output_unit, FMT="(T2,A,F20.10)") &
               " Convergence limit for max. spin dens.diff.= ", opt_embed%conv_max_spin

         ENDIF

         IF (conv_max_diff) THEN
            WRITE (UNIT=output_unit, FMT="(T2,2A)") &
               " Convergence in max. density diff.    =     ", &
               "             YES"
         ELSE
            WRITE (UNIT=output_unit, FMT="(T2,2A)") &
               " Convergence in max. density diff.    =     ", &
               "              NO"
         END IF

         ! Integrated abs. value of density
         WRITE (UNIT=output_unit, FMT="(T2,A,F20.10)") &
            " Integrated density difference             = ", int_diff
         WRITE (UNIT=output_unit, FMT="(T2,A,F20.10)") &
            " Conv. limit for integrated density diff.  = ", opt_embed%conv_int
         IF (opt_embed%open_shell_embed) THEN
            WRITE (UNIT=output_unit, FMT="(T2,A,F20.10)") &
               " Integrated spin density difference        = ", int_diff_spin
            WRITE (UNIT=output_unit, FMT="(T2,A,F20.10)") &
               " Conv. limit for integrated spin dens.diff.= ", opt_embed%conv_int_spin
         ENDIF

         IF (conv_int_diff) THEN
            WRITE (UNIT=output_unit, FMT="(T2,2A)") &
               " Convergence in integrated density diff.    =     ", &
               "             YES"
         ELSE
            WRITE (UNIT=output_unit, FMT="(T2,2A)") &
               " Convergence in integrated density diff.    =     ", &
               "              NO"
         END IF
      ENDIF

      IF ((opt_embed%converged) .AND. (output_unit > 0)) THEN
         WRITE (UNIT=output_unit, FMT="(/,T2,A)") REPEAT("*", 79)
         WRITE (UNIT=output_unit, FMT="(T2,A,T25,A,T78,A)") &
            "***", "EMBEDDING POTENTIAL OPTIMIZATION COMPLETED", "***"
         WRITE (UNIT=output_unit, FMT="(T2,A)") REPEAT("*", 79)
      END IF

   END SUBROUTINE conv_check_embed

END MODULE optimize_embedding_potential
