module testVMpi
  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 setup(this)
      class (MpiTestMethod), intent(inout) :: this
      character(len=strLen) :: caseFile

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

      call setMpiReal()

      write(caseFile,'(A,I0,A)') 'cmiD_', this%getNumProcesses()-1, '.xml'


      if(this%getProcessRank() < (this%getNumProcesses()-1)) 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, caseFile, .false.)
          call initGam2Volt(g2vCommMpi, gamAppMpi, getMpiF08Communicator(this), 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), caseFile)
      endif

  end subroutine setup

  @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

  @test(npes=[2,5,9])
  subroutine testStepUpdate(this)
      class (MpiTestMethod), intent(inout) :: this

      character(len=strLen) :: checkMessage

      if(isGamera) then
          ! setting fixed values to gamera to ensure updates occur
          gamAppMpi%Model%t = 1.23
          gamAppMpi%Model%ts = 178
          g2vCommMpi%time = 0
          g2vCommMpi%MJD = 0
          g2vCommMpi%ts = 0

          call localStepVoltronTime(g2vCommMpi, gamAppMpi)
          call performStepVoltron(g2vCommMpi,gamAppMpi)

          write (checkMessage,'(A,I0)') 'voltron time wrong on gamera rank ',this%getProcessRank()
          @assertEqual(1.23*gamAppMpi%Model%Units%gT0, g2vCommMpi%time, trim(checkMessage))
          write (checkMessage,'(A,I0)') 'MJD wrong on gamera rank ',this%getProcessRank()
          @assertLessThanOrEqual(abs(gamAppMpi%Model%MJD0 + g2vCommMpi%time/(60.0*60.0*24.0) - g2vCommMpi%MJD), 1e-5_rp, trim(checkMessage))
          write (checkMessage,'(A,I0)') 'voltron ts wrong on gamera rank ',this%getProcessRank()
          @assertEqual(178, g2vCommMpi%ts, trim(checkMessage))
      else
          ! setting fixed values to voltron to ensure updates occur
          voltAppMpi%gAppLocal%Model%t = 0
          voltAppMpi%gAppLocal%Model%ts = 0
          voltAppMpi%time = 0
          voltAppMpi%MJD = 0
          voltAppMpi%ts = 0

          call stepVoltron_mpi(voltAppMpi)

          @assertEqual(1.23, voltAppMpi%gAppLocal%Model%t, 'gamera time value wrong on voltron')
          @assertEqual(178, voltAppMpi%gAppLocal%Model%ts, 'gamera ts value wrong on voltron')
          @assertEqual(1.23*voltAppMpi%gAppLocal%Model%Units%gT0, voltAppMpi%time, 'voltron time value wrong on voltron')
          @assertLessThanOrEqual(abs(voltAppMpi%gAppLocal%Model%MJD0 + voltAppMpi%time/(60.0*60.0*24.0) - voltAppMpi%MJD), 1e-5_rp, 'MJD value wrong on voltron')
          @assertEqual(178, voltAppMpi%ts, 'voltron ts value wrong on voltron')
      endif

  end subroutine testStepUpdate

  @test(npes=[2,5,9])
  subroutine testVolt2MhdShallowCopy(this)
      class (MpiTestMethod), intent(inout) :: this

      integer :: i,j,k,l,jIndexStart,jIndexEnd
      real(rp) :: checkVal
      character(len=strLen) :: checkMessage
      real(rp) :: iVal, jVal, kVal, lVal

      ! coefficients chosen to allow easy interpretation of final values
      iVal = 1
      jVal = 20
      kVal = 400
      lVal = 8000

      if(isGamera) then
          if(gamAppMpi%Grid%hasLowerBC(IDIR)) then
              SELECT type(iiBC=>gamAppMpi%Grid%externalBCs(INI)%p)
                  TYPE IS (IonInnerBC_T)

                      ! clear all cells
                      iiBC%inEijk(:,:,:,:) = 0
                      iiBC%inExyz(:,:,:,:) = 0

                      ! receive the data from voltron
                      call recvShallowData(g2vCommMpi, gamAppMpi)

                      ! check inExyz values
                      do i=1,g2vCommMpi%PsiSh
                          do j=gamAppMpi%Grid%jsg,gamAppMpi%Grid%jeg
                              do k=gamAppMpi%Grid%ksg,gamAppMpi%Grid%keg
                                  do l=1,NDIM
                                      checkVal = lVal*l + &
                                          kVal*(k+gamAppMpi%Grid%ijkShift(KDIR)) + &
                                          jVal*(j+gamAppMpi%Grid%ijkShift(JDIR)) + &
                                          iVal*(i+gamAppMpi%Grid%ijkShift(IDIR))
                                      write (checkMessage,'(A,I0,A,I0,A,I0,A,I0,A)') 'gamera inExyz wrong at (',i,',',j,',',k,',',l,')'
                                      @assertEqual(checkVal, iiBC%inExyz(i,j,k,l), trim(checkMessage))
                                  enddo
                              enddo
                          enddo
                      enddo

                      ! check inEijk values
                      do i=1,g2vCommMpi%PsiSh+1
                          do j=gamAppMpi%Grid%jsg,gamAppMpi%Grid%jeg+1
                              do k=gamAppMpi%Grid%ksg,gamAppMpi%Grid%keg+1
                                  do l=1,NDIM
                                      checkVal = lVal*l + &
                                          kVal*(k+gamAppMpi%Grid%ijkShift(KDIR)) + &
                                          jVal*(j+gamAppMpi%Grid%ijkShift(JDIR)) + &
                                          iVal*(i+gamAppMpi%Grid%ijkShift(IDIR))
                                      write (checkMessage,'(A,I0,A,I0,A,I0,A,I0,A)') 'gamera inEijk wrong at (',i,',',j,',',k,',',l,')'
                                      @assertEqual(checkVal, iiBC%inEijk(i,j,k,l), trim(checkMessage))
                                  enddo
                              enddo
                          enddo
                      enddo
                  CLASS DEFAULT
                      @assertEqual(0,1,'No Remix BC on Gamera')
              END SELECT
          else
              ! everyone has to call to receive the data from voltron
              call recvShallowData(g2vCommMpi, gamAppMpi)
          endif
      else
          SELECT type(iiBC=>voltAppMpi%gAppLocal%Grid%externalBCs(INI)%p)
              TYPE IS (IonInnerBC_T)
                  ! clear all cells
                  iiBC%inExyz(:,:,:,:) = 0
                  iiBC%inEijk(:,:,:,:) = 0

                  ! set inExyz
                  do i=1,voltAppMpi%mix2mhd%PsiShells
                      do j=voltAppMpi%gAppLocal%Grid%jsg,voltAppMpi%gAppLocal%Grid%jeg
                          do k=voltAppMpi%gAppLocal%Grid%ksg,voltAppMpi%gAppLocal%Grid%keg
                              do l=1,NDIM
                                  iiBC%inExyz(i,j,k,l) = lVal*l + &
                                      iVal*i + &
                                      jVal*j + &
                                      kVal*k
                              enddo
                          enddo
                      enddo
                  enddo

                  ! set inEijk
                  do i=1,voltAppMpi%mix2mhd%PsiShells+1
                      do j=voltAppMpi%gAppLocal%Grid%jsg,voltAppMpi%gAppLocal%Grid%jeg+1
                          do k=voltAppMpi%gAppLocal%Grid%ksg,voltAppMpi%gAppLocal%Grid%keg+1
                              do l=1,NDIM
                                  iiBC%inEijk(i,j,k,l) = lVal*l + &
                                      iVal*i + &
                                      jVal*j + &
                                      kVal*k
                              enddo
                          enddo
                      enddo
                  enddo

                  ! send the data to gamera
                  call sendShallowData_mpi(voltAppMpi)
              CLASS DEFAULT
                  @assertEqual(0,1,'No Remix BC on Voltron')
          END SELECT
      endif

  end subroutine testVolt2MhdShallowCopy

  @test(npes=[2,5,9])
  subroutine VMPIcorotationTest(this)
      class (MpiTestMethod), intent(inout) :: this

      integer :: j,k,n,ip,it,jp,kp,h,ig,np
      real(rp) :: xc,yc,zc,deltaSum,rval,lambda,Psi0,RIn,errorThresh
      real(rp), dimension(NDIM) :: rHatP,eHat,xyz,Exyz,ccEijk

      !type(IOVAR_T), dimension(MAXMIXIOVAR) :: debugHDF

      !Setting some values for inner shell
      RIn  = 1.0
      Psi0 = 10.0

      if(isGamera) then
          if(gamAppMpi%Grid%hasLowerBC(IDIR)) then
              SELECT type(iiBC=>gamAppMpi%Grid%externalBCs(INI)%p)
                  TYPE IS (IonInnerBC_T)

                      ! clear all cells
                      iiBC%inEijk(:,:,:,:) = 0
                      iiBC%inExyz(:,:,:,:) = 0

                      ! receive the data from voltron
                      call recvShallowData(g2vCommMpi, gamAppMpi)

                      ! verify that the gamera field (iiBC%inEijk, iiBC%inExyz, or both) have the expected values here
                      ! inExyz should be pointed exactly radially at all locations
                      deltaSum = 0.0_rp
                      !Only loop over active cells
                      do k=gamAppMpi%Grid%ks,gamAppMpi%Grid%ke
                          do j=gamAppMpi%Grid%js,gamAppMpi%Grid%je

                              !Map to active ip,jp,kp (i=Grid%is => ip=Grid%is)
                              call lfmIJKcc(gamAppMpi%Model,gamAppMpi%Grid,gamAppMpi%Grid%is,j,k,ip,jp,kp)

                              do n=1,gamAppMpi%Model%Ng
                                      ig = gamAppMpi%Grid%is-n
                                      np = gamAppMpi%Model%nG-n+1

                                      call cellCenter(gamAppMpi%Grid, ig,j,k,xc,yc,zc)
                                      rHatP = normVec([xc,yc,zc])

                                      eHat = normVec(iiBC%inExyz(np,jp,kp,:))
                                      !write(*,*) 'ig,j,k,np = ',ig,j,k,np
                                      !write(*,*) 'eIJK / eXYZ = ', iiBC%inEijk(np,jp,kp,:), iiBC%inExyz(np,jp,kp,:)

                                      deltaSum = deltaSum + abs(norm2(rHatP-eHat))
                              end do
                          end do
                      end do

                       ! remix doesn't go below a certain lat, causing discontinuity. high error
                       errorThresh = 2*(1+gamAppMpi%Grid%ke-gamAppMpi%Grid%ks)*(1+gamAppMpi%Grid%je-gamAppMpi%Grid%js)*gamAppMpi%Model%Ng
                      @assertLessThanOrEqual(deltaSum, errorThresh, 'Magnetic field from remix was not pointing radially')

                      !Now test Eijk fields
                      deltaSum = 0.0
                      !Only loop over active cells
                      do k=gamAppMpi%Grid%ks,gamAppMpi%Grid%ke
                          do j=gamAppMpi%Grid%js,gamAppMpi%Grid%je
                              !Map to active ip,jp,kp (i=Grid%is => ip=Grid%is)
                              call lfmIJKcc(gamAppMpi%Model,gamAppMpi%Grid,gamAppMpi%Grid%is,j,k,ip,jp,kp)
                              do n=1,gamAppMpi%Model%Ng
                                  ig = gamAppMpi%Grid%is-n
                                  np = gamAppMpi%Model%nG-n+1
                                  !Convert Eijk (edges) to cc-Eijk
                                  ccEijk(IDIR) = 0.25*(iiBC%inEijk(np,j,k,IDIR)+iiBC%inEijk(np,j,k+1,IDIR)+iiBC%inEijk(np,j+1,k,IDIR)+iiBC%inEijk(np,j+1,k+1,IDIR))
                                  ccEijk(JDIR) = 0.25*(iiBC%inEijk(np,j,k,JDIR)+iiBC%inEijk(np,j,k+1,JDIR)+iiBC%inEijk(np+1,j,k,JDIR)+iiBC%inEijk(np+1,j,k+1,JDIR))
                                  ccEijk(KDIR) = 0.25*(iiBC%inEijk(np,j,k,KDIR)+iiBC%inEijk(np+1,j,k,KDIR)+iiBC%inEijk(np,j+1,k,KDIR)+iiBC%inEijk(np+1,j+1,k,KDIR))
                                  !Now convert cc-Eijk to Exyz
                                  Exyz = ccEijk2Exyz(gamAppMpi%Model,gamAppMpi%Grid,ccEijk,ig,j,k)
                                  !Exyz above is approximation to inExyz, ie testing whether or not Eijk and Exyz are consistent
                                  deltaSum = deltaSum + norm2(Exyz-iiBC%inExyz(np,jp,kp,:))
                              enddo
                          enddo
                    enddo

                    ! remix doesn't go below a certain lat, causing discontinuity. high error
                    errorThresh = 2.5e-2*(1+gamAppMpi%Grid%ke-gamAppMpi%Grid%ks)*(1+gamAppMpi%Grid%je-gamAppMpi%Grid%js)*gamAppMpi%Model%Ng
                    @assertLessThanOrEqual(deltaSum, errorThresh, 'inEijk values did not match inExyz values')

                  CLASS DEFAULT
                      @assertEqual(0,1,'No Remix BC on Gamera')
              END SELECT
          else
              ! everyone has to call to receive the data from voltron
              call recvShallowData(g2vCommMpi, gamAppMpi)
          endif
      else
          associate(rApp=>voltAppMpi%remixApp)
          do h=1,2
              do ip=1,rApp%ion(h)%g%Np
                  do it=1,rApp%ion(h)%g%Nt
                      rval = rApp%ion(h)%g%r(ip,it)
                      lambda = (PI/2.0_rp) - rApp%ion(h)%g%t(ip,it)
                      !write(*,*) 'h / ip / it, r/lam = ',h,ip,it,rval,lambda
                      rApp%ion(h)%st%vars(ip,it,POT) = Psi0*cos(lambda)*cos(lambda)/RIn
                  enddo
              enddo
          enddo
          end associate

          !Can add a test of the interpolation to Gamera shells here?
          call mapRemixToGamera(voltAppMpi%mix2mhd, voltAppMpi%remixApp)

          call convertRemixToGamera(voltAppMpi%mix2mhd, voltAppMpi%remixApp, voltAppMpi%gAppLocal, .false.) !Don't add corotation potential          

          ! send the data to gamera
          call sendShallowData_mpi(voltAppMpi)

          SELECT type(iiBC=>voltAppMpi%gAppLocal%Grid%externalBCs(INI)%p)
              TYPE IS (IonInnerBC_T)

                  !call CheckAndKill('voltrontest.h5')
                  !call AddOutVar(debugHDF,'remixN',voltAppMpi%remixApp%ion(1)%st%vars(:,:,POT))
                  !call AddOutVar(debugHDF,'remixS',voltAppMpi%remixApp%ion(2)%st%vars(:,:,POT))
                  !call AddOutVar(debugHDF,'mixOutput',voltAppMpi%mix2mhd%mixOutput)
                  !call AddOutVar(debugHDF,'gPsi',voltAppMpi%mix2mhd%gPsi)
                  !call AddOutVar(debugHDF,'inExyz',iiBC%inExyz)
                  !call AddOutVar(debugHDF,'inEijk',iiBC%inEijk)
                  !call WriteVars(debugHDF,.false.,'voltrontest.h5')

                  ! check that inExyz is pointed exactly radially at all locations
                  deltaSum = 0.0_rp
                  !Only loop over active cells
                  do k=voltAppMpi%gAppLocal%Grid%ks,voltAppMpi%gAppLocal%Grid%ke
                      do j=voltAppMpi%gAppLocal%Grid%js,voltAppMpi%gAppLocal%Grid%je

                          !Map to active ip,jp,kp (i=Grid%is => ip=Grid%is)
                          call lfmIJKcc(voltAppMpi%gAppLocal%Model,voltAppMpi%gAppLocal%Grid,voltAppMpi%gAppLocal%Grid%is,j,k,ip,jp,kp)

                          do n=1,voltAppMpi%gAppLocal%Model%Ng
                              ig = voltAppMpi%gAppLocal%Grid%is-n
                              np = voltAppMpi%gAppLocal%Model%nG-n+1

                              call cellCenter(voltAppMpi%gAppLocal%Grid, ig,j,k,xc,yc,zc)
                              rHatP = normVec([xc,yc,zc])

                              eHat = normVec(iiBC%inExyz(np,jp,kp,:))
                              !write(*,*) 'ig,j,k,np = ',ig,j,k,np
                              !write(*,*) 'eIJK / eXYZ = ', iiBC%inEijk(np,jp,kp,:), iiBC%inExyz(np,jp,kp,:)

                              deltaSum = deltaSum + abs(norm2(rHatP-eHat))
                          end do
                      end do
                  end do

                  ! remix doesn't go below a certain lat, causing discontinuity. high error
                  errorThresh = 2*(1+voltAppMpi%gAppLocal%Grid%ke-voltAppMpi%gAppLocal%Grid%ks)*(1+voltAppMpi%gAppLocal%Grid%je-voltAppMpi%gAppLocal%Grid%js)*voltAppMpi%gAppLocal%Model%Ng
                  @assertLessThanOrEqual(deltaSum, errorThresh, 'VOLTRON Magnetic field from remix was not pointing radially')

                  !Now test Eijk fields
                  deltaSum = 0.0
                  !Only loop over active cells
                  do k=voltAppMpi%gAppLocal%Grid%ks,voltAppMpi%gAppLocal%Grid%ke
                      do j=voltAppMpi%gAppLocal%Grid%js,voltAppMpi%gAppLocal%Grid%je
                          !Map to active ip,jp,kp (i=Grid%is => ip=Grid%is)
                          call lfmIJKcc(voltAppMpi%gAppLocal%Model,voltAppMpi%gAppLocal%Grid,voltAppMpi%gAppLocal%Grid%is,j,k,ip,jp,kp)
                          do n=1,voltAppMpi%gAppLocal%Model%Ng
                              ig = voltAppMpi%gAppLocal%Grid%is-n
                              np = voltAppMpi%gAppLocal%Model%nG-n+1
                              !Convert Eijk (edges) to cc-Eijk
                              ccEijk(IDIR) = 0.25*(iiBC%inEijk(np,j,k,IDIR)+iiBC%inEijk(np,j,k+1,IDIR)+iiBC%inEijk(np,j+1,k,IDIR)+iiBC%inEijk(np,j+1,k+1,IDIR))
                              ccEijk(JDIR) = 0.25*(iiBC%inEijk(np,j,k,JDIR)+iiBC%inEijk(np,j,k+1,JDIR)+iiBC%inEijk(np+1,j,k,JDIR)+iiBC%inEijk(np+1,j,k+1,JDIR))
                              ccEijk(KDIR) = 0.25*(iiBC%inEijk(np,j,k,KDIR)+iiBC%inEijk(np+1,j,k,KDIR)+iiBC%inEijk(np,j+1,k,KDIR)+iiBC%inEijk(np+1,j+1,k,KDIR))
                              !Now convert cc-Eijk to Exyz
                              Exyz = ccEijk2Exyz(voltAppMpi%gAppLocal%Model,voltAppMpi%gAppLocal%Grid,ccEijk,ig,j,k)
                              !Exyz above is approximation to inExyz, ie testing whether or not Eijk and Exyz are consistent
                              deltaSum = deltaSum + norm2(Exyz-iiBC%inExyz(np,jp,kp,:))
                          enddo
                      enddo
                  enddo

                  ! remix doesn't go below a certain lat, causing discontinuity. high error
                  errorThresh = 2.5e-2*(1+voltAppMpi%gAppLocal%Grid%ke-voltAppMpi%gAppLocal%Grid%ks)*(1+voltAppMpi%gAppLocal%Grid%je-voltAppMpi%gAppLocal%Grid%js)*voltAppMpi%gAppLocal%Model%Ng
                  @assertLessThanOrEqual(deltaSum, errorThresh, 'VOLTRON inEijk values did not match inExyz values')
              CLASS DEFAULT
                  @assertEqual(0,1,'No Remix BC on Voltron')
          END SELECT

      endif
  end subroutine VMPIcorotationTest

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

      real(rp) :: eijkSum, exyzSum

      if(isGamera) then

          if(gamAppMpi%Grid%hasLowerBC(IDIR)) then
              SELECT type(iiBC=>gamAppMpi%Grid%externalBCs(INI)%p)
                  TYPE IS (IonInnerBC_T)

                      ! confirm cells are non-zero now
                      eijkSum = sum(abs(iiBC%inEijk))
                      exyzSum = sum(abs(iiBC%inExyz))
                      @assertGreaterThan(eijkSum, 0.0_rp, 'inEijk is all zeroes at start')
                      @assertGreaterThan(exyzSum, 0.0_rp, 'inExyz is all zeroes at start')
              END SELECT
          endif

          do while (g2vCommMpi%time < 30.0)
              call stepGamera_mpi(gamAppMpi)
              call localStepVoltronTime(g2vCommMpi, gamAppMpi)
              if( (g2vCommMpi%time >= g2vCommMpi%DeepT) .or. &
                (g2vCommMpi%time >= 30.0)) then
                  call performStepVoltron(g2vCommMpi,gamAppMpi)
                  if ( g2vCommMpi%time >= g2vCommMpi%DeepT) then
                      call performDeepUpdate(g2vCommMpi, gamAppMpi)
                  endif
              endif
          end do

          if(gamAppMpi%Grid%hasLowerBC(IDIR)) then
              SELECT type(iiBC=>gamAppMpi%Grid%externalBCs(INI)%p)
                  TYPE IS (IonInnerBC_T)

                      ! confirm cells have changed
                      @assertGreaterThan(abs(sum(abs(iiBC%inEijk)) - eijkSum), 0.0_rp, 'inEijk has not changed')
                      @assertGreaterThan(abs(sum(abs(iiBC%inExyz)) - exyzSum), 0.0_rp, 'inExyz has not changed')
               END SELECT
          endif
      else
          do while (voltAppMpi%time < 30.0)
              if(gameraStepReady(voltAppMpi)) then
                  call stepVoltron_mpi(voltAppMpi)
                  write (*,'(a,f8.3)') '    Time = ',voltAppMpi%time
                  if (voltAppMpi%time >= voltAppMpi%DeepT) then
                      call DeepUpdate_mpi(voltAppMpi, voltAppMpi%time)
                  endif
              elseif(deepInProgress(voltAppMpi)) then
                  call doDeepBlock(voltAppMpi)
              else
                  call waitForGameraStep(voltAppMpi)
              endif
          end do
      endif

  end subroutine remixEnabledWithoutDeepTest

end module testVMpi

