! See copyright notice in the COPYRIGHT file.
! ****************************************************************************** !
!> This module keeps all information about the fluid
!!
!! In this module, all infos about the fluid is collected.
!! The file type [[mus_fluid_type]] contains all relevant information.
!! This includes physical parameters such as the viscosity and reference
!! density. Also, LBM-specific parameters such as the relaxation rates are
!! defined.
!!
module mus_fluid_module

  ! include treelm modules
  use env_module,           only: rk, single_k, eps_single_k, labelLen
  use tem_tools_module,     only: tem_horizontalSpacer
  use tem_aux_module,       only: tem_abort
  use tem_param_module,     only: cs2, cs2inv, rho0
  use tem_logging_module,   only: logUnit
  use tem_time_module,      only: tem_time_type
  use tem_temporal_module,  only: tem_temporal_for, &
    &                             tem_temporal_type, tem_load_temporal
  use tem_spacetime_fun_module, only: tem_spacetime_fun_type, &
    &                                 tem_load_spacetime
  use tem_construction_module, only: tem_levelDesc_type
  use tem_grow_array_module,   only: init

  ! include aotus modules
  use aotus_module,     only: flu_State, aot_get_val, aoterr_Fatal,            &
    &                         aoterr_NonExistent, aoterr_WrongType
  use aot_table_module, only: aot_table_open, aot_table_close
  use aot_out_module,   only: aot_out_type, aot_out_val, aot_out_open_table,   &
    &                         aot_out_close_table

  ! include musubi modules
  use mus_physics_module,       only: mus_physics_type
  use mus_nonNewtonian_module,  only: mus_nNwtn_type, mus_nNwtn_load, &
    &                                 mus_nNwtn_dump2outUnit, mus_nNwtn_save2lua
  use mus_scheme_header_module, only: mus_scheme_header_type  
  use mus_moments_type_module,  only: mus_moment_type
  use mus_derivedQuantities_module2, only: getNonEqFac_intp
  use mus_relaxationParam_module, only: mus_viscosity_type,              &
    &                                   mus_init_relaxParam,             &
    &                                   mus_update_relaxParamFromViscSTfun
  use mus_mrtRelaxation_module,   only: mus_mrt_type, mus_proc_mrt,     &
    &                                   mus_set_mrtRelaxation,          &
    &                                   mus_proc_intp_nonEqScalingFacs, &
    &                                   mus_assign_mrt_ptr,                &
    &                                   mus_assign_intp_nonEqScalingFacs_ptr

  implicit none

  private

  public :: mus_fluid_type
  public :: mus_load_fluid
  public :: mus_fluid_save2lua
  ! public :: set_omegas_acoustic
  ! public :: set_omegas_diffusive
  public :: mus_update_relaxation
  public :: mus_alloc_mrt
  public :: mus_init_fluid
  public :: mus_init_fluid_new

  !> Large Eddy Turbulence (LES) model type
  type les_type
    real(kind=rk) :: c_s = 0.17_rk     !< Smagorinsky constant
  end type les_type

  !> collection of properties of the fluid
  type mus_fluid_type

    logical :: active = .false.       !< is this object a fluid?

    !> Value of omage for minimum level
    real(kind=rk) :: omega = 1.7_rk

    !> Magic value for TRT collision model
    !! Lambda = ( 1/omega_+ - 0.5 ) * ( 1/omega_- - 0.5 )
    real(kind=rk) :: lambda = 0.25_rk

    !> level-wise reference omegas
    !! allocated in mus_init_fluid
    !! Set in routine: set_omegas_acoustic or set_omegas_diffusive
    real(kind=rk), allocatable :: omLvl(:)
    !> level-wise bulk omegas, used as relaxation in mrt model
    real(kind=rk), allocatable :: omLvl_bulk(:)
    ! \todo KM: move dx and dt level into separate type and store in param_type
    !> dt in LB unit, dtLvl(minLevel) = 1.0
    real(kind=rk), allocatable :: dtLvl(:)
    !> dx in LB unit, dxLvl(minLevel) = 1.0
    real(kind=rk), allocatable :: dxLvl(:)
    !> Level wise viscosity
    real(kind=rk), allocatable :: visc(:)
    !> Level wise bulk viscosity
    real(kind=rk), allocatable :: visc_bulk(:)

    !> negative omega for TRT
    real(kind=rk), allocatable :: wN(:)

    real(kind=rk) :: kine_viscosityLB  !< lattice kinematic shear viscosity 
    real(kind=rk) :: kine_viscosity    !< physical kinematic shear viscosity 

    !> lattice bulk viscosity for each level
    real(kind=rk) :: bulk_viscosityLB   !< lattice bulk viscosity
    real(kind=rk) :: bulk_viscosity     !< physical bulk viscosity

    !> background velocity for acoustic equation and lee
    real(kind=rk) :: u0_x   = 0.0_rk
    real(kind=rk) :: u0_y   = 0.0_rk
    real(kind=rk) :: u0_z   = 0.0_rk

    !> background theta0 for acoustic equation and lee
    real(kind=rk) :: theta0 = 0.0_rk

    real(kind=rk) :: force(3)        !< forcing term

    !> Multi Relaxation parameters
    type( mus_mrt_type ), allocatable :: mrt(:)
    !> Large Eddy Simulation parameters
    type( les_type ) :: les

    !> nonNewtonian fluid parameter
    type( mus_nNwtn_type ) :: nNwtn

    !> temporal ramping for visoscity
    type(tem_temporal_type) :: visc_ramping

    !> function pointer to get MRT diagonal relaxation matrix
    procedure(mus_proc_mrt), nopass, pointer :: mrtPtr => null()

    !> function pointer to get MRT diagonal relaxation matrix
    procedure(mus_proc_intp_nonEqScalingFacs), nopass, pointer &
      & :: nonEqScalingFacs => null()

    !> is true only for weakly compressible model with MRT collision and 
    !! 3 dimension layout
    !! i.e. scheme kind is lbm and scheme relaxation is mrt
    logical :: requireOmegaBulk = .false.

    !> kinematic viscosity
    !! \todo KM: implement interpolation routine for constant viscosity
    type(mus_viscosity_type) :: viscKine

    !> Bulk viscosity
    !! KM: Use constant bulk viscosity
    !type(mus_viscosity_type) :: viscBulk

  end type mus_fluid_type

