!Driver for mpi decomposed Gamera coupled with Voltron (remix only)

program voltron_mpix
    use clocks
    use gamapp_mpi
    use voltapp_mpi
    use gam2VoltComm_mpi
    use output
    use voltio
    use uservoltic
    use mpi_f08
    use xml_input

    implicit none

    ! this driver requires that each rank be either Gamera OR Voltron, never both
    logical :: isGamera = .true.

    ! gamera data
    type(gamAppMpi_T) :: gApp
    type(gam2VoltCommMpi_T) :: g2vComm

    ! voltron data
    type(voltAppMpi_T) :: vApp

    procedure(StateIC_T), pointer :: userInitFunc => initUser

    integer :: ierror, length, provided, worldSize, worldRank, numHelpers
    type(MPI_Comm) :: gamComm, voltComm
    integer :: required=MPI_THREAD_MULTIPLE
    character( len = MPI_MAX_ERROR_STRING) :: message
    character(len=strLen) :: inpXML, helpersBuf
    logical :: useHelpers, helperQuit
    integer(KIND=MPI_AN_MYADDR) :: tagMax
    logical :: tagSet
    integer :: divideSize,i

    ! initialize MPI
    !Set up MPI with or without thread support
#ifdef _OPENMP
    call MPI_INIT_THREAD(required,provided,ierror)
    if(provided < required) then
        print *,"Not support for MPI_THREAD_MULTIPLE, aborting!"
        call abort()
    end if
    !print *,"MPI + OpenMP !!!"
#else
    !print *," MPI without threading"
    call MPI_INIT(ierror)
