!> This module contains functions for MRT relaxation paramater for different
!! stencil layouts.
!! NOTE: The order of relaxation entries is consistent with moments 
!! transformation matrix used in compute kernel.
!! author: Kannan Masilamani
module mus_mrtRelaxation_module
  use env_module,     only: rk
  use tem_aux_module, only: tem_abort

  use mus_moments_type_module,  only: mus_moment_type
  use mus_scheme_header_module,      only: mus_scheme_header_type
  use mus_derivedQuantities_module2, only: getNonEqFac_intp

  implicit none
  private
 
  public :: mus_mrt_type
  public :: mus_set_mrtRelaxation
  public :: mus_alloc_mrt
  public :: mus_intp_getNonEqScalingFacs
  public :: mus_assign_mrt_ptr
  public :: mus_assign_intp_nonEqScalingFacs_ptr
  public :: mus_proc_mrt
  public :: mus_proc_intp_nonEqScalingFacs

  !> Multiple Relaxation Time (MRT) model type
  !! For D3Q19,
  !! s_mrt(2)  = omLvl_bulk(Level) * fac          --> s_bulk
  !! s_mrt(10) = s_mrt(12) = s_mrt(14:16) = omega --> s_omega
  !! s_mrt(others) = 1.0
  type mus_mrt_type
    !> Relaxation times for each moment
    real(kind=rk), allocatable :: s_mrt(:)
    !> transformed relaxation matrix-moments factor
    !! omegaMoments = (Moments^-1.s_mrt.Moments)
    !!               .(I+(Moments^-1.s_mrt.Moments)/2.0)^-1
    real(kind=rk), allocatable :: omegaMoments(:,:)
  end type mus_mrt_type
 
  !> function pointers to obtain relaxation matrix for MRT collision operator
  abstract interface
    pure function mus_proc_mrt(omegaKine, omegaBulk, QQ) result (s_mrt)
      import :: rk

      !> omega related to kinematic viscosity
      real(kind=rk), intent(in) :: omegaKine

      !> omega related to bulk viscosity for compressible model
      real(kind=rk), intent(in) :: omegaBulk
    
      !> number of directions
      integer, intent(in) :: QQ

      !> output: MRT diagonal relaxation parameter
      real(kind=rk) :: s_mrt(QQ)
    end function mus_proc_mrt  

    pure function mus_proc_intp_nonEqScalingFacs(omegaKine_SRC, omegaKine_TGT, &
      &                                 omegaBulk_SRC, omegaBulk_TGT, &
      &                                 scaleFac, QQ) result (nonEqScalingFacs)
      import :: rk

      !> source viscous omega
      real(kind=rk), intent(in) :: omegaKine_SRC
      !> target viscous omega
      real(kind=rk), intent(in) :: omegaKine_TGT
      !> source bulk viscous omega
      real(kind=rk), intent(in) :: omegaBulk_SRC
      !> target bulk viscous omega
      real(kind=rk), intent(in) :: omegaBulk_TGT
      !> factor for omega and non-conserved moments
      real(kind=rk), intent(in) :: scaleFac
      !> number of stencil directions
      integer, intent(in) :: QQ
      !> output: nonequilibrium scaling facs
      real(kind=rk) :: nonEqScalingFacs(QQ)

    end function mus_proc_intp_nonEqScalingFacs 
  end interface

contains
  
! ****************************************************************************** !
  !> Allocate mrt array
  subroutine mus_alloc_mrt( mrt, QQ )
    type(mus_mrt_type), intent(out) :: mrt
    integer, intent(in) :: QQ

    allocate( mrt%s_mrt(QQ) )
    allocate( mrt%omegaMoments(QQ,QQ) )
  end subroutine mus_alloc_mrt
