# AccuSNV pipeline: preprocess -> map -> call candidate SNVs -> AccuSNV CNN.
# One workflow for both local and SLURM runs; only the launch profile differs.

import os
import re
import sys
import glob
import gzip
import pickle
import subprocess
import numpy as np
import pandas as pd
from itertools import compress
from snakemake.logging import logger
from snakemake.exceptions import WorkflowError

## AccuSNV pipeline imports
from accusnv import log as accusnv_log
from accusnv.preprocessing import utils


# Get conda env if specified
if config.get("env", ""):
    ENV = f"{config['env']} ; "
else:
    ENV = ""

aligner = config.get("aligner", "bwa")
# Cores for the two rules that can use more than one. On SLURM this is also what the executor
# requests as cpus-per-task, because it takes that from each rule's threads.
CORES = int(config.get("cores", 4))
IDX_SUFFIX = ".bwt" if aligner == "bwa" else ".1.bt2"
# samclip discards reads with soft-clipped ends, which bwa produces around indels and repeat edges
# and which turn into false SNVs. bowtie2 is run end-to-end, so it has nothing to clip.
CLIP = "" if config.get("skip_samclip") or aligner != "bwa" else " | samclip --ref {params.ref} --max 0"
outdir = config['outdir'].rstrip("/")

# --- Logging ---------------------------------------------------------------
# Each job writes its own pair of log files in <outdir>/logs, named for its rule and the sample or
# group it ran on, rather than every job appending to one shared pair. The run's two log files are
# the merge of those, made when the run finishes. Every rule starts its shell with job_log(), which
# names its files in $L; Python steps are then handed $L with LOG_ARGS and log themselves, and
# rules that only run external tools write a summary line with note() and send the tool's output
# to $L.full.log.

FULL_LOG = accusnv_log.paths(outdir, config)[1]
LOG_ARGS = '--log "$L"'
TOOL_OUT = '>> "$L.full.log" 2>&1'    # tool writes its data to files; all its output is chatter
TOOL_ERR = '2>> "$L.full.log"'        # tool writes its data to stdout, so only capture stderr

# The tags naming a job's log files, in the same form as its output filenames.
SAMPLE_TAG = "{wildcards.sampleID}_ref_{wildcards.reference}"
GROUP_TAG = "group_{wildcards.cladeID}"

def job_log(step, tag=None):
    '''Shell prefix putting this job's log files in $L, for LOG_ARGS, TOOL_OUT and note().'''
    return f'L="{accusnv_log.stem(outdir, step, tag)}" ; '

def note(message):
    '''Shell fragment that writes one summary-log line, for rules that only run external tools.'''
    return f'python -m accusnv.log "$L" "{message}" ; '

################# SAMPLE SHEET AND OUTPUTS #################

# --- Sample sheet -----------------------------------------------------------

# Read the sample sheet with Pandas. We will always reference sample info with this dataframe.
# Indexed on Sample because that is what the {sampleID} wildcard carries.

sample_table = pd.read_csv(config["sample_table"], keep_default_na=False).set_index("Sample", drop=False)
sample_table['Outgroup'] = sample_table['Outgroup'].astype(int)
sample_table['Group'] = sample_table['Group'].astype(str)
sample_groups = sorted(set(sample_table['Group']))

# A sample can be listed on several rows, e.g. as the outgroup of more than one group. Those rows
# all name the same reads, so the trimming and mapping rules run once per sample and reference.
def sample_row(sampleID):
    '''The first row for a sample. Its reads and Type are the same on all of them.'''
    return sample_table.loc[[sampleID]].iloc[0]

# The reference FASTA per reference name, so rules carrying the {reference} wildcard can look it
# up without going through a sample (a sample in two groups can have two references).
REF_FASTA = dict(zip(sample_table['Reference'], sample_table['Reference_FASTA']))

# --- Dynamic resources -----------------------------------------------------
# Every job starts in its first memory/runtime tier; a retry bumps it to the next tier up.
sys.path.insert(0, workflow.basedir)
import resources as res

def _genome_len(fasta):
    fai = fasta + '.fai'
    if os.path.exists(fai):
        return sum(int(l.split('\t')[1]) for l in open(fai))
    return sum(len(l.strip()) for l in open(fasta) if not l.startswith('>'))

GENOME_LEN = {f: _genome_len(f) for f in sample_table['Reference_FASTA'].unique()}


# --- Run banner ------------------------------------------------------------
# These hooks fire once in the main Snakemake process, so they describe the run as a whole. On
# SLURM each job re-reads this file, which is why the overview lives here and not at top level.

onstart:
    log = accusnv_log.setup('workflow', accusnv_log.stem(outdir, 'workflow'))
    log.info('Workflow starting: %d samples in %d group(s), aligner %s',
             len(sample_table), len(sample_groups), aligner)
    for clade in sample_groups:
        rows = sample_table[sample_table['Group'] == clade]
        ingroup = rows[rows['Outgroup'] == 0]
        ref = rows['Reference_FASTA'].iloc[0]
        log.info('Group %s: %d samples (%d ingroup, %d outgroup) against %s (%s bp)',
                 clade, len(rows), len(ingroup), len(rows) - len(ingroup),
                 rows['Reference'].iloc[0], f"{GENOME_LEN[ref]:,}")
        log.debug('Group %s reference FASTA: %s', clade, ref)
        log.debug('Group %s samples: %s', clade, ', '.join(rows['Sample']))