#endif

    if(ierror /= MPI_Success) then
        call MPI_Error_string( ierror, message, length, ierror)
        print *,message(1:length)
        call mpi_Abort(MPI_COMM_WORLD, 1, ierror)
    end if

    ! initialize mpi data type
    call setMpiReal()

    call initClocks()

    call mpi_comm_get_attr(MPI_COMM_WORLD, MPI_TAG_UB, tagMax, tagSet, ierror)
    !print *, 'Tag Upper-Bound = ', tagMax

    gApp%Model%isLoud = .true.    

    ! need to know how many voltron helpers there are
    call getIDeckStr(inpXML)
    call ReadXmlImmediate(trim(inpXML),'/Kaiju/Voltron/Helpers/useHelpers',helpersBuf,'F',.false.)
    read(helpersBuf,*) useHelpers
    call ReadXmlImmediate(trim(inpXML),'/Kaiju/Voltron/Helpers/numHelpers',helpersBuf,'0',.false.)
    read(helpersBuf,*) numHelpers
    if(.not. useHelpers) numHelpers = 0
    call ReadXmlImmediate(trim(inpXML),'/Kaiju/Voltron/coupling/doGCM',helpersBuf,'F',.false.)
    read(helpersBuf,*) vApp%doGCM

    ! create a new MPI communicator for just Gamera
    !    for now this is always all ranks excep the last one (which is reserved for voltron)
    call MPI_Comm_Size(MPI_COMM_WORLD, worldSize, ierror)
    if(ierror /= MPI_Success) then
        call MPI_Error_string( ierror, message, length, ierror)
        print *,message(1:length)
        call mpi_Abort(MPI_COMM_WORLD, 1, ierror)
    end if
    call MPI_Comm_Rank(MPI_COMM_WORLD, worldRank, ierror)
    if(ierror /= MPI_Success) then
        call MPI_Error_string( ierror, message, length, ierror)
        print *,message(1:length)
        call mpi_Abort(MPI_COMM_WORLD, 1, ierror)
    end if
    if(worldRank .ge. (worldSize-1-numHelpers)) then
        ! voltron rank
        isGamera = .false.
        call MPI_Comm_Split(MPI_COMM_WORLD, voltId, worldRank, voltComm, ierror)
        if(ierror /= MPI_Success) then
            call MPI_Error_string( ierror, message, length, ierror)
            print *,message(1:length)
            call mpi_Abort(MPI_COMM_WORLD, 1, ierror)
        end if
    else
        ! gamera rank
        isGamera = .true.
        call MPI_Comm_Split(MPI_COMM_WORLD, gamId, worldRank, gamComm, ierror)
        if(ierror /= MPI_Success) then
            call MPI_Error_string( ierror, message, length, ierror)
            print *,message(1:length)
            call mpi_Abort(MPI_COMM_WORLD, 1, ierror)
        end if
    endif

    if (worldRank .eq. worldSize-1-numHelpers) then
        call MPI_Comm_Split(MPI_COMM_WORLD, mageId, 0, vApp%mageCplComm, ierror)
        call MPI_Comm_Rank(vApp%mageCplComm, vApp%voltCplRank, ierror)
        call MPI_Comm_Size(vApp%mageCplComm, vApp%CplSize, ierror)
      if(vApp%mageCplComm /= MPI_COMM_NULL) then
        print *,'VOLTRON has created an mageCplComm with ', vApp%CplSize-1, ' other app(s)'
        print *,'VOLTRON using mageCplComm tag ', mageId
        write(*,*) "VOLTRON CPLCOMM: ",vApp%mageCplComm,vApp%voltCplRank

        ! Tell everyone who I am
        if (.not.allocated(vApp%IAm)) allocate(vApp%IAm(vApp%CplSize))
        vApp%IAm(vApp%voltCplRank+1) = voltId
        do i=1,vApp%CplSize
          call MPI_Bcast(vApp%IAm(i),1,MPI_INTEGER,i-1,vApp%mageCplComm,ierror)
        enddo

        ! IAm array starts at 1, MPI Ranks start at 0
        do i=1,vApp%CplSize
          ! Assign rank if match
          select case (vApp%IAm(i))
            case (voltId)
              if (i-1 .ne. vApp%voltCplRank) then
                write(*,*) "I AM NOT MYSELF:",i-1,vApp%voltCplRank
              endif
            case (gamId)
              write(*,*) "Gam not involved in MAGE2MAGE yet"
            case (rcmId)
              write(*,*) "RCM not involved in MAGE2MAGE yet"
            case (hidraNId)
              vApp%hidraNCplRank = i-1
              write(*,*) "Volt coupling to hidraN"
            case (hidraSId)
              vApp%hidraSCplRank = i-1
              write(*,*) "Volt coupling to hidraS"
            case (hidraId)
              vApp%hidraCplRank = i-1
              write(*,*) "Volt coupling to hidra"
            case (tgcmId)
              vApp%gcmCplRank = i-1
              write(*,*) "Volt coupling to TIEGCM"
            case default
              write(*,*) "Volt does not know about this Coupling ID: ", vApp%IAm(i)
          end select
        enddo

      endif
      if(vApp%CplSize == 1 .and. (.not. vApp%doGCM)) then
        write(*,*) "VOLTRON: We're not coupling to a GCM"
        call mpi_comm_free(vApp%mageCplComm, ierror)
        if(ierror /= MPI_Success) then
          call MPI_Error_string( ierror, message, length, ierror)
          print *,message(1:length)
          call mpi_Abort(MPI_COMM_WORLD, 1, ierror)
        end if
        vApp%mageCplComm = MPI_COMM_NULL
      endif
      if (vApp%CplSize == 1 .and. vApp%doGCM) then
        write(*,*) "VOLTRON: Coupling to GCM Failed."
        call mpi_Abort(MPI_COMM_WORLD, 1, ierror)
      endif
    else
        call MPI_Comm_Split(MPI_COMM_WORLD, MPI_UNDEFINED, worldRank, vApp%mageCplComm, ierror)
    endif
    if(ierror /= MPI_Success) then
        call MPI_Error_string( ierror, message, length, ierror)
        print *,message(1:length)
        call mpi_Abort(MPI_COMM_WORLD, 1, ierror)
    end if

    if(isGamera) then
        call Tic("Omega", .true.)
        call initGamera_mpi(gApp,userInitFunc,gamComm,doIO=.false.)
        call initGam2Volt(g2vComm,gApp,MPI_COMM_WORLD)
        call Toc("Omega", .true.)

        do while (g2vComm%time < g2vComm%tFin)
            !Start root timer
            call Tic("Omega", .true.)
        
            !Advance Gamera MHD
            call stepGamera_mpi(gApp)

            !Update local gamera time units
            call localStepVoltronTime(g2vComm, gApp)

            !Tell gamera to synchronize with voltron (and wait for it if necessary)
            ! If it needs to do deep coupling, file io, restart io, or
            ! Terminate the simulation
            ! Don't synchronize for console or timing output, not worth it (?)
            if( (g2vComm%time >= g2vComm%DeepT) .or. &
                gApp%Model%IO%doRestart(gApp%Model%t) .or. &
                gApp%Model%IO%doOutput(gApp%Model%t) .or. &
                (g2vComm%time >= g2vComm%tFin)) then
                !Do any updates to Voltron
                call Tic("StepVoltron")
                call performStepVoltron(g2vComm,gApp)
                call Toc("StepVoltron")
        
                !Coupling
                if ( g2vComm%time >= g2vComm%DeepT ) then
                    call Tic("Coupling", .true.)
                    call performDeepUpdate(g2vComm, gApp)
                    call Toc("Coupling", .true.)
                endif
            endif

            !IO checks
            call Tic("IO", .true.)
            !NOTE: Does this need to be duplicated from gamera_mpix or can both be done w/ a single subroutine?
            
            !Console output
            if (gApp%Model%IO%doConsole(g2vComm%ts)) then
                !Using console output from Gamera
                call consoleOutput_mpi(gApp)
                if (gApp%Model%IO%doTimerOut .and. debugPrintingRank(gApp)) then
                    write (*,'(a,I3,a,I3,a,I3,a)') "Rank (I,J,K) (", &
                        gApp%Grid%Ri,",", &
                        gApp%Grid%Rj,",", &
                        gApp%Grid%Rk,") is printing debug clock info"
                    call printClocks()
                endif
                call cleanClocks()
            elseif (gApp%Model%IO%doTimer(g2vComm%ts)) then
                if (gApp%Model%IO%doTimerOut .and. debugPrintingRank(gApp)) then
                    write (*,'(a,I3,a,I3,a,I3,a)') "Rank (I,J,K) (", &
                        gApp%Grid%Ri,",", &
                        gApp%Grid%Rj,",", &
                        gApp%Grid%Rk,") is printing debug clock info"
                    call printClocks()
                endif
                call cleanClocks()
            endif

            !Data output
            if (gApp%Model%IO%doOutput(gApp%Model%t)) then
                call fOutput(gApp%Model, gApp%Grid, gApp%State)
            endif
            !Restart output
            if (gApp%Model%IO%doRestart(gApp%Model%t)) then
                call resOutput(gApp%Model, gApp%Grid, gApp%oState, gApp%State)
            endif

            call Toc("IO", .true.)

            call Toc("Omega", .true.)

        end do

        call endGam2VoltWaits(g2vComm, gApp)

    else 
        ! voltron
        call Tic("Omega", .true.)
        call initVoltron_mpi(vApp, userInitFunc, voltComm, MPI_COMM_WORLD)
        call Toc("Omega", .true.)

        if(vApp%amHelper) then
            ! do helper loop
            helperQuit = .false.
            do while(.not. helperQuit)
                call Tic("Omega", .true.)
                call helpVoltron(vApp, helperQuit)
                call Toc("Omega", .true.)
            end do
        else
            ! voltron run loop
            do while (vApp%time < vApp%tFin)
                !Start root timer
                call Tic("Omega", .true.)
                !If coupling from Gamera is ready
                if(gameraStepReady(vApp)) then
                    !Do any updates to Voltron
                    call Tic("StepVoltronAndWait")
                    call stepVoltron_mpi(vApp)
                    call Toc("StepVoltronAndWait")

                    !Coupling
                    call Tic("Coupling", .true.)
                    if (vApp%time >= vApp%DeepT ) then
                        call DeepUpdate_mpi(vApp, vApp%time)
                    endif
                    call Toc("Coupling", .true.)

                    !IO checks
                    call Tic("IO", .true.)
                    !Console output
                    if (vApp%IO%doConsole(vApp%ts)) then
                        call consoleOutputVOnly(vApp,vApp%gAppLocal,vApp%gAppLocal%Model%MJD0)
                        if (vApp%IO%doTimerOut) then
                            call printClocks()
                        endif
                        call cleanClocks()
                    elseif (vApp%IO%doTimer(vApp%ts)) then
                        if (vApp%IO%doTimerOut) then
                            call printClocks()
                        endif
                        call cleanClocks()
                    endif

                    !Data output
                    if (vApp%IO%doOutput(vApp%time)) then
                        call fOutputVOnly(vApp,vApp%gAppLocal)
                    endif
                    !Restart output
                    if (vApp%IO%doRestart(vApp%time)) then
                        call resOutputVOnly(vApp,vApp%gAppLocal)
                    endif

                    call Toc("IO", .true.)

                elseif(deepInProgress(vApp)) then
                    ! If we did not couple, check for deep work to be done
                    call Tic("Coupling", .true.)
                    call doDeepBlock(vApp)
                    call Toc("Coupling", .true.)
                else
                    ! Gamera wasn't ready and we don't have deep work to do, wait for gamera
                    call waitForGameraStep(vApp)
                endif
                call Toc("Omega", .true.)
            enddo
            ! if using helpers, tell them to quit
            if(vApp%useHelpers) call vhReqHelperQuit(vApp)
        endif ! end all voltron loops
        call endVoltronWaits(vApp)
    endif

    call MPI_FINALIZE(ierror)
    write(*,*) "Fin"
    

end program voltron_mpix

