!> @file
!!   Routines to precondition wavefunctions
!! @author
!!    Copyright (C) 2005-2011 BigDFT group
!!    This file is distributed under the terms of the
!!    GNU General Public License, see ~/COPYING file
!!    or http://www.gnu.org/copyleft/gpl.txt .
!!    For the list of contributors, see ~/AUTHORS


!> Calls the preconditioner for each orbital treated by the processor
subroutine preconditionall(orbs,lr,hx,hy,hz,ncong,hpsi,gnrm,gnrm_zero)
  use module_precisions
  use module_types
  use locregs
  use locreg_operations, only: confpot_data, workarrays_quartic_convolutions,workarr_precond
  use at_domain, only: domain_geocode
  use wrapper_linalg
  use module_bigdft_mpi
  implicit none
  integer, intent(in) :: ncong
  real(gp), intent(in) :: hx,hy,hz
  type(locreg_descriptors), intent(in) :: lr
  type(orbitals_data), intent(in) :: orbs
  real(dp), intent(out) :: gnrm,gnrm_zero
  real(wp), dimension(lr%wfd%nvctr_c+7*lr%wfd%nvctr_f,orbs%nspinor,orbs%norbp), intent(inout) :: hpsi
  !local variables
  type(confpot_data) :: confdata
  type(workarrays_quartic_convolutions) :: lin_prec_conv_work
  type(workarr_precond) :: lin_prec_work
  integer :: iorb,inds,ncplx,ikpt,jorb
  real(wp) :: cprecr,scpr,evalmax,eval_zero
  real(gp) :: kx,ky,kz

  ! Preconditions all orbitals belonging to iproc
  !and calculate the norm of the residue

  ! norm of gradient
  gnrm=0.0_dp
  !norm of gradient of unoccupied orbitals
  gnrm_zero=0.0_dp

  confdata%potorder = 0
  confdata%prefac = 0._gp
  !commented out, never used
!   evalmax=orbs%eval(orbs%isorb+1)
!   do iorb=1,orbs%norbp
!     evalmax=max(orbs%eval(orbs%isorb+iorb),evalmax)
!   enddo
!   call MPI_ALLREDUCE(evalmax,eval_zero,1,mpidtypd,&
!        MPI_MAX,bigdft_mpi%mpi_comm,ierr)


  if (orbs%norbp >0) ikpt=orbs%iokpt(1)
  do iorb=1,orbs%norbp
     !if it is the first orbital or the k-point has changed calculate the max
     if (orbs%iokpt(iorb) /= ikpt .or. iorb == 1) then
        !the eval array contains all the values
        !take the max for all k-points
        !one may think to take the max per k-point
        evalmax=orbs%eval((orbs%iokpt(iorb)-1)*orbs%norb+1)
        do jorb=1,orbs%norb
           evalmax=max(orbs%eval((orbs%iokpt(iorb)-1)*orbs%norb+jorb),evalmax)
        enddo
        eval_zero=evalmax
        ikpt=orbs%iokpt(iorb)
     end if

     !indo=(iorb-1)*nspinor+1
     !loop over the spinorial components
     !k-point values, if present
     kx=orbs%kpts(1,orbs%iokpt(iorb))
     ky=orbs%kpts(2,orbs%iokpt(iorb))
     kz=orbs%kpts(3,orbs%iokpt(iorb))
!       print *, iorb, orbs%kpts(1,orbs%iokpt(iorb)), orbs%kpts(2,orbs%iokpt(iorb)), orbs%kpts(3,orbs%iokpt(iorb))

     !real k-point different from Gamma still not implemented
     if (kx**2+ky**2+kz**2 > 0.0_gp .or. orbs%nspinor==2 ) then
        ncplx=2
     else
        ncplx=1
     end if

     do inds=1,orbs%nspinor,ncplx

        call precondition_ket(ncong, &
             confdata, ncplx, (/ hx, hy, hz /), (/ kx,ky,kz /), lr, &
             orbs%eval(orbs%isorb+iorb), eval_zero, hpsi(1,inds,iorb), scpr, &
             lin_prec_conv_work, lin_prec_work)

        if (orbs%occup(orbs%isorb+iorb) == 0.0_gp) then
           gnrm_zero=gnrm_zero+orbs%kwgts(orbs%iokpt(iorb))*scpr**2
        else
           !write(17,*)'iorb,gnrm',orbs%isorb+iorb,scpr**2
           gnrm=gnrm+orbs%kwgts(orbs%iokpt(iorb))*scpr**2
        end if
        !write(*,*) 'preconditionall: verbosity ',verbose
