#import matplotlib
from matplotlib.patches import Circle, Ellipse
from pylab import gca, gcf, gci, axis, histogram2d, hist
from numpy import array, append, flatnonzero

import numpy as np
import pylab as plt
from matplotlib.ticker import FixedFormatter

class PlotSequence(object):
	def __init__(self, basefn, format='%02i', suffix='png',
				 suffixes=None):
		self.ploti = 0
		self.basefn = basefn
		self.format = format
		if suffixes is None:
			self.suffixes = [suffix]
		else:
			self.suffixes = suffixes
	def skip(self, n=1):
		self.ploti += n
	def skipto(self, n):
		self.ploti = n

	def getnextlist(self):
		lst = ['%s-%s.%s' % (self.basefn, self.format % self.ploti, suff)
				for suff in self.suffixes]
		self.ploti += 1
		return lst

	def getnext(self):
		lst = self.getnextlist()
		if len(lst) == 1:
			return lst[0]
		return lst

	def savefig(self, **kwargs):
		import pylab as plt
		for fn in self.getnextlist():
			plt.savefig(fn, **kwargs)
			print 'saved', fn

def loghist(x, y, nbins=100,
			hot=True, doclf=True, docolorbar=True, lo=0.3,
			imshowargs={},
			clampxlo=False, clampxlo_val=None, clampxlo_to=None,
			clampxhi=False, clampxhi_val=None, clampxhi_to=None,
			clampylo=False, clampylo_val=None, clampylo_to=None,
			clampyhi=False, clampyhi_val=None, clampyhi_to=None,
			clamp=None, clamp_to=None,
			**kwargs):
	#np.seterr(all='warn')
	if doclf:
		plt.clf()
	myargs = kwargs.copy()
	if not 'bins' in myargs:
		myargs['bins'] = nbins

	rng = kwargs.get('range', None)
	x = np.array(x)
	y = np.array(y)

	if not (np.all(np.isfinite(x)) and np.all(np.isfinite(y))):
		K = np.flatnonzero(np.isfinite(x) * np.isfinite(y))
		print 'loghist: cutting to', len(K), 'of', len(x), 'finite values'
		x = x[K]
		y = y[K]
		
	if clamp is True:
		clamp = rng
	if clamp is not None:
		((clampxlo_val, clampxhi_val),(clampylo_val, clampyhi_val)) = clamp
	if clamp_to is not None:
		((clampxlo_to, clampxhi_to),(clampylo_to, clampyhi_to)) = clamp_to
	if clampxlo:
		if clampxlo_val is None:
			if rng is None:
				raise RuntimeError('clampxlo, but no clampxlo_val or range')
			clampxlo_val = rng[0][0]
	if clampxlo_val is not None:
		if clampxlo_to is None:
			clampxlo_to = clampxlo_val
		x[x < clampxlo_val] = clampxlo_to
	if clampxhi:
		if clampxhi_val is None:
			if rng is None:
				raise RuntimeError('clampxhi, but no clampxhi_val or range')
			clampxhi_val = rng[0][1]
	if clampxhi_val is not None:
		if clampxhi_to is None:
			clampxhi_to = clampxhi_val
		x[x > clampxhi_val] = clampxhi_to
	if clampylo:
		if clampylo_val is None:
			if rng is None:
				raise RuntimeError('clampylo, but no clampylo_val or range')
			clampylo_val = rng[1][0]
	if clampylo_val is not None:
		if clampylo_to is None:
			clampylo_to = clampylo_val
		y[y < clampylo_val] = clampylo_to
	if clampyhi:
		if clampyhi_val is None:
			if rng is None:
				raise RuntimeError('clampyhi, but no clampyhi_val or range')
			clampyhi_val = rng[1][1]
	if clampyhi_val is not None:
		if clampyhi_to is None:
			clampyhi_to = clampyhi_val
		y[y > clampyhi_val] = clampyhi_to




	(H,xe,ye) = np.histogram2d(x, y, **myargs)

	L = np.log10(np.maximum(lo, H.T))
	myargs = dict(extent=(min(xe), max(xe), min(ye), max(ye)),
				  aspect='auto',
				  interpolation='nearest', origin='lower')
	myargs.update(imshowargs)
	plt.imshow(L, **myargs)
	if hot:
		plt.hot()
	if docolorbar:
		r = [np.log10(lo)] + range(int(np.ceil(L.max())))
		# print 'loghist: L max', L.max(), 'r', r
		plt.colorbar(ticks=r, format=FixedFormatter(
			['0'] + ['%i'%(10**ri) for ri in r[1:]]))
	#set_fp_err()
	return H, xe, ye

