#!/usr/bin/env python

import sys
import time

from iv import *
from qt import *

class CircularColorMap(object):
    """Circular color map to create a packed color array in view of data
    transfer to Coin3D or Open Inventor.
    
    See: http://astronomy.swin.edu.au/~pbourke/colour/colourramp/
    """

    def __init__(self):
        """Create the lookup table.

        See: http://astronomy.swin.edu.au/~pbourke/colour/colourramp/
        """
        n = 6*256
        r = zeros(n, UInt32)
        g = zeros(n, UInt32)
        b = zeros(n, UInt32)
        up = arange(0, 256, 1, UInt32)
        down = arange(255, -1, -1, UInt32)
        r1, g1, b1 = (0, 1, 1) # cyan
        for i, (r2, g2, b2) in enumerate((
            (0, 0, 1), # blue
            (1, 0, 1), # magenta
            (1, 0, 0), # red
            (1, 1, 0), # yellow
            (0, 1, 0), # green
            (0, 1, 1), # cyan
            )):
            s = slice(i*256, (i+1)*256)
            if r1:
                if r2: r[s] = 255
                else: r[s] = down
            elif r2: r[s] = up
            if g1:
                if g2: g[s] = 255
                else: g[s] = down
            elif g2: g[s] = up
            if b1:
                if b2: b[s] = 255
                else: b[s] = down
            elif b2: b[s] = up
            r1, g1, b1 = r2, g2, b2

        self.__m = (r << 16)  + (g << 8) + b
        self.__m <<= 8
        self.__m += 255

    # __init__()

    def colors(self, data, datamin = None, datamax = None):
        """Map data to a packed color array.
        """
        if datamin is None:
            datamin = minimum.reduce(data)
        if datamax is None:
            datamax = maximum.reduce(data)
        if datamax < datamin:
            datamax, datamin = datamin, datamax
        indices = (((len(self.__m)-1)/(datamax-datamin)) * data).astype(Int32)
        return take(self.__m, indices)

    # colors()

# class CircularColorMap


def sphericalHarmonics(nu, nv, colorMap):
    """Calculate the vertices and colors of a 'spherical harmonics' in view of
    data transfer to Coin3D or Open Inventor.
    
    See: http://astronomy.swin.edu.au/~pbourke/surfaces/sphericalh
    """

    # Coin generates a warning when nu and/or nv are even.
    assert(nu % 2 == 1)
    assert(nv % 2 == 1)

    tick = time.time()
    i = arange(nu*nv)
    u = i % nu
    u %= nu
    u = pi*u/(nu-1)   # phi
    v = i / nu
    v %= nv
    v = 2*pi*v/(nu-1) # theta   
    m = (4, 3, 2, 3, 6, 2, 6, 4)

    r = sin(m[0]*u)**m[1]+cos(m[2]*u)**m[3]+sin(m[4]*v)**m[5]+cos(m[6]*v)**m[7]
    xyzs = zeros((nu*nv, 3), Float)
    xyzs[:, 0] = r*sin(u)*cos(v)
    xyzs[:, 1] = r*sin(u)*sin(v)
    xyzs[:, 2] = r*cos(u)

    colors = colorMap.colors(v, 0.0, 2*pi)
    message = 'Calculating %s vertices and colors took %s seconds' % (
        nu*nv, time.time()-tick)
    
    return xyzs, colors, message

# sphericalHarmonics()


def makeSurface(nu, nv, xyzs, colors):
    """Create a scene graph from vertex and color data."""

    result = SoSeparator()

    shapeHints = SoShapeHints()
    shapeHints.vertexOrdering.setValue(SoShapeHints.COUNTERCLOCKWISE)
    shapeHints.shapeType.setValue(SoShapeHints.UNKNOWN_SHAPE_TYPE)
    result.addChild(shapeHints)

    vertexProperty = SoVertexProperty()
    vertexProperty.materialBinding.setValue(
        SoMaterialBinding.PER_VERTEX_INDEXED)

    vertexProperty.orderedRGBA.setValues(0, colors)
    vertexProperty.vertex.setValues(0, xyzs)

    # Define the QuadMesh.
    quadMesh = SoQuadMesh()
    quadMesh.verticesPerRow.setValue(nu)
    quadMesh.verticesPerColumn.setValue(nv)
    quadMesh.vertexProperty.setValue(vertexProperty)
    result.addChild(quadMesh)

    return result

# makeSurface()


class InventorMainWindow(QMainWindow):

    def __init__(self):
        QMainWindow.__init__(
            self, None, 'SoQt Main Window Example', Qt.WDestructiveClose)

        self.colorMap = CircularColorMap()
        self.setCaption("Spherical Harmonics")

        # Since SoQt widgets are not derived from QWidget,
        # it is sometimes usefull to give them a empty shell QWidget as parent:
        shell = QWidget(self)
        self.viewer = SoQtExaminerViewer(shell)
        # and now
        self.setCentralWidget(shell)
        # instead of self.setCentralWidget(self.viewer) which raises an error.

        self.viewer.setFeedbackVisibility(True)

        toolBar = QToolBar(self)
        QLabel('size:', toolBar)
        sizeComboBox = QComboBox(toolBar)
        for text in ('17 by 17',
                     '33 by 33',
                     '65 by 65',
                     '127 by 127',
                     '257 by 257',
                     '513 by 513',
                     '1025 by 1025',
                     '2047 by 2047',
                     '4097 by 4097'):
            sizeComboBox.insertItem(text)
        sizeComboBox.setCurrentItem(0)
        sizeComboBox.setMaximumWidth(150)
        toolBar.addSeparator()
        self.newSurface(0)
        
        self.connect(sizeComboBox, SIGNAL('activated(int)'), self.newSurface)
        self.resize(600, 400)

    # __init__()
        
    def newSurface(self, index):
        nu = 2**(4+index) + 1
        nv = nu
        xyzs, colors, message = sphericalHarmonics(nu, nv, self.colorMap)
        root = makeSurface(nu, nv, xyzs, colors)
        self.viewer.setSceneGraph(root)
        self.viewer.viewAll()
        self.statusBar().message(message)

    # newSurface()

# InventorMainWindow()


def main():
    app = QApplication(sys.argv)
    SoQt.init(None)
    demo = InventorMainWindow()
    demo.show()

    app.connect(app, SIGNAL('lastWindowClosed()'), app, SLOT('quit()'))
    SoQt.mainLoop()

# main()


if __name__ == '__main__':
    
    if len(sys.argv) != 2:
        print 'Usage: python numpy'
        print 'where numpy must be Numeric, numarray, or numpy'
        sys.exit(1)

    if sys.argv[1] == 'Numeric':
        from Numeric import *
    elif sys.argv[1] == 'numarray':
        from numarray import *
    elif sys.argv[1] == 'numpy':
        from numpy import *

    main()

# Local Variables: ***
# mode: python ***
# End: ***