MODULE water_balance
!-----------------------------------------------------------------------
! NAME
!     water_balance
!
! DESCRIPTION
!     Balance the water exchanges between the atmosphere and all the surface
!     water reservoirs of the planet.
!
! AUTHORS & DATE
!     JB Clement, 08/2026
!
! NOTES
!
!-----------------------------------------------------------------------

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

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

! PARAMETERS
! ----------
! Range outside which the closure factor is considered suspicious
real(qp), parameter, private :: beta_min_warn = 0.1_qp
real(qp), parameter, private :: beta_max_warn = 10._qp

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

!=======================================================================
SUBROUTINE balance_h2o_atm_fluxes(M_subl,M_evap,M_cond,M_precip,M_runoff,beta,stopcrit)
!-----------------------------------------------------------------------
! NAME
!     balance_h2o_atm_fluxes
!
! DESCRIPTION
!     Compute the closure factor of the H2O fluxes leaving the atmosphere.
!
! AUTHORS & DATE
!     JB Clement, 08/2026
!
! NOTES
!     The PEM does not carry the atmospheric water column. It assumes it
!     stays constant over a time step, so the net exchange between the
!     atmosphere and the surface must vanish.
!
!     The closure is therefore global so that water can migrate between
!     the surface ice and the lakes:
!         - the fluxes leaving the surface (surface ice sublimation, lake
!           evaporation) are the forcing because they are set by the
!           radiative surface balance of the PCM. They are only bounded
!           by the water actually available;
!         - the fluxes leaving the atmosphere (surface ice condensation,
!           lake precipitation, land runoff) are the only ones to be
!           corrected.
!
!     A single closure factor is used so the redistribution is multiplicative.
!     Thus, the deposition pattern predicted by the PCM is preserved. No
!     physical bound can be broken because the closure factor is applied
!     once to all the fluxes leaving the atmosphere.
!
!     Known limitations of the closure:
!       - the soil water bucket is not a reservoir of the PEM;
!       - the evaporation of a lake is evaluated on the water it holds
!         before precipitation and runoff are added. So a nearly empty
!         lake under a heavy rain under-evaporates;
!       - a lake kept fixed (e.g. a prescribed ocean) is an untracked
!         infinite source/sink: it is excluded from the closure.
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry,      only: ngrid
use hydrology,     only: do_hydrology
use stopping_crit, only: stopFlags
use utility,       only: real2str
use display,       only: print_msg, LVL_NFO, LVL_WRN

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

! ARGUMENTS
! ---------
real(qp),        intent(in)    :: M_subl   ! Water mass sublimating from the surface ice [kg]
real(qp),        intent(in)    :: M_evap   ! Water mass evaporating from the lakes [kg]
real(qp),        intent(in)    :: M_cond   ! Water mass condensing onto the surface ice [kg]
real(qp),        intent(in)    :: M_precip ! Water mass precipitating onto the lakes [kg]
real(qp),        intent(in)    :: M_runoff ! Water mass running off into the lakes [kg]
real(qp),        intent(out)   :: beta     ! Closure factor of the fluxes leaving the atmosphere
type(stopFlags), intent(inout) :: stopcrit

! LOCAL VARIABLES
! ---------------
real(qp) :: M_in  ! Water mass entering the atmosphere [kg]
real(qp) :: M_out ! Water mass leaving the atmosphere [kg]

! CODE
! ----
call print_msg('> Balancing water fluxes with the atmosphere',LVL_NFO)

! In 1D, the surface ice should be the only reservoir and there is no other grid point to balance with
if (ngrid == 1 .and. .not. do_hydrology) then
    beta = 1._qp
    call print_msg('Only one surface ice reservoir in 1D: the water closure factor is kept at 1.',LVL_NFO)
    return
end if

M_in = M_subl + M_evap
M_out = M_cond + M_precip + M_runoff

if (M_out > 0._qp) then
    beta = M_in/M_out
else
    beta = 1._qp
    if (M_in > 0._qp) then
        call print_msg('Water leaves the atmosphere nowhere on the planet: the water closure factor cannot be computed!',LVL_WRN)
        stopcrit%h2o_flux_balance_unclosed = .true.
    else
        call print_msg('No water exchange with the atmosphere: the water closure factor is kept at 1.',LVL_NFO)
    end if
end if

! Diagnostics
call print_msg('Surface ice sublimation          [kg] = '//real2str(real(M_subl,dp)),LVL_NFO)
call print_msg('Surface ice condensation         [kg] = '//real2str(real(M_cond,dp)),LVL_NFO)
call print_msg('Lake evaporation                 [kg] = '//real2str(real(M_evap,dp)),LVL_NFO)
call print_msg('Lake precipitation               [kg] = '//real2str(real(M_precip,dp)),LVL_NFO)
call print_msg('Runoff into the lakes            [kg] = '//real2str(real(M_runoff,dp)),LVL_NFO)
call print_msg('H2O mass entering the atmosphere [kg] = '//real2str(real(M_in,dp)),LVL_NFO)
call print_msg('H2O mass leaving the atmosphere  [kg] = '//real2str(real(M_out,dp)),LVL_NFO)
call print_msg('Water closure factor beta             = '//real2str(real(beta,dp)),LVL_NFO)

if (beta < beta_min_warn .or. beta > beta_max_warn) call print_msg('The water closure factor is far from 1: the water cycle of the PCM is strongly out of balance!',LVL_WRN)

END SUBROUTINE balance_h2o_atm_fluxes
!=======================================================================

END MODULE water_balance