! 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_output_diag_phys_t object and related functions

!> Handles output of physical property diagnostics
module monarch_output_diag_phys
! MONARCH always has MPI
#define MONARCH_USE_MPI

  use monarch_aerosol_representation_pointer, only :                         &
                                         monarch_aerosol_group_collection_ptr
  use monarch_constants,                      only : mint, mflt, mdbl
  use monarch_state_changer,                  only : monarch_state_array_ptr
  use monarch_output_diagnostics,             only :                         &
                                         monarch_output_diagnostics_t

  implicit none
  private

  public :: monarch_output_diag_phys_t

  !> Output physical properties
  type, extends(monarch_output_diagnostics_t) ::                             &
      monarch_output_diag_phys_t
  private
    !> Set of pointers to state arrays holding the diagnostics
    type(monarch_state_array_ptr), allocatable :: array_ptr( : )
  contains
    !> Register space in the model state for the diagnostics
    procedure :: register
    !> Update output diagnostic once all MONARCH ids have been initialized
    procedure :: update
    !> Perform column-level calculations
    procedure :: calculate_column
    !> Perform grid cell-level calculations
    procedure :: calculate_grid_cell
    !> Determine the number of bytes required to pack the given object
    procedure :: pack_size
    !> Pack the given object onto the buffer, advancing position
    procedure :: bin_pack
    !> Unpack the given object from the buffer, advancing position
    procedure :: bin_unpack
    !> Output the diagnostics
    procedure :: output
    !> Print out the configuration of the output diagnostic
    procedure :: print => do_print
  end type monarch_output_diag_phys_t

  !> Constructor for the physical property diagnostics
  interface monarch_output_diag_phys_t
    procedure :: constructor
  end interface monarch_output_diag_phys_t

