! 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.

Subroutine dataout3d
  use precision
  use global_parameters
  use field_array
  use equilibrium
  use spline_function, only: spx, spz, spb,spf
  use system_env, only: MPI_SUM, MPI_COMM_WORLD, ierror
  implicit none
! #ifndef ADIOS2
  include 'netcdf.inc'
! #endif

  integer i,j,k,n,ij,istatus,ncid,dataid2(3),dimid(2),dataid1(2)
  real(lk) data3d(mgrid),radial(mpsi+1),&
       r,theta,zeta,theta0
  character(len=1) cdum(3)
  character(len=9) vdum
  character(len=100) filedum
!!! variables needed for data_3d section, added by Javier
  character(len=64) cdum0,cdum1
  integer :: nprint
! HDF5 declarations

#ifndef ADIOS2
! we don't need to gather data to mype==0 when turn on ADIOS2
  real(lk) :: bdum(mgrid,mtoroidal),xdum(mgrid,mtoroidal),zdum(mgrid,mtoroidal),ydum(mgrid,mtoroidal),ndum(mgrid,mtoroidal)
#endif
  real(lk) :: bdumdom(mgrid),xdumdom(mgrid),ydumdom(mgrid),zdumdom(mgrid),ndumdom(mgrid)
  integer ii,jt,kj,icount,ipsi
  real(lk) :: pdum,tdum,zetadum,rdum,fdum


!!!!!!!!!!!!!
!!! Sec 0: data_3d added by Javierhn. 
!!!        First of all, prepare 3d mesh and potential variables for output.
!!!!!!!!!!!!!

  !javierhn: 3D writting of mesh and potential 
  !The firt time, the mesh and 'b' are printed. Then, every time we called this
  !funtion the values of phi are printed in a new file (mesh is not printed
  !again). There are two kinds of output: volume and flux surface.
  !TODO:Printing all this data is computationally expensive. ADIOS or some MPI
  !writing should be considered.
 
  ipsi=diag_flux !radial location of the flux surface data

!!! part 1: compute mesh (rdum,zdum,fdum) and the magnetic field bdum at the first diag step.

  if (.not. data3d_init_output) then  !only write the mesh and equilibirum Bfield the first time
    !each quantity is calculated only in its toroidal domain-> quantitydom
    !Then we will gather all of them in to the total quantity
    bdumdom=0.0_lk
    zdumdom=0.0_lk
#ifndef ADIOS2
! we don't need to gather data to mype==0 when turn on ADIOS2
    bdum=0.0_lk
    xdum=0.0_lk
    ydum=0.0_lk
    zdum=0.0_lk
#endif

    !TODO: OPENACC. But it needs to declare the functions
    !TODO: check tdum when fielddir=1 or 3
    !$omp parallel do private(i,j,ij,zetadum,tdum,pdum,rdum,fdum)
    do i = 0, mpsi
       pdum=psimesh(i)
       do j=0,mtheta(i)
          tdum=pi2*real(j)/real(mtheta(i))
          zetadum=zeta1
          if(fielddir==1 .or. fielddir==3)then
            tdum=modulo(tdum + (zetadum-pi2)*qtinv(i),pi2)
          else
            tdum=modulo(tdum + zetadum*qtinv(i),pi2)
          endif

          ij=igrid(i)+j
! 2D cases:
#ifndef _TOROIDAL3D
          bdumdom(ij)=spb(pdum,tdum)  !magnetic field
          rdum=spx(pdum,tdum)  !Major radius (toroidal coord)
          zdumdom(ij)=spz(pdum,tdum)  !Zeta (toroidal coordinates)
          fdum=spf(pdum,tdum) + zetadum  !phi (toroidal coordinates)
! 3D
#else
          bdumdom(ij)=spb(pdum,tdum,zetadum)  
          rdum=spx(pdum,tdum,zetadum)  
          zdumdom(ij)=spz(pdum,tdum,zetadum)  
          fdum=spf(pdum,tdum,zetadum)+zetadum
