! Copyright 2008-2021 Earth Sciences Department, BSC-CNS
!
! This file is part of MONARCH.
!
! MONARCH is free software: you can redistribute it and/or modify
! it under the terms of the GNU General Public License as published by
! the Free Software Foundation, either version 3 of the License, or
! (at your option) any later version.
!
! MONARCH is distributed in the hope that it will be useful,
! but WITHOUT ANY WARRANTY; without even the implied warranty of
! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
! GNU General Public License for more details.
!
! You should have received a copy of the GNU General Public License
! along with MONARCH. If not, see <http://www.gnu.org/licenses/>.

!> \file
!> The monarch_wet_deposition_t object and related functions

!> \page config_wet_deposition Wet deposition configuration file
!!
!! Wet deposition files are in \c json format and structured as follows:
!!
!! \code{.json}
!! {
!!   "wet-dep-data": {
!!     "resolved clouds" : {
!!       "calculate" : true,
!!       "output rates" : true
!!     },
!!     "sub-grid scale clouds" : {
!!       "calculate" : true,
!!       "output rates" : true
!!     },
!!     "vertical mixing" : {
!!       "calculate" : true,
!!       "output rates" : true
!!     }
!!   }
!! }
!! \endcode
!!
!! The \c wet-dep-data, each of the three calculation types, and the
!! \c calculate key-value pairs are required. When a \c calculate flag is set
!! to \c false, no calculations of that type are performed. If all
!! calculation types are set to not be calculated, wet deposition reactions
!! in PartMC will retain their default values (zero unless otherwise
!! specified).
!!
!! The \c output \c rates key-value pairs are optional and indicate whether
!! calculated rates of that type are output to the history file. They default
!! to \c false.

!> Functions to calculate wet deposition
module monarch_wet_deposition
! MONARCH always has MPI and JSON
#define MONARCH_USE_MPI

#ifdef DO_PARTMC
! PartMC always has JSON when used in MONARCH
#define PMC_USE_JSON

! Prefix to use in labelling PMC wet deposition rxns for updates
!( This can be any integer, but must be unique to wet deposition )
#define PMC_WET_DEP_PREFIX 8675309

  use monarch_constants,        only : mint, mflt, mdbl
  use monarch_state_changer,    only : monarch_state_array_ptr
  use monarch_util,             only : string_t
  use pmc_phlex_core,           only : phlex_core_t
  use pmc_property,             only : property_t
  use pmc_rxn_wet_deposition,   only : rxn_update_data_wet_deposition_rate_t

  implicit none
  private

  public :: monarch_wet_deposition_t

  !> Number of possible output arrays
  integer(kind=mint), parameter :: NUM_RATE_ARRAYS = 3
  !> Index for output array for resolved cloud loss rates
  integer(kind=mint), parameter :: RESOLVED_CLOUD_LOSS_RATE = 1
  !> Index for output array for sub-grid scale cloud loss rates
  integer(kind=mint), parameter :: SUB_GRID_SCALE_LOSS_RATE = 2
  !> Index for output array for relative change rates from vertical mixing
  integer(kind=mint), parameter :: VERTICAL_MIXING_RATE     = 3

  !> PartMC wet deposition
  !!
  !! This type provides functions to calculate wet deposition rates and update
  !! the PartMC-phlex first-order loss reactions accordingly
  type monarch_wet_deposition_t
    private
    !> Phlex-chem core
    type(phlex_core_t), pointer :: phlex_core => null( )
    !> Number of vertical model layers
    integer(kind=mint) :: n_layers
    !> Number of aqueous phases subject to removal by rainout
    integer(kind=mint) :: n_aqueous_phase
    !> Number of species involved in vertical mixing
    integer(kind=mint) :: n_mixing_species
    !> Wet deposition rate update nugget
    type(rxn_update_data_wet_deposition_rate_t) :: rate_update
    !> Mixing species concentrations
    real(kind=mflt), allocatable :: column_conc( :, : )
    !> Resolved-cloud loss rate [s-1]
    real(kind=mdbl), allocatable :: resolved_cloud_loss_rate( :,:,: )
    !> Sub-grid-scale cloud loss rate [s-1]
    real(kind=mflt), allocatable :: sub_grid_scale_cloud_loss_rate( :,:,: )
    !> Property set
    type(property_t), pointer :: property_set => null( )
    !> Flag indicating whether to calculate wet deposition rates and mixing
    logical, dimension(NUM_RATE_ARRAYS) :: calculate_rates = .false.
    !> Flag for output of wet deposition rates
    logical, dimension(NUM_RATE_ARRAYS) :: output_rates = .false.
    !> Species labels for outputting mixing rates
    type(string_t), allocatable :: species_labels( : )
    !> State arrays for outputing wet deposition rates
    type(monarch_state_array_ptr), dimension(NUM_RATE_ARRAYS) :: rate_array_ptr
  contains
    !> Calculate the wet deposition for one vertical column saving loss rates
    !! and performing vertical mixing on the column
    procedure :: calc_rates_and_mix
    !> Update the deposition rates for one grid cell
    procedure :: update_rxn_rates
    !> Determine the number of bytes required to pack the given object
    procedure :: pack_size
    !> Packs the given object in the buffer, advancing position
    procedure :: bin_pack
    !> Unpacks the given object from the buffer, advancing position
    procedure :: bin_unpack
    !> Print the wet deposition data
    procedure :: print => do_print
    !> Load wet deposition reaction properties from PartMC
    procedure, private :: load
    !> Set up output arrays
    procedure, private :: setup_output_arrays
    !> Initialize the wet deposition module
    procedure, private :: initialize
    !> Calculate loss rate for resolved clouds
    procedure, private :: calc_resolved_clouds
    !> Calculate loss rate for sub-grid-scale clouds and mixing factors
    procedure, private :: calc_sub_grid_scale_clouds
    !> Finalize the wet deposition object
    final :: finalize
  end type monarch_wet_deposition_t

  !> PartMC wet deposition constructor
  interface monarch_wet_deposition_t
    procedure :: constructor
  end interface monarch_wet_deposition_t

