# Imports from JFreeChart
from org.jfree.chart             import ChartFactory
from org.jfree.chart             import ChartPanel
from org.jfree.chart.plot        import PlotOrientation
from org.jfree.chart.plot        import XYPlot
from org.jfree.data.xy           import XYSeries
from org.jfree.data.xy           import XYSeriesCollection
from javax.swing                 import JFrame, WindowConstants
from java.awt                    import Dimension, Color, Font, BasicStroke
from org.jfree.chart.axis		 import NumberTickUnit, LogarithmicAxis
from java.text					 import DecimalFormat
from Statistics                  import average, median, differential
import math


# TODO Generalize this function
def analyzeLegends(legends):
	maxIndex = 1
	for legend in legends:
		if legend.startswith("legend"):
			idx = int(legend.split("legend")[1])
			if idx > maxIndex:
				maxIndex = idx
	return "legend%d" % (maxIndex + 1)

class Plot(object):
	"""Plots with JFreeChart
	
	Usage:
		p = Plot("title")
	
		# Plot a scatter plot 
		>>> p.plot(x, y, title="Title", legend="legend", xlabel="x", ylabel="y", legend="legend", lineWidth=2, lineStyle="line")

		# Close the plot
		>>> p.close()

		# Clear the frame
		>>> p.clear()
	
		# x is an iterable for the abscissa
		# y is an iterable for the ordinates  

	"""

	STROKE_MITER_LIMIT = 1.0
	STROKE_DASH_PHASE  = 0.0
	LINE_STYLES = {"dash": [5.0, 20.0], "dot": [2.0, 3.0]}
	PLOT_METHODS = {"line": "createXYLineChart", "scatter": "createScatterPlot"}
	frames = []

	def __makeColor(self, colorName):
		"""Make the corresponding attribute of Color from colorName
		
		Example:
			p.__makeColor("red") # Color.RED

		"""
		try:
			return getattr(Color, colorName)
		except AttributeError:
			print("%s is not a valid color.\n" % colorName)
			print("Color set to red.")
			return getattr(Color, "RED")

	def __init__(self, title="Title"):
		super(Plot, self).__init__()
		self.title = title
		self.frame = JFrame(title=self.title)
		self.frame.setDefaultCloseOperation(WindowConstants.DISPOSE_ON_CLOSE)
		self.manageFrames()
		self.lineStyle = "line"
		self.legend = None
		self.legends = []
		self.stroke = None
		self.lineColor = self.__makeColor("red")

	def manageFrames(self):
		if self.title in [frame.title for frame in self.frames]:
			i = filter(lambda a: self.frames[a].title == self.title, range(len(self.frames)))[0]
			self.frames[i].setVisible(False)
			print("Updated frame %s." % self.frames[i].title)
			del(self.frames[i])
			self.frames.append(self.frame)
		else:
			self.frames.append(self.frame)

	def setLineColor(self, idx, lineColor):
		self.lineColor = self.__makeColor(lineColor)
		self.chart.getPlot().getRenderer().setSeriesPaint(idx, self.lineColor)

	def getLineWidth(self):
		return self.lineWidth

	def setLineWidth(self, idx, lineWidth):
		self.lineWidth = lineWidth
		self.setStroke(idx, self.lineWidth, self.lineStyle)

	def getLineStyle(self):
		return self.lineStyle

	def setLineStyle(self, idx, lineStyle):
		self.lineStyle = lineStyle
		self.setStroke(idx, self.lineWidth, self.lineStyle)

	def getStroke(self):
		return self.stroke

	def setStroke(self, idx, lineWidth, lineStyle):
		# For some strange reason, the lineStyle "line" doesn't work on some 
		# Mac Book pro
		if lineStyle != "line":
			self.stroke = BasicStroke(lineWidth,
								 BasicStroke.CAP_BUTT,
								 BasicStroke.JOIN_ROUND,
								 self.STROKE_MITER_LIMIT,
								 self.LINE_STYLES[lineStyle],
								 self.STROKE_DASH_PHASE)
		else:
			self.stroke = BasicStroke(lineWidth)

		self.chart.getPlot().getRenderer().setSeriesStroke(idx, self.stroke)

	def setLegend(self, legend):
		if legend is None:
			if "legend1" not in self.legends:
				legend = "legend1"
			else:
				legend = analyzeLegends(self.legends)

		self.legends.append(legend)
		self.legend = legend

	def setLine(self, idx, lineWidth, lineStyle, lineColor):
		self.setLineWidth(idx, lineWidth)
		self.setLineStyle(idx, lineStyle)
		if lineColor:
			self.setLineColor(idx, lineColor)

	def plot(self, x, y, xlabel="x", ylabel="y", legend=None, lineWidth=1, lineStyle="line", lineColor=None, plotType="line"):

		if plotType not in ("line", "scatter"):
			msg = "Choose a valid type of plot. Available types are 'line' and 'scatter'"
			raise ValueError(msg)

		selectedPlotType = self.PLOT_METHODS[plotType]

		n = len(x)
		self.setLegend(legend)
		series = XYSeries(self.legend)

		for i in range(n):
			series.add(x[i], y[i])

		self.dataset = XYSeriesCollection()
		self.dataset.addSeries(series)
		createPlot = getattr(ChartFactory, selectedPlotType)
		self.chart = createPlot(self.title, xlabel, ylabel, self.dataset,
											   PlotOrientation.VERTICAL,
											   True, 
											   True, 
											   False)

		# Set the background of the plot
		self.chart.getPlot().setBackgroundPaint(Color.WHITE)

		# RangeAxis (Y-axis) must 
		# - exclude 0 when autoRange is on
		# - set the best NumberTickUnit when the distance between 2 points 
		# is too high in y
		rangeAxis = self.chart.getPlot().getRangeAxis()
		rangeAxis.setAutoRangeIncludesZero(False)
		# rangeAxis.setRange(math.floor(min(y)), math.ceil(max(y)))

		averagedDistance = average(differential(sorted(y)))
		averagedDistance = float("%.2g" % averagedDistance)

		# TODO Refine this arbitrary max value
		if averagedDistance > 0 and math.log10(averagedDistance) > 5:
			rangeAxis.setTickUnit(NumberTickUnit(averagedDistance))

			logAxis = LogarithmicAxis("YYYY")
			numFormat = DecimalFormat("0.##E0")
			rangeAxis.setNumberFormatOverride(numFormat)

		# Create the chart panel
		self.chartPanel = ChartPanel(self.chart)

		# Set the line
		self.setLine(0, lineWidth, lineStyle, lineColor)

		# Set the frame
		self.frame.setDefaultCloseOperation(JFrame.DISPOSE_ON_CLOSE)
		self.frame.setContentPane(self.chartPanel)
		self.frame.pack()
		self.frame.setVisible(True)

	def add(self, x, y, legend=None, lineWidth=1, lineStyle="line", lineColor=None):
		"""Add the given data to the plot"""
		self.setLegend(legend)
		series = XYSeries(self.legend)
		for i in range(len(x)):
			series.add(x[i], y[i])

		self.dataset.addSeries(series)

		# Set the line
		idx = self.chart.getPlot().getDataset().indexOf(self.legend)
		self.setLine(idx, lineWidth, lineStyle, lineColor)

	def clear(self):
		self.dataset.removeAllSeries()

	def close(self):
		self.frame.setVisible(False)

	def closeAll(self):
		for frame in self.frames:
			frame.setVisible(False)
		self.frames = []
		

# Demo

# Useful functions
#import math
#_sin = lambda x: map(math.sin, x)
#_cos = lambda x: map(math.cos, x)
#x = [i * 0.1 for i in range(100)]
#
#p = Plot("My plot")
#p.plot(x=x, y=_sin(x), xlabel="xaxis", ylabel="yaxis", legend="sin(x)", lineStyle="dash", lineWidth=1, plotType="scatter")
#p.add(x, _cos(x), lineColor="blue", legend="cos(x)", lineWidth=1, lineStyle="dash")
#
#p2 = Plot("My other plot")
#p2.plot(x=x, y=_sin(x), xlabel="xaxis", ylabel="yaxis", legend="sin(x)", lineStyle="dash", lineWidth=1, plotType="scatter")

# vim: set noet fdm=marker:
