module adios2_common
  use mpi
  use adios2
  use simple_error_handler
  use global_parameters
  implicit none

  ! Main adios2 handle
  type(adios2_adios)  :: adios2_obj

  ! IO objects
  type(adios2_io)     :: adios2_eqplot_io
  type(adios2_io)     :: adios2_restart_io
  type(adios2_io)     :: adios2_restart_io_read
  type(adios2_io)     :: adios2_hist_restart_io
  type(adios2_io)     :: adios2_diag_hist_io
  type(adios2_io)     :: adios2_diag_data1d_io
  type(adios2_io)     :: adios2_data1d_restart_io
  type(adios2_io)     :: adios2_snapshot_io
  type(adios2_io)     :: adios2_data3d_io
  type(adios2_io)     :: adios2_phi_io
  type(adios2_io)     :: adios2_ptrack_io
  type(adios2_io)     :: adios2_poincare_io
  type(adios2_io)     :: adios2_history_io
  
  ! Unique IDs for all IO objects
  character(len=64)   :: adios2_eqplot_io_id
  character(len=64)   :: adios2_restart_io_id
  character(len=64)   :: adios2_restart_io_read_id
  character(len=64)   :: adios2_snapshot_io_id
  character(len=64)   :: adios2_diag_hist_io_id
  character(len=64)   :: adios2_hist_restart_io_id
  character(len=64)   :: adios2_diag_data1d_io_id
  character(len=64)   :: adios2_data1d_restart_io_id
  character(len=64)   :: adios2_data3d_id
  character(len=64)   :: adios2_phi_io_id
  character(len=64)   :: adios2_ptrack_io_id
  character(len=64)   :: adios2_poincare_io_id
  character(len=64)   :: adios2_history_io_id

  ! Bool representing if the IO object is initialized
  logical             :: adios2_eqplot_io_initialized
  logical             :: adios2_restart_io_initialized
  logical             :: adios2_snapshot_io_initialized
  logical             :: adios2_data3d_io_initialized
  logical             :: adios2_phi_io_initialized
  logical             :: adios2_ptrack_io_initialized
  logical             :: adios2_poincare_io_initialized
  logical             :: adios2_history_io_initialized

  ! Engine objects
  type(adios2_engine) :: adios2_eqplot_e
  type(adios2_engine) :: snapshot_engine
  type(adios2_engine) :: adios2_diag_hist_e
  type(adios2_engine) :: adios2_diag_data1d_e
  type(adios2_engine) :: adios2_data3d_e
  type(adios2_engine) :: adios2_phi_e
  type(adios2_engine) :: adios2_ptrack_e
  type(adios2_engine) :: adios2_poincare_e
  type(adios2_engine) :: adios2_history_e
  

contains


  !----------------------------------------------------------------------------!
  subroutine adios2_init_common()
    implicit none
  
    integer :: ierr, mpi_rank
    
    call adios2_init(adios2_obj, "gtc_adios2_config.xml", mpi_comm_world, adios2_debug_mode_on, ierr)
    call error_check(ierr, "adios2_init", .true.)
    
    ! Initialize eqplot IO object
    adios2_eqplot_io_initialized = .false.
    adios2_eqplot_io_id = "eqplot_io"

    ! Initialize history IO object (history_md.bp)
    adios2_history_io_initialized = .false.
    adios2_history_io_id = "history_io"
    
    ! Initialize restart IO object
    adios2_restart_io_initialized = .false.
    adios2_restart_io_id = "restart_io"
    adios2_restart_io_read_id = "restart_io_read"

    ! Initialize snapshot IO object
    adios2_snapshot_io_initialized = .false.
    adios2_snapshot_io_id = "snapshot_io"

    ! Initialize diagnosis history io object
    adios2_diag_hist_io_id = "diag_history_io"

    ! Initialize the restart obj for history files
    adios2_hist_restart_io_id = "hist_restart_io"
    
    ! Initialize diagnosis data1d io object
    adios2_diag_data1d_io_id = "diag_data1d_io"

    ! Initialize big data io object
    adios2_data3d_io_initialized = .false.
    adios2_data3d_id = "dataout3d_io"

    ! Initialize phishot io object
    adios2_phi_io_initialized = .false.
    adios2_phi_io_id = "phishot_io"

    ! Initialize particle tracking io object
    adios2_ptrack_io_initialized = .false.
    adios2_ptrack_io_id = "write_tracked_particles_io"

    ! Initialize poincare io object
    adios2_poincare_io_initialized = .false.
    adios2_poincare_io_id = "poincare_io"

    ! Declare the I/O objects for history.bp and history_restart.bp
    call adios2_declare_io (adios2_diag_hist_io,        adios2_obj, adios2_diag_hist_io_id,        ierr)
    call adios2_declare_io (adios2_hist_restart_io,     adios2_obj, adios2_hist_restart_io_id,     ierr)

    ! Declare the I/O objects for data1d.bp and data1d_restart.bp
    call adios2_declare_io (adios2_diag_data1d_io,      adios2_obj, adios2_diag_data1d_io_id,      ierr)
    call adios2_declare_io (adios2_data1d_restart_io,   adios2_obj, adios2_data1d_restart_io_id,   ierr)
  
  end subroutine adios2_init_common


  !----------------------------------------------------------------------------!
  subroutine adios2_finalize_common()
    implicit none

    integer :: ierr

    ! master,  close thy open engines
    if (mype .eq. 0) then
      ! Close eqplot engine
      if (adios2_eqplot_io_initialized) call adios2_close(adios2_eqplot_e, ierr)

      ! Close history engine
      if (adios2_history_io_initialized) call adios2_close(adios2_history_e, ierr)
    
      ! Close snapshot engine
      if (adios2_snapshot_io_initialized) call adios2_close(snapshot_engine, ierr)

      ! Close diagnosis history engine
      call adios2_close(adios2_diag_hist_e, ierr)

      ! Close diagnosis data1d engine
      call adios2_close(adios2_diag_data1d_e, ierr)

      ! Close phi engine
      if (adios2_phi_io_initialized) call adios2_close(adios2_phi_e, ierr)
    endif

    ! Close phi_3d engine
    if(adios2_data3d_io_initialized) call adios2_close (adios2_data3d_e, ierr)

    ! Close tracked particle engine
    if (adios2_ptrack_io_initialized) call adios2_close (adios2_ptrack_e, ierr)

    ! Close poincare engine
    if (adios2_poincare_io_initialized) call adios2_close (adios2_poincare_e, ierr)

  end subroutine adios2_finalize_common

end module adios2_common

