MODULE mp2m_diagnostics

!=======================================================================
!   Purpose
!   -------
!   Microphysics diagnostics variables module.
!
!   The module contains all the variables related to output diagnostics of the microphysics two moments.
!   Such variables are not (and should not be) used as input in the GCM except for output data writing.
!
!   The module also contains three methods:
!      - inimuphy_diag(ngrid,nlayer,nq,pq,int2ext)
!      - alpha_s(k)
!      - alpha_f(k)
!
!   Authors
!   -------
!      Bruno de Batz de Trenquelléon (11/2024)
!=======================================================================

  USE tracer_h
  
  IMPLICIT NONE

  REAL(kind=8), ALLOCATABLE, DIMENSION(:)     :: mp2m_aer_s_prec ! Spherical aerosols precipitations (kg.m-2.s-1).
  REAL(kind=8), ALLOCATABLE, DIMENSION(:)     :: mp2m_aer_f_prec ! Fractal aerosols precipitations (kg.m-2.s-1).
  REAL(kind=8), ALLOCATABLE, DIMENSION(:,:)   :: mp2m_aer_s_w    ! Spherical aerosol settling velocity (m.s-1).
  REAL(kind=8), ALLOCATABLE, DIMENSION(:,:)   :: mp2m_aer_f_w    ! Fractal aerosol settling velocity (m.s-1).
  REAL(kind=8), ALLOCATABLE, DIMENSION(:,:)   :: mp2m_aer_s_flux ! Spherical aerosol mass flux (kg.m-2.s-1).
  REAL(kind=8), ALLOCATABLE, DIMENSION(:,:)   :: mp2m_aer_f_flux ! Fractal aerosol mass flux (kg.m-2.s-1).
  REAL(kind=8), ALLOCATABLE, DIMENSION(:,:)   :: mp2m_rc_sph     ! Spherical mode characteristic radius, i.e. bulk radius (m).
  REAL(kind=8), ALLOCATABLE, DIMENSION(:,:)   :: mp2m_rc_fra     ! Fractal mode characteristic radius, i.e. bulk radius (m).
  
  !$OMP THREADPRIVATE(mp2m_aer_s_prec,mp2m_aer_f_prec,mp2m_aer_s_w,mp2m_aer_f_w,mp2m_aer_s_flux,mp2m_aer_f_flux)
  !$OMP THREADPRIVATE(mp2m_rc_sph,mp2m_rc_fra)

  CONTAINS

  SUBROUTINE inimufi_diag(ngrid,nlayer,nq,pq,int2ext)
    !! Initialize the variables associated to microphysics diagnostics.
    !!
    use callkeys_mod, only : haze_rc_prod, haze_rm

    INTEGER, INTENT(in) :: ngrid  ! Number of points of the horizontal grid.
    INTEGER, INTENT(in) :: nlayer ! Number of points of the vertical grid (layers).
    INTEGER, INTENT(in) :: nq     ! Number of advected fields.

    REAL,INTENT(in) :: pq(ngrid,nlayer,nq)   ! Tracers (X/kg_of_air).
    REAL,INTENT(in) :: int2ext(ngrid,nlayer) ! Intensive to extensive factor (kg_air/m3).

    ! Local variables:
    !~~~~~~~~~~~~~~~~~
    REAL, DIMENSION(:,:), ALLOCATABLE :: m0as ! 0th order moment of the spherical mode (m-3).
    REAL, DIMENSION(:,:), ALLOCATABLE :: m3as ! 3rd order moment of the spherical mode (m3.m-3).
    REAL, DIMENSION(:,:), ALLOCATABLE :: m0af ! 0th order moment of the fractal mode (m-3).
    REAL, DIMENSION(:,:), ALLOCATABLE :: m3af ! 3rd order moment of the fractal mode (m3.m-3).
    
    IF (.NOT.ALLOCATED(mp2m_aer_s_prec)) ALLOCATE(mp2m_aer_s_prec(ngrid))
    IF (.NOT.ALLOCATED(mp2m_aer_f_prec)) ALLOCATE(mp2m_aer_f_prec(ngrid))
    IF (.NOT.ALLOCATED(mp2m_aer_s_w))    ALLOCATE(mp2m_aer_s_w(ngrid,nlayer))
    IF (.NOT.ALLOCATED(mp2m_aer_f_w))    ALLOCATE(mp2m_aer_f_w(ngrid,nlayer))
    IF (.NOT.ALLOCATED(mp2m_aer_s_flux)) ALLOCATE(mp2m_aer_s_flux(ngrid,nlayer))
    IF (.NOT.ALLOCATED(mp2m_aer_f_flux)) ALLOCATE(mp2m_aer_f_flux(ngrid,nlayer))
    IF (.NOT.ALLOCATED(mp2m_rc_sph))     ALLOCATE(mp2m_rc_sph(ngrid,nlayer))
    IF (.NOT.ALLOCATED(mp2m_rc_fra))     ALLOCATE(mp2m_rc_fra(ngrid,nlayer))
    
    mp2m_aer_s_prec(:)   = 0d0
    mp2m_aer_f_prec(:)   = 0d0
    mp2m_aer_s_w(:,:)    = 0d0
    mp2m_aer_f_w(:,:)    = 0d0
    mp2m_aer_s_flux(:,:) = 0d0
    mp2m_aer_f_flux(:,:) = 0d0
    
    ! Initializes the radius to their first value (useful for radiative transfer).
    m0as = pq(:,:,micro_indx(1)) * int2ext
    m3as = pq(:,:,micro_indx(2)) * int2ext
    m0af = pq(:,:,micro_indx(3)) * int2ext
    m3af = pq(:,:,micro_indx(4)) * int2ext

    WHERE(m0as > 1e-10 .AND. m3as > (1e-10*alpha_s(3.)*haze_rc_prod**3))
      mp2m_rc_sph = (m3as / (m0as*alpha_s(3.)))**(1./3.)
    ELSEWHERE
      mp2m_rc_sph = 0d0
    ENDWHERE

    WHERE(m0af > 1e-10 .AND. m3af > (1e-10*alpha_f(3.)*haze_rm**3))
      mp2m_rc_fra = (m3af / (m0af*alpha_f(3.)))**(1./3.)
    ELSEWHERE
      mp2m_rc_fra = 0d0
    ENDWHERE
  
    CONTAINS

    FUNCTION alpha_s(k) RESULT (res)
      !! Inter-moment relation for spherical aerosols size distribution law.
      !! Mk / M0 = rc^k . alpha(k)
      !!
      use mp2m_intgcm
      REAL, INTENT(in) :: k ! k order of the moment.
      REAL :: sigma         ! Standard deviation.
      REAL :: res           ! Alpha value.
    
      ! Titan's case
      !~~~~~~~~~~~~~
      ! res = SUM(dexp(mm_asp%a*k**2 + mm_asp%b*k + mm_asp%c))

      ! Pluto's case
      !~~~~~~~~~~~~~
      sigma = 0.2
      res = exp(k**2 * sigma**2 / 2.)
    END FUNCTION alpha_s

    FUNCTION alpha_f(k) RESULT (res)
      !! Inter-moment relation for fractal aerosols size distribution law.
      !! Mk / M0 = rc^k . alpha(k)
      !!
      use mp2m_intgcm
      REAL, INTENT(in) :: k ! k order of the moment.
      REAL :: sigma         ! Standard deviation.
      REAL :: res           ! Alpha value.
    
      ! Titan's case
      !~~~~~~~~~~~~~
      ! res = SUM(dexp(mm_asp%a*k**2 + mm_asp%b*k + mm_asp%c))

      ! Pluto's case
      !~~~~~~~~~~~~~
      sigma = 0.35
      res = exp(k**2 * sigma**2 / 2.)
    END FUNCTION alpha_f
  END SUBROUTINE inimufi_diag

END MODULE mp2m_diagnostics