#endif
          if((fielddir==2).or.(fielddir==3))fdum=-fdum

          !translate to cartesian coordinates(X,Y,Z)
          xdumdom(ij)=rdum*cos(fdum)
          ydumdom(ij)=rdum*sin(fdum)
          !Z was already calculated

        enddo
     enddo
  endif

!!! part 2: compute the potential variable ndum for every diag step.

    !each time we write the potential (not the mesh!)

    ndumdom=0.0_lk
#ifndef ADIOS2
! we don't need to gather data to mype==0 when turn on ADIOS2
    ndum=0.0_lk
#endif

!interpolation disabled
!    !$omp parallel do private(i,j,k,wz,zetadum,tdum,adum,jt,kj,wt)
!    do i=1,nzeach
!       wz=real(i,lk)/real(nzeach,lk)
!       zetadum=zeta0+wz*deltaz  !zeta angle (Boozer)
!       do j = 1, ntheta
!          tdum=deltheta*real(j-1) ! theta
!          do k=1,mpsi

!              if(fielddir==1 .or. fielddir==3)then
!                adum=modulo(tdum - (zetadum-pi2)*qtinv(k),pi2) ! alpha
!              else
!                adum=modulo(tdum - zetadum*qtinv(k),pi2) ! alpha
!              endif

!              jt=max(0,min(mtheta(k)-1,int(adum/deltat(k))))
!              kj=igrid(k)+jt
!              wt=adum/deltat(k)-real(jt)
!             ndumdom(k,j,i)=((1.0_lk-wt)*phi(1,kj)+wt*phi(1,kj+1))*wz&
!                 +(1.0_lk-wz)*((1.0_lk-wt)*phi(0,kj)+wt*phi(0,kj+1))
!          enddo
!       enddo
!    enddo

    !$omp parallel do private(i,j,ij)
    do i = 0, mpsi
       do j=0,mtheta(i)
          ij=igrid(i)+j
          ndumdom(ij)=phi(1,ij)
     enddo
   enddo

#ifdef ADIOS2
!!!!!!!!!!!!!
!!! Sec 1: adios2 output part, enabled with ADIOS2 option in Makefile.
!!!!!!!!!!!!!

  call dataout3d_adios2(mgrid,bdumdom,xdumdom,ydumdom,zdumdom,ndumdom)

#elif oldNETCDF
!!!!!!!!!!!!!
!!! Sec 2: old netcdf part, disabled now
!!!!!!!!!!!!!

!!!netcdf 3d data
  if(.not. data3d_init_output)then
     if(mype==0)then
        do i=0,mpsi
           radial(i+1)=(rg0+deltar*real(i))/rho0
        enddo
! write run-dimension
        filedum='phi_dir/RUNdimen.ncd'
! open netcdf data file
        istatus=nf_create(trim(filedum),nf_clobber,ncid)

! define data array dimension
        istatus=nf_def_dim(ncid,'scalar',1,dimid(1))
        istatus=nf_def_dim(ncid,'flux-surfaces',mpsi+1,dimid(2))
! define data array id
        istatus=nf_def_var(ncid,'PE-number',nf_int,1,dimid(1),dataid1(1))
        istatus=nf_def_var(ncid,'flux-surface-number',nf_int,1,dimid(1),dataid1(2))
        istatus=nf_def_var(ncid,'radial-grids',nf_real,1,dimid(2),dataid2(1))
        istatus=nf_def_var(ncid,'poloidal-grids',nf_int,1,dimid(2),dataid2(2))
        istatus=nf_def_var(ncid,'index-shift',nf_int,1,dimid(2),dataid2(3))
! end of netcdf definition
        istatus=nf_enddef(ncid)

! write data
        istatus=nf_put_var_int(ncid,dataid1(1),numberpe)
        istatus=nf_put_var_int(ncid,dataid1(2),mpsi+1)
        istatus=nf_put_var_real(ncid,dataid2(1),radial)
        istatus=nf_put_var_int(ncid,dataid2(2),mtheta+1)
        istatus=nf_put_var_int(ncid,dataid2(3),itran)
