module SLB4MPI_work_stealing_load_balancer_m
  use mpi
  implicit none
  private

  integer, parameter :: BITS_IN_BYTE = 8

  type :: work_stealing_load_balancer_t
    integer(MPI_INTEGER_KIND) :: communicator !< MPI communicator
    integer(MPI_INTEGER_KIND) :: rank         !< Rank of the process
    integer(MPI_INTEGER_KIND) :: root = 0     !< Rank of the root process
    integer(MPI_INTEGER_KIND) :: nranks       !< Number of processes (group size)
    integer(8) :: lower_bound                 !< Lower bound of range
    integer(8) :: upper_bound                 !< Upper bound of range
    integer(8) :: min_chunk_size              !< Minimal chunk size for job
    integer(8) :: max_chunk_size              !< Maximal chunk size for job
    integer(MPI_INTEGER_KIND) :: window_num_active  !< number of active threads
    integer(MPI_INTEGER_KIND) :: window_bounds      !< lower and upper bounds of rank
    integer(MPI_INTEGER_KIND) :: window_actual_rank !< actual rank to compute (for fast look up)
    integer(MPI_INTEGER_KIND) :: window_done        !< status of thread
    integer(MPI_INTEGER_KIND) :: actual_rank
    logical :: done = .false.
  contains
    procedure :: initialize
    procedure :: get_range
    procedure :: clean
  end type

  interface

  end interface

  public :: work_stealing_load_balancer_t

