module dynPatchStateUpdaterMod

  !---------------------------------------------------------------------------
  !
  ! !DESCRIPTION:
  ! Class for adjusting patch-level (aboveground) state variables due to transient patch
  ! areas.
  !
  ! In each time step, the object should be set up with:
  !
  !    - call patch_state_updater%set_old_weights (before dyn subgrid weight updates)
  !
  !    - call patch_state_updater%set_new_weights (after dyn subgrid weight updates)
  !
  ! Then it can be used to update each state variable with a call to:
  !
  !    - call patch_state_updater%update_patch_state
  !
  ! !USES:
#include "shr_assert.h"
  use shr_kind_mod         , only : r8 => shr_kind_r8
  use shr_infnan_mod       , only : nan => shr_infnan_nan
  use shr_log_mod          , only : errMsg => shr_log_errMsg
  use decompMod            , only : bounds_type, BOUNDS_LEVEL_PROC
  use VegetationType       , only : veg_pp
  use ColumnType           , only : col_pp
  use elm_varpar           , only : mxpft
  use abortutils           , only : endrun
  use elm_varcon           , only : spval
  !
  implicit none
  private
  !
  ! !PUBLIC TYPES:
  public :: patch_state_updater_type

  ! Public routines
  public :: set_old_patch_weights     ! set weights before dyn subgrid updates
  public :: set_new_patch_weights     ! set weights after dyn subgrid updates

  ! Update a patch-level state variable and compute associated fluxes based on changing
  ! patch areas
  public :: update_patch_state

  ! Update a patch-level state variable and compute associated fluxes based on
  ! changing patch areas, with flux out partitioned into two fluxes. Partitioning is
  ! based on pft type.
  public :: update_patch_state_partition_flux_by_type

  type :: patch_state_updater_type

     real(r8), pointer :: pwtgcell_old(:) => null() ! old patch weights on the gridcell
     real(r8), pointer :: pwtgcell_new(:) => null()! new patch weights on the gridcell

     real(r8), pointer :: cwtgcell_old(:) => null()! old column weights on the gridcell

     ! (pwtgcell_new - pwtgcell_old) from last call to set_new_weights
     real(r8), pointer :: dwt(:) => null()

     ! (pwtgcell_old / pwtgcell_new) from last call to set_new_weights; only valid for
     ! growing patches
     real(r8), pointer :: growing_old_fraction(:) => null()

     ! (dwt / pwtgcell_new) from last call to set_new_weights; only valid for growing
     ! patches
     real(r8), pointer :: growing_new_fraction(:) => null()

   contains
     ! Public routines
     !procedure, public :: set_old_patch_weights     ! set weights before dyn subgrid updates
     !procedure, public :: set_new_patch_weights     ! set weights after dyn subgrid updates



     ! returns a patch-level logical array that is true wherever the patch weight was zero
     ! prior to weight updates
     procedure, public :: old_weight_was_zero

     ! returns a patch-level logical array that is true wherever the patch grew in this
     ! time step
     procedure, public :: patch_grew

     ! returns a patch-level logical array that is true wherever a patch newly has
     ! non-zero weight in this time step
     procedure, public :: patch_initiating

  end type patch_state_updater_type

  interface patch_state_updater_type
     module procedure constructor
  end interface patch_state_updater_type

  ! Update a patch-level state variable and compute associated fluxes based on changing
  ! patch areas

  character(len=*), parameter, private :: sourcefile = &
       __FILE__

