#!/usr/bin/python
__author__		= "Sander Granneman"
__copyright__	= "Copyright 2021"
__version__		= "0.6.5"
__credits__		= ["Sander Granneman"]
__email__		= "Sander.Granneman@ed.ac.uk"
__status__		= "beta"

from ruffus import *
from ruffus.cmdline import MESSAGE
from collections import defaultdict
import ruffus.cmdline as cmdline
import subprocess
import platform
import glob
import re
import os
import argparse
import sys
import time

DEFAULTADAPT = "NAGATCGGAAGAGCACACG"

parser = cmdline.get_argparse(description="CRAC pipeline for processing paired-end CRAC data")
parser.add_argument("-f",dest="forwardreads",help="the path to your fastq files with forward reads.",metavar="data_1_1.fastq data_2_1.fastq ...",nargs="*",default=None)
parser.add_argument("-r",dest="reversereads",help="the path to your fastq files with reverse reads.",metavar="data_1_2.fastq data_2_2.fastq ...",nargs="*",default=None)
parser.add_argument("-g","--gtf",dest="gtf",help="the path to your gtf annotation file",metavar="rRNA.gtf",default=None)
parser.add_argument("-c","--chromosome",dest="chromosome",help="the path to your chromosome length file",metavar="chromosome_lengths.txt",default=None)
parser.add_argument("-i","--ignoremutations",dest="ignoremuts",action="store_true",help="to tell the novoalign parser to ignore mutations. Useful when analyzing ChemModSeq data. Default is OFF",default=False)
parser.add_argument("--nocollapse",dest="nocollapse",action="store_true",help="skips the trimming and collapsing step of the analysis. Note that files should have a .fasta file extension! Default is OFF",default=False)
parser.add_argument("--novoindex",dest="novoindex",help="the path to your novoindex file",metavar="yeast.novoindex",default=None)
parser.add_argument("--name",dest="name",help="provide a single word describing the run. Default is 'analysis' with a time stamp",default="analysis_%s" % time.strftime("%d%m%Y"))
parser.add_argument("-b","--barcodes",dest="barcodes",help="the path to the file containing the list of barcodes. If you do not provide a barcode file, the demultiplexing step will be skipped",metavar="barcodes.txt",default=None)
parser.add_argument("--search",dest="search",choices=["forward","reverse"],help="to search for barcodes in forward OR reverse reads. Default is forward read data",default="forward")
parser.add_argument("-a","--adapter",dest="adapter",help="the path to the file containing the 3' adapter sequences for trimming the reads using flexbar. If you do not provide adapter sequences, the trimming step will be skipped",default=None)
parser.add_argument("-a2","--adapter2",dest="adapter2",help="use this option to trim any barcode/adapter sequences from the reverse read. Important when not using TrueSeq adapters and processing CRAC data! Default=None",default=None)
parser.add_argument("--truseq",dest="truseq",action="store_true",help="add this flag if your library was prepared using TruSeq kits. NOTE! This requires Flexbar version 3.4.0 or later! Default is False",default=False)
parser.add_argument("-m","--mismatches",dest="mismatches",type=int,help="indicate how many mismatches you allow for demultiplexing. Default is 1",default=1)
parser.add_argument("-p","--processors",dest="processors",type=int,help="indicate how many processors you want to use for analyses. Default is 8",default=8)
args = parser.parse_args()

if args.barcodes and not args.adapter:
	parser.error("Program assumes adapter trimming and demultiplexing happends at the same time!\nWhen using the --barcodes flag you also have to use --adapter flag.\n")

def getFileBaseName(filename):
	""" removes path and extension from file name """
	return os.path.splitext(os.path.basename(filename))[0]

def getBarcodeInfo(barcodefile):
	return ["%s.fastq" % "_".join(line.strip().split()) for line in open(barcodefile,"r").readlines()]

def runFlexBar(inputfiles,outputfiles):
	""" runs Flexbar on the data to remove the adapter sequence from the forward reads """
	if args.adapter and not args.adapter2:
		cmd = "flexbar -r '%s' -p '%s' -qf i1.8 -n 10 -ao 7 --adapters '%s' --output-reads '%s' --output-reads2 '%s' -qt 30 -ap ON" % (inputfiles[0],inputfiles[1],args.adapter,outputfiles[0],outputfiles[1])
	elif args.adapter and args.adapter2:
		cmd = "flexbar -r '%s' -p '%s' -qf i1.8 -n 10 -ao 7 --adapters '%s' --adapters2 '%s' --output-reads '%s' --output-reads2 '%s' -qt 30 -ap ON" % (inputfiles[0],inputfiles[1],args.adapter,args.adapter2,outputfiles[0],outputfiles[1])
	elif args.truseq:
		cmd = "flexbar -r '%s' -p '%s' -qf i1.8 -n 10 -ao 7 -aa TruSeq --output-reads '%s' --output-reads2 '%s' -qt 30 -ap ON" % (inputfiles[0],inputfiles[1],outputfiles[0],outputfiles[1])
	else:
		cmd = "flexbar -r '%s' -p '%s' -qf i1.8 -n 10 -ao 7 --output-reads '%s' --output-reads2 '%s' -qt 30 -ap ON" % (inputfiles[0],inputfiles[1],outputfiles[0],outputfiles[1])
	logger.info(cmd)
	os.system(cmd)

