program test_compression
  use yaml_output
  use f_unittests
  use wrapper_MPI

  implicit none

  type(mpi_environment) :: env
  
  call mpiinit()

  call f_lib_initialize()

  env = mpi_environment_comm()  

  if (env%iproc == 0) call yaml_sequence_open("Testing type(wavefunctions_descriptors)")
  call run(null)
  call run(allocateWfd)
  call run(deallocateWfd)
  call run(initFromFullGrid)
  call run(initFromSubGrid)
  call run(initFromDisk)
  call run(extractFromBox)
  call run(extractFromSphere)
  call run(equals)
  call run(sendReceive, env)
  call run(broadcast, env)
  call run(deepCopy)
  if (env%iproc == 0) call yaml_sequence_close()
  
  if (env%iproc == 0) call yaml_sequence_open("Testing type(wfd_to_wfd)")
  call run(nullWfdToWfd)
  call run(tolrStrategy)
  call run(initTolr)
  call run(freeTolr)
  if (env%iproc == 0) call yaml_sequence_close()
  
  if (env%iproc == 0) call yaml_sequence_open("Testing operations")
  call run(prDotPsi)
  call run(cprojDot)
  call run(cprojPrPsi)
  if (env%iproc == 0) call yaml_sequence_close()

  if (env%iproc == 0) then
     call f_lib_finalize()
  else
     call f_lib_finalize_noreport()
  end if

  call release_mpi_environment(env)
  call mpifinalize()

