! This file is part of GTC version 4.4 
! GTC version 4.4 is released under the 3-Clause BSD license:

! Copyright (c) 2002,2010,2016, GTC Team (team leader: Zhihong Lin, zhihongl@uci.edu)
! All rights reserved.

! Redistribution and use in source and binary forms, with or without 
! modification, are permitted provided that the following conditions are met:

! 1. Redistributions of source code must retain the above copyright notice, 
!    this list of conditions and the following disclaimer.

! 2. Redistributions in binary form must reproduce the above copyright notice, 
!    this list of conditions and the following disclaimer in the documentation 
!    and/or other materials provided with the distribution.

! 3. Neither the name of the GTC Team nor the names of its contributors may be 
!    used to endorse or promote products derived from this software without 
!    specific prior written permission.

!> \file
!> Time diagnostics subroutines

!==================================
!|| New Diagnostic I/O Interface ||
!==================================
!> Diagnostic data using new file I/O, netCDF or ADIOS 

!> New diagnostic module that provides all diagnostic related functionalities
module diagnostic
  ! Currently only support netCDF files
  ! Future extensions can be made as sub-modules
  use precision, only: lk
  use global_parameters, only: gtcout, r0, b0, utime, etemp0, eden0, rho0,&
    istep, tstep, mstepall, mpsi, nhybrid
  use equilibrium, only: ped
  use particle_array, only: zonale, meshne, meshte
  use field_array, only: phi00, psimesh
  use spline_function, only: sprpsi
#ifdef ADIOS2
  use adios2
#endif
  implicit none

  private
  public init_diagnostic, write_history, calc_history
#ifdef ADIOS2
  public write_history_adios2
#endif

  ! reserved diagnostic file names
  character(len=*),parameter:: history_file='history.nc'
  ! netCDF IDs 
  ! file IDs
  integer :: history_id
  ! dimension IDs
  integer :: dim_time, dim_mpsi, dim_mgrid, dim_mtoroidal
  ! coordinate variable IDs
  integer :: psi_id, r_id
  ! equilibrium variable IDs
  integer :: ne0_id, te0_id
  ! time series data variable IDs
  integer :: t_id, zonal_ne_id, zonal_phi_id
  ! current time in simulation
  real(lk) :: t
  save
  contains

    !-------------------------------------------------------------------------
#ifdef ADIOS2
    subroutine write_history_adios2
        use adios2_common
        implicit none

        real(lk) :: r(0:mpsi), psimesh_tmp(0:mpsi)
        integer  :: i, ierr
        type(adios2_variable), save :: var_zonale, var_phi00, var_meshte, var_meshne, &
                                       var_t, var_r, var_psi
        type(adios2_attribute)      :: att_psi_unit, att_psi_scale, att_r_unit, att_r_scale, &
                                       att_t_unit, att_t_scale, att_ne0_unit, att_ne0_scale, &
                                       att_te0_unit, att_te0_scale, att_zonalne_unit,        &
                                       att_zonalphi_unit, att_zonalphi_scale

        ! Define variables and open file. This block is executed once
        if (.not. adios2_history_io_initialized) then
            call adios2_declare_io (adios2_history_io, adios2_obj, adios2_history_io_id, ierr)

            call adios2_define_variable (var_psi, adios2_history_io,    &
                "psi", adios2_type_real_lk, 1, 1_8*shape(psimesh),         &
                1_8*(/0/), 1_8*shape(psimesh), .true., ierr)
            
            call adios2_define_variable (var_r, adios2_history_io, "r", &
                adios2_type_real_lk, 1, 1_8*shape(r), 1_8*(/0/),           &
                1_8*shape(r), .true., ierr)

            call adios2_define_variable (var_meshte, adios2_history_io, &
                "Te0", adios2_type_real_lk , 1, 1_8 * shape(meshte),    &
                1_8 * (/0/), 1_8 * shape(meshte), .true., ierr)

            call adios2_define_variable (var_meshne, adios2_history_io, &
                "ne0", adios2_type_real_lk , 1, 1_8 * shape(meshne),    &
                1_8 * (/0/), 1_8 * shape(meshne), .true., ierr)

            if (nhybrid > 0) then
                call adios2_define_variable (var_zonale, adios2_history_io, &
                    "zonal_ne", adios2_type_real_lk , 1, 1_8 * shape(zonale),    &
                    1_8 * (/0/), 1_8 * shape(zonale), .true., ierr)
            endif

            call adios2_define_variable (var_phi00, adios2_history_io,  &
                "zonal_phi", adios2_type_real_lk , 1, 1_8 * shape(phi00),      &
                1_8 * (/0/), 1_8 * shape(phi00), .true., ierr)
            
            call adios2_define_variable (var_t, adios2_history_io, "t", &
                adios2_type_real_lk, ierr)

            call adios2_define_attribute (att_psi_unit, adios2_history_io, "psi unit", &
              "Maxwell", ierr)
            
            call adios2_define_attribute (att_psi_scale, adios2_history_io, "psi scale factor", &
              ped*r0*r0*b0, ierr)
            
            call adios2_define_attribute (att_r_unit, adios2_history_io, "r unit", &
              "cm", ierr)
            
            call adios2_define_attribute (att_r_scale, adios2_history_io, "r scale factor", &
              r0, ierr)
            
            call adios2_define_attribute (att_t_unit, adios2_history_io, "t unit", &
              "s", ierr)
            
            call adios2_define_attribute (att_t_scale, adios2_history_io, "t scale factor", &
              utime, ierr)
            
            call adios2_define_attribute (att_ne0_unit, adios2_history_io, "ne0 unit", &
              "cm^-3", ierr)
            
            call adios2_define_attribute (att_ne0_scale, adios2_history_io, "ne0 scale factor", &
              eden0, ierr)
            
            call adios2_define_attribute (att_te0_unit, adios2_history_io, "te0 unit", &
              "eV", ierr)
            
            call adios2_define_attribute (att_te0_scale, adios2_history_io, "te0 scale factor", &
              etemp0, ierr)
            
            call adios2_define_attribute (att_zonalne_unit, adios2_history_io, "zonal_ne unit", &
              "local_ne0", ierr)
            
            call adios2_define_attribute (att_zonalphi_unit, adios2_history_io, "zonal_phi unit", &
              "Te0/e", ierr)
            
            call adios2_define_attribute (att_zonalphi_scale, adios2_history_io, "zonal_phi scale factor", &
              1.0_lk/(rho0*rho0), ierr)
            
            ! Open file
            call adios2_open (adios2_history_e, adios2_history_io,        &
                "history_md.bp", adios2_mode_write, mpi_comm_self, ierr)

            adios2_history_io_initialized = .true.
            
            ! Write the coordinate variables
            do i=0, mpsi
                r(i) = sprpsi(psimesh(i))
            enddo

            psimesh_tmp = psimesh/ped

            ! The variables below are written once
            call adios2_begin_step (adios2_history_e, adios2_step_mode_append, ierr)
            call adios2_put (adios2_history_e, var_r, r, ierr)
            call adios2_put (adios2_history_e, var_psi, psimesh_tmp, ierr)
            call adios2_put (adios2_history_e, var_meshne, meshne, ierr)
            call adios2_put (adios2_history_e, var_meshte, meshte, ierr)
            call adios2_end_step (adios2_history_e, ierr)

        endif !if .not. adios2_history_io_initialized
        
        ! Write t, zonale, phi00 for all calls to write_history
        call adios2_begin_step (adios2_history_e, adios2_step_mode_append, ierr)
        call adios2_put (adios2_history_e, var_t, t, ierr)

        if (nhybrid > 0) call adios2_put (adios2_history_e, var_zonale, zonale, ierr)

        call adios2_put (adios2_history_e, var_phi00, phi00, ierr)
        call adios2_end_step (adios2_history_e, ierr)

    end subroutine write_history_adios2
#endif
    !-------------------------------------------------------------------------



    !> initialize I/O files and data arrays, write out some equilibrium data
    subroutine init_diagnostic

    ! Execute this subroutine only if ADIOS2 is not enabled
      
! #ifndef ADIOS2
      use netcdf_wrapper

      integer :: radial_time_dims(2)
      real(lk) :: r(0:mpsi)
      integer :: i

      call check_nf( nf90_create(history_file, NF90_CLOBBER, history_id))
      call check_nf( nf90_def_dim(history_id, 'mpsi', mpsi+1, dim_mpsi) )
      ! define the unlimited time dimension
      call check_nf( nf90_def_dim(history_id, 'time', NF90_UNLIMITED,&
        dim_time) )

      ! set the array shape
      radial_time_dims = (/ dim_mpsi, dim_time /)

      ! define coordiantes variables
      ! poloidal flux
      call check_nf( nf90_def_var(history_id, 'psi', NC_REAL, dim_mpsi, psi_id) )
      call check_nf( nf90_put_att(history_id, psi_id, "unit", "Maxwell") )
      call check_nf( nf90_put_att(history_id, psi_id, "scale_factor",ped*r0*r0*b0 ) )

      call check_nf( nf90_def_var(history_id, 'r', NC_REAL, dim_mpsi, r_id) )
      call check_nf( nf90_put_att(history_id, r_id, "unit", "cm") )
      call check_nf( nf90_put_att(history_id, r_id, "scale_factor",r0 ) )

      call check_nf( nf90_def_var(history_id, 't', NC_REAL, dim_time, t_id) )
      call check_nf( nf90_put_att(history_id, t_id, "unit", "s") )
      call check_nf( nf90_put_att(history_id, t_id, "scale_factor",utime ) )

      ! define relevant equilibrium variables
      ! equilibrium density
      call check_nf( nf90_def_var(history_id, 'ne0', NC_REAL, dim_mpsi, ne0_id) )
      call check_nf( nf90_put_att(history_id, ne0_id, "unit", "cm^-3") )
      call check_nf( nf90_put_att(history_id, ne0_id, "scale_factor",eden0 ) )

      ! equilibrium temperature
      call check_nf( nf90_def_var(history_id, 'Te0', NC_REAL, dim_mpsi, te0_id) )
      call check_nf( nf90_put_att(history_id, te0_id, "unit", "eV") )
      call check_nf( nf90_put_att(history_id, te0_id, "scale_factor",etemp0 ) )

      ! define data variables
      call check_nf( nf90_def_var(history_id, 'zonal_ne', NC_REAL,&
        radial_time_dims, zonal_ne_id) )
      call check_nf( nf90_put_att(history_id, zonal_ne_id, "unit", "local_ne0") )

      call check_nf( nf90_def_var(history_id, 'zonal_phi', NC_REAL,&
        radial_time_dims, zonal_phi_id) )
      call check_nf( nf90_put_att(history_id, zonal_phi_id, "unit", "Te0/e") )
      call check_nf( nf90_put_att(history_id, zonal_phi_id, "scale_factor",&
        1.0_lk/(rho0*rho0)) )


      call check_nf( nf90_enddef(history_id) )

      ! write the coordinate variables
      call check_nf( nf90_put_var(history_id, psi_id, psimesh/ped ) )
      do i=0, mpsi
        r(i) = sprpsi(psimesh(i))
      enddo
      call check_nf( nf90_put_var(history_id, r_id, r ) )

      ! write out relevant equilibrium quantities
      call check_nf( nf90_put_var(history_id, ne0_id, meshne ) )
      call check_nf( nf90_put_var(history_id, Te0_id, meshte ) )

      call check_nf( nf90_close(history_id) )
