import sys, string
import os.path
import unique
import math
import time
import gene_associations
import OBO_import
import statistics
import random
################# Parse directory files
dirfile = unique

def filepath(filename):
    dir=os.path.dirname(dirfile.__file__)       #directory file is input as a variable under the main            
    fn=os.path.join(dir,filename)
    fn = string.replace(fn,'\\library.zip','') ###py2exe on some systems, searches for all files in the library file, eroneously
    return fn

def read_directory(sub_dir):
    dirfile = unique
    dir=os.path.dirname(dirfile.__file__)
    dir = string.replace(dir,'\\library.zip','') ###py2exe on some systems, searches for all files in the library file, eroneously
    dir_list = os.listdir(dir + sub_dir)
    #add in code to prevent folder names from being included
    dir_list2 = [] 
    for entry in dir_list:
        if entry[-4:] == ".txt" or entry[-4:] == ".csv" or entry[-4:] == ".TXT":
            dir_list2.append(entry)
    return dir_list2

###### Classes ######
class GrabFiles:
    def setdirectory(self,value):
        self.data = value
    def display(self):
        print self.data
    def searchdirectory(self,search_term):
        #self is an instance while self.data is the value of the instance
        file_dir,file = getDirectoryFiles(self.data,str(search_term))
        if len(file)<1: print search_term,'not found'
        return file_dir,file
    
def getDirectoryFiles(import_dir, search_term):
    exact_file = ''; exact_file_dir=''
    dir_list = read_directory(import_dir)  #send a sub_directory to a function to identify all files in a directory
    for data in dir_list:    #loop through each file in the directory to output results
        affy_data_dir = import_dir[1:]+'/'+data
        if search_term in affy_data_dir: exact_file_dir = affy_data_dir; exact_file = data
    return exact_file_dir,exact_file

################# Import and Annotate Data
def eliminate_redundant_dict_values(database):
    db1={}
    for key in database: list = unique.unique(database[key]); list.sort(); db1[key] = list
    return db1

def swapKeyValues(db):
    swapped={}
    for key in db:
        values = db[key]
        for value in values:
            try: swapped[value].append(key)
            except KeyError: swapped[value] = [key]
    return swapped

def identifyGeneFiles(import_dir,gene_file):
    split_name = string.split(gene_file,'.')
    e = GrabFiles(); e.setdirectory(import_dir)
    dir_files = read_directory(import_dir)
    try:
        index = int(split_name[0])
        index_str = str(index)+'.'
        denominator_file_dir,denominator_file = e.searchdirectory(index_str)
    except ValueError: denominator_file_dir =''
    if len(dir_files)==1 or denominator_file_dir=='':
        try: denominator_file_dir,denominator_file = e.searchdirectory(dir_files[0])
        except IndexError: print "WARNING: No denominator file included in the GeneQuery/DenominatorGenes directory. To proceed, place all denominator IDs in a file in that directory. Exiting program."; inp = sys.stdin.readline(); sys.exit()
    return denominator_file_dir

def associateInputSourceWithGene(source_to_gene,source_id_list):
    gene_db={}; count_null = 0
    for source_id in source_id_list:
        try:
            gene_ids = source_to_gene[source_id]
            for gene_id in gene_ids:
                try: gene_db[gene_id].append(source_id)
                except KeyError: gene_db[gene_id] = [source_id]
        except KeyError: count_null+=1
    #print count_null, 'source IDs not imported'
    return gene_db

def cleanUpLine(line):
    line = string.replace(line,'\n','')
    line = string.replace(line,'\c','')
    data = string.replace(line,'\r','')
    data = string.replace(data,'"','')
    return data

def importVersionData(dir):
    filename = dir+'version.txt'; fn=filepath(filename)
    for line in open(fn,'r').readlines():
        data = cleanUpLine(line)
        OBO_version, OBO_date = string.split(data,'\t')
    return OBO_date

def checkDenominatorMatchesInput(input_gene_list,denominator_gene_list,gene_file):
    for id in input_gene_list:
        try: null = denominator_gene_list[id] ###this object was changed from a list to a dictionary for efficiency
        except KeyError: ###Only occurs if an input ID is NOT found in the denominator
            print id
            print '\nWARNING!!!! Job stopped... Denominator gene list does not match the input gene list for',gene_file
            print 'Please use an appropriate denominator list (e.g. all probe sets from the microarray experiment or all known species gene IDs)'; sys.stdin.readline(); sys.exit()

