MODULE tracers
!-----------------------------------------------------------------------
! NAME
!     tracers
!
! DESCRIPTION
!     Tracer species management for tracking.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

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

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

! PARAMETERS
! ----------
character(11),                                parameter :: traceurdef_name = 'traceur.def'
integer(di),                                  protected :: nq        ! Number of tracers
character(30), dimension(:),     allocatable, protected :: qnames    ! Names of tracers
real(dp),      dimension(:,:,:), allocatable, protected :: q_PCM     ! Tracers in the "start.nc" at the beginning
integer(di),                                  protected :: iPCM_qh2o ! Index for H2O vapor tracer from PCM
real(dp),      dimension(:),     allocatable, protected :: mmol      ! Molar masses of tracers [g/mol]

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

!=======================================================================
SUBROUTINE ini_tracers()
!-----------------------------------------------------------------------
! NAME
!     ini_tracers
!
! DESCRIPTION
!     Allocate tracer molar mass array.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

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

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

! LOCAL VARIABLES
! ---------------
logical(k4) :: here
integer(di) :: i, ierr, funit

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

! Allocation
if (.not. allocated(mmol)) allocate(mmol(nq))
if (.not. allocated(qnames)) allocate(qnames(nq))
if (.not. allocated(q_PCM)) allocate(q_PCM(ngrid,nlayer,nq))

