# This file contains rules to benchmark pangraph performances as a function of input dataset size

import json

# extract config for accuracy rules
SB_config = config["size-benchmark"]

# extract simulation parameters from config files
SB_hgt = SB_config["hgt"]
SB_snps = SB_config["snps"]
SB_inv = SB_config["inv"]
SB_del = SB_config["del"]
SB_ker_opt = SB_config["kernel-options"]
SB_trials = list(range(SB_config["trials"]))

# list of n. of isolates and genome sizes to test
with open(SB_config["sim-params"], "r") as f:
    SB_sim_params = json.load(f)
SB_Ns = SB_sim_params["N"]
SB_Ls = SB_sim_params["L"]


# generate synthetic data to benchmark pangraph performances as a function of dataset size
rule SB_generate_data:
    message:
        "generating size-benchmark sequence with N={wildcards.N}, L={wildcards.L}, trial={wildcards.n}"
    output:
        "size-benchmark/sequences/seqs_{N,[0-9]+}_{L,[0-9]+}_{n,[0-9]+}.fa",
    params:
        hgt=SB_hgt,
        snps=SB_snps,
        inv=SB_inv,
        delet=SB_del,
        T=lambda w: int(w.N) // 2,
    shell:
        """
        julia -t 1 --project=. workflow_scripts/make-sequence.jl \
            -N {wildcards.N} -L {wildcards.L} \
            | julia -t 1 --project=./.. ./../src/PanGraph.jl generate \
            -m {params.snps} -r {params.hgt} -t {params.T} \
            -i {params.inv} -d {params.delet} \
            > {output}
        """


# build pangraphs and time the command. Save the results in a txt file
rule SB_time_pangraph:
    message:
        "size-benchmark time pangraph construction N={wildcards.N}, L={wildcards.L}, trial={wildcards.n}"
    input:
        rules.SB_generate_data.output,
    output:
        "size-benchmark/timings/time_{N,[0-9]+}_{L,[0-9]+}_{n,[0-9]+}.txt",
    params:
        ker_opt=SB_ker_opt,
    conda:
        "../conda_envs/pangraph_build_env.yml"
    shell:
        """
        echo "N = {wildcards.N}" > {output}
        echo "L = {wildcards.L}" >> {output}
        echo "trial = {wildcards.n}" >> {output}
        export JULIA_NUM_THREADS=8
        /usr/bin/time --verbose -o {output} -a pangraph build \
            --circular {params.ker_opt} \
            {input} > /dev/null
        """


# Parse all txt files from the previous rules and collect the results in a csv dataframe
rule SB_summary_dataframe:
    message:
        "creating summary dataframe for pangraph performace vs dataset size"
    input:
        expand(
            "size-benchmark/timings/time_{N}_{L}_{n}.txt",
            N=SB_Ns,
            L=SB_Ls,
            n=SB_trials,
        ),
    output:
        "size-benchmark/summary/size_vs_time.csv",
    conda:
        "../conda_envs/bioinfo_env.yml"
    shell:
        """
        python3 workflow_scripts/size_benchmark_summary_dataframe.py --csv {output} --txt {input}
        """


# Plots summary statistics from the csv dataframe generated by the previous rule
rule SB_summary_plot:
    message:
        "creating benchmark plot for time vs size"
    input:
        rules.SB_summary_dataframe.output,
    output:
        "figs/size-benchmark/size_vs_time.pdf",
    conda:
        "../conda_envs/bioinfo_env.yml"
    shell:
        """
        python3 workflow_scripts/size_benchmark_summary_plot.py --csv {input} --pdf {output}
        """


# rule to generate all relevant plots for the size benchmark
rule SB_all:
    input:
        rules.SB_summary_plot.output,
