module ips_driver

use driver_module
use IMOD_UTL, only : imod_utl_capf
use imod_idf_par, only : idfobj
use m_main_info
use imod_utl, only: imod_utl_closeunits, imod_utl_has_ext, &
    imod_utl_printtext,imod_utl_getunit, imod_utl_getslash
use pks_imod_utl, only: pks_imod_utl_idfmerge_init, &
    pks_imod_utl_write_idfmergefile, pks_imod_utl_idfmerge ! PKS
use pks_imod_utl, only: pks_imod_utl_iarmwp_xch_read
use, intrinsic :: ieee_exceptions

implicit none
  
logical          :: doTimeLoop
logical          :: ok, usests, usestsmodflow, stsave, strestore
logical          :: llpf, psolved, lpks, lidfmerge
logical          :: converged, mf_steadystate

integer          :: date, hour, minute, second
integer          :: retValMF2005, mf_igrid, exitcode
integer          :: tsc, ncvgerr ! debug variable
integer          :: ios, lunc, lunsts, lc, ivcl, jvcl, iarg, jarg, narg

double precision :: currentTime, endOfCurrentTimeStepMF2005
double precision :: endOfCurrentTimeStep, nodataTime, BeginOfNextTimeStepMF2005

character (len=1024) :: modwd2, simwd2, root, idfmergefile

contains

! Initialise simulation
!---------------------------------------------------------------------------------------
subroutine driver_init_ips( root, inputfile )

implicit none
character(len=*), intent(in)     :: root, inputfile

! parameters
integer, parameter :: nsubmax = 1000

! functions
logical             pks7mpimasterwrite
integer             osd_open2,cfn_length
double precision    cfn_mjd_nodata

! General
logical    :: usemodflow, usemetaswap, usetransol, usemozart, mozstsave
logical    :: lrunfile, lnamfile, lipest, lss, lrf
logical    :: lwstdo

integer    :: isub, nsub, nnsub, mf_ngrid, n!, exitcode

character                              :: tekens(1)*1
character                              :: del*1, cont*1, comm*1
character (len=1)                      :: slash
character (len=32)                     :: compcls
character (len=50), dimension(nsubmax) :: submstr=''
character (len=1024)                   :: wd, modwd1, simwd1, &
                                          mozwd, infile, rfopt, rfroot
character (len=1024)                   :: record, modrecord
character (len=1024)                   :: compfile, str

real               :: hnoflo
real, dimension(4) :: imodusebox
double precision   :: currentTimeMF2005



! program section
! ------------------------------------------------------------------------------
CALL IEEE_SET_HALTING_MODE (IEEE_DIVIDE_BY_ZERO, .TRUE.)


call pks7mpiini1(lwstdo) ! PKS
call pks7mpiactive(lpks) ! PKS

! Evaluate iMOD license
if(pks7mpimasterwrite()) call imod_license()
call pks7mpibarrier() ! PKS

! ... init
exitcode    = 0            ! when 0: ok
lunc        = 9            ! unit number for components file
lunsts      = 8            ! <=0: sts2init() will assign a free unit number
!  >0: unit number will be used
comm        = '!'          ! comment sign for components input file
cont        = char(92)     ! continuation mark for components input file records
nodataTime  = cfn_mjd_nodata()

nsub = 0
NCVGERR = 0

! ... process initialisation
! ... init components
!     ios: 0: OK
!         -1: commandline argument not found
!        <-1: ERROR
! define virtual commandline

modwd1  = root
simwd1  = root
mozwd   = root

! get slash
call imod_utl_getslash(slash)

! run-file
lrunfile    = .false.
lnamfile    = .false.
usemodflow  = .false.
usemetaswap = .false.
lipest      = .false.
lidfmerge   = .false.

