MODULE surf_ice
!-----------------------------------------------------------------------
! NAME
!     surf_ice
!
! DESCRIPTION
!     Surface ice management.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

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

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

! PARAMETERS
! ----------
real(dp),                                 parameter :: rho_co2ice = 1650._dp ! Density of CO2 ice [kg.m-3]
real(dp),                                 parameter :: rho_h2oice = 920._dp  ! Density of H2O ice [kg.m-3]
real(dp),                                 protected :: threshold_h2oice_cap  ! Threshold to consider H2O ice as infinite reservoir [kg.m-2]
real(dp),                                 protected :: h2oice_huge_ini       ! Initial value for huge reservoirs of H2O ice [kg/m^2]
logical(k4), dimension(:),   allocatable, protected :: is_h2o_perice_PCM     ! Flag for the presence of perennial H2O ice in the PCM at the beginning [kg.m-2]
real(dp),    dimension(:,:), allocatable, protected :: co2_perice_PCM        ! Perennial CO2 ice in the PCM at the beginning [kg.m-2]

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

!=======================================================================
SUBROUTINE ini_surf_ice()
!-----------------------------------------------------------------------
! NAME
!     ini_surf_ice
!
! DESCRIPTION
!     Initialize surface ice arrays.
!
! AUTHORS & DATE
!     JB Clement 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

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

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

! CODE
! ----
if (.not. allocated(is_h2o_perice_PCM)) allocate(is_h2o_perice_PCM(ngrid))
if (.not. allocated(co2_perice_PCM)) allocate(co2_perice_PCM(ngrid,nslope))

END SUBROUTINE ini_surf_ice
!=======================================================================

!=======================================================================
SUBROUTINE end_surf_ice()
!-----------------------------------------------------------------------
! NAME
!     end_surf_ice
!
! DESCRIPTION
!     Deallocate surface ice arrays.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

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

! CODE
! ----
if (allocated(is_h2o_perice_PCM)) deallocate(is_h2o_perice_PCM)
if (allocated(co2_perice_PCM)) deallocate(co2_perice_PCM)

END SUBROUTINE end_surf_ice
!=======================================================================

!=======================================================================
SUBROUTINE set_surf_ice_config(threshold_h2oice_cap_in,h2oice_huge_ini_in)
!-----------------------------------------------------------------------
! NAME
!     set_surf_ice_config
!
! DESCRIPTION
!     Setter for 'surf_ice' configuration parameters.
!
! AUTHORS & DATE
!     JB Clement, 02/2026
!
! NOTES
!
!-----------------------------------------------------------------------

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

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

! ARGUMENTS
! ---------
real(dp), intent(in) :: threshold_h2oice_cap_in, h2oice_huge_ini_in