contains

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Create a new monarch_wet_deposition_t object
  !!
  !! Create a new monarch_wet_deposition_t object at the beginning of the model
  !! run for each process. This constructor should be called with an already-
  !! initialized phlex_core_t object. This should be called with the
  !! initialize flag set to true on the primary process, and false on the
  !! other processes.
  function constructor( phlex_core, diag, state_changer, config_file )       &
      result( new_obj )

    use monarch_diagnostics,                only : monarch_diagnostics_t
    use monarch_state_changer,              only : monarch_state_changer_t
    use monarch_util,                       only : monarch_dimensions_t,     &
                                                   assert_msg

    !> A new PartMC wet deposition object
    type(monarch_wet_deposition_t), pointer :: new_obj
    !> Phlex core
    type(phlex_core_t), target, intent(inout) :: phlex_core
    !> Diagnostics object attached to internal state with model dimensions
    !! Should only be included in calls from the primary process
    type(monarch_diagnostics_t), intent(in), optional :: diag
    !> Staet changer object with registered state arrays
    !! Should only be included in calls from the primary process
    type(monarch_state_changer_t), intent(inout), optional :: state_changer
    !> Path to the wet deposition configuration file
    !! Should only be included in calls from the primary process
    character(len=:), allocatable, optional, intent(in) :: config_file

    ! model domain dimensions
    type(monarch_dimensions_t) :: model_domain

    ! Create a new wet deposition object
    allocate( new_obj )

    ! Attach to the phlex_core_t object
    new_obj%phlex_core => phlex_core

    ! Find the depositing species properties on the primary process
    if( present( diag ) ) then

      call assert_msg( 537529094, present( state_changer ),                  &
                       "Missing state changer in wet deposition set up" )

      ! Save the number of vertical layers
      model_domain = diag%total_domain_dimensions( )
      new_obj%n_layers = model_domain%upper%vertical_idx( )

      call assert_msg( 197740772, present( config_file ),                    &
                       "Missing configuration file for wet deposition" )

      ! Load the interface data
      call new_obj%load( config_file )

      ! Initialize the wet deposition module on the primary process
      call new_obj%initialize( diag )

      ! Set up and needed output arrays
      call new_obj%setup_output_arrays( state_changer )

    end if

  end function constructor

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Run the wet deposition module for a given set of conditions for one
  !! vertical column.
  !!
  !! This subroutine calculated the loss rate for dissolved species and
  !! performs vertical mixing on the column
  subroutine calc_rates_and_mix( this, state_changer, column,                &
      last_time_step )

    use monarch_column,                     only : monarch_column_t
    use monarch_last_time_step,             only : monarch_last_time_step_t
    use monarch_state_changer,              only : monarch_state_changer_t

    !> Wet deposition object
    class(monarch_wet_deposition_t), intent(inout) :: this
    !> State changer object
    type(monarch_state_changer_t), intent(inout) :: state_changer
    !> State object for column to calculate wet deposition for
    type(monarch_column_t), intent(inout) :: column
    !> Values from the last time step
    type(monarch_last_time_step_t), intent(in) :: last_time_step

    ! Only recalculate on the precipitation timestep
    if( column%diag%do_precipitation_now( ) ) then
      call this%calc_resolved_clouds( column, last_time_step, state_changer )
      call this%calc_sub_grid_scale_clouds( column, last_time_step,          &
                                            state_changer )
    end if

  end subroutine calc_rates_and_mix

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Update wet deposition reaction rates for a single grid cell
  subroutine update_rxn_rates( this, indices )

    use monarch_state_changer,              only : monarch_state_changer_t
    use monarch_util,                       only : monarch_point_t,          &
                                                   assert_msg, to_string

    !> Wet deposition object
    class(monarch_wet_deposition_t), intent(inout) :: this
    !> Grid cell indices
    type(monarch_point_t), intent(in) :: indices

    integer(kind=mint) :: i_rate

    ! Do nothing if no rates are calculated
    if( .not. this%calculate_rates( RESOLVED_CLOUD_LOSS_RATE ) .and.         &
        .not. this%calculate_rates( SUB_GRID_SCALE_LOSS_RATE ) .and.         &
        .not. this%calculate_rates( VERTICAL_MIXING_RATE ) ) return

    ! Update the loss rate constant for aqueous species
    do i_rate = 1, this%n_aqueous_phase
#ifdef DEBUG
      call assert_msg( 933837023,                                            &
                       this%resolved_cloud_loss_rate(                        &
                                     indices%vertical_idx( ) ) .ge. 0.0 .and.&
                       this%resolved_cloud_loss_rate(                        &
                                     indices%vertical_idx( ) ) .lt. 1.e20,   &
                       "Bad resolved cloud loss rate for layer "//           &
                       to_string( indices%vertical_idx( ) )//": "//          &
                       to_string( this%resolved_cloud_loss_rate(             &
                                     indices%vertical_idx( ) ) ) )
      call assert_msg( 898343759,                                            &
                       this%sub_grid_scale_cloud_loss_rate(                  &
                                     indices%vertical_idx( ) ) .ge. 0.0 .and.&
                       this%sub_grid_scale_cloud_loss_rate(                  &
                                     indices%vertical_idx( ) ) .lt. 1.e20,   &
                       "Bad sub grid scale cloud loss rate for layer "//     &
                       to_string( indices%vertical_idx( ) )//": "//          &
                       to_string( this%sub_grid_scale_cloud_loss_rate(       &
                                     indices%vertical_idx( ) ) ) )
