PROGRAM bench_alltoallw

    USE, INTRINSIC:: iso_c_binding
    USE mpi

    IMPLICIT NONE

    INTEGER                                     :: N
    INTEGER                                     :: ix, iy, iz
    INTEGER                                     :: mpi_err, rank, nproc
    INTEGER                                     :: dtype, dtype_row, dtype_mat
    INTEGER                                     :: sizeofdtype
    INTEGER, DIMENSION(:), ALLOCATABLE          :: counts_send, counts_recv,&
                                                 & displs_send, displs_recv,&
                                                 & dtypes_send, dtypes_recv
    REAL(KIND=8), DIMENSION(:,:,:), ALLOCATABLE :: data_in, data_out
    CHARACTER(LEN=200)                          :: tmp_str
    INTEGER                                     :: nwarmup, nrepeats, irepeat
    REAL(KIND=8), DIMENSION(2)                  :: timers

    ! Parameters
    nrepeats = 10 ! Number of runs that will be timed
    nwarmup = 4 ! Number of runs to warm up the network
    dtype = MPI_DOUBLE_PRECISION

    ! MPI initialization
    CALL MPI_INIT(mpi_err)
    CALL MPI_COMM_RANK(MPI_COMM_WORLD, rank, mpi_err)
    CALL MPI_COMM_SIZE(MPI_COMM_WORLD, nproc, mpi_err)

    ! Read array size from command line
    IF(COMMAND_ARGUMENT_COUNT().NE.1) THEN
        WRITE(*,*) 'Error, incorrect number of command-line arguments'
        WRITE(*,*) 'Usage: ./bench N'
        CALL EXIT(1)
    END IF
    CALL GET_COMMAND_ARGUMENT(1, tmp_str)
    READ(tmp_str, *) N

    ! Initialize arrays
    ALLOCATE(data_in(N, N, N))
    ALLOCATE(data_out(N, N, N))
    data_out(:,:,:) = 0.0d0
    CALL RANDOM_NUMBER(data_in)