! check error
        if (istatus .ne. nf_noerr) then
           print *, nf_strerror(istatus)
        endif
! close netcdf data file
        istatus=nf_close(ncid)
     endif !if(mype==0)

! grid coordinates,each PE writes to a separate file 
     cdum(1)='X'
     cdum(2)='Y'
     cdum(3)='Z'
     do n=1,3
        zeta=zeta1
        do i=0,mpsi
           r=rg0+deltar*real(i)
           do j=0,mtheta(i)
              ij=igrid(i)+j
              theta=deltat(i)*real(j)+zeta*qtinv(i)
              theta0=theta
! grid coordinates (x,y,z), use geometry center as the origin 
              if(n==1)then
                 data3d(ij)=cos(zeta)*(1.0+r*cos(theta0))
              elseif(n==2)then
                 data3d(ij)=-sin(zeta)*(1.0+r*cos(theta0))
              else
                 data3d(ij)=r*sin(theta0)
              endif
           enddo
        enddo
        
! coordinate data file name
        if(myrank_toroidal < 10)then
           write(filedum,'("phi_dir/NCD",a1,"coor.00",i1)')cdum(n),myrank_toroidal
        elseif(myrank_toroidal < 100)then
           write(filedum,'("phi_dir/NCD",a1,"coor.0",i2)')cdum(n),myrank_toroidal
        else
           write(filedum,'("phi_dir/NCD",a1,"coor.",i3)')cdum(n),myrank_toroidal
        endif
        
! variable name
        write(vdum,'(a1,"coordina")')cdum(n)

! open netcdf data file
        istatus=nf_create(trim(filedum),nf_clobber,ncid)
! define data array dimension
        istatus=nf_def_dim(ncid,'poloidal_grid',mgrid,dimid(1))
        istatus=nf_def_dim(ncid,'toroidal_grid',1,dimid(2))
! define data array id
        istatus=nf_def_var(ncid,vdum,nf_real,2,dimid,dataid1(1))
! end of netcdf definition
        istatus=nf_enddef(ncid)
! write data
        istatus=nf_put_var_real(ncid,dataid1(1),data3d)
        istatus=nf_close(ncid)        
     enddo
     data3d_init_output = .true.
  endif !if(.not. data3d_init_output)

! potential data file
  if(myrank_partd==0)then
     do i=0,mpsi
        do j=0,mtheta(i)
           ij=igrid(i)+j
           data3d(ij)=phi(1,ij)
        enddo
     enddo
     
     write(filedum,'("phi_dir/PHI_",i0,"_",i0,".ncd")')(mstepall+istep),myrank_toroidal

   ! variable name
     write(vdum,'("Potential")')

   ! open netcdf data file
     istatus=nf_create(trim(filedum),nf_clobber,ncid)
   ! define data array dimension
     istatus=nf_def_dim(ncid,'poloidal_grid',mgrid,dimid(1))
     istatus=nf_def_dim(ncid,'toroidal_grid',1,dimid(2))
   ! define data array id
     istatus=nf_def_var(ncid,vdum,nf_real,2,dimid,dataid1(1))
   ! end of netcdf definition
     istatus=nf_enddef(ncid)
   ! write data
     istatus=nf_put_var_real(ncid,dataid1(1),data3d)
     istatus=nf_close(ncid)
  endif


#else  
!!!!!!!!!!!!!
!!! Sec 3: data_3d added by Javierhn, default data_3d output format, *.out
!!!!!!!!!!!!!

  !javierhn: 3D writting of mesh and potential 
  !The firt time, the mesh and 'b' are printed. Then, every time we called this
  !funtion the values of phi are printed in a new file (mesh is not printed
  !again). There are two kinds of output: volume and flux surface.
  !TODO:Printing all this data is computationally expensive. ADIOS or some MPI
  !writing should be considered.

