!test reproduce crash of pzgetri when N=leading dimension=200
!should be run on 5 MPI nodes
module globals
implicit none

  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 eye(m,n,context,mat,desc)
    implicit none
    integer, intent(in) :: m, n !row,col length
    integer, intent(in) :: context ! blacs context.
    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))
       print*, "eye: matrix mat is allocated for context < 0 m= ", m,"n=",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))
    print*, "eye: matrix mat is allocated for dim lm=",lm, "ln=",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)
  
    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 eye
  
  !======================================================================
  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
  
    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

  include 'mpif.h'   ! Use 8-byte intel mpi library

  ! this is the integer type for mpi control values 
  integer, parameter :: mpi_int_kind = selected_int_kind(12) !8-byte integer

  ! local variables
!  integer(mpi_int_kind) :: rank, ierr, nproc, version, subversion
  integer :: rank, ierr, nproc, version, subversion
!  integer(mpi_int_kind) :: mpi_int_size
  integer :: mpi_int_size   
  logical :: master
  integer :: icontext, gcontext, itmp, i, info

  complex(db), allocatable :: zz(:,:)
  integer ::  zz_desc(DLEN_)

  integer, parameter :: num_matrices = 4
  integer  :: mat_to_node(2,num_matrices)

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

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

  integer,parameter :: qr_size = 32 
  integer,parameter :: zz_size = 200 ! MPI_Recv error occurs
  !integer,parameter :: zz_size = 400 ! MPI_Recv error does not occur

  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)
  master = rank .eq. 0
#if defined TBMKL_DEBUG
  p_stage = 1
#endif

  write(stdout,*) 'rank is', rank
  if( master ) then
     call mpi_get_version(version,subversion,ierr)
     write(stdout,'(2(A,I1.1))') ' MPI Version : ', version, '.', subversion
  
     call mpi_type_size(MPI_INTEGER,mpi_int_size,ierr)
     write(stdout,'(A,I)') ' MPI Integer Size : ', mpi_int_size
     
     call mkl_get_version_string(mversion)
     write(stdout,*) mversion
#if defined TBMKL_DEBUG
   ! for debug with gdb: attach to process and type 'set var p_stage=2' in gdb
     do
         if (p_stage.eq.2) exit
     end do
#endif
  endif
  ! create global context
  call CreateContextFromRange(0,nproc-1,rank,nproc,gcontext)

  ! create local contexts. Matrix 1-3 are on single nodes, Matrix 4 is on two nodes
  mat_to_node(1,1) = 0; mat_to_node(2,1) = 0
  mat_to_node(1,2) = 1; mat_to_node(2,2) = 1
  mat_to_node(1,3) = 2; mat_to_node(2,3) = 2
  mat_to_node(1,4) = 3; mat_to_node(2,4) = 4
  icontext = -1
  do i = 1,num_matrices
     ! Use lapack for single-node matrices is use_lapack is true
     if ( use_lapack .and. (mat_to_node(2,i) - mat_to_node(1,i) + 1  .eq. 1) ) cycle !single-node matrix
     call CreateContextFromRange(mat_to_node(1,i),mat_to_node(2,i),rank,nproc,itmp)
     if ( rank .ge. mat_to_node(1,i) .and. rank .le. mat_to_node(2,i) ) icontext = itmp
  end do
  if ( icontext .lt. 0 ) write(stdout,*) 'Local context not assigned ', icontext
  
  ! invert a matrix
  if ( master ) write(stdout,*) 'Starting inverse'
  if ( master ) write(stdout,*) 'Inverse Matrix dimension', zz_size
  print*, " ----- before eye: input matrix - zz "
  print*, "                   zz_desc= ", zz_desc 
  call eye(zz_size,zz_size,gcontext,zz,zz_desc)
   print*, " ----- after call eye for matrix - zz"
   print*, "                   zz_desc= ", zz_desc
  call pzgetrf(zz_size, zz_size, zz, 1, 1, zz_desc,pivot, info)
  if ( master ) write(stdout,*) 'Done with LU'

  !if increase leading dimension test passes
  !zz_desc(9)=zz_size+32

  call pzgetri(zz_size, zz, 1, 1, zz_desc, pivot, work, lwork, iwork, liwork, info)
  if ( master ) write(stdout,*) 'Done with inverse'

  deallocate(zz)
  call mpi_finalize(ierr)

end program test

