# Path: augur/export_v2.py
# def convert_tree_to_json_structure(node, metadata, div=0):
#     """
#     converts the Biopython tree structure to a dictionary that can
#     be written to file as a json. This is called recursively.
#     Creates the name property & divergence on each node
# 
#     input
#         node -- node for which top level dict is produced.
#         div  -- cumulative divergence (root = 0). False → divergence won't be exported.
# 
#     returns
#         tree in JSON structure
#         list of strains
#     """
# 
#     # Does the tree have divergence? (BEAST trees may not)
#     # only calculate this for the root node!
#     if div == 0 and 'mutation_length' not in metadata[node.name] and 'branch_length' not in metadata[node.name]:
#         div = False
# 
#     node_struct = {'name': node.name, 'node_attrs': {}, 'branch_attrs': {}}
#     if div is not False: # div=0 is ok
#         node_struct["node_attrs"]["div"] = div
# 
#     if node.clades:
#         node_struct["children"] = []
#         for child in node.clades:
#             if div is False:
#                 cdiv=False
#             else:
#                 if 'mutation_length' in metadata[child.name]:
#                     cdiv = div + metadata[child.name]['mutation_length']
#                 elif 'branch_length' in metadata[child.name]:
#                     cdiv = div + metadata[child.name]['branch_length']
#                 else:
#                     print("ERROR: Cannot find branch length information for %s"%(child.name))
# 
#             node_struct["children"].append(convert_tree_to_json_structure(child, metadata, div=cdiv))
# 
#     return node_struct
# 
# Path: augur/validate.py
# class ValidateError(Exception):
#     pass
# 
# Path: augur/validate_export.py
# def ensure_no_duplicate_names(root, ValidateError):
#     """
#     Check that all node names are identical, which is required for auspice (v2) JSONs.
#     """
#     names = set()
#     def recurse(node):
#         if node["name"] in names:
#             raise ValidateError(f"Node {node['name']} appears multiple times in the tree.")
#         names.add(node["name"])
#         if "children" in node:
#             [recurse(child) for child in node["children"]]
#     recurse(root)


import Bio.Phylo
import pytest
import sys
from io import StringIO
from pathlib import Path
from augur.export_v2 import convert_tree_to_json_structure
from augur.validate import ValidateError
from augur.validate_export import ensure_no_duplicate_names

# we assume (and assert) that this script is running from the tests/ directory
sys.path.append(str(Path(__file__).parent.parent.parent))



class TestValidateExport():
    def test_export_without_duplicate_names(self):
        # Create a tree with unique tip names.
        tree = Bio.Phylo.read(StringIO("root(A, internal(B, C))"), "newick")
        metadata = {"A": {}, "B": {}, "C": {}, "root": {}, "internal": {}}
        root = convert_tree_to_json_structure(tree.root, metadata)
        ensure_no_duplicate_names(root, ValidateError)

    def test_export_with_duplicate_names(self):
        # Create a tree with duplicate tip names.
        tree = Bio.Phylo.read(StringIO("root(A, internal(B, B))"), "newick")
        metadata = {"A": {}, "B": {}, "root": {}, "internal": {}}
        root = convert_tree_to_json_structure(tree.root, metadata)

        with pytest.raises(ValidateError):

# Path: augur/export_v2.py
# def convert_tree_to_json_structure(node, metadata, div=0):
#     """
#     converts the Biopython tree structure to a dictionary that can
#     be written to file as a json. This is called recursively.
#     Creates the name property & divergence on each node
# 
#     input
#         node -- node for which top level dict is produced.
#         div  -- cumulative divergence (root = 0). False → divergence won't be exported.
# 
#     returns
#         tree in JSON structure
#         list of strains
#     """
# 
#     # Does the tree have divergence? (BEAST trees may not)
#     # only calculate this for the root node!
#     if div == 0 and 'mutation_length' not in metadata[node.name] and 'branch_length' not in metadata[node.name]:
#         div = False
# 
#     node_struct = {'name': node.name, 'node_attrs': {}, 'branch_attrs': {}}
#     if div is not False: # div=0 is ok
#         node_struct["node_attrs"]["div"] = div
# 
#     if node.clades:
#         node_struct["children"] = []
#         for child in node.clades:
#             if div is False:
#                 cdiv=False
#             else:
#                 if 'mutation_length' in metadata[child.name]:
#                     cdiv = div + metadata[child.name]['mutation_length']
#                 elif 'branch_length' in metadata[child.name]:
#                     cdiv = div + metadata[child.name]['branch_length']
#                 else:
#                     print("ERROR: Cannot find branch length information for %s"%(child.name))
# 
#             node_struct["children"].append(convert_tree_to_json_structure(child, metadata, div=cdiv))
# 
#     return node_struct
# 
# Path: augur/validate.py
# class ValidateError(Exception):
#     pass
# 
# Path: augur/validate_export.py
# def ensure_no_duplicate_names(root, ValidateError):
#     """
#     Check that all node names are identical, which is required for auspice (v2) JSONs.
#     """
#     names = set()
#     def recurse(node):
#         if node["name"] in names:
#             raise ValidateError(f"Node {node['name']} appears multiple times in the tree.")
#         names.add(node["name"])
#         if "children" in node:
#             [recurse(child) for child in node["children"]]
#     recurse(root)


# Path: tests/test_validate_export.py
import Bio.Phylo
import pytest
import sys
from io import StringIO
from pathlib import Path
from augur.export_v2 import convert_tree_to_json_structure
from augur.validate import ValidateError
from augur.validate_export import ensure_no_duplicate_names


# we assume (and assert) that this script is running from the tests/ directory
sys.path.append(str(Path(__file__).parent.parent.parent))



class TestValidateExport():
    def test_export_without_duplicate_names(self):
        # Create a tree with unique tip names.
        tree = Bio.Phylo.read(StringIO("root(A, internal(B, C))"), "newick")
        metadata = {"A": {}, "B": {}, "C": {}, "root": {}, "internal": {}}
        root = convert_tree_to_json_structure(tree.root, metadata)
        ensure_no_duplicate_names(root, ValidateError)

    def test_export_with_duplicate_names(self):
        # Create a tree with duplicate tip names.
        tree = Bio.Phylo.read(StringIO("root(A, internal(B, B))"), "newick")
        metadata = {"A": {}, "B": {}, "root": {}, "internal": {}}
        root = convert_tree_to_json_structure(tree.root, metadata)

        with pytest.raises(ValidateError):

            ensure_no_duplicate_names(root, ValidateError)