import adios2
from mpi4py import MPI
import numpy

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

adios = adios2.ADIOS(MPI.COMM_WORLD)

# read in only gem_den_3d and gem_field_3d
# xgc type data is really large and we don't really care about it for testing
readIO = adios.DeclareIO("reader")
readIO.SetEngine('bp4')
reader = readIO.Open("gem_xgc.3d.bp", adios2.Mode.Read)

denVar = readIO.InquireVariable("gem_den_3d")
densityR = numpy.ones(denVar.SelectionSize(), dtype=numpy.float64)
reader.Get(denVar, densityR, adios2.Mode.Sync)

fieldVar = readIO.InquireVariable("gem_field_3d")
fieldR = numpy.ones(fieldVar.SelectionSize(), dtype=numpy.float64)
reader.Get(fieldVar, fieldR, adios2.Mode.Sync)

reader.PerformGets()
reader.Close()

#write out the two fields
writer = readIO.Open("gem.3d.bp", adios2.Mode.Write)
writer.Put(denVar, densityR, adios2.Mode.Sync)
writer.Put(fieldVar, fieldR, adios2.Mode.Sync)
writer.PerformPuts()
writer.Close()