! #endif
    end subroutine init_diagnostic

    !> prepare time series data for current time step
    subroutine calc_history
      ! Currently, only the time needs to be calculated here
      t = (istep+mstepall)*tstep
    end subroutine calc_history

    !> output time series data
    subroutine write_history
      use netcdf_wrapper

      ! current length of time record
      integer:: numTimes
      ! open the history file
      call check_nf( nf90_open(history_file, NF90_WRITE, history_id) )
      ! refresh the variable IDs
      call check_nf( nf90_inq_varid(history_id, "t", t_id) )
      if(nhybrid>0) then
        call check_nf( nf90_inq_varid(history_id, "zonal_ne", zonal_ne_id) )
      endif
      call check_nf( nf90_inq_varid(history_id, "zonal_phi", zonal_phi_id) )

      ! check the current time record length
      call check_nf( nf90_inq_dimid(history_id, "time", dim_time) )
      call check_nf( nf90_inquire_dimension(history_id, dim_time, len=numTimes) )

      ! write new time slice to the file
      ! write the new time
      call check_nf( nf90_put_var(history_id, t_id, t, start=(/numTimes+1/)) )
      if(nhybrid > 0) then
        call check_nf( nf90_put_var(history_id, zonal_ne_id, zonale, &
          start=(/1, numTimes+1/),count=(/mpsi+1 , 1/)) )
      endif
      call check_nf( nf90_put_var(history_id, zonal_phi_id, phi00, start=(/1, numTimes+1/),&
        count=(/mpsi+1 , 1/)) )

      call check_nf(nf90_close(history_id))
    end subroutine write_history

end module diagnostic

!> wrapper for both current and new diagnostic subroutines
subroutine diagnosis
  use global_parameters, only: mype,mpsi,magnetic
#ifdef ADIOS2
  use diagnostic, only: calc_history, write_history, write_history_adios2
#else
  use diagnostic, only: calc_history, write_history
#endif
  use field_array
  use equilibrium
  implicit none

  integer i,j,ij,isp
  real(lk) pdum,q_inv,dpx,dp2,g,ri,gqi_inv,dptdt,dptdz

  ! Computation of the MHD displacement based on Z. Wang et al. PoP 2015, Eq. (A10)
  if(magnetic==1) then
     do i=0,mpsi
        pdum=psimesh(i)
        q_inv=1/qmesh(i)

        isp=max(1,min(lsp-1,ceiling(pdum*spdpsi_inv)))
        dpx=pdum-spdpsi*real(isp-1)
        dp2=dpx*dpx
        g=   gpsi(1,isp)   +gpsi(2,isp)*dpx   +gpsi(3,isp)*dp2
        ri=  cpsi(1,isp)   +cpsi(2,isp)*dpx   +cpsi(3,isp)*dp2
        gqi_inv=1.0/(g*qmesh(i)+ri)

        do j=0,mtheta(i)
           ij=igrid(i)+j
           dptdt=gradphi(2,1,ij) ! contain only non-zonal
           dptdz=gradphi(3,1,ij)-gradphi(2,1,ij)*q_inv  ! contain only non-zonal

           xi_MHD(ij) = xi_MHD(ij) + gqi_inv*sqrt(gdpp(ij))*(ri*dptdz-g*dptdt)
        enddo
     enddo
  endif

  ! current diagnosis written in history.out
  call current_diagnosis
  ! new diagnosics written in history.nc
  call calc_history
  if (mype == 0) then
#ifdef ADIOS2
    call write_history_adios2
#endif
    call write_history
  endif
end subroutine diagnosis

subroutine current_diagnosis
  use system_env
  use utility, only: check_allocation, mstat
  use precision, only: mpi_Rsize, mpi_Csize
  use global_parameters
  use particle_array
  use field_array
  use conservative_scheme,only:sapara_na
  use antenna, only: phi_ext
#ifdef ADIOS2
  use adios2
  use adios2_common
  use adios2_history_diagnostics
  use adios2_data1d_diagnostics
#endif
  use conservative_scheme,only:sapara_na
  implicit none

  integer,parameter :: nfield=3,mfdiag=4,mfdata1d=2
  integer i,j,k,ii,ip,icount,mtgrid,nf,mtest,ierr
  real(lk) fieldmode(2,modes,nfield),fieldtime(mfdiag,nfield),partdata(mpdiag,nspecies),&
       field00(0:mpsi,nfield),fieldrms(0:mpsi,nfield),rmstmp(0:mpsi,nfield)
  real(lk),dimension(:,:,:),allocatable :: fieldgrid
  real(lk),dimension(:),allocatable :: kapatmti,kapatmte,kapatmtf,kapatmtfe,kapatmni,kapatmne,kapatmnf,kapatmnfe
  save kapatmti,kapatmte,kapatmtf,kapatmtfe,kapatmni,kapatmne,kapatmnf,kapatmnfe

  integer  :: ndstep,ndata,mpsi_tmp,step_index
  real(lk) :: tstep_tmp

! open output files history.out and data1d.out
  if(mype==0 .and. istep==ndiag)call opendiag(mpdiag,nfield,mfdiag,modes,mpdata1d,mfdata1d)
  call MPI_BCAST(mstepall,1,MPI_INTEGER,0,MPI_COMM_WORLD,ierror)
!!! particle diagnosis: main ion, electron and/or EP, imputities, ...
! record the initial temperature/density gradient for normalization
  if(istep==ndiag)then
     allocate(kapatmti(0:mpsi),kapatmte(0:mpsi),kapatmtf(0:mpsi),kapatmtfe(0:mpsi),kapatmni(0:mpsi),&
     kapatmne(0:mpsi),kapatmnf(0:mpsi),kapatmnfe(0:mpsi),STAT=mstat)
     call check_allocation(mstat, "kapam ararys in diagnosis.F90")
     do i=0,mpsi
        kapatmti(i)=max(abs(kapati(i)),1.0e-2_lk)
        kapatmte(i)=max(abs(kapate(i)),1.0e-2_lk)
        kapatmtf(i)=max(abs(kapatf(i)),1.0e-2_lk)
        kapatmtfe(i)=max(abs(kapatfe(i)),1.0e-2_lk)
        kapatmni(i)=max(abs(kapani(i)),1.0e-2_lk)
        kapatmne(i)=max(abs(kapane(i)),1.0e-2_lk)
        kapatmnf(i)=max(abs(kapanf(i)),1.0e-2_lk)
        kapatmnfe(i)=max(abs(kapanfe(i)),1.0e-2_lk)
     enddo

  endif

!!convert heat and particle flux to Bohm unit with normalization on middle flux surface
  diagion(7)=diagion(7)/(rho0*rho0*kapatmni(diag_flux)*gpsi200(diag_flux))
  diagion(9)=diagion(9)/(rho0*rho0*kapatmti(diag_flux)*gpsi200(diag_flux))

  if(nhybrid>0)then
     diagelectron(7)=diagelectron(7)*tfracn/(rho0*rho0*kapatmne(diag_flux)*gpsi200(diag_flux))
     diagelectron(9)=diagelectron(9)*tfracn/(rho0*rho0*kapatmte(diag_flux)*gpsi200(diag_flux))
  endif
  if(fload>0)then
     diagfast(7)=diagfast(7)/(rho0*rho0*kapatmnf(diag_flux)*gpsi200(diag_flux))
     diagfast(9)=diagfast(9)/(rho0*rho0*kapatmtf(diag_flux)*gpsi200(diag_flux))
  endif
 
  if(feload>0)then
     diagfaste(7)=diagfaste(7)*fetfracn/(rho0*rho0*kapatmnfe(diag_flux)*gpsi200(diag_flux))
     diagfaste(9)=diagfaste(9)*fetfracn/(rho0*rho0*kapatmtfe(diag_flux)*gpsi200(diag_flux))
  endif

  partdata=0.0_lk
  do ip=1,nspecies
     if(ip==1)then
        partdata(:,ip)=diagion/diagion(11) !first species is thermal ion

     elseif(ip==2)then
        if(nhybrid>0)then
           partdata(:,ip)=diagelectron/diagelectron(11) !second species is electron
        elseif(feload>0)then
           partdata(:,ip)=diagfaste/diagfaste(11) !second species is fast electron
        else
           partdata(:,ip)=diagfast/diagfast(11) !second species is fast ion if electron and fast electron not loaded
        endif

     elseif(ip==3)then
        if(feload>0)then
          partdata(:,ip)=diagfaste/diagfaste(11) !third species is fast electron
        else
          partdata(:,ip)=diagfast/diagfast(11) !third species is fast ion
        endif
     endif
  enddo

