##################################################################################################
##################################################################################################
##################################################################################################
# presplit_map.py: written by Aaron Lun
# written 7 February, 2014

import argparse
class myformat(argparse.RawDescriptionHelpFormatter, argparse.ArgumentDefaultsHelpFormatter):
    pass

parser = argparse.ArgumentParser(description="""Align chimeric reads from a Hi-C experiment using a presplitting approach:
 a) Reads are split into segments at a known 'ligation signature' using cutadapt
 b) Segments are aligned with Bowtie2 in end-to-end mode with very high sensitivity
 c) All results for each pair of libraries are stitched into a single BAM file
""",

epilog="""Notes:
 - No processing is done to remove non-uniquely mapped reads. Filtering on MAPQ
   or other associated tags (e.g., XS) is recommended before downstream analyses.
 - The segment corresponding to the 5' end of each read will always be considered the
   primary alignment (i.e., all other segments are marked as secondary alignments).
""", formatter_class=myformat)

parser.add_argument("-G", type=str, dest="genome", help="path to bowtie2 genome index, prefix only", required=True)
parser.add_argument("-1", type=str, dest="fq1", help="path to FASTQ file containing the first read of the pair", required=True)
parser.add_argument("-2", type=str, dest="fq2", help="path to FASTQ file containing the second read of the pair", required=True)
parser.add_argument("-P", type=int, dest="phred", help="phred offset for the quality scores in the FASTQ files", choices=[33, 64], default=33)
parser.add_argument("--sig", type=str, dest="ligation", help="sequence of ligation signature (i.e. after blunt-end ligation of filled-in overhangs)", required=True)
parser.add_argument("-o", type=str, dest="output", help="path to the output BAM file", required=True)
parser.add_argument("--cmd", type=str, dest="cmd", help="path to/command for bowtie2 executable, with or without flags", default="bowtie2")
parser.add_argument("--cut", type=str, dest="cut", help="path to/command for cutadapt executable, with or without flags", default="cutadapt")
args = parser.parse_args()

## Setting up imports.
import pysam
from subprocess import Popen, PIPE
import os

from Bio import SeqIO
import gzip

import tempfile
tmpdir=tempfile.mkdtemp(dir=".")

## Additional processing of arguments before proceeding.
ligseq=args.ligation
liglen=len(ligseq)
if liglen%2!=0:
    raise ValueError("ligation signature must be even in length")

## Assembling the bowtie2 command.
import re
bwtcmd=re.split("\s+", args.cmd)+["--phred"+str(args.phred), "-x", args.genome]
cutcmd=re.split("\s+", args.cut)

####################################################################################################

def add_hard_clip(xread, is5, altlen):
    """A function that decides where to add the hard clipping (passes by reference)."""
    if not xread.is_unmapped:
        xcigar=list(xread.cigar)

        if xread.is_reverse:
            addfront=is5
        else:
            addfront=not is5

        if addfront:
            xcigar.insert(0, (5, altlen))
        else:
            xcigar.append((5, altlen))

        xread.cigar=tuple(xcigar)

    return

is_gz=re.compile("\\.gz$")
def open_handle(fname, write=False):
    """A function that returns a Gzip file handle or a normal file handle, depending on the file name."""
    if (is_gz.search(fname)):
        if write:
            mode='wt'
        else:
            mode='rt'

        return gzip.open(fname, mode)
    else:
        if write:
            mode='w'
        else:
            mode='rU'

        return open(fname, mode)

####################################################################################################

