#!/usr/bin/python
# ***************************************************************
# Name:      PHYLOmap.py
# Purpose:   Draws phylogenetic tree with a heatmap based on interactive tree of life webservice (http://itol.embl.de)
# Dependencies: urllib2_file (https://github.com/seisen/urllib2_file)
# Version:   0.4
# History:   0.1 -> 0.2 (Included support for LCAClassifier http://code.google.com/p/lcaclassifier/)
#            0.2 -> 0.3 (Included support for RDP Classifier http://sourceforge.net/projects/rdp-classifier/files/ + modified bash one-liners at different taxonomic levels)
#            0.3 -> 0.4 (Included more options for exported pdf)
# Authors:   Umer Zeeshan Ijaz (Umer.Ijaz@glasgow.ac.uk)
#                 http://userweb.eng.gla.ac.uk/umer.ijaz
# Created:   2014-03-26
# License:   Copyright (c) 2014 Computational Microbial Genomics Group, University of Glasgow, UK
#
#            This program 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/>.
# **************************************************************/     

import urllib2,getopt
import urllib2_file
import sys
import os
import re
import subprocess
import math

def run_prog(prog):
    p = subprocess.Popen(prog, shell=True, stdout=None, stderr=None)
    retval = p.wait()

def usage():
   print 'Usage:'
   print '\tpython PHYLOmap.py -i <input_file> -o <output_folder> [OPTIONS]'
   print '\nOptions:'
   print '\t -s (--sqrt_transformation)     Apply square-root transformation to the data before generating the heatmap (Default: False)'
   print '\t -w (--heatmap_width) NUM       Width of heatmap in pixels (Default: 20)'
   print '\t -l (--label) STR               Label of the figure (Default: \"PHYLOmap_label\")'
   print '\t -t (--title_tree) STR          Title of the generated tree (Default: \"PHYLOmap_tree\")'
   print '\t -x (--min_color) STR           Minimum color in the heatmap (Default: \"#FFFFFF\")'
   print '\t -y (--mid_color) STR           Mid color in the heatmap (Default: \"#99CCFF\")'
   print '\t -z (--max_color) STR           Maximum color in the heatmap (Default: \"#007FFF\")'
   print '\t -r (--tree_file) STR           Newick tree file incase you dont want to use NCBI\'s taxonomy (Default:\'\')'
   print '\t --display_mode STR             Tree display mode; \"circular\" or \"normal\" (Default: \"circular\")'
   print '\t --font_size NUM                Font size to be used for leaf labels (Default: 10)'
   print '\t --line_width NUM               Line width in pixels (Default: 1)'
   print '\t --scale_factor NUM             Default horizontal tree scale will be multiplied with this value (Default: 0.5)'

