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: expand('{{sample}}/0.basecall/raw_calls/{foo}/sequencing_summary.txt', foo = [os.path.basename(f).split('.')[0] for f in fast5_files])
	output: '{sample}/0.basecall/{sample}.fq'
	shell:
		"find {sample}/0.basecall/raw_calls/*/*.fastq | xargs cat > {output}"

	shell: "NanoPlot --fastq {input} -t {threads} --color green  -o " + "{s}/0.basecall/nanoplots".format(s = sample)

rule assemble_canu:
	#Run canu. This can be run either in distributed fashion on the cluster or in a single job. If run in distributed fashion (usedgrid is set to 'True')
	#then canu must be installed in the user's environment. In this case, the singualarity image is not used. Canu arguments are passed through from
	#the config within the 'canu_args' key. 'gridOptions' are also passed through from the config and instruct Canu on any additional parameters
	#needed to run on the cluster.
	input: rules.basecall_final.output
	output:
		'{sample}/1.assemble/assemble_{genome_size}/{sample}_{genome_size}.contigs.fasta',
		#'{sample}/1.assemble/assemble_{genome_size}/{sample}_{genome_size}.correctedReads.fasta.gz'
	threads: 1
	resources:
		mem=100,
		time=80
	singularity: singularity_image if config['usegrid'] != 'True' and config['usegrid'] != True else '' #switch between image and local environment for canu depending on whether cluster execution is required
	shell:
		"canu -p {sample}_{wildcards.genome_size} -d {sample}/1.assemble/assemble_{wildcards.genome_size}/ -nanopore-raw {input} " +
		config['canu_args'] +
		" stopOnReadQuality=false genomeSize={wildcards.genome_size} " +
		"useGrid={grid} gridOptions='{opts}'".format(
			grid = config['usegrid'],
			opts = config['grid_options']
		)

rule assemble_flye:
	input: rules.basecall_final.output
	output:
		"{sample}/1.assemble/assemble_{genome_size}/assembly.fasta"

shell:
	"flye -t {threads} --nano-raw {input} --meta -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]} > {sample}/{wildcards.sequence}.tigs.toremove
		grep -vf {sample}/{wildcards.sequence}.tigs.toremove {input[1]} | cut -f1 | xargs samtools faidx {input[2]} >> {output[0]}
		rm {sample}/{wildcards.sequence}.tigs.toremove
	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[0]}; bwa mem -t {threads} {input} | 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 {wildcards.range} | 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]} {wildcards.range} > {params.bam}
		samtools index {params.bam}
		samtools faidx {input[0]} $(echo {wildcards.range}| 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 {wildcards.range} \
			--unpaired {params.bam} --output {sample}_{wildcards.range} --outdir {params.subrun_folder} \
			--vcf --nostrays --mindepth 1
		bgzip {params.subrun_folder}/{sample}_{wildcards.range}.vcf
		tabix -fp vcf {params.subrun_folder}/{sample}_{wildcards.range}.vcf.gz
		"""

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

	#get properly sorted intervals
	ls {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 {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 {sample}/2.polish/pilon/sub_runs/foo/{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}"