# -*- coding: utf-8 -*-
# 
#   $Id: Plotting.py 5752 2018-05-09 16:38:35Z jhill $
#   
#   This file is part of the BCPy2000 framework, a Python framework for
#   implementing modules that run on top of the BCI2000 <http://bci2000.org/>
#   platform, for the purpose of realtime biosignal processing.
# 
#   Copyright (C) 2007-11  Jeremy Hill, Thomas Schreiner,
#                          Christian Puzicha, Jason Farquhar
#   
#   bcpy2000@bci2000.org
#   
#   The BCPy2000 framework is free software: you can redistribute it
#   and/or modify it under the terms of the GNU General Public License
#   as published by the Free Software Foundation, either version 3 of
#   the License, or (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
__all__ = [
	'plot', 'plotsig',
	'imagesc', 'scatterplot', 'scatcmp', 'stem',
	'colorbar', 'rmcolorbar',
	'make_cmap', 'complement_cmap', 'reverse_cmap', 'show_cmaps',
	'subplots',
	'curve',
	'cax',
	'superimpose', 'look', 'hotfig', 
]

import numpy
from .Basic import getfs,fft2ap,fft
from .NumTools import isnumpyarray, project
from functools import reduce

def filterdict(d, exact=(), startswith=None, endswith=None):
	f = dict([(k,v) for k,v in list(d.items()) if k in exact or (startswith != None and k.startswith(startswith)) or (endswith != None and k.endswith(endswith))])
	for k in f: d.pop(k)
	return f

def plot(*pargs,**kwargs):
	"""
	A wrapper around pylab.plot that reduces dimensionality-related fiddliness.
	
	plot(x, y)
	plot(y)
	
	where either x or y, independent of each other, can specify multiple lines.
	
	Additional options and their defaults:
		axis = 0         Along which dimension can one step while staying on
		                 the same graphical line?
		stem = False     If True, do a stem plot instead of an ordinary line
		grid = True      Whether to turn on the grid.
		balance = None   If None, do nothing. If a numeric value, balance the
		                 y axis limits around this value.
		aspect = None    If 'auto', 'image' or 'equal', apply that aspect mode
		hold = False     If False, clear the axes before plotting.
		drawnow = True   If True, execute pylab.draw() after plotting.
	"""###
	hold = kwargs.pop('hold', False)
	axis = kwargs.pop('axis', 0)
	drawnow = kwargs.pop('drawnow', True)
	grid = kwargs.pop('grid', True)
	balance = kwargs.pop('balance', None)
	stem = kwargs.pop('stem', False)
	aspect = kwargs.pop('aspect', None)
	axkwargs = filterdict(kwargs, ['title', 'xlabel', 'ylabel', 'ylim', 'xlim'])

	pargs = list(pargs) # makes a copy, at least of the list container
	allvec = True
	for i in range(len(pargs)):
		if isinstance(pargs[i], (tuple,list)):
			pargs[i] = numpy.array(pargs[i],dtype=numpy.float64)
		if isinstance(pargs[i], numpy.ndarray):
			allvec &= (max(pargs[i].shape) == numpy.prod(pargs[i].shape))
			if len(pargs[i].shape) > 1:
				pargs[i] = project(pargs[i],axis).swapaxes(0,axis)
			isinf = numpy.isinf(pargs[i])
			if numpy.any(isinf):
				pargs[i] = pargs[i].copy()
				pargs[i][isinf] = numpy.nan # crude workaround---pylab.plot can't cope with infinite values
	if allvec:
		for i in range(len(pargs)):
			if isinstance(pargs[i], numpy.matrix):  pargs[i] = pargs[i].A
			if isinstance(pargs[i], numpy.ndarray): pargs[i] = pargs[i].flatten()	
	pargs = tuple(pargs)

	pylab = load_pylab()
	if not hold: # undo a few known hangovers before clearing
		ax = pylab.gca()
		ax.set_aspect('auto') # image aspect ratio
		ax.set_ylim(sorted(ax.get_ylim())) # ydir reversed
		rmcolorbar(drawnow=False) # colorbar
		pylab.cla()
	
	if stem:
		if len(pargs) == 1: pargs = (list(range(pargs[0].shape[axis])),) + pargs
		p = pylab.stem(*pargs, **kwargs)
		x = pargs[0]
		xl = numpy.r_[float(x[0]), float(x[-1])]
		if len(x) == 1: xl += [-0.5, 0.5]
		else: xl += numpy.r_[float(x[0]-x[1]), float(x[-1]-x[-2])] / 2.0
		pylab.gca().set_xlim(xl)
	else:
		p = pylab.plot(*pargs,**kwargs)
	pylab.grid(grid)
	ax = pylab.gca()
	if balance != None:
		yl = numpy.array(ax.get_ylim())
		yl = balance + numpy.array([-1,+1]) * max(abs(yl - balance))
		ax.set_ylim(yl)
	if aspect != None: ax.set_aspect({'image':'equal'}.get(aspect, aspect))
	ax.set(**axkwargs)
	if drawnow: pylab.draw()
	return p
	
