import vtk

# This script shows a cube and a cylinder in a vtkMultiBlockDataSet
# using four different mapper configurations and three different cell
# opacity scenarios:
# 1. a vtkCompositeDataMapper2 mapping the vtkMultiBlockDataSet
# 2. a vtkCompositeDataMapper mapping the vtkMultiBlockDataSet
# 3. a vtkPolyDataMapper mapping all blocks appended into a single vtkPolyData
# 4. two vtkPolyDataMappers each mapping one of the blocks
#
# Opacity scenarios demonstrated are:
# A. opaque cylinder, opaque block
# B. opaque cylinder, transparent block
# C. transparent cylinder, transparent block

opaque_red=vtk.vtkApplyColors()
opaque_red.SetDefaultCellOpacity(1.)
opaque_red.SetDefaultCellColor(1.,0,0)
opaque_red.SetCellColorOutputArrayName('scalars')

transparent_red=vtk.vtkApplyColors()
transparent_red.SetDefaultCellOpacity(.5)
transparent_red.SetDefaultCellColor(1.,0,0)
transparent_red.SetCellColorOutputArrayName('scalars')

opaque_blue=vtk.vtkApplyColors()
opaque_blue.SetDefaultCellOpacity(1.)
opaque_blue.SetDefaultCellColor(0,0,1.)
opaque_blue.SetCellColorOutputArrayName('scalars')

transparent_blue=vtk.vtkApplyColors()
transparent_blue.SetDefaultCellOpacity(.5)
transparent_blue.SetDefaultCellColor(0,0,1.)
transparent_blue.SetCellColorOutputArrayName('scalars')

# ------------------------------------------------------------------
def process(inpoly,filters):
    """
    Build a pipeline from a list of filters and an input vtkPolyData
    Update and return the output
    """
    port=None
    for f in filters:
        if not port:
            f.SetInputData(inpoly)
        else:
            f.SetInputConnection(port)
        port=f.GetOutputPort()
    filters[-1].Update()
    outpoly = vtk.vtkPolyData()
    outpoly.ShallowCopy(filters[-1].GetOutputDataObject(0))
    return outpoly

# ------------------------------------------------------------------
def cylinder():
    """Creates a cylinder vtkPolyData"""
    poly=vtk.vtkPolyData()
    cylinder = vtk.vtkCylinderSource()
    cylinder.SetResolution(20)
    cylinder.SetHeight(2)
    cylinder.SetRadius(0.4)
    cylinder.Update()
    poly.ShallowCopy(cylinder.GetOutput())
    return poly

# ------------------------------------------------------------------
def cube():
    """Creates a cube vtkPolyData"""
    poly=vtk.vtkPolyData()
    cube = vtk.vtkCubeSource()
    cube.Update()
    poly.ShallowCopy(cube.GetOutput())
    return poly

# ------------------------------------------------------------------
def leafObjects( obj ):
    """ Returns a list with all the leaf objects of a multiblock dataset """
    if obj.IsA("vtkDataSet"):
        return [obj]
    leafs=list()
    n = obj.GetNumberOfBlocks()
    for i in range(0,n):
        leafs = leafs + leafObjects( obj.GetBlock(i) )
    return leafs
# ------------------------------------------------------------------

def setScalarMode(mapper):
    """Setup the mapper to use our cell field array"""
    mapper.ScalarVisibilityOn()
    mapper.SetScalarModeToUseCellFieldData()
    mapper.SelectColorArray('scalars')

def addLabel(ren,text):
    """ Add a label actor to the renderer"""
    if text:
        label = vtk.vtkTextActor()
        label.SetInput(text)
        label.SetPosition(3,3)
        label.GetTextProperty().SetFontSize(14)
        label.GetTextProperty().SetColor(0,0,0)
        ren.AddActor2D(label)
# ------------------------------------------------------------------
def compositePolyDataMapper2( objects,text ):
    """
    Render each object by an actor using a vtkCompositePolyDataMapper2.
    If the object is not a vtkMultiBlockDataSet create one with a single block
    """
    fg = vtk.vtkRenderer()
    for obj in objects:
        m = vtk.vtkCompositePolyDataMapper2()
        if obj.IsA("vtkMultiBlockDataSet"):
            m.SetInputDataObject(obj)
        else:
            mbd=vtk.vtkMultiBlockDataSet()
            mbd.SetBlock(0,obj)
            m.SetInputDataObject(obj)
        setScalarMode(m)
        a = vtk.vtkActor()
        a.SetMapper(m)
        fg.AddActor(a)
    addLabel(fg,text)
    print(f"Using {fg.GetActors().GetNumberOfItems()} vtkCompositePolyDataMapper2 actors")
    return fg