def generateMAPPFinderScores(species_title,species_id,source,mod_db,system_codes,permute):
    global source_data; source_data = source; global mod; mod = mod_db
    global species_code; species_code = species_id
    global species_name; species_name = species_title; global gene_to_mapp
    global permutations; permutations = permute; previous_denominator_file_dir = ''
    global eliminate_redundant_genes; eliminate_redundant_genes = 'yes'
    global permuted_z_scores; global go_annotations
    global original_go_z_score_data; global original_mapp_z_score_data 
    
    start_time = time.time()    
    go_annotations = OBO_import.importPreviousGOAnnotations()
    
    ### Import Gene-to-Nested-GO and Gene-to-MAPP Associations
    try:
        gene_to_go = gene_associations.importGeneGOData(species_code,mod,'nested')
        go_to_gene = OBO_import.swapKeyValues(gene_to_go)
    except IOError: print "Warning!!! The MOD you have selected:",mod,"is missing the appropriate relationship files necessary to run GO-Elite.  Either replace the missing MOD files in Databases/"+species_code+' sub-directories or select a different MOD at run-time. Exiting program.'; inp = sys.stdin.readline(); sys.exit()
    ### Since MAPP tables can be provided by the user, allow the file to be missing
    try: gene_to_mapp = gene_associations.importGeneMAPPData(species_code,mod)
    except IOError: gene_to_mapp = {}
    mapp_to_gene = OBO_import.swapKeyValues(gene_to_mapp)

    OBO_date = importVersionData('OBO/')
    import_dir = '/input/GenesToQuery/'+species_code
    m = GrabFiles(); m.setdirectory(import_dir)
    dir_list = read_directory(import_dir)  #send a sub_directory to a function to identify all files in a directory
    for mappfinder_input in dir_list:    #loop through each file in the directory
        permuted_z_scores={}; original_go_z_score_data={}; original_mapp_z_score_data={}
        gene_file_dir, gene_file = m.searchdirectory(mappfinder_input)
        ###Import Input gene/source-id lists
        input_gene_list,source_data_input = gene_associations.importUIDsForMAPPFinderQuery('input/GenesToQuery/'+species_code+'/'+gene_file,system_codes,'no'); input_count = len(input_gene_list)
        try:
            error_warning = "The directory /input/GenesToQuery/"+species_code+"/DenominatorGenes is not found. Please create the directory and place an appropriate denominator file or files in it. Exiting program."
            try: denominator_file_dir = identifyGeneFiles('/input/GenesToQuery/'+species_code+'/DenominatorGenes',gene_file) ###input is in input\Genes, denominator in
            except WindowsError: print error_warning; sys.stdin.readline(); sys.exit()
        except IOError: print error_warning; sys.stdin.readline(); sys.exit()
        if denominator_file_dir == previous_denominator_file_dir: denom_file_status = 'old'
        else: denom_file_status = 'new'
        if denom_file_status == 'new':
            denominator_gene_list,source_data_denom = gene_associations.importUIDsForMAPPFinderQuery(denominator_file_dir,system_codes,'no'); denom_count = len(denominator_gene_list)
            original_denominator_gene_list=[]
            for id in denominator_gene_list: original_denominator_gene_list.append(id) ###need this to be a valid list not dictionary for permutation analysis
        if len(source_data_input)>0: source_data = source_data_input ###over-ride source_data if a source was identified from the input file
        if source_data != mod:
            mod_source = mod+'-'+source_data
            #checkDenominatorMatchesInput(input_gene_list,denominator_gene_list,gene_file) ###This is checked for the source IDs not associated MOD IDs
            gene_to_source_id = gene_associations.getGeneToUid(species_code,mod_source)
            source_to_gene = OBO_import.swapKeyValues(gene_to_source_id)
            ###Replace input lists with corresponding MOD IDs
            input_gene_list = associateInputSourceWithGene(source_to_gene,input_gene_list)
            if denom_file_status == 'new':
                #source_data = 'Ensembl'  ###This tends to better replicate MAPPFinder version 2.0 statistics
                denominator_gene_list = associateInputSourceWithGene(source_to_gene,denominator_gene_list)
        checkDenominatorMatchesInput(input_gene_list,denominator_gene_list,gene_file) ###This is for only the associated MOD IDs
        
        input_gene_count = len(input_gene_list) ###Count number of genes associated with source input IDs
        if eliminate_redundant_genes == 'yes':
            print 'Attempting to count source IDs that redudantly linked to genes, only once'
        if len(input_gene_list)>0 and len(denominator_gene_list)>0:
            ###Calculate primary z-scores for GO terms
            go_input_gene_count,Rg,input_linked_go = countGenesInPathway(input_gene_list,gene_to_go,'yes')
            if denom_file_status == 'new':
                go_denominator_gene_count,Ng,denom_linked_go = countGenesInPathway(denominator_gene_list,gene_to_go,'yes')
            print Ng,"unique genes, linked to GO and in dataset and", Rg, "unique GO linked genes matching criterion."
            calculateZScores(go_input_gene_count,go_denominator_gene_count,Ng,Rg,go_to_gene,'GO')
            
            ###Calculate primary z-scores for GenMAPP MAPPs
            mapp_input_gene_count,Rm,input_linked_mapp = countGenesInPathway(input_gene_list,gene_to_mapp,'yes')
            mapp_denominator_gene_count,Nm,denom_linked_mapp = countGenesInPathway(denominator_gene_list,gene_to_mapp,'yes')
            print Nm,"unique genes, linked to GenMAPP MAPPs and in dataset and", Rm, "unique GenMAPP MAPPs linked genes matching criterion."
            calculateZScores(mapp_input_gene_count,mapp_denominator_gene_count,Nm,Rm,mapp_to_gene,'MAPP')
            end_time = time.time(); time_diff = int(end_time-start_time)
            print "Primary Z-scores for GO and MAPPs calculated in %d seconds" % time_diff

            ###Begin GO Permute Analysis
            global k; k=0
            start_time = time.time(); print 'Permuting the Gene Ontology analysis %d times' % permutations
            x=0; permute_inputs=[]; permute_go_inputs=[] #; denominator_gene_list = unique.unique_db(denominator_gene_list)
            while x<permutations:
                try: permute_input_list = random.sample(original_denominator_gene_list,input_count); x+=1
                except ValueError: print 'Input count>Denominator',len(original_denominator_gene_list), input_count,'\n','terminating'; sys.stdin.readline(); sys.exit()
                #permute_input_list = random.sample(denominator_gene_list,len(input_gene_list)); x+=1
                #permute_input_list = random.shuffle(original_denominator_gene_list); x+=1; permute_input_list = permute_input_list[:input_count]
                if source_data!=mod:  ###Store the randomly choosen input lists for GenMAPP MAPP Permutation analysis
                    permute_input_list = associateInputSourceWithGene(source_to_gene,permute_input_list)
                    if len(permute_input_list)>len(input_gene_list): k+=1
                permute_inputs.append(permute_input_list) 
                permute_go_input_gene_count,null,null = countGenesInPathway(permute_input_list,gene_to_go,'no'); permute_input_list=[]
                permute_go_inputs.append(permute_go_input_gene_count)
            calculatePermuteZScores(permute_go_inputs,go_denominator_gene_count,Ng,Rg)
            calculatePermuteStats(original_go_z_score_data)
            adjustPermuteStats(original_go_z_score_data)
            end_time = time.time(); time_diff = int(end_time-start_time)
            print "Permuted p-values for GO calculated in %d seconds" % time_diff
            go_headers = formatHeaders(gene_file,input_count,input_linked_go,denom_count,denom_linked_go,Rg,Ng,'GO',OBO_date)
            exportPathwayData(original_go_z_score_data,gene_file,go_headers,'GO')

            start_time = time.time(); print 'Permuting the GenMAPP MAPP analysis %d times' % permutations
            permute_mapp_inputs=[]
            ###Begin GenMAPP MAPP Permute Analysis
            for permute_input_list in permute_inputs:
                permute_mapp_input_gene_count,null,null = countGenesInPathway(permute_input_list,gene_to_mapp,'no')
                permute_mapp_inputs.append(permute_mapp_input_gene_count)
            calculatePermuteZScores(permute_mapp_inputs,mapp_denominator_gene_count,Nm,Rm)
            calculatePermuteStats(original_mapp_z_score_data)
            adjustPermuteStats(original_mapp_z_score_data)
            end_time = time.time(); time_diff = int(end_time-start_time)
            print "Permuted p-values for GenMAPP MAPPs calculated in %d seconds" % time_diff
            mapp_headers = formatHeaders(gene_file,input_count,input_linked_mapp,denom_count,denom_linked_mapp,Rm,Nm,'MAPP',OBO_date)
            exportPathwayData(original_mapp_z_score_data,gene_file,mapp_headers,'local')
            
            previous_denominator_file_dir = denominator_file_dir
            permute_inputs=[]; permute_go_inputs=[]; permute_mapp_inputs=[]
            go_input_gene_count=[]; mapp_input_gene_count=[]
        else:
            if len(input_gene_list)==0: print 'WARNING!!!! The length of the input file',mappfinder_input,'is zero. Length of input ID list is',input_count,'for system:',[source_data_input],[source]
            if len(denominator_gene_list)==0: print 'WARNING!!!! The length of the denominator file',denominator_file_dir,'is zero. Length of input ID list is',denom_count,'for system:',[source_data_denom],[source]

