from java.awt import *
from visad import FlatField,Display,Linear3DSet,RealTupleType,FunctionType,ScalarMap,RealType
from visad import Integer3DSet,DataReferenceImpl
from java.lang import Math
from jhplot import *


c1=HVisAd("3D display")
display=c1.getDisplay()
ren=c1.getRender()
ren.setBackgroundColor(Color.white)


red = RealType("RED", None, None)
green = RealType("GREEN", None, None)
blue = RealType("BLUE", None, None)
domain_tuple = RealTupleType(red, green, blue)

# The independent variable
rgbVal = RealType.getRealType("RGB_VALUE", None, None)
# Create a FunctionType (domain_tuple -> range_tuple )
# Use FunctionType(MathType domain, MathType range)
func_domain_rgbVal = FunctionType( domain_tuple, rgbVal)

# Create the domain Set
# Integer3DSet(MathType type, int lengthX, int lengthY, int lengthZ)
NCOLS,NROWS,NLEVS = 8,8,8
domain_set = Integer3DSet(domain_tuple, NROWS, NCOLS, NLEVS )

flat_samples=[]
# Our 'flat' array
# double[][] flat_samples = new double[1][NCOLS * NROWS * NLEVS]
# Fill our 'flat' array with the rgbVal values
# by looping over NCOLS and NROWS
# Note the use of an index variable, indicating the order of the samples

index = 0
sss=[]
for  l in range(NLEVS):
    for c in range(NCOLS):
        for r in range(NROWS):
	      sss.append(index) 
	      index=index+1

flat_samples=[]
flat_samples.append([sss][0])

# Create a FlatField
# Use FlatField(FunctionType type, Set domain_set)
vals_ff = FlatField( func_domain_rgbVal, domain_set)

#  ...and put the rgbVal values above into it
#  Note the argument false, meaning that the array won't be copied
vals_ff.setSamples( flat_samples, False )


# Get display's graphics mode control and draw scales
dispGMC = display.getGraphicsModeControl()
dispGMC.setScaleEnable(True)
dispGMC.setPointSize(4.0)

redXMap = ScalarMap(red,   Display.XAxis )
greenYMap = ScalarMap(green, Display.YAxis )
blueZMap = ScalarMap( blue, Display.ZAxis )

redXMap.setScalarName("The RED Component")
greenYMap.setScalarName("The GREEN Component")
blueZMap.setScalarName("The BLUE Component")

redMap = ScalarMap( red,  Display.Red )
greenMap = ScalarMap( green,  Display.Green )
blueMap = ScalarMap( blue,  Display.Blue )

# Add maps to display
display.addMap( redXMap )
display.addMap( greenYMap )
display.addMap( blueZMap )

display.addMap( redMap )
display.addMap( greenMap )
display.addMap( blueMap)

def colorToFloats(c):
    rgb = [0.5,0.5,0.5]
    if(c != None):
      rgb=[]
      rgb.append(c.getRed()/255.0)
      rgb.append(c.getGreen()/255.0)
      rgb.append(c.getBlue()/255.0)
    return rgb

# Set axes colors
r = colorToFloats(Color.red)
redXMap.setScaleColor( r )

g = colorToFloats(Color.green)
greenYMap.setScaleColor( g )

b = colorToFloats(Color.blue)
blueZMap.setScaleColor( b )

# Create a data reference and set the FlatField as our data
data_ref = DataReferenceImpl("data_ref")
data_ref.setData( vals_ff )

# Add reference to display
display.addReference( data_ref )

c1.visible()