MODULE mass_balance
!-----------------------------------------------------------------------
! NAME
!     mass_balance
!
! DESCRIPTION
!     Accounting of the mass budget of a condensable species over a PEM
!     run.
!
! AUTHORS & DATE
!     JB Clement, 08/2026
!
! NOTES
!     The budget of a condensable species is described by:
!         - the "reservoirs", which hold an absolute mass [kg] that is
!           re-evaluated at each PEM iteration (surface ice, atmosphere,
!           lake water, adsorbed regolith, etc);
!         - the "exchanges" with the atmosphere, which are cumulated
!           masses [kg] counted positively when they are removed from the
!           atmosphere. They are only needed when the atmospheric content
!           of the species is not itself a tracked reservoir. For example,
!           the PEM does not know the atmospheric water column and assumes
!           (enforces) that it stays constant.
!     Two quantities are then monitored:
!         - the drift = sum(reservoirs) - sum(reservoirs_ini) -
!           sum(exchanges), which must be zero. It measures the amount of
!           mass that the model created or destroyed;
!         - the net exchange = sum(exchanges), which must be zero. if
!           the atmospheric content is to remain constant.
!     For CO2, the atmosphere is a reservoir and there is no exchange
!     term, so the drift alone describes the budget.
!-----------------------------------------------------------------------

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

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

! PARAMETERS
! ----------
integer(di),  parameter           :: nterm_max = 10_di                ! Maximum number of reservoirs/exchanges in a budget
integer(di),  parameter, private  :: len_name = 20_di                 ! Length of the name of a reservoir/exchange
integer(di),  parameter, private  :: len_label = 30_di                ! Length of the left-justified label of a reported line
character(9), parameter, private  :: fmt_mass = '(es16.8)'            ! Format to display a mass [kg]
character(9), parameter, private  :: fmt_pct = '(f16.6)'              ! Format to display a percentage [%]
real(dp),     protected           :: tol_mass_balance = 0.01_dp       ! Relative tolerance on the mass balance drift [-]
logical(k4),  protected           :: stop_on_mass_imbalance = .false. ! Flag to stop the PEM when the drift exceeds the tolerance

! TYPES
! -----
type, private :: budget_term
    character(len_name) :: name = ''     ! Name of the term
    real(qp)            :: m_ini = 0._qp ! Initial mass [kg] (reservoirs only)
    real(qp)            :: m = 0._qp     ! Current mass or cumulated exchange [kg]
end type budget_term

type mass_budget
    character(8)                            :: species = ''         ! Name of the condensable species ('CO2', 'H2O', etc)
    logical(k4)                             :: check_drift = .true. ! Flag to test the drift against the tolerance
    integer(di)                             :: nresvr = 0_di        ! Number of declared reservoirs
    integer(di)                             :: nexch = 0_di         ! Number of declared exchanges
    type(budget_term), dimension(nterm_max) :: resvr                ! Reservoirs
    type(budget_term), dimension(nterm_max) :: exch                 ! Exchanges with the atmosphere (> 0 <=> removed from the atmosphere)
contains
    procedure          :: add_resvr  => mass_budget_add_resvr
    procedure          :: add_exch   => mass_budget_add_exch
    procedure          :: set_resvr  => mass_budget_set_resvr
    procedure          :: accum_exch => mass_budget_accum_exch
    procedure          :: check      => mass_budget_check
    procedure          :: report     => mass_budget_report
    procedure          :: write_diag => mass_budget_write_diag
    procedure, private :: totals     => mass_budget_totals
end type mass_budget

! INTERFACES
! ----------
interface mass_surf
    module procedure mass_surf_slope, mass_surf_cell
end interface mass_surf

private :: mass_surf_slope, mass_surf_cell

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

!=======================================================================
SUBROUTINE set_mass_balance_config(tol_mass_balance_in,stop_on_mass_imbalance_in)
!-----------------------------------------------------------------------
! NAME
!     set_mass_balance_config
!
! DESCRIPTION
!     Setter for "mass_balance" configuration parameters.
!
! AUTHORS & DATE
!     JB Clement, 08/2026
!
! NOTES
!
!-----------------------------------------------------------------------

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

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

