import glob
import pandas as pd

"""
Author: Christiane Hassenrueck
Affiliation: Leibniz Institute for Baltic Sea Research Warnemuende (IOW)
Aim: 
"""

# specify config file
# change this path to the location of the config file with your project specific settings
configfile: "/bio/Common_repositories/workflow_templates/Amplicon_dada2_MiSeq/config/config.yaml"

# specify wildcards for samples and paths of input file names
sample_table = pd.read_table(config["sample_list"], sep="\t", lineterminator="\n", header=None, names = ["SID", "path_R1", "path_R2"])
SAMPLES = sample_table.SID.to_list()
PATH_R1 = sample_table.set_index('SID')['path_R1'].to_dict()
PATH_R2 = sample_table.set_index('SID')['path_R2'].to_dict()

# snakemake rules

rule all:
	input:
		plot_filt = config["adir"] + "/Intermediate_results/" + config["analysis"] + "/Parameter_screening_" + config["project"] + ".pdf",
		nseq = config["adir"] + "/Intermediate_results/" + config["analysis"] + "/nseqs_" + config["project"] + ".txt",
		checksums = config["adir"] + "/Intermediate_results/" + config["analysis"] + "/md5_checksums_" + config["project"] + ".txt"

def R1_fun(wildcards):
	return PATH_R1[wildcards.sample]

def R2_fun(wildcards):
	return PATH_R2[wildcards.sample]

rule primer_clipping:
	input:
		fq_R1 = R1_fun,
		fq_R2 = R2_fun
	output:
		fq_fr_R1 = config["adir"] + "/Validated_data/" + config["analysis"] + "/{sample}_clip_fr_R1.fastq.gz",
		fq_fr_R2 = config["adir"] + "/Validated_data/" + config["analysis"] + "/{sample}_clip_fr_R2.fastq.gz",
		fq_rf_R1 = config["adir"] + "/Validated_data/" + config["analysis"] + "/{sample}_clip_rf_R1.fastq.gz",
		fq_rf_R2 = config["adir"] + "/Validated_data/" + config["analysis"] + "/{sample}_clip_rf_R2.fastq.gz"
	params:
		fwd = config["fwd"],
		rev = config["rev"],
		ofwd = config["ofwd"],
		orev = config["orev"],
		error = config["error"]
	conda:
		config["wdir"] + "/envs/cutadapt.yaml"
	log:
		fr = config["adir"] + "/Intermediate_results/Logfiles/" + config["analysis"] + "/{sample}.clip_fr.log",
		rf = config["adir"] + "/Intermediate_results/Logfiles/" + config["analysis"] + "/{sample}.clip_rf.log"
	shell:
		"""
		cutadapt -j 1 --no-indels -e {params.error} -g "{params.fwd};o={params.ofwd}" -G "{params.rev};o={params.orev}" -m 50 --discard-untrimmed -o {output.fq_fr_R1} -p {output.fq_fr_R2} {input.fq_R1} {input.fq_R2} > {log.fr} 2>&1
		cutadapt -j 1 --no-indels -e {params.error} -g "{params.rev};o={params.orev}" -G "{params.fwd};o={params.ofwd}" -m 50 --discard-untrimmed -o {output.fq_rf_R1} -p {output.fq_rf_R2} {input.fq_R1} {input.fq_R2} > {log.rf} 2>&1
		"""

rule count_seqs:
	input:
		clipped = expand(config["adir"] + "/Validated_data/" + config["analysis"] + "/{sample}_clip_fr_R1.fastq.gz", sample = SAMPLES)
	output:
		nseq = config["adir"] + "/Intermediate_results/" + config["analysis"] + "/nseqs_" + config["project"] + ".txt"
	params:
		sample_list = config["sample_list"],
		adir = config["adir"],
		analysis = config["analysis"]
	conda:
		config["wdir"] + "/envs/parallel.yaml"
	threads:
		config["sample_threads"]
	shell:
		"""
		cut -f2 {params.sample_list} | parallel -k -j {threads} "bzcat {{}} | awk 'END {{print NR/4}}'" | paste <(cut -f1 {params.sample_list}) - > {params.adir}/Intermediate_results/{params.analysis}/tmp1
		cut -f1 {params.sample_list} | parallel -k -j {threads} "zcat {params.adir}/Validated_data/{params.analysis}/{{}}_clip_fr_R1.fastq.gz | awk 'END {{print NR/4}}'" | paste {params.adir}/Intermediate_results/{params.analysis}/tmp1 - > {params.adir}/Intermediate_results/{params.analysis}/tmp2
		cut -f1 {params.sample_list} | parallel -k -j {threads} "zcat {params.adir}/Validated_data/{params.analysis}/{{}}_clip_rf_R1.fastq.gz | awk 'END {{print NR/4}}'" | paste {params.adir}/Intermediate_results/{params.analysis}/tmp2 - > {params.adir}/Intermediate_results/{params.analysis}/tmp3
		echo -e 'SID\\tDemux\\tClipped_fr\\tClipped_rf' | cat - "{params.adir}/Intermediate_results/{params.analysis}/tmp3" > {output.nseq}
		rm {params.adir}/Intermediate_results/{params.analysis}/tmp*
		"""