def main(argv):

   #=== Default parameters ===#
   input_file=''
   tree_file=''
   output_folder=''
   sqrt_transformation=False
   data_max_value=sys.float_info.min
   data_min_value=sys.float_info.max
   heatmap_width="20"
   label_figure='PHYLOmap_label'
   post_data={}
   min_color="#FFFFFF"
   mid_color="#99CCFF"
   max_color="#007FFF"
   title_tree="PHYLOmap_tree"
   display_mode="circular"
   font_size="10"
   line_width="1"
   scale_factor="0.5"
   #/=== Default parameters ===#

   try:       
	opts, args =getopt.getopt(argv,"hsi:o:w:l:t:x:y:z:r:",["input_file=","output_folder=","sqrt_transformation=","heatmap_width=","label=","min_color=","mid_color=","max_color=","title_tree=","tree_file=","display_mode=","font_size=","line_width=","scale_factor="])    
   except getopt.GetoptError:       
	usage()       
	sys.exit(2)   
   for opt, arg in opts:       
	if opt == '-h':          
		usage()          
		sys.exit()       
   	elif opt in ("-i", "--input_file"):          
		input_file = arg       
        elif opt in ("-o", "--output_folder"):          
                output_folder = arg
	elif opt in ("-s", "--sqrt_transformation"):
		sqrt_transformation=True   
        elif opt in ("-w", "--heatmap_width"):
                heatmap_width=str(arg)
        elif opt in ("-l", "--label"):
                label_figure=arg
        elif opt in ("-x","--min_color"):
                min_color=arg
        elif opt in ("-y","--mid_color"):
                mid_color=arg
        elif opt in ("-z","--max_color"):
                max_color=arg
        elif opt in ("-t","--title_tree"):
                title_tree=arg
	elif opt in ("-r","--tree_file"):
		tree_file=arg
	elif opt in ("--display_mode"):
		display_mode=arg
        elif opt in ("--font_size"):
                font_size=str(arg)
        elif opt in ("--line_width"):
                line_width=str(arg)
        elif opt in ("--scale_factor"):
                scale_factor=str(arg)

   if input_file=='' or output_folder=='':
	usage()
	sys.exit(2)

   if os.path.isdir(output_folder):
	print "ERROR - Folder "+output_folder+" already exists. Remove it and try again!"
	sys.exit(2)
   else:
	run_prog("mkdir "+output_folder)

   scientific_names=[]

   ins=open(input_file,"r")
   out=open(output_folder+'/'+'data.csv','w')
   record_count=0
   for line in ins:
	record=line.split(",")
	if record_count>0:
		scientific_names.append(record[0].replace(' ','_'))
		record[0]=record[0].replace(' ','_')
                if sqrt_transformation==True:
                        record[1:]=[str(math.sqrt(float(x))) for x in record[1:]]
			record[len(record)-1]=record[len(record)-1]+'\n'
		out.write(",".join(record))
		record_max=max(map(float,record[1:]))
		record_min=min(map(float,record[1:]))
		if record_max>data_max_value:
			data_max_value=record_max
		if record_min<data_min_value:
			data_min_value=record_min
	else:
		record[0]='LABELS'
		out.write(",".join(record))
	record_count=record_count+1
   ins.close()
   #Bug in file upload and so have to insert a blank line
   out.write("\n")
   out.close()
   print "Generated the data file: "+output_folder+'/'+'data.csv'
   if tree_file=='':
   	print ''
   	print 'Creating the upload params'	
   	post_data['ncbiIDs']=" ".join(scientific_names)
   	print 'ncbiIDs : '+post_data['ncbiIDs']
   	try:
		u = urllib2.urlopen('http://itol.embl.de/ncbi_tree_generator.cgi', post_data)
   	except urllib2.HTTPError, error:
		print "ERROR: code=%s message=%s" % (error.code, error.msg)
		sys.exit(1)
   	#print "HTTP_CODE=%s" % u.code
   	print ''
   	print 'Uploading NCBI\'s scientific names to http://itol.embl.de/ncbi_tree_generator.cgi. This may take some time depending on how many scientific names are there and how much load there is on the itol server'
   	#This is a hackjob, extracting the newick tree from the webpage itself
   	matchObj=re.match(r'.*<textarea readonly id=\'nameExp\'>(.*?)</textarea>.*',u.read().replace("\n",""),re.M|re.I)
   	if matchObj:
		out=open(output_folder+'/'+'data.nwk','w')
		out.write(matchObj.group(1))
		out.close()
		print "Generated the tree in newick format: "+output_folder+'/'+'data.nwk'	
   	else:
		print "ERROR - There seems to be something wrong with your file. Unable to generate newick tree"
		sys.exit(2)
   else:
	run_prog("cp "+tree_file+" "+output_folder+"/"+'data.nwk')
	print "Copied "+tree_file+" to "+output_folder

   print ''
   print 'Creating the upload params'

   #Create the Itol class
   test = Itol()

   #Set the tree file
   test.add_variable('treeFile',output_folder+'/'+'data.nwk')
   #Add parameters
   test.add_variable('treeName',title_tree)
   test.add_variable('treeFormat','newick')
   test.add_variable('dataset1File',output_folder+'/'+'data.csv')
   test.add_variable('dataset1Label',label_figure)
   test.add_variable('dataset1Separator','comma')
   test.add_variable('dataset1Type','heatmap')
   test.add_variable('dataset1MinPointColor',min_color)
   test.add_variable('dataset1MidPointColor',mid_color)
   test.add_variable('dataset1MaxPointColor',max_color)
   test.add_variable('dataset1HeatmapBoxWidth',heatmap_width)
   test.add_variable('dataset1MinPointValue',str(math.floor(data_min_value)))
   test.add_variable('dataset1MidPointValue',str(math.floor((data_min_value+data_max_value)/2)))
   test.add_variable('dataset1MaxPointValue',str(math.ceil(data_max_value)))
   # Check parameters
   test.print_variables()
   #Submit the tree
   print ''
   print 'Uploading the tree to http://itol.embl.de/batch_uploader.cgi. This may take some time depending on how large the tree is and how much load there is on the itol server'
   good_upload = test.upload()
   if good_upload == False:
        print 'There was an error:'+test.comm.upload_output
        sys.exit(1)

   #Read the tree ID
   print 'Tree ID: '+str(test.comm.tree_id)
   #Read the iTOL API return statement
   print 'iTOL output: '+str(test.comm.upload_output)
   #Website to be redirected to iTOL tree
   print 'Tree Web Page URL: '+test.get_webpage()
   # Warnings associated with the upload
   print "\n".join(test.comm.warnings)

   print ''
   print 'Downloading data from http://itol.embl.de/batch_downloader.cgi. This may take some time depending on how large the tree is and how much load there is on the itol server'
   # Export a pre-made tree to pdf
   itol_exporter = test.get_itol_export()
   export_location = output_folder+'/'+'data.pdf'

   print ''
   print 'Creating export params'	
   itol_exporter.set_export_param_value('displayMode', display_mode)
   itol_exporter.set_export_param_value('format', 'pdf')
   itol_exporter.set_export_param_value('fontSize', font_size)	
   itol_exporter.set_export_param_value('scaleFactor',scale_factor)	
   itol_exporter.set_export_param_value('lineWidth',line_width)	
   itol_exporter.set_export_param_value('omitDashedLines', '1')	
   itol_exporter.set_export_param_value('datasetList','dataset1')
   itol_exporter.print_variables()
   itol_exporter.export(export_location)
   print ''	
   print 'Exported tree to',export_location


