"""
Bar chart demo with pairs of bars grouped for easy comparison.
"""
import numpy as np

import matplotlib
matplotlib.use('Agg') # Must be before importing matplotlib.pyplot or pylab!

import matplotlib.pyplot as plt
import os

#remove the insertion codes from Chothia codes
def stripChothia(res):
	icode = "ZXCVBNMLKJHGFDSAWQERTYUIOP"
	if res[len(res)-1] in icode:
		res = res[0:(len(res)-1)]
	return res

#Check if a residue is a CDR residue according to the given definition (deff)
def is_CDR(res,deff):
	
	for CDR in definitions[deff]:
		if res in definitions[deff][CDR]:
			return CDR

	return False

##MAIN##
#AA mapping
AAtable = {
'ALA': 'A',
'ARG': 'R',
'ASN': 'N',
'ASP': 'D',
'CYS': 'C',
'GLU': 'E',
'GLN': 'Q',
'GLY': 'G',
'HIS': 'H',
'ILE': 'I',
'LEU': 'L',
'LYS': 'K',
'MET': 'M',
'PHE': 'F',
'PRO': 'P',
'SER': 'S',
'THR': 'T',
'TRP': 'W',
'TYR': 'Y',
'VAL': 'V'}

#CDR Definitions - Chothia+2 on either side.
definitions = {		
		
		"chothia" : {"L1" : ["L22","L23","L24", "L25", "L26", "L27", "L28", "L29", "L30", "L31", "L32", "L33", "L34","L35","L36"], 
				  "L2" : ["L48","L49","L50", "L51", "L52", "L53", "L54", "L55", "L56","L57","L58"],
				  "L3" : ["L87","L88","L89", "L90", "L91", "L92", "L93", "L94", "L95", "L96", "L97","L98","L99"], 
				  "H1" : ["H24","H25","H26", "H27", "H28", "H29", "H30", "H31", "H32","H33","H34"], 
				  "H2" : ["H50","H51","H52", "H53", "H54", "H55", "H56","H57","H58"] ,
				  "H3" : ["H93","H94","H95", "H96", "H97", "H98", "H99", "H100", "H101", "H102","H103","H104"]} 
		
		}

def barplot(groups,labels,colors=['g','b','r','y'],title='Title',ylabel='Scores',xlabel='Chothia Position',results_images=""):


	n_groups = 0
	for entry in groups:
		n_groups = len(groups[entry])
		#get the first one - its users fault if subsequent vectors have more than n groups.	
		break

	fig, ax = plt.subplots()

	index = np.arange(n_groups)
	bar_width = 0.15

	opacity = 0.4
	error_config = {'ecolor': '0.3'}

	ind = 0
	plt.ylim([0,150])
	plt.xlim([0,len(labels)])
	print(groups)
	for group in groups:
		
		rects1 = plt.bar(index+(bar_width*ind), groups[group], bar_width,
				 alpha=opacity,
				 color=colors[ind],
				 error_kw=error_config,
				 label=group)
		ind+=1

	#Add line markers
	plt.axhline(y=100,linewidth=1,color='b')
	plt.axhline(y=80,linewidth=1,color='b')
	plt.axhline(y=60,linewidth=1,color='b')
	plt.axhline(y=40,linewidth=1,color='b')
	plt.axhline(y=20,linewidth=1,color='b')	

	plt.xlabel(xlabel)
	plt.ylabel(ylabel)
	plt.title(title)
	plt.xticks(index + bar_width, labels,rotation='vertical')
	plt.legend(loc=2)

	plt.tight_layout()
	#plt.show()
	#If you want to save the figure rather than show it COMMENTOUT plt.show() above and uncomment the plt.savefit below
	plt.savefig(results_images+"/"+title+'.png')

#Get the background mapping computed using the Chothia Profile bit of PACE
def get_background(chid):
	f = open(os.path.dirname(os.path.realpath(__file__))+'/background_files/'+chid+'.txt','r')
	mapping = dict()
	f_mapping = dict()
	for line in f:
		line = line.strip()
		line = line.split(" ")
		chothia_id = line[0].strip()
		value = line[1]
		freq = line[2]		
		mapping[chothia_id] = int(value)
		f_mapping[chothia_id] = int(freq)
	return mapping,f_mapping

