!!!    This program sloves Lid Driven Cavity Flow problem using Lattice Boltzmann Method
!!!    Lattice Boltzmann Equation with MRT-LBE model

module commondata
    implicit none

    integer, parameter :: total_nx = 65, total_ny = 65, total_nz = 65
    integer, parameter :: nx_block = 2, ny_block = 3, nz_block = 3
    integer :: nx, ny, nz
    integer :: block_x, block_y, block_z
    real(8), parameter :: Reynolds=1000.0d0
    real(8), parameter :: rho0=1.0d0
    real(8), parameter :: U0=0.1d0
    real(8), parameter :: tauf=U0*dble(total_nx)/Reynolds*3.0d0+0.5d0

    integer :: itc
    integer, parameter :: itc_max=INT(50000000)

    real(8) :: errorU
    real(8), parameter :: eps=1e-6

    real(8) :: xp(0:total_nx+1), yp(0:total_ny+1), zp(0:total_nz+1)
    real(8), allocatable :: u(:, :, :), v(:, :, :), w(:, :, :)
    real(8), allocatable :: rho(:, :, :)
    real(8), allocatable :: up(:, :, :), vp(:, :, :), wp(:, :, :)

    real(8), allocatable :: f(:, :, :, :), f_post(:, :, :, :)

    integer :: idx
    ! real(8), parameter :: omega(0:18) = (/ 1.0d0/3.0d0, &
    !         1.0d0/18.0d0, 1.0d0/18.0d0, 1.0d0/18.0d0, 1.0d0/18.0d0, 1.0d0/18.0d0, 1.0d0/18.0d0, &
    !         1.0d0/36.0d0, 1.0d0/36.0d0, 1.0d0/36.0d0, 1.0d0/36.0d0, 1.0d0/36.0d0, 1.0d0/36.0d0, &
    !         1.0d0/36.0d0, 1.0d0/36.0d0, 1.0d0/36.0d0, 1.0d0/36.0d0, 1.0d0/36.0d0, 1.0d0/36.0d0 /)
    real(8), parameter :: omega(0:18) = (/ 1.0d0/3.0d0, &
                                        (1.0d0/18.0d0, idx = 1, 6), &
                                        (1.0d0/36.0d0, idx = 7, 18) /)
    integer, parameter :: ex(0:18) = (/ 0,  &
                                        1, -1,  0,  0,  0,  0, &
                                        1, -1,  1, -1,  1, -1,  1, -1,  0,  0,  0,  0 /)
    integer, parameter :: ey(0:18) = (/ 0, &
                                        0,  0,  1, -1,  0,  0, &
                                        1,  1, -1, -1,  0,  0,  0,  0,  1, -1,  1, -1 /)
    integer, parameter :: ez(0:18) = (/ 0, &
                                        0,  0,  0,  0,  1, -1, &
                                        0,  0,  0,  0,  1,  1, -1, -1,  1,  1, -1, -1 /)

    real(8), parameter :: Snu=1.0d0/tauf, Sq=8.0d0*(2.0d0*tauf-1.0d0)/(8.0d0*tauf-1.0d0)
    integer :: rc, rank, num_process
end module commondata

! block arrangement is same as the data arrangement

program main
    use mpi
    use commondata
    implicit none
    real(8) :: start, finish
    real(8) :: start_time, end_time
    integer :: name_len
    character(len=MPI_MAX_PROCESSOR_NAME) :: processor_name

    call MPI_Init(rc)

    call MPI_Comm_size(MPI_COMM_WORLD, num_process, rc)
    call MPI_Comm_rank(MPI_COMM_WORLD, rank, rc)
    call MPI_Get_processor_name(processor_name, name_len, rc)

    if (rank == 0) then
        write(*,*) "Using ", num_process, "processers."
        write(*,*) "nx_block = ", nx_block
        write(*,*) "ny_block = ", ny_block
        write(*,*) "nz_block = ", nz_block
    endif

    if (nx_block * ny_block * nz_block /= num_process) then
        if (rank == 0) then
            write(*,*) "Please make sure nx_block * ny_block * nz_block = num_process."
        endif
        stop
    endif

    ! block index --- block start from index 1 but rank start from rank 0
    block_x = mod(rank, nx_block) + 1
    block_y = mod(rank, nx_block * ny_block) / nx_block + 1
    block_z = rank / (nx_block * ny_block) + 1

    ! write(*,"('Rank ', I3, ' has block index = (', I3, ',', I3, ',', I3, ').')") rank, block_x, block_y, block_z

    !!! zone divide
    nx = total_nx / nx_block
    ny = total_ny / ny_block
    nz = total_nz / nz_block
    ! handle non-divisible part
    if(block_x <= mod(total_nx, nx_block)) then
        nx = nx + 1
    endif
    if(block_y <= mod(total_ny, ny_block)) then
        ny = ny + 1
    endif
    if(block_z <= mod(total_nz, nz_block)) then
        nz = nz + 1
    endif

    ! write(*,"('Rank ', I3, ' has nx*ny*nz = ', I3, 'x', I3, 'x', I3)") rank, nx, ny, nz

    call initial()

    ! call output_field()
    ! call output_velocity()
    call output()

    call CPU_TIME(start)

    call MPI_Barrier(MPI_COMM_WORLD, rc)
    start_time = MPI_Wtime()

    do while((errorU.GT.eps).AND.(itc.LE.itc_max))

        itc = itc+1

        call collision()

        call message_passing()

        call streaming()

        call bounceback()

        call macro()

        if(MOD(itc,2000).EQ.0) then
            call check()
            ! call output()
        endif

    enddo

    call MPI_Barrier(MPI_COMM_WORLD, rc)
    end_time = MPI_Wtime()

    call CPU_TIME(finish)

    if (rank == 0) then
        write(*,*) "Time (CPU) = ", real(finish-start), "s"
        write(*,*) "Wall time = ", real(end_time - start_time), "s"
    endif

    itc = itc+1

    call output()

    if (rank == 0) then
        write(*,*) "Deallocate Array..."
    endif

    deallocate(u)
    deallocate(v)
    deallocate(w)
    deallocate(rho)
    deallocate(up)
    deallocate(vp)
    deallocate(wp)
    deallocate(f)
    deallocate(f_post)

     if (rank == 0) then
        write(*,*) "    "

        write(*,*) "Successfully: DNS completed!"
    endif

    call MPI_Finalize(rc)

end program main

subroutine initial()
    use commondata
    implicit none
    integer :: i, j, k
    integer :: alpha
    real(8) :: un(0:18)
    real(8) :: us2

    itc = 0
    errorU = 100.0d0

    if (rank == 0) then
        write(*,*) "nx=",total_nx,", ny=",total_ny, ",  nz = ",total_nz
        write(*,*) "Reynolds=",real(Reynolds)
        write(*,*) "U0=",real(U0),",    tauf=",real(tauf)
    endif

    xp(0) = 0.0d0
    xp(total_nx+1) = dble(total_nx)
    do i=1,total_nx
        xp(i) = dble(i)-0.5d0
    enddo

    yp(0) = 0.0d0
    yp(total_ny+1) = dble(total_ny)
    do j=1,total_ny
        yp(j) = dble(j)-0.5d0
    enddo

    zp(0) = 0.0d0
    zp(total_nz+1) = dble(total_nz)
    do k=1,total_nz
        zp(k) = dble(k)-0.5d0
    enddo

    allocate (u(nx, ny, nz))
    allocate (v(nx, ny, nz))
    allocate (w(nx, ny, nz))
    allocate (rho(nx, ny, nz))
    allocate (up(nx, ny, nz))
    allocate (vp(nx, ny, nz))
    allocate (wp(nx, ny, nz))

    allocate (f(0:18, nx, ny, nz))
    allocate (f_post(0:18, 0:nx+1, 0:ny+1, 0:nz+1))

    rho = rho0
    u = 0.0d0
    v = 0.0d0
    w = 0.0d0
    up = 0.0d0
    vp = 0.0d0
    wp = 0.0d0

    if (block_z == nz_block) then
        do j = 1, ny
            do i = 1, nx
                u(i, j, nz) = U0
            enddo
        enddo
    endif

    do k = 1, nz
        do j = 1, ny
            do i = 1, nx
                us2 = u(i,j,k)*u(i,j,k) + v(i,j,k)*v(i,j,k) + w(i,j,k)*w(i,j,k)
                do alpha = 0, 18
                    un(alpha) = u(i,j,k)*ex(alpha) + v(i,j,k)*ey(alpha) + w(i,j,k)*ez(alpha)
                    f(alpha,i,j,k) = rho(i,j,k)*omega(alpha)*(1.0d0+3.0d0*un(alpha)+4.5d0*un(alpha)*un(alpha)-1.5d0*us2)
                enddo
            enddo
        enddo
    enddo

    return
end subroutine initial

