import os
import hashlib

# Prepwork

def get_genome_file_name(wildcards):
    return '../data/reference_genomes/'+'_'.join(x.lower() for x in wildcards.species.split())+'.fasta'

def get_unique_download_dir(wildcards):
    unique_id = hashlib.md5(wildcards.species.encode()).hexdigest()[:8]
    return config["tmp_dir"]+'/'+unique_id

def format_cli_arg(wildcards):
    # Replaces underscores with spaces for the CLI command
    return wildcards.species.replace("_", " ")

rule pull_ncbi_datasets_cli:
    output:
        config["datasets_binary"]
    params:
        arch=config["arch"]
    threads: 1
    shell:
        """
        curl https://ftp.ncbi.nlm.nih.gov/pub/datasets/command-line/v2/linux-{params.arch}/datasets -o {output}
        chmod a+x {output}
        """

rule pull_dorado_models:
    output:
        directory(f"{config['dorado_model_dir']}/{config['fast_model_name']}"),
        directory(f"{config['dorado_model_dir']}/{config['hac_model_name']}")
    threads: 1
    shell:
        """
        dorado download --models-directory {config[dorado_model_dir]} --model {config[fast_model_name]}
        dorado download --models-directory {config[dorado_model_dir]} --model {config[hac_model_name]}
        """

rule download_reference_genome:
    input:
        config["datasets_binary"]
    output:
        temp('../data/reference_genomes/{species}.fasta')
    params:
        api_key=os.environ['NCBI_API_KEY'],
        wd=get_unique_download_dir,
        ncbi_tax_name=lambda w: format_cli_arg(w)
    threads: 1
    shell:
        """
        echo {wildcards.species}
        rm -rf {params.wd}
        mkdir {params.wd}
        {config[datasets_binary]} download genome taxon "{params.ncbi_tax_name}" --reference --no-progressbar --filename {params.wd}/{wildcards.species}.zip
        unzip -d {params.wd} -o {params.wd}/{wildcards.species}.zip 
        mv {params.wd}/ncbi_dataset/data/*/*fna {output}
        rm -rf {params.wd}
        """

rule concatenate_reference_genomes:
    input:
        expand("../data/reference_genomes/{species}.fasta",species=config["species_list"])
    output:
        "../data/reference_genomes/full_reference.fasta"
    threads: 1
    shell:
        """
        cat {input} > {output}
        """

rule create_minimap_index:
    input: 
        "../data/reference_genomes/full_reference.fasta"
    output:
        "../data/reference_genomes/full_reference.mmi"
    threads:
        32
    conda:
        "../envs/minimap.yaml"
    shell:
        """
        minimap2 -x map-ont -t {threads} -d {output} {input}
        """