### only do test for overlapped ground trouth and experiment cluster ###
### go annotation clean up ###
import pandas as pd 
import numpy as np 
import random
from collections import defaultdict
from scipy.stats import hypergeom
import json
import os
from collections import Counter as ct 
from statsmodels.stats import multitest as mt

## define a function to calculate the probability of 
## N = total number of common genes from the network and GO term
## n1 = number of genes in a cluster
## n2 = number of genes in a GO term
def hyperG(N, n1, n2, x):
	x = x - 1
	p = hypergeom.sf(x, N, n1, n2)
	p = float(np.format_float_scientific(p, precision = 4))
	return (p)


### Check the case of GO terms ##
def CluGO(fn1, fn2, outName, path, fn3, m):
	## load Cluster file start ##
	clu_dic = defaultdict(dict)
	with open(fn2) as f:
		line = f.readline()
		cnt = 0
		while line:
			line = line.rstrip('\n').split('\t')
			clu_dic[cnt] = line[:len(line)]
			line = f.readline()
			cnt += 1

	clu_cnt = []
	c1 = 0
	c2 = 0
	for i in clu_dic:
		clu_cnt.append(len(clu_dic[i]))
		if len(clu_dic[i]) == 1:
			c1 += 1
		if len(clu_dic[i]) == 2:
			c2 += 1

	outF = open(os.path.join(path, outName+"_cluster_size.txt"), "w")
	outF.write(','.join(map(str, clu_cnt)))
	outF.close()

	print ("Total number of clusters: " + str(len(clu_dic)))
	print ("Number of clusters of size 1: " + str(c1))
	print ("Number of clusters of size 2: " + str(c2))
	print ("Number of clusters of size >= 3: " + str(len(clu_dic) - c1 - c2))
	## load Cluster file end ##

	## load network start ##
	df = pd.read_csv(fn1, delimiter = '\t', header = None)

	gene_list = df.iloc[:, 0].tolist() + df.iloc[:, 1].tolist()
	gene_list=list(set(gene_list))

	## load network end ##

	## read GO term File start ##
	go = pd.read_csv(fn3, header = None, sep = "\t")

	## calculate full size of each go term
	head = go.iloc[:, 0].tolist()
	go_ID = go.iloc[:, 1].tolist()

	# dictionary for GO IDs ###
	dic1_go = defaultdict(dict)
	for i in range(go.shape[0]):
		if dic1_go[go.iloc[i, 1]] == {}:
			dic1_go[go.iloc[i, 1]] = []
		dic1_go[go.iloc[i, 1]].append(go.iloc[i, 0])

	dic1_go = dict(dic1_go)
	#print (dic1_go)

	# dictionary for GO IDs with at least m items ##
	dic2_go = {}
	for i in dic1_go.keys():
		if len(dic1_go[i]) >= m:
			A = list(set(dic1_go[i]).intersection(gene_list))
			if len(A) >= m:
				dic2_go[i] = dic1_go[i]

	#print ("Total number of GO terms: " + str(len(dic1_go)))
	#print ("Total number of GO terms of size >= %s" %m + ": " + str(len(dic2_go)))

	GeneID_go = []
	ID_cnt1 = [] ## which follows the order of adding the key
	GO_term = []

	for i in dic2_go:
		GO_term.append(i)
		ID_cnt1.append(len(dic2_go[i]))
		for j in dic2_go[i]:
			GeneID_go.append(j)

	GeneID_common = list(set(gene_list).intersection(GeneID_go))

	print ("Total number of overlap between Ground Truth and Network Nodes: " + str(len(GeneID_common)) + "\n")

	outF = open(os.path.join(path, outName+"_go_3_size.txt"), "w")
	outF.write(','.join(map(str, ID_cnt1)))
	outF.close()

	### read GO term File end ##

	### only consider genes in the clusters that are annotated by the GO Term, where these GO term annotates
	# at least m genes that are in the network ##
	clu_dic1 = defaultdict(dict)
	for i in clu_dic:
		A = list(set(clu_dic[i]).intersection(GeneID_common))
		clu_dic1[i] = A

	### put all the clusters that a gene belong to as dictionary start ##
	gene_clu_dic = defaultdict(dict)
	for i in clu_dic.keys():
		for j in clu_dic[i]:
			if gene_clu_dic[j] == {}:
				gene_clu_dic[j] = []
			gene_clu_dic[j].append(i)

	### put all the GO term that a gene belong to as dictionary start ##
	gene_go_dic = defaultdict(dict)
	for i in dic2_go.keys():
		for j in dic2_go[i]:
			if gene_go_dic[j] == {}:
				gene_go_dic[j] = []
			gene_go_dic[j].append(i)

	#### print gene-GO association list ###
	outF = open(os.path.join(path, outName+"_gene_GO_asso.txt"), "w")
	for i in gene_go_dic.keys():
		for j in gene_go_dic[i]:
			outF.write('\t'.join([i, j]))
			outF.write('\n')
	outF.close()

	## number of genes annoated by number of GO terms ##
	gene_cnt = []
	for i in gene_go_dic.keys():
		if len(list(set([i]).intersection(gene_list))) > 0:
			gene_cnt.append(len(gene_go_dic[i]))

	outF = open(os.path.join(path, outName+"_gene_go.txt"), "w")
	outF.write(','.join(map(str, gene_cnt)))
	outF.close()

	return clu_dic, gene_list, dic1_go, dic2_go, GO_term, GeneID_common, clu_dic1, gene_go_dic, gene_clu_dic