!           write(*,*)'iorb,gnrm',orbs%isorb+iorb,scpr**2
!     print *,iorb,inds,dot(lr%wfd%nvctr_c+7*lr%wfd%nvctr_f, hpsi(1,inds,iorb),1,hpsi(1,inds,iorb),1)
!     print *,iorb,inds+1,dot(lr%wfd%nvctr_c+7*lr%wfd%nvctr_f, hpsi(1,inds+1,iorb),1,hpsi(1,inds+1,iorb),1)
     end do
  enddo

END SUBROUTINE preconditionall

!!$subroutine preconditioner(ob)
!!$  call f_routine(id='preconditioner')
!!$
!!$  ! Precondition all orbitals belonging to iproc
!!$  !and calculate the norm of the residue
!!$  ! norm of gradient
!!$  gnrm=0.0_dp
!!$  !norm of gradient of unoccupied orbitals
!!$  gnrm_zero=0.0_dp
!!$
!!$  !prepare the arrays for the
!!$  if (verbose >= 3) then
!!$     gnrmp = f_malloc(max(ob%orbs%norbp, 1),id='gnrmp')
!!$  end if
!!$
!!$  it=orbital_basis_iterator(ob)
!!$  do while(ket_next_kpt(it))
!!$     evalmax=ob%orbs%eval((it%ikpt-1)*orbs%norb+1)
!!$     do jorb=2,ob%orbs%norb
!!$        evalmax=max(ob%orbs%eval((it%ikpt-1)*orbs%norb+jorb),evalmax)
!!$     enddo
!!$     eval_zero=evalmax
!!$     do while(ket_next(it,ikpt=it%ikpt))
!!$        !overrid the association of the ket with the hpsi array
!!$        it%phi_wvl=>ob_ket_map(hpsi,it)
!!$        !real k-point different from Gamma still not implemented
!!$        ncplx= merge(2,1,any(it%kpoint /= 0.0_gp) .or. it%nspinor==2)
!!$        gnrm_orb=0.0_wp
!!$        do inds=1,it%nspinor,ncplx
!!$           hpsi_ptr=>ob_subket_ptr(it,inds)
!!$           call precondition_ket(ncong,it%confdata,ncplx,hgrids,it%kpoint,it%lr,&
!!$                it%ob%orbs%eval(it%iorb),eval_zero,hpsi_ptr,scpr,&
!!$                lin_prec_conv_work(it%iorbp),lin_prec_work(it%iorbp))
!!$           if (it%occup == 0.0_gp) then
!!$              gnrm_zero=gnrm_zero+it%kwgt*scpr**2
!!$           else
!!$              gnrm=gnrm+it%kwgt*scpr**2
!!$           end if
!!$           gnrm_orb=gnrm_orb+scpr
!!$        end do
!!$        if (verbose =>3) gnrmp(it%iorbp)=gnrm_orb
!!$     end do
!!$  end do
!!$
!!$  !gather the results of the gnrm per orbital in the case of high verbosity
!!$  if (verbose >= 3) then
!!$     gnrms = f_malloc0(ob%orbs%norb*ob%orbs%nkpts,id='gnrms')
!!$     !prepare displacements arrays
!!$     ncntdsp = f_malloc((/ nproc, 2 /),id='ncntdsp')
!!$     ncntdsp(1,2)=0
!!$     ncntdsp(1,1)=ob%orbs%norb_par(0,0)
!!$     do jproc=1,nproc-1
!!$        ncntdsp(jproc+1,2)=ncntdsp(jproc,2)+ncntdsp(jproc,1)
!!$        ncntdsp(jproc+1,1)=ob%orbs%norb_par(jproc,0)
!!$     end do
!!$     !root mpi task collects the data
!!$     if (bigdft_mpi%nproc > 1) then
!!$        call MPI_GATHERV(gnrmp(1),ob%orbs%norbp,mpidtypw,gnrms(1),ncntdsp(1,1),&
!!$             ncntdsp(1,2),mpidtypw,0,bigdft_mpi%mpi_comm,ierr)
!!$     else
!!$        call f_memcpy(src=gnrmp,dest=gnrms)
!!$     end if
!!$
!!$     !write the values per orbitals
!!$     if (bigdft_mpi%iproc ==0) call write_gnrms(ob%orbs%nkpts,ob%orbs%norb,gnrms)
!!$
!!$     call f_free(ncntdsp)
!!$     call f_free(gnrms)
!!$     call f_free(gnrmp)
!!$  end if
!!$
!!$
!!$  call f_release_routine()
!!$
!!$end subroutine preconditioner

