import sys
import os
import argparse
import pysam

fastaf = snakemake.input[0]
out = open(snakemake.output[0], 'w')

#knobs
smooth_gap_width = 150000
contig_edge_margin = 150000
min_smoothed_aln_len = 10000
min_aln_len = 5000

#align
print("Aligning...".format(t=str(snakemake.threads)))
os.system("nucmer -p {delta} -b 4000 -l 2000 --maxmatch {fa} {fa}".format( #-t {threads} back to mummer3, no more multithreading : (
	threads = str(snakemake.threads),
	delta = snakemake.params['delta'],
	fa = fastaf
))
os.system("show-coords -T {delta}.delta -L 2000 | sed '1,5d' > coords.tsv".format(
	delta = snakemake.params['delta']
)) #the 1,5d gets rid of the header and identity hit as well
print("Trimming genome...")

max_tiglen = 0

with open('coords.tsv') as coords:
	lines = coords.readlines()
	smoothed_lines = []
	if len(lines) > 0:
		aln_start_line = lines[0].strip().split("\t")
		prev_line = lines[0].strip().split("\t")

		#goal: identify corner-cutting parallel off-diagonals.
	#	print(prev_line)

		for l in lines[1:] + [lines[0]]: #just so we don't miss the last item
			s = l.strip().split("\t")
			#print(s)

			tigname = s[-1] #store the name of the tig.  this repeats nugatorily.

			if int(s[1]) > max_tiglen: #keep track of the largest alignment coordinate found, used as the tig length
				max_tiglen = int(s[1])

			if int(s[0]) > int(s[1]): #ignore inversions
	#				print('ignoring')
				continue

			if int(s[1]) - int(s[0]) < min_aln_len:
		#		print('ignoring')
				continue

			if abs(int(s[0]) - int(prev_line[1])) < smooth_gap_width and \
				abs(int(s[2]) - int(prev_line[3])) < smooth_gap_width:
			#	print("joining")
				pass
			else:
	#				print("terminating")
				newline = aln_start_line
				newline[1] = prev_line[1]
				newline[3] = prev_line[3]
				if int(newline[1]) - int(newline[0]) > min_smoothed_aln_len:
	#				print("storing")
					smoothed_lines.append(newline)
				aln_start_line = s
			prev_line = s


#print("output")
#for s in smoothed_lines:
#	print(s)

if len(smoothed_lines) > 0:
	#is it a corner cutting parallel off-diagonal?
	if int(smoothed_lines[0][0]) < contig_edge_margin and int(smoothed_lines[0][3]) > max_tiglen - contig_edge_margin:
		if int(smoothed_lines[-1][2]) < contig_edge_margin and int(smoothed_lines[-1][1]) > max_tiglen - contig_edge_margin:
			#out.write("It's overcircularized!")

			out.write(tigname + ":" + smoothed_lines[0][0] + '-' + smoothed_lines[-1][0] + "\n")

out.write("done\n")

    shell:
        "guppy_basecaller --cpu_threads_per_caller {threads} -i {params.in_dir} -s {params.out_dir} " +
        "--flowcell {fc} --kit {k}" .format(fc=config['flowcell'], k = config['kit'])

rule basecall_final:
    #Collate called bases into a single fastq file
    input: lambda wildcards: expand('{{sample}}/0.basecall/raw_calls/{foo}/sequencing_summary.txt', foo = [os.path.basename(f).split('.')[0] for f in fast5_files_dict[wildcards.sample]])
    output: '{sample}/0.basecall/{sample}.fq'
    params:
        sample="{sample}"

shell: "NanoPlot --fastq {input} -t {threads} --color green  -o  {wildcards.sample}/0.basecall/nanoplots"

    shell:
        "canu -p {sample}_{wildcards.genome_size} -d {sample}/1.assemble/assemble_{wildcards.genome_size}/ -nanopore-raw {input} " +
        config['canu_args'] +
        " genomeSize={wildcards.genome_size} " + #  stopOnReadQuality=false
        "useGrid={grid} gridOptions='{opts}'".format(
            grid = config['usegrid'],
            opts = config['grid_options']
        )

rule assemble_flye:
    input: lambda wildcards: fq_files_dict[wildcards.sample]
    output:
        "{sample}/1.assemble/assemble_{genome_size}/assembly.fasta"

shell:
    "flye -t {threads} --meta --nano-raw {input} -o {wildcards.sample}/1.assemble/assemble_{wildcards.genome_size}/ -g {wildcards.genome_size}"

shell: """
    bedtools makewindows -g {input[1]} -w {params.window_width} | join - {input[1]} | tr ' ' '\t' | \
    cut -f1-4 | awk '{{if ($2 > {params.window_width} && $3 < $4 - {params.window_width} && $4 > {params.min_tig_size}) print $0}}' | \
    xargs -P 16 -l bash -c '
     htsbox samview {input[2]} $0:$1-$1 -p | \
     cut -f8,9 | awk "{{if (\$1 < $1 - ({params.window_width}/2) && \$2 > $1 + ({params.window_width}/2)) print \$0}}" | wc -l | \
     paste <(echo $0) <(echo $1) - ' | awk '{{if ($3 < 2) print $0}}
    ' > {output}
    """

