import json, argparse

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="remove time info",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--input-node-data", type=str, required=True, help="input data"
    )
    parser.add_argument(
        "--output-node-data",
        type=str,
        metavar="JSON",
        required=True,
        help="output Auspice JSON",
    )
    args = parser.parse_args()
    with open(args.input_node_data) as fh:
        data = json.load(fh)
    new_node_data = {}
    for name, node in data["nodes"].items():
        old_clade_name = node["clade_membership"]
        outbreak_name = ""
        lineage_name = ""

        # if it starts with clade -> it's a clade
        # if it starts with outbreak -> it's outbreak, need to look up clade
        # if it starts with lineage -> it's clade IIb, outbreak hMPXV-1
        if old_clade_name.startswith("clade"):
            clade_name = old_clade_name.split()[1]
        # Need to set up clade dictionary for when we have other outbreaks
        # if old_clade_name.startswith('outbreak'):
        #     outbreak_name = old_clade_name.split()[1]
        #     clade_name = clade[outbreak_name]
        elif old_clade_name.startswith("outgroup"):
            clade_name = "outgroup"
        else:
            clade_name = "IIb"
            outbreak_name = "hMPXV-1"
            lineage_name = old_clade_name

        new_node_data[name] = {
            "clade_membership": clade_name,
            "outbreak": outbreak_name,
            "lineage": lineage_name,
        }
        if "clade_annotation" in node:
            new_node_data[name]["clade_annotation"] = node["clade_annotation"]
            if node["clade_annotation"] == "A":
                new_node_data[name]["clade_annotation"] = "hMPXV-1 A"

    data["nodes"] = new_node_data
    with open(args.output_node_data, "w") as fh:
        json.dump(data, fh)

from collections import defaultdict
import json, argparse




if __name__=="__main__":
    parser = argparse.ArgumentParser(
        description="calculate mutation context json",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )

    parser.add_argument('--tree', type=str, required=True, help="tree file")
    parser.add_argument('--mutations', type=str, required=True, help="mutations")
    parser.add_argument('--output', type=str, metavar="JSON", required=True, help="output Auspice JSON")
    args = parser.parse_args()


    with open(args.mutations) as fh:
        data = json.load(fh)['nodes']


    terminal_muts = defaultdict(lambda: defaultdict(int))
    internal_muts = defaultdict(lambda: defaultdict(int))

    node_data = {}

    for name, node in data.items():
        GA_count = 0
        CT_count = 0
        total_muts = 0
        for mut in node["muts"]:
            a, pos, d = mut[0], int(mut[1:-1]), mut[-1]
            if a in 'ACGT' and d in 'ACGT':
                total_muts += 1
                if a+d == 'GA':
                    GA_count += 1
                elif a+d == 'CT':
                    CT_count += 1
        GA_CT_count = GA_count + CT_count
        if total_muts:
            node_data[name] = {"GA_CT_fraction": GA_CT_count/total_muts}
        else:
            node_data[name] = {"GA_CT_fraction": None }


        dinuc_count = 0
        if GA_CT_count:
            #node_data[name]["CT_fraction"] = CT_count/GA_CT_count
            for mut in node["muts"]:
                a, pos, d = mut[0], int(mut[1:-1]), mut[-1]
                if a in 'ACGT' and d in 'ACGT':
                    if a+d == 'GA' and node['sequence'][pos]=='A':
                        dinuc_count+=1
                    elif a+d == 'CT' and node['sequence'][pos-2]=='T':
                        dinuc_count+=1
            node_data[name]["dinuc_context_fraction"] = dinuc_count/GA_CT_count
        else:
            node_data[name]["dinuc_context_fraction"] = None
            #node_data[name]["CT_fraction"] = None

    with open(args.output, 'w') as fh:
        json.dump({"nodes":node_data}, fh)

import argparse
import json
from collections import defaultdict


def sample_date(node):
    """
    Returns the sample date in numeric form.
    In the future, we could examine the 'raw_date' attr here to decide whether to ignore
    some sequences, as 'numdate' is the inferred (timetree) date which can hide
    uncertainty in actual sampling date,
    """
    if "raw_date" not in node: # internal node or tip with no date info
        return
    return node['numdate']



if __name__=="__main__":
    parser = argparse.ArgumentParser(
        description="remove time info",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )

    parser.add_argument('--input-node-data', type=str, required=True, help="input data")
    parser.add_argument('--output-node-data', type=str, metavar="JSON", required=True, help="output Auspice JSON")
    args = parser.parse_args()


    with open(args.input_node_data) as fh:
        data = json.load(fh)

    new_node_data = {}

    for name, node in data["nodes"].items():
        try:
            new_node_data[name] = {
                "mutation_length": node["mutation_length"],
                "branch_length": node["branch_length"]
            }
            sdate = sample_date(node)
            if sdate:
                new_node_data[name]["sample_date"] = sdate
        except KeyError:
            # internal node or tip with no date info
            pass

    data["nodes"] = new_node_data
    with open(args.output_node_data, 'w') as fh:
        json.dump(data, fh)

