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

! **************************************************************************************************
!> \brief   DBCSR data operations
!> \author  Urban Borstnik
!> \date    2010-02-18
!> \version 0.9
!>
!> <b>Modification history:</b>
!> - 2010-02-18 Moved from dbcsr_util
! **************************************************************************************************
MODULE dbcsr_data_operations

   USE dbcsr_block_operations,          ONLY: dbcsr_block_transpose,&
                                              dbcsr_data_copy,&
                                              dbcsr_data_set
   USE dbcsr_data_methods,              ONLY: dbcsr_data_get_size,&
                                              dbcsr_data_get_size_referenced,&
                                              dbcsr_data_hold,&
                                              dbcsr_data_release,&
                                              dbcsr_data_set_size_referenced,&
                                              dbcsr_get_data
   USE dbcsr_types,                     ONLY: dbcsr_data_obj,&
                                              dbcsr_type,&
                                              dbcsr_type_complex_4,&
                                              dbcsr_type_complex_8,&
                                              dbcsr_type_real_4,&
                                              dbcsr_type_real_8
   USE dbcsr_util,                      ONLY: sgn
   USE kinds,                           ONLY: real_4,&
                                              real_8
#include "../../base/base_uses.f90"

!$ USE OMP_LIB, ONLY: omp_get_max_threads, omp_get_thread_num, omp_get_num_threads

   IMPLICIT NONE

   PRIVATE

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

   PUBLIC :: dbcsr_data_copyall, dbcsr_data_convert, &
             dbcsr_copy_sort_data, &
             dbcsr_sort_data
   PUBLIC :: dbcsr_switch_data_area

CONTAINS

! **************************************************************************************************
!> \brief Sets the data area of a matrix
!> \param[in,out] matrix     matrix for which to set the data area
!> \param[in] data_area      data area to set
!> \param[out] previous_data_area  (optional) previous data area
! **************************************************************************************************
   SUBROUTINE dbcsr_switch_data_area(matrix, data_area, previous_data_area)
      TYPE(dbcsr_type), INTENT(INOUT)                    :: matrix
      TYPE(dbcsr_data_obj), INTENT(IN)                   :: data_area
      TYPE(dbcsr_data_obj), INTENT(OUT), OPTIONAL        :: previous_data_area

      CHARACTER(len=*), PARAMETER :: routineN = 'dbcsr_switch_data_area', &
         routineP = moduleN//':'//routineN

      INTEGER                                            :: handle

!   ---------------------------------------------------------------------------

      CALL timeset(routineN, handle)
      IF (PRESENT(previous_data_area)) THEN
         previous_data_area = matrix%data_area
      ELSE
         CALL dbcsr_data_release(matrix%data_area)
      ENDIF
      matrix%data_area = data_area
      CALL dbcsr_data_hold(matrix%data_area)
      CALL timestop(handle)
   END SUBROUTINE dbcsr_switch_data_area

! **************************************************************************************************
!> \brief Copies a data area, deep by default.
!> \param[in,out] target_area     target data area
!> \param[in]     source_area     source data area
!> \param[in]     shallow         shallow copy (default is deep)
! **************************************************************************************************
   SUBROUTINE dbcsr_data_copyall(target_area, source_area, shallow)
      TYPE(dbcsr_data_obj), INTENT(INOUT)                :: target_area
      TYPE(dbcsr_data_obj), INTENT(IN)                   :: source_area
      LOGICAL, INTENT(IN), OPTIONAL                      :: shallow

      CHARACTER(len=*), PARAMETER :: routineN = 'dbcsr_data_copyall', &
         routineP = moduleN//':'//routineN

      INTEGER                                            :: i, n
      LOGICAL                                            :: shallow_copy

!   ---------------------------------------------------------------------------

!      CALL timeset(routineN, error_handle)
      IF (.NOT. ASSOCIATED(source_area%d)) &
         CPABORT("Attempt to copy unassigned data")
      IF (source_area%d%refcount .LE. 0) &
         CPWARN("Attempt to copy unheld data")
      shallow_copy = .FALSE.
      IF (PRESENT(shallow)) shallow_copy = shallow
      IF (shallow_copy) THEN
         target_area = source_area
         CALL dbcsr_data_hold(target_area)
      ELSE
         IF (.NOT. ASSOCIATED(target_area%d)) &
            CPABORT("Target area does not exist.")
         CALL dbcsr_data_set_size_referenced(target_area, &
                                             dbcsr_data_get_size_referenced(source_area))
         n = dbcsr_data_get_size_referenced(source_area)
         SELECT CASE (target_area%d%data_type)
         CASE (dbcsr_type_real_4)
