! The program is adpated from NElIB ScaLAPACK user guide
! Description: http://netlib.org/scalapack/slug/node28.html#SECTION04230000000000000000
! Source code: http://netlib.org/scalapack/slug/node27.html#SECTION04220000000000000000
! A better description: http://acts.nersc.gov/scalapack/hands-on/exercise4/description.html
!
! The major differences between this code and the NETLIB example are
!  1. This code is written in f90, while the original is in f77
!  2. This code initializes matrices A and B on one process and 
!     sends sub-matrices to all processes. Details on the 
!     sub-mtrices received in each process can be found in the
!     description (see the link above).
!  3. The original code use predefined 2x3 grid. Consequently, it 
!     only runs on 6 cores. This code can run on arbitrary number
!     of nodes by calling MPI_DIMS_CREATE
!
!  Compile the code with
!      mpiifort -mkl=cluster mypdgesvdriver.f90 -o mypdgesvdriver.exe
!
! Supercomputing Facility 
! Texas A&M University
! Programmer: Ping Luo

program pdgesvdriver
use mpi
implicit none

! Distributed matrix parameters: matrix size is MxN (M=N), 
! number of RHS is NRHS, block sizes are MB ad NB
integer,parameter :: M=9, N=9, NRHS=1, MB=2, NB=2
! Distributed matrices
real*8,allocatable :: A(:,:), B(:,:)
!
! Local Matrices parameters
integer :: Mlocal, Nlocal, NRHS_local
! Local matrices
real*8,allocatable :: Alocal(:,:), Blocal(:,:) 
!
! Temprary workspace and its sizes
integer :: Mwork, Nwork
real*8,allocatable :: work(:,:)
!
! Descriptors for the distributed matrices A and B
! Avaialbe on each process; contains information about the distributed
! global matrix and the local matrix
integer,parameter :: DLEN_=9
real*8 :: desca(DLEN_), descb(DLEN_)
!
! Local variables for MPI and BLACS that handle communication
integer,parameter :: NDIMS = 2
integer,dimension(1:NDIMS) :: dims        ! for 2D dimensions
integer :: myrank, nprocs, nprow, npcol
integer :: myrow, mycol                   ! equiv. to dims(ndims)
integer,parameter :: ROW_SRC=0, COL_SRC=0 ! The root process that 
                                          ! serves as the origin of
                                          ! a couple of operations
integer :: context                        ! BLACS context
!
! Local variables for ScaLAPACK subroutine PDGESV
real*8,allocatable :: ipiv(:)
!
! Compute the number of rows or columns of a block-cyclically distributed
! matrix contained in a process
integer,external :: numroc
!
! Return status from subroutine calls
integer :: info
!
integer :: i, j 



! Define process grid
call BLACS_PINFO(myrank, nprocs)      ! can be replaced with
                                      ! call MPI_INIT(info)
                                      ! call MPI_COMM_RANK(MPI_COMM_WORLD, myrank, info)
                                      ! call MPI_COMM_SIZE(MPI_COMM_WORLD, nprocs, info)
call MPI_DIMS_CREATE(nprocs, NDIMS, dims, info)
nprow = dims(1)
npcol = dims(2)
call BLACS_GET(-1, 0, context)
call BLACS_GRIDINIT(context, 'Row-major', nprow, npcol)
call BLACS_GRIDINFO(context, nprow, npcol, myrow, mycol)
!
Mlocal     = numroc(M,    MB, myrow, ROW_SRC, nprow)
Nlocal     = numroc(N,    NB, mycol, COL_SRC, npcol)
NRHS_local = numroc(NRHS, NB, mycol, COL_SRC, npcol)
allocate(Alocal(Mlocal,Nlocal),Blocal(Mlocal,NRHS_local))
Alocal=0.0d0
Blocal=0.0d0
!
print *, myrank, myrow, mycol
print *, myrank, mlocal, nlocal, nrhs_local 