!!! part 1: output mesh (xdum,ydum,zdum) and the magnetic field bdum

  if (.not. data3d_init_output) then  !only write the mesh and equilibirum Bfield the first time

    ! gather all 3d mesh info into mype==0
    icount=mgrid
    call MPI_GATHER(bdumdom,icount,mpi_Rsize,bdum,icount,mpi_Rsize,0,toroidal_comm,ierror)
    call MPI_GATHER(xdumdom,icount,mpi_Rsize,xdum,icount,mpi_Rsize,0,toroidal_comm,ierror)
    call MPI_GATHER(ydumdom,icount,mpi_Rsize,ydum,icount,mpi_Rsize,0,toroidal_comm,ierror)
    call MPI_GATHER(zdumdom,icount,mpi_Rsize,zdum,icount,mpi_Rsize,0,toroidal_comm,ierror)


    if (mype==0) then
      open(113,file='phi_dir/bmesh_vol.out',status='replace')
      write(113,102) mgrid,mtoroidal
      open(114,file='phi_dir/bmesh_surf.out',status='replace')
      write(114,102) mtheta(ipsi)+1,mtoroidal

      do k=1,mtoroidal
         do i=0,mpsi
            do j=0,mtheta(i)
               ij=igrid(i)+j
               ! write data for all the 3d grids 
               write(113,101) bdum(ij,k),xdum(ij,k),ydum(ij,k),zdum(ij,k)
               ! write data for a specific psi index, i.e. a flux surface only
               if (i==ipsi) write(114,101) bdum(ij,k),xdum(ij,k),ydum(ij,k),zdum(ij,k)
            enddo
         enddo
      enddo
      close(113)
      close(114)
    endif

  endif

!!! part 2: output the potential variable ndum for every diag step.

    ! gather all 3d potential info into mype==0
    icount=mgrid
    call MPI_GATHER(ndumdom,icount,mpi_Rsize,ndum,icount,mpi_Rsize,0,toroidal_comm,ierror)

    if (mype==0) then
      nprint=mstepall+istep
      write(cdum0,'(i7.7,".out")')nprint
      cdum1='phi_dir/phi_vol_'//trim(cdum0)
      open(115,file=cdum1,status='replace')

      write(cdum0,'(i7.7,".out")')nprint
      cdum1='phi_dir/phi_surf_'//trim(cdum0)
      open(116,file=cdum1,status='replace')

      do k=1,mtoroidal
         do i=0,mpsi
            do j=0,mtheta(i)
               ij=igrid(i)+j
               write(115,103) ndum(ij,k)  !write the potential only, mesh has been written above
               if(i==ipsi) write(116,103) ndum(ij,k)  !write potential at specific psi index, i.e. a flux surface only
            enddo
         enddo
      enddo
      close(115)
      close(116)
    endif

!!! synchronize all the MPI progresses
  call MPI_BARRIER(MPI_COMM_WORLD,ierror)


101 format(4(e11.4,1x))
102 format(7i6)
103 format(e11.4)


#endif 
end subroutine dataout3d


#ifdef ADIOS2
!----------------------------------------------------------------------------!
subroutine dataout3d_adios2(mdim,bdumdom,xdumdom,ydumdom,zdumdom,ndumdom)
  use precision, only: lk
  use global_parameters
  use field_array
  use simple_error_handler
  use adios2_common
  implicit none

! declaration of the dummy arguments
  integer :: mdim
  real(lk),dimension(mdim) :: bdumdom,xdumdom,ydumdom,zdumdom,ndumdom
 