# Both hooks write to a second pair of files, so that their message lands at the end of the merged
# logs rather than back with the overview above. Both then merge the per-job logs, so a run started
# with snakemake directly still ends with the two log files; a run started with the accusnv CLI is
# merged again there, after its own last message.

onsuccess:
    accusnv_log.setup('workflow', accusnv_log.stem(outdir, 'workflow', 'finished')).info(
        'Workflow finished successfully.')
    accusnv_log.merge(outdir, config)

onerror:
    accusnv_log.setup('workflow', accusnv_log.stem(outdir, 'workflow', 'finished')).error(
        'Workflow failed. Each failed job is logged above with the reason and the tier it ran '
        'at; its tool output is in %s, and on its own in %s', FULL_LOG, os.path.join(outdir, 'logs'))
    accusnv_log.merge(outdir, config)


# --- Outputs ---------------------------------------------------------------


# One SNV table for each Group, plus whichever downstream analyses are switched on.
targets = expand(outdir + "/group_{cladeID}_snv_table_final.tsv", cladeID=sample_groups)
targets += expand(outdir + "/group_{cladeID}_snv_table_invariant_positions.tsv", cladeID=sample_groups)
if not config.get("downstream_only", False):
    # Reads the mapping-stage VCFs, which --downstream_only is meant not to touch.
    targets += expand(outdir + "/group_{cladeID}_snv_table_rejected_upstream.tsv", cladeID=sample_groups)

if not config['skip_all_downstream']:
    analysis_dir = outdir + "/3-Analysis/group_{cladeID}/"
    if not config['skip_dnds']:
        targets += expand(analysis_dir + "dNdS_out/data_dNdS.npz", cladeID=sample_groups)
    if not config['skip_report']:
        targets += expand(outdir + "/group_{cladeID}_snv_dashboard.html", cladeID=sample_groups)
    if not config['skip_trees']:
        targets += expand(outdir + "/group_{cladeID}_snv_tree_final.nwk.tree", cladeID=sample_groups)

rule all:
    input:
        targets


################# READ PREPROCESSING #################


def get_read_files(wildcards):
    '''Returns the R1 and R2 files (if present) for a given sample for cutadapt.'''

    sample_data = sample_row(wildcards.sampleID)
    if sample_data['Type'].upper() == 'PE':
        if not (sample_data['Read1_file_path'] and sample_data['Read2_file_path']):
            raise WorkflowError("Error: this row is marked as PE data but has no Read2 file: {}".format(wildcards.sampleID))
        return {"r1": sample_data['Read1_file_path'], "r2": sample_data['Read2_file_path']}
    else:
        logger.warning("Note: running this sample as single end: {}".format(wildcards.sampleID))
        if sample_data['Read2_file_path']:
            raise WorkflowError("Error: this sample is not marked as Type='PE', but it has a Read2 file: {}".format(wildcards.sampleID))
        return {"r1": sample_data['Read1_file_path']}

rule cutadapt:
    ''' Trim the 3' adapter from raw reads (single- or paired-end) with cutadapt. 
    cutadapt auto-detects gzip in input.'''

    input:
        unpack(get_read_files)
    params:
        adapter  = config.get("adapter_sequence", "CTGTCTCTTAT"),  # 3' adapter trimmed by cutadapt
    output:
        fq1 = outdir + "/1-Mapping/alignment/trimmed_filtered_reads/{sampleID}_R1_trimmed.fastq.gz",
        fq2 = outdir + "/1-Mapping/alignment/trimmed_filtered_reads/{sampleID}_R2_trimmed.fastq.gz"
    threads: CORES
    resources:
        mem_mb  = lambda wc, attempt: res.mem('cutadapt', attempt),
        runtime = lambda wc, attempt: res.runtime('cutadapt', attempt),
    run:
        layout = "paired-end" if len(input) == 2 else "single-end"
        L = job_log('cutadapt', wildcards.sampleID)
        shell(ENV + L + note(f"{wildcards.sampleID}: trimming adapter {params.adapter} from {layout} reads"))
        if len(input) == 2: # Paired end reads
            shell(ENV + L +
                "cutadapt -a {params.adapter} -A {params.adapter} --cores={threads} "
                "-o {output.fq1} -p {output.fq2} {input.r1} {input.r2} " + TOOL_OUT)
        else: # Single end reads
            shell(ENV + L +
                "cutadapt -a {params.adapter} --cores={threads} -o {output.fq1} {input.r1} " + TOOL_OUT)

            # Create empty r2 file for snakemake
            shell(ENV + "touch {output.fq2}")


def get_read_files_trimmed(wildcards):
    '''Returns the trimmed R1 and R2 files (if present) for a given sample for sickle.
    '''
    sample_data = sample_row(wildcards.sampleID)
    if sample_data['Type'].upper() == 'PE':
        return {"r1": rules.cutadapt.output.fq1, "r2": rules.cutadapt.output.fq2}
    else:
        return {"r1": rules.cutadapt.output.fq1}