!!! field diagnosis: phi, a_para, fluid_ne, ...
  ii=igrid(diag_flux)
  mtgrid=mtheta(diag_flux)
  fieldtime=0.0_lk 
  fieldmode=0.0_lk
  allocate(fieldgrid(0:1,0:mtgrid,nfield), STAT=mstat)
  call check_allocation(mstat, "fieldgrid in diagnosis.F90")
  
  do nf=1,nfield
     if(nf==1)then
       ! external antenna phi is subtracted from the diagnostic
       ! Note that phi_ext is normalized to Te/e which is the diagnostic unit,
       ! not the GTC unit. So it must be put outside the (/rho0**2)
       ! normalization.
       ! if antenna is not used, phi_ext is simply 0
        fieldgrid(:,:,nf)= phi(:,ii:ii+mtgrid)/(rho0*rho0) - phi_ext(:,ii:ii+mtgrid)
!time history of field quantity at theta=zeta=0 & i=iflux
        fieldtime(1,nf)=phi(0,ii)/(rho0*rho0)-phi_ext(0,ii)
        fieldtime(2,nf)=phip00(diag_flux)/rho0 
        fieldtime(3,nf)=sqrt(sum(phip00(0:mpsi)**2)/real(mpsi+1))/rho0
!$omp parallel do private(i)
        do i=0,mpsi
           field00(i,nf)=phi00(i)/(rho0*rho0)
           fieldrms(i,nf)=sum((phi(0,igrid(i):igrid(i)+mtheta(i)-1)/(rho0**2) - &
                            phi_ext(0,igrid(i):igrid(i)+mtheta(i)-1) )**2)
        enddo

     elseif(nf==2)then
! apara is renormalized in such a way that it has the same
! amplitude as phi in ideal shear Alfven waves
        fieldgrid(:,:,nf)=sapara(:,ii:ii+mtgrid)/(rho0*sqrt(betae*aion))
        fieldtime(1,nf)=sapara(0,ii)/(rho0*sqrt(betae*aion))
        fieldtime(2,nf)=apara00(diag_flux)/(rho0*sqrt(betae*aion))
        fieldtime(3,nf)=sqrt(sum(apara00(0:mpsi)**2)/real(mpsi+1))/(rho0*sqrt(betae*aion)) 
!$omp parallel do private(i)
        do i=0,mpsi
           field00(i,nf)=apara00(i)/(rho0*sqrt(betae*aion))
           fieldrms(i,nf)=sum(sapara(0,igrid(i):igrid(i)+mtheta(i)-1)**2)/(rho0*rho0*betae*aion)
        enddo

     elseif(nf==3)then
        if(cs_method==0)then
!Output electron density for hybrid model
          fieldgrid(:,:,nf)=sfluidne(:,ii:ii+mtgrid)
          fieldtime(1,nf)=sfluidne(0,ii)
          fieldtime(2,nf)=fluidne00(diag_flux)
          fieldtime(3,nf)=sqrt(sum(fluidne00(0:mpsi)**2)/real(mpsi+1))
!$omp parallel do private(i)
          do i=0,mpsi
             field00(i,nf)=fluidne00(i)
             fieldrms(i,nf)=sum(sfluidne(0,igrid(i):igrid(i)+mtheta(i)-1)**2)
          enddo
        elseif(cs_method==1)then
!Output non-adiabatic vector potential for conservative model
          fieldgrid(:,:,nf)=sapara_na(:,ii:ii+mtgrid)/(rho0*sqrt(betae*aion))
          fieldtime(1,nf)=sapara_na(0,ii)/(rho0*sqrt(betae*aion))
          fieldtime(2,nf)=apara00(diag_flux)/(rho0*sqrt(betae*aion))
          fieldtime(3,nf)=sqrt(sum(apara00(0:mpsi)**2)/real(mpsi+1))/(rho0*sqrt(betae*aion))
!$omp parallel do private(i)
          do i=0,mpsi
             field00(i,nf)=apara00(i)/(rho0*sqrt(betae*aion))
             fieldrms(i,nf)=sum(sapara_na(0,igrid(i):igrid(i)+mtheta(i)-1)**2)/(rho0*rho0*betae*aion)
          enddo
        endif
     endif
  enddo

!Volume-averaged RMS
  icount=nfield*(mpsi+1)
  call MPI_REDUCE(fieldrms,rmstmp,icount,mpi_Rsize,MPI_SUM,0,toroidal_COMM,ierror)
  fieldrms=rmstmp
  do nf=1,nfield
     fieldtime(4,nf)=sqrt(sum(fieldrms(:,nf))/real(mtoroidal*sum(mtheta)))
!$omp parallel do private(i)
     do i=0,mpsi
        fieldrms(i,nf)=sqrt(fieldrms(i,nf)/real(mtoroidal*mtheta(i)))
     enddo
  enddo

  call spectrum(nfield,modes,nmodes,mmodes,diag_flux,mtgrid,fieldgrid,fieldmode)

  deallocate(fieldgrid)

  if(mype==0)then
#ifdef ADIOS2
    ndata=(nspecies*mpdiag+nfield*(2*modes+mfdiag))
    tstep_tmp = tstep*ndiag
    call adios2_begin_step (adios2_diag_hist_e, adios2_step_mode_append, ierr)
    if (irun.eq.0) then
      if ( (mype.eq.0) .and. (istep.eq.ndiag) ) then
        call adios2_put (adios2_diag_hist_e, t_hist_diag%nspecies, nspecies,  ierr)
        call adios2_put (adios2_diag_hist_e, t_hist_diag%mpdiag,   mpdiag,    ierr)
        call adios2_put (adios2_diag_hist_e, t_hist_diag%nfield,   nfield,    ierr)
        call adios2_put (adios2_diag_hist_e, t_hist_diag%modes,    modes,     ierr)
        call adios2_put (adios2_diag_hist_e, t_hist_diag%mfdiag,   mfdiag,    ierr)
        call adios2_put (adios2_diag_hist_e, t_hist_diag%tstep,    tstep_tmp, ierr)
      endif
    endif
    step_index = mstepall + mstep/ndiag
    ! call adios2_put (adios2_diag_hist_e, t_hist_diag%numsteps, step_index,ierr)
    call adios2_put (adios2_diag_hist_e, t_hist_diag%partdata,  partdata,  ierr)
    call adios2_put (adios2_diag_hist_e, t_hist_diag%fieldtime, fieldtime, ierr)
    call adios2_put (adios2_diag_hist_e, t_hist_diag%fieldmode, fieldmode, ierr)
    call adios2_end_step (adios2_diag_hist_e, ierr)
#endif
! #else
! write particle and field data to history file
     do i=1,nspecies
        do j=1,mpdiag
           write(iodiag,102)partdata(j,i)
        enddo
     enddo
     do i=1,nfield
        do j=1,mfdiag
           write(iodiag,102)fieldtime(j,i)
        enddo
     enddo
     do i=1,nfield
        do j=1,modes
           write(iodiag,102)fieldmode(1,j,i),fieldmode(2,j,i)
        enddo
     enddo
     flush(iodiag)
! #endif

! convert 1d heat and particle diffusivity to Bohm unit
     do i=1,mpdata1d
         if(i==1)data1di(:,i)=data1di(:,i)/(rho0*rho0*kapatmni*gpsi200)
         if(i==2)data1di(:,i)=data1di(:,i)/(rho0*rho0*kapatmti*gpsi200)
         if(nspecies>1)then
            if(nhybrid>0)then
               if(i==1)data1de(:,i)=data1de(:,i)*tfrac/(rho0*rho0*kapatmne*gpsi200)
               if(i==2)data1de(:,i)=data1de(:,i)*tfrac/(rho0*rho0*kapatmte*gpsi200)
               if(i==3)data1de(:,i)=data1de(:,i)/sqrt(gpsi200)
               if(i==4)data1de(:,i)=data1de(:,i)/sqrt(gpsi200)
            endif

            if(fload>0)then
               if(i==1)data1df(:,i)=data1df(:,i)/(rho0*rho0*kapatmnf*gpsi200)
               if(i==2)data1df(:,i)=data1df(:,i)/(rho0*rho0*kapatmtf*gpsi200)
               if(i==3)data1df(:,i)=data1df(:,i)/sqrt(gpsi200)
               if(i==4)data1df(:,i)=data1df(:,i)/sqrt(gpsi200)
            endif
             
              if(feload>0)then
               if(i==1)data1dfe(:,i)=data1dfe(:,i)*fetfrac/(rho0*rho0*kapatmnfe*gpsi200)
               if(i==2)data1dfe(:,i)=data1dfe(:,i)*fetfrac/(rho0*rho0*kapatmtfe*gpsi200)
               if(i==3)data1dfe(:,i)=data1dfe(:,i)/sqrt(gpsi200)
               if(i==4)data1dfe(:,i)=data1dfe(:,i)/sqrt(gpsi200)
            endif

         endif
     enddo

! write radial-time data to data1d.out
#ifdef ADIOS2
    call adios2_begin_step (adios2_diag_data1d_e, adios2_step_mode_append,  ierr)
    if (irun.eq.0) then
      if ( (mype.eq.0) .and. (istep.eq.ndiag) ) then
        mpsi_tmp = mpsi+1
        call adios2_put (adios2_diag_data1d_e, t_data1d_diag%nspecies, nspecies, ierr)
        call adios2_put (adios2_diag_data1d_e, t_data1d_diag%mpsi,     mpsi_tmp, ierr)
        call adios2_put (adios2_diag_data1d_e, t_data1d_diag%nfield,   nfield,   ierr)
        call adios2_put (adios2_diag_data1d_e, t_data1d_diag%nhybrid,  nhybrid,  ierr)
        call adios2_put (adios2_diag_data1d_e, t_data1d_diag%mpdata1d, mpdata1d, ierr)
        call adios2_put (adios2_diag_data1d_e, t_data1d_diag%mfdata1d, mfdata1d, ierr)
      endif
    endif
    call adios2_put   (adios2_diag_data1d_e, t_data1d_diag%data1di,  data1di,  ierr)
    call adios2_put   (adios2_diag_data1d_e, t_data1d_diag%field00,  field00,  ierr)
    call adios2_put   (adios2_diag_data1d_e, t_data1d_diag%fieldrms, fieldrms, ierr)

    if(nspecies>1)then
      if(nhybrid>0) call adios2_put (adios2_diag_data1d_e, t_data1d_diag%data1de,  data1de,  ierr)
      if(fload>0)   call adios2_put (adios2_diag_data1d_e, t_data1d_diag%data1df,  data1df,  ierr)
      if(feload>0)  call adios2_put (adios2_diag_data1d_e, t_data1d_diag%data1dfe, data1dfe, ierr)
    endif
    
    call adios2_end_step   (adios2_diag_data1d_e, ierr)
