from java.awt import *
from visad import Display,RealTupleType,ScalarMap,RealType
from visad import SI,FunctionType,FlatField,Linear2DSet,Irregular3DSet,DataReferenceImpl
from jhplot import *
from java.util import *
from java.lang import *

#The domain quantities longitude and latitude
# and the dependent quantity temperature

#  get 2D display
c1=HVisAd("2D example",False)
display = c1.getDisplay()
ren=c1.getRender() 
ren.setBackgroundColor(Color.white) 
display.getGraphicsModeControl().setLineWidth(1.0)

#Create the quantities
# Use RealType(String name)
latitude = RealType("latitude")
longitude = RealType("longitude")

domain_tuple = RealTupleType(latitude, longitude)
temperature = RealType("temperature", SI.kelvin, None)
isoTemperature = RealType("isoTemperature", SI.kelvin, None)

# Create a FunctionType (domain_tuple -> temperature )
# Use FunctionType(MathType domain, MathType range)
func_domain_temp = FunctionType( domain_tuple, temperature)

# ... the same for isoTemperature
func_domain_isoTemp =FunctionType( domain_tuple, isoTemperature)

# Create the domain Set
# Use LinearDSet(MathType type, double first1, double last1, int lengthX,
#		double first2, double last2, int lengthY)
NCOLS = 50
NROWS = NCOLS
domain_set = Linear2DSet(domain_tuple, -Math.PI, Math.PI, NROWS, -Math.PI, Math.PI, NCOLS) 

# Get the Set samples to facilitate the calculations
set_samples = domain_set.getSamples( True )

# The actual temperature values are stored in this array
# float[1][ number_of_samples ]
flat_samples = []

# We fill our 'flat' array with the generated values
# by looping over NCOLS and NROWS
# Note the use of an index variable, indicating the order of the samples
sss=[]
for c in range(NCOLS):
   for r in range(NROWS):
     f=((Math.sin( 0.50*set_samples[0][ c * NROWS + r ])  ) * Math.cos( set_samples[1][ c * NROWS + r ] ) )  
     # f=c*r+ set_samples[1][ c * NROWS + r ]
     sss.append(f)
     # set_samples[0][c * NROWS + r]=f

flat_samples.append([sss][0])


# Create the FlatFields
# Use FlatField(FunctionType type, Set domain_set)
# For the colored image
vals_ff = FlatField( func_domain_temp, domain_set)

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


# ...and for the isocontours
iso_vals_ff =FlatField( func_domain_isoTemp, domain_set)


# Get the values from the temperature FlatField
# create flat_isoVals array for clarity's sake
# "false" argument means "don't copy"
flat_isoVals = vals_ff.getFloats(False)

# ...and put the values above into it
#/ Note the argument false, meaning that the array won't be copied again
iso_vals_ff.setSamples( flat_isoVals, False )


#  Create Display and its maps
#  Get display's graphics mode control and draw scales
dispGMC = display.getGraphicsModeControl()
dispGMC.setScaleEnable(True)



# Create the ScalarMaps: latitude to YAxis, longitude to XAxis and
# temperature to RGB and
# isoTemperature to IsoContour
# Use ScalarMap(ScalarType scalar, DisplayRealType display_scalar)
latMap = ScalarMap( latitude, Display.YAxis )
lonMap = ScalarMap( longitude, Display.XAxis )
latMap.getAxisScale().setColor(Color.black)
lonMap.getAxisScale().setColor(Color.black)
latMap.getAxisScale().setLabelSize(18)
lonMap.getAxisScale().setLabelSize(18)
# latMap.getAxisScale().setSnapToBox(True)
# lonMap.getAxisScale().setSnapToBox(True)

tempIsoMap = ScalarMap( isoTemperature,  Display.IsoContour )
tempRGBMap = ScalarMap( temperature,  Display.RGB )


# Add maps to display
display.addMap( latMap )
display.addMap( lonMap )

display.addMap( tempIsoMap )
display.addMap( tempRGBMap )

# The ContourControl
# Note that we get the control from the IsoContour map
isoControl = tempIsoMap.getControl()

# Define some parameters for contour lines
interval = 0.125  # interval between lines
lowValue = -0.5   # lowest value
highValue = 1     # highest value
base = -1.       #  starting at this base value

# ...and set the lines with the method
isoControl.setContourInterval(interval, lowValue, highValue, base)
isoControl.enableLabels(True)

# Create data references and set the FlatField as our data
data_ref = DataReferenceImpl("data_ref")
iso_data_ref = DataReferenceImpl("iso_data_ref")

data_ref.setData( vals_ff )
iso_data_ref.setData( iso_vals_ff )

# Add reference to display
display.addReference( data_ref )
display.addReference( iso_data_ref )

# show it
c1.visible()