#!/usr/bin/env python
# -*- coding:UTF-8 -*-
'''
Note:
1) -g: group names of your samples separated by comma. e.g. si-Ctrl,si-eIF3e. And the first one best to be the control group.
2) -r: replicate names of each group. groups separated by "__", and replicates separated by comma. e.g. si_ctrl_1_80S,si_ctrl_2_80S,si_ctrl_3_80S__si_3e_1_80S,si_3e_2_80S,si_3e_3_80S
'''


import pandas as pd
import numpy as np
import matplotlib
import sys
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.backends.backend_pdf import PdfPages
import seaborn as sns
from collections import defaultdict
from optparse import OptionParser
from .__init__ import __version__

def create_parser_for_density_on_each_kind_of_AA_codon():
	'''argument parser.'''
	usage="usage: python %prog [options]"
	parser=OptionParser(usage=usage,version=__version__)
	parser.add_option('-i','--input',action='store',type='string',dest='density_file',
		help='Input file in dataframe format (pd.DataFrame). Generated by RiboDensityAtEachKindAAOrCodon.py')
	parser.add_option('-o','--output',action='store',type='string',dest='output_prefix',
		help='Output files')
	parser.add_option('-g','--group',action="store",type="string",dest="group_name",
		help="Group name of each group separated by comma. e.g. 'si-control,si-eIF3e'")
	parser.add_option('-r','--replicate',action="store",type="string",dest="replicate_name",
		help="Replicate name of each group separated by comma. e.g. 'si_3e_1_80S,si_3e_2_80S__si_cttl_1_80S,si_ctrl_2_80S'")
	parser.add_option("--ymax",action="store",type="float",dest="ymax",default=None,
		help="The max of ylim. default=%default")
	parser.add_option("--ymin",action="store",type="float",dest="ymin",default=None,
		help="The min of ylim. default=%default")
	parser.add_option("--level",action="store",type='string',dest='level',default="codon",
		help="Control the plot in codon level or AA level. default=%default")
	return parser

def calculate_mean_density(data,groups,replicates,output_prefix):
	'''calculate the mean values among different replicates. the input data have N+2 columns, N is the number of samples
	1) col 1: AAs
	2) col2: codons
	3) col3-colN: different samples
	4) suppose we just have two groups (if you have more than two groups, split it up and analysis every two groups.[one is control and the other is treated])
	'''
	label_dict={}
	data_dict={}
	data_mean_dict=defaultdict(dict)
	AAs=data.iloc[:,0]
	codons=data.iloc[:,1]
	for g,r in zip(groups,replicates):
		label_dict[g]=r.strip().split(',')
	## separate the data into different groups
	for g in groups:
		data_dict[g]=data.loc[:,label_dict[g]]

	for g in groups:
		for i in np.arange(data.shape[0]):
			density=np.mean(data_dict[g].loc[i,label_dict[g]])
			data_mean_dict[g][i]=density
	## transform the dict to a python dataframe
	for g in groups:
		data_mean_dict[g]=pd.DataFrame(data_mean_dict[g],index=[g]).T
	## concatenate different data frame
	data_mean=pd.concat([v for v in data_mean_dict.values()],axis=1)
	data_mean=pd.concat((AAs,codons,data_mean),axis=1)
	## write the mean density file
	data_mean.to_csv(output_prefix+"_mean_codon_density_dataframe.txt",sep="\t",index=0)
	return data_mean