rule screen_filter_settings:
	input:
		clipped = expand(config["adir"] + "/Validated_data/" + config["analysis"] + "/{sample}_clip_fr_R1.fastq.gz", sample = SAMPLES)
	output:
		plot_filt = config["adir"] + "/Intermediate_results/" + config["analysis"] + "/Parameter_screening_" + config["project"] + ".pdf",
		plot_qual_R1 = config["adir"] + "/Intermediate_results/" + config["analysis"] + "/QualityProfile_" + config["project"] + "_R1.pdf",
		plot_qual_R2 = config["adir"] + "/Intermediate_results/" + config["analysis"] + "/QualityProfile_" + config["project"] + "_R2.pdf"
	params:
		script = config["wdir"] + "/scripts/trimming.R",
		truncLen_R1_min = config["truncLen_R1_min"],
		truncLen_R1_max = config["truncLen_R1_max"],
		fragment_max = config["fragment_max"],
		maxEE_min = config["maxEE_min"],
		maxEE_max = config["maxEE_max"],
		project = config["project"],
		sample_list = config["sample_list"],
		adir = config["adir"],
		analysis = config["analysis"]
	conda:
		config["wdir"] + "/envs/r.yaml"
	threads:
		config["sample_threads"]
	log:
		config["adir"] + "/Intermediate_results/Logfiles/" + config["analysis"] + "/trimming_R.log"
	shell:
		"""
		{params.script} -n {params.project} -s {params.sample_list} -l {params.truncLen_R1_min} -L {params.truncLen_R1_max} -f {params.fragment_max} -e {params.maxEE_min} -E {params.maxEE_max} -d {params.adir} -a {params.analysis} -c {threads} &>> {log}
		"""

rule prepare_checksums:
	input:
		clipped_fr_R1 = expand(config["adir"] + "/Validated_data/" + config["analysis"] + "/{sample}_clip_fr_R1.fastq.gz", sample = SAMPLES),
		clipped_fr_R2 = expand(config["adir"] + "/Validated_data/" + config["analysis"] + "/{sample}_clip_fr_R2.fastq.gz", sample = SAMPLES),
		clipped_rf_R1 = expand(config["adir"] + "/Validated_data/" + config["analysis"] + "/{sample}_clip_rf_R1.fastq.gz", sample = SAMPLES),
		clipped_rf_R2 = expand(config["adir"] + "/Validated_data/" + config["analysis"] + "/{sample}_clip_rf_R2.fastq.gz", sample = SAMPLES)
	output:
		checksums = config["adir"] + "/Intermediate_results/" + config["analysis"] + "/md5_checksums_" + config["project"] + ".txt"
	params:
		script = config["wdir"] + "/scripts/checksums.R",
		project = config["project"],
		sample_list = config["sample_list"],
		adir = config["adir"],
		analysis = config["analysis"]
	conda:
		config["wdir"] + "/envs/r.yaml"
	shell:
		"""
		md5sum {input.clipped_fr_R1} > "{params.adir}/Intermediate_results/{params.analysis}/tmp_fr_R1"
		md5sum {input.clipped_fr_R2} > "{params.adir}/Intermediate_results/{params.analysis}/tmp_fr_R2"
		md5sum {input.clipped_rf_R1} > "{params.adir}/Intermediate_results/{params.analysis}/tmp_rf_R1"
		md5sum {input.clipped_rf_R2} > "{params.adir}/Intermediate_results/{params.analysis}/tmp_rf_R2"
		{params.script} -f "{params.adir}/Intermediate_results/{params.analysis}/tmp_fr_R1" -F "{params.adir}/Intermediate_results/{params.analysis}/tmp_fr_R2" -r "{params.adir}/Intermediate_results/{params.analysis}/tmp_rf_R1" -R "{params.adir}/Intermediate_results/{params.analysis}/tmp_rf_R2" -s {params.sample_list} -o {output.checksums}
		rm {params.adir}/Intermediate_results/{params.analysis}/tmp_*
		"""


