! See copyright notice in the COPYRIGHT file.
! ****************************************************************************** !
!> This module assigns function pointer to calculate turbulent viscosity 
!! according to turbulence model and scheme definition
module mus_turb_viscosity_module
  ! include treelm modules 
  use env_module,          only: rk
  use tem_aux_module,      only: tem_abort
 
  use mus_scheme_header_module, only: mus_scheme_header_type  
  use mus_turbulence_module,    only: mus_turbulence_type, &
    &                                 mus_turbulence_config_type
  use mus_scheme_layout_module,  only: mus_scheme_layout_type
  use mus_WALE_module,          only: mus_turbVisc_WALE_2D, &
    &                                 mus_turbVisc_WALE_3D
  use mus_Vreman_module,        only: mus_turbVisc_Vreman_2D, &
    &                                 mus_turbVisc_Vreman_3D
  use mus_Smagorinsky_module,   only: mus_turbVisc_Smagorinsky_fromGradU3D, &
    & mus_turbVisc_Smagorinsky_fromGradU2D,         &
    & mus_turbVisc_Smagorinsky_fromGradU3D_incomp,  &
    & mus_turbVisc_Smagorinsky_fromGradU2D_incomp,  &
    & mus_turbVisc_Smagorinsky_fromPreColPDF,        &
    & mus_turbVisc_Smagorinsky_fromPreColPDF_incomp
  
  implicit none
  private

  public :: mus_assign_turbVisc_ptr
  
contains
  
  ! ************************************************************************** !
  !> This routine assigns function pointer to compute turbulence viscosity
  !! based on turbulence model and scheme header definition
  subroutine mus_assign_turbVisc_ptr(schemeHeader, turb ) 
    ! --------------------------------------------------------------------------
    !> Scheme header information
    type(mus_scheme_header_type), intent(in) :: schemeHeader
    
    !> turbulence type
    type(mus_turbulence_type), intent(inout) :: turb
    ! --------------------------------------------------------------------------
    turb%viscFunc%calcFromGradU => mus_turbVisc_fromGradU_dummy
    turb%viscFunc%calcFromPreColPDF => mus_turbVisc_fromPreColPDF_dummy

    select case(trim(schemeHeader%layout))
    case ('d2q9')
      select case(trim(turb%config%model))
      case ('smagorinsky')
        select case (trim(schemeHeader%kind))
        case ('lbm', 'lbm_nNwtn')
          if (turb%config%compSR_fromPDF) then
            turb%viscFunc%calcFromPreColPDF &
              & => mus_turbVisc_Smagorinsky_fromPreColPDF
          else
            turb%viscFunc%calcFromGradU => mus_turbVisc_Smagorinsky_fromGradU2D 
          end if
        case ('lbm_incomp', 'lbm_incomp_nNwtn')
          if (turb%config%compSR_fromPDF) then
            turb%viscFunc%calcFromPreColPDF                        &
              & => mus_turbVisc_Smagorinsky_fromPreColPDF_incomp
          else
            turb%viscFunc%calcFromGradU                        &
              & => mus_turbVisc_Smagorinsky_fromGradU2D_incomp
          end if
        case default
          call tem_abort('Error: Unknown scheme kind for turbulence '&
            &          //'viscosity ptr')
        end select

      case ('wale')
        turb%viscFunc%calcFromGradU => mus_turbVisc_WALE_2D
      case ('vreman')
        turb%viscFunc%calcFromGradU => mus_turbVisc_Vreman_2D 
      end select

    case ('d3q15','d3q19','d3q27')
      select case(trim(turb%config%model))
      case ('smagorinsky')
        select case (trim(schemeHeader%kind))
        case ('lbm', 'lbm_nNwtn')
          if (turb%config%compSR_fromPDF) then
            turb%viscFunc%calcFromPreColPDF &
              & => mus_turbVisc_Smagorinsky_fromPreColPDF
          else
            turb%viscFunc%calcFromGradU => mus_turbVisc_Smagorinsky_fromGradU3D 
          end if
        case ('lbm_incomp', 'lbm_incomp_nNwtn')
          if (turb%config%compSR_fromPDF) then
            turb%viscFunc%calcFromPreColPDF                        &
              & => mus_turbVisc_Smagorinsky_fromPreColPDF_incomp
          else
            turb%viscFunc%calcFromGradU                        &
              & => mus_turbVisc_Smagorinsky_fromGradU3D_incomp
          end if
        case default
          call tem_abort('Error: Unknown scheme kind for turbulence '&
            &          //'viscosity ptr')
        end select

      case ('wale')
        turb%viscFunc%calcFromGradU => mus_turbVisc_WALE_3D
      case ('vreman')
        turb%viscFunc%calcFromGradU => mus_turbVisc_Vreman_3D 
      end select
    case default
      call tem_abort('Error: Unknown layout for turbulence viscosity ptr')
    end select
  end subroutine mus_assign_turbVisc_ptr
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> Dummy function for turbulent viscosity from Gradu procedure
  pure function mus_turbVisc_fromGradU_dummy(turbConfig, gradU, dxL, dtL) result(res)
    ! --------------------------------------------------------------------------
    !> Contains turbulenct coefficients
    type(mus_turbulence_config_type), intent(in) :: turbConfig
    !> gradient of velocity field 
    real(kind=rk), intent(in) :: gradU(:,:)
    !> current level lattice element size
    real(kind=rk), intent(in) :: dxL
    !> current level lattice time step size
    real(kind=rk), intent(in) :: dtL
    !> output: eddy viscosity of an element
    real(kind=rk) :: res
    ! --------------------------------------------------------------------------
    !call tem_abort('DUMMY function to calculate turbVisc from gradU')
  end function mus_turbVisc_fromGradU_dummy
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> Dummy function to compute turbulent viscosity from PDF
  pure function mus_turbVisc_fromPreColPDF_dummy(turbConfig, f_preCol, layout, dxL, &
    & dtL, viscKine) result(res)
    ! --------------------------------------------------------------------------
    !> Contains turbulenct coefficients
    type(mus_turbulence_config_type), intent(in) :: turbConfig
    !> pre-collision pdf
    real(kind=rk), intent(in) :: f_preCol(:)
    !> scheme layout
    type(mus_scheme_layout_type), intent(in) :: layout
    !> current level element size
    real(kind=rk), intent(in) :: dxL
    !> current level lattice time step size
    real(kind=rk), intent(in) :: dtL
    !> Background kinematic viscosity
    real(kind=rk), intent(in) :: viscKine
    !> output: eddy viscosity of an element
    real(kind=rk) :: res
    ! --------------------------------------------------------------------------
    !call tem_abort('DUMMY function to calculate turbVisc from PreCol PDF')
  end function mus_turbVisc_fromPreColPDF_dummy
  ! ************************************************************************** !

end module mus_turb_viscosity_module
