import adios2
from mpi4py import MPI
import numpy
import sys
import os

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
schemaOutput = False
if len(sys.argv) == 2:
    if sys.argv[1] == "attributes":
        attrOutput = True
    elif sys.argv[1] == "schema":
        schemaOutput = True

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

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

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

if schemaOutput:
    filename = os.path.join(os.path.dirname(os.path.dirname(__file__)), "tests", "vtk-uns-grid-2.json")
    with open(filename) as f:
        bpIO.DefineAttribute("fides/schema", f.read())
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
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 = None
if attrOutput:
    bpFileWriter = bpIO.Open("tris-blocks-attr.bp", adios2.Mode.Write)
elif schemaOutput:
    bpFileWriter = bpIO.Open("tris-blocks-schema-as-json-string.bp", adios2.Mode.Write)
else:
    bpFileWriter = bpIO.Open("tris-blocks.bp", adios2.Mode.Write)
bpFileWriter.Put(conVar, vtkm_conn, adios2.Mode.Sync)
bpFileWriter.Put(ptsVar, pts, adios2.Mode.Sync)
bpFileWriter.Put(dpotVar, dpot, adios2.Mode.Sync)
bpFileWriter.Put(dpotVar2, dpot, adios2.Mode.Sync)
bpFileWriter.Close()
