###############################################################################
##
## Copyright (C) 2006-2011, University of Utah. 
## All rights reserved.
## Contact: contact@vistrails.org
##
## This file is part of VisTrails.
##
## "Redistribution and use in source and binary forms, with or without 
## modification, are permitted provided that the following conditions are met:
##
##  - Redistributions of source code must retain the above copyright notice, 
##    this list of conditions and the following disclaimer.
##  - Redistributions in binary form must reproduce the above copyright 
##    notice, this list of conditions and the following disclaimer in the 
##    documentation and/or other materials provided with the distribution.
##  - Neither the name of the University of Utah nor the names of its 
##    contributors may be used to endorse or promote products derived from 
##    this software without specific prior written permission.
##
## THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 
## AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, 
## THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 
## PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR 
## CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 
## EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 
## PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; 
## OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, 
## WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR 
## OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF 
## ADVISED OF THE POSSIBILITY OF SUCH DAMAGE."
##
###############################################################################

##############################################################################
# Transfer Function Widget for VTK

from PyQt4 import QtCore, QtGui
from gui.modules.constant_configuration import ConstantWidgetMixin
from core.modules.basic_modules import new_constant, init_constant, Module
from core.modules.module_registry import get_module_registry
from core.system import get_elementtree_library
ElementTree = get_elementtree_library()
from core.utils.color import ColorByName
import vtk
import math
import pickle
import copy

################################################################################
# etc

def clamp(v, mn, mx, eps=0.0):
    mne = mn + eps
    mxe = mx - eps
    if v < mne: return mn
    if v > mxe: return mx
    return v
# Because of a Qt bug see 
# http://bugreports.qt.nokia.com/browse/QTBUG-17985
# We cannot set the scene from 0 to 1. In this case we will set it
# 4000 x 4000 with GLOBAL_SCALE. When the bug is fixed, just set it to 1.0

GLOBAL_SCALE = 4000.0
##############################################################################
# Transfer Function object

class TransferFunction(object):

    def __init__(self):
        self._min_range = 0.0
        self._max_range = 1.0
        self._pts = []

    def set_range(self, mn, mx):
        self._min_range = mn
        self._max_range = mx

    def set_on_vtk_volume_property(self, vtk_volume_property):
        # Builds the opacity and color functions
        of = vtk.vtkPiecewiseFunction()
        cf = vtk.vtkColorTransferFunction()
        vp = vtk_volume_property
        for pt in self._pts:
            (scalar, opacity, color) = pt
            # Map scalar to tf range
            s = self._min_range + (self._max_range - self._min_range) * scalar
            of.AddPoint(s, opacity)
            cf.AddRGBPoint(s, color[0], color[1], color[2])
        vp.SetScalarOpacity(of)
        vp.SetColor(cf)

    def get_vtk_transfer_functions(self):
        of = vtk.vtkPiecewiseFunction()
        cf = vtk.vtkColorTransferFunction()
        for pt in self._pts:
            (scalar, opacity, color) = pt
            # Map scalar to tf range
            s = self._min_range + (self._max_range - self._min_range) * scalar
            of.AddPoint(s, opacity)
            cf.AddRGBPoint(s, color[0], color[1], color[2])
        return (of,cf)
    
    def add_point(self, scalar, opacity, color):
        self._pts.append((scalar, opacity, color))
        self._pts.sort()

    def get_value(self, scalar):
        """get_value(scalar): returns the opacity and color
        linearly interpolated at the value. Useful for
        adding knots."""
        ix = 0
        while ix < len(self._pts) and self._pts[ix][0] > scalar:
            ix += 1
        if ix == 0:
            return (self._pts[0][1], self._pts[0][2])
        elif ix == len(self._pts):
            return (self._pts[-1][1], self._pts[-1][2])
        else:
            u = ((self._pts[ix][0] - scalar) /
                 (self._pts[ix][0] - self._pts[ix-1][0]))
            do = self._pts[ix][1] - self._pts[ix-1][1]
            dr = self._pts[ix][2][0] - self._pts[ix-1][2][0]
            dg = self._pts[ix][2][1] - self._pts[ix-1][2][1]
            db = self._pts[ix][2][2] - self._pts[ix-1][2][2]
            return (self._pts[ix-1][1] + u * do,
                    (self._pts[ix-1][2][0] + u * dr,
                     self._pts[ix-1][2][1] + u * dg,
                     self._pts[ix-1][2][2] + u * db))

    def __copy__(self):
        result = TransferFunction()
        result._min_range = self._min_range
        result._max_range = self._max_range
        result._pts = copy.copy(self._pts)
        return result

    def __eq__(self, other):
        if type(other) != type(self):
            return False
        if self._min_range != other._min_range:
            return False
        if self._max_range != other._max_range:
            return False
        for my_pt, other_pt in zip(self._pts, other._pts):
            if my_pt != other_pt:
                return False
        return True

    def __ne__(self, other):
        return not self.__eq__(other)

    def serialize(self, node=None):
        """serialize(node: ElementTree.Element) -> str
        Convert this object to an XML representation in string format.
        """
        if node is None:
            node = ElementTree.Element('transfer_function')
            
        node.set('min_range', str(self._min_range))
        node.set('max_range', str(self._max_range))
        for pt in self._pts:
            ptNode = ElementTree.SubElement(node, 'point')
            ptNode.set('scalar', str(pt[0]))
            ptNode.set('opacity', str(pt[1]))
            color = pt[2]
            colorNode = ElementTree.SubElement(ptNode, 'color')
            colorNode.set('R', str(color[0]))
            colorNode.set('G', str(color[1]))
            colorNode.set('B', str(color[2]))
            
        return ElementTree.tostring(node)
   
    @staticmethod         
    def parse(strNode):
        """parse(strNode: str) -> TransferFunction
        Parses a string representing a TransferFunction and returns a 
        TransferFunction object 
        """
        try:
            node = ElementTree.fromstring(strNode)
        except SyntaxError:
            #it was serialized using pickle
            tf = pickle.loads(strNode.decode('hex'))
            tf._pts.sort()
            return tf
        
        if node.tag != 'transfer_function':
            return None
        #read attributes
        tf = TransferFunction()
        tf._min_range = float(node.get('min_range', "0.0")) 
        tf._max_range = float(node.get('max_range', "1.0"))
        for ptNode in node.getchildren():
            if ptNode.tag == 'point':
                scalar = float(ptNode.get('scalar','-1.0'))
                opacity = float(ptNode.get('opacity', '1.0'))
                for colorNode in ptNode.getchildren():
                    if colorNode.tag == 'color':
                        color = (float(colorNode.get('R','0.0')),
                                 float(colorNode.get('G','0.0')),
                                 float(colorNode.get('B','0.0')))
                        break
                tf._pts.append((scalar,opacity,color))
        tf._pts.sort()
        return tf
    
##############################################################################
# Graphics Items

class TransferFunctionPoint(QtGui.QGraphicsEllipseItem):

    selection_pens = { True: QtGui.QPen(QtGui.QBrush(
        QtGui.QColor(*(ColorByName.get_int('goldenrod_medium')))),GLOBAL_SCALE * 0.012),
                       False: QtGui.QPen() }

    def __init__(self, scalar, opacity, color, parent=None):
        QtGui.QGraphicsEllipseItem.__init__(self, parent)
        self._scalar = scalar
        self._opacity = opacity
        self._color = QtGui.QColor(color[0]*255,
                                   color[1]*255,