#!/usr/bin/env python
# -*- coding:UTF-8 -*-
## lifj
## 2018.10.11

'''
process the files generated by calculate_pausing_score.py. There only exists two groups,
the first one is the control groups and the secondary group is the treated groups. e.g -g si-ctrl,si-eIF3e
1) load those files
2) merge those files
3) calculate the mean pausing scores from different replicates
4) output a position weight matrix used for motif logo
input: directory where the puasing score file exist which generated by calculate_pausing_score.py
.....: a data mean score file requires:
1) index: motifs
2) col0: pausing score of control groups
3) col1: pausing score of treat groups
4) col2: ratio values using col1 divided by col0
output:
position weight matrix
'''

import pandas as pd
import numpy as np
import os
import sys
from functools import reduce
from itertools import groupby
from collections import defaultdict
from optparse import OptionParser
from .__init__ import __version__

def create_parse_for_pausing_score_process():
	'''Arguments parse'''
	usage="usage: python %prog [options]" +"\n" + __doc__+"\n"
	parser=OptionParser(usage=usage,version=__version__)
	parser.add_option("-i","--input",action="store",type="string",dest="score_files",help="The input files are pausing scores generated by PausingScore.py. All files are separated by comma. \nOr a data_mean_score with three columns: col0--control pausing score;col1--treated pausing score;col2--ratio. and its index must be motifs")
	parser.add_option("-o","--output",action="store",type="string",dest="output_prefix",help="Prefix of output files.[required]")
	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('--pausing_score_filter',action="store",type="float",dest="pausing_score_filter",default=10,help="The filter criteria for pausing score on treat samples.default=%default")
	parser.add_option('--ratio_filter',action="store",type='float',dest="ratio_filter",default=2,help="The filter criteria for ratio values. default=%default")
	parser.add_option('--mode',action="store",type='string',dest="input_mode",default='raw',help="The input is the raw density generated by PausingScore.py or a file processed by ProcessPusingScore.py with three columns---[ctrl0 treat1 ratio2]. You can chose raw or ratio. default=%default")
	return parser


def load_pausing_scores(inputFiles,groups,replicates,output_prefix):
	'''according to the files generated by calculate_pausing_score.py, caclulate the mean pausing score of different replicates'''
	label_dict={}
	data_dict=defaultdict(list)
	data_mean_dict=defaultdict()
	replicates_list=','.join(replicates).strip().split(',')
	replicates_files={r:f for r,f in zip(replicates_list,inputFiles)}
	for g,r in zip(groups,replicates):
		label_dict[g]=r.strip().split(',')
	for g in groups:
		for r in label_dict[g]:
			tmp=pd.read_csv(replicates_files[r],sep="\t")
			motifs=tmp.iloc[:,0]
			data_dict[g].append(tmp.iloc[:,1])
	for g in groups:
		data_mean_dict[g]=pd.DataFrame(reduce(lambda x,y: x+y,data_dict[g])/len(data_dict[g]),columns=[g])

	data_mean=pd.concat(data_mean_dict.values(),axis=1)
	data_mean=pd.concat((motifs,data_mean),axis=1)
	data_mean.to_csv(output_prefix+"_mean_pausing_score.txt",sep="\t",index=0)
	return data_mean

def process_mean_data(data_mean,output_prefix):
	data_mean_score=data_mean.groupby('motif').sum()
	data_mean_score=data_mean_score.sort_values(by=data_mean_score.columns[1],ascending=False)
	## filtering zeros
	data_mean_score=data_mean_score[data_mean_score>0]
	data_mean_score=data_mean_score.dropna()
	## filtering motifs whose length less than 5 (EPA + 2 spaces)
	motifs=[motif for motif in data_mean_score.index  if len(motif)==3]
	data_mean_score=data_mean_score.loc[motifs,:]
	## calculate the weight
	data_mean_score['ratio']=data_mean_score.iloc[:,1]/data_mean_score.iloc[:,0]
	data_mean_score.to_csv(output_prefix+"_mean_pausing_score_filtred.txt",sep="\t")

	return data_mean_score



def calculate_pwm(data_mean_score,output_prefix,ratio_filter,pausing_score_filter):
	'''according to the mean pausing score and the motifs, calculate the position weight martix'''

	data_mean_score_filtered=data_mean_score[(data_mean_score['ratio']>=ratio_filter) & (data_mean_score.iloc[:,1]>=pausing_score_filter)]
	if data_mean_score_filtered.shape[0] <1:
		raise IOError("There is no AAA motifs left for too strict filter parameters! Please reset your filter parameters!")
	data_mean_score_filtered.to_csv(output_prefix+"_motifs_used_for_pwm.txt",sep="\t")

	## PWM
	PWM=pd.DataFrame()
	PWM['motif']=data_mean_score_filtered.index
	PWM['weight']=data_mean_score_filtered['ratio'].values
	# PWM['E'],PWM['P'],PWM['A']=PWM['motif'].str.split(' ').str
	PWM['E'],PWM['P'],PWM['A']=PWM['motif'].str
	pwm_E=PWM[['E','weight']].groupby('E').sum()
	pwm_P=PWM[['P','weight']].groupby('P').sum()
	pwm_A=PWM[['A','weight']].groupby('A').sum()
	pwm=pd.merge(pwm_E,pwm_P,how='left',left_index=True,right_index=True)
	pwm=pd.merge(pwm,pwm_A,how='left',left_index=True,right_index=True)
	pwm.columns=['E','P','A']
	pwm=pwm.fillna(0)
	pwm['E']=pwm['E']/sum(pwm['E'])
	pwm['P']=pwm['P']/sum(pwm['P'])
	pwm['A']=pwm['A']/sum(pwm['A'])
	pwm=pwm.T
	pwm=pwm.reset_index(drop=True)
	pwm.to_csv(output_prefix+"_pwm.txt",sep="\t")


def parse_args_for_pwm_calculation():
	'''Arguments parse'''
	parser=create_parse_for_pausing_score_process()
	(options,args)=parser.parse_args()
	(inputFiles,output_prefix,groups,replicates,pausing_score_filter,ratio_filter)=(options.score_files.strip().split(","),options.output_prefix,options.group_name.strip().split(','),options.replicate_name.strip().split('__'),options.pausing_score_filter,options.ratio_filter)
	replicates_list=','.join(replicates).strip().split(',')
	if options.input_mode == 'raw':
		if len(inputFiles)!=len(replicates_list):
			raise IOError("The number of input files is not consistent with numbers of replicates, please check it!")
		print("Start calculating the mean pausing score.",file=sys.stderr)
		data_mean=load_pausing_scores(inputFiles,groups,replicates,output_prefix)
		data_mean_score=process_mean_data(data_mean,output_prefix)
		print("Start calculating the position weight matrix...",file=sys.stderr)
		calculate_pwm(data_mean_score,output_prefix,ratio_filter,pausing_score_filter)
		print("Finish the step of pwm calculation",file=sys.stderr)
	elif options.input_mode == 'ratio':
		data_mean_score=pd.read_csv(inputFiles[0],sep="\t",index_col=0)
		print("Start calculating the position weight matrix...",file=sys.stderr)
		calculate_pwm(data_mean_score,output_prefix,ratio_filter,pausing_score_filter)
		print("Finish the step of pwm calculation",file=sys.stderr)
	else:
		raise IOError("Please input a correct --mode parameter.[raw or ratio]")

def main():
	parse_args_for_pwm_calculation()

if __name__ == "__main__":
	main()




