! See copyright notice in the COPYRIGHT file.
! ***************************************************************************** !
!> author: Kannan Masilamani
!! Module containing subroutines for building MUSUBI specific source 
!! variables
!!
module mus_source_var_module
  use, intrinsic :: iso_c_binding, only: c_ptr, c_f_pointer

  ! include treelm modules
  use env_module,               only: rk, long_k, labelLen, solSpecLen
  use tem_aux_module,           only: tem_abort
  use tem_varSys_module,        only: tem_varSys_type,               &
    &                                 tem_varSys_append_derVar,      &
    &                                 tem_varSys_proc_point,         &
    &                                 tem_varSys_proc_element,       &
    &                                 tem_varSys_proc_setParams,     &
    &                                 tem_varSys_proc_getParams,     &
    &                                 tem_varSys_proc_setupIndices,  &
    &                                 tem_varSys_proc_getValOfIndex, &
    &                                 tem_varSys_dump
  use tem_varMap_module,        only: tem_possible_variable_type, &
    &                                 init, append, truncate,     &
    &                                 tem_variable_loadMapping
  use tem_stringKeyValuePair_module, only: init, truncate, &
    &                                      grw_stringKeyValuePairArray_type
  use tem_stencil_module,       only: tem_stencilHeader_type
  use tem_time_module,          only: tem_time_type
  use treelmesh_module,         only: treelmesh_type
  use tem_logging_module,       only: logUnit

  ! include musubi modules
  use mus_scheme_header_module, only: mus_scheme_header_type
  use mus_physics_module,       only: mus_convertFac_type
  use mus_derVarPos_type_module, only: mus_derVarPos_type

  ! include aotus modules
  use aotus_module,    only: flu_State

  implicit none
  private

  public :: mus_source_type
  public :: mus_source_op_type
  public :: mus_create_poss_srcVar
  public :: mus_load_source_var

  ! *************************************************************************** !
  !> Contains source elements position in state array and idx to access 
  !! data variable refered in config file.
  !! This type is defined for each level
  type mus_source_elems_type
    !> Number of source elements on this level.
    !! nFluids + nGhosts
    integer :: nElems

    !> Position of elements in state array to apply source terms.
    !! Position in state array is same as position in total list
    !! Size: nElems
    integer, allocatable :: posInTotal(:)

    !> Index to access point data type to retrieve values from variable 
    !! refered for source variable
    integer, allocatable :: idx(:)

    !> source field value obtained from ST_fun data variable.
    !! Filled only for elements where source is active i.e. elements in
    !! posInTotal. 
    !! size: nElems*nComponents
    !! \todo KM: might be not neccessary
!KM!    real(kind=rk), allocatable :: val(:)
  end type mus_source_elems_type
  ! *************************************************************************** !

  ! *************************************************************************** !
  !> Description contains list of elements on which source is active and
  !! function pointer to update source
  type mus_source_op_type
    !> Position of this source term variable in the varSys
    integer :: srcTerm_varPos

    !> Position of data variable provided in config file in the varSys
    integer :: data_varPos

    !> Contains source elements position in state array/total list for 
    !! each level
    type(mus_source_elems_type), allocatable :: elemLvl(:)

    !> Function to update state with source term
    procedure(proc_apply_source), pointer :: applySrc => null()

    !> to use source field array
    logical :: useSrcFieldVal = .false.

    !> name of the source variable
    character(len=labelLen) :: varname

    !> Function pointer to append source field to auxilary variable
    procedure(proc_addSrcToAuxField), pointer :: addSrcToAuxField => null()

  end type mus_source_op_type
  ! *************************************************************************** !


  ! *************************************************************************** !
  !> Description of musubi source type 
  type mus_source_type
    !> Contains source elements position in tree%treeID and
    !! function pointer to update source
    !! Size: varDict%nVals
    type(mus_source_op_type), allocatable :: method(:)

    !> Dictionary of source variable with 
    !! varDict%val()%key is the name of source variable and
    !! varDict%val()%value is the name of variable provided for the key
    type(grw_stringKeyValuePairArray_type) :: varDict
  end type mus_source_type
  ! *************************************************************************** !


  ! *************************************************************************** !
  abstract interface
    !> Abstract interface to update state with source terms
    subroutine proc_apply_source( fun, inState, outState, neigh, auxField, &
      & nPdfSize, iLevel, varSys, time, phyConvFac, derVarPos )
      import :: rk, mus_source_op_type, tem_varSys_type, tem_time_type, &
        &       mus_convertFac_type, mus_derVarPos_type

      !> Description of method to update source
      class(mus_source_op_type), intent(in) :: fun

      !> input  pdf vector
      !! \todo KM: instate is passed to compute auxField.
      !! Since auxField is precomputed from instate and passed to this routine.
      !! instate can be removed
      real(kind=rk), intent(in)  ::  inState(:)
  
      !> output pdf vector
      real(kind=rk), intent(inout) :: outState(:)
  
      !> connectivity Array corresponding to state vector
      integer,intent(in) :: neigh(:)

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

      !> number of elements in state Array
      integer, intent(in) :: nPdfSize

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

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

      !> Point in time at which to evaluate the variable.
      type(tem_time_type), intent(in)  :: time

      !> 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(:)
    end subroutine proc_apply_source

    !> Interface to add source to auxField vars if addSrcToAuxField is true 
    !! in source_op_type for all nSolve elements (nFluids+nGhostFromCoarser
    !! +nGhostFromFiner). Halo elements are exchanged
    subroutine proc_addSrcToAuxField(fun, auxField, iLevel, time, varSys, &
      & phyConvFac, derVarPos)
      import :: rk, tem_varSys_type, mus_derVarPos_type, mus_convertFac_type, &
        &       tem_time_type, mus_source_op_type

      !> Description of method to update source
      class(mus_source_op_type), intent(in) :: fun
      !> output auxField array
      real(kind=rk), intent(inout)          :: auxField(:)
      !> current level
      integer, intent(in)                   :: iLevel
      !> current timing information
      type(tem_time_type), intent(in)       :: time
      !> variable system definition
      type(tem_varSys_type), intent(in)     :: varSys
      !> 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(:)
    end subroutine proc_addSrcToAuxField 

  end interface
  ! *************************************************************************** !