contains

! ****************************************************************************** !
  !> Read in the fluid property from the LUA parameter file
  subroutine mus_load_fluid( me, conf, parent, minLevel, physics, schemeKind )
    !--------------------------------------------------------------------------
    !> fluid type
    type( mus_fluid_type ),intent(out) :: me
    !> lua state
    type( flu_state )              :: conf
    !> parent handle
    integer, intent(in), optional  :: parent
    !> global pdf info
    integer, intent(in) :: minLevel
    !> physics type to convert physics to lattice unit or vice versa
    type( mus_physics_type ), intent(in) :: physics
    !> scheme kind
    character(len=*), intent(in) :: schemeKind
    ! --------------------------------------------------------------------------
    integer :: fluid_table
    integer :: iError
    integer :: vecError(3)
    ! ---------------------------------------------------------------------------


    ! if fluid informations in scheme table parentHandle /= 0
    if( present(parent) ) then
      call aot_table_open( L       = conf,                                     &
        &                  parent  = parent,                                   &
        &                  thandle = fluid_table,                              &
        &                  key     = 'fluid' )
    else
      call aot_table_open( L=conf, thandle = fluid_table, key = 'fluid' )
    end if

    if ( fluid_table == 0 ) then
      write(logUnit(1),*)'No fluid table defined'
      call tem_abort()
    endif

    write(logUnit(1),*) 'Loading fluid informations'
    ! Activating the current fluid object
    me%active = .true.

    ! load lambda (magic value for trt model)
    call aot_get_val(L       = conf,                                           &
      &              thandle = fluid_table,                                    &
      &              key     = 'lambda',                                       &
      &              val     = me%lambda,                                      &
      &              ErrCode = iError )

    ! load omega
    call aot_get_val(L       = conf,                                           &
      &              thandle = fluid_table,                                    &
      &              key     = 'omega',                                        &
      &              val     = me%omega,                                       &
      &              ErrCode = iError )

    if (btest(iError, aoterr_NonExistent)) then
      ! omega not defined, try load shear viscosity
      write(logUnit(1),*) 'omega not defined. Load kine_shear_viscosity'
      call aot_get_val(L = conf, thandle = fluid_table,                        &
        &              key = 'kine_shear_viscosity',                           &
        &              val = me%kine_viscosity,                                &
        &              ErrCode = iError)
      if (btest(iError, aoterr_NonExistent)) then 
        write(logUnit(1),*) 'ERROR: specify either omega or kine_shear_viscosity'
        call tem_abort()
      end if
      ! convert physical viscosity to LB viscosity, and calculate omega
      me%kine_viscosityLB = me%kine_viscosity/physics%fac(minLevel)%visc
      me%omega = 1._rk / ( cs2inv*me%kine_viscosityLB + 0.5_rk )
    else
      ! calculate LB viscosity from omega, and convert it to physical viscosity
      me%kine_viscosityLB = ( 1._rk/me%omega - 0.5_rk )*cs2
      me%kine_viscosity = me%kine_viscosityLB*physics%fac(minLevel)%visc
    endif

    call aot_get_val( L       = conf,                                          &
      &               thandle = fluid_table,                                   &
      &               key     = 'les_cs',                                      &
      &               val     = me%les%c_s,                                    &
      &               default = 0.17_rk,                                       &
      &               ErrCode = iError )

    ! get bulk viscosity and convert it to LB
    call aot_get_val( L       = conf,                                          &
      &               thandle = fluid_table,                                   &
      &               key     = 'bulk_viscosity',                              &
      &               val     = me%bulk_viscosity,                             &
      &               ErrCode = iError )
    me%bulk_viscosityLB = me%bulk_viscosity/physics%fac(minLevel)%visc

    ! calculate bulk viscosity from kinetic viscosity
    if (btest(iError, aoterr_NonExistent)) then
      me%bulk_viscosityLB = 2._rk / 3._rk * me%kine_viscosityLB
      me%bulk_viscosity = me%bulk_viscosityLB                                  &
        &               * physics%fac(minLevel)%visc
    endif

    ! load forces
    call aot_get_val( L       = conf,                                          &
      &               thandle = fluid_table,                                   &
      &               key     = 'force',                                       &
      &               val     = me%force,                                      &
      &               default = [0._rk,0._rk,0._rk],                           &
      &               ErrCode = vecError )

    ! convert force to lattice
    me%force = me%force / physics%fac(minLevel)%force

    ! load nonNewtonian fluid feature
    call mus_nNwtn_load( me     = me%nNwtn,   &
      &                  conf   = conf,       &
      &                  parent = fluid_table )

    ! Omega ramping
    call tem_load_temporal( me     = me%visc_ramping, &
      &                     conf   = conf,            &
      &                     parent = fluid_table,     &
      &                     key    = 'visc_ramping'   )

    ! load u0_x
    call aot_get_val(L       = conf,                                           &
      &              thandle = fluid_table,                                    &
      &              key     = 'u0_x',                                        &
      &              val     = me%u0_x,                                       &
      &              default = 0._rk,                                         &
      &              ErrCode = iError )

    ! load u0_y
    call aot_get_val(L       = conf,                                           &
      &              thandle = fluid_table,                                    &
      &              key     = 'u0_y',                                        &
      &              val     = me%u0_y,                                       &
      &              default = 0._rk,                                         &
      &              ErrCode = iError )

    ! load u0_z
    call aot_get_val(L       = conf,                                           &
      &              thandle = fluid_table,                                    &
      &              key     = 'u0_z',                                        &
      &              val     = me%u0_z,                                       &
      &              default = 0._rk,                                         &
      &              ErrCode = iError )

    ! load theta0
    call aot_get_val(L       = conf,        &
      &              thandle = fluid_table, &
      &              key     = 'theta0',    &
      &              val     = me%theta0,   &
      &              default = 1._rk/3._rk, &
      &              ErrCode = iError       )

    ! load kinematic viscosity as spacetime function
    call tem_load_spacetime( me      = me%viscKine%STfun,      &
      &                      conf    = conf,                   &
      &                      parent  = fluid_table,            &
      &                      key     = 'kinematic_viscosity',  &
      &                      errCode = iError                  )

    ! load bulk viscosity only for compressible model
