class ColorbarWidget(QWidget):

    def __init__(self, parent=None):
        super(ColorbarWidget, self).__init__(parent)
        fig = Figure()
        rect = 0.25, 0.05, 0.1, 0.90
        self.cb_axes = fig.add_axes(rect)
        self.canvas = FigureCanvas(fig)
        self.setLayout(QVBoxLayout())
        self.layout().addWidget(self.canvas)
        self.button = QPushButton("Update")
        self.layout().addWidget(self.button)
        self.button.pressed.connect(self._update_cb_scale)

        self._create_colorbar(fig)

    def _create_colorbar(self, fig):
        self.mappable = ScalarMappable(norm=SymLogNorm(0.0001, 1,vmin=-10., vmax=10000.),
                                  cmap=DEFAULT_CMAP)
        self.mappable.set_array([])
        fig.colorbar(self.mappable, ax=self.cb_axes, cax=self.cb_axes)

    def _update_cb_scale(self):
        self.mappable.colorbar.remove()
        rect = 0.25, 0.05, 0.1, 0.90
        self.cb_axes = self.canvas.figure.add_axes(rect)
        self.mappable = ScalarMappable(Normalize(30, 4300),
                                   cmap=DEFAULT_CMAP)
        self.mappable.set_array([])
        self.canvas.figure.colorbar(self.mappable, ax=self.cb_axes, cax=self.cb_axes)
        self.canvas.draw()

def scatter3d(x,y,z, cs, colorsMap='jet'):
    cm = plt.get_cmap(colorsMap)
    cNorm = Normalize(vmin=min(cs), vmax=max(cs))
    scalarMap = ScalarMappable(norm=cNorm, cmap=cm)
    ax.scatter(x, y, z, c=scalarMap.to_rgba(cs), s=5, linewidth=0)
    scalarMap.set_array(cs)
    plt.show()

def create_heatmap(xs, ys, imageSize, blobSize, cmap):
    blob = Image.new('RGBA', (blobSize * 2, blobSize * 2), '#000000')
    blob.putalpha(0)
    colour = 255 / int(math.sqrt(len(xs)))
    draw = ImageDraw.Draw(blob)
    draw.ellipse((blobSize / 2, blobSize / 2, blobSize * 1.5, blobSize * 1.5),
                 fill=(colour, colour, colour))
    blob = blob.filter(ImageFilter.GaussianBlur(radius=blobSize / 2))
    heat = Image.new('RGBA', (imageSize, imageSize), '#000000')
    heat.putalpha(0)
    xScale = float(imageSize - 1) / (max(xs) - min(xs))
    yScale = float(imageSize - 1) / (min(ys) - max(ys))
    xOff = min(xs)
    yOff = max(ys)
    for i in range(len(xs)):
        xPos = int((xs[i] - xOff) * xScale)
        yPos = int((ys[i] - yOff) * yScale)
        blobLoc = Image.new('RGBA', (imageSize, imageSize), '#000000')
        blobLoc.putalpha(0)
        blobLoc.paste(blob, (xPos - blobSize, yPos - blobSize), blob)
        heat = ImageChops.add(heat, blobLoc)

    norm = Normalize(vmin=min(min(heat.getdata())),
                     vmax=max(max(heat.getdata())))
    sm = ScalarMappable(norm, cmap)
    heatArray = pil_to_array(heat)
    rgba = sm.to_rgba(heatArray[:, :, 0], bytes=True)
    rgba[:, :, 3] = heatArray[:, :, 3]
    coloured = Image.fromarray(rgba, 'RGBA')

    return coloured

 def get_color_map(self, levels):
     """Returns gradient of color from green to red.
     """
     sm = ScalarMappable(cmap='RdYlGn_r')
     normed_levels = levels / np.max(levels)
     colors = 255 * sm.to_rgba(normed_levels)[:, :3]
     return ['#%02x%02x%02x' % (r, g, b) for r,g,b in colors]

def make_graph(data):
	from matplotlib.cm import jet, ScalarMappable
	from matplotlib.colors import Normalize
	g = nx.Graph()
	cnorm = Normalize(vmin=1, vmax=241)
	smap = ScalarMappable(norm=cnorm, cmap=jet)
	edge_list = []
	for k in data:
		tk = k.split('_')
		if len(tk) != 5:
			g.add_node(k)
		else:
			a,b = tk[1],tk[3]
			g.add_node(a)
			g.add_node(b)
			g.add_edge(a,b)
	pos = nx.spring_layout(g)
	nxcols,glabels = [],{}
	for i,node in enumerate(g.nodes()):
		if '_' not in node:
			nxcols.append(smap.to_rgba(int(node)))
		else:
			nxcols.append((0,1,0,0))
	nx.draw_networkx_nodes(g,pos,node_color=nxcols)
	nx.draw_networkx_labels(g,pos)
	nx.draw_networkx_edges(g,pos)
	plt.show()
	return 

