"""

This module displays cone glyphs scaled and colored as per the
vector or scalar data on a cut plane.  This will work for any
dataset.

This code is distributed under the conditions of the GPL (GNU General
Public License).  Look at the file LICENSE included with the
distribution for details.

Copyright (C) 2001, Prabhu Ramachandran.
"""

__author__ = "Prabhu Ramachandran <prabhu_r@users.sf.net>"
__version__ = "$Revision: 1.4 $"
__date__ = "$Date: 2001/10/02 14:37:32 $"

import Base.Objects, Common
import Tkinter, tkColorChooser, math
import vtkpython
import vtkPipeline.vtkMethodParser

debug = Common.debug

class VectorCutPlane (Base.Objects.CutPlaneModule):

    """ This module displays cone glyphs scaled and colored as per the
    vector or scalar data on cut plane.  This will work for any
    dataset. """

    def __init__ (self, mod_m): 
        debug ("In VectorCutPlane::__init__ ()")
        Common.state.busy ()
        Base.Objects.CutPlaneModule.__init__ (self, mod_m)
        self.cut = vtkpython.vtkCutter ()
        self.glyph = vtkpython.vtkGlyphSource2D ()
	self.cone = vtkpython.vtkConeSource ()
	self.cones = vtkpython.vtkGlyph3D ()
        self.mapper = self.map = vtkpython.vtkPolyDataMapper ()
        self.actor = self.act = vtkpython.vtkActor ()
        # used to orient the cone properly
        self.cone_trfm = vtkpython.vtkTransformFilter ()
        self.cone_trfm.SetTransform (vtkpython.vtkTransform ())
        self.data_out = self.mod_m.GetOutput ()

        # Point of glyph that is attached -- -1 is tail, 0 is center,
        # 1 is head.
        self.glyph_pos = -1 
	self.scale = 1.0
        self.color_mode = 2 #2 is vector, 1 is scalar, -1 none
        self._initialize ()
        self._gui_init ()
        self.renwin.Render ()
        Common.state.idle ()

    def __del__ (self): 
        debug ("In VectorCutPlane::__del__ ()")
        if self.act:
            self.renwin.remove_actors (self.act)
        self.renwin.Render ()

    def _initialize (self):
        debug ("In VectorCutPlane::_initialize ()")
        self.cut.SetInput (self.data_out)
        self.cut.SetCutFunction (self.plane)
        self.glyph.SetGlyphTypeToArrow ()
        self.glyph.SetFilled (0)

	self.cone.SetResolution (5)
	self.cone.SetHeight (1)
	self.cone.SetRadius (0.25)
        self.cone_trfm.SetInput (self.cone.GetOutput ())
        
        self.orient_glyph (1.0, 0.0)
        
	self.cones.SetInput (self.cut.GetOutput ())
	self.cones.SetSource (self.glyph.GetOutput ())
	self.cones.SetScaleFactor (self.scale)
	self.cones.SetScaleModeToScaleByVector ()
	self.cones.SetColorModeToColorByVector ()
	self.cones.SetClamping (1)
        dr = self.mod_m.get_vector_data_range ()
        self.cones.SetRange (dr)

	self.map.SetInput (self.cones.GetOutput ())
        self.act.SetMapper (self.map)
        self.map.SetScalarRange (dr)
        self.map.SetLookupTable (self.mod_m.get_vector_lut ())
        self.act.GetProperty ().SetLineWidth (2)
        self.act.GetProperty ().BackfaceCullingOff ()
        self.act.GetProperty ().FrontfaceCullingOff ()
        apply (self.act.GetProperty ().SetColor, Common.config.fg_color)
        self.center = self.data_out.GetCenter ()
        self.plane.SetOrigin (self.center)
        self.plane.SetNormal (0.0, 0.0, 1.0)
        self.set_def_step_size ()
        self.slider_pos = 0
        self.renwin.add_actors (self.act)
        # used for the pipeline browser
        self.pipe_objs = self.act
        
    def set_def_step_size (self):
        debug ("In VectorCutPlane::set_def_step_size ()")
        out = self.data_out
        bnd = out.GetBounds ()
        n = out.GetNumberOfCells ()
        if not n:
            n = out.GetNumberOfPoints ()
        data_2d = 0
        vol = 1.0
        for i in range (0,3):
            l = abs (bnd[2*i+1] - bnd[2*i])
            if l == 0.0:
                data_2d=1
            else:
                vol = vol*l
        
        l = vol/n
        if data_2d == 1:
            self.step_size = math.sqrt (l)
        else:
            self.step_size = math.pow (l, 1.0/3.0)
        
    def _gui_init (self): 
        debug ("In VectorCutPlane::_gui_init ()")
	self.step_var.set (self.step_size)
	self.n_step_var.set (10)
        self.glyph_var = Tkinter.IntVar ()
        self.glyph_var.set (self.glyph.GetGlyphType ())
        self._auto_sweep_init ()
        self.sweep_step.set (1)

    def SetInput (self, source): 
        debug ("In VectorCutPlane::SetInput ()")
        Common.state.busy ()
        self.data_out = source
        self.cut.SetInput (self.data_out)
        self.do_color_mode ()
        Common.state.idle ()

    def save_config (self, file): 
        debug ("In VectorCutPlane::save_config ()")
        file.write ("%d, %d, %d, %f, %d\n"%(self.glyph_var.get (),
                                        self.slider_pos,
                                        self.n_step_var.get (),
                                        self.step_size, self.glyph_pos))
        p = vtkPipeline.vtkMethodParser.VtkPickler ()
        for obj in (self.plane, self.cut, self.cone, self.glyph, self.cones,
                    self.map, self.act, self.act.GetProperty ()):
            p.dump (obj, file)

    def load_config (self, file): 
        debug ("In VectorCutPlane::load_config ()")
        val = file.readline ()
        glyph_pos = self.glyph_pos
        try:
            g_t, self.slider_pos, n_s, self.step_size,glyph_pos = eval (val)
        except ValueError: # old format 
            g_t, self.slider_pos, n_s, self.step_size = eval (val)
           
        self.glyph_var.set (g_t)
        self.n_step_var.set (n_s)
        p = vtkPipeline.vtkMethodParser.VtkPickler ()
        for obj in (self.plane, self.cut, self.cone, self.glyph, self.cones,
                    self.map, self.act, self.act.GetProperty ()):
            p.load (obj, file)

        self.set_glyph_mode ()
        self.color_mode = self.cones.GetColorMode ()
        if self.color_mode == 0: # I Dislike Scale colouring.
            self.color_mode = 2
        if self.map.GetScalarVisibility () == 0:
            self.color_mode = -1
        self.do_color_mode ()
        self.change_glyph_pos (glyph_pos)

    def config_changed (self): 
        debug ("In VectorCutPlane::config_changed ()")
        apply (self.act.GetProperty ().SetColor, Common.config.fg_color)

    def make_main_gui (self): 
        debug ("In VectorCutPlane::make_main_gui ()")
        self.make_cut_plane_gui ()
        self.make_cone_prop_gui ()
        self.make_choice_gui ()
        self.make_actor_gui (color=0, scalar=0, compact=1)
        self.make_auto_sweep_gui ()

    def make_choice_gui (self):
        debug ("In VectorCutPlane::make_choice_gui ()")
	frame = Tkinter.Frame (self.root, relief='ridge', bd=2)
	frame.pack (side='top', fill='both', expand=1)
        self.make_glyph_gui (frame)
        self.make_color_mode_gui (frame)

    def make_cone_prop_gui (self):
        debug ("In VectorCutPlane::make_cone_prop_gui ()")
	frame = Tkinter.Frame (self.root)
	frame.pack (side='top', fill='both', expand=1)
        self.make_glyph_position_gui (frame)
        self.make_cone_gui (frame)

    def make_glyph_gui (self, master):
        debug ("In VectorCutPlane::make_glyph_gui ()")
	frame = Tkinter.Frame (master, relief='ridge', bd=2)
	frame.pack (side='left', fill='both', expand=1)
        rw = 0
	rb = Tkinter.Radiobutton (frame, text="Normal Arrow",
                                  variable=self.glyph_var, value=9,
                                  command=self.set_glyph_mode)
        rb.grid (row=rw, column=0, sticky='w')
        rw = rw + 1
	rb = Tkinter.Radiobutton (frame, text="Thick Arrow",
                                  variable=self.glyph_var, value=10,
                                  command=self.set_glyph_mode)
        rb.grid (row=rw, column=0, sticky='w')
        rw = rw + 1
	rb = Tkinter.Radiobutton (frame, text="Cone",
                                  variable=self.glyph_var, value=-1,
                                  command=self.set_glyph_mode)
        rb.grid (row=rw, column=0, sticky='w')
        rw = rw + 1
        
    def make_color_mode_gui (self, master):
        debug ("In VectorCutPlane::make_color_mode_gui ()")
	frame = Tkinter.Frame (master, relief="ridge", bd=2)
	frame.pack (side='left', fill='both', expand=1)
	self.color_var = Tkinter.IntVar ()
        self.color_var.set (self.color_mode)
	rb = Tkinter.Radiobutton (frame, text="No Coloring",
                                  variable=self.color_var, value=-1,
                                  command=self.set_color_mode_gui)
	rb.grid (row=0, column=0, sticky='w')
	rb = Tkinter.Radiobutton (frame, text="Scalar Coloring",
                                  variable=self.color_var, value=1,
                                  command=self.set_color_mode_gui)
	rb.grid (row=1, column=0, sticky='w')
	rb = Tkinter.Radiobutton (frame, text="Vector Coloring",
                                  variable=self.color_var, value=2,
                                  command=self.set_color_mode_gui)
	rb.grid (row=2, column=0, sticky='w')
	but = Tkinter.Button (frame, text="Change Color",
                              command=self.set_actor_color)
	but.grid (row=3, column=0, sticky='ew')

    def make_cone_gui (self, master):
        debug ("In VectorCutPlane::make_cone_gui ()")
	frame = Tkinter.Frame (master, relief='ridge', bd=2)
	frame.pack (side='left', fill='both', expand=1)

        rw = 0
	labl = Tkinter.Label (frame, text="Vector scale factor:")
	labl.grid (row=rw, column=0, sticky='w')
	self.scale_var = Tkinter.DoubleVar ()
	self.scale_var.set (self.cones.GetScaleFactor ())
	scale_entr = Tkinter.Entry (frame, width=5, relief='sunken',
                                    textvariable=self.scale_var)
	scale_entr.grid (row=rw, column=1, sticky='we')
	scale_entr.bind ("<Return>", self.change_scale)
        rw = rw + 1

	labl = Tkinter.Label (frame, text="Cone resolution:")
	labl.grid (row=rw, column=0, sticky='w')
	self.cone_res_var = Tkinter.IntVar ()
	self.cone_res_var.set (self.cone.GetResolution ())
        self.cone_entry = []
	entr = Tkinter.Entry (frame, width=5, relief='sunken',
                              state='disabled',
                              textvariable=self.cone_res_var)
	entr.grid (row=rw, column=1, sticky='we')
	entr.bind ("<Return>", self.change_cone)
        self.cone_entry.append (entr)
        rw = rw + 1

	labl = Tkinter.Label (frame, text="Cone height:")
	labl.grid (row=rw, column=0, sticky='w')
	self.cone_ht_var = Tkinter.DoubleVar ()
	self.cone_ht_var.set (self.cone.GetHeight ())
	entr = Tkinter.Entry (frame, width=5, relief='sunken', 
                              state='disabled',
                              textvariable=self.cone_ht_var)
	entr.grid (row=rw, column=1, sticky='we')
	entr.bind ("<Return>", self.change_cone)
        self.cone_entry.append (entr)
        rw = rw + 1

	labl = Tkinter.Label (frame, text="Cone radius:")
	labl.grid (row=rw, column=0, sticky='w')
	self.cone_rad_var = Tkinter.DoubleVar ()
	self.cone_rad_var.set (self.cone.GetRadius ())
	entr = Tkinter.Entry (frame, width=5, relief='sunken', 
                              state='disabled',
                              textvariable=self.cone_rad_var)
	entr.grid (row=rw, column=1, sticky='we')
	entr.bind ("<Return>", self.change_cone)  
        self.cone_entry.append (entr)
        self.update_cone_entry ()

    def make_glyph_position_gui (self, master):
        debug ("In VectorCutPlane::make_glyph_position_gui ()")
	frame = Tkinter.Frame (master, relief='ridge', bd=2)
	frame.pack (side='left', fill='both', expand=1)

        self.glyph_pos_var = Tkinter.IntVar ()
        self.glyph_pos_var.set (self.glyph_pos)
        rw = 0
	labl = Tkinter.Label (frame, text="Glyph position:")
	labl.grid (row=rw, column=0, sticky='ew')
        rw = rw + 1        
	rb = Tkinter.Radiobutton (frame, text="Tail",
                                  variable=self.glyph_pos_var, value=-1,
                                  command=self.set_glyph_pos)
        rb.grid (row=rw, column=0, sticky='w')
        rw = rw + 1
	rb = Tkinter.Radiobutton (frame, text="Center",
                                  variable=self.glyph_pos_var, value=0,
                                  command=self.set_glyph_pos)
        rb.grid (row=rw, column=0, sticky='w')
        rw = rw + 1
	rb = Tkinter.Radiobutton (frame, text="Head",
                                  variable=self.glyph_pos_var, value=1,
                                  command=self.set_glyph_pos)
        rb.grid (row=rw, column=0, sticky='w')
        rw = rw + 1        

    def change_scale (self, event=None):
        debug ("In VectorCutPlane::change_scale ()")
        Common.state.busy ()
	self.cones.SetScaleFactor (self.scale_var.get ())
	self.renwin.Render ()
        Common.state.idle ()

    def change_cone (self, event=None):
        debug ("In VectorCutPlane::change_cone ()")
        Common.state.busy ()
	res = self.cone_res_var.get ()
	ht = self.cone_ht_var.get ()
	rad = self.cone_rad_var.get ()
        self.orient_glyph (ht, self.cone.GetHeight ())
	self.cone.SetResolution (res)
	self.cone.SetHeight (ht)
	self.cone.SetRadius (rad)
	self.renwin.Render ()    
        Common.state.idle ()
    
    def set_color_mode_gui (self, event=None):
        "This sets up the data to setup the actual coloring mode."
        debug ("In VectorCutPlane::set_color_mode_gui ()")
        Common.state.busy ()
        self.color_mode = self.color_var.get ()
        self.do_color_mode ()
        self.renwin.Render ()
        Common.state.idle ()

    def do_color_mode (self):
        debug ("In VectorCutPlane::do_color_mode ()")
        self.cones.SetRange (self.mod_m.get_vector_data_range ())
        if self.color_mode == 2: # Vector Coloring
            dr = self.mod_m.get_vector_data_range ()
            self.cones.SetColorModeToColorByVector ()
            self.map.ScalarVisibilityOn ()
            self.map.SetScalarRange (dr)
            self.map.SetLookupTable (self.mod_m.get_vector_lut ())
        elif self.color_mode == 1: # Scalar Coloring
            self.cones.SetColorModeToColorByScalar ()
            self.map.ScalarVisibilityOn ()
            dr = self.mod_m.get_scalar_data_range ()
            self.map.SetScalarRange (dr)
            self.map.SetLookupTable (self.mod_m.get_scalar_lut ())
        elif self.color_mode == -1: # No colouring
            self.map.ScalarVisibilityOff ()

    def set_actor_color (self, event=None):
        debug ("In VectorCutPlane::set_actor_color ()")
	clr = self.act.GetProperty ().GetColor ()
	init_clr = "#%02x%02x%02x"%(clr[0]*255, clr[1]*255, clr[2]*255)
	color = tkColorChooser.askcolor (title="Change object color", 
					 initialcolor=init_clr)
	if color[1] is not None:
            Common.state.busy ()
	    clr = Common.tk_2_vtk_color (color[0])
	    apply (self.act.GetProperty ().SetColor, clr)
	    self.renwin.Render ()        
            Common.state.idle ()
            if self.color_mode != -1:
                msg = "Warning: This setting will have effect only if "\
                      "the coloring mode is set to 'No Coloring'."
                Common.print_err (msg)

    def do_sweep (self, event=None):
        debug ("In VectorCutPlane::do_sweep ()")
        if self.sweep_var.get ():
            val = int (1000*self.sweep_delay.get ())
            self.root.after (val, self.update_sweep)

    def update_sweep (self):
        debug ("In VectorCutPlane::update_sweep ()")
        if self.sweep_var.get ():
            d_pos = self.sweep_step.get ()
            pos = self.slider_pos + d_pos
            if (d_pos > 0) and (pos > self.n_step_var.get ()):
                pos = -self.n_step_var.get ()
            elif (d_pos < 0) and (pos < -self.n_step_var.get ()):
                pos = self.n_step_var.get ()
            self.slider[2].set (pos)
            self.change_slider ()            
            val = int (1000*self.sweep_delay.get ())
            self.root.after (val, self.update_sweep)

    def set_glyph_mode (self, event=None):
        debug ("In VectorCutPlane::set_glyph_mode ()")
        Common.state.busy ()
        val = self.glyph_var.get ()
        if val == 9: # normal arrow
            self.glyph.SetGlyphTypeToArrow ()
            self.glyph.SetFilled (0)
            self.cones.SetSource (self.glyph.GetOutput ())
        elif val == 10: # Thick arrow
            self.glyph.SetGlyphTypeToThickArrow ()
            self.glyph.SetFilled (1)
            self.cones.SetSource (self.glyph.GetOutput ())
        elif val == -1: # Cone
            self.cones.SetSource (self.cone_trfm.GetOutput ())

        self.update_cone_entry ()
        self.renwin.Render ()
        Common.state.idle ()

    def update_cone_entry (self):
        debug ("In VectorCutPlane::update_cone_entry ()")
        try:
            val = self.glyph_var.get ()
        except AttributeError:
            pass
        else:
            try:
                self.cone_entry[0]
            except AttributeError:
                pass
            else:
                for entry in self.cone_entry:
                    if val == -1:
                        entry.config (state='normal')
                    else:
                        entry.config (state='disabled')

    def set_glyph_pos (self, event=None):
        debug ("In VectorCutPlane::change_glyph_pos ()")        
        Common.state.busy ()
        self.change_glyph_pos (self.glyph_pos_var.get ())
        self.renwin.Render ()
        Common.state.idle ()        
        
    def change_glyph_pos (self, new_pos):
        debug ("In VectorCutPlane::change_glyph_pos ()")
        old_pos = self.glyph_pos
        if old_pos == new_pos:
            return
        else:
            self._center_cone()
            self.glyph_pos = new_pos
            self.orient_glyph (self.cone.GetHeight (), 0.0)

    def _center_cone (self):
        debug ("In VectorCutPlane::orient_glyph ()")
        tr = self.cone_trfm.GetTransform ()
        if self.glyph_pos == -1:
            tr.Translate (-self.cone.GetHeight ()*0.5, 0.0, 0.0)
        elif self.glyph_pos == 1:
            tr.Translate (self.cone.GetHeight ()*0.5, 0.0, 0.0)
    
    def orient_glyph (self, new_ht, old_ht):
        debug ("In VectorCutPlane::orient_glyph ()")
        tr = self.cone_trfm.GetTransform ()
        if self.glyph_pos == -1:
            tr.Translate ((new_ht - old_ht)*0.5, 0.0, 0.0)
            self.glyph.SetCenter (0.5, 0.0, 0.0)
        elif self.glyph_pos == 1:
            tr.Translate (-(new_ht - old_ht)*0.5, 0.0, 0.0)
            self.glyph.SetCenter (-0.5, 0.0, 0.0)
        else:
            self.glyph.SetCenter (0.0, 0.0, 0.0)