! local variables
  integer :: i, ij0, ij1, mdim_surf, mpsi_tmp, ierr
  real(lk) b_3d(mdim), X_coordinates(mdim), Y_coordinates(mdim), Z_coordinates(mdim), data3d_potential(mdim),&
           b_surf(0:mtheta(diag_flux)), X_surf(0:mtheta(diag_flux)), Y_surf(0:mtheta(diag_flux)), Z_surf(0:mtheta(diag_flux)), phi_surf(0:mtheta(diag_flux)),&
           radial(mpsi+1)
  integer :: mtheta_tmp(0:mpsi), current_step

  type(adios2_variable),save :: var_pe_number, var_flux_surface_number, var_radial_grids
  type(adios2_variable),save :: var_poloidal_grids, var_index_shift, var_mgrid, var_tr_comm_size
  type(adios2_variable),save :: var_x_coordinates, var_y_coordinates, var_z_coordinates, var_b_3d
  type(adios2_variable),save :: var_phi, var_igrid, var_step_id
  type(adios2_variable),save :: var_x_surf, var_y_surf, var_z_surf, var_b_surf, var_phi_surf
  type(adios2_attribute)     :: att_flux_num, att_poloidal, att_indexshift, att_3dcoord

  integer   :: shape2(2), tr_comm_size
   
  !!! tr_comm_size records the # of toroidal planes
  tr_comm_size = mtoroidal

  ! data length for the variables on the diag flux surface
  mdim_surf = mtheta(diag_flux)+1