! ****************************************************************************** !

  ! ************************************************************************** !
  !> This routine sets mrt relaxtion matrix according to scheme header 
  !! definition
  subroutine mus_set_mrtRelaxation(mrt, omegaKine, omegaBulk, schemeHeader, &
    &                              moment)
    ! --------------------------------------------------------------------------
    type(mus_mrt_type), intent(inout) :: mrt
    real(kind=rk), intent(in) :: omegaKine, omegaBulk
    type(mus_scheme_header_type), intent(in) :: schemeHeader
    type(mus_moment_type), intent(in) :: moment
    ! --------------------------------------------------------------------------
    integer :: iDir, QQ
    real(kind=rk) :: tmpMat(moment%toPDF%nEntries(1),moment%toPDF%nEntries(2))
    ! --------------------------------------------------------------------------
    QQ = size(mrt%s_mrt)
    ! set mrt relaxation parameters
    select case ( trim(schemeHeader%relaxation) )
    case ( 'mrt_bgk' )
      ! set all relaxation parameters to the same omega value
      mrt%s_mrt = omegaKine

    case ( 'mrt', 'mrt_explicit', 'mrt_les', 'mrt_les_explicit', 'mrt_les_new' )
      select case( trim(schemeHeader%layout) )
      case('d2q9' )
        mrt%s_mrt(:)   = mrt_d2q9(omegaKine, omegaBulk, QQ) 
      case('d3q15')
        if (trim(schemeHeader%kind) == 'lbm_incomp' ) then
          mrt%s_mrt(:) = mrt_d3q15_incomp(omegaKine, omegaBulk, QQ)
        else
          mrt%s_mrt(:) = mrt_d3q15(omegaKine, omegaBulk, QQ)
        end if

      case('d3q19')
        ! set relaxation by bulk viscostiy
        ! set to 1.19 for incompressible model
        if (trim(schemeHeader%kind) == 'lbm_incomp' ) then
          mrt%s_mrt(:) = mrt_d3q19_incomp(omegaKine, omegaBulk, QQ)
        else
          mrt%s_mrt(:) = mrt_d3q19(omegaKine, omegaBulk, QQ)
        end if
      case('d3q27')
        !! K. Suga et al. A D3Q27 MRT LBM for turbulent flows,
        !! Computers & Mathematics with Application 69 (2015) 518-529
        ! bulk viscosity
        if (trim(schemeHeader%kind) == 'lbm_incomp' ) then
          mrt%s_mrt(:) = mrt_d3q27_incomp(omegaKine, omegaBulk, QQ)
        else
          mrt%s_mrt(:) = mrt_d3q27(omegaKine, omegaBulk, QQ)
        end if
      case default
        call tem_abort('Error: Unknown layout for for mrt relaxation')
      end select ! d2q9 or d3q19?
    case default
      ! set all relaxation parameters to the same omega value
      mrt%s_mrt(:) = omegaKine
    end select ! mrt or mrt_bgk?

    tmpMat = 0.0_rk
    do iDir = 1, QQ 
      tmpMat(iDir, iDir) = mrt%s_mrt(iDir)
    end do
    mrt%omegaMoments = matmul( matmul( moment%toPDF%A, tmpMat), &
      &                        moment%toMoments%A )

  end subroutine mus_set_mrtRelaxation
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> This function returns scaling factor for nonequilibrium moments for
  !! transformation from coarser to finer elements and vice versa.
  !! To get scaling factor to convert fine to coarse, scaleFac = 2.0
  !! and to convert coarse to fine, scaleFac = 0.5.
  !! The order of return fac must be consistent with transformation matrix
  !! used in compute kernels
  pure function mus_intp_getNonEqScalingFacs( schemeHeader,              &
    & viscOmegaSRC, viscOmegaTGT, bulkOmegaSRC, bulkOmegaTGT, scaleFac, QQ )   &
    & result (nonEqScalingFacs)
    ! --------------------------------------------------------------------------
    type(mus_scheme_header_type), intent(in) :: schemeHeader
    !> source viscous omega
    real(kind=rk), intent(in) :: viscOmegaSRC
    !> target viscous omega
    real(kind=rk), intent(in) :: viscOmegaTGT
    !> source bulk viscous omega
    real(kind=rk), intent(in) :: bulkOmegaSRC
    !> target bulk viscous omega
    real(kind=rk), intent(in) :: bulkOmegaTGT
    !> factor for omega and non-conserved moments
    real(kind=rk), intent(in) :: scaleFac
    !> number of stencil directions
    integer, intent(in) :: QQ
    ! --------------------------------------------------------------------------
    real(kind=rk) :: nonEqScalingFacs(QQ)
    real(kind=rk) :: viscFac, bulkFac
    ! --------------------------------------------------------------------------
    ! kinematic viscosity fac 
    viscFac = scaleFac*getNonEqFac_intp( omegaS   = viscOmegaSRC, &
      &                                  omegaT   = viscOmegaTGT  ) 

    ! bulk omega is same in all levels for incompressible model
    if (bulkOmegaSRC /= bulkOmegaTGT) then
      ! bulk viscosity fac
      bulkFac = scaleFac*getNonEqFac_intp( omegaS   = bulkOmegaSRC, &
        &                                  omegaT   = bulkOmegaTGT  ) 
    else
      bulkFac = scaleFac
    end if

    ! Initialize all factors with scaleFac and overwrite viscous and 
    ! conserved moments according to stencil
    nonEqScalingFacs(:) = scaleFac

    ! Set non-conserved moments for mrt, viscosus moments
    select case ( trim(schemeHeader%relaxation) )
    case ( 'mrt', 'mrt_explicit', 'mrt_les', 'mrt_les_explicit', 'mrt_les_new' )
      select case(trim(schemeHeader%layout))
      case('d2q9')
        nonEqScalingFacs(2) = bulkFac
        nonEqScalingFacs(8:9) = viscFac
      case('d3q15')
        nonEqScalingFacs(2) = bulkFac
        nonEqScalingFacs(10:14) = viscFac
      case('d3q19')
        nonEqScalingFacs(2) = bulkFac
        nonEqScalingFacs(10) = viscFac
        nonEqScalingFacs(12) = viscFac
        nonEqScalingFacs(14:16) = viscFac
      case('d3q27')
        nonEqScalingFacs(5) = bulkFac
        nonEqScalingFacs(6:10) = viscFac
      end select
    case default  
      ! For non-mrt relaxation set viscFac to all except conserved moments
      nonEqScalingFacs(:) = viscFac
    end select

    ! set conserved moments to 1.0
    select case(trim(schemeHeader%layout))
    case('d2q9')
      nonEqScalingFacs(1) = 1.0_rk ! density
      nonEqScalingFacs(4) = 1.0_rk ! momX
      nonEqScalingFacs(6) = 1.0_rk ! momY
    case('d3q15','d3q19')
      nonEqScalingFacs(1) = 1.0_rk ! density
      nonEqScalingFacs(4) = 1.0_rk ! momX 
      nonEqScalingFacs(6) = 1.0_rk ! momY 
      nonEqScalingFacs(8) = 1.0_rk ! momZ 
    case('d3q27')
      nonEqScalingFacs(1:4) = 1.0_rk ! density, momX, momY, momZ
    end select

  end function mus_intp_getNonEqScalingFacs
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> This function returns mrt function pointer according to scheme definition
  function mus_assign_mrt_ptr(schemeHeader) result(mrtPtr)
    ! --------------------------------------------------------------------------
    !> Scheme header information
    type(mus_scheme_header_type), intent(in) :: schemeHeader
    !> mrt function pointer
    procedure(mus_proc_mrt), pointer :: mrtPtr
    ! --------------------------------------------------------------------------
    mrtPtr => null()
    select case (trim(schemeHeader%relaxation))
    case ( 'mrt', 'mrt_explicit', 'mrt_les', 'mrt_les_explicit', 'mrt_les_new' )
      select case (trim(schemeHeader%kind))
      case ('lbm', 'lbm_nNwtn')
        select case (trim(schemeHeader%layout))
        case ('d2q9')
          mrtPtr => mrt_d2q9
        case ('d3q15')
          mrtPtr => mrt_d3q15
        case ('d3q19')
          mrtPtr => mrt_d3q19
        case ('d3q27')
          mrtPtr => mrt_d3q27
        case default
          call tem_abort('Error: Unknown layout for mrt relaxation')
        end select
      case ('lbm_incomp', 'lbm_incomp_nNwtn')
        select case (trim(schemeHeader%layout))
        case ('d2q9')
          mrtPtr => mrt_d2q9
        case ('d3q15')
          mrtPtr => mrt_d3q15_incomp
        case ('d3q19')
          mrtPtr => mrt_d3q19_incomp
        case ('d3q27')
          mrtPtr => mrt_d3q27_incomp
        case default
          call tem_abort('Error: Unknown layout for mrt relaxation')
        end select
      case default
        call tem_abort('Error: Unknown scheme kind for mrt relaxation')
      end select
    case default 
      ! set all relaxation paramter to same omega value as default
      mrtPtr => mrt_bgk
    end select  
  end function mus_assign_mrt_ptr
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> This function returns function pointer of nonEquilibrium scaling 
  !! for interpolation according to scheme definition
  function mus_assign_intp_nonEqScalingFacs_ptr(schemeHeader) &
    & result(nEqScalFacPtr)
    ! --------------------------------------------------------------------------
    !> Scheme header information
    type(mus_scheme_header_type), intent(in) :: schemeHeader
    !> mrt function pointer
    procedure(mus_proc_intp_nonEqScalingFacs), pointer :: nEqScalFacPtr
    ! --------------------------------------------------------------------------
    nEqScalFacPtr => null()
    select case (trim(schemeHeader%relaxation))
    case ( 'mrt', 'mrt_explicit', 'mrt_les', 'mrt_les_explicit', 'mrt_les_new' )
      select case (trim(schemeHeader%kind))
      case ('lbm', 'lbm_nNwtn')
        select case (trim(schemeHeader%layout))
        case ('d2q9')
          nEqScalFacPtr => nonEqScalingFacs_d2q9
        case ('d3q15')
          nEqScalFacPtr => nonEqScalingFacs_d3q15
        case ('d3q19')
          nEqScalFacPtr => nonEqScalingFacs_d3q19
        case ('d3q27')
          nEqScalFacPtr => nonEqScalingFacs_d3q27
        case default
          call tem_abort('Error: Unknown layout for into nonEqScalingFacs')
        end select
      case ('lbm_incomp', 'lbm_incomp_nNwtn')
        select case (trim(schemeHeader%layout))
        case ('d2q9')
          nEqScalFacPtr => nonEqScalingFacs_d2q9
        case ('d3q15')
          nEqScalFacPtr => nonEqScalingFacs_d3q15_incomp
        case ('d3q19')
          nEqScalFacPtr => nonEqScalingFacs_d3q19_incomp
        case ('d3q27')
          nEqScalFacPtr => nonEqScalingFacs_d3q27_incomp
        case default
          call tem_abort('Error: Unknown layout for intp nonEqScalingFacs')
        end select
      case default
        call tem_abort('Error: Unknown scheme kind for intp nonEqScalingFacs')
      end select
    case ('bgk', 'bgk_les','bgk_les_new','bgk_pl', 'trt') 
      ! scale all non-conserved moments with viscFac
      select case (trim(schemeHeader%layout))
      case ('d2q9')
        nEqScalFacPtr => nonEqScalingFacs_bgk_d2q9
      case ('d3q15','d3q19')
        nEqScalFacPtr => nonEqScalingFacs_bgk_d3q19
      case ('d3q27')
        nEqScalFacPtr => nonEqScalingFacs_bgk_d3q27
      case default
        call tem_abort('Error: Unknown layout for intp nonEqScalingFacs')
      end select
    case default
      call tem_abort('Error: Unknown scheme relaxation for intp nonEqScalingFacs')
    end select  

  end function mus_assign_intp_nonEqScalingFacs_ptr
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> This function returns mrt relaxation diagonal matrix for d2q9 layout
  !! Parameters are taken from:
  !! Lallemand, P., & Luo, L. (2000). 
  !! Theory of the lattice boltzmann method: dispersion, dissipation, 
  !! isotropy, galilean invariance, and stability. Physical Review. E, 
  !! Statistical Physics, Plasmas, Fluids, and Related Interdisciplinary 
  !! Topics, 61(6 Pt A), 6546–62.
  !!
  !! Another paper for d2q9 with different set of parameters 
  !! Chen, S., Peng, C., Teng, Y., Wang, L. P., & Zhang, K. (2016). 
  !! Improving lattice Boltzmann simulation of moving particles in a 
  !! viscous flow using local grid refinement. Computers and Fluids, 
  !! 136, 228–246.
  pure function mrt_d2q9(omegaKine, omegaBulk, QQ) result(s_mrt)
    ! --------------------------------------------------------------------------
    !> omega related to kinematic viscosity
    real(kind=rk), intent(in) :: omegaKine
    !> omega related to bulk viscosity
    real(kind=rk), intent(in) :: omegaBulk
    !> number of directions
    integer, intent(in) :: QQ
    !> output mrt diagonal matrix
    real(kind=rk) :: s_mrt(QQ)
    ! --------------------------------------------------------------------------
    s_mrt = 0.0_rk
    ! set relaxation by bulk viscostiy
    ! set to 1.63 for both compressible and incompressible model
    ! KM: Tried omegaBulk for compressible model and it is unstable
    s_mrt(2) = 1.63_rk
    s_mrt(3)   = 1.14_rk
    s_mrt(5)   = 1.92_rk
    s_mrt(7)   = 1.92_rk
    s_mrt(8:9) = omegaKine

  end function mrt_d2q9
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> This function returns mrt relaxation diagonal matrix for d3q15 layout
  !! Parameters are taken from:
  !! D’Humières, D., Ginzburg, I., Krafczyk, M., Lallemand, P., & Luo, L.-S. 
  !! (2002). Multiple-relaxation-time lattice Boltzmann models in three 
  !! dimensions. Philosophical Transactions. Series A, Mathematical, 
  !! Physical, and Engineering Sciences, 360(1792), 437–51. 
  pure function mrt_d3q15(omegaKine, omegaBulk, QQ) result(s_mrt)
    ! --------------------------------------------------------------------------
    !> omega related to kinematic viscosity
    real(kind=rk), intent(in) :: omegaKine
    !> omega related to bulk viscosity
    real(kind=rk), intent(in) :: omegaBulk
    !> number of directions
    integer, intent(in) :: QQ
    !> output mrt diagonal matrix
    real(kind=rk) :: s_mrt(QQ)
    ! --------------------------------------------------------------------------
    s_mrt = 0.0_rk
    s_mrt(2)     = omegaBulk
    s_mrt(3)     = 1.20_rk
    s_mrt(5)     = 1.60_rk
    s_mrt(7)     = 1.60_rk
    s_mrt(9)     = 1.60_rk
    s_mrt(10:14) = omegaKine
    s_mrt(15)    = 1.20_rk

  end function mrt_d3q15  
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> This function returns mrt relaxation diagonal matrix for d3q15 layout
  !! incompressible model.
  !! Parameters are taken from:
  !! D’Humières, D., Ginzburg, I., Krafczyk, M., Lallemand, P., & Luo, L.-S. 
  !! (2002). Multiple-relaxation-time lattice Boltzmann models in three 
  !! dimensions. Philosophical Transactions. Series A, Mathematical, 
  !! Physical, and Engineering Sciences, 360(1792), 437–51. 
  pure function mrt_d3q15_incomp(omegaKine, omegaBulk, QQ) result(s_mrt)
    ! --------------------------------------------------------------------------
    !> omega related to kinematic viscosity
    real(kind=rk), intent(in) :: omegaKine
    !> omega related to bulk viscosity
    real(kind=rk), intent(in) :: omegaBulk
    !> number of directions
    integer, intent(in) :: QQ
    !> output mrt diagonal matrix
    real(kind=rk) :: s_mrt(QQ)
    ! --------------------------------------------------------------------------
    s_mrt = 0.0_rk
    s_mrt(2)     = 1.60_rk
    s_mrt(3)     = 1.20_rk
    s_mrt(5)     = 1.60_rk
    s_mrt(7)     = 1.60_rk
    s_mrt(9)     = 1.60_rk
    s_mrt(10:14) = omegaKine
    s_mrt(15)    = 1.20_rk

  end function mrt_d3q15_incomp
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> This function returns mrt relaxation diagonal matrix for d3q19 layout
  !! Parameters are taken from:
  !! D’Humières, D., Ginzburg, I., Krafczyk, M., Lallemand, P., & Luo, L.-S. 
  !! (2002). Multiple-relaxation-time lattice Boltzmann models in three 
  !! dimensions. Philosophical Transactions. Series A, Mathematical, 
  !! Physical, and Engineering Sciences, 360(1792), 437–51. 
  pure function mrt_d3q19(omegaKine, omegaBulk, QQ) result(s_mrt)
    ! --------------------------------------------------------------------------
    !> omega related to kinematic viscosity
    real(kind=rk), intent(in) :: omegaKine
    !> omega related to bulk viscosity
    real(kind=rk), intent(in) :: omegaBulk
    !> number of directions
    integer, intent(in) :: QQ
    !> output mrt diagonal matrix
    real(kind=rk) :: s_mrt(QQ)
    ! --------------------------------------------------------------------------
    s_mrt = 0.0_rk
    s_mrt(2)     = omegaBulk
    s_mrt(3)     = 1.40_rk
    s_mrt(5)     = 1.20_rk
    s_mrt(7)     = 1.20_rk
    s_mrt(9)     = 1.20_rk
    s_mrt(10)    = omegaKine
    s_mrt(11)    = 1.40_rk
    s_mrt(12)    = omegaKine
    s_mrt(13)    = 1.40_rk
    s_mrt(14:16) = omegaKine
    s_mrt(17:19) = 1.98_rk

  end function mrt_d3q19
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> This function returns mrt relaxation diagonal matrix for d3q19 layout
  !! incompressible model. omegaBulk is not used for incompressible model
  !! Parameters are taken from:
  !! D’Humières, D., Ginzburg, I., Krafczyk, M., Lallemand, P., & Luo, L.-S. 
  !! (2002). Multiple-relaxation-time lattice Boltzmann models in three 
  !! dimensions. Philosophical Transactions. Series A, Mathematical, 
  !! Physical, and Engineering Sciences, 360(1792), 437–51. 
  pure function mrt_d3q19_incomp(omegaKine, omegaBulk, QQ) result(s_mrt)
    ! --------------------------------------------------------------------------
    !> omega related to kinematic viscosity
    real(kind=rk), intent(in) :: omegaKine
    !> omega related to bulk viscosity
    real(kind=rk), intent(in) :: omegaBulk
    !> number of directions
    integer, intent(in) :: QQ
    !> output mrt diagonal matrix
    real(kind=rk) :: s_mrt(QQ)
    ! --------------------------------------------------------------------------
    s_mrt = 0.0_rk
    s_mrt(2)     = 1.19_rk
    s_mrt(3)     = 1.40_rk
    s_mrt(5)     = 1.20_rk
    s_mrt(7)     = 1.20_rk
    s_mrt(9)     = 1.20_rk
    s_mrt(10)    = omegaKine
    s_mrt(11)    = 1.40_rk
    s_mrt(12)    = omegaKine
    s_mrt(13)    = 1.40_rk
    s_mrt(14:16) = omegaKine
    s_mrt(17:19) = 1.98_rk

  end function mrt_d3q19_incomp 
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> This function returns mrt relaxation diagonal matrix for d3q27 layout
  pure function mrt_d3q27(omegaKine, omegaBulk, QQ) result(s_mrt)
    ! --------------------------------------------------------------------------
    !> omega related to kinematic viscosity
    real(kind=rk), intent(in) :: omegaKine
    !> omega related to bulk viscosity
    real(kind=rk), intent(in) :: omegaBulk
    !> number of directions
    integer, intent(in) :: QQ
    !> output mrt diagonal matrix
    real(kind=rk) :: s_mrt(QQ)
    ! --------------------------------------------------------------------------
    s_mrt = 0.0_rk
    s_mrt(5)     = omegaBulk
    s_mrt(6:10)  = omegaKine
    s_mrt(11:13) = 1.5_rk
    s_mrt(14:16) = 1.83_rk
    s_mrt(17)    = 1.4_rk
    s_mrt(18)    = 1.61_rk
    s_mrt(19:23) = 1.98_rk
    s_mrt(24:26) = 1.74_rk
    s_mrt(27)    = 1.74_rk

  end function mrt_d3q27
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> This function returns mrt relaxation diagonal matrix for d3q27 layout
  !! incompressible model. omegaBulk is not used for incompressible model
  pure function mrt_d3q27_incomp(omegaKine, omegaBulk, QQ) result(s_mrt)
    ! --------------------------------------------------------------------------
    !> omega related to kinematic viscosity
    real(kind=rk), intent(in) :: omegaKine
    !> omega related to bulk viscosity
    real(kind=rk), intent(in) :: omegaBulk
    !> number of directions
    integer, intent(in) :: QQ
    !> output mrt diagonal matrix
    real(kind=rk) :: s_mrt(QQ)
    ! --------------------------------------------------------------------------
    s_mrt = 0.0_rk
    s_mrt(5)     = 1.19_rk
    s_mrt(6:10)  = omegaKine
    s_mrt(11:13) = 1.5_rk
    s_mrt(14:16) = 1.83_rk
    s_mrt(17)    = 1.4_rk
    s_mrt(18)    = 1.61_rk
    s_mrt(19:23) = 1.98_rk
    s_mrt(24:26) = 1.74_rk
    s_mrt(27)    = 1.74_rk

  end function mrt_d3q27_incomp 
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> set all relaxation parameter to same omega, results in bgk collision
  pure function mrt_bgk(omegaKine, omegaBulk, QQ) result(s_mrt)
    ! --------------------------------------------------------------------------
    !> omega related to kinematic viscosity
    real(kind=rk), intent(in) :: omegaKine
    !> omega related to bulk viscosity
    real(kind=rk), intent(in) :: omegaBulk
    !> number of directions
    integer, intent(in) :: QQ
    !> output mrt diagonal matrix
    real(kind=rk) :: s_mrt(QQ)
    ! --------------------------------------------------------------------------
    s_mrt(:) = omegaKine
  end function mrt_bgk
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> This function returns nonequilibrium scaling factor for d2q9 layout
  pure function nonEqScalingFacs_d2q9(omegaKine_SRC, omegaKine_TGT, &
    &                            omegaBulk_SRC, omegaBulk_TGT, &
    &                            scaleFac, QQ) result (nonEqScalingFacs)
    ! --------------------------------------------------------------------------
    !> source viscous omega
    real(kind=rk), intent(in) :: omegaKine_SRC
    !> target viscous omega
    real(kind=rk), intent(in) :: omegaKine_TGT
    !> source bulk viscous omega
    real(kind=rk), intent(in) :: omegaBulk_SRC
    !> target bulk viscous omega
    real(kind=rk), intent(in) :: omegaBulk_TGT
    !> factor for omega and non-conserved moments
    real(kind=rk), intent(in) :: scaleFac
    !> number of stencil directions
    integer, intent(in) :: QQ
    !> output: nonequilibrium scaling facs
    real(kind=rk) :: nonEqScalingFacs(QQ)
    ! --------------------------------------------------------------------------
    real(kind=rk) :: viscFac
    ! --------------------------------------------------------------------------

    ! kinematic viscosity fac 
    viscFac = scaleFac*getNonEqFac_intp( omegaS   = omegaKine_SRC, &
      &                                  omegaT   = omegaKine_TGT  )

    ! Initialize all factors with scaleFac and overwrite viscous omegas 
    ! and conserved moments 
    nonEqScalingFacs(:) = scaleFac

    nonEqScalingFacs(1) = 1.0_rk ! density
    nonEqScalingFacs(4) = 1.0_rk ! momX
    nonEqScalingFacs(6) = 1.0_rk ! momY
    nonEqScalingFacs(8:9) = viscFac
  end function nonEqScalingFacs_d2q9
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> This function returns nonequilibrium scaling factor for d3q15 layout
  pure function nonEqScalingFacs_d3q15(omegaKine_SRC, omegaKine_TGT, &
    &                             omegaBulk_SRC, omegaBulk_TGT, &
    &                             scaleFac, QQ) result (nonEqScalingFacs)
    ! --------------------------------------------------------------------------
    !> source viscous omega
    real(kind=rk), intent(in) :: omegaKine_SRC
    !> target viscous omega
    real(kind=rk), intent(in) :: omegaKine_TGT
    !> source bulk viscous omega
    real(kind=rk), intent(in) :: omegaBulk_SRC
    !> target bulk viscous omega
    real(kind=rk), intent(in) :: omegaBulk_TGT
    !> factor for omega and non-conserved moments
    real(kind=rk), intent(in) :: scaleFac
    !> number of stencil directions
    integer, intent(in) :: QQ
    !> output: nonequilibrium scaling facs
    real(kind=rk) :: nonEqScalingFacs(QQ)
    ! --------------------------------------------------------------------------
    real(kind=rk) :: viscFac, bulkFac
    ! --------------------------------------------------------------------------

    ! kinematic viscosity fac 
    viscFac = scaleFac*getNonEqFac_intp( omegaS   = omegaKine_SRC, &
      &                                  omegaT   = omegaKine_TGT  )

    ! bulk viscosity fac
    bulkFac = scaleFac*getNonEqFac_intp( omegaS   = omegaBulk_SRC, &
      &                                  omegaT   = omegaBulk_TGT  ) 

    ! Initialize all factors with scaleFac and overwrite viscous omegas 
    ! and conserved moments 
    nonEqScalingFacs(:) = scaleFac

    nonEqScalingFacs(1) = 1.0_rk ! density
    nonEqScalingFacs(4) = 1.0_rk ! momX 
    nonEqScalingFacs(6) = 1.0_rk ! momY 
    nonEqScalingFacs(8) = 1.0_rk ! momZ 
    nonEqScalingFacs(2) = bulkFac
    nonEqScalingFacs(10:14) = viscFac

  end function nonEqScalingFacs_d3q15
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> This function returns nonequilibrium scaling factor for d3q19 layout
  pure function nonEqScalingFacs_d3q19(omegaKine_SRC, omegaKine_TGT, &
    &                             omegaBulk_SRC, omegaBulk_TGT, &
    &                             scaleFac, QQ) result (nonEqScalingFacs)
    ! --------------------------------------------------------------------------
    !> source viscous omega
    real(kind=rk), intent(in) :: omegaKine_SRC
    !> target viscous omega
    real(kind=rk), intent(in) :: omegaKine_TGT
    !> source bulk viscous omega
    real(kind=rk), intent(in) :: omegaBulk_SRC
    !> target bulk viscous omega
    real(kind=rk), intent(in) :: omegaBulk_TGT
    !> factor for omega and non-conserved moments
    real(kind=rk), intent(in) :: scaleFac
    !> number of stencil directions
    integer, intent(in) :: QQ
    !> output: nonequilibrium scaling facs
    real(kind=rk) :: nonEqScalingFacs(QQ)
    ! --------------------------------------------------------------------------
    real(kind=rk) :: viscFac, bulkFac
    ! --------------------------------------------------------------------------

    ! kinematic viscosity fac 
    viscFac = scaleFac*getNonEqFac_intp( omegaS   = omegaKine_SRC, &
      &                                  omegaT   = omegaKine_TGT  )

    ! bulk viscosity fac
    bulkFac = scaleFac*getNonEqFac_intp( omegaS   = omegaBulk_SRC, &
      &                                  omegaT   = omegaBulk_TGT  ) 

    ! Initialize all factors with scaleFac and overwrite viscous omegas 
    ! and conserved moments 
    nonEqScalingFacs(:) = scaleFac

    nonEqScalingFacs(1) = 1.0_rk ! density
    nonEqScalingFacs(4) = 1.0_rk ! momX 
    nonEqScalingFacs(6) = 1.0_rk ! momY 
    nonEqScalingFacs(8) = 1.0_rk ! momZ 
    nonEqScalingFacs(2) = bulkFac
    nonEqScalingFacs(10) = viscFac
    nonEqScalingFacs(12) = viscFac
    nonEqScalingFacs(14:16) = viscFac

  end function nonEqScalingFacs_d3q19
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> This function returns nonequilibrium scaling factor for d3q27 layout
  pure function nonEqScalingFacs_d3q27(omegaKine_SRC, omegaKine_TGT, &
    &                             omegaBulk_SRC, omegaBulk_TGT, &
    &                             scaleFac, QQ) result (nonEqScalingFacs)
    ! --------------------------------------------------------------------------
    !> source viscous omega
    real(kind=rk), intent(in) :: omegaKine_SRC
    !> target viscous omega
    real(kind=rk), intent(in) :: omegaKine_TGT
    !> source bulk viscous omega
    real(kind=rk), intent(in) :: omegaBulk_SRC
    !> target bulk viscous omega
    real(kind=rk), intent(in) :: omegaBulk_TGT
    !> factor for omega and non-conserved moments
    real(kind=rk), intent(in) :: scaleFac
    !> number of stencil directions
    integer, intent(in) :: QQ
    !> output: nonequilibrium scaling facs
    real(kind=rk) :: nonEqScalingFacs(QQ)
    ! --------------------------------------------------------------------------
    real(kind=rk) :: viscFac, bulkFac
    ! --------------------------------------------------------------------------

    ! kinematic viscosity fac 
    viscFac = scaleFac*getNonEqFac_intp( omegaS   = omegaKine_SRC, &
      &                                  omegaT   = omegaKine_TGT  )

    ! bulk viscosity fac
    bulkFac = scaleFac*getNonEqFac_intp( omegaS   = omegaBulk_SRC, &
      &                                  omegaT   = omegaBulk_TGT  ) 

    ! Initialize all factors with scaleFac and overwrite viscous omegas 
    ! and conserved moments 
    nonEqScalingFacs(:) = scaleFac

    nonEqScalingFacs(1) = 1.0_rk ! density
    nonEqScalingFacs(2) = 1.0_rk ! momX 
    nonEqScalingFacs(3) = 1.0_rk ! momY 
    nonEqScalingFacs(4) = 1.0_rk ! momZ 
    nonEqScalingFacs(5) = bulkFac
    nonEqScalingFacs(6:10) = viscFac

  end function nonEqScalingFacs_d3q27
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> This function returns nonequilibrium scaling factor for d3q15 layout
  !! incompressible model
  pure function nonEqScalingFacs_d3q15_incomp(omegaKine_SRC, omegaKine_TGT, &
    &                             omegaBulk_SRC, omegaBulk_TGT, &
    &                             scaleFac, QQ) result (nonEqScalingFacs)
    ! --------------------------------------------------------------------------
    !> source viscous omega
    real(kind=rk), intent(in) :: omegaKine_SRC
    !> target viscous omega
    real(kind=rk), intent(in) :: omegaKine_TGT
    !> source bulk viscous omega
    real(kind=rk), intent(in) :: omegaBulk_SRC
    !> target bulk viscous omega
    real(kind=rk), intent(in) :: omegaBulk_TGT
    !> factor for omega and non-conserved moments
    real(kind=rk), intent(in) :: scaleFac
    !> number of stencil directions
    integer, intent(in) :: QQ
    !> output: nonequilibrium scaling facs
    real(kind=rk) :: nonEqScalingFacs(QQ)
    ! --------------------------------------------------------------------------
    real(kind=rk) :: viscFac
    ! --------------------------------------------------------------------------

    ! kinematic viscosity fac 
    viscFac = scaleFac*getNonEqFac_intp( omegaS   = omegaKine_SRC, &
      &                                  omegaT   = omegaKine_TGT  )

    ! Initialize all factors with scaleFac and overwrite viscous omegas 
    ! and conserved moments 
    nonEqScalingFacs(:) = scaleFac

    nonEqScalingFacs(1) = 1.0_rk ! density
    nonEqScalingFacs(4) = 1.0_rk ! momX 
    nonEqScalingFacs(6) = 1.0_rk ! momY 
    nonEqScalingFacs(8) = 1.0_rk ! momZ 
    nonEqScalingFacs(10:14) = viscFac

  end function nonEqScalingFacs_d3q15_incomp
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> This function returns nonequilibrium scaling factor for d3q19 layout
  !! incompressible model
  pure function nonEqScalingFacs_d3q19_incomp(omegaKine_SRC, omegaKine_TGT, &
    &                             omegaBulk_SRC, omegaBulk_TGT, &
    &                             scaleFac, QQ) result (nonEqScalingFacs)
    ! --------------------------------------------------------------------------
    !> source viscous omega
    real(kind=rk), intent(in) :: omegaKine_SRC
    !> target viscous omega
    real(kind=rk), intent(in) :: omegaKine_TGT
    !> source bulk viscous omega
    real(kind=rk), intent(in) :: omegaBulk_SRC
    !> target bulk viscous omega
    real(kind=rk), intent(in) :: omegaBulk_TGT
    !> factor for omega and non-conserved moments
    real(kind=rk), intent(in) :: scaleFac
    !> number of stencil directions
    integer, intent(in) :: QQ
    !> output: nonequilibrium scaling facs
    real(kind=rk) :: nonEqScalingFacs(QQ)
    ! --------------------------------------------------------------------------
    real(kind=rk) :: viscFac
    ! --------------------------------------------------------------------------

    ! kinematic viscosity fac 
    viscFac = scaleFac*getNonEqFac_intp( omegaS   = omegaKine_SRC, &
      &                                  omegaT   = omegaKine_TGT  )

    ! Initialize all factors with scaleFac and overwrite viscous omegas 
    ! and conserved moments 
    nonEqScalingFacs(:) = scaleFac

    nonEqScalingFacs(1) = 1.0_rk ! density
    nonEqScalingFacs(4) = 1.0_rk ! momX 
    nonEqScalingFacs(6) = 1.0_rk ! momY 
    nonEqScalingFacs(8) = 1.0_rk ! momZ 
    nonEqScalingFacs(10) = viscFac
    nonEqScalingFacs(12) = viscFac
    nonEqScalingFacs(14:16) = viscFac

  end function nonEqScalingFacs_d3q19_incomp
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> This function returns nonequilibrium scaling factor for d3q27 layout
  !! incompressible model
  pure function nonEqScalingFacs_d3q27_incomp(omegaKine_SRC, omegaKine_TGT, &
    &                             omegaBulk_SRC, omegaBulk_TGT, &
    &                             scaleFac, QQ) result (nonEqScalingFacs)
    ! --------------------------------------------------------------------------
    !> source viscous omega
    real(kind=rk), intent(in) :: omegaKine_SRC
    !> target viscous omega
    real(kind=rk), intent(in) :: omegaKine_TGT
    !> source bulk viscous omega
    real(kind=rk), intent(in) :: omegaBulk_SRC
    !> target bulk viscous omega
    real(kind=rk), intent(in) :: omegaBulk_TGT
    !> factor for omega and non-conserved moments
    real(kind=rk), intent(in) :: scaleFac
    !> number of stencil directions
    integer, intent(in) :: QQ
    !> output: nonequilibrium scaling facs
    real(kind=rk) :: nonEqScalingFacs(QQ)
    ! --------------------------------------------------------------------------
    real(kind=rk) :: viscFac
    ! --------------------------------------------------------------------------

    ! kinematic viscosity fac 
    viscFac = scaleFac*getNonEqFac_intp( omegaS   = omegaKine_SRC, &
      &                                  omegaT   = omegaKine_TGT  )

    ! Initialize all factors with scaleFac and overwrite viscous omegas 
    ! and conserved moments 
    nonEqScalingFacs(:) = scaleFac

    nonEqScalingFacs(1) = 1.0_rk ! density
    nonEqScalingFacs(2) = 1.0_rk ! momX 
    nonEqScalingFacs(3) = 1.0_rk ! momY 
    nonEqScalingFacs(4) = 1.0_rk ! momZ 
    nonEqScalingFacs(6:10) = viscFac

  end function nonEqScalingFacs_d3q27_incomp
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> This function returns nonequilibrium scaling factor for bgk model
  pure function nonEqScalingFacs_bgk_d2q9(omegaKine_SRC, omegaKine_TGT, &
    &                                omegaBulk_SRC, omegaBulk_TGT, &
    &                                scaleFac, QQ) result (nonEqScalingFacs)
    ! --------------------------------------------------------------------------
    !> source viscous omega
    real(kind=rk), intent(in) :: omegaKine_SRC
    !> target viscous omega
    real(kind=rk), intent(in) :: omegaKine_TGT
    !> source bulk viscous omega
    real(kind=rk), intent(in) :: omegaBulk_SRC
    !> target bulk viscous omega
    real(kind=rk), intent(in) :: omegaBulk_TGT
    !> factor for omega and non-conserved moments
    real(kind=rk), intent(in) :: scaleFac
    !> number of stencil directions
    integer, intent(in) :: QQ
    !> output: nonequilibrium scaling facs
    real(kind=rk) :: nonEqScalingFacs(QQ)
    ! --------------------------------------------------------------------------
    real(kind=rk) :: viscFac
    ! --------------------------------------------------------------------------

    ! kinematic viscosity fac 
    viscFac = scaleFac*getNonEqFac_intp( omegaS   = omegaKine_SRC, &
      &                                  omegaT   = omegaKine_TGT  )

    ! scale all except non-conserved moments
    nonEqScalingFacs(:) = viscFac

    ! conserved moments
    nonEqScalingFacs(1) = 1.0_rk ! density
    nonEqScalingFacs(4) = 1.0_rk ! momX
    nonEqScalingFacs(6) = 1.0_rk ! momY

  end function nonEqScalingFacs_bgk_d2q9
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> This function returns nonequilibrium scaling factor for bgk model
  pure function nonEqScalingFacs_bgk_d3q19(omegaKine_SRC, omegaKine_TGT, &
    &                                 omegaBulk_SRC, omegaBulk_TGT, &
    &                                 scaleFac, QQ) result (nonEqScalingFacs)
    ! --------------------------------------------------------------------------
    !> source viscous omega
    real(kind=rk), intent(in) :: omegaKine_SRC
    !> target viscous omega
    real(kind=rk), intent(in) :: omegaKine_TGT
    !> source bulk viscous omega
    real(kind=rk), intent(in) :: omegaBulk_SRC
    !> target bulk viscous omega
    real(kind=rk), intent(in) :: omegaBulk_TGT
    !> factor for omega and non-conserved moments
    real(kind=rk), intent(in) :: scaleFac
    !> number of stencil directions
    integer, intent(in) :: QQ
    !> output: nonequilibrium scaling facs
    real(kind=rk) :: nonEqScalingFacs(QQ)
    ! --------------------------------------------------------------------------
    real(kind=rk) :: viscFac
    ! --------------------------------------------------------------------------

    ! kinematic viscosity fac 
    viscFac = scaleFac*getNonEqFac_intp( omegaS   = omegaKine_SRC, &
      &                                  omegaT   = omegaKine_TGT  )

    ! scale all except non-conserved moments
    nonEqScalingFacs(:) = viscFac

    ! conserved moments
    nonEqScalingFacs(1) = 1.0_rk ! density
    nonEqScalingFacs(4) = 1.0_rk ! momX
    nonEqScalingFacs(6) = 1.0_rk ! momY
    nonEqScalingFacs(8) = 1.0_rk ! momZ 

  end function nonEqScalingFacs_bgk_d3q19
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> This function returns nonequilibrium scaling factor for bgk model
  pure function nonEqScalingFacs_bgk_d3q27(omegaKine_SRC, omegaKine_TGT, &
    &                                 omegaBulk_SRC, omegaBulk_TGT, &
    &                                 scaleFac, QQ) result (nonEqScalingFacs)
    ! --------------------------------------------------------------------------
    !> source viscous omega
    real(kind=rk), intent(in) :: omegaKine_SRC
    !> target viscous omega
    real(kind=rk), intent(in) :: omegaKine_TGT
    !> source bulk viscous omega
    real(kind=rk), intent(in) :: omegaBulk_SRC
    !> target bulk viscous omega
    real(kind=rk), intent(in) :: omegaBulk_TGT
    !> factor for omega and non-conserved moments
    real(kind=rk), intent(in) :: scaleFac
    !> number of stencil directions
    integer, intent(in) :: QQ
    !> output: nonequilibrium scaling facs
    real(kind=rk) :: nonEqScalingFacs(QQ)
    ! --------------------------------------------------------------------------
    real(kind=rk) :: viscFac
    ! --------------------------------------------------------------------------

    ! kinematic viscosity fac 
    viscFac = scaleFac*getNonEqFac_intp( omegaS   = omegaKine_SRC, &
      &                                  omegaT   = omegaKine_TGT  )

    ! scale all except non-conserved moments
    nonEqScalingFacs(:) = viscFac

    ! conserved moments
    nonEqScalingFacs(1:4) = 1.0_rk ! density, momX, momY, momZ

  end function nonEqScalingFacs_bgk_d3q27
  ! ************************************************************************** !

end module mus_mrtRelaxation_module
