import numpy as np
from vtk.numpy_interface import algorithms as algs
from vtk.numpy_interface import dataset_adapter as dsa


# "Boilerplate" code to get distribution information
executive = self.GetExecutive()
outInfo = executive.GetOutputInformation(0)
pid = outInfo.Get(executive.UPDATE_PIECE_NUMBER())

exts = [executive.UPDATE_EXTENT().Get(outInfo, i) for i in xrange(6)]
exts_slice_x = slice(exts[0], exts[1]+1)
exts_slice_y = slice(exts[2], exts[3]+1)
exts_slice_z = slice(exts[4], exts[5]+1)
exts_slice = (exts_slice_x, exts_slice_y, exts_slice_z)
whole = [executive.WHOLE_EXTENT().Get(outInfo, i) for i in xrange(6)]

dims = [exts[1]-exts[0]+1, exts[3]-exts[2]+1, exts[5]-exts[4]+1]
global_dims = [whole[1]-whole[0]+1, whole[3]-whole[2]+1, whole[5]-whole[4]+1]

output.SetExtent(exts)

# Read the binary cartesian file function
def readcart(fname, dims):
    dims = np.array(dims)  # Easier to use np array to get count of items
    with open(fname, 'rb') as f:
        # Fortran records have INTEGER*4 headers which are useless here
        f.seek(4)

        data = np.fromfile(f, dtype=np.float32, count=np.prod(dims), sep="")
        # Order="F" is the key here
        return data.reshape(dims, order="F")

# Construct grid
xaxis = np.linspace(-1, 1, global_dims[0])[exts_slice_x]
yaxis = np.linspace(-1, 1, global_dims[1])[exts_slice_y]
zaxis = np.linspace(-1, 1, global_dims[2])[exts_slice_z]

output.SetXCoordinates(dsa.numpyTovtkDataArray(xaxis , "X"))
output.SetYCoordinates(dsa.numpyTovtkDataArray(yaxis , "Y"))
output.SetZCoordinates(dsa.numpyTovtkDataArray(zaxis , "Z"))

# Get data
if directory[-1] != r"/":
    directory = directory + "/"
xname = directory + "field_bx.{}.out".format(filenum)
yname = directory + "field_by.{}.out".format(filenum)
zname = directory + "field_bz.{}.out".format(filenum)


# Reshaping is because I need two axes to hstack and get (3n, 3). If I ravel
# I get (3n,) with hstack
# Must use "F" order again as reshaping only changes access order, not
# memory location (implying that every reshaping must use "F" order)
bx = readcart(xname, global_dims)[exts_slice].reshape(-1, 1, order="F")
by = readcart(yname, global_dims)[exts_slice].reshape(-1, 1, order="F")
bz = readcart(zname, global_dims)[exts_slice].reshape(-1, 1, order="F")

b = np.hstack((bx, by, bz))

# Convert to VTK and give names
array = dsa.numpyTovtkDataArray(b, "B")
array.SetComponentName(0, "B_x")
array.SetComponentName(1, "B_y")
array.SetComponentName(2, "B_z")

# Associate data to grid
output.GetPointData().AddArray(array)