module module_bigdft_profiling
  use dynamic_memory, only: f_routine, f_release_routine, f_malloc_dump_status
  use time_profiling, only: TIMING_UNINITIALIZED, f_timing_stop, f_timing_reset, f_timing

  implicit none

  public

  !> Timing categories
  character(len=*), parameter, private :: tgrp_pot='Potential'
  integer, save :: TCAT_EXCHANGECORR=TIMING_UNINITIALIZED
  integer, parameter, private :: ncls_max=6,ncat_bigdft=155   ! define timimg categories and classes
  character(len=*), parameter, private :: tgrp_paw='PAW'
  character(len=*), parameter, private :: tgrp_io='IO'
  integer, save :: TCAT_LIBPAW    = TIMING_UNINITIALIZED
  integer, save :: TCAT_PAW_DIJ   = TIMING_UNINITIALIZED
  integer, save :: TCAT_PAW_RHOIJ = TIMING_UNINITIALIZED
  integer, save :: TCAT_IO_MULTIPOLES = TIMING_UNINITIALIZED
  character(len=14), dimension(ncls_max), parameter, private :: clss = (/ &
       'Communications'    ,  &
       'Convolutions  '    ,  &
       'Linear Algebra'    ,  &
       'Other         '    ,  &
       !      'Potential     '    ,  &
       'Initialization'    ,  &
       'Finalization  '    /)
  character(len=14), dimension(3,ncat_bigdft), parameter, private :: cats = reshape((/ &
                                !       Name           Class       Operation Kind
       'ReformatWaves ','Initialization' ,'Small Convol  ' ,  &  !< Reformatting of input waves
       'CrtDescriptors','Initialization' ,'RMA Pattern   ' ,  &  !< Calculation of descriptor arrays
       'CrtLocPot     ','Initialization' ,'Miscellaneous ' ,  &  !< Calculation of local potential
       'CrtProjectors ','Initialization' ,'RMA Pattern   ' ,  &  !< Calculation of projectors
       'CrtPcProjects ','Initialization' ,'RMA Pattern   ' ,  &  !< Calculation of preconditioning projectors
       'CrtPawProjects','Initialization' ,'RMA Pattern   ' ,  &  !< Calculation of abscalc-pawprojectors
       'ApplyLocPotKin','Convolutions  ' ,'OpenCL ported ' ,  &  !< Application of PSP, kinetic energy
       'ApplyProj     ','Other         ' ,'RMA pattern   ' ,  &  !< Application of nonlocal PSP
       'Precondition  ','Convolutions  ' ,'OpenCL ported ' ,  &  !< Precondtioning
       'Rho_comput    ','Convolutions  ' ,'OpenCL ported ' ,  &  !< Calculation of charge density (sumrho) computation
       'Rho_commun    ','Communications' ,'AllReduce grid' ,  &  !< Calculation of charge density (sumrho) communication
       'Pot_commun    ','Communications' ,'AllGathrv grid' ,  &  !< Communication of potential
       'Pot_comm start','Communications' ,'MPI_types/_get' ,  &  !< Communication of potential
       'Un-TransSwitch','Other         ' ,'RMA pattern   ' ,  &  !< Transposition of wavefunction, computation
       'Un-TransComm  ','Communications' ,'ALLtoALLV     ' ,  &  !< Transposition of wavefunction, communication
       'GramS_comput  ','Linear Algebra' ,'DPOTRF        ' ,  &  !< Gram Schmidt computation
       'GramS_commun  ','Communications' ,'ALLReduce orbs' ,  &  !< Gram Schmidt communication
       'LagrM_comput  ','Linear Algebra' ,'DGEMM         ' ,  &  !< Lagrange Multipliers computation
       'LagrM_commun  ','Communications' ,'ALLReduce orbs' ,  &  !< Lagrange Multipliers communication
       'Diis          ','Other         ' ,'Other         ' ,  &
                                !       'PSolv_comput  ','Potential     ' ,'3D FFT        ' ,  &
                                !       'PSolv_commun  ','Communications' ,'ALLtoALL      ' ,  &
                                !       'PSolvKernel   ','Initialization' ,'Miscellaneous ' ,  &
       !      'Exchangecorr  ','Potential     ' ,'Miscellaneous ' ,  &
       'Forces        ','Finalization  ' ,'Miscellaneous ' ,  &
       'Tail          ','Finalization  ' ,'Miscellaneous ' ,  &
       'Loewdin_comput','Linear Algebra' ,'              ' ,  &
       'Loewdin_commun','Communications' ,'ALLReduce orbs' ,  &
       'Chol_commun   ','Communications' ,'              ' ,  &
       'Chol_comput   ','Linear Algebra' ,'ALLReduce orbs' ,  &
       'GS/Chol_comput','Linear Algebra' ,'              ' ,  &
       'GS/Chol_commun','Communications' ,'ALLReduce orbs' ,  &
       'Input_comput  ','Initialization' ,'Miscellaneous ' ,  &
       'Input_commun  ','Communications' ,'ALLtoALL+Reduc' ,  &
       'Davidson      ','Finalization  ' ,'Complete SCF  ' ,  &
       'check_IG      ','Initialization' ,'Linear Scaling' ,  &
       'constrc_locreg','Initialization' ,'Miscellaneous ' ,  &
       'wavefunction  ','Initialization' ,'Miscellaneous ' ,  &
       'create_nlpspd ','Initialization' ,'RMA pattern   ' ,  &
       'p2pOrtho_post ','Communications' ,'irecv / irsend' ,  &
       'p2pOrtho_wait ','Communications' ,'mpi_waitany   ' ,  &
       'lovrlp_comm   ','Communications' ,'mpi_allgatherv' ,  &
       'lovrlp_comp   ','Linear Algebra' ,'many ddots    ' ,  &
       'lovrlp_compr  ','Other         ' ,'cut out zeros ' ,  &
       'lovrlp_uncompr','Other         ' ,'insert zeros  ' ,  &
       'extract_orbs  ','Other         ' ,'copy to sendb ' ,  &
       'lovrlp^-1/2   ','Linear Algebra' ,'exact or appr ' ,  &
       'lovrlp^-1/2old','Linear Algebra' ,'exact or appr ' ,  &
       'lovrlp^-1/2com','Linear Algebra' ,'exact or appr ' ,  &
       'lovrlp^-1/2par','Linear Algebra' ,'exact or appr ' ,  &
       'build_lincomb ','Linear Algebra' ,'many daxpy    ' ,  &
       'p2pSumrho_wait','Communications' ,'mpi_test/wait ' ,  &
       'sumrho_TMB    ','Other         ' ,'port to GPU?  ' ,  &
       'TMB_kernel    ','Linear Algebra' ,'dgemm         ' ,  &
       'diagonal_seq  ','Linear Algebra' ,'dsygv         ' ,  &
       'lovrlp^-1     ','Linear Algebra' ,'exact or appr ' ,  &
       'lagmat_orthoco','Linear Algebra' ,'dgemm seq/par ' ,  &
       'optimize_DIIS ','Other         ' ,'Other         ' ,  &
       'optimize_SD   ','Other         ' ,'Other         ' ,  &
       'mix_linear    ','Other         ' ,'Other         ' ,  &
       'mix_DIIS      ','Other         ' ,'Other         ' ,  &
       'ig_matric_comm','Communications' ,'mpi p2p       ' ,  &
       'wf_signals    ','Communications' ,'Socket transf.' ,  &
       'energs_signals','Communications' ,'Socket transf.' ,  &
       'rhov_signals  ','Communications' ,'Socket transf.' ,  &
       'init_locregs  ','Initialization' ,'Miscellaneous ' ,  &
       'init_commSumro','Initialization' ,'Miscellaneous ' ,  &
       'init_commPot  ','Initialization' ,'Miscellaneous ' ,  &
       'init_commOrtho','Initialization' ,'Miscellaneous ' ,  &
       'init_inguess  ','Initialization' ,'Miscellaneous ' ,  &
       'init_matrCompr','Initialization' ,'Miscellaneous ' ,  &
       'init_collcomm ','Initialization' ,'Miscellaneous ' ,  &
       'init_collco_sr','Initialization' ,'Miscellaneous ' ,  &
       'init_orbs_lin ','Initialization' ,'Miscellaneous ' ,  &
       'init_repart   ','Initialization' ,'Miscellaneous ' ,  &
       'initMatmulComp','Initialization' ,'Miscellaneous ' ,  &
       'Pot_after_comm','Other         ' ,'global_to_loca' ,  &
                                !       'Init to Zero  ','Other         ' ,'Memset        ' ,  &
       'calc_kernel   ','Other         ' ,'Miscellaneous ' ,  &
       'commun_kernel ','Communications' ,'mpi_allgatherv' ,  &
       'getlocbasinit ','Other         ' ,'Miscellaneous ' ,  &
       'updatelocreg1 ','Other         ' ,'Miscellaneous ' ,  &
       'linscalinit   ','Other         ' ,'Miscellaneous ' ,  &
       'commbasis4dens','Communications' ,'Miscellaneous ' ,  &
       'buildgrad_mcpy','Other         ' ,'Miscellaneous ' ,  &
       'buildgrad_comm','Communications' ,'Allgatherv    ' ,  &
       'allocommsumrho','Communications' ,'Miscellaneous ' ,  &
       'ovrlptransComp','Other         ' ,'Miscellaneous ' ,  &
       'ovrlptransComm','Communications' ,'mpi_allreduce ' ,  &
       'lincombtrans  ','Other         ' ,'Miscellaneous ' ,  &
       'glsynchham1   ','Communications' ,'load balancing' ,  &
       'glsynchham2   ','Communications' ,'load balancing' ,  &
       'gauss_proj    ','Other         ' ,'Miscellaneous ' ,  &
       'sumrho_allred ','Communications' ,'mpiallred     ' ,  &
       'deallocprec   ','Other         ' ,'Miscellaneous ' ,  &
       'large2small   ','Other         ' ,'Miscellaneous ' ,  &
       'small2large   ','Other         ' ,'Miscellaneous ' ,  &
       'renormCoefCom1','Linear Algebra' ,'Miscellaneous ' ,  &
       'renormCoefCom2','Linear Algebra' ,'Miscellaneous ' ,  &
       'renormCoefComm','Communications' ,'Miscellaneous ' ,  &
       'waitAllgatKern','Other         ' ,'Miscellaneous ' ,  &
       'UnBlockPot    ','Other         ' ,'Overlap comms ' ,  &
       'UnBlockDen    ','Other         ' ,'Overlap comms ' ,  &
       'global_local  ','Initialization' ,'Unknown       ' ,  &
                                !       'AllocationProf','Other         ' ,'Allocate arrs ' ,  &
       'dirmin_lagmat1','Linear Algebra' ,'grad calc     ' ,  &
       'dirmin_lagmat2','Linear Algebra' ,'allgatherv    ' ,  &
       'dirmin_dgesv  ','Linear Algebra' ,'dgesv/pdgesv  ' ,  &
       'dirmin_sddiis ','Linear Algebra' ,'Miscellaneous ' ,  &
       'dirmin_allgat ','Linear Algebra' ,'allgatherv    ' ,  &
       'dirmin_sdfit  ','Linear Algebra' ,'allgatherv etc' ,  &
       'chebyshev_comp','Linear Algebra' ,'matmul/matadd ' ,  &
       'chebyshev_comm','Communications' ,'allreduce     ' ,  &
       'chebyshev_coef','Other         ' ,'Miscellaneous ' ,  &
       'FOE_auxiliary ','Other         ' ,'Miscellaneous ' ,  &
       'FOE_init      ','Other         ' ,'Miscellaneous ' ,  &
       'compressd_mcpy','Other         ' ,'Miscellaneous ' ,  &
       'compressd_comm','Communications' ,'Allgatherv    ' ,  &
       'foe_aux_mcpy  ','Other         ' ,'Miscellaneous ' ,  &
       'foe_aux_comm  ','Communications' ,'Allgatherv    ' ,  &
       'norm_trans    ','Other         ' ,'Miscellaneous ' ,  &
       'misc          ','Other         ' ,'Miscellaneous ' ,  &
       'sparse_copy   ','Other         ' ,'Miscellaneous ' ,  &
       'constraineddft','Other         ' ,'Miscellaneous ' ,  &
       'transfer_int  ','Other         ' ,'Miscellaneous ' ,  &
       'Reformatting  ','Initialization' ,'Interpolation ' ,  &
       'restart_wvl   ','Initialization' ,'inguess    rst' ,  &
       'restart_rsp   ','Initialization' ,'inguess    rst' ,  &
       'check_sumrho  ','Initialization' ,'unitary check ' ,  &
       'check_pot     ','Initialization' ,'unitary check ' ,  &
       'ApplyLocPot   ','Convolutions  ' ,'OpenCL ported ' ,  &
       'ApplyLocKin   ','Convolutions  ' ,'OpenCL ported ' ,  &
       'kernel_init   ','Other         ' ,'Fragment calc ' ,  &
       'calc_energy   ','Linear Algebra' ,'allred etc    ' ,  &
       'new_pulay_corr','Other         ' ,'Pulay forces  ' ,  &
       'dev_from_unity','Other         ' ,'Miscellaneous ' ,  &
       'ks_residue    ','Linear Algebra' ,'Miscellaneous ' ,  &
       'weightanalysis','Linear Algebra' ,'Fragment calc ' ,  &
       'tmbrestart    ','Initialization' ,'Miscellaneous ' ,  &
       'readtmbfiles  ','Initialization' ,'Miscellaneous ' ,  &
       'readisffiles  ','Initialization' ,'Miscellaneous ' ,  &
       'purify_kernel ','Linear Algebra' ,'dgemm         ' ,  &
       'potential_dims','Other         ' ,'auxiliary     ' ,  &
       'sparse_matmul ','Linear Algebra' ,'self-made     ' ,  &
       'transform_matr','Other         ' ,'small to large' ,  &
       'calctrace_comp','Other         ' ,'Miscellaneous ' ,  &
       'calctrace_comm','Communications' ,'allreduce     ' ,  &
       'determinespars','Other         ' ,'Miscellaneous ' ,  &
       'inittaskgroup ','Other         ' ,'Miscellaneous ' ,  &
       'write_matrices','Other         ' ,'dump to disk  ' ,  &
       'transformspars','Other         ' ,'Miscellaneous ' ,  &
       'matrix_extents','Other         ' ,'Miscellaneous ' ,  &
       'lin_inputguess','Other         ' ,'Miscellaneous ' ,  &
       'ionic_energy  ','Other         ' ,'Miscellaneous ' ,  &
       'dgemm_parallel','Linear Algebra' ,'(Sca)LAPACK   ' ,  &
       'dsyev_parallel','Linear Algebra' ,'(Sca)LAPACK   ' ,  &
       'dsygv_parallel','Linear Algebra' ,'(Sca)LAPACK   ' ,  &
       'dgesv_parallel','Linear Algebra' ,'(Sca)LAPACK   ' ,  &
       'dpotrf_paralle','Linear Algebra' ,'(Sca)LAPACK   ' ,  &
       'dpotri_paralle','Linear Algebra' ,'(Sca)LAPACK   ' ,  &
       'plotting      ','Other         ' ,'Miscellaneous ' /),(/3,ncat_bigdft/))
  integer, dimension(ncat_bigdft), private, save :: cat_ids !< id of the categories to be converted

