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

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

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

! PARAMETERS
! ----------
integer(di), parameter :: WEIGHT_UNIFORM   = 1_di ! Equal weight for all adjustable points
integer(di), parameter :: WEIGHT_ABSFLUX   = 2_di ! Weight by absolute local H2O surface flux
integer(di), parameter :: WEIGHT_CAPACITY  = 3_di ! Weight by available local adjustment room
integer(di), parameter :: WEIGHT_COMBINED  = 4_di ! Weight by current absolute flux and adjustment room
integer(di), parameter :: WEIGHT_RESERVOIR = 5_di ! Weight by locally available H2O ice mass
integer(di), parameter :: WEIGHT_DIRECTION = 6_di ! Weight by directional capacity (up/down)
integer(di), private   :: weight_method = WEIGHT_ABSFLUX ! Default method for balancing
logical(k4), private   :: correct_condens_only = .true.  ! Correct only the reservoirs receiving water from the atmosphere

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

!=======================================================================
SUBROUTINE balance_h2o_atm_fluxes(delta_h2o_ads,delta_icetable,h2o_ice,d_h2oice_ini,d_h2oice,stopcrit,weight_type,condens_only)
!-----------------------------------------------------------------------
! NAME
!     balance_h2o_atm_fluxes
!
! DESCRIPTION
!     Balance H2O fluxes from and into atmosphere across reservoirs.
!
! AUTHORS & DATE
!     JB Clement, 04/2025-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 reservoirs:
!         - the fluxes leaving the surface (surface ice sublimation) are
!           the forcing because they are set by the radiative surface
!           balance of the PCM. They are only bounded by the ice actually
!           available;
!         - the fluxes leaving the atmosphere (surface ice condensation)
!           are the only ones to be corrected.
!     The option 'correct_condens_only' enforces this behavior. Otherwise,
!     the correction is spread over every reservoir, the sublimating ones
!     included.
!
!     The redistribution of the residual imbalance is additive, done
!     iteratively with a weight given by 'weight_method'. No physical
!     bound can be broken because a reservoir reaching one of its bounds
!     is deactivated.
!     The method 'WEIGHT_ABSFLUX' is the default because it reduces to a
!     multiplicative closure factor with 'correct_condens_only = .true.'
!     so that the deposition pattern predicted by the PCM is preserved,
!     converges in one iteration and respects the physical bounds by
!     construction.
!
!     Known limitations of the closure:
!       - the ice table and the adsorbed regolith exchange H2O with the
!         atmosphere without being compensated here.
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use evolution,     only: get_dt_yr
use geometry,      only: ngrid, nslope, cell_area
use slopes,        only: slope_weight
use stopping_crit, only: stopFlags
use utility,       only: real2str, int2str
use display,       only: print_msg, LVL_NFO, LVL_WRN

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

! ARGUMENTS
! ---------
real(dp), dimension(:),   intent(in)           :: delta_h2o_ads, delta_icetable
real(dp), dimension(:,:), intent(in)           :: h2o_ice
real(dp), dimension(:,:), intent(in)           :: d_h2oice_ini ! Reference tendency given by the PCM [kg/m2/y]
integer(di),              intent(in), optional :: weight_type
logical(k4),              intent(in), optional :: condens_only
real(dp), dimension(:,:), intent(out)          :: d_h2oice     ! Effective tendency of the current time step [kg/m2/y]
type(stopFlags),          intent(inout)        :: stopcrit

! LOCAL VARIABLES
! ---------------
integer(di), parameter :: max_iter = 50_di ! Maximum number of iterations for the balancing procedure
integer(di)                              :: i, islope, iter
integer(di)                              :: weight_type_used
logical(k4)                              :: condens_only_used
real(dp)                                 :: dt_yr ! Time step in Planetary years
real(dp)                                 :: flux_scale, tol_flux
real(qp)                                 :: B_flux, B_flux_surfice, B_flux_nonsurfice, Delta, Delta_used, wsum
real(qp)                                 :: B_flux_raw ! Raw budget of H2O exchange with the atmosphere [kg/y]
real(qp)                                 :: M_subl     ! H2O mass rate sublimating from the surface ice [kg/y]
real(qp)                                 :: M_cond     ! H2O mass rate condensing onto the surface ice [kg/y]
real(dp),    dimension(:,:), allocatable :: w, d_corr, flux_new
real(dp),    dimension(:,:), allocatable :: flux_min, flux_max
real(dp),    dimension(:,:), allocatable :: capacity, capacity_up, capacity_down
real(dp),    dimension(:,:), allocatable :: area ! True area of each sub-slope [m2]
logical(k4), dimension(:,:), allocatable :: active, adjustable, is_condensing

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

! Start over from the reference tendency (per-timestep correction)
d_h2oice(:,:) = d_h2oice_ini(:,:)