allnames=[]
dumpf=open(os.devnull, "w")
for x, curf in enumerate([args.fq1, args.fq2]):

    # Setting up custom names for the output file.
    fname, fext = os.path.splitext(os.path.basename(curf))
    is_first=(x==0)
    outbam=os.path.join(tmpdir, fname+".bam")
    allnames.append(outbam)

    ####################################################
    ##################### SPLITTING ####################
    ####################################################

    # Trimming the FASTQ file prior to input.
    init_split=os.path.join(tmpdir, "splitted_temp.fastq")
    cut_proc=Popen(cutcmd+["-o", init_split, "-a", ligseq, curf], stdout=dumpf, stderr=PIPE)
    cut_out, cut_err=cut_proc.communicate()
    if cut_proc.returncode:
        raise SystemError("cutadapt failed to trim reads\n"+cut_err)

    # Pulling out file handles.
    tempUs=os.path.join(tmpdir, "unsplit.fastq")
    temp5=os.path.join(tmpdir, "5.fastq")
    temp3=os.path.join(tmpdir, "3.fastq")

    with open_handle(init_split) as splitted, \
            open_handle(curf) as original_0, \
            open_handle(tempUs, True) as outputUs, \
            open_handle(temp5, True) as output5, \
            open_handle(temp3, True) as output3:

        # Runs through the temporary and the original, and adds back half the ligation signature length.
        # Also records the 3' sequence for realignment. If it wasn't split, then we record that
        # in a separate file as well. We also keep a log telling us what order everything was in before
        # we started cutting things up.
        counter=1
        original=SeqIO.parse(original_0, "fastq")

        for record in SeqIO.parse(splitted, "fastq"):
            orecord=next(original)

            junction=len(record)+liglen//2
            if junction < len(orecord):
                SeqIO.write(orecord[:junction], output5, "fastq")
                SeqIO.write(orecord[junction:], output3, "fastq")
            else:
                SeqIO.write(orecord, outputUs, "fastq")

        try:
            next(original)
            raise IOError("mismatching number of lines in cutadapt output and original file")
        except StopIteration:
            pass

    os.remove(init_split)

    ####################################################
    ###################### MAPPING #####################
    ####################################################

    # Mapping both ends, separately. Need reordering to stitch split reads back together.
    mapped5=os.path.join(tmpdir, "5.sam")
    mapped3=os.path.join(tmpdir, "3.sam")
    for dex in range(2):
        if not dex:
            curfile=temp5
            curout=mapped5
        else:
            curfile=temp3
            curout=mapped3

        map_proc=Popen(bwtcmd+["--reorder", "--very-sensitive", "-U", curfile, "-S", curout], stdout=dumpf, stderr=PIPE)
        map_out, map_err=map_proc.communicate()
        if map_proc.returncode:
            raise SystemError("bowtie2 failed for presplit alignment\n"+map_err)

        os.remove(curfile)

    # Mapping any unsplit reads with local alignment.
    mappedUs=os.path.join(tmpdir, "us.sam")
    map_proc=Popen(bwtcmd+["--local", "--very-sensitive-local", "-U", tempUs, "-S", mappedUs], stdout=dumpf, stderr=PIPE)
    if map_proc.wait():
        map_out, map_err=map_proc.communicate()
        raise SystemError("bowtie2 failed for unsplit alignment\n"+map_err)

    os.remove(tempUs)

    ####################################################
    ################## RE-INTEGRATION ##################
    ####################################################

    # Running through and re-integerating split and unsplit reads into a single file.
    with pysam.Samfile(mapped3, "r") as sin3, \
            pysam.Samfile(mapped5, "r") as sin5, \
            pysam.Samfile(mappedUs, "r") as sinU, \
            pysam.Samfile(outbam, "wb", template=sin3) as sout:

        for read3 in sin3:
            try:
                read5=next(sin5)
            except StopIteration:
                raise IOError("ran out of split reads to add")
    
            if read5.qname!=read3.qname:
                raise ValueError("names of split segments do not match")
    
            len5=len(read5.seq)
            len3=len(read3.seq)
    
            # We need to set appropriate flags for first/second read of a pair,
            # and for primary/secondary alignments. 5' reads are marked as primaries.
            read3.is_paired=read5.is_paired=True
            read3.is_read1=read5.is_read1=is_first
            read5.is_read2=read5.is_read2=not is_first
    
            add_hard_clip(read5, True, len3)
            sout.write(read5)
            read3.is_supplementary=True
            add_hard_clip(read3, False, len5)
            sout.write(read3)
   
        try:
            next(sin5)
            raise IOError("mismatching number of alignments in 3' and 5' SAM files")
        except StopIteration:
            pass
 
        # Merging the unsplit reads.
        for okread in sinU:
            okread.is_paired=True
            okread.is_read1=is_first
            okread.is_read2=not is_first
            sout.write(okread)

    os.remove(mapped5)
    os.remove(mapped3)
    os.remove(mappedUs)

    # Sorting the resulting file. We do it this way because the original
    # file wasn't guaranteed to be sorted. To try to merge the three files
    # (split, 5' and 3'), we'd need to assume some sorting order and I'm
    # not willing to do that. Sorting afterward enforces the 'samtools'
    # name ordering over anything that might have been there originally.
    bsorted=os.path.join(tmpdir, "sorted.bam")
    pysam.sort("-o", bsorted, "-n", outbam)
    os.rename(bsorted, outbam)

####################################################################################################
# Stitching two files together to reform a single BAM file.

cmd=["-n", "-f", args.output]+allnames
pysam.merge(*cmd)
for x in allnames:
    os.remove(x)

# Mopping up.
import shutil
shutil.rmtree(tmpdir)
dumpf.close()

####################################################################################################


