import logging
import os
from pathlib import Path
from itertools import product

LOGGER = logging.getLogger('snakemake.logging')

if not 'VPIPE_BASEDIR' in dir():
    VPIPE_BASEDIR = workflow.basedir
LOGGER.info('VPIPE_BASEDIR = %s', VPIPE_BASEDIR)

# Include config file parser and globals
VPIPE_BENCH = True
include: "rules/benchmark.smk"
functions = srcdir("functions.sh")

ALIGNERS = config.benchmark['aligners'].split(',')
SNV_CALLERS = config.benchmark['snv_callers'].split(',')
BENCH_RUNS = list(product(ALIGNERS, SNV_CALLERS))

configs = {}
for opts in BENCH_RUNS:
    dirname = '-'.join(opts)
    configs[dirname] = VPIPE_CONFIG()
    configs[dirname].general['aligner'] = opts[0]
    configs[dirname].general['snv_caller'] = opts[1]
BENCH_DIRS = ','.join([k for k in configs.keys()])

# DUMMY RULE
RESULTS = [f'{bench_dir}/vpipeBench.lock' for bench_dir in configs.keys()]
if len(ALIGNERS) > 1:
    RESULTS.append(
        [f'{bench_dir}/aligner_tests.lock' for bench_dir in configs.keys()]
    )
if len(SNV_CALLERS) > 1:
    RESULTS.append(
        [f'{bench_dir}/caller_tests.lock' for bench_dir in configs.keys()]
    )

rule allbenchRunner:
    input:
        RESULTS

localrules: prepare_run, prepare_config, prepare_testing_run

rule run_simBench:
    output:
        "test_data/test_data.lock",
    params:
        scratch = '1250',
        mem = config.run_simBench['mem'],
        time = config.run_simBench['time'],
        BASEDIR = VPIPE_BASEDIR,
        SAMPLES_BASE = config.input['samples_file'],
        CONFIG_FILE = "test_data/vpipe.config",
        SAMPLES = f"test_data/{config.input['samples_file']}",
        SNAKEMAKE_OPTIONS = lambda wildcards: config.benchmark['snakemake_options'],
        SNAKEMAKE_OPTS_ENVVAR = lambda wildcards: os.environ.get('SNAKEMAKE_OPTS', ''),
        SNAKEMAKE = config.applications['snakemake'],
    log:
        outfile = "logs/run_simBench.out.log",
        errfile = "logs/run_simBench.err.log",
    threads:
        1
    shell:
        """
        # create symbolic links: samples-file and config file
        ln -fs ${{PWD}}/vpipe.config {params.CONFIG_FILE}
        ln -fs ${{PWD}}/{params.SAMPLES_BASE} {params.SAMPLES}
        touch -h {params.CONFIG_FILE}
        touch -h {params.SAMPLES}

        # copy reference file(s). Expected location current directory or
        # sub-directories.
        # Alternative: read references from samples-file and config-file
        REFERENCES=$( find -L . -path ./test_data -prune -o -type f \( -iname \*.fasta -o -iname \*.fa \) -print )
        echo ${{REFERENCES}}
        for i in ${{REFERENCES}}; do
            DIRNAME=$(dirname $i)
            BASENAME=$(basename $i)
            mkdir -p test_data/${{DIRNAME}}
            # copy references files to conserve timestamps
            cp -va ${{i}} test_data/${{DIRNAME}} >> {log.outfile}
        done

        # change to working directory
        CWD=${{PWD}}
        cd test_data/

        if [[ -z "{params.SNAKEMAKE_OPTS_ENVVAR}" ]]; then
            {params.SNAKEMAKE} -s {params.BASEDIR}/vpipeBench.snake simulated_data \
                {params.SNAKEMAKE_OPTIONS}
        else
            {params.SNAKEMAKE} -s {params.BASEDIR}/vpipeBench.snake simulated_data \
                {params.SNAKEMAKE_OPTS_ENVVAR}
        fi

        # Change back to working directory
        cd ${{CWD}}
        touch {output}
        """