def reshape_mean_data(data_mean,groups):
	''' this requires groups only have two conditions: the first one is control group and the second one is the treat group'''
	Non_polar_AA=['A','P','G','V','W','I','F','M','L']
	Polar_AA=['T','Y','Q','C','N','S']
	Positive_AA=['K','H','R']
	Negative_AA=['D','E']
	Non_polar_AA_color=['gray']*len(Non_polar_AA)
	Polar_AA_color=["purple"]*len(Polar_AA)
	Negative_AA_color=['blue']*len(Negative_AA)
	Positive_AA_color=['red']*len(Positive_AA)
	my_color=Non_polar_AA_color+Polar_AA_color+Negative_AA_color+Positive_AA_color
	data_mean_sorted=data_mean.sort_values('AA',inplace=False)
	ctrl_mean=data_mean_sorted[str(groups[0])].groupby(data_mean_sorted['AA']).sum()
	treat_mean=data_mean_sorted[str(groups[1])].groupby(data_mean_sorted['AA']).sum()
	data_mean_AA=pd.concat((ctrl_mean,treat_mean),axis=1)
	data_mean_AA['ratio']=(data_mean_AA.iloc[:,1]-data_mean_AA.iloc[:,0])/data_mean_AA.iloc[:,0]

	Non_polar_AA_data=data_mean_AA.loc[Non_polar_AA,:]
	Non_polar_AA_data=Non_polar_AA_data.sort_values(by="ratio",ascending=False)
	Polar_AA_data=data_mean_AA.loc[Polar_AA,:]
	Polar_AA_data=Polar_AA_data.sort_values(by="ratio",ascending=False)
	Negative_AA_data=data_mean_AA.loc[Negative_AA,:]
	Negative_AA_data=Negative_AA_data.sort_values(by="ratio",ascending=False)
	Positive_AA_data=data_mean_AA.loc[Positive_AA,:]
	Positive_AA_data=Positive_AA_data.sort_values(by="ratio",ascending=False)
	data_mean_AA_reshaped=pd.concat((Non_polar_AA_data,Polar_AA_data,Negative_AA_data,Positive_AA_data),axis=0)
	return data_mean_AA_reshaped,my_color