contains

  !> initialize the timing categories for BigDFT runs.
  !! It is of course assumed that f_lib_initialize has already been called
  subroutine bigdft_init_timing_categories()
    use Poisson_Solver, only: PS_initialize_timing_categories
    use sparsematrix_base
    use time_profiling
    implicit none
    !local variables
    integer :: icls,icat

    !initialize categories for the Poisson Solver
    call PS_initialize_timing_categories()

    !initialize groups
    call f_timing_category_group(tgrp_pot,'Operations for local potential construction (mainly XC)')
    call f_timing_category_group(tgrp_paw,'Operations done for PAW treatment')
    call f_timing_category_group(tgrp_io,'Operations related to I/O')

    do icls=2,ncls_max
       call f_timing_category_group(trim(clss(icls)),'Empty description for the moment')
    end do

    !define the timing categories for exchange and correlation
    call f_timing_category('Exchange-Correlation',tgrp_pot,&
         'Operations needed to construct local XC potential',&
         TCAT_EXCHANGECORR)

    !define the categories for PAW
    call f_timing_category('libPAW',tgrp_paw, 'Operations done inside libPAW',&
         TCAT_LIBPAW)
    call f_timing_category('paw dij',tgrp_paw, 'Computation of PAW dij terms',&
         TCAT_PAW_DIJ)
    call f_timing_category('paw rhoij',tgrp_paw, 'Computation of PAW rhoij terms',&
         TCAT_PAW_RHOIJ)

    ! define the categories for the I/O
    call f_timing_category('dump multipoles', tgrp_io, 'Output of the multipoles to the standard output ',&
         TCAT_IO_MULTIPOLES)


    !! little by little, these categories should be transformed in the
    !! new scheme dictated by f_timing API in time_profiling module of f_lib.

    !initialize categories
    do icat=1,ncat_bigdft
       call f_timing_category(trim(cats(1,icat)),trim(cats(2,icat)),trim(cats(3,icat)),&
            cat_ids(icat))
    end do

    ! Initialize sparse matrix timings
    call sparsematrix_initialize_timing_categories()

  end subroutine bigdft_init_timing_categories
  
  !> Routine to convert timing categories from the old scheme to the API of f_lib
  !! as soon as the timing categories are identified with their ID, this routine should disappear
  subroutine find_category(category,cat_id)
    use dictionaries
    use time_profiling
    implicit none
    character(len=*), intent(in) :: category
    integer, intent(out) :: cat_id !< id of the found category
    !local variables
    integer :: i
    !controls if the category exists
    cat_id=0
    do i=1,ncat_bigdft
       if (trim(category) == trim(cats(1,i))) then
          cat_id=cat_ids(i)
          exit
       endif
    enddo
    if (cat_id==0) then
       call f_err_throw('Timing routine error,'//&
            ' requested category '//trim(category)//' has not been found',&
            err_id=TIMING_INVALID)
