Snakefile_WHO

configfile: "config/config.json"

passages = ['cell', 'egg']
assays = ['hi', 'fra']
centers = ['who', 'cdc', 'crick', 'niid', 'vidrl']
resolutions = ['2y', '6y']
lineages = ['h3n2' , 'h1n1pdm', 'vic', 'yam']
segments = ['ha', 'na']

include: "Snakefile_base"
localrules: all_builds, download_all, download_titers, download_sequences


def all_builds(w):
    builds = []
    for lineage in lineages:
        assays = ['hi', 'fra'] if lineage=='h3n2' else ['hi']
        for file_suffix in ['tree', 'meta', 'frequencies', 'titers', 'titer-tree-model', 'titer-sub-model', 'entropy', 'sequences']:
           builds.extend(
                expand("auspice-who/flu_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}_{f}.json",
                        center=centers, lineage=lineages, segment=segments,
                        resolution=resolutions, assay=assays, passage=passages, f=[file_suffix]))
    return builds


def region_translations(w):
    genes = gene_names(w)
    return ["results/full-aaseq-%s_%s_%s_%s_%s.fasta"%(g, w.region, w.lineage, w.segment, w.resolution)
            for g in genes]

rule all_who:
    input:
        all_builds

# separate rule for interaction with fauna
rule download_all:
    input:
        titers = expand("data/{center}_{lineage}_{passage}_{assay}_titers.tsv",
                         lineage=lineages, center=centers, assay=assays, passage=passages),
        sequences = expand("data/{lineage}_{segment}.fasta", lineage=lineages, segment=segments)


for seg, genes in genes_to_translate.items():
    rule:
        input:
            metadata = rules.parse.output.metadata,
            sequences = rules.parse.output.sequences,
            exclude = files.outliers,
            reference = files.reference
        params:
            genes=genes,
            region="{region}"
        output:
            alignments = expand("results/full-aaseq-{gene}_{{region}}_{{lineage}}_{{segment}}_{{resolution}}.fasta",
                                gene=genes)
        shell:
            """
            python3 scripts/full_region_alignments.py  --sequences {input.sequences}\
                                                 --metadata {input.metadata} \
                                                 --exclude {input.exclude} \
                                                 --genes {params.genes} \
                                                 --region {params.region:q} \
                                                 --resolution {wildcards.resolution} \
                                                 --reference {input.reference} \
                                                 --output {output.alignments:q}
            """


rule complete_mutation_frequencies_by_region:
    input:
        metadata = rules.parse.output.metadata,
        alignment = region_translations
    params:
        genes = gene_names,
        min_date = min_date,
        max_date = max_date,
        min_freq = 0.003,
        pivot_interval = pivot_interval,
        stiffness = 20,
        inertia = 0.2
    output:
        mut_freq = "results/mutation_frequencies_{region}_{lineage}_{segment}_{resolution}.json"
    shell:
        """
        augur frequencies --method diffusion \
                          --alignments {input.alignment:q} \
                          --metadata {input.metadata} \
                          --gene-names {params.genes} \
                          --pivot-interval {params.pivot_interval} \
                          --stiffness {params.stiffness} \
                          --inertia {params.inertia} \
                          --ignore-char X \
                          --min-date {params.min_date} \
                          --max-date {params.max_date} \
                          --minimal-frequency {params.min_freq} \
                          --output {output.mut_freq:q}
        """

rule global_mutation_frequencies:
    input:
        frequencies = expand("results/mutation_frequencies_{region}_{{lineage}}_{{segment}}_{{resolution}}.json",
                             region = frequency_regions),
        tree_freq = rules.tree_frequencies.output
    params:
        regions = frequency_regions
    output:
        auspice="auspice-who/flu_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}_frequencies.json",
        augur="results/flu_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}_frequencies.json"
    shell:
        '''
        python3 scripts/global_frequencies.py --region-frequencies {input.frequencies:q} \
                                              --tree-frequencies {input.tree_freq} \
                                              --regions {params.regions:q} \
                                              --output-auspice {output.auspice} \
                                              --output-augur {output.augur}
        '''


rule scores:
    input:
        metadata = rules.parse.output.metadata,
        tree = rules.refine.output.tree
    output:
        scores = "results/scores_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}.json"
    shell:
        """
        python3 scripts/scores.py  --metadata {input.metadata} \
                                  --tree {input.tree} \
                                  --output {output}
        """


