! See copyright notice in the COPYRIGHT file.
! ****************************************************************************** !
!> This module contains function to compute eddy viscosity for
!! Smagorinsky les turbulence model. 
!! author: Kannan Masilamani
module mus_Smagorinsky_module
  ! treelm modules
  use env_module,           only: rk, labelLen
  use tem_param_module,     only: div1_2, div1_3, div2_3, cs2, rho0
  use tem_tools_module,     only: tem_horizontalSpacer
  use tem_aux_module,       only: tem_abort
  use tem_logging_module,   only: logUnit
  
  ! musubi modules
  use mus_turbulence_module, only: mus_turbulence_config_type
  use mus_scheme_layout_module, only: mus_scheme_layout_type
  use mus_derivedQuantities_module2, only: secondMom, getEquilibrium, &
    &                                      getEquilibriumIncomp
 
  implicit none
  private

  public :: mus_turbVisc_Smagorinsky_fromGradU3D
  public :: mus_turbVisc_Smagorinsky_fromGradU2D
  public :: mus_turbVisc_Smagorinsky_fromGradU3D_incomp
  public :: mus_turbVisc_Smagorinsky_fromGradU2D_incomp
  public :: mus_turbVisc_Smagorinsky_fromPreColPDF
  public :: mus_turbVisc_Smagorinsky_fromPreColPDF_incomp

