#=======================================================================
# Author: Donovan Parks
#
# Extended error bar plot.
#
# Copyright 2011 Donovan Parks
#
# This file is part of STAMP.
#
# STAMP 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.
#
# STAMP 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 STAMP.  If not, see <http://www.gnu.org/licenses/>.
#=======================================================================

from PyQt4 import QtGui, QtCore

import sys
import math
import numpy as np
from mpl_toolkits.axes_grid import make_axes_locatable, Size

from stamp.plugins.groups.AbstractGroupPlotPlugin import AbstractGroupPlotPlugin, TestWindow, ConfigureDialog
from stamp.plugins.groups.plots.configGUI.extendedErrorBarUI import Ui_ExtendedErrorBarDialog
from stamp.metagenomics import TableHelper

from matplotlib.patches import Rectangle
		
class ExtendedErrorBar(AbstractGroupPlotPlugin):
	'''
	Extended error bar plot.
	'''	 
	def __init__(self, preferences, parent=None):
		AbstractGroupPlotPlugin.__init__(self, preferences, parent)
		
		self.name = 'Extended error bar'
		self.type = 'Statistical'
		
		self.bSupportsHighlight = True
		self.bPlotFeaturesIndividually = False
		
		self.settings = preferences['Settings']
		self.figWidth = self.settings.value('group: ' + self.name + '/width', 7.0).toDouble()[0]
		self.figHeightPerRow = self.settings.value('group: ' + self.name + '/row height', 0.2).toDouble()[0]
		self.sortingField = self.settings.value('group: ' + self.name + '/field', 'p-values').toString()
		self.bShowBarPlot = self.settings.value('group: ' + self.name + '/sequences subplot', True).toBool()
		self.bShowPValueLabels = self.settings.value('group: ' + self.name + '/p-value labels', True).toBool()
		self.bShowCorrectedPvalues = self.settings.value('group: ' + self.name + '/show corrected p-values', True).toBool()
		self.bCustomLimits = self.settings.value('group: ' + self.name + '/use custom limits', False).toBool()
		self.minX = self.settings.value('group: ' + self.name + '/minimum', 0.0).toDouble()[0]
		self.maxX = self.settings.value('group: ' + self.name + '/maximum', 1.0).toDouble()[0]
		self.markerSize = self.settings.value('group: ' + self.name + '/marker size', 30).toInt()[0]
		self.bShowStdDev = self.settings.value('group: ' + self.name + '/show std. dev.', False).toBool()
		self.endCapSize = self.settings.value('group: ' + self.name + '/end cap size', 0.0).toInt()[0]
		self.legendPos = self.settings.value('group: ' + self.name + '/legend position', -1).toInt()[0]

	def mirrorProperties(self, plotToCopy):
		self.name = plotToCopy.name
		
		self.figWidth = plotToCopy.figWidth
		self.figHeightPerRow = plotToCopy.figHeightPerRow
		
		self.sortingField = plotToCopy.sortingField
		
		self.bShowBarPlot = plotToCopy.bShowBarPlot
		self.bShowPValueLabels = plotToCopy.bShowPValueLabels
		
		self.bShowCorrectedPvalues = plotToCopy.bShowCorrectedPvalues
		
		self.bCustomLimits = plotToCopy.bCustomLimits
		self.minX = plotToCopy.minX
		self.maxX = plotToCopy.maxX
		
		self.markerSize = plotToCopy.markerSize
		self.bShowStdDev = plotToCopy.bShowStdDev
		self.endCapSize = plotToCopy.endCapSize
		
		self.legendPos = plotToCopy.legendPos
	
	def plot(self, profile, statsResults):
		# *** Check if there is sufficient data to generate the plot
		if len(statsResults.activeData) <= 0:
			self.emptyAxis()			
			return
		
		features = statsResults.getColumn('Features')
		if len(features) > 200:
			QtGui.QApplication.instance().setOverrideCursor(QtGui.QCursor(QtCore.Qt.ArrowCursor))
			reply = QtGui.QMessageBox.question(self, 'Continue?', 'Profile contains ' + str(len(features)) + ' features. ' +
																		'It may take several seconds to generate this plot. We recommend filtering your profile first. ' + 
																		'Do you wish to continue?', QtGui.QMessageBox.Yes, QtGui.QMessageBox.No)
			QtGui.QApplication.instance().restoreOverrideCursor()
			if reply == QtGui.QMessageBox.No:
				self.emptyAxis()	
				return

		
		# *** Colour of plot elements
		axesColour = str(self.preferences['Axes colour'].name())
		group1Colour = str(self.preferences['Group colours'][profile.groupName1].name())
		group2Colour = str(self.preferences['Group colours'][profile.groupName2].name())
						
		# *** Colour of plot elements
		highlightColor = (0.9, 0.9, 0.9)
		
		# *** Sort data
		if self.sortingField == 'p-values':
			statsResults.activeData = TableHelper.SortTable(statsResults.activeData,\
																												[statsResults.dataHeadings['pValues']], False)
		elif self.sortingField == 'Effect sizes':
			statsResults.activeData = TableHelper.SortTable(statsResults.activeData,\
																												[statsResults.dataHeadings['EffectSize']], 
																												True, True, False)
			
		elif self.sortingField == 'Feature labels':
			statsResults.activeData = Ta
