! See copyright notice in the COPYRIGHT file.
! ****************************************************************************** !
!> author: Manuel Hasert
!! Auxiliary functionality
!!
module mus_aux_module

  ! include musubi modules
  use mus_param_module,              only: mus_param_type
  use mus_geom_module,               only: mus_geom_type
  use mus_tools_module,              only: perform_checks,                     &
    &                                      check_streaming_layout!,             &
    !&                                      mus_writeSolverSpecInfo
  use mus_interpolate_module,        only: mus_init_interpolate
  use mus_source_module,             only: mus_init_sourceTerms
  use mus_transport_var_module,      only: mus_init_transport_var
  use mus_scheme_type_module,        only: mus_scheme_type
  use mus_geomIncr_module,           only: mus_geomIncr, mus_init_geomIncr
  use mus_interpolate_tools_module,  only: debug_dependencies, dump_intpLists
  use mus_time_module,               only: mus_time_homogenize,                &
    &                                      mus_timeControl_homogenize
  use mus_fluid_module,              only: mus_update_relaxation, &
    &                                      mus_init_fluid_new
  use mus_mrtRelaxation_module,      only: mus_alloc_mrt   
  use mus_relaxationParam_module,    only: mus_update_relaxParamKine,       &
    &                                      mus_update_viscKine,             &
    &                                      mus_update_relaxParamFromViscSTfun
  use mus_field_module,              only: setParameters_multispecies
  use mus_tracking_module,           only: mus_init_tracker
  use mus_restart_module,            only: mus_writeRestart
  use mus_timer_module,              only: mus_timerHandles
  use mus_physics_module,            only: mus_physics_type

  ! include treelm modules
  use env_module,               only: rk, PathLen, pathSep, long_k! , newUnit
  use treelmesh_module,         only: treelmesh_type
  use tem_aux_module,           only: tem_print_execInfo, utc_date_string, &
    &                                 tem_abort
  use tem_debug_module,         only: dbgUnit
  use tem_restart_module,       only: tem_init_restart
  use tem_tools_module,         only: tem_horizontalSpacer
  ! use tem_balance_module,       only: tem_balance_initTracks
  use tem_tracking_module,      only: tem_tracker
  use tem_convergence_module,   only: tem_convergence_check,                   &
    &                                 tem_init_convergence
  use tem_timeControl_module,   only: tem_timeControl_check,                   &
    &                                 tem_timeControl_update
  use tem_simControl_module,    only: tem_simControl_syncUpdate
  use tem_time_module,          only: tem_time_dump, tem_time_type
  use tem_debug_module,         only: main_debug
  use tem_solveHead_module,     only: tem_solveHead_type
  use tem_logging_module,       only: logUnit
  use tem_global_module,        only: tem_global_type
  use tem_depend_module,        only: tem_init_depend
  use tem_spacetime_fun_module, only: tem_create_subTree_of_st_funList
  use tem_dyn_array_module,     only: dyn_intArray_type
  use tem_timer_module,         only: tem_getTimerVal
  use tem_general_module,       only: tem_general_type
  use tem_operation_var_module, only: tem_opVar_reduction_transient_update


  implicit none
  private

  public :: check_flow_status
  public :: mus_init_aux
  public :: mus_update_relaxParams
  public :: mus_banner
  public :: mus_dumpData

contains

