!+
! Module sc_track_mod_mpi
!
! This module consolidates all of the native MPI calls and code in the sc_track_mod_mpi
! program.  
!-
MODULE sc_track_mod_mpi

USE precision_def
USE bmad !needed for err_exit

IMPLICIT NONE
INCLUDE 'mpif.h'

!status and mpistatus are used to send and receive codes from the MPI daemon.
INTEGER status(MPI_STATUS_SIZE),mpistatus(MPI_STATUS_SIZE)

!myrank is populated by the mympi_initialize and contains the rank number.
!nslave contains the number of workers in the cluster.  Typically #nodes-1
INTEGER myrank, nslave
INTEGER MPI_DATATYPE_COORD

!Subroutines
PUBLIC mympi_initialize         !how the executible announces itself to the MPI daemon
PUBLIC mympi_bmad_parser        !front end to bmad parser
PUBLIC mympi_shutdown
PRIVATE mympi_register_derived_types

!Data
PRIVATE myrank
PRIVATE status, mpistatus

CONTAINS

!+
! Subroutine mympi_initialize(master)
!
! First mpi call made by a program.  This subroutine calls mpi_init, which announces
! the executible to the MPI daemon.  mpi_comm_rank asks the MPI daemon to tell the executible
! what its rank number is.  One node is given rank 0 and becomes the master.  All other nodes
! receive numbers counting up from 1 and become workers.
! mympi_register_derived_types tells the MPI daemon about the structures we will be sending.
!
! Input:
!   None
! Output:
!   master:  LOGICAL, INTENT(OUT): Set to true if master, false otherwise.
!-
SUBROUTINE mympi_initialize(spokesnode)
  LOGICAL, INTENT(OUT) :: spokesnode
  INTEGER i
  INTEGER mpierr
  INTEGER cluster_size

  CALL mpi_init(mpierr)                             ! Introduce yourself to the MPI daemon
  CALL mpi_comm_rank(MPI_COMM_WORLD,myrank,mpierr)  ! Get your rank number, store in myrank.  Master is rank 0.
  CALL mympi_register_derived_types()               ! Tell the daemon about derived types that will be 
                                                    ! shared among the cluster
  DO i=1,MPI_STATUS_SIZE
    status(i) = 1
    mpistatus(i) = 1
  ENDDO

  !Check that cluster has at least two nodes
  IF(myrank .eq. 0) THEN
    CALL mpi_comm_size(MPI_COMM_WORLD,cluster_size,mpierr)
    nslave=cluster_size-1
    IF(nslave .eq. 0) THEN
      WRITE(*,*) "NOTICE: only one node found in cluster."
!      STOP
    ENDIF
  ENDIF

  !node zero is the spokesnode
  IF(myrank .eq. 0) THEN
    spokesnode=.true.
  ELSE
    spokesnode=.false.
  ENDIF

END SUBROUTINE mympi_initialize

!+
! Subroutine mympi_get_rank(myrank,numnodes,spokesnode)
!
! Front-end to mpi_comm_rank.  Returns the rank in the mpi cluster of the calling process.
! If rank is zero, spokesnode is true, otherwise it is false.
!
! Input:
!        none
! Output:
!        myrank     -- INTEGER: rank of calling process
!        spokesnode -- LOGICAL: true if rank 0
!-
SUBROUTINE mympi_get_rank(myrank,numnodes,spokesnode)
  INTEGER myrank
  INTEGER numnodes
  LOGICAL spokesnode
  INTEGER mpierr

  CALL mpi_comm_rank(MPI_COMM_WORLD,myrank,mpierr)
  CALL mpi_comm_size(MPI_COMM_WORLD,numnodes,mpierr)

  IF(myrank .eq. 0) THEN
    spokesnode = .true.
  ELSE
    spokesnode = .false.
  ENDIF
END SUBROUTINE mympi_get_rank

!+
! Subroutine mympi_bmad_parser(lat_file,lat)
!
! Front-end to the bmad_parser subroutine.  If node 0, then bmad_parser is called immediately.
! Otherwise, wait until node 0 call to bmad_parser is complete.
!
! This subroutine avoids a problem where multiple processes may try to write a new digested file
! simultaneously.
!
! Input:
!   lat_file    -- CHARACTER(130), INTENT(IN): filename of lattice
! Output:
!   lat         -- TYPE(lat_struct): parsed lattice
!-
SUBROUTINE mympi_bmad_parser(lat_file, lat)
  USE bmad
  CHARACTER(130), INTENT(IN) :: lat_file
  TYPE(lat_struct) lat
  INTEGER mpierr

  ! Node 0 calls bmad parser first, which will write the digested file if it does not exist.
  ! After completing bmad parser, it reaches the barrier, signaling the other nodes to call bmad parser.
  IF(myrank .eq. 0) THEN
    CALL fullfilename(lat_file, lat_file)
    CALL bmad_parser(lat_file, lat)
    CALL mpi_barrier(MPI_COMM_WORLD,mpierr)  ! Each process waits at barrier till all processes have reached the barrier
  ELSE
    CALL mpi_barrier(MPI_COMM_WORLD,mpierr)  ! Each process waits at barrier till all processes have reached the barrier
    CALL fullfilename(lat_file, lat_file)
    CALL bmad_parser(lat_file, lat)
  ENDIF

END SUBROUTINE mympi_bmad_parser

!+
! Subroutine mympi_register_derived_types()
!
! Because this program uses structures, rather than native data types, for communication,
! the MPI daemon needs to be told about form of the structure.
! Later when mpi_send and mpi_recv are called, the integer parameters ApertureJob_mpi_id
! and ApertureResult_mpi_id are used to tell MPI which data type is being sent.
!
! Input:
!   None
! Output:
!   None
!-
SUBROUTINE mympi_register_derived_types()
  INTEGER mpierr
  INTEGER oldtypes(0:1),blockcounts(0:1),offsets(0:1),extent

  !Setup mpi type for coord_struct
  offsets(0)=0
  oldtypes(0)=MPI_DOUBLE_PRECISION
  blockcounts(0)=6
  CALL mpi_type_struct(1,blockcounts,offsets,oldtypes,MPI_DATATYPE_COORD,mpierr)
  CALL mpi_type_commit(MPI_DATATYPE_COORD,mpierr)
END SUBROUTINE mympi_register_derived_types

!+
! Subroutine mympi_shutdown()
!
! This routine prepares for program termination by calling mpi_finalize, which disconnects
! the process from the MPI daemon.
!
! Input:
!   None
! Output:
!   None
!-
SUBROUTINE mympi_shutdown()
  INTEGER mpierr
  CALL mpi_finalize(mpierr)
END SUBROUTINE mympi_shutdown

END MODULE sc_track_mod_mpi