#endif
! #else
    ! write radial-time data to data1d.out
     write(iodata1d,102)data1di
     if(nspecies>1)then
        if(nhybrid>0)write(iodata1d,102)data1de
        if(fload>0)write(iodata1d,102)data1df
        if(feload>0)write(iodata1d,102)data1dfe
     endif

     write(iodata1d,102)field00
     write(iodata1d,102)fieldrms
     flush(iodata1d)
! #endif

    ! monitor: rms of phi, apara, fluidne, heatflux of ion & electron
     write(gtcout,"(I7,6e15.6)")istep+mstepall,(fieldtime(4,i),i=1,nfield),&
          (partdata(9,i),i=1,nspecies)
     flush(gtcout)
  endif !if mype==0

101 format(i6)
102 format(e13.6)

end subroutine current_diagnosis
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! calculate (n,m) mode amplitude for diagnostics
subroutine spectrum(nfield,modes,nmodes,mmodes,diag_flux,mtgrid,fieldgrid,fieldmode)
  use system_env 
  use precision, only: mpi_Rsize, mpi_Csize, lk
  use global_parameters, only: pi2,fielddir,pi2_inv,toroidal_comm,myrank_toroidal,mtdiag,&
  mtoroidal
  use field_array,only: zeta0,deltaz,qtinv
  implicit none

  integer nfield,modes,nmodes(modes),mmodes(modes),diag_flux,mtgrid
  real(lk) fieldgrid(0:1,0:mtgrid,nfield),fieldmode(2,modes,nfield)
  integer i,j,k,jt,kz,mzeach,jpe,indp,indt,indp1,indt1,mteach,icount,nf
  real(lk) wt,dt,wz,zdum,tdum,xz(mtdiag),fieldflux(mtdiag/mtoroidal,mtdiag,nfield),&
       eachzeta((mtdiag/mtoroidal)*(mtdiag/mtoroidal)*nfield),&
       allzeta((mtdiag/mtoroidal)*mtdiag*nfield)
  complex(lk) yzeta(nfield,(mtdiag/mtoroidal)*modes),ytheta(nfield,mtdiag*modes),&
       yz(mtdiag/2+1),ye(mtdiag) 

! mesh in magnetic coordinates
  mzeach=mtdiag/mtoroidal
  mteach=mzeach
  dt=pi2/real(mtdiag)

! Interpolate on a flux surface from fieldline coordinates to magnetic coordinates. 
! Use mtdiag for both poloidal and toroidal grid points.
!$omp parallel do private(kz,nf,j,tdum,zdum,jt,wt,wz)           
  do kz=1,mzeach
     wz=real(kz)/real(mzeach)
     zdum=zeta0+wz*deltaz
     do nf=1,nfield
        do j=1,mtdiag
           if(fielddir==1 .or. fielddir==3)then
             tdum=pi2_inv*modulo(dt*real(j)-(zdum-pi2)*qtinv(diag_flux),pi2)*real(mtgrid)
           else
             tdum=pi2_inv*modulo(dt*real(j)-zdum*qtinv(diag_flux),pi2)*real(mtgrid)
           endif
           jt=max(0,min(mtgrid-1,int(tdum)))
           wt=tdum-real(jt)
           
           fieldflux(kz,j,nf)=wz*((1.0_lk-wt)*fieldgrid(1,jt,nf)+wt*fieldgrid(1,jt+1,nf))+&
                (1.0_lk-wz)*((1.0_lk-wt)*fieldgrid(0,jt,nf)+wt*fieldgrid(0,jt+1,nf))
        enddo
     enddo
  enddo
 
! transpose 2-d matrix from (mtoroidal,mzeach) to (1,mtoroidal*mzeach)
  allzeta=0.0_lk
  icount=mteach*mzeach*nfield  
  do jpe=0,mtoroidal-1

!$omp parallel do private(j,nf,k,jt,indt,indp1,indp)
     do j=1,mteach
        jt=jpe*mteach+j
        indt=(j-1)*mzeach
        do nf=1,nfield
           indp1=indt+(nf-1)*mteach*mzeach
           do k=1,mzeach
              indp=indp1+k
              eachzeta(indp)=fieldflux(k,jt,nf)
           enddo
        enddo
     enddo
     
     call MPI_GATHER(eachzeta,icount,mpi_Rsize,allzeta,icount,mpi_Rsize,jpe,toroidal_comm,ierror)
  enddo
    
! transform to k space
  yz=0.0_lk
  do j=1,mteach
     indt1=(j-1)*mzeach
     do nf=1,nfield
        indt=indt1+(nf-1)*mteach*mzeach

!$omp parallel do private(kz,k,indp)      
        do kz=0,mtoroidal-1
           do k=1,mzeach
              indp=kz*icount+indt+k
              xz(kz*mzeach+k)=allzeta(indp)
           enddo
        enddo
      
        call fftr1d(1,mtdiag,1.0_lk,xz,yz,2)

! record toroidal mode amplitude for diagnostic
        do kz=1,modes
           yzeta(nf,j+mteach*(kz-1))=yz(nmodes(kz)+1)
        enddo

     enddo
  enddo

! gather toroidal mode amplitude for calculation of poloidal mode amplitudes
  ytheta=0.0_lk
  icount=mteach*modes*nfield
  call MPI_GATHER(yzeta,icount,mpi_Csize,ytheta,icount,mpi_Csize,0,toroidal_comm,ierror)

  icount=mteach*modes
  if(myrank_toroidal == 0)then
     do nf=1,nfield
        do kz=1,modes

!$omp parallel do private(i,j)     
           do i=0,mtoroidal-1
              do j=1,mteach
                 ye(j+i*mteach)=ytheta(nf,j+(kz-1)*mteach+i*icount)
              enddo
           enddo

           call fftc1d(1,mtdiag,1.0_lk,ye)
           
           if(fielddir==1 .or. fielddir==3)then
             fieldmode(1,kz,nf)=real(ye(mmodes(kz)+1))/(real(mtdiag)**2)
             fieldmode(2,kz,nf)=aimag(ye(mmodes(kz)+1))/(real(mtdiag)**2)
           else
             fieldmode(1,kz,nf)=real(ye(mtdiag-mmodes(kz)+1))/(real(mtdiag)**2)
             fieldmode(2,kz,nf)=aimag(ye(mtdiag-mmodes(kz)+1))/(real(mtdiag)**2)
           endif
        enddo
     enddo
  endif

end subroutine spectrum
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! filter out some poloidal modes in linear simulation
subroutine mfilter_frc(farray)
  use global_parameters
  use field_array
  implicit none

  integer i,j,k,ii,ij,ip,jt,kz,ipe,kpe,jtp,&
       icount,idest,isource,l,isp,maxmmode
  real(lk) farray(0:1,mgrid),farray2d(mthetamax,mpsi),&
       ffilter(mtoroidal/2+1),xz(mtoroidal)
  complex(lk) yz(mtoroidal/2+1)

 !2d theta-psi array
 !$omp parallel do private(i,j,ij)
  do i=1,mpsi-1
     do j=1,mthetamax
        ij=igrid(i)+j
        farray2d(j,i)=farray(1,ij)
     enddo
  enddo

  ffilter=0.0_lk
  if(nfilter==2)then
      ffilter(mmodes+1)=1.0_lk
  endif

  yz=0.0_lk
  xz=0.0_lk
 
  !$omp parallel do private(i,xz,yz)
  do i=1,mpsi-1
       xz=farray2d(:,i)
       call fftr1d(1,mthetamax,1.0_lk,xz,yz,2)
       yz=yz*ffilter
       call fftr1d(-1,mthetamax,1.0_lk,xz,yz,2)
       farray2d(:,i)=xz
  enddo

  !2d theta-psi array
  !$omp parallel do private(i,j,ij)
  do i=1,mpsi-1
     do j=1,mthetamax
        ij=igrid(i)+j
        farray(1,ij)=farray2d(j,i)
     enddo
  enddo
 
  !$omp parallel do private(i)
  do i=1,mpsi-1
        farray(1,igrid(i))=farray(1,igrid(i)+mtheta(i)) ! poloidal BC
  enddo


end subroutine mfilter_frc
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! cpu and gpu wrapper
subroutine filter_cpu(farray)
  use precision, only:lk
  use global_parameters, only: mgrid
#ifdef _OPENACC
  use openacc
#endif
  implicit none
  logical tmp
  real(lk),dimension(0:1,mgrid) :: farray

#ifdef _OPENACC
  if(acc_is_present(farray))then
    tmp = .false.
    !$acc update device(farray)
  else
    tmp = .true.
    !$acc enter data copyin(farray)
  endif
#endif
  call filter_t(farray)
#ifdef _OPENACC
  if(tmp)then
    !$acc exit data copyout(farray)
    !$acc exit data delete(farray)
  else
    !$acc update host(farray)
  endif
#endif
end subroutine filter_cpu
subroutine filter_gpu(farray)
  call filter_t(farray)
end subroutine filter_gpu
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! filter out some toroidal (and optionally poloidal) modes in linear simulation
subroutine filter_t(farray)
  use system_env
  use precision, only: mpi_Rsize
  use global_parameters
  use field_array
  use interfaces
#ifdef TIME_MPI
  use mpi_timing
#endif  
  implicit none
  
  integer i,j,k,ii,ij,ip,jt,kz,jpe,indp,indt,indp1,indt1,meachtheta,mteach,mzeach,jtp,&
       icount,idest,isource,isendtag,l,irecvtag,joffset,mteachpd,&
       istatus(MPI_STATUS_SIZE)
  real(lk) wt,r,dt,wz,zdum,tdum,ptemp,sendl(mgrid),sendr(mgrid),recvl(mgrid),recvr(mgrid),&
       phism(mgrid),pleft(mthetamax),pright(mthetamax),phitmp(mgrid),ffilter(mtdiag/2+1),&
       xz(mtdiag),farray(0:1,mgrid),&
       phifluxpartd(mpsi,mtdiag/npartdom,mtdiag/mtoroidal),&
       phifluxlz(mpsi,mtdiag/numberpe,mtdiag),&
       phiflux(mpsi,mtdiag),&
       allzeta((mtdiag/numberpe)*mtdiag*mpsi)
  !$acc declare create(phifluxpartd,farray,allzeta,phiflux,phifluxlz,ffilter)
  complex(lk) yz(mtdiag/2+1)
     
