# ----------------------------------------------------------------------------
# Copyright (c) 2013--, scikit-bio development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file LICENSE.txt, distributed with this software.
# ----------------------------------------------------------------------------

from heapq import heapify, heappop

import numpy as np

from skbio.tree import TreeNode
from ._c_me import (
    _preorder,
    _postorder,
    _insert_taxon,
    _avgdist_taxon,
    _bal_avgdist_taxon,
    _avgdist_d2_insert,
    _bal_avgdist_insert,
    _ols_lengths,
    _ols_lengths_d2,
    _bal_lengths,
    _ols_min_branch_d2,
    _bal_min_branch,
    _avgdist_matrix,
    _bal_avgdist_matrix,
    _avgdist_swap,
    _bal_avgdist_swap,
    _ols_all_swaps,
    _ols_corner_swaps,
    _bal_all_swaps,
    _bal_avgdist_insert_p,
)
from ._utils import _check_dm, _check_dm_tree


def gme(dm, neg_as_zero=True):
    r"""Perform greedy minimum evolution (GME) for phylogenetic reconstruction.

    .. versionadded:: 0.6.2

    .. versionchanged:: 0.6.3
        Computational efficiency significantly improved.

    Parameters
    ----------
    dm : skbio.DistanceMatrix
        Input distance matrix containing distances between taxa.
    neg_as_zero : bool, optional
        If True (default), convert negative branch lengths into zeros.

    Returns
    -------
    TreeNode
        Reconstructed phylogenetic tree.

    See Also
    --------
    bme
    nni
    nj

    Notes
    -----
    Greedy Minimum Evolution (GME) [1]_ is a distance-based algorithm for phylogenetic
    reconstruction. It utilizes the minimum evolution (ME) principle [2]_ for selecting
    a tree topology with the lowest sum of branch lengths, as calculated using an
    ordinary least squares (OLS) framework [3]_.

    GME is *O*\(*n*:sup:`2`) in time and *O*\(*n*) in space, making it a more scalable
    method than multiple alternatives (e.g., :func:`nj` and :func:`bme`). Therefore it
    is suitable for reconstructing very large phylogenetic trees.

    GME generates an unrooted tree with variable tip heights and potentially some
    negative branch lengths. The first taxon in the distance matrix is always placed as
    a child of the root node. See the notes of :func:`nj` for how to deal with unrooted
    trees and negative branch lengths.

    A GME-generated tree may be further improved by executing the FastNNI algorithm
    implemented in :func:`nni` (with ``balanced=False``).

    A similar but less scalable algorithm using the balanced instead of OLS framework
    is provided in :func:`bme`.

    The same methods underlying :func:`gme`, :func:`bme` and :func:`nni` were also
    provided by the software package FastME [4]_.

    .. note::
        These scikit-bio functions were implemented following the original paper [1]_.
        It is not guaranteed that they will precisely mirror FastME's output. Although
        in practices they usually generate identical or equally optimal phylogenetic
        trees as FastME does.

    References
    ----------
    .. [1] Desper, R., & Gascuel, O. (2002). Fast and accurate phylogeny reconstruction
       algorithms based on the minimum-evolution principle. J Comput Biol, 9(5),
       687-705.

    .. [2] Rzhetsky, A., & Nei, M. (1993). Theoretical foundation of the
       minimum-evolution method of phylogenetic inference. Mol Biol Evol, 10(5),
       1073-1095.

    .. [3] Cavalli-Sforza, L. L., & Edwards, A. W. (1967). Phylogenetic analysis.
       Models and estimation procedures. American journal of human genetics,
       19(3 Pt 1), 233.

    .. [4] Lefort, V., Desper, R., & Gascuel, O. (2015). FastME 2.0: a comprehensive,
       accurate, and fast distance-based phylogeny inference program. Mol Biol Evol,
       32(10), 2798-2800.

    Examples
    --------
    Define a new distance matrix object describing the distances between five taxa:
    human, monkey, pig, rat, and chicken.

    >>> from skbio import DistanceMatrix
    >>> from skbio.tree import gme

    >>> dm = DistanceMatrix([[0, 0.02,  0.18,  0.34,  0.55],
    ...                      [0.02,  0, 0.19, 0.35,  0.55],
    ...                      [0.18, 0.19,  0,  0.34,  0.54],
    ...                      [0.34, 0.35,  0.34,  0,  0.62],
    ...                      [0.55,  0.55,  0.54,  0.62,  0]],
    ...                     ['human','monkey','pig','rat','chicken'])

    Perform Greedy Minimum Evoltuion (GME) and construct the minimum evolution tree
    representing the relationship between those taxa. This is returned as a TreeNode
    object.

    >>> tree = gme(dm)
    >>> print(tree.ascii_art())
              /-monkey
             |
             |          /-pig
             |---------|
    ---------|         |          /-rat
             |          \--------|
             |                    \-chicken
             |
              \-human

    """
    _check_dm(dm)

    # reconstruct tree topology and branch lengths using GME
    tree, lens = _gme(dm.data)

    if neg_as_zero:
        lens[lens < 0] = 0

    return _to_treenode(tree, dm.ids, lens, unroot=True)


