## Automatically adapted for numpy.oldnumeric Jul 23, 2007 by 

#############################################################################
#
# Author: Anna Omelchenko
#
# Copyright: M. Sanner TSRI
#
#############################################################################

#
#$Header: /opt/cvs/Volume/Pvv/transferCommands.py,v 1.4 2007/09/27 20:33:00 annao Exp $
#
#$Id: transferCommands.py,v 1.4 2007/09/27 20:33:00 annao Exp $
#


"""This module implements a command to edit the color and opacity
transfer functions.
The function is drawn at the top of the interface as a series of
circles (points) connected by line segments. The X location of each point
indicates the data value, the Y location indicates the opacity of
that data value. Clicking the left mouse button on a point selects it.
Dragging the point with the left mouse button modifies the opacity. Left
clicking or moving the circle in the area labeled Hue/Saturation
changes the color of the selected point. The red-green-blue-opacity
values can be directly typed in the text entries.
Shift-left button click in the transfer function window adds a new set of
four points (a shape) to the function. Middle clicking and draggind a
point of a shape moves the shape left or right in the window.
Double left click on a line connecting two points adds a point to the
shape. Dragging a point up out of the window deletes it.
The 'Split Interval' button allows to zoom in on a particular interval of
values. The user must enter starting and ending value points of a
new interval in the pop up form. This will result in splitting the
initial interval of values into two or three intervals that will appear
in separate transfer function windows.
The 'Reset' button sets the lookup table to grayscale-ramp color values
and linear opacity values.
The 'Save TF' button allows to save current color and opacity values to
a file.
The 'Save LUT' button allows to save the lookup table to a .lut file, which
can be used later to restore current transfer function.
The 'Load LUT' button loads a .lut file.
An ISO value can be added by pressing 'Add new' button.
In the pop up dialog the user can specify the scalar value,
opacity and RGB values for the new ISO point.
The RGBA values of selected ISOval can be typed in the corresponding
entry fields."""

from ViewerFramework.VFCommand import Command,CommandGUI
import numpy.oldnumeric as Numeric
import Pmw
import string, os
from mglutil.gui.BasicWidgets.Tk import tablemaker
import Tkinter
import tkFileDialog
from DejaVu import colorTool
import traceback

volGeom = None

class TransferCommand(Command):
    """Command to edit the color and opacity transfer functions."""

    def __init__(self):
        Command.__init__(self)
        self.isDisplayed=0
        self.isWithdrawn = 0
        self.root = None
        self.maxval = 0
        self.minval = 0
        self.alphamaxval = 4095
        self.lut = None
        self.tm = None
        self.file = None
        self.modifcallback = None # function that need to be set
                                  # by other application (ie PyARTK)
                                  # and will be call if new transfer
                                  # function is apply to volume
                                  


    def SetCallback(self,func):
        """set a callback fuction"""
        assert callable(func)
        self.modifcallback = func
        
    def OnAddVolumeToViewer(self):
        """Called when a volume is loaded into the viewer. Builds
        or updates the input form (GUI) and withdraws it"""
        
        if not volGeom.volume: return
        if self.isNewForm():
                self.root = Tkinter.Toplevel(self.master)
                self.root.withdraw()
                self.root.title('Transfer Function')
                self.isDisplayed=0
                if self.lut:
                    self.lut.Release()
                if self.tm:
                    del self.tm
                    #remove entries from undoCmdStack list
                    if len(self.vf.undoCmdStack):
                        ent = []
                        for c in self.vf.undoCmdStack:
                            if len(c) > 1:
                                if c[1] == 'loadLUT' or c[1]=='split interval':
                                    ent.append(c)
                        if len(ent):
                            for e in ent:
                                self.vf.undoCmdStack.remove(e)
                            self.vf.undo.setLabel()
                self.buildInputForm()
                self.isWithdrawn = 1
                
    def isNewForm(self):
        """Check if we need a new input form when a new volume is
        loaded. """
        #Has to be implemented by a derived class. 
        pass
    
    def onAddCmdToViewer(self):
        """Called when the command is loaded (added) to the viewer."""
        #vi = self.vf.GUI.VIEWER
        self.master = self.vf.GUI.ROOT
        volGeom.onaddVolume_list.append(self)
        if not volGeom.volume:
            return
        if self.isWithdrawn:
            return
        #self.root=Tkinter.Toplevel()
        self.root = Tkinter.Toplevel(self.master)
        self.root.withdraw()
        self.root.title('Transfer Function')
        self.buildInputForm()
        self.isWithdrawn = 1


    def buildInputForm(self):
        """Places Tkinter widgets on the input form."""
        self.set_maxminval()
        self.tm = tablemaker.TableManager(self.vf, self.root,
                                          xmaxval = self.maxval-1,
                                          xminval = self.minval,
                                          ymaxval = self.alphamaxval,
                                          alphaCallback=self.setAlpha,
                                          colorCallback=self.setColors)

        self.tm.execute_split = self.execute_split
        self.tm.load_lutfile = self.load_lutfile
        font = '-*-Helvetica-Bold-R-Normal-*-*-160-*-*-*-*-*-*'