subroutine collision()
    use commondata
    implicit none
    integer :: i, j, k
    integer :: alpha
    !---------------------------
    real(8) :: s(0:18)
    real(8) :: m(0:18)
    real(8) :: m_post(0:18)
    real(8) :: meq(0:18)
    !---------------------------
    ! real(8) :: us2, un, feq
    !---------------------------

    do k = 1, nz
        do j = 1, ny
            do i = 1, nx
    ! loop body, re-indented for readability

    m(0) = f(0,i,j,k) &
        + f(1,i,j,k) + f(2,i,j,k) + f(3,i,j,k) + f(4,i,j,k) + f(5,i,j,k) + f(6,i,j,k) &
        + f(7,i,j,k) + f(8,i,j,k) + f(9,i,j,k) + f(10,i,j,k) + f(11,i,j,k) + f(12,i,j,k) &
        + f(13,i,j,k) + f(14,i,j,k) + f(15,i,j,k) + f(16,i,j,k) + f(17,i,j,k) + f(18,i,j,k)
    m(1) = -30.0d0 * f(0,i,j,k) &
        - 11.0d0 * (f(1,i,j,k) + f(2,i,j,k) + f(3,i,j,k) + f(4,i,j,k) + f(5,i,j,k) + f(6,i,j,k)) &
        + 8.0d0 * (f(7,i,j,k) + f(8,i,j,k) + f(9,i,j,k) + f(10,i,j,k) + f(11,i,j,k) + f(12,i,j,k)) &
        + 8.0d0 * (f(13,i,j,k) + f(14,i,j,k) + f(15,i,j,k) + f(16,i,j,k) + f(17,i,j,k) + f(18,i,j,k))
    m(2) = 12.0d0 * f(0,i,j,k) &
        - 4.0d0 * (f(1,i,j,k) + f(2,i,j,k) + f(3,i,j,k) + f(4,i,j,k) + f(5,i,j,k) + f(6,i,j,k)) &
        + f(7,i,j,k) + f(8,i,j,k) + f(9,i,j,k) + f(10,i,j,k) + f(11,i,j,k) + f(12,i,j,k) &
        + f(13,i,j,k) + f(14,i,j,k) + f(15,i,j,k) + f(16,i,j,k) + f(17,i,j,k) + f(18,i,j,k)
    m(3) = f(1,i,j,k) - f(2,i,j,k) &
        + f(7,i,j,k) - f(8,i,j,k) + f(9,i,j,k) - f(10,i,j,k) + f(11,i,j,k) - f(12,i,j,k) &
        + f(13,i,j,k) - f(14,i,j,k)
    m(4) = -4.0d0 * (f(1,i,j,k) - f(2,i,j,k)) &
        + f(7,i,j,k) - f(8,i,j,k) + f(9,i,j,k) - f(10,i,j,k) + f(11,i,j,k) - f(12,i,j,k) &
        + f(13,i,j,k) - f(14,i,j,k)
    m(5) = f(3,i,j,k) - f(4,i,j,k) &
        + f(7,i,j,k) + f(8,i,j,k) - f(9,i,j,k) - f(10,i,j,k) &
        + f(15,i,j,k) - f(16,i,j,k) + f(17,i,j,k) - f(18,i,j,k)
    m(6) = -4.0d0 * (f(3,i,j,k) - f(4,i,j,k)) &
        + f(7,i,j,k) + f(8,i,j,k) - f(9,i,j,k) - f(10,i,j,k) &
        + f(15,i,j,k) - f(16,i,j,k) + f(17,i,j,k) - f(18,i,j,k)
    m(7) = f(5,i,j,k) - f(6,i,j,k) &
        + f(11,i,j,k) + f(12,i,j,k) &
        - f(13,i,j,k) - f(14,i,j,k) + f(15,i,j,k) + f(16,i,j,k) - f(17,i,j,k) - f(18,i,j,k)
    m(8) = -4.0d0 * (f(5,i,j,k) - f(6,i,j,k)) &
        + f(11,i,j,k) + f(12,i,j,k) &
        - f(13,i,j,k) - f(14,i,j,k) + f(15,i,j,k) + f(16,i,j,k) - f(17,i,j,k) - f(18,i,j,k)
    m(9) = 2.0d0 * (f(1,i,j,k) + f(2,i,j,k)) - f(3,i,j,k) - f(4,i,j,k) - f(5,i,j,k) - f(6,i,j,k) &
        + f(7,i,j,k) + f(8,i,j,k) + f(9,i,j,k) + f(10,i,j,k) + f(11,i,j,k) + f(12,i,j,k) &
        + f(13,i,j,k) + f(14,i,j,k) - 2.0 * (f(15,i,j,k) + f(16,i,j,k) + f(17,i,j,k) + f(18,i,j,k))
    m(10) = -4.0d0 * (f(1,i,j,k) + f(2,i,j,k)) + 2.0d0 * (f(3,i,j,k) + f(4,i,j,k) + f(5,i,j,k) + f(6,i,j,k)) &
        + f(7,i,j,k) + f(8,i,j,k) + f(9,i,j,k) + f(10,i,j,k) + f(11,i,j,k) + f(12,i,j,k) &
        + f(13,i,j,k) + f(14,i,j,k) - 2.0 * (f(15,i,j,k) + f(16,i,j,k) + f(17,i,j,k) + f(18,i,j,k))
    m(11) = f(3,i,j,k) + f(4,i,j,k) - f(5,i,j,k) - f(6,i,j,k) &
        + f(7,i,j,k) + f(8,i,j,k) + f(9,i,j,k) + f(10,i,j,k) - f(11,i,j,k) - f(12,i,j,k) &
        - f(13,i,j,k) - f(14,i,j,k)
    m(12) = -2.0d0 * (f(3,i,j,k) + f(4,i,j,k) - f(5,i,j,k) - f(6,i,j,k)) &
        + f(7,i,j,k) + f(8,i,j,k) + f(9,i,j,k) + f(10,i,j,k) - f(11,i,j,k) - f(12,i,j,k) &
        - f(13,i,j,k) - f(14,i,j,k)
    m(13) = f(7,i,j,k) - f(8,i,j,k) - f(9,i,j,k) + f(10,i,j,k)
    m(14) = f(15,i,j,k) - f(16,i,j,k) - f(17,i,j,k) + f(18,i,j,k)
    m(15) = f(11,i,j,k) - f(12,i,j,k) - f(13,i,j,k) + f(14,i,j,k)
    m(16) = f(7,i,j,k) - f(8,i,j,k) + f(9,i,j,k) - f(10,i,j,k) - f(11,i,j,k) + f(12,i,j,k) &
        - f(13,i,j,k) + f(14,i,j,k)
    m(17) = -f(7,i,j,k) - f(8,i,j,k) + f(9,i,j,k) + f(10,i,j,k) &
        + f(15,i,j,k) - f(16,i,j,k) + f(17,i,j,k) - f(18,i,j,k)
    m(18) = f(11,i,j,k) + f(12,i,j,k) &
        - f(13,i,j,k) - f(14,i,j,k) - f(15,i,j,k) - f(16,i,j,k) + f(17,i,j,k) + f(18,i,j,k)

        meq(0) = rho(i,j,k)
        meq(1) = rho(i,j,k) * (-11.0d0 + 19.0d0 * (u(i,j,k)*u(i,j,k) + v(i,j,k)*v(i,j,k) + w(i,j,k)*w(i,j,k)))
        meq(2) = rho(i,j,k) * (3.0d0 - 11.0d0/2.0d0 * (u(i,j,k)*u(i,j,k) + v(i,j,k)*v(i,j,k) + w(i,j,k)*w(i,j,k)))
        meq(3) = rho(i,j,k) * u(i,j,k)
        meq(4) = -2.0d0/3.0d0 * rho(i,j,k) * u(i,j,k)
        meq(5) = rho(i,j,k) * v(i,j,k)
        meq(6) = -2.0d0/3.0d0 * rho(i,j,k) * v(i,j,k)
        meq(7) = rho(i,j,k) * w(i,j,k)
        meq(8) = -2.0d0/3.0d0 * rho(i,j,k) * w(i,j,k)
        meq(9) = rho(i,j,k) * (2.0d0 * u(i,j,k)*u(i,j,k) - v(i,j,k)*v(i,j,k) - w(i,j,k)*w(i,j,k))
        meq(10) = -1.0d0/2.0d0 * rho(i,j,k) * (2.0d0 * u(i,j,k)*u(i,j,k) - v(i,j,k)*v(i,j,k) - w(i,j,k)*w(i,j,k))
        meq(11) = rho(i,j,k) * (v(i,j,k)*v(i,j,k) - w(i,j,k)*w(i,j,k))
        meq(12) = -1.0d0/2.0d0 * (v(i,j,k)*v(i,j,k) - w(i,j,k)*w(i,j,k))
        meq(13) = rho(i,j,k) * u(i,j,k) * v(i,j,k)
        meq(14) = rho(i,j,k) * v(i,j,k) * w(i,j,k)
        meq(15) = rho(i,j,k) * u(i,j,k) * w(i,j,k)
        meq(16) = 0.0d0
        meq(17) = 0.0d0
        meq(18) = 0.0d0

        s(0) = 0.0d0    !! s_{\rho}
        s(1) = Snu      !! s_{e}
        s(2) = Snu      !! s_{epsilon}
        s(3) = 0.0d0    !! s_{j}
        s(4) = Sq       !! s_{q}
        s(5) = 0.0d0    !! s_{j}
        s(6) = Sq       !! s_{q}
        s(7) = 0.0d0    !! s_{j}
        s(8) = Sq       !! s_{q}
        s(9) = Snu      !! s_{\nu}
        s(10) = Snu     !! s_{\pi}
        s(11) = Snu     !! s_{\nu}
        s(12) = Snu     !! s_{\pi}
        s(13) = Snu     !! s_{nu}
        s(14) = Snu     !! s_{nu}
        s(15) = Snu     !! s_{nu}
        s(16) = Sq      !! s_{m}
        s(17) = Sq      !! s_{m}
        s(18) = Sq      !! s_{m}

        do alpha=0,18
            m_post(alpha) = m(alpha)-s(alpha)*(m(alpha)-meq(alpha))
        enddo

    f_post(0,i,j,k) = m_post(0)/19.0d0 - 5.0d0/399.0d0*m_post(1) + m_post(2)/21.0d0

    ! ------------
    f_post(1,i,j,k) = m_post(0)/19.0d0 - 11.0d0/2394.0d0*m_post(1) - m_post(2)/63.0d0 + m_post(3)/10.0d0 &
        - m_post(4)/10.0d0 + m_post(9)/18.0d0 - m_post(10)/18.0d0

    f_post(2,i,j,k) = m_post(0)/19.0d0 - 11.0d0/2394.0d0*m_post(1) - m_post(2)/63.0d0 - m_post(3)/10.0d0 &
        + m_post(4)/10.0d0 + m_post(9)/18.0d0 - m_post(10)/18.0d0

    ! ------------
    f_post(3,i,j,k) = m_post(0)/19.0d0 - 11.0d0/2394.0d0*m_post(1) - m_post(2)/63.0d0 + m_post(5)/10.0d0 &
        - m_post(6)/10.0d0 - m_post(9)/36.0d0 + m_post(10)/36.0d0 + m_post(11)/12.0d0 - m_post(12)/12.0d0

    f_post(4,i,j,k) = m_post(0)/19.0d0 - 11.0d0/2394.0d0*m_post(1) - m_post(2)/63.0d0 - m_post(5)/10.0d0 &
        + m_post(6)/10.0d0 - m_post(9)/36.0d0 + m_post(10)/36.0d0 + m_post(11)/12.0d0 - m_post(12)/12.0d0

    f_post(5,i,j,k) = m_post(0)/19.0d0 - 11.0d0/2394.0d0*m_post(1) - m_post(2)/63.0d0 + m_post(7)/10.0d0 &
        - m_post(8)/10.0d0 - m_post(9)/36.0d0 + m_post(10)/36.0d0 - m_post(11)/12.0d0 + m_post(12)/12.0d0

    f_post(6,i,j,k) = m_post(0)/19.0d0 - 11.0d0/2394.0d0*m_post(1) - m_post(2)/63.0d0 - m_post(7)/10.0d0 &
        + m_post(8)/10.0d0 - m_post(9)/36.0d0 + m_post(10)/36.0d0 - m_post(11)/12.0d0 + m_post(12)/12.0d0

    ! ------------
    f_post(7,i,j,k) = m_post(0)/19.0d0 + 4.0d0/1197.0d0*m_post(1) + m_post(2)/252.0d0 + m_post(3)/10.0d0 &
        + m_post(4)/40.0d0 + m_post(5)/10.0d0 + m_post(6)/40.0d0 + m_post(9)/36.0d0 + m_post(10)/72.0d0 &
        + m_post(11)/12.0d0 + m_post(12)/24.0d0 + m_post(13)/4.0d0 + m_post(16)/8.0d0 - m_post(17)/8.0d0

    f_post(8,i,j,k) = m_post(0)/19.0d0 + 4.0d0/1197.0d0*m_post(1) + m_post(2)/252.0d0 - m_post(3)/10.0d0 &
        - m_post(4)/40.0d0 + m_post(5)/10.0d0 + m_post(6)/40.0d0 + m_post(9)/36.0d0 + m_post(10)/72.0d0 &
        + m_post(11)/12.0d0 + m_post(12)/24.0d0 - m_post(13)/4.0d0 - m_post(16)/8.0d0 - m_post(17)/8.0d0

    f_post(9,i,j,k) = m_post(0)/19.0d0 + 4.0d0/1197.0d0*m_post(1) + m_post(2)/252.0d0 + m_post(3)/10.0d0 &
        + m_post(4)/40.0d0 - m_post(5)/10.0d0 - m_post(6)/40.0d0 + m_post(9)/36.0d0 + m_post(10)/72.0d0 &
        + m_post(11)/12.0d0 + m_post(12)/24.0d0 - m_post(13)/4.0d0 + m_post(16)/8.0d0 + m_post(17)/8.0d0

    f_post(10,i,j,k) = m_post(0)/19.0d0 + 4.0d0/1197.0d0*m_post(1) + m_post(2)/252.0d0 - m_post(3)/10.0d0 &
        - m_post(4)/40.0d0 - m_post(5)/10.0d0 - m_post(6)/40.0d0 + m_post(9)/36.0d0 + m_post(10)/72.0d0 &
        + m_post(11)/12.0d0 + m_post(12)/24.0d0 + m_post(13)/4.0d0 - m_post(16)/8.0d0 + m_post(17)/8.0d0

    ! -----------
    f_post(11,i,j,k) = m_post(0)/19.0d0 + 4.0d0/1197.0d0*m_post(1) + m_post(2)/252.0d0 + m_post(3)/10.0d0 &
        + m_post(4)/40.0d0 + m_post(7)/10.0d0 + m_post(8)/40.0d0 + m_post(9)/36.0d0 + m_post(10)/72.0d0 &
        - m_post(11)/12.0d0 - m_post(12)/24.0d0 + m_post(15)/4.0d0 - m_post(16)/8.0d0 + m_post(18)/8.0d0

    f_post(12,i,j,k) = m_post(0)/19.0d0 + 4.0d0/1197.0d0*m_post(1) + m_post(2)/252.0d0 - m_post(3)/10.0d0 &
        - m_post(4)/40.0d0 + m_post(7)/10.0d0 + m_post(8)/40.0d0 + m_post(9)/36.0d0 + m_post(10)/72.0d0 &
        - m_post(11)/12.0d0 - m_post(12)/24.0d0 - m_post(15)/4.0d0 + m_post(16)/8.0d0 + m_post(18)/8.0d0

    f_post(13,i,j,k) = m_post(0)/19.0d0 + 4.0d0/1197.0d0*m_post(1) + m_post(2)/252.0d0 + m_post(3)/10.0d0 &
        + m_post(4)/40.0d0 - m_post(7)/10.0d0 - m_post(8)/40.0d0 + m_post(9)/36.0d0 + m_post(10)/72.0d0 &
        - m_post(11)/12.0d0 - m_post(12)/24.0d0 - m_post(15)/4.0d0 - m_post(16)/8.0d0 - m_post(18)/8.0d0

    f_post(14,i,j,k) = m_post(0)/19.0d0 + 4.0d0/1197.0d0*m_post(1) + m_post(2)/252.0d0 - m_post(3)/10.0d0 &
        - m_post(4)/40.0d0 - m_post(7)/10.0d0 - m_post(8)/40.0d0 + m_post(9)/36.0d0 + m_post(10)/72.0d0 &
        - m_post(11)/12.0d0 - m_post(12)/24.0d0 + m_post(15)/4.0d0 + m_post(16)/8.0d0 - m_post(18)/8.0d0

    ! -----------
    f_post(15,i,j,k) = m_post(0)/19.0d0 + 4.0d0/1197.0d0*m_post(1) + m_post(2)/252.0d0 + m_post(5)/10.0d0 &
        + m_post(6)/40.0d0 + m_post(7)/10.0d0 + m_post(8)/40.0d0 - m_post(9)/18.0d0 - m_post(10)/36.0d0 &
        + m_post(14)/4.0d0 + m_post(17)/8.0d0 - m_post(18)/8.0d0

    f_post(16,i,j,k) = m_post(0)/19.0d0 + 4.0d0/1197.0d0*m_post(1) + m_post(2)/252.0d0 - m_post(5)/10.0d0 &
        - m_post(6)/40.0d0 + m_post(7)/10.0d0 + m_post(8)/40.0d0 - m_post(9)/18.0d0 - m_post(10)/36.0d0 &
        - m_post(14)/4.0d0 - m_post(17)/8.0d0 - m_post(18)/8.0d0

    f_post(17,i,j,k) = m_post(0)/19.0d0 + 4.0d0/1197.0d0*m_post(1) + m_post(2)/252.0d0 + m_post(5)/10.0d0 &
        + m_post(6)/40.0d0 - m_post(7)/10.0d0 - m_post(8)/40.0d0 - m_post(9)/18.0d0 - m_post(10)/36.0d0 &
        - m_post(14)/4.0d0 + m_post(17)/8.0d0 + m_post(18)/8.0d0

    f_post(18,i,j,k) = m_post(0)/19.0d0 + 4.0d0/1197.0d0*m_post(1) + m_post(2)/252.0d0 - m_post(5)/10.0d0 &
        - m_post(6)/40.0d0 - m_post(7)/10.0d0 - m_post(8)/40.0d0 - m_post(9)/18.0d0 - m_post(10)/36.0d0 &
        + m_post(14)/4.0d0 - m_post(17)/8.0d0 + m_post(18)/8.0d0

                ! us2 = u(i,j,k)*u(i,j,k) + v(i,j,k)*v(i,j,k) + w(i,j,k)*w(i,j,k)
                ! do alpha = 0, 18
                !     un = u(i,j,k)*ex(alpha) + v(i,j,k)*ey(alpha) + w(i,j,k)*ez(alpha)
                !     feq = rho(i,j,k) * omega(alpha) * (1.0d0 + 3.0d0*un + 4.5d0*un*un - 1.5d0*us2)

                !     f_post(alpha,i,j,k) = f(alpha,i,j,k) - Snu*(f(alpha,i,j,k) - feq)
                ! enddo

            enddo
        enddo
    enddo

    return
end subroutine collision

subroutine streaming()
    use commondata
    implicit none
    integer :: i, j, k
    integer :: ip, jp, kp
    integer :: alpha

    do k = 1, nz
        do j = 1, ny
            do i = 1, nx
                do alpha=0,18
                    ip = i - ex(alpha)
                    jp = j - ey(alpha)
                    kp = k - ez(alpha)

                    f(alpha,i,j,k) = f_post(alpha,ip,jp,kp)
                enddo
            enddo
        enddo
    enddo

    return
end subroutine streaming

subroutine bounceback()
    use commondata
    implicit none
    integer :: i, j, k

    if (block_x == 1) then
        do k = 1, nz
            do j = 1, ny
                ! (i = 1)
                f(1,1,j,k) = f_post(2,1,j,k)
                f(7,1,j,k) = f_post(10,1,j,k)
                f(9,1,j,k) = f_post(8,1,j,k)
                f(11,1,j,k) = f_post(14,1,j,k)
                f(13,1,j,k) = f_post(12,1,j,k)
            enddo
        enddo
    endif

    if (block_x == nx_block) then
        do k = 1, nz
            do j = 1, ny
                !(i = nx)
                f(2,nx,j,k) = f_post(1,nx,j,k)
                f(10,nx,j,k) = f_post(7,nx,j,k)
                f(8,nx,j,k) = f_post(9,nx,j,k)
                f(14,nx,j,k) = f_post(11,nx,j,k)
                f(12,nx,j,k) = f_post(13,nx,j,k)
            enddo
        enddo
    endif

    if (block_y == 1) then
        do k = 1, nz
            do i = 1, nx
                !(j = 1)
                f(3,i,1,k) = f_post(4,i,1,k)
                f(7,i,1,k) = f_post(10,i,1,k)
                f(8,i,1,k) = f_post(9,i,1,k)
                f(15,i,1,k) = f_post(18,i,1,k)
                f(17,i,1,k) = f_post(16,i,1,k)
            enddo
        enddo
    endif

    if (block_y == ny_block) then
        do k = 1, nz
            do i = 1, nx
                !(j = ny)
                f(4,i,ny,k) = f_post(3,i,ny,k)
                f(10,i,ny,k) = f_post(7,i,ny,k)
                f(9,i,ny,k) = f_post(8,i,ny,k)
                f(18,i,ny,k) = f_post(15,i,ny,k)
                f(16,i,ny,k) = f_post(17,i,ny,k)
            enddo
        enddo
    endif

    if (block_z == 1) then
        do j = 1, ny
            do i = 1, nx
                ! (k = 1)
                f(5,i,j,1) = f_post(6,i,j,1)
                f(11,i,j,1) = f_post(14,i,j,1)
                f(12,i,j,1) = f_post(13,i,j,1)
                f(15,i,j,1) = f_post(18,i,j,1)
                f(16,i,j,1) = f_post(17,i,j,1)
            enddo
        enddo
    endif

    if (block_z == nz_block) then
        do j = 1, ny
            do i = 1, nx
                ! (k = nz)
                f(6,i,j,nz) = f_post(5,i,j,nz)
                f(14,i,j,nz) = f_post(11,i,j,nz) - rho(i,j,nz) / 6.0d0 * (U0)
                f(13,i,j,nz) = f_post(12,i,j,nz) - rho(i,j,nz) / 6.0d0 * (-U0)
                f(18,i,j,nz) = f_post(15,i,j,nz)
                f(17,i,j,nz) = f_post(16,i,j,nz)
            enddo
        enddo
    endif

    return