! periodic BC in poloidal and toroidal
  CALL PERIODICITY_GPU(farray)
!by zhs nfilter > 2, use zhs's filter
  if(nfilter > 2)then
     !$acc update host(farray) 
     call mnfilter(farray)
     !$acc update device(farray)
     return
   endif

! mesh in magnetic coordinates
  mzeach=mtdiag/mtoroidal
  mteach=mzeach/npartdom
  mteachpd=(mtdiag/npartdom)
  dt=2.0_lk*pi/real(mtdiag)

  ffilter=0.0_lk
  ffilter(nmodes+1)=1.0_lk
  !$acc update device(ffilter)
  !$acc parallel loop gang vector present(allzeta)
  do i=1,(mtdiag/numberpe)*mtdiag*mpsi
    allzeta(i)=0.0_lk
  enddo
  !$acc end parallel

  joffset=mteachpd*myrank_partd

! Interpolate on a flux surface from fieldline coordinates to magnetic
  ! coordinates. Use mtdiag for both poloidal and toroidal grid points.
#ifdef _OPENACC
  !$acc parallel loop gang vector present(qtinv,mtheta,igrid,farray) &
  !$acc& present(phifluxpartd) collapse(2)
#else
  !$omp parallel do private(kz,i,j,wz,ii,zdum,tdum,jt,wt)
#endif  
  do kz=1,mzeach     
    do i=1,mpsi
      !$acc loop seq
      do j=1,mteachpd
        wz=real(kz)/real(mzeach)
        zdum=zeta0+wz*deltaz            
        ii=igrid(i)           
        if(fielddir==1 .or. fielddir==3)then
          tdum=pi2_inv*modulo(dt*real(j+joffset)-(zdum-pi2)*qtinv(i),pi2)*real(mtheta(i))
        else
          tdum=pi2_inv*modulo(dt*real(j+joffset)-zdum*qtinv(i),pi2)*real(mtheta(i))
        endif
        jt=max(0,min(mtheta(i)-1,int(tdum)))
        wt=tdum-real(jt)
        phifluxpartd(i,j,kz)=((1.0_lk-wt)*farray(1,ii+jt)+wt*farray(1,ii+jt+1))*&
            wz+(1.0_lk-wz)*((1.0_lk-wt)*farray(0,ii+jt)+wt*farray(0,ii+jt+1))    
        
      enddo
    enddo
  enddo
  !$acc end parallel
! transpose 2-d matrix from (mtoroidal,mzeach) to (1,mtoroidal*mzeach)
  icount=mteach*mzeach*mpsi
#ifdef _OPENACC
  !$acc parallel loop gang vector present(phifluxpartd,allzeta) collapse(4)
#else  
  !$omp parallel do private(jpe,j,i,k,jt,indt,indp1,indp)
#endif
  do jpe=0,mtoroidal-1
    do j=1,mteach
      do k=1,mzeach
        do i=1,mpsi
          jt=jpe*mteach+j
          indt=(j-1)*mpsi            
          indp1=indt+(k-1)*mpsi*mteach               
          indp=indp1+i
          allzeta(indp+jpe*icount)=phifluxpartd(i,jt,k)
        enddo
      enddo
    enddo
  enddo
  !$acc end parallel

#ifdef TIME_MPI
  call timer_start()
#endif
#ifdef CUDA_MPI
  !$acc host_data use_device(allzeta,phifluxlz)  
#else
  !$acc update host(allzeta)  
#endif  
  call MPI_ALLTOALL(allzeta,icount,mpi_Rsize,phifluxlz,icount,&
      mpi_Rsize,toroidal_comm,ierror)
#ifdef CUDA_MPI  
  !$acc end host_data
#else
  !$acc update device(phifluxlz)
#endif  
  
#ifdef TIME_MPI
  filter_alltoall_time=filter_alltoall_time+timer_get_et()  
  filter_alltoall_count=filter_alltoall_count+1
  filter_alltoall_size=filter_alltoall_size+mtoroidal*icount*lk
#endif  
  
! transform to k space
  yz=0.0_lk
#ifdef _OPENACC
  !$acc host_data use_device(phifluxlz,ffilter)
  call cufft_filter1(phifluxlz,ffilter,mtdiag,mpsi,mteach)
  !$acc end host_data
#else
!$omp parallel do private(j,i,kz,k,indt1,indt,indp,xz,yz)
  do j=1,mteach
    do i=1,mpsi
      xz=phifluxlz(i,j,:)
      call fftr1d(1,mtdiag,1.0_lk,xz,yz,2)
      ! linear run only keep a few modes
      yz=ffilter*yz
      ! transform back to real space
      call fftr1d(-1,mtdiag,1.0_lk,xz,yz,2)
      phifluxlz(i,j,:)=xz
    enddo
  enddo
#endif  

#ifdef TIME_MPI
  call timer_start()
#endif

#ifdef CUDA_MPI
  !$acc host_data use_device(phifluxlz,allzeta)
#else
  !$acc update host(phifluxlz)
#endif  
  call MPI_ALLTOALL(phifluxlz,icount,mpi_Rsize,allzeta,&
      icount,mpi_Rsize,toroidal_comm,ierror)
#ifdef CUDA_MPI  
  !$acc end host_data
#else
  !$acc update device(allzeta)
#endif  
#ifdef TIME_MPI
  filter_alltoall_time=filter_alltoall_time+timer_get_et()  
  filter_alltoall_count=filter_alltoall_count+1
  filter_alltoall_size=filter_alltoall_size+mtoroidal*icount*lk
#endif  
  
#ifdef _OPENACC
  !$acc parallel loop gang vector present(allzeta,phifluxpartd) collapse(4)
#else  
  !$omp parallel do private(jpe,j,i,k,jt,indt,indp1,indp)
#endif  
  do jpe=0,mtoroidal-1
    do j=1,mteach
      do k=1,mzeach
        do i=1,mpsi
          jt=jpe*mteach+j
          indt=(j-1)*mpsi          
          indp1=indt+(k-1)*mpsi*mteach
          indp=indp1+i
          phifluxpartd(i,jt,k)=allzeta(indp+jpe*icount)          
        enddo
      enddo
    enddo
  enddo
#ifdef TIME_MPI
  call timer_start()
#endif

#ifdef CUDA_MPI  
  !$acc host_data use_device(phifluxpartd,phiflux)
#else
  !$acc update host(phifluxpartd(:,:,mzeach))
#endif
  call MPI_ALLGATHER(phifluxpartd(:,:,mzeach),mpsi*mteachpd,mpi_Rsize,phiflux,mpsi*mteachpd,mpi_Rsize,partd_comm,ierror)
#ifdef CUDA_MPI  
  !$acc end host_data
#else
  !$acc update device(phiflux)
#endif  
#ifdef TIME_MPI
  filter_allgather_time=filter_allgather_time+timer_get_et()  
  filter_allgather_count=filter_allgather_count+1
  filter_allgather_size=filter_allgather_size+npartdom*mpsi*mteachpd*lk
#endif    
  
  if(nfilter==2)then ! filter poloidal modes
    ffilter=0.0_lk
    ffilter(mmodes+1)=1.0_lk
    !$acc update device(ffilter)

#ifdef _OPENACC
    !$acc host_data use_device(phiflux,ffilter)
    call cufft_filter2(phiflux,ffilter,mtdiag,mpsi)
    !$acc end host_data
#else
    do i=1,mpsi
      do j=1,mtdiag
        xz(j)=phiflux(i,j)
      enddo
      call fftr1d(1,mtdiag,1.0_lk,xz,yz,2)
      yz=ffilter*yz
      call fftr1d(-1,mtdiag,1.0_lk,xz,yz,2)
      do j=1,mtdiag
        phiflux(i,j)=xz(j)
      enddo
    enddo
#endif 
  endif

  ! interpolate field from magnetic coordinates to fieldline coordinates
#ifdef _OPENACC
  !$acc parallel loop gang present(igrid,mtheta,qtinv,deltat,phiflux,farray)
#else
  !$omp parallel do private(i,j,ii,tdum,jt,wt,jtp)
#endif  
  do i=1,mpsi
    ii=igrid(i)
    !$acc loop vector
    do j=1,mtheta(i)
      if(fielddir==1 .or. fielddir==3)then
        tdum=pi2_inv*modulo(deltat(i)*real(j)+(zeta1-pi2)*qtinv(i),pi2)*real(mtdiag)
      else
        tdum=pi2_inv*modulo(deltat(i)*real(j)+zeta1*qtinv(i),pi2)*real(mtdiag)
      endif
      jt=max(0,min(mtdiag-1,int(tdum)))
      wt=tdum-real(jt)
      jtp=jt+1
      if(jt==0)jt=mtdiag

      farray(1,ii+j)=wt*phiflux(i,jtp)+(1.0-wt)*phiflux(i,jt)
    enddo
  enddo

  call periodicity_gpu(farray)

end subroutine filter_t
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
subroutine mnfilter(farray)
  use system_env
  use precision, only: mpi_Rsize
  use global_parameters
  use field_array
  implicit none

  integer i,j,k,ii,idx,m_start,m_end,ij,ip,jt,kz,jpe,indp,indt,indp1,indt1,meachtheta,mteach,mzeach,jtp,&
       icount,idest,isource,isendtag,l,irecvtag,joffset,mteachpd,&
       istatus(MPI_STATUS_SIZE)
  real(lk) wt,r,dt,wz,zdum,tdum,ptemp,sendl(mgrid),sendr(mgrid),recvl(mgrid),recvr(mgrid),&
       phism(mgrid),pleft(mthetamax),pright(mthetamax),phitmp(mgrid),ffilter(mtdiag/2+1),&
       xz(mtdiag),farray(0:1,mgrid),&
       phifluxpartd(mtdiag/mtoroidal,mpsi,mtdiag/npartdom),&
       phifluxlz(mpsi,mtdiag/numberpe,mtdiag),&
       phiflux(mpsi,mtdiag),&
       allzeta((mtdiag/numberpe)*mtdiag*mpsi)

