Trees¶

class Tree:
    """A rooted tree with named leaves.

    Construct via :meth:`from_newick`, :meth:`read`, or
    :meth:`from_linkage`; the leaf order (:meth:`leaf_names`) defines the
    plot row order when the tree is passed to
    :meth:`dot_explorer.DotPlotter.plot` or
    :func:`dot_explorer.plot_similarity_heatmap`.
    """

    def __init__(self, root: TreeNode):
        self.root = root

    # ------------------------------------------------------------------
    # Construction
    # ------------------------------------------------------------------

    @classmethod
    def from_newick(cls, text: str) -> 'Tree':
        """Parse a newick string (IQ-TREE ``.treefile`` compatible).

        Handles quoted labels (``'...'`` with ``''`` escapes), bracketed
        comments, missing branch lengths, and internal support labels in
        both the ``)95:0.1`` and ``)label/95:0.1`` styles.

        Parameters
        ----------
        text : str
            Newick source; everything after the first ``;`` is ignored.

        Returns
        -------
        Tree
            The parsed tree.

        Raises
        ------
        ValueError
            On malformed newick (unbalanced parentheses, unterminated
            quotes, empty input, or duplicate leaf names).
        """
        root = _parse_newick(text)
        tree = cls(root)
        names = tree.leaf_names()
        if not names:
            raise ValueError('newick tree has no labelled leaves')
        from collections import Counter

        dupes = sorted(n for n, c in Counter(names).items() if c > 1)
        if dupes:
            raise ValueError(f'duplicate leaf names in tree: {dupes}')
        return tree

    @classmethod
    def read(cls, path: str | Path) -> 'Tree':
        """Read a newick tree from *path* (``.nwk``/``.newick``/``.treefile``/``.txt``).

        Parameters
        ----------
        path : str or Path
            File containing a single newick tree.

        Returns
        -------
        Tree
            The parsed tree.
        """
        return cls.from_newick(Path(path).read_text())

    @classmethod
    def from_linkage(cls, Z: 'np.ndarray', labels: Sequence[str]) -> 'Tree':
        """Build a tree from a scipy ``linkage`` matrix.

        Branch lengths are derived from merge heights, so cutting the
        tree at distance ``d`` from the tips matches
        ``scipy.cluster.hierarchy.fcluster(Z, t=d, criterion='distance')``.
        The leaf order matches scipy's default (unsorted) ``dendrogram``.

        Parameters
        ----------
        Z : numpy.ndarray
            ``(n - 1, 4)`` linkage matrix from
            :func:`scipy.cluster.hierarchy.linkage`.
        labels : sequence of str
            Names of the ``n`` original observations, in the order they
            were given to ``linkage``.

        Returns
        -------
        Tree
            Ultrametric tree with one leaf per label.
        """
        n = len(labels)
        if len(Z) != n - 1:
            raise ValueError(
                f'linkage matrix has {len(Z)} rows; expected {n - 1} for {n} labels'
            )
        nodes: dict[int, TreeNode] = {
            i: TreeNode(name=str(labels[i])) for i in range(n)
        }
        heights: dict[int, float] = dict.fromkeys(range(n), 0.0)
        for k, (a, b, height, _count) in enumerate(Z):
            a, b = int(a), int(b)
            parent = TreeNode(children=[nodes[a], nodes[b]])
            for child in (a, b):
                nodes[child].length = float(height) - heights[child]
            nodes[n + k] = parent
            heights[n + k] = float(height)
        return cls(nodes[n + len(Z) - 1])

    # ------------------------------------------------------------------
    # Queries
    # ------------------------------------------------------------------

    def _postorder(self) -> Iterator[TreeNode]:
        """Yield nodes children-first (iterative; safe for deep trees)."""
        stack: list[tuple[TreeNode, bool]] = [(self.root, False)]
        while stack:
            node, expanded = stack.pop()
            if expanded:
                yield node
            else:
                stack.append((node, True))
                for child in reversed(node.children):
                    stack.append((child, False))

    def leaves(self) -> list[TreeNode]:
        """Return the leaf nodes in left-to-right (plot top-to-bottom) order."""
        out = []
        stack = [self.root]
        while stack:
            node = stack.pop()
            if node.is_leaf:
                out.append(node)
            else:
                stack.extend(reversed(node.children))
        return out

    def leaf_names(self) -> list[str]:
        """Return leaf labels in left-to-right (plot top-to-bottom) order."""
        return [leaf.name or '' for leaf in self.leaves()]

    def max_depth(self) -> float:
        """Return the largest root-to-leaf cumulative branch length.

        When the tree carries no branch lengths at all, every edge counts
        as 1.0; otherwise missing lengths count as 0.0.
        """
        default = 1.0 if not self._has_lengths() else 0.0
        deepest = 0.0
        stack: list[tuple[TreeNode, float]] = [(self.root, 0.0)]
        while stack:
            node, depth = stack.pop()
            if node.is_leaf:
                deepest = max(deepest, depth)
            for child in node.children:
                length = child.length if child.length is not None else default
                stack.append((child, depth + length))
        return deepest

    def _has_lengths(self) -> bool:
        return any(node.length is not None for node in self._postorder())

    def validate_labels(self, seq_names: Sequence[str]) -> None:
        """Check that tree tips and sequence names match exactly.

        Parameters
        ----------
        seq_names : sequence of str
            The sequence names the tree will be plotted against.

        Raises
        ------
        ValueError
            When any tip lacks a sequence or any sequence lacks a tip;
            the message lists both directions of the mismatch.
        """
        tips = set(self.leaf_names())
        seqs = set(seq_names)
        problems = []
        extra_tips = sorted(tips - seqs)
        if extra_tips:
            problems.append(f'tree tips with no matching sequence: {extra_tips}')
        missing_tips = sorted(seqs - tips)
        if missing_tips:
            problems.append(f'sequences with no tip in the tree: {missing_tips}')
        if problems:
            raise ValueError(
                'tree tip labels do not match sequence names — ' + '; '.join(problems)
            )

@dataclass
class TreeNode:
    """A node in a :class:`Tree`.

    Attributes
    ----------
    name : str or None
        Leaf label, or internal-node label when present (IQ-TREE writes
        support values here).
    length : float or None
        Branch length to the parent, or ``None`` when the newick source
        omitted it.
    support : float or None
        Numeric support value parsed from an internal-node label
        (``)95:`` or ``)label/95:`` styles), when recognisable.
    children : list of TreeNode
        Child nodes; empty for leaves.
    """

    name: str | None = None
    length: float | None = None
    support: float | None = None
    children: list['TreeNode'] = field(default_factory=list)

    @property
    def is_leaf(self) -> bool:
        """Return True when this node has no children."""
        return not self.children