from os import listdir
from os.path import isfile, join

in_dir = "/home/gk/analysis/intrahost/raw_reads/field_samples/zika"
out_dir = "/home/gk/analysis/2018.11.22.1542874525.ZI"
bed = "/home/gk/db/zika/amplicons/zika_primers.bed"
ref="/home/gk/db/zika/ZIKV_RIO.fasta"
index="/home/gk/db/zika/ZIKV_RIO"
pair_information="/home/gk/db/zika/amplicons/pair_information.tsv"

primer_csv = "/home/gk/db/zika/amplicons/zika_primers.csv"


#in_dir = "/home/gk/analysis/intrahost/raw_reads/field_samples/wnv"

#bed = "/home/gk/db/wnv/WNV_400.bed"
#ref="/home/gk/db/wnv/JF703162.fa"
#index="/home/gk/db/wnv/JF703162"
#pair_information="/home/gk/db/wnv/pair_information.tsv"

#primer_csv = "/home/gk/db/wnv/wnv_primers.csv"

# File name format of replicates - ZI26_a_R{1-2}.fastq.gz ZI26_b_R{1-2}.fastq.gz
lst = listdir(in_dir)
files = [f for f in lst  if isfile(join(in_dir, f)) and ".fastq.gz" in f]
files.sort()

SAMPLES = []
for i in files:
    SAMPLES.append(i.split("_")[0])
SAMPLES = list(set(SAMPLES))

SAMPLES=SAMPLES[:2]

REPS=["a","b"]


rule all:
    input:
        expand("{out_dir}/masked_variants/{sample}.filtered.tsv", out_dir = out_dir, sample = SAMPLES)

rule filter_variants:
    input:
        expand("{{out_dir}}/masked_variants/{{sample}}_{rep}.tsv", rep=REPS)
    output:
        "{out_dir}/masked_variants/{sample}.filtered.tsv"
    shell:
        """
        ivar filtervariants -p {output} {input}
        """

rule call_variants_post_removal:
    input:
        "{out_dir}/masked/{sample}_{rep}.masked.sorted.bam",
        "{out_dir}/consensus/{sample}.fa"
    output:
        "{out_dir}/masked_variants/{sample}_{rep}.tsv"
    shell:
        """
        samtools mpileup -A -d 0 --reference {input[1]} -Q 0 -F 0 {input[0]} | ivar variants -p {output} -t 0.03
        """

rule remove_reads:
    input:
        "{out_dir}/realigned/{sample}_{rep}.realigned.trimmed.sorted.bam",
        "{out_dir}/{sample}_{rep}_masked_primer_names.txt",
        "{out_dir}/bed/{sample}.bed"
    output:
        "{out_dir}/masked/{sample}_{rep}.masked.sorted.bam"
    shell:
        """
        mkdir -p $(dirname {output})/
        ivar removereads -i {input[0]} -p {wildcards.out_dir}/masked/{wildcards.sample}_{wildcards.rep}.masked -t {input[1]} -b {input[2]}
        samtools sort -T {wildcards.sample}_{wildcards.rep}_remove_reads -o {output} {wildcards.out_dir}/masked/{wildcards.sample}_{wildcards.rep}.masked.bam
        samtools index {output}
        """

rule get_masked:
    input:
        "{out_dir}/primer_mismatches/{sample}.tsv",
        "{out_dir}/bed/{sample}.bed",
        "{pair_information}".format(pair_information = pair_information)
    output:
        "{out_dir}/{sample}_{rep}_masked_primer_names.txt"
    shell:
        "ivar getmasked -i {input[0]} -b {input[1]}  -f {input[2]} -p {output}"

rule call_variants_in_primer:
    input:
        "{out_dir}/bed/{sample}.bam",
        "{out_dir}/consensus/{sample}.fa"
    output:
        "{out_dir}/primer_mismatches/{sample}.tsv"
    shell:
        "samtools mpileup -A -d 0 --reference {input[1]} -Q 0 -F 0 {input[0]} | ivar variants -p {output} -t 0.03"

