import os
import yaml

rule all:
    input:
        f"{config['output']}config.yaml",
        f"{config['output']}log/umap.done"
        
rule create_config:
    output:
        f"{config['output']}log/create_config.done",
        f"{config['output']}config.yaml"
    script:
        "scripts/createconfig.py"

rule kb_count:
    input:
        f"{config['output']}config.yaml",
        f"{config['output']}log/create_config.done"
    output:
        f"{config['output']}log/kb.done"
    shell:
        """
        echo -e "\033[0;33mStarting to quantify the data...\033[0m"

        # Check for whitelist
        if [ "{config[whitelist]}" = "None" ] || [ -z "{config[whitelist]}" ]; then
            output_log=$(kb count -i {config[index]} -g {config[transcripts]} --overwrite -o {config[output]} -x {config[technology]} {config[sample1]} {config[sample2]} --h5ad 2> /dev/null)
        else
            output_log=$(kb count -i {config[index]} -g {config[transcripts]} --overwrite -o {config[output]} -x {config[technology]} -w {config[whitelist]} {config[sample1]} {config[sample2]} --h5ad 2> /dev/null)
        fi
        
        # Check for "no reads pseudoaligned" in the log
        if echo "$output_log" | grep -q "no reads pseudoaligned"; then
            echo -e "\033[0;31mERROR: kb count failed! No reads pseudoaligned. Check your sample or reference index.\033[0m"
            exit 1
        fi

        touch {output}
        echo -e "\033[0;32mQuantifying the data is done!\033[0m"
        mkdir {config[output]}kb-python/
        mv {config[output]}*.* {config[output]}kb-python/
        mv {config[output]}counts_unfiltered/ {config[output]}kb-python/
        mv {config[output]}kb-python/config.yaml {config[output]}
        """

rule analysis:
    input:
        f"{config['output']}config.yaml",
        f"{config['output']}log/kb.done"
    output:
        f"{config['output']}log/analysis.txt"
    params:
        configfile = lambda wildcards: f"{config['output']}config.yaml"
    script:
        "scripts/analysis.py"

rule multimap:
    input:
        file_viral_accessions = f"{config['output']}log/analysis.txt"
    output:
        f"{config['output']}log/multimap.done"
    params:
        configfile = lambda wildcards: f"{config['output']}config.yaml",
        snakefile_dir = lambda wildcards: workflow.basedir
    script:
        "scripts/multimap.py"

rule detection:
    input:
        file_viral_accessions = f"{config['output']}log/analysis.txt",
        multimap_done = f"{config['output']}log/multimap.done"
    output:
        f"{config['output']}log/detection.done"
    params:
        configfile = lambda wildcards: f"{config['output']}config.yaml",
        snakefile_dir = lambda wildcards: workflow.basedir
    script:
        "scripts/detection.py"

rule umap:
    input: 
        f"{config['output']}config.yaml",
        f"{config['output']}log/detection.done"
    output:
        f"{config['output']}log/umap.done"
    params:
        configfile = lambda wildcards: f"{config['output']}config.yaml"
    script:
        "scripts/umap.py"