! In 1D, the surface ice is the only reservoir and there is no other grid point to balance with
if (ngrid == 1) then
    ! Make sure we don't sublimate more ice than available
    where (d_h2oice(:,:) < 0._dp .and. h2o_ice(:,:) + d_h2oice(:,:)*dt_yr < 0._dp) d_h2oice(:,:) = -h2o_ice(:,:)/dt_yr
    call print_msg('Only one surface ice reservoir in 1D: the H2O fluxes are only clipped by the available ice.',LVL_NFO)
    return
end if

allocate(w(ngrid,nslope),d_corr(ngrid,nslope),flux_new(ngrid,nslope))
allocate(flux_min(ngrid,nslope),flux_max(ngrid,nslope))
allocate(capacity(ngrid,nslope),capacity_up(ngrid,nslope),capacity_down(ngrid,nslope))
allocate(area(ngrid,nslope))
allocate(active(ngrid,nslope),adjustable(ngrid,nslope),is_condensing(ngrid,nslope))

! Build physical bounds and geometry weights
B_flux_raw = 0._qp
do i = 1,ngrid
    do islope = 1,nslope
        ! The balance has to be closed on masses, not on surface densities, so every flux is area-weighted
        area(i,islope) = cell_area(i)*slope_weight(i,islope)
        ! Raw imbalance before any bound or correction
        B_flux_raw = B_flux_raw + abs(real(d_h2oice(i,islope)*area(i,islope),qp))
        flux_min(i,islope) = -h2o_ice(i,islope)/dt_yr ! Maximum sublimation rate (limited by available ice)
        flux_max(i,islope) = 1.e10_dp ! Maximum deposition rate (unlimited)
        if (d_h2oice(i,islope) > 0._dp) then
            flux_min(i,islope) = 0._dp
        else if (d_h2oice(i,islope) < 0._dp) then
            flux_max(i,islope) = 0._dp
        end if
        capacity(i,islope) = flux_max(i,islope) - flux_min(i,islope)
    end do
end do
active(:,:) = .true.

! Enforce physical bounds before starting redistribution
where (d_h2oice(:,:) > flux_max(:,:))
    d_h2oice(:,:) = flux_max(:,:)
else where (d_h2oice(:,:) < flux_min(:,:))
    d_h2oice(:,:) = flux_min(:,:)
end where

! Choose weighting method
weight_type_used = weight_method
if (present(weight_type)) then
    if (weight_type >= WEIGHT_UNIFORM .and. weight_type <= WEIGHT_DIRECTION) then
        weight_method = weight_type
        weight_type_used = weight_type
    else
        call print_msg('Unknown H2O flux weighting method. Falling back to WEIGHT_ABSFLUX.',LVL_WRN)
        weight_type_used = WEIGHT_ABSFLUX
    end if
end if

! Choose the reservoirs on which the correction is applied
condens_only_used = correct_condens_only
if (present(condens_only)) then
    correct_condens_only = condens_only
    condens_only_used = condens_only
end if

! Compute initial imbalance
B_flux_nonsurfice = 0._qp
M_subl = 0._qp
M_cond = 0._qp
do i = 1,ngrid
    ! Non-surface ice fluxes (adsorption and icetable) are not included because they can lead to unrealistic adjustments of surface ice reservoirs
    ! if we try to balance them together. For instance, when using the equilibrium computation of ice table, the initial imbalance can be very large.
    !B_flux_nonsurfice = B_flux_nonsurfice + real((delta_h2o_ads(i) + delta_icetable(i))*cell_area(i),qp)/real(dt_yr,qp)
    do islope = 1,nslope
        is_condensing(i,islope) = d_h2oice(i,islope) > 0._dp
        if (is_condensing(i,islope)) then
            M_cond = M_cond + real(d_h2oice(i,islope)*area(i,islope),qp)
        else
            M_subl = M_subl - real(d_h2oice(i,islope)*area(i,islope),qp)
        end if
    end do
