m_a2_utils.f90

module m_a2_utils
  use m_a2_types

  implicit none
  private

  ! Public subroutines
  public :: a2_loop_box
  public :: a2_loop_box_arg
  public :: a2_loop_boxes
  public :: a2_loop_boxes_arg
  public :: a2_tree_clear_cc
  public :: a2_box_clear_cc
  public :: a2_tree_clear_ghostcells
  public :: a2_box_clear_ghostcells
  public :: a2_box_add_cc
  public :: a2_box_sub_cc
  public :: a2_tree_times_cc
  public :: a2_tree_apply
  public :: a2_box_times_cc
  public :: a2_box_lincomb_cc
  public :: a2_box_copy_cc_to
  public :: a2_box_copy_cc
  public :: a2_boxes_copy_cc
  public :: a2_tree_copy_cc
  public :: a2_reduction
  public :: a2_reduction_vec
  public :: a2_reduction_loc
  public :: a2_tree_max_cc
  public :: a2_tree_min_cc
  public :: a2_tree_max_fc
  public :: a2_tree_sum_cc
  public :: a2_box_copy_fc
  public :: a2_boxes_copy_fc
  public :: a2_tree_copy_fc

  ! Public functions
  public :: a2_get_id_at
  public :: a2_get_loc
  public :: a2_r_loc
  public :: a2_r_inside
  public :: a2_n_cell