shell: """
    if [ -s {input[0]} ] #if the input is nonempty...
    then
        cat {input[0]} | grep -v ^# | sort -k1,1 -k2,2g | join - <(cat {input[1]} | sort -k1,1 -k2,2g) | \
        awk '{{
            if ($1 == prev_tig){{
             print($1,prev_coord,$2)
             }}
            else{{
                if (prev_len > 0){{
                    print(prev_tig,prev_coord,prev_len)
                }}
                print($1,"1",$2)
                }}
            prev_tig = $1
            prev_coord = $2
            prev_len = $4
            }}
            END {{ print(prev_tig,prev_coord,prev_len)
        }}' | sed "s/\(.*\) \(.*\)\ \(.*\)/\\1:\\2-\\3/g" |  xargs samtools faidx {input[2]} \
        | cut -f1 -d ':' | awk '(/^>/ && s[$0]++){{$0=$0\"_\"s[$0]}}1;' > {output[0]}

        cut -f1 {input[0]} > {wildcards.sample}/{wildcards.sequence}.tigs.toremove

        # DM: temp fix to skip misassembly if formatting fails
        x=$(cat {input[1]} | wc -l)
        if [ $x -gt 1 ]
        then
            grep -vf {wildcards.sample}/{wildcards.sequence}.tigs.toremove {input[1]} | cut -f1 | xargs samtools faidx {input[2]} >> {output[0]}
            rm {wildcards.sample}/{wildcards.sequence}.tigs.toremove
        else
            cp {input[2]} {output}
            echo "Misassembly error: skipping misassembly removal"
        fi

    else
        cp {input[2]} {output}
    fi
    """

shell: """
    racon -m 8 -x -6 -g -8 -w 500 -t {threads} {input} > {output}
    """

shell: """
    medaka consensus {input[1]} {output} --model r941_flip213 --threads {threads} --regions {wildcards.range}
    """

shell:
    "bwa index {input.ref}; bwa mem -t {threads} {input.ref} {input.reads} | \
        samtools sort --threads {threads} > {output}"

shell: """
    # set env var $i to be the smallest read subset decimal (in increments of 0.1, with a couple very low values thrown in, too)
    # sufficient to generate at least 40x coverage depth of the target sequence, or
    # 1 if 40x coverage cannot be achieved with the available read data

    for i in 0.01 0.05 $(seq 0.1 0.1 1);
    do
       cov=$(samtools view {input[1]} -s $i -h {params.myrange} | samtools depth - | cut -f3 | awk '{{sum+=$1}}END{{print sum/(NR+1)}}')
       if [ $(echo $cov'>'{params.target_coverage}|bc) -eq 1 ]
       then
           break
       fi
    done
    echo Using $i x of total coverage;

    samtools view -h -O BAM -s $i {input[1]} {params.myrange} > {params.bam}
    samtools index {params.bam}
    samtools faidx {input[0]} $(echo {params.myrange}| cut -f1 -d ':') | cut -f1 -d ':' > {params.fa}
    java -Xmx{params.java_mem}G -jar $(which pilon | sed 's/\/pilon//g')/../share/pilon*/pilon*.jar \
        --genome {params.fa} --targets {params.myrange} \
        --unpaired {params.bam} --output {params.sample}_{params.myrange} --outdir {params.subrun_folder} \
        --vcf --nostrays --mindepth 1
    bgzip {params.subrun_folder}/{params.sample}_{params.myrange}.vcf
    tabix -fp vcf {params.subrun_folder}/{params.sample}_{params.myrange}.vcf.gz
    """

shell: """
    #workaround!  Snakemake was causing the vcf's to appear newer than the indices, which tabix didn't like
    touch {params.sample}/2.polish/pilon/sub_runs/*/*.vcf.gz.tbi

    #get properly sorted intervals
    ls {params.sample}/2.polish/pilon/ranges/ | tr ':-' '\t' | sort -k1,1 -k2,2g | awk '{{print $1,":",$2,"-",$3}}' | tr -d ' ' > sorted_ranges.tmp

    #get header
    (zcat {params.sample}/2.polish/pilon/sub_runs/*/*.vcf.gz | head -1000 | grep ^#

    #get corrections within each range (omitting DUP records, which bcftools can't understand)
    cat sorted_ranges.tmp | xargs -n 1 -I foo sh -c "
    tabix {params.sample}/2.polish/pilon/sub_runs/foo/{params.sample}_foo.vcf.gz foo | grep -v '0/0' | grep -v DUP" |

    #sort by position
    sort -k1,1 -k2,2g) |

    #compress and store
    bgzip > {output} || true

    #index
    tabix -p vcf {output}
    """

shell:
    "minimap2 -t {threads} -x map-ont {input} > {output}"