#!/usr/bin/env python3
##import dependency libraries
import sys,time,glob,os,pickle,fileinput,argparse
from subprocess import Popen
from operator import itemgetter
import gc, fileinput
import numpy as np
from scipy import special 
#import pandas as pd
##############################################################
## basic function library
def read_file(tempFile,linesOraw='lines',quiet=False):
    if not quiet:
        print('reading',tempFile)
    f=open(tempFile,'r')
    if linesOraw=='lines':
        lines=f.readlines()
        for i in range(0,len(lines)):
            lines[i]=lines[i].strip('\n')
    elif linesOraw=='raw':
        lines=f.read()
    f.close()
    return(lines)

def make_file(contents,path):
    f=open(path,'w')
    if isinstance(contents,list):
        f.writelines(contents)
    elif isinstance(contents,str):
        f.write(contents)
    f.close()

    
def flatten_2D_table(table,delim):
    #print(type(table))
    if str(type(table))=="<class 'numpy.ndarray'>":
        out=[]
        for i in range(0,len(table)):
            out.append([])
            for j in range(0,len(table[i])):
                try:
                    str(table[i][j])
                except:
                    print(table[i][j])
                else:
                    out[i].append(str(table[i][j]))
            out[i]=delim.join(out[i])+'\n'
        return(out)
    else:
        for i in range(0,len(table)):
            for j in range(0,len(table[i])):
                try:
                    str(table[i][j])
                except:
                    print(table[i][j])
                else:
                    table[i][j]=str(table[i][j])
            table[i]=delim.join(table[i])+'\n'
    #print(table[0])
        return(table)

def strip_split(line, delim = '\t'):
    return(line.strip('\n').split(delim))

def make_table(lines,delim):
    for i in range(0,len(lines)):
        lines[i]=lines[i].strip()
        lines[i]=lines[i].split(delim)
        for j in range(0,len(lines[i])):
            try:
                float(lines[i][j])
            except:
                lines[i][j]=lines[i][j].replace('"','')
            else:
                lines[i][j]=float(lines[i][j])
    return(lines)


def get_file_path(in_path):
    in_path = in_path.split('/')
    in_path = in_path[:-1]
    in_path = '/'.join(in_path)
    return(in_path+'/')


def read_table(file, sep='\t'):
    return(make_table(read_file(file,'lines'),sep))
    
def write_table(table, out_file, sep = '\t'):
    make_file(flatten_2D_table(table,sep), out_file)
    

def import_dict(f):
    f=open(f,'rb')
    d=pickle.load(f)
    f.close()
    return(d)

def save_dict(d,path):
    f=open(path,'wb')
    pickle.dump(d,f)
    f.close()

def cmd(in_message, com=True):
    print(in_message)
    time.sleep(.25)
    if com:
        Popen(in_message,shell=True).communicate()
    else:
        Popen(in_message,shell=True)



##############################################################



###############################################################
###############################################################

parser = argparse.ArgumentParser()


## global arguments
parser.add_argument(
	'-in_dir','-in','-i','-input',
	dest='in_dir',
	type=str)
parser.add_argument(
	'-out_dir','-o','-out',
	dest='out',
	type=str)

args = parser.parse_args()
###############################################################
###############################################################
## functions for calculating the final ranking metric
def norm(vect):
	vect = vect - min(vect)
	vect = vect / max(vect)
	return(vect)

def kl_from_gaus(vect,fudge=1e-4,bins=5):
    ## first make the default Gaussian distribution probability vector
    vect = np.array(vect)
    # gaus_bins = np.histogram(np.random.normal(size=max([10000,len(vect)])),bins=len(vect),density=True)[0]
    # observed = np.histogram(norm(vect), bins = len(vect))[0]
    gaus_bins = np.histogram(np.random.normal(size=max([10000,len(vect)])),bins=bins,density=True)[0]
    observed = np.histogram(norm(vect), bins = bins)[0]

    ## do the kl divergence, but add the fudge factor to prevent inf
    full_kl = special.kl_div(gaus_bins+fudge,observed+fudge)
    return(np.sum(full_kl))