rule sickle:
    ''' Quality-trim reads with sickle. 
    Currently, reads without a pair are not dropped, but sent to an _unpaired.fastq.gz file
    which is just for the user to examine (we do not map these).
    '''
    input:
        unpack(get_read_files_trimmed)
    params:
        qual = config.get("sickle_quality", 20),
        readlen = config.get("sickle_min_length", 50)
    output:
        fq1 = outdir + "/1-Mapping/alignment/trimmed_filtered_reads/{sampleID}_R1_filtered.fastq.gz",
        fq2 = outdir + "/1-Mapping/alignment/trimmed_filtered_reads/{sampleID}_R2_filtered.fastq.gz",
        fqS = outdir + "/1-Mapping/alignment/trimmed_filtered_reads/{sampleID}_unpaired.fastq.gz",
    resources:
        mem_mb  = lambda wc, attempt: res.mem('sickle', attempt),
        runtime = lambda wc, attempt: res.runtime('sickle', attempt),
    run:
        L = job_log('sickle', wildcards.sampleID)
        shell(ENV + L + note(f"{wildcards.sampleID}: quality trimming at Q{params.qual}, "
                             f"dropping reads under {params.readlen} bp"))
        if len(input) == 2: # Paired end reads
            shell(ENV + L +
            "sickle pe -f {input.r1} -r {input.r2} -t sanger -o {output.fq1} -p {output.fq2} -s {output.fqS}"
            " -g -q {params.qual} -l {params.readlen} -x -n " + TOOL_OUT)
            shell("  rm {input.r1} {input.r2}")

        else: #Single end reads
            shell(ENV + L +
            "sickle se -f {input.r1} -t sanger -o {output.fq1} -g -q {params.qual} -l {params.readlen} -x -n " + TOOL_OUT)

            # Create empty r2 file for snakemake
            shell(ENV + "   touch {output.fq2} {output.fqS}")
            shell(ENV + "   rm {input.r1}")

################# MAPPING #################


rule create_mapping_index:
    ''' creates the mapping index for the reference FASTA(s) if one does not exist
    '''
    input:
        fasta = "{fasta}"
    output:
        idx = "{fasta}" + IDX_SUFFIX
    resources:
        mem_mb  = lambda wc, attempt: res.mem('create_mapping_index', attempt),
        runtime = lambda wc, attempt: res.runtime('create_mapping_index', attempt),
    run:
        L = job_log('create_mapping_index', os.path.basename(input.fasta))
        shell(ENV + L + note(f"Building {aligner} index for {input.fasta}"))
        if aligner == "bwa":
            shell(ENV + L + "bwa index {input.fasta} " + TOOL_OUT)
        else:
            shell(ENV + L + "bowtie2-build -q {input.fasta} {input.fasta} " + TOOL_OUT)


rule mapping:
    '''Runs the alignment (either bwa or bowtie2) against the reference for each sample'''

    input:
        r1  = rules.sickle.output.fq1,
        r2  = rules.sickle.output.fq2,
        ref_index = lambda sample: REF_FASTA[sample.reference] + IDX_SUFFIX,
        # samclip reads the .fai; the pileup needs it later anyway, so always wait for it here.
        fai = lambda sample: REF_FASTA[sample.reference] + ".fai",
    params:
        ref    = lambda sample: REF_FASTA[sample.reference],
        paired = lambda sample: sample_row(sample.sampleID)["Type"].upper(),
        maxins = config.get("bowtie2_maxins", 2000),
    output:
        # temp(): SAMs are large and only sam2bam reads them, so Snakemake deletes each one as
        # soon as its BAM is made, and also if either rule fails partway.
        sam = temp(outdir + "/1-Mapping/alignment/{sampleID}_ref_{reference}_aligned.sam"),
    threads: CORES
    resources:
        mem_mb  = lambda wc, attempt: res.mem('mapping', attempt),
        runtime = lambda wc, attempt: res.runtime('mapping', attempt),
    run:
        L = job_log('mapping', f"{wildcards.sampleID}_ref_{wildcards.reference}")
        shell(ENV + L + note(f"{wildcards.sampleID}: mapping {params.paired} reads to "
                             f"{wildcards.reference} with {aligner} on {threads} threads"))
        if aligner == "bwa":
            if params.paired == 'PE':
                shell(ENV + L + "bwa mem -t {threads} {params.ref} {input.r1} {input.r2}" + CLIP + " > {output.sam} " + TOOL_ERR)
            else:
                shell(ENV + L + "bwa mem -t {threads} {params.ref} {input.r1}" + CLIP + " > {output.sam} " + TOOL_ERR)
        else: #bowtie2
            if params.paired == 'PE':
                shell(ENV + L + "bowtie2 --threads {threads} -X {params.maxins} --no-unal -x {params.ref} -1 {input.r1} -2 {input.r2} --no-mixed --dovetail -S {output.sam} " + TOOL_ERR)
            else:
                shell(ENV + L + "bowtie2 --threads {threads} -X {params.maxins} --no-unal -x {params.ref} -U {input.r1} -S {output.sam} " + TOOL_ERR)