def digit_to_rgb(X, scaling=3, shape = (), cmap = 'binary'):
    '''
    Takes as input an intensity array and produces a rgb image due to some color map

    Parameters
    ----------

    X : numpy.ndarray
        intensity matrix as array of shape [M x N]

    scaling : int
        optional. positive integer value > 0

    shape: tuple or list of its , length = 2
        optional. if not given, X is reshaped to be square.

    cmap : str
        name of color map of choice. default is 'binary'

    Returns
    -------

    image : numpy.ndarray
        three-dimensional array of shape [scaling*H x scaling*W x 3] , where H*W == M*N
    '''

    sm = ScalarMappable(cmap = cmap)
    image = sm.to_rgba(enlarge_image(vec2im(X,shape), scaling))[:,:,0:3]
    return image

 def getcolor(self, V, F):
     dS = numpy.empty(len(F))
     for i, f in enumerate(F):
         v = V[f][0]
         dS[i] = v[0] * v[0] + v[1] * v[1] + v[2] * v[2]
     cmap = ScalarMappable(cmap='jet')
     cmap.set_array(dS)
     return cmap, cmap.to_rgba(dS)

def _check_cmap_rgb_vals(vals, cmap, vmin=0, vmax=1):
    """Helper function to check RGB values of color images"""
    from matplotlib.colors import Normalize
    from matplotlib.cm import ScalarMappable
    norm = Normalize(vmin, vmax)
    sm = ScalarMappable(norm=norm, cmap=cmap)
    for val, rgb_expected in vals:
        rgb_actual = sm.to_rgba(val)[:-1]
        assert_allclose(rgb_actual, rgb_expected, atol=1e-5)

def plot_events(mapobj,axisobj,catalog,label= None, color='depth', pretty = False, colormap=None, 
             llat = -90, ulat = 90, llon = -180, ulon = 180, figsize=(16,24), 
             par_range = (-90., 120., 30.), mer_range = (0., 360., 60.),
             showHour = False, M_above = 0.0, location = 'World', min_size=1, max_size=8,**kwargs):

    '''Simplified version of plot_event'''

    import matplotlib.pyplot as plt
    from matplotlib.colors import Normalize
    from matplotlib.cm import ScalarMappable

    lats, lons, mags, times, labels, colors = get_event_info(catalog, M_above, llat, ulat, llon, ulon, color, label)
                    
    min_color = min(colors)
    max_color = max(colors)

    if colormap is None:
        if color == "date":
             colormap = plt.get_cmap()
        else:
            # Choose green->yellow->red for the depth encoding.
             colormap = plt.get_cmap("RdYlGn_r")
                
    scal_map = ScalarMappable(norm=Normalize(min_color, max_color),
                                  cmap=colormap)
    scal_map.set_array(np.linspace(0, 1, 1))

    x, y = mapobj(lons, lats)

    min_mag = 0
    max_mag = 10
    if len(mags) > 1:
        frac = [(_i - min_mag) / (max_mag - min_mag) for _i in mags]
        magnitude_size = [(_i * (max_size - min_size)) ** 2 for _i in frac]
        #magnitude_size = [(_i * min_size) for _i in mags]
        #print magnitude_size
        colors_plot = [scal_map.to_rgba(c) for c in colors]
    else:
        magnitude_size = 15.0 ** 2
        colors_plot = "red"

    quakes = mapobj.scatter(x, y, marker='o', s=magnitude_size, c=colors_plot, zorder=10)
    #mapobj.drawmapboundary(fill_color='aqua')
    #mapobj.drawparallels(np.arange(-90,90,30),labels=[1,0,0,0])
    #mapobj.drawmeridians(np.arange(mapobj.lonmin,mapobj.lonmax+30,60),labels=[0,0,0,1])

    # if len(mags) > 1:
    #     cb = mpl.colorbar.ColorbarBase(ax=axisobj, cmap=colormap, orientation='vertical')
    #     cb.set_ticks([0, 0.25, 0.5, 0.75, 1.0])
    #     color_range = max_color - min_color
    #     cb.set_ticklabels([_i.strftime('%Y-%b-%d') if color == "date" else '%.1fkm' % (_i)
    #             for _i in [min_color, min_color + color_range * 0.25,
    #                    min_color + color_range * 0.50,
    #                    min_color + color_range * 0.75, max_color]])

    return quakes

 def _create_colorbar(self, cmap, ncolors, labels, **kwargs):    
     norm = BoundaryNorm(range(0, ncolors), cmap.N)
     mappable = ScalarMappable(cmap=cmap, norm=norm)
     mappable.set_array([])
     mappable.set_clim(-0.5, ncolors+0.5)
     colorbar = plt.colorbar(mappable, **kwargs)
     colorbar.set_ticks(np.linspace(0, ncolors, ncolors+1)+0.5)
     colorbar.set_ticklabels(range(0, ncolors))
     colorbar.set_ticklabels(labels)
     return colorbar

    def plotTciData(self, ax):
        rdata = self.elecTree.getNode('\electrons::top.tci.results:rad').data()
        sm = ScalarMappable()
        rhoColor = sm.to_rgba(-rdata)

        for i in range(1,11):
            pnode = self.elecTree.getNode('\electrons::top.tci.results:nl_%02d' % i)
            ax.plot(pnode.dim_of().data(), pnode.data(), c = rhoColor[i-1])

        ax.set_ylabel('ne')

