#!/usr/bin/env python
import sys
import os
import struct
from optparse import OptionParser

# Read command line arguments
parser = OptionParser()
parser.add_option("-z", "--z", dest="z_coord", default=0.0,
                  help="Create a slice at z = ZCOORD", metavar="ZCOORD")
(options, args) = parser.parse_args()

if len(args) != 2:
    print "Two arguments required; in_file and out_file"
    sys.exit(1)

in_file = args[0]
out_file = args[1]

# Read 10 header lines from the VTK file. We assume a valid and very
# precise format for the file.  Behaviour is undefined if the file is
# malformed or does not precisely correspond to the format we expect.
f = open(in_file, 'r')
for i in range(1,11): 
    line = f.readline().rstrip("\n")
    # Parse these lines
    words = line.split()
    key = words[0]

    if key == 'ASCII':
        mode = "ascii"

    if key == 'BINARY':
        mode = "binary"

    if key == 'DIMENSIONS':
        dims = map(int, words[1:])
        [nx,ny,nz] = dims

    if key == 'ORIGIN':
        origin = map(float, words[1:])
        [ox,oy,oz] = origin

    if key == 'SPACING':
        spacing = map(float, words[1:])
        [dx,dy,dz] = spacing

    if key == 'POINT_DATA':
        point_data = int(words[1])

    if key == 'SCALARS':
        scalars = words[1:]
        [var, type, num] = words[1:4]
        if type == "float":
            binary_size = 4
        elif type == "double":
            binary_size = 8
        else:
            print "Unrecognized data type: %s" % type

    if key == 'LOOKUP_TABLE':
        lookup_table = words[1:]


iz = int(round((float(options.z_coord) - oz) / dz))

if iz < 0 or iz >= nz:
    print "Coordinate z = %f is not within file %s" % in_file
    sys.exit(1)

# To read a slice in the xy plane, we should just have to read a
# consecutive sequence of nx * ny values

# How many values to skip before the slice
skip_values = iz * (nx * ny)

# Read all the remaining words (whitespace-separated strings) in the file
def read_words(f):
    words = []
    for line in f:
        words[len(words):] = line.split()
    return words

if mode == "binary":
    f.seek(skip_values * binary_size, 1) # Avoid having to read the data we don't need
else:
    floats = read_words(f)

g = open(out_file, 'w')

# Write out the sliced version of the file
g.write('# vtk DataFile Version 2.0\n')
g.write('Generated by h5tovtk, unchunked by unchunk, sliced by vtkslice\n')
if mode == "binary":
    g.write('BINARY\n')
else:
    g.write('ASCII\n')
g.write('DATASET STRUCTURED_POINTS\n')
g.write('DIMENSIONS %d %d %d\n' % (nx, ny, 1))
g.write('ORIGIN %f %f %f\n' % (ox, oy, oz + iz * dz))
g.write('SPACING %f %f %f\n' % (dx, dy, dz))
g.write('POINT_DATA %d\n' % (nx * ny))
g.write('SCALARS %s %s %s\n' % (var, type, num))
g.write('LOOKUP_TABLE default\n')

if mode == "binary":
    for i in range(0, nx * ny):
        bytes = f.read(binary_size)
        g.write(bytes)
        #    x = struct.unpack(">f", bytes)[0]
else:
    for j in range(0, ny):
        for i in range(0, nx):
            g.write(floats[skip_values+j*nx+i])
            if i < nx - 1:
                g.write(" ")
            else:
                g.write("\n")

g.close()
f.close()