def bme(dm, neg_as_zero=True, **kwargs):
    r"""Perform balanced minimum evolution (BME) for phylogenetic reconstruction.

    .. versionadded:: 0.6.3

    Parameters
    ----------
    dm : skbio.DistanceMatrix
        Input distance matrix containing distances between taxa.
    neg_as_zero : bool, optional
        If True (default), convert negative branch lengths into zeros.

    Returns
    -------
    TreeNode
        Reconstructed phylogenetic tree.

    See Also
    --------
    gme
    nni
    nj

    Notes
    -----
    Balanced Minimum Evolution (BME) [1]_ is a refinement of the distance-based minimum
    evolution problem where the average distances between subtrees are independent of
    the sizes of the subtrees. This is referred to as a balanced (or simply BME)
    framework [2]_, as in contrast to the OLS framework used by GME (:func:`gme`).

    The BME algorithm implemented here uses a similar greedy algorithm as implemented
    in :func:`gme`, but less scalable due to the need to update subtree distances as
    the tree topology changes. The algorithm is sub-*O*\(*n*:sup:`3`) in time and
    *O*\(*n*:sup:`2`) in space.

    Refer to :func:`gme` for the format of the output tree and subsequent treatments.

    A BME-generated tree may be further improved by executing the BNNI algorithm
    implemented in :func:`nni` (with ``balanced=True``).

    The same method was provided by FastME [3]_. See :func:`gme` for notes on this.

    .. note::
        Experimental feature: Add ``parallel=True`` will enable parallelization,
        which may increase the performance of the algorithm. This feature may not
        be stable and may be modified without notice in the future.

    References
    ----------
    .. [1] Desper, R., & Gascuel, O. (2002). Fast and accurate phylogeny reconstruction
       algorithms based on the minimum-evolution principle. J Comput Biol, 9(5),
       687-705.

    .. [2] Pauplin, Y. (2000). Direct calculation of a tree length using a distance
       matrix. J Mol Evol, 51, 41-47.

    .. [3] Lefort, V., Desper, R., & Gascuel, O. (2015). FastME 2.0: a comprehensive,
       accurate, and fast distance-based phylogeny inference program. Mol Biol Evol,
       32(10), 2798-2800.

    Examples
    --------
    Define a new distance matrix object describing the distances between five taxa:
    human, monkey, pig, rat, and chicken.

    >>> from skbio import DistanceMatrix
    >>> from skbio.tree import bme

    >>> dm = DistanceMatrix([[0, 0.02,  0.18,  0.34,  0.55],
    ...                      [0.02,  0, 0.19, 0.35,  0.55],
    ...                      [0.18, 0.19,  0,  0.34,  0.54],
    ...                      [0.34, 0.35,  0.34,  0,  0.62],
    ...                      [0.55,  0.55,  0.54,  0.62,  0]],
    ...                     ['human','monkey','pig','rat','chicken'])

    Perform Balanced Minimum Evoltuion (BME) and construct the minimum evolution tree
    representing the relationship between those taxa. This is returned as a TreeNode
    object.

    >>> tree = bme(dm)
    >>> print(tree.ascii_art())
              /-monkey
             |
             |          /-pig
             |---------|
    ---------|         |          /-rat
             |          \--------|
             |                    \-chicken
             |
              \-human

    """
    _check_dm(dm)

    # reconstruct tree topology and branch lengths using BME
    tree, lens = _bme(dm.data, **kwargs)

    if neg_as_zero:
        lens[lens < 0] = 0

    return _to_treenode(tree, dm.ids, lens, unroot=True)


def nni(tree, dm, balanced=True, neg_as_zero=True):
    r"""Perform nearest neighbor interchange (NNI) to improve a phylogenetic tree.

    .. versionadded:: 0.6.2

    .. versionchanged:: 0.6.3
        Computational efficiency significantly improved.

    Parameters
    ----------
    tree : skbio.TreeNode
        Input phylogenetic tree to be rearranged.
    dm : skbio.DistanceMatrix
        Input distance matrix containing distances between taxa.
    balanced : bool, optional
        Use the OLS framework (False) or the balanced framework (True, default).
    neg_as_zero : bool, optional
        If True (default), convert negative branch lengths into zeros.

    Returns
    -------
    TreeNode
        Rearranged phylogenetic tree.

    Notes
    -----
    Nearest neighbor interchange (NNI) is a method for tree rearrangement aiming at
    optimizing a given tree topology according to certain criteria. It iteratively
    swaps neighboring branches that could result in a better tree, until no such swaps
    remain in the optimized tree.

    This function performs NNI for the minimum evolution (ME) problem on phylogenetic
    trees. The implementation is based on [1]_. Two versions of the method are
    provided: the FastNNI algorithm (``balanced=False``) based on an ordinary least
    squares (OLS) framework (see also :func:`gme`) and the BNNI algorithm
    (``balanced=True``) based on a balanced framework (see also :func:`bme`). The
    former is more scalable than the latter.

    The input tree may be rooted or unrooted, and it must be strictly bifurcating. One
    can apply :meth:`~TreeNode.is_bifurcating` or :meth:`~TreeNode.bifurcate` to check
    or make a bifurcating tree. The set of taxa in the tree (accessible via
    :meth:`~TreeNode.subset`) must match those in the distance matrix.

    The output tree is unrooted, with the first taxon in the distance matrix placed as
    a child of the root node. Branch lengths are assigned according to the framework of
    choice. See also :func:`nj` regarding rooting an unrooted tree and dealing with
    potential negative branch lengths.

    The same methods were provided by FastME [2]_. See :func:`gme` for notes on this.

    References
    ----------
    .. [1] Desper, R., & Gascuel, O. (2002). Fast and accurate phylogeny reconstruction
       algorithms based on the minimum-evolution principle. J Comput Biol, 9(5),
       687-705.

    .. [2] Lefort, V., Desper, R., & Gascuel, O. (2015). FastME 2.0: a comprehensive,
       accurate, and fast distance-based phylogeny inference program. Mol Biol Evol,
       32(10), 2798-2800.

    Examples
    --------
    >>> from skbio import DistanceMatrix, TreeNode
    >>> from skbio.tree import nni

    Define a distance matrix describing the distances between five taxa.

    >>> dm = DistanceMatrix([[0, 0.02,  0.18,  0.34,  0.55],
    ...                      [0.02,  0, 0.19, 0.35,  0.55],
    ...                      [0.18, 0.19,  0,  0.34,  0.54],
    ...                      [0.34, 0.35,  0.34,  0,  0.62],
    ...                      [0.55,  0.55,  0.54,  0.62,  0]],
    ...                     ['human','monkey','pig','rat','chicken'])

    Provide a tree to be rearranged. The tree must be bifurcating. Branches lengths are
    not required.

    >>> tree = TreeNode.read(["((human,chicken),(rat,monkey),pig);"])
    >>> print(tree.ascii_art())
                        /-human
              /--------|
             |          \-chicken
             |
    ---------|          /-rat
             |---------|
             |          \-monkey
             |
              \-pig

    Perform nearest neighbor interchange (NNI). This will rearrange tree topolgy to
    better approach the minimum evolution criterion.

    >>> tree = nni(tree, dm)
    >>> print(tree.ascii_art())
                        /-pig
              /--------|
             |         |          /-chicken
             |          \--------|
    ---------|                    \-rat
             |
             |--monkey
             |
              \-human

    Besides rearranging tree topology, estimated branch lengths are assigned to the
    tree.

    >>> rat = tree.find('rat')
    >>> print(rat.length)
    0.20875

    """
    _check_dm_tree(dm, tree)

    # generate tree array
    taxa = dm.ids
    tree, preodr, postodr = _root_from_treenode(tree, taxa)

    # allocate lengths
    lens = np.empty(len(tree), dtype=float)

    # perform BNNI or FastNNI
    func = _bnni if balanced else _fastnni
    func(dm.data, tree, preodr, postodr, lens)

    if neg_as_zero:
        lens[lens < 0] = 0

    # generate TreeNode object
    return _to_treenode(tree, taxa, lens, unroot=True)


