import adios2
from mpi4py import MPI
import numpy

def readVar(readIO, reader, varNm, varType) :
  print(varNm)
  var = readIO.InquireVariable(varNm)
  arr = numpy.ones(var.SelectionSize(), dtype=varType)
  reader.Get(var, arr, adios2.Mode.Sync)
  return (var, arr)

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('gtc.bp', adios2.Mode.Read)

(bVar, bArr) = readVar(readIO, reader, 'b', numpy.float32)
(phiVar, phiArr) = readVar(readIO, reader, 'phi', numpy.float32)
(xVar, xArr) = readVar(readIO, reader, 'x', numpy.float32)
(yVar, yArr) = readVar(readIO, reader, 'y', numpy.float32)
(zVar, zArr) = readVar(readIO, reader, 'z', numpy.float32)
(igVar, igArr) = readVar(readIO, reader, 'igrid', numpy.int32)
(shiftVar, shiftArr) = readVar(readIO, reader, 'index-shift', numpy.int32)
(numPlnVar, numPlnArr) = readVar(readIO, reader, 'numPlanes', numpy.int32)
(numPtsVar, numPtsArr) = readVar(readIO, reader, 'numPtsPerPlane', numpy.int32)

reader.PerformGets()
reader.Close()

writer = readIO.Open('gtc.2.bp', adios2.Mode.Write)
#write the grid information.
writer.Put(xVar, xArr, adios2.Mode.Sync)
writer.Put(yVar, yArr, adios2.Mode.Sync)
writer.Put(zVar, zArr, adios2.Mode.Sync)
writer.Put(igVar, igArr, adios2.Mode.Sync)
writer.Put(shiftVar, shiftArr, adios2.Mode.Sync)
writer.Put(numPlnVar, numPlnArr, adios2.Mode.Sync)
writer.Put(numPtsVar, numPtsArr, adios2.Mode.Sync)

#write out two steps for each of the two fields.
writer.BeginStep()
writer.Put(bVar, bArr, adios2.Mode.Sync)
writer.Put(phiVar, phiArr, adios2.Mode.Sync)
writer.EndStep()

writer.BeginStep()
writer.Put(bVar, bArr, adios2.Mode.Sync)
writer.Put(phiVar, phiArr, adios2.Mode.Sync)
writer.EndStep()

writer.Close()