!by zhs
  integer mype_old,mark_array(num_modes),localpe(num_modes),localmt(num_modes),label_array(0:1000)
  real(lk) field2d(mtdiag,mtdiag,mpsi/mtoroidal),ffiltern(mtdiag/2+1),ffilterm(mtdiag),ffiltern0(mtdiag)

  complex(lk) yz(mtdiag/2+1),zz(mtdiag),cfield2d1(mpsi,mtdiag/numberpe,mtdiag),&
  cfield2d2(mpsi,mtdiag/numberpe,mtdiag),cphifluxpartd(mtdiag/mtoroidal,mpsi,mtdiag/npartdom),&
  allcfield(mtdiag*mtdiag*mpsi/numberpe)

! mesh in magnetic coordinates
  mzeach=mtdiag/mtoroidal
  mteach=mzeach/npartdom
  mteachpd=(mtdiag/npartdom)
  dt=2.0_lk*pi/real(mtdiag)

  ffilter=0.0_lk
  ffiltern=0.0_lk
  ffilterm=0.0_lk
  ffiltern0=0.0_lk
  ffilter(nmodes+1)=1.0_lk
  ffiltern(nmodes+1)=1.0_lk
!  ffiltern(1)=1.0    !keep the n=0 mode
!  ffiltern0(mtdiag-1)=1.0 !keep a few m mode for n=0
!  ffiltern0(3)=1.0

!by fks
!first step, initilize mark_array for selected unique n in nmodes,
!eg. if nmodes=(7,8,9,10,10,10,10,11),  mark_array=(1,1,1,1,0,0,0,1)
!the dupicated n is only used one time for the poloidal filter  
label_array=0
mark_array=0

do i=1,num_modes
   j=nmodes(i)
   if(label_array(j)==0)then
       label_array(j)=1
       mark_array(i)=1
   endif
enddo

!set the localmt and localpe array according to nmodes
do i=1,num_modes
! explicit set localmt array
  localmt(i)=mod(nmodes(i)+1,mteach)
  localpe(i)=(nmodes(i)+1-localmt(i))/mteach
  if(localmt(i) == 0)then
     localmt(i)=mteach
     localpe(i)=localpe(i)-1
  endif
enddo

!transback to the old mype MPI layout since the mnfilter subroutine uses the old form by default
!old:
!  particle_domain_location=mod(mype,npartdom)
!  toroidal_domain_location=mype/npartdom
!new:
!  particle_domain_location=mype/mtoroidal
!  toroidal_domain_location=mod(mype,mtoroidal)

mype_old=particle_domain_location+npartdom*toroidal_domain_location


  allzeta=0.0_lk

  joffset=mteachpd*myrank_partd

! Interpolate on a flux surface from fieldline coordinates to magnetic
! coordinates. Use mtdiag for both poloidal and toroidal grid points.
!$omp parallel do private(kz,i,j,wz,ii,zdum,tdum,jt,wt)
  do kz=1,mzeach
     wz=real(kz)/real(mzeach)
     zdum=zeta0+wz*deltaz
     do i=1,mpsi
        ii=igrid(i)
        do j=1,mteachpd
           if(fielddir==1 .or. fielddir==3)then
             tdum=pi2_inv*modulo(dt*real(j+joffset)-(zdum-pi2)*qtinv(i),pi2)*real(mtheta(i))
           else
             tdum=pi2_inv*modulo(dt*real(j+joffset)-zdum*qtinv(i),pi2)*real(mtheta(i))
           endif
           jt=max(0,min(mtheta(i)-1,int(tdum)))
           wt=tdum-real(jt)
           phifluxpartd(kz,i,j)=((1.0_lk-wt)*farray(1,ii+jt)+wt*farray(1,ii+jt+1))*&
                wz+(1.0_lk-wz)*((1.0_lk-wt)*farray(0,ii+jt)+wt*farray(0,ii+jt+1))
        enddo
     enddo
  enddo

! transpose 2-d matrix from (mtoroidal,mzeach) to (1,mtoroidal*mzeach)
  allzeta=0.0_lk
  icount=mteach*mzeach*mpsi
 !$omp parallel do private(jpe,j,i,k,jt,indt,indp1,indp)
   do jpe=0,mtoroidal-1
      do j=1,mteach
         jt=jpe*mteach+j
         indt=(j-1)*mpsi
         do k=1,mzeach
            indp1=indt+(k-1)*mpsi*mteach
            do i=1,mpsi
               indp=indp1+i
               allzeta(indp+jpe*icount)=phifluxpartd(k,i,jt)
            enddo
         enddo
      enddo
   enddo

  call MPI_ALLTOALL(allzeta,icount,mpi_Rsize,phifluxlz,icount,&
          mpi_Rsize,toroidal_comm,ierror)

! transform to k space
  yz=0.0_lk
  cfield2d1=0.0_lk
!$omp parallel do private(j,i,kz,k,indt1,indt,indp,xz,yz)
  do j=1,mteach
     do i=1,mpsi
        xz=phifluxlz(i,j,:)
        call fftr1d(1,mtdiag,1.0_lk,xz,yz,2)
! linear run only keep a few modes
        yz=ffiltern*yz
        do k=1,mtdiag
           if(k<mtdiag/2+2)then
              cfield2d1(i,j,k)=yz(k)
           else
              cfield2d1(i,j,k)=0.0_lk
           endif
        enddo
     enddo
  enddo

  allcfield=0.0_lk  !similar to allzeta structure
  call MPI_ALLTOALL(cfield2d1,icount,mpi_complex,allcfield,&
          icount,mpi_complex,toroidal_comm,ierror)

!transform back to (mtoroidal,mzetach)
  cphifluxpartd=0.0_lk  !similar to phifluxpartd
  !$omp parallel do private(jpe,j,i,k,jt,indt,indp1,indp)
  do jpe=0,mtoroidal-1
    do j=1,mteach
      jt=jpe*mteach+j
      indt=(j-1)*mpsi
      do k=1,mzeach
        indp1=indt+(k-1)*mpsi*mteach
        do i=1,mpsi
          indp=indp1+i
          cphifluxpartd(k,i,jt)=allcfield(indp+jpe*icount)
        enddo
      enddo
    enddo
  enddo

! transpose 2-d matrix from (mtoroidal,mzeach) to (mtoroidal*mzeach,1)
  allcfield=0.0_lk
  icount=mteach*mteachpd*mpsi !here replace mzeach by mteachpd
 !$omp parallel do private(jpe,j,i,k,jt,indt,indp1,indp)
   do jpe=0,npartdom-1      ! zeta in one pe is divided by jpe(kpe)
      do j=1,mteach         ! zeta in one jpe is mteach(j as k)
         jt=jpe*mteach+j  
         indt=(j-1)*mpsi
         do k=1,mteachpd    ! theta in one partdom (k as j)
            indp1=indt+(k-1)*mpsi*mteach
            do i=1,mpsi
               indp=indp1+i
               allcfield(indp+jpe*icount)=cphifluxpartd(jt,i,k)
            enddo
         enddo
      enddo
   enddo

  cfield2d1=0.0_lk
  call MPI_ALLTOALL(allcfield,icount,mpi_complex,cfield2d1,&
          icount,mpi_complex,partd_comm,ierror)

! transform from k to k space
  zz=0.0_lk
  cfield2d2=0.0_lk
!~~~~~~~~~~~~~~~~~~new scheme for fft in 2d,valid for multi mode~~~~~~~~~~~~~
do idx=1,num_modes
! the duplicated n in nmodes will be only processed once
  if(mype_old==localpe(idx) .and. mark_array(idx)==1)then
     j=localmt(idx)
     ! reuse ffilterm
     ffilterm=0.0_lk
	 
	 ! the resonant mode, m=n*q
     ii = nint(nmodes(idx)*qmesh(diag_flux))
	 
	if(nfilter .GE. 3)then
        ! for nfilter>=3, only keep (nfilter-3)*2+1 modes centered at the resonant mode,
        ! change m_start and m_end for your own choice of m if needed
        m_start=max(0,ii-(nfilter-3))
        m_end=ii+(nfilter-3)
        ffilterm(mtdiag+1-m_end:mtdiag+1-m_start)=1.0_lk
     endif
	 
!$omp parallel do private(i,k,zz)
     do i=1,mpsi
        do k=1,mtdiag
           zz(k)=cfield2d1(i,j,k)
        enddo
!fft in m direction
        call fftc1d(1,mtdiag,1.0_lk,zz)
        do k=1,mtdiag
           zz(k)=ffilterm(k)*zz(k)
        enddo
! transform back to k_n space
        call fftc1d(-1,mtdiag,1.0_lk,zz)
        do k=1,mtdiag
           cfield2d2(i,j,k)=zz(k)
        enddo
     enddo
  endif
enddo  
  
!to control the n=0,m/=0 harmonics in NL simulation
  if(mype_old==0)then
    j=1          
!$omp parallel do private(i,k,zz)
     do i=1,mpsi
        do k=1,mtdiag
           zz(k)=cfield2d1(i,j,k)
        enddo
!fft in m direction
        call fftc1d(1,mtdiag,1.0_lk,zz)
        do k=1,mtdiag
           zz(k)=ffiltern0(k)*zz(k)
        enddo
! transform back to k_n space
        call fftc1d(-1,mtdiag,1.0_lk,zz)
        do k=1,mtdiag
           cfield2d2(i,j,k)=zz(k)
        enddo
     enddo
  endif

!~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

!~~~~~~~~~~~~~~~original scheme for fft in 2d,valid for single n mode~~~~~~
!!$omp parallel do private(j,i,k,zz)
!  do j=1,mteach
!     do i=1,mpsi
!        do k=1,mtdiag
!           zz(k)=cfield2d1(i,j,k)
!        enddo
!!fft in m direction
!        call fftc1d(1,mtdiag,1.0,zz)
!        do k=1,mtdiag
!           zz(k)=ffilterm(k)*zz(k)
!        enddo
!! transform back to k_n space
!        call fftc1d(-1,mtdiag,1.0,zz)
!        do k=1,mtdiag
!           cfield2d2(i,j,k)=zz(k)
!        enddo
!     enddo
!  enddo
!~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