################# CANDIDATE MUTATIONS #################
# This section converts and filters the mapped SAM to tabular input formats 
# for AccuSNV - this pipeline is largely derived from WideVariant.

rule sam2bam:
    ''' Convert sam to a sorted, indexed bam with duplicates removed.'''

    input:
        sam = rules.mapping.output.sam
    params:
        optdist = config.get("markdup_optical_distance", 100),
        # samtools -@ counts threads beyond the first, so pass one less than we asked SLURM for.
        extra_threads = lambda wc, threads: threads - 1,
        # sort -m is per thread; the half leaves room for sort overshooting -m and markdup's buffer.
        sortmem = lambda wc, threads, resources: f"{int(0.5 * resources.mem_mb / threads)}M",
    output:
        # The three intermediate BAMs are outputs rather than params so that Snakemake deletes
        # them when the rule finishes and also when it fails partway; the stats file is kept.
        bam     = temp(outdir + "/1-Mapping/alignment/{sampleID}_ref_{reference}.bam"),
        mate    = temp(outdir + "/1-Mapping/alignment/{sampleID}_ref_{reference}.fixmates.bam"),
        matesrt = temp(outdir + "/1-Mapping/alignment/{sampleID}_ref_{reference}.sortedmates.bam"),
        # samtools sort names its own spill files and decides at run time how many to write, so they
        # go in a directory Snakemake owns instead. A killed job used to leave them next to the BAMs,
        # where its retry could not overwrite them and silently merged the stale ones instead.
        sorttmp = temp(directory(outdir + "/1-Mapping/alignment/{sampleID}_ref_{reference}.sorttmp")),
        stats   = outdir + "/1-Mapping/alignment/{sampleID}_ref_{reference}.bam.stats.txt",
        final_bam = outdir + "/1-Mapping/alignment/{sampleID}_ref_{reference}_aligned.sorted.bam",
        bai = outdir + "/1-Mapping/alignment/{sampleID}_ref_{reference}_aligned.sorted.bam.bai",
    threads: CORES
    resources:
        mem_mb  = lambda wc, attempt: res.mem('sam2bam', attempt),
        runtime = lambda wc, attempt: res.runtime('sam2bam', attempt),
    shell:
        ENV + job_log('sam2bam', SAMPLE_TAG) +
        note("{wildcards.sampleID}: sorting alignments and marking duplicates "
             "(optical distance {params.optdist}) on {threads} threads") +
        # -u keeps the piped BAM uncompressed, so view stays cheap and sort gets the threads.
        "mkdir -p {output.sorttmp} && "
        "samtools view -uS {input.sam} | samtools sort -@ {params.extra_threads} -m {params.sortmem} -n -T {output.sorttmp}/byname - -o {output.bam} " + TOOL_OUT + " && "
        "samtools fixmate -@ {params.extra_threads} -m {output.bam} {output.mate} " + TOOL_OUT + " && "
        "samtools sort -@ {params.extra_threads} -m {params.sortmem} -T {output.sorttmp}/bypos -o {output.matesrt} {output.mate} " + TOOL_OUT + " && "
        "samtools markdup -@ {params.extra_threads} -r -s -f {output.stats} -d {params.optdist} -m s {output.matesrt} {output.final_bam} " + TOOL_OUT + " && "
        "samtools index -@ {params.extra_threads} {output.final_bam} {output.bai} " + TOOL_OUT


rule samtools_idx:
    ''' Create the .fai index samtools needs to read the reference FASTA during pileup.'''
    input:
        fasta = "{fasta}",
    params:
        # The wildcard here is a path, so the log files are named after the FASTA alone.
        tag = lambda sample: os.path.basename(sample.fasta),
    output:
        fai = "{fasta}.fai",
    shell:
        ENV + job_log('samtools_idx', "{params.tag}") +
        note("Indexing reference FASTA {input.fasta}") +
        "samtools faidx {input.fasta} " + TOOL_OUT + " ;"


rule mpileup2vcf:
    # Produce a pileup and a SNP-only variant VCF for one sample with samtools/bcftools.
    input:
        bam     = rules.sam2bam.output.final_bam,
        fai     = lambda sample: REF_FASTA[sample.reference] + ".fai",
    params:
        ref      = lambda sample: REF_FASTA[sample.reference],
        minmapq  = config.get("mpileup_min_mapq", 30),
        maxdepth = config.get("mpileup_max_depth", 3000),
        minaf    = config.get("variant_min_af", 0.75),
    output:
        pileup   = outdir + "/1-Mapping/vcf/{sampleID}_ref_{reference}.pileup",
        variants = outdir + "/1-Mapping/vcf/{sampleID}_ref_{reference}.variant.vcf.gz",
        strain   = outdir + "/1-Mapping/vcf/{sampleID}_ref_{reference}.strain.vcf.gz",
    resources:
        mem_mb  = lambda wc, attempt: res.mem('mpileup2vcf', attempt),
        runtime = lambda wc, attempt: res.runtime('mpileup2vcf', attempt),
    shell:
        ENV + job_log('mpileup2vcf', SAMPLE_TAG) +
        note("{wildcards.sampleID}: calling variants (min mapping quality {params.minmapq}, "
             "max depth {params.maxdepth}, min alt-allele fraction {params.minaf})") +
        "samtools mpileup -q{params.minmapq} -x -s -O -d{params.maxdepth} -f {params.ref} {input.bam} > {output.pileup} " + TOOL_ERR + " && "
        "bcftools mpileup -q{params.minmapq} --annotate FORMAT/SP -d{params.maxdepth} -f {params.ref} {input.bam} -Ou " + TOOL_ERR + " | bcftools call -c -Oz -o {output.strain} " + TOOL_ERR + " && "
        "bcftools view -Oz -v snps -q {params.minaf} {output.strain} > {output.variants} " + TOOL_ERR + " && "
        "tabix -p vcf {output.variants} " + TOOL_OUT

