!>
!> A module to store adios2 variables used in diagnosis.F90 (current_diagnosis,
!> opendiag*, copy_history/data1d_* subroutines
!>
#ifdef ADIOS2
module adios2_history_diagnostics
  use adios2
  use adios2_common
  implicit none

  public

  type t_history_diagnosis
    type(adios2_variable) :: numsteps, nspecies, mpdiag, nfield, &
                             modes, mfdiag, tstep
    type(adios2_variable) :: partdata, fieldtime, fieldmode
  end type t_history_diagnosis

  type (t_history_diagnosis) :: t_hist_diag

  contains


  !> -------------------------------------------------------------------------
  subroutine define_vars(t_hist_diag, io_obj, mpdiag, nfield, mfdiag, modes)
    use adios2
    use precision
    implicit none

    type(t_history_diagnosis),  intent(inout) :: t_hist_diag
    type(adios2_io),            intent(inout) :: io_obj
    integer,                    intent(in)    :: mpdiag,nfield,mfdiag,modes
    integer                                   :: ierr

    call adios2_define_variable (t_hist_diag%numsteps, io_obj, 'num_steps',   adios2_type_integer4, ierr)
    call adios2_define_variable (t_hist_diag%nspecies, io_obj, 'nspecies',    adios2_type_integer4, ierr)
    call adios2_define_variable (t_hist_diag%mpdiag,   io_obj, 'mpdiag',      adios2_type_integer4, ierr)
    call adios2_define_variable (t_hist_diag%nfield,   io_obj, 'nfield',      adios2_type_integer4, ierr)
    call adios2_define_variable (t_hist_diag%modes,    io_obj, 'modes',       adios2_type_integer4, ierr)
    call adios2_define_variable (t_hist_diag%mfdiag,   io_obj, 'mfdiag',      adios2_type_integer4, ierr)
    call adios2_define_variable (t_hist_diag%tstep,    io_obj, 'tstep*ndiag', adios2_type_real_lk,  ierr)

    call adios2_define_variable (t_hist_diag%partdata,  io_obj, 'partdata', adios2_type_real_lk, 2, &
      1_8 * (/mpdiag,nspecies/), (/0_8,0_8/), 1_8 * (/mpdiag,nspecies/), .true., ierr)
    call adios2_define_variable (t_hist_diag%fieldtime, io_obj, 'fieldtime', adios2_type_real_lk, 2, &
      1_8 * (/mfdiag,nfield/), (/0_8,0_8/), 1_8 * (/mfdiag,nfield/), .true., ierr)
    call adios2_define_variable (t_hist_diag%fieldmode, io_obj, 'fieldmode', adios2_type_real_lk, 3, &
      1_8 * (/2,modes,nfield/), (/0_8,0_8,0_8/), 1_8 * (/2,modes,nfield/), .true., ierr)

  end subroutine define_vars



  !> -------------------------------------------------------------------------
  !> Inquire all variables
  !> -------------------------------------------------------------------------
  subroutine inq_vars (t_hist_diag, io_obj)
    use adios2
    implicit none

    type(t_history_diagnosis), intent(inout) :: t_hist_diag
    type(adios2_io),           intent(inout) :: io_obj
    integer                                  :: ierr
    
    call adios2_inquire_variable (t_hist_diag%numsteps,  io_obj, "num_steps",   ierr)
    call adios2_inquire_variable (t_hist_diag%nspecies,  io_obj, "nspecies",    ierr)
    call adios2_inquire_variable (t_hist_diag%mpdiag,    io_obj, "mpdiag",      ierr)
    call adios2_inquire_variable (t_hist_diag%nfield,    io_obj, "nfield",      ierr)
    call adios2_inquire_variable (t_hist_diag%modes,     io_obj, "modes",       ierr)
    call adios2_inquire_variable (t_hist_diag%mfdiag,    io_obj, "mfdiag",      ierr)
    call adios2_inquire_variable (t_hist_diag%tstep,     io_obj, "tstep*ndiag", ierr)
    call adios2_inquire_variable (t_hist_diag%partdata,  io_obj, "partdata",    ierr)
    call adios2_inquire_variable (t_hist_diag%fieldtime, io_obj, "fieldtime",   ierr)
    call adios2_inquire_variable (t_hist_diag%fieldmode, io_obj, "fieldmode",   ierr)
  end subroutine inq_vars


  !> -------------------------------------------------------------------------
  !> Copy data between the history and restart file or vice versa
  !> Both files should already be open, and the io object should be empty
  !> -------------------------------------------------------------------------
  subroutine copy_data (src_e, dest_e, io_obj, t_hist, nsteps_in_file)
    use mpi
    use adios2
    use simple_error_handler
    implicit none

    type(adios2_engine),       intent(inout) :: src_e, dest_e
    type(adios2_io),           intent(inout) :: io_obj
    type(t_history_diagnosis), intent(inout) :: t_hist
    integer,                   intent(inout) :: nsteps_in_file

    integer                                  :: nspecies, mpdiag, nfield, modes, mfdiag
    real(lk)                                 :: tstep_ndiag
    real(lk), allocatable                    :: partdata_in(:,:), fieldtime_in(:,:), fieldmode_in(:,:,:)

    integer                                  :: i, ndims, ierr
    integer (kind=8)                         :: nsteps_in_file8
    integer (kind=8), allocatable            :: shape_dims(:)

    ! Inq vars from the source file
    call inq_vars (t_hist, io_obj)

    ! Get the no. of steps in the source file
    ! Just query partdata to get the no. of timesteps
    call adios2_variable_steps (nsteps_in_file8, t_hist%partdata, ierr)
    nsteps_in_file = nsteps_in_file8

    ! Loop over nsteps
    do i=1, nsteps_in_file
      call adios2_begin_step (src_e,  adios2_step_mode_read,   ierr)
      call adios2_begin_step (dest_e, adios2_step_mode_append, ierr)

      ! Read scalars and allocate arrays in the first step
      if (i .eq. 1) then
        ! Read scalars
        call adios2_get (src_e, t_hist%nspecies, nspecies,    adios2_mode_sync, ierr)
        call adios2_get (src_e, t_hist%mpdiag,   mpdiag,      adios2_mode_sync, ierr)
        call adios2_get (src_e, t_hist%nfield,   nfield,      adios2_mode_sync, ierr)
        call adios2_get (src_e, t_hist%modes,    modes,       adios2_mode_sync, ierr)
        call adios2_get (src_e, t_hist%mfdiag,   mfdiag,      adios2_mode_sync, ierr)
        call adios2_get (src_e, t_hist%tstep,    tstep_ndiag, adios2_mode_sync, ierr)

        ! Write scalars
        call adios2_put (dest_e, t_hist%nspecies, nspecies,    ierr)
        call adios2_put (dest_e, t_hist%mpdiag,   mpdiag,      ierr)
        call adios2_put (dest_e, t_hist%nfield,   nfield,      ierr)
        call adios2_put (dest_e, t_hist%modes,    modes,       ierr)
        call adios2_put (dest_e, t_hist%mfdiag,   mfdiag,      ierr)
        call adios2_put (dest_e, t_hist%tstep,    tstep_ndiag, ierr)

        ! Get the shape of the arrays and allocate them
        call adios2_variable_shape (shape_dims, ndims, t_hist%partdata, ierr)
        allocate(partdata_in(shape_dims(1), shape_dims(2)), stat=ierr)
        call error_check (ierr, "allocation failure")
        deallocate(shape_dims)

        call adios2_variable_shape (shape_dims, ndims, t_hist%fieldtime, ierr)
        allocate(fieldtime_in(shape_dims(1), shape_dims(2)), stat=ierr)
        call error_check (ierr, "allocation failure")
        deallocate(shape_dims)

        call adios2_variable_shape (shape_dims, ndims, t_hist%fieldmode, ierr)
        allocate(fieldmode_in(shape_dims(1), shape_dims(2), shape_dims(3)), stat=ierr)
        call error_check (ierr, "allocation failure")
        deallocate(shape_dims)

      endif
      
      ! Read arrays
      call adios2_get (src_e, t_hist%partdata,  partdata_in,  ierr)
      call adios2_get (src_e, t_hist%fieldtime, fieldtime_in, ierr)
      call adios2_get (src_e, t_hist%fieldmode, fieldmode_in, ierr)
      
      ! End step in the src file to complete the read operations
      call adios2_end_step (src_e, ierr)

      ! Write arrays
      call adios2_put (dest_e, t_hist%partdata,  partdata_in,  ierr)
      call adios2_put (dest_e, t_hist%fieldtime, fieldtime_in, ierr)
      call adios2_put (dest_e, t_hist%fieldmode, fieldmode_in, ierr)

      ! End step for the dest file
      call adios2_end_step (dest_e, ierr)
    enddo

    ! Free them variables
    if (allocated(partdata_in))  deallocate (partdata_in)
    if (allocated(fieldtime_in)) deallocate (fieldtime_in)
    if (allocated(fieldmode_in)) deallocate (fieldmode_in)

  end subroutine copy_data

end module adios2_history_diagnostics
#endif