!!!!!!!!!!!!!
!!! Sec 1: define adios2 variables, must be done at first for each restart run
!!!!!!!!!!!!!

  !@TODO: CHECK if this should be istepall or mstepall
  if(.not. adios2_data3d_io_initialized)then
    ! Declare IO object
    call adios2_declare_io (adios2_data3d_io, adios2_obj, adios2_data3d_id, ierr)

    ! Define variables
    if (mype .eq. 0) then
      call adios2_define_variable (var_pe_number, adios2_data3d_io, 'num-processes', adios2_type_integer4, ierr)
      
      call adios2_define_variable (var_flux_surface_number, adios2_data3d_io, 'flux-surface-number', adios2_type_integer4, ierr)
      
      call adios2_define_variable (var_radial_grids, adios2_data3d_io, "radial-grids", adios2_type_real_lk, 1,     &
        1_8 * (/ size(radial) /), 1_8 * (/ 0 /),  1_8 * (/ size(radial) /), .true., ierr)
      
      call adios2_define_variable (var_poloidal_grids, adios2_data3d_io, "poloidal-grids", adios2_type_integer4, 1, &
        1_8 * (/ size(mtheta) /), 1_8 * (/ 0 /),  1_8 * (/ size(mtheta) /), .true., ierr)
      
      call adios2_define_variable (var_index_shift, adios2_data3d_io, "index-shift", adios2_type_integer4, 1,       &
        1_8 * (/ size(itran) /), 1_8 * (/ 0 /), 1_8 * (/ size(itran) /), .true., ierr)

      call adios2_define_variable (var_igrid, adios2_data3d_io, "igrid", adios2_type_integer4, 1,       &
        1_8 * (/ size(igrid) /), 1_8 * (/ 0 /), 1_8 * (/ size(igrid) /), .true., ierr)
      
      call adios2_define_variable (var_mgrid, adios2_data3d_io, 'mgrid', adios2_type_integer4, ierr)

      call adios2_define_variable (var_tr_comm_size, adios2_data3d_io, 'toroidal_root_comm_size', adios2_type_integer4, ierr)

      call adios2_define_variable (var_step_id, adios2_data3d_io, 'step_id', adios2_type_integer4, 1, &
        1_8 * (/ 1 /), 1_8 * (/ 0 /), 1_8 * (/ 1 /), .true., ierr)

      ! Define attributes
      call adios2_define_attribute (att_flux_num, adios2_data3d_io, "flux-surface-number maps to", "mpsi+1", ierr)
      call adios2_define_attribute (att_poloidal, adios2_data3d_io, "poloidal-grids maps to", "mtheta+1", ierr)
      
      call adios2_define_attribute (att_indexshift, adios2_data3d_io, "index-shift maps to", "itran", ierr)

    endif ! if mype .eq. 0

    !!! this is written individually for each mtoroidal plane
    if (myrank_partd .eq. 0) then

      !!! the 3d volume variables of all grids

      call adios2_define_variable (var_phi, adios2_data3d_io, "phi", adios2_type_real_lk, 2, &
        1_8 * (/ mdim, tr_comm_size /), 1_8 * (/ 0, myrank_toroidal /), 1_8 * (/ mdim, 1 /), .true., ierr)
          
      call adios2_define_variable (var_x_coordinates, adios2_data3d_io, "x", adios2_type_real_lk, 2, &
        1_8 * (/ mdim, tr_comm_size /), 1_8 * (/ 0, myrank_toroidal /), 1_8 * (/ mdim, 1 /), .true., ierr)
    
      call adios2_define_variable (var_y_coordinates, adios2_data3d_io, "y", adios2_type_real_lk, 2, &
        1_8 * (/ mdim, tr_comm_size /), 1_8 * (/ 0, myrank_toroidal /), 1_8 * (/ mdim, 1 /), .true., ierr)
    
      call adios2_define_variable (var_z_coordinates, adios2_data3d_io, "z", adios2_type_real_lk, 2, &
        1_8 * (/ mdim, tr_comm_size /), 1_8 * (/ 0, myrank_toroidal /), 1_8 * (/ mdim, 1 /), .true., ierr)

      call adios2_define_variable (var_b_3d, adios2_data3d_io, "b", adios2_type_real_lk, 2, &
        1_8 * (/ mdim, tr_comm_size /), 1_8 * (/ 0, myrank_toroidal /), 1_8 * (/ mdim, 1 /), .true., ierr)
    
      !!! the surface variables of the diag flux location
      
      call adios2_define_variable (var_phi_surf, adios2_data3d_io, "phi_surf", adios2_type_real_lk, 2, &
        1_8 * (/ mdim_surf, tr_comm_size /), 1_8 * (/ 0, myrank_toroidal /), 1_8 * (/ mdim_surf, 1 /), .true., ierr)
          
      ! This is written by all 
      call adios2_define_variable (var_x_surf, adios2_data3d_io, "x_surf", adios2_type_real_lk, 2, &
        1_8 * (/ mdim_surf, tr_comm_size /), 1_8 * (/ 0, myrank_toroidal /), 1_8 * (/ mdim_surf, 1 /), .true., ierr)
    
      call adios2_define_variable (var_y_surf, adios2_data3d_io, "y_surf", adios2_type_real_lk, 2, &
        1_8 * (/ mdim_surf, tr_comm_size /), 1_8 * (/ 0, myrank_toroidal /), 1_8 * (/ mdim_surf, 1 /), .true., ierr)
    
      call adios2_define_variable (var_z_surf, adios2_data3d_io, "z_surf", adios2_type_real_lk, 2, &
        1_8 * (/ mdim_surf, tr_comm_size /), 1_8 * (/ 0, myrank_toroidal /), 1_8 * (/ mdim_surf, 1 /), .true., ierr)

      call adios2_define_variable (var_b_surf, adios2_data3d_io, "b_surf", adios2_type_real_lk, 2, &
        1_8 * (/ mdim_surf, tr_comm_size /), 1_8 * (/ 0, myrank_toroidal /), 1_8 * (/ mdim_surf, 1 /), .true., ierr)
    endif

    call adios2_define_attribute (att_3dcoord, adios2_data3d_io, "Dimensions of B and X, Y, Z coordinate arrays", &
        "(mgrid, mtoroidal) - Fortran order", ierr)

    !!! Open file
    call adios2_open (adios2_data3d_e, adios2_data3d_io, "phi_3d.bp", adios2_mode_append, mpi_comm_world, ierr)

    adios2_data3d_io_initialized=.true.
  endif ! if (istep == ndiag)