! CODE
! ----
threshold_h2oice_cap = threshold_h2oice_cap_in
h2oice_huge_ini = h2oice_huge_ini_in
call print_msg('threshold_h2oice_cap = '//real2str(threshold_h2oice_cap))
call print_msg('h2oice_huge_ini      = '//real2str(h2oice_huge_ini))

END SUBROUTINE set_surf_ice_config
!=======================================================================

!=======================================================================
SUBROUTINE set_co2_perice_PCM(co2_perice_PCM_in)
!-----------------------------------------------------------------------
! NAME
!     set_co2_perice_PCM
!
! DESCRIPTION
!     Setter for 'co2_perice_PCM'.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

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

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

! CODE
! ----
co2_perice_PCM(:,:) = co2_perice_PCM_in(:,:)

END SUBROUTINE set_co2_perice_PCM
!=======================================================================

!=======================================================================
SUBROUTINE set_is_h2o_perice_PCM(is_h2o_perice_PCM_in)
!-----------------------------------------------------------------------
! NAME
!     set_is_h2o_perice_PCM
!
! DESCRIPTION
!     Setter for 'is_h2o_perice_PCM'.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

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

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

! CODE
! ----
where(is_h2o_perice_PCM_in > 0.5_dp)
    is_h2o_perice_PCM = .true.
else where
    is_h2o_perice_PCM = .false.
end where

END SUBROUTINE set_is_h2o_perice_PCM
!=======================================================================

!=======================================================================
SUBROUTINE build4PCM_perice(h2o_ice,co2_ice,is_h2o_perice,h2o_ice4PCM,co2_ice4PCM)
!-----------------------------------------------------------------------
! NAME
!     build4PCM_perice
!
! DESCRIPTION
!     Reconstructs perennial ice and frost for the PCM.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry, only: ngrid
use frost,    only: h2o_frost4PCM
use slopes,   only: subslope_dist, def_slope_mean
use maths,    only: pi
use display,  only: print_msg

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

! ARGUMENTS
! ---------
real(dp),    dimension(:,:), intent(inout) :: h2o_ice, co2_ice
logical(k4), dimension(:),   intent(out)   :: is_h2o_perice            ! H2O perennial ice flag
real(dp),    dimension(:,:), intent(out)   :: h2o_ice4PCM, co2_ice4PCM ! Ice for PCM

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

! CODE
! ----
call print_msg('> Building surface ice/frost for the PCM')
co2_ice4PCM(:,:) = co2_ice(:,:)
h2o_ice4PCM(:,:) = 0._dp ! Because in the Mars PCM, only the variation of perennial H2O ice is monitored, not the absolute quantity
do i = 1,ngrid
    ! Is H2O ice still considered as an infinite reservoir for the PCM?
    if (sum(h2o_ice(i,:)*subslope_dist(i,:)/cos(pi*def_slope_mean(:)/180._dp)) > threshold_h2oice_cap) then
        ! There is enough ice to be considered as an infinite reservoir
        is_h2o_perice(i) = .true.
    else
        ! Too little ice to be considered as an infinite reservoir so ice is transferred to the frost
        is_h2o_perice(i) = .false.
        h2o_frost4PCM(i,:) = h2o_frost4PCM(i,:) + h2o_ice(i,:)
        h2o_ice(i,:) = 0._dp
    end if
end do

END SUBROUTINE build4PCM_perice
!=======================================================================

!=======================================================================
SUBROUTINE evolve_co2ice(co2_ice,d_co2ice,zshift_surf)
!-----------------------------------------------------------------------
! NAME
!     evolve_co2ice
!
! DESCRIPTION
!     Compute the evolution of CO2 ice over one year.
!
! AUTHORS & DATE
!     R. Vandemeulebrouck
!     L. Lange
!     JB Clement, 2023-2025
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry,  only: ngrid, nslope
use evolution, only: dt
use display,   only: print_msg

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

! ARGUMENTS
! ---------
real(dp), dimension(:,:), intent(inout) :: co2_ice
real(dp), dimension(:,:), intent(inout) :: d_co2ice
real(dp), dimension(:,:), intent(out)   :: zshift_surf

! LOCAL VARIABLES
! ---------------
real(dp), dimension(ngrid,nslope) :: co2_ice_old ! Old density of CO2 ice

! CODE
! ----
! Evolution of CO2 ice for each physical point
call print_msg('> Evolution of CO2 ice')

co2_ice_old = co2_ice
co2_ice = co2_ice + d_co2ice*dt
where (co2_ice < 0._dp)
    co2_ice = 0._dp
    d_co2ice = -co2_ice_old/dt
end where

zshift_surf = d_co2ice*dt/rho_co2ice

END SUBROUTINE evolve_co2ice
!=======================================================================

!=======================================================================
SUBROUTINE evolve_h2oice(delta_h2o_ads,delta_icetable,h2o_ice,d_h2oice,zshift_surf,stopcrit)
!-----------------------------------------------------------------------
! NAME
!     evolve_h2oice
!
! DESCRIPTION
!     Compute the evolution of H2O ice over one year.
!
! AUTHORS & DATE
!     R. Vandemeulebrouck
!     L. Lange
!     JB Clement, 2023-2025
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry,      only: ngrid, nslope
use evolution,     only: dt
use stopping_crit, only: stopping_crit_h2o, stopFlags
use display,       only: print_msg

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

! ARGUMENTS
! ---------
real(dp), dimension(:),   intent(in)    :: delta_h2o_ads, delta_icetable
real(dp), dimension(:,:), intent(inout) :: h2o_ice, d_h2oice
type(stopFlags),          intent(inout) :: stopcrit
real(dp), dimension(:,:), intent(out)   :: zshift_surf

! LOCAL VARIABLES
! ---------------
real(qp)                          :: S_atm_2_h2o, S_h2o_2_atm, S_atm_2_h2oice, S_h2oice_2_atm ! H2O balance variables
real(dp), dimension(ngrid,nslope) :: d_h2oice_new                                             ! Tendencies computed to keep balance

! CODE
! ----
call print_msg('> Evolution of H2O ice')

if (ngrid == 1) then ! In 1D
    h2o_ice = h2o_ice + d_h2oice*dt
    zshift_surf = d_h2oice*dt/rho_h2oice
else ! In 3D
    call stopping_crit_h2o(delta_h2o_ads,delta_icetable,h2o_ice,d_h2oice,S_atm_2_h2o,S_h2o_2_atm,S_atm_2_h2oice,S_h2oice_2_atm,stopcrit)
    if (stopcrit%stop_code() > 0) return

    call balance_h2oice_reservoirs(S_atm_2_h2o,S_h2o_2_atm,S_atm_2_h2oice,S_h2oice_2_atm,h2o_ice,d_h2oice,d_h2oice_new)
    h2o_ice = h2o_ice + d_h2oice_new*dt
    zshift_surf = d_h2oice_new*dt/rho_h2oice
end if

END SUBROUTINE evolve_h2oice
!=======================================================================

!=======================================================================
SUBROUTINE balance_h2oice_reservoirs(S_atm_2_h2o,S_h2o_2_atm,S_atm_2_h2oice,S_h2oice_2_atm,h2o_ice,d_h2oice,d_h2oice_new)
!-----------------------------------------------------------------------
! NAME
!     balance_h2oice_reservoirs
!
! DESCRIPTION
!     Balance H2O flux from and into atmosphere across reservoirs.
!
! AUTHORS & DATE
!     JB Clement, 2025
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use evolution, only: dt
use geometry,  only: ngrid, nslope

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

! ARGUMENTS
! ---------
real(qp),                 intent(in)    :: S_atm_2_h2o, S_h2o_2_atm, S_atm_2_h2oice, S_h2oice_2_atm
real(dp), dimension(:,:), intent(in)    :: h2o_ice
real(dp), dimension(:,:), intent(inout) :: d_h2oice
real(dp), dimension(:,:), intent(out)   :: d_h2oice_new

! LOCAL VARIABLES
! ---------------
integer(di) :: i, islope
real(qp)    :: S_target, S_target_subl_h2oice, S_target_cond_h2oice, S_ghostice, d_target ! Balance variables

! CODE
! ----
S_target = (S_atm_2_h2o + S_h2o_2_atm)/2._dp
S_target_cond_h2oice = S_atm_2_h2oice + S_target - S_atm_2_h2o
S_target_subl_h2oice = S_h2oice_2_atm + S_target - S_h2o_2_atm

d_h2oice_new = 0._dp
S_ghostice = 0._qp
do i = 1,ngrid
    do islope = 1,nslope
        if (d_h2oice(i,islope) > 0._dp) then ! Condensing
            d_h2oice_new(i,islope) = d_h2oice_new(i,islope)*real(S_target_cond_h2oice/S_atm_2_h2oice,dp)
        else if (d_h2oice(i,islope) < 0._dp) then ! Sublimating
            d_target = d_h2oice(i,islope)*real(S_target_subl_h2oice/S_h2oice_2_atm,dp)
            if (abs(d_target*dt) < h2o_ice(i,islope)) then ! Enough ice to sublimate everything
                d_h2oice_new(i,islope) = real(d_target,dp)
            else ! Not enough ice to sublimate everything
                ! We sublimate what we can
                d_h2oice_new(i,islope) = h2o_ice(i,islope)/dt
                ! It means the tendency is zero next time
                d_h2oice(i,islope) = 0._dp
                ! We compute the amount of H2O ice that we could not make sublimate
                S_ghostice = S_ghostice + abs(d_target*dt) - h2o_ice(i,islope)
            end if
        end if
    end do
end do

! We need to remove this ice unable to sublimate from places where ice condenseds in order to keep balance
where (d_h2oice_new > 0._dp) d_h2oice_new = d_h2oice_new*real((S_target_cond_h2oice - S_ghostice)/S_target_cond_h2oice,dp)

END SUBROUTINE balance_h2oice_reservoirs
!=======================================================================

END MODULE surf_ice