!KM!    if (trim(schemeKind)=='lbm' .or. trim(schemeKind)=='lbm_nNwtn') then
!KM!      me%requireOmegaBulk = .true.
!KM!      call tem_load_spacetime( me      = me%viscBulk%STfun,      &
!KM!        &                      conf    = conf,                   &
!KM!        &                      parent  = fluid_table,            &
!KM!        &                      key     = 'bulk_viscosity_new',   &
!KM!        &                      errCode = iError                  )
!KM!    else
!KM!      me%requireOmegaBulk = .false.
!KM!    end if

    call aot_table_close( L=conf, thandle=fluid_table )
    call tem_horizontalSpacer(fUnit = logUnit(1))

    call mus_fluid_dump( me, logUnit(1) )

  end subroutine mus_load_fluid
  ! **************************************************************************** !

  ! **************************************************************************** !
  subroutine mus_fluid_dump( me, outUnit )
    ! --------------------------------------------------------------------------
    type( mus_fluid_type ), intent(inout) :: me !< fluid type
    integer, intent(in) :: outUnit
    ! --------------------------------------------------------------------------

    write(outUnit,"(A)")       'Fluid properties:'
    write(outUnit,"(A,F10.7)") '  Relaxation Parameter Omega: ', me%omega
    write(outUnit,"(A,F10.7)") '  Reference density rho0: ', rho0
    write(outUnit,"(A,F10.7)") '  Shear viscosity Physical: ', me%kine_viscosity
    write(outUnit,"(A,F10.7)") '  Shear viscosity LB: ', me%kine_viscosityLB
    write(outUnit,"(A,F10.7)") '  Bulk viscosity physical:   ', me%bulk_viscosity
    write(outUnit,"(A,F10.7)") '  Bulk viscosityLB: ', me%bulk_viscosityLB
    write(outUnit,"(A,F10.7)") '  Smagorinsky Cs: ', me%les%c_s
    if( maxval( abs( me%force )) > 0._rk )  &
      & write(logUnit(1),*) ' Forcing:             ', real(me%force)

    ! Dump nonNewtonian parameters to outUnit
    call mus_nNwtn_dump2outUnit( me%nNwtn, outUnit )

  end subroutine mus_fluid_dump
  ! **************************************************************************** !