!transform back from (mtoroidal*mzeach,1) to (mtoroidal,mzeach)
  allcfield=0.0_lk
  call MPI_ALLTOALL(cfield2d2,icount,mpi_complex,allcfield,&
          icount,mpi_complex,partd_comm,ierror)

 !$omp parallel do private(jpe,j,i,k,jt,indt,indp1,indp)
   do jpe=0,npartdom-1      ! zeta in one pe is divided by jpe(kpe)
      do j=1,mteach         ! zeta in one jpe is mteach(j as k)
         jt=jpe*mteach+j
         indt=(j-1)*mpsi
         do k=1,mteachpd    ! theta in one partdom (k as j)
            indp1=indt+(k-1)*mpsi*mteach
            do i=1,mpsi
               indp=indp1+i
               cphifluxpartd(jt,i,k)=allcfield(indp+jpe*icount)
            enddo
         enddo
      enddo
   enddo

!transform back from (mtoroidal,mzeach) to (1,mtoroidal*mzeach)
  icount=mteach*mzeach*mpsi
  !$omp parallel do private(jpe,j,i,k,jt,indt,indp1,indp)
  do jpe=0,mtoroidal-1
    do j=1,mteach
      jt=jpe*mteach+j
      indt=(j-1)*mpsi
      do k=1,mzeach
        indp1=indt+(k-1)*mpsi*mteach
        do i=1,mpsi
          indp=indp1+i
          allcfield(indp+jpe*icount)=cphifluxpartd(k,i,jt)
        enddo
      enddo
    enddo
  enddo

  cfield2d1=0.0_lk
  call MPI_ALLTOALL(allcfield,icount,mpi_complex,cfield2d1,&
          icount,mpi_complex,toroidal_comm,ierror)

! transform back to real space
  xz=0.0_lk
!$omp parallel do private(j,i,kz,k,xz,yz)
  do j=1,mteach
     do i=1,mpsi
        do k=1,mtdiag/2+1
           yz(k)=cfield2d1(i,j,k)
        enddo
        call fftr1d(-1,mtdiag,1.0_lk,xz,yz,2)
        phifluxlz(i,j,:)=xz
     enddo
  enddo

  call MPI_ALLTOALL(phifluxlz,icount,mpi_Rsize,allzeta,icount,&
          mpi_Rsize,toroidal_comm,ierror)

! transpose 2-d matrix from (1,mtoroidal*mzeach) to (mtoroidal,mzeach)
 !$omp parallel do private(jpe,j,i,k,jt,indt,indp1,indp)
   do jpe=0,mtoroidal-1
      do j=1,mteach
         jt=jpe*mteach+j
         indt=(j-1)*mpsi
         do k=1,mzeach
            indp1=indt+(k-1)*mpsi*mteach
            do i=1,mpsi
               indp=indp1+i
               phifluxpartd(k,i,jt)=allzeta(indp+jpe*icount)
            enddo
         enddo
      enddo
   enddo

  call MPI_ALLGATHER(phifluxpartd(mzeach,:,:),mpsi*mteachpd,mpi_Rsize,phiflux,mpsi*mteachpd,mpi_Rsize,partd_comm,ierror)

! interpolate field from magnetic coordinates to fieldline coordinates
!$omp parallel do private(i,j,ii,tdum,jt,wt,jtp)
  do i=1,mpsi
     ii=igrid(i)
     do j=1,mtheta(i)
        if(fielddir==1 .or. fielddir==3)then
          tdum=pi2_inv*modulo(deltat(i)*real(j)+(zeta1-pi2)*qtinv(i),pi2)*real(mtdiag)
        else
          tdum=pi2_inv*modulo(deltat(i)*real(j)+zeta1*qtinv(i),pi2)*real(mtdiag)
        endif
        jt=max(0,min(mtdiag-1,int(tdum)))
        wt=tdum-real(jt)
        jtp=jt+1
        if(jt==0)jt=mtdiag

        farray(1,ii+j)=wt*phiflux(i,jtp)+(1.0_lk-wt)*phiflux(i,jt)
     enddo
  enddo

  call periodicity_cpu(farray)

end subroutine mnfilter

subroutine mfilter(farray)
  use system_env
  use precision, only: mpi_Rsize
  use global_parameters
  use field_array
  implicit none

  integer i,j,k,ii,ij,ip,jt,kz,ipe,kpe,indp,indt,indp1,indt1,meachtheta,mteach,mzeach,jtp,&
       icount,idest,isource,isendtag,l,irecvtag,istatus(MPI_STATUS_SIZE)
  real(lk) wt,r,dt,wz,zdum,tdum,ptemp,sendl(mgrid),sendr(mgrid),recvl(mgrid),recvr(mgrid),&
       phism(mgrid),pleft(mthetamax),pright(mthetamax),phitmp(mgrid),ffilter(mtdiag/2+1),&
       xz(mtdiag),farray(0:1,mgrid),phiflux(mtdiag/mtoroidal,mtdiag,mpsi),&
       allzeta2(mtdiag*mtdiag*mpsi/mtoroidal),allzeta(mtdiag*mtdiag*mpsi/mtoroidal)
       
!by zhs
  integer mpeach
  real(lk) field2d(mtdiag,mtdiag,mpsi/mtoroidal),ffilters(mtdiag/2+1,mtdiag)

  complex(lk) yz(mtdiag/2+1),zz(mtdiag),cfield2d(mtdiag/2+1,mtdiag)

! mesh in magnetic coordinates
  mzeach=mtdiag/mtoroidal
  mteach=mzeach
  mpeach=mpsi/mtoroidal

  dt=2.0_lk*pi/real(mtdiag)

  ffilter=0.0_lk
  ffilters=0.0_lk
!by zhs use the first mode. when nfilter>3, mult m
  ffilter(nmodes(1)+1)=1.0_lk
  ffilters(nmodes(1)+1,mtdiag+1-mmodes(1))=1.0_lk
  if(nfilter>3)then  !nfilter>3, include m-1,m+1 in the simulation
    ffilters(nmodes(1)+1,mtdiag-mmodes(1))=1.0_lk
    ffilters(nmodes(1)+1,mtdiag+2-mmodes(1))=1.0_lk
  endif
!  ffilters(nmodes(1)+1,mtdiag-1-mmodes(1))=1.0
!  ffilters(nmodes(1)+1,mtdiag+3-mmodes(1))=1.0

  field2d=0.0_lk
  cfield2d=0.0_lk
! Interpolate on a flux surface from fieldline coordinates to magnetic
! coordinates. Use mtdiag for both poloidal and toroidal grid points.
!$omp parallel do private(kz,i,j,wz,ii,zdum,tdum,jt,wt)
  do kz=1,mzeach
     wz=real(kz)/real(mzeach)
     zdum=zeta0+wz*deltaz
     do i=1,mpsi
        ii=igrid(i)
        do j=1,mtdiag
           if(fielddir==1 .or. fielddir==3)then
             tdum=pi2_inv*modulo(dt*real(j)-(zdum-pi2)*qtinv(i),pi2)*real(mtheta(i))
           else
             tdum=pi2_inv*modulo(dt*real(j)-zdum*qtinv(i),pi2)*real(mtheta(i))
           endif
           jt=max(0,min(mtheta(i)-1,int(tdum)))
           wt=tdum-real(jt)

           phiflux(kz,j,i)=((1.0_lk-wt)*farray(1,ii+jt)+wt*farray(1,ii+jt+1))*&
                wz+(1.0_lk-wz)*((1.0_lk-wt)*farray(0,ii+jt)+wt*farray(0,ii+jt+1))
        enddo
     enddo
  enddo

! transpose 2-d matrix from (mtoroidal,mzeach) to (1,mtoroidal*mzeach)
  allzeta=0.0_lk
  icount=mtdiag*mzeach*mpeach

!$omp parallel do private(ipe,j,i,k,ip,indp,indp1,indt)
  do ipe=0,mtoroidal-1
     do i=1,mpeach
        ip=mpeach*ipe+i
        indp1=(i-1)*mtdiag*mzeach
        do j=1,mtdiag
           indt=indp1+(j-1)*mzeach
           do k=1,mzeach
              indp=indt+k
              allzeta(indp+ipe*icount)=phiflux(k,j,ip)
           enddo
       enddo
     enddo
  enddo
     call MPI_ALLTOALL(allzeta,icount,mpi_Rsize,allzeta2,icount,&
          mpi_Rsize,toroidal_comm,ierror)


!transpose 1-d array (mtdiag*mtdiag,mpeach) allzeta to 2-d array field2d(mtdiag,mtdiag,mpeach)
!$omp parallel do private(j,i,kz,k,kpe,indp,indp1,indt)
  do kpe=0,mtoroidal-1
     do i=1,mpeach
        indp1=kpe*icount+(i-1)*mtdiag*mzeach
        do j=1,mtdiag
           indt=indp1+(j-1)*mzeach
           do k=1,mzeach
              indp=indt+k
              kz=kpe*mzeach+k
              field2d(kz,j,i)=allzeta2(indp)
           enddo
        enddo
     enddo
  enddo

! transform to k space
  yz=0.0_lk
  zz=0.0_lk

!filter n and m mode
!$omp parallel do private(j,i,k,xz,yz,zz,cfield2d)
     do i=1,mpeach
        do j=1,mtdiag
           do k=1,mtdiag
                 xz(k)=field2d(k,j,i)
           enddo
! fft in n direction
           call fftr1d(1,mtdiag,1.0_lk,xz,yz,2)
           do k=1,mtdiag/2+1
              cfield2d(k,j)=yz(k)
           enddo
        enddo
        do k=1,mtdiag/2+1
           do j=1,mtdiag
              zz(j)=cfield2d(k,j)
           enddo
! fft in m direction
           call fftc1d(1,mtdiag,1.0_lk,zz)
           do j=1,mtdiag
              cfield2d(k,j)=zz(j)
           enddo
        enddo

        cfield2d(:,:)=ffilters(:,:)*cfield2d(:,:)