def plot_density_for_each_sample_codon_level(data,data_mean,groups,replicates,output_prefix,ymin,ymax):
	''' this requires groups only have two conditions: the first one is control group and the second one is the treat group'''
	plt.rc('font',weight='bold')
	data_sorted=data.sort_values('AA',inplace=False)
	data_mean_sorted=data_mean.sort_values('AA',inplace=False)
	data_mean_sorted['ratio']=(data_mean_sorted.iloc[:,3]-data_mean_sorted.iloc[:,2])/data_mean_sorted.iloc[:,2]
	codons_sort=data_sorted.iloc[:,1].values
	codons_mean_sort=data_mean_sorted.iloc[:,1].values
	replicates=','.join(replicates).split(',')
	with PdfPages(output_prefix + "_density_on_each_kind_of_codon.pdf") as pdf:
		for r in replicates:
			fig=plt.figure(figsize=(16,8))
			ax=fig.add_subplot(111)
			sns.barplot(x=np.arange(len(codons_sort)),y=data_sorted.loc[:,r],palette='husl',ax=ax)
			ax.set_xlabel("codons",fontdict={'size':20,'weight':"bold","family":"Arial"})
			ax.set_ylabel("Relative ribosome density",fontdict={'size':20,'weight':"bold","family":"Arial"})
			ax.spines['top'].set_visible(False)
			ax.spines['right'].set_visible(False)
			ax.spines['bottom'].set_linewidth(2)
			ax.spines['left'].set_linewidth(2)
			if not ymin and not ymax:
				pass
			elif not ymin and ymax:
				ax.set_ylim(0,ymax)
			else:
				pass
			ax.tick_params(which="both",width=2)
			plt.xticks(np.arange(len(codons_sort)),codons_sort,rotation=270,fontsize=12,fontweight="bold")
			for a,b,c in zip(np.arange(len(codons_sort)),data_sorted.loc[:,r],data_sorted['AA']):
				if b>0:
					plt.text(a,b,c,ha='center',va='bottom',fontdict={'size':10,'weight':"bold","family":"Arial"})
				else:
					plt.text(a,float(b)-0.012,c,ha='center',va='bottom',fontdict={'size':10,'weight':"bold","family":"Arial"})
			plt.title(r,fontdict={'size':20,'weight':"bold","family":"Arial"})
			# plt.tight_layout()
			pdf.savefig(fig)
			plt.close()
		for r in data_mean_sorted.columns[2:]:
			fig=plt.figure(figsize=(16,8))
			ax=fig.add_subplot(111)
			sns.barplot(x=np.arange(len(codons_mean_sort)),y=data_mean_sorted.loc[:,r],palette='husl',ax=ax)
			ax.set_xlabel("codons",fontdict={'size':20,'weight':"bold","family":"Arial"})
			min_ylim=min(plt.ylim())
			if r == 'ratio':
				ax.axhline(0,color='k',clip_on=False,linewidth=1)
				ax.set_ylabel("Difference of ribosome density.\n"+ "("+str(groups[1])+"-"+str(groups[0])+")"+"/"+str(groups[0]),fontdict={'size':20,'weight':"bold","family":"Arial"})
				if not ymin and not ymax:
					pass
				elif not ymin and ymax:
					ax.set_ylim(0,ymax)
				elif ymin and not ymax:
					raise IOError("Please offer the ymax parameter as well!")
				elif ymin and ymax:
					ax.set_ylim(ymin,ymax)
				else:
					raise IOError("Please enter correct ymin and ymax parameters!")
				for a,b,c in zip(np.arange(len(codons_mean_sort)),data_mean_sorted.loc[:,r],data_mean_sorted['AA']):
					if b>0:
						plt.text(a,b,c,ha='center',va='bottom',fontdict={'size':10,'weight':"bold","family":"Arial"})
					else:
						plt.text(a,float(b)-0.012,c,ha='center',va='bottom',fontdict={'size':10,'weight':"bold","family":"Arial"})
					plt.vlines(a,ymin=min_ylim,ymax=0,color='gray',linestyles='dashed',alpha=0.5)
			else:
				ax.set_ylabel("Relative ribosome density",fontdict={'size':20,'weight':"bold","family":"Arial"})
				if not ymin and not ymax:
					pass
				elif not ymin and ymax:
					ax.set_ylim(0,ymax)
				else:
					pass
				for a,b,c in zip(np.arange(len(codons_mean_sort)),data_mean_sorted.loc[:,r],data_mean_sorted['AA']):
					if b>0:
						plt.text(a,b,c,ha='center',va='bottom',fontdict={'size':10,'weight':"bold","family":"Arial"})
					else:
						plt.text(a,float(b)-0.012,c,ha='center',va='bottom',fontdict={'size':10,'weight':"bold","family":"Arial"})
			ax.spines['top'].set_visible(False)
			ax.spines['right'].set_visible(False)
			ax.spines['bottom'].set_linewidth(2)
			ax.spines['left'].set_linewidth(2)
			ax.tick_params(which="both",width=2)
			plt.xticks(np.arange(len(codons_mean_sort)),codons_mean_sort,rotation=270,fontsize=12,fontweight="bold")
			plt.title(r,fontdict={'size':20,'weight':"bold","family":"Arial"})
			pdf.savefig(fig)
			plt.close()