contains

  subroutine null(label)
    use compression
    implicit none
    character(len = *), intent(out) :: label

    type(wavefunctions_descriptors) :: wfd

    label = "Nullify"
    call nullify_wfd(wfd)
    call verify(wfd == wfd_null(), "equals to null wfd")
  end subroutine null

  subroutine allocateWfd(label)
    use compression
    implicit none
    character(len = *), intent(out) :: label

    type(wavefunctions_descriptors) :: wfd

    label = "Allocate"
    call nullify_wfd(wfd)
    wfd%nseg_c = 5
    wfd%nseg_f = 7

    call allocate_wfd(wfd)
    call verify(associated(wfd%buffer), "associated buffer")
    call compare(size(wfd%buffer), 72, "buffer size, global = not set")
    call deallocate_wfd(wfd)

    call allocate_wfd(wfd, global = .false.)
    call verify(associated(wfd%buffer), "associated buffer")
    call compare(size(wfd%buffer), 72, "buffer size, global = .false.")
    call verify(associated(wfd%keyglob), "associated keyglob")
    call compare(size(wfd%keyglob), 24, "buffer size, global = .false.")
    call verify(associated(wfd%keyvglob), "associated keyvglob")
    call compare(size(wfd%keyvglob), 12, "buffer size, global = .false.")
    call verify(associated(wfd%keygloc), "associated keygloc")
    call compare(size(wfd%keyglob), 24, "buffer size, global = .false.")
    call verify(associated(wfd%keyvloc), "associated keyvloc")
    call compare(size(wfd%keyvglob), 12, "buffer size, global = .false.")
    call deallocate_wfd(wfd)

    call allocate_wfd(wfd, global = .true.)
    call verify(associated(wfd%buffer), "associated buffer")
    call compare(size(wfd%buffer), 36, "buffer size, global = .true.")
    call verify(associated(wfd%keyglob), "associated keyglob")
    call compare(size(wfd%keyglob), 24, "buffer size, global = .false.")
    call verify(associated(wfd%keyvglob), "associated keyvglob")
    call compare(size(wfd%keyvglob), 12, "buffer size, global = .false.")
    call verify(associated(wfd%keygloc, target = wfd%keyglob), "associated keygloc")
    call verify(associated(wfd%keyvloc, target = wfd%keyvglob), "associated keyvloc")
    call deallocate_wfd(wfd)    
  end subroutine allocateWfd

  subroutine deallocateWfd(label)
    use compression
    implicit none
    character(len = *), intent(out) :: label

    type(wavefunctions_descriptors) :: wfd

    label = "Deallocate"
    call nullify_wfd(wfd)
    wfd%nseg_c = 5
    wfd%nseg_f = 7

    call allocate_wfd(wfd)
    call verify(associated(wfd%buffer), "associated buffer")
    call deallocate_wfd(wfd)

    call verify(.not. associated(wfd%keyglob), "associated keyglob")
    call verify(.not. associated(wfd%keyvglob), "associated keyvglob")
    call verify(.not. associated(wfd%keygloc), "associated keygloc")
    call verify(.not. associated(wfd%keyvloc), "associated keyvloc")
    call verify(.not. associated(wfd%buffer), "associated buffer")
  end subroutine deallocateWfd

  subroutine initFromFullGrid(label)
    use compression
    use f_precisions, only: f_byte
    implicit none
    character(len = *), intent(out) :: label

    type(wavefunctions_descriptors) :: wfd
    integer, parameter :: n1 = 10, n2 = 11, n3 = 13
    logical(f_byte), dimension(0:n1,0:n2,0:n3) :: logrid_c, logrid_f

    label = "Init from full grid"

    logrid_c = .false.
    logrid_f = .false.
    call init_wfd_from_full_grids(wfd, n1, n2, n3, logrid_c, logrid_f)
    call compare(wfd%nseg_c, 0, "nseg_c for no points")
    call compare(wfd%nseg_f, 0, "nseg_f for no points")
    call compare(wfd%nvctr_c, 0, "nvctr_c for no points")
    call compare(wfd%nvctr_f, 0, "nvctr_f for no points")
    call deallocate_wfd(wfd)

    logrid_c(2:5, 7, 9) = .true.
    logrid_c(4:8, 8, 9) = .true.
    logrid_c(3, 8, 10) = .true.
    call init_wfd_from_full_grids(wfd, n1, n2, n3, logrid_c, logrid_f)
    call verify(wfd_is_global(wfd), "global grid coarse only")
    call compare(wfd%nseg_c, 3, "nseg_c")
    call compare(wfd%nseg_f, 0, "nseg_f")
    call compare(wfd%nvctr_c, 10, "nvctr_c")
    call compare(wfd%nvctr_f, 0, "nvctr_f")
    call compare(wfd%keyglob(:, 1), (/ 1268, 1271 /), "keyglob seg 1")
    call compare(wfd%keyglob(:, 2), (/ 1281, 1285 /), "keyglob seg 2")
    call compare(wfd%keyglob(:, 3), (/ 1412, 1412 /), "keyglob seg 3")
    call compare(wfd%keyvglob, (/ 1, 5, 10 /), "keyvglob")
    call deallocate_wfd(wfd)

    logrid_f(2:3, 7, 9) = .true.
    logrid_f(5:6, 8, 9) = .true.
    call init_wfd_from_full_grids(wfd, n1, n2, n3, logrid_c, logrid_f)
    call verify(wfd_is_global(wfd), "global grid")
    call compare(wfd%nseg_c, 3, "nseg_c")
    call compare(wfd%nseg_f, 2, "nseg_f")
    call compare(wfd%nvctr_c, 10, "nvctr_c")
    call compare(wfd%nvctr_f, 4, "nvctr_f")
    call compare(wfd%keyglob(:, 1), (/ 1268, 1271 /), "keyglob seg 1")
    call compare(wfd%keyglob(:, 2), (/ 1281, 1285 /), "keyglob seg 2")
    call compare(wfd%keyglob(:, 3), (/ 1412, 1412 /), "keyglob seg 3")
    call compare(wfd%keyglob(:, 4), (/ 1268, 1269 /), "keyglob fine seg 1")
    call compare(wfd%keyglob(:, 5), (/ 1282, 1283 /), "keyglob fine seg 2")
    call compare(wfd%keyvglob, (/ 1, 5, 10, 1, 3 /), "keyvglob")
    call deallocate_wfd(wfd)
  end subroutine initFromFullGrid

  subroutine initFromSubGrid(label)
    use compression
    use f_precisions, only: f_byte
    implicit none
    character(len = *), intent(out) :: label

    type(wavefunctions_descriptors) :: wfd
    integer, parameter :: n1 = 10, n2 = 11, n3 = 13
    integer, dimension(2, 3) :: nbox_c, nbox_f, nbox_l
    logical(f_byte), dimension(0:n1,0:n2,0:n3) :: logrid_c, logrid_f

    label = "Init from sub grid"

    nbox_c = 0
    nbox_f = 0

    logrid_c = .false.
    logrid_f = .false.
    call init_wfd_from_sub_grids(wfd, n1, n2, n3, logrid_c, nbox_c, logrid_f, nbox_f)
    call compare(wfd%nseg_c, 0, "nseg_c for no points")
    call compare(wfd%nseg_f, 0, "nseg_f for no points")
    call compare(wfd%nvctr_c, 0, "nvctr_c for no points")
    call compare(wfd%nvctr_f, 0, "nvctr_f for no points")

    logrid_c(2:5, 7, 9) = .true.
    logrid_c(4:8, 8, 9) = .true.
    logrid_c(3, 8, 10) = .true.
    nbox_c(:, 1) = (/ 1, 3 /)
    nbox_c(:, 2) = (/ 6, 8 /)
    nbox_c(:, 3) = (/ 8, 10 /)
    nbox_l = nbox_c ! keyloc should be referenced to the coarse region
    call init_wfd_from_sub_grids(wfd, n1, n2, n3, logrid_c, nbox_c, logrid_f, nbox_f, &
         do_allocate = .true., nbox_local = nbox_l)
    call verify(.not. wfd_is_global(wfd), "global grid coarse only")
    call compare(wfd%nseg_c, 2, "nseg_c")
    call compare(wfd%nseg_f, 0, "nseg_f")
    call compare(wfd%nvctr_c, 3, "nvctr_c")
    call compare(wfd%nvctr_f, 0, "nvctr_f")
    call compare(wfd%keyglob(:, 1), (/ 1268, 1269 /), "keyglob seg 1")
    call compare(wfd%keyglob(:, 2), (/ 1412, 1412 /), "keyglob seg 2")
    call compare(wfd%keyvglob, (/ 1, 3 /), "keyvglob")
    call compare(wfd%keygloc(:, 1), (/ 14, 15 /), "keygloc seg 1")
    call compare(wfd%keygloc(:, 2), (/ 27, 27 /), "keygloc seg 2")
    call compare(wfd%keyvloc, wfd%keyvglob, "keyvloc")
    call deallocate_wfd(wfd)

    logrid_f(2:3, 7, 9) = .true.
    logrid_f(5:6, 8, 9) = .true.
    nbox_f(:, 1) = (/ 2, 3 /)
    nbox_f(:, 2) = (/ 7, 8 /)
    nbox_f(:, 3) = (/ 9, 9 /)
    call init_wfd_from_sub_grids(wfd, n1, n2, n3, logrid_c, nbox_c, logrid_f, nbox_f, &
         do_allocate = .true., nbox_local = nbox_l)
    call verify(.not. wfd_is_global(wfd), "global grid")
    call compare(wfd%nseg_c, 2, "nseg_c")
    call compare(wfd%nseg_f, 1, "nseg_f")
    call compare(wfd%nvctr_c, 3, "nvctr_c")
    call compare(wfd%nvctr_f, 2, "nvctr_f")
    call compare(wfd%keyglob(:, 1), (/ 1268, 1269 /), "keyglob seg 1")
    call compare(wfd%keyglob(:, 2), (/ 1412, 1412 /), "keyglob seg 2")
    call compare(wfd%keyglob(:, 3), (/ 1268, 1269 /), "keyglob seg 2")
    call compare(wfd%keyvglob, (/ 1, 3, 1 /), "keyvglob")
    call compare(wfd%keygloc(:, 1), (/ 14, 15 /), "keygloc seg 1")
    call compare(wfd%keygloc(:, 2), (/ 27, 27 /), "keygloc seg 2")
    call compare(wfd%keygloc(:, 3), (/ 14, 15 /), "keygloc seg 2")
    call compare(wfd%keyvloc, wfd%keyvglob, "keyvloc")
    call deallocate_wfd(wfd)
  end subroutine initFromSubGrid

  subroutine initFromDisk(label)
    use compression
    implicit none
    character(len = *), intent(out) :: label

    integer, parameter :: unit = 456
    type(wavefunctions_descriptors) :: wfdG, wfd
    integer, parameter :: n1 = 10, n2 = 11, n3 = 13
    logical(f_byte), dimension(0:n1,0:n2,0:n3) :: logrid_c, logrid_f
    integer :: iseg
    
    label = "Init from disk"

    wfdG = wfdSample(n1, n2, n3, logrid_c, logrid_f)

    open(unit, action = "WRITE", form = "UNFORMATTED")
    write(unit) wfdG%nvctr_c, wfdG%nvctr_f
    write(unit) wfdG%nseg_c, wfdG%nseg_f
    do iseg=1,wfdG%nseg_c
       write(unit) wfdG%keygloc(1:2,iseg), wfdG%keyglob(1:2,iseg), wfdG%keyvloc(iseg), wfdG%keyvglob(iseg)
    end do
    do iseg=wfdG%nseg_c + 1,wfdG%nseg_c + wfdG%nseg_f
       write(unit) wfdG%keygloc(1:2,iseg), wfdG%keyglob(1:2,iseg), wfdG%keyvloc(iseg), wfdG%keyvglob(iseg)
    end do
    close(unit)

    open(unit, action = "READ", form = "UNFORMATTED")
    call init_wfd_from_disk(wfd, unit)
    close(unit)
    call verify(wfd == wfdG, "unformatted")
    call deallocate_wfd(wfd)

    open(unit, action = "WRITE", form = "FORMATTED")
    write(unit, *) wfdG%nvctr_c, wfdG%nvctr_f
    write(unit, *) wfdG%nseg_c, wfdG%nseg_f
    do iseg=1,wfdG%nseg_c
       write(unit, *) wfdG%keygloc(1:2,iseg), wfdG%keyglob(1:2,iseg), wfdG%keyvloc(iseg), wfdG%keyvglob(iseg)
    end do
    do iseg=wfdG%nseg_c + 1,wfdG%nseg_c + wfdG%nseg_f
       write(unit, *) wfdG%keygloc(1:2,iseg), wfdG%keyglob(1:2,iseg), wfdG%keyvloc(iseg), wfdG%keyvglob(iseg)
    end do
    close(unit)

    open(unit, action = "READ", form = "FORMATTED")
    call init_wfd_from_disk_with_format(wfd, unit)
    close(unit)
    call verify(wfd == wfdG, "formatted")
    call deallocate_wfd(wfd)

    call deallocate_wfd(wfdG)
  end subroutine initFromDisk

  subroutine extractFromBox(label)
    use compression
    use f_precisions, only: f_byte
    implicit none
    character(len = *), intent(out) :: label

    type(wavefunctions_descriptors) :: wfd, wfdG
    integer, parameter :: n1 = 10, n2 = 11, n3 = 13
    integer, dimension(2, 3) :: nbox
    logical(f_byte), dimension(0:n1,0:n2,0:n3) :: logrid_c, logrid_f
    integer, dimension(3) :: outofzone

    label = "Extract from box"

    wfdG = wfdSample(n1, n2, n3, logrid_c, logrid_f)

    nbox(:, 1) = (/ 1, 3 /)
    nbox(:, 2) = (/ 6, 8 /)
    nbox(:, 3) = (/ 8, 10 /)
    outofzone = 0
    call extract_from_box(wfd, n1, n2, n3, wfdG, nbox, outofzone)
    call verify(.not. wfd_is_global(wfd), "global grid")
    call compare(wfd%nseg_c, 2, "nseg_c")
    call compare(wfd%nseg_f, 1, "nseg_f")
    call compare(wfd%nvctr_c, 3, "nvctr_c")
    call compare(wfd%nvctr_f, 2, "nvctr_f")
    call compare(wfd%keyglob(:, 1), (/ 1268, 1269 /), "keyglob seg 1")
    call compare(wfd%keyglob(:, 2), (/ 1412, 1412 /), "keyglob seg 2")
    call compare(wfd%keyglob(:, 3), (/ 1268, 1269 /), "keyglob seg 2")
    call compare(wfd%keyvglob, (/ 1, 3, 1 /), "keyvglob")
    call compare(wfd%keygloc(:, 1), (/ 14, 15 /), "keygloc seg 1")
    call compare(wfd%keygloc(:, 2), (/ 27, 27 /), "keygloc seg 2")
    call compare(wfd%keygloc(:, 3), (/ 14, 15 /), "keygloc seg 2")
    call compare(wfd%keyvloc, wfd%keyvglob, "keyvloc")
    call deallocate_wfd(wfd)

    ! Todo: need to test the outofzone argument.
    
    call deallocate_wfd(wfdG)
  end subroutine extractFromBox

  subroutine extractFromSphere(label)
    use compression
    use f_precisions, only: f_byte
    use liborbs_precisions
    implicit none
    character(len = *), intent(out) :: label

    type(wavefunctions_descriptors) :: wfd, wfdG
    integer, parameter :: n1 = 10, n2 = 11, n3 = 13
    real(gp), parameter :: hx = 0.2_gp, hy = 0.2_gp, hz = 0.2_gp
    integer, dimension(2, 3) :: nbox
    logical(f_byte), dimension(0:n1,0:n2,0:n3) :: logrid_c, logrid_f

    label = "Extract from sphere"

    wfdG = wfdSample(n1, n2, n3, logrid_c, logrid_f)
    
    nbox(:, 1) = (/ 1, 3 /)
    nbox(:, 2) = (/ 6, 8 /)
    nbox(:, 3) = (/ 8, 10 /)
    call extract_from_sphere(wfd, n1, n2, n3, 0, 0, 0, .false., .false., .false., &
         hx, hy, hz, wfdG, (/ 0.4_gp, 1.4_gp, 1.8_gp /), 0.35_gp, nbox)
    call verify(.not. wfd_is_global(wfd), "global grid")
    call compare(wfd%nseg_c, 2, "nseg_c")
    call compare(wfd%nseg_f, 1, "nseg_f")
    call compare(wfd%nvctr_c, 3, "nvctr_c")
    call compare(wfd%nvctr_f, 2, "nvctr_f")
    call compare(wfd%keyglob(:, 1), (/ 1268, 1269 /), "keyglob seg 1")
    call compare(wfd%keyglob(:, 2), (/ 1412, 1412 /), "keyglob seg 2")
    call compare(wfd%keyglob(:, 3), (/ 1268, 1269 /), "keyglob seg 2")
    call compare(wfd%keyvglob, (/ 1, 3, 1 /), "keyvglob")
    call compare(wfd%keygloc(:, 1), (/ 14, 15 /), "keygloc seg 1")
    call compare(wfd%keygloc(:, 2), (/ 27, 27 /), "keygloc seg 2")
    call compare(wfd%keygloc(:, 3), (/ 14, 15 /), "keygloc seg 2")
    call compare(wfd%keyvloc, wfd%keyvglob, "keyvloc")
    call deallocate_wfd(wfd)

    ! Todo: need to test other periodicities.

    call deallocate_wfd(wfdG)
  end subroutine extractFromSphere

  function wfdSample(n1, n2, n3, logrid_c, logrid_f, global) result(wfd)
    use compression
    use f_precisions, only: f_byte
    implicit none
    integer, intent(in) :: n1, n2, n3
    logical(f_byte), dimension(0:n1,0:n2,0:n3), intent(out) :: logrid_c, logrid_f
    logical, intent(in), optional :: global
    type(wavefunctions_descriptors) :: wfd
    integer, dimension(2, 3) :: nbox_c, nbox_f, nbox_l
    
    logrid_c = .false.
    logrid_f = .false.
    logrid_c(2:5, 7, 9) = .true.
    logrid_c(4:8, 8, 9) = .true.
    logrid_c(3, 8, 10) = .true.
    logrid_f(2:3, 7, 9) = .true.
    logrid_f(5:6, 8, 9) = .true.
    if (present(global)) then
       if (.not. global) then
          nbox_c(:, 1) = (/ 1, 5 /)
          nbox_c(:, 2) = (/ 6, 8 /)
          nbox_c(:, 3) = (/ 8, 10 /)
          nbox_l = nbox_c
          nbox_f(:, 1) = (/ 2, 5 /)
          nbox_f(:, 2) = (/ 7, 8 /)
          nbox_f(:, 3) = (/ 9, 9 /)
          call init_wfd_from_sub_grids(wfd, n1, n2, n3, logrid_c, nbox_c, logrid_f, nbox_f, &
               do_allocate = .true., nbox_local = nbox_l)
       else
          call init_wfd_from_full_grids(wfd, n1, n2, n3, logrid_c, logrid_f)
       end if
    else
       call init_wfd_from_full_grids(wfd, n1, n2, n3, logrid_c, logrid_f)
    end if
  end function wfdSample

  function wfdPSample(n1, n2, n3, logrid_c, logrid_f, global) result(wfd)
    use compression
    use f_precisions, only: f_byte
    implicit none
    integer, intent(in) :: n1, n2, n3
    logical(f_byte), dimension(0:n1,0:n2,0:n3), intent(out) :: logrid_c, logrid_f
    logical, intent(in), optional :: global
    type(wavefunctions_descriptors) :: wfd
    integer, dimension(2, 3) :: nbox_c, nbox_f, nbox_l
    
    logrid_c = .false.
    logrid_f = .false.
    logrid_c(3:7, 8, 9) = .true.
    logrid_c(6:10, 9, 9) = .true.
    logrid_c(5, 9, 10) = .true.
    logrid_f(4:5, 8, 9) = .true.
    logrid_f(7:8, 9, 9) = .true.
    if (present(global)) then
       if (.not. global) then
          nbox_c(:, 1) = (/ 3, 5 /)
          nbox_c(:, 2) = (/ 7, 9 /)
          nbox_c(:, 3) = (/ 8, 10 /)
          nbox_l = nbox_c
          nbox_f(:, 1) = (/ 4, 5 /)
          nbox_f(:, 2) = (/ 8, 9 /)
          nbox_f(:, 3) = (/ 9, 9 /)
          call init_wfd_from_sub_grids(wfd, n1, n2, n3, logrid_c, nbox_c, logrid_f, nbox_f, &
               do_allocate = .true., nbox_local = nbox_l)
       else
          call init_wfd_from_full_grids(wfd, n1, n2, n3, logrid_c, logrid_f)
       end if
    else
       call init_wfd_from_full_grids(wfd, n1, n2, n3, logrid_c, logrid_f)
    end if
  end function wfdPSample

  subroutine equals(label)
    use compression
    use f_precisions, only: f_byte
    implicit none
    character(len = *), intent(out) :: label

    type(wavefunctions_descriptors) :: wfd1, wfd2
    integer, parameter :: n1 = 10, n2 = 11, n3 = 13
    logical(f_byte), dimension(0:n1,0:n2,0:n3) :: logrid_c, logrid_f

    label = "Equal operator"

    wfd1 = wfdSample(n1, n2, n3, logrid_c, logrid_f)

    call init_wfd_from_full_grids(wfd2, n1, n2, n3, logrid_c, logrid_f)
    call verify(wfd1 == wfd2, "equal case")
    call deallocate_wfd(wfd2)

    logrid_f = .false.
    logrid_f(2:3, 7, 9) = .true.
    call init_wfd_from_full_grids(wfd2, n1, n2, n3, logrid_c, logrid_f)
    call verify(.not. wfd1 == wfd2, "different on nseg_f")
    call deallocate_wfd(wfd2)

    logrid_c = .false.
    logrid_f = .false.
    logrid_c(2:5, 7, 9) = .true.
    logrid_c(4:8, 8, 9) = .true.
    logrid_f(2:3, 7, 9) = .true.
    logrid_f(5:6, 8, 9) = .true.
    call init_wfd_from_full_grids(wfd2, n1, n2, n3, logrid_c, logrid_f)
    call verify(.not. wfd1 == wfd2, "different on nseg_c")
    call deallocate_wfd(wfd2)

    call deallocate_wfd(wfd1)
  end subroutine equals

  subroutine sendReceive(label, env)
    use compression
    use f_precisions, only: f_byte
    implicit none
    character(len = *), intent(out) :: label
    type(mpi_environment), intent(in) :: env

    type(wavefunctions_descriptors) :: wfdRef, wfd
    integer, parameter :: n1 = 10, n2 = 11, n3 = 13
    logical(f_byte), dimension(0:n1,0:n2,0:n3) :: logrid_c, logrid_f
    integer :: req

    label = "MPI key send/receive"

    if (env%nproc < 2) then
       call skip("not an MPI run")
       return
    end if

    ! Global case
    wfdRef = wfdSample(n1, n2, n3, logrid_c, logrid_f)
    if (env%iproc == 1) then
       call send_wfd_keys(wfdRef, 0, 123, req, env%mpi_comm)
    else if (env%iproc == 0) then
       wfd = wfdRef
       call allocate_wfd(wfd, .true.)
       call recv_wfd_keys(wfd, 1, 123, req, env%mpi_comm)
       call fmpi_wait(req)
       call verify(wfd == wfdRef, "receiving keys")
       call deallocate_wfd(wfd)
    end if
    call deallocate_wfd(wfdRef)

    ! Local case
    wfdRef = wfdSample(n1, n2, n3, logrid_c, logrid_f, global = .false.)
    if (env%iproc == 1) then
       call send_wfd_keys(wfdRef, 0, 123, req, env%mpi_comm)
    else if (env%iproc == 0) then
       wfd = wfdRef
       call allocate_wfd(wfd)
       call recv_wfd_keys(wfd, 1, 123, req, env%mpi_comm)
       call fmpi_wait(req)
       call verify(wfd == wfdRef, "receiving keys")
       call deallocate_wfd(wfd)
    end if
    call deallocate_wfd(wfdRef)
  end subroutine sendReceive

  subroutine broadcast(label, env)
    use compression
    use f_precisions, only: f_byte
    implicit none
    character(len = *), intent(out) :: label
    type(mpi_environment), intent(in) :: env

    type(wavefunctions_descriptors) :: wfdRef, wfd
    integer, parameter :: n1 = 10, n2 = 11, n3 = 13
    logical(f_byte), dimension(0:n1,0:n2,0:n3) :: logrid_c, logrid_f

    label = "MPI key broadcast"

    if (env%nproc < 2) then
       call skip("not an MPI run")
       return
    end if

    ! Global case
    wfd = wfdSample(n1, n2, n3, logrid_c, logrid_f, global = .true.)
    if (env%iproc /= 1) then
       call deallocate_wfd(wfd)
    end if
    call wfd_keys_broadcast(wfd, 1, env%mpi_comm, .true.)
    if (env%iproc == 0) then
       wfdRef = wfdSample(n1, n2, n3, logrid_c, logrid_f)
       call verify(wfd == wfdRef, "receiving keys")
       call deallocate_wfd(wfdRef)
    end if
    call deallocate_wfd(wfd)

    ! Local case
    wfd = wfdSample(n1, n2, n3, logrid_c, logrid_f, global = .false.)
    if (env%iproc /= 1) then
       call deallocate_wfd(wfd)
    end if
    call wfd_keys_broadcast(wfd, 1, env%mpi_comm)
    if (env%iproc == 0) then
       wfdRef = wfdSample(n1, n2, n3, logrid_c, logrid_f, global = .false.)
       call verify(wfd == wfdRef, "receiving keys")
       call deallocate_wfd(wfdRef)
    end if
    call deallocate_wfd(wfd)
  end subroutine broadcast

  subroutine deepCopy(label)
    use compression
    use f_precisions, only: f_byte
    implicit none
    character(len = *), intent(out) :: label

    type(wavefunctions_descriptors) :: wfdRef, wfd
    integer, parameter :: n1 = 10, n2 = 11, n3 = 13
    logical(f_byte), dimension(0:n1,0:n2,0:n3) :: logrid_c, logrid_f

    label = "Deep copy"

    ! Global case
    wfdRef = wfdSample(n1, n2, n3, logrid_c, logrid_f, global = .true.)
    call copy_wfd(wfdRef, wfd)
    call verify(wfd == wfdRef, "global case")
    call verify(associated(wfd%buffer) .and. &
         .not. associated(wfd%buffer, target = wfdRef%buffer), "buffer deep copy")
    call deallocate_wfd(wfd)
    call deallocate_wfd(wfdRef)

    ! Local case
    wfdRef = wfdSample(n1, n2, n3, logrid_c, logrid_f, global = .false.)
    call copy_wfd(wfdRef, wfd)
    call verify(wfd == wfdRef, "local case")
    call verify(associated(wfd%buffer) .and. &
         .not. associated(wfd%buffer, target = wfdRef%buffer), "buffer deep copy")
    call deallocate_wfd(wfd)
    call deallocate_wfd(wfdRef)
  end subroutine deepCopy

  subroutine nullWfdToWfd(label)
    use compression
    implicit none
    character(len = *), intent(out) :: label

    type(wfd_to_wfd) :: wfd

    label = "Nullify"
    call nullify_wfd_to_wfd(wfd)
    call verify(wfd == wfd_to_wfd_null(), "equals to null wfd_to_wfd")
    call deallocate_wfd_to_wfd(wfd)
  end subroutine nullWfdToWfd

  subroutine tolrStrategy(label)
    use compression
    implicit none
    character(len = *), intent(out) :: label

    type(wfd_to_wfd) :: wfd

    label = "Set strategy"
    call nullify_wfd_to_wfd(wfd)
    call verify(wfd_to_wfd_skip(wfd), "skip")
    call tolr_set_strategy(wfd, "MASK_PACK")
    call verify(wfd_to_wfd_mask_pack(wfd), "mask pack")
    call tolr_set_strategy(wfd, "MASK")
    call verify(wfd_to_wfd_mask(wfd), "mask")
    call tolr_set_strategy(wfd, "KEYS")
    call verify(wfd_to_wfd_keys(wfd), "keys")
    call tolr_set_strategy(wfd, "KEYS_PACK")
    call verify(wfd_to_wfd_keys_pack(wfd), "keys pack")
    call deallocate_wfd_to_wfd(wfd)
  end subroutine tolrStrategy

  subroutine sampleTolr(tolr, wfd, wfdp)
    use compression
    use dynamic_memory
    implicit none

    type(wfd_to_wfd), intent(out) :: tolr
    type(wavefunctions_descriptors), intent(out) :: wfd, wfdp

    integer, parameter :: n1 = 10, n2 = 11, n3 = 13
    logical(f_byte), dimension(0:n1,0:n2,0:n3) :: logrid_c, logrid_f
    integer, dimension(:), allocatable :: keyag_lin_cf, nbsegs_cf

    wfd = wfdSample(n1, n2, n3, logrid_c, logrid_f, global = .false.)
    keyag_lin_cf = f_malloc(wfd%nseg_c + wfd%nseg_f, "keyag_lin_cf")

    wfdp = wfdPSample(n1, n2, n3, logrid_c, logrid_f, global = .false.)
    nbsegs_cf = f_malloc(wfdp%nseg_c + wfdp%nseg_f, "nbsegs_cf")

    call nullify_wfd_to_wfd(tolr)
    call init_tolr(tolr, wfd, wfdp, keyag_lin_cf, nbsegs_cf)
    call f_free(nbsegs_cf)
    call f_free(keyag_lin_cf)
  end subroutine sampleTolr

  subroutine initTolr(label)
    use compression
    use dynamic_memory
    implicit none
    character(len = *), intent(out) :: label

    type(wfd_to_wfd) :: tolr
    type(wavefunctions_descriptors) :: wfd, wfdp

    label = "Init"

    call sampleTolr(tolr, wfd, wfdp)
    call compare(tolr%nmseg_c, 1, "nmseg coarse")
    call compare(tolr%nmseg_f, 1, "nmseg fine")
    call verify(associated(tolr%mask), "associated mask")
    call compare(tolr%mask, reshape((/ 2, 5, 2, 1, 3, 2 /), (/ 3, 2 /)), "mask")

    call deallocate_wfd(wfdp)
    call deallocate_wfd(wfd)
    call deallocate_wfd_to_wfd(tolr)
  end subroutine initTolr

  subroutine freeTolr(label)
    use compression
    use dynamic_memory
    implicit none
    character(len = *), intent(out) :: label

    type(wfd_to_wfd) :: tolr
    type(wavefunctions_descriptors) :: wfd, wfdp

    label = "Free"

    call sampleTolr(tolr, wfd, wfdp)
    call deallocate_wfd(wfdp)
    call deallocate_wfd(wfd)

    call deallocate_wfd_to_wfd(tolr)
    call verify(.not. associated(tolr%mask), "mask")
  end subroutine freeTolr

  subroutine prDotPsi(label)
    use liborbs_precisions
    use compression
    use dynamic_memory
    implicit none
    character(len = *), intent(out) :: label

    type(wfd_to_wfd) :: tolr
    type(wavefunctions_descriptors) :: wfd, wfdp

    integer, parameter :: nproj = 3, npsi = 2
    integer, parameter :: ncplx = 2
    real(wp), dimension(:, :, :), allocatable :: psi, proj
    real(wp), dimension(:, :), allocatable :: work
    real(wp), dimension(ncplx, npsi, ncplx, nproj) :: scpr
    real(wp), dimension(ncplx, npsi, nproj) :: cproj

    label = "Proj . Psi"

    call sampleTolr(tolr, wfd, wfdp)

    proj = f_malloc((/wfdp%nvctr_c + 7 * wfdp%nvctr_f, ncplx, nproj/), "proj")
    psi = f_malloc((/wfd%nvctr_c + 7 * wfd%nvctr_f, ncplx, npsi/), "psi")
    work = f_malloc((/wfd%nvctr_c + 7 * wfd%nvctr_f, ncplx * npsi/), "work")
    psi = 1_wp
    proj = 1_wp

    call tolr_set_strategy(tolr, "MASK_PACK")
    call pr_dot_psi(ncplx, nproj, wfdp, proj, ncplx, npsi, wfd, psi, tolr, &
         work, scpr, cproj)
    call verify(all(scpr == 9_wp), "scpr mask pack")
    call verify(all(cproj(1, :, :) == 18_wp), "real(cproj) mask pack")
    call verify(all(cproj(2, :, :) == 0_wp), "imag(cproj) mask pack")

    call tolr_set_strategy(tolr, "MASK")
    call pr_dot_psi(ncplx, nproj, wfdp, proj, ncplx, npsi, wfd, psi, tolr, &
         work, scpr, cproj)
    call verify(all(scpr == 9_wp), "scpr mask")
    call verify(all(cproj(1, :, :) == 18_wp), "real(cproj) mask")
    call verify(all(cproj(2, :, :) == 0_wp), "imag(cproj) mask")

    call tolr_set_strategy(tolr, "KEYS_PACK")
    call pr_dot_psi(ncplx, nproj, wfdp, proj, ncplx, npsi, wfd, psi, tolr, &
         work, scpr, cproj)
    call verify(all(scpr == 9_wp), "scpr key pack")
    call verify(all(cproj(1, :, :) == 18_wp), "real(cproj) key pack")
    call verify(all(cproj(2, :, :) == 0_wp), "imag(cproj) key pack")

    call tolr_set_strategy(tolr, "KEYS")
    call pr_dot_psi(ncplx, nproj, wfdp, proj, ncplx, npsi, wfd, psi, tolr, &
         work, scpr, cproj)
    call verify(all(scpr == 9_wp), "scpr key")
    call verify(all(cproj(1, :, :) == 18_wp), "real(cproj) key")
    call verify(all(cproj(2, :, :) == 0_wp), "imag(cproj) key")

    call f_free(proj)
    call f_free(psi)
    call f_free(work)

    call deallocate_wfd_to_wfd(tolr)
    call deallocate_wfd(wfdp)
    call deallocate_wfd(wfd)
  end subroutine prDotPsi
  
  subroutine cprojDot(label)
    use liborbs_precisions
    use compression
    use dynamic_memory
    implicit none
    character(len = *), intent(out) :: label

    integer, parameter :: nproj = 3, npsi = 2
    integer, parameter :: ncplx = 1
    real(wp), dimension(ncplx, npsi, ncplx, nproj) :: scpr
    real(wp), dimension(ncplx, npsi, nproj) :: cproj1, cproj2
    real(wp) :: energy

    label = "Cproj . Cproj"

    cproj1 = 1_wp
    cproj2 = 2_wp
    scpr = 5_wp
    call cproj_dot(ncplx, nproj, ncplx, npsi, scpr, cproj1, cproj2, energy)
    call compare(energy, 60._wp, "energy")
  end subroutine cprojDot

  subroutine cprojPrPsi(label)
    use liborbs_precisions
    use compression
    use dynamic_memory
    implicit none
    character(len = *), intent(out) :: label

    type(wfd_to_wfd) :: tolr
    type(wavefunctions_descriptors) :: wfd, wfdp

    integer, parameter :: nproj = 2, npsi = 1
    integer, parameter :: ncplx = 2
    real(wp), dimension(:, :, :), allocatable :: psi, proj
    real(wp), dimension(:, :), allocatable :: work
    real(wp), dimension(ncplx, npsi, ncplx, nproj) :: scpr
    real(wp), dimension(ncplx, npsi, nproj) :: cproj

    label = "Cproj . Proj . Psi"

    call sampleTolr(tolr, wfd, wfdp)

    proj = f_malloc((/wfdp%nvctr_c + 7 * wfdp%nvctr_f, ncplx, nproj/), "proj")
    psi = f_malloc((/wfd%nvctr_c + 7 * wfd%nvctr_f, ncplx, npsi/), "psi")
    work = f_malloc((/wfd%nvctr_c + 7 * wfd%nvctr_f, ncplx * npsi/), "work")
    psi = 1_wp
    proj = 1_wp
    cproj = 2_wp

    call tolr_set_strategy(tolr, "MASK_PACK")
    call cproj_pr_p_psi(cproj, ncplx, nproj, wfdp, proj, ncplx, npsi, wfd, psi, tolr, &
         work, scpr)
    call compare(reshape(scpr, (/size(scpr)/)), (/ 2._wp, 2._wp, -2._wp, 2._wp, &
         2._wp, 2._wp, -2._wp, 2._wp /), "scpr")

    call tolr_set_strategy(tolr, "KEYS_PACK")
    call cproj_pr_p_psi(cproj, ncplx, nproj, wfdp, proj, ncplx, npsi, wfd, psi, tolr, &
         work, scpr)
    call compare(reshape(scpr, (/size(scpr)/)), (/ 2._wp, 2._wp, -2._wp, 2._wp, &
         2._wp, 2._wp, -2._wp, 2._wp /), "scpr")

    call f_free(proj)
    call f_free(psi)
    call f_free(work)

    call deallocate_wfd_to_wfd(tolr)
    call deallocate_wfd(wfdp)
    call deallocate_wfd(wfd)
  end subroutine cprojPrPsi

end program test_compression