end subroutine bounceback

subroutine macro()
    use commondata
    implicit none
    integer :: i, j, k, alpha

    rho = 0.0d0
    u = 0.0d0
    v = 0.0d0
    w = 0.0d0
    do k = 1, nz
        do j = 1, ny
            do i = 1, nx
                do alpha = 0, 18
                    rho(i,j,k) = rho(i,j,k) + f(alpha,i,j,k)
                    u(i,j,k) = u(i,j,k) + f(alpha,i,j,k) * ex(alpha)
                    v(i,j,k) = v(i,j,k) + f(alpha,i,j,k) * ey(alpha)
                    w(i,j,k) = w(i,j,k) + f(alpha,i,j,k) * ez(alpha)
                enddo

                u(i,j,k) = u(i,j,k)/rho(i,j,k)
                v(i,j,k) = v(i,j,k)/rho(i,j,k)
                w(i,j,k) = w(i,j,k)/rho(i,j,k)
            enddo
        enddo
    enddo

    return
end subroutine macro

subroutine check()
    use mpi
    use commondata
    implicit none
    integer :: i, j, k
    real(8) :: error1, error2
    real(8) :: total_error1, total_error2

    error1 = 0.0d0
    error2 = 0.0d0

    do k = 1, nz
        do j = 1, ny
            do i = 1, nx
                error1 = error1 + (u(i,j,k)-up(i,j,k))*(u(i,j,k)-up(i,j,k))+(v(i,j,k)-vp(i,j,k))*(v(i,j,k)-vp(i,j,k))
                error2 = error2 + u(i,j,k)*u(i,j,k)+v(i,j,k)*v(i,j,k) + w(i,j,k)*w(i,j,k)
            enddo
        enddo
    enddo

    up = u
    vp = v
    wp = w

    call MPI_Barrier(MPI_COMM_WORLD, rc)

    call MPI_ALLreduce(error1, total_error1, 1, MPI_DOUBLE_PRECISION, MPI_SUM, MPI_COMM_WORLD, rc)
    call MPI_ALLreduce(error2, total_error2, 1, MPI_DOUBLE_PRECISION, MPI_SUM, MPI_COMM_WORLD, rc)

    errorU = sqrt(total_error1)/sqrt(total_error2)

    if (rank == 0) then
        write(*,*) itc,' ',errorU
    endif

    return
end subroutine check

subroutine message_passing()
    use mpi
    use commondata
    implicit none
    integer :: i, j, k
    ! surface data to send:   fps_send_i(5,j,k) is data at surface i=1 and i = nx
    real(8) :: fps_send_i(5,ny,nz), fps_send_j(5,nx,nz), fps_send_k(5,nx,ny)
    ! surface data to receive
    real(8) :: fps_recv_i(5,ny,nz), fps_recv_j(5,nx,nz), fps_recv_k(5,nx,ny)

    ! line data to send  ---  fpl_send_z(k) is the line parallel to axis z
    real(8) :: fpl_send_z(nz), fpl_send_y(ny), fpl_send_x(nx)
    real(8) :: fpl_recv_z(nz), fpl_recv_y(ny), fpl_recv_x(nx)

    ! we don't need to exchange fp for cornor points

    !!! /* ------------- exchange message with surfaces -----------------
    !!!    message tag(for sender):  1 -> i=1  /  2 -> i=nx   /  3 -> j=1  /
    !!!                              4 -> j=ny   /   5 -> k=1    /   6 -> k=nz  */

    ! /* exchange message with i=1 surface --- send then receive*/
    if (block_x > 1) then
        ! collect data to send (i=1)
        do k = 1, nz
            do j = 1, ny
                fps_send_i(1,j,k) = f_post(2,1,j,k)
                fps_send_i(2,j,k) = f_post(8,1,j,k)
                fps_send_i(3,j,k) = f_post(10,1,j,k)
                fps_send_i(4,j,k) = f_post(12,1,j,k)
                fps_send_i(5,j,k) = f_post(14,1,j,k)
            enddo
        enddo

        call MPI_Send(fps_send_i, 5*ny*nz, MPI_DOUBLE_PRECISION, rank-1, 1, MPI_COMM_WORLD, rc)
        call MPI_Recv(fps_recv_i, 5*ny*nz, MPI_DOUBLE_PRECISION, rank-1, 2, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)

        ! scatter the received data (i=0)
        do k = 1, nz
            do j = 1, ny
                f_post(1,0,j,k) = fps_recv_i(1,j,k)
                f_post(7,0,j,k) = fps_recv_i(2,j,k)
                f_post(9,0,j,k) = fps_recv_i(3,j,k)
                f_post(11,0,j,k) = fps_recv_i(4,j,k)
                f_post(13,0,j,k) = fps_recv_i(5,j,k)
            enddo
        enddo
    endif

    ! /* exchange message with i=nx surface --- receive then send*/
    if (block_x < nx_block) then
        ! collect data to send (i=nx)
        do k = 1, nz
            do j = 1, ny
                fps_send_i(1,j,k) = f_post(1,nx,j,k)
                fps_send_i(2,j,k) = f_post(7,nx,j,k)
                fps_send_i(3,j,k) = f_post(9,nx,j,k)
                fps_send_i(4,j,k) = f_post(11,nx,j,k)
                fps_send_i(5,j,k) = f_post(13,nx,j,k)
            enddo
        enddo

        call MPI_Recv(fps_recv_i, 5*ny*nz, MPI_DOUBLE_PRECISION, rank+1, 1, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)
        call MPI_Send(fps_send_i, 5*ny*nz, MPI_DOUBLE_PRECISION, rank+1, 2, MPI_COMM_WORLD, rc)

        ! scatter the received data (i=nx+1)
        do k = 1, nz
            do j = 1, ny
                f_post(2,nx+1,j,k) = fps_recv_i(1,j,k)
                f_post(8,nx+1,j,k) = fps_recv_i(2,j,k)
                f_post(10,nx+1,j,k) = fps_recv_i(3,j,k)
                f_post(12,nx+1,j,k) = fps_recv_i(4,j,k)
                f_post(14,nx+1,j,k) = fps_recv_i(5,j,k)
            enddo
        enddo
    endif

    ! /* exchange message with j=1 surface --- send then receive*/
    if (block_y > 1) then
        ! collect data to send (j=1)
        do k = 1, nz
            do i = 1, nx
                fps_send_j(1,i,k) = f_post(4,i,1,k)
                fps_send_j(2,i,k) = f_post(9,i,1,k)
                fps_send_j(3,i,k) = f_post(10,i,1,k)
                fps_send_j(4,i,k) = f_post(16,i,1,k)
                fps_send_j(5,i,k) = f_post(18,i,1,k)
            enddo
        enddo

        call MPI_Send(fps_send_j, 5*nx*nz, MPI_DOUBLE_PRECISION, rank-nx_block, 3, MPI_COMM_WORLD, rc)
        call MPI_Recv(fps_recv_j, 5*nx*nz, MPI_DOUBLE_PRECISION, rank-nx_block, 4, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)

        ! scatter the received data (j=0)
        do k = 1, nz
            do i = 1, nx
                f_post(3,i,0,k) = fps_recv_j(1,i,k)
                f_post(7,i,0,k) = fps_recv_j(2,i,k)
                f_post(8,i,0,k) = fps_recv_j(3,i,k)
                f_post(15,i,0,k) = fps_recv_j(4,i,k)
                f_post(17,i,0,k) = fps_recv_j(5,i,k)
            enddo
        enddo
    endif

    ! /* exchange message with j=ny surface --- receive then send */
    if (block_y < ny_block) then
        ! collect data to send (j=ny)
        do k = 1, nz
            do i = 1, nx
                fps_send_j(1,i,k) = f_post(3,i,ny,k)
                fps_send_j(2,i,k) = f_post(7,i,ny,k)
                fps_send_j(3,i,k) = f_post(8,i,ny,k)
                fps_send_j(4,i,k) = f_post(15,i,ny,k)
                fps_send_j(5,i,k) = f_post(17,i,ny,k)
            enddo
        enddo

        call MPI_Send(fps_send_j, 5*nx*nz, MPI_DOUBLE_PRECISION, rank+nx_block, 4, MPI_COMM_WORLD, rc)
        call MPI_Recv(fps_recv_j, 5*nx*nz, MPI_DOUBLE_PRECISION, rank+nx_block, 3, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)

        ! scatter the received data (j=ny+1)
        do k = 1, nz
            do i = 1, nx
                f_post(4,i,ny+1,k) = fps_recv_j(1,i,k)
                f_post(9,i,ny+1,k) = fps_recv_j(2,i,k)
                f_post(10,i,ny+1,k) = fps_recv_j(3,i,k)
                f_post(16,i,ny+1,k) = fps_recv_j(4,i,k)
                f_post(18,i,ny+1,k) = fps_recv_j(5,i,k)
            enddo
        enddo
    endif

    ! /* exchange message with k=1 surface --- send then receive */
    if (block_z > 1) then
        ! collect data to send (k=1)
        do j = 1, ny
            do i = 1, nx
                fps_send_k(1,i,j) = f_post(6,i,j,1)
                fps_send_k(2,i,j) = f_post(13,i,j,1)
                fps_send_k(3,i,j) = f_post(14,i,j,1)
                fps_send_k(4,i,j) = f_post(17,i,j,1)
                fps_send_k(5,i,j) = f_post(18,i,j,1)
            enddo
        enddo

        call MPI_Send(fps_send_k, 5*nx*ny, MPI_DOUBLE_PRECISION, rank - nx_block*ny_block, 5, MPI_COMM_WORLD, rc)
        call MPI_Recv(fps_recv_k, 5*nx*ny, MPI_DOUBLE_PRECISION, rank - nx_block*ny_block, 6, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)

        ! scatter the received data (k=0)
        do j = 1, ny
            do i = 1, nx
                f_post(5,i,j,0) = fps_recv_k(1,i,j)
                f_post(11,i,j,0) = fps_recv_k(2,i,j)
                f_post(12,i,j,0) = fps_recv_k(3,i,j)
                f_post(15,i,j,0) = fps_recv_k(4,i,j)
                f_post(16,i,j,0) = fps_recv_k(5,i,j)
            enddo
        enddo
    endif

    ! /* exchange message with k=nz surface --- receive then send */
    if (block_z < nz_block) then
        ! collect data to send (k=nz)
        do j = 1, ny
            do i = 1, nx
                fps_send_k(1,i,j) = f_post(5,i,j,nz)
                fps_send_k(2,i,j) = f_post(11,i,j,nz)
                fps_send_k(3,i,j) = f_post(12,i,j,nz)
                fps_send_k(4,i,j) = f_post(15,i,j,nz)
                fps_send_k(5,i,j) = f_post(16,i,j,nz)
            enddo
        enddo

        call MPI_Send(fps_send_k, 5*nx*ny, MPI_DOUBLE_PRECISION, rank + nx_block*ny_block, 6, MPI_COMM_WORLD, rc)
        call MPI_Recv(fps_recv_k, 5*nx*ny, MPI_DOUBLE_PRECISION, rank + nx_block*ny_block, 5, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)

        ! scatter the received data(k=nz+1)
        do j = 1, ny
            do i = 1, nx
                f_post(6,i,j,nz+1) = fps_recv_k(1,i,j)
                f_post(13,i,j,nz+1) = fps_recv_k(2,i,j)
                f_post(14,i,j,nz+1) = fps_recv_k(3,i,j)
                f_post(17,i,j,nz+1) = fps_recv_k(4,i,j)
                f_post(18,i,j,nz+1) = fps_recv_k(5,i,j)
            enddo
        enddo
    endif

    !!! exchange messages with boundary lines -- message tag(for sender), velocity index in velocity model
    ! /* exchange with (nx, ny, k), message tag:7 */
    if (block_x < nx_block .AND. block_y < ny_block) then
        ! collect data to send (i=nx, j=ny)
        do k = 1, nz
            fpl_send_z(k) = f_post(7, nx, ny, k)
        enddo

        call MPI_Send(fpl_send_z, nz, MPI_DOUBLE_PRECISION, rank + nx_block + 1, 7, MPI_COMM_WORLD, rc)
        call MPI_Recv(fpl_recv_z, nz, MPI_DOUBLE_PRECISION, rank + nx_block + 1, 10, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)

        ! scatter the received data (i=nx+1, j=ny+1)
        do k = 1, nz
            f_post(10, nx+1, ny+1, k) = fpl_recv_z(k)
        enddo
    endif

    ! /* exchange with (1, ny, k), message tag:8 */
    if (block_x > 1 .AND. block_y < ny_block) then
        ! collect data to send (i=1, j=ny)
        do k = 1, nz
            fpl_send_z(k) = f_post(8, 1, ny, k)
        enddo

        call MPI_Send(fpl_send_z, nz, MPI_DOUBLE_PRECISION, rank + nx_block - 1, 8, MPI_COMM_WORLD, rc)
        call MPI_Recv(fpl_recv_z, nz, MPI_DOUBLE_PRECISION, rank + nx_block - 1, 9, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)

        ! scatter the received data (i=0, j=ny+1)
        do k = 1, nz
            f_post(9, 0, ny+1, k) = fpl_recv_z(k)
        enddo
    endif

    ! /* exchange with (nx, 1, k), message tag:9 */
    if (block_x < nx_block .AND. block_y > 1) then
        ! collect data to send (i=nx, j=1)
        do k = 1, nz
            fpl_send_z(k) = f_post(9, nx, 1, k)
        enddo

        call MPI_Send(fpl_send_z, nz, MPI_DOUBLE_PRECISION, rank - nx_block + 1, 9, MPI_COMM_WORLD, rc)
        call MPI_Recv(fpl_recv_z, nz, MPI_DOUBLE_PRECISION, rank - nx_block + 1, 8, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)

        ! scatter the received data (i=nx+1, j=0)
        do k = 1, nz
            f_post(8, nx+1, 0, k) = fpl_recv_z(k)
        enddo
    endif

    ! /* exchange with (1, 1, k), message tag:10 */
    if (block_x > 1 .AND. block_y > 1) then
        ! collect data to send (i=1, j=1)
        do k = 1, nz
            fpl_send_z(k) = f_post(10, 1, 1, k)
        enddo

        call MPI_Send(fpl_send_z, nz, MPI_DOUBLE_PRECISION, rank - nx_block - 1, 10, MPI_COMM_WORLD, rc)
        call MPI_Recv(fpl_recv_z, nz, MPI_DOUBLE_PRECISION, rank - nx_block - 1, 7, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)

        ! scatter the received data (i=0, j=0)
        do k = 1, nz
            f_post(7, 0, 0, k) = fpl_recv_z(k)
        enddo
    endif

    ! /* exchange with (nx, j, nz), message tag:11 */
    if (block_x < nx_block .AND. block_z < nz_block) then
        ! collect data to send (i=nx, k=nz)
        do j = 1, ny
            fpl_send_y(j) = f_post(11, nx, j, nz)
        enddo

        call MPI_Send(fpl_send_y, ny, MPI_DOUBLE_PRECISION, rank + nx_block*ny_block + 1, 11, MPI_COMM_WORLD, rc)
        call MPI_Recv(fpl_recv_y, ny, MPI_DOUBLE_PRECISION, rank + nx_block*ny_block + 1, 14, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)

        ! scatter the received data (i=nx+1, k=nz+1)
        do j = 1, ny
            f_post(14, nx+1, j, nz+1) = fpl_recv_y(j)
        enddo
    endif

    ! /* exchange with (1, j, nz), message tag:12 */
    if (block_x > 1 .AND. block_z < nz_block) then
        ! collect data to send (i=1, k=nz)
        do j = 1, ny
            fpl_send_y(j) = f_post(12, 1, j, nz)
        enddo

        call MPI_Send(fpl_send_y, ny, MPI_DOUBLE_PRECISION, rank + nx_block*ny_block - 1, 12, MPI_COMM_WORLD, rc)
        call MPI_Recv(fpl_recv_y, ny, MPI_DOUBLE_PRECISION, rank + nx_block*ny_block - 1, 13, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)

        ! scatter the received data (i=0, k=nz+1)
        do j = 1, ny
            f_post(13, 0, j, nz+1) = fpl_recv_y(j)
        enddo
    endif

    ! /* exchange with (nx, j, 1), message tag:13 */
    if (block_x < nx_block .AND. block_z > 1) then
        ! collect data to send (i=nx, k=1)
        do j = 1, ny
            fpl_send_y(j) = f_post(13, nx, j, 1)
        enddo

        call MPI_Send(fpl_send_y, ny, MPI_DOUBLE_PRECISION, rank - nx_block*ny_block + 1, 13, MPI_COMM_WORLD, rc)
        call MPI_Recv(fpl_recv_y, ny, MPI_DOUBLE_PRECISION, rank - nx_block*ny_block + 1, 12, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)

        ! scatter the received data (i=nx+1, k=0)
        do j = 1, ny
            f_post(12, nx+1, j, 0) = fpl_recv_y(j)
        enddo
    endif

    ! /* exchange with (1, j, 1), message tag:14 */
    if (block_x > 1 .AND. block_z > 1) then
        ! collect data to send (i=1, k=1)
        do j = 1, ny
            fpl_send_y(j) = f_post(14, 1, j, 1)
        enddo

        call MPI_Send(fpl_send_y, ny, MPI_DOUBLE_PRECISION, rank - nx_block*ny_block - 1, 14, MPI_COMM_WORLD, rc)
        call MPI_Recv(fpl_recv_y, ny, MPI_DOUBLE_PRECISION, rank - nx_block*ny_block - 1, 11, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)

        ! scatter the received data (i=0, k=0)
        do j = 1, ny
            f_post(11, 0, j, 0) = fpl_recv_y(j)
        enddo
    endif

    ! /* exchange with (i, ny, nz), message tag:15 */
    if (block_y < ny_block .AND. block_z < nz_block) then
        ! collect data to send (j=ny, k=nz)
        do i = 1, nx
            fpl_send_x(i) = f_post(15, i, ny, nz)
        enddo

        call MPI_Send(fpl_send_x, nx, MPI_DOUBLE_PRECISION, rank + nx_block + nx_block*ny_block, 15, MPI_COMM_WORLD, rc)
        call MPI_Recv(fpl_recv_x, nx, MPI_DOUBLE_PRECISION, rank + nx_block + nx_block*ny_block, 18, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)

        ! scatter the received data (j=ny+1, k=nz+1)
        do i = 1, nx
            f_post(18, i, ny+1, nz+1) = fpl_recv_x(i)
        enddo
    endif

    ! /* exchange with (i, 1, nz), message tag:16 */
    if (block_y > 1 .AND. block_z < nz_block) then
        ! collect data to send (j=1, k=nz)
        do i = 1, nx
            fpl_send_x(i) = f_post(16, i, 1, nz)
        enddo

        call MPI_Send(fpl_send_x, nx, MPI_DOUBLE_PRECISION, rank - nx_block + nx_block*ny_block, 16, MPI_COMM_WORLD, rc)
        call MPI_Recv(fpl_recv_x, nx, MPI_DOUBLE_PRECISION, rank - nx_block + nx_block*ny_block, 17, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)

        ! scatter the received data (j=0, k=nz+1)
        do i = 1, nx
            f_post(17, i, 0, nz+1) = fpl_recv_x(i)
        enddo
    endif

    ! /* exchange with (i, ny, 1), message tag:17 */
    if (block_y < ny_block .AND. block_z > 1) then
        ! collect data to send (j=ny, k=1)
        do i = 1, nx
            fpl_send_x(i) = f_post(17, i, ny, 1)
        enddo

        call MPI_Send(fpl_send_x, nx, MPI_DOUBLE_PRECISION, rank + nx_block - nx_block*ny_block, 17, MPI_COMM_WORLD, rc)
        call MPI_Recv(fpl_recv_x, nx, MPI_DOUBLE_PRECISION, rank + nx_block - nx_block*ny_block, 16, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)

        ! scatter the received data (j=ny+1, k=0)
        do i = 1, nx
            f_post(16, i, ny+1, 0) = fpl_recv_x(i)
        enddo
    endif

    ! /* exchange with (i, 1, 1), message tag:18 */
    if (block_y > 1 .AND. block_z > 1) then
        ! collect data to send (j=1, k=1)
        do i = 1, nx
            fpl_send_x(i) = f_post(18, i, 1, 1)
        enddo

        call MPI_Send(fpl_send_x, nx, MPI_DOUBLE_PRECISION, rank - nx_block - nx_block*ny_block, 18, MPI_COMM_WORLD, rc)
        call MPI_Recv(fpl_recv_x, nx, MPI_DOUBLE_PRECISION, rank - nx_block - nx_block*ny_block, 15, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)

        ! scatter the received data (j=0, k=0)
        do i = 1, nx
            f_post(15, i, 0, 0) = fpl_recv_x(i)
        enddo
    endif

