include: "base.snakefile"

rule all:
    input:
        hidden_nodes = "auspice/v2_hidden-nodes.json",
        time_only = "auspice/v2_time-only-tree.json",
        div_only = "auspice/v2_div-only-tree.json",
        bool_metadata = "auspice/v2_boolean-metadata.json",
        no_transmission_lines = "auspice/v2_no-transmission-lines.json",
        default_layout = "auspice/v2_default-layout.json",
        footer = "auspice/v2_custom-footer.json",
        filter_not_color = "auspice/v2_filters-not-colors.json"
rule files:
    params:
        hidden_nodes = "data/hidden_nodes.json",
        boolean_metadata = "data/metadata_boolean.tsv",
        footer_description = "config/footer-description.md"
files = rules.files.params

rule refine_with_temporal_information:
    message: "Refining tree with temporal information"
    input:
        tree = rules.tree.output.tree,
        alignment = rules.align.output,
        metadata = rules.parse.output.metadata
    output:
        tree = "results/tree_temporal.nwk",
        node_data = "results/temporal_branch_lengths.json"
    shell:
        """
        augur refine \
            --tree {input.tree} \
            --alignment {input.alignment} \
            --metadata {input.metadata} \
            --output-tree {output.tree} \
            --output-node-data {output.node_data} \
            --timetree \
            --coalescent opt \
            --date-confidence \
            --date-inference marginal \
            --clock-filter-iqd 4
        """

rule refine_without_temporal_information:
    message: "Refining tree with temporal information"
    input:
        tree = rules.tree.output.tree,
        alignment = rules.align.output,
        metadata = rules.parse.output.metadata
    output:
        tree = "results/tree_div_only.nwk",
        node_data = "results/div_only_branch_lengths.json"
    shell:
        """
        augur refine \
            --tree {input.tree} \
            --alignment {input.alignment} \
            --metadata {input.metadata} \
            --output-tree {output.tree} \
            --output-node-data {output.node_data} \
            --keep-root
        """

rule export_with_hidden_nodes:
    message: "Exporting (v2) temporal tree for auspice with hidden nodes"
    input:
        tree = {rules.refine_with_temporal_information.output.tree},
        branch_lengths = {rules.refine_with_temporal_information.output.node_data},
        hidden_nodes = {files.hidden_nodes},
        metadata = {rules.parse.output.metadata},
        colors = {base_files.colors},
        config = {base_files.auspice_config_v2}
    output:
        auspice = rules.all.input.hidden_nodes,
    shell:
        """
        augur export v2 \
            --tree {input.tree} \
            --node-data {input.branch_lengths} {input.hidden_nodes} \
            --metadata {input.metadata} \
            --colors {input.colors} \
            --auspice-config {input.config} \
            --title "Tree with hidden nodes (both div + temporal)" \
            --output {output.auspice}
        """

rule make_node_data_time_only:
    message: "Removing divergence information from {input.branch_lengths}"
    input:
        branch_lengths = {rules.refine_with_temporal_information.output.node_data}
    output:
        node_data = "results/time_only_branch_lengths.json"
    run:
        import json
        with open(input.branch_lengths[0]) as f:
            data = json.load(f)
        for d in data["nodes"].values():
            d.pop("branch_length", None)
            d.pop("clock_length", None)
            d.pop("mutation_length", None)
        with open(output.node_data, 'w') as f:
            json.dump(data, f, indent=2, sort_keys=True)

rule export_div_only_tree:
    message: "Exporting (v2) div-only tree"
    input:
        tree = {rules.refine_without_temporal_information.output.tree},
        branch_lengths = {rules.refine_without_temporal_information.output.node_data},
        metadata = {rules.parse.output.metadata},
        colors = {base_files.colors},
        config = {base_files.auspice_config_v2}
    output:
        auspice = rules.all.input.div_only,
    shell:
        """
        augur export v2 \
            --tree {input.tree} \
            --node-data {input.branch_lengths} \
            --metadata {input.metadata} \
            --colors {input.colors} \
            --auspice-config {input.config} \
            --title "Divergence-only tree" \
            --output {output.auspice}
        """