def _gme(dm):
    r"""Perform greedy minimum evolution (GME) for phylogenetic reconstruction.

    Parameters
    ----------
    dm : ndarray of float of shape (m, m)
        Input distance matrix containing distances between taxa.

    Returns
    -------
    tree : ndarray of int of shape (2m - 3, 4)
        Reconstructed tree topology by GME.
    lens : ndarray of float of shape (2m - 3,)
        Estimated branch lengths by OLS.

    Notes
    -----
    The GME algorithm involves growing an unrooted tree by iteratively inserting taxa.
    In the beginning it initializes a triplet tree with the first three taxa:

          0
          |
          x
         / \
        1   2

    Taxon 0 is placed at the root position, with its unique descendant x serving as
    connector for taxa 1 and 2 as its left and right children. Next, taxon 3 will
    be inserted into a branch that suffices the minimum evolution criterion. Let's say
    the branch connecting x and 2. The tree becomes:

          0
          |
          x
         / \
        1   y
           / \
          2   3

    So on so forth until all taxa are included in the tree.

    Finally, branch lengths are estimated using an ordinary least squares (OLS)
    framework.

    ---

    This implementation uses an array-based tree structure to improve efficiency. See
    :func:`_check_tree` for details of this structure. Basically, the root node (taxon
    0) is omitted, making the tree strictly binary. Instead, the unique descendant (x)
    is now treated as the root and marked with taxon 0:

          0
         / \
        1   y
           / \
          2   3

    The root is always the first row in the tree array. Other nodes are appended to the
    tree array in the same order as they are added to the tree.

    ---

    For an input distance matrix with m taxa, the output tree will have 2m - 3 nodes.
    This memory space is pre-allocated to improve efficiency. The algorithm doesn't
    remake arrays during iteration, but repeatedly uses the already allocated space.
    During an interation involving x nodes, the first x positions of each array are
    occupied, while the remaining positions are disregarded. There is no need to reset
    array values after each iteration.

    Since they are created de novo, all arrays are C-continuous. This permits further
    optimization in the Cython code.

    """
    # number of taxa
    m = dm.shape[0]

    # number of nodes in the final tree
    n = 2 * m - 3

    # Pre-allocate memory spaces:
    # tree structure
    tree, preodr, postodr = _allocate_tree(n)

    # average distances between distant-2 subtrees
    ad2 = np.empty((n, 2), dtype=float)

    # average distances from a taxon to each subtree
    adk = np.empty((n, 2), dtype=float)

    # branch lengths or length changes
    lens = np.empty((n,), dtype=float)

    # Initialize 3-taxon tree.
    _init_tree(dm, tree, preodr, postodr, ad2, matrix=False)

    # Iteratively add taxa to the tree.
    for k in range(3, m):
        # Calculate average distances from new taxon to existing subtrees.
        _avgdist_taxon(adk, k, dm, tree, preodr, postodr)

        # Find a branch with minimum length change.
        target = _ols_min_branch_d2(lens, ad2, adk, tree, preodr)

        # Update average distances between distant-2 subtrees.
        _avgdist_d2_insert(ad2, target, adk, tree, preodr)

        # Insert new taxon into tree.
        _insert_taxon(k, target, tree, preodr, postodr, use_depth=False)

    # Calculate branch lengths using an OLS framework.
    _ols_lengths_d2(lens, ad2, tree)

    return tree, lens


def _bme(dm, parallel=False):
    r"""Perform balanced minimum evolution (BME) for phylogenetic reconstruction.

    Parameters
    ----------
    dm : ndarray of float of shape (m, m)
        Input distance matrix containing distances between taxa.

    Returns
    -------
    tree : ndarray of int of shape (2m - 3, 4)
        Reconstructed tree topology by BME.
    lens : ndarray of float of shape (2m - 3,)
        Estimated branch lengths by a balanced framework.

    Notes
    -----
    The BME algorithm resembles GME but uses a balanced instead of OLS framework.
    See :func:`_gme` for details of the latter.

    The main differences are: 1) BME requires a full matrix stored and iteratively
    updated, which is less efficient than GME in time and space. 2) BME doesn't rely
    on subtree sizes, which is simpler and more efficient than GME in individual
    calculations (but this cannot compensate for the former).

    """
    func = _bal_avgdist_insert_p if parallel else _bal_avgdist_insert

    # numbers of taxa and nodes in the tree
    m = dm.shape[0]
    n = 2 * m - 3

    # Pre-allocate memory spaces:
    # tree structure
    tree, preodr, postodr = _allocate_tree(n)

    # average distances between all subtrees
    # (This is the dominant factor in BME, and what makes it more expensive than GME.)
    adm = np.empty((n, n), dtype=float)

    # average distances from a taxon to each subtree
    adk = np.empty((n, 2), dtype=float)

    # branch lengths or length changes
    lens = np.empty((n,), dtype=float)

    # a stack for traversal operations
    stack = np.empty((n,), dtype=int)

    # initialize 3-taxon tree
    _init_tree(dm, tree, preodr, postodr, adm, matrix=True)

    # Pre-calculate negative powers of 2.
    powers = np.ldexp(1.0, -np.arange(m))

    # Iteratively add taxa to the tree.
    for k in range(3, m):
        # Calculate balanced average distances from new taxon to existing subtrees.
        _bal_avgdist_taxon(adk, k, dm, tree, preodr, postodr)

        # Find the branch with minimum length change.
        target = _bal_min_branch(lens, adm, adk, tree, preodr)

        # Update balanced average distance matrix between all subtrees.
        func(adm, target, adk, tree, preodr, postodr, powers, stack)

        # Insert new taxon into tree.
        _insert_taxon(k, target, tree, preodr, postodr, use_depth=True)

    # Calculate branch lengths using a balanced framework.
    _bal_lengths(lens, adm, tree)

    return tree, lens


