import adios2
from mpi4py import MPI
import numpy
import sys

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

adios = adios2.ADIOS(MPI.COMM_WORLD)
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

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

if attrOutput:
    bpIO.DefineAttribute("Fides_Data_Model", "unstructured_single")
    bpIO.DefineAttribute("Fides_Cell_Type", "triangle")
    bpIO.DefineAttribute("Fides_Time_Variable", "time")

    varList = ["dpot", "dpot2"]
    bpIO.DefineAttribute("Fides_Variable_List", varList)

    varAssoc = ["points", "points"]
    bpIO.DefineAttribute("Fides_Variable_Associations", varAssoc)

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)

time = numpy.ndarray((1))
time[0] = 0.1
timeVar = bpIO.DefineVariable("time", time)

bpFileWriter = None
if attrOutput:
    bpFileWriter = bpIO.Open("tris-blocks-time-attr.bp", adios2.Mode.Write)
else:
    bpFileWriter = bpIO.Open("tris-blocks-time.bp", adios2.Mode.Write)

for i in range(5):
    dpot = dpot*1.1
    # if i == 2:
    #     print(numpy.min(dpot), numpy.max(dpot))
    bpFileWriter.BeginStep()
    if i == 0:
        bpFileWriter.Put(conVar, vtkm_conn, adios2.Mode.Sync)
        bpFileWriter.Put(ptsVar, pts, adios2.Mode.Sync)
    bpFileWriter.Put(timeVar, time, adios2.Mode.Sync)
    bpFileWriter.Put(dpotVar, dpot, adios2.Mode.Sync)
    bpFileWriter.Put(dpotVar2, dpot, adios2.Mode.Sync)
    bpFileWriter.EndStep()
    time[0] += 1.0
bpFileWriter.Close()