!> Generalized for the Linearscaling code
subroutine preconditionall2(iproc,nproc,orbs,Lzd,hx,hy,hz,ncong,npsidim,hpsi,confdatarr,gnrm,gnrm_zero,&
                            linear_precond_convol_workarrays, linear_precond_workarrays)
  use module_precisions
  use module_types
  use Poisson_Solver, except_dp => dp, except_gp => gp
  use module_bigdft_output
  use locregs
  use locreg_operations
  use at_domain, only: domain_geocode
  use module_bigdft_mpi
  use module_bigdft_errors
  use module_bigdft_arrays
  use module_bigdft_profiling
  use wrapper_linalg
  implicit none
  integer, intent(in) :: iproc,nproc,ncong,npsidim
  real(gp), intent(in) :: hx,hy,hz
  type(local_zone_descriptors), intent(in) :: Lzd
  type(orbitals_data), intent(in) :: orbs
  real(dp), intent(out) :: gnrm,gnrm_zero
  real(wp), dimension(npsidim), intent(inout) :: hpsi
  type(confpot_data), dimension(orbs%norbp), intent(in) :: confdatarr !< used in the linear scaling but also for the cubic case
  type(workarrays_quartic_convolutions),dimension(orbs%norbp),intent(inout),optional :: linear_precond_convol_workarrays !< convolution workarrays for the linear case
  type(workarr_precond),dimension(orbs%norbp),intent(inout),optional :: linear_precond_workarrays !< workarrays for the linear case
  !local variables
  character(len=*), parameter :: subname='preconditionall2'
  integer :: iorb,inds,ncplx,ikpt,jorb,ist,ilr,ierr,jproc
  real(wp) :: cprecr,scpr,evalmax,eval_zero,gnrm_orb
  real(gp) :: kx,ky,kz
!!$  integer :: i_stat,i_all,ispinor,nbox
!!$  logical, parameter :: newp=.true.
!!$  real(gp) :: eh_fake,monop
!!$  real(wp), dimension(:), allocatable :: hpsir
!!$  type(coulomb_operator) :: G_Helmholtz
!!$  real(wp), dimension(:,:), allocatable :: gnrm_per_orb
!!$  type(workarr_sumrho) :: w
  integer, dimension(:,:), allocatable :: ncntdsp
  real(wp), dimension(:), allocatable :: gnrms,gnrmp
  type(workarrays_quartic_convolutions) :: w_convol
  type(workarr_precond) :: w_precond
  !debug
!!$  type(atoms_data) atoms_fake
!!$  integer :: iter=0
!!$  iter=iter+1

  call f_routine(id=subname)

  ! Preconditions all orbitals belonging to iproc
  !and calculate the norm of the residue
  ! norm of gradient
  gnrm=0.0_dp
  !norm of gradient of unoccupied orbitals
  gnrm_zero=0.0_dp

  !prepare the arrays for the
  if (get_verbose_level() >= 3) then
     gnrmp = f_malloc(max(orbs%norbp, 1),id='gnrmp')
  end if

!!$  if (newp) then
!!$     ilr=1
!!$     hpsir=f_malloc(Lzd%Llr(ilr)%d%n1i*Lzd%Llr(ilr)%d%n2i*Lzd%Llr(ilr)%d%n3i,id='hpsir',routine_id=subname)
!!$     call initialize_work_arrays_sumrho(1,Lzd%Llr(ilr),.true.,w)
!!$  end if
  !if (iproc.eq. 0 .and. verbose.ge.3) write(*,*) ' '
  ist = 0
  if (orbs%norbp >0) ikpt=orbs%iokpt(1)
  do iorb=1,orbs%norbp
     ilr = orbs%inwhichlocreg(iorb+orbs%isorb)
     !if it is the first orbital or the k-point has changed calculate the max
     if (orbs%iokpt(iorb) /= ikpt .or. iorb == 1) then
        !the eval array contains all the values
        !take the max for all k-points
        !one may think to take the max per k-point
        evalmax=orbs%eval((orbs%iokpt(iorb)-1)*orbs%norb+1)
        do jorb=1,orbs%norb
           evalmax=max(orbs%eval((orbs%iokpt(iorb)-1)*orbs%norb+jorb),evalmax)
        enddo
        eval_zero=evalmax
        ikpt=orbs%iokpt(iorb)
     end if
     !print *,'iorb,eval,evalmax',iorb+orbs%isorb,orbs%eval(iorb+orbs%isorb),eval_zero
     !indo=(iorb-1)*nspinor+1
     !loop over the spinorial components
     !k-point values, if present
     kx=orbs%kpts(1,orbs%iokpt(iorb))
     ky=orbs%kpts(2,orbs%iokpt(iorb))
     kz=orbs%kpts(3,orbs%iokpt(iorb))