## P value calculation ###

def Pval(clu_dic1, dic2_go, N, GO_term, a):
	#### Hypergeometric test for each clusters comparing with every GO term ####
	pval = defaultdict(dict)
	P = [] ## for p values
	for i in clu_dic1:
		A = clu_dic1[i]
		for j in dic2_go:
			B = dic2_go[j]
			x = list(set(A).intersection(B))
			p = hyperG(N, len(A), len(B), len(x))
			P.append(p)
	
	#Q = mt.fdrcorrection(P, alpha = a, method = 'indep')[1].tolist()

	P_c = []
	for i in P:
		if i != 1:
			P_c.append(i)

	Q1 = mt.fdrcorrection(P_c, alpha=a, method='indep')[1].tolist()
	Q = []
	j = 0
	for i in P:
		if i == 1:
			Q.append(i)
		else:
			Q.append(Q1[j])
			j = j + 1

	k = 0
	for i in clu_dic1:
		A = clu_dic1[i]
		for j in dic2_go:
			B = dic2_go[j]
			x = list(set(A).intersection(B))
			if pval[i] == {}:
				pval[i] = []
			if len(A) == 0:
				a = [j, N,len(A), len(B), len(x), P[k], Q[k], 0, len(x)*1.0/len(B)]
			elif len(B) == 0:
				a = [j, N,len(A), len(B), len(x), P[k], Q[k], len(x)*1.0/len(A), 0]
			elif len(A) == 0 and len(B) == 0:
				a = [j, N,len(A), len(B), len(x), P[k], Q[k], 0, 0]
			else:
				a = [j, N,len(A), len(B), len(x), P[k], Q[k], len(x)*1.0/len(A), len(x)*1.0/len(B)]
			pval[i].append(a)
			k += 1

	pmin_indx = []

	for i in pval:
		A = []
		for j in range(len(pval[i])):
			A.append(pval[i][j][6])
		ind = np.argmin(A)
		B = ['min-pval', GO_term[ind]]
		B.append(pval[i][ind])
		pmin_indx.append(ind)
		pval[i].append(B)

	return pval