! tranform back to real space
        do k=1,mtdiag/2+1
           do j=1,mtdiag
              zz(j)=cfield2d(k,j)
           enddo
           call fftc1d(-1,mtdiag,1.0_lk,zz)
           do j=1,mtdiag
              cfield2d(k,j)=zz(j)
           enddo
        enddo

        do j=1,mtdiag
           do k=1,mtdiag/2+1
              yz(k)=cfield2d(k,j)
           enddo
           call fftr1d(-1,mtdiag,1.0_lk,xz,yz,2)
           do k=1,mtdiag
              field2d(k,j,i)=xz(k)
           enddo
        enddo
     enddo

! tranpose from field2d back to allzeta
!$omp parallel do private(j,i,kz,k,kpe,indp,indt,indp1)
  do kpe=0,mtoroidal-1
     do i=1,mpeach
        indp1=kpe*icount+(i-1)*mtdiag*mzeach
        do j=1,mtdiag
           indt=indp1+(j-1)*mzeach
           do k=1,mzeach
              indp=indt+k
              kz=kpe*mzeach+k
              allzeta2(indp)=field2d(kz,j,i)
           enddo
        enddo
     enddo
  enddo

  call MPI_ALLTOALL(allzeta2,icount,mpi_Rsize,allzeta,&
          icount,mpi_Rsize,toroidal_comm,ierror)

!$omp parallel do private(ipe,j,i,k,ip,indp,indt,indp1)
  do ipe=0,mtoroidal-1
     do i=1,mpeach
        ip=mpeach*ipe+i
        indp1=(i-1)*mtdiag*mzeach
        do j=1,mtdiag
           indt=indp1+(j-1)*mzeach
           do k=1,mzeach
              indp=indt+k
              phiflux(k,j,ip)=allzeta(indp+ipe*icount)
           enddo
        enddo
     enddo
  enddo

! interpolate field from magnetic coordinates to fieldline coordinates
!$omp parallel do private(i,j,ii,tdum,jt,wt,jtp)
  do i=1,mpsi
     ii=igrid(i)
     do j=1,mtheta(i)
        if(fielddir==1 .or. fielddir==3)then
          tdum=pi2_inv*modulo(deltat(i)*real(j)+(zeta1-pi2)*qtinv(i),pi2)*real(mtdiag)
        else
          tdum=pi2_inv*modulo(deltat(i)*real(j)+zeta1*qtinv(i),pi2)*real(mtdiag)
        endif
        jt=max(0,min(mtdiag-1,int(tdum)))
        wt=tdum-real(jt)
        jtp=jt+1
        if(jt==0)jt=mtdiag

        farray(1,ii+j)=wt*phiflux(mzeach,jtp,i)+(1.0_lk-wt)*phiflux(mzeach,jt,i)
     enddo
  enddo

  call periodicity_cpu(farray)

end subroutine mfilter
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!


!> ---------------------------------------------------------------------------
#ifdef ADIOS2
!> ---------------------------------------------------------------------------

!> ---------------------------------------------------------------------------
!> opendiag for the history file
!> ---------------------------------------------------------------------------
subroutine opendiag_history_adios2(mpdiag,nfield,mfdiag,modes)
  use adios2
  use adios2_common
  use adios2_history_diagnostics
  use mpi
  use global_parameters
  implicit none
  
  integer, intent(in)           :: mpdiag,nfield,mfdiag,modes
  integer                       :: ierr, prev_num_steps
  real(lk)                      :: tstep_tmp
  character(len=128)            :: src
  type(adios2_engine)           :: r_engine

  ! Open history file for writing
  call adios2_open (adios2_diag_hist_e, adios2_diag_hist_io, "history.bp", &
    adios2_mode_write, mpi_comm_self, ierr)

  ! Open the previous history restart file
  if(irun .gt. 0) then
    
    if(mod(irest-1,2)==0)then
      src = "./restart_dir1/history_restart.bp"//char(0)
    else
      src = "./restart_dir2/history_restart.bp"//char(0)
    endif

    ! Open restart file
    call adios2_open (r_engine, adios2_diag_hist_io, trim(src), adios2_mode_read, &
      mpi_comm_self, ierr)

    ! Copy history_restart.bp into history.bp
    call copy_data (r_engine, adios2_diag_hist_e, adios2_diag_hist_io, t_hist_diag, prev_num_steps)
    mstepall = prev_num_steps

    ! Close the restart file
    call adios2_close (r_engine, ierr)
  
  else ! if irun .lt. 0

    ! Define vars in the adios io object
    call define_vars (t_hist_diag, adios2_diag_hist_io, mpdiag, nfield, mfdiag, modes)
  
  endif

end subroutine opendiag_history_adios2


!> ---------------------------------------------------------------------------
!> opendiag for the data1d file
!> ---------------------------------------------------------------------------
subroutine opendiag_data1d_adios2 (nfield, mpdata1d, mfdata1d)
  use adios2
  use adios2_common
  use adios2_data1d_diagnostics
  use mpi
  use global_parameters
  use particle_array, only: fload, feload
  implicit none
  
  integer, intent(in)           :: nfield, mpdata1d, mfdata1d
  integer                       :: ierr
  character(len=128)            :: src
  type(adios2_engine)           :: r_engine

  ! Open data1d file for writing
  call adios2_open (adios2_diag_data1d_e, adios2_diag_data1d_io, "data1d.bp", &
    adios2_mode_write, mpi_comm_self, ierr)

  ! Open the previous data1d restart file
  if(irun .gt. 0) then
    
    if(mod(irest-1,2)==0)then
      src = "./restart_dir1/data1d_restart.bp"//char(0)
    else
      src = "./restart_dir2/data1d_restart.bp"//char(0)
    endif

    ! Open restart file
    call adios2_open (r_engine, adios2_diag_data1d_io, trim(src), adios2_mode_read, &
      mpi_comm_self, ierr)

    ! Copy data1d_restart.bp into data1d.bp
    call copy_data (r_engine, adios2_diag_data1d_e, adios2_diag_data1d_io, t_data1d_diag)

    ! Close the restart file
    call adios2_close (r_engine, ierr)
  
  else ! if irun .lt. 0

    ! Define vars in the adios io object
    call define_vars (t_data1d_diag, adios2_diag_data1d_io, nspecies, mpsi, &
      nfield, nhybrid, mpdata1d, mfdata1d, fload, feload)
  
  endif

end subroutine opendiag_data1d_adios2


!> ---------------------------------------------------------------------------
!> generic opendiag implementation for adios
!> ---------------------------------------------------------------------------
subroutine opendiag_adios2(mpdiag,nfield,mfdiag,modes,mpdata1d,mfdata1d)
  use global_parameters
  use particle_array, only: fload, feload
  implicit none

  integer, intent(in) :: mpdiag,nfield,mfdiag,modes,mpdata1d,mfdata1d
  integer             :: tmpfd, prev_diag_start_step,stat
  character(len=32)   :: prev_diag_start_step_str

  if(irun==0) mstepall=0
  
  call opendiag_history_adios2 (mpdiag, nfield, mfdiag, modes)
  call opendiag_data1d_adios2  (nfield, mpdata1d, mfdata1d)

  ! calculate current total time step        
  if (irun .gt. 0) mstepall=mstepall*ndiag

end subroutine opendiag_adios2

!> ---------------------------------------------------------------------------
! End of the ifdef ADIOS2 block
#endif
!> ---------------------------------------------------------------------------


!> ---------------------------------------------------------------------------
subroutine opendiag(mpdiag,nfield,mfdiag,modes,mpdata1d,mfdata1d)
  use global_parameters
  implicit none

  integer mpdiag,nfield,mfdiag,modes,mpdata1d,mfdata1d
  integer i,j,iotmp,ndstep,ndata
  real(lk) tdum
  character(len=100) cdum

  ! # of time steps and # of data added to history file
  ndstep=mstep/ndiag ! no. of diagnosis steps in this file, this is a local variable
  ndata=(nspecies*mpdiag+nfield*(2*modes+mfdiag))

#ifdef ADIOS2
  call opendiag_adios2(mpdiag,nfield,mfdiag,modes,mpdata1d,mfdata1d)
! #else
#endif

  iodiag=123
  iotmp=456
  iodata1d=789

  if(irun==0)then
     mstepall=0
     ! open time history file in new run
     open(iodiag,file='history.out',status='replace')
     write(iodiag,101)ndstep,nspecies,mpdiag,nfield,modes,mfdiag
     write(iodiag,102)tstep*ndiag

     ! open output file for radial-time data in new run
     open(iodata1d,file='data1d.out',status='replace')            
     write(iodata1d,101)ndstep,mpsi+1,nspecies,nhybrid,mpdata1d,nfield,mfdata1d

  else ! if irun!=0

     ! find history file from previous run
     irest=irest-1
     if(mod(irest,2)==0)then
        cdum="restart_dir1/history_restart.out"
     else
        cdum="restart_dir2/history_restart.out"
     endif

     open(iotmp,file=trim(cdum),status='old')
     rewind(iotmp)

     ! open time history file for restart run
     open(iodiag,file='history.out',status='replace')

     ! # of time steps
     read(iotmp,101)mstepall
     write(iodiag,101)mstepall+ndstep
     do i=1,5
        read(iotmp,101)j
        write(iodiag,101)j
     enddo
     
     ! copy restart history data file
     do i=0,ndata*mstepall
        read(iotmp,102)tdum
        write(iodiag,102)tdum
     enddo
     close(iotmp)

     !!copy restart data1d.out file
     open(iodata1d,file='data1d.out',status='replace')
     if(mod(irest,2)==0)then
        cdum="restart_dir1/data1d_restart.out"
     else
        cdum="restart_dir2/data1d_restart.out"
     endif
     
     irest=irest+1
     
     open(iotmp,file=trim(cdum),status='old')
     rewind(iotmp)
     read(iotmp,101)j
     write(iodata1d,101)mstepall+ndstep
     do i=1,6
        read(iotmp,101)j
        write(iodata1d,101)j
     enddo

     ndata=(mpsi+1)*(nspecies*mpdata1d+nfield*mfdata1d)
     do i=1,ndata*mstepall
        read(iotmp,102)tdum
        write(iodata1d,102)tdum
     enddo
     close(iotmp)

     !!calculate current total time step        
     mstepall=mstepall*ndiag
     
  endif !if irun==0
! #endif

101 format(i6)
102 format(e13.6)

end subroutine opendiag