! ****************************************************************************** !
  !> Set the omegas according to the time step setting
  !! LB viscosity is calculated by:
  !!  nu = ( 1/omega - 0.5 ) / 3 * ( (dx)^2 / dt )
  !! Assuming:
  !!   me%omega, me%bulk_viscosityLB
  !!   physics%fac(minLevel:maxLevel)%visc
  !! are already set before calling this routine.
  subroutine set_omegas_acoustic( me, physics, minLevel, maxLevel )
    ! --------------------------------------------------------------------------
    type( mus_fluid_type ),intent(inout) :: me
    !> physics type to convert physLB unit
    type( mus_physics_type ), intent(in) :: physics
    integer, intent(in) :: minLevel
    integer, intent(in) :: maxLevel
    ! --------------------------------------------------------------------------
    integer :: iLevel
    ! --------------------------------------------------------------------------

    ! set LB viscosity and omega for each level
    do iLevel = minLevel, maxLevel
      me%dtLvl( iLevel ) = 1._rk / (2**(iLevel-minLevel))
      me%dxLvl( iLevel ) = 1._rk / (2**(iLevel-minLevel))

      ! omega is adjusted to maintain constant LB viscosity over levels
      me%omLvl( iLevel ) = 1._rk / &
        &  ( 0.5_rk + (1._rk/me%omega - 0.5_rk) * (2**(iLevel-minLevel)) )
      me%visc( iLevel ) = cs2 * ( 1.0_rk/me%omLvl( iLevel ) - 0.5_rk )

      ! bulk omega
      me%visc_bulk( iLevel) = me%bulk_viscosityLB *           &
        &                     physics%fac( minLevel )%visc /  &
        &                     physics%fac( iLevel )%visc
      me%omLvl_bulk( iLevel ) = 2.0_rk/(9.0_rk * me%visc_bulk(iLevel) + 1.0_rk )

    end do

  end subroutine set_omegas_acoustic
