program diagonal


use mpi

implicit none

    INTEGER :: rank, size, ierr
    INTEGER :: i,j
    INTEGER, DIMENSION (:,:), ALLOCATABLE :: matrix
    INTEGER :: diag

    !$ MPI Initialization 
    CALL MPI_Init(ierr)
    CALL MPI_Comm_size(MPI_COMM_WORLD,size,ierr)
    CALL MPI_Comm_rank(MPI_COMM_WORLD,rank,ierr)

    !$ Matrix initialization 
    ALLOCATE (matrix(size,size))

    DO j=1,size
        DO i=1,size
            IF (i.EQ.j) THEN  
            matrix(i,j) = rank
            ELSE
            matrix(i,j) = 0
            ENDIF
        ENDDO
    ENDDO

    !$ Print rank 0 matrix (should be filled with 0s)
    IF (rank.EQ.0) THEN
       WRITE (*,*) 'Rank 0 matrix before communication:'
       DO i=1,size
           PRINT *, matrix(i,:)
       ENDDO
    ENDIF 

    !$ Diagonal datatype vector creation:
    ! - use the function MPI_Type_vector to obtain the 'diag' type handle
	! - do not forget to commit the new data type
		---
		
    !$ Communication: rank 0 gathers all the diagonals from the other ranks and stores them in the column corresponding to the
    !$ sending rank. 
    ! - use proper sending and receiving data types
		---

    !$ Print rank 0 matrix after communication (each element should be its column number)
    IF (rank.EQ.0) THEN
       WRITE (*,*) ' '
       WRITE (*,*) 'Rank 0 matrix after communication:'
       DO i=1,size
           PRINT *, matrix(i,:)
       ENDDO
    ENDIF

    !$ Remember to free the datatype!
		---
    DEALLOCATE(matrix)

    CALL MPI_Finalize(ierr)

end program diagonal