def cm2colors(N, cmap='autumn'):
    """Takes N evenly spaced colors out of cmap and returns a
    list of rgb values"""
    values = range(N)
    cNorm  = Normalize(vmin=0, vmax=N-1)
    scalarMap = ScalarMappable(norm=cNorm, cmap=cmap)
    colors = []
    for i in xrange(N):
        colors.append(scalarMap.to_rgba(values[i]))
    return colors

def plot_net_layerwise(net, x_spacing=5, y_spacing=10, colors=[], use_labels=True, ax=None, cmap='gist_heat', cbar=False, positions={}):
	if not colors:
		colors = [1] * net.size()
	args = {
		'ax' : ax,
		'node_color' : colors,
		'nodelist' : net.nodes(), # ensure that same order is used throughout for parallel data like colors
		'vmin' : 0,
		'vmax' : 1,
		'cmap' : cmap
	}

	if not positions:
		# compute layer-wise positions of nodes (distance from roots)
		nodes_by_layer = defaultdict(lambda: [])
		def add_to_layer(n,l):
			nodes_by_layer[l].append(n)
		net.bfs_traverse(net.get_roots(), add_to_layer)


		positions = {}
		for l, nodes in nodes_by_layer.iteritems():
			y = -l*y_spacing
			# reorder layer lexicographically
			nodes.sort(key=lambda n: n.get_name())
			width = (len(nodes)-1) * x_spacing
			for i,n in enumerate(nodes):
				x = x_spacing*i - width/2
				positions[n] = (x,y)
	args['pos'] = positions

	if use_labels:
		labels = {n:n.get_name() for n in net.iter_nodes()}
		args['labels'] = labels

	if ax is None:
		ax = plt.figure().add_subplot(1,1,1)
	nxg = net_to_digraph(net)
	nx.draw_networkx(nxg, **args)
	ax.tick_params(axis='x', which='both', bottom='off', top='off', labelbottom='off')
	ax.tick_params(axis='y', which='both', left='off', right='off', labelleft='off')
	
	if cbar:
		color_map = ScalarMappable(cmap=cmap)
		color_map.set_clim(vmin=0, vmax=1)
		color_map.set_array(np.array([0,1]))
		plt.colorbar(color_map, ax=ax)

	ax.set_aspect('equal')
	# zoom out slightly to avoid cropping issues with nodes
	xl = ax.get_xlim()
	yl = ax.get_ylim()
	ax.set_xlim(xl[0]-x_spacing/2, xl[1]+x_spacing/2)
	ax.set_ylim(yl[0]-y_spacing/2, yl[1]+y_spacing/2)

  def _update_scatter(self):
    # Updates the scatter points for changes in *tidx* or *xidx*. This 
    # just changes the face color
    # 
    # CALL THIS AFTER *_update_image*
    #
    if len(self.data_sets) < 2:
      return

    sm = ScalarMappable(norm=self.cbar.norm,cmap=self.cbar.get_cmap())
    colors = sm.to_rgba(self.data_sets[1][self.tidx,:,2])
    self.scatter.set_facecolors(colors)

    def densityplot2(self,modelname='Model',refname='Ref',units = 'mmol m-3',sub = 'med'):
        '''
        opectool like density plot
        
        Ref is in x axis
        Model  in y axis
        
        Args
         - *modelname* (optional) string , default ='Model'
         - *refname*   (optional) string , default ='Ref'
         - *units*     (optional) string , default ='mmol m-3'
         - *sub*       (optional) string , default ='med'
        
        
        Returns: a matplotlib Figure object and a matplotlib Axes object
        '''
        
        fig, ax = plt.subplots()
        plt.title('%s Density plot of %s and %s\nNumber of considered matchups: %s' % (sub, modelname, refname, self.number()))
        cmap = 'spectral_r'
        axis_min = min(self.Ref.min(),self.Model.min())
        axis_max = max(self.Ref.max(),self.Model.max())
        extent = [axis_min, axis_max, axis_min, axis_max]

        hexbin = ax.hexbin(self.Ref, self.Model, bins=None, extent=extent, cmap=cmap)
        data = hexbin.get_array().astype(np.int32)
        MAX = data.max()

        for nticks in range(10,2,-1):
            float_array=np.linspace(0,MAX,nticks)
            int___array = float_array.astype(np.int32)
            if np.all(float_array == int___array ):
                break

        mappable = ScalarMappable(cmap=cmap)
        mappable.set_array(data)
        #fig.colorbar(mappable, ticks = int___array, ax=ax)
        cbar = fig.colorbar(mappable, ax=ax)
        labels = cbar.ax.get_yticklabels()
        FloatNumberFlag = False
        for label in labels:
            numstr = str(label.get_text())
            if numstr.find(".") > -1:
                FloatNumberFlag = True

        if FloatNumberFlag:
            cbar.remove()
            cbar = fig.colorbar(mappable, ticks = int___array, ax=ax)

        ax.set_xlabel('%s %s' % (refname,  units))
        ax.set_ylabel('%s %s' % (modelname,units))
        ax.grid()
        return fig,ax