! ****************************************************************************** !

! ****************************************************************************** !
  !> Set the omegas according to diffusive scaling
  subroutine set_omegas_diffusive( me, physics, minLevel, maxLevel )
    ! --------------------------------------------------------------------------
    type( mus_fluid_type ),intent(inout) :: me
    !> physics type to convert physLB unit
    type( mus_physics_type ), intent(in) :: physics
    integer, intent(in) :: minLevel
    integer, intent(in) :: maxLevel
    ! --------------------------------------------------------------------------
    integer :: iLevel
    ! --------------------------------------------------------------------------

    me%omLvl(:) = me%omega
    me%visc(:) = cs2 * ( 1.0_rk/me%omega - 0.5_rk )

    ! set LB viscosity and omega for each level
    do iLevel = minLevel, maxLevel

      me%dtLvl( iLevel ) = 1._rk / (4**(iLevel-minLevel))
      me%dxLvl( iLevel ) = 1._rk / (2**(iLevel-minLevel))

      ! @todo: How to set bulk omega by diffusive scaling?
      me%visc_bulk(iLevel) = me%bulk_viscosityLB *          &
        &                    physics%fac( minLevel )%visc / &
        &                    physics%fac( iLevel )%visc
      me%omLvl_bulk( iLevel ) = 2.0_rk/(9.0_rk * me%visc_bulk(iLevel) + 1.0_rk )

    end do

  end subroutine set_omegas_diffusive
! ****************************************************************************** !

! ****************************************************************************** !
  !> Check if omega is in the range of 1.0 - 1.99.\n
  !! check if omega is close to 1.0, which may cause shear stress to be 0.\n
  !! And cross check whether omega at each level is set correctly
  !! by computing physical viscosity from omega and check it with
  !! specified physical kinematic viscosity.\n
  !! nu = (1/omega - 0.5)*cs2
  subroutine check_omegas( me, physics, minLevel, maxLevel )
    ! ---------------------------------------------------------------------------
    type( mus_fluid_type ),intent(inout) :: me !< fluid type
    !> physics type to convert physLB unit
    type( mus_physics_type ), intent(in) :: physics
    integer, intent(in) :: minLevel
    integer, intent(in) :: maxLevel
    ! ---------------------------------------------------------------------------
    integer :: iLevel
    real(kind=rk) :: viscPhy_loc, error_rel
    ! --------------------------------------------------------------------------

    do iLevel = minLevel, maxLevel

      ! check if omega in the range of [1.0, 1.99]
      if( me%omLvl(iLevel) < 1.0_rk) then
        write(logUnit(1),"(A,I0)")' Attention: Omega on level: ', iLevel
        write(logUnit(1),"(A)")   ' is chosen to be under-relaxing (<1.0).'
        write(logUnit(1),"(A)")   ' This might cause non-physical behavior.'
      elseif( abs(me%omLvl(iLevel) - 1.0_rk ) < 1000 * eps_single_k ) then
        write(logUnit(1),"(A,I0)")' Attention: Omega on level: ', iLevel
        write(logUnit(1),"(A)")   ' is chosen close to 1.0'
        write(logUnit(1),"(A)")   ' This can cause program to crash '&
          &                     //'when it calculates strain rate or shear stress.'
      endif

      ! calculate physical viscosity from omega
      viscPhy_loc = (1.0_rk/me%omLvl(iLevel) - 0.5_rk) * cs2 &
        &           * physics%fac(iLevel)%visc
