#include "test-utils.F90"
#include "language-support.F90"

module prif_teams_test_m
# include "test-uses-alloc.F90"
    use prif
    use julienne_m, only: test_description_t, test_diagnosis_t, test_result_t, test_t, string_t, usher &
      ,operator(.also.), operator(.isAtLeast.), operator(.isAtMost.), operator(.equalsExpected.), operator(//)

    implicit none
    private
    public :: prif_teams_test_t

    type, extends(test_t) :: prif_teams_test_t
    contains
      procedure, nopass, non_overridable :: subject
      procedure, nopass, non_overridable :: results
    end type

    integer :: cleanup_count = 0

contains
    pure function subject()
      character(len=:), allocatable :: subject
      subject = "PRIF Teams"
    end function

    function results() result(test_results)
        type(test_result_t), allocatable :: test_results(:)
        type(prif_teams_test_t) prif_teams_test

        allocate(test_results, source = prif_teams_test%run([ &
            test_description_t("creating, changing to, and allocating coarrays", usher(check_teams)) &
        ]))
    end function

    function check_teams() result(diag)
        type(test_diagnosis_t) :: diag

        integer :: dummy_element, i
        integer(c_int) :: initial_num_imgs, num_imgs, me, me_child, x
        integer(c_size_t) :: element_size
        integer(c_int64_t) :: which_team, n
        integer, parameter :: num_coarrays = 10
        type(prif_coarray_handle) :: initial_coarray, coarrays(num_coarrays)
        type(c_ptr) :: allocated_memory
        type(prif_team_type) :: team, initial_team, t, team_self

        diag = .true.

        call prif_num_images(num_images=initial_num_imgs)
        ALSO2(initial_num_imgs .isAtLeast. 1,  "invalid prif_num_images")

        call prif_this_image_no_coarray(this_image=me)
        ALSO2(me .isAtLeast. 1, "invalid prif_this_image")
        ALSO2(me .isAtMost. initial_num_imgs, "invalid prif_this_image")

        n = 0 ! clear outputs
        call prif_team_number(team_number=n)
        ALSO2(n .equalsExpected. -1_c_int64_t, "Initial team number is -1")

        n = 0 ! clear outputs
        call prif_get_team(team=initial_team)
        call prif_team_number(team=initial_team, team_number=n)
        ALSO2(n .equalsExpected. -1_c_int64_t, "prif_get_team retrieval of current initial team")
          
        ! ensure prif_sync_team is usable
        call prif_sync_team(team=initial_team)

        x = 0 ! clear outputs
        call prif_num_images_with_team(team=initial_team, num_images=x)
        ALSO2(x .equalsExpected. initial_num_imgs, "prif_num_images in initial team")
          
        x = 0 ! clear outputs
        call prif_num_images_with_team_number(team_number=-1_c_int64_t, num_images=x)
        ALSO2(x .equalsExpected. initial_num_imgs, "prif_num_images_with_team_number in initial team")

        x = 0 ! clear outputs
        call prif_this_image_no_coarray(team=initial_team, this_image=x)
        ALSO2(x .equalsExpected. me, "prif_this_image_no_coarray in initial team")
          
        t = prif_team_type() ; n = 0 ! clear outputs
        call prif_get_team(level=PRIF_INITIAL_TEAM, team=t)
        call prif_team_number(team=t, team_number=n)
        ALSO2(n .equalsExpected. -1_c_int64_t, "prif_get_team(PRIF_INITIAL_TEAM) retrieves initial team")
          
        t = prif_team_type() ; n = 0 ! clear outputs
        call prif_get_team(level=PRIF_CURRENT_TEAM, team=t)
        call prif_team_number(team=t, team_number=n)
        ALSO2(n .equalsExpected. -1_c_int64_t, "prif_get_team(PRIF_CURRENT_TEAM) retrieval of initial team when current team is initial team")

        t = prif_team_type() ; n = 0 ! clear outputs
        call prif_get_team(level=PRIF_PARENT_TEAM, team=t)
        call prif_team_number(team=t, team_number=n)
        ALSO2(n .equalsExpected. -1_c_int64_t, "prif_get_team(PRIF_PARENT_TEAM) retrieves initial team when parent team is initial team")
          
        element_size = int(storage_size(dummy_element)/8, c_size_t)
        call prif_allocate_coarray( &
            lcobounds = [1_c_int64_t], &
            ucobounds = [integer(c_int64_t)::], &
            size_in_bytes = element_size, &
            final_func = c_null_funptr, &
            coarray_handle = initial_coarray, &
            allocated_memory = allocated_memory)
        n = 0 ! clear outputs
        call prif_ucobound_with_dim(initial_coarray, 1, n)
        ALSO(n .equalsExpected. int(initial_num_imgs, c_int64_t))

        which_team = merge(1_c_int64_t, 2_c_int64_t, mod(me, 2) == 0)
        call prif_form_team(team_number = which_team, team = team)
        call prif_change_team(team)
            call prif_num_images(num_images=num_imgs)
            ALSO2(num_imgs .equalsExpected.  initial_num_imgs/2 + mod(initial_num_imgs,2)*(int(which_team)-1), "Team has correct number of images")

            ! ensure prif_sync_team is usable
            call prif_sync_team(team=team)
            call prif_sync_team(team=initial_team)

            x = 0 ! clear outputs
            call prif_num_images_with_team(team=team, num_images=x)
            ALSO2(x .equalsExpected. num_imgs, "prif_num_images works with team")

            x = 0 ! clear outputs
            call prif_num_images_with_team_number(team_number=which_team, num_images=x)
            ALSO2 (x .equalsExpected. num_imgs, "prif_num_images_with_team_number works with current team")

            call prif_this_image_no_coarray(this_image=me_child)
            ALSO2(me_child .equalsExpected. (me - 1)/2 + 1, "prif_this_image is valid")
              

            x = 0 ! clear outputs
            call prif_this_image_no_coarray(team=team, this_image=x)
            ALSO2(x .equalsExpected. me_child, "prif_this_image is valid")

            n = 0 ! clear outputs
            call prif_team_number(team_number=n)
            ALSO2(n .equalsExpected. which_team, "Correct current team number")
    
            n = 0 ! clear outputs
            call prif_team_number(team=team, team_number=n)
            ALSO2(n .equalsExpected. which_team, "Correct current team number")
    
            t = prif_team_type() ; n = 0 ! clear outputs
            call prif_get_team(team=t)
            call prif_team_number(team=t, team_number=n)
            ALSO2(n .equalsExpected. which_team, "prif_get_team retrieves current team")
    
            t = prif_team_type() ; n = 0 ! clear outputs
            call prif_get_team(level=PRIF_INITIAL_TEAM, team=t)
            call prif_team_number(team=t, team_number=n)
            ALSO2(n .equalsExpected. -1_c_int64_t, "prif_get_team(PRIF_INITIAL_TEAM) retrieves initial team")
    
            t = prif_team_type() ; n = 0 ! clear outputs
            call prif_get_team(level=PRIF_CURRENT_TEAM, team=t)
            call prif_team_number(team=t, team_number=n)
            ALSO2(n .equalsExpected. which_team, "prif_get_team(PRIF_CURRENT_TEAM) retrieves current team")
    
            t = prif_team_type() ; n = 0 ! clear outputs
            call prif_get_team(level=PRIF_PARENT_TEAM, team=t)
            call prif_team_number(team=t, team_number=n)
            ALSO2(n .equalsExpected. -1_c_int64_t, "prif_get_team(PRIF_PARENT_TEAM) retrieves initial team when parent team is initial team")

            x = 0 ! clear outputs
            call prif_num_images_with_team(team=initial_team, num_images=x)
            ALSO2(x .equalsExpected. initial_num_imgs, "prif_num_images works with initial team")

            x = 0 ! clear outputs
            call prif_num_images_with_team_number(team_number=-1_c_int64_t, num_images=x)
            ALSO2(x .equalsExpected. initial_num_imgs, "prif_num_images_with_team_number works with initial team")

            x = 0 ! clear outputs
            call prif_this_image_no_coarray(team=initial_team, this_image=x)
            ALSO2(x .equalsExpected. me, "prif_this_image_no_coarray works with initial team")

            n = 0 ! clear outputs
            call prif_ucobound_with_dim(initial_coarray, 1, n)
            ALSO(n .equalsExpected. int(num_imgs, c_int64_t))

            ALSO(cleanup_count .equalsExpected. 0)
            do i = 1, num_coarrays
                call prif_allocate_coarray( &
                    lcobounds = [1_c_int64_t], &
                    ucobounds = [integer(c_int64_t)::], &
                    size_in_bytes = element_size, &
#if HAVE_FINAL_FUNC_SUPPORT
                    final_func = c_funloc(coarray_cleanup), &
#  define CHECK_COUNT(n) ALSO(cleanup_count .equalsExpected. n)
#else
                    final_func = c_null_funptr, &
#  define CHECK_COUNT(n) 
#endif
                    coarray_handle = coarrays(i), &
                    allocated_memory = allocated_memory)
            end do
            CHECK_COUNT(0)
            call prif_deallocate_coarrays(coarrays(4:4))
            call prif_deallocate_coarrays(coarrays(2:2))
            call prif_deallocate_coarray(coarrays(7))
            CHECK_COUNT(3)

            call prif_form_team(team_number = int(me,c_int64_t), team = team_self)
            x = 0 ! clear outputs
            call prif_num_images_with_team(team=team_self, num_images=x)
            ALSO2(x .equalsExpected. 1, "prif_num_images works with team_self")
            call prif_change_team(team_self)
                x = 0 ! clear outputs
                call prif_num_images(num_images=x)
                ALSO2(x .equalsExpected. 1, "prif_num_images works in team_self")
                x = 0 ! clear outputs
                call prif_this_image_no_coarray(this_image=x)
                ALSO2(x .equalsExpected. 1, "prif_this_image is valid in team_self")
                x = 0 ! clear outputs
                call prif_this_image_no_coarray(team=team, this_image=x)
                ALSO2(x .equalsExpected. me_child, "prif_this_image is valid")
                x = 0 ! clear outputs
                call prif_this_image_no_coarray(team=initial_team, this_image=x)
                ALSO2(x .equalsExpected. me, "prif_this_image is valid")

                ! ensure prif_sync_team is usable
                call prif_sync_team(team=team)
                call prif_sync_team(team=initial_team)
                call prif_sync_team(team=team_self)

                n = 0 ! clear outputs
                call prif_ucobound_with_dim(initial_coarray, 1, n)
                ALSO(n .equalsExpected. 1_c_int64_t)
                n = 0 ! clear outputs
                call prif_ucobound_with_dim(coarrays(3), 1, n)
                ALSO(n .equalsExpected. 1_c_int64_t)

                CHECK_COUNT(3)
            call prif_end_team()

            CHECK_COUNT(3)
        call prif_end_team()
        CHECK_COUNT(num_coarrays)

        ! ensure prif_sync_team is usable
        call prif_sync_team(team=team)
        call prif_sync_team(team=initial_team)

        t = prif_team_type() ; n = 0 ! clear outputs
        call prif_get_team(team=t)
        call prif_team_number(team=t, team_number=n)
        ALSO2(n .equalsExpected. -1_c_int64_t, "prif_end_team restores initial team")

    end function

#if HAVE_FINAL_FUNC_SUPPORT
    subroutine coarray_cleanup(handle, stat, errmsg) bind(C)
      type(prif_coarray_handle), pointer, intent(in) :: handle
      integer(c_int), intent(out) :: stat
      character(len=:), intent(out), allocatable :: errmsg

      cleanup_count = cleanup_count + 1
      stat = 0
    end subroutine
#endif
end module prif_teams_test_m