#endif
      call this%rate_update%set_rate( PMC_WET_DEP_PREFIX + i_rate,           &
           real( this%resolved_cloud_loss_rate( indices%WE_idx( ),           &
                                                indices%SN_idx( ),           &
                                                indices%vertical_idx( ) ) +  &
           this%sub_grid_scale_cloud_loss_rate( indices%WE_idx( ),           &
                                                indices%SN_idx( ),           &
                                                indices%vertical_idx( ) ),   &
           kind=mdbl ) )
      call this%phlex_core%update_rxn_data( this%rate_update )
    end do

  end subroutine update_rxn_rates

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Determine the size of a binary required to pack the PMC wet deposition
  !! module
  integer(kind=mint)function pack_size( this, comm )

    use monarch_mpi,                   only : mnrc_mpi_pack_size_integer,    &
                                       mnrc_mpi_pack_size_logical,           &
                                       mnrc_mpi_pack_size_string_t_array

    !> Wet deposition object
    class(monarch_wet_deposition_t), intent(in) :: this
    !> MPI communicator
    integer(kind=mint), intent(in) :: comm

    integer(kind=mint) :: i_flag

    pack_size =                                                              &
        mnrc_mpi_pack_size_integer( this%n_layers            , comm )      + &
        mnrc_mpi_pack_size_integer( this%n_aqueous_phase     , comm )      + &
        mnrc_mpi_pack_size_integer( this%n_mixing_species    , comm )      + &
        mnrc_mpi_pack_size_string_t_array( this%species_labels, comm )
    do i_flag = 1, NUM_RATE_ARRAYS
      pack_size = pack_size +                                                &
          mnrc_mpi_pack_size_logical( this%output_rates( i_flag ), comm )  + &
          mnrc_mpi_pack_size_logical( this%calculate_rates( i_flag ), comm )
    end do

  end function pack_size

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Pack the given object to the buffer, advancing position
  subroutine bin_pack( this, buffer, pos, comm )
#ifdef MONARCH_USE_MPI
    use monarch_mpi,                        only : mnrc_mpi_pack_integer,    &
                                            mnrc_mpi_pack_logical,           &
                                            mnrc_mpi_pack_string_t_array
    use monarch_util,                       only : assert
#endif

    !> Wet deposition object
    class(monarch_wet_deposition_t), intent(in) :: this
    !> Memory buffer
    character, intent(inout) :: buffer( : )
    !> Current buffer position
    integer(kind=mint), intent(inout) :: pos
    !> MPI communicator
    integer(kind=mint), intent(in) :: comm

#ifdef MONARCH_USE_MPI
    integer(kind=mint) :: prev_position, i_flag

    prev_position = pos
    call mnrc_mpi_pack_integer( buffer, pos, this%n_layers            , comm )
    call mnrc_mpi_pack_integer( buffer, pos, this%n_aqueous_phase     , comm )
    call mnrc_mpi_pack_integer( buffer, pos, this%n_mixing_species    , comm )
    call mnrc_mpi_pack_string_t_array( buffer, pos, this%species_labels,comm )
    do i_flag = 1, NUM_RATE_ARRAYS
      call mnrc_mpi_pack_logical( buffer, pos, this%output_rates( i_flag ),  &
                                  comm )
      call mnrc_mpi_pack_logical( buffer, pos,                               &
                                  this%calculate_rates( i_flag ), comm )
    end do
    call assert( 479241856, &
         pos - prev_position <= this%pack_size( comm ) )

#endif

  end subroutine bin_pack

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Unpack the given object to the buffer, advancing position
  subroutine bin_unpack( this, buffer, pos, comm, diag, state_changer )

    use monarch_diagnostics,                only : monarch_diagnostics_t
    use monarch_state_changer,              only : monarch_state_changer_t
#ifdef MONARCH_USE_MPI
    use monarch_mpi,                        only : mnrc_mpi_unpack_integer,  &
                                            mnrc_mpi_unpack_logical,         &
                                            mnrc_mpi_unpack_string_t_array
    use monarch_util,                       only : assert
#endif

    !> Wet deposition object
    class(monarch_wet_deposition_t), intent(inout) :: this
    !> Memory buffer
    character, intent(inout) :: buffer( : )
    !> Current buffer position
    integer(kind=mint), intent(inout) :: pos
    !> MPI communicator
    integer(kind=mint), intent(in) :: comm
    !> Diagnostics
    type(monarch_diagnostics_t), intent(in) :: diag
    !> State changer
    type(monarch_state_changer_t), intent(inout) :: state_changer

#ifdef MONARCH_USE_MPI
    integer(kind=mint) :: prev_position, i_flag

    prev_position = pos
    call mnrc_mpi_unpack_integer( buffer, pos, this%n_layers          , comm )
    call mnrc_mpi_unpack_integer( buffer, pos, this%n_aqueous_phase   , comm )
    call mnrc_mpi_unpack_integer( buffer, pos, this%n_mixing_species  , comm )
    call mnrc_mpi_unpack_string_t_array(buffer, pos,this%species_labels,comm )
    do i_flag = 1, NUM_RATE_ARRAYS
      call mnrc_mpi_unpack_logical( buffer, pos, this%output_rates( i_flag ),&
                                    comm )
      call mnrc_mpi_unpack_logical( buffer, pos,                             &
                                    this%calculate_rates( i_flag ), comm )
    end do
    call assert( 634058134, &
         pos - prev_position <= this%pack_size( comm ) )

    ! Initialize the wet deposition object on this process
    call this%initialize( diag )

    ! Setup any needed output arrays
    call this%setup_output_arrays( state_changer )