def plotsig(x, samplingfreq_hz=None, hold=False, axis=0, welch=0, **kwargs):
	"""
	Makes two subplots, showing time-series <x> in the upper panel and its
	amplitude spectrum in the lower panel.  Set <hold> in order to re-use
	a previous figure.
	
	Any additional keyword arguments are passed through to pylab.plot for
	both subplots.
	"""###
	fs = getfs(samplingfreq_hz)
	if fs==None: fs = getfs(x,2.0)
	if hasattr(x, 'x'): x = x.x
	elif hasattr(x, 'y'): x = x.y
	
	if not isnumpyarray(x):
		axis = 0
		if isinstance(x[0], list) or isinstance(x[0], tuple): axis = 1
		x = numpy.array(x,dtype='float')

	xwin = x = project(x,axis).swapaxes(0, axis)
	nsamp = x.shape[0]
	
	class Unfinished(Exception): pass
	if welch==1: xwin = x * project(hanning(nsamp),len(x.shape)-1)
	elif welch > 0: raise Unfinished("Welch periodogram not yet implemented")
	
	t = numpy.arange(0, nsamp) / float(fs)
	ap = fft2ap(fft(xwin,axis=0),samplingfreq_hz=fs,axis=0)
	f = ap['freq_hz']
	a = 20.0 * numpy.log10(ap['amplitude'])

	pylab = load_pylab()
	if not hold: pylab.clf()
	
	pylab.subplot(2,1,1)
	h1 = pylab.plot(t,x,**kwargs)
	ax = pylab.gca()
	ax.set_xlim(t[0], t[-1])
	ax.xaxis.grid(True)
	ax.yaxis.grid(True)
	
	pylab.subplot(2,1,2)
	a[numpy.isinf(a)] = numpy.nan # crude workaround---pylab.plot can't cope with infinite values
	h2 = pylab.plot(f,a,**kwargs)
	ax = pylab.gca()
	ax.set_xlim(f[0], f[-1])
	ax.xaxis.grid(True)
	ax.yaxis.grid(True)

	pylab.draw()

def wavplotsig(self, *pargs, **kwargs): plotsig(self, *pargs, **kwargs)
try:
	try: from WavTools.Base import wav
	except ImportError: from BCPy2000.WavTools.Base import wav
	wav.plotsig = wavplotsig
except:
	pass

