from pling.utils import get_pling_root_dir

OUTPUTPATH = config["output_dir"]
dcj_threshold = config["dcj_dist_threshold"]

def smash():
    if config.get("sourmash_only", False):
        return f"--sourmash_only --sourmash --smash_threshold " + str(config["seq_containment_distance"])
    elif config.get("sourmash", False) and not config.get("sourmash_only", False):
        return f"--sourmash --smash_threshold " + config.get("sourmash_threshold", "1")
    else:
        return ""

def get_batching_resources(type):
    if config.get("sourmash", False):
        threads = config["sourmash_threads"]
        mem = config["sourmash_mem"]
    else:
        threads = 1
        mem = 4000
    if type == "threads":
        return threads
    elif type == "mem":
        return mem

def get_containmentpath():
    if config.get("sourmash_only", False):
        return f"{OUTPUTPATH}/containment/all_pairs_containment_distance.tsv"
    else:
        f"{OUTPUTPATH}/containment/not_pairs_containment_distance.tsv"

def get_prev_pling():
    if config.get("previous_pling",False):
        blub = config["previous_pling"].strip().split(",")
        blah = " ".join([f"{bl}/dcj_thresh_{dcj_threshold}_graph/objects" for bl in blub])
        return f"--prev_typing {blah}"
    else:
        return ""

rule all:
    input:
        f"{OUTPUTPATH}/batches"

rule get_batches:
    output:
        directory(f"{OUTPUTPATH}/batches")
    params:
        outputpath = OUTPUTPATH,
        genomes_list = config["genomes_list"],
        batch_size = config["batch_size"],
        sourmash = smash(),
        containmentpath = get_containmentpath(),
        pling_root_dir = get_pling_root_dir(),
        prev_pling = get_prev_pling()
    threads: get_batching_resources("threads")
    resources:
        mem_mb=lambda wildcards, attempt: get_batching_resources("mem")*attempt
    shell:
        """
        PYTHONPATH={params.pling_root_dir} python {params.pling_root_dir}/pling/batching/get_batches.py \
            --genomes_list {params.genomes_list} \
            --batch_size {params.batch_size} \
            --outputpath {params.outputpath} \
            {params.sourmash} \
            --containmentpath {params.containmentpath} \
            {params.prev_pling}
        """