import pandas as pd
import argparse
from Bio import SeqIO

if __name__=="__main__":
    parser = argparse.ArgumentParser(
        description="Reverse-complement reverse-complemented sequence",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )

    parser.add_argument('--metadata', type=str, required=True, help="input metadata")
    parser.add_argument('--sequences', type=str, required=True, help="input sequences")
    parser.add_argument('--output', type=str, required=True, help="output sequences")
    args = parser.parse_args()

    metadata = pd.read_csv(args.metadata, sep='\t')

    # Read in fasta file
    with open(args.sequences, 'r') as f_in:
        with open(args.output, 'w') as f_out:
            for seq in SeqIO.parse(f_in, 'fasta'):
                # Check if metadata['reverse'] is True
                try:
                    if metadata.loc[metadata['strain'] == seq.id, 'reverse'].values[0] == True:
                        # Reverse-complement sequence
                        seq.seq = seq.seq.reverse_complement()
                        print("Reverse-complementing sequence:", seq.id)
                except:
                    print("No reverse complement for:", seq.id)

                # Write sequences to file
                SeqIO.write(seq, f_out, 'fasta')

import pandas as pd
import json, argparse
from augur.io import read_metadata

def replace_name_recursive(node, lookup):
    if node["name"] in lookup:
        node["name"] = lookup[node["name"]]

    if "children" in node:
        for child in node["children"]:
            replace_name_recursive(child, lookup)

if __name__=="__main__":
    parser = argparse.ArgumentParser(
        description="Swaps out the strain names in the Auspice JSON with the final strain name",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )

    parser.add_argument('--input-auspice-json', type=str, required=True, help="input auspice_json")
    parser.add_argument('--metadata', type=str, required=True, help="input data")
    parser.add_argument('--metadata-id-columns', nargs="+", help="names of possible metadata columns containing identifier information, ordered by priority. Only one ID column will be inferred.")
    parser.add_argument('--display-strain-name', type=str, required=True, help="field to use as strain name in auspice")
    parser.add_argument('--output', type=str, metavar="JSON", required=True, help="output Auspice JSON")
    args = parser.parse_args()

    metadata = read_metadata(args.metadata, id_columns=args.metadata_id_columns)
    name_lookup = {}
    for ri, row in metadata.iterrows():
        strain_id = row.name
        name_lookup[strain_id] = args.display_strain_name if pd.isna(row[args.display_strain_name]) else row[args.display_strain_name]

    with open(args.input_auspice_json, 'r') as fh:
        data = json.load(fh)

    replace_name_recursive(data['tree'], name_lookup)

    with open(args.output, 'w') as fh:
        json.dump(data, fh)

shell:
    """
    cp {input.auspice_json} {output.auspice_json}
    cp {input.root_sequence} {output.root_sequence_json}
    """

shell:
    "rm -rfv {params}"

shell:
    "rm -rfv {params}"

shell:
    """
    augur filter \
        --metadata {input.metadata} \
        --metadata-id-columns {params.strain_id} \
        --sequences {input.sequences} \
        --include-where strain=MK783032 strain=MK783030 \
        --group-by clade lineage \
        --subsample-max-sequences 50 \
        --subsample-seed 0 \
        --output-metadata {output.metadata} \
        --output-sequences {output.sequences}
    """

shell:
    """
    augur filter \
        --sequences {input.sequences} \
        --metadata {input.metadata} \
        --metadata-id-columns {params.strain_id} \
        --exclude {input.exclude} \
        --output-sequences {output.sequences} \
        --output-metadata {output.metadata} \
        --min-date {params.min_date} \
        --min-length {params.min_length} \
        --exclude-where QC_rare_mutations=bad \
        --output-log {output.log}
    """

shell:
    """
    augur filter \
        --sequences {input.sequences} \
        --metadata {input.metadata} \
        --metadata-id-columns {params.strain_id} \
        --output-strains {output.strains} \
        {params.group_by} \
        {params.sequences_per_group} \
        {params.exclude} \
        {params.other_filters} \
        --output-log {output.log}
    """

shell:
    """
    augur filter \
        --metadata-id-columns {params.strain_id} \
        --sequences {input.sequences} \
        --metadata {input.metadata} \
        --exclude-all \
        --include {input.strains} {input.include}\
        --output-sequences {output.sequences} \
        --output-metadata {output.metadata}
    """

shell:
    """
    python3 scripts/reverse_reversed_sequences.py \
        --metadata {input.metadata} \
        --sequences {input.sequences} \
        --output {output}
    """