! DEBUG output
! write(logUnit(1),*) ' iLevel: ', iLevel
! write(logUnit(1),*) ' omega: ', fluid%omLvl(iLevel)
! write(logUnit(1),*) ' viscFac: ', physics%fac(iLevel)%visc
! write(logUnit(1),*) ' fluid kine viscosity: ', fluid%kine_viscosity

      ! calculate relative error
      error_rel = (me%kine_viscosity-viscPhy_loc) / me%kine_viscosity
      if( abs(real(error_rel, kind=single_k)) > eps_single_k) then
        write(logUnit(1),*) 'Error: Physical kinematic viscosity computed '//  &
          &                 'from omega'
        write(logUnit(1),*) '       does not match with specified physical '// &
          &                 'kinematic viscosity on level:', iLevel
        write(logUnit(1),*) '       Physical kinematic shear viscosity: ',     &
          &                 me%kine_viscosity
        write(logUnit(1),*) '       Calculated phy kinematic shear viscosity: '&
          &                 , viscPhy_loc
        write(logUnit(1),*) '       Defined omega: ', me%omLvl(iLevel)
        call tem_abort()
      end if
    end do ! iLevel

  end subroutine check_omegas
! ****************************************************************************** !


  ! ************************************************************************** !
  !> write fluid prop into a lua file
  !!
  subroutine mus_fluid_save2lua( me, conf )
    ! ---------------------------------------------------------------------------
    !> single fluid type
    type( mus_fluid_type ), intent(in) :: me
    type( aot_out_type ) :: conf
    ! ---------------------------------------------------------------------------

    call aot_out_open_table( put_conf = conf, tname = 'fluid' )

    call aot_out_val( put_conf = conf,    &
      &               vname    = 'omega', &
      &               val      = me%omega )
    call aot_out_val( put_conf = conf,               &
      &               vname    = 'bulk_viscosityLB', &
      &               val      = me%bulk_viscosityLB )
    call aot_out_val( put_conf = conf,               &
      &               vname    = 'kine_viscosityLB', &
      &               val      = me%kine_viscosityLB )
    call aot_out_val( put_conf = conf,    &
      &               vname    = 'omLvl', &
      &               val      = me%omLvl )

    if ( me%nNwtn%active ) call mus_nNwtn_save2lua( me%nNwtn, conf )

    call aot_out_close_table( put_conf = conf )

  end subroutine mus_fluid_save2lua
  ! **************************************************************************** !