end do
B_flux_surfice = M_cond - M_subl
B_flux = B_flux_nonsurfice + B_flux_surfice
Delta = -B_flux
flux_scale = max(1._dp,abs(real(B_flux,dp))) ! Scale flux imbalance to get a relevant tolerance for the balancing procedure (e.g., if imbalance is very small, we don't want to require an even smaller imbalance to stop iterating)
tol_flux = max(1.e-12_dp,1.e-10_dp*flux_scale) ! Tolerance for closing the flux balance, scaled to the initial imbalance

! Iterative correction
do iter = 1,max_iter
    ! Stopping criterion
    if (abs(Delta) <= real(tol_flux,qp)) exit

    ! Prepare correction
    w(:,:) = 0._dp
    ! Compute directional capacities
    capacity_up(:,:) = flux_max(:,:) - d_h2oice(:,:)
    capacity_down(:,:) = d_h2oice(:,:) - flux_min(:,:)

    ! Compute adjustable mask
    adjustable(:,:) = active(:,:)
    if (condens_only_used) then
        where (.not. is_condensing(:,:)) adjustable(:,:) = .false.
    end if
    if (Delta > 0._qp) then
        where (capacity_up <= eps) adjustable = .false.
    else
        where (capacity_down <= eps) adjustable = .false.
    end if
    if (.not. any(adjustable)) exit

    ! Compute weights
    select case (weight_type_used)
        case (WEIGHT_UNIFORM) ! Flat redistribution across adjustable cells
            where (adjustable) w = 1._dp
        case (WEIGHT_ABSFLUX) ! Prioritize cells already carrying strong fluxes
            where (adjustable) w = abs(d_h2oice)
        case (WEIGHT_CAPACITY) ! Prioritize cells with large admissible adjustment room
            where (adjustable) w = capacity
        case (WEIGHT_COMBINED) ! Mix current activity (abs flux) and admissible room
            where (adjustable) w = capacity*max(abs(d_h2oice),eps)
        case (WEIGHT_RESERVOIR) ! Favor points with larger local ice reservoirs
            where (adjustable) w = max(h2o_ice,eps)
        case (WEIGHT_DIRECTION) ! Follow directional room needed by the sign of Delta
            if (Delta > 0._qp) then
                where (adjustable) w = max(capacity_up,eps)
            else
                where (adjustable) w = max(capacity_down,eps)
            end if
        case default
            where (adjustable) w = abs(d_h2oice)
    end select

    wsum = sum(real(w,qp)*real(area,qp),mask = adjustable)
    if (wsum <= 0._qp) exit ! Fallback all weights are zero

    ! Apply correction
    flux_new(:,:) = d_h2oice(:,:)
    where (adjustable) flux_new(:,:) = d_h2oice(:,:) + real(Delta*real(w,qp)/wsum,dp)

    ! Apply bounds
    where (flux_new > flux_max)
        flux_new(:,:) = flux_max(:,:)
        where (adjustable) active(:,:) = .false.
    else where (flux_new < flux_min)
        flux_new(:,:) = flux_min(:,:)
        where (adjustable) active(:,:) = .false.
    end where

    ! Compute effective correction and update fluxes
    d_corr(:,:) = flux_new(:,:) - d_h2oice(:,:)
    d_h2oice(:,:) = flux_new(:,:)

    ! Update imbalance
    Delta_used = sum(real(d_corr,qp)*real(area,qp))
    Delta = Delta - Delta_used

    ! Test if the balancing procedure progresses
    if (abs(Delta_used) < real(tol_flux,qp)) exit
end do

! Diagnostics
call print_msg('Surface ice sublimation   [kg/y] = '//real2str(M_subl),LVL_NFO)
call print_msg('Surface ice condensation  [kg/y] = '//real2str(M_cond),LVL_NFO)
call print_msg('Initial raw imbalance     [kg/y] = '//real2str(B_flux_raw),LVL_NFO)
call print_msg('Initial bounded imbalance [kg/y] = '//real2str(B_flux),LVL_NFO)
call print_msg('Residual imbalance        [kg/y] = '//real2str(Delta),LVL_NFO)
call print_msg('Number of balancing iterations   = '//int2str(iter),LVL_NFO)
call print_msg('Number of condensing reservoirs  = '//int2str(count(is_condensing)),LVL_NFO)
call print_msg('Number of deactivated reservoirs = '//int2str(count(.not. active)),LVL_NFO)
if (condens_only_used) then
    call print_msg('The correction is restricted to the condensing reservoirs.',LVL_NFO)
else
    call print_msg('The correction is spread over all the surface ice reservoirs.',LVL_NFO)
end if

if (abs(Delta) > real(tol_flux,qp)) then
    call print_msg('Unable to close the water flux balance on surface ice reservoirs.',LVL_WRN)
    call print_msg('Weighting method used = '//int2str(weight_type_used),LVL_WRN)
    call print_msg('Surface ice     part of the imbalance [kg/y] = '//real2str(B_flux_surfice),LVL_WRN)
    call print_msg('Non-surface ice part of the imbalance [kg/y] = '//real2str(B_flux_nonsurfice),LVL_WRN)
    if (condens_only_used .and. count(is_condensing) == 0) call print_msg('Water leaves the atmosphere nowhere on the planet: the correction has no reservoir to be applied on!',LVL_WRN)
    stopcrit%h2o_flux_balance_unclosed = .true.
end if

! Cleanup
deallocate(w,d_corr,flux_new,flux_min,flux_max)
deallocate(capacity,capacity_up,capacity_down,area)
deallocate(active,adjustable,is_condensing)

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

END MODULE water_balance