contains

  ! ========================================================================
  ! Constructors
  ! ========================================================================

  !-----------------------------------------------------------------------
  function constructor(bounds) result(this)
    !
    ! !DESCRIPTION:
    ! Initialize a patch_state_updater_type object
    !
    ! !USES:
    !
    ! !ARGUMENTS:
    type(patch_state_updater_type) :: this  ! function result
    type(bounds_type), intent(in) :: bounds
    !
    ! !LOCAL VARIABLES:
    integer :: begp, endp
    integer :: begc, endc

    character(len=*), parameter :: subname = 'constructor'
    !-----------------------------------------------------------------------

    SHR_ASSERT(bounds%level == BOUNDS_LEVEL_PROC, errMsg(sourcefile, __LINE__))

    begp = bounds%begp
    endp = bounds%endp
    begc = bounds%begc
    endc = bounds%endc

    allocate(this%pwtgcell_old(begp:endp))
    this%pwtgcell_old(:) = spval
    allocate(this%pwtgcell_new(begp:endp))
    this%pwtgcell_new(:) = spval
    allocate(this%cwtgcell_old(begc:endc))
    this%cwtgcell_old(:) = spval
    allocate(this%dwt(begp:endp))
    this%dwt(:) = spval
    allocate(this%growing_old_fraction(begp:endp))
    this%growing_old_fraction(:) = spval
    allocate(this%growing_new_fraction(begp:endp))
    this%growing_new_fraction(:) = spval

  end function constructor

  ! ========================================================================
  ! Public methods
  ! ========================================================================

  !-----------------------------------------------------------------------
  subroutine set_old_patch_weights(this, bounds)
    !
    ! !DESCRIPTION:
    ! Set subgrid weights before dyn subgrid updates
    !
    ! !USES:
    !
    ! !ARGUMENTS:
    class(patch_state_updater_type), intent(inout) :: this
    type(bounds_type), intent(in) :: bounds
    !
    ! !LOCAL VARIABLES:
    integer :: p
    integer :: c

    character(len=*), parameter :: subname = 'set_old_weights'
    !-----------------------------------------------------------------------

    do p = bounds%begp, bounds%endp
       c = veg_pp%column(p)
       this%pwtgcell_old(p) = veg_pp%wtgcell(p)
       this%cwtgcell_old(c) = col_pp%wtgcell(c)
    end do

  end subroutine set_old_patch_weights

  !-----------------------------------------------------------------------
  subroutine set_new_patch_weights(this, bounds)
    !
    ! !DESCRIPTION:
    ! Set subgrid weights after dyn subgrid updates
    !
    ! !USES:
    !$acc routine seq
    ! !ARGUMENTS:
    class(patch_state_updater_type), intent(inout) :: this
    type(bounds_type), intent(in) :: bounds
    !
    ! !LOCAL VARIABLES:
    integer :: p

    !character(len=*), parameter :: subname = 'set_new_weights'
    !-----------------------------------------------------------------------

    do p = bounds%begp, bounds%endp
       this%pwtgcell_new(p) = veg_pp%wtgcell(p)
       this%dwt(p)          = this%pwtgcell_new(p) - this%pwtgcell_old(p)
       if (this%dwt(p) > 0._r8) then
          this%growing_old_fraction(p) = this%pwtgcell_old(p)/ this%pwtgcell_new(p)
          this%growing_new_fraction(p) = this%dwt(p)         / this%pwtgcell_new(p)
       else
          ! These values are unused in this case, but set them to something reasonable for
          ! safety. (We could set them to NaN, but that requires a more expensive
          ! subroutine call, using the shr_infnan_mod infrastructure.)
          this%growing_old_fraction(p) = 1._r8
          this%growing_new_fraction(p) = 0._r8
       end if
    end do

  end subroutine set_new_patch_weights

  !-----------------------------------------------------------------------
  subroutine update_patch_state(this, p, c, &
       var, flux_out_col_area, flux_out_grc_area, &
       seed, seed_addition)
    !
    ! !DESCRIPTION:
    ! Update a patch-level state variable and compute associated fluxes based on changing
    ! patch areas
    !
    ! For growing patches, this subroutine adjusts var. For shrinking patches, this
    ! subroutine accumulates flux in flux_out_col_area.
    !
    ! Changes are only made within the given filter. Note that this filter should include
    ! inactive as well as active patches, so that it includes patches that just became
    ! inactive.
    !
    ! !USES:
    !$acc routine seq
    ! !ARGUMENTS:
    type(patch_state_updater_type), intent(in) :: this
    integer , value  , intent(in) :: p, c
    real(r8), intent(inout) :: var     ! patch-level state variable

    ! Accumulated flux from shrinking areas, expressed as mass per unit area COLUMN, using
    ! the OLD column weight. (The use of the old column weight is appropriate if the
    ! fluxes from shrinking patches are applied to column-level state variables BEFORE
    ! doing column-level state adjustments via the column_state_updater). For shrinking
    ! areas, this is given as a NEGATIVE quantity. Often you will provide one of
    ! flux_out_col_area or flux_out_grc_area, but it is okay to provide both, or to
    ! provide neither if you don't need to track the flux out from this state variable.
    real(r8), intent(inout), optional :: flux_out_col_area

    ! Accumulated flux from shrinking areas, expressed as mass per unit area GRIDCELL. For
    ! shrinking areas, this is given as a NEGATIVE quantity. Often you will provide one of
    ! flux_out_col_area or flux_out_grc_area, but it is okay to provide both, or to
    ! provide neither if you don't need to track the flux out from this state variable.
    real(r8), intent(inout), optional :: flux_out_grc_area

    ! If provided, this gives some 'seed' amount added to the state in the area into
    ! which each growing patch grows. The value is ignored for patches that are either
    ! constant or shrinking in area.
    real(r8), intent(in), optional :: seed

    ! If provided, this accumulates the amount of seed added to each patch. This gives
    ! seed(p) * dwt(p). This can only be provided if seed is provided. Even though this is
    ! a patch-level array, it is expressed as mass per unit area GRIDCELL.
    real(r8), intent(inout), optional :: seed_addition
    !
    character(len=*), parameter :: subname = 'update_patch_state'
    !-----------------------------------------------------------------------

    if (present(seed_addition)) then
       if (.not. present(seed)) then
          call endrun(subname//' ERROR: seed_addition can only be provided if seed is provided')
       end if
    end if

    if (this%dwt(p) > 0._r8) then
      var = var * this%growing_old_fraction(p)
      if (present(seed)) then
         var = var + seed * this%growing_new_fraction(p)
         if (present(seed_addition)) then
            seed_addition = seed_addition + seed * this%dwt(p)
         end if
      end if

    else if (this%dwt(p) < 0._r8) then
      if (present(flux_out_grc_area)) then
         flux_out_grc_area = flux_out_grc_area + var * this%dwt(p)
      end if
      if (present(flux_out_col_area)) then
         ! No need to check for divide by 0 here: If dwt < 0 then we must have had
         ! cwtgcell_old > 0.
         flux_out_col_area = flux_out_col_area + &
              var * (this%dwt(p) / this%cwtgcell_old(c))
      end if
    end if

  end subroutine update_patch_state

  !-----------------------------------------------------------------------
  subroutine update_patch_state_partition_flux_by_type(this, p,c, &
       flux1_out_dest, flux1_fraction_by_pft_type, &
       var, flux1_out, flux2_out, &
       seed, seed_addition)
    !
    ! !DESCRIPTION:
    ! Update a patch-level state variable and compute associated fluxes based on
    ! changing patch areas, with flux out partitioned into two fluxes. Partitioning is
    ! based on pft type.
    !
    ! !USES:
    !$acc routine seq
    ! !ARGUMENTS:
    type(patch_state_updater_type), intent(in) :: this
    integer   , value, intent(in) :: p,c
    character(len=1), intent(in) :: flux1_out_dest  ! flux1_out to column or grid (default)
    real(r8), intent(in) :: flux1_fraction_by_pft_type( 0: ) ! fraction of flux that goes into flux1_out, indexed by pft type
    real(r8), intent(inout) :: var  ! patch-level state variable

    ! Accumulated fluxes from shrinking areas. For shrinking areas, these are given as
    ! NEGATIVE quantities. Even though these are patch-level arrays, they are expressed
    ! as mass per unit area GRIDCELL (so these are equivalent to the flux_out_grc_area
    ! argument in the main update_patch_state routine).
    real(r8), intent(inout) :: flux1_out
    real(r8), intent(inout) :: flux2_out

    ! If provided, this gives some 'seed' amount added to the state in the area into
    ! which each growing patch grows. The value is ignored for patches that are either
    ! constant or shrinking in area.
    real(r8), intent(in), optional :: seed

    ! If provided, this accumulates the amount of seed added to each patch. This gives
    ! seed(p) * dwt(p). This can only be provided if seed is provided. Even though this is
    ! a patch-level array, it is expressed as mass per unit area GRIDCELL.
    real(r8), intent(inout), optional :: seed_addition
    !
    ! !LOCAL VARIABLES:
    real(r8) :: total_flux_out
    real(r8) :: my_flux1_fraction
    !-----------------------------------------------------------------------


    total_flux_out  = 0._r8
    if (flux1_out_dest=='c') then
       call update_patch_state(this, p, c, &
          var, flux_out_col_area = total_flux_out, &
          seed = seed, seed_addition = seed_addition)
    else
       call update_patch_state(this,p,c, &
          var, flux_out_grc_area = total_flux_out, &
          seed = seed, seed_addition = seed_addition)
    end if

    my_flux1_fraction = flux1_fraction_by_pft_type(veg_pp%itype(p))
    flux1_out = flux1_out + total_flux_out * my_flux1_fraction
    flux2_out = flux2_out + total_flux_out * (1._r8 - my_flux1_fraction)

  end subroutine update_patch_state_partition_flux_by_type


  !-----------------------------------------------------------------------
  function old_weight_was_zero(this, bounds)
    !
    ! !DESCRIPTION:
    ! Returns a patch-level logical array that is true wherever the patch weight was zero
    ! prior to weight updates
    !
    ! !USES:
    !
    ! !ARGUMENTS:
    class(patch_state_updater_type), intent(in) :: this
    type(bounds_type), intent(in) :: bounds
    logical :: old_weight_was_zero(bounds%begp:bounds%endp)  ! function result
    !
    ! !LOCAL VARIABLES:
    integer :: p

    character(len=*), parameter :: subname = 'old_weight_was_zero'
    !-----------------------------------------------------------------------

    do p = bounds%begp, bounds%endp
       old_weight_was_zero(p) = (this%pwtgcell_old(p) == 0._r8)
    end do

  end function old_weight_was_zero

  !-----------------------------------------------------------------------
  function patch_grew(this, bounds)
    !
    ! !DESCRIPTION:
    ! Returns a patch-level logical array that is true wherever the patch grew in this
    ! time step
    !
    ! !USES:
    !
    ! !ARGUMENTS:
    class(patch_state_updater_type), intent(in) :: this
    type(bounds_type), intent(in) :: bounds
    logical :: patch_grew(bounds%begp:bounds%endp)  ! function result
    !
    ! !LOCAL VARIABLES:
    integer :: p

    character(len=*), parameter :: subname = 'patch_grew'
    !-----------------------------------------------------------------------

    do p = bounds%begp, bounds%endp
       patch_grew(p) = (this%dwt(p) > 0._r8)
    end do

  end function patch_grew

  !-----------------------------------------------------------------------
  function patch_initiating(this, bounds)
    !
    ! !DESCRIPTION:
    ! Returns a patch-level logical array wherever the patch is initiating - i.e., growing
    ! from zero area to non-zero area - in this time step
    !
    ! !USES:
    !
    ! !ARGUMENTS:
    class(patch_state_updater_type), intent(in) :: this
    type(bounds_type), intent(in) :: bounds
    logical :: patch_initiating(bounds%begp:bounds%endp)  ! function result
    !
    ! !LOCAL VARIABLES:
    integer :: p

    character(len=*), parameter :: subname = 'patch_initiating'
    !-----------------------------------------------------------------------

    do p = bounds%begp, bounds%endp
       patch_initiating(p) = ( &
            this%pwtgcell_old(p) == 0._r8 .and. &
            this%pwtgcell_new(p) > 0._r8)
    end do

  end function patch_initiating


end module dynPatchStateUpdaterMod
