import pickle
import HTSeq
import math
import sys
import scipy
from scipy import stats
from scipy.stats import ttest_ind
from scipy.stats import pearsonr
import numpy as np
import matplotlib as mpl
mpl.use('agg')
import matplotlib.pyplot as plt
import os

#Read expression files
refe={}
sal={}
for subdir, dirs, files in os.walk('.'):
		
	for f in files:
		if f.find('.pkl')!=-1:
			with open(f,'rb') as input:
				expresion=pickle.load(input)
				
			if f.find('MCF7')!=-1: 
				refe[f[:-4]]=expresion
			if f.find('BCC')!=-1:
				sal[f[:-4]]=expresion

print(sal)
print(refe)

#Read all the sequenced genes

genes=[]

for s in refe:
	for g in refe[s]:
		if g not in genes:
			genes.append(g)

for s in sal:
	for g in sal[s]:
		if g not in genes:
			genes.append(g)

#differential expression
pvals={}
mlogpvals={}
lograts={}
dif={}

for ge in genes:
	x=[]
	for s in refe:
		if ge in refe[s]:
			x.append(refe[s][ge])
		else:
			x.append(0)
	y=[]
	for s in sal:
		if ge in sal[s]:
			y.append(sal[s][ge])
		else:
			y.append(0)

	t,pv=ttest_ind(x,y,equal_var=False)
	if math.isnan(pv)==False:
		pvals[ge]=pv
		mlogpvals[ge]=-math.log10(pvals[ge])
		dif[ge]=np.mean(y)-np.mean(x)
	if np.mean(x)>0 and np.mean(y)>0:
		lograts[ge]=math.log(float(np.mean(y))/float(np.mean(x)),2)

import operator

sorted_vals = sorted(pvals.items(), key=operator.itemgetter(1))

#Correct for multiple testing
n=len(sorted_vals)
print(n)
hvals={}
for i in range(n):
	
	hvals[sorted_vals[i][0]]=float(n)*float(sorted_vals[i][1])/float(i+1)

for j in range(n-1):
	for i in range(n-1):
		if hvals[sorted_vals[i][0]]>hvals[sorted_vals[i+1][0]]:
			hvals[sorted_vals[i][0]]=hvals[sorted_vals[i+1][0]]	
	
			



upre=[]

#Differential expression results
with open("Upregulated.txt","w") as text_file:

	for g in hvals:
		if hvals[g]<0.01 and dif[g]>0:
			upre.append(g)
			if g not in lograts:
				line=g+'\t'+str(pvals[g])+'\t'+str(hvals[g])+'\t'+str(dif[g])+'\r'
				text_file.write(line)
				
			else:
				if lograts[g]>1:
					line=g+'\t'+g+'\t'+str(pvals[g])+'\t'+str(hvals[g])+'\t'+str(dif[g])+'\r'
					text_file.write(line)

with open("Downregulated.txt","w") as text_file:

	for g in hvals:
		if hvals[g]<0.01 and dif[g]<0:
			if g not in lograts:
				line=g+'\t'+str(pvals[g])+'\t'+str(hvals[g])+'\t'+str(dif[g])+'\r'
				text_file.write(line)
			else:
				
				if lograts[g]<-1:
					line=g+'\t'+g+'\t'+str(pvals[g])+'\t'+str(hvals[g])+'\t'+str(dif[g])+'\r'
					text_file.write(line)

with open("AllTheGenes.txt","w") as text_file:
	for g in hvals:
		if g not in lograts:
			line=g+'\t'+str(pvals[g])+'\t'+str(hvals[g])+'\t'+str(dif[g])+'\r'
			text_file.write(line)
		else:
			line=g+'\t'+g+'\t'+str(pvals[g])+'\t'+str(hvals[g])+'\t'+str(dif[g])+'\r'
			text_file.write(line)



x=[]
y=[]
for g in lograts:
	x.append(lograts[g])
	y.append(mlogpvals[g])

xu=[]
yu=[]
for g in lograts:
	if hvals[g]<0.01 and lograts[g]>1:
		xu.append(lograts[g])
		yu.append(mlogpvals[g])

xd=[]
yd=[]
for g in lograts:
	if hvals[g]<0.01 and lograts[g]<-1:
		xd.append(lograts[g])
		yd.append(mlogpvals[g])
		
		
plt.scatter(x,y,s=3, c='black')
plt.scatter(xu,yu,s=5, c='red')
plt.scatter(xd,yd,s=5, c='blue')

plt.xlabel('log2(FC)')
plt.ylabel('-log10(p-val)')
plt.savefig('volcano.png')
plt.close()

with open('pvals.pkl','wb') as output:
	pickle.dump(pvals, output)

with open('dif.pkl','wb') as output:
	pickle.dump(dif, output)