!       print *, iorb, orbs%kpts(1,orbs%iokpt(iorb)), orbs%kpts(2,orbs%iokpt(iorb)), orbs%kpts(3,orbs%iokpt(iorb))

     !real k-point different from Gamma still not implemented
     if (kx**2+ky**2+kz**2 > 0.0_gp .or. orbs%nspinor==2) then
        ncplx=2
     else
        ncplx=1
     end if

     gnrm_orb=0.0_wp
     do inds=1,orbs%nspinor,ncplx

!!$        print *,'plotting gradient iorb,initial',iorb,&
!!$             nrm2(Lzd%Llr(ilr)%wfd%nvctr_c+7*Lzd%Llr(ilr)%wfd%nvctr_f,hpsi(1+ist),1)
!!$        atoms_fake=atoms_null()
!!$        atoms_fake%nat=0
!!$        atoms_fake%geocode='F'
!!$        atoms_fake%alat1=Lzd%hgrids(1)*(Lzd%Llr(ilr)%d%n1+1)
!!$        atoms_fake%alat2=Lzd%hgrids(2)*(Lzd%Llr(ilr)%d%n2+1)
!!$        atoms_fake%alat3=Lzd%hgrids(3)*(Lzd%Llr(ilr)%d%n3+1)
!!$        call plot_wf('G'//trim(adjustl(yaml_toa(iorb)))//'-'//trim(adjustl(yaml_toa(iter))),1,&
!!$             atoms_fake,1.0_gp,Lzd%llr(ilr),&
!!$             Lzd%hgrids(1),Lzd%hgrids(2),Lzd%hgrids(3),(/0.0_gp,0.0_gp,0.0_gp/),hpsi(1+ist))
!!$
!!$        if (newp) then
!!$           call cprecr_from_eval(Lzd%Llr(ilr)%geocode,eval_zero,orbs%eval(orbs%isorb+iorb),cprecr)
!!$           !alternative preconditioner
!!$           call vscal(ncplx*(Lzd%Llr(ilr)%wfd%nvctr_c+7*Lzd%Llr(ilr)%wfd%nvctr_f),0.5_gp/pi_param,hpsi(1+ist),1)
!!$           call daub_to_isf(Lzd%Llr(ilr),w,hpsi(1+ist),hpsir(1))
!!$           !sequential kernel
!!$           G_Helmholtz=pkernel_init(.false.,0,1,0,Lzd%Llr(ilr)%geocode,&
!!$                (/Lzd%Llr(ilr)%d%n1i,Lzd%Llr(ilr)%d%n2i,Lzd%Llr(ilr)%d%n3i/),&
!!$                0.5_gp*Lzd%hgrids,16,mu0_screening=sqrt(2.0_gp*abs(cprecr)))
!!$
!!$           call pkernel_set(G_Helmholtz,.true.)
!!$
!!$           !apply it to the gradient to smooth it
!!$           call H_potential('D',G_Helmholtz,hpsir(1),hpsir(1),eh_fake,0.d0,.false.)
!!$           call pkernel_free(G_Helmholtz,subname)
!!$           !convert the gradient back to the locreg
!!$           call isf_to_daub(Lzd%Llr(ilr),w,hpsir(1),hpsi(1+ist))
!!$
!!$        end if

        if(confdatarr(iorb)%prefac > 0.0_gp .or. confdatarr(iorb)%potorder > 0)then
           if (.not.present(linear_precond_convol_workarrays)) then
              call f_err_throw("linear_precond_convol_workarrays must be present when calling the linear preconditioner", &
                   err_name='BIGDFT_RUNTIME_ERROR')
           end if
           if (.not.present(linear_precond_workarrays)) then
              call f_err_throw("linear_precond_workarrays must be present when calling the linear preconditioner", &
                   err_name='BIGDFT_RUNTIME_ERROR')
           end if
           w_convol = linear_precond_convol_workarrays(iorb)
           w_precond = linear_precond_workarrays(iorb)
        end if
        call precondition_ket(ncong, confdatarr(iorb), ncplx, &
             (/ hx, hy, hz /), (/ kx,ky,kz /), Lzd%Llr(ilr), &
             orbs%eval(orbs%isorb+iorb), eval_zero, hpsi(1+ist), scpr, &
             w_convol, w_precond)

        if (orbs%occup(orbs%isorb+iorb) == 0.0_gp) then
           gnrm_zero=gnrm_zero+orbs%kwgts(orbs%iokpt(iorb))*scpr**2
        else
           !write(*,*)'iorb,gnrm',orbs%isorb+iorb,scpr**2,ilr
           gnrm=gnrm+orbs%kwgts(orbs%iokpt(iorb))*scpr**2
        end if
        if (get_verbose_level() >= 3) then
           gnrm_orb=gnrm_orb+scpr
           if (inds+ncplx-1==orbs%nspinor) gnrmp(iorb)=gnrm_orb
           !write(*,*) 'iorb,gnrm,ilr',orbs%isorb+iorb,scpr,ilr,gnrm_orb,iproc
        end if