def formatHeaders(gene_file,input_count,input_linked,denom_count,denom_linked,R,N,pathway_type,OBO_date):
    headers = []
    headers.append('GO-Elite MAPPFinder Results')
    headers.append('File:')
    headers.append('Table:')
    headers.append('Database: Based on OBO-Database version: '+OBO_date)
    headers.append('colors:')
    t = time.localtime(); dt = str(t[1])+'/'+str(t[2])+'/'+str(t[0])
    headers.append(dt)
    headers.append(species_name)
    headers.append('Pvalues = true')
    headers.append('Calculation Summary:')
    headers.append(str(input_count)+' '+source_data+' source identifiers supplied in the input file:'+gene_file)
    headers.append(str(input_linked)+' source identifiers meeting the filter linked to a '+mod+' ID.')
    headers.append(str(R)+' genes meeting the criterion linked to a GO term.')
    headers.append(str(denom_count)+' source identifiers in this dataset.')
    headers.append(str(denom_linked)+' source identifiers linked to a '+mod+' ID.')
    headers.append(str(N)+' Genes linked to a GO term.')
    headers.append('The z score is based on an N of '+str(N)+' and a R of '+str(R)+' distinct genes in the GO.\n')
    if pathway_type == 'GO':
        title = ['GOID','GO Name','GO Type','Number Changed Local','Number Measured Local','Number in GO Local','Percent Changed Local','Percent Present Local','Number Changed','Number Measured','Number in GO','Percent Changed','Percent Present','Z Score','PermuteP','AdjustedP']
        title = string.join(title,'\t'); headers.append(title)
    else:
        title = ['MAPP Name','Number Changed','Number Measured','Number on MAPP','Percent Changed','Percent Present','Z Score','PermuteP','AdjustedP']
        title = string.join(title,'\t'); headers.append(title)
    header_str = string.join(headers,'\n')
    return header_str+'\n'
                                                               
