#------------------------------------------------------------------------------
# convert_to_vdb.py - Create OpenVDB (.vdb) volume data file from
# image data in argument JSON file. JSON file is generated using
# "VTK To OpenVDB Exporter" node in BVTKNodes (Blender add-on).
#
# Run example: python3 convert_to_vdb.py volume_00001.json
#
# Requirement: pyopenvdb module must be available to Python
#
# If you get error like:
#     "libjemalloc.so.2: cannot allocate memory in static TLS block"
# then prepend command with LD_PRELOAD:
#     LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libjemalloc.so.2 python3 convert_to_vdb.py volume_00001.json
#------------------------------------------------------------------------------

import sys
import json

try:
    import pyopenvdb
    vdb = pyopenvdb
except ImportError:
    print("ERROR: Python can't import pyopenvdb. Please see "
    + "VTK To OpenVDB Exporter node documentation for more information.")
    raise


def create_grid(background_value, dims, array, grid_name, atype):
    """
    Return an OpenVDB grid with dimensions dims containing data from array.
    """

    if array == None:
        return None

    if atype == 'scalar':
        grid = vdb.FloatGrid(background_value)
    elif atype == 'vector':
        grid = vdb.Vec3SGrid()
    else:
        raise TypeError("Unknown type %s" % str(atype))

    grid.gridClass = vdb.GridClass.FOG_VOLUME
    grid.name = grid_name

    # Create grid by looping over points with accessor
    acc = grid.getAccessor()
    for i in range(dims[0]):
        for j in range(dims[1]):
            for k in range(dims[2]):
                idx = i + j*dims[0] + k*dims[0]*dims[1]
                value = array[idx]
                if atype == 'scalar':
                    if value == None:
                        continue
                    if value > background_value:
                        acc.setValueOn((i, j, k), value)

                elif atype == 'vector':
                    if value[0] == None:
                        continue
                    acc.setValueOn((i, j, k), value)
    return grid


def count_active_voxels(grids):
    """
    Counts the number of active voxels in list of OpenVDB grids
    """

    n = 0
    for grid in grids:
        n += grid.activeVoxelCount()
    return n


def create_vdb(vdb_filename, background_value, dims, density_data,
               color_data, flame_data, temperature_data):
    """
    Create vdb file from argument data.
    """

    density_grid = create_grid(background_value, dims, density_data, 'density', 'scalar')
    color_grid = create_grid(background_value, dims, color_data, 'color', 'vector')
    flame_grid = create_grid(background_value, dims, flame_data, 'flame', 'scalar')
    temperature_grid = create_grid(background_value, dims, temperature_data, 'temperature', 'scalar')

    grids = [density_grid, color_grid, flame_grid, temperature_grid]
    grids = [g for g in grids if g is not None]
    nvoxels = count_active_voxels(grids)
    vdb.write(vdb_filename, grids=grids)
    print("%s: %d grids, %d active voxels exported." % (vdb_filename, len(grids), nvoxels))


# Main program: Process all arguments

for filename in sys.argv[1:]:
    with open(filename, "r") as read_file:
        background_value, dims, density_data, color_data, flame_data, \
            temperature_data = json.load(read_file)
        vdb_filename = filename.replace('.json', '.vdb')
        create_vdb(vdb_filename, background_value, dims, density_data,
                   color_data, flame_data, temperature_data)