! ****************************************************************************** !
  !> This routine performs several tasks: geometry increment, time updating,
  !! tracking, density checking, restart
  subroutine check_flow_status( scheme, geometry, general, restart_triggered )
    ! --------------------------------------------------------------------------
    !> containers for the different schemes
    type(mus_scheme_type), target, intent(inout) :: scheme
    !> geometry infomation
    type(mus_geom_type),intent(in) :: geometry
    !> Global parameters
    type( tem_general_type ), intent(inout)  :: general
    !>
    logical, intent(inout) :: restart_triggered

    ! --------------------------------------------------------------------------

    ! Call the geometry increment module which  performs solidification
    ! or liquification based on certain criteria read from the Lua file
    if( geometry%dynamicGeom ) then
      call mus_geomIncr(geometry = geometry, scheme = scheme,                  &
        &               commPattern = general%commPattern,                     &
        &               general = general                                      )
    end if ! Time for geom incr

    ! perform density check for all schemes
    ! Perform run-time checks for total density, NaNs
    call perform_checks( scheme, geometry%tree%global%minLevel, &
      &                  geometry%tree%global%maxLevel, general)

    ! check for convergence only if abortCriteria is set for steady_state
    if( general%simControl%abortCriteria%steady_state) then
      ! check for steady state convergence based on convergence criteria
      ! defined in convergence table
      call tem_convergence_check(                                   &
        &                 me     = general%simControl%abortCriteria &
        &                                 %convergence,             &
        &                 time   = general%simControl%now,          &
        &                 status = general%simControl%status,       &
        &                 varSys = scheme%varSys,                   &
        &                 tree   = geometry%tree                    )
    end if

    ! Update time reduction variable
    if (scheme%redTransVarMap%varPos%nVals > 0) then
      call tem_opVar_reduction_transient_update(                            &
        &      redTransVarPos = scheme%redTransVarMap%varPos                &
        &                       %val(1:scheme%redTransVarMap%varPos%nVals), &
        &      varSys         = scheme%varSys,                              &
        &      tree           = geometry%tree,                              &
        &      time           = general%simControl%now                      )
    end if

    ! in musubi, advance time step separtely in control_routine so no need
    ! to pass dt to syncUpdate
    ! This sync update: check for stop file and time
    ! control interval trigger, communicate status bits and update timeControl
    call tem_simControl_syncUpdate( me   = general%simControl, &
      &                             proc = general%proc        )

    ! Dump tracking and restart if they are active
    call mus_dumpData( scheme   = scheme,                   &
      &                tree     = geometry%tree,            &
      &      restart_triggered  = restart_triggered,        &
      &                general  = general,                  &
      &            levelPointer = geometry%levelPointer     )

  end subroutine check_flow_status
! ****************************************************************************** !



! ****************************************************************************** !
  !> Initialize musubi solverHead and print musubi banner to screen
  subroutine mus_banner( solver )
    ! ---------------------------------------------------------------------------
    !> solver definition
    type(tem_solveHead_type), intent(in) :: solver
    ! ---------------------------------------------------------------------------
    character(len=26) :: dat_string
    ! ---------------------------------------------------------------------------
    write(logUnit(1),"(A)")''
    write(logUnit(1),"(A)")" .___  ___.  __    __       _______. __    __  .______    __  "
    write(logUnit(1),"(A)")" |   \/   | |  |  |  |     /       ||  |  |  | |   _  \  |  | "
    write(logUnit(1),"(A)")" |  \  /  | |  |  |  |    |   (----`|  |  |  | |  |_)  | |  | "
    write(logUnit(1),"(A)")" |  |\/|  | |  |  |  |     \   \    |  |  |  | |   _  <  |  | "
    write(logUnit(1),"(A)")" |  |  |  | |  `--'  | .----)   |   |  `--'  | |  |_)  | |  | "
    write(logUnit(1),"(A)")" |__|  |__|  \______/  |_______/     \______/  |______/  |_"  &
      &              //trim(solver%version)
    write(logUnit(1),"(A)")''
    write(logUnit(1),"(A)")" (C) 2012 German Research School for Simulation Sciences"
    write(logUnit(1),"(A)")" "
    call tem_print_execInfo()
    write(logUnit(1),"(A)")''
    dat_string = utc_date_string()
    write(logUnit(1),"(A)")"Run at: "//trim(dat_string)
    write(logUnit(1),"(A)")''

  end subroutine mus_banner
! ****************************************************************************** !