bleHelper.SortTableStrCol(statsResults.activeData,\
																												statsResults.dataHeadings['Features'], False)

		features = statsResults.getColumn('Features')	# get sorted feature labe
ls
					
		# *** Create lists for each quantity of interest
		if statsResults.multCompCorrection.method == 'False discovery rate':
			pValueTitle = 'q-value'
		else:
			pValueTitle = 'p-value'

		if self.bShowCorrectedPvalues:
			pValueLabels = statsResults.getColumnAsStr('pValuesCorrected')
			if statsResults.multCompCorrection.method != 'No correction':
				pValueTitle += ' (corrected)'
		else:
			pValueLabels = statsResults.getColumnAsStr('pValues')
			
		effectSizes = statsResults.getColumn('EffectSize')
		
		lowerCIs = statsResults.getColumn('LowerCI')
		upperCIs = statsResults.getColumn('UpperCI')
		ciTitle = ('%.3g' % (statsResults.oneMinusAlpha()*100)) + '% confidence intervals'
			
		# *** Truncate feature labels
		highlightedFeatures = list(self.preferences['Highlighted group features'])
		if self.preferences['Truncate feature names']:
			length = self.preferences['Length of truncated feature names']
			
			for i in xrange(0, len(features)):
				if len(features[i]) > length+3:
					features[i] = features[i][0:length] + '...'
								
			for i in xrange(0, len(highlightedFeatures)):
				if len(highlightedFeatures[i]) > length+3:
					highlightedFeatures[i] = highlightedFeatures[i][0:length] + '...'
					
		# *** Check that there is at least one significant feature
		if len(features) <= 0:
			self.emptyAxis('No significant features')
			return
				
		# *** Adjust effect size for axis scale
		dominateInSample2 = []
		percentage1 = []
		percentage2 = []
		for i in xrange(0, len(effectSizes)):
			if statsResults.bConfIntervRatio:
				if effectSizes[i] < 1:
					# mirror CI across y-axis
					effectSizes[i] = 1.0 / effectSizes[i]
					lowerCI = effectSizes[i] - (1.0 / upperCIs[i])
					upperCI = (1.0 / lowerCIs[i]) - effectSizes[i]

					lowerCIs[i] = lowerCI
					upperCIs[i] = upperCI

					dominateInSample2.append(i)
				else:
					lowerCIs[i] = effectSizes[i] - lowerCIs[i]
					upperCIs[i] = upperCIs[i] - effectSizes[i] 
			else:
				lowerCIs[i] = effectSizes[i] - lowerCIs[i]
				upperCIs[i] = upperCIs[i] - effectSizes[i]
				if effectSizes[i] < 0.0:
					dominateInSample2.append(i)

		# *** Set figure size
		if self.legendPos == 3 or self.legendPos == 4 or self.legendPos == 8: # bottom legend
			heightBottomLabels = 0.56	# inches
		else:
			heightBottomLabels = 0.4	# inches
			
		heightTopLabels = 0.25
		plotHeight = self.figHeightPerRow*len(features) 
		self.imageWidth = self.figWidth
		self.imageHeight = plotHeight + heightBottomLabels + heightTopLabels
		if self.imageWidth > 256 or self.imageHeight > 256:
				QtGui.QApplication.instance().setOverrideCursor(QtGui.QCursor(QtCore.Qt.ArrowCursor))
				self.