def _fastnni(dm, tree, preodr, postodr, lens):
    r"""Perform fast nearest neighbor interchange (FastNNI) on a tree.

    To improve the tree under the minimum evolution (ME) criterion using an OLS
    framework on a distance matrix between taxa.

    This algorithm was implemented following Section 2.2 and Appendix 4 of Desper and
    Gascuel (2002). Basically, it creates a full average distance matrix between all
    pairs of subtrees, evaluates all possible swaps, then it iteratively perform
    swaps that lead to the maximum reduction of overall branch lengths. During this
    process, it needs to update the evaluation of the four corner branches after one
    branch is swapped.

    ---

    The swaps that can reduce overall branch length are stored in a heap. The paper
    suggests using a maxheap, but the following code uses Python's heapq which is a
    minheap. Therefore the length changes are negated (in contrast to `_bnni`). Each
    item in the heap is a tuple consisting of:

        0. Length change (negative; smaller is better)
        1. Corresponding node index
        2. Corresponding side (0: left, 1: right) child to be swapped with sibling

    For example, (-0.15, 7, 1) means that swapping the right child and the sibling of
    node 7 will reduce the overall branch length by 0.15.

    Meanwhile, the length change of each node is stored in `lens`. For root and tips,
    this value is 0. For internal nodes, this value is negative or 0. Only negative
    values are stored in the heap. 0 means this node is not useful.

    Additionally, the "side" information is stored in column 7 of the tree array.

    When a swap is updated, the new length change (if negative) is pushed into the
    heap, but the old length change won't be deleted from the heap (which would be
    inefficient as O(n)). Instead, a "lazy deletion" mechanism is adopted: When the
    negative-most swap is popped out of the heap, it is checked against the stored
    length change (in `lens`) and side (in `tree[:, 7]`). If they don't match, then
    it means that it is outdated and should be skipped. Otherwise, it is considered
    relevant and the swap will be performed.

    TODO: heapq is a Python module and may not be efficient. Tests showed that the
    code isn't significantly faster than a naive `np.argmax` on then entire `lens`.
    Consider optimization.

    """
    n = tree.shape[0]
    stack = np.empty(n, dtype=int)

    # Calculate average distances between all pairs of subtrees.
    adm = np.empty((n, n))
    _avgdist_matrix(adm, dm, tree, preodr, postodr)

    # Calculate length changes of all possible swaps.
    _ols_all_swaps(lens, tree, adm)

    # Create a heap that stores negative length changes and their nodes and sides.
    heap = [(lens[i], i, tree[i, 7]) for i in np.nonzero(lens)[0]]
    heapify(heap)

    # Iteratively swap branches until there is no more beneficial swap.
    while heap:
        # Pop the swap that reduces the most length.
        L, target, side = heappop(heap)

        # Skip if this swap is outdated.
        if L != lens[target] or side != tree[target, 7]:
            continue

        # Reset this swap.
        lens[target] = 0

        # Swap the branches in the tree.
        _swap_branches(target, side, tree, preodr, stack, use_depth=False)

        # Update average distances after swapping.
        _avgdist_swap(adm, target, side, tree)

        # Update length reductions of swaps of the four corner branches.
        _ols_corner_swaps(target, heap, lens, tree, adm)

    # Calculate branch lengths using an OLS framework.
    _ols_lengths(lens, adm, tree)


def _bnni(dm, tree, preodr, postodr, lens):
    r"""Perform balanced nearest neighbor interchange (BNNI) on a tree.

    To improve the tree under the balanced minimum evolution (BME) criterion given a
    distance matrix between taxa.

    Implemented according to Section 3.2 and Appendix 5 of Desper and Gascuel (2002).
    Basically, BNNI is similar to FastNNI (see `_fastnni`), but it needs to update
    more than the four corner branches after each swap, as balanced average distances
    are dependent on the topology of the tree.

    ---

    The original paper doesn't explain how to store and update positive swaps. Because
    a large proportion of swaps need to be re-evaluated during every iteration, one
    would expect that a heap structure as suggested for the FastNNI algorithm may not
    be efficient. The current code simply stores all swaps (positive or not) and runs
    `np.argmax` on all of them during each iteration.

    """
    n = tree.shape[0]
    stack = np.empty(n, dtype=int)

    # Calculate balanced average distances between all pairs of subtrees.
    adm = np.empty((n, n))
    _bal_avgdist_matrix(adm, dm, tree, preodr, postodr)

    # Pre-calculate negative powers of 2.
    powers = np.ldexp(1.0, -np.arange(dm.shape[0]))

    # Initialize branch swapping information.
    gains, sides, nodes = _init_swaps(tree)

    # Iteratively swap branches until there is no more beneficial swap.
    while True:
        # Calculate length reductions of all possible swaps.
        _bal_all_swaps(gains, sides, nodes, adm, tree)

        # Find the swap with the maximum length reduction, and stop if non-positive.
        branch = gains.argmax()
        if (gain := gains[branch]) <= 0:
            break
        side = sides[branch]
        target = nodes[branch]

        # Swap the branches in the tree.
        _swap_branches(target, side, tree, preodr, stack, use_depth=True)

        # Update balanced average distances after swapping.
        _bal_avgdist_swap(adm, target, side, tree, preodr, powers, stack)

    # Calculate branch lengths using a balanced framework.
    _bal_lengths(lens, adm, tree)