infile = inputfile
!call cfn_vcl_arg(ivcl,1,infile,n)
if (imod_utl_has_ext(infile,'nam')) lnamfile= .true.
if (imod_utl_has_ext(infile,'run')) lrunfile= .true.
if (lrunfile .or. lnamfile) then ! run-file or nam-file
 usemodflow=.true.
 modrecord=''
 if (lnamfile) write(modrecord,'(a,1x,3a)') '-namfile','"',trim(infile),'"'
 if (lrunfile) then
  write(modrecord,'(a,1x,3a)') '-runfile','"',trim(infile),'"'
  if (narg.eq.3) then ! runfile options
   call cfn_vcl_arg(ivcl,2,rfopt,n)
   write(modrecord,'(a,1x,a,1x,a)') trim(modrecord),'-rfopt',trim(rfopt)
   call cfn_vcl_arg(ivcl,3,rfopt,n)
   write(modrecord,'(a,1x,a)') trim(modrecord),trim(rfopt)
  end if
 end if
 wd = ''
 call rf2mf_prg(lrunfile,lipest,lidfmerge,modrecord,usemetaswap, &
     submstr,nsub,nsubmax,wd,imodusebox,rfroot)
 modwd1 = trim(wd)
 call osd_s_filename(modwd1)
end if

usests        = .false.
usestsmodflow = .false.
usetransol    = .false.
usemozart     = .false.

rt = driverGetRunType(usemodflow,usemetaswap,usetransol,usemozart,usests)
str = 'Running mode: '//trim(rtdesc(rt))
if(lpks) str = trim(str)//' (Parallel Krylov Solver activated)'
if (pks7mpimasterwrite()) write(*,'(79(''*''),/,(1x,a),/,79(''*''))') trim(str)

! check state-save options
if (usests.and..not.usestsmodflow) call driverChk(.false.,&
 'MODFLOW state-save not activated.')
if (.not.usests.and.usestsmodflow) call driverChk(.false.,&
 'MODFLOW state-save activated.')

nnsub = max(1,nsub)
! ################################## SUBMODEL LOOP ############################
submodelloop: do isub = 1,nnsub
! ################################## PEST LOOP ################################
 call pks_imod_utl_idfmerge_init() ! PKS

 if (lrunfile) then
  if (nsub.gt.0) then
   write(modwd2,'(5a)') trim(modwd1),slash,trim(submstr(isub)),slash,'mf2005_tmp'
   write(*,'(50(''+''),/,1x,a,1x,a,1x,a,i3.3,a,i3.3,a,/,50(''+''))') &
       'Computing for submodel:',trim(submstr(isub)),'(',isub,'/',nsub,')'
  else
   write(modwd2,'(3a)') trim(modwd1),slash,'mf2005_tmp'
  end if
 else
     modwd2 = modwd1
     simwd2 = simwd1
 end if

 ! PKS IARMWP option
 call pks_imod_utl_iarmwp_xch_read(modwd2)

 timestep    = 1
 stsave      = .false.
 strestore   = .false.

 ! init return values
 retValMF2005      = 0
 ! time values
 currentTime       = nodataTime
 currentTimeMF2005 = 0.d0

 ! init all components TODO!!!!!!!
 if (usemodflow) then
  call osd_chdir(modwd2)
  if (usestsmodflow) call sts2init(usestsmodflow,lunsts)
  call mf2005_initComponent(modrecord,retValMF2005)
  ! get number of grids
  ok = mf2005_PutNumberOfGrids(mf_ngrid); call driverChk(ok,&
   'mf2005_PutNumberOfGrids')
  if (mf_ngrid.ne.1) call driverChk(.false.,&
   'More than one MODFLOW grids is not supported.'); mf_igrid = 1
 end if

 ! check if MODFLOW is activated
 if (retValMF2005.ne.0) call driverChk(.false.,'MODFLOW initialization failed.')

 ! get starting date for this simulation
 call mf2005_getCurrentTime(currentTimeMF2005,retValMF2005)
 currentTime = currentTimeMF2005

 ! start sts
 if (usestsmodflow) call sts2start(currentTime)

 ! read data for entire simulation
 call osd_chdir(modwd2)
 call mf2005_initSimulation(currentTime, retValMF2005)


 if (retValMF2005.ne.0) exitcode = -12

 ! #### BEGIN EXCHANGE: BeforeInitSimulationMozart #############################

 ! get flag if BCF is used or LPF
 ok = mf2005_PutLPFActive(mf_igrid, llpf); call driverChk(ok,&
  'mf2005_PutLPFActive')
 ! get flag if simulation is steady-state or transient
 ok = mf2005_PutSimulationType(mf_igrid, lss); call driverChk(ok,&
  'mf2005_PutSimulationType')
 ! get hnoflo
 ok = mf2005_PutHeadNoFlo(mf_igrid, hnoflo);  call driverChk(ok,&
'mf2005_PutHeadNoFlo')

 ! #### BEGIN EXCHANGE: BeforeTimestep #########################################
 if (rt.eq.rtmodsim) then
  ok = mf2005_PutSimulationType(mf_igrid, mf_steadystate)
  ! MODFLOW puts the simulation type (transient/steady-state)
  call driverChk(ok,'mf2005_PutSimulationType')
  if (mf_steadystate) call driverChk(ok,'MODFLOW may not be steady-state')
 end if

 ! Create mappings
 if (rt.eq.rtmodsim) call driverXchIniMapModSim()

 ! ###### END EXCHANGE: BeforeTimestep #########################################

 ! ... timestep
 timestep = 0

 tsc=0 ! debug: TimeStepCycle

 stsave    = .false.
 mozstsave = .false.
 strestore = .false.
