spanAverage

!****************************************************************************************
!****************************************************************************************
! MODULE types_vars
! Sets precision and common constants
MODULE types_vars
  USE mpi
  IMPLICIT NONE

  ! Symbolic names for kind types of 4-, 2- and 1-byte integers:
  INTEGER, PARAMETER :: I4B = SELECTED_INT_KIND(9)
  INTEGER, PARAMETER :: I2B = SELECTED_INT_KIND(4)
  INTEGER, PARAMETER :: I1B = SELECTED_INT_KIND(2)
  ! Symbolic names for kind types of single- and double-precison reals
  INTEGER, PARAMETER :: SP = KIND(1.0)
  INTEGER, PARAMETER :: DP = KIND(1.0D0)
  ! Symbolic names for kind types of single- and double-precison complex
  INTEGER, PARAMETER :: SPC = KIND((1.0,1.0))
  INTEGER, PARAMETER :: DPC = KIND((1.0D0,1.0D0))
  ! Symbolic name for kind type of default logical
  INTEGER, PARAMETER :: LOGIC = KIND(.true.)
  ! Frequently used mathematical constants (with precision to spare)
  REAL(DP), PARAMETER :: zero  = 0.0_dp
  REAL(DP), PARAMETER :: half  = 0.5_dp
  REAL(DP), PARAMETER :: one   = 1.0_dp
  REAL(DP), PARAMETER :: two   = 2.0_dp
  REAL(DP), PARAMETER :: three = 3.0_dp
  REAL(DP), PARAMETER :: four  = 4.0_dp
  REAL(DP), PARAMETER :: pi    = 3.141592653589793238462643383279502884197_dp
  REAL(DP), PARAMETER :: pio2  = 1.57079632679489661923132169163975144209858_dp
  REAL(DP), PARAMETER :: twopi = 6.283185307179586476925286766559005768394_dp

END MODULE types_vars

!****************************************************************************************
!****************************************************************************************
! MODULE variables
! Contains subroutines to set variables, allocate, and deallocate memory
MODULE variables
  USE types_vars
  USE mpi
  IMPLICIT NONE

  ! MPI variables
  INTEGER :: ierr, myid, status(MPI_STATUS_SIZE)
  INTEGER :: nprocs, nprocs_x, nprocs_y, nprocs_z
  INTEGER :: XVEC_TYPE, YCONT_TYPE
  INTEGER :: xsrc, xdst, ysrc, ydst, zsrc, zdst
  INTEGER :: comm3d, comm3d_rank
  LOGICAL :: isperiodic(3), reorder
  INTEGER :: dims(3), coords(3)
  INTEGER, ALLOCATABLE, DIMENSION(:) :: all_coords
  ! Grid
  INTEGER :: nx, ny, nz
  REAL(DP) :: xmin, xmax, ymin, ymax, zmin, zmax
  REAL(DP) :: dx, dy, dz
  REAL(DP), ALLOCATABLE, DIMENSION(:,:,:) :: x, y, z
  ! Buffer points
  INTEGER :: bufx, bufy, bufz
  ! Solution
  REAL(DP), ALLOCATABLE, DIMENSION(:,:,:) :: u, u_init
  ! Spanwise average stuff
  REAL(DP), ALLOCATABLE, DIMENSION(:,:) :: cpu_x
  REAL(DP), ALLOCATABLE, DIMENSION(:,:) :: u_2d_x_avg