contains

  subroutine a2_loop_box(tree, my_procedure, leaves_only)
    type(a2_t), intent(inout)     :: tree
    procedure(a2_subr)           :: my_procedure
    logical, intent(in), optional :: leaves_only
    logical                       :: leaves
    integer                       :: lvl, i, id

    leaves = .false.; if (present(leaves_only)) leaves = leaves_only
    if (.not. tree%ready) stop "a2_loop_box: set_base has not been called"

    !$omp parallel private(lvl, i, id)
    do lvl = lbound(tree%lvls, 1), tree%highest_lvl
       if (leaves) then
          !$omp do
          do i = 1, size(tree%lvls(lvl)%leaves)
             id = tree%lvls(lvl)%leaves(i)
             call my_procedure(tree%boxes(id))
          end do
          !$omp end do
       else
          !$omp do
          do i = 1, size(tree%lvls(lvl)%ids)
             id = tree%lvls(lvl)%ids(i)
             call my_procedure(tree%boxes(id))
          end do
          !$omp end do
       end if
    end do
    !$omp end parallel
  end subroutine a2_loop_box

  subroutine a2_loop_box_arg(tree, my_procedure, rarg, leaves_only)
    type(a2_t), intent(inout)     :: tree
    procedure(a2_subr_arg)        :: my_procedure
    real(dp), intent(in)          :: rarg(:)
    logical, intent(in), optional :: leaves_only
    logical                       :: leaves
    integer                       :: lvl, i, id

    if (.not. tree%ready) stop "Tree not ready"
    leaves = .false.; if (present(leaves_only)) leaves = leaves_only

    !$omp parallel private(lvl, i, id)
    do lvl = lbound(tree%lvls, 1), tree%highest_lvl
       if (leaves) then
          !$omp do
          do i = 1, size(tree%lvls(lvl)%leaves)
             id = tree%lvls(lvl)%leaves(i)
             call my_procedure(tree%boxes(id), rarg)
          end do
          !$omp end do
       else
          !$omp do
          do i = 1, size(tree%lvls(lvl)%ids)
             id = tree%lvls(lvl)%ids(i)
             call my_procedure(tree%boxes(id), rarg)
          end do
          !$omp end do
       end if
    end do
    !$omp end parallel
  end subroutine a2_loop_box_arg

  subroutine a2_loop_boxes(tree, my_procedure, leaves_only)
    type(a2_t), intent(inout)     :: tree
    procedure(a2_subr_boxes)      :: my_procedure
    logical, intent(in), optional :: leaves_only
    logical                       :: leaves
    integer                       :: lvl, i, id

    if (.not. tree%ready) stop "Tree not ready"
    leaves = .false.; if (present(leaves_only)) leaves = leaves_only

    !$omp parallel private(lvl, i, id)
    do lvl = lbound(tree%lvls, 1), tree%highest_lvl
       if (leaves) then
          !$omp do
          do i = 1, size(tree%lvls(lvl)%leaves)
             id = tree%lvls(lvl)%leaves(i)
             call my_procedure(tree%boxes, id)
          end do
          !$omp end do
       else
          !$omp do
          do i = 1, size(tree%lvls(lvl)%ids)
             id = tree%lvls(lvl)%ids(i)
             call my_procedure(tree%boxes, id)
          end do
          !$omp end do
       end if
    end do
    !$omp end parallel
  end subroutine a2_loop_boxes

  subroutine a2_loop_boxes_arg(tree, my_procedure, rarg, leaves_only)
    type(a2_t), intent(inout)    :: tree
    procedure(a2_subr_boxes_arg) :: my_procedure
    real(dp), intent(in)         :: rarg(:)
    logical, intent(in), optional :: leaves_only
    logical                       :: leaves
    integer                      :: lvl, i, id

    if (.not. tree%ready) stop "Tree not ready"
    leaves = .false.; if (present(leaves_only)) leaves = leaves_only

    !$omp parallel private(lvl, i, id)
    do lvl = lbound(tree%lvls, 1), tree%highest_lvl
       if (leaves) then
          !$omp do
          do i = 1, size(tree%lvls(lvl)%leaves)
             id = tree%lvls(lvl)%leaves(i)
             call my_procedure(tree%boxes, id, rarg)
          end do
          !$omp end do
       else
          !$omp do
          do i = 1, size(tree%lvls(lvl)%ids)
             id = tree%lvls(lvl)%ids(i)
             call my_procedure(tree%boxes, id, rarg)
          end do
          !$omp end do
       end if
    end do
    !$omp end parallel
  end subroutine a2_loop_boxes_arg

  pure function a2_r_inside(box, r, d) result(inside)
    type(box2_t), intent(in)       :: box
    real(dp), intent(in)           :: r(2)
    real(dp), intent(in), optional :: d
    real(dp)                       :: r_max(2)
    logical                        :: inside

    r_max = box%r_min + box%dr * box%n_cell
    if (present(d)) then
       inside = all(r+d >= box%r_min) .and. all(r-d <= r_max)
    else
       inside = all(r >= box%r_min) .and. all(r <= r_max)
    end if
  end function a2_r_inside

  pure function a2_get_id_at(tree, rr, highest_lvl, guess) result(id)
    type(a2_t), intent(in)       :: tree
    real(dp), intent(in)          :: rr(2)
    integer, intent(in), optional :: highest_lvl
    integer, intent(in), optional :: guess
    integer                       :: id

    integer                       :: i, i_ch, lvl_max

    if (present(guess)) then
       id = guess
       if (id > 0 .and. id < tree%highest_id) then
          if (tree%boxes(id)%in_use .and. &
               a2_r_inside(tree%boxes(id), rr)) then

             ! Jump into potential children for as long as possible
             do while (a2_has_children(tree%boxes(id)))
                i_ch = child_that_contains(tree%boxes(id), rr)
                id = tree%boxes(id)%children(i_ch)
             end do

             ! return             ! Exit routine
          end if
       end if
    end if

    lvl_max = tree%lvl_limit
    if (present(highest_lvl)) lvl_max = highest_lvl

    ! Find lvl 1 box that includes rr
    do i = 1, size(tree%lvls(1)%ids)
       id = tree%lvls(1)%ids(i)
       if (a2_r_inside(tree%boxes(id), rr)) exit
    end do

    if (i > size(tree%lvls(1)%ids)) then
       ! Not inside any box
       id = -1
    else
       ! Jump into children for as long as possible
       do
          if (tree%boxes(id)%lvl >= lvl_max .or. &
               .not. a2_has_children(tree%boxes(id))) exit
          i_ch = child_that_contains(tree%boxes(id), rr)
          id = tree%boxes(id)%children(i_ch)
       end do
    end if

  end function a2_get_id_at

  pure function a2_get_loc(tree, rr, highest_lvl, guess) result(loc)
    type(a2_t), intent(in)       :: tree
    real(dp), intent(in)          :: rr(2)
    integer, intent(in), optional :: highest_lvl
    integer, intent(in), optional :: guess
    type(a2_loc_t)               :: loc

    loc%id = a2_get_id_at(tree, rr, highest_lvl, guess)

    if (loc%id == -1) then
       loc%ix = -1
    else
       loc%ix = a2_cc_ix(tree%boxes(loc%id), rr)

       ! Fix indices for points exactly on the boundaries of a box (which could
       ! get a ghost cell index)
       where (loc%ix < 1) loc%ix = 1
       where (loc%ix > tree%n_cell) loc%ix = tree%n_cell
    end if
  end function a2_get_loc

  pure function child_that_contains(box, rr) result(i_ch)
    type(box2_t), intent(in) :: box
    real(dp), intent(in)      :: rr(2)
    integer                   :: i_ch
    real(dp)                  :: center(2)

    i_ch   = 1
    center = box%r_min + box%dr * ishft(box%n_cell, -1)

    if (rr(1) > center(1)) i_ch = i_ch + 1
    if (rr(2) > center(2)) i_ch = i_ch + 2
  end function child_that_contains

  pure function a2_r_loc(tree, loc) result(r)
    type(a2_t), intent(in)     :: tree
    type(a2_loc_t), intent(in) :: loc
    real(dp)                   :: r(2)
    r = tree%boxes(loc%id)%r_min + &
         (loc%ix-0.5_dp) * tree%boxes(loc%id)%dr
  end function a2_r_loc

  subroutine a2_tree_clear_cc(tree, iv)
    type(a2_t), intent(inout) :: tree
    integer, intent(in)        :: iv
    integer                    :: lvl, i, id

    !$omp parallel private(lvl, i, id)
    do lvl = lbound(tree%lvls, 1), tree%highest_lvl
       !$omp do
       do i = 1, size(tree%lvls(lvl)%ids)
          id = tree%lvls(lvl)%ids(i)
          call a2_box_clear_cc(tree%boxes(id), iv)
       end do
       !$omp end do
    end do
    !$omp end parallel
  end subroutine a2_tree_clear_cc

  subroutine a2_box_clear_cc(box, iv)
    type(box2_t), intent(inout) :: box
    integer, intent(in)         :: iv
    box%cc(:,:, iv) = 0
  end subroutine a2_box_clear_cc

  subroutine a2_tree_clear_ghostcells(tree, iv)
    type(a2_t), intent(inout) :: tree
    integer, intent(in)        :: iv
    integer                    :: lvl, i, id

    !$omp parallel private(lvl, i, id)
    do lvl = lbound(tree%lvls, 1), tree%highest_lvl
       !$omp do
       do i = 1, size(tree%lvls(lvl)%ids)
          id = tree%lvls(lvl)%ids(i)
          call a2_box_clear_ghostcells(tree%boxes(id), iv)
       end do
       !$omp end do
    end do
    !$omp end parallel
  end subroutine a2_tree_clear_ghostcells

  subroutine a2_box_clear_ghostcells(box, iv)
    type(box2_t), intent(inout) :: box
    integer, intent(in)          :: iv
    integer                      :: nc

    nc = box%n_cell

    box%cc(0, :, iv)       = 0
    box%cc(nc+1, :, iv)    = 0
    box%cc(:, 0, iv)       = 0
    box%cc(:, nc+1, iv)    = 0
  end subroutine a2_box_clear_ghostcells

  subroutine a2_box_add_cc(box, iv_from, iv_to)
    type(box2_t), intent(inout) :: box
    integer, intent(in)         :: iv_from, iv_to
    box%cc(:,:, iv_to) = box%cc(:,:, iv_to) + box%cc(:,:, iv_from)
  end subroutine a2_box_add_cc

  subroutine a2_box_sub_cc(box, iv_from, iv_to)
    type(box2_t), intent(inout) :: box
    integer, intent(in)         :: iv_from, iv_to
    box%cc(:,:, iv_to) = box%cc(:,:, iv_to) - box%cc(:,:, iv_from)
  end subroutine a2_box_sub_cc

  subroutine a2_tree_apply(tree, iv_a, iv_b, op, eps)
    type(a2_t), intent(inout)     :: tree
    integer, intent(in)            :: iv_a, iv_b
    character(len=*), intent(in)   :: op
    real(dp), intent(in), optional :: eps
    integer                        :: lvl, i, id
    real(dp)                       :: use_eps

    use_eps = sqrt(tiny(1.0_dp))
    if (present(eps)) use_eps = eps

    !$omp parallel private(lvl, i, id)
    do lvl = lbound(tree%lvls, 1), tree%highest_lvl
       !$omp do
       do i = 1, size(tree%lvls(lvl)%ids)
          id = tree%lvls(lvl)%ids(i)
          select case (op)
          case ('+')
             tree%boxes(id)%cc(:, :, iv_a) = &
                  tree%boxes(id)%cc(:, :, iv_a) + &
                  tree%boxes(id)%cc(:, :, iv_b)
          case ('*')
             tree%boxes(id)%cc(:, :, iv_a) = &
                  tree%boxes(id)%cc(:, :, iv_a) * &
                  tree%boxes(id)%cc(:, :, iv_b)
          case ('/')
             tree%boxes(id)%cc(:, :, iv_a) = &
                  tree%boxes(id)%cc(:, :, iv_a) / &
                  max(tree%boxes(id)%cc(:, :, iv_b), eps)
          case default
             error stop "a2_tree_apply: unknown operand"
          end select
       end do
       !$omp end do
    end do
    !$omp end parallel
  end subroutine a2_tree_apply

  subroutine a2_tree_times_cc(tree, ivs, facs)
    type(a2_t), intent(inout) :: tree
    integer, intent(in)        :: ivs(:)
    real(dp), intent(in)       :: facs(:)
    integer                    :: lvl, i, id, n

    if (size(ivs) /= size(facs)) &
         error stop "a2_times_cc: invalid array size"

    !$omp parallel private(lvl, i, id, n)
    do lvl = lbound(tree%lvls, 1), tree%highest_lvl
       !$omp do
       do i = 1, size(tree%lvls(lvl)%ids)
          id = tree%lvls(lvl)%ids(i)
          do n = 1, size(ivs)
             call a2_box_times_cc(tree%boxes(id), facs(n), ivs(n))
          end do
       end do
       !$omp end do
    end do
    !$omp end parallel
  end subroutine a2_tree_times_cc

  subroutine a2_box_times_cc(box, a, iv)
    type(box2_t), intent(inout) :: box
    real(dp), intent(in)        :: a
    integer, intent(in)         :: iv
    box%cc(:,:, iv) = a * box%cc(:,:, iv)
  end subroutine a2_box_times_cc

  subroutine a2_box_lincomb_cc(box, a, iv_a, b, iv_b)
    type(box2_t), intent(inout) :: box
    real(dp), intent(in)        :: a, b
    integer, intent(in)         :: iv_a, iv_b
    box%cc(:,:, iv_b) = a * box%cc(:,:, iv_a) + b * box%cc(:,:, iv_b)
  end subroutine a2_box_lincomb_cc

  subroutine a2_box_copy_cc_to(box_from, iv_from, box_to, iv_to)
    type(box2_t), intent(in)    :: box_from
    type(box2_t), intent(inout) :: box_to
    integer, intent(in)         :: iv_from, iv_to
    box_to%cc(:,:, iv_to) = box_from%cc(:,:, iv_from)
  end subroutine a2_box_copy_cc_to

  subroutine a2_box_copy_cc(box, iv_from, iv_to)
    type(box2_t), intent(inout) :: box
    integer, intent(in)         :: iv_from, iv_to
    box%cc(:,:, iv_to) = box%cc(:,:, iv_from)
  end subroutine a2_box_copy_cc

  subroutine a2_boxes_copy_cc(boxes, ids, iv_from, iv_to)
    type(box2_t), intent(inout) :: boxes(:)
    integer, intent(in)         :: ids(:), iv_from, iv_to
    integer                     :: i

    !$omp parallel do
    do i = 1, size(ids)
       call a2_box_copy_cc(boxes(ids(i)), iv_from, iv_to)
    end do
    !$omp end parallel do
  end subroutine a2_boxes_copy_cc

  subroutine a2_tree_copy_cc(tree, iv_from, iv_to)
    type(a2_t), intent(inout) :: tree
    integer, intent(in)       :: iv_from, iv_to
    integer                   :: lvl

    do lvl = lbound(tree%lvls, 1), tree%highest_lvl
       call a2_boxes_copy_cc(tree%boxes, tree%lvls(lvl)%ids, iv_from, iv_to)
    end do
  end subroutine a2_tree_copy_cc

  subroutine a2_reduction(tree, box_func, reduction, init_val, out_val)
    type(a2_t), intent(in) :: tree
    real(dp), intent(in)   :: init_val
    real(dp), intent(out)  :: out_val
    real(dp)               :: tmp, my_val
    integer                :: i, id, lvl

    interface

       real(dp) function box_func(box)
         import
         type(box2_t), intent(in) :: box
       end function box_func

       real(dp) function reduction(a, b)
         import
         real(dp), intent(in) :: a, b
       end function reduction
    end interface

    if (.not. tree%ready) stop "Tree not ready"
    out_val = init_val
    my_val  = init_val

    !$omp parallel private(lvl, i, id, tmp) firstprivate(my_val)
    do lvl = lbound(tree%lvls, 1), tree%highest_lvl
       !$omp do
       do i = 1, size(tree%lvls(lvl)%leaves)
          id = tree%lvls(lvl)%leaves(i)
          tmp = box_func(tree%boxes(id))
          my_val = reduction(tmp, my_val)
       end do
       !$omp end do
    end do

    !$omp critical
    out_val = reduction(my_val, out_val)
    !$omp end critical
    !$omp end parallel
  end subroutine a2_reduction

  subroutine a2_reduction_vec(tree, box_func, reduction, init_val, &
       out_val, n_vals)
    type(a2_t), intent(in) :: tree
    integer, intent(in)     :: n_vals
    real(dp), intent(in)    :: init_val(n_vals)
    real(dp), intent(out)   :: out_val(n_vals)
    real(dp)                :: tmp(n_vals), my_val(n_vals)
    integer                 :: i, id, lvl

    interface

       function box_func(box, n_vals) result(vec)
         import
         type(box2_t), intent(in) :: box
         integer, intent(in)       :: n_vals
         real(dp)                  :: vec(n_vals)
       end function box_func

       function reduction(vec_1, vec_2, n_vals) result(vec)
         import
         integer, intent(in)  :: n_vals
         real(dp), intent(in) :: vec_1(n_vals), vec_2(n_vals)
         real(dp)             :: vec(n_vals)
       end function reduction
    end interface

    if (.not. tree%ready) stop "Tree not ready"
    out_val = init_val
    my_val  = init_val

    !$omp parallel private(lvl, i, id, tmp) firstprivate(my_val)
    do lvl = lbound(tree%lvls, 1), tree%highest_lvl
       !$omp do
       do i = 1, size(tree%lvls(lvl)%leaves)
          id = tree%lvls(lvl)%leaves(i)
          tmp = box_func(tree%boxes(id), n_vals)
          my_val = reduction(tmp, my_val, n_vals)
       end do
       !$omp end do
    end do

    !$omp critical
    out_val = reduction(my_val, out_val, n_vals)
    !$omp end critical
    !$omp end parallel
  end subroutine a2_reduction_vec

  subroutine a2_reduction_loc(tree, iv, box_subr, reduction, &
       init_val, out_val, out_loc)
    type(a2_t), intent(in)      :: tree
    integer, intent(in)          :: iv
    real(dp), intent(in)         :: init_val
    real(dp), intent(out)        :: out_val
    type(a2_loc_t), intent(out) :: out_loc
    real(dp)                     :: tmp, new_val, my_val
    integer                      :: i, id, lvl, tmp_ix(2)
    type(a2_loc_t)              :: my_loc

    interface

       subroutine box_subr(box, iv, val, ix)
         import
         type(box2_t), intent(in) :: box
         integer, intent(in)       :: iv
         real(dp), intent(out)     :: val
         integer, intent(out)      :: ix(2)
       end subroutine box_subr

       real(dp) function reduction(a, b)
         import
         real(dp), intent(in)      :: a, b
       end function reduction
    end interface

    if (.not. tree%ready) stop "Tree not ready"
    out_val   = init_val
    my_val    = init_val
    my_loc%id = -1
    my_loc%ix = -1

    !$omp parallel private(lvl, i, id, tmp, tmp_ix, new_val) &
    !$omp firstprivate(my_val, my_loc)
    do lvl = lbound(tree%lvls, 1), tree%highest_lvl
       !$omp do
       do i = 1, size(tree%lvls(lvl)%leaves)
          id = tree%lvls(lvl)%leaves(i)
          call box_subr(tree%boxes(id), iv, tmp, tmp_ix)
          new_val = reduction(tmp, my_val)
          if (abs(new_val - my_val) > 0) then
             my_loc%id = id
             my_loc%ix = tmp_ix
             my_val = tmp
          end if
       end do
       !$omp end do
    end do

    !$omp critical
    new_val = reduction(my_val, out_val)
    if (abs(new_val - out_val) > 0) then
       out_loc%id = my_loc%id
       out_loc%ix = my_loc%ix
       out_val = my_val
    end if
    !$omp end critical
    !$omp end parallel
  end subroutine a2_reduction_loc

  subroutine a2_tree_max_cc(tree, iv, cc_max, loc)
    type(a2_t), intent(in)      :: tree
    integer, intent(in)          :: iv
    real(dp), intent(out)        :: cc_max
    type(a2_loc_t), intent(out), optional :: loc
    type(a2_loc_t)                        :: tmp_loc

    call a2_reduction_loc(tree, iv, box_max_cc, reduce_max, &
         -huge(1.0_dp)/10, cc_max, tmp_loc)
    if (present(loc)) loc = tmp_loc
  end subroutine a2_tree_max_cc

  subroutine a2_tree_min_cc(tree, iv, cc_min, loc)
    type(a2_t), intent(in)      :: tree
    integer, intent(in)          :: iv
    real(dp), intent(out)        :: cc_min
    type(a2_loc_t), intent(out), optional :: loc
    type(a2_loc_t)                        :: tmp_loc

    call a2_reduction_loc(tree, iv, box_min_cc, reduce_min, &
         huge(1.0_dp)/10, cc_min, tmp_loc)

    if (present(loc)) loc = tmp_loc
  end subroutine a2_tree_min_cc

  subroutine a2_tree_max_fc(tree, dim, iv, fc_max, loc)
    type(a2_t), intent(in)      :: tree
    integer, intent(in)          :: dim
    integer, intent(in)          :: iv
    real(dp), intent(out)        :: fc_max
    type(a2_loc_t), intent(out), optional :: loc
    type(a2_loc_t)                        :: tmp_loc
    integer                                :: dim_iv

    ! Encode dim and iv in a single variable
    dim_iv = (dim-1) * tree%n_var_face + iv - 1

    call a2_reduction_loc(tree, dim_iv, box_max_fc, reduce_max, &
         -huge(1.0_dp)/10, fc_max, tmp_loc)
    if (present(loc)) loc = tmp_loc
  end subroutine a2_tree_max_fc

  subroutine box_max_cc(box, iv, val, ix)
    type(box2_t), intent(in) :: box
    integer, intent(in)       :: iv
    real(dp), intent(out)     :: val
    integer, intent(out)      :: ix(2)
    integer                   :: nc

    nc = box%n_cell
    ix = maxloc(box%cc(1:nc, 1:nc, iv))
    val = box%cc(ix(1), ix(2), iv)
  end subroutine box_max_cc

  subroutine box_min_cc(box, iv, val, ix)
    type(box2_t), intent(in) :: box
    integer, intent(in)       :: iv
    real(dp), intent(out)     :: val
    integer, intent(out)      :: ix(2)
    integer                   :: nc

    nc = box%n_cell
    ix = minloc(box%cc(1:nc, 1:nc, iv))
    val = box%cc(ix(1), ix(2), iv)
  end subroutine box_min_cc

  subroutine box_max_fc(box, dim_iv, val, ix)
    type(box2_t), intent(in) :: box
    integer, intent(in)       :: dim_iv
    real(dp), intent(out)     :: val
    integer, intent(out)      :: ix(2)
    integer                   :: dim, iv, n_fc, nc, dix(2)

    nc     = box%n_cell
    dix(:) = 0

    n_fc = size(box%fc, 4)

    ! Decode dim and iv
    dim      = dim_iv / n_fc + 1
    iv       = dim_iv - (dim-1) * n_fc + 1
    ! Also include fluxes at 'upper' boundary
    dix(dim) = 1
    ix       = maxloc(box%fc(1:nc+dix(1), 1:nc+dix(2), dim, iv))
    val      = box%fc(ix(1), ix(2), dim, iv)
  end subroutine box_max_fc

  real(dp) function reduce_max(a, b)
    real(dp), intent(in) :: a, b
    reduce_max = max(a, b)
  end function reduce_max

  real(dp) function reduce_min(a, b)
    real(dp), intent(in) :: a, b
    reduce_min = min(a, b)
  end function reduce_min

  subroutine a2_tree_sum_cc(tree, iv, cc_sum)
    type(a2_t), intent(in) :: tree
    integer, intent(in)    :: iv
    real(dp), intent(out)  :: cc_sum
    real(dp)               :: tmp, my_sum, fac
    integer                :: i, id, lvl, nc

    if (.not. tree%ready) stop "Tree not ready"
    my_sum = 0

    !$omp parallel reduction(+: my_sum) private(lvl, i, id, nc, tmp, fac)
    do lvl = lbound(tree%lvls, 1), tree%highest_lvl
       fac = a2_lvl_dr(tree, lvl)**2

       !$omp do
       do i = 1, size(tree%lvls(lvl)%leaves)
          id = tree%lvls(lvl)%leaves(i)
          nc = tree%boxes(id)%n_cell
          if (tree%coord_t == af_cyl) then
             tmp = sum_2pr_box(tree%boxes(id), iv)
          else
             tmp = sum(tree%boxes(id)%cc(1:nc, 1:nc, iv))
          end if
          my_sum = my_sum + fac * tmp
       end do
       !$omp end do
    end do
    !$omp end parallel

    cc_sum = my_sum

  contains

    ! Sum of 2 * pi * r * values
    pure function sum_2pr_box(box, iv) result(res)
      type(box2_t), intent(in) :: box
      integer, intent(in)      :: iv
      real(dp), parameter      :: twopi = 2 * acos(-1.0_dp)
      real(dp)                 :: res
      integer                  :: i, j, nc

      res = 0
      nc  = box%n_cell

      do j = 1, nc
         do i = 1, nc
            res = res + box%cc(i, j, iv) * a2_cyl_radius_cc(box, i)
         end do
      end do
      res = res * twopi
    end function sum_2pr_box
  end subroutine a2_tree_sum_cc

  subroutine a2_box_copy_fc(box, iv_from, iv_to)
    type(box2_t), intent(inout) :: box
    integer, intent(in)         :: iv_from
    integer, intent(in)         :: iv_to
    box%fc(:,:,:, iv_to) = box%fc(:,:,:, iv_from)
  end subroutine a2_box_copy_fc

  subroutine a2_boxes_copy_fc(boxes, ids, iv_from, iv_to)
    type(box2_t), intent(inout) :: boxes(:)
    integer, intent(in)         :: ids(:)
    integer, intent(in)         :: iv_from
    integer, intent(in)         :: iv_to
    integer                     :: i

    !$omp parallel do
    do i = 1, size(ids)
       call a2_box_copy_fc(boxes(ids(i)), iv_from, iv_to)
    end do
    !$omp end parallel do
  end subroutine a2_boxes_copy_fc

  subroutine a2_tree_copy_fc(tree, iv_from, iv_to)
    type(a2_t), intent(inout) :: tree
    integer, intent(in)       :: iv_from
    integer, intent(in)       :: iv_to
    integer                   :: lvl

    if (.not. tree%ready) stop "Tree not ready"
    do lvl = lbound(tree%lvls, 1), tree%highest_lvl
       call a2_boxes_copy_fc(tree%boxes, tree%lvls(lvl)%ids, iv_from, iv_to)
    end do
  end subroutine a2_tree_copy_fc

  pure function a2_n_cell(tree, lvl) result(n_cell)
    type(a2_t), intent(in) :: tree
    integer, intent(in)    :: lvl
    integer                :: n_cell

    if (lvl >= 1) then
       n_cell = tree%n_cell
    else
       n_cell = tree%n_cell / (2**(1-lvl))
    end if
  end function a2_n_cell

end module m_a2_utils