rule realigned_trim_primer_quality:
    input:
        "{out_dir}/realigned/{sample}_{rep}.realigned.sorted.bam",
        "{out_dir}/bed/{sample}.bed"
    output:
        "{out_dir}/realigned/{sample}_{rep}.realigned.trimmed.sorted.bam",
    shell:
        """
        mkdir -p $(dirname {output})/
        ivar trim -b {input[1]} -p {wildcards.out_dir}/realigned/{wildcards.sample}_{wildcards.rep}.realigned.trimmed -i {input[0]}
        samtools sort -T {wildcards.sample}_{wildcards.rep}_retrim -o {output} {wildcards.out_dir}/realigned/{wildcards.sample}_{wildcards.rep}.realigned.trimmed.bam
        samtools index {output}
        """

rule create_bed:
    input:
        "{out_dir}/bed/{sample}.bam"
    output:
        "{out_dir}/bed/{sample}.bed"
    shell:
        """
        bedtools bamtobed -i {input} > {output}
        """

rule create_primer_bam:
    input:
        "{out_dir}/index/{sample}.bwt",
        "{out_dir}/bed/primers.fa"
    output:
        "{out_dir}/bed/{sample}.bam"
    params:
        index= "{out_dir}/index/{sample}"
    shell:
        """
        bwa mem -k 5 -T 16 {params.index} {input[1]} | samtools view -bS -F 4 -o {output}
        """

rule create_primer_fasta:
    input:
        "{csv}".format(csv = primer_csv)
    output:
        "{out_dir}/bed/primers.fa"
    shell:
        "cut -f 1,2 -d ',' {input} | sed 's/^/>/g' | tr ',' '\n' > {output}"

rule align_to_consensus:
    input:
        "{out_dir}/index/{sample}.bwt",
        expand("{in_dir}/{{sample}}_{{rep}}_R{readno}.fastq.gz", readno = [1,2], in_dir = in_dir)
    output:
        "{out_dir}/realigned/{sample}_{rep}.realigned.sorted.bam"
    shell:
        """
        bwa mem {wildcards.out_dir}/index/{wildcards.sample} {input[1]} {input[2]} | samtools view -F 4 -Sb | samtools sort -T {wildcards.sample}_{wildcards.rep}_align_consensus -o {output}
        samtools index {output}
        """

rule call_consensus:
    input:
        "{out_dir}/merged_bams/{sample}.merged.bam"
    output:
        "{out_dir}/consensus/{sample}.fa",
        "{out_dir}/index/{sample}.bwt"
    shell:
        """
        samtools mpileup -A -d 0 -Q 0 -F 0 {input} | ivar consensus -p {output[0]}
        bwa index -p {wildcards.out_dir}/index/{wildcards.sample} {output[0]}
        """

rule merge_replicates:
    input:
        expand("{{out_dir}}/trimmed/{{sample}}_{rep}.trimmed.sorted.bam", rep=REPS)
    output:
        "{out_dir}/merged_bams/{sample}.merged.bam"
    shell:
        """
        mkdir -p $(dirname {output}/)
        samtools merge {output} {input}
        """

rule trim_primer_quality:
    input:
        "{out_dir}/aligned/{sample}_{rep}.sorted.bam",
        "{b}".format(b = bed)
    output:
        "{out_dir}/trimmed/{sample}_{rep}.trimmed.sorted.bam"
    shell:
        """
        mkdir -p {wildcards.out_dir}/trimmed/
        ivar trim -b {input[1]} -p {wildcards.out_dir}/trimmed/{wildcards.sample}_{wildcards.rep}.trimmed -i {input[0]}
        samtools sort -T {wildcards.sample}_{wildcards.rep}_trim -o {output} {wildcards.out_dir}/trimmed/{wildcards.sample}_{wildcards.rep}.trimmed.bam
        samtools index {output}
        """

rule align_reads:
    input:
        expand("{in_dir}/{{sample}}_{{rep}}_R{num}.fastq.gz", in_dir = in_dir, num=[1,2])
    params:
        index="{i}".format(i = index)
    output:
        "{out_dir}/aligned/{sample}_{rep}.sorted.bam"
    shell:
        """
        mkdir -p $(dirname {output})/
        bwa mem {params.index} {input[0]} {input[1]} | samtools view -F 4 -Sb | samtools sort -T {wildcards.sample}_{wildcards.rep}_align -o {output}
        samtools index {output}
        """
