import numpy
from numpy import *
from CGLutil import vrml

#=====================
maxHeight = 64.
# maxHeight does not affect the plot unless the largest bucket value
# is higher than maxHeight in which case the output is scaled so that
# the "tallest spire" is only maxHeight units high.
#=====================

# -----------------------------------------------------------------------------
# Return VRML string that draws a set of boxes.
def vr_plot(bucketCounts):                   
    baseColorCodes = [[2,3,3,3,3,3,3,3,3,3,3,2,2,2,1,1,0,0,0,0,0,1,1,1,1,1,0,0,0,0,0,1,1,1,1,2],
                      [2,3,3,3,3,3,3,3,3,3,3,3,2,2,2,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1,2,2],
                      [2,3,3,3,3,3,3,3,3,3,3,3,2,2,2,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,1,1,2,2,2],
                      [2,2,3,3,3,3,3,3,3,3,3,3,3,2,2,2,2,1,1,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1,2,2],
                      [2,2,3,3,3,3,3,3,3,3,3,3,3,3,2,2,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1,1,2],
                      [2,2,3,3,3,3,3,3,3,3,3,3,3,2,2,2,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1,2],
                      [2,2,2,3,3,3,3,3,3,3,3,3,2,2,2,2,1,1,1,0,0,0,1,1,1,1,1,1,0,0,0,0,0,1,1,2],
                      [2,2,2,2,3,3,3,3,3,3,3,2,2,2,2,2,2,1,1,1,1,1,1,1,1,1,1,1,0,0,0,0,0,1,1,2],
                      [2,2,2,2,2,3,2,2,3,2,2,2,2,2,2,2,2,1,1,1,1,1,1,1,2,2,1,1,0,0,0,0,0,1,1,2],
                      [2,2,2,2,2,2,2,2,2,2,2,2,2,2,1,1,1,1,1,1,1,2,2,2,2,1,1,1,0,0,0,0,0,1,1,1],
                      [2,2,2,2,2,2,2,2,2,2,2,2,1,1,1,1,1,1,1,1,1,2,2,2,2,1,1,1,0,0,0,0,0,1,1,1],
                      [1,2,2,2,2,2,2,2,2,2,2,2,2,1,1,0,0,0,0,1,1,1,2,2,2,1,1,1,1,0,0,0,0,1,1,1],
                      [1,1,2,2,2,2,2,2,2,2,2,2,2,1,1,0,0,0,0,1,1,2,2,2,2,2,2,1,1,0,0,0,0,1,1,1],
                      [1,1,2,2,2,2,2,2,2,2,2,2,1,1,1,0,0,0,0,1,1,2,2,3,2,2,1,1,1,0,0,0,0,1,1,1],
                      [1,2,2,2,2,2,2,2,2,2,2,2,1,1,1,0,0,0,0,1,1,1,2,3,3,2,1,1,1,0,0,0,0,1,1,1],
                      [1,2,2,2,2,2,2,3,2,2,2,2,2,1,1,1,0,0,0,1,1,1,2,2,2,2,2,1,1,0,0,0,0,1,1,1],
                      [1,2,2,2,2,2,3,3,3,3,2,2,2,1,1,1,1,0,0,0,1,1,1,2,2,2,1,1,1,0,0,0,0,1,1,1],
                      [1,2,2,2,2,2,3,3,3,3,3,2,2,2,1,1,1,0,0,0,1,1,1,1,2,2,2,1,1,0,0,0,0,1,1,1],
                      [1,2,2,2,2,3,3,3,3,3,3,3,2,2,2,1,1,1,0,0,1,1,1,2,1,2,2,1,1,0,0,0,0,1,1,1],
                      [1,2,2,2,2,2,3,3,3,3,3,3,3,2,2,1,1,1,1,0,1,1,1,1,1,1,1,1,1,0,0,0,0,1,1,1],
                      [1,1,2,2,2,2,3,3,3,3,3,3,3,3,2,2,1,1,1,0,1,1,1,1,1,1,1,1,1,0,0,0,0,1,1,1],
                      [1,1,2,2,2,2,2,3,3,3,3,3,3,3,2,2,2,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1],
                      [1,1,1,2,2,2,2,2,3,3,3,3,3,3,3,2,2,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1],
                      [1,2,2,2,2,2,2,2,2,3,3,3,3,3,3,2,2,2,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1],
                      [2,2,1,2,2,2,2,2,2,2,2,3,3,3,3,2,2,2,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1],
                      [1,1,1,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,1,1,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1],
                      [1,1,1,1,1,2,2,2,2,2,2,2,2,2,2,2,2,2,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1],
                      [0,1,1,1,1,2,2,2,2,2,2,2,2,2,2,1,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
                      [0,0,0,1,1,1,1,1,1,1,1,2,1,1,1,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
                      [0,0,0,1,1,1,1,2,2,1,1,1,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
                      [1,1,1,1,1,2,2,1,1,1,1,1,1,1,0,0,0,0,0,0,0,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0],
                      [1,1,1,1,1,1,2,2,2,2,1,1,1,1,0,0,0,0,0,0,0,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0],
                      [1,2,2,2,1,2,2,2,2,2,1,1,1,1,1,0,0,0,0,0,0,1,1,2,1,1,1,0,0,0,0,0,0,0,0,0],
                      [1,2,2,2,2,2,2,2,2,2,2,1,1,1,1,0,0,0,0,0,0,1,1,2,1,1,1,0,0,0,0,0,0,1,1,1],
                      [1,2,2,2,2,2,2,2,2,2,2,2,2,1,1,0,0,0,0,0,0,1,1,2,2,1,1,0,0,0,0,0,0,1,1,1],
                      [2,2,2,2,2,2,2,2,2,2,2,2,2,1,1,0,0,0,0,0,0,1,1,1,1,1,1,0,0,0,0,0,0,1,1,2]]

    wrl = vrml.Transform()
    
    baseColors = ((1.0, 1.0, 0.5), (1.0, 0.8, 0.3), (0.8, 0.6, 0.1), (1.0, 0.0, 0.0))
    baseHeights = (0.6, 0.7, 0.8, 1.00)
    
    for ii in range(36):
        for jj in range(36):
            # Put in cubes making up the base:
            bCC = baseColorCodes[ii][jj]
            translateBaseCube = vrml.Transform(translation = tuple([jj, 35-ii, 0.0]))
            base = vrml.Box(size = (1, 1, baseHeights[bCC]), color = baseColors[bCC])
            translateBaseCube.addChild(base)
            wrl.addChild(translateBaseCube)

            # Put in boxes for bucket values:
            boxHeight = bucketCounts[ii][jj]
            translateBucketBox = vrml.Transform(translation = tuple([jj, 35-ii, boxHeight/2.0]))
            bucketBox =vrml.Box(size = (1, 1, boxHeight), color = baseColors[bCC])
            translateBucketBox.addChild(bucketBox)
            wrl.addChild(translateBucketBox)

    transYaxis = vrml.Transform(translation = tuple([17.5, 17.0, 0.5]))
    yAxisBox = vrml.Box(size = (0.1, 37, .1), color = (0,0,0))
    transYaxis.addChild(yAxisBox)
    wrl.addChild(transYaxis)

    trYticks = []
    for ii in range(9):
        xpos = -0.5 + ii*4.5
        trYticks.append(vrml.Transform(translation = tuple([xpos, 17.0, 0.4])))
        yTickBox = vrml.Box(size = (0.06, 37, .1), color = (.2,.2,.2))
        trYticks[ii].addChild(yTickBox)
        wrl.addChild(trYticks[ii])
        
    transXaxis = vrml.Transform(translation = tuple([17.0, 17.5, 0.5]))
    xAxisBox = vrml.Box(size = (37, 0.1, .1), color = (0,0,0))
    transXaxis.addChild(xAxisBox)
    wrl.addChild(transXaxis)

    trXticks = []
    for ii in range(9):
        ypos = -0.5 + ii*4.5
        trXticks.append(vrml.Transform(translation = tuple([17.0, ypos, 0.4])))
        xTickBox = vrml.Box(size = (37, 0.06, .1), color = (.2,.2,.2))
        trXticks[ii].addChild(xTickBox)
        wrl.addChild(trXticks[ii])

    
    return vrml.vrml(wrl)

# -----------------------------------------------------------------------------
# Put in axes as a separate model.
def setUpAxesLabels():
    from BuildStructure import _newModel, _newResidue
    from chimera.molEdit import addAtom
    from chimera import Element, Point


    ax_model = _newModel("Axes")
    res = _newResidue(ax_model, "Axes")

    xPosOffset = -2.0

    addAtom('Phi', Element('Phi'), res, Point(36.0, 18.0, 0.0))
    addAtom('Psi', Element('Psi'), res, Point(18.0, 36.0, 1.0))

    addAtom('  0', Element('  0'), res, Point(xPosOffset+18.0, -4.0, 0.6))
    addAtom('  0', Element('  0'), res, Point(-4.0, 18.0, 0.6))

    pAxVals = (' 45', ' 90', ' 135', ' 180')
    nAxVals = ('-45', '-90', '-135', '-180')
    for ii in range(4):
        addAtom(pAxVals[ii], Element(pAxVals[ii]), res, Point(xPosOffset+22.5 + 4.5*ii, -4.0, 0.6))
        addAtom(nAxVals[3-ii], Element(nAxVals[3-ii]), res, Point(xPosOffset-0.5 + 4.5*ii, -4.0, 0.6))
        addAtom(pAxVals[ii], Element(pAxVals[ii]), res, Point(-4.0, 22.5 + 4.5*ii,  0.6))
    for ii in range(3):
        addAtom(nAxVals[2-ii], Element(nAxVals[2-ii]), res, Point(-4.0, 4. + 4.5*ii,  0.6))

    for ats in res.atoms:
        ats.radius = 0.0001

    return

#====================================== Code above this line given to students.


testCaseCS483  =     [[2,3,3,3,3,3,3,3,3,3,3,2,2,2,1,1,0,0,0,0,0,1,1,1,1,1,0,0,0,0,0,1,1,1,1,2],
                      [2,3,3,3,3,3,3,3,3,3,3,3,2,2,2,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1,2,2],
                      [2,3,3,3,3,3,3,3,3,3,3,3,2,2,2,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,1,1,2,2,2],
                      [2,2,3,3,3,3,3,3,3,3,3,3,3,2,2,2,2,1,1,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1,2,2],
                      [2,2,3,3,3,3,3,3,3,3,3,3,3,3,2,2,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1,1,2],
                      [2,2,3,3,3,3,3,3,3,3,3,3,3,2,2,2,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1,2],
                      [2,2,2,3,3,3,3,3,3,3,3,3,2,2,2,2,1,1,1,0,0,0,1,1,1,1,1,1,0,0,0,0,0,1,1,2],
                      [2,2,2,2,3,3,3,3,3,3,3,2,2,2,2,2,2,1,1,1,1,1,1,1,1,1,1,1,0,0,0,0,0,1,1,2],
                      [2,2,2,2,2,3,2,2,3,2,2,2,2,2,2,2,2,1,1,1,1,1,1,1,2,2,1,1,0,0,0,0,0,1,1,2],
                      [2,2,2,2,2,2,2,2,2,2,2,2,2,2,1,1,1,1,1,1,1,2,2,2,2,1,1,1,0,0,0,0,0,1,1,1],
                      [2,2,2,2,2,2,2,2,2,2,2,2,1,1,1,1,1,1,1,1,1,2,2,2,2,1,1,1,0,0,0,0,0,1,1,1],
                      [1,2,2,2,2,2,2,2,2,2,2,2,2,1,1,0,0,0,0,1,1,1,2,2,2,1,1,1,1,0,0,0,0,1,1,1],
                      [1,1,2,2,2,2,2,2,2,2,2,2,2,1,1,0,0,0,0,1,1,2,2,2,2,2,2,1,1,0,0,0,0,1,1,1],
                      [1,1,2,2,2,2,2,2,2,2,2,2,1,1,1,0,0,0,0,1,1,2,2,3,2,2,1,1,1,0,0,0,0,1,1,1],
                      [1,2,2,2,2,2,2,2,2,2,2,2,1,1,1,0,0,0,0,1,1,1,2,3,3,2,1,1,1,0,0,0,0,1,1,1],
                      [1,2,2,2,2,2,2,3,2,2,2,2,2,1,1,1,0,0,0,1,1,1,2,2,2,2,2,1,1,0,0,0,0,1,1,1],
                      [1,2,2,2,2,2,3,3,3,3,2,2,2,1,1,1,1,0,0,0,1,1,1,2,2,2,1,1,1,0,0,0,0,1,1,1],
                      [1,2,2,2,2,2,3,3,3,3,3,2,2,2,1,1,1,0,0,0,1,1,1,1,2,2,2,1,1,0,0,0,0,1,1,1],
                      [1,2,2,2,2,3,3,3,3,3,3,3,2,2,2,1,1,1,0,0,1,1,1,2,1,2,2,1,1,0,0,0,0,1,1,1],
                      [1,2,2,2,2,2,3,3,3,3,3,3,3,2,2,1,1,1,1,0,1,1,1,1,1,1,1,1,1,0,0,0,0,1,1,1],
                      [1,1,2,2,2,2,3,3,3,3,3,3,3,3,2,2,1,1,1,0,1,1,1,1,1,1,1,1,1,0,0,0,0,1,1,1],
                      [1,1,2,2,2,2,2,3,3,3,3,3,3,3,2,2,2,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1],
                      [1,1,1,2,2,2,2,2,3,3,3,3,3,3,3,2,2,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1],
                      [1,2,2,2,2,2,2,2,2,3,3,3,3,3,3,2,2,2,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1],
                      [2,2,1,2,2,2,2,2,2,2,2,3,3,3,3,2,2,2,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1],
                      [1,1,1,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,1,1,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1],
                      [1,1,1,1,1,2,2,2,2,2,2,2,2,2,2,2,2,2,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1],
                      [0,1,1,1,1,2,2,2,2,2,2,2,2,2,2,1,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
                      [0,0,0,1,1,1,1,1,1,1,1,2,1,1,1,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
                      [0,0,0,1,1,1,1,2,2,1,1,1,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
                      [1,1,1,1,1,2,2,1,1,1,1,1,1,1,0,0,0,0,0,0,0,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0],
                      [1,1,1,1,1,1,2,2,2,2,1,1,1,1,0,0,0,0,0,0,0,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0],
                      [1,2,2,2,1,2,2,2,2,2,1,1,1,1,1,0,0,0,0,0,0,1,1,2,1,1,1,0,0,0,0,0,0,0,0,0],
                      [1,2,2,2,2,2,2,2,2,2,2,1,1,1,1,0,0,0,0,0,0,1,1,2,1,1,1,0,0,0,0,0,0,1,1,1],
                      [1,2,2,2,2,2,2,2,2,2,2,2,2,1,1,0,0,0,0,0,0,1,1,2,2,1,1,0,0,0,0,0,0,1,1,1],
                      [2,2,2,2,2,2,2,2,2,2,2,2,2,1,1,0,0,0,0,0,0,1,1,1,1,1,1,0,0,0,0,0,0,1,1,2]]

# To get rid of the protein in your display you can enable the following line.  Here prot was my model[0].
#prot.destroy()

vvrr = vr_plot(testCaseCS483)
vrlmModel = chimera.openModels.open(vvrr, 'VRML')[0]
setUpAxesLabels()
from chimera import runCommand
runCommand("set bg_color white")
runCommand("label")
runCommand("color black ,la")

chimera.viewer.viewAll()