module testCoupling
  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

      call setMpiReal()

  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

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

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

      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

  subroutine runApplication(tFin)
      real(rp), intent(in) :: tFin

      if(allocated(gamCplMpi)) then
          ! adjust tFin time to gamera units
          call gamCplMpi%AdvanceModel(tFin/gamCplMpi%Model%Units%gT0 - gamCplMpi%Model%t)
      else
          do while (voltAppMpi%time < tFin)
              call stepVoltron_mpi(voltAppMpi, tFin - voltAppMpi%time)
          end do
      endif

  end subroutine runApplication

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

      character(len=strLen) :: caseFile
      write(caseFile,'(A,I0,A)') 'cmiD_deep_', this%getNumProcesses()-1, '_res60.xml'
      call initializeWithCaseXml(this, caseFile)

      if(allocated(gamCplMpi)) then
          ! run case for 20 seconds past restart
          call runApplication(81.0_rp)
      else
          write (*,'(a,I0)') 'Testing Serial ',this%getNumProcesses()

          !adjust coupling parameters, must be MPI Gamera Coupler
          SELECT type(cpl=>voltAppMpi%gApp)
              TYPE IS (gamCouplerMpi_volt_T)
                  call endVoltronWaits(voltAppMpi)
                  voltAppMPi%doSerialMHD = .true.
                  cpl%doAsyncCoupling = .false.
              CLASS DEFAULT
                  @assertTrue(.false., "Voltron Allocated non-mpi Gamera coupler for MPI Voltron Coupling Test. Failure")
          ENDSELECT

          ! run case for 20 seconds past restart
          call runApplication(81.0_rp)
      endif

  end subroutine testSerialCoupling

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

      character(len=strLen) :: caseFile
      write(caseFile,'(A,I0,A)') 'cmiD_deep_', this%getNumProcesses()-1, '_res60.xml'
      call initializeWithCaseXml(this, caseFile)

      if(allocated(gamCplMpi)) then
          ! run case for 20 seconds past restart
          call runApplication(81.0_rp)
      else
          write (*,'(a,I0)') 'Testing Concurrent ',this%getNumProcesses()

          !adjust coupling parameters, must be MPI Gamera Coupler
          SELECT type(cpl=>voltAppMpi%gApp)
              TYPE IS (gamCouplerMpi_volt_T)
                  call endVoltronWaits(voltAppMpi)
                  voltAppMpi%doSerialMHD = .false.
                  cpl%doAsyncCoupling = .false.
              CLASS DEFAULT
                  @assertTrue(.false., "Voltron Allocated non-mpi Gamera coupler for MPI Voltron Coupling Test. Failure")
          ENDSELECT

          ! run case for 20 seconds past restart
          call runApplication(81.0_rp)
      endif

  end subroutine testConcCoupling

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

      character(len=strLen) :: caseFile
      write(caseFile,'(A,I0,A)') 'cmiD_deep_', this%getNumProcesses()-1, '_res60.xml'
      call initializeWithCaseXml(this, caseFile)

      if(allocated(gamCplMpi)) then
          ! run case for 20 seconds past restart
          call runApplication(81.0_rp)
      else
          write (*,'(a,I0)') 'Testing Async ',this%getNumProcesses()

          !adjust coupling parameters, must be MPI Gamera Coupler
          SELECT type(cpl=>voltAppMpi%gApp)
              TYPE IS (gamCouplerMpi_volt_T)
                  call endVoltronWaits(voltAppMpi)
                  voltAppMpi%doSerialMHD = .false.
                  cpl%doAsyncCoupling = .true.
              CLASS DEFAULT
                  @assertTrue(.false., "Voltron Allocated non-mpi Gamera coupler for MPI Voltron Coupling Test. Failure")
          ENDSELECT

          ! run case for 20 seconds past restart
          call runApplication(81.0_rp)
      endif

  end subroutine testAsyncCoupling

  !Temporary disabled
  !@test(npes=[9])
  subroutine testAdjustDeepCoupling(this)
      class (MpiTestMethod), intent(inout) :: this

      character(len=strLen) :: caseFile
      write(caseFile,'(A,I0,A)') 'cmiD_deep_', this%getNumProcesses()-1, '_res0.xml'
      call initializeWithCaseXml(this, caseFile)

      if(allocated(gamCplMpi)) then
          ! run case for 20 seconds past restart
          call runApplication(80.0_rp)
      else
          write (*,'(a,I0)') 'Testing DenDT Deep ',this%getNumProcesses()

          !adjust coupling parameters, must be MPI Gamera Coupler
          SELECT type(cpl=>voltAppMpi%gApp)
              TYPE IS (gamCouplerMpi_volt_T)
                  call endVoltronWaits(voltAppMpi)
                  voltAppMpi%doSerialMHD = .false.
                  cpl%doAsyncCoupling = .true.
              CLASS DEFAULT
                  @assertTrue(.false., "Voltron Allocated non-mpi Gamera coupler for MPI Voltron Coupling Test. Failure")
          ENDSELECT

          
          ! run case for 10 seconds past restart with low DeepDT
          voltAppMpi%DeepDT = 5.0_rp
          call runApplication(70.0_rp)
          @assertEqual(5.0_rp, voltAppMpi%DeepDT, trim("Voltron DeepDT not 5.0 after first couplings"))

          ! now run the case until 20 seconds past restart with high DeepDT
          call resetDeepCoupling(voltAppMpi, 10.0_rp)
          voltAppMpi%DeepDT = 10.0_rp
          call runApplication(80.0_rp)
          @assertEqual(10.0_rp, voltAppMpi%DeepDT, trim("Voltron DeepDT not 10.0 after second couplings"))

      endif

  end subroutine testAdjustDeepCoupling

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

      integer :: Nbase, Nx2, Nx4
      real(rp) :: x2err, x4err, Ndiv, approxShells
      character(len=strLen) :: caseFile
      write(caseFile,'(A,I0,A)') 'cmiD_deep_', this%getNumProcesses()-1, '_res0.xml'
      call initializeWithCaseXml(this, caseFile)


      if(allocated(gamCplMpi)) then
          ! run case for 10 seconds past spinup
          call runApplication(11.0_rp)
      else
          write (*,'(a,I0)') 'Testing Quick Squish ',this%getNumProcesses()

          ! run case for 10 seconds past spinup
          call runApplication(11.0_rp)

          ! now check the quick squish erroers
          call CheckQuickSquishError(voltAppMpi, voltAppMpi%gApp, Nbase, Nx2, Nx4, x2Err, x4Err)

          ! Ensure that the number of valid cells decreases as quick squish interpolation increases
          @assertTrue(Nx2 <= Nbase, trim("X2 Quick Squish has more valid cells than normal squish. This should not be possible."))
          @assertTrue(Nx4 <= Nbase, trim("X4 Quick Squish has more valid cells than normal squish. This should not be possible."))
          @assertTrue(Nx4 <= Nx2,   trim("X4 Quick Squish has more valid cells than X2 Quick Squish. This should not be possible."))

          ! Per cell error metric
          @assertLessThanOrEqual(x2Err, 5e-3_rp, trim("X2 Quick Squish Error per cell is too large"))
          @assertLessThanOrEqual(x4Err, 5e-2_rp, trim("X4 Quick Squish Error per cell is too large"))

          ! Approximate how many i layers of actual data there are if the data were spherical
          approxShells = Nbase
          approxShells = approxShells / (voltAppMpi%ebTrcApp%ebState%ebGr%ke+2-voltAppMpi%ebTrcApp%ebState%ebGr%ks)
          approxShells = approxShells / (voltAppMpi%ebTrcApp%ebState%ebGr%je+2-voltAppMpi%ebTrcApp%ebState%ebGr%js)
          ! If the data were spherical, squish x2 would reduce it by half of a layer on average. It's not spherical, so increase the error margin by 10
          Ndiv = Nbase-Nx2
          @assertLessThanOrEqual(Ndiv/Nbase, 5.0_rp/approxShells, trim("X2 Quick Squish has too many cells missing values compared to normal squish."))
          ! If the data were spherical, squish x4 would reduce it by 1.5 of a layer on average. It's not spherical, so increase the error margin by 10
          Ndiv = Nbase-Nx4
          @assertLessThanOrEqual(Ndiv/Nbase, 15.0_rp/approxShells, trim("X4 Quick Squish has too many cells missing values compared to normal squish."))

      endif

  end subroutine testQuickSquish

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

      integer :: Nbase, Nx2, Nx4
      real(rp) :: x2err, x4err, Ndiv, approxShells
      character(len=strLen) :: caseFile
      write(caseFile,'(A,I0,A)') 'cmiD_deep_', this%getNumProcesses()-1, '_res600.xml'
      call initializeWithCaseXml(this, caseFile)

      if(allocated(gamCplMpi)) then
          ! run case for 10 seconds past restart
          call runApplication(611.0_rp)
      else
          write (*,'(a,I0)') 'Testing Quick Squish Storm ',this%getNumProcesses()

          ! run case for 10 seconds past restart
          call runApplication(611.0_rp)

          ! now check the quick squish erroers
          call CheckQuickSquishError(voltAppMpi, voltAppMpi%gApp, Nbase, Nx2, Nx4, x2Err, x4Err)

          ! Ensure that the number of valid cells decreases as quick squish interpolation increases
          @assertTrue(Nx2 <= Nbase, trim("X2 Quick Squish has more valid cells than normal squish. This should not be possible."))
          @assertTrue(Nx4 <= Nbase, trim("X4 Quick Squish has more valid cells than normal squish. This should not be possible."))
          @assertTrue(Nx4 <= Nx2,   trim("X4 Quick Squish has more valid cells than X2 Quick Squish. This should not be possible."))

          ! Per cell error metric
          @assertLessThanOrEqual(x2Err, 5e-3_rp, trim("X2 Quick Squish Error per cell is too large"))
          @assertLessThanOrEqual(x4Err, 5e-2_rp, trim("X4 Quick Squish Error per cell is too large"))

          ! Approximate how many i layers of actual data there are if the data were spherical
          approxShells = Nbase
          approxShells = approxShells / (voltAppMpi%ebTrcApp%ebState%ebGr%ke+2-voltAppMpi%ebTrcApp%ebState%ebGr%ks)
          approxShells = approxShells / (voltAppMpi%ebTrcApp%ebState%ebGr%je+2-voltAppMpi%ebTrcApp%ebState%ebGr%js)
          ! If the data were spherical, squish x2 would reduce it by half of a layer on average. It's not spherical, so increase the error margin by 10
          Ndiv = Nbase-Nx2
          @assertLessThanOrEqual(Ndiv/Nbase, 5.0_rp/approxShells, trim("X2 Quick Squish has too many cells missing values compared to normal squish."))
          ! If the data were spherical, squish x4 would reduce it by 1.5 of a layer on average. It's not spherical, so increase the error margin by 10
          Ndiv = Nbase-Nx4
          @assertLessThanOrEqual(Ndiv/Nbase, 15.0_rp/approxShells, trim("X4 Quick Squish has too many cells missing values compared to normal squish."))

      endif

  end subroutine testQuickSquishStorm

end module testCoupling

