MODULE atmosphere
!-----------------------------------------------------------------------
! NAME
!     atmosphere
!
! DESCRIPTION
!     Contains global parameters used for the atmosphere.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

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

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

! PARAMETERS
! ----------
character(9),                          parameter :: z2sigdef_name = 'z2sig.def'
logical(k4),                           protected :: hybrid_alt_coord   ! Flag fo hybrid coordinates
real(dp), dimension(:),   allocatable, protected :: ap             ! Hybrid (pressure contribution) coordinate at layer interfaces [Pa]
real(dp), dimension(:),   allocatable, protected :: bp             ! Hybrid (sigma contribution) coordinate at layer interfaces
real(dp), dimension(:),   allocatable, protected :: ps_PCM         ! Surface pressure in the "start.nc" at the beginning [Pa]
real(dp), dimension(:,:), allocatable, protected :: teta_PCM       ! Potential temperature in the "start.nc" at the beginning [K]
real(dp), dimension(:,:), allocatable, protected :: u_PCM          ! Zonal wind [m/s]
real(dp), dimension(:,:), allocatable, protected :: v_PCM          ! Meridional wind [m/s]
real(dp)                                         :: CO2cond_ps_PCM ! Coefficient to control the surface pressure change

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

!=======================================================================
SUBROUTINE ini_atmosphere()
!-----------------------------------------------------------------------
! NAME
!     ini_atmosphere
!
! DESCRIPTION
!     Initialize the parameters of module 'atmosphere'.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

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

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

! CODE
! ----
if (.not. allocated(ap)) allocate(ap(nlayer + 1))
if (.not. allocated(bp)) allocate(bp(nlayer + 1))
if (.not. allocated(ps_PCM)) allocate(ps_PCM(ngrid))
if (.not. allocated(teta_PCM)) allocate(teta_PCM(ngrid,nlayer))
if (.not. allocated(u_PCM)) allocate(u_PCM(ngrid,nlayer))
if (.not. allocated(v_PCM)) allocate(v_PCM(ngrid,nlayer))

END SUBROUTINE ini_atmosphere
!=======================================================================

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

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

! CODE
! ----
if (allocated(ap)) deallocate(ap)
if (allocated(bp)) deallocate(bp)
if (allocated(ps_PCM)) deallocate(ps_PCM)
if (allocated(teta_PCM)) deallocate(teta_PCM)
if (allocated(u_PCM)) deallocate(u_PCM)
if (allocated(v_PCM)) deallocate(v_PCM)

END SUBROUTINE end_atmosphere
!=======================================================================

!=======================================================================
SUBROUTINE set_atmosphere_config(hybrid_alt_coord_in)
!-----------------------------------------------------------------------
! NAME
!     set_atmosphere_config
!
! DESCRIPTION
!     Setter for 'atmosphere' configuration parameters.
!
! AUTHORS & DATE
!     JB Clement, 02/2026
!
! NOTES
!
!-----------------------------------------------------------------------

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

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

! ARGUMENTS
! ---------
logical(k4), intent(in) :: hybrid_alt_coord_in

