# AccuSNV pipeline
# Steps are: 
# 1. preprocess then map reads
# 2. read pileups and call candidate SNVs 
# 3. AccuSNV CNN classification
# 4. Annotate, calculate dNdS, create tree, create dashboard

import os
import sys
import gzip
import pickle
import numpy as np
import pandas as pd
from snakemake.logging import logger
from snakemake.exceptions import WorkflowError

## AccuSNV pipeline imports
from accusnv import log as accusnv_log


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

aligner = config.get("aligner", "bwa")
# Mumber of for the two rules that use more than 1
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. 

FULL_LOG = accusnv_log.paths(outdir, config)[1]
LOG_ARGS = '--log "$L"'
TOOL_OUT = '>> "$L.full.log" 2>&1'    
TOOL_ERR = '2>> "$L.full.log"'        # when tools write data to stdout, 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 for this job's log files.'''
    return f'L="{accusnv_log.stem(outdir, step, tag)}" ; '

def note(message):
    '''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. Retrying 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('>'))


# --- Run ------------------------------------------------------------
# These commands are run 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']))

# onsuccess and onerror write to a second pair of files 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 still ends with the two log files.

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 + " && "
        # New versions of bcftools dropped a default -Q13, so having this restores old WideVariant
        # BQ filtering.
        "bcftools mpileup -q{params.minmapq} -Q13 --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.'''

    from accusnv.preprocessing import utils
    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)
contig_edge_bp                 = config.get("contig_edge_bp", 100)
max_edge_strand_imbalance      = config.get("max_edge_strand_imbalance", 0.3)
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} "
            "--contig_edge_bp {contig_edge_bp} "
            "--max_edge_strand_imbalance {max_edge_strand_imbalance} "
            "--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} ;"