!!$        print *,'iorb,newgradient',iorb,scpr,nrm2(ncplx*(Lzd%Llr(ilr)%wfd%nvctr_c+7*Lzd%Llr(ilr)%wfd%nvctr_f),hpsi(1+ist),1),eh_fake
!!$        call plot_wf('GPnew'//trim(adjustl(yaml_toa(iorb)))//'-'//trim(adjustl(yaml_toa(iter))),1,&
!!$             atoms_fake,1.0_gp,Lzd%llr(ilr),&
!!$             Lzd%hgrids(1),Lzd%hgrids(2),Lzd%hgrids(3),(/0.0_gp,0.0_gp,0.0_gp/),hpsi(1+ist))


       ist = ist + (Lzd%Llr(ilr)%wfd%nvctr_c+7*Lzd%Llr(ilr)%wfd%nvctr_f)*ncplx
!     print *,iorb,inds,dot(lr%wfd%nvctr_c+7*lr%wfd%nvctr_f, hpsi(1,inds,iorb),1,hpsi(1,inds,iorb),1)
!     print *,iorb,inds+1,dot(lr%wfd%nvctr_c+7*lr%wfd%nvctr_f, hpsi(1,inds+1,iorb),1,hpsi(1,inds+1,iorb),1)
    end do

 enddo

!!$ if (newp) then
!!$    call deallocate_work_arrays_sumrho(w)
!!$    call f_free(hpsir)
!!$ end if
  !gather the results of the gnrm per orbital in the case of high verbosity
  if (get_verbose_level() >= 3) then
     gnrms = f_malloc0(orbs%norb*orbs%nkpts,id='gnrms')
     !prepare displacements arrays
     ncntdsp = f_malloc((/ nproc, 2 /),id='ncntdsp')
     ncntdsp(1,2)=0
     ncntdsp(1,1)=orbs%norb_par(0,0)
     do jproc=1,nproc-1
        ncntdsp(jproc+1,2)=ncntdsp(jproc,2)+ncntdsp(jproc,1)
        ncntdsp(jproc+1,1)=orbs%norb_par(jproc,0)
     end do
     !call f_zero(orbs%norb*orbs%nkpts,gnrms(1))
     !root mpi task collects the data
     if (nproc > 1) then
        call MPI_GATHERV(gnrmp(1),orbs%norbp,mpidtypw,gnrms(1),ncntdsp(1,1),&
             ncntdsp(1,2),mpidtypw,0,bigdft_mpi%mpi_comm,ierr)
     else
        call vcopy(orbs%norb*orbs%nkpts,gnrmp(1),1,gnrms(1),1)
     end if

     !if (iproc ==0) print *,'ciao',gnrmp,orbs%nspinor

     !write the values per orbitals
     if (iproc ==0) call write_gnrms(orbs%nkpts,orbs%norb,gnrms)

     call f_free(ncntdsp)
     call f_free(gnrms)
     call f_free(gnrmp)
  end if

  call f_release_routine()

END SUBROUTINE preconditionall2