!$OMP           PARALLEL DO DEFAULT(NONE) PRIVATE(I) SHARED(target_area,source_area,n)
            DO i = 1, n
               target_area%d%r_sp(i) = source_area%d%r_sp(i)
            ENDDO
         CASE (dbcsr_type_real_8)
!$OMP           PARALLEL DO DEFAULT(NONE) PRIVATE(I) SHARED(target_area,source_area,n)
            DO i = 1, n
               target_area%d%r_dp(i) = source_area%d%r_dp(i)
            ENDDO
         CASE (dbcsr_type_complex_4)
!$OMP           PARALLEL DO DEFAULT(NONE) PRIVATE(I) SHARED(target_area,source_area,n)
            DO i = 1, n
               target_area%d%c_sp(i) = source_area%d%c_sp(i)
            ENDDO
         CASE (dbcsr_type_complex_8)
!$OMP           PARALLEL DO DEFAULT(NONE) PRIVATE(I) SHARED(target_area,source_area,n)
            DO i = 1, n
               target_area%d%c_dp(i) = source_area%d%c_dp(i)
            ENDDO
         CASE default
            CPABORT("Invalid data type")
         END SELECT
      ENDIF
!      CALL timestop(error_handle)
   END SUBROUTINE dbcsr_data_copyall

! **************************************************************************************************
!> \brief Copies a data area, converting data type
!> \param[inout] target_area     target data area
!> \param[in]    source_area     source data area
!> \param[in]    drop_real       (optional) drops real part of complex
!>                               numbers instead of the imaginary part; default
!>                               is false
!> \param[in]    multiply_by_i   (optional) converts real to complex by placing
!>                               into imaginary instead of real part
! **************************************************************************************************
   SUBROUTINE dbcsr_data_convert(target_area, source_area, drop_real, &
                                 multiply_by_i)
      TYPE(dbcsr_data_obj), INTENT(INOUT)                :: target_area
      TYPE(dbcsr_data_obj), INTENT(IN)                   :: source_area
      LOGICAL, INTENT(IN), OPTIONAL                      :: drop_real, multiply_by_i

      CHARACTER(len=*), PARAMETER :: routineN = 'dbcsr_data_convert', &
         routineP = moduleN//':'//routineN

      COMPLEX(KIND=real_4), DIMENSION(:), POINTER        :: s_data_c, t_data_c
      COMPLEX(KIND=real_8), DIMENSION(:), POINTER        :: s_data_z, t_data_z
      INTEGER                                            :: n, ns, nt
      LOGICAL                                            :: keep_real, noimult
      REAL(KIND=real_4), DIMENSION(:), POINTER           :: s_data_r, t_data_r
      REAL(KIND=real_8), DIMENSION(:), POINTER           :: s_data_d, t_data_d