! ****************************************************************************** !
  !> Init auxiliary features such as interpolation boundaries, restart and 
  !! the tracker
  subroutine mus_init_aux( scheme, geometry, params)
    ! ---------------------------------------------------------------------------
    !> containers for the different schemes
    type(mus_scheme_type), target, intent(inout) :: scheme
    !> geometry information
    type(mus_geom_type), intent(inout)           :: geometry
    !> global parameters
    type(mus_param_type),intent(inout)           :: params
    ! ---------------------------------------------------------------------------
    integer :: iField, iLevel, iConv
    integer :: minLevel, maxLevel
    ! ---------------------------------------------------------------------------
    minLevel = geometry%tree%global%minLevel
    maxLevel = geometry%tree%global%maxLevel

    ! When the restart is read from separate restart table, we start the 
    ! simulation from the time step given in restart file. In the case,
    ! when restart is read from initial condition table, the simulation start
    ! time step is taken as the one defined in configuration file
    call mus_time_homogenize( me = params%general%simControl%now,              &
      &                       dt = params%physics%dtLvl( maxLevel ),           &
      &              readRestart = params%general%restart%controller           &
      &                             %readRestart )
    call mus_timeControl_homogenize(                                          &
      &                       me     = params%general%simControl%timeControl, &
      &                       dt     = params%physics%dtLvl( maxLevel ),      &
      &                       reqInt = params%reqInterval                     )

    ! De-activate the check density for time step 1
    ! \todo KM: is this call necessary with new timeControl
    call tem_timeControl_update( me   = params%general%simControl%timeControl, &
      &                          now  = params%general%simControl%now,         &
      &                          comm = params%general%proc%comm )

    !> initialize fluid type which contains relaxation parameter
    !! and function pointers to get mrt paramter and nonEqScaling factor
    !! for interpolation
    select case( trim(scheme%header%kind) )
    case('lbm', 'lbm_incomp', 'lbm_nNwtn', 'lbm_incomp_nNwtn', &
      &  'isotherm_acEq')
      if (scheme%nFields > 1) then
        call tem_abort('chosen scheme kind supports only one field')
      end if
      ! initialize fluid viscosity relaxation paramters 
      call mus_init_fluid_new(                            &
        & me           = scheme%field(1)%fieldProp%fluid, &
        & physics      = params%physics,                  &
        & schemeHeader = scheme%header,                   &
        & minLevel     = minLevel,                        &
        & maxLevel     = maxLevel,                        &
        & nSolve       = scheme%pdf(:)%nElems_solve,      &
        & levelDesc    = scheme%levelDesc(:),             &
        & tNow         = params%general%simControl%now    )
    end select  

    ! \todo KM: remove old mus_alloc_mrt.
    ! also mus_setParameters called twice at initial time once here and
    ! in the beginning of mus_control_module
    if ( scheme%header%kind(1:3) == 'lbm' ) then
      do iLevel = minLevel, maxLevel
        ! allocate fluid mrt array
        call mus_alloc_mrt( mrt = scheme%field(1)%fieldProp%   &
          &                                fluid%mrt(iLevel),  &
          &                 QQ    = scheme%layout%fStencil%QQ  )
        ! set relaxation parameter for MRT
        ! \todo remove
        !KM!call mus_setParameters( scheme, iLevel, params%general%simControl%now, &
        !KM!  &                     params%physics )

      end do
    end if

    ! create subTree for all spacetime function in the linked list of
    ! spacetime function
    call tem_create_subTree_of_st_funList(     &
      &       me      = scheme%st_funList,     &
      &       tree    = geometry%tree,         &
      &       bc_prop = geometry%boundary,     &
      &       stencil = scheme%layout%fStencil )

    ! initialize the source terms for all fields and global source
    call mus_init_sourceTerms( field        = scheme%field(:),            &
      &                        nFields      = scheme%nFields,             &
      &                        globSrc      = scheme%globSrc,             &
      &                        varSys       = scheme%varSys,              &
      &                        tree         = geometry%tree,              &
      &                        nElems_solve = scheme%pdf(:)%nElems_solve, &
      &                        levelDesc    = scheme%levelDesc       )

    ! initialize transport variables like velocity for passive scalar
    call mus_init_transport_var( me           = scheme%transVar,            &
      &                          varSys       = scheme%varSys,              &
      &                          tree         = geometry%tree,              &
      &                          nElems_solve = scheme%pdf(:)%nElems_solve, &
      &                          levelDesc    = scheme%levelDesc       )

    ! verify correct settings for the streaming layout
    call check_streaming_layout( minLevel, maxLevel )


    ! dynamic load balance time control homogenize
    if ( params%general%balance%dynamic ) then
      call mus_timeControl_homogenize(                     &
        &     me     = params%general%balance%timeControl, &
        &     dt     = params%physics%dtLvl( maxLevel ),   &
        &     reqInt = params%reqInterval                  )
    endif

    if ( ( params%general%restart%controller%writeRestart  .or. &
      &   params%general%restart%controller%readRestart ) ) then

      ! initialize the restart
      write(logUnit(1),*) 'Initializing restart...'
      !TG: deactivated because solver specific info given throug config file
      !TG:! write scheme info to restart solver specific conf
      !TG:! only root needs to write it since only root write header file
      !TG:if ( params%general%restart%controller%writeRestart ) then
      !TG:  call mus_writeSolverSpecInfo( scheme  = scheme,                   &
      !TG:    &                           params  = params,                   &
      !TG:    &                           rank    = params%general%proc%rank, &
      !TG:    &                           outUnit = params%solSpec_unit       )
      !TG:end if

      call tem_init_restart( me           = params%general%restart, &
        &                    solver       = params%general%solver,  &
        ! &                    varSys       = scheme%varSys,          &
        &                    varMap       = scheme%stateVarMap,     &
        &                    tree         = geometry%tree           )

      call mus_timeControl_homogenize(                               &
        &    me     = params%general%restart%controller%timeControl, &
        &    dt     = params%physics%dtLvl( maxLevel ),              &
        &    reqInt = params%reqInterval                             )
    end if


    ! initialize tracking objects.
    call mus_init_tracker( scheme    = scheme,   &
      &                    geometry  = geometry, &
      &                    params    = params    )

    ! convergence objects
    if ( params%general%simControl%abortCriteria%steady_state ) then
      write(logUnit(1),*) 'Initializing convergence...'

      do iConv = 1, size( params%general%simControl%abortCriteria%convergence)
        call mus_timeControl_homogenize(                             &
          &       me = params%general%simControl%abortCriteria       &
          &                  %convergence(iConv)%header%timeControl, &
          &       dt     = params%physics%dtLvl( maxLevel ),         &
          &       reqInt = params%reqInterval                        )
      end do

      call tem_init_convergence( me       = params%general%simControl        &
        &                                         %abortCriteria%convergence,&
        &                        tree     = geometry%tree,                   &
        &                        bc_prop  = geometry%boundary,               &
        &                        stencil  = scheme%layout%fStencil,          &
        &                        globProc = params%general%proc,             &
        &                        varSys   = scheme%varSys                    )
    end if

    if( minLevel /= maxlevel ) then
      write(logUnit(1),"(A)") 'Initializing interpolation...'
      ! initialize the interpolation
      call mus_init_interpolate(                             &
        &             intp        = scheme%intp,             &
        &             levelDesc   = scheme%levelDesc,        &
        &             scheme_kind = scheme%header%kind,      &
        &             stencil     = scheme%layout%fStencil,  &
        &             minLevel    = minLevel,                &
        &             maxLevel    = maxLevel,                &
        &             turbActive  = scheme%turbulence%active )
    end if
    if( main_debug%debugDependencies ) then
      call debug_dependencies( scheme%intp, scheme%levelDesc, &
        &                      geometry%tree, params%general%proc%rank )
      call dump_intpLists( minLevel, maxLevel, scheme%intp%config%order,  &
        &                  scheme%levelDesc, params%general%proc%rank )
    end if

    if( geometry%dynamicGeom ) then
      write(logUnit(1),*) 'Initializing geomIncr ...'
      call mus_init_geomIncr( me     = geometry%geomIncr,                &
        &                     varSys = scheme%varSys,                    &
        &                     dt     = params%physics%dtLvl( maxLevel ), &
        &                     reqInt = params%reqInterval                )
      write(logUnit(1),*) 'Done initializing geomIncr.'
    end if ! Time for geom incr

    call tem_horizontalSpacer( after = 1, fUnit = logUnit(1) )

  end subroutine mus_init_aux