rule vcf2quals:
    ''' Reads the strain VCF and creates an FQ quality score vector
    for all positions in that genome. '''
    input:
        strain = rules.mpileup2vcf.output.strain
    params:
        ref = lambda sample: REF_FASTA[sample.reference]
    output:
        quals = outdir + "/1-Mapping/quals/{sampleID}_ref_{reference}.quals.pickle.gz"
    resources:
        mem_mb  = lambda wc, attempt: res.mem('vcf2quals', attempt),
        runtime = lambda wc, attempt: res.runtime('vcf2quals', attempt),
    shell:
        ENV + job_log('vcf2quals', SAMPLE_TAG) +
        "python -m accusnv.preprocessing.vcf2quals_snakemake -i {input.strain} -r {params.ref} -o {output.quals} "
        "--sample {wildcards.sampleID} " + LOG_ARGS

rule variants2positions:
    ''' Reads the variants VCF and creates a vector of bcftools SNP positions.'''
    input:
        variants = rules.mpileup2vcf.output.variants
    params:
        ref = lambda sample: REF_FASTA[sample.reference],
        maxfq = config.get("max_fq", -30)
    output:
        positions = outdir + "/1-Mapping/quals/{sampleID}_ref_{reference}.positions.pickle"
    resources:
        mem_mb  = lambda wc, attempt: res.mem('variants2positions', attempt),
        runtime = lambda wc, attempt: res.runtime('variants2positions', attempt),
    shell:
        ENV + job_log('variants2positions', SAMPLE_TAG) +
        "python -m accusnv.preprocessing.variants2positions -i {input.variants} -o {output.positions} "
        "-r {params.ref} -q {params.maxfq} --sample {wildcards.sampleID} " + LOG_ARGS + " ;"

rule pileup2diversity:
    ''' Creates the per-position diversity array for the whole genome, then deletes the pileup file.
    Per-base coverage is not written here: nothing read it, and the same numbers (correctly summed
    over all samples) are in group_<G>_coverage_matrix_raw.npz.
        coverage = outdir + "/1-Mapping/diversity/{sampleID}_ref_{reference}.coverage.pickle.gz",
    '''
    input:
        pileup = rules.mpileup2vcf.output.pileup,
    params:
        ref = lambda sample: REF_FASTA[sample.reference],
    output:
        diversity = outdir + "/1-Mapping/diversity/{sampleID}_ref_{reference}.diversity.pickle.gz",
    resources:
        mem_mb  = lambda wc, attempt: res.mem('pileup2diversity', attempt),
        runtime = lambda wc, attempt: res.runtime('pileup2diversity', attempt),
    run:
        shell(ENV + job_log('pileup2diversity', SAMPLE_TAG) +
        "python -m accusnv.preprocessing.pileup2diversity -i {input.pileup} -r {params.ref} "
        "-o {output.diversity} --sample {wildcards.sampleID} " + LOG_ARGS + " ; ")

        ## if successful, delete the big pileup file
        if os.path.exists(output.diversity) and os.path.getsize(output.diversity) > 0:
            os.remove(input.pileup)

def combine_positions(positions_files, output_p_file, ref_genome, cladeID):
    ''' Function called by combine_positions rule to create clade-wide positions lists.'''
    log = accusnv_log.setup('combine_positions', accusnv_log.stem(outdir, 'combine_positions', f'group_{cladeID}'))
    log.info('Group %s: merging candidate variant positions from %d ingroup samples',
             cladeID, len(positions_files))
    chr_starts = utils.genomestats(ref_genome)[0]
    positions = np.array([], dtype=int)
    for path in positions_files:
        with gzip.open(path, "rb") as f:
            chrpos = pickle.load(f)
        log.debug('%s contributed %d variant positions', os.path.basename(path), chrpos.shape[0])
        if chrpos.shape[0]:
            if len(chr_starts) == 1:
                idx = chrpos[:, 1]
            else:
                idx = chr_starts[chrpos[:, 0] - 1] + chrpos[:, 1]
            positions = np.union1d(positions, idx)
    with open(output_p_file, "wb") as wf:
        pickle.dump(positions, wf)
    log.info('Group %s: %d distinct candidate SNV positions to genotype in every sample',
             cladeID, len(positions))

def clade_paths(cladeID, pattern, ingroup_only=False):
    '''All per-sample paths for one clade, from an upstream rule's output pattern.

    ingroup_only drops the outgroup samples. Only the candidate position list wants that:
    outgroup samples do not contribute candidate positions, so asking for theirs would run
    variants2positions on them just to get a file nothing reads.'''
    rows = sample_table[sample_table['Group'] == cladeID]
    if ingroup_only:
        rows = rows[rows['Outgroup'] == 0]

    return expand(pattern, zip,
        sampleID=rows['Sample'].tolist(),
        reference=rows['Reference'].tolist())

