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

!> Handles output of diagnostics
module monarch_diagnostics_writer
! MONARCH always has MPI
#define MONARCH_USE_MPI

  use monarch_column,                       only : monarch_column_t
  use monarch_constants,                    only : mint, mflt, mdbl
  use monarch_output_diagnostics_pointer,   only :                           &
                                            monarch_output_diagnostics_ptr

  implicit none
  private

  public :: monarch_diagnostics_writer_t

  !> Diagnostics writer
  !!
  !! Handles calculation, storage and output of diagnostics
  type :: monarch_diagnostics_writer_t
  private
    !> Column object for use during calculations
    type(monarch_column_t), pointer :: column => null( )
    !> Set of output diagnostics
    class(monarch_output_diagnostics_ptr), allocatable :: members( : )
  contains
    !> Load diagnostic output input data
#ifdef DO_PARTMC
    generic :: load => load_no_partmc,                                       &
                       load_partmc
#else
    generic :: load => load_no_partmc
#endif
    !> Load diagnostic output input data when not using PartMC for chemistry
    procedure :: load_no_partmc
#ifdef DO_PARTMC
    !> Load diagnostic output input data when using PartMC for chemistry
    procedure :: load_partmc
#endif
    !> Load member of the output diagnostics
    procedure, private :: load_members
    !> Register space in the model state for output diagnostics
    procedure :: register
    !> Update members after all MONARCH indices have been initialized
    procedure :: update_diagnostics
    !> Calculate diagnostics for the current time step
    procedure :: calculate
    !> Output diagnostics
    procedure :: output
    !> Determine the number of bytes required to pack the given object
    procedure :: pack_size
    !> Pack the given object in the buffer, advancing position
    procedure :: bin_pack
    !> Unpack the given object from the buffer, advancing position
    procedure :: bin_unpack
    !> Print out the diagnostic writer configuration
    procedure :: print => do_print
  end type monarch_diagnostics_writer_t

  !> Constructor for the diagnostics writer
  interface  monarch_diagnostics_writer_t
    procedure :: constructor
  end interface monarch_diagnostics_writer_t

contains

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

  !> Constructor for the diagnostics writer
  function constructor( ) result( new_obj )

    !> New diagnostics writer object
    type(monarch_diagnostics_writer_t), pointer :: new_obj

    allocate( new_obj )

  end function constructor

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

  !> Load a set the set of diagnostics to calculate and output
  !!
  !! This is a constructor for use when not using PartMC for chemistry
  !! (This only needs to be called on the integration tasks)
  subroutine load_no_partmc( this, diag, aero_rep )

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

    !> Diagnostics writer
    class(monarch_diagnostics_writer_t), intent(inout) :: this
    !> Diagnostics object attached to solver internal state
    type(monarch_diagnostics_t), intent(in) :: diag
    !> Aerosol representation
    class(monarch_aerosol_representation_t), intent(in) :: aero_rep

    call assert_msg( 231675689, .not. ( diag%use_partmc_for_chemistry( )     &
                                        .and. diag%is_chemistry_on( ) ),     &
                     "Called wrong diagnostics writer constructor." )

    ! Set up a column object
    this%column => monarch_column_t( diag )

    ! Load the members
    call this%load_members( diag, aero_rep )

  end subroutine load_no_partmc

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
#ifdef DO_PARTMC
  !> Load a set the set of diagnostics to calcualte and output
  !! (This only needs to be called on the integration tasks)
  subroutine load_partmc( this, phlex_core, species_map, diag, aero_rep )

    use monarch_species_map,                only : monarch_species_map_t
    use pmc_phlex_core,                     only : phlex_core_t
    use monarch_aerosol_representation,     only :                           &
                                            monarch_aerosol_representation_t
    use monarch_diagnostics,                only : monarch_diagnostics_t
    use monarch_util,                       only : assert_msg

    !> Diagnostics writer
    class(monarch_diagnostics_writer_t), intent(inout) :: this
    !> Phlex-core
    type(phlex_core_t), intent(in) :: phlex_core
    !> PartMC-Phlex <-> MONARCH species map
    type(monarch_species_map_t), intent(in) :: species_map
    !> Diagnostics object attached to solver internal state
    type(monarch_diagnostics_t), intent(in) :: diag
    !> Aerosol representation
    class(monarch_aerosol_representation_t), intent(in) :: aero_rep

    call assert_msg( 402284888, diag%use_partmc_for_chemistry( )             &
                                .and. diag%is_chemistry_on( ),               &
                     "Called wrong diagnostics writer constructor." )

    ! Set up a column object
    this%column => monarch_column_t( phlex_core, species_map, diag )

    ! Load the members
    call this%load_members( diag, aero_rep )

  end subroutine load_partmc