end subroutine message_passing

subroutine output()
    use mpi
    use commondata
    implicit none
    integer :: i, j, k
    integer :: p_rank, num(1:6) ,dx = 0, dy = 0, dz = 0
    real(8), allocatable :: total_u(:, :, :), total_v(:, :, :), total_w(:, :, :), total_rho(:, :, :)
    real(8), allocatable :: tmp_u(:, :, :), tmp_v(:, :, :), tmp_w(:, :, :), tmp_rho(:, :, :)

    ! use rank 0 to receive data and output the results

    if (rank > 0) then  !!! ----  rank != 0 send data
        ! collect the rank information
        num(1) = nx
        num(2) = ny
        num(3) = nz
        num(4) = block_x
        num(5) = block_y
        num(6) = block_z
        ! send to rank 0
        call MPI_Send(num, 6, MPI_INTEGER, 0, 0, MPI_COMM_WORLD, rc)    ! rank information
        call MPI_Send(u, nx*ny*nz, MPI_DOUBLE_PRECISION, 0, 1, MPI_COMM_WORLD, rc)
        call MPI_Send(v, nx*ny*nz, MPI_DOUBLE_PRECISION, 0, 2, MPI_COMM_WORLD, rc)
        call MPI_Send(w, nx*ny*nz, MPI_DOUBLE_PRECISION, 0, 3, MPI_COMM_WORLD, rc)
        call MPI_Send(rho, nx*ny*nz, MPI_DOUBLE_PRECISION, 0, 4, MPI_COMM_WORLD, rc)
    else    !!! ---- rank 0 collect data
        ! allocate array
        allocate(total_u(total_nx, total_ny, total_nz))
        allocate(total_v(total_nx, total_ny, total_nz))
        allocate(total_w(total_nx, total_ny, total_nz))
        allocate(total_rho(total_nx, total_ny, total_nz))

        ! collect data from rank 0
        do k = 1, nz
            do j = 1, ny
                do i = 1, nx
                    total_u(i, j, k) = u(i, j, k)
                    total_v(i, j, k) = v(i, j, k)
                    total_w(i, j, k) = w(i, j, k)
                    total_rho(i, j, k) = rho(i, j, k)
                enddo
            enddo
        enddo

        ! collect data from all other processors
        do p_rank = 1, num_process-1
            ! receive the rank information
            call MPI_Recv(num, 6, MPI_INTEGER, p_rank, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)

            ! creat buffer
            allocate(tmp_u(num(1), num(2), num(3)))
            allocate(tmp_v(num(1), num(2), num(3)))
            allocate(tmp_w(num(1), num(2), num(3)))
            allocate(tmp_rho(num(1), num(2), num(3)))

            ! receive data
            call MPI_Recv(tmp_u, num(1) * num(2) * num(3), MPI_DOUBLE_PRECISION, p_rank, 1, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)
            call MPI_Recv(tmp_v, num(1) * num(2) * num(3), MPI_DOUBLE_PRECISION, p_rank, 2, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)
            call MPI_Recv(tmp_w, num(1) * num(2) * num(3), MPI_DOUBLE_PRECISION, p_rank, 3, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)
            call MPI_Recv(tmp_rho, num(1) * num(2) * num(3), MPI_DOUBLE_PRECISION, p_rank, 4, MPI_COMM_WORLD, MPI_STATUS_IGNORE, rc)

            ! determine the origin
            dx = num(1) * (num(4) - 1)  !! nx * (block_x - 1)
            if (num(1) == total_nx / nx_block) then
                dx = dx + MOD(total_nx, nx_block)
            endif

            dy = num(2) * (num(5) - 1)  !! ny * (block_y - 1)
            if (num(2) == total_ny / ny_block) then
                dy = dy + MOD(total_ny, ny_block)
            endif

            dz = num(3) * (num(6) - 1)  !! nz * (block_z - 1)
            if (num(3) == total_nz / nz_block) then
                dz = dz + MOD(total_nz, nz_block)
            endif

            ! assign data
            do k = 1, num(3)
                do j = 1, num(2)
                    do i = 1, num(1)
                        total_u(i + dx, j + dy, k + dz) = tmp_u(i, j, k)
                        total_v(i + dx, j + dy, k + dz) = tmp_v(i, j, k)
                        total_w(i + dx, j + dy, k + dz) = tmp_w(i, j, k)
                        total_rho(i + dx, j + dy, k + dz) = tmp_rho(i, j, k)
                    enddo
                enddo
            enddo

            ! de-allocate buffer arrays
            deallocate(tmp_u)
            deallocate(tmp_v)
            deallocate(tmp_w)
            deallocate(tmp_rho)
        enddo

        ! after collect total_* data, then output
        ! call output_ASCII(xp, yp, zp, total_u, total_v, total_w, total_rho, total_nx, total_ny, total_nz, itc)
        ! call output_binary(total_u, total_v, total_w, total_rho, total_nx, total_ny, total_nz, itc)
        call output_Tecplot(xp, yp, zp, total_u, total_v, total_w, total_rho, total_nx, total_ny, total_nz, itc)
        call getVelocity(xp, yp, zp, total_u, total_v, total_w, total_nx, total_ny, total_nz, U0, itc)

        ! de-allocate total arrays
        deallocate(total_u)
        deallocate(total_v)
        deallocate(total_w)
        deallocate(total_rho)
    endif

end subroutine output