def exportPathwayData(original_pathway_z_score_data,gene_file,headers,pathway_type):
    new_file = 'input//MAPPFinder//'+gene_file[:-4]+'-'+pathway_type+'.txt'
    global sort_results
    fn=filepath(new_file); data = open(fn,'w'); data.write(headers); sort_results=[]
    print "Results for",len(original_pathway_z_score_data),"pathways exported to",new_file
    for pathway in original_pathway_z_score_data:
        zsd=original_pathway_z_score_data[pathway]
        try: results = [zsd.Changed(), zsd.Measured(), zsd.InPathway(), zsd.PercentChanged(), zsd.PercentPresent(), zsd.ZScore(), zsd.PermuteP(), zsd.AdjP()]
        except AttributeError: print pathway,len(permuted_z_scores[pathway]);kill
        try: ###This is unnecessary, unless using the non-nested GO associations (which can have out of sync GOIDs)
            if pathway_type == 'GO':
                s = go_annotations[pathway]
                annotations = [s.GOID(),s.GOName(),s.GOType()]; annotations += ['']*5; results = annotations + results 
            else:
                results = [pathway] + results
            results = string.join(results,'\t') + '\n'
            sort_results.append([float(zsd.PermuteP()),-1/float(zsd.Measured()),results])
        except KeyError: null = []
    
    sort_results.sort()
    for values in sort_results:
        results = values[2]
        data.write(results)
    data.close()
    