contains


  ! *************************************************************************** !
  !> Routine initialize possible source variable depends on scheme kind
  subroutine mus_create_poss_srcVar(poss_srcVar, schemeHeader)
    ! --------------------------------------------------------------------------!
    !> possible source variables
    type(tem_possible_variable_type), intent(out) :: poss_srcVar

    !> Identifier of the scheme
    type(mus_scheme_header_type), intent(in) :: schemeHeader
    ! --------------------------------------------------------------------------!
    write(logUnit(10),*) 'Creating possible source terms '
    call init(me = poss_srcVar, length = 2 )

    select case(trim(schemeHeader%kind))
    case ('lbm', 'lbm_incomp')
      ! Second order force term
      call append(me          = poss_srcVar, &
        &         varName     = 'force',     &
        &         nComponents = 3            )
      ! first order force term for small time steps 
      call append(me          = poss_srcVar,      &
        &         varName     = 'force_explicit', &
        &         nComponents = 3                 )
      ! Omega sponge layer  
      call append(me          = poss_srcVar,    &
        &         varName     = 'sponge_omega', &
        &         nComponents = 1               )
      ! Sponge in terms of viscosity  
      call append(me          = poss_srcVar,    &
        &         varName     = 'sponge_visc',  &
        &         nComponents = 1               )
    case ('passive_scalar')
      call append(me          = poss_srcVar,       &
        &         varName     = 'equal_injection', &
        &         nComponents = 1                  )

      call append(me          = poss_srcVar, &
        &         varName     = 'injection', &
        &         nComponents = 1            )
    case ('nernst_planck')
      call append(me          = poss_srcVar,      &
        &         varName     = 'electric_field', &
        &         nComponents = 3                 )

    case ('multi-species_liquid')
      call append(me          = poss_srcVar,     &
        &         varName     = 'gravity_field', &
        &         nComponents = 3                )

      call append(me          = poss_srcVar,      &
        &         varName     = 'electric_field', &
        &         nComponents = 3                 )

      call append(me          = poss_srcVar,                  &
        &         varName     = 'electric_field_1order_diff', &
        &         nComponents = 3                             )

      call append(me          = poss_srcVar,                   &
        &         varName     = 'electric_field_1order_const', &
        &         nComponents = 3                              )

      call append(me          = poss_srcVar,                 &
        &         varName     = 'electric_field_1order_old', &
        &         nComponents = 3                            )

      call append(me          = poss_srcVar,             &
        &         varName     = 'electric_field_2order_nodiff', &
        &         nComponents = 3                        )

      call append(me          = poss_srcVar,             &
        &         varName     = 'electric_field_2order_new', &
        &         nComponents = 3                        )
    case ('poisson')
      call append(me          = poss_srcVar,      &
        &         varName     = 'charge_density', &
        &         nComponents = 1                 )
    case default
      write(logUnit(1),*) 'No possible source term defined for scheme kind:' &
        &                 //trim(schemeHeader%kind)
    end select
    call truncate(poss_srcVar)

  end subroutine mus_create_poss_srcVar
  ! *************************************************************************** !


  ! ***************************************************************************!
  !> Routine load musubi source terms for given key.
  !! key is glob_source or source
  subroutine mus_load_source_var(me, possVars, conf, parent, key, varSys)
    ! --------------------------------------------------------------------------!
    !> Source variable type to initialize
    type(mus_source_type), intent(out) :: me
    !> possible source variables
    type(tem_possible_variable_type), intent(in) :: possVars
    !> flu state
    type( flu_State ) :: conf
    !> parent handle if scheme table is defined
    integer, intent(in), optional :: parent
    !> key to load source term
    character(len=*), intent(in) :: key
    !> Global variable system
    type(tem_varSys_type), intent(inout) :: varSys
    ! --------------------------------------------------------------------------!
    ! initialize growing array stringKeyValuePair
    call init( me = me%varDict )

    ! load the global source variables
    call tem_variable_loadMapping( possVars = possVars,   &
      &                            conf     = conf,       &
      &                            parent   = parent,     &
      &                            key      = key,        &
      &                            varDict  = me%varDict, &
      &                            varSys   = varSys      )

    call truncate( me = me%varDict )

  end subroutine mus_load_source_var
  ! ***************************************************************************!

end module mus_source_var_module
! ***************************************************************************** !