def clade_reference(cladeID):
    '''The reference FASTA for one clade. Every sample in a clade shares it.'''
    return sample_table[sample_table['Group'] == cladeID]['Reference_FASTA'].iloc[0]

rule upstream_rejects:
    ''' Records the positions bcftools called but dropped before candidate selection, and
    which of the two gates dropped them. Nothing reads this; it exists so those removals are
    not invisible. '''
    input:
        strain   = rules.mpileup2vcf.output.strain,
        variants = rules.mpileup2vcf.output.variants,
    params:
        maxfq = config.get("max_fq", -30),
    output:
        rejects = outdir + "/1-Mapping/vcf/{sampleID}_ref_{reference}.upstream_rejects.tsv",
    resources:
        mem_mb  = lambda wc, attempt: res.mem('variants2positions', attempt),
        runtime = lambda wc, attempt: res.runtime('variants2positions', attempt),
    shell:
        ENV + job_log('upstream_rejects', SAMPLE_TAG) +
        "python -m accusnv.preprocessing.upstream_rejects --strain {input.strain} "
        "--variant {input.variants} -q {params.maxfq} -o {output.rejects} "
        "--sample {wildcards.sampleID} " + LOG_ARGS + " ;"


rule combine_upstream_rejects:
    ''' One table per group of everything dropped before candidate selection. '''
    input:
        rejects = lambda sample: clade_paths(sample.cladeID, rules.upstream_rejects.output.rejects, ingroup_only=True),
    output:
        combined = outdir + "/group_{cladeID}_snv_table_rejected_upstream.tsv",
    resources:
        mem_mb  = lambda wc, attempt: res.mem('combine_upstream_rejects', attempt),
        runtime = lambda wc, attempt: res.runtime('combine_upstream_rejects', attempt),
    shell:
        ENV + job_log('combine_upstream_rejects', GROUP_TAG) +
        "python -m accusnv.preprocessing.upstream_rejects --combine {input.rejects} "
        "-o {output.combined} --group {wildcards.cladeID} " + LOG_ARGS + " ;"


rule combine_positions:
    ''' Merge per-sample variants into a combined list per each clade. '''
    input:
        positions = lambda sample: clade_paths(sample.cladeID, rules.variants2positions.output.positions, ingroup_only=True),
    params:
        ref      = lambda sample: clade_reference(sample.cladeID),
    output:
        allpositions = outdir + "/2-SNV-filtering/raw_tables/group_{cladeID}_allpositions.pickle",
    resources:
        mem_mb  = lambda wc, attempt: res.mem('combine_positions', attempt),
        runtime = lambda wc, attempt: res.runtime('combine_positions', attempt),
    run:
        combine_positions(input.positions, output.allpositions, params.ref, wildcards.cladeID)

rule candidate_mutation_table:
    ''' Combines all sample to creates three tables for each group (clade) in the set.
    cmt: Candidate mutation table.
    cov_raw: coverage information for all positions in the genome.
    cov_norm: z-score normalized coverage for all positions in the genome.
    '''
    input:
        allpositions = rules.combine_positions.output.allpositions,
        diversity = lambda sample: clade_paths(sample.cladeID, rules.pileup2diversity.output.diversity),
        quals  = lambda sample: clade_paths(sample.cladeID, rules.vcf2quals.output.quals),
    params:
        names    = lambda sample: sample_table[sample_table['Group'] == sample.cladeID]['Sample'].tolist(),
        outgroup = lambda sample: sample_table[sample_table['Group'] == sample.cladeID]['Outgroup'].tolist(),
    output:
        cmt      = outdir + "/2-SNV-filtering/raw_tables/group_{cladeID}_candidate_mutation_table.npz",
        cov_raw  = outdir + "/2-SNV-filtering/raw_tables/group_{cladeID}_coverage_matrix_raw.npz",
        cov_norm = outdir + "/2-SNV-filtering/raw_tables/group_{cladeID}_coverage_matrix_norm.npz",
    resources:
        mem_mb  = lambda wc, attempt: res.mem('candidate_mutation_table', attempt),
        runtime = lambda wc, attempt: res.runtime('candidate_mutation_table', attempt),
    run:
        from accusnv.preprocessing.build_candidate_mutation_table import build_candidate_mutation_table
        accusnv_log.setup('candidate_mutation_table',
                          accusnv_log.stem(outdir, 'candidate_mutation_table', f'group_{wildcards.cladeID}'))
        build_candidate_mutation_table(input.allpositions, params.names, params.outgroup, input.quals, input.diversity,
                                       output.cmt, output.cov_raw, output.cov_norm, True, True, 40,
                                       wildcards.cladeID)