def swapKeyValuesTuple(db):
    swapped={}
    for key in db:
        values = tuple(db[key]) ###If the value is not a list, make a list
        swapped[values] = [key]
    swapped = eliminate_redundant_dict_values(swapped)
    return swapped

class ZScoreData:
    def __init__(self,pathway,changed,measured,zscore,null_z,in_pathway):
        self._pathway = pathway; self._changed = changed; self._measured = measured
        self._zscore = zscore; self._null_z = null_z; self._in_pathway = in_pathway
    def PathwayID(self): return self._pathway
    def Changed(self): return str(self._changed)
    def Measured(self): return str(self._measured)
    def InPathway(self): return str(self._in_pathway)
    def ZScore(self): return str(self._zscore)
    def SetP(self,p): self._permute_p = p
    def PermuteP(self): return str(self._permute_p)
    def SetAdjP(self,adjp): self._adj_p = adjp
    def AdjP(self): return str(self._adj_p)
    def PercentChanged(self):
        try: pc = float(self.Changed())/float(self.Measured())*100
        except ZeroDivisionError: pc = 0
        return str(pc)
    def PercentPresent(self):
        try: pp = float(self.Measured())/float(self.InPathway())*100
        except ZeroDivisionError: pp = 0
        return str(pp)
    def NullZ(self): return self._null_z
    def Report(self):
        output = self.PathwayID()
        return output
    def __repr__(self): return self.Report()
    
def calculateZScores(pathway_input_gene_count,pathway_denominator_gene_count,N,R,pathway_db,pathway_type):
    """where N is the total number of genes measured: 
    R is the total number of genes meeting the criterion:
    n is the total number of genes in this specific MAPP: 
    r is the number of genes meeting the criterion in this MAPP: """
    for pathway in pathway_db:
        if pathway in pathway_denominator_gene_count:
            n = pathway_denominator_gene_count[pathway]
            if pathway in pathway_input_gene_count: r = pathway_input_gene_count[pathway]
            else: r = 0
        else: n = 0; r = 0
        if n != 0:
            z = statistics.zscore(r,n,N,R)
            null_z = statistics.zscore(0,n,N,R)
            genes_in_pathway = len(pathway_db[pathway])
            zsd = ZScoreData(pathway,r,n,z,null_z,genes_in_pathway)
            if pathway_type == 'GO': original_go_z_score_data[pathway] = zsd
            else: original_mapp_z_score_data[pathway] = zsd
            permuted_z_scores[pathway] = [z]
            #if '06878' in pathway: print pathway, z, null_z, r,n, N, R;kill

def calculatePermuteZScores(permute_pathway_inputs,pathway_denominator_gene_count,N,R):
    for pathway_input_gene_count in permute_pathway_inputs:  
        for pathway in pathway_input_gene_count:
            r = pathway_input_gene_count[pathway]
            n = pathway_denominator_gene_count[pathway]
            z = statistics.zscore(r,n,N,R)
            permuted_z_scores[pathway].append(abs(z))
            #if pathway == '0005488':
            #a.append(r)
                
def calculatePermuteStats(original_pathway_z_score_data):
    for pathway in original_pathway_z_score_data:
        zsd = original_pathway_z_score_data[pathway]
        z = abs(permuted_z_scores[pathway][0])
        permute_scores = permuted_z_scores[pathway][1:] ###Exclude the true value
        nullz = zsd.NullZ()
        if abs(nullz) == z: ###Only add the nullz values if they can count towards the p-value (if equal to the original z)
            null_z_to_add = permutations - len(permute_scores)
            permute_scores+=[abs(nullz)]*null_z_to_add ###Add null_z's in proportion to the amount of times there were not genes found for that pathway
        if len(permute_scores)>0: p = permute_p(permute_scores,z)  
        else: p = 1
        #if p>1: p=1
        zsd.SetP(p)