def plothist(x, y, nbins=100, log=False,
			 doclf=True, docolorbar=True, dohot=True,
			 plo=None, phi=None,
			 imshowargs={}, **hist2dargs):
	if log:
		return loghist(x, y, nbins=nbins, doclf=doclf, docolorbar=docolorbar,
					   dohot=dohot, imshowargs=imshowargs) #, **kwargs)
					   
	if doclf:
		plt.clf()
	(H,xe,ye) = np.histogram2d(x, y, nbins, **hist2dargs)
	myargs = dict(extent=(min(xe), max(xe), min(ye), max(ye)),
				  aspect='auto',
				  interpolation='nearest', origin='lower')
	vmin = None
	if plo is not None:
		vmin = np.percentile(H.ravel(), plo)
		myargs.update(vmin=vmin)
	if phi is not None:
		vmin = imshowargs.get('vmin', vmin)
		vmax = np.percentile(H.ravel(), phi)
		if vmax != vmin:
			myargs.update(vmax=vmax)
	myargs.update(imshowargs)
	plt.imshow(H.T, **myargs)
	if dohot:
		plt.hot()
	if docolorbar:
		plt.colorbar()
	return H, xe, ye

def setRadecAxes(ramin, ramax, decmin, decmax):
	rl,rh = ramin,ramax
	dl,dh = decmin,decmax
	rascale = np.cos(np.deg2rad((dl+dh)/2.))
	ax = [rh,rl, dl,dh]
	plt.axis(ax)
	plt.gca().set_aspect(1./rascale, adjustable='box', anchor='C')
	plt.xlabel('RA (deg)')
	plt.ylabel('Dec (deg)')
	return ax

import matplotlib.colors as mc
class ArcsinhNormalize(mc.Normalize):
 	def __init__(self, mean=None, std=None, **kwargs):
 		self.mean = mean
 		self.std = std
		mc.Normalize.__init__(self, **kwargs)

	def _map(self, X, out=None):
		Y = (X - self.mean) / self.std
		args = (Y,)
		if out is not None:
			args = args + (out,)
		return np.arcsinh(*args)

	def __call__(self, value, clip=None):
		# copied from Normalize since it's not easy to subclass
		if clip is None:
			clip = self.clip
		result, is_scalar = self.process_value(value)
		self.autoscale_None(result)
		vmin, vmax = self.vmin, self.vmax
		if vmin > vmax:
			raise ValueError("minvalue must be less than or equal to maxvalue")
		elif vmin == vmax:
			result.fill(0)	 # Or should it be all masked?	Or 0.5?
		else:
			vmin = float(vmin)
			vmax = float(vmax)
			if clip:
				mask = ma.getmask(result)
				result = ma.array(np.clip(result.filled(vmax), vmin, vmax), mask=mask)
			# ma division is very slow; we can take a shortcut
			resdat = result.data
			self._map(resdat, resdat)
			vmin = self._map(vmin)
			vmax = self._map(vmax)
			resdat -= vmin
			resdat /= (vmax - vmin)
			result = np.ma.array(resdat, mask=result.mask, copy=False)
		if is_scalar:
			result = result[0]
		return result


from matplotlib.colors import LinearSegmentedColormap

# a colormap that goes from white to black: the opposite of matplotlib.gray()
antigray = LinearSegmentedColormap('antigray',
								   {'red':   ((0., 1, 1), (1., 0, 0)),
									'green': ((0., 1, 1), (1., 0, 0)),
									'blue':  ((0., 1, 1), (1., 0, 0))})

bluegrayred = LinearSegmentedColormap('bluegrayred',
									  {'red':   ((0., -1, 0),
												 (1., 1, -1)),
									   'green': ((0., -1,   0),
												 (0.5,0.5, 0.5),
												 (1., 0, -1)),
									   'blue':  ((0., -1, 1),
												 (1., 0, -1))})