#Gets the CDR annotation
def get_cdr_annotation(pdb):
	
	seenids = []
	cdrs = dict()
	for line in open(pdb,'r'):
		line = line.strip()
		chain = line[21]
		sid = line[22:27].replace(" ","")
		ch_id = chain+sid
		aa = AAtable[line[17:21].replace(" ","")]
		if ch_id in seenids:
			continue
		else:
			seenids.append(ch_id)
		bfactor = int(line[60:63].replace(" ",""))
		cdr = is_CDR(stripChothia(ch_id),"chothia")
		if (cdr != False):
			if cdr not in cdrs:
				cdrs[cdr] = []
			cdrs[cdr].append((ch_id+aa,bfactor))
			#print chain,sid,is_CDR(stripChothia(ch_id),"chothia"),bfactor
	return cdrs# zip(*cdrs['H1'])[0]
	
def translate_results(results_folder):
	
	results_images = results_folder+"/CDR_results/"
	if not os.path.exists(results_images):
		os.mkdir(results_images)
	
	ipatch = get_cdr_annotation(results_folder+'/ABipatch_normalized.pdb')
	docking = get_cdr_annotation(results_folder+'/Docking_normalized.pdb')
	#Background files	
	bg_H,freq_H = get_background('H')
	bg_L,freq_L = get_background('L')
	for cdr in ipatch:
		#Constrain the map
		bg_cdr = []
		freq_cdr = []
		for entry in ipatch[cdr]:
			chothiaid = str(entry[0][0:(len(entry[0])-1)].strip())
						
			if (chothiaid[0]=='H'):
				if chothiaid in bg_H:
					val=bg_H[chothiaid]
					freq = freq_H[chothiaid]
				else:
					val=0
			else:
				if chothiaid in bg_L:
					val=bg_L[chothiaid]
					freq = freq_L[chothiaid]
				else:
					val=0
			bg_cdr.append((entry[0],val))
			if freq>90:
				freq_cdr.append((entry[0],0))
			else:
				freq_cdr.append((entry[0],freq))
		ipatch_uz = list(zip(*ipatch[cdr]))
		docking_uz = list(zip(*docking[cdr]))
		background = list(zip(*bg_cdr))
		frequency = list(zip(*freq_cdr))
		#print len((ipatch_uz[0])),len((ipatch_uz[1])),len((background[0])),len((background[1]))
		groups = {'Ab i-Patch':ipatch_uz[1],'Docking':docking_uz[1],'Background':background[1],'Frequency':frequency[1]}
		labels = ipatch_uz[0]
		barplot(groups,labels,title=cdr,results_images=results_images)

#Main#
#os.system('rm -rf /home/krawczyk/test2/Ralph/Results/VR1/CDR_results/')
#os.system('rm -rf /home/krawczyk/test2/Ralph/Results/VR4347/CDR_results/')
#os.system('rm -rf /home/krawczyk/test2/Ralph/Results/VR4380/CDR_results/')
#os.system('rm -rf /home/krawczyk/test2/Ralph/Results/VR4431/CDR_results/')
#os.system('rm -rf /home/krawczyk/test2/Ralph/Results/VR4498/CDR_results/')
#os.system('rm -rf /home/krawczyk/test2/Ralph/Results/VR6/CDR_results/')

#translate_results("/home/icruser/public_html/csb/PACE/web_front/USR/downloads/20140606100925_2132735254/output_folder/")
#translate_results("/home/krawczyk/test2/Ralph/Results/VR4347")
#translate_results("/home/krawczyk/test2/Ralph/Results/VR4380")
#translate_results("/home/krawczyk/test2/Ralph/Results/VR4431")
#translate_results("/home/krawczyk/test2/Ralph/Results/VR4498")
#translate_results("/home/krawczyk/test2/Ralph/Results/VR6")
	
#MAIN - demo - run: python barchart_demo.py
#groups = {'Ab i-Patch':(20, 35, 30, 35, 27),'Docking':(25, 32, 34, 20, 25)}
#labels = ('H26', 'H27', 'H28', 'H29', 'H30')
#barplot(groups,labels,title="CDR-H1")