def plot_density_for_each_sample_AA_level(data,data_mean_reshaped,groups,replicates,colors,output_prefix,ymin,ymax):
	plt.rc('font',weight='bold')
	data_sorted=data.sort_values('AA',inplace=False)
	data_sorted=data_sorted.groupby(data_sorted['AA']).sum()
	replicates=','.join(replicates).split(',')
	AAs=data_sorted.index.values
	AAs_index=data_mean_reshaped.index
	with PdfPages(output_prefix + "_density_on_each_kind_of_AA.pdf") as pdf:
		for r in replicates:
			fig=plt.figure(figsize=(16,8))
			ax=fig.add_subplot(111)
			sns.barplot(x=np.arange(len(AAs)),y=data_sorted.loc[:,r],palette='husl',ax=ax,alpha=0.9)
			ax.set_xlabel("AAs",fontdict={'size':20,'weight':"bold","family":"Arial"})
			ax.set_ylabel("Relative ribosome density",fontdict={'size':20,'weight':"bold","family":"Arial"})
			ax.spines['top'].set_visible(False)
			ax.spines['right'].set_visible(False)
			ax.spines['bottom'].set_linewidth(2)
			ax.spines['left'].set_linewidth(2)
			if not ymin and not ymax:
				pass
			elif not ymin and ymax:
				ax.set_ylim(0,ymax)
			else:
				pass
			ax.tick_params(which="both",width=2)
			plt.xticks(np.arange(len(AAs)),AAs,fontsize=12,fontweight='bold')
			plt.title(r,fontdict={'size':20,'weight':"bold","family":"Arial"})
			# plt.tight_layout()
			pdf.savefig(fig)
			plt.close()
		for r in data_mean_reshaped.columns:
			fig=plt.figure(figsize=(16,8))
			ax=fig.add_subplot(111)
			labels_legend=['non polar','polar','positively charged','negatively charged']
			colors_legend=['gray','purple','red','blue']
			patches = [ mpatches.Patch(color=colors_legend[i], label=labels_legend[i] ) for i in range(len(colors_legend)) ]
			sns.barplot(x=np.arange(len(AAs_index)),y=data_mean_reshaped.loc[:,r],palette=colors,ax=ax,alpha=0.9)
			ax.set_xlabel("AAs",fontdict={'size':20,'weight':"bold","family":"Arial"})
			if r == 'ratio':
				ax.axhline(0,color='k',clip_on=False,linewidth=1)
				ax.set_ylabel("Difference of ribosome density.\n"+ "("+str(groups[1])+"-"+str(groups[0])+")"+"/"+str(groups[0]),fontdict={'size':20,'weight':"bold","family":"Arial"})
				if not ymin and not ymax:
					pass
				elif not ymin and ymax:
					ax.set_ylim(0,ymax)
				elif ymin and not ymax:
					raise IOError("Please offer the ymax parameter as well!")
				elif ymin and ymax:
					ax.set_ylim(ymin,ymax)
				else:
					raise IOError("Please enter correct ymin and ymax parameters!")
			else:
				ax.set_ylabel("Relative ribosome density",fontdict={'size':20,'weight':"bold","family":"Arial"})
				if not ymin and not ymax:
					pass
				elif not ymin and ymax:
					ax.set_ylim(0,ymax)
				else:
					pass
			ax.spines['top'].set_visible(False)
			ax.spines['right'].set_visible(False)
			ax.spines['bottom'].set_linewidth(2)
			ax.spines['left'].set_linewidth(2)
			ax.tick_params(which="both",width=2)
			plt.xticks(np.arange(len(AAs_index)),AAs_index,fontsize=12,fontweight="bold")
			plt.title(r,fontdict={'size':20,'weight':"bold","family":"Arial"})
			plt.legend(handles=patches,loc='best',prop={'size':20,'weight':"bold","family":"Arial"})
			# plt.tight_layout()
			pdf.savefig(fig)
			plt.close()


def main():
	parser=create_parser_for_density_on_each_kind_of_AA_codon()
	(options,args)=parser.parse_args()
	(data,output_prefix,groups,replicates,level,ymin,ymax)=(options.density_file,options.output_prefix,options.group_name.strip().split(','),
	options.replicate_name.strip().split('__'),options.level,options.ymin,options.ymax)
	##  input data
	data=pd.read_csv(data,sep='\t')
	data_mean=calculate_mean_density(data,groups,replicates,output_prefix)
	data_mean_reshaped,colors_for_AA=reshape_mean_data(data_mean,groups)
	if level == 'codon':
		print("Start the step of plotting...",file=sys.stderr)
		plot_density_for_each_sample_codon_level(data,data_mean,groups,replicates,output_prefix,ymin,ymax)
		print('finishing the plot step!',file=sys.stderr)
	elif level == 'AA':
		print("Start the step of plotting...",file=sys.stderr)
		plot_density_for_each_sample_AA_level(data,data_mean_reshaped,groups,replicates,colors_for_AA,output_prefix,ymin,ymax)
		print('finishing the plot step!',file=sys.stderr)
	else:
		raise IOError("Please reset your --level parameter.[AA or codon]")


if __name__=="__main__":
	main()