end do submodelloop

return
end subroutine driver_init_ips

! Perform simulation
!---------------------------------------------------------------------------------------
subroutine driver_perform_iter_ips()
implicit none

logical     :: pks7mpimasterwrite, endOfSimulation, convergedMF2005

! convergence
convergedMF2005   = .true.
endOfSimulation   = .false.

exitcode = 0

do while (.not.endOfSimulation .and. exitcode.eq.0) ! MOZART-loop
 do while ((.not.endOfSimulation .and. exitcode.eq.0))
  doTimeLoop = .true.
  do while(doTimeLoop .and. exitcode.eq.0)
   tsc=tsc+1
   timestep = timestep + 1

   call cfn_mjd2datehms(currentTime,date,hour,minute,second)
   !##### BEGIN EXCHANGE: BeginOfTimestep ########################################
   if (rt.eq.rtmodsim) then
    ok = mf2005_PutHeads(mf_igrid,XchModSimModCells,XchModSimModNID,&
        XchModSimModHeads,mv)
    ! Put groundwater heads
    call DriverChk(ok,'mf2005_PutHeads')
    XchModSimModHeads = mv
   end if
   !####### END EXCHANGE: BeginOfTimestep ########################################

   ! perform state save, phase 1 (before reading and writing data)
   if (usests) then
    call osd_chdir(modwd2)
    if (usestsmodflow) call sts2saverestore(currentTime,stsave,strestore,1)
    if (rt.eq.rtmodsim) then
        call osd_chdir(simwd2)
    end if
   end if

   ! init timestep
   call mf2005_prepareTimestep(currentTime,stsave,retValMF2005)
   call osd_chdir(modwd2)
   call mf2005_initTimeStep(currentTime,retValMF2005)
   if (retValMF2005.ne.0) exitcode = -15

   ! get end of current timestep
   call mf2005_getEndOfCurrentTimeStep(endOfCurrentTimeStepMF2005,retValMF2005)
   if (retValMF2005.ne.0) exitcode = -17
   endOfCurrentTimeStep = endOfCurrentTimeStepMF2005

   ! extra check
   if (endOfCurrentTimeStep.eq.nodataTime) call driverChk(.false., &
      'endOfCurrentTimeStep = nodataTime')
   if (pks7mpimasterwrite()) &
    call mf2005_writeTimeStep(tsc,date,hour,minute,second)

   ! iteration
   converged =.false.
   do while (.not. converged .and. exitcode.eq.0)

    if (rt.eq.rtmodsim) then
     XchModSimSimUnsZFlux = mv
     if (.not.mf_steadystate) then
      if (.not.llpf) then ! BCF
       ok = mf2005_GetStorageFactor(mf_igrid,XchModSimSimStrFct,&
           XchModSimModNID,XchSim2ModIdx,XchSim2ModOff,mv)
       call DriverChk(ok,'mf2005_GetStorageFactor') ! Get storage factor
      else ! LPF
       ok = mf2005_GetStorageFactorLPF(mf_igrid,XchModSimSimStrFct,&
           XchModSimModNID,XchSim2ModIdx,XchSim2ModOff,mv)
       call DriverChk(ok,'mf2005_GetStorageFactorLPF') ! Get storage factor
      end if
      XchModSimSimStrFct = mv
     end if
    end if
    !####### END EXCHANGE: AfterSolveMetaSwap #####################################

    call mf2005_performIter(retValMF2005, psolved)

    if (retValMF2005.ne.0) exitcode = -26

    !##### BEGIN EXCHANGE: AfterSolve #############################################
    if (exitcode.eq.0) then
     if (rt.eq.rtmodsim) then
      ok = mf2005_PutHeads(mf_igrid,XchModSimModCells,&
           XchModSimModNID,XchModSimModHeads,mv)
      ! MODFLOW put groundwater heads
      call DriverChk(ok,'mf2005_PutHeads')
      XchModSimModHeads = mv
     end if
    end if
    !####### END EXCHANGE: AfterSolve #############################################

    ! finish iteration
    call mf2005_finishIter(convergedMF2005,retValMF2005)

    ! get next iteration information
    converged = .true.
    converged = converged .and. convergedMF2005
    if (rt.eq.rtmodsim .or. rt.eq.rtmodsimtranmoz)&
        converged = converged !.and. convergedMetaSwap

    if (retValMF2005.ne.0    ) exitcode = -261
   enddo
   if (.not.converged) then
    if (pks7mpimasterwrite()) write(*,*) ' Model did not converge!',exitcode
   end if

   ! ... write results
   if (converged .and. exitcode.eq.0) then
    call osd_chdir(modwd2)
    call mf2005_finishTimestep(retValMF2005,psolved,ncvgerr)
    if (retValMF2005.ne.0) exitcode = -31
   endif

   if (exitcode.eq.0) then
    ! perform state save, phase 2 (after reading and writing data)
    if (usestsmodflow) call sts2saverestore(currentTime,stsave,strestore,2)
   endif

   ! ... next timestep
   if (exitcode.ne.0) then ! ERROR, do not continue
    endOfSimulation=.true.
   else
    ! get new time information
    ! Let Modflow be leading. don't ask the other components
    call mf2005_getBeginOfNextTimeStep(BeginOfNextTimeStepMF2005,retValMF2005)
    ! find errors
    if (retValMF2005.eq.-1) then
     ! end of simulation for Modflow
     endOfSimulation =.true.
     retValMF2005    = 0
    endif
    if (retValMF2005.ne.0) exitcode = -31
    if (exitcode.ne.0)     endOfSimulation=.true.
    ! find start time of next timestep
    ! this may be a time in the 'future', the 'past' or the 'current' time
    ! nodata value of time-value means: 
    ! this was the last timestep for this component
    if (exitcode.eq.0) then
     currentTime=nodataTime
     ! check for each component (ok, only modflow available now)
     if (BeginOfNextTimeStepMF2005.ne.nodataTime) then
      if (currentTime.eq.nodataTime) then
       currentTime=BeginOfNextTimeStepMF2005
      else
       currentTime=min(currentTime,BeginOfNextTimeStepMF2005)
      endif
     endif
     if (currentTime.eq.nodataTime.and.rt.ne.rtmod) endOfSimulation=.true.
     !                if (currentTime.eq.nodataTime) endOfSimulation=.true.
    endif
   endif
   doTimeLoop = .not.endOfSimulation
  enddo ! time loop
 end do ! Mozart loop
