import sys
import random
import argparse

def parse_args():
    '''
    Gives options
    '''
    parser = argparse.ArgumentParser(description='Saves reads below a alignment threshold and discards all others')
    parser.add_argument('-k', help='Alignment number cutoff')
    parser.add_argument('--paired-end', dest='paired_ended', action='store_true', help='Data is paired-end')
    args = parser.parse_args()
    alignment_cutoff = int(args.k)
    paired_ended = args.paired_ended

    return alignment_cutoff, paired_ended


if __name__ == "__main__":
    '''
    Runs the filtering step of choosing multimapped reads
    '''

    [alignment_cutoff, paired_ended] = parse_args()

    if paired_ended:
        alignment_cutoff = int(alignment_cutoff) * 2

    # Store each line in sam file as a list of reads, 
    # where each read is a list of elements to easily 
    # modify or grab things
    current_reads = [] 
    current_qname = ''

    for line in sys.stdin:

        read_elems = line.strip().split('\t')

        if read_elems[0].startswith('@'):
            sys.stdout.write(line)
            continue

        # Keep taking lines that have the same qname
        if read_elems[0] == current_qname:
            # Add line to current reads
            current_reads.append(line)
            pass
        else:
            # Discard if there are more than the alignment cutoff
            if len(current_reads) >= alignment_cutoff:
                current_reads = [line]
                current_qname = read_elems[0]
            elif len(current_reads) > 0:
                # Just output all reads, which are then filtered with
                # samtools
                for read in current_reads:
                    sys.stdout.write(str(read))

                # And then discard
                current_reads = [line]
                current_qname = read_elems[0]
            else:
                # First read in file
                current_reads.append(line)
                current_qname = read_elems[0]

import os
import math
import sys
import argparse
import gzip
import numpy as np
from collections import namedtuple
import pysam
import re
import multiprocessing as mp
import time
from multiprocess_utils import listener,init_mp,close_mp
start_time = time.time()

def parseArgs() :
    # dir of source code
    srcpath=sys.argv[0]
    srcdir=os.path.dirname(os.path.abspath(srcpath))
    # parser
    parser = argparse.ArgumentParser(description='make wig file')
    parser.add_argument('-v','--verbose', action='store_true',default=False,
            help="verbose output")
    parser.add_argument('-i','--input',type=str,required=False,default="stdin",
            help="input file (methylation freq, default stdin)")
    parser.add_argument('-o','--output',type=argparse.FileType('w'),required=False, 
            default=sys.stdout,help="output file (default stout)")
    parser.add_argument('-t','--threshold',type=int,required=False,
            default=3,help="coverage thresdhold (default 3)")
    parser.add_argument('-c','--coverage',type=str,required=False,
            help="output path for coverage file (optional)")
    args = parser.parse_args()
    args.srcdir=srcdir
    return args

if __name__=="__main__":
    args=parseArgs()
    if args.coverage :
        covout = open(args.coverage,'w')
        print("track type=wiggle_0",file=covout)
        def printfxn(prechrom,chrom,pos,meth,cov,out,covout) :
            if chrom != prechrom :
                prechrom = fields[0]
                print("variableStep chrom={}".format(chrom),file=out)
                print("variableStep chrom={}".format(chrom),file=covout)
            print("{}\t{}".format(pos,meth/cov),file=out)
            print("{}\t{}".format(pos,cov),file=covout)
            return prechrom
    else :
        covout = None
        def printfxn(prechrom,chrom,pos,meth,cov,out,covout) :
            if chrom != prechrom :
                prechrom = fields[0]
                print("variableStep chrom={}".format(chrom),file=out)
            print("{}\t{}".format(pos,meth/cov),file=out)
            return prechrom
    if args.input == "stdin" :
        if args.verbose : print("reading from stdin",file=sys.stderr)
        in_fh = sys.stdin
    elif args.input.split(".")[-1] == "gz" :
        if args.verbose : 
            print("reading gzipped input {}".format(args.input),file=sys.stderr)
        in_fh = gzip.open(args.input,'rt')
    else :
        if args.verbose : print("reading {}".format(args.input),file=sys.stderr)
        in_fh = open(args.input)
    i = 0
    chrom=""
    # header
    print("track type=wiggle_0",file=args.output)
    for line in in_fh :
        fields = line.strip().split("\t")
        meth = int(fields[3])
        unmeth = int(fields[4])
        cov = meth+unmeth
        if cov < args.threshold : continue
        chrom = printfxn(chrom,fields[0],fields[1],meth,cov,args.output,covout)
        i += 1
        if args.verbose : 
            if i % 100000 == 0 :
                print("processed {} lines".format(i),file=sys.stderr)

    in_fh.close()

shell: 
	"trim_galore {input} "
	"-o {wildcards.dir}/fastq_trimmed &> {log}"

shell:
	"bowtie2-build {input} {params} --threads {threads} &> {log}"

shell:
	"picard -Xmx8G -Xms256M -XX:ParallelGCThreads={threads} "
	"-Djava.io.tmpdir={wildcards.dir}/tmp "
	"MarkDuplicates INPUT={input} OUTPUT={output} "
	"METRICS_FILE={wildcards.dir}/bam/{wildcards.sample}.dupmetrics.txt "
	"VALIDATION_STRINGENCY=LENIENT ASSUME_SORTED=true "
	"REMOVE_DUPLICATES=false &> {log}"

shell:
	"bedtools bamtobed -i {input} | "
	"awk 'BEGIN{{OFS=\"\t\"}}{{$4=\"N\";$5=\"1000\";print $0 }}' | "
	"sort -k1,1 -k2,2n -T {wildcards.dir}/tmp | "
	"bgzip > {output}"

shell:
	"macs2 callpeak -t {input} -f BED "
	"-n {wildcards.dir}/peaks/allsamples "
	"-g mm -p 0.01 --shift -75 --extsize 150 "
	"--nomodel -B --SPMR --keep-dup all --call-summits &> {log} && "
	"{params}/scripts/atacseq_bed_to_saf.sh "
	"{wildcards.dir}/peaks/allsamples_peaks.narrowPeak > {output}"

shell:
	"featureCounts --verbose -a {input.peaks} "
	"-o {output} -F SAF -T {threads} {input.reads} &> {log}"

shell:
	"ngmlr -x ont --bam-fix "
	"-t {threads} -r {params} -q {input} 2> {log} | "
	"samtools view -q 20 -b - | "
	"samtools sort -T bam/{wildcards.sample}.sorting "
	"-o {output} && "
	"samtools index {output}"

shell:
	"minimap2 -x map-ont -d {output} {input}"

shell:
	"nanopolish index --verbose -d {wildcards.pre} "

shell:
	"python {params}/nanopore-methylation-utilities/mtsv2bedGraph.py -g {input.fa} -i {input.meth} | "
	"sort -T tmp -k1,1 -k2,2n | "
	"bgzip > {output} && "
	"tabix -p bed {output}"

shell: 
	"wigToBigWig {input} {output}"

shell: 
	"sniffles -m {input} -v {output} -t {threads} "
	"--tmp_file sv/{wildcards.sample}.sniffles.tmp " 
	"-s 2  &> {log}"

shell:
	"bcftools sort -o {output} -T sv {input}"

shell:
	"echo {input} | tr \" \" \"\\n\" > {output.raw} && "
	"SURVIVOR merge {output.raw} 1000 1 1 -1 -1 -1 {output.vcf}"