! ****************************************************************************** !


! ****************************************************************************** !
  !> Set relaxation parameters for MRT
  !!
  subroutine mus_update_relaxParams( scheme, iLevel, tNow, physics )
    ! ---------------------------------------------------------------------------
    !> scheme type
    type(mus_scheme_type)          :: scheme
    !> level
    integer, intent(in)            :: iLevel
    !> global parameters
    type(tem_time_type),intent(in) :: tNow
    !> contains factors to convert physical to lattice unit and vice versa
    type(mus_physics_type), intent(in) :: physics
    ! ---------------------------------------------------------------------------

    select case(trim(scheme%header%kind))
    case('lbm', 'lbm_incomp')
      ! \todo KM: remove the old update relaxation after fixing all compute 
      ! kernels
      call mus_update_relaxation( me         = scheme%field(1)%fieldProp%fluid,&
        &                         moment     = scheme%layout%moment,           &
        &                         schemeHeader = scheme%header,                &
        &                         tNow       = tNow,                           &
        &                         level      = iLevel                          )

      ! Update kinematic viscosity from STfun and calculate turbulent viscosity
      ! from velocity gradient or nonEqPDF
      call mus_update_viscKine(                                              &
        & viscKine          = scheme%field(1)%fieldProp%fluid%viscKine,      &
        & state             = scheme%state(iLevel)%val(:,                    &
        &                            scheme%pdf(iLevel)%nNow),               &
        & neigh             = scheme%pdf(iLevel)%neigh,                      &
        & auxField          = scheme%auxField(iLevel)%val(:),                &
        & nSize             = scheme%pdf(iLevel)%nSize,                      &
        & nFluids           = scheme%pdf(iLevel)%nElems_fluid,               &
        & nGhostFromCoarser = scheme%pdf(iLevel)%nElems_ghostFromCoarser,    &
        & nGhostFromFiner   = scheme%pdf(iLevel)%nElems_ghostFromFiner,      &
        & baryOfTotal       = scheme%levelDesc(iLevel)%baryOfTotal,          &
        & varSys            = scheme%varSys,                                 &
        & iLevel            = iLevel,                                        &
        & layout            = scheme%layout,                                 &
        & tNow              = tnow,                                          &
        & viscRef           = physics%fac(iLevel)%visc,                      &
        & dxL               = scheme%field(1)%fieldProp%fluid%dxLvl(iLevel), &
        & dtL               = scheme%field(1)%fieldProp%fluid%dtLvl(iLevel), &
        & derVarPos         = scheme%derVarPos(1),                           &
        & turb              = scheme%turbulence                              )

      ! update fluid relaxation parameter omega kine
      call mus_update_relaxParamKine(                             &
        & viscKine    = scheme%field(1)%fieldProp%fluid%viscKine, &
        & turb        = scheme%turbulence,                        &
        & nSolve      = scheme%pdf(iLevel)%nElems_solve,          &
        & iLevel      = iLevel                                    )