################# AccuSNV CNN #################
min_cov_samp = config["min_cov_samp"]
min_cov_filt = config["min_cov_filt"]
exclude_samp = config["exclude_samp"]
recomb_flag  = "0" if config["skip_recombination"] else "1"
call_min_major_allele_freq     = config.get("call_min_major_allele_freq", 0.85)
call_min_qual                  = config.get("call_min_qual", 30)
call_max_indel_frac            = config.get("call_max_indel_frac", 0.33)
min_mut_qual                   = config.get("min_mut_qual", 1)
max_frac_ambiguous_samples     = config.get("max_frac_ambiguous_samples", 1)
min_median_coverage_position   = config.get("min_median_coverage_position", 5)
max_mean_copynum               = config.get("max_mean_copynum", 4)
max_max_copynum                = config.get("max_max_copynum", 7)
fast_mode_positions            = config.get("fast_mode_positions", 100000)
rebuild_sample_count           = config.get("rebuild_sample_count", 20)
rebuild_cutoff_many            = config.get("rebuild_cutoff_many", 0.1)
rebuild_cutoff_few             = config.get("rebuild_cutoff_few", 0.25)
recomb_distance_bp             = config.get("recomb_distance_bp", 1000)
recomb_corr_threshold          = config.get("recomb_corr_threshold", 0.75)
annotate_min_major_allele_freq = config.get("annotate_min_major_allele_freq", 0.75)

# Positions the user wants left out of, or kept in, the SNV set whatever the filters decided.
# Applied by the annotation, tree and dashboard stages rather than by the calling stage, so that
# they take effect in --downstream_only runs too, where the calling stage does not re-run.
exclude_positions = config.get("exclude_positions", "")
include_positions = config.get("include_positions", "")
position_flags = f'--exclude_positions "{exclude_positions}" ' if exclude_positions else ""
position_flags += f'--include_positions "{include_positions}" ' if include_positions else ""

# --downstream_only rebuilds only the annotation + evolutionary analyses from the SNV calls
# already in the output dir. ancient() marks those calls as given, so Snakemake takes them off
# disk instead of walking back into the mapping and calling stages to regenerate them. Anything
# genuinely missing is still built, which is what makes it useful for adding an analysis an
# earlier run skipped. .get so older pipeline.yaml files without the key still load.
# In --downstream_only mode the calling stage must not be re-run, so annotation's input is
# marked ancient(). Only that boundary is marked: the links between the downstream stages stay
# normal, so re-annotating cascades into dN/dS, the tree and the dashboard.
reuse = ancient if config.get("downstream_only", False) else (lambda path: path)

rule calling_accusnv:
    # Run the AccuSNV CNN + filters, select candidate SNVs, and flag potentially recombinant SNVs.
    input:
        cmt     = rules.candidate_mutation_table.output.cmt,
        cov_raw = rules.candidate_mutation_table.output.cov_raw,
    params:
        ref = lambda sample: clade_reference(sample.cladeID),
        out = outdir + "/2-SNV-filtering/group_{cladeID}/",
    output:
        state    = outdir + "/2-SNV-filtering/group_{cladeID}/_snv_state.npz",
        npz      = outdir + "/2-SNV-filtering/group_{cladeID}/candidate_mutation_table_final.npz",
        invariant = outdir + "/2-SNV-filtering/group_{cladeID}/snv_table_invariant_positions.tsv",
        deliverable_invariant = outdir + "/group_{cladeID}_snv_table_invariant_positions.tsv",
    resources:
        mem_mb  = lambda wc, attempt: res.mem('calling_accusnv', attempt),
        runtime = lambda wc, attempt: res.runtime('calling_accusnv', attempt),
    shell:
        ENV + job_log('calling_accusnv', GROUP_TAG) + "python -m accusnv.accusnv.accusnv "
            "-i {input.cmt} -c {input.cov_raw} -r {params.ref} -o {params.out} "
            "-s {min_cov_samp} -v {min_cov_filt} -e {exclude_samp} -m {recomb_flag} "
            "--min_maf_call {call_min_major_allele_freq} "
            "--min_qual_call {call_min_qual} "
            "--max_indel_frac {call_max_indel_frac} "
            "--min_mut_qual {min_mut_qual} "
            "--max_frac_ambiguous_samples {max_frac_ambiguous_samples} "
            "--min_median_coverage_position {min_median_coverage_position} "
            "--max_mean_copynum {max_mean_copynum} "
            "--max_max_copynum {max_max_copynum} "
            "--fast_mode_positions {fast_mode_positions} "
            "--rebuild_sample_count {rebuild_sample_count} "
            "--rebuild_cutoff_many {rebuild_cutoff_many} "
            "--rebuild_cutoff_few {rebuild_cutoff_few} "
            "--recomb_distance {recomb_distance_bp} "
            "--recomb_corr {recomb_corr_threshold} "
            "--group group_{wildcards.cladeID} " + LOG_ARGS + " && "
            "cp {output.invariant} {output.deliverable_invariant} ;"

