module test_variable_m
    !! Test of module variable_mk
    use pfunit
    use MPI
    implicit none
contains

    @test(npes = [4])
    subroutine test_variable(this)
        !! Tests variable data structure
        use precision_grillix_m, only : GP
        use helpers_m, only : almost_equal
        use comm_handler_m, only : comm_handler_t
        use variable_m, only : variable_t
        
        class (MpiTestMethod), intent(inout) :: this
        integer :: comm_world, rank, nprocs, ierr
        
        type(comm_handler_t) :: comm_handler
        logical :: fine
        integer :: i, k, ndim, nhalo_fwd, nhalo_bwd, ihalo, src
        real(GP) :: rval
        integer, dimension(:), allocatable :: ndim_halos
        type(variable_t) :: var
        real(GP), dimension(:), pointer :: uh => null()
        
        comm_world = this%getMpiCommunicator()
        call MPI_comm_rank(comm_world, rank, ierr)
        call MPI_comm_size(comm_world, nprocs, ierr)

        if (rank == 0) then     
            write(*,*)''
            write(*,*)''
            write(*,'(A80)')'test_variable -------------------------------------------------------------'
        endif

        call comm_handler%init(comm_world, 4, 1)

        ! Define dimensions, varying across ranks
        ndim = 5 + rank*(-1)**rank
        allocate(ndim_halos(0:nprocs-1))
        do k = 0, nprocs-1
            ndim_halos(k) = 5 + k*(-1)**k
        enddo
        
        ! Defines some variable and halos
        call var%init(comm_handler, ndim, staggered=.false.)
        do i = 1, ndim
            var%vals(i) = 1.0_GP * (10 * rank + i)
        enddo
        
        nhalo_fwd = 6
        nhalo_bwd = 7
        call var%create_halos(comm_handler, nhalo_fwd, nhalo_bwd)
        
        ! Check dimension
        @assertEqual(var%get_ndim(), ndim)
                
        ! Check halo communication
        call var%start_comm_halos(comm_handler)
        call var%finalize_comm_halos(comm_handler)

        do ihalo = -nhalo_bwd, nhalo_fwd
            src = rank + ihalo
            src = modulo(src, nprocs)
            
            uh => var%get_halo(comm_handler, ihalo)
            
            @assertEqual(size(uh), ndim_halos(src))
            
            do i = 1, ndim_halos(src)
                rval = 1.0_GP * (10 * src + i)
                fine = almost_equal(uh(i), rval, rtol=0.0_GP, atol=1.0E-10_GP)
                @assertTrue(fine)
                if (ihalo == -1) then
                    @assertEqual(size(var%hbwd), ndim_halos(src))
            
                    fine = almost_equal(var%hbwd(i), rval, &
                                        rtol=0.0_GP, atol=1.0E-10_GP) 
                    @assertTrue(fine)
                endif
                if (ihalo == 1) then
                    @assertEqual(size(var%hfwd), ndim_halos(src))
                    
                    fine = almost_equal(var%hfwd(i), rval, &
                                        rtol=0.0_GP, atol=1.0E-10_GP) 
                    @assertTrue(fine)
                endif
            enddo
            
            uh => null()
        enddo
        
        if (rank == 0) then     
            write(*,'(A80)')'test_variable -------------------------------------------------------------'
            write(*,*)''
            write(*,*)''
        endif

    end subroutine
    
    @test(npes = [4])
    subroutine test_variable_set(this)
        !! Tests set_blob_toroidal and set_blob_aligned of variables
        use precision_grillix_m, only : GP
        use helpers_m, only : almost_equal
        use helpers_mesh_sample_m, only : mesh_sample_circular
        use constants_m, only : TWO_PI
        use comm_handler_m, only : comm_handler_t
        use equilibrium_m, only : equilibrium_t
        use mesh_cart_m, only : mesh_cart_t
        use variable_m, only : variable_t
        
        class (MpiTestMethod), intent(inout) :: this
        integer :: comm_world, rank, nprocs, ierr
        
        type(comm_handler_t) :: comm_handler
        class(equilibrium_t), allocatable :: equi
        type(mesh_cart_t) :: mesh
        type(variable_t) :: var
        real(GP) :: x0, y0, phi0, wx, wy, wphi, wpar, amp
        real(GP) :: x, y, phi, dphi, ref
        integer :: l, ic, jc
        logical :: fine
               
        comm_world = this%getMpiCommunicator()
        call MPI_comm_rank(comm_world, rank, ierr)
        call MPI_comm_size(comm_world, nprocs, ierr)

        if (rank == 0) then     
            write(*,*)''
            write(*,*)''
            write(*,'(A80)')'test_variable_set ---------------------------------------------------------'
        endif

        call comm_handler%init(comm_world, 4, 1)
        dphi = TWO_PI / comm_handler%get_nplanes()
        phi = dphi * comm_handler%get_plane()


        ! Create sample equilibrium and mesh 
        call mesh_sample_circular(equi, phi, mesh)
        call var%init(comm_handler, ndim = mesh%get_n_points(), &
                      staggered=.false.)
        
        ! Toroidally aligned blob
        x0 = 0.3_GP
        y0 = 0.05_GP
        phi0 = 1.0_GP
        wx = 0.02_GP
        wy = 0.03_GP
        wphi = 1.0_GP
        amp = 2.0_GP
        
        call var%set_blob_toroidal(mesh, x0, y0, phi0, wx, wy, wphi, amp)
        
        do l = 1, mesh%get_n_points()
            x = mesh%get_x(l)
            y = mesh%get_y(l)
            
            ref = amp * exp( -(x-x0)**2 / wx**2 ) &
                  * exp( -(y-y0)**2 / wy**2 ) &
                  * exp( -(phi-phi0)**2 / wphi**2 )
            
            fine = almost_equal(var%vals(l), ref, rtol=0.0_GP, atol=1.0E-10_GP)
            @assertTrue(fine)
        
        enddo
        
        ! Field aligned blob
        x0 = 0.3_GP
        y0 = 0.05_GP
        phi0 = 0.0_GP
        wx = 0.02_GP
        wy = 0.03_GP
        wpar = 1.5_GP
        amp = 2.0_GP
        call var%set_blob_aligned(equi, mesh, x0, y0, phi0, wx, wy, wpar, amp)
        
        ! Check individual values close to maximum on plane
        select case(comm_handler%get_plane())
            case(0)
                ic = 88
                jc = 57
                ref = 1.913057478206_GP
            case(1)
                ic = 83
                jc = 71
                ref = 0.655272210341_GP
            case(2)
                ic = 73
                jc = 82
                ref = 0.024260822117_GP
            case(3)
                ic = 88
                jc = 42
                ref =  0.655145840497_GP
            case default
                @assertTrue(.false.)        
        end select
        l = mesh%cart_to_mesh_index(ic, jc) 
        
        fine = almost_equal(var%vals(l), ref, rtol=0.0_GP, atol=1.0E-10_GP) 
        @assertTrue(fine)
                
        if (rank == 0) then     
            write(*,'(A80)')'test_variable_set ---------------------------------------------------------'
            write(*,*)''
            write(*,*)''
        endif

    end subroutine
    
end module