#endif

  end subroutine bin_unpack

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Print the wet deposition data
  subroutine do_print( this, file_unit )

    use monarch_util,                       only : to_string

    !> Wet deposition object
    class(monarch_wet_deposition_t), intent(in) :: this
    !> File unit to write to
    integer(kind=mint), intent(in) :: file_unit

    character(len=*), parameter :: fmt_blank_line = "('')"

    write(file_unit,fmt_blank_line)
    write(file_unit,*) "*****************************"
    write(file_unit,*) "** PMC Wet Deposition Data **"
    write(file_unit,*) "*****************************"
    write(file_unit,fmt_blank_line)
    if( this%calculate_rates( RESOLVED_CLOUD_LOSS_RATE ) .and.         &
        this%calculate_rates( SUB_GRID_SCALE_LOSS_RATE ) .and.         &
        this%calculate_rates( VERTICAL_MIXING_RATE ) )  then
    write(file_unit,*) "Number of vertical layers:                   ",      &
                       this%n_layers
    write(file_unit,*) "Number of aqueous phases subject to rainout: ",      &
                       this%n_aqueous_phase
    write(file_unit,*) "Number of species subject to vertical mixing:",      &
                       this%n_mixing_species
    if( this%output_rates( RESOLVED_CLOUD_LOSS_RATE ) )                      &
      write(file_unit,*) "Outputing resolved cloud loss rates [s-1]"
    if( this%output_rates( SUB_GRID_SCALE_LOSS_RATE ) )                      &
      write(file_unit,*) "Outputing sub-grid scale cloud loss rates [s-1]"
    if( this%output_rates( VERTICAL_MIXING_RATE ) )                          &
      write(file_unit,*) "Outputing relative change rates from vertical "//  &
                         "mixing"
    write(file_unit,fmt_blank_line)
    else
      write(file_unit,*) "No wet deposition rates calculated - default "//   &
                         "rates will be used"
    end if
    write(file_unit,*) "*********************************"
    write(file_unit,*) "** End PMC Wet Deposition Data **"
    write(file_unit,*) "*********************************"
    write(file_unit,fmt_blank_line)

  end subroutine do_print

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Finalize the wet deposition object
  subroutine finalize( this )

    use monarch_CONVCLD_ACM,                          only : CONVCLD_FINALIZE

    !> Wet deposition object
    type(monarch_wet_deposition_t), intent(inout) :: this

    ! The PartMC <-> MONARCH interface owns the phlex_core object
    if( associated( this%phlex_core                    ) ) &
        this%phlex_core => null( )
    if( allocated( this%resolved_cloud_loss_rate       ) ) &
        deallocate( this%resolved_cloud_loss_rate        )
    if( allocated( this%sub_grid_scale_cloud_loss_rate ) ) &
        deallocate( this%sub_grid_scale_cloud_loss_rate  )
    if( allocated( this%species_labels                 ) ) &
        deallocate( this%species_labels                  )
    if( associated( this%property_set                  ) ) &
        deallocate( this%property_set                    )

    call CONVCLD_FINALIZE( )

  end subroutine finalize

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Load wet deposition reaction properties from PartMC
  !!
  !! The set deposition species are found in the PartMC mechanism and numbered
  !! in order of being found to create the set. Required properties are loaded
  !! from PartMC input data for chemical species.
  subroutine load( this, config_file )

#ifdef PMC_USE_JSON
    use json_module,                        only : json_core, json_file,     &
                                                   json_value, json_ck
#endif
    use monarch_util,                       only : assert_msg, die_msg
    use pmc_chem_spec_data,                 only : chem_spec_data_t
    use pmc_mechanism_data,                 only : mechanism_data_t
    use pmc_rxn_data,                       only : rxn_data_t
    use pmc_rxn_wet_deposition,             only : rxn_wet_deposition_t
    use pmc_util,                           only : pmc_string_t => string_t

    !> Wet deposition object
    class(monarch_wet_deposition_t), intent(inout) :: this
    !> Path to the PMC <-> wet deposition interface configuration file
    character(len=:), allocatable, intent(in) :: config_file