!   ---------------------------------------------------------------------------

      IF (.NOT. ASSOCIATED(source_area%d)) &
         CPWARN("Attempt to copy unassigned data")
      IF (source_area%d%refcount .LE. 0) &
         CPWARN("Attempt to copy unheld data")
      IF (.NOT. ASSOCIATED(source_area%d)) THEN
         RETURN
      ENDIF
      keep_real = .TRUE.
      IF (PRESENT(drop_real)) keep_real = .NOT. drop_real
      noimult = .TRUE.
      IF (PRESENT(multiply_by_i)) noimult = .NOT. multiply_by_i
      ns = dbcsr_data_get_size_referenced(source_area)
      nt = dbcsr_data_get_size_referenced(target_area)
      n = MIN(ns, nt)
      IF (n .GT. 0) THEN
         SELECT CASE (source_area%d%data_type)
         CASE (dbcsr_type_real_8)
            CALL dbcsr_get_data(source_area, s_data_d)
            SELECT CASE (target_area%d%data_type)
            CASE (dbcsr_type_real_8)
               CALL dbcsr_get_data(target_area, t_data_d)
               t_data_d(1:n) = s_data_d(1:n)
            CASE (dbcsr_type_real_4)
               CALL dbcsr_get_data(target_area, t_data_r)
               t_data_r(1:n) = REAL(s_data_d(1:n), KIND=real_4)
            CASE (dbcsr_type_complex_8)
               CALL dbcsr_get_data(target_area, t_data_z)
               IF (noimult) THEN
                  t_data_z(1:n) = CMPLX(s_data_d(1:n), KIND=real_8)
               ELSE
                  t_data_z(1:n) = CMPLX(0.0, s_data_d(1:n), KIND=real_8)
               ENDIF
            CASE (dbcsr_type_complex_4)
               CALL dbcsr_get_data(target_area, t_data_c)
               IF (noimult) THEN
                  t_data_c(1:n) = CMPLX(s_data_d(1:n), KIND=real_4)
               ELSE
                  t_data_c(1:n) = CMPLX(0.0, s_data_d(1:n), KIND=real_4)
               ENDIF
            CASE default
               CPABORT("Invalid data type")
            END SELECT
         CASE (dbcsr_type_real_4)
            CALL dbcsr_get_data(source_area, s_data_r)
            SELECT CASE (target_area%d%data_type)
            CASE (dbcsr_type_real_8)
               CALL dbcsr_get_data(target_area, t_data_d)
               t_data_d(1:n) = REAL(s_data_r(1:n), KIND=real_8)
            CASE (dbcsr_type_real_4)
               CALL dbcsr_get_data(target_area, t_data_r)
               t_data_r(1:n) = s_data_r(1:n)
            CASE (dbcsr_type_complex_8)
               CALL dbcsr_get_data(target_area, t_data_z)
               IF (noimult) THEN
                  t_data_z(1:n) = CMPLX(s_data_r(1:n), KIND=real_8)
               ELSE
                  t_data_z(1:n) = CMPLX(0.0, s_data_r(1:n), KIND=real_8)
               ENDIF
            CASE (dbcsr_type_complex_4)
               CALL dbcsr_get_data(target_area, t_data_c)
               IF (noimult) THEN
                  t_data_c(1:n) = CMPLX(s_data_r(1:n), KIND=real_4)
               ELSE
                  t_data_c(1:n) = CMPLX(0.0, s_data_r(1:n), KIND=real_4)
               ENDIF
            CASE default
               CPABORT("Invalid data type")
            END SELECT
         CASE (dbcsr_type_complex_8)
            CALL dbcsr_get_data(source_area, s_data_z)
            SELECT CASE (target_area%d%data_type)
            CASE (dbcsr_type_real_8)
               CALL dbcsr_get_data(target_area, t_data_d)
               IF (keep_real) THEN
                  t_data_d(1:n) = REAL(s_data_z(1:n), KIND=real_8)
               ELSE
                  t_data_d(1:n) = AIMAG(s_data_z(1:n))
               ENDIF
            CASE (dbcsr_type_real_4)
               CALL dbcsr_get_data(target_area, t_data_r)
               IF (keep_real) THEN
                  t_data_r(1:n) = REAL(s_data_z(1:n), KIND=real_4)
               ELSE
                  t_data_r(1:n) = REAL(AIMAG(s_data_z(1:n)), KIND=real_4)
               ENDIF
            CASE (dbcsr_type_complex_8)
               CALL dbcsr_get_data(target_area, t_data_z)
               t_data_z(1:n) = s_data_z(1:n)
            CASE (dbcsr_type_complex_4)
               CALL dbcsr_get_data(target_area, t_data_c)
               t_data_c(1:n) = CMPLX(s_data_z(1:n), KIND=real_4)
            CASE default
               CPABORT("Invalid data type")
            END SELECT
         CASE (dbcsr_type_complex_4)
            CALL dbcsr_get_data(source_area, s_data_c)
            SELECT CASE (target_area%d%data_type)
            CASE (dbcsr_type_real_8)
               CALL dbcsr_get_data(target_area, t_data_d)
               IF (keep_real) THEN
                  t_data_d(1:n) = REAL(s_data_c(1:n), KIND=real_8)
               ELSE
                  t_data_d(1:n) = REAL(AIMAG(s_data_c(1:n)), KIND=real_8)
               ENDIF
            CASE (dbcsr_type_real_4)
               CALL dbcsr_get_data(target_area, t_data_r)
               IF (keep_real) THEN
                  t_data_r(1:n) = REAL(s_data_c(1:n), KIND=real_4)
               ELSE
                  t_data_r(1:n) = AIMAG(s_data_c(1:n))
               ENDIF
            CASE (dbcsr_type_complex_8)
               CALL dbcsr_get_data(target_area, t_data_z)
               t_data_z(1:n) = CMPLX(s_data_c(1:n), KIND=real_8)
            CASE (dbcsr_type_complex_4)
               CALL dbcsr_get_data(target_area, t_data_c)
               t_data_c(1:n) = s_data_c(1:n)
            CASE default
               CPABORT("Invalid data type")
            END SELECT
         CASE default
            CPABORT("Invalid data type")
         END SELECT
      ENDIF
   END SUBROUTINE dbcsr_data_convert

