MODULE reservoir_init
!-----------------------------------------------------------------------
! NAME
!     reservoir_init
!
! DESCRIPTION
!     Module to initialize reservoirs.
!
! AUTHORS & DATE
!     JB Clement, 06/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use numerics, only: dp, di, li, k4

! DECLARATION
! -----------
implicit none

contains
!+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

!=======================================================================
SUBROUTINE setup_hydro_PCM(c1,cr)
!-----------------------------------------------------------------------
! NAME
!     setup_hydro_PCM
!
! DESCRIPTION
!     Initialize the hydrology variables for the PCM.
!
! AUTHORS & DATE
!     JB Clement, 06/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry,  only: ngrid
use display,   only: print_ini, print_simple_end, print_msg, LVL_NFO
use io_netcdf, only: open_nc, close_nc, get_dim_nc, def_var_nc, put_var_nc, startfi_name
use config,    only: read_hydrology_config, rundef_name
use geometry,  only: ini_geometry, end_geometry, dim_init
use hydrology, only: do_hydrology, ini_hydrology, end_hydrology, build4PCM_hydrology, starthydro_name
use utility,   only: int2str
use stoppage,  only: stop_clean

! DECLARATION
! -----------
implicit none

! ARGUMENTS
! ---------
integer(li), intent(in) :: c1, cr

! LOCAL VARIABLES
! ---------------
logical(k4)                            :: here
integer(li)                            :: c2 ! Count of processor clock
integer(di)                            :: cstat, dimid, vdim
real(dp),    dimension(:), allocatable :: peren_h2o_liq_frac4PCM
logical(k4), dimension(:), allocatable :: is_ocean4PCM

! CODE
! ----
! Header
call print_ini('Reservoir initialization for PCM: setting up hydrology')

! Check if file is already initialized
inquire(file = starthydro_name,exist = here)
if (here) call stop_clean(__FILE__,__LINE__,'The file"'//starthydro_name//'" already exists, remove it if you want to initialize hydrology reservoirs for the PCM!',1)

! Read the PEM parameters
call print_msg('> Reading "'//rundef_name//'" for hydrology parameters',LVL_NFO)
call read_hydrology_config()
if (.not. do_hydrology) call stop_clean(__FILE__,__LINE__,"'do_hydrology' must be true!",1)

! Initialize the dimensions
call ini_geometry()
if (.not. dim_init) call stop_clean(__FILE__,__LINE__,"dimensions of module 'geometry' were not correctly initialized.",1)

! Initialize hydrology (create "starthydro.nc")
call ini_hydrology()

! Build the hydrology variables for the PCM
call print_msg('',LVL_NFO)
call print_msg('********* Set-up *********',LVL_NFO)
allocate(peren_h2o_liq_frac4PCM(ngrid))
allocate(is_ocean4PCM(ngrid))
call build4PCM_hydrology(peren_h2o_liq_frac4PCM,is_ocean4PCM)

! Check if "startfi.nc" exists
inquire(file = startfi_name,exist = here)
if (.not. here) call stop_clean(__FILE__,__LINE__,'cannot find required file "'//startfi_name//'"!',1)

! Copy "startfi.nc" into "restartfi.nc"
call print_msg('> Writing "re'//startfi_name//'"',LVL_NFO)
call execute_command_line('\cp -f '//startfi_name//' re'//startfi_name,cmdstat = cstat)
if (cstat > 0) then
    call stop_clean(__FILE__,__LINE__,'command execution failed!',1)
else if (cstat < 0) then
    call stop_clean(__FILE__,__LINE__,'command execution not supported!',1)
end if

! Rewrite the variables computed by the PEM (and create them beforehand if they do not exist)
call open_nc('re'//startfi_name,'write')
call get_dim_nc('physical_points',vdim,dimid = dimid)
if (vdim /= ngrid) call stop_clean(__FILE__,__LINE__,'dimension ''physical_points'' in "'//startfi_name//'" is not equal to ''ngrid'' ('//int2str(vdim)//'/='//int2str(ngrid)//')!',1)

call def_var_nc('peren_h2o_liq_frac','Perennial h2o liquid area surface coverage fraction','',(/dimid/))
call put_var_nc('peren_h2o_liq_frac',peren_h2o_liq_frac4PCM)

call def_var_nc('is_ocean','Ocean mask','',(/dimid/))
call put_var_nc('is_ocean',merge(1._dp,0._dp,is_ocean4PCM))

! Close
call close_nc('re'//startfi_name)

! Finalization
deallocate(peren_h2o_liq_frac4PCM)
deallocate(is_ocean4PCM)
call end_hydrology()
call end_geometry()

! Footer
call system_clock(c2)
call print_simple_end(real((c2 - c1),dp)/real(cr,dp),'Hydrology initialized in "'//'re'//startfi_name//'" for the PCM.')

END SUBROUTINE setup_hydro_PCM
!=======================================================================

END MODULE reservoir_init