from latch.types.metadata import SnakemakeMetadata, SnakemakeFileParameter
from latch.types.directory import LatchDir
from latch.types.metadata import LatchAuthor, LatchMetadata, LatchParameter

SnakemakeMetadata(
    output_dir = LatchDir("latch:///sample_output"),
    display_name="snakemake_tutorial_workflow",
    author=LatchAuthor(
            name="Kenneth",
    ),
    parameters={
        "samples" : SnakemakeFileParameter(
                display_name="Sample Input Directory",
                description="A directory full of FastQ files",
                type=LatchDir,
                path=Path("data/samples"),
        ),
        "ref_genome" : SnakemakeFileParameter(
                display_name="Indexed Reference Genome",
                description="A directory with a reference Fasta file and the 6 index files produced from `bwa index`",
                type=LatchDir,
                path=Path("genome"),
        ),
    },
)

SAMPLES = ["A", "B"]


rule all:
    input:
        "plots/quals.svg"


rule bwa_map:
    input:
        "genome/genome.fa",
        "data/samples/{sample}.fastq",
        "genome/genome.fa.amb",
        "genome/genome.fa.ann", 
        "genome/genome.fa.bwt", 
        "genome/genome.fa.fai", 
        "genome/genome.fa.pac", 
        "genome/genome.fa.sa"
    output:
        "mapped_reads/{sample}.bam"
    shell:
        "bwa mem genome/genome.fa data/samples/{wildcards.sample}.fastq | samtools view -Sb - > {output}"


rule samtools_sort:
    input:
        "mapped_reads/{sample}.bam"
    output:
        "sorted_reads/{sample}.bam"
    shell:
        "samtools sort -T sorted_reads/{wildcards.sample} "
        "-O bam {input} > {output}"


rule samtools_index:
    input:
        "sorted_reads/{sample}.bam"
    output:
        "sorted_reads/{sample}.bam.bai"
    shell:
        "samtools index {input}"


rule bcftools_call:
    input:
        fa="genome/genome.fa",
        bam=expand("sorted_reads/{sample}.bam", sample=SAMPLES),
        bai=expand("sorted_reads/{sample}.bam.bai", sample=SAMPLES)
    output:
        "calls/all.vcf"
    shell:
        "bcftools mpileup -f {input.fa} {input.bam} | "
        "bcftools call -mv - > {output}"


rule plot_quals:
    input:
        "calls/all.vcf"
    output:
        "plots/quals.svg"
    script:
        "scripts/plot-quals.py"