def _check_tree(tree, preodr, postodr):
    r"""Check the integrity of an array-based tree structure.

    The greedy algorithms implemented in this module use a special array-based tree
    structure to improve efficiency.

    In this 2-D array, rows represent nodes in the order of addition to the tree,
    with the root placed at row 0. There are eight columns:

        0. Left child index, or 0 if a tip.
        1. Right child index, or taxon index if a tip.
        2. Parent index, or 0 if root.
        3. Sibling index, or 0 if root.
        4. Size, i.e., number of taxa descending from the node, or 1 if a tip.
        5. Depth, i.e., number of branches constituting the path to root.
        6. Preorder index.
        7. Postorder index.

    Meanwhile, node indices in preorder and postorder are stored in two separate 1-D
    arrays to facilitate traversals.

    Note: Columns 2-7 and the two separate arrays can be derived from columns 0 and 1.
    They are explicitly stored because of their frequent usage during the calculation.
    Also, they are updated iteratively as the tree grows, which is more efficient than
    de novo calculation on an existing tree.

    ---

    For example, a tree with four taxa is like (see :func:`_gme`):

          0
         / \
        1   y
           / \
          2   3

    This tree can be represented by the following 2-D array:

        [[1, 2, 0, 0, 3, 0, 0, 4],
         [0, 1, 0, 2, 1, 1, 1, 0],
         [3, 4, 0, 1, 2, 1, 2, 3],
         [0, 2, 2, 4, 1, 2, 3, 1],
         [0, 3, 2, 3, 1, 2, 4, 2]]

    The separate pre- and postorder arrays are:

        [0, 1, 2, 3, 4]
        [1, 3, 4, 2, 0]

    ---

    This function ensures that a tree array is valid. It is for test purpose, but not
    used in the actual greedy algorithms.

    """
    assert tree.shape[1] == 8
    n = tree[0, 4] * 2 - 1
    assert (tree[:n, 6].argsort() == preodr[:n]).all()
    assert (tree[:n, 7].argsort() == postodr[:n]).all()

    for i in range(n):
        left, right, parent, sibling, size, depth, preidx, postidx = tree[i]
        if left == 0:
            assert i == 0 or right != 0
        else:
            assert tree[left, 2] == tree[right, 2] == i
            assert tree[left, 3] == right
            assert tree[right, 3] == left

    sizes = np.zeros((n,), dtype=int)
    for i in range(n):
        node = postodr[i]
        if tree[node, 0] == 0:
            sizes[node] = 1
        else:
            sizes[node] = sizes[tree[node, :2]].sum()
    assert (sizes == tree[:n, 4]).all()

    depths = np.zeros((n,), dtype=int)
    for i in range(1, n):
        node = preodr[i]
        depths[node] = depths[tree[node, 2]] + 1
    assert (depths == tree[:n, 5]).all()

    stack = np.zeros((n,), dtype=int)
    _preorder(exp := np.zeros((n,), dtype=int), tree, stack)
    assert (exp == preodr[:n]).all()
    _postorder(exp := np.zeros((n,), dtype=int), tree, stack)
    assert (exp == postodr[:n]).all()


def _allocate_tree(n):
    r"""Pre-allocate memory space for an array-based tree structure.

    Parameters
    ----------
    n : int
        Total number of nodes in the tree.

    Returns
    -------
    (n, 8) ndarray of int
        Tree structure.
    (n,) ndarray of int
        Nodes in preorder.
    (n,) ndarray of int
        Nodes in postorder.

    See Also
    --------
    _check_tree

    """
    tree = np.empty((n, 8), dtype=int)
    preodr = np.empty((n,), dtype=int)
    postodr = np.empty((n,), dtype=int)
    return tree, preodr, postodr


def _to_treenode(tree, taxa, lens=None, unroot=False):
    r"""Convert an array-based tree structure into a TreeNode object.

    Parameters
    ----------
    tree : (n, 2+) array_like of int
        Tree structure.
    taxa : list of str
        Taxon names in order.
    lens : (n,) array_like of float, optional
        Branch lengths.
    unroot : bool, optional
        If True, generate an unrooted tree in which the root node has three children,
        the last of which is the first taxon. Otherwise leave the first taxon at the
        root node.

    Returns
    -------
    TreeNode
        Converted TreeNode object.

    See Also
    --------
    TreeNode

    Notes
    -----
    This function only needs the first two columns (left child, right child / taxon
    name) of a tree array to generate a TreeNode object.

    """
    nodes = [TreeNode(taxa[0])] + [
        TreeNode(None if a else taxa[b]) for a, b in tree[1:, :2]
    ]
    if lens is not None:
        for node, L in zip(nodes, lens):
            node.length = L
    for node, (left, right, *_) in zip(nodes, tree):
        if left:
            node.extend([nodes[left], nodes[right]], uncache=False)
    tree = nodes[0]
    if unroot:
        tree.append(TreeNode(taxa[0], length=tree.length), uncache=False)
        tree.name = None
        tree.length = None
    return tree