def imagesc(img, x=None, y=None, hold=False, drawnow=True, aspect='image', balance=None, colorbar=False, grid=False, ydir='reverse', **kwargs):
	kwargs['interpolation'] = kwargs.get('interpolation', 'nearest')
	kwargs['cmap'] = kwargs.get('cmap', 'kelvin_i')
	cbkwargs = filterdict(kwargs, startswith='colorbar')
	cbkwargs = dict([(k[len('colorbar'):],v) for k,v in list(cbkwargs.items())])
	axlabelkwargs = filterdict(kwargs, ['title', 'xlabel', 'ylabel'])
	auto_aspect = {'equal':False, 'image':False, 'auto':True}.get(aspect)
	if auto_aspect is None: raise ValueError('aspect should be "image" or "auto"')
	pylab = load_pylab()
	ax = pylab.gca()
	if not hold: ax.cla()
	
	img = numpy.asarray(img)

	xlab,ylab = None,None
	if isinstance(x, (tuple,list,numpy.ndarray)) and isinstance(x[0], str): xlab,x = x,None
	if isinstance(y, (tuple,list,numpy.ndarray)) and isinstance(y[0], str): ylab,y = y,None
	if x is None: x = numpy.arange(img.shape[1], dtype=numpy.float64)
	if y is None: y = numpy.arange(img.shape[0], dtype=numpy.float64)
	x = numpy.asarray(x).flatten()
	y = numpy.asarray(y).flatten()
	if len(x) == 1: xl = [ x[ 0 ] - 0.5, x[ 0 ] + 0.5 ]	
	else: xl = [x[0] - 0.5 * (x[1]-x[0]),   x[-1] + 0.5 * (x[-1]-x[-2])]
	if len(y) == 1: yl = [ y[ 0 ] - 0.5, y[ 0 ] + 0.5 ]	
	else: yl = [y[0] - 0.5 * (y[1]-y[0]),   y[-1] + 0.5 * (y[-1]-y[-2])]
	tol = 1e-10
	regular = ( len(x) == 1 or (numpy.diff(x).ptp() < tol) ) and ( len(y) == 1 or (numpy.diff(y).ptp() < tol) )
	if ydir == 'reverse': yl = yl[::-1]; origin='upper'
	else: origin='lower'
	if not regular: raise ValueError('x and y data must be regularly spaced') # NonUniformImage just doesn't play nice, with ydir or with clim scaling
	h = pylab.imshow(img, extent=xl+yl, origin=origin, **kwargs)		
	if xlab != None: ax.set_xticks(x); ax.set_xticklabels(list(xlab))
	if ylab != None: ax.set_yticks(y); ax.set_yticklabels(list(ylab))
	if auto_aspect: ax.set_aspect('auto')
	else: ax.set_aspect('equal')
	pylab.grid(grid)
	ax.set_xlim(xl)
	ax.set_ylim(yl)
	ax.set(**axlabelkwargs)
	if balance != None:
		c = numpy.array(h.get_clim())
		c = balance + numpy.array([-1,+1]) * max(abs(c - balance))
		h.set_clim(c)
	if colorbar: _colorbar(parent_axes=ax, **cbkwargs)
	else: rmcolorbar(parent_axes=ax)
	if drawnow: pylab.draw()
	return h

def scatterplot(*pargs, **kwargs):
	"""
	scatterplot(xy)         # scatterplots xy[:,1].flat against xy[:,0].flat
	scatterplot(xy1, 'r*')
	scatterplot(xy1, 'r*', xy2, 'bs')
	scatterplot(xy1, 'r*', xy2, 'bs', markersize=10, hold=False, drawnow=True)
	
	scatterplot(x, y)      # scatterplots y.flat against x.flat
	scatterplot(x1, y1, 'r*')
	scatterplot(x1, y1, 'r*', x2, y2, 'bs')
	scatterplot(x1, y1, 'r*', x2, y2, 'bs', markersize=10, hold=False, drawnow=True)
	"""
	x = None
	y = None
	plotargs = []
	connected = []
	if not isinstance(pargs[-1], str):
		pargs = list(pargs) + ['*']
	for i,arg in enumerate(pargs):
		if isinstance(arg, str):
			if y is None:
				if max( x.shape ) == x.size: x, y = numpy.arange( x.size ), x.flatten()
				else: x,y = x[:, 0], x[:, 1]
			plotargs += [x.flatten(), y.flatten(), arg]			
			connected.append('-' in arg)
			x = y = None
		elif x is None: x = arg
		elif y is None: y = arg
	h = plot(*plotargs, **kwargs)
	for hi in zip(h, connected):
		if not connected: hi.set_linestyle('None')
	return h
	
