import adios2
from mpi4py import MPI
import numpy
import sys

adios = adios2.ADIOS(MPI.COMM_WORLD)
bpIO = adios.DeclareIO("BPFile")
bpIO.SetEngine('bp3')

import vtk

sps = vtk.vtkSphereSource()
sps.Update()
pd = sps.GetOutput()
tris = pd.GetPolys()
triIds = tris.GetData()

from vtk.numpy_interface import dataset_adapter as dsa

attrOutput = False
if len(sys.argv) == 2 and sys.argv[1] == "attributes":
    attrOutput = True

if attrOutput:
    bpIO.DefineAttribute("Fides_Data_Model", "unstructured")

tris = dsa.vtkDataArrayToVTKArray(triIds)
vtkm_conn = tris.reshape([96,4])[:,1:4].flatten()
print(vtkm_conn)
conVar = bpIO.DefineVariable("connectivity", vtkm_conn, [288], [0], [288], adios2.ConstantDims)

cellTypes = numpy.ones(96, dtype=numpy.uint8) * 5
typesVar = bpIO.DefineVariable("cell_types", cellTypes, [96], [0], [96], adios2.ConstantDims)

numVerts = numpy.ones(96, dtype=numpy.int32) * 3
numVar = bpIO.DefineVariable("num_verts", numVerts, [96], [0], [96], adios2.ConstantDims)

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

bpFileWriter = None
if attrOutput:
    bpFileWriter = bpIO.Open("tris-explicit-attr.bp", adios2.Mode.Write)
else:
    bpFileWriter = bpIO.Open("tris-explicit.bp", adios2.Mode.Write)
bpFileWriter.Put(conVar, vtkm_conn, adios2.Mode.Sync)
bpFileWriter.Put(typesVar, cellTypes, adios2.Mode.Sync)
bpFileWriter.Put(numVar, numVerts, adios2.Mode.Sync)
bpFileWriter.Put(ptsVar, pts, adios2.Mode.Sync)
bpFileWriter.Close()