def _from_treenode(obj, taxmap, tree, preodr, postodr, pos=0, depth=0):
    r"""Convert a TreeNode object into an array-based tree structure.

    Parameters
    ----------
    tree : TreeNode
        Input tree. It must be strictly bifurcating.
    taxmap : dict of str : int
        Mapping of taxon names to indices.
    tree : (n, 8) ndarray of int
        Tree structure.
    preodr : (n,) ndarray of int
        Nodes in preorder.
    postodr : (n,) ndarray of int
        Nodes in postorder.
    pos : int, optional
        Position (axis 0) in the tree array for new data.
    depth : int, optional
        Depth of the root node.

    Raises
    ------
    ValueError
        If tree is not strictly bifurcating.

    See Also
    --------
    _allocate_tree
    _root_from_treenode

    Notes
    -----
    This function fills in pre-allocated tree arrays. It doesn't return any value.

    Nodes in the resulting tree array are ordered by preorder traversal. This may not
    agree with the order they were added to the original tree, if that ever happened.
    Therefore, one shouldn't directly compare the original and generated tree arrays.
    However, the topologies represented by the two tree arrays should be identical.

    The function doesn't fill branch lengths, which aren't required by the downstream
    algorithms. Adding this functionality should be straightforward.

    Parameters ``pos`` and ``depth`` are typically set by :func:`_root_from_treenode`.
    One doesn't need to specify them if working with this function alone. ``depth``
    also impacts the position in ``postodr``. See ``_root_from_treenode`` for
    rationales.

    """
    # determine position in postorder
    post_pos = pos - depth

    # create root node
    tree[pos, 2] = 0
    tree[pos, 3] = 0
    tree[pos, 5] = depth

    # if root node is a tip, just wrap up and return
    if not obj.children:
        tree[pos, 0] = 0
        tree[pos, 1] = taxmap[obj.name]
        tree[pos, 4] = 1
        preodr[pos] = tree[pos, 6] = pos
        postodr[post_pos] = pos
        tree[pos, 7] = post_pos
        return

    # perform preorder traversal, index nodes, fill parent (2), depth (5), taxon (1)
    # and length
    obj.i = pos
    pos_1 = pos + 1
    for i, node in enumerate(obj.preorder(include_self=False)):
        node.i = (ni := pos_1 + i)
        tree[ni, 2] = (pi := node.parent.i)
        tree[ni, 5] = tree[pi, 5] + 1
        if not node.children:
            tree[ni, 0] = 0
            tree[ni, 1] = taxmap[node.name]

    # number of nodes in the current tree
    n = i + 2

    # fill preorder and indices
    end = pos + n
    preodr[pos:end] = tree[pos:end, 6] = np.arange(pos, end)

    # perform postorder traversal, fill children (0 and 1), sibling (3), and size (4)
    for i, node in enumerate(obj.postorder()):
        postodr[post_pos + i] = (ni := node.i)
        tree[ni, 7] = post_pos + i
        if not node.children:
            tree[ni, 4] = 1
        else:
            try:
                left, right = node.children
            except ValueError:
                raise ValueError("Tree is not strictly bifurcating.")
            li, ri = left.i, right.i
            tree[ni, 0] = tree[ri, 3] = li
            tree[ni, 1] = tree[li, 3] = ri
            tree[ni, 4] = tree[li, 4] + tree[ri, 4]
            del left.i
            del right.i

    del obj.i


def _root_from_treenode(obj, taxa):
    r"""Convert a TreeNode object into a tree array rooted at the first taxon.

    Parameters
    ----------
    tree : TreeNode
        Input tree. It must be strictly bifurcating.
    taxa : list of str
        Taxon names in order.
    tree : (n, 8) ndarray of int
        Tree structure.

    Returns
    -------
    (n, 8) ndarray of int
        Tree structure.
    (n,) ndarray of int
        Nodes in preorder.
    (n,) ndarray of int
        Nodes in postorder.

    Raises
    ------
    ValueError
        If tree is not strictly bifurcating.

    See Also
    --------
    _from_treenode
    TreeNode.unrooted_move

    Notes
    -----
    This function locates the first taxon in the tree, walks from it to the root, and
    rotate nodes along the path such that the first taxon becomes the new root in the
    generated tree array. For example (p: parent, c: clade):

          root           taxon
         /  | \           /  \
        c3 p2 c4         c1  p2
          /  \      =>      /  \
         c2  p1           c2  root
            /  \              /  \
        taxon  c1            c3   c4

    Note that the original parent node always becomes the right child of the current
    node. The first parent (p1) is gone and replaced by the taxon. The output array is
    always preordered.

    """
    errmsg = "Tree is not strictly bifurcating."
    taxmap = {taxon: i for i, taxon in enumerate(taxa)}

    m = len(taxa)  # number of taxa
    n = m * 2 - 3  # number of nodes

    # allocate tree arrays
    tree, preodr, postodr = _allocate_tree(n)

    # position of the current node
    pos = 0

    # depth of the current node
    depth = 0

    # position of the previous node
    prev_pos = 0

    # position of the other child of the previous node in the path;
    # it will become the sibling of the current node
    sib_pos = 0

    # number of tips under the current node
    size = m - 1

    n_1 = n - 1

    # helper to add the current node to the array
    def _add_current():
        tree[pos, 2] = prev_pos
        tree[pos, 4] = size
        tree[pos, 5] = depth
        tree[pos, 6] = pos
        preodr[pos] = pos
        tree[pos, 7] = (post_pos := n_1 - depth)
        postodr[post_pos] = pos
        tree[prev_pos, 1] = pos
        tree[sib_pos, 3] = pos
        tree[pos, 3] = sib_pos

    # iterate over ancestors of the first taxon
    curr = obj.find(taxa[0])
    parent = None
    while True:
        prev = curr
        curr = curr.parent
        parent = curr.parent
        siblings = [x for x in curr.children if x is not prev]

        if not siblings:  # singleton branch
            raise ValueError(errmsg)
        n_sibs = len(siblings)
        if n_sibs == 1:
            other = siblings[0]

            # Current node is a regular internal node of the tree (the typical case):
            # Add the current node, then add the other child of it.
            #      parent            prev
            #        |                 |
            #      curr     =>       curr
            #     /   \             /   \
            # other   prev      other   parent
            if parent is not None:
                _add_current()

                prev_pos = pos
                pos += 1

                # add the other child clade
                _from_treenode(other, taxmap, tree, preodr, postodr, pos, depth + 1)

                tree[pos - 1, 0] = pos  # curr's left child
                tree[pos, 2] = pos - 1  # parent (now missing sibling)

                sib_pos = pos
                size -= (k := tree[pos, 4])
                pos += k * 2 - 1

            # Current node is the root of a rooted tree (i.e., it has two children):
            # Don't add the current node, instead add the other child.
            #      curr             prev
            #     /   \     =>      /
            # other   prev      other
            else:
                _from_treenode(other, taxmap, tree, preodr, postodr, pos, depth)

                tree[pos, 2] = prev_pos
                tree[prev_pos, 1] = pos
                tree[sib_pos, 3] = pos
                tree[pos, 3] = sib_pos

                break

        elif n_sibs == 2:
            # Current node is the root of an unrooted tree (i.e., it has three
            # children): Add the current node, then add the other two children of it.
            #                         prev
            #       curr               |
            #     /  |  \    =>       curr
            # left right prev        /   \
            #                     left   right
            if parent is None:
                _add_current()
                curr_pos = pos
                pos += 1

                left, right = siblings
                _from_treenode(left, taxmap, tree, preodr, postodr, pos, depth + 1)
                tree[pos, 2] = curr_pos
                left_pos = pos
                tree[curr_pos, 0] = pos
                pos += tree[pos, 4] * 2 - 1

                _from_treenode(right, taxmap, tree, preodr, postodr, pos, depth + 1)
                tree[pos, 2] = curr_pos
                tree[curr_pos, 1] = pos
                tree[left_pos, 3] = pos
                tree[pos, 3] = left_pos

                break

            else:
                raise ValueError(errmsg)
        else:
            raise ValueError(errmsg)

        # move up one level
        depth += 1

    return tree, preodr, postodr