def generalized_bar_chart(code_matrix, trans_names, code_names, show_it=True, show_trans_names=False, color_map = "jet", legend_labels = None, title=None, horizontal_grid = True):
    ldata = {}
    fig = pylab.figure(facecolor="white", figsize=(12, 4))
    fig.subplots_adjust(left=.05, bottom=.15, right=.98, top=.95)
    code_names = [c for c in range(code_matrix.shape[1])]
    for i, code in enumerate(range(len(code_names))):
        ldata[code] = [code_matrix[j, i] for j in range(len(trans_names))]
    ind = np.arange(len(trans_names))
    width = 1.0 / (len(code_names) + 1)
    traditional_colors = ['red', 'orange', 'yellow', 'green', 'blue', 'purple', 'black', 'grey', 'cyan', 'coral']

    ax = fig.add_subplot(111)
    if title is not None:
        ax.set_title(title, fontsize=10)

    if color_map == "AA_traditional_":
        lcolors = traditional_colors
    else:
        cNorm = mpl_Normalize(vmin = 1, vmax = len(code_names))
        comap = get_cmap(color_map)
        scalar_map = ScalarMappable(norm = cNorm, cmap = comap)
        lcolors = [scalar_map.to_rgba(idx + 1) for idx in range(len(code_names))]

    the_bars = []

    for c in range(len(code_names)):
        new_bars = ax.bar(ind + (c + .5) * width, ldata[code_names[c]], width, color=lcolors[c % (len(lcolors))])
        the_bars.append(new_bars[0])
        # bar_groups.append(bars)
    if show_trans_names:
        ax.set_xticks(ind + .5)
        ax.set_xticklabels(trans_names, size="x-small", rotation= -45)
    else:
        ax.grid(b = horizontal_grid, which = "major", axis = 'y')
        ax.set_xticks(ind + .5)
        ax.set_xticklabels(ind + 1, size="x-small")
        for i in ind[1:]:
            ax.axvline(x = i, linestyle = "--", linewidth = .25, color = 'black')

    if legend_labels != None:
        fontP =FontProperties()
        fontP.set_size('small')
        box = ax.get_position()
        ax.set_position([box.x0, box.y0, box.width * 0.825, box.height])
        # Put a legend to the right of the current axis
        ax.legend(the_bars, legend_labels, loc='center left', bbox_to_anchor=(1, 0.5), prop = fontP)
    ax.set_xlim(right=len(trans_names))

    if show_it:
        fig.show()
    return fig