# ------------------------------------------------------------------
def compositePolyDataMapper( objects,text ):
    """
    Render each object by an actor using a vtkCompositePolyDataMapper.
    If the object is not a vtkMultiBlockDataSet create one with a single block
    """
    fg = vtk.vtkRenderer()
    for obj in objects:
        m = vtk.vtkCompositePolyDataMapper()
        if obj.IsA("vtkMultiBlockDataSet"):
            m.SetInputDataObject(obj)
        else:
            mbd=vtk.vtkMultiBlockDataSet()
            mbd.SetBlock(0,obj)
            m.SetInputDataObject(obj)
        setScalarMode(m)
        a = vtk.vtkActor()
        a.SetMapper(m)
        fg.AddActor(a)
    addLabel(fg,text)
    print(f"Using {fg.GetActors().GetNumberOfItems()} vtkCompositePolyDataMapper actors")
    return fg

# ------------------------------------------------------------------
def appendIntoPolyDataMapper( objects,text ):
    """
    Render each object by an actor using a vtkPolyDataMapper.
    If the object is a vtkMultiBlockDataSet append all vtkPolyData blocks
    """
    fg = vtk.vtkRenderer()
    for obj in objects:
        m = vtk.vtkPolyDataMapper()
        if obj.IsA("vtkMultiBlockDataSet"):
            append = vtk.vtkAppendPolyData()
            for child in leafObjects(obj):
                append.AddInputData(child)
            append.Update()
            m.SetInputConnection(append.GetOutputPort())
        else:
            m.SetInputData(obj)
        setScalarMode(m)
        a = vtk.vtkActor()
        a.SetMapper(m)
        fg.AddActor(a)
    addLabel(fg,text)
    print(f"Using {fg.GetActors().GetNumberOfItems()} vtkPolyDataMapper actors")
    return fg


def leafPolyDataMapper( objects,text ):
    """
    Render each vtkPolyData by a separate actor using a vtkPolyDataMapper.
    If the object is a vtkMultiBlockDataSet create an actor/mapper for each block
    """
    fg = vtk.vtkRenderer()
    for obj in objects:
        for leaf in leafObjects(obj):
            m = vtk.vtkPolyDataMapper()
            m.SetInputData(leaf)
            setScalarMode(m)
            a = vtk.vtkActor()
            a.SetMapper(m)
            fg.AddActor(a)
    addLabel(fg,text)
    print(f"Using {fg.GetActors().GetNumberOfItems()} vtkPolyDataMapper actors")
    return fg


# ------------------------------------------------------------------
# Define pipelines for applying color and opacity
pipelines = [ ('opaque',           [(cube(),[opaque_blue]),      (cylinder(),[opaque_red])]),
              ('transparent_block',[(cube(),[transparent_blue]), (cylinder(),[opaque_red])]),
              ('transparent',      [(cube(),[transparent_blue]), (cylinder(),[transparent_red])]) ]

# Set up all renderers
renderers = list()
for pipeline in pipelines:
    mbd = vtk.vtkMultiBlockDataSet()
    label = pipeline[0]
    for pipe in pipeline[1]:
        mbd.SetBlock(mbd.GetNumberOfBlocks(),process(pipe[0],pipe[1]))

    renderers.append( [ compositePolyDataMapper2( [mbd],label + ' CPDM2' ),
                        compositePolyDataMapper( [mbd], label + ' CPDM' ),
                        appendIntoPolyDataMapper( [mbd], label + ' appended to PDM' ),
                        leafPolyDataMapper( [mbd], label + ' multiple PDMs' ) ] )
m = len(renderers)
n = len(renderers[0])
renwin = vtk.vtkRenderWindow()
renwin.SetSize(n*300,m*300)

i=0
for scenario in renderers:
    j = 0
    for r in scenario:
        r.SetViewport(float(j)/n,1.-float(i+1)/m,float(j+1)/n,1.-float(i)/m)
        bg = 1.-float(j+2)/(2*n)
        r.SetBackground(bg,bg,bg)
        renwin.AddRenderer( r )
        j = j+1
    i = i+1

iren = vtk.vtkRenderWindowInteractor()
iren.SetRenderWindow(renwin)
iren.Start()