! Copyright (c) 2019-2020 Kannan Masilamani <kannan.masilamani@uni-siegen.de>
! Copyright (c) 2019 Peter Vitt <peter.vitt2@uni-siegen.de>
!
! Redistribution and use in source and binary forms, with or without
! modification, are permitted provided that the following conditions are met:
!
! 1. Redistributions of source code must retain the above copyright notice,
! this list of conditions and the following disclaimer.
!
! 2. Redistributions in binary form must reproduce the above copyright notice,
! this list of conditions and the following disclaimer in the documentation
! and/or other materials provided with the distribution.
!
! THIS SOFTWARE IS PROVIDED BY THE UNIVERSITY OF SIEGEN “AS IS” AND ANY EXPRESS
! OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
! OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
! IN NO EVENT SHALL UNIVERSITY OF SIEGEN OR CONTRIBUTORS BE LIABLE FOR ANY
! DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
! (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
! LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
! ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
! (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
! SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
!> 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
?? include 'header/lbm_macros.inc'
module mus_mrtRelaxation_module
  use env_module,       only: rk
  use tem_aux_module,   only: tem_abort
  use tem_float_module, only: operator(.fne.)

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

  implicit none
  private

  public :: mus_assign_mrt_ptr
  public :: mus_assign_intp_nonEqScalingFacs_ptr
  public :: mus_proc_mrt
  public :: mus_proc_intp_nonEqScalingFacs
  public :: mrt_d2q9
  public :: mrt_d3q19
  public :: mrt_d3q27

  abstract interface
    !> function pointers to obtain relaxation matrix for MRT collision operator
    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

    !> function pointers to obtain nonEquilibrim scaling factor to convert
    !! nonEquilibrium in interpolation routines.
    pure function mus_proc_intp_nonEqScalingFacs(omegaKine_SRC, omegaKine_TGT, &
      &                                 omegaBulk_SRC, omegaBulk_TGT, &
      &                                 scaleFac, QQ) result (nonEqScalingFacs)
      import :: rk

      !> source kinematic viscosity omega
      real(kind=rk), intent(in) :: omegaKine_SRC
      !> target kinemaitc viscosity omega
      real(kind=rk), intent(in) :: omegaKine_TGT
      !> source bulk viscosity omega
      real(kind=rk), intent(in) :: omegaBulk_SRC
      !> target bulk viscosity 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


  ! ************************************************************************** !
  !> 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 viscosity omega
    real(kind=rk), intent(in) :: viscOmegaSRC
    !> target viscosity omega
    real(kind=rk), intent(in) :: viscOmegaTGT
    !> source bulk viscosity omega
    real(kind=rk), intent(in) :: bulkOmegaSRC
    !> target bulk viscosity 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 
??  if (PUSH) then ! for pre-collision PDF
      viscFac = scaleFac * viscOmegaSRC / viscOmegaTGT
??  else ! for post-collision PDF
      viscFac = scaleFac * viscOmegaSRC * (1.0_rk - viscOmegaTGT) &
        &     / ( (1.0_rk - viscOmegaSRC ) * viscOmegaTGT)
??  endif

    ! bulk omega is same in all levels for incompressible model
    if (bulkOmegaSRC .fne. bulkOmegaTGT) then
      ! bulk viscosity fac
??  if (PUSH) then ! for pre-collision PDF
      bulkFac = scaleFac * bulkOmegaSRC / bulkOmegaTGT
??  else ! for post-collision PDF
      bulkFac = scaleFac * bulkOmegaSRC * (1.0_rk - bulkOmegaTGT) &
        &     / ( (1.0_rk - bulkOmegaSRC ) * bulkOmegaTGT )
??  endif
    else
      bulkFac = scaleFac
    end if

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

    ! Set non-conserved moments for mrt, viscosus moments
    select case ( trim(schemeHeader%relaxation) )
    case ( 'mrt', 'mrt_generic' )
      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.
  !! In Jonas Toelke paper (2006) about MRT, the following notions are used:\n
  !!  s(a) = s(2)
  !!  s(b) = s(3)
  !!  s(c) = s(5) = s(7) = s(9)
  !!  s(d) = s(11) = s(13
  !!  s(e) = s(17) = s(18) = s(19)
  !!  s(w) = s(10) = s(12) = s(14) = s(15) = s(16)
  !! It is suggested that, for D3Q19,
  !!  s(a) = s(b) = s(c) = s(d) = s(e) = max( s(w), -1.0 )
  !! Notice that the collision matrix S used in this papar corresponds to
  !! -omega in BGK model, because it express the LB equation is slightly
  !! different way.
  !!
  subroutine mus_assign_mrt_ptr(mrtPtr, schemeHeader) 
    ! --------------------------------------------------------------------------
    !> mrt function pointer
    procedure(mus_proc_mrt), pointer :: mrtPtr
    !> Scheme header information
    type(mus_scheme_header_type), intent(in) :: schemeHeader
    ! --------------------------------------------------------------------------
    mrtPtr => null()
    select case (trim(schemeHeader%relaxation))
    case ( 'mrt', 'mrt_generic' )
      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 default
      ! set all relaxation paramter to same omega value as default
      mrtPtr => mrt_bgk
    end select
  end subroutine mus_assign_mrt_ptr
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> This function returns function pointer of nonEquilibrium scaling 
  !! for interpolation according to scheme definition
  subroutine mus_assign_intp_nonEqScalingFacs_ptr(nEqScalFacPtr, schemeHeader)
    ! --------------------------------------------------------------------------
    !> mrt function pointer
    procedure(mus_proc_intp_nonEqScalingFacs), pointer :: nEqScalFacPtr
    !> Scheme header information
    type(mus_scheme_header_type), intent(in) :: schemeHeader
    ! --------------------------------------------------------------------------
    nEqScalFacPtr => null()
    select case (trim(schemeHeader%relaxation))
    case ( 'mrt', 'mrt_generic' )
      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 default
      ! For all not-mrt relaxation, 
      ! 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
    end select

  end subroutine 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 
  !! viscosity 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 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 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
  ! ************************************************************************** !


  ! ************************************************************************** !
  !> 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 viscosity omega
    real(kind=rk), intent(in) :: omegaKine_SRC
    !> target viscosity omega
    real(kind=rk), intent(in) :: omegaKine_TGT
    !> source bulk viscosity omega
    real(kind=rk), intent(in) :: omegaBulk_SRC
    !> target bulk viscosity 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
??  if (PUSH) then ! for pre-collision PDF
      viscFac = scaleFac * omegaKine_SRC / omegaKine_TGT
??  else ! for post-collision PDF
      viscFac = scaleFac * omegaKine_SRC * (1.0_rk - omegaKine_TGT) &
        &     / ( (1.0_rk - omegaKine_SRC ) * omegaKine_TGT )
??  end if

    ! Initialize all factors with scaleFac and overwrite viscosity 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 viscosity omega
    real(kind=rk), intent(in) :: omegaKine_SRC
    !> target viscosity omega
    real(kind=rk), intent(in) :: omegaKine_TGT
    !> source bulk viscosity omega
    real(kind=rk), intent(in) :: omegaBulk_SRC
    !> target bulk viscosity 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 
??  if (PUSH) then ! for pre-collision PDF
      viscFac = scaleFac * omegaKine_SRC / omegaKine_TGT
??  else ! for post-collision PDF
      viscFac = scaleFac * omegaKine_SRC * (1.0_rk - omegaKine_TGT) &
        &     / ( (1.0_rk - omegaKine_SRC ) * omegaKine_TGT )
??  end if

    ! bulk viscosity fac
??  if (PUSH) then ! for pre-collision PDF
      bulkFac = scaleFac * omegaBulk_SRC / omegaBulk_TGT
??  else ! for post-collision PDF
      bulkFac = scaleFac * omegaBulk_SRC * (1.0_rk - omegaBulk_TGT) &
        &     / ( (1.0_rk - omegaBulk_SRC ) * omegaBulk_TGT )
??  end if

    ! Initialize all factors with scaleFac and overwrite viscosity 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 viscosity omega
    real(kind=rk), intent(in) :: omegaKine_SRC
    !> target viscosity omega
    real(kind=rk), intent(in) :: omegaKine_TGT
    !> source bulk viscosity omega
    real(kind=rk), intent(in) :: omegaBulk_SRC
    !> target bulk viscosity 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 
??  if (PUSH) then ! for pre-collision PDF
      viscFac = scaleFac * omegaKine_SRC / omegaKine_TGT
??  else ! for post-collision PDF
      viscFac = scaleFac * omegaKine_SRC * (1.0_rk - omegaKine_TGT) &
        &     / ( (1.0_rk - omegaKine_SRC ) * omegaKine_TGT )
??  end if

    ! bulk viscosity fac
??  if (PUSH) then ! for pre-collision PDF
      bulkFac = scaleFac * omegaBulk_SRC / omegaBulk_TGT
??  else ! for post-collision PDF
      bulkFac = scaleFac * omegaBulk_SRC * (1.0_rk - omegaBulk_TGT) &
        &     / ( (1.0_rk - omegaBulk_SRC ) * omegaBulk_TGT )
??  end if

    ! Initialize all factors with scaleFac and overwrite viscosity 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 viscosity omega
    real(kind=rk), intent(in) :: omegaKine_SRC
    !> target viscosity omega
    real(kind=rk), intent(in) :: omegaKine_TGT
    !> source bulk viscosity omega
    real(kind=rk), intent(in) :: omegaBulk_SRC
    !> target bulk viscosity 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 
??  if (PUSH) then ! for pre-collision PDF
      viscFac = scaleFac * omegaKine_SRC / omegaKine_TGT
??  else ! for post-collision PDF
      viscFac = scaleFac * omegaKine_SRC * (1.0_rk - omegaKine_TGT) &
        &     / ( (1.0_rk - omegaKine_SRC ) * omegaKine_TGT )
??  end if

    ! bulk viscosity fac
??  if (PUSH) then ! for pre-collision PDF
      bulkFac = scaleFac * omegaBulk_SRC / omegaBulk_TGT
??  else ! for post-collision PDF
      bulkFac = scaleFac * omegaBulk_SRC * (1.0_rk - omegaBulk_TGT) &
        &     / ( (1.0_rk - omegaBulk_SRC ) * omegaBulk_TGT )
??  end if

    ! Initialize all factors with scaleFac and overwrite viscosity 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 bgk model
  pure function nonEqScalingFacs_bgk_d2q9(omegaKine_SRC, omegaKine_TGT, &
    &                                omegaBulk_SRC, omegaBulk_TGT, &
    &                                scaleFac, QQ) result (nonEqScalingFacs)
    ! --------------------------------------------------------------------------
    !> source viscosity omega
    real(kind=rk), intent(in) :: omegaKine_SRC
    !> target viscosity omega
    real(kind=rk), intent(in) :: omegaKine_TGT
    !> source bulk viscosity omega
    real(kind=rk), intent(in) :: omegaBulk_SRC
    !> target bulk viscosity 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 
??  if (PUSH) then ! for pre-collision PDF
      viscFac = scaleFac * omegaKine_SRC / omegaKine_TGT
??  else ! for post-collision PDF
      viscFac = scaleFac * omegaKine_SRC * (1.0_rk - omegaKine_TGT) &
        &     / ( (1.0_rk - omegaKine_SRC ) * omegaKine_TGT )
??  end if

    ! 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 viscosity omega
    real(kind=rk), intent(in) :: omegaKine_SRC
    !> target viscosity omega
    real(kind=rk), intent(in) :: omegaKine_TGT
    !> source bulk viscosity omega
    real(kind=rk), intent(in) :: omegaBulk_SRC
    !> target bulk viscosity 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
??  if (PUSH) then ! for pre-collision PDF
      viscFac = scaleFac * omegaKine_SRC / omegaKine_TGT
??  else ! for post-collision PDF
      viscFac = scaleFac * omegaKine_SRC * (1.0_rk - omegaKine_TGT) &
        &     / ( (1.0_rk - omegaKine_SRC ) * omegaKine_TGT )
??  end if

    ! 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 viscosity omega
    real(kind=rk), intent(in) :: omegaKine_SRC
    !> target viscosity omega
    real(kind=rk), intent(in) :: omegaKine_TGT
    !> source bulk viscosity omega
    real(kind=rk), intent(in) :: omegaBulk_SRC
    !> target bulk viscosity 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
??  if (PUSH) then ! for pre-collision PDF
      viscFac = scaleFac * omegaKine_SRC / omegaKine_TGT
??  else ! for post-collision PDF
      viscFac = scaleFac * omegaKine_SRC * (1.0_rk - omegaKine_TGT) &
        &     / ( (1.0_rk - omegaKine_SRC ) * omegaKine_TGT )
??  end if

    ! 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