CONTAINS

  !****************************************************************************************
  ! SUBROUTINE memalloc
  ! Allocates memory for needed variables
  SUBROUTINE memalloc
    IMPLICIT NONE

    ! Cartesian grid
    ALLOCATE(x(0-bufx:nx+bufx,0-bufy:ny+bufy,0-bufz:nz+bufz))
    ALLOCATE(y(0-bufx:nx+bufx,0-bufy:ny+bufy,0-bufz:nz+bufz))
    ALLOCATE(z(0-bufx:nx+bufx,0-bufy:ny+bufy,0-bufz:nz+bufz))
    ! Solution variable
    ALLOCATE(u     (0-bufx:nx+bufx,0-bufy:ny+bufy,0-bufz:nz+bufz))
    ALLOCATE(u_init(0:nx,0:ny,0:nz))
    ! Spanwise average
    ALLOCATE(cpu_x(0:nprocs_x-1,0:nprocs_z-1))
    ALLOCATE(u_2d_x_avg(0:nx,0:ny))

  END SUBROUTINE memalloc

  !****************************************************************************************
  ! SUBROUTINE dealloc
  ! Deallocates memory
  SUBROUTINE dealloc
    IMPLICIT NONE

    DEALLOCATE(x)
    DEALLOCATE(y)
    DEALLOCATE(z)
    DEALLOCATE(u)
    DEALLOCATE(u_init)
    DEALLOCATE(cpu_x)
    DEALLOCATE(u_2d_x_avg)

  END SUBROUTINE dealloc

  !****************************************************************************************

END MODULE variables

!****************************************************************************************
!****************************************************************************************
! MODULE io
! Contains subroutines for input and output
MODULE io
  USE types_vars
  USE variables
  USE mpi
  IMPLICIT NONE

CONTAINS

  !****************************************************************************************
  ! SUBROUTINE input
  ! Reads in variables from input file
  SUBROUTINE input
    IMPLICIT NONE

    OPEN(unit=1,file='input.txt')
    READ(1,*) nx
    READ(1,*) ny
    READ(1,*) nz
    READ(1,*) xmin
    READ(1,*) xmax
    READ(1,*) ymin
    READ(1,*) ymax
    READ(1,*) zmin
    READ(1,*) zmax
    READ(1,*) nprocs
    READ(1,*) nprocs_x
    READ(1,*) nprocs_y
    READ(1,*) nprocs_z
    CLOSE(1)

    bufx = 1
    bufy = 1
    bufz = 1

  END SUBROUTINE input

  !****************************************************************************************
  ! SUBROUTINE tecplot_output
  ! Outputs data in tecplot ASCII readable format
  SUBROUTINE tecplot_output
    IMPLICIT NONE
    INTEGER :: i, j, k, ii
    CHARACTER(len=24) :: file_name, file_name2

    IF (myid < 10) THEN
       WRITE(file_name, 400) 'solution/solution_', myid, '.dat'
    ELSE IF (myid .GE. 10) THEN
       WRITE(file_name, 401) 'solution/solution_', myid, '.dat'
    END IF

400 FORMAT(A18, I1, A4)
401 FORMAT(A18, I2, A4)

    OPEN(unit=2, file = file_name, status='replace')
    WRITE(2,'(A)') 'VARIABLES = "x", "y", "z", "u_init"'
    WRITE(2,*) 'ZONE I=',nx+1,', J=',ny+1,', K=',nz+1,', ZONETYPE=ORDERED,'
    WRITE(2,*) 'DATAPACKING=POINT, SOLUTIONTIME=0.0'
    DO k = 0,nz
       DO j = 0,ny
          DO i = 0,nx
             WRITE(2,*) x(i,j,k), y(i,j,k), z(i,j,k), u_init(i,j,k)
          END DO
       END DO
    END DO

    CLOSE(2)

    DO ii = 0,nprocs_x-1
       IF (myid .EQ. cpu_x(i,0)) THEN
          IF (myid < 10) THEN
             WRITE(file_name2, 400) 'solution/solution2d_', myid, '.dat'
          ELSE IF (myid .GE. 10) THEN
             WRITE(file_name2, 401) 'solution/solution2d_', myid, '.dat'
          END IF

          OPEN(unit=3, file = file_name2, status='replace')
          WRITE(3,'(A)') 'VARIABLES = "x", "y", "u_2d"'
          WRITE(3,*) 'ZONE I=',nx+1,', J=',ny+1,', ZONETYPE=ORDERED,'
          WRITE(3,*) 'DATAPACKING=POINT, SOLUTIONTIME=0.0'
          DO j = 0,ny
             DO i = 0,nx
                WRITE(3,*) x(i,j,0), y(i,j,0), u_2d_x_avg(i,j)
             END DO
          END DO

          CLOSE(3)
       END IF
    END DO

  END SUBROUTINE tecplot_output

