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

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

! MODULE VARIABLES
! ----------------
real, dimension(:,:), allocatable :: h2o_ice ! H2O ice surface [kg.m-2]
real, dimension(:,:), allocatable :: co2_ice ! CO2 ice surface [kg.m-2]
real :: h2oice_cap_threshold                 ! Threshold to consider H2O ice as infinite reservoir [kg.m-2]

! MODULE PARAMETERS
! -----------------
real, parameter :: rho_co2ice = 1650. ! Density of CO2 ice [kg.m-3]
real, parameter :: rho_h2oice = 920.  ! Density of H2O ice [kg.m-3]

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

!=======================================================================
SUBROUTINE set_perice4PCM(ngrid,nslope,PCMfrost,is_h2o_perice,h2oice_PCM,co2ice_PCM)
!-----------------------------------------------------------------------
! NAME
!     set_perice4PCM
!
! DESCRIPTION
!     Reconstruct perennial ice for PCM from PEM ice computations.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use metamorphism,   only: iPCM_h2ofrost
use comslope_mod,   only: subslope_dist, def_slope_mean
use phys_constants, only: pi

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

! ARGUMENTS
! ---------
integer,                       intent(in)    :: ngrid, nslope
real, dimension(:,:,:),        intent(inout) :: PCMfrost               ! PCM frost
logical, dimension(ngrid),     intent(out)   :: is_h2o_perice          ! H2O perennial ice flag
real, dimension(ngrid,nslope), intent(out)   :: h2oice_PCM, co2ice_PCM ! Ice for PCM

! LOCAL VARIABLES
! ---------------
integer :: i

! CODE
! ----
write(*,*) '> Reconstructing perennial ic for the PCM'
co2ice_PCM(:,:) = co2_ice(:,:)
h2oice_PCM(:,:) = 0. ! 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.)) > h2oice_cap_threshold) 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.
        PCMfrost(i,iPCM_h2ofrost,:) = PCMfrost(i,iPCM_h2ofrost,:) + h2o_ice(i,:)
        h2o_ice(i,:) = 0.
    endif
enddo

END SUBROUTINE set_perice4PCM
!=======================================================================

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

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

! ARGUMENTS
! ---------
integer, intent(in) :: ngrid, nslope

! CODE
! ----
if (.not. allocated(h2o_ice)) allocate(h2o_ice(ngrid,nslope))
if (.not. allocated(co2_ice)) allocate(co2_ice(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(h2o_ice)) deallocate(h2o_ice)
if (allocated(co2_ice)) deallocate(co2_ice)

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

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

! DEPENDENCIES
! ------------
use evolution, only: dt

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

! ARGUMENTS
! ---------
integer,                       intent(in)    :: ngrid, nslope
real, dimension(ngrid,nslope), intent(inout) :: co2_ice     ! CO2 ice
real, dimension(ngrid,nslope), intent(inout) :: d_co2ice    ! Evolution of CO2 ice over one year
real, dimension(ngrid,nslope), intent(out)   :: zshift_surf ! Elevation shift for the surface [m]

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

! CODE
! ----
! Evolution of CO2 ice for each physical point
write(*,*) '> Evolution of CO2 ice'

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

zshift_surf = d_co2ice*dt/rho_co2ice

END SUBROUTINE evol_co2_ice
!=======================================================================

!=======================================================================
SUBROUTINE evol_h2o_ice(ngrid,nslope,cell_area,delta_h2o_adsorbed,delta_h2o_icetablesublim,h2o_ice,d_h2oice,zshift_surf,stopCrit)
!-----------------------------------------------------------------------
! NAME
!     evol_h2o_ice
!
! DESCRIPTION
!     Compute the evolution of H2O ice over one year.
!
! AUTHORS & DATE
!     R. Vandemeulebrouck
!     L. Lange
!     JB Clement, 2023-2025
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use evolution,     only: dt
use stopping_crit, only: stopping_crit_h2o, stopFlags

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

! ARGUMENTS
! ---------
integer,                       intent(in)    :: ngrid, nslope
real, dimension(ngrid),        intent(in)    :: cell_area                ! Area of each mesh grid (m^2)
real, dimension(ngrid),        intent(in)    :: delta_h2o_adsorbed       ! Mass of H2O adsorbed/desorbded in soil (kg/m^2)
real, dimension(ngrid),        intent(in)    :: delta_h2o_icetablesublim ! Mass condensed/sublimated at ice table (kg/m^2)
real, dimension(ngrid,nslope), intent(inout) :: h2o_ice                  ! H2O ice (kg/m^2)
real, dimension(ngrid,nslope), intent(inout) :: d_h2oice                 ! Tendency of H2O ice (kg/m^2/year)
type(stopFlags),               intent(inout) :: stopCrit                 ! Stopping criterion
real, dimension(ngrid,nslope), intent(out)   :: zshift_surf              ! Elevation shift for the surface [m]

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

! CODE
! ----
write(*,*) '> 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(ngrid,nslope,cell_area,delta_h2o_adsorbed,delta_h2o_icetablesublim,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(ngrid,nslope,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
endif

END SUBROUTINE evol_h2o_ice
!=======================================================================

!=======================================================================
SUBROUTINE balance_h2oice_reservoirs(ngrid,nslope,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

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

! ARGUMENTS
! ---------
integer,                       intent(in)    :: ngrid, nslope ! # of grid points, # of subslopes
real,                          intent(in)    :: S_atm_2_h2o, S_h2o_2_atm, S_atm_2_h2oice, S_h2oice_2_atm ! Variables to conserve H2O
real, dimension(ngrid,nslope), intent(in)    :: h2o_ice       ! H2O ice (kg/m^2)
real, dimension(ngrid,nslope), intent(inout) :: d_h2oice      ! Tendency of H2O ice (kg/m^2/year)
real, dimension(ngrid,nslope), intent(out)   :: d_h2oice_new  ! Adjusted tendencies to keep balance between donor and recipient reservoirs

! LOCAL VARIABLES
! ---------------
integer :: i, islope
real    :: 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.
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.
S_ghostice = 0.
do i = 1,ngrid
    do islope = 1,nslope
        if (d_h2oice(i,islope) > 0.) then ! Condensing
            d_h2oice_new(i,islope) = d_h2oice_new(i,islope)*S_target_cond_h2oice/S_atm_2_h2oice
        else if (d_h2oice(i,islope) < 0.) then ! Sublimating
            d_target = d_h2oice(i,islope)*S_target_subl_h2oice/S_h2oice_2_atm
            if (abs(d_target*dt) < h2o_ice(i,islope)) then ! Enough ice to sublimate everything
                d_h2oice_new(i,islope) = d_target
            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.
                ! 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)
            endif
        endif
    enddo
enddo

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

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

END MODULE surf_ice