! ****************************************************************************** !
  !> Update omega by temporal factor
  !! This routine is called every iteration to update relaxation parameters
  !! including omega and mrt.\n
  !! 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_update_relaxation( me, moment, schemeHeader, level, tNow )
    ! --------------------------------------------------------------------------
    type( mus_fluid_type ) :: me
    type(mus_moment_type), intent(in) :: moment
    type(mus_scheme_header_type), intent(in) :: schemeHeader
    integer, intent(in) :: level
    type( tem_time_type ),   intent(in) :: tNow
    ! --------------------------------------------------------------------------
    !> Relaxation parameter
    real(kind=rk) :: omega, fac, omega_bulk, omega1_2, omega1_6
    ! --------------------------------------------------------------------------

    ! calculate ramping omega based reference omega
    ! omega = me%omLvl( level )
    fac = tem_temporal_for( temporal = me%visc_ramping, &
      &                     time     = tNow )
    ! visc = cs2 * ( 1/w - 1/2 )
    omega = 1.0 / ( cs2inv * fac * me%visc( level ) + 0.5_rk )
    me%omLvl( level ) = omega

    omega_bulk = 2.0_rk / (9.0_rk * fac * me%visc_bulk(level) &
      &                     + 1.0_rk )
    me%omLvl_bulk( level ) = omega_bulk
    write(logUnit(10),"(A,I0)")    "Level: ", level
    write(logUnit(10),"(A,F10.5)") "Omega factor: ", fac
    write(logUnit(10),"(A,F10.5)") "Omega ramped: ", omega
    call mus_set_mrtRelaxation(mrt          = me%mrt(Level), &
      &                        omegaKine    = omega,         &
      &                        omegaBulk    = omega_bulk,    &
      &                        schemeHeader = schemeHeader,  &
      &                        moment       = moment         )

    ! Set scaling factor to convert nonEquilibrium moments from fine to coarse
    ! and vice versa
    
  end subroutine mus_update_relaxation
! ****************************************************************************** !

! ****************************************************************************** !
  !> 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
! ****************************************************************************** !

! ****************************************************************************** !
  subroutine mus_init_fluid( me, physics, scaling, minLevel, maxLevel )
    ! ---------------------------------------------------------------------------
    !> fluid type
    type( mus_fluid_type ), intent(inout) :: me
    !> physics type to convert physics to lattice unit or vice versa
    type( mus_physics_type ), intent(in) :: physics
    !> scaling type, used by setting omega
    character(len=labelLen), intent(in) :: scaling
    !> min and max level
    integer, intent(in) :: minLevel, maxLevel
    ! ---------------------------------------------------------------------------
    integer :: ii

    allocate( me%dtLvl( minLevel:maxLevel ) )
    allocate( me%dxLvl( minLevel:maxLevel ) )
    allocate( me%mrt  ( minLevel:maxLevel ) )
    allocate( me%omLvl( minLevel:maxLevel ) )
    allocate( me%omLvl_bulk( minLevel:maxLevel ) )
    allocate( me%visc( minLevel:maxLevel ) )
    allocate( me%visc_bulk( minLevel:maxLevel ) )
    allocate( me%wN( minLevel:maxLevel ) )

    ! assume me%omega, me%bulk_viscosityLB is set before here
    write(logUnit(1),"(A)") '  Setting omegas and viscosities for each level'
    if ( trim(scaling) == 'acoustic' ) then
      call set_omegas_acoustic( me = me, physics = physics, &
        &                       minLevel = minLevel, &
        &                       maxLevel = maxLevel )
    else
      call set_omegas_diffusive( me = me, physics = physics, &
        &                        minLevel = minLevel, &
        &                        maxLevel = maxLevel )
    end if

    do ii = minLevel, maxLevel
      me%wN(ii) = 1.0_rk / ( 0.5_rk + 2.0_rk * me%omLvl(ii) &
        &                  * me%lambda/(2.0_rk-me%omLvl(ii)) )
      write(logUnit(1),"(A, I0)")   '         level: ', ii
      write(logUnit(3),"(A, F6.4)") '          dxLB: ', me%dxLvl( ii )
      write(logUnit(3),"(A, F6.4)") '          dtLB: ', me%dtLvl( ii )
      write(logUnit(1),"(A, F6.4)") '         omega: ', me%omLvl( ii )
      write(logUnit(1),"(A, F6.4)") '    bulk omega: ', me%omLvl_bulk( ii )
      write(logUnit(1),"(A, F6.4)") '      omega(-): ', me%wN( ii )
    end do


    ! Check if omega is larger than 1.99 or smaller than 1.0
    ! Check if omega is close to 1.0, which may cause shear stress to be 0
    ! double check viscosity
    call check_omegas( me       = me,       &
      &                physics  = physics,  &
      &                minLevel = minLevel, &
      &                maxLevel = maxLevel  )
  end subroutine mus_init_fluid