def trimReads(intputfiles,outputfile):
	""" Flexbar doesn't always do a good job trimming the reverse read so this is done here.
	If the forward reads after flexbar trimming is shorter than the reverse, the reverse read
	will be trimmed from the 3' end to the same length."""
	first = open(inputfiles[0],"r")
	second = open(inputfiles[1],"r")
	secondout = open(outputfile,"w")
	while True:
		try:
			firstheader = next(first).strip()
			firstseq = next(first).strip()
			firstplus = next(first).strip()
			firstquality = next(first).strip()
			secondheader = next(second).strip()
			secondseq = next(second).strip()
			secondplus = next(second).strip()
			secondquality = next(second).strip()
			lengthfirst = len(firstseq)
			lengthsecond = len(secondseq)
			if lengthfirst != lengthsecond:
				secondseq = secondseq[:lengthfirst-1]
				secondquality = secondquality[:lengthfirst-1]
			secondout.write("%s\n%s\n+\n%s\n" %(secondheader,secondseq,secondquality))
		except StopIteration:
			secondout.close()
			break

def demultiplexSamples(inputfiles,outputfiles):
	""" demultiplexes all the samples """
	os.chdir(os.path.join(root_dir,"demultiplexed"))
	cmd = "pyBarcodeFilter.py -f '%s' -r '%s' -b '%s' -m '%s' --search '%s' " % (inputfiles[0],inputfiles[1],os.path.join(home_dir,args.barcodes),args.mismatches,args.search)
	logger.info(cmd)
	os.system(cmd)
	os.chdir(root_dir)

def Collapse(inputfiles,outputfiles):
	""" Collapses the data and then splits it again into fasta files """
	cmd = "pyFastqJoiner.py -f '%s' '%s' -c '|' | pyFastqDuplicateRemover.py | pyFastqSplitter.py -c '|' --file_type fasta -o '%s' '%s'" % (inputfiles[0],inputfiles[1],outputfiles[0],outputfiles[1])
	logger.info(cmd)
	os.system(cmd)

def alignReads(inputfiles,outputfile):
	""" Runs novoalign on all the collapsed files"""
	cmd = "novoalign -d '%s' -f '%s' '%s' -r Random > '%s'" % (args.novoindex,inputfiles[0],inputfiles[1],outputfile)
	if args.truseq:
		cmd = "novoalign -d '%s' -f '%s' '%s' -r Random > '%s'" % (args.novoindex,inputfiles[1],inputfiles[0],outputfile)  # reverse the file order for novoalignment with TrueSeq data!!!
	logger.info(cmd)
	os.system(cmd)

def runPyReadCounters(inputfile,outputfile):
	""" runs pyReadCounters on all the files """
	outputfile = re.search("(.*)_count_output_reads.gtf",outputfile).group(1)
	if args.ignoremuts:
		string = "--mutations=nomuts"
	else:
		string = ""
	cmd = "pyReadCounters.py -f '%s' --gtf '%s' -v -o '%s' '%s'" % (inputfile,args.gtf,outputfile,string)
	logger.info(cmd)
	os.system(cmd)

def runSecondPyReadCounters(inputfile,outputfile):
	""" runs pyReadCounters on all the files """
	outputfile = re.search("(.*)_count_output_cDNAs.gtf",outputfile).group(1)
	cmd = "pyReadCounters.py -f '%s' --gtf '%s' -v -o '%s' --mutations	nomuts --blocks" % (inputfile,args.gtf,outputfile)
	logger.info(cmd)
	os.system(cmd)

def makeCoverageSgrFiles(inputfile,outputfiles):
	""" takes the pyReadCounters gtf output files and generates coverage sgr files """
	outputfile = re.search("(.*)_plus_strand_reads.sgr",outputfiles[0]).group(1)
	cmd = "pyGTF2sgr.py --gtf '%s' --zeros --count -v --permillion -o '%s' -c '%s'" % (inputfile,outputfile,args.chromosome)
	logger.info(cmd)
	os.system(cmd)

def makeCoverageBedgraphFiles(inputfile,outputfiles):
	""" takes the pyReadCounters gtf output files and generates bedgraph files for viewing of the data in genome browsers """
	outputfile = re.search("(.*)_plus_strand_reads.bedgraph",outputfiles[0]).group(1)
	cmd = "pyGTF2bedGraph.py --gtf '%s' --count -v --permillion -o '%s' -c '%s'" % (inputfile,outputfile,args.chromosome)
	logger.info(cmd)
	os.system(cmd)