rule annotate_snvs:
    # Annotate mutations by type of mutation (N vs S and AA changes).
    input:
        state = reuse(rules.calling_accusnv.output.state),
    params:
        ref = lambda sample: clade_reference(sample.cladeID),
        out = outdir + "/2-SNV-filtering/group_{cladeID}/",
    output:
        unfiltered = outdir + "/2-SNV-filtering/group_{cladeID}/snv_table_unfiltered.tsv",
        final = outdir + "/2-SNV-filtering/group_{cladeID}/snv_table_final.tsv",
        deliverable_unfiltered = outdir + "/group_{cladeID}_snv_table_unfiltered.tsv",
        deliverable_final = outdir + "/group_{cladeID}_snv_table_final.tsv",
    resources:
        mem_mb  = lambda wc, attempt: res.mem('annotate_snvs', attempt),
        runtime = lambda wc, attempt: res.runtime('annotate_snvs', attempt),
    # && so that a failed stage fails the rule: the copies would otherwise succeed on
    # whatever is on disk and mask it.
    shell:
        ENV + job_log('annotate_snvs', GROUP_TAG) +
        "python -m accusnv.downstream.annotate -r {params.ref} -o {params.out} "
        "--maf {annotate_min_major_allele_freq} --group group_{wildcards.cladeID} "
        + position_flags + LOG_ARGS + " && "
        "cp {output.unfiltered} {output.deliverable_unfiltered} && "
        "cp {output.final} {output.deliverable_final} ;"


################# EVOLUTIONARY ANALYSES #################
# These perform downstream evolutionary analyses, 
# and read the group's tables from 2-SNV-filtering and write their results to 3-Analysis.

tree_flags = " ".join(flag for flag, on in (
    ("--skip_trees", config['skip_trees']),
    ("--use_nj_tree", config.get('use_nj_tree', False)),  # .get: absent from older pipeline.yaml files
    ("--build_snv_trees", config['build_snv_trees'])) if on)

rule dnds:
    # Calcualte genome-wide dN/dS from the annotated final table.
    input:
        final = rules.annotate_snvs.output.final,
    params:
        ref = lambda sample: clade_reference(sample.cladeID),
        source = outdir + "/2-SNV-filtering/group_{cladeID}/",
        out    = outdir + "/3-Analysis/group_{cladeID}/",
    output:
        dnds = outdir + "/3-Analysis/group_{cladeID}/dNdS_out/data_dNdS.npz",
    resources:
        mem_mb  = lambda wc, attempt: res.mem('dnds', attempt),
        runtime = lambda wc, attempt: res.runtime('dnds', attempt),
    shell:
        ENV + job_log('dnds', GROUP_TAG) +
        "python -m accusnv.downstream.dnds -r {params.ref} -i {params.source} -o {params.out} "
        "--group group_{wildcards.cladeID} " + LOG_ARGS + " ;"

rule build_tree:
    # Build parsimony tree (dnapars), the optional per-SNV trees, and the dMRCA tables.
    input:
        unfiltered = rules.annotate_snvs.output.unfiltered,
        final = rules.annotate_snvs.output.final,
    params:
        ref = lambda sample: clade_reference(sample.cladeID),
        source = outdir + "/2-SNV-filtering/group_{cladeID}/",
        out    = outdir + "/3-Analysis/group_{cladeID}/",
    output:
        # Groups with no SNVs will have an empty tree file.
        tree = outdir + "/3-Analysis/group_{cladeID}/phylogeny/snv_tree_final.nwk.tree",
        deliverable = outdir + "/group_{cladeID}_snv_tree_final.nwk.tree",
    resources:
        mem_mb  = lambda wc, attempt: res.mem('build_tree', attempt),
        runtime = lambda wc, attempt: res.runtime('build_tree', attempt),
    shell:
        ENV + job_log('build_tree', GROUP_TAG) +
        "python -m accusnv.downstream.tree_building -r {params.ref} -i {params.source} "
        "-o {params.out} -g group_{wildcards.cladeID} " + tree_flags + " " + position_flags + LOG_ARGS + " && "
        "cp {output.tree} {output.deliverable} ;"

# The dashboard shows the tree when there is one; with trees switched off it just leaves that
# panel out, so the tree is an input only when the run is actually building it.
tree_for_report = [] if config['skip_trees'] else [rules.build_tree.output.tree]

rule report_html:
    # Build the interactive SNV dashboard (filters, per-SNV evidence, gene lookups, tree) and its
    # per-SNV bar chart images.
    input:
        final = rules.annotate_snvs.output.final,
        tree = tree_for_report,
    params:
        ref = lambda sample: clade_reference(sample.cladeID),
        source = outdir + "/2-SNV-filtering/group_{cladeID}/",
        out    = outdir + "/3-Analysis/group_{cladeID}/",
        tree   = outdir + "/3-Analysis/group_{cladeID}/phylogeny/snv_tree_final.nwk.tree",
    output:
        dashboard = outdir + "/3-Analysis/group_{cladeID}/snv_dashboard.html",
        deliverable_dashboard = outdir + "/group_{cladeID}_snv_dashboard.html",
    resources:
        mem_mb  = lambda wc, attempt: res.mem('report_html', attempt),
        runtime = lambda wc, attempt: res.runtime('report_html', attempt),
    shell:
        ENV + job_log('report_html', GROUP_TAG) +
        "python -m accusnv.downstream.generate_dashboard -r {params.ref} -i {params.source} "
        "-o {params.out} -g group_{wildcards.cladeID} -t {params.tree} " + position_flags + LOG_ARGS + " && "
        "cp {output.dashboard} {output.deliverable_dashboard} ;"