! **************************************************************************************************
!> \brief Sorts the data in a matrix so that the data blocks follow
!>        sequentially and does various transposing options.
!>
!> As opposed to dbcsr_sort_data, this routine calculates block sizes
!> \param[out] blk_p          re-arranged block pointers reflecting the new
!>                            data order
!> \param[in] old_blk_p       current block pointers
!> \param[in] row_p index
!> \param[in] col_i index
!> \param[in] rbs sizes of the blocked rows
!> \param[in] cbs sizes of the blocked columns
!> \param[out] dst            sorted data
!> \param[in] src             existing unordered data
!> \param[in] mark_transposed           mark data as transposed by negating the
!>                                      blk_p index entries
!> \param[in] transpose_blocks          transpose data blocks
! **************************************************************************************************
   SUBROUTINE dbcsr_copy_sort_data(blk_p, old_blk_p, row_p, col_i, rbs, cbs, &
                                   dst, src, mark_transposed, transpose_blocks)
      INTEGER, DIMENSION(:), INTENT(INOUT)               :: blk_p
      INTEGER, DIMENSION(:), INTENT(IN)                  :: old_blk_p, row_p, col_i, rbs, cbs
      TYPE(dbcsr_data_obj), INTENT(INOUT)                :: dst
      TYPE(dbcsr_data_obj), INTENT(IN)                   :: src
      LOGICAL, INTENT(IN), OPTIONAL                      :: mark_transposed, transpose_blocks

      CHARACTER(len=*), PARAMETER :: routineN = 'dbcsr_copy_sort_data', &
         routineP = moduleN//':'//routineN

      INTEGER                                            :: blk, col_size, nblks, nrows, nze, &
                                                            nze_prev, row, row_size
      LOGICAL                                            :: mark, trb

!   ---------------------------------------------------------------------------
! Analyze parameters

      mark = .FALSE.
      IF (PRESENT(mark_transposed)) mark = mark_transposed
      trb = .FALSE.
      IF (PRESENT(transpose_blocks)) trb = transpose_blocks
      !
      nblks = SIZE(old_blk_p)
      nrows = SIZE(row_p)-1
      IF (SIZE(blk_p) < nblks) &
         CPABORT('Destination blk_p too small.')
      IF (nblks .GE. 1) &
         blk_p(1) = SGN(1, old_blk_p(1), mark)
      nze_prev = 0
      DO row = 1, nrows
         row_size = rbs(row)
         DO blk = row_p(row)+1, row_p(row+1)
            IF (old_blk_p(blk) .NE. 0) THEN
               col_size = cbs(col_i(blk))
               nze = row_size*col_size
               IF (blk .GT. 1) THEN
                  blk_p(blk) = SGN(ABS(blk_p(blk-1))+nze_prev, old_blk_p(blk), &
                                   mark)
               ENDIF
               IF (ABS(blk_p(blk))+nze-1 > dbcsr_data_get_size(dst)) &
                  CPABORT('Destination data space is too small.')
               IF (.NOT. trb) THEN
                  CALL dbcsr_data_copy(dst=dst, dst_lb=(/ABS(blk_p(blk))/), &
                                       dst_sizes=(/nze/), &
                                       src=src, src_lb=(/ABS(old_blk_p(blk))/), &
                                       src_sizes=(/nze/))
                  !CALL dbcsr_data_set (dst, ABS(blk_p(blk)), nze,&
                  !     src, source_lb=ABS(old_blk_p(blk)))
               ELSE
                  CALL dbcsr_block_transpose(dst, src, &
                                             col_size, row_size, &
                                             lb=ABS(blk_p(blk)), source_lb=ABS(old_blk_p(blk)))
               ENDIF
               nze_prev = nze
            ENDIF ! blk exists
         ENDDO ! blk
      ENDDO ! row
   END SUBROUTINE dbcsr_copy_sort_data

