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

sr = vtk.vtkXMLRectilinearGridReader()
sr.SetFileName("../tests/data/RectGrid2.vtr")
sr.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

rg = sr.GetOutput()
xcoords = dsa.vtkDataArrayToVTKArray(rg.GetXCoordinates())
ycoords = dsa.vtkDataArrayToVTKArray(rg.GetYCoordinates())
zcoords = dsa.vtkDataArrayToVTKArray(rg.GetZCoordinates())

xcoordsVar = bpIO.DefineVariable("x", xcoords, [wExt[1]+1], [ext[0]], [count[0]], adios2.ConstantDims)
ycoordsVar = bpIO.DefineVariable("y", ycoords, [wExt[3]+1], [ext[2]], [count[1]], adios2.ConstantDims)
zcoordsVar = bpIO.DefineVariable("z", zcoords, [wExt[5]+1], [ext[4]], [count[2]], adios2.ConstantDims)

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

if attrOutput:
    bpIO.DefineAttribute("Fides_Data_Model", "rectilinear")
    bpIO.DefineAttribute("Fides_Dimension_Variable", "scalars")

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

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

scalars = dsa.vtkDataArrayToVTKArray(rg.GetPointData().GetArray("scalars"))
scalars = scalars.reshape(rg.GetDimensions()[::-1], order='C')
scalarsVar = bpIO.DefineVariable("scalars", scalars, numpy.array(wExt[5::-2])+1, ext[4::-2], count[::-1], adios2.ConstantDims)

vectors = dsa.vtkDataArrayToVTKArray(rg.GetPointData().GetArray("vectors"))
vectors = vectors.reshape(rg.GetDimensions()[::-1]+(3,), order='C')
vectorsVar = bpIO.DefineVariable("vectors", vectors, numpy.array(wExt[5::-2]+(2,))+1, ext[4::-2]+(0,), count[::-1]+(3,), adios2.ConstantDims)

bpFileWriter = None
if attrOutput:
    bpFileWriter = bpIO.Open("rectilinear-attr.bp", adios2.Mode.Write)
else:
    bpFileWriter = bpIO.Open("rectilinear.bp", adios2.Mode.Write)
bpFileWriter.Put(xcoordsVar, xcoords, adios2.Mode.Sync)
bpFileWriter.Put(ycoordsVar, ycoords, adios2.Mode.Sync)
bpFileWriter.Put(zcoordsVar, zcoords, adios2.Mode.Sync)
bpFileWriter.Put(scalarsVar, scalars, adios2.Mode.Sync)
bpFileWriter.Put(vectorsVar, vectors, adios2.Mode.Sync)
bpFileWriter.EndStep()
bpFileWriter.Close()

# pts = dsa.vtkDataArrayToVTKArray(sg.GetPoints().GetData())
# points = pts.reshape(sg.GetDimensions()[::-1]+(3,), order='C')
# sp = points.shape
if rank == 0:
    print("wExt:" ,wExt[1:6:2])
    print(sr.GetOutput().GetExtent())
    print(numpy.array(wExt[5::-2])+1, ext[4::-2], count[::-1])
# 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"))
# density = dens.reshape(sg.GetDimensions()[::-1], order='C')
# sp = density.shape
# 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])
# 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()