contains
  !>
  !> @brief initialize work_stealing load balancer
  !>
  !> @param[in]           communicator   - MPI communicator on which load balancer will be used
  !> @param[in]           lower_bound    - lower bound of range
  !> @param[in]           upper_bound    - upper bound of range, upper_bound >= lower_bound
  !> @param[in, optional] min_chunk_size - minimal size of chank that can be associated with job, default: 1
  !> @param[in, optional] max_chunk_size - maximal size of chank that can be associated with job, default: upper_bound - lower_bound + 1
  !>
  subroutine initialize(lb, communicator, lower_bound, upper_bound, min_chunk_size, max_chunk_size)
    use, intrinsic :: iso_c_binding, only: c_ptr, c_f_pointer
    class(work_stealing_load_balancer_t), intent(inout) :: lb
    integer(MPI_INTEGER_KIND),      intent(in)    :: communicator
    integer(8),                     intent(in)    :: lower_bound, upper_bound
    integer(8),           optional, intent(in)    :: min_chunk_size, max_chunk_size

    integer(MPI_INTEGER_KIND) :: ierr
    integer(8) :: n_tasks, extra_tasks
    ! num active
    integer(MPI_ADDRESS_KIND) :: size_num_active
    integer(MPI_INTEGER_KIND) :: disp_unit_num_active
    type(c_ptr) :: baseaddr_num_active
    integer(MPI_INTEGER_KIND), pointer :: num_active
    ! actual rank
    integer(MPI_ADDRESS_KIND) :: size_actual_rank
    integer(MPI_INTEGER_KIND) :: disp_unit_actual_rank
    ! bounds
    integer(MPI_ADDRESS_KIND) :: size_bounds
    integer(MPI_INTEGER_KIND) :: disp_unit_bounds
    type(c_ptr) :: baseaddr_bounds
    integer(8), pointer :: bounds(:)
    integer(8), parameter :: bounds_2 = 0
    ! done
    integer(MPI_ADDRESS_KIND) :: size_done
    integer(MPI_INTEGER_KIND) :: disp_unit_done
    logical, parameter :: done = .false.

    size_num_active = storage_size(lb%nranks) / BITS_IN_BYTE
    disp_unit_num_active = size_num_active
    size_actual_rank = storage_size(lb%rank) / BITS_IN_BYTE
    disp_unit_actual_rank = size_actual_rank
    size_bounds = storage_size(bounds_2) / BITS_IN_BYTE * 2
    disp_unit_bounds = size_bounds / 2
    size_done = storage_size(done) / BITS_IN_BYTE
    disp_unit_done = size_done


    call MPI_Comm_rank(communicator, lb%rank, ierr)
    call MPI_Comm_size(communicator, lb%nranks, ierr)

    lb%communicator = communicator
    lb%lower_bound = lower_bound
    lb%upper_bound = upper_bound
    lb%min_chunk_size = 1
    if (present(min_chunk_size)) lb%min_chunk_size = min_chunk_size
    lb%max_chunk_size = max((lb%upper_bound - lb%lower_bound + 1) / lb%nranks, lb%min_chunk_size)
    if (present(max_chunk_size)) lb%max_chunk_size = max_chunk_size

    if (lb%min_chunk_size <= 0) error stop "Min chunk size must be greater than zero!"

    if (lb%lower_bound > lb%upper_bound) error stop "Upper bound is less than lower bound!"
    if (lb%min_chunk_size > lb%max_chunk_size) error stop "Max chunk size is less than min chunk size!"


    n_tasks = (upper_bound - lower_bound + 1) / lb%nranks
    extra_tasks = mod(upper_bound - lower_bound + 1, lb%nranks)
    if (n_tasks < min_chunk_size) then
        n_tasks = min_chunk_size
        extra_tasks = 0
    end if

    if (lb%rank < extra_tasks) then
        lb%lower_bound = lb%rank * (n_tasks + 1) + lower_bound
        lb%upper_bound = (lb%rank + 1) * (n_tasks + 1) + lower_bound - 1
    else
        lb%lower_bound = lb%rank * n_tasks + extra_tasks + lower_bound
        lb%upper_bound = (lb%rank + 1) * n_tasks + extra_tasks + lower_bound - 1
    end if

    lb%lower_bound = min(lb%lower_bound, upper_bound + 1)
    lb%upper_bound = min(lb%upper_bound, upper_bound)

    lb%actual_rank = lb%rank

    call MPI_Win_allocate(size_num_active, disp_unit_num_active, MPI_INFO_NULL, lb%communicator, baseaddr_num_active, lb%window_num_active, ierr)
    call MPI_Win_create(lb%actual_rank, size_actual_rank, disp_unit_actual_rank, MPI_INFO_NULL, lb%communicator, lb%window_actual_rank, ierr)
    call MPI_Win_allocate(size_bounds, disp_unit_bounds, MPI_INFO_NULL, lb%communicator, baseaddr_bounds, lb%window_bounds, ierr)
    call MPI_Win_create(lb%done, size_done, disp_unit_done, MPI_INFO_NULL, lb%communicator, lb%window_done, ierr)

    if (lb%rank == lb%root) then
      call c_f_pointer(baseaddr_num_active, num_active)
      num_active = lb%nranks
    end if
    call c_f_pointer(baseaddr_bounds, bounds, [2])
    bounds = [ lb%lower_bound, lb%upper_bound ]

    call MPI_Barrier(lb%communicator, ierr)

  end subroutine initialize

  !>
  !> @brief get range to compute
  !>
  !> @note upper_bound can be only less than or equal load_balancer_t%upper_bound
  !>       so, for last elements a batch with size less than min_chunk_size can be returned
  !>
  !> @param[out] lower_bound - lower bound of range to compute
  !> @param[out] upper_bound - upper bound of range to compute
  !> @return                 - is there something to compute?
  !>
  logical function get_range(lb, lower_bound, upper_bound) result(to_compute)
    class(work_stealing_load_balancer_t), intent(inout) :: lb
    integer(8), intent(out) :: lower_bound, upper_bound

    integer(MPI_INTEGER_KIND) :: ierr
    integer(MPI_INTEGER_KIND) :: num_active
    integer(8) :: bounds(2), hop_count, nohop_count
    integer(8) :: t1, t2, cr

    call system_clock(count_rate=cr)

    to_compute = .false.

    if (lb%actual_rank == lb%rank) then
        to_compute = .true.
        !> select range for computing
        call MPI_Win_lock(MPI_LOCK_EXCLUSIVE, lb%rank, 0_MPI_INTEGER_KIND, lb%window_bounds, ierr)
        call MPI_Get(bounds, 2_MPI_INTEGER_KIND, MPI_INTEGER8, lb%rank, 0_MPI_ADDRESS_KIND, 2_MPI_INTEGER_KIND, MPI_INTEGER8, lb%window_bounds, ierr)
        call MPI_Win_flush(lb%rank, lb%window_bounds, ierr)
        lb%lower_bound = bounds(1)
        lb%upper_bound = bounds(2)
        lower_bound = lb%lower_bound
        upper_bound = min(lower_bound + lb%max_chunk_size - 1, lb%upper_bound)
        bounds = [ upper_bound + 1, lb%upper_bound ]
        print '(A,I0,A,I0,A,I0)', 'Thr ', lb%rank, ' computes range from ', lower_bound, ' to ', upper_bound
        !> update lower bound
        call MPI_Accumulate(bounds, 2_MPI_INTEGER_KIND, MPI_INTEGER8, lb%rank, 0_MPI_ADDRESS_KIND, 2_MPI_INTEGER_KIND, MPI_INTEGER8, MPI_REPLACE, lb%window_bounds, ierr)
        call MPI_Win_unlock(lb%rank, lb%window_bounds, ierr)
        !> if it the last block, give that information in advance
        if (upper_bound + 1 > lb%upper_bound .and. .not.lb%done) then
            call MPI_Win_lock(MPI_LOCK_EXCLUSIVE, lb%root, 0_MPI_INTEGER_KIND, lb%window_num_active, ierr)
            call MPI_Fetch_and_op(-1_MPI_INTEGER_KIND, num_active, MPI_INTEGER4, lb%root, 0_MPI_ADDRESS_KIND, MPI_SUM, lb%window_num_active, ierr)
            call MPI_Win_unlock(lb%root, lb%window_num_active, ierr)
            lb%actual_rank = mod(lb%actual_rank + 1, lb%nranks)
            lb%done = .true.
        end if
        if (lower_bound > upper_bound) to_compute = .false.
    end if

    !> switch to other ranks
    hop_count = 0
    nohop_count = 0
    do while (.not.to_compute)
        block
            integer(MPI_INTEGER_KIND) :: compute_rank
            logical :: done
            !> get rank which computes lb%actual_rank


            call system_clock(count=t1)
            call MPI_Win_lock(MPI_LOCK_SHARED, lb%actual_rank, MPI_MODE_NOCHECK, lb%window_actual_rank, ierr)
            call system_clock(count=t2)
            print '(A,I0,A,I0)', 'lock0, MPI: ', lb%rank, ' delay: ', t2-t1


            call MPI_Get(compute_rank, 1_MPI_INTEGER_KIND, MPI_INTEGER4, lb%actual_rank, 0_MPI_ADDRESS_KIND, 1_MPI_INTEGER_KIND, MPI_INTEGER4, lb%window_actual_rank, ierr)
            call MPI_Win_unlock(lb%actual_rank, lb%window_actual_rank, ierr)
            !> check that compute_rank computes itself
            if (lb%actual_rank == compute_rank) then
                !> check that there is something to steal


                call system_clock(count=t1)
                call MPI_Win_lock(MPI_LOCK_SHARED, lb%actual_rank, MPI_MODE_NOCHECK, lb%window_done, ierr)
                call system_clock(count=t2)
                print '(A,I0,A,I0)', 'lock1, MPI: ', lb%rank, ' delay: ', t2-t1


                call MPI_Get(done, 1_MPI_INTEGER_KIND, MPI_LOGICAL, lb%actual_rank, 0_MPI_ADDRESS_KIND, 1_MPI_INTEGER_KIND, MPI_LOGICAL, lb%window_done, ierr)
                call MPI_Win_unlock(lb%actual_rank, lb%window_done, ierr)
                if (.not.done) then
                    !> try to steal min_chunk_size jobs from compute_rank


                    call system_clock(count=t1)
                    call MPI_Win_lock(MPI_LOCK_EXCLUSIVE, compute_rank, 0_MPI_INTEGER_KIND, lb%window_bounds, ierr)
                    call system_clock(count=t2)
                    print '(A,I0,A,I0)', 'lock2, MPI: ', lb%rank, ' delay: ', t2-t1


                    call MPI_Get(bounds, 2_MPI_INTEGER_KIND, MPI_INTEGER8, compute_rank, 0_MPI_ADDRESS_KIND, 2_MPI_INTEGER_KIND, MPI_INTEGER8, lb%window_bounds, ierr)
                    call MPI_Win_flush(compute_rank, lb%window_bounds, ierr)
                    upper_bound = bounds(2)
                    lower_bound = max(bounds(1), bounds(2) - lb%min_chunk_size + 1)
                    bounds = [ bounds(1), lower_bound - 1 ]
                    if (lower_bound <= upper_bound) to_compute = .true.
                    if (to_compute) then
                        print '(A,I0,A,I0,A,I0,A,I0)', 'Thr ', lb%rank, ' steals range from ', lower_bound, ' to ', upper_bound, ' from ', compute_rank
                    end if
                    if (to_compute) then
                        !> update upper bound
                        call MPI_Accumulate(bounds, 2_MPI_INTEGER_KIND, MPI_INTEGER8, compute_rank, 0_MPI_ADDRESS_KIND, 2_MPI_INTEGER_KIND, MPI_INTEGER8, MPI_REPLACE, lb%window_bounds, ierr)
                    end if
                    call MPI_Win_unlock(compute_rank, lb%window_bounds, ierr)
                    if (to_compute) return
                end if
            else
                !> switch to lb%actual_rank of compute_rank
                print '(A,I0)', 'hop MPI: ', lb%rank
                lb%actual_rank = compute_rank
                if (lb%actual_rank == lb%rank) exit
                !> but... to avoid infinity cycle, sometimes we will check that there is some jobs
                hop_count = hop_count + 1
                block
                    real(4) :: randval
                    call random_number(randval)
                    ! log(nranks) <- how often to check cond (each 1, 2, 3, ..)
                    ! 1 / log(nranks) <- probability
                    if (randval > 1._4 / log(real(lb%nranks, kind=4)) .and. hop_count < 20) cycle
                end block
            end if
            print '(A,I0)', 'no-hop MPI: ', lb%rank
            !> if could not steal job, check how many threads finished job


            call system_clock(count=t1)
            call MPI_Win_lock(MPI_LOCK_SHARED, lb%root, 0_MPI_INTEGER_KIND, lb%window_num_active, ierr)
            call system_clock(count=t2)
            print '(A,I0,A,I0)', 'lock3, MPI: ', lb%rank, ' delay: ', t2-t1


            call MPI_Get(num_active, 1_MPI_INTEGER_KIND, MPI_INTEGER4, lb%root, 0_MPI_ADDRESS_KIND, 1_MPI_INTEGER_KIND, MPI_INTEGER4, lb%window_num_active, ierr)
            call MPI_Win_unlock(lb%root, lb%window_num_active, ierr)
            if (num_active == 0) return
            lb%actual_rank = mod(lb%actual_rank + 1, lb%nranks)
            hop_count = 0
            nohop_count = nohop_count + 1
            if (nohop_count > lb%nranks) return
        end block
    end do

  end function get_range

  !>
  !> @brief work stealing load balancer destructor
  !>
  subroutine clean(lb)
    class(work_stealing_load_balancer_t), intent(inout) :: lb
    integer(MPI_INTEGER_KIND) :: ierr

    if (lb%window_num_active /= MPI_WIN_NULL) then
      call MPI_Win_free(lb%window_num_active, ierr)
    end if
    if (lb%window_actual_rank /= MPI_WIN_NULL) then
      call MPI_Win_free(lb%window_actual_rank, ierr)
    end if
    if (lb%window_bounds /= MPI_WIN_NULL) then
      call MPI_Win_free(lb%window_bounds, ierr)
    end if
    if (lb%window_done /= MPI_WIN_NULL) then
      call MPI_Win_free(lb%window_done, ierr)
    end if
  end subroutine clean
end module SLB4MPI_work_stealing_load_balancer_m
program main
  use MPI
  use SLB4MPI_work_stealing_load_balancer_m
  implicit none
  type(work_stealing_load_balancer_t) :: lb
  integer(8) :: bot, top, i, j, k
  integer(4) :: mpierr

  call MPI_init(mpierr)
  call lb%initialize(MPI_COMM_WORLD, 1_8, 300_8, 4_8, 8_8)
  do
    if (.not.lb%get_range(bot, top)) exit
    do j = bot, top
    block
      real(8), allocatable :: a(:), b(:), c
      i = 4 * (300 - j) + 200
      allocate(a(i),b(i),c)
      c = 0e0_8
      do k = 1, i * 100
      call random_number(a)
      call random_number(b)
      c = dot_product(a,b) + c
      end do
      write(7,'(I6,ES15.8)') j, c
      deallocate(a,b,c)
    end block
    end do
  end do
  print '(A)', "Done!"
  call lb%clean()
  call MPI_Finalize(mpierr)
end program main