if (myrow .eq. ROW_SRC .and. mycol .eq. COL_SRC) then
! Generate the global matrices in the root process
    allocate(A(M,N), B(M, NRHS))
    A = reshape((/19,-19,-19,-19,-19,-19,-19,-19,-19,  &
                  3,  3, -3, -3, -3, -3, -3,  3, -3,   &
                  1,  1,  1, -1, -1, -1, -1, -1, -1,   &
                  12, 12, 12, 12,-12,-12,-12,-12,-12,  &
                  1,  1,  1,  1,  1, -1, -1, -1, -1,   &
                  16, 16, 16, 16, 16, 16,-16,-16,-16,  &
                  1,  1,  1,  1,  1,  1,  1, -1, -1,   &
                  3,  3,  3,  3,  3,  3,  3,  3, -3,   &
                  11, 11, 11, 11, 11, 11, 11, 11, 11/),&
                  shape(A))
!
    B = reshape((/0, 0, 1, 0, 0, 0, 0, 0, 0/), shape(B))
! 
    do i=0, nprow-1
      do j=0, npcol-1

        ! skip if (i, j) is the elected process
        if (i .eq. ROW_SRC .and. j .eq. COL_SRC) cycle

        Mwork = numroc(M, MB, i, ROW_SRC, nprow)
        if (Mwork .ne. 0) then    
 
          Nwork = numroc(N, NB, j, COL_SRC, npcol)
          if (Nwork .ne. 0) then
            allocate(work(Mwork,Nwork))
            call create_local_ary_2D(nprow, npcol, i, j, A, MB, NB, work)
            call DGESD2D(context, Mwork, Nwork, work, Mwork, i, j)
            deallocate(work)
          endif

          Nwork = numroc(NRHS, NB, j, COL_SRC, npcol)
          if (Nwork .ne. 0) then         
            allocate(work(Mwork,Nwork))
            call create_local_ary_2D(nprow, npcol, i, j, B, MB, NB, work)
            call DGESD2D(context, Mwork, Nwork, work, Mwork, i, j)
            deallocate(work)
          endif
      
        endif
      enddo
    enddo

    ! Create local matrices for elected process
    call create_local_ary_2D(nprow, npcol, ROW_SRC, COL_SRC, A, MB, NB, Alocal)
    call create_local_ary_2D(nprow, npcol, ROW_SRC, COL_SRC, B, MB, NB, Blocal)
 
    ! The global matrices have been properly distributed to each
    ! process. They are not needed and can be destroyed.
    deallocate(A, B)
else
    if (Mlocal .ne. 0 .and. Nlocal     .ne. 0) then
        call DGERV2D(context, Mlocal, Nlocal,     Alocal, Mlocal, ROW_SRC, COL_SRC)
    endif
    if (Mlocal .ne. 0 .and. NRHS_local .ne. 0) then
        call DGERV2D(context, Mlocal, NRHS_local, Blocal, Mlocal, ROW_SRC, COL_SRC)
    endif
endif


! Initialize the array descriptiors for matrix A and B
call DESCINIT(desca, M, N,    MB, NB, ROW_SRC, COL_SRC, context, Mlocal, info)
call DESCINIT(descb, M, NRHS, MB, NB, ROW_SRC, COL_SRC, context, Mlocal, info)

! Call the solver
allocate(ipiv(Mlocal+MB))
call PDGESV(M, NRHS, Alocal, 1, 1, desca, ipiv, Blocal, 1, 1, descb, info)

if (NRHS_local .ne. 0) print *,myrow, mycol, Blocal

deallocate(Alocal, Blocal, ipiv)
10 continue

call BLACS_GRIDEXIT(context)
call BLACS_EXIT(0)        ! or call MPI_FINALIZE(info)


contains 
! matrix is distributed in a 2D block cyclic way
subroutine create_local_ary_2D(nprow, npcol, myrow, mycol, A, MB, NB, Alocal)
implicit none
integer,intent(in)  :: nprow, npcol, myrow, mycol, MB, NB
real*8 ,intent(in)  :: A(:,:)                 ! assumed shape array
real*8 ,intent(out) :: Alocal(:,:)            ! assumed shape array

integer :: M, N
integer :: i, j, gi, gj, li, lj


M = size(A, 1)
N = size(A, 2)

lj = 0
do j=mycol*NB+1, N, NB*npcol
  do gj=j, min(j+NB-1, N)
    li = 0
    lj = lj + 1
    do i=myrow*MB+1, M, MB*nprow
      do gi=i, min(i+MB-1, M)
        li = li+1
        Alocal(li,lj) = A(gi, gj)
      enddo
    enddo
  enddo
enddo

end subroutine create_local_ary_2D

end program pdgesvdriver