!!$     if (bigdft_mpi%iproc==0) &
!!$          call yaml_warning('Requested timing category '//trim(category)//&
!!$          ' has not been found')
       cat_id=TIMING_UNINITIALIZED
    end if
  end subroutine find_category

end module module_bigdft_profiling

!> old timing routine, should disappear as soon as the f_timing routine is called
subroutine timing(comm,category,action)
  use dictionaries, only: max_field_length,f_err_raise
  use yaml_strings, only: yaml_toa
  use module_bigdft_profiling
  use time_profiling
  use wrapper_MPI
  implicit none
  !Variables
  integer, intent(in) :: comm
  character(len=*), intent(in) :: category
  character(len=2), intent(in) :: action  
  !Local variables
  integer :: cat_id,ierr
  character(len=max_field_length) :: cattmp
  external :: gather_timings
  !this is to ensure that timing routines have been properly called
  !call check_initialization()

  !modification of the timing to see if it works
  select case(action)
  case('PR')
     !here iproc is the communicator
     call f_timing_checkpoint(ctr_name=category,mpi_comm=comm,nproc=mpisize(comm),&
          gather_routine=gather_timings)
  case default
     !find category in the old scheme
     call find_category(category,cat_id)

     if (cat_id /= TIMING_UNINITIALIZED) then
        call get_category_name(cat_id,cattmp)
        if (f_err_raise(trim(cattmp)/=trim(category),'Error in category '//&
             trim(yaml_toa(cat_id))//' (name='//trim(category)//' ), found '//&
             trim(cattmp)//' instead',err_name='BIGDFT_RUNTIME_ERROR')) return
        
        call f_timing(cat_id,action)
     end if
  end select

END SUBROUTINE timing