# x, y0, y1
_redgreen_data =  {'red':   ((0.,  -100,  1),
							 #(0.5,  0,  0),
							 #(0.5,  0.1, 0),
							 (0.49, 0.1, 0),
							 (0.491, 0, 0),
							 (0.51,  0, 0),
							 (0.511,  0, 0.1),
							 (1.,   0, -100)),
				   'green': ((0.,  -100,  0),
							 #(0.5,  0,  0),
							 #(0.5,  0,  0.1),
							 (0.49, 0.1, 0),
							 (0.491, 0, 0),
							 (0.51,  0, 0),
							 (0.511,  0, 0.1),
							 (1.,   1, -100)),
				   'blue':  ((0.,  -100,  0),
							 (1.,   0, -100))}
redgreen = LinearSegmentedColormap('redgreen',   _redgreen_data)

def hist_ints(x, step=1, **kwargs):
	'''
	Creates a histogram of integers.  The number of bins is set to the
	range of the data (+1).  That is, each integer gets its own bin.
	'''
	kwargs['bins'] = x.max()/step - x.min()/step + 1
	kwargs['range'] = ( (x.min()/int(step))*step - 0.5,
						((x.max()/int(step))*step + 0.5) )
	return hist(x, **kwargs)

def hist2d_with_outliers(x, y, xbins, ybins, nout):
	'''
	Creates a 2D histogram from the given data, and returns a list of
	the indices in the data of points that lie in low-occupancy cells
	(where the histogram counts is < "nout").

	The "xbins" and "ybins" arguments are passed to numpy.histogram2d.

	You probably want to show the histogram with:

	  (H, outliers, xe, ye) = hist2d_with_outliers(x, y, 10, 10, 10)
	  imshow(H, extent=(min(xe), max(xe), min(ye), max(ye)), aspect='auto')
	  plot(x[outliers], y[outliers], 'r.')

	Returns: (H, outliers, xe, ye)

	  H: 2D histogram image
	  outliers: array of integer indices of the outliers
	  xe: x edges chosen by histgram2d
	  ye: y edges chosen by histgram2d
	  
	'''
	# returns (density image, indices of outliers)
	(H,xe,ye) = histogram2d(x, y, (xbins,ybins))
	Out = array([]).astype(int)
	for i in range(len(xe)-1):
		for j in range(len(ye)-1):
			if H[i,j] > nout:
				continue
			if H[i,j] == 0:
				continue
			H[i,j] = 0
			Out = append(Out, flatnonzero((x >= xe[i]) *
										  (x <  xe[i+1]) *
										  (y >= ye[j]) *
										  (y <  ye[j+1])))
	return (H.T, Out, xe, ye)


# You probably want to set the keyword radius=R
def circle(xy=None, x=None, y=None, **kwargs):
	if xy is None:
		if x is None or y is None:
			raise 'circle: need x and y'
		xy = array([x,y])
	c = Circle(xy=xy, **kwargs)
	a=gca()
	c.set_clip_box(a.bbox)
	a.add_artist(c)
	return c
	
def ellipse(xy=None, x=None, y=None, **kwargs):
	if xy is None:
		if x is None or y is None:
			raise 'ellipse: need x and y'
		xy = array([x,y])
	c = Ellipse(xy=xy, **kwargs)
	a=gca()
	c.set_clip_box(a.bbox)
	a.add_artist(c)
	return c

# return (pixel width, pixel height) of the axes area.
def get_axes_pixel_size():
	dpi = gcf().get_dpi()
	figsize = gcf().get_size_inches()
	axpos = gca().get_position()
	pixw = figsize[0] * dpi * axpos.width
	pixh = figsize[1] * dpi * axpos.height
	return (pixw, pixh)

    # test:
	if False:
		figure(dpi=100)
		(w,h) = get_axes_pixel_size()
		# not clear why this is required...
		w += 1
		h += 1
		img = zeros((h,w))
		img[:,::2] = 1.
		img[::2,:] = 1.
		imshow(img, extent=(0,w,0,h), aspect='auto', cmap=antigray)
		xlim(0,w)
		ylim(0,h)
		savefig('imtest.png')
		sys.exit(0)

# returns (x data units per pixel, y data units per pixel)
# given the current plot range, figure size, and axes position.
def get_pixel_scales():
	a = axis()
	(pixw, pixh) = get_axes_pixel_size()
	return ((a[1]-a[0])/float(pixw), (a[3]-a[2])/float(pixh))

def set_image_color_percentiles(image, plo, phi):
	# hackery...
	I = image.copy().ravel()
	I.sort()
	N = len(I)
	mn = I[max(0, int(round(plo * N / 100.)))]
	mx = I[min(N-1, int(round(phi * N / 100.)))]
	gci().set_clim(mn, mx)
	return (mn,mx)

