! See copyright notice in the COPYRIGHT file.
! ***************************************************************************** !
!> author: Kannan Masilamani
!! Module containing subroutines for initialize Musubi source 
!! variables and update source terms
!!
module mus_source_module
  ! include treelm modules
  use mpi
  use env_module,               only: rk, long_k, long_k_mpi
  use tem_aux_module,           only: tem_abort
  use tem_varSys_module,        only: tem_varSys_type
  use tem_time_module,          only: tem_time_type
  use treelmesh_module,         only: treelmesh_type
  use tem_logging_module,       only: logUnit
  use tem_construction_module,  only: tem_levelDesc_type
  use tem_geometry_module,      only: tem_baryOfID
  use tem_timer_module,         only: tem_startTimer, tem_stopTimer
  use tem_tools_module,         only: tem_horizontalSpacer
  use tem_debug_module,         only: dbgUnit

  ! include musubi modules
  use mus_pdf_module,            only: pdf_data_type
  use mus_field_module,          only: mus_field_type
  use mus_source_var_module,     only: mus_source_type, mus_source_op_type
  use mus_physics_module,        only: mus_convertFac_type
  use mus_derVarPos_type_module, only: mus_derVarPos_type
  use mus_timer_module,          only: mus_timerHandles

  implicit none
  private

  public :: mus_init_sourceTerms
  public :: mus_apply_sourceTerms
  !public :: mus_update_sourceFields