#ifdef PMC_USE_JSON
    type(json_core), pointer :: json
    type(json_file) :: j_file
    type(json_value), pointer :: j_obj, j_next
    character(kind=json_ck, len=:), allocatable :: key, unicode_str_val
    integer(kind=mint) :: i_rxn, i_spec, n_rxn
    type(mechanism_data_t), pointer :: mechanism
    class(rxn_data_t), pointer :: rxn
    character(len=:), allocatable :: type_key, prop_key, str_val, owner_name
    real(kind=mdbl) :: real_val
    type(chem_spec_data_t), pointer :: chem_spec_data
    type(pmc_string_t), allocatable :: spec_names( : )
    logical :: found, logical_val
    type(property_t), pointer :: sub_prop

    ! Initialize the property set
    this%property_set => property_t( )

    ! Get a new json core
    allocate( json )

    ! Initialize the json objects
    j_obj  => null( )
    j_next => null( )

    ! Initialize the json file
    call j_file%initialize( )
    call j_file%get_core( json )
    call assert_msg( 345125683, allocated( config_file ),                    &
              "Received non-allocated string for file path" )
    call assert_msg( 957386126, trim( config_file ) .ne."",                  &
              "Received empty string for file path" )
    inquire( file=config_file, exist=found )
    call assert_msg( 399605469, found, "Cannot find file: "//                &
              config_file )
    call j_file%load_file( filename = config_file )

    ! Find the interface data
    call j_file%get( 'wet-dep-data', j_obj )

    ! Load the data to the property set
    call assert_msg( 551026628, associated( j_obj ),                         &
                     "Missing 'wet-dep-data' element in dry deposition "//   &
                     "configuration data file" )
    owner_name = "wet deposition interface"
    call this%property_set%load( json, j_obj, .true., owner_name )

    ! Get the resolved-cloud properties
    prop_key = "resoved clouds"
    call assert_msg( 806668039,                                              &
                     this%property_set%get_property_t( prop_key, sub_prop ), &
                     "Missing 'resolved clouds' properties in wet "//        &
                     "deposition configuration file" )
    prop_key = "calculate"
    call assert_msg( 280856701,                                              &
                     sub_prop%get_logical( prop_key, logical_val ),          &
                     "Missing 'calculate' flag in 'resolved clouds' "//      &
                     "in wet deposition configuration file" )
    this%calculate_rates( RESOLVED_CLOUD_LOSS_RATE ) = logical_val
    prop_key = "output rates"
    if( sub_prop%get_logical( prop_key, logical_val ) ) then
      this%output_rates( RESOLVED_CLOUD_LOSS_RATE ) = logical_val
    end if

    ! Get the sub-grid scale cloud properties
    prop_key = "sub-grid scale clouds"
    call assert_msg( 990999453,                                              &
                     this%property_set%get_property_t( prop_key, sub_prop ), &
                     "Missing 'sub-grid scale clouds' properties in wet "//  &
                     "deposition configuration file" )
    prop_key = "calculate"
    call assert_msg( 362946089,                                              &
                     sub_prop%get_logical( prop_key, logical_val ),          &
                     "Missing 'calculate' flag in 'sub-grid scal clouds' "// &
                     "in wet deposition configuration file" )
    this%calculate_rates( SUB_GRID_SCALE_LOSS_RATE ) = logical_val
    prop_key = "output rates"
    if( sub_prop%get_logical( prop_key, logical_val ) ) then
      this%output_rates( SUB_GRID_SCALE_LOSS_RATE ) = logical_val
    end if

    ! Get the vertical mixing properties
    prop_key = "vertical mixing"
    call assert_msg( 300295518,                                              &
                     this%property_set%get_property_t( prop_key, sub_prop ), &
                     "Missing 'vertical mixing' properties in wet "//        &
                     "deposition configuration file" )
    prop_key = "calculate"
    call assert_msg( 972300054,                                              &
                     sub_prop%get_logical( prop_key, logical_val ),          &
                     "Missing 'calculate' flag in 'vertical mixing' "//      &
                     "in wet deposition configuration file" )
    this%calculate_rates( VERTICAL_MIXING_RATE ) = logical_val
    prop_key = "output rates"
    if( sub_prop%get_logical( prop_key, logical_val ) ) then
      this%output_rates( VERTICAL_MIXING_RATE ) = logical_val
    end if

    ! Clean up the json objects
    call j_file%destroy( )
    call json%destroy( )
    deallocate( json )

    ! Verify the phlex core is initialize
    call assert_msg( 132780840, this%phlex_core%is_initialized( ),           &
                    "Trying to load properties for "//                       &
                    "monarch_wet_deposition_t with an "//                    &
                    "uninitialized phlex_core_t object" )

    ! FIXME This assumes MONARCH only uses one PMC mechanism at a time
    mechanism => this%phlex_core%mechanism( 1 )%val

    ! Save the number of mixing species
    ! FIXME Consider including filter for non-mixing species
    this%n_mixing_species = this%phlex_core%state_size( )

    ! Save the species names if outputting mixing rates
    if( this%output_rates( VERTICAL_MIXING_RATE ) ) then
      call assert_msg( 810291938,                                            &
                       this%phlex_core%get_chem_spec_data( chem_spec_data ), &
                       "Error getting chemical species data from" )
      spec_names = chem_spec_data%get_spec_names( )
      allocate( this%species_labels( size( spec_names ) ) )
      do i_spec = 1, size( this%species_labels )
        this%species_labels( i_spec )%string = spec_names( i_spec )%string
      end do
      deallocate( spec_names )
      chem_spec_data => null( )
    end if

    ! Count and get the names of species involved in wet deposition reactions
    n_rxn = 0
    type_key = "rxn id"
    do i_rxn = 1, mechanism%size( )

      ! Find all wet deposition reactions
      rxn => mechanism%get_rxn( i_rxn )
      select type( rxn_wet_dep => rxn )
        class is( rxn_wet_deposition_t )

          ! Find rain loss reactions
          if( rxn%property_set%get_string( type_key, str_val ) ) then
            if( str_val.eq."rain loss" ) then

              n_rxn = n_rxn + 1

              ! Save the wet deposition id for rate updates
              call rxn_wet_dep%set_rxn_id( PMC_WET_DEP_PREFIX + n_rxn )

            end if
          end if
      end select
    end do

    this%n_aqueous_phase = n_rxn
#else
    call die_msg( 120062295, "No JSON support for wet deposition "//         &
                             "configuration" )
