MODULE hydrology
!-----------------------------------------------------------------------
! NAME
!     hydrology
!
! DESCRIPTION
!     Module for managing surface water reservoirs.
!
! AUTHORS & DATE
!     A. Gauvain, 05/2026
!     JB Clement, 05/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use numerics, only: di, dp, qp, k4, eps

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

! VARIABLES
! ---------
real(dp), dimension(:), allocatable :: lake_depth  ! Depth of lakes [m]
real(dp), dimension(:), allocatable :: lake_area   ! Area of lakes [m2]
real(dp), dimension(:), allocatable :: lake_volume ! Volume of lakes [m3]
real(dp), dimension(:), allocatable :: lake_frac   ! Lake area fraction [m2/m2] (0-1)

! PARAMETERS
! ----------
! Run hydrology-related processes in PEM
character(13),  parameter          :: starthydro_name = 'starthydro.nc'
real(dp),       parameter          :: h2o_liq_density = 1000._dp ! Density of liquid water [kg/m3]
logical(k4),    protected          :: do_hydrology = .false.
character(256), protected, private :: hydrodata_name = 'hydrological_database.nc'
real(qp),       protected, private :: gel_water = 100._qp
integer(di),    protected, private :: init_gel_state = 1_di
real(qp),       protected, private :: gel_init_lon = 0._qp
real(qp),       protected, private :: gel_init_lat = 0._qp
logical(k4),    protected, private :: active_bypass = .true.

! Area covered by the depressions associated to each grid cell [m2]
real(dp), dimension(:), allocatable, protected :: dep_area_in_cell

! Perennial surface water area in the PCM at the beginning [m2]
real(dp), dimension(:), allocatable, protected :: peren_h2o_liq_frac_PCM

! Ocean mask in the PCM at the beginning
logical(k4), dimension(:), allocatable, protected :: is_ocean_PCM

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

!=======================================================================
SUBROUTINE set_hydrology_config(do_hydrology_in,hydrodata_path_in,gel_water_in,init_gel_state_in,gel_init_lon_in,gel_init_lat_in,active_bypass_in)
!-----------------------------------------------------------------------
! NAME
!     set_hydrology_config
!
! DESCRIPTION
!     Setter for "hydrology" configuration parameters.
!
! AUTHORS & DATE
!     A. Gauvain, 05/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use utility,  only: bool2str, real2str, int2str
use display,  only: print_msg, LVL_NFO
use stoppage, only: stop_clean

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

! ARGUMENTS
! ---------
logical(k4),    intent(in) :: do_hydrology_in, active_bypass_in
character(256), intent(in) :: hydrodata_path_in
real(dp),       intent(in) :: gel_water_in, gel_init_lon_in, gel_init_lat_in
integer(di),    intent(in) :: init_gel_state_in