contains


  ! ************************************************************************** !
  !> This routine does set_params and setupIndices for all sources terms
  !! by gathering points to apply souce term before.
  subroutine mus_init_sourceTerms( field, nFields, globSrc, varSys, &
    &                              tree, nElems_solve, levelDesc )
    ! -------------------------------------------------------------------------- 
    !> Number of fields
    integer, intent(in)                    :: nFields

    !> contains sources of all fields
    type(mus_field_type), intent(inout)   :: field(:)

    !> global source
    type(mus_source_type), intent(inout)   :: globSrc

    !> global variable system
    type(tem_varSys_type), intent(in)      :: varSys

    !> global treelm mesh
    type(treelmesh_type), intent(in)     :: tree
    
    !> Number of elements to solve in all levels
    !! nFluids + nGhosts
    integer, intent(in) :: nElems_solve(tree%global%minLevel:)
    
    !> Level descriptors
    type( tem_levelDesc_type ), intent(in) :: levelDesc(tree%global%minLevel:)
    ! -------------------------------------------------------------------------- 
    integer :: iLevel, iField, iSrc
    integer :: nSolve, minLevel, maxLevel
    ! -------------------------------------------------------------------------- 
    call tem_horizontalSpacer(fUnit = logUnit(1))
    write(logUnit(1),*)'Initializing Source terms ...'
    write(dbgUnit(3),*)'Initializing Source terms ...'
   
    ! immediately exit this routine if there are no source terms to apply
    if ( all(field(:)%source%varDict%nVals == 0) &
      & .and. globSrc%varDict%nVals == 0 ) then
      write(logUnit(1),*) 'No active source terms'
      return
    end if  

    minLevel = tree%global%minLevel
    maxLevel = tree%global%maxLevel

    ! allocate source element array
    do iField = 1, nFields
      do iSrc = 1, field(iField)%source%varDict%nVals
        allocate( field(iField)%source%method(iSrc)%elemLvl(minLevel:maxLevel) )
      end do  
    end do

    do iSrc = 1, globSrc%varDict%nVals
      allocate( globSrc%method(iSrc)%elemLvl(minLevel:maxlevel) )
    end do  

    ! get source parameter to store in data variable for coupling
    write(logUnit(10),*) ' Setup indices for sources'
    write(dbgUnit(3),*) ' Setup indices for source'
    do iLevel = minLevel, maxLevel
      write(logUnit(10),*) '  iLevel: ', iLevel
      write(dbgUnit(3),*) '  iLevel: ', iLevel


      ! Use barycenter of all elements to solve i.e fluid+ghost to apply
      ! source terms
      nSolve = nElems_solve(iLevel)

      do iField = 1, nFields
        write(dbgUnit(3),*) '  field source: ', iField
        ! Store idx for active field source variables
        call mus_setupIndices_forSrc( source     = field(iField)%source, &
          &                           varSys     = varSys,               &
          &                           tree       = tree,                 &
          &                           nSolve     = nSolve,               &
          &                           bary       = levelDesc(iLevel)     &
          &                                        %baryOfTotal,         &
          &                           iLevel     = iLevel                )
      end do

      write(dbgUnit(3),*) '  glob source: '
      ! Store idx for active glob source variables
      call mus_setupIndices_forSrc( source     = globSrc,          &
        &                           varSys     = varSys,           &
        &                           tree       = tree,             &
        &                           nSolve     = nSolve,           &
        &                           bary       = levelDesc(iLevel) &
        &                                        %baryOfTotal,     &
        &                           iLevel     = iLevel            )

    end do ! iLevel 
    call tem_horizontalSpacer(fUnit = logUnit(1))

  end subroutine mus_init_sourceTerms
  ! *************************************************************************** !


  ! *************************************************************************** !
  !> This routines does setup indices for given source within a field or 
  !! global. Index are stored for points which source term is active
  subroutine mus_setupIndices_forSrc(source, varSys, nSolve, bary, iLevel, &
    &                                tree)
    ! -------------------------------------------------------------------------- 
    !> Source term to fill in
    type(mus_source_type), intent(inout) :: source

    !> global variable system
    type(tem_varSys_type), intent(in)      :: varSys

    !> global treelm mesh
    type( treelmesh_type ), intent(in)     :: tree
    
    !> Number of elements to apply source term  on this level
    integer, intent(in) :: nSolve

    !> Space coordinates to apply source terms
    real(kind=rk), intent(in) :: bary(:,:)

    !> Current level
    integer, intent(in) :: iLevel
    ! -------------------------------------------------------------------------- 
    integer :: iVar, iElem, counter, src_nElems, nComps
    integer :: data_varPos
    integer, allocatable :: idx(:)
    integer(kind=long_k) :: nElems_var(source%varDict%nVals) 
    integer(kind=long_k) :: glob_nElems_var(source%varDict%nVals)
    ! -------------------------------------------------------------------------- 
    allocate(idx(nSolve))

    do iVar = 1, source%varDict%nVals
      idx = 0
      data_varPos = source%method(iVar)%data_varPos
      ! number of components of the source field
      nComps = varSys%method%val(data_varPos)%nComponents
      ! set params
      call varSys%method%val(data_varPos)%set_params( &
        & varSys   = varSys,                          &
        & instring = 'isSurface = false'              )
  
      call varSys%method%val(data_varPos)%setup_indices( &
        & varSys     = varSys,                           &
        & point      = bary,                             &
        & iLevel     = iLevel,                           &
        & tree       = tree,                             &
        & nPnts      = nSolve,                           &
        & idx        = idx                               )
  
      ! Store only valid idx to apply source term only on shapes defined
      ! for a data variable
      src_nElems = count(idx>0)
      nElems_var(iVar) = src_nElems
      source%method(iVar)%elemLvl(iLevel)%nElems = src_nElems 
      allocate(source%method(iVar)%elemLvl(iLevel)%posInTotal(src_nElems))
      allocate(source%method(iVar)%elemLvl(iLevel)%idx(src_nElems))