subroutine output_ASCII(xp, yp, zp, u, v, w, rho, nx, ny, nz, itc)
    implicit none
    integer, intent(in) :: nx, ny, nz, itc
    real(8), intent(in) :: xp(0:nx+1), yp(0:ny+1), zp(0:nz+1)
    real(8), intent(in) :: u(nx, ny, nz), v(nx, ny, nz), w(nx, ny, nz), rho(nx, ny, nz)
    integer :: i, j, k
    character(len=100) :: filename

    write(filename,*) itc
    filename = adjustl(filename)

    open(unit=77,file='MRTcavity-'//trim(filename)//'.dat',status='unknown')
        write(77,*) 'TITLE="Lid Driven Cavity"'
        write(77,*) 'VARIABLES="X" "Y" "Z" "U" "V" "W" "Pressure" '
        write(77,101) nx, ny, nz
        do k = 1, nz
            do j = 1, ny
                do i = 1, nx
                    write(77,100) xp(i), yp(j), zp(k), u(i,j,k), v(i,j,k), w(i,j,k), rho(i,j,k)/3.0d0
                enddo
            enddo
        enddo
100     format(1x,3(e11.4,' '),10(e13.6,' '))
101     format('ZONE',1x,'I=',1x,i5,2x,'J=',1x,i5,1x,'K=',1x,i5,1x,'F=POINT')
    close(77)

    return
end subroutine output_ASCII

!!!c--------------------------------
!!!c--------------------------------
subroutine output_Tecplot(xp, yp, zp, u, v, w, rho, nx, ny, nz, itc)
    implicit none
    integer, intent(in) :: nx, ny, nz, itc
    real(8), intent(in) :: xp(0:nx+1), yp(0:ny+1), zp(0:nz+1)
    real(8), intent(in) :: u(nx, ny, nz), v(nx, ny, nz), w(nx, ny, nz), rho(nx, ny, nz)
    integer :: i, j, k
    character(len=9) :: B2
    REAL(4) :: zoneMarker, eohMarker
    character(len=40) :: title
    character(len=40) :: V1,V2,V3,V4,V5,V6,V7
    character(len=40) :: zoneName

    write(B2,'(i9.9)') itc

    open(41,file='MRTcavity-'//B2//'.plt',form='binary')
    !---------------------------------------------
    zoneMarker= 299.0
    eohMarker = 357.0

    !I. HEAD SECTION--------------------------------------
    !c--Magic number, Version number
    write(41) "#!TDV101"

    !c--Integer value of 1
    write(41) 1

    Title="MyFirst"
    call dumpstring(title)

    !c-- Number of variables in this data file (here 5 variables)
    write(41) 7

    !c-- Variable names.
    V1 = 'X'
    call dumpstring(V1)
    V2 = 'Y'
    call dumpstring(V2)
    V3 = 'Z'
    call dumpstring(V3)
    V4 = 'U'
    call dumpstring(V4)
    V5 = 'V'
    call dumpstring(V5)
    V6 = 'W'
    call dumpstring(V6)
    V7 = 'Pressure'
    call dumpstring(V7)

    !c-----Zones-----------------------------

    !c--------Zone marker. Value = 299.0
    write(41) zoneMarker

    !--------Zone name.
    zoneName="ZONE 001"
    call dumpstring(zoneName)

    !---------Zone Color
    write(41) -1

    !---------ZoneType
    write(41) 0

    !---------DataPacking 0=Block, 1=Point
    write(41) 1

    !---------Specify Var Location. 0 = Do not specify, all data
    !---------is located at the nodes. 1 = Specify
    write(41) 0

    !---------Number of user defined face neighbor connections
    ! (value >= 0)
    write(41) 0

    !---------IMax,JMax,KMax
    write(41) nx
    write(41) ny
    write(41) nz

    !-----------1=Auxiliary name/value pair to follow
    !-----------0=No more Auxiliar name/value pairs.
    write(41) 0
    write(41) eohMarker

    !----zone ------------------------------------------------------------
    write(41) zoneMarker

    !--------variable data format, 1=Float, 2=Double, 3=LongInt,4=ShortInt, 5=Byte, 6=Bit
    write(41) 1
    write(41) 1
    write(41) 1
    write(41) 1
    write(41) 1
    write(41) 1
    write(41) 1

    !--------Has variable sharing 0 = no, 1 = yes.
    write(41) 0

    !----------Zone number to share connectivity list with (-1 = no
    ! sharing).
    write(41) -1

    !---------------------------------------------------------------------
    do k=1,nz
        do j=1,ny
            do i=1,nx
                write(41) real(xp(i))
                write(41) real(yp(j))
                write(41) real(zp(k))
                write(41) real(u(i,j,k))
                write(41) real(v(i,j,k))
                write(41) real(w(i,j,k))
                write(41) real(rho(i,j,k)/3.0d0)
            enddo
        enddo
    enddo
    close(41)
    !---------------------------------------------------------------------

    return
    end subroutine output_Tecplot
!!!c--------------------------------
    subroutine dumpstring(instring)
    implicit none
    character(len=40) instring
    integer :: stringLength
    integer :: ii
    integer :: I

    stringLength=LEN_TRIM(instring)
    do ii=1,stringLength
        I=ICHAR(instring(ii:ii))
        write(41) I
    end do
    write(41) 0

    return
    end subroutine dumpstring
!!!c--------------------------------
!!!c--------------------------------

subroutine getVelocity(xp, yp, zp, u, v, w, nx, ny, nz, U0, itc)
    implicit none
    integer, intent(in) :: nx, ny, nz, itc
    real(8), intent(in) :: xp(0:nx+1), yp(0:ny+1), zp(0:nz+1)
    real(8), intent(in) :: u(nx, ny, nz), v(nx, ny, nz), w(nx, ny, nz)
    real(8), intent(in) :: U0
    integer :: i, j, k
    integer :: nxHalf, nyHalf, nzHalf
    character(len=9) :: B2

    write(B2,'(i9.9)') itc

    nxHalf = (nx - 1) / 2 + 1
    nyHalf = (ny - 1) / 2 + 1
    nzHalf = (nz - 1) / 2 + 1

    open(unit=02,file='./u-z_'//B2//'.dat',status='unknown')
    do k=1,nz
        write(02,*) u(nxHalf, nyHalf, k)/U0, zp(k)/dble(nz)
    enddo
    close(02)

    open(unit=03,file='./x-w_'//B2//'.dat',status='unknown')
    do i=1,nx
        write(03,*) xp(i)/dble(nx), w(i, nyHalf, nzHalf)/U0
    enddo
    close(03)

    return
end subroutine getVelocity

subroutine output_binary(u, v, w, rho, nx, ny, nz, itc)
    implicit none
    integer, intent(in) :: nx, ny, nz, itc
    real(8), intent(in) :: u(nx, ny, nz), v(nx, ny, nz), w(nx, ny, nz), rho(nx, ny, nz)
    integer :: i, j, k
    character(len=100) :: filename

    write(filename,*) itc
    filename = adjustl(filename)

    open(unit=01,file='MRTcavity-'//trim(filename)//'.bin',form="unformatted",access="sequential")
    write(01) (((u(i,j,k), i = 1, nx), j = 1, ny), k = 1, nz)
    write(01) (((v(i,j,k), i = 1, nx), j = 1, ny), k = 1, nz)
    write(01) (((rho(i,j,k), i = 1, nx), j = 1, ny), k = 1, nz)
    close(01)

    return
end subroutine output_binary


! { dg-do run { xfail spu-*-* } }
! FAILs on SPU because of rounding error reading kinds.h
program main
  ! Test scalar masks for different intrinsics.
  real, dimension(2,2) :: a
  logical(kind=2) :: lo
  lo = .false.
  a(1,1) = 1.
  a(1,2) = -1.
  a(2,1) = 13.
  a(2,2) = -31.
  if (any (minloc (a, lo) /= 0)) STOP 1
  if (any (minloc (a, .true.) /= (/ 2, 2 /))) STOP 2
  if (any (minloc(a, 1, .true.) /= (/ 1, 2/))) STOP 3
  if (any (minloc(a, 1, lo ) /= (/ 0, 0/))) STOP 4

  if (any (maxloc (a, lo) /= 0)) STOP 5
  if (any (maxloc (a, .true.) /= (/ 2,1 /))) STOP 6
  if (any (maxloc(a, 1, .true.) /= (/ 2, 1/))) STOP 7
  if (any (maxloc(a, 1, lo) /= (/ 0, 0/))) STOP 8

  if (any (maxval(a, 1, lo) /= -HUGE(a))) STOP 9
  if (any (maxval(a, 1, .true.) /= (/13., -1./))) STOP 10
  if (any (minval(a, 1, lo) /= HUGE(a))) STOP 11
  if (any (minval(a, 1, .true.) /= (/1., -31./))) STOP 12

  if (any (product(a, 1, .true.) /= (/13., 31./))) STOP 13
  if (any (product(a, 1, lo ) /= (/1., 1./))) STOP 14

  if (any (sum(a, 1, .true.) /= (/14., -32./))) STOP 15
  if (any (sum(a, 1, lo) /= (/0., 0./))) STOP 16

end program main


! { dg-do run }
! PR 25577
! MVBITS didn't work correctly for integer types wider than a C int
! The testcase is based on the one Dale Ranta posted in the bug report
implicit none
integer(1) i1,j1
integer(2) i2,j2
integer(4) i4,j4
integer(8) i8,j8
integer ibits,n

ibits=bit_size(1_1)
do n=1,ibits
   i1=-1
   call mvbits(1_1, 0,n,i1,0)
   j1=-1-2_1**n+2
   if(i1.ne.j1)call abort
enddo
ibits=bit_size(1_2)
do n=1,ibits
   i2=-1
   call mvbits(1_2, 0,n,i2,0)
   j2=-1-2_2**n+2
   if(i2.ne.j2)call abort
enddo
ibits=bit_size(1_4)
do n=1,ibits
   i4=-1
   call mvbits(1_4, 0,n,i4,0)
   j4=-1-2_4**n+2
   if(i4.ne.j4)call abort
enddo
ibits=bit_size(1_8)
do n=1,ibits
   i8=-1
   call mvbits(1_8, 0,n,i8,0)
   j8=-1-2_8**n+2
   if(i8.ne.j8)call abort
enddo
end


!> \file general_func.f90
!> Contains module general_functions.

!>General equations that are used by hydrological models
MODULE GENERAL_FUNCTIONS

!Copyright 2012,2014,2016-2017 SMHI
!
!This file is part of HYPE.
!HYPE is free software: you can redistribute it and/or modify it under the terms of the Lesser GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version.
!HYPE is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the Lesser GNU General Public License for more details.
!You should have received a copy of the Lesser GNU General Public License along with HYPE. If not, see <http://www.gnu.org/licenses/>.

!-----------------------------------------------------------------------------------------

  IMPLICIT NONE
  PRIVATE
  PUBLIC :: tempfactor, &
            halfsatconcfactor, &
            moisturefactor, &
            exponential_moisturefactor, &
            simple_rating_curve,  &
            exponential_decay,  &
            sedimentation

CONTAINS

  ! >\brief Calculates a temperature factor
  !! Based on: Q10=2,reference rate is at 20 degrees and with
  !! thresholds at 0 and 5 degrees (from COUP-model).
  ! --------------------------------------------------------------
  REAL FUNCTION tempfactor(temp)

    REAL, INTENT(IN)  :: temp     !<current temperature
    !Local variables
    REAL fcn

    fcn = 2**((temp - 20.0) / 10.0)
    IF(temp < 5.0) fcn = fcn * (temp / 5.0)
    IF(temp < 0.)  fcn = 0.
    tempfactor = fcn

  END FUNCTION tempfactor

  !>\brief Calculates a concentration dependence factor
  !! Based on: half saturation function
  !----------------------------------------------------------------
  REAL FUNCTION halfsatconcfactor(conc,par)

    REAL, INTENT(IN)  :: conc     !<current concentration
    REAL, INTENT(IN)  :: par      !<half saturation concentration
    !Local variables
    REAL fcn

    fcn = conc / (conc + par)
    halfsatconcfactor = fcn

  END FUNCTION halfsatconcfactor

  !> Calculates a soil moisture dependence factor
  !-----------------------------------------------------------------------------
  REAL FUNCTION moisturefactor(sm,wp,pw,thickm,satsmf,exp,thetalow,thetaupp)

    REAL, INTENT(IN)  :: sm     !<soil moisture (mm)
    REAL, INTENT(IN)  :: wp     !<wilting point pore wolume (mm)
    REAL, INTENT(IN)  :: pw     !<total pore wolume (mm)
    REAL, INTENT(IN)  :: thickm !<thickness of soil layer (m)
    REAL, INTENT(IN)  :: satsmf !<saturated moisturefactor (satact)
    REAL, INTENT(IN)  :: exp    !<exponent of moisturefactor (thetapow)
    REAL, INTENT(IN)  :: thetalow !<low(?) moisture coefficient (thetalow)
    REAL, INTENT(IN)  :: thetaupp !<high(?) moisture coefficient (thetaupp)

    ! Local variables
    REAL thickness  !thickness of soil layer (mm)
    REAL smfcn      !soil moisture dependence factor

    !Initiations
    smfcn = 0.
    thickness = thickm * 1000.      !mm

    !Caclulate soilmoisture function value
    IF(thickness>0)THEN
      IF(sm >= pw) THEN
        smfcn = satsmf
      ELSEIF(sm <= wp) THEN
        smfcn=0.
      ELSE
        smfcn = min(1., (1-satsmf)*((pw-sm)/((thetaupp/100.)*thickness))**exp + satsmf, ((sm-wp)/((thetalow/100.)*thickness))**exp)
      ENDIF
    ENDIF

    moisturefactor = smfcn

  END FUNCTION moisturefactor

  !>Calculates an exponential soil moisture dependence factor
  !-----------------------------------------------------------------------------
  REAL FUNCTION exponential_moisturefactor(sm,pw,limpar,exp)

    REAL, INTENT(IN)  :: sm     !<soil moisture (mm)
    REAL, INTENT(IN)  :: pw     !<total pore wolume (mm)
    REAL, INTENT(IN)  :: limpar !<limitation parameter of moisturefactor (mm)
    REAL, INTENT(IN)  :: exp    !<exponent of moisturefactor

    ! Local variables
    REAL smfcn

    !Initiations
    smfcn = 0.

    !Calculate soilmoisture function value
    IF((sm/pw) > limpar)THEN
      smfcn = (((sm / pw)-limpar)/(1.0-limpar))**exp
    ENDIF

    exponential_moisturefactor = smfcn

  END FUNCTION exponential_moisturefactor

  !>Calculates momentanous flow from current water level with simple
  !>rating curve equation: Q = k*(w-w0)**p
  !-----------------------------------------------------------------------------
  REAL FUNCTION simple_rating_curve(wst,k,p,w0)

    !Argument declarations
    REAL, INTENT(IN) :: wst    !<current water level (m)
    REAL, INTENT(IN) :: k      !<rating curve coefficient
    REAL, INTENT(IN) :: p      !<rating curve exponent
    REAL, INTENT(IN) :: w0     !<rating curve threshold (m)

    !Local variables
    REAL flow      !discharge (m3/s)

    !Calculate discharge
    IF(wst > w0)THEN
      flow = k * (wst-w0)**p
    ELSE
      flow = 0.
    ENDIF

    simple_rating_curve = flow

  END FUNCTION simple_rating_curve

  !>Calculates factor of exponential decay from halftime
  !-----------------------------------------------------------------------------
  REAL FUNCTION exponential_decay(timesteps,halftime)

    !Argument declarations
    REAL, INTENT(IN) :: timesteps     !<in unit of half time
    REAL, INTENT(IN) :: halftime      !<half time of decay

    !Local variables
    REAL factor      !fraction left

    !Calculate discharge
    IF(halftime>0)THEN
      factor = EXP(-LOG(2.)*timesteps/halftime)
    ELSE
      WRITE(6,*) 'ERROR: half time parameter is zero. Check parameter file.'
      STOP 1
    ENDIF

    exponential_decay = factor

  END FUNCTION exponential_decay

  !>\brief Calculate sedimentation from sinking velocity
  !-----------------------------------------------------------------
  SUBROUTINE sedimentation(depth,conc,vel)

    !Argument declarations
    REAL, INTENT(IN)              :: depth      !<water depth (m)
    REAL, INTENT(INOUT)           :: conc       !<concentration of substance to sedimentate (mg/L or something)
    REAL, INTENT(IN)              :: vel        !<sinking velocity (m/timestep)

    !Local variables
    REAL pool      !pool in water (g/m2)
    REAL sedloss   !change (g/m2/timestep)

    !Calculate pool
    pool = depth * conc

    !Calculate sedimentation
    sedloss = vel * conc

    !Remove sedimentation from the pool and calculate new concentration
    IF(pool>sedloss)THEN
      pool = pool - sedloss
      conc = pool / depth
    ELSE
      conc = 0.
    ENDIF

  END SUBROUTINE sedimentation

  END MODULE GENERAL_FUNCTIONS


      double complex function A51mpmpp(j1,j2,j3,j4,j5,za,zb)
      implicit none
      include 'constants.f'
      include 'sprods_com.f'
      include 'zprods_decl.f'
      include 'epinv.f'
      include 'scale.f'
      double complex Vf,Vs,Ff,Fs,L0,L1,L2,Ls1,miA5tree,lnrat
      integer j1,j2,j3,j4,j5
C----Eq.(9,10) of hep-ph/9302280v1 of BDK multiplied by 16*pi^2*(-i)
C--- to give (16*pi^2)*(-i)*A^{[1/2]}_{5;1}
      miA5tree=za(j1,j3)**4
     & /(za(j1,j2)*za(j2,j3)*za(j3,j4)*za(j4,j5)*za(j5,j1))
      Vf=dcmplx(-2.5d0*epinv-2d0)
     & -0.5d0*(lnrat(musq,-s(j3,j4))+lnrat(musq,-s(j5,j1)))
      Vs=-Vf/3d0+dcmplx(2d0/9d0)
      Ff=
     & -za(j1,j3)**2*za(j4,j1)*zb(j2,j4)**2
     & /(za(j4,j5)*za(j5,j1)*s(j5,j1)**2)
     & *Ls1(-s(j2,j3),-s(j5,j1),-s(j3,j4),-s(j5,j1))

     & +za(j1,j3)**2*za(j5,j3)*zb(j2,j5)**2
     & /(za(j3,j4)*za(j4,j5)*s(j3,j4)**2)
     & *Ls1(-s(j1,j2),-s(j3,j4),-s(j5,j1),-s(j3,j4))

     & -za(j1,j3)**3*L0(-s(j3,j4),-s(j5,j1))
     & *(za(j1,j5)*zb(j5,j2)*za(j2,j3)-za(j3,j4)*zb(j4,j2)*za(j2,j1))
     & /(2d0*za(j1,j2)*za(j2,j3)*za(j3,j4)*za(j4,j5)*za(j5,j1)*s(j5,j1))
      Fs=
     & -za(j1,j2)*za(j2,j3)*za(j3,j4)*(za(j4,j1)*zb(j2,j4))**2
     & /(za(j4,j5)*za(j5,j1)*za(j2,j4)**2*s(j5,j1)**2)
     & *(2d0*Ls1(-s(j2,j3),-s(j5,j1),-s(j3,j4),-s(j5,j1))
     & +L1(-s(j2,j3),-s(j5,j1))+L1(-s(j3,j4),-s(j5,j1)))

     & +za(j3,j2)*za(j2,j1)*za(j1,j5)*(za(j5,j3)*zb(j2,j5))**2
     & /(za(j5,j4)*za(j4,j3)*za(j2,j5)**2*s(j3,j4)**2)
     & *(2d0*Ls1(-s(j1,j2),-s(j3,j4),-s(j5,j1),-s(j3,j4))
     & +L1(-s(j1,j2),-s(j3,j4))+L1(-s(j5,j1),-s(j3,j4)))

     &+2d0*za(j2,j3)**2*(za(j4,j1)*zb(j2,j4))**3*L2(-s(j2,j3),-s(j5,j1))
     & /(3d0*za(j4,j5)*za(j5,j1)*za(j2,j4)*s(j5,j1)**3)
     &-2d0*za(j2,j1)**2*(za(j5,j3)*zb(j2,j5))**3*L2(-s(j1,j2),-s(j3,j4))
     & /(3d0*za(j5,j4)*za(j4,j3)*za(j2,j5)*s(j3,j4)**3)

     & +L2(-s(j3,j4),-s(j5,j1))/s(j5,j1)**3
     * *(za(j1,j3)*zb(j2,j4)*zb(j2,j5)/(3d0*za(j4,j5))
     & *(za(j1,j5)*zb(j5,j2)*za(j2,j3)-za(j3,j4)*zb(j4,j2)*za(j2,j1))
     & +2d0*(za(j1,j2)*za(j3,j4))**2*za(j4,j1)*zb(j2,j4)**3
     & /(3d0*za(j4,j5)*za(j5,j1)*za(j2,j4))
     & -2d0*(za(j3,j2)*za(j1,j5))**2*za(j5,j3)*zb(j2,j5)**3
     & /(3d0*za(j5,j4)*za(j4,j3)*za(j2,j5)))

     & +za(j1,j3)**3*L0(-s(j3,j4),-s(j5,j1))
     & /(6d0*za(j1,j2)*za(j2,j3)*za(j3,j4)*za(j4,j5)*za(j5,j1)*s(j5,j1))
     & *(za(j1,j5)*zb(j5,j2)*za(j2,j3)-za(j3,j4)*zb(j4,j2)*za(j2,j1))

     & +(zb(j2,j4)*zb(j2,j5))**2
     & /(3d0*zb(j1,j2)*zb(j2,j3)*zb(j3,j4)*za(j4,j5)*zb(j5,j1))

     & -za(j1,j2)*za(j4,j1)**2*zb(j2,j4)**3
     & /(3d0*za(j4,j5)*za(j5,j1)*za(j2,j4)*zb(j2,j3)*zb(j3,j4)*s(j5,j1))

     & +za(j3,j2)*za(j5,j3)**2*zb(j2,j5)**3
     & /(3d0*za(j5,j4)*za(j4,j3)*za(j2,j5)*zb(j2,j1)*zb(j1,j5)*s(j3,j4))

     & +za(j1,j3)**2*zb(j2,j4)*zb(j2,j5)
     & /(6d0*s(j3,j4)*za(j4,j5)*s(j5,j1))

      A51mpmpp=-(Vf+Vs)*miA5tree-(Ff+Fs)
      return
      end


      FUNCTION torflux_deriv (x)
      USE stel_kinds
!     USE vmec_input, ONLY: af => aphi
C-----------------------------------------------
C   D u m m y   A r g u m e n t s
C-----------------------------------------------
      REAL(rprec) :: x, torflux_deriv
C-----------------------------------------------
!     TEMPORARILY DISABLED....
!     torflux_deriv = af(1) + x*(2*af(2) + x*(3*af(3) + x*(4*af(4) +
!    1           x*(5*af(5) + x*(6*af(6) + x*(7*af(7) + x*(8*af(8) +
!    2           x*(9*af(9) + x*10*af(10)))))))))
      torflux_deriv = 1

      END FUNCTION torflux_deriv


module geom_utils
use config, only: gtol
use math
implicit none

private
public :: dehomogenize_array1d, dehomogenize_array2d, tessellate_cp_net
public :: local_to_global_param, global_to_local_param
public :: is_curve_flat, is_surface_flat
public :: equivalence_array1d, barycentric_params, are_points_equal
public :: is_surface_closed
public :: invert_point_in_triangle
public :: angle_between_vecs

contains

subroutine dehomogenize_array1d(n, cpw, cp, w)
    !> Convert the 1D homogeneous control points cpw to
    !> dehomogenized coordinates (x/w, y/w, z/w).
    !> n - Number of control points - 1.
    !> cpw - Homogeneous control points.
    !> cp - Dehomogenized control points.
    !> w - Weights.

    !f2py intent(in) n, cpw
    !f2py intent(out) cp, w
    !f2py depend(n) cpw, cp, w

    ! Input
    integer, intent(in) :: n
    double precision, intent(in) :: cpw(0:n, 0:3)

    ! Output
    double precision, intent(out) :: cp(0:n, 0:2), w(0:n)

    ! Working
    integer :: i

    do i = 0, n
        w(i) = cpw(i, 3)
        cp(i, :) = cpw(i, :2) / w(i)
    end do

end subroutine dehomogenize_array1d

subroutine dehomogenize_array2d(n, m, cpw, cp, w)
    !> Convert the 2D homogeneous control points cpw to
    !> dehomogenized coordinates (x/w, y/w, z/w).
    !> n - Number of control points - 1 in u-direction.
    !> m - Number of control points - 1 in v-direction.
    !> cpw - Homogeneous control points.
    !> cp - Dehomogenized control points.
    !> w - Weights.

    !f2py intent(in) n, m, cpw
    !f2py intent(out) cp, w
    !f2py depend(n, m) cpw, cp, w

    ! Input
    integer, intent(in) :: n, m
    double precision, intent(in) :: cpw(0:n, 0:m, 0:3)

    ! Output
    double precision, intent(out) :: cp(0:n, 0:m, 0:2)
    double precision, intent(out) :: w(0:n, 0:m)

    ! Working
    integer :: i, j

    do j = 0, m
        do i = 0, n
            w(i, j) = cpw(i, j, 3)
            cp(i, j, :) = cpw(i, j, :2) / w(i, j)
        end do
    end do

end subroutine dehomogenize_array2d

subroutine tessellate_cp_net(n, m, cp, verts, triangles)
    !> Tessellate the control net of the surface as triangles.
    !> n - Number of control points - 1 in u-direction.
    !> m - Number of control points - 1 in v-direction.
    !> cp - Dehomogenized control points.
    !> verts - Vertices of triangles.
    !> triangles - Connectivity of triangles where each row is a
    !> triangle and each column is a vertex id.

    ! Input
    integer, intent(in) :: n, m
    double precision, intent(in) :: cp(0:n, 0:m, 0:2)

    ! Output
    integer, intent(out) :: triangles(0:2 * n * m - 1, 0:2)
    double precision, intent(out) :: verts(0:(n + 1) * (m + 1) - 1, 0:2)

    ! Working
    integer :: i, j, t, v, v1, v2, v3, v4
    integer :: vert_id(0:n, 0:m)

    v = 0
    do j = 0, m
        do i = 0, n
            verts(v, :) = cp(i, j, :)
            vert_id(i, j) = v
            v = v + 1
        end do
    end do
    t = 0
    do j = 0, m - 1
        do i = 0, n - 1
            v1 = vert_id(i, j)
            v2 = vert_id(i + 1, j)
            v3 = vert_id(i + 1, j + 1)
            v4 = vert_id(i, j + 1)
            triangles(t, :) = (/ v1, v2, v3 /)
            triangles(t + 1, :) = (/ v1, v3, v4 /)
            t = t + 2
        end do
    end do
end subroutine tessellate_cp_net

function local_to_global_param(a, b, ul) result(ug)
    !> Convert the local parameter to the global domain
    !> (a <= u <= b).
    !> a - Lower global parameter.
    !> b - Upper global parameter.
    !> ul - Local parameter.
    !> ug - Global parameter.

    ! Input
    double precision, intent(in) :: a, b, ul

    ! Output
    double precision :: ug

    ! Working
    double precision :: ui

    ui = ul
    if (ui .lt. 0.0d0) then
        ui = 0.0d0
    else if (ui .gt. 1.0d0) then
        ui = 1.0d0
    end if
    ug = a + ui * (b - a)

end function local_to_global_param

function global_to_local_param(a, b, ug) result(ul)
    !> Convert the global parameter to the local domain
    !> (0 <= u <= 1).
    !> a - Lower global parameter.
    !> b - Upper global parameter.
    !> ug - Global parameter.
    !> ul - Local parameter.

    ! Input
    double precision, intent(in) :: a, b, ug

    ! Output
    double precision :: ul

    ul = (ul - a) / (b - a)
    if (ul .lt. 0.0d0) then
        ul = 0.0d0
    else if (ul .gt. 1.0d0) then
        ul = 1.0d0
    end if

end function global_to_local_param

subroutine is_curve_flat(n, cp, tol, is_flat)
    !> Check flatness of curve.
    !> n - Number of control points - 1.
    !> cp - Dehomogenized control points.
    !> tol - Tolerance.
    !> is_flat - True if curve is flat, False if not.

    !f2py intent(in) n, cp, tol
    !f2py intent(out) is_flat
    !f2py depend(n) cp

    ! Input
    integer, intent(in) :: n
    double precision, intent(in) ::  tol
    double precision, intent(in) :: cp(0:n, 0:2)

    ! Output
    logical, intent(out) :: is_flat

    ! Working
    integer :: i
    double precision :: d, dline
    double precision :: v0(3), v1(3), v2(3), v3(3), vec(3)

    is_flat = .true.
    vec = cp(n, :) - cp(0, :)
    dline = norm(vec)

    if (dline .le. 0.0d0) then
        ! Curve must be closed. Use the distnace between points.
        do i = 1, n - 1
            v0 = cp(i, :) - cp(0, :)
            d = norm(v0)
            if (d .gt. tol) then
                is_flat = .false.
                return
            end if
        end do
        is_flat = .true.
        return
    end if

    do i = 1, n - 1
        v0 = cp(i, :) - cp(0, :)
        v1 = cp(i, :) - cp(n, :)
        v2 = cross(v0, v1)
        d = norm(v2) / dline
        if (d .gt. tol) then
            is_flat = .false.
            return
        end if
    end do
end subroutine is_curve_flat

subroutine is_surface_flat(n, m, cp, tol, is_flat)
    !> Check flatness of surface.
    !> n - Number of control points - 1 in u-direction.
    !> m - Number of control points - 1 in v-direction.
    !> cp - Dehomogenized control points.
    !> tol - Tolerance.
    !> is_flat - True if surface is flat, False if not.

    !f2py intent(in) n, m, cp, tol
    !f2py intent(out) is_flat
    !f2py depend(n, m) cp

    ! Input
    integer, intent(in) :: n, m
    double precision, intent(in) :: tol
    double precision, intent(in) :: cp(0:n, 0:m, 0:2)

    ! Output
    logical, intent(out) :: is_flat

    ! Working
    integer :: i, j
    double precision :: row(0:m, 0:2), col(0:n, 0:2)

    is_flat = .true.

    ! Columns
    do j = 0, m
        col = cp(:, j, :)
        call is_curve_flat(n, col, tol, is_flat)
        if (is_flat .eqv. .false.) then
            return
        end if
    end do

    ! Rows
    do i = 0, n
        row = cp(i, :, :)
        call is_curve_flat(m, row, tol, is_flat)
        if (is_flat .eqv. .false.) then
            return
        end if
    end do

end subroutine is_surface_flat

subroutine equivalence_array1d(d0, d1, array, tol, npts, eq_array)
    !> Equivalence points of an array.
    !> d0 - Dimension 1 of array (number of points).
    !> d1 - Dimension 2 of array (3 for dehomogenized array,
    !> 4 for homogeneous array).
    !> array - Input array.
    !> tol - Tolerance.
    !> npts - Number of unique points.
    !> eq_array - Array of unique points.

    !f2py intent(in) d0, d1, array, tol
    !f2py intent(out) npts, eq_array
    !f2py depend(d0, d1) array, eq_array

    ! Input
    integer, intent(in) :: d0, d1
    double precision, intent(in) ::  tol
    double precision, intent(in) :: array(d0, d1)

    ! Output
    integer, intent(out) :: npts
    double precision, intent(out) :: eq_array(d0, d1)

    ! Working
    logical :: unique
    integer :: i, j, k
    double precision :: dist, v(d1)

    eq_array(:, :) = 0.0d0
    eq_array(1, :) = array(1, :)
    npts = 1
    do i = 2, d0
        unique = .true.
        do j = 1, npts
            v = eq_array(j, :) - array(i, :)
            dist = 0.0d0
            do k = 1, d1
                dist = dist + v(k) * v(k)
            end do
            dist = dsqrt(dist)
            if (dist .le. tol) then
                unique = .false.
                exit
            end if
        end do
        if (unique .eqv. .true.) then
            npts = npts + 1
            eq_array(npts, :) = array(i, :)
        end if
    end do
end subroutine equivalence_array1d

subroutine barycentric_params(p, t, u, v)
    !> Determine the barycentric coordinates of a point with respect
    !> to the triangle defined by three vertices.
    !> p - Point.
    !> t - Vertices of triangle.
    !> u - Barycentric coordinate.
    !> v - Barycentric coordinate.

    ! Input
    double precision, intent(in) :: p(0:2), t(0:2, 0:2)

    ! Output
    double precision, intent(out) :: u, v

    ! Working
    double precision :: d01, d12, d02, dp1, dp2, denom
    double precision :: v01(0:2), v02(0:2), vp(0:2)

    v01 = t(1, :) - t(0, :)
    v02 = t(2, :) - t(0, :)
    vp = p - t(0, :)
    d01 = dot_product(v01, v01)
    d12 = dot_product(v01, v02)
    d02 = dot_product(v02, v02)
    dp1 = dot_product(vp, v01)
    dp2 = dot_product(vp, v02)
    denom = d01 * d02 - d12 * d12
    u = (d02 * dp1 - d12 * dp2) / denom
    v = (d01 * dp2 - d12 * dp1) / denom

    if (u .lt. 0.0d0) then
        u = 0.0d0
    elseif (u .gt. 1.0d0) then
        u = 1.0d0
    end if

    if (v .lt. 0.0d0) then
        v = 0.0d0
    elseif (v .gt. 1.0d0) then
        v = 1.0d0
    end if

end subroutine barycentric_params

function are_points_equal(p1, p2, tol) result(are_equal)
    !> Check to see if two points are coincident within a tolerance.
    !> p1 - A point.
    !> p2 - Other point.
    !> tol - Tolerance.
    !> are_equal - True if points are equal, False if not.

    ! Input
    double precision, intent(in) :: tol
    double precision, intent(in) :: p1(:), p2(:)

    ! Output
    logical :: are_equal

    ! Working
    double precision :: d

    are_equal = .false.
    d = norm(p1 - p2)
    if (d .le. tol) then
        are_equal = .true.
    end if

end function are_points_equal

subroutine is_surface_closed(n, m, cpw, uclosed, vclosed)
    !> Check to see if a surface is closed.
    !> n - Number of control points - 1 in u-direction.
    !> m - Number of control points - 1 in v-direction.
    !> cpw - Control points.
    !> uclosed - Is surface closed in u-direction.
    !> vclosed - Is surface closed in v-direction.

    ! Input
    integer, intent(in) :: n, m
    double precision, intent(in) :: cpw(0:n, 0:m, 0:3)

    ! Output
    logical, intent(out) :: uclosed, vclosed

    ! Working
    logical :: are_equal
    integer :: i
    double precision :: p1(0:3), p2(0:3)

    uclosed = .true.
    vclosed = .true.

    ! Check for u-direction by comparing control points at first and last rows.
    do i = 0, m
        p1 = cpw(0, i, :)
        p2 = cpw(n, i, :)
        are_equal = are_points_equal(p1, p2, gtol)
        if (are_equal .eqv. .false.) then
            uclosed = .false.
            exit
        end if
    end do

    ! Check for v-direction by comparing control points at first and last columns.
    do i = 0, n
        p1 = cpw(i, 0, :)
        p2 = cpw(i, m, :)
        are_equal = are_points_equal(p1, p2, gtol)
        if (are_equal .eqv. .false.) then
            vclosed = .false.
            exit
        end if
    end do

end subroutine is_surface_closed

subroutine invert_point_in_triangle(p, t3d, t2d, u, v)
    !> Find parameters for a point inside a triangle.
    !> p - Point.
    !> t3d - Vertices of triangle in 3-D.
    !> t2d - Vertices of triangle in 2-D.
    !> u - Point parameter in u-direction.
    !> v - Point parameter in v-direction.

    ! Input
    double precision, intent(in) :: p(0:2), t3d(0:2, 0:2), t2d(0:2, 0:1)

    ! Output
    double precision, intent(out) :: u, v

    ! Working
    double precision :: d01, d12, d02, dp1, dp2, denom, ub, vb
    double precision :: v01(0:2), v02(0:2), vp(0:2)

    ! Find barycentric parameters.
    v01 = t3d(1, :) - t3d(0, :)
    v02 = t3d(2, :) - t3d(0, :)
    vp = p - t3d(0, :)
    d01 = dot_product(v01, v01)
    d12 = dot_product(v01, v02)
    d02 = dot_product(v02, v02)
    dp1 = dot_product(vp, v01)
    dp2 = dot_product(vp, v02)
    denom = d01 * d02 - d12 * d12
    ub = (d02 * dp1 - d12 * dp2) / denom
    vb = (d01 * dp2 - d12 * dp1) / denom

    ! Find 2-D parameters.
    u = (1.0d0 - ub - vb) * t2d(0, 0) + ub * t2d(1, 0) + vb * t2d(2, 0)
    v = (1.0d0 - ub - vb) * t2d(0, 1) + ub * t2d(1, 1) + vb * t2d(2, 1)

end subroutine invert_point_in_triangle

function angle_between_vecs(v1, v2) result(angle)
    !> Calculate the angle between two vectors.
    !> v1 - Vector 1.
    !> v2 - Vector 2.

    ! Input
    double precision, intent(in) :: v1(:), v2(:)

    ! Output
    double precision :: angle

    ! Working
    double precision :: n1, n2, x, arad, pi

    n1 = norm(v1)
    n2 = norm(v2)
    x = dot_product(v1, v2) / (n1 * n2)
    if (x .gt. 1.0d0) then
        x = 1.0d0
    elseif (x .lt. -1.0d0) then
        x = -1.0d0
    end if
    arad = dacos(x)
    pi = 4.0d0 * datan(1.0d0)
    angle = arad * 180.0d0 / pi

end function angle_between_vecs

end module geom_utils


!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++!
!                          Futility Development Group                          !
!                             All rights reserved.                             !
!                                                                              !
! Futility is a jointly-maintained, open-source project between the University !
! of Michigan and Oak Ridge National Laboratory.  The copyright and license    !
! can be found in LICENSE.txt in the head directory of this repository.        !
!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++!
!> @brief A Fortran 2003 module implementing some basic search algorithms.
!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++!
MODULE Search
#include "Futility_DBC.h"
USE IntrType
USE Strings
USE ArrayUtils
USE Futility_DBC

IMPLICIT NONE
PRIVATE

PUBLIC :: upperBound
PUBLIC :: lowerBound
PUBLIC :: getFirstGreater
PUBLIC :: getFirstGreaterEqual
PUBLIC :: find
PUBLIC :: binarySearch

!TODO: Change names of these functions on MPACT side to the new versions below
!      and eliminate these old confusing names:
!> @brief Interface for getting "upper bounds"
!>
INTERFACE upperBound
  !> @copybrief Search::getFirstGreater_i
  !> @copydetails Search::getFirstGreater_i
  MODULE PROCEDURE getFirstGreater_i
  !> @copybrief Search::getFirstGreater_r
  !> @copydetails Search::getFirstGreater_r
  MODULE PROCEDURE getFirstGreater_r
ENDINTERFACE upperBound

!> @brief Interface for getting "lower bounds"
!>
INTERFACE lowerBound
  !> @copybrief Search::getFirstGreaterEqual_i
  !> @copydetails Search::getFirstGreaterEqual_i
  MODULE PROCEDURE getFirstGreaterEqual_i
  !> @copybrief Search::getFirstGreaterEqual_r
  !> @copydetails Search::getFirstGreaterEqual_r
  MODULE PROCEDURE getFirstGreaterEqual_r
ENDINTERFACE lowerBound

!> @brief Interface for getting the first entry of an array greater than the input
!>
INTERFACE getFirstGreater
  !> @copybrief Search::getFirstGreater_i
  !> @copydetails Search::getFirstGreater_i
  MODULE PROCEDURE getFirstGreater_i
  !> @copybrief Search::getFirstGreater_r
  !> @copydetails Search::getFirstGreater_r
  MODULE PROCEDURE getFirstGreater_r
ENDINTERFACE getFirstGreater

!> @brief Interface for getting the first entry of an array greater than or equal to the input
!>
INTERFACE getFirstGreaterEqual
  !> @copybrief Search::getFirstGreaterEqual_i
  !> @copydetails Search::getFirstGreaterEqual_i
  MODULE PROCEDURE getFirstGreaterEqual_i
  !> @copybrief Search::getFirstGreaterEqual_r
  !> @copydetails Search::getFirstGreaterEqual_r
  MODULE PROCEDURE getFirstGreaterEqual_r
ENDINTERFACE getFirstGreaterEqual

!> @brief Interface for finding entry of a 1D array using a linear search
!>
INTERFACE find
  !> @copybrief Search::linearSearch_i
  !> @copydetails Search::linearSearch_i
  MODULE PROCEDURE linearSearch_i
  !> @copybrief Search::linearSearch_r
  !> @copydetails Search::linearSearch_r
  MODULE PROCEDURE linearSearch_r
  !> @copybrief Search::linearSearch_str
  !> @copydetails Search::linearSearch_str
  MODULE PROCEDURE linearSearch_str
  !> @copybrief Search::linearSearch_strchar
  !> @copydetails Search::linearSearch_strchar
  MODULE PROCEDURE linearSearch_strchar
ENDINTERFACE find

!> @brief Interface for finding entry of a 1D array using a binary search
!>
INTERFACE binarySearch
  !> @copybrief Search::binarySearch_i
  !> @copydetails Search::binarySearch_i
  MODULE PROCEDURE binarySearch_i
  !> @copybrief Search::binarySearch_r
  !> @copydetails Search::binarySearch_r
  MODULE PROCEDURE binarySearch_r
ENDINTERFACE binarySearch
!
!===============================================================================
CONTAINS
!
!-------------------------------------------------------------------------------
!> @brief Returns the index of the first element of a list of integers that is
!>        greater than val
!> @param list 1D array of integers to search through. The list should be sorted in
!>        ascending order
!> @param val integer value to find the first entry of list that is greater than
!> @returns i the index of the first element of list that is greater than val
!>
!> NOTE: If val is greater than or equal to the largest element of list this
!>       returns SIZE(list)+1, and if val is smaller than the smallest element of
!>       list this returns 1.
!>
FUNCTION getFirstGreater_i(list,val) RESULT (i)
  INTEGER(SIK),INTENT(IN) :: list(:)
  INTEGER(SIK),INTENT(IN) :: val
  INTEGER(SIK) :: j
  !
  INTEGER(SIK) :: n,step,i

  REQUIRE(isIncreasing(list) .OR. SIZE(list) == 1)

  !Perform binary search
  j=1
  n=SIZE(list)
  DO WHILE(n>0)
    step=n/2
    i=j+step
    IF(.NOT. (val < list(i))) THEN
      i=i+1
      j=i
      n=n-step-1
    ELSE
      n=step
    ENDIF
  ENDDO

ENDFUNCTION getFirstGreater_i
!
!-------------------------------------------------------------------------------
!> @brief Returns the index of the first element of a list of reals that is
!>        greater than val
!> @param list 1D array of reals to search through. The list should be sorted in
!>        ascending order
!> @param val real value to find the first entry of list that is greater than
!> @param tol the comparison tolerance to use. Optional
!> @returns i the index of the first element of list that is greater than val
!>
!> NOTE: If val is greater than or equal to the largest element of list this
!>       returns SIZE(list)+1. If val is smaller than the smallest element of list
!>       this returns 1.
!>
FUNCTION getFirstGreater_r(list,val,tol) RESULT (i)
  REAL(SRK),INTENT(IN) :: list(:)
  REAL(SRK),INTENT(IN) :: val
  REAL(SRK),INTENT(IN),OPTIONAL :: tol
  INTEGER(SIK) :: j
  !
  INTEGER(SIK) :: n,step,i
  REAL(SRK) :: eps

  REQUIRE(isIncreasing(list) .OR. SIZE(list) == 1)

  !Set Soft comparison tolerance
  IF(PRESENT(tol)) THEN
    eps=tol
  ELSE
    eps=EPSREAL
  ENDIF

  !Perform binary search
  j=1
  n=SIZE(list)
  DO WHILE(n>0)
    step=n/2
    i=j+step
    IF(.NOT. SOFTLT(val,list(i),eps)) THEN
      i=i+1
      j=i
      n=n-step-1
    ELSE
      n=step
    ENDIF
  ENDDO

ENDFUNCTION getFirstGreater_r
!
!-------------------------------------------------------------------------------
!> @brief Return the index of the first element of a list that does not
!>        compare less than the passed value
!> @param list 1D array of integers to search through. The list should be sorted in
!>        ascending order
!> @param val integer value to find
!> @returns j the index of the element of list that is immediatly less than val
!>
!> NOTE: If val is greater than the largest element of list this returns SIZE(list)+1.
!>       If val is smaller than the smallest element of list this returns 1.
!>
FUNCTION getFirstGreaterEqual_i(list,val) RESULT(j)
  INTEGER(SIK),INTENT(IN) :: list(:)
  INTEGER(SIK),INTENT(IN) :: val
  INTEGER(SIK) :: j
  !
  INTEGER(SIK) :: n,step,i
  n=SIZE(list)
  j=1

  REQUIRE(isMonotonic(list) .OR. SIZE(list) == 1)

  !Check length and return if insufficient
  IF(n < 2) THEN
    IF(list(1) < val) j=2
    RETURN
  ENDIF

  !Perform binary search
  IF(list(1) < list(n)) THEN
    DO WHILE(n>0)
      step=n/2
      i=j+step
      IF(list(i) < val) THEN
        i=i+1
        j=i
        n=n-step-1
      ELSE
        n=step
      ENDIF
    ENDDO
  ELSE
    DO WHILE(n>0)
      step=n/2
      i=j+step
      IF(list(i) > val) THEN
        i=i+1
        j=i
        n=n-step-1
      ELSE
        n=step
      ENDIF
    ENDDO
  ENDIF

ENDFUNCTION getFirstGreaterEqual_i
!
!-------------------------------------------------------------------------------
!> @brief Return the index of the first element of a list that does not
!>        compare less than the passed value
!> @param list 1D array of reals to search through. The list should be sorted in
!>        ascending order
!> @param val real value to find
!> @param tol the comparison tolerance to use. Optional
!> @returns j the index of the element of list that is immediatly less than val
!>
!> NOTE: If val is greater than the largest element of list this returns SIZE(list)+1.
!>       If val is smaller than the smallest element of list this returns 1.
!>
FUNCTION getFirstGreaterEqual_r(list,val,tol) RESULT(j)
  REAL(SRK),INTENT(IN) :: list(:)
  REAL(SRK),INTENT(IN) :: val
  REAL(SRK),INTENT(IN),OPTIONAL :: tol
  INTEGER(SIK) :: j
  !
  PROCEDURE(SOFTLTwrap),POINTER :: condComp
  INTEGER(SIK) :: n,step,i
  REAL(SRK) :: eps
  n=SIZE(list)
  j=1

  REQUIRE(isMonotonic(list) .OR. SIZE(list) == 1)

  !Set Soft comparison tolerance
  IF(PRESENT(tol)) THEN
    eps=tol
  ELSE
    eps=EPSREAL
  ENDIF

  !Check length and return if insufficient
  IF(n < 2) THEN
    IF(SOFTLT(list(1),val,eps)) j=2
    RETURN
  ENDIF

  IF(list(1) < list(n)) THEN
    condComp => SOFTLTwrap
  ELSE
    condComp => SOFTGTwrap
  ENDIF

  !Perform binary search
  DO WHILE(n>0)
    step=n/2
    i=j+step
    IF(condComp(list(i),val,eps)) THEN
      i=i+1
      j=i
      n=n-step-1
    ELSE
      n=step
    ENDIF
  ENDDO

ENDFUNCTION getFirstGreaterEqual_r
!
!-------------------------------------------------------------------------------
!> @brief Return the index of the element of a list of integers that is equal to
!>        val using a linear search algorithm
!> @param list 1D array of integers to search through. Unlike for binary_search,
!>        sorting does not matter
!> @param val integer value to find in list
!> @returns i index of the element of list that equals val
!>
!> NOTE: If the value is not found this returns SIZE(list)+1
!>
FUNCTION linearSearch_i(list, val) RESULT(i)
  INTEGER(SIK),INTENT(IN) :: list(:)
  INTEGER(SIK),INTENT(IN) :: val
  INTEGER(SIK) :: i

  !Perform linear search
  DO i=1,SIZE(list)
    IF(list(i) == val) RETURN
  ENDDO

  i=SIZE(list)+1
ENDFUNCTION linearSearch_i
!
!-------------------------------------------------------------------------------
!> @brief Return the index of the element of a list of reals that is equal to val
!>        within tol using a linear search algorithm
!> @param list 1D array of reals to search through. Unlike for binary_search,
!>        sorting does not matter
!> @param val real value to find in list
!> @param tol the comparison tolerance to use. Optional
!> @returns i index of the element of list that equals val
!>
!> NOTE: If the value is not found this returns SIZE(list)+1
!>
FUNCTION linearSearch_r(list,val,tol) RESULT(i)
  REAL(SRK),INTENT(IN) :: list(:)
  REAL(SRK),INTENT(IN) :: val
  REAL(SRK),INTENT(IN),OPTIONAL :: tol
  INTEGER(SIK) :: i
  !
  REAL(SRK) :: eps

  !Set Soft comparison tolerance
  IF(PRESENT(tol)) THEN
    eps=tol
  ELSE
    eps=EPSREAL
  ENDIF

  !Perform linear search
  DO i=1,SIZE(list)
    IF(SOFTEQ(list(i), val, eps)) RETURN
  ENDDO

  i=SIZE(list)+1
ENDFUNCTION linearSearch_r
!
!-------------------------------------------------------------------------------
!> @brief Return the index of the element of a list of strings that is equal to
!>        val using a linear search algorithm
!> @param list 1D array of strings to search through. Unlike for binary_search,
!>        sorting does not matter
!> @param val string value to find in list
!> @returns i index of the element of list that equals val
!>
!> NOTE: If the value is not found this returns SIZE(list)+1
!>
FUNCTION linearSearch_str(list, val) RESULT(i)
  TYPE(StringType),INTENT(IN) :: list(:)
  TYPE(StringType),INTENT(IN) :: val
  INTEGER(SIK) :: i

  !Perform linear search
  DO i=1,SIZE(list)
    IF(list(i) == val) RETURN
  ENDDO

  i=SIZE(list)+1
ENDFUNCTION linearSearch_str
!
!-------------------------------------------------------------------------------
!> @brief Return the index of the element of a list of strings that is equal to
!>        val using a linear search algorithm
!> @param list 1D array of strings to search through. Unlike for binary_search,
!>        sorting does not matter
!> @param val character value to find in list
!> @returns i index of the element of list that equals val
!>
!> NOTE: If the value is not found this returns SIZE(list)+1
!>
FUNCTION linearSearch_strchar(list, val) RESULT(i)
  TYPE(StringType),INTENT(IN) :: list(:)
  CHARACTER(LEN=*),INTENT(IN) :: val
  INTEGER(SIK) :: i

  !Perform linear search
  DO i=1,SIZE(list)
    IF(list(i) == val) RETURN
  ENDDO

  i=SIZE(list)+1
ENDFUNCTION linearSearch_strchar
!
!-------------------------------------------------------------------------------
!> @brief Return the index of the element of a list of integers that is equal to
!>        val using a binary search algorithm
!> @param list 1D array of integers to search through. The list should be sorted in
!>        ascending order
!> @param val integer value to find in list
!> @returns i index of the element of list that equals val
!>
!> NOTE: If the value is not found this returns SIZE(list)+1
!>
FUNCTION binarySearch_i(list,val) RESULT(i)
  INTEGER(SIK),INTENT(IN) :: list(:)
  INTEGER(SIK),INTENT(IN) :: val
  INTEGER(SIK) :: i

  REQUIRE(isIncreasing(list) .OR. SIZE(list) == 1)

  !Find index of list entry at or just before val
  i=getFirstGreaterEqual(list,val)
  IF(i <= SIZE(list)) THEN
    IF(list(i) == val) THEN
      RETURN
    ENDIF
  ENDIF

  !If above equality check failed increment index by one
  i=SIZE(list)+1
ENDFUNCTION binarySearch_i
!
!-------------------------------------------------------------------------------
!> @brief Return the index of the element of a list of reals that is equal to val
!>        within tol using a binary search algorithm
!> @param list 1D array of reals to search through. The list should be sorted in
!>        ascending order
!> @param val real value to find in list
!> @param tol the comparison tolerance to use. Optional
!> @returns i index of the element of list that equals val
!>
!> NOTE: If the value is not found this returns SIZE(list)+1
!>
FUNCTION binarySearch_r(list,val,tol) RESULT(i)
  REAL(SRK),INTENT(IN) :: list(:)
  REAL(SRK),INTENT(IN) :: val
  REAL(SRK),INTENT(IN),OPTIONAL :: tol
  INTEGER(SIK) :: i
  !
  REAL(SRK) :: eps

  REQUIRE(isIncreasing(list) .OR. SIZE(list) == 1)

  !Set Soft comparison tolerance
  IF(PRESENT(tol)) THEN
    eps=tol
  ELSE
    eps=EPSREAL
  ENDIF

  !Find index of list entry at or just before val
  i=getFirstGreaterEqual(list,val,tol)
  IF(i <= SIZE(list)) THEN
    IF(SOFTEQ(list(i),val,eps)) THEN
      RETURN
    ENDIF
  ENDIF

  !If above equality check failed increment index by one
  i=SIZE(list)+1
ENDFUNCTION binarySearch_r
!
!-------------------------------------------------------------------------------
!> @brief Wrapper functions for the SOFT GT and LT functions to give them
!>        explicitly SRK type arguments and allow for precedure pointer usage
!> @param r1 first real to compare
!> @param r2 second real to comapre
!> @param tol the comparison tolerance to use.
!> @returns bool result of comparison
!>
FUNCTION SOFTLTwrap(r1,r2,tol) RESULT(bool)
  REAL(SRK),INTENT(IN) :: r1
  REAL(SRK),INTENT(IN) :: r2
  REAL(SRK),INTENT(IN) :: tol
  LOGICAL(SBK) :: bool

  bool=SOFTLT(r1,r2,tol)

ENDFUNCTION SOFTLTwrap
!
FUNCTION SOFTGTwrap(r1,r2,tol) RESULT(bool)
  REAL(SRK),INTENT(IN) :: r1
  REAL(SRK),INTENT(IN) :: r2
  REAL(SRK),INTENT(IN) :: tol
  LOGICAL(SBK) :: bool

  bool=SOFTGT(r1,r2,tol)

ENDFUNCTION SOFTGTwrap
!
ENDMODULE Search


module isolve
use types, only: dp
use utils, only: stop_error
use sparse, only: csr_matvec, csr_getvalue
use openmp, only: omp_get_wtime, omp_get_max_threads, with_openmp
implicit none
private
public solve_cg

contains

function solve_cg(Ap, Aj, Ax, b, x0, tol, maxiter, verbose) result(x)
! solves a system of equations A x = b with one right hand side
integer, intent(in) :: Ap(:), Aj(:)  ! coefficient matrix A in CSR format
real(dp), intent(in) :: Ax(:)        ! coefficient matrix A in CSR format
real(dp), intent(in) :: b(:)  ! right-hand-side A x = b
real(dp), intent(in) :: x0(:) ! Initial guess
real(dp), intent(in) :: tol ! Tolerance in residual
integer, intent(in) :: maxiter ! Maximum number of iterations
logical, intent(in), optional :: verbose
real(dp) :: x(size(b)) ! solution
real(dp), dimension(size(b)) :: r, p, Ap_, z
real(dp) :: r2, r2old, alpha, res_norm
real(dp) :: A_diag(size(x0))
real(dp) :: t1, t2, t1_omp, t2_omp
integer :: i
logical :: verbose_
verbose_ = .false.
if (present(verbose)) verbose_ = verbose
call precond_init()

x = x0
r = b - csr_matvec(Ap, Aj, Ax, x)
z = precond(r)
p = z
r2old = dot_product(r, z)
if (verbose_) then
    print "('Conjugate Gradient solver (using ', i0, ' threads)')", &
        omp_get_max_threads()
    print *, "Iter    Residual ||A x - b||"
end if
call cpu_time(t1)
t1_omp = omp_get_wtime()
do i = 1, maxiter
    Ap_ = csr_matvec(Ap, Aj, Ax, p)
    alpha = r2old / dot_product(p, Ap_)
    x = x + alpha * p
    r = r - alpha * Ap_
    res_norm = sqrt(dot_product(r, r))  ! Good approximation to ||A x - b||

    if (verbose_) then
        print "(i4, '      ', es10.2)", i, res_norm
    end if
    if (res_norm < tol) then
        t2_omp = omp_get_wtime()
        call cpu_time(t2)
        if (verbose_) then
            print *, "solve_cg() statistics:"
            if (with_openmp()) then
                print "('     Time (OpenMP walltime) [s]:     ', f11.6)", &
                    t2_omp-t1_omp
            end if
            print "('     Total CPU time on all cores [s]:', f11.6)", t2-t1
            if (with_openmp()) then
                print "('     Speedup:  ', f7.2, 'x')", (t2-t1) / (t2_omp-t1_omp)
            end if
            print "('     # threads:', i4)", omp_get_max_threads()
            print *, "    # iterations:", i
            print *, "    Matrix size:", size(x0)
            print *, "    # non-zeros:", size(Ax)
            r = csr_matvec(Ap, Aj, Ax, x) - b
            write(*, '(1x,a,es10.2)') "Solution vector residual ||A x - b||/||bv||: ",  sqrt(dot_product(r, r) / dot_product(b, b))
        end if
        return
    end if

    z = precond(r)
    r2 = dot_product(r, z)
    p = z + r2 / r2old * p
    r2old = r2
end do
call stop_error("Solution did not converge.")

contains

    subroutine precond_init()
    ! Precalculate the diagonal of the A matrix
    integer :: i
    do i = 1, size(x)
        A_diag(i) = csr_getvalue(Ap, Aj, Ax, i, i)
    end do
    end subroutine

    function precond(x) result(y)
    ! Calculates y = M^-1 x
    real(dp), intent(in) :: x(:)
    real(dp) :: y(size(x))
    ! No preconditioning:
    !y = x
    ! Jacobi preconditioning: M = diag(A):
    y = x / A_diag
    end function

end function

end module


!> Testing rank one updates.
program test_psyr_pher
  use, intrinsic :: iso_fortran_env, stdout => output_unit
  use test_common_module
  use libscalapackfx_module
  implicit none

  ! Block size (using an extremely small value for test purposes)
  integer, parameter :: bsize = 2

  call main()

contains

  subroutine main()
    type(blacsgrid) :: mygrid
    real(dp), allocatable :: xx(:,:), res(:,:)
    integer :: descx(DLEN_), descres(DLEN_)
    real(dp) :: alpha
    integer :: nprow, npcol, mm, nn, iproc, nproc
    integer :: ii

    ! Initialize blas and create a square processor grid
    call blacsfx_pinfo(iproc, nproc)
    do nprow = int(sqrt(real(nproc, dp))), nproc
      if (mod(nproc, nprow) == 0) then
        exit
      end if
    end do
    npcol = nproc / nprow
    call mygrid%initgrid(nprow, npcol)
    if (mygrid%master) then
      write(stdout, "(A,2(1X,I0))") "# processor grid:", nprow, npcol
      write(stdout, "(A,1X,I0)") "# block size:", bsize
    end if

    ! Read in matrix from disc.
    call readfromfile(mygrid, "hamsqr1.dat", bsize, bsize, xx, descx)
    mm = descx(M_)
    nn = descx(N_)
    if (mygrid%master) then
      write(stdout, "(A,2(1X,I0))") "# global matrix size:", mm, nn
      write(stdout, "(A,2(1X,I0))") "# local matrix size on master:",&
          & size(xx, dim=1), size(xx, dim=2)
    end if

    ! Do rank one update with all column vectors of the matrix.
    call scalafx_creatematrix(mygrid, mm, mm, bsize, bsize, res, descres)
    res(:,:) = 0.0_dp
    do ii = 1, mm
      alpha = real(ii, dp)
      call pblasfx_psyr(xx, descx, res, descres, alpha=alpha, ix=1, jx=ii,&
          & incx=1)
    end do

    ! Write results to disc.
    call writetofile(mygrid, "psyr_result.dat", res, descres)
    if (mygrid%master) then
      write(stdout, "(A)") "Result written to file 'psyr_result.dat'."
    end if

    ! Finish blacs.
    call blacsfx_exit()

  end subroutine main

end program test_psyr_pher


! chk_quadcomplex.f90 --
!     Check: does the compiler support complex numbers with larger precision than double precision?
!
!     Note: if not, then the compilation will fail (kind < 0)
!
program chk_quadcomplex
    implicit none

    integer, parameter :: quadruple = selected_real_kind(precision(1.0d0)+1)
    complex(quadruple) :: quadruple_value = (1.0_quadruple,1.0_quadruple)

    write(*,'(a,i0)') 'Precision for "quadruple" complex numbers: ', precision(quadruple_value)
    write(*,'(a,i0)') 'Range for "quadruple" complex numbers:     ', range(quadruple_value)
end program chk_quadcomplex


module ml_norm_module

  use multifab_module

  implicit none

  public :: ml_norm_inf, ml_norm_l2

contains

  function ml_norm_inf(mf, mask, local) result(r)
    type( multifab), intent(in)   :: mf(:)
    type(lmultifab), intent(in)   :: mask(:)
    logical, intent(in), optional :: local
    real(dp_t)                    :: r, r1
    integer                       :: n,nlevs
    nlevs = size(mf)
    r = norm_inf(mf(nlevs),local=local)
    do n = nlevs-1, 1, -1
       r1 = norm_inf(mf(n), mask(n), local=local)
       r = max(r1, r)
    end do
  end function ml_norm_inf

  function ml_norm_l2(mf, rr, mask) result(r)
    type( multifab), intent(in) :: mf(:)
    integer                     :: rr(:,:)
    type(lmultifab), intent(in) :: mask(:)
    real(dp_t)                  :: r
    integer                     :: n,nlevs
    nlevs = size(mf)
    r = norm_l2(mf(nlevs))**2
    do n = nlevs-1, 1, -1
       r =  r / product(rr(n,:)) &
          + norm_l2(mf(n), mask = mask(n))**2
    end do
    r = sqrt(r)
  end function ml_norm_l2

end module ml_norm_module