rule export_time_only_tree:
    message: "Exporting (v2) temporal-only tree"
    input:
        tree = {rules.refine_with_temporal_information.output.tree},
        branch_lengths = {rules.make_node_data_time_only.output.node_data},
        metadata = {rules.parse.output.metadata},
        colors = {base_files.colors},
        config = {base_files.auspice_config_v2}
    output:
        auspice = rules.all.input.time_only,
    shell:
        """
        augur export v2 \
            --tree {input.tree} \
            --node-data {input.branch_lengths} \
            --metadata {input.metadata} \
            --colors {input.colors} \
            --auspice-config {input.config} \
            --title "Temporal-only tree" \
            --output {output.auspice}
        """

rule export_with_boolean_metadata:
    message: "Exporting tree with boolean metadata"
    input:
        tree = {rules.refine_with_temporal_information.output.tree},
        branch_lengths = {rules.make_node_data_time_only.output.node_data},
        metadata = {files.boolean_metadata},
        colors = {base_files.colors},
        config = {base_files.auspice_config_v2}
    output:
        auspice = rules.all.input.bool_metadata
    shell:
        """
        augur export v2 \
            --tree {input.tree} \
            --node-data {input.branch_lengths} \
            --metadata {input.metadata} \
            --color-by-metadata top_half_True_False top_half_1_0 top_half_yes_no all_missing \
            --colors {input.colors} \
            --title "Boolean metadata traits" \
            --output {output.auspice};
        """

rule export_with_default_layout:
    message: "Exporting (v2) with default layout as clock"
    input:
        tree = {rules.refine_with_temporal_information.output.tree},
        branch_lengths = {rules.refine_with_temporal_information.output.node_data},
        metadata = {rules.parse.output.metadata},
        colors = {base_files.colors},
        config = "config/default-layout.json"
    output:
        auspice = rules.all.input.default_layout,
    shell:
        """
        augur export v2 \
            --tree {input.tree} \
            --node-data {input.branch_lengths} \
            --metadata {input.metadata} \
            --colors {input.colors} \
            --auspice-config {input.config} \
            --output {output.auspice}
        """

rule export_with_custom_footer_description:
    message: "Exporting (v2) with custom footer description"
    input:
        tree = {rules.refine_with_temporal_information.output.tree},
        branch_lengths = {rules.refine_with_temporal_information.output.node_data},
        metadata = {rules.parse.output.metadata},
        footer = {files.footer_description}
    output:
        auspice = rules.all.input.footer,
    shell:
        """
        augur export v2 \
            --tree {input.tree} \
            --node-data {input.branch_lengths} \
            --metadata {input.metadata} \
            --description {input.footer} \
            --title "Custom footer" \
            --output {output.auspice}
        """

rule export_filter_which_isnt_a_coloring:
    message: "Exporting (v2) with a filter which is not a coloring"
    input:
        tree = {rules.refine_with_temporal_information.output.tree},
        branch_lengths = {rules.refine_with_temporal_information.output.node_data},
        metadata = {rules.parse.output.metadata},
        config = "config/filters.json"
    output:
        auspice = rules.all.input.filter_not_color,
    shell:
        """
        augur export v2 \
            --tree {input.tree} \
            --node-data {input.branch_lengths} \
            --metadata {input.metadata} \
            --auspice-config {input.config} \
            --output {output.auspice}
        """

rule export_no_transmission_lines:
    message: "Exporting (v2) with transmission lines toggled off"
    input:
        tree = {rules.refine_with_temporal_information.output.tree},
        branch_lengths = {rules.refine_with_temporal_information.output.node_data},
        metadata = {rules.parse.output.metadata},
        config = "config/no-transmission-lines.json"
    output:
        auspice = rules.all.input.no_transmission_lines,
    shell:
        """
        augur export v2 \
            --tree {input.tree} \
            --node-data {input.branch_lengths} \
            --metadata {input.metadata} \
            --auspice-config {input.config} \
            --output {output.auspice}
        """

rule clean:
    message: "Removing auspice directory"
    shell:
        "rm -rfv auspice results"