END MODULE io

!****************************************************************************************
!****************************************************************************************
! MODULE mpi_wrapper
! Contains subroutines needed for MPI
MODULE mpi_wrapper
  USE types_vars
  USE variables
  USE io
  USE mpi
  IMPLICIT NONE

CONTAINS

  !****************************************************************************************
  ! SUBROUTINE init_parallel_environment
  ! Initializes the MPI environment and sets the domain decomposition up
  SUBROUTINE init_parallel_environment
    IMPLICIT NONE

    CALL MPI_INIT(ierr)                               ! initializes MPI environment
    CALL MPI_COMM_RANK(MPI_COMM_WORLD, myid, ierr)    ! gets the current rank, stores in myid
    CALL input                                        ! reads input file (so the number of points will be for each process!)
    CALL MPI_COMM_SIZE(MPI_COMM_WORLD, nprocs, ierr)  ! gets the number of processes, stores in nprocs (which is already read from input)
    CALL create_3d_communicator                       ! calls subroutine that sets up source/destination for each process
    CALL memalloc                                     ! allocate memory for each process
    CALL MPI_BARRIER(MPI_COMM_WORLD, ierr)            ! wait for each process in this communicator to finish

  END SUBROUTINE init_parallel_environment

  !****************************************************************************************
  ! SUBROUTINE create_3d_communicator
  ! Creates the source and destination for each processor to be able to communicate with neighbors
  SUBROUTINE create_3d_communicator
    IMPLICIT NONE

    dims(1) = nprocs_x
    dims(2) = nprocs_y
    dims(3) = nprocs_z
    isperiodic(1) = .FALSE.
    isperiodic(2) = .FALSE.
    isperiodic(3) = .FALSE.
    reorder = .TRUE.
    CALL MPI_CART_CREATE(MPI_COMM_WORLD, 3, dims, isperiodic, reorder, comm3d, ierr)
    CALL MPI_COMM_RANK(comm3d, comm3d_rank, ierr)
    CALL MPI_CART_COORDS(comm3d, comm3d_rank, 3, coords, ierr)
    CALL MPI_CART_SHIFT(comm3d, 0, 1, xsrc, xdst, ierr)
    CALL MPI_CART_SHIFT(comm3d, 1, 1, ysrc, ydst, ierr)
    CALL MPI_CART_SHIFT(comm3d, 2, 1, zsrc, zdst, ierr)
    IF(xsrc < 0) xsrc = -1
    IF(ysrc < 0) ysrc = -1
    IF(zsrc < 0) zsrc = -1
    IF(xdst < 0) xdst = -1
    IF(ydst < 0) ydst = -1
    IF(zdst < 0) zdst = -1
    PRINT*, 'myid = ', myid, 'xsrc = ', xsrc, 'xdst = ', xdst
    PRINT*, 'myid = ', myid, 'ysrc = ', ysrc, 'ydst = ', ydst
    PRINT*, 'myid = ', myid, 'zsrc = ', zsrc, 'zdst = ', zdst

  END SUBROUTINE create_3d_communicator

  !****************************************************************************************
  ! SUBROUTINE fill_buffer
  ! Fill the buffer zones for variables (ghost points)
  SUBROUTINE fill_buffer
    IMPLICIT NONE
    INTEGER :: i, j, k

    !~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    ! Xmin
    DO k = 1,bufx
       u(0-k,0:ny,0:nz)  = u(0,0:ny,0:nz)
    END DO
    !~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    ! Xmax
    DO k = 1,bufx
       u(nx+k,0:ny,0:nz) = u(nx,0:ny,0:nz)
    END DO
    !~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    ! Ymin
    DO k = 1,bufy
       u(0:nx,0-k,0:nz)  = u(0:nx,0,0:nz)
    END DO
    !~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    ! Ymax
    DO k = 1,bufy
       u(0:nx,ny+k,0:nz) = u(0:nx,ny,0:nz)
    END DO
    !~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    ! Zmin
    DO k = 1,bufz
       u(0:nx,0:ny,0-k)  = u(0:nx,0:ny,0)
    END DO
    !~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    ! Zmax
    DO k = 1,bufz
       u(0:nx,0:nz,nz+k) = u(0:nx,0:ny,nz)
    END DO


  END SUBROUTINE fill_buffer

  !****************************************************************************************
  ! SUBROUTINE close_parallel_environment
  ! Closes the parallel MPI environment
  SUBROUTINE close_parallel_environment
    IMPLICIT NONE

    CALL dealloc
    CALL MPI_FINALIZE(ierr)

  END SUBROUTINE close_parallel_environment