!!!!!!!!!!!!!
!!! Sec 2: write adios2 variables
!!!!!!!!!!!!!

  !!! write coordinates and b field at the initial step only.
  if(.not. data3d_init_output)then

    ! Begin step
    call adios2_begin_step (adios2_data3d_e, adios2_step_mode_append, ierr)

    if (mype .eq. 0) then
      do i=0,mpsi
         radial(i+1)=(rg0+deltar*real(i))/rho0
      enddo

      ! Cannot write mpsi+1 in adios2_put directly as it is a deferred call that will store the temporary pointer
      ! copy the value to a temp value and use it in adios2_put
      mpsi_tmp = mpsi+1
      ! Similar to mpsi, don't write mtheta+1 in the adios2_put call. Copy it to a tmp var first.
      mtheta_tmp = mtheta+1
      
      call adios2_put (adios2_data3d_e, var_pe_number,           numberpe,     ierr)
      call adios2_put (adios2_data3d_e, var_flux_surface_number, mpsi_tmp,     ierr)
      call adios2_put (adios2_data3d_e, var_mgrid,               mdim,        ierr)
      call adios2_put (adios2_data3d_e, var_tr_comm_size,        tr_comm_size, ierr)
      call adios2_put (adios2_data3d_e, var_igrid,               igrid,        ierr)
      call adios2_put (adios2_data3d_e, var_radial_grids,        radial,       ierr)
      call adios2_put (adios2_data3d_e, var_poloidal_grids,      mtheta_tmp,   ierr)
      call adios2_put (adios2_data3d_e, var_index_shift,         itran,        ierr)

    endif !if mype .eq. 0

    if (myrank_partd .eq. 0)then
      
      ! write data for all grids
      X_coordinates = xdumdom
      Y_coordinates = ydumdom
      Z_coordinates = zdumdom
      ! b field on the 3d grid
      b_3d = bdumdom
      call adios2_put (adios2_data3d_e, var_x_coordinates, X_coordinates, ierr)
      call adios2_put (adios2_data3d_e, var_y_coordinates, Y_coordinates, ierr)
      call adios2_put (adios2_data3d_e, var_z_coordinates, Z_coordinates, ierr)
      call adios2_put (adios2_data3d_e, var_b_3d, b_3d, ierr)

      ! write data for diag flux surface
      ij0=igrid(diag_flux)
      ij1=ij0+mtheta(diag_flux)
      x_surf = xdumdom(ij0:ij1)
      y_surf = ydumdom(ij0:ij1)
      z_surf = zdumdom(ij0:ij1)
      b_surf = bdumdom(ij0:ij1)
      call adios2_put (adios2_data3d_e, var_x_surf, x_surf, ierr)
      call adios2_put (adios2_data3d_e, var_y_surf, y_surf, ierr)
      call adios2_put (adios2_data3d_e, var_z_surf, z_surf, ierr)
      call adios2_put (adios2_data3d_e, var_b_surf, b_surf, ierr)

    endif

    call adios2_end_step (adios2_data3d_e, ierr)
    data3d_init_output = .true.
  endif ! if ( (.not. data3d_init_output) .and. irun == 0)

  !!! write phi for every diag step  

  ! Subsequent steps
  current_step = mstepall + istep
  call adios2_begin_step (adios2_data3d_e, adios2_step_mode_append, ierr)

  if (mype==0)then
    call adios2_put (adios2_data3d_e, var_step_id, current_step, ierr)
  endif
 
  ! potential data file
  if (myrank_partd .eq. 0) then
 
    data3d_potential = ndumdom
    call adios2_put(adios2_data3d_e, var_phi, data3d_potential, ierr)

    ij0=igrid(diag_flux)
    ij1=ij0+mtheta(diag_flux)
    phi_surf = ndumdom(ij0:ij1)
    call adios2_put(adios2_data3d_e, var_phi_surf, phi_surf, ierr)

  endif

  call adios2_end_step (adios2_data3d_e, ierr)
 
end subroutine dataout3d_adios2
#endif

