#!/usr/bin/env python
# -*- coding:UTF-8 -*-
'''
@Author: Li Fajin
@Date: 2019-08-28 15:07:53
@LastEditors: Li Fajin
@LastEditTime: 2019-08-30 16:34:02
@Description: This script is used for calculating mean ribosome density of each position for each transcript.
usage: python enrichmentMeanDensity.py -i rep1.txt,rep2.txt,rep3.txt -o outprefix
'''

import sys
import numpy as np
from functools import reduce
from collections import defaultdict
from optparse import OptionParser
from .__init__ import __version__

def create_parse_for_mean_enrichment():
	'''argument parser.'''
	usage="usage: %prog [options]" + '\n' + __doc__ + "\n"
	parser=OptionParser(usage=usage,version=__version__)
	parser.add_option("-i","--input",action="store",type="string",dest="inputFiles",help="Ribosome density files of different replicates separated by comma which could be generated by RiboDensityAtEachPosition.py. e.g. rep1.txt,rep2.txt")
	parser.add_option("-o","--output",action="store",type="string",dest="outprefix",help="Output prefix of the output file.")
	return parser
def get_density_dict(densityFile):
	density_dict={}
	with open(densityFile,'r') as f:
		i=0
		for line in f:
			i+=1
			if i==1:
				continue
			tmp=line.strip().split("\t")
			trans_id=tmp[0]
			density_dict[trans_id]=[float(i) for i in tmp[1:]]
	return density_dict

def flatten(xs):
	for x in xs:
		if isinstance(x,tuple):
			for xx in flatten(x):
				yield xx
		else:
			yield x

def meanDensity(inputFiles,outprefix):
	all_ratio_dict=defaultdict(list)
	for f in inputFiles:
		ratio_dict=get_density_dict(f)
		for trans,ratio in ratio_dict.items():
			all_ratio_dict[trans].append(ratio)
	with open(outprefix+"_mean_density.txt",'w') as fout:
		fout.write("%s\t%s\n" %("transcript","mean_density"))
		for trans,ratio in all_ratio_dict.items():
			mean_ratio=np.array(list(map(np.mean,[list(flatten(i)) for i in list(reduce(zip,all_ratio_dict[trans]))])))
			fout.write("%s\t" %(trans))
			for i in range(len(mean_ratio)):
				fout.write("%s\t" %str(mean_ratio[i]))
			fout.write("\n")

def main():
	parser=create_parse_for_mean_enrichment()
	(options,args)=parser.parse_args()
	if not options.inputFiles or not options.outprefix:
		raise IOError("Please set your -i and -o parameters!")
	inputFiles=options.inputFiles.strip().split(",")
	print("Start calculating mean density...",file=sys.stderr)
	meanDensity(inputFiles,options.outprefix)
	print("Finish the step of mean density calculation!",file=sys.stderr)

if __name__=="__main__":
	main()