contains
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Constructor for the physical property diagnostics
  function constructor( diag, aero_rep ) result( new_obj )

    use monarch_aerosol_representation,     only :                           &
                                            monarch_aerosol_representation_t
    use monarch_diagnostics,                only : monarch_diagnostics_t

    !> New output diagnostics object
    type(monarch_output_diag_phys_t), pointer :: new_obj
    !> Diagnostics object
    !! (Do not include when unpacking object from an mpi buffer)
    type(monarch_diagnostics_t), optional, intent(in) :: diag
    !> Aerosol representation
    !! (Do not include when unpacking object from an mpi buffer)
    class(monarch_aerosol_representation_t), optional, intent(in) :: aero_rep

    allocate( new_obj )

  end function constructor

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

  !> Register space in the model state for the diagnostics
  subroutine register( this, diag, state_changer )

    use monarch_diagnostics,                only : monarch_diagnostics_t
    use monarch_state_changer,              only : monarch_state_changer_t
    use monarch_util,                       only : string_t

    !> Output diagnostics
    class(monarch_output_diag_phys_t), intent(inout) :: this
    !> Diagnostics object
    type(monarch_diagnostics_t), intent(in) :: diag
    !> Model state changer
    type(monarch_state_changer_t), intent(inout) :: state_changer

    type(string_t), allocatable :: bin_labels( : )
    integer(kind=mint) :: i_bin

    ! Set up arrays for physical property outputs
    allocate( this%array_ptr( 6 ) )

    ! Register an array to hold vertical layer heights
    this%array_ptr( 1 ) = state_changer%register_per_cell_real(              &
                                              "layer_thickness",             & !<-- unique name for array
                                              .true. )                         !<-- output to history file
    call this%array_ptr( 1 )%set_netcdf_long_name( "Layer thickness" )
    call this%array_ptr( 1 )%set_netcdf_units( "m" )

    ! Register an array to hold mid-layer hydrostatic pressures
    this%array_ptr( 2 ) = state_changer%register_per_cell_real(              &
                                              "mid_layer_pressure",          & !<-- unique name for array
                                              .true. )                         !<-- output to history file
    call this%array_ptr( 2 )%set_netcdf_long_name(                           &
                                            "Mid-layer hydrostatic pressure" )
    call this%array_ptr( 2 )%set_netcdf_units( "Pa" )
    call this%array_ptr( 2 )%set_netcdf_standard_name( "air_pressure" )

    ! Register an array to hold the pressure at vertical interfaces
    this%array_ptr( 3 ) = state_changer%register_per_vertical_interface_real(&
                                              "interface_pressure",          & !<-- unique name for array
                                              .true. )                         !<-- output to history file
    call this%array_ptr( 3 )%set_netcdf_long_name(                           &
                                          "Interfacial hydrostatic pressure" )
    call this%array_ptr( 3 )%set_netcdf_units( "Pa" )

    ! Register an array to hold the relative humidity
    this%array_ptr( 4 ) = state_changer%register_per_cell_real(              &
                                              "relative_humidity",           & !<-- unique name for array
                                              .true. )                         !<-- output to history file
    call this%array_ptr( 4 )%set_netcdf_long_name( "Relative humidity" )
    call this%array_ptr( 4 )%set_netcdf_standard_name( "relative_humidity" )
    call this%array_ptr( 4 )%set_netcdf_units( "decimal fraction" )

    ! Register an array to hold the relative humidity at 2 m
    this%array_ptr( 5 ) = state_changer%register_per_column_real(            &
                                              "relative_humidity_2m",        & !<-- unique name for array
                                              .true. )                         !<-- output to history file
    call this%array_ptr( 5 )%set_netcdf_long_name(                           &
                                          "Relative humidity at 2 m" )
    call this%array_ptr( 5 )%set_netcdf_units( "decimal fraction" )

    ! Register an array to hold the mid-layer height above sea level
    this%array_ptr( 6 ) = state_changer%register_per_cell_real(              &
                                              "mid_layer_height",            & !<-- unique name for array
                                              .true. )                         !<-- output to history file
    call this%array_ptr( 6 )%set_netcdf_long_name(                           &
                                         "Mid-layer height above sea level" )
    call this%array_ptr( 6 )%set_netcdf_standard_name( "height" )
    call this%array_ptr( 6 )%set_netcdf_units( "m" )


  end subroutine register

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

  !> Update output diagnostics once all MONARCH ids have been initialized
  subroutine update( this, diag )

    use monarch_diagnostics,                only : monarch_diagnostics_t

    !> Output diagnostics
    class(monarch_output_diag_phys_t), intent(inout) :: this
    !> Diagnostics object
    type(monarch_diagnostics_t), intent(in) :: diag

    ! Nothing to do for this diagnostic

  end subroutine update

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

  !> Perform column-level calculations
  subroutine calculate_column( this, column, state_changer )

    use monarch_column,                     only : monarch_column_t
    use monarch_state_changer,              only : monarch_state_changer_t
    use monarch_util,                       only : monarch_point_t

    !> Output diagnostics
    class(monarch_output_diag_phys_t), intent(inout) :: this
    !> Column state
    type(monarch_column_t), intent(inout) :: column
    !> Model state changer
    type(monarch_state_changer_t), intent(inout) :: state_changer

    type(monarch_point_t) :: cell

    cell = column%indices

    call cell%set_vertical_idx( 1 )
    call cell%set_to_lower_interface( )
    do while( cell%vertical_idx( ) .le. column%n_layers )
      ! Save the lower interface hydrostatic pressure [Pa]
      call state_changer%save_property( this%array_ptr( 3 ),                 & !<-- interface pressure state array
          cell,                                                              & !<-- grid cell indices
          real( column%btw_layer_pressure__Pa( cell%vertical_idx( ) ),       &
                kind=mflt ) )                                                  !<-- interface pressure [Pa]
      call cell%set_vertical_idx( cell%vertical_idx( ) + 1 )
    end do

    ! Save the upper interface hydrostatic pressure of the highest cell [Pa]
    call cell%set_to_upper_interface( )
    call cell%set_vertical_idx( column%n_layers )
    call state_changer%save_property( this%array_ptr( 3 ),                   & !<-- interface pressure state array
           cell,                                                             & !<-- grid cell indices
           real( column%btw_layer_pressure__Pa( cell%vertical_idx( ) + 1 ),  &
                 kind=mflt ) )                                                 !<-- interface pressure [Pa]

    ! Save the relative humidity at 2-m
    call state_changer%save_property( this%array_ptr( 5 ),                   & !<-- RH at 2m state array
           cell,                                                             & !<-- grid cell indices
           real( column%relative_humidity_2m__dec( ), kind=mflt ) )            !<-- relative humidity @ 2m [dec]

  end subroutine calculate_column

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

  !> Perform grid cell-level calculations
  subroutine calculate_grid_cell( this, cell, diag, column, state_changer )

    use monarch_column,                     only : monarch_column_t
    use monarch_diagnostics,                only : monarch_diagnostics_t
    use monarch_state_changer,              only : monarch_state_changer_t
    use monarch_util,                       only : monarch_point_t

    !> Output diagnostics
    class(monarch_output_diag_phys_t), intent(inout) :: this
    !> Cell indices
    type(monarch_point_t), intent(in) :: cell
    !> Diagnostics object
    type(monarch_diagnostics_t), intent(in) :: diag
    !> Column object
    type(monarch_column_t), intent(in) :: column
    !> Model state changer
    type(monarch_state_changer_t), intent(inout) :: state_changer

    integer(kind=mint) :: i_group

    ! Save the layer thickness [m]
    call state_changer%save_property( this%array_ptr( 1 ),                   & !<-- layer height state array
                cell,                                                        & !<-- grid cell indices
                real( column%layer_thickness__m( cell%vertical_idx( ) ),     &
                      kind=mflt ) )                                            !<-- layer thickness [m]

    ! Save the mid-layer hydrostatic pressure [Pa]
    call state_changer%save_property( this%array_ptr( 2 ),                   & !<-- mid-layer pressure state array
                cell,                                                        & !<-- grid cell indices
                real( column%mid_layer_pressure__Pa( cell%vertical_idx( ) ), &
                      kind=mflt ) )                                            !<-- mid-layer-pressure [Pa]

    ! Save the relative humidity [decimal fraction]
    call state_changer%save_property( this%array_ptr( 4 ),                   & !<-- relative humidity state array
                cell,                                                        & !<-- grid cell indices
                real( column%relative_humidity__dec( cell%vertical_idx( ) ), &
                      kind=mflt ) )                                            !<-- relative humidity [decimal fraction]

    ! Save the mid-layer height above sea level [m]
    call state_changer%save_property( this%array_ptr( 6 ),                   & !<-- mid-layer height state array
                cell,                                                        & !<-- grid cell indices
                real( column%mid_layer_height__m( cell%vertical_idx( ) ),    &
                      kind=mflt ) )                                            !<-- mid-layer height above sea level [m]

  end subroutine calculate_grid_cell

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

  !> Output the diagnostics
  subroutine output( this, write_int_state, array_name, first_output )

    use module_WRITE_INTERNAL_STATE,        only : write_internal_state

    !> Output diagnostics
    class(monarch_output_diag_phys_t), intent(inout) :: this
    !> Write internal state
    type(write_internal_state), intent(in) :: write_int_state
    !> Name of array that is available for output
    character(len=*), intent(in) :: array_name
    !> Flag indicating whether the state arrays need to be defined in the
    !! NetCDF output file
    logical, intent(in) :: first_output

    integer(kind=mint) :: i_array

    ! See if the array is one that this output diagnostic registered
    do i_array = 1, size( this%array_ptr )
      call this%array_ptr( i_array )%output_to_netcdf( write_int_state,      &
                                                       array_name,           &
                                                       first_output )
    end do

  end subroutine output

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

  !> Determine the number of bytes required to pack the given object
  integer(kind=mint) function pack_size( this, comm )

    use monarch_mpi,               only : mnrc_mpi_pack_size_integer

    !> Output diagnostics
    class(monarch_output_diag_phys_t), intent(in) :: this
    !> MPI communicator
    integer(kind=mint), intent(in) :: comm

    integer(kind=mint) :: i_group, i_array

    pack_size = mnrc_mpi_pack_size_integer( 1, comm )
    do i_array = 1, size( this%array_ptr )
      pack_size = pack_size + this%array_ptr( i_array )%pack_size( 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
    use monarch_util,              only : assert
#endif

    !> Output diagnostics
    class(monarch_output_diag_phys_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_group, i_array

    prev_position = pos
    call mnrc_mpi_pack_integer( buffer, pos, size( this%array_ptr   ), comm )
    do i_array = 1, size( this%array_ptr )
      call this%array_ptr( i_array )%bin_pack( buffer, pos, comm )
    end do
    call assert( 279189350, &
         pos - prev_position <= this%pack_size( comm ) )
#endif

  end subroutine bin_pack

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

  !> Unpack the given object from the buffer, advancing position
  subroutine bin_unpack( this, buffer, pos, comm, aero_rep )
    use monarch_aerosol_representation,     only :                           &
                                            monarch_aerosol_representation_t
#ifdef MONARCH_USE_MPI
    use monarch_mpi,               only : mnrc_mpi_unpack_integer
    use monarch_util,              only : assert
#endif

    !> Output diagnostics
    class(monarch_output_diag_phys_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
    !> Aerosol representation
    class(monarch_aerosol_representation_t), target, intent(in) :: aero_rep

#ifdef MONARCH_USE_MPI
    integer(kind=mint) :: prev_position, i_group, n_groups, i_array, n_arrays

    call assert( 903883941, .not. allocated( this%array_ptr   ) )

    prev_position = pos
    call mnrc_mpi_unpack_integer( buffer, pos, n_arrays, comm )
    allocate( this%array_ptr(   n_arrays ) )
    do i_array = 1, n_arrays
      call this%array_ptr( i_array )%bin_unpack( buffer, pos, comm )
    end do
    call assert( 333669136, &
         pos - prev_position <= this%pack_size( comm ) )
#endif

  end subroutine bin_unpack

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

  !> Print out the configuration of the output diagnostic
  subroutine do_print( this, file_unit )

    !> Output diagnostics
    class(monarch_output_diag_phys_t), intent(in) :: this
    !> File unit
    integer(kind=mint), intent(in) :: file_unit

    integer(kind=mint) :: i_group

    write(file_unit,*) "  *** Physical property diagnostics ***"
    write(file_unit,*) "  Layer thickness [m] saved to:                 "//  &
                       "layer_thickness"
    write(file_unit,*) "  Mid-layer hydrostatic pressure [Pa] saved to: "//  &
                       "mid_layer_pressure"
    write(file_unit,*) "  Interface hydrostatic pressure [Pa] saved to: "//  &
                       "interface_pressure"
    write(file_unit,*) "  Relative humidity [decimal fraction] saved to: "// &
                       "relative_humidity"
    write(file_unit,*) "  Relative humidity at 2m [decimal fraction] "//     &
                       "saved to: relative_humidity_2m"
    write(file_unit,*) "  Mid-layer height above sea level [m] saved to: "// &
                       "mid_layer_height"

  end subroutine do_print

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

end module monarch_output_diag_phys