### checking if dependencies are installed

dependencies = ["flexbar","pyGTF2sgr.py","pyGTF2bedGraph.py","pyReadCounters.py","pyBarcodeFilter.py","pyFastqJoiner.py","pyFastqSplitter.py","novoalign","pyFastqDuplicateRemover.py"]

for i in dependencies:
	cmd = "where" if platform.system() == "Windows" else "which"
	try:
		subprocess.call([cmd, "%s"])
	except:
		sys.stderr.write("%s is not installed (properly) on your system. Please (re)install" % i)
		exit()

### start of pipeline

# setting the main working directory
home_dir = os.getcwd()
root_dir = "%s/%s" % (home_dir,args.name)
try:
	os.mkdir(root_dir)
except OSError:
	pass

### setting up logging
if not args.log_file:
	logfile = "%s/%s" % (root_dir,"log_file.txt")
else:
	logfile = "%s/%s" % (root_dir,args.log_file)
logger, logger_mutex = cmdline.setup_logging ("CRAC_pipeline_PE",logfile,10)

### setting the starting files

assert len(args.forwardreads) == len(args.reversereads), "ERROR! number of forward and reverse input files is not the same!\n"

args.forwardreads = [os.path.abspath(i) for i in args.forwardreads]
args.reversereads = [os.path.abspath(i) for i in args.reversereads]
startingfiles = zip(args.forwardreads,args.reversereads)

### checking if the correct pyBarcodeFilter version is installed. Need version 3.0 or higher!

result = subprocess.Popen(['pyBarcodeFilter.py', '--version'], stdout=subprocess.PIPE)
out,err = result.communicate()
if float(out) < 3.0:
	sys.stderr.write("To run this script you need to have pyBarcodeFilter version 3.0 or later installed. Please install the latest version of pyCRAC\n")
	exit()

### start of pipeline

pipeline = Pipeline(name="Paired-end data CRAC pipeline")

if (args.adapter or args.truseq) and args.barcodes:
	pipeline.transform(task_func = runFlexBar,
						input  = startingfiles,
						filter = formatter("^.+/([^/]+)_1.(san)?fastq$","^.+/([^/]+)_2.(san)?fastq$"),
						output = ["%s/flexbar_trimmed/{1[0]}_trimmed_1.fastq" % root_dir,"%s/flexbar_trimmed/{1[1]}_trimmed_2.fastq" % root_dir],
					).follows(pipeline.mkdir(os.path.join(root_dir,"flexbar_trimmed")))

#	pipeline.transform(task_func = trimReads,
#						input  = runFlexBar,
#						filter = formatter("^.+/([^/]+)_trimmed_1.fastq","^.+/([^/]+)_trimmed_2.fastq"),
#						output = ["%s/flexbar_trimmed/{1[0]}_trimmed_1.fastq" % root_dir,"%s/flexbar_trimmed/{1[1]}_extra_trimmed_2.fastq" % root_dir],
#
#					)

	pipeline.subdivide(task_func = demultiplexSamples,
						input  = runFlexBar,
						filter = formatter("^.+/([^/]+)_1.fastq","^.+/([^/]+)_2.fastq"),
						output = [["%s/demultiplexed/{1[0]}_1_%s" % (root_dir,barcodestring),"%s/demultiplexed/{1[1]}_2_%s" % (root_dir,barcodestring)] for barcodestring in getBarcodeInfo(args.barcodes)]
					).follows(pipeline.mkdir(os.path.join(root_dir,"demultiplexed")))

	if not args.nocollapse:
		pipeline.transform(task_func = Collapse,
							input  = demultiplexSamples,
							filter = formatter("^.+/([^/]+)_trimmed_1(?P<BARCODE>.*).fastq","^.+/([^/]+)_trimmed_2(?P<BARCODE>.*).fastq"),
							output = ["%s/collapsed/{1[0]}_trimmed{BARCODE[0]}_1.fasta" % root_dir,"%s/collapsed/{1[1]}_trimmed{BARCODE[1]}_2.fasta" % root_dir],
						).follows(pipeline.mkdir(os.path.join(root_dir,"collapsed")))

		startingfiles = Collapse

	elif args.nocollapse:
		startingfiles = demultiplexSamples

