module testHelpers
  use testHelperMpi
  use gamapp_mpi
  use voltapp_mpi
  use gam2voltcomm_mpi
  use uservoltic
  use ioH5

  implicit none

  logical :: isGamera = .false.
  type(gamAppMpi_T), allocatable :: gamAppMpi
  type(voltAppMpi_T), allocatable :: voltAppMpi
  type(gam2voltCommMpi_T), allocatable :: g2vCommMpi

contains

  @before
  subroutine emptySetup(this)
      class (MpiTestMethod), intent(inout) :: this
      ! does nothing, fighting with pfunit
  end subroutine

  subroutine manualSetup(this, caseFile)
      class (MpiTestMethod), intent(inout) :: this
      character(len=*), intent(in) :: caseFile

      integer :: ierror
      type(MPI_Comm) :: gamComm, voltComm

      call setMpiReal()

      ! hard-coded for 4 gamera ranks for these tests right now
      if(this%getProcessRank() < 4) then
          isGamera = .true.
          allocate(gamAppMpi)
          allocate(g2vCommMpi)

          ! make gamera-only mpi communicator
          call MPI_Comm_Split(getMpiF08Communicator(this), 0, this%getProcessRank(), gamComm, ierror)

          call initGamera_mpi(gamAppMpi, initUser, gamComm, trim(caseFile), .false.)
          call initGam2Volt(g2vCommMpi, gamAppMpi, getMpiF08Communicator(this), trim(caseFile))
      else
          isGamera = .false.
          allocate(voltAppMpi)

          ! make gamera-only mpi communicator
          call MPI_Comm_Split(getMpiF08Communicator(this), 1, this%getProcessRank(), voltComm, ierror)

          call initVoltron_mpi(voltAppMpi, initUser, voltComm, getMpiF08Communicator(this), trim(caseFile))
      endif

  end subroutine manualSetup

  @after
  subroutine teardown(this)
      class (MpiTestMethod), intent(inout) :: this

      if(allocated(voltAppMpi)) then
          call endVoltronWaits(voltAppMpi)
          deallocate(voltAppMpi)
      endif
      if(allocated(g2vCommMpi)) then
          call endGam2VoltWaits(g2vCommMpi, gamAppMpi)
          deallocate(g2vCommMpi)
      endif
      if(allocated(gamAppMpi)) deallocate(gamAppMpi)

  end subroutine teardown

  subroutine runApplication()
      logical :: helperQuit

      if(isGamera) then
          do while (g2vCommMpi%time < g2vCommMpi%tFin)
              call stepGamera_mpi(gamAppMpi)
              call localStepVoltronTime(g2vCommMpi, gamAppMpi)
              if( (g2vCommMpi%time >= g2vCommMpi%DeepT .and. g2vCommMpi%doDeep) .or. &
                  (g2vCommMpi%time >= g2vCommMpi%tFin)) then
                  call performStepVoltron(g2vCommMpi,gamAppMpi)
                  if ( g2vCommMpi%time >= g2vCommMpi%DeepT .and. g2vCommMpi%doDeep ) then
                      call performDeepUpdate(g2vCommMpi, gamAppMpi)
                  endif
              endif
          end do
      else
          if(voltAppMpi%amHelper) then
              helperQuit = .false.
              do while(.not. helperQuit)
                  call helpVoltron(voltAppMpi, helperQuit)
              end do
          else
              do while (voltAppMpi%time < voltAppMpi%tFin)
                  if(gameraStepReady(voltAppMpi)) then
                      call stepVoltron_mpi(voltAppMpi)
                      write (*,'(a,f8.3)') '    Time = ',voltAppMpi%time
                      if (voltAppMpi%time >= voltAppMpi%DeepT .and. voltAppMpi%doDeep ) then
                          call DeepUpdate_mpi(voltAppMpi, voltAppMpi%time)
                      endif
                  elseif(deepInProgress(voltAppMpi)) then
                      call doDeepBlock(voltAppMpi)
                  else
                      call waitForGameraStep(voltAppMpi)
                  endif
              enddo
              if(voltAppMpi%useHelpers) call vhReqHelperQuit(voltAppMpi)
          endif
      endif

  end subroutine runApplication

  @test(npes=[8])
  subroutine testHelpTimestepping(this)
      class (MpiTestMethod), intent(inout) :: this

      write (*,'(a,I0)') 'Testing HelpTimestepping ',this%getNumProcesses()

      call manualSetup(this, 'testHelpersStep_4.xml')

      if(isGamera) then
          call runApplication()
      else
          call runApplication()
      endif

  end subroutine

  @test(npes=[8])
  subroutine testHelpSquish(this)
      class (MpiTestMethod), intent(inout) :: this

      real(rp), dimension(:,:,:,:), allocatable :: distSquish
      real(rp), dimension(2) :: posErr
      real(rp) :: sqErr, maxSqErr
      integer :: i,j,k
      logical :: helperQuit

      ! debugging
      !type(IOVAR_T), dimension(12) :: IOVars

      write (*,'(a,I0)') 'Testing HelpSquish ',this%getNumProcesses()

      call manualSetup(this, 'testHelpersSquish_4.xml')

      if(isGamera) then
          call runApplication()
      else
          call runApplication()
      endif

      ! Now compare distributed squish results to local-only squish results
      if(.not. isGamera) then
          if(.not. voltAppMpi%amHelper) then
              
              ! call doImag once to update rTrc and nTrc before doing any squishing
              call PreDeep(voltAppMpi, voltAppMpi%gAppLocal)
              call DoImag(voltAppMpi)

              ! calculate distributed squish from most recent results
              call startDeep(voltAppMpi)
              do while(deepInProgress(voltAppMpi))
                  call doDeepBlock(voltAppMpi)
              enddo
              call vhReqHelperQuit(voltAppMpi) ! tell the helpers they're done
              allocate(distSquish,  MOLD=voltAppMpi%chmp2mhd%xyzSquish)
              distSquish = voltAppMpi%chmp2mhd%xyzSquish

              ! now calculate local-only squish
              voltAppMpi%ebTrcApp%ebSquish%myNumBlocks = -1 ! I do all blocks
              voltAppMpi%ebTrcApp%ebSquish%myFirstBlock = 1 ! ensure I start with the first block
              call DeepUpdate(voltAppMpi, voltAppMpi%gAppLocal)

              sqErr = 0
              maxSqErr = 0
              do i=voltAppMpi%ebTrcApp%ebState%ebGr%is,voltAppMpi%iDeep+1
                  do j=voltAppMpi%ebTrcApp%ebState%ebGr%js,voltAppMpi%ebTrcApp%ebState%ebGr%je+1
                      do k=voltAppMpi%ebTrcApp%ebState%ebGr%ks,voltAppMpi%ebTrcApp%ebState%ebGr%ke+1
                          posErr = abs(distSquish(i,j,k,:) - voltAppMpi%chmp2mhd%xyzSquish(i,j,k,:))
                          if(posErr(2) > PI) posErr(2) = 2*PI - posErr(2)
                          sqErr = sqErr + NORM2(posErr)
                          maxSqErr = MAX(maxSqErr, NORM2(posErr))
                      enddo
                  enddo
              enddo

              ! debugging...
              !call ClearIO(IOVars)
              !call AddOutVar(IOVars,"distSquish",distSquish)
              !call AddOutVar(IOVars,"localSquish",voltAppMpi%chmp2mhd%xyzSquish)
              !call WriteVars(IOVars,.false.,"SquishTestData.h5")
              ! ...debugging

              deallocate(distSquish)

              ! Floating point error rounding
              @assertLessThan(sqErr, 1.0e-15_rp, trim("Total Distributed Squish Error Too Large"))
              @assertLessThan(maxSqErr, 1.0e-15_rp, trim("Per Cell Distributed Squish Error Too Large"))

          else
              ! helpers
              helperQuit = .false.
              do while(.not. helperQuit)
                  call helpVoltron(voltAppMpi, helperQuit)
              end do
          endif
      endif

  end subroutine

  @test(npes=[8])
  subroutine testHelpSquishDip(this)
      class (MpiTestMethod), intent(inout) :: this

      real(rp), dimension(:,:,:,:), allocatable :: distSquish
      real(rp), dimension(2) :: posErr
      real(rp) :: sqErr, maxSqErr
      integer :: i,j,k
      logical :: helperQuit

      ! debugging
      !type(IOVAR_T), dimension(12) :: IOVars

      write (*,'(a,I0)') 'Testing HelpSquish ',this%getNumProcesses()

      call manualSetup(this, 'testHelpersSquish_4.xml')

      if(isGamera) then
          call runApplication()
      else
          voltAppMpi%ebTrcApp%ebModel%doDip = .true.
          call runApplication()
      endif

      ! Now compare distributed squish results to local-only squish results
      if(.not. isGamera) then
          if(.not. voltAppMpi%amHelper) then

              ! call doImag once to update rTrc and nTrc before doing any squishing
              call PreDeep(voltAppMpi, voltAppMpi%gAppLocal)
              call DoImag(voltAppMpi)

              ! calculate distributed squish from most recent results
              call startDeep(voltAppMpi)
              do while(deepInProgress(voltAppMpi))
                  call doDeepBlock(voltAppMpi)
              enddo
              call vhReqHelperQuit(voltAppMpi) ! tell the helpers they're done
              allocate(distSquish,  MOLD=voltAppMpi%chmp2mhd%xyzSquish)
              distSquish = voltAppMpi%chmp2mhd%xyzSquish

              ! now calculate local-only squish
              voltAppMpi%ebTrcApp%ebSquish%myNumBlocks = -1 ! I do all blocks
              voltAppMpi%ebTrcApp%ebSquish%myFirstBlock = 1 ! ensure I start with the first block
              call DeepUpdate(voltAppMpi, voltAppMpi%gAppLocal)

              sqErr = 0
              maxSqErr = 0
              do i=voltAppMpi%ebTrcApp%ebState%ebGr%is,voltAppMpi%iDeep+1
                  do j=voltAppMpi%ebTrcApp%ebState%ebGr%js,voltAppMpi%ebTrcApp%ebState%ebGr%je+1
                      do k=voltAppMpi%ebTrcApp%ebState%ebGr%ks,voltAppMpi%ebTrcApp%ebState%ebGr%ke+1
                          posErr = abs(distSquish(i,j,k,:) - voltAppMpi%chmp2mhd%xyzSquish(i,j,k,:))
                          if(posErr(2) > PI) posErr(2) = 2*PI - posErr(2)
                          sqErr = sqErr + NORM2(posErr)
                          maxSqErr = MAX(maxSqErr, NORM2(posErr))
                      enddo
                  enddo
              enddo

              ! debugging...
              !call ClearIO(IOVars)
              !call AddOutVar(IOVars,"distSquish",distSquish)
              !call AddOutVar(IOVars,"localSquish",voltAppMpi%chmp2mhd%xyzSquish)
              !call WriteVars(IOVars,.false.,"SquishTestData.h5")
              ! ...debugging

              deallocate(distSquish)

              ! Floating point error rounding
              @assertLessThan(sqErr, 1.0e-15_rp, trim("Total Distributed Squish Error Too Large"))
              @assertLessThan(maxSqErr, 1.0e-15_rp, trim("Per Cell Distributed Squish Error Too Large"))

          else
              ! helpers
              helperQuit = .false.
              do while(.not. helperQuit)
                  call helpVoltron(voltAppMpi, helperQuit)
              end do
          endif
      endif

  end subroutine

end module testHelpers