! CODE
! ----
hybrid_alt_coord = hybrid_alt_coord_in
call print_msg('hybrid_alt_coord = '//bool2str(hybrid_alt_coord))

END SUBROUTINE set_atmosphere_config
!=======================================================================

!=======================================================================
SUBROUTINE set_ap(ap_in)
!-----------------------------------------------------------------------
! NAME
!     set_ap
!
! DESCRIPTION
!     Setter for 'ap'.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

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

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

! CODE
! ----
ap(:) = ap_in(:)

END SUBROUTINE set_ap
!=======================================================================

!=======================================================================
SUBROUTINE set_bp(bp_in)
!-----------------------------------------------------------------------
! NAME
!     set_bp
!
! DESCRIPTION
!     Setter for 'bp'.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

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

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

! CODE
! ----
bp(:) = bp_in(:)

END SUBROUTINE set_bp
!=======================================================================

!=======================================================================
SUBROUTINE set_ps_PCM(ps_PCM_in)
!-----------------------------------------------------------------------
! NAME
!     set_ps_PCM
!
! DESCRIPTION
!     Setter for 'ps_PCM'.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry, only: dyngrd2vect, nlon, nlat, ngrid

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

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

! CODE
! ----
call dyngrd2vect(nlon + 1,nlat,ngrid,ps_PCM_in,ps_PCM)

END SUBROUTINE set_ps_PCM
!=======================================================================

!=======================================================================
SUBROUTINE set_teta_PCM(teta_PCM_in)
!-----------------------------------------------------------------------
! NAME
!     set_teta_PCM
!
! DESCRIPTION
!     Setter for 'teta_PCM'.
!
! AUTHORS & DATE
!     JB Clement, 01/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry, only: dyngrd2vect, nlon, nlat, ngrid, nlayer

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

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

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

! CODE
! ----
do l = 1,nlayer
    call dyngrd2vect(nlon + 1,nlat,ngrid,teta_PCM_in(:,:,l),teta_PCM(:,l))
end do

END SUBROUTINE set_teta_PCM
!=======================================================================

!=======================================================================
SUBROUTINE set_u_PCM(u_PCM_in)
!-----------------------------------------------------------------------
! NAME
!     set_u_PCM
!
! DESCRIPTION
!     Setter for 'u_PCM'.
!
! AUTHORS & DATE
!     JB Clement, 01/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry, only: dyngrd2vect, nlon, nlat, ngrid, nlayer

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

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

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

! CODE
! ----
do l = 1,nlayer
    call dyngrd2vect(nlon + 1,nlat,ngrid,u_PCM_in(:,:,l),u_PCM(:,l))
end do

END SUBROUTINE set_u_PCM
!=======================================================================

!=======================================================================
SUBROUTINE set_v_PCM(v_PCM_in)
!-----------------------------------------------------------------------
! NAME
!     set_v_PCM
!
! DESCRIPTION
!     Setter for 'v_PCM'.
!
! AUTHORS & DATE
!     JB Clement, 01/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry, only: dyngrd2vect, nlon, nlat, ngrid, nlayer

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

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

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

! CODE
! ----
do l = 1,nlayer
    call dyngrd2vect(nlon + 1,nlat,ngrid,v_PCM_in(:,:,l),v_PCM(:,l))
end do

END SUBROUTINE set_v_PCM
!=======================================================================

!=======================================================================
SUBROUTINE compute_alt_coord(pa,preff,ap,bp,aps,bps)
!-----------------------------------------------------------------------
! NAME
!     compute_alt_coord
!
! DESCRIPTION
!     Compute values of the coefficients to define the altitude
!     coordinates.
!
! AUTHORS & DATE
!     JB Clement, 02/2026
!
! NOTES
!     Taken from 'disvert_noterre' written by F. Forget, Y. Wanherdrick
!     and P. Levan in the dynamics.
!
!     'pa' and 'preff' are read from "start.nc"/"start1D.txt".
!
!     WARNING: to compute the values at mid_layer, there is an arbitrary
!     decision here to put the middle of the layer half-way (pressure-wise)
!     between boundaries.
!     In addition, for the last layer (the top of which is at zero pressure)
!     P(llm), we choose to put it with the same interval (in log(pressure))
!     from P(llm-1), than there is between between P(llm-1) and P(llm-2)
!     This choice is quite specific and must be compliant with the PCM!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry, only: nlayer
use display,  only: print_msg
use stoppage, only: stop_clean
use utility,  only: real2str

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

! ARGUMENTS
! ---------
real(dp)                                               :: pa, preff
real(dp), dimension(nlayer + 1), intent(out)           :: ap, bp
real(dp), dimension(nlayer),     intent(out), optional :: aps, bps

! LOCAL VAIRABLES
! ---------------
integer(di)                     :: i
real(dp)                        :: scaleheight, newsig
real(dp), dimension(nlayer + 1) :: sig

! CODE
! ----
! Read vertical grid definition
call read_z2sig(scaleheight,sig)
call print_msg('Pa    [Pa] = '//real2str(pa))
call print_msg('Preff [Pa] = '//real2str(preff))

! Compute ap and bp coefficients
if (hybrid_alt_coord) then
    call print_msg('> Defining hybrid altitude coordinates')
    do i = 1,nlayer
        call compute_hybrid_sig(sig(i),pa,preff,newsig)
        bp(i) = exp(1._dp - 1._dp/(newsig**2))
        ap(i) = pa*(newsig - bp(i))
    end do
    ap(nlayer + 1) = 0._dp
    bp(nlayer + 1) = 0._dp
else
    call print_msg('> Defining sigma altitude coordinates')
    ap(:) = 0._dp
    bp(1:nlayer) = sig(1:nlayer)
    bp(nlayer + 1) = 0._dp
end if

! Compute aps and bps coefficients (mid-layer) if asked
if (present(aps) .neqv. present(bps)) call stop_clean(__FILE__,__LINE__,"'aps' and 'bps' must be computed together!",1)
if (.not. present(aps)) return

do i = 1,nlayer - 1
    aps(i) =  0.5_dp*(ap(i) + ap(i + 1))
    bps(i) =  0.5_dp*(bp(i) + bp(i + 1))
end do
if (hybrid_alt_coord) then
    aps(nlayer) = aps(nlayer - 1)**2/aps(nlayer - 2)
    bps(nlayer) = 0.5_dp*(bp(nlayer) + bp(nlayer + 1))
else
    aps(nlayer) = 0._dp
    bps(nlayer) = bps(nlayer - 1)**2/bps(nlayer - 2)
end if

END SUBROUTINE compute_alt_coord
!=======================================================================

!=======================================================================
SUBROUTINE compute_hybrid_sig(sig,pa,preff,newsig)
!-----------------------------------------------------------------------
! NAME
!     compute_hybrid_sig
!
! DESCRIPTION
!     Compute modified sigma value to keep vertical coordinates described
!     in "z2sig.def" when defining hybrid coordinates.
!
! AUTHORS & DATE
!     JB Clement, 01/2026
!
! NOTES
!     Taken from 'sig_hybrid' written by F. Forget (2002) in the
!     dynamics.
!
!     Knowing sig (the "sigma" levels where we want to put the layers),
!     it solves 'newsig' such that:
!       (1 - pa/preff)*exp(1 - 1/newsig^2) + (pa/preff)*newsig = sig
!     which cannot be done analyticaly.
!     => needs for an iterative procedure
!
!     Information: when exp(1 - 1/x**2) becomes << x
!           x      exp(1 - 1/x**2)/x
!           1           1
!           0.68       0.5
!           0.5        1.E-1
!           0.391      1.E-2
!           0.333      1.E-3
!           0.295      1.E-4
!           0.269      1.E-5
!           0.248      1.E-6
!     => one can thus use newsig = sig*preff/pa if sig*preff/pa < 0.25
!-----------------------------------------------------------------------

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

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

! ARGUMENTS
! ---------
real(dp), intent(in)  :: sig, pa, preff
real(dp), intent(out) :: newsig

! LOCAL VARIABLES
! ---------------
real(dp)    :: x1, x2, F
integer(di) :: j

! CODE
! ----
newsig = sig
x1 = 0._dp
x2 = 1._dp

if (sig >= 1.) then
    newsig = sig
else if (sig*preff/pa >= 0.25_dp) then
    do j = 1,9999 ! Overkill maximum number of iterations
        F = ((1._dp - pa/preff)*exp(1._dp - 1._dp/newsig**2) + (pa/preff)*newsig)/sig

        if (F > 1.) then
            x2 = newsig
            newsig = 0.5_dp*(x1 + newsig)
        else
        x1 = newsig
        newsig = 0.5_dp*(x2 + newsig)
        end if
        ! Convergence is assumed if sig is within 0.01m (in pseudo-altitude) of the target value
        if (abs(10._dp * log(F)) < 1.e-5_dp) exit
    end do
else !if (sig*preff/pa <= 0.25_dp) then
    newsig = sig*preff/pa
end if

END SUBROUTINE compute_hybrid_sig
!=======================================================================

!=======================================================================
SUBROUTINE read_z2sig(scaleheight,sig)
!-----------------------------------------------------------------------
! NAME
!     read_z2sig
!
! DESCRIPTION
!     Read vertical grid definition from z2sig.def and compute sigma
!     levels.
!
! AUTHORS & DATE
!     JB Clement, 01/2026
!
! NOTES
!     Adapted from 'disvert_noterre' written by F. Forget, Y.
!     Wanherdrick and P. Levan in the dynamics.
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry, only: nlayer
use stoppage, only: stop_clean
use display,  only: print_msg
use utility,  only: real2str

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

! ARGUMENTS
! ---------
real(dp),                        intent(out) :: scaleheight
real(dp), dimension(nlayer + 1), intent(out) :: sig

! LOCAL VARIABLES
! ---------------
real(dp), dimension(nlayer) :: zsig
integer(di)                 :: l, ierr, funit
logical(k4)                 :: here

! CODE
! ----
! Open "z2sig.def"
inquire(file = z2sigdef_name,exist = here)
if (.not. here) call stop_clean(__FILE__,__LINE__,'cannot find required file "'//z2sigdef_name//'"!',1)
call print_msg('> Reading "'//z2sigdef_name//'"')
open(newunit = funit,file = z2sigdef_name,status = "old",action = "read",iostat = ierr)
if (ierr /= 0) call stop_clean(__FILE__,__LINE__,'error opening "'//z2sigdef_name//'"',ierr)

! Read scale height
read(funit,*,iostat = ierr) scaleheight
if (ierr /= 0) call stop_clean(__FILE__,__LINE__,"error reading 'scaleheight'!",1)

! Read pseudo-altitudes
do l = 1,nlayer
    read(funit,*,iostat = ierr) zsig(l)
    if (ierr /= 0) call stop_clean(__FILE__,__LINE__,'"'//z2sigdef_name//'" too short for the number of altitude levels!',1)
end do

! Check for extra lines in the file (warning only)
read(funit,*,iostat = ierr)
if (ierr == 0) call print_msg('Warning: "'//z2sigdef_name//'" has more lines than expected!')

! Compute sigma values
sig(1) = 1.
do l = 2,nlayer
    sig(l) = 0.5_dp*(exp(-zsig(l)/scaleheight) + exp(-zsig(l - 1)/scaleheight))
end do
sig(nlayer + 1) = 0._dp

! Close
close(funit)
call print_msg('Scale height [km] = '//real2str(scaleheight))

END SUBROUTINE read_z2sig
!=======================================================================

!=======================================================================
SUBROUTINE build4PCM_atmosphere(ps_avg,ps_dev,ps_avg_glob,ps_avg_glob_ini,ps4PCM,pa4PCM,preff4PCM,teta4PCM,air_mass4PCM)
!-----------------------------------------------------------------------
! NAME
!     build4PCM_atmosphere
!
! DESCRIPTION
!     Reconstructs the pressure, potential temperature and tha air mass
!     for the PCM.
!
! AUTHORS & DATE
!     JB Clement, 01/2026
!
! NOTES
!     Air mass computation taken from 'massdair' and 'pression' written
!     by P. Le Van and F. Hourdin in the dynamics.
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry, only: ngrid, nlayer, cell_area, vect2dyngrd
use physics,  only: g, rcp
use display,  only: print_msg

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

! ARGUMENTS
! ---------
real(dp), dimension(:),   intent(in)  :: ps_avg, ps_dev
real(dp),                 intent(in)  :: ps_avg_glob, ps_avg_glob_ini
real(dp), dimension(:,:), intent(out) :: air_mass4PCM, teta4PCM
real(dp), dimension(:),   intent(out) :: ps4PCM
real(dp),                 intent(out) :: pa4PCM, preff4PCM

! LOCAL VARIABLES
! ---------------
real(dp), dimension(ngrid,nlayer + 1) :: p ! 3D pressure field
integer(di)                           :: l

! CODE
! ----
call print_msg('> Building pressure for the PCM')
! The pressure deviation is rescaled to avoid disproportionate oscillations in case of huge change of average pressure during the PEM run
ps4PCM(:) = ps_avg(:) + ps_dev(:)*ps_avg_glob/ps_avg_glob_ini
pa4PCM = ps_avg_glob_ini/30.  ! For now the altitude grid is not changed
preff4PCM = ps_avg_glob_ini   ! For now the altitude grid is not changed

call print_msg('> Building potential temperature for the PCM')
do l = 1,nlayer
    ! Correction on teta due to surface pressure changes
    teta4PCM(:,l) = teta_PCM(:,l)*ps4PCM(:)**rcp
end do

call print_msg('> Building air mass for the PCM')
! Compute atmospheric pressure
do l = 1,nlayer + 1
    p(:,l) = ap(l) + bp(l)*ps4PCM(:)
end do

! Compute air mass
do l = 1,nlayer
    air_mass4PCM(:,l) = cell_area(:)*(p(:,l) - p(:,l + 1))/g
end do

END SUBROUTINE build4PCM_atmosphere
!=======================================================================

!=======================================================================
SUBROUTINE evolve_pressure(d_co2ice,delta_co2_ads,do_sorption,ps_avg_glob_old,ps_avg_glob,ps_avg)
!-----------------------------------------------------------------------
! NAME
!     evolve_pressure
!
! DESCRIPTION
!     Evolve pressure according to tendency.
!
! AUTHORS & DATE
!     JB Clement, 01/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry, only: ngrid, nslope, cell_area, total_surface
use slopes,   only: subslope_dist, def_slope_mean
use physics,  only: g
use maths,    only: pi
use display,  only: print_msg
use utility,  only: real2str

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

! ARGUMENTS
! ---------
real(dp), dimension(:,:), intent(in)    :: d_co2ice
real(dp), dimension(:),   intent(in)    :: delta_co2_ads
logical(k4),              intent(in)    :: do_sorption
real(dp),                 intent(inout) :: ps_avg_glob_old, ps_avg_glob
real(dp), dimension(:),   intent(inout) :: ps_avg

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

! CODE
! ----
call print_msg("> Evolving the surface pressure")
ps_avg_glob_old = ps_avg_glob
do i = 1,ngrid
    do islope = 1,nslope
        ps_avg_glob = ps_avg_glob - CO2cond_ps_PCM*g*cell_area(i)*d_co2ice(i,islope)*subslope_dist(i,islope)/cos(pi*def_slope_mean(islope)/180.)/total_surface
    end do
    if (do_sorption) ps_avg_glob = ps_avg_glob - g*cell_area(i)*delta_co2_ads(i)/total_surface
end do
call print_msg('Global average pressure (old|new) [Pa] = '//real2str(ps_avg_glob_old)//' | '//real2str(ps_avg_glob))
ps_avg(:) = ps_avg(:)*ps_avg_glob/ps_avg_glob_old

END SUBROUTINE evolve_pressure
!=======================================================================

END MODULE atmosphere