contains 

  ! ************************************************************************** !
  !> Calculate eddy viscosity with smagorinsky model for compressible model
  !! using gradient of velocity 
  !! The formula is taken from https://caefn.com/openfoam/smagorinsky-sgs-model
  !! nu_t = C_k delta sqrt(k_sgs)
  !! k_sgs = ((-b+sqrt(b^2+4ac))/2a)^2
  !! a = C_e/delta, b=2/3 tr(dev(Strain)), c = 2 C_k delta (dev(Strain):Strain)
  pure function mus_turbVisc_Smagorinsky_fromGradU3D(turbConfig, gradU, dxL, &
    & dtL) result(res)
    ! --------------------------------------------------------------------------
    !> Contains turbulenct coefficients
    type(mus_turbulence_config_type), intent(in) :: turbConfig
    !> gradient of velocity
    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: turbulent viscosity
    real(kind=rk) :: res
    ! --------------------------------------------------------------------------
    real(kind=rk) :: SR(6), devSR(3), oneThird_trSR, devSR_SR, tr_devSR
    real(kind=rk) :: sqrt_k_sgs
    real(kind=rk) :: a, b, c
    ! --------------------------------------------------------------------------
    ! symmetric strain rate tensors
    SR(1) = gradU(1,1)                      !S_XX
    SR(2) = gradU(2,2)                      !S_YY
    SR(3) = gradU(3,3)                      !S_ZZ
    SR(4) = (gradU(1,2)+gradU(2,1))*0.5_rk  !S_XY
    SR(5) = (gradU(2,3)+gradU(3,2))*0.5_rk  !S_YZ
    SR(6) = (gradU(1,3)+gradU(3,1))*0.5_rk  !S_XZ

    ! onethird of trace of strain rate
    oneThird_trSR = (SR(1) + SR(2) + SR(3))*div1_3
    ! deviatoric strain rate differs from strain rate only in diagonal terms
    devSR(1) = SR(1) - oneThird_trSR
    devSR(2) = SR(2) - oneThird_trSR
    devSR(3) = SR(3) - oneThird_trSR

    ! inner product of devSR:SR
    devSR_SR = SR(1)*devSR(1) + SR(2)*devSR(2) + SR(3)*devSR(3) &
      &      + 2.0_rk*( SR(4)**2 + SR(5)**2 + SR(6)**2)

    ! trace of deviatoric strain rate
    tr_devSR = devSR(1) + devSR(2) + devSR(3)

    ! parameters to compute subgrid scale kinetic energy
    a = turbConfig%coeff%C_e / dxL
    b = tr_devSR * div2_3
    c = 2.0_rk * turbConfig%coeff%C_k * dxL * devSR_SR

    ! subgrid scale kinetic energy
    sqrt_k_sgs = (-b + sqrt(b**2 + 4.0_rk*a*c))/(2.0_rk*a)

    ! subgrid scale turbulent viscosity normalized to current level
    res = (turbConfig%coeff%C_k * dxL * sqrt_k_sgs) / dtL

  end function mus_turbVisc_Smagorinsky_fromGradU3D
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> Calculate eddy viscosity with smagorinsky model for compressible model
  !! using gradient of velocity for 2D layout 
  pure function mus_turbVisc_Smagorinsky_fromGradU2D(turbConfig, gradU, dxL, &
    & dtL) result(res)
    ! --------------------------------------------------------------------------
    !> Contains turbulenct coefficients
    type(mus_turbulence_config_type), intent(in) :: turbConfig
    !> gradient of velocity
    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: turbulent viscosity
    real(kind=rk) :: res
    ! --------------------------------------------------------------------------
    real(kind=rk) :: SR(3), devSR(2), oneHalf_trSR, devSR_SR, tr_devSR
    real(kind=rk) :: sqrt_k_sgs
    real(kind=rk) :: a, b, c
    ! --------------------------------------------------------------------------
    ! symmetric strain rate tensors
    SR(1) = gradU(1,1)                      !S_XX
    SR(2) = gradU(2,2)                      !S_YY
    SR(3) = (gradU(1,2)+gradU(2,1))*0.5_rk  !S_XY

    ! one half of trace of strain rate
    oneHalf_trSR = (SR(1) + SR(2))*div1_2
    ! deviatoric strain rate differs from strain rate only in diagonal terms
    devSR(1) = SR(1) - oneHalf_trSR
    devSR(2) = SR(2) - oneHalf_trSR

    ! inner product of devSR:SR
    devSR_SR = SR(1)*devSR(1) + SR(2)*devSR(2) &
      &      + 2.0_rk*( SR(3)**2 )

    ! trace of deviatoric strain rate
    tr_devSR = devSR(1) + devSR(2)

    ! parameters to compute subgrid scale kinetic energy
    a = turbConfig%coeff%C_e / dxL
    b = tr_devSR 
    c = 2.0_rk * turbConfig%coeff%C_k * dxL * devSR_SR

    ! subgrid scale kinetic energy
    sqrt_k_sgs = (-b + sqrt(b**2 + 4.0_rk*a*c))/(2.0_rk*a)

    ! subgrid scale turbulent viscosity
    res = turbConfig%coeff%C_k * dxL * sqrt_k_sgs / dtL

  end function mus_turbVisc_Smagorinsky_fromGradU2D
  ! ************************************************************************** !


  ! ************************************************************************** !
  !> Calculate eddy viscosity with smagorinsky model for incompressible model
  !! using gradient of velocity 
  pure function mus_turbVisc_Smagorinsky_fromGradU3D_incomp(turbConfig, gradU, &
    & dxL, dtL) result(res)
    ! --------------------------------------------------------------------------
    !> Contains turbulenct coefficients
    type(mus_turbulence_config_type), intent(in) :: turbConfig
    !> gradient of velocity
    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: turbulent viscosity
    real(kind=rk) :: res
    ! --------------------------------------------------------------------------
    real(kind=rk) :: SR(6), SR_mag
    real(kind=rk) :: k_sgs
    ! --------------------------------------------------------------------------
    ! symmetric strain rate tensors
    SR(1) = gradU(1,1)                      !S_XX
    SR(2) = gradU(2,2)                      !S_YY
    SR(3) = gradU(3,3)                      !S_ZZ
    SR(4) = (gradU(1,2)+gradU(2,1))*0.5_rk  !S_XY
    SR(5) = (gradU(2,3)+gradU(3,2))*0.5_rk  !S_YZ
    SR(6) = (gradU(1,3)+gradU(3,1))*0.5_rk  !S_XZ

    ! magnitude of strain rate tensor, sqrt(2 S:S)
    SR_mag = sqrt(2.0_rk*( SR(1)**2 + SR(2)**2 + SR(3)**2 &
      &                  + 2.0_rk*(SR(4)**2 + SR(5)**2 + SR(6)**2)))

    ! subgrid scale kinetic energy of incompressible model
    k_sgs = turbConfig%coeff%C_k * dxL**2 * SR_mag**2 / turbConfig%coeff%C_e

    ! subgrid scale turbulent viscosity
    res = turbConfig%coeff%C_k * dxL * sqrt(k_sgs) / dtL

  end function mus_turbVisc_Smagorinsky_fromGradU3D_incomp
  ! ************************************************************************** !

  ! ************************************************************************** !
  !> Calculate eddy viscosity with smagorinsky model for incompressible model
  !! using gradient of velocity for 2D layout 
  pure function mus_turbVisc_Smagorinsky_fromGradU2D_incomp(turbConfig, gradU, &
    & dxL, dtL) result(res)
    ! --------------------------------------------------------------------------
    !> Contains turbulenct coefficients
    type(mus_turbulence_config_type), intent(in) :: turbConfig
    !> gradient of velocity
    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: turbulent viscosity
    real(kind=rk) :: res
    ! --------------------------------------------------------------------------
    real(kind=rk) :: SR(3), SR_mag
    real(kind=rk) :: k_sgs
    ! --------------------------------------------------------------------------
    ! symmetric strain rate tensors
    SR(1) = gradU(1,1)                      !S_XX
    SR(2) = gradU(2,2)                      !S_YY
    SR(3) = (gradU(1,2)+gradU(2,1))*0.5_rk  !S_XY

    ! magnitude of strain rate tensor, sqrt(2 S:S)
    SR_mag = sqrt(2.0_rk*( SR(1)**2 + SR(2)**2 + 2.0_rk*(SR(3)**2) ))

    ! subgrid scale kinetic energy of incompressible model
    k_sgs = turbConfig%coeff%C_k * dxL**2 * SR_mag**2 / turbConfig%coeff%C_e

    ! subgrid scale turbulent viscosity
    res = turbConfig%coeff%C_k * dxL * sqrt(k_sgs) / dtL

  end function mus_turbVisc_Smagorinsky_fromGradU2D_incomp
  ! ************************************************************************** !


  ! ************************************************************************** !
  !> Calculate eddy viscosity with smagorinsky model for compressible model
  !! using pre-collision PDF 
  pure function mus_turbVisc_Smagorinsky_fromPreColPDF(turbConfig, f_preCol, &
    & layout, dxL, dtL, viscKine) result(res)
    ! --------------------------------------------------------------------------
    !> Contains turbulenct coefficients
    type(mus_turbulence_config_type), intent(in) :: turbConfig
    !> precollision PDF
    real(kind=rk), intent(in) :: f_preCol(:)
    !> scheme layout
    type(mus_scheme_layout_type), intent(in) :: layout
    !> current level lattice element size
    real(kind=rk), intent(in) :: dxL
    !> current level lattice time step size
    real(kind=rk), intent(in) :: dtL
    !> Background kinematic viscosity divided by dtL
    real(kind=rk), intent(in) :: viscKine
    !> output: turbulent viscosity
    real(kind=rk) :: res
    ! --------------------------------------------------------------------------
    real(kind=rk) :: visc_coeff, rho, inv_rho, vel(3) 
    real(kind=rk) :: fEq(layout%fStencil%QQ), nEq(layout%fStencil%QQ)
    real(kind=rk) :: nEqTens(6), nEqTensMag, viscKineTerm, nEqTensTerm 
    ! --------------------------------------------------------------------------
    ! viscosity coeff
    visc_coeff = (turbConfig%coeff%C_s * dxL )**2.0_rk

    ! density
    rho = sum(f_preCol)
    inv_rho = 1.0_rk/rho
    ! velocity
    vel(1) = sum(f_preCol * layout%fStencil%cxDirRK(1,:)) * inv_rho
    vel(2) = sum(f_preCol * layout%fStencil%cxDirRK(2,:)) * inv_rho
    vel(3) = sum(f_preCol * layout%fStencil%cxDirRK(3,:)) * inv_rho

    ! Calculate the equilibrium distribution function
    fEq(:) = getEquilibrium( rho, vel, layout)

    ! Calculate the non-equilibrium part
    nEq(:) = f_preCol(:) - fEq(:)

    ! Now calculate the symmetric deviatoric second-order tensor of 
    ! nonEquilibrium part
    ! the static part cs2 I is usually neglected for weakly compressible flows
    ! however, in current implementation it is considered
    nEqTens(1) = sum( (layout%fStencil%cxcx(1,:) - cs2) * nEq)  !XX 
    nEqTens(2) = sum( (layout%fStencil%cxcx(2,:) - cs2) * nEq)  !YY
    nEqTens(3) = sum( (layout%fStencil%cxcx(3,:) - cs2) * nEq)  !ZZ
    !nEqTens(1) = sum( (layout%fStencil%cxcx(1,:) ) * nEq)  !XX 
    !nEqTens(2) = sum( (layout%fStencil%cxcx(2,:) ) * nEq)  !YY
    !nEqTens(3) = sum( (layout%fStencil%cxcx(3,:) ) * nEq)  !ZZ
    nEqTens(4) = sum( (layout%fStencil%cxcx(4,:) ) * nEq)  !XY
    nEqTens(5) = sum( (layout%fStencil%cxcx(5,:) ) * nEq)  !YZ
    nEqTens(6) = sum( (layout%fStencil%cxcx(6,:) ) * nEq)  !XZ

    ! magnitude of second-order tensor
    nEqTensMag = sqrt(2.0_rk*(nEqTens(1)**2 + nEqTens(2)**2 + nEqTens(3)**2 &
      &        + 2.0_rk*(nEqTens(4)**2 + nEqTens(5)**2 + nEqTens(6)**2) ) )

    ! turbulent viscosity
    !nu_t = (sqrt((v+cs2 dt/2)^2+2(Cs dx)^2|nEQ|/rho) - (v+cs2 dt/2))/2
    ! viscKine is scaled to current level dt so multiply viscKine with dt
    ! to get unit consistent
    viscKineTerm = (viscKine + div1_2 * cs2) * dtL
    nEqTensTerm = 2.0_rk * visc_coeff * nEqTensMag * inv_rho
    res =  (sqrt(viscKineTerm**2 + nEqTensTerm) - viscKineTerm)*div1_2 / dtL

  end function mus_turbVisc_Smagorinsky_fromPreColPDF
  ! ************************************************************************** !


  ! ************************************************************************** !
  !> Calculate eddy viscosity with smagorinsky model for incompressible model
  !! using pre-collision PDF 
  pure function mus_turbVisc_Smagorinsky_fromPreColPDF_incomp(turbConfig, &
    & f_preCol, layout, dxL, dtL, viscKine) result(res)
    ! --------------------------------------------------------------------------
    !> Contains turbulenct coefficients
    type(mus_turbulence_config_type), intent(in) :: turbConfig
    !> precollision PDF
    real(kind=rk), intent(in) :: f_preCol(:)
    !> scheme layout
    type(mus_scheme_layout_type), intent(in) :: layout
    !> current level lattice element size
    real(kind=rk), intent(in) :: dxL
    !> current level lattice time step size
    real(kind=rk), intent(in) :: dtL
    !> Background kinematic viscosity divided by dtL
    real(kind=rk), intent(in) :: viscKine
    !> output: turbulent viscosity
    real(kind=rk) :: res
    ! --------------------------------------------------------------------------
    real(kind=rk) :: visc_coeff, rho, inv_rho, vel(3) 
    real(kind=rk) :: fEq(layout%fStencil%QQ), nEq(layout%fStencil%QQ)
    real(kind=rk) :: nEqTens(6), nEqTensMag, viscKineTerm, nEqTensTerm 
    ! --------------------------------------------------------------------------
    ! viscosity coeff
    visc_coeff = (turbConfig%coeff%C_s * dxL )**2.0_rk

    ! density
    rho = sum(f_preCol)

    ! use reference density for incompressible model
    inv_rho = 1.0_rk/rho0
    ! velocity
    vel(1) = sum(f_preCol * layout%fStencil%cxDirRK(1,:)) * inv_rho
    vel(2) = sum(f_preCol * layout%fStencil%cxDirRK(2,:)) * inv_rho
    vel(3) = sum(f_preCol * layout%fStencil%cxDirRK(3,:)) * inv_rho

    ! Calculate the equilibrium distribution function
    fEq(:) = getEquilibriumIncomp( rho, vel, layout, rho0)

    ! Calculate the non-equilibrium part
    nEq(:) = f_preCol(:) - fEq(:)

    ! Now calculate the symmetric second-order tensor of nonEquilibrium part
    nEqTens = secondMom(layout%fStencil%cxcx, nEq, layout%fStencil%QQ)

    ! magnitude of second-order tensor
    nEqTensMag = sqrt(2.0_rk*(nEqTens(1)**2 + nEqTens(2)**2 + nEqTens(3)**2 &
      &          + 2.0_rk*(nEqTens(4)**2 + nEqTens(5)**2 + nEqTens(6)**2) ) )

    ! turbulent viscosity
    !nu_t = (sqrt((v+cs2 dt/2)^2+2(Cs dx)^2|nEQ|/rho) - (v+cs2 dt/2))/2
    ! viscKine is scaled to current level dt so multiply viscKine with dt
    ! to get unit consitent
    viscKineTerm = (viscKine + div1_2 * cs2) * dtL
    nEqTensTerm = 2.0_rk * visc_coeff * nEqTensMag * inv_rho
    res =  (sqrt(viscKineTerm**2 + nEqTensTerm) - viscKineTerm)*div1_2 / dtL

  end function mus_turbVisc_Smagorinsky_fromPreColPDF_incomp
  ! ************************************************************************** !

end module mus_Smagorinsky_module
