from cmaptools import truncate_colormap
from plot_quiver import plot_quiver
try:
	import pylab as p
	import matplotlib as mpl
	import matplotlib.pyplot as plt
	import numpy as  np
except ImportError:
	print "could not import pylab, matplotlib has not been installed, no plotting capabilities enabled"
	
def plot_unit(x,y,z,elements,data,is2d,isplanet,datatype,options,ax):
	"""
	PLOT_UNIT - unit plot, display data
	
	Usage:
	plot_unit(x,y,z,elements,data,is2d,isplanet,datatype,options)

	See also: PLOTMODEL, PLOT_MANAGER
	"""
	
	#edgecolor
	edgecolor=options.getfieldvalue('edgecolor','None')

	# colormap
	# {{{ give number of colorlevels and transparency
	colorlevels=options.getfieldvalue('colorlevels',128)
	alpha=options.getfieldvalue('alpha',1)
	# }}}
	# {{{ define wich colormap to use 
	try:
		defaultmap=plt.cm.get_cmap('viridis')
	except AttributeError:
		print("Viridis can't be found (probably too old Matplotlib) reverting to gnuplot colormap")
		defaultmap=truncate_colormap('gnuplot2',0.1,0.9,colorlevels)
	cmap=options.getfieldvalue('colormap',defaultmap)
	if options.exist('cmap_set_over'):
		over=options.getfieldvalue('cmap_set_over','0.5')
		cmap.set_over(over)
	if options.exist('cmap_set_under'):
		under=options.getfieldvalue('cmap_set_under','0.5')
		cmap.set_under(under)
	options.addfield('colormap',cmap)
	# }}}	
	# {{{ Get the colormap limits
	if options.exist('clim'):
		lims=options.getfieldvalue('clim',[np.amin(data),np.amax(data)])
	elif options.exist('caxis'):
		lims=options.getfieldvalue('caxis',[np.amin(data),np.amax(data)])
	else:
		if np.amin(data)==np.amax(data):
			lims=[np.amin(data)-0.5,np.amax(data)+0.5]
		else:
			lims=[np.amin(data),np.amax(data)]
	# }}}
	# {{{ Set the spread of the colormap (default is normal
	if options.exist('log'):
		norm = mpl.colors.LogNorm(vmin=lims[0], vmax=lims[1])
	else:
		norm = mpl.colors.Normalize(vmin=lims[0], vmax=lims[1])
	if options.exist('log'):
		norm = mpl.colors.LogNorm(vmin=lims[0], vmax=lims[1])
	else:
		norm = mpl.colors.Normalize(vmin=lims[0], vmax=lims[1])
	options.addfield('colornorm',norm)
	# }}}
	
	# Plot depending on the datatype
	# {{{ data are on elements
	if datatype==1:
		if is2d:
			if options.exist('mask'):
				triangles=mpl.tri.Triangulation(x,y,elements,data.mask)
			else:
				triangles=mpl.tri.Triangulation(x,y,elements)
			tri=ax.tripcolor(triangles,data,colorlevels,cmap=cmap,norm=norm,alpha=alpha,edgecolors=edgecolor)
		else:
			raise ValueError('plot_unit error: 3D element plot not supported yet')
		return 
	# }}}
	# {{{ data are on nodes
	elif datatype==2:
		if is2d:
			if options.exist('mask'):
				EltMask=np.asarray([np.any(np.in1d(index,np.where(data.mask))) for index in elements])
				triangles=mpl.tri.Triangulation(x,y,elements,EltMask)
			else:
				triangles=mpl.tri.Triangulation(x,y,elements)
			ax.tricontourf(triangles,data,colorlevels,cmap=cmap,norm=norm,alpha=alpha)
			if edgecolor != 'None':
				ax.triplot(x,y,elements,color=edgecolor)
		else:
			raise ValueError('plot_unit error: 3D node plot not supported yet')
		return
	# }}}
	# {{{ plotting quiver (TODO)
	elif datatype==3:
		if is2d:
			Q=plot_quiver(x,y,data,options,ax)
		else:
			raise ValueError('plot_unit error: 3D node plot not supported yet')
		return
	
	# }}}
	# {{{ plotting P1 Patch (TODO)
	elif datatype==4:
		print 'plot_unit message: P1 patch plot not implemented yet'
		return
	# }}}
	# {{{ plotting P0 Patch (TODO)
	elif datatype==5:
		print 'plot_unit message: P0 patch plot not implemented yet'
		return
	# }}}
	else:
		raise ValueError('datatype=%d not supported' % datatype)
    
