"""
Create a VTK file of vapor flux from a candis file
"""

from evtk.hl  import imageToVTK
from evtk.vtk import VtkFile, VtkImageData
import numpy as np

# module to read candis data files
import CandisConvert as cancon

basename = 'a021'
index = 262

## read data
#
cc = cancon.CandisConvert(basename)
cfile = cc.read(index)

r   = cfile.getField('r').data
rho = cfile.getField('rho').data
vz  = cfile.getField('vz').data

## calculate flux
#
iz = 90
rflux = r[:,:,iz]*rho[:,:,iz]*vz[:,:,iz]

## add an extra dimension to make it 3D for VTK
#
rflux = rflux.reshape(rflux.shape + (1,))
(nx,ny,nz) = rflux.shape

## setup cell information
#
ncells = nx*ny*nz
start   = (  0,  0,  0)
end     = ( nx, ny, nz)
origin  = (0.0,0.0,0.0)
spacing = (1.0,1.0,1.0)

w = VtkFile('rflux.' + basename + '.' + str(index), VtkImageData)
w.openGrid(start, end, origin, spacing)
w.openPiece(start, end)
w.openData('Cell', scalars = 'rflux')
w.addData('rflux', rflux)
w.closeData('Cell')
w.closePiece()
w.closeGrid()
w.appendData(data = rflux)
w.save()