! CODE
! ----
do_hydrology = do_hydrology_in
if (len_trim(adjustl(hydrodata_path_in)) > len(hydrodata_name)) then
    call stop_clean(__FILE__,__LINE__,'''hydrodata_name'' is too long!',1)
else
    hydrodata_name = trim(adjustl(hydrodata_path_in))
end if
gel_water = gel_water_in
init_gel_state = init_gel_state_in
gel_init_lon = gel_init_lon_in
gel_init_lat = gel_init_lat_in
active_bypass = active_bypass_in
call print_msg('do_hydrology   = '//bool2str(do_hydrology),LVL_NFO)
call print_msg('hydrodata_name = '//trim(adjustl(hydrodata_name)),LVL_NFO)
call print_msg('gel_water      = '//real2str(gel_water),LVL_NFO)
call print_msg('init_gel_state = '//int2str(init_gel_state),LVL_NFO)
if (init_gel_state == 2) then
    call print_msg('gel_init_lon   = '//real2str(gel_init_lon),LVL_NFO)
    call print_msg('gel_init_lat   = '//real2str(gel_init_lat),LVL_NFO)
end if
call print_msg('active_bypass  = '//bool2str(active_bypass),LVL_NFO)
if (gel_water <= 0._dp) call stop_clean(__FILE__,__LINE__,'''gel_water'' must be positive!',1)
if (init_gel_state < 1 .or. init_gel_state > 3) call stop_clean(__FILE__,__LINE__,'''init_gel_state'' is not a valid method (choose between 1 and 3)!',1)

END SUBROUTINE set_hydrology_config
!=======================================================================

!=======================================================================
SUBROUTINE ini_hydrology()
!-----------------------------------------------------------------------
! NAME
!     ini_hydrology
!
! DESCRIPTION
!     Allocate 'hydrology' module arrays.
!
! AUTHORS & DATE
!     JB Clement, 05/2026
!     A. Gauvain, 05/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry,        only: ngrid, igridfromlonlat
use stoppage,        only: stop_clean
use hydro_database,  only: load_hydro_database, clear_database_variables, load_diaghydro_iteration
use hydro_database,  only: ndep, label, child_list
use hydro_tree,      only: allocate_tree, build_tree, nodes
use hydro_core,      only: initialize_model, run_simulation
use hydro_diaghydro, only: allocate_outputs, deallocate_outputs
use utility,         only: int2str

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

! LOCAL VARIABLES
! ---------------
integer(di) :: i, ig
logical(k4) :: here

! CODE
! ----
if (.not. do_hydrology) return

inquire(file = trim(adjustl(hydrodata_name)),exist = here)
if (.not. here) call stop_clean(__FILE__,__LINE__,'cannot find required hydrological database file "'//trim(adjustl(hydrodata_name))//'"!',1)

call load_hydro_database(trim(adjustl(hydrodata_name)))
call allocate_tree(nodes)
call build_tree(nodes)
call clear_database_variables()
call allocate_outputs(ndep)
call initialize_model(starthydro_name,init_gel_state,gel_init_lon,gel_init_lat,gel_water,merge(1,0,active_bypass))
deallocate(label)

allocate(peren_h2o_liq_frac_PCM(ngrid))
allocate(is_ocean_PCM(ngrid))
allocate(dep_area_in_cell(ngrid))
allocate(lake_frac(ngrid))
allocate(lake_area(ngrid))
allocate(lake_depth(ngrid))
allocate(lake_volume(ngrid))
peren_h2o_liq_frac_PCM(:) = 0._dp
lake_frac(:) = 0._dp
lake_area(:) = 0._dp
lake_volume(:) = 0._dp
lake_depth(:) = 0._dp
dep_area_in_cell(:) = 0._dp

!$OMP PARALLEL DO PRIVATE(ig) SCHEDULE(static) REDUCTION(+:dep_area_in_cell)
do i = 1,ndep
    if (nodes(child_list(i))%is_leaf) then
        ! AG: These two next lines can be computed only once after the first call of initialize_model()
        ig = igridfromlonlat(nodes(child_list(i))%lake_lon,nodes(child_list(i))%lake_lat)
        nodes(child_list(i))%gcm_idx = ig
        dep_area_in_cell(ig) = dep_area_in_cell(ig) + real(nodes(child_list(i))%watershed_area,dp)
    end if
end do
!$OMP END PARALLEL DO
if (any(abs(dep_area_in_cell(:)) < eps)) call stop_clean(__FILE__,__LINE__,'at least one depression area associated to a grid cell is zero!',1)

END SUBROUTINE ini_hydrology
!=======================================================================

!=======================================================================
SUBROUTINE end_hydrology()
!-----------------------------------------------------------------------
! NAME
!     end_hydrology
!
! DESCRIPTION
!     Deallocate 'hydrology' module arrays.
!
! AUTHORS & DATE
!     JB Clement, 05/2026
!     A. Gauvain, 05/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use hydro_tree,      only: nodes
use hydro_diaghydro, only: deallocate_outputs
use hydro_core,      only: deallocate_nodes
use hydro_database,  only: longitude, latitude
use hydro_diaghydro, only: Z_out, V_out, A_out, Vout_out, Vin_out, active_out, ids_out

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

! CODE
! ----
if (allocated(peren_h2o_liq_frac_PCM)) deallocate(peren_h2o_liq_frac_PCM)
if (allocated(is_ocean_PCM)) deallocate(is_ocean_PCM)
if (allocated(dep_area_in_cell)) deallocate(dep_area_in_cell)
if (allocated(lake_frac)) deallocate(lake_frac)
if (allocated(lake_area)) deallocate(lake_area)
if (allocated(lake_depth)) deallocate(lake_depth)
if (allocated(lake_volume)) deallocate(lake_volume)
if (allocated(longitude)) deallocate(longitude)
if (allocated(latitude)) deallocate(latitude)
call deallocate_outputs()    ! hydro_inout.F90 - Diaghydro Module
call deallocate_nodes(nodes) ! hydro_tree.F90 - TreeNodeModule
if (allocated(Z_out)) deallocate(Z_out)
if (allocated(V_out)) deallocate(V_out)
if (allocated(A_out)) deallocate(A_out)
if (allocated(Vout_out)) deallocate(Vout_out)
if (allocated(Vin_out)) deallocate(Vin_out)
if (allocated(active_out)) deallocate(active_out)
if (allocated(ids_out)) deallocate(ids_out)

END SUBROUTINE end_hydrology
!=======================================================================

!=======================================================================
SUBROUTINE set_peren_h2o_liq_frac_PCM(peren_h2o_liq_frac_PCM_in)
!-----------------------------------------------------------------------
! NAME
!     set_peren_h2o_liq_frac_PCM
!
! DESCRIPTION
!     Setter for 'peren_h2o_liq_frac_PCM'.
!
! AUTHORS & DATE
!     JB Clement, 05/2026
!     A. Gauvain, 05/2026
!
! NOTES
!
!-----------------------------------------------------------------------

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

! ARGUMENTS
! ---------
real(dp), dimension(:), intent(in) :: peren_h2o_liq_frac_PCM_in

! CODE
! ----
peren_h2o_liq_frac_PCM(:) = peren_h2o_liq_frac_PCM_in(:)

END SUBROUTINE set_peren_h2o_liq_frac_PCM
!=======================================================================

!=======================================================================
SUBROUTINE set_is_ocean_PCM(is_ocean_PCM_in)
!-----------------------------------------------------------------------
! NAME
!     set_is_ocean_PCM
!
! DESCRIPTION
!     Setter for 'is_ocean_PCM'.
!
! AUTHORS & DATE
!     JB Clement, 05/2026
!
! NOTES
!
!-----------------------------------------------------------------------

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

! ARGUMENTS
! ---------
real(dp), dimension(:), intent(in) :: is_ocean_PCM_in

! CODE
! ----
is_ocean_PCM(:) = is_ocean_PCM_in(:) > 0.5_dp

END SUBROUTINE set_is_ocean_PCM
!=======================================================================

!=======================================================================
SUBROUTINE build4PCM_hydrology(lake_frac4PCM,is_ocean4PCM)
!-----------------------------------------------------------------------
! NAME
!     build4PCM_hydrology
!
! DESCRIPTION
!     Build perennial lake fraction field sent back to the PCM.
!
! AUTHORS & DATE
!     A. Gauvain, 05/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry, only: cell_area
use display,  only: print_msg, LVL_NFO
use utility,  only: real2str

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

! ARGUMENTS
! ---------
real(dp),    dimension(:), intent(out) :: lake_frac4PCM
logical(k4), dimension(:), intent(out) :: is_ocean4PCM

! CODE
! ----
! Build perennial lake fraction for the PCM
call compute_lake_coverage(lake_frac4PCM)

! Build ocean mask for the PCM
call print_msg('> Computing ocean mask',LVL_NFO)
is_ocean4PCM(:) = .false.
where (lake_frac4PCM(:) >= 0.9_dp) is_ocean4PCM(:) = .true. ! Grid cell is considered as ocean if lake coverage fraction >= 90%
call print_msg('Ocean global coverage fraction [%]: '//real2str(100._dp*sum(merge(1._dp,0._dp,is_ocean4PCM(:))*cell_area(:))/sum(cell_area(:))),LVL_NFO)

END SUBROUTINE build4PCM_hydrology
!=======================================================================

!=======================================================================
SUBROUTINE compute_lake_coverage(lake_frac_out)
!-----------------------------------------------------------------------
! NAME
!     compute_lake_coverage
!
! DESCRIPTION
!     Compute the perennial lake coverage fraction.
!
! AUTHORS & DATE
!     A. Gauvain, 06/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use hydro_core, only: ndep, nodes, child_list
use display,    only: print_msg, LVL_NFO
use utility,    only: real2str

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

! ARGUMENTS
! ---------
real(dp), dimension(:), intent(out) :: lake_frac_out

! LOCAL VARIABLES
! ---------------
integer(di) :: i, ig

! CODE
! ----
call print_msg('> Computing perennial lake coverage fraction',LVL_NFO)
lake_area(:) = 0._dp
lake_frac_out(:) = 0._dp

!$OMP PARALLEL DO PRIVATE(ig) SCHEDULE(static) REDUCTION(+:lake_area,lake_volume)
do i = 1,ndep
    if (nodes(child_list(i))%is_leaf) then
        ig = nodes(child_list(i))%gcm_idx
        lake_area(ig) = lake_area(ig) + real(nodes(nodes(child_list(i))%upper_active)%A &
                                             *nodes(child_list(i))%watershed_area       &
                                             /nodes(nodes(child_list(i))%upper_active)%watershed_area,dp)
    end if
end do
!$OMP END PARALLEL DO

lake_frac_out(:) = lake_area(:)/dep_area_in_cell(:)
where (lake_frac_out(:) < 0._dp) lake_frac_out(:) = 0._dp
where (lake_frac_out(:) > 1._dp) lake_frac_out(:) = 1._dp
call print_msg('Lake per-cell coverage fraction [%] (min|max): '//real2str(100._dp*minval(lake_frac_out))//' | '//real2str(100._dp*maxval(lake_frac_out)),LVL_NFO)

END SUBROUTINE compute_lake_coverage
!=======================================================================

!=======================================================================
SUBROUTINE compute_depression_volume(lake_volume_out)
!-----------------------------------------------------------------------
! NAME
!     compute_depression_volume
!
! DESCRIPTION
!     Compute the perennial lake coverage fraction.
!
! AUTHORS & DATE
!     A. Gauvain, 06/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use hydro_core, only: ndep, nodes, child_list
use display,    only: print_msg, LVL_NFO
use utility,    only: real2str
use hydro_core, only: compute_volume_above_depression

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

! ARGUMENTS
! ---------
real(dp), dimension(:), intent(out) :: lake_volume_out

! LOCAL VARIABLES
! ---------------
integer(di) :: i, ig

! CODE
! ----
lake_volume_out(:) = 0._dp

call print_msg('> Computing volume of depressions',LVL_NFO)

!$OMP PARALLEL DO SCHEDULE(static)
do i = 1,ndep
    if (nodes(child_list(i))%active == 1) then
        call compute_volume_above_depression(nodes(child_list(i))%child, nodes(child_list(i))%V, nodes(child_list(i))%A)
    end if
end do
!$OMP END PARALLEL DO

!$OMP PARALLEL DO PRIVATE(ig) SCHEDULE(static) REDUCTION(+:lake_volume_out)
do i = 1,ndep
    if (nodes(child_list(i))%is_leaf) then
        ig = nodes(child_list(i))%gcm_idx
        lake_volume_out(ig) = lake_volume_out(ig) +  real(nodes(child_list(i))%Vtop)
    end if
end do
!$OMP END PARALLEL DO

call print_msg('Lake per-cell volume [m3] (sum): '//real2str(sum(lake_volume_out)),LVL_NFO)

END SUBROUTINE compute_depression_volume
!=======================================================================

!=======================================================================
SUBROUTINE compute_lake_depth(lake_volume_in,dep_area_in_cell_in,lake_depth_out)
!-----------------------------------------------------------------------
! NAME
!     compute_lake_depth
!
! DESCRIPTION
!     Compute the per-cell lake depth.
!
! AUTHORS & DATE
!     A. Gauvain, 06/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use display,    only: print_msg, LVL_NFO
use utility,    only: real2str
use hydro_core, only: compute_volume_above_depression

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

! ARGUMENTS
! ---------
real(dp), dimension(:), intent(in)  :: lake_volume_in
real(dp), dimension(:), intent(in)  :: dep_area_in_cell_in
real(dp), dimension(:), intent(out) :: lake_depth_out

! CODE
! ----
call print_msg('> Computing lake depth',LVL_NFO)
lake_depth_out(:) = lake_volume_in(:)/dep_area_in_cell_in(:)
call print_msg('Lake per-cell depth [m] (min|max): '//real2str(minval(lake_depth_out))//' | '//real2str(maxval(lake_depth_out)),LVL_NFO)

END SUBROUTINE compute_lake_depth
!=======================================================================

!=======================================================================
SUBROUTINE evolve_hydrology(d_precip,d_evap,d_runoff,stopcrit)
!-----------------------------------------------------------------------
! NAME
!     evolve_hydrology
!
! DESCRIPTION
!     Evolve the surface water reservoirs.
!
! AUTHORS & DATE
!     A. Gauvain, 05/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use stopping_crit,  only: stopFlags
use hydro_core,     only: run_simulation
use stopping_crit,  only: stopping_crit_lakes
use hydro_database, only: ndep, child_list
use hydro_tree,     only: nodes

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

! ARGUMENTS
! ---------
real(dp), dimension(:,:), intent(in)    :: d_precip, d_evap, d_runoff
type(stopFlags),          intent(inout) :: stopcrit

! LOCAL VARIABLES
! ---------------
integer(di) :: i, ig

! CODE
! ----
!$OMP PARALLEL DO PRIVATE(ig) SCHEDULE(static)
do i = 1,ndep
    if (nodes(child_list(i))%is_leaf) then
        ig = nodes(child_list(i))%gcm_idx
        nodes(child_list(i))%E = d_evap(ig,1)
        nodes(child_list(i))%P = d_precip(ig,1)
        nodes(child_list(i))%Runoff = d_runoff(ig,1)
    end if
end do
!$OMP END PARALLEL DO

call run_simulation(starthydro_name)

call compute_lake_coverage(lake_frac)
call compute_depression_volume(lake_volume)
call compute_lake_depth(lake_volume,dep_area_in_cell,lake_depth)

call stopping_crit_lakes(peren_h2o_liq_frac_PCM,lake_frac,stopcrit)

END SUBROUTINE evolve_hydrology
!=======================================================================

!=======================================================================
SUBROUTINE compute_lake_water_mass_per_area(lake_water_mass_per_area)
!-----------------------------------------------------------------------
! NAME
!     compute_lake_water_mass_per_area
!
! DESCRIPTION
!     Convert the PEM hydrology reservoir volume in each grid cell into a
!     lake-area-normalized liquid-water mass [kg/m2].
!
! AUTHORS & DATE
!     C. Metz, 06/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry, only: ngrid

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

! ARGUMENTS
! ---------
real(dp), dimension(:), intent(out) :: lake_water_mass_per_area

! LOCAL VARIABLES
! ---------------
integer(di) :: ig

! CODE
! ----
lake_water_mass_per_area(:) = 0._dp
do ig = 1,ngrid
    if (lake_area(ig) > 0._dp) then
        lake_water_mass_per_area(ig) = lake_volume(ig)*h2o_liq_density/lake_area(ig)
    end if
end do

END SUBROUTINE compute_lake_water_mass_per_area
!=======================================================================

END MODULE hydrology