def scatcmp(a, b, hold=False, drawnow=True, markerorder=None, **kwargs):
	kwargs['linestyle'] = kwargs.get('linestyle', 'None')
	grid = kwargs.pop('grid', True)
	grid = {'on':True, 'off':False}.get(grid,grid)
	pylab = load_pylab()
	if not hold: pylab.cla()
	h = plot(a, b, drawnow=False, hold=hold, **kwargs)
	import matplotlib
	if 'marker' not in kwargs:
		if markerorder is None:
			mm = matplotlib.lines.lineMarkers
			markerorder = sorted([m for m in mm if isinstance(m, str) and len(m.strip())]) + sorted([m for m in mm if not isinstance(m, str)])
		for i,hi in enumerate(h): hi.set_marker(markerorder[i % len(markerorder)])
	ax = pylab.gca()
	lim = ax.get_xlim() + ax.get_ylim()
	lim = (min(lim),max(lim))
	ax.set_xlim(lim)
	ax.set_ylim(lim)
	ax.grid(grid)
	pylab.plot(lim, lim, linestyle='-', color=(0.6,0.6,0.6), scalex=False, scaley=False)
	if drawnow: pylab.draw()
	return h

def stem(*pargs, **kwargs):
	kwargs['stem'] = True
	return plot(*pargs, **kwargs)


def colorbar(mappable=None, parent_axes=None, **kwargs):
	"""
	Create a colorbar in such a way that it can be found again and reused, or removed.
	Or, if one already exists associated with the specified parent_axes, use and
	update the existing one.  Use rmcolorbar() to remove.  pylab, y u no do this?
	"""###
	pylab = load_pylab()
	if mappable is None:
		if parent_axes is None: mappable = pylab.gci()
		else: mappable = parent_axes.get_images()[-1]  # NB: doesn't find patch, surf, etc - only images. So to use this with non-images, supply the handle explicitly instead of relying on gca()
	ax = parent_axes
	if ax is None: ax = mappable.get_axes()
	cb = getattr(ax, 'colorbar', None)
	if cb is None:
		kwargs['ax'] = ax
		ax.oldposition = ax.get_position()
	else:
		kwargs['cax'] = cb.ax
		cb.ax.cla()
	title = kwargs.pop('title', None)
	ax.colorbar = pylab.colorbar(mappable, **kwargs)
	if title != None: ax.colorbar.ax.set_title(title)
	pylab.draw()
	return ax.colorbar
_colorbar = colorbar

def rmcolorbar(parent_axes=None, drawnow=True):
	"""
	Remove a colorbar created with colorbar() from this module
	"""###
	pylab = load_pylab()
	ax = parent_axes
	if ax is None: ax = pylab.gca()
	cb = getattr(ax, 'colorbar', None)
	if cb is None: return
	cb.ax.get_figure().delaxes(cb.ax)
	ax.colorbar = None
	if getattr(ax, 'oldposition', None) != None:
		ax.set_position(ax.oldposition)
	if drawnow: pylab.draw()

def load_pylab():
	try:
		import matplotlib,sys
		set_interactive = ('matplotlib.backends' not in sys.modules)
		if set_interactive: matplotlib.interactive(True)
		import pylab
	except:
		print(__name__, "module failed to import pylab: plotting methods will not work")
	else:
		if set_interactive and hasattr(pylab, 'ion'): pylab.ion()
		return pylab

def make_cmap(cmap, name, n=256):
	pylab = load_pylab(); import matplotlib
	cmap = matplotlib.colors.LinearSegmentedColormap(name, cmap, n)
	try: matplotlib.cm.register_cmap(name=name, cmap=cmap)
	except: print("failed to register colormap '%s'" % name)
	return cmap

def complement_cmap(cmap, name=None, n=256):
	pylab = load_pylab()
	if isinstance(cmap, str): cmap = pylab.cm.get_cmap(cmap)
	cmap = getattr(cmap, '_segmentdata', cmap)
	out = {}
	for k in ('red', 'green', 'blue'):
		v = numpy.asarray(cmap[k])
		v[:, 1:] = 1.0 - v[:, 1:]
		out[k] = v
	if name != None: out = make_cmap(out, name=name, n=n)
	return out