### print Results ##
def Print(pval, a, m, path, outName, clu_dic, clu_dic1, GO_term, dic2_go, gene_list):
	prec_list = []
	reca_list = []
	sigclu_size = []
	sigclu_no = []
	for i in pval:
		if pval[i][-1][2][6] < a:
			if len(clu_dic[i]) >= m:
				prec_list.append(pval[i][-1][2][7])
				reca_list.append(pval[i][-1][2][8])
				sigclu_size.append(pval[i][-1][2][2])
				sigclu_no.append(i)
	
	## 	print precison and list for all enriched cluster with smallest p values ##
	outF = open(os.path.join(path, outName+"_sig_precision.txt"), "w")
	outF.write(','.join(map(str, prec_list)))
	outF.close()
	
	outF = open(os.path.join(path, outName+"_sig_recall.txt"), "w")
	outF.write(','.join(map(str, reca_list)))
	outF.close()

	if len(prec_list) > 0:
		print ("Average PRECISION for enriched clusters: " + str(round(sum(prec_list)/float(len(prec_list)),5)) + "\n")
		print ("Average RECALL for enriched clusters: " + str(round(sum(reca_list)/float(len(reca_list)), 5)) + "\n")
	else:
		print ("Average PRECISION for enriched clusters: " + str(0) + "\n")
		print ("Average RECALL for enriched clusters: " + str(0) + "\n")
			
	prec_list = []
	reca_list = []
	sig_go = []
	for i in pval:
		if pval[i][-1][2][6] < a:
			for j in range(0, len(pval[i])-1):
				if pval[i][j][6] < a:
					if len(clu_dic[i]) >= m:
						A = [i, pval[i][j][7]]
						B = [i, pval[i][j][8]]
						C = [i, GO_term[j]]
						prec_list.append(A)
						reca_list.append(B)
						sig_go.append(C)

	## actual size ##
	clu_size_all = []
	clu_cover = []
	sig_clu_cover = []
	for i in sigclu_no:
		if len(clu_dic[i]) >= m:
			clu_size_all.append(len(clu_dic[i]))
			for j in clu_dic[i]:
				clu_cover.append(j)
			for k in clu_dic1[i]:
				sig_clu_cover.append(k)
	
	## 	print precison and list for all enriched cluster with all enriched GO Terms ##
	outF = open(os.path.join(path, outName+"_sig_allGO_precision.txt"), "w")
	for i in range(len(prec_list)):
		outF.write('\t'.join(map(str, prec_list[i])))
		outF.write('\n')
	outF.close()
	
	outF = open(os.path.join(path, outName+"_sig_allGO_recall.txt"), "w")
	for i in range(len(reca_list)):
		outF.write('\t'.join(map(str, reca_list[i])))
		outF.write('\n')
	outF.close()

	outF = open(os.path.join(path, outName+"_enriched.txt"), "w")

	outF.write('Enriched Cluster ID: ' + '\n')
	outF.write(','.join(map(str, sigclu_no)))
	outF.write('\n')

	outF.write('Number of genes in enriched clusters: ' + '\n')
	outF.write(','.join(map(str, clu_size_all)))
	outF.write('\n')

	outF.write('Number of genes (annotated by at least a GO term) in enriched clusters: ' + '\n')
	outF.write(','.join(map(str, sigclu_size)))
	outF.write('\n')

	outF.close()

	outF = open(os.path.join(path, outName+"_sig_cluster_GO.txt"), "w")
	for i in range(len(sig_go)):
		outF.write('\t'.join(map(str, sig_go[i])))
		outF.write('\n')
	outF.close()

	L = []
	L1 = []
	for i in range(len(sig_go)):
		L.append(sig_go[i][1])
		L1.append(sig_go[i][0])

	L = list(set(L))
	L1 = list(set(L1))

	r1 = 0
	for i in L1:
		if len(clu_dic[i]) >= m:
			r1 += 1

	print ("Number of genes in enriched clusters: " + str(len(set(clu_cover))) + '\n')
	print ("Number of genes (annotated by at least a GO term) in enriched clusters: " + str(len(set(sig_clu_cover))) + '\n')
	print ("Number of unique GO terms for enriched clusters: " + str(len(L)) + '\n')
	print ("Number of enriched clusters: " + str(len(L1)) + '\n')
	print ("Number of enriched clusters of >= %s" %m + ": " + str(r1))

	## Write the p values for all clusters and all GO terms to file ##
	outF = open(os.path.join(path, outName+"_hypergeom_all_clusters.txt"), "w")
	for i in pval:
		outF.write("Cluster %s" %i + "\n")
		for j in pval[i]:
			outF.write('\t'.join(map(str, j)))
			outF.write("\n")
	outF.close()

	## Write the p values for enriched clusters ##
	outF = open(os.path.join(path, outName+"_hypergeom_sig_clusters.txt"), "w")
	for i in pval:
		outF.write("Cluster %s" %i + "\n")
		for j in range(len(pval[i])-1):
			if pval[i][j][5] <= a:
				outF.write('\t'.join(map(str, pval[i][j])))
				outF.write("\n")
	outF.close()

	#############################################################################
	#####       Overlaps between Clusters Start                              ####
	#############################################################################
	### count overlaps between clusters ###
	outF = open(os.path.join(path, outName+"_absOverlap_clusters.txt"), "w")
	for i in clu_dic:
		outF.write("cluster_%s" %i)
		B1 = clu_dic[i]
		for j in clu_dic:
			B2 = clu_dic[j]
			C = list(set(B1).intersection(B2))
			outF.write("," + str(len(C)))
		outF.write('\n')
	outF.close()

	### count overlaps between clusters based on min(A, B)###
	outF = open(os.path.join(path, outName+"_minOverlap_clusters.txt"), "w")
	for i in clu_dic:
		outF.write("cluster_%s" %i)
		B1 = clu_dic[i]
		for j in clu_dic:
			B2 = clu_dic[j]
			C = list(set(B1).intersection(B2))
			D = round(1.0*len(C)/min(len(B1), len(B2)), 5)
			outF.write("," + str(D))
		outF.write('\n')
	outF.close()
	
	### count overlaps between clusters based on union(A, B)###
	outF = open(os.path.join(path, outName+"_unionOverlap_clusters.txt"), "w")
	for i in clu_dic:
		outF.write("cluster_%s" %i)
		B1 = clu_dic[i]
		for j in clu_dic:
			B2 = clu_dic[j]
			B = set(B1 + B2)
			C = list(set(B1).intersection(B2))
			D = round(1.0*len(C)/len(B),5)
			outF.write("," + str(D))
		outF.write('\n')
	outF.close()

	#############################################################################
	#####       Overlaps between Clusters End                                ####
	#############################################################################

	#############################################################################
	#####       Overlaps between GO Start.                                   ####
	#############################################################################

	### count overlaps between GO terms ###
	outF = open(os.path.join(path, outName+"_absOverlap_go.txt"), "w")
	for i in dic2_go:
		outF.write(i)
		B1 = dic2_go[i]
		for j in dic2_go:
			B2 = dic2_go[j]
			C = list(set(B1).intersection(B2))
			outF.write("," + str(len(C)))
		outF.write('\n')
	outF.close()

	### count overlaps between clusters based on min(A, B)###
	outF = open(os.path.join(path, outName+"_minOverlap_go.txt"), "w")
	for i in dic2_go:
		outF.write(i)
		B1 = dic2_go[i]
		for j in dic2_go:
			B2 = dic2_go[j]
			C = list(set(B1).intersection(B2))
			D = round(1.0*len(C)/min(len(B1), len(B2)), 5)
			outF.write("," + str(D))
		outF.write('\n')
	outF.close()
	
	### count overlaps between clusters based on union(A, B)###
	outF = open(os.path.join(path, outName+"_unionOverlap_go.txt"), "w")
	for i in dic2_go:
		outF.write(i)
		B1 = dic2_go[i]
		for j in dic2_go:
			B2 = dic2_go[j]
			B = set(B1 + B2)
			C = list(set(B1).intersection(B2))
			D = round(1.0*len(C)/len(B),5)
			outF.write("," + str(D))
		outF.write('\n')
	outF.close()

	#############################################################################
	#####       Overlaps between GO End                                      ####
	#############################################################################

	#############################################################################
	#####       Overlaps between Cluster and GO Start.                       ####
	#############################################################################

	### count overlaps between clusters and GO  ###
	outF = open(os.path.join(path, outName+"_absOverlap_clusters_go.txt"), "w")
	outF.write("*"+",")
	outF.write(','.join(map(str, dic2_go.keys())))
	outF.write('\n')
	for i in clu_dic:
		outF.write("cluster_%s" %i)
		B1 = clu_dic[i]
		for j in dic2_go:
			B2 = dic2_go[j]
			C = list(set(B1).intersection(B2))
			outF.write("," + str(len(C)))
		outF.write('\n')
	outF.close()

	### count overlaps between clusters and GO based on min(A, B)###
	outF = open(os.path.join(path, outName+"_minOverlap_clusters_go.txt"), "w")
	outF.write("*"+",")
	outF.write(','.join(map(str, dic2_go.keys())))
	outF.write('\n')
	for i in clu_dic:
		outF.write("cluster_%s" %i)
		B1 = clu_dic[i]
		for j in dic2_go:
			B2 = dic2_go[j]
			C = list(set(B1).intersection(B2))
			D = round(1.0*len(C)/min(len(B1), len(B2)), 5)
			outF.write("," + str(D))
		outF.write('\n')
	outF.close()

	### count overlaps between clusters and GO based on union(A, B)###
	outF = open(os.path.join(path, outName+"_unionOverlap_clusters_go.txt"), "w")
	outF.write("*"+",")
	outF.write(','.join(map(str, dic2_go.keys())))
	outF.write('\n')
	for i in clu_dic:
		outF.write("cluster_%s" %i)
		B1 = clu_dic[i]
		for j in dic2_go:
			B2 = dic2_go[j]
			B = set(B1 + B2)
			C = list(set(B1).intersection(B2))
			D = round(1.0*len(C)/len(B),5)
			outF.write("," + str(D))
		outF.write('\n')
	outF.close()

	#############################################################################
	#####       Overlaps between Cluster and GO End.                         ####
	#############################################################################


	## how many nodes in multiple clusters ##
	# use gene_list

	nodes = []
	for i in range(len(gene_list)):
		a = 0
		for j in clu_dic:
			if len(list(set(gene_list[i].split()).intersection(clu_dic[j]))) > 0:
				a = a + 1
		if a > 1:
			nodes.append(a)
	
	dic = ct(nodes)
	dic[1] = len(gene_list) - len(nodes)

	print ("Key is number of clusters, value is count. It means there is $value genes in $key clusters")
	print (str(dic)[8:-1])
	print ("Number of nodes in multiple clusters: " + str(len(nodes)))