!KM!      allocate(source%method(iVar)%elemLvl(iLevel)%val(src_nElems*nComps))

      counter = 0
      do iElem = 1, nSolve
        if (idx(iElem) > 0) then
          counter = counter + 1
          source%method(iVar)%elemLvl(iLevel)%posInTotal( counter ) = iElem
          source%method(iVar)%elemLvl(iLevel)%idx( counter ) = idx(iElem)
        end if
      end do
    end do !iVar

    !KM!call MPI_Reduce(nElems_var, glob_nElems_var, source%varDict%nVals, &
    !KM!  &             long_k_mpi, mpi_sum,                              &
    !KM!  &             0, tree%global%comm, ierror                        )
    do iVar = 1, source%varDict%nVals
      data_varPos = source%method(iVar)%data_varPos
      write(dbgUnit(3),*) 'Source iVar: ', iVar, &
        & trim(varSys%varName%val(data_varPos))
      write(dbgUnit(3),*) 'Total source nElems: ', nElems_var(iVar)
    end do  
  end subroutine mus_setupIndices_forSrc
  ! *************************************************************************** !


  ! *************************************************************************** !
  !> Apply all source terms i.e field specific source and global source on
  !! all fields.
  !!
  subroutine mus_apply_sourceTerms( field, nFields, globSrc, pdf, varSys,  &
    &                               iLevel, time, tree, phyConvFac, state, &
    &                               auxField, derVarPos )
    ! --------------------------------------------------------------------------
    !> Number of fields
    integer, intent(in)                :: nFields

    !> contains sources of all fields
    type(mus_field_type), intent(in)  :: field(nFields)

    !> global source
    type(mus_source_type), intent(in)  :: globSrc

    !> global tree information
    type(treelmesh_type), intent(in)   :: tree

    !> pdf datatype 
    type(pdf_data_type), intent(inout) :: pdf

    !> global variable system
    type(tem_varSys_type), intent(in)  :: varSys

    !> current level
    integer, intent(in)                :: iLevel

    !> current timing information
    type(tem_time_type), intent(in)    :: time

    !> state type containing the state vector to update
    real(kind=rk), intent(inout) :: state(:,:)

    !> auxField array
    real(kind=rk), intent(in) :: auxField(:)

    !> Physics conversion factor for current level
    type(mus_convertFac_type), intent(in) :: phyConvFac

    !> position of derived quantities in varsys
    type(mus_derVarPos_type), intent(in) :: derVarPos(:)
    ! -------------------------------------------------------------------------- 
    ! counter variables
    integer :: iVar, iField
    integer :: now, next
    ! -------------------------------------------------------------------------- 
    call tem_startTimer( timerHandle = mus_timerHandles%source )
    ! \todo create buffer for source to copy state values before applying 
    ! source term to avoid using state which overwritten by source variable

    now = pdf%nNow
    next = pdf%nNext

    ! loop over all active sourve variable to apply source
    do iField = 1, nFields
      do iVar = 1, field(iField)%source%varDict%nVals
        call field(iField)%source%method(iVar)%applySrc( &
          & inState    = state(:, now),                  &
          & outState   = state(:, next),                 &
          & neigh      = pdf%neigh(:),                   &
          & auxField   = auxField,                       &
          & nPdfSize   = pdf%nSize,                      &  
          & iLevel     = iLevel,                         &
          & varSys     = varSys,                         &
          & time       = time,                           &
          & derVarPos  = derVarPos,                      &
          & phyConvFac = phyConvFac                      )
      end do
    end do   

    ! apply global source
    do iVar = 1, globSrc%varDict%nVals
      call globSrc%method(iVar)%applySrc(  &
        & inState    = state(:, now),      &
        & outState   = state(:, next),     &
        & neigh      = pdf%neigh(:),       &
        & auxField   = auxField,           &
        & nPdfSize   = pdf%nSize,          &  
        & iLevel     = iLevel,             &
        & varSys     = varSys,             &
        & time       = time,               &
        & derVarPos  = derVarPos,          &
        & phyConvFac = phyConvFac          )
    end do

    call tem_stopTimer( timerHandle = mus_timerHandles%source )
  end subroutine mus_apply_sourceTerms
! ***************************************************************************** !

