import numpy
from numpy import *
from CGLutil import vrml

# -----------------------------------------------------------------------------
# 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)

    transVertAxis = vrml.Transform(translation = tuple([17.5, 17.5, 0.5]))
    vertAxisBox = vrml.Box(size = (0.1, 36, .1), color = (0,0,0))
    transVertAxis.addChild(vertAxisBox)
    wrl.addChild(transVertAxis)
        
    transHorzAxis = vrml.Transform(translation = tuple([17.5, 17.5, 0.5]))
    horzAxisBox = vrml.Box(size = (36, 0.1, .1), color = (0,0,0))
    transHorzAxis.addChild(horzAxisBox)
    wrl.addChild(transHorzAxis)
    
    return vrml.vrml(wrl)
#====================================== 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]]


    
vvrr = vr_plot(testCaseCS483)
vrlmModel = chimera.openModels.open(vvrr, 'VRML')[0]