#endif
  end subroutine load

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Initialize the wet deposition module
  subroutine initialize( this, diag )

    use monarch_CONVCLD_ACM,                only : CONVCLD_INIT
    use monarch_diagnostics,                only : monarch_diagnostics_t
    use monarch_util,                       only : assert_msg,               &
                                                   monarch_dimensions_t
    use pmc_rxn_factory,                    only : rxn_factory_t

    !> Wet deposition object
    class(monarch_wet_deposition_t), intent(inout) :: this
    !> Diagnostics object
    type(monarch_diagnostics_t), intent(in) :: diag

    type(rxn_factory_t) :: rxn_factory
    type(monarch_dimensions_t) :: task_dims

    ! Verify the phlex core is initialized
    call assert_msg( 750500267, this%phlex_core%is_initialized( ),           &
                    "Trying to initialize monarch_wet_deposition_t with "//  &
                    "an uninitialized phlex_core_t object." )

    ! Initialize the rate update nugget
    call rxn_factory%initialize_update_data( this%rate_update )

    ! Initialize the convective cloud arrays
    call CONVCLD_INIT( this%n_layers, this%n_mixing_species )

    ! Get the task integration domain
    task_dims = diag%task_integration_dimensions( )

    ! Set up the working arrays for this process
    allocate( this%column_conc( this%n_mixing_species, this%n_layers ) )
    allocate( this%resolved_cloud_loss_rate(                                 &
                   task_dims%lower%WE_idx( ) : task_dims%upper%WE_idx( ),    &
                   task_dims%lower%SN_idx( ) : task_dims%upper%SN_idx( ),    &
                   1 : this%n_layers ) )
    allocate( this%sub_grid_scale_cloud_loss_rate(                           &
                   task_dims%lower%WE_idx( ) : task_dims%upper%WE_idx( ),    &
                   task_dims%lower%SN_idx( ) : task_dims%upper%SN_idx( ),    &
                   1 : this%n_layers ) )

    this%resolved_cloud_loss_rate( :,:,: ) = real( 0.0, kind=mdbl )
    this%sub_grid_scale_cloud_loss_rate( :,:,: ) = real( 0.0, kind=mflt )

  end subroutine initialize

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Set up output arrays
  subroutine setup_output_arrays( this, state_changer )

    use monarch_state_changer,         only : monarch_state_changer_t

    !> Wet deposition object
    class(monarch_wet_deposition_t), intent(inout) :: this
    !> State changer
    type(monarch_state_changer_t), intent(inout) :: state_changer

    ! Register an array to hold resolved-cloud loss rates
    if( this%output_rates( RESOLVED_CLOUD_LOSS_RATE ) ) then
      this%rate_array_ptr( RESOLVED_CLOUD_LOSS_RATE ) =                      &
          state_changer%register_per_cell_real(                              &
              "resolved_cloud_loss_rate",                                    & !<-- unique name for the array
              .true. )                                                         !<-- output to history file
      call this%rate_array_ptr( RESOLVED_CLOUD_LOSS_RATE                     &
                                )%set_netcdf_long_name(                      &
                                "Wet deposition resolved-cloud loss rate" )
      call this%rate_array_ptr( RESOLVED_CLOUD_LOSS_RATE )%set_netcdf_units( &
                                "s-1" )
    end if

    ! Register an array to hold sub-grid scale cloud loss rates
    if( this%output_rates( SUB_GRID_SCALE_LOSS_RATE ) ) then
      this%rate_array_ptr( SUB_GRID_SCALE_LOSS_RATE ) =                      &
          state_changer%register_per_cell_real(                              &
              "sub_grid_cloud_loss_rate",                                    & !<-- unique name for the array
              .true. )                                                         !<-- output to history file
      call this%rate_array_ptr( SUB_GRID_SCALE_LOSS_RATE                     &
                             )%set_netcdf_long_name(                         &
                             "Wet deposition sub-grid scale cloud loss rate" )
      call this%rate_array_ptr( SUB_GRID_SCALE_LOSS_RATE )%set_netcdf_units( &
                             "s-1" )
    end if

    ! Register an array to hold relative change rates from vertical mixing
    if( this%output_rates( VERTICAL_MIXING_RATE ) ) then
      this%rate_array_ptr( VERTICAL_MIXING_RATE ) =                          &
          state_changer%register_per_cell_real(                              &
              "vertical_mixing_rate",                                        & !<-- unique name for the array
              .true.,                                                        & !<-- output to history file
              size( this%species_labels ),                                   & !<-- number of array elements
              this%species_labels )                                            !<-- label for each element
      call this%rate_array_ptr( RESOLVED_CLOUD_LOSS_RATE                     &
                                )%set_netcdf_long_name(                      &
                                "Wet deposition resolved-cloud loss rate" )
      call this%rate_array_ptr( RESOLVED_CLOUD_LOSS_RATE )%set_netcdf_units( &
                                "s-1" )
    end if

  end subroutine setup_output_arrays

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Calculate loss for resolved clouds
  subroutine calc_resolved_clouds( this, column, last_time_step,             &
      state_changer )

    use monarch_column,                     only : monarch_column_t
    use monarch_last_time_step,             only : monarch_last_time_step_t
    use monarch_state_changer,              only : monarch_state_changer_t
    use monarch_util,                       only : monarch_point_t,          &
                                                   assert_msg, to_string

    !> Wet deposition object
    class(monarch_wet_deposition_t), intent(inout) :: this
    !> State of column to calculated loss rates for
    type(monarch_column_t), intent(in) :: column
    !> Values from the last time step
    type(monarch_last_time_step_t), intent(in) :: last_time_step
    !> State changer
    type(monarch_state_changer_t), intent(inout) :: state_changer

    type(monarch_point_t) :: cell
    integer(kind=mint) :: i_layer

    ! TODO move to constants module
    real(kind=mdbl), parameter :: water_density = 1000.0 ! H2O density [kg m-3]

    real(kind=mdbl) :: non_convec_rain    ! non-convective rainfall for this timestep [m]
    real(kind=mdbl) :: cloud_lifetime     ! cloud lifetime [s]
    real(kind=mdbl) :: rainfall_rate      ! rainfall rate [m s-1]
    real(kind=mdbl) :: column_cloud_water ! total column cloud water [kg m-2]
    real(kind=mdbl) :: column_rain        ! total column rain water [kg m-2]
    real(kind=mdbl) :: column_graupel     ! total column graupel [kg m-2]
    real(kind=mdbl) :: column_snow        ! total column snow [kg m-2]
    real(kind=mdbl) :: column_ice         ! total column ice [kg m-2]
    real(kind=mdbl) :: column_total_water ! total column water [kg m-3]
    real(kind=mdbl) :: layer_air_density  ! air density in a layer [kg m-3]
    real(kind=mdbl) :: layer_height       ! layer height [m]
    real(kind=mdbl) :: layer_cloud_water  ! cloud water in a layer [kg m-3]
    real(kind=mdbl) :: layer_rain         ! rain water in a layer [kg m-3]
    real(kind=mdbl) :: layer_graupel      ! graupel in a layer [kg m-3]
    real(kind=mdbl) :: layer_snow         ! snow in a layer [kg m-3]
    real(kind=mdbl) :: layer_ice          ! ice in a layer [kg m-3]
    real(kind=mdbl) :: layer_total_water  ! total water in a layer [kg m-3]
    real(kind=mdbl) :: precip_flux        ! precipitation flux [m s-1]
    real(kind=mdbl) :: cloud_washout_time ! cloud washout time( s-1 )

    ! If not calculating resoved cloud loss, do nothing
    if( .not. this%calculate_rates( RESOLVED_CLOUD_LOSS_RATE ) ) return

    ! Set up the grid indices object
    cell = column%get_indices( )

    ! Reset the resolved-cloud loss rates
    this%resolved_cloud_loss_rate( cell%WE_idx( ), cell%SN_idx( ), : ) =     &
        real( 0.0, kind=mdbl )

    ! Get column-wide properties
    ! FIXME the cloud lifetime appears to be based on the model timesteps
    !       and not any properties of the cloud. Check why
    cloud_lifetime     = column%diag%physics_time_step__s( )
    non_convec_rain    = column%non_convective_precip__m( last_time_step )
    rainfall_rate      = non_convec_rain / cloud_lifetime
    column_cloud_water = column%cloud_water__kgm2( )
    column_rain        = column%rain__kgm2( )
    column_graupel     = column%graupel__kgm2( )
    column_snow        = column%snow__kgm2( )
    column_ice         = column%ice__kgm2( )
    column_total_water = column_cloud_water + column_rain + column_graupel   &
                         + column_snow + column_ice

    ! If there is no cloud, rain or graupel water, there is no loss
    if( column_cloud_water + column_rain + column_graupel .gt. 0.0 ) then

      ! Loop through the vertical layers
      do i_layer = 1, column%n_layers
        call cell%set_vertical_idx( i_layer )

        ! Get layer properties
        layer_air_density = column%diag%air_density__kgm3( cell )
        layer_height      = column%diag%mixed_layer_height__m( cell )

        ! Get the water concentrations in this layer
        layer_cloud_water = column%diag%cloud_water_mixing_ratio__kgkg( cell)&
                            * layer_air_density
        layer_rain        = column%diag%rain_mixing_ratio__kgkg( cell )      &
                            * layer_air_density
        layer_graupel     = column%diag%graupel_mixing_ratio__kgkg( cell )   &
                            * layer_air_density
        layer_snow        = column%diag%snow_mixing_ratio__kgkg( cell )      &
                            * layer_air_density
        layer_ice         = column%diag%cloud_ice_mixing_ratio__kgkg( cell ) &
                            * layer_air_density
        layer_total_water = layer_cloud_water + layer_rain + layer_graupel + &
                            layer_snow + layer_ice

        ! FIXME what is this filter for?
        if( layer_cloud_water + layer_rain + layer_graupel .gt.              &
             5.0e-5 * layer_air_density ) then

          ! Calculate the precipitation flux [m s-1]
          ! FIXME what is this condition for?
          if( column_snow + column_graupel .gt. 0.0 ) then
            precip_flux = rainfall_rate * layer_height *                     &
                          ( layer_rain + layer_snow + layer_graupel ) /      &
                          ( column_rain + column_snow + column_graupel )
          else
            precip_flux = rainfall_rate * layer_height *                     &
                          layer_total_water / column_total_water
          end if

          precip_flux = max( precip_flux, real( 1.0e-20, kind=mdbl ) )

          ! Calculate cloud washout time [s]
          cloud_washout_time =                                               &
                          layer_total_water *                                & ! <-- [kg m-3]
                          layer_height      /                                & ! <-- [m]
                          ( water_density   *                                & ! <-- [kg m-3]
                            precip_flux ) ! <-- [m s-1]
          cloud_washout_time = max( cloud_washout_time, cloud_lifetime )

          ! Calculate removal rate [s-1]
          this%resolved_cloud_loss_rate( cell%WE_idx( ),                     &
                                         cell%SN_idx( ),                     &
                                         i_layer ) = 1.0 / cloud_washout_time

