!>
!> A module to store adios2 variables used in diagnosis.F90 (current_diagnosis,
!> opendiag*, copy_history/data1d_* subroutines
!>
!> Assumptions: ndiag etc. do not change between runs
!>              all arrays have the same no. of timesteps
!>
#ifdef ADIOS2
module adios2_data1d_diagnostics
  use adios2
  use adios2_common
  implicit none

  public

  type t_data1d_diagnosis
    type(adios2_variable) :: numsteps, nspecies, mpsi, nfield, nhybrid, &
                             mpdata1d, mfdata1d
    type(adios2_variable) :: data1di, field00, fieldrms
    type(adios2_variable) :: data1de, data1df, data1dfe
  end type t_data1d_diagnosis

  type (t_data1d_diagnosis) :: t_data1d_diag

  contains


  !> -------------------------------------------------------------------------
  subroutine define_vars (t_data1d_diag, io_obj, nspecies, mpsi, nfield, nhybrid, &
                          mpdata1d, mfdata1d, fload, feload)
    use adios2
    use precision
    use particle_array, only: data1di, data1de, data1df, data1dfe
    implicit none

    type(t_data1d_diagnosis),   intent(inout) :: t_data1d_diag
    type(adios2_io),            intent(inout) :: io_obj
    integer,                    intent(in)    :: nspecies, mpsi, nfield, nhybrid
    integer,                    intent(in)    :: mpdata1d, mfdata1d, fload, feload
    integer                                   :: ierr

    call adios2_define_variable (t_data1d_diag%nspecies, io_obj, 'nspecies', adios2_type_integer4, ierr)
    call adios2_define_variable (t_data1d_diag%mpsi,     io_obj, 'mpsi+1',   adios2_type_integer4, ierr)
    call adios2_define_variable (t_data1d_diag%nfield,   io_obj, 'nfield',   adios2_type_integer4, ierr)
    call adios2_define_variable (t_data1d_diag%nhybrid,  io_obj, 'nhybrid',  adios2_type_integer4, ierr)
    call adios2_define_variable (t_data1d_diag%mpdata1d, io_obj, 'mpdata1d', adios2_type_integer4, ierr)
    call adios2_define_variable (t_data1d_diag%mfdata1d, io_obj, 'mfdata1d', adios2_type_integer4, ierr)

    call adios2_define_variable (t_data1d_diag%data1di, io_obj, 'data1di', adios2_type_real_lk, 2, &
      1_8 * shape(data1di), (/0_8,0_8/), 1_8 * shape(data1di), .true., ierr)
    call adios2_define_variable (t_data1d_diag%field00, io_obj, 'field00', adios2_type_real_lk, 2, &
      1_8 * (/mpsi+1,nfield/), (/0_8,0_8/), 1_8 * (/mpsi+1,nfield/), .true., ierr)
    call adios2_define_variable (t_data1d_diag%fieldrms, io_obj, 'fieldrms', adios2_type_real_lk, 2, &
      1_8 * (/mpsi+1,nfield/), (/0_8,0_8/), 1_8 * (/mpsi+1,nfield/), .true., ierr)
    
    if (nspecies>1) then
      if (nhybrid>0) then
        call adios2_define_variable (t_data1d_diag%data1de, io_obj, 'data1de', adios2_type_real_lk, 2, &
          1_8 * shape(data1de), (/0_8,0_8/), 1_8 * shape(data1de), .true., ierr)
      endif
      if (fload>0) then
        call adios2_define_variable (t_data1d_diag%data1df, io_obj, 'data1df', adios2_type_real_lk, 2, &
          1_8 * shape(data1df), (/0_8,0_8/), 1_8 * shape(data1df), .true., ierr)
      endif
      if(feload>0) then
        call adios2_define_variable (t_data1d_diag%data1dfe, io_obj, 'data1dfe', adios2_type_real_lk, 2, &
          1_8 * shape(data1dfe), (/0_8,0_8/), 1_8 * shape(data1dfe), .true., ierr)
      endif
    endif

  end subroutine define_vars



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

    type(t_data1d_diagnosis),  intent(inout) :: t_data1d_diag
    type(adios2_io),           intent(inout) :: io_obj
    integer                                  :: ierr
    
    call adios2_inquire_variable (t_data1d_diag%nspecies, io_obj, "nspecies", ierr)
    call adios2_inquire_variable (t_data1d_diag%mpsi,     io_obj, "mpsi+1",   ierr)
    call adios2_inquire_variable (t_data1d_diag%nfield,   io_obj, "nfield",   ierr)
    call adios2_inquire_variable (t_data1d_diag%nhybrid,  io_obj, "nhybrid",  ierr)
    call adios2_inquire_variable (t_data1d_diag%mpdata1d, io_obj, "mpdata1d", ierr)
    call adios2_inquire_variable (t_data1d_diag%mfdata1d, io_obj, "mfdata1d", ierr)
    call adios2_inquire_variable (t_data1d_diag%data1di,  io_obj, "data1di",  ierr)
    call adios2_inquire_variable (t_data1d_diag%field00,  io_obj, "field00",  ierr)
    call adios2_inquire_variable (t_data1d_diag%fieldrms, io_obj, "fieldrms", ierr)
    call adios2_inquire_variable (t_data1d_diag%data1de,  io_obj, "data1de",  ierr)
    call adios2_inquire_variable (t_data1d_diag%data1df,  io_obj, "data1df",  ierr)
    call adios2_inquire_variable (t_data1d_diag%data1dfe, io_obj, "data1dfe", ierr)
  end subroutine inq_vars


  !> -------------------------------------------------------------------------
  !> Copy data between the data1d 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_data1d)
    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_data1d_diagnosis),  intent(inout) :: t_data1d

    integer                                  :: nspecies, mpsi, nfield, nhybrid, mpdata1d, mfdata1d
    real(lk), allocatable                    :: data1di_in(:,:), field00_in(:,:), fieldrms_in(:,:)
    real(lk), allocatable                    :: data1de_in(:,:), data1df_in(:,:), data1dfe_in(:,:)

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

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

    ! Get the no. of steps in the source file
    ! Just query one of the arrays to get the no. of timesteps
    call adios2_variable_steps (nsteps_in_file8, t_data1d%data1di, 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_data1d%nspecies, nspecies, adios2_mode_sync, ierr)
        call adios2_get (src_e, t_data1d%mpsi,     mpsi,     adios2_mode_sync, ierr)
        call adios2_get (src_e, t_data1d%nfield,   nfield,   adios2_mode_sync, ierr)
        call adios2_get (src_e, t_data1d%nhybrid,  nhybrid,  adios2_mode_sync, ierr)
        call adios2_get (src_e, t_data1d%mpdata1d, mpdata1d, adios2_mode_sync, ierr)
        call adios2_get (src_e, t_data1d%mfdata1d, mfdata1d, adios2_mode_sync, ierr)

        ! Write scalars
        call adios2_put (dest_e, t_data1d%nspecies, nspecies, ierr)
        call adios2_put (dest_e, t_data1d%mpsi,     mpsi,     ierr)
        call adios2_put (dest_e, t_data1d%nfield,   nfield,   ierr)
        call adios2_put (dest_e, t_data1d%nhybrid,  nhybrid,  ierr)
        call adios2_put (dest_e, t_data1d%mpdata1d, mpdata1d, ierr)
        call adios2_put (dest_e, t_data1d%mfdata1d, mfdata1d, ierr)

        ! Get the shape of the arrays and allocate them
        call adios2_variable_shape (shape_dims, ndims, t_data1d%data1di, ierr)
        allocate(data1di_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_data1d%field00, ierr)
        allocate(field00_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_data1d%fieldrms, ierr)
        allocate(fieldrms_in(shape_dims(1), shape_dims(2)), stat=ierr)
        call error_check (ierr, "allocation failure")
        deallocate(shape_dims)

        if (t_data1d%data1de%valid) then
          call adios2_variable_shape (shape_dims, ndims, t_data1d%data1de, ierr)
          allocate(data1de_in(shape_dims(1), shape_dims(2)), stat=ierr)
          call error_check (ierr, "allocation failure")
          deallocate(shape_dims)
        endif

        if (t_data1d%data1df%valid) then
          call adios2_variable_shape (shape_dims, ndims, t_data1d%data1df, ierr)
          allocate(data1df_in(shape_dims(1), shape_dims(2)), stat=ierr)
          call error_check (ierr, "allocation failure")
          deallocate(shape_dims)
        endif

        if (t_data1d%data1dfe%valid) then
          call adios2_variable_shape (shape_dims, ndims, t_data1d%data1dfe, ierr)
          allocate(data1dfe_in(shape_dims(1), shape_dims(2)), stat=ierr)
          call error_check (ierr, "allocation failure")
          deallocate(shape_dims)
        endif

      endif ! first time step
      
      ! Read arrays
      call adios2_get (src_e, t_data1d%data1di,  data1di_in,  ierr)
      call adios2_get (src_e, t_data1d%field00,  field00_in,  ierr)
      call adios2_get (src_e, t_data1d%fieldrms, fieldrms_in, ierr)

      if (t_data1d%data1de%valid)  call adios2_get (src_e, t_data1d%data1de,  data1de_in,  ierr)
      if (t_data1d%data1df%valid)  call adios2_get (src_e, t_data1d%data1df,  data1df_in,  ierr)
      if (t_data1d%data1dfe%valid) call adios2_get (src_e, t_data1d%data1dfe, data1dfe_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_data1d%data1di,  data1di_in,  ierr)
      call adios2_put (dest_e, t_data1d%field00,  field00_in,  ierr)
      call adios2_put (dest_e, t_data1d%fieldrms, fieldrms_in, ierr)

      if (t_data1d%data1de%valid)  call adios2_put (dest_e, t_data1d%data1de,  data1de_in,  ierr)
      if (t_data1d%data1df%valid)  call adios2_put (dest_e, t_data1d%data1df,  data1df_in,  ierr)
      if (t_data1d%data1dfe%valid) call adios2_put (dest_e, t_data1d%data1dfe, data1dfe_in, ierr)

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

    ! Free them variables
    deallocate(data1di_in)
    deallocate(field00_in)
    deallocate(fieldrms_in)

    if (allocated(data1de_in))  deallocate(data1de_in)
    if (allocated(data1df_in))  deallocate(data1df_in)
    if (allocated(data1dfe_in)) deallocate(data1dfe_in)

  end subroutine copy_data

end module adios2_data1d_diagnostics
#endif