rule prepare_run:
    input:
        "test_data/test_data.lock",
    output:
        SAMPLES = "{bench_dir}" / Path(config.input['samples_file']),
        LOCK = "{bench_dir}/prepare_run.lock",
    params:
        SAMPLES = config.input['samples_file'],
        DATADIR = config.output['datadir'],
        REF = reference_file,
    log:
        outfile = "logs/{bench_dir}_prepare_run.out.log",
    shell:
        """
        # create symbolic link of samples-file
        ln -fs ${{PWD}}/{params.SAMPLES} {output.SAMPLES}
        touch -h {output.SAMPLES}

        # copy reference file(s)
        echo "Migrating reference files" > {log.outfile}
        mkdir -p {wildcards.bench_dir}/references
        cp -va test_data/references/ {wildcards.bench_dir} >> {log.outfile}
        cp -va test_data/{params.REF} {wildcards.bench_dir}/{params.REF} >> {log.outfile}

        # copy raw reads
        echo "Migrating raw data" >> {log.outfile}
        cp -va test_data/{params.DATADIR} {wildcards.bench_dir} >> {log.outfile}

        touch {output.LOCK}
        """

rule prepare_config:
    input:
        CONFIG_FILE = "vpipe.config",
    output:
        CONFIG_FILE = "{bench_dir}/vpipe.config",
    run:
        # Copy the config file with the necesary changes
        with open(output.CONFIG_FILE, 'w') as configfile:
            configs[wildcards.bench_dir].write(configfile)


# TODO: handle special cases, e.g. dry-run (touch and clean files)
#       clean rules, and --create-envs-only
rule run_vpipeBench:
    input:
        "{bench_dir}/vpipe.config",
        "{bench_dir}/prepare_run.lock"
    output:
        "{bench_dir}/vpipeBench.lock"
    params:
        scratch = '1250',
        mem = config.run_vpipeBench['mem'],
        time = config.run_vpipeBench['time'],
        BASEDIR = VPIPE_BASEDIR,
        SNAKEMAKE_OPTIONS = lambda wildcards: config.benchmark['snakemake_options'],
        SNAKEMAKE_OPTS_ENVVAR = lambda wildcards: os.environ.get('SNAKEMAKE_OPTS', ''),
        SNAKEMAKE = config.applications['snakemake']
    log:
        outfile = "logs/{bench_dir}_vpipeBench.out.log",
        errfile = "logs/{bench_dir}_vpipeBench.err.log",
    threads:
        1
    shell:
        """
        # Change to working directory
        CWD=${{PWD}}
        cd {wildcards.bench_dir}/

        if [[ -z "{params.SNAKEMAKE_OPTS_ENVVAR}" ]]; then
            {params.SNAKEMAKE}  -s {params.BASEDIR}/vpipeBench.snake \
                {params.SNAKEMAKE_OPTIONS}
        else
            {params.SNAKEMAKE} -s {params.BASEDIR}/vpipeBench.snake \
                {params.SNAKEMAKE_OPTS_ENVVAR}
        fi

        # Change back to working directory
        cd ${{CWD}}
        touch {output}
        """

rule alignment_intervals:
    input:
        LOCK = expand("{bench_dir}/vpipeBench.lock", bench_dir=configs.keys()),
        REF = reference_file,
    output:
        "variants/union_coverage_intervals.tsv",
        "variants/intersect_coverage_intervals.tsv",
        "variants/union_coverage_intervals_ShoRAH.tsv",
        "variants/intersect_coverage_intervals_ShoRAH.tsv",
    params:
        scratch = '1250',
        mem = config.alignment_intervals['mem'],
        time = config.alignment_intervals['time'],
        WINDOW_LEN = window_lengths,
        WINDOW_SHIFT = shifts,
        OUTDIR = "variants",
        DIRS = BENCH_DIRS,
        ALIGNMENT_INTERVALS = config.applications['alignmentIntervals'],
    log:
        outfile = "logs/alignment_intervals.out.log",
        errfile = "logs/alignment_intervals.out.log",
    threads:
        1
    shell:
        """
        {params.ALIGNMENT_INTERVALS} -r {input.REF} -o {params.OUTDIR} \
            -wl {params.WINDOW_LEN} -ws {params.WINDOW_SHIFT} -d {params.DIRS} > >(tee -a {log.outfile}) 2>&1
        """

rule prepare_testing_run:
    input:
        "variants/{file}.tsv",
    output:
        "{bench_dir}/variants/{file}.tsv",
    shell:
        """
        ln -fs ${{PWD}}/{input} {output}
        touch -h {output}
        """