#ifdef DEBUG
          call assert_msg( 923334006,                                        &
                     this%resolved_cloud_loss_rate( i_layer ) .ge. 0.0 .and. &
                     this%resolved_cloud_loss_rate( i_layer ) .lt. 1.e20,    &
                     "Bad resolved cloud loss rate in layer "//              &
                     to_string( i_layer )//": "//                            &
                     to_string( this%resolved_cloud_loss_rate(               &
                                cell%WE_idx( ), cell%SN_idx( ), i_layer ) ) )
#endif
        end if
      end do
    end if

    ! Output the removal rate if flagged to do so
    if( this%output_rates( RESOLVED_CLOUD_LOSS_RATE ) ) then
      do i_layer = 1, column%n_layers
        call cell%set_vertical_idx( i_layer )
        call state_changer%save_property(                                    &
            this%rate_array_ptr( RESOLVED_CLOUD_LOSS_RATE ),                 & !<-- state array pointer
            cell,                                                            & !<-- grid-cell indices
            real( this%resolved_cloud_loss_rate(                             &
                   cell%WE_idx( ), cell%SN_idx( ), i_layer ), kind=mflt ) )    !<-- loss rate [s-1]
      end do
    end if

  end subroutine calc_resolved_clouds

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Calculate sub-grid-scale cloud loss and perform mixing
  subroutine calc_sub_grid_scale_clouds( this, column, last_time_step,       &
      state_changer )

    use monarch_CONVCLD_ACM,                only : CONVCLD_ACM_MOD_PMC
    use monarch_column,                     only : monarch_column_t
    use monarch_last_time_step,             only : monarch_last_time_step_t
    use monarch_state_changer,              only : monarch_state_changer_t
    use monarch_util,                       only : monarch_point_t,          &
                                                   assert_msg, to_string
    use pmc_phlex_state,                    only : phlex_state_t

    !> Wet deposition object
    class(monarch_wet_deposition_t), intent(inout) :: this
    !> State of column to calculated loss rates for
    type(monarch_column_t), intent(inout) :: column
    !> Values from the last time step
    type(monarch_last_time_step_t), intent(in) :: last_time_step
    !> State changer
    type(monarch_state_changer_t), intent(inout) :: state_changer

    integer(kind=mint) :: i_layer, i_spec
    type(monarch_point_t) :: cell
    type(phlex_state_t), pointer :: phlex_state

    ! If not calculating sub-gridscale cloud loss and mixing, do nothing
    if( .not. this%calculate_rates( RESOLVED_CLOUD_LOSS_RATE ) .and.         &
        .not. this%calculate_rates( VERTICAL_MIXING_RATE ) ) return

    cell = column%get_indices()

    ! Update local column concentrations
    do i_layer = 1, column%n_layers
      call cell%set_vertical_idx( i_layer )
      phlex_state => column%get_phlex_state( cell )
