from PyQt4 import QtGui
import matplotlib.pyplot as plt

class GPSModule:
	""" This GPS Module provides analysis functions for the community composition
		on different taxonomic levels """
	helpText = """<b>Composition Analysis</b><br><br>
	Analyze the abundance composition of the dataset. This tool visualizes the
	abundances of a selected subset of the data (e.g. all <i>phyla</i>) in a 
	pie chart.<br><br>As a first step, the MicrobeGPS results have to be 
	scanned and prepared for the analysis. This step may take some time, but
	must be performed only once.<br><br>There are two ways to proceed:<br>
	(a) Select a taxonomic rank from the corresponding list. All taxa with 
	this rank are then listed below. Select the taxon that you want to profile.
	<br>(b) Select a taxon from the tree panel and click the <i>Get from tree
	</i> Button.<br><br>After pressing <i>Show composition</i>, the program
	will show the pie chart in the <i>Figure</i> panel. """		
		
	def __init__(self, GPS):
		self.name = 'Composition Analysis'
		self.GPS = GPS
		
		self.selectedTaxa = []
		self.selectedParent = ''
		
		self.initialAnalysisLabel = QtGui.QLabel('Initial setup')
		self.initialAnalysis = QtGui.QPushButton('Scan GPS results')
		self.initialAnalysis.setToolTip('Run the initial analysis required for the following steps. This may take some time depending on the amount of data.')
		self.initialAnalysis.clicked.connect(self.initial_analysis)
		
		self.taxRankBoxLabel = QtGui.QLabel('Select a Taxonomic rank')
		self.taxRankBox = QtGui.QComboBox()
		self.taxRankBox.currentIndexChanged.connect(self.select_tax_rank)
		self.taxGroupBoxLabel = QtGui.QLabel('Select group')
		self.taxGroupBox = QtGui.QComboBox()
		self.taxGroupBox.currentIndexChanged.connect(self.select_taxon)
		self.taxSelectedText = QtGui.QLabel('Found 0 [Not found]')
		self.getFromTree = QtGui.QPushButton('Get from tree')
		self.getFromTree.setToolTip('Use the currently selected taxon from the tree.')
		self.getFromTree.clicked.connect(self.select_taxon_from_tree)
		self.showComposition = QtGui.QPushButton('Show composition')
		self.showComposition.setToolTip('Calculate and visualize the taxonomic composition.')
		self.showComposition.clicked.connect(self.show_composition)		
		self.taxRankBox.setEnabled(False)
		self.taxGroupBox.setEnabled(False)
		self.getFromTree.setEnabled(False)
		self.showComposition.setEnabled(False)
		
		tabGrid = QtGui.QGridLayout()
		tabGrid.addWidget(self.initialAnalysisLabel,0,0)
		tabGrid.addWidget(self.initialAnalysis, 0,1)
		tabGrid.addWidget(self.taxRankBoxLabel,1,0)
		tabGrid.addWidget(self.taxRankBox, 1,1)
		tabGrid.addWidget(self.taxGroupBoxLabel,2,0)
		tabGrid.addWidget(self.taxGroupBox, 2,1)
		tabGrid.addWidget(self.getFromTree, 3,1)
		tabGrid.addWidget(self.taxSelectedText,4,0)
		tabGrid.addWidget(self.showComposition,4,1)
		tabWidget = QtGui.QWidget()
		tabWidget.setLayout(tabGrid)
		tabWidget.helpText = self.helpText
		
		GPS.toolsTab.addTab(tabWidget,"Comp. Analysis")
		

	def initial_analysis(self):
		""" Scan though the analyis results. Gather the reads of each reference 
		and propagate the information to the higher taxonomic ranks. """
		# first check if all required information is available
		if not hasattr(self.GPS,'taxonomy_nodes') or not hasattr(self.GPS,'taxonomy_ranks'):
			self.GPS.pr('<b><font color="DarkRed">Error in Composition Analysis:</font></b><br><i>    Could not run initial analysis. NCBI Taxonomy is missing!</i>',True)
			return
		if not hasattr(self.GPS,'taxonomy_names'):
			self.GPS.pr('<b><font color="DarkRed">Error in Composition Analysis:</font></b><br><i>    Could not run initial analysis. NCBI Names are missing!</i>',True)
			return
		if not hasattr(self.GPS,'references') or not hasattr(self.GPS,'reads'):
			self.GPS.pr('<b><font color="DarkRed">Error in Composition Analysis:</font></b><br><i>    Could not run initial analysis. Run MicrobeGPS analysis first!</i>',True)
			return
		
		# collect all read names for each reference
		self.reads = dict()
		for ref in self.GPS.references.itervalues():
			taxid = ref.name
			self.reads[taxid] = set()
			for trg in ref.targets.itervalues():
				self.reads[taxid].update(trg.reads) # add all read names to the set
				
		self.rank_tids = dict() # set of TaxIDs for a specific rank
		# calculate the read sets at the higher taxonomic ranks
		for taxid in self.reads.keys():
			# trace the lineage of each taxid and propagate the reads to the higher ranks
			lineage = [taxid]
			while True:
				current = lineage[-1]
				parent = self.GPS.taxonomy_nodes.get(current,None)
				lineage.append(parent)
				if not parent or parent == 1:
					break
			for tid in lineage[1:]:
				rank = self.GPS.taxonomy_ranks.get(tid,'[Not found]')
				tid_list = self.rank_tids.setdefault(rank,set())
				tid_list.add(tid)
				self.reads[tid] = self.reads[taxid].union(self.reads.get(tid,set()))
		
		# we are done with the analysis. Fill the Combo Box widgets accordingly
		self.taxRankBox.clear()
		self.taxRankBox.addItems(self.rank_tids.keys())
		self.taxRankBox.setEnabled(True)
		self.taxGroupBox.setEnabled(True)
		self.getFromTree.setEnabled(True)
		self.showComposition.setEnabled(True)
	
	def select_tax_rank(self,currentItem):
		if currentItem < 0:
			return
		rank = str(self.taxRankBox.currentText())
		# find all taxa of this rank in the dataset
		tax_ids = [tid for tid in self.rank_tids[rank]]
		tax_names = [self.GPS.taxonomy_names.get(tid,tid) for tid in tax_ids]
		tax_names.sort()
		# put names into combo box
		self.taxGroupBox.clear()
		self.taxGroupBox.addItems(tax_names)
		# update label
		self.taxGroupBoxLabel.setText('Select '+rank)

	def select_taxon(self,currentItem):
		if currentItem < 0:
			return
		taxon = str(self.taxGroupBox.currentText())
		children = [tid for tid in self.reads.keys() if self.GPS.taxonomy_names.get(self.GPS.taxonomy_nodes.get(tid))==taxon]
		self.taxSelectedText.setText('Found %i taxa below %s'%(len(children),taxon))
		self.selectedTaxa = children
		self.selectedParent = taxon

	def select_taxon_from_tree(self):
		if len(self.GPS.phyloTreeWidget.selectedItems()) == 0:
			return
		self.GPS.graphicsTab.setCurrentIndex(1)
		taxon = self.GPS.phyloTreeWidget.currentItem().text(0)
		children = [tid for tid in self.reads.keys() if self.GPS.taxonomy_names.get(self.GPS.taxonomy_nodes.get(tid))==taxon]
		self.taxSelectedText.setText('Found %i taxa below %s'%(len(children),taxon))
		self.selectedTaxa = children
		self.selectedParent = taxon
		
	def show_composition(self):
		if len(self.selectedTaxa) == 0:
			return
		tax_names = [self.GPS.taxonomy_names.get(tid,tid) for tid in self.selectedTaxa]
		tax_abundances = [len(self.reads[tid]) for tid in self.selectedTaxa]
		sort_order = sorted(range(len(tax_abundances)), key=tax_abundances.__getitem__)[::-1]
		tax_names = [tax_names[i] for i in sort_order]
		for i in range(7,len(tax_names)):
			tax_names[i] = ''
		tax_abundances = [tax_abundances[i] for i in sort_order]
		self.GPS.graphicsTab.setCurrentIndex(0)
		self.GPS.figure.clear()
		ax = self.GPS.figure.add_subplot(111,aspect=1)
		ax.set_title('Composition of %s'%str(self.selectedParent))
		ax.pie(tax_abundances,labels=tax_names)
		plt.tight_layout()
		self.GPS.canvas.draw()