def get_rank_metric(vect):
	sum_kl = kl_from_gaus(vect)
	# multiply by the range
	rank_metric = (max(vect)-min(vect)) * sum_kl
	return(rank_metric)



###############################################################
###############################################################
## if the gprofiler results were generated by cell_signals it will be in a different format than from pyminer_gprofile.py
## the first thing we have to do is figure out which one it is.
## PyMINEr will be a directory, cell signals will be a single file.

if os.path.isdir(args.in_dir):
    print(args.in_dir,'is a directory')
    pyminer = True
    cell_signals = False
elif os.path.isfile(args.in_dir):
    print(args.in_dir,'is a file')
    cell_signals = True
    pyminer = False
else:
    sys.exit("Huston, we have a problem. The file or directory that was passed through couldn't be found/read")


###############################################################
def get_unique_terms(term_ids):
        ## make the terms unique
    print('all terms',len(term_ids))
    term_ids = sorted(list(set(term_ids)))
    print('unqiue',len(term_ids))

    ## catelogue the indices of the unique terms
    term_index_hash = {j:i for i, j in enumerate(term_ids)}

    return(term_ids, term_index_hash)

def get_empty_annotations(term_ids):
    ## make the annotation leader columns
    annotations = []
    for i in range(0,len(term_ids)):
        annotations.append([term_ids[i],None, None])
    return(annotations)

###############################################################
if cell_signals:
    ## catelogue the terms and groups
    all_groups = []
    term_ids = []
    first = True
    for line in fileinput.input(args.in_dir):
        if first:
            first = False
        else:
            ## get the list of term ids and group ids
            temp_line = strip_split(line)
            all_groups.append(temp_line[0]+"_&_"+temp_line[1])
            term_ids.append(temp_line[10])
    fileinput.close()

    all_groups = sorted(list(set(all_groups)))
    all_groups_index_hash = {j:i for i, j in enumerate(all_groups)}
    print(all_groups)
    term_ids, term_index_hash = get_unique_terms(term_ids)

    ## if there aren't any results, there's really nothing left to do...
    if len(term_ids)==0:
        sys.exit('no gprofiler results...')

    ## get the annotations
    annotations = get_empty_annotations(term_ids)

    ## populate the annotations
    first = True
    for line in fileinput.input(args.in_dir):
        if first:
            first = False
        else:
            ## get the list of term ids and group ids
            temp_line = strip_split(line)
            term_row = term_index_hash[temp_line[10]]
            annotations[term_row][0]=temp_line[10]
            annotations[term_row][1]=temp_line[11]
            annotations[term_row][2]=temp_line[13]
    fileinput.close()

    for i in range(0,5):
        print(annotations[i])

    ## make the negative log10 p-value matrix
    neg_log_p = np.zeros((len(term_ids),len(all_groups)))
    first = True
    for line in fileinput.input(args.in_dir):
        if first:
            first = False
        else:
            ## get the list of term ids and group ids
            temp_line = strip_split(line)
            temp_group = temp_line[0]+"_&_"+temp_line[1]
            temp_col_idx = all_groups_index_hash[temp_group]
            temp_row_idx = term_index_hash[temp_line[10]]
            temp_neg_log_p = -1*np.log10(float(temp_line[4]))
            neg_log_p[temp_row_idx,temp_col_idx] = temp_neg_log_p
    fileinput.close()     

