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

import vtk

sr = vtk.vtkXMLStructuredGridReader()
sr.SetFileName("../tests/data/multicomb_0.vts")

el = vtk.vtkElevationFilter()
el.SetInputConnection(sr.GetOutputPort())
tr = vtk.vtkThreshold()
tr.SetInputConnection(el.GetOutputPort())
tr.ThresholdByUpper(-1)
tr.UpdatePiece(rank, size, 0)

wExt = sr.GetExecutive().GetOutputInformation(0).Get(
    vtk.vtkStreamingDemandDrivenPipeline.WHOLE_EXTENT())
ext = sr.GetOutput().GetExtent()
count = sr.GetOutput().GetDimensions()


from vtk.numpy_interface import dataset_adapter as dsa

sg = sr.GetOutput()
pts = dsa.vtkDataArrayToVTKArray(sg.GetPoints().GetData())
# print(sg.GetDimensions())
points = pts.reshape(sg.GetDimensions()[::-1]+(3,), order='C')
sp = points.shape
# print(sp)
if rank == 1:
    print("wExt:" ,wExt[1:6:2])
    print(sr.GetOutput().GetExtent())
    print(numpy.array(wExt[5::-2])+1, ext[4::-2], count[::-1])
# print(sp, sg.GetDimensions())
# print(numpy.array(wExt[5::-2]+(2,))+1, ext[4::-2]+(0,), count[::-1]+(3,))
pointsVar = bpIO.DefineVariable("points", points, numpy.array(wExt[5::-2]+(2,))+1, ext[4::-2]+(0,), count[::-1]+(3,), adios2.ConstantDims)

dens = dsa.vtkDataArrayToVTKArray(sg.GetPointData().GetArray("Density"))
# print(sg.GetDimensions())
density = dens.reshape(sg.GetDimensions()[::-1], order='C')
sp = density.shape
# print(sp)
if rank == 1:
    print("wExt:" ,wExt[1:6:2])
    print(sr.GetOutput().GetExtent())
    print(numpy.array(wExt[5::-2])+1, ext[4::-2], count[::-1])
# print(sp, sg.GetDimensions())
densityVar = bpIO.DefineVariable("density", density, numpy.array(wExt[5::-2])+1, ext[4::-2], count[::-1], adios2.ConstantDims)

bpFileWriter = bpIO.Open("structured.bp", adios2.Mode.Write)
bpFileWriter.Put(pointsVar, points, adios2.Mode.Sync)
bpFileWriter.Put(densityVar, density, adios2.Mode.Sync)
bpFileWriter.Close()

# ioRead = adios.DeclareIO("ioReader")

# ibpStream = ioRead.Open("uns-uniform-pts.bp", adios2.Mode.Read)

# densityVar = ioRead.InquireVariable("density")

# densityR = numpy.ones(densityVar.SelectionSize(), dtype=numpy.float32)

# ibpStream.Get(densityVar, densityR, adios2.Mode.Sync)

# ibpStream.PerformGets()
#ibpStream.Close()

#print(densityR-dens)


# ug = tr.GetOutput()
# cells = ug.GetCells().GetData()
# cells = dsa.vtkDataArrayToVTKArray(cells)
# vtkm_conn = cells[1:9]
# print(vtkm_conn)

# conVar = bpIO.DefineVariable("connectivity", vtkm_conn, [8], [0], [8], adios2.ConstantDims)

# pts = dsa.vtkDataArrayToVTKArray(ug.GetPoints().GetData())
# sp = pts.shape
# ptsVar = bpIO.DefineVariable("points", pts, list(sp), [0, 0], list(sp), adios2.ConstantDims)
# dims = numpy.array([2,2,2], dtype=numpy.int64)
# dimsVar = bpIO.DefineVariable("dimensions", dims, [3], [0], [3], adios2.ConstantDims)

# bpFileWriter = bpIO.Open("uns-uniform-pts.bp", adios2.Mode.Write)
# bpFileWriter.Put(conVar, vtkm_conn, adios2.Mode.Sync)
# bpFileWriter.Put(ptsVar, pts, adios2.Mode.Sync)
# bpFileWriter.Put(dimsVar, dims, adios2.Mode.Sync)
# bpFileWriter.Close()