#=== Import Comm, Itol, and ItolExport classes from https://github.com/albertyw/itol-api ===#
'''Copyright (C) 2013 Albert Wang

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
'''

class Comm:
    """
    This class handles communication between the API and the iTOL servers
    This also processes and stores information returned from the server
    """
    def __init__(self):
        """
        Initialize
        """
        self.upload_url = 'http://itol.embl.de/batch_uploader.cgi'
        self.export_url = 'http://itol.embl.de/batch_downloader.cgi'
        self.upload_output = ''
        self.export_output = ''
        self.tree_id = ''
        self.warnings = []

    def upload_tree(self, params):
        """
        Submit the File to Itol using api at self.upload_url
        params is the dictionary of variables that will be uploaded
        """
        url_handle = urllib2.urlopen(self.upload_url, params)
        data = url_handle.read()
        url_handle.close()
        self.upload_output = data
        good_upload = self.parse_upload()
        return good_upload

    def parse_upload(self):
        """
        Parse the raw returned output for uploading to iTOL
        The output is read from self.upload_output
        @return: True if the tree is uploaded successfully or successfully with warnings; False if error occured
        """
        if self.upload_output.find('SUCCESS') != -1:
            # Success, possibly with warnings
            tree_id_start_pos = self.upload_output.rfind('SUCCESS')+9
        else:
            # Fatal Error
            self.warnings = [self.upload_output]
            return False
        self.warnings = self.upload_output[0:].strip().split("\n")
        self.tree_id = self.upload_output[tree_id_start_pos:].strip()
        return True

    def export_image(self, params):
        """
        Submit an export request to Itol using api at self.export_url
        @return: true if connection was established to server
        """
        url_handle = urllib2.urlopen(self.export_url, params)
        self.export_output = url_handle.read()
        url_handle.close()
        return self.export_output


