#!/usr/bin/env python

from vtk import *
import sys

def make_lut(sw, nw, se, ne):

    def lerp(u, v1, v2):
      return u * (v2 - v1) + v1

    def lerp_t(u, t1, t2):
      return tuple([lerp(u, v1, v2) for (v1, v2) in zip(t1, t2)])

    m = vtkMath()

    resolution = 256
    vtkLut = vtkLookupTable()
    vtkLut.SetNumberOfTableValues(resolution * resolution)
    vtkLut.SetTableRange(0, resolution * resolution)

    for i in xrange(resolution):
      u = i / float(resolution-1)
      lerp_north = lerp_t(u, nw, ne)
      lerp_south = lerp_t(u, sw, se)
      for j in xrange(resolution):
        v = j / float(resolution-1)
        uv_lerp = lerp_t(v, lerp_south, lerp_north)
        r, g, b = m.HSVToRGB(uv_lerp)
        vtkLut.SetTableValue(resolution * j + i, r, g, b, 1)
    vtkLut.Build()
    return vtkLut

reader = vtkDataSetReader()
reader.SetFileName(sys.argv[1])

im2c = vtkImageMapToColors()
im2c.SetLookupTable(make_lut(sw=(0.5, 0, 0.5),
                             nw=(0.5, 1, 1),
                             se=(1.0, 0, 0.5),
                             ne=(1.0, 1, 1)))
im2c.AddInputConnection(reader.GetOutputPort())

im2c2 = vtkImageMapToColors()
im2c2.SetLookupTable(make_lut(sw=(0.5, 0, 0.5),
                              nw=(0.5, 1, 1),
                              se=(1.0, 0, 0.5),
                              ne=(1.0, 1, 1)))
im2c2.AddInputConnection(reader.GetOutputPort())

probe = vtkProbeFilter()
probe.SetInputConnection(0, reader.GetOutputPort())
probe.SetInputConnection(1, im2c.GetOutputPort())

probe2 = vtkProbeFilter()
probe2.SetInputConnection(0, reader.GetOutputPort())
probe2.SetInputConnection(1, im2c2.GetOutputPort())

def add_renderer(data, bounds):
    dsm = vtkDataSetMapper()
    dsm.AddInputConnection(data.GetOutputPort())

    actor = vtkActor()
    actor.SetMapper(dsm)

    renderer = vtkRenderer()
    renderer.AddActor(actor)
    renderer.SetViewport(*bounds)

    return renderer

window = vtkRenderWindow()

# When both of the following lines are uncommented, one of the quads
# becomes gray
window.AddRenderer(add_renderer(probe, [0.0, 0.0, 0.5, 1.0]))
window.AddRenderer(add_renderer(probe2, [0.5, 0.0, 1.0, 1.0]))

interactor = vtkRenderWindowInteractor()
window.SetInteractor(interactor)

interactor.Start()
