# f(x,y) plotter
#  by Mark Dewing
#  version 0.1


import vtk
from math import *
from Tkinter import *
import vtk.tk.vtkTkRenderWindowInteractor


root = Tk()
display_frame = Frame(root)
render_frame = Frame(display_frame)
tools_frame = Frame(display_frame)

ren = vtk.vtkRenderer()
renWin = vtk.vtkRenderWindow()
renWin.AddRenderer(ren)

render_widget = vtk.tk.vtkTkRenderWindowInteractor.\
vtkTkRenderWindowInteractor(render_frame,rw=renWin,width=300,height=300)

func=".5*y*cos(3*x*pi)"

label = Label(tools_frame,text='Function f(x,y) = ')
label.pack(side=LEFT)

entry = Entry(tools_frame)
entry.pack(side=LEFT)
entry.insert(0,func)


plane = vtk.vtkPlaneSource()
plane.SetXResolution(100)
plane.SetYResolution(100)

transform = vtk.vtkTransform()

trans_filter = vtk.vtkTransformPolyDataFilter()
trans_filter.SetInput(plane.GetOutput())
trans_filter.SetTransform(transform)

func_filter = vtk.vtkProgrammableFilter()
func_filter.SetInput(trans_filter.GetOutput())



class plots:
 def plot_func(self):
   input =  func_filter.GetPolyDataInput()
   npts = input.GetNumberOfPoints()
   newpts = vtk.vtkPoints()
   vals = vtk.vtkFloatArray()

   for i in range(0,npts):
      coord = input.GetPoint(i)
      x,y = coord[:2]
      f=eval(self.func)
      #f = x*y
      newpts.InsertPoint(i,x,y,f)
      vals.InsertValue(i,f)

   func_filter.GetPolyDataOutput().CopyStructure(input)
   func_filter.GetPolyDataOutput().SetPoints(newpts)
   func_filter.GetPolyDataOutput().GetPointData().SetScalars(vals)

plt=plots()
plt.func = func

def change_entry():
   func=entry.get()
   plt.func = func
   print func
   func_filter.Modified()
   render_widget.Render()
   


button = Button(tools_frame,text='Update',command=change_entry)
button.pack(side=RIGHT)
tools_frame.pack(side=TOP)
    
func_filter.SetExecuteMethod(plt.plot_func)

warp = vtk.vtkWarpScalar()
warp.SetInput(func_filter.GetPolyDataOutput())
warp.XYPlaneOn()

plane_mapper = vtk.vtkPolyDataMapper()
plane_mapper.SetInput(warp.GetPolyDataOutput())
plane_mapper.SetScalarRange(func_filter.GetPolyDataOutput().GetScalarRange())

carpet = vtk.vtkActor()
carpet.SetMapper(plane_mapper)

axes = vtk.vtkAxes()
axes_mapper = vtk.vtkPolyDataMapper()
axes_mapper.SetInput(axes.GetOutput())
axes_actor = vtk.vtkActor()
axes_actor.SetMapper(axes_mapper)

ren.AddActor(carpet)
ren.AddActor(axes_actor)

camera = ren.GetActiveCamera();
camera.Dolly(.5)

display_frame.pack(expand=YES,fill=BOTH)
render_widget.pack(expand=YES,fill=BOTH)
render_frame.pack(expand=YES,fill=BOTH)
root.mainloop()