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, adios2.DebugON)
bpIO = adios.DeclareIO("BPFile")
bpIO.SetEngine('bp3')

import vtk

sps = vtk.vtkPSphereSource()
el = vtk.vtkElevationFilter()
el.SetLowPoint(-1, 0, 0)
el.SetInputConnection(sps.GetOutputPort())
el.UpdatePiece(rank, size, 0)
pd = el.GetOutput()
tris = pd.GetPolys()
triIds = tris.GetData()

from vtk.numpy_interface import dataset_adapter as dsa

tris = dsa.vtkDataArrayToVTKArray(triIds)
ntris = int(tris.shape[0] / 4)
vtkm_conn = tris.reshape([ntris,4])[:,1:4].flatten()
conVar = bpIO.DefineVariable("connectivity", vtkm_conn, [], [], [ntris*3], adios2.ConstantDims)

pts = dsa.vtkDataArrayToVTKArray(pd.GetPoints().GetData())
sp = pts.shape
# print(pts)
ptsVar = bpIO.DefineVariable("points", pts, [], [], list(sp), adios2.ConstantDims)

dpot = dsa.vtkDataArrayToVTKArray(pd.GetPointData().GetArray("Elevation"))
sp = dpot.shape
dpotVar = bpIO.DefineVariable("dpot", dpot, [], [], list(sp), adios2.ConstantDims)
dpotVar2 = bpIO.DefineVariable("dpot2", dpot, [], [], list(sp), adios2.ConstantDims)

bpFileWriter = bpIO.Open("tris-blocks-mesh.bp", adios2.Mode.Write)
bpFileWriter.BeginStep()
bpFileWriter.Put(conVar, vtkm_conn, adios2.Mode.Sync)
bpFileWriter.Put(ptsVar, pts, adios2.Mode.Sync)
bpFileWriter.EndStep()
bpFileWriter.Close()

bpFileWriter = bpIO.Open("tris-blocks-fields.bp", adios2.Mode.Write)
for i in range(5):
    dpot = dpot*1.1
    bpFileWriter.BeginStep()
    bpFileWriter.Put(dpotVar, dpot, adios2.Mode.Sync)
    bpFileWriter.Put(dpotVar2, dpot, adios2.Mode.Sync)
    bpFileWriter.EndStep()
bpFileWriter.Close()