def plot(countries,values,label='',clim=None,verbose=False):
    """
    Usage: worldmap.plot(countries, values [, label] [, clim])
    """
    countries_shp = shpreader.natural_earth(resolution='110m',category='cultural',
                                            name='admin_0_countries')
    ## Create a plot
    fig = plt.figure()
    ax = plt.axes(projection=ccrs.PlateCarree())
    ## Create a colormap
    cmap = plt.get_cmap('RdYlGn_r')
    if clim:
       vmin = clim[0]
       vmax = clim[1]
    else:
       val = values[np.isfinite(values)]
       mean = val.mean()
       std = val.std()
       vmin = mean-2*std
       vmax = mean+2*std
    norm = Normalize(vmin=vmin,vmax=vmax)
    smap = ScalarMappable(norm=norm,cmap=cmap)
    ax2 = fig.add_axes([0.3, 0.18, 0.4, 0.03])
    cbar = ColorbarBase(ax2,cmap=cmap,norm=norm,orientation='horizontal')
    cbar.set_label(label)
    ## Add countries to the map
    for country in shpreader.Reader(countries_shp).records():
        countrycode = country.attributes['adm0_a3']
        countryname = country.attributes['name_long']
        ## Check for country code consistency
        if countrycode == 'SDS': #South Sudan
           countrycode = 'SSD'
        elif countrycode == 'ROU': #Romania
           countrycode = 'ROM'
        elif countrycode == 'COD': #Dem. Rep. Congo
           countrycode = 'ZAR'
        elif countrycode == 'KOS': #Kosovo
           countrycode = 'KSV'
        if countrycode in countries:
           val = values[countries==countrycode]
           if np.isfinite(val):
              color = smap.to_rgba(val)
           else:
              color = 'grey'
        else:
           color = 'w'
           if verbose:
              print("No data available for "+countrycode+": "+countryname)
        ax.add_geometries(country.geometry,ccrs.PlateCarree(),facecolor=color,label=countryname)
    plt.show()

def custom_colorbar(cmap, ncolors, breaks, **kwargs):
    from matplotlib.colors import BoundaryNorm
    from matplotlib.cm import ScalarMappable
    import matplotlib.colors as mplc

    breaklabels = ['No Counts']+["> %d counts"%(perc) for perc in breaks[:-1]]

    norm = BoundaryNorm(range(0, ncolors), cmap.N)
    mappable = ScalarMappable(cmap=cmap, norm=norm)
    mappable.set_array([])
    mappable.set_clim(-0.5, ncolors+0.5)
    colorbar = plt.colorbar(mappable, **kwargs)
    colorbar.set_ticks(np.linspace(0, ncolors, ncolors+1)+0.5)
    colorbar.set_ticklabels(range(0, ncolors))
    colorbar.set_ticklabels(breaklabels)
    return colorbar

  def _init_scatter(self):
    # Plots the scatter points at the base of each vector showing the 
    # vertical deformation for the second data set. If there is only 
    # one data set then this function does nothing.
    # 
    # CALL THIS AFTER *_init_image*
    #
    if len(self.data_sets) < 2:
      self.scatter = None 
      return

    sm = ScalarMappable(norm=self.cbar.norm,cmap=self.cbar.get_cmap())
    # use scatter points to show z for second data set 
    colors = sm.to_rgba(self.data_sets[1][self.tidx,:,2])
    self.scatter = self.map_ax.scatter(self.x[:,0],self.x[:,1],
                                       c=colors,s=self.scatter_size,
                                       zorder=2,edgecolor=self.colors[1])

def generate_colors(desired_palette, num_desired_colors):
  """
  Generate an array of color strings, interpolated from the desired_palette.
  desired_palette is from palettable

  Conceptually, this takes a list of colors, and lets you generate any length
  of colors from that array.
  """
  cmap = desired_palette.mpl_colormap
  mappable = ScalarMappable(norm=Normalize(vmin=0, vmax=1), cmap=cmap)

  cols = []
  for i in range(1,num_desired_colors+1):
    (r,g,b,a) = mappable.to_rgba((i - 1) / num_desired_colors)
    cols.append(rgb_to_hex(map(lambda x: int(x*255), [r,g,b])))

  return cols