! See copyright notice in the COPYRIGHT file.
! **************************************************************************** !
!> author: Kannan Masilamani
!! This module contains routine to retrieve auxiliary field variables for 
!! getElement, getPoint, setupIndices and getValOfIndex.
!! Auxilary field variables are:
!!    * density and velocity for fluid 
!!    * species desity and velocity for multispecies
!!    * potential for poisson
!!
module mus_auxField_module
  use env_module,         only: rk
  use tem_param_module,   only: rho0, rho0Inv
  use tem_varSys_module,  only: tem_varSys_type, tem_varSys_op_type
  use tem_time_module,    only: tem_time_type
  use treelmesh_module,   only: treelmesh_type
  use tem_stencil_module, only: tem_stencilHeader_type
  use tem_comm_module,    only: tem_commpattern_type, tem_communication_type, &
    &                           tem_comm_init
  use tem_construction_module, only: tem_levelDesc_type
  use tem_time_module,         only: tem_time_type

  use mus_derVarPos_type_module, only: mus_derVarPos_type
  use mus_scheme_header_module,  only: mus_scheme_header_type

  implicit none
  private

  public :: mus_init_auxField
  public :: mus_auxFieldVar_type
  public :: mus_proc_calcAuxField

  !> Contains auxiliary field variable values per level and communication buffers
  type mus_auxFieldVar_type
    !> auxiliary field variable values computed from pre-collision PDF
    !! after PDF exchange
    !! Size: nSize*nScalars
    !! Element order is same as state array
    !! Access: (iElem-1)*nScalars + varSys%method%val(iVar)%auxField_varPos
    real(kind=rk), allocatable :: val(:)

    !> Local Fluids required by remote processes
    type( tem_communication_type ) :: sendBuffer
    !> Local ghostFromCoarser required by remote processes
    type( tem_communication_type ) :: sendBufferFromCoarser
    !> Local ghostFromFiner required by remote processes
    type( tem_communication_type ) :: sendBufferFromFiner
    !> My halos which are fluids on remote processes
    type( tem_communication_type ) :: recvBuffer
    !> My halos which are ghostFromCoarser on remote processes
    type( tem_communication_type ) :: recvBufferFromCoarser
    !> My halos which are ghostFromFiner on remote processes
    type( tem_communication_type ) :: recvBufferFromFiner
  end type mus_auxFieldVar_type

  abstract interface
    !> Interface to compute auxField vars i.e. conserved macroscopic moments
    !! from pre-collision PDF for fluid and ghostFromCoarser. 
    !! auxField on GhostFromFiner elements are interpolated and 
    !! halo elements are exchanged
    !! For Multicomponent models: in calcAuxField function, the velocity
    !! is computed on transformed PDF such that force term can be added to it
    !! in addSrcToAuxField routine. The auxField is updated with correct
    !! velocity field in compute kernel 
    !! i.e. velocity of original PDF is obtained by solving 
    !! linear equation system  in compute kernel
    subroutine mus_proc_calcAuxField(auxField, state, neigh, nSize, nSolve, &
      & iLevel, stencil, varSys, derVarPos)
      import :: rk, tem_varSys_type, tem_stencilHeader_type, mus_derVarPos_type 
      !> output auxField array
      real(kind=rk), intent(inout) :: auxField(:)
      !> input state array
      real(kind=rk), intent(in) :: state(:)
      !> connectivity array
      integer, intent(in) :: neigh(:)
      !> number of elements in the state array
      integer, intent(in) :: nSize
      !> number of fluid elements + ghostFromCoarser
      integer, intent(in) :: nSolve
      !> current level
      integer, intent(in) :: iLevel
      !> stencil header 
      type(tem_stencilHeader_type), intent(in) :: stencil
      !> variable system definition
      type(tem_varSys_type), intent(in) :: varSys
      !> position of derived quantities in varsys
      type(mus_derVarPos_type), intent(in) :: derVarPos(:)
    end subroutine mus_proc_calcAuxField  

  end interface

contains  

  ! ************************************************************************* !
  !> This routine initialize auxField var val array and communication buffers
  subroutine mus_init_auxField(me, levelDesc, pattern, nSize, nAuxScalars)
    ! --------------------------------------------------------------------- !
    !> Auxiliary field variable
    type(mus_auxFieldVar_type), intent(out) :: me
    !> levelDesc to access communication buffers of state array
    type(tem_levelDesc_type), intent(in) :: levelDesc
    !> communication pattern
    type(tem_commPattern_type), intent(in) :: pattern
    !> Number of elements in state array
    integer, intent(in) :: nSize
    !> Number of scalars in auxiliary variables
    integer, intent(in) :: nAuxScalars
    ! --------------------------------------------------------------------- !
    ! --------------------------------------------------------------------- !
    allocate(me%val(nSize * nAuxScalars))
    me%val(:) = -1000000.0_rk

    ! initialize send buffer
    call init_commBuffer(buffer_aux   = me%sendBuffer,       &
      &                  buffer_State = levelDesc%sendBuffer )

    ! initialize recv buffer
    call init_commBuffer(buffer_aux   = me%recvBuffer,       &
      &                  buffer_State = levelDesc%recvBuffer )

     ! initialize send buffer
    call init_commBuffer(buffer_aux   = me%sendBufferFromCoarser,       &
      &                  buffer_State = levelDesc%sendBufferFromCoarser )

    ! initialize recv buffer
    call init_commBuffer(buffer_aux   = me%recvBufferFromCoarser,       &
      &                  buffer_State = levelDesc%recvBufferFromCoarser )

      ! initialize send buffer
    call init_commBuffer(buffer_aux   = me%sendBufferFromFiner,       &
      &                  buffer_State = levelDesc%sendBufferFromFiner )

    ! initialize recv buffer
    call init_commBuffer(buffer_aux   = me%recvBufferFromFiner,       &
      &                  buffer_State = levelDesc%recvBufferFromFiner )

    contains

    ! ************************************************************************ !
    subroutine init_commBuffer(buffer_aux, buffer_state)
      !-------------------------------------------------------------------------
      !> communication buffer for velocity field
      type(tem_communication_type), intent(out) :: buffer_aux
      !> communication buffer of state array which is already initialized
      !! in tem_construction_module
      type(tem_communication_type), intent(in) :: buffer_state
      !-------------------------------------------------------------------------
      ! positions, in velocity vectors
      integer, allocatable :: pos(:)
      integer :: iProc, iElem, iScal, counter, elemPos
      !-------------------------------------------------------------------------
      ! copy information about target and source procs from pdf sendBuffer to
      ! velocity sendBuffer
      call tem_comm_init(buffer_aux, buffer_state%nProcs)
      buffer_aux%proc = buffer_state%proc
      buffer_aux%nElemsProc = buffer_state%nElemsProc

      allocate(pos(maxval(buffer_aux%nElemsProc*nAuxScalars)))
      
      do iProc = 1, buffer_aux%nProcs
        counter = 0
        do iElem = 1, buffer_aux%nElemsProc( iProc )
          ! halo element position in the levelDesc total list
          elemPos = buffer_state%elemPos( iProc )%val( iElem )
          do iScal = 1, nAuxScalars
            counter = counter + 1
            pos(counter) = (elemPos-1)*nAuxScalars + iScal
          end do
        end do
        ! copy position array to me%pos, allocate me%val array
        call pattern%initBuf_real( buffer_aux%buf_real( iProc ), pos, counter )
      end do 

      deallocate(pos)
    end subroutine init_commBuffer
    ! ************************************************************************ !

  end subroutine mus_init_auxField
  ! ************************************************************************* !


end module mus_auxField_module

