#!/usr/bin/env python3

from argparse import ArgumentParser

import h5py
import vtk

_min_version = (2, 0)


def voxel_to_vtk(voxel_file: str, output: str):

    # Read data from voxel file
    fh = h5py.File(voxel_file, "r")

    # check version
    version = tuple(fh.attrs["version"])
    if version < _min_version:
        old_version = ".".join(map(str, version))
        min_version = ".".join(map(str, _min_version))
        err_msg = (
            "This voxel file's version is {}. This script "
            "only supports voxel files with version {} or "
            "higher. Please generate a new voxel file using "
            "a newer version of OpenMC.".format(old_version, min_version)
        )
        raise ValueError(err_msg)

    dimension = fh.attrs["num_voxels"]
    width = fh.attrs["voxel_width"]
    lower_left = fh.attrs["lower_left"]

    nx, ny, nz = dimension

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

    # transpose data from OpenMC ordering (zyx) to VTK ordering (xyz)
    # and flatten to 1-D array
    print("Reading and translating data...")
    h5data = fh["data"][...]

    data = vtk.vtkIntArray()
    data.SetName("id")
    # set the array using the h5data array
    data.SetArray(h5data, h5data.size, True)
    # add data to image grid
    grid.GetCellData().AddArray(data)

    writer = vtk.vtkXMLImageDataWriter()
    if vtk.vtkVersion.GetVTKMajorVersion() > 5:
        writer.SetInputData(grid)
    else:
        writer.SetInput(grid)
    if not output.endswith(".vti"):
        output += ".vti"
    writer.SetFileName(output)
    print("Writing VTK file {}...".format(output))
    writer.Write()


if __name__ == "__main__":
    # Process command line arguments
    parser = ArgumentParser('Converts a voxel HDF5 file to a VTK file')
    parser.add_argument("voxel_file", help="Path to voxel h5 file")
    parser.add_argument(
        "-o",
        "--output",
        action="store",
        default="plot",
        help="Path to output VTK file.",
    )
    args = parser.parse_args()
    voxel_to_vtk(args.voxel_file, args.output)