def _init_tree(dm, tree, preodr, postodr, ads, matrix=False):
    """Initialize tree (triplet with taxa 0, 1 and 2)."""
    # triplet tree
    tree[:3] = [
        [1, 2, 0, 0, 2, 0, 0, 2],  # 0: root
        [0, 1, 0, 2, 1, 1, 1, 0],  # 1: left child
        [0, 2, 0, 1, 1, 1, 2, 1],  # 2: right child
    ]

    # nodes in pre- and postorder
    preodr[:3] = [0, 1, 2]
    postodr[:3] = [1, 2, 0]

    # average distance matrix between all subtrees
    if matrix:
        ads[0, 1] = ads[1, 0] = dm[0, 1]
        ads[0, 2] = ads[2, 0] = dm[0, 2]
        ads[1, 2] = ads[2, 1] = dm[1, 2]

    # average distances between distant-2 subtrees
    else:
        ads[1, 0] = dm[0, 1]
        ads[2, 0] = dm[0, 2]
        ads[1, 1] = ads[2, 1] = dm[1, 2]
        ads[0, 0] = ads[0, 1] = 0


def _insert_taxon_treenode(taxon, target, tree):
    r"""Insert a taxon between a target node and its parent.

    This function resembles :func:`_insert_taxon` but operates on a TreeNode object. It
    is for test purpose. It is not used in the actual algorithms.

    """
    link = TreeNode()
    tip = TreeNode(taxon)
    parent = target.parent
    if parent is None:  # root branch
        link.extend(tree.children)
        tree.append(link)
        tree.append(tip)
    else:  # other branch
        is_left = target is parent.children[0]
        link.append(target)
        link.append(tip)
        parent.append(link)
        if is_left:
            parent.children = parent.children[::-1]


def _avgdist_matrix_naive(adm, dm, tree, postodr):
    r"""Calculate a matrix of average distances between all pairs of subtrees.

    This function produces the same result as :func:`_avgdist_matrix`. However, it
    calculates all subtree-to-subtree distances based on the original taxon-to-taxon
    distances, as discussed in Eq. 1 of Desper and Gascuel (2002). instead of adopting
    a recursive strategy (Eq. 2). Therefore, it is significantly slower.

        d(A, B) = \frac{1}{|A||B|} \sum_{i \in A, j \in B} d(i, j)

    This algorithm is implemented as a reference and for comparison purpose. It is not
    used in the actual GME algorithm.

    """
    n = tree[0, 4] * 2 - 1
    taxas = {}
    for node in postodr[:n]:
        left, right = tree[node, :2]
        if not left:
            taxas[node] = frozenset([right])
        else:
            taxas[node] = taxas[left] | taxas[right]

    full = taxas[0] | frozenset([0])

    for a in range(n):
        a_taxa = taxas[a]
        a_taxa_r = full - a_taxa
        # adm[a, a] = dm[list(a_taxa)][:, list(a_taxa_r)].mean()  # self distance
        for b in range(a + 1, n):
            # If a is an ancestor (proper superset) of b, flip a's taxon set as the
            # the upper subtree of a will be considered.
            a_taxa_ = a_taxa_r if a_taxa > (b_taxa := taxas[b]) else a_taxa
            adm[a, b] = adm[b, a] = dm[list(a_taxa_)][:, list(b_taxa)].mean()
            # Because this is postorder traversal, a cannot be a descendant (proper
            # subset) of b.


def _avgdist_taxon_naive(adk, taxon, dm, tree, postodr):
    """Calculate average distances between a new taxon and existing subtrees.

    This function produces the same result as :func:`_avgdist_taxon`, but it
    calculates all taxon-to-subtree distances based on the original taxon-to-taxon
    distances, as discussed in Eq. 1 of Desper and Gascuel (2002), instead of
    adopting a recursive strategy (Eq. 2). Therefore, it is significantly slower.

    This algorithm is implemented as a reference and for comparison purpose, but not
    used in the actual GME algorithm.

    """
    n = tree[0, 4] * 2 - 1
    taxas = {}
    for node in postodr[:n]:
        left, right = tree[node, :2]
        if not left:
            taxas[node] = frozenset([right])
        else:
            taxas[node] = taxas[left] | taxas[right]

    full = taxas[0] | frozenset([0])

    dk = dm[taxon]
    for node in range(n):
        taxa_lower = taxas[node]
        adk[node, 0] = dk[list(taxa_lower)].mean()
        taxa_upper = full - taxa_lower
        adk[node, 1] = dk[list(taxa_upper)].mean()