##          Tkinter.Button(self.tm.f2, text='Load LUT',
##                         command=self.ask_open_file_cb,
##                         takefocus=0, font=font).pack(side=Tkinter.LEFT)
        f = Tkinter.Frame(self.root)
        f.pack(side=Tkinter.BOTTOM)
        Tkinter.Button(f, text='Dismiss',
               command=self.destroy, font=font).pack(side=Tkinter.TOP)

    def set_maxminval(self, minval=0):
        """Sets self.maxval to the number of LUT entries.
        Sets self.minval.
        """
        #Has to be implemented by a derived class. 
        pass


    def execute_split(self, entry_minval, entry_maxval, split=1):
        """Calls doitWrapper() when 'OK' or 'Cancel' button of the
        'Split dialog' is pressed."""
        self.doitWrapper(entry_minval, entry_maxval, split=split)

    def load_lutfile(self, file):
        """Calls doitWrapper() when the 'Load LUT' button is pressed. """
        
        self.doitWrapper(file)
        
    def doit(self, *args, **kw):
        """ Loads the selected .lut file or splits a value interval."""
        if not volGeom.volume:
            print "Volume has not been loaded."
            return
        if len(args) == 1:
            self.file = args[0]
            if args[0]:
                self.tm.load_file(args[0])
            else:
                if len(self.tm.data):
                    self.tm.undo_split()
        elif len(args) == 2 :
            if kw.get('split')==1:
                self.tm.split_cb(args[0],args[1])
            else:
                self.tm.undo_split()
                
    def setupUndoBefore(self, *args, **kw):
        """Memorizes current data before the doit() method is
        called """
        
        if not volGeom.volume:
            return
        tr = self
        if len(args) == 1:
            self.undoMenuString = 'loadLUT'