! **************************************************************************************************
!> \brief Sorts the data in a matrix so that the data blocks follow
!>        sequentially.
!> \param[out] blk_p          re-arranged block pointers reflecting the new
!>                            data order
!> \param[in] old_blk_p       current block pointers
!> \param[in] sizes           sizes of the data blocks
!> \param[out] dsts           sorted data
!> \param[in] src             existing unordered data
!> \param[in] srcs            (optional) multiple source areas
!> \param old_blk_d ...
! **************************************************************************************************
   SUBROUTINE dbcsr_sort_data(blk_p, old_blk_p, sizes, dsts, src, &
                              srcs, old_blk_d)
      INTEGER, DIMENSION(:), INTENT(INOUT)               :: blk_p
      INTEGER, DIMENSION(:), INTENT(IN)                  :: old_blk_p, sizes
      TYPE(dbcsr_data_obj), INTENT(INOUT)                :: dsts
      TYPE(dbcsr_data_obj), INTENT(IN)                   :: src
      TYPE(dbcsr_data_obj), DIMENSION(:), INTENT(IN), &
         OPTIONAL                                        :: srcs
      INTEGER, DIMENSION(:), INTENT(IN), OPTIONAL        :: old_blk_d

      CHARACTER(len=*), PARAMETER :: routineN = 'dbcsr_sort_data', &
         routineP = moduleN//':'//routineN

      INTEGER                                            :: handle, i, nblks
      LOGICAL                                            :: multidata

!   ---------------------------------------------------------------------------

      CALL timeset(routineN, handle)
      multidata = PRESENT(srcs) .AND. PRESENT(old_blk_d)
      nblks = SIZE(old_blk_p)
      IF (nblks .GT. 0) THEN
!$OMP        BARRIER
!$OMP        MASTER
         blk_p(1) = SIGN(1, old_blk_p(1))
         DO i = 2, nblks
            blk_p(i) = SIGN(ABS(blk_p(i-1))+sizes(i-1), old_blk_p(i))
         ENDDO
         CALL dbcsr_data_set_size_referenced(dsts, &
                                             ABS(blk_p(nblks))+sizes(nblks)-1)
!$OMP        END MASTER
!$OMP        BARRIER
!$OMP        DO
         DO i = 1, nblks
            IF (old_blk_p(i) .NE. 0) THEN
               IF (.NOT. multidata) THEN
                  CALL dbcsr_data_set(dsts, &
                                      ABS(blk_p(i)), sizes(i), &
                                      src, source_lb=ABS(old_blk_p(i)))
                  !dst(ABS(blk_p(i)):ABS(blk_p(i))+sizes(i)-1) =&
                  !     src(ABS(old_blk_p(i)):ABS(old_blk_p(i))+sizes(i)-1)
               ELSE
                  CALL dbcsr_data_set(dsts, &
                                      ABS(blk_p(i)), sizes(i), &
                                      srcs(old_blk_d(i)), source_lb=ABS(old_blk_p(i)))
                  !dst(ABS(blk_p(i)):ABS(blk_p(i))+sizes(i)-1) =&
                  !     srcs(old_blk_d(i))%d&
                  !     %r_dp(ABS(old_blk_p(i)):ABS(old_blk_p(i))+sizes(i)-1)
               ENDIF
            ENDIF
         ENDDO
!$OMP        END DO NOWAIT
      ENDIF
      CALL timestop(handle)
   END SUBROUTINE dbcsr_sort_data

END MODULE dbcsr_data_operations
