#!/usr/bin/env python3
import vtk

def get_program_parameters():
    import argparse
    description = 'Read and display ExodusII data from multiple files.'
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument('filenames', nargs='+', help="A required filename pattern, e.g, mug_*.e.")
    parser.add_argument('--variable', type=str, help="The nodal variable, e.g. 'convected'.")
    parser.add_argument('--timestep', type=int, help="The timestep to load.")
    args = parser.parse_args()
    return args

def main():
    colors = vtk.vtkNamedColors()

    # Input files and variable
    args = get_program_parameters()

    # Read Exodus Data
    readers = list()
    for filename in args.filenames:
        reader = vtk.vtkExodusIIReader()
        reader.SetFileName(filename)
        reader.UpdateInformation()
        reader.SetTimeStep(args.timestep)
        reader.SetAllArrayStatus(vtk.vtkExodusIIReader.NODAL, 1)  # enables all NODAL variables
        reader.Update()
        readers.append(reader)

    merge = vtk.vtkMultiBlockMergeFilter()
    for reader in readers:
        merge.AddInputConnection(reader.GetOutputPort(0))
    merge.Update()

    geometry = vtk.vtkCompositeDataGeometryFilter()
    geometry.SetInputConnection(0, merge.GetOutputPort(0))
    geometry.Update()

    # Mapper
    mapper = vtk.vtkPolyDataMapper()
    mapper.SetInputConnection(geometry.GetOutputPort())
    mapper.SelectColorArray(args.variable)
    mapper.SetScalarModeToUsePointFieldData()
    mapper.InterpolateScalarsBeforeMappingOn()

    # Actor
    actor = vtk.vtkActor()
    actor.SetMapper(mapper)

    # Renderer
    colors = vtk.vtkNamedColors()
    renderer = vtk.vtkRenderer()
    renderer.AddViewProp(actor)
    renderer.SetBackground(colors.GetColor3d('DimGray'))

    # Window and Interactor
    window = vtk.vtkRenderWindow()
    window.AddRenderer(renderer)
    window.SetSize(600, 600)
    window.SetWindowName('ReadExodusData')

    interactor = vtk.vtkRenderWindowInteractor()
    interactor.SetRenderWindow(window)
    interactor.Initialize()

    # Show the result
    window.Render()
    interactor.Start()

if __name__ == '__main__':
    main()