##              if kw.get('only') == 1:
##                  self.tm.saveData()
##              kw['only'] = not kw.get('only')
##              if self.file:
##                  self.addUndoCall((self.file,), kw, tr.name)
##              elif not self.file or len(self.tm.data) == 0:
##                  self.tm.saveData()
##                  kw['only'] = 0
##                  self.addUndoCall((args[0],), kw, tr.name)
            if not self.file or len(self.tm.data) == 0:
                self.tm.saveData()
            self.addUndoCall((self.file,), kw, tr.name)
        elif len(args) == 2:
            self.undoMenuString = 'split interval'
            if kw.get('split')==1:
                self.tm.saveData()
                #print "in setupUndoBefore, intervals_list:", self.tm.intervals_list
                #print "data['intervals']=", self.tm.data[-1]['intervals']
            kw['split'] = not kw.get('split')
            self.addUndoCall(args, kw, tr.name)

    def __call__(self, *args, **kw):
        """None <- Transfer(filename) -- loads new lookup table
           None <- Transfer(minval, maxval, split = 1) -- splits an interval
        filename: name of a file with .lut extension;
        minval:   min. splitting point(an int.);
        maxval:   max. splitting point(an int.) -- the length of the resulting
                  intervals must be 49 or more;
        split:   1 - to split an interval, 0 - to undo splitting. """

        if len(args)==1:
            file = args[0]
            if file:
                file = string.strip(file)
                if os.path.splitext(file)[-1] != '.lut':
                    print "Wrong file name", file
                    return
            apply(self.doitWrapper,(file,), kw)

        elif len(args)==2:
            minval = args[0]
            maxval = args[1]
            kw['split'] = kw.get('split', 1)
            if kw['split']:
                if (maxval-minval)<49:
                    print "(maxval - minval) must be greater than 49"
                    return
                parent_interval = None
                for i in self.tm.intervals_list:
                    if minval == i[0]:
                        if maxval == i[1]:
                            print "interval (%d,%d) exists" % (minval, maxval)
                            return
                        if (i[1]-maxval)>=49:
                            parent_interval = i
                    elif maxval == i[1]:
                        if (minval-i[0])>=49:
                            parent_interval = i
                    else:
                        if (minval-i[0])>=49 and (i[1]-maxval)>=49:
                            parent_interval = i
                if parent_interval :
                    self.tm.parent_interval = parent_interval
                    apply(self.doitWrapper,args, kw)
                else:
                    if minval < 0 or maxval > (self.maxval-1):
                        print "wrong minval or maxval: valid range: (%d,%d)"\
                              %(0,self.maxval-1)
                    print "wrong minval or maxval:the length of the resulting intervals must be 49 or more"
            elif kw['split'] == 0:
               apply(self.doitWrapper,args, kw) 

    def setAlpha(self, value):
        """Modifies the alpha(opacity) table entries for given interval
        of data values. value is a list: value[0] - starting index,
        value[1] - an integer array containing new alpha values.
        Each entry is in range 0...4095."""
        #Has to be implemented by a derived class. 
        pass

    def setColors(self, value):
        """Modifies the color table entries for given interval of data
        values. value is a list: value[0] - starting index,
        value[1] - an  array of RGB values.
        Each entry is in range 0.0 ... 1.0"""
        #Has to be implemented by a derived class. 
        pass
                                                   
    def destroy(self):
        """Withdraws the input form (GUI)."""
        self.root.withdraw()
        self.isWithdrawn = 1
        self.isDisplayed=0

    def guiCallback(self, event=None, *args, **kw):
        """Creates or deiconifies an input form (GUI). Called
        evry time the 'Transfer Function' button is pressed."""
        
        if not volGeom.volume:
            print "volume has not been loaded"
            return
        
        if self.isDisplayed:
            self.root.lift()
            return
        self.isDisplayed=1
        if self.isWithdrawn:
            self.root.deiconify()
            self.isWithdrawn = 0
        else:
            self.root=Tkinter.Toplevel(self.master)
            root = self.root
            root.title('Transfer Function')
            self.buildInputForm()

  



