module age_of_air_mod

!****************************************************************
!
!     Age of air module
!
!     Author: Maureen Cohen
!     -------
!
!     Version: 01/12/2023
!     See "Age of air as a diagnostic for transport timescales in global models,"
!     Geosci. Model Dev., 11, 3109–3130, 2018, https://doi.org/10.5194/gmd-11-3109-2018
!
!*****************************************************************

implicit none

logical,save :: ok_aoa ! flag to trigger runing with Age of Air scheme
logical,save :: reinit_aoa ! flag to reinitialize Age of air scheme
integer, save :: i_aoa ! index of age of air tracer in the tracer array
!$OMP THREADPRIVATE(ok_aoa,reinit_aoa,i_aoa)
real (kind=8), parameter :: forc_con = 1e-15 ! forcing constant in kg/kg/s
integer,save :: lev_aoa ! atmospheric level at which  forcing is introduced
!$OMP THREADPRIVATE(lev_aoa)
real, allocatable, save:: init_age(:,:) ! saved initial value of age(:,:)
!$OMP THREADPRIVATE(init_age)

!---------------------------------------------------------------------------- 
  contains
!----------------------------------------------------------------------------

subroutine age_of_air(aoa_tracer,klon,klev,itap,pdtphys,age_offset,aoa)

    implicit none

! Inputs    
    integer,intent(in) :: klon, klev           ! number of gridpoints and levels
    integer,intent(in) :: itap                 ! physics timestep counter
    real,intent(in) :: pdtphys                 ! physics timestep length (s)
    real,intent(in) :: age_offset(klon,klev)   ! Init value of age, from restartphy or reinialitised to 0
                                               ! This value is the same every time step, but changes between restarts
    
! In/out
    real,intent(inout) :: aoa_tracer(klon,klev) ! age of air tracer mixing ratio (kg/kg)
    
! Outputs
    real,intent(out) :: aoa(klon,klev) ! age of air value in seconds
    
! Local
    real :: ts_forc ! forcing added in this timestep in kg/kg
    real :: run_time ! elapsed simulation time in seconds

    
!******************Add forcing constant to level lev_aoa**********

    run_time = (itap+1)*pdtphys ! (number of physics timesteps * seconds per physics timestep) = seconds elapsed since start

    ts_forc = run_time*forc_con ! elapsed time x forcing constant = tracer source this timestep
    
    aoa_tracer(:,lev_aoa) = ts_forc
    
    aoa(:,:) = run_time - aoa_tracer(:,:)/forc_con + age_offset(:,:) ! calculate age of air in seconds

end subroutine age_of_air

!*****************************************************************
!
!     Extracts index of age of air tracer in the tracer array
!
!*****************************************************************
  subroutine aoa_ini(ok_chem, tr_scheme)
  use infotrac_phy, only: nqtot, tname
  use chemparam_mod, only: M_tr, type_tr
  implicit none
  
  integer,intent(in) :: tr_scheme
  logical,intent(in) :: ok_chem
  integer :: i
  character(len=*),parameter :: rname="aoa_ini"
  
  if (.not. ok_chem .and. (tr_scheme .eq. 0)) then
    allocate(M_tr(nqtot))    ! allocate M_tr only if this has not already been done in another routine
    allocate(type_tr(nqtot)) ! same for type_tr
  endif

! Initialise index of aoa to 0
  i_aoa = 0
  
  do i=1, nqtot                                 
    
    write(*,*) rname,' i',i
    write(*,*) rname,' tname(i): ',tname(i)
    
    select case(tname(i))
        case('aoa')
        i_aoa=i
        write(*,*) rname,' aoa: ',i_aoa
        M_tr(i_aoa) = 44.0
        write(*,*) rname,' M_tr(i_aoa): ', 44.0
        type_tr(i_aoa) = 10
        
    end select
  end do
  
  end subroutine aoa_ini
  
end module age_of_air_mod