module globals
  !use mpi ! uncomment for LP64 mpi library
implicit none

  include 'mpif.h'   ! Uncomment for ILP64 intel mpi library

  integer, parameter :: I4B = selected_int_kind( 8) !4-byte integer
  integer, parameter :: I8B = selected_int_kind(12) !8-byte integer

  ! this is the integer type for mpi control values 
  integer, parameter :: mpi_int_kind = I8B ! uncomment for ILP64 mpi library
  !integer, parameter :: mpi_int_kind = I4B  ! uncomment for LP64 mpi library

  integer, parameter :: db = kind(1.0d0)
  integer, parameter :: stdout = 6

  integer, parameter, public :: BLOCK_CYCLIC_2D = 1 
  integer, parameter, public :: DLEN_ = 9 
  integer, parameter, public :: DT_ = 1
  integer, parameter, public :: CTXT_ = 2
  integer, parameter, public :: M_ = 3
  integer, parameter, public :: N_ = 4
  integer, parameter, public :: MB_ = 5
  integer, parameter, public :: NB_ = 6
  integer, parameter, public :: RSRC_ = 7
  integer, parameter, public :: CSRC_ = 8
  integer, parameter, public :: LLD_ = 9

  contains

  !======================================================================
  ! Create a diagonal matrix with ones on the diagonal
  subroutine create_matrix(m,n,context,eye,mat,desc)
    implicit none
    integer, intent(in) :: m, n !row,col length
    integer, intent(in) :: context ! blacs context.
    logical, intent(in) :: eye ! diagonal = 1?
    complex(db), allocatable, intent(out) :: mat(:,:) ! matrix
    integer, intent(out) :: desc(9) ! scalapack descriptor
  
    integer :: lm, ln ! local row and col lengths
    integer :: nprow, npcol, myrow, mycol, info, i, j, gi, gj
    integer, parameter :: MB = 32
    integer, parameter :: NB = 32
  
    integer, external :: numroc, indxl2g

    if ( context .lt. 0 ) then ! not distributed
       allocate(mat(m,n))
       ! set diagonal to 1, 0 otherwise
       mat = cmplx(0._db,0._db,db)
       do i = 1, min(m,n)
          mat(i,i) = cmplx(1.0_db,0.0_db,db)
       end do
       return
    end if
  
    call blacs_gridinfo(context,nprow,npcol,myrow,mycol)
    lm = numroc(m,MB,myrow,0,nprow)
    ln = numroc(n,NB,mycol,0,npcol)
  
    allocate(mat(lm,ln))
    call descinit( desc, m, n, MB, NB, 0, 0, context, lm, info)
  
    ! set diagonal to 1, 0 otherwise
    mat = cmplx(0._db,0._db,db)
    if ( .not. eye ) return
  
    do j = 1, ln
        gj = indxl2g(j,NB,mycol,desc(CSRC_),npcol)
        do i = 1, lm
            gi = indxl2g(i,MB,myrow,desc(RSRC_),nprow)
            if ( gi .eq. gj ) mat(i,j) = cmplx(1.0_db,0.0_db,db)
        enddo
    end do
  
  end subroutine create_matrix
  
  !======================================================================
  subroutine CreateContextFromRange(nstart,nend,rank,totProcs,context)
    implicit none
    integer, intent(in) :: nstart !> mpi node start
    integer, intent(in) :: nend   !> mpi node end
    integer, intent(in) :: rank   !> node rank
    integer, intent(in) :: totProcs !> toal mpi processes
    integer, intent(out) :: context  !> the created context
    
    integer :: nprocs, nprow, npcol
    integer :: k, i, j, tmpcontxt
    integer :: imap(1,totProcs), nodes(totProcs) !single-row grid shape
  
    integer, external :: blacs_pnum

    nprocs =  nend - nstart + 1 ! number of processes in grid to create
  
    ! Set single-row grid shape 
    nprow = 1
    npcol = nprocs
  
    ! create temproray grid (if we don't and assign using system context, weird things may happen)
    call blacs_get( 0, 0, tmpcontxt ) ! get system context
    call blacs_gridinit( tmpcontxt, 'Row-major', 1, totProcs )
  
    ! Get the system process number of each node in the each context.
    ! Actually, this will only order nodes ascendingly if each context has a row major ordering and is ordered itself
    k = 1
    do i = nstart,nend
      nodes(k) = blacs_pnum(tmpcontxt, 0, i) ! this gets the process number at (0,k) within the process grid tmpcontxt
      k = k + 1
    end do
  
    ! create a row-major mapping
    k = 1
    do j = 1, npcol
       imap(1, j) = nodes(k) ! this gets the process number at (0,k) within the process grid tmpcontxt
       k = k + 1
    end do
    !write(stdout,*) imap
  
    call blacs_gridexit(tmpcontxt) ! Free temporary context
  
    !  ****   Apply the new mapping to form desired context ****
    call blacs_get( 0, 0, context ) ! get system context
    call blacs_gridmap( context, imap, nprow, nprow, npcol ) ! ictxt returns with the new grid context
  
  end subroutine CreateContextFromRange


end module globals
!======================================================================
program test
  use globals
  implicit none



  ! local variables
  integer(mpi_int_kind) :: rank_, ierr, nproc_, version, subversion
  integer(mpi_int_kind) :: mpi_int_size
  logical :: master
  integer :: icontextA, icontextB, gcontext, itmp, i, info

  integer :: rank, nproc
  integer(I4B) :: isize,v_,sv_

  complex(db), allocatable :: qq(:,:), xx(:,:), yy(:,:), zz(:,:)
  integer :: qq_desc(DLEN_), rr_desc(DLEN_)
  integer :: xx_desc(DLEN_), yy_desc(DLEN_), zz_desc(DLEN_)

  integer, parameter :: numA_matrices = 4
  integer, parameter :: numB_matrices = 2
  integer  :: matA_to_node(2,numA_matrices)
  integer  :: matB_to_node(2,numB_matrices)

  integer, parameter :: lwork = 25000, liwork = 25000, &
                        lpivot = 25000, ltau = 25000
  integer :: pivot(lpivot), iwork(liwork)
  complex(db) :: work(lwork), tau(ltau)

  ! IF use_lapack is .false., then no MPI_Recv error or hang occurs
  ! IF use_lapack is .true., then an MPI_Recv error or hang may occur for some matrix dimensions 
  logical, parameter :: use_lapack = .true.

  integer,parameter :: qr_size = 200 

  ! NO ERROR/HANG OCCURS UNLESS THE CALL TO PZUNGQR/UNZGQR IS MADE.

  !integer,parameter :: zz_size = 32   ! PZGEMM completes if use_lapack=.true.
  integer,parameter :: zz_size = 200  ! PZGEMM hangs if use_lapack=.true.
  !integer,parameter :: zz_size = 400 ! MPI_Recv error occurs if use_lapack=.true.

  character(len=198) :: mversion

  ! Initialize the MPI environment
  call mpi_init(ierr)
  call mpi_comm_rank(MPI_COMM_WORLD, rank_, ierr)
  call mpi_comm_size(MPI_COMM_WORLD, nproc_, ierr)
  rank = rank_; nproc = nproc_
  master = rank .eq. 0

  if( master ) then
     call mpi_get_version(version,subversion,ierr)
     write(stdout,'(2(A,I1.1))') ' MPI Version : ', version, '.', subversion

     call mpi_get_version(v_,sv_,ierr)
     write(stdout,'(2(A,I1.1))') ' MPI Version : ', v_, '.', sv_
  
     call mpi_type_size(MPI_INTEGER,mpi_int_size,ierr)
     write(stdout,'(A,I)') ' MPI Integer Size : ', mpi_int_size

     call mpi_type_size(MPI_INTEGER,isize,ierr)
     write(stdout,'(A,I)') ' MPI Integer Size : ', isize
     
     call mkl_get_version_string(mversion)
     write(stdout,*) mversion
     write(stdout,*) 'Number of nodes = ', nproc
  endif

  ! create global context
  call CreateContextFromRange(0,nproc-1,rank,nproc,gcontext)
  write(stdout,*) 'Global context is ', gcontext

  ! create local contexts. Matrix 1-4 are on single nodes,
  matA_to_node(1,1) = 0; matA_to_node(2,1) = 0
  matA_to_node(1,2) = 1; matA_to_node(2,2) = 1
  matA_to_node(1,3) = 2; matA_to_node(2,3) = 2
  matA_to_node(1,4) = 3; matA_to_node(2,4) = 4

  ! create local contexts. Matrix 1 is on a single node, Matrix 2 is on three nodes
  matB_to_node(1,1) = 0; matB_to_node(2,1) = 1
  matB_to_node(1,2) = 2; matB_to_node(2,2) = 4

  icontextA = -1
  do i = 1,numA_matrices
     ! Use lapack for single-node matrices is use_lapack is true
     if ( use_lapack .and. &
        (matA_to_node(2,i) - matA_to_node(1,i) + 1  .eq. 1) ) cycle !single-node matrix

     call CreateContextFromRange(matA_to_node(1,i),matA_to_node(2,i),rank,nproc,itmp)
     if ( rank .ge. matA_to_node(1,i) .and. rank .le. matA_to_node(2,i) ) icontextA = itmp
  end do
  if ( icontextA .lt. 0 ) write(stdout,*) 'Local context A not assigned ', icontextA

  icontextB = -1
  do i = 1,numB_matrices
     ! Use lapack for single-node matrices is use_lapack is true
     if ( use_lapack .and. &
        (matB_to_node(2,i) - matB_to_node(1,i) + 1  .eq. 1) ) cycle !single-node matrix

     call CreateContextFromRange(matB_to_node(1,i),matB_to_node(2,i),rank,nproc,itmp)
     if ( rank .ge. matB_to_node(1,i) .and. rank .le. matB_to_node(2,i) ) icontextB = itmp
  end do
  if ( icontextB .lt. 0 ) write(stdout,*) 'Local context B not assigned ', icontextB
  
  ! create diagonal matrices
  call create_matrix(qr_size,qr_size,icontextA,.true.,qq,qq_desc)

  if ( master ) write(stdout,*) 'Starting QR'
  if ( master ) write(stdout,*) 'QR Matrix dimension', qr_size
  if ( master ) write(stdout,*) 'Single node matrices use Lapack?', use_lapack
  write(stdout,'(A,2I5,A,I5)') ' QR local matrix size = ', size(qq,1), size(qq,2), ' on rank ', rank
  if ( icontextA .ge. 0 ) then
     call pzgeqrf(qr_size, qr_size, qq, 1, 1, qq_desc, tau, work, lwork, info)
     call pzungqr(qr_size, qr_size, qr_size,  qq, 1, 1, qq_desc, tau, work, lwork, info)    
  else
     call zgeqrf(qr_size, qr_size, qq, qr_size, tau, work, lwork, info)
     call zungqr(qr_size, qr_size, qr_size, qq, qr_size, tau, work, lwork, info)    
  end if
  call mpi_barrier(MPI_COMM_WORLD,ierr)
  if ( master ) write(stdout,*) 'Done with QR'

  ! invert a matrix
  if ( master ) write(stdout,*) 'Starting pzgemm'
  if ( master ) write(stdout,*) 'pzgemm matrix dimension', zz_size
  call create_matrix(zz_size,zz_size,gcontext,.false.,xx,xx_desc)
  call create_matrix(zz_size,zz_size,gcontext,.false.,yy,yy_desc)
  call create_matrix(zz_size,zz_size,gcontext,.false.,zz,zz_desc)
  write(stdout,'(A,2I5,A,I5)') ' PZGEMM local matrix size = ', size(qq,1), size(qq,2), ' on rank ', rank
  call pzgemm('N','N',zz_size,zz_size,zz_size,(1._db,0._db), xx,1,1,xx_desc, &
           yy,1,1,zz_desc, (0._db,0._db), zz,1,1,zz_desc)

  if ( master ) write(stdout,*) 'Done with pzgemm'

  deallocate(qq,xx,yy,zz)
  call mpi_finalize(ierr)

end program test