def reverse_cmap(cmap, name=None, n=256):
	pylab = load_pylab()
	if isinstance(cmap, str): cmap = pylab.cm.get_cmap(cmap)
	cmap = getattr(cmap, '_segmentdata', cmap)
	out = pylab.cm.revcmap(cmap)
	if name != None: out = make_cmap(out, name=name, n=n)
	return out
	
def show_cmaps(*pargs):
	pylab = load_pylab()
	cmaps = []
	for arg in pargs:
		if isinstance(arg, str): arg = arg.split()
		if isinstance(arg, (tuple,list)): cmaps += arg
		else: cmaps.append(arg)
	if len(cmaps) == 0: cmaps=sorted([m for m in pylab.cm.cmap_d if not m.endswith("_r")])
	pylab.clf()
	for i,cmap in enumerate(cmaps):
		pylab.subplot(len(cmaps), 1, i+1)
		if isinstance(cmap, str): cmap = pylab.cm.get_cmap(cmap)
		pylab.imshow([numpy.linspace(0.0,1.0,256,endpoint=True)], cmap=cmap)
		pylab.gca().set(aspect='auto', yticks=[0.0], yticklabels=[cmap.name])
	pylab.draw()

def subplots(r, c=None, fig=None, **kwargs):
	pylab = load_pylab()
	if fig is None: fig = pylab.gcf()
	if c is None:
		if isinstance(r, int):
			nPlots = r
			ar = fig.get_size_inches()
			ar = float(ar[0])/float(ar[1])
			layout = numpy.r_[nPlots/ar, nPlots*ar] ** 0.5
			i = numpy.argmin(abs(layout - numpy.round(layout)))
			r = int(round(layout[i]))
			c = int(numpy.ceil(nPlots/float(r)))
			if i == 1: r,c = c,r
			while r * (c-1) >= nPlots: c -= 1
			while (r-1) * c >= nPlots: r -= 1
		else:
			r,c = r
	i = 0
	ax = []
	for ri in range(r):
		row = []; ax.append(row)
		for ci in range(c):
			row.append(pylab.subplot(r, c, i+1, **kwargs))
			i += 1
	return numpy.array(ax)
	
def make_kelvin():
	kelvin_i = {
		'red': (
			(0.000, 0.0, 0.0, ),
			(0.350, 0.0, 0.0, ),
			(0.500, 1.0, 1.0, ),
			(0.890, 1.0, 1.0, ),
			(1.000, 0.5, 0.5, ),
		),
		'green': (
			(0.000, 0.0, 0.0, ),
			(0.125, 0.0, 0.0, ),
			(0.375, 1.0, 1.0, ),
			(0.640, 1.0, 1.0, ),
			(0.910, 0.0, 0.0, ),
			(1.000, 0.0, 0.0, ),
		),
		'blue': (
			(0.000, 0.5, 0.5, ),
			(0.110, 1.0, 1.0, ),
			(0.500, 1.0, 1.0, ),
			(0.650, 0.0, 0.0, ),
			(1.000, 0.0, 0.0, ),
		),
	}
	kelvin_i = make_cmap(kelvin_i, 'kelvin_i', 256)
	kelvin_r = complement_cmap(kelvin_i, 'kelvin_r', 256)
	kelvin   = reverse_cmap(kelvin_r, 'kelvin', 256)

make_kelvin()

def curve(x,y=None, hold=False, drawnow=True, **kwargs):
	pylab = load_pylab(); import matplotlib, matplotlib.path
	if not hold: pylab.cla()
	kwargs['facecolor'] = kwargs.get('facecolor', 'None')
	if y is None: xy = x
	else: xy = numpy.c_[x,y]
	codes = [matplotlib.path.Path.MOVETO] + [matplotlib.path.Path.CURVE4] * (len(xy)-1)
	path = matplotlib.path.Path(xy,codes)
	patch = matplotlib.patches.PathPatch(path, lw=2, **kwargs)
	pylab.gca().add_patch(patch)
	if drawnow: pylab.draw()