! ARGUMENTS
! ---------
real(dp),    intent(in) :: tol_mass_balance_in
logical(k4), intent(in) :: stop_on_mass_imbalance_in

! CODE
! ----
tol_mass_balance = tol_mass_balance_in
stop_on_mass_imbalance = stop_on_mass_imbalance_in
if (tol_mass_balance <= 0._dp) call stop_clean(__FILE__,__LINE__,'''tol_mass_balance'' must be positive!',1)

call print_msg('tol_mass_balance       = '//real2str(tol_mass_balance),LVL_NFO)
call print_msg('stop_on_mass_imbalance = '//bool2str(stop_on_mass_imbalance),LVL_NFO)

END SUBROUTINE set_mass_balance_config
!=======================================================================

!=======================================================================
FUNCTION mass_surf_slope(field) RESULT(m)
!-----------------------------------------------------------------------
! NAME
!     mass_surf_slope
!
! DESCRIPTION
!     Total mass held by a surface field defined on the sub-slopes.
!
! AUTHORS & DATE
!     JB Clement, 08/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry, only: ngrid, nslope, cell_area
use slopes,   only: slope_weight

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

! ARGUMENTS
! ---------
real(dp), dimension(:,:), intent(in) :: field ! Surface density of the species [kg/m2]

! RESULT
! ------
real(qp) :: m ! Total mass [kg]

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

! CODE
! ----
m = 0._qp
do i = 1,ngrid
    do islope = 1,nslope
        m = m + real(field(i,islope)*cell_area(i)*slope_weight(i,islope),qp)
    end do
end do

END FUNCTION mass_surf_slope
!=======================================================================

!=======================================================================
FUNCTION mass_surf_cell(field) RESULT(m)
!-----------------------------------------------------------------------
! NAME
!     mass_surf_cell
!
! DESCRIPTION
!     Total mass held by a surface field defined on the grid cell.
!
! AUTHORS & DATE
!     JB Clement, 08/2026
!
! NOTES
!
!-----------------------------------------------------------------------

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

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

! ARGUMENTS
! ---------
real(dp), dimension(:), intent(in) :: field ! Surface density of the species [kg/m2]

! RESULT
! ------
real(qp) :: m ! Total mass [kg]

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

! CODE
! ----
m = 0._qp
do i = 1,ngrid
    m = m + real(field(i)*cell_area(i),qp)
end do

END FUNCTION mass_surf_cell
!=======================================================================

!=======================================================================
FUNCTION mass_atm(ps) RESULT(m)
!-----------------------------------------------------------------------
! NAME
!     mass_atm
!
! DESCRIPTION
!     Total atmospheric mass deduced from the surface pressure.
!
! AUTHORS & DATE
!     JB Clement, 08/2026
!
! NOTES
!     The whole atmospheric column is assumed to be made of the species.
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry, only: ngrid, cell_area
use physics,  only: g

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

! ARGUMENTS
! ---------
real(dp), dimension(:), intent(in) :: ps ! Surface pressure [Pa]

! RESULT
! ------
real(qp) :: m ! Total mass [kg]

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

! CODE
! ----
m = 0._qp
do i = 1,ngrid
    m = m + real(cell_area(i)*ps(i)/g,qp)
end do

END FUNCTION mass_atm
!=======================================================================

!=======================================================================
FUNCTION find_term(terms,n,name) RESULT(k)
!-----------------------------------------------------------------------
! NAME
!     find_term
!
! DESCRIPTION
!     Index of the budget term for a given name.
!
! AUTHORS & DATE
!     JB Clement, 08/2026
!
! NOTES
!     The returned index is 0 when the name is not declared.
!-----------------------------------------------------------------------

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

! ARGUMENTS
! ---------
type(budget_term), dimension(:), intent(in) :: terms ! Declared terms
integer(di),                     intent(in) :: n     ! Number of declared terms
character(*),                    intent(in) :: name  ! Name of the term

! RESULT
! ------
integer(di) :: k ! Index of the term

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

! CODE
! ----
k = 0_di
do i = 1,n
    if (trim(terms(i)%name) == trim(name)) then
        k = i
        return
    end if
end do

END FUNCTION find_term
!=======================================================================

!=======================================================================
SUBROUTINE mass_budget_add_resvr(this,name,m)
!-----------------------------------------------------------------------
! NAME
!     mass_budget_add_resvr
!
! DESCRIPTION
!     Declare a reservoir in a budget and record its initial mass.
!
! AUTHORS & DATE
!     JB Clement, 08/2026
!
! NOTES
!
!-----------------------------------------------------------------------

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

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

! ARGUMENTS
! ---------
class(mass_budget), intent(inout) :: this
character(*),       intent(in)    :: name ! Name of the reservoir
real(qp),           intent(in)    :: m    ! Initial mass [kg]

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

! CODE
! ----
if (this%nresvr >= nterm_max) call stop_clean(__FILE__,__LINE__,'too many reservoirs in the '//trim(this%species)//' mass budget!',1)
if (len_trim(name) > len_name) call stop_clean(__FILE__,__LINE__,'the name of the reservoir "'//trim(name)//'" in the '//trim(this%species)//' mass budget exceeds '//int2str(len_name)//' characters!',1)
if (find_term(this%resvr,this%nresvr,name) > 0) call stop_clean(__FILE__,__LINE__,'duplicate reservoir "'//trim(name)//'" in the '//trim(this%species)//' mass budget!',1)
this%nresvr = this%nresvr + 1
k = this%nresvr
this%resvr(k)%name = name
this%resvr(k)%m_ini = m
this%resvr(k)%m = m
call print_msg('Initial '//trim(this%species)//' mass in "'//trim(name)//'" [kg] = '//real2str(m,fmt_mass),LVL_NFO)

END SUBROUTINE mass_budget_add_resvr
!=======================================================================

!=======================================================================
SUBROUTINE mass_budget_add_exch(this,name)
!-----------------------------------------------------------------------
! NAME
!     mass_budget_add_exch
!
! DESCRIPTION
!     Declare an exchange term with the atmosphere in a budget.
!
! AUTHORS & DATE
!     JB Clement, 08/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use utility,  only: int2str
use stoppage, only: stop_clean

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

! ARGUMENTS
! ---------
class(mass_budget), intent(inout) :: this
character(*),       intent(in)    :: name ! Name of the exchange

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

! CODE
! ----
if (this%nexch >= nterm_max) call stop_clean(__FILE__,__LINE__,'too many exchanges in the '//trim(this%species)//' mass budget!',1)
if (len_trim(name) > len_name) call stop_clean(__FILE__,__LINE__,'the name of the exchange "'//trim(name)//'" in the '//trim(this%species)//' mass budget exceeds '//int2str(len_name)//' characters!',1)
if (find_term(this%exch,this%nexch,name) > 0) call stop_clean(__FILE__,__LINE__,'duplicate exchange "'//trim(name)//'" in the '//trim(this%species)//' mass budget!',1)
this%nexch = this%nexch + 1
k = this%nexch
this%exch(k)%name = name
this%exch(k)%m = 0._qp

END SUBROUTINE mass_budget_add_exch
!=======================================================================

!=======================================================================
SUBROUTINE mass_budget_set_resvr(this,name,m)
!-----------------------------------------------------------------------
! NAME
!     mass_budget_set_resvr
!
! DESCRIPTION
!     Set the current mass of a declared reservoir.
!
! AUTHORS & DATE
!     JB Clement, 08/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use stoppage, only: stop_clean

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

! ARGUMENTS
! ---------
class(mass_budget), intent(inout) :: this
character(*),       intent(in)    :: name ! Name of the reservoir
real(qp),           intent(in)    :: m    ! Current mass [kg]

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

! CODE
! ----
k = find_term(this%resvr,this%nresvr,name)
if (k == 0) call stop_clean(__FILE__,__LINE__,'unknown reservoir "'//trim(name)//'" in the '//trim(this%species)//' mass budget!',1)
this%resvr(k)%m = m

END SUBROUTINE mass_budget_set_resvr
!=======================================================================

!=======================================================================
SUBROUTINE mass_budget_accum_exch(this,name,dm)
!-----------------------------------------------------------------------
! NAME
!     mass_budget_accum_exch
!
! DESCRIPTION
!     Accumulate an increment into a declared exchange term.
!
! AUTHORS & DATE
!     JB Clement, 08/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use stoppage, only: stop_clean

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

! ARGUMENTS
! ---------
class(mass_budget), intent(inout) :: this
character(*),       intent(in)    :: name ! Name of the exchange
real(qp),           intent(in)    :: dm   ! Increment over the PEM time step [kg] (> 0 <=> removed from the atmosphere)

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

! CODE
! ----
k = find_term(this%exch,this%nexch,name)
if (k == 0) call stop_clean(__FILE__,__LINE__,'unknown exchange "'//trim(name)//'" in the '//trim(this%species)//' mass budget!',1)
this%exch(k)%m = this%exch(k)%m + dm

END SUBROUTINE mass_budget_accum_exch
!=======================================================================

!=======================================================================
SUBROUTINE mass_budget_totals(this,m_ini,m,net,drift,rel_drift)
!-----------------------------------------------------------------------
! NAME
!     mass_budget_totals
!
! DESCRIPTION
!     Sum the quantities describing a budget.
!
! AUTHORS & DATE
!     JB Clement, 08/2026
!
! NOTES
!
!-----------------------------------------------------------------------

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

! ARGUMENTS
! ---------
class(mass_budget), intent(in)  :: this
real(qp),           intent(out) :: m_ini     ! Initial total mass held by the reservoirs [kg]
real(qp),           intent(out) :: m         ! Current total mass held by the reservoirs [kg]
real(qp),           intent(out) :: net       ! Net mass removed from the atmosphere since the beginning of the run [kg]
real(qp),           intent(out) :: drift     ! Mass created or destroyed by the model without being accounted for [kg]
real(qp),           intent(out) :: rel_drift ! Drift relative to the initial total mass [-]

! CODE
! ----
m_ini = sum(this%resvr(1:this%nresvr)%m_ini)
m = sum(this%resvr(1:this%nresvr)%m)
net = sum(this%exch(1:this%nexch)%m)
drift = m - m_ini - net
rel_drift = drift/max(m_ini,eps_qp) ! To avoid division by 0

END SUBROUTINE mass_budget_totals
!=======================================================================

!=======================================================================
FUNCTION mass_budget_check(this) RESULT(must_stop)
!-----------------------------------------------------------------------
! NAME
!     mass_budget_check
!
! DESCRIPTION
!     Report the mass balance of a budget and test it against the
!     tolerance.
!
! AUTHORS & DATE
!     JB Clement, 08/2026
!
! NOTES
!
!-----------------------------------------------------------------------

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

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

! ARGUMENTS
! ---------
class(mass_budget), intent(in) :: this

! RESULT
! ------
logical(k4) :: must_stop ! True if the drift exceeds the tolerance and the PEM has to stop

! LOCAL VARIABLES
! ---------------
real(qp)             :: m_ini, m, net, drift, rel_drift
character(len_label) :: label ! Left-justified and blank-padded label

! CODE
! ----
call this%totals(m_ini,m,net,drift,rel_drift)

call print_msg('# Mass balance for '//trim(this%species)//':',LVL_NFO)
call print_msg('### Total mass             (ini|now) [kg] = '//real2str(m_ini,fmt_mass)//' | '//real2str(m,fmt_mass),LVL_NFO)
if (this%nexch > 0) call print_msg('### Net exchange with the atmosphere [kg] = '//real2str(net,fmt_mass),LVL_NFO)
call print_msg('### Drift          (accepted|current) [%] = +/- '//real2str(100._dp*tol_mass_balance,fmt_pct)//' | '//real2str(100._qp*rel_drift,fmt_pct),LVL_NFO)

must_stop = .false.
if (this%check_drift .and. abs(rel_drift) > real(tol_mass_balance,qp)) then
    call print_msg('The '//trim(this%species)//' mass balance is not conserved!',LVL_WRN)
    call this%report(LVL_WRN)
    must_stop = stop_on_mass_imbalance
end if

END FUNCTION mass_budget_check
!=======================================================================

!=======================================================================
SUBROUTINE mass_budget_report(this,lvl)
!-----------------------------------------------------------------------
! NAME
!     mass_budget_report
!
! DESCRIPTION
!     Print each term of a budget to detail an imbalance.
!
! AUTHORS & DATE
!     JB Clement, 08/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use display, only: print_msg
use utility, only: real2str

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

! ARGUMENTS
! ---------
class(mass_budget), intent(in) :: this
integer(di),        intent(in) :: lvl ! Verbosity level of the report

! LOCAL VARIABLES
! ---------------
integer(di)          :: k
real(qp)             :: m_ini, m, net, drift, rel_drift
character(len_label) :: label ! Left-justified and blank-padded label

! CODE
! ----
call this%totals(m_ini,m,net,drift,rel_drift)

call print_msg('# Mass budget for '//trim(this%species)//':',lvl)
do k = 1,this%nresvr
    label = trim(this%resvr(k)%name)
    call print_msg('### '//label//' (ini|now) [kg] = '//real2str(this%resvr(k)%m_ini,fmt_mass)//' | '//real2str(this%resvr(k)%m,fmt_mass),lvl)
end do
call print_msg('### TOTAL (ini|now) [kg] = '//real2str(m_ini,fmt_mass)//' | '//real2str(m,fmt_mass),lvl)

if (this%nexch > 0) then
    do k = 1,this%nexch
        label = trim(this%exch(k)%name)
        call print_msg('### '//label//' [kg] = '//real2str(this%exch(k)%m,fmt_mass),lvl)
    end do
    call print_msg('### NET exchange      [kg] = '//real2str(net,fmt_mass),lvl)
end if

call print_msg('### DRIFT           [kg] = '//real2str(drift,fmt_mass),lvl)
call print_msg('### DRIFT [% of initial] = '//real2str(100._qp*rel_drift,fmt_pct),lvl)

END SUBROUTINE mass_budget_report
!=======================================================================

!=======================================================================
SUBROUTINE mass_budget_write_diag(this)
!-----------------------------------------------------------------------
! NAME
!     mass_budget_write_diag
!
! DESCRIPTION
!     Write the terms of a budget into the file "diagevo.nc".
!
! AUTHORS & DATE
!     JB Clement, 08/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use output,  only: write_diagevo
use utility, only: slugify

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

! ARGUMENTS
! ---------
class(mass_budget), intent(in) :: this

! LOCAL VARIABLES
! ---------------
integer(di)               :: k
real(qp)                  :: m_ini, m, net, drift, rel_drift
character(:), allocatable :: prefix

! CODE
! ----
call this%totals(m_ini,m,net,drift,rel_drift)

prefix = 'totmass_'//trim(slugify(this%species))//'_'
do k = 1,this%nresvr
    call write_diagevo(prefix//trim(slugify(this%resvr(k)%name)),trim(this%species)//' mass in '//trim(this%resvr(k)%name),'kg',real(this%resvr(k)%m,dp))
end do
do k = 1,this%nexch
    call write_diagevo(prefix//'exch_'//trim(slugify(this%exch(k)%name)),'Cumulated '//trim(this%species)//' exchange from '//trim(this%exch(k)%name),'kg',real(this%exch(k)%m,dp))
end do
call write_diagevo(prefix//'total','Total '//trim(this%species)//' mass','kg',real(m,dp))
call write_diagevo(prefix//'drift',trim(this%species)//' mass balance drift','kg',real(drift,dp))

END SUBROUTINE mass_budget_write_diag
!=======================================================================

END MODULE mass_balance