module testVMpi
  use testHelperMpi
  use voltapp_mpi
  use gamCouple_mpi_G2V
  use uservoltic
  use ioH5

  implicit none

  type(gamCouplerMpi_gam_T), allocatable :: gamCplMpi
  type(voltAppMpi_T), allocatable :: voltAppMpi

contains

  @before
  subroutine setup(this)
      class (MpiTestMethod), intent(inout) :: this
      character(len=strLen) :: caseFile

      integer :: ierror
      type(MPI_Comm) :: voltComm
      type(XML_Input_T) :: xmlInp

      call setMpiReal()

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


      if(this%getProcessRank() < (this%getNumProcesses()-1)) then
          allocate(gamCplMpi)

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

	  gamCplMpi%gOptionsCplMpiG%couplingPoolComm = getMpiF08Communicator(this)
          gamCplMpi%gOptions%userInitFunc => initUser
          gamCplMpi%gOptionsMpi%doIO = .false.
          xmlInp = New_XML_Input(trim(caseFile),'Kaiju',.true.)
          call gamCplMpi%InitModel(xmlInp)
      else
          allocate(voltAppMpi)

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

          voltAppMpi%vOptions%gamUserInitFunc => initUser
	  allocate(voltAppMpi%vOptionsMpi%couplingPoolComm)
          voltAppMpi%vOptionsMpi%couplingPoolComm = getMpiF08Communicator(this)
          call initVoltron_mpi(voltAppMpi, trim(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(gamCplMpi)) deallocate(gamCplMpi)

  end subroutine teardown

  @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(allocated(gamCplMpi)) then
          if(gamCplMpi%Grid%hasLowerBC(IDIR)) then
              SELECT type(iiBC=>gamCplMpi%Grid%externalBCs(INI)%p)
                  TYPE IS (IonInnerBC_T)

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

                      ! receive the data from voltron
                      call recvShallowCplDataMpi(gamCplMpi)

                      ! check inExyz values
                      do i=1,PsiSh
                          do j=gamCplMpi%Grid%jsg,gamCplMpi%Grid%jeg
                              do k=gamCplMpi%Grid%ksg,gamCplMpi%Grid%keg
                                  do l=1,NDIM
                                      checkVal = lVal*l + &
                                          kVal*(k+gamCplMpi%Grid%ijkShift(KDIR)) + &
                                          jVal*(j+gamCplMpi%Grid%ijkShift(JDIR)) + &
                                          iVal*(i+gamCplMpi%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,PsiSh+1
                          do j=gamCplMpi%Grid%jsg,gamCplMpi%Grid%jeg+1
                              do k=gamCplMpi%Grid%ksg,gamCplMpi%Grid%keg+1
                                  do l=1,NDIM
                                      checkVal = lVal*l + &
                                          kVal*(k+gamCplMpi%Grid%ijkShift(KDIR)) + &
                                          jVal*(j+gamCplMpi%Grid%ijkShift(JDIR)) + &
                                          iVal*(i+gamCplMpi%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 recvShallowCplDataMpi(gamCplMpi)
          endif
      else
          SELECT type(iiBC=>voltAppMpi%gApp%Grid%externalBCs(INI)%p)
              TYPE IS (IonInnerBC_T)
                  ! clear all cells
                  iiBC%inExyz(:,:,:,:) = 0
                  iiBC%inEijk(:,:,:,:) = 0

                  ! set inExyz
                  do i=1,PsiSh
                      do j=voltAppMpi%gApp%Grid%jsg,voltAppMpi%gApp%Grid%jeg
                          do k=voltAppMpi%gApp%Grid%ksg,voltAppMpi%gApp%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,PsiSh+1
                      do j=voltAppMpi%gApp%Grid%jsg,voltAppMpi%gApp%Grid%jeg+1
                          do k=voltAppMpi%gApp%Grid%ksg,voltAppMpi%gApp%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

              CLASS DEFAULT
                  @assertEqual(0,1,'No Remix BC on Voltron')
          END SELECT

          SELECT type(cpl=>voltAppMpi%gApp)
              TYPE IS (gamCouplerMpi_volt_T)
                  ! send the data to gamera
                  call sendShallowCplDataMpi(cpl)

              CLASS DEFAULT
                  @assertTrue(.false., "Voltron Allocated non-mpi Gamera coupler for MPI Voltron Coupling Test. Failure")
          ENDSELECT

      endif

  end subroutine testVolt2MhdShallowCopy

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

      real(rp) :: checkVal
      character(len=strLen) :: checkMessage

      checkVal = 123456.0_rp ! random value for testing with

      if(allocated(gamCplMpi)) then
          call recvCplTimeMpi(gamCplMpi)
          @assertEqual(checkVal/gamCplMpi%Model%Units%gT0, gamCplMpi%DeepT, "Time value transmitted from Voltron to Gamera was wrong")
      else
          SELECT type(cpl=>voltAppMpi%gApp)
              TYPE IS (gamCouplerMpi_volt_T)
                  ! send the data to gamera
                  call sendCplTimeMpi(cpl, checkVal)
              CLASS DEFAULT
                  @assertTrue(.false., "Voltron Allocated non-mpi Gamera coupler for MPI Voltron Coupling Test. Failure")
          ENDSELECT
      endif

  end subroutine testVolt2MhdTimeCopy

  @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(allocated(gamCplMpi)) then
          if(gamCplMpi%Grid%hasLowerBC(IDIR)) then
              SELECT type(iiBC=>gamCplMpi%Grid%externalBCs(INI)%p)
                  TYPE IS (IonInnerBC_T)

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

                      ! receive the data from voltron
                      call recvShallowCplDataMpi(gamCplMpi)

                      ! 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=gamCplMpi%Grid%ks,gamCplMpi%Grid%ke
                          do j=gamCplMpi%Grid%js,gamCplMpi%Grid%je

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

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

                                      call cellCenter(gamCplMpi%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+gamCplMpi%Grid%ke-gamCplMpi%Grid%ks)*(1+gamCplMpi%Grid%je-gamCplMpi%Grid%js)*gamCplMpi%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=gamCplMpi%Grid%ks,gamCplMpi%Grid%ke
                          do j=gamCplMpi%Grid%js,gamCplMpi%Grid%je
                              !Map to active ip,jp,kp (i=Grid%is => ip=Grid%is)
                              call lfmIJKcc(gamCplMpi%Model,gamCplMpi%Grid,gamCplMpi%Grid%is,j,k,ip,jp,kp)
                              do n=1,gamCplMpi%Model%Ng
                                  ig = gamCplMpi%Grid%is-n
                                  np = gamCplMpi%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(gamCplMpi%Model,gamCplMpi%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+gamCplMpi%Grid%ke-gamCplMpi%Grid%ks)*(1+gamCplMpi%Grid%je-gamCplMpi%Grid%js)*gamCplMpi%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 recvShallowCplDataMpi(gamCplMpi)
          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

          SELECT type(cpl=>voltAppMpi%gApp)
              TYPE IS (gamCouplerMpi_volt_T)
                  !Can add a test of the interpolation to Gamera shells here?
                  call mapRemixToGamera(cpl, voltAppMpi%remixApp)

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

                  ! send the data to gamera
                  call sendShallowCplDataMpi(cpl)
              CLASS DEFAULT
                  @assertTrue(.false., "Voltron Allocated non-mpi Gamera coupler for MPI Voltron Coupling Test. Failure")
          ENDSELECT

          SELECT type(iiBC=>voltAppMpi%gApp%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%gApp%Grid%ks,voltAppMpi%gApp%Grid%ke
                      do j=voltAppMpi%gApp%Grid%js,voltAppMpi%gApp%Grid%je

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

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

                              call cellCenter(voltAppMpi%gApp%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%gApp%Grid%ke-voltAppMpi%gApp%Grid%ks)*(1+voltAppMpi%gApp%Grid%je-voltAppMpi%gApp%Grid%js)*voltAppMpi%gApp%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%gApp%Grid%ks,voltAppMpi%gApp%Grid%ke
                      do j=voltAppMpi%gApp%Grid%js,voltAppMpi%gApp%Grid%je
                          !Map to active ip,jp,kp (i=Grid%is => ip=Grid%is)
                          call lfmIJKcc(voltAppMpi%gApp%Model,voltAppMpi%gApp%Grid,voltAppMpi%gApp%Grid%is,j,k,ip,jp,kp)
                          do n=1,voltAppMpi%gApp%Model%Ng
                              ig = voltAppMpi%gApp%Grid%is-n
                              np = voltAppMpi%gApp%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%gApp%Model,voltAppMpi%gApp%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%gApp%Grid%ke-voltAppMpi%gApp%Grid%ks)*(1+voltAppMpi%gApp%Grid%je-voltAppMpi%gApp%Grid%js)*voltAppMpi%gApp%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(allocated(gamCplMpi)) then

          if(gamCplMpi%Grid%hasLowerBC(IDIR)) then
              SELECT type(iiBC=>gamCplMpi%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

          call gamCplMpi%AdvanceModel(30.0/gamCplMpi%Model%Units%gT0 - gamCplMpi%Model%t)

          if(gamCplMpi%Grid%hasLowerBC(IDIR)) then
              SELECT type(iiBC=>gamCplMpi%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)
              call stepVoltron_mpi(voltAppMpi, 30.0_rp - voltAppMpi%time)
          end do
      endif

  end subroutine remixEnabledWithoutDeepTest

end module testVMpi