!KM!  ! *************************************************************************** !
!KM!  !> Update all source fields i.e field specific source and global source on
!KM!  !! all fields. 
!KM!  !!
!KM!  subroutine mus_update_sourceFields( field, nFields, globSrc, varSys, &
!KM!    &                               iLevel, time, phyConvFac )
!KM!    ! --------------------------------------------------------------------------
!KM!    !> Number of fields
!KM!    integer, intent(in)                :: nFields
!KM!
!KM!    !> contains sources of all fields
!KM!    type(mus_field_type), intent(inout)  :: field(nFields)
!KM!
!KM!    !> global source
!KM!    type(mus_source_type), intent(inout)  :: globSrc
!KM!
!KM!    !> global variable system
!KM!    type(tem_varSys_type), intent(in)  :: varSys
!KM!
!KM!    !> current level
!KM!    integer, intent(in)                :: iLevel
!KM!
!KM!    !> current timing information
!KM!    type(tem_time_type), intent(in)    :: time
!KM!
!KM!    !> Physics conversion factor for current level
!KM!    type(mus_convertFac_type), intent(in) :: phyConvFac
!KM!    ! -------------------------------------------------------------------------- 
!KM!    ! counter variables
!KM!    integer :: iVar, iField
!KM!    ! -------------------------------------------------------------------------- 
!KM!    call tem_startTimer( timerHandle = mus_timerHandles%source )
!KM!
!KM!    ! loop over all active sourve variable to apply source
!KM!    do iField = 1, nFields
!KM!      do iVar = 1, field(iField)%source%varDict%nVals
!KM!        if (field(iField)%source%method(iVar)%useSrcFieldVal) then
!KM!          call mus_update_singleSourceField(field(iField)%source%method(iVar))
!KM!        end if  
!KM!      end do
!KM!    end do   
!KM!
!KM!    ! apply global source
!KM!    do iVar = 1, globSrc%varDict%nVals
!KM!      if (globSrc%method(iVar)%useSrcFieldVal) then
!KM!        call mus_update_singleSourceField(globSrc%method(iVar))
!KM!      end if  
!KM!    end do
!KM!
!KM!    call tem_stopTimer( timerHandle = mus_timerHandles%source )
!KM!
!KM!  contains
!KM!    ! ************************************************************************ !
!KM!    !> This routine update single source field
!KM!    subroutine mus_update_singleSourceField(srcOp)
!KM!      ! ------------------------------------------------------------------------ 
!KM!      !> source data
!KM!      type(mus_source_op_type), intent(inout) :: srcOP
!KM!      ! ------------------------------------------------------------------------ 
!KM!      integer :: nElems
!KM!      ! ------------------------------------------------------------------------ 
!KM!      nElems = srcOp%elemLvl(ilevel)%nElems
!KM!  
!KM!      ! Get force which is refered in config file either its
!KM!      ! spacetime variable or operation variable
!KM!      call varSys%method%val(srcOp%data_varPos)%get_valOfIndex( &
!KM!        & varSys  = varSys,                                     &
!KM!        & time    = time,                                       &
!KM!        & iLevel  = iLevel,                                     &
!KM!        & idx     = srcOp%elemLvl(iLevel)%idx(1:nElems),        &
!KM!        & nVals   = nElems,                                     & 
!KM!        & res     = srcOp%elemLvl(iLevel)%val                   )
!KM!  
!KM!      select case (trim(srcOp%varname))
!KM!      case ('force','force_explicit')
!KM!        srcOp%elemLvl(iLevel)%val(:) = srcOp%elemLvl(iLevel)%val(:) &
!KM!          & / phyConvFac%body_force
!KM!      end select
!KM!  
!KM!    end subroutine mus_update_singleSourceField
!KM!  ! ************************************************************************** !
!KM!
!KM!  end subroutine mus_update_sourceFields
!KM!  ! ************************************************************************** !

end module mus_source_module
! ***************************************************************************** !