## Leave one out prediction start ###
def LOO(clu_dic, clu_dic1, gene_list, dic2_go, gene_go_dic, GeneID_common):
	### put all the clusters that a gene belong to as dictionary start ##
	for i in clu_dic.keys():
		for j in clu_dic[i]:
			if gene_clu_dic[j] == {}:
				gene_clu_dic[j] = []
			gene_clu_dic[j].append(i)
	## end ##

def LOO(path, outName, a, gene_list, GeneID_common, gene_clu_dic, dic2_go, clu_dic1):

	### get all the predictions ###
	outF = open(os.path.join(path, outName+"_gene_GO_pred.txt"), "w")
	P = []
	for i in GeneID_common:
		for j in gene_clu_dic[i]: ## loop all the clusters the gene belongs to ##
			for k in dic2_go.keys(): ## for given gene, perform hypergeom for every GO term ##
				n1 = list(set(clu_dic1[int(j)]) - set([i]))
				n2 = list(set(dic2_go[k]) - set([i])) ## remove the gene in the GO file
				x = list(set(n1).intersection(n2))
				pval = hyperG((len(GeneID_common)-1), len(n1), len(n2), len(x))
				P.append(pval)

	#Q = mt.fdrcorrection(P, alpha = a, method = 'indep')[1].tolist()
	P_c = []
	for i in P:
		if i != 1:
			P_c.append(i)

	Q1 = mt.fdrcorrection(P_c, alpha=a, method='indep')[1].tolist()
	Q = []
	j = 0
	for i in P:
			if i == 1:
				Q.append(i)
			else:
				Q.append(Q1[j])
				j += 1

	l = 0
	for i in GeneID_common:
		for j in gene_clu_dic[i]: ## loop all the clusters the gene belongs to ##
			for k in dic2_go.keys():
				if Q[l] < a:
					pred = [i, k]
					outF.write('\t'.join(pred))
					outF.write('\n')
					outF.flush()
				l += 1
	outF.close()

	## evaluation of predictions ###
	try:
		pred  = pd.read_csv(os.path.join(path, outName+"_gene_GO_pred.txt"), header = None, sep = '\t')
		
		truth  = pd.read_csv(os.path.join(path, outName+"_gene_GO_asso.txt"), header = None, sep = '\t')
		
		pred.columns = ['gene', 'GO']
		truth.columns = ['gene', 'GO']
		
		print ("# of genes predicted: " + str(len(set(pred['gene'].tolist()))))
		print ("# of GO term predicted: " + str(len(set(pred['GO'].tolist()))))
		print ("# of GO term we should have predicted: " + str(len(set(truth['GO'].tolist()))))
		
		pred['gene_go'] = pred[['gene', 'GO']].apply(lambda x: ''.join(x), axis=1)
		truth['gene_go'] = truth[['gene', 'GO']].apply(lambda x: ''.join(x), axis=1)
		
		pred_gene_go = pred['gene_go'].tolist()
		truth_gene_go = truth['gene_go'].tolist()
		
		pred_gene_go = list(set(pred_gene_go))
		truth_gene_go = list(set(truth_gene_go))
		
		print ("# of gene-GO predicted: " + str(len(pred_gene_go)))
		print ("# of gene_GO we should have predicted: " + str(len(truth_gene_go)))
		
		TP = len(list(set(pred_gene_go).intersection(truth_gene_go)))
		
		print ("# of true positives: " + str(TP))
		
		precision = round(TP*1.0/len(pred_gene_go), 5)
		recall = round(TP*1.0/len(truth_gene_go), 5)
		
		print ("PRECISION of LOO prediction is: " + str(precision))
		print ("RECALL of LOO prediction is: " + str(recall))
	except:
		print ("nothing to read of gene_GO_pred.txt file")