!    DO iz=1,N
!    DO iy=1,N
!    DO ix=1,N
!        data_in(ix,iy,iz) = ix-1 + N*(iy-1) + N*N*(iz-1)
!        WRITE(*,'(A5,3I5,F12.2)') 'IN', ix, iy, iz, data_in(ix,iy,iz)
!    END DO
!    END DO
!    END DO

    ! Initialize count, displacement, and data type arrays
    ALLOCATE(counts_send(nproc), counts_recv(nproc))
    ALLOCATE(displs_send(nproc), displs_recv(nproc))
    ALLOCATE(dtypes_send(nproc), dtypes_recv(nproc))
 
    ! Get size of base data type in bytes
    CALL MPI_TYPE_EXTENT(dtype, sizeofdtype, mpi_err)

    ! Sender data type
    !
    ! Collect all elements along y into a row
    CALL MPI_Type_vector(N, 1, N, dtype, dtype_row, mpi_err)
    ! Collect all rows into a matrix
    !CALL MPI_Type_hvector(N, 1, sizeofdtype, dtype_row, dtypes_send(1),&
    !                    & mpi_err)
    CALL MPI_Type_hvector(N, 1, sizeofdtype, dtype_row, dtype_mat,&
                        & mpi_err)
    ! Collect all matrices into the full 3D array
    CALL MPI_Type_hvector(N, 1, N*N*sizeofdtype, dtype_mat, dtypes_send(1),&
                         & mpi_err)
    CALL MPI_Type_commit(dtypes_send(1), mpi_err)
    counts_send(1) = 1
    displs_send(1) = 0

    ! Receiver data type
    !
    ! All elements are contiguous at the receiver side
    CALL MPI_Type_vector(N*N*N, 1, 1, dtype, dtypes_recv(1), mpi_err)
    CALL MPI_Type_commit(dtypes_recv(1), mpi_err)
    counts_recv(1) = 1
    displs_recv(1) = 0

    ! Send data back and forth

    ! Use Fortran transpose
    timers(:) = 0.0d0
    DO irepeat=1,nrepeats
        IF (irepeat.gt.nwarmup) timers(1) = timers(1) - MPI_WTIME()
        DO iz=1,N
            data_out(:,:,iz) = transpose(data_in(:,:,iz))
        END DO
        IF (irepeat.gt.nwarmup) timers(1) = timers(1) + MPI_WTIME()
        IF (irepeat.gt.nwarmup) timers(2) = timers(2) - MPI_WTIME()
        DO iz=1,N
            data_in(:,:,iz) = transpose(data_out(:,:,iz))
        END DO
        IF (irepeat.gt.nwarmup) timers(2) = timers(2) + MPI_WTIME()
    END DO
    IF (rank.eq.0) WRITE(*,'(A20,A5,I6,3(A5,F12.6))') &
        & 'TRANSPOSE', 'N', N,&
        & 'FWD', timers(1) / (nrepeats - nwarmup),&
        & 'BCK', timers(2) / (nrepeats - nwarmup),&
        & 'TOT', (timers(1) + timers(2)) / (nrepeats - nwarmup)

    ! Use Alltoall
    timers(:) = 0.0d0
    DO irepeat=1,nrepeats
        IF (irepeat.gt.nwarmup) timers(1) = timers(1) - MPI_WTIME()
        CALL MPI_Alltoall(data_in, 1, dtypes_send(1),&
                        & data_out, 1, dtypes_recv(1), MPI_COMM_WORLD, mpi_err)
        IF (irepeat.gt.nwarmup) timers(1) = timers(1) + MPI_WTIME()
        IF (irepeat.gt.nwarmup) timers(2) = timers(2) - MPI_WTIME()
        CALL MPI_Alltoall(data_out, 1, dtypes_recv(1),&
                        & data_in, 1, dtypes_send(1), MPI_COMM_WORLD, mpi_err)
        IF (irepeat.gt.nwarmup) timers(2) = timers(2) + MPI_WTIME()
    END DO
    IF (rank.eq.0) WRITE(*,'(A20,A5,I6,3(A5,F12.6))') &
        & 'ALLTOALL', 'N', N,&
        & 'FWD', timers(1) / (nrepeats - nwarmup),&
        & 'BCK', timers(2) / (nrepeats - nwarmup),&
        & 'TOT', (timers(1) + timers(2)) / (nrepeats - nwarmup)

    ! Use Alltoallw
    timers(:) = 0.0d0
    DO irepeat=1,nrepeats
        IF (irepeat.gt.nwarmup) timers(1) = timers(1) - MPI_WTIME()
        CALL MPI_Alltoallw(data_in,&
                         & counts_send,&
                         & displs_send,&
                         & dtypes_send,&
                         & data_out,&
                         & counts_recv,&
                         & displs_recv,&
                         & dtypes_recv,&
                         & MPI_COMM_WORLD, mpi_err)
        IF (irepeat.gt.nwarmup) timers(1) = timers(1) + MPI_WTIME()
        IF (irepeat.gt.nwarmup) timers(2) = timers(2) - MPI_WTIME()
        CALL MPI_Alltoallw(data_out,&
                         & counts_send,&
                         & displs_send,&
                         & dtypes_send,&
                         & data_in,&
                         & counts_recv,&
                         & displs_recv,&
                         & dtypes_recv,&
                         & MPI_COMM_WORLD, mpi_err)
        IF (irepeat.gt.nwarmup) timers(2) = timers(2) + MPI_WTIME()
    END DO
    IF (rank.eq.0) WRITE(*,'(A20,A5,I6,3(A5,F12.6))') &
        & 'ALLTOALLW', 'N', N,&
        & 'FWD', timers(1) / (nrepeats - nwarmup),&
        & 'BCK', timers(2) / (nrepeats - nwarmup),&
        & 'TOT', (timers(1) + timers(2)) / (nrepeats - nwarmup)

!    DO iz=1,N
!    DO ix=1,N
!    DO iy=1,N
!        write(*,'(A3,3I5,F12.2)') 'OUT', ix, iy, iz, data_out(iy,ix,iz)
!    END DO
!    END DO
!    END DO

    ! Clean up
    DEALLOCATE(data_in, data_out)
    DEALLOCATE(counts_recv, displs_recv, dtypes_recv)
    DEALLOCATE(counts_send, displs_send, dtypes_send)
    CALL MPI_Finalize(mpi_err)

END PROGRAM bench_alltoallw