END MODULE mpi_wrapper

!****************************************************************************************
!****************************************************************************************
! MODULE subroutines
! Contains all the miscellaneous subroutines
MODULE subroutines
  USE types_vars
  USE variables
  USE mpi_wrapper
  USE mpi
  IMPLICIT NONE

CONTAINS

  !****************************************************************************************
  !SUBROUTINE create_output_dir_if_needed
  !Create an output folder to store solution files
  SUBROUTINE create_output_dir_if_needed
    IMPLICIT NONE
    CHARACTER(len=100) :: dir_name1
    CHARACTER :: delimiter
    LOGICAL :: dir_exists

    dir_exists = .FALSE.
    dir_name1 = 'solution'
    dir_name1 = TRIM(dir_name1)
    CALL get_environment_variable('DELIMITER', delimiter)
    INQUIRE(file = dir_name1//'/', exist = dir_exists)
    IF (dir_exists .EQV. .TRUE.) THEN
       PRINT*, 'Output directory exists...'
    ELSE
       PRINT*, 'No output directory exists, creating one...'
       CALL system('mkdir '//TRIM(dir_name1))
    END IF

  END SUBROUTINE create_output_dir_if_needed

  !****************************************************************************************
  ! SUBROUTINE grid3d
  ! Creates a 3d cartesian grid
  SUBROUTINE grid3d
    IMPLICIT NONE
    INTEGER :: i, j, k
    INTEGER :: is, ie, js, je, ks, ke
    REAL(DP) :: lx, ly, lz
    is = 0-bufx; ie = nx+bufx
    js = 0-bufy; je = ny+bufy
    ks = 0-bufz; ke = nz+bufz

    lx = ABS(xmax-xmin)/nprocs_x
    ly = ABS(ymax-ymin)/nprocs_y
    lz = ABS(zmax-zmin)/nprocs_z
    dx = lx/DBLE(nx)
    dy = ly/DBLE(ny)
    dz = lz/DBLE(nz)

    x(is:ie,js:je,ks:ke) = 0.0d0
    y(is:ie,js:je,ks:ke) = 0.0d0
    z(is:ie,js:je,ks:ke) = 0.0d0
    DO k = ks,ke
       DO j = js,je
          DO i = is,ie
             x(i,j,k) = xmin + DBLE(coords(1))*lx + dx*DBLE(i)
             y(i,j,k) = ymin + DBLE(coords(2))*ly + dy*DBLE(j)
             z(i,j,k) = zmin + DBLE(coords(3))*lz + dz*DBLE(k)
          END DO
       END DO
    END DO

  END SUBROUTINE grid3d

  !****************************************************************************************
  ! SUBROUTINE init
  ! Initializes the solution
  SUBROUTINE init
    IMPLICIT NONE
    INTEGER :: i, j, k
    INTEGER :: is, ie, js, je, ks, ke
    is = 0-bufx; ie = nx+bufx
    js = 0-bufy; je = ny+bufy
    ks = 0-bufz; ke = nz+bufz

    DO k = ks,ke
       DO j = js,je
          DO i = is,ie
             u(i,j,k) = 2.d0*DCOS(DSIN(z(i,j,k)))*(DCOS(x(i,j,k) + y(i,j,k) + z(i,j,k)) - &
                  &    (1.d0 + x(i,j,k))*DSIN(x(i,j,k) + y(i,j,k) + z(i,j,k)))
          END DO
       END DO
    END DO

    u_init(0:nx,0:ny,0:nz) = u(0:nx,0:ny,0:nz)

  END SUBROUTINE init

  !****************************************************************************************
  ! SUBROUTINE spanAverage
  ! Compute spanwise average
  SUBROUTINE spanAverage
    IMPLICIT NONE
    INTEGER :: i, j, k
    INTEGER :: sendcnt, recvcnt
    REAL(DP), DIMENSION(0:nx,0:ny) :: u_2d
    REAL(DP), DIMENSION(0:nx,0:ny,0:nprocs_x-1,0:nprocs_z-1) :: u_2d_x_buf
    !~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    ! For each processor, compute a spanwise average
    u_2d(0:nx,0:ny) = 0.0d0
    DO k = 0,nz
       u_2d(0:nx,0:ny) = u_2d(0:nx,0:ny) + u(0:nx,0:ny,k)/DBLE(nz+1)
    END DO
    ! Now, we need to gather all spanwise averages into one buffer varaible
    ! for that streamwise location. First, we need to identify the row of
    ! spanwise CPUs for that streamwise location
    DO i = 0,nprocs_x-1
       DO k = 0,nprocs_z-1
          cpu_x(i,k) = (nprocs_x)*k+i
          PRINT*, cpu_x(i,k)
       END DO
    END DO
    ! Now we know the row of CPUs for each streamwise location. Now we need to
    ! send spanwise averages that are not on the first spanwise plane (right edge
    ! of domain when looking at inlet), to the CPU on the first spanwise plane
    ! at each respective streamwise location
    DO i = 0,nprocs_x-1
       DO k = 0,nprocs_z-1
          IF (k .NE. 0) THEN ! it's not on the first spanwise plane of processors. send to same streamwise location on first spanwise plane
             print*, 'about to send'
             CALL MPI_SEND(u_2d(0,0),sendcnt,MPI_DOUBLE_PRECISION,cpu_x(i,0),0,comm3d,ierr)
             ! CALL MPI_SENDRECV(u_2d(0,0), sendcnt, MPI_DOUBLE_PRECISION, &
             !      &            cpu_x(i,0), 0, &
             !      &            u_2d_x_buf(0,0,i,k), recvcnt, MPI_DOUBLE_PRECISION, &
             !      &            cpu_x(i,k), 0, comm3d, MPI_STATUS_IGNORE, ierr)
             PRINT*, 'SENT ', i, k
          ELSE IF (k .EQ. 0) THEN ! it's on the first spanwise plane of processors. receive all other spanwise averages at that streamwise location
             print*, 'about to receive'
             CALL MPI_RECV(u_2d_x_buf(0,0,i,k),recvcnt,MPI_DOUBLE_PRECISION,MPI_ANY_SOURCE,MPI_ANY_TAG, &
                  &        comm3d,MPI_STATUS_IGNORE,ierr)
             PRINT*, 'RECEIVED ', i
          END IF
       END DO
    END DO
    ! For each streamwise location, compute the spanwise average
    u_2d_x_avg(0:nx,0:ny) = 0.0d0
    DO i = 0,nprocs_x-1
       IF (myid .EQ. cpu_x(i,0)) THEN ! if the current CPU is on the first spanwise plane
          DO k = 0,nprocs_z-1
             u_2d_x_avg(0:nx,0:ny) = u_2d_x_avg(0:nx,0:ny) + u_2d_x_buf(0:nx,0:ny,i,k)/DBLE(nprocs_z)
          END DO
       END IF
    END DO
    ! Now for each streamwise location, you have the spanwise average. These need
    ! to be printed in separate .dat files to be post-processed

  END SUBROUTINE spanAverage

END MODULE subroutines

!****************************************************************************************
!****************************************************************************************
! PROGRAM main
PROGRAM main

  USE types_vars
  USE variables
  USE io
  USE mpi_wrapper
  USE subroutines
  USE mpi
  IMPLICIT NONE

  CALL init_parallel_environment
  CALL grid3d
  CALL init
  CALL fill_buffer
  IF (myid == 0) CALL create_output_dir_if_needed
  CALL spanAverage
  CALL tecplot_output
  CALL close_parallel_environment

END PROGRAM main