rule run_aligner_tests:
    input:
        "{bench_dir}/variants/union_coverage_intervals.tsv",
        "{bench_dir}/variants/intersect_coverage_intervals.tsv",
        "{bench_dir}/variants/union_coverage_intervals_ShoRAH.tsv",
        "{bench_dir}/variants/intersect_coverage_intervals_ShoRAH.tsv",
    output:
        "{bench_dir}/aligner_tests.lock"
    params:
        scratch = '1250',
        mem = config.run_tests['mem'],
        time = config.run_tests['time'],
        BASEDIR = VPIPE_BASEDIR,
        SNAKEMAKE_OPTIONS = lambda wildcards: config.benchmark['snakemake_options'],
        SNAKEMAKE_OPTS_ENVVAR = lambda wildcards: os.environ.get('SNAKEMAKE_OPTS', ''),
        SNAKEMAKE = config.applications['snakemake']
    log:
        outfile = "logs/{bench_dir}_aligner_tests.out.log",
        errfile = "logs/{bench_dir}_aligner_tests.err.log",
    threads:
        1
    shell:
        """
        # Change to working directory
        CWD=${{PWD}}
        cd {wildcards.bench_dir}/

        if [[ -z "{params.SNAKEMAKE_OPTS_ENVVAR}" ]]; then
            {params.SNAKEMAKE}  -s {params.BASEDIR}/vpipeBench.snake aligner_tests \
                {params.SNAKEMAKE_OPTIONS}
        else
            {params.SNAKEMAKE} -s {params.BASEDIR}/vpipeBench.snake aligner_tests \
                {params.SNAKEMAKE_OPTS_ENVVAR}
        fi

        # Change back to working directory
        cd ${{CWD}}
        touch {output}
        """

rule alignment_intervals_callers:
    input:
        LOCK = expand("{bench_dir}/vpipeBench.lock", bench_dir=configs.keys()),
        REF = reference_file,
    output:
        expand("variants/union_coverage_intervals_{aligner}.tsv", aligner=ALIGNERS),
        expand("variants/intersect_coverage_intervals_{aligner}.tsv", aligner=ALIGNERS)
    params:
        scratch = '1250',
        mem = config.alignment_intervals['mem'],
        time = config.alignment_intervals['time'],
        WINDOW_LEN = window_lengths,
        WINDOW_SHIFT = shifts,
        OUTDIR = "variants",
        DIRS = BENCH_DIRS,
        ALIGNMENT_INTERVALS = config.applications['alignmentIntervals'],
    log:
        outfile = "logs/alignment_intervals_callers.out.log",
        errfile = "logs/alignment_intervals_callers.out.log",
    threads:
        1
    shell:
        """
        {params.ALIGNMENT_INTERVALS} -r {input.REF} -o {params.OUTDIR} --caller \
            -wl {params.WINDOW_LEN} -ws {params.WINDOW_SHIFT} -d {params.DIRS} > >(tee -a {log.outfile}) 2>&1
        """

rule run_caller_tests:
    input:
        "{aligner}-{caller}/variants/union_coverage_intervals_{aligner}.tsv",
        "{aligner}-{caller}/variants/intersect_coverage_intervals_{aligner}.tsv",
    output:
        "{aligner}-{caller}/caller_tests.lock"
    params:
        scratch = '1250',
        mem = config.run_tests['mem'],
        time = config.run_tests['time'],
        BASEDIR = VPIPE_BASEDIR,
        SNAKEMAKE_OPTIONS = lambda wildcards: config.benchmark['snakemake_options'],
        SNAKEMAKE_OPTS_ENVVAR = lambda wildcards: os.environ.get('SNAKEMAKE_OPTS', ''),
        SNAKEMAKE = config.applications['snakemake']
    log:
        outfile = "logs/{aligner}-{caller}_caller_tests.out.log",
        errfile = "logs/{aligner}-{caller}_caller_tests.err.log",
    threads:
        1
    shell:
        """
        # Change to working directory
        CWD=${{PWD}}
        cd {wildcards.aligner}-{wildcards.caller}/

        if [[ -z "{params.SNAKEMAKE_OPTS_ENVVAR}" ]]; then
            {params.SNAKEMAKE}  -s {params.BASEDIR}/vpipeBench.snake snv_caller_tests \
                {params.SNAKEMAKE_OPTIONS}
        else
            {params.SNAKEMAKE} -s {params.BASEDIR}/vpipeBench.snake snv_caller_tests \
                {params.SNAKEMAKE_OPTS_ENVVAR}
        fi

        # Change back to working directory
        cd ${{CWD}}
        touch {output}
        """

rule simbenchclean:
    shell:
        """
        rm -rf test_data
        """