class ItolExport:
    """
    Instantiate the itolexport class with empty params and empty server
    """
    def __init__(self):
        """
        Instantiate class
        """
        self.params = dict({})
        self.comm = Comm()

    ###Setting Export Parameters
    def add_export_param_dict(self, param_dict):
        """
        Add a dictionary of parameters to the parameters to be used when exporting
        @param: dictionary of parameters to be used
        """
        self.params.update(param_dict)

    def set_export_param_value(self, key, value):
        """
        Add a value to the dictionary of parameters to be used when exporting
        @param: dictionary of parameters to be used
        """
        self.params[key] = value

    def get_export_params(self):
        """
        Get the dictionary of parameters to tbe used when exporting
        @return: export the Parameters
        """
        return self.params

    def print_variables(self):
        """
        Print the variables that have been set so far
        """
        for variable_name, variable_value in self.params.items():
        	print variable_name+': '+variable_value

    ###Do Exporting
    def export(self, export_location):
        """
        Call the export process
        Calling this directly assumes that the export filetype is already set in the export params
        @param filelocation: the location to write the export to
        @return: whether the export works
        """
        output = self.comm.export_image(self.params)
        file_handle = open(export_location,'w')
        file_handle.write(output)
        file_handle.close()

class Itol:
    """
    This class handles the main itol functionality
    """
    def __init__(self):
        """
        Initialize a few required variables
        """
        self.variables = dict()
        self.comm = Comm()


    def add_variable(self, variable_name, variable_value):
        """
        Add a variable and its value to this upload.  This function includes 
        some basic variable checking and should be used instead of directly 
        modifying the variables dictionary
        """
        # Variable checking
        if not isinstance(variable_name, str):
            raise TypeError('variable name is not a string')
        if not isinstance(variable_value, str):
            raise TypeError('variable value should be a string')
        if self.is_file(variable_name):
            if not os.path.isfile(variable_value):
                raise IOError('variable name '+variable_name+\
                    ' indicates value should be a file')
            variable_value = open(variable_value, 'r')
        # Add the variable
        self.variables[variable_name] = variable_value
        return True

    @staticmethod
    def is_file(variable_name):
        """
        This returns a boolean whether the string in variable_name is a file
        This is determined by looking at whether "File" is a substring of 
        variable_name; this assumes that variable_name is a string
        """
        if variable_name.find('File')!=-1:
            return True
        else:
            return False

    def upload(self):
        """
        Upload the variables to the iTOL server and return an ItolExport object
        """
        good_upload = self.comm.upload_tree(self.variables)
        if good_upload:
            return self.comm.tree_id
        else:
            self.comm.tree_id = 0
            return False

    def get_webpage(self):
        """
        Get the web page where you can download the Itol tree
        """
        webpage = "http://itol.embl.de/external.cgi?tree="+\
            str(self.comm.tree_id)+"&restore_saved=1"
        return webpage

    def get_itol_export(self):
        """
        Returns an instance of ItolExport in preparation of exporting from the generated tree
        @return: instance of ItolExport
        """
        itol_exporter = ItolExport()
        itol_exporter.set_export_param_value('tree', self.comm.tree_id)
        return itol_exporter

    def print_variables(self):
        """
        Print the variables that have been set so far
        """
        for variable_name, variable_value in self.variables.items():
            if isinstance(variable_value, file):
                print variable_name+': '+variable_value.name
            else:
                print variable_name+': '+variable_value

    def delete_variable(self, variable_name):
        """
        Remove a variable from the dictionary of set variables
        """
        if self.variables.has_key(variable_name):
            del self.variables[variable_name]

#/=== Import Comm, Itol, and ItolExport classes from https://github.com/albertyw/itol-api ===#

if __name__ == '__main__':
	main(sys.argv[1:])