#endif
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Load members of the output diagnostics
  subroutine load_members( this, diag, aero_rep )

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

    !> Diagnostics writer
    class(monarch_diagnostics_writer_t), intent(inout) :: this
    !> Diagnostics object attached to solver internal state
    type(monarch_diagnostics_t), intent(in) :: diag
    !> Aerosol representation
    class(monarch_aerosol_representation_t), intent(in) :: aero_rep

    ! Without PartMC-Phlex loaded, diagnostics must be hard coded because no
    ! json support is available
    ! TODO load members based on a configuration file

    allocate( this%members( 2 ) )
    call this%members( 1 )%new_by_name(     "aerosol optics", diag, aero_rep )
    call this%members( 2 )%new_by_name("physical properties", diag, aero_rep )
    ! call this%members( 3 )%new_by_name(          "debugging", diag, aero_rep )

  end subroutine load_members

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

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

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

    !> Diagnostics writer
    class(monarch_diagnostics_writer_t), intent(inout) :: this
    !> Diagnostics object attached to solver internal state
    type(monarch_diagnostics_t), target, intent(in) :: diag
    !> Model state changer
    type(monarch_state_changer_t), intent(inout) :: state_changer

    integer(kind=mint) :: i_member

    do i_member = 1, size( this%members )
      call this%members( i_member )%val%register( diag, state_changer )
    end do

  end subroutine register

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

  !> Update diagnostics after all MONARCH indices have been initialized
  subroutine update_diagnostics( this, diag )

    use monarch_diagnostics,                only : monarch_diagnostics_t

    !> Diagnostics writer
    class(monarch_diagnostics_writer_t), intent(inout) :: this
    !> Diagnostics object attached to solver internal state
    type(monarch_diagnostics_t), target, intent(in) :: diag

    integer(kind=mint) :: i_member

    do i_member = 1, size( this%members )
      call this%members( i_member )%val%update( diag )
    end do

  end subroutine update_diagnostics

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

  !> Calculate the diagnostics
  subroutine calculate( this, diag, state_changer )

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

    !> Diagnostics writer
    class(monarch_diagnostics_writer_t), intent(inout) :: this
    !> Diagnostics object attached to solver internal state
    type(monarch_diagnostics_t), target, intent(in) :: diag
    !> Model state changer
    type(monarch_state_changer_t), intent(inout) :: state_changer

    type(monarch_dimensions_t) :: task_dims
    type(monarch_point_t) :: cell
    integer(kind=mint) :: SN_idx, WE_idx, vert_idx, i_member

    ! Get the dimensions of the grid to calculate on this process
    task_dims = diag%task_integration_dimensions( )

    ! Create a cell indices object
    cell      = diag%new_point( )

    ! Loop through the grid cells
    ! (West -> East)
    do WE_idx = task_dims%lower%WE_idx( ), task_dims%upper%WE_idx( )
      call cell%set_WE_idx( WE_idx )

      ! (South -> North)
      do SN_idx = task_dims%lower%SN_idx( ), task_dims%upper%SN_idx( )
        call cell%set_SN_idx( SN_idx )

        ! Update the column state
        call this%column%update( diag, cell )

        ! Do column-level diagnostics calculations
        do i_member = 1, size( this%members )
          call this%members( i_member )%val%calculate_column( this%column,   &
                                                             state_changer )
        end do

        ! (Surface -> Top Layer)
        do vert_idx = task_dims%lower%vertical_idx( ),                       &
                      task_dims%upper%vertical_idx( )
          call cell%set_vertical_idx( vert_idx )

          ! Do grid cell-level diagnostic calculations
          do i_member = 1, size( this%members )
            call this%members( i_member )%val%calculate_grid_cell( cell,     &
                                                             diag,           &
                                                             this%column,    &
                                                             state_changer )
          end do

        end do
      end do
    end do

  end subroutine calculate

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

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

    use module_WRITE_INTERNAL_STATE,             only : write_internal_state

    !> Diagnostics writer
    class(monarch_diagnostics_writer_t), intent(inout) :: this
    !> Write internal state
    !! TODO Consider building a wrapper for this variable or allowing the
    !! diagnostics object to attach to it
    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_member

    do i_member = 1, size( this%members )
      call this%members( i_member )%val%output( 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

    !> Diagnostics writer
    class(monarch_diagnostics_writer_t), intent(in) :: this
    !> MPI communicator
    integer(kind=mint), intent(in) :: comm

    integer(kind=mint) :: i_member

    pack_size = mnrc_mpi_pack_size_integer( 1, comm )
    do i_member = 1, size( this%members )
      pack_size = pack_size + this%members( i_member )%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

    !> Diagnostics writer
    class(monarch_diagnostics_writer_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_member

    prev_position = pos
    call mnrc_mpi_pack_integer( buffer, pos, size( this%members ), comm )
    do i_member = 1, size( this%members )
      call this%members( i_member )%bin_pack( buffer, pos, comm )
    end do
    call assert( 111018163, &
         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, column )
    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

    !> Diagnostics writer
    class(monarch_diagnostics_writer_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
    !> Column object used for calculating diagnostics
    !! (Can be a null pointer for write-only tasks)
    type(monarch_column_t), pointer, intent(inout) :: column

#ifdef MONARCH_USE_MPI
    integer(kind=mint) :: prev_position, i_member, n_member, member_type

    this%column => column
    if( allocated( this%members ) ) deallocate( this%members )

    prev_position = pos
    call mnrc_mpi_unpack_integer( buffer, pos, n_member, comm )
    allocate( this%members( n_member ) )
    do i_member = 1, n_member
      call this%members( i_member )%bin_unpack( buffer, pos, comm, aero_rep )
    end do
    call assert( 763322356, &
         pos - prev_position <= this%pack_size( comm ) )
#endif

  end subroutine bin_unpack

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

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

    !> Diagnostics writer
    class(monarch_diagnostics_writer_t), intent(in) :: this
    !> File unit
    integer(kind=mint), intent(in) :: file_unit

    integer(kind=mint) :: i_member
    character(len=*), parameter :: fmt_blank_line = "('')"

    write(file_unit,fmt_blank_line)
    write(file_unit,*) "************************************************"
    write(file_unit,*) "*** MONARCH Diagnostics Writer Configuration ***"
    write(file_unit,*) "************************************************"
    do i_member = 1, size( this%members )
      write(file_unit,fmt_blank_line)
      write(file_unit,"('   *** Output Diagnostic ', I3, ' ***')")  i_member
      call this%members( i_member )%val%print( file_unit )
    end do
    write(file_unit,fmt_blank_line)
    write(file_unit,*) "****************************************************"
    write(file_unit,*) "*** End MONARCH Diagnostics Writer Configuration ***"
    write(file_unit,*) "****************************************************"

  end subroutine do_print

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

end module monarch_diagnostics_writer