class VLITransfer(TransferCommand):
    """Command to edit VLI transfer function (lookup table) to map
    the data into color and opacity values.
    """
    
    def isNewForm(self):
        """Check if we need a new input form when a new volume is
        loaded. """
        
        from Volume.Renderers.VLI import vli
        format = volGeom.volume.GetFormat()
        if format == vli.kVLIVoxelFormatUINT8:
            maxval = 256
        elif format == vli.kVLIVoxelFormatUINT12L:
            maxval = 4096
        elif format == vli.kVLIVoxelFormatUINT12U:
            maxval = 65536
        needForm = 0
        if self.root:
            if maxval != self.maxval:
                self.maxval = maxval
                needForm = 1
                self.root.destroy()
        else:
            needForm = 1
            if maxval != self.maxval:  self.maxval = maxval
        return needForm

    def set_maxminval(self, minval=0):
        """Sets self.maxval to the number of LUT entries.
        Creates an instance of VLILookupTable.
        """
        from Volume.Renderers.VLI import vli
        if self.maxval == 0:
            format = volGeom.volume.GetFormat()
            if format == vli.kVLIVoxelFormatUINT8: 
                self.maxval = 256
            elif format == vli.kVLIVoxelFormatUINT12L:
                self.maxval = 4096
            elif format == vli.kVLIVoxelFormatUINT12U:
                self.maxval = 65536
        self.lut = vli.VLILookupTable_Create(self.maxval)
        volGeom.context.SetLookupTable(self.lut)
        if minval:
            self.minval = minval


    def setAlpha(self, value):
        """Modifies the alpha(opacity) table entries for given interval
        of data values. value is a list: value[0] - starting index,
        value[1] - an integer array containing new alpha values.
        Each entry is in range 0...4095."""
        
        self.lut.setAlphaEntries(value[1], int(value[0]))
        vi = self.vf.GUI.VIEWER
        volGeom.context.SetLookupTable(self.lut)
        vi.Redraw()

    def setColors(self, value):
        """Modifies the color table entries for given interval of data
        values. value is a list: value[0] - starting index,
        value[1] - an  array of RGB values.
        Each entry is in range 0.0 ... 1.0"""
        
        arr_len = len(value[1])
        rgb = (value[1]*255).astype(Numeric.UnsignedInt8)
        self.lut.setColorEntries(value[0], rgb)
        volGeom.context.SetLookupTable(self.lut)
        self.vf.GUI.VIEWER.Redraw()

    def getTransferFunc(self):
        """Returns 3 lists containing LUT values, corresponding
        alphas and colors."""
        values = []
        alphas = []
        colors = []
        LUT = volGeom.context.GetLookupTable()
        for lut in self.tm.lut_list:
            values.extend(lut.values[1:-1])
            xminval = lut.xminval
            for v in lut.values[1:-1]:
                alphas.append(LUT.getAlphaEntry(v-xminval))
                colors.append([ LUT.getColorEntry(v-xminval,0),
                                LUT.getColorEntry(v-xminval,1),
                                LUT.getColorEntry(v-xminval,2)])
        return values, alphas, colors

    def getRGB(self,x1, x2):
        """Returns RGB values of the selected interval of LUT entries. """
        rgb = []
        LUT = volGeom.context.GetLookupTable()
        for i in range(x1, x2+1):
           rgb.append([ LUT.getColorEntry(i,0),
                        LUT.getColorEntry(i,1),
                        LUT.getColorEntry(i,2)])
        return rgb
                
##      def getLookupTable(self):
##          """build a list of (rgba) tuples for the current lookup table)"""
##          lut=volGeom.context.GetLookupTable()
##          size=lut.GetSize()
##          lutdata=[]
##          for i in range(size):
##              r=lut.getColorEntry(i,0)
##              g=lut.getColorEntry(i,1)
##              b=lut.getColorEntry(i,2)
##              a=lut.getAlphaEntry(i)
##              lutdata.append((r,g,b,a))
##          return lutdata
                
class UTVolRenTransfer(TransferCommand):
    """Command to edit transfer function (lookup table) to map
    the data into color and opacity values ( Univ. of Texas Volume Library). 
    """
    def isNewForm(self):
        """Check if we need a new input form when a new volume is
        loaded. """

        self.maxval = 256
        needForm = 0
        if self.root:
            if volGeom.byte_map is not None:
                volGeom.volume.uploadColorMap(volGeom.byte_map)
            else:
                volGeom.grayRamp()
        else:
            needForm = 1
            volGeom.grayRamp()
        return needForm
    
    def set_maxminval(self, minval = 1):
        """Sets self.maxval to the number of LUT entries."""
        self.maxval =  256
        self.alphamaxval = 255
        if minval:
            self.minval=minval

    def setAlpha(self, value):
        volGeom.setVolRenAlpha(value)
        try:
            if self.modifcallback:
                apply(self.modifcallback)
        except:
                print 'ERROR ************************************************'
                traceback.print_exc()
                return 'ERROR'

    def setColors(self, value):
        volGeom.setVolRenColors(value)
        try:
            if self.modifcallback:
                apply(self.modifcallback)
        except:
                print 'ERROR ************************************************'
                traceback.print_exc()
                return 'ERROR'
        
commandGUI = CommandGUI()
menuName = 'Transfer Function'
menuIndex = 2
commandNames = {'vli':'VLITransfer', 'utvolren': 'UTVolRenTransfer'}
commandList = [{'name': '', 'cmd': None,
                  'gui':None}]

def initModule(viewer):
    for dict in commandList:
        if dict['name']:
            viewer.addCommand( dict['cmd'], dict['name'], dict['gui'])