end do
call osd_chdir(modwd2)
if (exitcode.eq.0)  call mf2005_finishSimulation(retValMF2005)
if (retValMF2005.ne.0 ) exitcode = -61

return
end subroutine driver_perform_iter_ips

! Finalise
!---------------------------------------------------------------------------------------
subroutine driver_finalise_ips()
implicit none
logical   :: pks7mpimasterwrite

! close files and deallocate memory
call mf2005_closeDeallocate( retValMF2005 )

! list routine timer info
call cfn_rtt_list(0)

! Deallocate coupling arrays
call driverXchDeallocate()

! exit
if (exitcode.ne.0) then
 if (pks7mpimasterwrite()) write(unit=*,fmt=*) 'ERROR, exit code: ',exitcode
 call pks7mpifinalize()! PKS
 call exit(10)
endif

!-----end of program.
if(ncvgerr.gt.0) then
 if (pks7mpimasterwrite()) &
  write(*,*) 'FAILED TO MEET SOLVER CONVERGENCE CRITERIA ',ncvgerr,' TIME(S)'
end if

! next pest iteration
call imod_utl_closeunits()
! PKS IDF merge
call osd_chdir(root)
call pks_imod_utl_write_idfmergefile(modwd2,idfmergefile)
if (lidfmerge.and.lpks) call pks_imod_utl_idfmerge(idfmergefile)

return
end subroutine driver_finalise_ips
end module