#!/usr/bin/env python

from __future__ import division, print_function
import struct
import sys

import numpy as np
import h5py

def parse_options():
    """Process command line arguments"""

    from optparse import OptionParser
    usage = r"""%prog [options] <voxel_file>"""
    p = OptionParser(usage=usage)
    p.add_option('-o', '--output', action='store', dest='output',
                 default='plot', help='Path to output SILO or VTK file.')
    p.add_option('-v', '--vtk', action='store_true', dest='vtk',
                 default=False, help='Flag to convert to VTK instead of SILO.')
    parsed = p.parse_args()
    if not parsed[1]:
        p.print_help()
        return parsed
    return parsed


def main(filename, o):
    # Read data from voxel file
    fh = h5py.File(filename, 'r')
    dimension = fh['num_voxels'].value
    width = fh['voxel_width'].value
    lower_left = fh['lower_left'].value
    voxel_data = fh['data'].value

    nx, ny, nz = dimension
    upper_right = lower_left + width*dimension

    if o.vtk:
        try:
            import vtk
        except:
            print('The vtk python bindings do not appear to be installed '
                  'properly.\nOn Ubuntu: sudo apt-get install python-vtk\n'
                  'See: http://www.vtk.org/')
            return

        grid = vtk.vtkImageData()
        grid.SetDimensions(nx+1, ny+1, nz+1)
        grid.SetOrigin(*lower_left)
        grid.SetSpacing(*width)

        data = vtk.vtkDoubleArray()
        data.SetName("id")
        data.SetNumberOfTuples(nx*ny*nz)
        for x in range(nx):
            sys.stdout.write(" {0}%\r".format(int(x/nx*100)))
            sys.stdout.flush()
            for y in range(ny):
                for z in range(nz):
                    i = z*nx*ny + y*nx + x
                    data.SetValue(i, voxel_data[x,y,z])
        grid.GetCellData().AddArray(data)

        writer = vtk.vtkXMLImageDataWriter()
        if vtk.vtkVersion.GetVTKMajorVersion() > 5:
            writer.SetInputData(grid)
        else:
            writer.SetInput(grid)
        if not o.output.endswith(".vti"):
            o.output += ".vti"
        writer.SetFileName(o.output)
        writer.Write()

    else:
        try:
            import silomesh
        except:
            print('The silomesh package does not appear to be installed '
                  'properly.\nSee: https://github.com/nhorelik/silomesh/')
            return
        if not o.output.endswith(".silo"):
            o.output += ".silo"
        silomesh.init_silo(o.output)
        meshparams = list(map(int, dimension)) + list(map(float, lower_left)) + \
                     list(map(float, upper_right))
        silomesh.init_mesh('plot', *meshparams)
        silomesh.init_var("id")
        for x in range(nx):
            sys.stdout.write(" {0}%\r".format(int(x/nx*100)))
            sys.stdout.flush()
            for y in range(ny):
                for z in range(nz):
                    silomesh.set_value(float(voxel_data[x,y,z]),
                                       x + 1, y + 1, z + 1)
        print()
        silomesh.finalize_var()
        silomesh.finalize_mesh()
        silomesh.finalize_silo()


if __name__ == '__main__':
    (options, args) = parse_options()
    if args:
        main(args[0], options)