rule export_titers:
    input:
        sub = rules.titers_sub.output.titers_model,
        tree = rules.titers_tree.output.titers_model
    output:
        raw = "auspice-who/flu_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}_titers.json",
        tree = "auspice-who/flu_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}_titer-tree-model.json",
        sub = "auspice-who/flu_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}_titer-sub-model.json",
    run:
        import json
        with open(input.sub) as fh:
            sub = json.load(fh)

        with open(output.sub, 'wt') as sub_file:
            json.dump({'avidity': sub['avidity'],
                       'potency': sub['potency'],
                       'substitution': sub['substitution']},
                      sub_file, indent=1)

        with open(output.raw, 'wt') as raw_file:
            json.dump(sub['titers'], raw_file, indent=1)

        with open(input.tree) as fh:
            tree = json.load(fh)

        with open(output.tree, 'wt') as tree_file:
            json.dump({'avidity': tree['avidity'],
                       'potency': tree['potency'],
                       'dTiter': {k:v['dTiter'] for k,v in tree['nodes'].items()}},
                      tree_file, indent=1)

rule export_entropy:
    input:
        aln = rules.align.output.alignment,
        reference = files.reference,
    params:
        genes = gene_names
    output:
        "auspice-who/flu_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}_entropy.json",
    shell:
        '''
        python3 scripts/entropy.py --alignment {input.aln} \
                --genes {params.genes} \
                --reference-sequence {input.reference} \
                --output {output}
        '''

rule export_sequence_json:
    input:
        aln = rules.ancestral.output.node_data,
        tree = rules.refine.output.tree,
        aa_seqs = translations
    params:
        genes = gene_names
    output:
        "auspice-who/flu_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}_sequences.json",
    shell:
        '''
        python3 scripts/sequence_export.py --alignment {input.aln} \
                --genes {params.genes} \
                --tree {input.tree} \
                --translations {input.aa_seqs} \
                --output {output}
        '''


def _get_node_data_for_report_export(wildcards):
    """Return a list of node data files to include for a given build's wildcards.
    """
    # Define inputs shared by all builds.
    inputs = [
        rules.annotate_recency_of_submissions.output.node_data,
        rules.refine.output.node_data,
        rules.ancestral.output.node_data,
        rules.translate.output.node_data,
        rules.titers_tree.output.titers_model,
        rules.titers_sub.output.titers_model,
        rules.clades.output.clades,
        rules.traits.output.node_data,
        rules.lbi.output.lbi,
        rules.scores.output.scores,
    ]

    # glycosilation only makes sense for surface proteins
    if wildcards.segment in ['ha', 'na']:
        inputs.append(rules.glyc.output.glyc)

    # Only request a distance file for builds that have mask configurations
    # defined.
    if _get_build_distance_map_config(wildcards) is not None:
        inputs.append(rules.distances.output.distances)

    # Convert input files from wildcard strings to real file names.
    inputs = [input_file.format(**wildcards) for input_file in inputs]
    return inputs


rule export_who:
    input:
        tree = rules.refine.output.tree,
        metadata = rules.parse.output.metadata,
        auspice_config = "config/auspice_config_who_{lineage}.json",
        node_data = _get_node_data_for_report_export,
        colors = files.colors,
    output:
        tree = "auspice-who/flu_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}_tree.json",
        meta = "auspice-who/flu_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}_meta.json",
    shell:
        """
        augur export v1 \
            --tree {input.tree} \
            --metadata {input.metadata} \
            --node-data {input.node_data} \
            --auspice-config {input.auspice_config} \
            --colors {input.colors} \
            --output-tree {output.tree} \
            --output-meta {output.meta} \
            --minify-json
        """

rule targets:
    input:
        tree = "auspice-who/flu_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}_tree.json",
        meta = "auspice-who/flu_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}_meta.json",
        entropy = "auspice-who/flu_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}_entropy.json",
        frequencies = "auspice-who/flu_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}_frequencies.json",
        sequences = "auspice-who/flu_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}_sequences.json",
        titer_sub_model = "auspice-who/flu_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}_titer-sub-model.json",
        titer_tree_model = "auspice-who/flu_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}_titer-tree-model.json",
        titers = "auspice-who/flu_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}_titers.json"
    output:
        target = "targets/flu_{center}_{lineage}_{segment}_{resolution}_{passage}_{assay}"
    shell:
        '''
        touch {output.target}
        '''