import os


configfile: "config.yaml"


# Restrict wildcard expansion to prevent globbing through subdirectories
wildcard_constraints:
    sample="[^\/]+",


INPUT_DIR = config["input"]
OUTPUT_DIR = config["output"]

# Raise error if input dir does not exist
if not os.path.exists(INPUT_DIR):
    raise ValueError("Input directory does not exist.")

# Raise error if output dir exists and is not a directory
if os.path.exists(OUTPUT_DIR) and not os.path.isdir(OUTPUT_DIR):
    raise ValueError("Output directory is not a directory.")

# Check if input dir is apollo_mapping_output, if so, set input dir to variants/snps subfolder
if (
    os.path.exists(INPUT_DIR + "/variants/snps")
    and os.path.exists(INPUT_DIR + "/variants/raw")
    and os.path.exists(INPUT_DIR + "/variants/marked")
    and os.path.exists(INPUT_DIR + "/audit_trail")
    and os.path.exists(INPUT_DIR + "/clean_fastq")
):
    INPUT_DIR = INPUT_DIR + "/variants/snps"

(SAMPLES,) = glob_wildcards(INPUT_DIR + "/{sample}.vcf")

if len(SAMPLES) == 0:
    raise ValueError("No VCF files found in input directory.")


rule all:
    input:
        OUTPUT_DIR + "/tree.nwk",
        OUTPUT_DIR + "/vcf_snp.aln.fasta",
        OUTPUT_DIR + "/snp_distance_matrix.tsv",


rule index_vcf:
    input:
        INPUT_DIR + "/{sample}.vcf",
    output:
        vcf=OUTPUT_DIR + "/vcf/{sample}.vcf.gz",
        tbi=OUTPUT_DIR + "/vcf/{sample}.vcf.gz.tbi",
    conda:
        "envs/main.yaml"
    shell:
        """
        bgzip -c {input} > {output.vcf}
        tabix -p vcf {output.vcf}
        """


rule merge_vcf:
    input:
        vcf=expand(OUTPUT_DIR + "/vcf/{sample}.vcf.gz", sample=SAMPLES),
        tbi=expand(OUTPUT_DIR + "/vcf/{sample}.vcf.gz.tbi", sample=SAMPLES),
    output:
        OUTPUT_DIR + "/merged.vcf",
    conda:
        "envs/main.yaml"
    shell:
        """
        vcf-merge {input.vcf} > {output}
        """


rule make_tree:
    input:
        OUTPUT_DIR + "/merged.vcf",
    output:
        OUTPUT_DIR + "/tree.nwk",
    conda:
        "envs/vk.yaml"
    shell:
        """
        vk phylo tree upgma {input} > {output}
        """


rule make_aln:
    input:
        OUTPUT_DIR + "/merged.vcf",
    output:
        OUTPUT_DIR + "/vcf_snp.aln.fasta",
    conda:
        "envs/vk.yaml"
    shell:
        """
        vk phylo fasta {input} > {output}
        """


rule make_dist:
    input:
        OUTPUT_DIR + "/vcf_snp.aln.fasta",
    output:
        OUTPUT_DIR + "/snp_distance_matrix.tsv",
    conda:
        "envs/main.yaml"
    shell:
        """
        snp-dists {input} > {output}
        """