! ****************************************************************************** !

! ****************************************************************************** !
  !> This routine initilizes fluid visocity relaxation paramters
  subroutine mus_init_fluid_new(me, physics, schemeHeader, minLevel, maxLevel, &
    &                           nSolve, levelDesc, tNow )
    ! ---------------------------------------------------------------------------
    !> fluid type
    type(mus_fluid_type), intent(inout) :: me
    !> physics type to convert physics to lattice unit or vice versa
    type(mus_physics_type), intent(in) :: physics
    !> scheme header
    type(mus_scheme_header_type), intent(in) :: schemeHeader
    !> min and max level
    integer, intent(in) :: minLevel, maxLevel
    !> level descriptor
    type(tem_levelDesc_type), intent(in) :: levelDesc(minLevel:maxLevel)
    !> number of elements to solve per level (fluid+ghost)
    integer, intent(in) :: nSolve(minLevel:maxLevel)
    !> current simulation time
    type(tem_time_type),intent(in) :: tNow
    ! ---------------------------------------------------------------------------
    integer :: iLevel
    ! ---------------------------------------------------------------------------
    ! allocate array to store kinematic viscosity
    allocate(me%viscKine%dataOnLvl(minLevel:maxLevel))
    do iLevel = minLevel, maxLevel
      call init(me%viscKine%dataOnLvl(iLevel), nSolve(iLevel))
      me%viscKine%dataOnLvl(iLevel)%val = -100.0_rk
    end do

    ! New initialization
    ! allocate relaxation parameter array to size nSolve 
    call mus_init_relaxParam(omLvl    = me%viscKine%omLvl, &
      &                      minLevel = minLevel,          &
      &                      maxLevel = maxLevel,          &
      &                      nSolve   = nSolve             )

    ! assign function pointer
    me%mrtPtr => mus_assign_mrt_ptr(schemeHeader)

    !call mus_assign_mrt_ptr(schemeHeader, me%mrtPtr)
    me%nonEqScalingFacs => mus_assign_intp_nonEqScalingFacs_ptr(schemeHeader)

    ! intialize kinematic omega
    do iLevel = minLevel, maxLevel
      call mus_update_relaxParamFromViscSTfun(             &
        & omega       = me%viscKine%omLvl(iLevel)%val,     &
        & visc        = me%viscKine%dataOnLvl(iLevel)%val, & 
        & viscSTfun   = me%viscKine%STfun,                 &
        & viscRef     = physics%fac(iLevel)%visc,          &
        & nSolve      = nSolve(iLevel),                    &
        & baryOfTotal = levelDesc(iLevel)%baryOfTotal,     &
        & dtL         = me%dtLvl(iLevel),                  &
        & tNow        = tNow                               )
    end do  

    ! initialize bulk omega
!KM!    if (me%requireOmegaBulk) then
!KM!      call mus_init_relaxParam(omLvl    = me%viscBulk%omLvl, &
!KM!        &                      minLevel = minLevel,          &
!KM!        &                      maxLevel = maxLevel,          &
!KM!        &                      nSolve   = nSolve             )
!KM!      do iLevel = minLevel, maxLevel
!KM!        call mus_update_relaxParamFromViscSTfun(         &
!KM!          & omega       = me%viscBulk%omLvl(iLevel)%val, &
!KM!          & viscSTfun   = me%viscBulk%STfun,             &
!KM!          & viscRef     = physics%fac(iLevel)%visc,      &
!KM!          & nSolve      = nSolve(iLevel),                &
!KM!          & baryOfTotal = levelDesc(iLevel)%baryOfTotal, &
!KM!          & dtL         = me%dtLvl(iLevel),              &
!KM!          & tNow        = tNow                           )
!KM!      end do  
!KM!    end if  

  end subroutine mus_init_fluid_new
! ****************************************************************************** !

end module mus_fluid_module
! ****************************************************************************** !