def adjustPermuteStats(original_pathway_z_score_data):
 
    #1. Sort ascending the original input p value vector.  Call this spval.  Keep the original indecies so you can sort back.
    #2. Define a new vector called tmp.  tmp= spval.  tmp will contain the BH p values.
    #3. m is the length of tmp (also spval)
    #4. i=m-1
    #5  tmp[ i ]=min(tmp[i+1], min((m/i)*spval[ i ],1)) - second to last, last, last/second to last
    #6. i=m-2
    #7  tmp[ i ]=min(tmp[i+1], min((m/i)*spval[ i ],1))
    #8  repeat step 7 for m-3, m-4,... until i=1
    #9. sort tmp back to the original order of the input p values.
    
    global spval; spval=[]
    for pathway in original_pathway_z_score_data:
        zsd = original_pathway_z_score_data[pathway]
        p = float(zsd.PermuteP())
        spval.append([p,pathway])
        
    spval.sort(); tmp = spval; m = len(spval); i=m-2; x=0 ###Step 1-4
    
    while i > -1:
        tmp[i]=min(tmp[i+1][0], min((m/(i+1))*spval[i][0],1)),tmp[i][1]; i -= 1
        
    for (adjp,pathway) in tmp:
        zsd = original_pathway_z_score_data[pathway]
        zsd.SetAdjP(adjp)
        
def permute_p(null_list,true_value):
    y = 0; z = 0; x = permutations
    for value in null_list:
        if value >= true_value: y += 1
    #if true_value > 8: global a; a = null_list; print true_value,y,x;kill
    return (float(y)/float(x))  ###Multiply probabilty x2?

def countGenesInPathway(gene_list,gene_to_pathway,count_linked_source):
    pathway_count={}; associated_genes={}; linked_source={}
    ### Add genes to a dictionary of pathways to get unique counts (could count directly, but biased by redundant source-id associations with MOD)
    for gene in gene_list:
        if source_data != mod and eliminate_redundant_genes == 'yes':
            gene_id = tuple(gene_list[gene]) ### switches gene with list of source_ids (if made unique, decreased redundant association)
            if count_linked_source == 'yes':
                for id in gene_id: linked_source[id] = []
        else: gene_id = gene; linked_source[gene_id] = []
        if gene in gene_to_pathway:
            associated_genes[gene_id] = []
            pathways = gene_to_pathway[gene]
            for pathway in pathways:
                try: pathway_count[pathway].append(gene_id)
                except: pathway_count[pathway] = [gene_id]
    ### Count unique gene or source set associations per pathway
    #global pathway_redundancies; pathway_redundancies={}
    linked_count=0
    for pathway in pathway_count:
        unique_genes = unique.dictionary(pathway_count[pathway])  ###more efficient code for unique
        """###Code used to assess redundancy in probeset assignments
        if source_data != mod:
            redundant_source_ids={}
            for source_id_set in unique_genes:
                for source in source_id_set:
                    try: redundant_source_ids[source].append(source_id_set)
                    except KeyError: redundant_source_ids[source] = [source_id_set]
            for id in redundant_source_ids:
                redundant_source_ids_filtered={}
                if len(redundant_source_ids[id])>1:
                    redundant_source_ids_filtered[id] = redundant_source_ids[id]
                    try: pathway_redundancies[pathway].append([redundant_source_ids_filtered])
                    except KeyError: pathway_redundancies[pathway] = [redundant_source_ids_filtered]
        """
        count = len(unique_genes)
        pathway_count[pathway] = count
    unique_associated_gene_count = len(associated_genes)
    linked_count = len(linked_source)
    return pathway_count, unique_associated_gene_count, linked_count

if __name__ == '__main__':
    species_name = 'Homo Sapiens'; species_code = 'Hs'; source_data = 'Affymetrix'; mod = 'Ensembl'
    system_codes={}; system_codes['X'] = 'Affymetrix'#; system_codes['L'] = 'EntrezGene'
    generateMAPPFinderScores(species_name,species_code,source_data,mod,system_codes,20)
    
#!/usr/bin/python
###########################
#Program:	GO-elite.py
#Author:	Nathan Salomonis
#Date:		12/12/06
#Website:	http://www.genmapp.org
#Email:	nsalomonis@gmail.com
###########################


        
    