###############################################################
###############################################################
if pyminer:
    ## go through all of the files logging all of the pathways that appear

    """
    p-value: 2
    term id: 8
    term name: 11
    """

    os.chdir(args.in_dir)
    all_groups = []
    for f in glob.glob("*gprofiler.txt"):
    	#print(f)
    	all_groups.append(f)

    all_groups = sorted(all_groups)

    ###################################################################
    ## go through the files & grab the unique term ids
    term_ids = []
    for f in all_groups:
    	first = True
    	for line in fileinput.input(f):
    		if first:
    			first = False
    		else:
    			temp_line = strip_split(line)
    			#print(temp_line[2],temp_line[8],temp_line[11])
    			term_ids.append(temp_line[8])
    	fileinput.close()

    ## if there aren't any results, there's really nothing left to do...
    if len(term_ids)==0:
    	sys.exit('no gprofiler results...')


    term_ids, term_index_hash = get_unique_terms(term_ids)

    ###################################################################
    annotations = get_empty_annotations(term_ids)

    for f in all_groups:
        first = True
        for line in fileinput.input(f):
            if first:
                first = False
            else:
                temp_line = strip_split(line)
                term_idx = term_index_hash[temp_line[8]]
                annotations[term_idx][1] = temp_line[9]
                annotations[term_idx][2] = temp_line[11]

    for i in range(0,5):
    	print(annotations[i])


    ###################################################################

    ## make the neg log p-val matrix
    neg_log_p = np.zeros((len(term_ids),len(all_groups)))


    ## add the negative log10 p-values
    for f in range(0,len(all_groups)):
    	first = True
    	for line in fileinput.input(all_groups[f]):
    		if first:
    			first = False
    		else:
    			temp_line = strip_split(line)
    			term_idx = term_index_hash[temp_line[8]]
    			neg_log_p[term_idx, f] = -np.log10(float(temp_line[2]))


    print(neg_log_p)

## normalize the -log10(p-values) within each group
def lin_norm_rows(in_mat,min_range=0,max_range=1):
    in_mat = np.transpose(np.array(in_mat))
    in_mat = in_mat - np.min(in_mat, axis = 0)
    in_mat = in_mat / np.max(in_mat, axis = 0)
    in_mat[np.isnan(in_mat)]=0
    return(np.transpose(in_mat))

## normalize the neg_log_p
norm_neg_log_p = np.transpose(lin_norm_rows(np.transpose(neg_log_p)))


###################################################################
## prioritize them based on the sum kl-divergance*range
def get_full_kl(in_mat, fudge = 1e-4):
    global term_ids
    kl_metric=[]
    for i in range(0,len(term_ids)):
       kl_metric.append(get_rank_metric(in_mat[i,:]))
    ## normalize kl metric
    kl_metric = kl_metric - min(kl_metric) + fudge
    kl_metric = kl_metric / max(kl_metric)

    return(kl_metric)

kl_metric = get_full_kl(neg_log_p)

reordering = sorted(range(len(kl_metric)), key=lambda k: kl_metric[k])[::-1]

## do the same for the individual class weights
individual_class_weights = norm_neg_log_p * np.transpose(np.array([kl_metric]))
individual_class_weights = np.transpose(lin_norm_rows(np.transpose(individual_class_weights)))
kl_individual = get_full_kl(individual_class_weights)

individual_reordering = sorted(range(len(kl_individual)), key=lambda k: kl_individual[k])[::-1]


###################################################################


## make the final output list
out_table_head = [['importance_metric', 'term_id','source','term_name']+all_groups]
out_table = out_table_head[:]
out_table_individual_importance = out_table_head[:]

for i in range(0,len(reordering)):
    temp_idx = reordering[i]
    out_table.append([kl_metric[temp_idx]]+annotations[temp_idx]+neg_log_p[temp_idx,:].tolist())

for i in range(0,len(individual_reordering)):
    temp_idx = individual_reordering[i]
    out_table_individual_importance.append([kl_individual[temp_idx]]+annotations[temp_idx]+individual_class_weights[temp_idx,:].tolist())


## write the file
if args.out[-1]!='/':
    args.out+='/'
    if os.path.isdir(args.out):
        pass
    else:
        cmd("mkdir "+args.out)

write_table(out_table,args.out+"combined_neg_log10p_gprofiler.txt")

write_table(out_table_individual_importance,args.out+'individual_class_importance.txt')