#ifdef DEBUG
      do i_spec = 1, size( phlex_state%state_var )
        call assert_msg( 803183193,                                          &
                         phlex_state%state_var( i_spec ) .ge. 0.0 .and.      &
                         phlex_state%state_var( i_spec ) .le. 1.e20,         &
                         "Bad concentration in layer "//to_string( i_layer ) &
                         //" for species "//to_string( i_spec )//": "//      &
                         to_string( phlex_state%state_var( i_spec ) ) )
      end do
#endif
      this%column_conc( :, i_layer ) = phlex_state%state_var( : )
    end do

    ! Calculate convective cloud rain-out rate and do vertical mixing
    call CONVCLD_ACM_MOD_PMC( column, last_time_step, this%column_conc,      &
                              this%sub_grid_scale_cloud_loss_rate(           &
                                        cell%WE_idx( ), cell%SN_idx( ), : ) )

    ! Output the removal rate if flagged to do so
    if( this%output_rates( SUB_GRID_SCALE_LOSS_RATE ) ) then
      do i_layer = 1, column%n_layers
        call cell%set_vertical_idx( i_layer )
        call state_changer%save_property(                                    &
            this%rate_array_ptr( SUB_GRID_SCALE_LOSS_RATE ),                 & !<-- state array pointer
            cell,                                                            & !<-- grid-cell indices
            real( this%sub_grid_scale_cloud_loss_rate(                       &
                     cell%WE_idx( ), cell%SN_idx( ), i_layer ), kind=mflt ) )  !<-- loss rate [s-1]
      end do
    end if

    ! If doing vertical mixing but not sub-grid scale loss, reset the rates
    if( .not. this%calculate_rates( SUB_GRID_SCALE_LOSS_RATE ) ) then
      this%sub_grid_scale_cloud_loss_rate( cell%WE_idx( ),                   &
                                           cell%SN_idx( ),                   &
                                           : ) = real( 0.0, kind=mflt )
    end if

    ! Update state with mixed concentrations for the column
    if( this%calculate_rates( VERTICAL_MIXING_RATE ) ) then
      do i_layer = 1, column%n_layers
        call cell%set_vertical_idx( i_layer )
        phlex_state => column%get_phlex_state( cell )

        ! Output relative change rates from vertical mixing if flagged to do so
        if( this%output_rates( VERTICAL_MIXING_RATE ) ) then
          do i_spec = 1, size( phlex_state%state_var )
             call state_changer%save_property(                                  &
                this%rate_array_ptr( VERTICAL_MIXING_RATE ),                   & !<-- state array pointer
                cell,                                                          & !<-- grid cell indices
                i_spec,                                                        & !<-- species index
                real( 2.0 * ( - phlex_state%state_var( i_spec )              + &
                              this%column_conc( i_spec, i_layer ) )          / &
                      ( phlex_state%state_var( i_spec )                      + &
                        this%column_conc( i_spec, i_layer ) ), kind=mflt ) )
          end do
        end if

        ! Update the phlex-state
        phlex_state%state_var( : ) = this%column_conc( :, i_layer )
#ifdef DEBUG
        do i_spec = 1, size( phlex_state%state_var )
          call assert_msg( 803183193,                                          &
                           phlex_state%state_var( i_spec ) .ge. 0.0 .and.      &
                           phlex_state%state_var( i_spec ) .le. 1.e20,         &
                           "Bad concentration in layer "//to_string( i_layer ) &
                         //" for species "//to_string( i_spec )//": "//      &
                         to_string( phlex_state%state_var( i_spec ) ) )
        end do
#endif
      end do
    end if

  end subroutine calc_sub_grid_scale_clouds

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
#endif // DO_PARTMC
end module monarch_wet_deposition