!KM!      ! update fluid relaxation paramter omega bulk
!KM!      if (scheme%field(1)%fieldProp%fluid%requireOmegaBulk) then
!KM!        call mus_update_relaxParamFromViscSTfun(                        &
!KM!          & omega      = scheme%field(1)%fieldProp%fluid%viscBulk       &
!KM!          &                             %omLvl(iLevel)%val(:),          &
!KM!          & viscSTfun   = scheme%field(1)%fieldProp%fluid%viscBulk      &
!KM!          &                              %STfun,                        &
!KM!          & nSolve      = scheme%pdf(iLevel)%nElems_solve,              &
!KM!          & baryOfTotal = scheme%levelDesc(iLevel)%baryOfTotal,         &
!KM!          & tNow        = tnow,                                         &
!KM!          & viscRef     = physics%fac(iLevel)%visc,                     &
!KM!          & dtL         = scheme%field(1)%fieldProp%fluid%dtLvl(iLevel) )
!KM!      end if

    case('multi-species_gas','multi-species_liquid')
      call setParameters_multispecies( field   = scheme%field(:), &
        &                              nFields = scheme%nFields,  &
        &                              mixture = scheme%mixture,  &
        &                              header  = scheme%header,   &
        &                              layout  = scheme%layout,   &
        &                              iLevel  = iLevel,          &
        &                              tNow    = tNow )
    end select

  end subroutine mus_update_relaxParams
! ****************************************************************************** !


! ****************************************************************************** !
  !> This routine dumps tracking and restart when timeControl is active
  subroutine mus_dumpData( scheme, tree, levelPointer, general, restart_triggered )
    ! ---------------------------------------------------------------------------
    !> scheme type
    type( mus_scheme_type ), intent(inout) :: scheme
    !> Treelmesh data
    type(treelmesh_type), intent(in)       :: tree
    !> Level Pointer
    integer, intent(in) :: levelPointer(:)
    !> Global parameters
    type( tem_general_type ), intent(inout)  :: general
    !> 
    logical, intent(inout) :: restart_triggered
    ! ---------------------------------------------------------------------------

    if ( scheme%track%control%active ) then
      call tem_tracker( track          = scheme%track,       &
        &               simControl     = general%simControl, &
        &               tree           = tree,               &
        &               varSys         = scheme%varSys       )
    end if

    if( general%restart%controller%writeRestart ) then
      ! check time control and update it if triggered
      call tem_timeControl_check(                                       &
        &               me     = general%restart%controller%timeControl,&
        &               now    = general%simControl%now,                &
        &               comm   = general%proc%comm,                     &
        &               triggered = restart_triggered )

      if ( restart_triggered ) then
        write(logUnit(1),*) ' Writing restart at:'
        call tem_time_dump( general%simControl%now, logUnit(1) )
        call mus_writeRestart( levelPointer = levelPointer,                  &
          &                    restart      = general%restart,               &
          &                    scheme       = scheme,                        &
          &                    tree         = tree,                          &
          &                    timing       = general%simControl%now,        &
          &                    timerHandle  = mus_timerHandles%wRestart      )
        write(logUnit(1),*) 'Done writing'
      end if
    end if

  end subroutine mus_dumpData
! ****************************************************************************** !


end module mus_aux_module
! ****************************************************************************** !