def cax( handles=None, whichAx='c', lim=None, min=None, max=None, include=None, common=False, balance=None, midpoint=None, auto=False ):
	import pylab
	
	def search_handles( handles, reclevel=0 ):
		out = []
		if reclevel >= 10: zark # allows `debug` of apparent indefinite recursion
		try: handles = list( handles )
		except: handles = [ handles ]
		for h in handles:
			if isinstance( h, ( int, float ) ):
				if h == 0: h = pylab.get_fignums()
				elif pylab.fignum_exists( h ): h = pylab.figure( h )
			while hasattr( h, 'axes' ) and h.axes is not h: h = h.axes   # this was intended to grab the children of a figure, but it can also grab the parent of an AxesImage....
			try: h = list( h )
			except: out.append( h )
			else: out.extend( search_handles( h, reclevel=reclevel+1 ) )
		return out
	if handles in [ 'gca' ]: handles = pylab.gca()
	elif handles in [ None, 'gcf' ]: handles = pylab.gcf()
	handles = search_handles( handles )
	if whichAx == 'c': handles = reduce( list.__add__, [ ax.images for ax in handles ] )  #... with the unintended side-effect that if you pass an AxesImage handle, the function will affect it and any of its siblings on the same axes
	
	all = []
	
	for handle in handles:
		getter = getattr( handle, 'get_' + whichAx + 'lim' )
		setter = getattr( handle, 'set_' + whichAx + 'lim' )
		if auto:
			if whichAx == 'c': handle.autoscale()
			else: handle.autoscale( enable=True, axis=whichAx, tight=True )
		v = list( getter() )
		if lim != None: v = [ numpy.min( lim ), numpy.max( lim ) ]
		if min != None: v[ 0 ] = min
		if max != None: v[ 1 ] = max
		if include != None: v = v + list( include ); v = [ numpy.min( v ), numpy.max( v ) ]
		if balance and midpoint is None: midpoint = 0.0
		if balance == False: midpoint = None
		if midpoint != None:
			hw = numpy.max( [ abs( x - midpoint ) for x in v ] )
			v = [ midpoint - hw, midpoint + hw ]
		setter( v )
		all.extend( v )
		
	if common:
		v = [ numpy.min( all ), numpy.max( all ) ]
		for handle in handles:
			setter = getattr( handle, 'set_' + whichAx + 'lim' )
			setter( v )
	pylab.draw()

def superimpose( figs='all' ):
	import pylab
	if figs == 'all': figs = pylab.get_fignums()
	if not isinstance( figs, ( tuple, list ) ): figs = [ figs ]
	figs = list( figs )
	for i, f in enumerate( figs ):
		if isinstance( f, int ): figs[ i ] = pylab.figure( f )
	geom = figs[ 0 ].canvas.manager.window.geometry()
	for fig in figs[ 1: ]: fig.canvas.manager.window.geometry( geom )
	pylab.draw()

def look( fig=None ):
	import pylab
	if fig is None and len( pylab.get_fignums() ): fig = pylab.gcf()
	if isinstance( fig, int ): fig = pylab.figure( fig )
	try: fig.canvas._tkcanvas.focus_force()
	except:
		try: fig.canvas.manager.window.focus_force() # works but then removes focus from whichever part of the figure is listening to keystrokes :-(
		except: pass

def hotfig( figs='all', align=False ):
	import pylab
	if figs == 'all': figs = pylab.get_fignums()
	if not isinstance( figs, ( tuple, list ) ): figs = [ figs ]
	figs = list( figs )
	for i, f in enumerate( figs ):
		if isinstance( f, int ): figs[ i ] = pylab.figure( f )
	if align == True: align = figs
	if align: superimpose( align )
	def kp(event):
		if not isinstance( event.key, str ): return
		key = event.key.lower()
		if key.startswith( 'alt+' ): key = key[ 4: ]
		codes = list( '1234567890QWERTYUIOP'.lower() )
		if key in codes: target = int( codes.index( key ) ) + 1
		elif key == '[': target = event.canvas.figure.number - 1
		elif key == ']': target = event.canvas.figure.number + 1
		else: target = None
		if target != None and pylab.fignum_exists( target ): look( target )
		
	for fig in figs: fig.canvas.mpl_connect( 'key_press_event', kp )
	look()