def _init_swaps(tree):
    """Initialize branch swapping information.

    It will create three 1-D arrays to store information of all internal branches of
    the tree:

    - `gains`: Overall tree length reduction (larger is better).
    - `sides`: Which side (left: 0, right: 1) of the target node should be swapped.
    - `nodes`: Corresponding node index of the branch.

    Meanwhile, column 7 of the tree array will be filled with the branch index of the
    each node, if it corresponds to one. This column was previous used to store
    postorder index, but that information is not needed during swapping.

    For a tree with n taxa, there should be n - 3 internal branches. The association
    between nodes and internal branches are fixed, no matter how tree is swapped.

    """
    # number of internal branches
    n = tree[0, 4] - 2
    gains = np.zeros((n,), dtype=float)
    sides = np.empty((n,), dtype=int)
    nodes = np.empty((n,), dtype=int)

    # identify all internal branches
    branch = 0
    for node in range(1, tree.shape[0]):
        if tree[node, 0]:
            tree[node, 7] = branch
            nodes[branch] = node
            branch += 1

    return gains, sides, nodes


def _swap_branches_treenode(node1, node2):
    """Swap two nodes in a TreeNode object.

    This function is for testing purpose.

    """
    parent1, parent2 = node1.parent, node2.parent
    parent1.append(node2, False)
    parent2.append(node1, False)


def _swap_branches(target, side, tree, preodr, stack, use_depth=True):
    r"""Swap one child of the target node with its sibling.

                 |                         |
               parent                    parent
               /   \                     /   \
           target  sibling    =>     target  child
            /  \                      /  \
        other  child              other  sibling

    Target must be an internal node. Root and tips are not allowed.

    This function currently doesn't handle postorder, as subsequent algorithms don't
    need this information.

    This function is specifically designed for nearest neighbor interchange (NNI). It
    may be generalized to swapping non-neighbor branches, but that is not currently
    implemented.

    """
    # This function can potentially optimized using Cython. See _insert_taxon.

    child = tree[target, side]
    other = tree[target, 1 - side]  # other child
    parent = tree[target, 2]
    sibling = tree[target, 3]

    c_size = tree[child, 4]
    o_size = tree[other, 4]
    s_size = tree[sibling, 4]

    # update connections
    tree[target, side] = sibling
    tree[target, 3] = child
    tree[target, 4] += s_size - c_size

    p_side = int(tree[parent, 0] == target)  # sibling side of parent
    tree[parent, p_side] = child

    tree[sibling, 2] = target
    tree[sibling, 3] = other

    tree[child, 2] = parent
    tree[child, 3] = target

    tree[other, 3] = sibling

    # locate the clades under the relevant nodes
    c_start = tree[child, 6]
    c_width = c_size * 2 - 1
    c_end = c_start + c_width
    c_clade = preodr[c_start:c_end]

    s_start = tree[sibling, 6]
    s_width = s_size * 2 - 1
    s_end = s_start + s_width
    s_clade = preodr[s_start:s_end]

    o_start = tree[other, 6]
    o_width = o_size * 2 - 1
    o_end = o_start + o_width
    o_clade = preodr[o_start:o_end]

    # update depth (if needed)
    if use_depth:
        tree[c_clade, 5] -= 1
        tree[s_clade, 5] += 1

    # update preorder (naive)
    # _preorder(preodr, tree, stack)
    # tree[:n, 6] = np.argsort(preodr[:n])

    # update postorder (naive)
    # _postorder(postodr, tree, stack)
    # tree[:n, 7] = np.argsort(postodr[:n])

    # update preorder
    # sibling is the left child of parent
    if p_side == 0:
        tree[target, 6] += c_width - s_width
        stack[:c_width] = c_clade
        stack[c_width] = target
        c1_width = c_width + 1

        # sibling_clade, target, child_clade, other_clade =>
        # child_clade, target, sibling_clade, other_clade
        if side == 0:
            tree[s_clade, 6] += c_width + 1
            tree[c_clade, 6] -= s_width + 1
            preodr[(sc1_start := s_start + c1_width) : c_end] = s_clade
            preodr[s_start:sc1_start] = stack[:c1_width]

        # sibling_clade, target, other_clade, child_clade =>
        # child_clade, target, other_clade, sibling_clade
        else:
            tree[s_clade, 6] += c_width + o_width + 1
            tree[c_clade, 6] -= s_width + o_width + 1
            tree[o_clade, 6] += c_width - s_width
            stack[c1_width : (c1o_width := c1_width + o_width)] = o_clade
            preodr[(sc1o_start := s_start + c1o_width) : c_end] = s_clade
            preodr[s_start:sc1o_start] = stack[:c1o_width]

    # sibling is the right child of parent
    else:
        stack[:s_width] = s_clade

        # target, child_clade, other_clade, sibling_clade =>
        # target, sibling_clade, other_clade, child_clade
        if side == 0:
            tree[c_clade, 6] += (so_width := s_width + o_width)
            tree[s_clade, 6] -= c_width + o_width
            tree[o_clade, 6] += s_width - c_width
            stack[s_width:so_width] = o_clade
            preodr[(cso_start := c_start + so_width) : s_end] = c_clade
            preodr[c_start:cso_start] = stack[:so_width]

        # target, other_clade, child_clade, sibling_clade =>
        # target, other_clade, sibling_clade, child_clade
        else:
            tree[c_clade, 6] += s_width
            tree[s_clade, 6] -= c_width
            preodr[(cs_start := c_start + s_width) : s_end] = c_clade
            preodr[c_start:cs_start] = stack[:s_width]