shell:
    """
    augur mask \
        --sequences {input.sequences} \
        --mask {input.mask} \
        --mask-from-beginning {params.from_start} \
        --mask-from-end {params.from_end} --output {output}
    """

shell:
    """
    augur tree \
        --alignment {input.alignment} \
        --exclude-sites {input.tree_mask} \
        --tree-builder-args="-redo" \
        --output {output.tree} \
        --nthreads {threads}
    """

shell:
    """
    python3 scripts/fix_tree.py \
        --alignment {input.alignment} \
        --input-tree {input.tree} \
        {params.root} \
        --output {output.tree}
    """

shell:
    """
    augur refine \
        --tree {input.tree} \
        --alignment {input.alignment} \
        --metadata {input.metadata} \
        --metadata-id-columns {params.strain_id} \
        --output-tree {output.tree} \
        --timetree \
        --root {params.root} \
        --precision 3 \
        --keep-polytomies \
        {params.clock_rate} \
        {params.clock_std_dev} \
        --output-node-data {output.node_data} \
        --coalescent {params.coalescent} \
        --date-inference {params.date_inference} \
        --date-confidence \
        --clock-filter-iqd {params.clock_filter_iqd}
    """

shell:
    """
    augur ancestral \
        --tree {input.tree} \
        --alignment {input.alignment} \
        --output-node-data {output.node_data} \
        --inference {params.inference}
    """

shell:
    """
    augur translate \
        --tree {input.tree} \
        --ancestral-sequences {input.node_data} \
        --reference-sequence {input.genemap} \
        --output {output.node_data}
    """

shell:
    """
    augur traits \
        --tree {input.tree} \
        --metadata {input.metadata} \
        --metadata-id-columns {params.strain_id} \
        --output {output.node_data} \
        --columns {params.columns} \
        --confidence \
        --sampling-bias-correction {params.sampling_bias_correction}
    """

shell:
    """
    augur clades \
        --tree {input.tree} \
        --mutations {input.nuc_muts} {input.aa_muts} \
        --clades {input.clades} \
        --output-node-data {output.node_data} 2>&1 | tee {log}
    """

shell:
    """
    python scripts/clades_renaming.py \
    --input-node-data {input} \
    --output-node-data {output.node_data}
    """

shell:
    """
    python3 scripts/mutation_context.py \
        --tree {input.tree} \
        --mutations {input.node_data} \
        --output {output.node_data}
    """

shell:
    """
    python3 scripts/remove_timeinfo.py --input-node-data {input} --output-node-data {output}
    """

shell:
    """
    python3 scripts/construct-recency-from-submission-date.py \
        --metadata {input.metadata} \
        --metadata-id-columns {params.strain_id} \
        --output {output} 2>&1
    """

shell:
    """
    python3 scripts/assign-colors.py \
        --ordering {input.ordering} \
        --color-schemes {input.color_schemes} \
        --output {output.colors} \
        --metadata {input.metadata} 2>&1
    """

shell:
    """
    augur export v2 \
        --tree {input.tree} \
        --metadata {input.metadata} \
        --metadata-id-columns {params.strain_id} \
        --node-data {input.branch_lengths} {input.nt_muts} {input.aa_muts} {input.mutation_context} {input.clades} {input.recency}\
        --colors {input.colors} \
        --lat-longs {input.lat_longs} \
        --description {input.description} \
        --auspice-config {input.auspice_config} \
        --include-root-sequence \
        --output {output.auspice_json}
    """

shell:
    """
    python3 scripts/set_final_strain_name.py --metadata {input.metadata} \
            --metadata-id-columns {params.strain_id} \
            --input-auspice-json {input.auspice_json} \
            --display-strain-name {params.display_strain_field} \
            --output {output.auspice_json}
    cp {input.root_sequence} {output.root_sequence}
    """

shell:
    """
    curl https://mpox-lapis.genspectrum.org/v1/sample/fasta --output {output.sequences}
    """

shell:
    """
    curl https://mpox-lapis.genspectrum.org/v1/sample/details?dataFormat=csv | \
        tr -d "\r" |
        sed -E 's/("([^"]*)")?,/\\2\\t/g' > {output.metadata}
    """

shell:
    """
    nextstrain remote upload {params.deploy_url} {input}
    """

shell:
    """
    ./bin/notify-on-deploy {params.deploy_url} {params.slack_ts}
    """

shell:
    """
    curl -fsSL --compressed {params.sequences_url:q} --output {output.sequences}
    curl -fsSL --compressed {params.metadata_url:q} --output {output.metadata}
    """

shell:
    """
    gzip --decompress --keep {input.metadata}
    xz --decompress --keep {input.sequences}
    """