elif (args.adapter or args.truseq) and	(not args.adapter2 and not args.barcodes):
	pipeline.transform(task_func = runFlexBar,
						input  = startingfiles,
						filter = formatter("^.+/([^/]+)_1.(san)?fastq$","^.+/([^/]+)_2.(san)?fastq$"),
						output = ["%s/flexbar_trimmed/{1[0]}_trimmed_1.fastq" % root_dir,"%s/flexbar_trimmed/{1[1]}_trimmed_2.fastq" % root_dir],
					).follows(pipeline.mkdir(os.path.join(root_dir,"flexbar_trimmed")))

	if not args.nocollapse:
		pipeline.transform(task_func = Collapse,
							input  = startingfiles,
							filter = formatter("^.+/([^/]+).fastq$","^.+/([^/]+).fastq$"),
							output = ["%s/collapsed/{1[0]}.fasta" % root_dir,"%s/collapsed/{1[1]}.fasta" % root_dir],
						).follows(pipeline.mkdir(os.path.join(root_dir,"collapsed")))
		startingfiles = Collapse
	elif args.nocollapse:
		startingfiles = runFlexBar

elif (args.adapter and args.adapter2) and not args.barcodes:
	pipeline.transform(task_func = runFlexBar,
						input  = startingfiles,
						filter = formatter("^.+/([^/]+)_1.(san)?fastq$","^.+/([^/]+)_2.(san)?fastq$"),
						output = ["%s/flexbar_trimmed/{1[0]}_trimmed_1.fastq" % root_dir,"%s/flexbar_trimmed/{1[1]}_trimmed_2.fastq" % root_dir],
					).follows(pipeline.mkdir(os.path.join(root_dir,"flexbar_trimmed")))

	if not args.nocollapse:
		pipeline.transform(task_func = Collapse,
							input  = startingfiles,
							filter = formatter("^.+/([^/]+).fastq$","^.+/([^/]+).fastq$"),
							output = ["%s/collapsed/{1[0]}.fasta" % root_dir,"%s/collapsed/{1[1]}.fasta" % root_dir],
						).follows(pipeline.mkdir(os.path.join(root_dir,"collapsed")))
		startingfiles = Collapse
	elif args.nocollapse:
		startingfiles = runFlexBar

elif not (args.adapter or args.truseq) and not args.barcodes:
	if not args.nocollapse:
		pipeline.transform(task_func = Collapse,
							input  = startingfiles,
							filter = formatter("^.+/([^/]+).fastq$","^.+/([^/]+).fastq$"),
							output = ["%s/collapsed/{1[0]}.fasta" % root_dir,"%s/collapsed/{1[1]}.fasta" % root_dir],
						).follows(pipeline.mkdir(os.path.join(root_dir,"collapsed")))

		startingfiles = Collapse


pipeline.transform(task_func = alignReads,
					input  = startingfiles,
					filter = formatter("^.+/([^/]+)(_1)?.fast[aq]","^.+/([^/]+)(_2)?.fast[aq]"),
					output = "%s/novo_files/{1[0]}.novo" % root_dir,
				).follows(pipeline.mkdir(os.path.join(root_dir,"novo_files")))

pipeline.transform(task_func = runPyReadCounters,
					input  = alignReads,
					filter = formatter("^.+/([^/]+).novo"),
					output = "%s/pyReadCounters_analyses/{1[0]}_count_output_reads.gtf" % root_dir,
				).follows(pipeline.mkdir(os.path.join(root_dir,"pyReadCounters_analyses")),alignReads)

pipeline.subdivide(task_func = makeCoverageBedgraphFiles,
					input  = runPyReadCounters,
					filter = formatter("^.+/([^/]+)_count_output_reads.gtf"),
					output = ["%s/bedgraph_files/{1[0]}_plus_strand_reads.bedgraph" % root_dir,"%s/bedgraph_files/{1[0]}_minus_strand_reads.bedgraph" % root_dir]
				).follows(pipeline.mkdir(os.path.join(root_dir,"bedgraph_files")),runPyReadCounters)

pipeline.transform(task_func = runSecondPyReadCounters,
					input  = alignReads,
					filter = formatter("^.+/([^/]+).novo"),
					output = "%s/pyReadCounters_blocks_nomuts_analyses/{1[0]}_blocks_nomuts_count_output_cDNAs.gtf" % root_dir,
				).follows(pipeline.mkdir(os.path.join(root_dir,"pyReadCounters_blocks_nomuts_analyses")),alignReads)

pipeline.subdivide(task_func = makeCoverageBedgraphFiles,
					input  = runSecondPyReadCounters,
					filter = formatter("^.+/([^/]+)_count_output_cDNAs.gtf"),
					output = ["%s/bedgraph_files_blocks_no_muts/{1[0]}_plus_strand_cDNAs.bedgraph" % root_dir,"%s/bedgraph_files/{1[0]}_minus_strand_cDNAs.bedgraph" % root_dir]
				).follows(pipeline.mkdir(os.path.join(root_dir,"bedgraph_files_blocks_no_muts")),runSecondPyReadCounters)


pipeline_run(multiprocess=args.processors,verbose=5)