! Initialize the names of tracers
do i = 1,nq
    read(funit,*,iostat = ierr) qnames(i)
    if (ierr /= 0) call stop_clean(__FILE__,__LINE__,'error reading the names of tracers in "'//traceurdef_name//'"!',1)
end do
close(funit)

! Initialize the indices of tracers
iPCM_qh2o = -1
do i = 1,nq
    if (qnames(i) == "h2o_vap") then
        iPCM_qh2o = i
        mmol(i) = 18._dp
    end if
end do

! Checking if everything has been found
if (iPCM_qh2o < 0) call stop_clean(__FILE__,__LINE__,'H2O vapour index not found!',1)

END SUBROUTINE ini_tracers
!=======================================================================

!=======================================================================
SUBROUTINE end_tracers()
!-----------------------------------------------------------------------
! NAME
!     end_tracers
!
! DESCRIPTION
!     Deallocate tracer molar mass array.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

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

! CODE
! ----
if (allocated(mmol)) deallocate(mmol)
if (allocated(qnames)) deallocate(qnames)
if (allocated(q_PCM)) deallocate(q_PCM)

END SUBROUTINE end_tracers
!=======================================================================

!=======================================================================
SUBROUTINE set_q_PCM(q_PCM_in,iq)
!-----------------------------------------------------------------------
! NAME
!     set_q_PCM
!
! DESCRIPTION
!     Setter for 'q_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) :: q_PCM_in
integer(di),                intent(in) :: iq

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

! CODE
! ----
do l = 1,size(q_PCM_in,3)
    call dyngrd2vect(nlon + 1,nlat,ngrid,q_PCM_in(:,:,l),q_PCM(:,l,iq))
end do

END SUBROUTINE set_q_PCM
!=======================================================================

!=======================================================================
SUBROUTINE adapt_tracers2pressure(ps_avg_glob_old,ps_avg_glob,ps_ts,q_h2o_ts,q_co2_ts)
!-----------------------------------------------------------------------
! NAME
!     adapt_tracers2pressure
!
! DESCRIPTION
!     Adapt the timeseries of tracers VMR to the new pressure.
!
! AUTHORS & DATE
!     JB Clement, 01/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry,   only: ngrid, nday
use atmosphere, only: ap, bp
use display,    only: print_msg

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

! ARGUMENTS
! ---------
real(dp),                        intent(in)    :: ps_avg_glob_old, ps_avg_glob
real(dp), dimension(ngrid,nday), intent(inout) :: ps_ts, q_h2o_ts, q_co2_ts

! LOCAL VARIABLES
! ---------------
real(dp), dimension(ngrid,nday) :: plev1, plev2, plev1_old, plev2_old

! CODE
! ----
call print_msg("> Adapting the timeseries of tracers VMR to the new pressure")

! Build the first levels of pressure
plev1_old(:,:) = ap(1) + bp(1)*ps_ts(:,:)
plev2_old(:,:) = ap(2) + bp(2)*ps_ts(:,:)

! Evolve timeseries of surface pressure
ps_ts(:,:) = ps_ts(:,:)*ps_avg_glob/ps_avg_glob_old

! Build the first levels of pressure
plev1(:,:) = ap(1) + bp(1)*ps_ts(:,:)
plev2(:,:) = ap(2) + bp(2)*ps_ts(:,:)

! Adapt the timeseries of VMR tracers
! H2O
q_h2o_ts(:,:) = q_h2o_ts(:,:)*(plev1_old(:,:) - plev2_old(:,:))/(plev1(:,:) - plev2(:,:))
where (q_h2o_ts < 0._dp) q_h2o_ts = 0._dp
where (q_h2o_ts > 1._dp) q_h2o_ts = 1._dp
! CO2
q_co2_ts(:,:) = q_co2_ts(:,:)*(plev1_old(:,:) - plev2_old(:,:))/(plev1(:,:) - plev2(:,:))
where (q_co2_ts < 0._dp) q_co2_ts = 0._dp
where (q_co2_ts > 1._dp) q_co2_ts = 1._dp

END SUBROUTINE adapt_tracers2pressure
!=======================================================================

!=======================================================================
SUBROUTINE build4PCM_tracers(ps4PCM,q4PCM)
!-----------------------------------------------------------------------
! NAME
!     build4PCM_tracers
!
! DESCRIPTION
!     Reconstructs the tracer VMRs for the PCM.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!     C. Metz, 01/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use atmosphere, only: ap, bp, ps_PCM
use geometry,   only: ngrid, nlayer
use display,    only: print_msg

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

! ARGUMENTS
! ---------
real(dp), dimension(:),     intent(in)  :: ps4PCM
real(dp), dimension(:,:,:), intent(out) :: q4PCM

! LOCAL VARIABLES
! ---------------
integer(di)                           :: i, l, iq
real(dp)                              :: extra_mass
real(dp), dimension(:,:), allocatable :: dz_plev_PCM, dz_plev_new

! CODE
! ----
call print_msg('> Building tracers for the PCM')
allocate(dz_plev_PCM(ngrid,nlayer),dz_plev_new(ngrid,nlayer))

! Compute thickness of pressure levels (= atmospheric layers)
do l = 1,nlayer
    dz_plev_PCM(:,l) = ap(l) + bp(l)*ps_PCM(:) - ap(l + 1) - bp(l + 1)*ps_PCM(:)
    dz_plev_new(:,l) = ap(l) + bp(l)*ps_PCM(:) - ap(l + 1) - bp(l + 1)*ps4PCM(:)
end do

! Reconstruct tracers
do iq = 1,nq
    do l = 1,nlayer
        do i = 1,ngrid
            ! Standard scaling
            q4PCM(i,l,iq) = q_PCM(i,l,iq)*dz_plev_PCM(i,l)/dz_plev_new(i,l)

            ! CO2 logic: account for the dry air mass change
            if (trim(qnames(iq)) == "co2") q4PCM(i,l,iq) = q4PCM(i,l,iq) + (dz_plev_new(i,l) - dz_plev_PCM(i,l)) / dz_plev_new(i,l)
        end do
    end do
end do

! Mass conservation and clipping
do iq = 1,nq
    do i = 1,ngrid
        do l = 1,nlayer
            ! Clipping negative values
            if (q4PCM(i,l,iq) < 0._dp) q4PCM(i,l,iq) = 0._dp

            ! VMR limit check (ignore number density tracers)
            if (q4PCM(i,l,iq) > 1._dp .and.              &
                (qnames(iq) /= "dust_number") .and.      &
                (qnames(iq) /= "ccn_number") .and.       &
                (qnames(iq) /= "stormdust_number") .and. &
                (qnames(iq) /= "topdust_number")) then

                if (l < nlayer) then
                    extra_mass = (q4PCM(i,l,iq) - 1.)*dz_plev_new(i,l) ! extra_mass units: [VMR*pressure]
                    q4PCM(i,l,iq) = 1._dp

                    ! Transfer extra mass to the layer above
                    ! Note: We divide by dz_plev_new(i,l + 1) to convert back to VMR
                    q4PCM(i,l + 1,iq) = q4PCM(i,l + 1,iq) + extra_mass/dz_plev_new(i,l + 1)
                else
                    q4PCM(i,l,iq) = 1._dp
                end if
            end if
        end do
    end do
end do

deallocate(dz_plev_PCM,dz_plev_new)

END SUBROUTINE build4PCM_tracers
!=======================================================================

END MODULE tracers