

--- ../../tskit/python/setup.py ---

import os
import platform

import numpy
from setuptools import Extension
from setuptools import setup

IS_WINDOWS = platform.system() == "Windows"


libdir = "lib"
kastore_dir = os.path.join(libdir, "subprojects", "kastore")
tsk_source_files = [
    "core.c",
    "tables.c",
    "trees.c",
    "genotypes.c",
    "stats.c",
    "convert.c",
    "haplotype_matching.c",
]
sources = (
    ["_tskitmodule.c"]
    + [os.path.join(libdir, "tskit", f) for f in tsk_source_files]
    + [os.path.join(kastore_dir, "kastore.c")]
)

defines = []
libraries = []
if IS_WINDOWS:
    libraries.append("Advapi32")
    defines.append(("WIN32", None))

_tskit_module = Extension(
    "_tskit",
    sources=sources,
    extra_compile_args=["-std=c99"],
    libraries=libraries,
    define_macros=defines,
    include_dirs=["lwt_interface", libdir, kastore_dir, numpy.get_include()],
)

setup(
    ext_modules=[_tskit_module],
)


--- ../../tskit/python/stress_lowlevel.py ---

import curses
import os
import random
import resource
import sys
import time
import tracemalloc
from contextlib import redirect_stdout

import pytest

"""
Code to stress the low-level API as much as possible to expose
any memory leaks or error handling issues.
"""


def main(stdscr):
    if len(sys.argv) > 1:
        args = sys.argv[1:]
    else:
        args = ["-n0", "tests/test_lowlevel.py"]

    class StressPlugin:
        def __init__(self):
            self.max_rss = 0
            self.max_rss_iter = 0
            self.min_rss = 1e100
            self.iteration = 0
            self.last_print = time.time()
            self.memory_start = None

        def pytest_sessionstart(self):
            if self.memory_start is None:
                tracemalloc.start()
                self.memory_start = tracemalloc.take_snapshot()

        def pytest_sessionfinish(self):
            memory_current = tracemalloc.take_snapshot()
            rusage = resource.getrusage(resource.RUSAGE_SELF)
            if self.max_rss < rusage.ru_maxrss:
                self.max_rss = rusage.ru_maxrss
                self.max_rss_iter = self.iteration
            if self.min_rss > rusage.ru_maxrss:
                self.min_rss = rusage.ru_maxrss

            # We don't want to flood stdout, so we rate-limit to 1 per second.
            if time.time() - self.last_print > 1:
                stdscr.clear()
                rows, cols = stdscr.getmaxyx()
                stdscr.addstr(
                    0,
                    0,
                    "iter\tRSS\tmin\tmax\tmax@iter"[: cols - 1],
                )
                stdscr.addstr(
                    1,
                    0,
                    "\t".join(
                        map(
                            str,
                            [
                                self.iteration,
                                rusage.ru_maxrss,
                                self.min_rss,
                                self.max_rss,
                                self.max_rss_iter,
                            ],
                        )
                    )[: cols - 1],
                )
                stats = memory_current.compare_to(self.memory_start, "traceback")
                for i, stat in enumerate(stats[: rows - 3], 1):
                    stdscr.addstr(i + 2, 0, str(stat)[: cols - 1])
                self.last_print = time.time()
                stdscr.refresh()
                self.iteration += 1

    plugin = StressPlugin()
    while True:
        # We don't want any random variation in the amount of memory
        # used from test-to-test.
        random.seed(1)
        with open(os.devnull, "w") as devnull:
            with redirect_stdout(devnull):
                result = pytest.main(args, plugins=[plugin])
        if result != 0:
            exit("TESTS FAILED")


if __name__ == "__main__":
    stdscr = curses.initscr()
    curses.noecho()
    curses.cbreak()

    try:
        main(stdscr)
    finally:
        curses.echo()
        curses.nocbreak()
        curses.endwin()


--- ../../tskit/python/tests/test_wright_fisher.py ---


"""
Test various functions using messy tables output by a forwards-time simulator.
"""
import itertools
import random

import msprime
import numpy as np
import numpy.testing as nt
import pytest

import tests as tests
import tests.tsutil as tsutil
import tskit


class WrightFisherSimulator:
    """
    SIMPLE simulation of `num_pops` bisexual, haploid Wright-Fisher populations
    of size `N` for `ngens` generations, in which each individual survives with
    probability `survival` and only those who die are replaced. If `num_pops` is
    greater than 1, the individual to be replaced has a chance `mig_rate` of
    being the offspring of nodes from a different and randomly chosen
    population. If `num_loci` is None, the chromosome is 1.0 Morgans long. If
    `num_loci` not None, a discrete recombination model is used where
    breakpoints are chosen uniformly from 1 to `num_loci` - 1. If
    `deep_history` is True, a history to coalescence of just one population of
    `self.N` samples is added at the beginning.
    """

    def __init__(
        self,
        N,
        survival=0.0,
        seed=None,
        deep_history=True,
        debug=False,
        initial_generation_samples=False,
        num_loci=None,
        num_pops=1,
        mig_rate=0.0,
        record_migrations=False,
        record_individuals=True,
    ):
        self.N = N
        self.num_pops = num_pops
        self.num_loci = num_loci
        self.survival = survival
        self.mig_rate = mig_rate
        self.record_migrations = record_migrations
        self.record_individuals = record_individuals
        self.deep_history = deep_history
        self.debug = debug
        self.initial_generation_samples = initial_generation_samples
        self.seed = seed
        self.rng = random.Random(seed)

    def random_breakpoint(self):
        if self.num_loci is None:
            return min(1.0, max(0.0, 2 * self.rng.random() - 0.5))
        else:
            return self.rng.randint(1, self.num_loci - 1)

    def run(self, ngens):
        L = 1
        if self.num_loci is not None:
            L = self.num_loci
        tables = tskit.TableCollection(sequence_length=L)
        for _ in range(self.num_pops):
            tables.populations.add_row()
        if self.deep_history:
            # initial population
            population_configurations = [
                msprime.PopulationConfiguration(sample_size=self.N)
            ]
            init_ts = msprime.simulate(
                population_configurations=population_configurations,
                recombination_rate=1.0,
                length=L,
                random_seed=self.seed,
            )
            init_tables = init_ts.dump_tables()
            flags = init_tables.nodes.flags
            if not self.initial_generation_samples:
                flags = np.zeros_like(init_tables.nodes.flags)
            tables.nodes.set_columns(time=init_tables.nodes.time + ngens, flags=flags)
            tables.edges.set_columns(
                left=init_tables.edges.left,
                right=init_tables.edges.right,
                parent=init_tables.edges.parent,
                child=init_tables.edges.child,
            )
        else:
            flags = 0
            if self.initial_generation_samples:
                flags = tskit.NODE_IS_SAMPLE
            for p in range(self.num_pops):
                for _ in range(self.N):
                    individual = -1
                    if self.record_individuals:
                        individual = tables.individuals.add_row(parents=[-1, -1])
                    tables.nodes.add_row(
                        flags=flags, time=ngens, population=p, individual=individual
                    )

        pops = [
            list(range(p * self.N, (p * self.N) + self.N)) for p in range(self.num_pops)
        ]
        pop_ids = list(range(self.num_pops))
        for t in range(ngens - 1, -1, -1):
            if self.debug:
                print("t:", t)
                print("pops:", pops)
            dead = [[self.rng.random() > self.survival for _ in pop] for pop in pops]
            # sample these first so that all parents are from the previous gen
            parent_pop = []
            new_parents = []
            for p in pop_ids:
                w = [
                    1 - self.mig_rate if i == p else self.mig_rate / (self.num_pops - 1)
                    for i in pop_ids
                ]
                parent_pop.append(self.rng.choices(pop_ids, w, k=sum(dead[p])))
                new_parents.append(
                    [
                        self.rng.choices(pops[parent_pop[p][k]], k=2)
                        for k in range(sum(dead[p]))
                    ]
                )

            if self.debug:
                for p in pop_ids:
                    print("Replacing", sum(dead[p]), "individuals from pop", p)
            for p in pop_ids:
                k = 0
                for j in range(self.N):
                    if dead[p][j]:
                        lparent, rparent = new_parents[p][k]
                        individual = -1
                        if self.record_individuals:
                            individual = tables.individuals.add_row(
                                parents=[
                                    tables.nodes[lparent].individual,
                                    tables.nodes[rparent].individual,
                                ]
                            )
                        offspring = tables.nodes.add_row(
                            time=t, population=p, individual=individual
                        )
                        if parent_pop[p][k] != p and self.record_migrations:
                            tables.migrations.add_row(
                                left=0.0,
                                right=L,
                                node=offspring,
                                source=parent_pop[p][k],
                                dest=p,
                                time=t,
                            )
                        k += 1
                        bp = self.random_breakpoint()
                        if self.debug:
                            print("--->", offspring, lparent, rparent, bp)
                        pops[p][j] = offspring
                        if bp > 0.0:
                            tables.edges.add_row(
                                left=0.0, right=bp, parent=lparent, child=offspring
                            )
                        if bp < L:
                            tables.edges.add_row(
                                left=bp, right=L, parent=rparent, child=offspring
                            )

        if self.debug:
            print("Done! Final pop:")
            print(pops)
        flags = tables.nodes.flags
        flattened = [n for pop in pops for n in pop]
        flags[flattened] = tskit.NODE_IS_SAMPLE
        tables.nodes.flags = flags
        tables.time_units = "generations"
        return tables


def wf_sim(
    N,
    ngens,
    survival=0.0,
    deep_history=True,
    debug=False,
    seed=None,
    initial_generation_samples=False,
    num_loci=None,
    num_pops=1,
    mig_rate=0.0,
    record_migrations=False,
    record_individuals=True,
):
    sim = WrightFisherSimulator(
        N,
        survival=survival,
        deep_history=deep_history,
        debug=debug,
        seed=seed,
        initial_generation_samples=initial_generation_samples,
        num_loci=num_loci,
        num_pops=num_pops,
        mig_rate=mig_rate,
        record_migrations=record_migrations,
        record_individuals=record_individuals,
    )
    return sim.run(ngens)


class TestSimulation:
    """
    Tests that the simulations produce the output we expect.
    """

    random_seed = 5678

    def test_one_gen_multipop_mig_no_deep(self):
        tables = wf_sim(
            N=5,
            ngens=1,
            num_pops=4,
            mig_rate=1.0,
            deep_history=False,
            seed=self.random_seed,
            record_migrations=True,
        )
        assert tables.time_units == "generations"
        assert tables.nodes.num_rows == 5 * 4 * (1 + 1)
        assert tables.edges.num_rows > 0
        assert tables.migrations.num_rows == 5 * 4
        assert tables.individuals.num_rows == tables.nodes.num_rows

    def test_multipop_mig_deep(self):
        N = 10
        ngens = 20
        num_pops = 3
        tables = wf_sim(
            N=N,
            ngens=ngens,
            num_pops=num_pops,
            mig_rate=1.0,
            seed=self.random_seed,
            record_migrations=True,
        )
        assert tables.nodes.num_rows > (num_pops * N * ngens) + N
        assert tables.edges.num_rows > 0
        assert tables.sites.num_rows == 0
        assert tables.mutations.num_rows == 0
        assert tables.migrations.num_rows >= N * num_pops * ngens
        assert tables.populations.num_rows == num_pops
        assert tables.individuals.num_rows >= num_pops * N * ngens

        # sort does not support mig
        tables.migrations.clear()
        # making sure trees are valid
        tables.sort()
        tables.simplify()
        ts = tables.tree_sequence()
        sample_pops = tables.nodes.population[ts.samples()]
        assert np.unique(sample_pops).size == num_pops

    def test_multipop_mig_no_deep(self):
        N = 5
        ngens = 5
        num_pops = 2
        tables = wf_sim(
            N=N,
            ngens=ngens,
            num_pops=num_pops,
            mig_rate=1.0,
            deep_history=False,
            seed=self.random_seed,
            record_migrations=True,
        )
        assert tables.nodes.num_rows == num_pops * N * (ngens + 1)
        assert tables.edges.num_rows > 0
        assert tables.sites.num_rows == 0
        assert tables.mutations.num_rows == 0
        assert tables.migrations.num_rows == N * num_pops * ngens
        assert tables.populations.num_rows == num_pops
        assert tables.individuals.num_rows == tables.nodes.num_rows
        # FIXME this is no longer needed.
        # sort does not support mig
        tables.migrations.clear()
        # making sure trees are valid
        tables.sort()
        tables.simplify()
        ts = tables.tree_sequence()
        sample_pops = tables.nodes.population[ts.samples()]
        assert np.unique(sample_pops).size == num_pops

    def test_non_overlapping_generations(self):
        tables = wf_sim(N=10, ngens=10, survival=0.0, seed=self.random_seed)
        assert tables.nodes.num_rows > 0
        assert tables.edges.num_rows > 0
        assert tables.sites.num_rows == 0
        assert tables.mutations.num_rows == 0
        assert tables.migrations.num_rows == 0
        assert tables.individuals.num_rows > 0
        tables.sort()
        tables.simplify()
        ts = tables.tree_sequence()
        # All trees should have exactly one root and all internal nodes should
        # have arity > 1
        for tree in ts.trees():
            assert tree.num_roots == 1
            leaves = set(tree.leaves(tree.root))
            assert leaves == set(ts.samples())
            for u in tree.nodes():
                if tree.is_internal(u):
                    assert len(tree.children(u)) > 1

    def test_overlapping_generations(self):
        tables = wf_sim(N=30, ngens=10, survival=0.85, seed=self.random_seed)
        assert tables.nodes.num_rows > 0
        assert tables.edges.num_rows > 0
        assert tables.sites.num_rows == 0
        assert tables.mutations.num_rows == 0
        assert tables.migrations.num_rows == 0
        assert tables.individuals.num_rows > 0
        tables.sort()
        tables.simplify()
        ts = tables.tree_sequence()
        for tree in ts.trees():
            assert tree.num_roots == 1

    def test_one_generation_no_deep_history(self):
        N = 20
        tables = wf_sim(N=N, ngens=1, deep_history=False, seed=self.random_seed)
        assert tables.nodes.num_rows == 2 * N
        assert tables.edges.num_rows > 0
        assert tables.sites.num_rows == 0
        assert tables.mutations.num_rows == 0
        assert tables.migrations.num_rows == 0
        assert tables.individuals.num_rows > 0
        tables.sort()
        tables.simplify()
        ts = tables.tree_sequence()
        assert tables.nodes.num_rows > 0
        assert tables.edges.num_rows > 0
        ts = tables.tree_sequence()
        for tree in ts.trees():
            all_samples = set()
            for root in tree.roots:
                root_samples = set(tree.samples(root))
                assert len(root_samples & all_samples) == 0
                all_samples |= root_samples
            assert all_samples == set(ts.samples())

    def test_many_generations_no_deep_history(self):
        N = 10
        ngens = 100
        tables = wf_sim(N=N, ngens=ngens, deep_history=False, seed=self.random_seed)
        assert tables.nodes.num_rows == N * (ngens + 1)
        assert tables.edges.num_rows > 0
        assert tables.sites.num_rows == 0
        assert tables.mutations.num_rows == 0
        assert tables.migrations.num_rows == 0
        assert tables.individuals.num_rows > 0
        tables.sort()
        tables.simplify()
        ts = tables.tree_sequence()
        assert tables.nodes.num_rows > 0
        assert tables.edges.num_rows > 0
        ts = tables.tree_sequence()
        # We are assuming that everything has coalesced and we have single-root trees
        for tree in ts.trees():
            assert tree.num_roots == 1

    def test_with_mutations(self):
        N = 10
        ngens = 100
        tables = wf_sim(N=N, ngens=ngens, deep_history=False, seed=self.random_seed)
        tables.sort()
        ts = tables.tree_sequence()
        ts = tsutil.jukes_cantor(ts, 10, 0.1, seed=self.random_seed)
        tables = ts.tables
        assert tables.sites.num_rows > 0
        assert tables.mutations.num_rows > 0
        samples = np.where(tables.nodes.flags == tskit.NODE_IS_SAMPLE)[0].astype(
            np.int32
        )
        tables.sort()
        tables.simplify(samples)
        assert tables.nodes.num_rows > 0
        assert tables.edges.num_rows > 0
        assert tables.nodes.num_rows > 0
        assert tables.edges.num_rows > 0
        assert tables.sites.num_rows > 0
        assert tables.mutations.num_rows > 0
        ts = tables.tree_sequence()
        assert ts.sample_size == N
        for hap in ts.haplotypes():
            assert len(hap) == ts.num_sites

    def test_with_recurrent_mutations(self):
        # actually with only ONE site, at 0.0
        N = 10
        ngens = 100
        tables = wf_sim(N=N, ngens=ngens, deep_history=False, seed=self.random_seed)
        tables.sort()
        ts = tables.tree_sequence()
        ts = tsutil.jukes_cantor(ts, 1, 10, seed=self.random_seed)
        tables = ts.tables
        assert tables.sites.num_rows == 1
        assert tables.mutations.num_rows > 0
        # before simplify
        for h in ts.haplotypes():
            assert len(h) == 1
        # after simplify
        tables.sort()
        tables.simplify()
        assert tables.nodes.num_rows > 0
        assert tables.edges.num_rows > 0
        assert tables.sites.num_rows == 1
        assert tables.mutations.num_rows > 0
        ts = tables.tree_sequence()
        assert ts.sample_size == N
        for hap in ts.haplotypes():
            assert len(hap) == ts.num_sites

    def test_record_individuals_initial_state(self):
        N = 10
        tables = wf_sim(N=N, ngens=0, seed=12345, deep_history=False)
        tables.sort()
        assert len(tables.individuals) == N
        assert len(tables.nodes) == N
        for individual in list(tables.individuals)[:N]:
            assert list(individual.parents) == [-1, -1]
        for j, node in enumerate(tables.nodes):
            assert node.individual == j

    def test_record_individuals(self):
        N = 10
        tables = wf_sim(N=N, ngens=10, seed=12345, deep_history=False)
        assert len(tables.individuals) == len(tables.nodes)
        for node_id, individual in enumerate(tables.nodes.individual):
            assert node_id == individual
        tables.sort()
        ts = tables.tree_sequence()
        for tree in ts.trees():
            for u in tree.nodes():
                individual = ts.individual(ts.node(u).individual)
                parent_node = tree.parent(u)
                if parent_node != tskit.NULL:
                    parent_individual = ts.individual(ts.node(parent_node).individual)
                    assert parent_individual.id in individual.parents


def get_wf_sims(seed):
    wf_sims = []
    for N in [5, 10, 20]:
        for surv in [0.0, 0.5, 0.9]:
            for mut in [0.01, 1.0]:
                for nloci in [1, 2, 3]:
                    tables = wf_sim(N=N, ngens=N, survival=surv, seed=seed)
                    tables.sort()
                    ts = tables.tree_sequence()
                    ts = tsutil.jukes_cantor(ts, num_sites=nloci, mu=mut, seed=seed)
                    wf_sims.append(ts)
    return wf_sims


# List of simulations used to parametrize tests.
wf_sims = get_wf_sims(1234)


class TestSimplify:
    """
    Tests for simplify on cases generated by the Wright-Fisher simulator.
    """

    def verify_simplify(self, ts, new_ts, samples, node_map):
        """
        Check that trees in `ts` match `new_ts` using the specified node_map.
        Modified from `verify_simplify_topology`.  Also check that the `parent`
        and `time` column in the MutationTable is correct.
        """
        # check trees agree at these points
        locs = [random.random() for _ in range(20)]
        locs += random.sample(list(ts.breakpoints())[:-1], min(20, ts.num_trees))
        locs.sort()
        old_trees = ts.trees()
        new_trees = new_ts.trees()
        old_right = -1
        new_right = -1
        for loc in locs:
            while old_right <= loc:
                old_tree = next(old_trees)
                old_left, old_right = old_tree.get_interval()
            assert old_left <= loc < old_right
            while new_right <= loc:
                new_tree = next(new_trees)
                new_left, new_right = new_tree.get_interval()
            assert new_left <= loc < new_right
            # print("comparing trees")
            # print("interval:", old_tree.interval)
            # print(old_tree.draw(format="unicode"))
            # print("interval:", new_tree.interval)
            # print(new_tree.draw(format="unicode"))
            pairs = itertools.islice(itertools.combinations(samples, 2), 500)
            for pair in pairs:
                mapped_pair = [node_map[u] for u in pair]
                mrca1 = old_tree.get_mrca(*pair)
                assert mrca1 != tskit.NULL
                mrca2 = new_tree.get_mrca(*mapped_pair)
                assert mrca2 != tskit.NULL
                assert node_map[mrca1] == mrca2
        mut_parent = tsutil.compute_mutation_parent(ts=ts)
        nt.assert_equal(mut_parent, ts.tables.mutations.parent)

    def verify_haplotypes(self, ts, samples):
        """
        Check that haplotypes are unchanged by simplify.
        """
        sub_ts, node_map = ts.simplify(samples, map_nodes=True, filter_sites=False)
        # Sites tables should be equal
        assert ts.tables.sites == sub_ts.tables.sites
        sub_haplotypes = dict(zip(sub_ts.samples(), sub_ts.haplotypes()))
        all_haplotypes = dict(zip(ts.samples(), ts.haplotypes()))
        mapped_ids = []
        for node_id, h in all_haplotypes.items():
            mapped_node_id = node_map[node_id]
            if mapped_node_id in sub_haplotypes:
                assert h == sub_haplotypes[mapped_node_id]
                mapped_ids.append(mapped_node_id)
        assert sorted(mapped_ids) == sorted(sub_ts.samples())

    @pytest.mark.parametrize("ts", wf_sims)
    def test_python_simplify_all_samples(self, ts):
        s = tests.Simplifier(ts, ts.samples())
        py_full_ts, py_full_map = s.simplify()
        full_ts, full_map = ts.simplify(ts.samples(), map_nodes=True)
        assert all(py_full_map == full_map)
        full_ts.tables.assert_equals(py_full_ts.tables, ignore_provenance=True)

    @pytest.mark.parametrize("ts", wf_sims)
    @pytest.mark.parametrize("nsamples", [2, 5, 10])
    def test_python_simplify_sample_subset(self, ts, nsamples):
        sub_samples = random.sample(list(ts.samples()), min(nsamples, ts.sample_size))
        s = tests.Simplifier(ts, sub_samples)
        py_small_ts, py_small_map = s.simplify()
        small_ts, small_map = ts.simplify(samples=sub_samples, map_nodes=True)
        small_ts.tables.assert_equals(py_small_ts.tables, ignore_provenance=True)
        self.verify_simplify(ts, small_ts, sub_samples, small_map)
        self.verify_haplotypes(ts, samples=sub_samples)

    @pytest.mark.parametrize("ts", wf_sims)
    @pytest.mark.parametrize("nsamples", [2, 5, 10])
    def test_simplify_tables(self, ts, nsamples):
        tables = ts.dump_tables()
        sub_samples = random.sample(list(ts.samples()), min(nsamples, ts.num_samples))
        node_map = tables.simplify(samples=sub_samples)
        small_ts = tables.tree_sequence()
        other_tables = small_ts.dump_tables()
        tables.assert_equals(other_tables, ignore_provenance=True)
        self.verify_simplify(ts, small_ts, sub_samples, node_map)

    @pytest.mark.parametrize("ts", wf_sims)
    @pytest.mark.parametrize("nsamples", [2, 5])
    def test_simplify_keep_unary(self, ts, nsamples):
        np.random.seed(123)
        ts = tsutil.mark_metadata(ts, "nodes")
        sub_samples = random.sample(list(ts.samples()), min(nsamples, ts.num_samples))
        random_nodes = np.random.choice(ts.num_nodes, ts.num_nodes // 2)
        ts = tsutil.insert_individuals(ts, random_nodes)
        ts = tsutil.mark_metadata(ts, "individuals")

        for params in [{}, {"keep_unary": True}, {"keep_unary_in_individuals": True}]:
            sts = ts.simplify(sub_samples, **params)
            # check samples match
            assert sts.num_samples == len(sub_samples)
            for n, sn in zip(sub_samples, sts.samples()):
                assert ts.node(n).metadata == sts.node(sn).metadata

            # check that nodes are correctly retained: only nodes ancestral to
            # retained samples, and: by default, only coalescent events; if
            # keep_unary_in_individuals then also nodes in individuals; if
            # keep_unary then all such nodes.
            for t in ts.trees(tracked_samples=sub_samples):
                st = sts.at(t.interval.left)
                visited = [False for _ in sts.nodes()]
                for n, sn in zip(sub_samples, sts.samples()):
                    last_n = t.num_tracked_samples(n)
                    while n != tskit.NULL:
                        ind = ts.node(n).individual
                        keep = False
                        if t.num_tracked_samples(n) > last_n:
                            # a coalescent node
                            keep = True
                        if "keep_unary_in_individuals" in params and ind != tskit.NULL:
                            keep = True
                        if "keep_unary" in params:
                            keep = True
                        if (n in sub_samples) or keep:
                            visited[sn] = True
                            assert sn != tskit.NULL
                            assert ts.node(n).metadata == sts.node(sn).metadata
                            assert t.num_tracked_samples(n) == st.num_samples(sn)
                            if ind != tskit.NULL:
                                sind = sts.node(sn).individual
                                assert sind != tskit.NULL
                                assert (
                                    ts.individual(ind).metadata
                                    == sts.individual(sind).metadata
                                )
                            sn = st.parent(sn)
                        last_n = t.num_tracked_samples(n)
                        n = t.parent(n)
                st_nodes = list(st.nodes())
                for k, v in enumerate(visited):
                    assert v == (k in st_nodes)


--- ../../tskit/python/tests/conftest.py ---


"""
Configuration and fixtures for pytest. Only put test-suite wide fixtures in here. Module
specific fixtures should live in their modules.

To use a fixture in a test simply refer to it by name as an argument. This is called
dependency injection. Note that all fixtures should have the suffix "_fixture" to make
it clear in test code.

For example to use the `ts` fixture (a tree sequence with data in all fields) in a test:

class TestClass:
    def test_something(self, ts_fixture):
        assert ts_fixture.some_method() == expected

Fixtures can be parameterised etc. see https://docs.pytest.org/en/stable/fixture.html

Note that fixtures have a "scope" for example `ts_fixture` below is only created once
per test session and re-used for subsequent tests.
"""
import msprime
import pytest
from pytest import fixture

from . import tsutil


def pytest_addoption(parser):
    """
    Add options, e.g. to skip tests marked with `@pytest.mark.slow`
    """
    parser.addoption(
        "--skip-slow", action="store_true", default=False, help="Skip slow tests"
    )
    parser.addoption(
        "--overwrite-expected-visualizations",
        action="store_true",
        default=False,
        help="Overwrite the expected viz files in tests/data/svg/",
    )
    parser.addoption(
        "--draw-svg-debug-box",
        action="store_true",
        default=False,
        help="To help debugging, draw lines around the plotboxes in SVG output files",
    )


def pytest_configure(config):
    """
    Add docs on the "slow" marker
    """
    config.addinivalue_line("markers", "slow: mark test as slow to run")


def pytest_collection_modifyitems(config, items):
    if config.getoption("--skip-slow"):
        skip_slow = pytest.mark.skip(reason="--skip-slow specified")
        for item in items:
            if "slow" in item.keywords:
                item.add_marker(skip_slow)


@fixture
def overwrite_viz(request):
    return request.config.getoption("--overwrite-expected-visualizations")


@fixture
def draw_plotbox(request):
    return request.config.getoption("--draw-svg-debug-box")


@fixture(scope="session")
def simple_degree1_ts_fixture():
    return msprime.simulate(10, random_seed=42)


@fixture(scope="session")
def simple_degree2_ts_fixture():
    ts = msprime.simulate(10, recombination_rate=0.2, random_seed=42)
    assert ts.num_trees == 2
    return ts


@fixture(scope="session")
def ts_fixture():
    """
    A tree sequence with data in all fields
    """
    return tsutil.all_fields_ts()


@fixture(scope="session")
def ts_fixture_for_simplify():
    """
    A tree sequence with data in all fields execpt edge metadata and migrations
    """
    return tsutil.all_fields_ts(edge_metadata=False, migrations=False)


@fixture(scope="session")
def replicate_ts_fixture():
    """
    A list of tree sequences
    """
    return list(msprime.simulate(10, num_replicates=10, random_seed=42))


--- ../../tskit/python/tests/test_util.py ---

 
"""
Tests for functions in util.py
"""
import collections
import itertools
import math
import pickle
import textwrap

import numpy as np
import pytest
from numpy.testing import assert_array_equal

import tests.tsutil as tsutil
import tskit
import tskit.util as util
from tskit import UNKNOWN_TIME


class TestCanonicalJSON:
    def test_canonical_json(self):
        assert util.canonical_json([3, 2, 1]) == "[3,2,1]"
        assert (
            util.canonical_json(collections.OrderedDict(c=3, b=2, a=1))
            == '{"a":1,"b":2,"c":3}'
        )
        assert (
            util.canonical_json(
                collections.OrderedDict(
                    c="3",
                    b=collections.OrderedDict(
                        {
                            "b": 1,
                            "z": {},
                            " space": 42,
                            "1": "number",
                            "_": "underscore",
                        }
                    ),
                    a="1",
                )
            )
            == '{"a":"1","b":{" space":42,"1":"number",'
            '"_":"underscore","b":1,"z":{}},"c":"3"}'
        )


class TestUnknownTime:
    def test_unknown_time_bad_types(self):
        with pytest.raises(ValueError):
            util.is_unknown_time("bad")
        with pytest.raises(ValueError):
            util.is_unknown_time(np.array(["bad"]))
        with pytest.raises(ValueError):
            util.is_unknown_time(["bad"])

    def test_unknown_time_scalar(self):
        assert math.isnan(UNKNOWN_TIME)
        assert util.is_unknown_time(UNKNOWN_TIME)
        assert not util.is_unknown_time(math.nan)
        assert not util.is_unknown_time(np.nan)
        assert not util.is_unknown_time(0)
        assert not util.is_unknown_time(math.inf)
        assert not util.is_unknown_time(1)
        assert not util.is_unknown_time(None)
        assert not util.is_unknown_time([None])

    def test_unknown_time_array(self):
        test_arrays = (
            [],
            [True],
            [False],
            [True, False] * 5,
            [[True], [False]],
            [[[True, False], [True, False]], [[False, True], [True, False]]],
        )
        for spec in test_arrays:
            spec = np.asarray(spec, dtype=bool)
            array = np.zeros(shape=spec.shape)
            array[spec] = UNKNOWN_TIME
            assert_array_equal(spec, util.is_unknown_time(array))

        weird_array = [0, UNKNOWN_TIME, np.nan, 1, math.inf]
        assert_array_equal(
            [False, True, False, False, False], util.is_unknown_time(weird_array)
        )


class TestNumpyArrayCasting:
    """
    Tests that the safe_np_int_cast() function works.
    """

    dtypes_to_test = [np.int32, np.uint32, np.int8, np.uint8]

    def test_basic_arrays(self):
        # Simple array
        for dtype in self.dtypes_to_test:
            target = np.array([0, 1], dtype=dtype)
            for test_array in [[0, 1], (0, 1), np.array([0, 1]), target]:
                converted = util.safe_np_int_cast(test_array, dtype=dtype)
                # Use pickle to test exact equality including dtype
                assert pickle.dumps(converted) == pickle.dumps(target)
            # Nested array
            target = np.array([[0, 1], [2, 3]], dtype=dtype)
            for test_array in [[[0, 1], [2, 3]], np.array([[0, 1], [2, 3]]), target]:
                converted = util.safe_np_int_cast(test_array, dtype=dtype)
                assert pickle.dumps(converted) == pickle.dumps(target)

    def test_copy(self):
        # Check that a copy is not returned if copy=False & the original matches
        # the specs
        for dtype in self.dtypes_to_test:
            for orig in (np.array([0, 1], dtype=dtype), np.array([], dtype=dtype)):
                converted = util.safe_np_int_cast(orig, dtype=dtype, copy=True)
                assert id(orig) != id(converted)
                converted = util.safe_np_int_cast(orig, dtype=dtype, copy=False)
                assert id(orig) == id(converted)
        for dtype in [d for d in self.dtypes_to_test if d != np.int64]:
            # non numpy arrays, or arrays of a different dtype don't get converted
            for orig in ([0, 1], np.array([0, 1], dtype=np.int64)):
                converted = util.safe_np_int_cast(orig, dtype=dtype, copy=False)
                assert id(orig) != id(converted)

    def test_empty_arrays(self):
        # Empty arrays of any type (including float) should be allowed
        for dtype in self.dtypes_to_test:
            target = np.array([], dtype=dtype)
            converted = util.safe_np_int_cast([], dtype=dtype)
            assert pickle.dumps(converted) == pickle.dumps(target)
            target = np.array([[]], dtype=dtype)
            converted = util.safe_np_int_cast([[]], dtype=dtype)
            assert pickle.dumps(converted) == pickle.dumps(target)

    def test_bad_types(self):
        # Shouldn't be able to convert a float (possibility of rounding error)
        for dtype in self.dtypes_to_test:
            for bad_type in [
                [0.1],
                ["str"],
                {},
                [{}],
                np.array([0, 1], dtype=float),
            ]:
                with pytest.raises(TypeError):
                    util.safe_np_int_cast(bad_type, dtype)

    def test_overflow(self):
        for dtype in self.dtypes_to_test:
            for bad_node in [np.iinfo(dtype).min - 1, np.iinfo(dtype).max + 1]:
                with pytest.raises(OverflowError):
                    util.safe_np_int_cast([0, bad_node], dtype)
                with pytest.raises(OverflowError):
                    util.safe_np_int_cast(np.array([0, bad_node]), dtype)
            for good_node in [np.iinfo(dtype).min, np.iinfo(dtype).max]:
                target = np.array([good_node], dtype=dtype)
                assert pickle.dumps(target) == pickle.dumps(
                    util.safe_np_int_cast([good_node], dtype)
                )
                assert pickle.dumps(target) == pickle.dumps(
                    util.safe_np_int_cast(np.array([good_node]), dtype)
                )

    def test_nonrectangular_input(self):
        bad_inputs = [
            [0, 1, [2]],
            [[0, 1, 2], []],
            [(0, 1, 2), [2, 3]],
            [(0, 1, 2), tuple()],
            [(0, 1, 2), (2,)],
            [(0, 1, 2), [2, 3]],
        ]
        for dtype in self.dtypes_to_test:
            for bad_input in bad_inputs:
                # On some platforms and Python / numpy versions, a ValueError
                # occurs instead
                with pytest.raises((TypeError, ValueError)):
                    util.safe_np_int_cast(bad_input, dtype)


class TestIntervalOps:
    """
    Test cases for the interval operations used in masks and slicing operations.
    """

    def test_bad_intervals(self):
        for bad_type in [{}, Exception]:
            with pytest.raises(TypeError):
                util.intervals_to_np_array(bad_type, 0, 1)
        for bad_depth in [[[[]]]]:
            with pytest.raises(ValueError):
                util.intervals_to_np_array(bad_depth, 0, 1)
        for bad_shape in [[[0], [0]], [[[0, 1, 2], [0, 1]]]]:
            with pytest.raises(ValueError):
                util.intervals_to_np_array(bad_shape, 0, 1)

        # Out of bounds
        with pytest.raises(ValueError):
            util.intervals_to_np_array([[-1, 0]], 0, 1)
        with pytest.raises(ValueError):
            util.intervals_to_np_array([[0, 1]], 1, 2)
        with pytest.raises(ValueError):
            util.intervals_to_np_array([[0, 1]], 0, 0.5)

        # Overlapping intervals
        with pytest.raises(ValueError):
            util.intervals_to_np_array([[0, 1], [0.9, 2.0]], 0, 10)

        # Empty intervals
        for bad_interval in [[0, 0], [1, 0]]:
            with pytest.raises(ValueError):
                util.intervals_to_np_array([bad_interval], 0, 10)

    def test_empty_interval_list(self):
        intervals = util.intervals_to_np_array([], 0, 10)
        assert len(intervals) == 0

    def test_negate_intervals(self):
        L = 10
        cases = [
            ([], [[0, L]]),
            ([[0, 5], [6, L]], [[5, 6]]),
            ([[0, 5]], [[5, L]]),
            ([[5, L]], [[0, 5]]),
            ([[0, 1], [2, 3], [3, 4], [5, 6]], [[1, 2], [4, 5], [6, L]]),
        ]
        for source, dest in cases:
            assert np.array_equal(util.negate_intervals(source, 0, L), dest)


class TestStringPacking:
    """
    Tests the code for packing and unpacking unicode string data into numpy arrays.
    """

    def test_simple_string_case(self):
        strings = ["hello", "world"]
        packed, offset = util.pack_strings(strings)
        assert list(offset) == [0, 5, 10]
        assert packed.shape == (10,)
        returned = util.unpack_strings(packed, offset)
        assert returned == strings

    def verify_packing(self, strings):
        packed, offset = util.pack_strings(strings)
        assert packed.dtype == np.int8
        assert offset.dtype == np.uint32
        assert packed.shape[0] == offset[-1]
        returned = util.unpack_strings(packed, offset)
        assert strings == returned

    def test_regular_cases(self):
        for n in range(10):
            strings = ["a" * j for j in range(n)]
            self.verify_packing(strings)

    def test_random_cases(self):
        for n in range(100):
            strings = [tsutil.random_strings(10) for _ in range(n)]
            self.verify_packing(strings)

    def test_unicode(self):
        self.verify_packing(["abcdé", "€"])


class TestBytePacking:
    """
    Tests the code for packing and unpacking binary data into numpy arrays.
    """

    def test_simple_string_case(self):
        strings = [b"hello", b"world"]
        packed, offset = util.pack_bytes(strings)
        assert list(offset) == [0, 5, 10]
        assert packed.shape == (10,)
        returned = util.unpack_bytes(packed, offset)
        assert returned == strings

    def verify_packing(self, data):
        packed, offset = util.pack_bytes(data)
        assert packed.dtype == np.int8
        assert offset.dtype == np.uint32
        assert packed.shape[0] == offset[-1]
        returned = util.unpack_bytes(packed, offset)
        assert data == returned
        return returned

    def test_random_cases(self):
        for n in range(100):
            data = [tsutil.random_bytes(10) for _ in range(n)]
            self.verify_packing(data)

    def test_pickle_packing(self):
        data = [list(range(j)) for j in range(10)]
        # Pickle each of these in turn
        pickled = [pickle.dumps(d) for d in data]
        unpacked = self.verify_packing(pickled)
        unpickled = [pickle.loads(p) for p in unpacked]
        assert data == unpickled


class TestArrayPacking:
    """
    Tests the code for packing and unpacking numpy data into numpy arrays.
    """

    def test_simple_case(self):
        lists = [[0], [1.125, 1.25]]
        packed, offset = util.pack_arrays(lists)
        assert list(offset) == [0, 1, 3]
        assert list(packed) == [0, 1.125, 1.25]
        returned = util.unpack_arrays(packed, offset)
        for a1, a2 in itertools.zip_longest(lists, returned):
            assert a1 == list(a2)

    def verify_packing(self, data):
        packed, offset = util.pack_arrays(data)
        assert packed.dtype == np.float64
        assert offset.dtype == np.uint32
        assert packed.shape[0] == offset[-1]
        returned = util.unpack_arrays(packed, offset)
        for a1, a2 in itertools.zip_longest(data, returned):
            assert np.array_equal(a1, a2)
        return returned

    def test_regular_cases(self):
        for n in range(100):
            data = [np.arange(n) for _ in range(n)]
            self.verify_packing(data)
            data = [1 / (1 + np.arange(n)) for _ in range(n)]
            self.verify_packing(data)


@pytest.mark.parametrize(
    "value, expected",
    [
        (0, "0 Bytes"),
        (1, "1 Byte"),
        (300, "300 Bytes"),
        (3000, "2.9 KiB"),
        (3000000, "2.9 MiB"),
        (10**26 * 30, "2481.5 YiB"),
    ],
)
def test_naturalsize(value, expected):
    assert util.naturalsize(value) == expected
    if value != 0:
        assert util.naturalsize(-value) == "-" + expected
    else:
        assert util.naturalsize(-value) == expected


@pytest.mark.parametrize(
    "obj, expected",
    [
        (0, "Test: 0"),
        (
            {"a": 1},
            '<div><span class="tskit-details-label">Test:</span><details open>'
            "<summary>dict</summary>a: 1<br/></details></div>",
        ),
        (
            {"b": [1, 2, 3]},
            '<div><span class="tskit-details-label">Test:</span><details open>'
            '<summary>dict</summary><div><span class="tskit-details-label">b:'
            "</span><details ><summary>list</summary> 1<br/> 2<br/> 3<br/></"
            "details></div><br/></details></div>",
        ),
        (
            {"b": [1, 2, {"c": 1}]},
            '<div><span class="tskit-details-label">Test:</span><details open>'
            '<summary>dict</summary><div><span class="tskit-details-label">b:'
            "</span><details ><summary>list</summary> 1<br/> 2<br/><div><span"
            ' class="tskit-details-label"></span><details ><summary>dict</'
            "summary>c: 1<br/></details></div><br/></details></div><br/></"
            "details></div>",
        ),
        (
            {"a": "1", "b": "2"},
            '<div><span class="tskit-details-label">Test:</span><details open>'
            "<summary>dict</summary>a: 1<br/>b: 2<br/></details></div>",
        ),
        (
            {"a": "a very long thing that is broken in the output"},
            '<div><span class="tskit-details-label">Test:</span><details open>'
            "<summary>dict</summary>a: a very long thing that is<br/>broken in"
            " the output<br/></details></div>",
        ),
    ],
    ids=[
        "integer",
        "simple_dict",
        "dict_with_list",
        "nested_dict_list",
        "dict_with_strings",
        "dict_with_multiline_strings",
    ],
)
def test_obj_to_collapsed_html(obj, expected):
    assert (
        util.obj_to_collapsed_html(obj, "Test", 1).replace("  ", "").replace("\n", "")
        == expected
    )


def test_truncate_string_end():
    assert util.truncate_string_end("testing", 40) == "testing"
    assert util.truncate_string_end("testing", 7) == "testing"
    assert util.truncate_string_end("testing", 5) == "te..."


def test_render_metadata():
    assert util.render_metadata({}) == "{}"
    assert util.render_metadata("testing") == "testing"
    assert util.render_metadata(b"testing") == "b'testing'"
    assert util.render_metadata(b"testing", 6) == "b't..."
    assert util.render_metadata(b"") == ""


def test_unicode_table():
    assert (
        util.unicode_table(
            [["5", "6", "7", "8"], ["90", "10", "11", "12"]],
            header=["1", "2", "3", "4"],
        )
        == textwrap.dedent(
            """
           ╔══╤══╤══╤══╗
           ║1 │2 │3 │4 ║
           ╠══╪══╪══╪══╣
           ║5 │ 6│ 7│ 8║
           ╟──┼──┼──┼──╢
           ║90│10│11│12║
           ╚══╧══╧══╧══╝
        """
        )[1:]
    )

    assert (
        util.unicode_table(
            [
                ["1", "2", "3", "4"],
                ["5", "6", "7", "8"],
                "__skipped__",
                ["90", "10", "11", "12"],
            ],
            title="TITLE",
        )
        == textwrap.dedent(
            """
           ╔═══════════╗
           ║TITLE      ║
           ╠══╤══╤══╤══╣
           ║1 │ 2│ 3│ 4║
           ╟──┼──┼──┼──╢
           ║5 │ 6│ 7│ 8║
           ╟──┴──┴──┴──╢
           ║ rows skipp║
           ╟──┬──┬──┬──╢
           ║90│10│11│12║
           ╚══╧══╧══╧══╝
        """
        )[1:]
    )

    assert (
        util.unicode_table(
            [["1", "2", "3", "4"], ["5", "6", "7", "8"], ["90", "10", "11", "12"]],
            title="TITLE",
            row_separator=False,
        )
        == textwrap.dedent(
            """
           ╔═══════════╗
           ║TITLE      ║
           ╠══╤══╤══╤══╣
           ║1 │ 2│ 3│ 4║
           ║5 │ 6│ 7│ 8║
           ║90│10│11│12║
           ╚══╧══╧══╧══╝
        """
        )[1:]
    )


def test_unicode_table_column_alignments():
    assert (
        util.unicode_table(
            [["5", "6", "7", "8"], ["90", "10", "11", "12"]],
            header=["1", "2", "3", "4"],
            column_alignments="<>><",
        )
        == textwrap.dedent(
            """
           ╔══╤══╤══╤══╗
           ║1 │2 │3 │4 ║
           ╠══╪══╪══╪══╣
           ║5 │ 6│ 7│8 ║
           ╟──┼──┼──┼──╢
           ║90│10│11│12║
           ╚══╧══╧══╧══╝
        """
        )[1:]
    )


def test_set_printoptions():
    assert tskit._print_options == {"max_lines": 40}
    util.set_print_options(max_lines=None)
    assert tskit._print_options == {"max_lines": None}
    util.set_print_options(max_lines=40)
    assert tskit._print_options == {"max_lines": 40}
    with pytest.raises(TypeError):
        util.set_print_options(40)


class TestRandomNuceotides:
    @pytest.mark.parametrize("length", [0, 1, 10, 10.0, np.array([10])[0]])
    def test_length(self, length):
        s = tskit.random_nucleotides(length, seed=42)
        assert len(s) == length
        assert isinstance(s, str)

    def test_default_alphabet(self):
        s = tskit.random_nucleotides(100, seed=42)
        assert "".join(sorted(set(s))) == "ACGT"

    def test_length_keyword(self):
        s1 = tskit.random_nucleotides(length=10, seed=42)
        s2 = tskit.random_nucleotides(length=10, seed=42)
        assert s1 == s2

    def test_length_required(self):
        with pytest.raises(TypeError, match="required positional"):
            tskit.random_nucleotides()

    def test_seed_keyword_only(self):
        with pytest.raises(TypeError, match="1 positional"):
            tskit.random_nucleotides(10, 42)

    @pytest.mark.parametrize("seed", [1, 2, 3])
    def test_seed_equality(self, seed):
        s1 = tskit.random_nucleotides(10, seed=seed)
        s2 = tskit.random_nucleotides(10, seed=seed)
        assert s1 == s2

    def test_different_seed_not_equal(self):
        s1 = tskit.random_nucleotides(20, seed=1)
        s2 = tskit.random_nucleotides(20, seed=2)
        assert s1 != s2

    def test_no_seed_different_values(self):
        s1 = tskit.random_nucleotides(20)
        s2 = tskit.random_nucleotides(20)
        assert s1 != s2

    @pytest.mark.parametrize("length", ["0", 0.1, np.array([1.1])[0]])
    def test_length_bad_value(self, length):
        with pytest.raises(ValueError, match="must be an integer"):
            tskit.random_nucleotides(length)

    @pytest.mark.parametrize("length", [{}, None])
    def test_length_bad_type(self, length):
        with pytest.raises(TypeError, match="argument must be a string"):
            tskit.random_nucleotides(length)


--- ../../tskit/python/tests/simplify.py ---


"""
Python implementation of the simplify algorithm.
"""
import sys

import numpy as np
import portion

import tskit


def overlapping_segments(segments):
    """
    Returns an iterator over the (left, right, X) tuples describing the
    distinct overlapping segments in the specified set.
    """
    S = sorted(segments, key=lambda x: x.left)
    n = len(S)
    # Insert a sentinel at the end for convenience.
    S.append(Segment(sys.float_info.max, 0))
    right = S[0].left
    X = []
    j = 0
    while j < n:
        # Remove any elements of X with right <= left
        left = right
        X = [x for x in X if x.right > left]
        if len(X) == 0:
            left = S[j].left
        while j < n and S[j].left == left:
            X.append(S[j])
            j += 1
        j -= 1
        right = min(x.right for x in X)
        right = min(right, S[j + 1].left)
        yield left, right, X
        j += 1

    while len(X) > 0:
        left = right
        X = [x for x in X if x.right > left]
        if len(X) > 0:
            right = min(x.right for x in X)
            yield left, right, X


class Segment:
    """
    A class representing a single segment. Each segment has a left and right,
    denoting the loci over which it spans, a node and a next, giving the next
    in the chain.

    The node it records is the *output* node ID.
    """

    def __init__(self, left=None, right=None, node=None, next_segment=None):
        self.left = left
        self.right = right
        self.node = node
        self.next = next_segment

    def __str__(self):
        s = "({}-{}->{}:next={})".format(
            self.left, self.right, self.node, repr(self.next)
        )
        return s

    def __repr__(self):
        return repr((self.left, self.right, self.node))

    def __lt__(self, other):
        return (self.left, self.right, self.node) < (other.left, other.right, self.node)


class Simplifier:
    """
    Simplifies a tree sequence to its minimal representation given a subset
    of the leaves.
    """

    def __init__(
        self,
        ts,
        sample,
        reduce_to_site_topology=False,
        filter_sites=True,
        filter_populations=True,
        filter_individuals=True,
        keep_unary=False,
        keep_unary_in_individuals=False,
        keep_input_roots=False,
        filter_nodes=True,
        update_sample_flags=True,
    ):
        self.ts = ts
        self.n = len(sample)
        self.reduce_to_site_topology = reduce_to_site_topology
        self.sequence_length = ts.sequence_length
        self.filter_sites = filter_sites
        self.filter_populations = filter_populations
        self.filter_individuals = filter_individuals
        self.filter_nodes = filter_nodes
        self.update_sample_flags = update_sample_flags
        self.keep_unary = keep_unary
        self.keep_unary_in_individuals = keep_unary_in_individuals
        self.keep_input_roots = keep_input_roots
        self.num_mutations = ts.num_mutations
        self.input_sites = list(ts.sites())
        self.A_head = [None for _ in range(ts.num_nodes)]
        self.A_tail = [None for _ in range(ts.num_nodes)]
        self.tables = self.ts.tables.copy()
        self.tables.clear()
        self.edge_buffer = {}
        self.node_id_map = np.zeros(ts.num_nodes, dtype=np.int32) - 1
        self.is_sample = np.zeros(ts.num_nodes, dtype=np.int8)
        self.mutation_node_map = [-1 for _ in range(self.num_mutations)]
        self.samples = set(sample)
        self.sort_offset = -1
        # We keep a map of input nodes to mutations.
        self.mutation_map = [[] for _ in range(ts.num_nodes)]
        position = ts.sites_position
        site = ts.mutations_site
        node = ts.mutations_node
        for mutation_id in range(ts.num_mutations):
            site_position = position[site[mutation_id]]
            self.mutation_map[node[mutation_id]].append((site_position, mutation_id))

        for sample_id in sample:
            self.is_sample[sample_id] = 1

        if not self.filter_nodes:
            # NOTE In the C implementation we would really just not touch the
            # original tables.
            self.tables.nodes.replace_with(self.ts.tables.nodes)
            if self.update_sample_flags:
                flags = self.tables.nodes.flags
                # Zero out other sample flags
                flags = np.bitwise_and(
                    flags, np.uint32(~tskit.NODE_IS_SAMPLE & 0xFFFFFFFF)
                )
                flags[sample] |= tskit.NODE_IS_SAMPLE
                self.tables.nodes.flags = flags.astype(np.uint32)

            self.node_id_map[:] = np.arange(ts.num_nodes)
            for sample_id in sample:
                self.add_ancestry(sample_id, 0, self.sequence_length, sample_id)
        else:
            for sample_id in sample:
                output_id = self.record_node(sample_id)
                self.add_ancestry(sample_id, 0, self.sequence_length, output_id)

        self.position_lookup = None
        if self.reduce_to_site_topology:
            self.position_lookup = np.hstack([[0], position, [self.sequence_length]])

    def record_node(self, input_id):
        """
        Adds a new node to the output table corresponding to the specified input
        node ID.
        """
        node = self.ts.node(input_id)
        flags = node.flags
        if self.update_sample_flags:
            # Need to zero out the sample flag
            flags &= ~tskit.NODE_IS_SAMPLE
            if self.is_sample[input_id]:
                flags |= tskit.NODE_IS_SAMPLE
        output_id = self.tables.nodes.append(node.replace(flags=flags))
        self.node_id_map[input_id] = output_id
        return output_id

    def rewind_node(self, input_id, output_id):
        """
        Remove the mapping for the specified input and output node pair. This is
        done because there are no edges referring to the node.
        """
        assert output_id == len(self.tables.nodes) - 1
        assert output_id == self.node_id_map[input_id]
        self.tables.nodes.truncate(output_id)
        self.node_id_map[input_id] = -1

    def flush_edges(self):
        """
        Flush the edges to the output table after sorting and squashing
        any redundant records.
        """
        num_edges = 0
        for child in sorted(self.edge_buffer.keys()):
            for edge in self.edge_buffer[child]:
                self.tables.edges.append(edge)
                num_edges += 1
        self.edge_buffer.clear()
        return num_edges

    def record_edge(self, left, right, parent, child):
        """
        Adds an edge to the output list.
        """
        if self.reduce_to_site_topology:
            X = self.position_lookup
            left_index = np.searchsorted(X, left)
            right_index = np.searchsorted(X, right)
            # Find the smallest site position index greater than or equal to left
            # and right, i.e., slide each endpoint of an interval to the right
            # until they hit a site position. If both left and right map to the
            # the same position then we discard this edge. We also discard an
            # edge if left = 0 and right is less than the first site position.
            if left_index == right_index or (left_index == 0 and right_index == 1):
                return
            # Remap back to zero if the left end maps to the first site.
            if left_index == 1:
                left_index = 0
            left = X[left_index]
            right = X[right_index]
        if child not in self.edge_buffer:
            self.edge_buffer[child] = [tskit.Edge(left, right, parent, child)]
        else:
            last = self.edge_buffer[child][-1]
            if last.right == left:
                last.right = right
            else:
                self.edge_buffer[child].append(tskit.Edge(left, right, parent, child))

    def print_state(self):
        print(".................")
        print("Ancestors: ")
        num_nodes = len(self.A_tail)
        for j in range(num_nodes):
            print("\t", j, "->", end="")
            x = self.A_head[j]
            while x is not None:
                print(f"({x.left}-{x.right}->{x.node})", end="")
                x = x.next
            print()
        print("Mutation map:")
        for u in range(len(self.mutation_map)):
            v = self.mutation_map[u]
            if len(v) > 0:
                print("\t", u, "->", v)
        print("Node ID map: (input->output)")
        for input_id, output_id in enumerate(self.node_id_map):
            print("\t", input_id, "->", output_id)
        print("Mutation node map")
        for j in range(self.num_mutations):
            print("\t", j, "->", self.mutation_node_map[j])
        print("Output:")
        print(self.tables)
        self.check_state()

    def map_mutations(self, left, right, input_id, output_id):
        """
        Map any mutations for the input node ID on the
        interval to its output ID.
        """
        assert output_id != -1
        # TODO we should probably remove these as they are used.
        # Or else, binary search the list so it's quick.
        for x, mutation_id in self.mutation_map[input_id]:
            if left <= x < right:
                self.mutation_node_map[mutation_id] = output_id

    def add_ancestry(self, input_id, left, right, node):
        tail = self.A_tail[input_id]
        if tail is None:
            x = Segment(left, right, node)
            self.A_head[input_id] = x
            self.A_tail[input_id] = x
        else:
            if tail.right == left and tail.node == node:
                tail.right = right
            else:
                x = Segment(left, right, node)
                tail.next = x
                self.A_tail[input_id] = x

        self.map_mutations(left, right, input_id, node)

    def merge_labeled_ancestors(self, S, input_id):
        """
        All ancestry segments in S come together into a new parent.
        The new parent must be assigned and any overlapping segments coalesced.
        """
        output_id = self.node_id_map[input_id]
        is_sample = self.is_sample[input_id]
        if is_sample:
            # Free up the existing ancestry mapping.
            x = self.A_tail[input_id]
            assert x.left == 0 and x.right == self.sequence_length
            self.A_tail[input_id] = None
            self.A_head[input_id] = None

        prev_right = 0
        for left, right, X in overlapping_segments(S):
            if len(X) == 1:
                ancestry_node = X[0].node
                if is_sample:
                    self.record_edge(left, right, output_id, ancestry_node)
                    ancestry_node = output_id
                elif self.keep_unary or (
                    self.keep_unary_in_individuals
                    and self.ts.node(input_id).individual >= 0
                ):
                    if output_id == -1:
                        output_id = self.record_node(input_id)
                    self.record_edge(left, right, output_id, ancestry_node)
            else:
                if output_id == -1:
                    output_id = self.record_node(input_id)
                ancestry_node = output_id
                for x in X:
                    self.record_edge(left, right, output_id, x.node)
            if is_sample and left != prev_right:
                # Fill in any gaps in the ancestry for the sample
                self.add_ancestry(input_id, prev_right, left, output_id)
            if self.keep_unary or (
                self.keep_unary_in_individuals
                and self.ts.node(input_id).individual >= 0
            ):
                ancestry_node = output_id
            self.add_ancestry(input_id, left, right, ancestry_node)
            prev_right = right

        if is_sample and prev_right != self.sequence_length:
            # If a trailing gap exists in the sample ancestry, fill it in.
            self.add_ancestry(input_id, prev_right, self.sequence_length, output_id)
        if output_id != -1:
            num_edges = self.flush_edges()
            if self.filter_nodes and num_edges == 0 and not is_sample:
                self.rewind_node(input_id, output_id)

    def extract_ancestry(self, edge):
        S = []
        x = self.A_head[edge.child]

        x_head = None
        x_prev = None
        while x is not None:
            if x.right > edge.left and edge.right > x.left:
                y = Segment(max(x.left, edge.left), min(x.right, edge.right), x.node)
                # print("snip", y)
                S.append(y)
                assert x.left <= y.left
                assert x.right >= y.right
                seg_left = None
                seg_right = None
                if x.left != y.left:
                    seg_left = Segment(x.left, y.left, x.node)
                    if x_prev is None:
                        x_head = seg_left
                    else:
                        x_prev.next = seg_left
                    x_prev = seg_left
                if x.right != y.right:
                    x.left = y.right
                    seg_right = x
                else:
                    # Free x
                    seg_right = x.next
                if x_prev is None:
                    x_head = seg_right
                else:
                    x_prev.next = seg_right
                x = seg_right
            else:
                if x_prev is None:
                    x_head = x
                x_prev = x
                x = x.next
        # Note - we had some code to defragment segments in the output
        # chain here, but couldn't find an example where it needed to
        # be called. So, looks like squashing isn't necessary here.
        self.A_head[edge.child] = x_head
        self.A_tail[edge.child] = x_prev
        return S

    def process_parent_edges(self, edges):
        """
        Process all of the edges for a given parent.
        """
        assert len({e.parent for e in edges}) == 1
        parent = edges[0].parent
        S = []
        for edge in edges:
            S.extend(self.extract_ancestry(edge))
        self.merge_labeled_ancestors(S, parent)
        self.check_state()

    def finalise_sites(self):
        # Build a map from the old mutation IDs to new IDs. Any mutation that
        # has not been mapped to a node in the new tree sequence will be removed.
        mutation_id_map = [-1 for _ in range(self.num_mutations)]
        num_output_mutations = 0

        for site in self.ts.sites():
            num_output_site_mutations = 0
            for mut in site.mutations:
                mapped_node = self.mutation_node_map[mut.id]
                mapped_parent = -1
                if mut.parent != -1:
                    mapped_parent = mutation_id_map[mut.parent]
                if mapped_node != -1:
                    mutation_id_map[mut.id] = num_output_mutations
                    num_output_mutations += 1
                    num_output_site_mutations += 1
            output_site = True
            if self.filter_sites and num_output_site_mutations == 0:
                output_site = False

            if output_site:
                for mut in site.mutations:
                    if mutation_id_map[mut.id] != -1:
                        mapped_parent = -1
                        if mut.parent != -1:
                            mapped_parent = mutation_id_map[mut.parent]
                        self.tables.mutations.append(
                            mut.replace(
                                site=len(self.tables.sites),
                                node=self.mutation_node_map[mut.id],
                                parent=mapped_parent,
                            )
                        )
                self.tables.sites.append(site)

    def finalise_references(self):
        input_populations = self.ts.tables.populations
        population_id_map = np.arange(len(input_populations) + 1, dtype=np.int32)
        # Trick to ensure the null population gets mapped to itself.
        population_id_map[-1] = -1
        input_individuals = self.ts.tables.individuals
        individual_id_map = np.arange(len(input_individuals) + 1, dtype=np.int32)
        # Trick to ensure the null individual gets mapped to itself.
        individual_id_map[-1] = -1

        population_ref_count = np.ones(len(input_populations), dtype=int)
        if self.filter_populations:
            population_ref_count[:] = 0
            population_id_map[:] = -1
        individual_ref_count = np.ones(len(input_individuals), dtype=int)
        if self.filter_individuals:
            individual_ref_count[:] = 0
            individual_id_map[:] = -1

        for node in self.tables.nodes:
            if self.filter_populations and node.population != tskit.NULL:
                population_ref_count[node.population] += 1
            if self.filter_individuals and node.individual != tskit.NULL:
                individual_ref_count[node.individual] += 1

        for input_id, count in enumerate(population_ref_count):
            if count > 0:
                row = input_populations[input_id]
                output_id = self.tables.populations.append(row)
                population_id_map[input_id] = output_id
        for input_id, count in enumerate(individual_ref_count):
            if count > 0:
                row = input_individuals[input_id]
                output_id = self.tables.individuals.append(row)
                individual_id_map[input_id] = output_id

        # Remap the population ID references for nodes.
        nodes = self.tables.nodes
        nodes.set_columns(
            flags=nodes.flags,
            time=nodes.time,
            metadata=nodes.metadata,
            metadata_offset=nodes.metadata_offset,
            individual=individual_id_map[nodes.individual],
            population=population_id_map[nodes.population],
        )

        # Remap the parent ids of individuals
        individuals_copy = self.tables.individuals.copy()
        self.tables.individuals.clear()
        for row in individuals_copy:
            mapped_parents = []
            for p in row.parents:
                if p == -1:
                    mapped_parents.append(-1)
                else:
                    mapped_parents.append(individual_id_map[p])
            self.tables.individuals.append(row.replace(parents=mapped_parents))

        # We don't support migrations for now. We'll need to remap these as well.
        assert self.ts.num_migrations == 0

    def insert_input_roots(self):
        youngest_root_time = np.inf
        for input_id in range(len(self.node_id_map)):
            x = self.A_head[input_id]
            if x is not None:
                output_id = self.node_id_map[input_id]
                if output_id == -1:
                    output_id = self.record_node(input_id)
                while x is not None:
                    if x.node != output_id:
                        self.record_edge(x.left, x.right, output_id, x.node)
                        self.map_mutations(x.left, x.right, input_id, output_id)
                    x = x.next
                self.flush_edges()
                root_time = self.tables.nodes.time[output_id]
                if root_time < youngest_root_time:
                    youngest_root_time = root_time
        # We have to sort the edge table from the point where the edges
        # for the youngest root would be inserted.
        # Note: it would be nicer to do the sort here, but we have to
        # wait until the finalise_references method has been called to
        # make sure all the populations etc have been setup.
        node_time = self.tables.nodes.time
        edge_parent = self.tables.edges.parent
        offset = 0
        while (
            offset < len(self.tables.edges)
            and node_time[edge_parent[offset]] < youngest_root_time
        ):
            offset += 1
        self.sort_offset = offset

    def simplify(self):
        if self.ts.num_edges > 0:
            all_edges = list(self.ts.edges())
            edges = all_edges[:1]
            for e in all_edges[1:]:
                if e.parent != edges[0].parent:
                    self.process_parent_edges(edges)
                    edges = []
                edges.append(e)
            self.process_parent_edges(edges)
        if self.keep_input_roots:
            self.insert_input_roots()
        self.finalise_sites()
        self.finalise_references()
        if self.sort_offset != -1:
            self.tables.sort(edge_start=self.sort_offset)
        ts = self.tables.tree_sequence()
        return ts, self.node_id_map

    def check_state(self):
        # print("CHECK_STATE")
        all_ancestry = []
        num_nodes = len(self.A_head)
        for j in range(num_nodes):
            head = self.A_head[j]
            tail = self.A_tail[j]
            if head is None:
                assert tail is None
            else:
                x = head
                while x.next is not None:
                    assert x.right <= x.next.left
                    x = x.next
                assert x == tail
                x = head
                while x is not None:
                    assert x.left < x.right
                    all_ancestry.append(portion.openclosed(x.left, x.right))
                    if x.next is not None:
                        assert x.right <= x.next.left
                        # We should also not have any squashable segments.
                        if x.right == x.next.left:
                            assert x.node != x.next.node
                    x = x.next
        # Make sure we haven't lost ancestry.
        if len(all_ancestry) > 0:
            union = all_ancestry[0]
            for interval in all_ancestry[1:]:
                union = union.union(interval)
            assert union.atomic
            assert union.lower == 0
            assert union.upper == self.sequence_length


class AncestorMap:
    """
    Simplifies a tree sequence to show relationships between
    samples and a designated set of ancestors.
    """

    def __init__(self, ts, sample, ancestors):
        self.ts = ts
        self.samples = set(sample)
        assert (self.samples).issubset(set(range(0, ts.num_nodes)))
        self.ancestors = set(ancestors)
        assert (self.ancestors).issubset(set(range(0, ts.num_nodes)))
        self.table = tskit.EdgeTable()
        self.sequence_length = ts.sequence_length
        self.A_head = [None for _ in range(ts.num_nodes)]
        self.A_tail = [None for _ in range(ts.num_nodes)]
        for sample_id in sample:
            self.add_ancestry(0, self.sequence_length, sample_id, sample_id)
        self.edge_buffer = {}
        self.oldest_ancestor_time = max(ts.nodes_time[u] for u in ancestors)
        self.oldest_sample_time = max(ts.nodes_time[u] for u in sample)
        self.oldest_node_time = max(self.oldest_ancestor_time, self.oldest_sample_time)

    def link_ancestors(self):
        if self.ts.num_edges > 0:
            all_edges = list(self.ts.edges())
            edges = all_edges[:1]
            for e in all_edges[1:]:
                if self.ts.tables.nodes.time[e.parent] > self.oldest_node_time:
                    break
                if e.parent != edges[0].parent:
                    self.process_parent_edges(edges)
                    edges = []
                edges.append(e)
            self.process_parent_edges(edges)
        return self.table

    def process_parent_edges(self, edges):
        """
        Process all of the edges for a given parent.
        """
        assert len({e.parent for e in edges}) == 1
        parent = edges[0].parent
        S = []
        for edge in edges:
            x = self.A_head[edge.child]
            while x is not None:
                if x.right > edge.left and edge.right > x.left:
                    y = Segment(
                        max(x.left, edge.left), min(x.right, edge.right), x.node
                    )
                    S.append(y)
                x = x.next
        self.merge_labeled_ancestors(S, parent)
        self.check_state()

    def merge_labeled_ancestors(self, S, input_id):
        """
        All ancestry segments in S come together into a new parent.
        The new parent must be assigned and any overlapping segments coalesced.
        """
        is_sample = input_id in self.samples
        if is_sample:
            # Free up the existing ancestry mapping.
            x = self.A_tail[input_id]
            assert x.left == 0 and x.right == self.sequence_length
            self.A_tail[input_id] = None
            self.A_head[input_id] = None

        is_ancestor = input_id in self.ancestors
        prev_right = 0
        for left, right, X in overlapping_segments(S):
            if is_ancestor or is_sample:
                for x in X:
                    ancestry_node = x.node
                    self.record_edge(left, right, input_id, ancestry_node)
                self.add_ancestry(left, right, input_id, input_id)

                if is_sample and left != prev_right:
                    # Fill in any gaps in the ancestry for the sample.
                    self.add_ancestry(prev_right, left, input_id, input_id)

            else:
                for x in X:
                    ancestry_node = x.node
                    # Add sample ancestry for the currently-processed segment set.
                    self.add_ancestry(left, right, ancestry_node, input_id)
            prev_right = right

        if is_sample and prev_right != self.sequence_length:
            # If a trailing gap exists in the sample ancestry, fill it in.
            self.add_ancestry(prev_right, self.sequence_length, input_id, input_id)
        if input_id != -1:
            self.flush_edges()

    def record_edge(self, left, right, parent, child):
        """
        Adds an edge to the output list.
        """
        if child not in self.edge_buffer:
            self.edge_buffer[child] = [tskit.Edge(left, right, parent, child)]
        else:
            last = self.edge_buffer[child][-1]
            if last.right == left:
                last.right = right
            else:
                self.edge_buffer[child].append(tskit.Edge(left, right, parent, child))

    def add_ancestry(self, left, right, node, current_node):
        tail = self.A_tail[current_node]
        if tail is None:
            x = Segment(left, right, node)
            self.A_head[current_node] = x
            self.A_tail[current_node] = x
        else:
            if tail.right == left and tail.node == node:
                tail.right = right
            else:
                x = Segment(left, right, node)
                tail.next = x
                self.A_tail[current_node] = x

    def flush_edges(self):
        """
        Flush the edges to the output table after sorting and squashing
        any redundant records.
        """
        num_edges = 0
        for child in sorted(self.edge_buffer.keys()):
            for edge in self.edge_buffer[child]:
                self.table.append(edge)
                num_edges += 1
        self.edge_buffer.clear()
        return num_edges

    def check_state(self):
        num_nodes = len(self.A_head)
        for j in range(num_nodes):
            head = self.A_head[j]
            tail = self.A_tail[j]
            if head is None:
                assert tail is None
            else:
                x = head
                while x.next is not None:
                    x = x.next
                assert x == tail
                x = head.next
                while x is not None:
                    assert x.left < x.right
                    if x.next is not None:
                        if self.ancestors is None:
                            assert x.right <= x.next.left
                        # We should also not have any squashable segments.
                        if x.right == x.next.left:
                            assert x.node != x.next.node
                    x = x.next

    def print_state(self):
        print(".................")
        print("Ancestors: ")
        num_nodes = len(self.A_tail)
        for j in range(num_nodes):
            print("\t", j, "->", end="")
            x = self.A_head[j]
            while x is not None:
                print(f"({x.left}-{x.right}->{x.node})", end="")
                x = x.next
            print()
        print("Output:")
        print(self.table)
        self.check_state()


if __name__ == "__main__":
    # Simple CLI for running simplifier/ancestor mapping above.
    class_to_implement = sys.argv[1]
    assert class_to_implement == "Simplifier" or class_to_implement == "AncestorMap"
    ts = tskit.load(sys.argv[2])

    if class_to_implement == "Simplifier":
        samples = list(map(int, sys.argv[3:]))

        print("When keep_unary = True:")
        s = Simplifier(ts, samples, keep_unary=True)
        # s.print_state()
        tss, _ = s.simplify()
        tables = tss.dump_tables()
        print(tables.nodes)
        print(tables.edges)
        print(tables.sites)
        print(tables.mutations)

        print("\nWhen keep_unary = False")
        s = Simplifier(ts, samples, keep_unary=False)
        # s.print_state()
        tss, _ = s.simplify()
        tables = tss.dump_tables()
        print(tables.nodes)
        print(tables.edges)
        print(tables.sites)
        print(tables.mutations)

        print("\nWhen keep_unary_in_individuals = True")
        s = Simplifier(ts, samples, keep_unary_in_individuals=True)
        # s.print_state()
        tss, _ = s.simplify()
        tables = tss.dump_tables()
        print(tables.nodes)
        print(tables.edges)
        print(tables.sites)
        print(tables.mutations)

    elif class_to_implement == "AncestorMap":
        samples = sys.argv[3]
        samples = samples.split(",")
        samples = list(map(int, samples))

        ancestors = sys.argv[4]
        ancestors = ancestors.split(",")
        ancestors = list(map(int, ancestors))

        s = AncestorMap(ts, samples, ancestors)
        tss = s.link_ancestors()
        # tables = tss.dump_tables()
        # print(tables.nodes)
        print(tss)


--- ../../tskit/python/tests/test_tree_positioning.py ---


"""
Tests for tree iterator schemes. Mostly used to develop the incremental
iterator infrastructure.
"""
import msprime
import numpy as np
import pytest

import tests
import tskit
from tests import tsutil
from tests.test_highlevel import get_example_tree_sequences

# ↑ See https://github.com/tskit-dev/tskit/issues/1804 for when
# we can remove this.


class StatefulTree:
    """
    Just enough functionality to mimic the low-level tree implementation
    for testing of forward/backward moving.
    """

    def __init__(self, ts):
        self.ts = ts
        self.tree_pos = tsutil.TreePosition(ts)
        self.parent = [-1 for _ in range(ts.num_nodes)]

    def __str__(self):
        s = f"parent: {self.parent}\nposition:\n"
        for line in str(self.tree_pos).splitlines():
            s += f"\t{line}\n"
        return s

    def assert_equal(self, other):
        assert self.parent == other.parent
        assert self.tree_pos.index == other.tree_pos.index
        assert self.tree_pos.interval == other.tree_pos.interval

    def next(self):  # NOQA: A003
        valid = self.tree_pos.next()
        if valid:
            for j in range(self.tree_pos.out_range.start, self.tree_pos.out_range.stop):
                e = self.tree_pos.out_range.order[j]
                c = self.ts.edges_child[e]
                self.parent[c] = -1
            for j in range(self.tree_pos.in_range.start, self.tree_pos.in_range.stop):
                e = self.tree_pos.in_range.order[j]
                c = self.ts.edges_child[e]
                p = self.ts.edges_parent[e]
                self.parent[c] = p
        return valid

    def prev(self):
        valid = self.tree_pos.prev()
        if valid:
            for j in range(
                self.tree_pos.out_range.start, self.tree_pos.out_range.stop, -1
            ):
                e = self.tree_pos.out_range.order[j]
                c = self.ts.edges_child[e]
                self.parent[c] = -1
            for j in range(
                self.tree_pos.in_range.start, self.tree_pos.in_range.stop, -1
            ):
                e = self.tree_pos.in_range.order[j]
                c = self.ts.edges_child[e]
                p = self.ts.edges_parent[e]
                self.parent[c] = p
        return valid

    def iter_forward(self, index):
        while self.tree_pos.index != index:
            self.next()

    def seek_forward(self, index):
        old_left, old_right = self.tree_pos.interval
        self.tree_pos.seek_forward(index)
        left, right = self.tree_pos.interval
        for j in range(self.tree_pos.out_range.start, self.tree_pos.out_range.stop):
            e = self.tree_pos.out_range.order[j]
            e_left = self.ts.edges_left[e]
            # We only need to remove an edge if it's in the current tree, which
            # can only happen if the edge's left coord is < the current tree's
            # right coordinate.
            if e_left < old_right:
                c = self.ts.edges_child[e]
                assert self.parent[c] != -1
                self.parent[c] = -1
            assert e_left < left
        for j in range(self.tree_pos.in_range.start, self.tree_pos.in_range.stop):
            e = self.tree_pos.in_range.order[j]
            if self.ts.edges_left[e] <= left < self.ts.edges_right[e]:
                c = self.ts.edges_child[e]
                p = self.ts.edges_parent[e]
                self.parent[c] = p
            else:
                a = self.tree_pos.in_range.start
                b = self.tree_pos.in_range.stop
                # The first and last indexes in the range should always be valid
                # for the tree.
                assert a < j < b - 1

    def seek_backward(self, index):
        old_left, old_right = self.tree_pos.interval
        self.tree_pos.seek_backward(index)
        left, right = self.tree_pos.interval
        for j in range(self.tree_pos.out_range.start, self.tree_pos.out_range.stop, -1):
            e = self.tree_pos.out_range.order[j]
            e_right = self.ts.edges_right[e]
            # We only need to remove an edge if it's in the current tree, which
            # can only happen if the edge's right coord is >= the current tree's
            # right coordinate.
            if e_right >= old_right:
                c = self.ts.edges_child[e]
                assert self.parent[c] != -1
                self.parent[c] = -1
            assert e_right > right
        for j in range(self.tree_pos.in_range.start, self.tree_pos.in_range.stop, -1):
            e = self.tree_pos.in_range.order[j]
            if self.ts.edges_right[e] >= right > self.ts.edges_left[e]:
                c = self.ts.edges_child[e]
                p = self.ts.edges_parent[e]
                self.parent[c] = p
            else:
                a = self.tree_pos.in_range.start
                b = self.tree_pos.in_range.stop
                # The first and last indexes in the range should always be valid
                # for the tree.
                assert a > j > b + 1

    def iter_backward(self, index):
        while self.tree_pos.index != index:
            self.prev()


def check_iters_forward(ts):
    alg_t_output = tsutil.algorithm_T(ts)
    lib_tree = tskit.Tree(ts)
    tree_pos = tsutil.TreePosition(ts)
    sample_count = np.zeros(ts.num_nodes, dtype=int)
    sample_count[ts.samples()] = 1
    parent1 = [-1 for _ in range(ts.num_nodes)]
    i = 0
    lib_tree.next()
    while tree_pos.next():
        out_times = []
        for j in range(tree_pos.out_range.start, tree_pos.out_range.stop):
            e = tree_pos.out_range.order[j]
            c = ts.edges_child[e]
            p = ts.edges_parent[e]
            out_times.append(ts.nodes_time[p])
            parent1[c] = -1
        in_times = []
        for j in range(tree_pos.in_range.start, tree_pos.in_range.stop):
            e = tree_pos.in_range.order[j]
            c = ts.edges_child[e]
            p = ts.edges_parent[e]
            in_times.append(ts.nodes_time[p])
            parent1[c] = p
        # We must visit the edges in *increasing* time order on the way in,
        # and *decreasing* order on the way out. Otherwise we get quadratic
        # behaviour for algorithms that need to propagate changes up to the
        # root.
        assert out_times == sorted(out_times, reverse=True)
        assert in_times == sorted(in_times)

        interval, parent2 = next(alg_t_output)
        assert list(interval) == list(tree_pos.interval)
        assert parent1 == parent2

        assert lib_tree.index == i
        assert list(lib_tree.interval) == list(interval)
        assert list(lib_tree.parent_array[:-1]) == parent1

        lib_tree.next()
        i += 1
    assert i == ts.num_trees
    assert lib_tree.index == -1
    assert next(alg_t_output, None) is None


def check_iters_back(ts):
    alg_t_output = [
        (list(interval), list(parent)) for interval, parent in tsutil.algorithm_T(ts)
    ]
    i = len(alg_t_output) - 1

    lib_tree = tskit.Tree(ts)
    tree_pos = tsutil.TreePosition(ts)
    parent1 = [-1 for _ in range(ts.num_nodes)]

    lib_tree.last()

    while tree_pos.prev():
        # print(tree_pos.out_range)
        out_times = []
        for j in range(tree_pos.out_range.start, tree_pos.out_range.stop, -1):
            e = tree_pos.out_range.order[j]
            c = ts.edges_child[e]
            p = ts.edges_parent[e]
            out_times.append(ts.nodes_time[p])
            parent1[c] = -1
        in_times = []
        for j in range(tree_pos.in_range.start, tree_pos.in_range.stop, -1):
            e = tree_pos.in_range.order[j]
            c = ts.edges_child[e]
            p = ts.edges_parent[e]
            in_times.append(ts.nodes_time[p])
            parent1[c] = p

        # We must visit the edges in *increasing* time order on the way in,
        # and *decreasing* order on the way out. Otherwise we get quadratic
        # behaviour for algorithms that need to propagate changes up to the
        # root.
        assert out_times == sorted(out_times, reverse=True)
        assert in_times == sorted(in_times)

        interval, parent2 = alg_t_output[i]
        assert list(interval) == list(tree_pos.interval)
        assert parent1 == parent2

        assert lib_tree.index == i
        assert list(lib_tree.interval) == list(interval)
        assert list(lib_tree.parent_array[:-1]) == parent1

        lib_tree.prev()
        i -= 1

    assert lib_tree.index == -1
    assert i == -1


def check_forward_back_sweep(ts):
    alg_t_output = [
        (list(interval), list(parent)) for interval, parent in tsutil.algorithm_T(ts)
    ]
    for j in range(ts.num_trees - 1):
        tree = StatefulTree(ts)
        # Seek forward to j
        k = 0
        while k <= j:
            tree.next()
            interval, parent = alg_t_output[k]
            assert tree.tree_pos.index == k
            assert list(tree.tree_pos.interval) == interval
            assert parent == tree.parent
            k += 1
        k = j
        # And back to zero
        while k >= 0:
            interval, parent = alg_t_output[k]
            assert tree.tree_pos.index == k
            assert list(tree.tree_pos.interval) == interval
            assert parent == tree.parent
            tree.prev()
            k -= 1


def check_seek_forward_out_range_is_empty(ts, index):
    tree = StatefulTree(ts)
    tree.seek_forward(index)
    assert tree.tree_pos.out_range.start == tree.tree_pos.out_range.stop
    tree.iter_backward(-1)
    tree.seek_forward(index)
    assert tree.tree_pos.out_range.start == tree.tree_pos.out_range.stop


def check_seek_backward_out_range_is_empty(ts, index):
    tree = StatefulTree(ts)
    tree.seek_backward(index)
    assert tree.tree_pos.out_range.start == tree.tree_pos.out_range.stop
    tree.iter_forward(-1)
    tree.seek_backward(index)
    assert tree.tree_pos.out_range.start == tree.tree_pos.out_range.stop


def check_seek_forward_from_null(ts, index):
    tree1 = StatefulTree(ts)
    tree1.seek_forward(index)
    tree2 = StatefulTree(ts)
    tree2.iter_forward(index)
    tree1.assert_equal(tree2)


def check_seek_backward_from_null(ts, index):
    tree1 = StatefulTree(ts)
    tree1.seek_backward(index)
    tree2 = StatefulTree(ts)
    tree2.iter_backward(index)
    tree1.assert_equal(tree2)


def check_seek_forward_from_first(ts, index):
    tree1 = StatefulTree(ts)
    tree1.next()
    tree1.seek_forward(index)
    tree2 = StatefulTree(ts)
    tree2.iter_forward(index)
    tree1.assert_equal(tree2)


def check_seek_backward_from_last(ts, index):
    tree1 = StatefulTree(ts)
    tree1.prev()
    tree1.seek_backward(index)
    tree2 = StatefulTree(ts)
    tree2.iter_backward(index)


class TestDirectionSwitching:
    # 2.00┊       ┊   4   ┊   4   ┊   4   ┊
    #     ┊       ┊ ┏━┻┓  ┊  ┏┻━┓ ┊  ┏┻━┓ ┊
    # 1.00┊   3   ┊ ┃  3  ┊  3  ┃ ┊  3  ┃ ┊
    #     ┊ ┏━╋━┓ ┊ ┃ ┏┻┓ ┊ ┏┻┓ ┃ ┊ ┏┻┓ ┃ ┊
    # 0.00┊ 0 1 2 ┊ 0 1 2 ┊ 0 2 1 ┊ 0 1 2 ┊
    #     0       1       2       3       4
    # index   0       1       2       3
    def ts(self):
        return tsutil.all_trees_ts(3)

    @pytest.mark.parametrize("index", [0, 1, 2, 3])
    def test_iter_backward_matches_iter_forward(self, index):
        ts = self.ts()
        tree1 = StatefulTree(ts)
        tree1.iter_forward(index)
        tree2 = StatefulTree(ts)
        tree2.iter_backward(index)
        tree1.assert_equal(tree2)

    @pytest.mark.parametrize("index", [1, 2, 3])
    def test_prev_from_seek_forward(self, index):
        tree1 = StatefulTree(self.ts())
        tree1.seek_forward(index)
        tree1.prev()
        tree2 = StatefulTree(self.ts())
        tree2.seek_forward(index - 1)
        tree1.assert_equal(tree2)

    @pytest.mark.parametrize("index", [1, 2, 3])
    def test_seek_forward_from_prev(self, index):
        tree1 = StatefulTree(self.ts())
        tree1.iter_forward(index)
        tree1.prev()
        tree1.seek_forward(index)
        tree2 = StatefulTree(self.ts())
        tree2.iter_forward(index)
        tree1.assert_equal(tree2)

    @pytest.mark.parametrize("index", [0, 1, 2, 3])
    def test_seek_forward_from_null(self, index):
        ts = self.ts()
        check_seek_forward_from_null(ts, index)

    def test_seek_forward_next_null(self):
        tree1 = StatefulTree(self.ts())
        tree1.seek_forward(3)
        tree1.next()
        assert tree1.tree_pos.index == -1
        assert list(tree1.tree_pos.interval) == [0, 0]

    @pytest.mark.parametrize("index", [0, 1, 2])
    def test_next_from_seek_backward(self, index):
        tree1 = StatefulTree(self.ts())
        tree1.seek_backward(index)
        tree1.next()
        tree2 = StatefulTree(self.ts())
        tree2.seek_backward(index + 1)
        tree1.assert_equal(tree2)

    @pytest.mark.parametrize("index", [0, 1, 2])
    def test_seek_backward_from_next(self, index):
        tree1 = StatefulTree(self.ts())
        tree1.iter_backward(index)
        tree1.next()
        tree1.seek_backward(index)
        tree2 = StatefulTree(self.ts())
        tree2.iter_backward(index)
        tree1.assert_equal(tree2)

    @pytest.mark.parametrize("index", [0, 1, 2, 3])
    def test_seek_backward_from_null(self, index):
        ts = self.ts()
        check_seek_backward_from_null(ts, index)

    def test_seek_backward_prev_null(self):
        tree1 = StatefulTree(self.ts())
        tree1.seek_backward(0)
        tree1.prev()
        assert tree1.tree_pos.index == -1
        assert list(tree1.tree_pos.interval) == [0, 0]

    @pytest.mark.parametrize("index", [0, 1, 2, 3])
    def test_seek_forward_out_range_is_empty(self, index):
        ts = self.ts()
        check_seek_forward_out_range_is_empty(ts, index)

    @pytest.mark.parametrize("index", [0, 1, 2, 3])
    def test_seek_backward_out_range_is_empty(self, index):
        ts = self.ts()
        check_seek_backward_out_range_is_empty(ts, index)


class TestTreePositionStep:
    def ts(self):
        return tsutil.all_trees_ts(3)

    @pytest.mark.parametrize("index", [0, 1, 2])
    def test_tree_position_step_forward(self, index):
        ts = self.ts()
        tree1_pos = tsutil.TreePosition(ts)
        tree1_pos.seek_forward(index)
        tree1_pos.step(direction=1)
        tree2_pos = tsutil.TreePosition(ts)
        tree2_pos.seek_forward(index + 1)
        tree1_pos.assert_equal(tree2_pos)

    @pytest.mark.parametrize("index", [1, 2, 3])
    def test_tree_position_step_backward(self, index):
        ts = self.ts()
        tree1_pos = tsutil.TreePosition(ts)
        tree1_pos.seek_backward(index)
        tree1_pos.step(direction=-1)
        tree2_pos = tsutil.TreePosition(ts)
        tree2_pos.seek_backward(index - 1)
        tree1_pos.assert_equal(tree2_pos)

    def test_tree_position_step_invalid_direction(self):
        ts = self.ts()
        # Test for unallowed direction
        with pytest.raises(ValueError, match="Direction must be FORWARD"):
            tsutil.TreePosition(ts).step(direction="foo")


class TestSeeking:
    @tests.cached_example
    def ts(self):
        ts = tsutil.all_trees_ts(4)
        assert ts.num_trees == 26
        return ts

    @pytest.mark.parametrize("index", range(26))
    def test_seek_forward_from_null(self, index):
        ts = self.ts()
        check_seek_forward_from_null(ts, index)

    @pytest.mark.parametrize("index", range(1, 26))
    def test_seek_forward_from_first(self, index):
        ts = self.ts()
        check_seek_forward_from_first(ts, index)

    @pytest.mark.parametrize("index", range(1, 26))
    def test_seek_last_from_index(self, index):
        ts = self.ts()
        tree1 = StatefulTree(ts)
        tree1.iter_forward(index)
        tree1.seek_forward(ts.num_trees - 1)
        tree2 = StatefulTree(ts)
        tree2.prev()
        tree1.assert_equal(tree2)

    @pytest.mark.parametrize("index", range(26))
    def test_seek_backward_from_null(self, index):
        ts = self.ts()
        check_seek_backward_from_null(ts, index)

    @pytest.mark.parametrize("index", range(0, 25))
    def test_seek_backward_from_last(self, index):
        ts = self.ts()
        check_seek_backward_from_last(ts, index)

    @pytest.mark.parametrize("index", range(0, 25))
    def test_seek_first_from_index(self, index):
        ts = self.ts()
        tree1 = StatefulTree(ts)
        tree1.iter_backward(index)
        tree1.seek_backward(0)
        tree2 = StatefulTree(ts)
        tree2.next()
        tree1.assert_equal(tree2)

    @pytest.mark.parametrize("index", range(26))
    def test_seek_forward_out_range_is_empty(self, index):
        ts = self.ts()
        check_seek_forward_out_range_is_empty(ts, index)

    @pytest.mark.parametrize("index", range(26))
    def test_seek_backward_out_range_is_empty(self, index):
        ts = self.ts()
        check_seek_backward_out_range_is_empty(ts, index)


class TestAllTreesTs:
    @pytest.mark.parametrize("n", [2, 3, 4])
    def test_forward_full(self, n):
        ts = tsutil.all_trees_ts(n)
        check_iters_forward(ts)

    @pytest.mark.parametrize("n", [2, 3, 4])
    def test_back_full(self, n):
        ts = tsutil.all_trees_ts(n)
        check_iters_back(ts)

    @pytest.mark.parametrize("n", [2, 3, 4])
    def test_forward_back(self, n):
        ts = tsutil.all_trees_ts(n)
        check_forward_back_sweep(ts)


class TestManyTreesSimulationExample:
    @tests.cached_example
    def ts(self):
        ts = msprime.sim_ancestry(
            10, sequence_length=1000, recombination_rate=0.1, random_seed=1234
        )
        assert ts.num_trees > 250
        return ts

    @pytest.mark.parametrize("index", [1, 5, 10, 50, 100])
    def test_seek_forward_from_null(self, index):
        ts = self.ts()
        check_seek_forward_from_null(ts, index)

    @pytest.mark.parametrize("num_trees", [1, 5, 10, 50, 100])
    def test_seek_forward_from_mid(self, num_trees):
        ts = self.ts()
        start_index = ts.num_trees // 2
        dest_index = min(start_index + num_trees, ts.num_trees - 1)
        tree1 = StatefulTree(ts)
        tree1.iter_forward(start_index)
        tree1.seek_forward(dest_index)
        tree2 = StatefulTree(ts)
        tree2.iter_forward(dest_index)
        tree1.assert_equal(tree2)

    @pytest.mark.parametrize("index", [1, 5, 10, 50, 100])
    def test_seek_backward_from_null(self, index):
        ts = self.ts()
        check_seek_backward_from_null(ts, index)

    @pytest.mark.parametrize("num_trees", [1, 5, 10, 50, 100])
    def test_seek_backward_from_mid(self, num_trees):
        ts = self.ts()
        start_index = ts.num_trees // 2
        dest_index = max(start_index - num_trees, 0)
        tree1 = StatefulTree(ts)
        tree1.iter_backward(start_index)
        tree1.seek_backward(dest_index)
        tree2 = StatefulTree(ts)
        tree2.iter_backward(dest_index)

    @pytest.mark.parametrize("index", [1, 5, 10, 50, 100])
    def test_seek_forward_out_range_is_empty(self, index):
        ts = self.ts()
        check_seek_forward_out_range_is_empty(ts, index)

    @pytest.mark.parametrize("index", [1, 5, 10, 50, 100])
    def test_seek_backward_out_range_is_empty(self, index):
        ts = self.ts()
        check_seek_backward_out_range_is_empty(ts, index)

    def test_forward_full(self):
        check_iters_forward(self.ts())

    def test_back_full(self):
        check_iters_back(self.ts())


class TestSuiteExamples:
    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_forward_full(self, ts):
        check_iters_forward(ts)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_back_full(self, ts):
        check_iters_back(ts)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_seek_forward_from_null(self, ts):
        index = ts.num_trees // 2
        check_seek_forward_from_null(ts, index)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_seek_forward_from_first(self, ts):
        index = ts.num_trees - 1
        check_seek_forward_from_first(ts, index)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_seek_backward_from_null(self, ts):
        index = ts.num_trees // 2
        check_seek_backward_from_null(ts, index)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_seek_backward_from_last(self, ts):
        index = 0
        check_seek_backward_from_last(ts, index)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_seek_forward_out_range_is_empty(self, ts):
        index = ts.num_trees // 2
        check_seek_forward_out_range_is_empty(ts, index)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_seek_backward_out_range_is_empty(self, ts):
        index = ts.num_trees // 2
        check_seek_backward_out_range_is_empty(ts, index)


--- ../../tskit/python/tests/test_ms.py ---


"""
Test cases for ms output in tskit. All of these tests have separate versions
for the cases of single replicate and multiple replicates. This is because
msprime.simulate generates a tree_sequence object if the num_replicates argument
is not used but an iterator over tree_sequences if the num_replicates argument
is used.
"""
import collections
import itertools
import os
import tempfile
import unittest

import msprime

import tskit as ts

length = 1e2
mutation_rate = 1e-2
num_replicates = 3


def get_ms_file_quantity(ms_file, quantity):
    quantities = {}
    num_replicates = 0
    num_sites = []
    num_positions = []
    num_haplotypes = []
    genotypes = []
    positions = []
    gens = []
    for line in ms_file:
        if len(line.split()) > 0:
            if line.split()[0] == "//":
                num_replicates = num_replicates + 1
                num_haplotypes.append(0)
                if len(gens) > 0:
                    genotypes.append(gens)
                    gens = []
            if line.split()[0] == "segsites:":
                num_sites.append(int(line.split()[1]))
            if line.split()[0] == "positions:":
                num_positions.append(len(line.split()) - 1)
                positions.append(line[11:].rstrip())
            if (
                line[0:2] == "00"
                or line[0:2] == "01"
                or line[0:2] == "10"
                or line[0:2] == "11"
            ):
                num_haplotypes[-1] = num_haplotypes[-1] + 1
                gens.append(line.rstrip())
    genotypes.append(gens)
    quantities["num_replicates"] = num_replicates
    quantities["num_sites"] = num_sites
    quantities["num_positions"] = num_positions
    quantities["num_haplotypes"] = num_haplotypes
    quantities["genotypes"] = genotypes
    quantities["positions"] = positions

    return quantities[quantity]


class TestNumReplicates(unittest.TestCase):
    """
    Tests that the number of replicates written out is the same as
    the number of replicates simulated
    """

    def verify_num_replicates(self, tree_seq, num_replicates):
        if isinstance(tree_seq, collections.abc.Iterable):
            with tempfile.TemporaryDirectory() as temp_dir:
                ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
                with open(ms_file_path, "w") as f:
                    ts.write_ms(
                        tree_seq,
                        f,
                        num_replicates=num_replicates,
                    )
                with open(ms_file_path) as handle:
                    num_replicates_file = get_ms_file_quantity(handle, "num_replicates")
        else:
            with tempfile.TemporaryDirectory() as temp_dir:
                ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
                with open(ms_file_path, "w") as f:
                    ts.write_ms(tree_seq, f)
                with open(ms_file_path) as handle:
                    num_replicates_file = get_ms_file_quantity(handle, "num_replicates")
        self.assertEqual(num_replicates, num_replicates_file)

    def test_num_replicates(self):
        tree_seq = msprime.simulate(
            25, length=length, mutation_rate=mutation_rate, random_seed=123
        )
        self.verify_num_replicates(tree_seq, 1)

    def test_num_replicates_multiple(self):
        tree_seq = msprime.simulate(
            25,
            length=length,
            mutation_rate=mutation_rate,
            random_seed=123,
            num_replicates=num_replicates,
        )
        self.verify_num_replicates(tree_seq, num_replicates)


class TestNumHaplotypes(unittest.TestCase):
    """
    Tests that the number of haplotypes output are the same as the
    number of individuals simulated.
    """

    def verify_num_haplotypes(self, tree_seq, tree_seq2, num_replicates):
        if isinstance(tree_seq, collections.abc.Iterable):
            with tempfile.TemporaryDirectory() as temp_dir:
                ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
                with open(ms_file_path, "w") as f:
                    ts.write_ms(
                        tree_seq,
                        f,
                        num_replicates=num_replicates,
                    )
                with open(ms_file_path) as handle:
                    num_haplotypes = get_ms_file_quantity(handle, "num_haplotypes")
            j = 0
            for ts_indv in tree_seq2:
                self.assertEqual(ts_indv.num_samples, num_haplotypes[j])
                j = j + 1
        else:
            with tempfile.TemporaryDirectory() as temp_dir:
                ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
                with open(ms_file_path, "w") as f:
                    ts.write_ms(tree_seq, f)
                with open(ms_file_path) as handle:
                    num_haplotypes = get_ms_file_quantity(handle, "num_haplotypes")
            self.assertEqual(tree_seq.num_samples, num_haplotypes[0])

    def test_num_haplotypes(self):
        tree_seq = msprime.simulate(
            25, length=length, mutation_rate=mutation_rate, random_seed=123
        )
        self.verify_num_haplotypes(tree_seq, tree_seq, 1)

    def test_num_haplotypes_replicates(self):
        tree_seq = msprime.simulate(
            25,
            length=length,
            mutation_rate=mutation_rate,
            random_seed=123,
            num_replicates=num_replicates,
        )
        tree_seq, tree_seq2 = itertools.tee(tree_seq)
        self.verify_num_haplotypes(tree_seq, tree_seq2, num_replicates)


class TestNumSites(unittest.TestCase):
    """
    Tests that the number of sites written out as well as the length
    of the positions list match the number of variants in the tree sequence
    """

    def verify_num_sites(self, tree_seq, tree_seq2, num_replicates):
        if isinstance(tree_seq, collections.abc.Iterable):
            with tempfile.TemporaryDirectory() as temp_dir:
                ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
                with open(ms_file_path, "w") as f:
                    ts.write_ms(
                        tree_seq,
                        f,
                        num_replicates=num_replicates,
                    )
                with open(ms_file_path) as handle:
                    num_sites = get_ms_file_quantity(handle, "num_sites")
                with open(ms_file_path) as handle:
                    num_positions = get_ms_file_quantity(handle, "num_positions")
            j = 0
            for ts_indv in tree_seq2:
                self.assertEqual(ts_indv.num_sites, num_sites[j])
                self.assertEqual(ts_indv.num_sites, num_positions[j])
                j = j + 1
        else:
            with tempfile.TemporaryDirectory() as temp_dir:
                ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
                with open(ms_file_path, "w") as f:
                    ts.write_ms(tree_seq, f)
                with open(ms_file_path) as handle:
                    num_sites = get_ms_file_quantity(handle, "num_sites")
                with open(ms_file_path) as handle:
                    num_positions = get_ms_file_quantity(handle, "num_positions")
            self.assertEqual(tree_seq.num_sites, num_sites[0])
            self.assertEqual(tree_seq.num_sites, num_positions[0])

    def test_num_sites(self):
        tree_seq = msprime.simulate(
            25, length=length, mutation_rate=mutation_rate, random_seed=123
        )
        self.verify_num_sites(tree_seq, tree_seq, 1)

    def test_num_sites_replicates(self):
        tree_seq = msprime.simulate(
            25,
            length=length,
            mutation_rate=mutation_rate,
            random_seed=123,
            num_replicates=num_replicates,
        )
        tree_seq, tree_seq2 = itertools.tee(tree_seq)
        self.verify_num_sites(tree_seq, tree_seq2, num_replicates)


class TestGenotypes(unittest.TestCase):
    """
    Tests that the haplotypes written out are the same as the haplotypes generated.
    """

    def get_genotypes(self, tree_seq):
        genotypes = tree_seq.genotype_matrix()
        gens_array = []
        for k in range(tree_seq.num_samples):
            tmp_str = "".join(map(str, genotypes[:, k]))
            gens_array.append(tmp_str)
        return gens_array

    def verify_genotypes(self, tree_seq, tree_seq2, num_replicates):
        if isinstance(tree_seq, collections.abc.Iterable):
            with tempfile.TemporaryDirectory() as temp_dir:
                ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
                with open(ms_file_path, "w") as f:
                    ts.write_ms(
                        tree_seq,
                        f,
                        num_replicates=num_replicates,
                    )
                with open(ms_file_path) as handle:
                    genotypes = get_ms_file_quantity(handle, "genotypes")
            j = 0
            for ts_indv in tree_seq2:
                self.assertEqual(self.get_genotypes(ts_indv), genotypes[j])
                j = j + 1
        else:
            with tempfile.TemporaryDirectory() as temp_dir:
                ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
                with open(ms_file_path, "w") as f:
                    ts.write_ms(tree_seq, f)
                with open(ms_file_path) as handle:
                    genotypes = get_ms_file_quantity(handle, "genotypes")
                self.assertEqual(self.get_genotypes(tree_seq), genotypes[0])

    def test_genotypes(self):
        tree_seq = msprime.simulate(
            25, length=length, mutation_rate=mutation_rate, random_seed=123
        )
        self.verify_genotypes(tree_seq, tree_seq, 1)

    def test_genotypes_replicates(self):
        tree_seq = msprime.simulate(
            25,
            length=length,
            mutation_rate=mutation_rate,
            random_seed=123,
            num_replicates=num_replicates,
        )
        tree_seq, tree_seq2 = itertools.tee(tree_seq)
        self.verify_genotypes(tree_seq, tree_seq2, num_replicates)


class TestPositions(unittest.TestCase):
    """
    Tests that the positions for the mutations written out are the same as the
    positions generated.
    """

    def get_positions(self, tree_seq):
        positions = []
        for i in range(tree_seq.num_sites):
            positions.append(
                f"{tree_seq.site(i).position / tree_seq.sequence_length:.4f}"
            )
        positions = " ".join(positions)
        return positions

    def verify_positions(self, tree_seq, tree_seq2, num_replicates):
        if isinstance(tree_seq, collections.abc.Iterable):
            with tempfile.TemporaryDirectory() as temp_dir:
                ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
                with open(ms_file_path, "w") as f:
                    ts.write_ms(
                        tree_seq,
                        f,
                        num_replicates=num_replicates,
                    )
                with open(ms_file_path) as handle:
                    positions = get_ms_file_quantity(handle, "positions")
            j = 0
            for ts_indv in tree_seq2:
                self.assertEqual(self.get_positions(ts_indv), positions[j])
                j = j + 1
        else:
            with tempfile.TemporaryDirectory() as temp_dir:
                ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
                with open(ms_file_path, "w") as f:
                    ts.write_ms(tree_seq, f)
                with open(ms_file_path) as handle:
                    positions = get_ms_file_quantity(handle, "positions")
            self.assertEqual(self.get_positions(tree_seq), positions[0])

    def test_positions(self):
        tree_seq = msprime.simulate(
            25, length=length, mutation_rate=mutation_rate, random_seed=123
        )
        self.verify_positions(tree_seq, tree_seq, 1)

    def test_positions_replicates(self):
        tree_seq = msprime.simulate(
            25,
            length=length,
            mutation_rate=mutation_rate,
            random_seed=123,
            num_replicates=num_replicates,
        )
        tree_seq, tree_seq2 = itertools.tee(tree_seq)
        self.verify_positions(tree_seq, tree_seq2, num_replicates)


--- ../../tskit/python/tests/test_extend_haplotypes.py ---

import msprime
import numpy as np
import pytest

import _tskit
import tests.test_wright_fisher as wf
import tskit
from tests import tsutil
from tests.test_highlevel import get_example_tree_sequences

# ↑ See https://github.com/tskit-dev/tskit/issues/1804 for when
# we can remove this.


def _slide_mutation_nodes_up(ts, mutations):
    # adjusts mutations' nodes to place each mutation on the correct edge given
    # their time; requires mutation times be nonmissing and the mutation times
    # be >= their nodes' times.

    assert np.all(~tskit.is_unknown_time(mutations.time)), "times must be known"
    new_nodes = mutations.node.copy()

    mut = 0
    for tree in ts.trees():
        _, right = tree.interval
        while (
            mut < mutations.num_rows and ts.sites_position[mutations.site[mut]] < right
        ):
            t = mutations.time[mut]
            c = mutations.node[mut]
            p = tree.parent(c)
            assert ts.nodes_time[c] <= t
            while p != -1 and ts.nodes_time[p] <= t:
                c = p
                p = tree.parent(c)
            assert ts.nodes_time[c] <= t
            if p != -1:
                assert t < ts.nodes_time[p]
            new_nodes[mut] = c
            mut += 1

    # in C the node column can be edited in place
    new_mutations = mutations.copy()
    new_mutations.clear()
    for mut, n in zip(mutations, new_nodes):
        new_mutations.append(mut.replace(node=n))

    return new_mutations


def print_edge_list(head, edges, left, right):
    print("Edge list:")
    for j, (e, x) in enumerate(head):
        print(
            f"  {j}: {e} ({x}); "
            + (
                f"{edges.child[e]}->{edges.parent[e]} on [{left[e]}, {right[e]})"
                if e >= 0
                else "(null)"
            )
        )
    print(f"length = {len(head)}")


class HaplotypeExtender:
    def __init__(self, ts, forwards):
        """
        Below we will iterate through the trees, either to the left or the right,
        keeping the following state consistent:
        - we are moving from a previous tree, last_tree, to new one, next_tree
        - here: the position that separates the last_tree from the next_tree
        - (here, there): the segment covered by next_tree
        - edges_out: edges to be removed from last_tree to get next_tree
        - parent_out: the forest induced by edges_out, a subset of last_tree
        - edges_in: edges to be added to last_tree to get next_tree
        - parent_in: the forest induced by edges_in, a subset of next_tree
        - next_degree: the degree of each node in next_tree
        - next_nodes_edge: for each node, the edge above it in next_tree
        - last_degree: the degree of each node in last_tree
        - last_nodes_edge: for each node, the edge above it in last_tree
        Except: each of edges_in and edges_out is of the form e, x, and the
        label x>0 if the edge is postponed to the next segment.
        The label is x=1 for postponed edges, and x=2 for new edges.
        In other words:
        - elements e, x of edges_out with x=0 are in last_tree but not next_tree
        - elements e, x of edges_in with x=0 are in next_tree but not last_tree
        - elements e, x of edges_out with x=1 are in both trees,
            and hence don't count for parent_out
        - elements e, x of edges_in with x=1 are in neither,
            and hence don't count for parent_in
        - elements e, x for edges_out with x=2 have just been added, and so ought
            to count towards the next tree, but we have to put them in edges out
            because they'll be removed next time.
        Notes:
        - things having to do with last_tree do not change,
          but things having to do with next_tree might change as we go along
        - parent_out and parent_in do not refer to the *entire* last/next_tree,
          but rather to *only* the edges_in/edges_out
        Edges in can have one of three things happen to them:
        1. they get added to the next tree, as usual;
        2. they get postponed to the tree after the next tree,
            and are thus part of edges_in again next time;
        3. they get postponed but run out of span so they dissapear entirely.
        Edges out are similarly of four varieties:
        0. they are also in case (3) of edges_in, i.e., their extent was modified
            when they were in edges_in so that they now have left=right;
        1. they get removed from the last tree, as usual;
        2. they get extended to the next tree,
            and are thus part of edges_out again next time;
        3. they are in fact a newly added edge, and so are part of edges_out next time.
        """
        self.ts = ts
        self.edges = ts.tables.edges.copy()
        self.new_left = ts.edges_left.copy()
        self.new_right = ts.edges_right.copy()
        self.last_degree = np.full(ts.num_nodes, 0, dtype="int")
        self.next_degree = np.full(ts.num_nodes, 0, dtype="int")
        self.parent_out = np.full(ts.num_nodes, -1, dtype="int")
        self.parent_in = np.full(ts.num_nodes, -1, dtype="int")
        self.not_sample = [not n.is_sample() for n in ts.nodes()]
        self.next_nodes_edge = np.full(ts.num_nodes, -1, dtype="int")
        self.last_nodes_edge = np.full(ts.num_nodes, -1, dtype="int")

        if forwards:
            self.direction = 1
            # in C we can just modify these in place, but in
            # python they are (silently) immutable
            self.near_side = list(self.new_left)
            self.far_side = list(self.new_right)
        else:
            self.direction = -1
            self.near_side = list(self.new_right)
            self.far_side = list(self.new_left)

        self.edges_out = []
        self.edges_in = []

    def print_state(self):
        print("~~~~~~~~~~~~~~~~~~~~~~~~")
        print("edges in:", self.edges_in)
        print("parent out:")
        for j, pj in enumerate(self.parent_out):
            print(f"   {j}: {pj}")
        print("parent in:")
        for j, pj in enumerate(self.parent_in):
            print(f"   {j}: {pj}")
        print("edges out:", self.edges_out)
        print("parent out:", self.parent_out)
        print("last nodes edge:")
        for j, ej in enumerate(self.last_nodes_edge):
            print(
                f"   {j}: {ej}, "
                + (
                    "(null)"
                    if ej == -1
                    else (
                        f"({self.edges.child[ej]}->{self.edges.parent[ej]}, "
                        "{self.near_side[ej]}-{self.far_side[ej]}"
                    )
                )
            )
        for e, _ in self.edges_out:
            print(
                "edge out:   ",
                "e =",
                e,
                "c =",
                self.edges.child[e],
                "p =",
                self.edges.parent[e],
                self.near_side[e],
                self.far_side[e],
            )

    def next_tree(self, tree_pos):
        # Clear out non-extended or postponed edges:
        # Note: maintaining parent_out is a bit tricky, because
        # if an edge from p->c has been extended, entirely replacing
        # another edge from p'->c, then both edges may be in edges_out,
        # and we only want to include the *first* one.

        for e, x in self.edges_out:
            self.parent_out[self.edges.child[e]] = tskit.NULL
            if x > 1:
                # this is needed to catch newly-created edges
                self.last_nodes_edge[self.edges.child[e]] = e
                self.last_degree[self.edges.child[e]] += 1
                self.last_degree[self.edges.parent[e]] += 1
            elif x == 0 and self.near_side[e] != self.far_side[e]:
                self.last_nodes_edge[self.edges.child[e]] = tskit.NULL
                self.last_degree[self.edges.child[e]] -= 1
                self.last_degree[self.edges.parent[e]] -= 1
        tmp = []
        for e, x in self.edges_out:
            if x > 0:
                tmp.append([e, 0])
        self.edges_out = tmp
        for e, x in self.edges_in:
            self.parent_in[self.edges.child[e]] = tskit.NULL
            if x == 0 and self.near_side[e] != self.far_side[e]:
                assert self.last_nodes_edge[self.edges.child[e]] == tskit.NULL
                self.last_nodes_edge[self.edges.child[e]] = e
                self.last_degree[self.edges.child[e]] += 1
                self.last_degree[self.edges.parent[e]] += 1
        tmp = []
        for e, x in self.edges_in:
            if x > 0:
                tmp.append([e, 0])
        self.edges_in = tmp

        # done cleanup from last tree transition;
        # now we update the state to reflect the current tree transition
        for j in range(
            tree_pos.out_range.start, tree_pos.out_range.stop, self.direction
        ):
            e = tree_pos.out_range.order[j]
            if (self.parent_out[self.edges.child[e]] == tskit.NULL) and (
                self.near_side[e] != self.far_side[e]
            ):
                self.edges_out.append([e, 0])

        for e, _ in self.edges_out:
            self.parent_out[self.edges.child[e]] = self.edges.parent[e]
            self.next_nodes_edge[self.edges.child[e]] = tskit.NULL
            self.next_degree[self.edges.child[e]] -= 1
            self.next_degree[self.edges.parent[e]] -= 1

        for j in range(tree_pos.in_range.start, tree_pos.in_range.stop, self.direction):
            e = tree_pos.in_range.order[j]
            self.edges_in.append([e, 0])

        for e, _ in self.edges_in:
            self.parent_in[self.edges.child[e]] = self.edges.parent[e]
            assert self.next_nodes_edge[self.edges.child[e]] == tskit.NULL
            self.next_nodes_edge[self.edges.child[e]] = e
            self.next_degree[self.edges.child[e]] += 1
            self.next_degree[self.edges.parent[e]] += 1

    def check_state_at(self, pos, before, degree, nodes_edge):
        # if before=True then we construct the state at epsilon-on-near-side-of `pos`,
        # otherwise, at epsilon-on-far-side-of `pos`.
        check_degree = np.zeros(self.ts.num_nodes, dtype="int")
        check_nodes_edge = np.full(self.ts.num_nodes, -1, dtype="int")
        assert len(self.near_side) == self.edges.num_rows
        assert len(self.far_side) == self.edges.num_rows
        for j, (e, l, r) in enumerate(zip(self.edges, self.near_side, self.far_side)):
            overlaps = (l != r) and (
                ((pos - l) * (r - pos) > 0)
                or (r == pos and before)
                or (l == pos and not before)
            )
            if overlaps:
                check_degree[e.child] += 1
                check_degree[e.parent] += 1
                assert check_nodes_edge[e.child] == tskit.NULL
                check_nodes_edge[e.child] = j
        np.testing.assert_equal(check_nodes_edge, nodes_edge)
        np.testing.assert_equal(check_degree, degree)

    def check_parent(self, parent, edge_ids):
        temp_parent = np.full(self.ts.num_nodes, -1, dtype="int")
        for j in edge_ids:
            c = self.edges.child[j]
            p = self.edges.parent[j]
            temp_parent[c] = p
        np.testing.assert_equal(temp_parent, parent)

    def check_state(self, here):
        for e, x in self.edges_in:
            assert x == 0
            assert self.near_side[e] != self.far_side[e]
        for e, x in self.edges_out:
            assert x == 0
            assert self.near_side[e] != self.far_side[e]
        self.check_state_at(here, False, self.next_degree, self.next_nodes_edge)
        self.check_state_at(here, True, self.last_degree, self.last_nodes_edge)
        self.check_parent(self.parent_in, [j for j, x in self.edges_in if x == 0])
        self.check_parent(self.parent_out, [j for j, x in self.edges_out if x == 0])

    def add_or_extend_edge(self, new_parent, child, left, right):
        there = right if (self.direction == 1) else left
        old_edge = self.next_nodes_edge[child]
        if old_edge != tskit.NULL:
            old_parent = self.edges.parent[old_edge]
        else:
            old_parent = tskit.NULL
        if new_parent != old_parent:
            # if our new edge is in edges_out, it should be extended
            if self.parent_out[child] == new_parent:
                e_out = self.last_nodes_edge[child]
                assert e_out >= 0
                assert self.edges.child[e_out] == child
                assert self.edges.parent[e_out] == new_parent
                self.far_side[e_out] = there
                assert self.near_side[e_out] != self.far_side[e_out]
                for ex_out in self.edges_out:
                    if ex_out[0] == e_out:
                        break
                assert ex_out[0] == e_out
                ex_out[1] = 1
            else:
                e_out = self.add_edge(new_parent, child, left, right)
                self.edges_out.append([e_out, 2])
            # If we're replacing the edge above this node, it must be in edges_in;
            # note that this assertion excludes the case that we're interrupting
            # an existing edge.
            assert (self.next_nodes_edge[child] == tskit.NULL) or (
                self.next_nodes_edge[child] in [e for e, _ in self.edges_in]
            )
            self.next_nodes_edge[child] = e_out
            self.next_degree[child] += 1
            self.next_degree[new_parent] += 1
            self.parent_out[child] = tskit.NULL
            if old_edge != tskit.NULL:
                for ex_in in self.edges_in:
                    e_in = ex_in[0]
                    if e_in == old_edge and (ex_in[1] == 0):
                        self.near_side[e_in] = there
                        if self.far_side[e_in] != there:
                            ex_in[1] = 1
                        self.next_nodes_edge[child] = tskit.NULL
                        self.next_degree[child] -= 1
                        self.next_degree[self.parent_in[child]] -= 1
                        self.parent_in[child] = tskit.NULL

    def add_edge(self, parent, child, left, right):
        new_id = self.edges.add_row(parent=parent, child=child, left=left, right=right)
        # this appending should not be necessary in C
        if self.direction == 1:
            self.near_side.append(left)
            self.far_side.append(right)
        else:
            self.near_side.append(right)
            self.far_side.append(left)
        return new_id

    def mergeable(self, c):
        # returns a finite number of new edges needed
        # if the paths in parent_in and parent_out
        # up through nodes that aren't in the other tree
        # end at the same place and don't have conflicting times;
        # otherwise, returns Inf
        p_out = self.parent_out[c]
        p_in = self.parent_in[c]
        t_out = np.inf if p_out == tskit.NULL else self.ts.nodes_time[p_out]
        t_in = np.inf if p_in == tskit.NULL else self.ts.nodes_time[p_in]
        child = c
        num_new_edges = 0
        num_extended = 0
        while True:
            climb_in = (
                p_in != tskit.NULL
                and self.last_degree[p_in] == 0
                and self.not_sample[p_in]
                and t_in < t_out
            )
            climb_out = (
                p_out != tskit.NULL
                and self.next_degree[p_out] == 0
                and self.not_sample[p_out]
                and t_out < t_in
            )
            if climb_in:
                if self.parent_in[child] != p_in and self.parent_out[child] != p_in:
                    num_new_edges += 1
                child = p_in
                p_in = self.parent_in[p_in]
                t_in = np.inf if p_in == tskit.NULL else self.ts.nodes_time[p_in]
            elif climb_out:
                if self.parent_in[child] != p_out and self.parent_out[child] != p_out:
                    num_new_edges += 1
                child = p_out
                p_out = self.parent_out[p_out]
                t_out = np.inf if p_out == tskit.NULL else self.ts.nodes_time[p_out]
                num_extended += 1
            else:
                break
        if num_extended == 0 or p_in != p_out or p_in == tskit.NULL:
            num_new_edges = np.inf
        return num_new_edges

    def merge_paths(self, c, left, right):
        p_out = self.parent_out[c]
        p_in = self.parent_in[c]
        t_out = self.ts.nodes_time[p_out]
        t_in = self.ts.nodes_time[p_in]
        child = c
        while True:
            climb_in = (
                p_in != tskit.NULL
                and self.last_degree[p_in] == 0
                and self.not_sample[p_in]
                and t_in < t_out
            )
            climb_out = (
                p_out != tskit.NULL
                and self.next_degree[p_out] == 0
                and self.not_sample[p_out]
                and t_out < t_in
            )
            if climb_in:
                self.add_or_extend_edge(p_in, child, left, right)
                child = p_in
                p_in = self.parent_in[p_in]
                t_in = np.inf if p_in == tskit.NULL else self.ts.nodes_time[p_in]
            elif climb_out:
                self.add_or_extend_edge(p_out, child, left, right)
                child = p_out
                p_out = self.parent_out[p_out]
                t_out = np.inf if p_out == tskit.NULL else self.ts.nodes_time[p_out]
            else:
                break
        assert p_out == p_in
        self.add_or_extend_edge(p_out, child, left, right)

    def extend_haplotypes(self):
        tree_pos = tsutil.TreePosition(self.ts)
        if self.direction == 1:
            valid = tree_pos.next()
        else:
            valid = tree_pos.prev()
        while valid:
            left, right = tree_pos.interval
            # there = right if self.direction == 1 else left
            here = left if self.direction == 1 else right
            self.next_tree(tree_pos)
            self.check_state(here)
            max_new_edges = 0
            next_max_new_edges = np.inf
            while max_new_edges < np.inf:
                for e_in, x in self.edges_in:
                    if x == 0:
                        c = self.edges.child[e_in]
                        assert self.next_degree[c] > 0
                        if self.last_degree[c] > 0:
                            ne = self.mergeable(c)
                            if ne <= max_new_edges:
                                self.merge_paths(c, left, right)
                            else:
                                next_max_new_edges = min(ne, next_max_new_edges)
                max_new_edges = next_max_new_edges
                next_max_new_edges = np.inf
            # end of loop, next tree
            if self.direction == 1:
                valid = tree_pos.next()
            else:
                valid = tree_pos.prev()
        if self.direction == 1:
            self.new_left = np.array(self.near_side)
            self.new_right = np.array(self.far_side)
        else:
            self.new_right = np.array(self.near_side)
            self.new_left = np.array(self.far_side)
        # Get rid of adjacent, identical edges
        keep = np.full(self.edges.num_rows, True, dtype=bool)
        for j in range(self.edges.num_rows - 1):
            if (
                self.edges.parent[j] == self.edges.parent[j + 1]
                and self.edges.child[j] == self.edges.child[j + 1]
                and self.new_right[j] == self.new_left[j + 1]
            ):
                self.new_right[j] = self.new_right[j + 1]
                self.new_left[j + 1] = self.new_right[j + 1]
        for j in range(self.edges.num_rows):
            left = self.new_left[j]
            right = self.new_right[j]
            if left < right:
                self.edges[j] = self.edges[j].replace(left=left, right=right)
            else:
                keep[j] = False
        self.edges.keep_rows(keep)


def extend_haplotypes(ts, max_iter=10):
    tables = ts.dump_tables()
    mutations = tables.mutations.copy()
    tables.mutations.clear()

    last_num_edges = ts.num_edges
    for _ in range(max_iter):
        for forwards in [True, False]:
            extender = HaplotypeExtender(ts, forwards=forwards)
            extender.extend_haplotypes()
            tables.edges.replace_with(extender.edges)
            tables.sort()
            tables.build_index()
            ts = tables.tree_sequence()
        if ts.num_edges == last_num_edges:
            break
        else:
            last_num_edges = ts.num_edges

    tables = ts.dump_tables()
    mutations = _slide_mutation_nodes_up(ts, mutations)
    tables.mutations.replace_with(mutations)
    tables.sort()
    ts = tables.tree_sequence()
    return ts


def _path_pairs(tree):
    for c in tree.postorder():
        p = tree.parent(c)
        while p != tskit.NULL:
            yield (c, p)
            p = tree.parent(p)


def _path_up(c, p, tree, include_parent=False):
    # path from c up to p in tree, not including c or p
    c = tree.parent(c)
    while c != p and c != tskit.NULL:
        yield c
        c = tree.parent(c)
    assert c == p
    if include_parent:
        yield p


def _path_up_pairs(c, p, tree, others):
    # others should be a list of nodes
    otherdict = {tree.time(n): n for n in others}
    ot = min(otherdict)
    for n in _path_up(c, p, tree, include_parent=True):
        nt = tree.time(n)
        while ot < nt:
            on = otherdict.pop(ot)
            yield c, on
            c = on
            if len(otherdict) > 0:
                ot = min(otherdict)
            else:
                ot = np.inf
        yield c, n
        c = n
    assert n == p
    assert len(otherdict) == 0


def _path_overlaps(c, p, tree1, tree2):
    for n in _path_up(c, p, tree1):
        if n in tree2.nodes():
            return True
    return False


def _paths_mergeable(c, p, tree1, tree2):
    # checks that the nodes between c and p in each tree
    # are not present in the other tree
    # and their sets of times are disjoint
    nodes1 = set(tree1.nodes())
    nodes2 = set(tree2.nodes())
    assert c in nodes1, f"child node {c} not in tree1"
    assert p in nodes1, f"parent node {p} not in tree1"
    assert c in nodes2, f"child node {c} not in tree2"
    assert p in nodes2, f"parent node {p} not in tree2"
    path1 = set(_path_up(c, p, tree1))
    path2 = set(_path_up(c, p, tree2))
    times1 = {tree1.time(n) for n in path1}
    times2 = {tree2.time(n) for n in path2}
    return (
        (not _path_overlaps(c, p, tree1, tree2))
        and (not _path_overlaps(c, p, tree2, tree1))
        and len(times1.intersection(times2)) == 0
    )


def _extend_nodes(ts, interval, extendable):
    tables = ts.dump_tables()
    tables.edges.clear()
    mutations = tables.mutations.copy()
    tables.mutations.clear()
    left, right = interval
    # print("=================")
    # print("extending", left, right)
    extend_above = {}  # gives the new child->parent mapping
    todo_edges = np.repeat(True, ts.num_edges)
    tree = ts.at(left)
    for c, p, others in extendable:
        # print("c:", c, "p:", p, "others:", others)
        others_not_done_yet = set(others) - set(extend_above)
        if len(others_not_done_yet) > 0:
            for cn, pn in _path_up_pairs(c, p, tree, others_not_done_yet):
                if cn not in extend_above:
                    assert cn not in extend_above
                    extend_above[cn] = pn
    for c, p in extend_above.items():
        e = tree.edge(c)
        if e == tskit.NULL or ts.edge(e).parent != p:
            # print("adding", c, p)
            tables.edges.add_row(child=c, parent=p, left=left, right=right)
            if e != tskit.NULL:
                edge = ts.edge(e)
                # adjust endpoints on existing edge
                for el, er in [
                    (max(edge.left, right), edge.right),
                    (edge.left, min(edge.right, left)),
                ]:
                    if el < er:
                        # print("replacing", edge, el, er)
                        tables.edges.append(edge.replace(left=el, right=er))
                todo_edges[e] = False
    for todo, edge in zip(todo_edges, ts.edges()):
        if todo:
            # print("retaining", edge)
            tables.edges.append(edge)
    tables.sort()
    ts = tables.tree_sequence()
    mutations = _slide_mutation_nodes_up(ts, mutations)
    tables.mutations.replace_with(mutations)
    tables.sort()
    return tables.tree_sequence()


def _naive_pass(ts, direction):
    assert direction in (-1, +1)
    num_trees = ts.num_trees
    if direction == +1:
        indexes = range(0, num_trees - 1, 1)
    else:
        indexes = range(num_trees - 1, 0, -1)
    for tj in indexes:
        extendable = []
        this_tree = ts.at_index(tj)
        next_tree = ts.at_index(tj + direction)
        # print("-----------", this_tree.index)
        # print(this_tree.draw_text())
        # print(next_tree.draw_text())
        for c, p in _path_pairs(this_tree):
            if (
                p != this_tree.parent(c)
                and p in next_tree.nodes()
                and c in next_tree.nodes(p)
            ):
                # print(c, p, " and ", list(next_tree.nodes(p)))
                if _paths_mergeable(c, p, this_tree, next_tree):
                    extendable.append((c, p, list(_path_up(c, p, this_tree))))
        # print("extending to", extendable)
        ts = _extend_nodes(ts, next_tree.interval, extendable)
        assert num_trees == ts.num_trees
    return ts


def naive_extend_haplotypes(ts, max_iter=20):
    for _ in range(max_iter):
        ets = _naive_pass(ts, +1)
        ets = _naive_pass(ets, -1)
        if ets == ts:
            break
        ts = ets
    return ts


class TestExtendThings:
    """
    Common utilities in the two classes below.
    """

    def verify_simplify_equality(self, ts, ets):
        assert ts.num_nodes == ets.num_nodes
        assert ts.num_samples == ets.num_samples
        t = ts.simplify().tables
        et = ets.simplify().tables
        et.assert_equals(t, ignore_provenance=True)
        assert np.all(ts.genotype_matrix() == ets.genotype_matrix())

    def naive_verify(self, ts):
        ets = naive_extend_haplotypes(ts)
        self.verify_simplify_equality(ts, ets)


class TestExtendHaplotypes(TestExtendThings):
    """
    Test the 'extend_haplotypes' method.
    """

    def get_example1(self):
        # 15.00|         |   13    |         |
        #      |         |    |    |         |
        # 12.00|   10    |   10    |    10   |
        #      |  +-+-+  |  +-+-+  |   +-+-+ |
        # 10.00|  8   |  |  |   |  |   8   | |
        #      |  |   |  |  |   |  |  ++-+ | |
        # 8.00 |  |   |  | 11  12  |  |  | | |
        #      |  |   |  |  |   |  |  |  | | |
        # 6.00 |  |   |  |  7   |  |  |  | | |
        #      |  |   |  |  |   |  |  |  | | |
        # 4.00 |  6   9  |  |   |  |  |  | | |
        #      |  |   |  |  |   |  |  |  | | |
        # 1.00 |  4   5  |  4   5  |  4  | 5 |
        #      | +++ +++ | +++ +++ | +++ | | |
        # 0.00 | 0 1 2 3 | 0 1 2 3 | 0 1 2 3 |
        #      0         3         6         9
        node_times = {
            0: 0,
            1: 0,
            2: 0,
            3: 0,
            4: 1,
            5: 1,
            6: 4,
            7: 6,
            8: 10,
            9: 4,
            10: 12,
            11: 8,
            12: 8,
            13: 15,
        }
        # (p,c,l,r)
        edges = [
            (4, 0, 0, 9),
            (4, 1, 0, 9),
            (5, 2, 0, 6),
            (5, 3, 0, 9),
            (6, 4, 0, 3),
            (9, 5, 0, 3),
            (7, 4, 3, 6),
            (11, 7, 3, 6),
            (12, 5, 3, 6),
            (8, 2, 6, 9),
            (8, 4, 6, 9),
            (8, 6, 0, 3),
            (10, 5, 6, 9),
            (10, 8, 0, 3),
            (10, 8, 6, 9),
            (10, 9, 0, 3),
            (10, 11, 3, 6),
            (10, 12, 3, 6),
            (13, 10, 3, 6),
        ]
        extended_edges = [
            (4, 0, 0.0, 9.0),
            (4, 1, 0.0, 9.0),
            (5, 2, 0.0, 6.0),
            (5, 3, 0.0, 9.0),
            (6, 4, 0.0, 9.0),
            (9, 5, 0.0, 9.0),
            (7, 6, 0.0, 9.0),
            (11, 7, 0.0, 9.0),
            (12, 9, 0.0, 9.0),
            (8, 2, 6.0, 9.0),
            (8, 11, 0.0, 9.0),
            (10, 8, 0.0, 9.0),
            (10, 12, 0.0, 9.0),
            (13, 10, 3.0, 6.0),
        ]
        samples = list(np.arange(4))
        tables = tskit.TableCollection(sequence_length=9)
        for (
            n,
            t,
        ) in node_times.items():
            flags = tskit.NODE_IS_SAMPLE if n in samples else 0
            tables.nodes.add_row(time=t, flags=flags)
        for p, c, l, r in edges:
            tables.edges.add_row(parent=p, child=c, left=l, right=r)
        ts = tables.tree_sequence()
        tables.edges.clear()
        for p, c, l, r in extended_edges:
            tables.edges.add_row(parent=p, child=c, left=l, right=r)
        ets = tables.tree_sequence()
        assert ts.num_edges == 19
        assert ets.num_edges == 14
        return ts, ets

    def get_example2(self):
        # 12.00|                     |          21         |                     |
        #      |                     |      +----+-----+   |                     |
        # 11.00|            20       |      |          |   |            20       |
        #      |        +----+---+   |      |          |   |        +----+---+   |
        # 10.00|        |       19   |      |         19   |        |       19   |
        #      |        |       ++-+ |      |        +-+-+ |        |       ++-+ |
        # 9.00 |       18       |  | |     18        |   | |       18       |  | |
        #      |     +--+--+    |  | |   +--+--+     |   | |     +--+--+    |  | |
        # 8.00 |     |     |    |  | |   |     |     |   | |    17     |    |  | |
        #      |     |     |    |  | |   |     |     |   | |   +-+-+   |    |  | |
        # 7.00 |     |     |   16  | |   |     |    16   | |   |   |   |    |  | |
        #      |     |     |   +++ | |   |     |   +-++  | |   |   |   |    |  | |
        # 6.00 |    15     |   | | | |   |     |   |  |  | |   |   |   |    |  | |
        #      |   +-+-+   |   | | | |   |     |   |  |  | |   |   |   |    |  | |
        # 5.00 |   |   |  14   | | | |   |    14   |  |  | |   |   |  14    |  | |
        #      |   |   |  ++-+ | | | |   |    ++-+ |  |  | |   |   |  ++-+  |  | |
        # 4.00 |  13   |  |  | | | | |  13    |  | |  |  | |  13   |  |  |  |  | |
        #      |  ++-+ |  |  | | | | |  ++-+  |  | |  |  | |  ++-+ |  |  |  |  | |
        # 3.00 |  |  | |  |  | | | | |  |  |  |  | | 12  | |  |  | |  |  | 12  | |
        #      |  |  | |  |  | | | | |  |  |  |  | | +++ | |  |  | |  |  | +++ | |
        # 2.00 | 11  | |  |  | | | | | 11  |  |  | | | | | | 11  | |  |  | | | | |
        #      | +++ | |  |  | | | | | +++ |  |  | | | | | | +++ | |  |  | | | | |
        # 1.00 | | | | | 10  | | | | | | | | 10  | | | | | | | | | | 10  | | | | |
        #      | | | | | +++ | | | | | | | | +++ | | | | | | | | | | +++ | | | | |
        # 0.00 | 0 7 4 9 2 5 6 1 3 8 | 0 7 4 2 5 6 1 3 9 8 | 0 7 4 1 2 5 6 3 9 8 |
        #      0                     3                     6                     9
        node_times = {
            0: 0,
            1: 0,
            2: 0,
            3: 0,
            4: 0,
            5: 0,
            6: 0,
            7: 0,
            8: 0,
            9: 0,
            10: 1,
            11: 2,
            12: 3,
            13: 4,
            14: 5,
            15: 6,
            16: 7,
            17: 8,
            18: 9,
            19: 10,
            20: 11,
            21: 12,
        }
        # (p,c,l,r)
        edges = [
            (10, 2, 0, 9),
            (10, 5, 0, 9),
            (11, 0, 0, 9),
            (11, 7, 0, 9),
            (12, 3, 3, 9),
            (12, 9, 3, 9),
            (13, 4, 0, 9),
            (13, 11, 0, 9),
            (14, 6, 0, 9),
            (14, 10, 0, 9),
            (15, 9, 0, 3),
            (15, 13, 0, 3),
            (16, 1, 0, 6),
            (16, 3, 0, 3),
            (16, 12, 3, 6),
            (17, 1, 6, 9),
            (17, 13, 6, 9),
            (18, 13, 3, 6),
            (18, 14, 0, 9),
            (18, 15, 0, 3),
            (18, 17, 6, 9),
            (19, 8, 0, 9),
            (19, 12, 6, 9),
            (19, 16, 0, 6),
            (20, 18, 0, 3),
            (20, 18, 6, 9),
            (20, 19, 0, 3),
            (20, 19, 6, 9),
            (21, 18, 3, 6),
            (21, 19, 3, 6),
        ]
        extended_edges = [
            (10, 2, 0.0, 9.0),
            (10, 5, 0.0, 9.0),
            (11, 0, 0.0, 9.0),
            (11, 7, 0.0, 9.0),
            (12, 3, 0.0, 9.0),
            (12, 9, 3.0, 9.0),
            (13, 4, 0.0, 9.0),
            (13, 11, 0.0, 9.0),
            (14, 6, 0.0, 9.0),
            (14, 10, 0.0, 9.0),
            (15, 9, 0.0, 3.0),
            (15, 13, 0.0, 9.0),
            (16, 1, 0.0, 6.0),
            (16, 12, 0.0, 9.0),
            (17, 1, 6.0, 9.0),
            (17, 15, 0.0, 9.0),
            (18, 14, 0.0, 9.0),
            (18, 17, 0.0, 9.0),
            (19, 8, 0.0, 9.0),
            (19, 16, 0.0, 9.0),
            (20, 18, 0.0, 3.0),
            (20, 18, 6.0, 9.0),
            (20, 19, 0.0, 3.0),
            (20, 19, 6.0, 9.0),
            (21, 18, 3.0, 6.0),
            (21, 19, 3.0, 6.0),
        ]
        samples = list(np.arange(10))
        tables = tskit.TableCollection(sequence_length=9)
        for (
            n,
            t,
        ) in node_times.items():
            flags = tskit.NODE_IS_SAMPLE if n in samples else 0
            tables.nodes.add_row(time=t, flags=flags)
        for p, c, l, r in edges:
            tables.edges.add_row(parent=p, child=c, left=l, right=r)
        ts = tables.tree_sequence()
        tables.edges.clear()
        for p, c, l, r in extended_edges:
            tables.edges.add_row(parent=p, child=c, left=l, right=r)
        ets = tables.tree_sequence()
        assert ts.num_edges == 30
        assert ets.num_edges == 26
        return ts, ets

    def get_example3(self):
        # Here is the full tree; extend edges should be able to
        # recover all unary nodes after simplification:
        #
        #       9         9         9          9
        #     +-+-+    +--+--+  +---+---+  +-+-+--+
        #     8   |    8     |  8   |   |  8 | |  |
        #     |   |  +-+-+   |  |   |   |  | | |  |
        #     7   |  |   7   |  |   7   |  | | |  7
        #   +-+-+ |  | +-++  |  | +-++  |  | | |  |
        #   6   | |  | |  6  |  | |  6  |  | | |  6
        # +-++  | |  | |  |  |  | |  |  |  | | |  |
        # 1  0  2 3  1 2  0  3  1 2  0  3  1 2 3  0
        #   +++          +++        +++          +++
        #   4 5          4 5        4 5          4 5
        #
        samples = [0, 1, 2, 3, 4, 5]
        node_times = [1, 1, 1, 1, 0, 0, 2, 3, 4, 5]
        # (p, c, l, r)
        edges = [
            (0, 4, 0, 10),
            (0, 5, 0, 10),
            (6, 0, 0, 10),
            (6, 1, 0, 3),
            (7, 2, 0, 7),
            (7, 6, 0, 10),
            (8, 1, 3, 10),
            (8, 7, 0, 5),
            (9, 2, 7, 10),
            (9, 3, 0, 10),
            (9, 7, 5, 10),
            (9, 8, 0, 10),
        ]
        tables = tskit.TableCollection(sequence_length=10)
        for n, t in enumerate(node_times):
            flags = tskit.NODE_IS_SAMPLE if n in samples else 0
            tables.nodes.add_row(time=t, flags=flags)
        for p, c, l, r in edges:
            tables.edges.add_row(parent=p, child=c, left=l, right=r)
        ets = tables.tree_sequence()
        ts = ets.simplify()
        assert ts.num_edges == 16
        assert ets.num_edges == 12
        return ts, ets

    def get_example4(self):
        # 7 and 8 should be extended to the whole sequence;
        # and also 5 to the second tree
        #
        #    6          6      6         6
        #  +-+-+      +-+-+  +-+-+     +-+-+
        #  |   |      7   |  |   8     |   |
        #  |   |     ++-+ |  | +-++    |   |
        #  4   5     4  | |  4 |  5    4   5
        # +++ +++   +++ | |  | | +++  +++ +++
        # 0 1 2 3   0 1 2 3  0 1 2 3  0 1 2 3
        node_times = (0, 0, 0, 0, 1, 1, 3, 2, 2)
        samples = (0, 1, 2, 3)
        # (p, c, l, r)
        extended_edges = [
            (4, 0, 0, 10),
            (4, 1, 0, 5),
            (4, 1, 7, 10),
            (5, 2, 0, 2),
            (5, 2, 5, 10),
            (5, 3, 0, 10),
            (7, 2, 2, 5),
            (7, 4, 0, 10),
            (8, 1, 5, 7),
            (8, 5, 0, 10),
            (6, 7, 0, 10),
            (6, 8, 0, 10),
        ]
        edges = [
            (4, 0, 0, 10),
            (4, 1, 0, 5),
            (4, 1, 7, 10),
            (5, 2, 0, 2),
            (5, 2, 5, 10),
            (5, 3, 0, 2),
            (5, 3, 5, 10),
            (7, 2, 2, 5),
            (7, 4, 2, 5),
            (8, 1, 5, 7),
            (8, 5, 5, 7),
            (6, 3, 2, 5),
            (6, 4, 0, 2),
            (6, 4, 5, 10),
            (6, 5, 0, 2),
            (6, 5, 7, 10),
            (6, 7, 2, 5),
            (6, 8, 5, 7),
        ]
        tables = tskit.TableCollection(sequence_length=10)
        tables.sort()
        for n, t in enumerate(node_times):
            flags = tskit.NODE_IS_SAMPLE if n in samples else 0
            tables.nodes.add_row(time=t, flags=flags)
        for p, c, l, r in edges:
            tables.edges.add_row(parent=p, child=c, left=l, right=r)
        ts = tables.tree_sequence()
        tables.edges.clear()
        for p, c, l, r in extended_edges:
            tables.edges.add_row(parent=p, child=c, left=l, right=r)
        ets = tables.tree_sequence()
        assert ts.num_edges == 18
        assert ets.num_edges == 12
        return ts, ets

    def get_example5(self):
        # This is an example where new edges are added
        # on both forwards and back passes
        # 4.00┊   ┊ 4 ┊ 4 ┊
        #     ┊   ┊ ┃ ┊ ┃ ┊
        # 3.00┊ 2 ┊ ┃ ┊ 2 ┊
        #     ┊ ┃ ┊ ┃ ┊ ┃ ┊
        # 2.00┊ ┃ ┊ 3 ┊ ┃ ┊
        #     ┊ ┃ ┊ ┃ ┊ ┃ ┊
        # 1.00┊ 1 ┊ ┃ ┊ ┃ ┊
        #     ┊ ┃ ┊ ┃ ┊ ┃ ┊
        # 0.00┊ 0 ┊ 0 ┊ 0 ┊
        #     0   2   4   6
        node_times = (0, 1, 3, 2, 4)
        samples = (0,)
        # (p, c, l, r)
        edges = [
            (1, 0, 0, 2),
            (2, 1, 0, 2),
            (3, 0, 2, 4),
            (4, 3, 2, 4),
            (4, 2, 4, 6),
            (2, 0, 4, 6),
        ]
        extended_edges = [
            (1, 0, 0, 6),
            (3, 1, 0, 6),
            (2, 3, 0, 6),
            (4, 2, 2, 6),
        ]
        site_positions = (3,)
        # site, node, derived_state, time
        mutations = [
            (0, 4, 5, 4.5),
            (0, 3, 4, 3.5),
            (0, 3, 3, 2.5),
            (0, 0, 2, 1.5),
            (0, 0, 1, 0.5),
        ]
        extended_mutations_node = [4, 2, 3, 1, 0]
        tables = tskit.TableCollection(sequence_length=6)
        for n, t in enumerate(node_times):
            flags = tskit.NODE_IS_SAMPLE if n in samples else 0
            tables.nodes.add_row(time=t, flags=flags)
        for p, c, l, r in edges:
            tables.edges.add_row(parent=p, child=c, left=l, right=r)
        for x in site_positions:
            tables.sites.add_row(ancestral_state="0", position=x)
        for s, n, d, t in mutations:
            tables.mutations.add_row(site=s, node=n, derived_state=str(d), time=t)
        tables.sort()
        tables.build_index()
        tables.compute_mutation_parents()
        ts = tables.tree_sequence()
        tables.edges.clear()
        for p, c, l, r in extended_edges:
            tables.edges.add_row(parent=p, child=c, left=l, right=r)
        tables.sort()
        tables.mutations.clear()
        for (s, _, d, t), n in zip(mutations, extended_mutations_node):
            tables.mutations.add_row(site=s, node=n, derived_state=str(d), time=t)
        tables.build_index()
        tables.compute_mutation_parents()
        ets = tables.tree_sequence()
        return ts, ets

    def get_example(self, j):
        if j == 1:
            ts, ets = self.get_example1()
        elif j == 2:
            ts, ets = self.get_example2()
        elif j == 3:
            ts, ets = self.get_example3()
        elif j == 4:
            ts, ets = self.get_example4()
        elif j == 5:
            ts, ets = self.get_example5()
        else:
            raise ValueError
        return ts, ets

    def verify_extend_haplotypes(self, ts, max_iter=10):
        ets = ts.extend_haplotypes(max_iter=max_iter)
        py_ets = extend_haplotypes(ts, max_iter=max_iter)
        ets.tables.assert_equals(py_ets.tables, ignore_provenance=True)
        self.verify_simplify_equality(ts, ets)

    def test_runs(self):
        ts = msprime.simulate(5, mutation_rate=1.0, random_seed=126)
        self.verify_extend_haplotypes(ts)
        self.naive_verify(ts)

    @pytest.mark.parametrize("j", [1, 2, 3, 4, 5])
    def test_example(self, j):
        ts, correct_ets = self.get_example(j)
        test_ets = ts.extend_haplotypes()
        test_ets.tables.assert_equals(correct_ets.tables, ignore_provenance=True)
        self.verify_extend_haplotypes(ts)
        self.naive_verify(ts)

    @pytest.mark.parametrize("j", [1, 2, 3, 4, 5])
    def test_redundant_breakpoitns(self, j):
        ts, correct_ets = self.get_example(j)
        ts = tsutil.insert_redundant_breakpoints(ts)
        test_ets = ts.extend_haplotypes()
        test_ets.tables.assert_equals(correct_ets.tables, ignore_provenance=True)
        self.verify_extend_haplotypes(ts)
        self.naive_verify(ts)

    def test_migrations_disallowed(self):
        ts = msprime.simulate(5, mutation_rate=1.0, random_seed=126)
        tables = ts.dump_tables()
        tables.populations.add_row()
        tables.populations.add_row()
        tables.migrations.add_row(0, 1, 0, 0, 1, 0)
        ts = tables.tree_sequence()
        with pytest.raises(
            _tskit.LibraryError, match="TSK_ERR_MIGRATIONS_NOT_SUPPORTED"
        ):
            _ = ts.extend_haplotypes()

    def test_unknown_times(self):
        ts = msprime.simulate(5, mutation_rate=1.0, random_seed=126)
        tables = ts.dump_tables()
        tables.mutations.clear()
        for mut in ts.mutations():
            tables.mutations.append(mut.replace(time=tskit.UNKNOWN_TIME))
        ts = tables.tree_sequence()
        with pytest.raises(
            _tskit.LibraryError, match="TSK_ERR_DISALLOWED_UNKNOWN_MUTATION_TIME"
        ):
            _ = ts.extend_haplotypes()

    def test_max_iter(self):
        ts = msprime.simulate(5, random_seed=126)
        with pytest.raises(_tskit.LibraryError, match="positive"):
            ets = ts.extend_haplotypes(max_iter=0)
        with pytest.raises(_tskit.LibraryError, match="positive"):
            ets = ts.extend_haplotypes(max_iter=-1)
        ets = ts.extend_haplotypes(max_iter=1)
        et = ets.extend_haplotypes(max_iter=1).dump_tables()
        eet = ets.extend_haplotypes(max_iter=2).dump_tables()
        eet.assert_equals(et)

    def test_very_simple(self):
        samples = [0]
        node_times = [0, 1, 2, 3]
        # (p, c, l, r)
        edges = [
            (1, 0, 0, 1),
            (2, 0, 1, 2),
            (2, 1, 0, 1),
            (3, 0, 2, 3),
            (3, 2, 0, 2),
        ]
        correct_edges = [
            (1, 0, 0, 3),
            (2, 1, 0, 3),
            (3, 2, 0, 3),
        ]
        tables = tskit.TableCollection(sequence_length=3)
        for n, t in enumerate(node_times):
            flags = tskit.NODE_IS_SAMPLE if n in samples else 0
            tables.nodes.add_row(time=t, flags=flags)
        for p, c, l, r in edges:
            tables.edges.add_row(parent=p, child=c, left=l, right=r)
        ts = tables.tree_sequence()
        ets = extend_haplotypes(ts)
        etables = ets.tables
        correct_tables = etables.copy()
        correct_tables.edges.clear()
        for p, c, l, r in correct_edges:
            correct_tables.edges.add_row(parent=p, child=c, left=l, right=r)
        etables.assert_equals(correct_tables, ignore_provenance=True)
        self.naive_verify(ts)

    def test_internal_samples(self):
        # Now we should have the same but not extend 5 (where * is),
        # since 5 is a sample; nor 8 because it's extension depends on 5
        #
        #    6         6      6         6
        #  +-+-+     +-+-+  +-+-+     +-+-+
        #  7   *     7   *  7   8     7   8
        #  |   |    ++-+ |  | +-++    |   |
        #  4   5    4  | *  4 |  5    4   5
        # +++ +++  +++ | |  | | +++  +++ +++
        # 0 1 2 3  0 1 2 3  0 1 2 3  0 1 2 3
        #
        node_times = {
            0: 0,
            1: 0,
            2: 0,
            3: 0,
            4: 1.0,
            5: 1.0,
            6: 3.0,
            7: 2.0,
            8: 2.0,
        }
        # (p, c, l, r)
        edges = [
            (4, 0, 0, 10),
            (4, 1, 0, 5),
            (4, 1, 7, 10),
            (5, 2, 0, 2),
            (5, 2, 5, 10),
            (5, 3, 0, 2),
            (5, 3, 5, 10),
            (7, 2, 2, 5),
            (7, 4, 0, 10),
            (8, 1, 5, 7),
            (8, 5, 5, 10),
            (6, 3, 2, 5),
            (6, 5, 0, 2),
            (6, 7, 0, 10),
            (6, 8, 5, 10),
        ]
        tables = tskit.TableCollection(sequence_length=10)
        samples = [0, 1, 2, 3, 5]
        for n, t in node_times.items():
            flags = tskit.NODE_IS_SAMPLE if n in samples else 0
            tables.nodes.add_row(time=t, flags=flags)
        for p, c, l, r in edges:
            tables.edges.add_row(parent=p, child=c, left=l, right=r)
        ts = tables.tree_sequence()
        ets = extend_haplotypes(ts)
        # nothing should have happened
        ets.tables.assert_equals(tables)
        self.verify_extend_haplotypes(ts)
        self.naive_verify(ts)

    @pytest.mark.parametrize("seed", [3, 4, 5, 6])
    def test_wf(self, seed):
        tables = wf.wf_sim(N=6, ngens=9, num_loci=100, deep_history=False, seed=seed)
        tables.sort()
        ts = tables.tree_sequence().simplify()
        self.verify_extend_haplotypes(ts)
        self.naive_verify(ts)


class TestExamples(TestExtendThings):
    """
    Compare the ts method with local implementation.
    """

    def check(self, ts):
        if np.any(tskit.is_unknown_time(ts.mutations_time)):
            tables = ts.dump_tables()
            tables.compute_mutation_times()
            ts = tables.tree_sequence()
        py_ets = extend_haplotypes(ts)
        self.verify_simplify_equality(ts, py_ets)
        lib_ts = ts.extend_haplotypes()
        lib_ts.tables.assert_equals(py_ets.tables)
        assert np.all(ts.genotype_matrix() == lib_ts.genotype_matrix())
        sts = ts.simplify()
        lib_sts = lib_ts.simplify()
        lib_sts.tables.assert_equals(sts.tables, ignore_provenance=True)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_suite_examples_defaults(self, ts):
        if ts.num_migrations == 0:
            self.check(ts)
        else:
            pass
            with pytest.raises(
                _tskit.LibraryError, match="TSK_ERR_MIGRATIONS_NOT_SUPPORTED"
            ):
                _ = ts.extend_haplotypes()

    @pytest.mark.parametrize("n", [3, 4, 5])
    def test_all_trees_ts(self, n):
        ts = tsutil.all_trees_ts(n)
        self.check(ts)


--- ../../tskit/python/tests/test_cli.py ---


"""
Test cases for the command line interfaces to tskit
"""
import io
import os
import sys
import tempfile
import unittest
from unittest import mock

import h5py
import msprime
import pytest

import tskit
import tskit.cli as cli
from . import tsutil


class TestException(Exception):
    __test__ = False
    """
    Custom exception we can throw for testing.
    """


def capture_output(func, *args, **kwargs):
    """
    Runs the specified function and arguments, and returns the
    tuple (stdout, stderr) as strings.
    """
    buffer_class = io.BytesIO
    if sys.version_info[0] == 3:
        buffer_class = io.StringIO
    stdout = sys.stdout
    sys.stdout = buffer_class()
    stderr = sys.stderr
    sys.stderr = buffer_class()

    try:
        # Recent versions of MacOS seem to have issues with us calling signal
        # during tests.
        with mock.patch("signal.signal"):
            func(*args, **kwargs)
        stdout_output = sys.stdout.getvalue()
        stderr_output = sys.stderr.getvalue()
    finally:
        sys.stdout.close()
        sys.stdout = stdout
        sys.stderr.close()
        sys.stderr = stderr
    return stdout_output, stderr_output


class TestCli(unittest.TestCase):
    """
    Superclass of tests for the CLI needing temp files.
    """

    def setUp(self):
        fd, self.temp_file = tempfile.mkstemp(prefix="tsk_cli_testcase_")
        os.close(fd)

    def tearDown(self):
        os.unlink(self.temp_file)


class TestTskitArgumentParser:
    """
    Tests for the argument parsers in msp.
    """

    def test_individuals_default_values(self):
        parser = cli.get_tskit_parser()
        cmd = "individuals"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence])
        assert args.tree_sequence == tree_sequence
        assert args.precision == 6

    def test_individuals_short_args(self):
        parser = cli.get_tskit_parser()
        cmd = "individuals"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence, "-p", "8"])
        assert args.tree_sequence == tree_sequence
        assert args.precision == 8

    def test_individuals_long_args(self):
        parser = cli.get_tskit_parser()
        cmd = "individuals"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence, "--precision", "5"])
        assert args.tree_sequence == tree_sequence
        assert args.precision == 5

    def test_nodes_default_values(self):
        parser = cli.get_tskit_parser()
        cmd = "nodes"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence])
        assert args.tree_sequence == tree_sequence
        assert args.precision == 6

    def test_nodes_short_args(self):
        parser = cli.get_tskit_parser()
        cmd = "nodes"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence, "-p", "8"])
        assert args.tree_sequence == tree_sequence
        assert args.precision == 8

    def test_nodes_long_args(self):
        parser = cli.get_tskit_parser()
        cmd = "nodes"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence, "--precision", "5"])
        assert args.tree_sequence == tree_sequence
        assert args.precision == 5

    def test_edges_default_values(self):
        parser = cli.get_tskit_parser()
        cmd = "edges"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence])
        assert args.tree_sequence == tree_sequence
        assert args.precision == 6

    def test_edges_short_args(self):
        parser = cli.get_tskit_parser()
        cmd = "edges"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence, "-p", "8"])
        assert args.tree_sequence == tree_sequence
        assert args.precision == 8

    def test_edges_long_args(self):
        parser = cli.get_tskit_parser()
        cmd = "edges"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence, "--precision", "5"])
        assert args.tree_sequence == tree_sequence
        assert args.precision == 5

    def test_sites_default_values(self):
        parser = cli.get_tskit_parser()
        cmd = "sites"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence])
        assert args.tree_sequence == tree_sequence
        assert args.precision == 6

    def test_sites_short_args(self):
        parser = cli.get_tskit_parser()
        cmd = "sites"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence, "-p", "8"])
        assert args.tree_sequence == tree_sequence
        assert args.precision == 8

    def test_sites_long_args(self):
        parser = cli.get_tskit_parser()
        cmd = "sites"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence, "--precision", "5"])
        assert args.tree_sequence == tree_sequence
        assert args.precision == 5

    def test_mutations_default_values(self):
        parser = cli.get_tskit_parser()
        cmd = "mutations"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence])
        assert args.tree_sequence == tree_sequence
        assert args.precision == 6

    def test_mutations_short_args(self):
        parser = cli.get_tskit_parser()
        cmd = "mutations"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence, "-p", "4"])
        assert args.tree_sequence == tree_sequence
        assert args.precision == 4

    def test_mutations_long_args(self):
        parser = cli.get_tskit_parser()
        cmd = "mutations"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence, "--precision", "9"])
        assert args.tree_sequence == tree_sequence
        assert args.precision == 9

    def test_provenances_default_values(self):
        parser = cli.get_tskit_parser()
        cmd = "provenances"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence])
        assert args.tree_sequence == tree_sequence
        assert not args.human

    def test_provenances_short_args(self):
        parser = cli.get_tskit_parser()
        cmd = "provenances"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence, "-H"])
        assert args.tree_sequence == tree_sequence
        assert args.human

    def test_provenances_long_args(self):
        parser = cli.get_tskit_parser()
        cmd = "provenances"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence, "--human"])
        assert args.tree_sequence == tree_sequence
        assert args.human

    @pytest.mark.skip(reason="fasta output temporarily disabled")
    def test_fasta_default_values(self):
        parser = cli.get_tskit_parser()
        cmd = "fasta"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence])
        assert args.tree_sequence == tree_sequence
        assert args.wrap == 60

    @pytest.mark.skip(reason="fasta output temporarily disabled")
    def test_fasta_short_args(self):
        parser = cli.get_tskit_parser()
        cmd = "fasta"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence, "-w", "100"])
        assert args.tree_sequence == tree_sequence
        assert args.wrap == 100

    @pytest.mark.skip(reason="fasta output temporarily disabled")
    def test_fasta_long_args(self):
        parser = cli.get_tskit_parser()
        cmd = "fasta"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence, "--wrap", "50"])
        assert args.tree_sequence == tree_sequence
        assert args.wrap == 50

    @pytest.mark.parametrize(
        "flags,expected",
        (
            [[], None],
            [["-P", "2"], 2],
            [["--ploidy", "5"], 5],
        ),
    )
    def test_vcf_ploidy(self, flags, expected):
        parser = cli.get_tskit_parser()
        cmd = "vcf"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence, *flags])
        assert args.tree_sequence == tree_sequence
        assert args.ploidy == expected

    @pytest.mark.parametrize(
        "flags,expected",
        (
            [[], "1"],
            [["-c", "chrX"], "chrX"],
            [["--contig-id", "chr20"], "chr20"],
        ),
    )
    def test_vcf_contig_id(self, flags, expected):
        parser = cli.get_tskit_parser()
        cmd = "vcf"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence, *flags])
        assert args.tree_sequence == tree_sequence
        assert args.contig_id == expected

    @pytest.mark.parametrize(
        "flags,expected",
        (
            [[], False],
            [["-0"], True],
            [["--allow-position-zero"], True],
        ),
    )
    def test_vcf_allow_position_zero(self, flags, expected):
        parser = cli.get_tskit_parser()
        cmd = "vcf"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence, *flags])
        assert args.tree_sequence == tree_sequence
        assert args.allow_position_zero == expected

    def test_upgrade_default_values(self):
        parser = cli.get_tskit_parser()
        cmd = "upgrade"
        source = "in.trees"
        destination = "out.trees"
        args = parser.parse_args([cmd, source, destination])
        assert args.source == source
        assert args.destination == destination
        assert not args.remove_duplicate_positions

    def test_info_default_values(self):
        parser = cli.get_tskit_parser()
        cmd = "info"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence])
        assert args.tree_sequence == tree_sequence

    def test_populations_default_values(self):
        parser = cli.get_tskit_parser()
        cmd = "populations"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence])
        assert args.tree_sequence == tree_sequence

    def test_migrations_default_values(self):
        parser = cli.get_tskit_parser()
        cmd = "migrations"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence])
        assert args.tree_sequence == tree_sequence

    def test_migrations_short_args(self):
        parser = cli.get_tskit_parser()
        cmd = "migrations"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence, "-p", "2"])
        assert args.tree_sequence == tree_sequence
        assert args.precision == 2

    def test_migrations_long_args(self):
        parser = cli.get_tskit_parser()
        cmd = "migrations"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence, "--precision", "5"])
        assert args.tree_sequence == tree_sequence
        assert args.precision == 5

    def test_trees_default_values(self):
        parser = cli.get_tskit_parser()
        cmd = "trees"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence])
        assert args.tree_sequence == tree_sequence
        assert args.precision == 6
        assert not args.draw

    def test_trees_short_args(self):
        parser = cli.get_tskit_parser()
        cmd = "trees"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence, "-d", "-p", "8"])
        assert args.tree_sequence == tree_sequence
        assert args.precision == 8
        assert args.draw

    def test_trees_long_args(self):
        parser = cli.get_tskit_parser()
        cmd = "trees"
        tree_sequence = "test.trees"
        args = parser.parse_args([cmd, tree_sequence, "--precision", "5", "--draw"])
        assert args.tree_sequence == tree_sequence
        assert args.precision == 5
        assert args.draw


class TestTskitConversionOutput(unittest.TestCase):
    """
    Tests the output of tskit to ensure it's correct.
    """

    @classmethod
    def setUpClass(cls):
        ts = msprime.simulate(
            length=1,
            recombination_rate=2,
            mutation_rate=2,
            random_seed=1,
            migration_matrix=[[0, 1], [1, 0]],
            population_configurations=[
                msprime.PopulationConfiguration(5) for _ in range(2)
            ],
            record_migrations=True,
        )
        assert ts.num_migrations > 0
        cls._tree_sequence = tsutil.insert_random_ploidy_individuals(
            ts, samples_only=True
        )
        fd, cls._tree_sequence_file = tempfile.mkstemp(
            prefix="tsk_cli", suffix=".trees"
        )
        os.close(fd)
        cls._tree_sequence.dump(cls._tree_sequence_file)

    @classmethod
    def tearDownClass(cls):
        os.unlink(cls._tree_sequence_file)

    def verify_individuals(self, output_individuals, precision):
        with tempfile.TemporaryFile("w+") as f:
            self._tree_sequence.dump_text(individuals=f, precision=precision)
            f.seek(0)
            output = f.read().splitlines()
        assert output == output_individuals

    def test_individuals(self):
        cmd = "individuals"
        precision = 8
        stdout, stderr = capture_output(
            cli.tskit_main, [cmd, self._tree_sequence_file, "-p", str(precision)]
        )
        assert len(stderr) == 0
        output_individuals = stdout.splitlines()
        self.verify_individuals(output_individuals, precision)

    def verify_nodes(self, output_nodes, precision):
        with tempfile.TemporaryFile("w+") as f:
            self._tree_sequence.dump_text(nodes=f, precision=precision)
            f.seek(0)
            output = f.read().splitlines()
        assert output == output_nodes

    def test_nodes(self):
        cmd = "nodes"
        precision = 8
        stdout, stderr = capture_output(
            cli.tskit_main, [cmd, self._tree_sequence_file, "-p", str(precision)]
        )
        assert len(stderr) == 0
        output_nodes = stdout.splitlines()
        self.verify_nodes(output_nodes, precision)

    def verify_edges(self, output_edges, precision):
        with tempfile.TemporaryFile("w+") as f:
            self._tree_sequence.dump_text(edges=f, precision=precision)
            f.seek(0)
            output = f.read().splitlines()
        assert output == output_edges

    def test_edges(self):
        cmd = "edges"
        precision = 8
        stdout, stderr = capture_output(
            cli.tskit_main, [cmd, self._tree_sequence_file, "-p", str(precision)]
        )
        assert len(stderr) == 0
        output_edges = stdout.splitlines()
        self.verify_edges(output_edges, precision)

    def verify_sites(self, output_sites, precision):
        with tempfile.TemporaryFile("w+") as f:
            self._tree_sequence.dump_text(sites=f, precision=precision)
            f.seek(0)
            output = f.read().splitlines()
        assert output == output_sites

    def test_sites(self):
        cmd = "sites"
        precision = 8
        stdout, stderr = capture_output(
            cli.tskit_main, [cmd, self._tree_sequence_file, "-p", str(precision)]
        )
        assert len(stderr) == 0
        output_sites = stdout.splitlines()
        self.verify_sites(output_sites, precision)

    def verify_mutations(self, output_mutations, precision):
        with tempfile.TemporaryFile("w+") as f:
            self._tree_sequence.dump_text(mutations=f, precision=precision)
            f.seek(0)
            output = f.read().splitlines()
        assert output == output_mutations

    def test_mutations(self):
        cmd = "mutations"
        precision = 8
        stdout, stderr = capture_output(
            cli.tskit_main, [cmd, self._tree_sequence_file, "-p", str(precision)]
        )
        assert len(stderr) == 0
        output_mutations = stdout.splitlines()
        self.verify_mutations(output_mutations, precision)

    def verify_migrations(self, output_migrations, precision):
        with tempfile.TemporaryFile("w+") as f:
            self._tree_sequence.dump_text(migrations=f, precision=precision)
            f.seek(0)
            output = f.read().splitlines()
        assert output == output_migrations

    def test_migrations(self):
        cmd = "migrations"
        precision = 4
        stdout, stderr = capture_output(
            cli.tskit_main, [cmd, self._tree_sequence_file, "-p", str(precision)]
        )
        assert len(stderr) == 0
        output_migrations = stdout.splitlines()
        self.verify_migrations(output_migrations, precision)

    def verify_provenances(self, output_provenances):
        with tempfile.TemporaryFile("w+") as f:
            self._tree_sequence.dump_text(provenances=f)
            f.seek(0)
            output = f.read().splitlines()
        assert output == output_provenances

    def test_provenances(self):
        cmd = "provenances"
        stdout, stderr = capture_output(cli.tskit_main, [cmd, self._tree_sequence_file])
        assert len(stderr) == 0
        output_provenances = stdout.splitlines()
        self.verify_provenances(output_provenances)

    def test_provenances_human(self):
        cmd = "provenances"
        stdout, stderr = capture_output(
            cli.tskit_main, [cmd, "-H", self._tree_sequence_file]
        )
        assert len(stderr) == 0
        output_provenances = stdout.splitlines()
        # TODO Check the actual output here.
        assert len(output_provenances) > 0

    def verify_fasta(self, output_fasta):
        with tempfile.TemporaryFile("w+") as f:
            self._tree_sequence.write_fasta(f)
            f.seek(0)
            fasta = f.read()
        assert output_fasta == fasta

    @pytest.mark.skip(reason="fasta output temporarily disabled")
    def test_fasta(self):
        cmd = "fasta"
        stdout, stderr = capture_output(cli.tskit_main, [cmd, self._tree_sequence_file])
        assert len(stderr) == 0
        self.verify_fasta(stdout)

    def verify_vcf(self, output_vcf):
        with tempfile.TemporaryFile("w+") as f:
            self._tree_sequence.write_vcf(f, allow_position_zero=True)
            f.seek(0)
            vcf = f.read()
        assert output_vcf == vcf

    def test_vcf(self):
        cmd = "vcf"
        stdout, stderr = capture_output(
            cli.tskit_main, [cmd, "-0", self._tree_sequence_file]
        )
        assert len(stderr) == 0
        self.verify_vcf(stdout)

    def verify_info(self, ts, output_info):
        assert str(ts) == output_info

    def test_info(self):
        cmd = "info"
        stdout, stderr = capture_output(cli.tskit_main, [cmd, self._tree_sequence_file])
        assert len(stderr) == 0
        ts = tskit.load(self._tree_sequence_file)
        self.verify_info(ts, stdout[:-1])

    def test_trees_no_draw(self):
        cmd = "trees"
        stdout, stderr = capture_output(cli.tskit_main, [cmd, self._tree_sequence_file])
        assert len(stderr) == 0
        ts = tskit.load(self._tree_sequence_file)
        assert len(stdout.splitlines()) == 3 * ts.num_trees

    def test_trees_draw(self):
        cmd = "trees"
        stdout, stderr = capture_output(
            cli.tskit_main, [cmd, "-d", self._tree_sequence_file]
        )
        assert len(stderr) == 0
        ts = tskit.load(self._tree_sequence_file)
        assert len(stdout.splitlines()) > 3 * ts.num_trees


class TestVCFZeroPosition:
    """
    Tests that we can write VCF files with position 0.
    """

    def test_zero_position(self, tmp_path):
        ts = msprime.simulate(10, mutation_rate=1, random_seed=1)
        ts.dump(tmp_path / "test.trees")
        with pytest.raises(ValueError):
            capture_output(cli.tskit_main, ["vcf", str(tmp_path / "test.trees")])
        stdout, stderr = capture_output(
            cli.tskit_main, ["vcf", "-0", str(tmp_path / "test.trees")]
        )
        assert len(stderr) == 0


class TestBadFile:
    """
    Tests that we deal with IO errors appropriately.
    """

    def verify(self, command):
        with mock.patch("sys.exit", side_effect=TestException) as mocked_exit:
            with pytest.raises(TestException):
                capture_output(cli.tskit_main, ["info", "/no/such/file"])
            mocked_exit.assert_called_once_with(
                "Load error: [Errno 2] No such file or directory: '/no/such/file'"
            )

    def test_info(self):
        self.verify("info")

    def test_fasta(self):
        self.verify("fasta")

    def test_vcf(self):
        self.verify("vcf")

    def test_nodes(self):
        self.verify("nodes")

    def test_edges(self):
        self.verify("edges")

    def test_sites(self):
        self.verify("sites")

    def test_mutations(self):
        self.verify("mutations")

    def test_migrations(self):
        self.verify("migrations")

    def test_provenances(self):
        self.verify("provenances")


class TestUpgrade(TestCli):
    """
    Tests the results of the upgrade operation to ensure they are
    correct.
    """

    def setUp(self):
        fd, self.legacy_file_name = tempfile.mkstemp(prefix="msp_cli", suffix=".trees")
        os.close(fd)
        fd, self.current_file_name = tempfile.mkstemp(prefix="msp_cli", suffix=".trees")
        os.close(fd)

    def tearDown(self):
        os.unlink(self.legacy_file_name)
        os.unlink(self.current_file_name)

    def test_conversion(self):
        ts1 = msprime.simulate(10)
        for version in [2, 3]:
            tskit.dump_legacy(ts1, self.legacy_file_name, version=version)
            stdout, stderr = capture_output(
                cli.tskit_main,
                ["upgrade", self.legacy_file_name, self.current_file_name],
            )
            ts2 = tskit.load(self.current_file_name)
            assert stdout == ""
            assert stderr == ""
            # Quick checks to ensure we have the right tree sequence.
            # More thorough checks are done elsewhere.
            assert ts1.get_sample_size() == ts2.get_sample_size()
            assert ts1.num_edges == ts2.num_edges
            assert ts1.get_num_trees() == ts2.get_num_trees()

    def test_duplicate_positions(self):
        ts = msprime.simulate(10, mutation_rate=10)
        for version in [2, 3]:
            tskit.dump_legacy(ts, self.legacy_file_name, version=version)
            root = h5py.File(self.legacy_file_name, "r+")
            root["mutations/position"][:] = 0
            root.close()
            stdout, stderr = capture_output(
                cli.tskit_main,
                ["upgrade", "-d", self.legacy_file_name, self.current_file_name],
            )
            assert stdout == ""
            tsp = tskit.load(self.current_file_name)
            assert tsp.sample_size == ts.sample_size
            assert tsp.num_sites == 1

    def test_duplicate_positions_error(self):
        ts = msprime.simulate(10, mutation_rate=10)
        for version in [2, 3]:
            tskit.dump_legacy(ts, self.legacy_file_name, version=version)
            root = h5py.File(self.legacy_file_name, "r+")
            root["mutations/position"][:] = 0
            root.close()
            with mock.patch("sys.exit", side_effect=TestException) as mocked_exit:
                with pytest.raises(TestException):
                    capture_output(
                        cli.tskit_main,
                        ["upgrade", self.legacy_file_name, self.current_file_name],
                    )
                assert mocked_exit.call_count == 1


--- ../../tskit/python/tests/test_phylo_formats.py ---


"""
Tests for phylogenetics export functions, newick, nexus, FASTA etc.
"""
import functools
import io
import textwrap

import dendropy
import msprime
import newick
import numpy as np
import pytest
from Bio import SeqIO

import tests
import tskit
from tests.test_highlevel import get_example_tree_sequences

# ↑ See https://github.com/tskit-dev/tskit/issues/1804 for when
# we can remove this. The example_ts here is intended to be the
# basic tree sequence which should give a meaningful result for
# most operations. Probably rename it to ``examples.simple_ts()``
# or something.


@functools.lru_cache(maxsize=100)
def alignment_example(sequence_length, include_reference=True):
    ts = msprime.sim_ancestry(
        samples=5, sequence_length=sequence_length, random_seed=123
    )
    ts = msprime.sim_mutations(ts, rate=0.1, random_seed=1234)
    tables = ts.dump_tables()
    if include_reference:
        tables.reference_sequence.data = tskit.random_nucleotides(
            ts.sequence_length, seed=1234
        )
    ts = tables.tree_sequence()
    assert ts.num_sites > 5
    return ts


@tests.cached_example
def missing_data_example():
    # 2.00┊   4     ┊
    #     ┊ ┏━┻┓    ┊
    # 1.00┊ ┃  3    ┊
    #     ┊ ┃ ┏┻┓   ┊
    # 0.00┊ 0 1 2 5 ┊
    #     0        10
    #      |      |
    #  pos 2      9
    #  anc A      T
    ts = tskit.Tree.generate_balanced(3, span=10).tree_sequence
    tables = ts.dump_tables()
    tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
    tables.sites.add_row(2, ancestral_state="A")
    tables.sites.add_row(9, ancestral_state="T")
    tables.mutations.add_row(site=0, node=0, derived_state="G")
    tables.mutations.add_row(site=1, node=3, derived_state="C")
    return tables.tree_sequence()


def alignment_map(ts, **kwargs):
    alignments = ts.alignments(**kwargs)
    return {f"n{u}": alignment for u, alignment in zip(ts.samples(), alignments)}


def assert_fully_labelled_trees_equal(tree, root, node_labels, dpy_tree):
    """
    Checks the the specified fully-labelled tree rooted at the specified
    node is equivalent to the specified Dendropy tree.
    """
    label_map = {}
    for node in dpy_tree:
        label_map[str(node.taxon.label)] = node

    for u in tree.nodes(root, order="postorder"):
        # Consume the nodes in the dendropy node map one-by-one
        dpy_node = label_map.pop(node_labels[u])
        parent = tree.parent(u)
        if parent == tskit.NULL:
            assert dpy_node.edge_length is None
            assert dpy_node.parent_node is None
        else:
            assert tree.branch_length(u) == pytest.approx(dpy_node.edge_length)
            assert dpy_node.parent_node is label_map[node_labels[parent]]
    assert len(label_map) == 0


def assert_sample_labelled_trees_equal(tree, dpy_tree):
    """
    Checks that the specified trees are equivalent, where the dendropy tree
    only has labels identifying the samples.
    """
    for sample in tree.samples():
        dpy_node = dpy_tree.find_node_with_taxon_label(f"n{sample}")
        # Check the branch length paths to root are equal
        p1 = []
        u = sample
        while tree.parent(u) != tskit.NULL:
            p1.append(tree.branch_length(u))
            u = tree.parent(u)
        p2 = []
        while dpy_node.parent_node is not None:
            p2.append(dpy_node.edge_length)
            dpy_node = dpy_node.parent_node
        assert len(p1) == len(p2)
        np.testing.assert_array_almost_equal(p1, p2)


def assert_dpy_tree_list_equal(ts, tree_list):
    """
    Check that the nexus-encoded tree list output from tskit is
    parsed correctly by dendropy.
    """
    assert ts.num_trees == len(tree_list)
    for tsk_tree, dpy_tree in zip(ts.trees(), tree_list):
        # We're specifying that the tree is rooted.
        assert dpy_tree.is_rooted
        assert dpy_tree.label.startswith("t")
        left, right = map(float, dpy_tree.label[1:].split("^"))
        assert tsk_tree.interval.left == pytest.approx(left)
        assert tsk_tree.interval.right == pytest.approx(right)
        assert_sample_labelled_trees_equal(tsk_tree, dpy_tree)


class TestBackendsGiveIdenticalOutput:
    # At the default precision of 17 we should get identical results between
    # the two backends as there's no rounding done. In general, we can't
    # depend on this, though, since rounding may be done differently by the
    # Python and C library implementations.
    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_default_precision(self, ts):
        for tree in ts.trees():
            if tree.has_single_root:
                assert tree.as_newick() == tree.as_newick(
                    node_labels={u: f"n{u}" for u in tree.samples()}
                )


class TestNewickRoundTrip:
    """
    Test that the newick formats can round-trip the data under various
    assumptions.
    """

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_leaf_labels_newick_lib(self, ts):
        for tree in ts.trees():
            # Multiroot trees raise an error
            for root in tree.roots:
                leaf_labels = {u: f"n{u}" for u in tree.leaves(root)}
                ns = tree.newick(
                    root=root,
                    precision=16,
                    node_labels=leaf_labels,
                )
                newick_tree = newick.loads(
                    ns, length_parser=lambda x: None if x is None else float(x)
                )[0]
                leaf_names = newick_tree.get_leaf_names()
                assert sorted(leaf_names) == sorted(leaf_labels.values())
                for u in tree.leaves(root):
                    name = leaf_labels[u]
                    node = newick_tree.get_node(name)
                    while u != root:
                        assert node.length == pytest.approx(tree.branch_length(u))
                        node = node.ancestor
                        u = tree.parent(u)
                    assert node.ancestor is None

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_all_node_labels_dendropy(self, ts):
        node_labels = {u: f"n{u}" for u in range(ts.num_nodes)}
        for tree in ts.trees():
            # Multiroot trees raise an error
            for root in tree.roots:
                ns = tree.newick(
                    root=root,
                    precision=16,
                    node_labels=node_labels,
                )
                dpy_tree = dendropy.Tree.get(
                    data=ns, suppress_internal_node_taxa=False, schema="newick"
                )
                assert_fully_labelled_trees_equal(tree, root, node_labels, dpy_tree)


class TestNexusTreeRoundTrip:
    """
    Test that the nexus format can round-trip tree data under various
    assumptions.
    """

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_dendropy_defaults(self, ts):
        if any(tree.num_roots != 1 for tree in ts.trees()):
            with pytest.raises(ValueError, match="single root"):
                ts.as_nexus(include_alignments=False)
        else:
            nexus = ts.as_nexus(include_alignments=False)
            tree_list = dendropy.TreeList()
            tree_list.read(
                data=nexus,
                schema="nexus",
                suppress_internal_node_taxa=False,
            )
            assert_dpy_tree_list_equal(ts, tree_list)


class TestNexusIncludeSections:
    """
    Test if we include the sections as expected.
    """

    @tests.cached_example
    def ts(self):
        # 2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊
        # 1.00┊ ┃  3  ┊
        #     ┊ ┃ ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0      10
        #      |    |
        #  pos 2    9
        #  anc A    T
        ts = tskit.Tree.generate_balanced(3, span=10).tree_sequence
        tables = ts.dump_tables()
        tables.sites.add_row(2, ancestral_state="A")
        tables.sites.add_row(9, ancestral_state="T")
        tables.mutations.add_row(site=0, node=0, derived_state="G")
        tables.mutations.add_row(site=1, node=3, derived_state="C")
        return tables.tree_sequence()

    def test_nexus_default(self):
        ref = "ACTGACTGAC"
        expected = textwrap.dedent(
            """\
            #NEXUS
            BEGIN TAXA;
              DIMENSIONS NTAX=3;
              TAXLABELS n0 n1 n2;
            END;
            BEGIN DATA;
              DIMENSIONS NCHAR=10;
              FORMAT DATATYPE=DNA MISSING=?;
              MATRIX
                n0 ACGGACTGAT
                n1 ACAGACTGAC
                n2 ACAGACTGAC
              ;
            END;
            BEGIN TREES;
              TREE t0^10 = [&R] (n0:2,(n1:1,n2:1):1);
            END;
            """
        )
        assert expected == self.ts().as_nexus(reference_sequence=ref)

    def test_nexus_no_trees(self):
        ref = "ACTGACTGAC"
        expected = textwrap.dedent(
            """\
            #NEXUS
            BEGIN TAXA;
              DIMENSIONS NTAX=3;
              TAXLABELS n0 n1 n2;
            END;
            BEGIN DATA;
              DIMENSIONS NCHAR=10;
              FORMAT DATATYPE=DNA MISSING=?;
              MATRIX
                n0 ACGGACTGAT
                n1 ACAGACTGAC
                n2 ACAGACTGAC
              ;
            END;
            """
        )
        assert expected == self.ts().as_nexus(
            reference_sequence=ref, include_trees=False
        )

    def test_nexus_no_alignments(self):
        expected = textwrap.dedent(
            """\
            #NEXUS
            BEGIN TAXA;
              DIMENSIONS NTAX=3;
              TAXLABELS n0 n1 n2;
            END;
            BEGIN TREES;
              TREE t0^10 = [&R] (n0:2,(n1:1,n2:1):1);
            END;
            """
        )
        assert expected == self.ts().as_nexus(include_alignments=False)

    def test_nexus_no_trees_or_alignments(self):
        expected = textwrap.dedent(
            """\
            #NEXUS
            BEGIN TAXA;
              DIMENSIONS NTAX=3;
              TAXLABELS n0 n1 n2;
            END;
            """
        )
        assert expected == self.ts().as_nexus(
            include_trees=False, include_alignments=False
        )


class TestNewickCodePaths:
    """
    Test that the different code paths we use under the hood lead to
    identical results.
    """

    # NOTE this probabably won't work in general because the C and
    # Python code paths using different rounding algorithms.

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_default_newick(self, ts):
        for tree in ts.trees():
            for root in tree.roots:
                ns1 = tree.newick(root=root)
                node_labels = {u: str(u + 1) for u in tree.leaves()}
                ns2 = tree.newick(root=root, node_labels=node_labels)
                assert ns1 == ns2

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_default_as_newick(self, ts):
        for tree in ts.trees():
            for root in tree.roots:
                ns1 = tree.as_newick(root=root)
                node_labels = {u: f"n{u}" for u in tree.tree_sequence.samples()}
                ns2 = tree.as_newick(root=root, node_labels=node_labels)
                assert ns1 == ns2


class TestBalancedBinaryExample:
    #   4
    # ┏━┻┓
    # ┃  3
    # ┃ ┏┻┓
    # 0 1 2
    @tests.cached_example
    def tree(self):
        return tskit.Tree.generate_balanced(3)

    def test_newick_default(self):
        s = (
            "(1:2.00000000000000,(2:1.00000000000000,3:1.00000000000000)"
            ":1.00000000000000);"
        )
        assert self.tree().newick() == s

    def test_as_newick_default(self):
        s = "(n0:2,(n1:1,n2:1):1);"
        assert self.tree().as_newick() == s

    def test_newick_zero_precision(self):
        s = "(1:2,(2:1,3:1):1);"
        assert self.tree().newick(precision=0) == s

    def test_as_newick_zero_precision(self):
        s = "(n0:2,(n1:1,n2:1):1);"
        assert self.tree().as_newick(precision=0) == s

    def test_as_newick_precision_1(self):
        s = "(n0:2.0,(n1:1.0,n2:1.0):1.0);"
        assert self.tree().as_newick(precision=1) == s

    def test_as_newick_precision_1_explicit_labels(self):
        tree = self.tree()
        s = "(x0:2.0,(x1:1.0,x2:1.0):1.0);"
        node_labels = {u: f"x{u}" for u in tree.samples()}
        assert tree.as_newick(precision=1, node_labels=node_labels) == s

    def test_newick_no_branch_lengths(self):
        s = "(1,(2,3));"
        assert self.tree().newick(include_branch_lengths=False) == s

    def test_as_newick_no_branch_lengths(self):
        s = "(n0,(n1,n2));"
        assert self.tree().as_newick(include_branch_lengths=False) == s

    def test_newick_all_node_labels(self):
        s = "(0:2,(1:1,2:1)3:1)4;"
        node_labels = {u: str(u) for u in self.tree().nodes()}
        ns = self.tree().newick(precision=0, node_labels=node_labels)
        assert s == ns

    def test_as_newick_all_node_labels(self):
        s = "(0:2,(1:1,2:1)3:1)4;"
        node_labels = {u: str(u) for u in self.tree().nodes()}
        ns = self.tree().as_newick(node_labels=node_labels)
        assert s == ns

    def test_as_newick_variable_length_node_labels(self):
        s = "(:2,(1:1,22:1)333:1)4444;"
        node_labels = {u: str(u) * u for u in self.tree().nodes()}
        ns = self.tree().as_newick(node_labels=node_labels)
        assert s == ns

    def test_as_newick_empty_node_labels(self):
        s = "(:2,(:1,:1):1);"
        ns = self.tree().as_newick(node_labels={})
        assert s == ns

    def test_newick_partial_node_labels(self):
        s = "(0:2,(1:1,2:1)3:1);"
        node_labels = {u: str(u) for u in self.tree().preorder()[1:]}
        ns = self.tree().newick(precision=0, node_labels=node_labels)
        assert s == ns

    def test_newick_root(self):
        s = "(2:1,3:1);"
        assert self.tree().newick(root=3, precision=0) == s

    def test_as_nexus_default(self):
        ts = self.tree().tree_sequence
        expected = textwrap.dedent(
            """\
            #NEXUS
            BEGIN TAXA;
              DIMENSIONS NTAX=3;
              TAXLABELS n0 n1 n2;
            END;
            BEGIN TREES;
              TREE t0^1 = [&R] (n0:2,(n1:1,n2:1):1);
            END;
        """
        )
        assert ts.as_nexus() == expected

    def test_as_nexus_precision_1(self):
        ts = self.tree().tree_sequence
        expected = textwrap.dedent(
            """\
            #NEXUS
            BEGIN TAXA;
              DIMENSIONS NTAX=3;
              TAXLABELS n0 n1 n2;
            END;
            BEGIN TREES;
              TREE t0.0^1.0 = [&R] (n0:2.0,(n1:1.0,n2:1.0):1.0);
            END;
        """
        )
        assert ts.as_nexus(precision=1) == expected


class TestFractionalBranchLengths:
    # 0.67┊   4   ┊
    #     ┊ ┏━┻┓  ┊
    # 0.33┊ ┃  3  ┊
    #     ┊ ┃ ┏┻┓ ┊
    # 0.00┊ 0 1 2 ┊
    #     0       1
    @tests.cached_example
    def tree(self):
        return tskit.Tree.generate_balanced(3, branch_length=1 / 3)

    def test_newick_default(self):
        s = (
            "(1:0.66666666666667,(2:0.33333333333333,3:0.33333333333333)"
            ":0.33333333333333);"
        )
        assert self.tree().newick() == s

    def test_as_newick_default(self):
        s = (
            "(n0:0.66666666666666663,(n1:0.33333333333333331,"
            "n2:0.33333333333333331):0.33333333333333331);"
        )
        assert self.tree().as_newick() == s

    def test_c_and_py_output_equal(self):
        t = self.tree()
        assert t.as_newick() == t.as_newick(
            node_labels={u: f"n{u}" for u in t.samples()}
        )

    def test_as_newick_precision_3(self):
        s = "(n0:0.667,(n1:0.333,n2:0.333):0.333);"
        assert self.tree().as_newick(precision=3) == s

    def test_newick_precision_3(self):
        s = "(1:0.667,(2:0.333,3:0.333):0.333);"
        assert self.tree().newick(precision=3) == s

    def test_as_newick_precision_3_labels(self):
        node_labels = {u: f"n{u}" * 3 for u in self.tree().nodes()}
        s = "(n0n0n0:0.667,(n1n1n1:0.333,n2n2n2:0.333)n3n3n3:0.333)n4n4n4;"
        assert self.tree().as_newick(precision=3, node_labels=node_labels) == s


class TestLargeBranchLengths:
    # 2000000000.00┊   4   ┊
    #              ┊ ┏━┻┓  ┊
    # 1000000000.00┊ ┃  3  ┊
    #              ┊ ┃ ┏┻┓ ┊
    # 0.00         ┊ 0 1 2 ┊
    #              0       1
    @tests.cached_example
    def tree(self):
        return tskit.Tree.generate_balanced(3, branch_length=1e9)

    def test_newick_default(self):
        s = (
            "(1:2000000000.00000000000000,(2:1000000000.00000000000000,"
            "3:1000000000.00000000000000):1000000000.00000000000000);"
        )
        assert self.tree().newick() == s

    def test_as_newick_default(self):
        s = "(n0:2000000000,(n1:1000000000,n2:1000000000):1000000000);"
        assert self.tree().as_newick() == s

    def test_newick_precision_3(self):
        s = "(1:2000000000.000,(2:1000000000.000,3:1000000000.000):1000000000.000);"
        assert self.tree().newick(precision=3) == s

    def test_as_newick_precision_3(self):
        s = "(n0:2000000000.000,(n1:1000000000.000,n2:1000000000.000):1000000000.000);"
        assert self.tree().as_newick(precision=3) == s


class TestInternalSampleExample:
    #   4
    # ┏━┻┓
    # ┃ *3*
    # ┃ ┏┻┓
    # 0 1 2
    # Leaves are samples but 3 is also a sample.
    @tests.cached_example
    def tree(self):
        tables = tskit.Tree.generate_balanced(3).tree_sequence.dump_tables()
        flags = tables.nodes.flags
        flags[3] = 1
        tables.nodes.flags = flags
        return tables.tree_sequence().first()

    def test_newick_default(self):
        # Old newick method doesn't do anything with internal sample
        s = (
            "(1:2.00000000000000,(2:1.00000000000000,3:1.00000000000000)"
            ":1.00000000000000);"
        )
        assert self.tree().newick() == s

    def test_as_newick_default(self):
        # Samples are labelled by default, not leaves.
        s = "(n0:2,(n1:1,n2:1)n3:1);"
        assert self.tree().as_newick() == s

    def test_dendropy_parsing(self):
        dpy_tree = dendropy.Tree.get(
            data=self.tree().as_newick(),
            schema="newick",
            suppress_internal_node_taxa=False,
            rooting="default-rooted",
        )
        # Just check that we can correctly parse out the internal sample.
        # More exhaustive testing of properties is done elsewhere.
        n3 = dpy_tree.find_node_with_taxon_label("n3")
        n1 = dpy_tree.find_node_with_taxon_label("n1")
        assert n1.parent_node is n3
        n2 = dpy_tree.find_node_with_taxon_label("n2")
        assert n2.parent_node is n3


class TestAncientSampleExample:
    #     8
    #  ┏━━┻━┓
    #  5    7
    # ┏┻┓ ┏━┻┓
    # 0 1 ┃  6
    #     ┃ ┏┻┓
    #     2 3 4
    @tests.cached_example
    def tree(self):
        tables = tskit.Tree.generate_balanced(5).tree_sequence.dump_tables()
        time = tables.nodes.time
        time[0] = 1
        time[1] = 1
        time[5] = 2
        tables.nodes.time = time
        tables.sort()
        return tables.tree_sequence().first()

    def test_as_newick(self):
        s = "((n0:1,n1:1):1,(n2:2,(n3:1,n4:1):1):1);"
        assert self.tree().as_newick() == s

    def test_newick(self):
        s = "((1:1,2:1):1,(3:2,(4:1,5:1):1):1);"
        assert self.tree().newick(precision=0) == s


class TestNonSampleLeafExample:
    #   4
    # ┏━┻┓
    # ┃  3
    # ┃ ┏┻┓
    # |0|1 2
    # Leaf 0 is *not* a sample
    @tests.cached_example
    def tree(self):
        tables = tskit.Tree.generate_balanced(3).tree_sequence.dump_tables()
        flags = tables.nodes.flags
        flags[0] = 0
        tables.nodes.flags = flags
        return tables.tree_sequence().first()

    def test_newick(self):
        # newick method doesn't think about samples at all.
        s = "(1:2,(2:1,3:1):1);"
        assert self.tree().newick(precision=0) == s

    def test_as_newick_default(self):
        # We don't label node 0 even though it's a leaf.
        s = "(:2,(n1:1,n2:1):1);"
        assert self.tree().as_newick() == s

    def test_dendropy_parsing(self):
        # This odd topology parses OK with dendropy
        dpy_tree = dendropy.Tree.get(
            data=self.tree().as_newick(),
            schema="newick",
            suppress_internal_node_taxa=False,
            rooting="default-rooted",
        )
        n1 = dpy_tree.find_node_with_taxon_label("n1")
        assert n1 is not None
        n2 = dpy_tree.find_node_with_taxon_label("n2")
        assert n2 is not None
        leaves = dpy_tree.leaf_nodes()
        assert len(leaves) == 3
        leaves = set(leaves)
        leaves.remove(n1)
        leaves.remove(n2)
        n0 = leaves.pop()
        assert n0.taxon is None

    def test_newick_lib_parsing(self):
        newick_tree = newick.loads(self.tree().as_newick())[0]
        leaf_names = newick_tree.get_leaf_names()
        assert len(leaf_names) == 3
        assert "n1" in leaf_names
        assert "n2" in leaf_names
        assert None in leaf_names


class TestNonBinaryExample:
    # 2.00┊        12         ┊
    #     ┊   ┏━━━━━╋━━━━━┓   ┊
    # 1.00┊   9    10    11   ┊
    #     ┊ ┏━╋━┓ ┏━╋━┓ ┏━╋━┓ ┊
    # 0.00┊ 0 1 2 3 4 5 6 7 8 ┊
    #     0                   1
    @tests.cached_example
    def tree(self):
        return tskit.Tree.generate_balanced(9, arity=3)

    def test_as_newick(self):
        s = "((n0:1,n1:1,n2:1):1,(n3:1,n4:1,n5:1):1,(n6:1,n7:1,n8:1):1);"
        assert self.tree().as_newick() == s

    def test_newick(self):
        s = "((1:1,2:1,3:1):1,(4:1,5:1,6:1):1,(7:1,8:1,9:1):1);"
        assert self.tree().newick(precision=0) == s


class TestMultiRootExample:
    #
    # 1.00┊   9    10    11   ┊
    #     ┊ ┏━╋━┓ ┏━╋━┓ ┏━╋━┓ ┊
    # 0.00┊ 0 1 2 3 4 5 6 7 8 ┊
    #     0                   1
    @tests.cached_example
    def tree(self):
        tables = tskit.Tree.generate_balanced(9, arity=3).tree_sequence.dump_tables()
        edges = tables.edges.copy()
        tables.edges.clear()
        for edge in edges:
            if edge.parent != 12:
                tables.edges.append(edge)
        return tables.tree_sequence().first()

    def test_as_newick_fails(self):
        with pytest.raises(ValueError, match="single root"):
            self.tree().as_newick()

    def test_newick_fails(self):
        with pytest.raises(ValueError, match="single root"):
            self.tree().newick()

    def test_as_newick_per_root(self):
        t = self.tree()
        assert t.as_newick(root=9) == "(n0:1,n1:1,n2:1);"
        assert t.as_newick(root=10) == "(n3:1,n4:1,n5:1);"
        assert t.as_newick(root=11) == "(n6:1,n7:1,n8:1);"

    def test_newick_per_root(self):
        t = self.tree()
        assert t.newick(root=9, precision=0) == "(1:1,2:1,3:1);"
        assert t.newick(root=10, precision=0) == "(4:1,5:1,6:1);"
        assert t.newick(root=11, precision=0) == "(7:1,8:1,9:1);"


class TestLineTree:
    # 3.00┊ 3 ┊
    #     ┊ ┃ ┊
    # 2.00┊ 2 ┊
    #     ┊ ┃ ┊
    # 1.00┊ 1 ┊
    #     ┊ ┃ ┊
    # 0.00┊ 0 ┊
    #     0   1

    @tests.cached_example
    def tree(self):
        tables = tskit.TableCollection(1.0)
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        for j in range(3):
            tables.nodes.add_row(flags=0, time=j + 1)
            tables.edges.add_row(left=0, right=1, parent=j + 1, child=j)
        tables.sort()
        return tables.tree_sequence().first()

    def test_newick(self):
        s = "(((1:1.00000000000000):1.00000000000000):1.00000000000000);"
        assert s == self.tree().newick()

    def test_as_newick(self):
        s = "(((n0:1):1):1);"
        assert s == self.tree().as_newick()

    def test_dendropy_parsing(self):
        dpy_tree = dendropy.Tree.get(
            data=self.tree().as_newick(),
            schema="newick",
            rooting="default-rooted",
        )
        n0 = dpy_tree.find_node_with_taxon_label("n0")
        assert n0 is not None
        assert n0.edge_length == 1


class TestEmptyTree:
    # The empty tree sequence has no nodes and so there's zero roots.
    # This gets caught by the "has_single_root" error check, which is
    # probably not right (we should just return the empty string).
    # It's not an important corner case though, so probably not worth
    # worrying about.
    def tree(self):
        tables = tskit.TableCollection(1.0)
        return tables.tree_sequence().first()

    def test_newick(self):
        with pytest.raises(ValueError, match="single root"):
            self.tree().newick()

    def test_as_newick(self):
        with pytest.raises(ValueError, match="single root"):
            self.tree().as_newick()

    def test_as_nexus(self):
        with pytest.raises(ValueError, match="single root"):
            self.tree().tree_sequence.as_nexus()


class TestSingleNodeTree:
    def tree(self):
        tables = tskit.TableCollection(1.0)
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        return tables.tree_sequence().first()

    def test_newick(self):
        assert self.tree().newick() == "1;"

    def test_as_newick(self):
        assert self.tree().as_newick() == "n0;"

    def test_as_newick_labels(self):
        assert self.tree().as_newick(node_labels={0: "ABCDE"}) == "ABCDE;"


class TestIntegerTreeSequence:
    # 3.00┊   5   ┊       ┊
    #     ┊ ┏━┻┓  ┊       ┊
    # 2.00┊ ┃  4  ┊   4   ┊
    #     ┊ ┃ ┏┻┓ ┊  ┏┻━┓ ┊
    # 1.00┊ ┃ ┃ ┃ ┊  3  ┃ ┊
    #     ┊ ┃ ┃ ┃ ┊ ┏┻┓ ┃ ┊
    # 0.00┊ 0 1 2 ┊ 0 2 1 ┊
    #     0       2      10
    @tests.cached_example
    def ts(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       0           1
        4       0           2
        5       0           3
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        2.0     10      3       0
        2.0     10      3       2
        0.0     10      4       1
        0.0     2.0     4       2
        2.0     10      4       3
        0.0     2.0     5       0
        0.0     2.0     5       4
        """
        )
        return tskit.load_text(nodes=nodes, edges=edges, strict=False)

    def test_nexus_defaults(self):
        ts = self.ts()
        expected = textwrap.dedent(
            """\
            #NEXUS
            BEGIN TAXA;
              DIMENSIONS NTAX=3;
              TAXLABELS n0 n1 n2;
            END;
            BEGIN TREES;
              TREE t0^2 = [&R] (n0:3,(n1:2,n2:2):1);
              TREE t2^10 = [&R] (n1:2,(n0:1,n2:1):1);
            END;
            """
        )
        assert ts.as_nexus() == expected

    def test_nexus_precision_2(self):
        ts = self.ts()
        expected = textwrap.dedent(
            """\
            #NEXUS
            BEGIN TAXA;
              DIMENSIONS NTAX=3;
              TAXLABELS n0 n1 n2;
            END;
            BEGIN TREES;
              TREE t0.00^2.00 = [&R] (n0:3.00,(n1:2.00,n2:2.00):1.00);
              TREE t2.00^10.00 = [&R] (n1:2.00,(n0:1.00,n2:1.00):1.00);
            END;
            """
        )
        assert ts.as_nexus(precision=2) == expected

    @pytest.mark.parametrize("precision", [None, 0, 1, 3])
    def test_file_version_identical(self, precision):
        ts = self.ts()
        out = io.StringIO()
        ts.write_nexus(out, precision=precision)
        assert out.getvalue() == ts.as_nexus(precision=precision)


class TestFloatTimeTreeSequence:
    # 3.25┊   5   ┊       ┊
    #     ┊ ┏━┻┓  ┊       ┊
    # 2.00┊ ┃  4  ┊   4   ┊
    #     ┊ ┃ ┏┻┓ ┊  ┏┻━┓ ┊
    # 1.00┊ ┃ ┃ ┃ ┊  3  ┃ ┊
    #     ┊ ┃ ┃ ┃ ┊ ┏┻┓ ┃ ┊
    # 0.00┊ 0 1 2 ┊ 0 2 1 ┊
    #     0       2      10
    @tests.cached_example
    def ts(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       0           1
        4       0           2
        5       0           3.25
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        2.0     10      3       0
        2.0     10      3       2
        0.0     10      4       1
        0.0     2.0     4       2
        2.0     10      4       3
        0.0     2.0     5       0
        0.0     2.0     5       4
        """
        )
        return tskit.load_text(nodes=nodes, edges=edges, strict=False)

    def test_nexus_defaults(self):
        ts = self.ts()
        expected = textwrap.dedent(
            """\
            #NEXUS
            BEGIN TAXA;
              DIMENSIONS NTAX=3;
              TAXLABELS n0 n1 n2;
            END;
            BEGIN TREES;
              TREE t0^2 = [&R] (n0:3.25000000000000000,(n1:2.00000000000000000,n2:2.00000000000000000):1.25000000000000000);
              TREE t2^10 = [&R] (n1:2.00000000000000000,(n0:1.00000000000000000,n2:1.00000000000000000):1.00000000000000000);
            END;
            """  # noqa: B950
        )
        assert ts.as_nexus() == expected

    def test_nexus_precision_2(self):
        ts = self.ts()
        expected = textwrap.dedent(
            """\
            #NEXUS
            BEGIN TAXA;
              DIMENSIONS NTAX=3;
              TAXLABELS n0 n1 n2;
            END;
            BEGIN TREES;
              TREE t0.00^2.00 = [&R] (n0:3.25,(n1:2.00,n2:2.00):1.25);
              TREE t2.00^10.00 = [&R] (n1:2.00,(n0:1.00,n2:1.00):1.00);
            END;
            """
        )
        assert ts.as_nexus(precision=2) == expected


class TestFloatPositionTreeSequence:
    # 3.00┊   5   ┊       ┊
    #     ┊ ┏━┻┓  ┊       ┊
    # 2.00┊ ┃  4  ┊   4   ┊
    #     ┊ ┃ ┏┻┓ ┊  ┏┻━┓ ┊
    # 1.00┊ ┃ ┃ ┃ ┊  3  ┃ ┊
    #     ┊ ┃ ┃ ┃ ┊ ┏┻┓ ┃ ┊
    # 0.00┊ 0 1 2 ┊ 0 2 1 ┊
    #     0      2.5      10
    @tests.cached_example
    def ts(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       0           1
        4       0           2
        5       0           3
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        2.5     10      3       0
        2.5     10      3       2
        0.0     10      4       1
        0.0     2.5     4       2
        2.5     10      4       3
        0.0     2.5     5       0
        0.0     2.5     5       4
        """
        )
        return tskit.load_text(nodes=nodes, edges=edges, strict=False)

    def test_nexus_defaults(self):
        ts = self.ts()
        expected = textwrap.dedent(
            """\
            #NEXUS
            BEGIN TAXA;
              DIMENSIONS NTAX=3;
              TAXLABELS n0 n1 n2;
            END;
            BEGIN TREES;
              TREE t0.00000000000000000^2.50000000000000000 = [&R] (n0:3,(n1:2,n2:2):1);
              TREE t2.50000000000000000^10.00000000000000000 = [&R] (n1:2,(n0:1,n2:1):1);
            END;
            """  # noqa: B950
        )
        assert ts.as_nexus() == expected

    def test_nexus_precision_2(self):
        ts = self.ts()
        expected = textwrap.dedent(
            """\
            #NEXUS
            BEGIN TAXA;
              DIMENSIONS NTAX=3;
              TAXLABELS n0 n1 n2;
            END;
            BEGIN TREES;
              TREE t0.00^2.50 = [&R] (n0:3.00,(n1:2.00,n2:2.00):1.00);
              TREE t2.50^10.00 = [&R] (n1:2.00,(n0:1.00,n2:1.00):1.00);
            END;
            """
        )
        assert ts.as_nexus(precision=2) == expected


def test_newick_buffer_too_small_bug():
    nodes = io.StringIO(
        """\
    id  is_sample   population individual time
    0       1       0       -1      0.00000000000000
    1       1       0       -1      0.00000000000000
    2       1       0       -1      0.00000000000000
    3       1       0       -1      0.00000000000000
    4       0       0       -1      0.21204940078588
    5       0       0       -1      0.38445004304611
    6       0       0       -1      0.83130278081275
    """
    )
    edges = io.StringIO(
        """\
    id      left            right           parent  child
    0       0.00000000      1.00000000      4       0
    1       0.00000000      1.00000000      4       2
    2       0.00000000      1.00000000      5       1
    3       0.00000000      1.00000000      5       3
    4       0.00000000      1.00000000      6       4
    5       0.00000000      1.00000000      6       5
    """
    )
    ts = tskit.load_text(nodes, edges, sequence_length=1, strict=False)
    tree = ts.first()
    for precision in range(18):
        newick_c = tree.newick(precision=precision)
        node_labels = {u: str(u + 1) for u in ts.samples()}
        newick_py = tree.newick(precision=precision, node_labels=node_labels)
        assert newick_c == newick_py


class TestWrapText:
    def test_even_split(self):
        example = "ABCDEFGH"
        result = list(tskit.text_formats.wrap_text(example, 4))
        assert result == ["ABCD", "EFGH"]

    def test_non_even_split(self):
        example = "ABCDEFGH"
        result = list(tskit.text_formats.wrap_text(example, 3))
        assert result == ["ABC", "DEF", "GH"]

    def test_width_one(self):
        example = "ABCDEFGH"
        result = list(tskit.text_formats.wrap_text(example, 1))
        assert result == ["A", "B", "C", "D", "E", "F", "G", "H"]

    def test_width_full_length(self):
        example = "ABCDEFGH"
        result = list(tskit.text_formats.wrap_text(example, 8))
        assert result == ["ABCDEFGH"]

    def test_width_more_than_length(self):
        example = "ABCDEFGH"
        result = list(tskit.text_formats.wrap_text(example, 100))
        assert result == ["ABCDEFGH"]

    def test_width_0(self):
        example = "ABCDEFGH"
        result = list(tskit.text_formats.wrap_text(example, 0))
        assert result == ["ABCDEFGH"]

    @pytest.mark.parametrize("width", [-1, -2, -8, -100])
    def test_width_negative(self, width):
        # Just documenting that the current implementation works for negative
        # values fine.
        example = "ABCDEFGH"
        result = list(tskit.text_formats.wrap_text(example, width))
        assert result == ["ABCDEFGH"]


class TestFastaLineLength:
    """
    Tests if the fasta file produced has the correct line lengths for
    default, custom, and no-wrapping options.
    """

    def verify_line_length(self, length, wrap_width=60):
        # set up data
        ts = alignment_example(length)
        output = io.StringIO()
        ts.write_fasta(output, wrap_width=wrap_width)
        output.seek(0)

        # check if length perfectly divisible by wrap_width or not and thus
        # expected line lengths
        no_hanging_line = True
        if wrap_width == 0:
            lines_expect = 1
            # for easier code in testing function, redefine wrap_width as
            # full length, ok as called write already
            wrap_width = length
        elif length % wrap_width == 0:
            lines_expect = length // wrap_width
        else:
            lines_expect = length // wrap_width + 1
            extra_line_length = length % wrap_width
            no_hanging_line = False

        seq_line_counter = 0
        id_lines = 0
        for line in output:
            # testing correct characters per sequence line
            if line[0] != ">":
                seq_line_counter += 1
                line_chars = len(line.strip("\n"))
                # test full default width lines
                if seq_line_counter < lines_expect:
                    assert wrap_width == line_chars
                elif no_hanging_line:
                    assert wrap_width == line_chars
                # test extra line if not perfectly divided by wrap_width
                else:
                    assert extra_line_length == line_chars
            # testing correct number of lines per sequence and correct num sequences
            else:
                id_lines += 1
                if seq_line_counter > 0:
                    assert lines_expect == seq_line_counter
                    seq_line_counter = 0
        assert id_lines == ts.num_samples

    def test_wrap_length_default_easy(self):
        # default wrap width (60) perfectly divides sequence length
        self.verify_line_length(length=300)

    def test_wrap_length_default_harder(self):
        # default wrap_width imperfectly divides sequence length
        self.verify_line_length(length=280)

    def test_wrap_length_custom_easy(self):
        # custom wrap_width, perfectly divides
        self.verify_line_length(length=100, wrap_width=20)

    def test_wrap_length_custom_harder(self):
        # custom wrap_width, imperfectly divides
        self.verify_line_length(length=100, wrap_width=30)

    def test_wrap_length_no_wrap(self):
        # no wrapping set by wrap_width = 0
        self.verify_line_length(length=100, wrap_width=0)

    def test_negative_wrap(self):
        ts = alignment_example(10)
        with pytest.raises(ValueError, match="non-negative integer"):
            ts.as_fasta(wrap_width=-1)

    def test_floating_wrap(self):
        ts = alignment_example(10)
        with pytest.raises(ValueError):
            ts.as_fasta(wrap_width=1.1)

    def test_numpy_wrap(self):
        ts = alignment_example(10)
        x1 = ts.as_fasta(wrap_width=4)
        x2 = ts.as_fasta(wrap_width=np.array([4.0])[0])
        assert x1 == x2


class TestFileTextOutputEqual:
    @tests.cached_example
    def ts(self):
        return alignment_example(20)

    def test_fasta_defaults(self):
        ts = self.ts()
        buff = io.StringIO()
        ts.write_fasta(buff)
        assert buff.getvalue() == ts.as_fasta()

    def test_fasta_wrap_width(self):
        ts = self.ts()
        buff = io.StringIO()
        ts.write_fasta(buff, wrap_width=4)
        assert buff.getvalue() == ts.as_fasta(wrap_width=4)

    def test_nexus_defaults(self):
        ts = self.ts()
        buff = io.StringIO()
        ts.write_nexus(buff)
        assert buff.getvalue() == ts.as_nexus()

    def test_nexus_precision(self):
        ts = self.ts()
        buff = io.StringIO()
        ts.write_nexus(buff, precision=2)
        assert buff.getvalue() == ts.as_nexus(precision=2)


class TestFlexibleFileArgFasta:
    @tests.cached_example
    def ts(self):
        return alignment_example(20)

    def test_pathlib(self, tmp_path):
        path = tmp_path / "file.fa"
        ts = self.ts()
        ts.write_fasta(path)
        with open(path) as f:
            assert f.read() == ts.as_fasta()

    def test_path_str(self, tmp_path):
        path = str(tmp_path / "file.fa")
        ts = self.ts()
        ts.write_fasta(path)
        with open(path) as f:
            assert f.read() == ts.as_fasta()

    def test_fileobj(self, tmp_path):
        path = tmp_path / "file.fa"
        ts = self.ts()
        with open(path, "w") as f:
            ts.write_fasta(f)
        with open(path) as f:
            assert f.read() == ts.as_fasta()


class TestFlexibleFileArgNexus:
    @tests.cached_example
    def ts(self):
        return alignment_example(20)

    def test_pathlib(self, tmp_path):
        path = tmp_path / "file.nex"
        ts = self.ts()
        ts.write_nexus(path)
        with open(path) as f:
            assert f.read() == ts.as_nexus()

    def test_path_str(self, tmp_path):
        path = str(tmp_path / "file.nex")
        ts = self.ts()
        ts.write_nexus(path)
        with open(path) as f:
            assert f.read() == ts.as_nexus()

    def test_fileobj(self, tmp_path):
        path = tmp_path / "file.nex"
        ts = self.ts()
        with open(path, "w") as f:
            ts.write_nexus(f)
        with open(path) as f:
            assert f.read() == ts.as_nexus()


def get_alignment_map(ts, reference_sequence=None):
    alignments = ts.alignments(reference_sequence=reference_sequence)
    return {f"n{u}": alignment for u, alignment in zip(ts.samples(), alignments)}


class TestFastaBioPythonRoundTrip:
    """
    Tests that output from our code is read in by available software packages
    Here test for compatability with biopython processing - Bio.SeqIO
    """

    def verify(self, ts, wrap_width=60, reference_sequence=None):
        text = ts.as_fasta(wrap_width=wrap_width, reference_sequence=reference_sequence)
        bio_map = {
            k: v.seq
            for k, v in SeqIO.to_dict(SeqIO.parse(io.StringIO(text), "fasta")).items()
        }
        assert bio_map == get_alignment_map(ts, reference_sequence)

    def test_equal_lines(self):
        # sequence length perfectly divisible by wrap_width
        ts = alignment_example(300)
        self.verify(ts)

    def test_unequal_lines(self):
        # sequence length not perfectly divisible by wrap_width
        ts = alignment_example(280)
        self.verify(ts)

    def test_unwrapped(self):
        # sequences not wrapped
        ts = alignment_example(300)
        self.verify(ts, wrap_width=0)

    def test_A_reference(self):
        ts = alignment_example(20)
        self.verify(ts, reference_sequence="A" * 20)

    @pytest.mark.skip("Missing data in alignments: #1896")
    def test_missing_data(self):
        self.verify(missing_data_example())


class TestFastaDendropyRoundTrip:
    def parse(self, fasta):
        d = dendropy.DnaCharacterMatrix.get(data=fasta, schema="fasta")
        return {str(k.label): str(v) for k, v in d.items()}

    def test_wrapped(self):
        ts = alignment_example(300)
        text = ts.as_fasta()
        alignment_map = self.parse(text)
        assert get_alignment_map(ts) == alignment_map

    def test_unwrapped(self):
        ts = alignment_example(300)
        text = ts.as_fasta(wrap_width=0)
        alignment_map = self.parse(text)
        assert get_alignment_map(ts) == alignment_map

    def test_no_reference(self):
        ts = alignment_example(100, include_reference=False)
        text = ts.as_fasta()
        alignment_map = self.parse(text)
        assert get_alignment_map(ts) == alignment_map

    @pytest.mark.skip("Missing data in alignments: #1896")
    def test_missing_data(self):
        ts = missing_data_example()
        text = ts.as_fasta()
        alignment_map = self.parse(text)
        assert get_alignment_map(ts) == alignment_map


class TestDendropyMissingReference:
    # 2.00┊   4   ┊
    #     ┊ ┏━┻┓  ┊
    # 1.00┊ ┃  3  ┊
    #     ┊ ┃ ┏┻┓ ┊
    # 0.00┊ 0 1 2 ┊
    #     0       10
    #      |     |
    #  pos 2     9
    #  anc A     T

    def ts(self):
        ts = tskit.Tree.generate_balanced(3, span=10).tree_sequence
        tables = ts.dump_tables()
        tables.sites.add_row(2, ancestral_state="A")
        tables.sites.add_row(9, ancestral_state="T")
        tables.mutations.add_row(site=0, node=0, derived_state="G")
        tables.mutations.add_row(site=1, node=3, derived_state="C")
        return tables.tree_sequence()

    def assert_missing_data_encoded(self, d):
        assert d.sequence_size == 10
        assert str(d["n0"][2]) == "G"
        assert str(d["n0"][9]) == "T"
        assert str(d["n1"][2]) == "A"
        assert str(d["n1"][9]) == "C"
        assert str(d["n2"][2]) == "A"
        assert str(d["n2"][9]) == "C"
        for a in d.values():
            for j in range(d.sequence_size):
                if j in [2, 9]:
                    assert (
                        a[j].state_denomination
                        == dendropy.StateAlphabet.FUNDAMENTAL_STATE
                    )
                else:
                    assert (
                        a[j].state_denomination
                        == dendropy.StateAlphabet.AMBIGUOUS_STATE
                    )

    def test_fasta(self):
        ts = self.ts()
        text = ts.as_fasta()
        d = dendropy.DnaCharacterMatrix.get(data=text, schema="fasta")
        self.assert_missing_data_encoded(d)
        assert str(d["n0"][0]) == "N"

    def test_fasta_missing_question(self):
        ts = self.ts()
        text = ts.as_fasta(missing_data_character="?")
        d = dendropy.DnaCharacterMatrix.get(data=text, schema="fasta")
        self.assert_missing_data_encoded(d)
        assert str(d["n0"][0]) == "?"

    def test_nexus(self):
        ts = self.ts()
        text = ts.as_nexus()
        d = dendropy.DnaCharacterMatrix.get(data=text, schema="nexus")
        self.assert_missing_data_encoded(d)
        assert str(d["n0"][0]) == "?"

    def test_nexus_missing_N(self):
        ts = self.ts()
        text = ts.as_nexus(missing_data_character="N")
        d = dendropy.DnaCharacterMatrix.get(data=text, schema="nexus")
        self.assert_missing_data_encoded(d)
        assert str(d["n0"][0]) == "N"


@pytest.mark.skip("Missing data in alignments: #1896")
class TestDendropyMissingData:
    """
    Test that we detect missing data correctly in dendropy under
    various combinations of options.
    """

    # 2.00┊   4     ┊
    #     ┊ ┏━┻┓    ┊
    # 1.00┊ ┃  3    ┊
    #     ┊ ┃ ┏┻┓   ┊
    # 0.00┊ 0 1 2 5 ┊
    #     0        10
    #      |      |
    #  pos 2      9
    #  anc A      T

    def ts(self):
        ts = tskit.Tree.generate_balanced(3, span=10).tree_sequence
        tables = ts.dump_tables()
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        tables.sites.add_row(2, ancestral_state="A")
        tables.sites.add_row(9, ancestral_state="T")
        tables.mutations.add_row(site=0, node=0, derived_state="G")
        tables.mutations.add_row(site=1, node=3, derived_state="C")
        return tables.tree_sequence()

    def assert_missing_data_encoded_A_ref(self, d):
        assert d.sequence_size == 10
        assert str(d["n0"]) == "AAGAAAAAAT"
        assert str(d["n1"]) == "AAAAAAAAAC"
        assert str(d["n2"]) == "AAAAAAAAAC"
        for a in [d["n0"], d["n1"], d["n2"]]:
            assert all(
                a[j].state_denomination == dendropy.StateAlphabet.FUNDAMENTAL_STATE
                for j in range(10)
            )
        # Do we detect that we have an ambiguous state for the missing sample?

        a5 = d["n5"]
        # a5 is missing along the full length of the genome, so all sites are
        # missing.
        assert all(
            a5.state_denomination == dendropy.StateAlphabet.AMBIGUOUS_STATE
            for j in range(10)
        )

    def test_fasta_defaults_A_ref(self):
        ts = self.ts()
        ref = "A" * int(ts.sequence_length)
        text = ts.as_fasta(reference_sequence=ref)
        d = dendropy.DnaCharacterMatrix.get(data=text, schema="fasta")
        self.assert_missing_data_encoded_A_ref(d)

    def test_nexus_defaults_A_ref(self):
        ts = self.ts()
        ref = "A" * int(ts.sequence_length)
        text = ts.as_nexus(reference_sequence=ref, include_trees=False)
        d = dendropy.DnaCharacterMatrix.get(data=text, schema="nexus")
        self.assert_missing_data_encoded_A_ref(d)


--- ../../tskit/python/tests/test_tree_stats.py ---


"""
Test cases for generalized statistic computation.
"""
import collections
import contextlib
import functools
import io
import itertools
import random

import msprime
import numpy as np
import numpy.testing as nt
import pytest

import tests.test_wright_fisher as wf
import tests.tsutil as tsutil
import tskit
import tskit.exceptions as exceptions

np.random.seed(5)


def cached_np(func):
    """
    Decorator to speed up functions that take numpy arrays as positional
    arguments that get called a lot with the same arguments.

    # See https://github.com/tskit-dev/tskit/issues/1856 for more info
    """
    cache = {}

    def f(*args):
        nonlocal cache
        key = tuple(x.tobytes() for x in args)
        if key not in cache:
            cache[key] = func(*args)
        return cache[key]

    return f


def subset_combos(*args, p=0.5, min_tests=3):
    # We have too many tests, combinatorially; so we will run a random subset
    # of them, using this function, below. If we don't set a seed, a different
    # random set is run each time. Ensures that at least min_tests are run.
    # Uncomment this line to run all tests (takes about an hour):
    p = 1.0
    num_tests = 0
    skipped_tests = []
    # total_tests = 0
    for x in itertools.product(*args):
        # total_tests = total_tests + 1
        if np.random.uniform() < p:
            num_tests += 1
            yield x
        elif len(skipped_tests) < min_tests:
            skipped_tests.append(x)
        elif np.random.uniform() < 0.1:
            skipped_tests[np.random.randint(min_tests)] = x
    while num_tests < min_tests:
        yield skipped_tests.pop()
        num_tests += 1
    # print("tests", num_tests)
    assert num_tests >= min_tests


def path_length(tr, x, y):
    L = 0
    if x >= 0 and y >= 0:
        mrca = tr.mrca(x, y)
    else:
        mrca = -1
    for u in x, y:
        while u != mrca:
            L += tr.branch_length(u)
            u = tr.parent(u)
    return L


@contextlib.contextmanager
def suppress_division_by_zero_warning():
    with np.errstate(invalid="ignore", divide="ignore"):
        yield


##############################
# Branch general stat algorithms
##############################


def windowed_tree_stat(ts, stat, windows, span_normalise=True):
    shape = list(stat.shape)
    shape[0] = len(windows) - 1
    A = np.zeros(shape)

    tree_breakpoints = np.array(list(ts.breakpoints()))
    tree_index = 0
    for j in range(len(windows) - 1):
        w_left = windows[j]
        w_right = windows[j + 1]
        while True:
            t_left = tree_breakpoints[tree_index]
            t_right = tree_breakpoints[tree_index + 1]
            left = max(t_left, w_left)
            right = min(t_right, w_right)
            weight = max(0.0, (right - left) / (t_right - t_left))
            A[j] += stat[tree_index] * weight
            assert left != right
            if t_right <= w_right:
                tree_index += 1
                # TODO This is inelegant - should include this in the case below
                if t_right == w_right:
                    break
            else:
                break
    if span_normalise:
        # re-normalize by window lengths
        window_lengths = np.diff(windows)
        for j in range(len(windows) - 1):
            A[j] /= window_lengths[j]
    return A


def naive_branch_general_stat(
    ts, w, f, windows=None, polarised=False, span_normalise=True
):
    # NOTE: does not behave correctly for unpolarised stats
    # with non-ancestral material.
    if windows is None:
        windows = [0.0, ts.sequence_length]
    n, k = w.shape
    # hack to determine m
    m = len(f(w[0]))
    total = np.sum(w, axis=0)

    sigma = np.zeros((ts.num_trees, m))
    for tree in ts.trees():
        x = np.zeros((ts.num_nodes, k))
        x[ts.samples()] = w
        for u in tree.nodes(order="postorder"):
            for v in tree.children(u):
                x[u] += x[v]
        if polarised:
            s = sum(tree.branch_length(u) * f(x[u]) for u in tree.nodes())
        else:
            s = sum(
                tree.branch_length(u) * (f(x[u]) + f(total - x[u]))
                for u in tree.nodes()
            )
        sigma[tree.index] = s * tree.span
    if isinstance(windows, str) and windows == "trees":
        # need to average across the windows
        if span_normalise:
            for j, tree in enumerate(ts.trees()):
                sigma[j] /= tree.span
        return sigma
    else:
        return windowed_tree_stat(ts, sigma, windows, span_normalise=span_normalise)


def branch_general_stat(
    ts, sample_weights, summary_func, windows=None, polarised=False, span_normalise=True
):
    """
    Efficient implementation of the algorithm used as the basis for the
    underlying C version.
    """
    n, state_dim = sample_weights.shape
    windows = ts.parse_windows(windows)
    num_windows = windows.shape[0] - 1

    # Determine result_dim
    result_dim = len(summary_func(sample_weights[0]))
    result = np.zeros((num_windows, result_dim))
    state = np.zeros((ts.num_nodes, state_dim))
    state[ts.samples()] = sample_weights
    total_weight = np.sum(sample_weights, axis=0)

    time = ts.tables.nodes.time
    parent = np.zeros(ts.num_nodes, dtype=np.int32) - 1
    branch_length = np.zeros(ts.num_nodes)
    # The value of summary_func(u) for every node.
    summary = np.zeros((ts.num_nodes, result_dim))
    # The result for the current tree *not* weighted by span.
    running_sum = np.zeros(result_dim)

    def polarised_summary(u):
        s = summary_func(state[u])
        if not polarised:
            s += summary_func(total_weight - state[u])
        return s

    for u in range(ts.num_nodes):
        summary[u] = polarised_summary(u)

    window_index = 0
    for (t_left, t_right), edges_out, edges_in in ts.edge_diffs():
        for edge in edges_out:
            u = edge.child
            running_sum -= branch_length[u] * summary[u]
            u = edge.parent
            while u != -1:
                running_sum -= branch_length[u] * summary[u]
                state[u] -= state[edge.child]
                summary[u] = polarised_summary(u)
                running_sum += branch_length[u] * summary[u]
                u = parent[u]
            parent[edge.child] = -1
            branch_length[edge.child] = 0

        for edge in edges_in:
            parent[edge.child] = edge.parent
            branch_length[edge.child] = time[edge.parent] - time[edge.child]
            u = edge.child
            running_sum += branch_length[u] * summary[u]
            u = edge.parent
            while u != -1:
                running_sum -= branch_length[u] * summary[u]
                state[u] += state[edge.child]
                summary[u] = polarised_summary(u)
                running_sum += branch_length[u] * summary[u]
                u = parent[u]

        # Update the windows
        assert window_index < num_windows
        while windows[window_index] < t_right:
            w_left = windows[window_index]
            w_right = windows[window_index + 1]
            left = max(t_left, w_left)
            right = min(t_right, w_right)
            span = right - left
            assert span > 0
            result[window_index] += running_sum * span
            if w_right <= t_right:
                window_index += 1
            else:
                # This interval crosses a tree boundary, so we update it again in the
                # for the next tree
                break

    # print("window_index:", window_index, windows.shape)
    assert window_index == windows.shape[0] - 1
    if span_normalise:
        for j in range(num_windows):
            result[j] /= windows[j + 1] - windows[j]
    return result


##############################
# Site general stat algorithms
##############################


def windowed_sitewise_stat(ts, sigma, windows, span_normalise=True):
    M = sigma.shape[1]
    A = np.zeros((len(windows) - 1, M))
    window = 0
    for site in ts.sites():
        while windows[window + 1] <= site.position:
            window += 1
        assert windows[window] <= site.position < windows[window + 1]
        A[window] += sigma[site.id]
    if span_normalise:
        diff = np.zeros((A.shape[0], 1))
        diff[:, 0] = np.diff(windows).T
        A /= diff
    return A


def naive_site_general_stat(
    ts, W, f, windows=None, polarised=False, span_normalise=True
):
    n, K = W.shape
    # Hack to determine M
    M = len(f(W[0]))
    sigma = np.zeros((ts.num_sites, M))
    for tree in ts.trees():
        X = np.zeros((ts.num_nodes, K))
        X[ts.samples()] = W
        for u in tree.nodes(order="postorder"):
            for v in tree.children(u):
                X[u] += X[v]
        for site in tree.sites():
            state_map = collections.defaultdict(functools.partial(np.zeros, K))
            state_map[site.ancestral_state] = sum(X[root] for root in tree.roots)
            for mutation in site.mutations:
                state_map[mutation.derived_state] += X[mutation.node]
                if mutation.parent != tskit.NULL:
                    parent = site.mutations[mutation.parent - site.mutations[0].id]
                    state_map[parent.derived_state] -= X[mutation.node]
                else:
                    state_map[site.ancestral_state] -= X[mutation.node]
            if polarised:
                del state_map[site.ancestral_state]
            sigma[site.id] += sum(map(f, state_map.values()))
    return windowed_sitewise_stat(
        ts, sigma, ts.parse_windows(windows), span_normalise=span_normalise
    )


def site_general_stat(
    ts, sample_weights, summary_func, windows=None, polarised=False, span_normalise=True
):
    """
    Problem: 'sites' is different that the other windowing options
    because if we output by site we don't want to normalize by length of the window.
    Solution: we pass an argument "normalize", to the windowing function.
    """
    windows = ts.parse_windows(windows)
    num_windows = windows.shape[0] - 1
    n, state_dim = sample_weights.shape
    # Determine result_dim
    (result_dim,) = summary_func(sample_weights[0]).shape
    result = np.zeros((num_windows, result_dim))
    state = np.zeros((ts.num_nodes, state_dim))
    state[ts.samples()] = sample_weights
    total_weight = np.sum(sample_weights, axis=0)

    site_index = 0
    mutation_index = 0
    window_index = 0
    sites = ts.tables.sites
    mutations = ts.tables.mutations
    parent = np.zeros(ts.num_nodes, dtype=np.int32) - 1
    for (left, right), edges_out, edges_in in ts.edge_diffs():
        for edge in edges_out:
            u = edge.parent
            while u != -1:
                state[u] -= state[edge.child]
                u = parent[u]
            parent[edge.child] = -1
        for edge in edges_in:
            parent[edge.child] = edge.parent
            u = edge.parent
            while u != -1:
                state[u] += state[edge.child]
                u = parent[u]
        while site_index < len(sites) and sites.position[site_index] < right:
            assert left <= sites.position[site_index]
            ancestral_state = sites[site_index].ancestral_state
            allele_state = collections.defaultdict(
                functools.partial(np.zeros, state_dim)
            )
            allele_state[ancestral_state][:] = total_weight
            while (
                mutation_index < len(mutations)
                and mutations[mutation_index].site == site_index
            ):
                mutation = mutations[mutation_index]
                allele_state[mutation.derived_state] += state[mutation.node]
                if mutation.parent != -1:
                    parent_allele = mutations[mutation.parent].derived_state
                    allele_state[parent_allele] -= state[mutation.node]
                else:
                    allele_state[ancestral_state] -= state[mutation.node]
                mutation_index += 1
            if polarised:
                del allele_state[ancestral_state]

            pos = sites.position[site_index]
            while windows[window_index + 1] <= pos:
                window_index += 1
            assert windows[window_index] <= pos < windows[window_index + 1]
            site_result = result[window_index]

            for _allele, value in allele_state.items():
                site_result += summary_func(value)
            site_index += 1
    if span_normalise:
        for j in range(num_windows):
            span = windows[j + 1] - windows[j]
            result[j] /= span
    return result


##############################
# Node general stat algorithms
##############################


def naive_node_general_stat(
    ts, W, f, windows=None, polarised=False, span_normalise=True
):
    windows = ts.parse_windows(windows)
    n, K = W.shape
    M = f(W[0]).shape[0]
    total = np.sum(W, axis=0)
    sigma = np.zeros((ts.num_trees, ts.num_nodes, M))
    for tree in ts.trees():
        X = np.zeros((ts.num_nodes, K))
        X[ts.samples()] = W
        for u in tree.nodes(order="postorder"):
            for v in tree.children(u):
                X[u] += X[v]
        s = np.zeros((ts.num_nodes, M))
        for u in range(ts.num_nodes):
            s[u] = f(X[u])
            if not polarised:
                s[u] += f(total - X[u])
        sigma[tree.index] = s * tree.span
    return windowed_tree_stat(ts, sigma, windows, span_normalise=span_normalise)


def node_general_stat(
    ts, sample_weights, summary_func, windows=None, polarised=False, span_normalise=True
):
    """
    Efficient implementation of the algorithm used as the basis for the
    underlying C version.
    """
    n, state_dim = sample_weights.shape
    windows = ts.parse_windows(windows)
    num_windows = windows.shape[0] - 1
    result_dim = summary_func(sample_weights[0]).shape[0]
    result = np.zeros((num_windows, ts.num_nodes, result_dim))
    state = np.zeros((ts.num_nodes, state_dim))
    state[ts.samples()] = sample_weights
    total_weight = np.sum(sample_weights, axis=0)

    def node_summary(u):
        s = summary_func(state[u])
        if not polarised:
            s += summary_func(total_weight - state[u])
        return s

    window_index = 0
    parent = np.zeros(ts.num_nodes, dtype=np.int32) - 1
    # contains summary_func(state[u]) for each node
    current_values = np.zeros((ts.num_nodes, result_dim))
    for u in range(ts.num_nodes):
        current_values[u] = node_summary(u)
    # contains the location of the last time we updated the output for a node.
    last_update = np.zeros((ts.num_nodes, 1))
    for (t_left, t_right), edges_out, edges_in in ts.edge_diffs():
        for edge in edges_out:
            u = edge.child
            v = edge.parent
            while v != -1:
                result[window_index, v] += (t_left - last_update[v]) * current_values[v]
                last_update[v] = t_left
                state[v] -= state[u]
                current_values[v] = node_summary(v)
                v = parent[v]
            parent[u] = -1

        for edge in edges_in:
            u = edge.child
            v = edge.parent
            parent[u] = v
            while v != -1:
                result[window_index, v] += (t_left - last_update[v]) * current_values[v]
                last_update[v] = t_left
                state[v] += state[u]
                current_values[v] = node_summary(v)
                v = parent[v]

        # Update the windows
        while window_index < num_windows and windows[window_index + 1] <= t_right:
            w_right = windows[window_index + 1]
            # Flush the contribution of all nodes to the current window.
            for u in range(ts.num_nodes):
                result[window_index, u] += (w_right - last_update[u]) * current_values[
                    u
                ]
                last_update[u] = w_right
            window_index += 1

    assert window_index == windows.shape[0] - 1
    if span_normalise:
        for j in range(num_windows):
            result[j] /= windows[j + 1] - windows[j]
    return result


def general_stat(
    ts,
    sample_weights,
    summary_func,
    windows=None,
    polarised=False,
    mode="site",
    span_normalise=True,
):
    """
    General iterface for algorithms above. Directly corresponds to the interface
    for TreeSequence.general_stat.
    """
    method_map = {
        "site": site_general_stat,
        "node": node_general_stat,
        "branch": branch_general_stat,
    }
    return method_map[mode](
        ts,
        sample_weights,
        summary_func,
        windows=windows,
        polarised=polarised,
        span_normalise=span_normalise,
    )


def upper_tri_to_matrix(x):
    """
    Given x, a vector of entries of the upper triangle of a matrix
    in row-major order, including the diagonal, return the corresponding matrix.
    """
    # n^2 + n = 2 u => n = (-1 + sqrt(1 + 8*u))/2
    n = int((np.sqrt(1 + 8 * len(x)) - 1) / 2.0)
    out = np.ones((n, n))
    k = 0
    for i in range(n):
        for j in range(i, n):
            out[i, j] = out[j, i] = x[k]
            k += 1
    return out


##################################
# Test cases
##################################


class StatsTestCase:
    """
    Provides convenience functions.
    """

    def assertListAlmostEqual(self, x, y):
        assert len(x) == len(y)
        for a, b in zip(x, y):
            self.assertAlmostEqual(a, b)

    def assertArrayEqual(self, x, y):
        nt.assert_equal(x, y)

    def assertArrayAlmostEqual(self, x, y, atol=1e-6, rtol=1e-7):
        nt.assert_allclose(x, y, atol=atol, rtol=rtol)

    def identity_f(self, ts):
        return lambda x: x * (x < ts.num_samples)

    def cumsum_f(self, ts):
        return lambda x: np.cumsum(x) * (x < ts.num_samples)

    def sum_f(self, ts, k=1):
        return lambda x: np.array([sum(x) * (sum(x) < 2 * ts.num_samples)] * k)


class TopologyExamplesMixin:
    """
    Defines a set of test cases on different example tree sequence topologies.
    Derived classes need to define a 'verify' function which will perform the
    actual tests.
    """

    def test_single_tree_sequence_length(self):
        ts = msprime.simulate(6, length=10, random_seed=1)
        self.verify(ts)

    def test_single_tree_multiple_roots(self):
        ts = msprime.simulate(8, random_seed=1, end_time=0.5)
        assert ts.first().num_roots > 1
        self.verify(ts)

    def test_many_trees(self):
        ts = msprime.simulate(6, recombination_rate=2, random_seed=1)
        assert ts.num_trees > 2
        self.verify(ts)

    def test_short_sequence_length(self):
        ts = msprime.simulate(6, length=0.5, recombination_rate=2, random_seed=1)
        assert ts.num_trees > 2
        self.verify(ts)

    def test_wright_fisher_unsimplified(self):
        tables = wf.wf_sim(
            4,
            5,
            seed=1,
            deep_history=True,
            initial_generation_samples=False,
            num_loci=5,
        )
        tables.sort()
        ts = tables.tree_sequence()
        self.verify(ts)

    @pytest.mark.slow
    def test_wright_fisher_initial_generation(self):
        tables = wf.wf_sim(
            6, 5, seed=3, deep_history=True, initial_generation_samples=True, num_loci=2
        )
        tables.sort()
        tables.simplify()
        ts = tables.tree_sequence()
        self.verify(ts)

    def test_wright_fisher_initial_generation_no_deep_history(self):
        tables = wf.wf_sim(
            6,
            15,
            seed=202,
            deep_history=False,
            initial_generation_samples=True,
            num_loci=5,
        )
        tables.sort()
        tables.simplify()
        ts = tables.tree_sequence()
        self.verify(ts)

    def test_wright_fisher_unsimplified_multiple_roots(self):
        tables = wf.wf_sim(
            6,
            5,
            seed=1,
            deep_history=False,
            initial_generation_samples=False,
            num_loci=4,
        )
        tables.sort()
        ts = tables.tree_sequence()
        self.verify(ts)

    def test_wright_fisher_simplified(self):
        tables = wf.wf_sim(
            5,
            8,
            seed=1,
            deep_history=True,
            initial_generation_samples=False,
            num_loci=5,
        )
        tables.sort()
        ts = tables.tree_sequence().simplify()
        self.verify(ts)

    def test_wright_fisher_simplified_multiple_roots(self):
        tables = wf.wf_sim(
            6,
            8,
            seed=1,
            deep_history=False,
            initial_generation_samples=False,
            num_loci=3,
        )
        tables.sort()
        ts = tables.tree_sequence().simplify()
        self.verify(ts)

    def test_empty_ts(self):
        tables = tskit.TableCollection(1.0)
        tables.nodes.add_row(1, 0)
        tables.nodes.add_row(1, 0)
        tables.nodes.add_row(1, 0)
        tables.nodes.add_row(1, 0)
        ts = tables.tree_sequence()
        self.verify(ts)

    def test_non_sample_ancestry(self):
        tables = tskit.TableCollection(1.0)
        tables.nodes.add_row(1, 0)
        tables.nodes.add_row(1, 0)
        tables.nodes.add_row(0, 1)
        tables.nodes.add_row(0, 0)  # 3 is a leaf but not a sample.
        # Make sure we have 4 samples for the tests.
        tables.nodes.add_row(1, 1)
        tables.nodes.add_row(1, 1)
        tables.edges.add_row(0, 1, 2, 0)
        tables.edges.add_row(0, 1, 2, 1)
        tables.edges.add_row(0, 1, 4, 3)
        ts = tables.tree_sequence()
        self.verify(ts)


class MutatedTopologyExamplesMixin:
    """
    Defines a set of test cases on different example tree sequence topologies.
    Derived classes need to define a 'verify' function which will perform the
    actual tests.
    """

    def test_single_tree_no_sites(self):
        ts = msprime.simulate(6, random_seed=1)
        assert ts.num_sites == 0
        self.verify(ts)

    def test_ghost_allele(self):
        tables = tskit.TableCollection(1)
        tables.nodes.add_row(flags=1, time=0)
        tables.nodes.add_row(flags=1, time=0)
        tables.nodes.add_row(flags=0, time=1)
        tables.edges.add_row(0, 1, 2, 0)
        tables.edges.add_row(0, 1, 2, 1)
        tables.sites.add_row(position=0.5, ancestral_state="A")
        # Make sure there's 4 samples
        tables.nodes.add_row(flags=1, time=0)
        tables.nodes.add_row(flags=1, time=0)
        # The ghost mutation that's never seen in the genotypes
        tables.mutations.add_row(site=0, node=0, derived_state="T")
        tables.mutations.add_row(site=0, node=0, derived_state="G", parent=0)
        ts = tables.tree_sequence()
        self.verify(ts)

    def test_ghost_allele_all_ancestral(self):
        tables = tskit.TableCollection(1)
        tables.nodes.add_row(flags=1, time=0)
        tables.nodes.add_row(flags=1, time=0)
        tables.nodes.add_row(flags=0, time=1)
        # Make sure there's 4 samples
        tables.nodes.add_row(flags=1, time=0)
        tables.nodes.add_row(flags=1, time=0)
        tables.edges.add_row(0, 1, 2, 0)
        tables.edges.add_row(0, 1, 2, 1)
        tables.sites.add_row(position=0.5, ancestral_state="A")
        tables.mutations.add_row(site=0, node=0, derived_state="T")
        # Mutate back to the ancestral state so that all genotypes are zero
        tables.mutations.add_row(site=0, node=0, derived_state="A", parent=0)
        ts = tables.tree_sequence()
        self.verify(ts)

    def test_non_sample_ancestry(self):
        # 2.00┊       5   ┊
        #     ┊    ┏━━┻━┓ ┊
        # 1.00┊    4    ┃ ┊
        #     ┊ ┏━┳┻┳━┓ ┃ ┊
        # 0.00┊ 0 1 2 3 6 ┊
        #    0.00        1.00
        tables = tskit.TableCollection(1)
        # Four sample nodes
        for j in range(4):
            tables.nodes.add_row(flags=1, time=0)
            tables.edges.add_row(0, 1, 4, j)
        # Their MRCA, 4, joins to older ancestor 5
        tables.nodes.add_row(flags=0, time=1)
        tables.nodes.add_row(flags=0, time=2)
        tables.edges.add_row(0, 1, 5, 4)
        # Which has non-sample leaf at time 0
        tables.nodes.add_row(flags=0, time=0)
        tables.edges.add_row(0, 1, 5, 6)
        # Two sites with mutations. One over the MRCA of the
        # samples so it's fixed at 1 and one over the non sample
        # leaf so that samples are fixed at zero.
        tables.sites.add_row(position=0.25, ancestral_state="0")
        tables.sites.add_row(position=0.5, ancestral_state="0")
        tables.mutations.add_row(site=0, node=4, derived_state="1")
        tables.mutations.add_row(site=1, node=6, derived_state="1")
        ts = tables.tree_sequence()
        self.verify(ts)

    def test_single_tree_infinite_sites(self):
        ts = msprime.simulate(6, random_seed=1, mutation_rate=1)
        assert ts.num_sites > 0
        self.verify(ts)

    def test_single_tree_sites_no_mutations(self):
        ts = msprime.simulate(6, random_seed=1)
        tables = ts.dump_tables()
        tables.sites.add_row(0.1, "a")
        tables.sites.add_row(0.2, "aaa")
        self.verify(tables.tree_sequence())

    def test_single_tree_jukes_cantor(self):
        ts = msprime.simulate(6, random_seed=1, mutation_rate=1)
        ts = tsutil.jukes_cantor(ts, 20, 1, seed=10)
        self.verify(ts)

    def test_single_tree_single_site_many_silent(self):
        ts = msprime.simulate(6, random_seed=1)
        ts = tsutil.jukes_cantor(ts, 1, 20, seed=10)
        self.verify(ts)

    def test_single_tree_multichar_mutations(self):
        ts = msprime.simulate(6, random_seed=1, mutation_rate=1)
        ts = tsutil.insert_multichar_mutations(ts)
        self.verify(ts)

    def test_many_trees_infinite_sites(self):
        ts = msprime.simulate(6, recombination_rate=2, mutation_rate=2, random_seed=1)
        assert ts.num_sites > 0
        assert ts.num_trees > 2
        self.verify(ts)

    def test_many_trees_sequence_length_infinite_sites(self):
        for L in [0.5, 1.5, 3.3333]:
            ts = msprime.simulate(
                6, length=L, recombination_rate=2, mutation_rate=1, random_seed=1
            )
            self.verify(ts)

    def test_wright_fisher_unsimplified(self):
        tables = wf.wf_sim(
            4,
            5,
            seed=1,
            deep_history=True,
            initial_generation_samples=False,
            num_loci=10,
        )
        tables.sort()
        ts = msprime.mutate(tables.tree_sequence(), rate=0.05, random_seed=234)
        assert ts.num_sites > 0
        self.verify(ts)

    def test_wright_fisher_initial_generation(self):
        tables = wf.wf_sim(
            6, 5, seed=3, deep_history=True, initial_generation_samples=True, num_loci=2
        )
        tables.sort()
        tables.simplify()
        ts = msprime.mutate(tables.tree_sequence(), rate=0.08, random_seed=2)
        assert ts.num_sites > 0
        self.verify(ts)

    def test_wright_fisher_initial_generation_no_deep_history(self):
        tables = wf.wf_sim(
            7,
            15,
            seed=202,
            deep_history=False,
            initial_generation_samples=True,
            num_loci=5,
        )
        tables.sort()
        tables.simplify()
        ts = msprime.mutate(tables.tree_sequence(), rate=0.01, random_seed=2)
        assert ts.num_sites > 0
        self.verify(ts)

    def test_wright_fisher_unsimplified_multiple_roots(self):
        tables = wf.wf_sim(
            8,
            15,
            seed=1,
            deep_history=False,
            initial_generation_samples=False,
            num_loci=20,
        )
        tables.sort()
        ts = msprime.mutate(tables.tree_sequence(), rate=0.01, random_seed=2)
        assert ts.num_sites > 0
        self.verify(ts)

    def test_wright_fisher_simplified(self):
        tables = wf.wf_sim(
            9,
            10,
            seed=1,
            deep_history=True,
            initial_generation_samples=False,
            num_loci=5,
        )
        tables.sort()
        ts = tables.tree_sequence().simplify()
        ts = tsutil.jukes_cantor(ts, 10, 0.01, seed=1)
        assert ts.num_sites > 0
        self.verify(ts)

    def test_empty_ts(self):
        tables = tskit.TableCollection(1.0)
        for _ in range(10):
            tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
        ts = tables.tree_sequence()
        self.verify(ts)


def example_sample_sets(ts, min_size=1):
    """
    Generate a series of example sample sets from the specfied tree sequence. The
    number of sample sets returned in each example must be at least min_size
    """
    samples = ts.samples()
    np.random.shuffle(samples)
    splits = np.array_split(samples, min_size)
    yield splits
    yield [[s] for s in samples]
    if min_size == 1:
        yield [samples[:1]]
    if ts.num_samples > 2 and min_size <= 2:
        yield [samples[:2], samples[2:]]
    if ts.num_samples > 4 and min_size <= 2:
        yield [samples[:2], samples[2:4]]
    if ts.num_samples > 7 and min_size <= 4:
        yield [samples[:2], samples[2:4], samples[4:6], samples[6:]]
    if ts.num_samples > 8 and min_size <= 4:
        yield [samples[:2], samples[2:4], samples[4:6], samples[6:8]]


def example_sample_set_index_pairs(sample_sets):
    assert len(sample_sets) >= 2
    yield [(0, 1)]
    yield [(1, 0), (0, 1)]
    if len(sample_sets) > 2:
        yield [(0, 1), (1, 2), (0, 2)]


def example_sample_set_index_triples(sample_sets):
    assert len(sample_sets) >= 3
    yield [(0, 1, 2)]
    yield [(0, 2, 1), (2, 1, 0)]
    if len(sample_sets) > 3:
        yield [(3, 0, 1), (0, 2, 3), (1, 2, 3)]


def example_sample_set_index_quads(sample_sets):
    assert len(sample_sets) >= 4
    yield [(0, 1, 2, 3)]
    yield [(0, 1, 2, 3), (3, 2, 1, 0)]
    yield [(0, 1, 2, 3), (3, 2, 1, 0), (1, 2, 3, 0)]


def example_windows(ts):
    """
    Generate a series of example windows for the specified tree sequence.
    """
    L = ts.sequence_length
    yield "sites"
    yield "trees"
    yield [0, L]
    yield ts.breakpoints(as_array=True)
    yield np.linspace(0, L, num=10)


class WeightStatsMixin:
    """
    Implements the verify method and dispatches it to verify_weighted_stat
    for a representative set of sample sets and windows.
    """

    def example_weights(self, ts, min_size=1):
        """
        Generate a series of example weights from the specfied tree sequence.
        """
        np.random.seed(46)
        for k in [min_size, min_size + 1, min_size + 10]:
            W = 1.0 + np.zeros((ts.num_samples, k))
            W[0, :] = 2.0
            yield W
            for j in range(k):
                W[:, j] = np.random.exponential(1, ts.num_samples)
            yield W
            for j in range(k):
                W[:, j] = np.random.normal(0, 1, ts.num_samples)
            yield W

    def transform_weights(self, W):
        """
        Specific methods will need to transform weights
        before passing them to general_stat.
        """
        return W

    def verify(self, ts):
        for W, windows in subset_combos(
            self.example_weights(ts), example_windows(ts), p=0.1
        ):
            self.verify_weighted_stat(ts, W, windows=windows)

    def verify_definition(self, ts, W, windows, summary_func, ts_method, definition):
        # general_stat will need an extra column for p
        gW = self.transform_weights(W)

        def wrapped_summary_func(x):
            with suppress_division_by_zero_warning():
                return summary_func(x)

        # Determine output_dim of the function
        M = len(wrapped_summary_func(gW[0]))
        for sn in [True, False]:
            sigma1 = ts.general_stat(
                gW, wrapped_summary_func, M, windows, mode=self.mode, span_normalise=sn
            )
            sigma2 = general_stat(
                ts, gW, wrapped_summary_func, windows, mode=self.mode, span_normalise=sn
            )
            sigma3 = ts_method(W, windows=windows, mode=self.mode, span_normalise=sn)
            sigma4 = definition(
                ts, W, windows=windows, mode=self.mode, span_normalise=sn
            )

            assert sigma1.shape == sigma2.shape
            assert sigma1.shape == sigma3.shape
            assert sigma1.shape == sigma4.shape
            self.assertArrayAlmostEqual(sigma1, sigma2)
            self.assertArrayAlmostEqual(sigma1, sigma3)
            self.assertArrayAlmostEqual(sigma1, sigma4)


class SampleSetStatsMixin:
    """
    Implements the verify method and dispatches it to verify_sample_sets
    for a representative set of sample sets and windows.
    """

    def verify(self, ts):
        for sample_sets, windows in subset_combos(
            example_sample_sets(ts), example_windows(ts)
        ):
            self.verify_sample_sets(ts, sample_sets, windows=windows)

    def verify_definition(
        self, ts, sample_sets, windows, summary_func, ts_method, definition
    ):
        W = np.array([[u in A for A in sample_sets] for u in ts.samples()], dtype=float)

        def wrapped_summary_func(x):
            with suppress_division_by_zero_warning():
                return summary_func(x)

        for sn in [True, False]:
            # Determine output_dim of the function
            M = len(summary_func(W[0]))
            sigma1 = ts.general_stat(
                W, wrapped_summary_func, M, windows, mode=self.mode, span_normalise=sn
            )
            sigma2 = general_stat(
                ts, W, wrapped_summary_func, windows, mode=self.mode, span_normalise=sn
            )
            sigma3 = ts_method(
                sample_sets, windows=windows, mode=self.mode, span_normalise=sn
            )
            sigma4 = definition(
                ts, sample_sets, windows=windows, mode=self.mode, span_normalise=sn
            )

            assert sigma1.shape == sigma2.shape
            assert sigma1.shape == sigma3.shape
            assert sigma1.shape == sigma4.shape
            self.assertArrayAlmostEqual(sigma1, sigma2)
            self.assertArrayAlmostEqual(sigma1, sigma3)
            self.assertArrayAlmostEqual(sigma1, sigma4)


class KWaySampleSetStatsMixin(SampleSetStatsMixin):
    """
    Defines the verify definition method, which comparse the results from
    several different ways of defining and computing the same statistic.
    """

    def verify_definition(
        self, ts, sample_sets, indexes, windows, summary_func, ts_method, definition
    ):
        def wrapped_summary_func(x):
            with suppress_division_by_zero_warning():
                return summary_func(x)

        W = np.array([[u in A for A in sample_sets] for u in ts.samples()], dtype=float)
        # Determine output_dim of the function
        M = len(wrapped_summary_func(W[0]))
        sigma1 = ts.general_stat(W, wrapped_summary_func, M, windows, mode=self.mode)
        sigma2 = general_stat(ts, W, wrapped_summary_func, windows, mode=self.mode)
        sigma3 = ts_method(
            sample_sets, indexes=indexes, windows=windows, mode=self.mode
        )
        sigma4 = definition(
            ts, sample_sets, indexes=indexes, windows=windows, mode=self.mode
        )

        assert sigma1.shape == sigma2.shape
        assert sigma1.shape == sigma3.shape
        assert sigma1.shape == sigma4.shape
        self.assertArrayAlmostEqual(sigma1, sigma2)
        self.assertArrayAlmostEqual(sigma1, sigma3)
        self.assertArrayAlmostEqual(sigma1, sigma4)


class TwoWaySampleSetStatsMixin(KWaySampleSetStatsMixin):
    """
    Implements the verify method and dispatches it to verify_sample_sets_indexes,
    which gives a representative sample of sample set indexes.
    """

    def verify(self, ts):
        for sample_sets, windows in subset_combos(
            example_sample_sets(ts, min_size=2), example_windows(ts)
        ):
            for indexes in example_sample_set_index_pairs(sample_sets):
                self.verify_sample_sets_indexes(ts, sample_sets, indexes, windows)


class ThreeWaySampleSetStatsMixin(KWaySampleSetStatsMixin):
    """
    Implements the verify method and dispatches it to verify_sample_sets_indexes,
    which gives a representative sample of sample set indexes.
    """

    def verify(self, ts):
        for sample_sets, windows in subset_combos(
            example_sample_sets(ts, min_size=3), example_windows(ts)
        ):
            for indexes in example_sample_set_index_triples(sample_sets):
                self.verify_sample_sets_indexes(ts, sample_sets, indexes, windows)


class FourWaySampleSetStatsMixin(KWaySampleSetStatsMixin):
    """
    Implements the verify method and dispatches it to verify_sample_sets_indexes,
    which gives a representative sample of sample set indexes.
    """

    def verify(self, ts):
        for sample_sets, windows in subset_combos(
            example_sample_sets(ts, min_size=4), example_windows(ts)
        ):
            for indexes in example_sample_set_index_quads(sample_sets):
                self.verify_sample_sets_indexes(ts, sample_sets, indexes, windows)


############################################
# Diversity
############################################


def site_diversity(ts, sample_sets, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(sample_sets)))
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        haps = ts.genotype_matrix(isolated_as_missing=False).T
        site_positions = [x.position for x in ts.sites()]
        for i, X in enumerate(sample_sets):
            S = 0
            site_in_window = False
            denom = np.float64(len(X) * (len(X) - 1))
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    site_in_window = True
                    for x in X:
                        for y in set(X) - {x}:
                            x_index = np.where(samples == x)[0][0]
                            y_index = np.where(samples == y)[0][0]
                            if haps[x_index][k] != haps[y_index][k]:
                                # x|y
                                S += 1
            if site_in_window:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def branch_diversity(ts, sample_sets, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(sample_sets)))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, X in enumerate(sample_sets):
            S = 0
            denom = np.float64(len(X) * (len(X) - 1))
            has_trees = False
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                if tr.total_branch_length > 0:
                    has_trees = True
                SS = 0
                for x in X:
                    for y in set(X) - {x}:
                        SS += path_length(tr, x, y)
                S += SS * (min(end, tr.interval.right) - max(begin, tr.interval.left))
            if has_trees:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def node_diversity(ts, sample_sets, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    K = len(sample_sets)
    out = np.zeros((len(windows) - 1, ts.num_nodes, K))
    for k in range(K):
        X = sample_sets[k]
        for j in range(len(windows) - 1):
            begin = windows[j]
            end = windows[j + 1]
            tX = len(X)
            denom = np.float64(len(X) * (len(X) - 1))
            S = np.zeros(ts.num_nodes)
            for tr in ts.trees(tracked_samples=X):
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in tr.nodes():
                    # count number of pairwise paths going through u
                    n = tr.num_tracked_samples(u)
                    SS[u] += 2 * n * (tX - n)
                S += SS * (min(end, tr.interval.right) - max(begin, tr.interval.left))
            with suppress_division_by_zero_warning():
                out[j, :, k] = S / denom
            if span_normalise:
                out[j, :, k] /= end - begin
    return out


def diversity(ts, sample_sets, windows=None, mode="site", span_normalise=True):
    """
    Computes average pairwise diversity between two random choices from x
    over the window specified.
    """
    method_map = {
        "site": site_diversity,
        "node": node_diversity,
        "branch": branch_diversity,
    }
    return method_map[mode](
        ts, sample_sets, windows=windows, span_normalise=span_normalise
    )


class TestDiversity(StatsTestCase, SampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_sample_sets(self, ts, sample_sets, windows):
        n = np.array([len(x) for x in sample_sets])

        def f(x):
            with np.errstate(invalid="ignore", divide="ignore"):
                return x * (n - x) / (n * (n - 1))

        self.verify_definition(ts, sample_sets, windows, f, ts.diversity, diversity)


class TestBranchDiversity(TestDiversity, TopologyExamplesMixin):
    mode = "branch"


class TestNodeDiversity(TestDiversity, TopologyExamplesMixin):
    mode = "node"


class TestSiteDiversity(TestDiversity, MutatedTopologyExamplesMixin):
    mode = "site"


############################################
# Segregating sites
############################################


def site_segregating_sites(ts, sample_sets, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(sample_sets)))
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        haps = ts.genotype_matrix(isolated_as_missing=False)
        site_positions = [x.position for x in ts.sites()]
        for i, X in enumerate(sample_sets):
            set_X = set(X)
            X_index = np.where(np.fromiter((s in set_X for s in samples), dtype=bool))[
                0
            ]
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    num_alleles = len(set(haps[k, X_index]))
                    out[j][i] += num_alleles - 1
            if span_normalise:
                out[j][i] /= end - begin
    return out


def branch_segregating_sites(ts, sample_sets, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(sample_sets)))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, X in enumerate(sample_sets):
            tX = len(X)
            for tr in ts.trees(tracked_samples=X):
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                SS = 0
                for u in tr.nodes():
                    nX = tr.num_tracked_samples(u)
                    if nX > 0 and nX < tX:
                        SS += tr.branch_length(u)
                out[j][i] += SS * (
                    min(end, tr.interval.right) - max(begin, tr.interval.left)
                )
            if span_normalise:
                out[j][i] /= end - begin
    return out


def node_segregating_sites(ts, sample_sets, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    K = len(sample_sets)
    out = np.zeros((len(windows) - 1, ts.num_nodes, K))
    for k in range(K):
        X = sample_sets[k]
        for j in range(len(windows) - 1):
            begin = windows[j]
            end = windows[j + 1]
            tX = len(X)
            S = np.zeros(ts.num_nodes)
            for tr in ts.trees(tracked_samples=X):
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in tr.nodes():
                    nX = tr.num_tracked_samples(u)
                    SS[u] = (nX > 0) and (nX < tX)
                S += SS * (min(end, tr.interval.right) - max(begin, tr.interval.left))
            out[j, :, k] = S
            if span_normalise:
                out[j, :, k] /= end - begin
    return out


def segregating_sites(ts, sample_sets, windows=None, mode="site", span_normalise=True):
    """
    Computes the density of segregating sites over the window specified.
    """
    method_map = {
        "site": site_segregating_sites,
        "node": node_segregating_sites,
        "branch": branch_segregating_sites,
    }
    return method_map[mode](
        ts, sample_sets, windows=windows, span_normalise=span_normalise
    )


class TestSegregatingSites(StatsTestCase, SampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_sample_sets(self, ts, sample_sets, windows):
        n = np.array([len(x) for x in sample_sets])

        # this works because sum_{i=1}^k (1-p_i) = k-1
        def f(x):
            return (x > 0) * (1 - x / n)

        self.verify_definition(
            ts, sample_sets, windows, f, ts.segregating_sites, segregating_sites
        )


class TestBranchSegregatingSites(TestSegregatingSites, TopologyExamplesMixin):
    mode = "branch"


class TestNodeSegregatingSites(TestSegregatingSites, TopologyExamplesMixin):
    mode = "node"


class TestSiteSegregatingSites(TestSegregatingSites, MutatedTopologyExamplesMixin):
    mode = "site"


class TestBranchSegregatingSitesProperties(StatsTestCase, TopologyExamplesMixin):
    def verify(self, ts):
        windows = ts.breakpoints(as_array=True)
        # If we split by tree, this should always be equal to the total
        # branch length. The definition of total_branch_length here is slightly
        # tricky: it's the sum of all branch lengths that subtend between 0
        # and n samples. This differs from the built-in total_branch_length
        # function, which just sums that total branch length reachable from
        # roots.
        tbl_tree = [
            sum(
                tree.branch_length(u)
                for u in tree.nodes()
                if 0 < tree.num_samples(u) < ts.num_samples
            )
            for tree in ts.trees()
        ]
        # We must span_normalise, because these values are always weighted
        # by the span, so we're effectively cancelling out this contribution
        tbl = ts.segregating_sites([ts.samples()], windows=windows, mode="branch")
        tbl = tbl.reshape(tbl.shape[:-1])
        self.assertArrayAlmostEqual(tbl_tree, tbl)


############################################
# Tajima's D
############################################


def site_tajimas_d(ts, sample_sets, windows=None):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(sample_sets)))
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        haps = ts.genotype_matrix(isolated_as_missing=False)
        site_positions = [x.position for x in ts.sites()]
        n = np.array([len(X) for X in sample_sets])
        for i, X in enumerate(sample_sets):
            nn = n[i]
            S = 0
            T = 0
            set_X = set(X)
            X_index = np.where(np.fromiter((s in set_X for s in samples), dtype=bool))[
                0
            ]
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    hX = haps[k, X_index]
                    alleles = set(hX)
                    num_alleles = len(alleles)
                    n_alleles = [np.sum(hX == a) for a in alleles]
                    S += num_alleles - 1
                    for k in n_alleles:
                        with suppress_division_by_zero_warning():
                            T += k * (nn - k) / (nn * (nn - 1))
            with suppress_division_by_zero_warning():
                a1 = np.sum(1 / np.arange(1, nn))  # this is h in the main version
                a2 = np.sum(1 / np.arange(1, nn) ** 2)  # this is g
                b1 = (nn + 1) / (3 * (nn - 1))
                b2 = 2 * (nn**2 + nn + 3) / (9 * nn * (nn - 1))
                c1 = b1 - 1 / a1
                c2 = b2 - (nn + 2) / (a1 * nn) + a2 / a1**2
                e1 = c1 / a1  # this is a
                e2 = c2 / (a1**2 + a2)  # this is b
                out[j][i] = (T - S / a1) / np.sqrt(e1 * S + e2 * S * (S - 1))
    return out


def tajimas_d(ts, sample_sets, windows=None, mode="site", span_normalise=True):
    method_map = {"site": site_tajimas_d}
    return method_map[mode](
        ts, sample_sets, windows=windows, span_normalise=span_normalise
    )


class TestTajimasD(StatsTestCase, SampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify(self, ts):
        # only check per-site
        for sample_sets in example_sample_sets(ts, min_size=1):
            self.verify_persite_tajimas_d(ts, sample_sets)

    def get_windows(self, ts):
        yield "sites"
        yield [0, ts.sequence_length]
        yield np.arange(0, 1.1, 0.1) * ts.sequence_length

    def verify_persite_tajimas_d(self, ts, sample_sets):
        for windows in self.get_windows(ts):
            sigma1 = ts.Tajimas_D(sample_sets, windows=windows, mode=self.mode)
            sigma2 = site_tajimas_d(ts, sample_sets, windows=windows)
            assert sigma1.shape == sigma2.shape
            self.assertArrayAlmostEqual(sigma1, sigma2)


class TestSiteTajimasD(TestTajimasD, MutatedTopologyExamplesMixin):
    mode = "site"


############################################
# Y1
############################################


def branch_Y1(ts, sample_sets, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(sample_sets)))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, X in enumerate(sample_sets):
            S = 0
            denom = np.float64(len(X) * (len(X) - 1) * (len(X) - 2))
            has_trees = False
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                if tr.total_branch_length > 0:
                    has_trees = True
                this_length = min(end, tr.interval.right) - max(begin, tr.interval.left)
                for x in X:
                    for y in set(X) - {x}:
                        for z in set(X) - {x, y}:
                            xy_mrca = tr.mrca(x, y)
                            xz_mrca = tr.mrca(x, z)
                            yz_mrca = tr.mrca(y, z)
                            if xy_mrca == xz_mrca:
                                #   /\
                                #  / /\
                                # x y  z
                                S += path_length(tr, x, yz_mrca) * this_length
                            elif xy_mrca == yz_mrca:
                                #   /\
                                #  / /\
                                # y x  z
                                S += path_length(tr, x, xz_mrca) * this_length
                            elif xz_mrca == yz_mrca:
                                #   /\
                                #  / /\
                                # z x  y
                                S += path_length(tr, x, xy_mrca) * this_length
            if has_trees:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def site_Y1(ts, sample_sets, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(sample_sets)))
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        haps = ts.genotype_matrix(isolated_as_missing=False).T
        site_positions = [x.position for x in ts.sites()]
        for i, X in enumerate(sample_sets):
            S = 0
            denom = np.float64(len(X) * (len(X) - 1) * (len(X) - 2))
            site_in_window = False
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    site_in_window = True
                    for x in X:
                        x_index = np.where(samples == x)[0][0]
                        for y in set(X) - {x}:
                            y_index = np.where(samples == y)[0][0]
                            for z in set(X) - {x, y}:
                                z_index = np.where(samples == z)[0][0]
                                condition = (
                                    haps[x_index, k] != haps[y_index, k]
                                    and haps[x_index, k] != haps[z_index, k]
                                )
                                if condition:
                                    # x|yz
                                    S += 1
            if site_in_window:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def node_Y1(ts, sample_sets, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    K = len(sample_sets)
    out = np.zeros((len(windows) - 1, ts.num_nodes, K))
    for k in range(K):
        X = sample_sets[k]
        for j in range(len(windows) - 1):
            begin = windows[j]
            end = windows[j + 1]
            tX = len(X)
            denom = np.float64(tX * (tX - 1) * (tX - 2))
            S = np.zeros(ts.num_nodes)
            for tr in ts.trees(tracked_samples=X):
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in tr.nodes():
                    # count number of paths above a but not b,c
                    n = tr.num_tracked_samples(u)
                    SS[u] += n * (tX - n) * (tX - n - 1) + (tX - n) * n * (n - 1)
                S += SS * (min(end, tr.interval.right) - max(begin, tr.interval.left))
            with suppress_division_by_zero_warning():
                out[j, :, k] = S / denom
            if span_normalise:
                out[j, :, k] /= end - begin
    return out


def Y1(ts, sample_sets, windows=None, mode="site", span_normalise=True):
    windows = ts.parse_windows(windows)
    method_map = {"site": site_Y1, "node": node_Y1, "branch": branch_Y1}
    return method_map[mode](
        ts, sample_sets, windows=windows, span_normalise=span_normalise
    )


class TestY1(StatsTestCase, SampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_sample_sets(self, ts, sample_sets, windows):
        n = np.array([len(x) for x in sample_sets])
        denom = n * (n - 1) * (n - 2)

        def f(x):
            with np.errstate(invalid="ignore", divide="ignore"):
                return x * (n - x) * (n - x - 1) / denom

        self.verify_definition(ts, sample_sets, windows, f, ts.Y1, Y1)


class TestBranchY1(TestY1, TopologyExamplesMixin):
    mode = "branch"


class TestNodeY1(TestY1, TopologyExamplesMixin):
    mode = "node"


class TestSiteY1(TestY1, MutatedTopologyExamplesMixin):
    mode = "site"


############################################
# Divergence
############################################


def site_divergence(ts, sample_sets, indexes, windows=None, span_normalise=True):
    out = np.zeros((len(windows) - 1, len(indexes)))
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        haps = ts.genotype_matrix(isolated_as_missing=False).T
        site_positions = [x.position for x in ts.sites()]
        for i, (ix, iy) in enumerate(indexes):
            X = sample_sets[ix]
            Y = sample_sets[iy]
            denom = np.float64(len(X) * len(Y))
            site_in_window = False
            S = 0
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    site_in_window = True
                    for x in X:
                        x_index = np.where(samples == x)[0][0]
                        for y in Y:
                            y_index = np.where(samples == y)[0][0]
                            if haps[x_index][k] != haps[y_index][k]:
                                # x|y
                                S += 1
            if site_in_window:
                with np.errstate(invalid="ignore", divide="ignore"):
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def branch_divergence(ts, sample_sets, indexes, windows=None, span_normalise=True):
    out = np.zeros((len(windows) - 1, len(indexes)))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, (ix, iy) in enumerate(indexes):
            X = sample_sets[ix]
            Y = sample_sets[iy]
            denom = np.float64(len(X) * len(Y))
            has_trees = False
            S = 0
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                if tr.total_branch_length > 0:
                    has_trees = True
                SS = 0
                for x in X:
                    for y in Y:
                        SS += path_length(tr, x, y)
                S += SS * (min(end, tr.interval.right) - max(begin, tr.interval.left))
            if has_trees:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def node_divergence(ts, sample_sets, indexes, windows=None, span_normalise=True):
    out = np.zeros((len(windows) - 1, ts.num_nodes, len(indexes)))
    for i, (ix, iy) in enumerate(indexes):
        X = sample_sets[ix]
        Y = sample_sets[iy]
        tX = len(X)
        tY = len(Y)
        denom = np.float64(len(X) * len(Y))
        for j in range(len(windows) - 1):
            begin = windows[j]
            end = windows[j + 1]
            S = np.zeros(ts.num_nodes)
            for t1, t2 in zip(ts.trees(tracked_samples=X), ts.trees(tracked_samples=Y)):
                if t1.interval.right <= begin:
                    continue
                if t1.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in t1.nodes():
                    # count number of pairwise paths going through u
                    nX = t1.num_tracked_samples(u)
                    nY = t2.num_tracked_samples(u)
                    SS[u] += nX * (tY - nY) + (tX - nX) * nY
                S += SS * (min(end, t1.interval.right) - max(begin, t1.interval.left))
            with suppress_division_by_zero_warning():
                out[j, :, i] = S / denom
            if span_normalise:
                out[j, :, i] /= end - begin
    return out


def divergence(
    ts, sample_sets, indexes=None, windows=None, mode="site", span_normalise=True
):
    """
    Computes average pairwise divergence between two random choices from x
    over the window specified.
    """
    windows = ts.parse_windows(windows)
    if indexes is None:
        indexes = [(0, 1)]
    method_map = {
        "site": site_divergence,
        "node": node_divergence,
        "branch": branch_divergence,
    }
    return method_map[mode](
        ts, sample_sets, indexes=indexes, windows=windows, span_normalise=span_normalise
    )


class TestDivergence(StatsTestCase, TwoWaySampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_sample_sets_indexes(self, ts, sample_sets, indexes, windows):
        n = np.array([len(x) for x in sample_sets])

        denom = np.array([n[i] * (n[j] - (i == j)) for i, j in indexes])

        def f(x):
            numer = np.array([(x[i] * (n[j] - x[j])) for i, j in indexes])
            return numer / denom

        self.verify_definition(
            ts, sample_sets, indexes, windows, f, ts.divergence, divergence
        )


class TestBranchDivergence(TestDivergence, TopologyExamplesMixin):
    mode = "branch"


class TestNodeDivergence(TestDivergence, TopologyExamplesMixin):
    mode = "node"


class TestSiteDivergence(TestDivergence, MutatedTopologyExamplesMixin):
    mode = "site"


############################################
# Genetic relatedness
############################################


def site_genetic_relatedness(
    ts,
    sample_sets,
    indexes,
    windows=None,
    span_normalise=True,
    polarised=True,
    proportion=True,
    centre=True,
):
    if windows is None:
        windows = [0.0, ts.sequence_length]
    out = np.zeros((len(windows) - 1, len(indexes)))
    all_samples = np.array(list({u for s in sample_sets for u in s}))
    denom = np.ones(len(windows))
    if proportion:
        denom = ts.segregating_sites(
            sample_sets=all_samples,
            windows=windows,
            mode="site",
            span_normalise=span_normalise,
        )
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for vv in zip(
            *[
                ts.variants(left=begin, right=end, samples=x, isolated_as_missing=False)
                for x in sample_sets
            ]
        ):
            ancestral_state = vv[0].site.ancestral_state
            alleles = vv[0].alleles
            ff = [v.frequencies() for v in vv]
            for a in alleles:
                mean_f = sum([f[a] for f in ff]) / len(ff)
                for i, (ix, iy) in enumerate(indexes):
                    fx = ff[ix][a]
                    fy = ff[iy][a]
                    if not (polarised and a == ancestral_state):
                        if centre:
                            out[j][i] += (fx - mean_f) * (fy - mean_f)
                        else:
                            out[j][i] += fx * fy
        for i in range(len(indexes)):
            with np.errstate(invalid="ignore", divide="ignore"):
                out[j][i] /= denom[j]
            if span_normalise:
                out[j][i] /= end - begin
    return out


def branch_genetic_relatedness(
    ts,
    sample_sets,
    indexes,
    windows=None,
    span_normalise=True,
    polarised=True,
    proportion=True,
    centre=True,
):
    if windows is None:
        windows = [0.0, ts.sequence_length]
    out = np.zeros((len(windows) - 1, len(indexes)))
    all_samples = np.array(list({u for s in sample_sets for u in s}))
    denom = np.ones(len(windows))
    if proportion:
        denom = ts.segregating_sites(
            sample_sets=all_samples,
            windows=windows,
            mode="branch",
            span_normalise=span_normalise,
        )
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for tr in ts.trees():
            if tr.interval.right <= begin:
                continue
            if tr.interval.left >= end:
                break
            span = min(end, tr.interval.right) - max(begin, tr.interval.left)
            # iterating over tr.nodes will miss nodes unreachable from samples
            for v in range(ts.num_nodes):
                area = tr.branch_length(v) * span
                freqs = [
                    sum([tr.is_descendant(u, v) for u in x]) / len(x)
                    for x in sample_sets
                ]
                mean_freq = sum(freqs) / len(freqs)
                for i, (ix, iy) in enumerate(indexes):
                    fx = freqs[ix]
                    fy = freqs[iy]
                    if centre:
                        out[j][i] += area * (fx - mean_freq) * (fy - mean_freq)
                        if not polarised:
                            out[j][i] += (
                                area
                                * (1 - fx - (1 - mean_freq))
                                * (1 - fy - (1 - mean_freq))
                            )
                    else:
                        out[j][i] += area * fx * fy
                        if not polarised:
                            out[j][i] += area * (1 - fx) * (1 - fy)
        for i in range(len(indexes)):
            with np.errstate(invalid="ignore", divide="ignore"):
                out[j][i] /= denom[j]
            if span_normalise:
                out[j][i] /= end - begin
    return out


def node_genetic_relatedness(
    ts,
    sample_sets,
    indexes,
    windows=None,
    span_normalise=True,
    proportion=True,
    centre=True,
    polarised=True,
):
    if windows is None:
        windows = [0.0, ts.sequence_length]
    out = np.zeros((len(windows) - 1, ts.num_nodes, len(indexes)))
    all_samples = np.array(list({u for s in sample_sets for u in s}))
    denom = np.ones((len(windows), ts.num_nodes))
    if proportion:
        denom = ts.segregating_sites(
            sample_sets=all_samples,
            windows=windows,
            mode="node",
            span_normalise=span_normalise,
        )
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for tr in ts.trees():
            if tr.interval.right <= begin:
                continue
            if tr.interval.left >= end:
                break
            span = min(end, tr.interval.right) - max(begin, tr.interval.left)
            for v in range(ts.num_nodes):
                freqs = [
                    sum([tr.is_descendant(u, v) for u in x]) / len(x)
                    for x in sample_sets
                ]
                mean_freq = sum(freqs) / len(freqs)
                for i, (ix, iy) in enumerate(indexes):
                    fx = freqs[ix]
                    fy = freqs[iy]
                    if centre:
                        out[j][v][i] += span * (fx - mean_freq) * (fy - mean_freq)
                        if not polarised:
                            out[j][v][i] += (
                                span
                                * (1 - fx - (1 - mean_freq))
                                * (1 - fy - (1 - mean_freq))
                            )
                    else:
                        out[j][v][i] += span * fx * fy
                        if not polarised:
                            out[j][v][i] += span * (1 - fx) * (1 - fy)
        for i in range(len(indexes)):
            for v in ts.nodes():
                iV = v.id
                with np.errstate(invalid="ignore", divide="ignore"):
                    out[j, iV, i] /= denom[j, iV]
                if span_normalise:
                    out[j, iV, i] /= end - begin
    return out


def genetic_relatedness(
    ts,
    sample_sets,
    indexes=None,
    windows=None,
    mode="site",
    span_normalise=True,
    proportion=True,
    centre=True,
    polarised=True,
):
    """
    Computes genetic relatedness between two random choices from x
    over the window specified.
    """
    windows = ts.parse_windows(windows)
    if indexes is None:
        indexes = [(0, 1)]
    method_map = {
        "site": site_genetic_relatedness,
        "node": node_genetic_relatedness,
        "branch": branch_genetic_relatedness,
    }
    return method_map[mode](
        ts,
        sample_sets,
        indexes=indexes,
        windows=windows,
        span_normalise=span_normalise,
        polarised=polarised,
        proportion=proportion,
        centre=centre,
    )


class TestGeneticRelatedness(StatsTestCase, TwoWaySampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_definition(
        self,
        ts,
        sample_sets,
        indexes,
        windows,
        summary_func,
        ts_method,
        definition,
        proportion,
        polarised=True,
        centre=True,
    ):
        def wrapped_summary_func(x):
            with suppress_division_by_zero_warning():
                return summary_func(x)

        W = np.array([[u in A for A in sample_sets] for u in ts.samples()], dtype=float)
        # Determine output_dim of the function
        M = len(wrapped_summary_func(W[0]))
        denom = 1
        if proportion:
            all_samples = list({u for s in sample_sets for u in s})
            denom = ts.segregating_sites(
                sample_sets=[all_samples], windows=windows, mode=self.mode
            )

        with np.errstate(divide="ignore", invalid="ignore"):
            sigma1 = (
                ts.general_stat(
                    W,
                    wrapped_summary_func,
                    M,
                    windows,
                    mode=self.mode,
                    strict=centre,
                    polarised=polarised,
                )
                / denom
            )
        sigma2 = ts_method(
            sample_sets,
            indexes=indexes,
            windows=windows,
            mode=self.mode,
            proportion=proportion,
            centre=centre,
            polarised=polarised,
        )
        sigma3 = definition(
            ts,
            sample_sets,
            indexes=indexes,
            windows=windows,
            mode=self.mode,
            proportion=proportion,
            centre=centre,
            polarised=polarised,
        )
        assert sigma1.shape == sigma2.shape
        assert sigma1.shape == sigma3.shape
        self.assertArrayAlmostEqual(sigma1, sigma2)
        self.assertArrayAlmostEqual(sigma1, sigma3)

    def verify_sample_sets_indexes(self, ts, sample_sets, indexes, windows):
        n = np.array([len(x) for x in sample_sets])

        def f_noncentred(x):
            p = x / n
            return np.array([p[i] * p[j] for i, j in indexes])

        def f_centred(x):
            p = x / n
            mp = np.mean(p)
            return np.array([(p[i] - mp) * (p[j] - mp) for i, j in indexes])

        for proportion in [True, False]:
            self.verify_definition(
                ts,
                sample_sets,
                indexes,
                windows,
                f_centred,
                ts.genetic_relatedness,
                genetic_relatedness,
                proportion,
            )

        for centre, polarised in [
            (True, True),
            (False, True),
            (True, False),
            (False, False),
        ]:
            f = f_centred if centre else f_noncentred
            self.verify_definition(
                ts,
                sample_sets,
                indexes,
                windows,
                f,
                ts.genetic_relatedness,
                genetic_relatedness,
                proportion=False,
                centre=centre,
                polarised=polarised,
            )

    @pytest.mark.parametrize("proportion", [None, True, False])
    def test_shapes(self, proportion):
        # exclude this test in the parent class
        if self.mode is None:
            return
        ts = msprime.sim_ancestry(
            8,
            random_seed=1,
            end_time=10,
            sequence_length=10,
            population_size=10,
            recombination_rate=0.02,
        )
        ts = msprime.sim_mutations(ts, rate=0.01, random_seed=2)
        x = ts.genetic_relatedness(
            sample_sets=[[0, 1, 2], [3]],
            indexes=None,
            windows=None,
            mode=self.mode,
            proportion=proportion,
        )
        if self.mode == "node":
            assert x.shape == (ts.num_nodes,)
        else:
            assert x.shape == ()
        x = ts.genetic_relatedness(
            sample_sets=[[0, 1, 2], [3]],
            indexes=[(0, 1)],
            windows=None,
            mode=self.mode,
            proportion=proportion,
        )
        if self.mode == "node":
            assert x.shape == (ts.num_nodes, 1)
        else:
            assert x.shape == (1,)
        x = ts.genetic_relatedness(
            sample_sets=[[0, 1, 2], [3]],
            indexes=[(0, 1)],
            windows=[0, 10],
            mode=self.mode,
            proportion=proportion,
        )
        if self.mode == "node":
            assert x.shape == (1, ts.num_nodes, 1)
        else:
            assert x.shape == (1, 1)
        x = ts.genetic_relatedness(
            sample_sets=[[0, 1, 2], [3]],
            indexes=[(0, 1)],
            windows=[0, 5, 10],
            mode=self.mode,
            proportion=proportion,
        )
        if self.mode == "node":
            assert x.shape == (2, ts.num_nodes, 1)
        else:
            assert x.shape == (2, 1)
        x = ts.genetic_relatedness(
            sample_sets=[[0, 1, 2], [3]],
            indexes=None,
            windows=[0, 5, 10],
            mode=self.mode,
            proportion=proportion,
        )
        if self.mode == "node":
            assert x.shape == (2, ts.num_nodes)
        else:
            assert x.shape == (2,)
        x = ts.genetic_relatedness(
            sample_sets=[[0, 1, 2], [3], [4, 5]],
            indexes=[(0, 1), (1, 2)],
            windows=[0, 5, 9, 10],
            mode=self.mode,
            proportion=proportion,
        )
        if self.mode == "node":
            assert x.shape == (3, ts.num_nodes, 2)
        else:
            assert x.shape == (3, 2)
        x = ts.genetic_relatedness(
            sample_sets=[[0, 1, 2], [3], [4, 5]],
            indexes=[(0, 1), (1, 2)],
            windows=None,
            mode=self.mode,
            proportion=proportion,
        )
        if self.mode == "node":
            assert x.shape == (ts.num_nodes, 2)
        else:
            assert x.shape == (2,)


class TestBranchGeneticRelatedness(TestGeneticRelatedness, TopologyExamplesMixin):
    mode = "branch"

    @pytest.mark.parametrize("polarised", [True, False])
    def test_simple_tree_noncentred(self, polarised):
        # 2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊
        # 1.00┊ ┃  3  ┊
        #     ┊ ┃ ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        ts = tskit.Tree.generate_balanced(3).tree_sequence
        indexes = [(0, 0), (0, 1), (1, 1), (1, 2), (2, 2)]
        sample_sets = [[0], [1], [2]]
        if polarised:
            A = np.array(
                [
                    2,  # (0, 0)
                    0,  # (0, 1)
                    2,  # (1, 1)
                    1,  # (1, 2),
                    2,  # (2, 2)
                ]
            )
        else:
            A = np.array(
                [
                    (2 + 3),  # (0, 0)
                    (0 + 1),  # (0, 1)
                    (2 + 3),  # (1, 1)
                    (1 + 2),  # (1, 2),
                    (2 + 3),  # (2, 2)
                ]
            )
        B = branch_genetic_relatedness(
            ts,
            sample_sets=sample_sets,
            indexes=indexes,
            polarised=polarised,
            proportion=False,
            centre=False,
        ).squeeze()
        C = ts.genetic_relatedness(
            sample_sets=sample_sets,
            indexes=indexes,
            mode="branch",
            polarised=polarised,
            proportion=False,
            centre=False,
        ).squeeze()
        self.assertArrayAlmostEqual(A, B)
        self.assertArrayAlmostEqual(A, C)


class TestNodeGeneticRelatedness(TestGeneticRelatedness, TopologyExamplesMixin):
    mode = "node"


class TestSiteGeneticRelatedness(TestGeneticRelatedness, MutatedTopologyExamplesMixin):
    mode = "site"

    def test_match_K_c0(self):
        # This test checks that ts.genetic_relatedness() matches K_c0
        # from Speed & Balding (2014) https://www.nature.com/articles/nrg3821
        ts = msprime.simulate(
            10, mutation_rate=0.01, length=100, recombination_rate=0.01, random_seed=23
        )
        samples = [u for u in ts.samples()]
        sample_sets = [[0, 1], [2, 3], [4, 5]]
        all_samples = list({u for s in sample_sets for u in s})
        sample_ind = [samples.index(x) for x in all_samples]
        indexes = [(0, 0), (0, 1), (0, 2), (1, 1), (1, 2), (2, 2)]
        A = ts.genetic_relatedness(
            sample_sets, indexes=indexes, mode="site", span_normalise=False
        )
        # Genotype covariance as in Speed and Balding
        G = ts.genotype_matrix().T
        G = G[sample_ind]
        G_centered = G - G.mean(axis=0)
        B = np.zeros(len(indexes))
        for i, (ix, iy) in enumerate(indexes):
            x1 = sample_sets[ix][0]
            x2 = sample_sets[ix][1]
            y1 = sample_sets[iy][0]
            y2 = sample_sets[iy][1]
            B[i] = (
                (G_centered[x1] + G_centered[x2])
                @ (G_centered[y1] + G_centered[y2])
                / ts.segregating_sites(sample_sets=all_samples, span_normalise=False)
            )
        self.assertArrayAlmostEqual(4 * A, B)


############################################
# Genetic relatedness weighted
############################################


def genetic_relatedness_matrix(
    ts, sample_sets, windows=None, mode="site", polarised=True, centre=True
):
    n = len(sample_sets)
    indexes = [
        (n1, n2) for n1, n2 in itertools.combinations_with_replacement(range(n), 2)
    ]
    if windows is None:
        if mode == "node":
            n_nodes = ts.num_nodes
            K = np.zeros((n_nodes, n, n))
            out = ts.genetic_relatedness(
                sample_sets,
                indexes,
                mode=mode,
                proportion=False,
                span_normalise=True,
                polarised=polarised,
                centre=centre,
            )
            for node in range(n_nodes):
                this_K = np.zeros((n, n))
                this_K[np.triu_indices(n)] = out[node, :]
                this_K = this_K + np.triu(this_K, 1).transpose()
                K[node, :, :] = this_K
        else:
            K = np.zeros((n, n))
            K[np.triu_indices(n)] = ts.genetic_relatedness(
                sample_sets,
                indexes,
                mode=mode,
                proportion=False,
                span_normalise=True,
                centre=centre,
                polarised=polarised,
            )
            K = K + np.triu(K, 1).transpose()
    else:
        windows = ts.parse_windows(windows)
        n_windows = len(windows) - 1
        out = ts.genetic_relatedness(
            sample_sets,
            indexes,
            mode=mode,
            windows=windows,
            proportion=False,
            span_normalise=True,
            polarised=polarised,
            centre=centre,
        )
        if mode == "node":
            n_nodes = ts.num_nodes
            K = np.zeros((n_windows, n_nodes, n, n))
            for win in range(n_windows):
                for node in range(n_nodes):
                    K_this = np.zeros((n, n))
                    K_this[np.triu_indices(n)] = out[win, node, :]
                    K_this = K_this + np.triu(K_this, 1).transpose()
                    K[win, node, :, :] = K_this
        else:
            K = np.zeros((n_windows, n, n))
            for win in range(n_windows):
                K_this = np.zeros((n, n))
                K_this[np.triu_indices(n)] = out[win, :]
                K_this = K_this + np.triu(K_this, 1).transpose()
                K[win, :, :] = K_this
    return K


def genetic_relatedness_weighted(
    ts, W, indexes, windows=None, mode="site", polarised=True, centre=True
):
    if centre:
        W_mean = W.mean(axis=0)
        W = W - W_mean
    sample_sets = [[u] for u in ts.samples()]
    K = genetic_relatedness_matrix(
        ts, sample_sets, windows=windows, mode=mode, centre=centre, polarised=polarised
    )
    n_indexes = len(indexes)
    n_nodes = ts.num_nodes
    if windows is None:
        if mode == "node":
            out = np.zeros((n_nodes, n_indexes))
        else:
            out = np.zeros(n_indexes)
    else:
        windows = ts.parse_windows(windows)
        n_windows = len(windows) - 1
        if mode == "node":
            out = np.zeros((n_windows, n_nodes, n_indexes))
        else:
            out = np.zeros((n_windows, n_indexes))
    for pair in range(n_indexes):
        i1 = indexes[pair][0]
        i2 = indexes[pair][1]
        if windows is None:
            if mode == "node":
                for node in range(n_nodes):
                    this_K = K[node, :, :]
                    out[node, pair] = W[:, i1] @ this_K @ W[:, i2]
            else:
                out[pair] = W[:, i1] @ K @ W[:, i2]
        else:
            for win in range(n_windows):
                if mode == "node":
                    for node in range(n_nodes):
                        this_K = K[win, node, :, :]
                        out[win, node, pair] = W[:, i1] @ this_K @ W[:, i2]
                else:
                    this_K = K[win, :, :]
                    out[win, pair] = W[:, i1] @ this_K @ W[:, i2]
    return out


def example_index_pairs(weights):
    assert weights.shape[1] >= 2
    yield [(0, 1)]
    yield [(1, 0), (0, 1)]
    if weights.shape[1] > 2:
        yield [(0, 1), (1, 2), (0, 2)]


class TestGeneticRelatednessWeighted(StatsTestCase, WeightStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_definition(
        self,
        ts,
        W,
        indexes,
        windows,
        summary_func,
        ts_method,
        definition,
        polarised=True,
        centre=True,
    ):
        # Determine output_dim of the function
        M = len(indexes)

        sigma1 = ts.general_stat(
            W,
            summary_func,
            M,
            windows,
            mode=self.mode,
            span_normalise=True,
            strict=centre,
            polarised=polarised,
        )
        sigma2 = general_stat(
            ts,
            W,
            summary_func,
            windows,
            mode=self.mode,
            span_normalise=True,
            polarised=polarised,
        )

        sigma3 = ts_method(
            W,
            indexes=indexes,
            windows=windows,
            mode=self.mode,
            polarised=polarised,
            centre=centre,
        )
        sigma4 = definition(
            ts,
            W,
            indexes=indexes,
            windows=windows,
            mode=self.mode,
            polarised=polarised,
            centre=centre,
        )
        assert sigma1.shape == sigma2.shape
        assert sigma1.shape == sigma3.shape
        assert sigma1.shape == sigma4.shape
        self.assertArrayAlmostEqual(sigma1, sigma2)
        self.assertArrayAlmostEqual(sigma1, sigma3)
        self.assertArrayAlmostEqual(sigma1, sigma4)

    def verify(self, ts):
        for W, windows in subset_combos(
            self.example_weights(ts, min_size=2), example_windows(ts), p=0.1
        ):
            for indexes in example_index_pairs(W):
                self.verify_weighted_stat(ts, W, indexes, windows)

    def verify_weighted_stat(self, ts, W, indexes, windows):
        n = W.shape[0]
        K = W.shape[1]
        WW = np.column_stack([W, np.ones(n) / n])
        W_sum = WW.sum(axis=0)

        def f_noncentred(x):
            return np.array([x[i] * x[j] for i, j in indexes])

        def f_centred(x):
            pn = x[K]
            return np.array(
                [(x[i] - W_sum[i] * pn) * (x[j] - W_sum[j] * pn) for i, j in indexes]
            )

        for centre, polarised in [
            (True, True),
            (False, True),
            (True, False),
            (False, False),
        ]:
            f = f_centred if centre else f_noncentred
            self.verify_definition(
                ts,
                WW,
                indexes,
                windows,
                f,
                ts.genetic_relatedness_weighted,
                genetic_relatedness_weighted,
                centre=centre,
                polarised=polarised,
            )


class TestBranchGeneticRelatednessWeighted(
    TestGeneticRelatednessWeighted, TopologyExamplesMixin
):
    mode = "branch"


class TestNodeGeneticRelatednessWeighted(
    TestGeneticRelatednessWeighted, TopologyExamplesMixin
):
    mode = "node"


class TestSiteGeneticRelatednessWeighted(
    TestGeneticRelatednessWeighted, MutatedTopologyExamplesMixin
):
    mode = "site"


# NOTE: these classes don't follow the same (anti)-patterns as used elsewhere as they
# were added in several years afterwards.


class TestGeneticRelatednessWeightedSimpleExamples:
    # Values verified against the simple implementations above
    site_value = 22.24
    branch_value = 29.44

    def fixture(self):
        ts = tskit.Tree.generate_balanced(5).tree_sequence
        # Abitrary weights that give non-zero results
        W = np.zeros((ts.num_samples, 2))
        W[0, :] = 1
        W[1, :] = 2
        return tsutil.insert_branch_sites(ts), W

    def test_no_arguments_site(self):
        ts, W = self.fixture()
        X = ts.genetic_relatedness_weighted(W, mode="site")
        assert X.shape == tuple()
        nt.assert_almost_equal(X, self.site_value)

    def test_windows_site(self):
        ts, W = self.fixture()
        X = ts.genetic_relatedness_weighted(W, mode="site", windows=[0, 1 - 1e-12, 1])
        assert X.shape == (2,)
        nt.assert_almost_equal(X[0], self.site_value)
        nt.assert_almost_equal(X[1], 0)

    def test_no_arguments_branch(self):
        ts, W = self.fixture()
        X = ts.genetic_relatedness_weighted(W, mode="branch")
        assert X.shape == tuple()
        nt.assert_almost_equal(X, self.branch_value)

    def test_windows_branch(self):
        ts, W = self.fixture()
        X = ts.genetic_relatedness_weighted(W, mode="branch", windows=[0, 0.5, 1])
        assert X.shape == (2,)
        nt.assert_almost_equal(X, self.branch_value)

    def test_indexes_1D(self):
        ts, W = self.fixture()
        indexes = [0, 1]
        X = ts.genetic_relatedness_weighted(W, indexes, mode="branch")
        assert X.shape == tuple()
        nt.assert_almost_equal(X, self.branch_value)

    def test_indexes_2D(self):
        ts, W = self.fixture()
        indexes = [[0, 1]]
        X = ts.genetic_relatedness_weighted(W, indexes, mode="branch")
        assert X.shape == (1,)
        nt.assert_almost_equal(X, self.branch_value)

    def test_indexes_2D_windows(self):
        ts, W = self.fixture()
        indexes = [[0, 1], [0, 1]]
        X = ts.genetic_relatedness_weighted(
            W, indexes, windows=[0, 0.5, 1], mode="branch"
        )
        assert X.shape == (2, 2)
        nt.assert_almost_equal(X, self.branch_value)


class TestGeneticRelatednessWeightedErrors:
    def ts(self):
        return tskit.Tree.generate_balanced(3).tree_sequence

    @pytest.mark.parametrize("W", [[0], np.array([0]), np.zeros(100)])
    def test_bad_weight_size(self, W):
        with pytest.raises(ValueError, match="First trait dimension"):
            self.ts().genetic_relatedness_weighted(W)

    @pytest.mark.parametrize("cols", [1, 3])
    def test_no_indexes_with_non_2_cols(self, cols):
        ts = self.ts()
        W = np.zeros((ts.num_samples, cols))
        with pytest.raises(ValueError, match="Must specify indexes"):
            ts.genetic_relatedness_weighted(W)

    @pytest.mark.parametrize("indexes", [[], [[0]], [[0, 0, 0]], [[[0], [0], [0]]]])
    def test_bad_index_shapes(self, indexes):
        ts = self.ts()
        W = np.zeros((ts.num_samples, 2))
        with pytest.raises(ValueError, match="Indexes must be convertable to a 2D"):
            ts.genetic_relatedness_weighted(W, indexes=indexes)


############################################
# Fst
############################################


def single_site_Fst(ts, sample_sets, indexes):
    """
    Compute single-site Fst, which between two groups with frequencies p and q is
      1 - 2 * (p (1-p) + q(1-q)) / ( p(1-p) + q(1-q) + p(1-q) + q(1-p) )
    or in the multiallelic case, replacing p(1-p) with the sum over alleles of p(1-p),
    and adjusted for sampling without replacement.
    """
    # TODO: what to do in this case?
    if ts.num_sites == 0:
        out = np.array([np.repeat(np.nan, len(indexes))])
        return out
    out = np.zeros((ts.num_sites, len(indexes)))
    samples = ts.samples()
    # TODO deal with missing data properly.
    for j, v in enumerate(ts.variants(isolated_as_missing=False)):
        for i, (ix, iy) in enumerate(indexes):
            g = v.genotypes
            X = sample_sets[ix]
            Y = sample_sets[iy]
            gX = [a for k, a in zip(samples, g) if k in X]
            gY = [a for k, a in zip(samples, g) if k in Y]
            nX = len(X)
            nY = len(Y)
            dX = dY = dXY = 0
            for a in set(g):
                fX = np.sum(gX == a)
                fY = np.sum(gY == a)
                with suppress_division_by_zero_warning():
                    dX += fX * (nX - fX) / (nX * (nX - 1))
                    dY += fY * (nY - fY) / (nY * (nY - 1))
                    dXY += (fX * (nY - fY) + (nX - fX) * fY) / (2 * nX * nY)
            with suppress_division_by_zero_warning():
                out[j][i] = 1 - 2 * (dX + dY) / (dX + dY + 2 * dXY)
    return out


class TestFst(StatsTestCase, TwoWaySampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify(self, ts):
        # only check per-site
        for sample_sets in example_sample_sets(ts, min_size=2):
            for indexes in example_sample_set_index_pairs(sample_sets):
                self.verify_persite_Fst(ts, sample_sets, indexes)

    def verify_persite_Fst(self, ts, sample_sets, indexes):
        sigma1 = ts.Fst(
            sample_sets,
            indexes=indexes,
            windows="sites",
            mode=self.mode,
            span_normalise=False,
        )
        sigma2 = single_site_Fst(ts, sample_sets, indexes)
        assert sigma1.shape == sigma2.shape
        self.assertArrayAlmostEqual(sigma1, sigma2)


class FstInterfaceMixin:
    def test_interface(self):
        ts = msprime.simulate(10, mutation_rate=0.0)
        sample_sets = [[0, 1, 2], [6, 7], [4]]
        with pytest.raises(ValueError):
            ts.Fst(sample_sets, mode=self.mode)
        with pytest.raises(ValueError):
            ts.Fst(sample_sets, indexes=[(0, 1, 2), (3, 4, 5)], mode=self.mode)
        with pytest.raises(tskit.LibraryError):
            ts.Fst(sample_sets, indexes=[(0, 1), (0, 20)])
        sigma1 = ts.Fst(sample_sets, indexes=[(0, 1)], mode=self.mode)
        sigma2 = ts.Fst(sample_sets, indexes=[(0, 1), (0, 2), (1, 2)], mode=self.mode)
        self.assertArrayAlmostEqual(sigma1[..., 0], sigma2[..., 0])


class TestSiteFst(TestFst, MutatedTopologyExamplesMixin, FstInterfaceMixin):
    mode = "site"


# Since Fst is defined using diversity and divergence, we don't seriously
# test it for correctness for node and branch, and only test the interface.


class TestNodeFst(StatsTestCase, FstInterfaceMixin):
    mode = "node"


class TestBranchFst(StatsTestCase, FstInterfaceMixin):
    mode = "node"


############################################
# Y2
############################################


def branch_Y2(ts, sample_sets, indexes, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(indexes)))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, (ix, iy) in enumerate(indexes):
            X = sample_sets[ix]
            Y = sample_sets[iy]
            denom = np.float64(len(X) * len(Y) * (len(Y) - 1))
            has_trees = False
            S = 0
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                if tr.total_branch_length > 0:
                    has_trees = True
                this_length = min(end, tr.interval.right) - max(begin, tr.interval.left)
                for x in X:
                    for y in Y:
                        for z in set(Y) - {y}:
                            xy_mrca = tr.mrca(x, y)
                            xz_mrca = tr.mrca(x, z)
                            yz_mrca = tr.mrca(y, z)
                            if xy_mrca == xz_mrca:
                                #   /\
                                #  / /\
                                # x y  z
                                S += path_length(tr, x, yz_mrca) * this_length
                            elif xy_mrca == yz_mrca:
                                #   /\
                                #  / /\
                                # y x  z
                                S += path_length(tr, x, xz_mrca) * this_length
                            elif xz_mrca == yz_mrca:
                                #   /\
                                #  / /\
                                # z x  y
                                S += path_length(tr, x, xy_mrca) * this_length
            if has_trees:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def site_Y2(ts, sample_sets, indexes, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    samples = ts.samples()
    out = np.zeros((len(windows) - 1, len(indexes)))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        haps = ts.genotype_matrix(isolated_as_missing=False).T
        site_positions = [x.position for x in ts.sites()]
        for i, (ix, iy) in enumerate(indexes):
            X = sample_sets[ix]
            Y = sample_sets[iy]
            denom = np.float64(len(X) * len(Y) * (len(Y) - 1))
            S = 0
            site_in_window = False
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    site_in_window = True
                    for x in X:
                        x_index = np.where(samples == x)[0][0]
                        for y in Y:
                            y_index = np.where(samples == y)[0][0]
                            for z in set(Y) - {y}:
                                z_index = np.where(samples == z)[0][0]
                                condition = (
                                    haps[x_index, k] != haps[y_index, k]
                                    and haps[x_index, k] != haps[z_index, k]
                                )
                                if condition:
                                    # x|yz
                                    S += 1
            if site_in_window:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def node_Y2(ts, sample_sets, indexes, windows=None, span_normalise=True):
    out = np.zeros((len(windows) - 1, ts.num_nodes, len(indexes)))
    for i, (ix, iy) in enumerate(indexes):
        X = sample_sets[ix]
        Y = sample_sets[iy]
        tX = len(X)
        tY = len(Y)
        denom = np.float64(tX * tY * (tY - 1))
        for j in range(len(windows) - 1):
            begin = windows[j]
            end = windows[j + 1]
            S = np.zeros(ts.num_nodes)
            for t1, t2 in zip(ts.trees(tracked_samples=X), ts.trees(tracked_samples=Y)):
                if t1.interval.right <= begin:
                    continue
                if t1.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in t1.nodes():
                    # count number of pairwise paths going through u
                    nX = t1.num_tracked_samples(u)
                    nY = t2.num_tracked_samples(u)
                    SS[u] += nX * (tY - nY) * (tY - nY - 1) + (tX - nX) * nY * (nY - 1)
                S += SS * (min(end, t1.interval.right) - max(begin, t1.interval.left))
            with suppress_division_by_zero_warning():
                out[j, :, i] = S / denom
            if span_normalise:
                out[j, :, i] /= end - begin
    return out


def Y2(ts, sample_sets, indexes=None, windows=None, mode="site", span_normalise=True):
    windows = ts.parse_windows(windows)
    if indexes is None:
        indexes = [(0, 1)]
    method_map = {"site": site_Y2, "node": node_Y2, "branch": branch_Y2}
    return method_map[mode](
        ts, sample_sets, indexes=indexes, windows=windows, span_normalise=span_normalise
    )


class TestY2(StatsTestCase, TwoWaySampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_sample_sets_indexes(self, ts, sample_sets, indexes, windows):
        n = np.array([len(x) for x in sample_sets])

        denom = np.array([n[i] * n[j] * (n[j] - 1) for i, j in indexes])

        def f(x):
            numer = np.array(
                [(x[i] * (n[j] - x[j]) * (n[j] - x[j] - 1)) for i, j in indexes]
            )
            return numer / denom

        self.verify_definition(ts, sample_sets, indexes, windows, f, ts.Y2, Y2)


class TestBranchY2(TestY2, TopologyExamplesMixin):
    mode = "branch"


class TestNodeY2(TestY2, TopologyExamplesMixin):
    mode = "node"


class TestSiteY2(TestY2, MutatedTopologyExamplesMixin):
    mode = "site"


############################################
# Y3
############################################


def branch_Y3(ts, sample_sets, indexes, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(indexes)))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, (ix, iy, iz) in enumerate(indexes):
            S = 0
            X = sample_sets[ix]
            Y = sample_sets[iy]
            Z = sample_sets[iz]
            denom = np.float64(len(X) * len(Y) * len(Z))
            has_trees = False
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                if tr.total_branch_length > 0:
                    has_trees = True
                this_length = min(end, tr.interval.right) - max(begin, tr.interval.left)
                for x in X:
                    for y in Y:
                        for z in Z:
                            xy_mrca = tr.mrca(x, y)
                            xz_mrca = tr.mrca(x, z)
                            yz_mrca = tr.mrca(y, z)
                            if xy_mrca == xz_mrca:
                                #   /\
                                #  / /\
                                # x y  z
                                S += path_length(tr, x, yz_mrca) * this_length
                            elif xy_mrca == yz_mrca:
                                #   /\
                                #  / /\
                                # y x  z
                                S += path_length(tr, x, xz_mrca) * this_length
                            elif xz_mrca == yz_mrca:
                                #   /\
                                #  / /\
                                # z x  y
                                S += path_length(tr, x, xy_mrca) * this_length
            if has_trees:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def site_Y3(ts, sample_sets, indexes, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(indexes)))
    haps = ts.genotype_matrix(isolated_as_missing=False).T
    site_positions = ts.tables.sites.position
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, (ix, iy, iz) in enumerate(indexes):
            X = sample_sets[ix]
            Y = sample_sets[iy]
            Z = sample_sets[iz]
            denom = np.float64(len(X) * len(Y) * len(Z))
            S = 0
            site_in_window = False
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    site_in_window = True
                    for x in X:
                        x_index = np.where(samples == x)[0][0]
                        for y in Y:
                            y_index = np.where(samples == y)[0][0]
                            for z in Z:
                                z_index = np.where(samples == z)[0][0]
                                if (haps[x_index][k] != haps[y_index][k]) and (
                                    haps[x_index][k] != haps[z_index][k]
                                ):
                                    # x|yz
                                    with suppress_division_by_zero_warning():
                                        S += 1
            if site_in_window:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def node_Y3(ts, sample_sets, indexes, windows=None, span_normalise=True):
    out = np.zeros((len(windows) - 1, ts.num_nodes, len(indexes)))
    for i, (ix, iy, iz) in enumerate(indexes):
        X = sample_sets[ix]
        Y = sample_sets[iy]
        Z = sample_sets[iz]
        tX = len(X)
        tY = len(Y)
        tZ = len(Z)
        denom = np.float64(tX * tY * tZ)
        for j in range(len(windows) - 1):
            begin = windows[j]
            end = windows[j + 1]
            S = np.zeros(ts.num_nodes)
            for t1, t2, t3 in zip(
                ts.trees(tracked_samples=X),
                ts.trees(tracked_samples=Y),
                ts.trees(tracked_samples=Z),
            ):
                if t1.interval.right <= begin:
                    continue
                if t1.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in t1.nodes():
                    # count number of pairwise paths going through u
                    nX = t1.num_tracked_samples(u)
                    nY = t2.num_tracked_samples(u)
                    nZ = t3.num_tracked_samples(u)
                    SS[u] += nX * (tY - nY) * (tZ - nZ) + (tX - nX) * nY * nZ
                S += SS * (min(end, t1.interval.right) - max(begin, t1.interval.left))
            with suppress_division_by_zero_warning():
                out[j, :, i] = S / denom
            if span_normalise:
                out[j, :, i] /= end - begin
    return out


def Y3(ts, sample_sets, indexes=None, windows=None, mode="site", span_normalise=True):
    windows = ts.parse_windows(windows)
    if indexes is None:
        indexes = [(0, 1, 2)]
    method_map = {"site": site_Y3, "node": node_Y3, "branch": branch_Y3}
    return method_map[mode](
        ts, sample_sets, indexes=indexes, windows=windows, span_normalise=span_normalise
    )


class TestY3(StatsTestCase, ThreeWaySampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_sample_sets_indexes(self, ts, sample_sets, indexes, windows):
        n = np.array([len(x) for x in sample_sets])
        denom = np.array([n[i] * n[j] * n[k] for i, j, k in indexes])

        def f(x):
            numer = np.array(
                [x[i] * (n[j] - x[j]) * (n[k] - x[k]) for i, j, k in indexes]
            )
            return numer / denom

        self.verify_definition(ts, sample_sets, indexes, windows, f, ts.Y3, Y3)


class TestBranchY3(TestY3, TopologyExamplesMixin):
    mode = "branch"


class TestNodeY3(TestY3, TopologyExamplesMixin):
    mode = "node"


class TestSiteY3(TestY3, MutatedTopologyExamplesMixin):
    mode = "site"


############################################
# f2
############################################


def branch_f2(ts, sample_sets, indexes, windows=None, span_normalise=True):
    # this is f4(A,B;A,B) but drawing distinct samples from A and B
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(indexes)))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, (ia, ib) in enumerate(indexes):
            A = sample_sets[ia]
            B = sample_sets[ib]
            denom = np.float64(len(A) * (len(A) - 1) * len(B) * (len(B) - 1))
            has_trees = False
            S = 0
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                if tr.total_branch_length > 0:
                    has_trees = True
                this_length = min(end, tr.interval.right) - max(begin, tr.interval.left)
                SS = 0
                for a in A:
                    for b in B:
                        for c in set(A) - {a}:
                            for d in set(B) - {b}:
                                with suppress_division_by_zero_warning():
                                    SS += path_length(tr, tr.mrca(a, c), tr.mrca(b, d))
                                    SS -= path_length(tr, tr.mrca(a, d), tr.mrca(b, c))
                S += SS * this_length
            if has_trees:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def site_f2(ts, sample_sets, indexes, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(indexes)))
    samples = ts.samples()
    haps = ts.genotype_matrix(isolated_as_missing=False).T
    site_positions = ts.tables.sites.position
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, (iA, iB) in enumerate(indexes):
            A = sample_sets[iA]
            B = sample_sets[iB]
            denom = np.float64(len(A) * (len(A) - 1) * len(B) * (len(B) - 1))
            S = 0
            site_in_window = False
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    site_in_window = True
                    for a in A:
                        a_index = np.where(samples == a)[0][0]
                        for b in B:
                            b_index = np.where(samples == b)[0][0]
                            for c in set(A) - {a}:
                                c_index = np.where(samples == c)[0][0]
                                for d in set(B) - {b}:
                                    d_index = np.where(samples == d)[0][0]
                                    if (
                                        (haps[a_index][k] == haps[c_index][k])
                                        and (haps[a_index][k] != haps[d_index][k])
                                        and (haps[a_index][k] != haps[b_index][k])
                                    ):
                                        # ac|bd
                                        S += 1
                                    elif (
                                        (haps[a_index][k] == haps[d_index][k])
                                        and (haps[a_index][k] != haps[c_index][k])
                                        and (haps[a_index][k] != haps[b_index][k])
                                    ):
                                        # ad|bc
                                        S -= 1
            if site_in_window:
                with np.errstate(invalid="ignore", divide="ignore"):
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def node_f2(ts, sample_sets, indexes, windows=None, span_normalise=True):
    out = np.zeros((len(windows) - 1, ts.num_nodes, len(indexes)))
    for i, (ia, ib) in enumerate(indexes):
        A = sample_sets[ia]
        B = sample_sets[ib]
        tA = len(A)
        tB = len(B)
        denom = np.float64(tA * (tA - 1) * tB * (tB - 1))
        for j in range(len(windows) - 1):
            begin = windows[j]
            end = windows[j + 1]
            S = np.zeros(ts.num_nodes)
            for t1, t2 in zip(ts.trees(tracked_samples=A), ts.trees(tracked_samples=B)):
                if t1.interval.right <= begin:
                    continue
                if t1.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in t1.nodes():
                    # count number of pairwise paths going through u
                    nA = t1.num_tracked_samples(u)
                    nB = t2.num_tracked_samples(u)
                    # xy|uv - xv|uy with x,y in A, u, v in B
                    SS[u] += nA * (nA - 1) * (tB - nB) * (tB - nB - 1) + (tA - nA) * (
                        tA - nA - 1
                    ) * nB * (nB - 1)
                    SS[u] -= 2 * nA * nB * (tA - nA) * (tB - nB)
                S += SS * (min(end, t1.interval.right) - max(begin, t1.interval.left))
            with suppress_division_by_zero_warning():
                out[j, :, i] = S / denom
            if span_normalise:
                out[j, :, i] /= end - begin
    return out


def f2(ts, sample_sets, indexes=None, windows=None, mode="site", span_normalise=True):
    """
    Patterson's f2 statistic definitions.
    """
    windows = ts.parse_windows(windows)
    if indexes is None:
        indexes = [(0, 1)]
    method_map = {"site": site_f2, "node": node_f2, "branch": branch_f2}
    return method_map[mode](
        ts, sample_sets, indexes=indexes, windows=windows, span_normalise=span_normalise
    )


class Testf2(StatsTestCase, TwoWaySampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_sample_sets_indexes(self, ts, sample_sets, indexes, windows):
        n = np.array([len(x) for x in sample_sets])

        denom = np.array([n[i] * (n[i] - 1) * n[j] * (n[j] - 1) for i, j in indexes])

        def f(x):
            numer = np.array(
                [
                    x[i] * (x[i] - 1) * (n[j] - x[j]) * (n[j] - x[j] - 1)
                    - x[i] * (n[i] - x[i]) * (n[j] - x[j]) * x[j]
                    for i, j in indexes
                ]
            )
            return numer / denom

        self.verify_definition(ts, sample_sets, indexes, windows, f, ts.f2, f2)


class TestBranchf2(Testf2, TopologyExamplesMixin):
    mode = "branch"


class TestNodef2(Testf2, TopologyExamplesMixin):
    mode = "node"


class TestSitef2(Testf2, MutatedTopologyExamplesMixin):
    mode = "site"


############################################
# f3
############################################


def branch_f3(ts, sample_sets, indexes, windows=None, span_normalise=True):
    # this is f4(A,B;A,C) but drawing distinct samples from A
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(indexes)))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, (ia, ib, ic) in enumerate(indexes):
            A = sample_sets[ia]
            B = sample_sets[ib]
            C = sample_sets[ic]
            denom = np.float64(len(A) * (len(A) - 1) * len(B) * len(C))
            has_trees = False
            S = 0
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                if tr.total_branch_length > 0:
                    has_trees = True
                this_length = min(end, tr.interval.right) - max(begin, tr.interval.left)
                SS = 0
                for a in A:
                    for b in B:
                        for c in set(A) - {a}:
                            for d in C:
                                SS += path_length(tr, tr.mrca(a, c), tr.mrca(b, d))
                                SS -= path_length(tr, tr.mrca(a, d), tr.mrca(b, c))
                S += SS * this_length
            if has_trees:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def site_f3(ts, sample_sets, indexes, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(indexes)))
    samples = ts.samples()
    haps = ts.genotype_matrix(isolated_as_missing=False).T
    site_positions = ts.tables.sites.position
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, (iA, iB, iC) in enumerate(indexes):
            A = sample_sets[iA]
            B = sample_sets[iB]
            C = sample_sets[iC]
            denom = np.float64(len(A) * (len(A) - 1) * len(B) * len(C))
            S = 0
            site_in_window = False
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    site_in_window = True
                    for a in A:
                        a_index = np.where(samples == a)[0][0]
                        for b in B:
                            b_index = np.where(samples == b)[0][0]
                            for c in set(A) - {a}:
                                c_index = np.where(samples == c)[0][0]
                                for d in C:
                                    d_index = np.where(samples == d)[0][0]
                                    if (
                                        (haps[a_index][k] == haps[c_index][k])
                                        and (haps[a_index][k] != haps[d_index][k])
                                        and (haps[a_index][k] != haps[b_index][k])
                                    ):
                                        # ac|bd
                                        S += 1
                                    elif (
                                        (haps[a_index][k] == haps[d_index][k])
                                        and (haps[a_index][k] != haps[c_index][k])
                                        and (haps[a_index][k] != haps[b_index][k])
                                    ):
                                        # ad|bc
                                        S -= 1
            if site_in_window:
                with np.errstate(invalid="ignore", divide="ignore"):
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def node_f3(ts, sample_sets, indexes, windows=None, span_normalise=True):
    out = np.zeros((len(windows) - 1, ts.num_nodes, len(indexes)))
    for i, (iA, iB, iC) in enumerate(indexes):
        A = sample_sets[iA]
        B = sample_sets[iB]
        C = sample_sets[iC]
        tA = len(A)
        tB = len(B)
        tC = len(C)
        denom = np.float64(tA * (tA - 1) * tB * tC)
        for j in range(len(windows) - 1):
            begin = windows[j]
            end = windows[j + 1]
            S = np.zeros(ts.num_nodes)
            for t1, t2, t3 in zip(
                ts.trees(tracked_samples=A),
                ts.trees(tracked_samples=B),
                ts.trees(tracked_samples=C),
            ):
                if t1.interval.right <= begin:
                    continue
                if t1.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in t1.nodes():
                    # count number of pairwise paths going through u
                    nA = t1.num_tracked_samples(u)
                    nB = t2.num_tracked_samples(u)
                    nC = t3.num_tracked_samples(u)
                    # xy|uv - xv|uy with x,y in A, u in B and v in C
                    SS[u] += (
                        nA * (nA - 1) * (tB - nB) * (tC - nC)
                        + (tA - nA) * (tA - nA - 1) * nB * nC
                    )
                    SS[u] -= (
                        nA * nC * (tA - nA) * (tB - nB)
                        + (tA - nA) * (tC - nC) * nA * nB
                    )
                S += SS * (min(end, t1.interval.right) - max(begin, t1.interval.left))
            with suppress_division_by_zero_warning():
                out[j, :, i] = S / denom
            if span_normalise:
                out[j, :, i] /= end - begin
    return out


def f3(ts, sample_sets, indexes=None, windows=None, mode="site", span_normalise=True):
    """
    Patterson's f3 statistic definitions.
    """
    windows = ts.parse_windows(windows)
    if indexes is None:
        indexes = [(0, 1, 2)]
    method_map = {"site": site_f3, "node": node_f3, "branch": branch_f3}
    return method_map[mode](
        ts, sample_sets, indexes=indexes, windows=windows, span_normalise=span_normalise
    )


class Testf3(StatsTestCase, ThreeWaySampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_sample_sets_indexes(self, ts, sample_sets, indexes, windows):
        n = np.array([len(x) for x in sample_sets])
        denom = np.array([n[i] * (n[i] - 1) * n[j] * n[k] for i, j, k in indexes])

        def f(x):
            numer = np.array(
                [
                    x[i] * (x[i] - 1) * (n[j] - x[j]) * (n[k] - x[k])
                    - x[i] * (n[i] - x[i]) * (n[j] - x[j]) * x[k]
                    for i, j, k in indexes
                ]
            )
            return numer / denom

        self.verify_definition(ts, sample_sets, indexes, windows, f, ts.f3, f3)


class TestBranchf3(Testf3, TopologyExamplesMixin):
    mode = "branch"


class TestNodef3(Testf3, TopologyExamplesMixin):
    mode = "node"


class TestSitef3(Testf3, MutatedTopologyExamplesMixin):
    mode = "site"


############################################
# f4
############################################


def branch_f4(ts, sample_sets, indexes, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(indexes)))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, (iA, iB, iC, iD) in enumerate(indexes):
            A = sample_sets[iA]
            B = sample_sets[iB]
            C = sample_sets[iC]
            D = sample_sets[iD]
            denom = np.float64(len(A) * len(B) * len(C) * len(D))
            has_trees = False
            S = 0
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                if tr.total_branch_length > 0:
                    has_trees = True
                this_length = min(end, tr.interval.right) - max(begin, tr.interval.left)
                SS = 0
                for a in A:
                    for b in B:
                        for c in C:
                            for d in D:
                                with suppress_division_by_zero_warning():
                                    SS += path_length(tr, tr.mrca(a, c), tr.mrca(b, d))
                                    SS -= path_length(tr, tr.mrca(a, d), tr.mrca(b, c))
                S += SS * this_length
            if has_trees:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def site_f4(ts, sample_sets, indexes, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    samples = ts.samples()
    haps = ts.genotype_matrix(isolated_as_missing=False).T
    site_positions = ts.tables.sites.position
    out = np.zeros((len(windows) - 1, len(indexes)))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, (iA, iB, iC, iD) in enumerate(indexes):
            A = sample_sets[iA]
            B = sample_sets[iB]
            C = sample_sets[iC]
            D = sample_sets[iD]
            denom = np.float64(len(A) * len(B) * len(C) * len(D))
            S = 0
            site_in_window = False
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    site_in_window = True
                    for a in A:
                        a_index = np.where(samples == a)[0][0]
                        for b in B:
                            b_index = np.where(samples == b)[0][0]
                            for c in C:
                                c_index = np.where(samples == c)[0][0]
                                for d in D:
                                    d_index = np.where(samples == d)[0][0]
                                    if (
                                        (haps[a_index][k] == haps[c_index][k])
                                        and (haps[a_index][k] != haps[d_index][k])
                                        and (haps[a_index][k] != haps[b_index][k])
                                    ):
                                        # ac|bd
                                        S += 1
                                    elif (
                                        (haps[a_index][k] == haps[d_index][k])
                                        and (haps[a_index][k] != haps[c_index][k])
                                        and (haps[a_index][k] != haps[b_index][k])
                                    ):
                                        # ad|bc
                                        S -= 1
            if site_in_window:
                with np.errstate(invalid="ignore", divide="ignore"):
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def node_f4(ts, sample_sets, indexes, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, ts.num_nodes, len(indexes)))
    for i, (iA, iB, iC, iD) in enumerate(indexes):
        A = sample_sets[iA]
        B = sample_sets[iB]
        C = sample_sets[iC]
        D = sample_sets[iD]
        tA = len(A)
        tB = len(B)
        tC = len(C)
        tD = len(D)
        denom = np.float64(tA * tB * tC * tD)
        for j in range(len(windows) - 1):
            begin = windows[j]
            end = windows[j + 1]
            S = np.zeros(ts.num_nodes)
            for t1, t2, t3, t4 in zip(
                ts.trees(tracked_samples=A),
                ts.trees(tracked_samples=B),
                ts.trees(tracked_samples=C),
                ts.trees(tracked_samples=D),
            ):
                if t1.interval.right <= begin:
                    continue
                if t1.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in t1.nodes():
                    # count number of pairwise paths going through u
                    nA = t1.num_tracked_samples(u)
                    nB = t2.num_tracked_samples(u)
                    nC = t3.num_tracked_samples(u)
                    nD = t4.num_tracked_samples(u)
                    # ac|bd - ad|bc
                    SS[u] += (
                        nA * nC * (tB - nB) * (tD - nD)
                        + (tA - nA) * (tC - nC) * nB * nD
                    )
                    SS[u] -= (
                        nA * nD * (tB - nB) * (tC - nC)
                        + (tA - nA) * (tD - nD) * nB * nC
                    )
                S += SS * (min(end, t1.interval.right) - max(begin, t1.interval.left))
            with suppress_division_by_zero_warning():
                out[j, :, i] = S / denom
            if span_normalise:
                out[j, :, i] /= end - begin
    return out


def f4(ts, sample_sets, indexes=None, windows=None, mode="site", span_normalise=True):
    """
    Patterson's f4 statistic definitions.
    """
    if indexes is None:
        indexes = [(0, 1, 2, 3)]
    method_map = {"site": site_f4, "node": node_f4, "branch": branch_f4}
    return method_map[mode](
        ts, sample_sets, indexes=indexes, windows=windows, span_normalise=span_normalise
    )


class Testf4(StatsTestCase, FourWaySampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_sample_sets_indexes(self, ts, sample_sets, indexes, windows):
        n = np.array([len(x) for x in sample_sets])
        denom = np.array([n[i] * n[j] * n[k] * n[l] for i, j, k, l in indexes])

        def f(x):
            numer = np.array(
                [
                    x[i] * x[k] * (n[j] - x[j]) * (n[l] - x[l])
                    - x[i] * x[l] * (n[j] - x[j]) * (n[k] - x[k])
                    for i, j, k, l in indexes
                ]
            )
            return numer / denom

        self.verify_definition(ts, sample_sets, indexes, windows, f, ts.f4, f4)


class TestBranchf4(Testf4, TopologyExamplesMixin):
    mode = "branch"


class TestNodef4(Testf4, TopologyExamplesMixin):
    mode = "node"


class TestSitef4(Testf4, MutatedTopologyExamplesMixin):
    mode = "site"


############################################
# Allele frequency spectrum
############################################


def fold(x, dims):
    """
    Folds the specified coordinates.
    """
    x = np.array(x, dtype=int)
    dims = np.array(dims, dtype=int)
    k = len(dims)
    n = np.sum(dims - 1) / 2
    s = np.sum(x)
    while s == n and k > 0:
        k -= 1
        assert k >= 0
        n -= (dims[k] - 1) / 2
        s -= x[k]
    if s > n:
        x = dims - 1 - x
    assert np.all(x >= 0)
    return tuple(x)


def foldit(A):
    B = np.zeros(A.shape)
    dims = A.shape
    inds = [range(k) for k in dims]
    for ij in itertools.product(*inds):
        nij = fold(ij, dims)
        B[nij] += A[ij]
    return B


def fold_windowed(X):
    Y = np.zeros(X.shape)
    for k in range(X.shape[0]):
        Y[k] = foldit(X[k])
    return Y


class TestFold:
    """
    Tests for the fold operation used in the AFS.
    """

    def test_examples(self):
        A = np.arange(12)
        Af = np.array(
            [11.0, 11.0, 11.0, 11.0, 11.0, 11.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        )

        assert np.all(foldit(A) == Af)

        B = A.copy().reshape(3, 4)
        Bf = np.array(
            [[11.0, 11.0, 11.0, 0.0], [11.0, 11.0, 0.0, 0.0], [11.0, 0.0, 0.0, 0.0]]
        )
        assert np.all(foldit(B) == Bf)

        C = A.copy().reshape(3, 2, 2)
        Cf = np.array(
            [
                [[11.0, 11.0], [11.0, 11.0]],
                [[11.0, 11.0], [0.0, 0.0]],
                [[0.0, 0.0], [0.0, 0.0]],
            ]
        )
        assert np.all(foldit(C) == Cf)

        D = np.arange(9).reshape((3, 3))
        Df = np.array([[8.0, 8.0, 8.0], [8.0, 4.0, 0.0], [0.0, 0.0, 0.0]])
        assert np.all(foldit(D) == Df)

        E = np.arange(9)
        Ef = np.array([8.0, 8.0, 8.0, 8.0, 4.0, 0.0, 0.0, 0.0, 0.0])
        assert np.all(foldit(E) == Ef)

    def test_branch_folded(self):
        ts = msprime.sim_ancestry(10, random_seed=1, sequence_length=10)
        folded = ts.allele_frequency_spectrum(
            windows=[0, 5, 8, 9, 10], mode="branch", polarised=False
        )
        unfolded = ts.allele_frequency_spectrum(
            windows=[0, 5, 8, 9, 10], mode="branch", polarised=True
        )
        assert np.allclose(fold_windowed(unfolded), folded)

    def test_site_folded(self):
        ts = msprime.sim_ancestry(10, random_seed=1, sequence_length=10)
        ts = msprime.sim_mutations(ts, rate=1, random_seed=1, discrete_genome=False)
        for s in ts.sites():
            assert len(s.mutations) == 1
        folded = ts.allele_frequency_spectrum(
            windows=[0, 5, 8, 9, 10], mode="site", polarised=False, span_normalise=False
        )
        unfolded = ts.allele_frequency_spectrum(
            windows=[0, 5, 8, 9, 10], mode="site", polarised=True, span_normalise=False
        )
        assert np.allclose(fold_windowed(unfolded), folded)


def naive_site_allele_frequency_spectrum(
    ts, sample_sets, windows=None, polarised=False, span_normalise=True
):
    """
    The joint allele frequency spectrum for sites.
    """
    windows = ts.parse_windows(windows)
    num_windows = len(windows) - 1
    out_dim = [1 + len(sample_set) for sample_set in sample_sets]
    out = np.zeros([num_windows] + out_dim)
    G = ts.genotype_matrix(isolated_as_missing=False)
    samples = ts.samples()
    # Indexes of the samples within the sample sets into the samples array.
    sample_set_indexes = [
        np.array([np.where(x == samples)[0][0] for x in sample_set])
        for sample_set in sample_sets
    ]
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for site in ts.sites():
            S = np.zeros(out_dim)
            if begin <= site.position < end:
                g = G[site.id]
                alleles = np.unique(g)

                # Any site monomorphic across all samples does not contribute
                if len(alleles) == 1:
                    continue

                # For each allele, count the number present in each sample set.
                count = {
                    allele: np.zeros(len(sample_sets), dtype=int) for allele in alleles
                }
                for k, sample_set in enumerate(sample_set_indexes):
                    allele_counts = zip(*np.unique(g[sample_set], return_counts=True))
                    for allele, c in allele_counts:
                        count[allele][k] = c
                increment = 0.5
                if polarised:
                    increment = 1
                    # Remove the contribution of the ancestral state
                    if 0 in count:
                        del count[0]
                for allele_count in count.values():
                    x = tuple(allele_count)
                    if not polarised:
                        x = fold(x, out_dim)
                    S[x] += increment
            if span_normalise:
                S /= end - begin
            out[j, :] += S
    return out


def naive_branch_allele_frequency_spectrum(
    ts, sample_sets, windows=None, polarised=False, span_normalise=True
):
    """
    The joint allele frequency spectrum for branches.
    """
    windows = ts.parse_windows(windows)
    num_windows = len(windows) - 1
    out_dim = [1 + len(sample_set) for sample_set in sample_sets]
    out = np.zeros([num_windows] + out_dim)
    for j in range(num_windows):
        begin = windows[j]
        end = windows[j + 1]
        S = np.zeros(out_dim)
        trees = [
            next(ts.trees(tracked_samples=sample_set)) for sample_set in sample_sets
        ]
        t = trees[0]
        while True:
            tr_len = min(end, t.interval.right) - max(begin, t.interval.left)
            if tr_len > 0:
                for node in t.nodes():
                    if 0 < t.num_samples(node) < ts.num_samples:
                        x = [tree.num_tracked_samples(node) for tree in trees]
                        # Note x must be a tuple for indexing to work
                        if not polarised:
                            x = fold(x, out_dim)
                        S[tuple(x)] += t.branch_length(node) * tr_len

            # Advance the trees
            more = [tree.next() for tree in trees]
            assert len(set(more)) == 1
            if not more[0]:
                break
        if span_normalise:
            S /= end - begin
        out[j, :] = S
    return out


def naive_allele_frequency_spectrum(
    ts, sample_sets, windows=None, polarised=False, mode="site", span_normalise=True
):
    """
    Naive definition of the generalised site frequency spectrum.
    """
    method_map = {
        "site": naive_site_allele_frequency_spectrum,
        "branch": naive_branch_allele_frequency_spectrum,
    }
    return method_map[mode](
        ts,
        sample_sets,
        windows=windows,
        polarised=polarised,
        span_normalise=span_normalise,
    )


def branch_allele_frequency_spectrum(
    ts, sample_sets, windows, polarised=False, span_normalise=True
):
    """
    Efficient implementation of the algorithm used as the basis for the
    underlying C version.
    """
    num_sample_sets = len(sample_sets)
    windows = ts.parse_windows(windows)
    num_windows = windows.shape[0] - 1
    out_dim = [1 + len(sample_set) for sample_set in sample_sets]
    time = ts.tables.nodes.time

    result = np.zeros([num_windows] + out_dim)
    # Number of nodes in sample_set j ancestral to each node u.
    count = np.zeros((ts.num_nodes, num_sample_sets + 1), dtype=np.uint32)
    for j in range(num_sample_sets):
        count[sample_sets[j], j] = 1
    # The last column counts across all samples
    count[ts.samples(), -1] = 1
    # contains the location of the last time we updated the output for a node.
    last_update = np.zeros(ts.num_nodes)
    window_index = 0
    parent = np.zeros(ts.num_nodes, dtype=np.int32) - 1
    branch_length = np.zeros(ts.num_nodes)
    tree_index = 0

    def update_result(window_index, u, right):
        if 0 < count[u, -1] < ts.num_samples:
            x = (right - last_update[u]) * branch_length[u]
            c = count[u, :num_sample_sets]
            if not polarised:
                c = fold(c, out_dim)
            index = tuple([window_index] + list(c))
            result[index] += x
        last_update[u] = right

    for (t_left, t_right), edges_out, edges_in in ts.edge_diffs():
        for edge in edges_out:
            u = edge.child
            v = edge.parent
            update_result(window_index, u, t_left)
            while v != -1:
                update_result(window_index, v, t_left)
                count[v] -= count[u]
                v = parent[v]
            parent[u] = -1
            branch_length[u] = 0

        for edge in edges_in:
            u = edge.child
            v = edge.parent
            parent[u] = v
            branch_length[u] = time[v] - time[u]
            while v != -1:
                update_result(window_index, v, t_left)
                count[v] += count[u]
                v = parent[v]

        # Update the windows
        while window_index < num_windows and windows[window_index + 1] <= t_right:
            w_right = windows[window_index + 1]
            # This seems like a bad idea as we incur a O(N) cost for each window,
            # where N is the number of nodes.  It might be hard to do much better
            # though, since we can't help but incur O(|sample_set|) cost at each window
            # which we'll assume is O(n), and for large n, N isn't much larger than n.
            # For K > 1 dimensions, the cost of the scan through the nodes is much
            # less than the O(n^K) required to copy (if n is large and K is small).
            # We could keep track of the roots and do a tree traversal, bringing this
            # down to O(n), but this adds a lot of complexity and memory and I'm
            # fairly confident would be slower overall. We could keep a set of
            # non-zero branches, but this would add a O(log n) cost to each edge
            # insertion and removal and a lot of complexity to the C implementation.
            for u in range(ts.num_nodes):
                update_result(window_index, u, w_right)
            window_index += 1
        tree_index += 1

    assert window_index == windows.shape[0] - 1
    if span_normalise:
        for j in range(num_windows):
            result[j] /= windows[j + 1] - windows[j]
    return result


def site_allele_frequency_spectrum(
    ts, sample_sets, windows, polarised=False, span_normalise=True
):
    """
    Efficient implementation of the algorithm used as the basis for the
    underlying C version.
    """
    windows = ts.parse_windows(windows)
    num_windows = windows.shape[0] - 1
    out_dim = [1 + len(sample_set) for sample_set in sample_sets]

    result = np.zeros([num_windows] + out_dim)
    # Add an extra sample set to count across all samples
    sample_sets = list(sample_sets) + [ts.samples()]
    # Number of nodes in sample_set j ancestral to each node u.
    count = np.zeros((ts.num_nodes, len(sample_sets)), dtype=np.uint32)
    for j in range(len(sample_sets)):
        count[sample_sets[j], j] = 1

    site_index = 0
    mutation_index = 0
    window_index = 0
    sites = ts.tables.sites
    mutations = ts.tables.mutations
    parent = np.zeros(ts.num_nodes, dtype=np.int32) - 1
    for (t_left, t_right), edges_out, edges_in in ts.edge_diffs():
        for edge in edges_out:
            u = edge.child
            v = edge.parent
            while v != -1:
                count[v] -= count[u]
                v = parent[v]
            parent[u] = -1

        for edge in edges_in:
            u = edge.child
            v = edge.parent
            parent[u] = v
            while v != -1:
                count[v] += count[u]
                v = parent[v]

        while site_index < len(sites) and sites.position[site_index] < t_right:
            assert t_left <= sites.position[site_index]
            ancestral_state = sites[site_index].ancestral_state
            allele_count = collections.defaultdict(
                functools.partial(np.zeros, len(sample_sets), dtype=int)
            )
            allele_count[ancestral_state][:] = [
                len(sample_set) for sample_set in sample_sets
            ]
            while (
                mutation_index < len(mutations)
                and mutations[mutation_index].site == site_index
            ):
                mutation = mutations[mutation_index]
                allele_count[mutation.derived_state] += count[mutation.node]
                if mutation.parent != -1:
                    parent_allele = mutations[mutation.parent].derived_state
                    allele_count[parent_allele] -= count[mutation.node]
                else:
                    allele_count[ancestral_state] -= count[mutation.node]
                mutation_index += 1

            pos = sites.position[site_index]
            while windows[window_index + 1] <= pos:
                window_index += 1
            assert windows[window_index] <= pos < windows[window_index + 1]
            site_result = result[window_index]

            for allele, c in dict(allele_count).items():
                # Any allele monomorphic across all samples does not
                # contribute to the AFS
                if 0 == c[-1] or c[-1] == ts.num_samples:
                    del allele_count[allele]
            if polarised and ancestral_state in allele_count:
                del allele_count[ancestral_state]

            increment = 1 if polarised else 0.5
            for _allele, c in allele_count.items():
                x = tuple(c[:-1])
                if not polarised:
                    x = fold(x, out_dim)
                site_result[x] += increment
            site_index += 1

    if span_normalise:
        for j in range(num_windows):
            span = windows[j + 1] - windows[j]
            result[j] /= span
    return result


def allele_frequency_spectrum(
    ts, sample_sets, windows=None, polarised=False, mode="site", span_normalise=True
):
    """
    Generalised site frequency spectrum.
    """
    method_map = {
        "site": site_allele_frequency_spectrum,
        "branch": branch_allele_frequency_spectrum,
    }
    return method_map[mode](
        ts,
        sample_sets,
        windows=windows,
        polarised=polarised,
        span_normalise=span_normalise,
    )


class TestAlleleFrequencySpectrum(StatsTestCase, SampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_single_sample_set(self, ts):
        L = ts.sequence_length
        samples = ts.samples()
        a1 = ts.allele_frequency_spectrum(mode=self.mode)
        a2 = ts.allele_frequency_spectrum([samples], mode=self.mode)
        self.assertArrayEqual(a1, a2)
        for windows in [None, (0, L), (0, L / 2, L)]:
            a1 = ts.allele_frequency_spectrum(mode=self.mode, windows=windows)
            a2 = ts.allele_frequency_spectrum(
                [samples], mode=self.mode, windows=windows
            )
            self.assertArrayEqual(a1, a2)
        for polarised in [True, False]:
            a1 = ts.allele_frequency_spectrum(mode=self.mode, polarised=polarised)
            a2 = ts.allele_frequency_spectrum(
                [samples], mode=self.mode, polarised=polarised
            )
            self.assertArrayEqual(a1, a2)
        for span_normalise in [True, False]:
            a1 = ts.allele_frequency_spectrum(
                mode=self.mode, span_normalise=span_normalise
            )
            a2 = ts.allele_frequency_spectrum(
                [samples], mode=self.mode, span_normalise=span_normalise
            )
            self.assertArrayEqual(a1, a2)

    def verify_sample_sets(self, ts, sample_sets, windows):
        # print(ts.genotype_matrix())
        # print(ts.draw_text())
        # print("sample_sets = ", sample_sets)
        windows = ts.parse_windows(windows)
        for span_normalise, polarised in itertools.product(
            [True, False], [True, False]
        ):
            sfs1 = naive_allele_frequency_spectrum(
                ts,
                sample_sets,
                windows,
                mode=self.mode,
                polarised=polarised,
                span_normalise=span_normalise,
            )
            sfs2 = allele_frequency_spectrum(
                ts,
                sample_sets,
                windows,
                mode=self.mode,
                polarised=polarised,
                span_normalise=span_normalise,
            )
            sfs3 = ts.allele_frequency_spectrum(
                sample_sets,
                windows,
                mode=self.mode,
                polarised=polarised,
                span_normalise=span_normalise,
            )
            assert sfs1.shape[0] == len(windows) - 1
            assert len(sfs1.shape) == len(sample_sets) + 1
            for j, sample_set in enumerate(sample_sets):
                n = 1 + len(sample_set)
                assert sfs1.shape[j + 1] == n

            assert len(sfs1.shape) == len(sample_sets) + 1
            assert sfs1.shape == sfs2.shape
            assert sfs1.shape == sfs3.shape
            if not np.allclose(sfs1, sfs3):
                print()
                print("sample sets", sample_sets)
                print("simple", sfs1)
                print("effic ", sfs2)
                print("ts    ", sfs3)
            self.assertArrayAlmostEqual(sfs1, sfs2)
            self.assertArrayAlmostEqual(sfs1, sfs3)


class TestBranchAlleleFrequencySpectrum(
    TestAlleleFrequencySpectrum, TopologyExamplesMixin
):
    mode = "branch"

    def test_simple_example(self):
        ts = msprime.simulate(6, recombination_rate=0.1, random_seed=1)
        self.verify_single_sample_set(ts)

        self.verify_sample_sets(ts, [range(6)], [0, 1])
        self.verify_sample_sets(ts, [[0, 1]], [0, 1])
        self.verify_sample_sets(ts, [[0, 1], [2, 3]], [0, 1])
        self.verify_sample_sets(ts, [[0, 1, 2, 3, 4, 5]], [0, 1])
        self.verify_sample_sets(ts, [[0, 1, 2], [3, 4, 5]], [0, 1])
        self.verify_sample_sets(ts, [[0, 1], [2, 3], [4, 5]], [0, 1])


class TestSiteAlleleFrequencySpectrum(
    TestAlleleFrequencySpectrum, MutatedTopologyExamplesMixin
):
    mode = "site"

    def test_simple_example(self):
        ts = msprime.simulate(6, mutation_rate=0.2, random_seed=1)
        self.verify_single_sample_set(ts)

        self.verify_sample_sets(ts, [[0]], [0, 1])
        self.verify_sample_sets(ts, [[0, 1, 2, 3, 4, 5]], [0, 1])
        self.verify_sample_sets(ts, [[0, 1, 2], [3, 4, 5]], [0, 1])
        self.verify_sample_sets(ts, [[0, 1], [2, 3], [4, 5]], [0, 1])


class TestBranchAlleleFrequencySpectrumProperties(StatsTestCase, TopologyExamplesMixin):
    def verify(self, ts):
        # If we split by tree, the sum of the AFS should be equal to the
        # tree total branch length in each window
        windows = ts.breakpoints(as_array=True)
        S = ts.samples()
        examples = [
            [S],
            [S[:1]],
            [S[:-1]],
            [S[:1], S[1:]],
            [S[:1], S[:-1]],
        ]
        if len(S) > 2:
            examples += [[S[:1], S[2:], S[:3]]]
        # This is the same definition that we use for segregating_sites
        tbl = [
            sum(
                tree.branch_length(u)
                for u in tree.nodes()
                if 0 < tree.num_samples(u) < ts.num_samples
            )
            for tree in ts.trees()
        ]
        for polarised in [True, False]:
            for sample_sets in examples:
                afs = ts.allele_frequency_spectrum(
                    sample_sets,
                    windows=windows,
                    mode="branch",
                    polarised=polarised,
                    span_normalise=True,
                )
                afs_sum = [np.sum(window) for window in afs]
                self.assertArrayAlmostEqual(afs_sum, tbl)


############################################
# End of specific stats tests.
############################################


class TestWindowedTreeStat(StatsTestCase):
    """
    Tests that the treewise windowing function defined here has the correct
    behaviour.
    """

    # TODO add more tests here covering the various windowing possibilities.
    def get_tree_sequence(self):
        ts = msprime.simulate(10, recombination_rate=2, random_seed=1)
        assert ts.num_trees > 3
        return ts

    def test_all_trees(self):
        ts = self.get_tree_sequence()
        A1 = np.ones((ts.num_trees, 1))
        windows = np.array(list(ts.breakpoints()))
        A2 = windowed_tree_stat(ts, A1, windows)
        # print("breakpoints = ", windows)
        # print(A2)
        assert A1.shape == A2.shape
        # JK: I don't understand what we're computing here, this normalisation
        # seems pretty weird.
        # for tree in ts.trees():
        #     self.assertAlmostEqual(A2[tree.index, 0], tree.span / ts.sequence_length)

    def test_single_interval(self):
        ts = self.get_tree_sequence()
        A1 = np.ones((ts.num_trees, 1))
        windows = np.array([0, ts.sequence_length])
        A2 = windowed_tree_stat(ts, A1, windows)
        assert A2.shape == (1, 1)
        # TODO: Test output


class TestSampleSets(StatsTestCase):
    """
    Tests that passing sample sets in various ways gets interpreted correctly.
    """

    def get_example_ts(self):
        ts = msprime.simulate(10, mutation_rate=1, recombination_rate=1, random_seed=2)
        assert ts.num_mutations > 0
        return ts

    def test_duplicate_samples(self):
        ts = self.get_example_ts()
        for bad_set in [[1, 1], [1, 2, 1], list(range(10)) + [9]]:
            with pytest.raises(exceptions.LibraryError):
                ts.diversity([bad_set])
            with pytest.raises(exceptions.LibraryError):
                ts.divergence([[0, 1], bad_set])
            with pytest.raises(ValueError):
                ts.sample_count_stat([bad_set], self.identity_f(ts), 1)

    def test_empty_sample_set(self):
        ts = self.get_example_ts()
        with pytest.raises(ValueError):
            ts.diversity([[]])
        for bad_sample_sets in [[[], []], [[1], []], [[1, 2], [1], []]]:
            with pytest.raises(ValueError):
                ts.diversity(bad_sample_sets)
            with pytest.raises(ValueError):
                ts.divergence(bad_sample_sets)
            with pytest.raises(ValueError):
                ts.sample_count_stat(bad_sample_sets, self.identity_f(ts), 1)

    def test_non_samples(self):
        ts = self.get_example_ts()
        with pytest.raises(exceptions.LibraryError):
            ts.diversity([[ts.num_samples]])

        with pytest.raises(exceptions.LibraryError):
            ts.divergence([[ts.num_samples], [1, 2]])

        with pytest.raises(ValueError):
            ts.sample_count_stat([[ts.num_samples]], self.identity_f(ts), 1)

    def test_span_normalise(self):
        np.random.seed(92)
        ts = self.get_example_ts()
        sample_sets = [[0, 1], [2, 3, 4], [5, 6]]
        windows = ts.sequence_length * np.random.uniform(size=10)
        windows.sort()
        windows[0] = 0.0
        windows[-1] = ts.sequence_length
        n = np.array([len(u) for u in sample_sets])

        def f(x):
            return x * (x < n)

        # Determine output_dim of the function
        for mode in ("site", "branch", "node"):
            sigma1 = ts.sample_count_stat(sample_sets, f, 3, windows=windows, mode=mode)
            sigma2 = ts.sample_count_stat(
                sample_sets, f, 3, windows=windows, mode=mode, span_normalise=True
            )
            sigma3 = ts.sample_count_stat(
                sample_sets, f, 3, windows=windows, mode=mode, span_normalise=False
            )
            denom = np.diff(windows)[:, np.newaxis]
            if mode == "node":
                denom = np.diff(windows)[:, np.newaxis, np.newaxis]

            assert sigma1.shape == sigma2.shape
            assert sigma1.shape == sigma3.shape
            self.assertArrayAlmostEqual(sigma1, sigma2)
            self.assertArrayAlmostEqual(sigma1, sigma3 / denom)


class TestSampleSetIndexes(StatsTestCase):
    """
    Tests that we get the correct behaviour from the indexes argument to
    k-way stats functions.
    """

    def get_example_ts(self):
        ts = msprime.simulate(10, mutation_rate=1, random_seed=1)
        assert ts.num_mutations > 0
        return ts

    def test_2_way_default(self):
        ts = self.get_example_ts()
        sample_sets = np.array_split(ts.samples(), 2)
        S1 = ts.divergence(sample_sets)
        S2 = divergence(ts, sample_sets)[0, 0]
        S3 = ts.divergence(sample_sets, [0, 1])
        assert S1.shape == S2.shape
        self.assertArrayAlmostEqual(S1, S2)
        self.assertArrayAlmostEqual(S1, S3)
        sample_sets = np.array_split(ts.samples(), 3)
        with pytest.raises(ValueError):
            _ = ts.divergence(sample_sets)
        with pytest.raises(ValueError):
            _ = ts.divergence([sample_sets[0]])

    def test_3_way_default(self):
        ts = self.get_example_ts()
        sample_sets = np.array_split(ts.samples(), 3)
        S1 = ts.f3(sample_sets)
        S2 = f3(ts, sample_sets)[0, 0]
        S3 = ts.f3(sample_sets, [0, 1, 2])
        assert S1.shape == S2.shape
        self.assertArrayAlmostEqual(S1, S2)
        self.assertArrayAlmostEqual(S1, S3)
        sample_sets = np.array_split(ts.samples(), 4)
        with pytest.raises(ValueError):
            _ = ts.f3(sample_sets)

    def test_4_way_default(self):
        ts = self.get_example_ts()
        sample_sets = np.array_split(ts.samples(), 4)
        S1 = ts.f4(sample_sets)
        S2 = f4(ts, sample_sets)
        S3 = ts.f4(sample_sets, [0, 1, 2, 3])
        assert S1.shape == S3.shape
        self.assertArrayAlmostEqual(S1, S2)
        self.assertArrayAlmostEqual(S1, S3)
        sample_sets = np.array_split(ts.samples(), 5)
        with pytest.raises(ValueError):
            _ = ts.f4(sample_sets)

    def test_2_way_combinations(self):
        ts = self.get_example_ts()
        sample_sets = np.array_split(ts.samples(), 4)
        pairs = list(itertools.combinations(range(4), 2))
        for k in range(1, len(pairs)):
            S1 = ts.divergence(sample_sets, pairs[:k])
            S2 = divergence(ts, sample_sets, pairs[:k])[0]
            assert S1.shape[-1] == k
            assert S1.shape == S2.shape
            self.assertArrayAlmostEqual(S1, S2)

    def test_3_way_combinations(self):
        ts = self.get_example_ts()
        sample_sets = np.array_split(ts.samples(), 5)
        triples = list(itertools.combinations(range(5), 3))
        for k in range(1, len(triples)):
            S1 = ts.Y3(sample_sets, triples[:k])
            S2 = Y3(ts, sample_sets, triples[:k])[0]
            assert S1.shape[-1] == k
            assert S1.shape == S2.shape
            self.assertArrayAlmostEqual(S1, S2)

    def test_4_way_combinations(self):
        ts = self.get_example_ts()
        sample_sets = np.array_split(ts.samples(), 5)
        quads = list(itertools.combinations(range(5), 4))
        for k in range(1, len(quads)):
            S1 = ts.f4(sample_sets, quads[:k], windows=[0, ts.sequence_length])
            S2 = f4(ts, sample_sets, quads[:k])
            assert S1.shape[-1] == k
            assert S2.shape == S2.shape
            self.assertArrayAlmostEqual(S1, S2)

    def test_errors(self):
        ts = self.get_example_ts()
        sample_sets = np.array_split(ts.samples(), 2)
        with pytest.raises(ValueError):
            ts.divergence(sample_sets, indexes=[])
        with pytest.raises(ValueError):
            ts.divergence(sample_sets, indexes=[(1, 1, 1)])
        with pytest.raises(exceptions.LibraryError):
            ts.divergence(sample_sets, indexes=[(1, 2)])


class TestGeneralStatInterface(StatsTestCase):
    """
    Tests for the basic interface for general_stats.
    """

    def get_tree_sequence(self):
        ts = msprime.simulate(10, recombination_rate=2, mutation_rate=2, random_seed=1)
        return ts

    def test_function_cannot_update_state(self):
        ts = self.get_tree_sequence()

        def f(x):
            out = x.copy()
            x[:] = 0.0
            return out

        def g(x):
            return x

        x = ts.sample_count_stat(
            [ts.samples()],
            f,
            output_dim=1,
            strict=False,
            mode="node",
            span_normalise=False,
        )
        y = ts.sample_count_stat(
            [ts.samples()],
            g,
            output_dim=1,
            strict=False,
            mode="node",
            span_normalise=False,
        )
        self.assertArrayEqual(x, y)

    def test_default_mode(self):
        ts = msprime.simulate(10, recombination_rate=1, random_seed=2)
        W = np.ones((ts.num_samples, 2))
        sigma1 = ts.general_stat(W, self.identity_f(ts), W.shape[1])
        sigma2 = ts.general_stat(W, self.identity_f(ts), W.shape[1], mode="site")
        self.assertArrayEqual(sigma1, sigma2)

    def test_bad_mode(self):
        ts = msprime.simulate(10, recombination_rate=1, random_seed=2)
        W = np.ones((ts.num_samples, 2))
        for bad_mode in ["", "MODE", "x" * 8192]:
            with pytest.raises(ValueError):
                ts.general_stat(W, self.identity_f(ts), W.shape[1], mode=bad_mode)

    def test_bad_window_strings(self):
        ts = self.get_tree_sequence()
        with pytest.raises(ValueError):
            ts.diversity([ts.samples()], mode="site", windows="abc")
        with pytest.raises(ValueError):
            ts.diversity([ts.samples()], mode="site", windows="")
        with pytest.raises(ValueError):
            ts.diversity([ts.samples()], mode="tree", windows="abc")

    def test_bad_summary_function(self):
        ts = self.get_tree_sequence()
        W = np.ones((ts.num_samples, 3))
        with pytest.raises(ValueError):
            ts.general_stat(W, lambda x: x, 3, windows="sites")
        with pytest.raises(ValueError):
            ts.general_stat(W, lambda x: np.array([1.0]), 1, windows="sites")

    def test_nonnumpy_summary_function(self):
        ts = self.get_tree_sequence()
        W = np.ones((ts.num_samples, 3))
        sigma1 = ts.general_stat(W, lambda x: [0.0], 1)
        sigma2 = ts.general_stat(W, lambda x: np.array([0.0]), 1)
        self.assertArrayEqual(sigma1, sigma2)


class TestGeneralBranchStats(StatsTestCase):
    """
    Tests for general branch stats (using functions and arbitrary weights)
    """

    def compare_general_stat(self, ts, W, f, windows=None, polarised=False):
        # Determine output_dim of the function
        M = len(f(W[0]))
        sigma1 = naive_branch_general_stat(ts, W, f, windows, polarised=polarised)
        sigma2 = ts.general_stat(W, f, M, windows, polarised=polarised, mode="branch")
        sigma3 = branch_general_stat(ts, W, f, windows, polarised=polarised)
        assert sigma1.shape == sigma2.shape
        assert sigma1.shape == sigma3.shape
        self.assertArrayAlmostEqual(sigma1, sigma2)
        self.assertArrayAlmostEqual(sigma1, sigma3)
        return sigma1

    def test_simple_identity_f_w_zeros(self):
        ts = msprime.simulate(12, recombination_rate=3, random_seed=2)
        W = np.zeros((ts.num_samples, 3))
        for polarised in [True, False]:
            sigma = self.compare_general_stat(
                ts, W, self.identity_f(ts), windows="trees", polarised=polarised
            )
            assert sigma.shape == (ts.num_trees, W.shape[1])
            assert np.all(sigma == 0)

    def test_simple_identity_f_w_ones(self):
        ts = msprime.simulate(10, recombination_rate=1, random_seed=2)
        W = np.ones((ts.num_samples, 2))
        sigma = self.compare_general_stat(
            ts, W, self.identity_f(ts), windows="trees", polarised=True
        )
        assert sigma.shape == (ts.num_trees, W.shape[1])
        # A W of 1 for every node and identity f counts the samples in the subtree
        # if polarised is True.
        for tree in ts.trees():
            s = sum(tree.num_samples(u) * tree.branch_length(u) for u in tree.nodes())
            assert np.allclose(sigma[tree.index], s)

    def test_simple_cumsum_f_w_ones(self):
        ts = msprime.simulate(13, recombination_rate=1, random_seed=2)
        W = np.ones((ts.num_samples, 8))
        for polarised in [True, False]:
            sigma = self.compare_general_stat(
                ts, W, self.cumsum_f(ts), windows="trees", polarised=polarised
            )
            assert sigma.shape == (ts.num_trees, W.shape[1])

    def test_simple_cumsum_f_w_ones_many_windows(self):
        ts = msprime.simulate(15, recombination_rate=3, random_seed=3)
        assert ts.num_trees > 3
        windows = np.linspace(0, ts.sequence_length, num=ts.num_trees * 10)
        W = np.ones((ts.num_samples, 3))
        sigma = self.compare_general_stat(ts, W, self.cumsum_f(ts), windows=windows)
        assert sigma.shape == (windows.shape[0] - 1, W.shape[1])

    def test_windows_equal_to_ts_breakpoints(self):
        ts = msprime.simulate(14, recombination_rate=1, random_seed=2)
        W = np.ones((ts.num_samples, 1))
        for polarised in [True, False]:
            sigma_no_windows = self.compare_general_stat(
                ts, W, self.cumsum_f(ts), windows="trees", polarised=polarised
            )
            assert sigma_no_windows.shape == (ts.num_trees, W.shape[1])
            sigma_windows = self.compare_general_stat(
                ts,
                W,
                self.cumsum_f(ts),
                windows=ts.breakpoints(as_array=True),
                polarised=polarised,
            )
            assert sigma_windows.shape == sigma_no_windows.shape
            assert np.allclose(sigma_windows.shape, sigma_no_windows.shape)

    def test_single_tree_windows(self):
        ts = msprime.simulate(15, random_seed=2, length=100)
        W = np.ones((ts.num_samples, 2))
        f = self.sum_f(ts)
        # for num_windows in range(1, 10):
        for num_windows in [2]:
            windows = np.linspace(0, ts.sequence_length, num=num_windows + 1)
            sigma = self.compare_general_stat(ts, W, f, windows)
            assert sigma.shape == (num_windows, 1)

    def test_simple_identity_f_w_zeros_windows(self):
        ts = msprime.simulate(15, recombination_rate=3, random_seed=2)
        W = np.zeros((ts.num_samples, 3))
        f = self.identity_f(ts)
        windows = np.linspace(0, ts.sequence_length, num=11)
        for polarised in [True, False]:
            sigma = self.compare_general_stat(ts, W, f, windows, polarised=polarised)
            assert sigma.shape == (10, W.shape[1])
            assert np.all(sigma == 0)

    def test_nonstrict_nonancestral_material(self):
        # 0 is a sample, 1 is not
        #
        # 2.00┊  2  ┊ 2 ┊  2  ┊
        #     ┊ ┏┻┓ ┊ ┃ ┊ ┏┻┓ ┊
        # 1.00┊ ┃ 1 ┊ 1 ┊ ┃ 1 ┊
        #     ┊ ┃   ┊ ┃ ┊ ┃   ┊
        # 0.00┊ 0   ┊ 0 ┊ 0   ┊
        #     0     1   2     3

        tables = tskit.TableCollection(sequence_length=3)

        node_times = [0, 1, 2]
        samples = [0]
        for n, t in enumerate(node_times):
            tables.nodes.add_row(
                time=t, flags=tskit.NODE_IS_SAMPLE if n in samples else 0
            )

        # p, c, l, r
        edges = [
            (1, 0, 1, 2),
            (2, 0, 0, 1),
            (2, 0, 2, 3),
            (2, 1, 0, 3),
        ]
        for p, c, l, r in edges:
            tables.edges.add_row(parent=p, child=c, left=l, right=r)

        # this makes it so 'site' mode counts branches
        for x in range(int(tables.sequence_length)):
            for n in range(tables.nodes.num_rows - 1):
                offset = n / tables.nodes.num_rows
                s = tables.sites.add_row(position=x + offset, ancestral_state="0")
                tables.mutations.add_row(site=s, node=n, derived_state="1")

        ts = tables.tree_sequence()

        def f(x):
            return x

        for polarised, mode, answer in [
            (True, "branch", 6),
            (True, "site", 4),
            (False, "branch", 8),
            (False, "site", 6),
        ]:
            (stat,) = ts.sample_count_stat(
                [[0]],
                f,
                1,
                strict=False,
                span_normalise=False,
                polarised=polarised,
                mode=mode,
            )
            assert stat == answer


class TestGeneralSiteStats(StatsTestCase):
    """
    Tests for general site stats (using functions and arbitrary weights)
    """

    def compare_general_stat(self, ts, W, f, windows=None, polarised=False):
        # Determine output_dim of the function
        M = len(f(W[0]))
        sigma1 = naive_site_general_stat(ts, W, f, windows, polarised=polarised)
        sigma2 = ts.general_stat(W, f, M, windows, polarised=polarised, mode="site")
        sigma3 = site_general_stat(ts, W, f, windows, polarised=polarised)
        assert sigma1.shape == sigma2.shape
        assert sigma1.shape == sigma3.shape
        self.assertArrayAlmostEqual(sigma1, sigma2)
        self.assertArrayAlmostEqual(sigma1, sigma3)
        return sigma1

    def test_identity_f_W_0_multiple_alleles(self):
        ts = msprime.simulate(20, recombination_rate=0, random_seed=2)
        ts = tsutil.jukes_cantor(ts, 20, 1, seed=10)
        W = np.zeros((ts.num_samples, 3))
        for polarised in [True, False]:
            sigma = self.compare_general_stat(
                ts, W, self.identity_f(ts), windows="sites", polarised=polarised
            )
            assert sigma.shape == (ts.num_sites, W.shape[1])
            assert np.all(sigma == 0)

    def test_identity_f_W_0_multiple_alleles_windows(self):
        ts = msprime.simulate(34, recombination_rate=0, random_seed=2)
        ts = tsutil.jukes_cantor(ts, 20, 1, seed=10)
        W = np.zeros((ts.num_samples, 3))
        windows = np.linspace(0, 1, num=11)
        for polarised in [True, False]:
            sigma = self.compare_general_stat(
                ts, W, self.identity_f(ts), windows=windows, polarised=polarised
            )
            assert sigma.shape == (windows.shape[0] - 1, W.shape[1])
            assert np.all(sigma == 0)

    def test_cumsum_f_W_1_multiple_alleles(self):
        ts = msprime.simulate(3, recombination_rate=2, random_seed=2)
        ts = tsutil.jukes_cantor(ts, 20, 1, seed=10)
        W = np.ones((ts.num_samples, 3))
        for polarised in [True, False]:
            sigma = self.compare_general_stat(
                ts, W, self.cumsum_f(ts), windows="sites", polarised=polarised
            )
            assert sigma.shape == (ts.num_sites, W.shape[1])

    def test_cumsum_f_W_1_two_alleles(self):
        ts = msprime.simulate(33, recombination_rate=1, mutation_rate=2, random_seed=1)
        W = np.ones((ts.num_samples, 5))
        for polarised in [True, False]:
            sigma = self.compare_general_stat(
                ts, W, self.cumsum_f(ts), windows="sites", polarised=polarised
            )
            assert sigma.shape == (ts.num_sites, W.shape[1])


class TestGeneralNodeStats(StatsTestCase):
    """
    Tests for general node stats (using functions and arbitrary weights)
    """

    def compare_general_stat(self, ts, W, f, windows=None, polarised=False):
        # Determine output_dim of the function
        M = len(f(W[0]))
        sigma1 = naive_node_general_stat(ts, W, f, windows, polarised=polarised)
        sigma2 = ts.general_stat(W, f, M, windows, polarised=polarised, mode="node")
        sigma3 = node_general_stat(ts, W, f, windows, polarised=polarised)
        assert sigma1.shape == sigma2.shape
        assert sigma1.shape == sigma3.shape
        self.assertArrayAlmostEqual(sigma1, sigma2)
        self.assertArrayAlmostEqual(sigma1, sigma3)
        return sigma1

    def test_simple_sum_f_w_zeros(self):
        ts = msprime.simulate(12, recombination_rate=3, random_seed=2)
        W = np.zeros((ts.num_samples, 3))
        for polarised in [True, False]:
            sigma = self.compare_general_stat(
                ts, W, self.identity_f(ts), windows="trees", polarised=polarised
            )
            assert sigma.shape == (ts.num_trees, ts.num_nodes, 3)
            assert np.all(sigma == 0)

    def test_simple_sum_f_w_ones(self):
        ts = msprime.simulate(44, recombination_rate=1, random_seed=2)
        W = np.ones((ts.num_samples, 2))
        f = self.sum_f(ts)
        sigma = self.compare_general_stat(ts, W, f, windows="trees", polarised=True)
        assert sigma.shape == (ts.num_trees, ts.num_nodes, 1)
        # Drop the last dimension
        sigma = sigma.reshape((ts.num_trees, ts.num_nodes))
        # A W of 1 for every node and f(x)=sum(x) counts the samples in the subtree
        # times 2 if polarised is True.
        for tree in ts.trees():
            s = np.array(
                [
                    tree.num_samples(u) if tree.num_samples(u) < ts.num_samples else 0
                    for u in range(ts.num_nodes)
                ]
            )
            self.assertArrayAlmostEqual(sigma[tree.index], 2 * s)

    def test_simple_sum_f_w_ones_notstrict(self):
        ts = msprime.simulate(44, recombination_rate=1, random_seed=2)
        W = np.ones((ts.num_samples, 2))
        sigma = ts.general_stat(
            W,
            lambda x: np.array([np.sum(x)]),
            1,
            windows="trees",
            polarised=True,
            mode="node",
            strict=False,
        )
        assert sigma.shape == (ts.num_trees, ts.num_nodes, 1)
        # Drop the last dimension
        sigma = sigma.reshape((ts.num_trees, ts.num_nodes))
        # A W of 1 for every node and f(x)=sum(x) counts the samples in the subtree
        # times 2 if polarised is True.
        for tree in ts.trees():
            s = np.array([tree.num_samples(u) for u in range(ts.num_nodes)])
            self.assertArrayAlmostEqual(sigma[tree.index], 2 * s)

    def test_small_tree_windows_polarised(self):
        ts = msprime.simulate(4, recombination_rate=0.5, random_seed=2)
        assert ts.num_trees > 1
        W = np.ones((ts.num_samples, 1))
        sigma = self.compare_general_stat(
            ts,
            W,
            self.cumsum_f(ts),
            windows=ts.breakpoints(as_array=True),
            polarised=True,
        )
        assert sigma.shape == (ts.num_trees, ts.num_nodes, 1)

    def test_one_window_polarised(self):
        ts = msprime.simulate(4, recombination_rate=1, random_seed=2)
        W = np.ones((ts.num_samples, 1))
        sigma = self.compare_general_stat(
            ts, W, self.cumsum_f(ts), windows=[0, ts.sequence_length], polarised=True
        )
        assert sigma.shape == (1, ts.num_nodes, W.shape[1])

    def test_one_window_unpolarised(self):
        ts = msprime.simulate(4, recombination_rate=1, random_seed=2)
        W = np.ones((ts.num_samples, 2))
        sigma = self.compare_general_stat(
            ts, W, self.cumsum_f(ts), windows=[0, ts.sequence_length], polarised=False
        )
        assert sigma.shape == (1, ts.num_nodes, 2)

    def test_many_windows(self):
        ts = msprime.simulate(24, recombination_rate=3, random_seed=2)
        W = np.ones((ts.num_samples, 3))
        for k in [1, ts.num_trees // 2, ts.num_trees, ts.num_trees * 2]:
            windows = np.linspace(0, 1, num=k + 1)
            for polarised in [True]:
                sigma = self.compare_general_stat(
                    ts, W, self.cumsum_f(ts), windows=windows, polarised=polarised
                )
            assert sigma.shape == (k, ts.num_nodes, 3)

    def test_one_tree(self):
        ts = msprime.simulate(10, random_seed=3)
        W = np.ones((ts.num_samples, 2))
        f = self.sum_f(ts, k=2)
        sigma = self.compare_general_stat(ts, W, f, windows=[0, 1], polarised=True)
        assert sigma.shape == (1, ts.num_nodes, 2)
        # A W of 1 for every node and f(x)=sum(x) counts the samples in the subtree
        # times 2 if polarised is True.
        tree = ts.first()
        s = np.array(
            [
                tree.num_samples(u) if tree.num_samples(u) < ts.num_samples else 0
                for u in range(ts.num_nodes)
            ]
        )
        self.assertArrayAlmostEqual(sigma[tree.index, :, 0], 2 * s)
        self.assertArrayAlmostEqual(sigma[tree.index, :, 1], 2 * s)


##############################
# Trait covariance
##############################


@cached_np
def covsq(x, y):
    cov = np.dot(x - np.mean(x), y - np.mean(y)) / (len(x) - 1)
    return cov * cov


@cached_np
def corsq(x, y):
    vx = covsq(x, x)
    vy = covsq(y, y)
    # sqrt is because vx and vy are *squared* variances
    return covsq(x, y) / np.sqrt(vx * vy)


def site_trait_covariance(ts, W, windows=None, span_normalise=True):
    """
    For each site, computes the covariance between the columns of W and the genotypes.
    """
    windows = ts.parse_windows(windows)
    n, K = W.shape
    assert n == ts.num_samples
    out = np.zeros((len(windows) - 1, K))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        haps = ts.genotype_matrix(isolated_as_missing=False)
        site_positions = [x.position for x in ts.sites()]
        for i in range(K):
            w = W[:, i].copy()
            w -= np.mean(w)
            S = 0
            site_in_window = False
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    site_in_window = True
                    hX = haps[k]
                    alleles = set(hX)
                    for a in alleles:
                        S += covsq(w, hX == a) / 2
            if site_in_window:
                out[j, i] = S
                if span_normalise:
                    out[j, i] /= end - begin
    return out


def branch_trait_covariance(ts, W, windows=None, span_normalise=True):
    """
    For each branch, computes the covariance between the columns of W and the split
    induced by the branch, multiplied by the length of the branch.
    """
    windows = ts.parse_windows(windows)
    n, K = W.shape
    assert n == ts.num_samples
    out = np.zeros((len(windows) - 1, K))
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i in range(K):
            w = W[:, i].copy()
            w -= np.mean(w)
            S = 0
            has_trees = False
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                if tr.total_branch_length > 0:
                    has_trees = True
                SS = 0
                for u in range(ts.num_nodes):
                    tree_samples = set(tr.samples(u))
                    below = np.fromiter(
                        (s in tree_samples for s in samples), dtype=bool
                    )
                    branch_length = tr.branch_length(u)
                    SS += covsq(w, below) * branch_length
                S += SS * (min(end, tr.interval.right) - max(begin, tr.interval.left))
            if has_trees:
                out[j, i] = S
                if span_normalise:
                    out[j, i] /= end - begin
    return out


def node_trait_covariance(ts, W, windows=None, span_normalise=True):
    """
    For each node, computes the covariance between the columns of W and the split
    induced by above/below the node.
    """
    windows = ts.parse_windows(windows)
    n, K = W.shape
    assert n == ts.num_samples
    out = np.zeros((len(windows) - 1, ts.num_nodes, K))
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i in range(K):
            w = W[:, i].copy()
            w -= np.mean(w)
            S = np.zeros(ts.num_nodes)
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in range(ts.num_nodes):
                    tree_samples = set(tr.samples(u))
                    below = np.fromiter(
                        (s in tree_samples for s in samples), dtype=bool
                    )
                    SS[u] += covsq(w, below)
                S += SS * (min(end, tr.interval.right) - max(begin, tr.interval.left))
            out[j, :, i] = S
            if span_normalise:
                out[j, :, i] /= end - begin
    return out


def trait_covariance(ts, W, windows=None, mode="site", span_normalise=True):
    method_map = {
        "site": site_trait_covariance,
        "node": node_trait_covariance,
        "branch": branch_trait_covariance,
    }
    return method_map[mode](ts, W, windows=windows, span_normalise=span_normalise)


class TestTraitCovariance(StatsTestCase, WeightStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def get_example_ts(self):
        ts = msprime.simulate(10, mutation_rate=1, recombination_rate=2, random_seed=1)
        assert ts.num_mutations > 0
        return ts

    def transform_weights(self, W):
        """
        Need centered weights to compare to general stats.
        """
        W -= np.mean(W, axis=0)
        return W

    def verify_weighted_stat(self, ts, W, windows):
        n = W.shape[0]

        def f(x):
            return (x**2) / (2 * (n - 1) * (n - 1))

        self.verify_definition(ts, W, windows, f, ts.trait_covariance, trait_covariance)

    def verify_interface(self, ts, ts_method):
        W = np.array([np.arange(ts.num_samples)]).T
        sigma1 = ts_method(W, mode=self.mode)
        sigma2 = ts_method(W, windows=None, mode=self.mode)
        sigma3 = ts_method(W, windows=[0.0, ts.sequence_length], mode=self.mode)
        assert sigma1.shape == sigma2.shape
        self.assertArrayAlmostEqual(sigma1, sigma2)
        self.assertArrayAlmostEqual(sigma1, sigma3[0])

    def verify_centering(self, ts, method, ts_method):
        # Since weights are mean-centered, adding a constant shouldn't change anything.
        ts = self.get_example_ts()
        for W, windows in subset_combos(
            self.example_weights(ts), example_windows(ts), p=0.1
        ):
            shift = np.arange(1, W.shape[1] + 1)
            sigma1 = ts_method(W, windows=windows, mode=self.mode)
            sigma2 = ts_method(W + shift, windows=windows, mode=self.mode)
            sigma3 = method(ts, W, windows=windows, mode=self.mode)
            sigma4 = method(ts, W + shift, windows=windows, mode=self.mode)
            assert sigma1.shape == sigma2.shape
            assert sigma1.shape == sigma3.shape
            assert sigma1.shape == sigma4.shape
            self.assertArrayAlmostEqual(sigma1, sigma2)
            self.assertArrayAlmostEqual(sigma1, sigma3)
            self.assertArrayAlmostEqual(sigma1, sigma4)


class TraitCovarianceMixin:
    def test_interface(self):
        ts = self.get_example_ts()
        self.verify_interface(ts, ts.trait_covariance)

    def test_normalisation(self):
        ts = self.get_example_ts()
        self.verify_centering(ts, trait_covariance, ts.trait_covariance)

    def test_errors(self):
        ts = self.get_example_ts()
        W = np.ones((ts.num_samples, 2))
        # W must have the right number of rows
        with pytest.raises(ValueError):
            ts.trait_correlation(W[1:, :])


@pytest.mark.slow
class TestBranchTraitCovariance(
    TestTraitCovariance, TopologyExamplesMixin, TraitCovarianceMixin
):
    mode = "branch"


@pytest.mark.slow
class TestNodeTraitCovariance(
    TestTraitCovariance, TopologyExamplesMixin, TraitCovarianceMixin
):
    mode = "node"


class TestSiteTraitCovariance(
    TestTraitCovariance, MutatedTopologyExamplesMixin, TraitCovarianceMixin
):
    mode = "site"


##############################
# Trait correlation
##############################


def site_trait_correlation(ts, W, windows=None, span_normalise=True):
    """
    For each site, computes the correlation between the columns of W and the genotypes.
    """
    windows = ts.parse_windows(windows)
    n, K = W.shape
    assert n == ts.num_samples
    out = np.zeros((len(windows) - 1, K))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        haps = ts.genotype_matrix(isolated_as_missing=False)
        site_positions = [x.position for x in ts.sites()]
        for i in range(K):
            w = W[:, i].copy()
            w -= np.mean(w)
            w /= np.std(w) * np.sqrt(len(w) / (len(w) - 1))
            S = 0
            site_in_window = False
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    site_in_window = True
                    hX = haps[k]
                    alleles = set(hX)
                    for a in alleles:
                        p = np.mean(hX == a)
                        if p > 0 and p < 1:
                            # S += sum(w[hX == a])**2 / (2 * (p * (1 - p)))
                            S += corsq(w, hX == a) / 2
            if site_in_window:
                out[j, i] = S
                if span_normalise:
                    out[j, i] /= end - begin
    return out


def branch_trait_correlation(ts, W, windows=None, span_normalise=True):
    """
    For each branch, computes the correlation between the columns of W and the split
    induced by the branch, multiplied by the length of the branch.
    """
    windows = ts.parse_windows(windows)
    n, K = W.shape
    assert n == ts.num_samples
    out = np.zeros((len(windows) - 1, K))
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i in range(K):
            w = W[:, i].copy()
            w -= np.mean(w)
            w /= np.std(w) * np.sqrt(len(w) / (len(w) - 1))
            S = 0
            has_trees = False
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                if tr.total_branch_length > 0:
                    has_trees = True
                SS = 0
                for u in range(ts.num_nodes):
                    tree_samples = set(tr.samples(u))
                    below = np.fromiter(
                        (s in tree_samples for s in samples), dtype=bool
                    )
                    p = np.mean(below)
                    if p > 0 and p < 1:
                        branch_length = tr.branch_length(u)
                        # SS += ((sum(w[below])**2 +
                        #         sum(w[np.logical_not(below)])**2) * branch_length
                        #        / (2 * (p * (1 - p))))
                        SS += corsq(w, below) * branch_length
                S += SS * (min(end, tr.interval.right) - max(begin, tr.interval.left))
            if has_trees:
                out[j, i] = S
                if span_normalise:
                    out[j, i] /= end - begin
    return out


def node_trait_correlation(ts, W, windows=None, span_normalise=True):
    """
    For each node, computes the correlation between the columns of W and the split
    induced by above/below the node.
    """
    windows = ts.parse_windows(windows)
    n, K = W.shape
    assert n == ts.num_samples
    out = np.zeros((len(windows) - 1, ts.num_nodes, K))
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i in range(K):
            w = W[:, i].copy()
            w -= np.mean(w)
            w /= np.std(w) * np.sqrt(len(w) / (len(w) - 1))
            S = np.zeros(ts.num_nodes)
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in range(ts.num_nodes):
                    tree_samples = set(tr.samples(u))
                    below = np.fromiter(
                        (s in tree_samples for s in samples), dtype=bool
                    )
                    p = np.mean(below)
                    if p > 0 and p < 1:
                        # SS[u] += sum(w[below])**2 / 2
                        # SS[u] += sum(w[np.logical_not(below)])**2 / 2
                        # SS[u] /= (p * (1 - p))
                        SS[u] += corsq(w, below)
                S += SS * (min(end, tr.interval.right) - max(begin, tr.interval.left))
            out[j, :, i] = S
            if span_normalise:
                out[j, :, i] /= end - begin
    return out


def trait_correlation(ts, W, windows=None, mode="site", span_normalise=True):
    method_map = {
        "site": site_trait_correlation,
        "node": node_trait_correlation,
        "branch": branch_trait_correlation,
    }
    return method_map[mode](ts, W, windows=windows, span_normalise=span_normalise)


class TestTraitCorrelation(TestTraitCovariance):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def transform_weights(self, W):
        """
        Need standardised weights to compare to general stats,
        and also an extra column to compute allele frequencies.
        """
        W -= np.mean(W, axis=0)
        n = W.shape[0]
        with suppress_division_by_zero_warning():
            W /= np.std(W, axis=0) * np.sqrt(n / (n - 1))
        return np.column_stack((W, np.ones(W.shape[0]) / W.shape[0]))

    def verify_weighted_stat(self, ts, W, windows):
        n = W.shape[0]

        def f(x):
            p = x[-1]
            if p > 0 and p < 1:
                return (x[:-1] ** 2) / (2 * (p * (1 - p)) * n * (n - 1))
            else:
                return x[:-1] * 0.0

        self.verify_definition(
            ts, W, windows, f, ts.trait_correlation, trait_correlation
        )

    def test_errors(self):
        ts = self.get_example_ts()
        # columns of W must have positive SD
        W = np.ones((ts.num_samples, 2))
        with pytest.raises(ValueError):
            ts.trait_correlation(W)
        # W must have the right number of rows
        with pytest.raises(ValueError):
            ts.trait_correlation(W[1:, :])

    def verify_standardising(self, ts, method, ts_method):
        """
        Since weights are standardised, multiplying by a constant shouldn't
        change anything.
        """
        for W, windows in subset_combos(
            self.example_weights(ts), example_windows(ts), p=0.1
        ):
            scale = np.arange(1, W.shape[1] + 1)
            sigma1 = ts_method(W, windows=windows, mode=self.mode)
            sigma2 = ts_method(W * scale, windows=windows, mode=self.mode)
            sigma3 = method(ts, W, windows=windows, mode=self.mode)
            sigma4 = method(ts, W * scale, windows=windows, mode=self.mode)
            assert sigma1.shape == sigma2.shape
            self.assertArrayAlmostEqual(sigma1, sigma2)
            self.assertArrayAlmostEqual(sigma1, sigma3)
            self.assertArrayAlmostEqual(sigma1, sigma4)


class TraitCorrelationMixin:
    def test_interface(self):
        ts = self.get_example_ts()
        self.verify_interface(ts, ts.trait_correlation)

    def test_normalisation(self):
        ts = self.get_example_ts()
        self.verify_centering(ts, trait_correlation, ts.trait_correlation)
        self.verify_standardising(ts, trait_correlation, ts.trait_correlation)


@pytest.mark.slow
class TestBranchTraitCorrelation(
    TestTraitCorrelation, TopologyExamplesMixin, TraitCorrelationMixin
):
    mode = "branch"


@pytest.mark.slow
class TestNodeTraitCorrelation(
    TestTraitCorrelation, TopologyExamplesMixin, TraitCorrelationMixin
):
    mode = "node"


class TestSiteTraitCorrelation(
    TestTraitCorrelation, MutatedTopologyExamplesMixin, TraitCorrelationMixin
):
    mode = "site"


##############################
# Trait linear_model
##############################

# Quick hack to speed up the tests a bit. We're running this linear_model
# function a gazillion times with similar arguments, so it's worth
# caching the results. We could use functools.lru_cache, but the arguments
# are numpy arrays and so we have to do something different. This
# reduces runtime from ~40 seconds to ~7 seconds on the
# TestNodeTraitLinearModel.
# See https://github.com/tskit-dev/tskit/issues/1856 for more info
_lm_cache = {}


@cached_np
def linear_model(y, x, z):
    key = (y.tobytes(), x.tobytes(), z.tobytes())
    if key not in _lm_cache:
        _lm_cache[key] = _linear_model(y, x, z)
    return _lm_cache[key]


def _linear_model(y, x, z):
    """
    Returns the squared coefficient of x in the least-squares linear model
    :   y ~ x + z
    where x and y are vectors and z is a matrix.
    Note that if z is None then the output is
      cor(x, y) * sd(y) / sd(x) = cov(x, y) / (sd(x) ** 2) .
    """
    # add the constant vector to z
    if z is None:
        z = np.ones((len(x), 1))
    else:
        xz = np.column_stack([z, np.ones((len(x), 1))])
        if np.linalg.matrix_rank(xz) == xz.shape[1]:
            z = xz
    xz = np.column_stack([x, z])
    # check if y is sufficiently independent of the subspace spanned by xz
    Pz = np.matmul(z, np.linalg.pinv(z))
    Py = np.matmul(Pz, y)
    denom = np.sum((y - Py) ** 2)
    if np.linalg.matrix_rank(xz) < xz.shape[1] or denom < 1e-8:
        return 0.0
    else:
        coefs, _, _, _ = np.linalg.lstsq(xz, y, rcond=None)
        return coefs[0] * coefs[0]


def site_trait_linear_model(ts, W, Z, windows=None, span_normalise=True):
    """
    For each site, and for each trait w (column of W), computes the coefficient
    of site in the linear model:
      w ~ site + Z
    """
    windows = ts.parse_windows(windows)
    n, K = W.shape
    assert n == ts.num_samples
    out = np.zeros((len(windows) - 1, K))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        haps = ts.genotype_matrix(isolated_as_missing=False)
        site_positions = [x.position for x in ts.sites()]
        for i in range(K):
            w = W[:, i]
            S = 0
            site_in_window = False
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    site_in_window = True
                    hX = haps[k]
                    alleles = set(hX)
                    for a in alleles:
                        p = np.mean(hX == a)
                        if p > 0 and p < 1:
                            S += linear_model(w, hX == a, Z) / 2
            if site_in_window:
                out[j, i] = S
                if span_normalise:
                    out[j, i] /= end - begin
    return out


def branch_trait_linear_model(ts, W, Z, windows=None, span_normalise=True):
    """
    For each branch, fits the linear_model of each column of W onto the split
    induced by the branch and the covariates Z, multiplied by the length of the branch,
    returning the squared coefficient of the column of W.
    """
    windows = ts.parse_windows(windows)
    n, K = W.shape
    assert n == ts.num_samples
    out = np.zeros((len(windows) - 1, K))
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i in range(K):
            w = W[:, i]
            S = 0
            has_trees = False
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                if tr.total_branch_length > 0:
                    has_trees = True
                SS = 0
                for u in range(ts.num_nodes):
                    tree_samples = set(tr.samples(u))
                    below = np.fromiter(
                        (s in tree_samples for s in samples), dtype=bool
                    )
                    branch_length = tr.branch_length(u)
                    SS += linear_model(w, below, Z) * branch_length
                S += SS * (min(end, tr.interval.right) - max(begin, tr.interval.left))
            if has_trees:
                out[j, i] = S
                if span_normalise:
                    out[j, i] /= end - begin
    return out


def node_trait_linear_model(ts, W, Z, windows=None, span_normalise=True):
    """
    For each node, fits the linear model of each columns of W on the split
    induced by above/below the node and the covariates Z, returning the squared
    coefficient of the column of W.
    """
    windows = ts.parse_windows(windows)
    n, K = W.shape
    assert n == ts.num_samples
    out = np.zeros((len(windows) - 1, ts.num_nodes, K))
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i in range(K):
            w = W[:, i]
            S = np.zeros(ts.num_nodes)
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in range(ts.num_nodes):
                    tree_samples = set(tr.samples(u))
                    below = np.fromiter(
                        (s in tree_samples for s in samples), dtype=bool
                    )
                    SS[u] += linear_model(w, below, Z)
                S += SS * (min(end, tr.interval.right) - max(begin, tr.interval.left))
            out[j, :, i] = S
            if span_normalise:
                out[j, :, i] /= end - begin
    return out


def trait_linear_model(ts, W, Z, windows=None, mode="site", span_normalise=True):
    method_map = {
        "site": site_trait_linear_model,
        "node": node_trait_linear_model,
        "branch": branch_trait_linear_model,
    }
    return method_map[mode](ts, W, Z, windows=windows, span_normalise=span_normalise)


class TestTraitLinearModel(StatsTestCase, WeightStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def get_example_ts(self):
        ts = msprime.simulate(10, mutation_rate=1, recombination_rate=2, random_seed=1)
        assert ts.num_mutations > 0
        return ts

    def example_covariates(self, ts):
        np.random.seed(999)
        N = ts.num_samples
        for k in [1, 2, 5]:
            k = min(k, ts.num_samples)
            Z = np.ones((N, k))
            Z[1, :] = np.arange(k, 2 * k)
            yield Z
            for j in range(k):
                Z[:, j] = np.random.normal(0, 1, N)
            yield Z

    def transform_weights(self, W, Z):
        n = W.shape[0]
        return np.column_stack([W, Z, np.ones((n, 1))])

    def transform_covariates(self, Z):
        tZ = np.column_stack([Z, np.ones((Z.shape[0], 1))])
        if np.linalg.matrix_rank(tZ) == tZ.shape[1]:
            Z = tZ
        assert np.linalg.matrix_rank(Z) == Z.shape[1]
        K = np.linalg.cholesky(np.matmul(Z.T, Z)).T
        Z = np.matmul(Z, np.linalg.inv(K))
        return Z

    def verify(self, ts):
        for W, Z, windows in subset_combos(
            self.example_weights(ts),
            self.example_covariates(ts),
            example_windows(ts),
            p=0.04,
        ):
            self.verify_trait_linear_model(ts, W, Z, windows=windows)

    def verify_trait_linear_model(self, ts, W, Z, windows):
        n, result_dim = W.shape
        tZ = self.transform_covariates(Z)
        n, k = tZ.shape
        V = np.matmul(W.T, tZ)

        def f(x):
            m = x[-1]
            a = np.zeros(result_dim)
            for i in range(result_dim):
                # print("i=", i, "result_dim=", result_dim, "m=", m, "x=", x)
                # print("V=", V)
                if m > 0 and m < ts.num_samples:
                    v = V[i, :]
                    a[i] = x[i]
                    denom = m
                    for j in range(k):
                        xx = x[result_dim + j]
                        a[i] -= xx * v[j]
                        denom -= xx * xx
                    if abs(denom) < 1e-8:
                        a[i] = 0.0
                    else:
                        a[i] /= denom
                else:
                    a[i] = 0.0
            # print("out", a*a/2)
            return a * a / 2

        # general_stat will need Z added, and an extra column for m
        gW = self.transform_weights(W, tZ)

        def wrapped_summary_func(x):
            with suppress_division_by_zero_warning():
                return f(x)

        # Determine output_dim of the function
        M = len(wrapped_summary_func(gW[0]))
        for sn in [True, False]:
            sigma1 = ts.general_stat(
                gW, wrapped_summary_func, M, windows, mode=self.mode, span_normalise=sn
            )
            sigma2 = general_stat(
                ts, gW, wrapped_summary_func, windows, mode=self.mode, span_normalise=sn
            )
            sigma3 = ts.trait_linear_model(
                W, Z, windows=windows, mode=self.mode, span_normalise=sn
            )
            sigma4 = trait_linear_model(
                ts, W, Z, windows=windows, mode=self.mode, span_normalise=sn
            )

            assert sigma1.shape == sigma2.shape
            assert sigma1.shape == sigma3.shape
            assert sigma1.shape == sigma4.shape
            self.assertArrayAlmostEqual(sigma1, sigma2)
            self.assertArrayAlmostEqual(sigma1, sigma3)
            self.assertArrayAlmostEqual(sigma1, sigma4)


class TraitLinearModelMixin:
    def test_interface(self):
        ts = self.get_example_ts()
        W = np.array([np.arange(ts.num_samples)]).T
        Z = np.ones((ts.num_samples, 1))
        sigma1 = ts.trait_linear_model(W, Z=Z, mode=self.mode)
        sigma2 = ts.trait_linear_model(W, Z=Z, windows=None, mode=self.mode)
        sigma3 = ts.trait_linear_model(
            W, Z=Z, windows=[0.0, ts.sequence_length], mode=self.mode
        )
        sigma4 = ts.trait_linear_model(
            W, Z=None, windows=[0.0, ts.sequence_length], mode=self.mode
        )
        assert sigma1.shape == sigma2.shape
        assert sigma3.shape[0] == 1
        assert sigma1.shape == sigma3.shape[1:]
        assert sigma1.shape == sigma4.shape[1:]
        self.assertArrayAlmostEqual(sigma1, sigma2)
        self.assertArrayAlmostEqual(sigma1, sigma3[0])
        self.assertArrayAlmostEqual(sigma1, sigma4[0])

    def test_errors(self):
        ts = self.get_example_ts()
        W = np.array([np.arange(ts.num_samples)]).T
        Z = np.ones((ts.num_samples, 1))
        # singular covariates
        with pytest.raises(ValueError):
            ts.trait_linear_model(
                W,
                np.ones((ts.num_samples, 2)),
                mode=self.mode,
            )
        # wrong dimensions of W
        with pytest.raises(ValueError):
            ts.trait_linear_model(W[1:, :], Z, mode=self.mode)
        # wrong dimensions of Z
        with pytest.raises(ValueError):
            ts.trait_linear_model(W, Z[1:, :], mode=self.mode)

    def test_deprecation(self):
        ts = self.get_example_ts()
        W = np.array([np.arange(ts.num_samples)]).T
        Z = np.ones((ts.num_samples, 1))
        with pytest.warns(FutureWarning):
            ts.trait_regression(W, Z=Z, mode=self.mode)


@pytest.mark.slow
class TestBranchTraitLinearModel(
    TestTraitLinearModel, TopologyExamplesMixin, TraitLinearModelMixin
):
    mode = "branch"


@pytest.mark.slow
class TestNodeTraitLinearModel(
    TestTraitLinearModel, TopologyExamplesMixin, TraitLinearModelMixin
):
    mode = "node"


class TestSiteTraitLinearModel(
    TestTraitLinearModel, MutatedTopologyExamplesMixin, TraitLinearModelMixin
):
    mode = "site"


##############################
# Sample set statistics
##############################


@pytest.mark.skip(reason="Broken - need to port tests")
class SampleSetStatTestCase(StatsTestCase):
    """
    Provides checks for testing of sample set-based statistics.  Actual testing
    is done by derived classes, which should have attributes `stat_type` and `rng`.
    This works by using parallel structure between different statistic "modes",
    in tree sequence methods (with stat_type=X) and python stat calculators as
    implemented here (with StatCalculator.X).
    """

    random_seed = 123456

    def compare_sfs(self, ts, tree_fn, sample_sets, tsc_fn):
        for sample_set in sample_sets:
            windows = [
                k * ts.sequence_length / 20
                for k in [0] + sorted(self.rng.sample(range(1, 20), 4)) + [20]
            ]
            win_args = [
                {"begin": windows[i], "end": windows[i + 1]}
                for i in range(len(windows) - 1)
            ]
            tree_vals = [tree_fn(sample_set, **b) for b in win_args]

            tsc_vals = tsc_fn(sample_set, windows)
            assert len(tsc_vals) == len(windows) - 1
            for i in range(len(windows) - 1):
                self.assertListAlmostEqual(tsc_vals[i], tree_vals[i])

    def check_sfs_interface(self, ts):
        samples = ts.samples()

        # empty sample sets will raise an error
        with pytest.raises(ValueError):
            ts.site_frequency_spectrum([], self.stat_type)
        # sample_sets must be lists without repeated elements
        with pytest.raises(ValueError):
            ts.site_frequency_spectrum(
                [samples[2], samples[2]],
                self.stat_type,
            )
        # and must all be samples
        with pytest.raises(ValueError):
            ts.site_frequency_spectrum(
                [samples[0], max(samples) + 1],
                self.stat_type,
            )
        # windows must start at 0.0, be increasing, and extend to the end
        with pytest.raises(ValueError):
            ts.site_frequency_spectrum(
                samples[0:2],
                [0.1, ts.sequence_length],
                self.stat_type,
            )
        with pytest.raises(ValueError):
            ts.site_frequency_spectrum(
                samples[0:2],
                [0.0, 0.8 * ts.sequence_length],
                self.stat_type,
            )
        with pytest.raises(ValueError):
            ts.site_frequency_spectrum(
                samples[0:2],
                [
                    0.0,
                    0.8 * ts.sequence_length,
                    0.4 * ts.sequence_length,
                    ts.sequence_length,
                ],
                self.stat_type,
            )

    def check_sfs(self, ts):
        # check site frequency spectrum
        self.check_sfs_interface(ts)
        A = [
            self.rng.sample(list(ts.samples()), 2),
            self.rng.sample(list(ts.samples()), 4),
            self.rng.sample(list(ts.samples()), 8),
            self.rng.sample(list(ts.samples()), 10),
            self.rng.sample(list(ts.samples()), 12),
        ]
        py_tsc = self.py_stat_class(ts)

        self.compare_sfs(
            ts, py_tsc.site_frequency_spectrum, A, ts.site_frequency_spectrum
        )


class BranchSampleSetStatsTestCase(SampleSetStatTestCase):
    """
    Tests of branch statistic computation with sample sets,
    mostly running the checks in SampleSetStatTestCase.
    """

    def setUp(self):
        self.rng = random.Random(self.random_seed)
        self.stat_type = "branch"

    def get_ts(self):
        for N in [12, 15, 20]:
            yield msprime.simulate(
                N, random_seed=self.random_seed, recombination_rate=10
            )

    @pytest.mark.skip(reason="Skipping SFS.")
    def test_sfs_interface(self):
        ts = msprime.simulate(10)
        tsc = tskit.BranchStatCalculator(ts)

        # Duplicated samples raise an error
        with pytest.raises(ValueError):
            tsc.site_frequency_spectrum([1, 1])
        with pytest.raises(ValueError):
            tsc.site_frequency_spectrum([])
        with pytest.raises(ValueError):
            tsc.site_frequency_spectrum([0, 11])
        # Check for bad windows
        for bad_start in [-1, 1, 1e-7]:
            with pytest.raises(ValueError):
                tsc.site_frequency_spectrum(
                    [1, 2],
                    [bad_start, ts.sequence_length],
                )
        for bad_end in [0, ts.sequence_length - 1, ts.sequence_length + 1]:
            with pytest.raises(ValueError):
                tsc.site_frequency_spectrum([1, 2], [0, bad_end])
        # Windows must be increasing.
        with pytest.raises(ValueError):
            tsc.site_frequency_spectrum([1, 2], [0, 1, 1])

    @pytest.mark.skip(reason="No SFS.")
    def test_branch_sfs(self):
        for ts in self.get_ts():
            self.check_sfs(ts)


class SpecificTreesTestCase(StatsTestCase):
    """
    Some particular cases, that are easy to see and debug.
    """

    seed = 21

    def test_case_1(self):
        # With mutations:
        #
        # 1.0          6
        # 0.7         / \                                    5
        #            /   X                                  / \
        # 0.5       X     4                4               /   4
        #          /     / \              / \             /   X X
        # 0.4     X     X   \            X   3           X   /   \
        #        /     /     X          /   / X         /   /     \
        # 0.0   0     1       2        1   0   2       0   1       2
        #          (0.0, 0.2),        (0.2, 0.8),       (0.8, 1.0)
        #
        # haplotypes:
        # site:  0   1 2  3 4         5    6              7   8 9
        # 0:     0   1 0  1 0         0    0              1   0 0
        # 1:     1   0 0  0 1         1    0              0   1 0
        # 2:     1   0 1  0 0         0    1              0   0 1
        branch_true_diversity_01 = 2 * (
            1 * (0.2 - 0) + 0.5 * (0.8 - 0.2) + 0.7 * (1.0 - 0.8)
        )
        branch_true_diversity_02 = 2 * (
            1 * (0.2 - 0) + 0.4 * (0.8 - 0.2) + 0.7 * (1.0 - 0.8)
        )
        branch_true_diversity_12 = 2 * (
            0.5 * (0.2 - 0) + 0.5 * (0.8 - 0.2) + 0.5 * (1.0 - 0.8)
        )
        branch_true_Y = 0.2 * (1 + 0.5) + 0.6 * (0.4) + 0.2 * (0.7 + 0.2)
        site_true_Y = 3 + 0 + 1
        node_true_diversity_012 = (
            np.array(
                [
                    0.2 * np.array([2, 2, 2, 0, 2, 0, 0])
                    + 0.6 * np.array([2, 2, 2, 2, 0, 0, 0])
                    + 0.2 * np.array([2, 2, 2, 0, 2, 0, 0])
                ]
            )
            / 3
        )
        node_true_divergence_0_12 = (
            np.array(
                [
                    0.2 * np.array([2, 1, 1, 0, 2, 0, 0])
                    + 0.6 * np.array([2, 1, 1, 1, 0, 0, 0])
                    + 0.2 * np.array([2, 1, 1, 0, 2, 0, 0])
                ]
            )
            / 2
        )
        haplotypes = np.array(
            [
                [0, 1, 1],
                [1, 0, 0],
                [0, 0, 1],
                [1, 0, 0],
                [0, 1, 0],
                [0, 1, 0],
                [0, 0, 1],
                [1, 0, 0],
                [0, 1, 0],
                [0, 0, 1],
            ]
        )
        traits = np.array([[1, 2, 3, 0], [-5, 0, 1, 1], [3, 4, 1.2, 2]])
        # nb: verified the following with R
        true_cov = (
            np.cov(haplotypes, traits.T)[: haplotypes.shape[0], haplotypes.shape[0] :]
            ** 2
        )
        true_cor = (
            np.corrcoef(haplotypes, traits.T)[
                : haplotypes.shape[0], haplotypes.shape[0] :
            ]
            ** 2
        )
        cov02 = np.cov(np.array([1, 0, 1]), traits.T)[:1, 1:] ** 2
        true_branch_cov = (
            true_cov[1, :] * 1.0 * 0.2
            + true_cov[4, :] * 0.5 * 0.2  # branch 0, tree 0
            + true_cov[2, :] * 0.5 * 0.2  # branch 1, tree 0
            + true_cov[0, :] * 0.5 * 0.2  # branch 2, tree 0
            + true_cov[1, :] * 0.4 * 0.6  # branch 4, tree 0
            + true_cov[4, :] * 0.5 * 0.6  # branch 0, tree 1
            + true_cov[2, :] * 0.4 * 0.6  # branch 1, tree 1
            + cov02 * 0.1 * 0.6  # branch 2, tree 1
            + true_cov[1, :] * 0.7 * 0.2  # branch 3, tree 1
            + true_cov[4, :] * 0.5 * 0.2  # branch 0, tree 2
            + true_cov[2, :] * 0.5 * 0.2  # branch 1, tree 2
            + true_cov[0, :] * 0.2 * 0.2  # branch 2, tree 2
        )  # branch 4, tree 2
        cor02 = np.corrcoef(np.array([1, 0, 1]), traits.T)[:1, 1:] ** 2
        true_branch_cor = (
            true_cor[1, :] * 1.0 * 0.2
            + true_cor[4, :] * 0.5 * 0.2  # branch 0, tree 0
            + true_cor[2, :] * 0.5 * 0.2  # branch 1, tree 0
            + true_cor[0, :] * 0.5 * 0.2  # branch 2, tree 0
            + true_cor[1, :] * 0.4 * 0.6  # branch 4, tree 0
            + true_cor[4, :] * 0.5 * 0.6  # branch 0, tree 1
            + true_cor[2, :] * 0.4 * 0.6  # branch 1, tree 1
            + cor02 * 0.1 * 0.6  # branch 2, tree 1
            + true_cor[1, :] * 0.7 * 0.2  # branch 3, tree 1
            + true_cor[4, :] * 0.5 * 0.2  # branch 0, tree 2
            + true_cor[2, :] * 0.5 * 0.2  # branch 1, tree 2
            + true_cor[0, :] * 0.2 * 0.2  # branch 2, tree 2
        )  # branch 4, tree 2

        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       0           0.4
        4       0           0.5
        5       0           0.7
        6       0           1.0
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.2     0.8     3       0,2
        0.0     0.2     4       1,2
        0.2     0.8     4       1,3
        0.8     1.0     4       1,2
        0.8     1.0     5       0,4
        0.0     0.2     6       0,4
        """
        )
        sites = io.StringIO(
            """\
        id  position    ancestral_state
        0   0.05        0
        1   0.1         0
        2   0.11        0
        3   0.15        0
        4   0.151       0
        5   0.3         0
        6   0.6         0
        7   0.9         0
        8   0.95        0
        9   0.951       0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    derived_state
        0       4       1
        1       0       1
        2       2       1
        3       0       1
        4       1       1
        5       1       1
        6       2       1
        7       0       1
        8       1       1
        9       2       1
        """
        )
        ts = tskit.load_text(
            nodes=nodes, edges=edges, sites=sites, mutations=mutations, strict=False
        )

        # diversity between 0 and 1
        A = [[0], [1]]
        n = [len(a) for a in A]

        def f(x):
            return np.array(
                [float(x[0] * (n[1] - x[1]) + (n[0] - x[0]) * x[1]) / (2 * n[0] * n[1])]
            )

        # tree lengths:
        mode = "branch"
        self.assertAlmostEqual(
            divergence(ts, [[0], [1]], [(0, 1)], mode=mode), branch_true_diversity_01
        )
        self.assertAlmostEqual(
            ts.divergence([[0], [1]], [(0, 1)], mode=mode), branch_true_diversity_01
        )
        self.assertAlmostEqual(
            ts.sample_count_stat(A, f, 1, mode=mode)[0], branch_true_diversity_01
        )
        self.assertAlmostEqual(
            ts.diversity([[0, 1]], mode=mode)[0], branch_true_diversity_01
        )

        # mean diversity between [0, 1] and [0, 2]:
        branch_true_mean_diversity = (
            0
            + branch_true_diversity_02
            + branch_true_diversity_01
            + branch_true_diversity_12
        ) / 4
        A = [[0, 1], [0, 2]]
        n = [len(a) for a in A]

        def f(x):
            return np.array([float(x[0] * (n[1] - x[1]) + (n[0] - x[0]) * x[1]) / 8.0])

        # tree lengths:
        self.assertAlmostEqual(
            divergence(ts, [A[0], A[1]], [(0, 1)], mode=mode),
            branch_true_mean_diversity,
        )
        self.assertAlmostEqual(
            ts.divergence([A[0], A[1]], [(0, 1)], mode=mode), branch_true_mean_diversity
        )
        self.assertAlmostEqual(
            ts.sample_count_stat(A, f, 1, mode=mode)[0], branch_true_mean_diversity
        )

        # Y-statistic for (0/12)
        A = [[0], [1, 2]]

        def f(x):
            return np.array(
                [
                    float(
                        ((x[0] == 1) and (x[1] == 0)) or ((x[0] == 0) and (x[1] == 2))
                    )
                    / 2.0
                ]
            )

        # tree lengths:
        bts_Y = ts.Y3([[0], [1], [2]], mode=mode)
        py_bsc_Y = Y3(ts, [[0], [1], [2]], [(0, 1, 2)], windows=[0.0, 1.0], mode=mode)
        self.assertArrayAlmostEqual(bts_Y, branch_true_Y)
        self.assertArrayAlmostEqual(py_bsc_Y, branch_true_Y)
        self.assertArrayAlmostEqual(
            ts.sample_count_stat(A, f, 1, mode=mode)[0], branch_true_Y
        )

        mode = "site"
        # sites, Y:
        sts_Y = ts.Y3([[0], [1], [2]], mode=mode)
        py_ssc_Y = Y3(ts, [[0], [1], [2]], [(0, 1, 2)], windows=[0.0, 1.0], mode=mode)
        self.assertArrayAlmostEqual(sts_Y, site_true_Y)
        self.assertArrayAlmostEqual(py_ssc_Y, site_true_Y)
        self.assertArrayAlmostEqual(
            ts.sample_count_stat(A, f, 1, mode=mode)[0], site_true_Y
        )

        A = [[0, 1, 2]]
        n = 3
        W = np.array([[u in A[0]] for u in ts.samples()], dtype=float)

        def f(x):
            return np.array([x[0] * (n - x[0]) / (n * (n - 1))])

        mode = "node"
        # nodes, diversity in [0,1,2]
        nodes_div_012 = ts.diversity([[0, 1, 2]], mode=mode).reshape((1, 7))
        py_nodes_div_012 = diversity(ts, [[0, 1, 2]], mode=mode).reshape((1, 7))
        py_general_nodes_div_012 = general_stat(ts, W, f, mode=mode).reshape((1, 7))
        self.assertArrayAlmostEqual(py_nodes_div_012, node_true_diversity_012)
        self.assertArrayAlmostEqual(py_general_nodes_div_012, node_true_diversity_012)
        self.assertArrayAlmostEqual(nodes_div_012, node_true_diversity_012)

        # nodes, divergence [0] to [1,2]
        nodes_div_0_12 = ts.divergence([[0], [1, 2]], mode=mode).reshape((1, 7))
        py_nodes_div_0_12 = divergence(ts, [[0], [1, 2]], mode=mode).reshape((1, 7))
        self.assertArrayAlmostEqual(nodes_div_0_12, node_true_divergence_0_12)
        self.assertArrayAlmostEqual(py_nodes_div_0_12, node_true_divergence_0_12)

        # covariance and correlation
        ts_sitewise_cov = ts.trait_covariance(
            traits, mode="site", windows="sites", span_normalise=False
        )
        py_sitewise_cov = site_trait_covariance(
            ts, traits, windows="sites", span_normalise=False
        )
        self.assertArrayAlmostEqual(py_sitewise_cov, true_cov)
        self.assertArrayAlmostEqual(ts_sitewise_cov, true_cov)
        ts_sitewise_cor = ts.trait_correlation(
            traits, mode="site", windows="sites", span_normalise=False
        )
        py_sitewise_cor = site_trait_correlation(
            ts, traits, windows="sites", span_normalise=False
        )
        self.assertArrayAlmostEqual(py_sitewise_cor, true_cor)
        self.assertArrayAlmostEqual(ts_sitewise_cor, true_cor)
        # mean
        ts_mean_cov = ts.trait_covariance(
            traits, mode="site", windows=[0, ts.sequence_length]
        )
        py_mean_cov = site_trait_covariance(ts, traits)
        self.assertArrayAlmostEqual(ts_mean_cov, np.array([np.sum(true_cov, axis=0)]))
        self.assertArrayAlmostEqual(ts_mean_cov, py_mean_cov)
        ts_mean_cor = ts.trait_correlation(
            traits, mode="site", windows=[0, ts.sequence_length]
        )
        py_mean_cor = site_trait_correlation(ts, traits)
        self.assertArrayAlmostEqual(ts_mean_cor, np.array([np.sum(true_cor, axis=0)]))
        self.assertArrayAlmostEqual(ts_mean_cor, py_mean_cor)
        # mode = 'branch'
        ts_mean_cov = ts.trait_covariance(
            traits, mode="branch", windows=[0, ts.sequence_length]
        )
        py_mean_cov = branch_trait_covariance(ts, traits)
        self.assertArrayAlmostEqual(ts_mean_cov, true_branch_cov)
        self.assertArrayAlmostEqual(ts_mean_cov, py_mean_cov)
        ts_mean_cor = ts.trait_correlation(
            traits, mode="branch", windows=[0, ts.sequence_length]
        )
        py_mean_cor = branch_trait_correlation(ts, traits)
        self.assertArrayAlmostEqual(ts_mean_cor, true_branch_cor)
        self.assertArrayAlmostEqual(ts_mean_cor, py_mean_cor)

        # trait linear_model:
        # r = cor * sd(y) / sd(x) = cov / var(x)
        # geno_var = allele_freqs * (1 - allele_freqs) * (3 / (3 - 1))
        geno_var = np.var(haplotypes, axis=1) * (3 / (3 - 1))
        trait_var = np.var(traits, axis=0) * (3 / (3 - 1))
        py_r = trait_linear_model(
            ts, traits, None, mode="site", windows="sites", span_normalise=False
        )
        ts_r = ts.trait_linear_model(
            traits, None, mode="site", windows="sites", span_normalise=False
        )
        self.assertArrayAlmostEqual(py_r, ts_r)
        self.assertArrayAlmostEqual(true_cov, py_r * (geno_var[:, np.newaxis] ** 2))
        self.assertArrayAlmostEqual(
            true_cor, ts_r * geno_var[:, np.newaxis] / trait_var
        )

    def test_case_odds_and_ends(self):
        # Tests having (a) the first site after the first window, and
        # (b) no samples having the ancestral state.
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           0.5
        3       0           1.0
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.0     0.5     2       0,1
        0.5     1.0     3       0,1
        """
        )
        sites = io.StringIO(
            """\
        id  position    ancestral_state
        0   0.65        0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    derived_state   parent
        0       0       1               -1
        0       1       2               -1
        """
        )
        ts = tskit.load_text(
            nodes=nodes, edges=edges, sites=sites, mutations=mutations, strict=False
        )

        mode = "site"
        py_div = divergence(
            ts, [[0], [1]], indexes=[(0, 1)], windows=[0.0, 0.5, 1.0], mode=mode
        )
        div = ts.divergence(
            [[0], [1]], indexes=[(0, 1)], windows=[0.0, 0.5, 1.0], mode=mode
        )
        self.assertArrayEqual(py_div, div)

    def test_case_four_taxa(self):
        #
        # 1.0          7
        # 0.7         / \                                    6
        #            /   \                                  / \
        # 0.5       /     5              5                 /   5
        #          /     / \            / \__             /   / \
        # 0.4     /     8   \          8     4           /   8   \
        #        /     / \   \        / \   / \         /   / \   \
        # 0.0   0     1   3   2      1   3 0   2       0   1   3   2
        #          (0.0, 0.2),        (0.2, 0.8),       (0.8, 2.5)

        # f4(0, 1, 2, 3): (0 -> 1)(2 -> 3)
        branch_true_f4_0123 = (0.1 * 0.2 + (0.1 + 0.1) * 0.6 + 0.1 * 1.7) / 2.5
        windows = [0.0, 0.4, 2.5]
        branch_true_f4_0123_windowed = np.array(
            [
                (0.1 * 0.2 + (0.1 + 0.1) * 0.2) / 0.4,
                ((0.1 + 0.1) * 0.4 + 0.1 * 1.7) / 2.1,
            ]
        )
        # f4(0, 3, 2, 1): (0 -> 3)(2 -> 1)
        branch_true_f4_0321 = (0.1 * 0.2 + (0.1 + 0.1) * 0.6 + 0.1 * 1.7) / 2.5
        # f2([0,2], [1,3]) = (1/2) (f4(0,1,2,3) + f4(0,3,2,1))
        branch_true_f2_02_13 = (branch_true_f4_0123 + branch_true_f4_0321) / 2
        # diversity([0,1,2,3])
        branch_true_diversity_windowed = (2 / 6) * np.array(
            [
                [
                    (
                        0.2 * (1 + 1 + 1 + 0.5 + 0.4 + 0.5)
                        + (0.4 - 0.2) * (0.5 + 0.4 + 0.5 + 0.5 + 0.4 + 0.5)
                    )
                    / 0.4
                ],
                [
                    (
                        (0.8 - 0.4) * (0.5 + 0.4 + 0.5 + 0.5 + 0.4 + 0.5)
                        + (2.5 - 0.8) * (0.7 + 0.7 + 0.7 + 0.5 + 0.4 + 0.5)
                    )
                    / (2.5 - 0.4)
                ],
            ]
        )

        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       1           0
        4       0           0.4
        5       0           0.5
        6       0           0.7
        7       0           1.0
        8       0           0.4
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.0     2.5     8       1,3
        0.2     0.8     4       0,2
        0.0     0.2     5       8,2
        0.2     0.8     5       8,4
        0.8     2.5     5       8,2
        0.8     2.5     6       0,5
        0.0     0.2     7       0,5
        """
        )
        sites = io.StringIO(
            """\
        id  position    ancestral_state
        """
        )
        mutations = io.StringIO(
            """\
        site    node    derived_state   parent
        """
        )
        ts = tskit.load_text(
            nodes=nodes, edges=edges, sites=sites, mutations=mutations, strict=False
        )

        mode = "branch"
        A = [[0], [1], [2], [3]]
        self.assertAlmostEqual(branch_true_f4_0123, f4(ts, A, mode=mode)[0][0])
        self.assertAlmostEqual(branch_true_f4_0123, ts.f4(A, mode=mode))
        self.assertArrayAlmostEqual(
            branch_true_f4_0123_windowed, ts.f4(A, windows=windows, mode=mode).flatten()
        )
        A = [[0], [3], [2], [1]]
        self.assertAlmostEqual(
            branch_true_f4_0321, f4(ts, A, [(0, 1, 2, 3)], mode=mode)[0][0]
        )
        self.assertAlmostEqual(branch_true_f4_0321, ts.f4(A, mode=mode))
        A = [[0], [2], [1], [3]]
        self.assertAlmostEqual(0.0, f4(ts, A, [(0, 1, 2, 3)], mode=mode)[0])
        self.assertAlmostEqual(0.0, ts.f4(A, mode=mode))
        A = [[0, 2], [1, 3]]
        self.assertAlmostEqual(
            branch_true_f2_02_13, f2(ts, A, [(0, 1)], mode=mode)[0][0]
        )
        self.assertAlmostEqual(branch_true_f2_02_13, ts.f2(A, mode=mode))

        # diversity
        A = [[0, 1, 2, 3]]
        self.assertArrayAlmostEqual(
            branch_true_diversity_windowed, diversity(ts, A, windows=windows, mode=mode)
        )
        self.assertArrayAlmostEqual(
            branch_true_diversity_windowed, ts.diversity(A, windows=windows, mode=mode)
        )

    def test_case_recurrent_muts(self):
        # With mutations:
        #
        # 1.0          6
        # 0.7         / \                                    5
        #           (0)  \                                  /(6)
        # 0.5      (1)    4                4               /   4
        #          /     / \              / \             /  (7|8)
        # 0.4    (2)   (3)  \           (4)  3           /   /   \
        #        /     /     \          /   /(5)        /   /     \
        # 0.0   0     1       2        1   0   2       0   1       2
        #          (0.0, 0.2),        (0.2, 0.8),       (0.8, 1.0)
        # genotypes:
        #       0     2       0        1   0   1       0   2       3
        site_true_Y = 0 + 1 + 1

        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       0           0.4
        4       0           0.5
        5       0           0.7
        6       0           1.0
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.2     0.8     3       0,2
        0.0     0.2     4       1,2
        0.2     0.8     4       1,3
        0.8     1.0     4       1,2
        0.8     1.0     5       0,4
        0.0     0.2     6       0,4
        """
        )
        sites = io.StringIO(
            """\
        id  position    ancestral_state
        0   0.05        0
        1   0.3         0
        2   0.9         0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    derived_state   parent
        0       0       1               -1
        0       0       2               0
        0       0       0               1
        0       1       2               -1
        1       1       1               -1
        1       2       1               -1
        2       4       1               -1
        2       1       2               6
        2       2       3               6
        """
        )
        ts = tskit.load_text(
            nodes=nodes, edges=edges, sites=sites, mutations=mutations, strict=False
        )

        # Y3:
        site_tsc_Y = ts.Y3([[0], [1], [2]], mode="site")
        py_ssc_Y = Y3(ts, [[0], [1], [2]], [(0, 1, 2)], windows=[0.0, 1.0], mode="site")
        self.assertAlmostEqual(site_tsc_Y, site_true_Y)
        self.assertAlmostEqual(py_ssc_Y, site_true_Y)

    def test_case_2(self):
        # Here are the trees:
        # t                  |              |              |             |            |
        #
        # 0       --3--      |     --3--    |     --3--    |    --3--    |    --3--   |
        #        /  |  \     |    /  |  \   |    /     \   |   /     \   |   /     \  |
        # 1     4   |   5    |   4   |   5  |   4       5  |  4       5  |  4       5 |
        #       |\ / \ /|    |   |\   \     |   |\     /   |  |\     /   |  |\     /| |
        # 2     | 6   7 |    |   | 6   7    |   | 6   7    |  | 6   7    |  | 6   7 | |
        #       | |\ /| |    |   *  \  |    |   |  \  |    |  |  \       |  |  \    | |
        # 3     | | 8 | |    |   |   8 *    |   |   8 |    |  |   8      |  |   8   | |
        #       | |/ \| |    |   |  /  |    |   |  /  |    |  |  / \     |  |  / \  | |
        # 4     | 9  10 |    |   * 9  10    |   | 9  10    |  | 9  10    |  | 9  10 | |
        #       |/ \ / \|    |   |  \   \   |   |  \   \   |  |  \   \   |  |  \    | |
        # 5     0   1   2    |   0   1   2  |   0   1   2  |  0   1   2  |  0   1   2 |
        #
        #                    |   0.0 - 0.1  |   0.1 - 0.2  |  0.2 - 0.4  |  0.4 - 0.5 |
        # ... continued:
        # t                  |             |             |             |
        #
        # 0         --3--    |    --3--    |    --3--    |    --3--    |    --3--
        #          /     \   |   /     \   |   /     \   |   /     \   |   /  |  \
        # 1       4       5  |  4       5  |  4       5  |  4       5  |  4   |   5
        #         |\     /|  |   \     /|  |   \     /|  |   \     /|  |     /   /|
        # 2       | 6   7 |  |    6   7 |  |    6   7 |  |    6   7 |  |    6   7 |
        #         |  *    *  |     \    |  |       *  |  |    |  /  |  |    |  /  |
        # 3  ...  |   8   |  |      8   |  |      8   |  |    | 8   |  |    | 8   |
        #         |  / \  |  |     / \  |  |     * \  |  |    |  \  |  |    |  \  |
        # 4       | 9  10 |  |    9  10 |  |    9  10 |  |    9  10 |  |    9  10 |
        #         |    /  |  |   /   /  |  |   /   /  |  |   /   /  |  |   /   /  |
        # 5       0   1   2  |  0   1   2  |  0   1   2  |  0   1   2  |  0   1   2
        #
        #         0.5 - 0.6  |  0.6 - 0.7  |  0.7 - 0.8  |  0.8 - 0.9  |  0.9 - 1.0
        #
        # Above, subsequent mutations are backmutations.

        # divergence betw 0 and 1
        branch_true_diversity_01 = 2 * (0.6 * 4 + 0.2 * 2 + 0.2 * 5)
        # divergence betw 1 and 2
        branch_true_diversity_12 = 2 * (0.2 * 5 + 0.2 * 2 + 0.3 * 5 + 0.3 * 4)
        # divergence betw 0 and 2
        branch_true_diversity_02 = 2 * (0.2 * 5 + 0.2 * 4 + 0.3 * 5 + 0.1 * 4 + 0.2 * 5)
        # Y(0;1, 2)
        branch_true_Y = 0.2 * 4 + 0.2 * (4 + 2) + 0.2 * 4 + 0.2 * 2 + 0.2 * (5 + 1)

        # site stats
        # Y(0;1, 2)
        site_true_Y = 1

        nodes = io.StringIO(
            """\
        is_sample       time    population
        1       0.000000        0
        1       0.000000        0
        1       0.000000        0
        0       5.000000        0
        0       4.000000        0
        0       4.000000        0
        0       3.000000        0
        0       3.000000        0
        0       2.000000        0
        0       1.000000        0
        0       1.000000        0
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.500000        1.000000        10      1
        0.000000        0.400000        10      2
        0.600000        1.000000        9       0
        0.000000        0.500000        9       1
        0.800000        1.000000        8       10
        0.200000        0.800000        8       9,10
        0.000000        0.200000        8       9
        0.700000        1.000000        7       8
        0.000000        0.200000        7       10
        0.800000        1.000000        6       9
        0.000000        0.700000        6       8
        0.400000        1.000000        5       2,7
        0.100000        0.400000        5       7
        0.600000        0.900000        4       6
        0.000000        0.600000        4       0,6
        0.900000        1.000000        3       4,5,6
        0.100000        0.900000        3       4,5
        0.000000        0.100000        3       4,5,7
        """
        )
        sites = io.StringIO(
            """\
        id  position    ancestral_state
        0   0.0         0
        1   0.55        0
        2   0.75        0
        3   0.85        0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    derived_state   parent
        0       0       1               -1
        0       10      1               -1
        0       0       0               0
        1       8       1               -1
        1       2       1               -1
        2       8       1               -1
        2       9       0               5
        """
        )
        ts = tskit.load_text(
            nodes=nodes, edges=edges, sites=sites, mutations=mutations, strict=False
        )

        def f(x):
            return np.array([float(x[0] == 1) / 2.0])

        # divergence between 0 and 1
        mode = "branch"
        for A, truth in zip(
            [[[0, 1]], [[1, 2]], [[0, 2]]],
            [
                branch_true_diversity_01,
                branch_true_diversity_12,
                branch_true_diversity_02,
            ],
        ):
            self.assertAlmostEqual(diversity(ts, A, mode=mode)[0][0], truth)
            self.assertAlmostEqual(ts.sample_count_stat(A, f, 1, mode=mode)[0], truth)
            self.assertAlmostEqual(ts.diversity(A, mode="branch")[0], truth)

        # Y-statistic for (0/12)
        A = [[0], [1, 2]]

        def f(x):
            return np.array(
                [
                    float(
                        ((x[0] == 1) and (x[1] == 0)) or ((x[0] == 0) and (x[1] == 2))
                    )
                    / 2.0
                ]
            )

        # tree lengths:
        self.assertArrayAlmostEqual(
            Y3(ts, [[0], [1], [2]], [(0, 1, 2)], mode=mode), branch_true_Y
        )
        self.assertArrayAlmostEqual(
            ts.Y3([[0], [1], [2]], [(0, 1, 2)], mode=mode), branch_true_Y
        )
        self.assertArrayAlmostEqual(
            ts.sample_count_stat(A, f, 1, mode=mode)[0], branch_true_Y
        )

        # sites:
        mode = "site"
        site_tsc_Y = ts.Y3([[0], [1], [2]], mode=mode)
        py_ssc_Y = Y3(ts, [[0], [1], [2]], [(0, 1, 2)], windows=[0.0, 1.0])
        self.assertAlmostEqual(site_tsc_Y, site_true_Y)
        self.assertAlmostEqual(py_ssc_Y, site_true_Y)
        self.assertAlmostEqual(ts.sample_count_stat(A, f, 1, mode=mode)[0], site_true_Y)


class TestOutputDimensions(StatsTestCase):
    """
    Tests for the dimension stripping behaviour of the stats functions.
    """

    def get_example_ts(self):
        ts = msprime.simulate(10, mutation_rate=1, random_seed=1)
        assert ts.num_sites > 1
        return ts

    def test_one_way_no_window_scalar_stat(self):
        ts = self.get_example_ts()
        x = ts.diversity()
        assert isinstance(x, np.floating)

    def test_one_way_one_list_scalar_stat(self):
        ts = self.get_example_ts()
        x = ts.diversity(sample_sets=list(ts.samples()))
        assert isinstance(x, np.floating)

    def test_one_way_nested_list_not_scalar_stat(self):
        ts = self.get_example_ts()
        x = ts.diversity(sample_sets=[list(ts.samples())])
        assert x.shape == (1,)

    def test_one_way_one_window_scalar_stat(self):
        ts = self.get_example_ts()
        x = ts.diversity(windows=[0, ts.sequence_length])
        assert x.shape == (1,)
        for samples in (None, list(ts.samples())):
            x = ts.diversity(sample_sets=samples, windows=[0, ts.sequence_length])
            assert x.shape == (1,)

    def test_multi_way_no_window_scalar_stat(self):
        ts = self.get_example_ts()
        n = ts.num_samples
        x = ts.f2(
            sample_sets=[
                [i for i in range(0, int(n / 2))],
                [i for i in range(int(n / 2), n)],
            ]
        )
        assert isinstance(x, np.floating)

    def test_multi_way_one_window_not_scalar_stat(self):
        ts = self.get_example_ts()
        n = ts.num_samples
        x = ts.f2(
            sample_sets=[
                [i for i in range(0, int(n / 2))],
                [i for i in range(int(n / 2), n)],
            ],
            windows=[0, ts.sequence_length],
        )
        assert x.shape == (1,)

    def test_multi_way_no_indexes_scalar_stat(self):
        ts = self.get_example_ts()
        n = ts.num_samples
        x = ts.f2(
            sample_sets=[
                [i for i in range(0, int(n / 2))],
                [i for i in range(int(n / 2), n)],
            ],
        )
        assert isinstance(x, np.floating)

    def test_multi_way_indexes_not_scalar_stat(self):
        ts = self.get_example_ts()
        n = ts.num_samples
        x = ts.f2(
            sample_sets=[
                [i for i in range(0, int(n / 2))],
                [i for i in range(int(n / 2), n)],
            ],
            indexes=[(0, 1)],
        )
        assert x.shape == (1,)

    def test_afs_default_windows(self):
        ts = self.get_example_ts()
        n = ts.num_samples
        A = ts.samples()[:4]
        B = ts.samples()[6:]
        for mode in ["site", "branch"]:
            x = ts.allele_frequency_spectrum(mode=mode)
            # x is a 1D numpy array with n + 1 values
            assert x.shape == (n + 1,)
            self.assertArrayEqual(
                x, ts.allele_frequency_spectrum([ts.samples()], mode=mode)
            )
            x = ts.allele_frequency_spectrum([A, B], mode=mode)
            assert x.shape == (len(A) + 1, len(B) + 1)

    def test_afs_windows(self):
        ts = self.get_example_ts()
        L = ts.sequence_length

        windows = [0, L / 4, L / 2, L]
        A = ts.samples()[:4]
        B = ts.samples()[6:]
        for mode in ["site", "branch"]:
            x = ts.allele_frequency_spectrum([A, B], windows=windows, mode=mode)
            assert x.shape == (3, len(A) + 1, len(B) + 1)

            x = ts.allele_frequency_spectrum([A], windows=windows, mode=mode)
            assert x.shape == (3, len(A) + 1)

            x = ts.allele_frequency_spectrum(windows=windows, mode=mode)
            # Default returns this for all samples
            assert x.shape == (3, ts.num_samples + 1)
            y = ts.allele_frequency_spectrum([ts.samples()], windows=windows, mode=mode)
            self.assertArrayEqual(x, y)

    def test_one_way_stat_default_windows(self):
        ts = self.get_example_ts()
        # Use diversity as the example one-way stat.
        for mode in ["site", "branch"]:
            x = ts.diversity(mode=mode)
            # x is a zero-d numpy value
            assert np.shape(x) == tuple()
            assert x == float(x)
            assert x == ts.diversity(ts.samples(), mode=mode)
            self.assertArrayEqual([x], ts.diversity([ts.samples()], mode=mode))

        mode = "node"
        x = ts.diversity(mode=mode)
        # x is a 1D numpy array with N values
        assert x.shape == (ts.num_nodes,)
        self.assertArrayEqual(x, ts.diversity(ts.samples(), mode=mode))
        y = ts.diversity([ts.samples()], mode=mode)
        # We're adding on the *last* dimension, so must reshape
        self.assertArrayEqual(x.reshape(ts.num_nodes, 1), y)

    def verify_one_way_stat_windows(self, ts, method):
        L = ts.sequence_length
        N = ts.num_nodes

        windows = [0, L / 4, L / 2, 0.75 * L, L]
        A = ts.samples()[:6]
        B = ts.samples()[6:]
        for mode in ["site", "branch"]:
            x = method([A, B], windows=windows, mode=mode)
            # Four windows, 2 sets.
            assert x.shape == (4, 2)

            x = method([A], windows=windows, mode=mode)
            # Four windows, 1 sets.
            assert x.shape == (4, 1)

            x = method(A, windows=windows, mode=mode)
            # Dropping the outer list removes the last dimension
            assert x.shape == (4,)

            x = method(windows=windows, mode=mode)
            # Default returns this for all samples
            assert x.shape == (4,)
            y = method(ts.samples(), windows=windows, mode=mode)
            self.assertArrayEqual(x, y)

        mode = "node"
        x = method([A, B], windows=windows, mode=mode)
        # Four windows, N nodes and 2 sets.
        assert x.shape == (4, N, 2)

        x = method([A], windows=windows, mode=mode)
        # Four windows, N nodes and 1 set.
        assert x.shape == (4, N, 1)

        x = method(A, windows=windows, mode=mode)
        # Drop the outer list, so we lose the last dimension
        assert x.shape == (4, N)

        x = method(windows=windows, mode=mode)
        # The default sample sets also drops the last dimension
        assert x.shape == (4, N)

        assert ts.num_trees == 1
        # In this example, we know that the trees are all the same so check this
        # for sanity.
        self.assertArrayEqual(x[0], x[1])
        self.assertArrayEqual(x[0], x[2])

    def test_diversity_windows(self):
        ts = self.get_example_ts()
        self.verify_one_way_stat_windows(ts, ts.diversity)

    def test_Tajimas_D_windows(self):
        ts = self.get_example_ts()
        self.verify_one_way_stat_windows(ts, ts.Tajimas_D)

    def test_segregating_sites_windows(self):
        ts = self.get_example_ts()
        self.verify_one_way_stat_windows(ts, ts.segregating_sites)

    def test_two_way_stat_default_windows(self):
        ts = self.get_example_ts()
        # Use divergence as the example one-way stat.
        A = ts.samples()[:6]
        B = ts.samples()[6:]
        for mode in ["site", "branch"]:
            x = ts.divergence([A, B], mode=mode)
            # x is a zero-d numpy value
            assert np.shape(x) == tuple()
            assert x == float(x)
            # If indexes is a 1D array, we also drop the outer dimension
            assert x == ts.divergence([A, B, A], indexes=[0, 1], mode=mode)
            # But, if it's a 2D array we keep the outer dimension
            assert [x] == ts.divergence([A, B], indexes=[[0, 1]], mode=mode)

        mode = "node"
        x = ts.divergence([A, B], mode=mode)
        # x is a 1D numpy array with N values
        assert x.shape == (ts.num_nodes,)
        self.assertArrayEqual(x, ts.divergence([A, B], indexes=[0, 1], mode=mode))
        y = ts.divergence([A, B], indexes=[[0, 1]], mode=mode)
        # We're adding on the *last* dimension, so must reshape
        self.assertArrayEqual(x.reshape(ts.num_nodes, 1), y)

    def verify_two_way_stat_windows(self, ts, method):
        L = ts.sequence_length
        N = ts.num_nodes

        windows = [0, L / 4, L / 2, L]
        A = ts.samples()[:7]
        B = ts.samples()[7:]
        for mode in ["site", "branch"]:
            x = method([A, B, A], indexes=[[0, 1], [0, 2]], windows=windows, mode=mode)
            # Three windows, 2 pairs
            assert x.shape == (3, 2)

            x = method([A, B], indexes=[[0, 1]], windows=windows, mode=mode)
            # Three windows, 1 pair
            assert x.shape == (3, 1)

            x = method([A, B], indexes=[0, 1], windows=windows, mode=mode)
            # Dropping the outer list removes the last dimension
            assert x.shape == (3,)

            y = method([A, B], windows=windows, mode=mode)
            assert y.shape == (3,)
            self.assertArrayEqual(x, y)

        mode = "node"
        x = method([A, B], indexes=[[0, 1], [0, 1]], windows=windows, mode=mode)
        # Three windows, N nodes and 2 pairs
        assert x.shape == (3, N, 2)

        x = method([A, B], indexes=[[0, 1]], windows=windows, mode=mode)
        # Three windows, N nodes and 1 pairs
        assert x.shape == (3, N, 1)

        x = method([A, B], indexes=[0, 1], windows=windows, mode=mode)
        # Drop the outer list, so we lose the last dimension
        assert x.shape == (3, N)

        x = method([A, B], windows=windows, mode=mode)
        # The default sample sets also drops the last dimension
        assert x.shape == (3, N)

        assert ts.num_trees == 1
        # In this example, we know that the trees are all the same so check this
        # for sanity.
        self.assertArrayEqual(x[0], x[1])
        self.assertArrayEqual(x[0], x[2])

    def test_divergence_windows(self):
        ts = self.get_example_ts()
        self.verify_two_way_stat_windows(ts, ts.divergence)

    def test_Fst_windows(self):
        ts = self.get_example_ts()
        self.verify_two_way_stat_windows(ts, ts.Fst)

    def test_f2_windows(self):
        ts = self.get_example_ts()
        self.verify_two_way_stat_windows(ts, ts.f2)

    def verify_three_way_stat_windows(self, ts, method):
        L = ts.sequence_length
        N = ts.num_nodes

        windows = [0, L / 4, L / 2, L]
        A = ts.samples()[:2]
        B = ts.samples()[2:4]
        C = ts.samples()[4:]
        for mode in ["site", "branch"]:
            x = method(
                [A, B, C], indexes=[[0, 1, 2], [0, 2, 1]], windows=windows, mode=mode
            )
            # Three windows, 2 triple
            assert x.shape == (3, 2)

            x = method([A, B, C], indexes=[[0, 1, 2]], windows=windows, mode=mode)
            # Three windows, 1 triple
            assert x.shape == (3, 1)

            x = method([A, B, C], indexes=[0, 1, 2], windows=windows, mode=mode)
            # Dropping the outer list removes the last dimension
            assert x.shape == (3,)

            y = method([A, B, C], windows=windows, mode=mode)
            assert y.shape == (3,)
            self.assertArrayEqual(x, y)

        mode = "node"
        x = method(
            [A, B, C], indexes=[[0, 1, 2], [0, 2, 1]], windows=windows, mode=mode
        )
        # Three windows, N nodes and 2 triples
        assert x.shape == (3, N, 2)

        x = method([A, B, C], indexes=[[0, 1, 2]], windows=windows, mode=mode)
        # Three windows, N nodes and 1 triples
        assert x.shape == (3, N, 1)

        x = method([A, B, C], indexes=[0, 1, 2], windows=windows, mode=mode)
        # Drop the outer list, so we lose the last dimension
        assert x.shape == (3, N)

        x = method([A, B, C], windows=windows, mode=mode)
        # The default sample sets also drops the last dimension
        assert x.shape == (3, N)

        assert ts.num_trees == 1
        # In this example, we know that the trees are all the same so check this
        # for sanity.
        self.assertArrayEqual(x[0], x[1])
        self.assertArrayEqual(x[0], x[2])

    def test_Y3_windows(self):
        ts = self.get_example_ts()
        self.verify_three_way_stat_windows(ts, ts.Y3)

    def test_f3_windows(self):
        ts = self.get_example_ts()
        self.verify_three_way_stat_windows(ts, ts.f3)


class TestTimeUncalibratedErrors:
    def test_uncalibrated_time_allele_frequency_spectrum(self, ts_fixture):
        ts_fixture.allele_frequency_spectrum(mode="branch")
        tables = ts_fixture.dump_tables()
        tables.time_units = tskit.TIME_UNITS_UNCALIBRATED
        ts_uncalibrated = tables.tree_sequence()
        ts_uncalibrated.allele_frequency_spectrum(mode="site")
        with pytest.raises(
            tskit.LibraryError,
            match="Statistics using branch lengths cannot be calculated when time_units"
            " is 'uncalibrated'",
        ):
            ts_uncalibrated.allele_frequency_spectrum(mode="branch")

    def test_uncalibrated_time_general_stat(self, ts_fixture):
        W = np.ones((ts_fixture.num_samples, 2))
        ts_fixture.general_stat(
            W, lambda x: x * (x < ts_fixture.num_samples), W.shape[1], mode="branch"
        )
        tables = ts_fixture.dump_tables()
        tables.time_units = tskit.TIME_UNITS_UNCALIBRATED
        ts_uncalibrated = tables.tree_sequence()
        ts_uncalibrated.general_stat(
            W, lambda x: x * (x < ts_uncalibrated.num_samples), W.shape[1], mode="site"
        )
        with pytest.raises(
            tskit.LibraryError,
            match="Statistics using branch lengths cannot be calculated when time_units"
            " is 'uncalibrated'",
        ):
            ts_uncalibrated.general_stat(
                W,
                lambda x: x * (x < ts_uncalibrated.num_samples),
                W.shape[1],
                mode="branch",
            )


class TestGeneralStatCallbackErrors:
    def test_zero_d(self, ts_fixture):
        def f_0d(_):
            return 0

        msg = "Array returned by general_stat callback is 0 dimensional; must be 1D"
        with pytest.raises(ValueError, match=msg):
            ts_fixture.sample_count_stat(
                sample_sets=[ts_fixture.samples()], f=f_0d, output_dim=1, strict=False
            )

    def test_two_d(self, ts_fixture):
        def f_2d(x):
            return np.array([x])

        msg = "Array returned by general_stat callback is 2 dimensional; must be 1D"
        with pytest.raises(ValueError, match=msg):
            ts_fixture.sample_count_stat(
                sample_sets=[ts_fixture.samples()], f=f_2d, output_dim=1, strict=False
            )

    def test_wrong_length(self, ts_fixture):
        def f_too_long(_):
            return np.array([0, 0])

        msg = "Array returned by general_stat callback is of length 2; must be 1"
        with pytest.raises(ValueError, match=msg):
            ts_fixture.sample_count_stat(
                sample_sets=[ts_fixture.samples()],
                f=f_too_long,
                output_dim=1,
                strict=False,
            )


--- ../../tskit/python/tests/test_tables.py ---


"""
Test cases for the low-level tables used to transfer information
between simulations and the tree sequence.
"""
import dataclasses
import io
import json
import math
import pathlib
import pickle
import platform
import random
import re
import struct
import time
import unittest
import warnings

import kastore
import msprime
import numpy as np
import pytest

import _tskit
import tests.test_wright_fisher as wf
import tests.tsutil as tsutil
import tskit
import tskit.exceptions as exceptions
import tskit.metadata as metadata


class Column:
    def __init__(self, name):
        self.name = name


class Int32Column(Column):
    def get_input(self, n):
        return 1 + np.arange(n, dtype=np.int32)


class UInt8Column(Column):
    def get_input(self, n):
        return 2 + np.arange(n, dtype=np.uint8)


class UInt32Column(Column):
    def get_input(self, n):
        return 3 + np.arange(n, dtype=np.uint32)


class CharColumn(Column):
    def get_input(self, n):
        rng = np.random.RandomState(42)
        return rng.randint(low=65, high=122, size=n, dtype=np.int8)


class DoubleColumn(Column):
    def get_input(self, n):
        return 4 + np.arange(n, dtype=np.float64)


class CommonTestsMixin:
    """
    Abstract base class for common table tests. Because of the design of unittest,
    we have to make this a mixin.
    """

    def make_input_data(self, num_rows):
        rng = np.random.RandomState(42)
        input_data = {col.name: col.get_input(num_rows) for col in self.columns}
        for list_col, offset_col in self.ragged_list_columns:
            lengths = rng.randint(low=0, high=10, size=num_rows)
            input_data[list_col.name] = list_col.get_input(sum(lengths))
            input_data[offset_col.name] = np.zeros(num_rows + 1, dtype=np.uint64)
            input_data[offset_col.name][1:] = np.cumsum(lengths, dtype=np.uint64)
        return input_data

    def make_transposed_input_data(self, num_rows):
        cols = self.make_input_data(num_rows)
        return [
            {
                col: (
                    data[j]
                    if len(data) == num_rows
                    else (
                        bytes(
                            data[
                                cols[f"{col}_offset"][j] : cols[f"{col}_offset"][j + 1]
                            ]
                        )
                        if "metadata" in col
                        else data[
                            cols[f"{col}_offset"][j] : cols[f"{col}_offset"][j + 1]
                        ]
                    )
                )
                for col, data in cols.items()
                if "offset" not in col
            }
            for j in range(num_rows)
        ]

    @pytest.fixture
    def test_rows(self, scope="session"):
        test_rows = self.make_transposed_input_data(10)
        # Annoyingly we have to tweak some types as once added to a row and then put in
        # an error message things come out differently
        for n in range(10):
            for col in test_rows[n].keys():
                if col in ["timestamp", "record", "ancestral_state", "derived_state"]:
                    test_rows[n][col] = bytes(test_rows[n][col]).decode("ascii")
        return test_rows

    @pytest.fixture
    def table(self, test_rows):
        table = self.table_class()
        for row in test_rows:
            table.add_row(**row)
        return table

    @pytest.fixture
    def table_5row(self, test_rows):
        table_5row = self.table_class()
        for row in test_rows[:5]:
            table_5row.add_row(**row)
        return table_5row

    def test_asdict(self, table, test_rows):
        for table_row, test_row in zip(table, test_rows):
            for k, v in table_row.asdict().items():
                if isinstance(v, np.ndarray):
                    assert np.array_equal(v, test_row[k])
                else:
                    assert v == test_row[k]

    def test_max_rows_increment(self):
        for bad_value in [-1, -(2**10)]:
            with pytest.raises(ValueError):
                self.table_class(max_rows_increment=bad_value)
        for v in [1, 100, 256]:
            table = self.table_class(max_rows_increment=v)
            assert table.max_rows_increment == v
        # Setting zero implies doubling
        table = self.table_class()
        assert table.max_rows_increment == 0
        table = self.table_class(max_rows_increment=1024)
        assert table.max_rows_increment == 1024
        table = self.table_class(max_rows_increment=0)
        assert table.max_rows_increment == 0

    def test_low_level_get_row(self):
        # Tests the low-level get_row interface to ensure we're getting coverage.
        t = self.table_class()
        with pytest.raises(TypeError):
            t.ll_table.get_row()
        with pytest.raises(TypeError):
            t.ll_table.get_row("row")
        with pytest.raises(_tskit.LibraryError):
            t.ll_table.get_row(1)

    def test_low_level_equals(self):
        # Tests the low-level equals interface to ensure we're getting coverage.
        t = self.table_class()
        with pytest.raises(TypeError):
            t.ll_table.equals()
        with pytest.raises(TypeError):
            t.ll_table.equals(None)

    def test_low_level_set_columns(self):
        t = self.table_class()
        with pytest.raises(TypeError):
            t.ll_table.set_columns(None)
        with pytest.raises(TypeError):
            t.ll_table.append_columns(None)

    def test_input_parameters_errors(self):
        assert len(self.input_parameters) > 0
        for param, _ in self.input_parameters:
            for bad_value in [-1, -(2**10)]:
                with pytest.raises(ValueError):
                    self.table_class(**{param: bad_value})
            for bad_type in [None, ValueError, "ser"]:
                with pytest.raises(TypeError):
                    self.table_class(**{param: bad_type})

    def test_input_parameter_values(self):
        assert len(self.input_parameters) > 0
        for param, _ in self.input_parameters:
            for v in [1, 100, 256]:
                table = self.table_class(**{param: v})
                assert getattr(table, param) == v

    def test_set_columns_string_errors(self):
        inputs = {c.name: c.get_input(1) for c in self.columns}
        for list_col, offset_col in self.ragged_list_columns:
            value = list_col.get_input(1)
            inputs[list_col.name] = value
            inputs[offset_col.name] = [0, 1]
        # Make sure this works.
        table = self.table_class()
        table.set_columns(**inputs)
        for list_col, offset_col in self.ragged_list_columns:
            kwargs = dict(inputs)
            del kwargs[list_col.name]
            with pytest.raises(TypeError):
                table.set_columns(**kwargs)
            kwargs = dict(inputs)
            del kwargs[offset_col.name]
            with pytest.raises(TypeError):
                table.set_columns(**kwargs)

    def test_set_columns_interface(self):
        kwargs = self.make_input_data(1)
        # Make sure this works.
        table = self.table_class()
        table.set_columns(**kwargs)
        table.append_columns(**kwargs)
        for focal_col in self.columns:
            table = self.table_class()
            for bad_type in [Exception, tskit]:
                error_kwargs = dict(kwargs)
                error_kwargs[focal_col.name] = bad_type
                with pytest.raises(ValueError):
                    table.set_columns(**error_kwargs)
                with pytest.raises(ValueError):
                    table.append_columns(**error_kwargs)
            for bad_value in ["qwer", [0, "sd"]]:
                error_kwargs = dict(kwargs)
                error_kwargs[focal_col.name] = bad_value
                with pytest.raises(ValueError):
                    table.set_columns(**error_kwargs)
                with pytest.raises(ValueError):
                    table.append_columns(**error_kwargs)

    def test_set_columns_from_dict(self):
        kwargs = self.make_input_data(1)
        # Make sure this works.
        t1 = self.table_class()
        t1.set_columns(**kwargs)
        t2 = self.table_class()
        t2.set_columns(**t1.asdict())
        t1.assert_equals(t2)

    def test_set_columns_dimension(self):
        kwargs = self.make_input_data(1)
        table = self.table_class()
        table.set_columns(**kwargs)
        table.append_columns(**kwargs)
        for focal_col in self.columns:
            table = self.table_class()
            for bad_dims in [5, [[1], [1]], np.zeros((2, 2))]:
                error_kwargs = dict(kwargs)
                error_kwargs[focal_col.name] = bad_dims
                with pytest.raises(ValueError):
                    table.set_columns(**error_kwargs)
                with pytest.raises(ValueError):
                    table.append_columns(**error_kwargs)
        for _, offset_col in self.ragged_list_columns:
            error_kwargs = dict(kwargs)
            for bad_dims in [5, [[1], [1]], np.zeros((2, 2))]:
                error_kwargs[offset_col.name] = bad_dims
                with pytest.raises(ValueError):
                    table.set_columns(**error_kwargs)
                with pytest.raises(ValueError):
                    table.append_columns(**error_kwargs)
            # Empty offset columns are caught also
            error_kwargs[offset_col.name] = []
            with pytest.raises(ValueError):
                table.set_columns(**error_kwargs)

    def test_set_columns_input_sizes(self):
        input_data = self.make_input_data(100)
        col_map = {col.name: col for col in self.columns}
        for list_col, offset_col in self.ragged_list_columns:
            col_map[list_col.name] = list_col
            col_map[offset_col.name] = offset_col
        table = self.table_class()
        table.set_columns(**input_data)
        table.append_columns(**input_data)
        for equal_len_col_set in self.equal_len_columns:
            if len(equal_len_col_set) > 1:
                for col in equal_len_col_set:
                    kwargs = dict(input_data)
                    kwargs[col] = col_map[col].get_input(1)
                    with pytest.raises(ValueError):
                        table.set_columns(**kwargs)
                    with pytest.raises(ValueError):
                        table.append_columns(**kwargs)

    def test_set_read_only_attributes(self):
        table = self.table_class()
        with pytest.raises(AttributeError):
            table.num_rows = 10
        with pytest.raises(AttributeError):
            table.max_rows = 10
        for param, _default in self.input_parameters:
            with pytest.raises(AttributeError):
                setattr(table, param, 2)
        assert table.num_rows == 0
        assert len(table) == 0

    def test_set_column_attributes_empty(self):
        table = self.table_class()
        input_data = {col.name: col.get_input(0) for col in self.columns}
        for col, data in input_data.items():
            setattr(table, col, data)
            assert len(getattr(table, col)) == 0

    def test_set_column_attributes_data(self):
        table = self.table_class()
        for num_rows in [1, 10, 100]:
            input_data = self.make_input_data(num_rows)
            table.set_columns(**input_data)

            for list_col, offset_col in self.ragged_list_columns:
                list_data = input_data[list_col.name]
                assert np.array_equal(getattr(table, list_col.name), list_data)
                list_data += 1
                assert not np.array_equal(getattr(table, list_col.name), list_data)
                setattr(table, list_col.name, list_data)
                assert np.array_equal(getattr(table, list_col.name), list_data)
                list_value = getattr(table[0], list_col.name)
                assert len(list_value) == input_data[offset_col.name][1]

                # Reset the offsets so that all the full array is associated with the
                # first element.
                offset_data = np.zeros(num_rows + 1, dtype=np.uint32) + len(
                    input_data[list_col.name]
                )
                offset_data[0] = 0
                setattr(table, offset_col.name, offset_data)
                list_value = getattr(table[0], list_col.name)
                assert len(list_value) == len(input_data[list_col.name])

                del input_data[list_col.name]
                del input_data[offset_col.name]

            for col, data in input_data.items():
                assert np.array_equal(getattr(table, col), data)
                data += 1
                assert not np.array_equal(getattr(table, col), data)
                setattr(table, col, data)
                assert np.array_equal(getattr(table, col), data)

    def test_set_column_attributes_errors(self):
        table = self.table_class()
        num_rows = 10
        input_data = self.make_input_data(num_rows)
        table.set_columns(**input_data)

        for list_col, offset_col in self.ragged_list_columns:
            for bad_list_col in [[], input_data[list_col.name][:-1]]:
                with pytest.raises(ValueError):
                    setattr(table, list_col.name, bad_list_col)
            for bad_offset_col in [[], np.arange(num_rows + 2, dtype=np.uint32)]:
                with pytest.raises(ValueError):
                    setattr(table, offset_col.name, bad_offset_col)

            del input_data[list_col.name]
            del input_data[offset_col.name]

        for col, data in input_data.items():
            for bad_data in [[], data[:-1]]:
                with pytest.raises(ValueError):
                    setattr(table, col, bad_data)

        # Try to read a column that isn't there. (We can always write to new attributes
        # in Python, so there's nothing to test in that case.)
        with pytest.raises(AttributeError):
            _ = table.no_such_column

    def test_defaults(self):
        table = self.table_class()
        assert table.num_rows == 0
        assert len(table) == 0
        for param, default in self.input_parameters:
            assert getattr(table, param) == default
        for col in self.columns:
            array = getattr(table, col.name)
            assert array.shape == (0,)

    def test_add_row_data(self):
        for num_rows in [0, 10, 100]:
            table = self.table_class()
            for j, row in enumerate(self.make_transposed_input_data(num_rows)):
                k = table.add_row(**row)
                assert k == j
            for colname, input_array in self.make_input_data(num_rows).items():
                output_array = getattr(table, colname)
                assert input_array.shape == output_array.shape
                assert np.all(input_array == output_array)
            table.clear()
            assert table.num_rows == 0
            assert len(table) == 0

    def test_add_row_round_trip(self):
        for num_rows in [0, 10, 100]:
            input_data = self.make_input_data(num_rows)
            t1 = self.table_class()
            t1.set_columns(**input_data)
            for colname, input_array in input_data.items():
                output_array = getattr(t1, colname)
                assert input_array.shape == output_array.shape
                assert np.all(input_array == output_array)
            t2 = self.table_class()
            for row in list(t1):
                t2.add_row(**dataclasses.asdict(row))
            t1.assert_equals(t2)

    def test_append_row(self):
        for num_rows in [0, 10, 100]:
            table = self.table_class()
            for j, row in enumerate(self.make_transposed_input_data(num_rows)):
                k = table.append(table.row_class(**row))
                assert k == j
            for colname, input_array in self.make_input_data(num_rows).items():
                output_array = getattr(table, colname)
                assert input_array.shape == output_array.shape
                assert np.all(input_array == output_array)
            table.clear()
            assert table.num_rows == 0
            assert len(table) == 0

    def test_append_duck_type(self):
        class Duck:
            pass

        table = self.table_class()
        for j, row in enumerate(self.make_transposed_input_data(20)):
            duck = Duck()
            for k, v in row.items():
                setattr(duck, k, v)
            k = table.append(duck)
            assert k == j
        for colname, input_array in self.make_input_data(20).items():
            output_array = getattr(table, colname)
            assert np.array_equal(input_array, output_array)

    def test_append_error(self):
        class NotADuck:
            pass

        with pytest.raises(AttributeError, match="'NotADuck' object has no attribute"):
            self.table_class().append(NotADuck())

    def test_setitem(self):
        table = self.table_class()
        for row in self.make_transposed_input_data(10):
            table.append(table.row_class(**row))
        table2 = self.table_class()
        for row in self.make_transposed_input_data(20)[10:]:
            table2.append(table.row_class(**row))
        assert table != table2

        copy = table.copy()
        for j in range(10):
            table[j] = table[j]
        table.assert_equals(copy)

        for j in range(10):
            table[j] = table2[j]
        table.assert_equals(table2)

    def test_setitem_duck_type(self):
        class Duck:
            pass

        table = self.table_class()
        for row in self.make_transposed_input_data(10):
            table.append(table.row_class(**row))
        table2 = self.table_class()
        for row in self.make_transposed_input_data(20)[10:]:
            table2.append(table.row_class(**row))
        assert table != table2

        for j in range(10):
            duck = Duck()
            for k, v in dataclasses.asdict(table2[j]).items():
                setattr(duck, k, v)
            table[j] = duck
        table.assert_equals(table2)

    def test_setitem_error(self):
        class NotADuck:
            pass

        table = self.table_class()
        table.append(table.row_class(**self.make_transposed_input_data(1)[0]))
        with pytest.raises(AttributeError, match="'NotADuck' object has no attribute"):
            table[0] = NotADuck()

        with pytest.raises(IndexError, match="Index out of bounds"):
            self.table_class()[0] = table[0]
        with pytest.raises(IndexError, match="Index out of bounds"):
            self.table_class()[-1] = table[0]

        with pytest.raises(TypeError, match="Index must be integer"):
            self.table_class()[0.5] = table[0]
        with pytest.raises(TypeError, match="Index must be integer"):
            self.table_class()[None] = table[0]
        with pytest.raises(TypeError, match="Index must be integer"):
            self.table_class()[[1]] = table[0]

    def test_set_columns_data(self):
        for num_rows in [0, 10, 100, 1000]:
            input_data = {col.name: col.get_input(num_rows) for col in self.columns}
            offset_cols = set()
            for list_col, offset_col in self.ragged_list_columns:
                value = list_col.get_input(num_rows)
                input_data[list_col.name] = value
                input_data[offset_col.name] = np.arange(num_rows + 1, dtype=np.uint32)
                offset_cols.add(offset_col.name)
            table = self.table_class()
            for _ in range(5):
                table.set_columns(**input_data)
                for colname, input_array in input_data.items():
                    output_array = getattr(table, colname)
                    assert input_array.shape == output_array.shape
                    assert np.all(input_array == output_array)
                table.clear()
                assert table.num_rows == 0
                assert len(table) == 0
                for colname in input_data.keys():
                    if colname in offset_cols:
                        assert list(getattr(table, colname)) == [0]
                    else:
                        assert list(getattr(table, colname)) == []

    def test_truncate(self):
        num_rows = 100
        input_data = {col.name: col.get_input(num_rows) for col in self.columns}
        for list_col, offset_col in self.ragged_list_columns:
            value = list_col.get_input(2 * num_rows)
            input_data[list_col.name] = value
            input_data[offset_col.name] = 2 * np.arange(num_rows + 1, dtype=np.uint32)
        table = self.table_class()
        table.set_columns(**input_data)

        copy = table.copy()
        table.truncate(num_rows)
        assert copy == table

        for num_rows in [100, 10, 1]:
            table.truncate(num_rows)
            assert table.num_rows == num_rows
            assert len(table) == num_rows
            used = set()
            for list_col, offset_col in self.ragged_list_columns:
                offset = getattr(table, offset_col.name)
                assert offset.shape == (num_rows + 1,)
                assert np.array_equal(
                    input_data[offset_col.name][: num_rows + 1], offset
                )
                list_data = getattr(table, list_col.name)
                assert np.array_equal(
                    list_data, input_data[list_col.name][: offset[-1]]
                )
                used.add(offset_col.name)
                used.add(list_col.name)
            for name, data in input_data.items():
                if name not in used:
                    assert np.array_equal(data[:num_rows], getattr(table, name))

    def test_truncate_errors(self):
        num_rows = 10
        input_data = {col.name: col.get_input(num_rows) for col in self.columns}
        for list_col, offset_col in self.ragged_list_columns:
            value = list_col.get_input(2 * num_rows)
            input_data[list_col.name] = value
            input_data[offset_col.name] = 2 * np.arange(num_rows + 1, dtype=np.uint32)
        table = self.table_class()
        table.set_columns(**input_data)
        for bad_type in [None, 0.001, {}]:
            with pytest.raises(TypeError):
                table.truncate(bad_type)
        for bad_num_rows in [-1, num_rows + 1, 10**6]:
            with pytest.raises(ValueError):
                table.truncate(bad_num_rows)

    def test_append_columns_data(self):
        for num_rows in [0, 10, 100, 1000]:
            input_data = self.make_input_data(num_rows)
            offset_cols = set()
            for _, offset_col in self.ragged_list_columns:
                offset_cols.add(offset_col.name)
            table = self.table_class()
            for j in range(1, 10):
                table.append_columns(**input_data)
                for colname, values in input_data.items():
                    output_array = getattr(table, colname)
                    if colname in offset_cols:
                        input_array = np.zeros(j * num_rows + 1, dtype=np.uint32)
                        for k in range(j):
                            input_array[k * num_rows : (k + 1) * num_rows + 1] = (
                                k * values[-1]
                            ) + values
                        assert input_array.shape == output_array.shape
                    else:
                        input_array = np.hstack([values for _ in range(j)])
                        assert input_array.shape == output_array.shape
                    assert np.array_equal(input_array, output_array)
                assert table.num_rows == j * num_rows
                assert len(table) == j * num_rows

    def test_append_columns_max_rows(self):
        for num_rows in [0, 10, 100, 1000]:
            input_data = self.make_input_data(num_rows)
            for max_rows in [1, 8192]:
                table = self.table_class(max_rows_increment=max_rows)
                for j in range(1, 10):
                    table.append_columns(**input_data)
                    assert table.num_rows == j * num_rows
                    assert len(table) == j * num_rows
                    if table.num_rows == 0:
                        assert table.max_rows == 1
                    elif table.num_rows > max_rows + 1:
                        assert table.max_rows == max((max_rows * 2) + 1, table.num_rows)
                    else:
                        assert table.max_rows == max(max_rows + 1, table.num_rows)

    def test_keep_rows_data(self):
        input_data = self.make_input_data(100)
        t1 = self.table_class()
        t1.append_columns(**input_data)
        t2 = t1.copy()
        keep = np.ones(len(t1), dtype=bool)
        # Only keep even
        keep[::2] = 0
        t1.keep_rows(keep)
        keep_rows_definition(t2, keep)
        assert t1.equals(t2)

    def test_str(self):
        for num_rows in [0, 10]:
            input_data = self.make_input_data(num_rows)
            table = self.table_class()
            table.set_columns(**input_data)
            s = str(table)
            assert len(s.splitlines()) == num_rows + 4
        input_data = self.make_input_data(41)
        table = self.table_class()
        table.set_columns(**input_data)
        blank_meta_row = 39
        if "metadata" in input_data:
            table[blank_meta_row] = table[blank_meta_row].replace(metadata=b"")
        assert "1 rows skipped" in str(table)
        tskit.set_print_options(max_lines=None)
        assert "1 rows skipped" not in str(table)
        assert "b''" not in str(table)
        tskit.set_print_options(max_lines=40)
        tskit.MAX_LINES = 40

    def test_str_pos_time_integer(self):
        num_rows = 2
        identifiable_integers = [12345, 54321]
        identifiable_floats = [1.2345, 5.4321]
        table = self.table_class()
        for test_cols in [
            ["left", "right"],
            ["position"],
            ["time"],
        ]:  # only cols that get discretised
            input_data = self.make_input_data(num_rows)
            if all(col in input_data for col in test_cols):
                for i, col in enumerate(test_cols):
                    input_data[col] = [identifiable_floats[i]] * num_rows
                table.set_columns(**input_data)
                _, rows = table._text_header_and_rows()
                for row in rows:
                    assert f"{identifiable_floats[0]:.8f}" in row
                    assert f"{identifiable_integers[0]}" not in row
                for i, col in enumerate(test_cols):
                    input_data[col] = [identifiable_integers[i]] * num_rows
                table.set_columns(**input_data)
                _, rows = table._text_header_and_rows()
                for row in rows:
                    assert f"{identifiable_integers[0]:,}" in row
                    assert f"{identifiable_floats[0]:.8f}" not in row

    def test_repr_html(self):
        for num_rows in [0, 10, 40, 50]:
            input_data = {col.name: col.get_input(num_rows) for col in self.columns}
            for list_col, offset_col in self.ragged_list_columns:
                value = list_col.get_input(num_rows)
                input_data[list_col.name] = value
                input_data[offset_col.name] = np.arange(num_rows + 1, dtype=np.uint32)
            table = self.table_class()
            table.set_columns(**input_data)
            html = table._repr_html_()
            if num_rows == 50:
                assert len(html.splitlines()) == num_rows + 11
                assert (
                    "<em>10 rows skipped (tskit.set_print_options)</em>"
                    in html.split("</tr>")[21]
                )
            else:
                assert len(html.splitlines()) == num_rows + 20

    def test_copy(self):
        for num_rows in [0, 10]:
            input_data = self.make_input_data(num_rows)
            table = self.table_class()
            table.set_columns(**input_data)
            for _ in range(10):
                copy = table.copy()
                assert id(copy) != id(table)
                assert isinstance(copy, self.table_class)
                copy.assert_equals(table)
                table = copy

    def test_pickle(self):
        for num_rows in [0, 10, 100]:
            input_data = self.make_input_data(num_rows)
            table = self.table_class()
            table.set_columns(**input_data)
            pkl = pickle.dumps(table)
            new_table = pickle.loads(pkl)
            table.assert_equals(new_table)
            for protocol in range(pickle.HIGHEST_PROTOCOL + 1):
                pkl = pickle.dumps(table, protocol=protocol)
                new_table = pickle.loads(pkl)
                table.assert_equals(new_table)

    def test_equality(self):
        for num_rows in [1, 10, 100]:
            input_data = self.make_input_data(num_rows)
            t1 = self.table_class()
            t2 = self.table_class()
            assert t1 == t1
            assert t1 == t2
            assert t1 == t2
            assert not (t1 != t2)
            t1.set_columns(**input_data)
            assert t1 == t1
            assert t1 != t2
            assert t2 != t1
            t2.set_columns(**input_data)
            assert t1 == t2
            assert t2 == t2
            t2.clear()
            assert t1 != t2
            assert t2 != t1
            # Check each column in turn to see if we are correctly checking values.
            for col in self.columns:
                col_copy = np.copy(input_data[col.name])
                input_data_copy = dict(input_data)
                input_data_copy[col.name] = col_copy
                t2.set_columns(**input_data_copy)
                assert t1 == t2
                assert not (t1 != t2)
                assert t1[0] == t2[0]
                col_copy += 1
                t2.set_columns(**input_data_copy)
                assert t1 != t2
                assert t2 != t1
                assert t1[0] != t2[0]
                assert t1[0] != t2[0]
                assert t1[0] != []
            for list_col, offset_col in self.ragged_list_columns:
                value = list_col.get_input(num_rows)
                input_data_copy = dict(input_data)
                input_data_copy[list_col.name] = value + 1
                input_data_copy[offset_col.name] = np.arange(
                    num_rows + 1, dtype=np.uint32
                )
                t2.set_columns(**input_data_copy)
                assert t1 != t2
                assert t1[0] != t2[0]
                value = list_col.get_input(num_rows + 1)
                input_data_copy = dict(input_data)
                input_data_copy[list_col.name] = value
                input_data_copy[offset_col.name] = np.arange(
                    num_rows + 1, dtype=np.uint32
                )
                input_data_copy[offset_col.name][-1] = num_rows + 1
                t2.set_columns(**input_data_copy)
                assert t1 != t2
                assert t2 != t1
                assert t1[-1] != t2[-1]
            # Different types should always be unequal.
            assert t1 is not None
            assert t1 != []

    def test_nbytes(self):
        for num_rows in [0, 10, 100]:
            input_data = self.make_input_data(num_rows)
            table = self.table_class()
            table.set_columns(**input_data)
            # We don't have any metadata_schema here, so we can sum over the
            # columns directly.
            assert sum(col.nbytes for col in input_data.values()) == table.nbytes

    def test_bad_offsets(self):
        for num_rows in [10, 100]:
            input_data = self.make_input_data(num_rows)
            t = self.table_class()
            t.set_columns(**input_data)

            for _list_col, offset_col in self.ragged_list_columns:
                original_offset = np.copy(input_data[offset_col.name])
                # As numpy no longer allows conversion of out-of-bounds values, we
                # explictly cast first.
                input_data[offset_col.name][0] = np.array(-1).astype(
                    input_data[offset_col.name].dtype
                )
                with pytest.raises(ValueError):
                    t.set_columns(**input_data)
                input_data[offset_col.name] = np.copy(original_offset)
                t.set_columns(**input_data)
                input_data[offset_col.name][-1] = 0
                with pytest.raises(ValueError):
                    t.set_columns(**input_data)
                input_data[offset_col.name] = np.copy(original_offset)
                t.set_columns(**input_data)
                input_data[offset_col.name][num_rows // 2] = 2**31
                with pytest.raises(ValueError):
                    t.set_columns(**input_data)
                input_data[offset_col.name] = np.copy(original_offset)

                input_data[offset_col.name][0] = np.array(-1).astype(
                    input_data[offset_col.name].dtype
                )
                with pytest.raises(ValueError):
                    t.append_columns(**input_data)
                input_data[offset_col.name] = np.copy(original_offset)
                t.append_columns(**input_data)
                input_data[offset_col.name][-1] = 0
                with pytest.raises(ValueError):
                    t.append_columns(**input_data)
                input_data[offset_col.name] = np.copy(original_offset)
                t.append_columns(**input_data)
                input_data[offset_col.name][num_rows // 2] = 2**31
                with pytest.raises(ValueError):
                    t.append_columns(**input_data)
                input_data[offset_col.name] = np.copy(original_offset)

    def test_replace_with_wrong_class(self):
        t = self.table_class()
        with pytest.raises(TypeError, match="is required"):
            t.replace_with(tskit.BaseTable(None, None))


class MetadataTestsMixin:
    """
    Tests for column that have metadata columns.
    """

    metadata_schema = metadata.MetadataSchema(
        {
            "codec": "json",
            "title": "Example Metadata",
            "type": "object",
            "properties": {
                "one": {"type": "string"},
                "two": {"type": "number"},
                "three": {"type": "array"},
                "four": {"type": "boolean"},
            },
            "required": ["one", "two", "three", "four"],
            "additionalProperties": False,
        },
    )

    def metadata_example_data(self):
        try:
            self.val += 1
        except AttributeError:
            self.val = 0
        return {
            "one": "val one",
            "two": self.val,
            "three": list(range(self.val, self.val + 10)),
            "four": True,
        }

    def input_data_for_add_row(self):
        input_data = {col.name: col.get_input(1) for col in self.columns}
        kwargs = {col: data[0] for col, data in input_data.items()}
        for col in self.string_colnames:
            kwargs[col] = "x"
        for col in self.binary_colnames:
            kwargs[col] = b"x"
        return kwargs

    def test_random_metadata(self):
        for num_rows in [0, 10, 100]:
            input_data = self.make_input_data(num_rows)
            table = self.table_class()
            metadatas = [tsutil.random_bytes(10) for _ in range(num_rows)]
            metadata, metadata_offset = tskit.pack_bytes(metadatas)
            input_data["metadata"] = metadata
            input_data["metadata_offset"] = metadata_offset
            table.set_columns(**input_data)
            unpacked_metadatas = tskit.unpack_bytes(
                table.metadata, table.metadata_offset
            )
            assert metadatas == unpacked_metadatas

    def test_drop_metadata(self):
        for num_rows in [1, 10, 100]:
            input_data = self.make_input_data(num_rows)
            table_no_meta = self.table_class()
            table_with_meta = self.table_class()
            table_with_meta.set_columns(**input_data)
            if not getattr(self, "metadata_mandatory", False):
                del input_data["metadata"]
                del input_data["metadata_offset"]
            else:
                # Have to do this slightly circular way for the population
                # table because it requires metadata.
                input_data["metadata"] = []
                input_data["metadata_offset"][:] = 0
            table_no_meta.set_columns(**input_data)
            assert not table_no_meta.equals(table_with_meta)
            table_with_meta.drop_metadata()
            table_no_meta.assert_equals(table_with_meta)

    def test_optional_metadata(self):
        if not getattr(self, "metadata_mandatory", False):
            for num_rows in [0, 10, 100]:
                input_data = self.make_input_data(num_rows)
                table = self.table_class()
                del input_data["metadata"]
                del input_data["metadata_offset"]
                table.set_columns(**input_data)
                assert len(list(table.metadata)) == 0
                assert list(table.metadata_offset) == [0 for _ in range(num_rows + 1)]
                # Supplying None is the same not providing the column.
                input_data["metadata"] = None
                input_data["metadata_offset"] = None
                table.set_columns(**input_data)
                assert len(list(table.metadata)) == 0
                assert list(table.metadata_offset) == [0 for _ in range(num_rows + 1)]

    def test_packset_metadata(self):
        for num_rows in [0, 10, 100]:
            input_data = self.make_input_data(num_rows)
            table = self.table_class()
            table.set_columns(**input_data)
            metadatas = [tsutil.random_bytes(10) for _ in range(num_rows)]
            metadata, metadata_offset = tskit.pack_bytes(metadatas)
            table.packset_metadata(metadatas)
            assert np.array_equal(table.metadata, metadata)
            assert np.array_equal(table.metadata_offset, metadata_offset)

    def test_set_metadata_schema(self):
        metadata_schema2 = metadata.MetadataSchema({"codec": "json"})
        table = self.table_class()
        # Default is no-op metadata codec
        assert repr(table.metadata_schema) == repr(metadata.MetadataSchema(None))
        # Set
        table.metadata_schema = self.metadata_schema
        assert repr(table.metadata_schema) == repr(self.metadata_schema)
        # Overwrite
        table.metadata_schema = metadata_schema2
        assert repr(table.metadata_schema) == repr(metadata_schema2)
        # Remove
        table.metadata_schema = metadata.MetadataSchema(None)
        assert repr(table.metadata_schema) == repr(metadata.MetadataSchema(None))
        # Set after remove
        table.metadata_schema = self.metadata_schema
        assert repr(table.metadata_schema) == repr(self.metadata_schema)
        # Del should fail
        with pytest.raises(AttributeError):
            del table.metadata_schema
        # None should fail
        with pytest.raises(
            TypeError,
            match="Only instances of tskit.MetadataSchema can be assigned to "
            "metadata_schema, not <class 'NoneType'>",
        ):
            table.metadata_schema = None
        # And dict
        with pytest.raises(
            TypeError,
            match="Only instances of tskit.MetadataSchema can be assigned to "
            "metadata_schema, not <class 'dict'>",
        ):
            table.metadata_schema = {}

    def test_drop_metadata_with_schema(self):
        table = self.table_class()
        table.metadata_schema = metadata.MetadataSchema.permissive_json()
        data = self.input_data_for_add_row()
        data["metadata"] = {"a": "dict"}
        table.add_row(**data)
        assert table[0].metadata == {"a": "dict"}
        table.drop_metadata()
        assert table.metadata_schema == metadata.MetadataSchema.null()
        assert table[0].metadata == b""

    def test_drop_metadata_keep_schema(self):
        table = self.table_class()
        table.metadata_schema = metadata.MetadataSchema.permissive_json()
        data = self.input_data_for_add_row()
        data["metadata"] = {"a": "dict"}
        table.add_row(**data)
        assert table[0].metadata == {"a": "dict"}
        table.drop_metadata(keep_schema=True)
        assert table.metadata_schema == metadata.MetadataSchema.permissive_json()
        assert table[0].metadata == {}

    def test_default_metadata_schema(self):
        # Default should allow bytes as in pre-exisiting code
        table = self.table_class()
        table.add_row(
            **{**self.input_data_for_add_row(), "metadata": b"acceptable bytes"}
        )
        # Adding non-bytes metadata should error
        with pytest.raises(TypeError):
            table.add_row(
                **{
                    **self.input_data_for_add_row(),
                    "metadata": self.metadata_example_data(),
                }
            )

    def test_default_metadata_add_row(self):
        row_data = self.input_data_for_add_row()
        del row_data["metadata"]

        table = self.table_class()
        table.add_row(**row_data)
        assert table[0].metadata == b""
        assert table[0].metadata == table.metadata_schema.empty_value

        table = self.table_class()
        table.metadata_schema = tskit.MetadataSchema({"codec": "json"})
        table.add_row(**row_data)
        assert table[0].metadata == {}
        assert table[0].metadata == table.metadata_schema.empty_value

    def test_row_round_trip_metadata_schema(self):
        data = self.metadata_example_data()
        table = self.table_class()
        table.metadata_schema = self.metadata_schema
        table.add_row(**{**self.input_data_for_add_row(), "metadata": data})
        assert table[0].metadata == data

    def test_bad_row_metadata_schema(self):
        metadata = self.metadata_example_data()
        metadata["I really shouldn't be here"] = 6
        table = self.table_class()
        table.metadata_schema = self.metadata_schema
        with pytest.raises(exceptions.MetadataValidationError):
            table.add_row(**{**self.input_data_for_add_row(), "metadata": metadata})
        assert len(table) == 0

    def test_absent_metadata_with_required_schema(self):
        table = self.table_class()
        table.metadata_schema = self.metadata_schema
        input_data = self.input_data_for_add_row()
        del input_data["metadata"]
        with pytest.raises(exceptions.MetadataValidationError):
            table.add_row(**{**input_data})

    def test_unsupported_type(self):
        table = self.table_class()
        table.metadata_schema = metadata.MetadataSchema(
            {
                "codec": "json",
                "type": "object",
                "properties": {"an_array": {"type": "array"}},
            }
        )
        input_data = self.input_data_for_add_row()
        # Numpy is not a JSONSchema array
        input_data["metadata"] = {"an_array": np.arange(10)}
        with pytest.raises(exceptions.MetadataValidationError):
            table.add_row(**{**input_data})

    def test_round_trip_set_columns(self):
        for num_rows in [0, 10, 100]:
            table = self.table_class()
            table.metadata_schema = self.metadata_schema
            input_data = self.make_input_data(num_rows)
            del input_data["metadata"]
            del input_data["metadata_offset"]
            metadata_column = [self.metadata_example_data() for _ in range(num_rows)]
            encoded_metadata_column = [
                table.metadata_schema.validate_and_encode_row(r)
                for r in metadata_column
            ]
            packed_metadata, metadata_offset = tskit.util.pack_bytes(
                encoded_metadata_column
            )
            table.set_columns(
                metadata=packed_metadata, metadata_offset=metadata_offset, **input_data
            )
            table.append_columns(
                metadata=packed_metadata, metadata_offset=metadata_offset, **input_data
            )
            for j in range(num_rows):
                assert table[j].metadata == metadata_column[j]
                assert table[j + num_rows].metadata == metadata_column[j]

    @pytest.mark.parametrize(
        "codec",
        ["struct", "json"],
    )
    def test_set_null_metadata(self, codec):
        table = self.table_class()
        table.metadata_schema = metadata.MetadataSchema(
            {
                "codec": f"{codec}",
                "title": "Example Metadata",
                "type": ["object", "null"],
                "properties": {
                    "a": {"type": "number", "binaryFormat": "i"},
                },
                "required": ["a"],
                "additionalProperties": False,
            },
        )
        examples = [{"a": 4}, None]
        for md in examples:
            table.add_row(
                **{
                    **self.input_data_for_add_row(),
                    "metadata": md,
                }
            )
        assert table.num_rows == len(examples)
        for md, row in zip(examples, table):
            assert md == row.metadata

    # only json allows leaving out of optional entries
    def test_set_empty_metadata_json(self):
        table = self.table_class()
        table.metadata_schema = metadata.MetadataSchema(
            {
                "codec": "json",
                "title": "Example Metadata",
                "type": ["object", "null"],
                "properties": {
                    "a": {"type": "number", "binaryFormat": "i"},
                },
                "required": [],
                "additionalProperties": False,
            },
        )
        examples = [{"a": 4}, {}]
        for md in examples:
            table.add_row(
                **{
                    **self.input_data_for_add_row(),
                    "metadata": md,
                }
            )
        assert table.num_rows == len(examples)
        for md, row in zip(examples, table):
            assert md == row.metadata

    @pytest.mark.parametrize(
        "codec",
        ["struct", "json"],
    )
    def test_set_with_optional_properties(self, codec):
        table = self.table_class()
        table.metadata_schema = metadata.MetadataSchema(
            {
                "codec": f"{codec}",
                "title": "Example Metadata",
                "type": ["object", "null"],
                "properties": {
                    "a": {"type": "number", "binaryFormat": "i", "default": 0},
                },
                "additionalProperties": False,
            },
        )
        metadata_list = [{"a": 4}, None, {"a": 5}, {}]
        for md in metadata_list:
            table.add_row(
                **{
                    **self.input_data_for_add_row(),
                    "metadata": md,
                }
            )
        assert table.num_rows == len(metadata_list)
        for md, row in zip(metadata_list, table):
            # If None is allowed by the schema it gets used even in the presence of
            # default values.
            if isinstance(md, dict):
                defaults = {"a": 0}
                defaults.update(md)
                assert defaults == row.metadata
            else:
                assert md == row.metadata

    def test_copy_metadata_schema(self):
        table = self.table_class()
        assert table.metadata_schema == tskit.MetadataSchema(None)
        table.metadata_schema = tskit.MetadataSchema({"codec": "json"})
        copy = table.copy()
        table.assert_equals(copy)
        # Check this independently to check the schema cache was invalidated
        assert table.metadata_schema == copy.metadata_schema

        copy.metadata_schema = tskit.MetadataSchema(None)
        assert table.metadata_schema != copy.metadata_schema

    def test_set_columns_metadata_schema(self):
        table = self.table_class()
        table2 = self.table_class()
        ms = tskit.MetadataSchema({"codec": "json"})
        table2.metadata_schema = ms
        table.set_columns(**table2.asdict())
        assert table.metadata_schema == ms

    def verify_metadata_vector(self, table, key, dtype, default_value=9999):
        # this is just a hack for testing; the actual method
        # does this more elegantly
        has_default = default_value != 9999
        if has_default:
            md_vec = table.metadata_vector(
                key, default_value=default_value, dtype=dtype
            )
        else:
            md_vec = table.metadata_vector(key, dtype=dtype)
        assert isinstance(md_vec, np.ndarray)
        if dtype is not None:
            assert md_vec.dtype == np.dtype(dtype)
        assert len(md_vec) == table.num_rows
        if not isinstance(key, list):
            key = [key]
        for x, row in zip(md_vec, table):
            md = row.metadata
            for k in key:
                if k in md or not has_default:
                    md = md[k]
                else:
                    md = default_value
                    break
            assert np.all(np.asarray(md, dtype=dtype) == x)

    def test_metadata_vector_errors(self):
        table = self.table_class()
        ms = tskit.MetadataSchema({"codec": "json"})
        table.metadata_schema = ms
        table.add_row(
            **{
                **self.input_data_for_add_row(),
                "metadata": None,
            }
        )
        with pytest.raises(KeyError):
            _ = table.metadata_vector("x")
        metadata_list = [
            {"a": 4, "u": [1, 2]},
            {},
        ]
        for md in metadata_list:
            table.add_row(
                **{
                    **self.input_data_for_add_row(),
                    "metadata": md,
                }
            )
        with pytest.raises(KeyError):
            _ = table.metadata_vector("x")

        table.clear()
        metadata_list = [
            {"a": {"c": 5}, "u": [1, 2]},
            {"a": {"b": 6}},
        ]
        for md in metadata_list:
            table.add_row(
                **{
                    **self.input_data_for_add_row(),
                    "metadata": md,
                }
            )
        with pytest.raises(KeyError):
            _ = table.metadata_vector(["a", "x"])

    def test_metadata_vector_nodefault(self):
        table = self.table_class()
        ms = tskit.MetadataSchema({"codec": "json"})
        table.metadata_schema = ms
        metadata_list = [
            {"abc": 4, "u": [1, 2]},
            {"abc": 10, "u": [3, 4]},
            {"abc": -3, "b": {"c": 1}, "u": [5, 6]},
            {"abc": 1},
        ]
        for md in metadata_list:
            table.add_row(
                **{
                    **self.input_data_for_add_row(),
                    "metadata": md,
                }
            )
        # first the totally obvious test
        md_vec = table.metadata_vector("abc")
        assert np.all(np.equal(md_vec, [d["abc"] for d in metadata_list]))
        # now automated ones
        for dtype in [None, "int", "float", "object"]:
            self.verify_metadata_vector(
                table, key="abc", dtype=dtype, default_value=9999
            )
            self.verify_metadata_vector(
                table, key=["abc"], dtype=dtype, default_value=9999
            )

    def test_metadata_vector(self):
        table = self.table_class()
        ms = tskit.MetadataSchema({"codec": "json"})
        table.metadata_schema = ms
        metadata_list = [
            {"abc": 4, "u": [1, 2]},
            {"abc": 10, "u": [3, 4]},
            {"abc": -3, "b": {"c": 1}, "u": [5, 6]},
            {"b": {"c": 3.2}, "u": [7, 8]},
            {"b": {"x": 8.2}},
            {},
            None,
        ]
        for md in metadata_list:
            table.add_row(
                **{
                    **self.input_data_for_add_row(),
                    "metadata": md,
                }
            )
        # first the totally obvious test
        md_vec = table.metadata_vector("abc", default_value=0)
        assert np.all(
            np.equal(
                md_vec,
                [
                    d["abc"] if (d is not None and "abc" in d) else 0
                    for d in metadata_list
                ],
            )
        )

        # now some automated ones
        for dtype in [None, "int", "float", "object"]:
            self.verify_metadata_vector(table, key="abc", dtype=dtype, default_value=-1)
            self.verify_metadata_vector(
                table, key=["abc"], dtype=dtype, default_value=-1
            )
            self.verify_metadata_vector(table, key=["x"], dtype=dtype, default_value=-1)
            self.verify_metadata_vector(
                table, key=["b", "c"], dtype=dtype, default_value=-1
            )
        self.verify_metadata_vector(table, key=["b"], dtype="object", default_value=-1)
        self.verify_metadata_vector(table, key=["u"], dtype="int", default_value=[0, 0])
        # and finally we should get rectangular arrays when it makes sense
        md_vec = table.metadata_vector("u", default_value=[0, 0])
        assert md_vec.shape == (table.num_rows, 2)


class AssertEqualsMixin:
    def test_equal(self, table_5row, test_rows):
        table2 = self.table_class()
        for row in test_rows[:5]:
            table2.add_row(**row)
        table_5row.assert_equals(table2)

    def test_type(self, table_5row):
        with pytest.raises(
            AssertionError,
            match=f"Types differ: self={type(table_5row)} other=<class 'int'>",
        ):
            table_5row.assert_equals(42)

    def test_metadata_schema(self, table_5row):
        if hasattr(table_5row, "metadata_schema"):
            assert table_5row.metadata_schema == tskit.MetadataSchema(None)
            table2 = table_5row.copy()
            table2.metadata_schema = tskit.MetadataSchema({"codec": "json"})
            with pytest.raises(
                AssertionError,
                match=f"{type(table_5row).__name__} metadata schemas differ: ",
            ):
                table_5row.assert_equals(table2)
            table_5row.assert_equals(table2, ignore_metadata=True)

    def test_row_changes(self, table_5row, test_rows):
        for column_name in test_rows[0].keys():
            table2 = self.table_class()
            for row in test_rows[:4]:
                table2.add_row(**row)
            modified_row = {
                **test_rows[4],
                **{column_name: test_rows[5][column_name]},
            }
            table2.add_row(**modified_row)
            with pytest.raises(
                AssertionError,
                match=re.escape(
                    f"{type(table_5row).__name__} row 4 differs:\n"
                    f"self.{column_name}={test_rows[4][column_name]} "
                    f"other.{column_name}={test_rows[5][column_name]}"
                ),
            ):
                table_5row.assert_equals(table2)
            if column_name == "metadata":
                table_5row.assert_equals(table2, ignore_metadata=True)
            if column_name == "timestamp":
                table_5row.assert_equals(table2, ignore_timestamps=True)

        # Two columns differ, as we don't know the order in the error message
        # test for both independently
        for column_name, column_name2 in zip(
            list(test_rows[0].keys())[:-1], list(test_rows[0].keys())[1:]
        ):
            table2 = self.table_class()
            for row in test_rows[:4]:
                table2.add_row(**row)
            modified_row = {
                **test_rows[4],
                **{
                    column_name: test_rows[5][column_name],
                    column_name2: test_rows[5][column_name2],
                },
            }
            table2.add_row(**modified_row)
            with pytest.raises(
                AssertionError,
                match=re.escape(
                    f"self.{column_name}={test_rows[4][column_name]} "
                    f"other.{column_name}={test_rows[5][column_name]}"
                ),
            ):
                table_5row.assert_equals(table2)
            with pytest.raises(
                AssertionError,
                match=re.escape(
                    f"self.{column_name2}={test_rows[4][column_name2]} "
                    f"other.{column_name2}={test_rows[5][column_name2]}"
                ),
            ):
                table_5row.assert_equals(table2)

    def test_num_rows(self, table_5row, test_rows):
        table2 = self.table_class()
        for row in test_rows[:4]:
            table2.add_row(**row)
        with pytest.raises(
            AssertionError,
            match=f"{type(table_5row).__name__} number of rows differ: self=5 other=4",
        ):
            table_5row.assert_equals(table2)

    def test_metadata(self, table_5row, test_rows):
        if "metadata" in test_rows[0].keys():
            table2 = self.table_class()
            for row in test_rows[:4]:
                table2.add_row(**row)
            modified_row = {
                **test_rows[4],
                **{"metadata": test_rows[5]["metadata"]},
            }
            table2.add_row(**modified_row)
            with pytest.raises(
                AssertionError,
                match=re.escape(
                    f"{type(table_5row).__name__} row 4 differs:\n"
                    f"self.metadata={test_rows[4]['metadata']} "
                    f"other.metadata={test_rows[5]['metadata']}"
                ),
            ):
                table_5row.assert_equals(table2)
            table_5row.assert_equals(table2, ignore_metadata=True)

    def test_timestamp(self, table_5row, test_rows):
        if "timestamp" in test_rows[0].keys():
            table2 = self.table_class()
            for row in test_rows[:4]:
                table2.add_row(**row)
            modified_row = {
                **test_rows[4],
                **{"timestamp": test_rows[5]["timestamp"]},
            }
            table2.add_row(**modified_row)
            with pytest.raises(
                AssertionError,
                match=re.escape(
                    f"{type(table_5row).__name__} row 4 differs:\n"
                    f"self.timestamp={test_rows[4]['timestamp']} "
                    f"other.timestamp={test_rows[5]['timestamp']}"
                ),
            ):
                table_5row.assert_equals(table2)
            table_5row.assert_equals(table2, ignore_timestamps=True)


class FancyIndexingMixin:
    @pytest.mark.parametrize(
        "slic",
        [
            slice(None, None),
            slice(None, 3),
            slice(2, None),
            slice(1, 4),
            slice(4, 1),
            slice(1, 4, 2),
            slice(4, 1, 2),
            slice(4, 1, -1),
            slice(1, 4, -1),
            slice(3, None, -1),
            slice(None, 3, -1),
            slice(None, None, -2),
        ],
    )
    def test_slice(self, table, test_rows, slic):
        assert table.num_rows >= 5
        expected = table.copy()
        expected.truncate(0)
        for row in test_rows[slic]:
            expected.add_row(**row)
        table[slic].assert_equals(expected)

    @pytest.mark.parametrize(
        "mask",
        [
            [False] * 5,
            [True] * 5,
            [True] + [False] * 4,
            [False, True, False, True, True],
        ],
    )
    def test_boolean_array(self, table_5row, test_rows, mask):
        assert table_5row.num_rows >= 5
        expected = table_5row.copy()
        expected.truncate(0)
        for flag, row in zip(mask, test_rows[:5]):
            if flag:
                expected.add_row(**row)
        table_5row[mask].assert_equals(expected)

    @pytest.mark.parametrize(
        "index_array",
        [
            [],
            [0],
            [4],
            random.choices(range(5), k=100),
            np.array([0, 0, 0, 2], dtype=np.uint64),
            np.array([2, 4, 4, 0], dtype=np.int64),
            np.array([0, 0, 0, 2], dtype=np.uint32),
            np.array([2, 4, 4, 0], dtype=np.int32),
            np.array([4, 3, 4, 1], dtype=np.uint8),
            np.array([4, 3, 4, 1], dtype=np.int8),
        ],
    )
    def test_index_array(self, table_5row, index_array):
        assert table_5row.num_rows >= 5
        expected = table_5row.copy()
        expected.truncate(0)
        for index in index_array:
            expected.append(table_5row[index])
        table_5row[index_array].assert_equals(expected)
        table_5row[tuple(index_array)].assert_equals(expected)

    def test_index_range(self, table_5row):
        expected = table_5row.copy()
        expected.truncate(0)
        for index in range(2, 4):
            expected.append(table_5row[index])
        table_5row[range(2, 4)].assert_equals(expected)

    @pytest.mark.parametrize(
        "dtype",
        [
            np.float32,
            np.float64,
            object,
            str,
        ],
    )
    def test_bad_dtypes(self, table, dtype):
        with pytest.raises(TypeError):
            table[np.zeros((10,), dtype=np.float32)]

    @pytest.mark.parametrize(
        "dtype",
        [
            np.uint32,
            np.int64,
            np.uint64,
        ],
    )
    def test_bad_casts(self, table, dtype):
        with pytest.raises(OverflowError, match="Cannot convert safely"):
            table[np.asarray([np.iinfo(np.int32).max + 1], dtype=dtype)]

    def test_extrema(self, table):
        max_ = np.iinfo(np.int32).max
        with pytest.raises(OverflowError, match="Cannot convert safely"):
            table[[max_ + 1]]

        # Slice gets clipped to valid range
        copy = table.copy()
        copy.clear()
        table[max_ + 1 : max_ + 2].assert_equals(copy)

        with pytest.raises(OverflowError, match="Cannot convert safely"):
            table[range(max_ + 1, max_ + 2)]

    @pytest.mark.parametrize(
        "bad_shape",
        [
            [[0]],
            [[1, 2], [3, 4]],
        ],
    )
    def test_bad_shapes(self, table, bad_shape):
        with pytest.raises(ValueError, match="object too deep"):
            table[bad_shape]

    def test_bad_bool_length(self, table):
        with pytest.raises(
            IndexError, match="Boolean index must be same length as table"
        ):
            table[[False] * (len(table) + 1)]
        with pytest.raises(
            IndexError, match="Boolean index must be same length as table"
        ):
            table[[False]]

    def test_bad_indexes(self, table):
        with pytest.raises(_tskit.LibraryError, match="out of bounds"):
            table[[-1]]
        with pytest.raises(_tskit.LibraryError, match="out of bounds"):
            table[range(-5, 0)]
        with pytest.raises(_tskit.LibraryError, match="out of bounds"):
            table[[len(table)]]
        with pytest.raises(TypeError, match="Cannot cast"):
            table[[5.5]]
        with pytest.raises(TypeError, match="Cannot convert"):
            table[[None]]
        with pytest.raises(TypeError, match="not supported|did not contain"):
            table[["foobar"]]
        with pytest.raises(TypeError, match="Index must be integer, slice or iterable"):
            table[5.5]
        with pytest.raises(TypeError, match="Cannot convert to a rectangular array"):
            table[None]
        with pytest.raises(TypeError, match="not supported|did not contain"):
            table["foobar"]


common_tests = [
    CommonTestsMixin,
    MetadataTestsMixin,
    AssertEqualsMixin,
    FancyIndexingMixin,
]


class TestIndividualTable(*common_tests):
    columns = [UInt32Column("flags")]
    ragged_list_columns = [
        (DoubleColumn("location"), UInt32Column("location_offset")),
        (Int32Column("parents"), UInt32Column("parents_offset")),
        (CharColumn("metadata"), UInt32Column("metadata_offset")),
    ]
    string_colnames = []
    binary_colnames = ["metadata"]
    input_parameters = [("max_rows_increment", 0)]
    equal_len_columns = [["flags"]]
    table_class = tskit.IndividualTable

    def test_simple_example(self):
        t = tskit.IndividualTable()
        t.add_row(flags=0, location=[], parents=[], metadata=b"123")
        t.add_row(
            flags=1, location=(0, 1, 2, 3), parents=(4, 5, 6, 7), metadata=b"\xf0"
        )
        s = str(t)
        assert len(s) > 0
        assert len(t) == 2
        assert t[0].flags == 0
        assert list(t[0].location) == []
        assert list(t[0].parents) == []
        assert t[0].metadata == b"123"
        assert t[1].flags == 1
        assert list(t[1].location) == [0, 1, 2, 3]
        assert list(t[1].parents) == [4, 5, 6, 7]
        assert t[1].metadata == b"\xf0"
        with pytest.raises(IndexError):
            t.__getitem__(-4)

    def test_add_row_defaults(self):
        t = tskit.IndividualTable()
        assert t.add_row() == 0
        assert t.flags[0] == 0
        assert len(t.location) == 0
        assert t.location_offset[0] == 0
        assert len(t.parents) == 0
        assert t.parents_offset[0] == 0
        assert len(t.metadata) == 0
        assert t.metadata_offset[0] == 0

    def test_add_row_bad_data(self):
        t = tskit.IndividualTable()
        with pytest.raises(TypeError):
            t.add_row(flags="x")
        with pytest.raises(TypeError):
            t.add_row(metadata=123)
        with pytest.raises(ValueError):
            t.add_row(location="1234")
        with pytest.raises(ValueError):
            t.add_row(parents="forty-two")

    def test_packset_location(self):
        t = tskit.IndividualTable()
        t.add_row(flags=0)
        t.packset_location([[0.125, 2]])
        assert list(t[0].location) == [0.125, 2]
        t.add_row(flags=1)
        assert list(t[1].location) == []
        t.packset_location([[0], [1, 2, 3]])
        assert list(t[0].location) == [0]
        assert list(t[1].location) == [1, 2, 3]

    def test_packset_parents(self):
        t = tskit.IndividualTable()
        t.add_row(flags=0)
        t.packset_parents([[0, 2]])
        assert list(t[0].parents) == [0, 2]
        t.add_row(flags=1)
        assert list(t[1].parents) == []
        t.packset_parents([[0], [1, 2, 3]])
        assert list(t[0].parents) == [0]
        assert list(t[1].parents) == [1, 2, 3]

    def test_missing_time_equal_to_self(self):
        t = tskit.TableCollection(sequence_length=10)
        t.sites.add_row(position=1, ancestral_state="0")
        t.mutations.add_row(site=0, node=0, derived_state="1", time=tskit.UNKNOWN_TIME)
        assert t.mutations[0] == t.mutations[0]

    def test_various_not_equals(self):
        args = {
            "site": 0,
            "node": 0,
            "derived_state": "a",
            "parent": 0,
            "metadata": b"abc",
            "time": 0,
        }
        a = tskit.MutationTableRow(**args)
        assert a != []
        assert a != 12
        assert a is not None
        b = tskit.MutationTableRow(**args)
        assert a == b
        args["site"] = 2
        b = tskit.MutationTableRow(**args)
        assert a != b
        args["site"] = 0
        args["node"] = 2
        b = tskit.MutationTableRow(**args)
        assert a != b
        args["node"] = 0
        args["derived_state"] = "b"
        b = tskit.MutationTableRow(**args)
        assert a != b
        args["derived_state"] = "a"
        args["parent"] = 2
        b = tskit.MutationTableRow(**args)
        assert a != b
        args["parent"] = 0
        args["metadata"] = b""
        b = tskit.MutationTableRow(**args)
        assert a != b
        args["metadata"] = b"abc"
        args["time"] = 1
        b = tskit.MutationTableRow(**args)
        assert a != b
        args["time"] = 0
        args["time"] = tskit.UNKNOWN_TIME
        b = tskit.MutationTableRow(**args)
        assert a != b
        a = tskit.MutationTableRow(**args)
        assert a == b

    def test_keep_rows_data(self):
        input_data = self.make_input_data(100)
        t1 = self.table_class()
        # Set the parent column to -1s for this simple test as
        # we need to reason about reference integrity
        t1.append_columns(**input_data)
        t1.parents = np.full_like(t1.parents, -1)
        t2 = t1.copy()
        keep = np.ones(len(t1), dtype=bool)
        # Only keep even
        keep[::2] = 0
        t1.keep_rows(keep)
        keep_rows_definition(t2, keep)
        assert t1.equals(t2)


class TestNodeTable(*common_tests):
    columns = [
        UInt32Column("flags"),
        DoubleColumn("time"),
        Int32Column("individual"),
        Int32Column("population"),
    ]
    ragged_list_columns = [(CharColumn("metadata"), UInt32Column("metadata_offset"))]
    string_colnames = []
    binary_colnames = ["metadata"]
    input_parameters = [("max_rows_increment", 0)]
    equal_len_columns = [["time", "flags", "population"]]
    table_class = tskit.NodeTable

    def test_simple_example(self):
        t = tskit.NodeTable()
        t.add_row(flags=0, time=1, population=2, individual=0, metadata=b"123")
        t.add_row(flags=1, time=2, population=3, individual=1, metadata=b"\xf0")
        s = str(t)
        assert len(s) > 0
        assert len(t) == 2
        assert dataclasses.astuple(t[0]) == (0, 1, 2, 0, b"123")
        assert dataclasses.astuple(t[1]) == (1, 2, 3, 1, b"\xf0")
        assert t[0].flags == 0
        assert t[0].time == 1
        assert t[0].population == 2
        assert t[0].individual == 0
        assert t[0].metadata == b"123"
        assert t[0] == t[-2]
        assert t[1] == t[-1]
        with pytest.raises(IndexError):
            t.__getitem__(-3)

    def test_add_row_defaults(self):
        t = tskit.NodeTable()
        assert t.add_row() == 0
        assert t.time[0] == 0
        assert t.flags[0] == 0
        assert t.population[0] == tskit.NULL
        assert t.individual[0] == tskit.NULL
        assert len(t.metadata) == 0
        assert t.metadata_offset[0] == 0

    def test_optional_population(self):
        for num_rows in [0, 10, 100]:
            metadatas = [str(j) for j in range(num_rows)]
            metadata, metadata_offset = tskit.pack_strings(metadatas)
            flags = list(range(num_rows))
            time = list(range(num_rows))
            table = tskit.NodeTable()
            table.set_columns(
                metadata=metadata,
                metadata_offset=metadata_offset,
                flags=flags,
                time=time,
            )
            assert list(table.population) == [-1 for _ in range(num_rows)]
            assert list(table.flags) == flags
            assert list(table.time) == time
            assert list(table.metadata) == list(metadata)
            assert list(table.metadata_offset) == list(metadata_offset)
            table.set_columns(flags=flags, time=time, population=None)
            assert list(table.population) == [-1 for _ in range(num_rows)]
            assert list(table.flags) == flags
            assert list(table.time) == time

    def test_add_row_bad_data(self):
        t = tskit.NodeTable()
        with pytest.raises(TypeError):
            t.add_row(flags="x")
        with pytest.raises(TypeError):
            t.add_row(time="x")
        with pytest.raises(TypeError):
            t.add_row(individual="x")
        with pytest.raises(TypeError):
            t.add_row(population="x")
        with pytest.raises(TypeError):
            t.add_row(metadata=123)


class TestEdgeTable(*common_tests):
    columns = [
        DoubleColumn("left"),
        DoubleColumn("right"),
        Int32Column("parent"),
        Int32Column("child"),
    ]
    equal_len_columns = [["left", "right", "parent", "child"]]
    string_colnames = []
    binary_colnames = ["metadata"]
    ragged_list_columns = [(CharColumn("metadata"), UInt32Column("metadata_offset"))]
    input_parameters = [("max_rows_increment", 0)]
    table_class = tskit.EdgeTable

    def test_simple_example(self):
        t = tskit.EdgeTable()
        t.add_row(left=0, right=1, parent=2, child=3, metadata=b"123")
        t.add_row(1, 2, 3, 4, b"\xf0")
        assert len(t) == 2
        assert dataclasses.astuple(t[0]) == (0, 1, 2, 3, b"123")
        assert dataclasses.astuple(t[1]) == (1, 2, 3, 4, b"\xf0")
        assert t[0].left == 0
        assert t[0].right == 1
        assert t[0].parent == 2
        assert t[0].child == 3
        assert t[0].metadata == b"123"
        assert t[0] == t[-2]
        assert t[1] == t[-1]
        with pytest.raises(IndexError):
            t.__getitem__(-3)

    def test_add_row_defaults(self):
        t = tskit.EdgeTable()
        assert t.add_row(0, 0, 0, 0) == 0
        assert len(t.metadata) == 0
        assert t.metadata_offset[0] == 0

    def test_add_row_bad_data(self):
        t = tskit.EdgeTable()
        with pytest.raises(TypeError):
            t.add_row(left="x", right=0, parent=0, child=0)
        with pytest.raises(TypeError):
            t.add_row()
        with pytest.raises(TypeError):
            t.add_row(0, 0, 0, 0, metadata=123)


class TestSiteTable(*common_tests):
    columns = [DoubleColumn("position")]
    ragged_list_columns = [
        (CharColumn("ancestral_state"), UInt32Column("ancestral_state_offset")),
        (CharColumn("metadata"), UInt32Column("metadata_offset")),
    ]
    equal_len_columns = [["position"]]
    string_colnames = ["ancestral_state"]
    binary_colnames = ["metadata"]
    input_parameters = [("max_rows_increment", 0)]
    table_class = tskit.SiteTable

    def test_simple_example(self):
        t = tskit.SiteTable()
        t.add_row(position=0, ancestral_state="1", metadata=b"2")
        t.add_row(1, "2", b"\xf0")
        s = str(t)
        assert len(s) > 0
        assert len(t) == 2
        assert dataclasses.astuple(t[0]) == (0, "1", b"2")
        assert dataclasses.astuple(t[1]) == (1, "2", b"\xf0")
        assert t[0].position == 0
        assert t[0].ancestral_state == "1"
        assert t[0].metadata == b"2"
        assert t[0] == t[-2]
        assert t[1] == t[-1]
        with pytest.raises(IndexError):
            t.__getitem__(2)
        with pytest.raises(IndexError):
            t.__getitem__(-3)

    def test_add_row_bad_data(self):
        t = tskit.SiteTable()
        t.add_row(0, "A")
        with pytest.raises(TypeError):
            t.add_row("x", "A")
        with pytest.raises(TypeError):
            t.add_row(0, 0)
        with pytest.raises(TypeError):
            t.add_row(0, "A", metadata=[0, 1, 2])

    def test_packset_ancestral_state(self):
        for num_rows in [0, 10, 100]:
            input_data = self.make_input_data(num_rows)
            table = self.table_class()
            table.set_columns(**input_data)
            ancestral_states = [tsutil.random_strings(10) for _ in range(num_rows)]
            ancestral_state, ancestral_state_offset = tskit.pack_strings(
                ancestral_states
            )
            table.packset_ancestral_state(ancestral_states)
            assert np.array_equal(table.ancestral_state, ancestral_state)
            assert np.array_equal(table.ancestral_state_offset, ancestral_state_offset)


class TestMutationTable(*common_tests):
    columns = [
        Int32Column("site"),
        Int32Column("node"),
        DoubleColumn("time"),
        Int32Column("parent"),
    ]
    ragged_list_columns = [
        (CharColumn("derived_state"), UInt32Column("derived_state_offset")),
        (CharColumn("metadata"), UInt32Column("metadata_offset")),
    ]
    equal_len_columns = [["site", "node", "time"]]
    string_colnames = ["derived_state"]
    binary_colnames = ["metadata"]
    input_parameters = [("max_rows_increment", 0)]
    table_class = tskit.MutationTable

    def test_simple_example(self):
        t = tskit.MutationTable()
        t.add_row(site=0, node=1, derived_state="2", parent=3, metadata=b"4", time=5)
        t.add_row(1, 2, "3", 4, b"\xf0", 6)
        t.add_row(
            site=0,
            node=1,
            derived_state="2",
            parent=3,
            metadata=b"4",
            time=tskit.UNKNOWN_TIME,
        )
        s = str(t)
        assert len(s) > 0
        assert len(t) == 3
        assert dataclasses.astuple(t[0]) == (0, 1, "2", 3, b"4", 5)
        assert dataclasses.astuple(t[1]) == (1, 2, "3", 4, b"\xf0", 6)
        assert t[0].site == 0
        assert t[0].node == 1
        assert t[0].derived_state == "2"
        assert t[0].parent == 3
        assert t[0].metadata == b"4"
        assert t[0].time == 5
        assert t[0] == t[-3]
        assert t[1] == t[-2]
        assert t[2] == t[-1]
        with pytest.raises(IndexError):
            t.__getitem__(-4)

    def test_add_row_defaults(self):
        t = tskit.MutationTable()
        assert t.add_row(0, 0, "A", 0) == 0
        assert len(t.metadata) == 0
        assert t.metadata_offset[0] == 0
        assert tskit.is_unknown_time(t.time[0])

    def test_add_row_bad_data(self):
        t = tskit.MutationTable()
        t.add_row(0, 0, "A")
        with pytest.raises(TypeError):
            t.add_row("0", 0, "A")
        with pytest.raises(TypeError):
            t.add_row(0, "0", "A")
        with pytest.raises(TypeError):
            t.add_row(0, 0, "A", parent=None)
        with pytest.raises(TypeError):
            t.add_row(0, 0, "A", metadata=[0])
        with pytest.raises(TypeError):
            t.add_row(0, 0, "A", time="A")

    def test_packset_derived_state(self):
        for num_rows in [0, 10, 100]:
            input_data = self.make_input_data(num_rows)
            table = self.table_class()
            table.set_columns(**input_data)
            derived_states = [tsutil.random_strings(10) for _ in range(num_rows)]
            derived_state, derived_state_offset = tskit.pack_strings(derived_states)
            table.packset_derived_state(derived_states)
            assert np.array_equal(table.derived_state, derived_state)
            assert np.array_equal(table.derived_state_offset, derived_state_offset)

    def test_keep_rows_data(self):
        input_data = self.make_input_data(100)
        t1 = self.table_class()
        # Set the parent column to -1s for this simple test as
        # we need to reason about reference integrity
        t1.append_columns(**input_data)
        t1.parent = np.full_like(t1.parent, -1)
        t2 = t1.copy()
        keep = np.ones(len(t1), dtype=bool)
        # Only keep even
        keep[::2] = 0
        t1.keep_rows(keep)
        keep_rows_definition(t2, keep)
        assert t1.equals(t2)


class TestMigrationTable(*common_tests):
    columns = [
        DoubleColumn("left"),
        DoubleColumn("right"),
        Int32Column("node"),
        Int32Column("source"),
        Int32Column("dest"),
        DoubleColumn("time"),
    ]
    ragged_list_columns = [(CharColumn("metadata"), UInt32Column("metadata_offset"))]
    string_colnames = []
    binary_colnames = ["metadata"]
    input_parameters = [("max_rows_increment", 0)]
    equal_len_columns = [["left", "right", "node", "source", "dest", "time"]]
    table_class = tskit.MigrationTable

    def test_simple_example(self):
        t = tskit.MigrationTable()
        t.add_row(left=0, right=1, node=2, source=3, dest=4, time=5, metadata=b"123")
        t.add_row(1, 2, 3, 4, 5, 6, b"\xf0")
        assert len(t) == 2
        assert dataclasses.astuple(t[0]) == (0, 1, 2, 3, 4, 5, b"123")
        assert dataclasses.astuple(t[1]) == (1, 2, 3, 4, 5, 6, b"\xf0")
        assert t[0].left == 0
        assert t[0].right == 1
        assert t[0].node == 2
        assert t[0].source == 3
        assert t[0].dest == 4
        assert t[0].time == 5
        assert t[0].metadata == b"123"
        assert t[0] == t[-2]
        assert t[1] == t[-1]
        with pytest.raises(IndexError):
            t.__getitem__(-3)

    def test_add_row_defaults(self):
        t = tskit.MigrationTable()
        assert t.add_row(0, 0, 0, 0, 0, 0) == 0
        assert len(t.metadata) == 0
        assert t.metadata_offset[0] == 0

    def test_add_row_bad_data(self):
        t = tskit.MigrationTable()
        with pytest.raises(TypeError):
            t.add_row(left="x", right=0, node=0, source=0, dest=0, time=0)
        with pytest.raises(TypeError):
            t.add_row()
        with pytest.raises(TypeError):
            t.add_row(0, 0, 0, 0, 0, 0, metadata=123)


class TestProvenanceTable(CommonTestsMixin, AssertEqualsMixin):
    columns = []
    ragged_list_columns = [
        (CharColumn("timestamp"), UInt32Column("timestamp_offset")),
        (CharColumn("record"), UInt32Column("record_offset")),
    ]
    equal_len_columns = [[]]
    string_colnames = ["record", "timestamp"]
    binary_colnames = []
    input_parameters = [("max_rows_increment", 0)]
    table_class = tskit.ProvenanceTable

    def test_simple_example(self):
        t = tskit.ProvenanceTable()
        t.add_row(timestamp="0", record="1")
        t.add_row("2", "1")  # The orders are reversed for default timestamp.
        assert len(t) == 2
        assert dataclasses.astuple(t[0]) == ("0", "1")
        assert dataclasses.astuple(t[1]) == ("1", "2")
        assert t[0].timestamp == "0"
        assert t[0].record == "1"
        assert t[0] == t[-2]
        assert t[1] == t[-1]
        with pytest.raises(IndexError):
            t.__getitem__(-3)

    def test_add_row_bad_data(self):
        t = tskit.ProvenanceTable()
        t.add_row("a", "b")
        with pytest.raises(TypeError):
            t.add_row(0, "b")
        with pytest.raises(TypeError):
            t.add_row("a", 0)

    def test_packset_timestamp(self):
        t = tskit.ProvenanceTable()
        t.add_row(timestamp="0", record="1")
        t.add_row(timestamp="1", record="2")
        t.packset_timestamp(["AAAA", "BBBB"])
        assert t[0].timestamp == "AAAA"
        assert t[1].timestamp == "BBBB"

    def test_packset_record(self):
        t = tskit.ProvenanceTable()
        t.add_row(timestamp="0", record="1")
        t.add_row(timestamp="1", record="2")
        t.packset_record(["AAAA", "BBBB"])
        assert t[0].record == "AAAA"
        assert t[1].record == "BBBB"


class TestPopulationTable(*common_tests):
    metadata_mandatory = True
    columns = []
    ragged_list_columns = [(CharColumn("metadata"), UInt32Column("metadata_offset"))]
    equal_len_columns = [[]]
    string_colnames = []
    binary_colnames = ["metadata"]
    input_parameters = [("max_rows_increment", 0)]
    table_class = tskit.PopulationTable

    def test_simple_example(self):
        t = tskit.PopulationTable()
        t.add_row(metadata=b"\xf0")
        t.add_row(b"1")
        s = str(t)
        assert len(s) > 0
        assert len(t) == 2
        assert dataclasses.astuple(t[0]) == (b"\xf0",)
        assert t[0].metadata == b"\xf0"
        assert dataclasses.astuple(t[1]) == (b"1",)
        with pytest.raises(IndexError):
            t.__getitem__(-3)

    def test_add_row_defaults(self):
        t = tskit.PopulationTable()
        assert t.add_row() == 0
        assert len(t.metadata) == 0
        assert t.metadata_offset[0] == 0

    def test_add_row_bad_data(self):
        t = tskit.PopulationTable()
        t.add_row()
        with pytest.raises(TypeError):
            t.add_row(metadata=[0])


class TestTableCollectionIndexes:
    def test_index(self):
        i = np.arange(20)
        r = np.arange(20)[::-1]
        index = tskit.TableCollectionIndexes(
            edge_insertion_order=i, edge_removal_order=r
        )
        assert np.array_equal(index.edge_insertion_order, i)
        assert np.array_equal(index.edge_removal_order, r)
        d = index.asdict()
        assert np.array_equal(d["edge_insertion_order"], i)
        assert np.array_equal(d["edge_removal_order"], r)

        index = tskit.TableCollectionIndexes()
        assert index.edge_insertion_order is None
        assert index.edge_removal_order is None
        assert index.asdict() == {}


class TestSortTables:
    """
    Tests for the TableCollection.sort() and TableCollection.canonicalise() methods.
    """

    random_seed = 12345

    def verify_sort_equality(self, tables, seed):
        tables1 = tables.copy()
        tsutil.shuffle_tables(
            tables1,
            seed,
            shuffle_populations=False,
        )
        tables1.individuals.packset_metadata(
            [bytes(str(i), "utf-8") for i in range(tables1.individuals.num_rows)]
        )
        tables2 = tables1.copy()
        tables1.sort()
        tsutil.py_sort(tables2)

        # TODO - Check the sorted tables are valid ts, currently fails due to mutations
        # tables1.tree_sequence()
        # tables2.tree_sequence()

        tables1.assert_equals(tables2)

    def verify_canonical_equality(self, tables, seed):
        # Migrations not supported
        tables.migrations.clear()

        for ru in [True, False]:
            tsk_tables = tables.copy()
            tsutil.shuffle_tables(
                tsk_tables,
                seed,
            )
            py_tables = tsk_tables.copy()
            tsk_tables.canonicalise(remove_unreferenced=ru)
            tsutil.py_canonicalise(py_tables, remove_unreferenced=ru)
            tsk_tables.assert_equals(py_tables)

    def verify_sort_mutation_consistency(self, orig_tables, seed):
        tables = orig_tables.copy()
        mut_map = {s.position: [] for s in tables.sites}
        for mut in tables.mutations:
            mut_map[tables.sites[mut.site].position].append(
                (mut.node, mut.derived_state, mut.metadata)
            )
        tsutil.shuffle_tables(tables, seed, shuffle_populations=False)
        for mut in tables.mutations:
            site = tables.sites[mut.site]
            assert (mut.node, mut.derived_state, mut.metadata) in mut_map[site.position]
        tables.sort()
        for mut in tables.mutations:
            site = tables.sites[mut.site]
            assert (mut.node, mut.derived_state, mut.metadata) in mut_map[site.position]

    def verify_randomise_tables(self, orig_tables, seed):
        # Check we can shuffle everything and then put it back in canonical form
        tables = orig_tables.copy()
        tables.sort()
        sorted_tables = tables.copy()

        # First randomize only edges: this should work without canonical sorting.
        tsutil.shuffle_tables(
            tables,
            seed=seed,
            shuffle_edges=True,
            shuffle_populations=False,
            shuffle_individuals=False,
            shuffle_sites=False,
            shuffle_mutations=False,
        )
        tables.sort()
        tables.assert_equals(sorted_tables)

        # Now also randomize sites, mutations and individuals
        tables.canonicalise(remove_unreferenced=False)
        sorted_tables = tables.copy()
        tsutil.shuffle_tables(
            tables,
            seed=1234,
            shuffle_populations=False,
        )
        tables.canonicalise(remove_unreferenced=False)
        tables.assert_equals(sorted_tables)

        # Finally, randomize everything else
        tsutil.shuffle_tables(tables, seed=1234)
        tables.canonicalise(remove_unreferenced=False)
        tables.assert_equals(sorted_tables)

        # Check the canonicalised form meets the tree sequence requirements
        tables.tree_sequence()

    def verify_sort(self, tables, seed):
        self.verify_sort_equality(tables, seed)
        self.verify_canonical_equality(tables, seed)
        self.verify_sort_mutation_consistency(tables, seed)
        self.verify_randomise_tables(tables, seed)

    def verify_sort_offset(self, ts):
        """
        Verifies the behaviour of the edge_start offset value.
        """
        tables = ts.dump_tables()
        edges = tables.edges.copy()
        starts = [0]
        if len(edges) > 2:
            starts = [0, 1, len(edges) // 2, len(edges) - 2]
        for start in starts:
            # Unsort the edges starting from index start
            all_edges = list(ts.edges())
            keep = all_edges[:start]
            reversed_edges = all_edges[start:][::-1]
            all_edges = keep + reversed_edges
            tables.edges.clear()
            for e in all_edges:
                tables.edges.append(e)
            # Verify that import fails for reversed edges
            with pytest.raises(_tskit.LibraryError):
                tables.tree_sequence()
            # If we sort after the start value we should still fail.
            tables.sort(edge_start=start + 1)
            with pytest.raises(_tskit.LibraryError):
                tables.tree_sequence()
            # Sorting from the correct index should give us back the original table.
            tables.edges.clear()
            for e in all_edges:
                tables.edges.append(e)
            tables.sort(edge_start=start)
            # Verify the new and old edges are equal.
            assert edges == tables.edges

        tables.tree_sequence()
        if len(tables.mutations) > 2:
            mutations = tables.mutations.copy()
            tables.mutations.clear()
            for m in mutations[::-1]:
                tables.mutations.append(m)
            with pytest.raises(_tskit.LibraryError):
                tables.tree_sequence()
            tables.sort(
                0, site_start=len(tables.sites), mutation_start=len(tables.mutations)
            )
            with pytest.raises(_tskit.LibraryError):
                tables.tree_sequence()
            tables.sort(0)
            tables.tree_sequence()

    def get_wf_example(self, seed):
        tables = wf.wf_sim(
            6,
            3,
            num_pops=2,
            seed=seed,
            num_loci=3,
            record_migrations=True,
        )
        tables.sort()
        ts = tables.tree_sequence()
        return ts

    def test_wf_example(self):
        tables = wf.wf_sim(
            N=6,
            ngens=3,
            num_pops=2,
            mig_rate=1.0,
            deep_history=False,
            seed=42,
            record_migrations=True,
        )
        self.verify_sort(tables, 42)

    def test_single_tree_no_mutations(self):
        ts = msprime.simulate(10, random_seed=self.random_seed)
        self.verify_sort_offset(ts)
        self.verify_sort(ts.tables, 432)

    def test_single_tree_no_mutations_metadata(self):
        ts = msprime.simulate(10, random_seed=self.random_seed)
        ts = tsutil.add_random_metadata(ts, self.random_seed)
        self.verify_sort(ts.tables, 12)

    def test_many_trees_no_mutations(self):
        ts = msprime.simulate(10, recombination_rate=2, random_seed=self.random_seed)
        assert ts.num_trees > 2
        self.verify_sort_offset(ts)
        self.verify_sort(ts.tables, 31)

    def test_single_tree_mutations(self):
        ts = msprime.simulate(10, mutation_rate=2, random_seed=self.random_seed)
        assert ts.num_sites > 2
        self.verify_sort_offset(ts)
        self.verify_sort(ts.tables, 83)

    def test_single_tree_mutations_metadata(self):
        ts = msprime.simulate(10, mutation_rate=2, random_seed=self.random_seed)
        assert ts.num_sites > 2
        ts = tsutil.add_random_metadata(ts, self.random_seed)
        self.verify_sort(ts.tables, 384)

    def test_single_tree_multichar_mutations(self):
        ts = msprime.simulate(10, random_seed=self.random_seed)
        ts = tsutil.insert_multichar_mutations(ts, self.random_seed)
        self.verify_sort(ts.tables, 185)

    def test_single_tree_multichar_mutations_metadata(self):
        ts = msprime.simulate(10, random_seed=self.random_seed)
        ts = tsutil.insert_multichar_mutations(ts, self.random_seed)
        ts = tsutil.add_random_metadata(ts, self.random_seed)
        self.verify_sort(ts.tables, 2175)

    def test_many_trees_mutations(self):
        ts = msprime.simulate(
            10, recombination_rate=2, mutation_rate=2, random_seed=self.random_seed
        )
        assert ts.num_trees > 2
        assert ts.num_sites > 2
        self.verify_sort_offset(ts)
        self.verify_sort(ts.tables, 173)

    def test_many_trees_multichar_mutations(self):
        ts = msprime.simulate(10, recombination_rate=2, random_seed=self.random_seed)
        assert ts.num_trees > 2
        ts = tsutil.insert_multichar_mutations(ts, self.random_seed)
        self.verify_sort(ts.tables, 16)

    def test_many_trees_multichar_mutations_metadata(self):
        ts = msprime.simulate(10, recombination_rate=2, random_seed=self.random_seed)
        assert ts.num_trees > 2
        ts = tsutil.insert_multichar_mutations(ts, self.random_seed)
        ts = tsutil.add_random_metadata(ts, self.random_seed)
        self.verify_sort(ts.tables, 91)

    def get_nonbinary_example(self, mutation_rate):
        ts = msprime.simulate(
            sample_size=20,
            recombination_rate=10,
            random_seed=self.random_seed,
            mutation_rate=mutation_rate,
            demographic_events=[
                msprime.SimpleBottleneck(time=0.5, population=0, proportion=1)
            ],
        )
        # Make sure this really has some non-binary nodes
        found = False
        for e in ts.edgesets():
            if len(e.children) > 2:
                found = True
                break
        assert found
        assert ts.num_trees > 2
        return ts

    def test_nonbinary_trees(self):
        ts = self.get_nonbinary_example(mutation_rate=0)
        self.verify_sort_offset(ts)
        self.verify_sort(ts.tables, 9182)

    def test_nonbinary_trees_mutations(self):
        ts = self.get_nonbinary_example(mutation_rate=2)
        assert ts.num_trees > 2
        assert ts.num_sites > 2
        self.verify_sort_offset(ts)
        self.verify_sort(ts.tables, 44)

    def test_unknown_times(self):
        ts = self.get_wf_example(seed=486)
        ts = tsutil.insert_branch_mutations(ts, mutations_per_branch=2)
        ts = tsutil.remove_mutation_times(ts)
        self.verify_sort(ts.tables, 9182)

    def test_no_mutation_parents(self):
        # we should maintain relative order of mutations when all else fails:
        # these tables are not canonicalizable (since we don't sort on derived state)
        rng = random.Random(7000)
        alleles = ["A", "B", "C", "D", "E", "F", "G"]
        for t in [0.5, None]:
            rng.shuffle(alleles)
            tables = tskit.TableCollection(sequence_length=1)
            tables.nodes.add_row(time=0, flags=tskit.NODE_IS_SAMPLE)
            tables.sites.add_row(position=0, ancestral_state="")
            for a in alleles:
                tables.mutations.add_row(site=0, node=0, derived_state=a, time=t)
            tables_canonical = tables.copy()
            tables_canonical.canonicalise()
            tables.sort()
            for t in (tables, tables_canonical):
                for a, mut in zip(alleles, t.mutations):
                    assert a == mut.derived_state
                self.verify_sort_equality(t, 985)
                self.verify_sort_mutation_consistency(t, 985)

    def test_stable_individual_order(self):
        # canonical should retain individual order lacking any other information
        tables = tskit.TableCollection(sequence_length=100)
        for a in "arbol":
            tables.individuals.add_row(metadata=a.encode())
        tables2 = tables.copy()
        tables2.canonicalise(remove_unreferenced=False)
        tables.assert_equals(tables2)

    def test_discrete_times(self):
        ts = self.get_wf_example(seed=623)
        ts = tsutil.insert_discrete_time_mutations(ts)
        self.verify_sort(ts.tables, 9183)

    def test_incompatible_edges(self):
        ts1 = msprime.simulate(10, random_seed=self.random_seed)
        ts2 = msprime.simulate(20, random_seed=self.random_seed)
        tables1 = ts1.dump_tables()
        tables2 = ts2.dump_tables()
        tables2.edges.set_columns(**tables1.edges.asdict())
        # The edges in tables2 will refer to nodes that don't exist.
        with pytest.raises(_tskit.LibraryError):
            tables2.sort()
        with pytest.raises(_tskit.LibraryError):
            tables2.canonicalise()

    def test_incompatible_sites(self):
        ts1 = msprime.simulate(10, random_seed=self.random_seed)
        ts2 = msprime.simulate(10, mutation_rate=2, random_seed=self.random_seed)
        assert ts2.num_sites > 1
        tables1 = ts1.dump_tables()
        tables2 = ts2.dump_tables()
        # The mutations in tables2 will refer to sites that don't exist.
        tables1.mutations.set_columns(**tables2.mutations.asdict())
        with pytest.raises(_tskit.LibraryError):
            tables1.sort()
        with pytest.raises(_tskit.LibraryError):
            tables1.canonicalise()

    def test_incompatible_mutation_nodes(self):
        ts1 = msprime.simulate(2, random_seed=self.random_seed)
        ts2 = msprime.simulate(10, mutation_rate=2, random_seed=self.random_seed)
        assert ts2.num_sites > 1
        tables1 = ts1.dump_tables()
        tables2 = ts2.dump_tables()
        # The mutations in tables2 will refer to nodes that don't exist.
        tables1.sites.set_columns(**tables2.sites.asdict())
        tables1.mutations.set_columns(**tables2.mutations.asdict())
        with pytest.raises(_tskit.LibraryError):
            tables1.sort()
        with pytest.raises(_tskit.LibraryError):
            tables1.canonicalise()

    def test_empty_tables(self):
        tables = tskit.TableCollection(1)
        tables.sort()
        assert tables.nodes.num_rows == 0
        assert tables.edges.num_rows == 0
        assert tables.sites.num_rows == 0
        assert tables.mutations.num_rows == 0
        assert tables.migrations.num_rows == 0


class TestSortMigrations:
    """
    Tests that migrations are correctly ordered when sorting tables.
    """

    def test_msprime_output_unmodified(self):
        pop_configs = [msprime.PopulationConfiguration(5) for _ in range(3)]
        migration_matrix = [[0, 1, 1], [1, 0, 1], [1, 1, 0]]
        ts = msprime.simulate(
            recombination_rate=0.1,
            population_configurations=pop_configs,
            migration_matrix=migration_matrix,
            record_migrations=True,
            random_seed=1,
        )
        assert ts.num_migrations > 100
        tables = ts.tables.copy()
        tables.sort()
        tables.assert_equals(ts.tables, ignore_provenance=True)

    def test_full_sort_order(self):
        tables = tskit.TableCollection(1)
        for _ in range(3):
            tables.nodes.add_row()
            tables.populations.add_row()
        for left in [0, 0.5]:
            for a_time in range(3):
                for node in range(2):
                    tables.migrations.add_row(
                        left=left, right=1, node=node, source=0, dest=1, time=a_time
                    )
                    tables.migrations.add_row(
                        left=left, right=1, node=node, source=1, dest=2, time=a_time
                    )

        sorted_list = sorted(
            tables.migrations, key=lambda m: (m.time, m.source, m.dest, m.left, m.node)
        )
        assert sorted_list != list(tables.migrations)
        tables.sort()
        assert sorted_list == list(tables.migrations)


class TestSortMutations:
    """
    Tests that mutations are correctly ordered when sorting tables.
    """

    def test_sort_mutations_stability(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        """
        )
        sites = io.StringIO(
            """\
        position    ancestral_state
        0.1     0
        0.2     0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    derived_state   parent
        1       0       1               -1
        1       1       1               -1
        0       1       1               -1
        0       0       1               -1
        """
        )
        ts = tskit.load_text(
            nodes=nodes,
            edges=edges,
            sites=sites,
            mutations=mutations,
            sequence_length=1,
            strict=False,
        )
        # Load text automatically calls tables.sort(), so we can test the
        # output directly.
        sites = ts.tables.sites
        mutations = ts.tables.mutations
        assert len(sites) == 2
        assert len(mutations) == 4
        assert list(mutations.site) == [0, 0, 1, 1]
        assert list(mutations.node) == [1, 0, 0, 1]

    def test_sort_mutations_remap_parent_id(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        """
        )
        sites = io.StringIO(
            """\
        position    ancestral_state
        0.1     0
        0.2     0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    time    derived_state   parent
        1       0       0.5     1               -1
        1       0       0.25    0               0
        1       0       0       1               1
        0       0       0.5     1               -1
        0       0       0.125   0               3
        0       0       0       1               4
        """
        )
        ts = tskit.load_text(
            nodes=nodes,
            edges=edges,
            sites=sites,
            mutations=mutations,
            sequence_length=1,
            strict=False,
        )
        # Load text automatically calls sort tables, so we can test the
        # output directly.
        sites = ts.tables.sites
        mutations = ts.tables.mutations
        assert len(sites) == 2
        assert len(mutations) == 6
        assert list(mutations.site) == [0, 0, 0, 1, 1, 1]
        assert list(mutations.node) == [0, 0, 0, 0, 0, 0]
        assert list(mutations.time) == [0.5, 0.125, 0.0, 0.5, 0.25, 0.0]
        assert list(mutations.parent) == [-1, 0, 1, -1, 3, 4]

    def test_sort_mutations_bad_parent_id(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        """
        )
        sites = io.StringIO(
            """\
        position    ancestral_state
        0.1     0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    derived_state   parent
        1       0       1               -2
        """
        )
        with pytest.raises(ValueError):
            tskit.load_text(
                nodes=nodes,
                edges=edges,
                sites=sites,
                mutations=mutations,
                sequence_length=1,
                strict=False,
            )

    def test_sort_mutations_time(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           -6
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        """
        )
        sites = io.StringIO(
            """\
        position    ancestral_state
        0.1     0
        0.2     0
        0.3     0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    time    derived_state   parent
        2       0       4       a              -1
        2       0       -5      b              -1
        2       0       6       c              -1
        1       0       0.5     d              -1
        1       0       0.5     e              -1
        1       0       0.5     f              -1
        0       0       1       g              -1
        0       0       2       h              -1
        0       0       3       i              -1
        """
        )
        ts = tskit.load_text(
            nodes=nodes,
            edges=edges,
            sites=sites,
            mutations=mutations,
            sequence_length=1,
            strict=False,
        )
        # Load text automatically calls sort tables, so we can test the
        # output directly.
        sites = ts.tables.sites
        mutations = ts.tables.mutations
        assert len(sites) == 3
        assert len(mutations) == 9
        assert list(mutations.site) == [0, 0, 0, 1, 1, 1, 2, 2, 2]
        assert list(mutations.node) == [0, 0, 0, 0, 0, 0, 0, 0, 0]
        # Nans are not equal so swap in -1
        times = mutations.time
        times[np.isnan(times)] = -1
        assert list(times) == [3.0, 2.0, 1.0, 0.5, 0.5, 0.5, 6.0, 4.0, -5.0]
        assert list(mutations.derived_state) == list(
            map(ord, ["i", "h", "g", "d", "e", "f", "c", "a", "b"])
        )
        assert list(mutations.parent) == [-1, -1, -1, -1, -1, -1, -1, -1, -1]


class TestTablesToTreeSequence:
    """
    Tests for the .tree_sequence() method of a TableCollection.
    """

    def test_round_trip(self):
        a = msprime.simulate(5, mutation_rate=1, random_seed=42)
        tables = a.dump_tables()
        b = tables.tree_sequence()
        assert a.tables == b.tables


class TestMutationTimeErrors:
    def test_younger_than_node_below(self):
        ts = msprime.simulate(5, mutation_rate=1, random_seed=42)
        tables = ts.dump_tables()
        tables.mutations.time = np.zeros(len(tables.mutations.time), dtype=np.float64)
        with pytest.raises(
            _tskit.LibraryError,
            match="A mutation's time must be >= the node time, or be marked as"
            " 'unknown'",
        ):
            tables.tree_sequence()

    def test_older_than_node_above(self):
        ts = msprime.simulate(5, mutation_rate=1, random_seed=42)
        tables = ts.dump_tables()
        tables.mutations.time = (
            np.ones(len(tables.mutations.time), dtype=np.float64) * 42
        )
        with pytest.raises(
            _tskit.LibraryError,
            match="A mutation's time must be < the parent node of the edge on which it"
            " occurs, or be marked as 'unknown'",
        ):
            tables.tree_sequence()

    def test_older_than_parent_node(self):
        ts = msprime.simulate(
            10, random_seed=42, mutation_rate=0.0, recombination_rate=1.0
        )
        ts = tsutil.jukes_cantor(
            ts, num_sites=10, mu=1, multiple_per_node=False, seed=42
        )
        tables = ts.dump_tables()
        assert sum(tables.mutations.parent != -1) != 0
        # Make all times the node time
        times = tables.nodes.time[tables.mutations.node]
        # Then make mutations without a parent really old
        times[tables.mutations.parent == -1] = 64.0
        tables.mutations.time = times
        tables.sort()
        with pytest.raises(
            _tskit.LibraryError,
            match="A mutation's time must be < the parent node of the edge on which it"
            " occurs, or be marked as 'unknown'",
        ):
            tables.tree_sequence()

    def test_older_than_parent_mutation(self):
        ts = msprime.simulate(
            10, random_seed=42, mutation_rate=0.0, recombination_rate=1.0
        )
        ts = tsutil.jukes_cantor(
            ts, num_sites=10, mu=1, multiple_per_node=False, seed=42
        )
        tables = ts.dump_tables()
        tables.compute_mutation_times()
        assert sum(tables.mutations.parent != -1) != 0
        times = tables.mutations.time
        # Then make mutations without a parent really old
        times[tables.mutations.parent != -1] = 64.0
        tables.mutations.time = times
        with pytest.raises(_tskit.LibraryError):
            tables.tree_sequence()

    def test_unsorted_times(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        """
        )
        sites = io.StringIO(
            """\
        position    ancestral_state
        0.1     0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    time    derived_state   parent
        0       0       1       0              -1
        0       0       2       0              -1
        0       0       3       0              -1
        """
        )
        ts = tskit.load_text(
            nodes=nodes,
            edges=edges,
            sites=sites,
            mutations=mutations,
            sequence_length=1,
            strict=False,
        )
        tables = ts.dump_tables()
        tables.mutations.time = tables.mutations.time[::-1]
        with pytest.raises(
            _tskit.LibraryError,
            match="Mutations must be provided in non-decreasing site order and"
            " non-increasing"
            " time order within each site",
        ):
            tables.tree_sequence()

    def test_mixed_known_and_unknown(self):
        ts = msprime.simulate(
            10, random_seed=42, mutation_rate=0.0, recombination_rate=1.0
        )
        ts = tsutil.jukes_cantor(
            ts, num_sites=10, mu=1, multiple_per_node=False, seed=42
        )
        tables = ts.dump_tables()
        tables.compute_mutation_times()
        tables.sort()
        times = tables.mutations.time
        # Unknown times on diff sites pass
        times[(tables.mutations.site % 2) == 0] = tskit.UNKNOWN_TIME
        tables.mutations.time = times
        tables.tree_sequence()
        # Mixed known/unknown times on sites fail
        times[::2] = tskit.UNKNOWN_TIME
        tables.mutations.time = times
        with pytest.raises(
            _tskit.LibraryError,
            match="Mutation times must either be all marked 'unknown', or all be known "
            "values for any single site.",
        ):
            tables.tree_sequence()


class TestNanDoubleValues:
    """
    In some tables we need to guard against NaN/infinite values in the input.
    """

    def test_edge_coords(self):
        ts = msprime.simulate(5, mutation_rate=1, random_seed=42)

        tables = ts.dump_tables()
        bad_coords = tables.edges.left + float("inf")
        tables.edges.left = bad_coords
        with pytest.raises(_tskit.LibraryError):
            tables.tree_sequence()

        tables = ts.dump_tables()
        bad_coords = tables.edges.right + float("nan")
        tables.edges.right = bad_coords
        with pytest.raises(_tskit.LibraryError):
            tables.tree_sequence()

    def test_migrations(self):
        ts = msprime.simulate(5, mutation_rate=1, random_seed=42)

        tables = ts.dump_tables()
        tables.populations.add_row()
        tables.migrations.add_row(float("inf"), 1, time=0, node=0, source=0, dest=1)
        with pytest.raises(_tskit.LibraryError):
            tables.tree_sequence()

        tables = ts.dump_tables()
        tables.populations.add_row()
        tables.migrations.add_row(0, float("nan"), time=0, node=0, source=0, dest=1)
        with pytest.raises(_tskit.LibraryError):
            tables.tree_sequence()

        tables = ts.dump_tables()
        tables.populations.add_row()
        tables.migrations.add_row(0, 1, time=float("nan"), node=0, source=0, dest=1)
        with pytest.raises(_tskit.LibraryError):
            tables.tree_sequence()

    def test_site_positions(self):
        ts = msprime.simulate(5, mutation_rate=1, random_seed=42)
        tables = ts.dump_tables()
        bad_pos = tables.sites.position.copy()
        bad_pos[-1] = np.inf
        tables.sites.position = bad_pos
        with pytest.raises(_tskit.LibraryError):
            tables.tree_sequence()

    def test_node_times(self):
        ts = msprime.simulate(5, mutation_rate=1, random_seed=42)
        tables = ts.dump_tables()
        bad_times = tables.nodes.time.copy()
        bad_times[-1] = np.inf
        tables.nodes.time = bad_times
        with pytest.raises(_tskit.LibraryError, match="Times must be finite"):
            tables.tree_sequence()
        bad_times[-1] = math.nan
        tables.nodes.time = bad_times
        with pytest.raises(_tskit.LibraryError, match="Times must be finite"):
            tables.tree_sequence()

    def test_mutation_times(self):
        ts = msprime.simulate(5, mutation_rate=1, random_seed=42)
        tables = ts.dump_tables()
        bad_times = tables.mutations.time.copy()
        bad_times[-1] = np.inf
        tables.mutations.time = bad_times
        with pytest.raises(_tskit.LibraryError, match="Times must be finite"):
            tables.tree_sequence()
        bad_times = tables.mutations.time.copy()
        bad_times[-1] = math.nan
        tables.mutations.time = bad_times
        with pytest.raises(_tskit.LibraryError, match="Times must be finite"):
            tables.tree_sequence()

    def test_individual(self):
        ts = msprime.simulate(12, mutation_rate=1, random_seed=42)
        ts = tsutil.insert_random_ploidy_individuals(ts, seed=42)
        assert ts.num_individuals > 1
        tables = ts.dump_tables()
        bad_locations = tables.individuals.location.copy()
        bad_locations[0] = np.inf
        tables.individuals.location = bad_locations
        ts = tables.tree_sequence()


class TestSimplifyTables:
    """
    Tests for the simplify_tables function.
    """

    random_seed = 42

    def test_deprecated_zero_mutation_sites(self):
        ts = msprime.simulate(10, mutation_rate=1, random_seed=self.random_seed)
        tables = ts.dump_tables()
        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter("always")
            tables.simplify(ts.samples(), filter_zero_mutation_sites=True)
            assert len(w) == 1
            assert issubclass(w[-1].category, FutureWarning)

    def test_zero_mutation_sites(self):
        ts = msprime.simulate(10, mutation_rate=1, random_seed=self.random_seed)
        for filter_sites in [True, False]:
            t1 = ts.dump_tables()
            with pytest.warns(FutureWarning):
                t1.simplify([0, 1], filter_zero_mutation_sites=filter_sites)
            t2 = ts.dump_tables()
            t2.simplify([0, 1], filter_sites=filter_sites)
            t1.assert_equals(t2, ignore_provenance=True)
            if filter_sites:
                assert ts.num_sites > len(t1.sites)

    def test_full_samples(self):
        for n in [2, 10, 100, 1000]:
            ts = msprime.simulate(
                n, recombination_rate=1, mutation_rate=1, random_seed=self.random_seed
            )
            tables = ts.dump_tables()
            nodes_before = tables.nodes.copy()
            edges_before = tables.edges.copy()
            sites_before = tables.sites.copy()
            mutations_before = tables.mutations.copy()
            for samples in [None, list(ts.samples()), ts.samples()]:
                node_map = tables.simplify(samples=samples)
                assert node_map.shape == (len(nodes_before),)
                assert nodes_before == tables.nodes
                assert edges_before == tables.edges
                assert sites_before == tables.sites
                assert mutations_before == tables.mutations

    def test_bad_samples(self):
        n = 10
        ts = msprime.simulate(n, random_seed=self.random_seed)
        for bad_node in [-1, ts.num_nodes, 2**31 - 1]:
            tables = ts.dump_tables()
            with pytest.raises(_tskit.LibraryError):
                tables.simplify(samples=[0, bad_node])

    def test_bad_edge_ordering(self):
        ts = msprime.simulate(10, random_seed=self.random_seed)
        tables = ts.dump_tables()
        edges = tables.edges
        # Reversing the edges violates the ordering constraints.
        edges.set_columns(
            left=edges.left[::-1],
            right=edges.right[::-1],
            parent=edges.parent[::-1],
            child=edges.child[::-1],
        )
        with pytest.raises(_tskit.LibraryError):
            tables.simplify(samples=[0, 1])

    def test_bad_edges(self):
        ts = msprime.simulate(10, random_seed=self.random_seed)
        for bad_node in [-1, ts.num_nodes, ts.num_nodes + 1, 2**31 - 1]:
            # Bad parent node
            tables = ts.dump_tables()
            edges = tables.edges
            parent = edges.parent
            parent[0] = bad_node
            edges.set_columns(
                left=edges.left, right=edges.right, parent=parent, child=edges.child
            )
            with pytest.raises(_tskit.LibraryError):
                tables.simplify(samples=[0, 1])
            # Bad child node
            tables = ts.dump_tables()
            edges = tables.edges
            child = edges.child
            child[0] = bad_node
            edges.set_columns(
                left=edges.left, right=edges.right, parent=edges.parent, child=child
            )
            with pytest.raises(_tskit.LibraryError):
                tables.simplify(samples=[0, 1])
            # child == parent
            tables = ts.dump_tables()
            edges = tables.edges
            child = edges.child
            child[0] = edges.parent[0]
            edges.set_columns(
                left=edges.left, right=edges.right, parent=edges.parent, child=child
            )
            with pytest.raises(_tskit.LibraryError):
                tables.simplify(samples=[0, 1])
            # left == right
            tables = ts.dump_tables()
            edges = tables.edges
            left = edges.left
            left[0] = edges.right[0]
            edges.set_columns(
                left=left, right=edges.right, parent=edges.parent, child=edges.child
            )
            with pytest.raises(_tskit.LibraryError):
                tables.simplify(samples=[0, 1])
            # left > right
            tables = ts.dump_tables()
            edges = tables.edges
            left = edges.left
            left[0] = edges.right[0] + 1
            edges.set_columns(
                left=left, right=edges.right, parent=edges.parent, child=edges.child
            )
            with pytest.raises(_tskit.LibraryError):
                tables.simplify(samples=[0, 1])

    def test_bad_mutation_nodes(self):
        ts = msprime.simulate(10, random_seed=self.random_seed, mutation_rate=1)
        assert ts.num_mutations > 0
        for bad_node in [-1, ts.num_nodes, 2**31 - 1]:
            tables = ts.dump_tables()
            mutations = tables.mutations
            node = mutations.node
            node[0] = bad_node
            mutations.set_columns(
                site=mutations.site,
                node=node,
                derived_state=mutations.derived_state,
                derived_state_offset=mutations.derived_state_offset,
            )
            with pytest.raises(_tskit.LibraryError):
                tables.simplify(samples=[0, 1])

    def test_bad_mutation_sites(self):
        ts = msprime.simulate(10, random_seed=self.random_seed, mutation_rate=1)
        assert ts.num_mutations > 0
        for bad_site in [-1, ts.num_sites, 2**31 - 1]:
            tables = ts.dump_tables()
            mutations = tables.mutations
            site = mutations.site
            site[0] = bad_site
            mutations.set_columns(
                site=site,
                node=mutations.node,
                derived_state=mutations.derived_state,
                derived_state_offset=mutations.derived_state_offset,
            )
            with pytest.raises(_tskit.LibraryError):
                tables.simplify(samples=[0, 1])

    def test_bad_site_positions(self):
        ts = msprime.simulate(10, random_seed=self.random_seed, mutation_rate=1)
        assert ts.num_mutations > 0
        # Positions > sequence_length are valid, as we can have gaps at the end of
        # a tree sequence.
        for bad_position in [-1, -1e-6]:
            tables = ts.dump_tables()
            sites = tables.sites
            position = sites.position
            position[0] = bad_position
            sites.set_columns(
                position=position,
                ancestral_state=sites.ancestral_state,
                ancestral_state_offset=sites.ancestral_state_offset,
            )
            with pytest.raises(_tskit.LibraryError):
                tables.simplify(samples=[0, 1])

    def test_duplicate_positions(self):
        tables = tskit.TableCollection(sequence_length=1)
        tables.sites.add_row(0, ancestral_state="0")
        tables.sites.add_row(0, ancestral_state="0")
        with pytest.raises(_tskit.LibraryError):
            tables.simplify([])

    def test_samples_interface(self):
        ts = msprime.simulate(50, random_seed=1)
        for good_form in [[], [0, 1], (0, 1), np.array([0, 1], dtype=np.int32)]:
            tables = ts.dump_tables()
            tables.simplify(good_form)
        tables = ts.dump_tables()
        for bad_values in [[[[]]], np.array([[0, 1], [2, 3]], dtype=np.int32)]:
            with pytest.raises(ValueError):
                tables.simplify(bad_values)
        for bad_type in [[0.1], ["string"], {}, [{}]]:
            with pytest.raises(TypeError):
                tables.simplify(bad_type)
        # We only convert to int if we don't overflow
        for bad_node in [np.iinfo(np.int32).min - 1, np.iinfo(np.int32).max + 1]:
            with pytest.raises(OverflowError):
                tables.simplify(samples=np.array([0, bad_node]))

    @pytest.fixture(scope="session")
    def wf_sim_with_individual_metadata(self):
        tables = wf.wf_sim(
            9,
            10,
            seed=1,
            deep_history=False,
            initial_generation_samples=False,
            num_loci=5,
        )
        assert tables.individuals.num_rows > 50
        assert np.all(tables.nodes.individual >= 0)
        individuals_copy = tables.copy().individuals
        tables.individuals.clear()
        tables.individuals.metadata_schema = tskit.MetadataSchema({"codec": "json"})
        for i, individual in enumerate(individuals_copy):
            tables.individuals.add_row(
                flags=individual.flags,
                location=individual.location,
                parents=individual.parents,
                metadata={
                    "original_id": i,
                    "original_parents": [int(p) for p in individual.parents],
                },
            )
        tables.sort()
        return tables

    def test_individual_parent_mapping(self, wf_sim_with_individual_metadata):
        tables = wf_sim_with_individual_metadata.copy()
        tables.simplify()
        ts = tables.tree_sequence()
        for individual in tables.individuals:
            for parent, original_parent in zip(
                individual.parents, individual.metadata["original_parents"]
            ):
                if parent != tskit.NULL:
                    assert (
                        ts.individual(parent).metadata["original_id"] == original_parent
                    )
        assert set(tables.individuals.parents) != {tskit.NULL}

    def verify_complete_genetic_pedigree(self, tables):
        ts = tables.tree_sequence()
        for edge in ts.edges():
            child = ts.individual(ts.node(edge.child).individual)
            parent = ts.individual(ts.node(edge.parent).individual)
            assert parent.id in child.parents
            assert parent.metadata["original_id"] in child.metadata["original_parents"]

    def test_no_complete_genetic_pedigree(self, wf_sim_with_individual_metadata):
        tables = wf_sim_with_individual_metadata.copy()
        tables.simplify()  # Will remove intermediate individuals
        with pytest.raises(AssertionError):
            self.verify_complete_genetic_pedigree(tables)

    def test_complete_genetic_pedigree(self, wf_sim_with_individual_metadata):
        for params in [{"keep_unary": True}, {"keep_unary_in_individuals": True}]:
            tables = wf_sim_with_individual_metadata.copy()
            tables.simplify(**params)  # Keep intermediate individuals
            self.verify_complete_genetic_pedigree(tables)

    def test_shuffled_individual_parent_mapping(self, wf_sim_with_individual_metadata):
        tables = wf_sim_with_individual_metadata.copy()
        tsutil.shuffle_tables(
            tables,
            42,
            shuffle_edges=False,
            shuffle_populations=False,
            shuffle_individuals=True,
            shuffle_sites=False,
            shuffle_mutations=False,
            shuffle_migrations=False,
        )
        # Check we have a mixed up order
        tables2 = tables.copy()
        tables2.sort_individuals()
        with pytest.raises(AssertionError, match="IndividualTable row 0 differs"):
            tables.assert_equals(tables2)

        tables.simplify()
        metadata = [
            tables.individuals.metadata_schema.decode_row(m)
            for m in tskit.unpack_bytes(
                tables.individuals.metadata, tables.individuals.metadata_offset
            )
        ]
        for individual in tables.individuals:
            for parent, original_parent in zip(
                individual.parents, individual.metadata["original_parents"]
            ):
                if parent != tskit.NULL:
                    assert metadata[parent]["original_id"] == original_parent
        assert set(tables.individuals.parents) != {tskit.NULL}

    def test_individual_mapping(self):
        tables = wf.wf_sim(
            9,
            10,
            seed=1,
            deep_history=True,
            initial_generation_samples=False,
            num_loci=5,
        )
        assert tables.individuals.num_rows > 50
        node_md = []
        individual_md = [b""] * tables.individuals.num_rows
        for i, node in enumerate(tables.nodes):
            node_md.append(struct.pack("i", i))
            individual_md[node.individual] = struct.pack("i", i)
        tables.nodes.packset_metadata(node_md)
        tables.individuals.packset_metadata(individual_md)
        tables.sort()
        tables.simplify()
        ts = tables.tree_sequence()
        for node in tables.nodes:
            if node.individual != tskit.NULL:
                assert ts.individual(node.individual).metadata == node.metadata
        assert set(tables.individuals.parents) != {tskit.NULL}

    def test_bad_individuals(self, simple_degree1_ts_fixture):
        tables = simple_degree1_ts_fixture.dump_tables()
        tables.individuals.clear()
        tables.individuals.add_row(parents=[-2])
        with pytest.raises(tskit.LibraryError, match="Individual out of bounds"):
            tables.simplify()
        tables.individuals.clear()
        tables.individuals.add_row(parents=[0])
        with pytest.raises(
            tskit.LibraryError, match="Individuals cannot be their own parents"
        ):
            tables.simplify()

    def test_unsorted_individuals_ok(self, simple_degree1_ts_fixture):
        tables = simple_degree1_ts_fixture.dump_tables()
        tables.individuals.clear()
        tables.individuals.add_row(parents=[1])
        tables.individuals.add_row(parents=[-1])
        # we really just want to check that no error is thrown here
        tables.simplify()
        assert tables.individuals.num_rows == 0

    def test_filter_none(self, simple_degree1_ts_fixture):
        tables = simple_degree1_ts_fixture.simplify().dump_tables()
        tables.populations.add_row()
        tables.individuals.add_row()
        tables.sites.add_row(
            position=np.nextafter(tables.sequence_length, 0),
            ancestral_state="XXX",
        )
        orig_num_populations = len(tables.populations)
        orig_num_individuals = len(tables.individuals)
        orig_num_sites = len(tables.sites)

        tables.simplify(
            filter_populations=False, filter_individuals=False, filter_sites=False
        )
        assert len(tables.populations) == orig_num_populations
        assert len(tables.individuals) == orig_num_individuals
        assert len(tables.sites) == orig_num_sites

        tables.simplify(
            filter_populations=None, filter_individuals=None, filter_sites=None
        )
        assert len(tables.populations) < orig_num_populations
        assert len(tables.individuals) < orig_num_individuals
        assert len(tables.sites) < orig_num_sites


class TestTableCollection:
    """
    Tests for the convenience wrapper around a collection of related tables.
    """

    def test_table_references(self):
        ts = msprime.simulate(10, mutation_rate=2, random_seed=1)
        tables = ts.tables
        before_individuals = str(tables.individuals)
        individuals = tables.individuals
        before_nodes = str(tables.nodes)
        nodes = tables.nodes
        before_edges = str(tables.edges)
        edges = tables.edges
        before_migrations = str(tables.migrations)
        migrations = tables.migrations
        before_sites = str(tables.sites)
        sites = tables.sites
        before_mutations = str(tables.mutations)
        mutations = tables.mutations
        before_populations = str(tables.populations)
        populations = tables.populations
        before_provenances = str(tables.provenances)
        provenances = tables.provenances
        del tables
        assert str(individuals) == before_individuals
        assert str(nodes) == before_nodes
        assert str(edges) == before_edges
        assert str(migrations) == before_migrations
        assert str(sites) == before_sites
        assert str(mutations) == before_mutations
        assert str(populations) == before_populations
        assert str(provenances) == before_provenances

    def test_str(self):
        ts = msprime.simulate(10, random_seed=1)
        tables = ts.tables
        s = str(tables)
        assert len(s) > 0

    def test_nbytes_empty_tables(self):
        tables = tskit.TableCollection(1)
        assert tables.nbytes == 119

    def test_nbytes(self, tmp_path, ts_fixture):
        tables = ts_fixture.dump_tables()
        tables.dump(tmp_path / "tables")
        store = kastore.load(tmp_path / "tables")
        for v in store.values():
            # Check we really have data in every field
            assert v.nbytes > 0
        nbytes = sum(
            array.nbytes * 2 if "_offset" in name else array.nbytes
            for name, array in store.items()
            # nbytes is the size of asdict, so exclude file format items
            if name not in ["format/version", "format/name", "uuid"]
        )
        assert nbytes == tables.nbytes

    def test_asdict(self, ts_fixture):
        t = ts_fixture.dump_tables()
        d1 = {
            "encoding_version": (1, 6),
            "sequence_length": t.sequence_length,
            "metadata_schema": repr(t.metadata_schema),
            "metadata": t.metadata_schema.encode_row(t.metadata),
            "time_units": t.time_units,
            "individuals": t.individuals.asdict(),
            "populations": t.populations.asdict(),
            "nodes": t.nodes.asdict(),
            "edges": t.edges.asdict(),
            "sites": t.sites.asdict(),
            "mutations": t.mutations.asdict(),
            "migrations": t.migrations.asdict(),
            "provenances": t.provenances.asdict(),
            "indexes": t.indexes.asdict(),
            "reference_sequence": t.reference_sequence.asdict(),
        }
        d2 = t.asdict()
        assert set(d1.keys()) == set(d2.keys())
        t1 = tskit.TableCollection.fromdict(d1)
        t2 = tskit.TableCollection.fromdict(d2)
        t1.assert_equals(t2)
        assert t1.has_index()
        assert t2.has_index()

    @pytest.mark.parametrize("force_offset_64", [True, False])
    def test_asdict_force_offset_64(self, ts_fixture, force_offset_64):
        tables = ts_fixture.dump_tables()
        d = tables.asdict(force_offset_64=force_offset_64)
        for table in tables.table_name_map:
            for name, column in d[table].items():
                if name.endswith("_offset"):
                    if force_offset_64:
                        assert column.dtype == np.uint64
                    else:
                        assert column.dtype == np.uint32

    def test_asdict_force_offset_64_default(self, ts_fixture):
        tables = ts_fixture.dump_tables()
        d = tables.asdict()
        for table in tables.table_name_map:
            for name, column in d[table].items():
                if name.endswith("_offset"):
                    assert column.dtype == np.uint32

    def test_asdict_lifecycle(self, ts_fixture):
        tables = ts_fixture.dump_tables()
        tables_dict = tables.asdict()
        del tables
        tskit.TableCollection.fromdict(tables_dict).assert_equals(
            ts_fixture.dump_tables()
        )

    def test_from_dict(self, ts_fixture):
        t1 = ts_fixture.tables
        d = {
            "encoding_version": (1, 1),
            "sequence_length": t1.sequence_length,
            "metadata_schema": repr(t1.metadata_schema),
            "metadata": t1.metadata_schema.encode_row(t1.metadata),
            "time_units": t1.time_units,
            "individuals": t1.individuals.asdict(),
            "populations": t1.populations.asdict(),
            "nodes": t1.nodes.asdict(),
            "edges": t1.edges.asdict(),
            "sites": t1.sites.asdict(),
            "mutations": t1.mutations.asdict(),
            "migrations": t1.migrations.asdict(),
            "provenances": t1.provenances.asdict(),
            "indexes": t1.indexes.asdict(),
            "reference_sequence": t1.reference_sequence.asdict(),
        }
        t2 = tskit.TableCollection.fromdict(d)
        t1.assert_equals(t2)

    def test_roundtrip_dict(self, ts_fixture):
        t1 = ts_fixture.tables
        t2 = tskit.TableCollection.fromdict(t1.asdict())
        t1.assert_equals(t2)

    def test_table_name_map(self, ts_fixture):
        tables = ts_fixture.tables
        td1 = {
            "individuals": tables.individuals,
            "populations": tables.populations,
            "nodes": tables.nodes,
            "edges": tables.edges,
            "sites": tables.sites,
            "mutations": tables.mutations,
            "migrations": tables.migrations,
            "provenances": tables.provenances,
        }
        td2 = tables.table_name_map
        assert isinstance(td2, dict)
        assert set(td1.keys()) == set(td2.keys())
        for name in td2.keys():
            assert td1[name] == td2[name]
        assert td1 == td2

        # Deprecated in 0.4.1
        with pytest.warns(FutureWarning):
            td1 = tables.name_map
        assert td1 == td2

    def test_equals_empty(self):
        assert tskit.TableCollection() == tskit.TableCollection()

    def test_equals_sequence_length(self):
        assert tskit.TableCollection(sequence_length=1) != tskit.TableCollection(
            sequence_length=2
        )

    def test_copy(self, ts_fixture):
        t1 = ts_fixture.dump_tables()
        t2 = t1.copy()
        assert t1 is not t2
        t1.assert_equals(t2)
        t1.edges.clear()
        assert t1 != t2

    def test_clear_table(self, ts_fixture):
        tables = ts_fixture.dump_tables()
        tables.clear()
        data_tables = [t for t in tskit.TABLE_NAMES if t != "provenances"]
        for table in data_tables:
            assert getattr(tables, f"{table}").num_rows == 0
            assert repr(getattr(tables, f"{table}").metadata_schema) != ""
        assert tables.provenances.num_rows > 0
        assert len(tables.metadata) > 0
        assert repr(tables.metadata_schema) != ""

        tables.clear(clear_provenance=True)
        assert tables.provenances.num_rows == 0
        for table in data_tables:
            assert repr(getattr(tables, f"{table}").metadata_schema) != ""
        assert len(tables.metadata) > 0
        assert repr(tables.metadata_schema) != ""

        tables.clear(clear_metadata_schemas=True)
        for table in data_tables:
            assert repr(getattr(tables, f"{table}").metadata_schema) == ""
        assert len(tables.metadata) > 0
        assert repr(tables.metadata_schema) != 0

        tables.clear(clear_ts_metadata_and_schema=True)
        assert len(tables.metadata) == 0
        assert repr(tables.metadata_schema) == ""

    def test_equals(self):
        # Here we don't use the fixture as we'd like to run the same sim twice
        pop_configs = [msprime.PopulationConfiguration(5) for _ in range(2)]
        migration_matrix = [[0, 1], [1, 0]]
        t1 = msprime.simulate(
            population_configurations=pop_configs,
            migration_matrix=migration_matrix,
            mutation_rate=1,
            record_migrations=True,
            random_seed=1,
        ).dump_tables()
        t2 = msprime.simulate(
            population_configurations=pop_configs,
            migration_matrix=migration_matrix,
            mutation_rate=1,
            record_migrations=True,
            random_seed=1,
        ).dump_tables()
        assert t1 == t1
        assert t1 == t1.copy()
        assert t1.copy() == t1

        # The provenances may or may not be equal depending on the clock
        # precision for record. So clear them first.
        t1.provenances.clear()
        t2.provenances.clear()
        assert t1 == t2
        assert t2 == t1
        assert not (t1 != t2)

        t1.nodes.clear()
        assert t1 != t2
        t2.nodes.clear()
        assert t1 == t2

        t1.edges.clear()
        assert t1 != t2
        t2.edges.clear()
        assert t1 == t2

        t1.migrations.clear()
        assert t1 != t2
        t2.migrations.clear()
        assert t1 == t2

        t1.sites.clear()
        assert t1 != t2
        t2.sites.clear()
        assert t1 == t2

        t1.mutations.clear()
        assert t1 != t2
        t2.mutations.clear()
        assert t1 == t2

        t1.populations.clear()
        assert t1 != t2
        t2.populations.clear()
        assert t1 == t2

    def test_equals_options(self, ts_fixture):
        t1 = ts_fixture.dump_tables()
        t2 = t1.copy()

        t1.provenances.add_row("random stuff")
        assert not (t1 == t2)
        t1.assert_equals(t2, ignore_provenance=True)
        t2.assert_equals(t1, ignore_provenance=True)
        assert not (t1.equals(t2))
        assert not (t2.equals(t1))
        t1.provenances.clear()
        t2.provenances.clear()
        t1.assert_equals(t2)
        t2.assert_equals(t1)

        t1.metadata_schema = tskit.MetadataSchema({"codec": "json", "type": "object"})
        t1.metadata = {"hello": "world"}
        assert not t1.equals(t2)
        t1.assert_equals(t2, ignore_ts_metadata=True)
        assert not t2.equals(t1)
        t2.assert_equals(t1, ignore_ts_metadata=True)
        t2.metadata_schema = t1.metadata_schema
        assert not t1.equals(t2)
        t1.assert_equals(t2, ignore_ts_metadata=True)
        assert not t2.equals(t1)
        t2.assert_equals(t1, ignore_ts_metadata=True)

        t1.provenances.add_row("random stuff")
        assert not t1.equals(t2)
        assert not t1.equals(t2, ignore_ts_metadata=True)
        assert not t1.equals(t2, ignore_provenance=True)
        t1.assert_equals(t2, ignore_ts_metadata=True, ignore_provenance=True)

        t1.provenances.clear()
        t2.metadata = t1.metadata
        t1.assert_equals(t2)
        t2.assert_equals(t1)

        with pytest.raises(TypeError):
            t1.equals(t2, True)

    def test_sequence_length(self):
        for sequence_length in [0, 1, 100.1234]:
            tables = tskit.TableCollection(sequence_length=sequence_length)
            assert tables.sequence_length == sequence_length

    def test_uuid_simulation(self, ts_fixture):
        tables = ts_fixture.tables
        assert tables.file_uuid is None, None

    def test_uuid_empty(self):
        tables = tskit.TableCollection(sequence_length=1)
        assert tables.file_uuid is None, None

    def test_empty_indexes(self):
        tables = tskit.TableCollection(sequence_length=1)
        assert not tables.has_index()
        tables.build_index()
        assert tables.has_index()
        tables.drop_index()
        assert not tables.has_index()

    def test_index_unsorted(self):
        tables = tskit.TableCollection(sequence_length=1)
        tables.nodes.add_row(flags=1, time=0)
        tables.nodes.add_row(flags=1, time=0)
        tables.nodes.add_row(flags=1, time=0)
        tables.nodes.add_row(flags=0, time=1)
        tables.nodes.add_row(flags=0, time=2)
        tables.edges.add_row(0, 1, 3, 0)
        tables.edges.add_row(0, 1, 3, 1)
        tables.edges.add_row(0, 1, 4, 3)
        tables.edges.add_row(0, 1, 4, 2)

        assert not tables.has_index()
        with pytest.raises(tskit.LibraryError):
            tables.build_index()
        assert not tables.has_index()
        tables.sort()
        tables.build_index()
        assert tables.has_index()
        ts = tables.tree_sequence()
        assert ts.tables == tables

    def test_index_from_ts(self, ts_fixture):
        tables = ts_fixture.dump_tables()
        assert tables.has_index()
        tables.drop_index()
        assert not tables.has_index()
        ts = tables.tree_sequence()
        assert ts.tables == tables
        assert tables.has_index()

    def test_set_sequence_length_errors(self):
        tables = tskit.TableCollection(1)
        with pytest.raises(AttributeError):
            del tables.sequence_length
        for bad_value in ["asdf", None, []]:
            with pytest.raises(TypeError):
                tables.sequence_length = bad_value

    def test_set_sequence_length(self):
        tables = tskit.TableCollection(1)
        for value in [-1, 100, 2**32, 1e-6]:
            tables.sequence_length = value
            assert tables.sequence_length == value

    def test_bad_sequence_length(self, ts_fixture):
        tables = ts_fixture.dump_tables()
        assert tables.sequence_length == 5
        for value in [-1, 0, -0.99, 0.9999]:
            tables.sequence_length = value
            with pytest.raises(tskit.LibraryError):
                tables.tree_sequence()
            with pytest.raises(tskit.LibraryError):
                tables.sort()
            with pytest.raises(tskit.LibraryError):
                tables.build_index()
            with pytest.raises(tskit.LibraryError):
                tables.compute_mutation_parents()
            with pytest.raises(tskit.LibraryError):
                tables.simplify()
            assert tables.sequence_length == value

    def test_sequence_length_longer_than_edges(self, ts_fixture):
        tables = ts_fixture.dump_tables()
        tables.sequence_length = 20
        ts = tables.tree_sequence()
        assert ts.sequence_length == 20
        assert ts.num_trees == 6
        trees = ts.trees()
        tree = next(trees)
        assert len(tree.parent_dict) > 0
        for _ in range(5):
            tree = next(trees)
        assert len(tree.parent_dict) == 0

    def test_indexes(self, simple_degree1_ts_fixture):
        tc = tskit.TableCollection(sequence_length=1)
        assert tc.indexes == tskit.TableCollectionIndexes()
        tc = simple_degree1_ts_fixture.tables
        assert np.array_equal(
            tc.indexes.edge_insertion_order, np.arange(18, dtype=np.int32)
        )
        assert np.array_equal(
            tc.indexes.edge_removal_order, np.arange(18, dtype=np.int32)[::-1]
        )
        tc.drop_index()
        assert tc.indexes == tskit.TableCollectionIndexes()
        tc.build_index()
        assert np.array_equal(
            tc.indexes.edge_insertion_order, np.arange(18, dtype=np.int32)
        )
        assert np.array_equal(
            tc.indexes.edge_removal_order, np.arange(18, dtype=np.int32)[::-1]
        )

        modify_indexes = tskit.TableCollectionIndexes(
            edge_insertion_order=np.arange(42, 42 + 18, dtype=np.int32),
            edge_removal_order=np.arange(4242, 4242 + 18, dtype=np.int32),
        )
        tc.indexes = modify_indexes
        assert np.array_equal(
            tc.indexes.edge_insertion_order, np.arange(42, 42 + 18, dtype=np.int32)
        )
        assert np.array_equal(
            tc.indexes.edge_removal_order, np.arange(4242, 4242 + 18, dtype=np.int32)
        )

    def test_indexes_roundtrip(self, simple_degree1_ts_fixture):
        # Indexes shouldn't be made by roundtripping
        tables = tskit.TableCollection(sequence_length=1)
        assert not tables.has_index()
        assert not tskit.TableCollection.fromdict(tables.asdict()).has_index()

        tables = simple_degree1_ts_fixture.dump_tables()
        tables.drop_index()
        assert not tskit.TableCollection.fromdict(tables.asdict()).has_index()

    def test_asdict_lwt_concordance(self, ts_fixture):
        def check_concordance(d1, d2):
            assert set(d1.keys()) == set(d2.keys())
            for k1, v1 in d1.items():
                v2 = d2[k1]
                assert type(v1) is type(v2)
                if type(v1) is dict:
                    assert set(v1.keys()) == set(v2.keys())
                    for sk1, sv1 in v1.items():
                        sv2 = v2[sk1]
                        assert type(sv1) is type(sv2)
                        if isinstance(sv1, np.ndarray):
                            assert np.array_equal(sv1, sv2) or (
                                np.all(tskit.is_unknown_time(sv1))
                                and np.all(tskit.is_unknown_time(sv2))
                            )
                        elif type(sv1) in [bytes, str]:
                            assert sv1 == sv2
                        else:
                            raise AssertionError()

                else:
                    assert v1 == v2

        tables = ts_fixture.dump_tables()
        assert tables.has_index()
        lwt = _tskit.LightweightTableCollection()
        lwt.fromdict(tables.asdict())
        check_concordance(lwt.asdict(), tables.asdict())

        tables.drop_index()
        lwt = _tskit.LightweightTableCollection()
        lwt.fromdict(tables.asdict())
        check_concordance(lwt.asdict(), tables.asdict())

    def test_dump_pathlib(self, ts_fixture, tmp_path):
        path = pathlib.Path(tmp_path) / "tmp.trees"
        assert path.exists
        assert path.is_file
        tc = ts_fixture.dump_tables()
        tc.dump(path)
        other_tc = tskit.TableCollection.load(path)
        tc.assert_equals(other_tc)

    @pytest.mark.skipif(platform.system() == "Windows", reason="Windows doesn't raise")
    def test_dump_load_errors(self, ts_fixture):
        tc = ts_fixture.dump_tables()
        # Try to dump/load files we don't have access to or don't exist.
        for func in [tc.dump, tskit.TableCollection.load]:
            for f in ["/", "/test.trees", "/dir_does_not_exist/x.trees"]:
                with pytest.raises(OSError):
                    func(f)
                try:
                    func(f)
                except OSError as e:
                    message = str(e)
                    assert len(message) > 0
            f = "/" + 4000 * "x"
            with pytest.raises(OSError):
                func(f)
            try:
                func(f)
            except OSError as e:
                message = str(e)
            assert "File name too long" in message
            for bad_filename in [[], None, {}]:
                with pytest.raises(TypeError):
                    func(bad_filename)

    def test_set_table(self):
        tc = tskit.TableCollection()
        for name, table in tc.table_name_map.items():
            with pytest.raises(AttributeError, match="replace_with"):
                setattr(tc, name, table)


class TestEqualityOptions:
    def test_equals_provenance(self):
        t1 = msprime.simulate(10, random_seed=42).tables
        time.sleep(0.1)
        t2 = msprime.simulate(10, random_seed=42).tables
        # Timestamps should differ
        assert t1.provenances[-1].timestamp != t2.provenances[-1].timestamp
        assert not t1.equals(t2)
        t1.assert_equals(t2, ignore_timestamps=True)
        t1.assert_equals(t2, ignore_provenance=True)
        t1.assert_equals(t2, ignore_provenance=True, ignore_timestamps=True)

    def test_equals_node_metadata(self, ts_fixture):
        t1 = ts_fixture.dump_tables()
        t2 = t1.copy()
        t1.assert_equals(t2)
        t1.nodes.add_row(time=0, metadata={"a": "a"})
        t2.nodes.add_row(time=0, metadata={"a": "b"})
        assert not t1.nodes.equals(t2.nodes)
        assert not t1.equals(t2)
        assert t1.nodes.equals(t2.nodes, ignore_metadata=True)

    def test_equals_edge_metadata(self, ts_fixture):
        t1 = ts_fixture.dump_tables()
        child = t1.nodes.add_row(time=0)
        parent = t1.nodes.add_row(time=1)
        t2 = t1.copy()
        t1.assert_equals(t2)
        t1.edges.add_row(0, 1, parent, child, metadata={"a": "a"})
        t2.edges.add_row(0, 1, parent, child, metadata={"a": "b"})
        assert not t1.edges.equals(t2.edges)
        assert not t1.equals(t2)
        assert t1.edges.equals(t2.edges, ignore_metadata=True)
        t1.assert_equals(t2, ignore_metadata=True)

    def test_equals_migration_metadata(self, ts_fixture):
        t1 = ts_fixture.dump_tables()
        t2 = t1.copy()
        t1.assert_equals(t2)
        t1.migrations.add_row(
            0, 1, source=0, dest=1, node=0, time=0, metadata={"a": "a"}
        )
        t2.migrations.add_row(
            0, 1, source=0, dest=1, node=0, time=0, metadata={"a": "b"}
        )
        assert not t1.migrations.equals(t2.migrations)
        assert not t1.equals(t2)
        assert t1.migrations.equals(t2.migrations, ignore_metadata=True)
        t1.assert_equals(t2, ignore_metadata=True)

    def test_equals_site_metadata(self, ts_fixture):
        t1 = ts_fixture.dump_tables()
        t2 = t1.copy()
        t1.assert_equals(t2)
        t1.sites.add_row(0, "A", metadata={"a": "a"})
        t2.sites.add_row(0, "A", metadata={"a": "b"})
        assert not t1.sites.equals(t2.sites)
        assert not t1.equals(t2)
        assert t1.sites.equals(t2.sites, ignore_metadata=True)
        t1.assert_equals(t2, ignore_metadata=True)

    def test_equals_mutation_metadata(self, ts_fixture):
        t1 = ts_fixture.dump_tables()
        t2 = t1.copy()
        t1.assert_equals(t2)
        t1.mutations.add_row(0, 0, "A", metadata={"a": "a"})
        t2.mutations.add_row(0, 0, "A", metadata={"a": "b"})
        assert not t1.mutations.equals(t2.mutations)
        assert not t1.equals(t2)
        assert t1.mutations.equals(t2.mutations, ignore_metadata=True)
        t1.assert_equals(t2, ignore_metadata=True)

    def test_equals_population_metadata(self, ts_fixture):
        t1 = ts_fixture.dump_tables()
        t2 = t1.copy()
        t1.assert_equals(t2)
        t1.populations.add_row({"a": "a"})
        t2.populations.add_row({"a": "b"})
        assert not t1.populations.equals(t2.populations)
        assert not t1.equals(t2)
        t1.assert_equals(t2, ignore_metadata=True)


class TestTableCollectionAssertEquals:
    @pytest.fixture
    def t1(self, ts_fixture):
        return ts_fixture.dump_tables()

    @pytest.fixture
    def t2(self, ts_fixture):
        return ts_fixture.dump_tables()

    def test_equal(self, t1, t2):
        assert t1 is not t2
        t1.assert_equals(t2)

    def test_type(self, t1):
        with pytest.raises(
            AssertionError,
            match=re.escape(
                "Types differ: self=<class 'tskit.tables.TableCollection'> "
                "other=<class 'int'>"
            ),
        ):
            t1.assert_equals(42)

    def test_sequence_length(self, t1, t2):
        t2.sequence_length = 42
        with pytest.raises(
            AssertionError, match="Sequence Length differs: self=5.0 other=42.0"
        ):
            t1.assert_equals(t2)

    def test_metadata_schema(self, t1, t2):
        t2.metadata_schema = tskit.MetadataSchema(None)
        with pytest.raises(
            AssertionError,
            match=re.escape("Metadata schemas differ:"),
        ):
            t1.assert_equals(t2)
        t1.assert_equals(t2, ignore_metadata=True)
        t1.assert_equals(t2, ignore_ts_metadata=True)

    def test_metadata(self, t1, t2):
        t2.metadata = {"foo": "bar"}
        with pytest.raises(
            AssertionError,
            match=re.escape(
                "Metadata differs: self=Test metadata other={'foo': 'bar'}"
            ),
        ):
            t1.assert_equals(t2)
        t1.assert_equals(t2, ignore_metadata=True)
        t1.assert_equals(t2, ignore_ts_metadata=True)

    def test_time_units(self, t1, t2):
        t2.time_units = "microseconds"
        with pytest.raises(
            AssertionError,
            match=re.escape(
                "Time units differs: self=Test time units other=microseconds"
            ),
        ):
            t1.assert_equals(t2)

    @pytest.mark.parametrize("table_name", tskit.TableCollection(1).table_name_map)
    def test_tables(self, t1, t2, table_name):
        table = getattr(t2, table_name)
        table.truncate(0)
        with pytest.raises(
            AssertionError,
            match=f"{type(table).__name__} number of rows differ: "
            f"self={len(getattr(t1, table_name))} other=0",
        ):
            t1.assert_equals(t2)

    @pytest.mark.parametrize("table_name", tskit.TableCollection(1).table_name_map)
    def test_ignore_metadata(self, t1, t2, table_name):
        table = getattr(t2, table_name)
        if hasattr(table, "metadata_schema"):
            table.metadata_schema = tskit.MetadataSchema(None)
            with pytest.raises(
                AssertionError,
                match=re.escape(f"{type(table).__name__} metadata schemas differ:"),
            ):
                t1.assert_equals(t2)
            t1.assert_equals(t2, ignore_metadata=True)

    def test_ignore_provenance(self, t1, t2):
        t2.provenances.truncate(0)
        with pytest.raises(
            AssertionError,
            match="ProvenanceTable number of rows differ: self=5 other=0",
        ):
            t1.assert_equals(t2)
        with pytest.raises(
            AssertionError,
            match="ProvenanceTable number of rows differ: self=5 other=0",
        ):
            t1.assert_equals(t2, ignore_timestamps=True)

        t1.assert_equals(t2, ignore_provenance=True)

    def test_ignore_timestamps(self, t1, t2):
        table = t2.provenances
        timestamp = table.timestamp
        timestamp[0] = ord("F")
        table.set_columns(
            timestamp=timestamp,
            timestamp_offset=table.timestamp_offset,
            record=table.record,
            record_offset=table.record_offset,
        )
        with pytest.raises(
            AssertionError,
            match="ProvenanceTable row 0 differs:\n"
            "self.timestamp=.* other.timestamp=F.*",
        ):
            t1.assert_equals(t2)
        t1.assert_equals(t2, ignore_provenance=True)
        t1.assert_equals(t2, ignore_timestamps=True)

    def test_ignore_tables(self, t1, t2):
        t2.individuals.truncate(0)
        t2.nodes.truncate(0)
        t2.edges.truncate(0)
        t2.migrations.truncate(0)
        t2.sites.truncate(0)
        t2.mutations.truncate(0)
        t2.populations.truncate(0)
        with pytest.raises(
            AssertionError,
            match="EdgeTable number of rows differ: self=390 other=0",
        ):
            t1.assert_equals(t2)
        t1.assert_equals(t2, ignore_tables=True)

    def test_ignore_reference_sequence(self, t1, t2):
        t2.reference_sequence.clear()
        with pytest.raises(
            AssertionError,
            match=re.escape("Metadata schemas differ: "),
        ):
            t1.assert_equals(t2)
        t1.assert_equals(t2, ignore_reference_sequence=True)


class TestTableCollectionMethodSignatures:
    tc = msprime.simulate(10, random_seed=1234).dump_tables()

    def test_kwargs_only(self):
        with pytest.raises(TypeError):
            self.tc.simplify([], True)


class TestTableCollectionMetadata:
    metadata_schema = metadata.MetadataSchema(
        {
            "codec": "json",
            "title": "Example Metadata",
            "type": "object",
            "properties": {
                "one": {"type": "string"},
                "two": {"type": "number"},
                "three": {"type": "array"},
                "four": {"type": "boolean"},
            },
            "required": ["one", "two", "three", "four"],
            "additionalProperties": False,
        },
    )

    def metadata_example_data(self, val=0):
        return {
            "one": "val one",
            "two": val,
            "three": list(range(val, val + 10)),
            "four": True,
        }

    def test_set_metadata_schema(self):
        tc = tskit.TableCollection(1)
        metadata_schema2 = metadata.MetadataSchema({"codec": "json"})
        # Default is no-op metadata codec
        assert repr(tc.metadata_schema) == repr(metadata.MetadataSchema(None))
        # Set
        tc.metadata_schema = self.metadata_schema
        assert repr(tc.metadata_schema) == repr(self.metadata_schema)
        # Overwrite
        tc.metadata_schema = metadata_schema2
        assert repr(tc.metadata_schema) == repr(metadata_schema2)
        # Remove
        tc.metadata_schema = metadata.MetadataSchema(None)
        assert repr(tc.metadata_schema) == repr(metadata.MetadataSchema(None))
        # Set after remove
        tc.metadata_schema = self.metadata_schema
        assert repr(tc.metadata_schema) == repr(self.metadata_schema)
        # Del should fail
        with pytest.raises(AttributeError):
            del tc.metadata_schema
        # None should fail
        with pytest.raises(ValueError):
            tc.metadata_schema = None

    def test_set_metadata(self):
        tc = tskit.TableCollection(1)
        # Default is empty bytes
        assert tc.metadata == b""
        assert tc.metadata_bytes == b""

        tc.metadata_schema = self.metadata_schema
        md1 = self.metadata_example_data()
        md2 = self.metadata_example_data(val=2)
        # Set
        tc.metadata = md1
        assert tc.metadata == md1
        assert tc.metadata_bytes == tskit.canonical_json(md1).encode()
        # Overwrite
        tc.metadata = md2
        assert tc.metadata == md2
        assert tc.metadata_bytes == tskit.canonical_json(md2).encode()
        # Del should fail
        with pytest.raises(AttributeError):
            del tc.metadata
        with pytest.raises(AttributeError):
            del tc.metadata_bytes
        # None should fail
        with pytest.raises(exceptions.MetadataValidationError):
            tc.metadata = None
        # Setting bytes should fail
        with pytest.raises(AttributeError):
            tc.metadata_bytes = b"123"

    def test_set_time_units(self):
        tc = tskit.TableCollection(1)
        assert tc.time_units == tskit.TIME_UNITS_UNKNOWN

        ex1 = "years"
        ex2 = "generations"
        # Set
        tc.time_units = ex1
        assert tc.time_units == ex1
        # Overwrite
        tc.time_units = ex2
        assert tc.time_units == ex2
        # Del should fail
        with pytest.raises(AttributeError):
            del tc.time_units
        # None should fail
        with pytest.raises(TypeError):
            tc.time_units = None

    def test_default_metadata_schema(self):
        # Default should allow bytes
        tc = tskit.TableCollection(1)
        tc.metadata = b"acceptable bytes"
        assert tc.metadata == b"acceptable bytes"
        # Adding non-bytes metadata should error
        with pytest.raises(TypeError):
            tc.metadata = self.metadata_example_data()

    def test_round_trip_metadata(self):
        data = self.metadata_example_data()
        tc = tskit.TableCollection(1)
        tc.metadata_schema = self.metadata_schema
        tc.metadata = data
        assert tc.metadata == data
        assert tc.metadata_bytes == tskit.canonical_json(data).encode()

    def test_bad_metadata(self):
        metadata = self.metadata_example_data()
        metadata["I really shouldn't be here"] = 6
        tc = tskit.TableCollection(1)
        tc.metadata_schema = self.metadata_schema
        with pytest.raises(exceptions.MetadataValidationError):
            tc.metadata = metadata
        assert tc._ll_tables.metadata == b""


def add_table_collection_metadata(tc):
    tc.metadata_schema = tskit.MetadataSchema(
        {
            "codec": "struct",
            "type": "object",
            "properties": {"top-level": {"type": "string", "binaryFormat": "50p"}},
        }
    )
    tc.metadata = {"top-level": "top-level-metadata"}
    for table in tskit.TABLE_NAMES:
        t = getattr(tc, table)
        if hasattr(t, "metadata_schema"):
            t.packset_metadata([f"{table}-{i:10}".encode() for i in range(t.num_rows)])
            t.metadata_schema = tskit.MetadataSchema(
                {
                    "codec": "struct",
                    "type": "object",
                    "properties": {table: {"type": "string", "binaryFormat": "16p"}},
                }
            )


class TestTableCollectionPickle:
    """
    Tests that we can round-trip table collections through pickle.
    """

    def verify(self, tables):
        add_table_collection_metadata(tables)
        other_tables = pickle.loads(pickle.dumps(tables))
        tables.assert_equals(other_tables)

    def test_simple_simulation(self):
        ts = msprime.simulate(2, random_seed=1)
        self.verify(ts.dump_tables())

    def test_simulation_populations(self):
        ts = msprime.simulate(
            population_configurations=[
                msprime.PopulationConfiguration(10),
                msprime.PopulationConfiguration(10),
            ],
            migration_matrix=[[0, 1], [1, 0]],
            record_migrations=True,
            random_seed=1,
        )
        self.verify(ts.dump_tables())

    def test_simulation_sites(self):
        ts = msprime.simulate(12, random_seed=1, mutation_rate=5)
        assert ts.num_sites > 1
        self.verify(ts.dump_tables())

    def test_simulation_individuals(self):
        ts = msprime.simulate(100, random_seed=1)
        ts = tsutil.insert_random_ploidy_individuals(ts, seed=1)
        assert ts.num_individuals > 1
        self.verify(ts.dump_tables())

    def test_empty_tables(self):
        self.verify(tskit.TableCollection())


class TestDeduplicateSites:
    """
    Tests for the TableCollection.deduplicate_sites method.
    """

    def test_empty(self):
        tables = tskit.TableCollection(1)
        tables.deduplicate_sites()
        tables.assert_equals(tskit.TableCollection(1))

    def test_unsorted(self):
        tables = msprime.simulate(10, mutation_rate=1, random_seed=1).dump_tables()
        assert len(tables.sites) > 0
        position = tables.sites.position
        for _ in range(len(position) - 1):
            position = np.roll(position, 1)
            tables.sites.set_columns(
                position=position,
                ancestral_state=tables.sites.ancestral_state,
                ancestral_state_offset=tables.sites.ancestral_state_offset,
            )
            with pytest.raises(_tskit.LibraryError):
                tables.deduplicate_sites()

    def test_bad_position(self):
        for bad_position in [-1, -0.001]:
            tables = tskit.TableCollection()
            tables.sites.add_row(bad_position, "0")
            with pytest.raises(_tskit.LibraryError):
                tables.deduplicate_sites()

    def test_no_effect(self):
        t1 = msprime.simulate(10, mutation_rate=1, random_seed=1).dump_tables()
        t2 = msprime.simulate(10, mutation_rate=1, random_seed=1).dump_tables()
        assert len(t1.sites) > 0
        t1.deduplicate_sites()
        t1.assert_equals(t2, ignore_provenance=True)

    def test_same_sites(self):
        t1 = msprime.simulate(10, mutation_rate=1, random_seed=1).dump_tables()
        t2 = msprime.simulate(10, mutation_rate=1, random_seed=1).dump_tables()
        assert len(t1.sites) > 0
        t1.sites.append_columns(
            position=t1.sites.position,
            ancestral_state=t1.sites.ancestral_state,
            ancestral_state_offset=t1.sites.ancestral_state_offset,
        )
        assert len(t1.sites) == 2 * len(t2.sites)
        t1.sort()
        t1.deduplicate_sites()
        t1.assert_equals(t2, ignore_provenance=True)

    def test_order_maintained(self):
        t1 = tskit.TableCollection(1)
        t1.sites.add_row(position=0, ancestral_state="first")
        t1.sites.add_row(position=0, ancestral_state="second")
        t1.deduplicate_sites()
        assert len(t1.sites) == 1
        assert t1.sites.ancestral_state.tobytes() == b"first"

    def test_multichar_ancestral_state(self):
        ts = msprime.simulate(8, random_seed=3, mutation_rate=1)
        assert ts.num_sites > 2
        tables = ts.dump_tables()
        tables.sites.clear()
        tables.mutations.clear()
        for site in ts.sites():
            site_id = tables.sites.add_row(
                position=site.position, ancestral_state="A" * site.id
            )
            tables.sites.add_row(position=site.position, ancestral_state="0")
            for mutation in site.mutations:
                tables.mutations.add_row(
                    site=site_id, node=mutation.node, derived_state="T" * site.id
                )
        tables.deduplicate_sites()
        new_ts = tables.tree_sequence()
        assert new_ts.num_sites == ts.num_sites
        for site in new_ts.sites():
            assert site.ancestral_state == site.id * "A"

    def test_multichar_metadata(self):
        ts = msprime.simulate(8, random_seed=3, mutation_rate=1)
        assert ts.num_sites > 2
        tables = ts.dump_tables()
        tables.sites.clear()
        tables.mutations.clear()
        for site in ts.sites():
            site_id = tables.sites.add_row(
                position=site.position, ancestral_state="0", metadata=b"A" * site.id
            )
            tables.sites.add_row(position=site.position, ancestral_state="0")
            for mutation in site.mutations:
                tables.mutations.add_row(
                    site=site_id,
                    node=mutation.node,
                    derived_state="1",
                    metadata=b"T" * site.id,
                )
        tables.deduplicate_sites()
        new_ts = tables.tree_sequence()
        assert new_ts.num_sites == ts.num_sites
        for site in new_ts.sites():
            assert site.metadata == site.id * b"A"


class TestBaseTable:
    """
    Tests of the table superclass.
    """

    def test_set_columns_not_implemented(self):
        t = tskit.BaseTable(None, None)
        with pytest.raises(NotImplementedError):
            t.set_columns()

    def test_replace_with(self, ts_fixture):
        # Although replace_with is a BaseTable method, it is simpler to test it
        # on the subclasses directly, as some differ e.g. in having metadata schemas
        original_tables = ts_fixture.dump_tables()
        original_tables.nodes.metadata_schema = tskit.MetadataSchema.permissive_json()
        new_tables = ts_fixture.dump_tables()
        new_tables.clear(clear_provenance=True, clear_metadata_schemas=True)

        # write all the data back in again
        for name, table in new_tables.table_name_map.items():
            new_table = getattr(original_tables, name)
            table.replace_with(new_table)
        new_tables.assert_equals(original_tables)


class TestSubsetTables:
    """
    Tests for the TableCollection.subset method.
    """

    def get_msprime_example(self, sample_size=10, seed=1234):
        M = [[0.0, 0.1], [1.0, 0.0]]
        population_configurations = [
            msprime.PopulationConfiguration(sample_size=sample_size),
            msprime.PopulationConfiguration(sample_size=sample_size),
        ]
        ts = msprime.simulate(
            population_configurations=population_configurations,
            migration_matrix=M,
            length=2e5,
            recombination_rate=1e-8,
            mutation_rate=1e-7,
            record_migrations=False,
            random_seed=seed,
        )
        # adding metadata and locations
        ts = tsutil.add_random_metadata(ts, seed)
        ts = tsutil.insert_random_ploidy_individuals(ts, max_ploidy=1)
        return ts.tables

    def get_wf_example(self, N=5, ngens=2, seed=1249):
        tables = wf.wf_sim(N, N, num_pops=2, seed=seed)
        tables.sort()
        ts = tables.tree_sequence()
        ts = tsutil.jukes_cantor(ts, 1, 10, seed=seed)
        ts = tsutil.add_random_metadata(ts, seed)
        ts = tsutil.insert_random_ploidy_individuals(ts, max_ploidy=2)
        return ts.tables

    def get_examples(self, seed):
        yield self.get_msprime_example(seed=seed)
        yield self.get_wf_example(seed=seed)

    def verify_subset_equality(self, tables, nodes):
        for rp in [True, False]:
            for ru in [True, False]:
                py_sub = tables.copy()
                tsk_sub = tables.copy()
                tsutil.py_subset(
                    py_sub,
                    nodes,
                    record_provenance=False,
                    reorder_populations=rp,
                    remove_unreferenced=ru,
                )
                tsk_sub.subset(
                    nodes,
                    record_provenance=False,
                    reorder_populations=rp,
                    remove_unreferenced=ru,
                )
                py_sub.assert_equals(tsk_sub)

    def verify_subset(self, tables, nodes):
        self.verify_subset_equality(tables, nodes)
        subset = tables.copy()
        subset.subset(nodes, record_provenance=False)
        # adding one so the last element always maps to NULL (-1 -> -1)
        node_map = np.repeat(tskit.NULL, tables.nodes.num_rows + 1)
        indivs = []
        pops = []
        for k, n in enumerate(nodes):
            node_map[n] = k
            ind = tables.nodes[n].individual
            pop = tables.nodes[n].population
            if ind not in indivs and ind != tskit.NULL:
                indivs.append(ind)
            if pop not in pops and pop != tskit.NULL:
                pops.append(pop)
        indivs.sort()  # keep individuals in the same order
        ind_map = np.repeat(tskit.NULL, tables.individuals.num_rows + 1)
        ind_map[indivs] = np.arange(len(indivs), dtype="int32")
        pop_map = np.repeat(tskit.NULL, tables.populations.num_rows + 1)
        pop_map[pops] = np.arange(len(pops), dtype="int32")
        assert subset.nodes.num_rows == len(nodes)
        for k, n in zip(nodes, subset.nodes):
            nn = tables.nodes[k]
            assert nn.time == n.time
            assert nn.flags == n.flags
            assert nn.metadata == n.metadata
            assert ind_map[nn.individual] == n.individual
            assert pop_map[nn.population] == n.population
        assert subset.individuals.num_rows == len(indivs)
        for k, i in zip(indivs, subset.individuals):
            ii = tables.individuals[k]
            assert np.all(np.equal(ii.location, i.location))
            assert ii.metadata == i.metadata
            sub_parents = []
            for p in ii.parents:
                if p == tskit.NULL or ind_map[p] != tskit.NULL:
                    sub_parents.append(ind_map[p])
            assert np.all(np.equal(sub_parents, i.parents))
        assert subset.populations.num_rows == len(pops)
        for k, p in zip(pops, subset.populations):
            pp = tables.populations[k]
            assert pp == p
        # subset can reorder the edges: we need to check we have the same set
        edges = {
            e.replace(parent=node_map[e.parent], child=node_map[e.child])
            for e in tables.edges
            if e.parent in nodes and e.child in nodes
        }
        assert subset.edges.num_rows == len(edges)
        for e in edges:
            assert e in subset.edges
        muts = []
        sites = []
        for k, m in enumerate(tables.mutations):
            if m.node in nodes:
                muts.append(k)
                if m.site not in sites:
                    sites.append(m.site)
        site_map = np.repeat(-1, tables.sites.num_rows)
        site_map[sites] = np.arange(len(sites), dtype="int32")
        mutation_map = np.repeat(tskit.NULL, tables.mutations.num_rows + 1)
        mutation_map[muts] = np.arange(len(muts), dtype="int32")
        assert subset.sites.num_rows == len(sites)
        for k, s in zip(sites, subset.sites):
            ss = tables.sites[k]
            assert ss == s
        assert subset.mutations.num_rows == len(muts)
        for k, m in zip(muts, subset.mutations):
            mm = tables.mutations[k]
            assert mutation_map[mm.parent] == m.parent
            assert site_map[mm.site] == m.site
            assert node_map[mm.node] == m.node
            assert mm.derived_state == m.derived_state
            assert mm.metadata == m.metadata
        assert tables.migrations == subset.migrations
        assert tables.provenances == subset.provenances

    def test_ts_subset(self):
        nodes = np.array([0, 1])
        for tables in self.get_examples(83592):
            ts = tables.tree_sequence()
            tables2 = ts.subset(nodes, record_provenance=False).dump_tables()
            tables.subset(nodes, record_provenance=False)
            tables.assert_equals(tables2)

    def test_subset_all(self):
        # subsetting to everything shouldn't change things except the
        # individual and population ids in the node tables if there are gaps
        for tables in self.get_examples(123583):
            tables2 = tables.copy()
            tables2.subset(np.arange(tables.nodes.num_rows))
            tables.individuals.clear()
            tables2.individuals.clear()
            assert np.all(tables.nodes.time == tables2.nodes.time)
            assert np.all(tables.nodes.flags == tables2.nodes.flags)
            assert np.all(tables.nodes.population == tables2.nodes.population)
            assert np.all(tables.nodes.metadata == tables2.nodes.metadata)
            tables.nodes.clear()
            tables2.nodes.clear()
            tables.assert_equals(tables2, ignore_provenance=True)

    def test_shuffled_tables(self):
        # subset should work on even unsorted tables
        # (tested more thoroughly in TestSortTables)
        for tables in self.get_examples(95521):
            tables2 = tables.copy()
            tsutil.shuffle_tables(tables2, 7000)
            tables2.subset(
                np.arange(tables.nodes.num_rows),
                remove_unreferenced=False,
            )
            assert tables.nodes.num_rows == tables2.nodes.num_rows
            assert tables.individuals.num_rows == tables2.individuals.num_rows
            assert tables.populations.num_rows == tables2.populations.num_rows
            assert tables.edges.num_rows == tables2.edges.num_rows
            assert tables.sites.num_rows == tables2.sites.num_rows
            assert tables.mutations.num_rows == tables2.mutations.num_rows
            tables2 = tables.copy()
            tsutil.shuffle_tables(tables2, 7001)
            tables2.subset([])
            assert tables2.nodes.num_rows == 0
            assert tables2.individuals.num_rows == 0
            assert tables2.populations.num_rows == 0
            assert tables2.edges.num_rows == 0
            assert tables2.sites.num_rows == 0
            assert tables2.mutations.num_rows == 0

    def test_doesnt_reorder_individuals(self):
        tables = wf.wf_sim(N=5, ngens=5, num_pops=2, seed=123)
        tsutil.shuffle_tables(tables, 7000)
        tables2 = tables.copy()
        tables2.subset(np.arange(tables.nodes.num_rows))
        assert tables.individuals == tables2.individuals

    def test_random_subsets(self):
        rng = np.random.default_rng(1542)
        for tables in self.get_examples(9412):
            for n in [2, tables.nodes.num_rows - 10]:
                nodes = rng.choice(np.arange(tables.nodes.num_rows), n, replace=False)
                self.verify_subset(tables, nodes)

    def test_empty_nodes(self):
        for tables in self.get_examples(8724):
            subset = tables.copy()
            subset.subset(np.array([]), record_provenance=False)
            assert subset.nodes.num_rows == 0
            assert subset.edges.num_rows == 0
            assert subset.populations.num_rows == 0
            assert subset.individuals.num_rows == 0
            assert subset.migrations.num_rows == 0
            assert subset.sites.num_rows == 0
            assert subset.mutations.num_rows == 0
            assert subset.provenances == tables.provenances

    def test_no_remove_unreferenced(self):
        tables = tskit.TableCollection(sequence_length=10)
        tables.nodes.add_row(time=0, flags=tskit.NODE_IS_SAMPLE)
        tables.nodes.add_row(time=1)
        tables.edges.add_row(parent=1, child=0, left=0, right=10)
        for k in range(5):
            tables.sites.add_row(position=k, ancestral_state=str(k))
        # these are all unused, so should remain unchanged
        for k in range(5):
            tables.populations.add_row(metadata=str(k).encode())
        for k in range(5):
            tables.individuals.add_row(metadata=str(k).encode())
        sub_tables = tables.copy()
        sub_tables.subset([], remove_unreferenced=False)
        assert tables.sites == sub_tables.sites
        assert tables.populations == sub_tables.populations
        assert tables.individuals == sub_tables.individuals
        ts = tables.tree_sequence()
        sub_tables = ts.subset([], remove_unreferenced=False).tables
        assert tables.sites == sub_tables.sites
        assert tables.populations == sub_tables.populations
        assert tables.individuals == sub_tables.individuals

    def test_subset_reverse_all_nodes(self):
        ts = tskit.Tree.generate_comb(5).tree_sequence
        assert np.all(ts.samples() == np.arange(ts.num_samples))
        tables = ts.dump_tables()
        flipped_ids = np.flip(np.arange(tables.nodes.num_rows))
        self.verify_subset(tables, flipped_ids)
        # Now test the topology is the same
        tables.subset(flipped_ids)
        new_ts = tables.tree_sequence()
        assert set(new_ts.samples()) == set(flipped_ids[np.arange(ts.num_samples)])
        r1 = ts.first().rank()
        r2 = new_ts.first().rank()
        assert r1.shape == r2.shape
        assert r1.label != r2.label

    def test_subset_reverse_internal_nodes(self):
        ts = tskit.Tree.generate_balanced(5).tree_sequence
        internal_nodes = np.ones(ts.num_nodes, dtype=bool)
        internal_nodes[ts.samples()] = False
        tables = ts.dump_tables()
        node_ids = np.arange(tables.nodes.num_rows)
        node_ids[internal_nodes] = np.flip(node_ids[internal_nodes])
        self.verify_subset(tables, node_ids)
        # Now test the topology and the sample labels are the same
        tables.subset(node_ids)
        new_ts = tables.tree_sequence()
        assert np.any(new_ts.nodes_time != ts.nodes_time)
        assert new_ts.first().rank() == ts.first().rank()


class TestUnionTables(unittest.TestCase):
    """
    Tests for the TableCollection.union method
    """

    def get_msprime_example(self, sample_size, T, seed):
        # we assume after the split the ts are completely independent
        M = [[0, 0], [0, 0]]
        population_configurations = [
            msprime.PopulationConfiguration(sample_size=sample_size),
            msprime.PopulationConfiguration(sample_size=sample_size),
        ]
        demographic_events = [
            msprime.CensusEvent(time=T),
            msprime.MassMigration(T, source=1, dest=0, proportion=1),
        ]
        ts = msprime.simulate(
            population_configurations=population_configurations,
            demographic_events=demographic_events,
            migration_matrix=M,
            length=2e5,
            recombination_rate=1e-8,
            mutation_rate=1e-7,
            record_migrations=False,
            random_seed=seed,
        )
        ts = tsutil.add_random_metadata(ts, seed)
        ts = tsutil.insert_random_ploidy_individuals(
            ts, max_ploidy=1, samples_only=True
        )
        return ts

    def get_wf_example(self, N, T, seed):
        twopop_tables = wf.wf_sim(N, T, num_pops=2, seed=seed, deep_history=True)
        twopop_tables.sort()
        ts = twopop_tables.tree_sequence()
        ts = ts.simplify()
        ts = tsutil.jukes_cantor(ts, 1, 10, seed=seed)
        ts = tsutil.add_random_metadata(ts, seed)
        ts = tsutil.insert_random_ploidy_individuals(
            ts, max_ploidy=2, samples_only=True
        )
        return ts

    def split_example(self, ts, T):
        # splitting two pop ts *with no migration* into disjoint ts
        shared_nodes = [n.id for n in ts.nodes() if n.time >= T]
        pop1 = list(ts.samples(population=0))
        pop2 = list(ts.samples(population=1))
        tables1 = ts.simplify(shared_nodes + pop1, record_provenance=False).tables
        tables2 = ts.simplify(shared_nodes + pop2, record_provenance=False).tables
        node_mapping = [
            i if i < len(shared_nodes) else tskit.NULL
            for i in range(tables2.nodes.num_rows)
        ]
        # adding some metadata to one of the tables
        # union should disregard differences in metadata
        tables1.metadata_schema = tskit.MetadataSchema(
            {"codec": "json", "type": "object"}
        )
        tables1.metadata = {"hello": "world"}
        return tables1, tables2, node_mapping

    def verify_union(self, tables, other, node_mapping, add_populations=True):
        self.verify_union_consistency(tables, other, node_mapping)
        self.verify_union_equality(
            tables, other, node_mapping, add_populations=add_populations
        )

    def verify_union_equality(self, tables, other, node_mapping, add_populations=True):
        uni1 = tables.copy()
        uni2 = tables.copy()
        uni1.union(
            other,
            node_mapping,
            record_provenance=False,
            add_populations=add_populations,
        )
        tsutil.py_union(
            uni2,
            other,
            node_mapping,
            record_provenance=False,
            add_populations=add_populations,
        )
        uni1.assert_equals(uni2, ignore_provenance=True)
        # verifying that subsetting to original nodes return the same table
        orig_nodes = [j for i, j in enumerate(node_mapping) if j != tskit.NULL]
        uni1.subset(orig_nodes)
        # subsetting tables just to make sure order is the same
        tables.subset(orig_nodes)
        uni1.assert_equals(tables, ignore_provenance=True)

    def verify_union_consistency(self, tables, other, node_mapping):
        ts1 = tsutil.insert_unique_metadata(tables)
        ts2 = tsutil.insert_unique_metadata(other, offset=1000000)
        tsu = ts1.union(ts2, node_mapping, check_shared_equality=False)
        mapu = tsutil.metadata_map(tsu)
        for j, n1 in enumerate(ts1.nodes()):
            # nodes in ts1 should be preserved, in the same order
            nu = tsu.node(j)
            assert n1.metadata == nu.metadata
            if n1.individual == tskit.NULL:
                assert nu.individual == tskit.NULL
            else:
                assert (
                    ts1.individual(n1.individual).metadata
                    == tsu.individual(nu.individual).metadata
                )
        for j, k in enumerate(node_mapping):
            # nodes in ts2 should match if they are not in node mapping
            if k == tskit.NULL:
                n2 = ts2.node(j)
                md2 = n2.metadata
                assert md2 in mapu["nodes"]
                nu = tsu.node(mapu["nodes"][md2])
                if n2.individual == tskit.NULL:
                    assert nu.individual == tskit.NULL
                else:
                    assert (
                        ts2.individual(n2.individual).metadata
                        == tsu.individual(nu.individual).metadata
                    )
        for e1 in ts1.edges():
            # relationships between nodes in ts1 should be preserved
            p1, c1 = e1.parent, e1.child
            assert e1.metadata in mapu["edges"]
            eu = tsu.edge(mapu["edges"][e1.metadata])
            pu, cu = eu.parent, eu.child
            assert ts1.node(p1).metadata == tsu.node(pu).metadata
            assert ts1.node(c1).metadata == tsu.node(cu).metadata
        for e2 in ts2.edges():
            # relationships between nodes in ts2 should be preserved
            # if both are new nodes
            p2, c2 = e2.parent, e2.child
            if node_mapping[p2] == tskit.NULL and node_mapping[c2] == tskit.NULL:
                assert e2.metadata in mapu["edges"]
                eu = tsu.edge(mapu["edges"][e2.metadata])
                pu, cu = eu.parent, eu.child
                assert ts2.node(p2).metadata == tsu.node(pu).metadata
                assert ts2.node(c2).metadata == tsu.node(cu).metadata

        for i1 in ts1.individuals():
            # individuals in ts1 should be preserved
            assert i1.metadata in mapu["individuals"]
            iu = tsu.individual(mapu["individuals"][i1.metadata])
            assert len(i1.parents) == len(iu.parents)
            for p1, pu in zip(i1.parents, iu.parents):
                if p1 == tskit.NULL:
                    assert pu == tskit.NULL
                else:
                    assert ts1.individual(p1).metadata == tsu.individual(pu).metadata
        # how should individual metadata from ts2 map to ts1
        # and only individuals without shared nodes should be added
        indivs21 = {}
        new_indivs2 = [True for _ in ts2.individuals()]
        for j, k in enumerate(node_mapping):
            n = ts2.node(j)
            if n.individual != tskit.NULL:
                i2 = ts2.individual(n.individual)
                if k == tskit.NULL:
                    indivs21[i2.metadata] = i2.metadata
                else:
                    new_indivs2[n.individual] = False
                    i1 = ts1.individual(ts1.node(k).individual)
                    if i2.metadata in indivs21:
                        assert indivs21[i2.metadata] == i1.metadata
                    else:
                        indivs21[i2.metadata] = i1.metadata
        for i2 in ts2.individuals():
            if new_indivs2[i2.id]:
                assert i2.metadata in mapu["individuals"]
                iu = tsu.individual(mapu["individuals"][i2.metadata])
                assert np.sum(i2.parents == tskit.NULL) == np.sum(
                    iu.parents == tskit.NULL
                )
                md2 = [
                    ts2.individual(p).metadata for p in i2.parents if p != tskit.NULL
                ]
                md2u = [indivs21[md] for md in md2]
                mdu = [
                    tsu.individual(p).metadata for p in iu.parents if p != tskit.NULL
                ]
                assert set(md2u) == set(mdu)
            else:
                # the individual *should* be there, but by a different name
                assert i2.metadata not in mapu["individuals"]
                assert indivs21[i2.metadata] in mapu["individuals"]
        for m1 in ts1.mutations():
            # all mutations in ts1 should be present
            assert m1.metadata in mapu["mutations"]
            mu = tsu.mutation(mapu["mutations"][m1.metadata])
            assert m1.derived_state == mu.derived_state
            assert m1.node == mu.node
            if tskit.is_unknown_time(m1.time):
                assert tskit.is_unknown_time(mu.time)
            else:
                assert m1.time == mu.time
            assert ts1.site(m1.site).position == tsu.site(mu.site).position
        for m2 in ts2.mutations():
            # and those in ts2 if their node has been added
            if node_mapping[m2.node] == tskit.NULL:
                assert m2.metadata in mapu["mutations"]
                mu = tsu.mutation(mapu["mutations"][m2.metadata])
                assert m2.derived_state == mu.derived_state
                assert ts2.node(m2.node).metadata == tsu.node(mu.node).metadata
                if tskit.is_unknown_time(m2.time):
                    assert tskit.is_unknown_time(mu.time)
                else:
                    assert m2.time == mu.time
                assert ts2.site(m2.site).position == tsu.site(mu.site).position
        for s1 in ts1.sites():
            assert s1.metadata in mapu["sites"]
            su = tsu.site(mapu["sites"][s1.metadata])
            assert s1.position == su.position
            assert s1.ancestral_state == su.ancestral_state
        for s2 in ts2.sites():
            if s2.position not in ts1.tables.sites.position:
                assert s2.metadata in mapu["sites"]
                su = tsu.site(mapu["sites"][s2.metadata])
                assert s2.position == su.position
                assert s2.ancestral_state == su.ancestral_state
        # check mutation parents
        tables_union = tsu.tables
        tables_union.compute_mutation_parents()
        assert tables_union.mutations == tsu.tables.mutations

    def test_union_empty(self):
        tables = self.get_msprime_example(sample_size=3, T=2, seed=9328).dump_tables()
        tables.sort()
        empty_tables = tables.copy()
        for table in empty_tables.table_name_map.keys():
            getattr(empty_tables, table).clear()
        uni = tables.copy()
        uni.union(empty_tables, [])
        tables.assert_equals(uni, ignore_provenance=True)

    def test_noshared_example(self):
        ts1 = self.get_msprime_example(sample_size=3, T=2, seed=9328)
        ts2 = self.get_msprime_example(sample_size=3, T=2, seed=2125)
        node_mapping = np.full(ts2.num_nodes, tskit.NULL, dtype="int32")
        uni1 = ts1.union(ts2, node_mapping, record_provenance=False)
        uni2_tables = ts1.dump_tables()
        tsutil.py_union(uni2_tables, ts2.tables, node_mapping, record_provenance=False)
        assert uni1.tables == uni2_tables

    def test_all_shared_example(self):
        tables = self.get_wf_example(N=5, T=5, seed=11349).dump_tables()
        tables.sort()
        uni = tables.copy()
        node_mapping = np.arange(tables.nodes.num_rows)
        uni.union(tables, node_mapping, record_provenance=False)
        uni.assert_equals(tables)

    def test_no_add_pop(self):
        self.verify_union(
            *self.split_example(self.get_msprime_example(10, 10, seed=135), 10),
            add_populations=False,
        )
        self.verify_union(
            *self.split_example(self.get_wf_example(10, 10, seed=157), 10),
            add_populations=False,
        )

    def test_provenance(self):
        tables, other, node_mapping = self.split_example(
            self.get_msprime_example(5, T=2, seed=928), 2
        )
        tables_copy = tables.copy()
        tables.union(other, node_mapping)
        uni_other_dict = json.loads(tables.provenances[-1].record)["parameters"][
            "other"
        ]
        recovered_prov_table = tskit.ProvenanceTable()
        assert len(uni_other_dict["timestamp"]) == len(uni_other_dict["record"])
        for timestamp, record in zip(
            uni_other_dict["timestamp"], uni_other_dict["record"]
        ):
            recovered_prov_table.add_row(record, timestamp)
        assert recovered_prov_table == other.provenances
        tables.provenances.truncate(tables.provenances.num_rows - 1)
        assert tables.provenances == tables_copy.provenances

    def test_examples(self):
        for N in [2, 4, 5]:
            for T in [2, 5, 20]:
                for mut_times in [True, False]:
                    with self.subTest(N=N, T=T):
                        ts = self.get_msprime_example(N, T=T, seed=888)
                        if mut_times:
                            tables = ts.tables
                            tables.compute_mutation_times()
                            ts = tables.tree_sequence()
                        self.verify_union(*self.split_example(ts, T))
                        ts = self.get_wf_example(N=N, T=T, seed=827)
                        if mut_times:
                            tables = ts.tables
                            tables.compute_mutation_times()
                            ts = tables.tree_sequence()
                        self.verify_union(*self.split_example(ts, T))


class TestTableSetitemMetadata:
    @pytest.mark.parametrize("table_name", tskit.TABLE_NAMES)
    def test_setitem_metadata(self, ts_fixture, table_name):
        table = getattr(ts_fixture.tables, table_name)
        if hasattr(table, "metadata_schema"):
            assert table.metadata_schema == tskit.MetadataSchema({"codec": "json"})
            assert table[0].metadata != table[1].metadata
            table[0] = table[1]
            assert table[0] == table[1]


def keep_rows_definition(table, keep):
    id_map = np.full(len(table), -1, np.int32)
    copy = table.copy()
    table.clear()
    for j, row in enumerate(copy):
        if keep[j]:
            id_map[j] = len(table)
            table.append(row)
    return id_map


class KeepRowsBaseTest:
    # Simple tests assuming that rows aren't self-referential

    def test_keep_all(self, ts_fixture):
        table = self.get_table(ts_fixture)
        before = table.copy()
        table.keep_rows(np.ones(len(table), dtype=bool))
        assert table.equals(before)

    def test_keep_none(self, ts_fixture):
        table = self.get_table(ts_fixture)
        table.keep_rows(np.zeros(len(table), dtype=bool))
        assert len(table) == 0

    def check_keep_rows(self, table, keep):
        copy = table.copy()
        id_map1 = keep_rows_definition(copy, keep)
        id_map2 = table.keep_rows(keep)
        table.assert_equals(copy)
        np.testing.assert_array_equal(id_map1, id_map2)

    def test_keep_even(self, ts_fixture):
        table = self.get_table(ts_fixture)
        keep = np.ones(len(table), dtype=bool)
        keep[1::2] = 0
        self.check_keep_rows(table, keep)

    def test_keep_odd(self, ts_fixture):
        table = self.get_table(ts_fixture)
        keep = np.ones(len(table), dtype=bool)
        keep[::2] = 0
        self.check_keep_rows(table, keep)

    def test_keep_first(self, ts_fixture):
        table = self.get_table(ts_fixture)
        keep = np.zeros(len(table), dtype=bool)
        keep[0] = 1
        self.check_keep_rows(table, keep)
        assert len(table) == 1

    def test_keep_last(self, ts_fixture):
        table = self.get_table(ts_fixture)
        keep = np.zeros(len(table), dtype=bool)
        keep[-1] = 1
        self.check_keep_rows(table, keep)
        assert len(table) == 1

    @pytest.mark.parametrize("dtype", [np.int32, int, np.float32])
    def test_bad_array_dtype(self, ts_fixture, dtype):
        table = self.get_table(ts_fixture)
        keep = np.zeros(len(table), dtype=dtype)
        with pytest.raises(TypeError, match="Cannot cast array"):
            table.keep_rows(keep)

    @pytest.mark.parametrize("truthy", [False, 0, "", None])
    def test_python_falsey_input(self, ts_fixture, truthy):
        table = self.get_table(ts_fixture)
        keep = [truthy] * len(table)
        self.check_keep_rows(table, keep)
        assert len(table) == 0

    @pytest.mark.parametrize("truthy", [True, 1, "string", 1e-6])
    def test_python_truey_input(self, ts_fixture, truthy):
        table = self.get_table(ts_fixture)
        n = len(table)
        keep = [truthy] * len(table)
        self.check_keep_rows(table, keep)
        assert len(table) == n

    @pytest.mark.parametrize("offset", [-1, 1, 100])
    def test_bad_length(self, ts_fixture, offset):
        table = self.get_table(ts_fixture)
        keep = [True] * (len(table) + offset)
        match_str = f"need:{len(table)}, got:{len(table) + offset}"
        with pytest.raises(ValueError, match=match_str):
            table.keep_rows(keep)

    @pytest.mark.parametrize("bad_type", [False, 0, None])
    def test_non_list_input(self, ts_fixture, bad_type):
        table = self.get_table(ts_fixture)
        with pytest.raises(TypeError, match="has no len"):
            table.keep_rows(bad_type)


class TestNodeTableKeepRows(KeepRowsBaseTest):
    def get_table(self, ts):
        return ts.dump_tables().nodes


class TestEdgeTableKeepRows(KeepRowsBaseTest):
    def get_table(self, ts):
        return ts.dump_tables().edges


class TestSiteTableKeepRows(KeepRowsBaseTest):
    def get_table(self, ts):
        return ts.dump_tables().sites


class TestMigrationTableKeepRows(KeepRowsBaseTest):
    def get_table(self, ts):
        return ts.dump_tables().migrations


class TestPopulationTableKeepRows(KeepRowsBaseTest):
    def get_table(self, ts):
        return ts.dump_tables().populations


class TestProvenanceTableKeepRows(KeepRowsBaseTest):
    def get_table(self, ts):
        return ts.dump_tables().provenances


# Null out the self-referential columns (this is why the tests are structed via
# classes rather than pytest parametrize.


class TestIndividualTableKeepRows(KeepRowsBaseTest):
    def get_table(self, ts):
        table = ts.dump_tables().individuals
        table.parents = np.zeros_like(table.parents) - 1
        return table

    def check_keep_rows(self, table, keep):
        copy = table.copy()
        id_map1 = keep_rows_definition(copy, keep)
        for j, row in enumerate(copy):
            parents = [p if p == tskit.NULL else id_map1[p] for p in row.parents]
            copy[j] = row.replace(parents=parents)
        id_map2 = table.keep_rows(keep)
        table.assert_equals(copy)
        np.testing.assert_array_equal(id_map1, id_map2)

    def test_delete_unreferenced(self, ts_fixture):
        table = ts_fixture.dump_tables().individuals
        ref_count = np.zeros(len(table))
        for row in table:
            for parent in row.parents:
                ref_count[parent] += 1
        self.check_keep_rows(table, ref_count > 0)


class TestMutationTableKeepRows(KeepRowsBaseTest):
    def get_table(self, ts):
        table = ts.dump_tables().mutations
        table.parent = np.zeros_like(table.parent) - 1
        return table

    def check_keep_rows(self, table, keep):
        copy = table.copy()
        id_map1 = keep_rows_definition(copy, keep)
        for j, row in enumerate(copy):
            if row.parent != tskit.NULL:
                copy[j] = row.replace(parent=id_map1[row.parent])
        id_map2 = table.keep_rows(keep)
        table.assert_equals(copy)
        np.testing.assert_array_equal(id_map1, id_map2)

    def test_delete_unreferenced(self, ts_fixture):
        table = ts_fixture.dump_tables().mutations
        parent = table.parent.copy()
        parent[parent == tskit.NULL] = len(table)
        references = np.bincount(parent)
        self.check_keep_rows(table, references[:-1] > 0)

    def test_error_on_bad_ids(self, ts_fixture):
        table = ts_fixture.dump_tables().mutations
        table.add_row(site=0, node=0, derived_state="A", parent=10000)
        before = table.copy()
        with pytest.raises(tskit.LibraryError, match="TSK_ERR_MUTATION_OUT_OF_BOUNDS"):
            table.keep_rows(np.ones(len(table), dtype=bool))
        table.assert_equals(before)


class TestKeepRowsExamples:
    """
    Some examples of how to use the keep_rows method in an idiomatic
    and efficient way.

    TODO these should be converted into documentation examples when we
    write an "examples" section for table editing.
    """

    def test_detach_subtree(self):
        # 2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊
        # 1.00┊ ┃  3  ┊
        #     ┊ ┃ ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        ts = tskit.Tree.generate_balanced(3).tree_sequence
        tables = ts.dump_tables()
        tables.edges.keep_rows(tables.edges.child != 3)

        # 2.00┊ 4     ┊
        #     ┊ ┃     ┊
        # 1.00┊ ┃  3  ┊
        #     ┊ ┃ ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        ts = tables.tree_sequence()
        assert ts.num_trees == 1
        assert ts.first().parent_dict == {0: 4, 1: 3, 2: 3}

    def test_delete_older_edges(self):
        # 2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊
        # 1.00┊ ┃  3  ┊
        #     ┊ ┃ ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        ts = tskit.Tree.generate_balanced(3).tree_sequence
        tables = ts.dump_tables()
        tables.edges.keep_rows(tables.nodes.time[tables.edges.parent] <= 1)

        # 2.00┊       ┊
        #     ┊       ┊
        # 1.00┊    3  ┊
        #     ┊   ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        ts = tables.tree_sequence()
        assert ts.num_trees == 1
        assert ts.first().parent_dict == {1: 3, 2: 3}

    def test_delete_unreferenced_nodes(self):
        # 2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊
        # 1.00┊ ┃  3  ┊
        #     ┊ ┃ ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        ts = tskit.Tree.generate_balanced(3).tree_sequence
        tables = ts.dump_tables()
        edges = tables.edges
        nodes = tables.nodes
        edges.keep_rows(nodes.time[edges.parent] <= 1)
        # 2.00┊       ┊
        #     ┊       ┊
        # 1.00┊    3  ┊
        #     ┊   ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        ref_count = np.bincount(edges.child, minlength=len(nodes))
        ref_count += np.bincount(edges.parent, minlength=len(nodes))
        assert list(ref_count) == [0, 1, 1, 2, 0]
        id_map = nodes.keep_rows(ref_count > 0)
        assert list(id_map) == [-1, 0, 1, 2, -1]
        assert len(nodes) == 3
        # Remap the edges IDs
        edges.child = id_map[edges.child]
        edges.parent = id_map[edges.parent]
        ts = tables.tree_sequence()
        assert ts.num_trees == 1
        assert ts.first().parent_dict == {0: 2, 1: 2}

    def test_mutation_ids_auto_remapped(self):
        mutations = tskit.MutationTable()
        # Add 5 initial rows with no parents
        for j in range(5):
            mutations.add_row(site=j, node=j, derived_state=f"{j}")
        # Now 5 more in a chain
        last = -1
        for j in range(5):
            last = mutations.add_row(
                site=10 + j, node=10 + j, parent=last, derived_state=f"{j}"
            )

        # ╔══╤════╤════╤════╤═════════════╤══════╤════════╗
        # ║id│site│node│time│derived_state│parent│metadata║
        # ╠══╪════╪════╪════╪═════════════╪══════╪════════╣
        # ║0 │   0│   0│ nan│            0│    -1│        ║
        # ║1 │   1│   1│ nan│            1│    -1│        ║
        # ║2 │   2│   2│ nan│            2│    -1│        ║
        # ║3 │   3│   3│ nan│            3│    -1│        ║
        # ║4 │   4│   4│ nan│            4│    -1│        ║
        # ║5 │  10│  10│ nan│            0│    -1│        ║
        # ║6 │  11│  11│ nan│            1│     5│        ║
        # ║7 │  12│  12│ nan│            2│     6│        ║
        # ║8 │  13│  13│ nan│            3│     7│        ║
        # ║9 │  14│  14│ nan│            4│     8│        ║
        # ╚══╧════╧════╧════╧═════════════╧══════╧════════╝

        keep = np.ones(len(mutations), dtype=bool)
        keep[:5] = False
        mutations.keep_rows(keep)

        # ╔══╤════╤════╤════╤═════════════╤══════╤════════╗
        # ║id│site│node│time│derived_state│parent│metadata║
        # ╠══╪════╪════╪════╪═════════════╪══════╪════════╣
        # ║0 │  10│  10│ nan│            0│    -1│        ║
        # ║1 │  11│  11│ nan│            1│     0│        ║
        # ║2 │  12│  12│ nan│            2│     1│        ║
        # ║3 │  13│  13│ nan│            3│     2│        ║
        # ║4 │  14│  14│ nan│            4│     3│        ║
        # ╚══╧════╧════╧════╧═════════════╧══════╧════════╝
        assert list(mutations.site) == [10, 11, 12, 13, 14]
        assert list(mutations.node) == [10, 11, 12, 13, 14]
        assert list(mutations.parent) == [-1, 0, 1, 2, 3]

    def test_individual_ids_auto_remapped(self):
        individuals = tskit.IndividualTable()
        # Add some rows with missing parents in different forms
        individuals.add_row()
        individuals.add_row(parents=[-1])
        individuals.add_row(parents=[-1, -1])
        # Now 5 more in a chain
        last = -1
        for _ in range(5):
            last = individuals.add_row(parents=[last])
        last = individuals.add_row(parents=[last, last])

        # ╔══╤═════╤════════╤═══════╤════════╗
        # ║id│flags│location│parents│metadata║
        # ╠══╪═════╪════════╪═══════╪════════╣
        # ║0 │    0│        │       │        ║
        # ║1 │    0│        │     -1│        ║
        # ║2 │    0│        │ -1, -1│        ║
        # ║3 │    0│        │     -1│        ║
        # ║4 │    0│        │      3│        ║
        # ║5 │    0│        │      4│        ║
        # ║6 │    0│        │      5│        ║
        # ║7 │    0│        │      6│        ║
        # ║8 │    0│        │   7, 7│        ║
        # ╚══╧═════╧════════╧═══════╧════════╝

        keep = np.ones(len(individuals), dtype=bool)
        # Only delete one row
        keep[1] = False
        individuals.keep_rows(keep)

        # ╔══╤═════╤════════╤═══════╤════════╗
        # ║id│flags│location│parents│metadata║
        # ╠══╪═════╪════════╪═══════╪════════╣
        # ║0 │    0│        │       │        ║
        # ║1 │    0│        │ -1, -1│        ║
        # ║2 │    0│        │     -1│        ║
        # ║3 │    0│        │      2│        ║
        # ║4 │    0│        │      3│        ║
        # ║5 │    0│        │      4│        ║
        # ║6 │    0│        │      5│        ║
        # ║7 │    0│        │   6, 6│        ║
        # ╚══╧═════╧════════╧═══════╧════════╝
        parents = [list(ind.parents) for ind in individuals]
        assert parents == [[], [-1, -1], [-1], [2], [3], [4], [5], [6, 6]]


--- ../../tskit/python/tests/test_stats.py ---


"""
Test cases for stats calculations in tskit.
"""
import contextlib
import io

import msprime
import numpy as np
import pytest

import _tskit
import tests
import tests.test_wright_fisher as wf
import tests.tsutil as tsutil
import tskit


@contextlib.contextmanager
def suppress_division_by_zero_warning():
    with np.errstate(invalid="ignore", divide="ignore"):
        yield


def get_r2_matrix(ts):
    """
    Simple site-based version assuming biallic sites.
    """
    A = np.zeros((ts.num_sites, ts.num_sites))
    G = ts.genotype_matrix()
    n = ts.num_samples
    for a in range(ts.num_sites):
        A[a, a] = 1
        fA = np.sum(G[a] != 0) / n
        for b in range(a + 1, ts.num_sites):
            fB = np.sum(G[b] != 0) / n
            nAB = np.sum(np.logical_and(G[a] != 0, G[b] != 0))
            fAB = nAB / n
            D = fAB - fA * fB
            denom = fA * fB * (1 - fA) * (1 - fB)
            A[a, b] = D * D
            with suppress_division_by_zero_warning():
                A[a, b] /= denom
            A[b, a] = A[a, b]
    return A


def _compute_r2(tree, n, f_a, site_b):
    assert len(site_b.mutations) == 1
    assert site_b.ancestral_state != site_b.mutations[0].derived_state
    f_b = tree.num_samples(site_b.mutations[0].node) / n
    f_ab = tree.num_tracked_samples(site_b.mutations[0].node) / n
    D2 = (f_ab - f_a * f_b) ** 2
    denom = f_a * f_b * (1 - f_a) * (1 - f_b)
    if denom == 0:
        return np.nan
    return D2 / denom


def ts_r2(ts, a, b):
    """
    Returns the r2 value between sites a and b in the specified tree sequence.
    """
    a, b = (a, b) if a < b else (b, a)
    site_a = ts.site(a)
    site_b = ts.site(b)
    assert len(site_a.mutations) == 1
    assert len(site_b.mutations) == 1
    n = ts.num_samples
    tree = ts.at(site_a.position)
    a_samples = list(tree.samples(site_a.mutations[0].node))
    f_a = len(a_samples) / n
    tree = ts.at(site_b.position, tracked_samples=a_samples)
    return _compute_r2(tree, n, f_a, site_b)


class LdArrayCalculator:
    """
    Utility class to help organise the state required when tracking all
    the different termination conditions.
    """

    def __init__(self, ts, focal_site_id, direction, max_sites, max_distance):
        self.ts = ts
        self.focal_site = ts.site(focal_site_id)
        self.direction = direction
        self.max_sites = max_sites
        self.max_distance = max_distance
        self.result = []
        self.tree = None

    def _check_site(self, site):
        assert len(site.mutations) == 1
        assert site.ancestral_state != site.mutations[0].derived_state

    def _compute_and_append(self, target_site):
        self._check_site(target_site)

        distance = abs(target_site.position - self.focal_site.position)
        if distance > self.max_distance or len(self.result) >= self.max_sites:
            return True
        r2 = _compute_r2(
            self.tree, self.ts.num_samples, self.focal_frequency, target_site
        )
        self.result.append(r2)
        return False

    def _compute_forward(self):
        done = False
        for site in self.tree.sites():
            if site.id > self.focal_site.id:
                done = self._compute_and_append(site)
                if done:
                    break
        while self.tree.next() and not done:
            for site in self.tree.sites():
                done = self._compute_and_append(site)
                if done:
                    break

    def _compute_backward(self):
        done = False
        for site in reversed(list(self.tree.sites())):
            if site.id < self.focal_site.id:
                done = self._compute_and_append(site)
                if done:
                    break
        while self.tree.prev() and not done:
            for site in reversed(list(self.tree.sites())):
                done = self._compute_and_append(site)
                if done:
                    break

    def run(self):
        self._check_site(self.focal_site)
        self.tree = self.ts.at(self.focal_site.position)
        a_samples = list(self.tree.samples(self.focal_site.mutations[0].node))
        self.focal_frequency = len(a_samples) / self.ts.num_samples

        # Now set the tracked samples on the tree. We don't have a python
        # API for doing this, so we just create a new tree.
        self.tree = self.ts.at(self.focal_site.position, tracked_samples=a_samples)
        if self.direction == 1:
            self._compute_forward()
        else:
            self._compute_backward()
        return np.array(self.result)


def ts_r2_array(ts, a, *, direction=1, max_sites=None, max_distance=None):
    max_sites = ts.num_sites if max_sites is None else max_sites
    max_distance = np.inf if max_distance is None else max_distance
    calc = LdArrayCalculator(ts, a, direction, max_sites, max_distance)
    return calc.run()


class TestLdSingleTree:
    # 2.00┊   4   ┊
    #     ┊ ┏━┻┓  ┊
    # 1.00┊ ┃  3  ┊
    #     ┊ ┃ ┏┻┓ ┊
    # 0.00┊ 0 1 2 ┊
    #     0      10
    #      | |  |
    #  pos 2 4  9
    # node 0 1  0
    @tests.cached_example
    def ts(self):
        ts = tskit.Tree.generate_balanced(3, span=10).tree_sequence
        tables = ts.dump_tables()
        tables.sites.add_row(2, ancestral_state="A")
        tables.sites.add_row(4, ancestral_state="A")
        tables.sites.add_row(9, ancestral_state="T")
        tables.mutations.add_row(site=0, node=0, derived_state="G")
        tables.mutations.add_row(site=1, node=3, derived_state="C")
        tables.mutations.add_row(site=2, node=0, derived_state="G")
        return tables.tree_sequence()

    @pytest.mark.parametrize(["a", "b", "expected"], [(0, 0, 1), (0, 1, 1), (0, 2, 1)])
    def test_r2(self, a, b, expected):
        ts = self.ts()
        A = get_r2_matrix(ts)
        ldc = tskit.LdCalculator(ts)
        assert ldc.r2(a, b) == pytest.approx(expected)
        assert ts_r2(ts, a, b) == pytest.approx(expected)
        assert A[a, b] == pytest.approx(expected)
        assert ldc.r2(b, a) == pytest.approx(expected)
        assert ts_r2(ts, b, a) == pytest.approx(expected)
        assert A[b, a] == pytest.approx(expected)

    @pytest.mark.parametrize("a", [0, 1, 2])
    @pytest.mark.parametrize("direction", [1, -1])
    def test_r2_array(self, a, direction):
        ts = self.ts()
        ldc = tskit.LdCalculator(ts)
        lib_a = ldc.r2_array(a, direction=direction)
        py_a = ts_r2_array(ts, a, direction=direction)
        np.testing.assert_array_almost_equal(lib_a, py_a)


class TestLdFixedSites:
    # 2.00┊   4   ┊
    #     ┊ ┏━┻┓  ┊
    # 1.00┊ ┃  3  ┊
    #     ┊ ┃ ┏┻┓ ┊
    # 0.00┊ 0 1 2 ┊
    #     0      10
    #      | |  |
    #  pos 2 4  9
    # node 0 1  0
    @tests.cached_example
    def ts(self):
        ts = tskit.Tree.generate_balanced(3, span=10).tree_sequence
        tables = ts.dump_tables()
        # First and last mutations are over the root
        tables.sites.add_row(2, ancestral_state="A")
        tables.sites.add_row(4, ancestral_state="A")
        tables.sites.add_row(9, ancestral_state="T")
        tables.mutations.add_row(site=0, node=4, derived_state="G")
        tables.mutations.add_row(site=1, node=3, derived_state="C")
        tables.mutations.add_row(site=2, node=4, derived_state="G")
        return tables.tree_sequence()

    def test_r2_fixed_fixed(self):
        ts = self.ts()
        A = get_r2_matrix(ts)
        ldc = tskit.LdCalculator(ts)
        assert np.isnan(ldc.r2(0, 2))
        assert np.isnan(ts_r2(ts, 0, 2))
        assert np.isnan(A[0, 2])

    def test_r2_fixed_non_fixed(self):
        ts = self.ts()
        A = get_r2_matrix(ts)
        ldc = tskit.LdCalculator(ts)
        assert np.isnan(ldc.r2(0, 1))
        assert np.isnan(ts_r2(ts, 0, 1))
        assert np.isnan(A[0, 1])

    def test_r2_non_fixed_fixed(self):
        ts = self.ts()
        A = get_r2_matrix(ts)
        ldc = tskit.LdCalculator(ts)
        assert np.isnan(ldc.r2(1, 0))
        assert np.isnan(ts_r2(ts, 1, 0))
        assert np.isnan(A[1, 0])


class BaseTestLd:
    """
    Define a set of tests for LD calculations. Subclasses should be
    concrete examples with at least two sites which implement a
    method ts() which returns the tree sequence and the full LD
    matrix.
    """

    def test_r2_all_pairs(self):
        ts, A = self.ts()
        ldc = tskit.LdCalculator(ts)
        for j in range(ts.num_sites):
            for k in range(ts.num_sites):
                r2 = A[j, k]
                assert ldc.r2(j, k) == pytest.approx(r2)
                assert ts_r2(ts, j, k) == pytest.approx(r2)

    def test_r2_array_first_site_forward(self):
        ts, A = self.ts()
        ldc = tskit.LdCalculator(ts)
        A1 = ldc.r2_array(0, direction=1)
        A2 = ts_r2_array(ts, 0, direction=1)
        np.testing.assert_array_almost_equal(A2, A[0, 1:])
        np.testing.assert_array_almost_equal(A1, A2)

    def test_r2_array_mid_forward(self):
        ts, A = self.ts()
        ldc = tskit.LdCalculator(ts)
        site = ts.num_sites // 2
        A1 = ldc.r2_array(site, direction=1)
        A2 = ts_r2_array(ts, site, direction=1)
        np.testing.assert_array_almost_equal(A2, A[site, site + 1 :])
        np.testing.assert_array_almost_equal(A1, A2)

    def test_r2_array_first_site_forward_max_sites(self):
        ts, A = self.ts()
        ldc = tskit.LdCalculator(ts)
        A1 = ldc.r2_array(0, direction=1, max_sites=2)
        A2 = ts_r2_array(ts, 0, direction=1, max_sites=2)
        np.testing.assert_array_almost_equal(A2, A[0, 1:3])
        np.testing.assert_array_almost_equal(A1, A2)

    def test_r2_array_first_site_forward_max_distance(self):
        ts, _ = self.ts()
        ldc = tskit.LdCalculator(ts)
        A1 = ldc.r2_array(0, direction=1, max_distance=3)
        A2 = ts_r2_array(ts, 0, direction=1, max_distance=3)
        np.testing.assert_array_almost_equal(A1, A2)

    def test_r2_array_last_site_backward(self):
        ts, A = self.ts()
        ldc = tskit.LdCalculator(ts)
        a = ts.num_sites - 1
        A1 = ldc.r2_array(a, direction=-1)
        A2 = ts_r2_array(ts, a, direction=-1)
        np.testing.assert_array_almost_equal(A2, A[-1, :-1][::-1])
        np.testing.assert_array_almost_equal(A1, A2)

    def test_r2_array_mid_backward(self):
        ts, A = self.ts()
        ldc = tskit.LdCalculator(ts)
        site = ts.num_sites // 2
        A1 = ldc.r2_array(site, direction=-1)
        A2 = ts_r2_array(ts, site, direction=-1)
        np.testing.assert_array_almost_equal(A2, A[site, :site][::-1])
        np.testing.assert_array_almost_equal(A1, A2)

    def test_r2_array_last_site_backward_max_sites(self):
        ts, A = self.ts()
        ldc = tskit.LdCalculator(ts)
        a = ts.num_sites - 1
        A1 = ldc.r2_array(a, direction=-1, max_sites=2)
        A2 = ts_r2_array(ts, a, direction=-1, max_sites=2)
        np.testing.assert_array_almost_equal(A2, A[-1, -3:-1][::-1])
        np.testing.assert_array_almost_equal(A1, A2)

    def test_r2_array_last_site_backward_max_distance(self):
        ts, _ = self.ts()
        ldc = tskit.LdCalculator(ts)
        a = ts.num_sites - 1
        A1 = ldc.r2_array(a, direction=-1, max_distance=3)
        A2 = ts_r2_array(ts, a, direction=-1, max_distance=3)
        np.testing.assert_array_almost_equal(A1, A2)

    @pytest.mark.parametrize("max_sites", [0, 1, 2])
    def test_r2_array_forward_max_sites_zero(self, max_sites):
        ts, A = self.ts()
        ldc = tskit.LdCalculator(ts)
        site = ts.num_sites // 2
        A1 = ldc.r2_array(site, direction=1, max_sites=max_sites)
        assert A1.shape[0] == max_sites
        A2 = ts_r2_array(ts, site, direction=1, max_sites=max_sites)
        assert A2.shape[0] == max_sites

    @pytest.mark.parametrize("max_sites", [0, 1, 2])
    def test_r2_array_backward_max_sites_zero(self, max_sites):
        ts, A = self.ts()
        ldc = tskit.LdCalculator(ts)
        site = ts.num_sites // 2
        A1 = ldc.r2_array(site, direction=-1, max_sites=max_sites)
        assert A1.shape[0] == max_sites
        A2 = ts_r2_array(ts, site, direction=-1, max_sites=max_sites)
        assert A2.shape[0] == max_sites


class TestLdOneSitePerTree(BaseTestLd):
    @tests.cached_example
    def ts(self):
        ts = msprime.sim_ancestry(
            5, sequence_length=10, recombination_rate=0.1, random_seed=1234
        )
        assert ts.num_trees > 3

        tables = ts.dump_tables()
        for tree in ts.trees():
            site = tables.sites.add_row(tree.interval[0], ancestral_state="A")
            # Put the mutation somewhere deep in the tree
            node = tree.preorder()[2]
            tables.mutations.add_row(site=site, node=node, derived_state="B")
        ts = tables.tree_sequence()
        # Return the full f2 matrix also
        return ts, get_r2_matrix(ts)


class TestLdAllSitesOneTree(BaseTestLd):
    @tests.cached_example
    def ts(self):
        ts = msprime.sim_ancestry(
            5, sequence_length=10, recombination_rate=0.1, random_seed=1234
        )
        assert ts.num_trees > 3

        tables = ts.dump_tables()
        tree = ts.at(5)
        pos = np.linspace(tree.interval[0], tree.interval[1], num=10, endpoint=False)
        for x, node in zip(pos, tree.preorder()[1:]):
            site = tables.sites.add_row(x, ancestral_state="A")
            tables.mutations.add_row(site=site, node=node, derived_state="B")
        ts = tables.tree_sequence()
        return ts, get_r2_matrix(ts)


class TestLdSitesEveryOtherTree(BaseTestLd):
    @tests.cached_example
    def ts(self):
        ts = msprime.sim_ancestry(
            5, sequence_length=20, recombination_rate=0.1, random_seed=1234
        )
        assert ts.num_trees > 5

        tables = ts.dump_tables()
        for tree in ts.trees():
            if tree.index % 2 == 0:
                pos = np.linspace(*tree.interval, num=2, endpoint=False)
                for x, node in zip(pos, tree.preorder()[1:]):
                    site = tables.sites.add_row(x, ancestral_state="A")
                    tables.mutations.add_row(site=site, node=node, derived_state="B")
        ts = tables.tree_sequence()
        return ts, get_r2_matrix(ts)


class TestLdErrors:
    def test_multi_mutations(self):
        tables = tskit.TableCollection(2)
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        tables.sites.add_row(position=0, ancestral_state="A")
        tables.sites.add_row(position=1, ancestral_state="A")
        tables.mutations.add_row(site=0, node=0, derived_state="C")
        tables.mutations.add_row(site=0, node=0, derived_state="T", parent=0)
        tables.mutations.add_row(site=1, node=0, derived_state="C")
        ts = tables.tree_sequence()
        ldc = tskit.LdCalculator(ts)
        with pytest.raises(tskit.LibraryError, match="Only infinite sites mutations"):
            ldc.r2(0, 1)
        with pytest.raises(tskit.LibraryError, match="Only infinite sites mutations"):
            ldc.r2(1, 0)

    @pytest.mark.parametrize("state", ["", "A", "AAAA", "💩"])
    def test_silent_mutations(self, state):
        tables = tskit.TableCollection(2)
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        tables.sites.add_row(position=0, ancestral_state=state)
        tables.sites.add_row(position=1, ancestral_state="A")
        tables.mutations.add_row(site=0, node=0, derived_state=state)
        tables.mutations.add_row(site=1, node=0, derived_state="C")
        ts = tables.tree_sequence()
        ldc = tskit.LdCalculator(ts)
        with pytest.raises(tskit.LibraryError, match="Silent mutations not supported"):
            ldc.r2(0, 1)
        with pytest.raises(tskit.LibraryError, match="Silent mutations not supported"):
            ldc.r2(1, 0)


class TestLdCalculator:
    """
    Tests for the LdCalculator class.
    """

    num_test_sites = 50

    def verify_matrix(self, ts):
        m = ts.get_num_sites()
        ldc = tskit.LdCalculator(ts)
        A = ldc.get_r2_matrix()
        assert A.shape == (m, m)
        B = get_r2_matrix(ts)
        assert np.allclose(A, B)

        # Now look at each row in turn, and verify it's the same
        # when we use get_r2 directly.
        for j in range(m):
            a = ldc.get_r2_array(j, direction=tskit.FORWARD)
            b = A[j, j + 1 :]
            assert a.shape[0] == m - j - 1
            assert b.shape[0] == m - j - 1
            assert np.allclose(a, b)
            a = ldc.get_r2_array(j, direction=tskit.REVERSE)
            b = A[j, :j]
            assert a.shape[0] == j
            assert b.shape[0] == j
            assert np.allclose(a[::-1], b)

        # Now check every cell in the matrix in turn.
        for j in range(m):
            for k in range(m):
                assert ldc.get_r2(j, k) == pytest.approx(A[j, k])

    def verify_max_distance(self, ts):
        """
        Verifies that the max_distance parameter works as expected.
        """
        mutations = list(ts.mutations())
        ldc = tskit.LdCalculator(ts)
        A = ldc.get_r2_matrix()
        j = len(mutations) // 2
        for k in range(j):
            x = (
                ts.site(mutations[j + k].site).position
                - ts.site(mutations[j].site).position
            )
            a = ldc.get_r2_array(j, max_distance=x)
            assert a.shape[0] == k
            assert np.allclose(A[j, j + 1 : j + 1 + k], a)
            x = (
                ts.site(mutations[j].site).position
                - ts.site(mutations[j - k].site).position
            )
            a = ldc.get_r2_array(j, max_distance=x, direction=tskit.REVERSE)
            assert a.shape[0] == k
            assert np.allclose(A[j, j - k : j], a[::-1])
        L = ts.get_sequence_length()
        m = len(mutations)
        a = ldc.get_r2_array(0, max_distance=L)
        assert a.shape[0] == m - 1
        assert np.allclose(A[0, 1:], a)
        a = ldc.get_r2_array(m - 1, max_distance=L, direction=tskit.REVERSE)
        assert a.shape[0] == m - 1
        assert np.allclose(A[m - 1, :-1], a[::-1])

    def verify_max_mutations(self, ts):
        """
        Verifies that the max mutations parameter works as expected.
        """
        mutations = list(ts.mutations())
        ldc = tskit.LdCalculator(ts)
        A = ldc.get_r2_matrix()
        j = len(mutations) // 2
        for k in range(j):
            a = ldc.get_r2_array(j, max_mutations=k)
            assert a.shape[0] == k
            assert np.allclose(A[j, j + 1 : j + 1 + k], a)
            a = ldc.get_r2_array(j, max_mutations=k, direction=tskit.REVERSE)
            assert a.shape[0] == k
            assert np.allclose(A[j, j - k : j], a[::-1])

    def test_single_tree_simulated_mutations(self):
        ts = msprime.simulate(20, mutation_rate=10, random_seed=15)
        ts = tsutil.subsample_sites(ts, self.num_test_sites)
        self.verify_matrix(ts)
        self.verify_max_distance(ts)

    def test_deprecated_get_aliases(self):
        ts = msprime.simulate(20, mutation_rate=10, random_seed=15)
        ts = tsutil.subsample_sites(ts, self.num_test_sites)
        ldc = tskit.LdCalculator(ts)
        A = ldc.get_r2_matrix()
        B = ldc.r2_matrix()
        assert np.array_equal(A, B)
        a = ldc.get_r2_array(0)
        b = ldc.r2_array(0)
        assert np.array_equal(a, b)
        assert ldc.get_r2(0, 1) == ldc.r2(0, 1)

    def test_deprecated_max_mutations_alias(self):
        ts = msprime.simulate(2, mutation_rate=0.1, random_seed=15)
        ldc = tskit.LdCalculator(ts)
        with pytest.raises(ValueError, match="deprecated synonym"):
            ldc.r2_array(0, max_sites=1, max_mutations=1)

    def test_single_tree_regular_mutations(self):
        ts = msprime.simulate(self.num_test_sites, length=self.num_test_sites)
        ts = tsutil.insert_branch_mutations(ts)
        # We don't support back mutations, so this should fail.
        with pytest.raises(_tskit.LibraryError):
            self.verify_matrix(ts)
        with pytest.raises(_tskit.LibraryError):
            self.verify_max_distance(ts)

    def test_tree_sequence_regular_mutations(self):
        ts = msprime.simulate(
            self.num_test_sites, recombination_rate=1, length=self.num_test_sites
        )
        assert ts.get_num_trees() > 10
        t = ts.dump_tables()
        t.sites.reset()
        t.mutations.reset()
        for j in range(self.num_test_sites):
            site_id = len(t.sites)
            t.sites.add_row(position=j, ancestral_state="0")
            t.mutations.add_row(site=site_id, derived_state="1", node=j)
        ts = t.tree_sequence()
        self.verify_matrix(ts)
        self.verify_max_distance(ts)

    def test_tree_sequence_simulated_mutations(self):
        ts = msprime.simulate(20, mutation_rate=10, recombination_rate=10)
        assert ts.get_num_trees() > 10
        ts = tsutil.subsample_sites(ts, self.num_test_sites)
        self.verify_matrix(ts)
        self.verify_max_distance(ts)
        self.verify_max_mutations(ts)


def set_partitions(collection):
    """
    Returns an iterator over all partitions of the specified set.

    From https://stackoverflow.com/questions/19368375/set-partitions-in-python
    """
    if len(collection) == 1:
        yield [collection]
    else:
        first = collection[0]
        for smaller in set_partitions(collection[1:]):
            for n, subset in enumerate(smaller):
                yield smaller[:n] + [[first] + subset] + smaller[n + 1 :]
            yield [[first]] + smaller


def naive_mean_descendants(ts, reference_sets):
    """
    Straightforward implementation of mean sample ancestry by iterating
    over the trees and nodes in each tree.
    """
    # TODO generalise this to allow arbitrary nodes, not just samples.
    C = np.zeros((ts.num_nodes, len(reference_sets)))
    T = np.zeros(ts.num_nodes)
    tree_iters = [ts.trees(tracked_samples=sample_set) for sample_set in reference_sets]
    for _ in range(ts.num_trees):
        trees = [next(tree_iter) for tree_iter in tree_iters]
        span = trees[0].span
        for node in trees[0].nodes():
            num_samples = trees[0].num_samples(node)
            if num_samples > 0:
                for j, tree in enumerate(trees):
                    C[node, j] += span * tree.num_tracked_samples(node)
                T[node] += span
    for node in range(ts.num_nodes):
        if T[node] > 0:
            C[node] /= T[node]
    return C


class TestMeanDescendants:
    """
    Tests the TreeSequence.mean_descendants method.
    """

    def verify(self, ts, reference_sets):
        C1 = naive_mean_descendants(ts, reference_sets)
        C2 = tsutil.mean_descendants(ts, reference_sets)
        C3 = ts.mean_descendants(reference_sets)
        assert C1.shape == C2.shape
        assert np.allclose(C1, C2)
        assert np.allclose(C1, C3)
        return C1

    def test_two_populations_high_migration(self):
        ts = msprime.simulate(
            population_configurations=[
                msprime.PopulationConfiguration(8),
                msprime.PopulationConfiguration(8),
            ],
            migration_matrix=[[0, 1], [1, 0]],
            recombination_rate=3,
            random_seed=5,
        )
        assert ts.num_trees > 1
        self.verify(ts, [ts.samples(0), ts.samples(1)])

    def test_single_tree(self):
        ts = msprime.simulate(6, random_seed=1)
        S = [range(3), range(3, 6)]
        C = self.verify(ts, S)
        for j, samples in enumerate(S):
            tree = next(ts.trees(tracked_samples=samples))
            for u in tree.nodes():
                assert tree.num_tracked_samples(u) == C[u, j]

    def test_single_tree_partial_samples(self):
        ts = msprime.simulate(6, random_seed=1)
        S = [range(3), range(3, 4)]
        C = self.verify(ts, S)
        for j, samples in enumerate(S):
            tree = next(ts.trees(tracked_samples=samples))
            for u in tree.nodes():
                assert tree.num_tracked_samples(u) == C[u, j]

    def test_single_tree_all_sample_sets(self):
        ts = msprime.simulate(6, random_seed=1)
        for S in set_partitions(list(range(ts.num_samples))):
            C = self.verify(ts, S)
            for j, samples in enumerate(S):
                tree = next(ts.trees(tracked_samples=samples))
                for u in tree.nodes():
                    assert tree.num_tracked_samples(u) == C[u, j]

    def test_many_trees_all_sample_sets(self):
        ts = msprime.simulate(6, recombination_rate=2, random_seed=1)
        assert ts.num_trees > 2
        for S in set_partitions(list(range(ts.num_samples))):
            self.verify(ts, S)

    def test_wright_fisher_unsimplified_all_sample_sets(self):
        tables = wf.wf_sim(
            4,
            5,
            seed=1,
            deep_history=False,
            initial_generation_samples=False,
            num_loci=10,
        )
        tables.sort()
        ts = tables.tree_sequence()
        for S in set_partitions(list(ts.samples())):
            self.verify(ts, S)

    def test_wright_fisher_unsimplified(self):
        tables = wf.wf_sim(
            20,
            15,
            seed=1,
            deep_history=False,
            initial_generation_samples=False,
            num_loci=20,
        )
        tables.sort()
        ts = tables.tree_sequence()
        samples = ts.samples()
        self.verify(ts, [samples[:10], samples[10:]])

    def test_wright_fisher_simplified(self):
        tables = wf.wf_sim(
            30,
            10,
            seed=1,
            deep_history=False,
            initial_generation_samples=False,
            num_loci=5,
        )
        tables.sort()
        ts = tables.tree_sequence()
        samples = ts.samples()
        self.verify(ts, [samples[:10], samples[10:]])


def naive_genealogical_nearest_neighbours(ts, focal, reference_sets):
    # Make sure everything is a sample so we can use the tracked_samples option.
    # This is a limitation of the current API.
    tables = ts.dump_tables()
    tables.nodes.set_columns(
        flags=np.ones_like(tables.nodes.flags), time=tables.nodes.time
    )
    ts = tables.tree_sequence()

    A = np.zeros((len(focal), len(reference_sets)))
    L = np.zeros(len(focal))
    reference_set_map = np.zeros(ts.num_nodes, dtype=int) - 1
    for k, ref_set in enumerate(reference_sets):
        for u in ref_set:
            reference_set_map[u] = k
    tree_iters = [
        ts.trees(tracked_samples=reference_nodes) for reference_nodes in reference_sets
    ]
    for _ in range(ts.num_trees):
        trees = list(map(next, tree_iters))
        length = trees[0].interval.right - trees[0].interval.left
        for j, u in enumerate(focal):
            focal_node_set = reference_set_map[u]
            # delta(u) = 1 if u exists in any of the reference sets; 0 otherwise
            delta = int(focal_node_set != -1)
            v = u
            while v != tskit.NULL:
                total = sum(tree.num_tracked_samples(v) for tree in trees)
                if total > delta:
                    break
                v = trees[0].parent(v)
            if v != tskit.NULL:
                for k, tree in enumerate(trees):
                    # If the focal node is in the current set, we subtract its
                    # contribution from the numerator
                    n = tree.num_tracked_samples(v) - (k == focal_node_set)
                    # If the focal node is in *any* reference set, we subtract its
                    # contribution from the demoninator.
                    A[j, k] += length * n / (total - delta)
                L[j] += length
    # Normalise by the accumulated value for each focal node.
    index = L > 0
    L = L[index]
    L = L.reshape((L.shape[0], 1))
    A[index, :] /= L
    return A


def parse_time_windows(ts, time_windows):
    if time_windows is None:
        time_windows = [0.0, ts.max_root_time]
    return np.array(time_windows)


def windowed_genealogical_nearest_neighbours(
    ts,
    focal,
    reference_sets,
    windows=None,
    time_windows=None,
    span_normalise=True,
    time_normalise=True,
):
    """
    genealogical_nearest_neighbours with support for span- and time-based windows
    """
    reference_set_map = np.full(ts.num_nodes, tskit.NULL, dtype=int)
    for k, reference_set in enumerate(reference_sets):
        for u in reference_set:
            if reference_set_map[u] != tskit.NULL:
                raise ValueError("Duplicate value in reference sets")
            reference_set_map[u] = k

    windows_used = windows is not None
    time_windows_used = time_windows is not None
    windows = ts.parse_windows(windows)
    num_windows = windows.shape[0] - 1
    time_windows = parse_time_windows(ts, time_windows)
    num_time_windows = time_windows.shape[0] - 1
    A = np.zeros((num_windows, num_time_windows, len(focal), len(reference_sets)))
    K = len(reference_sets)
    parent = np.full(ts.num_nodes, tskit.NULL, dtype=int)
    sample_count = np.zeros((ts.num_nodes, K), dtype=int)
    time = ts.tables.nodes.time
    norm = np.zeros((num_windows, num_time_windows, len(focal)))

    # Set the initial conditions.
    for j in range(K):
        sample_count[reference_sets[j], j] = 1

    window_index = 0
    for (t_left, t_right), edges_out, edges_in in ts.edge_diffs():
        for edge in edges_out:
            parent[edge.child] = tskit.NULL
            v = edge.parent
            while v != tskit.NULL:
                sample_count[v] -= sample_count[edge.child]
                v = parent[v]
        for edge in edges_in:
            parent[edge.child] = edge.parent
            v = edge.parent
            while v != tskit.NULL:
                sample_count[v] += sample_count[edge.child]
                v = parent[v]

        # Update the windows
        assert window_index < num_windows
        while windows[window_index] < t_right and window_index + 1 <= num_windows:
            w_left = windows[window_index]
            w_right = windows[window_index + 1]
            left = max(t_left, w_left)
            right = min(t_right, w_right)
            span = right - left
            # Process this tree.
            for j, u in enumerate(focal):
                focal_reference_set = reference_set_map[u]
                delta = int(focal_reference_set != tskit.NULL)
                p = u
                while p != tskit.NULL:
                    total = np.sum(sample_count[p])
                    if total > delta:
                        break
                    p = parent[p]
                if p != tskit.NULL:
                    scale = span / (total - delta)
                    time_index = np.searchsorted(time_windows, time[p]) - 1
                    if 0 <= time_index < num_time_windows:
                        for k in range(len(reference_sets)):
                            n = sample_count[p, k] - int(focal_reference_set == k)
                            A[window_index, time_index, j, k] += n * scale
                        norm[window_index, time_index, j] += span
            assert span > 0
            if w_right <= t_right:
                window_index += 1
            else:
                # This interval crosses a tree boundary, so we update it again
                # in the next tree
                break

    # Reshape norm depending on normalization selected
    # Return NaN when normalisation value is 0
    if span_normalise and time_normalise:
        reshaped_norm = norm.reshape((num_windows, num_time_windows, len(focal), 1))
    elif span_normalise and not time_normalise:
        norm = np.sum(norm, axis=1)
        reshaped_norm = norm.reshape((num_windows, 1, len(focal), 1))
    elif time_normalise and not span_normalise:
        norm = np.sum(norm, axis=0)
        reshaped_norm = norm.reshape((1, num_time_windows, len(focal), 1))

    with np.errstate(invalid="ignore", divide="ignore"):
        A /= reshaped_norm
    A[np.all(A == 0, axis=3)] = np.nan

    # Remove dimension for windows and/or time_windows if parameter is None
    if not windows_used and time_windows_used:
        A = A.reshape((num_time_windows, len(focal), len(reference_sets)))
    elif not time_windows_used and windows_used:
        A = A.reshape((num_windows, len(focal), len(reference_sets)))
    elif not windows_used and not time_windows_used:
        A = A.reshape((len(focal), len(reference_sets)))
    return A


class TestGenealogicalNearestNeighbours:
    """
    Tests the TreeSequence.genealogical_nearest_neighbours method.
    """

    #
    #          8
    #         / \
    #        /   \
    #       /     \
    #      7       \
    #     / \       6
    #    /   5     / \
    #   /   / \   /   \
    #  4   0   1 2     3
    small_tree_ex_nodes = """\
    id      is_sample   population      time
    0       1       0               0.00000000000000
    1       1       0               0.00000000000000
    2       1       0               0.00000000000000
    3       1       0               0.00000000000000
    4       1       0               0.00000000000000
    5       0       0               0.14567111023387
    6       0       0               0.21385545626353
    7       0       0               0.43508024345063
    8       0       0               1.60156352971203
    """
    small_tree_ex_edges = """\
    id      left            right           parent  child
    0       0.00000000      1.00000000      5       0,1
    1       0.00000000      1.00000000      6       2,3
    2       0.00000000      1.00000000      7       4,5
    3       0.00000000      1.00000000      8       6,7
    """

    def verify(self, ts, reference_sets, focal=None):
        if focal is None:
            focal = [u for refset in reference_sets for u in refset]
        A1 = naive_genealogical_nearest_neighbours(ts, focal, reference_sets)
        A2 = tsutil.genealogical_nearest_neighbours(ts, focal, reference_sets)
        A3 = ts.genealogical_nearest_neighbours(focal, reference_sets)
        A4 = ts.genealogical_nearest_neighbours(focal, reference_sets, num_threads=3)
        A5 = windowed_genealogical_nearest_neighbours(ts, focal, reference_sets)
        assert np.array_equal(A3, A4)
        assert A1.shape == A2.shape
        assert A1.shape == A3.shape
        assert np.allclose(A1, A2)
        assert np.allclose(A1, A3)
        mask = ~np.isnan(A5)
        assert np.sum(mask) > 0 or ts.num_edges == 0
        assert np.allclose(A1[mask], A5[mask])
        assert np.allclose(A5[mask], A2[mask])
        assert np.allclose(A5[mask], A3[mask])

        if ts.num_edges > 0 and all(ts.node(u).is_sample() for u in focal):
            # When the focal nodes are samples, we can assert some stronger properties.
            assert np.allclose(np.sum(A1, axis=1), 1)
            assert np.allclose(np.sum(A5, axis=1), 1)
        return A1

    def test_simple_example_all_samples(self):
        ts = tskit.load_text(
            nodes=io.StringIO(self.small_tree_ex_nodes),
            edges=io.StringIO(self.small_tree_ex_edges),
            strict=False,
        )
        A = self.verify(ts, [[0, 1], [2, 3, 4]], [0])
        assert list(A[0]) == [1, 0]
        A = self.verify(ts, [[0, 1], [2, 3, 4]], [4])
        assert list(A[0]) == [1, 0]
        A = self.verify(ts, [[0, 1], [2, 3, 4]], [2])
        assert list(A[0]) == [0, 1]
        A = self.verify(ts, [[0, 2], [1, 3, 4]], [0])
        assert list(A[0]) == [0, 1]
        A = self.verify(ts, [[0, 2], [1, 3, 4]], [4])
        assert list(A[0]) == [0.5, 0.5]

    def test_simple_example_missing_samples(self):
        ts = tskit.load_text(
            nodes=io.StringIO(self.small_tree_ex_nodes),
            edges=io.StringIO(self.small_tree_ex_edges),
            strict=False,
        )
        A = self.verify(ts, [[0, 1], [2, 4]], [3])
        assert list(A[0]) == [0, 1]
        A = self.verify(ts, [[0, 1], [2, 4]], [2])
        assert np.allclose(A[0], [2 / 3, 1 / 3])

    def test_simple_example_internal_focal_node(self):
        ts = tskit.load_text(
            nodes=io.StringIO(self.small_tree_ex_nodes),
            edges=io.StringIO(self.small_tree_ex_edges),
            strict=False,
        )
        focal = [7]  # An internal node
        reference_sets = [[4, 0, 1], [2, 3]]
        GNN = naive_genealogical_nearest_neighbours(ts, focal, reference_sets)
        assert np.allclose(GNN[0], np.array([1.0, 0.0]))
        GNN = tsutil.genealogical_nearest_neighbours(ts, focal, reference_sets)
        assert np.allclose(GNN[0], np.array([1.0, 0.0]))
        GNN = ts.genealogical_nearest_neighbours(focal, reference_sets)
        assert np.allclose(GNN[0], np.array([1.0, 0.0]))
        focal = [8]  # The root
        GNN = naive_genealogical_nearest_neighbours(ts, focal, reference_sets)
        assert np.allclose(GNN[0], np.array([0.6, 0.4]))
        GNN = tsutil.genealogical_nearest_neighbours(ts, focal, reference_sets)
        assert np.allclose(GNN[0], np.array([0.6, 0.4]))
        GNN = ts.genealogical_nearest_neighbours(focal, reference_sets)
        assert np.allclose(GNN[0], np.array([0.6, 0.4]))

    def test_two_populations_high_migration(self):
        ts = msprime.simulate(
            population_configurations=[
                msprime.PopulationConfiguration(18),
                msprime.PopulationConfiguration(18),
            ],
            migration_matrix=[[0, 1], [1, 0]],
            recombination_rate=8,
            random_seed=5,
        )
        assert ts.num_trees > 1
        self.verify(ts, [ts.samples(0), ts.samples(1)])

    def test_single_tree(self):
        ts = msprime.simulate(6, random_seed=1)
        S = [range(3), range(3, 6)]
        self.verify(ts, S)

    def test_single_tree_internal_reference_sets(self):
        ts = msprime.simulate(10, random_seed=1)
        tree = ts.first()
        S = [[u] for u in tree.children(tree.root)]
        self.verify(ts, S, ts.samples())

    def test_single_tree_all_nodes(self):
        ts = msprime.simulate(10, random_seed=1)
        S = [np.arange(ts.num_nodes, dtype=np.int32)]
        self.verify(ts, S, np.arange(ts.num_nodes, dtype=np.int32))

    def test_single_tree_partial_samples(self):
        ts = msprime.simulate(6, random_seed=1)
        S = [range(3), range(3, 4)]
        self.verify(ts, S)

    def test_single_tree_all_sample_sets(self):
        ts = msprime.simulate(6, random_seed=1)
        for S in set_partitions(list(range(ts.num_samples))):
            self.verify(ts, S)

    def test_many_trees_all_sample_sets(self):
        ts = msprime.simulate(6, recombination_rate=2, random_seed=1)
        assert ts.num_trees > 2
        for S in set_partitions(list(range(ts.num_samples))):
            self.verify(ts, S)

    def test_many_trees_sequence_length(self):
        for L in [0.5, 1.5, 3.3333]:
            ts = msprime.simulate(6, length=L, recombination_rate=2, random_seed=1)
            self.verify(ts, [range(3), range(3, 6)])

    def test_many_trees_all_nodes(self):
        ts = msprime.simulate(6, length=4, recombination_rate=2, random_seed=1)
        S = [np.arange(ts.num_nodes, dtype=np.int32)]
        self.verify(ts, S, np.arange(ts.num_nodes, dtype=np.int32))

    def test_wright_fisher_unsimplified_all_sample_sets(self):
        tables = wf.wf_sim(
            4,
            5,
            seed=1,
            deep_history=True,
            initial_generation_samples=False,
            num_loci=10,
        )
        tables.sort()
        ts = tables.tree_sequence()
        for S in set_partitions(list(ts.samples())):
            self.verify(ts, S)

    def test_wright_fisher_unsimplified(self):
        tables = wf.wf_sim(
            20,
            15,
            seed=1,
            deep_history=True,
            initial_generation_samples=False,
            num_loci=20,
        )
        tables.sort()
        ts = tables.tree_sequence()
        samples = ts.samples()
        self.verify(ts, [samples[:10], samples[10:]])

    def test_wright_fisher_initial_generation(self):
        tables = wf.wf_sim(
            20,
            15,
            seed=1,
            deep_history=True,
            initial_generation_samples=True,
            num_loci=20,
        )
        tables.sort()
        tables.simplify()
        ts = tables.tree_sequence()
        samples = ts.samples()
        founders = [u for u in samples if ts.node(u).time > 0]
        samples = [u for u in samples if ts.node(u).time == 0]
        self.verify(ts, [founders[:10], founders[10:]], samples)

    def test_wright_fisher_initial_generation_no_deep_history(self):
        tables = wf.wf_sim(
            20,
            15,
            seed=2,
            deep_history=False,
            initial_generation_samples=True,
            num_loci=20,
        )
        tables.sort()
        tables.simplify()
        ts = tables.tree_sequence()
        samples = ts.samples()
        founders = [u for u in samples if ts.node(u).time > 0]
        samples = [u for u in samples if ts.node(u).time == 0]
        self.verify(ts, [founders[:10], founders[10:]], samples)

    def test_wright_fisher_unsimplified_multiple_roots(self):
        tables = wf.wf_sim(
            20,
            15,
            seed=1,
            deep_history=False,
            initial_generation_samples=False,
            num_loci=20,
        )
        tables.sort()
        ts = tables.tree_sequence()
        samples = ts.samples()
        self.verify(ts, [samples[:10], samples[10:]])

    def test_wright_fisher_simplified(self):
        tables = wf.wf_sim(
            31,
            10,
            seed=1,
            deep_history=True,
            initial_generation_samples=False,
            num_loci=5,
        )
        tables.sort()
        ts = tables.tree_sequence().simplify()
        samples = ts.samples()
        self.verify(ts, [samples[:10], samples[10:]])

    def test_wright_fisher_simplified_multiple_roots(self):
        tables = wf.wf_sim(
            31,
            10,
            seed=1,
            deep_history=False,
            initial_generation_samples=False,
            num_loci=5,
        )
        tables.sort()
        ts = tables.tree_sequence()
        samples = ts.samples()
        self.verify(ts, [samples[:10], samples[10:]])

    def test_empty_ts(self):
        tables = tskit.TableCollection(1.0)
        tables.nodes.add_row(1, 0)
        tables.nodes.add_row(1, 0)
        ts = tables.tree_sequence()
        self.verify(ts, [[0], [1]])


class TestWindowedGenealogicalNearestNeighbours(TestGenealogicalNearestNeighbours):
    """
    Tests the TreeSequence.genealogical_nearest_neighbours method.
    """

    #               .    5
    #               .   / \
    #        4      .  |   4
    #       / \     .  |   |\
    #      3   \    .  |   | \
    #     / \   \   .  |   |  \
    #   [0] [1] [2] . [0] [1] [2]
    #
    two_tree_nodes = """\
    id      is_sample   time
    0       1           0
    1       1           0
    2       1           0
    3       0           1
    4       0           2
    5       0           3
    """
    two_tree_edges = """\
    left    right   parent  child
    0       0.2     3       0,1
    0       1       4       2
    0       0.2     4       3
    0.2     1       4       1
    0.2     1       5       0,4
    """

    def get_two_tree_ts(self):
        ts = tskit.load_text(
            nodes=io.StringIO(self.two_tree_nodes),
            edges=io.StringIO(self.two_tree_edges),
            strict=False,
        )
        return ts

    def verify(self, ts, reference_sets, focal=None, windows=None, time_windows=None):
        if focal is None:
            focal = [u for refset in reference_sets for u in refset]
        gnn = windowed_genealogical_nearest_neighbours(
            ts, focal, reference_sets, windows, time_windows
        )
        if windows is not None:
            windows_len = len(windows) - 1
        if time_windows is not None:
            time_windows_len = len(time_windows) - 1
        if windows is None and time_windows is None:
            assert np.array_equal(gnn.shape, [len(focal), len(reference_sets)])
        elif windows is None and time_windows is not None:
            assert np.array_equal(
                gnn.shape, [time_windows_len, len(focal), len(reference_sets)]
            )
        elif windows is not None and time_windows is None:
            assert np.array_equal(
                gnn.shape, [windows_len, len(focal), len(reference_sets)]
            )
        else:
            assert np.array_equal(
                gnn.shape,
                [windows_len, time_windows_len, len(focal), len(reference_sets)],
            )

        return gnn

    def test_one_tree_windows(self):
        ts = tskit.load_text(
            nodes=io.StringIO(self.small_tree_ex_nodes),
            edges=io.StringIO(self.small_tree_ex_edges),
            strict=False,
        )
        A = self.verify(ts, [[0, 1], [2, 3, 4]], [0], [0, 1])
        assert np.allclose(A, [[[[1, 0]]]])
        A = self.verify(ts, [[0, 1], [2, 3, 4]], [0], [0, 0.5, 1])
        assert np.allclose(A, [[[[1.0, 0.0]]], [[[1.0, 0.0]]]])
        A = self.verify(ts, [[0, 1], [2, 3, 4]], [0], [0, 0.5, 0.6, 1])
        assert np.allclose(A, [[[[1.0, 0.0]]], [[[1.0, 0.0]]], [[[1.0, 0.0]]]])

    def test_two_tree_windows(self):
        ts = self.get_two_tree_ts()
        A = self.verify(ts, [[0, 1], [2]], [0], [0, 1])
        assert np.allclose(A, [[[0.6, 0.4]]])
        A = self.verify(ts, [[0, 1], [2]], [0], [0, 0.2, 1])
        assert np.allclose(A, [[[1.0, 0.0]], [[0.5, 0.5]]])
        A = self.verify(ts, [[0, 1], [2]], [0], [0, 0.2, 0.5, 1])
        assert np.allclose(A, [[[1.0, 0.0]], [[0.5, 0.5]], [[0.5, 0.5]]])

    def test_one_tree_time_windows(self):
        ts = tskit.load_text(
            nodes=io.StringIO(self.small_tree_ex_nodes),
            edges=io.StringIO(self.small_tree_ex_edges),
            strict=False,
        )
        A = self.verify(ts, [[0, 1], [2, 3, 4]], [0], None, [0, ts.max_root_time])
        assert np.allclose(A, [[[1, 0]]])
        A = self.verify(ts, [[0, 1], [2, 3, 4]], [0], None, [1, 2])
        assert np.allclose(A, [[[np.nan, np.nan]]], equal_nan=True)
        A = self.verify(ts, [[0, 1], [2, 3, 4]], [0], None, [0, 0.1])
        assert np.allclose(A, [[[np.nan, np.nan]]], equal_nan=True)

    def test_two_tree_time_windows(self):
        ts = self.get_two_tree_ts()
        A = self.verify(ts, [[0, 1], [2]], [0], None, [0, ts.max_root_time])
        assert np.allclose(A, [[[0.6, 0.4]]])
        A = self.verify(ts, [[0, 1], [2]], [0], None, [0, 1.1, ts.max_root_time])
        assert np.allclose(A, [[[1.0, 0.0]], [[0.5, 0.5]]])
        A = self.verify(ts, [[0, 1], [2]], [0], None, [0, 0.5, 1])
        assert np.allclose(A, [[[np.nan, np.nan]], [[1.0, 0.0]]], equal_nan=True)
        A = self.verify(ts, [[0, 1], [2]], [0], None, [1, ts.max_root_time, 10])
        assert np.allclose(A, [[[0.5, 0.5]], [[np.nan, np.nan]]], equal_nan=True)

    def test_one_tree_windows_time_windows(self):
        ts = tskit.load_text(
            nodes=io.StringIO(self.small_tree_ex_nodes),
            edges=io.StringIO(self.small_tree_ex_edges),
            strict=False,
        )
        A = self.verify(ts, [[0, 1], [2, 3, 4]], [0], [0, 1], [0, ts.max_root_time])
        assert np.allclose(A, [[[[1, 0]]]])
        A = self.verify(
            ts, [[0, 1], [2, 3, 4]], [0], [0, 0.2, 1], [0, 1.1, ts.max_root_time]
        )
        assert np.allclose(
            A,
            [
                [[[1.0, 0.0]], [[np.nan, np.nan]]],
                [[[1.0, 0.0]], [[np.nan, np.nan]]],
            ],
            equal_nan=True,
        )
        A = self.verify(ts, [[0, 1], [2, 3, 4]], [0], [0, 0.2], [0, 0.5, 1])
        assert np.allclose(A, [[[[1.0, 0.0]], [[np.nan, np.nan]]]], equal_nan=True)
        A = self.verify(
            ts, [[0, 1], [2, 3, 4]], [0], [0, 0.2, 1, 1.5], [0, ts.max_root_time, 10]
        )
        assert np.allclose(
            A,
            [
                [[[1.0, 0.0]], [[np.nan, np.nan]]],
                [[[1.0, 0.0]], [[np.nan, np.nan]]],
                [[[np.nan, np.nan]], [[np.nan, np.nan]]],
            ],
            equal_nan=True,
        )

    def test_two_tree_windows_time_windows(self):
        ts = self.get_two_tree_ts()
        A = self.verify(ts, [[0, 1], [2]], [0], [0, 1], [0, ts.max_root_time])
        assert np.allclose(A, [[[[0.6, 0.4]]]])
        A = self.verify(ts, [[0, 1], [2]], [0], [0, 0.2, 1], [0, 1.1, ts.max_root_time])
        assert np.allclose(
            A,
            [
                [[[1.0, 0.0]], [[np.nan, np.nan]]],
                [[[np.nan, np.nan]], [[0.5, 0.5]]],
            ],
            equal_nan=True,
        )
        A = self.verify(ts, [[0, 1], [2, 3, 4]], [0], [0, 0.2, 1], [0, 0.5, 1])
        assert np.allclose(
            A,
            [
                [[[np.nan, np.nan]], [[0.5, 0.5]]],
                [[[np.nan, np.nan]], [[np.nan, np.nan]]],
            ],
            equal_nan=True,
        )

    def test_span_normalise(self):
        ts = self.get_two_tree_ts()
        sample_sets = [[0, 1], [2]]
        focal = [0]
        np.random.seed(5)
        windows = ts.sequence_length * np.array([0.2, 0.4, 0.6, 0.8, 1])
        windows.sort()
        windows[0] = 0.0
        windows[-1] = ts.sequence_length

        result1 = windowed_genealogical_nearest_neighbours(
            ts, focal, sample_sets, windows
        )
        result2 = windowed_genealogical_nearest_neighbours(
            ts, focal, sample_sets, windows, span_normalise=True
        )
        result3 = windowed_genealogical_nearest_neighbours(
            ts, focal, sample_sets, windows, span_normalise=False
        )
        denom = np.diff(windows)[:, np.newaxis, np.newaxis]

        # Test the dimensions are correct
        assert np.array_equal(result1.shape, result2.shape)
        assert np.array_equal(result2.shape, result3.shape)

        # Test normalisation is correct
        assert np.allclose(result1, result2)
        assert np.allclose(result1, result3 / denom)

        # If span_normalised, then sum over all reference sets should be 1
        assert np.allclose(np.sum(result1, axis=2), 1)
        assert np.allclose(np.sum(result2, axis=2), 1)
        # If not span_normalised, then sum over all value is 1
        assert np.allclose(result3.sum(), 1)

    def test_time_normalise(self):
        """
        Testing time_normalise is trickier than span_normalise, as the norm
        depends on the span of the nearest neighbours found in each time grid.
        In this small example, we check which grid nodes 3 and 5 fall in, and use their
        spans to check the time_normalisation.
        """
        ts = self.get_two_tree_ts()
        sample_sets = [[0, 1], [2]]
        focal = [0]
        oldest_node = ts.max_root_time
        time_windows = oldest_node * np.array([0.2, 0.4, 0.6, 0.8, 1])
        time_windows.sort()
        time_windows[0] = 0.0
        time_windows[-1] = oldest_node

        # Determine output_dim of the function
        result1 = windowed_genealogical_nearest_neighbours(
            ts, focal, sample_sets, windows=None, time_windows=time_windows
        )
        result2 = windowed_genealogical_nearest_neighbours(
            ts,
            focal,
            sample_sets,
            windows=None,
            time_windows=time_windows,
            time_normalise=True,
        )
        result3 = windowed_genealogical_nearest_neighbours(
            ts,
            focal,
            sample_sets,
            windows=None,
            time_windows=time_windows,
            time_normalise=False,
        )
        denom = np.zeros(len(time_windows) - 1)
        time_index_3 = np.searchsorted(time_windows, ts.tables.nodes.time[3]) - 1
        time_index_5 = np.searchsorted(time_windows, ts.tables.nodes.time[5]) - 1
        denom[time_index_3] += 0.2
        denom[time_index_5] += 0.8

        # Avoid division by zero
        denom[denom == 0] = 1
        denom = denom[:, np.newaxis, np.newaxis]

        # Test the dimensions are correct
        assert np.array_equal(result1.shape, result2.shape)
        assert np.array_equal(result2.shape, result3.shape)

        # Test normalisation is correct
        assert np.allclose(result1, result2, equal_nan=True)
        assert np.allclose(result1, result3 / denom, equal_nan=True)

        # If time_normalised, then sum over all reference sets should be 1
        # Mask out time intervals that sum to 0
        result1_dim_sum = np.sum(result1, axis=2)
        mask = ~(np.isnan(result1_dim_sum))
        assert np.allclose(
            result1_dim_sum[mask],
            np.ones((len(result1_dim_sum), len(focal)))[mask],
            equal_nan=True,
        )
        result2_dim_sum = np.sum(result2, axis=2)
        mask = ~(np.isnan(result2_dim_sum))
        assert np.allclose(
            result2_dim_sum[mask],
            np.ones((len(result2_dim_sum), len(focal)))[mask],
            equal_nan=True,
        )
        # If not span_normalised, then sum over all value is 1
        assert np.allclose(np.nansum(result3), 1)


def exact_genealogical_nearest_neighbours(ts, focal, reference_sets):
    # Same as above, except we return the per-tree value for a single node.

    # Make sure everyhing is a sample so we can use the tracked_samples option.
    # This is a limitation of the current API.
    tables = ts.dump_tables()
    tables.nodes.set_columns(
        flags=np.ones_like(tables.nodes.flags), time=tables.nodes.time
    )
    ts = tables.tree_sequence()

    A = np.zeros((len(reference_sets), ts.num_trees))
    L = np.zeros(ts.num_trees)
    reference_set_map = np.zeros(ts.num_nodes, dtype=int) - 1
    for k, ref_set in enumerate(reference_sets):
        for u in ref_set:
            reference_set_map[u] = k
    tree_iters = [
        ts.trees(tracked_samples=reference_nodes) for reference_nodes in reference_sets
    ]
    u = focal
    focal_node_set = reference_set_map[u]
    # delta(u) = 1 if u exists in any of the reference sets; 0 otherwise
    delta = int(focal_node_set != -1)
    for _ in range(ts.num_trees):
        trees = list(map(next, tree_iters))
        v = trees[0].parent(u)
        while v != tskit.NULL:
            total = sum(tree.num_tracked_samples(v) for tree in trees)
            if total > delta:
                break
            v = trees[0].parent(v)
        if v != tskit.NULL:
            # The length is only reported where the statistic is defined.
            L[trees[0].index] = trees[0].interval.right - trees[0].interval.left
            for k, tree in enumerate(trees):
                # If the focal node is in the current set, we subtract its
                # contribution from the numerator
                n = tree.num_tracked_samples(v) - (k == focal_node_set)
                # If the focal node is in *any* reference set, we subtract its
                # contribution from the demoninator.
                A[k, tree.index] = n / (total - delta)
    return A, L


def local_gnn(ts, focal, reference_sets):
    # Temporary implementation of the treewise GNN.
    reference_set_map = np.zeros(ts.num_nodes, dtype=int) - 1
    for k, reference_set in enumerate(reference_sets):
        for u in reference_set:
            if reference_set_map[u] != -1:
                raise ValueError("Duplicate value in reference sets")
            reference_set_map[u] = k

    K = len(reference_sets)
    A = np.zeros((len(focal), ts.num_trees, K))
    lefts = np.zeros(ts.num_trees, dtype=float)
    rights = np.zeros(ts.num_trees, dtype=float)
    parent = np.zeros(ts.num_nodes, dtype=int) - 1
    sample_count = np.zeros((ts.num_nodes, K), dtype=int)

    # Set the intitial conditions.
    for j in range(K):
        sample_count[reference_sets[j], j] = 1

    for t, ((left, right), edges_out, edges_in) in enumerate(ts.edge_diffs()):
        for edge in edges_out:
            parent[edge.child] = -1
            v = edge.parent
            while v != -1:
                sample_count[v] -= sample_count[edge.child]
                v = parent[v]
        for edge in edges_in:
            parent[edge.child] = edge.parent
            v = edge.parent
            while v != -1:
                sample_count[v] += sample_count[edge.child]
                v = parent[v]

        # Process this tree.
        for j, u in enumerate(focal):
            focal_reference_set = reference_set_map[u]
            delta = int(focal_reference_set != -1)
            p = parent[u]
            lefts[t] = left
            rights[t] = right
            while p != tskit.NULL:
                total = np.sum(sample_count[p])
                if total > delta:
                    break
                p = parent[p]
            if p != tskit.NULL:
                scale = 1 / (total - delta)
                for k, _reference_set in enumerate(reference_sets):
                    n = sample_count[p, k] - int(focal_reference_set == k)
                    A[j, t, k] = n * scale
    return (A, lefts, rights)


class TestExactGenealogicalNearestNeighbours(TestGenealogicalNearestNeighbours):
    # This is a work in progress - these tests will be adapted to use the
    # treewise GNN when it's implemented.

    def verify(self, ts, reference_sets, focal=None):
        if focal is None:
            focal = [u for refset in reference_sets for u in refset]
        A = ts.genealogical_nearest_neighbours(focal, reference_sets)

        G, lefts, rights = local_gnn(ts, focal, reference_sets)
        for tree in ts.trees():
            assert lefts[tree.index] == tree.interval.left
            assert rights[tree.index] == tree.interval.right

        for j, u in enumerate(focal):
            T, L = exact_genealogical_nearest_neighbours(ts, u, reference_sets)
            assert np.allclose(G[j], T.T)
            # Ignore the cases where the node has no GNNs
            if np.sum(L) > 0:
                mean = np.sum(T * L, axis=1) / np.sum(L)
                assert np.allclose(mean, A[j])
        return A


--- ../../tskit/python/tests/test_file_format.py ---


"""
Test cases for tskit's file format.
"""
import json
import os
import sys
import tempfile
import unittest
import uuid as _uuid
from unittest import mock

import h5py
import kastore
import msprime
import numpy as np
import pytest
import tszip as tszip

import tests.tsutil as tsutil
import tskit
import tskit.exceptions as exceptions

CURRENT_FILE_MAJOR = 12
CURRENT_FILE_MINOR = 7

test_data_dir = os.path.join(os.path.dirname(__file__), "data")


def single_locus_no_mutation_example():
    return msprime.simulate(10, random_seed=10)


def single_locus_with_mutation_example():
    return msprime.simulate(10, mutation_rate=10, random_seed=11)


def multi_locus_with_mutation_example():
    return msprime.simulate(
        10, recombination_rate=1, length=10, mutation_rate=10, random_seed=2
    )


def recurrent_mutation_example():
    ts = msprime.simulate(10, recombination_rate=1, length=10, random_seed=2)
    return tsutil.insert_branch_mutations(ts)


def general_mutation_example():
    ts = msprime.simulate(10, recombination_rate=1, length=10, random_seed=2)
    tables = ts.dump_tables()
    tables.sites.add_row(position=0, ancestral_state="A", metadata=b"{}")
    tables.sites.add_row(position=1, ancestral_state="C", metadata=b"{'id':1}")
    tables.mutations.add_row(site=0, node=0, derived_state="T")
    tables.mutations.add_row(site=1, node=0, derived_state="G")
    return tables.tree_sequence()


def multichar_mutation_example():
    ts = msprime.simulate(10, recombination_rate=1, length=10, random_seed=2)
    return tsutil.insert_multichar_mutations(ts)


def migration_example():
    n = 10
    t = 1
    population_configurations = [
        msprime.PopulationConfiguration(n // 2),
        msprime.PopulationConfiguration(n // 2),
        msprime.PopulationConfiguration(0),
    ]
    demographic_events = [
        msprime.MassMigration(time=t, source=0, destination=2),
        msprime.MassMigration(time=t, source=1, destination=2),
    ]
    ts = msprime.simulate(
        population_configurations=population_configurations,
        demographic_events=demographic_events,
        random_seed=1,
        mutation_rate=1,
        record_migrations=True,
    )
    tables = ts.dump_tables()
    for j in range(n):
        tables.individuals.add_row(flags=j, location=(j, j), parents=(j - 1, j - 1))
    return tables.tree_sequence()


def bottleneck_example():
    return msprime.simulate(
        random_seed=1,
        sample_size=100,
        recombination_rate=0.1,
        length=10,
        demographic_events=[
            msprime.SimpleBottleneck(time=0.01, population=0, proportion=0.75)
        ],
    )


def historical_sample_example():
    return msprime.simulate(
        recombination_rate=0.1,
        length=10,
        random_seed=1,
        samples=[(0, j) for j in range(10)],
    )


def no_provenance_example():
    ts = msprime.simulate(10, random_seed=1)
    tables = ts.dump_tables()
    tables.provenances.clear()
    return tables.tree_sequence()


def provenance_timestamp_only_example():
    ts = msprime.simulate(10, random_seed=1)
    tables = ts.dump_tables()
    provenances = tskit.ProvenanceTable()
    provenances.add_row(timestamp="12345", record="")
    return tables.tree_sequence()


def node_metadata_example():
    ts = msprime.simulate(
        sample_size=100, recombination_rate=0.1, length=10, random_seed=1
    )
    tables = ts.dump_tables()
    metadatas = [f"n_{u}" for u in range(ts.num_nodes)]
    packed, offset = tskit.pack_strings(metadatas)
    tables.nodes.set_columns(
        metadata=packed,
        metadata_offset=offset,
        flags=tables.nodes.flags,
        time=tables.nodes.time,
    )
    return tables.tree_sequence()


def site_metadata_example():
    ts = msprime.simulate(10, length=10, random_seed=2)
    tables = ts.dump_tables()
    for j in range(10):
        tables.sites.add_row(j, ancestral_state="a", metadata=b"1234")
    return tables.tree_sequence()


def mutation_metadata_example():
    ts = msprime.simulate(10, length=10, random_seed=2)
    tables = ts.dump_tables()
    tables.sites.add_row(0, ancestral_state="a")
    for j in range(10):
        tables.mutations.add_row(site=0, node=j, derived_state="t", metadata=b"1234")
    return tables.tree_sequence()


def migration_metadata_example():
    ts = migration_example()
    tables = ts.dump_tables()
    metadatas = [f"n_{u}" for u in range(ts.num_migrations)]
    packed, offset = tskit.pack_strings(metadatas)
    tables.migrations.set_columns(
        metadata=packed,
        metadata_offset=offset,
        left=tables.migrations.left,
        right=tables.migrations.right,
        source=tables.migrations.source,
        dest=tables.migrations.dest,
        node=tables.migrations.node,
        time=tables.migrations.time,
    )
    return tables.tree_sequence()


def edge_metadata_example():
    ts = msprime.simulate(
        sample_size=100, recombination_rate=0.1, length=10, random_seed=1
    )
    tables = ts.dump_tables()
    metadatas = [f"edge_{u}" for u in range(ts.num_edges)]
    packed, offset = tskit.pack_strings(metadatas)
    tables.edges.set_columns(
        metadata=packed,
        metadata_offset=offset,
        left=tables.edges.left,
        right=tables.edges.right,
        child=tables.edges.child,
        parent=tables.edges.parent,
    )
    return tables.tree_sequence()


class TestFileFormat(unittest.TestCase):
    """
    Superclass of file format tests.
    """

    def setUp(self):
        fd, self.temp_file = tempfile.mkstemp(prefix="msp_file_test_")
        os.close(fd)

    def tearDown(self):
        os.unlink(self.temp_file)


class TestLoadLegacyExamples(TestFileFormat):
    """
    Tests using the saved legacy file examples to ensure we can load them.
    """

    def verify_tree_sequence(self, ts):
        # Just some quick checks to make sure the tree sequence makes sense.
        assert ts.sample_size > 0
        assert ts.num_edges > 0
        assert ts.num_sites > 0
        assert ts.num_mutations > 0
        assert ts.sequence_length > 0
        for t in ts.trees():
            left, right = t.interval
            assert right > left
            for site in t.sites():
                assert left <= site.position < right
                for mut in site.mutations:
                    assert mut.site == site.id

    def verify_0_3_3(self, ts):
        for table in tskit.TABLE_NAMES:
            t = getattr(ts.tables, table)
            assert t.num_rows > 0
            if hasattr(t, "metadata_schema"):
                assert t.metadata_schema == tskit.MetadataSchema({"codec": "json"})
                assert t[2].metadata == f"n_{table}_2"
        assert ts.tables.has_index()

    def test_format_too_old_raised_for_hdf5(self):
        files = [
            "msprime-0.3.0_v2.0.hdf5",
            "msprime-0.4.0_v3.1.hdf5",
            "msprime-0.5.0_v10.0.hdf5",
        ]
        for filename in files:
            path = os.path.join(test_data_dir, "hdf5-formats", filename)

            with pytest.raises(
                exceptions.FileFormatError,
                match="appears to be in HDF5 format",
            ):
                tskit.load(path)
            with pytest.raises(
                exceptions.FileFormatError,
                match="appears to be in HDF5 format",
            ):
                tskit.TableCollection.load(path)

    def test_msprime_v_0_5_0(self):
        path = os.path.join(test_data_dir, "hdf5-formats", "msprime-0.5.0_v10.0.hdf5")
        ts = tskit.load_legacy(path)
        self.verify_tree_sequence(ts)

    def test_msprime_v_0_4_0(self):
        path = os.path.join(test_data_dir, "hdf5-formats", "msprime-0.4.0_v3.1.hdf5")
        ts = tskit.load_legacy(path)
        self.verify_tree_sequence(ts)

    def test_msprime_v_0_3_0(self):
        path = os.path.join(test_data_dir, "hdf5-formats", "msprime-0.3.0_v2.0.hdf5")
        ts = tskit.load_legacy(path)
        self.verify_tree_sequence(ts)

    def test_tskit_v_0_3_3(self):
        path = os.path.join(test_data_dir, "old-formats", "tskit-0.3.3.trees")
        ts = tskit.load(path)
        self.verify_tree_sequence(ts)


class TestRoundTrip(TestFileFormat):
    """
    Tests if we can round trip convert a tree sequence in memory
    through a V2 file format and a V3 format.
    """

    def verify_tree_sequences_equal(self, ts, tsp, simplify=True):
        assert ts.sequence_length == tsp.sequence_length
        t1 = ts.dump_tables()
        # We need to sort and squash the edges in the new format because it
        # has gone through an edgesets representation. Simplest way to do this
        # is to call simplify.
        if simplify:
            t2 = tsp.simplify().tables
        else:
            t2 = tsp.tables
        assert t1.nodes == t2.nodes
        assert t1.edges == t2.edges
        assert t1.sites == t2.sites
        # The old formats can't represent mutation times so null them out.
        t1.mutations.time = np.full_like(t1.mutations.time, tskit.UNKNOWN_TIME)
        assert t1.mutations == t2.mutations

    def verify_round_trip(self, ts, version):
        tskit.dump_legacy(ts, self.temp_file, version=version)
        tsp = tskit.load_legacy(self.temp_file)
        simplify = version < 10
        self.verify_tree_sequences_equal(ts, tsp, simplify=simplify)
        tsp.dump(self.temp_file)
        tsp = tskit.load(self.temp_file)
        self.verify_tree_sequences_equal(ts, tsp, simplify=simplify)
        for provenance in tsp.provenances():
            tskit.validate_provenance(json.loads(provenance.record))

    def verify_round_trip_no_legacy(self, ts):
        ts.dump(self.temp_file)
        tsp = tskit.load(self.temp_file)
        self.verify_tree_sequences_equal(ts, tsp, simplify=False)
        for provenance in tsp.provenances():
            tskit.validate_provenance(json.loads(provenance.record))

    def verify_malformed_json_v2(self, ts, group_name, attr, bad_json):
        tskit.dump_legacy(ts, self.temp_file, 2)
        # Write some bad JSON to the provenance string.
        root = h5py.File(self.temp_file, "r+")
        group = root[group_name]
        group.attrs[attr] = bad_json
        root.close()
        tsp = tskit.load_legacy(self.temp_file)
        self.verify_tree_sequences_equal(ts, tsp)

    def test_malformed_json_v2(self):
        ts = multi_locus_with_mutation_example()
        for group_name in ["trees", "mutations"]:
            for attr in ["environment", "parameters"]:
                for bad_json in ["", "{", "{},"]:
                    self.verify_malformed_json_v2(ts, group_name, attr, bad_json)

    def test_single_locus_no_mutation(self):
        self.verify_round_trip(single_locus_no_mutation_example(), 2)
        self.verify_round_trip(single_locus_no_mutation_example(), 3)
        self.verify_round_trip(single_locus_no_mutation_example(), 10)

    def test_single_locus_with_mutation(self):
        self.verify_round_trip(single_locus_with_mutation_example(), 2)
        self.verify_round_trip(single_locus_with_mutation_example(), 3)
        self.verify_round_trip(single_locus_with_mutation_example(), 10)

    def test_multi_locus_with_mutation(self):
        self.verify_round_trip(multi_locus_with_mutation_example(), 2)
        self.verify_round_trip(multi_locus_with_mutation_example(), 3)
        self.verify_round_trip(multi_locus_with_mutation_example(), 10)

    def test_migration_example(self):
        self.verify_round_trip(migration_example(), 2)
        self.verify_round_trip(migration_example(), 3)
        self.verify_round_trip(migration_example(), 10)

    def test_bottleneck_example(self):
        self.verify_round_trip(migration_example(), 3)
        self.verify_round_trip(migration_example(), 10)

    def test_no_provenance(self):
        self.verify_round_trip(no_provenance_example(), 10)

    def test_provenance_timestamp_only(self):
        self.verify_round_trip(provenance_timestamp_only_example(), 10)

    def test_recurrent_mutation_example(self):
        ts = recurrent_mutation_example()
        for version in [2, 3]:
            with pytest.raises(ValueError):
                tskit.dump_legacy(ts, self.temp_file, version)
        self.verify_round_trip(ts, 10)

    def test_general_mutation_example(self):
        ts = general_mutation_example()
        for version in [2, 3]:
            with pytest.raises(ValueError):
                tskit.dump_legacy(ts, self.temp_file, version)
        self.verify_round_trip(ts, 10)

    def test_node_metadata_example(self):
        self.verify_round_trip(node_metadata_example(), 10)

    def test_site_metadata_example(self):
        self.verify_round_trip(site_metadata_example(), 10)

    def test_mutation_metadata_example(self):
        self.verify_round_trip(mutation_metadata_example(), 10)

    def test_migration_metadata_example(self):
        self.verify_round_trip(migration_metadata_example(), 10)

    def test_edge_metadata_example(self):
        # metadata for edges was introduced
        self.verify_round_trip_no_legacy(edge_metadata_example())

    def test_multichar_mutation_example(self):
        self.verify_round_trip(multichar_mutation_example(), 10)

    def test_empty_file(self):
        tables = tskit.TableCollection(sequence_length=3)
        self.verify_round_trip(tables.tree_sequence(), 10)

    def test_zero_edges(self):
        tables = tskit.TableCollection(sequence_length=3)
        tables.nodes.add_row(time=0)
        self.verify_round_trip(tables.tree_sequence(), 10)

    def test_v2_no_samples(self):
        ts = multi_locus_with_mutation_example()
        tskit.dump_legacy(ts, self.temp_file, version=2)
        root = h5py.File(self.temp_file, "r+")
        del root["samples"]
        root.close()
        tsp = tskit.load_legacy(self.temp_file)
        self.verify_tree_sequences_equal(ts, tsp)

    def test_duplicate_mutation_positions_single_value(self):
        ts = multi_locus_with_mutation_example()
        for version in [2, 3]:
            tskit.dump_legacy(ts, self.temp_file, version=version)
            root = h5py.File(self.temp_file, "r+")
            root["mutations/position"][:] = 0
            root.close()
            with pytest.raises(tskit.DuplicatePositionsError):
                tskit.load_legacy(self.temp_file)
            tsp = tskit.load_legacy(self.temp_file, remove_duplicate_positions=True)
            assert tsp.num_sites == 1
            sites = list(tsp.sites())
            assert sites[0].position == 0

    def test_duplicate_mutation_positions(self):
        ts = multi_locus_with_mutation_example()
        for version in [2, 3]:
            tskit.dump_legacy(ts, self.temp_file, version=version)
            root = h5py.File(self.temp_file, "r+")
            position = np.array(root["mutations/position"])
            position[0] = position[1]
            root["mutations/position"][:] = position
            root.close()
            with pytest.raises(tskit.DuplicatePositionsError):
                tskit.load_legacy(self.temp_file)
            tsp = tskit.load_legacy(self.temp_file, remove_duplicate_positions=True)
            assert tsp.num_sites == position.shape[0] - 1
            position_after = list(s.position for s in tsp.sites())
            assert list(position[1:]) == position_after


class TestErrors(TestFileFormat):
    """
    Test various API errors.
    """

    def test_v2_non_binary_records(self):
        demographic_events = [
            msprime.SimpleBottleneck(time=0.01, population=0, proportion=1)
        ]
        ts = msprime.simulate(
            sample_size=10, demographic_events=demographic_events, random_seed=1
        )
        with pytest.raises(ValueError):
            tskit.dump_legacy(ts, self.temp_file, 2)

    def test_unsupported_version(self):
        ts = msprime.simulate(10)
        with pytest.raises(ValueError):
            tskit.dump_legacy(ts, self.temp_file, version=4)
        # Cannot read current files.
        ts.dump(self.temp_file)
        # Catch Exception here because h5py throws different exceptions on py2 and py3
        with pytest.raises(Exception):  # noqa B017
            tskit.load_legacy(self.temp_file)

    def test_no_version_number(self):
        root = h5py.File(self.temp_file, "w")
        root.attrs["x"] = 0
        root.close()
        with pytest.raises(ValueError):
            tskit.load_legacy(self.temp_file)

    def test_unknown_legacy_version(self):
        root = h5py.File(self.temp_file, "w")
        root.attrs["format_version"] = (1024, 0)  # Arbitrary unknown version
        root.close()
        with pytest.raises(ValueError):
            tskit.load_legacy(self.temp_file)

    def test_no_h5py(self):
        ts = msprime.simulate(10)
        path = os.path.join(test_data_dir, "hdf5-formats", "msprime-0.3.0_v2.0.hdf5")
        msg = (
            "Legacy formats require h5py. Install via `pip install h5py` or"
            " `conda install h5py`"
        )
        with mock.patch.dict(sys.modules, {"h5py": None}):
            with pytest.raises(ImportError, match=msg):
                tskit.load_legacy(path)
            with pytest.raises(ImportError, match=msg):
                tskit.dump_legacy(ts, path)

    def test_tszip_file(self):
        ts = msprime.simulate(5)
        tszip.compress(ts, self.temp_file)
        with pytest.raises(tskit.FileFormatError, match="appears to be in zip format"):
            tskit.load(self.temp_file)
        with pytest.raises(tskit.FileFormatError, match="appears to be in zip format"):
            tskit.TableCollection.load(self.temp_file)


class TestDumpFormat(TestFileFormat):
    """
    Tests on the on-disk file format.
    """

    def verify_keys(self, ts):
        keys = [
            "edges/child",
            "edges/left",
            "edges/metadata",
            "edges/metadata_offset",
            "edges/metadata_schema",
            "edges/parent",
            "edges/right",
            "format/name",
            "format/version",
            "indexes/edge_insertion_order",
            "indexes/edge_removal_order",
            "individuals/flags",
            "individuals/location",
            "individuals/location_offset",
            "individuals/metadata",
            "individuals/metadata_offset",
            "individuals/metadata_schema",
            "individuals/parents",
            "individuals/parents_offset",
            "metadata",
            "metadata_schema",
            "migrations/dest",
            "migrations/left",
            "migrations/metadata",
            "migrations/metadata_offset",
            "migrations/metadata_schema",
            "migrations/node",
            "migrations/right",
            "migrations/source",
            "migrations/time",
            "mutations/derived_state",
            "mutations/derived_state_offset",
            "mutations/metadata",
            "mutations/metadata_offset",
            "mutations/metadata_schema",
            "mutations/node",
            "mutations/parent",
            "mutations/site",
            "mutations/time",
            "nodes/flags",
            "nodes/individual",
            "nodes/metadata",
            "nodes/metadata_offset",
            "nodes/metadata_schema",
            "nodes/population",
            "nodes/time",
            "populations/metadata",
            "populations/metadata_offset",
            "populations/metadata_schema",
            "provenances/record",
            "provenances/record_offset",
            "provenances/timestamp",
            "provenances/timestamp_offset",
            "sequence_length",
            "sites/ancestral_state",
            "sites/ancestral_state_offset",
            "sites/metadata",
            "sites/metadata_offset",
            "sites/metadata_schema",
            "sites/position",
            "time_units",
            "uuid",
        ]
        ts.dump(self.temp_file)
        store = kastore.load(self.temp_file)
        assert sorted(list(store.keys())) == keys

    def verify_uuid(self, ts, uuid):
        assert len(uuid) == 36
        # Check that the UUID is well-formed.
        parsed = _uuid.UUID("{" + uuid + "}")
        assert str(parsed) == uuid
        assert uuid == ts.file_uuid

    def verify_dump_format(self, ts):
        ts.dump(self.temp_file)
        assert os.path.exists(self.temp_file)
        assert os.path.getsize(self.temp_file) > 0
        self.verify_keys(ts)

        store = kastore.load(self.temp_file)
        # Check the basic root attributes
        format_name = store["format/name"]
        assert np.array_equal(
            np.array(bytearray(b"tskit.trees"), dtype=np.int8), format_name
        )
        format_version = store["format/version"]
        assert format_version[0] == CURRENT_FILE_MAJOR
        assert format_version[1] == CURRENT_FILE_MINOR
        assert ts.sequence_length == store["sequence_length"][0]
        assert repr(ts.metadata_schema) == "".join(store["metadata_schema"].astype("U"))

        # Load another copy from file so we can check the uuid.
        other_ts = tskit.load(self.temp_file)
        self.verify_uuid(other_ts, store["uuid"].tobytes().decode())

        tables = ts.tables

        assert np.array_equal(tables.metadata, b"".join(store["metadata"]))
        assert np.array_equal(tables.individuals.flags, store["individuals/flags"])
        assert np.array_equal(
            tables.individuals.location, store["individuals/location"]
        )
        assert np.array_equal(
            tables.individuals.location_offset, store["individuals/location_offset"]
        )
        assert np.array_equal(tables.individuals.parents, store["individuals/parents"])
        assert np.array_equal(
            tables.individuals.parents_offset, store["individuals/parents_offset"]
        )
        assert np.array_equal(
            tables.individuals.metadata, store["individuals/metadata"]
        )
        assert np.array_equal(
            tables.individuals.metadata_offset, store["individuals/metadata_offset"]
        )
        assert repr(tables.individuals.metadata_schema) == "".join(
            store["individuals/metadata_schema"].astype("U")
        )

        assert np.array_equal(tables.nodes.flags, store["nodes/flags"])
        assert np.array_equal(tables.nodes.time, store["nodes/time"])
        assert np.array_equal(tables.nodes.population, store["nodes/population"])
        assert np.array_equal(tables.nodes.individual, store["nodes/individual"])
        assert np.array_equal(tables.nodes.metadata, store["nodes/metadata"])
        assert np.array_equal(
            tables.nodes.metadata_offset, store["nodes/metadata_offset"]
        )
        assert repr(tables.nodes.metadata_schema) == "".join(
            store["nodes/metadata_schema"].astype("U")
        )

        assert np.array_equal(tables.edges.left, store["edges/left"])
        assert np.array_equal(tables.edges.right, store["edges/right"])
        assert np.array_equal(tables.edges.parent, store["edges/parent"])
        assert np.array_equal(tables.edges.child, store["edges/child"])
        assert np.array_equal(tables.edges.metadata, store["edges/metadata"])
        assert np.array_equal(
            tables.edges.metadata_offset, store["edges/metadata_offset"]
        )
        assert repr(tables.edges.metadata_schema) == "".join(
            store["edges/metadata_schema"].astype("U")
        )

        left = tables.edges.left
        right = tables.edges.right
        parent = tables.edges.parent
        child = tables.edges.child
        time = tables.nodes.time
        in_order = sorted(
            range(ts.num_edges),
            key=lambda j: (left[j], time[parent[j]], parent[j], child[j]),
        )
        out_order = sorted(
            range(ts.num_edges),
            key=lambda j: (right[j], -time[parent[j]], -parent[j], -child[j]),
        )
        assert np.array_equal(
            np.array(in_order, dtype=np.int32),
            store["indexes/edge_insertion_order"],
        )
        assert np.array_equal(
            np.array(out_order, dtype=np.int32), store["indexes/edge_removal_order"]
        )

        assert np.array_equal(tables.migrations.left, store["migrations/left"])
        assert np.array_equal(tables.migrations.right, store["migrations/right"])
        assert np.array_equal(tables.migrations.node, store["migrations/node"])
        assert np.array_equal(tables.migrations.source, store["migrations/source"])
        assert np.array_equal(tables.migrations.dest, store["migrations/dest"])
        assert np.array_equal(tables.migrations.time, store["migrations/time"])
        assert np.array_equal(tables.migrations.metadata, store["migrations/metadata"])
        assert np.array_equal(
            tables.migrations.metadata_offset, store["migrations/metadata_offset"]
        )
        assert repr(tables.migrations.metadata_schema) == "".join(
            store["migrations/metadata_schema"].astype("U")
        )

        assert np.array_equal(tables.sites.position, store["sites/position"])
        assert np.array_equal(
            tables.sites.ancestral_state, store["sites/ancestral_state"]
        )
        assert np.array_equal(
            tables.sites.ancestral_state_offset,
            store["sites/ancestral_state_offset"],
        )
        assert np.array_equal(tables.sites.metadata, store["sites/metadata"])
        assert np.array_equal(
            tables.sites.metadata_offset, store["sites/metadata_offset"]
        )
        assert repr(tables.sites.metadata_schema) == "".join(
            store["sites/metadata_schema"].astype("U")
        )

        assert np.array_equal(tables.mutations.site, store["mutations/site"])
        assert np.array_equal(tables.mutations.node, store["mutations/node"])
        # Default mutation time is a NaN value so we want to check for
        # bit equality, not numeric equality
        assert tables.mutations.time.tobytes() == store["mutations/time"].tobytes()
        assert np.array_equal(tables.mutations.parent, store["mutations/parent"])
        assert np.array_equal(
            tables.mutations.derived_state, store["mutations/derived_state"]
        )
        assert np.array_equal(
            tables.mutations.derived_state_offset,
            store["mutations/derived_state_offset"],
        )
        assert np.array_equal(tables.mutations.metadata, store["mutations/metadata"])
        assert np.array_equal(
            tables.mutations.metadata_offset, store["mutations/metadata_offset"]
        )
        assert repr(tables.mutations.metadata_schema) == "".join(
            store["mutations/metadata_schema"].astype("U")
        )

        assert np.array_equal(
            tables.populations.metadata, store["populations/metadata"]
        )
        assert np.array_equal(
            tables.populations.metadata_offset, store["populations/metadata_offset"]
        )
        assert repr(tables.populations.metadata_schema) == "".join(
            store["populations/metadata_schema"].astype("U")
        )

        assert np.array_equal(tables.provenances.record, store["provenances/record"])
        assert np.array_equal(
            tables.provenances.record_offset, store["provenances/record_offset"]
        )
        assert np.array_equal(
            tables.provenances.timestamp, store["provenances/timestamp"]
        )
        assert np.array_equal(
            tables.provenances.timestamp_offset,
            store["provenances/timestamp_offset"],
        )

        store.close()

    def test_single_locus_no_mutation(self):
        self.verify_dump_format(single_locus_no_mutation_example())

    def test_single_locus_with_mutation(self):
        self.verify_dump_format(single_locus_with_mutation_example())

    def test_multi_locus_with_mutation(self):
        self.verify_dump_format(multi_locus_with_mutation_example())

    def test_migration_example(self):
        self.verify_dump_format(migration_example())

    def test_bottleneck_example(self):
        self.verify_dump_format(bottleneck_example())

    def test_historical_sample_example(self):
        self.verify_dump_format(historical_sample_example())

    def test_node_metadata_example(self):
        self.verify_dump_format(node_metadata_example())

    def test_edge_metadata_example(self):
        self.verify_dump_format(edge_metadata_example())

    def test_site_metadata_example(self):
        self.verify_dump_format(site_metadata_example())

    def test_mutation_metadata_example(self):
        self.verify_dump_format(mutation_metadata_example())

    def test_migration_metadata_example(self):
        self.verify_dump_format(migration_metadata_example())

    def test_general_mutation_example(self):
        self.verify_dump_format(general_mutation_example())

    def test_multichar_mutation_example(self):
        self.verify_dump_format(multichar_mutation_example())


class TestUuid(TestFileFormat):
    """
    Basic tests for the UUID generation.
    """

    def test_different_files_same_ts(self):
        ts = msprime.simulate(10)
        uuids = []
        for _ in range(10):
            ts.dump(self.temp_file)
            with kastore.load(self.temp_file) as store:
                uuids.append(store["uuid"].tobytes().decode())
        assert len(uuids) == len(set(uuids))


class TestOptionalColumns(TestFileFormat):
    """
    Checks that optional columns in the file format are correctly handled.
    """

    def test_empty_edge_metadata(self):
        ts1 = migration_example()
        ts1.dump(self.temp_file)
        ts2 = tskit.load(self.temp_file)
        assert ts1.tables == ts2.tables
        assert len(ts1.tables.edges.metadata) == 0

        with kastore.load(self.temp_file) as store:
            all_data = dict(store)
        del all_data["edges/metadata"]
        del all_data["edges/metadata_offset"]
        kastore.dump(all_data, self.temp_file)
        ts3 = tskit.load(self.temp_file)
        assert ts1.tables == ts3.tables

    def test_empty_migration_metadata(self):
        ts1 = migration_example()
        ts1.dump(self.temp_file)
        ts2 = tskit.load(self.temp_file)
        assert ts1.tables == ts2.tables
        assert len(ts1.tables.migrations.metadata) == 0

        with kastore.load(self.temp_file) as store:
            all_data = dict(store)
        del all_data["migrations/metadata"]
        del all_data["migrations/metadata_offset"]
        kastore.dump(all_data, self.temp_file)
        ts3 = tskit.load(self.temp_file)
        assert ts1.tables == ts3.tables

    def test_empty_mutation_time(self):
        ts1 = migration_example()
        ts1.dump(self.temp_file)
        ts2 = tskit.load(self.temp_file)
        assert ts1.tables == ts2.tables
        assert len(ts1.tables.mutations.metadata) == 0
        with kastore.load(self.temp_file) as store:
            all_data = dict(store)
        del all_data["mutations/time"]
        kastore.dump(all_data, self.temp_file)
        ts3 = tskit.load(self.temp_file)
        # Null out the time column
        t1 = ts1.dump_tables()
        t1.mutations.time = np.full_like(t1.mutations.time, tskit.UNKNOWN_TIME)
        t1.assert_equals(ts3.tables)

    def test_empty_individual_parents(self):
        ts1 = migration_example()
        ts1.dump(self.temp_file)
        ts2 = tskit.load(self.temp_file)
        assert ts1.tables == ts2.tables
        assert len(ts1.tables.individuals.parents) > 0
        with kastore.load(self.temp_file) as store:
            all_data = dict(store)
        del all_data["individuals/parents"]
        del all_data["individuals/parents_offset"]
        kastore.dump(all_data, self.temp_file)
        ts3 = tskit.load(self.temp_file)
        tables = ts1.dump_tables()
        tables.individuals.packset_parents(
            [
                [],
            ]
            * tables.individuals.num_rows
        )
        tables.assert_equals(ts3.tables)


class TestReferenceSequence:
    def test_fixture_has_reference_sequence(self, ts_fixture):
        assert ts_fixture.has_reference_sequence()

    def test_round_trip(self, ts_fixture, tmp_path):
        ts1 = ts_fixture
        temp_file = tmp_path / "tmp.trees"
        ts1.dump(temp_file)
        ts2 = tskit.load(temp_file)
        ts1.tables.assert_equals(ts2.tables)

    def test_no_reference_sequence(self, ts_fixture, tmp_path):
        ts1 = ts_fixture
        temp_file = tmp_path / "tmp.trees"
        ts1.dump(temp_file)
        with kastore.load(temp_file) as store:
            all_data = dict(store)
        del all_data["reference_sequence/metadata_schema"]
        del all_data["reference_sequence/metadata"]
        del all_data["reference_sequence/data"]
        del all_data["reference_sequence/url"]
        for key in all_data.keys():
            assert not key.startswith("reference_sequence")
        kastore.dump(all_data, temp_file)
        ts2 = tskit.load(temp_file)
        assert not ts2.has_reference_sequence()
        tables = ts2.dump_tables()
        tables.reference_sequence = ts1.reference_sequence
        tables.assert_equals(ts1.tables)

    @pytest.mark.parametrize("attr", ["data", "url"])
    def test_missing_attr(self, ts_fixture, tmp_path, attr):
        ts1 = ts_fixture
        temp_file = tmp_path / "tmp.trees"
        ts1.dump(temp_file)
        with kastore.load(temp_file) as store:
            all_data = dict(store)
        del all_data[f"reference_sequence/{attr}"]
        kastore.dump(all_data, temp_file)
        ts2 = tskit.load(temp_file)
        assert ts2.has_reference_sequence
        assert getattr(ts2.reference_sequence, attr) == ""

    def test_missing_metadata(self, ts_fixture, tmp_path):
        ts1 = ts_fixture
        temp_file = tmp_path / "tmp.trees"
        ts1.dump(temp_file)
        with kastore.load(temp_file) as store:
            all_data = dict(store)
        del all_data["reference_sequence/metadata"]
        kastore.dump(all_data, temp_file)
        ts2 = tskit.load(temp_file)
        assert ts2.has_reference_sequence
        assert ts2.reference_sequence.metadata_bytes == b""

    def test_missing_metadata_schema(self, ts_fixture, tmp_path):
        ts1 = ts_fixture
        temp_file = tmp_path / "tmp.trees"
        ts1.dump(temp_file)
        with kastore.load(temp_file) as store:
            all_data = dict(store)
        del all_data["reference_sequence/metadata_schema"]
        kastore.dump(all_data, temp_file)
        ts2 = tskit.load(temp_file)
        assert ts2.has_reference_sequence
        assert repr(ts2.reference_sequence.metadata_schema) == ""


class TestFileFormatErrors(TestFileFormat):
    """
    Tests for errors in the HDF5 format.
    """

    def verify_missing_fields(self, ts):
        ts.dump(self.temp_file)
        with kastore.load(self.temp_file) as store:
            all_data = dict(store)
        for key in all_data.keys():
            # We skip these keys as they are optional
            if "metadata_schema" not in key and key not in [
                "metadata",
                "time_units",
                "mutations/time",
            ]:
                data = dict(all_data)
                del data[key]
                kastore.dump(data, self.temp_file)
                with pytest.raises(
                    (exceptions.FileFormatError, exceptions.LibraryError)
                ):
                    tskit.load(self.temp_file)

    def test_missing_fields(self):
        self.verify_missing_fields(migration_example())

    def verify_equal_length_columns(self, ts, table):
        ts.dump(self.temp_file)
        with kastore.load(self.temp_file) as store:
            all_data = dict(store)
        table_cols = [
            colname for colname in all_data.keys() if colname.startswith(table)
        ]
        # Remove all the 'offset' columns
        for col in list(table_cols):
            if col.endswith("_offset"):
                main_col = col[: col.index("_offset")]
                table_cols.remove(main_col)
                table_cols.remove(col)
            if "metadata_schema" in col:
                table_cols.remove(col)
        # Remaining columns should all be the same length
        for col in table_cols:
            for bad_val in [[], all_data[col][:-1]]:
                data = dict(all_data)
                data[col] = bad_val
                kastore.dump(data, self.temp_file)
                with pytest.raises(exceptions.FileFormatError):
                    tskit.load(self.temp_file)

    def test_equal_length_columns(self):
        ts = migration_example()
        for table in ["nodes", "edges", "migrations", "sites", "mutations"]:
            self.verify_equal_length_columns(ts, table)

    def verify_offset_columns(self, ts):
        ts.dump(self.temp_file)
        with kastore.load(self.temp_file) as store:
            all_data = dict(store)
        offset_col_pairs = []
        for col in all_data.keys():
            if col.endswith("_offset"):
                main_col = col[: col.index("_offset")]
                offset_col_pairs.append((main_col, col))
        for col, offset_col in offset_col_pairs:
            num_rows = len(all_data[offset_col]) - 1
            data = dict(all_data)
            # Check bad lengths of the offset col
            for bad_col_length in [[], range(2 * num_rows)]:
                data[offset_col] = bad_col_length
                kastore.dump(data, self.temp_file)
                with pytest.raises(exceptions.FileFormatError):
                    tskit.load(self.temp_file)

            # Check for a bad offset
            data = dict(all_data)
            original_offset = data[offset_col]
            original_col = data[col]
            data[offset_col] = np.zeros_like(original_offset)
            data[col] = np.zeros(10, dtype=original_col.dtype)
            kastore.dump(data, self.temp_file)
            with pytest.raises(exceptions.LibraryError):
                tskit.load(self.temp_file)

    def test_offset_columns(self):
        ts = migration_example()
        self.verify_offset_columns(ts)

    def test_index_columns(self):
        ts = migration_example()
        ts.dump(self.temp_file)
        with kastore.load(self.temp_file) as store:
            all_data = dict(store)

        edge_removal_order = "indexes/edge_removal_order"
        edge_insertion_order = "indexes/edge_insertion_order"

        data = dict(all_data)
        del data[edge_removal_order]
        del data[edge_insertion_order]
        kastore.dump(data, self.temp_file)
        with pytest.raises(exceptions.LibraryError):
            tskit.load(self.temp_file)

        data = dict(all_data)
        del data[edge_removal_order]
        kastore.dump(data, self.temp_file)
        with pytest.raises(exceptions.LibraryError):
            tskit.load(self.temp_file)

        data = dict(all_data)
        del data[edge_insertion_order]
        kastore.dump(data, self.temp_file)
        with pytest.raises(exceptions.LibraryError):
            tskit.load(self.temp_file)

        data = dict(all_data)
        data[edge_insertion_order] = data[edge_insertion_order][:1]
        kastore.dump(data, self.temp_file)
        with pytest.raises(exceptions.FileFormatError):
            tskit.load(self.temp_file)

        data = dict(all_data)
        data[edge_removal_order] = data[edge_removal_order][:1]
        kastore.dump(data, self.temp_file)
        with pytest.raises(exceptions.FileFormatError):
            tskit.load(self.temp_file)

    def test_load_empty_kastore(self):
        kastore.dump({}, self.temp_file)
        with pytest.raises(exceptions.LibraryError):
            tskit.load(self.temp_file)

    def test_load_non_tskit_hdf5(self):
        with h5py.File(self.temp_file, "w") as root:
            root["x"] = np.zeros(10)
        with pytest.raises(exceptions.FileFormatError):
            tskit.load(self.temp_file)

    def test_old_version_load_error(self):
        ts = msprime.simulate(10, random_seed=1)
        for bad_version in [(0, 1), (0, 8), (2, 0), (CURRENT_FILE_MAJOR - 1, 0)]:
            ts.dump(self.temp_file)
            with kastore.load(self.temp_file) as store:
                data = dict(store)
            data["format/version"] = np.array(bad_version, dtype=np.uint32)
            kastore.dump(data, self.temp_file)
            with pytest.raises(tskit.VersionTooOldError):
                tskit.load(self.temp_file)

    def test_new_version_load_error(self):
        ts = msprime.simulate(10, random_seed=1)
        for bad_version in [(CURRENT_FILE_MAJOR + j, 0) for j in range(1, 5)]:
            ts.dump(self.temp_file)
            with kastore.load(self.temp_file) as store:
                data = dict(store)
            data["format/version"] = np.array(bad_version, dtype=np.uint32)
            kastore.dump(data, self.temp_file)
            with pytest.raises(tskit.VersionTooNewError):
                tskit.load(self.temp_file)

    def test_format_name_error(self):
        ts = msprime.simulate(10)
        for bad_name in ["tskit.tree", "tskit.treesAndOther", "", "x" * 100]:
            ts.dump(self.temp_file)
            with kastore.load(self.temp_file) as store:
                data = dict(store)
            data["format/name"] = np.array(bytearray(bad_name.encode()), dtype=np.int8)
            kastore.dump(data, self.temp_file)
            with pytest.raises(exceptions.FileFormatError):
                tskit.load(self.temp_file)

    def test_load_bad_formats(self):
        # try loading a bunch of files in various formats.
        # First, check the empty file.
        with pytest.raises(EOFError):
            tskit.load(self.temp_file)
        # Now some ascii text
        with open(self.temp_file, "wb") as f:
            f.write(b"Some ASCII text")
        with pytest.raises(exceptions.FileFormatError):
            tskit.load(self.temp_file)
        # Now write 8k of random bytes
        with open(self.temp_file, "wb") as f:
            f.write(os.urandom(8192))
        with pytest.raises(exceptions.FileFormatError):
            tskit.load(self.temp_file)

    def test_load_bad_formats_fileobj(self):
        def load():
            with open(self.temp_file, "rb") as f:
                tskit.load(f)

        with pytest.raises(EOFError):
            load()
        with open(self.temp_file, "wb") as f:
            f.write(b"Some ASCII text")
        with pytest.raises(exceptions.FileFormatError):
            load()


def assert_tables_empty(tables):
    for table in tables.table_name_map.values():
        assert len(table) == 0


class TestSkipTables:
    """
    Test `skip_tables` flag to TreeSequence.load() and TableCollection.load().
    """

    def test_ts_read_path_interface(self, tmp_path, ts_fixture):
        # Check the fixture has metadata and a schema
        assert ts_fixture.metadata_schema is not None
        assert len(ts_fixture.metadata) > 0
        save_path = tmp_path / "tmp.trees"
        ts_fixture.dump(save_path)
        ts_no_tables = tskit.load(save_path, skip_tables=True)
        assert not ts_no_tables.equals(ts_fixture)
        assert ts_no_tables.equals(ts_fixture, ignore_tables=True)
        assert_tables_empty(ts_no_tables.tables)

    def test_ts_read_one_stream(self, tmp_path, ts_fixture):
        save_path = tmp_path / "tmp.trees"
        ts_fixture.dump(save_path)
        with open(save_path, "rb") as f:
            ts_no_tables = tskit.load(f, skip_tables=True)
        assert not ts_no_tables.equals(ts_fixture)
        assert ts_no_tables.equals(ts_fixture, ignore_tables=True)
        assert_tables_empty(ts_no_tables.tables)

    def test_ts_twofile_stream_noskip(self, tmp_path, ts_fixture):
        save_path = tmp_path / "tmp.trees"
        with open(save_path, "wb") as f:
            ts_fixture.dump(f)
            ts_fixture.dump(f)
        with open(save_path, "rb") as f:
            ts1 = tskit.load(f)
            ts2 = tskit.load(f)
        assert ts_fixture.equals(ts1)
        assert ts_fixture.equals(ts2)

    def test_ts_twofile_stream_fails(self, tmp_path, ts_fixture):
        # We can't skip_tables while reading from a stream
        save_path = tmp_path / "tmp.trees"
        with open(save_path, "wb") as f:
            ts_fixture.dump(f)
            ts_fixture.dump(f)
        with open(save_path, "rb") as f:
            tskit.load(f, skip_tables=True)
            with pytest.raises(exceptions.FileFormatError):
                tskit.load(f)

    def test_table_collection_load_path(self, tmp_path, ts_fixture):
        save_path = tmp_path / "tmp.trees"
        ts_fixture.dump(save_path)
        tables_skipped = tskit.TableCollection.load(save_path, skip_tables=True)
        tables = ts_fixture.tables
        assert not tables_skipped.equals(tables)
        assert tables_skipped.equals(tables, ignore_tables=True)
        assert_tables_empty(tables_skipped)

    def test_table_collection_load_stream(self, tmp_path, ts_fixture):
        save_path = tmp_path / "tmp.trees"
        ts_fixture.dump(save_path)
        with open(save_path, "rb") as f:
            tables_skipped = tskit.TableCollection.load(f, skip_tables=True)
        tables = ts_fixture.tables
        assert not tables_skipped.equals(tables)
        assert tables_skipped.equals(tables, ignore_tables=True)
        assert_tables_empty(tables_skipped)


class TestSkipReferenceSequence:
    """
    Test `skip_reference_sequence` flag to TreeSequence.load() and
    TableCollection.load().
    """

    def test_ts_load_path(self, tmp_path, ts_fixture):
        assert ts_fixture.has_reference_sequence()
        save_path = tmp_path / "tmp.trees"
        ts_fixture.dump(save_path)
        ts_no_refseq = tskit.load(save_path, skip_reference_sequence=True)
        assert not ts_no_refseq.equals(ts_fixture)
        assert ts_no_refseq.equals(ts_fixture, ignore_reference_sequence=True)
        assert not ts_no_refseq.has_reference_sequence()

    def test_ts_load_stream(self, tmp_path, ts_fixture):
        save_path = tmp_path / "tmp.trees"
        ts_fixture.dump(save_path)
        with open(save_path, "rb") as f:
            ts_no_refseq = tskit.load(f, skip_reference_sequence=True)
        assert not ts_no_refseq.equals(ts_fixture)
        assert ts_no_refseq.equals(ts_fixture, ignore_reference_sequence=True)
        assert not ts_no_refseq.has_reference_sequence()

    def test_ts_twofile_stream_fails(self, tmp_path, ts_fixture):
        # We can't skip_reference_sequence while reading from a stream
        save_path = tmp_path / "tmp.trees"
        with open(save_path, "wb") as f:
            ts_fixture.dump(f)
            ts_fixture.dump(f)
        with open(save_path, "rb") as f:
            tskit.load(f, skip_reference_sequence=True)
            with pytest.raises(exceptions.FileFormatError):
                tskit.load(f)

    def test_table_collection_load_path(self, tmp_path, ts_fixture):
        save_path = tmp_path / "tmp.trees"
        ts_fixture.dump(save_path)
        tables_no_refseq = tskit.TableCollection.load(
            save_path, skip_reference_sequence=True
        )
        tables = ts_fixture.tables
        assert not tables_no_refseq.equals(tables)
        assert tables_no_refseq.equals(tables, ignore_reference_sequence=True)
        assert not tables_no_refseq.has_reference_sequence()

    def test_table_collection_load_stream(self, tmp_path, ts_fixture):
        save_path = tmp_path / "tmp.trees"
        ts_fixture.dump(save_path)
        with open(save_path, "rb") as f:
            tables_no_refseq = tskit.TableCollection.load(
                f, skip_reference_sequence=True
            )
        tables = ts_fixture.tables
        assert not tables_no_refseq.equals(tables)
        assert tables_no_refseq.equals(tables, ignore_reference_sequence=True)
        assert not tables_no_refseq.has_reference_sequence()


--- ../../tskit/python/tests/__init__.py ---

 
import base64

import tskit
from . import tsutil
from .simplify import *  # NOQA

# TODO remove this code and refactor elsewhere.


class PythonTree:
    """
    Presents the same interface as the Tree object for testing. This
    is tightly coupled with the PythonTreeSequence object below which updates
    the internal structures during iteration.
    """

    def __init__(self, num_nodes):
        self.num_nodes = num_nodes
        self.parent = [tskit.NULL for _ in range(num_nodes)]
        self.left_child = [tskit.NULL for _ in range(num_nodes)]
        self.right_child = [tskit.NULL for _ in range(num_nodes)]
        self.left_sib = [tskit.NULL for _ in range(num_nodes)]
        self.right_sib = [tskit.NULL for _ in range(num_nodes)]
        self.num_children = [0 for _ in range(num_nodes)]
        self.edge = [tskit.NULL for _ in range(num_nodes)]
        self.left = 0
        self.right = 0
        self.index = -1
        self.left_root = -1
        # We need a sites function, so this name is taken.
        self.site_list = []

    @classmethod
    def from_tree(cls, tree):
        ret = PythonTree(tree.tree_sequence.num_nodes)
        ret.left, ret.right = tree.get_interval()
        ret.site_list = list(tree.sites())
        ret.index = tree.get_index()
        ret.left_root = tree.left_root
        ret.tree = tree
        for u in range(ret.num_nodes):
            ret.parent[u] = tree.parent(u)
            ret.left_child[u] = tree.left_child(u)
            ret.right_child[u] = tree.right_child(u)
            ret.left_sib[u] = tree.left_sib(u)
            ret.right_sib[u] = tree.right_sib(u)
            ret.num_children[u] = tree.num_children(u)
            ret.edge[u] = tree.edge(u)
        assert ret == tree
        return ret

    @property
    def roots(self):
        u = self.left_root
        roots = []
        while u != tskit.NULL:
            roots.append(u)
            u = self.right_sib[u]
        return roots

    def children(self, u):
        v = self.left_child[u]
        ret = []
        while v != tskit.NULL:
            ret.append(v)
            v = self.right_sib[v]
        return ret

    def get_interval(self):
        return self.left, self.right

    def get_parent(self, node):
        return self.parent[node]

    def get_children(self, node):
        return self.children[node]

    def get_index(self):
        return self.index

    def get_parent_dict(self):
        d = {
            u: self.parent[u]
            for u in range(self.num_nodes)
            if self.parent[u] != tskit.NULL
        }
        return d

    def sites(self):
        return iter(self.site_list)

    def __eq__(self, other):
        return (
            self.get_parent_dict() == other.get_parent_dict()
            and self.get_interval() == other.get_interval()
            and self.roots == other.roots
            and self.get_index() == other.get_index()
            and list(self.sites()) == list(other.sites())
        )

    def __ne__(self, other):
        return not self.__eq__(other)


class PythonTreeSequence:
    """
    A python implementation of the TreeSequence object.

    TODO this class is of limited use now and should be factored out as
    part of a drive towards more modular versions of the tests currently
    in tests_highlevel.py.
    """

    def __init__(self, tree_sequence, breakpoints=None):
        self._tree_sequence = tree_sequence
        self._sites = []
        # TODO this code here is expressed in terms of the low-level
        # tree sequence for legacy reasons. It probably makes more sense
        # to describe it in terms of the tables now if we want to have an
        # independent implementation.
        ll_ts = self._tree_sequence._ll_tree_sequence

        def make_mutation(id_):
            (
                site,
                node,
                derived_state,
                parent,
                metadata,
                time,
                edge,
            ) = ll_ts.get_mutation(id_)
            return tskit.Mutation(
                id=id_,
                site=site,
                node=node,
                time=time,
                derived_state=derived_state,
                parent=parent,
                metadata=metadata,
                edge=edge,
                metadata_decoder=tskit.metadata.parse_metadata_schema(
                    ll_ts.get_table_metadata_schemas().mutation
                ).decode_row,
            )

        for j in range(tree_sequence.num_sites):
            pos, ancestral_state, ll_mutations, id_, metadata = ll_ts.get_site(j)
            self._sites.append(
                tskit.Site(
                    id=id_,
                    position=pos,
                    ancestral_state=ancestral_state,
                    mutations=[make_mutation(ll_mut) for ll_mut in ll_mutations],
                    metadata=metadata,
                    metadata_decoder=tskit.metadata.parse_metadata_schema(
                        ll_ts.get_table_metadata_schemas().site
                    ).decode_row,
                )
            )

    def trees(self):
        pt = PythonTree(self._tree_sequence.get_num_nodes())
        pt.index = 0
        for (left, right), rtt in tsutil.algorithm_R(self._tree_sequence):
            pt.parent[:] = rtt.parent
            pt.left_child[:] = rtt.left_child
            pt.right_child[:] = rtt.right_child
            pt.left_sib[:] = rtt.left_sib
            pt.right_sib[:] = rtt.right_sib
            pt.num_children[:] = rtt.num_children
            pt.edge[:] = rtt.edge
            pt.left_root = rtt.left_child[-1]
            pt.left = left
            pt.right = right
            # Add in all the sites
            pt.site_list = [
                site for site in self._sites if left <= site.position < right
            ]
            yield pt
            pt.index += 1
        pt.index = -1


class MRCACalculator:
    """
    Class to that allows us to compute the nearest common ancestor of arbitrary
    nodes in an oriented forest.

    This is an implementation of Schieber and Vishkin's nearest common ancestor
    algorithm from TAOCP volume 4A, pg.164-167 [K11]_. Preprocesses the
    input tree into a sideways heap in O(n) time and processes queries for the
    nearest common ancestor between an arbitary pair of nodes in O(1) time.

    :param oriented_forest: the input oriented forest
    :type oriented_forest: list of integers
    """

    LAMBDA = 0

    def __init__(self, oriented_forest):
        # We turn this oriened forest into a 1 based array by adding 1
        # to everything
        converted = [0] + [x + 1 for x in oriented_forest]
        self.__preprocess(converted)

    def __preprocess(self, oriented_forest):
        """
        Preprocess the oriented forest, so that we can answer mrca queries
        in constant time.
        """
        n = len(oriented_forest)
        child = [self.LAMBDA for i in range(n)]
        parent = [self.LAMBDA for i in range(n)]
        sib = [self.LAMBDA for i in range(n)]
        self.__lambda = [0 for i in range(n)]
        self.__pi = [0 for i in range(n)]
        self.__tau = [0 for i in range(n)]
        self.__beta = [0 for i in range(n)]
        self.__alpha = [0 for i in range(n)]
        for u in range(n):
            v = oriented_forest[u]
            sib[u] = child[v]
            child[v] = u
            parent[u] = v
        p = child[self.LAMBDA]
        n = 0
        self.__lambda[0] = -1
        while p != self.LAMBDA:
            notDone = True
            while notDone:
                n += 1
                self.__pi[p] = n
                self.__tau[n] = self.LAMBDA
                self.__lambda[n] = 1 + self.__lambda[n >> 1]
                if child[p] != self.LAMBDA:
                    p = child[p]
                else:
                    notDone = False
            self.__beta[p] = n
            notDone = True
            while notDone:
                self.__tau[self.__beta[p]] = parent[p]
                if sib[p] != self.LAMBDA:
                    p = sib[p]
                    notDone = False
                else:
                    p = parent[p]
                    if p != self.LAMBDA:
                        h = self.__lambda[n & -self.__pi[p]]
                        self.__beta[p] = ((n >> h) | 1) << h
                    else:
                        notDone = False
        # Begin the second traversal
        self.__lambda[0] = self.__lambda[n]
        self.__pi[self.LAMBDA] = 0
        self.__beta[self.LAMBDA] = 0
        self.__alpha[self.LAMBDA] = 0
        p = child[self.LAMBDA]
        while p != self.LAMBDA:
            notDone = True
            while notDone:
                a = self.__alpha[parent[p]] | (self.__beta[p] & -self.__beta[p])
                self.__alpha[p] = a
                if child[p] != self.LAMBDA:
                    p = child[p]
                else:
                    notDone = False
            notDone = True
            while notDone:
                if sib[p] != self.LAMBDA:
                    p = sib[p]
                    notDone = False
                else:
                    p = parent[p]
                    notDone = p != self.LAMBDA

    def get_mrca(self, x, y):
        """
        Returns the most recent common ancestor of the nodes x and y,
        or -1 if the nodes belong to different trees.

        :param x: the first node
        :param y: the second node
        :return: the MRCA of nodes x and y
        """
        # WE need to rescale here because SV expects 1-based arrays.
        return self._sv_mrca(x + 1, y + 1) - 1

    def _sv_mrca(self, x, y):
        if self.__beta[x] <= self.__beta[y]:
            h = self.__lambda[self.__beta[y] & -self.__beta[x]]
        else:
            h = self.__lambda[self.__beta[x] & -self.__beta[y]]
        k = self.__alpha[x] & self.__alpha[y] & -(1 << h)
        h = self.__lambda[k & -k]
        j = ((self.__beta[x] >> h) | 1) << h
        if j == self.__beta[x]:
            xhat = x
        else:
            ell = self.__lambda[self.__alpha[x] & ((1 << h) - 1)]
            xhat = self.__tau[((self.__beta[x] >> ell) | 1) << ell]
        if j == self.__beta[y]:
            yhat = y
        else:
            ell = self.__lambda[self.__alpha[y] & ((1 << h) - 1)]
            yhat = self.__tau[((self.__beta[y] >> ell) | 1) << ell]
        if self.__pi[xhat] <= self.__pi[yhat]:
            z = xhat
        else:
            z = yhat
        return z


def base64_encode(metadata):
    """
    Returns the specified metadata bytes object encoded as an ASCII-safe
    string.
    """
    return base64.b64encode(metadata).decode("utf8")


def cached_example(ts_func):
    """
    Utility decorator to cache the result of a single function call
    returning a tree sequence example.
    """
    cache = None

    def f(*args):
        nonlocal cache
        if cache is None:
            cache = ts_func(*args)
        return cache

    return f


--- ../../tskit/python/tests/test_highlevel.py ---


"""
Test cases for the high level interface to tskit.
"""
import collections
import dataclasses
import decimal
import functools
import inspect
import io
import itertools
import json
import math
import os
import pathlib
import pickle
import platform
import random
import re
import tempfile
import unittest
import uuid as _uuid
import warnings
from xml.etree import ElementTree

import kastore
import msprime
import networkx as nx
import numpy as np
import pytest
from numpy.testing import assert_array_equal

import _tskit
import tests as tests
import tests.simplify as simplify
import tests.tsutil as tsutil
import tskit
import tskit.util as util
from tskit import UNKNOWN_TIME


def traversal_preorder(tree, root=None):
    roots = tree.roots if root is None else [root]
    for node in roots:
        yield node
        for child in tree.children(node):
            yield from traversal_preorder(tree, child)


def traversal_postorder(tree, root=None):
    roots = tree.roots if root is None else [root]
    for node in roots:
        for child in tree.children(node):
            yield from traversal_postorder(tree, child)
        yield node


def traversal_inorder(tree, root=None):
    roots = tree.roots if root is None else [root]
    for node in roots:
        children = list(tree.children(node))
        half = len(children) // 2
        for child in children[:half]:
            yield from traversal_inorder(tree, child)
        yield node
        for child in children[half:]:
            yield from traversal_inorder(tree, child)


def traversal_levelorder(tree, root=None):
    yield from sorted(list(tree.nodes(root)), key=lambda u: tree.depth(u))


def _traversal_minlex_postorder(tree, u):
    """
    For a given input ID u, this function returns a tuple whose first value
    is the minimum leaf node ID under node u, and whose second value is
    a list containing the minlex postorder for the subtree rooted at node u.
    The first value is needed for sorting, and the second value is what
    finally gets returned.
    """
    children = tree.children(u)
    if len(children) > 0:
        children_return = [_traversal_minlex_postorder(tree, c) for c in children]
        # sorts by first value, which is the minimum leaf node ID
        children_return.sort(key=lambda x: x[0])
        minlex_postorder = []
        for _, child_minlex_postorder in children_return:
            minlex_postorder.extend(child_minlex_postorder)
        minlex_postorder.extend([u])
        return (children_return[0][0], minlex_postorder)
    else:
        return (u, [u])


def traversal_minlex_postorder(tree, root=None):
    roots = tree.roots if root is None else [root]
    root_lists = [_traversal_minlex_postorder(tree, node) for node in roots]
    for _, node_list in sorted(root_lists, key=lambda x: x[0]):
        yield from node_list


def traversal_timeasc(tree, root=None):
    yield from sorted(tree.nodes(root), key=lambda u: (tree.time(u), u))


def traversal_timedesc(tree, root=None):
    yield from sorted(tree.nodes(root), key=lambda u: (tree.time(u), u), reverse=True)


traversal_map = {
    "preorder": traversal_preorder,
    "postorder": traversal_postorder,
    "inorder": traversal_inorder,
    "levelorder": traversal_levelorder,
    "breadthfirst": traversal_levelorder,
    "minlex_postorder": traversal_minlex_postorder,
    "timeasc": traversal_timeasc,
    "timedesc": traversal_timedesc,
}


def insert_uniform_mutations(tables, num_mutations, nodes):
    """
    Returns n evenly mutations over the specified list of nodes.
    """
    for j in range(num_mutations):
        tables.sites.add_row(
            position=j * (tables.sequence_length / num_mutations),
            ancestral_state="0",
            metadata=json.dumps({"index": j}).encode(),
        )
        tables.mutations.add_row(
            site=j,
            derived_state="1",
            node=nodes[j % len(nodes)],
            metadata=json.dumps({"index": j}).encode(),
        )


def get_table_collection_copy(tables, sequence_length):
    """
    Returns a copy of the specified table collection with the specified
    sequence length.
    """
    table_dict = tables.asdict()
    table_dict["sequence_length"] = sequence_length
    return tskit.TableCollection.fromdict(table_dict)


def insert_gap(ts, position, length):
    """
    Inserts a gap of the specified size into the specified tree sequence.
    This involves: (1) breaking all edges that intersect with this point;
    and (2) shifting all coordinates greater than this value up by the
    gap length.
    """
    new_edges = []
    for e in ts.edges():
        if e.left < position < e.right:
            new_edges.append([e.left, position, e.parent, e.child])
            new_edges.append([position, e.right, e.parent, e.child])
        else:
            new_edges.append([e.left, e.right, e.parent, e.child])

    # Now shift up all coordinates.
    for e in new_edges:
        # Left coordinates == position get shifted
        if e[0] >= position:
            e[0] += length
        # Right coordinates == position do not get shifted
        if e[1] > position:
            e[1] += length
    tables = ts.dump_tables()
    L = ts.sequence_length + length
    tables = get_table_collection_copy(tables, L)
    tables.edges.clear()
    tables.sites.clear()
    tables.mutations.clear()
    for left, right, parent, child in new_edges:
        tables.edges.add_row(left, right, parent, child)
    tables.sort()
    # Throw in a bunch of mutations over the whole sequence on the samples.
    insert_uniform_mutations(tables, 100, list(ts.samples()))
    return tables.tree_sequence()


@functools.lru_cache
def get_gap_examples(custom_max=None):
    """
    Returns example tree sequences that contain gaps within the list of
    edges.
    """
    ret = []
    if custom_max is None:
        n_list = [20, 10]
    else:
        n_list = [custom_max, custom_max // 2]

    ts = msprime.simulate(n_list[0], random_seed=56, recombination_rate=1)

    assert ts.num_trees > 1

    gap = 0.0125
    for x in [0, 0.1, 0.5, 0.75]:
        ts = insert_gap(ts, x, gap)
        found = False
        for t in ts.trees():
            if t.interval.left == x:
                assert t.interval.right == x + gap
                assert len(t.parent_dict) == 0
                found = True
        assert found
        ret.append((f"gap_{x}", ts))
    # Give an example with a gap at the end.
    ts = msprime.simulate(n_list[1], random_seed=5, recombination_rate=1)
    tables = get_table_collection_copy(ts.dump_tables(), 2)
    tables.sites.clear()
    tables.mutations.clear()
    insert_uniform_mutations(tables, 100, list(ts.samples()))
    ret.append(("gap_at_end", tables.tree_sequence()))
    return ret


@functools.lru_cache
def get_internal_samples_examples():
    """
    Returns example tree sequences with internal samples.
    """
    ret = []
    n = 5
    ts = msprime.simulate(n, random_seed=10, mutation_rate=5)
    assert ts.num_mutations > 0
    tables = ts.dump_tables()
    nodes = tables.nodes
    flags = nodes.flags
    # Set all nodes to be samples.
    flags[:] = tskit.NODE_IS_SAMPLE
    nodes.flags = flags
    ret.append(("all_nodes_samples", tables.tree_sequence()))

    # Set just internal nodes to be samples.
    flags[:] = 0
    flags[n:] = tskit.NODE_IS_SAMPLE
    nodes.flags = flags
    ret.append(("internal_nodes_samples", tables.tree_sequence()))

    # Set a mixture of internal and leaf samples.
    flags[:] = 0
    flags[n // 2 : n + n // 2] = tskit.NODE_IS_SAMPLE
    nodes.flags = flags
    ret.append(("mixed_internal_leaf_samples", tables.tree_sequence()))
    return ret


@functools.lru_cache
def get_decapitated_examples(custom_max=None):
    """
    Returns example tree sequences in which the oldest edges have been removed.
    """
    ret = []
    if custom_max is None:
        n_list = [10, 20]
    else:
        n_list = [custom_max // 2, custom_max]
    ts = msprime.simulate(n_list[0], random_seed=1234)
    # yield ts.decapitate(ts.tables.nodes.time[-1] / 2)
    ts = msprime.simulate(n_list[1], recombination_rate=1, random_seed=1234)
    assert ts.num_trees > 2
    ret.append(("decapitate_recomb", ts.decapitate(ts.tables.nodes.time[-1] / 4)))
    return ret


def get_bottleneck_examples(custom_max=None):
    """
    Returns an iterator of example tree sequences with nonbinary trees.
    """
    bottlenecks = [
        msprime.SimpleBottleneck(0.01, 0, proportion=0.05),
        msprime.SimpleBottleneck(0.02, 0, proportion=0.25),
        msprime.SimpleBottleneck(0.03, 0, proportion=1),
    ]
    if custom_max is None:
        n_list = [3, 10, 100]
    else:
        n_list = [i * custom_max // 3 for i in range(1, 4)]
    for n in n_list:
        ts = msprime.simulate(
            n,
            length=100,
            recombination_rate=1,
            demographic_events=bottlenecks,
            random_seed=n,
        )
        yield (f"bottleneck_n={n}", ts)


def get_back_mutation_examples():
    """
    Returns an iterator of example tree sequences with nonbinary trees.
    """
    ts = msprime.simulate(10, random_seed=1)
    for j in [1, 2, 3]:
        yield tsutil.insert_branch_mutations(ts, mutations_per_branch=j)
    for ts in get_bottleneck_examples():
        yield tsutil.insert_branch_mutations(ts)


def make_example_tree_sequences(custom_max=None):
    yield from get_decapitated_examples(custom_max=custom_max)
    yield from get_gap_examples(custom_max=custom_max)
    yield from get_internal_samples_examples()
    seed = 1
    if custom_max is None:
        n_list = [2, 3, 10, 100]
    else:
        n_list = [i * custom_max // 4 for i in range(1, 5)]
    for n in n_list:
        for m in [1, 2, 32]:
            for rho in [0, 0.1, 0.5]:
                recomb_map = msprime.RecombinationMap.uniform_map(m, rho, num_loci=m)
                ts = msprime.simulate(
                    recombination_map=recomb_map,
                    mutation_rate=0.1,
                    random_seed=seed,
                    population_configurations=[
                        msprime.PopulationConfiguration(n),
                        msprime.PopulationConfiguration(0),
                    ],
                    migration_matrix=[[0, 1], [1, 0]],
                )
                ts = tsutil.insert_random_ploidy_individuals(ts, 4, seed=seed)
                yield (
                    f"n={n}_m={m}_rho={rho}",
                    tsutil.add_random_metadata(ts, seed=seed),
                )
                seed += 1
    for name, ts in get_bottleneck_examples(custom_max=custom_max):
        yield (
            f"{name}_mutated",
            msprime.mutate(
                ts,
                rate=0.1,
                random_seed=seed,
                model=msprime.InfiniteSites(msprime.NUCLEOTIDES),
            ),
        )
    ts = tskit.Tree.generate_balanced(8).tree_sequence
    yield ("rev_node_order", ts.subset(np.arange(ts.num_nodes - 1, -1, -1)))
    ts = msprime.sim_ancestry(
        8, sequence_length=40, recombination_rate=0.1, random_seed=seed
    )
    tables = ts.dump_tables()
    tables.populations.metadata_schema = tskit.MetadataSchema(None)
    ts = tables.tree_sequence()
    assert ts.num_trees > 1
    yield (
        "back_mutations",
        tsutil.insert_branch_mutations(ts, mutations_per_branch=2),
    )
    ts = tsutil.insert_multichar_mutations(ts)
    yield ("multichar", ts)
    yield ("multichar_no_metadata", tsutil.add_random_metadata(ts))
    tables = ts.dump_tables()
    tables.nodes.flags = np.zeros_like(tables.nodes.flags)
    yield ("no_samples", tables.tree_sequence())  # no samples
    tables = ts.dump_tables()
    tables.edges.clear()
    yield ("empty_tree", tables.tree_sequence())  # empty tree
    yield (
        "empty_ts",
        tskit.TableCollection(sequence_length=1).tree_sequence(),
    )  # empty tree seq
    yield ("all_fields", tsutil.all_fields_ts())


_examples = tuple(make_example_tree_sequences(custom_max=None))


def get_example_tree_sequences(pytest_params=True, custom_max=None):
    if pytest_params:
        return [pytest.param(ts, id=name) for name, ts in _examples]
    else:
        return [ts for _, ts in _examples]


def simple_get_pairwise_diversity(haplotypes):
    """
    Returns the value of pi for the specified haplotypes.
    """
    # Very simplistic algorithm...
    n = len(haplotypes)
    pi = 0
    for k in range(n):
        for j in range(k):
            for u, v in zip(haplotypes[j], haplotypes[k]):
                pi += u != v
    return 2 * pi / (n * (n - 1))


def simplify_tree_sequence(ts, samples, filter_sites=True):
    """
    Simple tree-by-tree algorithm to get a simplify of a tree sequence.
    """
    s = simplify.Simplifier(ts, samples, filter_sites=filter_sites)
    return s.simplify()


def oriented_forests(n):
    """
    Implementation of Algorithm O from TAOCP section 7.2.1.6.
    Generates all canonical n-node oriented forests.
    """
    p = [k - 1 for k in range(0, n + 1)]
    k = 1
    while k != 0:
        yield p
        if p[n] > 0:
            p[n] = p[p[n]]
            yield p
        k = n
        while k > 0 and p[k] == 0:
            k -= 1
        if k != 0:
            j = p[k]
            d = k - j
            not_done = True
            while not_done:
                if p[k - d] == p[j]:
                    p[k] = p[j]
                else:
                    p[k] = p[k - d] + d
                if k == n:
                    not_done = False
                else:
                    k += 1


def get_mrca(pi, x, y):
    """
    Returns the most recent common ancestor of nodes x and y in the
    oriented forest pi.
    """
    x_parents = [x]
    j = x
    while j != 0:
        j = pi[j]
        x_parents.append(j)
    y_parents = {y: None}
    j = y
    while j != 0:
        j = pi[j]
        y_parents[j] = None
    # We have the complete list of parents for x and y back to root.
    mrca = 0
    j = 0
    while x_parents[j] not in y_parents:
        j += 1
    mrca = x_parents[j]
    return mrca


def get_samples(ts, time=None, population=None):
    samples = []
    for node in ts.nodes():
        keep = bool(node.is_sample())
        if time is not None:
            if isinstance(time, (int, float)):
                keep &= np.isclose(node.time, time)
            if isinstance(time, (tuple, list, np.ndarray)):
                keep &= node.time >= time[0]
                keep &= node.time < time[1]
        if population is not None:
            keep &= node.population == population
        if keep:
            samples.append(node.id)
    return np.array(samples)


class TestTreeTraversals:
    def test_bad_traversal_order(self, simple_degree2_ts_fixture):
        tree = simple_degree2_ts_fixture.first()
        for bad_order in ["pre", "post", "preorderorder", ("x",), b"preorder"]:
            with pytest.raises(ValueError, match="Traversal order"):
                tree.nodes(order=bad_order)

    @pytest.mark.parametrize("order", list(traversal_map.keys()))
    def test_returned_types(self, order):
        ts = msprime.sim_ancestry(2, random_seed=234)
        tree = ts.first()
        iterator = tree.nodes(order=order)
        assert isinstance(iterator, collections.abc.Iterable)
        lst = list(iterator)
        assert len(lst) > 0
        for u in lst:
            assert isinstance(u, int)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    @pytest.mark.parametrize("order", list(traversal_map.keys()))
    def test_traversals_virtual_root(self, ts, order):
        tree = ts.first()
        node_list2 = list(traversal_map[order](tree, tree.virtual_root))
        node_list1 = list(tree.nodes(tree.virtual_root, order=order))
        assert tree.virtual_root in node_list1
        assert node_list1 == node_list2

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    @pytest.mark.parametrize("order", list(traversal_map.keys()))
    def test_traversals(self, ts, order):
        tree = next(ts.trees())
        traverser = traversal_map[order]
        node_list1 = list(tree.nodes(order=order))
        node_list2 = list(traverser(tree))
        assert node_list1 == node_list2

    def test_binary_example(self):
        t = tskit.Tree.generate_balanced(5)
        #     8
        #  ┏━━┻━┓
        #  ┃    7
        #  ┃  ┏━┻┓
        #  5  ┃  6
        # ┏┻┓ ┃ ┏┻┓
        # 0 1 2 3 4

        def f(node=None, order=None):
            return list(t.nodes(node, order))

        assert f(order="preorder") == [8, 5, 0, 1, 7, 2, 6, 3, 4]
        assert f(order="postorder") == [0, 1, 5, 2, 3, 4, 6, 7, 8]
        assert f(order="inorder") == [0, 5, 1, 8, 2, 7, 3, 6, 4]
        assert f(order="levelorder") == [8, 5, 7, 0, 1, 2, 6, 3, 4]
        assert f(order="breadthfirst") == [8, 5, 7, 0, 1, 2, 6, 3, 4]
        assert f(order="timeasc") == [0, 1, 2, 3, 4, 5, 6, 7, 8]
        assert f(order="timedesc") == [8, 7, 6, 5, 4, 3, 2, 1, 0]
        assert f(order="minlex_postorder") == [0, 1, 5, 2, 3, 4, 6, 7, 8]

        q = t.virtual_root
        assert f(q, order="preorder") == [q, 8, 5, 0, 1, 7, 2, 6, 3, 4]
        assert f(q, order="postorder") == [0, 1, 5, 2, 3, 4, 6, 7, 8, q]
        assert f(q, order="inorder") == [q, 0, 5, 1, 8, 2, 7, 3, 6, 4]
        assert f(q, order="levelorder") == [q, 8, 5, 7, 0, 1, 2, 6, 3, 4]
        assert f(q, order="breadthfirst") == [q, 8, 5, 7, 0, 1, 2, 6, 3, 4]
        assert f(q, order="timeasc") == [0, 1, 2, 3, 4, 5, 6, 7, 8, q]
        assert f(q, order="timedesc") == [q, 8, 7, 6, 5, 4, 3, 2, 1, 0]
        assert f(q, order="minlex_postorder") == [0, 1, 5, 2, 3, 4, 6, 7, 8, q]

        assert f(7, order="preorder") == [7, 2, 6, 3, 4]
        assert f(7, order="postorder") == [2, 3, 4, 6, 7]
        assert f(7, order="inorder") == [2, 7, 3, 6, 4]
        assert f(7, order="levelorder") == [7, 2, 6, 3, 4]
        assert f(7, order="breadthfirst") == [7, 2, 6, 3, 4]
        assert f(7, order="timeasc") == [2, 3, 4, 6, 7]
        assert f(7, order="timedesc") == [7, 6, 4, 3, 2]
        assert f(7, order="minlex_postorder") == [2, 3, 4, 6, 7]

    def test_ternary_example(self):
        t = tskit.Tree.generate_balanced(7, arity=3)
        #      10
        #  ┏━━━┳┻━━━┓
        #  7   8    9
        # ┏┻┓ ┏┻┓ ┏━╋━┓
        # 0 1 2 3 4 5 6

        def f(node=None, order=None):
            return list(t.nodes(node, order))

        assert f(order="preorder") == [10, 7, 0, 1, 8, 2, 3, 9, 4, 5, 6]
        assert f(order="postorder") == [0, 1, 7, 2, 3, 8, 4, 5, 6, 9, 10]
        assert f(order="inorder") == [0, 7, 1, 10, 2, 8, 3, 4, 9, 5, 6]
        assert f(order="levelorder") == [10, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6]
        assert f(order="breadthfirst") == [10, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6]
        assert f(order="timeasc") == list(range(11))
        assert f(order="timedesc") == list(reversed(range(11)))
        assert f(order="minlex_postorder") == [0, 1, 7, 2, 3, 8, 4, 5, 6, 9, 10]

        q = t.virtual_root
        assert f(q, order="preorder") == [q, 10, 7, 0, 1, 8, 2, 3, 9, 4, 5, 6]
        assert f(q, order="postorder") == [0, 1, 7, 2, 3, 8, 4, 5, 6, 9, 10, q]
        assert f(q, order="inorder") == [q, 0, 7, 1, 10, 2, 8, 3, 4, 9, 5, 6]
        assert f(q, order="levelorder") == [q, 10, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6]
        assert f(q, order="breadthfirst") == [q, 10, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6]
        assert f(q, order="timeasc") == list(range(12))
        assert f(q, order="timedesc") == list(reversed(range(12)))
        assert f(q, order="minlex_postorder") == [0, 1, 7, 2, 3, 8, 4, 5, 6, 9, 10, q]

        assert f(9, order="preorder") == [9, 4, 5, 6]
        assert f(9, order="postorder") == [4, 5, 6, 9]
        assert f(9, order="inorder") == [4, 9, 5, 6]
        assert f(9, order="levelorder") == [9, 4, 5, 6]
        assert f(9, order="breadthfirst") == [9, 4, 5, 6]
        assert f(9, order="timeasc") == [4, 5, 6, 9]
        assert f(9, order="timedesc") == [9, 6, 5, 4]
        assert f(9, order="minlex_postorder") == [4, 5, 6, 9]

    def test_multiroot_example(self):
        tables = tskit.Tree.generate_balanced(7, arity=3).tree_sequence.dump_tables()
        tables.edges.truncate(len(tables.edges) - 3)
        t = tables.tree_sequence().first()

        #  7   8    9
        # ┏┻┓ ┏┻┓ ┏━╋━┓
        # 0 1 2 3 4 5 6
        def f(node=None, order=None):
            return list(t.nodes(node, order))

        assert f(order="preorder") == [7, 0, 1, 8, 2, 3, 9, 4, 5, 6]
        assert f(order="postorder") == [0, 1, 7, 2, 3, 8, 4, 5, 6, 9]
        assert f(order="inorder") == [0, 7, 1, 2, 8, 3, 4, 9, 5, 6]
        assert f(order="levelorder") == [7, 8, 9, 0, 1, 2, 3, 4, 5, 6]
        assert f(order="breadthfirst") == [7, 8, 9, 0, 1, 2, 3, 4, 5, 6]
        assert f(order="timeasc") == list(range(10))
        assert f(order="timedesc") == list(reversed(range(10)))
        assert f(order="minlex_postorder") == [0, 1, 7, 2, 3, 8, 4, 5, 6, 9]

        q = t.virtual_root
        assert f(q, order="preorder") == [q, 7, 0, 1, 8, 2, 3, 9, 4, 5, 6]
        assert f(q, order="postorder") == [0, 1, 7, 2, 3, 8, 4, 5, 6, 9, q]
        assert f(q, order="inorder") == [0, 7, 1, q, 2, 8, 3, 4, 9, 5, 6]
        assert f(q, order="levelorder") == [q, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6]
        assert f(q, order="breadthfirst") == [q, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6]
        assert f(q, order="timeasc") == list(range(10)) + [q]
        assert f(q, order="timedesc") == [q] + list(reversed(range(10)))
        assert f(q, order="minlex_postorder") == [0, 1, 7, 2, 3, 8, 4, 5, 6, 9, q]

        assert f(9, order="preorder") == [9, 4, 5, 6]
        assert f(9, order="postorder") == [4, 5, 6, 9]
        assert f(9, order="inorder") == [4, 9, 5, 6]
        assert f(9, order="levelorder") == [9, 4, 5, 6]
        assert f(9, order="breadthfirst") == [9, 4, 5, 6]
        assert f(9, order="minlex_postorder") == [4, 5, 6, 9]
        assert f(9, order="timeasc") == [4, 5, 6, 9]
        assert f(9, order="timedesc") == [9, 6, 5, 4]

    def test_multiroot_non_lexical_example(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time    population  individual  metadata
        0   1   0.000000    0   -1  b''
        1   1   0.000000    0   -1  b''
        2   1   0.000000    0   -1  b''
        3   1   0.000000    0   -1  b''
        4   1   0.000000    0   -1  b''
        5   1   0.000000    0   -1  b''
        6   1   0.000000    0   -1  b''
        7   1   0.000000    0   -1  b''
        8   1   0.000000    0   -1  b''
        9   1   0.000000    0   -1  b''
        10  0   0.047734    0   -1  b''
        11  0   0.061603    0   -1  b''
        12  0   0.189503    0   -1  b''
        13  0   0.275885    0   -1  b''
        14  0   0.518301    0   -1  b''
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.000000    10000.000000    10  0
        0.000000    10000.000000    10  2
        0.000000    10000.000000    11  9
        0.000000    10000.000000    11  10
        0.000000    10000.000000    12  3
        0.000000    10000.000000    12  7
        0.000000    10000.000000    13  5
        0.000000    10000.000000    13  11
        0.000000    10000.000000    14  1
        0.000000    10000.000000    14  8
        """
        )
        ts = tskit.load_text(
            nodes, edges, sequence_length=10000, strict=False, base64_metadata=False
        )
        t = ts.first()

        # Note: this is drawn out in "tree" order.
        #                 14
        #                 ┏┻┓
        #          13     ┃ ┃
        #         ┏━┻━┓   ┃ ┃
        #     12  ┃   ┃   ┃ ┃
        #     ┏┻┓ ┃   ┃   ┃ ┃
        #     ┃ ┃ ┃  11   ┃ ┃
        #     ┃ ┃ ┃ ┏━┻┓  ┃ ┃
        #     ┃ ┃ ┃ ┃ 10  ┃ ┃
        #     ┃ ┃ ┃ ┃ ┏┻┓ ┃ ┃
        # 4 6 3 7 5 9 0 2 1 8

        def f(node=None, order=None):
            return list(t.nodes(node, order))

        pre = f(order="preorder")
        post = f(order="postorder")
        inord = f(order="inorder")
        level = f(order="levelorder")
        breadth = f(order="breadthfirst")
        timeasc = f(order="timeasc")
        timedesc = f(order="timedesc")
        minlex = f(order="minlex_postorder")
        assert pre == [4, 6, 12, 3, 7, 13, 5, 11, 9, 10, 0, 2, 14, 1, 8]
        assert post == [4, 6, 3, 7, 12, 5, 9, 0, 2, 10, 11, 13, 1, 8, 14]
        assert inord == [4, 6, 3, 12, 7, 5, 13, 9, 11, 0, 10, 2, 1, 14, 8]
        assert level == [4, 6, 12, 13, 14, 3, 7, 5, 11, 1, 8, 9, 10, 0, 2]
        assert breadth == [4, 6, 12, 13, 14, 3, 7, 5, 11, 1, 8, 9, 10, 0, 2]
        assert timeasc == list(range(15))
        assert timedesc == list(reversed(range(15)))

        # And the minlex tree:
        #         14
        #         ┏┻┓
        #    13   ┃ ┃
        #   ┏━┻━┓ ┃ ┃
        #   ┃   ┃ ┃ ┃ 12
        #   ┃   ┃ ┃ ┃ ┏┻┓
        #  11   ┃ ┃ ┃ ┃ ┃
        #  ┏┻━┓ ┃ ┃ ┃ ┃ ┃
        # 10  ┃ ┃ ┃ ┃ ┃ ┃
        # ┏┻┓ ┃ ┃ ┃ ┃ ┃ ┃
        # 0 2 9 5 1 8 3 7 4 6
        assert minlex == [0, 2, 10, 9, 11, 5, 13, 1, 8, 14, 3, 7, 12, 4, 6]

    @pytest.mark.parametrize(
        ["order", "expected"],
        [
            ("preorder", [[9, 6, 2, 3, 7, 4, 5, 0, 1], [10, 4, 8, 5, 0, 1, 6, 2, 3]]),
            ("inorder", [[2, 6, 3, 9, 4, 7, 0, 5, 1], [4, 10, 0, 5, 1, 8, 2, 6, 3]]),
            ("postorder", [[2, 3, 6, 4, 0, 1, 5, 7, 9], [4, 0, 1, 5, 2, 3, 6, 8, 10]]),
            ("levelorder", [[9, 6, 7, 2, 3, 4, 5, 0, 1], [10, 4, 8, 5, 6, 0, 1, 2, 3]]),
            (
                "breadthfirst",
                [[9, 6, 7, 2, 3, 4, 5, 0, 1], [10, 4, 8, 5, 6, 0, 1, 2, 3]],
            ),
            ("timeasc", [[0, 1, 2, 3, 4, 5, 6, 7, 9], [0, 1, 2, 3, 4, 5, 6, 8, 10]]),
            ("timedesc", [[9, 7, 6, 5, 4, 3, 2, 1, 0], [10, 8, 6, 5, 4, 3, 2, 1, 0]]),
            (
                "minlex_postorder",
                [[0, 1, 5, 4, 7, 2, 3, 6, 9], [0, 1, 5, 2, 3, 6, 8, 4, 10]],
            ),
        ],
    )
    def test_ts_example(self, order, expected):
        # 1.20┊           ┊  10       ┊
        #     ┊           ┊ ┏━┻━━┓    ┊
        # 0.90┊     9     ┊ ┃    ┃    ┊
        #     ┊  ┏━━┻━┓   ┊ ┃    ┃    ┊
        # 0.60┊  ┃    ┃   ┊ ┃    8    ┊
        #     ┊  ┃    ┃   ┊ ┃  ┏━┻━┓  ┊
        # 0.44┊  ┃    7   ┊ ┃  ┃   ┃  ┊
        #     ┊  ┃  ┏━┻┓  ┊ ┃  ┃   ┃  ┊
        # 0.21┊  6  ┃  ┃  ┊ ┃  ┃   6  ┊
        #     ┊ ┏┻┓ ┃  ┃  ┊ ┃  ┃  ┏┻┓ ┊
        # 0.15┊ ┃ ┃ ┃  5  ┊ ┃  5  ┃ ┃ ┊
        #     ┊ ┃ ┃ ┃ ┏┻┓ ┊ ┃ ┏┻┓ ┃ ┃ ┊
        # 0.00┊ 2 3 4 0 1 ┊ 4 0 1 2 3 ┊
        #   0.00        0.50        1.00
        nodes = """\
        id      is_sample   population      time
        0       1       0               0.00000000000000
        1       1       0               0.00000000000000
        2       1       0               0.00000000000000
        3       1       0               0.00000000000000
        4       1       0               0.00000000000000
        5       0       0               0.14567111023387
        6       0       0               0.21385545626353
        7       0       0               0.43508024345063
        8       0       0               0.60156352971203
        9       0       0               0.90000000000000
        10      0       0               1.20000000000000
        """
        edges = """\
        id      left            right           parent  child
        0       0.00000000      1.00000000      5       0,1
        1       0.00000000      1.00000000      6       2,3
        2       0.00000000      0.50000000      7       4,5
        3       0.50000000      1.00000000      8       5,6
        4       0.00000000      0.50000000      9       6,7
        5       0.50000000      1.00000000      10      4,8
        """
        ts = tskit.load_text(
            nodes=io.StringIO(nodes), edges=io.StringIO(edges), strict=False
        )
        tree_orders = [list(tree.nodes(order=order)) for tree in ts.trees()]
        assert tree_orders == expected

    def test_polytomy_inorder(self):
        """
        If there are N children, current inorder traversal first visits
        floor(N/2) children, then the parent, then the remaining children.
        Here we explicitly test that behaviour.
        """
        #
        #    __4__
        #   / / \ \
        #  0 1   2 3
        #
        nodes_polytomy_4 = """\
        id      is_sample   population      time
        0       1       0               0.00000000000000
        1       1       0               0.00000000000000
        2       1       0               0.00000000000000
        3       1       0               0.00000000000000
        4       0       0               1.00000000000000
        """
        edges_polytomy_4 = """\
        id      left            right           parent  child
        0       0.00000000      1.00000000      4       0,1,2,3
        """
        #
        #    __5__
        #   / /|\ \
        #  0 1 2 3 4
        #
        nodes_polytomy_5 = """\
        id      is_sample   population      time
        0       1       0               0.00000000000000
        1       1       0               0.00000000000000
        2       1       0               0.00000000000000
        3       1       0               0.00000000000000
        4       1       0               0.00000000000000
        5       0       0               1.00000000000000
        """
        edges_polytomy_5 = """\
        id      left            right           parent  child
        0       0.00000000      1.00000000      5       0,1,2,3,4
        """
        for nodes_string, edges_string, expected_result in [
            [nodes_polytomy_4, edges_polytomy_4, [[0, 1, 4, 2, 3]]],
            [nodes_polytomy_5, edges_polytomy_5, [[0, 1, 5, 2, 3, 4]]],
        ]:
            ts = tskit.load_text(
                nodes=io.StringIO(nodes_string),
                edges=io.StringIO(edges_string),
                strict=False,
            )
            tree_orders = []
            for tree in ts.trees():
                tree_orders.append(list(tree.nodes(order="inorder")))
            assert tree_orders == expected_result

    def test_minlex_postorder_multiple_roots(self):
        #
        #    10    8     9     11
        #   / \   / \   / \   / \
        #  5   3 2   4 6   7 1   0
        #
        nodes_string = """\
        id      is_sample   population      time
        0       1       0               0.00000000000000
        1       1       0               0.00000000000000
        2       1       0               0.00000000000000
        3       1       0               0.00000000000000
        4       1       0               0.00000000000000
        5       1       0               0.00000000000000
        6       1       0               0.00000000000000
        7       1       0               0.00000000000000
        8       0       0               1.00000000000000
        9       0       0               1.00000000000000
        10      0       0               1.00000000000000
        11      0       0               1.00000000000000
        """
        edges_string = """\
        id      left            right           parent  child
        0       0.00000000      1.00000000      8       2,4
        1       0.00000000      1.00000000      9       6,7
        2       0.00000000      1.00000000      10      5,3
        3       0.00000000      1.00000000      11      1,0
        """
        expected_result = [[0, 1, 11, 2, 4, 8, 3, 5, 10, 6, 7, 9]]
        ts = tskit.load_text(
            nodes=io.StringIO(nodes_string),
            edges=io.StringIO(edges_string),
            strict=False,
        )
        tree_orders = []
        for tree in ts.trees():
            tree_orders.append(list(tree.nodes(order="minlex_postorder")))
        assert tree_orders == expected_result


class TestMRCA:
    """
    Test both the tree.mrca and tree.tmrca methods.
    """

    t = tskit.Tree.generate_balanced(3)
    #  4
    # ┏━┻┓
    # ┃  3
    # ┃ ┏┻┓
    # 0 1 2

    @pytest.mark.parametrize("args, expected", [((2, 1), 3), ((0, 1, 2), 4)])
    def test_two_or_more_args(self, args, expected):
        assert self.t.mrca(*args) == expected
        assert self.t.tmrca(*args) == self.t.tree_sequence.nodes_time[expected]

    def test_less_than_two_args(self):
        with pytest.raises(ValueError):
            self.t.mrca(1)
        with pytest.raises(ValueError):
            self.t.tmrca(1)

    def test_no_args(self):
        with pytest.raises(ValueError):
            self.t.mrca()
        with pytest.raises(ValueError):
            self.t.tmrca()

    def test_same_args(self):
        assert self.t.mrca(0, 0, 0, 0) == 0
        assert self.t.tmrca(0, 0, 0, 0) == self.t.tree_sequence.nodes_time[0]

    def test_different_tree_levels(self):
        assert self.t.mrca(0, 3) == 4
        assert self.t.tmrca(0, 3) == self.t.tree_sequence.nodes_time[4]

    def test_out_of_bounds_args(self):
        with pytest.raises(ValueError):
            self.t.mrca(0, 6)
        with pytest.raises(ValueError):
            self.t.tmrca(0, 6)

    def test_virtual_root_arg(self):
        assert self.t.mrca(0, 5) == 5
        assert np.isposinf(self.t.tmrca(0, 5))

    def test_multiple_roots(self):
        ts = tskit.Tree.generate_balanced(10).tree_sequence
        ts = ts.delete_intervals([ts.first().interval])
        assert ts.first().mrca(*ts.samples()) == tskit.NULL
        # We decided to raise an error for tmrca here, rather than report inf
        # see https://github.com/tskit-dev/tskit/issues/2801
        with pytest.raises(ValueError, match="do not share a common ancestor"):
            ts.first().tmrca(0, 6)


class TestPathLength:
    t = tskit.Tree.generate_balanced(9)
    #         16
    #    ┏━━━━┻━━━┓
    #    ┃       15
    #    ┃     ┏━━┻━┓
    #   11     ┃   14
    #  ┏━┻━┓   ┃  ┏━┻┓
    #  9  10  12  ┃ 13
    # ┏┻┓ ┏┻┓ ┏┻┓ ┃ ┏┻┓
    # 0 1 2 3 4 5 6 7 8

    def test_tmrca_leaf(self):
        assert self.t.path_length(0, 16) == 3
        assert self.t.path_length(16, 0) == 3
        assert self.t.path_length(7, 16) == 4

    def test_equal_depth(self):
        assert self.t.path_length(5, 16) == self.t.depth(5)

    def test_two_leaves(self):
        assert self.t.path_length(0, 8) == 7

    def test_two_leaves_depth(self):
        assert self.t.path_length(0, 8) == self.t.depth(0) + self.t.depth(8)

    @pytest.mark.parametrize("args", [[], [1], [1, 2, 3]])
    def test_bad_num_args(self, args):
        with pytest.raises(TypeError):
            self.t.path_length(*args)

    @pytest.mark.parametrize("bad_arg", [[], "1"])
    def test_bad_arg_type(self, bad_arg):
        with pytest.raises(TypeError):
            self.t.path_length(0, bad_arg)
        with pytest.raises(TypeError):
            self.t.path_length(bad_arg, 0)

    def test_same_args(self):
        assert self.t.path_length(10, 10) == 0

    def test_different_tree_levels(self):
        assert self.t.path_length(1, 10) == 3

    def test_out_of_bounds_args(self):
        with pytest.raises(ValueError):
            self.t.path_length(0, 20)

    @pytest.mark.parametrize("u", range(17))
    def test_virtual_root_arg(self, u):
        assert self.t.path_length(u, self.t.virtual_root) == self.t.depth(u) + 1
        assert self.t.path_length(self.t.virtual_root, u) == self.t.depth(u) + 1

    def test_both_args_virtual_root(self):
        assert self.t.path_length(self.t.virtual_root, self.t.virtual_root) == 0

    def test_no_mrca(self):
        tree = self.t.copy()
        tree.clear()
        assert math.isinf(tree.path_length(0, 1))


class TestMRCACalculator:
    """
    Class to test the Schieber-Vishkin algorithm.

    These tests are included here as we use the MRCA calculator below in
    our tests.
    """

    def test_all_oriented_forests(self):
        # Runs through all possible oriented forests and checks all possible
        # node pairs using an inferior algorithm.
        for n in range(2, 9):
            for pi in oriented_forests(n):
                sv = tests.MRCACalculator(pi)
                for j in range(1, n + 1):
                    for k in range(1, j + 1):
                        mrca = get_mrca(pi, j, k)
                        assert mrca == sv.get_mrca(j, k)


class HighLevelTestCase:
    """
    Superclass of tests on the high level interface.
    """

    def verify_tree_mrcas(self, st):
        # Check the mrcas
        oriented_forest = [st.get_parent(j) for j in range(st.tree_sequence.num_nodes)]
        mrca_calc = tests.MRCACalculator(oriented_forest)
        # We've done exhaustive tests elsewhere, no need to go
        # through the combinations.
        for j in range(st.tree_sequence.num_nodes):
            mrca = st.get_mrca(0, j)
            assert mrca == mrca_calc.get_mrca(0, j)
            if mrca != tskit.NULL:
                assert st.get_time(mrca) == st.get_tmrca(0, j)

    def verify_tree_branch_lengths(self, tree):
        for u in tree.tree_sequence.samples():
            while tree.parent(u) != tskit.NULL:
                length = tree.time(tree.parent(u)) - tree.time(u)
                assert length > 0.0
                assert tree.branch_length(u) == length
                u = tree.parent(u)
            assert tree.parent(u) == tskit.NULL
            assert tree.branch_length(u) == 0

    def verify_tree_structure(self, st):
        roots = set()
        for u in st.samples():
            # verify the path to root
            assert st.is_sample(u)
            times = []
            while st.get_parent(u) != tskit.NULL:
                v = st.get_parent(u)
                times.append(st.get_time(v))
                assert st.get_time(v) >= 0.0
                assert u in st.get_children(v)
                u = v
            roots.add(u)
            assert times == sorted(times)
        assert sorted(list(roots)) == sorted(st.roots)
        assert len(st.roots) == st.num_roots
        u = st.left_root
        roots = []
        while u != tskit.NULL:
            roots.append(u)
            u = st.right_sib(u)
        assert roots == st.roots
        # To a top-down traversal, and make sure we meet all the samples.
        samples = []
        for root in st.roots:
            stack = [root]
            while len(stack) > 0:
                u = stack.pop()
                assert u != tskit.NULL
                if st.is_sample(u):
                    samples.append(u)
                if st.is_leaf(u):
                    assert len(st.get_children(u)) == 0
                else:
                    for c in reversed(st.get_children(u)):
                        stack.append(c)
                # Check that we get the correct number of samples at each
                # node.
                assert st.get_num_samples(u) == len(list(st.samples(u)))
                assert st.get_num_tracked_samples(u) == 0
        assert sorted(samples) == sorted(st.samples())
        # Check the parent dict
        pi = st.get_parent_dict()
        for root in st.roots:
            assert root not in pi
        for k, v in pi.items():
            assert st.get_parent(k) == v
        assert st.num_samples() == len(samples)
        assert sorted(st.samples()) == sorted(samples)

    def verify_tree_depths(self, st):
        for root in st.roots:
            stack = [(root, 0)]
            while len(stack) > 0:
                u, depth = stack.pop()
                assert st.depth(u) == depth
                for c in st.children(u):
                    stack.append((c, depth + 1))

    def verify_tree(self, st):
        self.verify_tree_mrcas(st)
        self.verify_tree_branch_lengths(st)
        self.verify_tree_structure(st)
        self.verify_tree_depths(st)

    def verify_trees(self, ts):
        pts = tests.PythonTreeSequence(ts)
        iter1 = ts.trees()
        iter2 = pts.trees()
        length = 0
        num_trees = 0
        breakpoints = [0]
        for st1, st2 in zip(iter1, iter2):
            assert st1.get_sample_size() == ts.get_sample_size()
            roots = set()
            for u in ts.samples():
                root = u
                while st1.get_parent(root) != tskit.NULL:
                    root = st1.get_parent(root)
                roots.add(root)
            assert st1.left_root == st2.left_root
            assert sorted(list(roots)) == sorted(st1.roots)
            assert st1.roots == st2.roots
            if len(roots) == 0:
                assert st1.root == tskit.NULL
            elif len(roots) == 1:
                assert st1.root == list(roots)[0]
            else:
                with pytest.raises(ValueError):
                    st1.root
            assert st2 == st1
            assert not (st2 != st1)
            left, right = st1.get_interval()
            breakpoints.append(right)
            assert left == pytest.approx(length)
            assert left >= 0
            assert right > left
            assert right <= ts.get_sequence_length()
            length += right - left
            self.verify_tree(st1)
            num_trees += 1
        with pytest.raises(StopIteration):
            next(iter1)
        with pytest.raises(StopIteration):
            next(iter2)
        assert ts.get_num_trees() == num_trees
        assert breakpoints == list(ts.breakpoints())
        assert length == pytest.approx(ts.get_sequence_length())


class TestNumpySamples:
    """
    Tests that we correctly handle samples as numpy arrays when passed to
    various methods.
    """

    def get_tree_sequence(self, num_demes=4, times=None, n=40):
        if times is None:
            times = [0]
        return msprime.simulate(
            samples=[
                msprime.Sample(time=t, population=j % num_demes)
                for j in range(n)
                for t in times
            ],
            population_configurations=[
                msprime.PopulationConfiguration() for _ in range(num_demes)
            ],
            migration_matrix=[
                [int(j != k) for j in range(num_demes)] for k in range(num_demes)
            ],
            random_seed=1,
            mutation_rate=10,
        )

    def test_samples(self):
        d = 4
        ts = self.get_tree_sequence(d)
        assert np.array_equal(ts.samples(), np.arange(ts.num_samples, dtype=np.int32))
        total = 0
        for pop in range(d):
            subsample = ts.samples(pop)
            total += subsample.shape[0]
            assert np.array_equal(subsample, ts.samples(population=pop))
            assert list(subsample) == [
                node.id
                for node in ts.nodes()
                if node.population == pop and node.is_sample()
            ]
        assert total == ts.num_samples

    @pytest.mark.parametrize("time", [0, 0.1, 1 / 3, 1 / 4, 5 / 7])
    def test_samples_time(self, time):
        ts = self.get_tree_sequence(num_demes=2, n=20, times=[time, 0.2, 1, 15])
        assert np.array_equal(get_samples(ts, time=time), ts.samples(time=time))
        for population in (None, 0):
            assert np.array_equal(
                get_samples(ts, time=time, population=population),
                ts.samples(time=time, population=population),
            )

    @pytest.mark.parametrize(
        "time_interval",
        [
            [0, 0.1],
            (0, 1 / 3),
            np.array([1 / 4, 2 / 3]),
            (0.345, 5 / 7),
            (-1, 1),
        ],
    )
    def test_samples_time_interval(self, time_interval):
        rng = np.random.default_rng(seed=931)
        times = rng.uniform(low=time_interval[0], high=2 * time_interval[1], size=20)
        ts = self.get_tree_sequence(num_demes=2, n=1, times=times)
        assert np.array_equal(
            get_samples(ts, time=time_interval),
            ts.samples(time=time_interval),
        )
        for population in (None, 0):
            assert np.array_equal(
                get_samples(ts, time=time_interval, population=population),
                ts.samples(time=time_interval, population=population),
            )

    def test_samples_example(self):
        tables = tskit.TableCollection(sequence_length=10)
        time = [0, 0, 1, 1, 1, 3, 3.00001, 3.0 - 0.0001, 1 / 3]
        pops = [1, 3, 1, 2, 1, 1, 1, 3, 1]
        for _ in range(max(pops) + 1):
            tables.populations.add_row()
        for t, p in zip(time, pops):
            tables.nodes.add_row(
                flags=tskit.NODE_IS_SAMPLE,
                time=t,
                population=p,
            )
        # add not-samples also
        for t, p in zip(time, pops):
            tables.nodes.add_row(
                flags=0,
                time=t,
                population=p,
            )
        ts = tables.tree_sequence()
        assert np.array_equal(
            ts.samples(),
            np.arange(len(time)),
        )
        assert np.array_equal(
            ts.samples(time=[0, np.inf]),
            np.arange(len(time)),
        )
        assert np.array_equal(
            ts.samples(time=0),
            [0, 1],
        )
        # default tolerance is 1e-5
        assert np.array_equal(
            ts.samples(time=0.3333333),
            [8],
        )
        assert np.array_equal(
            ts.samples(time=3),
            [5, 6],
        )
        assert np.array_equal(
            ts.samples(time=1),
            [2, 3, 4],
        )
        assert np.array_equal(
            ts.samples(time=1, population=2),
            [3],
        )
        assert np.array_equal(
            ts.samples(population=0),
            [],
        )
        assert np.array_equal(
            ts.samples(population=1),
            [0, 2, 4, 5, 6, 8],
        )
        assert np.array_equal(
            ts.samples(population=2),
            [3],
        )
        assert np.array_equal(
            ts.samples(time=[0, 3]),
            [0, 1, 2, 3, 4, 7, 8],
        )
        # note tuple instead of array
        assert np.array_equal(
            ts.samples(time=(1, 3)),
            [2, 3, 4, 7],
        )
        assert np.array_equal(
            ts.samples(time=[0, 3], population=1),
            [0, 2, 4, 8],
        )
        assert np.array_equal(
            ts.samples(time=[0.333333, 3]),
            [2, 3, 4, 7, 8],
        )
        assert np.array_equal(
            ts.samples(time=[100, np.inf]),
            [],
        )
        assert np.array_equal(
            ts.samples(time=-1),
            [],
        )
        assert np.array_equal(
            ts.samples(time=[-100, 100]),
            np.arange(len(time)),
        )
        assert np.array_equal(
            ts.samples(time=[-100, -1]),
            [],
        )

    def test_samples_time_errors(self):
        ts = self.get_tree_sequence(4)
        # error incorrect types
        with pytest.raises(ValueError):
            ts.samples(time="s")
        with pytest.raises(ValueError):
            ts.samples(time=[])
        with pytest.raises(ValueError):
            ts.samples(time=np.array([1, 2, 3]))
        with pytest.raises(ValueError):
            ts.samples(time=(1, 2, 3))
        # error using min and max switched
        with pytest.raises(ValueError):
            ts.samples(time=(2.4, 1))

    def test_samples_args(self, ts_fixture):
        ts_fixture.samples(1)
        with pytest.raises(TypeError, match="takes from 1 to 2 positional arguments"):
            ts_fixture.samples(1, 2)

    def test_genotype_matrix_indexing(self):
        num_demes = 4
        ts = self.get_tree_sequence(num_demes)
        G = ts.genotype_matrix()
        for d in range(num_demes):
            samples = ts.samples(population=d)
            total = 0
            for tree in ts.trees(tracked_samples=samples):
                for mutation in tree.mutations():
                    total += tree.num_tracked_samples(mutation.node)
            assert total == np.sum(G[:, samples])

    def test_genotype_indexing(self):
        num_demes = 6
        ts = self.get_tree_sequence(num_demes)
        for d in range(num_demes):
            samples = ts.samples(population=d)
            total = 0
            for tree in ts.trees(tracked_samples=samples):
                for mutation in tree.mutations():
                    total += tree.num_tracked_samples(mutation.node)
            other_total = 0
            for variant in ts.variants():
                other_total += np.sum(variant.genotypes[samples])
            assert total == other_total

    def test_pairwise_diversity(self):
        num_demes = 6
        ts = self.get_tree_sequence(num_demes)
        pi1 = ts.pairwise_diversity(ts.samples())
        pi2 = ts.pairwise_diversity()
        assert pi1 == pi2
        for d in range(num_demes):
            samples = ts.samples(population=d)
            pi1 = ts.pairwise_diversity(samples)
            pi2 = ts.pairwise_diversity(list(samples))
            assert pi1 == pi2

    def test_simplify(self):
        num_demes = 3
        ts = self.get_tree_sequence(num_demes)
        sts = ts.simplify(samples=ts.samples())
        assert ts.num_samples == sts.num_samples
        for d in range(num_demes):
            samples = ts.samples(population=d)
            sts = ts.simplify(samples=samples)
            assert sts.num_samples == samples.shape[0]


class TestTreeSequence(HighLevelTestCase):
    """
    Tests for the tree sequence object.
    """

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_row_getter(self, ts):
        for table_name, table in ts.tables_dict.items():
            sequence = getattr(ts, table_name)()
            element_name = table_name[:-1]  # cut off the "s": "edges" -> "edge"
            element_accessor = getattr(ts, element_name)
            for i, n in enumerate(sequence):
                assert element_accessor(i) == n
                assert element_accessor(-(table.num_rows - i)) == n
            with pytest.raises(IndexError):
                element_accessor(table.num_rows)
            with pytest.raises(IndexError):
                element_accessor(-(table.num_rows + 1))

    @pytest.mark.parametrize("index", [0.1, float(0), None, np.array([0, 1]), np.inf])
    def test_bad_row_getter(self, index, simple_degree2_ts_fixture):
        for table_name in simple_degree2_ts_fixture.tables_dict.keys():
            element_name = table_name[:-1]  # cut off the "s": "edges" -> "edge"
            element_accessor = getattr(simple_degree2_ts_fixture, element_name)
            if element_name == "site" and index is None:
                # special case
                match = "id or position must be provided"
            else:
                match = "integer type"
            with pytest.raises(TypeError, match=match):
                element_accessor(index)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_discrete_genome(self, ts):
        def is_discrete(a):
            return np.all(np.floor(a) == a)

        tables = ts.tables
        discrete_genome = (
            is_discrete([tables.sequence_length])
            and is_discrete(tables.edges.left)
            and is_discrete(tables.edges.right)
            and is_discrete(tables.sites.position)
            and is_discrete(tables.migrations.left)
            and is_discrete(tables.migrations.right)
        )
        assert ts.discrete_genome == discrete_genome

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_discrete_time(self, ts):
        def is_discrete(a):
            return np.all(np.logical_or(np.floor(a) == a, tskit.is_unknown_time(a)))

        tables = ts.tables
        discrete_time = (
            is_discrete(tables.nodes.time)
            and is_discrete(tables.mutations.time)
            and is_discrete(tables.migrations.time)
        )
        assert ts.discrete_time == discrete_time

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_trees(self, ts):
        self.verify_trees(ts)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_mutations(self, ts):
        self.verify_mutations(ts)

    def verify_pairwise_diversity(self, ts):
        haplotypes = ts.genotype_matrix(isolated_as_missing=False).T
        if ts.num_samples == 0:
            with pytest.raises(ValueError, match="at least one element"):
                ts.get_pairwise_diversity()
            return
        pi1 = ts.get_pairwise_diversity()
        pi2 = simple_get_pairwise_diversity(haplotypes)
        assert pi1 == pytest.approx(pi2)
        assert pi1 >= 0.0
        assert not math.isnan(pi1)
        # Check for a subsample.
        num_samples = ts.get_sample_size() // 2 + 1
        samples = list(ts.samples())[:num_samples]
        pi1 = ts.get_pairwise_diversity(samples)
        pi2 = simple_get_pairwise_diversity([haplotypes[j] for j in range(num_samples)])
        assert pi1 == pytest.approx(pi2)
        assert pi1 >= 0.0
        assert not math.isnan(pi1)

    @pytest.mark.slow
    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_pairwise_diversity(self, ts):
        self.verify_pairwise_diversity(ts)

    @pytest.mark.parametrize("order", ["abc", 0, 1, False])
    def test_bad_node_iteration_order(self, order):
        ts = tskit.TableCollection(1).tree_sequence()
        with pytest.raises(ValueError, match="order"):
            ts.nodes(order=order)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_node_iteration_order(self, ts):
        order = [n.id for n in ts.nodes()]
        assert order == list(range(ts.num_nodes))
        order = [n.id for n in ts.nodes(order="id")]
        assert order == list(range(ts.num_nodes))
        order = np.array([n.id for n in ts.nodes(order="timeasc")], dtype=int)
        assert np.all(ts.nodes_time[order] == np.sort(ts.nodes_time))
        # Check it conforms to the order of parents in the edge table
        parent_only_order = order[np.isin(order, ts.edges_parent)]
        edge_parents = np.concatenate(
            (ts.edges_parent[:-1][np.diff(ts.edges_parent) != 0], ts.edges_parent[-1:])
        )
        assert np.all(parent_only_order == edge_parents)

    def verify_edgesets(self, ts):
        """
        Verifies that the edgesets we return are equivalent to the original edges.
        """
        new_edges = []
        for edgeset in ts.edgesets():
            assert edgeset.children == sorted(edgeset.children)
            assert len(edgeset.children) > 0
            for child in edgeset.children:
                new_edges.append(
                    tskit.Edge(edgeset.left, edgeset.right, edgeset.parent, child)
                )
        # squash the edges.
        t = ts.dump_tables().nodes.time
        new_edges.sort(key=lambda e: (t[e.parent], e.parent, e.child, e.left))

        squashed = []
        if len(new_edges) > 0:
            last_e = new_edges[0]
            for e in new_edges[1:]:
                condition = (
                    e.parent != last_e.parent
                    or e.child != last_e.child
                    or e.left != last_e.right
                )
                if condition:
                    squashed.append(last_e)
                    last_e = e
                last_e.right = e.right
            squashed.append(last_e)
            # reset the IDs
            for i, e in enumerate(squashed):
                e.id = i
        edges = list(ts.edges())
        assert len(squashed) == len(edges)
        assert edges == squashed

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_edge_ids(self, ts):
        for index, edge in enumerate(ts.edges()):
            assert edge.id == index

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_edge_span_property(self, ts):
        for edge in ts.edges():
            assert edge.span == edge.right - edge.left

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_edge_interval_property(self, ts):
        for edge in ts.edges():
            assert edge.interval == (edge.left, edge.right)
        if ts.num_trees == 1 and ts.num_edges > 0:
            for edge in ts.edges():
                assert edge.interval == ts.first().interval

    def test_edgesets(self):
        tested = False
        # We manual loop in this test to test the example tree sequences are working
        for ts in get_example_tree_sequences(pytest_params=False):
            # Can't get edgesets with metadata
            if ts.tables.edges.metadata_schema == tskit.MetadataSchema(None):
                self.verify_edgesets(ts)
                tested = True
        assert tested

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_breakpoints(self, ts):
        breakpoints = ts.breakpoints(as_array=True)
        assert breakpoints.shape == (ts.num_trees + 1,)
        other = np.fromiter(iter([0] + [t.interval.right for t in ts.trees()]), float)
        assert np.array_equal(other, breakpoints)
        # in case downstream code has
        for j, x in enumerate(ts.breakpoints()):
            assert breakpoints[j] == x
            assert isinstance(x, float)
        assert j == ts.num_trees

    def verify_coalescence_records(self, ts):
        """
        Checks that the coalescence records we output are correct.
        """
        edgesets = list(ts.edgesets())
        records = list(ts.records())
        assert len(edgesets) == len(records)
        for edgeset, record in zip(edgesets, records):
            assert edgeset.left == record.left
            assert edgeset.right == record.right
            assert edgeset.parent == record.node
            assert edgeset.children == record.children
            parent = ts.node(edgeset.parent)
            assert parent.time == record.time
            assert parent.population == record.population

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_coalescence_records(self, ts):
        self.verify_coalescence_records(ts)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_compute_mutation_parent(self, ts):
        tables = ts.dump_tables()
        before = tables.mutations.parent[:]
        tables.compute_mutation_parents()
        parent = ts.tables.mutations.parent
        assert np.array_equal(parent, before)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_compute_mutation_time(self, ts):
        tables = ts.dump_tables()
        python_time = tsutil.compute_mutation_times(ts)
        tables.compute_mutation_times()
        assert np.allclose(python_time, tables.mutations.time, rtol=1e-10, atol=1e-10)
        # Check we have valid times
        tables.tree_sequence()

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_tracked_samples(self, ts):
        # Should be empty list by default.
        for tree in ts.trees():
            assert tree.num_tracked_samples() == 0
            for u in tree.nodes():
                assert tree.num_tracked_samples(u) == 0
        samples = list(ts.samples())
        tracked_samples = samples[:2]
        for tree in ts.trees(tracked_samples=tracked_samples):
            nu = [0 for j in range(ts.num_nodes)]
            assert tree.num_tracked_samples() == len(tracked_samples)
            for j in tracked_samples:
                u = j
                while u != tskit.NULL:
                    nu[u] += 1
                    u = tree.parent(u)
            for u, count in enumerate(nu):
                assert tree.num_tracked_samples(u) == count
            assert tree.num_tracked_samples(tree.virtual_root) == len(tracked_samples)

    def test_tracked_samples_is_first_arg(self):
        ts = tskit.Tree.generate_balanced(6).tree_sequence
        samples = [0, 1, 2]
        tree = next(ts.trees(samples))
        assert tree.num_tracked_samples() == 3

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_deprecated_sample_aliases(self, ts):
        # Ensure that we get the same results from the various combinations
        # of leaf_lists, sample_lists etc.
        samples = list(ts.samples())[:2]
        # tracked leaves/samples
        trees_new = ts.trees(tracked_samples=samples)
        trees_old = ts.trees(tracked_leaves=samples)
        for t_new, t_old in zip(trees_new, trees_old):
            for u in t_new.nodes():
                assert t_new.num_tracked_samples(u) == t_old.get_num_tracked_leaves(u)
        trees_new = ts.trees()
        trees_old = ts.trees()
        for t_new, t_old in zip(trees_new, trees_old):
            for u in t_new.nodes():
                assert t_new.num_samples(u) == t_old.get_num_leaves(u)
                assert list(t_new.samples(u)) == list(t_old.get_leaves(u))
        for on in [True, False]:
            trees_new = ts.trees(sample_lists=on)
            trees_old = ts.trees(leaf_lists=on)
            for t_new, t_old in zip(trees_new, trees_old):
                for u in t_new.nodes():
                    assert t_new.num_samples(u) == t_old.get_num_leaves(u)
                    assert list(t_new.samples(u)) == list(t_old.get_leaves(u))

    def verify_samples(self, ts):
        # We should get the same list of samples if we use the low-level
        # sample lists or a simple traversal.
        samples1 = []
        for t in ts.trees(sample_lists=False):
            samples1.append(list(t.samples()))
        samples2 = []
        for t in ts.trees(sample_lists=True):
            samples2.append(list(t.samples()))
        assert samples1 == samples2

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_samples(self, ts):
        self.verify_samples(ts)
        pops = {node.population for node in ts.nodes()}
        for pop in pops:
            subsample = ts.samples(pop)
            assert np.array_equal(subsample, ts.samples(population=pop))
            assert np.array_equal(subsample, ts.samples(population_id=pop))
            assert list(subsample) == [
                node.id
                for node in ts.nodes()
                if node.population == pop and node.is_sample()
            ]
        with pytest.raises(ValueError):
            ts.samples(population=0, population_id=0)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_first_last(self, ts):
        for kwargs in [{}, {"tracked_samples": ts.samples()}]:
            t1 = ts.first(**kwargs)
            t2 = next(ts.trees())
            assert not (t1 is t2)
            assert t1.parent_dict == t2.parent_dict
            assert t1.index == 0
            if "tracked_samples" in kwargs:
                assert t1.num_tracked_samples() == ts.num_samples
            else:
                assert t1.num_tracked_samples() == 0

            t1 = ts.last(**kwargs)
            t2 = next(reversed(ts.trees()))
            assert not (t1 is t2)
            assert t1.parent_dict == t2.parent_dict
            assert t1.index == ts.num_trees - 1
            if "tracked_samples" in kwargs:
                assert t1.num_tracked_samples() == ts.num_samples
            else:
                assert t1.num_tracked_samples() == 0

    def test_trees_interface(self):
        # Use a tree sequence guaranteed to have node 0 as the first sample node
        ts = tskit.Tree.generate_balanced(10).tree_sequence
        for t in ts.trees():
            assert t.get_num_samples(0) == 1
            assert t.get_num_tracked_samples(0) == 0
            assert list(t.samples(0)) == [0]
            assert t.tree_sequence is ts

        for t in ts.trees(tracked_samples=[0]):
            assert t.get_num_samples(0) == 1
            assert t.get_num_tracked_samples(0) == 1
            assert list(t.samples(0)) == [0]

        for t in ts.trees(sample_lists=True):
            assert t.get_num_samples(0) == 1
            assert t.get_num_tracked_samples(0) == 0
            assert list(t.samples(0)) == [0]

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_get_pairwise_diversity(self, ts):
        with pytest.raises(ValueError, match="at least one element"):
            ts.get_pairwise_diversity([])
        samples = list(ts.samples())
        if len(samples) == 0:
            with pytest.raises(
                ValueError, match="Sample sets must contain at least one element"
            ):
                ts.get_pairwise_diversity()
        else:
            assert ts.get_pairwise_diversity() == ts.get_pairwise_diversity(samples)
            assert ts.get_pairwise_diversity(samples[:2]) == ts.get_pairwise_diversity(
                list(reversed(samples[:2]))
            )

    def test_populations(self):
        more_than_zero = False
        for ts in get_example_tree_sequences(pytest_params=False):
            N = ts.num_populations
            if N > 0:
                more_than_zero = True
            pops = list(ts.populations())
            assert len(pops) == N
            for j in range(N):
                assert pops[j] == ts.population(j)
                assert pops[j].id == j
        assert more_than_zero

    def test_individuals(self):
        more_than_zero = False
        mapped_to_nodes = False
        for ts in get_example_tree_sequences(pytest_params=False):
            ind_node_map = collections.defaultdict(list)
            for node in ts.nodes():
                if node.individual != tskit.NULL:
                    ind_node_map[node.individual].append(node.id)
            if len(ind_node_map) > 0:
                mapped_to_nodes = True
            N = ts.num_individuals
            if N > 0:
                more_than_zero = True
            inds = list(ts.individuals())
            assert len(inds) == N
            for j in range(N):
                assert inds[j] == ts.individual(j)
                assert inds[j].id == j
                assert isinstance(inds[j].parents, np.ndarray)
                assert isinstance(inds[j].location, np.ndarray)
                assert isinstance(inds[j].nodes, np.ndarray)
                assert ind_node_map[j] == list(inds[j].nodes)

        assert more_than_zero
        assert mapped_to_nodes

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_get_population(self, ts):
        # Deprecated interface for ts.node(id).population
        N = ts.get_num_nodes()
        with pytest.raises(ValueError):
            ts.get_population(-1)
        with pytest.raises(ValueError):
            ts.get_population(N)
        with pytest.raises(ValueError):
            ts.get_population(N + 1)
        for node in range(N):
            assert ts.get_population(node) == ts.node(node).population

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_get_time(self, ts):
        # Deprecated interface for ts.node(id).time
        N = ts.get_num_nodes()
        with pytest.raises(ValueError):
            ts.get_time(-1)
        with pytest.raises(ValueError):
            ts.get_time(N)
        with pytest.raises(ValueError):
            ts.get_time(N + 1)
        for u in range(N):
            assert ts.get_time(u) == ts.node(u).time

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_max_root_time(self, ts):
        oldest = None
        for tree in ts.trees():
            for root in tree.roots:
                oldest = (
                    tree.time(root) if oldest is None else max(oldest, tree.time(root))
                )
        if oldest is None:
            assert pytest.raises(ValueError, match="max()")
        else:
            assert oldest == ts.max_root_time

    def test_max_root_time_corner_cases(self):
        tables = tskit.TableCollection(1)
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=1)
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=2)
        tables.nodes.add_row(flags=0, time=3)
        assert tables.tree_sequence().max_root_time == 2
        tables.edges.add_row(0, 1, 1, 0)
        assert tables.tree_sequence().max_root_time == 2
        tables.edges.add_row(0, 1, 3, 1)
        assert tables.tree_sequence().max_root_time == 3

    def test_subset_reverse_all_nodes(self):
        ts = tskit.Tree.generate_comb(5).tree_sequence
        assert np.all(ts.samples() == np.arange(ts.num_samples))
        flipped_ids = np.flip(np.arange(ts.num_nodes))
        new_ts = ts.subset(flipped_ids)
        assert set(new_ts.samples()) == set(flipped_ids[np.arange(ts.num_samples)])
        r1 = ts.first().rank()
        r2 = new_ts.first().rank()
        assert r1.shape == r2.shape
        assert r1.label != r2.label

    def test_subset_reverse_internal_nodes(self):
        ts = tskit.Tree.generate_balanced(5).tree_sequence
        internal_nodes = np.ones(ts.num_nodes, dtype=bool)
        internal_nodes[ts.samples()] = False
        node_ids = np.arange(ts.num_nodes)
        node_ids[internal_nodes] = np.flip(node_ids[internal_nodes])
        new_ts = ts.subset(node_ids)
        assert np.any(new_ts.nodes_time != ts.nodes_time)
        assert new_ts.first().rank() == ts.first().rank()

    def test_deprecated_apis(self):
        ts = msprime.simulate(10, random_seed=1)
        assert ts.get_ll_tree_sequence() == ts.ll_tree_sequence
        assert ts.get_sample_size() == ts.sample_size
        assert ts.get_sample_size() == ts.num_samples
        assert ts.get_sequence_length() == ts.sequence_length
        assert ts.get_num_trees() == ts.num_trees
        assert ts.get_num_mutations() == ts.num_mutations
        assert ts.get_num_nodes() == ts.num_nodes
        assert ts.get_pairwise_diversity() == ts.pairwise_diversity()
        samples = ts.samples()
        assert ts.get_pairwise_diversity(samples) == ts.pairwise_diversity(samples)
        assert np.array_equal(ts.get_samples(), ts.samples())

    def test_sites(self):
        some_sites = False
        for ts in get_example_tree_sequences(pytest_params=False):
            tables = ts.dump_tables()
            sites = tables.sites
            mutations = tables.mutations
            assert ts.num_sites == len(sites)
            assert ts.num_mutations == len(mutations)
            previous_pos = -1
            mutation_index = 0
            ancestral_state = tskit.unpack_strings(
                sites.ancestral_state, sites.ancestral_state_offset
            )
            derived_state = tskit.unpack_strings(
                mutations.derived_state, mutations.derived_state_offset
            )

            for index, site in enumerate(ts.sites()):
                s2 = ts.site(site.id)
                assert s2 == site
                s3 = ts.site(position=site.position)
                assert s3 == site
                assert site.position == sites.position[index]
                assert site.position > previous_pos
                previous_pos = site.position
                assert ancestral_state[index] == site.ancestral_state
                assert site.id == index
                for mutation in site.mutations:
                    m2 = ts.mutation(mutation.id)
                    assert m2 == mutation
                    assert mutation.site == site.id
                    assert mutation.site == mutations.site[mutation_index]
                    assert mutation.node == mutations.node[mutation_index]
                    assert mutation.parent == mutations.parent[mutation_index]
                    assert mutation.id == mutation_index
                    assert derived_state[mutation_index] == mutation.derived_state
                    mutation_index += 1
                some_sites = True
            total_sites = 0
            for tree in ts.trees():
                assert len(list(tree.sites())) == tree.num_sites
                total_sites += tree.num_sites
            assert ts.num_sites == total_sites
            assert mutation_index == len(mutations)
        assert some_sites

    def verify_mutations(self, ts):
        other_mutations = []
        for site in ts.sites():
            for mutation in site.mutations:
                other_mutations.append(mutation)
        mutations = list(ts.mutations())
        assert ts.num_mutations == len(other_mutations)
        assert ts.num_mutations == len(mutations)
        for mut, other_mut in zip(mutations, other_mutations):
            assert mut == other_mut

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_sites_mutations(self, ts):
        # Check that the mutations iterator returns the correct values.
        self.verify_mutations(ts)

    def test_removed_methods(self):
        ts = tskit.TableCollection(1).tree_sequence()
        with pytest.raises(NotImplementedError):
            ts.get_num_records()
        with pytest.raises(NotImplementedError):
            ts.diffs()
        with pytest.raises(NotImplementedError):
            ts.newick_trees()
        with pytest.raises(NotImplementedError):
            ts.to_nexus()

    def test_dump_pathlib(self, ts_fixture, tmp_path):
        path = tmp_path / "tmp.trees"
        assert path.exists
        assert path.is_file
        ts_fixture.dump(path)
        other_ts = tskit.load(path)
        assert ts_fixture.tables == other_ts.tables

    @pytest.mark.skipif(platform.system() == "Windows", reason="Windows doesn't raise")
    def test_dump_load_errors(self):
        ts = msprime.simulate(5, random_seed=1)
        # Try to dump/load files we don't have access to or don't exist.
        for func in [ts.dump, tskit.load]:
            for f in ["/", "/test.trees", "/dir_does_not_exist/x.trees"]:
                with pytest.raises(OSError):
                    func(f)
                try:
                    func(f)
                except OSError as e:
                    message = str(e)
                    assert len(message) > 0
            f = "/" + 4000 * "x"
            with pytest.raises(OSError):
                func(f)
            try:
                func(f)
            except OSError as e:
                message = str(e)
            assert "File name too long" in message
            for bad_filename in [[], None, {}]:
                with pytest.raises(TypeError):
                    func(bad_filename)

    def test_zlib_compression_warning(self, ts_fixture, tmp_path):
        temp_file = tmp_path / "tmp.trees"
        with warnings.catch_warnings(record=True) as w:
            ts_fixture.dump(temp_file, zlib_compression=True)
            assert len(w) == 1
            assert issubclass(w[0].category, RuntimeWarning)
        with warnings.catch_warnings(record=True) as w:
            ts_fixture.dump(temp_file, zlib_compression=False)
            assert len(w) == 0

    def test_tables_sequence_length_round_trip(self):
        for sequence_length in [0.1, 1, 10, 100]:
            ts = msprime.simulate(5, length=sequence_length, random_seed=1)
            assert ts.sequence_length == sequence_length
            tables = ts.tables
            assert tables.sequence_length == sequence_length
            new_ts = tables.tree_sequence()
            assert new_ts.sequence_length == sequence_length

    def test_migrations(self):
        ts = msprime.simulate(
            population_configurations=[
                msprime.PopulationConfiguration(10),
                msprime.PopulationConfiguration(10),
            ],
            migration_matrix=[[0, 1], [1, 0]],
            random_seed=2,
            record_migrations=True,
        )
        assert ts.num_migrations > 0
        migrations = list(ts.migrations())
        assert len(migrations) == ts.num_migrations
        for migration in migrations:
            assert migration.source in [0, 1]
            assert migration.dest in [0, 1]
            assert migration.time > 0
            assert migration.left == 0
            assert migration.right == 1
            assert 0 <= migration.node < ts.num_nodes

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_len_trees(self, ts):
        tree_iter = ts.trees()
        assert len(tree_iter) == ts.num_trees

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_list(self, ts):
        for kwargs in [{}, {"tracked_samples": ts.samples()}]:
            tree_list = ts.aslist(**kwargs)
            assert len(tree_list) == ts.num_trees
            assert len(set(map(id, tree_list))) == ts.num_trees
            for index, tree in enumerate(tree_list):
                assert index == tree.index
            for t1, t2 in zip(tree_list, ts.trees(**kwargs)):
                assert t1 == t2
                assert t1.parent_dict == t2.parent_dict
                if "tracked_samples" in kwargs:
                    assert t1.num_tracked_samples() == ts.num_samples
                    assert t2.num_tracked_samples() == ts.num_samples
                else:
                    assert t1.num_tracked_samples() == 0
                    assert t2.num_tracked_samples() == 0

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_reversed_trees(self, ts):
        index = ts.num_trees - 1
        tree_list = ts.aslist()
        for tree in reversed(ts.trees()):
            assert tree.index == index
            t2 = tree_list[index]
            assert tree.interval == t2.interval
            assert tree.parent_dict == t2.parent_dict
            index -= 1

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_at_index(self, ts):
        for kwargs in [{}, {"tracked_samples": ts.samples()}]:
            tree_list = ts.aslist(**kwargs)
            for index in list(range(ts.num_trees)) + [-1]:
                t1 = tree_list[index]
                t2 = ts.at_index(index, **kwargs)
                assert t1 == t2
                assert t1.interval == t2.interval
                assert t1.parent_dict == t2.parent_dict
                if "tracked_samples" in kwargs:
                    assert t2.num_tracked_samples() == ts.num_samples
                else:
                    assert t2.num_tracked_samples() == 0

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_at(self, ts):
        for kwargs in [{}, {"tracked_samples": ts.samples()}]:
            tree_list = ts.aslist(**kwargs)
            for t1 in tree_list:
                left, right = t1.interval
                mid = left + (right - left) / 2
                for pos in [left, left + 1e-9, mid, right - 1e-9]:
                    t2 = ts.at(pos, **kwargs)
                    assert t1 == t2
                    assert t1.interval == t2.interval
                    assert t1.parent_dict == t2.parent_dict
                if right < ts.sequence_length:
                    t2 = ts.at(right, **kwargs)
                    t3 = tree_list[t1.index + 1]
                    assert t3 == t2
                    assert t3.interval == t2.interval
                    assert t3.parent_dict == t2.parent_dict
                if "tracked_samples" in kwargs:
                    assert t2.num_tracked_samples() == ts.num_samples
                else:
                    assert t2.num_tracked_samples() == 0

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_sequence_iteration(self, ts):
        for table_name in ts.tables_dict.keys():
            sequence = getattr(ts, table_name)()
            length = getattr(ts, "num_" + table_name)
            # Test __iter__
            i = None
            for i, n in enumerate(sequence):
                assert i == n.id
            if i is not None:
                assert n.id == (length - 1 if length else 0)
            if table_name == "mutations":
                # Mutations are not currently sequences, so have no len or idx access
                with pytest.raises(TypeError):
                    len(sequence)
                if length != 0:
                    with pytest.raises(TypeError):
                        sequence[0]
            else:
                # Test __len__
                assert len(sequence) == length
                # Test __getitem__ on the last item in the sequence
                if length != 0:
                    assert sequence[length - 1] == n  # +ive indexing
                    assert sequence[-1] == n  # -ive indexing
                with pytest.raises(IndexError):
                    sequence[length]
                # Test reverse
                i = None
                for i, n in enumerate(reversed(sequence)):
                    assert i == length - 1 - n.id
                if i is not None:
                    assert n.id == 0

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_load_tables(self, ts):
        tables = ts.dump_tables()
        tables.drop_index()

        # Tables not in tc not rebuilt as per default, so error
        with pytest.raises(
            _tskit.LibraryError, match="Table collection must be indexed"
        ):
            assert tskit.TreeSequence.load_tables(tables).dump_tables().has_index()

        # Tables not in tc, but rebuilt
        assert (
            tskit.TreeSequence.load_tables(tables, build_indexes=True)
            .dump_tables()
            .has_index()
        )

        tables.build_index()
        # Tables in tc, not rebuilt
        assert (
            tskit.TreeSequence.load_tables(tables, build_indexes=False)
            .dump_tables()
            .has_index()
        )
        # Tables in tc, and rebuilt
        assert tskit.TreeSequence.load_tables(tables).dump_tables().has_index()

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_html_repr(self, ts):
        html = ts._repr_html_()
        # Parse to check valid
        ElementTree.fromstring(html)
        assert len(html) > 5000
        assert f"<tr><td>Trees</td><td>{ts.num_trees:,}</td></tr>" in html
        assert f"<tr><td>Time Units</td><td>{ts.time_units}</td></tr>" in html
        for table in ts.tables.table_name_map:
            assert f"<td>{table.capitalize()}</td>" in html
        if ts.num_provenances > 0:
            assert (
                f"<td>{json.loads(ts.provenance(0).record)['software']['name']}</td>"
                in html
            )

    def test_bad_provenance(self, ts_fixture):
        tables = ts_fixture.dump_tables()
        tables.provenances.add_row("bad", "bad")
        ts = tables.tree_sequence()
        assert "Could not parse provenance" in ts._repr_html_()

    def test_provenance_summary_html(self, ts_fixture):
        tables = ts_fixture.tables
        for _ in range(20):
            # Add a row with isotimestamp
            tables.provenances.add_row("foo", "bar")
        assert "... 15 more" in tables.tree_sequence()._repr_html_()

    def test_html_repr_limit(self, ts_fixture):
        tables = ts_fixture.tables
        d = {n: n for n in range(50)}
        d[0] = "N" * 200
        tables.metadata = d
        ts = tables.tree_sequence()
        assert "... and 20 more" in ts._repr_html_()
        assert "NN..." in ts._repr_html_()

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_str(self, ts):
        s = str(ts)
        assert len(s) > 999
        assert re.search(rf"║Trees *│ *{ts.num_trees}║", s)
        assert re.search(rf"║Time Units *│ *{ts.time_units}║", s)
        for table in ts.tables.table_name_map:
            assert re.search(rf"║{table.capitalize()} *│", s)

    @pytest.mark.skip("FIXME nbytes")
    def test_nbytes(self, tmp_path, ts_fixture):
        ts_fixture.dump(tmp_path / "tables")
        store = kastore.load(tmp_path / "tables")
        for v in store.values():
            # Check we really have data in every field
            assert v.nbytes > 0
        nbytes = sum(
            array.nbytes
            for name, array in store.items()
            # nbytes is the size of asdict, so exclude file format items
            if name not in ["format/version", "format/name", "uuid"]
        )
        assert nbytes == ts_fixture.nbytes

    def test_equals(self):
        # Here we don't use the fixture as we'd like to run the same sim twice
        pop_configs = [msprime.PopulationConfiguration(5) for _ in range(2)]
        migration_matrix = [[0, 1], [1, 0]]
        t1 = msprime.simulate(
            population_configurations=pop_configs,
            migration_matrix=migration_matrix,
            mutation_rate=1,
            record_migrations=True,
            random_seed=1,
        )
        t2 = msprime.simulate(
            population_configurations=pop_configs,
            migration_matrix=migration_matrix,
            mutation_rate=1,
            record_migrations=True,
            random_seed=1,
        )

        assert t1 == t1
        assert t1 == t1.dump_tables().tree_sequence()
        assert t1.dump_tables().tree_sequence() == t1

        # The provenances may or may not be equal depending on the clock
        # precision for record. So clear them first.
        tb1 = t1.dump_tables()
        tb2 = t2.dump_tables()
        tb1.provenances.clear()
        tb2.provenances.clear()
        t1 = tb1.tree_sequence()
        t2 = tb2.tree_sequence()

        assert t1 == t2
        assert t1 == t2
        assert not (t1 != t2)
        # We don't do more as this is the same code path as TableCollection.__eq__

    def test_equals_options(self, ts_fixture):
        t1 = ts_fixture
        # Take a copy
        t2 = ts_fixture.dump_tables().tree_sequence()

        def modify(ts, func):
            tc = ts.dump_tables()
            func(tc)
            return tc.tree_sequence()

        t1 = modify(t1, lambda tc: tc.provenances.add_row("random stuff"))
        assert not (t1 == t2)
        assert t1.equals(t2, ignore_provenance=True)
        assert t2.equals(t1, ignore_provenance=True)
        assert not (t1.equals(t2))
        assert not (t2.equals(t1))
        t1 = modify(t1, lambda tc: tc.provenances.clear())
        t2 = modify(t2, lambda tc: tc.provenances.clear())
        assert t1.equals(t2)
        assert t2.equals(t1)

        tc = t1.dump_tables()
        tc.metadata_schema = tskit.MetadataSchema({"codec": "json", "type": "object"})
        t1 = tc.tree_sequence()
        tc = t1.dump_tables()
        tc.metadata = {"hello": "world"}
        t1 = tc.tree_sequence()

        assert not t1.equals(t2)
        assert t1.equals(t2, ignore_ts_metadata=True)
        assert not t2.equals(t1)
        assert t2.equals(t1, ignore_ts_metadata=True)
        tc = t2.dump_tables()
        tc.metadata_schema = t1.metadata_schema
        t2 = tc.tree_sequence()
        assert not t1.equals(t2)
        assert t1.equals(t2, ignore_ts_metadata=True)
        assert not t2.equals(t1)
        assert t2.equals(t1, ignore_ts_metadata=True)

        t1 = modify(t1, lambda tc: tc.provenances.add_row("random stuff"))
        assert not t1.equals(t2)
        assert not t1.equals(t2, ignore_ts_metadata=True)
        assert not t1.equals(t2, ignore_provenance=True)
        assert t1.equals(t2, ignore_ts_metadata=True, ignore_provenance=True)

        t1 = modify(t1, lambda tc: tc.provenances.clear())
        t2 = modify(t2, lambda tc: setattr(tc, "metadata", t1.metadata))  # noqa: B010
        assert t1.equals(t2)
        assert t2.equals(t1)

        # Empty out tables to test ignore_tables flag
        tc = t2.dump_tables()
        tc.individuals.truncate(0)
        tc.nodes.truncate(0)
        tc.edges.truncate(0)
        tc.migrations.truncate(0)
        tc.sites.truncate(0)
        tc.mutations.truncate(0)
        tc.populations.truncate(0)
        t2 = tc.tree_sequence()
        assert not t1.equals(t2)
        assert t1.equals(t2, ignore_tables=True)
        # Empty out reference to test ignore_reference_sequence flag
        tc = t1.dump_tables()
        tc.reference_sequence.clear()
        t2 = tc.tree_sequence()
        assert not t1.equals(t2)
        assert t1.equals(t2, ignore_reference_sequence=True)
        # Make t1 and t2 equal again
        t2 = t1.dump_tables().tree_sequence()
        assert t1.equals(t2)
        assert t2.equals(t1)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_tree_node_edges(self, ts):
        edge_visited = np.zeros(ts.num_edges, dtype=bool)
        for tree in ts.trees():
            mapping = tree.edge_array
            node_mapped = mapping >= 0
            edge_visited[mapping[node_mapped]] = True
            # Note that tree.nodes() does not necessarily list all the nodes
            # in the tree topology, only the ones that descend from a root.
            # Therefore if not all the topological trees in a single `Tree` have
            # a root, we can have edges above nodes that are not listed. This
            # happens, for example, in a tree with no sample nodes.
            assert np.sum(node_mapped) >= len(list(tree.nodes())) - tree.num_roots
            for u in tree.nodes():
                if tree.parent(u) == tskit.NULL:
                    assert mapping[u] == tskit.NULL
                else:
                    edge = ts.edge(mapping[u])
                    assert edge.child == u
                    assert edge.left <= tree.interval.left
                    assert edge.right >= tree.interval.right
        assert np.all(edge_visited)

    def verify_individual_vectors(self, ts):
        verify_times = np.repeat(np.nan, ts.num_individuals)
        verify_populations = np.repeat(tskit.NULL, ts.num_individuals)
        for ind in ts.individuals():
            if len(ind.nodes) > 0:
                t = {ts.node(n).time for n in ind.nodes}
                p = {ts.node(n).population for n in ind.nodes}
                assert len(t) <= 1
                assert len(p) <= 1
                verify_times[ind.id] = t.pop()
                verify_populations[ind.id] = p.pop()

        times = ts.individuals_time
        populations = ts.individuals_population
        assert np.array_equal(times, verify_times, equal_nan=True)
        assert np.array_equal(populations, verify_populations, equal_nan=True)
        times2 = ts.individuals_time
        populations2 = ts.individuals_population
        assert np.array_equal(times, times2, equal_nan=True)
        assert np.array_equal(populations, populations2, equal_nan=True)
        # check aliases also
        times3 = ts.individual_times
        populations3 = ts.individual_populations
        assert np.array_equal(times, times3, equal_nan=True)
        assert np.array_equal(populations, populations3, equal_nan=True)

    def test_individuals_population_errors(self):
        t = tskit.TableCollection(sequence_length=1)
        t.individuals.add_row()
        t.individuals.add_row()
        for j in range(2):
            t.populations.add_row()
            t.nodes.add_row(time=0, population=j, individual=0)
        ts = t.tree_sequence()
        with pytest.raises(
            _tskit.LibraryError, match="TSK_ERR_INDIVIDUAL_POPULATION_MISMATCH"
        ):
            _ = ts.individuals_population
        # inconsistent but NULL populations are also an error
        t.nodes.clear()
        t.nodes.add_row(time=0, population=1, individual=0)
        t.nodes.add_row(time=0, population=tskit.NULL, individual=0)
        ts = t.tree_sequence()
        with pytest.raises(
            _tskit.LibraryError, match="TSK_ERR_INDIVIDUAL_POPULATION_MISMATCH"
        ):
            _ = ts.individuals_population
        t.nodes.clear()
        t.nodes.add_row(time=0, population=tskit.NULL, individual=1)
        t.nodes.add_row(time=0, population=0, individual=1)
        ts = t.tree_sequence()
        with pytest.raises(
            _tskit.LibraryError, match="TSK_ERR_INDIVIDUAL_POPULATION_MISMATCH"
        ):
            _ = ts.individuals_population

    def test_individuals_time_errors(self):
        t = tskit.TableCollection(sequence_length=1)
        t.individuals.add_row()
        for j in range(2):
            t.nodes.add_row(time=j, individual=0)
        ts = t.tree_sequence()
        with pytest.raises(
            _tskit.LibraryError, match="TSK_ERR_INDIVIDUAL_TIME_MISMATCH"
        ):
            _ = ts.individuals_time

    @pytest.mark.parametrize("n", [1, 10])
    def test_individual_vectors(self, n):
        d = msprime.Demography.island_model([10] * n, 0.1)
        ts = msprime.sim_ancestry(
            {pop.name: 10 for pop in d.populations},
            demography=d,
            random_seed=100 + n,
            model="dtwf",
        )
        ts = tsutil.insert_random_consistent_individuals(ts, seed=100 + n)
        assert ts.num_individuals > 10
        self.verify_individual_vectors(ts)

    def test_individuals_location_errors(self):
        t = tskit.TableCollection(sequence_length=1)
        t.individuals.add_row(location=[1.0, 2.0])
        t.individuals.add_row(location=[0.0])
        ts = t.tree_sequence()
        with pytest.raises(ValueError, match="locations"):
            _ = ts.individuals_location

        t.clear()
        t.individuals.add_row(location=[1.0, 2.0])
        t.individuals.add_row(location=[])
        t.individuals.add_row(location=[1.0, 2.0])
        t.individuals.add_row(location=[])
        ts = t.tree_sequence()
        with pytest.raises(ValueError, match="locations"):
            _ = ts.individuals_location

    @pytest.mark.parametrize("nlocs", [0, 1, 4])
    @pytest.mark.parametrize("num_indivs", [0, 3])
    def test_individuals_location(self, nlocs, num_indivs):
        t = tskit.TableCollection(sequence_length=1)
        locs = np.array([j + np.arange(nlocs) for j in range(num_indivs)])
        if len(locs) == 0:
            locs = locs.reshape((num_indivs, 0))
        for j in range(num_indivs):
            t.individuals.add_row(location=locs[j])
        ts = t.tree_sequence()
        ts_locs = ts.individuals_location
        assert locs.shape == ts_locs.shape
        assert np.array_equal(locs, ts_locs)
        locs2 = ts.individuals_location
        assert np.array_equal(ts_locs, locs2)
        # test alias
        locs3 = ts.individual_locations
        assert np.array_equal(ts_locs, locs3)

    def verify_individual_properties(self, ts):
        for ind in ts.individuals():
            times = [ts.node(n).time for n in ind.nodes]
            if len(set(times)) > 1:
                with pytest.raises(ValueError, match="mis-matched times"):
                    _ = ind.time
            elif len(times) == 0:
                assert tskit.is_unknown_time(ind.time)
            else:
                assert len(set(times)) == 1
                assert times[0] == ind.time
                # test accessing more than once in case we mess up with {}.pop()
                assert times[0] == ind.time
            pops = [ts.node(n).population for n in ind.nodes]
            if len(set(pops)) > 1:
                with pytest.raises(ValueError, match="mis-matched populations"):
                    _ = ind.population
            elif len(pops) == 0:
                assert ind.population is tskit.NULL
            else:
                assert len(set(pops)) == 1
                assert ind.population == pops[0]
                # test accessing more than once in case we mess up with {}.pop()
                assert ind.population == pops[0]

    def test_individual_getter_population(self):
        tables = tskit.TableCollection(sequence_length=1)
        for _ in range(2):
            tables.populations.add_row()
        pop_list = [
            ((), tskit.NULL),
            ((tskit.NULL,), tskit.NULL),
            ((1,), 1),
            ((1, 1, 1), 1),
            ((tskit.NULL, 1), "ERR"),
            ((0, tskit.NULL), "ERR"),
            ((0, 1), "ERR"),
        ]
        for pops, _ in pop_list:
            j = tables.individuals.add_row()
            for p in pops:
                tables.nodes.add_row(time=0, population=p, individual=j)
        ts = tables.tree_sequence()
        for ind, (_, p) in zip(ts.individuals(), pop_list):
            if p == "ERR":
                with pytest.raises(ValueError, match="mis-matched populations"):
                    _ = ind.population
            else:
                assert p == ind.population

    def test_individual_getter_time(self):
        tables = tskit.TableCollection(sequence_length=1)
        time_list = [
            ((), tskit.UNKNOWN_TIME),
            ((0.0,), 0.0),
            ((1, 1, 1), 1),
            ((4.0, 1), "ERR"),
            ((0, 4.0), "ERR"),
        ]
        for times, _ in time_list:
            j = tables.individuals.add_row()
            for t in times:
                tables.nodes.add_row(time=t, individual=j)
        ts = tables.tree_sequence()
        for ind, (_, t) in zip(ts.individuals(), time_list):
            if t == "ERR":
                with pytest.raises(ValueError, match="mis-matched times"):
                    _ = ind.time
            elif tskit.is_unknown_time(t):
                assert tskit.is_unknown_time(ind.time)
            else:
                assert t == ind.time

    @pytest.mark.parametrize("n", [1, 10])
    def test_individual_properties(self, n):
        # tests for the .time and .population attributes of
        # the Individual class
        d = msprime.Demography.island_model([10] * n, 0.1)
        ts = msprime.sim_ancestry(
            {pop.name: int(150 / n) for pop in d.populations},
            demography=d,
            random_seed=100 + n,
            model="dtwf",
        )
        ts = tsutil.insert_random_consistent_individuals(ts, seed=100 + n)
        assert ts.num_individuals > 10
        self.verify_individual_properties(ts)
        ts = tsutil.insert_random_ploidy_individuals(ts, seed=100 + n)
        assert ts.num_individuals > 10
        self.verify_individual_properties(ts)

    @pytest.mark.parametrize(
        "array",
        [
            "individuals_flags",
            "nodes_time",
            "nodes_flags",
            "nodes_population",
            "nodes_individual",
            "edges_left",
            "edges_right",
            "edges_parent",
            "edges_child",
            "sites_position",
            "mutations_site",
            "mutations_node",
            "mutations_parent",
            "mutations_time",
            "migrations_left",
            "migrations_right",
            "migrations_node",
            "migrations_source",
            "migrations_dest",
            "migrations_time",
            "indexes_edge_insertion_order",
            "indexes_edge_removal_order",
        ],
    )
    def test_array_attr_properties(self, ts_fixture, array):
        ts = ts_fixture
        a = getattr(ts, array)
        assert isinstance(a, np.ndarray)
        with pytest.raises(AttributeError):
            setattr(ts, array, None)
        with pytest.raises(AttributeError):
            delattr(ts, array)
        with pytest.raises(ValueError, match="read-only"):
            a[:] = 1

    def test_arrays_equal_to_tables(self, ts_fixture):
        ts = ts_fixture
        tables = ts.tables

        assert_array_equal(ts.individuals_flags, tables.individuals.flags)

        assert_array_equal(ts.nodes_flags, tables.nodes.flags)
        assert_array_equal(ts.nodes_population, tables.nodes.population)
        assert_array_equal(ts.nodes_time, tables.nodes.time)
        assert_array_equal(ts.nodes_individual, tables.nodes.individual)

        assert_array_equal(ts.edges_left, tables.edges.left)
        assert_array_equal(ts.edges_right, tables.edges.right)
        assert_array_equal(ts.edges_parent, tables.edges.parent)
        assert_array_equal(ts.edges_child, tables.edges.child)

        assert_array_equal(ts.sites_position, tables.sites.position)

        assert_array_equal(ts.mutations_site, tables.mutations.site)
        assert_array_equal(ts.mutations_node, tables.mutations.node)
        assert_array_equal(ts.mutations_parent, tables.mutations.parent)
        assert_array_equal(ts.mutations_time, tables.mutations.time)

        assert_array_equal(ts.migrations_left, tables.migrations.left)
        assert_array_equal(ts.migrations_right, tables.migrations.right)
        assert_array_equal(ts.migrations_node, tables.migrations.node)
        assert_array_equal(ts.migrations_source, tables.migrations.source)
        assert_array_equal(ts.migrations_dest, tables.migrations.dest)
        assert_array_equal(ts.migrations_time, tables.migrations.time)

        assert_array_equal(
            ts.indexes_edge_insertion_order, tables.indexes.edge_insertion_order
        )
        assert_array_equal(
            ts.indexes_edge_removal_order, tables.indexes.edge_removal_order
        )

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_impute_unknown_mutations_time(self, ts):
        # Tests for method='min'
        imputed_time = ts.impute_unknown_mutations_time(method="min")
        mutations = ts.tables.mutations
        nodes_time = ts.nodes_time
        table_time = np.zeros(len(mutations))

        for mut_idx, mut in enumerate(mutations):
            if tskit.is_unknown_time(mut.time):
                node_time = nodes_time[mut.node]
                table_time[mut_idx] = node_time
            else:
                table_time[mut_idx] = mut.time

        assert np.allclose(imputed_time, table_time, rtol=1e-10, atol=1e-10)

        # Check we have valid times
        tables = ts.dump_tables()
        tables.mutations.time = imputed_time
        tables.sort()
        tables.tree_sequence()

        # Test for unallowed methods
        with pytest.raises(
            ValueError, match="Mutations time imputation method must be chosen"
        ):
            ts.impute_unknown_mutations_time(method="foobar")


class TestSimplify:
    # This class was factored out of the old TestHighlevel class 2022-12-13,
    # and is a mishmash of different testing paradigms. There is some valuable
    # testing done here, so it would be good to fully bring it up to date.

    def verify_simplify_provenance(self, ts):
        new_ts = ts.simplify()
        assert new_ts.num_provenances == ts.num_provenances + 1
        old = list(ts.provenances())
        new = list(new_ts.provenances())
        assert old == new[:-1]
        # TODO call verify_provenance on this.
        assert len(new[-1].timestamp) > 0
        assert len(new[-1].record) > 0

        new_ts = ts.simplify(record_provenance=False)
        assert new_ts.tables.provenances == ts.tables.provenances

    def verify_simplify_topology(self, ts, sample):
        new_ts, node_map = ts.simplify(sample, map_nodes=True)
        if len(sample) == 0:
            assert new_ts.num_nodes == 0
            assert new_ts.num_edges == 0
            assert new_ts.num_sites == 0
            assert new_ts.num_mutations == 0
        elif len(sample) == 1:
            assert new_ts.num_nodes == 1
            assert new_ts.num_edges == 0
        # The output samples should be 0...n
        assert new_ts.num_samples == len(sample)
        assert list(range(len(sample))) == list(new_ts.samples())
        for j in range(new_ts.num_samples):
            assert node_map[sample[j]] == j
        for u in range(ts.num_nodes):
            old_node = ts.node(u)
            if node_map[u] != tskit.NULL:
                new_node = new_ts.node(node_map[u])
                assert old_node.time == new_node.time
                assert old_node.population == new_node.population
                assert old_node.metadata == new_node.metadata
        for u in sample:
            old_node = ts.node(u)
            new_node = new_ts.node(node_map[u])
            assert old_node.flags == new_node.flags
            assert old_node.time == new_node.time
            assert old_node.population == new_node.population
            assert old_node.metadata == new_node.metadata
        old_trees = ts.trees()
        old_tree = next(old_trees)
        assert ts.get_num_trees() >= new_ts.get_num_trees()
        for new_tree in new_ts.trees():
            new_left, new_right = new_tree.get_interval()
            old_left, old_right = old_tree.get_interval()
            # Skip ahead on the old tree until new_left is within its interval
            while old_right <= new_left:
                old_tree = next(old_trees)
                old_left, old_right = old_tree.get_interval()
            # If the MRCA of all pairs of samples is the same, then we have the
            # same information. We limit this to at most 500 pairs
            pairs = itertools.islice(itertools.combinations(sample, 2), 500)
            for pair in pairs:
                mapped_pair = [node_map[u] for u in pair]
                mrca1 = old_tree.get_mrca(*pair)
                mrca2 = new_tree.get_mrca(*mapped_pair)
                if mrca1 == tskit.NULL:
                    assert mrca2 == mrca1
                else:
                    assert mrca2 == node_map[mrca1]
                    assert old_tree.get_time(mrca1) == new_tree.get_time(mrca2)
                    assert old_tree.get_population(mrca1) == new_tree.get_population(
                        mrca2
                    )

    def verify_simplify_equality(self, ts, sample):
        for filter_sites in [False, True]:
            s1, node_map1 = ts.simplify(
                sample, map_nodes=True, filter_sites=filter_sites
            )
            t1 = s1.dump_tables()
            s2, node_map2 = simplify_tree_sequence(
                ts, sample, filter_sites=filter_sites
            )
            t2 = s2.dump_tables()
            assert s1.num_samples == len(sample)
            assert s2.num_samples == len(sample)
            assert all(node_map1 == node_map2)
            assert t1.individuals == t2.individuals
            assert t1.nodes == t2.nodes
            assert t1.edges == t2.edges
            assert t1.migrations == t2.migrations
            assert t1.sites == t2.sites
            assert t1.mutations == t2.mutations
            assert t1.populations == t2.populations

    def verify_simplify_variants(self, ts, sample):
        subset = ts.simplify(sample)
        sample_map = {u: j for j, u in enumerate(ts.samples())}
        # Need to map IDs back to their sample indexes
        s = np.array([sample_map[u] for u in sample])
        # Build a map of genotypes by position
        full_genotypes = {}
        for variant in ts.variants(isolated_as_missing=False):
            alleles = [variant.alleles[g] for g in variant.genotypes]
            full_genotypes[variant.position] = alleles
        for variant in subset.variants(isolated_as_missing=False):
            if variant.position in full_genotypes:
                a1 = [full_genotypes[variant.position][u] for u in s]
                a2 = [variant.alleles[g] for g in variant.genotypes]
                assert a1 == a2

    def verify_tables_api_equality(self, ts):
        for samples in [None, list(ts.samples()), ts.samples()]:
            tables = ts.dump_tables()
            tables.simplify(samples=samples)
            tables.assert_equals(
                ts.simplify(samples=samples).tables, ignore_timestamps=True
            )

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_simplify_tables_equality(self, ts):
        # Can't simplify edges with metadata
        if ts.tables.edges.metadata_schema == tskit.MetadataSchema(schema=None):
            self.verify_tables_api_equality(ts)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_simplify_provenance(self, ts):
        # Can't simplify edges with metadata
        if ts.tables.edges.metadata_schema == tskit.MetadataSchema(schema=None):
            self.verify_simplify_provenance(ts)

    # TODO this test needs to be broken up into discrete bits, so that we can
    # test them independently. A way of getting a random-ish subset of samples
    # from the pytest param would be useful.
    @pytest.mark.slow
    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_simplify(self, ts):
        # Can't simplify edges with metadata
        if ts.tables.edges.metadata_schema == tskit.MetadataSchema(schema=None):
            n = ts.num_samples
            sample_sizes = {0}
            if n > 1:
                sample_sizes |= {1}
            if n > 2:
                sample_sizes |= {2, max(2, n // 2), n - 1}
            for k in sample_sizes:
                subset = random.sample(list(ts.samples()), k)
                self.verify_simplify_topology(ts, subset)
                self.verify_simplify_equality(ts, subset)
                self.verify_simplify_variants(ts, subset)

    def test_simplify_bugs(self):
        prefix = os.path.join(os.path.dirname(__file__), "data", "simplify-bugs")
        j = 1
        while True:
            nodes_file = os.path.join(prefix, f"{j:02d}-nodes.txt")
            if not os.path.exists(nodes_file):
                break
            edges_file = os.path.join(prefix, f"{j:02d}-edges.txt")
            sites_file = os.path.join(prefix, f"{j:02d}-sites.txt")
            mutations_file = os.path.join(prefix, f"{j:02d}-mutations.txt")
            with open(nodes_file) as nodes, open(edges_file) as edges, open(
                sites_file
            ) as sites, open(mutations_file) as mutations:
                ts = tskit.load_text(
                    nodes=nodes,
                    edges=edges,
                    sites=sites,
                    mutations=mutations,
                    strict=False,
                )
            samples = list(ts.samples())
            self.verify_simplify_equality(ts, samples)
            j += 1
        assert j > 1

    def test_simplify_migrations_fails(self):
        ts = msprime.simulate(
            population_configurations=[
                msprime.PopulationConfiguration(10),
                msprime.PopulationConfiguration(10),
            ],
            migration_matrix=[[0, 1], [1, 0]],
            random_seed=2,
            record_migrations=True,
        )
        assert ts.num_migrations > 0
        # We don't support simplify with migrations, so should fail.
        with pytest.raises(_tskit.LibraryError):
            ts.simplify()

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_no_update_sample_flags_no_filter_nodes(self, ts):
        # Can't simplify edges with metadata
        if ts.tables.edges.metadata_schema == tskit.MetadataSchema(schema=None):
            k = min(ts.num_samples, 3)
            subset = ts.samples()[:k]
            ts1 = ts.simplify(subset)
            ts2 = ts.simplify(subset, update_sample_flags=False, filter_nodes=False)
            assert ts1.num_samples == len(subset)
            assert ts2.num_samples == ts.num_samples
            assert ts1.num_edges == ts2.num_edges
            assert ts2.tables.nodes == ts.tables.nodes


class TestMinMaxTime:
    def get_example_tree_sequence(self, use_unknown_time):
        """
        Min time is set to 0.1.
        Max time is set to 2.0.
        """
        tables = tskit.TableCollection(sequence_length=2)
        tables.nodes.add_row(flags=1, time=0.1)
        tables.nodes.add_row(flags=1, time=0.1)
        tables.nodes.add_row(flags=1, time=0.1)
        tables.nodes.add_row(flags=0, time=1)
        tables.nodes.add_row(flags=0, time=2)
        tables.edges.add_row(left=0, right=2, parent=3, child=0)
        tables.edges.add_row(left=0, right=2, parent=3, child=1)
        tables.edges.add_row(left=0, right=2, parent=4, child=2)
        tables.edges.add_row(left=0, right=2, parent=4, child=3)
        tables.sites.add_row(position=0, ancestral_state="0")
        tables.sites.add_row(position=1, ancestral_state="0")
        if use_unknown_time:
            tables.mutations.add_row(
                site=0, node=2, derived_state="1", time=tskit.UNKNOWN_TIME
            )
            tables.mutations.add_row(
                site=1, node=3, derived_state="1", time=tskit.UNKNOWN_TIME
            )
        else:
            tables.mutations.add_row(site=0, node=2, derived_state="1", time=0.5)
            tables.mutations.add_row(site=1, node=3, derived_state="1", time=1.5)
        ts = tables.tree_sequence()
        return ts

    def get_empty_tree_sequence(self):
        """
        Min time is initialised to positive infinity.
        Max time is initialised to negative infinity.
        """
        tables = tskit.TableCollection(sequence_length=2)
        ts = tables.tree_sequence()
        return ts

    def test_example(self):
        ts = self.get_example_tree_sequence(use_unknown_time=False)
        expected_min_time = min(ts.nodes_time.min(), ts.mutations_time.min())
        expected_max_time = max(ts.nodes_time.max(), ts.mutations_time.max())
        assert ts.min_time == expected_min_time
        assert ts.max_time == expected_max_time

    def test_example_unknown_mutation_times(self):
        ts = self.get_example_tree_sequence(use_unknown_time=True)
        expected_min_time = ts.nodes_time.min()
        expected_max_time = ts.nodes_time.max()
        assert ts.min_time == expected_min_time
        assert ts.max_time == expected_max_time

    def test_empty(self):
        ts = self.get_empty_tree_sequence()
        assert ts.min_time == np.inf
        assert ts.max_time == -np.inf


class TestSiteAlleles:
    def test_no_mutations(self):
        tables = tskit.TableCollection(sequence_length=1)
        tables.sites.add_row(0, ancestral_state="")
        site = tables.tree_sequence().site(0)
        assert site.alleles == {""}

    @pytest.mark.parametrize("k", range(5))
    def test_k_mutations(self, k):
        tables = tskit.TableCollection(sequence_length=1)
        tables.sites.add_row(0, ancestral_state="ABC")
        tables.nodes.add_row(1, 0)
        tables.nodes.add_row(1, 0)  # will not have any mutations => missing
        for j in range(k):
            tables.mutations.add_row(site=0, node=0, derived_state=str(j))
        ts = tables.tree_sequence()
        variant = next(ts.variants())
        assert variant.has_missing_data
        assert len(variant.site.alleles) == k + 1
        assert "ABC" in variant.site.alleles
        assert variant.site.alleles == set(variant.alleles[:-1])


class TestEdgeDiffs:
    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_correct_trees_forward(self, ts):
        parent = np.full(ts.num_nodes + 1, tskit.NULL, dtype=np.int32)
        for edge_diff, tree in itertools.zip_longest(ts.edge_diffs(), ts.trees()):
            assert edge_diff.interval == tree.interval
            for edge in edge_diff.edges_out:
                parent[edge.child] = tskit.NULL
            for edge in edge_diff.edges_in:
                parent[edge.child] = edge.parent
            assert_array_equal(parent, tree.parent_array)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_correct_trees_reverse(self, ts):
        parent = np.full(ts.num_nodes + 1, tskit.NULL, dtype=np.int32)
        iterator = itertools.zip_longest(
            ts.edge_diffs(direction=tskit.REVERSE), reversed(ts.trees())
        )
        for edge_diff, tree in iterator:
            assert edge_diff.interval == tree.interval
            for edge in edge_diff.edges_out:
                parent[edge.child] = tskit.NULL
            for edge in edge_diff.edges_in:
                parent[edge.child] = edge.parent
            assert_array_equal(parent, tree.parent_array)

    def test_elements_are_like_named_tuple(self, simple_degree2_ts_fixture):
        for val in simple_degree2_ts_fixture.edge_diffs():
            assert len(val) == 3
            assert val[0] == val.interval
            assert val[1] == val.edges_out
            assert val[2] == val.edges_in

    @pytest.mark.parametrize("direction", [-6, "forward", None])
    def test_bad_direction(self, direction, simple_degree2_ts_fixture):
        ts = simple_degree2_ts_fixture
        with pytest.raises(ValueError, match="direction must be"):
            ts.edge_diffs(direction=direction)

    @pytest.mark.parametrize("direction", [tskit.FORWARD, tskit.REVERSE])
    def test_edge_properties(self, direction, simple_degree2_ts_fixture):
        ts = simple_degree2_ts_fixture
        edge_ids = set()
        for _, e_out, e_in in ts.edge_diffs(direction=direction):
            for edge in e_in:
                assert edge.id not in edge_ids
                edge_ids.add(edge.id)
                assert ts.edge(edge.id) == edge
            for edge in e_out:
                assert ts.edge(edge.id) == edge
        assert edge_ids == set(range(ts.num_edges))

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    @pytest.mark.parametrize("direction", [tskit.FORWARD, tskit.REVERSE])
    def test_include_terminal(self, ts, direction):
        edges = set()
        i = 0
        diffs = ts.edge_diffs(include_terminal=True, direction=direction)
        parent = np.full(ts.num_nodes + 1, tskit.NULL, dtype=np.int32)
        for (left, right), e_out, e_in in diffs:  # noqa: B007
            for e in e_out:
                edges.remove(e.id)
                parent[e.child] = tskit.NULL
            for e in e_in:
                edges.add(e.id)
                parent[e.child] = e.parent
            i += 1
        assert np.all(parent == tskit.NULL)
        assert i == ts.num_trees + 1
        assert len(edges) == 0
        # On last iteration, interval is empty
        if direction == tskit.FORWARD:
            assert left == ts.sequence_length
            assert right == ts.sequence_length
        else:
            assert left == 0
            assert right == 0


class TestTreeSequenceMethodSignatures:
    ts = msprime.simulate(10, random_seed=1234)

    def test_kwargs_only(self):
        with pytest.raises(TypeError, match="argument"):
            tskit.Tree(self.ts, [], True)
        with pytest.raises(TypeError, match="argument"):
            self.ts.trees([], True)
        with pytest.raises(TypeError, match="argument"):
            self.ts.haplotypes(True)
        with pytest.raises(TypeError, match="argument"):
            self.ts.variants(True)
        with pytest.raises(TypeError, match="argument"):
            self.ts.genotype_matrix(True)
        with pytest.raises(TypeError, match="argument"):
            self.ts.simplify([], True)
        with pytest.raises(TypeError, match="argument"):
            self.ts.draw_svg("filename", True)
        with pytest.raises(TypeError, match="argument"):
            tskit.TreeSequence.load_tables(tskit.TableCollection(1), True)

    def test_trees_params(self):
        """
        The initial .trees() iterator parameters should match those in Tree.__init__()
        """
        tree_class_params = list(inspect.signature(tskit.Tree).parameters.items())
        trees_iter_params = list(
            inspect.signature(tskit.TreeSequence.trees).parameters.items()
        )
        # Skip the first param, which is `tree_sequence` and `self` respectively
        tree_class_params = tree_class_params[1:]
        # The trees iterator has some extra (deprecated) aliases
        trees_iter_params = trees_iter_params[1:-3]
        assert trees_iter_params == tree_class_params


class TestTreeSequenceMetadata:
    metadata_tables = [
        "node",
        "edge",
        "site",
        "mutation",
        "migration",
        "individual",
        "population",
    ]
    metadata_schema = tskit.MetadataSchema(
        {
            "codec": "json",
            "title": "Example Metadata",
            "type": "object",
            "properties": {
                "table": {"type": "string"},
                "string_prop": {"type": "string"},
                "num_prop": {"type": "number"},
            },
            "required": ["table", "string_prop", "num_prop"],
            "additionalProperties": False,
        },
    )

    def test_tree_sequence_metadata_schema(self):
        tc = tskit.TableCollection(1)
        ts = tc.tree_sequence()
        assert repr(ts.metadata_schema) == repr(tskit.MetadataSchema(None))
        tc.metadata_schema = self.metadata_schema
        ts = tc.tree_sequence()
        assert repr(ts.metadata_schema) == repr(self.metadata_schema)
        with pytest.raises(AttributeError):
            del ts.metadata_schema
        with pytest.raises(AttributeError):
            ts.metadata_schema = tskit.MetadataSchema(None)

    def test_tree_sequence_metadata(self):
        tc = tskit.TableCollection(1)
        ts = tc.tree_sequence()
        assert ts.metadata == b""
        tc.metadata_schema = self.metadata_schema
        data = {
            "table": "tree-sequence",
            "string_prop": "stringy",
            "num_prop": 42,
        }
        tc.metadata = data
        ts = tc.tree_sequence()
        assert ts.metadata == data
        with pytest.raises(AttributeError):
            ts.metadata = {"should": "fail"}
        with pytest.raises(AttributeError):
            del ts.metadata

    def test_tree_sequence_time_units(self):
        tc = tskit.TableCollection(1)
        ts = tc.tree_sequence()
        assert ts.time_units == tskit.TIME_UNITS_UNKNOWN
        tc.time_units = "something else"
        ts = tc.tree_sequence()
        assert ts.time_units == "something else"
        with pytest.raises(AttributeError):
            del ts.time_units
        with pytest.raises(AttributeError):
            ts.time_units = "readonly"
        assert tskit.TIME_UNITS_UNKNOWN == "unknown"
        assert tskit.TIME_UNITS_UNCALIBRATED == "uncalibrated"

    def test_table_metadata_schemas(self):
        ts = msprime.simulate(5)
        for table in self.metadata_tables:
            tables = ts.dump_tables()
            # Set and read back a unique schema for each table
            schema = tskit.MetadataSchema({"codec": "json", "TEST": f"{table}-SCHEMA"})
            # Check via table API
            getattr(tables, f"{table}s").metadata_schema = schema
            assert repr(getattr(tables, f"{table}s").metadata_schema) == repr(schema)
            for other_table in self.metadata_tables:
                if other_table != table:
                    assert (
                        repr(getattr(tables, f"{other_table}s").metadata_schema) == ""
                    )
            # Check via tree-sequence API
            new_ts = tskit.TreeSequence.load_tables(tables)
            assert repr(getattr(new_ts.table_metadata_schemas, table)) == repr(schema)
            for other_table in self.metadata_tables:
                if other_table != table:
                    assert (
                        repr(getattr(new_ts.table_metadata_schemas, other_table)) == ""
                    )
            # Can't set schema via this API
            with pytest.raises(AttributeError):
                new_ts.table_metadata_schemas = {}
                # or modify the schema tuple return object
                with pytest.raises(dataclasses.exceptions.FrozenInstanceError):
                    setattr(
                        new_ts.table_metadata_schemas,
                        table,
                        tskit.MetadataSchema({"codec": "json"}),
                    )

    def test_table_metadata_round_trip_via_row_getters(self):
        # A tree sequence with all entities
        pop_configs = [msprime.PopulationConfiguration(5) for _ in range(2)]
        migration_matrix = [[0, 1], [1, 0]]
        ts = msprime.simulate(
            population_configurations=pop_configs,
            migration_matrix=migration_matrix,
            mutation_rate=1,
            record_migrations=True,
            random_seed=1,
        )
        tables = ts.dump_tables()
        tables.individuals.add_row(location=[1, 2, 3])
        tables.individuals.add_row(location=[4, 5, 6])
        ts = tables.tree_sequence()

        for table in self.metadata_tables:
            new_tables = ts.dump_tables()
            tables_copy = ts.dump_tables()
            table_obj = getattr(new_tables, f"{table}s")
            table_obj.metadata_schema = self.metadata_schema
            table_obj.clear()
            # Write back the rows, but adding unique metadata
            for j, row in enumerate(getattr(tables_copy, f"{table}s")):
                row_data = dataclasses.asdict(row)
                row_data["metadata"] = {
                    "table": table,
                    "string_prop": f"Row number{j}",
                    "num_prop": j,
                }
                table_obj.add_row(**row_data)
            new_ts = new_tables.tree_sequence()
            # Check that all tables have data otherwise we'll silently not check one
            assert getattr(new_ts, f"num_{table}s") > 0
            assert getattr(new_ts, f"num_{table}s") == getattr(ts, f"num_{table}s")
            for j, row in enumerate(getattr(new_ts, f"{table}s")()):
                assert row.metadata == {
                    "table": table,
                    "string_prop": f"Row number{row.id}",
                    "num_prop": row.id,
                }
                assert getattr(new_ts, f"{table}")(j).metadata == {
                    "table": table,
                    "string_prop": f"Row number{row.id}",
                    "num_prop": row.id,
                }


def test_pickle_round_trip(ts_fixture):
    for protocol in range(pickle.HIGHEST_PROTOCOL + 1):
        ts = pickle.loads(pickle.dumps(ts_fixture, protocol=protocol))
        assert ts.tables == ts_fixture.tables
        # Do some thing to check the ts is init'd properly
        ts.draw_text()


class TestFileUuid(HighLevelTestCase):
    """
    Tests that the file UUID attribute is handled correctly.
    """

    def validate(self, ts):
        with tempfile.TemporaryDirectory() as tempdir:
            temp_file = pathlib.Path(tempdir) / "tmp.trees"
            assert ts.file_uuid is None
            ts.dump(temp_file)
            other_ts = tskit.load(temp_file)
            assert other_ts.file_uuid is not None
            assert len(other_ts.file_uuid), 36
            uuid = other_ts.file_uuid
            other_ts = tskit.load(temp_file)
            assert other_ts.file_uuid == uuid
            assert ts.tables == other_ts.tables

            # Check that the UUID is well-formed.
            parsed = _uuid.UUID("{" + uuid + "}")
            assert str(parsed) == uuid

            # Save the same tree sequence to the file. We should get a different UUID.
            ts.dump(temp_file)
            other_ts = tskit.load(temp_file)
            assert other_ts.file_uuid is not None
            assert other_ts.file_uuid != uuid

            # Even saving a ts that has a UUID to another file changes the UUID
            old_uuid = other_ts.file_uuid
            other_ts.dump(temp_file)
            assert other_ts.file_uuid == old_uuid
            other_ts = tskit.load(temp_file)
            assert other_ts.file_uuid is not None
            assert other_ts.file_uuid != old_uuid

            # Tables dumped from this ts are a deep copy, so they don't have
            # the file_uuid.
            tables = other_ts.dump_tables()
            assert tables.file_uuid is None

            # For now, ts.tables also returns a deep copy. This will hopefully
            # change in the future though.
            assert ts.tables.file_uuid is None

    def test_simple_simulation(self):
        ts = msprime.simulate(2, random_seed=1)
        self.validate(ts)

    def test_empty_tables(self):
        tables = tskit.TableCollection(1)
        self.validate(tables.tree_sequence())


class TestTreeSequenceTextIO(HighLevelTestCase):
    """
    Tests for the tree sequence text IO.
    """

    def verify_nodes_format(self, ts, nodes_file, precision, base64_metadata):
        """
        Verifies that the nodes we output have the correct form.
        """

        def convert(v):
            return "{:.{}f}".format(v, precision)

        output_nodes = nodes_file.read().splitlines()
        assert len(output_nodes) - 1 == ts.num_nodes
        assert list(output_nodes[0].split()) == [
            "id",
            "is_sample",
            "time",
            "population",
            "individual",
            "metadata",
        ]
        for node, line in zip(ts.nodes(), output_nodes[1:]):
            splits = line.split("\t")
            assert str(node.id) == splits[0]
            assert str(node.is_sample()) == splits[1]
            assert convert(node.time) == splits[2]
            assert str(node.population) == splits[3]
            assert str(node.individual) == splits[4]
            if isinstance(node.metadata, bytes) and base64_metadata:
                assert tests.base64_encode(node.metadata) == splits[5]
            else:
                assert repr(node.metadata) == splits[5]

    def verify_edges_format(self, ts, edges_file, precision, base64_metadata):
        """
        Verifies that the edges we output have the correct form.
        """

        def convert(v):
            return "{:.{}f}".format(v, precision)

        output_edges = edges_file.read().splitlines()
        assert len(output_edges) - 1 == ts.num_edges
        assert list(output_edges[0].split()) == [
            "left",
            "right",
            "parent",
            "child",
            "metadata",
        ]
        for edge, line in zip(ts.edges(), output_edges[1:]):
            splits = line.split("\t")
            assert convert(edge.left) == splits[0]
            assert convert(edge.right) == splits[1]
            assert str(edge.parent) == splits[2]
            assert str(edge.child) == splits[3]
            if isinstance(edge.metadata, bytes) and base64_metadata:
                assert tests.base64_encode(edge.metadata) == splits[4]
            else:
                assert repr(edge.metadata) == splits[4]

    def verify_sites_format(self, ts, sites_file, precision, base64_metadata):
        """
        Verifies that the sites we output have the correct form.
        """

        def convert(v):
            return "{:.{}f}".format(v, precision)

        output_sites = sites_file.read().splitlines()
        assert len(output_sites) - 1 == ts.num_sites
        assert list(output_sites[0].split()) == [
            "position",
            "ancestral_state",
            "metadata",
        ]
        for site, line in zip(ts.sites(), output_sites[1:]):
            splits = line.split("\t")
            assert convert(site.position) == splits[0]
            assert site.ancestral_state == splits[1]
            if isinstance(site.metadata, bytes) and base64_metadata:
                assert tests.base64_encode(site.metadata) == splits[2]
            else:
                assert repr(site.metadata) == splits[2]

    def verify_mutations_format(self, ts, mutations_file, precision, base64_metadata):
        """
        Verifies that the mutations we output have the correct form.
        """

        def convert(v):
            return "{:.{}f}".format(v, precision)

        output_mutations = mutations_file.read().splitlines()
        assert len(output_mutations) - 1 == ts.num_mutations
        assert list(output_mutations[0].split()) == [
            "site",
            "node",
            "time",
            "derived_state",
            "parent",
            "metadata",
        ]
        mutations = [mut for site in ts.sites() for mut in site.mutations]
        for mutation, line in zip(mutations, output_mutations[1:]):
            splits = line.split("\t")
            assert str(mutation.site) == splits[0]
            assert str(mutation.node) == splits[1]
            assert (
                "unknown" if util.is_unknown_time(mutation.time) else str(mutation.time)
            ) == splits[2]
            assert str(mutation.derived_state) == splits[3]
            assert str(mutation.parent) == splits[4]
            if isinstance(mutation.metadata, bytes) and base64_metadata:
                assert tests.base64_encode(mutation.metadata) == splits[5]
            else:
                assert repr(mutation.metadata) == splits[5]

    def verify_individuals_format(
        self, ts, individuals_file, precision, base64_metadata
    ):
        """
        Verifies that the individuals we output have the correct form.
        """

        def convert(v):
            return "{:.{}f}".format(v, precision)

        output_individuals = individuals_file.read().splitlines()
        assert len(output_individuals) - 1 == ts.num_individuals
        assert list(output_individuals[0].split()) == [
            "id",
            "flags",
            "location",
            "parents",
            "metadata",
        ]
        for individual, line in zip(ts.individuals(), output_individuals[1:]):
            splits = line.split("\t")
            assert str(individual.id) == splits[0]
            assert str(individual.flags) == splits[1]
            assert ",".join(map(str, individual.location)) == splits[2]
            assert ",".join(map(str, individual.parents)) == splits[3]
            if isinstance(individual.metadata, bytes) and base64_metadata:
                assert tests.base64_encode(individual.metadata) == splits[4]
            else:
                assert repr(individual.metadata) == splits[4]

    def verify_populations_format(
        self, ts, populations_file, precision, base64_metadata
    ):
        """
        Verifies that the populations we output have the correct form.
        """

        def convert(v):
            return "{:.{}f}".format(v, precision)

        output_populations = populations_file.read().splitlines()
        assert len(output_populations) - 1 == ts.num_populations
        assert list(output_populations[0].split()) == [
            "id",
            "metadata",
        ]
        for population, line in zip(ts.populations(), output_populations[1:]):
            splits = line.split("\t")
            assert str(population.id) == splits[0]
            if isinstance(population.metadata, bytes) and base64_metadata:
                assert tests.base64_encode(population.metadata) == splits[1]
            else:
                assert repr(population.metadata) == splits[1]

    def verify_migrations_format(self, ts, migrations_file, precision, base64_metadata):
        """
        Verifies that the migrations we output have the correct form.
        """

        def convert(v):
            return "{:.{}f}".format(v, precision)

        output_migrations = migrations_file.read().splitlines()
        assert len(output_migrations) - 1 == ts.num_migrations
        assert list(output_migrations[0].split()) == [
            "left",
            "right",
            "node",
            "source",
            "dest",
            "time",
            "metadata",
        ]
        for migration, line in zip(ts.migrations(), output_migrations[1:]):
            splits = line.split("\t")
            assert str(migration.left) == splits[0]
            assert str(migration.right) == splits[1]
            assert str(migration.node) == splits[2]
            assert str(migration.source) == splits[3]
            assert str(migration.dest) == splits[4]
            assert str(migration.time) == splits[5]
            if isinstance(migration.metadata, bytes) and base64_metadata:
                assert tests.base64_encode(migration.metadata) == splits[6]
            else:
                assert repr(migration.metadata) == splits[6]

    @pytest.mark.parametrize(("precision", "base64_metadata"), [(2, True), (7, False)])
    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_output_format(self, precision, base64_metadata, ts):
        nodes_file = io.StringIO()
        edges_file = io.StringIO()
        sites_file = io.StringIO()
        mutations_file = io.StringIO()
        individuals_file = io.StringIO()
        populations_file = io.StringIO()
        migrations_file = io.StringIO()
        provenances_file = io.StringIO()
        ts.dump_text(
            nodes=nodes_file,
            edges=edges_file,
            sites=sites_file,
            mutations=mutations_file,
            individuals=individuals_file,
            populations=populations_file,
            migrations=migrations_file,
            provenances=provenances_file,
            precision=precision,
            base64_metadata=base64_metadata,
        )
        nodes_file.seek(0)
        edges_file.seek(0)
        sites_file.seek(0)
        mutations_file.seek(0)
        individuals_file.seek(0)
        populations_file.seek(0)
        migrations_file.seek(0)
        self.verify_nodes_format(ts, nodes_file, precision, base64_metadata)
        self.verify_edges_format(ts, edges_file, precision, base64_metadata)
        self.verify_sites_format(ts, sites_file, precision, base64_metadata)
        self.verify_mutations_format(ts, mutations_file, precision, base64_metadata)
        self.verify_individuals_format(ts, individuals_file, precision, base64_metadata)
        self.verify_populations_format(ts, populations_file, precision, base64_metadata)
        self.verify_migrations_format(ts, migrations_file, precision, base64_metadata)

    def verify_approximate_equality(self, ts1, ts2):
        """
        Verifies that the specified tree sequences are approximately
        equal, taking into account the error incurred in exporting to text.
        """
        assert ts1.sample_size == ts2.sample_size
        assert ts1.sequence_length == ts2.sequence_length
        assert ts1.num_nodes == ts2.num_nodes
        assert ts1.num_edges == ts2.num_edges
        assert ts1.num_sites == ts2.num_sites
        assert ts1.num_mutations == ts2.num_mutations
        assert ts1.num_populations == ts2.num_populations
        assert ts1.num_migrations == ts2.num_migrations

        checked = 0
        for n1, n2 in zip(ts1.nodes(), ts2.nodes()):
            assert n1.population == n2.population
            assert n1.metadata == n2.metadata
            assert n1.time == pytest.approx(n2.time)
            checked += 1
        assert checked == ts1.num_nodes

        checked = 0
        for r1, r2 in zip(ts1.edges(), ts2.edges()):
            checked += 1
            assert r1.left == pytest.approx(r2.left)
            assert r1.right == pytest.approx(r2.right)
            assert r1.parent == r2.parent
            assert r1.child == r2.child
        assert ts1.num_edges == checked

        checked = 0
        for s1, s2 in zip(ts1.sites(), ts2.sites()):
            checked += 1
            assert s1.position == pytest.approx(s2.position)
            assert s1.ancestral_state == s2.ancestral_state
            assert s1.metadata == s2.metadata
            assert s1.mutations == s2.mutations
        assert ts1.num_sites == checked

        checked = 0
        for s1, s2 in zip(ts1.mutations(), ts2.mutations()):
            checked += 1
            assert s1.site == s2.site
            assert s1.node == s2.node
            if not (math.isnan(s1.time) and math.isnan(s2.time)):
                assert s1.time == pytest.approx(s2.time)
            assert s1.derived_state == s2.derived_state
            assert s1.parent == s2.parent
            assert s1.metadata == s2.metadata
        assert ts1.num_mutations == checked

        checked = 0
        for s1, s2 in zip(ts1.migrations(), ts2.migrations()):
            checked += 1
            assert s1.left == s2.left
            assert s1.right == s2.right
            assert s1.node == s2.node
            assert s1.source == s2.source
            assert s1.dest == s2.dest
            assert s1.time == s2.time
            assert s1.metadata == s2.metadata
        assert ts1.num_migrations == checked

        # Check the trees
        check = 0
        for t1, t2 in zip(ts1.trees(), ts2.trees()):
            assert list(t1.nodes()) == list(t2.nodes())
            check += 1
        assert check == ts1.get_num_trees()

    @pytest.mark.parametrize("ts1", get_example_tree_sequences())
    def test_text_record_round_trip(self, ts1):
        # Can't round trip without the schema
        if ts1.tables.nodes.metadata_schema == tskit.MetadataSchema(None):
            nodes_file = io.StringIO()
            edges_file = io.StringIO()
            sites_file = io.StringIO()
            mutations_file = io.StringIO()
            individuals_file = io.StringIO()
            populations_file = io.StringIO()
            migrations_file = io.StringIO()
            ts1.dump_text(
                nodes=nodes_file,
                edges=edges_file,
                sites=sites_file,
                mutations=mutations_file,
                individuals=individuals_file,
                populations=populations_file,
                migrations=migrations_file,
                precision=16,
            )
            nodes_file.seek(0)
            edges_file.seek(0)
            sites_file.seek(0)
            mutations_file.seek(0)
            individuals_file.seek(0)
            populations_file.seek(0)
            migrations_file.seek(0)
            ts2 = tskit.load_text(
                nodes=nodes_file,
                edges=edges_file,
                sites=sites_file,
                mutations=mutations_file,
                individuals=individuals_file,
                populations=populations_file,
                migrations=migrations_file,
                sequence_length=ts1.sequence_length,
                strict=True,
            )
            self.verify_approximate_equality(ts1, ts2)

    def test_empty_files(self):
        nodes_file = io.StringIO("is_sample\ttime\n")
        edges_file = io.StringIO("left\tright\tparent\tchild\n")
        sites_file = io.StringIO("position\tancestral_state\n")
        mutations_file = io.StringIO("site\tnode\tderived_state\n")
        individuals_file = io.StringIO("flags\n")
        migrations_file = io.StringIO("left\tright\tnode\tsource\tdest\ttime\n")
        with pytest.raises(_tskit.LibraryError):
            tskit.load_text(
                nodes=nodes_file,
                edges=edges_file,
                sites=sites_file,
                mutations=mutations_file,
                individuals=individuals_file,
                migrations=migrations_file,
            )

    def test_empty_files_sequence_length(self):
        nodes_file = io.StringIO("is_sample\ttime\n")
        edges_file = io.StringIO("left\tright\tparent\tchild\n")
        sites_file = io.StringIO("position\tancestral_state\n")
        mutations_file = io.StringIO("site\tnode\tderived_state\n")
        individuals_file = io.StringIO("flags\n")
        migrations_file = io.StringIO("left\tright\tnode\tsource\tdest\ttime\n")
        ts = tskit.load_text(
            nodes=nodes_file,
            edges=edges_file,
            sites=sites_file,
            mutations=mutations_file,
            individuals=individuals_file,
            migrations=migrations_file,
            sequence_length=100,
        )
        assert ts.sequence_length == 100
        assert ts.num_nodes == 0
        assert ts.num_edges == 0
        assert ts.num_sites == 0
        assert ts.num_mutations == 0
        assert ts.num_individuals == 0
        assert ts.num_migrations == 0

    def test_load_text_no_populations(self):
        nodes_file = io.StringIO("is_sample\ttime\tpopulation\n1\t0\t2\n")
        edges_file = io.StringIO("left\tright\tparent\tchild\n")
        ts = tskit.load_text(nodes_file, edges_file, sequence_length=100)
        assert ts.num_nodes == 1
        assert ts.num_populations == 3

    def test_load_text_populations(self):
        nodes_file = io.StringIO("is_sample\ttime\tpopulation\n")
        edges_file = io.StringIO("left\tright\tparent\tchild\n")
        populations_file = io.StringIO("metadata\nmetadata_1\nmetadata_2\n")
        ts = tskit.load_text(
            nodes_file,
            edges_file,
            populations=populations_file,
            sequence_length=100,
            base64_metadata=False,
        )
        assert ts.num_populations == 2
        assert ts.tables.populations[0].metadata == b"metadata_1"
        assert ts.tables.populations[1].metadata == b"metadata_2"


class TestTree(HighLevelTestCase):
    """
    Some simple tests on the tree API.
    """

    def get_tree(self, sample_lists=False):
        ts = msprime.simulate(10, random_seed=1, mutation_rate=1, record_full_arg=True)
        return next(ts.trees(sample_lists=sample_lists))

    def verify_mutations(self, tree):
        assert tree.num_mutations > 0
        other_mutations = []
        for site in tree.sites():
            for mutation in site.mutations:
                other_mutations.append(mutation)
        mutations = list(tree.mutations())
        assert tree.num_mutations == len(other_mutations)
        assert tree.num_mutations == len(mutations)
        for mut, other_mut in zip(mutations, other_mutations):
            assert mut == other_mut

    def test_simple_mutations(self):
        tree = self.get_tree()
        self.verify_mutations(tree)

    def test_complex_mutations(self):
        ts = tsutil.insert_branch_mutations(msprime.simulate(10, random_seed=1))
        self.verify_mutations(ts.first())

    def test_str(self, ts_fixture):
        t = ts_fixture.first()
        assert isinstance(str(t), str)
        pattern = re.compile(
            r"""
            ╔═+╗\s*
            ║Tree.*?║\s*
            ╠═+╤═+╣\s*
            ║Index.*?│\s*\d+║\s*
            ╟─+┼─+╢\s*
            ║Interval.*?│\s*\d+-\d+\s*\(\d+\)║\s*
            ╟─+┼─+╢\s*
            ║Roots.*?│\s*\d+║\s*
            ╟─+┼─+╢\s*
            ║Nodes.*?│\s*\d+║\s*
            ╟─+┼─+╢\s*
            ║Sites.*?│\s*\d+║\s*
            ╟─+┼─+╢\s*
            ║Mutations.*?│\s*\d+║\s*
            ╟─+┼─+╢\s*
            ║Total\s*Branch\s*Length.*?│\s*[\d,]+\.\d+║\s*
            ╚═+╧═+╝\s*
            """,
            re.VERBOSE | re.DOTALL,
        )
        assert pattern.search(str(t))

    def test_html_repr(self, ts_fixture):
        html = ts_fixture.first()._repr_html_()
        # Parse to check valid
        ElementTree.fromstring(html)
        assert len(html) > 1900
        assert "<tr><td>Total Branch Length</td><td>" in html

    def test_samples(self):
        for sample_lists in [True, False]:
            t = self.get_tree(sample_lists)
            n = t.get_sample_size()
            all_samples = list(t.samples(t.get_root()))
            assert sorted(all_samples) == list(range(n))
            for j in range(n):
                assert list(t.samples(j)) == [j]

            def test_func(t, u):
                """
                Simple test definition of the traversal.
                """
                stack = [u]
                while len(stack) > 0:
                    v = stack.pop()
                    if t.is_sample(v):
                        yield v
                    if t.is_internal(v):
                        for c in reversed(t.get_children(v)):
                            stack.append(c)

            for u in t.nodes():
                l1 = list(t.samples(u))
                l2 = list(test_func(t, u))
                assert l1 == l2
                assert t.get_num_samples(u) == len(l1)

    def test_num_children(self):
        tree = self.get_tree()
        for u in tree.nodes():
            assert tree.num_children(u) == len(tree.children(u))

    def test_ancestors(self):
        tree = tskit.Tree.generate_balanced(10, arity=3)
        ancestors_arrays = {u: [] for u in np.arange(tree.tree_sequence.num_nodes)}
        ancestors_arrays[-1] = []
        for u in tree.nodes(order="preorder"):
            parent = tree.parent(u)
            if parent != tskit.NULL:
                ancestors_arrays[u] = [parent] + ancestors_arrays[tree.parent(u)]
        for u in tree.nodes():
            assert list(tree.ancestors(u)) == ancestors_arrays[u]

    def test_ancestors_empty(self):
        ts = tskit.Tree.generate_comb(10).tree_sequence
        tree = ts.delete_intervals([[0, 1]]).first()
        for u in ts.samples():
            assert len(list(tree.ancestors(u))) == 0

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_virtual_root_semantics(self, ts):
        for tree in ts.trees():
            assert math.isinf(tree.time(tree.virtual_root))
            assert tree.depth(tree.virtual_root) == -1
            assert tree.parent(tree.virtual_root) == -1
            assert list(tree.children(tree.virtual_root)) == tree.roots
            with pytest.raises(tskit.LibraryError, match="bounds"):
                tree.population(tree.virtual_root)

    def test_root_properties(self):
        tested = set()
        for ts in get_example_tree_sequences(pytest_params=False):
            for tree in ts.trees():
                if tree.has_single_root:
                    tested.add("single")
                    assert tree.num_roots == 1
                    assert tree.num_roots == 1
                    assert tree.root != tskit.NULL
                elif tree.has_multiple_roots:
                    tested.add("multiple")
                    assert tree.num_roots > 1
                    with pytest.raises(ValueError, match="More than one root exists"):
                        _ = tree.root
                else:
                    tested.add("zero")
                    assert tree.num_roots == 0
                    assert tree.root == tskit.NULL
        assert len(tested) == 3

    def test_as_dict_of_dicts(self):
        for ts in get_example_tree_sequences(pytest_params=False):
            tree = next(ts.trees())
            adj_dod = tree.as_dict_of_dicts()
            g = nx.DiGraph(adj_dod)

            self.verify_nx_graph_topology(tree, g)
            self.verify_nx_algorithm_equivalence(tree, g)
            self.verify_nx_for_tutorial_algorithms(tree, g)
        self.verify_nx_nearest_neighbor_search()

    def verify_nx_graph_topology(self, tree, g):
        assert set(tree.nodes()) == set(g.nodes)

        assert set(tree.roots) == {n for n in g.nodes if g.in_degree(n) == 0}

        assert set(tree.leaves()) == {n for n in g.nodes if g.out_degree(n) == 0}

        # test if tree has no in-degrees > 1
        if len(g) > 0:
            assert nx.is_branching(g)

    def verify_nx_algorithm_equivalence(self, tree, g):
        for root in tree.roots:
            assert nx.is_directed_acyclic_graph(g)

            # test descendants
            assert {u for u in tree.nodes() if tree.is_descendant(u, root)} == set(
                nx.descendants(g, root)
            ) | {root}

            # test MRCA
            if tree.tree_sequence.num_nodes < 20:
                for u, v in itertools.combinations(tree.nodes(), 2):
                    mrca = nx.lowest_common_ancestor(g, u, v)
                    if mrca is None:
                        mrca = -1
                    assert tree.mrca(u, v) == mrca

            # test node traversal modes
            assert list(tree.nodes(root=root, order="breadthfirst")) == [root] + [
                v for u, v in nx.bfs_edges(g, root)
            ]
            assert list(tree.nodes(root=root, order="preorder")) == list(
                nx.dfs_preorder_nodes(g, root)
            )

    def verify_nx_for_tutorial_algorithms(self, tree, g):
        # traversing upwards
        for u in tree.leaves():
            path = []
            v = u
            while v != tskit.NULL:
                path.append(v)
                v = tree.parent(v)

            assert set(path) == {u} | nx.ancestors(g, u)
            assert path == [u] + [
                n1 for n1, n2, _ in nx.edge_dfs(g, u, orientation="reverse")
            ]

        # traversals with information
        def preorder_dist(tree, root):
            stack = [(root, 0)]
            while len(stack) > 0:
                u, distance = stack.pop()
                yield u, distance
                for v in tree.children(u):
                    stack.append((v, distance + 1))

        for root in tree.roots:
            assert {
                k: v for k, v in preorder_dist(tree, root)
            } == nx.shortest_path_length(g, source=root)

        for root in tree.roots:
            # new traversal: measuring time between root and MRCA
            for u, v in itertools.combinations(nx.descendants(g, root), 2):
                mrca = tree.mrca(u, v)
                tmrca = tree.time(mrca)
                assert tree.time(root) - tmrca == pytest.approx(
                    nx.shortest_path_length(
                        g, source=root, target=mrca, weight="branch_length"
                    )
                )

    def verify_nx_nearest_neighbor_search(self):
        samples = [
            msprime.Sample(0, 0),
            msprime.Sample(0, 1),
            msprime.Sample(0, 20),
        ]
        ts = msprime.simulate(
            Ne=1e6,
            samples=samples,
            demographic_events=[
                msprime.PopulationParametersChange(
                    time=10, growth_rate=2, population_id=0
                ),
            ],
            random_seed=42,
        )

        tree = ts.first()
        g = nx.Graph(tree.as_dict_of_dicts())

        dist_dod = collections.defaultdict(dict)
        for source, target in itertools.combinations(tree.samples(), 2):
            dist_dod[source][target] = nx.shortest_path_length(
                g, source=source, target=target, weight="branch_length"
            )
            dist_dod[target][source] = dist_dod[source][target]

        nearest_neighbor_of = [min(dist_dod[u], key=dist_dod[u].get) for u in range(3)]
        assert [2, 2, 1] == [nearest_neighbor_of[u] for u in range(3)]

    def test_total_branch_length(self):
        # Note: this definition works when we have no non-sample branches.
        t1 = self.get_tree()
        bl = 0
        root = t1.get_root()
        for node in t1.nodes():
            if node != root:
                bl += t1.get_branch_length(node)
        assert bl > 0
        assert t1.get_total_branch_length() == pytest.approx(bl)

    def test_branch_length_empty_tree(self):
        tables = tskit.TableCollection(1)
        tables.nodes.add_row(flags=1, time=0)
        tables.nodes.add_row(flags=1, time=0)
        ts = tables.tree_sequence()
        assert ts.num_trees == 1
        tree = ts.first()
        assert tree.branch_length(0) == 0
        assert tree.branch_length(1) == 0
        assert tree.total_branch_length == 0

    @pytest.mark.parametrize("r_threshold", [0, -1])
    def test_bad_val_root_threshold(self, r_threshold):
        with pytest.raises(ValueError, match="greater than 0"):
            tskit.Tree.generate_balanced(2, root_threshold=r_threshold)

    @pytest.mark.parametrize("r_threshold", [None, 0.5, 1.5, np.inf])
    def test_bad_type_root_threshold(self, r_threshold):
        with pytest.raises(TypeError):
            tskit.Tree.generate_balanced(2, root_threshold=r_threshold)

    def test_simple_root_threshold(self):
        tree = tskit.Tree.generate_balanced(3, root_threshold=3)
        assert tree.num_roots == 1
        tree = tskit.Tree.generate_balanced(3, root_threshold=4)
        assert tree.num_roots == 0

    @pytest.mark.parametrize("root_threshold", [1, 2, 3])
    def test_is_root(self, root_threshold):
        # Make a tree with multiple roots with different numbers of samples under each
        ts = tskit.Tree.generate_balanced(5).tree_sequence
        ts = ts.decapitate(ts.max_root_time - 0.1)
        tables = ts.dump_tables()
        tables.nodes.add_row(flags=0)  # Isolated non-sample
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE)  # Isolated sample
        ts = tables.tree_sequence()
        assert {ts.first().num_samples(u) for u in ts.first().roots} == {1, 2, 3}
        tree = ts.first(root_threshold=root_threshold)
        roots = set(tree.roots)
        for u in range(ts.num_nodes):  # Will also test isolated nodes
            assert tree.is_root(u) == (u in roots)

    def test_is_descendant(self):
        def is_descendant(tree, u, v):
            path = []
            while u != tskit.NULL:
                path.append(u)
                u = tree.parent(u)
            return v in path

        tree = self.get_tree()
        for u, v in itertools.product(range(tree.tree_sequence.num_nodes), repeat=2):
            assert is_descendant(tree, u, v) == tree.is_descendant(u, v)
        # All nodes are descendents of themselves
        for u in range(tree.tree_sequence.num_nodes + 1):
            assert tree.is_descendant(u, u)
        for bad_node in [-1, -2, tree.tree_sequence.num_nodes + 1]:
            with pytest.raises(ValueError):
                tree.is_descendant(0, bad_node)
            with pytest.raises(ValueError):
                tree.is_descendant(bad_node, 0)
            with pytest.raises(ValueError):
                tree.is_descendant(bad_node, bad_node)

    def test_apis(self):
        # tree properties
        t1 = self.get_tree()
        assert t1.get_root() == t1.root
        assert t1.get_index() == t1.index
        assert t1.get_interval() == t1.interval
        assert t1.get_sample_size() == t1.sample_size
        assert t1.get_num_mutations() == t1.num_mutations
        assert t1.get_parent_dict() == t1.parent_dict
        assert t1.get_total_branch_length() == t1.total_branch_length
        assert t1.span == t1.interval.right - t1.interval.left
        assert t1.mid == t1.interval.left + (t1.interval.right - t1.interval.left) / 2
        # node properties
        root = t1.get_root()
        for node in t1.nodes():
            if node != root:
                assert t1.get_time(node) == t1.time(node)
                assert t1.get_parent(node) == t1.parent(node)
                assert t1.get_children(node) == t1.children(node)
                assert t1.get_population(node) == t1.population(node)
                assert t1.get_num_samples(node) == t1.num_samples(node)
                assert t1.get_branch_length(node) == t1.branch_length(node)
                assert t1.get_num_tracked_samples(node) == t1.num_tracked_samples(node)

        pairs = itertools.islice(itertools.combinations(t1.nodes(), 2), 50)
        for pair in pairs:
            assert t1.get_mrca(*pair) == t1.mrca(*pair)
            assert t1.get_tmrca(*pair) == t1.tmrca(*pair)

    @pytest.mark.filterwarnings("ignore::FutureWarning")
    def test_deprecated_apis(self):
        t1 = self.get_tree()
        assert t1.get_length() == t1.span
        assert t1.length == t1.span
        assert t1.num_nodes == t1.tree_sequence.num_nodes

    def test_deprecated_api_warnings(self):
        # Deprecated and will be removed
        t1 = self.get_tree()
        with pytest.warns(FutureWarning, match="Tree.tree_sequence.num_nodes"):
            t1.num_nodes

    def test_seek_index(self):
        ts = msprime.simulate(10, recombination_rate=3, length=5, random_seed=42)
        N = ts.num_trees
        assert ts.num_trees > 3
        tree = tskit.Tree(ts)
        for index in [0, N // 2, N - 1, 1]:
            fresh_tree = tskit.Tree(ts)
            assert fresh_tree.index == -1
            fresh_tree.seek_index(index)
            tree.seek_index(index)
            assert fresh_tree.index == index
            assert tree.index == index

        tree = tskit.Tree(ts)
        for index in [-1, -2, -N + 2, -N + 1, -N]:
            fresh_tree = tskit.Tree(ts)
            assert fresh_tree.index == -1
            fresh_tree.seek_index(index)
            tree.seek_index(index)
            assert fresh_tree.index == index + N
            assert tree.index == index + N
        with pytest.raises(IndexError):
            tree.seek_index(N)
        with pytest.raises(IndexError):
            tree.seek_index(N + 1)
        with pytest.raises(IndexError):
            tree.seek_index(-N - 1)
        with pytest.raises(IndexError):
            tree.seek_index(-N - 2)

    def test_first_last(self):
        ts = msprime.simulate(10, recombination_rate=3, length=2, random_seed=42)
        assert ts.num_trees > 3
        tree = tskit.Tree(ts)
        tree.first()
        assert tree.index == 0
        tree = tskit.Tree(ts)
        tree.last()
        assert tree.index == ts.num_trees - 1
        tree = tskit.Tree(ts)
        for _ in range(3):
            tree.last()
            assert tree.index == ts.num_trees - 1
            tree.first()
            assert tree.index == 0

    def test_eq_different_tree_sequence(self):
        ts = msprime.simulate(4, recombination_rate=1, length=2, random_seed=42)
        copy = ts.tables.tree_sequence()
        for tree1, tree2 in zip(ts.aslist(), copy.aslist()):
            assert tree1 != tree2

    def test_next_prev(self):
        ts = msprime.simulate(10, recombination_rate=3, length=3, random_seed=42)
        assert ts.num_trees > 5
        for index, tree in enumerate(ts.aslist()):
            assert tree.index == index
            j = index
            while tree.next():
                j += 1
                assert tree.index == j
            assert tree.index == -1
            assert j + 1 == ts.num_trees
        for index, tree in enumerate(ts.aslist()):
            assert tree.index == index
            j = index
            while tree.prev():
                j -= 1
                assert tree.index == j
            assert tree.index == -1
            assert j == 0
        tree.first()
        tree.prev()
        assert tree.index == -1
        tree.last()
        tree.next()
        assert tree.index == -1

    def test_interval(self):
        ts = msprime.simulate(10, recombination_rate=1, random_seed=1)
        assert ts.num_trees > 1
        breakpoints = list(ts.breakpoints())
        assert breakpoints[0] == 0
        assert breakpoints[-1] == ts.sequence_length
        for i, tree in enumerate(ts.trees()):
            assert tree.interval.left == pytest.approx(breakpoints[i])
            assert tree.interval.left == pytest.approx(breakpoints[i])
            assert tree.interval.right == pytest.approx(breakpoints[i + 1])
            assert tree.interval.right == pytest.approx(breakpoints[i + 1])
            assert tree.interval.span == pytest.approx(
                breakpoints[i + 1] - breakpoints[i]
            )
            assert tree.interval.mid == pytest.approx(
                breakpoints[i] + (breakpoints[i + 1] - breakpoints[i]) / 2
            )

    def verify_tree_arrays(self, tree):
        ts = tree.tree_sequence
        N = ts.num_nodes + 1
        assert tree.parent_array.shape == (N,)
        assert tree.left_child_array.shape == (N,)
        assert tree.right_child_array.shape == (N,)
        assert tree.left_sib_array.shape == (N,)
        assert tree.right_sib_array.shape == (N,)
        assert tree.num_children_array.shape == (N,)
        assert tree.edge_array.shape == (N,)
        for u in range(N):
            assert tree.parent(u) == tree.parent_array[u]
            assert tree.left_child(u) == tree.left_child_array[u]
            assert tree.right_child(u) == tree.right_child_array[u]
            assert tree.left_sib(u) == tree.left_sib_array[u]
            assert tree.right_sib(u) == tree.right_sib_array[u]
            assert tree.num_children(u) == tree.num_children_array[u]
            assert tree.edge(u) == tree.edge_array[u]

    def verify_tree_arrays_python_ts(self, ts):
        pts = tests.PythonTreeSequence(ts)
        iter1 = ts.trees()
        iter2 = pts.trees()
        for st1, st2 in zip(iter1, iter2):
            assert np.all(st1.parent_array == st2.parent)
            assert np.all(st1.left_child_array == st2.left_child)
            assert np.all(st1.right_child_array == st2.right_child)
            assert np.all(st1.left_sib_array == st2.left_sib)
            assert np.all(st1.right_sib_array == st2.right_sib)
            assert np.all(st1.num_children_array == st2.num_children)
            assert np.all(st1.edge_array == st2.edge)

    def test_tree_arrays(self):
        ts = msprime.simulate(10, recombination_rate=1, random_seed=1)
        assert ts.num_trees > 1
        self.verify_tree_arrays_python_ts(ts)
        for tree in ts.trees():
            self.verify_tree_arrays(tree)

    @pytest.mark.parametrize(
        "array",
        [
            "parent",
            "left_child",
            "right_child",
            "left_sib",
            "right_sib",
            "num_children",
            "edge",
        ],
    )
    def test_tree_array_properties(self, array):
        name = array + "_array"
        ts = msprime.simulate(10, random_seed=1)
        tree = ts.first()
        a = getattr(tree, name)
        assert getattr(tree, name) is a
        assert a.base is tree._ll_tree
        with pytest.raises(AttributeError):
            setattr(tree, name, None)
        with pytest.raises(AttributeError):
            delattr(tree, name)

    def verify_empty_tree(self, tree):
        ts = tree.tree_sequence
        assert tree.index == -1
        assert tree.parent_dict == {}
        for u in range(ts.num_nodes):
            assert tree.parent(u) == tskit.NULL
            assert tree.left_child(u) == tskit.NULL
            assert tree.right_child(u) == tskit.NULL
            assert tree.num_children(u) == 0
            assert tree.edge(u) == tskit.NULL
            if not ts.node(u).is_sample():
                assert tree.left_sib(u) == tskit.NULL
                assert tree.right_sib(u) == tskit.NULL
        # Samples should have left-sib right-sibs set
        samples = ts.samples()
        assert tree.left_root == samples[0]
        for j in range(ts.num_samples):
            if j > 0:
                assert tree.left_sib(samples[j]) == samples[j - 1]
            if j < ts.num_samples - 1:
                assert tree.right_sib(samples[j]) == samples[j + 1]
        self.verify_tree_arrays(tree)

    def test_empty_tree(self):
        ts = msprime.simulate(10, recombination_rate=3, length=3, random_seed=42)
        assert ts.num_trees > 5
        tree = tskit.Tree(ts)
        self.verify_empty_tree(tree)
        while tree.next():
            pass
        self.verify_empty_tree(tree)
        while tree.prev():
            pass
        self.verify_empty_tree(tree)

    def test_clear(self):
        ts = msprime.simulate(10, recombination_rate=3, length=3, random_seed=42)
        assert ts.num_trees > 5
        tree = tskit.Tree(ts)
        tree.first()
        tree.clear()
        self.verify_empty_tree(tree)
        tree.last()
        tree.clear()
        self.verify_empty_tree(tree)
        tree.seek_index(ts.num_trees // 2)
        tree.clear()
        self.verify_empty_tree(tree)

    def verify_trees_identical(self, t1, t2):
        assert t1.tree_sequence is t2.tree_sequence
        assert np.all(t1.parent_array == t2.parent_array)
        assert np.all(t1.left_child_array == t2.left_child_array)
        assert np.all(t1.right_child_array == t2.right_child_array)
        assert np.all(t1.left_sib_array == t2.left_sib_array)
        assert np.all(t1.right_sib_array == t2.right_sib_array)
        assert np.all(t1.num_children_array == t2.num_children_array)
        assert np.all(t1.edge_array == t2.edge_array)
        assert list(t1.sites()) == list(t2.sites())

    def test_copy_seek(self):
        ts = msprime.simulate(10, recombination_rate=3, length=3, random_seed=42)
        assert ts.num_trees > 5
        tree = tskit.Tree(ts)
        copy = tree.copy()
        self.verify_empty_tree(copy)
        while tree.next():
            copy = tree.copy()
            self.verify_trees_identical(tree, copy)
        while tree.prev():
            copy = tree.copy()
            self.verify_trees_identical(tree, copy)
        tree.clear()
        copy = tree.copy()
        tree.first()
        # Make sure the underlying arrays are different
        assert np.any(tree.parent_array != copy.parent_array)
        copy.first()
        while tree.index != -1:
            self.verify_trees_identical(tree, copy)
            assert tree.next() == copy.next()
        tree.last()
        copy.last()
        while tree.index != -1:
            self.verify_trees_identical(tree, copy)
            assert tree.prev() == copy.prev()
        # Seek to middle and two independent trees.
        tree.seek_index(ts.num_trees // 2)
        left_copy = tree.copy()
        right_copy = tree.copy()
        self.verify_trees_identical(tree, left_copy)
        self.verify_trees_identical(tree, right_copy)
        left_copy.prev()
        assert left_copy.index == tree.index - 1
        right_copy.next()
        assert right_copy.index == tree.index + 1

    def test_copy_tracked_samples(self):
        ts = msprime.simulate(10, recombination_rate=2, length=3, random_seed=42)
        tree = tskit.Tree(ts, tracked_samples=[0, 1])
        while tree.next():
            copy = tree.copy()
            for j in range(ts.num_nodes):
                assert tree.num_tracked_samples(j) == copy.num_tracked_samples(j)
        copy = tree.copy()
        while tree.next():
            copy.next()
            for j in range(ts.num_nodes):
                assert tree.num_tracked_samples(j) == copy.num_tracked_samples(j)

    def test_copy_multiple_roots(self):
        ts = msprime.simulate(20, recombination_rate=2, length=3, random_seed=42)
        ts = ts.decapitate(np.max(ts.tables.nodes.time) / 2)
        for root_threshold in [1, 2, 100]:
            tree = tskit.Tree(ts, root_threshold=root_threshold)
            copy = tree.copy()
            assert copy.roots == tree.roots
            assert copy.root_threshold == root_threshold
            while tree.next():
                copy = tree.copy()
                assert copy.roots == tree.roots
                assert copy.root_threshold == root_threshold
            copy = tree.copy()
            assert copy.roots == tree.roots
            assert copy.root_threshold == root_threshold

    def test_map_mutations(self):
        ts = msprime.simulate(5, random_seed=42)
        tree = ts.first()
        genotypes = np.zeros(5, dtype=np.int8)
        alleles = [str(j) for j in range(64)]
        ancestral_state, transitions = tree.map_mutations(genotypes, alleles)
        assert ancestral_state == "0"
        assert len(transitions) == 0
        for j in range(1, 64):
            genotypes[0] = j
            ancestral_state, transitions = tree.map_mutations(genotypes, alleles)
            assert ancestral_state == "0"
            assert len(transitions) == 1
        for j in range(64, 67):
            genotypes[0] = j
            with pytest.raises(ValueError):
                tree.map_mutations(genotypes, alleles)
        tree.map_mutations([0] * 5, alleles)
        tree.map_mutations(np.zeros(5, dtype=int), alleles)

    def test_sample_count_deprecated(self):
        ts = msprime.simulate(5, random_seed=42)
        with warnings.catch_warnings(record=True) as w:
            ts.trees(sample_counts=True)
            assert len(w) == 1
            assert issubclass(w[0].category, RuntimeWarning)

        with warnings.catch_warnings(record=True) as w:
            tskit.Tree(ts, sample_counts=False)
            assert len(w) == 1
            assert issubclass(w[0].category, RuntimeWarning)

    def test_node_edges(self):
        ts = msprime.simulate(5, recombination_rate=1, random_seed=42)
        assert ts.num_trees > 2
        edge_table = ts.tables.edges
        for tree in ts.trees():
            nodes = set(tree.nodes())
            midpoint = sum(tree.interval) / 2
            # mapping = tree._node_edges()
            mapping = tree.edge_array
            for node, edge in enumerate(mapping):
                if node in nodes and tree.parent(node) != tskit.NULL:
                    edge_above_node = np.where(
                        np.logical_and.reduce(
                            (
                                edge_table.child == node,
                                edge_table.left < midpoint,
                                edge_table.right > midpoint,
                            )
                        )
                    )[0]
                    assert len(edge_above_node) == 1
                    assert edge_above_node[0] == edge
                else:
                    assert edge == tskit.NULL


class TestSiblings:
    def test_balanced_binary_tree(self):
        t = tskit.Tree.generate_balanced(num_leaves=3)
        assert t.has_single_root
        # Nodes 0 to 2 are leaves
        for u in range(2):
            assert t.is_leaf(u)
        assert t.siblings(0) == (3,)
        assert t.siblings(1) == (2,)
        assert t.siblings(2) == (1,)
        # Node 3 is the internal node
        assert t.is_internal(3)
        assert t.siblings(3) == (0,)
        # Node 4 is the root
        assert 4 == t.root
        assert t.siblings(4) == tuple()
        # Node 5 is the virtual root
        assert 5 == t.virtual_root
        assert t.siblings(5) == tuple()

    def test_star(self):
        t = tskit.Tree.generate_star(num_leaves=3)
        assert t.has_single_root
        # Nodes 0 to 2 are leaves
        for u in range(2):
            assert t.is_leaf(u)
        assert t.siblings(0) == (1, 2)
        assert t.siblings(1) == (0, 2)
        assert t.siblings(2) == (0, 1)
        # Node 3 is the root
        assert 3 == t.root
        assert t.siblings(3) == tuple()
        # Node 4 is the virtual root
        assert 4 == t.virtual_root
        assert t.siblings(4) == tuple()

    def test_multiroot_tree(self):
        ts = tskit.Tree.generate_balanced(4, arity=2).tree_sequence
        t = ts.decapitate(ts.node(5).time).first()
        assert t.has_multiple_roots
        # Nodes 0 to 3 are leaves
        assert t.siblings(0) == (1,)
        assert t.siblings(1) == (0,)
        assert t.siblings(2) == (3,)
        assert t.siblings(3) == (2,)
        # Nodes 4 and 5 are both roots
        assert 4 in t.roots
        assert t.siblings(4) == (5,)
        assert 5 in t.roots
        assert t.siblings(5) == (4,)
        # Node 7 is the virtual root
        assert 7 == t.virtual_root
        assert t.siblings(7) == tuple()

    @pytest.mark.parametrize("flag,expected", [(0, ()), (1, (2,))])
    def test_isolated_node(self, flag, expected):
        tables = tskit.Tree.generate_balanced(2, arity=2).tree_sequence.dump_tables()
        tables.nodes.add_row(flags=flag)  # Add node 3
        t = tables.tree_sequence().first()
        assert t.is_isolated(3)
        assert t.siblings(3) == expected


class TestNodeOrdering(HighLevelTestCase):
    """
    Verify that we can use any node ordering for internal nodes
    and get the same topologies.
    """

    num_random_permutations = 10

    def verify_tree_sequences_equal(self, ts1, ts2, approximate=False):
        assert ts1.get_num_trees() == ts2.get_num_trees()
        assert ts1.get_sample_size() == ts2.get_sample_size()
        assert ts1.get_num_nodes() == ts2.get_num_nodes()
        j = 0
        for r1, r2 in zip(ts1.edges(), ts2.edges()):
            assert r1.parent == r2.parent
            assert r1.child == r2.child
            if approximate:
                assert r1.left == pytest.approx(r2.left)
                assert r1.right == pytest.approx(r2.right)
            else:
                assert r1.left == r2.left
                assert r1.right == r2.right
            j += 1
        assert ts1.num_edges == j
        j = 0
        for n1, n2 in zip(ts1.nodes(), ts2.nodes()):
            assert n1.metadata == n2.metadata
            assert n1.population == n2.population
            if approximate:
                assert n1.time == pytest.approx(n2.time)
            else:
                assert n1.time == n2.time
            j += 1
        assert ts1.num_nodes == j

    def verify_random_permutation(self, ts):
        n = ts.sample_size
        node_map = {}
        for j in range(n):
            node_map[j] = j
        internal_nodes = list(range(n, ts.num_nodes))
        random.shuffle(internal_nodes)
        for j, node in enumerate(internal_nodes):
            node_map[n + j] = node
        other_tables = tskit.TableCollection(ts.sequence_length)
        # Insert the new nodes into the table.
        inv_node_map = {v: k for k, v in node_map.items()}
        for j in range(ts.num_nodes):
            node = ts.node(inv_node_map[j])
            other_tables.nodes.append(node)
        for e in ts.edges():
            other_tables.edges.append(
                e.replace(parent=node_map[e.parent], child=node_map[e.child])
            )
        for _ in range(ts.num_populations):
            other_tables.populations.add_row()
        other_tables.sort()
        other_ts = other_tables.tree_sequence()

        assert ts.get_num_trees() == other_ts.get_num_trees()
        assert ts.get_sample_size() == other_ts.get_sample_size()
        assert ts.get_num_nodes() == other_ts.get_num_nodes()
        j = 0
        for t1, t2 in zip(ts.trees(), other_ts.trees()):
            # Verify the topologies are identical. We do this by traversing
            # upwards to the root for every sample and checking if we map to
            # the correct node and time.
            for u in range(n):
                v_orig = u
                v_map = u
                while v_orig != tskit.NULL:
                    assert node_map[v_orig] == v_map
                    assert t1.get_time(v_orig) == t2.get_time(v_map)
                    v_orig = t1.get_parent(v_orig)
                    v_map = t2.get_parent(v_map)
                assert v_orig == tskit.NULL
                assert v_map == tskit.NULL
            j += 1
        assert j == ts.get_num_trees()
        # Verify we can dump this new tree sequence OK.
        with tempfile.TemporaryDirectory() as tempdir:
            temp_file = pathlib.Path(tempdir) / "tmp.trees"
            other_ts.dump(temp_file)
            ts3 = tskit.load(temp_file)
        self.verify_tree_sequences_equal(other_ts, ts3)
        nodes_file = io.StringIO()
        edges_file = io.StringIO()
        # Also verify we can read the text version.
        other_ts.dump_text(nodes=nodes_file, edges=edges_file, precision=14)
        nodes_file.seek(0)
        edges_file.seek(0)
        ts3 = tskit.load_text(nodes_file, edges_file)
        self.verify_tree_sequences_equal(other_ts, ts3, True)

    def test_single_locus(self):
        ts = msprime.simulate(7)
        for _ in range(self.num_random_permutations):
            self.verify_random_permutation(ts)

    def test_multi_locus(self):
        ts = msprime.simulate(20, recombination_rate=10)
        for _ in range(self.num_random_permutations):
            self.verify_random_permutation(ts)

    def test_nonbinary(self):
        ts = msprime.simulate(
            sample_size=20,
            recombination_rate=10,
            demographic_events=[
                msprime.SimpleBottleneck(time=0.5, population=0, proportion=1)
            ],
        )
        # Make sure this really has some non-binary nodes
        found = False
        for t in ts.trees():
            for u in t.nodes():
                if len(t.children(u)) > 2:
                    found = True
                    break
            if found:
                break
        assert found
        for _ in range(self.num_random_permutations):
            self.verify_random_permutation(ts)


def assert_trees_identical(t1, t2):
    assert t1.tree_sequence == t2.tree_sequence
    assert t1.index == t2.index
    assert np.all(t1.parent_array == t2.parent_array)
    assert np.all(t1.left_child_array == t2.left_child_array)
    assert np.all(t1.left_sib_array == t2.left_sib_array)
    assert np.all(t1.right_child_array == t2.right_child_array)
    assert np.all(t1.right_sib_array == t2.right_sib_array)


def assert_same_tree_different_order(t1, t2):
    assert t1.tree_sequence == t2.tree_sequence
    assert t1.index == t2.index
    assert np.all(t1.parent_array == t2.parent_array)
    assert not np.all(t1.left_child_array == t2.left_child_array)


def seek(tree, x):
    """
    Python implementation of the seek algorithm. Useful for developing
    tests.
    """
    L = tree.tree_sequence.sequence_length
    t_l, t_r = tree.interval
    if x < t_l:
        # |-----|-----|========|---------|
        # 0     x    t_l      t_r        L
        distance_left = t_l - x
        distance_right = L - t_r + x
    else:
        # |------|========|------|-------|
        # 0     t_l      t_r     x       L
        distance_right = x - t_r
        distance_left = t_l + L - x
    if distance_right <= distance_left:
        while not (tree.interval.left <= x < tree.interval.right):
            tree.next()
    else:
        while not (tree.interval.left <= x < tree.interval.right):
            tree.prev()


class TestSeekDirection:
    """
    Test if we seek in the correct direction according to our hueristics.
    """

    # 2.00┊       ┊   4   ┊   4   ┊   4   ┊
    #     ┊       ┊ ┏━┻┓  ┊  ┏┻━┓ ┊  ┏┻━┓ ┊
    # 1.00┊   3   ┊ ┃  3  ┊  3  ┃ ┊  3  ┃ ┊
    #     ┊ ┏━╋━┓ ┊ ┃ ┏┻┓ ┊ ┏┻┓ ┃ ┊ ┏┻┓ ┃ ┊
    # 0.00┊ 0 1 2 ┊ 0 1 2 ┊ 0 2 1 ┊ 0 1 2 ┊
    #     0       1       2       3       4
    @tests.cached_example
    def ts(self):
        return tsutil.all_trees_ts(3)

    def get_tree_pair(self):
        ts = self.ts()
        t1 = tskit.Tree(ts)
        t2 = tskit.Tree(ts)
        # Note: for development we can monkeypatch in the Python implementation
        # above like this:
        # t2.seek = functools.partial(seek, t2)
        return t1, t2

    @pytest.mark.parametrize("index", range(4))
    def test_index_from_different_directions(self, index):
        # Check that we get different orderings of the children arrays
        # for all trees when we go in different directions.
        t1, t2 = self.get_tree_pair()
        while t1.index != index:
            t1.next()
        while t2.index != index:
            t2.prev()
        assert_same_tree_different_order(t1, t2)

    @pytest.mark.parametrize("position", [0, 1, 2, 3])
    def test_seek_from_null(self, position):
        t1, t2 = self.get_tree_pair()
        t1.clear()
        t1.seek(position)
        t2.first()
        t2.seek(position)
        assert_trees_identical(t1, t2)

    @pytest.mark.parametrize("index", range(3))
    def test_seek_next_tree(self, index):
        t1, t2 = self.get_tree_pair()
        while t1.index != index:
            t1.next()
            t2.next()
        t1.next()
        t2.seek(index + 1)
        assert_trees_identical(t1, t2)

    @pytest.mark.parametrize("index", [3, 2, 1])
    def test_seek_prev_tree(self, index):
        t1, t2 = self.get_tree_pair()
        while t1.index != index:
            t1.prev()
            t2.prev()
        t1.prev()
        t2.seek(index - 1)
        assert_trees_identical(t1, t2)

    def test_seek_1_from_0(self):
        t1, t2 = self.get_tree_pair()
        t1.first()
        t1.next()
        t2.first()
        t2.seek(1)
        assert_trees_identical(t1, t2)

    def test_seek_1_5_from_0(self):
        t1, t2 = self.get_tree_pair()
        t1.first()
        t1.next()
        t2.first()
        t2.seek(1.5)
        assert_trees_identical(t1, t2)

    def test_seek_1_5_from_1(self):
        t1, t2 = self.get_tree_pair()
        for _ in range(2):
            t1.next()
            t2.next()
        t2.seek(1.5)
        assert_trees_identical(t1, t2)

    def test_seek_3_from_null(self):
        t1, t2 = self.get_tree_pair()
        t1.last()
        t2.seek(3)
        assert_trees_identical(t1, t2)

    def test_seek_3_from_null_prev(self):
        t1, t2 = self.get_tree_pair()
        t1.last()
        t1.prev()
        t2.seek(3)
        t2.prev()
        assert_trees_identical(t1, t2)

    def test_seek_3_from_0(self):
        t1, t2 = self.get_tree_pair()
        t1.last()
        t2.first()
        t2.seek(3)
        assert_trees_identical(t1, t2)

    def test_seek_0_from_3(self):
        t1, t2 = self.get_tree_pair()
        t1.last()
        t1.first()
        t2.last()
        t2.seek(0)
        assert_trees_identical(t1, t2)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_seek_mid_null_and_middle(self, ts):
        breakpoints = ts.breakpoints(as_array=True)
        mid = breakpoints[:-1] + np.diff(breakpoints) / 2
        for index, x in enumerate(mid[:-1]):
            t1 = tskit.Tree(ts)
            t1.seek(x)
            # Also seek to this point manually to make sure we're not
            # reusing the seek from null under the hood.
            t2 = tskit.Tree(ts)
            if index <= ts.num_trees / 2:
                while t2.index != index:
                    t2.next()
            else:
                while t2.index != index:
                    t2.prev()
            assert t1.index == t2.index
            assert np.all(t1.parent_array == t2.parent_array)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_seek_last_then_prev(self, ts):
        t1 = tskit.Tree(ts)
        t1.seek(ts.sequence_length - 0.00001)
        assert t1.index == ts.num_trees - 1
        t2 = tskit.Tree(ts)
        t2.prev()
        assert_trees_identical(t1, t2)
        t1.prev()
        t2.prev()
        assert_trees_identical(t1, t2)


class TestSeek:
    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_new_seek_breakpoints(self, ts):
        breakpoints = ts.breakpoints(as_array=True)
        for index, left in enumerate(breakpoints[:-1]):
            tree = tskit.Tree(ts)
            tree.seek(left)
            assert tree.index == index

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_new_seek_mid(self, ts):
        breakpoints = ts.breakpoints(as_array=True)
        mid = breakpoints[:-1] + np.diff(breakpoints) / 2
        for index, left in enumerate(mid[:-1]):
            tree = tskit.Tree(ts)
            tree.seek(left)
            assert tree.index == index

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_same_seek_breakpoints(self, ts):
        breakpoints = ts.breakpoints(as_array=True)
        tree = tskit.Tree(ts)
        for index, left in enumerate(breakpoints[:-1]):
            tree.seek(left)
            assert tree.index == index

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_new_seek_breakpoints_reversed(self, ts):
        breakpoints = ts.breakpoints(as_array=True)
        for index, left in reversed(list(enumerate(breakpoints[:-1]))):
            tree = tskit.Tree(ts)
            tree.seek(left)
            assert tree.index == index

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_same_seek_breakpoints_reversed(self, ts):
        breakpoints = ts.breakpoints(as_array=True)
        tree = tskit.Tree(ts)
        for index, left in reversed(list(enumerate(breakpoints[:-1]))):
            tree.seek(left)
            assert tree.index == index

    def test_example(self):
        L = 10
        ts = msprime.simulate(10, recombination_rate=3, length=L, random_seed=42)
        assert ts.num_trees > 5
        same_tree = tskit.Tree(ts)
        for j in range(L):
            for tree in [same_tree, tskit.Tree(ts)]:
                tree.seek(j)
                index = tree.index
                assert tree.interval.left <= j < tree.interval.right
                tree.seek(tree.interval.left)
                assert tree.index == index
                if tree.interval.right < L:
                    tree.seek(tree.interval.right)
                    assert tree.index == index + 1
        for j in reversed(range(L)):
            for tree in [same_tree, tskit.Tree(ts)]:
                tree.seek(j)
                assert tree.interval.left <= j < tree.interval.right

    def test_errors(self, ts_fixture):
        L = ts_fixture.sequence_length
        tree = tskit.Tree(ts_fixture)
        for bad_position in [-1, L, L + 1, -L]:
            with pytest.raises(ValueError):
                tree.seek(bad_position)


class SimpleContainersMixin:
    """
    Tests for the SimpleContainer classes.
    """

    def test_equality(self):
        c1, c2 = self.get_instances(2)
        assert c1 == c1
        assert not (c1 == c2)
        assert not (c1 != c1)
        assert c1 != c2
        (c3,) = self.get_instances(1)
        assert c1 == c3
        assert not (c1 != c3)

    def test_repr(self):
        (c,) = self.get_instances(1)
        assert len(repr(c)) > 0


class SimpleContainersWithMetadataMixin:
    """
    Tests for the SimpleContainerWithMetadata classes.
    """

    def test_metadata(self):
        # Test decoding
        instances = self.get_instances(5)
        for j, inst in enumerate(instances):
            assert inst.metadata == ("x" * j) + "decoded"

        # Decoder doesn't effect equality
        (inst,) = self.get_instances(1)
        (inst2,) = self.get_instances(1)
        assert inst == inst2
        inst._metadata = "different"
        assert inst != inst2

    def test_decoder_run_once(self):
        # For a given instance, the decoded metadata should be cached, with the decoder
        # called once
        (inst,) = self.get_instances(1)
        times_run = 0

        # Hack in a tracing decoder
        def decoder(m):
            nonlocal times_run
            times_run += 1
            return m.decode() + "decoded"

        inst._metadata_decoder = decoder
        assert times_run == 0
        _ = inst.metadata
        assert times_run == 1
        _ = inst.metadata
        assert times_run == 1


class TestIndividualContainer(SimpleContainersMixin, SimpleContainersWithMetadataMixin):
    def get_instances(self, n):
        return [
            tskit.Individual(
                id=j,
                flags=j,
                location=[j],
                parents=[j],
                nodes=[j],
                metadata=b"x" * j,
                metadata_decoder=lambda m: m.decode() + "decoded",
            )
            for j in range(n)
        ]


class TestNodeContainer(SimpleContainersMixin, SimpleContainersWithMetadataMixin):
    def get_instances(self, n):
        return [
            tskit.Node(
                id=j,
                flags=j,
                time=j,
                population=j,
                individual=j,
                metadata=b"x" * j,
                metadata_decoder=lambda m: m.decode() + "decoded",
            )
            for j in range(n)
        ]


class TestEdgeContainer(SimpleContainersMixin, SimpleContainersWithMetadataMixin):
    def get_instances(self, n):
        return [
            tskit.Edge(
                left=j,
                right=j,
                parent=j,
                child=j,
                metadata=b"x" * j,
                metadata_decoder=lambda m: m.decode() + "decoded",
                id=j,
            )
            for j in range(n)
        ]


class TestSiteContainer(SimpleContainersMixin, SimpleContainersWithMetadataMixin):
    def get_instances(self, n):
        return [
            tskit.Site(
                id=j,
                position=j,
                ancestral_state="A" * j,
                mutations=TestMutationContainer().get_instances(j),
                metadata=b"x" * j,
                metadata_decoder=lambda m: m.decode() + "decoded",
            )
            for j in range(n)
        ]


class TestMutationContainer(SimpleContainersMixin, SimpleContainersWithMetadataMixin):
    def get_instances(self, n):
        return [
            tskit.Mutation(
                id=j,
                site=j,
                node=j,
                time=j,
                derived_state="A" * j,
                parent=j,
                metadata=b"x" * j,
                metadata_decoder=lambda m: m.decode() + "decoded",
            )
            for j in range(n)
        ]

    def test_nan_equality(self):
        a = tskit.Mutation(
            id=42,
            site=42,
            node=42,
            time=UNKNOWN_TIME,
            derived_state="A" * 42,
            parent=42,
            metadata=b"x" * 42,
            metadata_decoder=lambda m: m.decode() + "decoded",
        )
        b = tskit.Mutation(
            id=42,
            site=42,
            node=42,
            derived_state="A" * 42,
            parent=42,
            metadata=b"x" * 42,
            metadata_decoder=lambda m: m.decode() + "decoded",
        )
        c = tskit.Mutation(
            id=42,
            site=42,
            node=42,
            time=math.nan,
            derived_state="A" * 42,
            parent=42,
            metadata=b"x" * 42,
            metadata_decoder=lambda m: m.decode() + "decoded",
        )
        assert a == a
        assert a == b
        assert not (a == c)
        assert not (b == c)
        assert not (a != a)
        assert not (a != b)
        assert a != c
        assert c != c
        assert not (c == c)


class TestMigrationContainer(SimpleContainersMixin, SimpleContainersWithMetadataMixin):
    def get_instances(self, n):
        return [
            tskit.Migration(
                id=j,
                left=j,
                right=j,
                node=j,
                source=j,
                dest=j,
                time=j,
                metadata=b"x" * j,
                metadata_decoder=lambda m: m.decode() + "decoded",
            )
            for j in range(n)
        ]


class TestPopulationContainer(SimpleContainersMixin, SimpleContainersWithMetadataMixin):
    def get_instances(self, n):
        return [
            tskit.Population(
                id=j,
                metadata=b"x" * j,
                metadata_decoder=lambda m: m.decode() + "decoded",
            )
            for j in range(n)
        ]


class TestProvenanceContainer(SimpleContainersMixin):
    def get_instances(self, n):
        return [
            tskit.Provenance(id=j, timestamp="x" * j, record="y" * j) for j in range(n)
        ]


class TestEdgesetContainer(SimpleContainersMixin):
    def get_instances(self, n):
        return [tskit.Edgeset(left=j, right=j, parent=j, children=j) for j in range(n)]


class TestContainersAppend:
    def test_containers_append(self, ts_fixture):
        """
        Test that the containers work with `Table.append`
        """
        tables = ts_fixture.dump_tables()
        tables.clear(clear_provenance=True)
        for table_name in tskit.TABLE_NAMES:
            table = getattr(tables, table_name)
            for i in range(len(getattr(ts_fixture.tables, table_name))):
                table.append(getattr(ts_fixture, table_name[:-1])(i))
        ts_fixture.tables.assert_equals(tables)


class TestTskitConversionOutput(unittest.TestCase):
    """
    Tests conversion output to ensure it is correct.
    """

    @classmethod
    def setUpClass(cls):
        ts = msprime.simulate(
            length=1,
            recombination_rate=2,
            mutation_rate=2,
            random_seed=1,
            migration_matrix=[[0, 1], [1, 0]],
            population_configurations=[
                msprime.PopulationConfiguration(5) for _ in range(2)
            ],
            record_migrations=True,
        )
        assert ts.num_migrations > 0
        cls._tree_sequence = tsutil.insert_random_ploidy_individuals(ts)

    def test_macs(self):
        output = self._tree_sequence.to_macs().splitlines()
        assert output[0].startswith("COMMAND:")
        assert output[1].startswith("SEED:")
        assert len(output) == 2 + self._tree_sequence.get_num_mutations()
        n = self._tree_sequence.get_sample_size()
        m = self._tree_sequence.get_sequence_length()
        sites = list(self._tree_sequence.sites())
        haplotypes = list(self._tree_sequence.haplotypes())
        for site_id, line in enumerate(output[2:]):
            splits = line.split()
            assert splits[0] == "SITE:"
            assert int(splits[1]) == site_id
            position = sites[site_id].position / m
            self.assertAlmostEqual(float(splits[2]), position)
            col = splits[4]
            assert len(col) == n
            for j in range(n):
                assert col[j] == haplotypes[j][site_id]

    def test_macs_error(self):
        tables = tskit.TableCollection(1)
        tables.sites.add_row(position=0.5, ancestral_state="A")
        tables.nodes.add_row(time=1, flags=tskit.NODE_IS_SAMPLE)
        tables.mutations.add_row(node=0, site=0, derived_state="FOO")
        ts = tables.tree_sequence()
        with pytest.raises(
            ValueError, match="macs output only supports single letter alleles"
        ):
            ts.to_macs()


class TestTreeSequenceGetSite:
    """
    Tests for getting Site objects from a TreeSequence object
    by specifying the position.
    """

    def get_example_ts_discrete_coordinates(self):
        tables = tskit.TableCollection(sequence_length=10)
        tables.sites.add_row(position=3, ancestral_state="A")
        tables.sites.add_row(position=5, ancestral_state="C")
        tables.sites.add_row(position=7, ancestral_state="G")
        return tables.tree_sequence()

    def get_example_ts_continuous_coordinates(self):
        tables = tskit.TableCollection(sequence_length=10)
        tables.sites.add_row(position=0.5, ancestral_state="A")
        tables.sites.add_row(position=6.2, ancestral_state="C")
        tables.sites.add_row(position=8.3, ancestral_state="T")
        return tables.tree_sequence()

    def get_example_ts_without_sites(self):
        tables = tskit.TableCollection(sequence_length=10)
        return tables.tree_sequence()

    @pytest.mark.parametrize("id_", [0, 1, 2])
    def test_site_id(self, id_):
        ts = self.get_example_ts_discrete_coordinates()
        site = ts.site(id_)
        assert site.id == id_

    @pytest.mark.parametrize("position", [3, 5, 7])
    def test_position_discrete_coordinates(self, position):
        ts = self.get_example_ts_discrete_coordinates()
        site = ts.site(position=position)
        assert site.position == position

    @pytest.mark.parametrize("position", [0.5, 6.2, 8.3])
    def test_position_continuous_coordinates(self, position):
        ts = self.get_example_ts_continuous_coordinates()
        site = ts.site(position=position)
        assert site.position == position

    @pytest.mark.parametrize("position", [0, 2.999999999, 5.000000001, 9])
    def test_position_not_found(self, position):
        with pytest.raises(ValueError, match=r"There is no site at position"):
            ts = self.get_example_ts_discrete_coordinates()
            ts.site(position=position)

    @pytest.mark.parametrize(
        "position",
        [
            np.array([3], dtype=float)[0],
            np.array([3], dtype=int)[0],
            decimal.Decimal(3),
        ],
    )
    def test_position_good_type(self, position):
        ts = self.get_example_ts_discrete_coordinates()
        ts.site(position=position)

    def test_position_not_scalar(self):
        with pytest.raises(
            ValueError, match="Position must be provided as a scalar value."
        ):
            ts = self.get_example_ts_discrete_coordinates()
            ts.site(position=[1, 4, 8])

    @pytest.mark.parametrize("position", [-1, 10, 11])
    def test_position_out_of_bounds(self, position):
        with pytest.raises(
            ValueError,
            match="Position is beyond the coordinates defined by sequence length.",
        ):
            ts = self.get_example_ts_discrete_coordinates()
            ts.site(position=position)

    def test_query_position_siteless_ts(self):
        with pytest.raises(ValueError, match=r"There is no site at position"):
            ts = self.get_example_ts_without_sites()
            ts.site(position=1)

    def test_site_id_and_position_are_none(self):
        with pytest.raises(TypeError, match="Site id or position must be provided."):
            ts = self.get_example_ts_discrete_coordinates()
            ts.site(None, position=None)

    def test_site_id_and_position_are_specified(self):
        with pytest.raises(
            TypeError, match="Only one of site id or position needs to be provided."
        ):
            ts = self.get_example_ts_discrete_coordinates()
            ts.site(0, position=3)


def num_lineages_definition(tree, t):
    lineages = 0
    for u in tree.nodes():
        v = tree.parent(u)
        if v != tskit.NULL:
            if tree.time(u) <= t < tree.time(v):
                lineages += 1
    return lineages


class TestNumLineages:
    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_tree_midpoint_definition(self, ts):
        t = 0
        if ts.num_nodes > 0:
            t = np.max(ts.tables.nodes.time) / 2
        tree = ts.first()
        assert tree.num_lineages(t) == num_lineages_definition(tree, t)

    @pytest.mark.parametrize("t", [-np.inf, np.inf, np.nan])
    def test_nonfinite_time(self, t):
        tree = tskit.Tree.generate_balanced(2)
        with pytest.raises(tskit.LibraryError, match="NONFINITE"):
            tree.num_lineages(t)

    @pytest.mark.parametrize("t", [1, 1.0, np.array([1.0])[0]])
    def test_number_types(self, t):
        tree = tskit.Tree.generate_balanced(2)
        assert tree.num_lineages(t) == 0

    # 2.00┊        12         ┊
    #     ┊   ┏━━━━━╋━━━━━┓   ┊
    # 1.00┊   9    10    11   ┊
    #     ┊ ┏━╋━┓ ┏━╋━┓ ┏━╋━┓ ┊
    # 0.00┊ 0 1 2 3 4 5 6 7 8 ┊
    #     0                   1
    @pytest.mark.parametrize(
        ["t", "expected"],
        [
            (-0.00001, 0),
            (0, 9),
            (0.0000001, 9),
            (0.99999, 9),
            (1, 3),
            (1.999999, 3),
            (2, 0),
            (2.000001, 0),
        ],
    )
    def test_balanced_ternary(self, t, expected):
        tree = tskit.Tree.generate_balanced(9, arity=3)
        assert tree.num_lineages(t) == expected

    # 3.00┊            15     ┊
    #     ┊          ┏━━┻━┓   ┊
    # 2.00┊   11     ┃   14   ┊
    #     ┊  ┏━┻━┓   ┃  ┏━┻┓  ┊
    # 1.00┊  9  10  12  ┃ 13  ┊
    #     ┊ ┏┻┓ ┏┻┓ ┏┻┓ ┃ ┏┻┓ ┊
    # 0.00┊ 0 1 2 3 4 5 6 7 8 ┊
    #     0                   1
    @pytest.mark.parametrize(
        ["t", "expected"],
        [
            (-0.00001, 0),
            (0, 9),
            (0.0000001, 9),
            (0.99999, 9),
            (1, 5),
            (1.999999, 5),
            (2, 2),
            (2.000001, 2),
            (3.00000, 0),
            (5.00000, 0),
        ],
    )
    def test_multiroot_different_times(self, t, expected):
        tables = tskit.Tree.generate_balanced(9, arity=2).tree_sequence.dump_tables()
        edges = tables.edges.copy()
        tables.edges.clear()
        for edge in edges:
            if edge.parent != 16:
                tables.edges.append(edge)
        ts = tables.tree_sequence()
        tree = ts.first()
        assert tree.num_lineages(t) == expected

    # 4.00┊   8       ┊
    #     ┊ ┏━┻━┓     ┊
    # 3.00┊ 0   7     ┊
    #     ┊   ┏━┻━┓   ┊
    # 2.00┊   1   6   ┊
    #     ┊     ┏━┻┓  ┊
    # 1.00┊     2  5  ┊
    #     ┊       ┏┻┓ ┊
    # 0.00┊       3 4 ┊
    #     0           1
    @pytest.mark.parametrize(
        ["t", "expected"],
        [
            (-0.00001, 0),
            (0, 2),
            (1, 2),
            (2, 2),
            (3, 2),
            (3, 2),
            (4, 0),
        ],
    )
    def test_comb_different_leaf_times(self, t, expected):
        tables = tskit.Tree.generate_comb(5).tree_sequence.dump_tables()
        time = tables.nodes.time
        time[2] = 1
        time[1] = 2
        time[0] = 3
        tables.nodes.time = time
        ts = tables.tree_sequence()
        tree = ts.first()
        assert tree.num_lineages(t) == expected

    @pytest.mark.parametrize(
        ["t", "expected"],
        [
            (-0.00001, 0),
            (0, 0),
            (1, 0),
            (2, 0),
            (3, 0),
        ],
    )
    def test_missing_data_different_times(self, t, expected):
        tables = tskit.TableCollection(1)
        for j in range(3):
            tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=j)
        ts = tables.tree_sequence()
        tree = ts.first()
        assert tree.num_lineages(t) == expected


--- ../../tskit/python/tests/ibd.py ---


"""
Python implementation of the IBD-finding algorithms.
"""
import argparse
import collections

import numpy as np

import tskit


class Segment:
    """
    A class representing a single segment. Each segment has a left and right,
    denoting the loci over which it spans, a node and a next, giving the next
    in the chain.

    The node it records is the *output* node ID.
    """

    def __init__(self, left=None, right=None, node=None, next_seg=None):
        self.left = left
        self.right = right
        self.node = node
        self.next = next_seg

    def __str__(self):
        s = "({}-{}->{}:next={})".format(
            self.left, self.right, self.node, repr(self.next)
        )
        return s

    def __repr__(self):
        return repr((self.left, self.right, self.node))

    def __eq__(self, other):
        # NOTE: to simplify tests, we DON'T check for equality of 'next'.
        return (
            self.left == other.left
            and self.right == other.right
            and self.node == other.node
        )

    def __lt__(self, other):
        return (self.node, self.left, self.right) < (
            other.node,
            other.left,
            other.right,
        )


class SegmentList:
    """
    A class representing a list of segments that are descended from a given ancestral
    node via a particular child of the ancestor.
    Each SegmentList keeps track of the first and last segment in the list, head and
    tail.
    """

    def __init__(self, head=None, tail=None):
        self.head = head
        self.tail = tail

    def __str__(self):
        return repr(self)

    def __repr__(self):
        tuple_segs = []
        seg = self.head
        while seg is not None:
            tuple_segs.append((seg.left, seg.right, seg.node))
            seg = seg.next
        return repr(tuple_segs)

    def extend(self, seglist):
        """
        Extends this segment list with the segments in the specified list.
        """
        assert isinstance(seglist, SegmentList)
        if seglist.head is not None:
            if self.head is None:
                self.head = seglist.head
                self.tail = seglist.tail
            else:
                self.tail.next = seglist.head
                self.tail = seglist.tail

    def append(self, segment):
        """
        Append the specified segment to the end of this list.
        """
        assert isinstance(segment, Segment)
        if self.head is None:
            self.head = segment
            self.tail = segment
        else:
            self.tail.next = segment
            self.tail = segment


class IbdResult:
    """
    Class representing the IBD segments in a tree sequence for a given
    set of sample pairs.
    """

    def __init__(self):
        self.segments = collections.defaultdict(list)

    def __repr__(self):
        return repr(self.segments)

    def __str__(self):
        return repr(self)

    def add_segment_deprecated(self, a, b, seg):
        # The original version of add_segment that doesn't sort or squash.
        key = (a, b) if a < b else (b, a)
        self.segments[key].append(tskit.IdentitySegment(seg.left, seg.right, seg.node))

    def add_segment(self, a, b, seg):
        key = (a, b) if a < b else (b, a)

        # Get position and add into the correct position.
        current_segs = self.segments[key]
        num_segs = len(current_segs)

        if num_segs == 0:
            self.segments[key].append(
                tskit.IdentitySegment(seg.left, seg.right, seg.node)
            )
        else:
            # Find the position for the new segment.
            i = 0
            while (
                i < num_segs
                and current_segs[i].node <= seg.node
                and current_segs[i].right <= seg.left
            ):
                i += 1

            # Calculate boolean values that determine whether to squash
            # and if so, where.
            PUT_FIRST = False  # Insert segment at start of list.
            PUT_LAST = False  # Insert segment at end of list.
            SQUASH_LEFT = False  # Squash with the left segment.
            SQUASH_RIGHT = False  # Squash with the right segment.

            if i == 0:
                PUT_FIRST = True
            if i == num_segs:
                PUT_LAST = True
            if not PUT_FIRST:
                if (
                    current_segs[i - 1].node == seg.node
                    and current_segs[i - 1].right == seg.left
                ):
                    SQUASH_LEFT = True
            if not PUT_LAST:
                if (
                    seg.node == current_segs[i].node
                    and seg.right == current_segs[i].left
                ):
                    SQUASH_RIGHT = True

            # Insert the new segment and squash if needed.
            if SQUASH_LEFT and not SQUASH_RIGHT:
                current_segs[i - 1].right = seg.right
            elif SQUASH_RIGHT and not SQUASH_LEFT:
                current_segs[i].left = seg.left
            elif SQUASH_LEFT and SQUASH_RIGHT:
                # To squash twice, must pop one of the existing segments.
                current_segs[i - 1].right = current_segs[i].right
                current_segs.pop(i)
            else:
                self.segments[key].insert(
                    i, tskit.IdentitySegment(seg.left, seg.right, seg.node)
                )


class IbdFinder:
    """
    Finds all IBD relationships between specified sample pairs in a tree sequence.
    """

    def __init__(self, ts, *, within=None, between=None, min_span=0, max_time=None):
        self.ts = ts
        self.result = IbdResult()
        if within is not None and between is not None:
            raise ValueError("within and between are mutually exclusive")

        self.sample_set_id = np.zeros(ts.num_nodes, dtype=int) - 1
        self.finding_between = False
        if between is not None:
            self.finding_between = True
            for set_id, samples in enumerate(between):
                self.sample_set_id[samples] = set_id
        else:
            if within is None:
                within = ts.samples()
            self.sample_set_id[within] = 0
        self.min_span = min_span
        self.max_time = np.inf if max_time is None else max_time
        self.A = [SegmentList() for _ in range(ts.num_nodes)]  # Descendant segments
        for u in range(ts.num_nodes):
            if self.sample_set_id[u] != -1:
                self.A[u].append(Segment(0, ts.sequence_length, u))
        self.tables = self.ts.tables

    def print_state(self):
        print("IBD Finder")
        print("min_span = ", self.min_span)
        print("max_time   = ", self.max_time)
        print("finding_between = ", self.finding_between)
        print("u\tset_id\tA = ")
        for u, a in enumerate(self.A):
            print(u, self.sample_set_id[u], a, sep="\t")

    def run(self, squash=False):
        node_times = self.tables.nodes.time
        for e in self.ts.edges():
            time = node_times[e.parent]
            if time > self.max_time:
                # Stop looking for IBD segments once the
                # processed nodes are older than the max time.
                break
            child_segs = SegmentList()
            s = self.A[e.child].head
            while s is not None:
                intvl = (
                    max(e.left, s.left),
                    min(e.right, s.right),
                )
                # if intvl[1] - intvl[0] > self.min_span:
                child_segs.append(Segment(intvl[0], intvl[1], s.node))
                s = s.next
            self.record_ibd(e.parent, child_segs, squash=squash)
            self.A[e.parent].extend(child_segs)
        if self.min_span > 0:
            self.filter_by_min_span()
        return self.result.segments

    def record_ibd(self, current_parent, child_segs, squash):
        """
        Given the specified set of child segments for the current parent
        record the IBD segments that will occur as a result of adding these
        new segments into the existing list.
        """
        # Note the implementation here is O(n^2) because we have to compare
        # every segment with every other one. If the segments were stored in
        # left-to-right sorted order, we could avoid and merge them more
        # efficiently. There is some added complexity in doing this, however.
        seg0 = self.A[current_parent].head
        while seg0 is not None:
            seg1 = child_segs.head
            while seg1 is not None:
                left = max(seg0.left, seg1.left)
                right = min(seg0.right, seg1.right)
                # If there are any overlapping segments, record as a new
                # IBD relationship.
                if self.passes_filters(seg0.node, seg1.node, left, right):
                    if squash:
                        self.result.add_segment(
                            seg0.node,
                            seg1.node,
                            Segment(left, right, current_parent),
                        )
                    else:
                        self.result.add_segment_deprecated(
                            seg0.node,
                            seg1.node,
                            Segment(left, right, current_parent),
                        )
                seg1 = seg1.next
            seg0 = seg0.next

    def filter_by_min_span(self):
        """
        Remove any IBD segments that are smaller than min_span.
        Note that we can't do this until we have squashed the IBD segments
        """
        keys_to_pop = []
        for key in self.result.segments.keys():
            self.result.segments[key] = [
                s for s in self.result.segments[key] if s.right - s.left > self.min_span
            ]
            if len(self.result.segments[key]) == 0:
                keys_to_pop.append(key)

        # Remove any keys that now have no IBD segments.
        for key in keys_to_pop:
            self.result.segments.pop(key)

    def passes_filters(self, a, b, left, right):
        if a == b:
            return False
        if right - left <= 0:
            return False
        if self.finding_between:
            return self.sample_set_id[a] != self.sample_set_id[b]
        else:
            return True


if __name__ == "__main__":
    """
    A simple CLI for running IBDFinder on a command line from the `python`
    subdirectory. Basic usage:
    > python3 ./tests/ibd.py --infile test.trees
    """

    parser = argparse.ArgumentParser(
        description="Command line interface for the IBDFinder."
    )

    parser.add_argument(
        "--infile",
        type=str,
        dest="infile",
        nargs=1,
        metavar="IN_FILE",
        help="The tree sequence to be analysed.",
    )

    parser.add_argument(
        "--min-length",
        type=float,
        dest="min_span",
        nargs=1,
        metavar="MIN_LENGTH",
        help="Only segments longer than this cutoff will be returned.",
    )

    parser.add_argument(
        "--max-time",
        type=float,
        dest="max_time",
        nargs=1,
        metavar="MAX_TIME",
        help="Only segments younger this time will be returned.",
    )

    parser.add_argument(
        "--samples",
        type=int,
        dest="samples",
        nargs=2,
        metavar="SAMPLES",
        help="If provided, only this pair's IBD info is returned.",
    )

    args = parser.parse_args()
    ts = tskit.load(args.infile[0])
    if args.min_span is None:
        min_span = 0
    else:
        min_span = args.min_span[0]
    if args.max_time is None:
        max_time = None
    else:
        max_time = args.max_time[0]

    s = IbdFinder(ts, min_span=min_span, max_time=max_time)
    all_segs = s.run()

    if args.samples is None:
        print(all_segs)
    else:
        samples = args.samples
        print(all_segs[(samples[0], samples[1])])


--- ../../tskit/python/tests/test_table_transforms.py ---


"""
Test cases for table transformation operations like trim(), decapitate, etc.
"""
import decimal
import fractions
import io
import math

import numpy as np
import pytest

import tests
import tskit
import tskit.util as util
from tests.test_highlevel import get_example_tree_sequences

# ↑ See https://github.com/tskit-dev/tskit/issues/1804 for when
# we can remove this.


def delete_older_definition(tables, time):
    node_time = tables.nodes.time
    edges = tables.edges.copy()
    tables.edges.clear()
    for edge in edges:
        if node_time[edge.parent] <= time:
            tables.edges.append(edge)

    mutations = tables.mutations.copy()
    # Map of old ID -> new ID
    mutation_map = np.full(len(mutations), tskit.NULL, dtype=int)
    tables.mutations.clear()
    keep = []
    for j, mutation in enumerate(mutations):
        mutation_time = (
            node_time[mutation.node]
            if util.is_unknown_time(mutation.time)
            else mutation.time
        )
        if mutation_time < time:
            mutation_map[j] = len(keep)
            keep.append(mutation)
    # Not making assumptions about ordering, so do it in two passes.
    for mutation in keep:
        if mutation.parent != tskit.NULL:
            mutation = mutation.replace(parent=mutation_map[mutation.parent])
        tables.mutations.append(mutation)

    migrations = tables.migrations.copy()
    tables.migrations.clear()
    for migration in migrations:
        if migration.time < time:
            tables.migrations.append(migration)


class TestDeleteOlderExamples:
    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_definition(self, ts):
        time = 0 if ts.num_nodes == 0 else np.median(ts.tables.nodes.time)
        tables1 = ts.dump_tables()
        delete_older_definition(tables1, time)
        tables2 = ts.dump_tables()
        tables2.delete_older(time)
        tables1.assert_equals(tables2, ignore_provenance=True)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_mutation_parents(self, ts):
        time = 0 if ts.num_nodes == 0 else np.median(ts.tables.nodes.time)
        tables1 = ts.dump_tables()
        tables1.delete_older(time)
        tables2 = tables1.copy()
        tables2.build_index()
        tables2.compute_mutation_parents()
        tables1.assert_equals(tables2, ignore_provenance=True)


class TestDeleteOlderSimpleTree:
    # 2.00┊   4   ┊
    #     ┊ ┏━┻┓  ┊
    # 1.00┊ ┃  3  ┊
    #     ┊ ┃ ┏┻┓ ┊
    # 0.00┊ 0 1 2 ┊
    #     0       1
    def tables(self):
        # Don't cache this because we modify the result!
        tree = tskit.Tree.generate_balanced(3, branch_length=1)
        return tree.tree_sequence.dump_tables()

    @pytest.mark.parametrize("time", [0, -0.5, -100, 0.01, 0.999])
    def test_before_first_internal_node(self, time):
        tables = self.tables()
        before = tables.copy()
        tables.delete_older(time)
        ts = tables.tree_sequence()
        assert ts.num_trees == 1
        tree = ts.first()
        assert tree.num_roots == 3
        assert list(sorted(tree.roots)) == [0, 1, 2]
        assert before.nodes.equals(tables.nodes[: len(before.nodes)])
        assert len(tables.edges) == 0

    @pytest.mark.parametrize("time", [1, 1.01, 1.5, 1.999])
    def test_t1_to_2(self, time):
        #
        # 2.00┊       ┊
        #     ┊       ┊
        # 1.00┊    3  ┊
        #     ┊   ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        tables = self.tables()
        before = tables.copy()
        tables.delete_older(time)
        ts = tables.tree_sequence()
        assert ts.num_trees == 1
        tree = ts.first()
        assert tree.num_roots == 2
        assert list(sorted(tree.roots)) == [0, 3]
        assert len(tables.nodes) == 5
        assert before.nodes.equals(tables.nodes)

    @pytest.mark.parametrize("time", [2, 2.5, 1e9])
    def test_t2(self, time):
        tables = self.tables()
        before = tables.copy()
        tables.delete_older(time)
        tables.assert_equals(before, ignore_provenance=True)


class TestDeleteOlderSimpleTreeMutationExamples:
    def test_single_mutation_no_time(self):
        # 2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊
        # 1.00┊ ┃  3  ┊
        #     ┊ x ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        tree = tskit.Tree.generate_balanced(3, branch_length=1)
        tables = tree.tree_sequence.dump_tables()
        tables.sites.add_row(0, "A")
        tables.mutations.add_row(site=0, node=0, derived_state="T", metadata=b"1234")

        tables.delete_older(1)
        # 2.00┊       ┊
        #     ┊       ┊
        # 1.00┊    3  ┊
        #     ┊ x ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        assert len(tables.nodes) == 5
        mut = tables.mutations[0]
        assert mut.node == 0
        assert mut.derived_state == "T"
        assert mut.metadata == b"1234"
        assert tskit.is_unknown_time(mut.time)

    def test_single_mutation_before_time(self):
        # 2.00┊   4   ┊
        #     ┊ x━┻┓  ┊
        # 1.00┊ ┃  3  ┊
        #     ┊ ┃ ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        tree = tskit.Tree.generate_balanced(3, branch_length=1)
        tables = tree.tree_sequence.dump_tables()
        tables.sites.add_row(0, "A")
        tables.mutations.add_row(
            site=0, node=0, time=1.5, derived_state="T", metadata=b"1234"
        )
        tables.delete_older(1)
        # 2.00┊       ┊
        #     ┊       ┊
        # 1.00┊    3  ┊
        #     ┊   ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        assert len(tables.nodes) == 5
        assert len(tables.mutations) == 0

    def test_single_mutation_at_time(self):
        # 2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊
        # 1.00┊ x  3  ┊
        #     ┊ ┃ ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        tree = tskit.Tree.generate_balanced(3, branch_length=1)
        tables = tree.tree_sequence.dump_tables()
        tables.sites.add_row(0, "A")
        tables.mutations.add_row(
            site=0, node=0, time=1, derived_state="T", metadata=b"1234"
        )

        tables.delete_older(1)
        # 2.00┊       ┊
        #     ┊       ┊
        # 1.00┊    3  ┊
        #     ┊   ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        assert len(tables.nodes) == 5
        assert len(tables.mutations) == 0

    def test_multi_mutation_no_time(self):
        # 2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊
        # 1.00┊ x  3  ┊
        #     ┊ x ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        tree = tskit.Tree.generate_balanced(3, branch_length=1)
        tables = tree.tree_sequence.dump_tables()
        tables.sites.add_row(0, "A")
        tables.mutations.add_row(site=0, node=0, derived_state="T")
        tables.mutations.add_row(site=0, node=0, parent=0, derived_state="G")
        before = tables.copy()

        tables.delete_older(1)
        # 2.00┊   4   ┊
        #     ┊       ┊
        #     ┊    3  ┊
        #     ┊ x  ┃  ┊
        #     ┊ x ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        tables.mutations.assert_equals(before.mutations)

    def test_multi_mutation_out_of_order(self):
        # 2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊
        # 1.00┊ x  3  ┊
        #     ┊ x ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        tree = tskit.Tree.generate_balanced(3, branch_length=1)
        tables = tree.tree_sequence.dump_tables()
        tables.sites.add_row(0, "A")
        tables.mutations.add_row(site=0, node=0, parent=1, derived_state="G")
        tables.mutations.add_row(site=0, node=0, derived_state="T")
        before = tables.copy()
        with pytest.raises(tskit.LibraryError, match="PARENT_AFTER_CHILD"):
            tables.tree_sequence()

        tables.delete_older(1)
        # 2.00┊   4   ┊
        #     ┊       ┊
        #     ┊    3  ┊
        #     ┊ x  ┃  ┊
        #     ┊ x ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        tables.mutations.assert_equals(before.mutations)

    def test_mutation_not_on_branch(self):
        tables = tskit.TableCollection(1)
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        tables.sites.add_row(0, "A")
        tables.mutations.add_row(site=0, node=0, derived_state="T")
        before = tables.copy()
        tables.delete_older(0.01)
        tables.assert_equals(before, ignore_provenance=True)


class TestDeleteOlderSimpleTreeMigrationExamples:
    @tests.cached_example
    def ts(self):
        # 2.00┊   4   ┊
        #     ┊ o━┻┓  ┊
        # 1.00┊ o  3  ┊
        #     ┊ o ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        tree = tskit.Tree.generate_balanced(3, branch_length=1)
        tables = tree.tree_sequence.dump_tables()
        tables.populations.add_row()
        tables.populations.add_row()
        tables.migrations.add_row(source=0, dest=1, node=0, time=0.5, left=0, right=1)
        tables.migrations.add_row(source=1, dest=0, node=0, time=1.0, left=0, right=1)
        tables.migrations.add_row(source=0, dest=1, node=0, time=1.5, left=0, right=1)
        tables.compute_mutation_parents()
        ts = tables.tree_sequence()
        return ts

    def test_t099(self):
        tables = self.ts().dump_tables()
        tables.delete_older(0.99)
        assert len(tables.migrations) == 1
        assert tables.migrations[0].time == 0.5

    def test_t1(self):
        tables = self.ts().dump_tables()
        tables.delete_older(1)
        assert len(tables.migrations) == 1
        assert tables.migrations[0].time == 0.5

    @pytest.mark.parametrize("time", [1.51, 2.0, 2.5])
    def test_older(self, time):
        tables = self.ts().dump_tables()
        before = tables.copy()
        tables.delete_older(time)
        tables.migrations.assert_equals(before.migrations)


def split_edges_definition(ts, time, *, flags=0, population=None, metadata=None):
    population = -1 if population is None else population
    tables = ts.dump_tables()
    if ts.num_migrations > 0:
        raise ValueError("Migrations not supported")

    node_time = tables.nodes.time
    tables.edges.clear()
    split_edge = np.full(ts.num_edges, tskit.NULL, dtype=int)
    for edge in ts.edges():
        if node_time[edge.child] < time < node_time[edge.parent]:
            u = tables.nodes.add_row(
                flags=flags, time=time, population=population, metadata=metadata
            )
            tables.edges.append(edge.replace(parent=u))
            tables.edges.append(edge.replace(child=u))
            split_edge[edge.id] = u
        else:
            tables.edges.append(edge)

    tables.mutations.clear()
    for mutation in ts.mutations():
        mapped_node = tskit.NULL
        if mutation.edge != tskit.NULL:
            mapped_node = split_edge[mutation.edge]
        if mapped_node != tskit.NULL and mutation.time >= time:
            mutation = mutation.replace(node=mapped_node)
        tables.mutations.append(mutation)

    tables.sort()
    return tables.tree_sequence()


class TestSplitEdgesSimpleTree:
    # 2.00┊   4   ┊
    #     ┊ ┏━┻┓  ┊
    # 1.00┊ ┃  3  ┊
    #     ┊ ┃ ┏┻┓ ┊
    # 0.00┊ 0 1 2 ┊
    #     0       1
    @tests.cached_example
    def ts(self):
        return tskit.Tree.generate_balanced(3, branch_length=1).tree_sequence

    @pytest.mark.parametrize("time", [0.1, 0.5, 0.9])
    def test_lowest_branches(self, time):
        # 2.00┊   4   ┊    2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊        ┊ ┏━┻┓  ┊
        # 1.00┊ ┃  3  ┊    1.00┊ ┃  3  ┊
        #     ┊ ┃ ┏┻┓ ┊        ┊ ┃ ┏┻┓ ┊
        #     ┊ ┃ ┃ ┃ ┊      t ┊ 7 5 6 ┊
        #     ┊ ┃ ┃ ┃ ┊ ->     ┊ ┃ ┃ ┃ ┊
        # 0.00┊ 0 1 2 ┊    0.00┊ 0 1 2 ┊
        #     0       1        0       1
        before_ts = self.ts()
        ts = before_ts.split_edges(time)
        assert ts.num_nodes == 8
        assert all(ts.node(u).time == time for u in [5, 6, 7])
        assert ts.num_trees == 1
        assert ts.first().parent_dict == {0: 7, 1: 5, 2: 6, 5: 3, 6: 3, 7: 4, 3: 4}
        ts = ts.simplify()
        ts.tables.assert_equals(before_ts.tables, ignore_provenance=True)

    def test_same_time_as_node(self):
        # 2.00┊   4   ┊    2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊        ┊ ┏━┻┓  ┊
        # 1.00┊ ┃  3  ┊    1.00┊ 5  3  ┊
        #     ┊ ┃ ┏┻┓ ┊        ┊ ┃ ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊    0.00┊ 0 1 2 ┊
        #     0       1        0       1
        before_ts = self.ts()
        ts = before_ts.split_edges(1)
        assert ts.num_nodes == 6
        assert ts.node(5).time == 1
        assert ts.num_trees == 1
        assert ts.first().parent_dict == {0: 5, 1: 3, 2: 3, 5: 4, 3: 4}
        ts = ts.simplify()
        ts.tables.assert_equals(before_ts.tables, ignore_provenance=True)

    @pytest.mark.parametrize("time", [1.1, 1.5, 1.9])
    def test_top_branches(self, time):
        # 2.00┊   4   ┊    2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊        ┊ ┏━┻┓  ┊
        #     ┊ ┃  ┃  ┊      t ┊ 5  6  ┊
        #     ┊ ┃  ┃  ┊ ->     ┊ ┃  ┃  ┊
        # 1.00┊ ┃  3  ┊    1.00┊ ┃  3  ┊
        #     ┊ ┃ ┏┻┓ ┊        ┊ ┃ ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊    0.00┊ 0 1 2 ┊
        #     0       1        0       1

        before_ts = self.ts()
        ts = before_ts.split_edges(time)
        assert ts.num_nodes == 7
        assert all(ts.node(u).time == time for u in [5, 6])
        assert ts.num_trees == 1
        assert ts.first().parent_dict == {0: 5, 1: 3, 2: 3, 3: 6, 6: 4, 5: 4}
        ts = ts.simplify()
        ts.tables.assert_equals(before_ts.tables, ignore_provenance=True)

    @pytest.mark.parametrize("time", [0, 2])
    def test_at_leaf_or_root_time(self, time):
        split = self.ts().split_edges(time)
        split.tables.assert_equals(self.ts().tables, ignore_provenance=True)

    @pytest.mark.parametrize("time", [-1, 2.1])
    def test_outside_time_scales(self, time):
        split = self.ts().split_edges(time)
        split.tables.assert_equals(self.ts().tables, ignore_provenance=True)


class TestSplitEdgesSimpleTreeMutationExamples:
    def test_single_mutation_no_time(self):
        # 2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊
        # 1.00┊ ┃  3  ┊
        #     ┊ x ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        tree = tskit.Tree.generate_balanced(3, branch_length=1)
        tables = tree.tree_sequence.dump_tables()
        tables.sites.add_row(0, "A")
        tables.mutations.add_row(site=0, node=0, derived_state="T", metadata=b"1234")
        ts = tables.tree_sequence()

        ts_split = ts.split_edges(1)
        # 2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊
        # 1.00┊ 5  3  ┊
        #     ┊ x ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        assert ts_split.num_nodes == 6
        mut = ts_split.mutation(0)
        assert mut.node == 0
        assert mut.derived_state == "T"
        assert mut.metadata == b"1234"
        assert tskit.is_unknown_time(mut.time)

    def test_single_mutation_split_before_time(self):
        # 2.00┊   4   ┊
        #     ┊ x━┻┓  ┊
        # 1.00┊ ┃  3  ┊
        #     ┊ ┃ ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        tree = tskit.Tree.generate_balanced(3, branch_length=1)
        tables = tree.tree_sequence.dump_tables()
        tables.sites.add_row(0, "A")
        tables.mutations.add_row(
            site=0, node=0, time=1.5, derived_state="T", metadata=b"1234"
        )
        ts = tables.tree_sequence()

        ts_split = ts.split_edges(1)
        # 2.00┊   4   ┊
        #     ┊ x━┻┓  ┊
        # 1.00┊ 5  3  ┊
        #     ┊ ┃ ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        assert ts_split.num_nodes == 6
        mut = ts_split.mutation(0)
        assert mut.node == 5
        assert mut.derived_state == "T"
        assert mut.metadata == b"1234"
        assert mut.time == 1.5

    def test_single_mutation_split_at_time(self):
        # 2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊
        # 1.00┊ x  3  ┊
        #     ┊ ┃ ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        tree = tskit.Tree.generate_balanced(3, branch_length=1)
        tables = tree.tree_sequence.dump_tables()
        tables.sites.add_row(0, "A")
        tables.mutations.add_row(
            site=0, node=0, time=1, derived_state="T", metadata=b"1234"
        )
        ts = tables.tree_sequence()

        ts_split = ts.split_edges(1)
        # 2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊
        # 1.00┊ 5x 3  ┊
        #     ┊ ┃ ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        mut = ts_split.mutation(0)
        assert mut.node == 5
        assert mut.derived_state == "T"
        assert mut.metadata == b"1234"
        assert mut.time == 1.0

    def test_multi_mutation_no_time(self):
        # 2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊
        # 1.00┊ x  3  ┊
        #     ┊ x ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        tree = tskit.Tree.generate_balanced(3, branch_length=1)
        tables = tree.tree_sequence.dump_tables()
        tables.sites.add_row(0, "A")
        tables.mutations.add_row(site=0, node=0, derived_state="T")
        tables.mutations.add_row(site=0, node=0, parent=0, derived_state="G")
        ts = tables.tree_sequence()

        ts_split = ts.split_edges(1)
        # 2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊
        #     ┊ 5  3  ┊
        #     ┊ x  ┃  ┊
        #     ┊ x ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        ts_split.tables.mutations.assert_equals(tables.mutations)

    def test_multi_mutation_over_sample_time(self):
        # 2.00┊   4   ┊
        #     ┊ x━┻┓  ┊
        # 1.00┊ ┃  3  ┊
        #     ┊ x ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        tree = tskit.Tree.generate_balanced(3, branch_length=1)
        tables = tree.tree_sequence.dump_tables()
        tables.sites.add_row(0, "A")
        tables.mutations.add_row(site=0, node=0, time=1.01, derived_state="T")
        tables.mutations.add_row(site=0, node=0, time=0.99, parent=0, derived_state="G")
        ts = tables.tree_sequence()

        ts_split = ts.split_edges(1)
        # 2.00┊   4   ┊
        #     ┊ x━┻┓  ┊
        # 1.00┊ 5  3  ┊
        #     ┊ x ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        assert ts_split.num_mutations == 2

        mut = ts_split.mutation(0)
        assert mut.site == 0
        assert mut.node == 5
        assert mut.time == 1.01
        mut = ts_split.mutation(1)
        assert mut.site == 0
        assert mut.node == 0
        assert mut.time == 0.99

    def test_mutation_not_on_branch(self):
        tables = tskit.TableCollection(1)
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        tables.sites.add_row(0, "A")
        tables.mutations.add_row(site=0, node=0, derived_state="T")
        ts = tables.tree_sequence()
        tables.assert_equals(ts.split_edges(0).tables, ignore_provenance=True)


class TestSplitEdgesExamples:
    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_genotypes_round_trip(self, ts):
        time = 0 if ts.num_nodes == 0 else np.median(ts.tables.nodes.time)
        if ts.num_migrations == 0:
            split_ts = ts.split_edges(time)
            assert np.array_equal(split_ts.genotype_matrix(), ts.genotype_matrix())
        else:
            with pytest.raises(tskit.LibraryError):
                ts.split_edges(time)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    @pytest.mark.parametrize("population", [-1, None])
    def test_definition(self, ts, population):
        time = 0 if ts.num_nodes == 0 else np.median(ts.tables.nodes.time)
        if ts.num_migrations == 0:
            ts1 = split_edges_definition(ts, time, population=population)
            ts2 = ts.split_edges(time, population=population)
            ts1.tables.assert_equals(ts2.tables, ignore_provenance=True)


class TestSplitEdgesInterface:
    def test_migrations_fail(self, ts_fixture):
        assert ts_fixture.num_migrations > 0
        with pytest.raises(tskit.LibraryError, match="MIGRATIONS_NOT_SUPPORTED"):
            ts_fixture.split_edges(0)

    def test_population_out_of_bounds(self):
        tables = tskit.TableCollection(1)
        ts = tables.tree_sequence()
        with pytest.raises(tskit.LibraryError, match="POPULATION_OUT_OF_BOUNDS"):
            ts.split_edges(0, population=0)

    def test_bad_flags(self):
        ts = tskit.TableCollection(1).tree_sequence()
        with pytest.raises(TypeError):
            ts.split_edges(0, flags="asdf")

    def test_bad_metadata_no_schema(self):
        ts = tskit.TableCollection(1).tree_sequence()
        with pytest.raises(TypeError):
            ts.split_edges(0, metadata="asdf")

    def test_bad_metadata_json_schema(self):
        tables = tskit.TableCollection(1)
        tables.nodes.metadata_schema = tskit.MetadataSchema.permissive_json()
        ts = tables.tree_sequence()
        with pytest.raises(tskit.MetadataEncodingError):
            ts.split_edges(0, metadata=b"bytes")

    @pytest.mark.parametrize("time", [math.inf, np.inf, tskit.UNKNOWN_TIME, np.nan])
    def test_nonfinite_time(self, time):
        tables = tskit.TableCollection(1)
        ts = tables.tree_sequence()
        with pytest.raises(tskit.LibraryError, match="TIME_NONFINITE"):
            ts.split_edges(time)


class TestSplitEdgesNodeValues:
    @tests.cached_example
    def ts(self):
        tables = tskit.TableCollection(1)
        for _ in range(5):
            tables.populations.add_row()
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, population=0, time=0)
        tables.nodes.add_row(time=1)
        tables.edges.add_row(0, 1, 1, 0)
        return tables.tree_sequence()

    @tests.cached_example
    def ts_with_schema(self):
        tables = tskit.TableCollection(1)
        for _ in range(5):
            tables.populations.add_row()
        tables.nodes.metadata_schema = tskit.MetadataSchema.permissive_json()
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, population=0, time=0)
        tables.nodes.add_row(time=1)
        tables.edges.add_row(0, 1, 1, 0)
        return tables.tree_sequence()

    def test_default_population(self):
        ts = self.ts().split_edges(0.5)
        assert ts.node(2).population == -1

    @pytest.mark.parametrize("population", range(-1, 5))
    def test_specify_population(self, population):
        ts = self.ts().split_edges(0.5, population=population)
        assert ts.node(2).population == population

    def test_default_flags(self):
        ts = self.ts().split_edges(0.5)
        assert ts.node(2).flags == 0

    @pytest.mark.parametrize("flags", range(0, 5))
    def test_specify_flags(self, flags):
        ts = self.ts().split_edges(0.5, flags=flags)
        assert ts.node(2).flags == flags

    def test_default_metadata_no_schema(self):
        ts = self.ts().split_edges(0.5)
        assert ts.node(2).metadata == b""

    @pytest.mark.parametrize("metadata", [b"", b"some bytes"])
    def test_specify_metadata_no_schema(self, metadata):
        ts = self.ts().split_edges(0.5, metadata=metadata)
        assert ts.node(2).metadata == metadata

    def test_default_metadata_with_schema(self):
        ts = self.ts_with_schema().split_edges(0.5)
        assert ts.node(2).metadata == {}

    @pytest.mark.parametrize("metadata", [{}, {"some": "json"}])
    def test_specify_metadata_with_schema(self, metadata):
        ts = self.ts_with_schema().split_edges(0.5, metadata=metadata)
        assert ts.node(2).metadata == metadata


def decapitate_definition(ts, time, *, flags=0, population=None, metadata=None):
    """
    Simple loop implementation of the decapitate operation
    """
    population = -1 if population is None else population
    tables = ts.dump_tables()
    node_time = tables.nodes.time
    tables.edges.clear()
    for edge in ts.edges():
        if node_time[edge.parent] <= time:
            tables.edges.append(edge)
        elif node_time[edge.child] < time:
            new_parent = tables.nodes.add_row(
                time=time, population=population, flags=flags, metadata=metadata
            )
            tables.edges.append(edge.replace(parent=new_parent))

    tables.mutations.clear()
    for mutation in ts.mutations():
        mutation_time = (
            node_time[mutation.node]
            if util.is_unknown_time(mutation.time)
            else mutation.time
        )
        if mutation_time < time:
            tables.mutations.append(mutation.replace(parent=tskit.NULL))

    tables.migrations.clear()
    for migration in ts.migrations():
        if migration.time <= time:
            tables.migrations.append(migration)

    tables.build_index()
    tables.compute_mutation_parents()
    return tables.tree_sequence()


class TestDecapitateExamples:
    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_defaults(self, ts):
        time = 0 if ts.num_nodes == 0 else np.median(ts.tables.nodes.time)
        if ts.num_migrations == 0:
            decap1 = decapitate_definition(ts, time)
            decap2 = ts.decapitate(time)
            decap1.tables.assert_equals(decap2.tables, ignore_provenance=True)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_no_population(self, ts):
        time = 0 if ts.num_nodes == 0 else np.median(ts.tables.nodes.time)
        if ts.num_migrations == 0:
            decap1 = decapitate_definition(ts, time, population=-1)
            decap2 = ts.decapitate(time, population=-1)
            decap1.tables.assert_equals(decap2.tables, ignore_provenance=True)


class TestDecapitateSimpleTree:
    # 2.00┊   4   ┊
    #     ┊ ┏━┻┓  ┊
    # 1.00┊ ┃  3  ┊
    #     ┊ ┃ ┏┻┓ ┊
    # 0.00┊ 0 1 2 ┊
    #     0       1
    @tests.cached_example
    def ts(self):
        tree = tskit.Tree.generate_balanced(3, branch_length=1)
        return tree.tree_sequence

    @pytest.mark.parametrize("time", [0, -0.5, -100])
    def test_t0_or_before(self, time):
        before = self.ts()
        ts = before.decapitate(time)
        assert ts.num_trees == 1
        tree = ts.first()
        assert tree.num_roots == 3
        assert list(sorted(tree.roots)) == [0, 1, 2]
        assert before.tables.nodes.equals(ts.tables.nodes[: before.num_nodes])
        assert ts.num_edges == 0

    @pytest.mark.parametrize("time", [0.01, 0.5, 0.999])
    def test_t0_to_1(self, time):
        #
        # 2.00┊       ┊
        #     ┊       ┊
        # 0.99┊ 7 5 6 ┊
        #     ┊ ┃ ┃ ┃ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        before = self.ts()
        ts = before.decapitate(time)
        assert ts.num_trees == 1
        tree = ts.first()
        assert tree.num_roots == 3
        assert list(sorted(tree.roots)) == [5, 6, 7]
        assert ts.num_nodes == 8
        assert ts.tables.nodes[5].time == time
        assert ts.tables.nodes[6].time == time
        assert ts.tables.nodes[7].time == time

    def test_t1(self):
        #
        # 2.00┊       ┊
        #     ┊       ┊
        # 1.00┊ 5  3  ┊
        #     ┊ ┃ ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        before = self.ts()
        ts = before.decapitate(1)
        assert ts.num_trees == 1
        tree = ts.first()
        assert tree.num_roots == 2
        assert list(sorted(tree.roots)) == [3, 5]
        assert ts.num_nodes == 6
        assert ts.tables.nodes[5].time == 1

    @pytest.mark.parametrize("time", [1.01, 1.5, 1.999])
    def test_t1_to_2(self, time):
        # 2.00┊       ┊
        #     ┊       ┊
        # 1.01┊ 5  6  ┊
        #     ┊ ┃  ┃  ┊
        # 1.00┊ ┃  3  ┊
        #     ┊ ┃ ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #    0       1
        before = self.ts()
        ts = before.decapitate(time)
        assert ts.num_trees == 1
        tree = ts.first()
        assert tree.num_roots == 2
        assert list(sorted(tree.roots)) == [5, 6]
        assert ts.num_nodes == 7
        assert ts.tables.nodes[5].time == time
        assert ts.tables.nodes[6].time == time

    @pytest.mark.parametrize("time", [2, 2.5, 1e9])
    def test_t2(self, time):
        before = self.ts()
        ts = before.decapitate(time)
        ts.tables.assert_equals(before.tables, ignore_provenance=True)


class TestDecapitateSimpleTreeMutationExamples:
    def test_single_mutation_over_sample(self):
        # 2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊
        # 1.00┊ ┃  3  ┊
        #     ┊ x ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        tree = tskit.Tree.generate_balanced(3, branch_length=1)
        tables = tree.tree_sequence.dump_tables()
        tables.sites.add_row(0, "A")
        tables.mutations.add_row(site=0, node=0, derived_state="T")
        before = tables.tree_sequence()

        ts = before.decapitate(1)
        # 2.00┊       ┊
        #     ┊       ┊
        # 1.00┊ 5  3  ┊
        #     ┊ x ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        before.tables.mutations.assert_equals(ts.tables.mutations)
        assert list(before.alignments()) == list(ts.alignments())

    def test_single_mutation_at_decap_time(self):
        # 2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊
        # 1.00┊ x  3  ┊
        #     ┊ ┃ ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        tree = tskit.Tree.generate_balanced(3, branch_length=1)
        tables = tree.tree_sequence.dump_tables()
        tables.sites.add_row(0, "A")
        tables.mutations.add_row(site=0, node=0, time=1, derived_state="T")
        before = tables.tree_sequence()

        # Because the mutation is at exactly the decapitation time, we must
        # remove it, or it would violate the requirement that a mutation must
        # have a time less than that of the parent of the edge that its on.
        ts = before.decapitate(1)
        # 2.00┊       ┊
        #     ┊       ┊
        # 1.00┊ 5  3  ┊
        #     ┊ ┃ ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        assert ts.num_mutations == 0
        assert list(ts.alignments()) == ["A", "A", "A"]

    def test_multi_mutation_over_sample(self):
        # 2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊
        # 1.00┊ x  3  ┊
        #     ┊ x ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        tree = tskit.Tree.generate_balanced(3, branch_length=1)
        tables = tree.tree_sequence.dump_tables()
        tables.sites.add_row(0, "A")
        tables.mutations.add_row(site=0, node=0, derived_state="T")
        tables.mutations.add_row(site=0, node=0, parent=0, derived_state="G")
        before = tables.tree_sequence()

        ts = before.decapitate(1)
        # 2.00┊       ┊
        #     ┊ 5  3  ┊
        #     ┊ x  ┃  ┊
        #     ┊ x ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        before.tables.mutations.assert_equals(ts.tables.mutations)
        assert list(before.alignments()) == list(ts.alignments())

    def test_multi_mutation_over_sample_time(self):
        # 2.00┊   4   ┊
        #     ┊ x━┻┓  ┊
        # 1.00┊ ┃  3  ┊
        #     ┊ x ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        tree = tskit.Tree.generate_balanced(3, branch_length=1)
        tables = tree.tree_sequence.dump_tables()
        tables.sites.add_row(0, "A")
        tables.mutations.add_row(site=0, node=0, time=1.01, derived_state="T")
        tables.mutations.add_row(site=0, node=0, time=0.99, parent=0, derived_state="G")
        before = tables.tree_sequence()

        ts = before.decapitate(1)
        # 2.00┊       ┊
        #     ┊ 5  3  ┊
        #     ┊ ┃  ┃  ┊
        #     ┊ x ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        assert ts.num_mutations == 1
        # Alignments are equal because the ancestral mutation was silent anyway.
        assert list(before.alignments()) == list(ts.alignments())

    def test_multi_mutation_over_root(self):
        #         x
        # 2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊
        # 1.00┊ ┃  3  ┊
        #     ┊ x ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        tree = tskit.Tree.generate_balanced(3, branch_length=1)
        tables = tree.tree_sequence.dump_tables()
        tables.sites.add_row(0, "A")
        tables.mutations.add_row(site=0, node=4, derived_state="G")
        tables.mutations.add_row(site=0, node=0, parent=0, derived_state="T")
        before = tables.tree_sequence()

        ts = before.decapitate(1)
        # 2.00┊       ┊
        #     ┊ 5  3  ┊
        #     ┊ ┃  ┃  ┊
        #     ┊ x ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        assert ts.num_mutations == 1
        assert list(before.alignments()) == ["T", "G", "G"]
        # The states inherited by samples changes because we drop the old mutation
        assert list(ts.alignments()) == ["T", "A", "A"]


class TestDecapitateSimpleTsExample:
    # 9.08┊    9    ┊         ┊         ┊         ┊         ┊
    #     ┊  ┏━┻━┓  ┊         ┊         ┊         ┊         ┊
    # 6.57┊  ┃   ┃  ┊         ┊         ┊         ┊    8    ┊
    #     ┊  ┃   ┃  ┊         ┊         ┊         ┊  ┏━┻━┓  ┊
    # 5.31┊  ┃   ┃  ┊    7    ┊         ┊    7    ┊  ┃   ┃  ┊
    #     ┊  ┃   ┃  ┊  ┏━┻━┓  ┊         ┊  ┏━┻━┓  ┊  ┃   ┃  ┊
    # 1.75┊  ┃   ┃  ┊  ┃   ┃  ┊    6    ┊  ┃   ┃  ┊  ┃   ┃  ┊
    #     ┊  ┃   ┃  ┊  ┃   ┃  ┊  ┏━┻━┓  ┊  ┃   ┃  ┊  ┃   ┃  ┊
    # 1.11┊  ┃   5  ┊  ┃   5  ┊  ┃   5  ┊  ┃   5  ┊  ┃   5  ┊
    #     ┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊
    # 0.11┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊
    #     ┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊
    # 0.00┊ 0 1 2 3 ┊ 0 1 2 3 ┊ 0 1 2 3 ┊ 0 1 2 3 ┊ 0 1 2 3 ┊
    #   0.00      0.06      0.79      0.91      0.91      1.00

    @tests.cached_example
    def ts(self):
        nodes = io.StringIO(
            """\
        id      is_sample   population      individual      time    metadata
        0       1       0       -1      0
        1       1       0       -1      0
        2       1       0       -1      0
        3       1       0       -1      0
        4       0       0       -1      0.114
        5       0       0       -1      1.110
        6       0       0       -1      1.750
        7       0       0       -1      5.310
        8       0       0       -1      6.573
        9       0       0       -1      9.083
        """
        )
        edges = io.StringIO(
            """\
        id      left            right           parent  child
        0       0.00000000      1.00000000      4       0
        1       0.00000000      1.00000000      4       1
        2       0.00000000      1.00000000      5       2
        3       0.00000000      1.00000000      5       3
        4       0.79258618      0.90634460      6       4
        5       0.79258618      0.90634460      6       5
        6       0.05975243      0.79258618      7       4
        7       0.90634460      0.91029435      7       4
        8       0.05975243      0.79258618      7       5
        9       0.90634460      0.91029435      7       5
        10      0.91029435      1.00000000      8       4
        11      0.91029435      1.00000000      8       5
        12      0.00000000      0.05975243      9       4
        13      0.00000000      0.05975243      9       5
        """
        )
        sites = io.StringIO(
            """\
        position      ancestral_state
        0.05          A
        0.06          0
        0.3           C
        0.5           AAA
        0.91          T
        """
        )
        muts = io.StringIO(
            """\
        site   node    derived_state    parent    time
        0      9       T                -1        15
        0      9       GGG              0         9.1
        0      5       1                1         9
        1      4       C                -1        1.6
        1      4       G                3         1.5
        2      7       G                -1        10
        2      3       C                5         1
        4      3       G                -1        1
        """
        )
        ts = tskit.load_text(nodes, edges, sites=sites, mutations=muts, strict=False)
        return ts

    def test_at_time_of_5(self):
        # NOTE: we don't remember that the edge 4-7 was shared in trees 1 and 3.
        # 1.11┊  14  5  ┊ 11   5  ┊ 10   5  ┊  12  5  ┊  13  5  ┊
        #     ┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊
        # 0.11┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊
        #     ┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊
        # 0.00┊ 0 1 2 3 ┊ 0 1 2 3 ┊ 0 1 2 3 ┊ 0 1 2 3 ┊ 0 1 2 3 ┊
        #   0.00      0.06      0.79      0.91      0.91      1.00
        ts = self.ts().decapitate(1.110)
        assert ts.num_nodes == 15
        assert ts.num_trees == 5
        # Most mutations are older than this.
        assert ts.num_mutations == 2
        for u in range(10, 15):
            node = ts.node(u)
            assert node.time == 1.110
            assert node.flags == 0
        assert [set(tree.roots) for tree in ts.trees()] == [
            {5, 14},
            {11, 5},
            {10, 5},
            {12, 5},
            {13, 5},
        ]

    def test_at_time6(self):
        # 6   ┊ 12   13 ┊         ┊         ┊         ┊ 10   11 ┊
        # 5.31┊  ┃   ┃  ┊    7    ┊         ┊    7    ┊  ┃   ┃  ┊
        #     ┊  ┃   ┃  ┊  ┏━┻━┓  ┊         ┊  ┏━┻━┓  ┊  ┃   ┃  ┊
        # 1.75┊  ┃   ┃  ┊  ┃   ┃  ┊    6    ┊  ┃   ┃  ┊  ┃   ┃  ┊
        #     ┊  ┃   ┃  ┊  ┃   ┃  ┊  ┏━┻━┓  ┊  ┃   ┃  ┊  ┃   ┃  ┊
        # 1.11┊  ┃   5  ┊  ┃   5  ┊  ┃   5  ┊  ┃   5  ┊  ┃   5  ┊
        #     ┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊
        # 0.11┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊
        #     ┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊
        # 0.00┊ 0 1 2 3 ┊ 0 1 2 3 ┊ 0 1 2 3 ┊ 0 1 2 3 ┊ 0 1 2 3 ┊
        #   0.00      0.06      0.79      0.91      0.91      1.00
        ts = self.ts().decapitate(6)
        assert ts.num_nodes == 14
        assert ts.num_trees == 5
        assert ts.num_mutations == 4
        for u in range(10, 14):
            node = ts.node(u)
            assert node.time == 6
            assert node.flags == 0
        assert [set(tree.roots) for tree in ts.trees()] == [
            {12, 13},
            {7},
            {6},
            {7},
            {10, 11},
        ]


class TestDecapitateNodeValues:
    @tests.cached_example
    def ts(self):
        tables = tskit.TableCollection(1)
        for _ in range(5):
            tables.populations.add_row()
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, population=0, time=0)
        tables.nodes.add_row(time=1)
        tables.edges.add_row(0, 1, 1, 0)
        return tables.tree_sequence()

    @tests.cached_example
    def ts_with_schema(self):
        tables = tskit.TableCollection(1)
        for _ in range(5):
            tables.populations.add_row()
        tables.nodes.metadata_schema = tskit.MetadataSchema.permissive_json()
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, population=0, time=0)
        tables.nodes.add_row(time=1)
        tables.edges.add_row(0, 1, 1, 0)
        return tables.tree_sequence()

    def test_default_population(self):
        ts = self.ts().decapitate(0.5)
        assert ts.node(2).population == tskit.NULL

    @pytest.mark.parametrize("population", range(-1, 5))
    def test_specify_population(self, population):
        ts = self.ts().decapitate(0.5, population=population)
        assert ts.node(2).population == population

    def test_default_flags(self):
        ts = self.ts().decapitate(0.5)
        assert ts.node(2).flags == 0

    @pytest.mark.parametrize("flags", range(0, 5))
    def test_specify_flags(self, flags):
        ts = self.ts().decapitate(0.5, flags=flags)
        assert ts.node(2).flags == flags

    def test_default_metadata_no_schema(self):
        ts = self.ts().decapitate(0.5)
        assert ts.node(2).metadata == b""

    @pytest.mark.parametrize("metadata", [b"", b"some bytes"])
    def test_specify_metadata_no_schema(self, metadata):
        ts = self.ts().decapitate(0.5, metadata=metadata)
        assert ts.node(2).metadata == metadata

    def test_default_metadata_with_schema(self):
        ts = self.ts_with_schema().decapitate(0.5)
        assert ts.node(2).metadata == {}

    @pytest.mark.parametrize("metadata", [{}, {"some": "json"}])
    def test_specify_metadata_with_schema(self, metadata):
        ts = self.ts_with_schema().decapitate(0.5, metadata=metadata)
        assert ts.node(2).metadata == metadata


class TestDecapitateInterface:
    @tests.cached_example
    def ts(self):
        tree = tskit.Tree.generate_balanced(3, branch_length=1)
        return tree.tree_sequence

    @pytest.mark.parametrize("bad_type", ["x", "0.1", [], [0.1]])
    def test_bad_types(self, bad_type):
        with pytest.raises(TypeError, match="number"):
            self.ts().decapitate(bad_type)

    @pytest.mark.parametrize(
        "time", [1, 1.0, np.array([1])[0], fractions.Fraction(1, 1), decimal.Decimal(1)]
    )
    def test_number_types(self, time):
        expected = self.ts().decapitate(1)
        got = self.ts().decapitate(time)
        expected.tables.assert_equals(got.tables, ignore_timestamps=True)

    def test_migrations_not_supported(self, ts_fixture):
        with pytest.raises(tskit.LibraryError, match="MIGRATIONS_NOT_SUPPORTED"):
            ts_fixture.decapitate(0)

    def test_population_out_of_bounds(self):
        tables = tskit.TableCollection(1)
        ts = tables.tree_sequence()
        with pytest.raises(tskit.LibraryError, match="POPULATION_OUT_OF_BOUNDS"):
            ts.decapitate(0, population=0)

    def test_bad_flags(self):
        ts = tskit.TableCollection(1).tree_sequence()
        with pytest.raises(TypeError):
            ts.decapitate(0, flags="asdf")

    def test_bad_metadata_no_schema(self):
        ts = tskit.TableCollection(1).tree_sequence()
        with pytest.raises(TypeError):
            ts.decapitate(0, metadata="asdf")

    def test_bad_metadata_json_schema(self):
        tables = tskit.TableCollection(1)
        tables.nodes.metadata_schema = tskit.MetadataSchema.permissive_json()
        ts = tables.tree_sequence()
        with pytest.raises(tskit.MetadataEncodingError):
            ts.decapitate(0, metadata=b"bytes")

    @pytest.mark.parametrize("time", [math.inf, np.inf, tskit.UNKNOWN_TIME, np.nan])
    def test_nonfinite_time(self, time):
        tables = tskit.TableCollection(1)
        ts = tables.tree_sequence()
        with pytest.raises(tskit.LibraryError, match="TIME_NONFINITE"):
            ts.decapitate(time)


--- ../../tskit/python/tests/test_divmat.py ---


"""
Test cases for divergence matrix based pairwise stats
"""
import array
import collections
import functools

import msprime
import numpy as np
import pytest

import tskit
from tests import tsutil
from tests.test_highlevel import get_example_tree_sequences

# ↑ See https://github.com/tskit-dev/tskit/issues/1804 for when
# we can remove this.

DIVMAT_MODES = ["branch", "site"]

# NOTE: this implementation of Schieber-Vishkin algorithm is done like
# this so it's easy to run with numba. It would be more naturally
# packaged as a class. We don't actually use numba here, but it's
# handy to have a version of the SV code lying around that can be
# run directly with numba.


def sv_tables_init(parent_array):
    n = 1 + parent_array.shape[0]

    LAMBDA = 0
    # Triply-linked tree. FIXME we shouldn't need to build this as it's
    # available already in tskit
    child = np.zeros(n, dtype=np.int32)
    parent = np.zeros(n, dtype=np.int32)
    sib = np.zeros(n, dtype=np.int32)

    for j in range(n - 1):
        u = j + 1
        v = parent_array[j] + 1
        sib[u] = child[v]
        child[v] = u
        parent[u] = v

    lambd = np.zeros(n, dtype=np.int32)
    pi = np.zeros(n, dtype=np.int32)
    tau = np.zeros(n, dtype=np.int32)
    beta = np.zeros(n, dtype=np.int32)
    alpha = np.zeros(n, dtype=np.int32)

    p = child[LAMBDA]
    n = 0
    lambd[0] = -1
    while p != LAMBDA:
        while True:
            n += 1
            pi[p] = n
            tau[n] = LAMBDA
            lambd[n] = 1 + lambd[n >> 1]
            if child[p] != LAMBDA:
                p = child[p]
            else:
                break
        beta[p] = n
        while True:
            tau[beta[p]] = parent[p]
            if sib[p] != LAMBDA:
                p = sib[p]
                break
            else:
                p = parent[p]
                if p != LAMBDA:
                    h = lambd[n & -pi[p]]
                    beta[p] = ((n >> h) | 1) << h
                else:
                    break

    # Begin the second traversal
    lambd[0] = lambd[n]
    pi[LAMBDA] = 0
    beta[LAMBDA] = 0
    alpha[LAMBDA] = 0
    p = child[LAMBDA]
    while p != LAMBDA:
        while True:
            a = alpha[parent[p]] | (beta[p] & -beta[p])
            alpha[p] = a
            if child[p] != LAMBDA:
                p = child[p]
            else:
                break
        while True:
            if sib[p] != LAMBDA:
                p = sib[p]
                break
            else:
                p = parent[p]
                if p == LAMBDA:
                    break

    return lambd, pi, tau, beta, alpha


def _sv_mrca(x, y, lambd, pi, tau, beta, alpha):
    if beta[x] <= beta[y]:
        h = lambd[beta[y] & -beta[x]]
    else:
        h = lambd[beta[x] & -beta[y]]
    k = alpha[x] & alpha[y] & -(1 << h)
    h = lambd[k & -k]
    j = ((beta[x] >> h) | 1) << h
    if j == beta[x]:
        xhat = x
    else:
        ell = lambd[alpha[x] & ((1 << h) - 1)]
        xhat = tau[((beta[x] >> ell) | 1) << ell]
    if j == beta[y]:
        yhat = y
    else:
        ell = lambd[alpha[y] & ((1 << h) - 1)]
        yhat = tau[((beta[y] >> ell) | 1) << ell]
    if pi[xhat] <= pi[yhat]:
        z = xhat
    else:
        z = yhat
    return z


def sv_mrca(x, y, lambd, pi, tau, beta, alpha):
    # Convert to 1-based indexes
    return _sv_mrca(x + 1, y + 1, lambd, pi, tau, beta, alpha) - 1


def local_root(tree, u):
    while tree.parent(u) != tskit.NULL:
        u = tree.parent(u)
    return u


def span_normalise_windows(D, windows):
    assert len(D) == len(windows) - 1
    for j in range(len(windows) - 1):
        span = windows[j + 1] - windows[j]
        D[j] /= span


def sample_set_normalisation(sample_sets):
    n = len(sample_sets)
    C = np.zeros((n, n))
    for j in range(n):
        C[j, j] = len(sample_sets[j]) * (len(sample_sets[j]) - 1)
        for k in range(j + 1, n):
            C[j, k] = len(sample_sets[j]) * len(sample_sets[k])
            C[k, j] = C[j, k]
    # Avoid division by zero for singleton samplesets
    C[C == 0] = 1
    # print("C = ", C)
    return C


def branch_divergence_matrix(ts, sample_sets=None, windows=None, span_normalise=True):
    windows_specified = windows is not None
    windows = ts.parse_windows(windows)
    num_windows = len(windows) - 1

    n = len(sample_sets)
    D = np.zeros((num_windows, n, n))
    tree = tskit.Tree(ts)
    C = sample_set_normalisation(sample_sets)
    for i in range(num_windows):
        left = windows[i]
        right = windows[i + 1]
        # print(f"WINDOW {i} [{left}, {right})")
        tree.seek(left)
        # Iterate over the trees in this window
        while tree.interval.left < right and tree.index != -1:
            span_left = max(tree.interval.left, left)
            span_right = min(tree.interval.right, right)
            span = span_right - span_left
            # print(f"\ttree {tree.interval} [{span_left}, {span_right})")
            tables = sv_tables_init(tree.parent_array)
            for j in range(n):
                for u in sample_sets[j]:
                    for k in range(j, n):
                        for v in sample_sets[k]:
                            # The u=v case here contributes zero, not bothering
                            # to exclude it.
                            w = sv_mrca(u, v, *tables)
                            assert w == tree.mrca(u, v)
                            if w != tskit.NULL:
                                tu = ts.nodes_time[w] - ts.nodes_time[u]
                                tv = ts.nodes_time[w] - ts.nodes_time[v]
                            else:
                                tu = (
                                    ts.nodes_time[local_root(tree, u)]
                                    - ts.nodes_time[u]
                                )
                                tv = (
                                    ts.nodes_time[local_root(tree, v)]
                                    - ts.nodes_time[v]
                                )
                            d = (tu + tv) * span
                            D[i, j, k] += d
            tree.next()
        # Fill out symmetric triangle in the matrix, and get average
        for j in range(n):
            D[i, j, j] /= C[j, j]
            for k in range(j + 1, n):
                D[i, j, k] /= C[j, k]
                D[i, k, j] = D[i, j, k]
    if span_normalise:
        span_normalise_windows(D, windows)
    if not windows_specified:
        D = D[0]
    return D


def divergence_matrix(
    ts, windows=None, sample_sets=None, samples=None, mode="site", span_normalise=True
):
    assert mode in ["site", "branch"]
    if samples is not None and sample_sets is not None:
        raise ValueError("Cannot specify both")
    if samples is None and sample_sets is None:
        samples = ts.samples()
    if samples is not None:
        sample_sets = [[u] for u in samples]
    else:
        assert sample_sets is not None

    if mode == "site":
        return site_divergence_matrix(
            ts, sample_sets, windows=windows, span_normalise=span_normalise
        )
    else:
        return branch_divergence_matrix(
            ts, sample_sets, windows=windows, span_normalise=span_normalise
        )


def stats_api_divergence_matrix(ts, *args, **kwargs):
    return stats_api_matrix_method(ts, ts.divergence, *args, **kwargs)


def stats_api_genetic_relatedness_matrix(ts, *args, **kwargs):
    method = functools.partial(ts.genetic_relatedness, proportion=False)
    return stats_api_matrix_method(ts, method, *args, **kwargs)


def stats_api_matrix_method(
    ts,
    method,
    windows=None,
    samples=None,
    sample_sets=None,
    mode="site",
    span_normalise=True,
):
    if samples is not None and sample_sets is not None:
        raise ValueError("Cannot specify both")
    if samples is None and sample_sets is None:
        samples = ts.samples()
    if samples is not None:
        sample_sets = [[u] for u in samples]
    else:
        assert sample_sets is not None

    windows_specified = windows is not None
    windows = [0, ts.sequence_length] if windows is None else list(windows)
    num_windows = len(windows) - 1

    if len(sample_sets) == 0:
        # FIXME: the code general stat code doesn't seem to handle zero samples
        # case, need to identify MWE and file issue.
        if windows_specified:
            return np.zeros(shape=(num_windows, 0, 0))
        else:
            return np.zeros(shape=(0, 0))

    # FIXME We have to go through this annoying rigmarole because windows must start and
    # end with 0 and L. We should relax this requirement to just making the windows
    # contiguous, so that we just look at specific sections of the genome.
    drop = []
    if windows[0] != 0:
        windows = [0] + windows
        drop.append(0)
    if windows[-1] != ts.sequence_length:
        windows.append(ts.sequence_length)
        drop.append(-1)

    n = len(sample_sets)
    indexes = [(i, j) for i in range(n) for j in range(n)]
    X = method(
        sample_sets,
        indexes=indexes,
        mode=mode,
        span_normalise=span_normalise,
        windows=windows,
    )
    keep = np.ones(len(windows) - 1, dtype=bool)
    keep[drop] = False
    X = X[keep]
    # Quick hack to get the within singleton sampleset divergence=0
    X[np.isnan(X)] = 0
    out = X.reshape((X.shape[0], n, n))
    if not windows_specified:
        out = out[0]
    return out


def group_alleles(genotypes, num_alleles):
    n = genotypes.shape[0]
    A = np.zeros(n, dtype=int)
    offsets = np.zeros(num_alleles + 1, dtype=int)
    k = 0
    for a in range(num_alleles):
        offsets[a + 1] = offsets[a]
        for j in range(n):
            if genotypes[j] == a:
                offsets[a + 1] += 1
                A[k] = j
                k += 1
    return A, offsets


def site_divergence_matrix(ts, sample_sets, *, windows=None, span_normalise=True):
    windows_specified = windows is not None
    windows = ts.parse_windows(windows)
    num_windows = len(windows) - 1

    n = len(sample_sets)
    samples = []
    sample_set_index_map = []
    for j in range(n):
        for u in sample_sets[j]:
            samples.append(u)
            sample_set_index_map.append(j)
    C = sample_set_normalisation(sample_sets)
    D = np.zeros((num_windows, n, n))

    site_id = 0
    while site_id < ts.num_sites and ts.sites_position[site_id] < windows[0]:
        site_id += 1

    # Note we have to use isolated_as_missing here because we're working with
    # non-sample nodes. There are tricky problems here later with missing data.
    variant = tskit.Variant(ts, samples=samples, isolated_as_missing=False)
    for i in range(num_windows):
        left = windows[i]
        right = windows[i + 1]
        if site_id < ts.num_sites:
            assert ts.sites_position[site_id] >= left
        while site_id < ts.num_sites and ts.sites_position[site_id] < right:
            variant.decode(site_id)
            X, offsets = group_alleles(variant.genotypes, variant.num_alleles)
            for j in range(variant.num_alleles):
                A = X[offsets[j] : offsets[j + 1]]
                for k in range(j + 1, variant.num_alleles):
                    B = X[offsets[k] : offsets[k + 1]]
                    for a in A:
                        a_set_index = sample_set_index_map[a]
                        for b in B:
                            b_set_index = sample_set_index_map[b]
                            D[i, a_set_index, b_set_index] += 1
                            D[i, b_set_index, a_set_index] += 1
            site_id += 1
        D[i] /= C
    if span_normalise:
        span_normalise_windows(D, windows)
    if not windows_specified:
        D = D[0]
    return D


def check_divmat(
    ts,
    *,
    windows=None,
    samples=None,
    sample_sets=None,
    span_normalise=True,
    verbosity=0,
    compare_stats_api=True,
    compare_lib=True,
    mode="site",
):
    # print("samples = ", samples, sample_sets)
    # print(ts.draw_text())
    if verbosity > 1:
        print(ts.draw_text())

    D1 = divergence_matrix(
        ts,
        sample_sets=sample_sets,
        samples=samples,
        windows=windows,
        mode=mode,
        span_normalise=span_normalise,
    )
    if compare_stats_api:
        D2 = stats_api_divergence_matrix(
            ts,
            windows=windows,
            samples=samples,
            sample_sets=sample_sets,
            mode=mode,
            span_normalise=span_normalise,
        )
        # print("windows = ", windows)
        # print(D1)
        # print(D2)
        np.testing.assert_allclose(D1, D2)
        assert D1.shape == D2.shape
    if compare_lib:
        ids = None
        if sample_sets is not None:
            ids = sample_sets
        if samples is not None:
            ids = samples
        D3 = ts.divergence_matrix(
            ids,
            windows=windows,
            mode=mode,
            span_normalise=span_normalise,
        )
        # print()
        # np.set_printoptions(linewidth=500, precision=4)
        # print(D1)
        # print(D3)
        assert D1.shape == D3.shape
        np.testing.assert_allclose(D1, D3)

    return D1


class TestExamplesWithAnswer:
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_single_tree_zero_samples(self, mode):
        ts = tskit.Tree.generate_balanced(2).tree_sequence
        D = check_divmat(ts, samples=[], mode=mode)
        assert D.shape == (0, 0)

    @pytest.mark.parametrize("num_windows", [1, 2, 3, 5])
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_single_tree_zero_samples_windows(self, num_windows, mode):
        ts = tskit.Tree.generate_balanced(2).tree_sequence
        windows = np.linspace(0, ts.sequence_length, num=num_windows + 1)
        D = check_divmat(ts, samples=[], windows=windows, mode=mode)
        assert D.shape == (num_windows, 0, 0)

    @pytest.mark.parametrize("m", [0, 1, 2, 10])
    def test_single_tree_sites_per_branch(self, m):
        # 2.00┊    6    ┊
        #     ┊  ┏━┻━┓  ┊
        # 1.00┊  4   5  ┊
        #     ┊ ┏┻┓ ┏┻┓ ┊
        # 0.00┊ 0 1 2 3 ┊
        #     0         1
        ts = tskit.Tree.generate_balanced(4).tree_sequence
        ts = tsutil.insert_branch_sites(ts, m)
        D1 = check_divmat(ts, mode="site")
        D2 = np.array(
            [
                [0.0, 2.0, 4.0, 4.0],
                [2.0, 0.0, 4.0, 4.0],
                [4.0, 4.0, 0.0, 2.0],
                [4.0, 4.0, 2.0, 0.0],
            ]
        )
        np.testing.assert_array_equal(D1, m * D2)

    @pytest.mark.parametrize("n", [2, 3, 5])
    def test_single_tree_unique_sample_alleles(self, n):
        tables = tskit.Tree.generate_balanced(n).tree_sequence.dump_tables()
        tables.sites.add_row(position=0.5, ancestral_state="0")
        for j in range(n):
            tables.mutations.add_row(site=0, node=j, derived_state=f"{j + 1}")
        ts = tables.tree_sequence()
        D1 = check_divmat(ts, mode="site")
        D2 = np.ones((n, n))
        np.fill_diagonal(D2, 0)
        np.testing.assert_array_equal(D1, D2)

    @pytest.mark.parametrize("L", [0.1, 1, 2, 100])
    def test_single_tree_sequence_length(self, L):
        # 2.00┊    6    ┊
        #     ┊  ┏━┻━┓  ┊
        # 1.00┊  4   5  ┊
        #     ┊ ┏┻┓ ┏┻┓ ┊
        # 0.00┊ 0 1 2 3 ┊
        #     0         1
        ts = tskit.Tree.generate_balanced(4, span=L).tree_sequence
        D1 = check_divmat(ts, mode="branch", span_normalise=False)
        D2 = np.array(
            [
                [0.0, 2.0, 4.0, 4.0],
                [2.0, 0.0, 4.0, 4.0],
                [4.0, 4.0, 0.0, 2.0],
                [4.0, 4.0, 2.0, 0.0],
            ]
        )
        np.testing.assert_array_equal(D1, L * D2)

    @pytest.mark.parametrize("L", [0.1, 1, 2, 100])
    def test_single_tree_sequence_length_span_normalise(self, L):
        # 2.00┊    6    ┊
        #     ┊  ┏━┻━┓  ┊
        # 1.00┊  4   5  ┊
        #     ┊ ┏┻┓ ┏┻┓ ┊
        # 0.00┊ 0 1 2 3 ┊
        #     0         1
        ts = tskit.Tree.generate_balanced(4, span=L).tree_sequence
        D1 = check_divmat(ts, mode="branch", span_normalise=True)
        D2 = np.array(
            [
                [0.0, 2.0, 4.0, 4.0],
                [2.0, 0.0, 4.0, 4.0],
                [4.0, 4.0, 0.0, 2.0],
                [4.0, 4.0, 2.0, 0.0],
            ]
        )
        np.testing.assert_array_equal(D1, D2)

    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_single_tree_diploid_individuals(self, mode):
        # 2.00┊    6    ┊
        #     ┊  ┏━┻━┓  ┊
        # 1.00┊  4   5  ┊
        #     ┊ ┏┻┓ ┏┻┓ ┊
        # 0.00┊ 0 1 2 3 ┊
        #     0         1
        ts = tskit.Tree.generate_balanced(4).tree_sequence
        ts = tsutil.insert_branch_sites(ts)
        ts = tsutil.insert_individuals(ts, ploidy=2)
        D1 = check_divmat(
            ts,
            sample_sets=[ind.nodes for ind in ts.individuals()],
            mode=mode,
        )
        D2 = np.array([[2.0, 4.0], [4.0, 2.0]])
        np.testing.assert_array_equal(D1, D2)

    @pytest.mark.parametrize("num_windows", [1, 2, 3, 5])
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_single_tree_gap_at_end(self, num_windows, mode):
        # 2.00┊    6    ┊
        #     ┊  ┏━┻━┓  ┊
        # 1.00┊  4   5  ┊
        #     ┊ ┏┻┓ ┏┻┓ ┊
        # 0.00┊ 0 1 2 3 ┊ 0 1 2 3
        #     0         1         2
        ts = tskit.Tree.generate_balanced(4).tree_sequence
        ts = tsutil.insert_branch_sites(ts)
        tables = ts.dump_tables()
        tables.sequence_length = 2
        ts = tables.tree_sequence()
        windows = np.linspace(0, ts.sequence_length, num=num_windows + 1)
        D1 = check_divmat(ts, windows=windows, mode=mode, span_normalise=False)
        D1 = np.sum(D1, axis=0)
        D2 = np.array(
            [
                [0.0, 2.0, 4.0, 4.0],
                [2.0, 0.0, 4.0, 4.0],
                [4.0, 4.0, 0.0, 2.0],
                [4.0, 4.0, 2.0, 0.0],
            ]
        )
        np.testing.assert_array_equal(D1, D2)

    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_single_tree_subset_permuted_samples(self, mode):
        # 2.00┊    6    ┊
        #     ┊  ┏━┻━┓  ┊
        # 1.00┊  4   5  ┊
        #     ┊ ┏┻┓ ┏┻┓ ┊
        # 0.00┊ 0 1 2 3 ┊
        #     0         1
        ts = tskit.Tree.generate_balanced(4).tree_sequence
        ts = tsutil.insert_branch_sites(ts)
        D1 = check_divmat(ts, samples=[1, 2, 0], mode=mode)
        D2 = np.array(
            [
                [0.0, 4.0, 2.0],
                [4.0, 0.0, 4.0],
                [2.0, 4.0, 0.0],
            ]
        )
        np.testing.assert_array_equal(D1, D2)

    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_single_tree_mixed_non_sample_samples(self, mode):
        # 2.00┊    6    ┊
        #     ┊  ┏━┻━┓  ┊
        # 1.00┊  4   5  ┊
        #     ┊ ┏┻┓ ┏┻┓ ┊
        # 0.00┊ 0 1 2 3 ┊
        #     0         1
        ts = tskit.Tree.generate_balanced(4).tree_sequence
        ts = tsutil.insert_branch_sites(ts)
        with pytest.raises(tskit.LibraryError, match="TSK_ERR_BAD_SAMPLES"):
            ts.divergence_matrix([0, 5], mode=mode)

    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_single_tree_duplicate_samples(self, mode):
        # 2.00┊    6    ┊
        #     ┊  ┏━┻━┓  ┊
        # 1.00┊  4   5  ┊
        #     ┊ ┏┻┓ ┏┻┓ ┊
        # 0.00┊ 0 1 2 3 ┊
        #     0         1
        ts = tskit.Tree.generate_balanced(4).tree_sequence
        ts = tsutil.insert_branch_sites(ts)
        with pytest.raises(tskit.LibraryError, match="TSK_ERR_DUPLICATE_SAMPLE"):
            ts.divergence_matrix([0, 0, 1], mode=mode)

    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_single_tree_multiroot(self, mode):
        # 2.00┊         ┊
        #     ┊         ┊
        # 1.00┊  4   5  ┊
        #     ┊ ┏┻┓ ┏┻┓ ┊
        # 0.00┊ 0 1 2 3 ┊
        #     0         1
        ts = tskit.Tree.generate_balanced(4).tree_sequence
        ts = tsutil.insert_branch_sites(ts)
        ts = ts.decapitate(1)
        D1 = check_divmat(ts, mode=mode)
        D2 = np.array(
            [
                [0.0, 2.0, 2.0, 2.0],
                [2.0, 0.0, 2.0, 2.0],
                [2.0, 2.0, 0.0, 2.0],
                [2.0, 2.0, 2.0, 0.0],
            ]
        )
        np.testing.assert_array_equal(D1, D2)

    @pytest.mark.parametrize(
        ["left", "right"], [(0, 10), (1, 3), (3.25, 3.75), (5, 10)]
    )
    def test_single_tree_interval(self, left, right):
        # 2.00┊    6    ┊
        #     ┊  ┏━┻━┓  ┊
        # 1.00┊  4   5  ┊
        #     ┊ ┏┻┓ ┏┻┓ ┊
        # 0.00┊ 0 1 2 3 ┊
        #     0         1
        ts = tskit.Tree.generate_balanced(4, span=10).tree_sequence
        D1 = check_divmat(
            ts, windows=[left, right], mode="branch", span_normalise=False
        )
        D2 = np.array(
            [
                [0.0, 2.0, 4.0, 4.0],
                [2.0, 0.0, 4.0, 4.0],
                [4.0, 4.0, 0.0, 2.0],
                [4.0, 4.0, 2.0, 0.0],
            ]
        )
        np.testing.assert_array_equal(D1[0], (right - left) * D2)

    @pytest.mark.parametrize("num_windows", [1, 2, 3, 5, 11])
    def test_single_tree_equal_windows(self, num_windows):
        # 2.00┊    6    ┊
        #     ┊  ┏━┻━┓  ┊
        # 1.00┊  4   5  ┊
        #     ┊ ┏┻┓ ┏┻┓ ┊
        # 0.00┊ 0 1 2 3 ┊
        #     0         1
        ts = tskit.Tree.generate_balanced(4, span=10).tree_sequence
        windows = np.linspace(0, ts.sequence_length, num=num_windows + 1)
        x = ts.sequence_length / num_windows
        # print(windows)
        D1 = check_divmat(ts, windows=windows, mode="branch", span_normalise=False)
        assert D1.shape == (num_windows, 4, 4)
        D2 = np.array(
            [
                [0.0, 2.0, 4.0, 4.0],
                [2.0, 0.0, 4.0, 4.0],
                [4.0, 4.0, 0.0, 2.0],
                [4.0, 4.0, 2.0, 0.0],
            ]
        )
        for D in D1:
            np.testing.assert_array_almost_equal(D, x * D2)

    @pytest.mark.parametrize("n", [2, 3, 5])
    def test_single_tree_no_sites(self, n):
        ts = tskit.Tree.generate_balanced(n, span=10).tree_sequence
        D = check_divmat(ts, mode="site")
        np.testing.assert_array_equal(D, np.zeros((n, n)))


class TestExamples:
    @pytest.mark.parametrize(
        "interval", [(0, 26), (1, 3), (3.25, 13.75), (5, 10), (25.5, 26)]
    )
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    @pytest.mark.parametrize("span_normalise", [True, False])
    def test_all_trees_interval(self, interval, mode, span_normalise):
        ts = tsutil.all_trees_ts(4)
        ts = tsutil.insert_branch_sites(ts)
        assert ts.sequence_length == 26
        check_divmat(ts, windows=interval, mode=mode, span_normalise=span_normalise)

    @pytest.mark.parametrize(
        ["windows"],
        [
            ([0, 26],),
            ([0, 1, 2],),
            (list(range(27)),),
            ([5, 7, 9, 20],),
            ([5.1, 5.2, 5.3, 5.5, 6],),
            ([5.1, 5.2, 6.5],),
        ],
    )
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    @pytest.mark.parametrize("span_normalise", [True, False])
    def test_all_trees_windows(self, windows, mode, span_normalise):
        ts = tsutil.all_trees_ts(4)
        ts = tsutil.insert_branch_sites(ts)
        assert ts.sequence_length == 26
        D = check_divmat(ts, windows=windows, mode=mode, span_normalise=span_normalise)
        assert D.shape == (len(windows) - 1, 4, 4)

    @pytest.mark.parametrize("num_windows", [1, 5, 28])
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    @pytest.mark.parametrize("span_normalise", [True, False])
    def test_all_trees_windows_gap_at_end(self, num_windows, mode, span_normalise):
        tables = tsutil.all_trees_ts(4).dump_tables()
        tables.sequence_length = 30
        ts = tables.tree_sequence()
        ts = tsutil.insert_branch_sites(ts)
        assert ts.last().num_roots == 4
        windows = np.linspace(0, ts.sequence_length, num=num_windows + 1)
        check_divmat(ts, windows=windows, mode=mode, span_normalise=span_normalise)

    @pytest.mark.parametrize("n", [2, 3, 5])
    @pytest.mark.parametrize("seed", range(1, 4))
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_small_sims(self, n, seed, mode):
        ts = msprime.sim_ancestry(
            n,
            ploidy=1,
            sequence_length=1000,
            recombination_rate=0.01,
            random_seed=seed,
        )
        assert ts.num_trees >= 2
        ts = msprime.sim_mutations(
            ts, rate=0.1, discrete_genome=False, random_seed=seed
        )
        assert ts.num_mutations > 1
        check_divmat(ts, verbosity=0, mode=mode)

    @pytest.mark.parametrize("n", [2, 3, 5, 15])
    @pytest.mark.parametrize("num_windows", range(1, 5))
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_sims_windows(self, n, num_windows, mode):
        ts = msprime.sim_ancestry(
            n,
            ploidy=1,
            population_size=20,
            sequence_length=100,
            recombination_rate=0.01,
            random_seed=79234,
        )
        assert ts.num_trees >= 2
        ts = msprime.sim_mutations(
            ts,
            rate=0.01,
            discrete_genome=False,
            random_seed=1234,
        )
        assert ts.num_mutations >= 2
        windows = np.linspace(0, ts.sequence_length, num=num_windows + 1)
        check_divmat(ts, windows=windows, mode=mode)

    @pytest.mark.parametrize("n", [2, 3, 5, 15])
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_single_balanced_tree(self, n, mode):
        ts = tskit.Tree.generate_balanced(n).tree_sequence
        ts = tsutil.insert_branch_sites(ts)
        # print(ts.draw_text())
        check_divmat(ts, verbosity=0, mode=mode)

    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_internal_sample(self, mode):
        tables = tskit.Tree.generate_balanced(4).tree_sequence.dump_tables()
        flags = tables.nodes.flags
        flags[3] = 0
        flags[5] = tskit.NODE_IS_SAMPLE
        tables.nodes.flags = flags
        ts = tables.tree_sequence()
        ts = tsutil.insert_branch_sites(ts)
        check_divmat(ts, verbosity=0, mode=mode)

    @pytest.mark.parametrize("seed", range(1, 5))
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_one_internal_sample_sims(self, seed, mode):
        ts = msprime.sim_ancestry(
            10,
            ploidy=1,
            population_size=20,
            sequence_length=100,
            recombination_rate=0.01,
            random_seed=seed,
        )
        t = ts.dump_tables()
        # Add a new sample directly below another sample
        u = t.nodes.add_row(time=-1, flags=tskit.NODE_IS_SAMPLE)
        t.edges.add_row(parent=0, child=u, left=0, right=ts.sequence_length)
        t.sort()
        t.build_index()
        ts = t.tree_sequence()
        ts = tsutil.insert_branch_sites(ts)
        check_divmat(ts, mode=mode)

    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_missing_flanks(self, mode):
        ts = msprime.sim_ancestry(
            20,
            ploidy=1,
            population_size=20,
            sequence_length=100,
            recombination_rate=0.01,
            random_seed=1234,
        )
        assert ts.num_trees >= 2
        ts = ts.keep_intervals([[20, 80]])
        assert ts.first().interval == (0, 20)
        ts = tsutil.insert_branch_sites(ts)
        check_divmat(ts, mode=mode)

    @pytest.mark.parametrize("n", [2, 3, 10])
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_dangling_on_samples(self, n, mode):
        # Adding non sample branches below the samples does not alter
        # the overall divergence *between* the samples
        ts1 = tskit.Tree.generate_balanced(n).tree_sequence
        ts1 = tsutil.insert_branch_sites(ts1)
        D1 = check_divmat(ts1, mode=mode)
        tables = ts1.dump_tables()
        for u in ts1.samples():
            v = tables.nodes.add_row(time=-1)
            tables.edges.add_row(left=0, right=ts1.sequence_length, parent=u, child=v)
        tables.sort()
        tables.build_index()
        ts2 = tables.tree_sequence()
        D2 = check_divmat(ts2, mode=mode)
        np.testing.assert_array_almost_equal(D1, D2)

    @pytest.mark.parametrize("n", [2, 3, 10])
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_dangling_on_all(self, n, mode):
        # Adding non sample branches below the samples does not alter
        # the overall divergence *between* the samples
        ts1 = tskit.Tree.generate_balanced(n).tree_sequence
        ts1 = tsutil.insert_branch_sites(ts1)
        D1 = check_divmat(ts1, mode=mode)
        tables = ts1.dump_tables()
        for u in range(ts1.num_nodes):
            v = tables.nodes.add_row(time=-1)
            tables.edges.add_row(left=0, right=ts1.sequence_length, parent=u, child=v)
        tables.sort()
        tables.build_index()
        ts2 = tables.tree_sequence()
        D2 = check_divmat(ts2, mode=mode)
        np.testing.assert_array_almost_equal(D1, D2)

    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_disconnected_non_sample_topology(self, mode):
        # Adding non sample branches below the samples does not alter
        # the overall divergence *between* the samples
        ts1 = tskit.Tree.generate_balanced(5).tree_sequence
        ts1 = tsutil.insert_branch_sites(ts1)
        D1 = check_divmat(ts1, mode=mode)
        tables = ts1.dump_tables()
        # Add an extra bit of disconnected non-sample topology
        u = tables.nodes.add_row(time=0)
        v = tables.nodes.add_row(time=1)
        tables.edges.add_row(left=0, right=ts1.sequence_length, parent=v, child=u)
        tables.sort()
        tables.build_index()
        ts2 = tables.tree_sequence()
        D2 = check_divmat(ts2, mode=mode)
        np.testing.assert_array_almost_equal(D1, D2)


class TestSuiteExamples:
    """
    Compare the stats API method vs the library implementation for the
    suite test examples. Some of these examples are too large to run the
    Python code above on.
    """

    def check(
        self,
        ts,
        windows=None,
        sample_sets=None,
        num_threads=0,
        span_normalise=True,
        mode="branch",
    ):
        D1 = ts.divergence_matrix(
            sample_sets,
            windows=windows,
            num_threads=num_threads,
            mode=mode,
            span_normalise=span_normalise,
        )
        D2 = stats_api_divergence_matrix(
            ts,
            windows=windows,
            sample_sets=sample_sets,
            mode=mode,
            span_normalise=span_normalise,
        )
        assert D1.shape == D2.shape
        # np.set_printoptions(linewidth=500, precision=4)
        # print()
        # print(D1)
        # print(D2)
        if mode == "branch":
            # If we have missing data then parts of the divmat are defined to be zero,
            # so relative tolerances aren't useful. Because the stats API
            # method necessarily involves subtracting away all of the previous
            # values for an empty tree, there is a degree of numerical imprecision
            # here. This value for atol is what is needed to get the tests to
            # pass in practise.
            has_missing_data = any(tree._has_isolated_samples() for tree in ts.trees())
            atol = 1e-11 if has_missing_data else 0
            np.testing.assert_allclose(D1, D2, atol=atol)
        else:
            assert mode == "site"
            np.testing.assert_allclose(D1, D2)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_defaults(self, ts, mode):
        self.check(ts, mode=mode)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_subset_samples(self, ts, mode):
        n = min(ts.num_samples, 2)
        self.check(ts, sample_sets=[[u] for u in ts.samples()[:n]], mode=mode)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    @pytest.mark.parametrize("ploidy", [1, 2, 3])
    def test_ploidy_sample_sets(self, ts, mode, ploidy):
        if ts.num_samples >= 2 * ploidy:
            # Workaround limitations in the stats API
            sample_sets = np.array_split(ts.samples(), ts.num_samples // ploidy)
            self.check(ts, sample_sets=sample_sets, mode=mode)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    @pytest.mark.parametrize("span_normalise", [True, False])
    def test_windows(self, ts, mode, span_normalise):
        windows = np.linspace(0, ts.sequence_length, num=13)
        self.check(ts, windows=windows, mode=mode, span_normalise=span_normalise)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_threads_no_windows(self, ts, mode):
        self.check(ts, num_threads=5, mode=mode)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_threads_windows(self, ts, mode):
        windows = np.linspace(0, ts.sequence_length, num=11)
        self.check(ts, num_threads=5, windows=windows, mode=mode)


class TestThreadsNoWindows:
    def check(self, ts, num_threads, samples=None, mode=None, span_normalise=True):
        D1 = ts.divergence_matrix(
            samples, num_threads=0, mode=mode, span_normalise=span_normalise
        )
        D2 = ts.divergence_matrix(
            samples,
            num_threads=num_threads,
            mode=mode,
            span_normalise=span_normalise,
        )
        np.testing.assert_array_almost_equal(D1, D2)

    @pytest.mark.parametrize("num_threads", [1, 2, 3, 5, 26, 27])
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    @pytest.mark.parametrize("span_normalise", [True, False])
    def test_all_trees(self, num_threads, mode, span_normalise):
        ts = tsutil.all_trees_ts(4)
        assert ts.num_trees == 26
        self.check(ts, num_threads, mode=mode, span_normalise=span_normalise)

    @pytest.mark.parametrize("samples", [None, [0, 1]])
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_all_trees_samples(self, samples, mode):
        ts = tsutil.all_trees_ts(4)
        assert ts.num_trees == 26
        self.check(ts, 2, samples, mode=mode)

    @pytest.mark.parametrize("n", [2, 3, 5, 15])
    @pytest.mark.parametrize("num_threads", range(1, 5))
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_simple_sims(self, n, num_threads, mode):
        ts = msprime.sim_ancestry(
            n,
            ploidy=1,
            population_size=20,
            sequence_length=100,
            recombination_rate=0.01,
            random_seed=1234,
        )
        assert ts.num_trees >= 2
        self.check(ts, num_threads, mode=mode)


class TestThreadsWindows:
    def check(self, ts, num_threads, *, windows, samples=None, mode=None):
        D1 = ts.divergence_matrix(samples, num_threads=0, windows=windows, mode=mode)
        D2 = ts.divergence_matrix(
            samples, num_threads=num_threads, windows=windows, mode=mode
        )
        np.testing.assert_array_almost_equal(D1, D2)

    @pytest.mark.parametrize("num_threads", [1, 2, 3, 5, 26, 27])
    @pytest.mark.parametrize(
        ["windows"],
        [
            ([0, 26],),
            ([0, 1, 2],),
            (list(range(27)),),
            ([5, 7, 9, 20],),
            ([5.1, 5.2, 5.3, 5.5, 6],),
            ([5.1, 5.2, 6.5],),
            ("trees",),
            ("sites",),
        ],
    )
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_all_trees(self, num_threads, windows, mode):
        ts = tsutil.all_trees_ts(4)
        assert ts.num_trees == 26
        self.check(ts, num_threads, windows=windows, mode=mode)

    @pytest.mark.parametrize("samples", [None, [0, 1]])
    @pytest.mark.parametrize(
        ["windows"],
        [
            ([0, 26],),
            (None,),
            ("trees",),
            ("sites",),
        ],
    )
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_all_trees_samples(self, samples, windows, mode):
        ts = tsutil.all_trees_ts(4)
        self.check(ts, 2, windows=windows, samples=samples, mode=mode)

    @pytest.mark.parametrize("num_threads", range(1, 5))
    @pytest.mark.parametrize(
        ["windows"],
        [
            ([0, 100],),
            ([0, 50, 75, 95, 100],),
            ([50, 75, 95, 100],),
            ([0, 50, 75, 95],),
            (list(range(100)),),
            ("trees",),
            ("sites",),
        ],
    )
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_simple_sims(self, num_threads, windows, mode):
        ts = msprime.sim_ancestry(
            15,
            ploidy=1,
            population_size=20,
            sequence_length=100,
            recombination_rate=0.01,
            random_seed=1234,
        )
        assert ts.num_trees >= 2
        ts = msprime.sim_mutations(ts, rate=0.01, random_seed=1234)
        assert ts.num_mutations > 10
        self.check(ts, num_threads, windows=windows, mode=mode)


# NOTE these are tests that are for more general functionality that might
# get applied across many different functions, and so probably should be
# tested in another file. For now they're only used by divmat, so we can
# keep them here for simplificity.
class TestChunkByTree:
    # These are based on what we get from np.array_split, there's nothing
    # particularly critical about exactly how we portion things up.
    @pytest.mark.parametrize(
        ["num_chunks", "expected"],
        [
            (1, [[0, 26]]),
            (2, [[0, 13], [13, 26]]),
            (3, [[0, 9], [9, 18], [18, 26]]),
            (4, [[0, 7], [7, 14], [14, 20], [20, 26]]),
            (5, [[0, 6], [6, 11], [11, 16], [16, 21], [21, 26]]),
        ],
    )
    def test_all_trees_ts_26(self, num_chunks, expected):
        ts = tsutil.all_trees_ts(4)
        actual = ts._chunk_sequence_by_tree(num_chunks)
        np.testing.assert_equal(actual, expected)

    @pytest.mark.parametrize(
        ["num_chunks", "expected"],
        [
            (1, [[0, 4]]),
            (2, [[0, 2], [2, 4]]),
            (3, [[0, 2], [2, 3], [3, 4]]),
            (4, [[0, 1], [1, 2], [2, 3], [3, 4]]),
            (5, [[0, 1], [1, 2], [2, 3], [3, 4]]),
            (100, [[0, 1], [1, 2], [2, 3], [3, 4]]),
        ],
    )
    def test_all_trees_ts_4(self, num_chunks, expected):
        ts = tsutil.all_trees_ts(3)
        assert ts.num_trees == 4
        actual = ts._chunk_sequence_by_tree(num_chunks)
        np.testing.assert_equal(actual, expected)

    @pytest.mark.parametrize("span", [1, 2, 5, 0.3])
    @pytest.mark.parametrize(
        ["num_chunks", "expected"],
        [
            (1, [[0, 4]]),
            (2, [[0, 2], [2, 4]]),
            (3, [[0, 2], [2, 3], [3, 4]]),
            (4, [[0, 1], [1, 2], [2, 3], [3, 4]]),
            (5, [[0, 1], [1, 2], [2, 3], [3, 4]]),
            (100, [[0, 1], [1, 2], [2, 3], [3, 4]]),
        ],
    )
    def test_all_trees_ts_4_trees_span(self, span, num_chunks, expected):
        tables = tsutil.all_trees_ts(3).dump_tables()
        tables.edges.left *= span
        tables.edges.right *= span
        tables.sequence_length *= span
        ts = tables.tree_sequence()
        assert ts.num_trees == 4
        actual = ts._chunk_sequence_by_tree(num_chunks)
        np.testing.assert_equal(actual, np.array(expected) * span)

    @pytest.mark.parametrize("num_chunks", range(1, 5))
    def test_empty_ts(self, num_chunks):
        tables = tskit.TableCollection(1)
        ts = tables.tree_sequence()
        chunks = ts._chunk_sequence_by_tree(num_chunks)
        np.testing.assert_equal(chunks, [[0, 1]])

    @pytest.mark.parametrize("num_chunks", range(1, 5))
    def test_single_tree(self, num_chunks):
        L = 10
        ts = tskit.Tree.generate_balanced(2, span=L).tree_sequence
        chunks = ts._chunk_sequence_by_tree(num_chunks)
        np.testing.assert_equal(chunks, [[0, L]])

    @pytest.mark.parametrize("num_chunks", [0, -1, 0.5])
    def test_bad_chunks(self, num_chunks):
        ts = tskit.Tree.generate_balanced(2).tree_sequence
        with pytest.raises(ValueError, match="Number of chunks must be an integer > 0"):
            ts._chunk_sequence_by_tree(num_chunks)


class TestChunkWindows:
    # These are based on what we get from np.array_split, there's nothing
    # particularly critical about exactly how we portion things up.
    @pytest.mark.parametrize(
        ["windows", "num_chunks", "expected"],
        [
            ([0, 10], 1, [[0, 10]]),
            ([0, 10], 2, [[0, 10]]),
            ([0, 5, 10], 2, [[0, 5], [5, 10]]),
            ([0, 5, 6, 10], 2, [[0, 5, 6], [6, 10]]),
            ([0, 5, 6, 10], 3, [[0, 5], [5, 6], [6, 10]]),
        ],
    )
    def test_examples(self, windows, num_chunks, expected):
        actual = tskit.TreeSequence._chunk_windows(windows, num_chunks)
        np.testing.assert_equal(actual, expected)

    @pytest.mark.parametrize("num_chunks", [0, -1, 0.5])
    def test_bad_chunks(self, num_chunks):
        with pytest.raises(ValueError, match="Number of chunks must be an integer > 0"):
            tskit.TreeSequence._chunk_windows([0, 1], num_chunks)


class TestGroupAlleles:
    @pytest.mark.parametrize(
        ["G", "num_alleles", "A", "offsets"],
        [
            ([0, 1], 2, [0, 1], [0, 1, 2]),
            ([0, 1], 3, [0, 1], [0, 1, 2, 2]),
            ([0, 2], 3, [0, 1], [0, 1, 1, 2]),
            ([1, 0], 2, [1, 0], [0, 1, 2]),
            ([0, 0, 0, 1, 1, 1], 2, [0, 1, 2, 3, 4, 5], [0, 3, 6]),
            ([0, 0], 1, [0, 1], [0, 2]),
            ([2, 2], 3, [0, 1], [0, 0, 0, 2]),
        ],
    )
    def test_examples(self, G, num_alleles, A, offsets):
        A1, offsets1 = group_alleles(np.array(G), num_alleles)
        assert list(A) == list(A1)
        assert list(offsets) == list(offsets1)

    def test_simple_simulation(self):
        ts = msprime.sim_ancestry(
            15,
            ploidy=1,
            population_size=20,
            sequence_length=100,
            recombination_rate=0.01,
            random_seed=1234,
        )
        ts = msprime.sim_mutations(ts, rate=0.01, random_seed=1234)
        assert ts.num_mutations > 10
        for var in ts.variants():
            A, offsets = group_alleles(var.genotypes, var.num_alleles)
            allele_samples = [[] for _ in range(var.num_alleles)]
            for j, a in enumerate(var.genotypes):
                allele_samples[a].append(j)

            assert len(offsets) == var.num_alleles + 1
            assert offsets[0] == 0
            assert offsets[-1] == ts.num_samples
            assert np.all(np.diff(offsets) >= 0)
            for j in range(var.num_alleles):
                a = A[offsets[j] : offsets[j + 1]]
                assert list(a) == list(allele_samples[j])


class TestSampleSetParsing:
    @pytest.mark.parametrize(
        ["arg", "flattened", "sizes"],
        [
            ([], [], []),
            ([1], [1], [1]),
            ([1, 2], [1, 2], [1, 1]),
            ([[1, 2], [3, 4]], [1, 2, 3, 4], [2, 2]),
            (((1, 2), (3, 4)), [1, 2, 3, 4], [2, 2]),
            (np.array([[1, 2], [3, 4]]), [1, 2, 3, 4], [2, 2]),
            (np.array([1, 2]), [1, 2], [1, 1]),
            (np.array([1, 2], dtype=np.uint32), [1, 2], [1, 1]),
            (array.array("i", [1, 2]), [1, 2], [1, 1]),
            ([[1, 2], [3], [4]], [1, 2, 3, 4], [2, 1, 1]),
            ([[1], [2]], [1, 2], [1, 1]),
            ([[1, 1], [2]], [1, 1, 2], [2, 1]),
        ],
    )
    def test_good_args(self, arg, flattened, sizes):
        f, s = tskit.TreeSequence._parse_stat_matrix_sample_sets(arg)
        # print(f, s)
        assert isinstance(f, np.ndarray)
        assert f.dtype == np.int32
        assert isinstance(s, np.ndarray)
        assert s.dtype == np.uint64
        np.testing.assert_array_equal(f, flattened)
        np.testing.assert_array_equal(s, sizes)

    @pytest.mark.parametrize(
        "arg",
        [
            ["0", "1"],
            ["0", 1],
            [0, "1"],
            [0, {"a": "b"}],
        ],
    )
    def test_nested_bad_types(self, arg):
        with pytest.raises(TypeError):
            tskit.TreeSequence._parse_stat_matrix_sample_sets(arg)

    @pytest.mark.parametrize(
        "arg",
        [
            [[0], [[0, 0]]],
            [[[0, 0]], [0]],
            np.array([[[0, 0], [0, 0]]]),
        ],
    )
    def test_nested_arrays(self, arg):
        with pytest.raises(ValueError):
            tskit.TreeSequence._parse_stat_matrix_sample_sets(arg)

    @pytest.mark.parametrize("arg", ["", "string", "1", "[1, 2]", b"", "1234"])
    def test_string_args(self, arg):
        with pytest.raises(TypeError, match="ID specification cannot be"):
            tskit.TreeSequence._parse_stat_matrix_sample_sets(arg)

    @pytest.mark.parametrize(
        "arg",
        [
            {},
            {"a": "b"},
            collections.Counter(),
        ],
    )
    def test_dict_args(self, arg):
        with pytest.raises(TypeError, match="ID specification cannot be"):
            tskit.TreeSequence._parse_stat_matrix_sample_sets(arg)

    @pytest.mark.parametrize(
        "arg",
        [
            0,
            {0: 1},
            None,
            {"a": "b"},
            np.array([1.1]),
        ],
    )
    def test_bad_arg_types(self, arg):
        with pytest.raises(TypeError):
            tskit.TreeSequence._parse_stat_matrix_sample_sets(arg)


class TestGeneticRelatednessMatrix:
    def check(self, ts, mode, *, sample_sets=None, windows=None, span_normalise=True):
        # These are *only* expected to be the same
        # under infinite-sites mutations
        if mode == "site" and np.any([len(s.mutations) > 1 for s in ts.sites()]):
            ts = msprime.sim_mutations(
                ts,
                rate=100 / ts.segregating_sites(mode="branch", span_normalise=False),
                random_seed=123,
                discrete_genome=False,
                keep=False,
            )
        G1 = stats_api_genetic_relatedness_matrix(
            ts,
            mode=mode,
            sample_sets=sample_sets,
            windows=windows,
            span_normalise=span_normalise,
        )
        G2 = ts.genetic_relatedness_matrix(
            mode=mode,
            sample_sets=sample_sets,
            windows=windows,
            span_normalise=span_normalise,
        )
        np.testing.assert_array_almost_equal(G1, G2)

    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_single_tree(self, mode):
        # 2.00┊    6    ┊
        #     ┊  ┏━┻━┓  ┊
        # 1.00┊  4   5  ┊
        #     ┊ ┏┻┓ ┏┻┓ ┊
        # 0.00┊ 0 1 2 3 ┊
        #     0         1
        ts = tskit.Tree.generate_balanced(4).tree_sequence
        ts = tsutil.insert_branch_sites(ts)
        self.check(ts, mode)

    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_single_tree_sample_sets(self, mode):
        # 2.00┊    6    ┊
        #     ┊  ┏━┻━┓  ┊
        # 1.00┊  4   5  ┊
        #     ┊ ┏┻┓ ┏┻┓ ┊
        # 0.00┊ 0 1 2 3 ┊
        #     0         1
        ts = tskit.Tree.generate_balanced(4).tree_sequence
        ts = tsutil.insert_branch_sites(ts)
        self.check(ts, mode, sample_sets=[[0, 1], [2, 3]])

    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_single_tree_single_samples(self, mode):
        # 2.00┊    6    ┊
        #     ┊  ┏━┻━┓  ┊
        # 1.00┊  4   5  ┊
        #     ┊ ┏┻┓ ┏┻┓ ┊
        # 0.00┊ 0 1 2 3 ┊
        #     0         1
        ts = tskit.Tree.generate_balanced(4).tree_sequence
        ts = tsutil.insert_branch_sites(ts)
        self.check(ts, mode, sample_sets=[[0], [1]])
        self.check(ts, mode, sample_sets=[[0], [2]])
        self.check(ts, mode, sample_sets=[[0], [1], [2]])

    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_single_tree_windows(self, mode):
        # 2.00┊    6    ┊
        #     ┊  ┏━┻━┓  ┊
        # 1.00┊  4   5  ┊
        #     ┊ ┏┻┓ ┏┻┓ ┊
        # 0.00┊ 0 1 2 3 ┊
        #     0         1
        ts = tskit.Tree.generate_balanced(4).tree_sequence
        ts = tsutil.insert_branch_sites(ts)
        self.check(ts, mode, windows=[0, 0.5, 1])

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    def test_suite_defaults(self, ts, mode):
        self.check(ts, mode=mode)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    @pytest.mark.parametrize("span_normalise", [True, False])
    def test_suite_span_normalise(self, ts, mode, span_normalise):
        self.check(ts, mode=mode, span_normalise=span_normalise)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    @pytest.mark.parametrize("mode", DIVMAT_MODES)
    @pytest.mark.parametrize("num_sets", [2])  # [[2, 3, 4, 5])
    def test_suite_sample_sets(self, ts, mode, num_sets):
        if ts.num_samples >= num_sets:
            sample_sets = np.array_split(ts.samples(), num_sets)
            self.check(ts, sample_sets=sample_sets, mode=mode)


--- ../../tskit/python/tests/test_combinatorics.py ---


"""
Test cases for combinatorial algorithms.
"""
import collections
import io
import itertools
import json
import math
import random

import msprime
import numpy as np
import pytest

import tests.test_wright_fisher as wf
import tskit
import tskit.combinatorics as comb
from tests import test_stats
from tskit.combinatorics import Rank
from tskit.combinatorics import RankTree


class TestCombination:
    def test_combination_with_replacement_rank_unrank(self):
        for n in range(9):
            for k in range(n):
                nums = list(range(n))
                combs = itertools.combinations_with_replacement(nums, k)
                for exp_rank, c in enumerate(combs):
                    c = list(c)
                    actual_rank = comb.Combination.with_replacement_rank(c, n)
                    assert actual_rank == exp_rank
                    unranked = comb.Combination.with_replacement_unrank(exp_rank, n, k)
                    assert unranked == c

    def test_combination_rank_unrank(self):
        for n in range(11):
            for k in range(n):
                nums = list(range(n))
                for rank, c in enumerate(itertools.combinations(nums, k)):
                    c = list(c)
                    assert comb.Combination.rank(c, nums) == rank
                    assert comb.Combination.unrank(rank, nums, k) == c

    def test_combination_unrank_errors(self):
        self.verify_unrank_errors(1, 1, 1)
        self.verify_unrank_errors(2, 0, 1)

    def verify_unrank_errors(self, rank, n, k):
        with pytest.raises(ValueError):
            comb.Combination.unrank(rank, list(range(n)), k)


class TestPartition:
    def test_rule_asc(self):
        self.verify_rule_asc(1, [[1]])
        self.verify_rule_asc(2, [[1, 1], [2]])
        self.verify_rule_asc(3, [[1, 1, 1], [1, 2], [3]])
        self.verify_rule_asc(4, [[1, 1, 1, 1], [1, 1, 2], [1, 3], [2, 2], [4]])
        self.verify_rule_asc(
            5,
            [[1, 1, 1, 1, 1], [1, 1, 1, 2], [1, 1, 3], [1, 2, 2], [1, 4], [2, 3], [5]],
        )
        self.verify_rule_asc(
            6,
            [
                [1, 1, 1, 1, 1, 1],
                [1, 1, 1, 1, 2],
                [1, 1, 1, 3],
                [1, 1, 2, 2],
                [1, 1, 4],
                [1, 2, 3],
                [1, 5],
                [2, 2, 2],
                [2, 4],
                [3, 3],
                [6],
            ],
        )

    def verify_rule_asc(self, n, partitions):
        assert list(comb.rule_asc(n)) == partitions

    def test_partitions(self):
        assert list(comb.partitions(0)) == []
        for n in range(1, 7):
            assert list(comb.partitions(n)) == list(comb.rule_asc(n))[:-1]

    def test_group_partition(self):
        assert comb.group_partition([1]) == [[1]]
        assert comb.group_partition([1, 2]) == [[1], [2]]
        assert comb.group_partition([1, 1, 1]) == [[1, 1, 1]]
        assert comb.group_partition([1, 1, 2, 3, 3]) == [[1, 1], [2], [3, 3]]


class TestRankTree:
    @pytest.mark.parametrize("n", range(11))
    def test_num_shapes(self, n):
        all_trees = RankTree.all_unlabelled_trees(n)
        assert len(list(all_trees)) == comb.num_shapes(n)

    @pytest.mark.parametrize("n", range(2, 8))
    def test_num_labellings(self, n):
        for tree in RankTree.all_unlabelled_trees(n):
            tree = tree.label_unrank(0)
            tree2 = tree.to_tsk_tree()
            n_labellings = sum(1 for _ in RankTree.all_labellings(tree))
            assert n_labellings == RankTree.from_tsk_tree(tree2).num_labellings()

    def test_num_labelled_trees(self):
        # Number of leaf-labelled trees with n leaves on OEIS
        n_trees = [0, 1, 1, 4, 26, 236, 2752, 39208]
        for i, expected in zip(range(len(n_trees)), n_trees):
            actual = sum(1 for _ in RankTree.all_labelled_trees(i))
            assert actual == expected

    def test_all_labelled_trees_3(self):
        expected = ["(0,1,2);", "(0,(1,2));", "(1,(0,2));", "(2,(0,1));"]
        actual = [t.newick() for t in RankTree.all_labelled_trees(3)]
        assert expected == actual

    def test_all_labelled_trees_4(self):
        expected = [
            # 1 + 1 + 1 + 1 (partition of num leaves)
            "(0,1,2,3);",
            # 1 + 1 + 2
            "(0,1,(2,3));",
            "(0,2,(1,3));",
            "(0,3,(1,2));",
            "(1,2,(0,3));",
            "(1,3,(0,2));",
            "(2,3,(0,1));",
            # 1 + 3
            # partition of 3 = 1 + 1 + 1
            "(0,(1,2,3));",
            "(1,(0,2,3));",
            "(2,(0,1,3));",
            "(3,(0,1,2));",
            # partition of 3 = 1 + 2
            "(0,(1,(2,3)));",
            "(0,(2,(1,3)));",
            "(0,(3,(1,2)));",
            "(1,(0,(2,3)));",
            "(1,(2,(0,3)));",
            "(1,(3,(0,2)));",
            "(2,(0,(1,3)));",
            "(2,(1,(0,3)));",
            "(2,(3,(0,1)));",
            "(3,(0,(1,2)));",
            "(3,(1,(0,2)));",
            "(3,(2,(0,1)));",
            # 2 + 2
            "((0,1),(2,3));",
            "((0,2),(1,3));",
            "((0,3),(1,2));",
        ]
        actual = [t.newick() for t in RankTree.all_labelled_trees(4)]
        assert expected == actual

    def test_generate_trees_roundtrip(self):
        n = 5
        all_rank_trees = RankTree.all_labelled_trees(n)
        all_tsk_trees = tskit.all_trees(n)
        for rank_tree, tsk_tree in zip(all_rank_trees, all_tsk_trees):
            assert rank_tree == RankTree.from_tsk_tree(tsk_tree)

    def test_generate_treeseq_roundtrip(self):
        n = 5
        span = 9
        all_rank_trees = RankTree.all_labelled_trees(n)
        all_tsk_trees = tskit.all_trees(n, span=span)
        for rank_tree, tsk_tree in zip(all_rank_trees, all_tsk_trees):
            ts1 = tsk_tree.tree_sequence
            ts2 = rank_tree.to_tsk_tree(span=span).tree_sequence
            assert ts1.tables.equals(ts2.tables, ignore_provenance=True)

    def test_all_shapes_roundtrip(self):
        n = 5
        all_rank_tree_shapes = RankTree.all_unlabelled_trees(n)
        all_tsk_tree_shapes = tskit.all_tree_shapes(n)
        for rank_tree, tsk_tree in zip(all_rank_tree_shapes, all_tsk_tree_shapes):
            assert rank_tree.shape_equal(RankTree.from_tsk_tree(tsk_tree))

    def test_all_labellings_roundtrip(self):
        n = 5
        rank_tree = RankTree.unrank(n, (comb.num_shapes(n) - 1, 0))
        tsk_tree = rank_tree.to_tsk_tree()
        rank_tree_labellings = RankTree.all_labellings(rank_tree)
        tsk_tree_labellings = tskit.all_tree_labellings(tsk_tree)
        for rank_t, tsk_t in zip(rank_tree_labellings, tsk_tree_labellings):
            assert rank_t == RankTree.from_tsk_tree(tsk_t)

    @pytest.mark.parametrize("n", range(6))
    def test_unrank_labelled(self, n):
        for shape_rank, t in enumerate(RankTree.all_unlabelled_trees(n)):
            for label_rank, labelled_tree in enumerate(RankTree.all_labellings(t)):
                unranked = RankTree.unrank(n, (shape_rank, label_rank))
                assert labelled_tree == unranked

    @pytest.mark.parametrize("n", range(10))
    def test_unrank_unlabelled(self, n):
        for shape_rank in range(comb.num_shapes(n)):
            rank = Rank(shape_rank, 0)
            unranked = RankTree.unrank(n, rank)
            assert rank, unranked.rank()

            rank = (shape_rank, comb.num_labellings(n, shape_rank) - 1)
            unranked = RankTree.unrank(n, rank)
            assert rank, unranked.rank()

    def test_unrank_errors(self):
        self.verify_unrank_errors((-1, 0), 1)
        self.verify_unrank_errors((0, -1), 1)
        self.verify_unrank_errors((-1, 0), 2)
        self.verify_unrank_errors((0, -1), 2)
        self.verify_unrank_errors((-1, 0), 10)
        self.verify_unrank_errors((0, -1), 10)

        self.verify_unrank_errors((0, 1), 1)
        self.verify_unrank_errors((1, 0), 2)
        self.verify_unrank_errors((0, 1), 2)
        self.verify_unrank_errors((2, 0), 3)
        self.verify_unrank_errors((0, 1), 3)
        self.verify_unrank_errors((1, 3), 3)

        invalid_shape = (comb.num_shapes(10), 0)
        self.verify_unrank_errors(invalid_shape, 10)
        invalid_labelling = (0, comb.num_labellings(10, 0))
        self.verify_unrank_errors(invalid_labelling, 10)

    def verify_unrank_errors(self, rank, n):
        with pytest.raises(ValueError):
            RankTree.unrank(n, rank)
        with pytest.raises(ValueError):
            tskit.Tree.unrank(n, rank)

    @pytest.mark.parametrize("n", range(6))
    def test_shape_rank(self, n):
        for rank, tree in enumerate(RankTree.all_unlabelled_trees(n)):
            assert tree.shape_rank() == rank

    @pytest.mark.parametrize("n", range(6))
    def test_shape_unrank(self, n):
        for rank, tree in enumerate(RankTree.all_unlabelled_trees(n)):
            t = RankTree.shape_unrank(n, rank)
            assert tree.shape_equal(t)

    @pytest.mark.parametrize("n", range(2, 9))
    def test_shape_unrank_tsk_tree(self, n):
        for shape_rank, tree in enumerate(RankTree.all_unlabelled_trees(n)):
            tsk_tree = tskit.Tree.unrank(n, (shape_rank, 0))
            assert shape_rank == tree.shape_rank()
            shape_rank, _ = tsk_tree.rank()
            assert shape_rank == tree.shape_rank()

    @pytest.mark.parametrize("n", range(7))
    def test_label_rank(self, n):
        for tree in RankTree.all_unlabelled_trees(n):
            for rank, labelled_tree in enumerate(RankTree.all_labellings(tree)):
                assert labelled_tree.label_rank() == rank

    @pytest.mark.parametrize("n", range(7))
    def test_label_unrank(self, n):
        for shape_rank, tree in enumerate(RankTree.all_unlabelled_trees(n)):
            for label_rank, labelled_tree in enumerate(RankTree.all_labellings(tree)):
                rank = (shape_rank, label_rank)
                unranked = tree.label_unrank(label_rank)
                assert labelled_tree.rank() == rank
                assert unranked.rank() == rank

    def test_rank_names(self):
        shape = 1
        label = 0
        n = 3
        tree = tskit.Tree.unrank(n, (shape, label))
        rank = tree.rank()
        assert rank.shape == shape
        assert rank.label == label

    @pytest.mark.parametrize("n", range(6))
    def test_unrank_rank_round_trip(self, n):
        for shape_rank in range(comb.num_shapes(n)):
            tree = RankTree.shape_unrank(n, shape_rank)
            tree = tree.label_unrank(0)
            assert tree.shape_rank() == shape_rank
            for label_rank in range(tree.num_labellings()):
                tree = tree.label_unrank(label_rank)
                assert tree.label_rank() == label_rank
                tsk_tree = tree.label_unrank(label_rank).to_tsk_tree()
                _, tsk_label_rank = tsk_tree.rank()
                assert tsk_label_rank == label_rank

    def test_is_canonical(self):
        for n in range(7):
            for tree in RankTree.all_labelled_trees(n):
                assert tree.is_canonical()

        shape_not_canonical = RankTree(
            children=[
                RankTree(children=[], label=0),
                RankTree(
                    children=[
                        RankTree(
                            children=[
                                RankTree(children=[], label=1),
                                RankTree(children=[], label=2),
                            ]
                        ),
                        RankTree(children=[], label=3),
                    ]
                ),
            ]
        )
        assert not shape_not_canonical.is_canonical()

        labels_not_canonical = RankTree(
            children=[
                RankTree(children=[], label=0),
                RankTree(
                    children=[
                        RankTree(
                            children=[
                                RankTree(children=[], label=2),
                                RankTree(children=[], label=3),
                            ]
                        ),
                        RankTree(
                            children=[
                                RankTree(children=[], label=1),
                                RankTree(children=[], label=4),
                            ]
                        ),
                    ]
                ),
            ]
        )
        assert not labels_not_canonical.is_canonical()

    @pytest.mark.parametrize("n", range(7))
    def test_unranking_is_canonical(self, n):
        for shape_rank in range(comb.num_shapes(n)):
            for label_rank in range(comb.num_labellings(n, shape_rank)):
                t = RankTree.shape_unrank(n, shape_rank)
                assert t.is_canonical()
                t = t.label_unrank(label_rank)
                assert t.is_canonical()
                t = tskit.Tree.unrank(n, (shape_rank, label_rank))
                assert RankTree.from_tsk_tree(t).is_canonical()

    @pytest.mark.parametrize("n", range(5))
    def test_to_from_tsk_tree(self, n):
        for tree in RankTree.all_labelled_trees(n):
            assert tree.is_canonical()
            tsk_tree = tree.to_tsk_tree()
            reconstructed = RankTree.from_tsk_tree(tsk_tree)
            assert tree.is_canonical()
            assert tree == reconstructed

    @pytest.mark.parametrize("n", range(6))
    def test_to_tsk_tree_internal_nodes(self, n):
        branch_length = 1234
        for tree in RankTree.all_labelled_trees(n):
            tsk_tree = tree.to_tsk_tree(branch_length=branch_length)
            internal_nodes = [
                u for u in tsk_tree.nodes(order="postorder") if tsk_tree.is_internal(u)
            ]
            assert np.all(internal_nodes == n + np.arange(len(internal_nodes)))
            for u in tsk_tree.nodes():
                if tsk_tree.is_internal(u):
                    max_child_time = max(tsk_tree.time(v) for v in tsk_tree.children(u))
                    assert tsk_tree.time(u) == max_child_time + branch_length
                else:
                    assert tsk_tree.time(u) == 0

    def test_from_unary_tree(self):
        tables = tskit.TableCollection(sequence_length=1)
        c = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        p = tables.nodes.add_row(time=1)
        tables.edges.add_row(left=0, right=1, parent=p, child=c)

        t = tables.tree_sequence().first()
        with pytest.raises(ValueError):
            RankTree.from_tsk_tree(t)

    def test_to_tsk_tree_errors(self):
        alpha_tree = RankTree.unrank(3, (0, 0), ["A", "B", "C"])
        out_of_bounds_tree = RankTree.unrank(3, (0, 0), [2, 3, 4])
        with pytest.raises(ValueError):
            alpha_tree.to_tsk_tree()
        with pytest.raises(ValueError):
            out_of_bounds_tree.to_tsk_tree()

    def test_rank_errors_multiple_roots(self):
        tables = tskit.TableCollection(sequence_length=1.0)

        # Nodes
        sv = [True, True]
        tv = [0.0, 0.0]

        for is_sample, t in zip(sv, tv):
            flags = tskit.NODE_IS_SAMPLE if is_sample else 0
            tables.nodes.add_row(flags=flags, time=t)

        ts = tables.tree_sequence()
        with pytest.raises(ValueError):
            ts.first().rank()

    def test_span(self):
        n = 5
        span = 8
        # Create a start tree, with a single root
        tsk_tree = tskit.Tree.unrank(n, (0, 0), span=span)
        assert tsk_tree.tree_sequence.num_nodes == n + 1
        assert tsk_tree.interval.left == 0
        assert tsk_tree.interval.right == span
        assert tsk_tree.tree_sequence.sequence_length == span

    def test_big_trees(self):
        n = 14
        shape = 22
        labelling = 0
        tree = RankTree.unrank(n, (shape, labelling))
        tsk_tree = tskit.Tree.unrank(n, (shape, labelling))
        assert tree.rank() == tsk_tree.rank()

        n = 10
        shape = 95
        labelling = comb.num_labellings(n, shape) // 2
        tree = RankTree.unrank(n, (shape, labelling))
        tsk_tree = tskit.Tree.unrank(n, (shape, labelling))
        assert tree.rank() == tsk_tree.rank()

    def test_symmetrical_trees(self):
        for n in range(2, 18, 2):
            last_rank = comb.num_shapes(n) - 1
            t = RankTree.shape_unrank(n, last_rank)
            assert t.is_symmetrical()

    def test_equal(self):
        unlabelled_leaf = RankTree(children=[])
        assert unlabelled_leaf == unlabelled_leaf
        assert unlabelled_leaf.shape_equal(unlabelled_leaf)

        leaf_zero = RankTree(children=[], label=0)
        leaf_one = RankTree(children=[], label=1)
        leaf_two = RankTree(children=[], label=2)
        assert leaf_zero == leaf_zero
        assert leaf_zero != leaf_one
        assert leaf_zero.shape_equal(leaf_one)

        tree1 = RankTree(children=[leaf_zero, leaf_one])
        assert tree1 == tree1
        assert tree1 != unlabelled_leaf
        assert not tree1.shape_equal(unlabelled_leaf)

        tree2 = RankTree(children=[leaf_two, leaf_one])
        assert tree1 != tree2
        assert tree1.shape_equal(tree2)

    def test_is_symmetrical(self):
        unlabelled_leaf = RankTree(children=[])
        assert unlabelled_leaf.is_symmetrical()
        three_leaf_asym = RankTree(
            children=[
                unlabelled_leaf,
                RankTree(children=[unlabelled_leaf, unlabelled_leaf]),
            ]
        )
        assert not three_leaf_asym.is_symmetrical()
        six_leaf_sym = RankTree(children=[three_leaf_asym, three_leaf_asym])
        assert six_leaf_sym.is_symmetrical()


class TestPartialTopologyCounter:
    def test_add_sibling_topologies_simple(self):
        a = RankTree(children=[], label="A")
        b = RankTree(children=[], label="B")
        ab = RankTree(children=[a, b])

        a_counter = comb.TopologyCounter()
        a_counter["A"][a.rank()] = 1
        assert a_counter == comb.TopologyCounter.from_sample("A")

        b_counter = comb.TopologyCounter()
        b_counter["B"][b.rank()] = 1
        assert b_counter == comb.TopologyCounter.from_sample("B")

        partial_counter = comb.PartialTopologyCounter()
        partial_counter.add_sibling_topologies(a_counter)
        partial_counter.add_sibling_topologies(b_counter)

        expected = comb.TopologyCounter()
        expected["A"][a.rank()] = 1
        expected["B"][b.rank()] = 1
        expected["A", "B"][ab.rank()] = 1
        joined_counter = partial_counter.join_all_combinations()
        assert joined_counter == expected

    def test_add_sibling_topologies_polytomy(self):
        """
        Goes through the topology-merging step at the root
        of this tree:
                    |
                    |
            +----+-----+----+
            |    |     |    |
            |    |     |    |
            |    |     |  +---+
            |    |     |  |   |
            |    |     |  |   |
            A    A     B  A   C
        """
        partial_counter = comb.PartialTopologyCounter()
        a = RankTree(children=[], label="A")
        c = RankTree(children=[], label="C")
        ac = RankTree(children=[a, c])

        expected = collections.defaultdict(collections.Counter)

        a_counter = comb.TopologyCounter.from_sample("A")
        b_counter = comb.TopologyCounter.from_sample("B")
        ac_counter = comb.TopologyCounter()
        ac_counter["A"][a.rank()] = 1
        ac_counter["C"][c.rank()] = 1
        ac_counter["A", "C"][ac.rank()] = 1

        partial_counter.add_sibling_topologies(a_counter)
        expected[("A",)] = collections.Counter({((("A",), (0, 0)),): 1})
        assert partial_counter.partials == expected

        partial_counter.add_sibling_topologies(a_counter)
        expected[("A",)][((("A",), (0, 0)),)] += 1
        assert partial_counter.partials == expected

        partial_counter.add_sibling_topologies(b_counter)
        expected[("B",)][((("B",), (0, 0)),)] = 1
        expected[("A", "B")][((("A",), (0, 0)), (("B",), (0, 0)))] = 2
        assert partial_counter.partials == expected

        partial_counter.add_sibling_topologies(ac_counter)
        expected[("A",)][((("A",), (0, 0)),)] += 1
        expected[("C",)][((("C",), (0, 0)),)] = 1
        expected[("A", "B")][((("A",), (0, 0)), (("B",), (0, 0)))] += 1
        expected[("A", "C")][((("A",), (0, 0)), (("C",), (0, 0)))] = 2
        expected[("A", "C")][((("A", "C"), (0, 0)),)] = 1
        expected[("B", "C")][((("B",), (0, 0)), (("C",), (0, 0)))] = 1
        expected[("A", "B", "C")][
            ((("A",), (0, 0)), (("B",), (0, 0)), (("C",), (0, 0)))
        ] = 2
        expected[("A", "B", "C")][((("A", "C"), (0, 0)), (("B",), (0, 0)))] = 1
        assert partial_counter.partials == expected

        expected_topologies = comb.TopologyCounter()
        expected_topologies["A"][(0, 0)] = 3
        expected_topologies["B"][(0, 0)] = 1
        expected_topologies["C"][(0, 0)] = 1
        expected_topologies["A", "B"][(0, 0)] = 3
        expected_topologies["A", "C"][(0, 0)] = 3
        expected_topologies["B", "C"][(0, 0)] = 1
        expected_topologies["A", "B", "C"][(0, 0)] = 2
        expected_topologies["A", "B", "C"][(1, 1)] = 1
        joined_topologies = partial_counter.join_all_combinations()
        assert joined_topologies == expected_topologies

    def test_join_topologies(self):
        a = RankTree(children=[], label="A")
        b = RankTree(children=[], label="B")
        c = RankTree(children=[], label="C")
        a_tuple = (("A"), a.rank())
        b_tuple = (("B"), b.rank())
        c_tuple = (("C"), c.rank())
        ab_tuple = (("A", "B"), RankTree(children=[a, b]).rank())
        ac_tuple = (("A", "C"), RankTree(children=[a, c]).rank())
        bc_tuple = (("B", "C"), RankTree(children=[b, c]).rank())

        self.verify_join_topologies((a_tuple, b_tuple), (0, 0))
        self.verify_join_topologies((b_tuple, a_tuple), (0, 0))
        self.verify_join_topologies((b_tuple, c_tuple), (0, 0))

        self.verify_join_topologies((a_tuple, b_tuple, c_tuple), (0, 0))
        self.verify_join_topologies((a_tuple, bc_tuple), (1, 0))
        self.verify_join_topologies((b_tuple, ac_tuple), (1, 1))
        self.verify_join_topologies((c_tuple, ab_tuple), (1, 2))

    def verify_join_topologies(self, topologies, expected_topology):
        actual_topology = comb.PartialTopologyCounter.join_topologies(topologies)
        assert actual_topology == expected_topology


class TestCountTopologies:
    def verify_topologies(self, ts, sample_sets=None, expected=None):
        if sample_sets is None:
            sample_sets = [ts.samples(population=pop.id) for pop in ts.populations()]
        topologies = [t.count_topologies(sample_sets) for t in ts.trees()]
        inc_topologies = list(ts.count_topologies(sample_sets))
        # count_topologies calculates the embedded topologies for every
        # combination of populations, so we need to check the results
        # of subsampling for every combination.
        for num_sample_sets in range(1, len(sample_sets) + 1):
            for i, t in enumerate(ts.trees()):
                just_t = ts.keep_intervals([t.interval], simplify=False)
                for sample_set_indexes in itertools.combinations(
                    range(len(sample_sets)), num_sample_sets
                ):
                    actual_topologies = topologies[i][sample_set_indexes]
                    actual_inc_topologies = inc_topologies[i][sample_set_indexes]
                    if len(t.roots) == 1:
                        subsampled_topologies = self.subsample_topologies(
                            just_t, sample_sets, sample_set_indexes
                        )
                        assert actual_topologies == subsampled_topologies
                    if expected is not None:
                        assert actual_topologies == expected[i][sample_set_indexes]
                    assert actual_topologies == actual_inc_topologies

    def subsample_topologies(self, ts, sample_sets, sample_set_indexes):
        subsample_sets = [sample_sets[i] for i in sample_set_indexes]
        topologies = collections.Counter()
        for subsample in itertools.product(*subsample_sets):
            for pop_tree in ts.simplify(samples=subsample).trees():
                # regions before and after keep interval have all samples as roots
                # so don't count those
                # The single tree of interest should have one root
                if len(pop_tree.roots) == 1:
                    topologies[pop_tree.rank()] += 1
        return topologies

    def test_single_population(self):
        n = 10
        ts = msprime.simulate(n, recombination_rate=10)
        expected = comb.TopologyCounter()
        expected[0] = collections.Counter({(0, 0): n})
        self.verify_topologies(ts, expected=[expected] * ts.num_trees)

    def test_three_populations(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time    population  individual  metadata
        0   1   0.000000    0   -1
        1   1   0.000000    1   -1
        2   1   0.000000    1   -1
        3   1   0.000000    2   -1
        4   1   0.000000    2   -1
        5   1   0.000000    0   -1
        6   0   1.000000    0   -1
        7   0   2.000000    0   -1
        8   0   2.000000    0   -1
        9   0   3.000000    0   -1
        10  0   4.000000    0   -1
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.000000    1.000000    6  4
        0.000000    1.000000    6  5
        0.000000    1.000000    7  1
        0.000000    1.000000    7  2
        0.000000    1.000000    8  3
        0.000000    1.000000    8  6
        0.000000    1.000000    9  7
        0.000000    1.000000    9  8
        0.000000    1.000000    10  0
        0.000000    1.000000    10  9
        """
        )
        ts = tskit.load_text(
            nodes, edges, sequence_length=1, strict=False, base64_metadata=False
        )

        expected = comb.TopologyCounter()
        expected[0] = collections.Counter({(0, 0): 2})
        expected[1] = collections.Counter({(0, 0): 2})
        expected[2] = collections.Counter({(0, 0): 2})
        expected[0, 1] = collections.Counter({(0, 0): 4})
        expected[0, 2] = collections.Counter({(0, 0): 4})
        expected[1, 2] = collections.Counter({(0, 0): 4})
        expected[0, 1, 2] = collections.Counter({(1, 0): 4, (1, 1): 4})
        self.verify_topologies(ts, expected=[expected])

    def test_multiple_roots(self):
        tables = tskit.TableCollection(sequence_length=1.0)
        tables.populations.add_row()
        tables.populations.add_row()
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, population=0)
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, population=1)

        # Not samples so they are ignored
        tables.nodes.add_row(time=1)
        tables.nodes.add_row(time=1, population=1)

        expected = comb.TopologyCounter()
        expected[0] = collections.Counter({(0, 0): 1})
        expected[1] = collections.Counter({(0, 0): 1})
        self.verify_topologies(tables.tree_sequence(), expected=[expected])

    def test_no_sample_subtrees(self):
        tables = tskit.TableCollection(sequence_length=1.0)
        c1 = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        c2 = tables.nodes.add_row(time=0)
        c3 = tables.nodes.add_row(time=0)
        p1 = tables.nodes.add_row(time=1)
        p2 = tables.nodes.add_row(time=1)

        tables.edges.add_row(left=0, right=1, parent=p1, child=c2)
        tables.edges.add_row(left=0, right=1, parent=p1, child=c3)
        tables.edges.add_row(left=0, right=1, parent=p2, child=c1)

        expected = comb.TopologyCounter()
        expected[0] = collections.Counter({(0, 0): 1})
        self.verify_topologies(tables.tree_sequence(), expected=[expected])

    def test_no_full_topology(self):
        tables = tskit.TableCollection(sequence_length=1.0)
        tables.populations.add_row()
        tables.populations.add_row()
        tables.populations.add_row()
        child1 = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, population=0)
        child2 = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, population=1)
        parent = tables.nodes.add_row(time=1)
        tables.edges.add_row(left=0, right=1, parent=parent, child=child1)
        tables.edges.add_row(left=0, right=1, parent=parent, child=child2)

        # Left as root so there is no topology with all three populations
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, population=2)

        expected = comb.TopologyCounter()
        for pop_combo in [(0,), (1,), (2,), (0, 1)]:
            expected[pop_combo] = collections.Counter({(0, 0): 1})
        self.verify_topologies(tables.tree_sequence(), expected=[expected])

    def test_polytomies(self):
        tables = tskit.TableCollection(sequence_length=1.0)
        tables.populations.add_row()
        tables.populations.add_row()
        tables.populations.add_row()
        c1 = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, population=0)
        c2 = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, population=1)
        c3 = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, population=2)
        p = tables.nodes.add_row(time=1)
        tables.edges.add_row(left=0, right=1, parent=p, child=c1)
        tables.edges.add_row(left=0, right=1, parent=p, child=c2)
        tables.edges.add_row(left=0, right=1, parent=p, child=c3)

        expected = comb.TopologyCounter()
        for pop_combos in [0, 1, 2, (0, 1), (0, 2), (1, 2), (0, 1, 2)]:
            expected[pop_combos] = collections.Counter({(0, 0): 1})
        self.verify_topologies(tables.tree_sequence(), expected=[expected])

    def test_custom_key(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time    population  individual  metadata
        0   1   0.000000    0   -1
        1   1   0.000000    0   -1
        2   1   0.000000    0   -1
        3   1   0.000000    0   -1
        4   1   0.000000    0   -1
        5   0   1.000000    0   -1
        6   0   1.000000    0   -1
        7   0   2.000000    0   -1
        8   0   3.000000    0   -1
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.000000    1.000000    5  0
        0.000000    1.000000    5  1
        0.000000    1.000000    6  2
        0.000000    1.000000    6  3
        0.000000    1.000000    7  5
        0.000000    1.000000    7  6
        0.000000    1.000000    8  4
        0.000000    1.000000    8  7
        """
        )
        ts = tskit.load_text(
            nodes, edges, sequence_length=1, strict=False, base64_metadata=False
        )

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

        expected = comb.TopologyCounter()
        expected[0] = collections.Counter({(0, 0): 2})
        expected[1] = collections.Counter({(0, 0): 2})
        expected[2] = collections.Counter({(0, 0): 1})
        expected[0, 1] = collections.Counter({(0, 0): 4})
        expected[0, 2] = collections.Counter({(0, 0): 2})
        expected[1, 2] = collections.Counter({(0, 0): 2})
        expected[0, 1, 2] = collections.Counter({(1, 2): 4})

        tree_topologies = ts.first().count_topologies(sample_sets)
        treeseq_topologies = list(ts.count_topologies(sample_sets))
        assert tree_topologies == expected
        assert treeseq_topologies == [expected]

    def test_ignores_non_sample_leaves(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time    population  individual  metadata
        0   1   0.000000    0   -1
        1   0   0.000000    0   -1
        2   1   0.000000    0   -1
        3   0   0.000000    0   -1
        4   1   0.000000    0   -1
        5   0   1.000000    0   -1
        6   0   1.000000    0   -1
        7   0   2.000000    0   -1
        8   0   3.000000    0   -1
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.000000    1.000000    5  0
        0.000000    1.000000    5  1
        0.000000    1.000000    6  2
        0.000000    1.000000    6  3
        0.000000    1.000000    7  5
        0.000000    1.000000    7  6
        0.000000    1.000000    8  4
        0.000000    1.000000    8  7
        """
        )
        ts = tskit.load_text(
            nodes, edges, sequence_length=1, strict=False, base64_metadata=False
        )

        sample_sets = [[0], [2], [4]]

        expected = comb.TopologyCounter()
        expected[0] = collections.Counter({(0, 0): 1})
        expected[1] = collections.Counter({(0, 0): 1})
        expected[2] = collections.Counter({(0, 0): 1})
        expected[0, 1] = collections.Counter({(0, 0): 1})
        expected[0, 2] = collections.Counter({(0, 0): 1})
        expected[1, 2] = collections.Counter({(0, 0): 1})
        expected[0, 1, 2] = collections.Counter({(1, 2): 1})

        tree_topologies = ts.first().count_topologies(sample_sets)
        treeseq_topologies = list(ts.count_topologies(sample_sets))
        assert tree_topologies == expected
        assert treeseq_topologies == [expected]

    def test_internal_samples_errors(self):
        tables = tskit.TableCollection(sequence_length=1.0)

        c1 = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        c2 = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        p = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=1)

        tables.edges.add_row(left=0, right=1, parent=p, child=c1)
        tables.edges.add_row(left=0, right=1, parent=p, child=c2)

        self.verify_value_error(tables.tree_sequence())

    def test_non_sample_nodes_errors(self):
        tables = tskit.TableCollection(sequence_length=1.0)

        c1 = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        c2 = tables.nodes.add_row(time=0)
        p = tables.nodes.add_row(time=1)

        tables.edges.add_row(left=0, right=1, parent=p, child=c1)
        tables.edges.add_row(left=0, right=1, parent=p, child=c2)

        sample_sets = [[0], [1]]
        self.verify_value_error(tables.tree_sequence(), sample_sets)

        sample_sets = [[0], [tables.nodes.num_rows]]
        self.verify_node_out_of_bounds_error(tables.tree_sequence(), sample_sets)

    def verify_value_error(self, ts, sample_sets=None):
        with pytest.raises(ValueError):
            ts.first().count_topologies(sample_sets)
        with pytest.raises(ValueError):
            list(ts.count_topologies(sample_sets))

    def verify_node_out_of_bounds_error(self, ts, sample_sets=None):
        with pytest.raises(ValueError):
            ts.first().count_topologies(sample_sets)
        with pytest.raises(IndexError):
            list(ts.count_topologies(sample_sets))

    def test_standard_msprime_migrations(self):
        for num_populations in range(2, 5):
            samples = [5] * num_populations
            ts = self.simulate_multiple_populations(samples)
            self.verify_topologies(ts)

    def simulate_multiple_populations(self, sample_sizes):
        d = len(sample_sizes)
        M = 0.2
        m = M / (2 * (d - 1))

        migration_matrix = [
            [m if k < d and k == i + 1 else 0 for k in range(d)] for i in range(d)
        ]

        pop_configurations = [
            msprime.PopulationConfiguration(sample_size=size) for size in sample_sizes
        ]
        return msprime.simulate(
            population_configurations=pop_configurations,
            migration_matrix=migration_matrix,
            recombination_rate=0.1,
        )

    def test_msprime_dtwf(self):
        migration_matrix = np.zeros((4, 4))
        population_configurations = [
            msprime.PopulationConfiguration(
                sample_size=10, initial_size=10, growth_rate=0
            ),
            msprime.PopulationConfiguration(
                sample_size=10, initial_size=10, growth_rate=0
            ),
            msprime.PopulationConfiguration(
                sample_size=10, initial_size=10, growth_rate=0
            ),
            msprime.PopulationConfiguration(
                sample_size=0, initial_size=10, growth_rate=0
            ),
        ]
        demographic_events = [
            msprime.PopulationParametersChange(population=1, time=0.1, initial_size=5),
            msprime.PopulationParametersChange(population=0, time=0.2, initial_size=5),
            msprime.MassMigration(time=1.1, source=0, dest=2),
            msprime.MassMigration(time=1.2, source=1, dest=3),
            msprime.MigrationRateChange(time=2.1, rate=0.3, matrix_index=(2, 3)),
            msprime.MigrationRateChange(time=2.2, rate=0.3, matrix_index=(3, 2)),
        ]
        ts = msprime.simulate(
            migration_matrix=migration_matrix,
            population_configurations=population_configurations,
            demographic_events=demographic_events,
            random_seed=2,
            model="dtwf",
        )

        self.verify_topologies(ts)

    def test_forward_time_wright_fisher_unsimplified_all_sample_sets(self):
        tables = wf.wf_sim(
            4,
            5,
            seed=1,
            deep_history=False,
            initial_generation_samples=False,
            num_loci=10,
        )
        tables.sort()
        ts = tables.tree_sequence()
        for S in test_stats.set_partitions(list(ts.samples())):
            self.verify_topologies(ts, sample_sets=S)

    def test_forward_time_wright_fisher_unsimplified(self):
        tables = wf.wf_sim(
            20,
            15,
            seed=1,
            deep_history=False,
            initial_generation_samples=False,
            num_loci=20,
        )
        tables.sort()
        ts = tables.tree_sequence()
        samples = ts.samples()
        self.verify_topologies(ts, sample_sets=[samples[:10], samples[10:]])

    def test_forward_time_wright_fisher_simplified(self):
        tables = wf.wf_sim(
            30,
            10,
            seed=1,
            deep_history=False,
            initial_generation_samples=False,
            num_loci=5,
        )
        tables.sort()
        ts = tables.tree_sequence()
        samples = ts.samples()
        self.verify_topologies(ts, sample_sets=[samples[:10], samples[10:]])


class TestTreeNode:
    """
    Tests for the TreeNode class used to build simple trees in memory.
    """

    def verify_tree(self, root, labels):
        # Note this doesn't check any statistical properties of the returned
        # trees, just that a single instance returned in a valid binary tree.
        # Structural properties are best verified using the tskit API, and so
        # we test these properties elsewhere.
        stack = [root]
        num_nodes = 0
        recovered_labels = []
        while len(stack) > 0:
            node = stack.pop()
            num_nodes += 1
            if node.label is not None:
                assert len(node.children) == 0
                recovered_labels.append(node.label)
            for child in node.children:
                assert child.parent == node
                stack.append(child)
        assert sorted(recovered_labels) == list(labels)

    @pytest.mark.parametrize("n", range(1, 16))
    def test_random_binary_tree(self, n):
        rng = random.Random(32)
        labels = range(n)
        root = comb.TreeNode.random_binary_tree(labels, rng)
        self.verify_tree(root, range(n))

    @pytest.mark.parametrize("n", range(1, 16))
    def test_balanced_binary(self, n):
        root = comb.TreeNode.balanced_tree(range(n), 2)
        self.verify_tree(root, range(n))

    @pytest.mark.parametrize("arity", range(2, 8))
    def test_balanced_arity(self, arity):
        labels = range(30)
        root = comb.TreeNode.balanced_tree(labels, arity)
        self.verify_tree(root, labels)


def num_leaf_labelled_binary_trees(n):
    """
    Returns the number of leaf labelled binary trees with n leaves.

    TODO: this would probably be helpful to have in the combinatorics
    module.

    https://oeis.org/A005373/
    """
    return int(math.factorial(2 * n - 3) / (2 ** (n - 2) * math.factorial(n - 2)))


class TestPolytomySplitting:
    """
    Test the ability to randomly split polytomies
    """

    # A complex ts with polytomies
    #
    # 1.00┊    6      ┊      6    ┊       6   ┊           ┊      6    ┊
    #     ┊ ┏━┳┻┳━┓   ┊   ┏━┳┻┳━┓ ┊    ┏━━╋━┓ ┊           ┊   ┏━┳┻┳━┓ ┊
    # 0.50┊ 5 ┃ ┃ ┃   ┊   5 ┃ ┃ ┃ ┊    5  ┃ ┃ ┊      5    ┊   ┃ ┃ ┃ ┃ ┊
    #     ┊ ┃ ┃ ┃ ┃ . ┊   ┃ ┃ ┃ ┃ ┊ . ┏┻┓ ┃ ┃ ┊ . ┏━┳┻┳━┓ ┊ . ┃ ┃ ┃ ┃ ┊
    # 0.00┊ 0 2 3 4 1 ┊ 0 1 2 3 4 ┊ 0 1 2 3 4 ┊ 0 1 2 3 4 ┊ 0 1 2 3 4 ┊
    #   0.00        0.20        0.40        0.60        0.80        1.00
    nodes_polytomy_44344 = """\
    id      is_sample   population      time
    0       1           0               0.0
    1       1           0               0.0
    2       1           0               0.0
    3       1           0               0.0
    4       1           0               0.0
    5       0           0               0.5
    6       0           0               1.0
    """
    edges_polytomy_44344 = """\
    id      left     right    parent  child
    0       0.0      0.2      5       0
    1       0.0      0.8      5       1
    2       0.0      0.4      6       2
    3       0.4      0.8      5       2
    4       0.0      0.6      6       3,4
    5       0.0      0.6      6       5
    6       0.6      0.8      5       3,4
    7       0.8      1.0      6       1,2,3,4
    """

    def ts_polytomy_44344(self):
        return tskit.load_text(
            nodes=io.StringIO(self.nodes_polytomy_44344),
            edges=io.StringIO(self.edges_polytomy_44344),
            strict=False,
        )

    def verify_trees(self, source_tree, split_tree, epsilon=None):
        N = 0
        for u in split_tree.nodes():
            assert split_tree.num_children(u) < 3
            N += 1
            if u >= source_tree.tree_sequence.num_nodes:
                # This is a new node
                branch_length = split_tree.branch_length(u)
                if epsilon is not None:
                    assert epsilon == pytest.approx(branch_length)
                else:
                    assert branch_length > 0
                    assert 0 == pytest.approx(branch_length)

        assert N == len(list(split_tree.leaves())) * 2 - 1
        for u in source_tree.nodes():
            if source_tree.num_children(u) <= 2:
                assert source_tree.children(u) == split_tree.children(u)
            else:
                assert len(split_tree.children(u)) == 2

    @pytest.mark.parametrize("n", [2, 3, 4, 5, 6])
    def test_resolve_star(self, n):
        tree = tskit.Tree.generate_star(n)
        self.verify_trees(tree, tree.split_polytomies(random_seed=12))

    def test_large_epsilon(self):
        tree = tskit.Tree.generate_star(10, branch_length=100)
        eps = 10
        split = tree.split_polytomies(random_seed=12234, epsilon=eps)
        self.verify_trees(tree, split, epsilon=eps)

    def test_small_epsilon(self):
        tree = tskit.Tree.generate_star(10, branch_length=1e-20)
        eps = 1e-22
        split = tree.split_polytomies(random_seed=12234, epsilon=eps)
        self.verify_trees(tree, split, epsilon=eps)

    def test_nextafter_near_zero(self):
        tree = tskit.Tree.generate_star(3, branch_length=np.finfo(float).tiny)
        split = tree.split_polytomies(random_seed=234)
        self.verify_trees(tree, split)

    def test_nextafter_large_tree(self):
        tree = tskit.Tree.generate_star(100)
        split = tree.split_polytomies(random_seed=32)
        self.verify_trees(tree, split)
        for u in tree.nodes():
            if tree.parent(u) != tskit.NULL and not tree.is_leaf(u):
                parent_time = tree.time(tree.parent(u))
                child_time = tree.time(u)
                assert child_time == np.nextafter(parent_time, 0)
            if tree.is_leaf(u):
                assert tree.branch_length(u) == pytest.approx(1)

    def test_epsilon_near_one(self):
        tree = tskit.Tree.generate_star(3, branch_length=1)
        split = tree.split_polytomies(random_seed=234, epsilon=np.finfo(float).eps)
        self.verify_trees(tree, split)

    def verify_tree_sequence_splits(self, ts):
        n_poly = 0
        for e in ts.edgesets():
            if len(e.children) > 2:
                n_poly += 1
        assert n_poly > 3
        assert ts.num_trees > 3
        for tree in ts.trees():
            binary_tree = tree.split_polytomies(random_seed=11)
            assert binary_tree.interval == tree.interval
            for u in binary_tree.nodes():
                assert binary_tree.num_children(u) < 3
            for u in tree.nodes():
                assert binary_tree.time(u) == tree.time(u)
            resolved_ts = binary_tree.tree_sequence
            assert resolved_ts.sequence_length == ts.sequence_length
            assert resolved_ts.num_trees <= 3
            if tree.interval.left == 0:
                assert resolved_ts.num_trees == 2
                null_tree = resolved_ts.last()
                assert null_tree.num_roots == ts.num_samples
            elif tree.interval.right == ts.sequence_length:
                assert resolved_ts.num_trees == 2
                null_tree = resolved_ts.first()
                assert null_tree.num_roots == ts.num_samples
            else:
                null_tree = resolved_ts.first()
                assert null_tree.num_roots == ts.num_samples
                null_tree.next()
                assert null_tree.num_roots == tree.num_roots
                null_tree.next()
                assert null_tree.num_roots == ts.num_samples

    def test_complex_examples(self):
        self.verify_tree_sequence_splits(self.ts_polytomy_44344())

    def test_nonbinary_simulation(self):
        demographic_events = [
            msprime.SimpleBottleneck(time=1.0, population=0, proportion=0.95)
        ]
        ts = msprime.simulate(
            20,
            recombination_rate=10,
            mutation_rate=5,
            demographic_events=demographic_events,
            random_seed=7,
        )
        self.verify_tree_sequence_splits(ts)

    def test_seeds(self):
        base = tskit.Tree.generate_star(5)
        t1 = base.split_polytomies(random_seed=1234)
        t2 = base.split_polytomies(random_seed=1234)
        assert t1.tree_sequence.tables.equals(
            t2.tree_sequence.tables, ignore_timestamps=True
        )
        t2 = base.split_polytomies(random_seed=1)
        assert not t1.tree_sequence.tables.equals(
            t2.tree_sequence.tables, ignore_provenance=True
        )

    def test_internal_polytomy(self):
        #       9
        # ┏━┳━━━┻┳━━━━┓
        # ┃ ┃    8    ┃
        # ┃ ┃ ┏━━╋━━┓ ┃
        # ┃ ┃ ┃  7  ┃ ┃
        # ┃ ┃ ┃ ┏┻┓ ┃ ┃
        # 0 1 2 3 5 4 6
        t1 = tskit.Tree.unrank(7, (6, 25))
        t2 = t1.split_polytomies(random_seed=1234)
        assert t2.parent(3) == 7
        assert t2.parent(5) == 7
        assert t2.root == 9
        for u in t2.nodes():
            assert t2.num_children(u) in [0, 2]

    def test_binary_tree(self):
        t1 = msprime.simulate(10, random_seed=1234).first()
        t2 = t1.split_polytomies(random_seed=1234)
        tables = t1.tree_sequence.dump_tables()
        tables.assert_equals(t2.tree_sequence.tables, ignore_provenance=True)

    def test_bad_method(self):
        tree = tskit.Tree.generate_star(3)
        with pytest.raises(ValueError, match="Method"):
            tree.split_polytomies(method="something_else")

    @pytest.mark.parametrize("epsilon", [10, 1.1, 1.0])
    def test_epsilon_too_large(self, epsilon):
        tree = tskit.Tree.generate_star(3)
        msg = (
            "Cannot resolve the degree 3 polytomy rooted at node 3 "
            "with minimum time difference of 1.0 to the resolved leaves. "
            f"The fixed epsilon value of {epsilon} is too large, resulting in the "
            "parent time being less than the child time."
        )
        with pytest.raises(
            tskit.LibraryError,
            match=msg,
        ):
            tree.split_polytomies(epsilon=epsilon, random_seed=12)

    def test_epsilon_too_small(self):
        tree = tskit.Tree.generate_star(3)
        msg = (
            "Cannot resolve the degree 3 polytomy rooted at node 3 "
            "with minimum time difference of 1.0 to the resolved leaves. "
            "The fixed epsilon value of 0 is too small, resulting in the "
            "parent and child times being equal within the limits of "
            "numerical precision."
        )
        with pytest.raises(
            tskit.LibraryError,
            match=msg,
        ):
            tree.split_polytomies(epsilon=0, random_seed=12)

    def test_unsplittable_branch(self):
        branch_length = np.nextafter(0, 1)
        tree = tskit.Tree.generate_star(3, branch_length=branch_length)
        msg = (
            "Cannot resolve the degree 3 polytomy rooted at node 3 with "
            "minimum time difference of 5e-324 to the resolved leaves. "
            "The time difference between nodes is so small that more nodes "
            "cannot be inserted between within the limits of floating point "
            "precision."
        )
        with pytest.raises(
            tskit.LibraryError,
            match=msg,
        ):
            tree.split_polytomies(random_seed=12)

    def test_epsilon_for_mutations(self):
        tables = tskit.Tree.generate_star(3).tree_sequence.dump_tables()
        root_time = tables.nodes.time[-1]
        assert root_time == 1
        site = tables.sites.add_row(position=0.5, ancestral_state="0")
        tables.mutations.add_row(site=site, time=0.9, node=0, derived_state="1")
        tables.mutations.add_row(site=site, time=0.9, node=1, derived_state="1")
        tree = tables.tree_sequence().first()
        with pytest.raises(
            tskit.LibraryError,
            match="not small enough to create new nodes below a polytomy",
        ):
            tree.split_polytomies(epsilon=0.5, random_seed=123)

    def test_mutation_within_eps_parent(self):
        tables = tskit.Tree.generate_star(3).tree_sequence.dump_tables()
        site = tables.sites.add_row(position=0.5, ancestral_state="0")
        branch_length = np.nextafter(1, 0)
        tables.mutations.add_row(
            site=site, time=branch_length, node=0, derived_state="1"
        )
        tables.mutations.add_row(
            site=site, time=branch_length, node=1, derived_state="1"
        )
        tree = tables.tree_sequence().first()
        with pytest.raises(
            tskit.LibraryError,
            match="Cannot split polytomy: mutation with numerical precision",
        ):
            tree.split_polytomies(random_seed=123)

    def test_provenance(self):
        tree = tskit.Tree.generate_star(4)
        ts_split = tree.split_polytomies(random_seed=14).tree_sequence
        record = json.loads(ts_split.provenance(ts_split.num_provenances - 1).record)
        assert record["parameters"]["command"] == "split_polytomies"
        ts_split = tree.split_polytomies(
            random_seed=12, record_provenance=False
        ).tree_sequence
        record = json.loads(ts_split.provenance(ts_split.num_provenances - 1).record)
        assert record["parameters"]["command"] != "split_polytomies"

    def test_kwargs(self):
        tree = tskit.Tree.generate_star(4)
        split_tree = tree.split_polytomies(random_seed=14, tracked_samples=[0, 1])
        assert split_tree.num_tracked_samples() == 2

    @pytest.mark.slow
    @pytest.mark.parametrize("n", [3, 4, 5])
    def test_all_topologies(self, n):
        N = num_leaf_labelled_binary_trees(n)
        ranks = collections.Counter()
        for seed in range(20 * N):
            star = tskit.Tree.generate_star(n)
            random_tree = star.split_polytomies(random_seed=seed)
            ranks[random_tree.rank()] += 1
        # There are N possible binary trees here, we should have seen them
        # all with high probability after 20 N attempts.
        assert len(ranks) == N


class TreeGeneratorTestBase:
    """
    Abstract superclass of tree generator test methods.

    Concrete subclasses should defined "method_name" class variable.
    """

    def method(self, n, **kwargs):
        return getattr(tskit.Tree, self.method_name)(n, **kwargs)

    @pytest.mark.parametrize("n", range(2, 10))
    def test_leaves(self, n):
        tree = self.method(n)
        assert list(tree.leaves()) == list(range(n))

    def test_bad_n(self):
        for n in [-1, 0, np.array([1, 2])]:
            with pytest.raises(ValueError):
                self.method(n)
        for n in [None, "", []]:
            with pytest.raises(TypeError):
                self.method(n)

    def test_bad_span(self):
        with pytest.raises(tskit.LibraryError):
            self.method(2, span=0)

    def test_bad_branch_length(self):
        with pytest.raises(tskit.LibraryError):
            self.method(2, branch_length=0)

    @pytest.mark.parametrize("span", [0.1, 1, 100])
    def test_span(self, span):
        tree = self.method(5, span=span)
        assert tree.tree_sequence.sequence_length == span

    @pytest.mark.parametrize("branch_length", [0.25, 1, 100])
    def test_branch_length(self, branch_length):
        tree = self.method(5, branch_length=branch_length)
        for u in tree.nodes():
            if u != tree.root:
                assert tree.branch_length(u) >= branch_length

    def test_provenance(self):
        ts = self.method(2).tree_sequence
        assert ts.num_provenances == 1
        record = json.loads(ts.provenance(0).record)
        assert record["parameters"]["command"] == self.method_name
        ts = self.method(2, record_provenance=False).tree_sequence
        assert ts.num_provenances == 0

    @pytest.mark.parametrize("n", range(2, 10))
    def test_rank_unrank_round_trip(self, n):
        tree1 = self.method(n)
        rank = tree1.rank()
        tree2 = tskit.Tree.unrank(n, rank)
        tables1 = tree1.tree_sequence.tables
        tables2 = tree2.tree_sequence.tables
        tables1.assert_equals(tables2, ignore_provenance=True)

    def test_kwargs(self):
        tree = self.method(3, tracked_samples=[0, 1])
        assert tree.num_tracked_samples() == 2


class TestGenerateStar(TreeGeneratorTestBase):
    method_name = "generate_star"

    @pytest.mark.parametrize("n", range(2, 10))
    def test_unrank_equal(self, n):
        for extra_params in [{}, {"span": 2.5}, {"branch_length": 3}]:
            ts = tskit.Tree.generate_star(n, **extra_params).tree_sequence
            equiv_ts = tskit.Tree.unrank(n, (0, 0), **extra_params).tree_sequence
            assert ts.tables.equals(equiv_ts.tables, ignore_provenance=True)

    def test_branch_length_semantics(self):
        branch_length = 10
        ts = tskit.Tree.generate_star(7, branch_length=branch_length).tree_sequence
        time = ts.tables.nodes.time
        edges = ts.tables.edges
        length = time[edges.parent] - time[edges.child]
        assert np.all(length == branch_length)


class TestGenerateBalanced(TreeGeneratorTestBase):
    method_name = "generate_balanced"

    @pytest.mark.parametrize("arity", range(2, 10))
    def test_arity_leaves(self, arity):
        n = 20
        tree = tskit.Tree.generate_balanced(n, arity=arity)
        assert list(tree.leaves()) == list(range(n))

    @pytest.mark.parametrize("n", range(1, 13))
    def test_binary_unrank_equal(self, n):
        for extra_params in [{}, {"span": 2.5}, {"branch_length": 3}]:
            ts = tskit.Tree.generate_balanced(n, **extra_params).tree_sequence
            N = tskit.combinatorics.num_shapes(n)
            equiv_ts = tskit.Tree.unrank(n, (N - 1, 0), **extra_params).tree_sequence
            assert ts.tables.equals(equiv_ts.tables, ignore_provenance=True)

    @pytest.mark.parametrize(
        ("n", "arity"), [(2, 2), (8, 2), (27, 3), (29, 3), (11, 5), (5, 10)]
    )
    def test_rank_unrank_round_trip_arity(self, n, arity):
        tree1 = tskit.Tree.generate_balanced(n, arity=arity)
        rank = tree1.rank()
        tree2 = tskit.Tree.unrank(n, rank)
        tables1 = tree1.tree_sequence.tables
        tables2 = tree2.tree_sequence.tables
        tables1.assert_equals(tables2, ignore_provenance=True)

    def test_bad_arity(self):
        for arity in [-1, 0, 1]:
            with pytest.raises(ValueError):
                tskit.Tree.generate_balanced(10, arity=arity)

    def test_branch_length_semantics(self):
        branch_length = 10
        tree = tskit.Tree.generate_balanced(8, branch_length=branch_length)
        for u in tree.nodes():
            for v in tree.children(u):
                # Special case cause n is a power of 2
                assert tree.time(u) == tree.time(v) + branch_length


class TestGenerateRandomBinary(TreeGeneratorTestBase):
    method_name = "generate_random_binary"

    def method(self, n, **kwargs):
        return tskit.Tree.generate_random_binary(n, random_seed=53, **kwargs)

    @pytest.mark.slow
    @pytest.mark.parametrize("n", [3, 4, 5])
    def test_all_topologies(self, n):
        N = num_leaf_labelled_binary_trees(n)
        ranks = collections.Counter()
        for seed in range(20 * N):
            random_tree = tskit.Tree.generate_random_binary(n, random_seed=seed)
            ranks[random_tree.rank()] += 1
        # There are N possible binary trees here, we should have seen them
        # all with high probability after 20 N attempts.
        assert len(ranks) == N

    @pytest.mark.parametrize("n", range(2, 10))
    def test_leaves(self, n):
        tree = tskit.Tree.generate_random_binary(n, random_seed=1234)
        # The leaves should be a permutation of range(n)
        assert list(sorted(tree.leaves())) == list(range(n))

    @pytest.mark.parametrize("seed", range(1, 20))
    def test_rank_unrank_round_trip_seeds(self, seed):
        n = 10
        tree1 = tskit.Tree.generate_random_binary(n, random_seed=seed)
        rank = tree1.rank()
        tree2 = tskit.Tree.unrank(n, rank)
        tables1 = tree1.tree_sequence.tables
        tables2 = tree2.tree_sequence.tables
        tables1.assert_equals(tables2, ignore_provenance=True)


class TestGenerateComb(TreeGeneratorTestBase):
    method_name = "generate_comb"

    # Hard-code in some pre-computed ranks for the comb(n) tree.
    @pytest.mark.parametrize(["n", "rank"], [(2, 0), (3, 1), (4, 3), (5, 8), (6, 20)])
    def test_unrank_equal(self, n, rank):
        for extra_params in [{}, {"span": 2.5}, {"branch_length": 3}]:
            ts = tskit.Tree.generate_comb(n, **extra_params).tree_sequence
            equiv_ts = tskit.Tree.unrank(n, (rank, 0), **extra_params).tree_sequence
            assert ts.tables.equals(equiv_ts.tables, ignore_provenance=True)

    def test_branch_length_semantics(self):
        branch_length = 10
        tree = tskit.Tree.generate_comb(2, branch_length=branch_length)
        assert tree.time(tree.root) == branch_length


class TestEqualChunks:
    @pytest.mark.parametrize(("n", "k"), [(2, 1), (4, 2), (9, 3), (100, 10)])
    def test_evenly_divisible(self, n, k):
        lst = range(n)
        chunks = list(comb.equal_chunks(lst, k))
        assert len(chunks) == k
        for chunk in chunks:
            assert len(chunk) == n // k
        assert list(itertools.chain(*chunks)) == list(range(n))

    @pytest.mark.parametrize("n", range(1, 5))
    def test_one_chunk(self, n):
        lst = list(range(n))
        chunks = list(comb.equal_chunks(lst, 1))
        assert chunks == [lst]

    @pytest.mark.parametrize(("n", "k"), [(1, 2), (5, 6), (10, 20), (5, 100)])
    def test_empty_chunks(self, n, k):
        lst = range(n)
        chunks = list(comb.equal_chunks(lst, k))
        assert len(chunks) == n
        for chunk in chunks:
            assert len(chunk) == 1
        assert list(itertools.chain(*chunks)) == list(range(n))

    @pytest.mark.parametrize(("n", "k"), [(3, 2), (10, 3), (11, 5), (13, 10)])
    def test_trailing_chunk(self, n, k):
        lst = range(n)
        chunks = list(comb.equal_chunks(lst, k))
        assert len(chunks) == k
        assert list(itertools.chain(*chunks)) == list(range(n))

    def test_empty_list(self):
        assert len(list(comb.equal_chunks([], 1))) == 0
        assert len(list(comb.equal_chunks([], 2))) == 0

    def test_bad_num_chunks(self):
        for bad_num_chunks in [0, -1, 0.1]:
            with pytest.raises(ValueError):
                list(comb.equal_chunks([1], bad_num_chunks))


--- ../../tskit/python/tests/test_threads.py ---


"""
Test cases for threading enabled aspects of the API.
"""
import platform
import threading

import msprime
import numpy as np
import pytest

import tests.tsutil as tsutil
import tskit

IS_WINDOWS = platform.system() == "Windows"
IS_OSX = platform.system() == "Darwin"


def run_threads(worker, num_threads):
    results = [None for _ in range(num_threads)]
    threads = [
        threading.Thread(target=worker, args=(j, results)) for j in range(num_threads)
    ]
    for t in threads:
        t.start()
    for t in threads:
        t.join()
    return results


class TestLdCalculatorReplicates:
    """
    Tests the LdCalculator object to ensure we get correct results
    when using threads.
    """

    num_test_sites = 25

    def get_tree_sequence(self):
        ts = msprime.simulate(
            20, mutation_rate=10, recombination_rate=10, random_seed=8
        )
        return tsutil.subsample_sites(ts, self.num_test_sites)

    def test_get_r2_multiple_instances(self):
        # This is the nominal case where we have a separate LdCalculator
        # instance in each thread.
        ts = self.get_tree_sequence()
        ld_calc = tskit.LdCalculator(ts)
        A = ld_calc.get_r2_matrix()
        del ld_calc
        m = A.shape[0]

        def worker(thread_index, results):
            ld_calc = tskit.LdCalculator(ts)
            row = np.zeros(m)
            results[thread_index] = row
            for j in range(m):
                row[j] = ld_calc.get_r2(thread_index, j)

        results = run_threads(worker, m)
        for j in range(m):
            assert np.allclose(results[j], A[j])

    def test_get_r2_single_instance(self):
        # This is the degenerate case where we have a single LdCalculator
        # instance shared by the threads. We should have only one thread
        # actually executing get_r2() at one time.
        ts = self.get_tree_sequence()
        ld_calc = tskit.LdCalculator(ts)
        A = ld_calc.get_r2_matrix()
        m = A.shape[0]

        def worker(thread_index, results):
            row = np.zeros(m)
            results[thread_index] = row
            for j in range(m):
                row[j] = ld_calc.get_r2(thread_index, j)

        results = run_threads(worker, m)
        for j in range(m):
            assert np.allclose(results[j], A[j])

    def test_get_r2_array_multiple_instances(self):
        # This is the nominal case where we have a separate LdCalculator
        # instance in each thread.
        ts = self.get_tree_sequence()
        ld_calc = tskit.LdCalculator(ts)
        A = ld_calc.get_r2_matrix()
        m = A.shape[0]
        del ld_calc

        def worker(thread_index, results):
            ld_calc = tskit.LdCalculator(ts)
            results[thread_index] = np.array(ld_calc.get_r2_array(thread_index))

        results = run_threads(worker, m)
        for j in range(m):
            assert np.allclose(results[j], A[j, j + 1 :])

    def test_get_r2_array_single_instance(self):
        # This is the degenerate case where we have a single LdCalculator
        # instance shared by the threads. We should have only one thread
        # actually executing get_r2_array() at one time. Because the buffer
        # is shared by many different instances, we can't make any assertions
        # about the returned values --- they are essentially gibberish.
        # However, we shouldn't crash and burn, which is what this test
        # is here to check for.
        ts = self.get_tree_sequence()
        ld_calc = tskit.LdCalculator(ts)
        m = ts.get_num_mutations()

        def worker(thread_index, results):
            results[thread_index] = ld_calc.get_r2_array(thread_index).shape

        results = run_threads(worker, m)
        for j in range(m):
            assert results[j][0] == m - j - 1


# Temporarily skipping these on Windows and OSX See
# https://github.com/tskit-dev/tskit/issues/344
# https://github.com/tskit-dev/tskit/issues/1041
@pytest.mark.skipif(
    IS_WINDOWS or IS_OSX, reason="Can't test thread support on Windows."
)
class TestTables:
    """
    Tests to ensure that attempts to access tables in threads correctly
    raise an exception.
    """

    def get_tables(self):
        # TODO include migrations here.
        ts = msprime.simulate(
            100, mutation_rate=10, recombination_rate=10, random_seed=8
        )
        return ts.tables

    def run_multiple_writers(self, writer, num_writers=32):
        barrier = threading.Barrier(num_writers)

        def writer_proxy(thread_index, results):
            barrier.wait()
            # Attempts to operate on a table while locked should raise a RuntimeError
            try:
                writer(thread_index, results)
                results[thread_index] = 0
            except RuntimeError:
                results[thread_index] = 1

        results = run_threads(writer_proxy, num_writers)
        failures = sum(results)
        successes = num_writers - failures
        # Note: we would like to insist that #failures is > 0, but this is too
        # stochastic to guarantee for test purposes.
        assert failures >= 0
        assert successes > 0

    def run_failing_reader(self, writer, reader, num_readers=32):
        """
        Runs a test in which a single writer acceses some tables
        and a bunch of other threads try to read the data.
        """
        barrier = threading.Barrier(num_readers + 1)

        def writer_proxy():
            barrier.wait()
            writer()

        def reader_proxy(thread_index, results):
            barrier.wait()
            # Attempts to operate on a table while locked should raise a RuntimeError
            results[thread_index] = 0
            try:
                reader(thread_index, results)
            except RuntimeError:
                results[thread_index] = 1

        writer_thread = threading.Thread(target=writer_proxy)
        writer_thread.start()
        results = run_threads(reader_proxy, num_readers)
        writer_thread.join()

        failures = sum(results)
        successes = num_readers - failures
        # Note: we would like to insist that #failures is > 0, but this is too
        # stochastic to guarantee for test purposes.
        assert failures >= 0
        assert successes > 0

    def test_many_simplify_all_tables(self):
        tables = self.get_tables()

        def writer(thread_index, results):
            tables.simplify([0, 1])

        self.run_multiple_writers(writer)

    def test_many_sort(self):
        tables = self.get_tables()

        def writer(thread_index, results):
            tables.sort()

        self.run_multiple_writers(writer)

    def run_simplify_access_table(self, table_name, col_name):
        tables = self.get_tables()

        def writer():
            tables.simplify([0, 1])

        table = getattr(tables, table_name)

        def reader(thread_index, results):
            for _ in range(100):
                x = getattr(table, col_name)
                assert x.shape[0] == len(table)

        self.run_failing_reader(writer, reader)

    def run_sort_access_table(self, table_name, col_name):
        tables = self.get_tables()

        def writer():
            tables.sort()

        table = getattr(tables, table_name)

        def reader(thread_index, results):
            for _ in range(100):
                x = getattr(table, col_name)
                assert x.shape[0] == len(table)

        self.run_failing_reader(writer, reader)

    def test_simplify_access_nodes(self):
        self.run_simplify_access_table("nodes", "time")

    def test_simplify_access_edges(self):
        self.run_simplify_access_table("edges", "left")

    def test_simplify_access_sites(self):
        self.run_simplify_access_table("sites", "position")

    def test_simplify_access_mutations(self):
        self.run_simplify_access_table("mutations", "site")

    def test_sort_access_nodes(self):
        self.run_sort_access_table("nodes", "time")

    def test_sort_access_edges(self):
        self.run_sort_access_table("edges", "left")

    def test_sort_access_sites(self):
        self.run_sort_access_table("sites", "position")

    def test_sort_access_mutations(self):
        self.run_sort_access_table("mutations", "site")


--- ../../tskit/python/tests/test_metadata.py ---


"""
Tests for metadata handling.
"""
import collections
import io
import json
import os
import pickle
import pprint
import struct
import tempfile
import unittest
from unittest.mock import patch

import msgpack
import msprime
import numpy as np
import pytest

import tskit
import tskit.exceptions as exceptions
import tskit.metadata as metadata


class TestMetadataHdf5RoundTrip(unittest.TestCase):
    """
    Tests that we can encode metadata under various formats and this will
    successfully round-trip through the HDF5 format.
    """

    def setUp(self):
        fd, self.temp_file = tempfile.mkstemp(prefix="msp_hdf5meta_test_")
        os.close(fd)

    def tearDown(self):
        os.unlink(self.temp_file)

    def test_json(self):
        ts = msprime.simulate(10, random_seed=1)
        tables = ts.dump_tables()
        nodes = tables.nodes
        # For each node, we create some Python metadata that can be JSON encoded.
        metadata = [
            {"one": j, "two": 2 * j, "three": list(range(j))} for j in range(len(nodes))
        ]
        encoded, offset = tskit.pack_strings(map(json.dumps, metadata))
        nodes.set_columns(
            flags=nodes.flags,
            time=nodes.time,
            population=nodes.population,
            metadata_offset=offset,
            metadata=encoded,
        )
        assert np.array_equal(nodes.metadata_offset, offset)
        assert np.array_equal(nodes.metadata, encoded)
        ts1 = tables.tree_sequence()
        for j, node in enumerate(ts1.nodes()):
            decoded_metadata = json.loads(node.metadata.decode())
            assert decoded_metadata == metadata[j]
        ts1.dump(self.temp_file)
        ts2 = tskit.load(self.temp_file)
        assert ts1.tables.nodes == ts2.tables.nodes

    def test_pickle(self):
        ts = msprime.simulate(10, random_seed=1)
        tables = ts.dump_tables()
        # For each node, we create some Python metadata that can be pickled
        metadata = [
            {"one": j, "two": 2 * j, "three": list(range(j))}
            for j in range(ts.num_nodes)
        ]
        encoded, offset = tskit.pack_bytes(list(map(pickle.dumps, metadata)))
        tables.nodes.set_columns(
            flags=tables.nodes.flags,
            time=tables.nodes.time,
            population=tables.nodes.population,
            metadata_offset=offset,
            metadata=encoded,
        )
        assert np.array_equal(tables.nodes.metadata_offset, offset)
        assert np.array_equal(tables.nodes.metadata, encoded)
        ts1 = tables.tree_sequence()
        for j, node in enumerate(ts1.nodes()):
            decoded_metadata = pickle.loads(node.metadata)
            assert decoded_metadata == metadata[j]
        ts1.dump(self.temp_file)
        ts2 = tskit.load(self.temp_file)
        assert ts1.tables.nodes == ts2.tables.nodes


class ExampleMetadata:
    """
    Simple class that we can pickle/unpickle in metadata.
    """

    def __init__(self, one=None, two=None):
        self.one = one
        self.two = two


class TestMetadataPickleDecoding:
    """
    Tests in which use pickle.pickle to decode metadata in nodes, sites and mutations.
    """

    def test_nodes(self):
        tables = tskit.TableCollection(sequence_length=1)
        metadata = ExampleMetadata(one="node1", two="node2")
        pickled = pickle.dumps(metadata)
        tables.nodes.add_row(time=0.125, metadata=pickled)
        ts = tables.tree_sequence()
        node = ts.node(0)
        assert node.time == 0.125
        assert node.metadata == pickled
        unpickled = pickle.loads(node.metadata)
        assert unpickled.one == metadata.one
        assert unpickled.two == metadata.two

    def test_sites(self):
        tables = tskit.TableCollection(sequence_length=1)
        metadata = ExampleMetadata(one="node1", two="node2")
        pickled = pickle.dumps(metadata)
        tables.sites.add_row(position=0.1, ancestral_state="A", metadata=pickled)
        ts = tables.tree_sequence()
        site = ts.site(0)
        assert site.position == 0.1
        assert site.ancestral_state == "A"
        assert site.metadata == pickled
        unpickled = pickle.loads(site.metadata)
        assert unpickled.one == metadata.one
        assert unpickled.two == metadata.two

    def test_mutations(self):
        tables = tskit.TableCollection(sequence_length=1)
        metadata = ExampleMetadata(one="node1", two="node2")
        pickled = pickle.dumps(metadata)
        tables.nodes.add_row(time=0)
        tables.sites.add_row(position=0.1, ancestral_state="A")
        tables.mutations.add_row(site=0, node=0, derived_state="T", metadata=pickled)
        ts = tables.tree_sequence()
        mutation = ts.site(0).mutations[0]
        assert mutation.site == 0
        assert mutation.node == 0
        assert mutation.derived_state == "T"
        assert mutation.metadata == pickled
        unpickled = pickle.loads(mutation.metadata)
        assert unpickled.one == metadata.one
        assert unpickled.two == metadata.two


class TestLoadTextMetadata:
    """
    Tests that use the load_text interface.
    """

    def test_individuals(self):
        individuals = io.StringIO(
            """\
        id  flags location     parents  metadata
        0   1     0.0,1.0,0.0  -1,-1    abc
        1   1     1.0,2.0      0,0      XYZ+
        2   0     2.0,3.0,0.0  0,1      !@#$%^&*()
        """
        )
        i = tskit.parse_individuals(
            individuals, strict=False, encoding="utf8", base64_metadata=False
        )
        expected = [
            (1, [0.0, 1.0, 0.0], [-1, -1], "abc"),
            (1, [1.0, 2.0], [0, 0], "XYZ+"),
            (0, [2.0, 3.0, 0.0], [0, 1], "!@#$%^&*()"),
        ]
        for a, b in zip(expected, i):
            assert a[0] == b.flags
            assert len(a[1]) == len(b.location)
            for x, y in zip(a[1], b.location):
                assert x == y
            assert len(a[2]) == len(b.parents)
            for x, y in zip(a[2], b.parents):
                assert x == y
        assert a[3].encode("utf8") == b.metadata

    def test_nodes(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time    metadata
        0   1           0   abc
        1   1           0   XYZ+
        2   0           1   !@#$%^&*()
        """
        )
        n = tskit.parse_nodes(
            nodes, strict=False, encoding="utf8", base64_metadata=False
        )
        expected = ["abc", "XYZ+", "!@#$%^&*()"]
        for a, b in zip(expected, n):
            assert a.encode("utf8") == b.metadata

    def test_sites(self):
        sites = io.StringIO(
            """\
        position    ancestral_state metadata
        0.1 A   abc
        0.5 C   XYZ+
        0.8 G   !@#$%^&*()
        """
        )
        s = tskit.parse_sites(
            sites, strict=False, encoding="utf8", base64_metadata=False
        )
        expected = ["abc", "XYZ+", "!@#$%^&*()"]
        for a, b in zip(expected, s):
            assert a.encode("utf8") == b.metadata

    def test_mutations(self):
        mutations = io.StringIO(
            """\
        site    node    derived_state   metadata
        0   2   C   mno
        0   3   G   )(*&^%$#@!
        """
        )
        m = tskit.parse_mutations(
            mutations, strict=False, encoding="utf8", base64_metadata=False
        )
        expected = ["mno", ")(*&^%$#@!"]
        for a, b in zip(expected, m):
            assert a.encode("utf8") == b.metadata

    def test_populations(self):
        populations = io.StringIO(
            """\
        id    metadata
        0     mno
        1     )(*&^%$#@!
        """
        )
        p = tskit.parse_populations(
            populations, strict=False, encoding="utf8", base64_metadata=False
        )
        expected = ["mno", ")(*&^%$#@!"]
        for a, b in zip(expected, p):
            assert a.encode("utf8") == b.metadata

    @pytest.mark.parametrize(
        "base64_metadata,expected", [(True, ["pop", "gen"]), (False, ["cG9w", "Z2Vu"])]
    )
    def test_migrations(self, base64_metadata, expected):
        migrations = io.StringIO(
            """\
        left    right    node    source    dest    time    metadata
        10    100    0    3    4    123.0    cG9w
        150    360    1    1    2    307.0    Z2Vu
        """
        )
        m = tskit.parse_migrations(
            migrations, strict=False, encoding="utf8", base64_metadata=base64_metadata
        )
        for a, b in zip(expected, m):
            assert a.encode("utf8") == b.metadata


class TestMetadataModule:
    """
    Tests that use the metadata module
    """

    def test_metadata_schema(self):
        # Bad jsonschema
        with pytest.raises(exceptions.MetadataSchemaValidationError):
            metadata.MetadataSchema(
                {"codec": "json", "additionalProperties": "THIS ISN'T RIGHT"},
            )
        # Bad codec
        with pytest.raises(exceptions.MetadataSchemaValidationError):
            metadata.MetadataSchema({"codec": "morse-code"})
        # Missing codec
        with pytest.raises(exceptions.MetadataSchemaValidationError):
            metadata.MetadataSchema({})
        schema = {
            "codec": "json",
            "title": "Example Metadata",
            "type": "object",
            "properties": {"one": {"type": "string"}, "two": {"type": "number"}},
            "required": ["one", "two"],
            "additionalProperties": False,
        }
        ms = metadata.MetadataSchema(schema)
        assert repr(ms) == tskit.canonical_json(schema)
        # Missing required properties
        with pytest.raises(exceptions.MetadataValidationError):
            ms.validate_and_encode_row({})

    def test_schema_str(self):
        schema = {
            "codec": "json",
            "title": "Example Metadata",
            "type": "object",
            "properties": {"one": {"type": "string"}, "two": {"type": "number"}},
            "required": ["one", "two"],
            "additionalProperties": False,
        }
        assert (
            str(metadata.MetadataSchema(schema))
            == f"tskit.MetadataSchema(\n{pprint.pformat(schema)}\n)"
        )

    def test_register_codec(self):
        class TestCodec(metadata.AbstractMetadataCodec):
            pass

        metadata.register_metadata_codec(TestCodec, "test")
        assert TestCodec == metadata.codec_registry["test"]

    def test_parse(self):
        # Empty string gives MetaDataSchema with None codec
        ms = metadata.parse_metadata_schema("")
        assert isinstance(ms, metadata.MetadataSchema)
        assert ms.schema is None
        assert ms.asdict() is None

        # json gives MetaDataSchema with json codec
        ms = metadata.parse_metadata_schema(json.dumps({"codec": "json"}))
        assert isinstance(ms, metadata.MetadataSchema)
        assert ms.schema == {"codec": "json"}
        assert ms.asdict() == {"codec": "json"}
        # check we get a copy
        assert ms.asdict() is not ms._schema

        # Bad JSON gives error
        with pytest.raises(ValueError):
            metadata.parse_metadata_schema(json.dumps({"codec": "json"})[:-1])

    def test_canonical_string(self):
        schema = collections.OrderedDict(
            codec="json",
            title="Example Metadata",
            type="object",
            properties=collections.OrderedDict(
                one={"type": "string"}, two={"type": "number"}
            ),
            required=["one", "two"],
            additionalProperties=False,
        )
        schema2 = collections.OrderedDict(
            type="object",
            properties=collections.OrderedDict(
                two={"type": "number"}, one={"type": "string"}
            ),
            required=["one", "two"],
            additionalProperties=False,
            title="Example Metadata",
            codec="json",
        )
        assert json.dumps(schema) != json.dumps(schema2)
        assert repr(metadata.MetadataSchema(schema)) == repr(
            metadata.MetadataSchema(schema2)
        )

    def test_equality(self):
        schema = metadata.MetadataSchema(
            {
                "codec": "json",
                "title": "Example Metadata",
                "type": "object",
                "properties": {"one": {"type": "string"}, "two": {"type": "number"}},
                "required": ["one", "two"],
                "additionalProperties": False,
            }
        )
        schema_same = metadata.MetadataSchema(
            collections.OrderedDict(
                type="object",
                properties=collections.OrderedDict(
                    two={"type": "number"}, one={"type": "string"}
                ),
                required=["one", "two"],
                additionalProperties=False,
                title="Example Metadata",
                codec="json",
            )
        )
        schema_diff = metadata.MetadataSchema(
            {
                "codec": "json",
                "title": "Example Metadata",
                "type": "object",
                "properties": {"one": {"type": "string"}, "two": {"type": "string"}},
                "required": ["one", "two"],
                "additionalProperties": False,
            }
        )
        assert schema == schema
        assert not (schema != schema)
        assert schema == schema_same
        assert not (schema != schema_same)
        assert schema != schema_diff
        assert not (schema == schema_diff)

    def test_bad_top_level_type(self):
        for bad_type in ["array", "boolean", "integer", "null", "number", "string"]:
            schema = {
                "codec": "json",
                "type": bad_type,
            }
            with pytest.raises(exceptions.MetadataSchemaValidationError):
                metadata.MetadataSchema(schema)

    @pytest.mark.parametrize("codec", ["struct", "json"])
    def test_null_union_top_level(self, codec):
        schema = {
            "codec": f"{codec}",
            "type": ["object", "null"],
            "properties": {
                "one": {
                    "type": "string",
                    "binaryFormat": "1024s",
                    "nullTerminated": True,
                },
                "two": {"type": "number", "binaryFormat": "i"},
            },
        }
        ms = metadata.MetadataSchema(schema)
        row_data = {"one": "tree", "two": 5}
        assert ms.decode_row(ms.validate_and_encode_row(row_data)) == row_data
        assert ms.decode_row(ms.validate_and_encode_row(None)) is None

    def test_null_codec(self):
        ms = metadata.MetadataSchema(None)
        assert repr(ms) == ""
        row = b"Some binary data that tskit can't interpret "
        # Encode/decode are no-ops
        assert row == ms.validate_and_encode_row(row)
        assert row == ms.decode_row(row)
        # Only bytes validate
        with pytest.raises(TypeError):
            ms.validate_and_encode_row({})

    def test_json_codec(self):
        schema = {
            "codec": "json",
            "title": "Example Metadata",
            "type": "object",
            "properties": {"one": {"type": "string"}, "two": {"type": "number"}},
            "required": ["one", "two"],
            "additionalProperties": False,
        }
        ms = metadata.MetadataSchema(schema)
        # Valid row data
        row_data = {"one": "tree", "two": 5}
        assert (
            ms.validate_and_encode_row(row_data)
            == tskit.canonical_json(row_data).encode()
        )
        assert ms.decode_row(json.dumps(row_data).encode()) == row_data
        # Round trip
        assert ms.decode_row(ms.validate_and_encode_row(row_data)) == row_data
        # Test canonical encoding
        row_data = collections.OrderedDict(one="tree", two=5)
        row_data2 = collections.OrderedDict(two=5, one="tree")
        assert json.dumps(row_data) != json.dumps(row_data2)
        assert ms.validate_and_encode_row(row_data) == ms.validate_and_encode_row(
            row_data2
        )

    def test_msgpack_codec(self):
        class MsgPackCodec(metadata.AbstractMetadataCodec):
            def __init__(self, schema):
                pass

            def encode(self, obj):
                return msgpack.dumps(obj)

            def decode(self, encoded):
                return msgpack.loads(encoded)

        metadata.register_metadata_codec(MsgPackCodec, "msgpack")

        schema = {
            "codec": "msgpack",
            "title": "Example Metadata",
            "type": "object",
            "properties": {"one": {"type": "string"}, "two": {"type": "number"}},
            "required": ["one", "two"],
            "additionalProperties": False,
        }
        ms = metadata.MetadataSchema(schema)
        # Valid row data
        row_data = {"one": "tree", "two": 5}
        assert ms.validate_and_encode_row(row_data) == msgpack.dumps(row_data)
        assert ms.decode_row(msgpack.dumps(row_data)) == row_data
        # Round trip
        assert ms.decode_row(ms.validate_and_encode_row(row_data)) == row_data


class TestJSONCodec:
    def test_simple_default(self):
        schema = {
            "codec": "json",
            "type": "object",
            "properties": {"number": {"type": "number", "default": 5}},
        }
        ms = tskit.MetadataSchema(schema)
        assert ms.decode_row(b"") == {"number": 5}
        assert ms.decode_row(ms.validate_and_encode_row({})) == {"number": 5}
        assert ms.decode_row(ms.validate_and_encode_row({"number": 42})) == {
            "number": 42
        }

    def test_nested_default_error(self):
        schema = {
            "codec": "json",
            "type": "object",
            "properties": {
                "obj": {
                    "type": "object",
                    "properties": {
                        "nested_obj_no_default": {
                            "type": "object",
                            "properties": {},
                        },
                        "nested_obj": {
                            "type": "object",
                            "properties": {},
                            "default": {"foo": "bar"},
                        },
                    },
                }
            },
        }
        with pytest.raises(
            tskit.MetadataSchemaValidationError,
            match="Defaults can only be specified at the top level for JSON codec",
        ):
            tskit.MetadataSchema(schema)

    def test_bad_type_error(self):
        ms = tskit.MetadataSchema({"codec": "json"})
        with pytest.raises(
            exceptions.MetadataEncodingError,
            match="Could not encode metadata of type TableCollection",
        ):
            ms.validate_and_encode_row(tskit.TableCollection(1))

    def test_skip_validation(self):
        ms = tskit.MetadataSchema({"codec": "json"})
        assert ms._bypass_validation
        with patch.object(ms, "_validate_row", return_value=True) as mocked_validate:
            ms.validate_and_encode_row({})
            assert mocked_validate.call_count == 0

    def test_dont_skip_validation(self):
        ms = tskit.MetadataSchema({"codec": "json", "properties": {"foo": {}}})
        assert not ms._bypass_validation
        with patch.object(ms, "_validate_row", return_value=True) as mocked_validate:
            ms.validate_and_encode_row({})
            assert mocked_validate.call_count == 1

    def test_dont_skip_validation_other_codecs(self):
        ms = tskit.MetadataSchema(
            {
                "codec": "struct",
                "type": "object",
                "properties": {
                    "int": {"type": "number", "binaryFormat": "i"},
                },
            }
        )
        assert not ms._bypass_validation
        with patch.object(ms, "_validate_row", return_value=True) as mocked_validate:
            ms.validate_and_encode_row({"int": 1})
            assert mocked_validate.call_count == 1

    def test_zero_length(self):
        ms = tskit.MetadataSchema({"codec": "json"})
        assert ms.decode_row(b"") == {}


class TestStructCodec:
    def encode_decode(self, method_name, sub_schema, obj, buffer):
        assert (
            getattr(metadata.StructCodec, f"{method_name}_encode")(sub_schema)(obj)
            == buffer
        )
        assert (
            getattr(metadata.StructCodec, f"{method_name}_decode")(sub_schema)(
                iter(buffer)
            )
            == obj
        )

    def test_order_schema(self):
        # Make a guaranteed-unordered nested, schema
        schema = {
            "codec": "struct",
            "title": "Example Struct-encoded Metadata",
            "type": "object",
            "properties": collections.OrderedDict(
                [
                    ("d", {"type": "number", "binaryFormat": "L"}),
                    ("a", {"type": "string", "binaryFormat": "10s"}),
                    (
                        "f",
                        {
                            "type": "array",
                            "items": {
                                "type": "object",
                                "properties": collections.OrderedDict(
                                    [
                                        (
                                            "m",
                                            {
                                                "type": "number",
                                                "index": 0,
                                                "binaryFormat": "L",
                                            },
                                        ),
                                        (
                                            "n",
                                            {
                                                "type": "string",
                                                "index": -1000,
                                                "binaryFormat": "10s",
                                            },
                                        ),
                                        (
                                            "l",
                                            {
                                                "type": "string",
                                                "index": 1000,
                                                "binaryFormat": "10s",
                                            },
                                        ),
                                    ]
                                ),
                            },
                        },
                    ),
                    ("c", {"type": "string", "binaryFormat": "10s"}),
                    (
                        "h",
                        {
                            "type": "object",
                            "properties": collections.OrderedDict(
                                [
                                    (
                                        "i",
                                        {
                                            "type": "string",
                                            "index": 1000,
                                            "binaryFormat": "10s",
                                        },
                                    ),
                                    (
                                        "j",
                                        {
                                            "type": "string",
                                            "index": 567,
                                            "binaryFormat": "10s",
                                        },
                                    ),
                                    (
                                        "k",
                                        {
                                            "type": "number",
                                            "index": 567.5,
                                            "binaryFormat": "L",
                                        },
                                    ),
                                ]
                            ),
                        },
                    ),
                    ("e", {"type": "string", "binaryFormat": "10s"}),
                    ("g", {"type": "string", "binaryFormat": "10s"}),
                    ("b", {"type": "number", "binaryFormat": "L"}),
                ]
            ),
            "required": ["one", "two"],
            "additionalProperties": False,
        }
        schema_sorted = {
            "codec": "struct",
            "title": "Example Struct-encoded Metadata",
            "type": "object",
            "properties": collections.OrderedDict(
                [
                    ("a", {"type": "string", "binaryFormat": "10s"}),
                    ("b", {"type": "number", "binaryFormat": "L"}),
                    ("c", {"type": "string", "binaryFormat": "10s"}),
                    ("d", {"type": "number", "binaryFormat": "L"}),
                    ("e", {"type": "string", "binaryFormat": "10s"}),
                    (
                        "f",
                        {
                            "type": "array",
                            "items": {
                                "type": "object",
                                "properties": collections.OrderedDict(
                                    [
                                        (
                                            "n",
                                            {
                                                "type": "string",
                                                "index": -1000,
                                                "binaryFormat": "10s",
                                            },
                                        ),
                                        (
                                            "m",
                                            {
                                                "type": "number",
                                                "index": 0,
                                                "binaryFormat": "L",
                                            },
                                        ),
                                        (
                                            "l",
                                            {
                                                "type": "string",
                                                "index": 1000,
                                                "binaryFormat": "10s",
                                            },
                                        ),
                                    ]
                                ),
                            },
                        },
                    ),
                    ("g", {"type": "string", "binaryFormat": "10s"}),
                    (
                        "h",
                        {
                            "type": "object",
                            "properties": collections.OrderedDict(
                                [
                                    (
                                        "j",
                                        {
                                            "type": "string",
                                            "index": 567,
                                            "binaryFormat": "10s",
                                        },
                                    ),
                                    (
                                        "k",
                                        {
                                            "type": "number",
                                            "index": 567.5,
                                            "binaryFormat": "L",
                                        },
                                    ),
                                    (
                                        "i",
                                        {
                                            "type": "string",
                                            "index": 1000,
                                            "binaryFormat": "10s",
                                        },
                                    ),
                                ]
                            ),
                        },
                    ),
                ]
            ),
            "required": ["one", "two"],
            "additionalProperties": False,
        }
        assert metadata.StructCodec.order_by_index(schema) == schema_sorted

    def test_make_encode_and_decode(self):
        self.encode_decode(
            "make",
            {
                "type": "array",
                "arrayLengthFormat": "B",
                "items": {"type": "number", "binaryFormat": "b"},
            },
            list(range(5)),
            b"\x05\x00\x01\x02\x03\x04",
        )
        self.encode_decode(
            "make",
            {
                "type": "object",
                "properties": {
                    "a": {"type": "number", "binaryFormat": "b"},
                    "b": {"type": "string", "binaryFormat": "5p"},
                },
            },
            {"a": 5, "b": "FOO"},
            b"\x05\x03FOO\x00",
        )
        self.encode_decode(
            "make",
            {"type": "string", "binaryFormat": "10p"},
            "FOOBAR",
            b"\x06FOOBAR\x00\x00\x00",
        )
        self.encode_decode("make", {"type": "null"}, None, b"")
        self.encode_decode(
            "make", {"type": "boolean", "binaryFormat": "?"}, True, b"\x01"
        )
        self.encode_decode(
            "make", {"type": "integer", "binaryFormat": "b"}, -128, b"\x80"
        )
        self.encode_decode(
            "make",
            {"type": "number", "binaryFormat": "f"},
            42.424198150634766,
            b"a\xb2)B",
        )

    def test_make_array_encode_and_decode(self):
        # Default array length format is 'L'
        self.encode_decode(
            "make_array",
            {"type": "array", "items": {"type": "number", "binaryFormat": "b"}},
            list(range(5)),
            b"\x05\x00\x00\x00\x00\x01\x02\x03\x04",
        )
        self.encode_decode(
            "make_array",
            {
                "type": "array",
                "arrayLengthFormat": "H",
                "items": {"type": "number", "binaryFormat": "b"},
            },
            list(range(6)),
            b"\x06\x00\x00\x01\x02\x03\x04\x05",
        )
        self.encode_decode(
            "make_array",
            {
                "type": "array",
                "arrayLengthFormat": "B",
                "items": {"type": "number", "binaryFormat": "b"},
            },
            [],
            b"\x00",
        )
        sub_schema = {
            "type": "array",
            "arrayLengthFormat": "B",
            "items": {
                "type": "array",
                "arrayLengthFormat": "B",
                "items": {"type": "number", "binaryFormat": "b"},
            },
        }
        self.encode_decode("make_array", sub_schema, [], b"\x00")
        self.encode_decode("make_array", sub_schema, [[]], b"\x01\x00")
        self.encode_decode(
            "make_array", sub_schema, [[3, 4], [5]], b"\x02\x02\x03\x04\x01\x05"
        )

    def test_make_array_no_length_encoding_exhaust_buffer(self):
        self.encode_decode(
            "make_array",
            {
                "type": "array",
                "noLengthEncodingExhaustBuffer": True,
                "items": {"type": "number", "binaryFormat": "b"},
            },
            list(range(5)),
            b"\x00\x01\x02\x03\x04",
        )

        self.encode_decode(
            "make_array",
            {
                "type": "array",
                "noLengthEncodingExhaustBuffer": True,
                "items": {
                    "type": "object",
                    "properties": {
                        "a": {"type": "number", "binaryFormat": "b"},
                        "b": {"type": "number", "binaryFormat": "Q"},
                        "c": {"type": "number", "binaryFormat": "?"},
                        "d": {"type": "string", "binaryFormat": "5p"},
                    },
                },
            },
            [
                {
                    "a": 5 + i,
                    "b": 18446744073709551615 - i,
                    "c": (i // 2) == 0,
                    "d": "FOO",
                }
                for i in range(10)
            ],
            b"\x05\xff\xff\xff\xff\xff\xff\xff\xff\x01\x03FOO\x00"
            b"\x06\xfe\xff\xff\xff\xff\xff\xff\xff\x01\x03FOO\x00"
            b"\x07\xfd\xff\xff\xff\xff\xff\xff\xff\x00\x03FOO\x00"
            b"\x08\xfc\xff\xff\xff\xff\xff\xff\xff\x00\x03FOO\x00"
            b"\x09\xfb\xff\xff\xff\xff\xff\xff\xff\x00\x03FOO\x00"
            b"\x0a\xfa\xff\xff\xff\xff\xff\xff\xff\x00\x03FOO\x00"
            b"\x0b\xf9\xff\xff\xff\xff\xff\xff\xff\x00\x03FOO\x00"
            b"\x0c\xf8\xff\xff\xff\xff\xff\xff\xff\x00\x03FOO\x00"
            b"\x0d\xf7\xff\xff\xff\xff\xff\xff\xff\x00\x03FOO\x00"
            b"\x0e\xf6\xff\xff\xff\xff\xff\xff\xff\x00\x03FOO\x00",
        )

        # Other struct errors should still be raised
        schema = {
            "type": "array",
            "noLengthEncodingExhaustBuffer": True,
            "items": {"type": "number", "binaryFormat": "I'M NOT VALID"},
        }
        with pytest.raises(struct.error):
            metadata.StructCodec.make_array_encode(schema)(5)
        with pytest.raises(struct.error):
            metadata.StructCodec.make_array_decode(schema)(5)

    def test_make_object_encode_and_decode(self):
        self.encode_decode("make_object", {"type": "object", "properties": {}}, {}, b"")
        self.encode_decode(
            "make_object",
            {
                "type": "object",
                "properties": {
                    "a": {"type": "number", "binaryFormat": "b"},
                    "b": {"type": "number", "binaryFormat": "Q"},
                    "c": {"type": "number", "binaryFormat": "?"},
                    "d": {"type": "string", "binaryFormat": "5p"},
                },
            },
            {"a": 5, "b": 18446744073709551615, "c": True, "d": "FOO"},
            b"\x05\xff\xff\xff\xff\xff\xff\xff\xff\x01\x03FOO\x00",
        )
        self.encode_decode(
            "make_object",
            {
                "type": "object",
                "properties": {
                    "obj": {
                        "type": "object",
                        "properties": {
                            "a": {"type": "number", "binaryFormat": "b"},
                            "b": {"type": "number", "binaryFormat": "Q"},
                            "c": {"type": "number", "binaryFormat": "?"},
                            "d": {"type": "string", "binaryFormat": "5p"},
                        },
                    },
                },
            },
            {"obj": {"a": 5, "b": 18446744073709551615, "c": True, "d": "FOO"}},
            b"\x05\xff\xff\xff\xff\xff\xff\xff\xff\x01\x03FOO\x00",
        )

    def test_make_string_encode_and_decode(self):
        # Single byte
        self.encode_decode(
            "make_string", {"type": "string", "binaryFormat": "c"}, "a", b"a"
        )
        # With "s" encoding exactly the right size comes back fine
        self.encode_decode(
            "make_string", {"type": "string", "binaryFormat": "4s"}, "abcd", b"abcd"
        )
        # If too small gets truncated
        assert (
            metadata.StructCodec.make_string_encode(
                {"type": "string", "binaryFormat": "2s"}
            )("abcd")
            == b"ab"
        )
        # If too large gets padded - have to test separately as encode and decode are not
        # inverse of each other in this case
        assert (
            metadata.StructCodec.make_string_encode(
                {"type": "string", "binaryFormat": "6s"}
            )("abcd")
            == b"abcd\x00\x00"
        )
        # Too large getting decoded returns padding
        assert (
            metadata.StructCodec.make_string_decode(
                {"type": "string", "binaryFormat": "6s"}
            )(b"abcd\x00\x00")
            == "abcd\x00\x00"
        )
        assert (
            metadata.StructCodec.make_string_decode(
                {"type": "string", "binaryFormat": "6s", "nullTerminated": False}
            )(b"abcd\x00\x00")
            == "abcd\x00\x00"
        )
        # Unless we specify that the field is null-teminated
        self.encode_decode(
            "make_string",
            {"type": "string", "binaryFormat": "6s", "nullTerminated": True},
            "abcd",
            b"abcd\x00\x00",
        )
        # For "p" the padding is not returned, even if nullTerminated is False
        self.encode_decode(
            "make_string",
            {"type": "string", "binaryFormat": "8p"},
            "abcd",
            b"\x04abcd\x00\x00\x00",
        )

        # Unicode
        self.encode_decode(
            "make_string",
            {"type": "string", "binaryFormat": "6s", "nullTerminated": True},
            "💩",
            b"\xf0\x9f\x92\xa9\x00\x00",
        )
        self.encode_decode(
            "make_string",
            {
                "type": "string",
                "binaryFormat": "8s",
                "nullTerminated": True,
                "stringEncoding": "utf-16",
            },
            "💩",
            b"\xff\xfe=\xd8\xa9\xdc\x00\x00",
        )
        self.encode_decode(
            "make_string",
            {"type": "string", "binaryFormat": "9p", "stringEncoding": "utf-32"},
            "💩",
            b"\x08\xff\xfe\x00\x00\xa9\xf4\x01\x00",
        )

    def test_make_null_encode_and_decode(self):
        self.encode_decode("make_null", {"type": "null"}, None, b"")
        self.encode_decode(
            "make_null", {"type": "null", "binaryFormat": "x"}, None, b"\x00"
        )
        self.encode_decode(
            "make_null", {"type": "null", "binaryFormat": "3x"}, None, b"\x00\x00\x00"
        )

    def test_make_numeric_encode_and_decode(self):
        self.encode_decode(
            "make_numeric",
            {"type": "number", "binaryFormat": "f"},
            42.424198150634766,
            b"a\xb2)B",
        )
        self.encode_decode(
            "make_numeric", {"type": "integer", "binaryFormat": "b"}, 42, b"*"
        )

    def test_null_union_top_level(self):
        # This nested with mutiple values tests that the buffer length check has not
        # caused a list to past to sub-decoders
        schema = {
            "codec": "struct",
            "type": ["object", "null"],
            "properties": {
                "o": {
                    "type": "object",
                    "properties": {"x": {"type": "number", "binaryFormat": "d"}},
                },
                "a": {"type": "number", "binaryFormat": "d"},
                "b": {"type": "number", "binaryFormat": "d"},
            },
        }
        ms = metadata.MetadataSchema(schema)
        row_data = {"o": {"x": 5.5}, "a": 4, "b": 7}
        assert ms.decode_row(ms.validate_and_encode_row(row_data)) == row_data
        assert ms.decode_row(ms.validate_and_encode_row(None)) is None

    def test_default_values(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "int": {"type": "number", "binaryFormat": "b", "default": 42},
                "float": {"type": "number", "binaryFormat": "d"},
            },
        }
        ms = metadata.MetadataSchema(schema)
        row_data = {"float": 5.5}
        assert ms.validate_and_encode_row(row_data) == b"\x00\x00\x00\x00\x00\x00\x16@*"
        assert ms.decode_row(ms.validate_and_encode_row(row_data)) == {
            "float": 5.5,
            "int": 42,
        }

    def test_defaults_object_or_null(self):
        schema = {
            "codec": "struct",
            "type": ["object", "null"],
            "properties": {
                "int": {"type": "number", "binaryFormat": "b", "default": 42},
                "float": {"type": "number", "binaryFormat": "d"},
            },
        }
        ms = metadata.MetadataSchema(schema)
        row_data = {"float": 5.5}
        assert ms.validate_and_encode_row(row_data) == b"\x00\x00\x00\x00\x00\x00\x16@*"
        assert ms.decode_row(ms.validate_and_encode_row(row_data)) == {
            "float": 5.5,
            "int": 42,
        }
        assert ms.validate_and_encode_row(None) == b""
        assert ms.decode_row(b"") is None


class TestStructCodecRoundTrip:
    def round_trip(self, schema, row_data):
        ms = metadata.MetadataSchema(schema)
        assert ms.decode_row(ms.validate_and_encode_row(row_data)) == row_data

    def test_simple_types(self):
        for type_, binaryFormat, value in (
            ("number", "i", 5),
            ("number", "d", 5.5),
            ("string", "10p", "foobar"),
            ("boolean", "?", True),
            ("boolean", "?", False),
            ("null", "10x", None),
        ):
            schema = {
                "codec": "struct",
                "type": "object",
                "properties": {type_: {"type": type_, "binaryFormat": binaryFormat}},
            }
            self.round_trip(schema, {type_: value})

        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {"null": {"type": "null"}},
        }
        self.round_trip(schema, {"null": None})

    def test_flat_object(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "int": {"type": "number", "binaryFormat": "i"},
                "float": {"type": "number", "binaryFormat": "d"},
                "null": {"type": "null", "binaryFormat": "3x"},
                "str": {"type": "string", "binaryFormat": "10p"},
                "bool": {"type": "boolean", "binaryFormat": "?"},
            },
        }
        self.round_trip(
            schema, {"null": None, "bool": True, "float": 5.5, "int": 5, "str": "42"}
        )

    def test_nested_object(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "int": {"type": "number", "binaryFormat": "i"},
                "float": {"type": "number", "binaryFormat": "d"},
                "str": {"type": "string", "binaryFormat": "10p"},
                "bool": {"type": "boolean", "binaryFormat": "?"},
                "obj": {
                    "index": 5,
                    "type": "object",
                    "properties": {
                        "int": {"type": "number", "binaryFormat": "i"},
                        "float": {"type": "number", "binaryFormat": "d"},
                        "str": {"type": "string", "binaryFormat": "5p"},
                        "bool": {"type": "boolean", "binaryFormat": "?"},
                    },
                },
            },
        }
        self.round_trip(
            schema,
            {
                "bool": True,
                "float": 5.5,
                "int": 5,
                "str": "42",
                "obj": {"float": 5.78, "int": 9, "bool": False, "str": "41"},
            },
        )

    def test_flat_array(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "array": {
                    "type": "array",
                    "items": {"type": "number", "binaryFormat": "i"},
                }
            },
        }
        self.round_trip(schema, {"array": []})
        self.round_trip(schema, {"array": [1]})
        self.round_trip(schema, {"array": [1, 6, -900]})

        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "array": {
                    "type": "array",
                    "items": {"type": "number", "binaryFormat": "d"},
                }
            },
        }
        self.round_trip(schema, {"array": []})
        self.round_trip(schema, {"array": [1.5]})
        self.round_trip(schema, {"array": [1.5, 6.7, -900.00001]})

    def test_nested_array(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "array": {
                    "type": "array",
                    "items": {
                        "type": "array",
                        "items": {"type": "number", "binaryFormat": "i"},
                    },
                }
            },
        }
        self.round_trip(schema, {"array": [[]]})
        self.round_trip(schema, {"array": [[], []]})
        self.round_trip(schema, {"array": [[1]]})
        self.round_trip(schema, {"array": [[1, 6, -900]]})
        self.round_trip(schema, {"array": [[0, 987, 234903], [1, 6, -900]]})
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "array": {
                    "type": "array",
                    "items": {
                        "type": "array",
                        "items": {"type": "number", "binaryFormat": "d"},
                    },
                }
            },
        }
        self.round_trip(schema, {"array": [[]]})
        self.round_trip(schema, {"array": [[], []]})
        self.round_trip(schema, {"array": [[1.67]]})
        self.round_trip(schema, {"array": [[1.34, 6.56422, -900.0000006]]})
        self.round_trip(
            schema, {"array": [[0.0, 987.123, 234903.123], [1.1235, 6, -900]]}
        )

    def test_array_of_objects(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "array": {
                    "type": "array",
                    "items": {
                        "type": "object",
                        "properties": {
                            "int": {"type": "number", "binaryFormat": "i"},
                            "float": {"type": "number", "binaryFormat": "d"},
                            "padding": {"type": "null", "binaryFormat": "5x"},
                            "str": {"type": "string", "binaryFormat": "10p"},
                            "bool": {"type": "boolean", "binaryFormat": "?"},
                        },
                    },
                }
            },
        }
        self.round_trip(schema, {"array": []})
        self.round_trip(
            schema,
            {
                "array": [
                    {
                        "padding": None,
                        "float": 5.78,
                        "int": 9,
                        "bool": False,
                        "str": "41",
                    }
                ]
            },
        )
        self.round_trip(
            schema,
            {
                "array": [
                    {
                        "padding": None,
                        "float": 5.78,
                        "int": 9,
                        "bool": False,
                        "str": "41",
                    },
                    {
                        "str": "FOO",
                        "int": 7,
                        "bool": True,
                        "float": 45.7,
                        "padding": None,
                    },
                ],
            },
        )

    def test_object_with_array(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "int": {"type": "number", "binaryFormat": "i"},
                "arr": {
                    "index": 2,
                    "type": "array",
                    "items": {"type": "number", "binaryFormat": "i"},
                },
            },
        }
        self.round_trip(schema, {"int": 5, "arr": []})
        self.round_trip(schema, {"int": 5, "arr": [5]})
        self.round_trip(schema, {"arr": [5, 6, 7], "int": 5})

    def test_array_length_format(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "array": {
                    "type": "array",
                    "arrayLengthFormat": "B",
                    "items": {"type": "number", "binaryFormat": "H"},
                }
            },
        }
        self.round_trip(schema, {"array": []})
        self.round_trip(schema, {"array": [1]})
        self.round_trip(schema, {"array": list(range(255))})

    def test_string_encoding(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "string": {
                    "type": "string",
                    "stringEncoding": "utf-16",
                    "binaryFormat": "40p",
                }
            },
        }
        self.round_trip(schema, {"string": "Test string"})

    def test_ordering_of_fields(self):
        row_data = {
            "null": None,
            "bool": True,
            "float": -1.8440714901698642e18,
            "int": 5,
            "str": "foo",
        }
        alpha_ordered_encoded = b"\x01\xaa\xbb\xcc\xdd\x05\x00\x00\x00\x03foo"
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "null": {"type": "null", "binaryFormat": "3x"},
                "float": {"type": "number", "binaryFormat": "f"},
                "bool": {"type": "boolean", "binaryFormat": "?"},
                "int": {"type": "number", "binaryFormat": "b"},
                "str": {"type": "string", "binaryFormat": "4p"},
            },
        }
        alpha_ordered_encoded = b"\x01\xaa\xbb\xcc\xdd\x05\x00\x00\x00\x03foo"
        ms = metadata.MetadataSchema(schema)
        assert ms.validate_and_encode_row(row_data) == alpha_ordered_encoded
        assert ms.decode_row(alpha_ordered_encoded) == row_data
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "null": {"type": "null", "binaryFormat": "3x", "index": 0},
                "float": {"type": "number", "binaryFormat": "f", "index": 1},
                "bool": {"type": "boolean", "binaryFormat": "?", "index": 2},
                "int": {"type": "number", "binaryFormat": "b", "index": 3},
                "str": {"type": "string", "binaryFormat": "4p", "index": 4},
            },
        }
        index_order_encoded = b"\x00\x00\x00\xaa\xbb\xcc\xdd\x01\x05\x03foo"
        ms = metadata.MetadataSchema(schema)
        assert ms.validate_and_encode_row(row_data) == index_order_encoded
        assert ms.decode_row(index_order_encoded) == row_data


class TestStructCodecErrors:
    def encode(self, schema, row_data):
        ms = metadata.MetadataSchema(schema)
        ms.validate_and_encode_row(row_data)

    def test_missing_and_extra_property(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "int": {"type": "number", "binaryFormat": "i"},
                "float": {"type": "number", "binaryFormat": "d"},
            },
        }
        with pytest.raises(
            exceptions.MetadataValidationError, match="'int' is a required property"
        ):
            self.encode(schema, {"float": 5.5})
        with pytest.raises(
            exceptions.MetadataValidationError,
            match="Additional properties are not allowed",
        ):
            self.encode(
                schema, {"float": 5.5, "int": 9, "extra": "I really shouldn't be here"}
            )

    def test_bad_schema_union_type(self):
        schema = {"codec": "struct", "type": ["object", "number"], "binaryFormat": "d"}
        with pytest.raises(
            exceptions.MetadataSchemaValidationError, match="is not one of"
        ):
            metadata.MetadataSchema(schema)
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {"test": {"type": ["number", "string"], "binaryFormat": "d"}},
        }
        with pytest.raises(
            exceptions.MetadataSchemaValidationError, match="is not one of"
        ):
            metadata.MetadataSchema(schema)

    def test_bad_schema_hetrogeneous_array(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "array": {
                    "type": "array",
                    "items": [{"type": "number"}, {"type": "string"}],
                }
            },
        }
        with pytest.raises(
            exceptions.MetadataSchemaValidationError, match="is not of type 'object'"
        ):
            metadata.MetadataSchema(schema)

    def test_bad_binary_format(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {"int": {"type": "number", "binaryFormat": "int"}},
        }
        with pytest.raises(
            exceptions.MetadataSchemaValidationError, match="does not match"
        ):
            metadata.MetadataSchema(schema)
        # Can't specify endianness
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {"int": {"type": "number", "binaryFormat": ">b"}},
        }
        with pytest.raises(
            exceptions.MetadataSchemaValidationError, match="does not match"
        ):
            metadata.MetadataSchema(schema)
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {"null": {"type": "null", "binaryFormat": "l"}},
        }
        with pytest.raises(
            exceptions.MetadataSchemaValidationError,
            match="null type binaryFormat must be padding",
        ):
            metadata.MetadataSchema(schema)

    def test_bad_array_length_format(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {"array": {"type": "array", "arrayLengthFormat": "b"}},
        }
        with pytest.raises(
            exceptions.MetadataSchemaValidationError, match="does not match"
        ):
            metadata.MetadataSchema(schema)

    def test_missing_binary_format(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {"int": {"type": "number"}},
        }
        with pytest.raises(
            exceptions.MetadataSchemaValidationError,
            match="number type must have binaryFormat set",
        ):
            metadata.MetadataSchema(schema)

    def test_bad_string_encoding(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "string": {
                    "type": "string",
                    "binaryFormat": "5s",
                    "stringEncoding": 58,
                }
            },
        }
        with pytest.raises(
            exceptions.MetadataSchemaValidationError, match="is not of type"
        ):
            metadata.MetadataSchema(schema)

    def test_bad_null_terminated(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "string": {
                    "type": "string",
                    "binaryFormat": "5s",
                    "nullTerminated": 58,
                }
            },
        }
        with pytest.raises(
            exceptions.MetadataSchemaValidationError, match="is not of type"
        ):
            metadata.MetadataSchema(schema)

    def test_bad_no_length_encoding_exhaust_buffer(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "string": {
                    "type": "string",
                    "binaryFormat": "5s",
                    "noLengthEncodingExhaustBuffer": 58,
                }
            },
        }
        with pytest.raises(
            exceptions.MetadataSchemaValidationError, match="is not of type"
        ):
            metadata.MetadataSchema(schema)

    def test_too_long_array(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "array": {
                    "type": "array",
                    "arrayLengthFormat": "B",
                    "items": {"type": "number", "binaryFormat": "I"},
                },
            },
        }
        data = {"array": list(range(255))}
        metadata.MetadataSchema(schema).validate_and_encode_row(data)
        data2 = {"array": list(range(256))}
        with pytest.raises(
            ValueError,
            match="Couldn't pack array size - it is likely too long for the"
            " specified arrayLengthFormat",
        ):
            metadata.MetadataSchema(schema).validate_and_encode_row(data2)

    def test_additional_properties(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "additional_properties": True,
            "properties": {},
        }
        with pytest.raises(
            ValueError, match="Struct codec does not support additional_properties"
        ):
            metadata.MetadataSchema(schema)

    def test_unrequired_property_needs_default(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "int": {"type": "number", "binaryFormat": "i"},
                "float": {"type": "number", "binaryFormat": "d"},
            },
            "required": ["float"],
        }
        with pytest.raises(
            exceptions.MetadataSchemaValidationError,
            match="Optional property 'int' must have a default value",
        ):
            metadata.MetadataSchema(schema)

    def test_no_default_implies_required(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "int": {"type": "number", "binaryFormat": "i", "default": 5},
                "float": {"type": "number", "binaryFormat": "d"},
            },
        }
        self.encode(schema, {"float": 5.5})
        with pytest.raises(
            exceptions.MetadataValidationError, match="'float' is a required property"
        ):
            self.encode(schema, {})


class TestSLiMDecoding:
    """
    Test with byte strings copied from a SLiM tree sequence
    """

    def test_node(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "genomeID": {"type": "integer", "binaryFormat": "q", "index": 0},
                "isNull": {"type": "boolean", "binaryFormat": "?", "index": 1},
                "genomeType": {"type": "integer", "binaryFormat": "B", "index": 2},
            },
        }
        for example, expected in [
            (
                b"E,\x00\x00\x00\x00\x00\x00\x00\x01",
                {"genomeID": 11333, "genomeType": 1, "isNull": False},
            ),
            (
                b"\xdd.\x00\x00\x00\x00\x00\x00\x01\x00",
                {"genomeID": 11997, "genomeType": 0, "isNull": True},
            ),
        ]:
            assert metadata.MetadataSchema(schema).decode_row(example) == expected

    def test_individual(self):
        schema = {
            "codec": "struct",
            "type": ["object", "null"],
            "properties": {
                "pedigreeID": {"type": "integer", "binaryFormat": "q", "index": 1},
                "age": {"type": "integer", "binaryFormat": "i", "index": 2},
                "subpopulationID": {
                    "type": "integer",
                    "binaryFormat": "i",
                    "index": 3,
                },
                "sex": {"type": "integer", "binaryFormat": "i", "index": 4},
                "flags": {"type": "integer", "binaryFormat": "I", "index": 5},
            },
        }
        for example, expected in [
            (
                b"\x17\x99\x07\x00\x00\x00\x00\x00\x05\x00\x01\x00\x03\x00\x00\x00\x01"
                b"\x00\x00\x00\x00\x10\x00\x00",
                {
                    "age": 65541,
                    "flags": 4096,
                    "pedigreeID": 497943,
                    "sex": 1,
                    "subpopulationID": 3,
                },
            ),
            (b"", None),
            (
                b"\x18\x99\x07\x00\x00\x00\x00\x00\x05\x00\x00\x00\x01\x00\x00\x00\x01"
                b"\x00\x00\x00\x00\x00\x00\x00",
                {
                    "age": 5,
                    "flags": 0,
                    "pedigreeID": 497944,
                    "sex": 1,
                    "subpopulationID": 1,
                },
            ),
        ]:
            assert metadata.MetadataSchema(schema).decode_row(example) == expected

    def test_mutation(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "stacked_mutation_array": {
                    "type": "array",
                    "noLengthEncodingExhaustBuffer": True,
                    "items": {
                        "type": "object",
                        "properties": {
                            "mutationTypeID": {
                                "type": "integer",
                                "binaryFormat": "i",
                                "index": 1,
                            },
                            "selectionCoeff": {
                                "type": "number",
                                "binaryFormat": "f",
                                "index": 2,
                            },
                            "subpopulationID": {
                                "type": "integer",
                                "binaryFormat": "i",
                                "index": 3,
                            },
                            "originGeneration": {
                                "type": "integer",
                                "binaryFormat": "i",
                                "index": 4,
                            },
                            "nucleotide": {
                                "type": "integer",
                                "binaryFormat": "b",
                                "index": 5,
                            },
                        },
                    },
                }
            },
        }

        for example, expected in [
            (
                b"\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\xd8\x03\x00\x00\xff",
                [
                    {
                        "mutationTypeID": 1,
                        "selectionCoeff": 0.0,
                        "subpopulationID": 1,
                        "originGeneration": 984,
                        "nucleotide": -1,
                    }
                ],
            ),
            (
                b"\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\xc8\x03\x00\x00\xff"
                b"\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x94\x01\x00\x00\xff",
                [
                    {
                        "mutationTypeID": 1,
                        "selectionCoeff": 0.0,
                        "subpopulationID": 1,
                        "originGeneration": 968,
                        "nucleotide": -1,
                    },
                    {
                        "mutationTypeID": 1,
                        "selectionCoeff": 0.0,
                        "subpopulationID": 1,
                        "originGeneration": 404,
                        "nucleotide": -1,
                    },
                ],
            ),
            (
                b"\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\xd1\x03\x00\x00\xff"
                b"\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\xb1\x02\x00\x00\xff"
                b"\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\xdf\x01\x00\x00\xff"
                b"\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\xbc\x00\x00\x00\xff",
                [
                    {
                        "mutationTypeID": 1,
                        "selectionCoeff": 0.0,
                        "subpopulationID": 1,
                        "originGeneration": 977,
                        "nucleotide": -1,
                    },
                    {
                        "mutationTypeID": 1,
                        "selectionCoeff": 0.0,
                        "subpopulationID": 1,
                        "originGeneration": 689,
                        "nucleotide": -1,
                    },
                    {
                        "mutationTypeID": 1,
                        "selectionCoeff": 0.0,
                        "subpopulationID": 1,
                        "originGeneration": 479,
                        "nucleotide": -1,
                    },
                    {
                        "mutationTypeID": 1,
                        "selectionCoeff": 0.0,
                        "subpopulationID": 1,
                        "originGeneration": 188,
                        "nucleotide": -1,
                    },
                ],
            ),
        ]:
            assert (
                metadata.MetadataSchema(schema).decode_row(example)[
                    "stacked_mutation_array"
                ]
                == expected
            )

    def test_population(self):
        schema = {
            "codec": "struct",
            "type": "object",
            "properties": {
                "subpopulationID": {
                    "type": "integer",
                    "binaryFormat": "i",
                    "index": 0,
                },
                "femaleCloneFraction": {
                    "type": "number",
                    "binaryFormat": "d",
                    "index": 1,
                },
                "maleCloneFraction": {
                    "type": "number",
                    "binaryFormat": "d",
                    "index": 2,
                },
                "sexRatio": {"type": "number", "binaryFormat": "d", "index": 3},
                "boundsX0": {"type": "number", "binaryFormat": "d", "index": 4},
                "boundsX1": {"type": "number", "binaryFormat": "d", "index": 5},
                "boundsY0": {"type": "number", "binaryFormat": "d", "index": 6},
                "boundsY1": {"type": "number", "binaryFormat": "d", "index": 7},
                "boundsZ0": {"type": "number", "binaryFormat": "d", "index": 8},
                "boundsZ1": {"type": "number", "binaryFormat": "d", "index": 9},
                "migrationRecCount": {
                    "type": "integer",
                    "binaryFormat": "d",
                    "index": 10,
                },
            },
        }
        example = (
            b"\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
            b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
            b"\x00\x00\xe0?\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
            b"\x00\xf0?\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
            b"\xf0?\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\xf0"
            b"?\x00\x00\x00\x00"
        )
        expected = {
            "boundsX0": 0.5,
            "boundsX1": 0.0,
            "boundsY0": 1.0,
            "boundsY1": 0.0,
            "boundsZ0": 1.0,
            "boundsZ1": 0.0,
            "femaleCloneFraction": 0.0,
            "maleCloneFraction": 0.0,
            "migrationRecCount": 1.0,
            "sexRatio": 0.0,
            "subpopulationID": 1,
        }
        assert metadata.MetadataSchema(schema).decode_row(example) == expected


class TestTableCollectionEquality:
    def test_equality(self):
        ts = msprime.simulate(10, random_seed=42)
        tables = ts.dump_tables()
        tables2 = ts.dump_tables()
        schema = collections.OrderedDict(
            codec="json",
            title="Example Metadata",
            type="object",
            properties=collections.OrderedDict(
                one={"type": "string"}, two={"type": "number"}
            ),
            required=["one", "two"],
            additionalProperties=False,
        )
        schema2 = collections.OrderedDict(
            type="object",
            properties=collections.OrderedDict(
                two={"type": "number"}, one={"type": "string"}
            ),
            required=["one", "two"],
            additionalProperties=False,
            title="Example Metadata",
            codec="json",
        )
        tables.metadata_schema = metadata.MetadataSchema(schema)
        assert tables != tables2
        tables2.metadata_schema = metadata.MetadataSchema(schema2)
        tables.assert_equals(tables2)
        tables.metadata = collections.OrderedDict(one="tree", two=5)
        assert tables != tables2
        tables2.metadata = collections.OrderedDict(two=5, one="tree")
        tables.assert_equals(tables2)

    def test_fixing_uncanonical(self):
        ts = msprime.simulate(10, random_seed=42)
        tables = ts.dump_tables()
        schema = collections.OrderedDict(
            codec="json",
            title="Example Metadata",
            type="object",
            properties=collections.OrderedDict(
                one={"type": "string"}, two={"type": "number"}
            ),
            required=["one", "two"],
            additionalProperties=False,
        )
        # Set with low-level to emulate loading.
        tables._ll_tables.metadata_schema = json.dumps(schema)
        assert tables._ll_tables.metadata_schema != tskit.canonical_json(schema)
        tables.metadata_schema = tables.metadata_schema
        assert tables._ll_tables.metadata_schema == tskit.canonical_json(schema)


--- ../../tskit/python/tests/test_genotypes.py ---


"""
Test cases for generating genotypes/haplotypes.
"""
import itertools
import logging
import random
import re
import textwrap
from xml.etree import ElementTree

import msprime
import numpy as np
import pytest

import tests
import tests.test_wright_fisher as wf
import tests.tsutil as tsutil
import tskit
from tests.test_highlevel import get_example_tree_sequences
from tskit import exceptions
from tskit.genotypes import allele_remap

# ↑ See https://github.com/tskit-dev/tskit/issues/1804 for when
# we can remove this.

# TODO replace this with a call to
# example_tree_sequences(discrete_genome=True, snps_only=True)


@tests.cached_example
def get_example_discrete_genome_tree_sequences():
    ret = []
    for ts in get_example_tree_sequences(pytest_params=False):
        if ts.discrete_genome:
            snps = all(len(site.ancestral_state) == 1 for site in ts.sites()) and all(
                len(mut.derived_state) == 1 for mut in ts.mutations()
            )
            if snps:
                ret.append(ts)
    return ret


def naive_get_ancestral_haplotypes(ts):
    """
    Simple implementation using tree traversals. Note that this definition
    won't work when we have topology that's not reachable from a root,
    but this seems more trouble than it's worth dealing with.
    """
    A = np.zeros((ts.num_nodes, ts.num_sites), dtype=np.int8)
    A[:] = tskit.MISSING_DATA
    for t in ts.trees():
        for site in t.sites():
            alleles = {site.ancestral_state: 0}
            for u in t.nodes():
                A[u, site.id] = 0
            j = 1
            for mutation in site.mutations:
                if mutation.derived_state not in alleles:
                    alleles[mutation.derived_state] = j
                    j += 1
                for u in t.nodes(mutation.node):
                    A[u, site.id] = alleles[mutation.derived_state]
    return A


class TestGetAncestralHaplotypes:
    """
    Tests for the engine to the actual ancestors from a simulation.
    """

    def verify(self, ts):
        A = naive_get_ancestral_haplotypes(ts)
        # To detect missing data in ancestors we must set all nodes
        # to be samples
        tables = ts.dump_tables()
        nodes = tables.nodes
        flags = nodes.flags[:]
        flags[:] = 1
        nodes.set_columns(time=nodes.time, flags=flags)
        ts = tables.tree_sequence()
        B = ts.genotype_matrix().T
        assert np.array_equal(A, B)

    def test_single_tree(self):
        ts = msprime.simulate(5, mutation_rate=1, random_seed=234)
        self.verify(ts)

    def test_many_trees(self):
        ts = msprime.simulate(
            8, recombination_rate=10, mutation_rate=10, random_seed=234
        )
        assert ts.num_trees > 1
        assert ts.num_sites > 1
        self.verify(ts)

    def test_single_tree_jukes_cantor(self):
        ts = msprime.simulate(6, random_seed=1, mutation_rate=1)
        ts = tsutil.jukes_cantor(ts, 20, 1, seed=10)
        self.verify(ts)

    def test_single_tree_multichar_mutations(self):
        ts = msprime.simulate(6, random_seed=1, mutation_rate=1)
        ts = tsutil.insert_multichar_mutations(ts)
        self.verify(ts)

    def test_many_trees_infinite_sites(self):
        ts = msprime.simulate(6, recombination_rate=2, mutation_rate=2, random_seed=1)
        assert ts.num_sites > 0
        assert ts.num_trees > 2
        self.verify(ts)

    def test_wright_fisher_initial_generation(self):
        tables = wf.wf_sim(
            6, 5, seed=3, deep_history=True, initial_generation_samples=True, num_loci=2
        )
        tables.sort()
        tables.simplify()
        ts = msprime.mutate(tables.tree_sequence(), rate=0.08, random_seed=2)
        assert ts.num_sites > 0
        self.verify(ts)

    def test_wright_fisher_simplified(self):
        tables = wf.wf_sim(
            9,
            10,
            seed=1,
            deep_history=True,
            initial_generation_samples=False,
            num_loci=5,
        )
        tables.sort()
        ts = tables.tree_sequence().simplify()
        ts = msprime.mutate(ts, rate=0.2, random_seed=1234)
        assert ts.num_sites > 0
        self.verify(ts)

    def test_empty_ts(self):
        tables = tskit.TableCollection(1.0)
        for _ in range(10):
            tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
        ts = tables.tree_sequence()
        self.verify(ts)


def isolated_samples_genotype_matrix(ts):
    """
    Returns the genotype matrix for the specified tree sequence
    where isolated samples are marked with MISSING_DATA.
    """
    G = ts.genotype_matrix()
    samples = ts.samples()
    sample_index_map = np.zeros(ts.num_nodes, dtype=int) - 1
    for index, sample in enumerate(samples):
        sample_index_map[sample] = index
    for tree in ts.trees():
        for site in tree.sites():
            for root in tree.roots:
                # An isolated sample is any root that has no children.
                if tree.left_child(root) == -1:
                    assert sample_index_map[root] != -1
                    G[site.id, sample_index_map[root]] = -1
    return G


class TestVariantGenerator:
    """
    Tests the variants() method to ensure the output is consistent.
    """

    def get_tree_sequence(self):
        ts = msprime.simulate(
            10, length=10, recombination_rate=1, mutation_rate=10, random_seed=3
        )
        assert ts.get_num_mutations() > 10
        return ts

    def test_dtype(self):
        ts = self.get_tree_sequence()
        for var in ts.variants():
            assert var.genotypes.dtype == np.int32

    def test_dtype_conversion(self):
        # Check if we hit any issues if we assume the variants are uint8
        # as they were prior to version 0.2.0
        ts = self.get_tree_sequence()
        G = ts.genotype_matrix().astype(np.uint8)
        assert G.dtype == np.uint8
        for var in ts.variants():
            assert np.array_equal(G[var.index], var.genotypes)
            assert np.all(G[var.index] == var.genotypes)
            assert [var.alleles[g] for g in var.genotypes] == [
                var.alleles[g] for g in G[var.index]
            ]
            G[var.index, :] = var.genotypes
            assert np.array_equal(G[var.index], var.genotypes)

    def test_multichar_alleles(self):
        ts = tsutil.insert_multichar_mutations(self.get_tree_sequence())
        for var in ts.variants():
            assert len(var.alleles) == 2
            assert var.site.ancestral_state == var.alleles[0]
            assert var.site.mutations[0].derived_state == var.alleles[1]
            assert all(0 <= var.genotypes)
            assert all(var.genotypes <= 1)

    def test_many_alleles(self):
        ts = self.get_tree_sequence()
        tables = ts.dump_tables()
        tables.sites.clear()
        tables.mutations.clear()
        # This gives us a total of 360 permutations.
        alleles = list(map("".join, itertools.permutations("ABCDEF", 4)))
        assert len(alleles) > 127
        tables.sites.add_row(0, alleles[0])
        parent = -1
        num_alleles = 1
        for allele in alleles[1:]:
            ts = tables.tree_sequence()
            var = next(ts.variants())
            assert not var.has_missing_data
            assert var.num_alleles == num_alleles
            assert len(var.alleles) == num_alleles
            assert list(var.alleles) == alleles[:num_alleles]
            assert var.alleles[var.genotypes[0]] == alleles[num_alleles - 1]
            for u in ts.samples():
                if u != 0:
                    assert var.alleles[var.genotypes[u]] == alleles[0]
            tables.mutations.add_row(0, 0, allele, parent=parent)
            parent += 1
            num_alleles += 1

    def test_many_alleles_missing_data(self):
        ts = self.get_tree_sequence()
        tables = ts.dump_tables()
        tables.sites.clear()
        tables.mutations.clear()
        # Add an isolated sample
        tables.nodes.add_row(flags=1, time=0)
        # This gives us a total of 360 permutations.
        alleles = list(map("".join, itertools.permutations("ABCDEF", 4)))
        assert len(alleles) > 127
        tables.sites.add_row(0, alleles[0])
        parent = -1
        num_alleles = 1
        for allele in alleles[1:]:
            ts = tables.tree_sequence()

            var = next(ts.variants())
            assert var.has_missing_data
            assert var.num_alleles == num_alleles
            assert len(var.alleles) == num_alleles + 1
            assert list(var.alleles)[:-1] == alleles[:num_alleles]
            assert var.alleles[-1] is None
            assert var.alleles[var.genotypes[0]] == alleles[num_alleles - 1]
            assert var.genotypes[-1] == -1
            samples = ts.samples()
            for u in samples[:-1]:
                if u != 0:
                    assert var.alleles[var.genotypes[u]] == alleles[0]
            tables.mutations.add_row(0, 0, allele, parent=parent)
            parent += 1
            num_alleles += 1

    def test_site_information(self):
        ts = self.get_tree_sequence()
        for site, variant in zip(ts.sites(), ts.variants()):
            assert site.position == variant.position
            assert site == variant.site

    def test_no_mutations(self):
        ts = msprime.simulate(10)
        assert ts.get_num_mutations() == 0
        variants = list(ts.variants())
        assert len(variants) == 0

    @pytest.mark.parametrize("samples", [None, [1, 2], [2, 4], []])
    def test_genotype_matrix(self, samples):
        ts = self.get_tree_sequence()
        num_samples = ts.num_samples if samples is None else len(samples)
        G = np.empty((ts.num_sites, num_samples), dtype=np.int32)
        for v in ts.variants(samples=samples):
            G[v.index, :] = v.genotypes
        if samples is None:
            G2 = ts.genotype_matrix()
        else:
            G2 = ts.genotype_matrix(samples=samples)
        assert np.array_equal(G, G2)
        assert G2.dtype == np.int32

    def test_recurrent_mutations_over_samples(self):
        ts = self.get_tree_sequence()
        tables = ts.dump_tables()
        tables.sites.clear()
        tables.mutations.clear()
        num_sites = 5
        for j in range(num_sites):
            tables.sites.add_row(
                position=j * ts.sequence_length / num_sites, ancestral_state="0"
            )
            for u in range(ts.sample_size):
                tables.mutations.add_row(site=j, node=u, derived_state="1")
        ts = tables.tree_sequence()
        variants = list(ts.variants())
        assert len(variants) == num_sites
        for site, variant in zip(ts.sites(), variants):
            assert site.position == variant.position
            assert site == variant.site
            assert site.id == variant.index
            assert variant.alleles == ("0", "1")
            assert np.all(variant.genotypes == np.ones(ts.sample_size))

    def test_silent_mutations(self):
        ts = self.get_tree_sequence()
        tree = next(ts.trees())
        tables = ts.dump_tables()
        for u in tree.nodes():
            for sample in tree.samples(u):
                if sample != u:
                    tables.sites.clear()
                    tables.mutations.clear()
                    site = tables.sites.add_row(position=0, ancestral_state="0")
                    tables.mutations.add_row(site=site, node=u, derived_state="1")
                    tables.mutations.add_row(site=site, node=sample, derived_state="1")
                    ts_new = tables.tree_sequence()
                    assert all([v.genotypes[sample] == 1 for v in ts_new.variants()])

    def test_zero_samples(self):
        ts = self.get_tree_sequence()
        for var1, var2 in zip(ts.variants(), ts.variants(samples=[])):
            assert var1.site == var2.site
            assert var1.alleles == var2.alleles
            assert var2.genotypes.shape[0] == 0

    def test_samples(self):
        n = 4
        ts = msprime.simulate(
            n, length=5, recombination_rate=1, mutation_rate=5, random_seed=2
        )
        assert ts.num_sites > 1
        samples = list(range(n))
        # Generate all possible sample lists.
        for j in range(n + 1):
            for s in itertools.permutations(samples, j):
                s = np.array(s, dtype=np.int32)
                count = 0
                for var1, var2 in zip(ts.variants(), ts.variants(samples=s)):
                    assert var1.site == var2.site
                    assert var1.alleles == var2.alleles
                    assert np.array_equal(var1.samples, ts.samples())
                    assert np.array_equal(var2.samples, s)
                    assert var2.genotypes.shape == (len(s),)
                    assert np.array_equal(var1.genotypes[s], var2.genotypes)
                    count += 1
                assert count == ts.num_sites

    def test_samples_64bit(self):
        ts = msprime.simulate(4, length=5, mutation_rate=5, random_seed=2)
        s = np.where(ts.nodes_time == 0)[0]  # normally returns 64 bit ints
        next(ts.variants(samples=s))
        s = np.array(s, dtype=np.int64)  # cast just to make sure
        next(ts.variants(samples=s))

    def test_samples_missing_data(self):
        n = 4
        ts = msprime.simulate(
            n, length=5, recombination_rate=1, mutation_rate=5, random_seed=2
        )
        assert ts.num_sites > 1
        tables = ts.dump_tables()
        tables.delete_intervals([[0.5, 0.6]])
        tables.sites.add_row(0.5, ancestral_state="0")
        tables.sort()
        ts = tables.tree_sequence()
        samples = list(range(n))
        # Generate all possible sample lists.
        for j in range(1, n + 1):
            for s in itertools.permutations(samples, j):
                s = np.array(s, dtype=np.int32)
                count = 0
                for var1, var2 in zip(ts.variants(), ts.variants(samples=s)):
                    assert var1.site == var2.site
                    assert var1.alleles == var2.alleles
                    assert var2.genotypes.shape == (len(s),)
                    assert np.array_equal(var1.genotypes[s], var2.genotypes)
                    count += 1
                assert count == ts.num_sites

    def test_non_sample_samples(self):
        # We don't have to use sample nodes. This does make the terminology confusing
        # but it's probably still the best option.
        ts = msprime.simulate(
            10, length=5, recombination_rate=1, mutation_rate=5, random_seed=2
        )
        tables = ts.dump_tables()
        tables.nodes.set_columns(
            flags=np.zeros_like(tables.nodes.flags) + tskit.NODE_IS_SAMPLE,
            time=tables.nodes.time,
        )
        all_samples_ts = tables.tree_sequence()
        assert all_samples_ts.num_samples == ts.num_nodes

        count = 0
        samples = range(ts.num_nodes)
        for var1, var2 in zip(
            all_samples_ts.variants(isolated_as_missing=False),
            ts.variants(samples=samples, isolated_as_missing=False),
        ):
            assert var1.site == var2.site
            assert var1.alleles == var2.alleles
            assert var2.genotypes.shape == (len(samples),)
            assert np.array_equal(var1.genotypes, var2.genotypes)
            count += 1
        assert count == ts.num_sites

    def verify_jukes_cantor(self, ts):
        assert np.array_equal(ts.genotype_matrix(), ts.genotype_matrix())
        tree = ts.first()
        for variant in ts.variants():
            assert not variant.has_missing_data
            mutations = {
                mutation.node: mutation.derived_state
                for mutation in variant.site.mutations
            }
            for sample_index, u in enumerate(ts.samples()):
                while u not in mutations and u != tskit.NULL:
                    u = tree.parent(u)
                state1 = mutations.get(u, variant.site.ancestral_state)
                state2 = variant.alleles[variant.genotypes[sample_index]]
                assert state1 == state2

    def test_jukes_cantor_n5(self):
        ts = msprime.simulate(5, random_seed=2)
        ts = tsutil.jukes_cantor(ts, 5, 1, seed=2)
        self.verify_jukes_cantor(ts)

    def test_jukes_cantor_n20(self):
        ts = msprime.simulate(20, random_seed=2)
        ts = tsutil.jukes_cantor(ts, 5, 1, seed=2)
        self.verify_jukes_cantor(ts)

    def test_zero_edge_missing_data(self):
        ts = msprime.simulate(10, random_seed=2, mutation_rate=2)
        tables = ts.dump_tables()
        tables.keep_intervals([[0.25, 0.75]])
        # add some sites in the deleted regions
        tables.sites.add_row(0.1, "A")
        tables.sites.add_row(0.2, "A")
        tables.sites.add_row(0.8, "A")
        tables.sites.add_row(0.9, "A")
        tables.sort()
        ts = tables.tree_sequence()
        Gnm = ts.genotype_matrix(isolated_as_missing=False)
        assert np.all(Gnm[0] == 0)
        assert np.all(Gnm[1] == 0)
        assert np.all(Gnm[-1] == 0)
        assert np.all(Gnm[-2] == 0)
        Gm = isolated_samples_genotype_matrix(ts)
        assert np.all(Gm[0] == -1)
        assert np.all(Gm[1] == -1)
        assert np.all(Gm[-1] == -1)
        assert np.all(Gm[-2] == -1)
        Gm2 = ts.genotype_matrix(isolated_as_missing=True)
        assert np.array_equal(Gm, Gm2)

        # Test deprecated param

        with pytest.warns(FutureWarning):
            Gi = ts.genotype_matrix(impute_missing_data=True)
        assert np.array_equal(Gnm, Gi)
        with pytest.warns(FutureWarning):
            Gni = ts.genotype_matrix(impute_missing_data=False)
        assert np.array_equal(Gm, Gni)

        with pytest.warns(FutureWarning):
            G = ts.genotype_matrix(isolated_as_missing=False, impute_missing_data=True)
        assert np.array_equal(Gnm, G)
        with pytest.warns(FutureWarning):
            G = ts.genotype_matrix(isolated_as_missing=True, impute_missing_data=False)
        assert np.array_equal(Gm, G)

    def test_empty_ts_missing_data(self):
        tables = tskit.TableCollection(1.0)
        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
        tables.sites.add_row(0.5, "A")
        ts = tables.tree_sequence()
        variants = list(ts.variants())
        assert len(variants) == 1
        var = variants[0]
        assert var.alleles == ("A", None)
        assert var.num_alleles == 1
        assert np.all(var.genotypes == -1)

    def test_empty_ts_incomplete_samples(self):
        # https://github.com/tskit-dev/tskit/issues/776
        tables = tskit.TableCollection(1.0)
        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
        tables.sites.add_row(0.5, "A")
        ts = tables.tree_sequence()
        variants = list(ts.variants(samples=[0]))
        assert list(variants[0].genotypes) == [-1]
        variants = list(ts.variants(samples=[1]))
        assert list(variants[0].genotypes) == [-1]

    def test_missing_data_samples(self):
        tables = tskit.TableCollection(1.0)
        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
        tables.sites.add_row(0.5, "A")
        tables.mutations.add_row(0, 0, "T")
        ts = tables.tree_sequence()

        # If we have no samples we still get a list of variants.
        variants = list(ts.variants(samples=[]))
        assert len(variants[0].genotypes) == 0
        assert not variants[0].has_missing_data
        assert variants[0].alleles == ("A", "T")

        # If we have a single sample that's not missing, there's no
        # missing data.
        variants = list(ts.variants(samples=[0]))
        assert len(variants[0].genotypes) == 1
        assert variants[0].genotypes[0] == 1
        assert not variants[0].has_missing_data
        assert variants[0].alleles == ("A", "T")

        # If we have a single sample that is missing, there is
        # missing data.
        variants = list(ts.variants(samples=[1]))
        assert len(variants[0].genotypes) == 1
        assert variants[0].genotypes[0] == -1
        assert variants[0].has_missing_data
        assert variants[0].alleles == ("A", "T", None)

    def test_mutation_over_isolated_sample_not_missing(self):
        tables = tskit.TableCollection(1.0)
        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
        tables.sites.add_row(0.5, "A")
        tables.mutations.add_row(0, 0, "T")
        ts = tables.tree_sequence()
        variants = list(ts.variants())
        assert len(variants) == 1
        var = variants[0]
        assert var.alleles == ("A", "T", None)
        assert var.num_alleles == 2
        assert list(var.genotypes) == [1, -1]

    def test_multiple_mutations_over_isolated_sample(self):
        tables = tskit.TableCollection(1.0)
        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
        tables.sites.add_row(0.5, "A")
        tables.mutations.add_row(0, 0, "T")
        tables.mutations.add_row(0, 0, "G", parent=0)
        ts = tables.tree_sequence()
        variants = list(ts.variants())
        assert len(variants) == 1
        var = variants[0]
        assert var.alleles == ("A", "T", "G", None)
        assert var.num_alleles == 3
        assert len(var.site.mutations) == 2
        assert list(var.genotypes) == [2, -1]

    def test_snipped_tree_sequence_missing_data(self):
        ts = msprime.simulate(
            10, length=10, recombination_rate=0.1, mutation_rate=10, random_seed=3
        )
        tables = ts.dump_tables()
        tables.delete_intervals([[4, 6]], simplify=False)
        tables.sites.add_row(4, ancestral_state="0")
        tables.sites.add_row(5, ancestral_state="0")
        tables.sites.add_row(5.999999, ancestral_state="0")
        tables.sort()
        ts = tables.tree_sequence()
        G = ts.genotype_matrix()
        num_missing = 0
        for var in ts.variants():
            if 4 <= var.site.position < 6:
                assert var.has_missing_data
                assert np.all(var.genotypes == tskit.MISSING_DATA)
                num_missing += 1
            else:
                assert not var.has_missing_data
                assert np.all(var.genotypes != tskit.MISSING_DATA)
            assert np.array_equal(var.genotypes, G[var.site.id])
        assert num_missing == 3

        G = ts.genotype_matrix(isolated_as_missing=False)
        for var in ts.variants(isolated_as_missing=False):
            if 4 <= var.site.position < 6:
                assert not var.has_missing_data
                assert np.all(var.genotypes == 0)
            else:
                assert not var.has_missing_data
                assert np.all(var.genotypes != tskit.MISSING_DATA)
            assert np.array_equal(var.genotypes, G[var.site.id])

    def test_snipped_tree_sequence_mutations_over_isolated(self):
        ts = msprime.simulate(
            10, length=10, recombination_rate=0.1, mutation_rate=10, random_seed=3
        )
        tables = ts.dump_tables()
        tables.delete_intervals([[4, 6]], simplify=False)
        missing_site = tables.sites.add_row(4, ancestral_state="0")
        tables.mutations.add_row(missing_site, node=0, derived_state="1")
        # Add another site in which all the samples are marked with a mutation
        # to the ancestral state. Note: this would normally not be allowed because
        # there's not state change. However, this allows us to mark a sample
        # as not-missing, so it's an important feature.
        missing_site = tables.sites.add_row(5, ancestral_state="0")
        for u in range(10):
            tables.mutations.add_row(missing_site, node=u, derived_state="0")
        tables.sort()
        ts = tables.tree_sequence()
        G = ts.genotype_matrix()
        missing_found = False
        non_missing_found = False
        for var in ts.variants():
            if var.site.position == 4:
                assert var.has_missing_data
                assert var.genotypes[0] == 1
                assert np.all(var.genotypes[1:] == tskit.MISSING_DATA)
                missing_found += 1
            elif var.site.position == 5:
                assert not var.has_missing_data
                assert np.all(var.genotypes == 0)
                non_missing_found = 1
            else:
                assert not var.has_missing_data
                assert np.all(var.genotypes != tskit.MISSING_DATA)
            assert np.array_equal(var.genotypes, G[var.site.id])
        assert non_missing_found
        assert missing_found

    def get_missing_data_ts(self):
        tables = tskit.TableCollection(1.0)
        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
        s = tables.sites.add_row(0, "A")
        tables.mutations.add_row(site=s, derived_state="B", node=1)
        tables.mutations.add_row(site=s, derived_state="C", node=2)
        s = tables.sites.add_row(0.5, "")
        tables.mutations.add_row(site=s, derived_state="A long string", node=2)
        return tables.tree_sequence()

    def test_states(self):
        ts = self.get_missing_data_ts()
        v_iter = ts.variants(isolated_as_missing=False)
        v = next(v_iter)
        assert np.array_equal(v.states(), np.array(["A", "B", "C"]))
        v = next(v_iter)
        assert np.array_equal(v.states(), np.array(["", "", "A long string"]))
        # With no mssing data, it shouldn't matter if the missing string = an allele
        assert np.array_equal(
            v.states(missing_data_string=""), np.array(["", "", "A long string"])
        )

        v_iter = ts.variants(isolated_as_missing=True)
        v = next(v_iter)
        assert np.array_equal(v.states(), np.array(["N", "B", "C"]))
        v = next(v_iter)
        assert np.array_equal(v.states(), np.array(["N", "N", "A long string"]))
        assert np.array_equal(
            v.states(missing_data_string="MISSING"),
            np.array(["MISSING", "MISSING", "A long string"]),
        )

    @pytest.mark.parametrize("missing", [True, False])
    def test_states_haplotypes_equiv(self, missing):
        ts = msprime.sim_ancestry(2, sequence_length=20, random_seed=1)
        ts = msprime.sim_mutations(ts, rate=0.1, random_seed=1)
        assert ts.num_sites > 5
        tables = ts.dump_tables()
        tables.delete_intervals([[0, ts.site(4).position]])
        tables.sites.replace_with(ts.tables.sites)
        ts = tables.tree_sequence()
        states = np.array(
            [v.states() for v in ts.variants(isolated_as_missing=missing)]
        )
        for h1, h2 in zip(ts.haplotypes(isolated_as_missing=missing), states.T):
            assert h1 == "".join(h2)

    @pytest.mark.parametrize("s", ["", "A long string", True, np.nan, 0, -1])
    def test_bad_states(self, s):
        ts = self.get_missing_data_ts()
        v_iter = ts.variants(isolated_as_missing=True)
        v = next(v_iter)
        v = next(v_iter)
        match = "existing allele" if isinstance(s, str) else "not a string"
        with pytest.raises(ValueError, match=match):
            v.states(missing_data_string=s)


class TestLimitInterval:
    def test_simple_case(self, ts_fixture):
        ts = ts_fixture
        test_variant = tskit.Variant(ts)
        test_variant.decode(1)
        for v in ts.variants(left=ts.site(1).position, right=ts.site(2).position):
            # should only decode the first variant
            assert v.site.id == 1
            assert np.all(v.genotypes == test_variant.genotypes)
            assert v.alleles == test_variant.alleles

    @pytest.mark.parametrize(
        ["left", "expected"],
        [
            (None, [0, 1, 2, 3, 4]),
            (0, [0, 1, 2, 3, 4]),
            (0.999, [1, 2, 3, 4]),
            (1, [1, 2, 3, 4]),
            (3.999, [4]),
            (4, [4]),
            (4.00001, []),
            (4.99999, []),
            (np.array([4.99999])[0], []),
        ],
    )
    def test_left(self, left, expected):
        tables = tskit.TableCollection(5)
        for x in range(int(tables.sequence_length)):
            tables.sites.add_row(position=x, ancestral_state="A")
        ts = tables.tree_sequence()
        positions = [var.site.position for var in ts.variants(left=left)]
        assert positions == expected

    @pytest.mark.parametrize(
        ["right", "expected"],
        [
            (None, [0, 1, 2, 3, 4]),
            (5, [0, 1, 2, 3, 4]),
            (4.00001, [0, 1, 2, 3, 4]),
            (4.0, [0, 1, 2, 3]),
            (3.9999, [0, 1, 2, 3]),
            (0.00001, [0]),
            (np.array([1e-200])[0], [0]),
        ],
    )
    def test_right(self, right, expected):
        tables = tskit.TableCollection(5)
        for x in range(int(tables.sequence_length)):
            tables.sites.add_row(position=x, ancestral_state="A")
        ts = tables.tree_sequence()
        positions = [var.site.position for var in ts.variants(right=right)]
        assert positions == expected

    @pytest.mark.parametrize("bad_left", [-1, 10, 100, np.nan, np.inf, -np.inf])
    def test_bad_left(self, bad_left):
        ts = tskit.TableCollection(10).tree_sequence()
        with pytest.raises(ValueError, match="`left` not between"):
            list(ts.variants(left=bad_left))

    @pytest.mark.parametrize("bad_right", [-1, 0, 100, np.nan, np.inf, -np.inf])
    def test_bad_right(self, bad_right):
        ts = tskit.TableCollection(10).tree_sequence()
        with pytest.raises(ValueError, match="`right` not between"):
            list(ts.variants(right=bad_right))

    def test_bad_left_right(self):
        ts = tskit.TableCollection(10).tree_sequence()
        with pytest.raises(ValueError, match="must be less than"):
            list(ts.variants(left=1, right=1))


class TestHaplotypeGenerator:
    """
    Tests the haplotype generation code.
    """

    def verify_haplotypes(self, n, haplotypes):
        """
        Verify that the specified set of haplotypes is consistent.
        """
        assert len(haplotypes) == n
        m = len(haplotypes[0])
        for h in haplotypes:
            assert len(h) == m
        # Examine each column in H; we must have a mixture of 0s and 1s
        for k in range(m):
            zeros = 0
            ones = 0
            col = ""
            for j in range(n):
                b = haplotypes[j][k]
                zeros += b == "0"
                ones += b == "1"
                col += b
            assert zeros + ones == n

    def verify_tree_sequence(self, tree_sequence):
        n = tree_sequence.sample_size
        m = tree_sequence.num_sites
        haplotypes = list(tree_sequence.haplotypes())
        A = np.zeros((n, m), dtype="u1")
        B = np.zeros((n, m), dtype="u1")
        for j, h in enumerate(haplotypes):
            assert len(h) == m
            A[j] = np.frombuffer(h.encode("ascii"), np.uint8) - ord("0")
        for variant in tree_sequence.variants():
            B[:, variant.index] = variant.genotypes
        assert np.all(A == B)
        self.verify_haplotypes(n, haplotypes)

    def verify_simulation(self, n, m, r, theta):
        """
        Verifies a simulation for the specified parameters.
        """
        recomb_map = msprime.RecombinationMap.uniform_map(m, r, m)
        tree_sequence = msprime.simulate(
            n, recombination_map=recomb_map, mutation_rate=theta
        )
        self.verify_tree_sequence(tree_sequence)

    def test_random_parameters(self):
        num_random_sims = 10
        for _ in range(num_random_sims):
            n = random.randint(2, 50)
            m = random.randint(10, 200)
            r = random.random()
            theta = random.uniform(0, 2)
            self.verify_simulation(n, m, r, theta)

    def test_nonbinary_trees(self):
        bottlenecks = [
            msprime.SimpleBottleneck(0.01, 0, proportion=0.05),
            msprime.SimpleBottleneck(0.02, 0, proportion=0.25),
            msprime.SimpleBottleneck(0.03, 0, proportion=1),
        ]
        ts = msprime.simulate(
            10,
            length=100,
            recombination_rate=1,
            demographic_events=bottlenecks,
            random_seed=1,
        )
        self.verify_tree_sequence(ts)

    def test_acgt_mutations(self):
        ts = msprime.simulate(10, mutation_rate=10)
        assert ts.num_sites > 0
        tables = ts.tables
        sites = tables.sites
        mutations = tables.mutations
        sites.set_columns(
            position=sites.position,
            ancestral_state=np.zeros(ts.num_sites, dtype=np.int8) + ord("A"),
            ancestral_state_offset=np.arange(ts.num_sites + 1, dtype=np.uint32),
        )
        mutations.set_columns(
            site=mutations.site,
            node=mutations.node,
            derived_state=np.zeros(ts.num_sites, dtype=np.int8) + ord("T"),
            derived_state_offset=np.arange(ts.num_sites + 1, dtype=np.uint32),
        )
        tsp = tables.tree_sequence()
        H = [h.replace("0", "A").replace("1", "T") for h in ts.haplotypes()]
        assert H == list(tsp.haplotypes())

    def test_fails_multiletter_mutations(self):
        ts = msprime.simulate(10, random_seed=2)
        tables = ts.tables
        tables.sites.add_row(0, "ACTG")
        tsp = tables.tree_sequence()
        with pytest.raises(TypeError):
            list(tsp.haplotypes())

    def test_fails_deletion_mutations(self):
        ts = msprime.simulate(10, random_seed=2)
        tables = ts.tables
        tables.sites.add_row(0, "")
        tsp = tables.tree_sequence()
        with pytest.raises(TypeError):
            list(tsp.haplotypes())

    def test_nonascii_mutations(self):
        ts = msprime.simulate(10, random_seed=2)
        tables = ts.tables
        tables.sites.add_row(0, chr(169))  # Copyright symbol
        tsp = tables.tree_sequence()
        with pytest.raises(TypeError):
            list(tsp.haplotypes())

    def test_recurrent_mutations_over_samples(self):
        ts = msprime.simulate(10, random_seed=2)
        num_sites = 5
        tables = ts.dump_tables()
        for j in range(num_sites):
            tables.sites.add_row(
                position=j * ts.sequence_length / num_sites, ancestral_state="0"
            )
            for u in range(ts.sample_size):
                tables.mutations.add_row(site=j, node=u, derived_state="1")
        ts_new = tables.tree_sequence()
        ones = "1" * num_sites
        for h in ts_new.haplotypes():
            assert ones == h

    def test_silent_mutations(self):
        ts = msprime.simulate(10, random_seed=2)
        tables = ts.dump_tables()
        tree = next(ts.trees())
        for u in tree.children(tree.root):
            tables.sites.clear()
            tables.mutations.clear()
            site = tables.sites.add_row(position=0, ancestral_state="0")
            tables.mutations.add_row(site=site, node=u, derived_state="1")
            tables.mutations.add_row(site=site, node=tree.root, derived_state="1")
            ts_new = tables.tree_sequence()
            all(h == 1 for h in ts_new.haplotypes())

    def test_back_mutations(self):
        base_ts = msprime.simulate(10, random_seed=2)
        for j in [1, 2, 3]:
            ts = tsutil.insert_branch_mutations(base_ts, mutations_per_branch=j)
            self.verify_tree_sequence(ts)

    def test_missing_data(self):
        tables = tskit.TableCollection(1.0)
        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
        tables.sites.add_row(0.5, "A")
        ts = tables.tree_sequence()
        with pytest.raises(ValueError):
            list(ts.haplotypes(missing_data_character="A"))
        for c in ("-", ".", "a"):
            h = list(ts.haplotypes(missing_data_character=c))
            assert h == [c, c]
        h = list(ts.haplotypes(isolated_as_missing=True))
        assert h == ["N", "N"]
        h = list(ts.haplotypes(isolated_as_missing=False))
        assert h == ["A", "A"]
        h = list(ts.haplotypes())
        assert h == ["N", "N"]
        # Test deprecated method
        with pytest.warns(FutureWarning):
            h = list(ts.haplotypes(impute_missing_data=True))
        assert h == ["A", "A"]
        with pytest.warns(FutureWarning):
            h = list(ts.haplotypes(impute_missing_data=False))
        assert h == ["N", "N"]
        with pytest.warns(FutureWarning):
            h = list(ts.haplotypes(isolated_as_missing=True, impute_missing_data=True))
        assert h == ["N", "N"]
        with pytest.warns(FutureWarning):
            h = list(ts.haplotypes(isolated_as_missing=True, impute_missing_data=False))
        assert h == ["N", "N"]
        with pytest.warns(FutureWarning):
            h = list(ts.haplotypes(isolated_as_missing=False, impute_missing_data=True))
        assert h == ["A", "A"]
        with pytest.warns(FutureWarning):
            h = list(
                ts.haplotypes(isolated_as_missing=False, impute_missing_data=False)
            )
        assert h == ["A", "A"]

    def test_restrict_samples(self):
        tables = tskit.TableCollection(1.0)
        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
        tables.sites.add_row(0.5, "A")
        tables.mutations.add_row(0, 0, derived_state="B")
        ts = tables.tree_sequence()
        haplotypes = list(ts.haplotypes(samples=[0], isolated_as_missing=False))
        assert haplotypes == ["B"]
        haplotypes = list(ts.haplotypes(samples=[1], isolated_as_missing=False))
        assert haplotypes == ["A"]

    def test_restrict_positions(self):
        tables = tskit.TableCollection(1.0)
        tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
        tables.sites.add_row(0.1, "A")
        tables.sites.add_row(0.2, "B")
        tables.sites.add_row(0.3, "C")
        tables.sites.add_row(0.4, "D")
        ts = tables.tree_sequence()
        haplotypes = list(ts.haplotypes(left=0.2, right=0.4, isolated_as_missing=False))
        assert haplotypes == ["BC"]


class TestUserAlleles:
    """
    Tests the functionality of providing a user-specified allele mapping.
    """

    def test_simple_01(self):
        ts = msprime.simulate(10, mutation_rate=5, random_seed=2)
        assert ts.num_sites > 2
        G1 = ts.genotype_matrix()
        G2 = ts.genotype_matrix(alleles=("0", "1"))
        assert np.array_equal(G1, G2)
        for v1, v2 in itertools.zip_longest(
            ts.variants(), ts.variants(alleles=("0", "1"))
        ):
            assert v1.alleles == v2.alleles
            assert v1.site == v2.site
            assert np.array_equal(v1.genotypes, v2.genotypes)

    def test_simple_01_trailing_alleles(self):
        ts = msprime.simulate(10, mutation_rate=5, random_seed=2)
        assert ts.num_sites > 2
        G1 = ts.genotype_matrix()
        alleles = ("0", "1", "2", "xxxxx")
        G2 = ts.genotype_matrix(alleles=alleles)
        assert np.array_equal(G1, G2)
        for v1, v2 in itertools.zip_longest(
            ts.variants(), ts.variants(alleles=alleles)
        ):
            assert v2.alleles == alleles
            assert v1.site == v2.site
            assert np.array_equal(v1.genotypes, v2.genotypes)

    def test_simple_01_leading_alleles(self):
        ts = msprime.simulate(10, mutation_rate=5, random_seed=2)
        assert ts.num_sites > 2
        G1 = ts.genotype_matrix()
        alleles = ("A", "B", "C", "0", "1")
        G2 = ts.genotype_matrix(alleles=alleles)
        assert np.array_equal(G1 + 3, G2)
        for v1, v2 in itertools.zip_longest(
            ts.variants(), ts.variants(alleles=alleles)
        ):
            assert v2.alleles == alleles
            assert v1.site == v2.site
            assert np.array_equal(v1.genotypes + 3, v2.genotypes)

    def test_simple_01_duplicate_alleles(self):
        ts = msprime.simulate(10, mutation_rate=5, random_seed=2)
        assert ts.num_sites > 2
        G1 = ts.genotype_matrix()
        alleles = ("0", "0", "1")
        G2 = ts.genotype_matrix(alleles=alleles)
        index = np.where(G1 == 1)
        G1[index] = 2
        assert np.array_equal(G1, G2)
        for v1, v2 in itertools.zip_longest(
            ts.variants(), ts.variants(alleles=alleles)
        ):
            assert v2.alleles == alleles
            assert v1.site == v2.site
            g = np.array(v1.genotypes)
            index = np.where(g == 1)
            g[index] = 2
            assert np.array_equal(g, v2.genotypes)

    def test_simple_acgt(self):
        ts = msprime.simulate(10, random_seed=2)
        ts = msprime.mutate(
            ts, rate=4, random_seed=2, model=msprime.InfiniteSites(msprime.NUCLEOTIDES)
        )
        assert ts.num_sites > 2
        alleles = tskit.ALLELES_ACGT
        G = ts.genotype_matrix(alleles=alleles)
        for v1, v2 in itertools.zip_longest(
            ts.variants(), ts.variants(alleles=alleles)
        ):
            assert v2.alleles == alleles
            assert v1.site == v2.site
            h1 = "".join(v1.alleles[g] for g in v1.genotypes)
            h2 = "".join(v2.alleles[g] for g in v2.genotypes)
            assert h1 == h2
            assert np.array_equal(v2.genotypes, G[v1.site.id])

    def test_missing_alleles(self):
        ts = msprime.simulate(10, random_seed=2)
        ts = msprime.mutate(
            ts, rate=4, random_seed=2, model=msprime.InfiniteSites(msprime.NUCLEOTIDES)
        )
        assert ts.num_sites > 2
        bad_allele_examples = [
            tskit.ALLELES_01,
            tuple(["A"]),
            ("C", "T", "G"),
            ("AA", "C", "T", "G"),
            tuple(["ACTG"]),
        ]
        for bad_alleles in bad_allele_examples:
            with pytest.raises(exceptions.LibraryError):
                ts.genotype_matrix(alleles=bad_alleles)
            with pytest.raises(exceptions.LibraryError):
                list(ts.variants(alleles=bad_alleles))

    def test_too_many_alleles(self):
        ts = msprime.simulate(10, mutation_rate=5, random_seed=2)
        for n in range(128, 138):
            bad_alleles = tuple("0" for _ in range(n))
            with pytest.raises(exceptions.LibraryError):
                ts.genotype_matrix(alleles=bad_alleles)
            with pytest.raises(exceptions.LibraryError):
                list(ts.variants(alleles=bad_alleles))

    def test_zero_allele(self):
        ts = msprime.simulate(10, mutation_rate=5, random_seed=2)
        with pytest.raises(ValueError):
            ts.genotype_matrix(alleles=tuple())
        with pytest.raises(ValueError):
            list(ts.variants(alleles=tuple()))

    def test_missing_data(self):
        tables = tskit.TableCollection(1)
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        tables.sites.add_row(0.5, "0")
        tables.mutations.add_row(0, 0, "1")

        ts = tables.tree_sequence()
        for isolated_as_missing in [True, False]:
            G1 = ts.genotype_matrix(isolated_as_missing=isolated_as_missing)
            G2 = ts.genotype_matrix(
                isolated_as_missing=isolated_as_missing, alleles=tskit.ALLELES_01
            )
            assert np.array_equal(G1, G2)
            vars1 = ts.variants(isolated_as_missing=isolated_as_missing)
            vars2 = ts.variants(
                isolated_as_missing=isolated_as_missing, alleles=tskit.ALLELES_01
            )
            for v1, v2 in itertools.zip_longest(vars1, vars2):
                assert v2.alleles == v1.alleles
                assert v1.site == v2.site
                assert np.array_equal(v1.genotypes, v2.genotypes)


class TestUserAllelesRoundTrip:
    """
    Tests that we correctly produce haplotypes in a variety of situations for
    the user specified allele map encoding.
    """

    def verify(self, ts, alleles):
        for v1, v2 in itertools.zip_longest(
            ts.variants(), ts.variants(alleles=alleles)
        ):
            h1 = [v1.alleles[g] for g in v1.genotypes]
            h2 = [v2.alleles[g] for g in v2.genotypes]
            assert h1 == h2

    def test_simple_01(self):
        ts = msprime.simulate(5, mutation_rate=2, random_seed=3)
        assert ts.num_sites > 3
        valid_alleles = [
            tskit.ALLELES_01,
            ("0", "1", "xry"),
            ("xry", "0", "1", "xry"),
            tuple(str(j) for j in range(127)),
            tuple(["0" for j in range(126)] + ["1"]),
        ]
        for alleles in valid_alleles:
            self.verify(ts, alleles)

    def test_simple_acgt(self):
        ts = msprime.simulate(5, random_seed=3)
        ts = msprime.mutate(
            ts, rate=4, random_seed=3, model=msprime.InfiniteSites(msprime.NUCLEOTIDES)
        )
        assert ts.num_sites > 3
        valid_alleles = [
            tskit.ALLELES_ACGT,
            ("A", "C", "T", "G", "AAAAAAAAAAAAAA"),
            ("AA", "CC", "TT", "GG", "A", "C", "T", "G"),
        ]
        for alleles in valid_alleles:
            self.verify(ts, alleles)

    def test_jukes_cantor(self):
        ts = msprime.simulate(6, random_seed=1, mutation_rate=1)
        ts = tsutil.jukes_cantor(ts, 20, 1, seed=10)
        valid_alleles = [
            tskit.ALLELES_ACGT,
            ("A", "C", "T", "G", "AAAAAAAAAAAAAA"),
            ("AA", "CC", "TT", "GG", "A", "C", "T", "G"),
        ]
        for alleles in valid_alleles:
            self.verify(ts, alleles)

    def test_multichar_mutations(self):
        ts = msprime.simulate(6, random_seed=1, recombination_rate=2)
        ts = tsutil.insert_multichar_mutations(ts)
        assert ts.num_sites > 5
        all_alleles = set()
        for var in ts.variants():
            all_alleles.update(var.alleles)
        all_alleles = tuple(all_alleles)
        self.verify(ts, all_alleles)
        self.verify(ts, all_alleles[::-1])

    def test_simple_01_missing_data(self):
        ts = msprime.simulate(6, mutation_rate=2, random_seed=3)
        tables = ts.dump_tables()
        # Add another sample node. This will be missing data everywhere.
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        ts = tables.tree_sequence()
        assert ts.num_sites > 3
        valid_alleles = [
            tskit.ALLELES_01,
            ("0", "1", "xry"),
            ("xry", "0", "1", "xry"),
            tuple(str(j) for j in range(127)),
            tuple(["0" for j in range(126)] + ["1"]),
        ]
        for alleles in valid_alleles:
            self.verify(ts, alleles)


class TestBinaryTreeExample:
    # 2.00┊   4   ┊
    #     ┊ ┏━┻┓  ┊
    # 1.00┊ ┃  3  ┊
    #     ┊ ┃ ┏┻┓ ┊
    # 0.00┊ 0 1 2 ┊
    #     0      10
    #      |    |
    #  pos 2    9
    #  anc A    T
    @tests.cached_example
    def ts(self):
        ts = tskit.Tree.generate_balanced(3, span=10).tree_sequence
        tables = ts.dump_tables()
        tables.sites.add_row(2, ancestral_state="A")
        tables.sites.add_row(9, ancestral_state="T")
        tables.mutations.add_row(site=0, node=0, derived_state="G")
        tables.mutations.add_row(site=1, node=3, derived_state="C")
        return tables.tree_sequence()

    def test_haplotypes(self):
        H = list(self.ts().haplotypes())
        assert H[0] == "GT"
        assert H[1] == "AC"
        assert H[2] == "AC"

    def test_haplotypes_empty_interval(self):
        ts = self.ts()
        H = list(ts.haplotypes(left=4, right=5))
        assert H == ["", "", ""]

    def test_genotypes(self):
        G = self.ts().genotype_matrix()
        Gp = [[1, 0, 0], [0, 1, 1]]
        np.testing.assert_array_equal(G, Gp)

    def test_alignments_default(self):
        A = list(self.ts().alignments())
        assert A[0] == "NNGNNNNNNT"
        assert A[1] == "NNANNNNNNC"
        assert A[2] == "NNANNNNNNC"

    def test_alignments_restricted(self):
        ts = self.ts()
        samples = ts.samples()
        # Take the first 2 in reverse order
        A = list(ts.alignments(left=1, right=9, samples=samples[1::-1]))
        assert A[0] == "NANNNNNN"
        assert A[1] == "NGNNNNNN"

    def test_empty_samples(self):
        ts = self.ts()
        A = list(ts.alignments(samples=[]))
        assert len(A) == 0

    def test_non_sample_samples(self):
        ts = self.ts()
        with pytest.raises(tskit.LibraryError, match="MUST_IMPUTE_NON_SAMPLES"):
            list(ts.alignments(samples=[4]))

    def test_alignments_missing_data_char(self):
        A = list(self.ts().alignments(missing_data_character="x"))
        assert A[0] == "xxGxxxxxxT"
        assert A[1] == "xxAxxxxxxC"
        assert A[2] == "xxAxxxxxxC"

    def test_alignments_reference_sequence(self):
        ref = "0123456789"
        A = list(self.ts().alignments(reference_sequence=ref))
        assert A[0] == "01G345678T"
        assert A[1] == "01A345678C"
        assert A[2] == "01A345678C"

    def test_alignments_reference_sequence_embedded_null(self):
        # This is a total corner case, but just want to make sure
        # we do something sensible.
        ref = "0123" + "\0" + "56789"
        A = list(self.ts().alignments(reference_sequence=ref))
        assert A[0] == "01G3\x005678T"
        assert A[1] == "01A3\x005678C"
        assert A[2] == "01A3\x005678C"

    def test_fasta_default(self):
        expected = textwrap.dedent(
            """\
            >n0
            NNGNNNNNNT
            >n1
            NNANNNNNNC
            >n2
            NNANNNNNNC
            """
        )
        assert expected == self.ts().as_fasta()

    def test_fasta_missing_Q(self):
        expected = textwrap.dedent(
            """\
            >n0
            QQGQQQQQQT
            >n1
            QQAQQQQQQC
            >n2
            QQAQQQQQQC
            """
        )
        assert expected == self.ts().as_fasta(missing_data_character="Q")

    def test_fasta_reference_sequence(self):
        ref = "0123456789"
        expected = textwrap.dedent(
            """\
            >n0
            01G345678T
            >n1
            01A345678C
            >n2
            01A345678C
            """
        )
        assert expected == self.ts().as_fasta(reference_sequence=ref)

    def test_nexus_default(self):
        expected = textwrap.dedent(
            """\
            #NEXUS
            BEGIN TAXA;
              DIMENSIONS NTAX=3;
              TAXLABELS n0 n1 n2;
            END;
            BEGIN DATA;
              DIMENSIONS NCHAR=10;
              FORMAT DATATYPE=DNA MISSING=?;
              MATRIX
                n0 ??G??????T
                n1 ??A??????C
                n2 ??A??????C
              ;
            END;
            BEGIN TREES;
              TREE t0^10 = [&R] (n0:2,(n1:1,n2:1):1);
            END;
            """
        )
        assert expected == self.ts().as_nexus()

    def test_nexus_missing_N(self):
        expected = textwrap.dedent(
            """\
            #NEXUS
            BEGIN TAXA;
              DIMENSIONS NTAX=3;
              TAXLABELS n0 n1 n2;
            END;
            BEGIN DATA;
              DIMENSIONS NCHAR=10;
              FORMAT DATATYPE=DNA MISSING=N;
              MATRIX
                n0 NNGNNNNNNT
                n1 NNANNNNNNC
                n2 NNANNNNNNC
              ;
            END;
            BEGIN TREES;
              TREE t0^10 = [&R] (n0:2,(n1:1,n2:1):1);
            END;
            """
        )
        assert expected == self.ts().as_nexus(missing_data_character="N")

    def test_nexus_reference_sequence(self):
        ref = "0123456789"
        expected = textwrap.dedent(
            """\
            #NEXUS
            BEGIN TAXA;
              DIMENSIONS NTAX=3;
              TAXLABELS n0 n1 n2;
            END;
            BEGIN DATA;
              DIMENSIONS NCHAR=10;
              FORMAT DATATYPE=DNA MISSING=?;
              MATRIX
                n0 01G345678T
                n1 01A345678C
                n2 01A345678C
              ;
            END;
            BEGIN TREES;
              TREE t0^10 = [&R] (n0:2,(n1:1,n2:1):1);
            END;
            """
        )
        assert expected == self.ts().as_nexus(reference_sequence=ref)


class TestBinaryTreeWithReferenceExample:
    # 2.00┊   4   ┊
    #     ┊ ┏━┻┓  ┊
    # 1.00┊ ┃  3  ┊
    #     ┊ ┃ ┏┻┓ ┊
    # 0.00┊ 0 1 2 ┊
    #     0      10
    #      |    |
    #  pos 2    9
    #  anc A    T
    @tests.cached_example
    def ts(self):
        ts = tskit.Tree.generate_balanced(3, span=10).tree_sequence
        tables = ts.dump_tables()
        tables.sites.add_row(2, ancestral_state="A")
        tables.sites.add_row(9, ancestral_state="T")
        tables.mutations.add_row(site=0, node=0, derived_state="G")
        tables.mutations.add_row(site=1, node=3, derived_state="C")
        tables.reference_sequence.data = "ACGTACGTAC"
        return tables.tree_sequence()

    def test_alignments_default(self):
        A = list(self.ts().alignments())
        assert A[0] == "ACGTACGTAT"
        assert A[1] == "ACATACGTAC"
        assert A[2] == "ACATACGTAC"

    def test_alignments_missing_data_char(self):
        A = list(self.ts().alignments(missing_data_character="x"))
        assert A[0] == "ACGTACGTAT"
        assert A[1] == "ACATACGTAC"
        assert A[2] == "ACATACGTAC"

    def test_alignments_reference_sequence(self):
        ref = "0123456789"
        A = list(self.ts().alignments(reference_sequence=ref))
        assert A[0] == "01G345678T"
        assert A[1] == "01A345678C"
        assert A[2] == "01A345678C"

    def test_fasta_default(self):
        expected = textwrap.dedent(
            """\
            >n0
            ACGTACGTAT
            >n1
            ACATACGTAC
            >n2
            ACATACGTAC
            """
        )
        assert expected == self.ts().as_fasta()

    def test_fasta_reference_sequence(self):
        ref = "0123456789"
        expected = textwrap.dedent(
            """\
            >n0
            01G345678T
            >n1
            01A345678C
            >n2
            01A345678C
            """
        )
        assert expected == self.ts().as_fasta(reference_sequence=ref)

    def test_nexus_default(self):
        expected = textwrap.dedent(
            """\
            #NEXUS
            BEGIN TAXA;
              DIMENSIONS NTAX=3;
              TAXLABELS n0 n1 n2;
            END;
            BEGIN DATA;
              DIMENSIONS NCHAR=10;
              FORMAT DATATYPE=DNA MISSING=?;
              MATRIX
                n0 ACGTACGTAT
                n1 ACATACGTAC
                n2 ACATACGTAC
              ;
            END;
            BEGIN TREES;
              TREE t0^10 = [&R] (n0:2,(n1:1,n2:1):1);
            END;
            """
        )
        assert expected == self.ts().as_nexus()

    def test_nexus_reference_sequence(self):
        ref = "0123456789"
        expected = textwrap.dedent(
            """\
            #NEXUS
            BEGIN TAXA;
              DIMENSIONS NTAX=3;
              TAXLABELS n0 n1 n2;
            END;
            BEGIN DATA;
              DIMENSIONS NCHAR=10;
              FORMAT DATATYPE=DNA MISSING=?;
              MATRIX
                n0 01G345678T
                n1 01A345678C
                n2 01A345678C
              ;
            END;
            BEGIN TREES;
              TREE t0^10 = [&R] (n0:2,(n1:1,n2:1):1);
            END;
            """
        )
        assert expected == self.ts().as_nexus(reference_sequence=ref)


class TestMissingDataExample:
    # 2.00┊   4     ┊
    #     ┊ ┏━┻┓    ┊
    # 1.00┊ ┃  3    ┊
    #     ┊ ┃ ┏┻┓   ┊
    # 0.00┊ 0 1 2 5 ┊
    #     0        10
    #      |      |
    #  pos 2      9
    #  anc A      T
    @tests.cached_example
    def ts(self):
        ts = tskit.Tree.generate_balanced(3, span=10).tree_sequence
        tables = ts.dump_tables()
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        tables.sites.add_row(2, ancestral_state="A")
        tables.sites.add_row(9, ancestral_state="T")
        tables.mutations.add_row(site=0, node=0, derived_state="G")
        tables.mutations.add_row(site=1, node=3, derived_state="C")
        return tables.tree_sequence()

    def test_haplotypes(self):
        H = list(self.ts().haplotypes())
        assert H[0] == "GT"
        assert H[1] == "AC"
        assert H[2] == "AC"
        assert H[3] == "NN"

    def test_haplotypes_missing_data_char(self):
        H = list(self.ts().haplotypes(missing_data_character="?"))
        assert H[0] == "GT"
        assert H[1] == "AC"
        assert H[2] == "AC"
        assert H[3] == "??"

    def test_genotypes(self):
        G = self.ts().genotype_matrix()
        Gp = [[1, 0, 0, -1], [0, 1, 1, -1]]
        np.testing.assert_array_equal(G, Gp)

    @pytest.mark.skip("Missing data in alignments: #1896")
    def test_alignments_default(self):
        A = list(self.ts().alignments())
        assert A[0] == "NNGNNNNNNT"
        assert A[1] == "NNANNNNNNC"
        assert A[2] == "NNANNNNNNC"
        assert A[3] == "NNNNNNNNNN"

    def test_alignments_fails(self):
        # https://github.com/tskit-dev/tskit/issues/1896
        with pytest.raises(ValueError, match="1896"):
            next(self.ts().alignments())

    @pytest.mark.skip("Missing data in alignments: #1896")
    def test_alignments_impute_missing(self):
        ref = "N" * 10
        A = list(
            self.ts().alignments(reference_sequence=ref, isolated_as_missing=False)
        )
        assert A[0] == "NNGNNNNNNT"
        assert A[1] == "NNANNNNNNC"
        assert A[2] == "NNANNNNNNC"
        assert A[3] == "NNANNNNNNT"

    @pytest.mark.skip("Missing data in alignments: #1896")
    def test_alignments_missing_char(self):
        A = list(self.ts().alignments(missing_data_character="z"))
        assert A[0] == "zzGzzzzzzT"
        assert A[1] == "zzAzzzzzzC"
        assert A[2] == "zzAzzzzzzC"
        assert A[3] == "zzzzzzzzzz"

    @pytest.mark.skip("Missing data in alignments: #1896")
    def test_alignments_missing_char_ref(self):
        A = list(self.ts().alignments(missing_data_character="z"))
        assert A[0] == "NNGNNNNNNT"
        assert A[1] == "NNANNNNNNC"
        assert A[2] == "NNANNNNNNC"
        assert A[3] == "zzzzzzzzzz"

    @pytest.mark.skip("Missing data in alignments: #1896")
    def test_alignments_reference_sequence(self):
        ref = "0123456789"
        A = list(self.ts().alignments(reference_sequence=ref))
        assert A[0] == "01G345678T"
        assert A[1] == "01A345678C"
        assert A[2] == "01A345678C"
        assert A[3] == "NNNNNNNNNN"

    @pytest.mark.skip("Missing data in alignments: #1896")
    def test_alignments_reference_sequence_missing_data_char(self):
        ref = "0123456789"
        A = list(
            self.ts().alignments(reference_sequence=ref, missing_data_character="Q")
        )
        assert A[0] == "01G345678T"
        assert A[1] == "01A345678C"
        assert A[2] == "01A345678C"
        assert A[3] == "QQQQQQQQQQ"

    @pytest.mark.skip("Missing data in alignments: #1896")
    def test_fasta_reference_sequence(self):
        ref = "0123456789"
        expected = textwrap.dedent(
            """\
            >n0
            01G345678T
            >n1
            01A345678C
            >n2
            01A345678C
            >n5
            NNNNNNNNNN
            """
        )
        assert expected == self.ts().as_fasta(reference_sequence=ref)

    @pytest.mark.skip("Missing data in alignments: #1896")
    def test_fasta_reference_sequence_missing_data_char(self):
        ref = "0123456789"
        expected = textwrap.dedent(
            """\
            >n0
            01G345678T
            >n1
            01A345678C
            >n2
            01A345678C
            >n5
            QQQQQQQQQQ
            """
        )
        assert expected == self.ts().as_fasta(
            reference_sequence=ref, missing_data_character="Q"
        )

    @pytest.mark.skip("Missing data in alignments: #1896")
    def test_fasta_impute_missing(self):
        ref = "N" * 10
        expected = textwrap.dedent(
            """\
            >n0
            NNGNNNNNNT
            >n1
            NNANNNNNNC
            >n2
            NNANNNNNNC
            >n5
            NNANNNNNNT
            """
        )
        assert expected == self.ts().as_fasta(
            reference_sequence=ref, isolated_as_missing=False
        )

    # Note: the nexus tree output isn't compatible with our representation of
    # missing data as trees with isolated roots (newick parsers won't accept
    # this as valid input), so we set include_trees=False for these examples.
    @pytest.mark.skip("Missing data in alignments: #1896")
    def test_nexus_reference_sequence(self):
        ref = "0123456789"
        expected = textwrap.dedent(
            """\
            #NEXUS
            BEGIN TAXA;
              DIMENSIONS NTAX=4;
              TAXLABELS n0 n1 n2 n5;
            END;
            BEGIN DATA;
              DIMENSIONS NCHAR=10;
              FORMAT DATATYPE=DNA MISSING=?;
              MATRIX
                n0 01G345678T
                n1 01A345678C
                n2 01A345678C
                n5 ??????????
              ;
            END;
            """
        )
        assert expected == self.ts().as_nexus(
            reference_sequence=ref, include_trees=False
        )

    @pytest.mark.skip("Missing data in alignments: #1896")
    def test_nexus_reference_sequence_missing_data_char(self):
        ref = "0123456789"
        expected = textwrap.dedent(
            """\
            #NEXUS
            BEGIN TAXA;
              DIMENSIONS NTAX=4;
              TAXLABELS n0 n1 n2 n5;
            END;
            BEGIN DATA;
              DIMENSIONS NCHAR=10;
              FORMAT DATATYPE=DNA MISSING=Q;
              MATRIX
                n0 01G345678T
                n1 01A345678C
                n2 01A345678C
                n5 QQQQQQQQQQ
              ;
            END;
            """
        )
        assert expected == self.ts().as_nexus(
            reference_sequence=ref,
            missing_data_character="Q",
            include_trees=False,
        )

    @pytest.mark.skip("Missing data in alignments: #1896")
    def test_nexus_impute_missing(self):
        ref = "0123456789"
        expected = textwrap.dedent(
            """\
            #NEXUS
            BEGIN TAXA;
              DIMENSIONS NTAX=4;
              TAXLABELS n0 n1 n2 n5;
            END;
            BEGIN DATA;
              DIMENSIONS NCHAR=10;
              FORMAT DATATYPE=DNA MISSING=?;
              MATRIX
                n0 01G345678T
                n1 01A345678C
                n2 01A345678C
                n5 01A345678T
              ;
            END;
            """
        )
        assert expected == self.ts().as_nexus(
            reference_sequence=ref,
            isolated_as_missing=False,
            include_trees=False,
        )


class TestMultiRootExample:
    # 1.00┊  4   5  ┊
    #     ┊ ┏┻┓ ┏┻┓ ┊
    # 0.00┊ 0 1 2 3 ┊
    #     0        10
    #       |     |
    #  pos  2     8
    #  anc  G     C
    @tests.cached_example
    def ts(self):
        tree = tskit.Tree.generate_balanced(4, arity=2, span=10)
        tables = tree.tree_sequence.dump_tables()
        edges = tables.edges.copy()
        tables.edges.clear()
        for edge in edges:
            if edge.parent != 6:
                tables.edges.append(edge)
        tables.sites.add_row(2, ancestral_state="G")
        tables.sites.add_row(8, ancestral_state="C")
        tables.mutations.add_row(site=0, node=0, derived_state="T")
        tables.mutations.add_row(site=1, node=5, derived_state="A")
        return tables.tree_sequence()

    def test_haplotypes(self):
        H = list(self.ts().haplotypes())
        assert H[0] == "TC"
        assert H[1] == "GC"
        assert H[2] == "GA"
        assert H[3] == "GA"

    def test_genotypes(self):
        G = self.ts().genotype_matrix()
        Gp = [[1, 0, 0, 0], [0, 0, 1, 1]]
        np.testing.assert_array_equal(G, Gp)

    def test_alignments_default(self):
        A = list(self.ts().alignments())
        assert A[0] == "NNTNNNNNCN"
        assert A[1] == "NNGNNNNNCN"
        assert A[2] == "NNGNNNNNAN"
        assert A[3] == "NNGNNNNNAN"

    def test_alignments_N_ref(self):
        A = list(self.ts().alignments(reference_sequence="N" * 10))
        assert A[0] == "NNTNNNNNCN"
        assert A[1] == "NNGNNNNNCN"
        assert A[2] == "NNGNNNNNAN"
        assert A[3] == "NNGNNNNNAN"

    def test_fasta_reference_sequence(self):
        ref = "0123456789"
        expected = textwrap.dedent(
            """\
            >n0
            01T34567C9
            >n1
            01G34567C9
            >n2
            01G34567A9
            >n3
            01G34567A9
            """
        )
        assert expected == self.ts().as_fasta(reference_sequence=ref)


class TestAlignmentsErrors:
    @tests.cached_example
    def simplest_ts(self):
        tables = tskit.TableCollection(1)
        tables.nodes.add_row(flags=1, time=0)
        return tables.tree_sequence()

    def test_non_discrete_genome(self):
        ts = tskit.TableCollection(1.1).tree_sequence()
        assert not ts.discrete_genome
        with pytest.raises(ValueError, match="defined for discrete genomes"):
            list(ts.alignments())

    @pytest.mark.parametrize("ref_length", [1, 9, 11])
    def test_reference_length_mismatch(self, ref_length):
        tables = tskit.TableCollection(10)
        tables.reference_sequence.data = "A" * ref_length
        ts = tables.tree_sequence()
        if ref_length <= tables.sequence_length:
            with pytest.raises(ValueError, match="shorter than"):
                list(ts.alignments())
        else:
            # Longer reference sequences are allowed
            list(ts.alignments())

    @pytest.mark.parametrize("ref", ["", "xy"])
    def test_reference_sequence_length_mismatch(self, ref):
        ts = self.simplest_ts()
        with pytest.raises(ValueError, match="shorter than"):
            list(ts.alignments(reference_sequence=ref))

    @pytest.mark.parametrize("ref", ["À", "┃", "α"])
    def test_non_ascii_references(self, ref):
        ts = self.simplest_ts()
        with pytest.raises(UnicodeEncodeError):
            list(ts.alignments(reference_sequence=ref))

    @pytest.mark.parametrize("ref", ["À", "┃", "α"])
    def test_non_ascii_embedded_references(self, ref):
        tables = tskit.TableCollection(1)
        tables.nodes.add_row(flags=1, time=0)
        tables.reference_sequence.data = ref
        ts = tables.tree_sequence()
        with pytest.raises(UnicodeEncodeError):
            list(ts.alignments())

    @pytest.mark.parametrize("missing_data_char", ["À", "┃", "α"])
    def test_non_ascii_missing_data_char(self, missing_data_char):
        ts = self.simplest_ts()
        with pytest.raises(UnicodeEncodeError):
            list(ts.alignments(missing_data_character=missing_data_char))

    def test_bad_left(self):
        ts = tskit.TableCollection(10).tree_sequence()
        with pytest.raises(ValueError, match="integer"):
            list(ts.alignments(left=0.1))

    def test_bad_right(self):
        ts = tskit.TableCollection(10).tree_sequence()
        with pytest.raises(ValueError, match="integer"):
            list(ts.alignments(right=1.1))

    def test_bad_restricted(self):
        tables = tskit.TableCollection(10)
        tables.reference_sequence.data = "A" * 7
        ts = tables.tree_sequence()
        with pytest.raises(ValueError, match="sequence ends before"):
            list(ts.alignments(right=8))


class TestAlignmentExamples:
    @pytest.mark.parametrize("ts", get_example_discrete_genome_tree_sequences())
    def test_defaults(self, ts):
        has_missing_data = np.any(ts.genotype_matrix() == -1)
        if has_missing_data:
            with pytest.raises(ValueError, match="1896"):
                list(ts.alignments())
        else:
            A = list(ts.alignments())
            assert len(A) == ts.num_samples
            H = list(ts.haplotypes())
            pos = ts.tables.sites.position.astype(int)
            for a, h in map(np.array, zip(A, H)):
                last = 0
                for j, x in enumerate(pos):
                    assert a[last:x] == "N" * (x - last)
                    assert a[x] == h[j]
                    last = x + 1

    @pytest.mark.parametrize("ts", get_example_discrete_genome_tree_sequences())
    def test_reference_sequence(self, ts):
        ref = tskit.random_nucleotides(ts.sequence_length, seed=1234)
        has_missing_data = np.any(ts.genotype_matrix() == -1)
        if has_missing_data:
            with pytest.raises(ValueError, match="1896"):
                list(ts.alignments(reference_sequence=ref))
        else:
            A = list(ts.alignments(reference_sequence=ref))
            assert len(A) == ts.num_samples
            H = list(ts.haplotypes())
            pos = ts.tables.sites.position.astype(int)
            for a, h in map(np.array, zip(A, H)):
                last = 0
                for j, x in enumerate(pos):
                    assert a[last:x] == ref[last:x]
                    assert a[x] == h[j]
                    last = x + 1
                assert a[last:] == ref[last:]


#
# Tests for allele_remap
#
@pytest.mark.parametrize(
    "alleles_from, alleles_to, allele_map",
    [
        # Case 1: alleles_to is longer than alleles_from.
        (
            ["A", "C", "G", "T"],
            ["G", "C"],
            np.array([2, 1, 0, 3], dtype="uint32"),
        ),
        # Case 2: alleles_to is shorter than alleles_from.
        (
            ["G", "C"],
            ["A", "C", "G", "T"],
            np.array([2, 1], dtype="uint32"),
        ),
        # Case 3: alleles_to is empty.
        (
            ["A", "C", "G", "T"],
            [],
            np.array([0, 1, 2, 3], dtype="uint32"),
        ),
        # Case 4: alleles_from is empty.
        (
            [],
            ["A", "C", "G", "T"],
            np.array([], dtype="uint32"),
        ),
        # Case 5: Both lists are empty.
        (
            [],
            [],
            np.array([], dtype="uint32"),
        ),
        # Case 6: Both lists are tuples.
        (
            ("G", "C"),
            ("A", "C", "G", "T"),
            np.array([2, 1], dtype="uint32"),
        ),
        # Case 7: Both lists are numpy arrays.
        (
            np.array(("G", "C")),
            np.array(("A", "C", "G", "T")),
            np.array([2, 1], dtype="uint32"),
        ),
        # Case 8: Lists are of two different types.
        (
            np.array(("G", "C")),
            ["A", "C", "G", "T"],
            np.array([2, 1], dtype="uint32"),
        ),
        # Case 9: Lists contain elements of arbitrary types.
        (
            ["ABBA", "CDCD"],
            ["ABBA", "CDCD", "EFEF", "GG", 18],
            np.array([0, 1], dtype="uint32"),
        ),
        # Case 10: Lists contain unicode characters.
        (
            ["\u1F1E8", "\u1F1EC"],
            ["\u1F1EC", "\u1F1E8", "\u1F1E6", "\u1F1F3"],
            np.array([1, 0], dtype="uint32"),
        ),
    ],
)
def test_allele_remap(alleles_from, alleles_to, allele_map):
    assert np.array_equal(allele_map, allele_remap(alleles_from, alleles_to))


class TestVariant:
    # Much more in-depth testing of variant decoding is done via the ts.variants
    # method as it existed before this class was publicly creatable.
    def test_variant_init(self, ts_fixture):
        v = tskit.Variant(ts_fixture)
        assert np.array_equal(v.samples, np.array(ts_fixture.samples()))
        assert v.alleles == ()
        assert v.num_alleles == 0
        assert v.isolated_as_missing
        v = tskit.Variant(ts_fixture, samples=[43, 1])
        assert np.array_equal(v.samples, np.array([43, 1]))
        v = tskit.Variant(ts_fixture, alleles=("A", "💩"))
        assert v.alleles == ("A", "💩")
        v = tskit.Variant(ts_fixture, isolated_as_missing=False)
        assert not v.isolated_as_missing

    def test_not_decoded(self, ts_fixture):
        variant = tskit.Variant(ts_fixture)
        assert variant.index == tskit.NULL
        with pytest.raises(ValueError, match="not yet been decoded"):
            variant.site
        assert variant.alleles == ()
        with pytest.raises(ValueError, match="not yet been decoded"):
            assert variant.genotypes
        assert not variant.has_missing_data
        assert variant.num_alleles == 0
        with pytest.raises(ValueError, match="not yet been decoded"):
            variant.position
        assert np.array_equal(variant.samples, np.array(ts_fixture.samples()))

    def test_variant_decode(self, ts_fixture):
        v = tskit.Variant(ts_fixture)
        v.decode(2)
        assert v.index == 2
        assert np.array_equal(v.samples, np.array(ts_fixture.samples()))
        assert v.alleles == ("A", "T", "G", "C", None)
        # No need to check contents as done in other tests
        assert len(v.genotypes) == ts_fixture.num_samples

    def test_variant_num_missing(self, ts_fixture):
        variant = next(ts_fixture.variants())
        assert variant.num_missing > 0
        assert variant.num_missing == np.sum(variant.genotypes == -1)

    def test_variant_counts(self, ts_fixture):
        variant = next(ts_fixture.variants())
        assert len(variant.alleles) > 2
        assert None in variant.alleles
        counts = variant.counts()
        assert len(counts) == len(variant.alleles)
        assert np.sum(list(counts.values())) == ts_fixture.num_samples
        assert counts[None] == variant.num_missing
        assert ts_fixture.num_samples > variant.num_missing
        for i, v in enumerate(variant.alleles):
            if v is not None:
                assert np.sum(variant.genotypes == i) == counts[v]

    def test_variant_counts_empty(self):
        tables = tskit.TableCollection(sequence_length=1)
        tables.sites.add_row(0, ancestral_state="💩")
        ts = tables.tree_sequence()
        variant = next(ts.variants())
        assert len(variant.counts()) == 1
        assert variant.counts()["💩"] == 0

    def test_variant_simple_frequencies(self):
        simple_tree = tskit.Tree.generate_balanced(4)
        simple_ts = simple_tree.tree_sequence
        tables = simple_ts.dump_tables()
        tables.sites.add_row(position=0.3, ancestral_state="AS0")
        tables.sites.add_row(position=0.6, ancestral_state="AS1")
        tables.mutations.add_row(site=0, derived_state="DS0_0", node=0)
        tables.mutations.add_row(site=0, derived_state="DS0_3", node=3)
        tables.mutations.add_row(
            site=1, derived_state="DS1", node=simple_tree.parent(0)
        )
        ts = tables.tree_sequence()
        variant_0 = next(ts.variants())
        freqs = variant_0.frequencies()
        assert len(freqs) == 3
        assert np.allclose(freqs["AS0"], 0.5)
        assert np.allclose(freqs["DS0_0"], 0.25)
        assert np.allclose(freqs["DS0_3"], 0.25)
        variant_1 = next(ts.variants(left=0.5))
        freqs = variant_1.frequencies()
        assert len(freqs) == 2
        assert np.allclose(freqs["AS1"], 0.5)
        assert np.allclose(freqs["DS1"], 0.5)

    def test_variant_frequencies(self, ts_fixture):
        variant = next(ts_fixture.variants())
        assert variant.num_missing > 0
        freqs = variant.frequencies()
        assert len(freqs) == len(variant.alleles)
        assert np.allclose(np.sum(list(freqs.values())), 1)
        for i, v in enumerate(variant.alleles):
            if v is None:
                f = np.sum(variant.genotypes == tskit.NULL) / ts_fixture.num_samples
            else:
                f = np.sum(variant.genotypes == i) / ts_fixture.num_samples
            assert np.allclose(f, freqs[v])

        freqs = variant.frequencies(remove_missing=True)
        assert len(freqs) == len(variant.alleles) - 1
        for i, v in enumerate(variant.alleles[:-1]):
            f = np.sum(variant.genotypes == i) / (
                ts_fixture.num_samples - variant.num_missing
            )
            assert np.allclose(f, freqs[v])

    def test_variant_frequencies_limit_samples(self, ts_fixture):
        assert ts_fixture.num_samples > 1
        variant = next(ts_fixture.variants(samples=ts_fixture.samples()[0:1]))
        assert len(variant.genotypes) == 1
        allele = variant.alleles[variant.genotypes[0]]
        freqs = variant.frequencies()
        assert freqs[allele] == 1
        # should be one freq of 1 and all the rest zero
        assert list(freqs.values()).count(0) == len(freqs) - 1

    def test_variant_nonsample_freqs(self):
        simple_tree = tskit.Tree.generate_balanced(4)
        nonsample_node_left = simple_tree.parent(0)
        nonsample_node_right = simple_tree.parent(3)
        assert nonsample_node_left != nonsample_node_right
        simple_ts = simple_tree.tree_sequence
        tables = simple_ts.dump_tables()
        tables.sites.add_row(position=0, ancestral_state="As")
        tables.mutations.add_row(site=0, derived_state="Ds", node=nonsample_node_left)
        ts = tables.tree_sequence()
        samples = [nonsample_node_left, nonsample_node_right]
        samples += list(simple_tree.children(nonsample_node_right))
        variant = next(ts.variants(samples=samples, isolated_as_missing=False))
        freqs = variant.frequencies()
        assert np.allclose(freqs["Ds"], 0.25)  # Just nonsample_node_left has the Ds
        assert np.allclose(freqs["As"], 0.75)

    def test_variant_frequencies_no_samples(self, ts_fixture, caplog):
        tables = ts_fixture.dump_tables()
        tables.nodes.flags = np.zeros_like(tables.nodes.flags)
        ts = tables.tree_sequence()
        variant = next(ts.variants())
        assert ts.num_samples == 0
        with caplog.at_level(logging.WARNING):
            freqs = variant.frequencies()
            assert caplog.text.count("frequencies undefined") == 1
        assert np.all(np.isnan(list(freqs.values())))

    def test_variant_str(self):
        """
        Test using a simple dummy tree sequence for testing.
        It has only one tree and one site, whose variant has the alleles
        ('A', 'T', 'G', '💩', '', 'TAG', None).
        """
        tables = tskit.TableCollection(10)
        for _ in np.arange(6):
            tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        tables.sites.add_row(position=5, ancestral_state="A")
        tables.mutations.add_row(site=0, node=0, derived_state="T")
        tables.mutations.add_row(site=0, node=1, derived_state="G")
        tables.mutations.add_row(site=0, node=2, derived_state="💩")
        tables.mutations.add_row(site=0, node=3, derived_state="")
        tables.mutations.add_row(site=0, node=4, derived_state="TAG")
        ts = tables.tree_sequence()
        v = next(ts.variants())
        assert v.alleles == ("A", "T", "G", "💩", "", "TAG", None)
        assert isinstance(str(v), str)
        assert re.match(
            textwrap.dedent(
                r"""
                ╔═+╗
                ║Variant\s*║
                ╠═+╤═+╣
                ║Site id\s*│\s*0║
                ╟─+┼─+╢
                ║Site position\s*│\s*[0-9\.]+║
                ╟─+┼─+╢
                ║Number of samples\s*│\s*[0-9]+║
                ╟─+┼─+╢
                ║Number of alleles\s*│\s*[0-9]+║
                ╟─+┼─+╢
                ║Samples with allele \'A\'\s*│\s*[0-9]+\s*\([0-9\.]+\%\)║
                ╟─+┼─+╢
                ║Samples with allele \'T\'\s*│\s*[0-9]+\s*\([0-9\.]+\%\)║
                ╟─+┼─+╢
                ║Samples with allele \'G\'\s*│\s*[0-9]+\s*\([0-9\.]+\%\)║
                ╟─+┼─+╢
                ║Samples with allele \'💩\'\s*│\s*[0-9]+\s*\([0-9\.]+\%\)║
                ╟─+┼─+╢
                ║Samples with allele \'\'\s*│\s*[0-9]+\s*\([0-9\.]+\%\)║
                ╟─+┼─+╢
                ║Samples with allele \'TAG\'\s*│\s*[0-9]+\s*\([0-9\.]+\%\)║
                ╟─+┼─+╢
                ║Samples with allele missing\s*│\s*[0-9]+\s*\([0-9\.]+\%\)║
                ╟─+┼─+╢
                ║Has missing data\s*│\s*True║
                ╟─+┼─+╢
                ║Isolated as missing\s*│\s*True║
                ╚═+╧═+╝
                """[
                    1:
                ]
            ),
            str(v),
        )

    def test_variant_str_no_samples(self):
        tables = tskit.TableCollection(10)
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        tables.sites.add_row(position=5, ancestral_state="A")
        tables.mutations.add_row(site=0, node=0, derived_state="T")
        ts = tables.tree_sequence()
        v = next(ts.variants(samples=[]))
        for allele in v.alleles:
            if allele is not None:
                assert (
                    re.search(
                        rf"║Samples with allele '{allele}'\s*│\s*0\s*\(nan\%\)║", str(v)
                    )
                    is not None
                )

    def test_variant_str_no_site(self):
        tables = tskit.TableCollection(10)
        ts = tables.tree_sequence()
        v = tskit.Variant(ts)
        s = str(v)
        assert len(s.splitlines()) == 5
        assert (
            "This variant has not yet been decoded at a specific site, "
            + "call Variant.decode to set the site"
            in s
        )

    def test_variant_html_repr(self, ts_fixture):
        v = next(ts_fixture.variants())
        html = v._repr_html_()
        # Parse to check valid
        ElementTree.fromstring(html)
        assert len(html) > 1900

    def test_variant_html_repr_no_site(self):
        tables = tskit.TableCollection(10)
        ts = tables.tree_sequence()
        v = tskit.Variant(ts)
        html = v._repr_html_()
        ElementTree.fromstring(html)
        assert len(html) > 1600

    def test_variant_repr(self, ts_fixture):
        v = next(ts_fixture.variants())
        str_rep = repr(v)
        assert len(str_rep) > 0 and len(str_rep) < 10000
        assert re.search(r"\AVariant", str_rep)
        assert re.search(rf"\'site\': Site\(id={v.site.id}", str_rep)
        assert re.search(rf"position={v.position}", str_rep)
        alleles = re.escape("'alleles': " + str(v.alleles))
        assert re.search(rf"{alleles}", str_rep)
        assert re.search(r"\'genotypes\': array\(\[", str_rep)
        assert re.search(rf"position={v.position}", str_rep)
        assert re.search(rf"\'has_missing_data\': {v.has_missing_data}", str_rep)
        assert re.search(rf"\'isolated_as_missing\': {v.isolated_as_missing}", str_rep)


--- ../../tskit/python/tests/test_haplotype_matching.py ---


"""
Python implementation of the Li and Stephens forwards and backwards algorithms.
"""
import warnings

import lshmm as ls
import msprime
import numpy as np
import numpy.testing as nt
import pytest

import _tskit
import tskit
from tests import tsutil

MISSING = -1


def check_alleles(alleles, m):
    """
    Checks the specified allele list and returns a list of lists
    of alleles of length num_sites.
    If alleles is a 1D list of strings, assume that this list is used
    for each site and return num_sites copies of this list.
    Otherwise, raise a ValueError if alleles is not a list of length
    num_sites.
    """
    if isinstance(alleles[0], str):
        return [alleles for _ in range(m)], np.int8([len(alleles) for _ in range(m)])
    if len(alleles) != m:
        raise ValueError("Malformed alleles list")
    n_alleles = np.int8([(len(alleles_site)) for alleles_site in alleles])
    return alleles, n_alleles


def mirror_coordinates(ts):
    """
    Returns a copy of the specified tree sequence in which all
    coordinates x are transformed into L - x.
    """
    L = ts.sequence_length
    tables = ts.dump_tables()
    left = tables.edges.left
    right = tables.edges.right
    tables.edges.left = L - right
    tables.edges.right = L - left
    tables.sites.position = L - tables.sites.position  # + 1
    # TODO migrations.
    tables.sort()
    return tables.tree_sequence()


class ValueTransition:
    """Simple struct holding value transition values."""

    def __init__(self, tree_node=-1, value=-1, value_index=-1):
        self.tree_node = tree_node
        self.value = value
        self.value_index = value_index

    def copy(self):
        return ValueTransition(
            self.tree_node,
            self.value,
            self.value_index,
        )

    def __repr__(self):
        return repr(self.__dict__)

    def __str__(self):
        return repr(self)


class LsHmmAlgorithm:
    """
    Abstract superclass of Li and Stephens HMM algorithm.
    """

    def __init__(
        self, ts, rho, mu, alleles, n_alleles, precision=10, scale_mutation=False
    ):
        self.ts = ts
        self.mu = mu
        self.rho = rho
        self.precision = precision
        # The array of ValueTransitions.
        self.T = []
        # indexes in to the T array for each node.
        self.T_index = np.zeros(ts.num_nodes, dtype=int) - 1
        # The number of nodes underneath each element in the T array.
        self.N = np.zeros(ts.num_nodes, dtype=int)
        # Efficiently compute the allelic state at a site
        self.allelic_state = np.zeros(ts.num_nodes, dtype=int) - 1
        # TreePosition so we can can update T and T_index between trees.
        self.tree_pos = tsutil.TreePosition(ts)
        self.parent = np.zeros(self.ts.num_nodes, dtype=int) - 1
        self.tree = tskit.Tree(self.ts)
        self.output = None
        # Vector of the number of alleles at each site
        self.n_alleles = n_alleles
        self.alleles = alleles
        self.scale_mutation_based_on_n_alleles = scale_mutation

    def check_integrity(self):
        M = [st.tree_node for st in self.T if st.tree_node != -1]
        assert np.all(self.T_index[M] >= 0)
        index = np.ones_like(self.T_index, dtype=bool)
        index[M] = 0
        assert np.all(self.T_index[index] == -1)
        for j, st in enumerate(self.T):
            if st.tree_node != -1:
                assert j == self.T_index[st.tree_node]

    def compress(self):
        tree = self.tree
        T = self.T
        T_index = self.T_index

        values = np.unique(list(st.value if st.tree_node != -1 else 1e200 for st in T))
        for st in T:
            if st.tree_node != -1:
                st.value_index = np.searchsorted(values, st.value)

        child = np.zeros(len(values), dtype=int)
        num_values = len(values)
        value_count = np.zeros(num_values, dtype=int)

        def compute(u, parent_state):
            value_count[:] = 0
            for v in tree.children(u):
                child[:] = optimal_set[v]
                # If the set for a given child is empty, then we know it inherits
                # directly from the parent state and must be a singleton set.
                if np.sum(child) == 0:
                    child[parent_state] = 1
                for j in range(num_values):
                    value_count[j] += child[j]
            max_value_count = np.max(value_count)
            optimal_set[u, :] = 0
            optimal_set[u, value_count == max_value_count] = 1

        optimal_set = np.zeros((tree.tree_sequence.num_nodes, len(values)), dtype=int)
        t_node_time = [
            -1 if st.tree_node == -1 else tree.time(st.tree_node) for st in T
        ]
        order = np.argsort(t_node_time)
        for j in order:
            st = T[j]
            u = st.tree_node
            if u != -1:
                # Compute the value at this node
                state = st.value_index
                if tree.is_internal(u):
                    compute(u, state)
                else:
                    # A[u, state] = 1
                    optimal_set[u, state] = 1
                # Find parent state
                v = tree.parent(u)
                if v != -1:
                    while T_index[v] == -1:
                        v = tree.parent(v)
                    parent_state = T[T_index[v]].value_index
                    v = tree.parent(u)
                    while T_index[v] == -1:
                        compute(v, parent_state)
                        v = tree.parent(v)

        T_old = [st.copy() for st in T]
        T.clear()
        T_parent = []

        old_state = T_old[T_index[tree.root]].value_index
        new_state = np.argmax(optimal_set[tree.root])

        T.append(ValueTransition(tree_node=tree.root, value=values[new_state]))
        T_parent.append(-1)
        stack = [(tree.root, old_state, new_state, 0)]
        while len(stack) > 0:
            u, old_state, new_state, t_parent = stack.pop()
            for v in tree.children(u):
                old_child_state = old_state
                if T_index[v] != -1:
                    old_child_state = T_old[T_index[v]].value_index
                if np.sum(optimal_set[v]) > 0:
                    new_child_state = new_state
                    child_t_parent = t_parent

                    if optimal_set[v, new_state] == 0:
                        new_child_state = np.argmax(optimal_set[v])
                        child_t_parent = len(T)
                        T_parent.append(t_parent)
                        T.append(
                            ValueTransition(tree_node=v, value=values[new_child_state])
                        )
                    stack.append((v, old_child_state, new_child_state, child_t_parent))
                else:
                    if old_child_state != new_state:
                        T_parent.append(t_parent)
                        T.append(
                            ValueTransition(tree_node=v, value=values[old_child_state])
                        )

        for st in T_old:
            if st.tree_node != -1:
                T_index[st.tree_node] = -1
        for j, st in enumerate(T):
            T_index[st.tree_node] = j
            self.N[j] = tree.num_samples(st.tree_node)
        for j in range(len(T)):
            if T_parent[j] != -1:
                self.N[T_parent[j]] -= self.N[j]

    def update_tree(self, direction=tskit.FORWARD):
        """
        Update the internal data structures to move on to the next tree.
        """
        parent = self.parent
        T_index = self.T_index
        T = self.T
        if direction == tskit.FORWARD:
            self.tree_pos.next()
        else:
            self.tree_pos.prev()
        assert self.tree_pos.index == self.tree.index

        for j in range(
            self.tree_pos.out_range.start, self.tree_pos.out_range.stop, direction
        ):
            e = self.tree_pos.out_range.order[j]
            edge = self.ts.edge(e)
            u = edge.child
            if T_index[u] == -1:
                # Make sure the subtree we're detaching has an T_index-value at the root.
                while T_index[u] == -1:
                    u = parent[u]
                    assert u != -1
                T_index[edge.child] = len(T)
                T.append(
                    ValueTransition(tree_node=edge.child, value=T[T_index[u]].value)
                )
            parent[edge.child] = -1

        for j in range(
            self.tree_pos.in_range.start, self.tree_pos.in_range.stop, direction
        ):
            e = self.tree_pos.in_range.order[j]
            edge = self.ts.edge(e)
            parent[edge.child] = edge.parent
            u = edge.parent
            if parent[edge.parent] == -1:
                # Grafting onto a new root.
                if T_index[edge.parent] == -1:
                    T_index[edge.parent] = len(T)
                    T.append(
                        ValueTransition(
                            tree_node=edge.parent, value=T[T_index[edge.child]].value
                        )
                    )
            else:
                # Grafting into an existing subtree.
                while T_index[u] == -1:
                    u = parent[u]
                    assert u != -1
            assert T_index[u] != -1 and T_index[edge.child] != -1
            if T[T_index[u]].value == T[T_index[edge.child]].value:
                st = T[T_index[edge.child]]
                # Mark the lower ValueTransition as unused.
                st.value = -1
                st.tree_node = -1
                T_index[edge.child] = -1

        # We can have values left over still pointing to old roots. Remove
        for root in self.tree.roots:
            if T_index[root] != -1:
                # Use a special marker here to designate the real roots.
                T[T_index[root]].value_index = -2
        for vt in T:
            if vt.tree_node != -1:
                if parent[vt.tree_node] == -1 and vt.value_index != -2:
                    T_index[vt.tree_node] = -1
                    vt.tree_node = -1
                vt.value_index = -1

    def update_probabilities(self, site, haplotype_state):
        tree = self.tree
        T_index = self.T_index
        T = self.T
        alleles = self.alleles[site.id]
        allelic_state = self.allelic_state
        # Set the allelic_state for this site.
        allelic_state[tree.root] = alleles.index(site.ancestral_state)

        for mutation in site.mutations:
            u = mutation.node
            allelic_state[u] = alleles.index(mutation.derived_state)
            if T_index[u] == -1:
                while T_index[u] == tskit.NULL:
                    u = tree.parent(u)
                T_index[mutation.node] = len(T)
                T.append(
                    ValueTransition(tree_node=mutation.node, value=T[T_index[u]].value)
                )

        for st in T:
            u = st.tree_node
            if u != -1:
                # Get the allelic_state at u. TODO we can cache these states to
                # avoid some upward traversals.
                v = u
                while allelic_state[v] == -1:
                    v = tree.parent(v)
                    assert v != -1
                match = haplotype_state == allelic_state[v]
                is_query_missing = haplotype_state == MISSING
                # Note that the node u is used only by Viterbi
                st.value = self.compute_next_probability(
                    site.id, st.value, match, is_query_missing, u
                )

        # Unset the states
        allelic_state[tree.root] = -1
        for mutation in site.mutations:
            allelic_state[mutation.node] = -1

    def process_site(self, site, haplotype_state):
        self.update_probabilities(site, haplotype_state)
        self.compress()
        s = self.compute_normalisation_factor()
        for st in self.T:
            assert st.tree_node != tskit.NULL
            # if st.tree_node != tskit.NULL:
            st.value /= s
            st.value = round(st.value, self.precision)
        self.output.store_site(site.id, s, [(st.tree_node, st.value) for st in self.T])

    def compute_emission_proba(self, site_id, is_match):
        mu = self.mu[site_id]
        n_alleles = self.n_alleles[site_id]
        if self.scale_mutation_based_on_n_alleles:
            if is_match:
                # Scale mutation based on the number of alleles
                # - so the mutation rate is the mutation rate to one of the
                # alleles. The overall mutation rate is then
                # (n_alleles - 1) * mutation_rate.
                p_e = 1 - (n_alleles - 1) * mu
            else:
                p_e = mu - mu * (n_alleles == 1)
                # Added boolean in case we're at an invariant site
        else:
            # No scaling based on the number of alleles
            #  - so the mutation rate is the mutation rate to anything.
            # This means that we must rescale the mutation rate to a different
            # allele, by the number of alleles.
            if n_alleles == 1:  # In case we're at an invariant site
                if is_match:
                    p_e = 1
                else:
                    p_e = 0
            else:
                if is_match:
                    p_e = 1 - mu
                else:
                    p_e = mu / (n_alleles - 1)
        return p_e

    def initialise(self, value):
        self.tree.clear()
        for u in self.ts.samples():
            j = len(self.T)
            self.T_index[u] = j
            self.T.append(ValueTransition(tree_node=u, value=value))

    def run(self, h):
        n = self.ts.num_samples
        self.initialise(1 / n)
        while self.tree.next():
            self.update_tree()
            for site in self.tree.sites():
                self.process_site(site, h[site.id])
        return self.output

    def compute_normalisation_factor(self):
        raise NotImplementedError()

    def compute_next_probability(
        self, site_id, p_last, is_match, is_query_missing, node
    ):
        raise NotImplementedError()


class ForwardAlgorithm(LsHmmAlgorithm):
    """
    The Li and Stephens forward algorithm.
    """

    def __init__(
        self, ts, rho, mu, alleles, n_alleles, scale_mutation=False, precision=10
    ):
        super().__init__(
            ts,
            rho,
            mu,
            alleles,
            n_alleles,
            precision=precision,
            scale_mutation=scale_mutation,
        )
        self.output = CompressedMatrix(ts)

    def compute_normalisation_factor(self):
        s = 0
        for j, st in enumerate(self.T):
            assert st.tree_node != tskit.NULL
            # assert self.N[j] > 0
            s += self.N[j] * st.value
        return s

    def compute_next_probability(
        self, site_id, p_last, is_match, is_query_missing, node
    ):
        rho = self.rho[site_id]
        n = self.ts.num_samples
        if is_query_missing:
            p_e = 1.0
        else:
            p_e = self.compute_emission_proba(site_id, is_match)
        p_t = p_last * (1 - rho) + rho / n
        return p_t * p_e


class BackwardAlgorithm(ForwardAlgorithm):
    """
    The Li and Stephens backward algorithm.
    """

    def compute_next_probability(
        self, site_id, p_next, is_match, is_query_missing, node
    ):
        if is_query_missing:
            p_e = 1.0
        else:
            p_e = self.compute_emission_proba(site_id, is_match)
        return p_next * p_e

    def process_site(self, site, haplotype_state, s):
        # FIXME see nodes in the C code for why we have two calls to
        # compress
        # https://github.com/tskit-dev/tskit/issues/2803
        self.compress()
        self.output.store_site(
            site.id,
            s,
            [(st.tree_node, st.value) for st in self.T],
        )
        self.update_probabilities(site, haplotype_state)
        # FIXME see nodes in the C code for why we have two calls to
        # compress
        self.compress()
        b_last_sum = self.compute_normalisation_factor()
        n = self.ts.num_samples
        rho = self.rho[site.id]
        for st in self.T:
            if st.tree_node != tskit.NULL:
                st.value = rho * b_last_sum / n + (1 - rho) * st.value
                st.value /= s
                st.value = round(st.value, self.precision)

    def run(self, h, normalisation_factor):
        self.initialise(value=1)
        while self.tree.prev():
            self.update_tree(direction=tskit.REVERSE)
            for site in reversed(list(self.tree.sites())):
                self.process_site(site, h[site.id], normalisation_factor[site.id])
        return self.output


class ViterbiAlgorithm(LsHmmAlgorithm):
    """
    Runs the Li and Stephens Viterbi algorithm.
    """

    def __init__(
        self, ts, rho, mu, alleles, n_alleles, scale_mutation=False, precision=10
    ):
        super().__init__(
            ts,
            rho,
            mu,
            alleles,
            n_alleles,
            precision=precision,
            scale_mutation=scale_mutation,
        )
        self.output = ViterbiMatrix(ts)

    def compute_normalisation_factor(self):
        max_st = ValueTransition(value=-1)
        for st in self.T:
            assert st.tree_node != tskit.NULL
            if st.value > max_st.value:
                max_st = st
        if max_st.value == 0:
            raise ValueError(
                "Trying to match non-existent allele with zero mutation rate"
            )
        return max_st.value

    def compute_next_probability(
        self, site_id, p_last, is_match, is_query_missing, node
    ):
        rho = self.rho[site_id]
        n = self.ts.num_samples

        p_no_recomb = p_last * (1 - rho + rho / n)
        p_recomb = rho / n
        recombination_required = False
        if p_no_recomb > p_recomb:
            p_t = p_no_recomb
        else:
            p_t = p_recomb
            recombination_required = True
        self.output.add_recombination_required(site_id, node, recombination_required)

        if is_query_missing:
            p_e = 1.0
        else:
            p_e = self.compute_emission_proba(site_id, is_match)

        return p_t * p_e


def assert_compressed_matrices_equal(cm1, cm2):
    nt.assert_array_almost_equal(cm1.normalisation_factor, cm2.normalisation_factor)

    for j in range(cm1.num_sites):
        site1 = cm1.get_site(j)
        site2 = cm2.get_site(j)
        assert len(site1) == len(site2)
        site1 = dict(site1)
        site2 = dict(site2)

        assert set(site1.keys()) == set(site2.keys())
        for node in site1.keys():
            # TODO  the precision value should be used as a parameter here
            nt.assert_allclose(site1[node], site2[node], rtol=1e-5, atol=1e-8)


class CompressedMatrix:
    """
    Class representing a num_samples x num_sites matrix compressed by a
    tree sequence. Each site is represented by a set of (node, value)
    pairs, which act as "mutations", i.e., any sample that descends
    from a particular node will inherit that value (unless any other
    values are on the path).
    """

    def __init__(self, ts):
        self.ts = ts
        self.num_sites = ts.num_sites
        self.num_samples = ts.num_samples
        self.value_transitions = [None for _ in range(self.num_sites)]
        self.normalisation_factor = np.zeros(self.num_sites)

    def store_site(self, site, normalisation_factor, value_transitions):
        assert all(u >= 0 for u, _ in value_transitions)
        self.normalisation_factor[site] = normalisation_factor
        self.value_transitions[site] = value_transitions

    # Expose the same API as the low-level classes

    @property
    def num_transitions(self):
        a = [len(self.value_transitions[j]) for j in range(self.num_sites)]
        return np.array(a, dtype=np.int32)

    def get_site(self, site):
        return self.value_transitions[site]

    def decode(self):
        """
        Decodes the tree encoding of the values into an explicit
        matrix.
        """
        sample_index_map = np.zeros(self.ts.num_nodes, dtype=int) - 1
        sample_index_map[self.ts.samples()] = np.arange(self.ts.num_samples)
        A = np.zeros((self.num_sites, self.num_samples))
        for tree in self.ts.trees():
            for site in tree.sites():
                for node, value in self.value_transitions[site.id]:
                    for u in tree.samples(node):
                        j = sample_index_map[u]
                        A[site.id, j] = value
        return A


class ViterbiMatrix(CompressedMatrix):
    """
    Class representing the compressed Viterbi matrix.
    """

    def __init__(self, ts):
        super().__init__(ts)
        # Tuple containing the site, the node in the tree, and whether
        # recombination is required
        self.recombination_required = [(-1, 0, False)]

    def add_recombination_required(self, site, node, required):
        self.recombination_required.append((site, node, required))

    def choose_sample(self, site_id, tree):
        max_value = -1
        u = -1
        for node, value in self.value_transitions[site_id]:
            if value > max_value:
                max_value = value
                u = node
        assert u != -1

        transition_nodes = [u for (u, _) in self.value_transitions[site_id]]
        while not tree.is_sample(u):
            for v in tree.children(u):
                if v not in transition_nodes:
                    u = v
                    break
            else:
                raise AssertionError("could not find path")
        return u

    def traceback(self):
        # Run the traceback.
        m = self.ts.num_sites
        match = np.zeros(m, dtype=int)
        recombination_tree = np.zeros(self.ts.num_nodes, dtype=int) - 1
        tree = tskit.Tree(self.ts)
        tree.last()
        current_node = -1

        rr_index = len(self.recombination_required) - 1
        for site in reversed(self.ts.sites()):
            while tree.interval.left > site.position:
                tree.prev()
            assert tree.interval.left <= site.position < tree.interval.right

            # Fill in the recombination tree
            j = rr_index
            while self.recombination_required[j][0] == site.id:
                u, required = self.recombination_required[j][1:]
                recombination_tree[u] = required
                j -= 1

            if current_node == -1:
                current_node = self.choose_sample(site.id, tree)
            match[site.id] = current_node

            # Now traverse up the tree from the current node. The first marked node
            # we meet tells us whether we need to recombine.
            u = current_node
            while u != -1 and recombination_tree[u] == -1:
                u = tree.parent(u)

            assert u != -1
            if recombination_tree[u] == 1:
                # Need to switch at the next site.
                current_node = -1
            # Reset the nodes in the recombination tree.
            j = rr_index
            while self.recombination_required[j][0] == site.id:
                u, required = self.recombination_required[j][1:]
                recombination_tree[u] = -1
                j -= 1
            rr_index = j

        return match


def get_site_alleles(ts, h, alleles):
    if alleles is None:
        n_alleles = np.zeros(ts.num_sites, dtype=np.int8) - 1
        for j in range(ts.num_sites):
            uniq_alleles = np.unique(np.append(ts.genotype_matrix()[j, :], h[j]))
            uniq_alleles = uniq_alleles[uniq_alleles != MISSING]
            n_alleles[j] = len(uniq_alleles)
        assert np.all(n_alleles > 0)
        alleles = tskit.ALLELES_ACGT
        if len(set(alleles).intersection(next(ts.variants()).alleles)) == 0:
            alleles = tskit.ALLELES_01
            if len(set(alleles).intersection(next(ts.variants()).alleles)) == 0:
                raise ValueError(
                    """Alleles list could not be identified.
                    Please pass a list of lists of alleles of length m,
                    or a list of alleles (e.g. tskit.ALLELES_ACGT)"""
                )
        alleles = [alleles for _ in range(ts.num_sites)]
    else:
        alleles, n_alleles = check_alleles(alleles, ts.num_sites)
    return alleles, n_alleles


def ls_forward_tree(
    h, ts, rho, mu, precision=30, alleles=None, scale_mutation_based_on_n_alleles=False
):
    alleles, n_alleles = get_site_alleles(ts, h, alleles)
    fa = ForwardAlgorithm(
        ts,
        rho,
        mu,
        alleles,
        n_alleles,
        precision=precision,
        scale_mutation=scale_mutation_based_on_n_alleles,
    )
    return fa.run(h)


def ls_backward_tree(h, ts, rho, mu, normalisation_factor, precision=30, alleles=None):
    alleles, n_alleles = get_site_alleles(ts, h, alleles)
    ba = BackwardAlgorithm(
        ts,
        rho,
        mu,
        alleles,
        n_alleles,
        precision=precision,
    )
    return ba.run(h, normalisation_factor)


def ls_viterbi_tree(
    h, ts, rho, mu, precision=30, alleles=None, scale_mutation_based_on_n_alleles=False
):
    alleles, n_alleles = get_site_alleles(ts, h, alleles)
    va = ViterbiAlgorithm(
        ts,
        rho,
        mu,
        alleles,
        n_alleles,
        precision=precision,
        scale_mutation=scale_mutation_based_on_n_alleles,
    )
    return va.run(h)


class LSBase:
    """Superclass of Li and Stephens tests."""

    def example_haplotypes(self, ts):
        H = ts.genotype_matrix()
        s = H[:, 0].reshape(1, H.shape[0])
        H = H[:, 1:]

        haplotypes = [
            s,
            H[:, -1].reshape(1, H.shape[0]),
        ]
        s_tmp = s.copy()
        s_tmp[0, -1] = MISSING
        haplotypes.append(s_tmp)
        s_tmp = s.copy()
        s_tmp[0, ts.num_sites // 2] = MISSING
        haplotypes.append(s_tmp)
        s_tmp = s.copy()
        s_tmp[0, :] = MISSING
        haplotypes.append(s_tmp)

        return H, haplotypes

    def example_parameters_haplotypes(self, ts, seed=42):
        """Returns an iterator over combinations of haplotype,
        recombination and mutation rates."""
        np.random.seed(seed)
        H, haplotypes = self.example_haplotypes(ts)
        n = H.shape[1]
        m = ts.get_num_sites()

        # Here we have equal mutation and recombination
        r = np.zeros(m) + 0.01
        mu = np.zeros(m) + 0.01
        r[0] = 0

        for s in haplotypes:
            yield n, H, s, r, mu

        # FIXME removing these as tests are abominably slow.
        # We'll be refactoring all this to use pytest anyway, so let's not
        # worry too much about coverage for now.
        # # Mixture of random and extremes
        # rs = [np.zeros(m) + 0.999, np.zeros(m) + 1e-6, np.random.rand(m)]
        # mus = [np.zeros(m) + 0.33, np.zeros(m) + 1e-6, np.random.rand(m) * 0.33]

        # import itertools
        # for s, r, mu in itertools.product(haplotypes, rs, mus):
        #     r[0] = 0
        #     yield n, H, s, r, mu

    def assertAllClose(self, A, B):
        """Assert that all entries of two matrices are 'close'"""
        assert np.allclose(A, B, rtol=1e-5, atol=1e-8)

    # Define a bunch of very small tree-sequences for testing a collection
    # of parameters on
    def test_simple_n_10_no_recombination(self):
        ts = msprime.simulate(
            10, recombination_rate=0, mutation_rate=0.5, random_seed=42
        )
        assert ts.num_sites > 3
        self.verify(ts)

    def test_simple_n_10_no_recombination_high_mut(self):
        ts = msprime.simulate(10, recombination_rate=0, mutation_rate=3, random_seed=42)
        assert ts.num_sites > 3
        self.verify(ts)

    def test_simple_n_10_no_recombination_higher_mut(self):
        ts = msprime.simulate(20, recombination_rate=0, mutation_rate=3, random_seed=42)
        assert ts.num_sites > 3
        self.verify(ts)

    def test_simple_n_6(self):
        ts = msprime.simulate(6, recombination_rate=2, mutation_rate=7, random_seed=42)
        assert ts.num_sites > 5
        self.verify(ts)

    def test_simple_n_8(self):
        ts = msprime.simulate(8, recombination_rate=2, mutation_rate=5, random_seed=42)
        assert ts.num_sites > 5
        self.verify(ts)

    def test_simple_n_8_high_recombination(self):
        ts = msprime.simulate(8, recombination_rate=20, mutation_rate=5, random_seed=42)
        assert ts.num_trees > 15
        assert ts.num_sites > 5
        self.verify(ts)

    def test_simple_n_16(self):
        ts = msprime.simulate(16, recombination_rate=2, mutation_rate=5, random_seed=42)
        assert ts.num_sites > 5
        self.verify(ts)

    # # Define a bunch of very small tree-sequences for testing a collection
    # # of parameters on
    # def test_simple_n_10_no_recombination_blah(self):
    #     ts = msprime.sim_ancestry(
    #         samples=10,
    #         recombination_rate=0,
    #         random_seed=42,
    #         sequence_length=10,
    #         population_size=10000,
    #     )
    #     ts = msprime.sim_mutations(ts, rate=1e-5, random_seed=42)
    #     assert ts.num_sites > 3
    #     self.verify(ts)

    # def test_simple_n_6_blah(self):
    # ts = msprime.sim_ancestry(
    #     samples=6,
    #     recombination_rate=1e-4,
    #     random_seed=42,
    #     sequence_length=40,
    #     population_size=10000,
    # )
    # ts = msprime.sim_mutations(ts, rate=1e-3, random_seed=42)
    #     assert ts.num_sites > 5
    #     self.verify(ts)

    # def test_simple_n_8_blah(self):
    #     ts = msprime.sim_ancestry(
    #         samples=8,
    #         recombination_rate=1e-4,
    #         random_seed=42,
    #         sequence_length=20,
    #         population_size=10000,
    #     )
    #     ts = msprime.sim_mutations(ts, rate=1e-4, random_seed=42)
    #     assert ts.num_sites > 5
    #     assert ts.num_trees > 15
    #     self.verify(ts)

    # def test_simple_n_16_blah(self):
    #     ts = msprime.sim_ancestry(
    #         samples=16,
    #         recombination_rate=1e-2,
    #         random_seed=42,
    #         sequence_length=20,
    #         population_size=10000,
    #     )
    #     ts = msprime.sim_mutations(ts, rate=1e-4, random_seed=42)
    #     assert ts.num_sites > 5
    #     self.verify(ts)

    def verify(self, ts):
        raise NotImplementedError()


class FBAlgorithmBase(LSBase):
    """Base for forwards backwards algorithm tests."""


class VitAlgorithmBase(LSBase):
    """Base for viterbi algoritm tests."""


class TestMirroringHap(FBAlgorithmBase):
    """Tests that mirroring the tree sequence and running forwards and backwards
    algorithms gives the same log-likelihood of observing the data."""

    def verify(self, ts):
        for n, H, s, r, mu in self.example_parameters_haplotypes(ts):
            # Note, need to remove the first sample from the ts, and ensure that
            # invariant sites aren't removed.
            ts_check = ts.simplify(range(1, n + 1), filter_sites=False)
            cm = ls_forward_tree(s[0, :], ts_check, r, mu)
            ll_tree = np.sum(np.log10(cm.normalisation_factor))

            ts_check_mirror = mirror_coordinates(ts_check)
            r_flip = np.insert(np.flip(r)[:-1], 0, 0)
            cm_mirror = ls_forward_tree(
                np.flip(s[0, :]), ts_check_mirror, r_flip, np.flip(mu)
            )
            ll_mirror_tree = np.sum(np.log10(cm_mirror.normalisation_factor))
            self.assertAllClose(ll_tree, ll_mirror_tree)

            # Ensure that the decoded matrices are the same
            flipped_H = np.flip(H, axis=0)
            flipped_s = np.flip(s, axis=1)
            F_mirror_matrix, c, ll = ls.forwards(
                reference_panel=flipped_H,
                query=flipped_s,
                ploidy=1,
                prob_recombination=r_flip,
                prob_mutation=np.flip(mu),
                scale_mutation_rate=False,
            )

            self.assertAllClose(F_mirror_matrix, cm_mirror.decode())
            self.assertAllClose(ll, ll_tree)


class TestForwardHapTree(FBAlgorithmBase):
    """Tests that the tree algorithm computes the same forward matrix as the
    simple method."""

    def verify(self, ts):
        for n, H, s, r, mu in self.example_parameters_haplotypes(ts):
            for scale_mutation in [False, True]:
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")
                    # Warning from lshmm:
                    # Passed a vector of mutation rates, but rescaling each mutation
                    # rate conditional on the number of alleles
                    F, c, ll = ls.forwards(
                        reference_panel=H,
                        query=s,
                        ploidy=1,
                        prob_recombination=r,
                        prob_mutation=mu,
                        scale_mutation_rate=scale_mutation,
                    )
                # Note, need to remove the first sample from the ts, and ensure
                # that invariant sites aren't removed.
                ts_check = ts.simplify(range(1, n + 1), filter_sites=False)
                cm = ls_forward_tree(
                    s[0, :],
                    ts_check,
                    r,
                    mu,
                    scale_mutation_based_on_n_alleles=scale_mutation,
                )
                self.assertAllClose(cm.decode(), F)
                ll_tree = np.sum(np.log10(cm.normalisation_factor))
                self.assertAllClose(ll, ll_tree)


class TestForwardBackwardTree(FBAlgorithmBase):
    """Tests that the tree algorithm computes the same forward matrix as the
    simple method."""

    def verify(self, ts):
        for n, H, s, r, mu in self.example_parameters_haplotypes(ts):
            F, c, ll = ls.forwards(
                reference_panel=H,
                query=s,
                ploidy=1,
                prob_recombination=r,
                prob_mutation=mu,
                scale_mutation_rate=False,
            )
            B = ls.backwards(
                reference_panel=H,
                query=s,
                ploidy=1,
                normalisation_factor_from_forward=c,
                prob_recombination=r,
                prob_mutation=mu,
                scale_mutation_rate=False,
            )

            # Note, need to remove the first sample from the ts, and ensure that
            # invariant sites aren't removed.
            ts_check = ts.simplify(range(1, n + 1), filter_sites=False)
            c_f = ls_forward_tree(s[0, :], ts_check, r, mu)
            ll_tree = np.sum(np.log10(c_f.normalisation_factor))

            c_b = ls_backward_tree(
                s[0, :],
                ts_check,
                r,
                mu,
                c_f.normalisation_factor,
            )
            B_tree = c_b.decode()

            F_tree = c_f.decode()

            self.assertAllClose(B, B_tree)
            self.assertAllClose(F, F_tree)
            self.assertAllClose(ll, ll_tree)


class TestTreeViterbiHap(VitAlgorithmBase):
    """Test that we have the same log-likelihood between tree and matrix
    implementations"""

    def verify(self, ts):
        for n, H, s, r, mu in self.example_parameters_haplotypes(ts):
            path, ll = ls.viterbi(
                reference_panel=H,
                query=s,
                ploidy=1,
                prob_recombination=r,
                prob_mutation=mu,
                scale_mutation_rate=False,
            )
            ts_check = ts.simplify(range(1, n + 1), filter_sites=False)
            cm = ls_viterbi_tree(s[0, :], ts_check, r, mu)
            ll_tree = np.sum(np.log10(cm.normalisation_factor))
            self.assertAllClose(ll, ll_tree)

            # Now, need to ensure that the likelihood of the preferred path is
            # the same as ll_tree (and ll).
            path_tree = cm.traceback()
            ll_check = ls.path_loglik(
                reference_panel=H,
                query=s,
                ploidy=1,
                path=path_tree,
                prob_recombination=r,
                prob_mutation=mu,
                scale_mutation_rate=False,
            )
            self.assertAllClose(ll, ll_check)


# TODO add params to run the various checks
def check_viterbi(ts, h, recombination=None, mutation=None):
    h = np.array(h).astype(np.int8)
    m = ts.num_sites
    assert len(h) == m
    if recombination is None:
        recombination = np.zeros(ts.num_sites) + 1e-9
    if mutation is None:
        mutation = np.zeros(ts.num_sites)
    precision = 22

    G = ts.genotype_matrix()
    s = h.reshape(1, m)

    path, ll = ls.viterbi(
        reference_panel=G,
        query=s,
        ploidy=1,
        prob_recombination=recombination,
        prob_mutation=mutation,
        scale_mutation_rate=False,
    )
    assert np.isscalar(ll)

    cm = ls_viterbi_tree(h, ts, rho=recombination, mu=mutation)
    ll_tree = np.sum(np.log10(cm.normalisation_factor))
    assert np.isscalar(ll_tree)
    nt.assert_allclose(ll_tree, ll)

    # Check that the likelihood of the preferred path is
    # the same as ll_tree (and ll).
    path_tree = cm.traceback()
    ll_check = ls.path_loglik(
        reference_panel=G,
        query=s,
        ploidy=1,
        path=path_tree,
        prob_recombination=recombination,
        prob_mutation=mutation,
        scale_mutation_rate=False,
    )
    nt.assert_allclose(ll_check, ll)

    ll_ts = ts._ll_tree_sequence
    ls_hmm = _tskit.LsHmm(ll_ts, recombination, mutation, precision=precision)
    cm_lib = _tskit.ViterbiMatrix(ll_ts)
    ls_hmm.viterbi_matrix(h, cm_lib)
    path_lib = cm_lib.traceback()

    # Not true in general, but let's see how far it goes
    nt.assert_array_equal(path_lib, path_tree)

    nt.assert_allclose(cm_lib.normalisation_factor, cm.normalisation_factor)

    return path


# TODO add params to run the various checks
def check_forward_matrix(ts, h, recombination=None, mutation=None):
    precision = 22
    h = np.array(h).astype(np.int8)
    n = ts.num_samples
    m = ts.num_sites
    assert len(h) == m
    if recombination is None:
        recombination = np.zeros(ts.num_sites) + 1e-9
    if mutation is None:
        mutation = np.zeros(ts.num_sites)

    G = ts.genotype_matrix()
    s = h.reshape(1, m)

    F, c, ll = ls.forwards(
        reference_panel=G,
        query=s,
        ploidy=1,
        prob_recombination=recombination,
        prob_mutation=mutation,
        scale_mutation_rate=False,
    )
    assert F.shape == (m, n)
    assert c.shape == (m,)
    assert np.isscalar(ll)

    cm = ls_forward_tree(
        h, ts, recombination, mutation, scale_mutation_based_on_n_alleles=False
    )
    F2 = cm.decode()
    nt.assert_allclose(F, F2)
    nt.assert_allclose(c, cm.normalisation_factor)
    ll_tree = np.sum(np.log10(cm.normalisation_factor))
    nt.assert_allclose(ll_tree, ll)

    ll_ts = ts._ll_tree_sequence
    ls_hmm = _tskit.LsHmm(ll_ts, recombination, mutation, precision=precision)
    cm_lib = _tskit.CompressedMatrix(ll_ts)
    ls_hmm.forward_matrix(h, cm_lib)
    F3 = cm_lib.decode()

    assert_compressed_matrices_equal(cm, cm_lib)

    nt.assert_allclose(F, F3)
    nt.assert_allclose(c, cm_lib.normalisation_factor)
    return cm_lib


def check_backward_matrix(ts, h, forward_cm, recombination=None, mutation=None):
    precision = 22
    h = np.array(h).astype(np.int8)
    m = ts.num_sites
    assert len(h) == m
    if recombination is None:
        recombination = np.zeros(ts.num_sites) + 1e-9
    if mutation is None:
        mutation = np.zeros(ts.num_sites)

    G = ts.genotype_matrix()
    s = h.reshape(1, m)

    B = ls.backwards(
        reference_panel=G,
        query=s,
        ploidy=1,
        normalisation_factor_from_forward=forward_cm.normalisation_factor,
        prob_recombination=recombination,
        prob_mutation=mutation,
        scale_mutation_rate=False,
    )

    backward_cm = ls_backward_tree(
        h,
        ts,
        recombination,
        mutation,
        forward_cm.normalisation_factor,
        precision=precision,
    )
    nt.assert_array_equal(
        backward_cm.normalisation_factor, forward_cm.normalisation_factor
    )

    ll_ts = ts._ll_tree_sequence
    ls_hmm = _tskit.LsHmm(ll_ts, recombination, mutation, precision=precision)
    cm_lib = _tskit.CompressedMatrix(ll_ts)
    ls_hmm.backward_matrix(h, forward_cm.normalisation_factor, cm_lib)

    assert_compressed_matrices_equal(backward_cm, cm_lib)

    B_lib = cm_lib.decode()
    B_tree = backward_cm.decode()
    nt.assert_allclose(B_tree, B_lib)
    nt.assert_allclose(B, B_lib)


def add_unique_sample_mutations(ts, start=0):
    """
    Adds a mutation for each of the samples at equally spaced locations
    along the genome.
    """
    tables = ts.dump_tables()
    L = int(ts.sequence_length)
    assert L % ts.num_samples == 0
    gap = L // ts.num_samples
    x = start
    for u in ts.samples():
        site = tables.sites.add_row(position=x, ancestral_state="0")
        tables.mutations.add_row(site=site, derived_state="1", node=u)
        x += gap
    return tables.tree_sequence()


class TestSingleBalancedTreeExample:
    # 3.00┊    6    ┊
    #     ┊  ┏━┻━┓  ┊
    # 2.00┊  4   5  ┊
    #     ┊ ┏┻┓ ┏┻┓ ┊
    # 1.00┊ 0 1 2 3 ┊
    #     0         8

    @staticmethod
    def ts():
        return add_unique_sample_mutations(
            tskit.Tree.generate_balanced(4, span=8).tree_sequence,
            start=1,
        )

    @pytest.mark.parametrize("j", [0, 1, 2, 3])
    def test_match_sample(self, j):
        ts = self.ts()
        h = np.zeros(4)
        h[j] = 1
        path = check_viterbi(ts, h)
        nt.assert_array_equal([j, j, j, j], path)
        cm = check_forward_matrix(ts, h)
        check_backward_matrix(ts, h, cm)

    @pytest.mark.parametrize("j", [1, 2])
    def test_match_sample_missing_flanks(self, j):
        ts = self.ts()
        h = np.zeros(4)
        h[0] = -1
        h[-1] = -1
        h[j] = 1
        path = check_viterbi(ts, h)
        nt.assert_array_equal([j, j, j, j], path)
        cm = check_forward_matrix(ts, h)
        check_backward_matrix(ts, h, cm)

    def test_switch_each_sample(self):
        ts = self.ts()
        h = np.ones(4)
        path = check_viterbi(ts, h)
        nt.assert_array_equal([0, 1, 2, 3], path)
        cm = check_forward_matrix(ts, h)
        check_backward_matrix(ts, h, cm)

    def test_switch_each_sample_missing_flanks(self):
        ts = self.ts()
        h = np.ones(4)
        h[0] = -1
        h[-1] = -1
        path = check_viterbi(ts, h)
        nt.assert_array_equal([1, 1, 2, 2], path)
        cm = check_forward_matrix(ts, h)
        check_backward_matrix(ts, h, cm)

    def test_switch_each_sample_missing_middle(self):
        ts = self.ts()
        h = np.ones(4)
        h[1:3] = -1
        path = check_viterbi(ts, h)
        # Implementation of Viterbi switches at right-most position
        nt.assert_array_equal([0, 3, 3, 3], path)
        cm = check_forward_matrix(ts, h)
        check_backward_matrix(ts, h, cm)


class TestSimulationExamples:
    @pytest.mark.parametrize("n", [3, 10, 50])
    @pytest.mark.parametrize("L", [1, 10, 100])
    def test_continuous_genome(self, n, L):
        ts = msprime.simulate(
            n, length=L, recombination_rate=1, mutation_rate=1, random_seed=42
        )
        h = np.zeros(ts.num_sites, dtype=np.int8)
        # NOTE this is a bit slow at the moment but we can disable the Python
        # implementation once testing has been improved on smaller examples.
        # Add ``compare_py=False``to these calls.
        check_viterbi(ts, h)
        cm = check_forward_matrix(ts, h)
        check_backward_matrix(ts, h, cm)


--- ../../tskit/python/tests/test_fileobj.py ---

 
"""
Test cases for loading and dumping different types of files and streams
"""
import io
import multiprocessing
import os
import pathlib
import platform
import queue
import shutil
import socket
import socketserver
import tempfile
import traceback

import pytest
import tszip
from pytest import fixture

import tskit


IS_WINDOWS = platform.system() == "Windows"
IS_OSX = platform.system() == "Darwin"


class TestPath:
    @fixture
    def tempfile_name(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            yield f"{tmp_dir}/plain_path"

    def test_pathlib(self, ts_fixture, tempfile_name):
        ts_fixture.dump(tempfile_name)
        ts2 = tskit.load(tempfile_name)
        assert ts_fixture.tables == ts2.tables


class TestPathLib:
    @fixture
    def pathlib_tempfile(self):
        fd, path = tempfile.mkstemp(prefix="tskit_test_pathlib")
        os.close(fd)
        temp_file = pathlib.Path(path)
        yield temp_file
        temp_file.unlink()

    def test_pathlib(self, ts_fixture, pathlib_tempfile):
        ts_fixture.dump(pathlib_tempfile)
        ts2 = tskit.load(pathlib_tempfile)
        assert ts_fixture.tables == ts2.tables


class TestFileObj:
    @fixture
    def fileobj(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            with open(f"{tmp_dir}/fileobj", "wb") as f:
                yield f

    def test_fileobj(self, ts_fixture, fileobj):
        ts_fixture.dump(fileobj)
        fileobj.close()
        ts2 = tskit.load(fileobj.name)
        assert ts_fixture.tables == ts2.tables

    def test_fileobj_multi(self, replicate_ts_fixture, fileobj):
        file_offsets = []
        for ts in replicate_ts_fixture:
            ts.dump(fileobj)
            file_offsets.append(fileobj.tell())
        fileobj.close()
        with open(fileobj.name, "rb") as f:
            for ts, file_offset in zip(replicate_ts_fixture, file_offsets):
                ts2 = tskit.load(f)
                file_offset2 = f.tell()
                assert ts.tables == ts2.tables
                assert file_offset == file_offset2


class TestFileObjRW:
    @fixture
    def fileobj(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            pathlib.Path(f"{tmp_dir}/fileobj").touch()
            with open(f"{tmp_dir}/fileobj", "r+b") as f:
                yield f

    def test_fileobj(self, ts_fixture, fileobj):
        ts_fixture.dump(fileobj)
        fileobj.seek(0)
        ts2 = tskit.load(fileobj)
        assert ts_fixture.tables == ts2.tables

    def test_fileobj_multi(self, replicate_ts_fixture, fileobj):
        file_offsets = []
        for ts in replicate_ts_fixture:
            ts.dump(fileobj)
            file_offsets.append(fileobj.tell())
        fileobj.seek(0)
        for ts, file_offset in zip(replicate_ts_fixture, file_offsets):
            ts2 = tskit.load(fileobj)
            file_offset2 = fileobj.tell()
            assert ts.tables == ts2.tables
            assert file_offset == file_offset2


class TestFD:
    @fixture
    def fd(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            pathlib.Path(f"{tmp_dir}/fd").touch()
            with open(f"{tmp_dir}/fd", "r+b") as f:
                yield f.fileno()

    def test_fd(self, ts_fixture, fd):
        ts_fixture.dump(fd)
        os.lseek(fd, 0, os.SEEK_SET)
        ts2 = tskit.load(fd)
        assert ts_fixture.tables == ts2.tables

    def test_fd_multi(self, replicate_ts_fixture, fd):
        for ts in replicate_ts_fixture:
            ts.dump(fd)
        os.lseek(fd, 0, os.SEEK_SET)
        for ts in replicate_ts_fixture:
            ts2 = tskit.load(fd)
            assert ts.tables == ts2.tables


class TestUnsupportedObjects:
    def test_string_io(self, ts_fixture):
        with pytest.raises(io.UnsupportedOperation, match=r"fileno"):
            ts_fixture.dump(io.StringIO())
        with pytest.raises(io.UnsupportedOperation, match=r"fileno"):
            tskit.load(io.StringIO())
        with pytest.raises(io.UnsupportedOperation, match=r"fileno"):
            ts_fixture.dump(io.BytesIO())
        with pytest.raises(io.UnsupportedOperation, match=r"fileno"):
            tskit.load(io.BytesIO())


def dump_to_stream(q_err, q_in, file_out):
    """
    Get tree sequences from `q_in` and ts.dump() them to `file_out`.
    Uncaught exceptions are placed onto the `q_err` queue.
    """
    try:
        with open(file_out, "wb") as f:
            while True:
                ts = q_in.get()
                if ts is None:
                    break
                ts.dump(f)
    except Exception as exc:
        tb = traceback.format_exc()
        q_err.put((exc, tb))


def load_from_stream(q_err, q_out, file_in):
    """
    tskit.load() tree sequences from `file_in` and put them onto `q_out`.
    Uncaught exceptions are placed onto the `q_err` queue.
    """
    try:
        with open(file_in, "rb") as f:
            while True:
                try:
                    ts = tskit.load(f)
                except EOFError:
                    break
                q_out.put(ts)
    except Exception as exc:
        tb = traceback.format_exc()
        q_err.put((exc, tb))


def stream(fifo, ts_list):
    """
    data -> q_in -> ts.dump(fifo) -> tskit.load(fifo) -> q_out -> data_out
    """
    q_err = multiprocessing.Queue()
    q_in = multiprocessing.Queue()
    q_out = multiprocessing.Queue()
    proc1 = multiprocessing.Process(target=dump_to_stream, args=(q_err, q_in, fifo))
    proc2 = multiprocessing.Process(target=load_from_stream, args=(q_err, q_out, fifo))
    proc1.start()
    proc2.start()
    for data in ts_list:
        q_in.put(data)

    q_in.put(None)  # signal the process that we're done
    proc1.join(timeout=3)
    if not q_err.empty():
        # re-raise the first child exception
        exc, tb = q_err.get()
        print(tb)
        raise exc
    if proc1.is_alive():
        # prevent hang if proc1 failed to join
        proc1.terminate()
        proc2.terminate()
        raise RuntimeError("proc1 (ts.dump) failed to join")
    ts_list_out = []
    for _ in ts_list:
        try:
            data_out = q_out.get(timeout=3)
        except queue.Empty:
            # terminate proc2 so we don't hang
            proc2.terminate()
            raise
        ts_list_out.append(data_out)
    proc2.join(timeout=3)
    if proc2.is_alive():
        # prevent hang if proc2 failed to join
        proc2.terminate()
        raise RuntimeError("proc2 (tskit.load) failed to join")

    assert len(ts_list) == len(ts_list_out)
    for ts, ts_out in zip(ts_list, ts_list_out):
        assert ts.tables == ts_out.tables


@pytest.mark.skipif(IS_WINDOWS, reason="No FIFOs on Windows")
@pytest.mark.skipif(IS_OSX, reason="FIFO flakey on OS X, issue #1170")
class TestFIFO:
    @fixture
    def fifo(self):
        temp_dir = tempfile.mkdtemp(prefix="tsk_test_streaming")
        temp_fifo = os.path.join(temp_dir, "fifo")
        os.mkfifo(temp_fifo)
        yield temp_fifo
        shutil.rmtree(temp_dir)

    def test_single_stream(self, fifo, ts_fixture):
        stream(fifo, [ts_fixture])

    def test_multi_stream(self, fifo, replicate_ts_fixture):
        stream(fifo, replicate_ts_fixture)


ADDRESS = ("localhost", 10009)


class Server(socketserver.ThreadingTCPServer):
    allow_reuse_address = True


class StoreEchoHandler(socketserver.BaseRequestHandler):
    def handle(self):
        while True:
            try:
                ts = tskit.load(self.request.fileno())
            except EOFError:
                break
            ts.dump(self.request.fileno())
        self.server.shutdown()


def server_process(q):
    server = Server(ADDRESS, StoreEchoHandler)
    # Tell the client (on the other end of the queue) that it's OK to open
    # a connection
    q.put(None)
    server.serve_forever()


@pytest.mark.skipif(IS_WINDOWS or IS_OSX, reason="Errors on systems without proper fds")
class TestSocket:
    @fixture
    def client_fd(self):
        # Use a queue to synchronise the startup of the server and the client.
        q = multiprocessing.Queue()
        _server_process = multiprocessing.Process(target=server_process, args=(q,))
        _server_process.start()
        q.get(timeout=3)
        client = socket.create_connection(ADDRESS)
        yield client.fileno()
        client.close()
        _server_process.join(timeout=3)

    def verify_stream(self, ts_list, client_fd):
        for ts in ts_list:
            ts.dump(client_fd)
            echo_ts = tskit.load(client_fd)
            assert ts.tables == echo_ts.tables

    def test_single_then_multi(self, ts_fixture, replicate_ts_fixture, client_fd):
        self.verify_stream([ts_fixture], client_fd)
        self.verify_stream(replicate_ts_fixture, client_fd)


def write_to_fifo(path, file_path):
    with open(path, "wb") as fifo:
        with open(file_path, "rb") as file:
            fifo.write(file.read())


def read_from_fifo(path, expected_exception, error_text, read_func):
    with open(path) as fifo:
        with pytest.raises(expected_exception, match=error_text):
            read_func(fifo)


def write_and_read_from_fifo(fifo_path, file_path, expected_exception, error_text):
    os.mkfifo(fifo_path)
    for read_func in [tskit.load, tskit.TableCollection.load]:
        read_process = multiprocessing.Process(
            target=read_from_fifo,
            args=(fifo_path, expected_exception, error_text, read_func),
        )
        read_process.start()
        write_process = multiprocessing.Process(
            target=write_to_fifo, args=(fifo_path, file_path)
        )
        write_process.start()
        write_process.join(timeout=3)
        read_process.join(timeout=3)


@pytest.mark.skipif(IS_WINDOWS, reason="No FIFOs on Windows")
class TestBadStream:
    def test_bad_stream(self, tmp_path):
        fifo_path = tmp_path / "fifo"
        bad_file_path = tmp_path / "bad_file"
        bad_file_path.write_bytes(b"bad data")
        write_and_read_from_fifo(
            fifo_path, bad_file_path, tskit.FileFormatError, "not in kastore format"
        )

    def test_legacy_stream(self, tmp_path):
        fifo_path = tmp_path / "fifo"
        legacy_file_path = os.path.join(
            os.path.dirname(__file__), "data", "hdf5-formats", "msprime-0.3.0_v2.0.hdf5"
        )
        write_and_read_from_fifo(
            fifo_path, legacy_file_path, tskit.FileFormatError, "not in kastore format"
        )

    def test_tszip_stream(self, tmp_path, ts_fixture):
        fifo_path = tmp_path / "fifo"
        zip_file_path = tmp_path / "tszip_file"
        tszip.compress(ts_fixture, zip_file_path)
        write_and_read_from_fifo(
            fifo_path, zip_file_path, tskit.FileFormatError, "not in kastore format"
        )


--- ../../tskit/python/tests/test_intervals.py ---


"""
Test cases for the intervals module.
"""
import decimal
import fractions
import gzip
import io
import os
import pickle
import textwrap
import xml

import numpy as np
import pytest
from numpy.testing import assert_array_equal

import tskit


class TestRateMapErrors:
    @pytest.mark.parametrize(
        ("position", "rate"),
        [
            ([], []),
            ([0], []),
            ([0], [0]),
            ([1, 2], [0]),
            ([0, -1], [0]),
            ([0, 1], [-1]),
        ],
    )
    def test_bad_input(self, position, rate):
        with pytest.raises(ValueError):
            tskit.RateMap(position=position, rate=rate)

    def test_zero_length_interval(self):
        with pytest.raises(ValueError, match=r"at indexes \[2 4\]"):
            tskit.RateMap(position=[0, 1, 1, 2, 2, 3], rate=[0, 0, 0, 0, 0])

    def test_bad_length(self):
        positions = np.array([0, 1, 2])
        rates = np.array([0, 1, 2])
        with pytest.raises(ValueError, match="one less entry"):
            tskit.RateMap(position=positions, rate=rates)

    def test_bad_first_pos(self):
        positions = np.array([1, 2, 3])
        rates = np.array([1, 1])
        with pytest.raises(ValueError, match="First position"):
            tskit.RateMap(position=positions, rate=rates)

    def test_bad_rate(self):
        positions = np.array([0, 1, 2])
        rates = np.array([1, -1])
        with pytest.raises(ValueError, match="negative.*1"):
            tskit.RateMap(position=positions, rate=rates)

    def test_bad_rate_with_missing(self):
        positions = np.array([0, 1, 2])
        rates = np.array([np.nan, -1])
        with pytest.raises(ValueError, match="negative.*1"):
            tskit.RateMap(position=positions, rate=rates)

    def test_read_only(self):
        positions = np.array([0, 0.25, 0.5, 0.75, 1])
        rates = np.array([0.125, 0.25, 0.5, 0.75])  # 1 shorter than positions
        rate_map = tskit.RateMap(position=positions, rate=rates)
        assert np.all(rates == rate_map.rate)
        assert np.all(positions == rate_map.position)
        with pytest.raises(AttributeError):
            rate_map.rate = 2 * rate_map.rate
        with pytest.raises(AttributeError):
            rate_map.position = 2 * rate_map.position
        with pytest.raises(AttributeError):
            rate_map.left = 1234
        with pytest.raises(AttributeError):
            rate_map.right = 1234
        with pytest.raises(AttributeError):
            rate_map.mid = 1234
        with pytest.raises(ValueError):
            rate_map.rate[0] = 1
        with pytest.raises(ValueError):
            rate_map.position[0] = 1
        with pytest.raises(ValueError):
            rate_map.left[0] = 1
        with pytest.raises(ValueError):
            rate_map.mid[0] = 1
        with pytest.raises(ValueError):
            rate_map.right[0] = 1


class TestGetRateAllKnown:
    examples = [
        tskit.RateMap(position=[0, 1], rate=[0]),
        tskit.RateMap(position=[0, 1], rate=[0.1]),
        tskit.RateMap(position=[0, 1, 2], rate=[0.1, 0.2]),
        tskit.RateMap(position=[0, 1, 2], rate=[0, 0.2]),
        tskit.RateMap(position=[0, 1, 2], rate=[0.1, 1e-6]),
        tskit.RateMap(position=range(100), rate=range(99)),
    ]

    @pytest.mark.parametrize("rate_map", examples)
    def test_get_rate_mid(self, rate_map):
        rate = rate_map.get_rate(rate_map.mid)
        assert len(rate) == len(rate_map)
        for j in range(len(rate_map)):
            assert rate[j] == rate_map[rate_map.mid[j]]

    @pytest.mark.parametrize("rate_map", examples)
    def test_get_rate_left(self, rate_map):
        rate = rate_map.get_rate(rate_map.left)
        assert len(rate) == len(rate_map)
        for j in range(len(rate_map)):
            assert rate[j] == rate_map[rate_map.left[j]]

    @pytest.mark.parametrize("rate_map", examples)
    def test_get_rate_right(self, rate_map):
        rate = rate_map.get_rate(rate_map.right[:-1])
        assert len(rate) == len(rate_map) - 1
        for j in range(len(rate_map) - 1):
            assert rate[j] == rate_map[rate_map.right[j]]


class TestOperations:
    examples = [
        tskit.RateMap.uniform(sequence_length=1, rate=0),
        tskit.RateMap.uniform(sequence_length=1, rate=0.1),
        tskit.RateMap(position=[0, 1, 2], rate=[0.1, 0.2]),
        tskit.RateMap(position=[0, 1, 2], rate=[0, 0.2]),
        tskit.RateMap(position=[0, 1, 2], rate=[0.1, 1e-6]),
        tskit.RateMap(position=range(100), rate=range(99)),
        # Missing data
        tskit.RateMap(position=[0, 1, 2], rate=[np.nan, 0]),
        tskit.RateMap(position=[0, 1, 2], rate=[0, np.nan]),
        tskit.RateMap(position=[0, 1, 2, 3], rate=[0, np.nan, 1]),
    ]

    @pytest.mark.parametrize("rate_map", examples)
    def test_num_intervals(self, rate_map):
        assert rate_map.num_intervals == len(rate_map.rate)
        assert rate_map.num_missing_intervals == np.sum(np.isnan(rate_map.rate))
        assert rate_map.num_non_missing_intervals == np.sum(~np.isnan(rate_map.rate))

    @pytest.mark.parametrize("rate_map", examples)
    def test_mask_arrays(self, rate_map):
        assert_array_equal(rate_map.missing, np.isnan(rate_map.rate))
        assert_array_equal(rate_map.non_missing, ~np.isnan(rate_map.rate))

    @pytest.mark.parametrize("rate_map", examples)
    def test_missing_intervals(self, rate_map):
        missing = []
        for left, right, rate in zip(rate_map.left, rate_map.right, rate_map.rate):
            if np.isnan(rate):
                missing.append([left, right])
        if len(missing) == 0:
            assert len(rate_map.missing_intervals()) == 0
        else:
            assert_array_equal(missing, rate_map.missing_intervals())

    @pytest.mark.parametrize("rate_map", examples)
    def test_mean_rate(self, rate_map):
        total_span = 0
        total_mass = 0
        for span, mass in zip(rate_map.span, rate_map.mass):
            if not np.isnan(mass):
                total_span += span
                total_mass += mass
        assert total_mass / total_span == rate_map.mean_rate

    @pytest.mark.parametrize("rate_map", examples)
    def test_total_mass(self, rate_map):
        assert rate_map.total_mass == np.nansum(rate_map.mass)

    @pytest.mark.parametrize("rate_map", examples)
    def test_get_cumulative_mass(self, rate_map):
        assert list(rate_map.get_cumulative_mass([0])) == [0]
        assert list(rate_map.get_cumulative_mass([rate_map.sequence_length])) == [
            rate_map.total_mass
        ]
        assert_array_equal(
            rate_map.get_cumulative_mass(rate_map.right), np.nancumsum(rate_map.mass)
        )

    @pytest.mark.parametrize("rate_map", examples)
    def test_get_rate(self, rate_map):
        assert_array_equal(rate_map.get_rate([0]), rate_map.rate[0])
        assert_array_equal(
            rate_map.get_rate([rate_map.sequence_length - 1e-9]), rate_map.rate[-1]
        )
        assert_array_equal(rate_map.get_rate(rate_map.left), rate_map.rate)

    @pytest.mark.parametrize("rate_map", examples)
    def test_map_semantics(self, rate_map):
        assert len(rate_map) == rate_map.num_non_missing_intervals
        assert_array_equal(list(rate_map.keys()), rate_map.mid[rate_map.non_missing])
        for x in rate_map.left[rate_map.missing]:
            assert x not in rate_map
        for x in rate_map.mid[rate_map.missing]:
            assert x not in rate_map

    def test_asdict(self):
        rate_map = tskit.RateMap.uniform(sequence_length=2, rate=4)
        d = rate_map.asdict()
        assert_array_equal(d["position"], np.array([0.0, 2.0]))
        assert_array_equal(d["rate"], np.array([4.0]))


class TestFindIndex:
    def test_one_interval(self):
        rate_map = tskit.RateMap(position=[0, 10], rate=[0.1])
        for j in range(10):
            assert rate_map.find_index(j) == 0
        assert rate_map.find_index(0.0001) == 0
        assert rate_map.find_index(9.999) == 0

    def test_two_intervals(self):
        rate_map = tskit.RateMap(position=[0, 5, 10], rate=[0.1, 0.1])
        assert rate_map.find_index(0) == 0
        assert rate_map.find_index(0.0001) == 0
        assert rate_map.find_index(4.9999) == 0
        assert rate_map.find_index(5) == 1
        assert rate_map.find_index(5.1) == 1
        assert rate_map.find_index(7) == 1
        assert rate_map.find_index(9.999) == 1

    def test_three_intervals(self):
        rate_map = tskit.RateMap(position=[0, 5, 10, 15], rate=[0.1, 0.1, 0.1])
        assert rate_map.find_index(0) == 0
        assert rate_map.find_index(0.0001) == 0
        assert rate_map.find_index(4.9999) == 0
        assert rate_map.find_index(5) == 1
        assert rate_map.find_index(5.1) == 1
        assert rate_map.find_index(7) == 1
        assert rate_map.find_index(9.999) == 1
        assert rate_map.find_index(10) == 2
        assert rate_map.find_index(10.1) == 2
        assert rate_map.find_index(12) == 2
        assert rate_map.find_index(14.9999) == 2

    def test_out_of_bounds(self):
        rate_map = tskit.RateMap(position=[0, 10], rate=[0.1])
        for bad_value in [-1, -0.0001, 10, 10.0001, 1e9]:
            with pytest.raises(KeyError, match="out of bounds"):
                rate_map.find_index(bad_value)

    def test_input_types(self):
        rate_map = tskit.RateMap(position=[0, 10], rate=[0.1])
        assert rate_map.find_index(0) == 0
        assert rate_map.find_index(0.0) == 0
        assert rate_map.find_index(np.zeros(1)[0]) == 0


class TestSimpleExamples:
    def test_all_missing_one_interval(self):
        with pytest.raises(ValueError, match="missing data"):
            tskit.RateMap(position=[0, 10], rate=[np.nan])

    def test_all_missing_two_intervals(self):
        with pytest.raises(ValueError, match="missing data"):
            tskit.RateMap(position=[0, 5, 10], rate=[np.nan, np.nan])

    def test_count(self):
        rate_map = tskit.RateMap(position=[0, 5, 10], rate=[np.nan, 1])
        assert rate_map.num_intervals == 2
        assert rate_map.num_missing_intervals == 1
        assert rate_map.num_non_missing_intervals == 1

    def test_missing_arrays(self):
        rate_map = tskit.RateMap(position=[0, 5, 10], rate=[np.nan, 1])
        assert list(rate_map.missing) == [True, False]
        assert list(rate_map.non_missing) == [False, True]

    def test_missing_at_start_mean_rate(self):
        positions = np.array([0, 0.5, 1, 2])
        rates = np.array([np.nan, 0, 1])
        rate_map = tskit.RateMap(position=positions, rate=rates)
        assert np.isclose(rate_map.mean_rate, 1 / (1 + 0.5))

    def test_missing_at_end_mean_rate(self):
        positions = np.array([0, 1, 1.5, 2])
        rates = np.array([1, 0, np.nan])
        rate_map = tskit.RateMap(position=positions, rate=rates)
        assert np.isclose(rate_map.mean_rate, 1 / (1 + 0.5))

    def test_interval_properties_all_known(self):
        rate_map = tskit.RateMap(position=[0, 1, 2, 3], rate=[0.1, 0.2, 0.3])
        assert list(rate_map.left) == [0, 1, 2]
        assert list(rate_map.right) == [1, 2, 3]
        assert list(rate_map.mid) == [0.5, 1.5, 2.5]
        assert list(rate_map.span) == [1, 1, 1]
        assert list(rate_map.mass) == [0.1, 0.2, 0.3]

    def test_pickle_non_missing(self):
        r1 = tskit.RateMap(position=[0, 1, 2, 3], rate=[0.1, 0.2, 0.3])
        r2 = pickle.loads(pickle.dumps(r1))
        assert r1 == r2

    def test_pickle_missing(self):
        r1 = tskit.RateMap(position=[0, 1, 2, 3], rate=[0.1, np.nan, 0.3])
        r2 = pickle.loads(pickle.dumps(r1))
        assert r1 == r2

    def test_get_cumulative_mass_all_known(self):
        rate_map = tskit.RateMap(position=[0, 10, 20, 30], rate=[0.1, 0.2, 0.3])
        assert list(rate_map.mass) == [1, 2, 3]
        assert list(rate_map.get_cumulative_mass([10, 20, 30])) == [1, 3, 6]

    def test_cumulative_mass_missing(self):
        rate_map = tskit.RateMap(position=[0, 10, 20, 30], rate=[0.1, np.nan, 0.3])
        assert list(rate_map.get_cumulative_mass([10, 20, 30])) == [1, 1, 4]


class TestDisplay:
    def test_str(self):
        rate_map = tskit.RateMap(position=[0, 10], rate=[0.1])
        s = """\
        ╔════╤═════╤═══╤════╤════╗
        ║left│right│mid│span│rate║
        ╠════╪═════╪═══╪════╪════╣
        ║0   │10   │  5│  10│ 0.1║
        ╚════╧═════╧═══╧════╧════╝
        """
        assert textwrap.dedent(s) == str(rate_map)

    def test_str_scinot(self):
        rate_map = tskit.RateMap(position=[0, 10], rate=[0.000001])
        s = """\
        ╔════╤═════╤═══╤════╤═════╗
        ║left│right│mid│span│rate ║
        ╠════╪═════╪═══╪════╪═════╣
        ║0   │10   │  5│  10│1e-06║
        ╚════╧═════╧═══╧════╧═════╝
        """
        assert textwrap.dedent(s) == str(rate_map)

    def test_repr(self):
        rate_map = tskit.RateMap(position=[0, 10], rate=[0.1])
        s = "RateMap(position=array([ 0., 10.]), rate=array([0.1]))"
        assert repr(rate_map) == s

    def test_repr_html(self):
        rate_map = tskit.RateMap(position=[0, 10], rate=[0.1])
        html = rate_map._repr_html_()
        root = xml.etree.ElementTree.fromstring(html)
        assert root.tag == "div"
        table = root.find("table")
        rows = list(table.find("tbody"))
        assert len(rows) == 1

    def test_long_table(self):
        n = 100
        rate_map = tskit.RateMap(position=range(n + 1), rate=[0.1] * n)
        headers, data = rate_map._text_header_and_rows(limit=20)
        assert len(headers) == 5
        assert len(data) == 21
        # check some left values
        assert int(data[0][0]) == 0
        assert int(data[-1][0]) == n - 1

    def test_short_table(self):
        n = 10
        rate_map = tskit.RateMap(position=range(n + 1), rate=[0.1] * n)
        headers, data = rate_map._text_header_and_rows(limit=20)
        assert len(headers) == 5
        assert len(data) == n
        # check some left values.
        assert int(data[0][0]) == 0
        assert int(data[-1][0]) == n - 1


class TestRateMapIsMapping:
    def test_items(self):
        rate_map = tskit.RateMap(position=[0, 1, 2, 3], rate=[0.1, 0.2, 0.3])
        items = list(rate_map.items())
        assert items[0] == (0.5, 0.1)
        assert items[1] == (1.5, 0.2)
        assert items[2] == (2.5, 0.3)

    def test_keys(self):
        rate_map = tskit.RateMap(position=[0, 1, 2, 3], rate=[0.1, 0.2, 0.3])
        assert list(rate_map.keys()) == [0.5, 1.5, 2.5]

    def test_values(self):
        rate_map = tskit.RateMap(position=[0, 1, 2, 3], rate=[0.1, 0.2, 0.3])
        assert list(rate_map.values()) == [0.1, 0.2, 0.3]

    def test_in_points(self):
        rate_map = tskit.RateMap(position=[0, 1, 2, 3], rate=[0.1, 0.2, 0.3])
        # Any point within the map are True
        for x in [0, 0.5, 1, 2.9999]:
            assert x in rate_map
        # Points outside the map are False
        for x in [-1, -0.0001, 3, 3.1]:
            assert x not in rate_map

    def test_in_slices(self):
        rate_map = tskit.RateMap(position=[0, 1, 2, 3], rate=[0.1, 0.2, 0.3])
        # slices that are within the map are "in"
        for x in [slice(0, 0.5), slice(0, 1), slice(0, 2), slice(2, 3), slice(0, 3)]:
            assert x in rate_map
        # Any slice that doesn't fully intersect with the map "not in"
        assert slice(-0.001, 1) not in rate_map
        assert slice(0, 3.0001) not in rate_map
        assert slice(2.9999, 3.0001) not in rate_map
        assert slice(3, 4) not in rate_map
        assert slice(-2, -1) not in rate_map

    def test_other_types_not_in(self):
        rate_map = tskit.RateMap(position=[0, 1, 2, 3], rate=[0.1, 0.2, 0.3])
        for other_type in [None, "sdf", "123", {}, [], Exception]:
            assert other_type not in rate_map

    def test_len(self):
        rate_map = tskit.RateMap(position=[0, 1], rate=[0.1])
        assert len(rate_map) == 1
        rate_map = tskit.RateMap(position=[0, 1, 2], rate=[0.1, 0.2])
        assert len(rate_map) == 2
        rate_map = tskit.RateMap(position=[0, 1, 2, 3], rate=[0.1, 0.2, 0.3])
        assert len(rate_map) == 3

    def test_immutable(self):
        rate_map = tskit.RateMap(position=[0, 1], rate=[0.1])
        with pytest.raises(TypeError, match="item assignment"):
            rate_map[0] = 1
        with pytest.raises(TypeError, match="item deletion"):
            del rate_map[0]

    def test_eq(self):
        r1 = tskit.RateMap(position=[0, 1, 2], rate=[0.1, 0.2])
        r2 = tskit.RateMap(position=[0, 1, 2], rate=[0.1, 0.2])
        assert r1 == r1
        assert r1 == r2
        r2 = tskit.RateMap(position=[0, 1, 3], rate=[0.1, 0.2])
        assert r1 != r2
        assert tskit.RateMap(position=[0, 1], rate=[0.1]) != tskit.RateMap(
            position=[0, 1], rate=[0.2]
        )
        assert tskit.RateMap(position=[0, 1], rate=[0.1]) != tskit.RateMap(
            position=[0, 10], rate=[0.1]
        )

    def test_getitem_value(self):
        rate_map = tskit.RateMap(position=[0, 1, 2], rate=[0.1, 0.2])
        assert rate_map[0] == 0.1
        assert rate_map[0.5] == 0.1
        assert rate_map[1] == 0.2
        assert rate_map[1.5] == 0.2
        assert rate_map[1.999] == 0.2
        # Try other types
        assert rate_map[np.array([1], dtype=np.float32)[0]] == 0.2
        assert rate_map[np.array([1], dtype=np.int32)[0]] == 0.2
        assert rate_map[np.array([1], dtype=np.float64)[0]] == 0.2
        assert rate_map[1 / 2] == 0.1
        assert rate_map[fractions.Fraction(1, 3)] == 0.1
        assert rate_map[decimal.Decimal(1)] == 0.2

    def test_getitem_slice(self):
        r1 = tskit.RateMap(position=[0, 1, 2], rate=[0.1, 0.2])
        # The semantics of the slice() function are tested elsewhere.
        assert r1[:] == r1.copy()
        assert r1[:] is not r1
        assert r1[1:] == r1.slice(left=1)
        assert r1[:1.5] == r1.slice(right=1.5)
        assert r1[0.5:1.5] == r1.slice(left=0.5, right=1.5)

    def test_getitem_slice_step(self):
        r1 = tskit.RateMap(position=[0, 1, 2], rate=[0.1, 0.2])
        # Trying to set a "step" is a error
        with pytest.raises(TypeError, match="interval slicing"):
            r1[0:3:1]


class TestMappingMissingData:
    def test_get_missing(self):
        rate_map = tskit.RateMap(position=[0, 1, 2], rate=[np.nan, 0.2])
        with pytest.raises(KeyError, match="within a missing interval"):
            rate_map[0]
        with pytest.raises(KeyError, match="within a missing interval"):
            rate_map[0.999]

    def test_in_missing(self):
        rate_map = tskit.RateMap(position=[0, 1, 2], rate=[np.nan, 0.2])
        assert 0 not in rate_map
        assert 0.999 not in rate_map
        assert 1 in rate_map

    def test_keys_missing(self):
        rate_map = tskit.RateMap(position=[0, 1, 2], rate=[np.nan, 0.2])
        assert list(rate_map.keys()) == [1.5]


class TestGetIntermediates:
    def test_get_rate(self):
        positions = np.array([0, 1, 2])
        rates = np.array([1, 4])
        rate_map = tskit.RateMap(position=positions, rate=rates)
        assert np.all(rate_map.get_rate([0.5, 1.5]) == rates)

    def test_get_rate_out_of_bounds(self):
        positions = np.array([0, 1, 2])
        rates = np.array([1, 4])
        rate_map = tskit.RateMap(position=positions, rate=rates)
        with pytest.raises(ValueError, match="out of bounds"):
            rate_map.get_rate([1, -0.1])
        with pytest.raises(ValueError, match="out of bounds"):
            rate_map.get_rate([2])

    def test_get_cumulative_mass(self):
        positions = np.array([0, 1, 2])
        rates = np.array([1, 4])
        rate_map = tskit.RateMap(position=positions, rate=rates)
        assert np.allclose(rate_map.get_cumulative_mass([0.5, 1.5]), np.array([0.5, 3]))
        assert rate_map.get_cumulative_mass([2]) == rate_map.total_mass

    def test_get_bad_cumulative_mass(self):
        positions = np.array([0, 1, 2])
        rates = np.array([1, 4])
        rate_map = tskit.RateMap(position=positions, rate=rates)
        with pytest.raises(ValueError, match="positions"):
            rate_map.get_cumulative_mass([1, -0.1])
        with pytest.raises(ValueError, match="positions"):
            rate_map.get_cumulative_mass([1, 2.1])


class TestSlice:
    def test_slice_no_params(self):
        # test RateMap.slice(..., trim=False)
        a = tskit.RateMap(position=[0, 100, 200, 300, 400], rate=[0, 1, 2, 3])
        b = a.slice()
        assert a.sequence_length == b.sequence_length
        assert_array_equal(a.position, b.position)
        assert_array_equal(a.rate, b.rate)
        assert a == b

    def test_slice_left_examples(self):
        a = tskit.RateMap(position=[0, 100, 200, 300, 400], rate=[0, 1, 2, 3])
        b = a.slice(left=50)
        assert a.sequence_length == b.sequence_length
        assert_array_equal([0, 50, 100, 200, 300, 400], b.position)
        assert_array_equal([np.nan, 0, 1, 2, 3], b.rate)

        b = a.slice(left=100)
        assert a.sequence_length == b.sequence_length
        assert_array_equal([0, 100, 200, 300, 400], b.position)
        assert_array_equal([np.nan, 1, 2, 3], b.rate)

        b = a.slice(left=150)
        assert a.sequence_length == b.sequence_length
        assert_array_equal([0, 150, 200, 300, 400], b.position)
        assert_array_equal([np.nan, 1, 2, 3], b.rate)

    def test_slice_right_examples(self):
        a = tskit.RateMap(position=[0, 100, 200, 300, 400], rate=[0, 1, 2, 3])
        b = a.slice(right=300)
        assert a.sequence_length == b.sequence_length
        assert_array_equal([0, 100, 200, 300, 400], b.position)
        assert_array_equal([0, 1, 2, np.nan], b.rate)

        b = a.slice(right=250)
        assert a.sequence_length == b.sequence_length
        assert_array_equal([0, 100, 200, 250, 400], b.position)
        assert_array_equal([0, 1, 2, np.nan], b.rate)

    def test_slice_left_right_examples(self):
        a = tskit.RateMap(position=[0, 100, 200, 300, 400], rate=[0, 1, 2, 3])
        b = a.slice(left=50, right=300)
        assert a.sequence_length == b.sequence_length
        assert_array_equal([0, 50, 100, 200, 300, 400], b.position)
        assert_array_equal([np.nan, 0, 1, 2, np.nan], b.rate)

        b = a.slice(left=150, right=250)
        assert a.sequence_length == b.sequence_length
        assert_array_equal([0, 150, 200, 250, 400], b.position)
        assert_array_equal([np.nan, 1, 2, np.nan], b.rate)

        b = a.slice(left=150, right=300)
        assert a.sequence_length == b.sequence_length
        assert_array_equal([0, 150, 200, 300, 400], b.position)
        assert_array_equal([np.nan, 1, 2, np.nan], b.rate)

        b = a.slice(left=150, right=160)
        assert a.sequence_length == b.sequence_length
        assert_array_equal([0, 150, 160, 400], b.position)
        assert_array_equal([np.nan, 1, np.nan], b.rate)

    def test_slice_right_missing(self):
        # If we take a right-slice into a trailing missing region,
        # we should recover the same map.
        a = tskit.RateMap(position=[0, 100, 200, 300, 400], rate=[0, 1, 2, np.nan])
        b = a.slice(right=350)
        assert a.sequence_length == b.sequence_length
        assert_array_equal(a.position, b.position)
        assert_array_equal(a.rate, b.rate)

        b = a.slice(right=300)
        assert a.sequence_length == b.sequence_length
        assert_array_equal(a.position, b.position)
        assert_array_equal(a.rate, b.rate)

    def test_slice_left_missing(self):
        a = tskit.RateMap(position=[0, 100, 200, 300, 400], rate=[np.nan, 1, 2, 3])
        b = a.slice(left=50)
        assert a.sequence_length == b.sequence_length
        assert_array_equal(a.position, b.position)
        assert_array_equal(a.rate, b.rate)

        b = a.slice(left=100)
        assert a.sequence_length == b.sequence_length
        assert_array_equal(a.position, b.position)
        assert_array_equal(a.rate, b.rate)

    def test_slice_with_floats(self):
        #  test RateMap.slice(..., trim=False) with floats
        a = tskit.RateMap(
            position=[np.pi * x for x in [0, 100, 200, 300, 400]], rate=[0, 1, 2, 3]
        )
        b = a.slice(left=50 * np.pi)
        assert a.sequence_length == b.sequence_length
        assert_array_equal([0, 50 * np.pi] + list(a.position[1:]), b.position)
        assert_array_equal([np.nan] + list(a.rate), b.rate)

    def test_slice_trim_left(self):
        a = tskit.RateMap(position=[0, 100, 200, 300, 400], rate=[1, 2, 3, 4])
        b = a.slice(left=100, trim=True)
        assert b == tskit.RateMap(position=[0, 100, 200, 300], rate=[2, 3, 4])
        b = a.slice(left=50, trim=True)
        assert b == tskit.RateMap(position=[0, 50, 150, 250, 350], rate=[1, 2, 3, 4])

    def test_slice_trim_right(self):
        a = tskit.RateMap(position=[0, 100, 200, 300, 400], rate=[1, 2, 3, 4])
        b = a.slice(right=300, trim=True)
        assert b == tskit.RateMap(position=[0, 100, 200, 300], rate=[1, 2, 3])
        b = a.slice(right=350, trim=True)
        assert b == tskit.RateMap(position=[0, 100, 200, 300, 350], rate=[1, 2, 3, 4])

    def test_slice_error(self):
        recomb_map = tskit.RateMap(position=[0, 100], rate=[1])
        with pytest.raises(KeyError):
            recomb_map.slice(left=-1)
        with pytest.raises(KeyError):
            recomb_map.slice(right=-1)
        with pytest.raises(KeyError):
            recomb_map.slice(left=200)
        with pytest.raises(KeyError):
            recomb_map.slice(right=200)
        with pytest.raises(KeyError):
            recomb_map.slice(left=20, right=10)


class TestReadHapmap:
    def test_read_hapmap_simple(self):
        hapfile = io.StringIO(
            """\
            HEADER
            chr1 1 x 0
            chr1 2 x 0.000001 x
            chr1 3 x 0.000006 x x"""
        )
        rm = tskit.RateMap.read_hapmap(hapfile)
        assert_array_equal(rm.position, [0, 1, 2, 3])
        assert np.allclose(rm.rate, [np.nan, 1e-8, 5e-8], equal_nan=True)

    def test_read_hapmap_from_filename(self, tmp_path):
        with open(tmp_path / "hapfile.txt", "w") as hapfile:
            hapfile.write(
                """\
                HEADER
                chr1 1 x 0
                chr1 2 x 0.000001 x
                chr1 3 x 0.000006 x x"""
            )
        rm = tskit.RateMap.read_hapmap(tmp_path / "hapfile.txt")
        assert_array_equal(rm.position, [0, 1, 2, 3])
        assert np.allclose(rm.rate, [np.nan, 1e-8, 5e-8], equal_nan=True)

    @pytest.mark.filterwarnings("ignore:loadtxt")
    def test_read_hapmap_empty(self):
        hapfile = io.StringIO(
            """\
            HEADER"""
        )
        with pytest.raises(ValueError, match="Empty"):
            tskit.RateMap.read_hapmap(hapfile)

    def test_read_hapmap_col_pos(self):
        hapfile = io.StringIO(
            """\
            HEADER
            0 0
            0.000001 1 x
            0.000006 2 x x"""
        )
        rm = tskit.RateMap.read_hapmap(hapfile, position_col=1, map_col=0)
        assert_array_equal(rm.position, [0, 1, 2])
        assert np.allclose(rm.rate, [1e-8, 5e-8])

    def test_read_hapmap_map_and_rate(self):
        hapfile = io.StringIO(
            """\
            HEADER
            chr1 0 0 0
            chr1 1 1 0.000001 x
            chr1 2 2 0.000006 x x"""
        )
        with pytest.raises(ValueError, match="both rate_col and map_col"):
            tskit.RateMap.read_hapmap(hapfile, rate_col=2, map_col=3)

    def test_read_hapmap_duplicate_pos(self):
        hapfile = io.StringIO(
            """\
            HEADER
            0 0
            0.000001 1 x
            0.000006 2 x x"""
        )
        with pytest.raises(ValueError, match="same columns"):
            tskit.RateMap.read_hapmap(hapfile, map_col=1)

    def test_read_hapmap_nonzero_rate_start(self):
        hapfile = io.StringIO(
            """\
            HEADER
            chr1 1 5 x
            chr1 2 0 x x x"""
        )
        rm = tskit.RateMap.read_hapmap(hapfile, rate_col=2)
        assert_array_equal(rm.position, [0, 1, 2])
        assert_array_equal(rm.rate, [np.nan, 5e-8])

    def test_read_hapmap_nonzero_rate_end(self):
        hapfile = io.StringIO(
            """\
            HEADER
            chr1 0 5 x
            chr1 2 1 x x x"""
        )
        with pytest.raises(ValueError, match="last entry.*must be zero"):
            tskit.RateMap.read_hapmap(hapfile, rate_col=2)

    def test_read_hapmap_gzipped(self, tmp_path):
        hapfile = os.path.join(tmp_path, "hapmap.txt.gz")
        with gzip.GzipFile(hapfile, "wb") as gzfile:
            gzfile.write(b"HEADER\n")
            gzfile.write(b"chr1 0 1\n")
            gzfile.write(b"chr1 1 5.5\n")
            gzfile.write(b"chr1 2 0\n")
        rm = tskit.RateMap.read_hapmap(hapfile, rate_col=2)
        assert_array_equal(rm.position, [0, 1, 2])
        assert_array_equal(rm.rate, [1e-8, 5.5e-8])

    def test_read_hapmap_nonzero_map_start(self):
        hapfile = io.StringIO(
            """\
            HEADER
            chr1 1 x 0.000001
            chr1 2 x 0.000001 x
            chr1 3 x 0.000006 x x x"""
        )
        rm = tskit.RateMap.read_hapmap(hapfile)
        assert_array_equal(rm.position, [0, 1, 2, 3])
        assert np.allclose(rm.rate, [1e-8, 0, 5e-8])

    def test_read_hapmap_bad_nonzero_map_start(self):
        hapfile = io.StringIO(
            """\
            HEADER
            chr1 0 x 0.0000005
            chr1 1 x 0.000001 x
            chr1 2 x 0.000006 x x x"""
        )
        with pytest.raises(ValueError, match="start.*must be zero"):
            tskit.RateMap.read_hapmap(hapfile)

    def test_sequence_length(self):
        hapfile = io.StringIO(
            """\
            HEADER
            chr1 0 x 0
            chr1 1 x 0.000001 x
            chr1 2 x 0.000006 x x x"""
        )
        # test identical seq len
        rm = tskit.RateMap.read_hapmap(hapfile, sequence_length=2)
        assert_array_equal(rm.position, [0, 1, 2])
        assert np.allclose(rm.rate, [1e-8, 5e-8])

        hapfile.seek(0)
        rm = tskit.RateMap.read_hapmap(hapfile, sequence_length=10)
        assert_array_equal(rm.position, [0, 1, 2, 10])
        assert np.allclose(rm.rate, [1e-8, 5e-8, np.nan], equal_nan=True)

    def test_bad_sequence_length(self):
        hapfile = io.StringIO(
            """\
            HEADER
            chr1 0 x 0
            chr1 1 x 0.000001 x
            chr1 2 x 0.000006 x x x"""
        )
        with pytest.raises(ValueError, match="sequence_length"):
            tskit.RateMap.read_hapmap(hapfile, sequence_length=1.999)

    def test_no_header(self):
        data = """\
            chr1 0 x 0
            chr1 1 x 0.000001 x
            chr1 2 x 0.000006 x x x"""
        hapfile_noheader = io.StringIO(data)
        hapfile_header = io.StringIO("chr pos rate cM\n" + data)
        with pytest.raises(ValueError):
            tskit.RateMap.read_hapmap(hapfile_header, has_header=False)
        rm1 = tskit.RateMap.read_hapmap(hapfile_header)
        rm2 = tskit.RateMap.read_hapmap(hapfile_noheader, has_header=False)
        assert_array_equal(rm1.rate, rm2.rate)
        assert_array_equal(rm1.position, rm2.position)

    def test_hapmap_fragment(self):
        hapfile = io.StringIO(
            """\
            chr pos        rate                    cM
            1   4283592    3.79115663174456        0
            1   4361401    0.0664276817058413      0.294986106359414
            1   7979763   10.9082897515584         0.535345505591925
            1   8007051    0.0976780648822495      0.833010916332456
            1   8762788    0.0899929572085616      0.906829844052373
            1   9477943    0.0864382908650907      0.971188757364862
            1   9696341    4.76495005895746        0.990066707213216
            1   9752154    0.0864316558730679      1.25601286485381
            1   9881751    0.0                     1.26721414815999"""
        )
        rm1 = tskit.RateMap.read_hapmap(hapfile)
        hapfile.seek(0)
        rm2 = tskit.RateMap.read_hapmap(hapfile, rate_col=2)
        assert np.allclose(rm1.position, rm2.position)
        assert np.allclose(rm1.rate, rm2.rate, equal_nan=True)


--- ../../tskit/python/tests/test_text_formats.py ---


"""
Test cases for converting fam file to tskit
"""
import dataclasses
import tempfile
from dataclasses import asdict

import numpy as np
import pytest

import tskit


@dataclasses.dataclass
class FamEntry:
    fid: str = "0"
    iid: str = "0"
    pat: str = "0"
    mat: str = "0"
    sex: str = "0"
    phen: str = None

    def get_row(self, delimiter="\t"):
        return delimiter.join([x for x in asdict(self).values() if x is not None])


class TestParseFam:
    """
    Tests for the parse_fam function.
    """

    def get_parsed_fam(self, entries, delimiter="\t"):
        content = "\n".join([entry.get_row(delimiter=delimiter) for entry in entries])
        with tempfile.TemporaryFile() as f:
            f.write(bytes(content, "utf-8"))
            f.seek(0)
            return tskit.parse_fam(f)

    def test_empty_file(self):
        entries = []
        with pytest.warns(UserWarning):
            tb = self.get_parsed_fam(entries=entries)
        assert len(tb) == 0

    @pytest.mark.parametrize("iid", ["1", "a", "100", "abc"])
    def test_single_line(self, iid):
        entries = [FamEntry(iid=iid)]
        tb = self.get_parsed_fam(entries=entries)
        assert len(tb) == 1
        assert np.array_equal(tb[0].parents, [-1, -1])
        assert tb[0].metadata["plink_fid"] == "0"
        assert tb[0].metadata["plink_iid"] == str(iid)
        assert tb[0].metadata["sex"] == 0

    @pytest.mark.parametrize("iids", [("1", "2"), ("a", "b")])
    def test_multiple_line_file(self, iids):
        # test both integer and string IIDs
        iid1, iid2 = iids
        entries = [FamEntry(iid=iid1), FamEntry(iid=iid2)]
        tb = self.get_parsed_fam(entries=entries)
        assert len(tb) == 2
        for idx in range(2):
            assert np.array_equal(tb[idx].parents, [-1, -1])
            assert tb[idx].metadata["plink_fid"] == "0"
            assert tb[idx].metadata["plink_iid"] == str(entries[idx].iid)
            assert tb[idx].metadata["sex"] == 0

    @pytest.mark.parametrize("n_cols", range(1, 5))
    def test_insufficient_cols(self, n_cols):
        fields = list(asdict(FamEntry()))
        entry = FamEntry(iid="1")
        for field in fields[n_cols:]:
            entry.__setattr__(field, None)
        with pytest.raises(  # noqa B017
            Exception
        ):  # Have to be non-specific here as numpy 1.23 changed the exception type
            self.get_parsed_fam(entries=[entry])

    def test_unrelated_duplicate_iids(self):
        # Individuals have the same IID, but are in different families
        entries = [FamEntry(iid="1"), FamEntry(fid="1", iid="1")]
        tb = self.get_parsed_fam(entries=entries)
        assert len(tb) == 2
        assert tb[0].metadata["plink_fid"] == "0"
        assert tb[0].metadata["plink_iid"] == "1"
        assert tb[1].metadata["plink_fid"] == "1"
        assert tb[1].metadata["plink_iid"] == "1"

    def test_duplicate_rows(self):
        entries = [FamEntry(iid="1"), FamEntry(iid="1")]
        with pytest.raises(ValueError):
            self.get_parsed_fam(entries=entries)

    def test_space_delimited(self):
        entries = [FamEntry(iid="1")]
        tb = self.get_parsed_fam(entries=entries, delimiter=" ")
        assert np.array_equal(tb[0].parents, [-1, -1])
        assert tb[0].metadata["plink_fid"] == "0"
        assert tb[0].metadata["plink_iid"] == "1"
        assert tb[0].metadata["sex"] == 0

    def test_missing_phen_col(self):
        entries = [FamEntry(iid="1", phen="1")]
        tb = self.get_parsed_fam(entries=entries)

        entries = [FamEntry(iid="1")]  # remove last column (PHEN column)
        tb_missing = self.get_parsed_fam(entries=entries)

        assert tb == tb_missing

    @pytest.mark.parametrize("sex", [-2, 3, "F"])
    def test_bad_sex_value(self, sex):
        entries = [FamEntry(iid="1", sex=str(sex))]
        with pytest.raises(ValueError):
            self.get_parsed_fam(entries=entries)

    def test_empty_sex_value(self):
        entries = [FamEntry(iid="1", sex="")]
        with pytest.raises(  # noqa B017
            Exception
        ):  # Have to be non-specific here as numpy 1.23 changed the exception type
            self.get_parsed_fam(entries=entries)

    def test_single_family_map_parent_ids(self):
        # PAT is mapped if the individual exists in the dataset
        entries = [FamEntry(iid="1"), FamEntry(iid="2", pat="1")]
        tb = self.get_parsed_fam(entries=entries)
        assert np.array_equal(tb[1].parents, [0, -1])

        # MAT is mapped if the individual exists in the dataset
        entries = [FamEntry(iid="1"), FamEntry(iid="2", mat="1")]
        tb = self.get_parsed_fam(entries=entries)
        assert np.array_equal(tb[1].parents, [-1, 0])

        # both parent IDs are remapped if the both parents exist in the dataset
        entries = [
            FamEntry(iid="1"),
            FamEntry(iid="2"),
            FamEntry(iid="3", pat="1", mat="2"),
        ]
        tb = self.get_parsed_fam(entries=entries)
        assert np.array_equal(tb[2].parents, [0, 1])

    def test_missing_parent_id(self):
        # KeyError raised if at least one parent (PAT) does not exist in dataset
        entries = [
            FamEntry(iid="2"),
            FamEntry(iid="3", pat="1", mat="2"),
        ]
        with pytest.raises(KeyError):
            self.get_parsed_fam(entries=entries)

        # KeyError raised if at least one parent (MAT) does not exist in dataset
        entries = [
            FamEntry(iid="1"),
            FamEntry(iid="3", pat="1", mat="2"),
        ]
        with pytest.raises(KeyError):
            self.get_parsed_fam(entries=entries)

        # KeyError raised if both parents do not exist in dataset
        entries = [FamEntry(iid="1", pat="2", mat="3")]
        with pytest.raises(KeyError):
            self.get_parsed_fam(entries=entries)

    def test_multiple_family_map_parent_ids(self):
        # parents mapped correctly when the same parent ID is used in different families
        entries = [
            FamEntry(iid="2"),
            FamEntry(iid="1"),
            FamEntry(fid="1", iid="2"),
            FamEntry(fid="1", iid="1"),
            FamEntry(iid="3", pat="1", mat="2"),
            FamEntry(fid="1", iid="3", pat="1", mat="2"),
        ]
        tb = self.get_parsed_fam(entries=entries)
        for idx in range(4):
            assert np.array_equal(tb[idx].parents, [-1, -1])
        assert np.array_equal(tb[4].parents, [1, 0])
        assert np.array_equal(tb[5].parents, [3, 2])

        # KeyError raised when FID does not match, even if parent ID matches
        entries = [
            FamEntry(iid="2"),
            FamEntry(iid="1"),
            FamEntry(iid="3", pat="1", mat="2"),
            FamEntry(
                fid="1", iid="1", pat="2", mat="3"
            ),  # there is no parent with FID=1, IID=3
            FamEntry(fid="1", iid="2"),
        ]
        with pytest.raises(KeyError):
            self.get_parsed_fam(entries)

    def test_grandparents(self):
        entries = [
            FamEntry(iid="4"),
            FamEntry(iid="3"),
            FamEntry(iid="2"),
            FamEntry(iid="1"),
            FamEntry(iid="6", pat="3", mat="4"),
            FamEntry(iid="5", pat="1", mat="2"),
            FamEntry(iid="7", pat="5", mat="6"),
        ]
        tb = self.get_parsed_fam(entries=entries)
        assert np.array_equal(tb[4].parents, [1, 0])
        assert np.array_equal(tb[5].parents, [3, 2])
        assert np.array_equal(tb[6].parents, [5, 4])

    def test_children_before_parents(self, tmp_path):
        entries = [
            FamEntry(iid="1", pat="2", mat="3"),
            FamEntry(iid="2"),
            FamEntry(iid="3"),
        ]
        content = "\n".join([entry.get_row() for entry in entries])
        fam_path = f"{tmp_path}/test.fam"
        with open(fam_path, "w+") as f:
            f.write(content)
            f.seek(0)
            tb = tskit.parse_fam(f)

        tc = tskit.TableCollection(1)
        # Issue 1489 will make this better
        tc.individuals.metadata_schema = tb.metadata_schema
        for row in tb:
            tc.individuals.append(row)
        tc.tree_sequence()  # creating tree sequence should succeed


--- ../../tskit/python/tests/test_vcf.py ---


"""
Test cases for VCF output in tskit.
"""
import contextlib
import io
import math
import os
import tempfile
import textwrap

import msprime
import numpy as np
import pytest

import tests
import tests.test_wright_fisher as wf
import tskit
from tests import tsutil

# Pysam is not available on windows, so we don't make it mandatory here.
_pysam_imported = False
try:
    import pysam

    _pysam_imported = True
except ImportError:
    pass


@contextlib.contextmanager
def ts_to_pysam(ts, *args, **kwargs):
    """
    Returns a pysam VariantFile for the specified tree sequence and arguments.
    """
    with tempfile.TemporaryDirectory() as temp_dir:
        vcf_path = os.path.join(temp_dir, "file.vcf")
        with open(vcf_path, "w") as f:
            ts.write_vcf(f, *args, **kwargs, allow_position_zero=True)
        yield pysam.VariantFile(vcf_path)


def example_individuals(ts, ploidy=1):
    if ts.num_individuals == 0:
        yield None, ts.num_samples / ploidy
    else:
        yield None, ts.num_individuals
        yield list(range(ts.num_individuals)), ts.num_individuals
    if ts.num_individuals > 3:
        n = ts.num_individuals - 2
        yield list(range(n)), n
        yield 2 + np.random.choice(np.arange(n), n, replace=False), n


def legacy_write_vcf(tree_sequence, output, ploidy, contig_id):
    """
    Writes a VCF under the legacy conversion rules used in versions before 0.2.0.
    """
    if tree_sequence.get_sample_size() % ploidy != 0:
        raise ValueError("Sample size must a multiple of ploidy")
    n = tree_sequence.get_sample_size() // ploidy
    sample_names = [f"msp_{j}" for j in range(n)]
    last_pos = 0
    positions = []
    for variant in tree_sequence.variants():
        pos = int(round(variant.position))
        if pos <= last_pos:
            pos = last_pos + 1
        positions.append(pos)
        last_pos = pos
    contig_length = int(math.ceil(tree_sequence.get_sequence_length()))
    if len(positions) > 0:
        contig_length = max(positions[-1], contig_length)
    print("##fileformat=VCFv4.2", file=output)
    print(f"##source=tskit {tskit.__version__}", file=output)
    print('##FILTER=<ID=PASS,Description="All filters passed">', file=output)
    print(f"##contig=<ID={contig_id},length={contig_length}>", file=output)
    print('##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">', file=output)
    print(
        "#CHROM",
        "POS",
        "ID",
        "REF",
        "ALT",
        "QUAL",
        "FILTER",
        "INFO",
        "FORMAT",
        sep="\t",
        end="",
        file=output,
    )
    for sample_name in sample_names:
        print("\t", sample_name, sep="", end="", file=output)
    print(file=output)
    for variant in tree_sequence.variants():
        pos = positions[variant.index]
        site_id = variant.site.id
        assert variant.num_alleles == 2
        print(
            contig_id,
            pos,
            site_id,
            variant.alleles[0],
            variant.alleles[1],
            ".",
            "PASS",
            ".",
            "GT",
            sep="\t",
            end="",
            file=output,
        )
        for j in range(n):
            genotype = "|".join(
                str(g) for g in variant.genotypes[j * ploidy : j * ploidy + ploidy]
            )
            print("\t", genotype, end="", sep="", file=output)
        print(file=output)


class TestLegacyOutput:
    """
    Tests if the VCF file produced by the low level code is the
    same as one we generate here.
    """

    def verify(self, ts, ploidy=1, contig_id="1"):
        assert ts.num_sites > 0
        f = io.StringIO()
        legacy_write_vcf(ts, f, ploidy=ploidy, contig_id=contig_id)
        vcf1 = f.getvalue()

        num_individuals = ts.num_samples // ploidy
        individual_names = [f"msp_{j}" for j in range(num_individuals)]
        f = io.StringIO()
        ts.write_vcf(
            f,
            ploidy=ploidy,
            contig_id=contig_id,
            position_transform="legacy",
            individual_names=individual_names,
        )
        vcf2 = f.getvalue()
        assert vcf1 == vcf2

    def test_msprime_length_1(self):
        ts = msprime.simulate(10, mutation_rate=1, random_seed=666)
        self.verify(ts, ploidy=1)
        self.verify(ts, ploidy=2)
        self.verify(ts, ploidy=5)

    def test_msprime_length_10(self):
        ts = msprime.simulate(9, length=10, mutation_rate=0.1, random_seed=666)
        self.verify(ts, ploidy=1)
        self.verify(ts, ploidy=3)

    def test_contig_id(self):
        ts = msprime.simulate(10, mutation_rate=1, random_seed=666)
        self.verify(ts, ploidy=1, contig_id="X")
        self.verify(ts, ploidy=2, contig_id="X" * 10)


class ExamplesMixin:
    """
    Mixin defining tests on various example tree sequences.
    """

    def test_simple_infinite_sites_random_ploidy(self):
        ts = msprime.simulate(10, mutation_rate=1, random_seed=2)
        ts = tsutil.insert_random_ploidy_individuals(
            ts, min_ploidy=1, samples_only=True
        )
        assert ts.num_sites > 2
        self.verify(ts)

    def test_simple_infinite_sites_ploidy_2(self):
        ts = msprime.simulate(10, mutation_rate=1, random_seed=2)
        ts = tsutil.insert_individuals(ts, ploidy=2)
        assert ts.num_sites > 2
        self.verify(ts)

    def test_simple_infinite_sites_ploidy_2_reversed_samples(self):
        ts = msprime.simulate(10, mutation_rate=1, random_seed=2)
        samples = ts.samples()[::-1]
        ts = tsutil.insert_individuals(ts, nodes=samples, ploidy=2)
        assert ts.num_sites > 2
        self.verify(ts)

    def test_simple_jukes_cantor_random_ploidy(self):
        ts = msprime.simulate(10, random_seed=2)
        ts = tsutil.jukes_cantor(ts, num_sites=10, mu=1, seed=2)
        ts = tsutil.insert_random_ploidy_individuals(
            ts, min_ploidy=1, samples_only=True
        )
        self.verify(ts)

    def test_single_tree_multichar_mutations(self):
        ts = msprime.simulate(6, random_seed=1, mutation_rate=1)
        ts = tsutil.insert_multichar_mutations(ts)
        ts = tsutil.insert_individuals(ts, ploidy=2)
        self.verify(ts)

    def test_many_trees_infinite_sites(self):
        ts = msprime.simulate(6, recombination_rate=2, mutation_rate=2, random_seed=1)
        assert ts.num_sites > 0
        assert ts.num_trees > 2
        ts = tsutil.insert_individuals(ts, ploidy=2)
        self.verify(ts)

    def test_many_trees_sequence_length_infinite_sites(self):
        for L in [0.5, 1.5, 3.3333]:
            ts = msprime.simulate(
                6, length=L, recombination_rate=2, mutation_rate=1, random_seed=1
            )
            assert ts.num_sites > 0
            ts = tsutil.insert_individuals(ts, ploidy=2)
            self.verify(ts)

    def test_wright_fisher_unsimplified(self):
        tables = wf.wf_sim(
            4,
            5,
            seed=1,
            deep_history=True,
            initial_generation_samples=False,
            num_loci=10,
        )
        tables.sort()
        ts = msprime.mutate(tables.tree_sequence(), rate=0.05, random_seed=234)
        assert ts.num_sites > 0
        ts = tsutil.insert_individuals(ts, ploidy=4)
        self.verify(ts)

    def test_wright_fisher_initial_generation(self):
        tables = wf.wf_sim(
            6, 5, seed=3, deep_history=True, initial_generation_samples=True, num_loci=2
        )
        tables.sort()
        tables.simplify()
        ts = msprime.mutate(tables.tree_sequence(), rate=0.08, random_seed=2)
        assert ts.num_sites > 0
        ts = tsutil.insert_individuals(ts, ploidy=3)
        self.verify(ts)

    def test_wright_fisher_unsimplified_multiple_roots(self):
        tables = wf.wf_sim(
            8,
            15,
            seed=1,
            deep_history=False,
            initial_generation_samples=False,
            num_loci=20,
        )
        tables.sort()
        ts = msprime.mutate(tables.tree_sequence(), rate=0.006, random_seed=2)
        assert ts.num_sites > 0
        ts = tsutil.insert_individuals(ts, ploidy=2)
        self.verify(ts)

    def test_wright_fisher_simplified(self):
        tables = wf.wf_sim(
            9,
            10,
            seed=1,
            deep_history=True,
            initial_generation_samples=False,
            num_loci=5,
        )
        tables.sort()
        ts = tables.tree_sequence().simplify()
        ts = msprime.mutate(ts, rate=0.2, random_seed=1234)
        assert ts.num_sites > 0
        ts = tsutil.insert_individuals(ts, ploidy=3)
        self.verify(ts)


@pytest.mark.skipif(not _pysam_imported, reason="pysam not available")
class TestParseHeaderPysam(ExamplesMixin):
    """
    Test that pysam can parse the headers correctly.
    """

    def verify(self, ts):
        contig_id = "pysam"
        for indivs, num_indivs in example_individuals(ts):
            with ts_to_pysam(ts, contig_id=contig_id, individuals=indivs) as bcf_file:
                assert bcf_file.format == "VCF"
                assert bcf_file.version == (4, 2)
                header = bcf_file.header
                assert len(header.contigs) == 1
                contig = header.contigs[0]
                assert contig.name == contig_id
                assert contig.length > 0
                assert len(header.filters) == 1
                p = header.filters["PASS"]
                assert p.name == "PASS"
                assert p.description == "All filters passed"
                assert len(header.info) == 0
                assert len(header.formats) == 1
                fmt = header.formats["GT"]
                assert fmt.name == "GT"
                assert fmt.number == 1
                assert fmt.type == "String"
                assert fmt.description == "Genotype"
                assert len(bcf_file.header.samples) == num_indivs


class TestInterface:
    """
    Tests for the interface.
    """

    def test_bad_ploidy(self):
        ts = msprime.simulate(10, mutation_rate=0.1, random_seed=2)
        for bad_ploidy in [-1, 0]:
            with pytest.raises(ValueError, match="Ploidy must be >= 1"):
                ts.write_vcf(io.StringIO, bad_ploidy)
        # Non divisible
        for bad_ploidy in [3, 7]:
            with pytest.raises(
                ValueError, match="Sample size must be divisible by ploidy"
            ):
                ts.write_vcf(io.StringIO, bad_ploidy)

    def test_individuals_no_nodes_default_args(self):
        ts1 = msprime.simulate(10, mutation_rate=0.1, random_seed=2)
        tables = ts1.dump_tables()
        tables.individuals.add_row()
        ts2 = tables.tree_sequence()
        assert ts1.as_vcf(allow_position_zero=True) == ts2.as_vcf(
            allow_position_zero=True
        )

    def test_individuals_no_nodes_as_argument(self):
        ts1 = msprime.simulate(10, mutation_rate=0.1, random_seed=2)
        tables = ts1.dump_tables()
        tables.individuals.add_row()
        ts2 = tables.tree_sequence()
        with pytest.raises(ValueError, match="0 not associated with a node"):
            ts2.as_vcf(individuals=[0])

    def test_ploidy_with_sample_individuals(self):
        ts = msprime.sim_ancestry(3, random_seed=2)
        ts = tsutil.insert_branch_sites(ts)
        with pytest.raises(ValueError, match="Cannot specify ploidy when individuals"):
            ts.write_vcf(io.StringIO(), ploidy=2)

    def test_ploidy_with_no_node_individuals(self):
        ts1 = msprime.simulate(10, mutation_rate=0.1, random_seed=2)
        tables = ts1.dump_tables()
        tables.individuals.add_row()
        ts2 = tables.tree_sequence()
        with pytest.raises(ValueError, match="Cannot specify ploidy when individuals"):
            ts2.as_vcf(ploidy=2)

    def test_empty_individuals(self):
        ts = msprime.sim_ancestry(3, random_seed=2)
        ts = tsutil.insert_branch_sites(ts)
        with pytest.raises(ValueError, match="List of sample individuals empty"):
            ts.as_vcf(individuals=[])

    def test_duplicate_individuals(self):
        ts = msprime.sim_ancestry(3, random_seed=2)
        ts = tsutil.insert_branch_sites(ts)
        with pytest.raises(tskit.LibraryError, match="TSK_ERR_DUPLICATE_SAMPLE"):
            ts.as_vcf(individuals=[0, 0], allow_position_zero=True)

    def test_mixed_sample_non_sample_individuals(self):
        ts = msprime.sim_ancestry(3, random_seed=2)
        tables = ts.dump_tables()
        tables.individuals.add_row()
        # Add a reference to an individual from a non-sample
        individual = tables.nodes.individual
        individual[-1] = 0
        tables.nodes.individual = individual
        ts = tables.tree_sequence()
        ts = tsutil.insert_branch_sites(ts)
        with pytest.raises(
            ValueError, match="0 has nodes that are sample and non-sample"
        ):
            ts.as_vcf()
        # but it's OK if we run without the affected individual
        assert len(ts.as_vcf(individuals=[1, 2], allow_position_zero=True)) > 0

    def test_samples_with_and_without_individuals(self):
        ts = tskit.Tree.generate_balanced(3).tree_sequence
        tables = ts.dump_tables()
        tables.individuals.add_row()
        # Add a reference to an individual from one sample
        individual = tables.nodes.individual
        individual[0] = 0
        tables.nodes.individual = individual
        ts = tables.tree_sequence()
        ts = tsutil.insert_branch_sites(ts)
        with pytest.raises(
            ValueError, match="Sample nodes must either all be associated"
        ):
            ts.as_vcf()
        # But it's OK if explicitly specify that sample
        assert len(ts.as_vcf(individuals=[0], allow_position_zero=True)) > 0

    def test_bad_individuals(self):
        ts = msprime.simulate(10, mutation_rate=0.1, random_seed=2)
        ts = tsutil.insert_individuals(ts, ploidy=2)
        with pytest.raises(ValueError, match="Invalid individual IDs provided."):
            ts.write_vcf(io.StringIO(), individuals=[0, -1])
        with pytest.raises(ValueError, match="Invalid individual IDs provided."):
            ts.write_vcf(io.StringIO(), individuals=[1, 2, ts.num_individuals])

    def test_ploidy_positional(self):
        ts = msprime.simulate(2, mutation_rate=2, random_seed=1)
        assert ts.as_vcf(2, allow_position_zero=True) == ts.as_vcf(
            ploidy=2, allow_position_zero=True
        )

    def test_only_ploidy_positional(self):
        ts = msprime.simulate(2, mutation_rate=2, random_seed=1)
        with pytest.raises(TypeError, match="positional arguments"):
            assert ts.as_vcf(2, "chr2")


class TestLimitations:
    """
    Verify the correct error behaviour in cases we don't support.
    """

    def test_many_alleles(self):
        ts = msprime.simulate(20, random_seed=45)
        tables = ts.dump_tables()
        tables.sites.add_row(0.5, "0")
        # 9 alleles should be fine
        for j in range(8):
            tables.mutations.add_row(0, node=j, derived_state=str(j + 1))
        ts = tables.tree_sequence()
        ts.write_vcf(io.StringIO(), allow_position_zero=True)
        for j in range(9, 15):
            tables.mutations.add_row(0, node=j, derived_state=str(j))
            ts = tables.tree_sequence()
            with pytest.raises(
                ValueError, match="More than 9 alleles not currently supported"
            ):
                ts.write_vcf(io.StringIO(), allow_position_zero=True)


class TestPositionTransformErrors:
    """
    Tests what happens when we provide bad position transforms
    """

    def get_example_ts(self):
        ts = msprime.simulate(11, mutation_rate=1, random_seed=11)
        assert ts.num_sites > 1
        return ts

    def test_wrong_output_dimensions(self):
        ts = self.get_example_ts()
        for bad_func in [np.sum, lambda x: []]:
            with pytest.raises(ValueError):
                ts.write_vcf(io.StringIO(), position_transform=bad_func)

    def test_bad_func(self):
        ts = self.get_example_ts()
        for bad_func in ["", Exception]:
            with pytest.raises(TypeError):
                ts.write_vcf(io.StringIO(), position_transform=bad_func)


class TestZeroPositionErrors:
    """
    Tests for handling zero position sites
    """

    def test_zero_position_error(self):
        ts = msprime.sim_ancestry(3, random_seed=2, sequence_length=10)
        ts = msprime.sim_mutations(ts, rate=1, random_seed=2)
        assert ts.sites_position[0] == 0

        with pytest.raises(ValueError, match="A variant position of 0"):
            ts.write_vcf(io.StringIO())

        # Should succeed if we allow it, or the site is masked or transformed
        ts.write_vcf(io.StringIO(), allow_position_zero=True)
        ts.write_vcf(io.StringIO(), position_transform=lambda pos: [x + 1 for x in pos])
        mask = np.zeros(ts.num_sites, dtype=bool)
        mask[0] = True
        ts.write_vcf(io.StringIO(), site_mask=mask)

    def test_no_position_zero_ok(self):
        ts = msprime.sim_ancestry(3, random_seed=2, sequence_length=10)
        ts = msprime.sim_mutations(ts, rate=0.25, random_seed=4)
        assert ts.num_sites > 0
        assert ts.sites_position[0] != 0
        ts.write_vcf(io.StringIO(), allow_position_zero=True)
        ts.write_vcf(io.StringIO())


class TestIndividualNames:
    """
    Tests for the individual names argument.
    """

    def test_bad_length_individuals(self):
        ts = msprime.simulate(6, mutation_rate=2, random_seed=1)
        assert ts.num_sites > 0
        ts = tsutil.insert_individuals(ts, ploidy=2)
        with pytest.raises(
            ValueError,
            match="individual_names must have length equal to"
            " the number of individuals",
        ):
            ts.write_vcf(io.StringIO(), individual_names=[])
        with pytest.raises(
            ValueError,
            match="individual_names must have length equal to"
            " the number of individuals",
        ):
            ts.write_vcf(io.StringIO(), individual_names=["x" for _ in range(4)])
        with pytest.raises(
            ValueError,
            match="individual_names must have length equal to"
            " the number of individuals",
        ):
            ts.write_vcf(
                io.StringIO(),
                individuals=list(range(ts.num_individuals)),
                individual_names=["x" for _ in range(ts.num_individuals - 1)],
            )
        with pytest.raises(
            ValueError,
            match="individual_names must have length equal to"
            " the number of individuals",
        ):
            ts.write_vcf(
                io.StringIO(),
                individuals=list(range(ts.num_individuals - 1)),
                individual_names=["x" for _ in range(ts.num_individuals)],
            )

    def test_bad_length_ploidy(self):
        ts = msprime.simulate(6, mutation_rate=2, random_seed=1)
        assert ts.num_sites > 0
        with pytest.raises(
            ValueError,
            match="individual_names must have length equal to"
            " the number of individuals",
        ):
            ts.write_vcf(io.StringIO(), ploidy=2, individual_names=[])
        with pytest.raises(
            ValueError,
            match="individual_names must have length equal to"
            " the number of individuals",
        ):
            ts.write_vcf(
                io.StringIO(), ploidy=2, individual_names=["x" for _ in range(4)]
            )

    def test_bad_type(self):
        ts = msprime.simulate(2, mutation_rate=2, random_seed=1)
        with pytest.raises(
            TypeError, match="sequence item 0: expected str instance," " NoneType found"
        ):
            ts.write_vcf(
                io.StringIO(), individual_names=[None, "b"], allow_position_zero=True
            )
        with pytest.raises(
            TypeError, match="sequence item 0: expected str instance," " bytes found"
        ):
            ts.write_vcf(
                io.StringIO(), individual_names=[b"a", "b"], allow_position_zero=True
            )


def drop_header(s):
    return "\n".join(line for line in s.splitlines() if not line.startswith("##"))


class TestMasking:
    @tests.cached_example
    def ts(self):
        ts = tskit.Tree.generate_balanced(3, span=10).tree_sequence
        ts = tsutil.insert_branch_sites(ts)
        return ts

    @pytest.mark.parametrize("mask", [[True], np.zeros(5, dtype=bool), []])
    def test_site_mask_wrong_size(self, mask):
        with pytest.raises(ValueError, match="Site mask must be"):
            self.ts().as_vcf(site_mask=mask)

    @pytest.mark.parametrize("mask", [[[0, 1], [1, 0]], "abcd"])
    def test_site_mask_bad_type(self, mask):
        # converting to a bool array is pretty lax in what's allows.
        with pytest.raises(ValueError, match="Site mask must be"):
            self.ts().as_vcf(site_mask=mask)

    @pytest.mark.parametrize("mask", [[[0, 1], [1, 0]], "abcd"])
    def test_sample_mask_bad_type(self, mask):
        # converting to a bool array is pretty lax in what's allows.
        with pytest.raises(ValueError, match="Sample mask must be"):
            self.ts().as_vcf(sample_mask=mask, allow_position_zero=True)

    def test_no_masks(self):
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0\ttsk_1\ttsk_2
        1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t1\t0\t0
        1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t0\t1\t1
        1\t4\t2\t0\t1\t.\tPASS\t.\tGT\t0\t1\t0
        1\t6\t3\t0\t1\t.\tPASS\t.\tGT\t0\t0\t1"""
        expected = textwrap.dedent(s)
        assert drop_header(self.ts().as_vcf(allow_position_zero=True)) == expected

    def test_no_masks_triploid(self):
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0
        1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t1|0|0
        1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t0|1|1
        1\t4\t2\t0\t1\t.\tPASS\t.\tGT\t0|1|0
        1\t6\t3\t0\t1\t.\tPASS\t.\tGT\t0|0|1"""
        expected = textwrap.dedent(s)
        assert (
            drop_header(self.ts().as_vcf(ploidy=3, allow_position_zero=True))
            == expected
        )

    def test_site_0_masked(self):
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0\ttsk_1\ttsk_2
        1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t0\t1\t1
        1\t4\t2\t0\t1\t.\tPASS\t.\tGT\t0\t1\t0
        1\t6\t3\t0\t1\t.\tPASS\t.\tGT\t0\t0\t1"""
        expected = textwrap.dedent(s)
        actual = self.ts().as_vcf(
            site_mask=[True, False, False, False], allow_position_zero=True
        )
        assert drop_header(actual) == expected

    def test_site_0_masked_triploid(self):
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0
        1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t0|1|1
        1\t4\t2\t0\t1\t.\tPASS\t.\tGT\t0|1|0
        1\t6\t3\t0\t1\t.\tPASS\t.\tGT\t0|0|1"""
        expected = textwrap.dedent(s)
        actual = self.ts().as_vcf(
            ploidy=3, site_mask=[True, False, False, False], allow_position_zero=True
        )
        assert drop_header(actual) == expected

    def test_site_1_masked(self):
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0\ttsk_1\ttsk_2
        1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t1\t0\t0
        1\t4\t2\t0\t1\t.\tPASS\t.\tGT\t0\t1\t0
        1\t6\t3\t0\t1\t.\tPASS\t.\tGT\t0\t0\t1"""
        expected = textwrap.dedent(s)
        actual = self.ts().as_vcf(
            site_mask=[False, True, False, False], allow_position_zero=True
        )
        assert drop_header(actual) == expected

    def test_all_sites_masked(self):
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0\ttsk_1\ttsk_2"""
        expected = textwrap.dedent(s)
        actual = self.ts().as_vcf(
            site_mask=[True, True, True, True], allow_position_zero=True
        )
        assert drop_header(actual) == expected

    def test_all_sites_not_masked(self):
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0\ttsk_1\ttsk_2
        1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t1\t0\t0
        1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t0\t1\t1
        1\t4\t2\t0\t1\t.\tPASS\t.\tGT\t0\t1\t0
        1\t6\t3\t0\t1\t.\tPASS\t.\tGT\t0\t0\t1"""
        expected = textwrap.dedent(s)
        actual = self.ts().as_vcf(
            site_mask=[False, False, False, False], allow_position_zero=True
        )
        assert drop_header(actual) == expected

    @pytest.mark.parametrize(
        "mask",
        [[False, False, False], [0, 0, 0], lambda _: [False, False, False]],
    )
    def test_all_samples_not_masked(self, mask):
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0\ttsk_1\ttsk_2
        1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t1\t0\t0
        1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t0\t1\t1
        1\t4\t2\t0\t1\t.\tPASS\t.\tGT\t0\t1\t0
        1\t6\t3\t0\t1\t.\tPASS\t.\tGT\t0\t0\t1"""
        expected = textwrap.dedent(s)
        actual = self.ts().as_vcf(sample_mask=mask, allow_position_zero=True)
        assert drop_header(actual) == expected

    @pytest.mark.parametrize(
        "mask", [[True, False, False], [1, 0, 0], lambda _: [True, False, False]]
    )
    def test_sample_0_masked(self, mask):
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0\ttsk_1\ttsk_2
        1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t.\t0\t0
        1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t.\t1\t1
        1\t4\t2\t0\t1\t.\tPASS\t.\tGT\t.\t1\t0
        1\t6\t3\t0\t1\t.\tPASS\t.\tGT\t.\t0\t1"""
        expected = textwrap.dedent(s)
        actual = self.ts().as_vcf(sample_mask=mask, allow_position_zero=True)
        assert drop_header(actual) == expected

    @pytest.mark.parametrize(
        "mask", [[False, True, False], [0, 1, 0], lambda _: [False, True, False]]
    )
    def test_sample_1_masked(self, mask):
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0\ttsk_1\ttsk_2
        1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t1\t.\t0
        1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t0\t.\t1
        1\t4\t2\t0\t1\t.\tPASS\t.\tGT\t0\t.\t0
        1\t6\t3\t0\t1\t.\tPASS\t.\tGT\t0\t.\t1"""
        expected = textwrap.dedent(s)
        actual = self.ts().as_vcf(sample_mask=mask, allow_position_zero=True)
        assert drop_header(actual) == expected

    @pytest.mark.parametrize(
        "mask", [[True, True, True], [1, 1, 1], lambda _: [True, True, True]]
    )
    def test_all_samples_masked(self, mask):
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0\ttsk_1\ttsk_2
        1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t.\t.\t.
        1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t.\t.\t.
        1\t4\t2\t0\t1\t.\tPASS\t.\tGT\t.\t.\t.
        1\t6\t3\t0\t1\t.\tPASS\t.\tGT\t.\t.\t."""
        expected = textwrap.dedent(s)
        actual = self.ts().as_vcf(sample_mask=mask, allow_position_zero=True)
        assert drop_header(actual) == expected

    def test_all_functional_sample_mask(self):
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0\ttsk_1\ttsk_2
        1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t.\t0\t0
        1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t0\t.\t1
        1\t4\t2\t0\t1\t.\tPASS\t.\tGT\t0\t1\t.
        1\t6\t3\t0\t1\t.\tPASS\t.\tGT\t.\t0\t1"""

        def mask(variant):
            a = [0, 0, 0]
            a[variant.site.id % 3] = 1
            return a

        expected = textwrap.dedent(s)
        actual = self.ts().as_vcf(sample_mask=mask, allow_position_zero=True)
        assert drop_header(actual) == expected

    @pytest.mark.skipif(not _pysam_imported, reason="pysam not available")
    def test_mask_ok_with_pysam(self):
        with ts_to_pysam(self.ts(), sample_mask=[0, 0, 1]) as records:
            variants = list(records)
            assert len(variants) == 4
            samples = ["tsk_0", "tsk_1", "tsk_2"]
            gts = [variants[0].samples[key]["GT"] for key in samples]
            assert gts == [(1,), (0,), (None,)]

            gts = [variants[1].samples[key]["GT"] for key in samples]
            assert gts == [(0,), (1,), (None,)]

            gts = [variants[2].samples[key]["GT"] for key in samples]
            assert gts == [(0,), (1,), (None,)]

            gts = [variants[3].samples[key]["GT"] for key in samples]
            assert gts == [(0,), (0,), (None,)]


class TestMissingData:
    @tests.cached_example
    def ts(self):
        tables = tskit.Tree.generate_balanced(2, span=10).tree_sequence.dump_tables()
        tables.nodes.add_row(time=0, flags=tskit.NODE_IS_SAMPLE)
        ts = tsutil.insert_branch_sites(tables.tree_sequence())
        return ts

    def test_defaults(self):
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0\ttsk_1\ttsk_2
        1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t1\t0\t.
        1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t0\t1\t."""
        expected = textwrap.dedent(s)
        assert drop_header(self.ts().as_vcf(allow_position_zero=True)) == expected

    def test_isolated_as_missing_true(self):
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0\ttsk_1\ttsk_2
        1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t1\t0\t.
        1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t0\t1\t."""
        expected = textwrap.dedent(s)
        assert (
            drop_header(
                self.ts().as_vcf(isolated_as_missing=True, allow_position_zero=True)
            )
            == expected
        )

    def test_isolated_as_missing_false(self):
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0\ttsk_1\ttsk_2
        1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t1\t0\t0
        1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t0\t1\t0"""
        expected = textwrap.dedent(s)
        assert (
            drop_header(
                self.ts().as_vcf(isolated_as_missing=False, allow_position_zero=True)
            )
            == expected
        )

    @pytest.mark.skipif(not _pysam_imported, reason="pysam not available")
    def test_ok_with_pysam(self):
        with ts_to_pysam(self.ts(), sample_mask=[0, 0, 1]) as records:
            variants = list(records)
            assert len(variants) == 2
            samples = ["tsk_0", "tsk_1", "tsk_2"]
            gts = [variants[0].samples[key]["GT"] for key in samples]
            assert gts == [(1,), (0,), (None,)]

            gts = [variants[1].samples[key]["GT"] for key in samples]
            assert gts == [(0,), (1,), (None,)]


def drop_individuals(ts):
    tables = ts.dump_tables()
    individual = tables.nodes.individual
    individual[:] = -1
    tables.individuals.clear()
    tables.nodes.individual = individual
    return tables.tree_sequence()


class TestSampleOptions:
    @tests.cached_example
    def ts(self):
        ts = tskit.Tree.generate_balanced(3, span=10).tree_sequence
        ts = tsutil.insert_branch_sites(ts)
        tables = ts.dump_tables()
        tables.individuals.add_row()
        tables.individuals.add_row()
        individual = tables.nodes.individual
        # One diploid and one haploid, not in adjacent individuals
        individual[0] = 0
        individual[1] = 1
        individual[2] = 0
        tables.nodes.individual = individual
        return tables.tree_sequence()

    def test_no_individuals_defaults(self):
        ts = drop_individuals(self.ts())
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0\ttsk_1\ttsk_2
        1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t1\t0\t0
        1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t0\t1\t1
        1\t4\t2\t0\t1\t.\tPASS\t.\tGT\t0\t1\t0
        1\t6\t3\t0\t1\t.\tPASS\t.\tGT\t0\t0\t1"""
        expected = textwrap.dedent(s)
        assert drop_header(ts.as_vcf(allow_position_zero=True)) == expected

    def test_no_individuals_ploidy_3(self):
        ts = drop_individuals(self.ts())
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0
        1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t1|0|0
        1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t0|1|1
        1\t4\t2\t0\t1\t.\tPASS\t.\tGT\t0|1|0
        1\t6\t3\t0\t1\t.\tPASS\t.\tGT\t0|0|1"""
        expected = textwrap.dedent(s)
        assert drop_header(ts.as_vcf(ploidy=3, allow_position_zero=True)) == expected

    def test_no_individuals_ploidy_3_names(self):
        ts = drop_individuals(self.ts())
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tA
        1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t1|0|0
        1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t0|1|1
        1\t4\t2\t0\t1\t.\tPASS\t.\tGT\t0|1|0
        1\t6\t3\t0\t1\t.\tPASS\t.\tGT\t0|0|1"""
        expected = textwrap.dedent(s)
        assert (
            drop_header(
                ts.as_vcf(ploidy=3, individual_names="A", allow_position_zero=True)
            )
            == expected
        )

    def test_defaults(self):
        ts = self.ts()
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0\ttsk_1
        1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t1|0\t0
        1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t0|1\t1
        1\t4\t2\t0\t1\t.\tPASS\t.\tGT\t0|0\t1
        1\t6\t3\t0\t1\t.\tPASS\t.\tGT\t0|1\t0"""
        expected = textwrap.dedent(s)
        assert drop_header(ts.as_vcf(allow_position_zero=True)) == expected

    def test_individual_0(self):
        ts = self.ts()
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0
        1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t1|0
        1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t0|1
        1\t4\t2\t0\t1\t.\tPASS\t.\tGT\t0|0
        1\t6\t3\t0\t1\t.\tPASS\t.\tGT\t0|1"""
        expected = textwrap.dedent(s)
        assert (
            drop_header(ts.as_vcf(individuals=[0], allow_position_zero=True))
            == expected
        )

    def test_individual_1(self):
        ts = self.ts()
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0
        1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t0
        1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t1
        1\t4\t2\t0\t1\t.\tPASS\t.\tGT\t1
        1\t6\t3\t0\t1\t.\tPASS\t.\tGT\t0"""
        expected = textwrap.dedent(s)
        assert (
            drop_header(ts.as_vcf(individuals=[1], allow_position_zero=True))
            == expected
        )

    def test_reversed(self):
        ts = self.ts()
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ttsk_0\ttsk_1
        1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t0\t1|0
        1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t1\t0|1
        1\t4\t2\t0\t1\t.\tPASS\t.\tGT\t1\t0|0
        1\t6\t3\t0\t1\t.\tPASS\t.\tGT\t0\t0|1"""
        expected = textwrap.dedent(s)
        assert (
            drop_header(ts.as_vcf(individuals=[1, 0], allow_position_zero=True))
            == expected
        )

    def test_reversed_names(self):
        ts = self.ts()
        s = """\
        #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tA\tB
        1\t0\t0\t0\t1\t.\tPASS\t.\tGT\t0\t1|0
        1\t2\t1\t0\t1\t.\tPASS\t.\tGT\t1\t0|1
        1\t4\t2\t0\t1\t.\tPASS\t.\tGT\t1\t0|0
        1\t6\t3\t0\t1\t.\tPASS\t.\tGT\t0\t0|1"""
        expected = textwrap.dedent(s)
        assert (
            drop_header(
                ts.as_vcf(
                    individuals=[1, 0],
                    individual_names=["A", "B"],
                    allow_position_zero=True,
                ),
            )
            == expected
        )


--- ../../tskit/python/tests/test_provenance.py ---


"""
Tests for the provenance information attached to tree sequences.
"""
import json
import os
import platform
import sys
import time

try:
    import resource
except ImportError:
    resource = None  # resource absent on windows


import msprime
import pytest

import _tskit
import tskit
import tskit.provenance as provenance


def get_provenance(
    software_name="x",
    software_version="y",
    schema_version="1",
    environment=None,
    parameters=None,
):
    """
    Utility function to return a provenance document for testing.
    """
    document = {
        "schema_version": schema_version,
        "software": {"name": software_name, "version": software_version},
        "environment": {} if environment is None else environment,
        "parameters": {} if parameters is None else parameters,
    }
    return document


class TestSchema:
    """
    Tests for schema validation.
    """

    def test_empty(self):
        with pytest.raises(tskit.ProvenanceValidationError):
            tskit.validate_provenance({})

    def test_missing_keys(self):
        minimal = get_provenance()
        tskit.validate_provenance(minimal)
        for key in minimal.keys():
            copy = dict(minimal)
            del copy[key]
            with pytest.raises(tskit.ProvenanceValidationError):
                tskit.validate_provenance(copy)
        copy = dict(minimal)
        del copy["software"]["name"]
        with pytest.raises(tskit.ProvenanceValidationError):
            tskit.validate_provenance(copy)
        copy = dict(minimal)
        del copy["software"]["version"]
        with pytest.raises(tskit.ProvenanceValidationError):
            tskit.validate_provenance(copy)

    def test_software_types(self):
        for bad_type in [0, [1, 2, 3], {}]:
            doc = get_provenance(software_name=bad_type)
            with pytest.raises(tskit.ProvenanceValidationError):
                tskit.validate_provenance(doc)
            doc = get_provenance(software_version=bad_type)
            with pytest.raises(tskit.ProvenanceValidationError):
                tskit.validate_provenance(doc)

    def test_schema_version_empth(self):
        doc = get_provenance(schema_version="")
        with pytest.raises(tskit.ProvenanceValidationError):
            tskit.validate_provenance(doc)

    def test_software_empty_strings(self):
        doc = get_provenance(software_name="")
        with pytest.raises(tskit.ProvenanceValidationError):
            tskit.validate_provenance(doc)
        doc = get_provenance(software_version="")
        with pytest.raises(tskit.ProvenanceValidationError):
            tskit.validate_provenance(doc)

    def test_minimal(self):
        minimal = {
            "schema_version": "1",
            "software": {"name": "x", "version": "y"},
            "environment": {},
            "parameters": {},
        }
        tskit.validate_provenance(minimal)

    def test_extra_stuff(self):
        extra = {
            "you": "can",
            "schema_version": "1",
            "software": {"put": "anything", "name": "x", "version": "y"},
            "environment": {"extra": ["you", "want"]},
            "parameters": {"so": ["long", "its", "JSON", 0]},
        }
        tskit.validate_provenance(extra)

    def test_resources(self):
        resources = {
            "schema_version": "1",
            "software": {"name": "x", "version": "y"},
            "environment": {},
            "parameters": {},
            "resources": {
                "elapsed_time": 1,
                "user_time": 2,
                "sys_time": 3,
                "max_memory": 4,
            },
        }
        tskit.validate_provenance(resources)

    def test_resources_error(self):
        resources = {
            "schema_version": "1",
            "software": {"name": "x", "version": "y"},
            "environment": {},
            "parameters": {},
            "resources": {
                "elapsed_time": "1",
                "user_time": 2,
                "sys_time": 3,
                "max_memory": 4,
            },
        }
        with pytest.raises(tskit.ProvenanceValidationError):
            tskit.validate_provenance(resources)


class TestOutputProvenance:
    """
    Check that the schemas we produce in tskit are valid.
    """

    def test_simplify(self):
        ts = msprime.simulate(5, random_seed=1)
        ts = ts.simplify()
        prov = json.loads(ts.provenance(1).record)
        tskit.validate_provenance(prov)
        assert prov["parameters"]["command"] == "simplify"
        assert prov["environment"] == provenance.get_environment(include_tskit=False)
        assert prov["software"] == {"name": "tskit", "version": tskit.__version__}


class TestEnvironment:
    """
    Tests for the environment provenance.
    """

    def test_os(self):
        env = provenance.get_environment()
        os = {
            "system": platform.system(),
            "node": platform.node(),
            "release": platform.release(),
            "version": platform.version(),
            "machine": platform.machine(),
        }
        assert env["os"] == os

    def test_python(self):
        env = provenance.get_environment()
        python = {
            "implementation": platform.python_implementation(),
            "version": platform.python_version(),
        }
        assert env["python"] == python

    def test_libraries(self):
        kastore_lib = {"version": ".".join(map(str, _tskit.get_kastore_version()))}
        env = provenance.get_environment()
        assert {"kastore": kastore_lib, "tskit": {"version": tskit.__version__}} == env[
            "libraries"
        ]

        env = provenance.get_environment(include_tskit=False)
        assert {"kastore": kastore_lib} == env["libraries"]

        extra_libs = {"abc": [], "xyz": {"one": 1}}
        env = provenance.get_environment(include_tskit=False, extra_libs=extra_libs)
        libs = {"kastore": kastore_lib}
        libs.update(extra_libs)
        assert libs == env["libraries"]


class TestGetResources:
    def test_get_resources_keys(self):
        resources = provenance.get_resources(0)
        assert "elapsed_time" in resources
        assert "user_time" in resources
        assert "sys_time" in resources
        if resource is not None:
            assert "max_memory" in resources

    def test_get_resources_values(self):
        delta = 0.1
        t = time.time()
        resources = provenance.get_resources(t - delta)
        assert isinstance(resources["elapsed_time"], float)
        assert isinstance(resources["user_time"], float)
        assert isinstance(resources["sys_time"], float)
        assert resources["elapsed_time"] >= delta - 0.001
        assert resources["user_time"] > 0
        assert resources["sys_time"] > 0
        if resource is not None:
            assert isinstance(resources["max_memory"], int)
            assert resources["max_memory"] > 1024

    def test_get_resources_platform(self):
        resources = provenance.get_resources(0)
        if sys.platform != "darwin" and resource is not None:
            assert resources["max_memory"] % 1024 == 0


class TestGetSchema:
    """
    Ensure we return the correct JSON schema.
    """

    def test_file_equal(self):
        s1 = provenance.get_schema()
        base = os.path.join(os.path.dirname(__file__), "..", "tskit")
        with open(os.path.join(base, "provenance.schema.json")) as f:
            s2 = json.load(f)
        assert s1 == s2

    def test_caching(self):
        n = 10
        schemas = [provenance.get_schema() for _ in range(n)]
        # Ensure all the schemas are different objects.
        assert len(set(map(id, schemas))) == n
        # Ensure the schemas are all equal
        for j in range(n):
            assert schemas[0] == schemas[j]

    def test_form(self):
        s = provenance.get_schema()
        assert s["schema"] == "http://json-schema.org/draft-07/schema#"
        assert s["version"] == "1.1.0"


class TestTreeSeqEditMethods:
    """
    Ensure that tree sequence 'edit' methods correctly record themselves
    """

    def test_keep_delete_different(self):
        ts = msprime.simulate(5, random_seed=1)
        ts_keep = ts.keep_intervals([[0.25, 0.5]])
        ts_del = ts.delete_intervals([[0, 0.25], [0.5, 1.0]])
        assert ts_keep.num_provenances == ts_del.num_provenances
        for i, (p1, p2) in enumerate(zip(ts_keep.provenances(), ts_del.provenances())):
            if i == ts_keep.num_provenances - 1:
                # last one should be different
                assert p1.record != p2.record
            else:
                assert p1.record == p2.record


--- ../../tskit/python/tests/test_genotype_matching.py ---

import copy
import itertools

import lshmm as ls
import msprime
import numpy as np
import pytest

import tskit

EQUAL_BOTH_HOM = 4
UNEQUAL_BOTH_HOM = 0
BOTH_HET = 7
REF_HOM_OBS_HET = 1
REF_HET_OBS_HOM = 2

MISSING = -1


def mirror_coordinates(ts):
    """
    Returns a copy of the specified tree sequence in which all
    coordinates x are transformed into L - x.
    """
    L = ts.sequence_length
    tables = ts.dump_tables()
    left = tables.edges.left
    right = tables.edges.right
    tables.edges.left = L - right
    tables.edges.right = L - left
    tables.sites.position = L - tables.sites.position  # + 1
    # TODO migrations.
    tables.sort()
    return tables.tree_sequence()


class ValueTransition:
    """Simple struct holding value transition values."""

    def __init__(self, tree_node=-1, inner_summation=-1, value_list=-1, value_index=-1):
        self.tree_node = tree_node
        self.value_list = value_list
        self.inner_summation = inner_summation
        self.value_index = value_index

    def copy(self):
        return ValueTransition(
            self.tree_node,
            self.inner_summation,
            self.value_list,
            self.value_index,
        )

    def __repr__(self):
        return repr(self.__dict__)

    def __str__(self):
        return repr(self)


class InternalValueTransition:
    """Simple struct holding the internal value transition values."""

    def __init__(self, tree_node=-1, value=-1, inner_summation=-1):
        self.tree_node = tree_node
        self.value = value
        self.inner_summation = inner_summation

    def __repr__(self):
        return repr(self.__dict__)

    def __str__(self):
        return repr(self)


class LsHmmAlgorithm:
    """
    Abstract superclass of Li and Stephens HMM algorithm.
    """

    def __init__(self, ts, rho, mu, precision=30):
        self.ts = ts
        self.mu = mu
        self.rho = rho
        self.precision = precision
        # The array of ValueTransitions.
        self.T = []
        # indexes in to the T array for each node.
        self.T_index = np.zeros(ts.num_nodes, dtype=int) - 1
        # The number of nodes underneath each element in the T array.
        self.N = np.zeros(ts.num_nodes, dtype=int)
        # Efficiently compute the allelic state at a site
        self.allelic_state = np.zeros(ts.num_nodes, dtype=int) - 1
        # Diffs so we can can update T and T_index between trees.
        self.edge_diffs = self.ts.edge_diffs()
        self.parent = np.zeros(self.ts.num_nodes, dtype=int) - 1
        self.tree = tskit.Tree(self.ts)
        self.output = None

    def decode_site_dict(self):
        """
        Decodes the tree encoding of the values into an explicit
        matrix.
        """
        A = np.zeros((self.ts.num_samples, self.ts.num_samples))
        # To look at the inner summations too.
        B = np.zeros((self.ts.num_samples, self.ts.num_samples))
        f = {st.tree_node: st for st in self.T}

        for j1, u1 in enumerate(self.ts.samples()):
            while u1 not in f:
                u1 = self.tree.parent(u1)
            f1 = {st.tree_node: st for st in f[u1].value_list}
            for j2, u2 in enumerate(self.ts.samples()):
                while u2 not in f1:
                    u2 = self.tree.parent(u2)
                A[j1, j2] = f1[u2].value
                B[j1, j2] = f1[u2].inner_summation
        return A, B

    def check_integrity(self):
        M = [st.tree_node for st in self.T if st.tree_node != -1]
        assert np.all(self.T_index[M] >= 0)
        index = np.ones_like(self.T_index, dtype=bool)
        index[M] = 0
        assert np.all(self.T_index[index] == -1)
        for j, st in enumerate(self.T):
            if st.tree_node != -1:
                assert j == self.T_index[st.tree_node]

    def stupid_compress_dict(self):
        """
        Duncan created a compression that just runs parsimony so is
        guaranteed to work.
        """
        tree = self.tree
        T = self.T
        alleles_string_vec = np.zeros(tree.num_samples()).astype("object")
        genotypes = np.zeros(tree.num_samples(), dtype=int)
        genotype_index = 0
        mapping_back = {}

        node_map = {st.tree_node: st for st in self.T}

        for st1 in T:
            if st1.tree_node != -1:
                alleles_string_tmp = [
                    f"{st2.tree_node}:{st2.value:.16f}" for st2 in st1.value_list
                ]
                alleles_string = ",".join(alleles_string_tmp)
                # Add an extra element that tells me the alleles_string there.
                st1.alleles_string = alleles_string
                st1.genotype_index = genotype_index
                # assert alleles_string not in mapping_back
                if alleles_string not in mapping_back:
                    mapping_back[alleles_string] = {
                        "tree_node": st1.tree_node,
                        "value_list": st1.value_list,
                        "inner_summation": st1.inner_summation,
                    }
                genotype_index += 1

        for leaf in tree.samples():
            u = leaf
            while u not in node_map:
                u = tree.parent(u)
            genotypes[leaf] = node_map[u].genotype_index

        alleles_string_vec = []
        for st in T:
            if st.tree_node != -1:
                alleles_string_vec.append(st.alleles_string)

        ancestral_allele, mutations = tree.map_mutations(genotypes, alleles_string_vec)

        # Retain the old T_index, because the internal T that's passed up the tree will
        # retain this ordering.
        old_T_index = copy.deepcopy(self.T_index)
        self.T_index = np.zeros(tree.tree_sequence.num_nodes, dtype=int) - 1
        self.N = np.zeros(tree.tree_sequence.num_nodes, dtype=int)
        self.T.clear()

        # First, create T root.
        self.T_index[tree.root] = 0
        self.T.append(
            ValueTransition(
                tree_node=tree.root,
                value_list=[
                    InternalValueTransition(
                        tree_node=tree.root,
                        value=mapping_back[ancestral_allele]["value_list"][
                            old_T_index[mapping_back[ancestral_allele]["tree_node"]]
                        ].value,
                    )
                ],
            )
        )

        # Then create the rest of T, adding the root each time to value_list
        for i, mut in enumerate(mutations):
            self.T_index[mut.node] = i + 1
            self.T.append(
                ValueTransition(
                    tree_node=mut.node,
                    value_list=[
                        InternalValueTransition(
                            tree_node=tree.root,
                            value=mapping_back[mut.derived_state]["value_list"][
                                old_T_index[mapping_back[ancestral_allele]["tree_node"]]
                            ].value,
                        )
                    ],
                )
            )

        # First add to the root
        for mut in mutations:
            self.T[self.T_index[tree.root]].value_list.append(
                InternalValueTransition(
                    tree_node=mut.node,
                    value=mapping_back[ancestral_allele]["value_list"][
                        old_T_index[mapping_back[mut.derived_state]["tree_node"]]
                    ].value,
                )
            )

        # Then add the rest of T_internal to each internal T.
        for mut1 in mutations:
            for mut2 in mutations:
                self.T[self.T_index[mut1.node]].value_list.append(
                    InternalValueTransition(
                        tree_node=mut2.node,
                        value=mapping_back[mut1.derived_state]["value_list"][
                            old_T_index[mapping_back[mut2.derived_state]["tree_node"]]
                        ].value,
                    )
                )

        # General approach here is to use
        # mapping_back[mut.derived_state]['value_list'][
        #   old_T_index[mapping_back[mut2.derived_state]["tree_node"]
        # ] and append this to the T_inner.

        node_map = {st.tree_node: st for st in self.T}

        for u in tree.samples():
            while u not in node_map:
                u = tree.parent(u)
            self.N[self.T_index[u]] += 1

    def update_tree(self):
        """
        Update the internal data structures to move on to the next tree.
        """
        parent = self.parent
        T_index = self.T_index
        T = self.T
        _, edges_out, edges_in = next(self.edge_diffs)

        for edge in edges_out:
            u = edge.child
            if T_index[u] == -1:
                # Make sure the subtree we're detaching has an T_index-value at the root.
                while T_index[u] == -1:
                    u = parent[u]
                    assert u != -1
                T_index[edge.child] = len(T)
                T.append(
                    ValueTransition(
                        tree_node=edge.child,
                        value_list=copy.deepcopy(T[T_index[u]].value_list),
                    )
                )
                # Add on this extra node to each of the internal lists
                for st in T:
                    if not (st.value_list == tskit.NULL):
                        st.value_list.append(
                            InternalValueTransition(
                                tree_node=edge.child,
                                value=st.value_list.copy()[T_index[u]].value,
                            )
                        )
            parent[edge.child] = -1

        for edge in edges_in:
            parent[edge.child] = edge.parent
            u = edge.parent
            if parent[edge.parent] == -1:
                # Grafting onto a new root.
                if T_index[edge.parent] == -1:
                    T_index[edge.parent] = len(T)
                    T.append(
                        ValueTransition(
                            tree_node=edge.parent,
                            value_list=copy.deepcopy(T[T_index[edge.child]].value_list),
                        )
                    )
                    # Add on this extra node to each of the internal lists
                    for st in T:
                        if not (st.value_list == tskit.NULL):
                            st.value_list.append(
                                InternalValueTransition(
                                    tree_node=edge.parent,
                                    value=st.value_list.copy()[
                                        T_index[edge.child]
                                    ].value,
                                )
                            )
            else:
                # Grafting into an existing subtree.
                while T_index[u] == -1:
                    u = parent[u]
                    assert u != -1
            assert T_index[u] != -1 and T_index[edge.child] != -1
            if (
                T[T_index[u]].value_list == T[T_index[edge.child]].value_list
            ):  # DEV: is this fine?
                st = T[T_index[edge.child]]
                # Mark the lower ValueTransition as unused.
                st.value_list = -1
                # Also need to mark the corresponding InternalValueTransition as
                # unused for the remaining states
                for st2 in T:
                    if not (st2.value_list == tskit.NULL):
                        st2.value_list[T_index[edge.child]].value = -1
                        st2.value_list[T_index[edge.child]].tree_node = -1

                st.tree_node = -1
                T_index[edge.child] = -1

        # We can have values left over still pointing to old roots. Remove
        for root in self.tree.roots:
            if T_index[root] != -1:
                # Use a special marker here to designate the real roots.
                T[T_index[root]].value_index = -2

        for vt in T:
            if vt.tree_node != -1:
                if parent[vt.tree_node] == -1 and vt.value_index != -2:
                    # Also need to mark the corresponding InternalValueTransition as
                    # unused for the remaining states
                    for st2 in T:
                        if not (st2.value_list == tskit.NULL):
                            st2.value_list[T_index[vt.tree_node]].value = -1
                            st2.value_list[T_index[vt.tree_node]].tree_node = -1
                    T_index[vt.tree_node] = -1
                    vt.tree_node = -1
                vt.value_index = -1

        self.N = np.zeros(self.tree.tree_sequence.num_nodes, dtype=int)
        node_map = {st.tree_node: st for st in self.T}

        for u in self.tree.samples():
            while u not in node_map:
                u = self.tree.parent(u)
            self.N[self.T_index[u]] += 1

    def update_probabilities(self, site, genotype_state):
        tree = self.tree
        T_index = self.T_index
        T = self.T
        alleles = ["0", "1"]
        allelic_state = self.allelic_state
        # Set the allelic_state for this site.
        allelic_state[self.tree.root] = alleles.index(site.ancestral_state)
        normalisation_factor_inner = {}

        for st1 in T:
            if st1.tree_node != -1:
                normalisation_factor_inner[st1.tree_node] = (
                    self.compute_normalisation_factor_inner_dict(st1.tree_node)
                )

        for st1 in T:
            if st1.tree_node != -1:
                for st2 in st1.value_list:
                    if st2.tree_node != -1:
                        self.T[self.T_index[st1.tree_node]].value_list[
                            self.T_index[st2.tree_node]
                        ].inner_summation = self.inner_summation_evaluation(
                            normalisation_factor_inner,
                            st1.tree_node,
                            st2.tree_node,
                        )
                        # (
                        #     normalisation_factor_inner[st1.tree_node]
                        #     + normalisation_factor_inner[st2.tree_node]
                        # )

        for mutation in site.mutations:
            u = mutation.node
            allelic_state[u] = alleles.index(mutation.derived_state)
            if T_index[u] == -1:
                while T_index[u] == tskit.NULL:
                    u = tree.parent(u)
                T_index[mutation.node] = len(T)
                T.append(
                    ValueTransition(
                        tree_node=mutation.node,
                        value_list=copy.deepcopy(T[T_index[u]].value_list),
                    )  # DEV: is it possible to not use deepcopies?
                )
                for st in T:
                    if not (st.value_list == tskit.NULL):
                        st.value_list.append(
                            InternalValueTransition(
                                tree_node=mutation.node,
                                value=st.value_list.copy()[T_index[u]].value,
                                inner_summation=st.value_list.copy()[
                                    T_index[u]
                                ].inner_summation,
                            )
                        )

        # Get the allelic state at the leaves.
        allelic_state[: tree.num_samples()] = tree.tree_sequence.genotype_matrix()[
            site.id, :
        ]

        query_is_het = genotype_state == 1
        query_is_missing = genotype_state == MISSING

        for st1 in T:
            u1 = st1.tree_node

            if u1 != -1:
                # Get the allelic_state at u. TODO we can cache these states to
                # avoid some upward traversals.
                v1 = u1
                while allelic_state[v1] == -1:
                    v1 = tree.parent(v1)
                    assert v1 != -1

                for st2 in st1.value_list:
                    u2 = st2.tree_node
                    if u2 != -1:
                        # Get the allelic_state at u. TODO we can cache these states to
                        # avoid some upward traversals.
                        v2 = u2
                        while allelic_state[v2] == -1:
                            v2 = tree.parent(v2)
                            assert v2 != -1

                        genotype_template_state = allelic_state[v1] + allelic_state[v2]
                        match = genotype_state == genotype_template_state
                        template_is_het = genotype_template_state == 1
                        # Fill in the value at the combination of states: (s1, s2)
                        st2.value = self.compute_next_probability_dict(
                            site.id,
                            st2.value,
                            st2.inner_summation,
                            match,
                            template_is_het,
                            query_is_het,
                            query_is_missing,
                            u1,
                            u2,
                            # Last two are not used by forward-backward
                            # but required by Viterbi
                        )

                # This will ensure that allelic_state[:n] is filled
                genotype_template_state = (
                    allelic_state[v1] + allelic_state[: tree.num_samples()]
                )
                # These are vectors of length n (at internal nodes).
                match = genotype_state == genotype_template_state
                template_is_het = genotype_template_state == 1

        # Unset the states
        allelic_state[tree.root] = -1
        for mutation in site.mutations:
            allelic_state[mutation.node] = -1

    def process_site(
        self, site, genotype_state, forwards=True
    ):  # Note: forwards turned on for Viterbi
        if forwards:
            # Forwards algorithm
            self.update_probabilities(site, genotype_state)
            self.stupid_compress_dict()
            s1 = self.compute_normalisation_factor_dict()
            T = self.T

            for st in T:
                if st.tree_node != tskit.NULL:
                    # Need to loop through value copy, and normalise
                    for st2 in st.value_list:
                        st2.value /= s1
                        st2.value = np.round(st2.value, self.precision)

            self.output.store_site(
                site.id, s1, [(st.tree_node, st.value_list) for st in self.T]
            )
        else:
            # Backwards algorithm
            self.output.store_site(
                site.id,
                self.output.normalisation_factor[site.id],
                [(st.tree_node, st.value_list) for st in self.T],
            )
            self.update_probabilities(site, genotype_state)
            self.stupid_compress_dict()
            b_last_sum = self.compute_normalisation_factor_dict()  # (Double sum)

            normalisation_factor_inner = {}

            for st1 in self.T:
                if st1.tree_node != -1:
                    normalisation_factor_inner[st1.tree_node] = (
                        self.compute_normalisation_factor_inner_dict(st1.tree_node)
                    )

            for st1 in self.T:
                if st1.tree_node != -1:
                    for st2 in st1.value_list:
                        if st2.tree_node != -1:
                            self.T[self.T_index[st1.tree_node]].value_list[
                                self.T_index[st2.tree_node]
                            ].inner_summation = (
                                normalisation_factor_inner[st1.tree_node]
                                + normalisation_factor_inner[st2.tree_node]
                            )

            s = self.output.normalisation_factor[site.id]
            for st1 in self.T:
                if st1.tree_node != tskit.NULL:
                    for st2 in st1.value_list:
                        st2.value = (
                            ((self.rho[site.id] / self.ts.num_samples) ** 2)
                            * b_last_sum
                            + (1 - self.rho[site.id])
                            * (self.rho[site.id] / self.ts.num_samples)
                            * st2.inner_summation
                            + (1 - self.rho[site.id]) ** 2 * st2.value
                        )
                        st2.value /= s
                        st2.value = np.round(st2.value, self.precision)

    def run_forward(self, g):
        n = self.ts.num_samples
        self.tree.clear()
        for u in self.ts.samples():
            self.T_index[u] = len(self.T)
            self.T.append(ValueTransition(tree_node=u, value_list=[]))
            for v in self.ts.samples():
                self.T[self.T_index[u]].value_list.append(
                    InternalValueTransition(tree_node=v, value=(1 / n) ** 2)
                )

        while self.tree.next():
            self.update_tree()
            for site in self.tree.sites():
                self.process_site(site, g[site.id])

        return self.output

    def run_backward(self, g):
        self.tree.clear()
        for u in self.ts.samples():
            self.T_index[u] = len(self.T)
            self.T.append(ValueTransition(tree_node=u, value_list=[]))
            for v in self.ts.samples():
                self.T[self.T_index[u]].value_list.append(
                    InternalValueTransition(tree_node=v, value=1)
                )

        while self.tree.next():
            self.update_tree()
            for site in self.tree.sites():
                self.process_site(site, g[site.id], forwards=False)
        return self.output

    def compute_normalisation_factor_dict(self):
        raise NotImplementedError()

    def compute_next_probability_dict(
        self,
        site_id,
        p_last,
        inner_summation,
        is_match,
        template_is_het,
        query_is_het,
        query_is_missing,
        node_1,  # Note: these are only used in Viterbi (node_1 and node_2)
        node_2,
    ):
        raise NotImplementedError()


class CompressedMatrix:
    """
    Class representing a num_samples x num_sites matrix compressed by a
    tree sequence. Each site is represented by a set of (node, value)
    pairs, which act as "mutations", i.e., any sample that descends
    from a particular node will inherit that value (unless any other
    values are on the path).
    """

    def __init__(self, ts, normalisation_factor=None):
        self.ts = ts
        self.num_sites = ts.num_sites
        self.num_samples = ts.num_samples
        self.value_transitions = [None for _ in range(self.num_sites)]
        if normalisation_factor is None:
            self.normalisation_factor = np.zeros(self.num_sites)
        else:
            self.normalisation_factor = normalisation_factor
            assert len(self.normalisation_factor) == self.num_sites

    def store_site(self, site, normalisation_factor, value_transitions):
        self.normalisation_factor[site] = normalisation_factor
        self.value_transitions[site] = copy.deepcopy(value_transitions)

    # Expose the same API as the low-level classes

    @property
    def num_transitions(self):
        a = [len(self.value_transitions[j]) for j in range(self.num_sites)]
        return np.array(a, dtype=np.int32)

    def get_site(self, site):
        return self.value_transitions[site]

    def decode_site_dict(self, tree, site_id):
        """
        Decodes the tree encoding of the values into an explicit
        matrix.
        """
        A = np.zeros((self.num_samples, self.num_samples))
        f = dict(self.value_transitions[site_id])

        for j1, u1 in enumerate(self.ts.samples()):
            while u1 not in f:
                u1 = tree.parent(u1)
            f1 = {st.tree_node: st for st in f[u1]}
            for j2, u2 in enumerate(self.ts.samples()):
                while u2 not in f1:
                    u2 = tree.parent(u2)
                A[j1, j2] = f1[u2].value
        return A

    def decode(self):
        """
        Decodes the tree encoding of the values into an explicit
        matrix.
        """
        A = np.zeros((self.num_sites, self.num_samples, self.num_samples))
        for tree in self.ts.trees():
            for site in tree.sites():
                A[site.id] = self.decode_site_dict(tree, site.id)
        return A


class ForwardMatrix(CompressedMatrix):
    """Class representing a compressed forward matrix."""


class BackwardMatrix(CompressedMatrix):
    """Class representing a compressed forward matrix."""


class ViterbiMatrix(CompressedMatrix):
    """
    Class representing the compressed Viterbi matrix.
    """

    def __init__(self, ts):
        super().__init__(ts)
        # Tuples containing the site, the pair of nodes in the tree,
        # and whether recombination is required
        self.double_recombination_required = [(-1, 0, 0, False)]
        self.single_recombination_required = [(-1, 0, 0, False)]

    def add_single_recombination_required(self, site, node_s1, node_s2, required):
        self.single_recombination_required.append((site, node_s1, node_s2, required))

    def add_double_recombination_required(self, site, node_s1, node_s2, required):
        self.double_recombination_required.append((site, node_s1, node_s2, required))

    def choose_sample_double(self, site_id, tree):
        max_value = -1
        u1 = -1
        u2 = -1

        for node_s1, value_outer in self.value_transitions[site_id]:
            for value_list in value_outer:
                value_tmp = value_list
                if value_tmp.value > max_value:
                    max_value = value_tmp.value
                    u1 = node_s1
                    u2 = value_tmp.tree_node

        assert u1 != -1
        assert u2 != -1

        transition_nodes = [u_tmp for (u_tmp, _) in self.value_transitions[site_id]]

        while not tree.is_sample(u1):
            for v in tree.children(u1):
                if v not in transition_nodes:
                    u1 = v
                    break
            else:
                raise AssertionError("could not find path")

        while not tree.is_sample(u2):
            for v in tree.children(u2):
                if v not in transition_nodes:
                    u2 = v
                    break
            else:
                raise AssertionError("could not find path")

        return (u1, u2)

    def choose_sample_single(self, site_id, tree, current_nodes):
        # I want to find which is the max between any choice if I switch just u1,
        # and any choice if I switch just u2.
        node_map = {st[0]: st for st in self.value_transitions[site_id]}
        to_compute = (
            np.zeros(2, dtype=int) - 1
        )  # We have two to compute - one for each single switch set of possibilities.

        for i, v in enumerate(current_nodes):  # (u1, u2)
            while v not in node_map:
                v = tree.parent(v)
            to_compute[i] = v

        # Need to go to the (j1 :)th entries, and the (:,j2)the entries,
        # and pick the best.
        T_index = np.zeros(self.ts.num_nodes, dtype=int) - 1
        for j, st in enumerate(self.value_transitions[site_id]):
            T_index[st[0]] = j

        node_single_switch_maxes = np.zeros(2, dtype=int) - 1
        single_switch = np.zeros(2) - 1

        for i, node in enumerate(to_compute):
            value_list = self.value_transitions[site_id][T_index[node]][1]
            s_inner = 0
            for st in value_list:
                j = st.tree_node
                if j != -1:
                    max_st = st.value
                    max_arg = st.tree_node
                    if max_st > s_inner:
                        s_inner = max_st
                        s_arg = max_arg
            node_single_switch_maxes[i] = s_arg
            single_switch[i] = s_inner

        if np.argmax(single_switch) == 0:
            # u1 is fixed, and we switch u2
            u1 = current_nodes[0]
            current_nodes = (u1, node_single_switch_maxes[0])
        else:
            # u2 is fixed, and we switch u1.
            u2 = current_nodes[1]
            current_nodes = (node_single_switch_maxes[1], u2)

        u1 = current_nodes[0]
        u2 = current_nodes[1]

        # Find the collection of transition nodes to use to descend down the tree
        transition_nodes = [u for (u, _) in self.value_transitions[site_id]]

        # Traverse down to find a leaves.
        while not tree.is_sample(u1):
            for v in tree.children(u1):
                if v not in transition_nodes:
                    u1 = v
                    break
            else:
                raise AssertionError("could not find path")

        while not tree.is_sample(u2):
            for v in tree.children(u2):
                if v not in transition_nodes:
                    u2 = v
                    break
            else:
                raise AssertionError("could not find path")

        current_nodes = (u1, u2)

        return current_nodes

    def traceback(self):
        # Run the traceback.
        m = self.ts.num_sites
        match = np.zeros((m, 2), dtype=int)

        single_recombination_tree = (
            np.zeros((self.ts.num_nodes, self.ts.num_nodes), dtype=int) - 1
        )
        double_recombination_tree = (
            np.zeros((self.ts.num_nodes, self.ts.num_nodes), dtype=int) - 1
        )

        tree = tskit.Tree(self.ts)
        tree.last()
        double_switch = True
        current_nodes = (-1, -1)
        current_node_outer = current_nodes[0]

        rr_single_index = len(self.single_recombination_required) - 1
        rr_double_index = len(self.double_recombination_required) - 1

        for site in reversed(self.ts.sites()):
            while tree.interval.left > site.position:
                tree.prev()
            assert tree.interval.left <= site.position < tree.interval.right

            # Fill in the recombination single tree
            j_single = rr_single_index
            # The above starts from the end of all the recombination required
            # information, and includes all the information for the current site.
            while self.single_recombination_required[j_single][0] == site.id:
                u1, u2, required = self.single_recombination_required[j_single][1:]
                single_recombination_tree[u1, u2] = required
                j_single -= 1

            # Fill in the recombination double tree
            j_double = rr_double_index
            # The above starts from the end of all the recombination required
            # information, and includes all the information for the current site.
            while self.double_recombination_required[j_double][0] == site.id:
                u1, u2, required = self.double_recombination_required[j_double][1:]
                double_recombination_tree[u1, u2] = required
                j_double -= 1

            # Note - current nodes are the leaf nodes.
            if current_node_outer == -1:
                if double_switch:
                    current_nodes = self.choose_sample_double(site.id, tree)
                else:
                    current_nodes = self.choose_sample_single(
                        site.id, tree, current_nodes
                    )

            match[site.id, :] = current_nodes

            # Now traverse up the tree from the current node. The first marked node
            # we meet tells us whether we need to recombine.
            current_node_outer = current_nodes[0]
            u1 = current_node_outer
            u2 = current_nodes[1]

            # Just need to move up the tree to evaluate u1 and u2.
            if double_switch:
                while u1 != -1 and double_recombination_tree[u1, u1] == -1:
                    u1 = tree.parent(u1)

                while u2 != -1 and double_recombination_tree[u1, u2] == -1:
                    u2 = tree.parent(u2)
            else:
                while u1 != -1 and single_recombination_tree[u1, u1] == -1:
                    u1 = tree.parent(u1)

                while u2 != -1 and single_recombination_tree[u1, u2] == -1:
                    u2 = tree.parent(u2)

            assert u1 != -1
            assert u2 != -1

            if double_recombination_tree[u1, u2] == 1:
                # Need to double switch at the next site.
                current_node_outer = -1
                double_switch = True
            elif single_recombination_tree[u1, u2] == 1:
                # Need to single switch at the next site
                current_node_outer = -1
                double_switch = False

            # Reset the nodes in the double recombination tree.
            j = rr_single_index
            while self.single_recombination_required[j][0] == site.id:
                u1_tmp, u2_tmp, _ = self.single_recombination_required[j][1:]
                single_recombination_tree[u1_tmp, u2_tmp] = -1
                j -= 1
            rr_single_index = j

            # Reset the nodes in the single recombination tree.
            j = rr_double_index
            while self.double_recombination_required[j][0] == site.id:
                u1_tmp, u2_tmp, _ = self.double_recombination_required[j][1:]
                double_recombination_tree[u1_tmp, u2_tmp] = -1
                j -= 1
            rr_double_index = j

        return match


class ForwardAlgorithm(LsHmmAlgorithm):
    """Runs the Li and Stephens forward algorithm."""

    def __init__(self, ts, rho, mu, precision=30):
        super().__init__(ts, rho, mu, precision)
        self.output = ForwardMatrix(ts)

    def inner_summation_evaluation(
        self, normalisation_factor_inner, st1_tree_node, st2_tree_node
    ):
        return (
            normalisation_factor_inner[st1_tree_node]
            + normalisation_factor_inner[st2_tree_node]
        )

    def compute_normalisation_factor_dict(self):
        s = 0
        for j, st in enumerate(self.T):
            assert st.tree_node != tskit.NULL
            assert self.N[j] > 0
            s += self.N[j] * self.compute_normalisation_factor_inner_dict(st.tree_node)
        return s

    def compute_normalisation_factor_inner_dict(self, node):
        s_inner = 0
        F_previous = self.T[self.T_index[node]].value_list
        for st in F_previous:
            j = st.tree_node
            if j != -1:
                s_inner += self.N[self.T_index[j]] * st.value
        return s_inner

    def compute_next_probability_dict(
        self,
        site_id,
        p_last,
        inner_normalisation_factor,
        is_match,
        template_is_het,
        query_is_het,
        query_is_missing,
        node_1,
        node_2,
    ):
        rho = self.rho[site_id]
        mu = self.mu[site_id]
        n = self.ts.num_samples

        p_t = (
            (rho / n) ** 2
            + ((1 - rho) * (rho / n)) * inner_normalisation_factor
            + (1 - rho) ** 2 * p_last
        )

        if query_is_missing:
            p_e = 1
        else:
            query_is_hom = np.logical_not(query_is_het)
            template_is_hom = np.logical_not(template_is_het)

            equal_both_hom = np.logical_and(
                np.logical_and(is_match, template_is_hom), query_is_hom
            )
            unequal_both_hom = np.logical_and(
                np.logical_and(np.logical_not(is_match), template_is_hom), query_is_hom
            )
            both_het = np.logical_and(template_is_het, query_is_het)
            ref_hom_obs_het = np.logical_and(template_is_hom, query_is_het)
            ref_het_obs_hom = np.logical_and(template_is_het, query_is_hom)

            p_e = (
                equal_both_hom * (1 - mu) ** 2
                + unequal_both_hom * (mu**2)
                + ref_hom_obs_het * (2 * mu * (1 - mu))
                + ref_het_obs_hom * (mu * (1 - mu))
                + both_het * ((1 - mu) ** 2 + mu**2)
            )

        return p_t * p_e


class BackwardAlgorithm(LsHmmAlgorithm):
    """Runs the Li and Stephens forward algorithm."""

    def __init__(self, ts, rho, mu, normalisation_factor, precision=10):
        super().__init__(ts, rho, mu, precision)
        self.output = BackwardMatrix(ts, normalisation_factor)

    def inner_summation_evaluation(
        self, normalisation_factor_inner, st1_tree_node, st2_tree_node
    ):
        return (
            normalisation_factor_inner[st1_tree_node]
            + normalisation_factor_inner[st2_tree_node]
        )

    def compute_normalisation_factor_dict(self):
        s = 0
        for j, st in enumerate(self.T):
            assert st.tree_node != tskit.NULL
            assert self.N[j] > 0
            s += self.N[j] * self.compute_normalisation_factor_inner_dict(st.tree_node)
        return s

    def compute_normalisation_factor_inner_dict(self, node):
        s_inner = 0
        F_previous = self.T[self.T_index[node]].value_list
        for st in F_previous:
            j = st.tree_node
            if j != -1:
                s_inner += self.N[self.T_index[j]] * st.value
        return s_inner

    def compute_next_probability_dict(
        self,
        site_id,
        p_next,
        inner_normalisation_factor,
        is_match,
        template_is_het,
        query_is_het,
        query_is_missing,
        node_1,
        node_2,
    ):
        mu = self.mu[site_id]
        template_is_hom = np.logical_not(template_is_het)

        if query_is_missing:
            p_e = 1
        else:
            query_is_hom = np.logical_not(query_is_het)

            equal_both_hom = np.logical_and(
                np.logical_and(is_match, template_is_hom), query_is_hom
            )
            unequal_both_hom = np.logical_and(
                np.logical_and(np.logical_not(is_match), template_is_hom), query_is_hom
            )
            both_het = np.logical_and(template_is_het, query_is_het)
            ref_hom_obs_het = np.logical_and(template_is_hom, query_is_het)
            ref_het_obs_hom = np.logical_and(template_is_het, query_is_hom)

            p_e = (
                equal_both_hom * (1 - mu) ** 2
                + unequal_both_hom * (mu**2)
                + ref_hom_obs_het * (2 * mu * (1 - mu))
                + ref_het_obs_hom * (mu * (1 - mu))
                + both_het * ((1 - mu) ** 2 + mu**2)
            )

        return p_next * p_e


class ViterbiAlgorithm(LsHmmAlgorithm):
    """
    Runs the Li and Stephens Viterbi algorithm.
    """

    def __init__(self, ts, rho, mu, precision=10):
        super().__init__(ts, rho, mu, precision)
        self.output = ViterbiMatrix(ts)

    def inner_summation_evaluation(
        self, normalisation_factor_inner, st1_tree_node, st2_tree_node
    ):
        return max(
            normalisation_factor_inner[st1_tree_node],
            normalisation_factor_inner[st2_tree_node],
        )

    def compute_normalisation_factor_dict(self):
        s = 0
        for st in self.T:
            assert st.tree_node != tskit.NULL
            max_st = self.compute_normalisation_factor_inner_dict(st.tree_node)
            if max_st > s:
                s = max_st
        if s == 0:
            raise ValueError(
                "Trying to match non-existent allele with zero mutation rate"
            )
        return s

    def compute_normalisation_factor_inner_dict(self, node):
        s_inner = 0
        V_previous = self.T[self.T_index[node]].value_list
        for st in V_previous:
            j = st.tree_node
            if j != -1:
                max_st = st.value
                if max_st > s_inner:
                    s_inner = max_st

        return s_inner

    def compute_next_probability_dict(
        self,
        site_id,
        p_last,
        inner_normalisation_factor,
        is_match,
        template_is_het,
        query_is_het,
        query_is_missing,
        node_1,
        node_2,
    ):
        r = self.rho[site_id]
        mu = self.mu[site_id]
        n = self.ts.num_samples
        r_n = r / n

        double_recombination_required = False
        single_recombination_required = False

        if query_is_missing:
            p_e = 1
        else:
            template_is_hom = np.logical_not(template_is_het)
            query_is_hom = np.logical_not(query_is_het)
            equal_both_hom = np.logical_and(
                np.logical_and(is_match, template_is_hom), query_is_hom
            )
            unequal_both_hom = np.logical_and(
                np.logical_and(np.logical_not(is_match), template_is_hom), query_is_hom
            )
            both_het = np.logical_and(template_is_het, query_is_het)
            ref_hom_obs_het = np.logical_and(template_is_hom, query_is_het)
            ref_het_obs_hom = np.logical_and(template_is_het, query_is_hom)

            p_e = (
                equal_both_hom * (1 - mu) ** 2
                + unequal_both_hom * (mu**2)
                + ref_hom_obs_het * (2 * mu * (1 - mu))
                + ref_het_obs_hom * (mu * (1 - mu))
                + both_het * ((1 - mu) ** 2 + mu**2)
            )

        no_switch = (1 - r) ** 2 + 2 * (r_n * (1 - r)) + r_n**2
        single_switch = r_n * (1 - r) + r_n**2
        double_switch = r_n**2

        V_single_switch = inner_normalisation_factor
        p_t = p_last * no_switch
        single_switch_tmp = single_switch * V_single_switch

        if single_switch_tmp > double_switch:
            # Then single switch is the alternative
            if p_t < single_switch * V_single_switch:
                p_t = single_switch * V_single_switch
                single_recombination_required = True
        else:
            # Double switch is the alternative
            if p_t < double_switch:
                p_t = double_switch
                double_recombination_required = True

        self.output.add_single_recombination_required(
            site_id, node_1, node_2, single_recombination_required
        )
        self.output.add_double_recombination_required(
            site_id, node_1, node_2, double_recombination_required
        )

        return p_t * p_e


def ls_forward_tree(g, ts, rho, mu, precision=30):
    """Forward matrix computation based on a tree sequence."""
    fa = ForwardAlgorithm(ts, rho, mu, precision=precision)
    return fa.run_forward(g)


def ls_backward_tree(g, ts_mirror, rho, mu, normalisation_factor, precision=30):
    """Backward matrix computation based on a tree sequence."""
    ba = BackwardAlgorithm(
        ts_mirror, rho, mu, normalisation_factor, precision=precision
    )
    return ba.run_backward(g)


def ls_viterbi_tree(g, ts, rho, mu, precision=30):
    """
    Viterbi path computation based on a tree sequence.
    """
    va = ViterbiAlgorithm(ts, rho, mu, precision=precision)
    return va.run_forward(g)


class LSBase:
    """Superclass of Li and Stephens tests."""

    def genotype_emission(self, mu, m):
        # Define the emission probability matrix
        e = np.zeros((m, 8))
        e[:, EQUAL_BOTH_HOM] = (1 - mu) ** 2
        e[:, UNEQUAL_BOTH_HOM] = mu**2
        e[:, BOTH_HET] = (1 - mu) ** 2 + mu**2
        e[:, REF_HOM_OBS_HET] = 2 * mu * (1 - mu)
        e[:, REF_HET_OBS_HOM] = mu * (1 - mu)

        return e

    def example_genotypes(self, ts):
        H = ts.genotype_matrix()
        s = H[:, 0].reshape(1, H.shape[0]) + H[:, 1].reshape(1, H.shape[0])
        H = H[:, 2:]

        genotypes = [
            s,
            H[:, -1].reshape(1, H.shape[0]) + H[:, -2].reshape(1, H.shape[0]),
        ]

        s_tmp = s.copy()
        s_tmp[0, -1] = MISSING
        genotypes.append(s_tmp)
        s_tmp = s.copy()
        s_tmp[0, ts.num_sites // 2] = MISSING
        genotypes.append(s_tmp)
        s_tmp = s.copy()
        s_tmp[0, :] = MISSING
        genotypes.append(s_tmp)

        m = ts.get_num_sites()
        n = H.shape[1]

        G = np.zeros((m, n, n))
        for i in range(m):
            G[i, :, :] = np.add.outer(H[i, :], H[i, :])

        return H, G, genotypes

    def example_parameters_genotypes(self, ts, seed=42):
        np.random.seed(seed)
        H, G, genotypes = self.example_genotypes(ts)
        n = H.shape[1]
        m = ts.get_num_sites()

        # Here we have equal mutation and recombination
        r = np.zeros(m) + 0.01
        mu = np.zeros(m) + 0.01
        r[0] = 0

        e = self.genotype_emission(mu, m)

        for s in genotypes:
            yield n, m, G, s, e, r, mu

        # Mixture of random and extremes
        rs = [np.zeros(m) + 0.999, np.zeros(m) + 1e-6, np.random.rand(m)]
        mus = [np.zeros(m) + 0.33, np.zeros(m) + 1e-6, np.random.rand(m) * 0.33]

        s = genotypes[0]
        for r, mu in itertools.product(rs, mus):
            r[0] = 0
            e = self.genotype_emission(mu, m)
            yield n, m, G, s, e, r, mu

    def assertAllClose(self, A, B):
        """Assert that all entries of two matrices are 'close'"""
        assert np.allclose(A, B, rtol=1e-5, atol=1e-8)

    # Define a bunch of very small tree-sequences for testing a collection of
    # parameters on
    @pytest.mark.skip(reason="No plans to implement diploid LS HMM yet.")
    def test_simple_n_10_no_recombination(self):
        ts = msprime.simulate(
            10, recombination_rate=0, mutation_rate=0.5, random_seed=42
        )
        assert ts.num_sites > 3
        self.verify(ts)

    @pytest.mark.skip(reason="No plans to implement diploid LS HMM yet.")
    def test_simple_n_6(self):
        ts = msprime.simulate(6, recombination_rate=2, mutation_rate=7, random_seed=42)
        assert ts.num_sites > 5
        self.verify(ts)

    @pytest.mark.skip(reason="No plans to implement diploid LS HMM yet.")
    def test_simple_n_8_high_recombination(self):
        ts = msprime.simulate(8, recombination_rate=20, mutation_rate=5, random_seed=42)
        assert ts.num_trees > 15
        assert ts.num_sites > 5
        self.verify(ts)

    # FIXME Reducing the number of test cases here as they take a long time to run,
    # and we will want to refactor the test infrastructure when implementing these
    # diploid methods in the library.

    # def test_simple_n_10_no_recombination_high_mut(self):
    #     ts = msprime.simulate(
    #         10, recombination_rate=0, mutation_rate=3, random_seed=42)
    #     assert ts.num_sites > 3
    #     self.verify(ts)

    # def test_simple_n_10_no_recombination_higher_mut(self):
    #     ts = msprime.simulate(
    #         20, recombination_rate=0, mutation_rate=3, random_seed=42)
    #     assert ts.num_sites > 3
    #     self.verify(ts)

    # def test_simple_n_8(self):
    #     ts = msprime.simulate(
    #         8, recombination_rate=2, mutation_rate=5, random_seed=42)
    #     assert ts.num_sites > 5
    #     self.verify(ts)

    # def test_simple_n_16(self):
    #     ts = msprime.simulate(
    #         16, recombination_rate=2, mutation_rate=5, random_seed=42)
    #     assert ts.num_sites > 5
    #     self.verify(ts)

    def verify(self, ts):
        raise NotImplementedError()


class FBAlgorithmBase(LSBase):
    """Base for forwards backwards algorithm tests."""


class VitAlgorithmBase(LSBase):
    """Base for viterbi algoritm tests."""


class TestMirroringDipdict(FBAlgorithmBase):
    """Tests that mirroring the tree sequence and running forwards and backwards
    algorithms give the same log-likelihood of observing the data."""

    def verify(self, ts):
        for n, m, _, s, _, r, mu in self.example_parameters_genotypes(ts):
            # Note, need to remove the first sample from the ts, and ensure that
            # invariant sites aren't removed.
            ts_check, mapping = ts.simplify(
                range(1, n + 1), filter_sites=False, map_nodes=True
            )
            H_check = ts_check.genotype_matrix()
            G_check = np.zeros((m, n, n))
            for i in range(m):
                G_check[i, :, :] = np.add.outer(H_check[i, :], H_check[i, :])

            cm_d = ls_forward_tree(s[0, :], ts_check, r, mu)
            ll_tree = np.sum(np.log10(cm_d.normalisation_factor))

            ts_check_mirror = mirror_coordinates(ts_check)
            r_flip = np.insert(np.flip(r)[:-1], 0, 0)
            cm_mirror = ls_forward_tree(
                np.flip(s[0, :]), ts_check_mirror, r_flip, np.flip(mu)
            )
            ll_mirror_tree_dict = np.sum(np.log10(cm_mirror.normalisation_factor))

            self.assertAllClose(ll_tree, ll_mirror_tree_dict)

            # Ensure that the decoded matrices are the same
            flipped_H_check = np.flip(H_check, axis=0)
            flipped_s = np.flip(s, axis=1)

            F_mirror_matrix, c, ll = ls.forwards(
                flipped_H_check,
                flipped_s,
                ploidy=2,
                prob_recombination=r_flip,
                prob_mutation=np.flip(mu),
                scale_mutation_rate=False,
            )

            self.assertAllClose(F_mirror_matrix, cm_mirror.decode())


class TestForwardDipTree(FBAlgorithmBase):
    """Tests that the tree algorithm computes the same forward matrix as the simple
    method."""

    def verify(self, ts):
        for n, m, _, s, _, r, mu in self.example_parameters_genotypes(ts):
            # Note, need to remove the first sample from the ts, and ensure that
            # invariant sites aren't removed.
            ts_check, mapping = ts.simplify(
                range(1, n + 1), filter_sites=False, map_nodes=True
            )
            H_check = ts_check.genotype_matrix()
            G_check = np.zeros((m, n, n))
            for i in range(m):
                G_check[i, :, :] = np.add.outer(H_check[i, :], H_check[i, :])

            F, c, ll = ls.forwards(
                reference_panel=H_check,
                query=s,
                ploidy=2,
                prob_recombination=r,
                prob_mutation=mu,
                scale_mutation_rate=False,
            )
            cm_d = ls_forward_tree(s[0, :], ts_check, r, mu)
            self.assertAllClose(cm_d.decode(), F)
            ll_tree = np.sum(np.log10(cm_d.normalisation_factor))
            self.assertAllClose(ll, ll_tree)


class TestForwardBackwardTree(FBAlgorithmBase):
    """Tests that the tree algorithm computes the same forward matrix as the simple
    method."""

    def verify(self, ts):
        for n, m, _, s, _, r, mu in self.example_parameters_genotypes(ts):
            # Note, need to remove the first sample from the ts, and ensure that
            # invariant sites aren't removed.
            ts_check, mapping = ts.simplify(
                range(1, n + 1), filter_sites=False, map_nodes=True
            )
            H_check = ts_check.genotype_matrix()
            G_check = np.zeros((m, n, n))
            for i in range(m):
                G_check[i, :, :] = np.add.outer(H_check[i, :], H_check[i, :])

            F, c, ll = ls.forwards(
                reference_panel=H_check,
                query=s,
                ploidy=2,
                prob_recombination=r,
                prob_mutation=mu,
                scale_mutation_rate=False,
            )
            B = ls.backwards(
                reference_panel=H_check,
                query=s,
                ploidy=2,
                normalisation_factor_from_forward=c,
                prob_recombination=r,
                prob_mutation=mu,
                scale_mutation_rate=False,
            )

            # Note, need to remove the first sample from the ts, and ensure that
            # invariant sites aren't removed.
            ts_check = ts.simplify(range(1, n + 1), filter_sites=False)
            c_f = ls_forward_tree(s[0, :], ts_check, r, mu)
            ll_tree = np.sum(np.log10(c_f.normalisation_factor))

            ts_check_mirror = mirror_coordinates(ts_check)
            r_flip = np.flip(r)
            c_b = ls_backward_tree(
                np.flip(s[0, :]),
                ts_check_mirror,
                r_flip,
                np.flip(mu),
                np.flip(c_f.normalisation_factor),
            )
            B_tree = np.flip(c_b.decode(), axis=0)
            F_tree = c_f.decode()

            self.assertAllClose(B, B_tree)
            self.assertAllClose(F, F_tree)
            self.assertAllClose(ll, ll_tree)


class TestTreeViterbiDip(VitAlgorithmBase):
    """
    Test that we have the same log-likelihood between tree and matrix
    implementations
    """

    def verify(self, ts):
        for n, m, _, s, _, r, mu in self.example_parameters_genotypes(ts):
            # Note, need to remove the first sample from the ts, and ensure that
            # invariant sites aren't removed.
            ts_check, mapping = ts.simplify(
                range(1, n + 1), filter_sites=False, map_nodes=True
            )
            H_check = ts_check.genotype_matrix()
            G_check = np.zeros((m, n, n))
            for i in range(m):
                G_check[i, :, :] = np.add.outer(H_check[i, :], H_check[i, :])
            ts_check = ts.simplify(range(1, n + 1), filter_sites=False)

            phased_path, ll = ls.viterbi(
                reference_panel=H_check,
                query=s,
                ploidy=2,
                prob_recombination=r,
                prob_mutation=mu,
                scale_mutation_rate=False,
            )
            path_ll_matrix = ls.path_loglik(
                reference_panel=H_check,
                query=s,
                ploidy=2,
                path=phased_path,
                prob_recombination=r,
                prob_mutation=mu,
                scale_mutation_rate=False,
            )

            c_v = ls_viterbi_tree(s[0, :], ts_check, r, mu)
            ll_tree = np.sum(np.log10(c_v.normalisation_factor))

            # Attempt to get the path
            path_tree_dict = c_v.traceback()
            # Work out the likelihood of the proposed path
            path_ll_tree = ls.path_loglik(
                reference_panel=H_check,
                query=s,
                ploidy=2,
                path=np.transpose(path_tree_dict),
                prob_recombination=r,
                prob_mutation=mu,
                scale_mutation_rate=False,
            )

            self.assertAllClose(ll, ll_tree)
            self.assertAllClose(path_ll_tree, path_ll_matrix)


--- ../../tskit/python/tests/test_coalrate.py ---


"""
Test cases for coalescence rate calculation in tskit.
"""
import itertools

import msprime
import numpy as np
import pytest

import tests
import tskit
from tests import tsutil


def _single_tree_example(L, T):
    """
    For testing numerical issues with sequence scaling
    """
    tables = tskit.TableCollection(sequence_length=L)
    tables.nodes.set_columns(
        time=np.array([0.0] * 8 + [0.1, 0.2, 0.2, 0.6, 0.8, 1.0]) * T,
        flags=np.repeat([1, 0], [8, 6]).astype("uint32"),
    )
    tables.edges.set_columns(
        left=np.repeat([0], 13),
        right=np.repeat([L], 13),
        parent=np.array(
            [8, 8, 9, 9, 10, 10, 11, 11, 11, 12, 12, 13, 13], dtype="int32"
        ),
        child=np.array([1, 2, 3, 8, 0, 7, 4, 5, 10, 6, 11, 9, 12], dtype="int32"),
    )
    tables.populations.add_row()
    tables.populations.add_row()
    tables.nodes.population = np.array(
        [0, 1, 1, 1, 0, 0, 1, 0] + [tskit.NULL] * 6, dtype="int32"
    )
    return tables.tree_sequence()


# --- prototype --- #


def _pair_coalescence_weights(
    coalescing_pairs,
    nodes_time,
):
    return coalescing_pairs


def _pair_coalescence_rates(
    coalescing_pairs,
    nodes_time,
    time_windows,
):
    """
    Estimate pair coalescence rate from empirical CDF. `coalescing_pairs` and
    `nodes_time` are assumed to have been aggregated into time bins (by
    summation/averaging respectively). The terminal bin(s) use a different
    estimator (the mean time since the start of the first terminal bin).
    """
    assert time_windows.size - 1 == coalescing_pairs.size
    assert time_windows.size - 1 == nodes_time.size
    assert np.all(np.diff(time_windows) > 0)
    assert np.isfinite(time_windows[0])
    assert time_windows[-1] == np.inf
    num_time_windows = time_windows.size - 1
    coalescence_rate = np.full(num_time_windows, np.nan)
    coalesced = 0.0
    for j in np.arange(num_time_windows, 0, -1):  # find last window containing nodes
        if not np.isnan(nodes_time[j - 1]):
            break
    for i in range(j):
        a, b = time_windows[i : i + 2]
        assert 0.0 <= coalescing_pairs[i] <= 1.0
        if i + 1 == j:
            coalescence_rate[i] = 1 / (nodes_time[i] - a)
            break
        else:
            rate = -np.log(1 - coalescing_pairs[i] / (1 - coalesced)) / (b - a)
            assert rate >= 0
            coalescence_rate[i] = abs(rate)
        coalesced += coalescing_pairs[i]
    return coalescence_rate


def _pair_coalescence_quantiles(
    coalescing_pairs,
    nodes_time,
    quantiles,
):
    """
    Estimate `quantiles` of the distribution of `nodes_time` weighted by
    `coalescing_pairs`, by inverting the empirical CDF. Nodes are assumed
    to be sorted in ascending time order.
    """
    assert nodes_time.size == coalescing_pairs.size
    assert np.all(np.diff(quantiles) > 0)
    assert np.all(np.logical_and(0 <= quantiles, quantiles <= 1))
    num_nodes = coalescing_pairs.size
    num_quantiles = quantiles.size
    output = np.full(num_quantiles, np.nan)
    i, j = 0, 0
    coalesced = 0.0
    time = -np.inf
    while i < num_nodes:
        if coalescing_pairs[i] > 0:
            coalesced += coalescing_pairs[i]
            assert nodes_time[i] > time
            time = nodes_time[i]
            while j < num_quantiles and quantiles[j] <= coalesced:
                output[j] = time
                j += 1
        i += 1
    if quantiles[-1] == 1.0:
        output[-1] = time
    return output


def _pair_coalescence_stat(
    ts,
    summary_func,
    summary_func_dim,
    summary_func_kwargs,
    sample_sets=None,
    indexes=None,
    windows=None,
    time_windows=None,
    span_normalise=True,
    pair_normalise=False,
):
    """
    Apply `summary_func(node_weights, node_times, node_order, **summary_func_kwargs)` to
    the empirical distribution of pair coalescence times for each index / window.
    """

    if sample_sets is None:
        sample_sets = [list(ts.samples())]
    for s in sample_sets:
        if len(s) == 0:
            raise ValueError("Sample sets must contain at least one element")
        if not (min(s) >= 0 and max(s) < ts.num_nodes):
            raise ValueError("Sample is out of bounds")

    drop_middle_dimension = False
    if indexes is None:
        drop_middle_dimension = True
        if len(sample_sets) == 1:
            indexes = [(0, 0)]
        elif len(sample_sets) == 2:
            indexes = [(0, 1)]
        else:
            raise ValueError(
                "Must specify indexes if there are more than two sample sets"
            )
    for i in indexes:
        if not len(i) == 2:
            raise ValueError("Sample set indexes must be length two")
        if not (min(i) >= 0 and max(i) < len(sample_sets)):
            raise ValueError("Sample set index is out of bounds")

    drop_left_dimension = False
    if windows is None:
        drop_left_dimension = True
        windows = np.array([0.0, ts.sequence_length])
    if not (isinstance(windows, np.ndarray) and windows.size > 1):
        raise ValueError("Windows must be an array of breakpoints")
    if not (windows[0] == 0.0 and windows[-1] == ts.sequence_length):
        raise ValueError("First and last window breaks must be sequence boundary")
    if not np.all(np.diff(windows) > 0):
        raise ValueError("Window breaks must be strictly increasing")

    if isinstance(time_windows, str) and time_windows == "nodes":
        nodes_map = np.arange(ts.num_nodes)
        num_time_windows = ts.num_nodes
    else:
        if not (isinstance(time_windows, np.ndarray) and time_windows.size > 1):
            raise ValueError("Time windows must be an array of breakpoints")
        if not np.all(np.diff(time_windows) > 0):
            raise ValueError("Time windows must be strictly increasing")
        if ts.time_units == tskit.TIME_UNITS_UNCALIBRATED:
            raise ValueError("Time windows require calibrated node times")
        nodes_map = np.searchsorted(time_windows, ts.nodes_time, side="right") - 1
        nodes_oob = np.logical_or(nodes_map < 0, nodes_map >= time_windows.size)
        nodes_map[nodes_oob] = tskit.NULL
        num_time_windows = time_windows.size - 1

    num_nodes = ts.num_nodes
    num_windows = windows.size - 1
    num_sample_sets = len(sample_sets)
    num_indexes = len(indexes)

    edges_child = ts.edges_child
    edges_parent = ts.edges_parent
    nodes_time = ts.nodes_time
    sequence_length = ts.sequence_length
    output_size = summary_func_dim
    samples = np.concatenate(sample_sets)

    nodes_parent = np.full(num_nodes, tskit.NULL)
    nodes_sample = np.zeros((num_nodes, num_sample_sets))
    nodes_weight = np.zeros((num_time_windows, num_indexes))
    nodes_values = np.zeros((num_time_windows, num_indexes))
    coalescing_pairs = np.zeros((num_time_windows, num_indexes))
    coalescence_time = np.zeros((num_time_windows, num_indexes))
    output = np.zeros((num_windows, output_size, num_indexes))
    visited = np.full(num_nodes, False)

    total_pairs = np.zeros(num_indexes)
    sizes = [len(s) for s in sample_sets]
    for i, (j, k) in enumerate(indexes):
        if j == k:
            total_pairs[i] = sizes[j] * (sizes[k] - 1) / 2
        else:
            total_pairs[i] = sizes[j] * sizes[k]

    for i, s in enumerate(sample_sets):  # initialize
        nodes_sample[s, i] = 1
    sample_counts = nodes_sample.copy()

    w = 0
    position = tsutil.TreePosition(ts)
    while position.interval.right < sequence_length:
        position.next()
        left, right = position.interval.left, position.interval.right
        out_range, in_range = position.out_range, position.in_range
        remainder = sequence_length - left

        for b in range(out_range.start, out_range.stop):  # edges_out
            e = out_range.order[b]
            p = edges_parent[e]
            c = edges_child[e]
            nodes_parent[c] = tskit.NULL
            inside = sample_counts[c]
            while p != tskit.NULL:
                u = nodes_map[p]
                t = nodes_time[p]
                if u != tskit.NULL:
                    outside = sample_counts[p] - sample_counts[c] - nodes_sample[p]
                    for i, (j, k) in enumerate(indexes):
                        weight = inside[j] * outside[k]
                        if j != k:
                            weight += inside[k] * outside[j]
                        coalescing_pairs[u, i] -= weight * remainder
                        coalescence_time[u, i] -= weight * remainder * t
                c, p = p, nodes_parent[p]
            p = edges_parent[e]
            while p != tskit.NULL:
                sample_counts[p] -= inside
                p = nodes_parent[p]

        for b in range(in_range.start, in_range.stop):  # edges_in
            e = in_range.order[b]
            p = edges_parent[e]
            c = edges_child[e]
            nodes_parent[c] = p
            inside = sample_counts[c]
            while p != tskit.NULL:
                sample_counts[p] += inside
                p = nodes_parent[p]
            p = edges_parent[e]
            while p != tskit.NULL:
                u = nodes_map[p]
                t = nodes_time[p]
                if u != tskit.NULL:
                    outside = sample_counts[p] - sample_counts[c] - nodes_sample[p]
                    for i, (j, k) in enumerate(indexes):
                        weight = inside[j] * outside[k]
                        if j != k:
                            weight += inside[k] * outside[j]
                        coalescing_pairs[u, i] += weight * remainder
                        coalescence_time[u, i] += weight * remainder * t
                c, p = p, nodes_parent[p]

        while w < num_windows and windows[w + 1] <= right:  # flush window
            remainder = sequence_length - windows[w + 1]
            nodes_weight[:] = coalescing_pairs[:]
            nodes_values[:] = coalescence_time[:]
            coalescing_pairs[:] = 0.0
            coalescence_time[:] = 0.0
            for c in samples:
                p = nodes_parent[c]
                while not visited[c] and p != tskit.NULL:
                    u = nodes_map[p]
                    t = nodes_time[p]
                    if u != tskit.NULL:
                        inside = sample_counts[c]
                        outside = sample_counts[p] - sample_counts[c] - nodes_sample[p]
                        for i, (j, k) in enumerate(indexes):
                            weight = inside[j] * outside[k]
                            if j != k:
                                weight += inside[k] * outside[j]
                            x = weight * remainder / 2
                            nodes_weight[u, i] -= x
                            nodes_values[u, i] -= t * x
                            coalescing_pairs[u, i] += x
                            coalescence_time[u, i] += t * x
                    visited[c] = True
                    p, c = nodes_parent[p], p
            for c in samples:
                p = nodes_parent[c]
                while visited[c] and p != tskit.NULL:
                    visited[c] = False
                    p, c = nodes_parent[p], p
            for i in range(num_indexes):  # normalise values
                nonzero = nodes_weight[:, i] > 0
                nodes_values[nonzero, i] /= nodes_weight[nonzero, i]
                nodes_values[~nonzero, i] = np.nan
            if span_normalise:
                nodes_weight /= windows[w + 1] - windows[w]
            if pair_normalise:
                nodes_weight /= total_pairs[np.newaxis, :]
            for i in range(num_indexes):  # apply function to empirical distribution
                output[w, :, i] = summary_func(
                    nodes_weight[:, i],
                    nodes_values[:, i],
                    **summary_func_kwargs,
                )
            w += 1

    output = output.transpose(0, 2, 1)
    if drop_middle_dimension:
        output = output.squeeze(1)
    if drop_left_dimension:
        output = output.squeeze(0)

    return output


def proto_pair_coalescence_counts(
    ts,
    sample_sets=None,
    indexes=None,
    windows=None,
    span_normalise=True,
    pair_normalise=False,
    time_windows="nodes",
):
    """
    Prototype for ts.pair_coalescence_counts.

    Calculate the number of coalescing sample pairs per node, summed over
    trees and weighted by tree span.

    The number of coalescing pairs may be calculated within or between the
    non-overlapping lists of samples contained in `sample_sets`. In the
    latter case, pairs are counted if they have exactly one member in each
    of two sample sets. If `sample_sets` is omitted, a single group
    containing all samples is assumed.

    The argument `indexes` may be used to specify which pairs of sample
    sets to compute the statistic between, and in what order. If
    `indexes=None`, then `indexes` is assumed to equal `[(0,0)]` for a
    single sample set and `[(0,1)]` for two sample sets. For more than two
    sample sets, `indexes` must be explicitly passed.

    The argument `time_windows` may be used to count coalescence
    events within time intervals (if an array of breakpoints is supplied)
    rather than for individual nodes (the default).

    The output array has dimension `(windows, indexes, nodes)` with
    dimensions dropped when the corresponding argument is set to None.

    :param list sample_sets: A list of lists of Node IDs, specifying the
        groups of nodes to compute the statistic with, or None.
    :param list indexes: A list of 2-tuples, or None.
    :param list windows: An increasing list of breakpoints between the
        sequence windows to compute the statistic in, or None.
    :param bool span_normalise: Whether to divide the result by the span of
        the window (defaults to True).
    :param bool pair_normalise: Whether to divide the result by the total
        number of pairs for a given index (defaults to False).
    :param time_windows: Either a string "nodes" or an increasing
        list of breakpoints between time intervals.
    """

    if isinstance(time_windows, str) and time_windows == "nodes":
        summary_func_dim = ts.num_nodes
    else:
        if not (isinstance(time_windows, np.ndarray) and time_windows.size > 1):
            raise ValueError("Time windows must be an array of breakpoints")
        if not np.all(np.diff(time_windows) > 0):
            raise ValueError("Time windows must be strictly increasing")
        if ts.time_units == tskit.TIME_UNITS_UNCALIBRATED:
            raise ValueError("Time windows require calibrated node times")
        summary_func_dim = time_windows.size - 1

    summary_func = _pair_coalescence_weights
    summary_func_kwargs = {}

    return _pair_coalescence_stat(
        ts,
        summary_func=summary_func,
        summary_func_dim=summary_func_dim,
        summary_func_kwargs=summary_func_kwargs,
        sample_sets=sample_sets,
        indexes=indexes,
        windows=windows,
        time_windows=time_windows,
        span_normalise=span_normalise,
        pair_normalise=pair_normalise,
    )


def proto_pair_coalescence_rates(
    ts,
    time_windows,
    sample_sets=None,
    indexes=None,
    windows=None,
):
    r"""
    Prototype for ts.pair_coalescence_rates.

    Estimate the rate at which pairs of samples coalesce within time windows,
    from the empirical CDF of pair coalescence times.  Assuming that pair
    coalescence events follow a nonhomogeneous Poisson process, the empirical
    rate for a time window :math:`[a, b)` where `ecdf(b) < 1` is,

    ..math:

        log(1 - \frac{ecdf(b) - ecdf(a)}{1 - ecdf(a)}) / (a - b)

    If the last coalescence event is within `[a, b)` so that `ecdf(b) = 1`, then
    an estimate of the empirical rate is

    ..math:

        (\mathbb{E}[t | t > a] - a)^{-1}

    where :math:`\mathbb{E}[t | t < a]` is the average pair coalescence time
    conditional on coalescence after the start of the last epoch.

    The first breakpoint in `time_windows` must start at the age of the
    samples, and the last must end at infinity.

    Pair coalescence rates may be calculated within or between the
    non-overlapping lists of samples contained in `sample_sets`. In the
    latter case, pairs are counted if they have exactly one member in each
    of two sample sets. If `sample_sets` is omitted, a single group
    containing all samples is assumed.

    The argument `indexes` may be used to specify which pairs of sample
    sets to compute the statistic between, and in what order. If
    `indexes=None`, then `indexes` is assumed to equal `[(0,0)]` for a
    single sample set and `[(0,1)]` for two sample sets. For more than two
    sample sets, `indexes` must be explicitly passed.

    The output array has dimension `(windows, indexes, time_windows)` with
    dimensions dropped when the corresponding argument is set to None.

    :param time_windows: An increasing list of breakpoints between time
        intervals, starting at the age of the samples and ending at
        infinity.
    :param list sample_sets: A list of lists of Node IDs, specifying the
        groups of nodes to compute the statistic with, or None.
    :param list indexes: A list of 2-tuples, or None.
    :param list windows: An increasing list of breakpoints between the
        sequence windows to compute the statistic in, or None.
    """
    # TODO^^^

    if not (isinstance(time_windows, np.ndarray) and time_windows.size > 1):
        raise ValueError("Time windows must be an array of breakpoints")
    if not np.all(np.diff(time_windows) > 0):
        raise ValueError("Time windows must be strictly increasing")
    if ts.time_units == tskit.TIME_UNITS_UNCALIBRATED:
        raise ValueError("Time windows require calibrated node times")

    summary_func = _pair_coalescence_rates
    summary_func_dim = time_windows.size - 1
    summary_func_kwargs = {"time_windows": time_windows}

    return _pair_coalescence_stat(
        ts,
        summary_func=summary_func,
        summary_func_dim=summary_func_dim,
        summary_func_kwargs=summary_func_kwargs,
        sample_sets=sample_sets,
        indexes=indexes,
        windows=windows,
        time_windows=time_windows,
        span_normalise=True,
        pair_normalise=True,
    )


def proto_pair_coalescence_quantiles(
    ts,
    quantiles,
    sample_sets=None,
    indexes=None,
    windows=None,
):
    """
    Prototype for ts.pair_coalescence_quantiles.

    Estimate quantiles of pair coalescence times by inverting the empirical
    CDF. This is equivalent to the "inverted_cdf" method of `numpy.quantile`
    applied to node times, with weights proportional to the number of
    coalescing pairs per node (averaged over trees). The weights are calculated
    using `pair_coalescence_counts`.

    Quantiles of pair coalescence times may be calculated within or
    between the non-overlapping lists of samples contained in `sample_sets`. In
    the latter case, pairs are counted if they have exactly one member in each
    of two sample sets. If `sample_sets` is omitted, a single group containing
    all samples is assumed.

    The argument `indexes` may be used to specify which pairs of sample sets to
    compute coalescences between, and in what order. If `indexes=None`, then
    `indexes` is assumed to equal `[(0,0)]` for a single sample set and
    `[(0,1)]` for two sample sets. For more than two sample sets, `indexes`
    must be explicitly passed.

    The output array has dimension `(windows, indexes, quantiles)` with
    dimensions dropped when the corresponding argument is set to None.

    :param quantiles: A list of breakpoints between [0, 1].
    :param list sample_sets: A list of lists of Node IDs, specifying the
        groups of nodes to compute the statistic with, or None.
    :param list indexes: A list of 2-tuples, or None.
    :param list windows: An increasing list of breakpoints between the
        sequence windows to compute the statistic in, or None.
    """

    if not isinstance(quantiles, np.ndarray):
        raise ValueError("Quantiles must be an array of breakpoints")
    if not np.all(np.logical_and(quantiles >= 0, quantiles <= 1.0)):
        raise ValueError("Quantiles must be in [0, 1]")

    summary_func = _pair_coalescence_quantiles
    summary_func_dim = quantiles.size
    summary_func_kwargs = {"quantiles": quantiles}
    time_windows = np.append(
        np.unique(ts.nodes_time), np.inf
    )  # sort nodes in time order

    return _pair_coalescence_stat(
        ts,
        summary_func=summary_func,
        summary_func_dim=summary_func_dim,
        summary_func_kwargs=summary_func_kwargs,
        sample_sets=sample_sets,
        indexes=indexes,
        windows=windows,
        time_windows=time_windows,
        span_normalise=True,
        pair_normalise=True,
    )


# --- testing --- #


def naive_pair_coalescence_counts(ts, sample_set_0, sample_set_1):
    """
    Naive implementation of ts.pair_coalescence_counts.

    Count pairwise coalescences tree by tree, by enumerating nodes in each
    tree. For a binary node, the number of pairs of samples that coalesce in a
    given node is the product of the number of samples subtended by the left
    and right child. For higher arities, the count is summed over all possible
    pairs of children.
    """
    output = np.zeros(ts.num_nodes)
    for t in ts.trees():
        sample_counts = np.zeros((ts.num_nodes, 2), dtype=np.int32)
        pair_counts = np.zeros(ts.num_nodes)
        for p in t.postorder():
            samples = list(t.samples(p))
            sample_counts[p, 0] = np.intersect1d(samples, sample_set_0).size
            sample_counts[p, 1] = np.intersect1d(samples, sample_set_1).size
            for i, j in itertools.combinations(t.children(p), 2):
                pair_counts[p] += sample_counts[i, 0] * sample_counts[j, 1]
                pair_counts[p] += sample_counts[i, 1] * sample_counts[j, 0]
        output += pair_counts * t.span
    return output


def _numpy_weighted_quantile(values, weights, quantiles):
    """
    Requires numpy 2.0. Enforcing `weights > 0` avoids odd behaviour where
    numpy assigns the 0th quantile to the sample minimum, even if this minimum
    has zero weight.
    """
    assert np.all(weights >= 0.0)
    return np.quantile(
        values[weights > 0],
        quantiles,
        weights=weights[weights > 0] / weights.sum(),
        method="inverted_cdf",
    )


def _numpy_hazard_rate(values, weights, breaks):
    """
    Estimate hazard rate from empirical CDF over intervals
    """
    assert np.all(weights >= 0)
    assert np.all(np.diff(breaks) >= 0)
    assert np.isfinite(breaks[0])  # should equal sample time
    assert ~np.isfinite(breaks[-1])
    assert np.sum(weights) < 1.0 or np.isclose(np.sum(weights), 1.0)
    values = values[weights > 0]
    weights = weights[weights > 0]
    assert breaks[0] < np.min(values)
    max_value = np.max(values)
    rates = np.full(breaks.size - 1, np.nan)
    for i, (a, b) in enumerate(zip(breaks[:-1], breaks[1:])):
        if a < max_value <= b:  # terminal window
            keep = values >= a
            mean = np.sum(values[keep] * weights[keep]) / np.sum(weights[keep])
            rates[i] = 1.0 / (mean - a)
            break
        else:
            wa = np.sum(weights[values < a])
            wb = np.sum(weights[values < b])
            rates[i] = np.log(1 - (wb - wa) / (1 - wa)) / (b - a)
            assert rates[i] <= 0.0
            rates[i] = abs(rates[i])
    return rates


def convert_to_nonsuccinct(ts):
    """
    Give the edges and internal nodes in each tree distinct IDs
    """
    tables = tskit.TableCollection(sequence_length=ts.sequence_length)
    for _ in range(ts.num_populations):
        tables.populations.add_row()
    nodes_count = 0
    for n in ts.samples():
        tables.nodes.add_row(
            time=ts.nodes_time[n],
            flags=ts.nodes_flags[n],
            population=ts.nodes_population[n],
        )
        nodes_count += 1
    for t in ts.trees():
        nodes_map = {n: n for n in ts.samples()}
        for n in t.nodes():
            if t.num_samples(n) > 1:
                tables.nodes.add_row(
                    time=ts.nodes_time[n],
                    flags=ts.nodes_flags[n],
                    population=ts.nodes_population[n],
                )
                nodes_map[n] = nodes_count
                nodes_count += 1
        for n in t.nodes():
            if t.edge(n) != tskit.NULL:
                tables.edges.add_row(
                    parent=nodes_map[t.parent(n)],
                    child=nodes_map[n],
                    left=t.interval.left,
                    right=t.interval.right,
                )
    tables.sort()
    ts_unroll = tables.tree_sequence()
    assert nodes_count == ts_unroll.num_nodes
    return ts_unroll


class TestCoalescingPairsOneTree:
    """
    Test against worked example (single tree)
    """

    def example_ts(self):
        """
        10.0┊         13      ┊
            ┊       ┏━━┻━━┓   ┊
         8.0┊      12     ┃   ┊
            ┊     ┏━┻━┓   ┃   ┊
         6.0┊    11   ┃   ┃   ┊
            ┊  ┏━━╋━┓ ┃   ┃   ┊
         2.0┊ 10  ┃ ┃ ┃   9   ┊
            ┊ ┏┻┓ ┃ ┃ ┃  ┏┻━┓ ┊
         1.0┊ ┃ ┃ ┃ ┃ ┃  8  ┃ ┊
            ┊ ┃ ┃ ┃ ┃ ┃ ┏┻┓ ┃ ┊
         0.0┊ 0 7 4 5 6 1 2 3 ┊
            ┊ A A A A B B B B ┊
        """
        tables = tskit.TableCollection(sequence_length=100)
        tables.nodes.set_columns(
            time=np.array([0] * 8 + [1, 2, 2, 6, 8, 10]),
            flags=np.repeat([1, 0], [8, 6]).astype("uint32"),
        )
        tables.edges.set_columns(
            left=np.repeat([0], 13),
            right=np.repeat([100], 13),
            parent=np.array(
                [8, 8, 9, 9, 10, 10, 11, 11, 11, 12, 12, 13, 13], dtype="int32"
            ),
            child=np.array([1, 2, 3, 8, 0, 7, 4, 5, 10, 6, 11, 9, 12], dtype="int32"),
        )
        tables.populations.add_row()
        tables.populations.add_row()
        tables.nodes.population = np.array(
            [0, 1, 1, 1, 0, 0, 1, 0] + [tskit.NULL] * 6, dtype="int32"
        )
        return tables.tree_sequence()

    def test_total_pairs(self):
        """
        ┊         15 pairs ┊
        ┊       ┏━━┻━━┓    ┊
        ┊       4     ┃    ┊
        ┊     ┏━┻━┓   ┃    ┊
        ┊     5   ┃   ┃    ┊
        ┊  ┏━━╋━┓ ┃   ┃    ┊
        ┊  1  ┃ ┃ ┃   2    ┊
        ┊ ┏┻┓ ┃ ┃ ┃  ┏┻━┓  ┊
        ┊ ┃ ┃ ┃ ┃ ┃  1  ┃  ┊
        ┊ ┃ ┃ ┃ ┃ ┃ ┏┻┓ ┃  ┊
        ┊ 0 0 0 0 0 0 0 0  ┊
        """
        ts = self.example_ts()
        check = np.array([0.0] * 8 + [1, 2, 1, 5, 4, 15])
        implm = ts.pair_coalescence_counts()
        np.testing.assert_allclose(implm, check)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(ts)
        np.testing.assert_allclose(proto, check)

    def test_population_pairs(self):
        """
        ┊ AA       0 pairs ┊ AB      12 pairs ┊ BB       3 pairs ┊
        ┊       ┏━━┻━━┓    ┊       ┏━━┻━━┓    ┊       ┏━━┻━━┓    ┊
        ┊       0     ┃    ┊       4     ┃    ┊       0     ┃    ┊
        ┊     ┏━┻━┓   ┃    ┊     ┏━┻━┓   ┃    ┊     ┏━┻━┓   ┃    ┊
        ┊     5   ┃   ┃    ┊     0   ┃   ┃    ┊     0   ┃   ┃    ┊
        ┊  ┏━━╋━┓ ┃   ┃    ┊  ┏━━╋━┓ ┃   ┃    ┊  ┏━━╋━┓ ┃   ┃    ┊
        ┊  1  ┃ ┃ ┃   0    ┊  0  ┃ ┃ ┃   0    ┊  0  ┃ ┃ ┃   2    ┊
        ┊ ┏┻┓ ┃ ┃ ┃  ┏┻━┓  ┊ ┏┻┓ ┃ ┃ ┃  ┏┻━┓  ┊ ┏┻┓ ┃ ┃ ┃  ┏┻━┓  ┊
        ┊ ┃ ┃ ┃ ┃ ┃  0  ┃  ┊ ┃ ┃ ┃ ┃ ┃  0  ┃  ┊ ┃ ┃ ┃ ┃ ┃  1  ┃  ┊
        ┊ ┃ ┃ ┃ ┃ ┃ ┏┻┓ ┃  ┊ ┃ ┃ ┃ ┃ ┃ ┏┻┓ ┃  ┊ ┃ ┃ ┃ ┃ ┃ ┏┻┓ ┃  ┊
        ┊ A A A A B B B B  ┊ A A A A B B B B  ┊ A A A A B B B B  ┊
        """
        ts = self.example_ts()
        ss0 = np.flatnonzero(ts.nodes_population == 0)
        ss1 = np.flatnonzero(ts.nodes_population == 1)
        indexes = [(0, 0), (0, 1), (1, 1)]
        implm = ts.pair_coalescence_counts(sample_sets=[ss0, ss1], indexes=indexes)
        check = np.full(implm.shape, np.nan)
        check[0] = np.array([0.0] * 8 + [0, 0, 1, 5, 0, 0])
        check[1] = np.array([0.0] * 8 + [0, 0, 0, 0, 4, 12])
        check[2] = np.array([0.0] * 8 + [1, 2, 0, 0, 0, 3])
        np.testing.assert_allclose(implm, check)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(
            ts, sample_sets=[ss0, ss1], indexes=indexes
        )
        np.testing.assert_allclose(proto, check)

    def test_internal_samples(self):
        """
        ┊          Not     ┊         24 pairs ┊
        ┊       ┏━━┻━━┓    ┊       ┏━━┻━━┓    ┊
        ┊       N     ┃    ┊       5     ┃    ┊
        ┊     ┏━┻━┓   ┃    ┊     ┏━┻━┓   ┃    ┊
        ┊     S   ┃   ┃    ┊     5   ┃   ┃    ┊
        ┊  ┏━━╋━┓ ┃   ┃    ┊  ┏━━╋━┓ ┃   ┃    ┊
        ┊  N  ┃ ┃ ┃   Samp ┊  1  ┃ ┃ ┃   2    ┊
        ┊ ┏┻┓ ┃ ┃ ┃  ┏┻━┓  ┊ ┏┻┓ ┃ ┃ ┃  ┏┻━┓  ┊
        ┊ ┃ ┃ ┃ ┃ ┃  N  ┃  ┊ ┃ ┃ ┃ ┃ ┃  1  ┃  ┊
        ┊ ┃ ┃ ┃ ┃ ┃ ┏┻┓ ┃  ┊ ┃ ┃ ┃ ┃ ┃ ┏┻┓ ┃  ┊
        ┊ S S S S S S S S  ┊ 0 0 0 0 0 0 0 0  ┊
        """
        ts = self.example_ts()
        tables = ts.dump_tables()
        nodes_flags = tables.nodes.flags.copy()
        nodes_flags[9] = tskit.NODE_IS_SAMPLE
        nodes_flags[11] = tskit.NODE_IS_SAMPLE
        tables.nodes.flags = nodes_flags
        ts = tables.tree_sequence()
        assert ts.num_samples == 10
        implm = ts.pair_coalescence_counts(span_normalise=False)
        check = np.array([0] * 8 + [1, 2, 1, 5, 5, 24]) * ts.sequence_length
        np.testing.assert_allclose(implm, check)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(ts, span_normalise=False)
        np.testing.assert_allclose(proto, check)

    def test_windows(self):
        ts = self.example_ts()
        check = np.array([0.0] * 8 + [1, 2, 1, 5, 4, 15]) * ts.sequence_length / 2
        implm = ts.pair_coalescence_counts(
            windows=np.linspace(0, ts.sequence_length, 3), span_normalise=False
        )
        np.testing.assert_allclose(implm[0], check)
        np.testing.assert_allclose(implm[1], check)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(
            ts, windows=np.linspace(0, ts.sequence_length, 3), span_normalise=False
        )
        np.testing.assert_allclose(proto[0], check)
        np.testing.assert_allclose(proto[1], check)

    def test_time_windows(self):
        """
           ┊         15 pairs ┊
           ┊       ┏━━┻━━┓    ┊
           ┊       4     ┃    ┊
        7.0┊-----┏━┻━┓---┃----┊
           ┊     5   ┃   ┃    ┊
        5.0┊--┏━━╋━┓-┃---┃----┊
           ┊  1  ┃ ┃ ┃   2    ┊
           ┊ ┏┻┓ ┃ ┃ ┃  ┏┻━┓  ┊
           ┊ ┃ ┃ ┃ ┃ ┃  1  ┃  ┊
           ┊ ┃ ┃ ┃ ┃ ┃ ┏┻┓ ┃  ┊
        0.0┊ 0 0 0 0 0 0 0 0  ┊
        """
        ts = self.example_ts()
        time_windows = np.array([0.0, 5.0, 7.0, np.inf])
        check = np.array([4, 5, 19]) * ts.sequence_length
        implm = ts.pair_coalescence_counts(
            span_normalise=False, time_windows=time_windows
        )
        np.testing.assert_allclose(implm, check)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(
            ts, span_normalise=False, time_windows=time_windows
        )
        np.testing.assert_allclose(proto, check)

    def test_pair_normalise(self):
        ts = self.example_ts()
        ss0 = np.flatnonzero(ts.nodes_population == 0)
        ss1 = np.flatnonzero(ts.nodes_population == 1)
        indexes = [(0, 0), (0, 1), (1, 1)]
        implm = ts.pair_coalescence_counts(
            sample_sets=[ss0, ss1],
            indexes=indexes,
            pair_normalise=True,
        )
        check = np.full(implm.shape, np.nan)
        check[0] = np.array([0.0] * 8 + [0, 0, 1, 5, 0, 0])
        check[1] = np.array([0.0] * 8 + [0, 0, 0, 0, 4, 12])
        check[2] = np.array([0.0] * 8 + [1, 2, 0, 0, 0, 3])
        total_pairs = np.array([6, 16, 6])
        check /= total_pairs[:, np.newaxis]
        np.testing.assert_allclose(implm, check)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(
            ts,
            sample_sets=[ss0, ss1],
            indexes=indexes,
            pair_normalise=True,
        )
        np.testing.assert_allclose(proto, check)

    def test_multiple_roots(self):
        ts = self.example_ts().decapitate(6.0)
        implm = ts.pair_coalescence_counts(pair_normalise=True)
        total_pairs = ts.num_samples * (ts.num_samples - 1) / 2
        check = np.array([0.0] * 8 + [1, 2, 1, 5, 0, 0, 0, 0])
        check /= total_pairs
        np.testing.assert_allclose(implm, check)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(ts, pair_normalise=True)
        np.testing.assert_allclose(proto, check)


class TestCoalescingPairsTwoTree:
    """
    Test against worked example (two trees)
    """

    def example_ts(self, S, L):
        """
           0         S         L
        4.0┊   7     ┊   7     ┊
           ┊ ┏━┻━┓   ┊ ┏━┻━┓   ┊
        3.0┊ ┃   6   ┊ ┃   ┃   ┊
           ┊ ┃ ┏━┻┓  ┊ ┃   ┃   ┊
        2.0┊ ┃ ┃  5  ┊ ┃   5   ┊
           ┊ ┃ ┃ ┏┻┓ ┊ ┃  ┏┻━┓ ┊
        1.0┊ ┃ ┃ ┃ ┃ ┊ ┃  4  ┃ ┊
           ┊ ┃ ┃ ┃ ┃ ┊ ┃ ┏┻┓ ┃ ┊
        0.0┊ 0 1 2 3 ┊ 0 1 2 3 ┊
             A A B B   A A B B
        """
        tables = tskit.TableCollection(sequence_length=L)
        tables.nodes.set_columns(
            time=np.array([0, 0, 0, 0, 1.0, 2.0, 3.0, 4.0]),
            flags=np.array([1, 1, 1, 1, 0, 0, 0, 0], dtype="uint32"),
        )
        tables.edges.set_columns(
            left=np.array([S, S, 0, 0, S, 0, 0, 0, S, 0]),
            right=np.array([L, L, S, L, L, S, S, L, L, S]),
            parent=np.array([4, 4, 5, 5, 5, 6, 6, 7, 7, 7], dtype="int32"),
            child=np.array([1, 2, 2, 3, 4, 1, 5, 0, 5, 6], dtype="int32"),
        )
        return tables.tree_sequence()

    def test_total_pairs(self):
        """
        ┊   3 pairs   3     ┊
        ┊ ┏━┻━┓     ┏━┻━┓   ┊
        ┊ ┃   2     ┃   ┃   ┊
        ┊ ┃ ┏━┻┓    ┃   ┃   ┊
        ┊ ┃ ┃  1    ┃   2   ┊
        ┊ ┃ ┃ ┏┻┓   ┃  ┏┻━┓ ┊
        ┊ ┃ ┃ ┃ ┃   ┃  1  ┃ ┊
        ┊ ┃ ┃ ┃ ┃   ┃ ┏┻┓ ┃ ┊
        ┊ 0 0 0 0   0 0 0 0 ┊
        0         S         L
        """
        L, S = 1e8, 1.0
        ts = self.example_ts(S, L)
        implm = ts.pair_coalescence_counts(span_normalise=False)
        check = np.array([0] * 4 + [1 * (L - S), 2 * (L - S) + 1 * S, 2 * S, 3 * L])
        np.testing.assert_allclose(implm, check)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(ts, span_normalise=False)
        np.testing.assert_allclose(proto, check)

    def test_population_pairs(self):
        """
        ┊AA                 ┊AB                 ┊BB                 ┊
        ┊   1 pairs   1     ┊   2 pairs   2     ┊   0 pairs   0     ┊
        ┊ ┏━┻━┓     ┏━┻━┓   ┊ ┏━┻━┓     ┏━┻━┓   ┊ ┏━┻━┓     ┏━┻━┓   ┊
        ┊ ┃   0     ┃   ┃   ┊ ┃   2     ┃   ┃   ┊ ┃   0     ┃   ┃   ┊
        ┊ ┃ ┏━┻┓    ┃   ┃   ┊ ┃ ┏━┻┓    ┃   ┃   ┊ ┃ ┏━┻┓    ┃   ┃   ┊
        ┊ ┃ ┃  0    ┃   0   ┊ ┃ ┃  0    ┃   1   ┊ ┃ ┃  1    ┃   1   ┊
        ┊ ┃ ┃ ┏┻┓   ┃  ┏┻━┓ ┊ ┃ ┃ ┏┻┓   ┃  ┏┻━┓ ┊ ┃ ┃ ┏┻┓   ┃  ┏┻━┓ ┊
        ┊ ┃ ┃ ┃ ┃   ┃  0  ┃ ┊ ┃ ┃ ┃ ┃   ┃  1  ┃ ┊ ┃ ┃ ┃ ┃   ┃  0  ┃ ┊
        ┊ ┃ ┃ ┃ ┃   ┃ ┏┻┓ ┃ ┊ ┃ ┃ ┃ ┃   ┃ ┏┻┓ ┃ ┊ ┃ ┃ ┃ ┃   ┃ ┏┻┓ ┃ ┊
        ┊ A A B B   A A B B ┊ A A B B   A A B B ┊ A A B B   A A B B ┊
        0         S         L         S         L         S         L
        """
        L, S = 1e8, 1.0
        ts = self.example_ts(S, L)
        indexes = [(0, 0), (0, 1), (1, 1)]
        implm = ts.pair_coalescence_counts(
            sample_sets=[[0, 1], [2, 3]], indexes=indexes, span_normalise=False
        )
        check = np.empty(implm.shape)
        check[0] = np.array([0] * 4 + [0, 0, 0, 1 * L])
        check[1] = np.array([0] * 4 + [1 * (L - S), 1 * (L - S), 2 * S, 2 * L])
        check[2] = np.array([0] * 4 + [0, 1 * L, 0, 0])
        np.testing.assert_allclose(implm, check)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(
            ts, sample_sets=[[0, 1], [2, 3]], indexes=indexes, span_normalise=False
        )
        np.testing.assert_allclose(proto, check)

    def test_internal_samples(self):
        """
        ┊   Not       N     ┊   4 pairs   4     ┊
        ┊ ┏━┻━┓     ┏━┻━┓   ┊ ┏━┻━┓     ┏━┻━┓   ┊
        ┊ ┃   N     ┃   ┃   ┊ ┃   3     ┃   ┃   ┊
        ┊ ┃ ┏━┻┓    ┃   ┃   ┊ ┃ ┏━┻┓    ┃   ┃   ┊
        ┊ ┃ ┃  Samp ┃   S   ┊ ┃ ┃  1    ┃   2   ┊
        ┊ ┃ ┃ ┏┻┓   ┃  ┏┻━┓ ┊ ┃ ┃ ┏┻┓   ┃  ┏┻━┓ ┊
        ┊ ┃ ┃ ┃ ┃   ┃  N  ┃ ┊ ┃ ┃ ┃ ┃   ┃  1  ┃ ┊
        ┊ ┃ ┃ ┃ ┃   ┃ ┏┻┓ ┃ ┊ ┃ ┃ ┃ ┃   ┃ ┏┻┓ ┃ ┊
        ┊ S S S S   S S S S ┊ 0 0 0 0   0 0 0 0 ┊
        """
        L, S = 200, 100
        ts = self.example_ts(S, L)
        tables = ts.dump_tables()
        nodes_flags = tables.nodes.flags.copy()
        nodes_flags[5] = tskit.NODE_IS_SAMPLE
        tables.nodes.flags = nodes_flags
        ts = tables.tree_sequence()
        assert ts.num_samples == 5
        implm = ts.pair_coalescence_counts(span_normalise=False)
        check = np.array([0.0] * 4 + [(L - S), S + 2 * (L - S), 3 * S, 4 * L])
        np.testing.assert_allclose(implm, check)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(ts, span_normalise=False)
        np.testing.assert_allclose(proto, check)

    def test_windows(self):
        """
        ┊   3 pairs   3     ┊
        ┊ ┏━┻━┓     ┏━┻━┓   ┊
        ┊ ┃   2     ┃   ┃   ┊
        ┊ ┃ ┏━┻┓    ┃   ┃   ┊
        ┊ ┃ ┃  1    ┃   2   ┊
        ┊ ┃ ┃ ┏┻┓   ┃  ┏┻━┓ ┊
        ┊ ┃ ┃ ┃ ┃   ┃  1  ┃ ┊
        ┊ ┃ ┃ ┃ ┃   ┃ ┏┻┓ ┃ ┊
        ┊ 0 0 0 0   0 0 0 0 ┊
        0         S         L
        """
        L, S = 200, 100
        ts = self.example_ts(S, L)
        windows = np.array(list(ts.breakpoints()))
        check_0 = np.array([0.0] * 4 + [0, 1, 2, 3]) * S
        check_1 = np.array([0.0] * 4 + [1, 2, 0, 3]) * (L - S)
        implm = ts.pair_coalescence_counts(windows=windows, span_normalise=False)
        np.testing.assert_allclose(implm[0], check_0)
        np.testing.assert_allclose(implm[1], check_1)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(ts, windows=windows, span_normalise=False)
        np.testing.assert_allclose(proto[0], check_0)
        np.testing.assert_allclose(proto[1], check_1)

    def test_time_windows(self):
        """
           ┊   3 pairs   3     ┊
        3.5┊-┏━┻━┓---┊-┏━┻━┓---┊
           ┊ ┃   2   ┊ ┃   ┃   ┊
           ┊ ┃ ┏━┻┓  ┊ ┃   ┃   ┊
           ┊ ┃ ┃  1  ┊ ┃   2   ┊
        1.5┊-┃-┃-┏┻┓-┊-┃--┏┻━┓-┊
           ┊ ┃ ┃ ┃ ┃ ┊ ┃  1  ┃ ┊
           ┊ ┃ ┃ ┃ ┃ ┊ ┃ ┏┻┓ ┃ ┊
        0.0┊ 0 0 0 0 ┊ 0 0 0 0 ┊
           0         S         L
        """
        L, S = 200, 100
        ts = self.example_ts(S, L)
        time_windows = np.array([0.0, 1.5, 3.5, np.inf])
        windows = np.array(list(ts.breakpoints()))
        check_0 = np.array([0.0, 3.0, 3.0]) * S
        check_1 = np.array([1.0, 2.0, 3.0]) * (L - S)
        implm = ts.pair_coalescence_counts(
            span_normalise=False,
            windows=windows,
            time_windows=time_windows,
        )
        np.testing.assert_allclose(implm[0], check_0)
        np.testing.assert_allclose(implm[1], check_1)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(
            ts,
            span_normalise=False,
            windows=windows,
            time_windows=time_windows,
        )
        np.testing.assert_allclose(proto[0], check_0)
        np.testing.assert_allclose(proto[1], check_1)

    def test_pair_normalise(self):
        L, S = 200, 100
        ts = self.example_ts(S, L)
        indexes = [(0, 0), (0, 1), (1, 1)]
        implm = ts.pair_coalescence_counts(
            sample_sets=[[0, 1], [2, 3]],
            indexes=indexes,
            span_normalise=False,
            pair_normalise=True,
        )
        check = np.empty(implm.shape)
        check[0] = np.array([0] * 4 + [0, 0, 0, 1 * L])
        check[1] = np.array([0] * 4 + [1 * (L - S), 1 * (L - S), 2 * S, 2 * L])
        check[2] = np.array([0] * 4 + [0, 1 * L, 0, 0])
        total_pairs = np.array([1, 4, 1])
        check /= total_pairs[:, np.newaxis]
        np.testing.assert_allclose(implm, check)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(
            ts,
            sample_sets=[[0, 1], [2, 3]],
            indexes=indexes,
            span_normalise=False,
            pair_normalise=True,
        )
        np.testing.assert_allclose(proto, check)

    def test_multiple_roots(self):
        L, S = 200, 100
        ts = self.example_ts(S, L).decapitate(2.0)
        implm = ts.pair_coalescence_counts(pair_normalise=True, span_normalise=False)
        total_pairs = ts.num_samples * (ts.num_samples - 1) / 2
        check = np.array([0.0] * 4 + [1 * (L - S), 2 * (L - S) + 1 * S, 0, 0, 0, 0])
        check /= total_pairs
        np.testing.assert_allclose(implm, check)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(
            ts, pair_normalise=True, span_normalise=False
        )
        np.testing.assert_allclose(proto, check)


class TestCoalescingPairsSimulated:
    """
    Test against a naive implementation on simulated data.
    """

    @tests.cached_example
    def example_ts(self):
        n = 10
        model = msprime.BetaCoalescent(alpha=1.5)  # polytomies
        tables = msprime.sim_ancestry(
            samples=n,
            recombination_rate=1e-8,
            sequence_length=1e6,
            population_size=1e4,
            random_seed=1024,
            model=model,
        ).dump_tables()
        tables.populations.add_row(metadata={"name": "foo", "description": "bar"})
        tables.populations.add_row(metadata={"name": "bar", "description": "foo"})
        tables.nodes.population = np.repeat(
            [0, 1, 2, tskit.NULL],
            [n, n // 2, n - n // 2, tables.nodes.num_rows - 2 * n],
        ).astype("int32")
        ts = tables.tree_sequence()
        assert ts.num_trees > 1
        return ts

    @staticmethod
    def _check_total_pairs(ts, windows):
        samples = list(ts.samples())
        implm = ts.pair_coalescence_counts(windows=windows, span_normalise=False)
        dim = (windows.size - 1, ts.num_nodes)
        check = np.full(dim, np.nan)
        for w, (a, b) in enumerate(zip(windows[:-1], windows[1:])):
            tsw = ts.keep_intervals(np.array([[a, b]]), simplify=False)
            check[w] = naive_pair_coalescence_counts(tsw, samples, samples) / 2
        np.testing.assert_allclose(implm, check)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(ts, windows=windows, span_normalise=False)
        np.testing.assert_allclose(proto, check)

    @staticmethod
    def _check_subset_pairs(ts, windows):
        ss0 = np.flatnonzero(ts.nodes_population == 0)
        ss1 = np.flatnonzero(ts.nodes_population == 1)
        idx = [(0, 1), (1, 1), (0, 0)]
        implm = ts.pair_coalescence_counts(
            sample_sets=[ss0, ss1], indexes=idx, windows=windows, span_normalise=False
        )
        dim = (windows.size - 1, len(idx), ts.num_nodes)
        check = np.full(dim, np.nan)
        for w, (a, b) in enumerate(zip(windows[:-1], windows[1:])):
            tsw = ts.keep_intervals(np.array([[a, b]]), simplify=False)
            check[w, 0] = naive_pair_coalescence_counts(tsw, ss0, ss1)
            check[w, 1] = naive_pair_coalescence_counts(tsw, ss1, ss1) / 2
            check[w, 2] = naive_pair_coalescence_counts(tsw, ss0, ss0) / 2
        np.testing.assert_allclose(implm, check)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(
            ts,
            sample_sets=[ss0, ss1],
            indexes=idx,
            windows=windows,
            span_normalise=False,
        )
        np.testing.assert_allclose(proto, check)

    def test_sequence(self):
        ts = self.example_ts()
        windows = np.array([0.0, ts.sequence_length])
        self._check_total_pairs(ts, windows)
        self._check_subset_pairs(ts, windows)

    def test_missing_interval(self):
        """
        test case where three segments have all samples missing
        """
        ts = self.example_ts()
        windows = np.array([0.0, ts.sequence_length])
        intervals = np.array([[0.0, 0.1], [0.4, 0.6], [0.9, 1.0]]) * ts.sequence_length
        ts = ts.delete_intervals(intervals)
        self._check_total_pairs(ts, windows)
        self._check_subset_pairs(ts, windows)

    def test_missing_leaves(self):
        """
        test case where 1/2 of samples are missing
        """
        t = self.example_ts().dump_tables()
        ss0 = np.flatnonzero(t.nodes.population == 0)
        remove = np.isin(t.edges.child, ss0)
        assert np.any(remove)
        t.edges.set_columns(
            left=t.edges.left[~remove],
            right=t.edges.right[~remove],
            parent=t.edges.parent[~remove],
            child=t.edges.child[~remove],
        )
        t.sort()
        ts = t.tree_sequence()
        windows = np.array([0.0, ts.sequence_length])
        self._check_total_pairs(ts, windows)
        self._check_subset_pairs(ts, windows)

    def test_multiple_roots(self):
        """
        test case where all trees have multiple roots
        """
        ts = self.example_ts()
        ts = ts.decapitate(np.quantile(ts.nodes_time, 0.75))
        windows = np.array([0.0, ts.sequence_length])
        self._check_total_pairs(ts, windows)
        self._check_subset_pairs(ts, windows)

    def test_windows(self):
        ts = self.example_ts()
        windows = np.linspace(0.0, ts.sequence_length, 9)
        self._check_total_pairs(ts, windows)
        self._check_subset_pairs(ts, windows)

    def test_windows_are_trees(self):
        """
        test case where window breakpoints coincide with tree breakpoints
        """
        ts = self.example_ts()
        windows = np.array(list(ts.breakpoints()))
        self._check_total_pairs(ts, windows)
        self._check_subset_pairs(ts, windows)

    def test_windows_inside_trees(self):
        """
        test case where windows are nested within trees
        """
        ts = self.example_ts()
        windows = np.array(list(ts.breakpoints()))
        windows = np.sort(np.append(windows[:-1] / 2 + windows[1:] / 2, windows))
        self._check_total_pairs(ts, windows)
        self._check_subset_pairs(ts, windows)

    def test_nonsuccinct_sequence(self):
        """
        test case where each tree has distinct nodes
        """
        ts = convert_to_nonsuccinct(self.example_ts())
        windows = np.linspace(0, ts.sequence_length, 9)
        self._check_total_pairs(ts, windows)
        self._check_subset_pairs(ts, windows)

    def test_span_normalise(self):
        """
        test case where span is normalised
        """
        ts = self.example_ts()
        windows = np.array([0.0, 0.33, 1.0]) * ts.sequence_length
        window_size = np.diff(windows)
        implm = ts.pair_coalescence_counts(windows=windows, span_normalise=False)
        check = ts.pair_coalescence_counts(windows=windows) * window_size[:, np.newaxis]
        np.testing.assert_allclose(implm, check)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(ts, windows=windows, span_normalise=False)
        np.testing.assert_allclose(proto, check)

    def test_pair_normalise(self):
        ts = self.example_ts()
        windows = np.array([0.0, 0.33, 1.0]) * ts.sequence_length
        window_size = np.diff(windows)
        total_pairs = ts.num_samples * (ts.num_samples - 1) / 2
        implm = ts.pair_coalescence_counts(
            windows=windows, span_normalise=False, pair_normalise=True
        )
        check = ts.pair_coalescence_counts(windows=windows) * window_size[:, np.newaxis]
        check /= total_pairs
        np.testing.assert_allclose(implm, check)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(
            ts, windows=windows, span_normalise=False, pair_normalise=True
        )
        np.testing.assert_allclose(proto, check)

    def test_internal_nodes_are_samples(self):
        """
        test case where some samples are descendants of other samples
        """
        ts = self.example_ts()
        tables = ts.dump_tables()
        nodes_flags = tables.nodes.flags.copy()
        nodes_sample = np.arange(ts.num_samples, ts.num_nodes, 10)
        nodes_flags[nodes_sample] = tskit.NODE_IS_SAMPLE
        tables.nodes.flags = nodes_flags
        ts_modified = tables.tree_sequence()
        assert ts_modified.num_samples > ts.num_samples
        windows = np.linspace(0.0, 1.0, 9) * ts_modified.sequence_length
        self._check_total_pairs(ts_modified, windows)
        self._check_subset_pairs(ts_modified, windows)

    def test_time_windows(self):
        ts = self.example_ts()
        total_pair_count = ts.pair_coalescence_counts(
            time_windows=np.array([0.0, np.inf]),
            span_normalise=False,
        )[0]
        samples = list(ts.samples())
        time_windows = np.quantile(ts.nodes_time, [0.0, 0.25, 0.5, 0.75])
        time_windows = np.append(time_windows, np.inf)
        implm = ts.pair_coalescence_counts(
            span_normalise=False, time_windows=time_windows
        )
        assert np.isclose(np.sum(implm), total_pair_count)
        check = naive_pair_coalescence_counts(ts, samples, samples).squeeze() / 2
        nodes_map = np.searchsorted(time_windows, ts.nodes_time, side="right") - 1
        check = np.bincount(nodes_map, weights=check)
        np.testing.assert_allclose(implm, check)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(
            ts, span_normalise=False, time_windows=time_windows
        )
        assert np.isclose(np.sum(proto), total_pair_count)
        np.testing.assert_allclose(proto, check)

    def test_time_windows_truncated(self):
        """
        test case where some nodes fall outside of time bins
        """
        ts = self.example_ts()
        total_pair_count = ts.pair_coalescence_counts(
            time_windows=np.array([0.0, np.inf]),
            span_normalise=False,
        )[0]
        samples = list(ts.samples())
        time_windows = np.quantile(ts.nodes_time, [0.5, 0.75])
        assert time_windows[0] > 0.0
        time_windows = np.append(time_windows, np.inf)
        implm = ts.pair_coalescence_counts(
            span_normalise=False, time_windows=time_windows
        )
        assert np.sum(implm) < total_pair_count
        check = naive_pair_coalescence_counts(ts, samples, samples).squeeze() / 2
        nodes_map = np.searchsorted(time_windows, ts.nodes_time, side="right") - 1
        oob = np.logical_or(nodes_map < 0, nodes_map >= time_windows.size)
        check = np.bincount(nodes_map[~oob], weights=check[~oob])
        np.testing.assert_allclose(implm, check)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(
            ts, span_normalise=False, time_windows=time_windows
        )
        assert np.sum(proto) < total_pair_count
        np.testing.assert_allclose(proto, check)

    def test_time_windows_unique(self):
        ts = self.example_ts()
        total_pair_count = ts.pair_coalescence_counts(
            time_windows=np.array([0.0, np.inf]),
            span_normalise=False,
        )[0]
        samples = list(ts.samples())
        time_windows = np.unique(ts.nodes_time)
        time_windows = np.append(time_windows, np.inf)
        implm = ts.pair_coalescence_counts(
            span_normalise=False, time_windows=time_windows
        )
        assert np.isclose(np.sum(implm), total_pair_count)
        check = naive_pair_coalescence_counts(ts, samples, samples).squeeze() / 2
        nodes_map = np.searchsorted(time_windows, ts.nodes_time, side="right") - 1
        check = np.bincount(nodes_map, weights=check)
        np.testing.assert_allclose(implm, check)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(
            ts, span_normalise=False, time_windows=time_windows
        )
        assert np.isclose(np.sum(proto), total_pair_count)
        np.testing.assert_allclose(proto, check)

    def test_diversity(self):
        """
        test that weighted mean of node times equals branch diversity
        """
        ts = self.example_ts()
        windows = np.linspace(0.0, ts.sequence_length, 9)
        check = ts.diversity(mode="branch", windows=windows)
        implm = ts.pair_coalescence_counts(windows=windows)
        implm = 2 * (implm @ ts.nodes_time) / implm.sum(axis=1)
        np.testing.assert_allclose(implm, check)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(ts, windows=windows)
        proto = 2 * (proto @ ts.nodes_time) / proto.sum(axis=1)
        np.testing.assert_allclose(proto, check)

    def test_divergence(self):
        """
        test that weighted mean of node times equals branch divergence
        """
        ts = self.example_ts()
        ss0 = np.flatnonzero(ts.nodes_population == 0)
        ss1 = np.flatnonzero(ts.nodes_population == 1)
        windows = np.linspace(0.0, ts.sequence_length, 9)
        check = ts.divergence(sample_sets=[ss0, ss1], mode="branch", windows=windows)
        implm = ts.pair_coalescence_counts(sample_sets=[ss0, ss1], windows=windows)
        implm = 2 * (implm @ ts.nodes_time) / implm.sum(axis=1)
        np.testing.assert_allclose(implm, check)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_counts(
            ts, sample_sets=[ss0, ss1], windows=windows
        )
        proto = 2 * (proto @ ts.nodes_time) / proto.sum(axis=1)
        np.testing.assert_allclose(proto, check)


class TestCoalescingPairsUsage:
    """
    Test invalid inputs
    """

    @tests.cached_example
    def example_ts(self):
        return msprime.sim_ancestry(
            samples=10,
            recombination_rate=1e-8,
            sequence_length=1e5,
            population_size=1e4,
            random_seed=1024,
        )

    def test_bad_windows(self):
        ts = self.example_ts()
        with pytest.raises(ValueError, match="too small depth"):
            ts.pair_coalescence_counts(windows="whatever")
        with pytest.raises(ValueError, match="must have at least 2 elements"):
            ts.pair_coalescence_counts(windows=[0.0])
        with pytest.raises(tskit.LibraryError, match="must be increasing list"):
            ts.pair_coalescence_counts(
                windows=np.array([0.0, 0.3, 0.2, 1.0]) * ts.sequence_length
            )
        with pytest.raises(tskit.LibraryError, match="must be increasing list"):
            ts.pair_coalescence_counts(
                windows=np.array([0.0, 2.0]) * ts.sequence_length
            )

    def test_bad_sample_sets(self):
        ts = self.example_ts()
        with pytest.raises(tskit.LibraryError, match="out of bounds"):
            ts.pair_coalescence_counts(sample_sets=[[0, ts.num_nodes]])

    def test_bad_indexes(self):
        ts = self.example_ts()
        with pytest.raises(tskit.LibraryError, match="out of bounds"):
            ts.pair_coalescence_counts(indexes=[(0, 1)])
        with pytest.raises(ValueError, match="must be a k x 2 array"):
            ts.pair_coalescence_counts(indexes=[(0, 0, 0)])

    def test_no_indexes(self):
        ts = self.example_ts()
        ss = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
        with pytest.raises(ValueError, match="more than two sample sets"):
            ts.pair_coalescence_counts(sample_sets=ss)

    def test_oob_samples(self):
        ts = self.example_ts()
        sample_sets = [np.arange(ts.num_samples + 1)]
        with pytest.raises(tskit.LibraryError, match="are not samples"):
            ts.pair_coalescence_counts(sample_sets=sample_sets)

    def test_uncalibrated_time(self):
        tables = self.example_ts().dump_tables()
        tables.time_units = tskit.TIME_UNITS_UNCALIBRATED
        ts = tables.tree_sequence()
        with pytest.raises(ValueError, match="require calibrated node times"):
            ts.pair_coalescence_counts(time_windows=np.array([0.0, np.inf]))

    @pytest.mark.parametrize("time_windows", [[], [0.0], [[0.0, 1.0]], "whatever"])
    def test_bad_time_windows(self, time_windows):
        ts = self.example_ts()
        with pytest.raises(ValueError, match="too small depth"):
            ts.pair_coalescence_counts(time_windows="time_windows")

    def test_unsorted_time_windows(self):
        ts = self.example_ts()
        time_windows = np.array([0.0, 12.0, 6.0, np.inf])
        with pytest.raises(ValueError, match="monotonically increasing or decreasing"):
            ts.pair_coalescence_counts(time_windows=time_windows)

    def test_empty_time_windows(self):
        ts = self.example_ts()
        time_windows = [np.max(ts.nodes_time) + 1, np.max(ts.nodes_time) + 2]
        time_windows = np.append(time_windows, np.inf)
        with pytest.raises(ValueError, match="has null values for all nodes"):
            ts.pair_coalescence_counts(time_windows=time_windows)

    def test_output_dim(self):
        """
        test that output dimensions corresponding to None arguments are dropped
        """
        ts = self.example_ts()
        ss = [[0, 1, 2], [3, 4, 5]]
        implm = ts.pair_coalescence_counts(sample_sets=ss, windows=None, indexes=None)
        assert implm.shape == (ts.num_nodes,)
        windows = np.linspace(0.0, ts.sequence_length, 2)
        implm = ts.pair_coalescence_counts(
            sample_sets=ss, windows=windows, indexes=None
        )
        assert implm.shape == (1, ts.num_nodes)
        indexes = [(0, 1), (1, 1)]
        implm = ts.pair_coalescence_counts(
            sample_sets=ss, windows=windows, indexes=indexes
        )
        assert implm.shape == (1, 2, ts.num_nodes)
        implm = ts.pair_coalescence_counts(
            sample_sets=ss, windows=None, indexes=indexes
        )
        assert implm.shape == (2, ts.num_nodes)


class TestPairCoalescenceQuantiles:
    """
    Test quantile reduction
    """

    @tests.cached_example
    def example_ts(self):
        n = 10
        model = msprime.BetaCoalescent(alpha=1.5)  # polytomies
        tables = msprime.sim_ancestry(
            samples=n,
            recombination_rate=1e-8,
            sequence_length=1e6,
            population_size=1e4,
            random_seed=1024,
            model=model,
        ).dump_tables()
        tables.populations.add_row(metadata={"name": "foo", "description": "bar"})
        tables.nodes.population = np.repeat(
            [0, 1, tskit.NULL], [n, n, tables.nodes.num_rows - 2 * n]
        ).astype("int32")
        ts = tables.tree_sequence()
        assert ts.num_trees > 1
        return ts

    def test_quantiles(self):
        ts = self.example_ts()
        quantiles = np.linspace(0, 1, 10)
        weights = ts.pair_coalescence_counts()
        check = _numpy_weighted_quantile(ts.nodes_time, weights, quantiles)
        implm = ts.pair_coalescence_quantiles(quantiles)
        np.testing.assert_allclose(implm, check)
        # TODO: remove with prototype
        proto = proto_pair_coalescence_quantiles(ts, quantiles=quantiles)
        np.testing.assert_allclose(proto, check)

    def test_windows(self):
        ts = self.example_ts()
        quantiles = np.linspace(0, 1, 10)
        windows = np.array([0, 0.5, 1.0]) * ts.sequence_length
        implm = ts.pair_coalescence_quantiles(quantiles, windows=windows)
        weights = ts.pair_coalescence_counts(windows=windows)
        check = np.empty_like(implm)
        for i, w in enumerate(weights):
            check[i] = _numpy_weighted_quantile(ts.nodes_time, w, quantiles)
        np.testing.assert_allclose(implm, check)

    def test_sample_sets(self):
        ts = self.example_ts()
        sample_sets = [
            np.flatnonzero(ts.nodes_population[: ts.num_samples] == i) for i in range(2)
        ]
        quantiles = np.linspace(0, 1, 10)
        indexes = [(0, 1)]
        implm = ts.pair_coalescence_quantiles(
            quantiles, sample_sets=sample_sets, indexes=indexes
        )
        weights = ts.pair_coalescence_counts(sample_sets=sample_sets, indexes=indexes)
        check = _numpy_weighted_quantile(ts.nodes_time, weights.flatten(), quantiles)
        np.testing.assert_allclose(implm.flatten(), check)
        # check default
        implm = ts.pair_coalescence_quantiles(quantiles, sample_sets=sample_sets)
        np.testing.assert_allclose(implm, check)

    def test_observations_are_quantiles(self):
        """
        case where quantiles fall on observations
        """
        ts = self.example_ts()
        weights = ts.pair_coalescence_counts()
        quantiles = np.unique(weights / np.sum(weights))
        check = _numpy_weighted_quantile(ts.nodes_time, weights, quantiles)
        implm = ts.pair_coalescence_quantiles(quantiles)
        np.testing.assert_allclose(implm, check)

    def test_errors(self):
        ts = self.example_ts()
        sample_sets = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
        quantiles = np.linspace(0, 1, 10)
        with pytest.raises(ValueError, match="more than two sample sets"):
            ts.pair_coalescence_quantiles(quantiles, sample_sets=sample_sets)
        tables = ts.dump_tables()
        tables.time_units = tskit.TIME_UNITS_UNCALIBRATED
        with pytest.raises(ValueError, match="require calibrated node times"):
            tables.tree_sequence().pair_coalescence_quantiles(quantiles=np.array([0.5]))

    def test_long_sequence(self):
        ts = _single_tree_example(L=1e8, T=10)
        windows = np.linspace(0, ts.sequence_length, 100)
        time_windows = np.array([0, np.inf])
        # check that there is roundoff error present
        weights = ts.pair_coalescence_counts(
            windows=windows,
            time_windows=time_windows,
            pair_normalise=True,
            span_normalise=True,
        )
        assert np.all(np.isclose(weights, 1.0))
        assert not np.all(weights == 1.0)
        # check that we don't error out
        quantiles = np.linspace(0, 1, 10)
        quants = ts.pair_coalescence_quantiles(windows=windows, quantiles=quantiles)
        ck_quants = _numpy_weighted_quantile(
            ts.nodes_time,
            ts.pair_coalescence_counts(pair_normalise=True),
            quantiles,
        )
        np.testing.assert_allclose(quants, np.tile(ck_quants, (windows.size - 1, 1)))


class TestPairCoalescenceRates:
    """
    Test coalescence rate reduction
    """

    @tests.cached_example
    def example_ts(self):
        n = 10
        tables = msprime.sim_ancestry(
            samples=n,
            recombination_rate=1e-8,
            sequence_length=1e6,
            population_size=1e4,
            random_seed=1025,
        ).dump_tables()
        tables.populations.add_row(metadata={"name": "foo", "description": "bar"})
        tables.nodes.population = np.repeat(
            [0, 1, tskit.NULL], [n, n, tables.nodes.num_rows - 2 * n]
        ).astype("int32")
        ts = tables.tree_sequence()
        assert ts.num_trees > 1
        return ts

    def test_simulated(self):
        ts = self.example_ts()
        quantiles = np.linspace(0, 1, 5)
        weights = ts.pair_coalescence_counts(pair_normalise=True)
        breaks = _numpy_weighted_quantile(ts.nodes_time, weights, quantiles)
        breaks[0], breaks[-1] = 0.0, np.inf
        check = _numpy_hazard_rate(ts.nodes_time, weights, breaks)
        implm = ts.pair_coalescence_rates(breaks)
        np.testing.assert_allclose(implm, check)

    def test_windowed(self):
        ts = self.example_ts()
        quantiles = np.linspace(0, 1, 5)
        weights = ts.pair_coalescence_counts(pair_normalise=True)
        breaks = _numpy_weighted_quantile(ts.nodes_time, weights, quantiles)
        breaks[0], breaks[-1] = 0.0, np.inf
        windows = np.linspace(0, ts.sequence_length, 4)
        implm = ts.pair_coalescence_rates(breaks, windows=windows)
        check = np.empty_like(implm)
        weights = ts.pair_coalescence_counts(pair_normalise=True, windows=windows)
        for i, w in enumerate(weights):
            check[i] = _numpy_hazard_rate(ts.nodes_time, w, breaks)
        np.testing.assert_allclose(implm, check)

    def test_truncated(self):
        ts = self.example_ts()
        max_time = np.max(ts.nodes_time)
        breaks = np.array([0.0, 0.5, 1.0, 2, np.inf]) * np.ceil(max_time)
        weights = ts.pair_coalescence_counts(pair_normalise=True)
        check = _numpy_hazard_rate(ts.nodes_time, weights, breaks)
        implm = ts.pair_coalescence_rates(breaks)
        np.testing.assert_allclose(implm, check)

    def test_empty(self):
        ts = self.example_ts()
        i = ts.num_nodes // 2
        assert ts.nodes_time[i] < ts.nodes_time[i + 1]
        empty_time_window = [
            ts.nodes_time[i] * 0.75 + ts.nodes_time[i + 1] * 0.25,
            ts.nodes_time[i] * 0.25 + ts.nodes_time[i + 1] * 0.75,
        ]
        max_time = np.max(ts.nodes_time)
        breaks = np.array([0.0, *empty_time_window, max_time + 1, max_time + 2, np.inf])
        weights = ts.pair_coalescence_counts(pair_normalise=True)
        check = _numpy_hazard_rate(ts.nodes_time, weights, breaks)
        implm = ts.pair_coalescence_rates(breaks)
        np.testing.assert_allclose(implm, check)

    def test_single(self):
        ts = self.example_ts()
        breaks = np.array([0.0, np.inf])
        indexes = [(0, 0)]
        weights = ts.pair_coalescence_counts(pair_normalise=True)
        check = _numpy_hazard_rate(ts.nodes_time, weights, breaks).reshape(-1, 1)
        implm = ts.pair_coalescence_rates(breaks, indexes=indexes)
        np.testing.assert_allclose(implm, check)

    def test_indexes(self):
        ts = self.example_ts()
        breaks = np.array([0.0, np.inf])
        sample_sets = [[0, 1, 2], [3, 4, 5]]
        weights = ts.pair_coalescence_counts(
            sample_sets=sample_sets, pair_normalise=True
        )
        check = _numpy_hazard_rate(ts.nodes_time, weights, breaks)
        implm = ts.pair_coalescence_rates(breaks, sample_sets=sample_sets)
        np.testing.assert_allclose(implm, check)

    def test_errors(self):
        ts = self.example_ts()
        sample_sets = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
        time_windows = np.array([0, np.inf])
        with pytest.raises(ValueError, match="more than two sample sets"):
            ts.pair_coalescence_rates(time_windows, sample_sets=sample_sets)
        tables = ts.dump_tables()
        tables.time_units = tskit.TIME_UNITS_UNCALIBRATED
        with pytest.raises(ValueError, match="require calibrated node times"):
            tables.tree_sequence().pair_coalescence_rates(
                time_windows=np.array([0.0, np.inf])
            )

    def test_long_sequence(self):
        ts = _single_tree_example(L=1e8, T=10)
        windows = np.linspace(0, ts.sequence_length, 100)
        time_windows = np.array([0, np.inf])
        # check that there is roundoff error present
        weights = ts.pair_coalescence_counts(
            windows=windows,
            time_windows=time_windows,
            pair_normalise=True,
            span_normalise=True,
        )
        assert np.all(np.isclose(weights, 1.0))
        assert not np.all(weights == 1.0)
        # check that we don't error out
        rates = ts.pair_coalescence_rates(windows=windows, time_windows=time_windows)
        ck_rates = _numpy_hazard_rate(
            ts.nodes_time,
            ts.pair_coalescence_counts(pair_normalise=True),
            time_windows,
        )
        np.testing.assert_allclose(
            rates.flatten(), np.repeat(ck_rates, windows.size - 1)
        )


--- ../../tskit/python/tests/tsutil.py ---


"""
A collection of utilities to edit and construct tree sequences.
"""
import collections
import dataclasses
import functools
import json
import random
import string
import struct
import typing

import msprime
import numpy as np

import tskit
import tskit.provenance as provenance


def random_bytes(max_length):
    """
    Returns a random bytearray of the specified maximum length.
    """
    length = random.randint(0, max_length)
    return bytearray(random.randint(0, 255) for _ in range(length))


def random_strings(max_length):
    """
    Returns a random bytearray of the specified maximum length.
    """
    length = random.randint(0, max_length)
    return "".join(random.choice(string.printable) for _ in range(length))


def add_provenance(provenance_table, method_name):
    d = provenance.get_provenance_dict({"command": f"tsutil.{method_name}"})
    provenance_table.add_row(json.dumps(d))


def subsample_sites(ts, num_sites):
    """
    Returns a copy of the specified tree sequence with a random subsample of the
    specified number of sites.
    """
    t = ts.dump_tables()
    t.sites.reset()
    t.mutations.reset()
    sites_to_keep = set(random.sample(list(range(ts.num_sites)), num_sites))
    for site in ts.sites():
        if site.id in sites_to_keep:
            site_id = t.sites.append(site)
            for mutation in site.mutations:
                t.mutations.append(mutation.replace(site=site_id))
    add_provenance(t.provenances, "subsample_sites")
    return t.tree_sequence()


def insert_branch_mutations(ts, mutations_per_branch=1):
    """
    Returns a copy of the specified tree sequence with a mutation on every branch
    in every tree.
    """
    if mutations_per_branch == 0:
        return ts
    tables = ts.dump_tables()
    tables.sites.clear()
    tables.mutations.clear()
    for tree in ts.trees():
        site = tables.sites.add_row(position=tree.interval.left, ancestral_state="0")
        for root in tree.roots:
            state = {tskit.NULL: 0}
            mutation = {tskit.NULL: -1}
            stack = [root]
            while len(stack) > 0:
                u = stack.pop()
                stack.extend(tree.children(u))
                v = tree.parent(u)
                state[u] = state[v]
                parent = mutation[v]
                for _ in range(mutations_per_branch):
                    state[u] = (state[u] + 1) % 2
                    metadata = f"{len(tables.mutations)}".encode()
                    mutation[u] = tables.mutations.add_row(
                        site=site,
                        node=u,
                        derived_state=str(state[u]),
                        parent=parent,
                        metadata=metadata,
                    )
                    parent = mutation[u]
    add_provenance(tables.provenances, "insert_branch_mutations")
    return tables.tree_sequence()


def remove_mutation_times(ts):
    tables = ts.tables
    tables.mutations.time = np.full_like(tables.mutations.time, tskit.UNKNOWN_TIME)
    return tables.tree_sequence()


def insert_discrete_time_mutations(ts, num_times=4, num_sites=10):
    """
    Inserts mutations in the tree sequence at regularly-spaced num_sites
    positions, at only a discrete set of times (the same for all trees): at
    num_times times evenly spaced between 0 and the maximum time.
    """
    tables = ts.tables
    tables.sites.clear()
    tables.mutations.clear()
    height = max(t.time(t.roots[0]) for t in ts.trees())
    for j, pos in enumerate(np.linspace(0, tables.sequence_length, num_sites + 1)[:-1]):
        anc = "X" * j
        tables.sites.add_row(position=pos, ancestral_state=anc)
        t = ts.at(pos)
        for k, s in enumerate(np.linspace(0, height, num_times)):
            for n in t.nodes():
                if t.time(n) <= s and (
                    (t.parent(n) == tskit.NULL) or (t.time(t.parent(n)) > s)
                ):
                    tables.mutations.add_row(
                        site=j, node=n, derived_state=anc + str(k), time=s
                    )
                    k += 1
    tables.sort()
    tables.build_index()
    tables.compute_mutation_parents()
    return tables.tree_sequence()


def insert_branch_sites(ts, m=1):
    """
    Returns a copy of the specified tree sequence with m sites on every branch
    of every tree.
    """
    if m == 0:
        return ts
    tables = ts.dump_tables()
    tables.sites.clear()
    tables.mutations.clear()
    for tree in ts.trees():
        left, right = tree.interval
        delta = (right - left) / (m * len(list(tree.nodes())))
        x = left
        for u in tree.nodes():
            if tree.parent(u) != tskit.NULL:
                for _ in range(m):
                    site = tables.sites.add_row(position=x, ancestral_state="0")
                    tables.mutations.add_row(site=site, node=u, derived_state="1")
                    x += delta
    add_provenance(tables.provenances, "insert_branch_sites")
    return tables.tree_sequence()


def insert_multichar_mutations(ts, seed=1, max_len=10):
    """
    Returns a copy of the specified tree sequence with multiple chararacter
    mutations on a randomly chosen branch in every tree.
    """
    rng = random.Random(seed)
    letters = ["A", "C", "T", "G"]
    tables = ts.dump_tables()
    tables.sites.clear()
    tables.mutations.clear()
    for tree in ts.trees():
        ancestral_state = rng.choice(letters) * rng.randint(0, max_len)
        site = tables.sites.add_row(
            position=tree.interval.left, ancestral_state=ancestral_state
        )
        nodes = list(tree.nodes())
        nodes.remove(tree.root)
        u = rng.choice(nodes)
        derived_state = ancestral_state
        while ancestral_state == derived_state:
            derived_state = rng.choice(letters) * rng.randint(0, max_len)
        tables.mutations.add_row(site=site, node=u, derived_state=derived_state)
    add_provenance(tables.provenances, "insert_multichar_mutations")
    return tables.tree_sequence()


def insert_random_ploidy_individuals(
    ts, min_ploidy=0, max_ploidy=5, max_dimension=3, samples_only=True, seed=1
):
    """
    Takes random contiguous subsets of the samples and assigns them to individuals.
    Also creates random locations in variable dimensions in the unit interval,
    and assigns random parents (including NULL parents). Note that resulting
    individuals will often have nodes with inconsistent populations and/or time.
    """
    rng = random.Random(seed)
    if samples_only:
        node_ids = np.array(ts.samples(), dtype="int")
    else:
        node_ids = np.arange(ts.num_nodes)
    j = 0
    tables = ts.dump_tables()
    tables.individuals.clear()
    individual = tables.nodes.individual[:]
    individual[:] = tskit.NULL
    ind_id = -1
    while j < len(node_ids):
        ploidy = rng.randint(min_ploidy, max_ploidy)
        nodes = node_ids[j : min(j + ploidy, len(node_ids))]
        dimension = rng.randint(0, max_dimension)
        location = [rng.random() for _ in range(dimension)]
        parents = rng.sample(range(-1, 1 + ind_id), min(1 + ind_id, rng.randint(0, 3)))
        ind_id = tables.individuals.add_row(location=location, parents=parents)
        individual[nodes] = ind_id
        j += ploidy
    tables.nodes.individual = individual
    return tables.tree_sequence()


def insert_random_consistent_individuals(
    ts, min_ploidy=0, max_ploidy=5, min_dimension=0, max_dimension=3, seed=1
):
    """
    Takes random subsets of nodes having the same time and population and
    assigns them to individuals.  Also creates random locations in variable
    dimensions in the unit interval, and assigns random parents (including NULL
    parents).
    """
    rng = random.Random(seed)
    tables = ts.dump_tables()
    tables.individuals.clear()
    individual = tables.nodes.individual[:]
    individual[:] = tskit.NULL
    ind_id = -1
    pops = np.arange(ts.num_populations)
    for pop in pops:
        n = tables.nodes.population == pop
        times = np.unique(tables.nodes.time[n])
        for t in times:
            nn = np.where(np.logical_and(n, tables.nodes.time == t))[0]
            rng.shuffle(nn)
            j = 0
            while j < len(nn):
                ploidy = rng.randint(min_ploidy, max_ploidy)
                nodes = nn[j : min(j + ploidy, len(nn))]
                dimension = rng.randint(min_dimension, max_dimension)
                location = [rng.random() for _ in range(dimension)]
                parents = rng.sample(
                    range(-1, 1 + ind_id), min(1 + ind_id, rng.randint(0, 3))
                )
                ind_id = tables.individuals.add_row(location=location, parents=parents)
                individual[nodes] = ind_id
                j += ploidy
                j += rng.randint(0, 2)  # skip a random number
    tables.nodes.individual = individual
    return tables.tree_sequence()


def insert_individuals(ts, nodes=None, ploidy=1):
    """
    Inserts individuals into the tree sequence using the specified list
    of node (or use all sample nodes if None) with the specified ploidy by combining
    ploidy-sized chunks of the list. Add metadata to the individuals so we can
    track them
    """
    if nodes is None:
        nodes = ts.samples()
    assert len(nodes) % ploidy == 0  # To allow mixed ploidies we could comment this out
    tables = ts.dump_tables()
    tables.individuals.clear()
    individual = tables.nodes.individual[:]
    individual[:] = tskit.NULL
    j = 0
    while j < len(nodes):
        nodes_in_individual = nodes[j : min(len(nodes), j + ploidy)]
        # should we warn here if nodes[j : j + ploidy] are at different times?
        # probably not, as although this is unusual, it is actually allowed
        ind_id = tables.individuals.add_row(
            metadata=f"orig_id {tables.individuals.num_rows}".encode()
        )
        individual[nodes_in_individual] = ind_id
        j += ploidy
    tables.nodes.individual = individual
    return tables.tree_sequence()


def mark_metadata(ts, table_name, prefix="orig_id:"):
    """
    Add metadata to all rows of the form prefix + row_number
    """
    tables = ts.dump_tables()
    table = getattr(tables, table_name)
    table.packset_metadata([(prefix + str(i)).encode() for i in range(table.num_rows)])
    return tables.tree_sequence()


def permute_nodes(ts, node_map):
    """
    Returns a copy of the specified tree sequence such that the nodes are
    permuted according to the specified map.
    """
    tables = ts.dump_tables()
    tables.nodes.clear()
    tables.edges.clear()
    tables.mutations.clear()
    # Mapping from nodes in the new tree sequence back to nodes in the original
    reverse_map = [0 for _ in node_map]
    for j in range(ts.num_nodes):
        reverse_map[node_map[j]] = j
    old_nodes = list(ts.nodes())
    for j in range(ts.num_nodes):
        old_node = old_nodes[reverse_map[j]]
        tables.nodes.append(old_node)
    for edge in ts.edges():
        tables.edges.append(
            edge.replace(parent=node_map[edge.parent], child=node_map[edge.child])
        )
    for site in ts.sites():
        for mutation in site.mutations:
            tables.mutations.append(
                mutation.replace(site=site.id, node=node_map[mutation.node])
            )
    tables.sort()
    add_provenance(tables.provenances, "permute_nodes")
    return tables.tree_sequence()


def insert_redundant_breakpoints(ts):
    """
    Builds a new tree sequence containing redundant breakpoints.
    """
    tables = ts.dump_tables()
    tables.edges.reset()
    for r in ts.edges():
        x = r.left + (r.right - r.left) / 2
        tables.edges.append(r.replace(right=x))
        tables.edges.append(r.replace(left=x))
    add_provenance(tables.provenances, "insert_redundant_breakpoints")
    new_ts = tables.tree_sequence()
    assert new_ts.num_edges == 2 * ts.num_edges
    return new_ts


def single_childify(ts):
    """
    Builds a new equivalent tree sequence which contains an extra node in the
    middle of all existing branches.
    """
    tables = ts.dump_tables()

    mutations_above_node = collections.defaultdict(list)
    for mut in tables.mutations:
        mutations_above_node[mut.node].append(mut)

    mutations_on_edge = collections.defaultdict(list)
    for edge_idx, edge in enumerate(tables.edges):
        for mut in mutations_above_node[edge.child]:
            if edge.left <= tables.sites[mut.site].position < edge.right:
                mutations_on_edge[edge_idx].append(mut)

    time = tables.nodes.time[:]
    tables.edges.reset()
    tables.mutations.reset()
    for edge in ts.edges():
        # Insert a new node in between the parent and child.
        t = time[edge.child] + (time[edge.parent] - time[edge.child]) / 2
        u = tables.nodes.add_row(time=t)
        tables.edges.append(edge.replace(parent=u))
        tables.edges.append(edge.replace(child=u))
        for mut in mutations_on_edge[edge.id]:
            if mut.time < t:
                tables.mutations.append(mut)
            else:
                tables.mutations.append(mut.replace(node=u))
    tables.sort()
    add_provenance(tables.provenances, "insert_redundant_breakpoints")
    return tables.tree_sequence()


def add_random_metadata(ts, seed=1, max_length=10):
    """
    Returns a copy of the specified tree sequence with random metadata assigned
    to the nodes, sites and mutations.
    """
    tables = ts.dump_tables()
    np.random.seed(seed)

    length = np.random.randint(0, max_length, ts.num_nodes)
    offset = np.cumsum(np.hstack(([0], length)), dtype=np.uint32)
    # Older versions of numpy didn't have a dtype argument for randint, so
    # must use astype instead.
    metadata = np.random.randint(-127, 127, offset[-1]).astype(np.int8)
    nodes = tables.nodes
    nodes.set_columns(
        flags=nodes.flags,
        population=nodes.population,
        time=nodes.time,
        metadata_offset=offset,
        metadata=metadata,
        individual=nodes.individual,
    )

    length = np.random.randint(0, max_length, ts.num_sites)
    offset = np.cumsum(np.hstack(([0], length)), dtype=np.uint32)
    metadata = np.random.randint(-127, 127, offset[-1]).astype(np.int8)
    sites = tables.sites
    sites.set_columns(
        position=sites.position,
        ancestral_state=sites.ancestral_state,
        ancestral_state_offset=sites.ancestral_state_offset,
        metadata_offset=offset,
        metadata=metadata,
    )

    length = np.random.randint(0, max_length, ts.num_mutations)
    offset = np.cumsum(np.hstack(([0], length)), dtype=np.uint32)
    metadata = np.random.randint(-127, 127, offset[-1]).astype(np.int8)
    mutations = tables.mutations
    mutations.set_columns(
        site=mutations.site,
        node=mutations.node,
        time=mutations.time,
        parent=mutations.parent,
        derived_state=mutations.derived_state,
        derived_state_offset=mutations.derived_state_offset,
        metadata_offset=offset,
        metadata=metadata,
    )

    length = np.random.randint(0, max_length, ts.num_individuals)
    offset = np.cumsum(np.hstack(([0], length)), dtype=np.uint32)
    metadata = np.random.randint(-127, 127, offset[-1]).astype(np.int8)
    individuals = tables.individuals
    individuals.set_columns(
        flags=individuals.flags,
        location=individuals.location,
        location_offset=individuals.location_offset,
        parents=individuals.parents,
        parents_offset=individuals.parents_offset,
        metadata_offset=offset,
        metadata=metadata,
    )

    length = np.random.randint(0, max_length, ts.num_populations)
    offset = np.cumsum(np.hstack(([0], length)), dtype=np.uint32)
    metadata = np.random.randint(-127, 127, offset[-1]).astype(np.int8)
    populations = tables.populations
    populations.set_columns(metadata_offset=offset, metadata=metadata)

    add_provenance(tables.provenances, "add_random_metadata")
    ts = tables.tree_sequence()
    return ts


def jiggle_samples(ts):
    """
    Returns a copy of the specified tree sequence with the sample nodes switched
    around. The first n / 2 existing samples become non samples, and the last
    n / 2 node become samples.
    """
    tables = ts.dump_tables()
    nodes = tables.nodes
    flags = nodes.flags
    oldest_parent = tables.edges.parent[-1]
    n = ts.sample_size
    flags[: n // 2] = 0
    flags[oldest_parent - n // 2 : oldest_parent] = 1
    nodes.set_columns(flags, nodes.time)
    add_provenance(tables.provenances, "jiggle_samples")
    return tables.tree_sequence()


def generate_site_mutations(
    tree, position, mu, site_table, mutation_table, multiple_per_node=True
):
    """
    Generates mutations for the site at the specified position on the specified
    tree. Mutations happen at rate mu along each branch. The site and mutation
    information are recorded in the specified tables.  Note that this records
    more than one mutation per edge.
    """
    assert tree.interval.left <= position < tree.interval.right
    states = ["A", "C", "G", "T"]
    ancestral_state = random.choice(states)
    site_table.add_row(position, ancestral_state)
    site = site_table.num_rows - 1
    for root in tree.roots:
        stack = [(root, ancestral_state, tskit.NULL)]
        while len(stack) != 0:
            u, state, parent = stack.pop()
            if u != root:
                branch_length = tree.branch_length(u)
                x = random.expovariate(mu)
                new_state = state
                while x < branch_length:
                    new_state = random.choice(states)
                    if multiple_per_node:
                        mutation_table.add_row(site, u, new_state, parent)
                        parent = mutation_table.num_rows - 1
                        state = new_state
                    x += random.expovariate(mu)
                else:
                    if not multiple_per_node:
                        mutation_table.add_row(site, u, new_state, parent)
                        parent = mutation_table.num_rows - 1
                        state = new_state
            stack.extend(reversed([(v, state, parent) for v in tree.children(u)]))


def jukes_cantor(ts, num_sites, mu, multiple_per_node=True, seed=None):
    """
    Returns a copy of the specified tree sequence with Jukes-Cantor mutations
    applied at the specified rate at the specified number of sites. Site positions
    are chosen uniformly.
    """
    random.seed(seed)
    positions = [ts.sequence_length * random.random() for _ in range(num_sites)]
    positions.sort()
    tables = ts.dump_tables()
    tables.sites.clear()
    tables.mutations.clear()
    trees = ts.trees()
    t = next(trees)
    for position in positions:
        while position >= t.interval.right:
            t = next(trees)
        generate_site_mutations(
            t,
            position,
            mu,
            tables.sites,
            tables.mutations,
            multiple_per_node=multiple_per_node,
        )
    add_provenance(tables.provenances, "jukes_cantor")
    new_ts = tables.tree_sequence()
    return new_ts


def caterpillar_tree(n, num_sites=0, num_mutations=1):
    """
    Returns caterpillar tree with n samples. For each of the sites and
    path of at most n - 2 mutations are put down along the internal
    nodes. Each site gets exactly the same set of mutations.
    """
    if num_sites > 0 and num_mutations > n - 2:
        raise ValueError("At most n - 2 mutations allowed")
    tables = tskit.TableCollection(1)
    for _ in range(n):
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
    last_node = 0
    # Add the internal nodes
    for j in range(n - 1):
        u = tables.nodes.add_row(time=j + 1)
        tables.edges.add_row(0, tables.sequence_length, u, last_node)
        tables.edges.add_row(0, tables.sequence_length, u, j + 1)
        last_node = u
    for j in range(num_sites):
        tables.sites.add_row(position=(j + 1) / n, ancestral_state="0")
        node = 2 * n - 3
        state = 0
        for _ in range(num_mutations):
            state = (state + 1) % 2
            tables.mutations.add_row(site=j, derived_state=str(state), node=node)
            node -= 1

    tables.sort()
    tables.build_index()
    tables.compute_mutation_parents()
    return tables.tree_sequence()


def compute_mutation_parent(ts):
    """
    Compute the `parent` column of a MutationTable. Correct computation uses
    topological information in the nodes and edges, as well as the fact that
    each mutation must be listed after the mutation on whose background it
    occurred (i.e., its parent).

    :param TreeSequence ts: The tree sequence to compute for.  Need not
        have a valid mutation parent column.
    """
    mutation_parent = np.zeros(ts.num_mutations, dtype=np.int32) - 1
    # Maps nodes to the bottom mutation on each branch
    bottom_mutation = np.zeros(ts.num_nodes, dtype=np.int32) - 1
    for tree in ts.trees():
        for site in tree.sites():
            # Go forward through the mutations creating a mapping from the
            # mutations to the nodes. If we see more than one mutation
            # at a node, then these must be parents since we're assuming
            # they are in order.
            for mutation in site.mutations:
                if bottom_mutation[mutation.node] != tskit.NULL:
                    mutation_parent[mutation.id] = bottom_mutation[mutation.node]
                bottom_mutation[mutation.node] = mutation.id
            # There's no point in checking the first mutation, since this cannot
            # have a parent.
            for mutation in site.mutations[1:]:
                if mutation_parent[mutation.id] == tskit.NULL:
                    v = tree.parent(mutation.node)
                    # Traverse upwards until we find a another mutation or root.
                    while v != tskit.NULL and bottom_mutation[v] == tskit.NULL:
                        v = tree.parent(v)
                    if v != tskit.NULL:
                        mutation_parent[mutation.id] = bottom_mutation[v]
            # Reset the maps for the next site.
            for mutation in site.mutations:
                bottom_mutation[mutation.node] = tskit.NULL
            assert np.all(bottom_mutation == -1)
    return mutation_parent


def py_subset(
    tables,
    nodes,
    record_provenance=True,
    reorder_populations=True,
    remove_unreferenced=True,
):
    """
    Naive implementation of the TableCollection.subset method using the Python API.
    """
    if np.any(nodes > tables.nodes.num_rows) or np.any(nodes < 0):
        raise ValueError("Nodes out of bounds.")
    full = tables.copy()
    tables.clear()
    # mapping from old to new ids
    node_map = {}
    ind_map = {tskit.NULL: tskit.NULL}
    pop_map = {tskit.NULL: tskit.NULL}
    if not reorder_populations:
        for j, pop in enumerate(full.populations):
            pop_map[j] = j
            tables.populations.append(pop)
    # first build individual map
    if not remove_unreferenced:
        keep_ind = [True for _ in full.individuals]
    else:
        keep_ind = [False for _ in full.individuals]
        for old_id in nodes:
            i = full.nodes[old_id].individual
            if i != tskit.NULL:
                keep_ind[i] = True
    new_ind_id = 0
    for j, k in enumerate(keep_ind):
        if k:
            ind_map[j] = new_ind_id
            new_ind_id += 1
    # now the individual table
    for j, k in enumerate(keep_ind):
        if k:
            ind = full.individuals[j]
            new_ind_id = tables.individuals.append(
                ind.replace(parents=[ind_map[i] for i in ind.parents if i in ind_map])
            )

            assert new_ind_id == ind_map[j]

    for old_id in nodes:
        node = full.nodes[old_id]
        if node.population not in pop_map and node.population != tskit.NULL:
            pop = full.populations[node.population]
            new_pop_id = tables.populations.append(pop)
            pop_map[node.population] = new_pop_id
        new_id = tables.nodes.append(
            node.replace(
                population=pop_map[node.population],
                individual=ind_map[node.individual],
            )
        )
        node_map[old_id] = new_id
    if not remove_unreferenced:
        for j, ind in enumerate(full.populations):
            if j not in pop_map:
                pop_map[j] = tables.populations.append(ind)
    for edge in full.edges:
        if edge.child in nodes and edge.parent in nodes:
            tables.edges.append(
                edge.replace(parent=node_map[edge.parent], child=node_map[edge.child])
            )
    if full.migrations.num_rows > 0:
        raise ValueError("Migrations are currently not supported in this operation.")
    site_map = {}
    if not remove_unreferenced:
        for j, site in enumerate(full.sites):
            site_map[j] = tables.sites.append(site)
    mutation_map = {tskit.NULL: tskit.NULL}
    for i, mut in enumerate(full.mutations):
        if mut.node in nodes:
            if mut.site not in site_map:
                site = full.sites[mut.site]
                new_site = tables.sites.append(site)
                site_map[mut.site] = new_site
            new_mut = tables.mutations.append(
                mut.replace(
                    site=site_map[mut.site],
                    node=node_map[mut.node],
                    parent=mutation_map.get(mut.parent, tskit.NULL),
                )
            )
            mutation_map[i] = new_mut

    tables.sort()


def py_union(tables, other, nodes, record_provenance=True, add_populations=True):
    """
    Python implementation of TableCollection.union().
    """
    # mappings of id in other to new id in tables
    # the +1 is to take care of mapping tskit.NULL(-1) to tskit.NULL
    pop_map = [tskit.NULL for _ in range(other.populations.num_rows + 1)]
    ind_map = [tskit.NULL for _ in range(other.individuals.num_rows + 1)]
    node_map = [tskit.NULL for _ in range(other.nodes.num_rows + 1)]
    site_map = [tskit.NULL for _ in range(other.sites.num_rows + 1)]
    mut_map = [tskit.NULL for _ in range(other.mutations.num_rows + 1)]
    original_num_individuals = tables.individuals.num_rows

    for other_id, node in enumerate(other.nodes):
        if nodes[other_id] != tskit.NULL and node.individual != tskit.NULL:
            ind_map[node.individual] = tables.nodes[nodes[other_id]].individual

    for other_id, node in enumerate(other.nodes):
        if nodes[other_id] != tskit.NULL:
            node_map[other_id] = nodes[other_id]
        else:
            if ind_map[node.individual] == tskit.NULL and node.individual != tskit.NULL:
                ind = other.individuals[node.individual]
                ind_id = tables.individuals.append(ind)
                ind_map[node.individual] = ind_id
            if pop_map[node.population] == tskit.NULL and node.population != tskit.NULL:
                if not add_populations:
                    pop_map[node.population] = node.population
                else:
                    pop = other.populations[node.population]
                    pop_id = tables.populations.append(pop)
                    pop_map[node.population] = pop_id
            node_id = tables.nodes.append(
                node.replace(
                    population=pop_map[node.population],
                    individual=ind_map[node.individual],
                )
            )
            node_map[other_id] = node_id
    individuals = tables.individuals
    new_parents = individuals.parents
    for i in range(
        individuals.parents_offset[original_num_individuals], len(individuals.parents)
    ):
        new_parents[i] = ind_map[individuals.parents[i]]
    individuals.parents = new_parents
    for edge in other.edges:
        if (nodes[edge.parent] == tskit.NULL) or (nodes[edge.child] == tskit.NULL):
            tables.edges.append(
                edge.replace(parent=node_map[edge.parent], child=node_map[edge.child])
            )
    for other_id, mut in enumerate(other.mutations):
        if nodes[mut.node] == tskit.NULL:
            # add site: may already be in tables, but we deduplicate
            if site_map[mut.site] == tskit.NULL:
                site = other.sites[mut.site]
                site_id = tables.sites.append(site)
                site_map[mut.site] = site_id
            mut_id = tables.mutations.append(
                mut.replace(
                    site=site_map[mut.site],
                    node=node_map[mut.node],
                    parent=tskit.NULL,
                )
            )
            mut_map[other_id] = mut_id
    # migration table
    # grafting provenance table
    if record_provenance:
        pass
    # sorting, deduplicating sites, and re-computing mutation parents
    tables.sort()
    tables.deduplicate_sites()
    # need to sort again since after deduplicating sites, mutations may not be
    # sorted by time within sites
    tables.sort()
    tables.build_index()
    tables.compute_mutation_parents()


def compute_mutation_times(ts):
    """
    Compute the `time` column of a MutationTable in a TableCollection.
    Finds the set of mutations on an edge that share a site and spreads
    the times evenly over the edge.

    :param TreeSequence ts: The tree sequence to compute for.  Need not
        have a valid mutation time column.
    """
    tables = ts.dump_tables()
    mutations = tables.mutations

    mutations_above_node = collections.defaultdict(list)
    for mut_idx, mut in enumerate(mutations):
        mutations_above_node[mut.node].append((mut_idx, mut))

    mutations_at_site_on_edge = collections.defaultdict(list)
    for edge_idx, edge in enumerate(tables.edges):
        for mut_idx, mut in mutations_above_node[edge.child]:
            if edge.left <= tables.sites[mut.site].position < edge.right:
                mutations_at_site_on_edge[(mut.site, edge_idx)].append(mut_idx)

    edges = tables.edges
    nodes = tables.nodes
    times = np.full(len(mutations), -1, dtype=np.float64)
    for (_, edge_idx), edge_mutations in mutations_at_site_on_edge.items():
        start_time = nodes[edges[edge_idx].child].time
        end_time = nodes[edges[edge_idx].parent].time
        duration = end_time - start_time
        for i, mut_idx in enumerate(edge_mutations):
            times[mut_idx] = end_time - (
                duration * ((i + 1) / (len(edge_mutations) + 1))
            )

    # Mutations not on a edge (i.e. above a root) get given their node's time
    for i in range(len(mutations)):
        if times[i] == -1:
            times[i] = nodes[mutations[i].node].time
    tables.mutations.time = times
    tables.sort()
    return tables.mutations.time


def shuffle_tables(
    tables,
    seed,
    shuffle_edges=True,
    shuffle_populations=True,
    shuffle_individuals=True,
    shuffle_sites=True,
    shuffle_mutations=True,
    shuffle_migrations=True,
    keep_mutation_parent_order=False,
):
    """
    Randomizes the order of rows in (possibly) all except the Node table.  Note
    that if mutations are completely shuffled, then TableCollection.sort() will
    not necessarily produce valid tables (unless all mutation times are present
    and distinct), since currently only canonicalise puts parent mutations
    before children.  However, setting keep_mutation_parent_order to True will
    maintain the order of mutations within each site.

    :param TableCollection tables: The table collection that is shuffled (in place).
    """
    rng = random.Random(seed)
    orig = tables.copy()
    tables.nodes.clear()
    tables.individuals.clear()
    tables.populations.clear()
    tables.edges.clear()
    tables.sites.clear()
    tables.mutations.clear()
    tables.drop_index()
    # populations
    randomised_pops = list(enumerate(orig.populations))
    if shuffle_populations:
        rng.shuffle(randomised_pops)
    pop_id_map = {tskit.NULL: tskit.NULL}
    for j, p in randomised_pops:
        pop_id_map[j] = tables.populations.append(p)
    # individuals
    randomised_inds = list(enumerate(orig.individuals))
    if shuffle_individuals:
        rng.shuffle(randomised_inds)
    ind_id_map = {tskit.NULL: tskit.NULL}
    for j, i in randomised_inds:
        ind_id_map[j] = tables.individuals.append(i)
    tables.individuals.parents = [
        tskit.NULL if i == tskit.NULL else ind_id_map[i]
        for i in tables.individuals.parents
    ]
    # nodes (same order, but remapped populations and individuals)
    for n in orig.nodes:
        tables.nodes.append(
            n.replace(
                population=pop_id_map[n.population],
                individual=ind_id_map[n.individual],
            )
        )
    # edges
    randomised_edges = list(orig.edges)
    if shuffle_edges:
        rng.shuffle(randomised_edges)
    for e in randomised_edges:
        tables.edges.append(e)
    # migrations
    randomised_migrations = list(orig.migrations)
    if shuffle_migrations:
        rng.shuffle(randomised_migrations)
    for m in randomised_migrations:
        tables.migrations.append(
            m.replace(source=pop_id_map[m.source], dest=pop_id_map[m.dest])
        )
    # sites
    randomised_sites = list(enumerate(orig.sites))
    if shuffle_sites:
        rng.shuffle(randomised_sites)
    site_id_map = {}
    for j, s in randomised_sites:
        site_id_map[j] = tables.sites.append(s)
    # mutations
    randomised_mutations = list(enumerate(orig.mutations))
    if shuffle_mutations:
        if keep_mutation_parent_order:
            # randomise *except* keeping parent mutations before children
            mut_site_order = [mut.site for mut in orig.mutations]
            rng.shuffle(mut_site_order)
            mut_by_site = {s: [] for s in mut_site_order}
            for j, m in enumerate(orig.mutations):
                mut_by_site[m.site].insert(0, (j, m))
            randomised_mutations = []
            for s in mut_site_order:
                randomised_mutations.append(mut_by_site[s].pop())
        else:
            rng.shuffle(randomised_mutations)
    mut_id_map = {tskit.NULL: tskit.NULL}
    for j, (k, _) in enumerate(randomised_mutations):
        mut_id_map[k] = j
    for _, m in randomised_mutations:
        tables.mutations.append(
            m.replace(site=site_id_map[m.site], parent=mut_id_map[m.parent])
        )
    if keep_mutation_parent_order:
        assert np.all(tables.mutations.parent < np.arange(tables.mutations.num_rows))
    return tables


def cmp_site(i, j, tables):
    ret = tables.sites.position[i] - tables.sites.position[j]
    if ret == 0:
        ret = i - j
    return ret


def cmp_mutation_canonical(i, j, tables, site_order, num_descendants=None):
    site_i = tables.mutations.site[i]
    site_j = tables.mutations.site[j]
    ret = site_order[site_i] - site_order[site_j]
    if (
        ret == 0
        and (not tskit.is_unknown_time(tables.mutations.time[i]))
        and (not tskit.is_unknown_time(tables.mutations.time[j]))
    ):
        ret = tables.mutations.time[j] - tables.mutations.time[i]
    if ret == 0:
        ret = num_descendants[j] - num_descendants[i]
    if ret == 0:
        ret = tables.mutations.node[i] - tables.mutations.node[j]
    if ret == 0:
        ret = i - j
    return ret


def cmp_mutation(i, j, tables, site_order):
    site_i = tables.mutations.site[i]
    site_j = tables.mutations.site[j]
    ret = site_order[site_i] - site_order[site_j]
    if (
        ret == 0
        and (not tskit.is_unknown_time(tables.mutations.time[i]))
        and (not tskit.is_unknown_time(tables.mutations.time[j]))
    ):
        ret = tables.mutations.time[j] - tables.mutations.time[i]
    if ret == 0:
        ret = i - j
    return ret


def cmp_edge(i, j, tables):
    ret = (
        tables.nodes.time[tables.edges.parent[i]]
        - tables.nodes.time[tables.edges.parent[j]]
    )
    if ret == 0:
        ret = tables.edges.parent[i] - tables.edges.parent[j]
    if ret == 0:
        ret = tables.edges.child[i] - tables.edges.child[j]
    if ret == 0:
        ret = tables.edges.left[i] - tables.edges.left[j]
    return ret


def cmp_migration(i, j, tables):
    ret = tables.migrations.time[i] - tables.migrations.time[j]
    if ret == 0:
        ret = tables.migrations.source[i] - tables.migrations.source[j]
    if ret == 0:
        ret = tables.migrations.dest[i] - tables.migrations.dest[j]
    if ret == 0:
        ret = tables.migrations.left[i] - tables.migrations.left[j]
    if ret == 0:
        ret = tables.migrations.node[i] - tables.migrations.node[j]
    return ret


def cmp_individual_canonical(i, j, tables, num_descendants):
    ret = num_descendants[j] - num_descendants[i]
    if ret == 0:
        node_i = node_j = tables.nodes.num_rows
        ni = np.where(tables.nodes.individual == i)[0]
        if len(ni) > 0:
            node_i = np.min(ni)
        nj = np.where(tables.nodes.individual == j)[0]
        if len(nj) > 0:
            node_j = np.min(nj)
        ret = node_i - node_j
    if ret == 0:
        ret = i - j
    return ret


def compute_mutation_num_descendants(tables):
    mutations = tables.mutations
    num_descendants = np.zeros(mutations.num_rows)
    for p in mutations.parent:
        while p != tskit.NULL:
            num_descendants[p] += 1
            p = mutations.parent[p]
    return num_descendants


def compute_individual_num_descendants(tables):
    # adapted from sort_individual_table
    individuals = tables.individuals
    num_individuals = individuals.num_rows
    num_descendants = np.zeros((num_individuals,), np.int64)

    # First find the set of individuals that have no children
    # by creating an array of incoming edge counts
    incoming_edge_count = np.zeros((num_individuals,), np.int64)
    for parent in individuals.parents:
        if parent != tskit.NULL:
            incoming_edge_count[parent] += 1
    todo = np.full((num_individuals + 1,), -1, np.int64)
    current_todo = 0
    todo_insertion_point = 0
    for individual, num_edges in enumerate(incoming_edge_count):
        if num_edges == 0:
            todo[todo_insertion_point] = individual
            todo_insertion_point += 1

    # Now process individuals from the set that have no children, updating their
    # parents' information as we go, and adding their parents to the list if
    # this was their last child
    while todo[current_todo] != -1:
        individual = todo[current_todo]
        current_todo += 1
        for parent in individuals.parents[
            individuals.parents_offset[individual] : individuals.parents_offset[
                individual + 1
            ]
        ]:
            if parent != tskit.NULL:
                incoming_edge_count[parent] -= 1
                num_descendants[parent] += 1 + num_descendants[individual]
                if incoming_edge_count[parent] == 0:
                    todo[todo_insertion_point] = parent
                    todo_insertion_point += 1

    if num_individuals > 0:
        assert np.min(incoming_edge_count) >= 0
        if np.max(incoming_edge_count) > 0:
            raise ValueError("Individual pedigree has cycles")
    return num_descendants


def py_canonicalise(tables, remove_unreferenced=True):
    tables.subset(
        np.arange(tables.nodes.num_rows),
        record_provenance=False,
        remove_unreferenced=remove_unreferenced,
    )
    py_sort(tables, canonical=True)


def py_sort(tables, canonical=False):
    copy = tables.copy()
    tables.edges.clear()
    tables.sites.clear()
    tables.mutations.clear()
    tables.migrations.clear()
    edge_key = functools.cmp_to_key(lambda a, b: cmp_edge(a, b, tables=copy))
    sorted_edges = sorted(range(copy.edges.num_rows), key=edge_key)
    site_key = functools.cmp_to_key(lambda a, b: cmp_site(a, b, tables=copy))
    sorted_sites = sorted(range(copy.sites.num_rows), key=site_key)
    site_id_map = {k: j for j, k in enumerate(sorted_sites)}
    site_order = np.argsort(sorted_sites)
    if canonical:
        mut_num_descendants = compute_mutation_num_descendants(copy)
        mut_key = functools.cmp_to_key(
            lambda a, b: cmp_mutation_canonical(
                a,
                b,
                tables=copy,
                site_order=site_order,
                num_descendants=mut_num_descendants,
            )
        )
    else:
        mut_key = functools.cmp_to_key(
            lambda a, b: cmp_mutation(
                a,
                b,
                tables=copy,
                site_order=site_order,
            )
        )
    sorted_muts = sorted(range(copy.mutations.num_rows), key=mut_key)
    mut_id_map = {k: j for j, k in enumerate(sorted_muts)}
    mut_id_map[tskit.NULL] = tskit.NULL
    mig_key = functools.cmp_to_key(lambda a, b: cmp_migration(a, b, tables=copy))
    sorted_migs = sorted(range(copy.migrations.num_rows), key=mig_key)
    for edge_id in sorted_edges:
        tables.edges.append(copy.edges[edge_id])
    for site_id in sorted_sites:
        tables.sites.append(copy.sites[site_id])
    for mut_id in sorted_muts:
        tables.mutations.append(
            copy.mutations[mut_id].replace(
                site=site_id_map[copy.mutations[mut_id].site],
                parent=mut_id_map[copy.mutations[mut_id].parent],
            )
        )
    for mig_id in sorted_migs:
        tables.migrations.append(copy.migrations[mig_id])

    # individuals
    if canonical:
        tables.individuals.clear()
        ind_num_descendants = compute_individual_num_descendants(copy)
        ind_key = functools.cmp_to_key(
            lambda a, b: cmp_individual_canonical(
                a,
                b,
                tables=copy,
                num_descendants=ind_num_descendants,
            )
        )
        sorted_inds = sorted(range(copy.individuals.num_rows), key=ind_key)
        ind_id_map = {k: j for j, k in enumerate(sorted_inds)}
        ind_id_map[tskit.NULL] = tskit.NULL
        for ind_id in sorted_inds:
            tables.individuals.append(
                copy.individuals[ind_id].replace(
                    parents=[ind_id_map[p] for p in copy.individuals[ind_id].parents],
                )
            )
        tables.nodes.individual = [ind_id_map[i] for i in tables.nodes.individual]


def algorithm_T(ts):
    """
    Simple implementation of algorithm T from the PLOS paper, taking into
    account tree sequences with gaps and other complexities.
    """
    sequence_length = ts.sequence_length
    edges = list(ts.edges())
    M = len(edges)
    time = [ts.node(edge.parent).time for edge in edges]
    in_order = sorted(
        range(M),
        key=lambda j: (edges[j].left, time[j], edges[j].parent, edges[j].child),
    )
    out_order = sorted(
        range(M),
        key=lambda j: (edges[j].right, -time[j], -edges[j].parent, -edges[j].child),
    )
    j = 0
    k = 0
    left = 0
    parent = [-1 for _ in range(ts.num_nodes)]
    while j < M or left < sequence_length:
        while k < M and edges[out_order[k]].right == left:
            edge = edges[out_order[k]]
            parent[edge.child] = -1
            k += 1
        while j < M and edges[in_order[j]].left == left:
            edge = edges[in_order[j]]
            parent[edge.child] = edge.parent
            j += 1
        right = sequence_length
        if j < M:
            right = min(right, edges[in_order[j]].left)
        if k < M:
            right = min(right, edges[out_order[k]].right)
        yield (left, right), parent
        left = right


class QuintuplyLinkedTree:
    def __init__(self, n, root_threshold=1):
        self.root_threshold = root_threshold
        self.parent = np.zeros(n + 1, dtype=np.int32) - 1
        self.left_child = np.zeros(n + 1, dtype=np.int32) - 1
        self.right_child = np.zeros(n + 1, dtype=np.int32) - 1
        self.left_sib = np.zeros(n + 1, dtype=np.int32) - 1
        self.right_sib = np.zeros(n + 1, dtype=np.int32) - 1
        self.num_samples = np.zeros(n + 1, dtype=np.int32)
        self.num_edges = 0
        self.num_children = np.zeros(n + 1, dtype=np.int32)
        self.edge = np.zeros(n + 1, dtype=np.int32) - 1

    def __str__(self):
        s = "id\tparent\tlchild\trchild\tlsib\trsib\tnsamp\tnchild\tedge\n"
        for j in range(len(self.parent)):
            s += (
                f"{j}\t{self.parent[j]}\t"
                f"{self.left_child[j]}\t{self.right_child[j]}\t"
                f"{self.left_sib[j]}\t{self.right_sib[j]}\t"
                f"{self.num_samples[j]}\t"
                f"{self.num_children[j]}\t"
                f"{self.edge[j]}\n"
            )
        return s

    def roots(self):
        roots = []
        u = self.left_child[-1]
        while u != -1:
            roots.append(u)
            u = self.right_sib[u]
        return roots

    def remove_branch(self, p, c):
        lsib = self.left_sib[c]
        rsib = self.right_sib[c]
        if lsib == -1:
            self.left_child[p] = rsib
        else:
            self.right_sib[lsib] = rsib
        if rsib == -1:
            self.right_child[p] = lsib
        else:
            self.left_sib[rsib] = lsib
        self.parent[c] = -1
        self.left_sib[c] = -1
        self.right_sib[c] = -1
        self.num_children[p] -= 1

    def insert_branch(self, p, c):
        assert self.parent[c] == -1, "contradictory edges"
        self.parent[c] = p
        u = self.right_child[p]
        if u == -1:
            self.left_child[p] = c
            self.left_sib[c] = -1
            self.right_sib[c] = -1
        else:
            self.right_sib[u] = c
            self.left_sib[c] = u
            self.right_sib[c] = -1
        self.right_child[p] = c
        self.num_children[p] += 1

    def is_potential_root(self, u):
        return self.num_samples[u] >= self.root_threshold

    # Note we cheat a bit here and use the -1 == last element semantics from Python.
    # We could use self.insert_branch(N, root) and then set self.parent[root] = -1.
    def insert_root(self, root):
        self.insert_branch(-1, root)

    def remove_root(self, root):
        self.remove_branch(-1, root)

    def remove_edge(self, edge):
        self.remove_branch(edge.parent, edge.child)
        self.num_edges -= 1
        self.edge[edge.child] = -1

        u = edge.parent
        while u != -1:
            path_end = u
            path_end_was_root = self.is_potential_root(u)
            self.num_samples[u] -= self.num_samples[edge.child]
            u = self.parent[u]

        if path_end_was_root and not self.is_potential_root(path_end):
            self.remove_root(path_end)
        if self.is_potential_root(edge.child):
            self.insert_root(edge.child)

    def insert_edge(self, edge):
        u = edge.parent
        while u != -1:
            path_end = u
            path_end_was_root = self.is_potential_root(u)
            self.num_samples[u] += self.num_samples[edge.child]
            u = self.parent[u]

        if self.is_potential_root(edge.child):
            self.remove_root(edge.child)
        if self.is_potential_root(path_end) and not path_end_was_root:
            self.insert_root(path_end)

        self.insert_branch(edge.parent, edge.child)
        self.num_edges += 1
        self.edge[edge.child] = edge.id


def algorithm_R(ts, root_threshold=1):
    """
    Quintuply linked tree with root tracking.
    """
    sequence_length = ts.sequence_length
    N = ts.num_nodes
    M = ts.num_edges
    tree = QuintuplyLinkedTree(N, root_threshold=root_threshold)
    edges = list(ts.edges())
    in_order = ts.tables.indexes.edge_insertion_order
    out_order = ts.tables.indexes.edge_removal_order

    # Initialise the tree
    for u in ts.samples():
        tree.num_samples[u] = 1
        if tree.is_potential_root(u):
            tree.insert_root(u)

    j = 0
    k = 0
    left = 0
    while j < M or left < sequence_length:
        while k < M and edges[out_order[k]].right == left:
            tree.remove_edge(edges[out_order[k]])
            k += 1
        while j < M and edges[in_order[j]].left == left:
            tree.insert_edge(edges[in_order[j]])
            j += 1
        right = sequence_length
        if j < M:
            right = min(right, edges[in_order[j]].left)
        if k < M:
            right = min(right, edges[out_order[k]].right)
        yield (left, right), tree
        left = right


class SampleListTree:
    """
    Straightforward implementation of the quintuply linked tree for developing
    and testing the sample lists feature.

    NOTE: The interface is pretty awkward; it's not intended for anything other
    than testing.
    """

    def __init__(self, tree_sequence, tracked_samples=None):
        self.tree_sequence = tree_sequence
        num_nodes = tree_sequence.num_nodes
        # Quintuply linked tree.
        self.parent = [-1 for _ in range(num_nodes)]
        self.left_sib = [-1 for _ in range(num_nodes)]
        self.right_sib = [-1 for _ in range(num_nodes)]
        self.left_child = [-1 for _ in range(num_nodes)]
        self.right_child = [-1 for _ in range(num_nodes)]
        self.left_sample = [-1 for _ in range(num_nodes)]
        self.right_sample = [-1 for _ in range(num_nodes)]
        # This is too long, but it's convenient for printing.
        self.next_sample = [-1 for _ in range(num_nodes)]

        self.sample_index_map = [-1 for _ in range(num_nodes)]
        samples = tracked_samples
        if tracked_samples is None:
            samples = list(tree_sequence.samples())
        for j in range(len(samples)):
            u = samples[j]
            self.sample_index_map[u] = j
            self.left_sample[u] = j
            self.right_sample[u] = j

    def __str__(self):
        fmt = "{:<5}{:>8}{:>8}{:>8}{:>8}{:>8}{:>8}{:>8}{:>8}\n"
        s = fmt.format(
            "node",
            "parent",
            "lsib",
            "rsib",
            "lchild",
            "rchild",
            "nsamp",
            "lsamp",
            "rsamp",
        )
        for u in range(self.tree_sequence.num_nodes):
            s += fmt.format(
                u,
                self.parent[u],
                self.left_sib[u],
                self.right_sib[u],
                self.left_child[u],
                self.right_child[u],
                self.next_sample[u],
                self.left_sample[u],
                self.right_sample[u],
            )
        # Strip off trailing newline
        return s[:-1]

    def remove_edge(self, edge):
        p = edge.parent
        c = edge.child
        lsib = self.left_sib[c]
        rsib = self.right_sib[c]
        if lsib == -1:
            self.left_child[p] = rsib
        else:
            self.right_sib[lsib] = rsib
        if rsib == -1:
            self.right_child[p] = lsib
        else:
            self.left_sib[rsib] = lsib
        self.parent[c] = -1
        self.left_sib[c] = -1
        self.right_sib[c] = -1

    def insert_edge(self, edge):
        p = edge.parent
        c = edge.child
        assert self.parent[c] == -1, "contradictory edges"
        self.parent[c] = p
        u = self.right_child[p]
        if u == -1:
            self.left_child[p] = c
            self.left_sib[c] = -1
            self.right_sib[c] = -1
        else:
            self.right_sib[u] = c
            self.left_sib[c] = u
            self.right_sib[c] = -1
        self.right_child[p] = c

    def update_sample_list(self, parent):
        # This can surely be done more efficiently and elegantly. We are iterating
        # up the tree and iterating over all the siblings of the nodes we visit,
        # rebuilding the links as we go. This results in visiting the same nodes
        # over again, which if we have nodes with many siblings will surely be
        # expensive. Another consequence of the current approach is that the
        # next pointer contains an arbitrary value for the rightmost sample of
        # every root. This should point to NULL ideally, but it's quite tricky
        # to do in practise. It's easier to have a slightly uglier iteration
        # over samples.
        #
        # In the future it would be good have a more efficient version of this
        # algorithm using next and prev pointers that we keep up to date at all
        # times, and which we use to patch the lists together more efficiently.
        u = parent
        while u != -1:
            sample_index = self.sample_index_map[u]
            if sample_index != -1:
                self.right_sample[u] = self.left_sample[u]
            else:
                self.right_sample[u] = -1
                self.left_sample[u] = -1
            v = self.left_child[u]
            while v != -1:
                if self.left_sample[v] != -1:
                    assert self.right_sample[v] != -1
                    if self.left_sample[u] == -1:
                        self.left_sample[u] = self.left_sample[v]
                        self.right_sample[u] = self.right_sample[v]
                    else:
                        self.next_sample[self.right_sample[u]] = self.left_sample[v]
                        self.right_sample[u] = self.right_sample[v]
                v = self.right_sib[v]
            u = self.parent[u]

    def sample_lists(self):
        """
        Iterate over the the trees in this tree sequence, yielding the (left, right)
        interval tuples. The tree state is maintained internally.

        See note above about the cruddiness of this interface.
        """
        ts = self.tree_sequence
        sequence_length = ts.sequence_length
        edges = list(ts.edges())
        M = len(edges)
        in_order = ts.tables.indexes.edge_insertion_order
        out_order = ts.tables.indexes.edge_removal_order
        j = 0
        k = 0
        left = 0

        while j < M or left < sequence_length:
            while k < M and edges[out_order[k]].right == left:
                edge = edges[out_order[k]]
                self.remove_edge(edge)
                self.update_sample_list(edge.parent)
                k += 1
            while j < M and edges[in_order[j]].left == left:
                edge = edges[in_order[j]]
                self.insert_edge(edge)
                self.update_sample_list(edge.parent)
                j += 1
            right = sequence_length
            if j < M:
                right = min(right, edges[in_order[j]].left)
            if k < M:
                right = min(right, edges[out_order[k]].right)
            yield left, right
            left = right


class LegacyRootThresholdTree:
    """
    Implementation of the quintuply linked tree with root tracking using the
    pre C 1.0/Python 0.4.0 algorithm. We keep this version around to make sure
    that we can be clear what the differences in the semantics of the new
    and old versions are.

    NOTE: The interface is pretty awkward; it's not intended for anything other
    than testing.
    """

    def __init__(self, tree_sequence, root_threshold=1):
        self.tree_sequence = tree_sequence
        self.root_threshold = root_threshold
        num_nodes = tree_sequence.num_nodes
        # Quintuply linked tree.
        self.parent = [-1 for _ in range(num_nodes)]
        self.left_sib = [-1 for _ in range(num_nodes)]
        self.right_sib = [-1 for _ in range(num_nodes)]
        self.left_child = [-1 for _ in range(num_nodes)]
        self.right_child = [-1 for _ in range(num_nodes)]
        self.num_samples = [0 for _ in range(num_nodes)]
        self.left_root = -1
        for u in tree_sequence.samples()[::-1]:
            self.num_samples[u] = 1
            if self.root_threshold == 1:
                self.add_root(u)

    def __str__(self):
        fmt = "{:<5}{:>8}{:>8}{:>8}{:>8}{:>8}{:>8}\n"
        s = f"roots = {self.roots()}\n"
        s += fmt.format("node", "parent", "lsib", "rsib", "lchild", "rchild", "nsamp")
        for u in range(self.tree_sequence.num_nodes):
            s += fmt.format(
                u,
                self.parent[u],
                self.left_sib[u],
                self.right_sib[u],
                self.left_child[u],
                self.right_child[u],
                self.num_samples[u],
            )
        # Strip off trailing newline
        return s[:-1]

    def is_root(self, u):
        return self.num_samples[u] >= self.root_threshold

    def roots(self):
        roots = []
        u = self.left_root
        while u != -1:
            roots.append(u)
            u = self.right_sib[u]
        return roots

    def add_root(self, root):
        if self.left_root != tskit.NULL:
            lroot = self.left_sib[self.left_root]
            if lroot != tskit.NULL:
                self.right_sib[lroot] = root
            self.left_sib[root] = lroot
            self.left_sib[self.left_root] = root
        self.right_sib[root] = self.left_root
        self.left_root = root

    def remove_root(self, root):
        lroot = self.left_sib[root]
        rroot = self.right_sib[root]
        self.left_root = tskit.NULL
        if lroot != tskit.NULL:
            self.right_sib[lroot] = rroot
            self.left_root = lroot
        if rroot != tskit.NULL:
            self.left_sib[rroot] = lroot
            self.left_root = rroot
        self.left_sib[root] = tskit.NULL
        self.right_sib[root] = tskit.NULL

    def remove_edge(self, edge):
        p = edge.parent
        c = edge.child
        lsib = self.left_sib[c]
        rsib = self.right_sib[c]
        if lsib == -1:
            self.left_child[p] = rsib
        else:
            self.right_sib[lsib] = rsib
        if rsib == -1:
            self.right_child[p] = lsib
        else:
            self.left_sib[rsib] = lsib
        self.parent[c] = -1
        self.left_sib[c] = -1
        self.right_sib[c] = -1

        u = edge.parent
        while u != -1:
            path_end = u
            path_end_was_root = self.is_root(u)
            self.num_samples[u] -= self.num_samples[c]
            u = self.parent[u]
        if path_end_was_root and not self.is_root(path_end):
            self.remove_root(path_end)
        if self.is_root(c):
            self.add_root(c)

    def insert_edge(self, edge):
        p = edge.parent
        c = edge.child
        assert self.parent[c] == -1, "contradictory edges"
        self.parent[c] = p
        u = self.right_child[p]
        lsib = self.left_sib[c]
        rsib = self.right_sib[c]
        if u == -1:
            self.left_child[p] = c
            self.left_sib[c] = -1
            self.right_sib[c] = -1
        else:
            self.right_sib[u] = c
            self.left_sib[c] = u
            self.right_sib[c] = -1
        self.right_child[p] = c

        u = edge.parent
        while u != -1:
            path_end = u
            path_end_was_root = self.is_root(u)
            self.num_samples[u] += self.num_samples[c]
            u = self.parent[u]

        if self.is_root(c):
            if path_end_was_root:
                # Remove c from root list.
                # Note: we don't use the remove_root function here because
                # it assumes that the node is at the end of a path
                self.left_root = tskit.NULL
                if lsib != tskit.NULL:
                    self.right_sib[lsib] = rsib
                    self.left_root = lsib
                if rsib != tskit.NULL:
                    self.left_sib[rsib] = lsib
                    self.left_root = rsib
            else:
                # Replace c with path_end in the root list
                if lsib != tskit.NULL:
                    self.right_sib[lsib] = path_end
                if rsib != tskit.NULL:
                    self.left_sib[rsib] = path_end
                self.left_sib[path_end] = lsib
                self.right_sib[path_end] = rsib
                self.left_root = path_end
        else:
            if self.is_root(path_end) and not path_end_was_root:
                self.add_root(path_end)

    def iterate(self):
        """
        Iterate over the the trees in this tree sequence, yielding the (left, right)
        interval tuples. The tree state is maintained internally.
        """
        ts = self.tree_sequence
        sequence_length = ts.sequence_length
        edges = list(ts.edges())
        M = len(edges)
        in_order = ts.tables.indexes.edge_insertion_order
        out_order = ts.tables.indexes.edge_removal_order
        j = 0
        k = 0
        left = 0

        while j < M or left < sequence_length:
            while k < M and edges[out_order[k]].right == left:
                edge = edges[out_order[k]]
                self.remove_edge(edge)
                k += 1
            while j < M and edges[in_order[j]].left == left:
                edge = edges[in_order[j]]
                self.insert_edge(edge)
                j += 1
            if self.left_root != tskit.NULL:
                while self.left_sib[self.left_root] != tskit.NULL:
                    self.left_root = self.left_sib[self.left_root]
            right = sequence_length
            if j < M:
                right = min(right, edges[in_order[j]].left)
            if k < M:
                right = min(right, edges[out_order[k]].right)
            yield left, right
            left = right


FORWARD = 1
REVERSE = -1


@dataclasses.dataclass
class Interval:
    left: float
    right: float

    def __iter__(self):
        yield self.left
        yield self.right


@dataclasses.dataclass
class EdgeRange:
    start: int
    stop: int
    order: typing.List


class TreePosition:
    def __init__(self, ts):
        self.ts = ts
        self.index = -1
        self.direction = 0
        self.interval = Interval(0, 0)
        self.in_range = EdgeRange(0, 0, None)
        self.out_range = EdgeRange(0, 0, None)

    def __str__(self):
        s = f"index: {self.index}\ninterval: {self.interval}\n"
        s += f"direction: {self.direction}\n"
        s += f"in_range: {self.in_range}\n"
        s += f"out_range: {self.out_range}\n"
        return s

    def assert_equal(self, other):
        assert self.index == other.index
        assert self.direction == other.direction
        assert self.interval == other.interval

    def set_null(self):
        self.index = -1
        self.interval.left = 0
        self.interval.right = 0

    def next(self):  # NOQA: A003
        M = self.ts.num_edges
        breakpoints = self.ts.breakpoints(as_array=True)
        left_coords = self.ts.edges_left
        left_order = self.ts.indexes_edge_insertion_order
        right_coords = self.ts.edges_right
        right_order = self.ts.indexes_edge_removal_order

        if self.index == -1:
            self.interval.right = 0
            self.out_range.stop = 0
            self.in_range.stop = 0
            self.direction = FORWARD

        if self.direction == FORWARD:
            left_current_index = self.in_range.stop
            right_current_index = self.out_range.stop
        else:
            left_current_index = self.out_range.stop + 1
            right_current_index = self.in_range.stop + 1

        left = self.interval.right

        j = right_current_index
        self.out_range.start = j
        while j < M and right_coords[right_order[j]] == left:
            j += 1
        self.out_range.stop = j
        self.out_range.order = right_order

        j = left_current_index
        self.in_range.start = j
        while j < M and left_coords[left_order[j]] == left:
            j += 1
        self.in_range.stop = j
        self.in_range.order = left_order

        self.direction = FORWARD
        self.index += 1
        if self.index == self.ts.num_trees:
            self.set_null()
        else:
            self.interval.left = left
            self.interval.right = breakpoints[self.index + 1]
        return self.index != -1

    def prev(self):
        M = self.ts.num_edges
        breakpoints = self.ts.breakpoints(as_array=True)
        right_coords = self.ts.edges_right
        right_order = self.ts.indexes_edge_removal_order
        left_coords = self.ts.edges_left
        left_order = self.ts.indexes_edge_insertion_order

        if self.index == -1:
            self.index = self.ts.num_trees
            self.interval.left = self.ts.sequence_length
            self.in_range.stop = M - 1
            self.out_range.stop = M - 1
            self.direction = REVERSE

        if self.direction == REVERSE:
            left_current_index = self.out_range.stop
            right_current_index = self.in_range.stop
        else:
            left_current_index = self.in_range.stop - 1
            right_current_index = self.out_range.stop - 1

        right = self.interval.left

        j = left_current_index
        self.out_range.start = j
        while j >= 0 and left_coords[left_order[j]] == right:
            j -= 1
        self.out_range.stop = j
        self.out_range.order = left_order

        j = right_current_index
        self.in_range.start = j
        while j >= 0 and right_coords[right_order[j]] == right:
            j -= 1
        self.in_range.stop = j
        self.in_range.order = right_order

        self.direction = REVERSE
        self.index -= 1
        if self.index == -1:
            self.set_null()
        else:
            self.interval.left = breakpoints[self.index]
            self.interval.right = right
        return self.index != -1

    def seek_forward(self, index):
        # NOTE this is still in development and not fully tested.
        assert index >= self.index and index < self.ts.num_trees
        M = self.ts.num_edges
        breakpoints = self.ts.breakpoints(as_array=True)
        left_coords = self.ts.edges_left
        left_order = self.ts.indexes_edge_insertion_order
        right_coords = self.ts.edges_right
        right_order = self.ts.indexes_edge_removal_order

        if self.index == -1:
            self.interval.right = 0
            self.out_range.stop = 0
            self.in_range.stop = 0
            self.direction = FORWARD

        if self.direction == FORWARD:
            left_current_index = self.in_range.stop
            right_current_index = self.out_range.stop
        else:
            left_current_index = self.out_range.stop + 1
            right_current_index = self.in_range.stop + 1

        self.direction = FORWARD
        left = breakpoints[index]

        # The range of edges we need consider for removal starts
        # at the current right index and ends at the first edge
        # where the right coordinate is equal to the new tree's
        # left coordinate.
        j = right_current_index
        self.out_range.start = j
        # TODO This could be done with binary search
        while j < M and right_coords[right_order[j]] <= left:
            j += 1
        self.out_range.stop = j

        if self.index == -1:
            # No edges, so out_range should be empty
            self.out_range.start = self.out_range.stop

        # The range of edges we need to consider for the new tree
        # must have right coordinate > left
        j = left_current_index
        while j < M and right_coords[left_order[j]] <= left:
            j += 1
        self.in_range.start = j
        # TODO this could be done with a binary search
        while j < M and left_coords[left_order[j]] <= left:
            j += 1
        self.in_range.stop = j

        self.interval.left = left
        self.interval.right = breakpoints[index + 1]
        self.out_range.order = right_order
        self.in_range.order = left_order
        self.index = index

    def seek_backward(self, index):
        # NOTE this is still in development and not fully tested.
        assert index >= 0
        M = self.ts.num_edges
        breakpoints = self.ts.breakpoints(as_array=True)
        left_coords = self.ts.edges_left
        left_order = self.ts.indexes_edge_insertion_order
        right_coords = self.ts.edges_right
        right_order = self.ts.indexes_edge_removal_order

        if self.index == -1:
            assert index < self.ts.num_trees
            self.index = self.ts.num_trees
            self.interval.left = self.ts.sequence_length
            self.in_range.stop = M - 1
            self.out_range.stop = M - 1
            self.direction = REVERSE
        else:
            assert index <= self.index

        if self.direction == REVERSE:
            left_current_index = self.out_range.stop
            right_current_index = self.in_range.stop
        else:
            left_current_index = self.in_range.stop - 1
            right_current_index = self.out_range.stop - 1

        self.direction = REVERSE
        right = breakpoints[index + 1]

        # The range of edges we need consider for removal starts
        # at the current left index and ends at the first edge
        # where the left coordinate is equal to the new tree's
        # right coordinate.
        j = left_current_index
        self.out_range.start = j
        # TODO This could be done with binary search
        while j >= 0 and left_coords[left_order[j]] >= right:
            j -= 1
        self.out_range.stop = j

        if self.index == self.ts.num_trees:
            # No edges, so out_range should be empty
            self.out_range.start = self.out_range.stop

        # The range of edges we need to consider for the new tree
        # must have left coordinate < right
        j = right_current_index
        while j >= 0 and left_coords[right_order[j]] >= right:
            j -= 1
        self.in_range.start = j
        # We stop at the first edge with right coordinate < right
        while j >= 0 and right_coords[right_order[j]] >= right:
            j -= 1
        self.in_range.stop = j

        self.interval.right = right
        self.interval.left = breakpoints[index]
        self.out_range.order = left_order
        self.in_range.order = right_order
        self.index = index

    def step(self, direction):
        if direction == FORWARD:
            return self.next()
        elif direction == REVERSE:
            return self.prev()
        else:
            raise ValueError("Direction must be FORWARD (+1) or REVERSE (-1)")


def mean_descendants(ts, reference_sets):
    """
    Returns the mean number of nodes from the specified reference sets
    where the node is ancestral to at least one of the reference nodes. Returns a
    ``(ts.num_nodes, len(reference_sets))`` dimensional numpy array.
    """
    # Check the inputs (could be done more efficiently here)
    all_reference_nodes = set()
    for reference_set in reference_sets:
        U = set(reference_set)
        if len(U) != len(reference_set):
            raise ValueError("Cannot have duplicate values within set")
        if len(all_reference_nodes & U) != 0:
            raise ValueError("Sample sets must be disjoint")
        all_reference_nodes |= U

    K = len(reference_sets)
    C = np.zeros((ts.num_nodes, K))
    parent = np.zeros(ts.num_nodes, dtype=int) - 1
    # The -1th element of ref_count is for all nodes in the reference set.
    ref_count = np.zeros((ts.num_nodes, K + 1), dtype=int)
    last_update = np.zeros(ts.num_nodes)
    total_span = np.zeros(ts.num_nodes)

    def update_counts(edge, left, sign):
        # Update the counts and statistics for a given node. Before we change the
        # node counts in the given direction, check to see if we need to update
        # statistics for that node. When a node count changes, we add the
        # accumulated statistic value for the span since that node was last updated.
        v = edge.parent
        while v != -1:
            if last_update[v] != left:
                if ref_count[v, K] > 0:
                    span = left - last_update[v]
                    C[v] += span * ref_count[v, :K]
                    total_span[v] += span
                last_update[v] = left
            ref_count[v] += sign * ref_count[edge.child]
            v = parent[v]

    # Set the intitial conditions.
    for j in range(K):
        ref_count[reference_sets[j], j] = 1
    ref_count[ts.samples(), K] = 1

    for (left, _right), edges_out, edges_in in ts.edge_diffs():
        for edge in edges_out:
            parent[edge.child] = -1
            update_counts(edge, left, -1)
        for edge in edges_in:
            parent[edge.child] = edge.parent
            update_counts(edge, left, +1)

    # Finally, add the stats for the last tree and divide by the total
    # span that each node was an ancestor to > 0 samples.
    for v in range(ts.num_nodes):
        if ref_count[v, K] > 0:
            span = ts.sequence_length - last_update[v]
            total_span[v] += span
            C[v] += span * ref_count[v, :K]
        if total_span[v] != 0:
            C[v] /= total_span[v]
    return C


def genealogical_nearest_neighbours(ts, focal, reference_sets):
    reference_set_map = np.zeros(ts.num_nodes, dtype=int) - 1
    for k, reference_set in enumerate(reference_sets):
        for u in reference_set:
            if reference_set_map[u] != -1:
                raise ValueError("Duplicate value in reference sets")
            reference_set_map[u] = k

    K = len(reference_sets)
    A = np.zeros((len(focal), K))
    L = np.zeros(len(focal))
    parent = np.zeros(ts.num_nodes, dtype=int) - 1
    sample_count = np.zeros((ts.num_nodes, K), dtype=int)

    # Set the initial conditions.
    for j in range(K):
        sample_count[reference_sets[j], j] = 1

    for (left, right), edges_out, edges_in in ts.edge_diffs():
        for edge in edges_out:
            parent[edge.child] = -1
            v = edge.parent
            while v != -1:
                sample_count[v] -= sample_count[edge.child]
                v = parent[v]
        for edge in edges_in:
            parent[edge.child] = edge.parent
            v = edge.parent
            while v != -1:
                sample_count[v] += sample_count[edge.child]
                v = parent[v]

        # Process this tree.
        for j, u in enumerate(focal):
            focal_reference_set = reference_set_map[u]
            delta = int(focal_reference_set != -1)
            p = u
            while p != tskit.NULL:
                total = np.sum(sample_count[p])
                if total > delta:
                    break
                p = parent[p]
            if p != tskit.NULL:
                span = right - left
                L[j] += span
                scale = span / (total - delta)
                for k, _reference_set in enumerate(reference_sets):
                    n = sample_count[p, k] - int(focal_reference_set == k)
                    A[j, k] += n * scale

    # Avoid division by zero
    L[L == 0] = 1
    A /= L.reshape((len(focal), 1))
    return A


def sort_individual_table(tables):
    """
    Sorts the individual table by parents-before-children.
    """

    individuals = tables.individuals
    num_individuals = individuals.num_rows

    # First find the set of individuals that have no children
    # by creating an array of incoming edge counts
    incoming_edge_count = np.zeros((num_individuals,), np.int64)
    for parent in individuals.parents:
        if parent != tskit.NULL:
            incoming_edge_count[parent] += 1

    todo = collections.deque()
    sorted_order = []
    for individual, num_edges in reversed(list(enumerate(incoming_edge_count))):
        if num_edges == 0:
            todo.append(individual)
            sorted_order.append(individual)
    # Now emit individuals from the set that have no children, removing their edges
    # as we go adding new individuals to the no children set.
    while len(todo) > 0:
        individual = todo.popleft()
        for parent in individuals[individual].parents:
            if parent != tskit.NULL:
                incoming_edge_count[parent] -= 1
                if incoming_edge_count[parent] == 0:
                    todo.append(parent)
                    sorted_order.append(parent)

    if np.sum(incoming_edge_count) > 0:
        raise ValueError("Individual pedigree has cycles")

    ind_id_map = {tskit.NULL: tskit.NULL}

    individuals_copy = tables.copy().individuals
    tables.individuals.clear()
    for row in reversed(sorted_order):
        ind_id_map[row] = tables.individuals.append(individuals_copy[row])
    tables.individuals.parents = [ind_id_map[i] for i in tables.individuals.parents]
    tables.nodes.individual = [ind_id_map[i] for i in tables.nodes.individual]

    return tables


def insert_unique_metadata(tables, table=None, offset=0):
    if isinstance(tables, tskit.TreeSequence):
        tables = tables.dump_tables()
    else:
        tables = tables.copy()
    if table is None:
        table = [t for t in tskit.TABLE_NAMES if t != "provenances"]
    for t in table:
        getattr(tables, t).packset_metadata(
            [struct.pack("I", offset + i) for i in range(getattr(tables, t).num_rows)]
        )
    return tables.tree_sequence()


def metadata_map(tables):
    # builds a mapping from metadata (as produced by insert_unique_metadata)
    # to ID for all the tables (except provenance)
    if isinstance(tables, tskit.TreeSequence):
        tables = tables.dump_tables()
    out = {}
    for t in [t for t in tskit.TABLE_NAMES if t != "provenances"]:
        out[t] = {}
        for j, x in enumerate(getattr(tables, t)):
            out[t][x.metadata] = j
    return out


@functools.lru_cache(maxsize=None)
def all_trees_ts(n):
    """
    Generate a tree sequence that corresponds to the lexicographic listing
    of all trees with n leaves (i.e. from tskit.all_trees(n)).

    Note: it would be nice to include a version of this in the combinatorics
    module at some point but the implementation is quite inefficient. Also
    it's not entirely clear that the way we're allocating node times is
    guaranteed to work.
    """
    tables = tskit.TableCollection(0)
    for _ in range(n):
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
    for j in range(1, n):
        tables.nodes.add_row(flags=0, time=j)

    L = 0
    for tree in tskit.all_trees(n):
        for u in tree.preorder()[1:]:
            tables.edges.add_row(L, L + 1, tree.parent(u), u)
        L += 1
    tables.sequence_length = L
    tables.sort()
    tables.simplify()
    return tables.tree_sequence()


def all_fields_ts(edge_metadata=True, migrations=True):
    """
    A tree sequence with data in all fields (except edge metadata is not set if
    edge_metadata is False and migrations are not defined if migrations is False
    (this is needed to test simplify, which doesn't allow either)

    """
    demography = msprime.Demography()
    demography.add_population(name="A", initial_size=10_000)
    demography.add_population(name="B", initial_size=5_000)
    demography.add_population(name="C", initial_size=1_000)
    demography.add_population(name="D", initial_size=500)
    demography.add_population(name="E", initial_size=100)
    demography.add_population_split(time=1000, derived=["A", "B"], ancestral="C")
    ts = msprime.sim_ancestry(
        samples={"A": 10, "B": 10},
        demography=demography,
        sequence_length=5,
        random_seed=42,
        recombination_rate=1,
        record_migrations=migrations,
        record_provenance=True,
    )
    ts = msprime.sim_mutations(ts, rate=0.001, random_seed=42)
    tables = ts.dump_tables()
    # Add locations to individuals
    individuals_copy = tables.individuals.copy()
    tables.individuals.clear()
    for i, individual in enumerate(individuals_copy):
        tables.individuals.append(
            individual.replace(flags=i, location=[i, i + 1], parents=[i - 1, i - 1])
        )
    # Ensure all columns have unique values
    nodes_copy = tables.nodes.copy()
    tables.nodes.clear()
    for i, node in enumerate(nodes_copy):
        tables.nodes.append(
            node.replace(
                flags=i,
                time=node.time + 0.00001 * i,
                individual=i % len(tables.individuals),
                population=i % len(tables.populations),
            )
        )
    if migrations:
        tables.migrations.add_row(left=0, right=1, node=21, source=1, dest=3, time=1001)

    # Add metadata
    for name, table in tables.table_name_map.items():
        if name == "provenances":
            continue
        if name == "migrations" and not migrations:
            continue
        if name == "edges" and not edge_metadata:
            continue
        table.metadata_schema = tskit.MetadataSchema.permissive_json()
        metadatas = [f'{{"foo":"n_{name}_{u}"}}' for u in range(len(table))]
        metadata, metadata_offset = tskit.pack_strings(metadatas)
        table.set_columns(
            **{
                **table.asdict(),
                "metadata": metadata,
                "metadata_offset": metadata_offset,
            }
        )
    tables.metadata_schema = tskit.MetadataSchema.permissive_json()
    tables.metadata = "Test metadata"
    tables.time_units = "Test time units"

    tables.reference_sequence.metadata_schema = tskit.MetadataSchema.permissive_json()
    tables.reference_sequence.metadata = "Test reference metadata"
    tables.reference_sequence.data = "A" * int(ts.sequence_length)
    # NOTE: it's unclear whether we'll want to have this set at the same time as
    # 'data', but it's useful to have something in all columns for now.
    tables.reference_sequence.url = "http://example.com/a_reference"

    # Add some more rows to provenance to have enough for testing.
    for i in range(3):
        tables.provenances.add_row(record="A", timestamp=str(i))

    return tables.tree_sequence()


--- ../../tskit/python/tests/test_version.py ---


"""
Test python package versioning
"""
from packaging.version import Version

from tskit import _version


class TestPythonVersion:
    """
    Test that the version is PEP440 compliant
    """

    def test_version(self):
        assert str(Version(_version.tskit_version)) == _version.tskit_version


--- ../../tskit/python/tests/test_reference_sequence.py ---


"""
Tests for reference sequence support.
"""
import pytest

import tskit


class TestTablesProperties:
    def test_initially_not_set(self):
        tables = tskit.TableCollection(1)
        assert not tables.has_reference_sequence()
        tables.reference_sequence.data = "ABCDEF"
        assert tables.reference_sequence.data == "ABCDEF"
        assert tables.has_reference_sequence()

    def test_does_not_have_reference_sequence_if_empty(self):
        tables = tskit.TableCollection(1)
        assert not tables.has_reference_sequence()
        tables.reference_sequence.data = ""
        assert not tables.has_reference_sequence()

    def test_same_object(self):
        tables = tskit.TableCollection(1)
        refseq = tables.reference_sequence
        tables.reference_sequence.data = "asdf"
        assert refseq.data == "asdf"
        # Not clear we want to do this, but keeping the same pattern as the
        # tables for now.
        assert tables.reference_sequence is not refseq

    def test_clear(self, ts_fixture):
        tables = ts_fixture.dump_tables()
        tables.reference_sequence.clear()
        assert not tables.has_reference_sequence()

    def test_write_object_fails_bad_type(self):
        tables = tskit.TableCollection(1)
        with pytest.raises(AttributeError):
            tables.reference_sequence = None

    def test_write_object(self, ts_fixture):
        tables = tskit.TableCollection(1)
        tables.reference_sequence = ts_fixture.reference_sequence
        tables.reference_sequence.assert_equals(ts_fixture.reference_sequence)

    def test_asdict_no_reference(self):
        tables = tskit.TableCollection(1)
        d = tables.asdict()
        assert "reference_sequence" not in d

    def test_asdict_reference_no_metadata(self):
        tables = tskit.TableCollection(1)
        tables.reference_sequence.data = "ABCDEF"
        d = tables.asdict()["reference_sequence"]
        assert d["data"] == "ABCDEF"
        assert d["url"] == ""
        assert "metadata" not in d
        assert "metadata_schema" not in d

    def test_asdict_reference_metadata(self):
        tables = tskit.TableCollection(1)
        tables.reference_sequence.metadata_schema = (
            tskit.MetadataSchema.permissive_json()
        )
        tables.reference_sequence.metadata = {"a": "ABCDEF"}
        d = tables.asdict()["reference_sequence"]
        assert d["data"] == ""
        assert d["url"] == ""
        assert d["metadata_schema"] == '{"codec":"json"}'
        assert d["metadata"] == b'{"a":"ABCDEF"}'

    def test_fromdict_reference_data(self):
        d = tskit.TableCollection(1).asdict()
        d["reference_sequence"] = {"data": "XYZ"}
        tables = tskit.TableCollection.fromdict(d)
        assert tables.has_reference_sequence()
        assert tables.reference_sequence.data == "XYZ"
        assert tables.reference_sequence.url == ""
        assert repr(tables.reference_sequence.metadata_schema) == ""
        assert tables.reference_sequence.metadata == b""

    def test_fromdict_reference_url(self):
        d = tskit.TableCollection(1).asdict()
        d["reference_sequence"] = {"url": "file://file.fasta"}
        tables = tskit.TableCollection.fromdict(d)
        assert tables.has_reference_sequence()
        assert tables.reference_sequence.data == ""
        assert tables.reference_sequence.url == "file://file.fasta"
        assert repr(tables.reference_sequence.metadata_schema) == ""
        assert tables.reference_sequence.metadata == b""

    def test_fromdict_reference_metadata(self):
        tables = tskit.TableCollection(1)
        tables.reference_sequence.metadata_schema = (
            tskit.MetadataSchema.permissive_json()
        )
        tables.reference_sequence.metadata = {"a": "ABCDEF"}
        tables = tskit.TableCollection.fromdict(tables.asdict())
        assert tables.has_reference_sequence()
        assert tables.reference_sequence.data == ""
        assert (
            tables.reference_sequence.metadata_schema
            == tskit.MetadataSchema.permissive_json()
        )
        assert tables.reference_sequence.metadata == {"a": "ABCDEF"}

    def test_fromdict_no_reference(self):
        d = tskit.TableCollection(1).asdict()
        tables = tskit.TableCollection.fromdict(d)
        assert not tables.has_reference_sequence()

    def test_fromdict_all_values_empty(self):
        d = tskit.TableCollection(1).asdict()
        d["reference_sequence"] = dict(
            data="", url="", metadata_schema="", metadata=b""
        )
        tables = tskit.TableCollection.fromdict(d)
        assert not tables.has_reference_sequence()


class TestSummaries:
    def test_repr(self):
        tables = tskit.TableCollection(1)
        refseq = tables.reference_sequence
        # TODO add better tests when summaries are updated
        assert repr(refseq).startswith("ReferenceSequence")


class TestEquals:
    def test_equal_self(self, ts_fixture):
        ts_fixture.reference_sequence.assert_equals(ts_fixture.reference_sequence)
        assert ts_fixture.reference_sequence == ts_fixture.reference_sequence
        assert not ts_fixture.reference_sequence != ts_fixture.reference_sequence
        assert ts_fixture.reference_sequence.equals(ts_fixture.reference_sequence)

    def test_equal_empty(self):
        tables = tskit.TableCollection(1)
        tables.reference_sequence.assert_equals(tables.reference_sequence)
        assert tables.reference_sequence == tables.reference_sequence
        assert tables.reference_sequence.equals(tables.reference_sequence)

    @pytest.mark.parametrize("attr", ["url", "data"])
    def test_unequal_attr_missing(self, ts_fixture, attr):
        t1 = ts_fixture.tables
        d = t1.asdict()
        del d["reference_sequence"][attr]
        t2 = tskit.TableCollection.fromdict(d)
        with pytest.raises(AssertionError, match=attr):
            t1.reference_sequence.assert_equals(t2.reference_sequence)
        assert t1.reference_sequence != t2.reference_sequence
        assert not t1.reference_sequence.equals(t2.reference_sequence)
        with pytest.raises(AssertionError, match=attr):
            t2.reference_sequence.assert_equals(t1.reference_sequence)
        assert t2.reference_sequence != t1.reference_sequence
        assert not t2.reference_sequence.equals(t1.reference_sequence)

    @pytest.mark.parametrize(
        ("attr", "val"),
        [
            ("url", "foo"),
            ("data", "bar"),
            ("metadata", {"json": "runs the world"}),
            ("metadata_schema", tskit.MetadataSchema(None)),
        ],
    )
    def test_different_not_equal(self, ts_fixture, attr, val):
        t1 = ts_fixture.dump_tables()
        t2 = t1.copy()
        setattr(t1.reference_sequence, attr, val)

        with pytest.raises(AssertionError):
            t1.reference_sequence.assert_equals(t2.reference_sequence)
        assert t1.reference_sequence != t2.reference_sequence
        assert not t1.reference_sequence.equals(t2.reference_sequence)
        with pytest.raises(AssertionError):
            t2.reference_sequence.assert_equals(t1.reference_sequence)
        assert t2.reference_sequence != t1.reference_sequence
        assert not t2.reference_sequence.equals(t1.reference_sequence)

    @pytest.mark.parametrize(
        ("attr", "val"),
        [
            ("metadata", {"json": "runs the world"}),
            ("metadata_schema", tskit.MetadataSchema(None)),
        ],
    )
    def test_different_but_ignore(self, ts_fixture, attr, val):
        t1 = ts_fixture.dump_tables()
        t2 = t1.copy()
        setattr(t1.reference_sequence, attr, val)

        with pytest.raises(AssertionError):
            t1.reference_sequence.assert_equals(t2.reference_sequence)
        assert t1.reference_sequence != t2.reference_sequence
        assert not t1.reference_sequence.equals(t2.reference_sequence)
        with pytest.raises(AssertionError):
            t2.reference_sequence.assert_equals(t1.reference_sequence)
        assert t2.reference_sequence != t1.reference_sequence
        assert not t2.reference_sequence.equals(t1.reference_sequence)

        t2.reference_sequence.assert_equals(t1.reference_sequence, ignore_metadata=True)
        assert t2.reference_sequence.equals(t1.reference_sequence, ignore_metadata=True)


class TestTreeSequenceProperties:
    @pytest.mark.parametrize("data", ["abcd", "🎄🌳🌴"])
    def test_data_inherited_from_tables(self, data):
        tables = tskit.TableCollection(1)
        tables.reference_sequence.data = data
        ts = tables.tree_sequence()
        assert ts.reference_sequence.data == data
        assert ts.has_reference_sequence()

    @pytest.mark.parametrize("url", ["http://xyx.z", "file://"])
    def test_url_inherited_from_tables(self, url):
        tables = tskit.TableCollection(1)
        tables.reference_sequence.url = url
        ts = tables.tree_sequence()
        assert ts.reference_sequence.url == url
        assert ts.has_reference_sequence()

    def test_no_reference_sequence(self):
        tables = tskit.TableCollection(1)
        ts = tables.tree_sequence()
        assert not ts.has_reference_sequence()
        assert ts.reference_sequence is None

    def test_write_data_fails(self):
        tables = tskit.TableCollection(1)
        tables.reference_sequence.data = "abc"
        ts = tables.tree_sequence()
        with pytest.raises(AttributeError, match="read-only"):
            ts.reference_sequence.data = "xyz"

    def test_write_url_fails(self):
        tables = tskit.TableCollection(1)
        tables.reference_sequence.data = "abc"
        ts = tables.tree_sequence()
        with pytest.raises(AttributeError, match="read-only"):
            ts.reference_sequence.url = "xyz"

    def test_write_metadata_fails(self):
        tables = tskit.TableCollection(1)
        tables.reference_sequence.data = "abc"
        ts = tables.tree_sequence()
        with pytest.raises(AttributeError, match="read-only"):
            # NOTE: it can be slightly confusing here because we try to encode
            # first, and so we don't get an AttributeError for all inputs.
            ts.reference_sequence.metadata = b"xyz"

    def test_write_metadata_schema_fails(self):
        tables = tskit.TableCollection(1)
        tables.reference_sequence.data = "abc"
        ts = tables.tree_sequence()
        with pytest.raises(AttributeError, match="read-only"):
            ts.reference_sequence.metadata_schema = (
                tskit.MetadataSchema.permissive_json()
            )

    def test_write_object_fails(self, ts_fixture):
        tables = tskit.TableCollection(1)
        ts = tables.tree_sequence()
        with pytest.raises(AttributeError):
            ts.reference_sequence = ts_fixture.reference_sequence


--- ../../tskit/python/tests/test_ibd.py ---

import collections
import io
import itertools

import msprime
import numpy as np
import pytest

import tests
import tests.ibd as ibd
import tests.test_wright_fisher as wf
import tskit
from tests.test_highlevel import get_example_tree_sequences

"""
Tests of IBD finding algorithms.
"""


# ↑ See https://github.com/tskit-dev/tskit/issues/1804 for when
# we can remove this. The example_ts here is intended to be the
# basic tree sequence which should give a meaningful result for
# most operations. Probably rename it to ``examples.simple_ts()``
# or something.


@tests.cached_example
def example_ts():
    return [msprime.sim_ancestry(2, random_seed=1)]


def ibd_segments(
    ts,
    *,
    within=None,
    between=None,
    min_span=0,
    max_time=None,
    compare_lib=True,
    print_c=False,
    print_py=False,
    squash=False
):
    """
    Calculates IBD segments using Python and converts output to lists of segments.
    Also compares result with C library.
    """
    ibd_f = ibd.IbdFinder(
        ts, within=within, between=between, max_time=max_time, min_span=min_span
    )
    ibd_segs = ibd_f.run(squash=squash)
    if print_py:
        print("Python output:\n")
        print(ibd_segs)
    # ibd_f.print_state()
    if compare_lib:
        c_out = ts.ibd_segments(
            within=within,
            between=between,
            max_time=max_time,
            min_span=min_span,
            store_segments=True,
        )
        if print_c:
            print("C output:\n")
            print(c_out)
        assert_ibd_equal(ibd_segs, c_out)
    return ibd_segs


def naive_ibd(ts, a, b):
    """
    Returns the IBD segments along the genome for a and b.
    """

    tree = ts.first()
    mrca = tree.mrca(a, b)
    last_mrca = mrca
    left = 0.0
    segs = []
    while tree.next():
        mrca = tree.mrca(a, b)
        if mrca != last_mrca:
            segs.append(tskit.IdentitySegment(left, tree.interval.left, last_mrca))
            left = tree.interval.left
            last_mrca = mrca
    segs.append(tskit.IdentitySegment(left, ts.sequence_length, last_mrca))

    # Filter out segments with no mrca
    return [seg for seg in segs if seg.node != -1]


def naive_ibd_all_pairs(ts, samples=None):
    samples = ts.samples() if samples is None else samples
    all_pairs_map = {
        (a, b): naive_ibd(ts, a, b) for a, b in itertools.combinations(samples, 2)
    }
    # Filter out pairs with empty segment lists
    return {key: value for key, value in all_pairs_map.items() if len(value) > 0}


class TestIbdDefinition:
    @pytest.mark.skip("help")
    @pytest.mark.xfail()
    @pytest.mark.parametrize("ts", get_example_tree_sequences(custom_max=15))
    def test_all_pairs(self, ts):
        if ts.num_samples > 10:
            samples = ts.samples()[:10]
            ts = ts.simplify(samples=samples)
        else:
            samples = ts.samples()
        ibd_lib = ts.ibd_segments(within=samples, store_segments=True)
        ibd_def = naive_ibd_all_pairs(ts, samples=samples)
        assert_ibd_equal(ibd_lib, ibd_def)

    @pytest.mark.skip("help")
    @pytest.mark.parametrize("ts", get_example_tree_sequences(custom_max=15))
    def test_all_pairs_python_only(self, ts):
        samples = ts.samples()[:10]
        ibd_pylib = ibd_segments(ts, within=samples, squash=True, compare_lib=False)
        ibd_def = naive_ibd_all_pairs(ts, samples=samples)
        assert_ibd_equal(ibd_pylib, ibd_def)

    @pytest.mark.skip("help")
    @pytest.mark.parametrize("N", [2, 5, 10])
    @pytest.mark.parametrize("T", [2, 5, 10])
    def test_wright_fisher_examples(self, N, T):
        tables = wf.wf_sim(N, T, deep_history=False, seed=42)
        tables.sort()
        # NB this is essential! We get spurious breakpoints otherwise
        tables.edges.squash()
        tables.sort()
        ts = tables.tree_sequence()
        ibd0 = ibd_segments(ts, squash=True, compare_lib=False)
        ibd1 = naive_ibd_all_pairs(ts)
        assert_ibd_equal(ibd0, ibd1)


class TestIbdImplementations:
    @pytest.mark.skip("help")
    @pytest.mark.xfail()
    @pytest.mark.parametrize("ts", get_example_tree_sequences(custom_max=15))
    def test_all_pairs(self, ts):
        # Automatically compares the two implementations
        samples = ts.samples()[:10]
        ts = ts.simplify(samples=samples)
        ibd_segments(ts, squash=True)


def assert_ibd_equal(dict1, dict2):
    """
    Verifies that two dictionaries have the same keys, and that
    the set of items corresponding to each key is identical.
    Used to check identical IBD output.
    """
    assert len(dict1) == len(dict2)
    for key, val in dict1.items():
        assert key in dict2
        assert len(val) == len(dict2[key])
        segs1 = list(sorted(val))
        segs2 = list(sorted(dict2[key]))
        assert segs1 == segs2


class TestIbdSingleBinaryTree:
    @tests.cached_example
    def ts(self):
        #
        # 2        4
        #         / \
        # 1      3   \
        #       / \   \
        # 0    0   1   2
        print("evaluating ts")
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       0           1
        4       0           2
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       3       0,1
        0       1       4       2,3
        """
        )
        return tskit.load_text(nodes=nodes, edges=edges, strict=False)

    # Basic test
    def test_defaults(self):
        true_segs = {
            (0, 1): [tskit.IdentitySegment(0.0, 1.0, 3)],
            (0, 2): [tskit.IdentitySegment(0.0, 1.0, 4)],
            (1, 2): [tskit.IdentitySegment(0.0, 1.0, 4)],
        }
        ibd_segs = ibd_segments(self.ts(), within=[0, 1, 2], squash=True)
        assert_ibd_equal(ibd_segs, true_segs)

    def test_within(self):
        true_segs = {
            (0, 1): [tskit.IdentitySegment(0.0, 1.0, 3)],
        }
        ibd_segs = ibd_segments(self.ts(), within=[0, 1], squash=True)
        assert_ibd_equal(ibd_segs, true_segs)

    def test_between_0_1(self):
        true_segs = {
            (0, 1): [tskit.IdentitySegment(0.0, 1.0, 3)],
        }
        ibd_segs = ibd_segments(self.ts(), between=[[0], [1]], squash=True)
        assert_ibd_equal(ibd_segs, true_segs)

    def test_between_0_2(self):
        true_segs = {
            (0, 2): [tskit.IdentitySegment(0.0, 1.0, 4)],
        }
        ibd_segs = ibd_segments(self.ts(), between=[[0], [2]], squash=True)
        assert_ibd_equal(ibd_segs, true_segs)

    def test_between_0_1_2(self):
        true_segs = {
            (0, 1): [tskit.IdentitySegment(0.0, 1.0, 3)],
            (0, 2): [tskit.IdentitySegment(0.0, 1.0, 4)],
            (1, 2): [tskit.IdentitySegment(0.0, 1.0, 4)],
        }
        ibd_segs = ibd_segments(self.ts(), between=[[0], [1], [2]], squash=True)
        assert_ibd_equal(ibd_segs, true_segs)

    def test_between_0_12(self):
        true_segs = {
            (0, 1): [tskit.IdentitySegment(0.0, 1.0, 3)],
            (0, 2): [tskit.IdentitySegment(0.0, 1.0, 4)],
        }
        ibd_segs = ibd_segments(self.ts(), between=[[0], [1, 2]], squash=True)
        assert_ibd_equal(ibd_segs, true_segs)

    def test_time(self):
        ibd_segs = ibd_segments(
            self.ts(),
            max_time=1.5,
            squash=True,
        )
        true_segs = {(0, 1): [tskit.IdentitySegment(0.0, 1.0, 3)]}
        assert_ibd_equal(ibd_segs, true_segs)

    def test_length(self):
        ibd_segs = ibd_segments(self.ts(), min_span=2, squash=True)
        assert_ibd_equal(ibd_segs, {})


class TestIbdInterface:
    @pytest.mark.parametrize("ts", example_ts())
    def test_input_errors_within(self, ts):
        with pytest.raises(tskit.LibraryError, match="Node out of bounds"):
            ts.ibd_segments(within=[-1])
        with pytest.raises(tskit.LibraryError, match="Duplicate sample value"):
            ts.ibd_segments(within=[0, 0])

    @pytest.mark.parametrize("ts", example_ts())
    def test_input_errors_between(self, ts):
        with pytest.raises(tskit.LibraryError, match="Node out of bounds"):
            ts.ibd_segments(between=[[0], [-1]])
        with pytest.raises(tskit.LibraryError, match="Duplicate sample"):
            ts.ibd_segments(between=[[0], [0]])

    @pytest.mark.parametrize("ts", example_ts())
    def test_within_between_mutually_exclusive(self, ts):
        with pytest.raises(ValueError, match="mutually exclusive"):
            ts.ibd_segments(within=[0], between=[1])

    @pytest.mark.parametrize("ts", example_ts())
    def test_tables_interface(self, ts):
        ibd_tab = ts.tables.ibd_segments(store_segments=True)
        ibd_ts = ts.ibd_segments(store_segments=True)
        assert ibd_tab == ibd_ts

    @pytest.mark.parametrize("ts", example_ts())
    def test_empty_within(self, ts):
        ibd = ts.ibd_segments(within=[], store_pairs=True)
        assert len(ibd) == 0

    @pytest.mark.parametrize("ts", example_ts())
    def test_empty_between(self, ts):
        ibd = ts.ibd_segments(between=[], store_pairs=True)
        assert len(ibd) == 0

    @pytest.mark.parametrize("ts", example_ts())
    def test_empty_in_between(self, ts):
        ibd = ts.ibd_segments(between=[[1, 2], []], store_pairs=True)
        assert len(ibd) == 0


class TestIbdTwoSamplesTwoTrees:
    # 2
    #             |     3
    # 1      2    |    / \
    #       / \   |   /   \
    # 0    0   1  |  0     1
    # |------------|----------|
    # 0.0          0.4        1.0

    @tests.cached_example
    def ts(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           1
        3       0           1.5
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       0.4     2       0,1
        0.4     1.0     3       0,1
        """
        )
        return tskit.load_text(nodes=nodes, edges=edges, strict=False)

    # Basic test
    def test_basic(self):
        ibd_segs = ibd_segments(self.ts(), squash=True)
        true_segs = {
            (0, 1): [
                tskit.IdentitySegment(0.0, 0.4, 2),
                tskit.IdentitySegment(0.4, 1.0, 3),
            ]
        }
        assert_ibd_equal(ibd_segs, true_segs)

    # Max time = 1.2
    def test_time(self):
        ibd_segs = ibd_segments(self.ts(), max_time=1.2, squash=True)
        true_segs = {(0, 1): [tskit.IdentitySegment(0.0, 0.4, 2)]}
        assert_ibd_equal(ibd_segs, true_segs)

    # Min length = 0.5
    def test_length(self):
        ibd_segs = ibd_segments(self.ts(), min_span=0.5, squash=True)
        true_segs = {(0, 1): [tskit.IdentitySegment(0.4, 1.0, 3)]}
        assert_ibd_equal(ibd_segs, true_segs)


class TestIbdUnrelatedSamples:
    #
    #    2   3
    #    |   |
    #    0   1

    @tests.cached_example
    def ts(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           1
        3       0           1
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       2       0
        0       1       3       1
        """
        )

        return tskit.load_text(nodes=nodes, edges=edges, strict=False)

    def test_basic(self):
        ibd_segs = ibd_segments(self.ts(), squash=True)
        assert len(ibd_segs) == 0

    def test_time(self):
        ibd_segs = ibd_segments(self.ts(), max_time=1.2, squash=True)
        assert len(ibd_segs) == 0

    def test_length(self):
        ibd_segs = ibd_segments(self.ts(), min_span=0.2, squash=True)
        assert len(ibd_segs) == 0


class TestIbdNoSamples:
    #
    #     2
    #    / \
    #   /   \
    #  /     \
    # (0)   (1)

    @tests.cached_example
    def ts(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       0           0
        1       0           0
        2       0           1
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       2       0
        0       1       2       1
        """
        )
        return tskit.load_text(nodes=nodes, edges=edges, strict=False)

    def test_defaults(self):
        result = ibd_segments(self.ts(), squash=True)
        assert len(result) == 0

    def test_specified_samples(self):
        ibd_segs = ibd_segments(self.ts(), within=[0, 1], squash=True)
        true_segs = {
            (0, 1): [
                tskit.IdentitySegment(0.0, 1, 2),
            ]
        }
        assert_ibd_equal(ibd_segs, true_segs)


class TestIbdSamplesAreDescendants:
    #
    # 4     5
    # |     |
    # 2     3
    # |     |
    # 0     1
    #
    @tests.cached_example
    def ts(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           1
        3       1           1
        4       0           2
        5       0           2
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       2       0
        0       1       3       1
        0       1       4       2
        0       1       5       3
        """
        )
        return tskit.load_text(nodes=nodes, edges=edges, strict=False)

    def test_basic(self):
        ibd_segs = ibd_segments(self.ts(), squash=True)
        true_segs = {
            (0, 2): [tskit.IdentitySegment(0.0, 1.0, 2)],
            (1, 3): [tskit.IdentitySegment(0.0, 1.0, 3)],
        }

        assert_ibd_equal(ibd_segs, true_segs)

    def test_input_within(self):
        ibd_segs = ibd_segments(self.ts(), within=[0, 2, 3, 5], squash=True)
        true_segs = {
            (0, 2): [tskit.IdentitySegment(0.0, 1.0, 2)],
            (3, 5): [tskit.IdentitySegment(0.0, 1.0, 5)],
        }
        assert_ibd_equal(ibd_segs, true_segs)

    def test_all_samples(self):
        # FIXME
        ibd_segs = ibd_segments(self.ts(), within=range(6), compare_lib=False)
        true_segs = {
            (0, 2): [tskit.IdentitySegment(0.0, 1.0, 2)],
            (0, 4): [tskit.IdentitySegment(0.0, 1.0, 4)],
            (2, 4): [tskit.IdentitySegment(0.0, 1.0, 4)],
            (1, 3): [tskit.IdentitySegment(0.0, 1.0, 3)],
            (1, 5): [tskit.IdentitySegment(0.0, 1.0, 5)],
            (3, 5): [tskit.IdentitySegment(0.0, 1.0, 5)],
        }
        assert_ibd_equal(ibd_segs, true_segs)


class TestIbdSimpleInternalSampleChain:
    #
    # 2
    # |
    # 1
    # |
    # 0
    #
    @tests.cached_example
    def ts(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           1
        2       1           2
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       1       0
        0       1       2       1
        """
        )
        return tskit.load_text(nodes=nodes, edges=edges, strict=False)

    def test_basic(self):
        # FIXME
        ibd_segs = ibd_segments(self.ts(), compare_lib=False, squash=True)
        true_segs = {
            (0, 1): [tskit.IdentitySegment(0.0, 1.0, 1)],
            (0, 2): [tskit.IdentitySegment(0.0, 1.0, 2)],
            (1, 2): [tskit.IdentitySegment(0.0, 1.0, 2)],
        }

        assert_ibd_equal(ibd_segs, true_segs)


class TestIbdDifferentPaths:
    #
    #        4       |      4       |        4
    #       / \      |     / \      |       / \
    #      /   \     |    /   3     |      /   \
    #     /     \    |   2     \    |     /     \
    #    /       \   |  /       \   |    /       \
    #   0         1  | 0         1  |   0         1
    #                |              |
    #                0.2            0.7

    @tests.cached_example
    def ts(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           1
        3       0           1.5
        4       0           2.5
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.2     0.7     2       0
        0.2     0.7     3       1
        0.0     0.2     4       0
        0.0     0.2     4       1
        0.7     1.0     4       0
        0.7     1.0     4       1
        0.2     0.7     4       2
        0.2     0.7     4       3
        """
        )
        return tskit.load_text(nodes=nodes, edges=edges, strict=False)

    def test_defaults(self):
        ibd_segs = ibd_segments(self.ts(), squash=True, compare_lib=False)
        true_segs = {
            (0, 1): [
                tskit.IdentitySegment(0.0, 1.0, 4),
            ]
        }
        assert_ibd_equal(ibd_segs, true_segs)

    def test_time(self):
        ibd_segs = ibd_segments(self.ts(), max_time=1.8)
        assert len(ibd_segs) == 0

    def test_length(self):
        ibd_segs = ibd_segments(self.ts(), min_span=0.6, squash=True, compare_lib=False)
        true_segs = {(0, 1): [tskit.IdentitySegment(0.0, 1.0, 4)]}
        assert_ibd_equal(ibd_segs, true_segs)


class TestIbdDifferentPaths2:
    #
    #        5         |
    #       / \        |
    #      /   4       |      4
    #     /   / \      |     / \
    #    /   /   \     |    /   \
    #   /   /     \    |   3     \
    #  /   /       \   |  / \     \
    # 0   1         2  | 0   2     1
    #                  |
    #                  0.2

    @tests.cached_example
    def ts(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       0           1
        4       0           2.5
        5       0           3.5
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.2     1.0     3       0
        0.2     1.0     3       2
        0.0     1.0     4       1
        0.0     0.2     4       2
        0.2     1.0     4       3
        0.0     0.2     5       0
        0.0     0.2     5       4
        """
        )
        return tskit.load_text(nodes=nodes, edges=edges, strict=False)

    def test_defaults(self):
        ibd_segs = ibd_segments(
            self.ts(), within=[1, 2], squash=True, compare_lib=False
        )
        true_segs = {
            (1, 2): [
                tskit.IdentitySegment(0.0, 1.0, 4),
            ],
        }
        assert_ibd_equal(ibd_segs, true_segs)

    def test_length(self):
        ibd_segs = ibd_segments(
            self.ts(), within=[1, 2], min_span=0.5, squash=True, compare_lib=False
        )
        true_segs = {
            (1, 2): [
                tskit.IdentitySegment(0.0, 1.0, 4),
            ],
        }
        assert_ibd_equal(ibd_segs, true_segs)


class TestIbdDifferentPaths3:
    # 2.00┊   4   ┊   4   ┊
    #     ┊ ┏━╋━┓ ┊ ┏━╋━┓ ┊
    # 1.00┊ 2 ┃ 3 ┊ 3 ┃ 2 ┊
    #     ┊ ┃ ┃   ┊ ┃ ┃   ┊
    # 0.00┊ 0 1   ┊ 0 1   ┊
    #     0       5      10
    @tests.cached_example
    def ts(self):
        t = tskit.TableCollection(sequence_length=10)
        t.nodes.add_row(flags=1, time=0)
        t.nodes.add_row(flags=1, time=0)
        t.nodes.add_row(flags=0, time=1)
        t.nodes.add_row(flags=0, time=1)
        t.nodes.add_row(flags=0, time=2)
        t.edges.add_row(parent=2, child=0, left=0, right=5)
        t.edges.add_row(parent=3, child=0, left=5, right=10)
        t.edges.add_row(parent=4, child=2, left=0, right=10)
        t.edges.add_row(parent=4, child=3, left=0, right=10)
        t.edges.add_row(parent=4, child=1, left=0, right=10)
        t.sort()
        return t.tree_sequence()

    def test_defaults(self):
        ibd_segs = ibd_segments(self.ts(), squash=True, compare_lib=False)
        true_segs = {
            (0, 1): [tskit.IdentitySegment(0, 10, 4)],
        }
        assert_ibd_equal(ibd_segs, true_segs)

    def test_length(self):
        ibd_segs = ibd_segments(self.ts(), min_span=0.6, squash=True, compare_lib=False)
        true_segs = {
            (0, 1): [tskit.IdentitySegment(0, 10, 4)],
        }
        assert_ibd_equal(ibd_segs, true_segs)


class TestIbdPolytomies:
    #
    #          5         |         5
    #         / \        |        / \
    #        4   \       |       4   \
    #       /|\   \      |      /|\   \
    #      / | \   \     |     / | \   \
    #     /  |  \   \    |    /  |  \   \
    #    /   |   \   \   |   /   |   \   \
    #   0    1    2   3  |  0    1    3   2
    #                    |
    #                   0.3

    @tests.cached_example
    def ts(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       1           0
        4       0           2.5
        5       0           3.5
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.0     1.0     4       0
        0.0     1.0     4       1
        0.0     0.3     4       2
        0.3     1.0     4       3
        0.3     1.0     5       2
        0.0     0.3     5       3
        0.0     1.0     5       4
        """
        )
        return tskit.load_text(nodes=nodes, edges=edges, strict=False)

    def test_defaults(self):
        ibd_segs = ibd_segments(self.ts(), squash=True, compare_lib=False)
        true_segs = {
            (0, 1): [tskit.IdentitySegment(0, 1, 4)],
            (0, 2): [
                tskit.IdentitySegment(0, 0.3, 4),
                tskit.IdentitySegment(0.3, 1, 5),
            ],
            (0, 3): [
                tskit.IdentitySegment(0, 0.3, 5),
                tskit.IdentitySegment(0.3, 1, 4),
            ],
            (1, 2): [
                tskit.IdentitySegment(0, 0.3, 4),
                tskit.IdentitySegment(0.3, 1, 5),
            ],
            (1, 3): [
                tskit.IdentitySegment(0, 0.3, 5),
                tskit.IdentitySegment(0.3, 1, 4),
            ],
            (2, 3): [
                tskit.IdentitySegment(0, 1, 5),
            ],
        }
        assert_ibd_equal(ibd_segs, true_segs)

    def test_time(self):
        ibd_segs = ibd_segments(self.ts(), max_time=3)
        true_segs = {
            (0, 1): [tskit.IdentitySegment(0, 1, 4)],
            (0, 2): [tskit.IdentitySegment(0, 0.3, 4)],
            (0, 3): [tskit.IdentitySegment(0.3, 1, 4)],
            (1, 2): [tskit.IdentitySegment(0, 0.3, 4)],
            (1, 3): [tskit.IdentitySegment(0.3, 1, 4)],
        }
        assert_ibd_equal(ibd_segs, true_segs)

    def test_length(self):
        ibd_segs = ibd_segments(self.ts(), min_span=0.5, squash=True, compare_lib=False)
        true_segs = {
            (0, 1): [tskit.IdentitySegment(0, 1, 4)],
            (0, 2): [tskit.IdentitySegment(0.3, 1, 5)],
            (0, 3): [tskit.IdentitySegment(0.3, 1, 4)],
            (1, 2): [tskit.IdentitySegment(0.3, 1, 5)],
            (1, 3): [tskit.IdentitySegment(0.3, 1, 4)],
            (2, 3): [tskit.IdentitySegment(0, 1, 5)],
        }
        (ibd_segs, true_segs)

    def test_input_within(self):
        ibd_segs = ibd_segments(self.ts(), within=[0, 1, 2])
        true_segs = {
            (0, 1): [tskit.IdentitySegment(0, 1, 4)],
            (0, 2): [
                tskit.IdentitySegment(0, 0.3, 4),
                tskit.IdentitySegment(0.3, 1, 5),
            ],
            (1, 2): [
                tskit.IdentitySegment(0, 0.3, 4),
                tskit.IdentitySegment(0.3, 1, 5),
            ],
        }
        assert_ibd_equal(ibd_segs, true_segs)


class TestIbdInternalSamples:
    #
    #
    #      3
    #     / \
    #    /   2
    #   /     \
    #  0      (1)

    @tests.cached_example
    def ts(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       0           0
        2       1           1
        3       0           2
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.0     1.0     2       1
        0.0     1.0     3       0
        0.0     1.0     3       2
        """
        )
        return tskit.load_text(nodes=nodes, edges=edges, strict=False)

    def test_defaults(self):
        ibd_segs = ibd_segments(self.ts())
        true_segs = {
            (0, 2): [tskit.IdentitySegment(0, 1, 3)],
        }
        assert_ibd_equal(ibd_segs, true_segs)


class TestIbdLengthThreshold:
    """
    Tests the behaviour of the min_span argument in niche cases.
    """

    # 2
    #             |     3
    # 1      2    |    / \
    #       / \   |   /   \
    # 0    0   1  |  0     1
    # |------------|----------|
    # 0.0          0.4        1.0

    @tests.cached_example
    def ts(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           1
        3       0           1.5
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       0.4     2       0,1
        0.4     1.0     3       0,1
        """
        )
        return tskit.load_text(nodes=nodes, edges=edges, strict=False)

    def test_length_exceeds_segment(self):
        ibd_segs = ibd_segments(self.ts(), min_span=1.1)
        assert_ibd_equal(ibd_segs, {})

    def test_length_is_negative(self):
        with pytest.raises(tskit.LibraryError):
            ibd_segments(self.ts(), min_span=-0.1)

    def test_equal_to_length(self):
        ibd_segs = ibd_segments(self.ts(), min_span=0.4)
        true_segs = {(0, 1): [tskit.IdentitySegment(0.4, 1.0, 3)]}
        assert_ibd_equal(ibd_segs, true_segs)


class TestIbdProperties:
    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_default_within_all_samples(self, ts):
        segs = ts.ibd_segments(store_pairs=True)
        for a, b in segs.keys():
            assert ts.node(a).is_sample()
            assert ts.node(b).is_sample()

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_within_subset(self, ts):
        samples = ts.samples()
        samples = samples[:3]
        segs = ts.ibd_segments(store_pairs=True, within=samples)
        for a, b in segs.keys():
            assert a in samples
            assert b in samples

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_between_two_subsets(self, ts):
        samples = ts.samples()
        k = len(samples) // 2
        A = samples[:k]
        B = samples[k:]
        segs = ts.ibd_segments(store_pairs=True, between=[A, B])
        for a, b in segs.keys():
            assert a in A
            assert b in B

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_between_same_segments_as_filtered_within_pair(self, ts):
        samples = ts.samples()[:10]
        all_segs = ts.ibd_segments(within=samples, store_segments=True)
        A = samples[1::2]
        B = samples[::2]
        between_segs = ts.ibd_segments(store_segments=True, between=[A, B])
        filtered_segs = collections.defaultdict(list)
        for (u, v), seglist in all_segs.items():
            if (u in A and v in B) or (v in A and u in B):
                filtered_segs[(u, v)] = seglist
        assert_ibd_equal(between_segs, filtered_segs)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_between_same_segments_as_filtered_within_triple(self, ts):
        samples = ts.samples()[:9]  # Limit the number of samples
        all_segs = ts.ibd_segments(within=samples, store_segments=True)
        A = samples[1::3]
        B = samples[2::3]
        C = samples[0::3]
        all_pairs = set()
        for set_pair in itertools.combinations([A, B, C], 2):
            for pair in itertools.product(*set_pair):
                all_pairs.add(tuple(sorted(pair)))
        between_segs = ts.ibd_segments(store_segments=True, between=[A, B, C])
        filtered_segs = collections.defaultdict(list)
        for pair, seglist in all_segs.items():
            if pair in all_pairs:
                filtered_segs[pair] = seglist
        assert_ibd_equal(between_segs, filtered_segs)


class TestIdentitySegments:
    """
    Test the IdentitySegments class interface.
    """

    def verify_segments(self, ts, ibd_segments):
        samples = set(ts.samples())
        for (a, b), segment_list in ibd_segments.items():
            assert a < b
            assert a in samples
            assert b in samples
            left = segment_list.left
            right = segment_list.right
            node = segment_list.node

            num_segments = 0
            total_span = 0
            for j, seg in enumerate(segment_list):
                assert isinstance(seg, tskit.IdentitySegment)
                total_span += seg.span
                num_segments += 1
                assert seg.span == seg.right - seg.left
                assert seg.left == left[j]
                assert seg.right == right[j]
                assert seg.node == node[j]
                assert 0 <= seg.node < ts.num_nodes
            assert total_span == segment_list.total_span
            assert num_segments == len(segment_list)

        total_span = sum(lst.total_span for lst in ibd_segments.values())
        np.testing.assert_allclose(ibd_segments.total_span, total_span)
        num_segments = sum(len(lst) for lst in ibd_segments.values())
        assert num_segments == ibd_segments.num_segments

    @pytest.mark.parametrize("store_segments", [True, False])
    @pytest.mark.parametrize("store_pairs", [True, False])
    def test_str(self, store_segments, store_pairs):
        ts = msprime.sim_ancestry(2, random_seed=2)
        result = ts.ibd_segments(store_segments=store_segments, store_pairs=store_pairs)
        s = str(result)
        assert "IdentitySegments" in s
        assert "max_time" in s
        assert "min_span" in s

    def test_repr_store_segments(self):
        ts = msprime.sim_ancestry(2, random_seed=2)
        result = ts.ibd_segments(store_segments=True)
        s = repr(result)
        assert s.startswith("IdentitySegments({")
        for lst in result.values():
            s = repr(lst)
            assert s.startswith("IdentitySegmentList([")

    def test_repr_without_store_segments(self):
        ts = msprime.sim_ancestry(2, random_seed=2)
        result = ts.ibd_segments(store_pairs=True)
        s = repr(result)
        assert s.startswith("<tskit.tables.IdentitySegments")
        result = ts.ibd_segments()
        s = repr(result)
        assert s.startswith("<tskit.tables.IdentitySegments")

    def test_store_segs_implies_store_pairs(self):
        ts = msprime.sim_ancestry(2, random_seed=2)
        result = ts.ibd_segments(store_pairs=True)
        assert result.num_pairs == 6
        result = ts.ibd_segments(store_segments=True)
        assert result.num_pairs == 6

    def test_operations_available_by_default(self):
        ts = msprime.sim_ancestry(2, random_seed=2)
        result = ts.ibd_segments()
        assert result.num_segments == 6
        assert result.total_span == 6
        with pytest.raises(tskit.IdentityPairsNotStoredError):
            _ = result.num_pairs
        with pytest.raises(tskit.IdentityPairsNotStoredError):
            _ = len(result)
        with pytest.raises(tskit.IdentityPairsNotStoredError):
            _ = result.pairs
        with pytest.raises(tskit.IdentityPairsNotStoredError):
            _ = result[0, 1]
        with pytest.raises(tskit.IdentityPairsNotStoredError):
            _ = list(result)
        with pytest.raises(tskit.IdentityPairsNotStoredError):
            _ = result == result
        # It's OK to when we compare with another type
        assert result != []

    def test_operations_available_store_pairs(self):
        ts = msprime.sim_ancestry(2, random_seed=2)
        result = ts.ibd_segments(store_pairs=True)
        assert result.num_segments == 6
        assert result.total_span == 6
        assert result.num_pairs == 6
        assert len(result) == 6
        assert result.pairs is not None
        seglist = result[0, 1]
        assert seglist.total_span == 1
        assert len(seglist) == 1
        with pytest.raises(tskit.IdentitySegmentsNotStoredError):
            _ = list(seglist)
        with pytest.raises(tskit.IdentitySegmentsNotStoredError):
            _ = seglist.left
        with pytest.raises(tskit.IdentitySegmentsNotStoredError):
            _ = seglist.right
        with pytest.raises(tskit.IdentitySegmentsNotStoredError):
            _ = seglist.node
        with pytest.raises(tskit.IdentitySegmentsNotStoredError):
            _ = seglist == seglist

    @pytest.mark.parametrize("n", [1, 2, 3])
    def test_pairs_all_samples(self, n):
        ts = msprime.sim_ancestry(n, random_seed=2)
        result = ts.ibd_segments(store_segments=True)
        pairs = np.array(list(itertools.combinations(ts.samples(), 2)))
        np.testing.assert_array_equal(pairs, result.pairs)
        self.verify_segments(ts, result)

    @pytest.mark.parametrize("n", [3, 4, 5])
    def test_pairs_subset(self, n):
        ts = msprime.sim_ancestry(n, random_seed=2)
        pairs = np.array([(0, 1), (0, 2), (1, 2)])
        result = ts.ibd_segments(within=[0, 1, 2], store_segments=True)
        np.testing.assert_array_equal(pairs, result.pairs)
        self.verify_segments(ts, result)

    @pytest.mark.parametrize("max_time", [0, 1, 10])
    def test_max_time(self, max_time):
        ts = msprime.sim_ancestry(2, random_seed=2)
        result = ts.ibd_segments(max_time=max_time, store_segments=True)
        assert result.max_time == max_time
        self.verify_segments(ts, result)

    def test_max_time_default(self):
        ts = msprime.sim_ancestry(2, random_seed=2)
        result = ts.ibd_segments(store_segments=True)
        assert np.isinf(result.max_time)
        self.verify_segments(ts, result)

    @pytest.mark.parametrize("min_span", [0, 1, 10])
    def test_min_span(self, min_span):
        ts = msprime.sim_ancestry(2, random_seed=2)
        result = ts.ibd_segments(min_span=min_span, store_segments=True)
        assert result.min_span == min_span
        self.verify_segments(ts, result)

    @pytest.mark.parametrize("min_span", [100, 101, 100000])
    def test_min_span_longer_than_seq_length(self, min_span):
        ts = msprime.sim_ancestry(
            100, recombination_rate=0.1, sequence_length=100, random_seed=2
        )
        result = ts.ibd_segments(min_span=min_span, store_segments=True)
        assert result.min_span == min_span
        assert result.num_segments == 0
        self.verify_segments(ts, result)

    def test_recombination_discrete(self):
        ts = msprime.sim_ancestry(
            10, sequence_length=100, recombination_rate=0.1, random_seed=2
        )
        assert ts.num_trees > 2
        result = ts.ibd_segments(store_segments=True)
        self.verify_segments(ts, result)

    def test_recombination_continuous(self):
        ts = msprime.sim_ancestry(
            10,
            recombination_rate=1,
            random_seed=2,
            discrete_genome=False,
            sequence_length=1,
        )
        assert ts.num_trees > 2
        result = ts.ibd_segments(store_segments=True)
        self.verify_segments(ts, result)

    def test_dict_interface(self):
        ts = msprime.sim_ancestry(5, random_seed=2)
        pairs = list(itertools.combinations(ts.samples(), 2))
        result = ts.ibd_segments(store_segments=True)
        assert len(result) == len(pairs)
        for pair in pairs:
            assert pair in result
            assert result[pair] is not None
        for k, v in result.items():
            assert k in pairs
            assert isinstance(v, tskit.IdentitySegmentList)


class TestIdentitySegmentsList:
    """
    Tests for the IdentitySegmentList class.
    """

    example_ts = msprime.sim_ancestry(
        3, sequence_length=100, recombination_rate=0.1, random_seed=2
    )

    def test_list_semantics(self):
        result = self.example_ts.ibd_segments(store_segments=True)
        assert len(result) > 0
        for seglist in result.values():
            lst = list(seglist)
            assert len(lst) == len(seglist)
            assert lst == list(seglist)

    def test_str(self):
        result = self.example_ts.ibd_segments(store_segments=True)
        seglist = list(result.values())[0]
        assert str(seglist).startswith("IdentitySegmentList")

    def test_repr(self):
        result = self.example_ts.ibd_segments(store_segments=True)
        seglist = list(result.values())[0]
        assert repr(seglist).startswith("IdentitySegmentList([IdentitySegment")

    def test_eq_semantics(self):
        result = self.example_ts.ibd_segments(store_segments=True)
        seglists = list(result.values())
        assert len(result) == len(seglists)
        assert len(seglists) > 1
        for seglist1, seglist2 in zip(result.values(), seglists):
            assert seglist1 == seglist2
            assert not (seglist1 != seglist2)
            assert seglist1 != result
            assert seglist1 != []
        # The chance of getting two identical seglists is miniscule
        for seglist in seglists[1:]:
            assert seglist != seglists[0]

    def test_eq_fails_without_store_segments(self):
        result = self.example_ts.ibd_segments(store_pairs=True)
        for seglist in result.values():
            with pytest.raises(tskit.IdentitySegmentsNotStoredError):
                _ = seglist == seglist
            # But it's OK when comparing to another type, since we know
            # it'll be False regardless
            assert seglist != []

    def test_list_contents(self):
        result = self.example_ts.ibd_segments(store_segments=True)
        assert len(result) > 0
        for seglist in result.values():
            for seg in seglist:
                assert isinstance(seg, tskit.IdentitySegment)


--- ../../tskit/python/tests/test_drawing.py ---


"""
Test cases for visualisation in tskit.
"""
import collections
import io
import logging
import math
import os
import pathlib
import platform
import re
import xml.dom.minidom
import xml.etree

import msprime
import numpy as np
import pytest
import xmlunittest

import tests.test_wright_fisher as wf
import tests.tsutil as tsutil
import tskit
from tskit import drawing


IS_WINDOWS = platform.system() == "Windows"


class TestTreeDraw:
    """
    Tests for the tree drawing functionality.
    TODO - the get_XXX_tree() functions should probably be placed in fixtures
    """

    def get_binary_tree(self):
        ts = msprime.simulate(10, random_seed=1, mutation_rate=1)
        return next(ts.trees())

    def get_nonbinary_ts(self):
        tables = wf.wf_sim(
            8,
            4,
            seed=1,
            deep_history=True,
            initial_generation_samples=False,
            num_loci=2,
        )
        tables.sort()
        ts = tables.tree_sequence().simplify()
        return tsutil.jukes_cantor(ts, 10, 0.025, seed=1)

    def get_nonbinary_tree(self):
        for t in self.get_nonbinary_ts().trees():
            for u in t.nodes():
                if len(t.children(u)) > 2:
                    return t
        raise AssertionError()

    def get_zero_edge_tree(self):
        tables = tskit.TableCollection(sequence_length=2)
        # These must be samples or we will have zero roots.
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        tables.sites.add_row(position=0, ancestral_state="0")
        tables.mutations.add_row(site=0, node=0, derived_state="1")
        tables.mutations.add_row(site=0, node=1, derived_state="1")
        return tables.tree_sequence().first()

    def get_zero_roots_tree(self):
        tables = tskit.TableCollection(sequence_length=2)
        # If we have no samples we have zero roots
        tables.nodes.add_row(time=0)
        tables.nodes.add_row(time=0)
        tables.nodes.add_row(time=1)
        tables.edges.add_row(0, 2, 2, 0)
        tables.edges.add_row(0, 2, 2, 1)
        tree = tables.tree_sequence().first()
        assert tree.num_roots == 0
        return tree

    def get_multiroot_tree(self):
        ts = msprime.simulate(15, random_seed=1)
        # Take off the top quarter of edges
        tables = ts.dump_tables()
        edges = tables.edges
        n = len(edges) - len(edges) // 4
        edges.set_columns(
            left=edges.left[:n],
            right=edges.right[:n],
            parent=edges.parent[:n],
            child=edges.child[:n],
        )
        ts = tables.tree_sequence()
        for t in ts.trees():
            if t.num_roots > 1:
                return t
        raise AssertionError()

    def get_mutations_over_roots_tree(self):
        ts = msprime.simulate(15, random_seed=1)
        ts = ts.decapitate(ts.tables.nodes.time[-1] / 2)
        tables = ts.dump_tables()
        delta = 1.0 / (ts.num_nodes + 1)
        x = 0
        for node in range(ts.num_nodes):
            site_id = tables.sites.add_row(x, ancestral_state="0")
            x += delta
            tables.mutations.add_row(site_id, node=node, derived_state="1")
        ts = tables.tree_sequence()
        tree = ts.first()
        assert any(tree.parent(mut.node) == tskit.NULL for mut in tree.mutations())
        return tree

    def get_unary_node_tree(self):
        ts = msprime.simulate(2, random_seed=1)
        tables = ts.dump_tables()
        edges = tables.edges
        # Take out all the edges except 1
        n = 1
        edges.set_columns(
            left=edges.left[:n],
            right=edges.right[:n],
            parent=edges.parent[:n],
            child=edges.child[:n],
        )
        ts = tables.tree_sequence()
        for t in ts.trees():
            for u in t.nodes():
                if len(t.children(u)) == 1:
                    return t
        raise AssertionError()

    def get_empty_tree(self):
        tables = tskit.TableCollection(sequence_length=1)
        ts = tables.tree_sequence()
        return next(ts.trees())

    def get_simple_ts(self, use_mutation_times=False):
        """
        return a simple tree seq that does not depend on msprime
        """
        nodes = io.StringIO(
            """\
        id      is_sample   population      individual      time    metadata
        0       1       0       -1      0
        1       1       0       -1      0
        2       1       0       -1      0
        3       1       0       -1      0
        4       0       0       -1      0.1145014598813
        5       0       0       -1      1.11067965364865
        6       0       0       -1      1.75005250750382
        7       0       0       -1      5.31067154311640
        8       0       0       -1      6.57331354884652
        9       0       0       -1      9.08308317451295
        """
        )
        edges = io.StringIO(
            """\
        id      left            right           parent  child
        0       0.00000000      1.00000000      4       0
        1       0.00000000      1.00000000      4       1
        2       0.00000000      1.00000000      5       2
        3       0.00000000      1.00000000      5       3
        4       0.79258618      0.90634460      6       4
        5       0.79258618      0.90634460      6       5
        6       0.05975243      0.79258618      7       4
        7       0.90634460      0.91029435      7       4
        8       0.05975243      0.79258618      7       5
        9       0.90634460      0.91029435      7       5
        10      0.91029435      1.00000000      8       4
        11      0.91029435      1.00000000      8       5
        12      0.00000000      0.05975243      9       4
        13      0.00000000      0.05975243      9       5
        """
        )
        sites = io.StringIO(
            """\
        position      ancestral_state
        0.05          A
        0.06          0
        0.3           Empty
        0.5           XXX
        0.91          T
        """
        )
        muts = io.StringIO(
            """\
        site   node    derived_state    parent    time
        0      9       T                -1        15
        0      9       G                0         9.1
        0      5       1                1         9
        1      4       C                -1        1.6
        1      4       G                3         1.5
        2      7       G                -1        10
        2      3       C                5         1
        4      3       G                -1        1
        """
        )
        ts = tskit.load_text(nodes, edges, sites=sites, mutations=muts, strict=False)
        if use_mutation_times:
            return ts
        tables = ts.dump_tables()
        tables.mutations.time = np.full_like(tables.mutations.time, tskit.UNKNOWN_TIME)
        return tables.tree_sequence()

    def get_ts_varying_min_times(self, *args, **kwargs):
        """
        Like get_simple_ts but return a tree sequence with negative times, and some trees
        with different min times (i.e. with dangling nonsample nodes at negative times)
        """
        ts = self.get_simple_ts(*args, **kwargs)
        tables = ts.dump_tables()
        time = tables.nodes.time
        time[time == 0] = 0.1
        time[3] = -9.99
        tables.nodes.time = time
        # set node 3 to be non-sample node lower than the rest
        flags = tables.nodes.flags
        flags[3] = 0
        tables.nodes.flags = flags
        edges = tables.edges
        assert edges[3].child == 3 and edges[3].parent == 5
        edges[3] = edges[3].replace(left=ts.breakpoints(True)[1])
        tables.sort()
        tables.nodes.flags = flags
        return tables.tree_sequence()

    def fail(self, *args, **kwargs):
        """
        Required for xmlunittest.XmlTestMixin to work with pytest not unittest
        """
        pytest.fail(*args, **kwargs)


def closest_left_node(tree, u):
    """
    Returns the node that is closest to u in a left-to-right sense.
    """
    ret = tskit.NULL
    while u != tskit.NULL and ret == tskit.NULL:
        ret = tree.left_sib(u)
        u = tree.parent(u)
    return ret


def get_left_neighbour(tree, traversal_order):
    """
    This is a less efficient version of the get_left_neighbour function in
    drawing.py.
    """
    # Note: roots are the children of -1 here.
    children = collections.defaultdict(list)
    for u in tree.nodes(order=traversal_order):
        parent = tree.parent(u)
        children[parent].append(u)

    left_neighbour = np.full(tree.tree_sequence.num_nodes, tskit.NULL, dtype=int)
    for u in tree.nodes():
        next_left = tskit.NULL
        child = u
        while child != tskit.NULL and next_left == tskit.NULL:
            parent = tree.parent(child)
            child_index = children[parent].index(child)
            if child_index > 0:
                next_left = children[parent][child_index - 1]
            child = parent
        left_neighbour[u] = next_left
    return left_neighbour


class TestClosestLeftNode(TestTreeDraw):
    """
    Tests the code for finding the closest left node in a tree.
    """

    def verify(self, tree):
        m1 = drawing.get_left_neighbour(tree, "postorder")
        m2 = get_left_neighbour(tree, "postorder")
        np.testing.assert_array_equal(m1, m2)
        for u in tree.nodes():
            assert m1[u] == closest_left_node(tree, u)

        m1 = drawing.get_left_neighbour(tree, "minlex_postorder")
        m2 = get_left_neighbour(tree, "minlex_postorder")
        np.testing.assert_array_equal(m1, m2)

    def test_2_binary(self):
        ts = msprime.simulate(2, random_seed=2)
        self.verify(ts.first())

    def test_5_binary(self):
        ts = msprime.simulate(5, random_seed=2)
        self.verify(ts.first())

    def test_10_binary(self):
        ts = msprime.simulate(10, random_seed=2)
        self.verify(ts.first())

    def test_20_binary(self):
        ts = msprime.simulate(20, random_seed=3)
        self.verify(ts.first())

    def test_nonbinary(self):
        self.verify(self.get_nonbinary_tree())

    def test_zero_edge(self):
        self.verify(self.get_zero_edge_tree())

    def test_zero_roots(self):
        self.verify(self.get_zero_roots_tree())

    def test_multiroot(self):
        self.verify(self.get_multiroot_tree())

    def test_left_child(self):
        t = self.get_nonbinary_tree()
        left_child = drawing.get_left_child(t, t.postorder())
        for u in t.nodes(order="postorder"):
            if t.num_children(u) > 0:
                assert left_child[u] == t.children(u)[0]

    def test_null_node_left_child(self):
        t = self.get_nonbinary_tree()
        arr = list(t.nodes(order="minlex_postorder"))
        left_child = drawing.get_left_child(t, arr)
        assert left_child[tskit.NULL] == tskit.NULL

    def test_leaf_node_left_child(self):
        t = self.get_nonbinary_tree()
        arr = list(t.nodes(order="minlex_postorder"))
        left_child = drawing.get_left_child(t, arr)
        for u in t.samples():
            assert left_child[u] == tskit.NULL


class TestOrder(TestTreeDraw):
    """
    Tests for using the different node orderings.
    """

    def test_bad_order(self):
        for bad_order in [("sdf"), "sdf", 1234, ""]:
            with pytest.raises(ValueError):
                drawing.check_order(bad_order)

    def test_default_order(self):
        traversal_order = drawing.check_order(None)
        assert traversal_order == "minlex_postorder"

    def test_order_mapping(self):
        assert drawing.check_order("tree") == "postorder"
        assert drawing.check_order("minlex") == "minlex_postorder"

    def test_tree_svg_variants(self):
        t = self.get_binary_tree()
        output1 = t.draw(format="svg")
        output2 = t.draw(format="svg", order="minlex")
        output3 = t.draw(format="svg", order="tree")
        # Default is minlex
        assert output1 == output2
        # tree is at least different to minlex
        assert output1 != output3
        # draw_svg gets the same results
        assert t.draw_svg() == output1
        assert t.draw_svg(order="minlex") == output1
        assert t.draw_svg(order="tree") == output3

    def test_tree_text_variants(self):
        t = self.get_binary_tree()
        output1 = t.draw(format="unicode")
        output2 = t.draw(format="unicode", order="minlex")
        output3 = t.draw(format="unicode", order="tree")
        # Default is minlex
        assert output1 == output2
        # tree is at least different to minlex
        assert output1 != output3
        # draw_text gets the same results
        assert t.draw_text() == output1
        assert t.draw_text(order="minlex") == output1
        assert t.draw_text(order="tree") == output3

    def test_tree_sequence_text_variants(self):
        ts = msprime.simulate(10, random_seed=2)
        output1 = ts.draw_text()
        output2 = ts.draw_text(order="minlex")
        output3 = ts.draw_text(order="tree")

        # Default is minlex
        assert output1 == output2
        # tree is at least different to minlex
        assert output1 != output3

    def test_tree_sequence_svg_variants(self):
        ts = msprime.simulate(10, random_seed=2)
        output1 = ts.draw_svg()
        output2 = ts.draw_svg(order="minlex")
        output3 = ts.draw_svg(order="tree")

        # Default is minlex
        assert output1 == output2
        # tree is at least different to minlex
        assert output1 != output3


class TestFormats(TestTreeDraw):
    """
    Tests that formats are recognised correctly.
    """

    def test_svg_variants(self):
        t = self.get_binary_tree()
        for svg in ["svg", "SVG", "sVg"]:
            output = t.draw(format=svg)
            root = xml.etree.ElementTree.fromstring(output)
            assert root.tag == "{http://www.w3.org/2000/svg}svg"

    def test_default(self):
        # Default is SVG
        t = self.get_binary_tree()
        output = t.draw(format=None)
        root = xml.etree.ElementTree.fromstring(output)
        assert root.tag == "{http://www.w3.org/2000/svg}svg"
        output = t.draw()
        root = xml.etree.ElementTree.fromstring(output)
        assert root.tag == "{http://www.w3.org/2000/svg}svg"

    def test_ascii_variants(self):
        t = self.get_binary_tree()
        for fmt in ["ascii", "ASCII", "AScii"]:
            output = t.draw(format=fmt)
            with pytest.raises(xml.etree.ElementTree.ParseError):
                xml.etree.ElementTree.fromstring(
                    output,
                )

    def test_unicode_variants(self):
        t = self.get_binary_tree()
        for fmt in ["unicode", "UNICODE", "uniCODE"]:
            output = t.draw(format=fmt)
            with pytest.raises(xml.etree.ElementTree.ParseError):
                xml.etree.ElementTree.fromstring(
                    output,
                )

    def test_bad_formats(self):
        t = self.get_binary_tree()
        for bad_format in ["", "ASC", "SV", "jpeg"]:
            with pytest.raises(ValueError):
                t.draw(format=bad_format)


class TestDrawText(TestTreeDraw):
    """
    Tests the ASCII tree drawing method.
    """

    drawing_format = "ascii"
    example_label = "XXX"

    def verify_basic_text(self, text):
        assert isinstance(text, str)
        # TODO surely something else we can verify about this...

    def test_draw_defaults(self):
        t = self.get_binary_tree()
        text = t.draw(format=self.drawing_format)
        self.verify_basic_text(text)

    def test_draw_nonbinary(self):
        t = self.get_nonbinary_tree()
        text = t.draw(format=self.drawing_format)
        self.verify_basic_text(text)

    def test_draw_multiroot(self):
        t = self.get_multiroot_tree()
        text = t.draw(format=self.drawing_format)
        self.verify_basic_text(text)

    def test_draw_mutations_over_roots(self):
        t = self.get_mutations_over_roots_tree()
        text = t.draw(format=self.drawing_format)
        self.verify_basic_text(text)

    def test_draw_unary(self):
        t = self.get_unary_node_tree()
        text = t.draw(format=self.drawing_format)
        self.verify_basic_text(text)

    def test_draw_empty_tree(self):
        t = self.get_empty_tree()
        with pytest.raises(ValueError):
            t.draw(format=self.drawing_format)

    def test_draw_zero_roots_tree(self):
        t = self.get_zero_roots_tree()
        with pytest.raises(ValueError):
            t.draw(format=self.drawing_format)

    def test_draw_zero_edge_tree(self):
        t = self.get_zero_edge_tree()
        text = t.draw(format=self.drawing_format)
        self.verify_basic_text(text)

    def test_even_num_children_tree(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time
        0   1           0
        1   1           1
        2   1           2
        3   1           1
        4   1           4
        5   1           5
        6   1           7
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       6       0
        0       1       6       1
        0       1       6       2
        0       1       6       3
        0       1       6       4
        0       1       6       5
        """
        )
        ts = tskit.load_text(nodes, edges, strict=False)
        t = next(ts.trees())
        text = t.draw(format=self.drawing_format)
        self.verify_basic_text(text)

    def test_odd_num_children_tree(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time
        0   1           0
        1   1           1
        2   1           2
        3   1           1
        4   1           4
        5   1           5
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       5       0
        0       1       5       1
        0       1       5       2
        0       1       5       3
        0       1       5       4
        """
        )
        ts = tskit.load_text(nodes, edges, strict=False)
        t = next(ts.trees())
        text = t.draw(format=self.drawing_format)
        self.verify_basic_text(text)

    def test_node_labels(self):
        t = self.get_binary_tree()
        labels = {u: self.example_label for u in t.nodes()}
        text = t.draw(format=self.drawing_format, node_labels=labels)
        self.verify_basic_text(text)
        j = 0
        for _ in t.nodes():
            j = text[j:].find(self.example_label)
            assert j != -1

    def test_long_internal_labels(self):
        t = self.get_binary_tree()
        labels = {u: "X" * 10 for u in t.nodes() if t.is_internal(u)}
        text = t.draw(format=self.drawing_format, node_labels=labels)
        self.verify_basic_text(text)

    def test_no_node_labels(self):
        t = self.get_binary_tree()
        labels = {}
        text = t.draw(format=self.drawing_format, node_labels=labels)
        self.verify_basic_text(text)
        for u in t.nodes():
            assert text.find(str(u)) == -1

    def test_unused_args(self):
        t = self.get_binary_tree()
        with pytest.raises(ValueError):
            t.draw(format=self.drawing_format, width=300)
        with pytest.raises(ValueError):
            t.draw(format=self.drawing_format, height=300)
        with pytest.raises(ValueError):
            t.draw(format=self.drawing_format, mutation_labels={})
        with pytest.raises(ValueError):
            t.draw(format=self.drawing_format, mutation_colours={})
        with pytest.raises(ValueError):
            t.draw(format=self.drawing_format, edge_colours={})
        with pytest.raises(ValueError):
            t.draw(format=self.drawing_format, node_colours={})
        with pytest.raises(ValueError):
            t.draw(format=self.drawing_format, max_time=1234)
        with pytest.raises(ValueError):
            t.draw(format=self.drawing_format, min_time=1234)
        with pytest.raises(ValueError):
            with pytest.warns(FutureWarning):
                t.draw(format=self.drawing_format, max_tree_height=1234)
        with pytest.raises(ValueError):
            t.draw(format=self.drawing_format, time_scale="time")
        with pytest.raises(ValueError):
            with pytest.warns(FutureWarning):
                t.draw(format=self.drawing_format, tree_height_scale="time")


class TestDrawUnicode(TestDrawText):
    """
    Tests the Unicode tree drawing method
    """

    drawing_format = "unicode"
    example_label = "\u20ac" * 10  # euro symbol


class TestDrawTextErrors:
    """
    Tests for errors occuring in tree drawing code.
    """

    def test_bad_orientation(self):
        t = msprime.simulate(5, mutation_rate=0.1, random_seed=2).first()
        for bad_orientation in ["", "leftright", "sdf"]:
            with pytest.raises(ValueError):
                t.draw_text(orientation=bad_orientation)


class TestDrawTextExamples(TestTreeDraw):
    """
    Verify that we get the correct rendering for some examples.
    """

    def verify_text_rendering(self, drawn, drawn_tree, debug=False):
        if debug:
            print("Drawn:")
            print(drawn)
            print("Expected:")
            print(drawn_tree)
        tree_lines = drawn_tree.splitlines()
        drawn_lines = drawn.splitlines()
        assert len(tree_lines) == len(drawn_lines)
        for l1, l2 in zip(tree_lines, drawn_lines):
            # Trailing white space isn't significant.
            assert l1.rstrip() == l2.rstrip()

    def test_simple_tree(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time
        0   1           0
        1   1           0
        2   1           2
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       2       0
        0       1       2       1
        """
        )
        tree = (
            # fmt: off
            " 2 \n"
            "┏┻┓\n"
            "0 1"
            # fmt: on
        )
        ts = tskit.load_text(nodes, edges, strict=False)
        t = next(ts.trees())
        drawn = t.draw(format="unicode", order="tree")
        self.verify_text_rendering(drawn, tree)
        drawn = t.draw_text()
        self.verify_text_rendering(drawn, tree)

        tree = (
            # fmt: off
            " 2 \n"
            "+++\n"
            "0 1\n"
            # fmt: on
        )
        drawn = t.draw_text(use_ascii=True, order="tree")
        self.verify_text_rendering(drawn, tree)

        tree = (
            # fmt: off
            " ┏0\n"
            "2┫  \n"
            " ┗1\n"
            # fmt: on
        )
        drawn = t.draw_text(orientation="left", order="tree")
        self.verify_text_rendering(drawn, tree)
        tree = (
            # fmt: off
            " +0\n"
            "2+  \n"
            " +1\n"
            # fmt: on
        )
        drawn = t.draw_text(orientation="left", use_ascii=True, order="tree")
        self.verify_text_rendering(drawn, tree)

    def test_simple_tree_long_label(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time
        0   1           0
        1   1           0
        2   1           2
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       2       0
        0       1       2       1
        """
        )
        tree = (
            # fmt: off
            "ABCDEF\n"
            "┏┻┓   \n"
            "0 1   \n"
            # fmt: on
        )
        ts = tskit.load_text(nodes, edges, strict=False)
        t = next(ts.trees())
        drawn = t.draw_text(node_labels={0: "0", 1: "1", 2: "ABCDEF"}, order="tree")
        self.verify_text_rendering(drawn, tree)

        tree = (
            # fmt: off
            "0┓      \n"
            " ┣ABCDEF\n"
            "1┛      \n"
            # fmt: on
        )
        drawn = t.draw_text(
            node_labels={0: "0", 1: "1", 2: "ABCDEF"}, orientation="right", order="tree"
        )
        self.verify_text_rendering(drawn, tree)

        drawn = t.draw_text(
            node_labels={0: "ABCDEF", 1: "1", 2: "2"}, orientation="right", order="tree"
        )
        tree = (
            # fmt: off
            "ABCDEF┓ \n"
            "      ┣2\n"
            "1━━━━━┛ \n"
            # fmt: on
        )
        self.verify_text_rendering(drawn, tree)

        tree = (
            # fmt: off
            "      ┏0\n"
            "ABCDEF┫ \n"
            "      ┗1\n"
            # fmt: on
        )
        drawn = t.draw_text(
            node_labels={0: "0", 1: "1", 2: "ABCDEF"}, orientation="left", order="tree"
        )
        self.verify_text_rendering(drawn, tree)

    def test_four_leaves(self):
        nodes = io.StringIO(
            """\
        id      is_sample   population      individual      time    metadata
        0       1       0       -1      0.00000000000000
        1       1       0       -1      0.00000000000000
        2       1       0       -1      0.00000000000000
        3       1       0       -1      0.00000000000000
        4       0       0       -1      0.26676079696421
        5       0       0       -1      1.48826948286480
        6       0       0       -1      2.91835007758007
        """
        )
        edges = io.StringIO(
            """\
        left            right           parent  child
        0.00000000      1.00000000      4       0
        0.00000000      1.00000000      4       3
        0.00000000      1.00000000      5       2
        0.00000000      1.00000000      5       4
        0.00000000      1.00000000      6       1
        0.00000000      1.00000000      6       5
        """
        )
        tree = (
            "  6     \n"
            "┏━┻━┓   \n"
            "┃   5   \n"
            "┃ ┏━┻┓  \n"
            "┃ ┃  4  \n"
            "┃ ┃ ┏┻┓ \n"
            "1 2 0 3 \n"
        )
        ts = tskit.load_text(nodes, edges, strict=False)
        t = ts.first()
        drawn = t.draw(format="unicode", order="tree")
        self.verify_text_rendering(drawn, tree)
        self.verify_text_rendering(t.draw_text(order="tree"), tree)

        drawn = t.draw_text(orientation="bottom", order="tree")
        tree = (
            "1 2 0 3\n"
            "┃ ┃ ┗┳┛\n"
            "┃ ┃  4 \n"
            "┃ ┗━┳┛ \n"
            "┃   5  \n"
            "┗━┳━┛  \n"
            "  6    \n"
        )
        self.verify_text_rendering(drawn, tree)

        tree = (
            " ┏━━━━1\n"
            " ┃     \n"
            "6┫ ┏━━2\n"
            " ┃ ┃   \n"
            " ┗5┫ ┏0\n"
            "   ┗4┫  \n"
            "     ┗3\n"
        )
        self.verify_text_rendering(t.draw_text(orientation="left", order="tree"), tree)

        tree = (
            "2.92┊   6     ┊\n"
            "    ┊ ┏━┻━┓   ┊\n"
            "1.49┊ ┃   5   ┊\n"
            "    ┊ ┃ ┏━┻┓  ┊\n"
            "0.27┊ ┃ ┃  4  ┊\n"
            "    ┊ ┃ ┃ ┏┻┓ ┊\n"
            "0.00┊ 1 2 0 3 ┊\n"
            "    0         1\n"
        )
        self.verify_text_rendering(ts.draw_text(order="tree"), tree)

        tree = (
            "  6    \n"
            "+-+-+  \n"
            "|   5  \n"
            "| +-++ \n"
            "| |  4 \n"
            "| | +++\n"
            "1 2 0 3\n"
        )
        drawn = t.draw(format="ascii", order="tree")
        self.verify_text_rendering(drawn, tree)

        tree = (
            "  6     \n"
            "┏━┻━┓   \n"
            "┃xxxxxxxxxx\n"
            "┃ ┏━┻┓  \n"
            "┃ ┃  4  \n"
            "┃ ┃ ┏┻┓ \n"
            "1 2 0 3 \n"
        )
        labels = {u: str(u) for u in t.nodes()}
        labels[5] = "xxxxxxxxxx"
        drawn = t.draw_text(node_labels=labels, order="tree")
        self.verify_text_rendering(drawn, tree)

        tree = (
            " ┏━━━━━━━━━━━━━1\n"
            " ┃              \n"
            "6┫          ┏━━2\n"
            " ┃          ┃   \n"
            " ┗xxxxxxxxxx┫ ┏0\n"
            "            ┗4┫ \n"
            "              ┗3\n"
        )
        drawn = t.draw_text(node_labels=labels, orientation="left", order="tree")
        self.verify_text_rendering(drawn, tree)

        tree = (
            "2.92┊   6         ┊\n"
            "    ┊ ┏━┻━┓       ┊\n"
            "1.49┊ ┃xxxxxxxxxx ┊\n"
            "    ┊ ┃ ┏━┻┓      ┊\n"
            "0.27┊ ┃ ┃  4      ┊\n"
            "    ┊ ┃ ┃ ┏┻┓     ┊\n"
            "0.00┊ 1 2 0 3     ┊\n"
            "    0             1\n"
        )
        drawn = ts.draw_text(node_labels=labels, order="tree")
        self.verify_text_rendering(drawn, tree)

    def test_trident_tree(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time
        0   1           0
        1   1           0
        2   1           0
        3   1           2
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       3       0
        0       1       3       1
        0       1       3       2
        """
        )
        tree = (
            # fmt: off
            "  3  \n"
            "┏━╋━┓\n"
            "0 1 2\n"
            # fmt: on
        )
        ts = tskit.load_text(nodes, edges, strict=False)
        t = next(ts.trees())
        drawn = t.draw(format="unicode", order="tree")
        self.verify_text_rendering(drawn, tree)
        self.verify_text_rendering(t.draw_text(), tree)

        tree = (
            # fmt: off
            " ┏0\n"
            " ┃\n"
            "3╋1\n"
            " ┃\n"
            " ┗2\n"
            # fmt: on
        )
        drawn = t.draw_text(orientation="left")
        self.verify_text_rendering(drawn, tree)

        tree = (
            # fmt: off
            "0┓\n"
            " ┃\n"
            "1╋3\n"
            " ┃\n"
            "2┛\n"
            # fmt: on
        )
        drawn = t.draw_text(orientation="right")
        self.verify_text_rendering(drawn, tree)

    def test_pitchfork_tree(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time
        0   1           0
        1   1           0
        2   1           0
        3   1           0
        4   1           2
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       4       0
        0       1       4       1
        0       1       4       2
        0       1       4       3
        """
        )
        ts = tskit.load_text(nodes, edges, strict=False)
        t = next(ts.trees())
        tree = (
            # fmt: off
            "   4   \n"
            "┏━┳┻┳━┓\n"
            "0 1 2 3\n"
            # fmt: on
        )
        drawn = t.draw(format="unicode", order="tree")
        self.verify_text_rendering(drawn, tree)
        self.verify_text_rendering(t.draw_text(), tree)

        # No labels
        tree = (
            # fmt: off
            "   ┃   \n"
            "┏━┳┻┳━┓\n"
            "┃ ┃ ┃ ┃\n"
            # fmt: on
        )
        drawn = t.draw(format="unicode", node_labels={}, order="tree")
        self.verify_text_rendering(drawn, tree)
        self.verify_text_rendering(t.draw_text(node_labels={}), tree)
        # Some labels
        tree = (
            # fmt: off
            "   ┃   \n"
            "┏━┳┻┳━┓\n"
            "0 ┃ ┃ 3\n"
            # fmt: on
        )
        labels = {0: "0", 3: "3"}
        drawn = t.draw(format="unicode", node_labels=labels, order="tree")
        self.verify_text_rendering(drawn, tree)
        self.verify_text_rendering(t.draw_text(node_labels=labels), tree)

        tree = (
            # fmt: off
            " ┏0\n"
            " ┃\n"
            " ┣1\n"
            "4┫\n"
            " ┣2\n"
            " ┃\n"
            " ┗3\n"
            # fmt: on
        )
        drawn = t.draw_text(orientation="left")
        self.verify_text_rendering(drawn, tree)

        tree = (
            # fmt: off
            "0┓\n"
            " ┃\n"
            "1┫\n"
            " ┣4\n"
            "2┫\n"
            " ┃\n"
            "3┛\n"
            # fmt: on
        )
        drawn = t.draw_text(orientation="right")
        self.verify_text_rendering(drawn, tree)

    def test_stick_tree(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time
        0   1           0
        1   1           1
        2   1           2
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       1       0
        0       1       2       1
        """
        )
        tree = (
            # fmt: off
            "2\n"
            "┃\n"
            "1\n"
            "┃\n"
            "0\n"
            # fmt: on
        )
        ts = tskit.load_text(nodes, edges, strict=False)
        t = next(ts.trees())
        drawn = t.draw(format="unicode", order="tree")
        self.verify_text_rendering(drawn, tree)
        self.verify_text_rendering(t.draw_text(), tree)

        tree = (
            # fmt: off
            "0\n"
            "┃\n"
            "1\n"
            "┃\n"
            "2\n"
            # fmt: on
        )
        drawn = t.draw_text(orientation="bottom")
        self.verify_text_rendering(drawn, tree)

        tree = "2━1━0\n"
        drawn = t.draw_text(orientation="left")
        self.verify_text_rendering(drawn, tree)

        tree = "0━1━2\n"
        drawn = t.draw_text(orientation="right")
        self.verify_text_rendering(drawn, tree)

    def test_draw_forky_tree(self):
        tree = (
            "      14            \n"
            "  ┏━━━━┻━━━━┓       \n"
            "  ┃        13       \n"
            "  ┃   ┏━┳━┳━╋━┳━━┓  \n"
            "  ┃   ┃ ┃ ┃ ┃ ┃ 12  \n"
            "  ┃   ┃ ┃ ┃ ┃ ┃ ┏┻┓ \n"
            " 11   ┃ ┃ ┃ ┃ ┃ ┃ ┃ \n"
            "┏━┻┓  ┃ ┃ ┃ ┃ ┃ ┃ ┃ \n"
            "┃ 10  ┃ ┃ ┃ ┃ ┃ ┃ ┃ \n"
            "┃ ┏┻┓ ┃ ┃ ┃ ┃ ┃ ┃ ┃ \n"
            "8 0 3 2 4 5 6 9 1 7 \n"
        )

        nodes = io.StringIO(
            """\
            id      is_sample   population      individual      time    metadata
            0       1       0       -1      0.00000000000000
            1       1       0       -1      0.00000000000000
            2       1       0       -1      0.00000000000000
            3       1       0       -1      0.00000000000000
            4       1       0       -1      0.00000000000000
            5       1       0       -1      0.00000000000000
            6       1       0       -1      0.00000000000000
            7       1       0       -1      0.00000000000000
            8       1       0       -1      0.00000000000000
            9       1       0       -1      0.00000000000000
            10      0       0       -1      0.02398248117831
            11      0       0       -1      0.17378680550869
            12      0       0       -1      0.19950200178411
            13      0       0       -1      0.20000000000000
            14      0       0       -1      5.68339203134457
        """
        )
        edges = io.StringIO(
            """\
            left            right           parent  child
            0.00000000      1.00000000      10      0
            0.00000000      1.00000000      10      3
            0.00000000      1.00000000      11      8
            0.00000000      1.00000000      11      10
            0.00000000      1.00000000      12      1
            0.00000000      1.00000000      12      7
            0.00000000      1.00000000      13      2
            0.00000000      1.00000000      13      4
            0.00000000      1.00000000      13      5
            0.00000000      1.00000000      13      6
            0.00000000      1.00000000      13      9
            0.00000000      1.00000000      13      12
            0.00000000      1.00000000      14      11
            0.00000000      1.00000000      14      13
        """
        )
        ts = tskit.load_text(nodes, edges, strict=False)
        t = next(ts.trees())
        drawn = t.draw(format="unicode", order="tree")
        self.verify_text_rendering(drawn, tree)
        self.verify_text_rendering(t.draw_text(order="tree"), tree)

        tree = (
            "        14              \n"
            "  ┏━━━━━━┻━━━━━━┓       \n"
            "  ┃            13       \n"
            "  ┃        ┏━┳━┳┻┳━┳━━┓ \n"
            "  ┃        ┃ ┃ ┃ ┃ ┃ 12 \n"
            "  ┃        ┃ ┃ ┃ ┃ ┃ ┏┻┓\n"
            "x11xxxxxxx ┃ ┃ ┃ ┃ ┃ ┃ ┃\n"
            "┏━┻┓       ┃ ┃ ┃ ┃ ┃ ┃ ┃\n"
            "┃ 10       ┃ ┃ ┃ ┃ ┃ ┃ ┃\n"
            "┃ ┏┻┓      ┃ ┃ ┃ ┃ ┃ ┃ ┃\n"
            "8 0 3      2 4 5 6 9 1 7\n"
        )
        labels = {u: str(u) for u in t.nodes()}
        labels[11] = "x11xxxxxxx"
        self.verify_text_rendering(t.draw_text(node_labels=labels, order="tree"), tree)

        tree = (
            "      14           \n"
            "  ┏━━━━┻━━━━┓      \n"
            "  ┃        13      \n"
            "  ┃    ┏━━┳━╋━┳━┳━┓\n"
            "  ┃   12  ┃ ┃ ┃ ┃ ┃\n"
            "  ┃   ┏┻┓ ┃ ┃ ┃ ┃ ┃\n"
            " 11   ┃ ┃ ┃ ┃ ┃ ┃ ┃\n"
            " ┏┻━┓ ┃ ┃ ┃ ┃ ┃ ┃ ┃\n"
            "10  ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃\n"
            "┏┻┓ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃\n"
            "0 3 8 1 7 2 4 5 6 9\n"
        )
        self.verify_text_rendering(t.draw_text(order="minlex"), tree)

    def test_draw_multiroot_forky_tree(self):
        tree = (
            "           13      \n"
            "      ┏━┳━┳━╋━┳━━┓ \n"
            "      ┃ ┃ ┃ ┃ ┃ 12 \n"
            "      ┃ ┃ ┃ ┃ ┃ ┏┻┓\n"
            " 11   ┃ ┃ ┃ ┃ ┃ ┃ ┃\n"
            "┏━┻┓  ┃ ┃ ┃ ┃ ┃ ┃ ┃\n"
            "┃ 10  ┃ ┃ ┃ ┃ ┃ ┃ ┃\n"
            "┃ ┏┻┓ ┃ ┃ ┃ ┃ ┃ ┃ ┃\n"
            "8 0 3 2 4 5 6 9 1 7\n"
        )
        nodes = io.StringIO(
            """\
            id      is_sample   population      individual      time    metadata
            0       1       0       -1      0.00000000000000
            1       1       0       -1      0.00000000000000
            2       1       0       -1      0.00000000000000
            3       1       0       -1      0.00000000000000
            4       1       0       -1      0.00000000000000
            5       1       0       -1      0.00000000000000
            6       1       0       -1      0.00000000000000
            7       1       0       -1      0.00000000000000
            8       1       0       -1      0.00000000000000
            9       1       0       -1      0.00000000000000
            10      0       0       -1      0.02398248117831
            11      0       0       -1      0.17378680550869
            12      0       0       -1      0.19950200178411
            13      0       0       -1      0.20000000000000
            14      0       0       -1      5.68339203134457
        """
        )
        edges = io.StringIO(
            """\
            left            right           parent  child
            0.00000000      1.00000000      10      0
            0.00000000      1.00000000      10      3
            0.00000000      1.00000000      11      8
            0.00000000      1.00000000      11      10
            0.00000000      1.00000000      12      1
            0.00000000      1.00000000      12      7
            0.00000000      1.00000000      13      2
            0.00000000      1.00000000      13      4
            0.00000000      1.00000000      13      5
            0.00000000      1.00000000      13      6
            0.00000000      1.00000000      13      9
            0.00000000      1.00000000      13      12
        """
        )
        ts = tskit.load_text(nodes, edges, strict=False)
        t = next(ts.trees())
        drawn = t.draw(format="unicode", order="tree")
        self.verify_text_rendering(drawn, tree)
        self.verify_text_rendering(t.draw_text(order="tree"), tree)

        tree = (
            "           13      \n"
            "       ┏━━┳━╋━┳━┳━┓\n"
            "      12  ┃ ┃ ┃ ┃ ┃\n"
            "      ┏┻┓ ┃ ┃ ┃ ┃ ┃\n"
            " 11   ┃ ┃ ┃ ┃ ┃ ┃ ┃\n"
            " ┏┻━┓ ┃ ┃ ┃ ┃ ┃ ┃ ┃\n"
            "10  ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃\n"
            "┏┻┓ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃\n"
            "0 3 8 1 7 2 4 5 6 9\n"
        )
        drawn = t.draw(format="unicode")
        self.verify_text_rendering(drawn, tree)
        self.verify_text_rendering(t.draw_text(), tree)
        self.verify_text_rendering(t.draw_text(order="minlex"), tree)

    def test_simple_tree_sequence(self):
        ts = self.get_simple_ts()
        ts_drawing = (
            "9.08┊    9    ┊         ┊         ┊         ┊         ┊\n"
            "    ┊  ┏━┻━┓  ┊         ┊         ┊         ┊         ┊\n"
            "6.57┊  ┃   ┃  ┊         ┊         ┊         ┊    8    ┊\n"
            "    ┊  ┃   ┃  ┊         ┊         ┊         ┊  ┏━┻━┓  ┊\n"
            "5.31┊  ┃   ┃  ┊    7    ┊         ┊    7    ┊  ┃   ┃  ┊\n"
            "    ┊  ┃   ┃  ┊  ┏━┻━┓  ┊         ┊  ┏━┻━┓  ┊  ┃   ┃  ┊\n"
            "1.75┊  ┃   ┃  ┊  ┃   ┃  ┊    6    ┊  ┃   ┃  ┊  ┃   ┃  ┊\n"
            "    ┊  ┃   ┃  ┊  ┃   ┃  ┊  ┏━┻━┓  ┊  ┃   ┃  ┊  ┃   ┃  ┊\n"
            "1.11┊  ┃   5  ┊  ┃   5  ┊  ┃   5  ┊  ┃   5  ┊  ┃   5  ┊\n"
            "    ┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊\n"
            "0.11┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊\n"
            "    ┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊\n"
            "0.00┊ 0 1 2 3 ┊ 0 1 2 3 ┊ 0 1 2 3 ┊ 0 1 2 3 ┊ 0 1 2 3 ┊\n"
            "  0.00      0.06      0.79      0.91      0.91      1.00\n"
        )
        self.verify_text_rendering(ts.draw_text(), ts_drawing)

        ts_drawing = (
            "9.08|    9    |         |         |         |         |\n"
            "    |  +-+-+  |         |         |         |         |\n"
            "6.57|  |   |  |         |         |         |    8    |\n"
            "    |  |   |  |         |         |         |  +-+-+  |\n"
            "5.31|  |   |  |    7    |         |    7    |  |   |  |\n"
            "    |  |   |  |  +-+-+  |         |  +-+-+  |  |   |  |\n"
            "1.75|  |   |  |  |   |  |    6    |  |   |  |  |   |  |\n"
            "    |  |   |  |  |   |  |  +-+-+  |  |   |  |  |   |  |\n"
            "1.11|  |   5  |  |   5  |  |   5  |  |   5  |  |   5  |\n"
            "    |  |  +++ |  |  +++ |  |  +++ |  |  +++ |  |  +++ |\n"
            "0.11|  4  | | |  4  | | |  4  | | |  4  | | |  4  | | |\n"
            "    | +++ | | | +++ | | | +++ | | | +++ | | | +++ | | |\n"
            "0.00| 0 1 2 3 | 0 1 2 3 | 0 1 2 3 | 0 1 2 3 | 0 1 2 3 |\n"
            "  0.00      0.06      0.79      0.91      0.91      1.00\n"
        )
        self.verify_text_rendering(ts.draw_text(use_ascii=True), ts_drawing)

        ts_drawing = (
            "┊    9    ┊         ┊         ┊         ┊         ┊\n"
            "┊  ┏━┻━┓  ┊         ┊         ┊         ┊         ┊\n"
            "┊  ┃   ┃  ┊         ┊         ┊         ┊    8    ┊\n"
            "┊  ┃   ┃  ┊         ┊         ┊         ┊  ┏━┻━┓  ┊\n"
            "┊  ┃   ┃  ┊    7    ┊         ┊    7    ┊  ┃   ┃  ┊\n"
            "┊  ┃   ┃  ┊  ┏━┻━┓  ┊         ┊  ┏━┻━┓  ┊  ┃   ┃  ┊\n"
            "┊  ┃   ┃  ┊  ┃   ┃  ┊    6    ┊  ┃   ┃  ┊  ┃   ┃  ┊\n"
            "┊  ┃   ┃  ┊  ┃   ┃  ┊  ┏━┻━┓  ┊  ┃   ┃  ┊  ┃   ┃  ┊\n"
            "┊  ┃   5  ┊  ┃   5  ┊  ┃   5  ┊  ┃   5  ┊  ┃   5  ┊\n"
            "┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊\n"
            "┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊\n"
            "┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊\n"
            "┊ 0 1 2 3 ┊ 0 1 2 3 ┊ 0 1 2 3 ┊ 0 1 2 3 ┊ 0 1 2 3 ┊\n"
            "0.00    0.06      0.79      0.91      0.91      1.00\n"
        )
        self.verify_text_rendering(ts.draw_text(time_label_format=""), ts_drawing)

        ts_drawing = (
            "┊    9    ┊         ┊         ┊         ┊         ┊\n"
            "┊  ┏━┻━┓  ┊         ┊         ┊         ┊         ┊\n"
            "┊  ┃   ┃  ┊         ┊         ┊         ┊    8    ┊\n"
            "┊  ┃   ┃  ┊         ┊         ┊         ┊  ┏━┻━┓  ┊\n"
            "┊  ┃   ┃  ┊    7    ┊         ┊    7    ┊  ┃   ┃  ┊\n"
            "┊  ┃   ┃  ┊  ┏━┻━┓  ┊         ┊  ┏━┻━┓  ┊  ┃   ┃  ┊\n"
            "┊  ┃   ┃  ┊  ┃   ┃  ┊    6    ┊  ┃   ┃  ┊  ┃   ┃  ┊\n"
            "┊  ┃   ┃  ┊  ┃   ┃  ┊  ┏━┻━┓  ┊  ┃   ┃  ┊  ┃   ┃  ┊\n"
            "┊  ┃   5  ┊  ┃   5  ┊  ┃   5  ┊  ┃   5  ┊  ┃   5  ┊\n"
            "┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊  ┃  ┏┻┓ ┊\n"
            "┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊  4  ┃ ┃ ┊\n"
            "┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┊\n"
            "┊ 0 1 2 3 ┊ 0 1 2 3 ┊ 0 1 2 3 ┊ 0 1 2 3 ┊ 0 1 2 3 ┊\n"
            "┊         ┊         ┊         ┊         ┊         ┊\n"
        )
        self.verify_text_rendering(
            ts.draw_text(time_label_format="", position_label_format=""), ts_drawing
        )

    def test_tree_sequence_non_minlex(self):
        nodes = io.StringIO(
            """\
            id      is_sample       time    population      individual      metadata
            0       1       0.000000        0       -1
            1       1       0.000000        0       -1
            2       1       0.000000        0       -1
            3       1       0.000000        0       -1
            4       1       0.000000        0       -1
            5       0       1.174545        0       -1
            6       0       1.207717        0       -1
            7       0       1.276422        0       -1
            8       0       1.613390        0       -1
            9       0       2.700069        0       -1
        """
        )
        edges = io.StringIO(
            """\
            left    right   parent  child
            0.000000        1.000000        5       0
            0.000000        1.000000        5       1
            0.000000        0.209330        6       4
            0.000000        0.209330        6       5
            0.000000        1.000000        7       2
            0.209330        1.000000        7       5
            0.000000        0.209330        7       6
            0.209330        1.000000        8       3
            0.209330        1.000000        8       4
            0.000000        0.209330        9       3
            0.000000        1.000000        9       7
            0.209330        1.000000        9       8
        """
        )

        ts = tskit.load_text(nodes, edges, strict=False)

        drawn_minlex = (
            "2.70┊       9   ┊     9     ┊\n"
            "    ┊     ┏━┻━┓ ┊   ┏━┻━━┓  ┊\n"
            "1.61┊     ┃   ┃ ┊   ┃    8  ┊\n"
            "    ┊     ┃   ┃ ┊   ┃   ┏┻┓ ┊\n"
            "1.28┊     7   ┃ ┊   7   ┃ ┃ ┊\n"
            "    ┊   ┏━┻━┓ ┃ ┊  ┏┻━┓ ┃ ┃ ┊\n"
            "1.21┊   6   ┃ ┃ ┊  ┃  ┃ ┃ ┃ ┊\n"
            "    ┊  ┏┻━┓ ┃ ┃ ┊  ┃  ┃ ┃ ┃ ┊\n"
            "1.17┊  5  ┃ ┃ ┃ ┊  5  ┃ ┃ ┃ ┊\n"
            "    ┊ ┏┻┓ ┃ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┃ ┊\n"
            "0.00┊ 0 1 4 2 3 ┊ 0 1 2 3 4 ┊\n"
            "  0.00        0.21        1.00\n"
        )
        self.verify_text_rendering(ts.draw_text(order="minlex"), drawn_minlex)
        self.verify_text_rendering(ts.draw_text(), drawn_minlex)

        drawn_tree = (
            "2.70┊   9       ┊     9     ┊\n"
            "    ┊ ┏━┻━┓     ┊   ┏━┻━━┓  ┊\n"
            "1.61┊ ┃   ┃     ┊   ┃    8  ┊\n"
            "    ┊ ┃   ┃     ┊   ┃   ┏┻┓ ┊\n"
            "1.28┊ ┃   7     ┊   7   ┃ ┃ ┊\n"
            "    ┊ ┃ ┏━┻━┓   ┊ ┏━┻┓  ┃ ┃ ┊\n"
            "1.21┊ ┃ ┃   6   ┊ ┃  ┃  ┃ ┃ ┊\n"
            "    ┊ ┃ ┃ ┏━┻┓  ┊ ┃  ┃  ┃ ┃ ┊\n"
            "1.17┊ ┃ ┃ ┃  5  ┊ ┃  5  ┃ ┃ ┊\n"
            "    ┊ ┃ ┃ ┃ ┏┻┓ ┊ ┃ ┏┻┓ ┃ ┃ ┊\n"
            "0.00┊ 3 2 4 0 1 ┊ 2 0 1 3 4 ┊\n"
            "  0.00        0.21        1.00\n"
        )
        self.verify_text_rendering(ts.draw_text(order="tree"), drawn_tree)

    def test_max_time(self):
        ts = self.get_simple_ts()
        tree = (
            "   9   \n"
            " ┏━┻━┓ \n"
            " ┃   ┃ \n"
            " ┃   ┃ \n"
            " ┃   ┃ \n"
            " ┃   ┃ \n"
            " ┃   ┃ \n"
            " ┃   ┃ \n"
            " ┃   5 \n"
            " ┃  ┏┻┓\n"
            " 4  ┃ ┃\n"
            "┏┻┓ ┃ ┃\n"
            "0 1 2 3\n"
        )
        t = ts.first()
        self.verify_text_rendering(t.draw_text(max_time="ts"), tree)

        tree = (
            "   9   \n"
            " ┏━┻━┓ \n"
            " ┃   5 \n"
            " ┃  ┏┻┓\n"
            " 4  ┃ ┃\n"
            "┏┻┓ ┃ ┃\n"
            "0 1 2 3\n"
        )
        t = ts.first()
        self.verify_text_rendering(t.draw_text(max_time="tree"), tree)
        for bad_max_time in [1, "sdfr", ""]:
            with pytest.raises(ValueError):
                t.draw_text(max_time=bad_max_time)

    def test_no_repr_svg(self):
        tree = self.get_simple_ts().first()
        output = tree.draw(format="unicode")
        with pytest.raises(AttributeError, match="no attribute"):
            output._repr_svg_()


class TestDrawSvgBase(TestTreeDraw, xmlunittest.XmlTestMixin):
    """
    Base class for testing the SVG tree drawing method
    """

    def verify_basic_svg(self, svg, width=200, height=200, num_trees=1, has_root=True):
        prefix = "{http://www.w3.org/2000/svg}"
        root = xml.etree.ElementTree.fromstring(svg)
        assert root.tag == prefix + "svg"
        assert width * num_trees == int(root.attrib["width"])
        assert height == int(root.attrib["height"])

        # Verify the class structure of the svg
        root_group = root.find(prefix + "g")
        assert "class" in root_group.attrib
        assert re.search(r"\b(tree|tree-sequence)\b", root_group.attrib["class"])
        first_plotbox = None
        if "tree-sequence" in root_group.attrib["class"]:
            trees = None
            for g in root_group.findall(prefix + "g"):
                if "trees" in g.attrib.get("class", ""):
                    trees = g
                    break
            assert trees is not None  # Must have found a trees group
            first_tree = trees.find(prefix + "g")
            assert "class" in first_tree.attrib
            assert re.search(r"\btree\b", first_tree.attrib["class"])
            for g in first_tree.findall(prefix + "g"):
                if "class" in g.attrib and re.search(r"\bplotbox\b", g.attrib["class"]):
                    first_plotbox = g
        else:
            for g in root_group.findall(prefix + "g"):
                if "class" in g.attrib and re.search(r"\bplotbox\b", g.attrib["class"]):
                    first_plotbox = g
        assert first_plotbox is not None
        # Check that we have edges, symbols, and labels groups
        groups = first_plotbox.findall(prefix + "g")
        assert len(groups) > 0
        for group in groups:
            assert "class" in group.attrib
            cls = group.attrib["class"]
            # if a subtree plot, the top of the displayed topology is not a local root
            if has_root:
                assert re.search(r"\broot\b", cls)
            else:
                assert not re.search(r"\broot\b", cls)


class TestDrawSvg(TestDrawSvgBase):
    """
    Simple testing for the draw_svg method
    """

    def test_repr_svg(self):
        ts = self.get_simple_ts()
        svg = ts.draw_svg()
        assert str(svg) == svg._repr_svg_()
        svg = ts.first().draw_svg()
        assert str(svg) == svg._repr_svg_()
        svg = ts.first().draw(format="svg")
        assert str(svg) == svg._repr_svg_()

    def test_draw_to_file(self, tmp_path):
        # NB: to view output files for testing changes to drawing code, it is possible
        # to save to a fixed directory using e.g. `pytest --basetemp=/tmp/svgtest ...`
        t = self.get_binary_tree()
        filename = tmp_path / "tree-draw.svg"
        svg = t.draw(path=filename)
        assert os.path.getsize(filename) > 0
        with open(filename) as tmp:
            other_svg = tmp.read()
        assert svg == other_svg

        filename = tmp_path / "tree-draw_svg.svg"
        svg = t.draw_svg(path=filename)
        assert os.path.getsize(filename) > 0
        with open(filename) as tmp:
            other_svg = tmp.read()
        self.verify_basic_svg(svg)
        self.verify_basic_svg(other_svg)

        filename = tmp_path / "ts-draw_svg.svg"
        ts = self.get_simple_ts()
        svg = ts.draw_svg(path=filename)
        assert os.path.getsize(filename) > 0
        with open(filename) as tmp:
            other_svg = tmp.read()
        self.verify_basic_svg(svg, num_trees=ts.num_trees)
        self.verify_basic_svg(other_svg, num_trees=ts.num_trees)

    def test_nonimplemented_base_class(self):
        ts = self.get_simple_ts()
        plot = drawing.SvgAxisPlot(
            ts, (100, 100), {}, "", "dummy-class", None, True, True
        )
        plot.set_spacing()
        with pytest.raises(NotImplementedError):
            plot.draw_x_axis(tick_positions=ts.breakpoints(as_array=True))

    def test_bad_tick_spacing(self):
        # Integer y_ticks to give auto-generated tick locs is not currently implemented
        t = self.get_binary_tree()
        with pytest.raises(TypeError):
            t.draw_svg(y_axis=True, y_ticks=6)
        ts = self.get_simple_ts()
        with pytest.raises(TypeError):
            ts.draw_svg(y_axis=True, y_ticks=6)

    def test_no_mixed_yscales(self):
        ts = self.get_simple_ts()
        with pytest.raises(ValueError, match="vary in timescale"):
            ts.draw_svg(y_axis=True, max_time="tree")

    def test_draw_defaults(self):
        t = self.get_binary_tree()
        svg = t.draw()
        self.verify_basic_svg(svg)
        svg = t.draw_svg()
        self.verify_basic_svg(svg)

    @pytest.mark.parametrize("y_axis", (True, False))
    @pytest.mark.parametrize("y_label", (True, False))
    @pytest.mark.parametrize(
        "time_scale",
        (
            "rank",
            "time",
        ),
    )
    @pytest.mark.parametrize("y_ticks", ([], [0, 1], None))
    @pytest.mark.parametrize("y_gridlines", (True, False))
    def test_draw_svg_y_axis_parameter_combos(
        self, y_axis, y_label, time_scale, y_ticks, y_gridlines
    ):
        t = self.get_binary_tree()
        svg = t.draw_svg(
            y_axis=y_axis,
            y_label=y_label,
            y_ticks=y_ticks,
            y_gridlines=y_gridlines,
            time_scale=time_scale,
        )
        self.verify_basic_svg(svg)
        ts = self.get_simple_ts()
        svg = ts.draw_svg(
            y_axis=y_axis,
            y_label=y_label,
            y_ticks=y_ticks,
            y_gridlines=y_gridlines,
            time_scale=time_scale,
        )
        self.verify_basic_svg(svg, width=200 * ts.num_trees)

    def test_draw_multiroot(self):
        t = self.get_multiroot_tree()
        svg = t.draw()
        self.verify_basic_svg(svg)
        svg = t.draw_svg()
        self.verify_basic_svg(svg)

    def test_draw_mutations_over_roots(self):
        t = self.get_mutations_over_roots_tree()
        with pytest.warns(UserWarning, match="nodes which are not present"):
            svg = t.draw()
            self.verify_basic_svg(svg)
        with pytest.warns(UserWarning, match="nodes which are not present"):
            svg = t.draw_svg()
            self.verify_basic_svg(svg)

    def test_draw_unary(self):
        t = self.get_unary_node_tree()
        svg = t.draw()
        self.verify_basic_svg(svg)
        svg = t.draw_svg()
        self.verify_basic_svg(svg)

    def test_draw_empty(self):
        t = self.get_empty_tree()
        with pytest.raises(ValueError):
            t.draw()
        with pytest.raises(ValueError):
            t.draw_svg()

    def test_draw_zero_roots(self):
        t = self.get_zero_roots_tree()
        with pytest.raises(ValueError):
            t.draw()
        with pytest.raises(ValueError):
            t.draw_svg()

    def test_draw_zero_edge(self):
        t = self.get_zero_edge_tree()
        svg = t.draw()
        self.verify_basic_svg(svg)
        svg = t.draw_svg()
        self.verify_basic_svg(svg)

    def test_mutations_present(self):
        t = self.get_binary_tree()
        assert t.tree_sequence.num_mutations > 0
        svg = t.draw()
        self.verify_basic_svg(svg)
        assert svg.count('class="mut') == t.tree_sequence.num_mutations
        svg = t.draw_svg()
        self.verify_basic_svg(svg)
        assert svg.count('class="mut') == t.tree_sequence.num_mutations
        svg = t.tree_sequence.draw_svg()
        self.verify_basic_svg(svg)
        assert 'class="site' in svg
        assert svg.count('class="site') == t.tree_sequence.num_sites

    def test_sites_omitted(self):
        t = self.get_binary_tree()
        assert t.tree_sequence.num_mutations > 0
        svg = t.draw(omit_sites=True)
        self.verify_basic_svg(svg)
        assert svg.count('class="mut') == 0
        svg = t.draw_svg(omit_sites=True)
        self.verify_basic_svg(svg)
        assert svg.count('class="mut') == 0
        svg = t.tree_sequence.draw_svg(omit_sites=True)
        self.verify_basic_svg(svg)
        assert svg.count('class="mut') == 0
        assert svg.count('class="site') == 0

    def test_width_height(self):
        t = self.get_binary_tree()
        w = 123
        h = 456
        svg = t.draw(width=w, height=h)
        self.verify_basic_svg(svg, w, h)
        svg = t.draw_svg(size=(w, h))
        self.verify_basic_svg(svg, w, h)

    def test_node_labels(self):
        t = self.get_binary_tree()
        labels = {u: "XXX" for u in t.nodes()}
        svg = t.draw(format="svg", node_labels=labels)
        self.verify_basic_svg(svg)
        assert svg.count("XXX") == t.tree_sequence.num_nodes
        svg = t.draw_svg(node_label_attrs={u: {"text": labels[u]} for u in t.nodes()})
        self.verify_basic_svg(svg)
        assert svg.count("XXX") == t.tree_sequence.num_nodes

    def test_one_node_label(self):
        t = self.get_binary_tree()
        labels = {0: "XXX"}
        svg = t.draw(format="svg", node_labels=labels)
        self.verify_basic_svg(svg)
        assert svg.count("XXX") == 1
        svg = t.draw_svg(node_label_attrs={0: {"text": "XXX"}})
        self.verify_basic_svg(svg)
        assert svg.count("XXX") == 1

    def test_no_node_labels(self):
        t = self.get_binary_tree()
        labels = {}
        svg = t.draw(format="svg", node_labels=labels)
        self.verify_basic_svg(svg)
        # Can't really test for much here if we don't understand the SVG

    def test_one_node_colour(self):
        t = self.get_binary_tree()
        colour = "rgb(0, 1, 2)"
        colours = {0: colour}
        svg = t.draw(format="svg", node_colours=colours)
        self.verify_basic_svg(svg)
        assert svg.count(f"fill:{colour}") == 1
        svg = t.draw_svg(node_attrs={0: {"fill": colour}})
        self.verify_basic_svg(svg)
        assert svg.count(f'fill="{colour}"') == 1

    def test_all_nodes_colour(self):
        t = self.get_binary_tree()
        colours = {u: f"rgb({u}, {u}, {u})" for u in t.nodes()}
        svg = t.draw(format="svg", node_colours=colours)
        self.verify_basic_svg(svg)
        for colour in colours.values():
            assert svg.count(f"fill:{colour}") == 1

        svg = t.draw_svg(node_attrs={u: {"fill": colours[u]} for u in t.nodes()})
        self.verify_basic_svg(svg)
        assert svg.count(f'fill="{colour}"') == 1
        for colour in colours.values():
            assert svg.count(f'fill="{colour}"') == 1

    def test_unplotted_node(self):
        t = self.get_binary_tree()
        colour = None
        colours = {0: colour}
        svg = t.draw(format="svg", node_colours=colours)
        svg_no_css = svg[svg.find("</style>") :]
        assert svg_no_css.count("opacity:0") == 1

    def test_one_edge_colour(self):
        t = self.get_binary_tree()
        colour = "rgb(0, 1, 2)"
        colours = {0: colour}
        svg = t.draw(format="svg", edge_colours=colours)
        self.verify_basic_svg(svg)
        assert svg.count(f"stroke:{colour}") > 0
        svg = t.draw_svg(edge_attrs={0: {"stroke": colour}})
        self.verify_basic_svg(svg)
        assert svg.count(f'stroke="{colour}"') == 1

    def test_one_mutation_label_colour(self):
        t = self.get_binary_tree()
        colour = "rgb(0, 1, 2)"
        svg = t.draw_svg(mutation_label_attrs={0: {"stroke": colour}})
        self.verify_basic_svg(svg)
        assert svg.count(f'stroke="{colour}"') == 1

    def test_bad_time_scale(self):
        t = self.get_binary_tree()
        for bad_scale in ["te", "asdf", "", [], b"23"]:
            with pytest.raises(ValueError):
                t.draw_svg(time_scale=bad_scale)
            with pytest.raises(ValueError):
                with pytest.warns(FutureWarning):
                    t.draw_svg(tree_height_scale=bad_scale)

    def test_bad_max_time(self):
        t = self.get_binary_tree()
        for bad_height in ["te", "asdf", "", [], b"23"]:
            with pytest.raises(ValueError):
                t.draw_svg(max_time=bad_height)
            with pytest.raises(ValueError):
                with pytest.warns(FutureWarning):
                    t.draw_svg(max_tree_height=bad_height)

    def test_bad_min_time(self):
        t = self.get_binary_tree()
        for bad_min in ["te", "asdf", "", [], b"23"]:
            with pytest.raises(ValueError):
                t.draw_svg(min_time=bad_min)
            with pytest.raises(ValueError):
                with pytest.warns(FutureWarning):
                    t.draw_svg(max_tree_height=bad_min)

    def test_bad_neg_log_time(self):
        t = self.get_ts_varying_min_times().at_index(1)
        assert min(t.time(u) for u in t.nodes()) < 0
        with pytest.raises(ValueError, match="negative times"):
            with np.errstate(invalid="ignore"):
                t.draw_svg(t.draw_svg(time_scale="log_time"))

    def test_time_scale_time_and_max_time(self):
        ts = msprime.simulate(5, recombination_rate=2, random_seed=2)
        t = ts.first()
        # The default should be the same as tree.
        svg1 = t.draw_svg(max_time="tree")
        self.verify_basic_svg(svg1)
        svg2 = t.draw_svg()
        assert svg1 == svg2
        svg3 = t.draw_svg(max_time="ts")
        assert svg1 != svg3
        svg4 = t.draw_svg(max_time=max(ts.tables.nodes.time))
        assert svg3 == svg4
        with pytest.warns(FutureWarning):
            svg5 = t.draw_svg(max_tree_height="tree")
        assert svg5 == svg1
        svg6 = t.draw_svg(max_time="tree", max_tree_height="i should be ignored")
        assert svg6 == svg1

    def test_time_scale_rank_and_max_time(self):
        # Make sure the rank height scale and max_time interact properly.
        ts = msprime.simulate(5, recombination_rate=2, random_seed=2)
        t = ts.first()
        # The default should be the same as tree.
        svg1 = t.draw_svg(max_time="tree", time_scale="rank", y_axis=True)
        self.verify_basic_svg(svg1)
        svg2 = t.draw_svg(time_scale="rank", y_axis=True)
        assert svg1 == svg2
        svg3 = t.draw_svg(max_time="ts", time_scale="rank", y_axis=True)
        assert svg1 != svg3
        self.verify_basic_svg(svg3)
        # Numeric max time not supported for rank scale.
        with pytest.raises(ValueError):
            t.draw_svg(max_time=2, time_scale="rank", y_axis=True)

    def test_min_tree_time(self):
        ts = self.get_ts_varying_min_times()
        t = ts.first()
        # The default should be the same as tree.
        svg1 = t.draw_svg(min_time="tree", y_axis=True)
        self.verify_basic_svg(svg1)
        svg2 = t.draw_svg(y_axis=True)
        assert svg1 == svg2
        svg3 = t.draw_svg(min_time="ts", y_axis=True)
        assert svg1 != svg3
        svg4 = t.draw_svg(min_time=min(ts.tables.nodes.time), y_axis=True)
        assert svg3 == svg4

    def test_min_ts_time(self):
        ts = self.get_ts_varying_min_times()
        svg1 = ts.draw_svg(y_axis=True)
        self.verify_basic_svg(svg1, width=200 * ts.num_trees)
        svg2 = ts.draw_svg(min_time="ts", y_axis=True)
        assert svg1 == svg2
        with pytest.raises(ValueError, match="vary in timescale"):
            ts.draw_svg(min_time="tree", y_axis=True)
        svg3 = ts.draw_svg(min_time=min(ts.tables.nodes.time), y_axis=True)
        assert svg2 == svg3

    #
    # TODO: update the tests below here to check the new SVG based interface.
    #
    def test_all_edges_colour(self):
        t = self.get_binary_tree()
        colours = {u: "rgb({u},255,{u})".format(u=u) for u in t.nodes() if u != t.root}
        svg = t.draw(format="svg", edge_colours=colours)
        self.verify_basic_svg(svg)
        for colour in colours.values():
            assert svg.count(f"stroke:{colour}") > 0

    def test_unplotted_edge(self):
        t = self.get_binary_tree()
        colour = None
        colours = {0: colour}
        svg = t.draw(format="svg", edge_colours=colours)
        self.verify_basic_svg(svg)
        svg_no_css = svg[svg.find("</style>") :]
        assert svg_no_css.count("opacity:0") == 1

    def test_mutations_unknown_time(self):
        ts = self.get_simple_ts(use_mutation_times=True)
        svg = ts.draw_svg()
        self.verify_basic_svg(svg, width=200 * ts.num_trees)
        assert "unknown_time" not in svg
        ts = self.get_simple_ts(use_mutation_times=False)
        svg = ts.draw_svg()
        self.verify_basic_svg(svg, width=200 * ts.num_trees)
        assert svg.count("unknown_time") == ts.num_mutations

    def test_mutation_labels(self):
        t = self.get_binary_tree()
        labels = {u.id: "XXX" for u in t.mutations()}
        svg = t.draw(format="svg", mutation_labels=labels)
        self.verify_basic_svg(svg)
        assert svg.count("XXX") == t.num_mutations

    def test_one_mutation_label(self):
        t = self.get_binary_tree()
        labels = {0: "XXX"}
        svg = t.draw(format="svg", mutation_labels=labels)
        self.verify_basic_svg(svg)
        assert svg.count("XXX") == 1

    def test_no_mutation_labels(self):
        t = self.get_binary_tree()
        labels = {}
        svg = t.draw(format="svg", mutation_labels=labels)
        self.verify_basic_svg(svg)
        # Can't really test for much here if we don't understand the SVG

    def test_one_mutation_colour(self):
        t = self.get_binary_tree()
        colour = "rgb(0, 1, 2)"
        colours = {0: colour}
        svg = t.draw(format="svg", mutation_colours=colours)
        self.verify_basic_svg(svg)
        assert svg.count(f"fill:{colour}") == 1

    def test_all_mutations_colour(self):
        t = self.get_binary_tree()
        colours = {
            mut.id: f"rgb({mut.id}, {mut.id}, {mut.id})" for mut in t.mutations()
        }
        svg = t.draw(format="svg", mutation_colours=colours)
        self.verify_basic_svg(svg)
        for colour in colours.values():
            assert svg.count(f"fill:{colour}") == 1

    def test_unplotted_mutation(self):
        t = self.get_binary_tree()
        colour = None
        colours = {0: colour}
        svg = t.draw(format="svg", mutation_colours=colours)
        self.verify_basic_svg(svg)
        svg_no_css = svg[svg.find("</style>") :]
        assert svg_no_css.count("fill-opacity:0") == 1

    @pytest.mark.parametrize("all_muts", [False, True])
    @pytest.mark.parametrize("x_axis", [False, True])
    def test_extra_mutations(self, all_muts, x_axis):
        # The simple_ts has 2 mutations on an edge which spans the whole ts
        # One mut is within tree 1, the other within tree 3
        ts = self.get_simple_ts()
        extra_mut_copies = 0
        if all_muts:
            extra_mut_copies = 2 if x_axis else 1
        extra_right = ts.at_index(1)
        svg = extra_right.draw_svg(all_edge_mutations=all_muts, x_axis=x_axis)
        self.verify_basic_svg(svg)
        svg_no_css = svg[svg.find("</style>") :]
        assert svg_no_css.count("extra") == 1 * extra_mut_copies

        extra_right_and_left = ts.at_index(2)
        svg = extra_right_and_left.draw_svg(all_edge_mutations=all_muts, x_axis=x_axis)
        self.verify_basic_svg(svg)
        svg_no_css = svg[svg.find("</style>") :]
        assert svg_no_css.count("extra") == 2 * extra_mut_copies

        extra_left = ts.at_index(3)
        svg = extra_left.draw_svg(all_edge_mutations=all_muts, x_axis=x_axis)
        self.verify_basic_svg(svg)
        svg_no_css = svg[svg.find("</style>") :]
        assert svg_no_css.count("extra") == 1 * extra_mut_copies

    def test_max_time(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time
        0   1           0
        1   1           0
        2   1           0
        3   0           1
        4   0           2
        5   0           3
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       5       2
        0       1       5       3
        1       2       4       2
        1       2       4       3
        0       2       3       0
        0       2       3       1
        """
        )
        ts = tskit.load_text(nodes, edges, strict=False)

        svg1 = ts.at_index(0).draw()
        svg2 = ts.at_index(1).draw()
        # if not scaled to ts, the edge above node 0 is of a different length in both
        # trees, because the root is at a different height. We expect a group like
        # <path class="edge" d="M 0 0 V -46 H 22.5" /><text>0</text>
        str_pos = svg1.find(">0<")
        snippet1 = svg1[svg1.rfind("edge", 0, str_pos) : str_pos]
        str_pos = svg2.find(">0<")
        snippet2 = svg2[svg2.rfind("edge", 0, str_pos) : str_pos]
        assert snippet1 != snippet2

        svg1 = ts.at_index(0).draw(max_time="ts")
        svg2 = ts.at_index(1).draw(max_time="ts")
        with pytest.warns(FutureWarning):
            svg3 = ts.at_index(1).draw(max_tree_height="ts")
        assert svg3 == svg2
        # when scaled, node 3 should be at the *same* height in both trees, so the edge
        # definition should be the same
        self.verify_basic_svg(svg1)
        self.verify_basic_svg(svg2)
        str_pos = svg1.find(">0<")
        snippet1 = svg1[svg1.rfind("edge", 0, str_pos) : str_pos]
        str_pos = svg2.find(">0<")
        snippet2 = svg2[svg2.rfind("edge", 0, str_pos) : str_pos]
        assert snippet1 == snippet2

    def test_min_time(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time
        0   0           -1.11
        1   1           2.22
        2   1           2.22
        3   0           3.33
        4   0           4.44
        5   0           5.55
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       5       2
        0       1       5       3
        1       2       4       2
        1       2       4       3
        0       1       3       0
        0       2       3       1
        """
        )
        ts = tskit.load_text(nodes, edges, strict=False)
        svg1a = ts.at_index(0).draw_svg(y_axis=True)
        svg1b = ts.at_index(0).draw_svg(y_axis=True, min_time="ts")
        svg2a = ts.at_index(1).draw_svg(y_axis=True)
        svg2b = ts.at_index(1).draw_svg(y_axis=True, min_time="ts")
        # axis should start at -1.11
        assert svg1a == svg1b
        assert ">-1.11<" in svg1a
        # 2nd tree should be different depending on whether min_time is "tree" or "ts"
        assert svg2a != svg2b
        assert ">-1.11<" not in svg2a
        assert ">-1.11<" not in svg2b

    def test_draw_sized_tree(self):
        tree = self.get_binary_tree()
        svg = tree.draw_svg(size=(600, 400))
        self.verify_basic_svg(svg, width=600, height=400)

    def test_canvas_size_tree(self):
        tree = self.get_binary_tree()
        svg1 = tree.draw_svg(size=(200, 200))
        svg2 = tree.draw_svg(size=(200, 200), canvas_size=(700, 500))
        self.verify_basic_svg(svg1, width=200, height=200)
        self.verify_basic_svg(svg2, width=700, height=500)
        # height and width are specified in the starting <svg> tag
        assert svg1.startswith("<svg")
        assert svg2.startswith("<svg")
        # after the close of the tag, the two strings should be the same
        assert svg1[svg1.find(">") :] == svg2[svg2.find(">") :]

    def test_draw_bad_sized_treebox(self):
        tree = self.get_binary_tree()
        with pytest.raises(ValueError, match="too small to fit"):
            # Too small for plotbox
            tree.draw_svg(size=(20, 20))

    def test_draw_bad_sized_tree(self):
        tree = self.get_binary_tree()
        with pytest.raises(ValueError, match="too small to allow space"):
            # Too small for standard-sized labels on tree
            tree.draw_svg(size=(50, 50))

    def test_draw_simple_ts(self):
        ts = msprime.simulate(5, recombination_rate=1, random_seed=1)
        svg = ts.draw_svg()
        self.verify_basic_svg(svg, width=200 * ts.num_trees)

    def test_draw_integer_breaks_ts(self):
        ts = msprime.sim_ancestry(
            5, sequence_length=10, recombination_rate=0.05, random_seed=1
        )
        assert ts.num_trees > 2
        svg = ts.draw_svg()
        self.verify_basic_svg(svg, width=200 * ts.num_trees)
        axis_pos = svg.find('class="x-axis"')
        for b in ts.breakpoints():
            assert b == round(b)
            assert svg.find(f">{b:.0f}<", axis_pos) != -1

    def test_draw_integer_times_ts(self):
        ts = msprime.sim_ancestry(
            5, population_size=5, sequence_length=10, model="dtwf", random_seed=1
        )
        svg = ts.draw_svg(y_axis=True)
        self.verify_basic_svg(svg, width=200 * ts.num_trees)
        axis_pos = svg.find('class="y-axis"')
        for t in ts.tables.nodes.time:
            assert t == round(t)
            assert svg.find(f">{t:.0f}<", axis_pos) != -1

    def test_draw_integer_times_tree(self):
        ts = msprime.sim_ancestry(
            5, population_size=5, sequence_length=10, model="dtwf", random_seed=1
        )
        svg = ts.first().draw_svg(y_axis=True)
        self.verify_basic_svg(svg, width=200 * ts.num_trees)
        axis_pos = svg.find('class="y-axis"')
        for t in ts.tables.nodes.time:
            assert t == round(t)
            assert svg.find(f">{t:.0f}<", axis_pos) != -1

    def test_draw_even_height_ts(self):
        ts = msprime.simulate(5, recombination_rate=1, random_seed=1)
        svg = ts.draw_svg(max_time="tree")
        self.verify_basic_svg(svg, width=200 * ts.num_trees)
        with pytest.warns(FutureWarning):
            svg = ts.draw_svg(max_tree_height="tree")
        self.verify_basic_svg(svg, width=200 * ts.num_trees)

    def test_draw_sized_ts(self):
        ts = msprime.simulate(5, recombination_rate=1, random_seed=1)
        svg = ts.draw_svg(size=(600, 400))
        self.verify_basic_svg(svg, width=600, height=400)

    def test_canvas_size_ts(self):
        ts = msprime.simulate(5, recombination_rate=1, random_seed=1)
        svg1 = ts.draw_svg(size=(600, 400))
        svg2 = ts.draw_svg(size=(600, 400), canvas_size=(1000, 500))
        self.verify_basic_svg(svg1, width=600, height=400)
        self.verify_basic_svg(svg2, width=1000, height=500)
        # height and width are specified in the starting <svg> tag
        assert svg1.startswith("<svg")
        assert svg2.startswith("<svg")
        # after the close of the tag, the two strings should be the same
        assert svg1[svg1.find(">") :] == svg2[svg2.find(">") :]

    def test_time_scale(self):
        ts = msprime.simulate(4, random_seed=2)
        svg = ts.draw_svg(time_scale="time")
        self.verify_basic_svg(svg)
        svg = ts.draw_svg(time_scale="log_time")
        self.verify_basic_svg(svg)
        with pytest.warns(FutureWarning):
            svg2 = ts.draw_svg(tree_height_scale="log_time")
        assert svg2 == svg
        svg = ts.draw_svg(time_scale="rank")
        self.verify_basic_svg(svg)
        svg3 = ts.draw_svg(time_scale="rank", tree_height_scale="ignore me please")
        assert svg3 == svg
        for bad_scale in [0, "", "NOT A SCALE"]:
            with pytest.raises(ValueError):
                ts.draw_svg(time_scale=bad_scale)
            with pytest.raises(ValueError):
                with pytest.warns(FutureWarning):
                    ts.draw_svg(tree_height_scale=bad_scale)

    def test_x_scale(self):
        ts = msprime.simulate(4, random_seed=2)
        svg = ts.draw_svg(x_scale="physical")
        self.verify_basic_svg(svg)
        svg = ts.draw_svg(x_scale="treewise")
        self.verify_basic_svg(svg)

    def test_bad_x_scale(self):
        ts = msprime.simulate(4, random_seed=2)
        for bad_x_scale in ["te", "asdf", "", [], b"23"]:
            with pytest.raises(ValueError):
                ts.draw_svg(x_scale=bad_x_scale)

    def test_x_axis(self):
        tree = msprime.simulate(4, random_seed=2).first()
        svg = tree.draw_svg(x_axis=True)
        svg_no_css = svg[svg.find("</style>") :]
        assert "Genome position" in svg_no_css
        assert svg_no_css.count("axes") == 1
        assert svg_no_css.count("x-axis") == 1
        assert svg_no_css.count("y-axis") == 0

    def test_y_axis(self):
        tree = self.get_simple_ts().first()
        for hscale, label in [
            (None, "Time"),
            ("time", "Time"),
            ("log_time", "Time"),
            ("rank", "Node time"),
        ]:
            svg = tree.draw_svg(y_axis=True, time_scale=hscale)
            if hscale is not None:
                with pytest.warns(FutureWarning):
                    svg2 = tree.draw_svg(y_axis=True, tree_height_scale=hscale)
                assert svg2 == svg
                svg3 = tree.draw_svg(
                    y_axis=True, time_scale=hscale, tree_height_scale="ignore me please"
                )
                assert svg3 == svg
            svg_no_css = svg[svg.find("</style>") :]
            assert label in svg_no_css
            assert svg_no_css.count("axes") == 1
            assert svg_no_css.count("x-axis") == 0
            assert svg_no_css.count("y-axis") == 1
            assert svg_no_css.count("ticks") == 1
            assert svg_no_css.count('class="tick"') == len(
                {tree.time(u) for u in tree.nodes()}
            )

    def test_y_axis_noticks(self):
        tree = msprime.simulate(4, random_seed=2).first()
        svg = tree.draw_svg(y_axis=True, y_label="Time", y_ticks=[])
        svg_no_css = svg[svg.find("</style>") :]
        assert svg_no_css.count("axes") == 1
        assert svg_no_css.count("x-axis") == 0
        assert svg_no_css.count("y-axis") == 1
        assert svg_no_css.count('"tick"') == 0

    def test_y_axis_tick_warning(sefl, caplog):
        tree = msprime.simulate(4, random_seed=2).first()
        upper = int(tree.time(tree.root))
        with caplog.at_level(logging.WARNING):
            tree.draw_svg(
                y_axis=True,
                y_label="Time",
                y_ticks={upper + 100: "above", upper / 3: "inside"},
            )
            assert (
                f"Ticks {{{upper+100}: 'above'}} lie outside the plotted axis"
                in caplog.text
            )
        with caplog.at_level(logging.WARNING):
            tree.draw_svg(
                y_axis=True, y_label="Time", y_ticks={upper / 2: "inside", -1: "below"}
            )
            assert "Ticks {-1: 'below'} lie outside the plotted axis" in caplog.text

    def test_symbol_size(self):
        tree = msprime.simulate(4, random_seed=2, mutation_rate=8).first()
        sz = 24
        svg = tree.draw_svg(symbol_size=sz)
        svg_no_css = svg[svg.find("</style>") :]
        num_mutations = len([_ for _ in tree.mutations()])
        num_nodes = len([_ for _ in tree.nodes()])
        # Squares have 'height="sz" width="sz"'
        assert svg_no_css.count(f'"{sz}"') == tree.num_samples() * 2
        # Circles define a radius like 'r="sz/2"'
        assert svg_no_css.count(f'r="{sz/2:g}"') == num_nodes - tree.num_samples()
        # Mutations draw a line on the cross using 'l sz,sz'
        assert svg_no_css.count(f"l {sz},{sz} ") == num_mutations

    def test_no_edges_invalid(self):
        full_ts = msprime.simulate(10, random_seed=2)
        tables = full_ts.dump_tables()
        tables.edges.clear()
        ts = tables.tree_sequence()
        with pytest.raises(ValueError, match="To plot an empty tree sequence"):
            ts.draw_svg()
        with pytest.raises(ValueError, match="To plot an empty tree sequence"):
            ts.draw_svg(x_lim=[None, 1])
        with pytest.raises(ValueError, match="To plot an empty tree sequence"):
            ts.draw_svg(x_lim=[0, None])

    def test_no_edges_show_empty(self):
        # Should be possible to print empty trees if xlim=[0, seq_len]
        full_ts = msprime.simulate(10, random_seed=2)
        tables = full_ts.dump_tables()
        tables.edges.clear()
        ts = tables.tree_sequence()
        for time_scale in ("time", "log_time", "rank"):
            # SVG should just be a row of 10 sample nodes
            svg = ts.draw_svg(time_scale=time_scale, x_lim=[0, ts.sequence_length])
            self.verify_basic_svg(svg)
            assert svg.count("<rect") == 10  # Sample nodes are rectangles
            assert svg.count('<path class="edge"') == 0
        svg = ts.draw_svg(force_root_branch=True, x_lim=[0, ts.sequence_length])
        self.verify_basic_svg(svg)
        assert svg.count("<rect") == 10
        assert svg.count('<path class="edge"') == 10

    def test_no_edges_with_muts(self):
        # If there is a mutation above a sample, the root branches should be there too
        # And we should be able to plot the "empty" tree because the region still has
        # mutations
        full_ts = msprime.simulate(10, mutation_rate=1, random_seed=2)
        tables = full_ts.dump_tables()
        tables.edges.clear()
        ts = tables.tree_sequence().simplify()
        assert ts.num_mutations > 0  # Should have some singletons
        svg = ts.draw_svg()
        self.verify_basic_svg(svg)
        assert svg.count("<rect") == 10
        assert svg.count('<path class="edge"') == 10
        assert svg.count('<path class="sym"') == ts.num_mutations
        assert svg.count('<line class="sym"') == ts.num_sites

    def test_empty_flanks(self):
        ts = msprime.simulate(10, random_seed=2, recombination_rate=0.1)
        assert ts.num_trees == 2
        assert 0.2 < ts.first().interval.right < 0.8
        degree_2_ts = ts.keep_intervals([[0.2, 0.8]])
        svg = degree_2_ts.draw_svg(y_axis=False)
        assert svg.count('class="tick"') == 3
        assert svg.count('<text class="lab">0.2') == 1
        assert svg.count('<text class="lab">0.8') == 1
        degree_1_ts = ts.keep_intervals([[0.05, 0.15]])
        svg = degree_1_ts.draw_svg(y_axis=False)
        assert svg.count('class="tick"') == 2
        assert svg.count('<text class="lab">0.05') == 1
        assert svg.count('<text class="lab">0.15') == 1

    def test_bad_xlim(self):
        ts = msprime.simulate(10, random_seed=2)
        svg = ts.draw_svg(x_lim=[None, None])
        self.verify_basic_svg(svg)
        with pytest.raises(ValueError, match="must be a list of length 2"):
            ts.draw_svg(x_lim=[0])
        with pytest.raises(TypeError, match="must be numeric"):
            ts.draw_svg(x_lim=[0, "a"])
        with pytest.raises(ValueError, match="must be less than"):
            ts.draw_svg(x_lim=[0.5, 0.5])
        with pytest.raises(ValueError, match="cannot be negative"):
            ts.draw_svg(x_lim=[-1, 0])
        with pytest.raises(ValueError, match="cannot be greater than"):
            ts.draw_svg(x_lim=[0, ts.sequence_length * 2])

    def test_xlim_on_empty(self):
        full_ts = msprime.simulate(10, random_seed=2)
        tables = full_ts.dump_tables()
        tables.edges.clear()
        ts = tables.tree_sequence()
        ts.draw_svg(x_lim=[0, ts.sequence_length])
        with pytest.raises(ValueError, match="whole region is empty"):
            ts.draw_svg(x_lim=[0, 0.9])

    def test_xlim_edge_cases(self):
        tables = msprime.simulate(10, random_seed=2, mutation_rate=10).dump_tables()
        # Delete edges but keep mutations
        old_sites = tables.sites.copy()
        tables.keep_intervals([[0.4, 0.6]], simplify=False)
        tables.sites.set_columns(**old_sites.asdict())
        ts = tables.tree_sequence().simplify(filter_sites=False)
        assert np.any(ts.tables.sites.position < 0.4)
        assert np.any(ts.tables.sites.position > 0.6)
        for x_lim in [None, (0, 1), (None, 1), (0, None)]:
            # All have sites in the deleted region, so should have all trees
            svg = ts.draw_svg(x_lim=x_lim)
            self.verify_basic_svg(svg, width=200 * 3)
            assert svg.count('class="tree ') == 3
        tables.sites.clear()
        tables.mutations.clear()
        ts = tables.tree_sequence().simplify()
        for x_lim, n_trees in {None: 1, (0, 1): 3, (None, 1): 2, (0, None): 2}.items():
            # No sites in the deleted region, so x_lim determines # plotted trees
            svg = ts.draw_svg(x_lim=x_lim)
            self.verify_basic_svg(svg, width=200 * n_trees)
            assert svg.count('class="tree ') == n_trees

    def test_xlim_maintains_tree_ids(self):
        ts = self.get_simple_ts()
        breaks = ts.breakpoints(as_array=True)
        svg = ts.draw_svg(x_lim=[breaks[1], breaks[4]])
        assert "t0" not in svg
        assert "t4" not in svg
        svg = ts.draw_svg(
            x_lim=[np.nextafter(breaks[1], 0), np.nextafter(breaks[4], 1)]
        )
        assert "t0" in svg
        assert "t4" in svg

    def test_xlim_maintains_site_and_mutation_ids(self):
        ts = self.get_simple_ts()
        breaks = ts.breakpoints(as_array=True)
        tree_svg = ts.at_index(1).draw_svg(x_axis=True)

        ts_svg = ts.draw_svg(x_lim=[breaks[1], breaks[2]])
        assert re.findall(r">\d+<", tree_svg) == re.findall(r">\d+<", ts_svg)  # labels
        for identifier in ["s", "m"]:
            tree_ids = re.findall(rf"{identifier}\d+", tree_svg)
            assert len(tree_ids) > 0
            ts_ids = re.findall(rf"{identifier}\d+", ts_svg)
            assert tree_ids == ts_ids

        site_pos0_in_tree1 = next(ts.at_index(1).sites()).position
        ts_svg = ts.draw_svg(x_lim=[site_pos0_in_tree1, breaks[2]])
        assert re.findall(r">\d+<", tree_svg) == re.findall(r">\d+<", ts_svg)  # labels
        for identifier in ["s", "m"]:
            tree_ids = re.findall(rf"{identifier}\d+", tree_svg)
            ts_ids = re.findall(rf"{identifier}\d+", ts_svg)
            assert tree_ids == ts_ids

        ts_svg = ts.draw_svg(x_lim=[np.nextafter(site_pos0_in_tree1, 1), breaks[2]])
        assert re.findall(r">\d+<", tree_svg) != re.findall(r">\d+<", ts_svg)  # labels
        for identifier in ["s", "m"]:
            tree_ids = re.findall(rf"{identifier}\d+", tree_svg)
            ts_ids = re.findall(rf"{identifier}\d+", ts_svg)
            assert tree_ids != ts_ids

    def test_xlim_with_ranks(self):
        ts = self.get_simple_ts()
        xlim = ts.breakpoints(as_array=True)[:2]  # plot first tree only
        svg = ts.draw_svg(x_lim=xlim, time_scale="rank", y_axis=True, y_gridlines=True)
        # excluding ".grid" in the stylesheet, there should be only 4 y-axis steps
        # for a 4 tip tree with all samples at 0: simplest check is to count gridlines
        assert len(re.findall(r"[^.]grid", svg)) == 4

    def test_half_truncated(self):
        ts = msprime.simulate(10, random_seed=2)
        ts = ts.delete_intervals([[0.4, 0.6]])
        svg = ts.draw_svg(x_lim=(0.5, 0.7), y_axis=False)
        # Only one tree and one tick shown (leftmost is an empty region)
        assert svg.count('class="tree ') == 1
        assert svg.count('class="tick"') == 1

    def test_tree_root_branch(self):
        # in the simple_ts, there are root mutations in the first tree but not the last
        ts = self.get_simple_ts()
        tree_with_root_mutations = ts.at_index(0)
        root1 = tree_with_root_mutations.root
        tree_without_root_mutations = ts.at_index(-1)
        root2 = tree_without_root_mutations.root
        svg1 = tree_with_root_mutations.draw_svg()
        svg2a = tree_without_root_mutations.draw_svg()
        svg2b = tree_without_root_mutations.draw_svg(force_root_branch=True)
        self.verify_basic_svg(svg1)
        self.verify_basic_svg(svg2a)
        self.verify_basic_svg(svg2b)
        # Last <path> should be the root branch, if it exists
        edge_str = '<path class="edge" d='
        str_pos1 = svg1.rfind(edge_str, 0, svg1.find(f">{root1}<"))
        str_pos2a = svg2a.rfind(edge_str, 0, svg2a.find(f">{root2}<"))
        str_pos2b = svg2b.rfind(edge_str, 0, svg2b.find(f">{root2}<"))
        snippet1 = svg1[str_pos1 + len(edge_str) : svg1.find(">", str_pos1)]
        snippet2a = svg2a[str_pos2a + len(edge_str) : svg2a.find(">", str_pos2a)]
        snippet2b = svg2b[str_pos2b + len(edge_str) : svg2b.find(">", str_pos2b)]
        assert snippet1.startswith('"M 0 0')
        assert snippet2a.startswith('"M 0 0')
        assert snippet2b.startswith('"M 0 0')
        assert "H 0" in snippet1
        assert not ("H 0" in snippet2a)  # No root branch
        assert "H 0" in snippet2b

    def test_debug_box(self):
        ts = self.get_simple_ts()
        svg = ts.first().draw_svg(debug_box=True)
        self.verify_basic_svg(svg)
        assert svg.count("outer_plotbox") == 1
        assert svg.count("inner_plotbox") == 1
        svg = ts.draw_svg(debug_box=True)
        self.verify_basic_svg(svg, width=200 * ts.num_trees)
        assert svg.count("outer_plotbox") == ts.num_trees + 1
        assert svg.count("inner_plotbox") == ts.num_trees + 1

    @pytest.mark.parametrize("max_trees", [-1, 0, 1])
    def test_bad_max_num_trees(self, max_trees):
        ts = self.get_simple_ts()
        with pytest.raises(ValueError, match="at least 2"):
            ts.draw_svg(max_num_trees=max_trees)

    @pytest.mark.parametrize("max_trees", [2, 4, 9])
    def test_max_num_trees(self, max_trees):
        ts = msprime.sim_ancestry(
            3, sequence_length=100, recombination_rate=0.1, random_seed=1
        )
        ts = msprime.sim_mutations(ts, rate=0.1, random_seed=1)
        assert ts.num_trees > 10
        num_sites = 0
        num_unplotted_sites = 0
        svg = ts.draw_svg(max_num_trees=max_trees)
        for tree in ts.trees():
            if (
                tree.index < (max_trees + 1) // 2
                or ts.num_trees - tree.index <= max_trees // 2
            ):
                num_sites += tree.num_sites
                assert re.search(rf"t{tree.index}[^\d]", svg) is not None
            else:
                assert re.search(rf"t{tree.index}[^\d]", svg) is None
                num_unplotted_sites += tree.num_sites
        assert num_unplotted_sites > 0
        site_strings_in_stylesheet = svg.count(".site")
        assert svg.count("site") - site_strings_in_stylesheet == num_sites
        self.verify_basic_svg(svg, width=200 * (max_trees + 1))

    def test_draw_tree_symbol_titles(self):
        tree = self.get_binary_tree()
        assert tree.tree_sequence.num_mutations > 0
        svg = tree.draw_svg(
            node_titles={u: f"NODE{u}$" for u in tree.nodes()},
            mutation_titles={m.id: f"MUT{m.id}$" for m in tree.mutations()},
        )
        for u in tree.nodes():
            assert svg.count(f"<title>NODE{u}$</title>") == 1
        for m in tree.mutations():
            assert svg.count(f"<title>MUT{m.id}$</title>") == 1
        self.verify_basic_svg(svg)

    def test_nodraw_x_axis(self):
        ts = msprime.sim_ancestry(
            1, sequence_length=100, recombination_rate=0.1, random_seed=1
        )
        svg = ts.first().draw_svg(x_axis=False, y_axis=False)
        assert 'class="x-axis"' not in svg

    def test_x_regions_ts(self):
        ts = msprime.sim_ancestry(
            3, sequence_length=100, recombination_rate=0.1, random_seed=1
        )
        regions = [(0, 10), (9, 20), (50, 90)]
        svg = ts.draw_svg(
            x_regions={r: f"reg{'ABC'[i]}" for i, r in enumerate(regions)}
        )
        self.verify_basic_svg(svg, width=200 * ts.num_trees)
        assert svg.count("x-regions") == 2  # one in stylesheet, one in svg
        assert svg.count("r0") == 1
        assert svg.count("r1") == 1
        assert svg.count("r2") == 1
        assert svg.count("r3") == 0
        assert svg.count("regA") == 1
        assert svg.count("regB") == 1
        assert svg.count("regC") == 1
        # "rect" string present for 6 samples in each tree + 3 regions + 1 in stylesheet
        assert svg.count("rect") == 6 * ts.num_trees + 3 + 1

    def test_x_regions_tree(self):
        ts = msprime.sim_ancestry(
            3, sequence_length=100, recombination_rate=0.1, random_seed=1
        )
        svg = ts.first().draw_svg(x_regions={(0, 10): "💩"})
        assert svg.count("💩") == 0
        svg = ts.first().draw_svg(x_axis=True, x_regions={(0, 10): "💩"})
        assert svg.count("💩") == 1

    def test_unsupported_x_regions(self):
        ts = msprime.sim_ancestry(
            1, sequence_length=100, recombination_rate=0.1, random_seed=1
        )
        ts.draw_svg(x_scale="treewise")
        with pytest.raises(ValueError, match="not supported for treewise"):
            ts.draw_svg(x_scale="treewise", x_regions={(0, 1): "bad"})

    def test_bad_x_regions(self):
        ts = msprime.sim_ancestry(
            1, sequence_length=100, recombination_rate=0.1, random_seed=1
        )
        with pytest.raises(ValueError, match="Invalid coordinates"):
            ts.draw_svg(x_regions={(-1, 1): "bad"})
        with pytest.raises(ValueError, match="Invalid coordinates"):
            ts.draw_svg(x_regions={(0, ts.sequence_length + 1): "bad"})
        with pytest.raises(ValueError, match="Invalid coordinates"):
            ts.draw_svg(x_regions={(1, 0): "bad"})

    def test_title(self):
        ts = msprime.sim_ancestry(1, sequence_length=100, random_seed=1)
        svg = ts.draw_svg(title="This is a title")
        assert "This is a title" in svg
        svg = ts.first().draw_svg(title="This is another title")
        assert "This is another title" in svg

    def test_bad_ts_order(self):
        ts = msprime.sim_ancestry(1, sequence_length=100, random_seed=1)
        with pytest.raises(ValueError, match="Unknown display order"):
            ts.draw_svg(order=(ts.first().nodes(order="minlex_postorder")))

    def test_good_tree_order(self):
        ts = msprime.sim_ancestry(1, sequence_length=100, random_seed=1)
        ts.first().draw_svg(order=(ts.first().nodes(order="minlex_postorder")))

    def test_nonpostorder_tree_order(self):
        tree = tskit.Tree.generate_balanced(10)
        with pytest.raises(ValueError, match="must be passed in postorder"):
            tree.draw_svg(order=(tree.nodes(order="preorder")))

    def test_only_subset_nodes_in_rank(self, caplog):
        tree = tskit.Tree.generate_comb(100)
        # Only show the last few tips of the comb. We should only use the ranks
        # from those tip times, so ticks > 5 should raise a warning
        with caplog.at_level(logging.WARNING):
            tree.draw_svg(
                order=tree.nodes(root=105, order="minlex_postorder"),
                time_scale="rank",
                y_axis=True,
                y_ticks=[0, 1, 6],
            )
            assert "lie outside the plotted axis" not in caplog.text
        with caplog.at_level(logging.WARNING):
            tree.draw_svg(
                order=tree.nodes(root=105, order="minlex_postorder"),
                time_scale="rank",
                y_axis=True,
                y_ticks=[0, 1, 10],
            )
            assert "Ticks {10: '10'} lie outside the plotted axis" in caplog.text


class TestDrawKnownSvg(TestDrawSvgBase):
    """
    Compare against known files
    """

    def verify_known_svg(self, svg, filename, save=False, **kwargs):
        # expected SVG files can be inspected in tests/data/svg/*.svg
        svg = xml.dom.minidom.parseString(
            svg
        ).toprettyxml()  # Prettify for easy viewing
        self.verify_basic_svg(svg, **kwargs)
        svg_fn = pathlib.Path(__file__).parent / "data" / "svg" / filename
        if save:
            logging.warning(f"Overwriting SVG file `{svg_fn}` with new version")
            with open(svg_fn, "w") as file:
                file.write(svg)
        with open(svg_fn, "rb") as file:
            expected_svg = file.read()
        self.assertXmlEquivalentOutputs(svg, expected_svg)

    def test_known_svg_tree_no_mut(self, overwrite_viz, draw_plotbox):
        tree = self.get_simple_ts().at_index(-1)
        svg = tree.draw_svg(
            root_svg_attributes={"id": "XYZ"},
            style=".edge {stroke: blue}",
            debug_box=draw_plotbox,
        )
        svg_no_css = svg[svg.find("</style>") :]
        assert svg_no_css.count("axes") == 0
        assert svg_no_css.count("x-axis") == 0
        assert svg_no_css.count("y-axis") == 0
        self.verify_known_svg(svg, "tree.svg", overwrite_viz)

    def test_known_svg_tree_x_axis(self, overwrite_viz, draw_plotbox):
        tree = self.get_simple_ts().at_index(1)
        svg = tree.draw_svg(
            x_axis=True,
            x_label="pos on genome",
            size=(400, 200),
            debug_box=draw_plotbox,
        )
        svg_no_css = svg[svg.find("</style>") :]
        assert svg_no_css.count("axes") == 1
        assert svg_no_css.count("x-axis") == 1
        assert svg_no_css.count("title") == 1
        assert svg_no_css.count("y-axis") == 0
        self.verify_known_svg(svg, "tree_x_axis.svg", overwrite_viz, width=400)

    def test_known_svg_tree_y_axis_rank(self, overwrite_viz, draw_plotbox):
        tree = self.get_simple_ts().at_index(1)
        label = "Time (relative steps)"
        svg = tree.draw_svg(
            y_axis=True,
            y_label=label,
            y_gridlines=True,
            time_scale="rank",
            style=".y-axis line.grid {stroke: #CCCCCC}",
            debug_box=draw_plotbox,
        )
        svg_no_css = svg[svg.find("</style>") :]
        node_times = [tree.time(u) for u in tree.nodes()]
        assert label in svg_no_css
        assert svg_no_css.count('class="grid"') == len(set(node_times))
        assert svg_no_css.count("axes") == 1
        assert svg_no_css.count("x-axis") == 0
        assert svg_no_css.count("y-axis") == 1
        assert svg_no_css.count("title") == 1
        self.verify_known_svg(svg, "tree_y_axis_rank.svg", overwrite_viz)

    def test_known_svg_tree_both_axes(self, overwrite_viz, draw_plotbox):
        tree = self.get_simple_ts().at_index(-1)
        svg = tree.draw_svg(x_axis=True, y_axis=True, debug_box=draw_plotbox)
        svg_no_css = svg[svg.find("</style>") :]
        assert svg_no_css.count("axes") == 1
        assert svg_no_css.count("x-axis") == 1
        assert svg_no_css.count("y-axis") == 1
        assert svg_no_css.count("title") == 2
        self.verify_known_svg(svg, "tree_both_axes.svg", overwrite_viz)

    def test_known_svg_tree_root_mut(self, overwrite_viz, draw_plotbox):
        tree = self.get_simple_ts().at_index(0)  # Tree 0 has a few mutations above root
        svg = tree.draw_svg(debug_box=draw_plotbox)
        self.verify_known_svg(svg, "tree_muts.svg", overwrite_viz)

    def test_known_svg_tree_mut_all_edge(self, overwrite_viz, draw_plotbox):
        tree = self.get_simple_ts().at_index(1)
        size = (300, 400)
        svg = tree.draw_svg(
            size=size,
            debug_box=draw_plotbox,
            all_edge_mutations=True,
            x_axis=True,
            title="All mutations tree: background shading shown",
        )
        self.verify_known_svg(
            svg, "tree_muts_all_edge.svg", overwrite_viz, width=size[0], height=size[1]
        )

    def test_known_svg_tree_timed_root_mut(self, overwrite_viz, draw_plotbox):
        tree = self.get_simple_ts(use_mutation_times=True).at_index(0)
        svg = tree.draw_svg(debug_box=draw_plotbox)
        self.verify_known_svg(svg, "tree_timed_muts.svg", overwrite_viz)

    def test_known_svg_ts(self, overwrite_viz, draw_plotbox):
        ts = self.get_simple_ts()
        svg = ts.draw_svg(debug_box=draw_plotbox)
        svg_no_css = svg[svg.find("</style>") :]
        assert svg_no_css.count("axes") == 1
        assert svg_no_css.count("x-axis") == 1
        assert svg_no_css.count("y-axis") == 0
        assert svg_no_css.count('class="site ') == ts.num_sites
        assert svg_no_css.count('class="mut ') == ts.num_mutations * 2
        self.verify_known_svg(svg, "ts.svg", overwrite_viz, width=200 * ts.num_trees)

    def test_known_svg_ts_title(self, overwrite_viz, draw_plotbox):
        ts = self.get_simple_ts()
        svg = ts.draw_svg(title="The main plot title", debug_box=draw_plotbox)
        self.verify_known_svg(
            svg, "ts_title.svg", overwrite_viz, width=200 * ts.num_trees
        )

    def test_known_svg_ts_no_axes(self, overwrite_viz, draw_plotbox):
        ts = self.get_simple_ts()
        svg = ts.draw_svg(x_axis=False, debug_box=draw_plotbox)
        svg_no_css = svg[svg.find("</style>") :]
        assert svg_no_css.count("axes") == 0
        assert svg_no_css.count("x-axis") == 0
        assert svg_no_css.count("y-axis") == 0
        assert 'class="site ' not in svg_no_css
        assert svg_no_css.count('class="mut ') == ts.num_mutations
        self.verify_known_svg(
            svg, "ts_no_axes.svg", overwrite_viz, width=200 * ts.num_trees
        )

    def test_known_svg_ts_internal_sample(self, overwrite_viz, draw_plotbox):
        ts = tsutil.jiggle_samples(self.get_simple_ts())
        svg = ts.draw_svg(
            root_svg_attributes={"id": "XYZ"},
            style="#XYZ .leaf .sym {fill: magenta} #XYZ .sample > .sym {fill: cyan}",
            debug_box=draw_plotbox,
        )
        self.verify_known_svg(
            svg, "internal_sample_ts.svg", overwrite_viz, width=200 * ts.num_trees
        )

    def test_known_svg_ts_highlighted_mut(self, overwrite_viz, draw_plotbox):
        ts = self.get_simple_ts()
        style = (
            ".edge {stroke: grey}"
            ".mut .sym{stroke:pink} .mut text{fill:pink}"
            ".mut.m2 .sym,.m2>line, .m2>.node .edge{stroke:red} .mut.m2 text{fill:red}"
            ".mut.m3 .sym,.m3>line, .m3>.node .edge{stroke:cyan} .mut.m3 text{fill:cyan}"
            ".mut.m4 .sym,.m4>line, .m4>.node .edge{stroke:blue} .mut.m4 text{fill:blue}"
        )
        svg = ts.draw_svg(style=style, debug_box=draw_plotbox)
        self.verify_known_svg(
            svg, "ts_mut_highlight.svg", overwrite_viz, width=200 * ts.num_trees
        )

    def test_known_svg_ts_rank(self, overwrite_viz, draw_plotbox):
        ts = self.get_simple_ts()
        svg1 = ts.draw_svg(time_scale="rank", y_axis=True, debug_box=draw_plotbox)
        ts = self.get_simple_ts(use_mutation_times=True)
        svg2 = ts.draw_svg(time_scale="rank", y_axis=True, debug_box=draw_plotbox)
        assert svg1.count('class="site ') == ts.num_sites
        assert svg1.count('class="mut ') == ts.num_mutations * 2
        assert svg1.replace(" unknown_time", "") == svg2  # Trim the unknown_time class
        self.verify_known_svg(
            svg1, "ts_rank.svg", overwrite_viz, width=200 * ts.num_trees
        )

    @pytest.mark.skip(reason="Fails on CI as OSX gives different random numbers")
    def test_known_svg_nonbinary_ts(self, overwrite_viz, draw_plotbox):
        ts = self.get_nonbinary_ts()
        svg = ts.draw_svg(time_scale="log_time", debug_box=draw_plotbox)
        assert svg.count('class="site ') == ts.num_sites
        assert svg.count('class="mut ') == ts.num_mutations * 2
        self.verify_known_svg(
            svg, "ts_nonbinary.svg", overwrite_viz, width=200 * ts.num_trees
        )

    def test_known_svg_ts_plain(self, overwrite_viz, draw_plotbox):
        """
        Plain style: no background shading and a variable scale X axis with no sites
        """
        ts = self.get_simple_ts()
        svg = ts.draw_svg(x_scale="treewise", debug_box=draw_plotbox)
        assert svg.count('class="site ') == 0
        assert svg.count('class="mut ') == ts.num_mutations
        self.verify_known_svg(
            svg, "ts_plain.svg", overwrite_viz, width=200 * ts.num_trees
        )

    def test_known_svg_ts_plain_no_xlab(self, overwrite_viz, draw_plotbox):
        """
        Plain style: no background shading and a variable scale X axis with no sites
        """
        ts = self.get_simple_ts()
        svg = ts.draw_svg(x_scale="treewise", x_label="", debug_box=draw_plotbox)
        assert "Genome position" not in svg
        self.verify_known_svg(
            svg, "ts_plain_no_xlab.svg", overwrite_viz, width=200 * ts.num_trees
        )

    def test_known_svg_ts_plain_y(self, overwrite_viz, draw_plotbox):
        """
        Plain style: no background shading and a variable scale X axis with no sites
        """
        ts = self.get_simple_ts()
        ticks = [0, 5, 10]
        svg = ts.draw_svg(
            x_scale="treewise",
            y_axis=True,
            y_ticks=ticks,
            y_gridlines=True,
            style=".y-axis line.grid {stroke: #CCCCCC}",
            debug_box=draw_plotbox,
        )
        self.verify_known_svg(
            svg, "ts_plain_y.svg", overwrite_viz, width=200 * ts.num_trees
        )

    def test_known_svg_ts_with_xlabel(self, overwrite_viz, draw_plotbox):
        """
        Style with X axis label
        """
        ts = self.get_simple_ts()
        x_label = "genomic position (bp)"
        svg = ts.draw_svg(x_label=x_label, debug_box=draw_plotbox)
        assert x_label in svg
        self.verify_known_svg(
            svg, "ts_xlabel.svg", overwrite_viz, width=200 * ts.num_trees
        )

    def test_known_svg_ts_y_axis(self, overwrite_viz, draw_plotbox):
        tables = self.get_simple_ts().dump_tables()
        # set units
        tables.time_units = "generations"
        ts = tables.tree_sequence()
        svg = ts.draw_svg(y_axis=True, debug_box=draw_plotbox)
        assert "Time ago (generations)" in svg
        self.verify_known_svg(
            svg, "ts_y_axis.svg", overwrite_viz, width=200 * ts.num_trees
        )

    def test_known_svg_ts_y_axis_regular(self, overwrite_viz, draw_plotbox):
        # This should have gridlines
        ts = self.get_simple_ts()
        ticks = np.arange(0, max(ts.tables.nodes.time), 1)
        svg = ts.draw_svg(
            y_axis=True, y_ticks=ticks, y_gridlines=True, debug_box=draw_plotbox
        )
        assert svg.count('class="grid"') == len(ticks)
        self.verify_known_svg(
            svg, "ts_y_axis_regular.svg", overwrite_viz, width=200 * ts.num_trees
        )

    def test_known_svg_ts_y_axis_log(self, overwrite_viz, draw_plotbox):
        ts = self.get_simple_ts()
        svg = ts.draw_svg(
            y_axis=True,
            y_label="Time (log scale)",
            time_scale="log_time",
            debug_box=draw_plotbox,
        )
        self.verify_known_svg(
            svg, "ts_y_axis_log.svg", overwrite_viz, width=200 * ts.num_trees
        )

    def test_known_svg_ts_mutation_times(self, overwrite_viz, draw_plotbox):
        ts = self.get_simple_ts(use_mutation_times=True)
        svg = ts.draw_svg(debug_box=draw_plotbox)
        assert svg.count('class="site ') == ts.num_sites
        assert svg.count('class="mut ') == ts.num_mutations * 2
        self.verify_known_svg(
            svg, "ts_mut_times.svg", overwrite_viz, width=200 * ts.num_trees
        )

    def test_known_svg_ts_titles(self, overwrite_viz, draw_plotbox):
        ts = self.get_simple_ts(use_mutation_times=True)
        svg = ts.draw_svg(
            node_titles={nd.id: f"NoDe{nd.id}!" for nd in ts.nodes()},
            mutation_titles={m.id: f"MuT{m.id}!" for m in ts.mutations()},
            debug_box=draw_plotbox,
        )
        for nd in ts.nodes():
            if nd.is_sample():
                assert svg.count(f"<title>NoDe{nd.id}!</title>") == ts.num_trees
            else:
                assert f"<title>NoDe{nd.id}!</title>" in svg
        for m in ts.mutations():
            assert (
                svg.count(f"<title>MuT{m.id}!</title>") == 2
            )  # 1 on tree, 1 on x-axis
        self.verify_known_svg(
            svg, "ts_mut_times_titles.svg", overwrite_viz, width=200 * ts.num_trees
        )

    def test_known_svg_ts_mutation_times_logscale(self, overwrite_viz, draw_plotbox):
        ts = self.get_simple_ts(use_mutation_times=True)
        svg = ts.draw_svg(time_scale="log_time", debug_box=draw_plotbox)
        assert svg.count('class="site ') == ts.num_sites
        assert svg.count('class="mut ') == ts.num_mutations * 2
        self.verify_known_svg(
            svg, "ts_mut_times_logscale.svg", overwrite_viz, width=200 * ts.num_trees
        )

    def test_known_svg_ts_mut_no_edges(self, overwrite_viz, draw_plotbox):
        # An example with some muts on axis but not on a visible node
        ts = msprime.simulate(10, random_seed=2, mutation_rate=1)
        tables = ts.dump_tables()
        tables.edges.clear()
        tables.mutations.time = np.full_like(tables.mutations.time, tskit.UNKNOWN_TIME)
        ts_no_edges = tables.tree_sequence()
        with pytest.warns(UserWarning, match="nodes which are not present"):
            svg = ts_no_edges.draw_svg(debug_box=draw_plotbox)
            self.verify_known_svg(
                svg,
                "ts_mutations_no_edges.svg",
                overwrite_viz,
                width=200 * ts.num_trees,
            )

    def test_known_svg_ts_timed_mut_no_edges(self, overwrite_viz, draw_plotbox):
        # An example with some muts on axis but not on a visible node
        ts = msprime.simulate(10, random_seed=2, mutation_rate=1)
        tables = ts.dump_tables()
        tables.edges.clear()
        tables.mutations.time = np.arange(
            ts.num_mutations, dtype=tables.mutations.time.dtype
        )
        ts_no_edges = tables.tree_sequence()

        with pytest.warns(UserWarning, match="nodes which are not present"):
            svg = ts_no_edges.draw_svg(debug_box=draw_plotbox)
            self.verify_known_svg(
                svg,
                "ts_mutations_timed_no_edges.svg",
                overwrite_viz,
                width=200 * ts.num_trees,
            )

    def test_known_svg_ts_multiroot(self, overwrite_viz, draw_plotbox):
        tables = wf.wf_sim(
            6,
            5,
            seed=1,
            deep_history=False,
            initial_generation_samples=False,
            num_loci=8,
        )
        tables.sort()
        ts = tables.tree_sequence().simplify()
        tables = tsutil.jukes_cantor(ts, 10, mu=0.1, seed=123).dump_tables()
        # Set unknown times, so we are msprime 0.7.4 and 1.0.0 compatible
        tables.mutations.time = np.full(tables.mutations.num_rows, tskit.UNKNOWN_TIME)
        svg = tables.tree_sequence().draw_svg(
            y_axis=True, y_gridlines=True, debug_box=draw_plotbox
        )
        self.verify_known_svg(
            svg, "ts_multiroot.svg", overwrite_viz, width=200 * ts.num_trees
        )
        assert "Time ago (generations)" in svg

    def test_known_svg_ts_xlim(self, overwrite_viz, draw_plotbox):
        ts = self.get_simple_ts()
        svg = ts.draw_svg(x_lim=[0.051, 0.9], debug_box=draw_plotbox)
        num_trees = sum(1 for b in ts.breakpoints() if 0.051 <= b < 0.9) + 1
        self.verify_known_svg(svg, "ts_x_lim.svg", overwrite_viz, width=200 * num_trees)

    @pytest.mark.skipif(IS_WINDOWS, reason="Msprime gives different result on Windows")
    def test_known_max_num_trees(self, overwrite_viz, draw_plotbox):
        max_trees = 5
        ts = msprime.sim_ancestry(
            3, sequence_length=100, recombination_rate=0.1, random_seed=1
        )
        ts = msprime.sim_mutations(ts, rate=0.01, random_seed=1)
        assert ts.num_trees > 10
        first_break = next(ts.trees()).interval.right
        # limit to just past the first tree
        svg = ts.draw_svg(
            max_num_trees=max_trees,
            x_lim=(first_break + 0.1, ts.sequence_length - 0.1),
            y_axis=True,
            time_scale="log_time",
            debug_box=draw_plotbox,
        )
        self.verify_known_svg(
            svg, "ts_max_trees.svg", overwrite_viz, width=200 * (max_trees + 1)
        )

    @pytest.mark.skipif(IS_WINDOWS, reason="Msprime gives different result on Windows")
    def test_known_max_num_trees_treewise(self, overwrite_viz, draw_plotbox):
        max_trees = 5
        ts = msprime.sim_ancestry(
            3, sequence_length=100, recombination_rate=0.1, random_seed=1
        )
        ts = msprime.sim_mutations(ts, rate=0.01, random_seed=1)
        assert ts.num_trees > 10
        first_break = next(ts.trees()).interval.right
        svg = ts.draw_svg(
            max_num_trees=max_trees,
            x_lim=(first_break + 0.1, ts.sequence_length - 0.1),
            y_axis=True,
            x_scale="treewise",
            debug_box=draw_plotbox,
        )
        self.verify_known_svg(
            svg, "ts_max_trees_treewise.svg", overwrite_viz, width=200 * (max_trees + 1)
        )

    def test_known_svg_tree_collapsed(self, overwrite_viz, draw_plotbox):
        tree = tskit.Tree.generate_balanced(8)
        remove_nodes = set()
        remove_nodes_below = {8, 13}
        for u in remove_nodes_below:
            subtree_nodes = set(tree.nodes(root=u)) - {u}
            remove_nodes.update(subtree_nodes)
        order = [
            u for u in tree.nodes(order="minlex_postorder") if u not in remove_nodes
        ]
        svg = tree.draw_svg(order=order, debug_box=draw_plotbox)
        assert svg.count("multi") == len(remove_nodes_below)
        assert svg.count(">+2<") == 1  # One tip has 2 samples below it
        assert svg.count(">+4<") == 1  # Another tip has 4 samples below it
        for u in order:
            assert f'n{u}"' in svg or f"n{u} " in svg
        for u in remove_nodes:
            assert f'n{u}"' not in svg and f"n{u} " not in svg
        self.verify_known_svg(svg, "tree_simple_collapsed.svg", overwrite_viz)

    def test_known_svg_tree_subtree(self, overwrite_viz, draw_plotbox):
        tree = tskit.Tree.generate_balanced(8)
        order = [u for u in tree.nodes(root=10, order="minlex_postorder")]
        # The balanced tree has all descendants of nodes 10 with IDs < 10
        assert np.all(np.array(order) <= 10)
        svg = tree.draw_svg(order=order, debug_box=draw_plotbox)
        for u in order:
            assert f'n{u}"' in svg or f"n{u} " in svg
        for u in set(tree.nodes()) - set(order):
            assert f'n{u}"' not in svg and f"n{u} " not in svg
        self.verify_known_svg(svg, "tree_subtree.svg", overwrite_viz, has_root=False)

    def test_known_svg_tree_subtrees_with_collapsed(self, overwrite_viz, draw_plotbox):
        # Two subtrees, one with a collapsed node below node 16
        tree = tskit.Tree.generate_balanced(16)
        roots = [22, 25]
        order = []
        remove_nodes_below = 16
        remove_nodes = set(tree.nodes(root=remove_nodes_below)) - {remove_nodes_below}
        for root in roots:
            order += [
                u
                for u in tree.nodes(root=root, order="minlex_postorder")
                if u not in remove_nodes
            ]
        svg = tree.draw_svg(order=order, debug_box=draw_plotbox)
        assert svg.count("multi") == 1  # One tip representing multiple nodes
        for u in order:
            assert f'n{u}"' in svg or f"n{u} " in svg
        for u in remove_nodes:
            assert f'n{u}"' not in svg and f"n{u} " not in svg
        self.verify_known_svg(
            svg, "tree_subtrees_with_collapsed.svg", overwrite_viz, has_root=False
        )

    def test_known_svg_tree_polytomy(self, overwrite_viz, draw_plotbox):
        tracked_nodes = [20, 24, 25, 27, 28, 29]
        tree = tskit.Tree.generate_balanced(30, arity=4)
        svg = tree.draw_svg(
            time_scale="rank",
            debug_box=draw_plotbox,
            size=(600, 200),
            style="".join(f".n{u} > .sym {{fill: cyan}}" for u in tracked_nodes + [39]),
        )
        self.verify_known_svg(
            svg, "tree_poly.svg", overwrite_viz, width=600, height=200
        )

    def test_known_svg_tree_polytomy_tracked(self, overwrite_viz, draw_plotbox):
        tracked_nodes = [20, 24, 25, 27, 28, 29]
        tree = tskit.Tree.generate_balanced(30, arity=4, tracked_samples=tracked_nodes)
        svg = tree.draw_svg(
            time_scale="rank",
            order=drawing._postorder_tracked_minlex_traversal(tree),
            debug_box=draw_plotbox,
            pack_untracked_polytomies=True,
            size=(600, 200),
            style="".join(f".n{u} > .sym {{fill: cyan}}" for u in tracked_nodes + [39]),
        )
        self.verify_known_svg(
            svg, "tree_poly_tracked.svg", overwrite_viz, width=600, height=200
        )

    def test_known_svg_tree_polytomy_tracked_collapse(
        self, overwrite_viz, draw_plotbox
    ):
        tracked_nodes = [20, 24, 25, 27, 28, 29]
        tree = tskit.Tree.generate_balanced(30, arity=4, tracked_samples=tracked_nodes)
        svg = tree.draw_svg(
            time_scale="rank",
            order=drawing._postorder_tracked_minlex_traversal(
                tree, collapse_tracked=True
            ),
            debug_box=draw_plotbox,
            size=(600, 200),
            pack_untracked_polytomies=True,
            style="".join(f".n{u} > .sym {{fill: cyan}}" for u in tracked_nodes + [39]),
        )
        self.verify_known_svg(
            svg, "tree_poly_tracked_collapse.svg", overwrite_viz, width=600, height=200
        )


class TestRounding:
    def test_rnd(self):
        assert 0 == drawing.rnd(0)
        assert math.inf == drawing.rnd(math.inf)
        assert 1 == drawing.rnd(1)
        assert 1.1 == drawing.rnd(1.1)
        assert 1.11111 == drawing.rnd(1.111111)
        assert 1111110 == drawing.rnd(1111111)
        assert 123.457 == drawing.rnd(123.4567)
        assert 123.456 == drawing.rnd(123.4564)


class TestDrawingTraversals:
    # TODO: test drawing._postorder_tracked_minlex_traversal and
    # drawing._postorder_tracked_node_traversal
    pass


--- ../../tskit/python/tests/test_utilities.py ---


"""
Tests for the various testing utilities.
"""
import msprime
import numpy as np
import pytest

import tests.tsutil as tsutil
import tskit


class TestJukesCantor:
    """
    Check that the we get useable tree sequences.
    """

    def verify(self, ts):
        tables = ts.dump_tables()
        tables.compute_mutation_parents()
        tables.assert_equals(ts.tables)
        # This will catch inconsistent mutations.
        assert ts.genotype_matrix() is not None

    def test_n10_multiroot(self):
        ts = msprime.simulate(10, random_seed=1)
        ts = ts.decapitate(np.max(ts.tables.nodes.time) / 2)
        ts = tsutil.jukes_cantor(ts, 1, 2, seed=7)
        self.verify(ts)

    def test_n50_multiroot(self):
        ts = msprime.simulate(50, random_seed=1)
        ts = ts.decapitate(np.max(ts.tables.nodes.time) / 2)
        ts = tsutil.jukes_cantor(ts, 5, 2, seed=2)
        self.verify(ts)

    def test_silent_mutations(self):
        ts = msprime.simulate(50, random_seed=1)
        ts = tsutil.jukes_cantor(ts, 5, 2, seed=2)
        num_silent = 0
        for m in ts.mutations():
            if (
                m.parent != -1
                and ts.mutation(m.parent).derived_state == m.derived_state
            ):
                num_silent += 1
        assert num_silent > 20


class TestCaterpillarTree:
    """
    Tests for the caterpillar tree method.
    """

    def verify(self, ts, n):
        assert ts.num_trees == 1
        assert ts.num_nodes == ts.num_samples * 2 - 1
        tree = ts.first()
        for j in range(1, n):
            assert tree.parent(j) == n + j - 1
        # This will catch inconsistent mutations.
        assert ts.genotype_matrix() is not None

    def test_n_2(self):
        ts = tsutil.caterpillar_tree(2)
        self.verify(ts, 2)

    def test_n_3(self):
        ts = tsutil.caterpillar_tree(3)
        self.verify(ts, 3)

    def test_n_50(self):
        ts = tsutil.caterpillar_tree(50)
        self.verify(ts, 50)

    def test_n_5_sites(self):
        ts = tsutil.caterpillar_tree(5, num_sites=4)
        self.verify(ts, 5)
        assert ts.num_sites == 4
        assert ts.num_mutations == 4
        assert list(ts.tables.sites.position) == [0.2, 0.4, 0.6, 0.8]
        ts = tsutil.caterpillar_tree(5, num_sites=1, num_mutations=1)
        assert ts.num_sites == 1
        assert ts.num_mutations == 1
        site = ts.site(0)
        assert site.mutations[0].node == 7

    def test_n_5_mutations(self):
        ts = tsutil.caterpillar_tree(5, num_sites=1, num_mutations=3)
        self.verify(ts, 5)
        assert ts.num_sites == 1
        assert ts.num_mutations == 3
        node = ts.tables.mutations.node
        assert list(node) == [7, 6, 5]

    def test_n_many_mutations(self):
        for n in range(10, 15):
            for num_mutations in range(0, n - 1):
                ts = tsutil.caterpillar_tree(
                    n, num_sites=1, num_mutations=num_mutations
                )
                self.verify(ts, n)
                assert ts.num_sites == 1
                assert ts.num_mutations == num_mutations
            for num_mutations in range(n - 1, n + 2):
                with pytest.raises(ValueError):
                    tsutil.caterpillar_tree(n, num_sites=1, num_mutations=num_mutations)


class TestInsertIndividuals:
    """
    Test that we insert individuals correctly.
    """

    def test_ploidy_1(self):
        ts = msprime.simulate(10, random_seed=1)
        assert ts.num_individuals == 0
        ts = tsutil.insert_individuals(ts, ploidy=1)
        assert ts.num_individuals == 10
        for j, ind in enumerate(ts.individuals()):
            assert list(ind.nodes) == [j]

    def test_ploidy_2(self):
        ts = msprime.simulate(10, random_seed=1)
        assert ts.num_individuals == 0
        ts = tsutil.insert_individuals(ts, ploidy=2)
        assert ts.num_individuals == 5
        for j, ind in enumerate(ts.individuals()):
            assert list(ind.nodes) == [2 * j, 2 * j + 1]

    def test_ploidy_2_reversed(self):
        ts = msprime.simulate(10, random_seed=1)
        assert ts.num_individuals == 0
        samples = ts.samples()[::-1]
        ts = tsutil.insert_individuals(ts, nodes=samples, ploidy=2)
        assert ts.num_individuals == 5
        for j, ind in enumerate(ts.individuals()):
            assert list(ind.nodes) == [samples[2 * j + 1], samples[2 * j]]


class TestSortIndividuals:
    def test_sort_individuals(self):
        tables = tskit.TableCollection()
        tables.individuals.add_row(parents=[1], metadata=b"0")
        tables.individuals.add_row(parents=[-1], metadata=b"1")
        tsutil.sort_individual_table(tables)
        assert tables.individuals.metadata.tobytes() == b"10"

        tables = tskit.TableCollection()
        tables.individuals.add_row(parents=[2, 3], metadata=b"0")
        tables.individuals.add_row(parents=[5], metadata=b"1")
        tables.individuals.add_row(parents=[-1], metadata=b"2")
        tables.individuals.add_row(parents=[-1], metadata=b"3")
        tables.individuals.add_row(parents=[3], metadata=b"4")
        tables.individuals.add_row(parents=[4], metadata=b"5")

        tsutil.sort_individual_table(tables)
        assert tables.individuals.metadata.tobytes() == b"342501"

        tables = tskit.TableCollection()
        tables.individuals.add_row(parents=[1], metadata=b"0")
        tables.individuals.add_row(parents=[0], metadata=b"1")
        with pytest.raises(ValueError, match="Individual pedigree has cycles"):
            tsutil.sort_individual_table(tables)


class TestQuintuplyLinkedTrees:
    def test_branch_operations_num_children(self):
        qlt = tsutil.QuintuplyLinkedTree(3)
        assert np.sum(qlt.num_children) == 0
        qlt.insert_branch(2, 0)
        assert qlt.num_children[2] == 1
        assert np.sum(qlt.num_children) == 1

        qlt.remove_branch(2, 0)
        assert qlt.num_children[2] == 0

    def test_edge_operations(self):
        tt = tskit.Tree.generate_balanced(3)
        tts = tt.tree_sequence

        for _, qlt in tsutil.algorithm_R(tts):
            assert np.sum(qlt.edge != -1) == tt.num_edges
            self.verify_tree_edges(qlt, tts)

    def verify_tree_edges(self, quintuply_linked_tree, tts):
        for edge in tts.edges():
            assert quintuply_linked_tree.edge[edge.child] == edge.id
            assert quintuply_linked_tree.parent[edge.child] == edge.parent


--- ../../tskit/python/tests/test_distance_metrics.py ---


"""
Tests for tree distance metrics.
"""
import io
import itertools
import math
import unittest

import dendropy
import msprime
import numpy as np
import pytest
from dendropy.calculate import treecompare

import _tskit
import tests
import tests.tsutil as tsutil
import tskit


def c_kc_distance(tree1, tree2, lambda_=0):
    """
    Simplified version of the naive_kc_distance() function above.
    Written without Python features to aid writing C implementation.
    """
    samples = tree1.tree_sequence.samples()
    if tree1.tree_sequence.num_samples != tree2.tree_sequence.num_samples:
        raise ValueError("Trees must have the same samples")
    for sample1, sample2 in zip(samples, tree2.tree_sequence.samples()):
        if sample1 != sample2:
            raise ValueError("Trees must have the same samples")
    if not len(tree1.roots) == len(tree2.roots) == 1:
        raise ValueError("Trees must have one root")
    for tree in [tree1, tree2]:
        for u in range(tree.tree_sequence.num_nodes):
            if tree.num_children(u) == 1:
                raise ValueError("Unary nodes are not supported")

    n = tree1.tree_sequence.num_samples
    vecs1 = KCVectors(n)
    fill_kc_vectors(tree1, vecs1)
    vecs2 = KCVectors(n)
    fill_kc_vectors(tree2, vecs2)
    return norm_kc_vectors(vecs1, vecs2, lambda_)


def naive_kc_distance(tree1, tree2, lambda_=0):
    """
    Returns the Kendall-Colijn distance between the specified pair of trees.
    lambda_ determines weight of topology vs branch lengths in calculating
    the distance. Set lambda_ at 0 to only consider topology, set at 1 to
    only consider branch lengths. See Kendall & Colijn (2016):
    https://academic.oup.com/mbe/article/33/10/2735/2925548
    """
    samples = tree1.tree_sequence.samples()
    if not np.array_equal(samples, tree2.tree_sequence.samples()):
        raise ValueError("Trees must have the same samples")
    if not len(tree1.roots) == len(tree2.roots) == 1:
        raise ValueError("Trees must have one root")
    for tree in [tree1, tree2]:
        for u in tree.nodes():
            if tree.num_children(u) == 1:
                raise ValueError("Unary nodes are not supported")

    n = samples.shape[0]
    N = (n * (n - 1)) // 2
    m = [np.zeros(N + n), np.zeros(N + n)]
    M = [np.zeros(N + n), np.zeros(N + n)]
    for tree_index, tree in enumerate([tree1, tree2]):
        for sample in range(n):
            m[tree_index][N + sample] = 1
            M[tree_index][N + sample] = tree.branch_length(sample)

        for n1, n2 in itertools.combinations(range(n), 2):
            mrca = tree.mrca(samples[n1], samples[n2])
            depth = 0
            u = tree.parent(mrca)
            while u != tskit.NULL:
                depth += 1
                u = tree.parent(u)
            pair_index = n1 * (n1 - 2 * n + 1) // -2 + n2 - n1 - 1
            m[tree_index][pair_index] = depth
            M[tree_index][pair_index] = tree.time(tree.root) - tree.time(mrca)

    return np.linalg.norm((1 - lambda_) * (m[0] - m[1]) + lambda_ * (M[0] - M[1]))


class KCVectors:
    """
    Manages the two vectors (m and M) of a tree used to compute the
    KC distance between trees. For any two samples, u and v,
    m and M capture the distance of mrca(u, v) to the root in
    number of edges and time, respectively.

    See Kendall & Colijn (2016):
    https://academic.oup.com/mbe/article/33/10/2735/2925548
    """

    def __init__(self, n):
        self.n = n
        self.N = (self.n * (self.n - 1)) // 2
        self.m = np.zeros(self.N + self.n)
        self.M = np.zeros(self.N + self.n)


def fill_kc_vectors(tree, kc_vecs):
    sample_index_map = np.zeros(tree.tree_sequence.num_nodes)
    for j, u in enumerate(tree.tree_sequence.samples()):
        sample_index_map[u] = j
    for root in tree.roots:
        stack = [(tree.root, 0)]
        while len(stack) > 0:
            u, depth = stack.pop()
            if tree.is_sample(u):
                time = tree.branch_length(u)
                update_kc_vectors_single_leaf(kc_vecs, u, time, sample_index_map)

            c1 = tree.left_child(u)
            while c1 != tskit.NULL:
                stack.append((c1, depth + 1))
                c2 = tree.right_sib(c1)
                while c2 != tskit.NULL:
                    update_kc_vectors_all_pairs(
                        tree, kc_vecs, c1, c2, depth, tree.time(root) - tree.time(u)
                    )
                    c2 = tree.right_sib(c2)
                c1 = tree.right_sib(c1)


def update_kc_vectors_single_leaf(kc_vecs, u, time, sample_index_map):
    u_index = int(sample_index_map[u])
    kc_vecs.m[kc_vecs.N + u_index] = 1
    kc_vecs.M[kc_vecs.N + u_index] = time


def update_kc_vectors_all_pairs(tree, kc_vecs, c1, c2, depth, time):
    s1_index = tree.left_sample(c1)
    while True:
        s2_index = tree.left_sample(c2)
        while True:
            update_kc_vectors_pair(kc_vecs, s1_index, s2_index, depth, time)
            if s2_index == tree.right_sample(c2):
                break
            s2_index = tree.next_sample(s2_index)
        if s1_index == tree.right_sample(c1):
            break
        s1_index = tree.next_sample(s1_index)


def update_kc_vectors_pair(kc_vecs, n1, n2, depth, time):
    if n1 > n2:
        n1, n2 = n2, n1
    pair_index = n2 - n1 - 1 + (-1 * n1 * (n1 - 2 * kc_vecs.n + 1)) // 2

    kc_vecs.m[pair_index] = depth
    kc_vecs.M[pair_index] = time


def norm_kc_vectors(kc_vecs1, kc_vecs2, lambda_):
    vT1 = 0
    vT2 = 0
    distance_sum = 0
    for i in range(kc_vecs1.n + kc_vecs1.N):
        vT1 = (kc_vecs1.m[i] * (1 - lambda_)) + (lambda_ * kc_vecs1.M[i])
        vT2 = (kc_vecs2.m[i] * (1 - lambda_)) + (lambda_ * kc_vecs2.M[i])
        distance_sum += (vT1 - vT2) ** 2

    return math.sqrt(distance_sum)


class TestKCMetric(unittest.TestCase):
    """
    Tests on the KC metric distances.
    """

    def test_same_tree_zero_distance(self):
        for n in range(2, 10):
            for seed in range(1, 10):
                ts = msprime.simulate(n, random_seed=seed)
                tree = next(ts.trees(sample_lists=True))
                assert naive_kc_distance(tree, tree) == 0
                assert c_kc_distance(tree, tree) == 0
                assert tree.kc_distance(tree) == 0
                ts = msprime.simulate(n, random_seed=seed)
                tree2 = next(ts.trees(sample_lists=True))
                assert naive_kc_distance(tree, tree2) == 0
                assert c_kc_distance(tree, tree2) == 0
                assert tree.kc_distance(tree2) == 0

    def test_sample_2_zero_distance(self):
        # All trees with 2 leaves must be equal distance from each other.
        for seed in range(1, 10):
            ts1 = msprime.simulate(2, random_seed=seed)
            tree1 = next(ts1.trees(sample_lists=True))
            ts2 = msprime.simulate(2, random_seed=seed + 1)
            tree2 = next(ts2.trees(sample_lists=True))
            assert naive_kc_distance(tree1, tree2, 0) == 0
            assert c_kc_distance(tree1, tree2, 0) == 0
            assert tree1.kc_distance(tree2, 0) == 0

    def test_different_samples_error(self):
        tree1 = next(msprime.simulate(10, random_seed=1).trees(sample_lists=True))
        tree2 = next(msprime.simulate(2, random_seed=1).trees(sample_lists=True))
        with pytest.raises(ValueError):
            naive_kc_distance(tree1, tree2)
        with pytest.raises(ValueError):
            c_kc_distance(tree1, tree2)
        with pytest.raises(_tskit.LibraryError):
            tree1.kc_distance(tree2)

        ts1 = msprime.simulate(10, random_seed=1)
        nmap = np.arange(0, ts1.num_nodes)[::-1]
        ts2 = tsutil.permute_nodes(ts1, nmap)
        tree1 = next(ts1.trees(sample_lists=True))
        tree2 = next(ts2.trees(sample_lists=True))
        with pytest.raises(ValueError):
            naive_kc_distance(tree1, tree2)
        with pytest.raises(ValueError):
            c_kc_distance(tree1, tree2)
        with pytest.raises(_tskit.LibraryError):
            tree1.kc_distance(tree2)

        unsimplified_ts = msprime.simulate(
            10, random_seed=1, recombination_rate=10, record_full_arg=True
        )
        trees = unsimplified_ts.trees(sample_lists=True)
        tree1 = next(trees)
        tree2 = next(trees)
        with pytest.raises(ValueError):
            naive_kc_distance(tree1, tree2)
        with pytest.raises(ValueError):
            c_kc_distance(tree1, tree2)
        with pytest.raises(_tskit.LibraryError):
            tree1.kc_distance(tree2)

    def validate_trees(self, n):
        for seed in range(1, 10):
            ts1 = msprime.simulate(n, random_seed=seed)
            ts2 = msprime.simulate(n, random_seed=seed + 1)
            tree1 = next(ts1.trees(sample_lists=True))
            tree2 = next(ts2.trees(sample_lists=True))
            kc1 = naive_kc_distance(tree1, tree2)
            kc2 = c_kc_distance(tree1, tree2)
            kc3 = tree1.kc_distance(tree2)
            self.assertAlmostEqual(kc1, kc2)
            self.assertAlmostEqual(kc1, kc3)
            self.assertAlmostEqual(kc1, naive_kc_distance(tree2, tree1))
            self.assertAlmostEqual(kc2, c_kc_distance(tree2, tree1))
            self.assertAlmostEqual(kc3, tree2.kc_distance(tree1))

    def test_sample_3(self):
        self.validate_trees(3)

    def test_sample_4(self):
        self.validate_trees(4)

    def test_sample_10(self):
        self.validate_trees(10)

    def test_sample_20(self):
        self.validate_trees(20)

    def validate_nonbinary_trees(self, n):
        demographic_events = [
            msprime.SimpleBottleneck(0.02, 0, proportion=0.25),
            msprime.SimpleBottleneck(0.2, 0, proportion=1),
        ]

        for seed in range(1, 10):
            ts = msprime.simulate(
                n, random_seed=seed, demographic_events=demographic_events
            )
            # Check if this is really nonbinary
            found = False
            for edgeset in ts.edgesets():
                if len(edgeset.children) > 2:
                    found = True
                    break
            assert found
            tree1 = next(ts.trees(sample_lists=True))

            ts = msprime.simulate(
                n, random_seed=seed + 1, demographic_events=demographic_events
            )
            tree2 = next(ts.trees(sample_lists=True))
            self.do_kc_distance(tree1, tree2)
            # compare to a binary tree also

            ts = msprime.simulate(n, random_seed=seed + 1)
            tree2 = next(ts.trees(sample_lists=True))
            self.do_kc_distance(tree1, tree2)

    def test_non_binary_sample_10(self):
        self.validate_nonbinary_trees(10)

    def test_non_binary_sample_20(self):
        self.validate_nonbinary_trees(20)

    def test_non_binary_sample_30(self):
        self.validate_nonbinary_trees(30)

    def verify_result(self, tree1, tree2, lambda_, result, places=None):
        kc1 = naive_kc_distance(tree1, tree2, lambda_)
        kc2 = c_kc_distance(tree1, tree2, lambda_)
        kc3 = tree1.kc_distance(tree2, lambda_)
        self.assertAlmostEqual(kc1, result, places=places)
        self.assertAlmostEqual(kc2, result, places=places)
        self.assertAlmostEqual(kc3, result, places=places)

        kc1 = naive_kc_distance(tree2, tree1, lambda_)
        kc2 = c_kc_distance(tree2, tree1, lambda_)
        kc3 = tree2.kc_distance(tree1, lambda_)
        self.assertAlmostEqual(kc1, result, places=places)
        self.assertAlmostEqual(kc2, result, places=places)
        self.assertAlmostEqual(kc3, result, places=places)

    def test_known_kc_sample_3(self):
        # Test with hardcoded known values
        tables_1 = tskit.TableCollection(sequence_length=1.0)
        tables_2 = tskit.TableCollection(sequence_length=1.0)

        # Nodes
        sv = [True, True, True, False, False]
        tv_1 = [0.0, 0.0, 0.0, 2.0, 3.0]
        tv_2 = [0.0, 0.0, 0.0, 4.0, 6.0]

        for is_sample, t1, t2 in zip(sv, tv_1, tv_2):
            flags = tskit.NODE_IS_SAMPLE if is_sample else 0
            tables_1.nodes.add_row(flags=flags, time=t1)
            tables_2.nodes.add_row(flags=flags, time=t2)

        # Edges
        lv = [0.0, 0.0, 0.0, 0.0]
        rv = [1.0, 1.0, 1.0, 1.0]
        pv = [3, 3, 4, 4]
        cv = [0, 1, 2, 3]

        for left, right, p, c in zip(lv, rv, pv, cv):
            tables_1.edges.add_row(left=left, right=right, parent=p, child=c)
            tables_2.edges.add_row(left=left, right=right, parent=p, child=c)

        tree_1 = next(tables_1.tree_sequence().trees(sample_lists=True))
        tree_2 = next(tables_2.tree_sequence().trees(sample_lists=True))
        self.verify_result(tree_1, tree_2, 0, 0)
        self.verify_result(tree_1, tree_2, 1, 4.243, places=3)

    def test_10_samples(self):
        nodes_1 = io.StringIO(
            """\
        id  is_sample   time    population  individual  metadata
        0   1   0.000000    0   -1  b''
        1   1   0.000000    0   -1  b''
        2   1   0.000000    0   -1  b''
        3   1   0.000000    0   -1  b''
        4   1   0.000000    0   -1  b''
        5   1   0.000000    0   -1  b''
        6   1   0.000000    0   -1  b''
        7   1   0.000000    0   -1  b''
        8   1   0.000000    0   -1  b''
        9   1   0.000000    0   -1  b''
        10  0   0.047734    0   -1  b''
        11  0   0.061603    0   -1  b''
        12  0   0.189503    0   -1  b''
        13  0   0.275885    0   -1  b''
        14  0   0.518301    0   -1  b''
        15  0   0.543143    0   -1  b''
        16  0   0.865193    0   -1  b''
        17  0   1.643658    0   -1  b''
        18  0   2.942350    0   -1  b''
        """
        )
        edges_1 = io.StringIO(
            """\
        left    right   parent  child
        0.000000    10000.000000    10  0
        0.000000    10000.000000    10  2
        0.000000    10000.000000    11  9
        0.000000    10000.000000    11  10
        0.000000    10000.000000    12  3
        0.000000    10000.000000    12  7
        0.000000    10000.000000    13  5
        0.000000    10000.000000    13  11
        0.000000    10000.000000    14  1
        0.000000    10000.000000    14  8
        0.000000    10000.000000    15  4
        0.000000    10000.000000    15  14
        0.000000    10000.000000    16  13
        0.000000    10000.000000    16  15
        0.000000    10000.000000    17  6
        0.000000    10000.000000    17  12
        0.000000    10000.000000    18  16
        0.000000    10000.000000    18  17
        """
        )
        ts_1 = tskit.load_text(
            nodes_1, edges_1, sequence_length=10000, strict=False, base64_metadata=False
        )
        nodes_2 = io.StringIO(
            """\
        id  is_sample   time    population  individual  metadata
        0   1   0.000000    0   -1  b''
        1   1   0.000000    0   -1  b''
        2   1   0.000000    0   -1  b''
        3   1   0.000000    0   -1  b''
        4   1   0.000000    0   -1  b''
        5   1   0.000000    0   -1  b''
        6   1   0.000000    0   -1  b''
        7   1   0.000000    0   -1  b''
        8   1   0.000000    0   -1  b''
        9   1   0.000000    0   -1  b''
        10  0   0.210194    0   -1  b''
        11  0   0.212217    0   -1  b''
        12  0   0.223341    0   -1  b''
        13  0   0.272703    0   -1  b''
        14  0   0.443553    0   -1  b''
        15  0   0.491653    0   -1  b''
        16  0   0.729369    0   -1  b''
        17  0   1.604113    0   -1  b''
        18  0   1.896332    0   -1  b''
        """
        )
        edges_2 = io.StringIO(
            """\
        left    right   parent  child
        0.000000    10000.000000    10  5
        0.000000    10000.000000    10  7
        0.000000    10000.000000    11  3
        0.000000    10000.000000    11  4
        0.000000    10000.000000    12  6
        0.000000    10000.000000    12  9
        0.000000    10000.000000    13  10
        0.000000    10000.000000    13  12
        0.000000    10000.000000    14  8
        0.000000    10000.000000    14  11
        0.000000    10000.000000    15  1
        0.000000    10000.000000    15  2
        0.000000    10000.000000    16  13
        0.000000    10000.000000    16  14
        0.000000    10000.000000    17  0
        0.000000    10000.000000    17  16
        0.000000    10000.000000    18  15
        0.000000    10000.000000    18  17
        """
        )
        ts_2 = tskit.load_text(
            nodes_2, edges_2, sequence_length=10000, strict=False, base64_metadata=False
        )

        tree_1 = next(ts_1.trees(sample_lists=True))
        tree_2 = next(ts_2.trees(sample_lists=True))
        self.verify_result(tree_1, tree_2, 0, 12.85, places=2)
        self.verify_result(tree_1, tree_2, 1, 10.64, places=2)

    def test_15_samples(self):
        nodes_1 = io.StringIO(
            """\
        id  is_sample   time    population  individual  metadata
        0   1   0.000000    0   -1
        1   1   0.000000    0   -1
        2   1   0.000000    0   -1
        3   1   0.000000    0   -1
        4   1   0.000000    0   -1
        5   1   0.000000    0   -1
        6   1   0.000000    0   -1
        7   1   0.000000    0   -1
        8   1   0.000000    0   -1
        9   1   0.000000    0   -1
        10  1   0.000000    0   -1
        11  1   0.000000    0   -1
        12  1   0.000000    0   -1
        13  1   0.000000    0   -1
        14  1   0.000000    0   -1
        15  0   0.026043    0   -1
        16  0   0.032662    0   -1
        17  0   0.072032    0   -1
        18  0   0.086792    0   -1
        19  0   0.130699    0   -1
        20  0   0.177640    0   -1
        21  0   0.199800    0   -1
        22  0   0.236391    0   -1
        23  0   0.342445    0   -1
        24  0   0.380356    0   -1
        25  0   0.438502    0   -1
        26  0   0.525632    0   -1
        27  0   1.180078    0   -1
        28  0   2.548099    0   -1
        """
        )
        edges_1 = io.StringIO(
            """\
        left    right   parent  child
        0.000000    10000.000000    15  6
        0.000000    10000.000000    15  13
        0.000000    10000.000000    16  1
        0.000000    10000.000000    16  4
        0.000000    10000.000000    17  0
        0.000000    10000.000000    17  7
        0.000000    10000.000000    18  2
        0.000000    10000.000000    18  17
        0.000000    10000.000000    19  5
        0.000000    10000.000000    19  9
        0.000000    10000.000000    20  12
        0.000000    10000.000000    20  15
        0.000000    10000.000000    21  8
        0.000000    10000.000000    21  20
        0.000000    10000.000000    22  11
        0.000000    10000.000000    22  21
        0.000000    10000.000000    23  10
        0.000000    10000.000000    23  22
        0.000000    10000.000000    24  14
        0.000000    10000.000000    24  16
        0.000000    10000.000000    25  18
        0.000000    10000.000000    25  19
        0.000000    10000.000000    26  23
        0.000000    10000.000000    26  24
        0.000000    10000.000000    27  25
        0.000000    10000.000000    27  26
        0.000000    10000.000000    28  3
        0.000000    10000.000000    28  27
        """
        )
        ts_1 = tskit.load_text(
            nodes_1, edges_1, sequence_length=10000, strict=False, base64_metadata=False
        )

        nodes_2 = io.StringIO(
            """\
        id  is_sample   time    population  individual  metadata
        0   1   0.000000    0   -1
        1   1   0.000000    0   -1
        2   1   0.000000    0   -1
        3   1   0.000000    0   -1
        4   1   0.000000    0   -1
        5   1   0.000000    0   -1
        6   1   0.000000    0   -1
        7   1   0.000000    0   -1
        8   1   0.000000    0   -1
        9   1   0.000000    0   -1
        10  1   0.000000    0   -1
        11  1   0.000000    0   -1
        12  1   0.000000    0   -1
        13  1   0.000000    0   -1
        14  1   0.000000    0   -1
        15  0   0.011443    0   -1
        16  0   0.055694    0   -1
        17  0   0.061677    0   -1
        18  0   0.063416    0   -1
        19  0   0.163014    0   -1
        20  0   0.223445    0   -1
        21  0   0.251724    0   -1
        22  0   0.268749    0   -1
        23  0   0.352039    0   -1
        24  0   0.356134    0   -1
        25  0   0.399454    0   -1
        26  0   0.409174    0   -1
        27  0   2.090839    0   -1
        28  0   3.772716    0   -1
        """
        )
        edges_2 = io.StringIO(
            """\
        left    right   parent  child
        0.000000    10000.000000    15  6
        0.000000    10000.000000    15  8
        0.000000    10000.000000    16  9
        0.000000    10000.000000    16  12
        0.000000    10000.000000    17  3
        0.000000    10000.000000    17  4
        0.000000    10000.000000    18  13
        0.000000    10000.000000    18  16
        0.000000    10000.000000    19  2
        0.000000    10000.000000    19  11
        0.000000    10000.000000    20  1
        0.000000    10000.000000    20  17
        0.000000    10000.000000    21  0
        0.000000    10000.000000    21  18
        0.000000    10000.000000    22  10
        0.000000    10000.000000    22  15
        0.000000    10000.000000    23  14
        0.000000    10000.000000    23  21
        0.000000    10000.000000    24  5
        0.000000    10000.000000    24  7
        0.000000    10000.000000    25  19
        0.000000    10000.000000    25  22
        0.000000    10000.000000    26  24
        0.000000    10000.000000    26  25
        0.000000    10000.000000    27  20
        0.000000    10000.000000    27  23
        0.000000    10000.000000    28  26
        0.000000    10000.000000    28  27
        """
        )
        ts_2 = tskit.load_text(
            nodes_2, edges_2, sequence_length=10000, strict=False, base64_metadata=False
        )

        tree_1 = next(ts_1.trees(sample_lists=True))
        tree_2 = next(ts_2.trees(sample_lists=True))

        self.verify_result(tree_1, tree_2, 0, 19.95, places=2)
        self.verify_result(tree_1, tree_2, 1, 17.74, places=2)

    def test_nobinary_trees(self):
        nodes_1 = io.StringIO(
            """\
        id  is_sample   time    population  individual  metadata
        0   1   0.000000    -1  -1   e30=
        1   1   0.000000    -1  -1   e30=
        2   1   0.000000    -1  -1   e30=
        3   1   0.000000    -1  -1   e30=
        4   1   0.000000    -1  -1   e30=
        5   1   0.000000    -1  -1   e30=
        6   1   0.000000    -1  -1   e30=
        7   1   0.000000    -1  -1   e30=
        8   1   0.000000    -1  -1   e30=
        9   1   0.000000    -1  -1
        10  1   0.000000    -1  -1
        11  1   0.000000    -1  -1
        12  1   0.000000    -1  -1
        13  1   0.000000    -1  -1
        14  1   0.000000    -1  -1
        15  0   2.000000    -1  -1
        16  0   4.000000    -1  -1
        17  0   11.000000   -1  -1
        18  0   12.000000   -1  -1
        """
        )
        edges_1 = io.StringIO(
            """\
        left    right   parent  child
        0.000000    10000.000000    15  8
        0.000000    10000.000000    15  10
        0.000000    10000.000000    16  6
        0.000000    10000.000000    16  12
        0.000000    10000.000000    16  15
        0.000000    10000.000000    17  0
        0.000000    10000.000000    17  1
        0.000000    10000.000000    17  2
        0.000000    10000.000000    17  3
        0.000000    10000.000000    17  4
        0.000000    10000.000000    17  5
        0.000000    10000.000000    17  7
        0.000000    10000.000000    17  9
        0.000000    10000.000000    17  11
        0.000000    10000.000000    17  13
        0.000000    10000.000000    17  14
        0.000000    10000.000000    18  16
        0.000000    10000.000000    18  17
        """
        )
        ts_1 = tskit.load_text(
            nodes_1, edges_1, sequence_length=10000, strict=False, base64_metadata=False
        )

        nodes_2 = io.StringIO(
            """\
        id  is_sample   time    population  individual  metadata
        0   1   0.000000    -1  -1   e30=
        1   1   0.000000    -1  -1   e30=
        2   1   0.000000    -1  -1   e30=
        3   1   0.000000    -1  -1   e30=
        4   1   0.000000    -1  -1   e30=
        5   1   0.000000    -1  -1   e30=
        6   1   0.000000    -1  -1   e30=
        7   1   0.000000    -1  -1   e30=
        8   1   0.000000    -1  -1   e30=
        9   1   0.000000    -1  -1   e30=
        10  1   0.000000    -1  -1  e30=
        11  1   0.000000    -1  -1  e30=
        12  1   0.000000    -1  -1  e30=
        13  1   0.000000    -1  -1  e30=
        14  1   0.000000    -1  -1  e30=
        15  0   2.000000    -1  -1
        16  0   2.000000    -1  -1
        17  0   3.000000    -1  -1
        18  0   3.000000    -1  -1
        19  0   4.000000    -1  -1
        20  0   4.000000    -1  -1
        21  0   11.000000   -1  -1
        22  0   12.000000   -1  -1
        """
        )
        edges_2 = io.StringIO(
            """\
        left    right   parent  child
        0.000000    10000.000000    15  12
        0.000000    10000.000000    15  14
        0.000000    10000.000000    16  0
        0.000000    10000.000000    16  7
        0.000000    10000.000000    17  6
        0.000000    10000.000000    17  15
        0.000000    10000.000000    18  4
        0.000000    10000.000000    18  8
        0.000000    10000.000000    18  13
        0.000000    10000.000000    19  11
        0.000000    10000.000000    19  18
        0.000000    10000.000000    20  1
        0.000000    10000.000000    20  5
        0.000000    10000.000000    20  9
        0.000000    10000.000000    20  10
        0.000000    10000.000000    21  2
        0.000000    10000.000000    21  3
        0.000000    10000.000000    21  16
        0.000000    10000.000000    21  17
        0.000000    10000.000000    21  20
        0.000000    10000.000000    22  19
        0.000000    10000.000000    22  21
        """
        )
        ts_2 = tskit.load_text(
            nodes_2, edges_2, sequence_length=10000, strict=False, base64_metadata=False
        )
        tree_1 = next(ts_1.trees(sample_lists=True))
        tree_2 = next(ts_2.trees(sample_lists=True))
        self.verify_result(tree_1, tree_2, 0, 9.434, places=3)
        self.verify_result(tree_1, tree_2, 1, 44, places=1)

    def test_multiple_roots(self):
        tables = tskit.TableCollection(sequence_length=1.0)

        # Nodes
        sv = [True, True]
        tv = [0.0, 0.0]

        for is_sample, t in zip(sv, tv):
            flags = tskit.NODE_IS_SAMPLE if is_sample else 0
            tables.nodes.add_row(flags=flags, time=t)

        ts = tables.tree_sequence()

        with pytest.raises(ValueError):
            naive_kc_distance(ts.first(), ts.first(), 0)
        with pytest.raises(ValueError):
            c_kc_distance(ts.first(), ts.first(), 0)
        with pytest.raises(_tskit.LibraryError):
            ts.first().kc_distance(ts.first(), 0)

    def do_kc_distance(self, t1, t2, lambda_=0):
        kc1 = naive_kc_distance(t1, t2, lambda_)
        kc2 = c_kc_distance(t1, t2, lambda_)
        kc3 = t1.kc_distance(t2, lambda_)
        self.assertAlmostEqual(kc1, kc2)
        self.assertAlmostEqual(kc1, kc3)

        kc1 = naive_kc_distance(t2, t1, lambda_)
        kc2 = c_kc_distance(t2, t1, lambda_)
        kc3 = t2.kc_distance(t1, lambda_)
        self.assertAlmostEqual(kc1, kc2)
        self.assertAlmostEqual(kc1, kc3)

    def test_non_initial_samples(self):
        ts1 = msprime.simulate(10, random_seed=1)
        nmap = np.arange(0, ts1.num_nodes)[::-1]
        ts2 = tsutil.permute_nodes(ts1, nmap)
        t1 = next(ts2.trees(sample_lists=True))
        t2 = next(ts2.trees(sample_lists=True))
        self.do_kc_distance(t1, t2)

    def test_internal_samples(self):
        ts1 = msprime.simulate(10, random_seed=1)
        ts2 = tsutil.jiggle_samples(ts1)
        t1 = next(ts2.trees(sample_lists=True))
        t2 = next(ts2.trees(sample_lists=True))

        naive_kc_distance(t1, t2)
        c_kc_distance(t1, t2)
        t1.kc_distance(t2)

    def test_root_sample(self):
        tables1 = tskit.TableCollection(sequence_length=1.0)
        tables1.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        only_root = next(tables1.tree_sequence().trees(sample_lists=True))
        assert only_root.kc_distance(only_root) == 0
        assert only_root.kc_distance(only_root, lambda_=1) == 0

    def test_non_sample_leaf(self):
        tables = tskit.TableCollection(sequence_length=1.0)
        c1 = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        c2 = tables.nodes.add_row(time=0)
        p = tables.nodes.add_row(time=1)
        tables.edges.add_row(left=0, right=1, parent=p, child=c1)
        tables.edges.add_row(left=0, right=1, parent=p, child=c2)
        ts = tables.tree_sequence()
        tree = next(ts.trees(sample_lists=True))
        assert ts.kc_distance(ts) == 0
        assert tree.kc_distance(tree) == 0

        # mirrored
        tables = tskit.TableCollection(sequence_length=1.0)
        c1 = tables.nodes.add_row(time=0)
        c2 = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        p = tables.nodes.add_row(time=1)
        tables.edges.add_row(left=0, right=1, parent=p, child=c1)
        tables.edges.add_row(left=0, right=1, parent=p, child=c2)
        ts = tables.tree_sequence()
        tree = next(ts.trees(sample_lists=True))
        assert ts.kc_distance(ts) == 0
        assert tree.kc_distance(tree) == 0

    def test_ignores_subtrees_with_no_samples(self):
        nodes_1 = io.StringIO(
            """\
        id  is_sample   time    population  individual  metadata
        0   0   0.000000    0   -1
        1   0   0.000000    0   -1
        2   0   0.000000    0   -1
        3   1   0.000000    0   -1
        4   0   0.000000    0   -1
        5   0   0.000000    0   -1
        6   1   1.000000    0   -1
        7   1   2.000000    0   -1
        8   0   2.000000    0   -1
        9   0   3.000000    0   -1
        """
        )
        edges_1 = io.StringIO(
            """\
        left    right   parent  child
        0.000000    1.000000    6  0
        0.000000    1.000000    6  1
        0.000000    1.000000    7  2
        0.000000    1.000000    7  6
        0.000000    1.000000    8  4
        0.000000    1.000000    8  5
        0.000000    1.000000    9  3
        0.000000    1.000000    9  7
        0.000000    1.000000    9  8
        """
        )
        redundant = tskit.load_text(
            nodes_1, edges_1, sequence_length=1, strict=False, base64_metadata=False
        )

        nodes_2 = io.StringIO(
            """\
        id  is_sample   time    population  individual  metadata
        0   0   0.000000    0   -1
        1   0   0.000000    0   -1
        2   0   0.000000    0   -1
        3   1   0.000000    0   -1
        4   0   0.000000    0   -1
        5   0   0.000000    0   -1
        6   1   1.000000    0   -1
        7   1   2.000000    0   -1
        8   0   2.000000    0   -1
        9   0   3.000000    0   -1
        """
        )
        edges_2 = io.StringIO(
            """\
        left    right   parent  child
        0.000000    1.000000    7  2
        0.000000    1.000000    7  6
        0.000000    1.000000    9  3
        0.000000    1.000000    9  7
        """
        )
        simplified = tskit.load_text(
            nodes_2, edges_2, sequence_length=1, strict=False, base64_metadata=False
        )
        assert redundant.kc_distance(simplified, 0) == 0
        assert redundant.kc_distance(simplified, 1) == 0


def ts_kc_distance(ts1, ts2, lambda_=0):
    check_kc_tree_sequence_inputs(ts1, ts2)

    total = 0
    left = 0
    tree1_iter = ts1.trees(sample_lists=True)
    tree1 = next(tree1_iter)
    for tree2 in ts2.trees(sample_lists=True):
        while tree1.interval.right < tree2.interval.right:
            span = tree1.interval.right - left
            total += tree1.kc_distance(tree2, lambda_) * span

            left = tree1.interval.right
            tree1 = next(tree1_iter)
        span = tree2.interval.right - left
        left = tree2.interval.right
        total += tree1.kc_distance(tree2, lambda_) * span

    return total / ts1.sequence_length


def ts_kc_distance_incremental(ts1, ts2, lambda_=0):
    check_kc_tree_sequence_inputs(ts1, ts2)

    sample_maps = [dict(), dict()]
    for i, ts in enumerate([ts1, ts2]):
        for j, u in enumerate(ts.samples()):
            sample_maps[i][u] = j

    total = 0
    left = 0

    t1_vecs = KCVectors(ts1.num_samples)
    t2_vecs = KCVectors(ts2.num_samples)

    t1_depths = np.zeros(ts1.num_nodes)
    t2_depths = np.zeros(ts2.num_nodes)

    edge_diffs_iter_1 = ts1.edge_diffs()
    tree_iter_1 = ts1.trees(sample_lists=True)
    t1, t1_diffs = next(tree_iter_1), next(edge_diffs_iter_1)
    update_kc_incremental(t1, t1_vecs, t1_diffs, sample_maps[0], t1_depths)
    for t2, t2_diffs in zip(ts2.trees(sample_lists=True), ts2.edge_diffs()):
        update_kc_incremental(t2, t2_vecs, t2_diffs, sample_maps[1], t2_depths)
        while t1_diffs[0][1] < t2_diffs[0][1]:
            span = t1_diffs[0][1] - left
            total += norm_kc_vectors(t1_vecs, t2_vecs, lambda_) * span

            left = t1_diffs[0][1]
            t1, t1_diffs = next(tree_iter_1), next(edge_diffs_iter_1)
            update_kc_incremental(t1, t1_vecs, t1_diffs, sample_maps[0], t1_depths)
        span = t2_diffs[0][1] - left
        left = t2_diffs[0][1]
        total += norm_kc_vectors(t1_vecs, t2_vecs, lambda_) * span

    return total / ts1.sequence_length


# tree is the result of removing/inserting the edges in edge_diffs
def update_kc_incremental(tree, kc, edge_diffs, sample_index_map, depths):
    _, edges_out, edges_in = edge_diffs

    # Update state of detached subtrees.
    for e in reversed(edges_out):
        u = e.child
        depths[u] = 0

        # Only update detached subtrees that remain detached. Otherwise,
        # they must be reattached by an incoming edge and will be
        # updated below. We're looking into the future here by seeing
        # that u remains detached after all the incoming edges are
        # inserted into `tree`.
        if tree.parent(u) == tskit.NULL:
            update_kc_subtree_state(tree, kc, u, sample_index_map, depths)

    # Propagate state change down into reattached subtrees.
    for e in reversed(edges_in):
        u = e.child
        assert depths[u] == 0
        depths[u] = depths[e.parent] + 1
        update_kc_subtree_state(tree, kc, u, sample_index_map, depths)

        # The per-leaf elements of KC only change when the edge directly
        # above the leaf changes, so are handled separately from the
        # propagated state used for leaf-pair elements.
        if tree.is_leaf(u):
            time = tree.branch_length(u)
            update_kc_vectors_single_leaf(kc, u, time, sample_index_map)


def update_kc_subtree_state(tree, kc, u, sample_index_map, depths):
    """
    Update the depths of the nodes in this subtree. When a leaf is hit,
    update the KC vector elements associated with that leaf.
    """
    stack = [u]
    while len(stack) > 0:
        v = stack.pop()
        if tree.is_leaf(v):
            update_kc_pairs_with_leaf(tree, kc, v, sample_index_map, depths)
        else:
            c = tree.left_child(v)
            while c != -1:
                # Terminate iteration at nodes that are currently considered
                # roots by the edge diffs. Nodes with a depth of 0 are
                # temporary root nodes made by breaking an outgoing edge
                # that have yet to be inserted by a later incoming edge.
                if depths[c] != 0:
                    depths[c] = depths[v] + 1
                    stack.append(c)
                c = tree.right_sib(c)


def update_kc_pairs_with_leaf(tree, kc, leaf, sample_index_map, depths):
    """
    Perform an upward traversal from `leaf` to the root, updating the KC
    vector elements for pairs of `leaf` with every other leaf in the tree.
    """
    root_time = tree.time(tree.root)
    p = tree.parent(leaf)
    c = leaf
    while p != -1:
        time = root_time - tree.time(p)
        depth = depths[p]
        for sibling in tree.children(p):
            if sibling != c:
                update_kc_vectors_all_pairs(tree, kc, leaf, sibling, depth, time)
        c, p = p, tree.parent(p)


def check_kc_tree_sequence_inputs(ts1, ts2):
    if not np.array_equal(ts1.samples(), ts2.samples()):
        raise ValueError("Trees must have the same samples")
    if ts1.sequence_length != ts2.sequence_length:
        raise ValueError("Can't compare with sequences of different lengths")

    tree1_iter = ts1.trees(sample_lists=True)
    tree1 = next(tree1_iter)
    for tree2 in ts2.trees(sample_lists=True):
        while tree1.interval.right < tree2.interval.right:
            check_kc_tree_inputs(tree1, tree2)
            tree1 = next(tree1_iter)
        check_kc_tree_inputs(tree1, tree2)


def check_kc_tree_inputs(tree1, tree2):
    if not len(tree1.roots) == len(tree2.roots) == 1:
        raise ValueError("Trees must have one root")
    for tree in [tree1, tree2]:
        for u in tree.nodes():
            if tree.num_children(u) == 1:
                raise ValueError("Unary nodes are not supported")


class TestKCSequenceMetric(unittest.TestCase):
    """
    Tests the KC Metric on a tree sequence.
    """

    def test_0_distance_from_self(self):
        ts = msprime.simulate(10)
        assert ts_kc_distance(ts, ts) == 0

    def verify_errors(self, ts1, ts2):
        with pytest.raises(ValueError):
            ts_kc_distance(ts1, ts2)
        with pytest.raises(ValueError):
            ts_kc_distance_incremental(ts1, ts2)
        with pytest.raises(_tskit.LibraryError):
            ts1.kc_distance(ts2)

    def test_errors_diff_seq_length(self):
        ts1 = msprime.simulate(10, length=1)
        ts2 = msprime.simulate(10, length=2)
        self.verify_errors(ts1, ts2)

    def test_errors_diff_num_samples(self):
        ts1 = msprime.simulate(10, length=1)
        ts2 = msprime.simulate(12, length=2)
        self.verify_errors(ts1, ts2)

    def test_errors_different_sample_lists(self):
        tables_1 = tskit.TableCollection(sequence_length=2.0)
        tables_2 = tskit.TableCollection(sequence_length=2.0)

        sv1 = [True, True, True, False, False]
        tv1 = [0.0, 0.0, 0.0, 1.0, 2.0]
        sv2 = [True, True, False, False, True]
        tv2 = [0.0, 0.0, 1.0, 2.0, 0.0]
        for is_sample, t in zip(sv1, tv1):
            flags = tskit.NODE_IS_SAMPLE if is_sample else 0
            tables_1.nodes.add_row(flags=flags, time=t)
        for is_sample, t in zip(sv2, tv2):
            flags = tskit.NODE_IS_SAMPLE if is_sample else 0
            tables_2.nodes.add_row(flags=flags, time=t)

        lv = [0.0, 0.0, 0.0, 0.0]
        rv = [1.0, 1.0, 1.0, 1.0]
        pv1 = [3, 3, 4, 4]
        cv1 = [0, 1, 2, 3]
        for left, right, p, c in zip(lv, rv, pv1, cv1):
            tables_1.edges.add_row(left=left, right=right, parent=p, child=c)

        pv2 = [2, 2, 3, 3]
        cv2 = [0, 1, 2, 4]
        for left, right, p, c in zip(lv, rv, pv2, cv2):
            tables_2.edges.add_row(left=left, right=right, parent=p, child=c)

        ts1 = tables_1.tree_sequence()
        ts2 = tables_2.tree_sequence()
        self.verify_errors(ts1, ts2)

        unsimplified_ts = msprime.simulate(
            10, random_seed=1, recombination_rate=10, record_full_arg=True
        )
        self.verify_errors(unsimplified_ts, unsimplified_ts)

    def test_errors_unary_nodes(self):
        tables = tskit.TableCollection(sequence_length=2.0)

        sv = [True, False, False]
        tv = [0.0, 1.0, 2.0]
        for is_sample, t in zip(sv, tv):
            flags = tskit.NODE_IS_SAMPLE if is_sample else 0
            tables.nodes.add_row(flags=flags, time=t)

        lv = [0.0, 0.0, 0.0]
        rv = [1.0, 1.0, 1.0]
        pv = [1, 2]
        cv = [0, 1]
        for left, right, p, c in zip(lv, rv, pv, cv):
            tables.edges.add_row(left=left, right=right, parent=p, child=c)

        ts = tables.tree_sequence()
        self.verify_errors(ts, ts)

    def test_errors_different_samples(self):
        ts1 = msprime.simulate(10, random_seed=1)
        ts2 = tsutil.jiggle_samples(ts1)
        self.verify_errors(ts1, ts2)

    def verify_result(self, ts1, ts2, lambda_, result, places=None):
        kc1 = ts_kc_distance(ts1, ts2, lambda_)
        kc2 = ts_kc_distance_incremental(ts1, ts2, lambda_)
        kc3 = ts1.kc_distance(ts2, lambda_)
        self.assertAlmostEqual(kc1, result, places=places)
        self.assertAlmostEqual(kc2, result, places=places)
        self.assertAlmostEqual(kc3, result, places=places)

        kc1 = ts_kc_distance(ts2, ts1, lambda_)
        kc2 = ts_kc_distance_incremental(ts2, ts1, lambda_)
        kc3 = ts2.kc_distance(ts1, lambda_)
        self.assertAlmostEqual(kc1, result, places=places)
        self.assertAlmostEqual(kc2, result, places=places)
        self.assertAlmostEqual(kc3, result, places=places)

    def verify_same_kc(self, ts1, ts2, lambda_=0):
        kc1 = ts_kc_distance(ts1, ts2, lambda_)
        kc2 = ts_kc_distance_incremental(ts1, ts2, lambda_)
        kc3 = ts1.kc_distance(ts2, lambda_)
        self.assertAlmostEqual(kc1, kc2)
        self.assertAlmostEqual(kc2, kc3)

        kc1 = ts_kc_distance(ts2, ts1, lambda_)
        kc2 = ts_kc_distance_incremental(ts2, ts1, lambda_)
        kc3 = ts2.kc_distance(ts1, lambda_)
        self.assertAlmostEqual(kc1, kc2)
        self.assertAlmostEqual(kc2, kc3)

    def validate_trees(self, n):
        for seed in range(1, 10):
            ts1 = msprime.simulate(n, random_seed=seed, recombination_rate=1)
            ts2 = msprime.simulate(n, random_seed=seed + 1, recombination_rate=1)
            self.verify_same_kc(ts2, ts1)
            self.verify_same_kc(ts1, ts2)
            self.verify_same_kc(ts1, ts1)  # Test sequences with equal breakpoints

    def test_sample_5(self):
        self.validate_trees(5)

    def test_sample_10(self):
        self.validate_trees(10)

    def test_sample_20(self):
        self.validate_trees(20)

    def validate_nonbinary_trees(self, n):
        demographic_events = [
            msprime.SimpleBottleneck(0.02, 0, proportion=0.25),
            msprime.SimpleBottleneck(0.2, 0, proportion=1),
        ]

        for seed in range(1, 10):
            ts1 = msprime.simulate(
                n,
                random_seed=seed,
                demographic_events=demographic_events,
                recombination_rate=1,
            )
            # Check if this is really nonbinary
            found = False
            for edgeset in ts1.edgesets():
                if len(edgeset.children) > 2:
                    found = True
                    break
            assert found

            ts2 = msprime.simulate(
                n,
                random_seed=seed + 1,
                demographic_events=demographic_events,
                recombination_rate=1,
            )
            self.verify_same_kc(ts1, ts2)

            # compare to a binary tree also
            ts2 = msprime.simulate(n, recombination_rate=1, random_seed=seed + 1)
            self.verify_same_kc(ts1, ts2)

    def test_non_binary_sample_10(self):
        self.validate_nonbinary_trees(10)

    def test_non_binary_sample_20(self):
        self.validate_nonbinary_trees(20)

    def test_permit_internal_samples(self):
        tables = tskit.TableCollection(1.0)
        tables.nodes.add_row(flags=1)
        tables.nodes.add_row(flags=1)
        tables.nodes.add_row(flags=1, time=1)
        tables.edges.add_row(0, 1, 2, 0)
        tables.edges.add_row(0, 1, 2, 1)
        ts = tables.tree_sequence()
        assert ts.kc_distance(ts) == 0
        assert ts_kc_distance_incremental(ts, ts) == 0

    def test_known_kc_sample_trees_different_shapes(self):
        tables_1 = tskit.TableCollection(sequence_length=2.0)
        tables_2 = tskit.TableCollection(sequence_length=2.0)

        # Nodes
        sv = [True, True, True, True, False, False, False]
        tv = [0.0, 0.0, 0.0, 0.0, 1.0, 2.0, 3.0]
        for is_sample, t in zip(sv, tv):
            flags = tskit.NODE_IS_SAMPLE if is_sample else 0
            tables_1.nodes.add_row(flags=flags, time=t)
            tables_2.nodes.add_row(flags=flags, time=t)

        # First tree edges
        pv1 = [4, 4, 5, 5, 6, 6, 5, 6]
        cv1 = [2, 3, 1, 4, 0, 5, 0, 4]
        lv1 = [0, 0, 0, 0, 0, 0, 1, 1]
        rv1 = [2, 2, 2, 1, 1, 2, 2, 2]

        # Second tree edges
        pv2 = [4, 4, 5, 5, 6, 6, 5, 6]
        cv2 = [2, 3, 0, 1, 4, 5, 4, 0]
        lv2 = [0, 0, 0, 0, 0, 0, 1, 1]
        rv2 = [2, 2, 1, 2, 1, 2, 2, 2]

        for left, right, p, c in zip(lv1, rv1, pv1, cv1):
            tables_1.edges.add_row(left=left, right=right, parent=p, child=c)
        for left, right, p, c in zip(lv2, rv2, pv2, cv2):
            tables_2.edges.add_row(left=left, right=right, parent=p, child=c)

        tables_1.sort()
        tables_2.sort()
        ts_1 = tables_1.tree_sequence()
        ts_2 = tables_2.tree_sequence()
        self.verify_result(ts_1, ts_2, 0, 2.0)

    def test_known_kc_sample_trees_same_shape_different_times(self):
        tables_1 = tskit.TableCollection(sequence_length=1.0)
        tables_2 = tskit.TableCollection(sequence_length=1.0)

        # Nodes
        sv = [True, True, True, False, False]
        tv_1 = [0.0, 0.0, 0.0, 2.0, 3.0]
        tv_2 = [0.0, 0.0, 0.0, 4.0, 6.0]

        for is_sample, t1, t2 in zip(sv, tv_1, tv_2):
            flags = tskit.NODE_IS_SAMPLE if is_sample else 0
            tables_1.nodes.add_row(flags=flags, time=t1)
            tables_2.nodes.add_row(flags=flags, time=t2)

        # Edges
        lv = [0.0, 0.0, 0.0, 0.0]
        rv = [1.0, 1.0, 1.0, 1.0]
        pv = [3, 3, 4, 4]
        cv = [0, 1, 2, 3]

        for left, right, p, c in zip(lv, rv, pv, cv):
            tables_1.edges.add_row(left=left, right=right, parent=p, child=c)
            tables_2.edges.add_row(left=left, right=right, parent=p, child=c)

        ts_1 = tables_1.tree_sequence()
        ts_2 = tables_2.tree_sequence()

        self.verify_result(ts_1, ts_2, 0, 0)
        self.verify_result(ts_1, ts_2, 1, 4.243, places=3)

    def test_known_kc_same_tree_twice_same_metric(self):
        tables_1 = tskit.TableCollection(sequence_length=2.0)
        tables_2 = tskit.TableCollection(sequence_length=2.0)

        # Nodes
        sv = [True, True, True, False, False]
        tv_1 = [0.0, 0.0, 0.0, 2.0, 3.0]
        tv_2 = [0.0, 0.0, 0.0, 4.0, 6.0]

        for is_sample, t1, t2 in zip(sv, tv_1, tv_2):
            flags = tskit.NODE_IS_SAMPLE if is_sample else 0
            tables_1.nodes.add_row(flags=flags, time=t1)
            tables_2.nodes.add_row(flags=flags, time=t2)

        # Edges
        pv = [3, 3, 4, 4]
        cv = [0, 1, 2, 3]

        for p, c in zip(pv, cv):
            tables_1.edges.add_row(left=0, right=1, parent=p, child=c)
            tables_1.edges.add_row(left=1, right=2, parent=p, child=c)
            tables_2.edges.add_row(left=0, right=0.5, parent=p, child=c)
            tables_2.edges.add_row(left=0.5, right=2, parent=p, child=c)

        ts_1 = tables_1.tree_sequence()
        ts_2 = tables_2.tree_sequence()
        self.verify_result(ts_1, ts_2, 0, 0)
        self.verify_result(ts_1, ts_2, 1, 4.243, places=3)

    def test_remove_root(self):
        tables_1 = tskit.TableCollection(sequence_length=10.0)
        tables_2 = tskit.TableCollection(sequence_length=10.0)

        # Nodes
        sv1 = [True, True, True, True, True, False, False, False, False, False]
        tv1 = [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0]

        sv2 = [True, True, True, True, True, False, False, False, False]
        tv2 = [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 2.0, 3.0, 4.0]

        for is_sample, t in zip(sv1, tv1):
            flags = tskit.NODE_IS_SAMPLE if is_sample else 0
            tables_1.nodes.add_row(flags=flags, time=t)
        for is_sample, t in zip(sv2, tv2):
            flags = tskit.NODE_IS_SAMPLE if is_sample else 0
            tables_2.nodes.add_row(flags=flags, time=t)

        # Edges
        pv1 = [5, 5, 6, 6, 7, 7, 8, 8, 8, 9, 9]
        cv1 = [0, 1, 3, 4, 2, 5, 2, 6, 7, 5, 8]
        lv1 = [0, 0, 0, 0, 5, 5, 0, 0, 5, 0, 0]
        rv1 = [10, 10, 10, 10, 10, 10, 5, 10, 10, 5, 5]

        pv2 = [5, 5, 6, 6, 7, 7, 8, 8]
        cv2 = [0, 1, 2, 3, 4, 5, 6, 7]
        lv2 = [0, 0, 0, 0, 0, 0, 0, 0]
        rv2 = [10, 10, 10, 10, 10, 10, 10, 10]

        for p, c, l, r in zip(pv1, cv1, lv1, rv1):
            tables_1.edges.add_row(left=l, right=r, parent=p, child=c)

        for p, c, l, r in zip(pv2, cv2, lv2, rv2):
            tables_2.edges.add_row(left=l, right=r, parent=p, child=c)

        ts_1 = tables_1.tree_sequence()
        ts_2 = tables_2.tree_sequence()
        distance = (math.sqrt(8) * 5 + math.sqrt(6) * 5) / 10
        self.verify_result(ts_1, ts_2, 0, distance)

    def test_ignores_subtrees_with_no_samples(self):
        nodes_1 = io.StringIO(
            """\
        id  is_sample   time    population  individual  metadata
        0   0   0.000000    0   -1
        1   0   0.000000    0   -1
        2   0   0.000000    0   -1
        3   1   0.000000    0   -1
        4   0   0.000000    0   -1
        5   0   0.000000    0   -1
        6   1   1.000000    0   -1
        7   1   2.000000    0   -1
        8   0   2.000000    0   -1
        9   0   3.000000    0   -1
        """
        )
        edges_1 = io.StringIO(
            """\
        left    right   parent  child
        0.000000    1.000000    6  0
        0.000000    1.000000    6  1
        0.000000    1.000000    7  2
        0.000000    1.000000    7  6
        0.000000    1.000000    8  4
        0.000000    1.000000    8  5
        0.000000    1.000000    9  3
        0.000000    1.000000    9  7
        0.000000    1.000000    9  8
        """
        )
        redundant = tskit.load_text(
            nodes_1, edges_1, sequence_length=1, strict=False, base64_metadata=False
        )

        nodes_2 = io.StringIO(
            """\
        id  is_sample   time    population  individual  metadata
        0   0   0.000000    0   -1
        1   0   0.000000    0   -1
        2   0   0.000000    0   -1
        3   1   0.000000    0   -1
        4   0   0.000000    0   -1
        5   0   0.000000    0   -1
        6   1   1.000000    0   -1
        7   1   2.000000    0   -1
        8   0   2.000000    0   -1
        9   0   3.000000    0   -1
        """
        )
        edges_2 = io.StringIO(
            """\
        left    right   parent  child
        0.000000    1.000000    7  2
        0.000000    1.000000    7  6
        0.000000    1.000000    9  3
        0.000000    1.000000    9  7
        """
        )
        simplified = tskit.load_text(
            nodes_2, edges_2, sequence_length=1, strict=False, base64_metadata=False
        )
        t1 = next(redundant.trees(sample_lists=True))
        t2 = next(simplified.trees(sample_lists=True))
        assert t1.kc_distance(t2, 0) == 0
        assert t1.kc_distance(t2, 1) == 0


# Test the RF distance metrics:
# TODO: integrate with the KC tests


class TestTreeSameSamples:
    # Tree1
    # 2.00┊    6    ┊
    #     ┊  ┏━┻━┓  ┊
    # 1.00┊  4   5  ┊
    #     ┊ ┏┻┓ ┏┻┓ ┊
    # 0.00┊ 0 1 2 3 ┊
    #     0         1
    #
    # Tree2
    # 3.00┊   6     ┊
    #     ┊ ┏━┻━┓   ┊
    # 2.00┊ ┃   5   ┊
    #     ┊ ┃ ┏━┻┓  ┊
    # 1.00┊ ┃ ┃  4  ┊
    #     ┊ ┃ ┃ ┏┻┓ ┊
    # 0.00┊ 0 1 2 3 ┊
    #     0         1

    @tests.cached_example
    def tree(self):
        return tskit.Tree.generate_balanced(4)

    @tests.cached_example
    def tree_other(self):
        return tskit.Tree.generate_comb(4)

    def test_rf_distance(self):
        assert self.tree().rf_distance(self.tree_other()) == 2


class TestTreeDifferentSamples:
    # Tree1
    # 2.00┊     6     ┊
    #     ┊   ┏━┻━┓   ┊
    # 1.00┊   4   5   ┊
    #     ┊  ┏┻┓ ┏┻┓  ┊
    # 0.00┊  0 1 2 3  ┊
    #     0           1
    #
    # Tree2
    # 4.00┊   8       ┊
    #     ┊ ┏━┻━┓     ┊
    # 3.00┊ ┃   7     ┊
    #     ┊ ┃ ┏━┻━┓   ┊
    # 2.00┊ ┃ ┃   6   ┊
    #     ┊ ┃ ┃ ┏━┻┓  ┊
    # 1.00┊ ┃ ┃ ┃  5  ┊
    #     ┊ ┃ ┃ ┃ ┏┻┓ ┊
    # 0.00┊ 0 1 2 3 4 ┊
    #     0           1

    @tests.cached_example
    def tree(self):
        return tskit.Tree.generate_balanced(4)

    @tests.cached_example
    def tree_other(self):
        return tskit.Tree.generate_comb(5)

    def test_rf_distance(self):
        assert self.tree().rf_distance(self.tree_other()) == 8


class TestTreeMultiRoots:
    # Tree1
    # 4.00┊        15             ┊
    #     ┊     ┏━━━┻━━━┓         ┊
    # 3.00┊     ┃      14         ┊
    #     ┊     ┃     ┏━┻━┓       ┊
    # 2.00┊    12     ┃  13       ┊
    #     ┊   ┏━┻━┓   ┃  ┏┻┓      ┊
    # 1.00┊   9  10   ┃  ┃ 11     ┊
    #     ┊  ┏┻┓ ┏┻┓ ┏┻┓ ┃ ┏┻┓    ┊
    # 0.00┊  0 1 2 3 4 5 6 7 8    ┊
    #     0                       1
    #
    # Tree2
    # 3.00┊              15       ┊
    #     ┊            ┏━━┻━┓     ┊
    # 2.00┊     11     ┃   14     ┊
    #     ┊    ┏━┻━┓   ┃  ┏━┻┓    ┊
    # 1.00┊    9  10  12  ┃ 13    ┊
    #     ┊   ┏┻┓ ┏┻┓ ┏┻┓ ┃ ┏┻┓   ┊
    # 0.00┊   0 1 2 3 4 5 6 7 8   ┊
    #     0                       1

    @tests.cached_example
    def tree(self):
        return tskit.Tree.generate_balanced(9)

    @tests.cached_example
    def tree_other(self):
        tables = tskit.Tree.generate_balanced(9, arity=2).tree_sequence.dump_tables()
        edges = tables.edges.copy()
        tables.edges.clear()
        for edge in edges:
            if edge.parent != 16:
                tables.edges.append(edge)
        return tables.tree_sequence().first()

    def test_rf_distance(self):
        with pytest.raises(ValueError):
            self.tree().rf_distance(self.tree_other())


class TestEmpty:
    @tests.cached_example
    def tree(self):
        tables = tskit.TableCollection(1)
        return tables.tree_sequence().first()

    @tests.cached_example
    def tree_other(self):
        tables = tskit.TableCollection(1)
        return tables.tree_sequence().first()

    def test_rf_distance(self):
        with pytest.raises(ValueError):
            self.tree().rf_distance(self.tree_other())


class TestTreeInNullState:
    @tests.cached_example
    def tsk_tree1(self):
        tree = tskit.Tree.generate_comb(5)
        tree.clear()
        return tree

    @tests.cached_example
    def tree_other(self):
        tree = tskit.Tree.generate_comb(5)
        tree.clear()
        return tree

    def test_rf_distance(self):
        with pytest.raises(ValueError):
            self.tsk_tree1().rf_distance(self.tree_other())


class TestAllRootsN5:
    @tests.cached_example
    def tree(self):
        tables = tskit.TableCollection(1)
        for _ in range(5):
            tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        return tables.tree_sequence().first()

    def test_rf_distance(self):
        with pytest.raises(ValueError, match="single root"):
            self.tree().rf_distance(self.tree())


class TestWithPackages:
    def to_dendropy(self, newick_data, tns):
        return dendropy.Tree.get(
            data=newick_data,
            schema="newick",
            rooting="force-rooted",
            taxon_namespace=tns,
        )

    def dendropy_rf_distance(self, tree1, tree2, weighted=False):
        tns = dendropy.TaxonNamespace()
        tree1 = self.to_dendropy(tree1.as_newick(), tns)
        tree2 = self.to_dendropy(tree2.as_newick(), tns)
        tree1.encode_bipartitions()
        tree2.encode_bipartitions()
        if weighted:
            return treecompare.weighted_robinson_foulds_distance(tree1, tree2)
        else:
            return treecompare.unweighted_robinson_foulds_distance(tree1, tree2)

    @pytest.mark.parametrize("n", [2, 3, 5, 10, 20])
    def test_rf_distance_against_dendropy(self, n):
        trees = []
        for seed in [42, 43]:
            ts = msprime.sim_ancestry(n, ploidy=1, random_seed=seed)
            trees.append(ts.first())
        rf1 = trees[0].rf_distance(trees[1])
        rf2 = self.dendropy_rf_distance(trees[0], trees[1])
        assert rf1 == rf2


class TestDistanceBetween:
    @pytest.mark.parametrize(
        ("u", "v"),
        itertools.combinations([0, 1, 2, 3], 2),
    )
    def test_distance_between_sample(self, u, v):
        ts = msprime.sim_ancestry(
            2, sequence_length=10, recombination_rate=0.1, random_seed=42
        )
        test_tree = ts.first()
        assert test_tree.distance_between(u, v) == pytest.approx(
            ts.diversity([u, v], mode="branch", windows="trees")[0]
        )

    def test_distance_between_same_node(self):
        ts = msprime.sim_ancestry(
            2, sequence_length=10, recombination_rate=0.1, random_seed=42
        )
        test_tree = ts.first()
        assert test_tree.distance_between(0, 0) == 0

    def test_distance_between_nodes(self):
        # 4.00┊   8       ┊
        #     ┊ ┏━┻━┓     ┊
        # 3.00┊ ┃   7     ┊
        #     ┊ ┃ ┏━┻━┓   ┊
        # 2.00┊ ┃ ┃   6   ┊
        #     ┊ ┃ ┃ ┏━┻┓  ┊
        # 1.00┊ ┃ ┃ ┃  5  ┊
        #     ┊ ┃ ┃ ┃ ┏┻┓ ┊
        # 0.00┊ 0 1 2 3 4 ┊
        #     0           1
        ts = tskit.Tree.generate_comb(5)
        assert ts.distance_between(1, 7) == 3.0
        assert ts.distance_between(6, 8) == 2.0

    def test_distance_between_invalid_nodes(self):
        ts = tskit.Tree.generate_comb(5)
        with pytest.raises(ValueError):
            ts.distance_between(0, 100)


--- ../../tskit/python/tests/test_relatedness_vector.py ---


"""
Test cases for matrix-vector product stats
"""
import msprime
import numpy as np
import pytest

import tskit
from tests import tsutil
from tests.test_highlevel import get_example_tree_sequences

# ↑ See https://github.com/tskit-dev/tskit/issues/1804 for when
# we can remove this.


# Implementation note: the class structure here, where we pass in all the
# needed arrays through the constructor was determined by an older version
# in which we used numba acceleration. We could just pass in a reference to
# the tree sequence now, but it is useful to keep track of exactly what we
# require, so leaving it as it is for now.
class RelatednessVector:
    def __init__(
        self,
        sample_weights,
        windows,
        num_nodes,
        samples,
        focal_nodes,
        nodes_time,
        edges_left,
        edges_right,
        edges_parent,
        edges_child,
        sequence_length,
        tree_pos,
        verbosity=0,
        internal_checks=False,
        centre=True,
    ):
        self.sample_weights = np.asarray(sample_weights, dtype=np.float64)
        self.num_weights = self.sample_weights.shape[1]
        self.windows = windows
        N = num_nodes
        self.parent = np.full(N, -1, dtype=np.int32)
        # Edges and indexes
        self.edges_left = edges_left
        self.edges_right = edges_right
        self.edges_parent = edges_parent
        self.edges_child = edges_child
        self.sequence_length = sequence_length
        self.nodes_time = nodes_time
        self.samples = samples
        self.focal_nodes = focal_nodes
        self.tree_pos = tree_pos
        self.position = windows[0]
        self.x = np.zeros(N, dtype=np.float64)
        self.w = np.zeros((N, self.num_weights), dtype=np.float64)
        self.v = np.zeros((N, self.num_weights), dtype=np.float64)
        self.verbosity = verbosity
        self.internal_checks = internal_checks
        self.centre = centre

        if self.centre:
            self.sample_weights -= np.mean(self.sample_weights, axis=0)

        for j, u in enumerate(samples):
            self.w[u] = self.sample_weights[j]

        if self.verbosity > 0:
            self.print_state("init")

    def print_state(self, msg=""):
        num_nodes = len(self.parent)
        print(f"..........{msg}................")
        print("tree_pos:")
        print(self.tree_pos)
        print(f"position = {self.position}")
        for j in range(num_nodes):
            st = f"{self.nodes_time[j]}"
            pt = (
                "NaN"
                if self.parent[j] == tskit.NULL
                else f"{self.nodes_time[self.parent[j]]}"
            )
            print(
                f"node {j} -> {self.parent[j]}: "
                f"z = ({pt} - {st})"
                f" * ({self.position} - {self.x[j]:.2})"
                f" * {','.join(map(str, self.w[j].round(2)))}"
                f" = {','.join(map(str, self.get_z(j).round(2)))}"
            )
            print(f"         value: {','.join(map(str, self.v[j].round(2)))}")
        roots = []
        fmt = "{:<6}{:>8}\t{}\t{}\t{}"
        s = f"roots = {roots}\n"
        s += (
            fmt.format(
                "node",
                "parent",
                "value",
                "weight",
                "z",
            )
            + "\n"
        )
        for u in range(num_nodes):
            u_str = f"{u}"
            s += (
                fmt.format(
                    u_str,
                    self.parent[u],
                    ",".join(map(str, self.v[u].round(2))),
                    ",".join(map(str, self.w[u].round(2))),
                    ",".join(map(str, self.get_z(u).round(2))),
                )
                + "\n"
            )
        print(s)

        print("Current state:")
        state = self.current_state()
        for j, x in enumerate(state):
            print(f"   {j}: {x}")
        print("..........................")

    def remove_edge(self, p, c):
        if self.verbosity > 0:
            self.print_state(f"remove {int(p), int(c)}")
        assert p != -1
        self.v[c] += self.get_z(c)
        self.x[c] = self.position
        self.parent[c] = -1
        self.adjust_path_up(p, c, -1)

    def insert_edge(self, p, c):
        if self.verbosity > 0:
            self.print_state(f"insert {int(p), int(c)}")
        assert p != -1
        assert self.parent[c] == -1, "contradictory edges"
        self.adjust_path_up(p, c, +1)
        self.x[c] = self.position
        self.parent[c] = p

    def adjust_path_up(self, p, c, sign):
        # sign = -1 for removing edges, +1 for adding
        while p != tskit.NULL:
            self.v[p] += self.get_z(p)
            self.x[p] = self.position
            self.v[c] -= sign * self.v[p]
            self.w[p] += sign * self.w[c]
            p = self.parent[p]

    def get_z(self, u):
        p = self.parent[u]
        if p == tskit.NULL:
            return np.zeros(self.num_weights, dtype=np.float64)
        time = self.nodes_time[p] - self.nodes_time[u]
        span = self.position - self.x[u]
        return time * span * self.w[u]

    def mrca(self, a, b):
        # just used for `current_state`
        aa = [a]
        while a != tskit.NULL:
            a = self.parent[a]
            aa.append(a)
        while b not in aa:
            b = self.parent[b]
        return b

    def write_output(self):
        """
        Compute and return the current state, zero-ing out
        all contributions (used for switching between windows).
        """
        n = len(self.focal_nodes)
        out = np.zeros((n, self.num_weights))
        for j, c in enumerate(self.focal_nodes):
            while c != tskit.NULL:
                if self.x[c] != self.position:
                    self.v[c] += self.get_z(c)
                    self.x[c] = self.position
                out[j] += self.v[c]
                c = self.parent[c]
        self.v *= 0.0
        return out

    def current_state(self):
        """
        Compute the current output, for debugging.
        """
        if self.verbosity > 2:
            print("---------------")
        n = len(self.focal_nodes)
        out = np.zeros((n, self.num_weights))
        for j, a in enumerate(self.focal_nodes):
            # edges on the path up from a
            pa = a
            while pa != tskit.NULL:
                if self.verbosity > 2:
                    print("edge:", pa, self.get_z(pa))
                out[j] += self.get_z(pa) + self.v[pa]
                pa = self.parent[pa]
        if self.verbosity > 2:
            print("---------------")
        return out

    def run(self):
        M = self.edges_left.shape[0]
        edges_left = self.edges_left
        edges_right = self.edges_right
        edges_parent = self.edges_parent
        edges_child = self.edges_child
        tree_pos = self.tree_pos
        in_order = tree_pos.in_range.order
        out_order = tree_pos.out_range.order
        num_windows = len(self.windows) - 1
        out = np.zeros(
            (num_windows, len(self.focal_nodes), self.sample_weights.shape[1])
        )

        m = 0
        self.position = self.windows[0]

        # seek to first window
        for j in range(tree_pos.in_range.start, tree_pos.in_range.stop, 1):
            e = in_order[j]
            if edges_left[e] <= self.position and self.position < edges_right[e]:
                p = edges_parent[e]
                c = edges_child[e]
                self.insert_edge(p, c)

        valid = tree_pos.next()
        j = tree_pos.in_range.start - 1
        k = tree_pos.out_range.start - 1
        while m < num_windows:
            if valid and self.position == tree_pos.interval.left:
                for k in range(tree_pos.out_range.start, tree_pos.out_range.stop, 1):
                    e = out_order[k]
                    p = edges_parent[e]
                    c = edges_child[e]
                    self.remove_edge(p, c)
                for j in range(tree_pos.in_range.start, tree_pos.in_range.stop, 1):
                    e = in_order[j]
                    p = edges_parent[e]
                    c = edges_child[e]
                    self.insert_edge(p, c)
                    assert self.parent[p] == tskit.NULL or self.x[p] == self.position
                valid = tree_pos.next()
            next_position = self.windows[m + 1]
            if j + 1 < M:
                next_position = min(next_position, edges_left[in_order[j + 1]])
            if k + 1 < M:
                next_position = min(next_position, edges_right[out_order[k + 1]])
            assert self.position < next_position
            self.position = next_position
            if self.position == self.windows[m + 1]:
                out[m] = self.write_output()
                m = m + 1

        if self.verbosity > 1:
            self.print_state()

        if self.centre:
            for m in range(num_windows):
                out[m] -= np.mean(out[m], axis=0)
        return out


def relatedness_vector(ts, sample_weights, windows=None, nodes=None, **kwargs):
    if len(sample_weights.shape) == 1:
        sample_weights = sample_weights[:, np.newaxis]
    if nodes is None:
        nodes = np.fromiter(ts.samples(), dtype=np.int32)
    drop_dimension = windows is None
    if drop_dimension:
        windows = [0, ts.sequence_length]

    tree_pos = tsutil.TreePosition(ts)
    breakpoints = np.fromiter(ts.breakpoints(), dtype="float")
    index = np.searchsorted(breakpoints, windows[0])
    if breakpoints[index] > windows[0]:
        index -= 1
    tree_pos.seek_forward(index)

    rv = RelatednessVector(
        sample_weights,
        windows,
        ts.num_nodes,
        samples=ts.samples(),
        focal_nodes=nodes,
        nodes_time=ts.nodes_time,
        edges_left=ts.edges_left,
        edges_right=ts.edges_right,
        edges_parent=ts.edges_parent,
        edges_child=ts.edges_child,
        sequence_length=ts.sequence_length,
        tree_pos=tree_pos,
        **kwargs,
    )
    out = rv.run()
    if drop_dimension:
        assert len(out.shape) == 3 and out.shape[0] == 1
        out = out[0]
    return out


def relatedness_matrix(ts, windows, centre, nodes=None):
    if nodes is None:
        keep_rows = np.arange(ts.num_samples)
        keep_cols = np.arange(ts.num_samples)
    else:
        orig_samples = list(ts.samples())
        extra_nodes = set(nodes).difference(set(orig_samples))
        tables = ts.dump_tables()
        tables.nodes.clear()
        for n in ts.nodes():
            if n.id in extra_nodes:
                n = n.replace(flags=n.flags | tskit.NODE_IS_SAMPLE)
            tables.nodes.append(n)
        ts = tables.tree_sequence()
        all_samples = list(ts.samples())
        keep_rows = np.array([all_samples.index(i) for i in nodes])
        keep_cols = np.array([all_samples.index(i) for i in orig_samples])

    use_windows = windows
    drop_first = windows is not None and windows[0] > 0
    if drop_first:
        use_windows = np.concatenate([[0], np.array(use_windows)])
    drop_last = windows is not None and windows[-1] < ts.sequence_length
    if drop_last:
        use_windows = np.concatenate([np.array(use_windows), [ts.sequence_length]])
    Sigma = ts.genetic_relatedness(
        sample_sets=[[i] for i in ts.samples()],
        indexes=[(i, j) for i in range(ts.num_samples) for j in range(ts.num_samples)],
        windows=use_windows,
        mode="branch",
        span_normalise=False,
        proportion=False,
        centre=centre,
    )
    if windows is not None:
        if drop_first:
            Sigma = Sigma[1:]
        if drop_last:
            Sigma = Sigma[:-1]
    nwin = 1 if windows is None else len(windows) - 1
    shape = (nwin, ts.num_samples, ts.num_samples)
    Sigma = Sigma.reshape(shape)
    out = np.array([x[np.ix_(keep_rows, keep_cols)] for x in Sigma])
    if windows is None:
        out = out[0]
    return out


def verify_relatedness_vector(
    ts, w, windows, *, internal_checks=False, verbosity=0, centre=True, nodes=None
):
    R1 = relatedness_vector(
        ts,
        sample_weights=w,
        windows=windows,
        internal_checks=internal_checks,
        verbosity=verbosity,
        centre=centre,
        nodes=nodes,
    )
    nrows = ts.num_samples if nodes is None else len(nodes)
    wvec = w if len(w.shape) > 1 else w[:, np.newaxis]
    Sigma = relatedness_matrix(ts, windows=windows, centre=centre, nodes=nodes)
    if windows is None:
        R2 = Sigma.dot(wvec)
    else:
        R2 = np.zeros((len(windows) - 1, nrows, wvec.shape[1]))
        for k in range(len(windows) - 1):
            R2[k] = Sigma[k].dot(wvec)
    R3 = ts.genetic_relatedness_vector(
        w, windows=windows, mode="branch", centre=centre, nodes=nodes
    )
    if verbosity > 0:
        print(ts.draw_text())
        print("weights:", w)
        print("windows:", windows)
        print("centre:", centre)
        print("here:", R1)
        print("with ts:", R2)
        print("with lib:", R3)
        print("Sigma:", Sigma)
    if windows is None:
        assert R1.shape == (nrows, wvec.shape[1])
    else:
        assert R1.shape == (len(windows) - 1, nrows, wvec.shape[1])
    np.testing.assert_allclose(R1, R2, atol=1e-10)
    np.testing.assert_allclose(R1, R3, atol=1e-10)
    return R1


def check_relatedness_vector(
    ts,
    n=2,
    num_windows=0,
    *,
    internal_checks=False,
    verbosity=0,
    seed=123,
    centre=True,
    do_nodes=True,
):
    rng = np.random.default_rng(seed=seed)
    if num_windows == 0:
        windows = None
    elif num_windows % 2 == 0:
        windows = np.linspace(
            0.2 * ts.sequence_length, 0.8 * ts.sequence_length, num_windows + 1
        )
    else:
        windows = np.linspace(0, ts.sequence_length, num_windows + 1)
    num_nodes_list = (0,) if (centre or not do_nodes) else (0, 3)
    for num_nodes in num_nodes_list:
        if num_nodes == 0:
            nodes = None
        else:
            nodes = rng.choice(ts.num_nodes, num_nodes, replace=False)
        for k in range(n):
            if k == 0:
                w = rng.normal(size=ts.num_samples)
            else:
                w = rng.normal(size=ts.num_samples * k).reshape((ts.num_samples, k))
            w = np.round(len(w) * w)
            R = verify_relatedness_vector(
                ts,
                w,
                windows,
                internal_checks=internal_checks,
                verbosity=verbosity,
                centre=centre,
                nodes=nodes,
            )
    return R


class TestExamples:

    def test_bad_weights(self):
        n = 5
        ts = msprime.sim_ancestry(
            n,
            ploidy=2,
            sequence_length=10,
            random_seed=123,
        )
        for bad_W in (None, [1], np.ones((3 * n, 2)), np.ones((n - 1, 2))):
            with pytest.raises(ValueError, match="number of samples"):
                ts.genetic_relatedness_vector(bad_W, mode="branch")

    def test_bad_windows(self):
        n = 5
        ts = msprime.sim_ancestry(
            n,
            ploidy=2,
            sequence_length=10,
            random_seed=123,
        )
        for bad_w in ([1], []):
            with pytest.raises(ValueError, match="Windows array"):
                ts.genetic_relatedness_vector(
                    np.ones(ts.num_samples), windows=bad_w, mode="branch"
                )

    def test_nodes_centred_error(self):
        ts = msprime.sim_ancestry(
            5,
            ploidy=2,
            sequence_length=10,
            random_seed=123,
        )
        with pytest.raises(ValueError, match="must have centre"):
            ts.genetic_relatedness_vector(
                np.ones(ts.num_samples), mode="branch", centre=True, nodes=[0, 1]
            )

    def test_bad_nodes(self):
        n = 5
        ts = msprime.sim_ancestry(
            n,
            ploidy=2,
            sequence_length=10,
            random_seed=123,
        )
        for bad_nodes in ([[]], "foo"):
            with pytest.raises(ValueError):
                ts.genetic_relatedness_vector(
                    np.ones(ts.num_samples),
                    mode="branch",
                    centre=False,
                    nodes=bad_nodes,
                )
        for bad_nodes in ([-1, 10], [3, 2 * ts.num_nodes]):
            with pytest.raises(tskit.LibraryError, match="TSK_ERR_NODE_OUT_OF_BOUNDS"):
                ts.genetic_relatedness_vector(
                    np.ones(ts.num_samples),
                    mode="branch",
                    centre=False,
                    nodes=bad_nodes,
                )

    def test_good_nodes(self):
        n = 5
        ts = msprime.sim_ancestry(
            n,
            ploidy=2,
            sequence_length=10,
            random_seed=123,
        )
        V0 = ts.genetic_relatedness_vector(
            np.ones(ts.num_samples), mode="branch", centre=False
        )
        V = ts.genetic_relatedness_vector(
            np.ones(ts.num_samples),
            mode="branch",
            centre=False,
            nodes=list(ts.samples()),
        )
        np.testing.assert_allclose(V0, V, atol=1e-13)
        V = ts.genetic_relatedness_vector(
            np.ones(ts.num_samples),
            mode="branch",
            centre=False,
            nodes=np.fromiter(ts.samples(), dtype=np.int32),
        )
        np.testing.assert_allclose(V0, V, atol=1e-13)
        V = ts.genetic_relatedness_vector(
            np.ones(ts.num_samples),
            mode="branch",
            centre=False,
            nodes=np.fromiter(ts.samples(), dtype=np.int64),
        )
        np.testing.assert_allclose(V0, V, atol=1e-13)
        V = ts.genetic_relatedness_vector(
            np.ones(ts.num_samples),
            mode="branch",
            centre=False,
            nodes=list(ts.samples())[:2],
        )
        np.testing.assert_allclose(V0[:2], V, atol=1e-13)

    @pytest.mark.parametrize("n", [2, 3, 5])
    @pytest.mark.parametrize("seed", range(1, 4))
    @pytest.mark.parametrize("centre", (True, False))
    @pytest.mark.parametrize("num_windows", (0, 1, 2, 3))
    def test_small_internal_checks(self, n, seed, centre, num_windows):
        ts = msprime.sim_ancestry(
            n,
            ploidy=1,
            sequence_length=1000,
            recombination_rate=0.01,
            random_seed=seed,
        )
        assert ts.num_trees >= 2
        check_relatedness_vector(
            ts, num_windows=num_windows, internal_checks=True, centre=centre
        )

    @pytest.mark.parametrize("n", [2, 3, 5, 15])
    @pytest.mark.parametrize("seed", range(1, 5))
    @pytest.mark.parametrize("centre", (True, False))
    @pytest.mark.parametrize("num_windows", (0, 1, 2, 3))
    def test_simple_sims(self, n, seed, centre, num_windows):
        ts = msprime.sim_ancestry(
            n,
            ploidy=1,
            population_size=20,
            sequence_length=100,
            recombination_rate=0.01,
            random_seed=seed,
        )
        assert ts.num_trees >= 2
        check_relatedness_vector(
            ts, num_windows=num_windows, centre=centre, verbosity=0
        )

    def test_simple_sims_windows(self):
        L = 100
        ts = msprime.sim_ancestry(
            5,
            ploidy=1,
            population_size=20,
            sequence_length=L,
            recombination_rate=0.01,
            random_seed=345,
        )
        assert ts.num_trees >= 2
        W = np.linspace(0, 1, 2 * ts.num_samples).reshape((ts.num_samples, 2))
        kwargs = {"centre": False, "mode": "branch"}
        total = ts.genetic_relatedness_vector(W, **kwargs)
        for windows in [[0, L], [0, L / 3, L / 2, L]]:
            pieces = ts.genetic_relatedness_vector(W, windows=windows, **kwargs)
            np.testing.assert_allclose(total, pieces.sum(axis=0), atol=1e-13)
            assert len(pieces) == len(windows) - 1
            for k in range(len(pieces)):
                piece = ts.genetic_relatedness_vector(
                    W, windows=windows[k : k + 2], **kwargs
                )
                assert piece.shape[0] == 1
                np.testing.assert_allclose(piece[0], pieces[k], atol=1e-13)

    @pytest.mark.parametrize("n", [2, 3, 5, 15])
    @pytest.mark.parametrize("centre", (True, False))
    def test_single_balanced_tree(self, n, centre):
        ts = tskit.Tree.generate_balanced(n).tree_sequence
        check_relatedness_vector(ts, internal_checks=True, centre=centre)

    @pytest.mark.parametrize("centre", (True, False))
    def test_internal_sample(self, centre):
        tables = tskit.Tree.generate_balanced(4).tree_sequence.dump_tables()
        flags = tables.nodes.flags
        flags[3] = 0
        flags[5] = tskit.NODE_IS_SAMPLE
        tables.nodes.flags = flags
        ts = tables.tree_sequence()
        check_relatedness_vector(ts, centre=centre)

    @pytest.mark.parametrize("seed", range(1, 5))
    @pytest.mark.parametrize("centre", (True, False))
    @pytest.mark.parametrize("num_windows", (0, 1, 2, 3))
    def test_one_internal_sample_sims(self, seed, centre, num_windows):
        ts = msprime.sim_ancestry(
            10,
            ploidy=1,
            population_size=20,
            sequence_length=100,
            recombination_rate=0.01,
            random_seed=seed,
        )
        t = ts.dump_tables()
        # Add a new sample directly below another sample
        u = t.nodes.add_row(time=-1, flags=tskit.NODE_IS_SAMPLE)
        t.edges.add_row(parent=0, child=u, left=0, right=ts.sequence_length)
        t.sort()
        t.build_index()
        ts = t.tree_sequence()
        check_relatedness_vector(ts, num_windows=num_windows, centre=centre)

    @pytest.mark.parametrize("centre", (True, False))
    @pytest.mark.parametrize("num_windows", (0, 1, 2, 3))
    def test_missing_flanks(self, centre, num_windows):
        ts = msprime.sim_ancestry(
            2,
            ploidy=1,
            population_size=10,
            sequence_length=100,
            recombination_rate=0.001,
            random_seed=1234,
        )
        assert ts.num_trees >= 2
        ts = ts.keep_intervals([[20, 80]])
        assert ts.first().interval == (0, 20)
        check_relatedness_vector(ts, num_windows=num_windows, centre=centre)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    @pytest.mark.parametrize("centre", (True, False))
    def test_suite_examples(self, ts, centre):
        if ts.num_samples > 0:
            check_relatedness_vector(ts, centre=centre)

    @pytest.mark.parametrize("n", [2, 3, 10])
    def test_dangling_on_samples(self, n):
        # Adding non sample branches below the samples does not alter
        # the overall divergence *between* the samples
        ts1 = tskit.Tree.generate_balanced(n).tree_sequence
        D1 = check_relatedness_vector(ts1, do_nodes=False)
        tables = ts1.dump_tables()
        for u in ts1.samples():
            v = tables.nodes.add_row(time=-1)
            tables.edges.add_row(left=0, right=ts1.sequence_length, parent=u, child=v)
        tables.sort()
        tables.build_index()
        ts2 = tables.tree_sequence()
        D2 = check_relatedness_vector(ts2, internal_checks=True, do_nodes=False)
        np.testing.assert_array_almost_equal(D1, D2)

    @pytest.mark.parametrize("n", [2, 3, 10])
    @pytest.mark.parametrize("centre", (True, False))
    def test_dangling_on_all(self, n, centre):
        # Adding non sample branches below the samples does not alter
        # the overall divergence *between* the samples
        ts1 = tskit.Tree.generate_balanced(n).tree_sequence
        D1 = check_relatedness_vector(ts1, centre=centre, do_nodes=False)
        tables = ts1.dump_tables()
        for u in range(ts1.num_nodes):
            v = tables.nodes.add_row(time=-1)
            tables.edges.add_row(left=0, right=ts1.sequence_length, parent=u, child=v)
        tables.sort()
        tables.build_index()
        ts2 = tables.tree_sequence()
        D2 = check_relatedness_vector(
            ts2, internal_checks=True, centre=centre, do_nodes=False
        )
        np.testing.assert_array_almost_equal(D1, D2)

    @pytest.mark.parametrize("centre", (True, False))
    def test_disconnected_non_sample_topology(self, centre):
        # Adding non sample branches below the samples does not alter
        # the overall divergence *between* the samples
        ts1 = tskit.Tree.generate_balanced(5).tree_sequence
        D1 = check_relatedness_vector(ts1, centre=centre, do_nodes=False)
        tables = ts1.dump_tables()
        # Add an extra bit of disconnected non-sample topology
        u = tables.nodes.add_row(time=0)
        v = tables.nodes.add_row(time=1)
        tables.edges.add_row(left=0, right=ts1.sequence_length, parent=v, child=u)
        tables.sort()
        tables.build_index()
        ts2 = tables.tree_sequence()
        D2 = check_relatedness_vector(
            ts2, internal_checks=True, centre=centre, do_nodes=False
        )
        np.testing.assert_array_almost_equal(D1, D2)


--- ../../tskit/python/tests/test_balance_metrics.py ---


"""
Tests for tree balance/imbalance metrics.
"""
import math

import numpy as np
import pytest

import tests
import tskit
from tests.test_highlevel import get_example_tree_sequences

# ↑ See https://github.com/tskit-dev/tskit/issues/1804 for when
# we can remove this.


def sackin_index_definition(tree):
    return sum(tree.depth(u) for u in tree.leaves())


def colless_index_definition(tree):
    is_binary = all(
        tree.num_children(u) == 2 for u in tree.nodes() if tree.is_internal(u)
    )
    if tree.num_roots != 1:
        raise ValueError("Colless index not defined for multiroot trees")
    if not is_binary:
        raise ValueError("Colless index not defined for nonbinary trees")

    return sum(
        abs(
            len(list(tree.leaves(tree.left_child(u))))
            - len(list(tree.leaves(tree.right_child(u))))
        )
        for u in tree.nodes()
        if tree.is_internal(u)
    )


def b1_index_definition(tree):
    return sum(
        1 / max(tree.path_length(n, leaf) for leaf in tree.leaves(n))
        for n in tree.nodes()
        if tree.parent(n) != tskit.NULL and tree.is_internal(n)
    )


def b2_index_definition(tree, base=10):
    if tree.num_roots != 1:
        raise ValueError("B2 index is only defined for trees with one root")
    proba = [
        np.prod([1 / tree.num_children(u) for u in tree.ancestors(leaf)])
        for leaf in tree.leaves()
    ]
    return -sum(p * math.log(p, base) for p in proba)


class TestDefinitions:
    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_sackin(self, ts):
        for tree in ts.trees():
            assert tree.sackin_index() == sackin_index_definition(tree)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_colless(self, ts):
        for tree in ts.trees():
            is_binary = all(
                tree.num_children(u) == 2 for u in tree.nodes() if tree.is_internal(u)
            )
            if tree.num_roots != 1 or not is_binary:
                with pytest.raises(tskit.LibraryError):
                    tree.colless_index()
                with pytest.raises(ValueError):
                    colless_index_definition(tree)
            else:
                assert tree.colless_index() == colless_index_definition(tree)

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_b1(self, ts):
        for tree in ts.trees():
            assert tree.b1_index() == pytest.approx(b1_index_definition(tree))

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    def test_b2(self, ts):
        for tree in ts.trees():
            if tree.num_roots != 1:
                with pytest.raises(tskit.LibraryError, match="MULTIROOT"):
                    tree.b2_index()
                with pytest.raises(ValueError):
                    b2_index_definition(tree)
            else:
                assert tree.b2_index() == pytest.approx(b2_index_definition(tree))

    @pytest.mark.parametrize("ts", get_example_tree_sequences())
    @pytest.mark.parametrize("base", [0.1, 1.1, 2, 10, math.e, np.array([3])[0]])
    def test_b2_base(self, ts, base):
        for tree in ts.trees():
            if tree.num_roots != 1:
                with pytest.raises(tskit.LibraryError, match="MULTIROOT"):
                    tree.b2_index(base)
                with pytest.raises(ValueError):
                    b2_index_definition(tree, base)
            else:
                assert tree.b2_index(base) == pytest.approx(
                    b2_index_definition(tree, base)
                )


class TestBalancedBinaryOdd:
    # 2.00┊   4   ┊
    #     ┊ ┏━┻┓  ┊
    # 1.00┊ ┃  3  ┊
    #     ┊ ┃ ┏┻┓ ┊
    # 0.00┊ 0 1 2 ┊
    #     0      1
    @tests.cached_example
    def tree(self):
        return tskit.Tree.generate_balanced(3)

    def test_sackin(self):
        assert self.tree().sackin_index() == 5

    def test_colless(self):
        assert self.tree().colless_index() == 1

    def test_b1(self):
        assert self.tree().b1_index() == 1

    def test_b2(self):
        assert self.tree().b2_index(base=10) == pytest.approx(0.4515, rel=1e-3)


class TestBalancedBinaryEven:
    # 2.00┊    6    ┊
    #     ┊  ┏━┻━┓  ┊
    # 1.00┊  4   5  ┊
    #     ┊ ┏┻┓ ┏┻┓ ┊
    # 0.00┊ 0 1 2 3 ┊
    #     0         1
    @tests.cached_example
    def tree(self):
        return tskit.Tree.generate_balanced(4)

    def test_sackin(self):
        assert self.tree().sackin_index() == 8

    def test_colless(self):
        assert self.tree().colless_index() == 0

    def test_b1(self):
        assert self.tree().b1_index() == 2

    def test_b2(self):
        assert self.tree().b2_index() == pytest.approx(0.602, rel=1e-3)

    @pytest.mark.parametrize(
        ("base", "expected"),
        [
            (2, 2),
            (3, 1.2618595071429148),
            (4, 1.0),
            (5, 0.8613531161467861),
            (10, 0.6020599913279623),
            (100, 0.30102999566398114),
            (1000000, 0.10034333188799373),
            (2.718281828459045, 1.3862943611198906),
        ],
    )
    def test_b2_base(self, base, expected):
        assert self.tree().b2_index(base) == expected

    @pytest.mark.parametrize("base", [0, -0.001, -1, -1e-6, -1e200])
    def test_b2_bad_base(self, base):
        with pytest.raises(ValueError, match="math domain"):
            self.tree().b2_index(base=base)

    def test_b2_base1(self):
        with pytest.raises(ZeroDivisionError):
            self.tree().b2_index(base=1)


class TestBalancedTernary:
    # 2.00┊        12         ┊
    #     ┊   ┏━━━━━╋━━━━━┓   ┊
    # 1.00┊   9    10    11   ┊
    #     ┊ ┏━╋━┓ ┏━╋━┓ ┏━╋━┓ ┊
    # 0.00┊ 0 1 2 3 4 5 6 7 8 ┊
    #     0                   1
    @tests.cached_example
    def tree(self):
        return tskit.Tree.generate_balanced(9, arity=3)

    def test_sackin(self):
        assert self.tree().sackin_index() == 18

    def test_colless(self):
        with pytest.raises(tskit.LibraryError, match="UNDEFINED_NONBINARY"):
            self.tree().colless_index()

    def test_b1(self):
        assert self.tree().b1_index() == 3

    def test_b2(self):
        assert self.tree().b2_index() == pytest.approx(0.954, rel=1e-3)


class TestStarN10:
    # 1.00┊         10          ┊
    #     ┊ ┏━┳━┳━┳━┳┻┳━┳━┳━┳━┓ ┊
    # 0.00┊ 0 1 2 3 4 5 6 7 8 9 ┊
    #     0                     1
    @tests.cached_example
    def tree(self):
        return tskit.Tree.generate_star(10)

    def test_sackin(self):
        assert self.tree().sackin_index() == 10

    def test_colless(self):
        with pytest.raises(tskit.LibraryError, match="UNDEFINED_NONBINARY"):
            self.tree().colless_index()

    def test_b1(self):
        assert self.tree().b1_index() == 0

    def test_b2(self):
        assert self.tree().b2_index() == pytest.approx(0.9999, rel=1e-3)


class TestCombN5:
    # 4.00┊   8       ┊
    #     ┊ ┏━┻━┓     ┊
    # 3.00┊ ┃   7     ┊
    #     ┊ ┃ ┏━┻━┓   ┊
    # 2.00┊ ┃ ┃   6   ┊
    #     ┊ ┃ ┃ ┏━┻┓  ┊
    # 1.00┊ ┃ ┃ ┃  5  ┊
    #     ┊ ┃ ┃ ┃ ┏┻┓ ┊
    # 0.00┊ 0 1 2 3 4 ┊
    #     0           1
    @tests.cached_example
    def tree(self):
        return tskit.Tree.generate_comb(5)

    def test_sackin(self):
        assert self.tree().sackin_index() == 14

    def test_colless(self):
        assert self.tree().colless_index() == 6

    def test_b1(self):
        assert self.tree().b1_index() == pytest.approx(1.833, rel=1e-3)

    def test_b2(self):
        assert self.tree().b2_index() == pytest.approx(0.564, rel=1e-3)


class TestMultiRootBinary:
    # 3.00┊            15     ┊
    #     ┊          ┏━━┻━┓   ┊
    # 2.00┊   11     ┃   14   ┊
    #     ┊  ┏━┻━┓   ┃  ┏━┻┓  ┊
    # 1.00┊  9  10  12  ┃ 13  ┊
    #     ┊ ┏┻┓ ┏┻┓ ┏┻┓ ┃ ┏┻┓ ┊
    # 0.00┊ 0 1 2 3 4 5 6 7 8 ┊
    #     0                   1
    @tests.cached_example
    def tree(self):
        tables = tskit.Tree.generate_balanced(9, arity=2).tree_sequence.dump_tables()
        edges = tables.edges.copy()
        tables.edges.clear()
        for edge in edges:
            if edge.parent != 16:
                tables.edges.append(edge)
        return tables.tree_sequence().first()

    def test_sackin(self):
        assert self.tree().sackin_index() == 20

    def test_colless(self):
        with pytest.raises(tskit.LibraryError, match="UNDEFINED_MULTIROOT"):
            self.tree().colless_index()

    def test_b1(self):
        assert self.tree().b1_index() == 4.5

    def test_b2(self):
        with pytest.raises(tskit.LibraryError, match="UNDEFINED_MULTIROOT"):
            self.tree().b2_index()


class TestEmpty:
    @tests.cached_example
    def tree(self):
        tables = tskit.TableCollection(1)
        return tables.tree_sequence().first()

    def test_sackin(self):
        assert self.tree().sackin_index() == 0

    def test_colless(self):
        with pytest.raises(tskit.LibraryError, match="UNDEFINED_MULTIROOT"):
            self.tree().colless_index()

    def test_b1(self):
        assert self.tree().b1_index() == 0

    def test_b2(self):
        with pytest.raises(tskit.LibraryError, match="UNDEFINED_MULTIROOT"):
            self.tree().b2_index()


class TestTreeInNullState:
    @tests.cached_example
    def tree(self):
        tree = tskit.Tree.generate_comb(5)
        tree.clear()
        return tree

    def test_sackin(self):
        assert self.tree().sackin_index() == 0

    def test_colless(self):
        with pytest.raises(tskit.LibraryError, match="UNDEFINED_MULTIROOT"):
            self.tree().colless_index()

    def test_b1(self):
        assert self.tree().b1_index() == 0

    def test_b2(self):
        with pytest.raises(tskit.LibraryError, match="UNDEFINED_MULTIROOT"):
            self.tree().b2_index()


class TestAllRootsN5:
    @tests.cached_example
    def tree(self):
        tables = tskit.TableCollection(1)
        for _ in range(5):
            tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        return tables.tree_sequence().first()

    def test_sackin(self):
        assert self.tree().sackin_index() == 0

    def test_colless(self):
        with pytest.raises(tskit.LibraryError, match="UNDEFINED_MULTIROOT"):
            self.tree().colless_index()

    def test_b1(self):
        assert self.tree().b1_index() == 0

    def test_b2(self):
        with pytest.raises(tskit.LibraryError, match="UNDEFINED_MULTIROOT"):
            self.tree().b2_index()


--- ../../tskit/python/tests/test_avl_tree.py ---


"""
A Python version of the C AVL tree code for development purposes.

Based on Knuth's AVL tree code in TAOCP volume 3, adapted from
https://commandlinefanatic.com/cgi-bin/showarticle.cgi?article=art070

Note there is a bug in that Python translation which is missing
P.B = 0 at the end of A9.
"""
from __future__ import annotations

import dataclasses
from typing import Any

import numpy as np
import pytest


# The nodes of the tree are assumed to contain KEY, LLINK, and RLINK fields.
# We also have a new field
#
# B(P) = balance factor of NODE(P)
#
# the height of the right subtree minus the height of the left subtree; this field
# always contains either +1, 0, or -1.  A special header node also appears at the top
# of the tree, in location HEAD; the value of RLINK(HEAD) is a pointer to the root
# of the tree, and LLINK(HEAD) is used to keep track of the overall height of the tree.
# We assume that the tree is nonempty, namely that RLINK(HEAD) != ^.


@dataclasses.dataclass(eq=False)
class Node:
    key: Any = None
    llink: Node = None
    rlink: Node = None
    balance: int = 0

    def __str__(self):
        llink = None if self.llink is None else self.llink.key
        rlink = None if self.rlink is None else self.rlink.key
        return (
            f"Node(key={self.key}, balance={self.balance}, "
            f"llink={llink}, rlink={rlink})"
        )


# For convenience in description, the algorithm uses the notation LINK(a,P)
# as a synonym for LLINK(P) if a = -1, and for RLINK(P) if a = +1.


def get_link(a, P):
    if a == -1:
        return P.llink
    else:
        return P.rlink


def set_link(a, P, val):
    if a == -1:
        P.llink = val
    else:
        P.rlink = val


class AvlTree:
    def __init__(self):
        self.head = Node()
        self.size = 0
        self.height = 0

    @property
    def root(self):
        return self.head.rlink

    def __str__(self):
        stack = [(self.head, 0)]
        s = f"size = {self.size} height = {self.height}\n"
        while len(stack) > 0:
            node, depth = stack.pop()
            s += ("  " * depth) + f"KEY={node.key} B={node.balance}\n"
            for child in [node.llink, node.rlink]:
                if child is not None:
                    stack.append((child, depth + 1))
        return s

    def ordered_keys(self):
        """
        Return the keys in sorted order. This is done by an in-order
        traversal of the nodes.
        """

        def inorder(node):
            if node is not None:
                yield from inorder(node.llink)
                yield node.key
                yield from inorder(node.rlink)

        yield from inorder(self.root)

    def search(self, key):
        P = self.root
        while P is not None:
            if key == P.key:
                break
            elif key < P.key:
                P = P.llink
            else:
                P = P.rlink
        return P

    def __insert_empty(self, key):
        self.head.rlink = Node()
        self.head.rlink.key = key
        self.size = 1
        self.height = 1
        return self.head.rlink

    def __insert(self, K):
        # A1. [Initialize.] Set T <- HEAD, S <- P <- RLINK(HEAD).
        # (The pointer variable P will move down the tree; S will point
        # to the place where rebalancing may be necessary, and
        # T always points to the parent of S.)
        T = self.head
        S = P = self.head.rlink

        # A2. [Compare.] If K < KEY(P), go to A3; if K > KEY(P), go to A4; and if
        # K = KEY(P), the search terminates successfully.
        while True:
            if K == P.key:
                return P
            elif K < P.key:
                # A3. [Move left.] Set Q <- LLINK(P). If Q = ^, set Q <= AVAIL and
                # LLINK(P) <- Q and go to step A5. Otherwise if B(Q) != 0, set T <-
                # P and S <- Q. Finally set P <- Q and return to step A2.
                Q = P.llink
                if Q is None:
                    Q = Node()
                    P.llink = Q
                    break
            # A4. [Move right.] Set Q <- RLINK(P). If Q = ^, set Q <= AVAIL and
            # RLINK(P) <- Q and go to step A5. Otherwise if B(Q) != 0, set T <- P
            # and S <- Q. Finally set P <- Q and return to step A2.
            elif K > P.key:
                Q = P.rlink
                if Q is None:
                    Q = Node()
                    P.rlink = Q
                    break
            if Q.balance != 0:
                T = P
                S = Q
            P = Q
        # A5. [Insert.] Set KEY(Q) <- K, LLINK(Q) <- RLINK(Q) <- ^, and B(Q) <- 0.
        Q.key = K
        Q.llink = Q.rlink = None
        Q.balance = 0

        # A6. [Adjust balance factors.] If K < KEY(S) set a <- -1, otherwise set a
        # <- +1. Then set R <- P <- LINK(a,S), and repeatedly do the following
        # operations zero or more times until P = Q: If K < KEY(P) set B(P) <- -1
        # and P <- LLINK(P); if K > KEY(P), set B(P) <- +1 and P <- RLINK(P).
        if K < S.key:
            a = -1
        else:
            a = 1
        R = P = get_link(a, S)
        while P != Q:
            if K < P.key:
                P.balance = -1
                P = P.llink
            elif K > P.key:
                P.balance = 1
                P = P.rlink

        # A7. [Balancing act.] Several cases now arise:
        #
        #  i) If B(S) = 0, set B(S) <- a, LLINK(HEAD) <- LLINK(HEAD) + 1, and
        #  terminate the algorithm.
        #
        if S.balance == 0:
            S.balance = a
            self.height += 1

        # ii) If B(S) = -a, set B(S) <- 0 and terminate the algorithm.

        elif S.balance == -a:
            S.balance = 0

        # iii) If B(S) = a, go to step A8 if B(R) = a, to A9 if B(R) = -a.
        else:
            if R.balance == a:
                # A8. [Single rotation.] Set P <- R, LINK(a,S) <- LINK(-a,R),
                # LINK(-a,R) <- S,B(S) <- B(R) <- 0. Go to A10.
                P = R
                set_link(a, S, get_link(-a, R))
                set_link(-a, R, S)
                S.balance = R.balance = 0
            elif R.balance == -a:
                # A9. [Double rotation.] Set P <- LINK(-a,R),
                #  LINK(-a,R) <- LINK(a,P),LINK(a,P) <- R, LINK(a,S)
                #  <- LINK(-a,P), LINK(-a,P) <- S. Now set
                #
                #               { (-a,0), if B(P) =  a;
                #  (B(S),B(R))<-{ ( 0,0), if B(P) =  0;
                #               { ( 0,a), if B(P) = -a;
                #
                #  and then set B(P) <- 0
                P = get_link(-a, R)
                set_link(-a, R, get_link(a, P))
                set_link(a, P, R)
                set_link(a, S, get_link(-a, P))
                set_link(-a, P, S)
                if P.balance == a:
                    S.balance = -a
                    R.balance = 0
                elif P.balance == 0:
                    S.balance = 0
                    R.balance = 0
                else:
                    S.balance = 0
                    R.balance = a
                P.balance = 0

            # A10. [Finishing touch.] If S = RLINK(T) then set RLINK(T) <- P,
            # otherwise set LLINK(T) <- P.
            if S == T.rlink:
                T.rlink = P
            else:
                T.llink = P

        return Q

    def insert(self, key):
        if self.size == 0:
            return self.__insert_empty(key)
        return self.__insert(key)


class TestAvlTree:
    def verify_tree(self, tree):
        """
        Check that the tree fits the AVL tree properties.
        """
        # The height of a node is its maximum distance to a leaf
        node_height = {}

        def compute_height(node):
            if node is None:
                return 0
            val = 1 + max([compute_height(node.llink), compute_height(node.rlink)])
            node_height[node] = val
            return val

        compute_height(tree.head.rlink)
        assert tree.height == max(node_height.values())
        assert tree.height == node_height[tree.head.rlink]
        # print(tree)

        # The balance factor B is the height of the right subtree
        # minus the height of the left subtree
        stack = [tree.head.rlink]
        while len(stack) > 0:
            node = stack.pop()
            # print(node, node_height[node])
            assert node.balance in [-1, 0, 1]
            lheight = None
            if node.llink is not None:
                lheight = node_height[node.llink]
                stack.append(node.llink)
            rheight = None
            if node.rlink is not None:
                rheight = node_height[node.rlink]
                stack.append(node.rlink)
            if lheight is not None and rheight is not None:
                balance_factor = rheight - lheight
                assert node.balance == balance_factor
            elif lheight is None and rheight is None:
                assert node_height[node] == 1
                assert node.balance == 0
            else:
                # if one child is None, the height of this node must be 2
                assert node_height[node] == 2
                if lheight is None:
                    assert node.balance == 1
                else:
                    assert node.balance == -1

    def verify(self, keys):
        tree = AvlTree()
        key_set = set()
        for k in keys:
            node = tree.search(k)
            if k in key_set:
                assert node is not None
                assert node.key == k
            else:
                assert node is None
            node = tree.insert(k)
            key_set.add(k)
            self.verify_tree(tree)
            assert tree.search(k) is node
        for k in range(100):
            node = tree.search(k)
            if k in key_set:
                assert node is not None
                assert node.key == k
            else:
                assert node is None
        ordered_keys = list(tree.ordered_keys())
        assert ordered_keys == list(sorted(set(keys)))

        # Implement the inorder on an existing list to mimic C algorithm
        l2 = [None for _ in ordered_keys]

        def visit(node, index, out):
            if node is None:
                return index
            index = visit(node.llink, index, out)
            out[index] = node.key
            return visit(node.rlink, index + 1, out)

        visit(tree.root, 0, l2)
        assert l2 == ordered_keys

    @pytest.mark.parametrize("n", [0, 1, 10, 33, 64, 127, 133])
    def test_sequential(self, n):
        self.verify(range(n))

    @pytest.mark.parametrize("n", [0, 1, 10, 33, 64, 127, 133])
    def test_sequential_reversed(self, n):
        self.verify(range(n)[::-1])

    @pytest.mark.parametrize("n", [0, 1, 10, 33, 64, 127, 133])
    def test_random_integers(self, n):
        rng = np.random.RandomState(42)
        values = rng.randint(-100, 100, size=n)
        self.verify(values)

    @pytest.mark.parametrize("n", [0, 1, 10, 33, 64, 127, 133])
    def test_random_floats(self, n):
        rng = np.random.RandomState(42)
        values = rng.random(size=n)
        self.verify(values)


--- ../../tskit/python/tests/test_lowlevel.py ---


"""
Test cases for the low level C interface to tskit.
"""
import collections
import gc
import inspect
import itertools
import os
import random
import tempfile

import msprime
import numpy as np
import pytest

import _tskit
import tskit


NON_UTF8_STRING = "\ud861\udd37"


def get_tracked_sample_counts(ts, st, tracked_samples):
    """
    Returns a list giving the number of samples in the specified list
    that are in the subtree rooted at each node.
    """
    nu = [0 for j in range(ts.get_num_nodes())]
    for j in tracked_samples:
        # Duplicates not permitted.
        assert nu[j] == 0
        u = j
        while u != _tskit.NULL:
            nu[u] += 1
            u = st.get_parent(u)
    return nu


def get_sample_counts(tree_sequence, st):
    """
    Returns a list of the sample node counts for the specified tree.
    """
    nu = [0 for j in range(tree_sequence.get_num_nodes())]
    for j in range(tree_sequence.get_num_samples()):
        u = j
        while u != _tskit.NULL:
            nu[u] += 1
            u = st.get_parent(u)
    return nu


class LowLevelTestCase:
    """
    Superclass of tests for the low-level interface.
    """

    def verify_tree_dict(self, n, pi):
        """
        Verifies that the specified tree in dict format is a
        consistent coalescent history for a sample of size n.
        """
        assert len(pi) <= 2 * n - 1
        # _tskit.NULL should not be a node
        assert _tskit.NULL not in pi
        # verify the root is equal for all samples
        root = 0
        while pi[root] != _tskit.NULL:
            root = pi[root]
        for j in range(n):
            k = j
            while pi[k] != _tskit.NULL:
                k = pi[k]
            assert k == root
        # 0 to n - 1 inclusive should always be nodes
        for j in range(n):
            assert j in pi
        num_children = collections.defaultdict(int)
        for j in pi.keys():
            num_children[pi[j]] += 1
        # nodes 0 to n are samples.
        for j in range(n):
            assert pi[j] != 0
            assert num_children[j] == 0
        # All non-sample nodes should be binary
        for j in pi.keys():
            if j > n:
                assert num_children[j] >= 2

    def get_example_tree_sequence(
        self, sample_size=10, length=1, mutation_rate=1, random_seed=1
    ):
        ts = msprime.simulate(
            sample_size,
            recombination_rate=0.1,
            mutation_rate=mutation_rate,
            random_seed=random_seed,
            length=length,
        )
        return ts.ll_tree_sequence

    def get_example_tree_sequences(self):
        yield self.get_example_tree_sequence()
        yield self.get_example_tree_sequence(2, 10)
        yield self.get_example_tree_sequence(20, 10)
        yield self.get_example_migration_tree_sequence()

    def get_example_migration_tree_sequence(self):
        pop_configs = [msprime.PopulationConfiguration(5) for _ in range(2)]
        migration_matrix = [[0, 1], [1, 0]]
        ts = msprime.simulate(
            population_configurations=pop_configs,
            migration_matrix=migration_matrix,
            mutation_rate=1,
            record_migrations=True,
            random_seed=1,
        )
        return ts.ll_tree_sequence

    def verify_iterator(self, iterator):
        """
        Checks that the specified non-empty iterator implements the
        iterator protocol correctly.
        """
        list_ = list(iterator)
        assert len(list_) > 0
        for _ in range(10):
            with pytest.raises(StopIteration):
                next(iterator)


class MetadataTestMixin:
    metadata_tables = [
        "node",
        "edge",
        "site",
        "mutation",
        "migration",
        "individual",
        "population",
    ]


class TestTableCollection(LowLevelTestCase):
    """
    Tests for the low-level TableCollection class
    """

    def test_skip_tables(self, tmp_path):
        tc = _tskit.TableCollection(1)
        self.get_example_tree_sequence().dump_tables(tc)
        with open(tmp_path / "tmp.trees", "wb") as f:
            tc.dump(f)

        for good_bool in [1, True]:
            with open(tmp_path / "tmp.trees", "rb") as f:
                tc_skip = _tskit.TableCollection()
                tc_skip.load(f, skip_tables=good_bool)
            assert not tc.equals(tc_skip)
            assert tc.equals(tc_skip, ignore_tables=True)

        for bad_bool in ["x", 0.5, {}]:
            with open(tmp_path / "tmp.trees", "rb") as f:
                tc_skip = _tskit.TableCollection()
                with pytest.raises(TypeError):
                    tc_skip.load(f, skip_tables=bad_bool)

    def test_skip_reference_sequence(self, tmp_path):
        tc = _tskit.TableCollection(1)
        self.get_example_tree_sequence().dump_tables(tc)
        tc.reference_sequence.data = "ACGT"
        with open(tmp_path / "tmp.trees", "wb") as f:
            tc.dump(f)

        for good_bool in [1, True]:
            with open(tmp_path / "tmp.trees", "rb") as f:
                tc_skip = _tskit.TableCollection()
                tc_skip.load(f, skip_reference_sequence=good_bool)
            assert not tc.equals(tc_skip)
            assert tc.equals(tc_skip, ignore_reference_sequence=True)

        for bad_bool in ["x", 0.5, {}]:
            with open(tmp_path / "tmp.trees", "rb") as f:
                tc_skip = _tskit.TableCollection()
                with pytest.raises(TypeError):
                    tc_skip.load(f, skip_reference_sequence=bad_bool)

    def test_file_errors(self):
        tc1 = _tskit.TableCollection(1)
        self.get_example_tree_sequence().dump_tables(tc1)

        def loader(*args):
            tc = _tskit.TableCollection(1)
            tc.load(*args)

        for func in [tc1.dump, loader]:
            with pytest.raises(TypeError):
                func()
            for bad_type in [None, [], {}]:
                with pytest.raises(TypeError):
                    func(bad_type)

    def test_file_format_eof_error(self, tmp_path):
        with open(tmp_path / "tmp.trees", "wb") as f:
            f.write(b"")
        with open(tmp_path / "tmp.trees", "rb") as f:
            tc2 = _tskit.TableCollection()
            with pytest.raises(EOFError):
                tc2.load(f)

    def test_file_format_kas_error(self, tmp_path):
        tc1 = _tskit.TableCollection(1)
        self.get_example_tree_sequence().dump_tables(tc1)
        with open(tmp_path / "tmp.trees", "wb") as f:
            tc1.dump(f)
        with open(tmp_path / "tmp.trees", "rb") as f:
            f.seek(1)
            tc2 = _tskit.TableCollection()
            with pytest.raises(_tskit.FileFormatError):
                tc2.load(f)

    def test_dump_equality(self, tmp_path):
        for ts in self.get_example_tree_sequences():
            tc = _tskit.TableCollection(sequence_length=ts.get_sequence_length())
            ts.dump_tables(tc)
            with open(tmp_path / "tmp.trees", "wb") as f:
                tc.dump(f)
            with open(tmp_path / "tmp.trees", "rb") as f:
                tc2 = _tskit.TableCollection()
                tc2.load(f)
            assert tc.equals(tc2)

    def test_reference_deletion(self):
        ts = msprime.simulate(10, mutation_rate=1, random_seed=1)
        tc = ts.tables._ll_tables
        # Get references to all the tables
        tables = [
            tc.individuals,
            tc.nodes,
            tc.edges,
            tc.migrations,
            tc.sites,
            tc.mutations,
            tc.populations,
            tc.provenances,
        ]
        del tc
        for _ in range(10):
            for table in tables:
                assert len(str(table)) > 0

    def test_set_sequence_length_errors(self):
        tables = _tskit.TableCollection(1)
        with pytest.raises(TypeError):
            del tables.sequence_length
        for bad_value in ["sdf", None, []]:
            with pytest.raises(TypeError):
                tables.sequence_length = bad_value

    def test_set_sequence_length(self):
        tables = _tskit.TableCollection(1)
        assert tables.sequence_length == 1
        for value in [-1, 1e6, 1e-22, 1000, 2**32, -10000]:
            tables.sequence_length = value
            assert tables.sequence_length == value

    def test_set_time_units_errors(self):
        tables = _tskit.TableCollection(1)
        with pytest.raises(AttributeError):
            del tables.time_units
        for bad_value in [b"no bytes", 59, 43.4, None, []]:
            with pytest.raises(TypeError):
                tables.time_units = bad_value

    def test_set_time_units(self):
        tables = _tskit.TableCollection(1)
        assert tables.time_units == tskit.TIME_UNITS_UNKNOWN
        for value in ["foo", "", "💩", "null char \0 in string"]:
            tables.time_units = value
            assert tables.time_units == value

    def test_set_metadata_errors(self):
        tables = _tskit.TableCollection(1)
        with pytest.raises(AttributeError):
            del tables.metadata
        for bad_value in ["no bytes", 59, 43.4, None, []]:
            with pytest.raises(TypeError):
                tables.metadata = bad_value

    def test_set_metadata(self):
        tables = _tskit.TableCollection(1)
        assert tables.metadata == b""
        for value in [b"foo", b"", "💩".encode(), b"null char \0 in string"]:
            tables.metadata = value
            tables.metadata_schema = "Test we have two separate fields"
            assert tables.metadata == value

    def test_set_metadata_schema_errors(self):
        tables = _tskit.TableCollection(1)
        with pytest.raises(AttributeError):
            del tables.metadata_schema
        for bad_value in [59, 43.4, None, []]:
            with pytest.raises(TypeError):
                tables.metadata_schema = bad_value

    def test_set_metadata_schema(self):
        tables = _tskit.TableCollection(1)
        assert tables.metadata_schema == ""
        for value in ["foo", "", "💩", "null char \0 in string"]:
            tables.metadata_schema = value
            tables.metadata = b"Test we have two separate fields"
            assert tables.metadata_schema == value

    def test_simplify_bad_args(self):
        ts = msprime.simulate(10, random_seed=1)
        tc = ts.tables._ll_tables
        with pytest.raises(TypeError):
            tc.simplify()
        with pytest.raises(ValueError):
            tc.simplify("asdf")
        with pytest.raises(TypeError):
            tc.simplify([0, 1], keep_unary="sdf")
        with pytest.raises(TypeError):
            tc.simplify([0, 1], keep_unary_in_individuals="abc")
        with pytest.raises(TypeError):
            tc.simplify([0, 1], keep_input_roots="sdf")
        with pytest.raises(TypeError):
            tc.simplify([0, 1], filter_populations="x")
        with pytest.raises(TypeError):
            tc.simplify([0, 1], filter_nodes="x")
        with pytest.raises(TypeError):
            tc.simplify([0, 1], update_sample_flags="x")
        with pytest.raises(_tskit.LibraryError):
            tc.simplify([0, -1])

    @pytest.mark.parametrize("value", [True, False])
    @pytest.mark.parametrize(
        "flag",
        [
            "filter_sites",
            "filter_populations",
            "filter_individuals",
            "filter_nodes",
            "update_sample_flags",
            "reduce_to_site_topology",
            "keep_unary",
            "keep_unary_in_individuals",
            "keep_input_roots",
        ],
    )
    def test_simplify_flags(self, flag, value):
        tables = _tskit.TableCollection(1)
        tables.simplify([], **{flag: value})

    def test_link_ancestors_bad_args(self):
        ts = msprime.simulate(10, random_seed=1)
        tc = ts.tables._ll_tables
        with pytest.raises(TypeError):
            tc.link_ancestors()
        with pytest.raises(TypeError):
            tc.link_ancestors([0, 1])
        with pytest.raises(ValueError):
            tc.link_ancestors(samples=[0, 1], ancestors="sdf")
        with pytest.raises(ValueError):
            tc.link_ancestors(samples="sdf", ancestors=[0, 1])
        with pytest.raises(_tskit.LibraryError):
            tc.link_ancestors(samples=[0, 1], ancestors=[11, -1])
        with pytest.raises(_tskit.LibraryError):
            tc.link_ancestors(samples=[0, -1], ancestors=[11])

    def test_link_ancestors(self):
        ts = msprime.simulate(2, random_seed=1)
        tc = ts.tables._ll_tables
        edges = tc.link_ancestors([0, 1], [3])
        assert isinstance(edges, _tskit.EdgeTable)
        del edges
        assert tc.edges.num_rows == 2

    def test_subset_bad_args(self):
        ts = msprime.simulate(10, random_seed=1)
        tc = ts.tables._ll_tables
        with pytest.raises(TypeError):
            tc.subset(np.array(["a"]))
        with pytest.raises(ValueError):
            tc.subset(np.array([[1], [2]], dtype="int32"))
        with pytest.raises(TypeError):
            tc.subset()
        with pytest.raises(_tskit.LibraryError):
            tc.subset(np.array([100, 200], dtype="int32"))

    def test_union_bad_args(self):
        ts = msprime.simulate(10, random_seed=1)
        tc = ts.tables._ll_tables
        tc2 = tc
        with pytest.raises(TypeError):
            tc.union(tc2, np.array(["a"]))
        with pytest.raises(ValueError):
            tc.union(tc2, np.array([0], dtype="int32"))
        with pytest.raises(TypeError):
            tc.union(tc2)
        with pytest.raises(TypeError):
            tc.union()
        node_mapping = np.arange(ts.num_nodes, dtype="int32")
        node_mapping[0] = 1200
        with pytest.raises(_tskit.LibraryError):
            tc.union(tc2, node_mapping)
        node_mapping = np.array(
            [node_mapping.tolist(), node_mapping.tolist()], dtype="int32"
        )
        with pytest.raises(ValueError):
            tc.union(tc2, node_mapping)
        with pytest.raises(ValueError):
            tc.union(tc2, np.array([[1], [2]], dtype="int32"))

    def test_equals_bad_args(self):
        ts = msprime.simulate(10, random_seed=1242)
        tc = ts.tables._ll_tables
        with pytest.raises(TypeError):
            tc.equals()
        with pytest.raises(TypeError):
            tc.equals(None)
        assert tc.equals(tc)
        with pytest.raises(TypeError):
            tc.equals(tc, no_such_arg=1)
        bad_bool = "x"
        with pytest.raises(TypeError):
            tc.equals(tc, ignore_metadata=bad_bool)
        with pytest.raises(TypeError):
            tc.equals(tc, ignore_ts_metadata=bad_bool)
        with pytest.raises(TypeError):
            tc.equals(tc, ignore_provenance=bad_bool)
        with pytest.raises(TypeError):
            tc.equals(tc, ignore_timestamps=bad_bool)
        with pytest.raises(TypeError):
            tc.equals(tc, ignore_tables=bad_bool)
        with pytest.raises(TypeError):
            tc.equals(tc, ignore_reference_sequence=bad_bool)

    def test_asdict(self):
        for ts in self.get_example_tree_sequences():
            tc = _tskit.TableCollection(sequence_length=ts.get_sequence_length())
            ts.dump_tables(tc)
            d = tc.asdict()
            # Method is tested extensively elsewhere, just basic sanity check here
            assert isinstance(d, dict)
            assert len(d) > 0

    def test_fromdict(self):
        for ts in self.get_example_tree_sequences():
            tc1 = _tskit.TableCollection(sequence_length=ts.get_sequence_length())
            ts.dump_tables(tc1)
            d = tc1.asdict()
            tc2 = _tskit.TableCollection(sequence_length=0)
            tc2.fromdict(d)
            assert tc1.equals(tc2)

    def test_asdict_bad_args(self):
        ts = msprime.simulate(10, random_seed=1242)
        tc = ts.tables._ll_tables
        for bad_type in [None, 0.1, "str"]:
            with pytest.raises(TypeError):
                tc.asdict(force_offset_64=bad_type)

    def test_fromdict_bad_args(self):
        tc = _tskit.TableCollection(0)
        for bad_type in [None, 0.1, "str"]:
            with pytest.raises(TypeError):
                tc.fromdict(bad_type)

    def test_sort_individuals(self):
        tc = _tskit.TableCollection(1)
        tc.sort_individuals()

    def test_delete_older_bad_args(self):
        tc = _tskit.TableCollection(1)
        self.get_example_tree_sequence().dump_tables(tc)
        with pytest.raises(TypeError):
            tc.delete_older()
        with pytest.raises(TypeError):
            tc.delete_older("1234")


class TestIbd:
    def test_uninitialised(self):
        result = _tskit.IdentitySegments.__new__(_tskit.IdentitySegments)
        with pytest.raises(SystemError):
            result.get(0, 1)
        with pytest.raises(SystemError):
            result.print_state()
        with pytest.raises(SystemError):
            result.num_segments
        with pytest.raises(SystemError):
            result.total_span
        with pytest.raises(SystemError):
            result.get_keys()

    def test_get_keys(self):
        ts = msprime.simulate(10, random_seed=1)
        tc = ts.tables._ll_tables
        pairs = [[0, 1], [0, 2], [1, 2]]
        result = tc.ibd_segments_within([0, 1, 2], store_pairs=True)
        np.testing.assert_array_equal(result.get_keys(), pairs)

    def test_store_pairs(self):
        ts = msprime.simulate(10, random_seed=1)
        tc = ts.tables._ll_tables
        # By default we can't get any information about pairs.
        result = tc.ibd_segments_within()
        with pytest.raises(_tskit.IdentityPairsNotStoredError):
            result.get_keys()
        with pytest.raises(_tskit.IdentityPairsNotStoredError):
            result.num_pairs
        with pytest.raises(_tskit.IdentityPairsNotStoredError):
            result.get(0, 1)

        num_pairs = 45
        result = tc.ibd_segments_within(store_pairs=True)
        assert len(result.get_keys()) == num_pairs
        assert result.num_pairs == num_pairs

        seglist = result.get(0, 1)
        assert seglist.num_segments == 1
        assert seglist.total_span == 1
        with pytest.raises(_tskit.IdentitySegmentsNotStoredError):
            seglist.node
        with pytest.raises(_tskit.IdentitySegmentsNotStoredError):
            seglist.left
        with pytest.raises(_tskit.IdentitySegmentsNotStoredError):
            seglist.right

    def test_within_all_pairs(self):
        ts = msprime.simulate(10, random_seed=1)
        tc = ts.tables._ll_tables
        num_pairs = ts.num_samples * (ts.num_samples - 1) / 2
        result = tc.ibd_segments_within(store_pairs=True)
        assert result.num_pairs == num_pairs
        pairs = np.array(list(itertools.combinations(range(ts.num_samples), 2)))
        np.testing.assert_array_equal(result.get_keys(), pairs)

    def test_between_all_pairs(self):
        ts = msprime.simulate(10, random_seed=1)
        tc = ts.tables._ll_tables
        result = tc.ibd_segments_between([5, 5], range(10), store_pairs=True)
        assert result.num_pairs == 25
        pairs = np.array(list(itertools.product(range(5), range(5, 10))))
        np.testing.assert_array_equal(result.get_keys(), pairs)

    def test_within_bad_args(self):
        ts = msprime.simulate(10, random_seed=1)
        tc = ts.tables._ll_tables
        for bad_samples in ["sdf", {}]:
            with pytest.raises(ValueError):
                tc.ibd_segments_within(bad_samples)
        # input array must be 1D
        with pytest.raises(ValueError):
            tc.ibd_segments_within([[[1], [1]]])
        for bad_float in ["sdf", None, {}]:
            with pytest.raises(TypeError):
                tc.ibd_segments_within(min_span=bad_float)
            with pytest.raises(TypeError):
                tc.ibd_segments_within(max_time=bad_float)
        with pytest.raises(_tskit.LibraryError):
            tc.ibd_segments_within(max_time=-1)
        with pytest.raises(_tskit.LibraryError):
            tc.ibd_segments_within(min_span=-1)

    def test_between_bad_args(self):
        ts = msprime.simulate(10, random_seed=1)
        tc = ts.tables._ll_tables
        with pytest.raises(TypeError):
            tc.ibd_segments_between()
        with pytest.raises(TypeError):
            tc.ibd_segments_between([1])

        with pytest.raises(ValueError):
            tc.ibd_segments_between("sdf", [1, 2])
        with pytest.raises(ValueError):
            tc.ibd_segments_between([1, 2], "sdf")
        # The sample_set parsing code is tested elsewhere, so just test
        # something basic.
        with pytest.raises(ValueError, match="Sum of sample_set_sizes"):
            tc.ibd_segments_between([1, 1], [1])
        for bad_float in ["sdf", None, {}]:
            with pytest.raises(TypeError):
                tc.ibd_segments_between([1, 1], [0, 1], min_span=bad_float)
            with pytest.raises(TypeError):
                tc.ibd_segments_between([1, 1], [0, 1], max_time=bad_float)
        with pytest.raises(_tskit.LibraryError):
            tc.ibd_segments_between([1, 1], [0, 1], min_span=-1)
        with pytest.raises(_tskit.LibraryError):
            tc.ibd_segments_between([1, 1], [0, 1], max_time=-1)
        with pytest.raises(_tskit.LibraryError, match="Duplicate sample"):
            tc.ibd_segments_between([1, 1], [0, 0])

    def test_get_output(self):
        ts = msprime.simulate(5, random_seed=1)
        tc = ts.tables._ll_tables
        pairs = [(0, 1), (2, 3)]
        result = tc.ibd_segments_within([0, 1, 2, 3], store_segments=True)
        assert isinstance(result, _tskit.IdentitySegments)
        for pair in pairs:
            value = result.get(*pair)
            assert isinstance(value, _tskit.IdentitySegmentList)
            assert value.num_segments == 1
            assert isinstance(value.left, np.ndarray)
            assert isinstance(value.right, np.ndarray)
            assert isinstance(value.node, np.ndarray)
            assert list(value.left) == [0]
            assert list(value.right) == [1]
            assert len(value.node) == 1
            assert value.num_segments == 1
            assert value.total_span == 1

    def test_get_bad_args(self):
        ts = msprime.simulate(10, random_seed=1)
        tc = ts.tables._ll_tables
        result = tc.ibd_segments_within([0, 1, 2], store_segments=True)
        with pytest.raises(TypeError):
            result.get()
        with pytest.raises(TypeError):
            result.get("0", 1)
        with pytest.raises(_tskit.LibraryError, match="Both nodes"):
            result.get(0, 0)
        with pytest.raises(_tskit.LibraryError, match="Node out of bounds"):
            result.get(-1, 0)
        with pytest.raises(_tskit.LibraryError, match="Node out of bounds"):
            result.get(0, 100)
        with pytest.raises(KeyError):
            result.get(0, 3)

    def test_print_state(self):
        ts = msprime.simulate(10, random_seed=1)
        tc = ts.tables._ll_tables
        result = tc.ibd_segments_within()
        with pytest.raises(TypeError):
            result.print_state()

        with tempfile.TemporaryFile("w+") as f:
            result.print_state(f)
            f.seek(0)
            output = f.read()
        assert len(output) > 0
        assert "IBD" in output

    def test_direct_instantiation(self):
        # Nobody should do this, but just in case
        result = _tskit.IdentitySegments()
        assert result.num_segments == 0
        assert result.total_span == 0
        with tempfile.TemporaryFile("w+") as f:
            result.print_state(f)
            f.seek(0)
            output = f.read()
        assert len(output) > 0
        assert "IBD" in output


class TestIdentitySegmentList:
    def test_direct_instantiation(self):
        # Nobody should do this, but just in case
        seglist = _tskit.IdentitySegmentList()
        attrs = ["num_segments", "total_span", "left", "right", "node"]
        for attr in attrs:
            with pytest.raises(SystemError, match="not initialised"):
                getattr(seglist, attr)

    def test_memory_management_within(self):
        ts = msprime.simulate(10, random_seed=1)
        tc = ts.tables._ll_tables
        result = tc.ibd_segments_within(store_segments=True)
        del ts, tc
        lst = result.get(0, 1)
        assert lst.num_segments == 1
        del result
        gc.collect()
        assert lst.num_segments == 1
        # Do some allocs to see if we're still working properly
        x = sum(list(range(1000)))
        assert x > 0
        assert lst.num_segments == 1

    def test_memory_management_between(self):
        ts = msprime.simulate(10, random_seed=1)
        tc = ts.tables._ll_tables
        result = tc.ibd_segments_between([2, 2], range(4), store_segments=True)
        del ts, tc
        lst = result.get(0, 2)
        assert lst.num_segments == 1
        del result
        gc.collect()
        assert lst.num_segments == 1
        # Do some allocs to see if we're still working properly
        x = sum(list(range(1000)))
        assert x > 0
        assert lst.num_segments == 1


class TestTableMethods:
    """
    Tests for the low-level table methods.
    """

    @pytest.mark.parametrize("table_name", tskit.TABLE_NAMES)
    def test_table_extend(self, table_name, ts_fixture):
        table = getattr(ts_fixture.tables, table_name)
        assert len(table) >= 5
        ll_table = table.ll_table
        table_copy = table.copy()

        ll_table.extend(table_copy.ll_table, row_indexes=[])
        assert table == table_copy

        ll_table.clear()
        ll_table.extend(table_copy.ll_table, row_indexes=range(len(table_copy)))
        assert table == table_copy

    @pytest.mark.parametrize("table_name", tskit.TABLE_NAMES)
    @pytest.mark.parametrize(
        ["row_indexes", "expected_rows"],
        [
            ([0], [0]),
            ([4] * 1000, [4] * 1000),
            ([4, 1, 3, 0, 0], [4, 1, 3, 0, 0]),
            (np.array([0, 1, 4], dtype=np.uint8), [0, 1, 4]),
            (np.array([3, 3, 3], dtype=np.uint16), [3, 3, 3]),
            (np.array([4, 2, 1], dtype=np.int8), [4, 2, 1]),
            (np.array([4, 2], dtype=np.int16), [4, 2]),
            (np.array([0, 1], dtype=np.int32), [0, 1]),
            (range(2, -1, -1), [2, 1, 0]),
        ],
    )
    def test_table_extend_types(
        self, ts_fixture, table_name, row_indexes, expected_rows
    ):
        table = getattr(ts_fixture.tables, table_name)
        assert len(table) >= 5
        ll_table = table.ll_table
        table_copy = table.copy()

        ll_table.extend(table_copy.ll_table, row_indexes=row_indexes)
        assert len(table) == len(table_copy) + len(expected_rows)
        for i, expected_row in enumerate(expected_rows):
            assert table[len(table_copy) + i] == table_copy[expected_row]

    @pytest.mark.parametrize("table_name", tskit.TABLE_NAMES)
    def test_table_keep_rows_errors(self, table_name, ts_fixture):
        table = getattr(ts_fixture.tables, table_name)
        n = len(table)
        ll_table = table.ll_table
        with pytest.raises(ValueError, match="must be of length"):
            ll_table.keep_rows(np.ones(n - 1, dtype=bool))
        with pytest.raises(ValueError, match="must be of length"):
            ll_table.keep_rows(np.ones(n + 1, dtype=bool))
        with pytest.raises(TypeError, match="Cannot cast"):
            ll_table.keep_rows(np.ones(n, dtype=int))

    @pytest.mark.parametrize("table_name", tskit.TABLE_NAMES)
    def test_table_keep_rows_all(self, table_name, ts_fixture):
        table = getattr(ts_fixture.tables, table_name)
        n = len(table)
        ll_table = table.ll_table
        a = ll_table.keep_rows(np.ones(n, dtype=bool))
        assert ll_table.num_rows == n
        assert a.shape == (n,)
        assert a.dtype == np.int32
        assert np.all(a == np.arange(n))

    @pytest.mark.parametrize("table_name", tskit.TABLE_NAMES)
    def test_table_keep_rows_none(self, table_name, ts_fixture):
        table = getattr(ts_fixture.tables, table_name)
        n = len(table)
        ll_table = table.ll_table
        a = ll_table.keep_rows(np.zeros(n, dtype=bool))
        assert ll_table.num_rows == 0
        assert a.shape == (n,)
        assert a.dtype == np.int32
        assert np.all(a == -1)

    def test_mutation_table_keep_rows_ref_error(self):
        table = _tskit.MutationTable()
        table.add_row(site=0, node=0, derived_state="A", parent=2)
        with pytest.raises(_tskit.LibraryError, match="TSK_ERR_MUTATION_OUT_OF_BOUNDS"):
            table.keep_rows([True])

    def test_individual_table_keep_rows_ref_error(self):
        table = _tskit.IndividualTable()
        table.add_row(parents=[2])
        with pytest.raises(
            _tskit.LibraryError, match="TSK_ERR_INDIVIDUAL_OUT_OF_BOUNDS"
        ):
            table.keep_rows([True])

    @pytest.mark.parametrize(
        ["table_name", "column_name"],
        [
            (t, c)
            for t in tskit.TABLE_NAMES
            for c in getattr(tskit, f"{t[:-1].capitalize()}Table").column_names
            if c[-7:] != "_offset"
        ],
    )
    def test_table_update(self, ts_fixture, table_name, column_name):
        table = getattr(ts_fixture.tables, table_name)
        copy = table.copy()
        ll_table = table.ll_table

        # Find the first row where this column differs to get a value to swap in
        other_row_index = -1
        for i, row in enumerate(table):
            if not np.array_equal(
                getattr(table[0], column_name), getattr(row, column_name)
            ):
                other_row_index = i
        assert other_row_index != -1

        # No-op update should not create a change
        args = ll_table.get_row(0)
        ll_table.update_row(0, *args)
        table.assert_equals(copy)

        # Modify the column under test in the first row
        new_args = list(ll_table.get_row(0))
        arg_index = list(inspect.signature(table.add_row).parameters.keys()).index(
            column_name
        )
        new_args[arg_index] = ll_table.get_row(other_row_index)[arg_index]
        ll_table.update_row(0, *new_args)
        for a, b in zip(ll_table.get_row(0), new_args):
            np.array_equal(a, b)

    def test_update_defaults(self):
        t = tskit.IndividualTable()
        assert t.add_row(flags=1, location=[1, 2], parents=[3, 4], metadata=b"FOO") == 0
        t.ll_table.update_row(0)
        assert t.flags[0] == 0
        assert len(t.location) == 0
        assert t.location_offset[0] == 0
        assert len(t.parents) == 0
        assert t.parents_offset[0] == 0
        assert len(t.metadata) == 0
        assert t.metadata_offset[0] == 0

        t = tskit.NodeTable()
        assert (
            t.add_row(flags=1, time=2, population=3, individual=4, metadata=b"FOO") == 0
        )
        t.ll_table.update_row(0)
        assert t.time[0] == 0
        assert t.flags[0] == 0
        assert t.population[0] == tskit.NULL
        assert t.individual[0] == tskit.NULL
        assert len(t.metadata) == 0
        assert t.metadata_offset[0] == 0

        t = tskit.EdgeTable()
        assert t.add_row(1, 2, 3, 4, metadata=b"FOO") == 0
        t.ll_table.update_row(0, 1, 2, 3, 4)
        assert len(t.metadata) == 0
        assert t.metadata_offset[0] == 0

        t = tskit.MigrationTable()
        assert t.add_row(1, 2, 3, 4, 5, 6, b"FOO") == 0
        t.ll_table.update_row(0, 1, 2, 3, 4, 5, 6)
        assert len(t.metadata) == 0
        assert t.metadata_offset[0] == 0

        t = tskit.MutationTable()
        assert t.add_row(1, 2, "A", 3, b"FOO", 4) == 0
        t.ll_table.update_row(0, 1, 2, "A", 3)
        assert len(t.metadata) == 0
        assert t.metadata_offset[0] == 0
        assert tskit.is_unknown_time(t.time[0])

        t = tskit.PopulationTable()
        assert t.add_row(b"FOO") == 0
        t.ll_table.update_row(0)
        assert len(t.metadata) == 0
        assert t.metadata_offset[0] == 0

    def test_update_bad_data(self):
        t = tskit.IndividualTable()
        t.add_row()
        with pytest.raises(TypeError):
            t.ll_table.update_row(0, flags="x")
        with pytest.raises(TypeError):
            t.ll_table.update_row(0, metadata=123)
        with pytest.raises(ValueError):
            t.ll_table.update_row(0, location="1234")
        with pytest.raises(ValueError):
            t.ll_table.update_row(0, parents="forty-two")

        t = tskit.NodeTable()
        t.add_row()
        with pytest.raises(TypeError):
            t.ll_table.update_row(0, flags="x")
        with pytest.raises(TypeError):
            t.ll_table.update_row(0, time="x")
        with pytest.raises(TypeError):
            t.ll_table.update_row(0, individual="x")
        with pytest.raises(TypeError):
            t.ll_table.update_row(0, population="x")
        with pytest.raises(TypeError):
            t.ll_table.update_row(0, metadata=123)

        t = tskit.EdgeTable()
        t.add_row(1, 2, 3, 4)
        with pytest.raises(TypeError):
            t.ll_table.update_row(0, left="x", right=0, parent=0, child=0)
        with pytest.raises(TypeError):
            t.ll_table.update_row(
                0,
            )
        with pytest.raises(TypeError):
            t.ll_table.update_row(0, 0, 0, 0, 0, metadata=123)

        t = tskit.SiteTable()
        t.add_row(0, "A")
        with pytest.raises(TypeError):
            t.ll_table.update_row(0, "x", "A")
        with pytest.raises(TypeError):
            t.ll_table.update_row(0, 0, 0)
        with pytest.raises(TypeError):
            t.ll_table.update_row(0, 0, "A", metadata=[0, 1, 2])

        t = tskit.MutationTable()
        t.add_row(0, 0, "A")
        with pytest.raises(TypeError):
            t.ll_table.update_row(0, "0", 0, "A")
        with pytest.raises(TypeError):
            t.ll_table.update_row(0, 0, "0", "A")
        with pytest.raises(TypeError):
            t.ll_table.update_row(0, 0, 0, "A", parent=None)
        with pytest.raises(TypeError):
            t.ll_table.update_row(0, 0, 0, "A", metadata=[0])
        with pytest.raises(TypeError):
            t.ll_table.update_row(0, 0, 0, "A", time="A")

        t = tskit.MigrationTable()
        with pytest.raises(TypeError):
            t.add_row(left="x", right=0, node=0, source=0, dest=0, time=0)
        with pytest.raises(TypeError):
            t.ll_table.update_row(
                0,
            )
        with pytest.raises(TypeError):
            t.ll_table.update_row(0, 0, 0, 0, 0, 0, 0, metadata=123)

        t = tskit.ProvenanceTable()
        t.add_row("a", "b")
        with pytest.raises(TypeError):
            t.ll_table.update_row(0, 0, "b")
        with pytest.raises(TypeError):
            t.ll_table.update_row(0, "a", 0)

        t = tskit.PopulationTable()
        t.add_row()
        with pytest.raises(TypeError):
            t.ll_table.update_row(0, metadata=[0])


class TestTableMethodsErrors:
    """
    Tests for the error handling of errors in the low-level tables.
    """

    def yield_tables(self, ts):
        for table in ts.tables.table_name_map.values():
            yield table.ll_table

    @pytest.mark.parametrize(
        "table_name",
        tskit.TABLE_NAMES,
    )
    def test_table_extend_bad_args(self, ts_fixture, table_name):
        table = getattr(ts_fixture.tables, table_name)
        ll_table = table.ll_table
        ll_table_copy = table.copy().ll_table

        with pytest.raises(
            _tskit.LibraryError,
            match="Tables can only be extended using rows from a different table",
        ):
            ll_table.extend(ll_table, row_indexes=[])
        with pytest.raises(TypeError):
            ll_table.extend(None, row_indexes=[])
        with pytest.raises(ValueError):
            ll_table.extend(ll_table_copy, row_indexes=5)
        with pytest.raises(TypeError):
            ll_table.extend(ll_table_copy, row_indexes=[None])
        with pytest.raises(ValueError, match="object too deep"):
            ll_table.extend(ll_table_copy, row_indexes=[[0, 1], [2, 3]])
        with pytest.raises(ValueError, match="object too deep"):
            ll_table.extend(ll_table_copy, row_indexes=[[0, 1]])
        with pytest.raises(_tskit.LibraryError, match="out of bounds"):
            ll_table.extend(ll_table_copy, row_indexes=[-1])
        with pytest.raises(_tskit.LibraryError, match="out of bounds"):
            ll_table.extend(ll_table_copy, row_indexes=[1000])
        with pytest.raises(_tskit.LibraryError, match="out of bounds"):
            ll_table.extend(ll_table_copy, row_indexes=range(10000000, 10000001))

        # Uncastable types
        for dtype in [np.uint32, np.int64, np.uint64, np.float32, np.float64]:
            with pytest.raises(TypeError, match="Cannot cast"):
                ll_table.extend(ll_table_copy, row_indexes=np.array([0], dtype=dtype))

    @pytest.mark.parametrize("table_name", tskit.TABLE_NAMES)
    def test_update_bad_row_index(self, ts_fixture, table_name):
        table = getattr(ts_fixture.tables, table_name)
        ll_table = table.ll_table
        row_data = ll_table.get_row(0)
        with pytest.raises(_tskit.LibraryError, match="out of bounds"):
            ll_table.update_row(-1, *row_data)
        with pytest.raises(ValueError, match="tskit ids must be"):
            ll_table.update_row(-42, *row_data)
        with pytest.raises(TypeError):
            ll_table.update_row([], *row_data)
        with pytest.raises(TypeError):
            ll_table.update_row("abc", *row_data)
        with pytest.raises(_tskit.LibraryError, match="out of bounds"):
            ll_table.update_row(10000, *row_data)
        with pytest.raises(OverflowError, match="Value too large for tskit id type"):
            ll_table.update_row(2**62, *row_data)

    def test_equals_bad_args(self, ts_fixture):
        for ll_table in self.yield_tables(ts_fixture):
            assert ll_table.equals(ll_table)
            with pytest.raises(TypeError):
                ll_table.equals(None)
            with pytest.raises(TypeError):
                ll_table.equals(ll_table, no_such_arg="")
            uninit_other = type(ll_table).__new__(type(ll_table))
            with pytest.raises(SystemError):
                ll_table.equals(uninit_other)

    def test_get_row_bad_args(self, ts_fixture):
        for ll_table in self.yield_tables(ts_fixture):
            assert ll_table.get_row(0) is not None
            with pytest.raises(TypeError):
                ll_table.get_row(no_such_arg="")

    @pytest.mark.parametrize("table", ["nodes", "individuals"])
    def test_flag_underflow_overflow(self, table):
        tables = _tskit.TableCollection(1)
        table = getattr(tables, table)
        table.add_row(flags=0)
        table.add_row(flags=(1 << 32) - 1)
        with pytest.raises(OverflowError, match="unsigned int32 >= than 2\\^32"):
            table.add_row(flags=1 << 32)
        with pytest.raises(OverflowError, match="int too big to convert"):
            table.add_row(flags=1 << 64)
        with pytest.raises(OverflowError, match="int too big to convert"):
            table.add_row(flags=1 << 256)
        with pytest.raises(
            ValueError, match="Can't convert negative value to unsigned int"
        ):
            table.add_row(flags=-1)

    def test_index(self):
        tc = msprime.simulate(10, random_seed=42).tables._ll_tables
        assert tc.indexes["edge_insertion_order"].dtype == np.int32
        assert tc.indexes["edge_removal_order"].dtype == np.int32
        assert np.array_equal(
            tc.indexes["edge_insertion_order"], np.arange(18, dtype=np.int32)
        )
        assert np.array_equal(
            tc.indexes["edge_removal_order"], np.arange(18, dtype=np.int32)[::-1]
        )
        tc.drop_index()
        assert tc.indexes == {}
        tc.build_index()
        assert np.array_equal(
            tc.indexes["edge_insertion_order"], np.arange(18, dtype=np.int32)
        )
        assert np.array_equal(
            tc.indexes["edge_removal_order"], np.arange(18, dtype=np.int32)[::-1]
        )

        modify_indexes = tc.indexes
        modify_indexes["edge_insertion_order"] = np.arange(42, 42 + 18, dtype=np.int32)
        modify_indexes["edge_removal_order"] = np.arange(
            4242, 4242 + 18, dtype=np.int32
        )
        tc.indexes = modify_indexes
        assert np.array_equal(
            tc.indexes["edge_insertion_order"], np.arange(42, 42 + 18, dtype=np.int32)
        )
        assert np.array_equal(
            tc.indexes["edge_removal_order"], np.arange(4242, 4242 + 18, dtype=np.int32)
        )

    def test_no_indexes(self):
        tc = msprime.simulate(10, random_seed=42).tables._ll_tables
        tc.drop_index()
        assert tc.indexes == {}

    def test_bad_indexes(self):
        tc = msprime.simulate(10, random_seed=42).tables._ll_tables
        for col in ("insertion", "removal"):
            d = tc.indexes
            d[f"edge_{col}_order"] = d[f"edge_{col}_order"][:-1]
            with pytest.raises(
                ValueError,
                match="^edge_insertion_order and"
                " edge_removal_order must be the same"
                " length$",
            ):
                tc.indexes = d
        d = tc.indexes
        for col in ("insertion", "removal"):
            d[f"edge_{col}_order"] = d[f"edge_{col}_order"][:-1]
        with pytest.raises(
            ValueError,
            match="^edge_insertion_order and edge_removal_order must be"
            " the same length as the number of edges$",
        ):
            tc.indexes = d

        # Both columns must be provided, if one is
        for col in ("insertion", "removal"):
            d = tc.indexes
            del d[f"edge_{col}_order"]
            with pytest.raises(
                TypeError,
                match="^edge_insertion_order and "
                "edge_removal_order must be specified "
                "together$",
            ):
                tc.indexes = d

        tc = msprime.simulate(
            10, recombination_rate=10, random_seed=42
        ).tables._ll_tables
        modify_indexes = tc.indexes
        shape = modify_indexes["edge_insertion_order"].shape
        modify_indexes["edge_insertion_order"] = np.zeros(shape, dtype=np.int32)
        modify_indexes["edge_removal_order"] = np.zeros(shape, dtype=np.int32)
        tc.indexes = modify_indexes
        ts = _tskit.TreeSequence()
        with pytest.raises(
            _tskit.LibraryError,
            match="TSK_ERR_TABLES_BAD_INDEXES",
        ):
            ts.load_tables(tc, build_indexes=False)

        modify_indexes["edge_insertion_order"] = np.full(shape, 2**30, dtype=np.int32)
        modify_indexes["edge_removal_order"] = np.full(shape, 2**30, dtype=np.int32)
        tc.indexes = modify_indexes
        ts = _tskit.TreeSequence()
        with pytest.raises(_tskit.LibraryError, match="Edge out of bounds"):
            ts.load_tables(tc, build_indexes=False)


class TestTreeSequence(LowLevelTestCase, MetadataTestMixin):
    """
    Tests for the low-level interface for the TreeSequence.
    """

    ARRAY_NAMES = [
        "individuals_flags",
        "nodes_time",
        "nodes_flags",
        "nodes_population",
        "nodes_individual",
        "edges_left",
        "edges_right",
        "edges_parent",
        "edges_child",
        "sites_position",
        "mutations_site",
        "mutations_node",
        "mutations_parent",
        "mutations_time",
        "migrations_left",
        "migrations_right",
        "migrations_node",
        "migrations_source",
        "migrations_dest",
        "migrations_time",
        "indexes_edge_insertion_order",
        "indexes_edge_removal_order",
    ]

    def setUp(self):
        fd, self.temp_file = tempfile.mkstemp(prefix="msp_ll_ts_")
        os.close(fd)

    def tearDown(self):
        os.unlink(self.temp_file)

    def test_skip_tables(self, tmp_path):
        ts = self.get_example_tree_sequence()
        with open(tmp_path / "tmp.trees", "wb") as f:
            ts.dump(f)
        tc = _tskit.TableCollection(1)
        ts.dump_tables(tc)

        for good_bool in [1, True]:
            with open(tmp_path / "tmp.trees", "rb") as f:
                ts_skip = _tskit.TreeSequence()
                ts_skip.load(f, skip_tables=good_bool)
            tc_skip = _tskit.TableCollection()
            ts_skip.dump_tables(tc_skip)
            assert not tc.equals(tc_skip)
            assert tc.equals(tc_skip, ignore_tables=True)

        for bad_bool in ["x", 0.5, {}]:
            with open(tmp_path / "tmp.trees", "rb") as f:
                ts_skip = _tskit.TreeSequence()
                with pytest.raises(TypeError):
                    ts_skip.load(f, skip_tables=bad_bool)

    def test_skip_reference_sequence(self, tmp_path):
        tc = _tskit.TableCollection(1)
        self.get_example_tree_sequence().dump_tables(tc)
        tc.reference_sequence.data = "ACGT"
        ts = _tskit.TreeSequence()
        ts.load_tables(tc, build_indexes=True)
        with open(tmp_path / "tmp.trees", "wb") as f:
            ts.dump(f)

        for good_bool in [1, True]:
            with open(tmp_path / "tmp.trees", "rb") as f:
                ts_skip = _tskit.TreeSequence()
                ts_skip.load(f, skip_reference_sequence=good_bool)
            tc_skip = _tskit.TableCollection()
            ts_skip.dump_tables(tc_skip)
            assert not tc.equals(tc_skip)
            assert tc.equals(tc_skip, ignore_reference_sequence=True)

        for bad_bool in ["x", 0.5, {}]:
            with open(tmp_path / "tmp.trees", "rb") as f:
                ts_skip = _tskit.TreeSequence()
                with pytest.raises(TypeError):
                    ts_skip.load(f, skip_reference_sequence=bad_bool)

    def test_file_errors(self):
        ts1 = self.get_example_tree_sequence()

        def loader(*args):
            ts2 = _tskit.TreeSequence()
            ts2.load(*args)

        for func in [ts1.dump, loader]:
            with pytest.raises(TypeError):
                func()
            for bad_type in [None, [], {}]:
                with pytest.raises(TypeError):
                    func(bad_type)

    def test_initial_state(self):
        # Check the initial state to make sure that it is empty.
        ts = _tskit.TreeSequence()
        with pytest.raises(ValueError):
            ts.get_num_samples()
        with pytest.raises(ValueError):
            ts.get_sequence_length()
        with pytest.raises(ValueError):
            ts.get_num_trees()
        with pytest.raises(ValueError):
            ts.get_num_edges()
        with pytest.raises(ValueError):
            ts.get_num_mutations()
        with pytest.raises(ValueError):
            ts.get_num_migrations()
        with pytest.raises(ValueError):
            ts.get_num_migrations()
        with pytest.raises(ValueError):
            ts.dump()

    def test_num_nodes(self):
        for ts in self.get_example_tree_sequences():
            max_node = 0
            for j in range(ts.get_num_edges()):
                _, _, parent, child, _ = ts.get_edge(j)
                for node in [parent, child]:
                    if node > max_node:
                        max_node = node
            assert max_node + 1 == ts.get_num_nodes()

    def test_dump_equality(self, tmp_path):
        for ts in self.get_example_tree_sequences():
            tables = _tskit.TableCollection(sequence_length=ts.get_sequence_length())
            ts.dump_tables(tables)
            tables.compute_mutation_times()
            ts = _tskit.TreeSequence()
            ts.load_tables(tables)
            with open(tmp_path / "temp.trees", "wb") as f:
                ts.dump(f)
            with open(tmp_path / "temp.trees", "rb") as f:
                ts2 = _tskit.TreeSequence()
                ts2.load(f)
            tc = _tskit.TableCollection(ts.get_sequence_length())
            ts.dump_tables(tc)
            tc2 = _tskit.TableCollection(ts2.get_sequence_length())
            ts2.dump_tables(tc2)
            assert tc.equals(tc2)

    def test_get_mutation_interface(self):
        for ts in self.get_example_tree_sequences():
            mutations = [ts.get_mutation(j) for j in range(ts.get_num_mutations())]
            assert len(mutations) == ts.get_num_mutations()
            # Check the form of the mutations
            for packed in mutations:
                site, node, derived_state, parent, metadata, time, edge = packed
                assert isinstance(site, int)
                assert isinstance(node, int)
                assert isinstance(derived_state, str)
                assert isinstance(parent, int)
                assert isinstance(metadata, bytes)
                assert isinstance(time, float)
                assert isinstance(edge, int)

    def test_get_edge_interface(self):
        for ts in self.get_example_tree_sequences():
            num_edges = ts.get_num_edges()
            # We don't accept Python negative indexes here.
            with pytest.raises(IndexError):
                ts.get_edge(-1)
            for j in [0, 10, 10**6]:
                with pytest.raises(IndexError):
                    ts.get_edge(num_edges + j)
            for x in [None, "", {}, []]:
                with pytest.raises(TypeError):
                    ts.get_edge(x)

    def test_get_node_interface(self):
        for ts in self.get_example_tree_sequences():
            num_nodes = ts.get_num_nodes()
            # We don't accept Python negative indexes here.
            with pytest.raises(IndexError):
                ts.get_node(-1)
            for j in [0, 10, 10**6]:
                with pytest.raises(IndexError):
                    ts.get_node(num_nodes + j)
            for x in [None, "", {}, []]:
                with pytest.raises(TypeError):
                    ts.get_node(x)

    def test_get_migration_interface(self):
        ts = self.get_example_migration_tree_sequence()
        for bad_type in ["", None, {}]:
            with pytest.raises(TypeError):
                ts.get_migration(bad_type)
        num_records = ts.get_num_migrations()
        # We don't accept Python negative indexes here.
        with pytest.raises(IndexError):
            ts.get_migration(-1)
        for j in [0, 10, 10**6]:
            with pytest.raises(IndexError):
                ts.get_migration(num_records + j)

    def test_get_samples(self):
        for ts in self.get_example_tree_sequences():
            # get_samples takes no arguments.
            with pytest.raises(TypeError):
                ts.get_samples(0)
            assert np.array_equal(
                np.arange(ts.get_num_samples(), dtype=np.int32), ts.get_samples()
            )

    def test_genealogical_nearest_neighbours(self):
        for ts in self.get_example_tree_sequences():
            with pytest.raises(TypeError):
                ts.genealogical_nearest_neighbours()
            with pytest.raises(TypeError):
                ts.genealogical_nearest_neighbours(focal=None)
            with pytest.raises(TypeError):
                ts.genealogical_nearest_neighbours(
                    focal=ts.get_samples(),
                    reference_sets={},
                )
            with pytest.raises(ValueError):
                ts.genealogical_nearest_neighbours(
                    focal=ts.get_samples(),
                    reference_sets=[],
                )

            bad_array_values = ["", {}, "x", [[[0], [1, 2]]]]
            for bad_array_value in bad_array_values:
                with pytest.raises(ValueError):
                    ts.genealogical_nearest_neighbours(
                        focal=bad_array_value,
                        reference_sets=[[0], [1]],
                    )
                with pytest.raises(ValueError):
                    ts.genealogical_nearest_neighbours(
                        focal=ts.get_samples(),
                        reference_sets=[[0], bad_array_value],
                    )
                with pytest.raises(ValueError):
                    ts.genealogical_nearest_neighbours(
                        focal=ts.get_samples(),
                        reference_sets=[bad_array_value],
                    )
            focal = ts.get_samples()
            A = ts.genealogical_nearest_neighbours(focal, [focal[2:], focal[:2]])
            assert A.shape == (len(focal), 2)

    def test_mean_descendants(self):
        for ts in self.get_example_tree_sequences():
            with pytest.raises(TypeError):
                ts.mean_descendants()
            with pytest.raises(TypeError):
                ts.mean_descendants(reference_sets={})
            with pytest.raises(ValueError):
                ts.mean_descendants(reference_sets=[])

            bad_array_values = ["", {}, "x", [[[0], [1, 2]]]]
            for bad_array_value in bad_array_values:
                with pytest.raises(ValueError):
                    ts.mean_descendants(
                        reference_sets=[[0], bad_array_value],
                    )
                with pytest.raises(ValueError):
                    ts.mean_descendants(reference_sets=[bad_array_value])
            focal = ts.get_samples()
            A = ts.mean_descendants([focal[2:], focal[:2]])
            assert A.shape == (ts.get_num_nodes(), 2)

    def test_metadata_schemas(self):
        tables = _tskit.TableCollection(1.0)
        # Set the schema
        for table_name in self.metadata_tables:
            table = getattr(tables, f"{table_name}s")
            table.metadata_schema = f"{table_name} test metadata schema"
        # Read back via ll tree sequence
        tables.build_index()
        ts = _tskit.TreeSequence()
        ts.load_tables(tables)
        schemas = ts.get_table_metadata_schemas()
        for table_name in self.metadata_tables:
            assert getattr(schemas, table_name) == f"{table_name} test metadata schema"
        # Clear and read back again
        for table_name in self.metadata_tables:
            getattr(tables, f"{table_name}s").metadata_schema = ""
        ts = _tskit.TreeSequence()
        ts.load_tables(tables)
        schemas = ts.get_table_metadata_schemas()
        for table_name in self.metadata_tables:
            assert getattr(schemas, table_name) == ""

    def test_metadata(self):
        tables = _tskit.TableCollection(1)
        tables.build_index()
        ts = _tskit.TreeSequence()
        ts.load_tables(tables)
        assert ts.get_metadata() == b""
        for value in [b"foo", b"", "💩".encode(), b"null char \0 in string"]:
            tables.metadata = value
            ts = _tskit.TreeSequence()
            ts.load_tables(tables)
            assert ts.get_metadata() == value

    def test_metadata_schema(self):
        tables = _tskit.TableCollection(1)
        tables.build_index()
        ts = _tskit.TreeSequence()
        ts.load_tables(tables)
        assert ts.get_metadata_schema() == ""
        for value in ["foo", "", "💩", "null char \0 in string"]:
            tables.metadata_schema = value
            ts = _tskit.TreeSequence()
            ts.load_tables(tables)
            assert ts.get_metadata_schema() == value

    def test_time_units(self):
        tables = _tskit.TableCollection(1)
        tables.build_index()
        ts = _tskit.TreeSequence()
        ts.load_tables(tables)
        assert ts.get_time_units() == tskit.TIME_UNITS_UNKNOWN
        for value in ["foo", "", "💩", "null char \0 in string"]:
            tables.time_units = value
            ts = _tskit.TreeSequence()
            ts.load_tables(tables)
            assert ts.get_time_units() == value

    def test_extend_haplotypes(self):
        ts = self.get_example_tree_sequence(6)
        ets2 = ts.extend_haplotypes(2)
        ets4 = ts.extend_haplotypes(4)
        assert ets2.get_num_nodes() == ts.get_num_nodes()
        assert ets4.get_num_nodes() == ts.get_num_nodes()

    def test_extend_haplotypes_bad_args(self):
        ts1 = self.get_example_tree_sequence(10)
        with pytest.raises(TypeError):
            ts1.extend_haplotypes()
        with pytest.raises(TypeError, match="an integer"):
            ts1.extend_haplotypes("sdf")
        with pytest.raises(_tskit.LibraryError, match="positive"):
            ts1.extend_haplotypes(0)
        with pytest.raises(_tskit.LibraryError, match="positive"):
            ts1.extend_haplotypes(-1)
        tsm = self.get_example_migration_tree_sequence()
        with pytest.raises(
            _tskit.LibraryError, match="TSK_ERR_MIGRATIONS_NOT_SUPPORTED"
        ):
            tsm.extend_haplotypes(1)

    @pytest.mark.parametrize(
        "stat_method_name",
        [
            "D_matrix",
            "D2_matrix",
            "r2_matrix",
            "D_prime_matrix",
            "r_matrix",
            "Dz_matrix",
            "pi2_matrix",
            "D2_unbiased_matrix",
            "Dz_unbiased_matrix",
            "pi2_unbiased_matrix",
        ],
    )
    def test_ld_matrix(self, stat_method_name):
        ts = self.get_example_tree_sequence(10)
        stat_method = getattr(ts, stat_method_name)

        ss = ts.get_samples()  # sample sets
        ss_sizes = np.array([len(ss)], dtype=np.uint32)
        row_sites = np.arange(ts.get_num_sites(), dtype=np.int32)
        col_sites = row_sites
        row_pos = ts.get_breakpoints()[:-1]
        col_pos = row_pos
        row_pos_list = list(map(float, ts.get_breakpoints()[:-1]))
        col_pos_list = row_pos_list
        row_sites_list = list(range(ts.get_num_sites()))
        col_sites_list = row_sites_list

        # happy path
        a = stat_method(ss_sizes, ss, row_sites, col_sites, None, None, "site")
        assert a.shape == (10, 10, 1)
        a = stat_method(
            ss_sizes, ss, row_sites_list, col_sites_list, None, None, "site"
        )
        assert a.shape == (10, 10, 1)
        a = stat_method(ss_sizes, ss, None, None, None, None, "site")
        assert a.shape == (10, 10, 1)

        a = stat_method(ss_sizes, ss, None, None, row_pos, col_pos, "branch")
        assert a.shape == (2, 2, 1)
        a = stat_method(ss_sizes, ss, None, None, row_pos_list, col_pos_list, "branch")
        assert a.shape == (2, 2, 1)
        a = stat_method(ss_sizes, ss, None, None, None, None, "branch")
        assert a.shape == (2, 2, 1)

        # CPython API errors
        with pytest.raises(ValueError, match="Sum of sample_set_sizes"):
            bad_ss = np.array([], dtype=np.int32)
            stat_method(ss_sizes, bad_ss, row_sites, col_sites, None, None, "site")
        with pytest.raises(TypeError, match="cast array data"):
            bad_ss = np.array(ts.get_samples(), dtype=np.uint32)
            stat_method(ss_sizes, bad_ss, row_sites, col_sites, None, None, "site")
        with pytest.raises(ValueError, match="Unrecognised stats mode"):
            stat_method(ss_sizes, ss, row_sites, col_sites, None, None, "bla")
        with pytest.raises(TypeError, match="at most"):
            stat_method(ss_sizes, ss, row_sites, col_sites, None, None, "site", "abc")
        with pytest.raises(ValueError, match="invalid literal"):
            bad_sites = ["abadsite", 0, 3, 2]
            stat_method(ss_sizes, ss, bad_sites, col_sites, None, None, "site")
        with pytest.raises(TypeError):
            bad_sites = [None, 0, 3, 2]
            stat_method(ss_sizes, ss, bad_sites, col_sites, None, None, "site")
        with pytest.raises(TypeError):
            bad_sites = [{}, 0, 3, 2]
            stat_method(ss_sizes, ss, bad_sites, col_sites, None, None, "site")
        with pytest.raises(TypeError, match="Cannot cast array data"):
            bad_sites = np.array([0, 1, 2], dtype=np.uint32)
            stat_method(ss_sizes, ss, bad_sites, col_sites, None, None, "site")
        with pytest.raises(ValueError, match="invalid literal"):
            bad_sites = ["abadsite", 0, 3, 2]
            stat_method(ss_sizes, ss, row_sites, bad_sites, None, None, "site")
        with pytest.raises(TypeError):
            bad_sites = [None, 0, 3, 2]
            stat_method(ss_sizes, ss, row_sites, bad_sites, None, None, "site")
        with pytest.raises(TypeError):
            bad_sites = [{}, 0, 3, 2]
            stat_method(ss_sizes, ss, row_sites, bad_sites, None, None, "site")
        with pytest.raises(TypeError, match="Cannot cast array data"):
            bad_sites = np.array([0, 1, 2], dtype=np.uint32)
            stat_method(ss_sizes, ss, row_sites, bad_sites, None, None, "site")
        with pytest.raises(ValueError):
            bad_pos = ["abadpos", 0.1, 0.2, 2.0]
            stat_method(ss_sizes, ss, None, None, bad_pos, col_pos, "branch")
        with pytest.raises(TypeError):
            bad_pos = [{}, 0.1, 0.2, 2.0]
            stat_method(ss_sizes, ss, None, None, bad_pos, col_pos, "branch")
        with pytest.raises(ValueError):
            bad_pos = ["abadpos", 0, 3, 2]
            stat_method(ss_sizes, ss, None, None, row_pos, bad_pos, "branch")
        with pytest.raises(TypeError):
            bad_pos = [{}, 0, 3, 2]
            stat_method(ss_sizes, ss, None, None, row_pos, bad_pos, "branch")
        with pytest.raises(ValueError, match="Cannot specify sites in branch mode"):
            stat_method(ss_sizes, ss, row_sites, col_sites, None, None, "branch")
        with pytest.raises(ValueError, match="Cannot specify positions in site mode"):
            stat_method(ss_sizes, ss, None, None, row_pos, col_pos, "site")
        # C API errors
        with pytest.raises(tskit.LibraryError, match="TSK_ERR_STAT_UNSORTED_SITES"):
            bad_sites = np.array([1, 0, 2], dtype=np.int32)
            stat_method(ss_sizes, ss, bad_sites, col_sites, None, None, "site")
        with pytest.raises(tskit.LibraryError, match="TSK_ERR_STAT_UNSORTED_SITES"):
            bad_sites = np.array([1, 0, 2], dtype=np.int32)
            stat_method(ss_sizes, ss, row_sites, bad_sites, None, None, "site")
        with pytest.raises(tskit.LibraryError, match="TSK_ERR_STAT_DUPLICATE_SITES"):
            bad_sites = np.array([1, 1, 2], dtype=np.int32)
            stat_method(ss_sizes, ss, bad_sites, col_sites, None, None, "site")
        with pytest.raises(tskit.LibraryError, match="TSK_ERR_STAT_DUPLICATE_SITES"):
            bad_sites = np.array([1, 1, 2], dtype=np.int32)
            stat_method(ss_sizes, ss, row_sites, bad_sites, None, None, "site")
        with pytest.raises(tskit.LibraryError, match="TSK_ERR_SITE_OUT_OF_BOUNDS"):
            bad_sites = np.array([-1, 0, 2], dtype=np.int32)
            stat_method(ss_sizes, ss, bad_sites, col_sites, None, None, "site")
        with pytest.raises(tskit.LibraryError, match="TSK_ERR_SITE_OUT_OF_BOUNDS"):
            bad_sites = np.array([-1, 0, 2], dtype=np.int32)
            stat_method(ss_sizes, ss, row_sites, bad_sites, None, None, "site")
        with pytest.raises(tskit.LibraryError, match="TSK_ERR_STAT_UNSORTED_POSITIONS"):
            bad_pos = np.array([0.7, 0, 0.8], dtype=np.float64)
            stat_method(ss_sizes, ss, None, None, bad_pos, col_pos, "branch")
        with pytest.raises(tskit.LibraryError, match="TSK_ERR_STAT_UNSORTED_POSITIONS"):
            bad_pos = np.array([0.7, 0, 0.8], dtype=np.float64)
            stat_method(ss_sizes, ss, None, None, row_pos, bad_pos, "branch")
        with pytest.raises(
            tskit.LibraryError, match="TSK_ERR_STAT_DUPLICATE_POSITIONS"
        ):
            bad_pos = np.array([0.7, 0.7, 0.8], dtype=np.float64)
            stat_method(ss_sizes, ss, None, None, bad_pos, col_pos, "branch")
        with pytest.raises(
            tskit.LibraryError, match="TSK_ERR_STAT_DUPLICATE_POSITIONS"
        ):
            bad_pos = np.array([0.7, 0.7, 0.8], dtype=np.float64)
            stat_method(ss_sizes, ss, None, None, row_pos, bad_pos, "branch")
        with pytest.raises(tskit.LibraryError, match="TSK_ERR_POSITION_OUT_OF_BOUNDS"):
            bad_pos = np.array([-0.1, 0.7, 0.8], dtype=np.float64)
            stat_method(ss_sizes, ss, None, None, bad_pos, col_pos, "branch")
        with pytest.raises(tskit.LibraryError, match="TSK_ERR_POSITION_OUT_OF_BOUNDS"):
            bad_pos = np.array([-0.1, 0.7, 0.8], dtype=np.float64)
            stat_method(ss_sizes, ss, None, None, row_pos, bad_pos, "branch")
        with pytest.raises(
            _tskit.LibraryError, match="TSK_ERR_INSUFFICIENT_SAMPLE_SETS"
        ):
            bad_ss = np.array([], dtype=np.int32)
            bad_ss_sizes = np.array([], dtype=np.uint32)
            stat_method(bad_ss_sizes, bad_ss, row_sites, col_sites, None, None, "site")
        with pytest.raises(
            _tskit.LibraryError, match="TSK_ERR_INSUFFICIENT_SAMPLE_SETS"
        ):
            bad_ss = np.array([], dtype=np.int32)
            bad_ss_sizes = np.array([], dtype=np.uint32)
            stat_method(bad_ss_sizes, bad_ss, None, None, row_pos, col_pos, "branch")
        with pytest.raises(_tskit.LibraryError, match="TSK_ERR_EMPTY_SAMPLE_SET"):
            bad_ss = np.array([], dtype=np.int32)
            bad_ss_sizes = np.array([0], dtype=np.uint32)
            stat_method(bad_ss_sizes, bad_ss, row_sites, col_sites, None, None, "site")
        with pytest.raises(_tskit.LibraryError, match="TSK_ERR_EMPTY_SAMPLE_SET"):
            bad_ss = np.array([], dtype=np.int32)
            bad_ss_sizes = np.array([0], dtype=np.uint32)
            stat_method(bad_ss_sizes, bad_ss, None, None, row_pos, col_pos, "branch")
        with pytest.raises(_tskit.LibraryError, match="TSK_ERR_NODE_OUT_OF_BOUNDS"):
            bad_ss = np.array([1000], dtype=np.int32)
            bad_ss_sizes = np.array([1], dtype=np.uint32)
            stat_method(bad_ss_sizes, bad_ss, row_sites, col_sites, None, None, "site")
        with pytest.raises(_tskit.LibraryError, match="TSK_ERR_NODE_OUT_OF_BOUNDS"):
            bad_ss = np.array([1000], dtype=np.int32)
            bad_ss_sizes = np.array([1], dtype=np.uint32)
            stat_method(bad_ss_sizes, bad_ss, None, None, row_pos, col_pos, "branch")
        with pytest.raises(_tskit.LibraryError, match="TSK_ERR_DUPLICATE_SAMPLE"):
            bad_ss = np.array([2, 2], dtype=np.int32)
            bad_ss_sizes = np.array([2], dtype=np.uint32)
            stat_method(bad_ss_sizes, bad_ss, row_sites, col_sites, None, None, "site")
        with pytest.raises(_tskit.LibraryError, match="TSK_ERR_DUPLICATE_SAMPLE"):
            bad_ss = np.array([2, 2], dtype=np.int32)
            bad_ss_sizes = np.array([2], dtype=np.uint32)
            stat_method(bad_ss_sizes, bad_ss, None, None, row_pos, col_pos, "branch")
        with pytest.raises(_tskit.LibraryError, match="TSK_ERR_UNSUPPORTED_STAT_MODE"):
            stat_method(ss_sizes, ss, col_sites, row_sites, None, None, "node")

    def test_kc_distance_errors(self):
        ts1 = self.get_example_tree_sequence(10)
        with pytest.raises(TypeError):
            ts1.get_kc_distance()
        with pytest.raises(TypeError):
            ts1.get_kc_distance(ts1)
        for bad_tree in [None, "tree", 0]:
            with pytest.raises(TypeError):
                ts1.get_kc_distance(bad_tree, lambda_=0)
        for bad_value in ["tree", [], None]:
            with pytest.raises(TypeError):
                ts1.get_kc_distance(ts1, lambda_=bad_value)

        # Different numbers of samples fail.
        ts2 = self.get_example_tree_sequence(11)
        self.verify_kc_library_error(ts1, ts2)

        # Different sequence lengths fail.
        ts2 = self.get_example_tree_sequence(10, length=11)
        self.verify_kc_library_error(ts1, ts2)

    def verify_kc_library_error(self, ts1, ts2):
        with pytest.raises(_tskit.LibraryError):
            ts1.get_kc_distance(ts2, 0)

    def test_kc_distance(self):
        ts1 = self.get_example_tree_sequence(10, random_seed=123456)
        ts2 = self.get_example_tree_sequence(10, random_seed=1234)
        for lambda_ in [-1, 0, 1, 1000, -1e300]:
            x1 = ts1.get_kc_distance(ts2, lambda_)
            x2 = ts2.get_kc_distance(ts1, lambda_)
            assert x1 == x2

    def test_divergence_matrix(self):
        n = 10
        ts = self.get_example_tree_sequence(n, random_seed=12)
        windows = [0, ts.get_sequence_length()]
        ids = np.arange(n, dtype=np.int32)
        sizes = np.ones(n, dtype=np.uint64)
        D = ts.divergence_matrix(windows, sizes, ids)
        assert D.shape == (1, n, n)
        D = ts.divergence_matrix(windows, sample_set_sizes=[1, 1], sample_sets=[0, 1])
        assert D.shape == (1, 2, 2)
        D = ts.divergence_matrix(
            windows, sample_set_sizes=[1, 1], sample_sets=[0, 1], span_normalise=True
        )
        assert D.shape == (1, 2, 2)

        for bad_node in [-1, -2, 1000]:
            with pytest.raises(_tskit.LibraryError, match="TSK_ERR_NODE_OUT_OF_BOUNDS"):
                ts.divergence_matrix(windows, [1, 1], [0, bad_node])
        with pytest.raises(ValueError, match="Sum of sample_set_sizes"):
            ts.divergence_matrix(windows, [1, 2], [0, 1])
        with pytest.raises((ValueError, OverflowError), match="Overflow|out of bounds"):
            ts.divergence_matrix(windows, [-1, 2], [0])

        with pytest.raises(TypeError, match="str"):
            ts.divergence_matrix(windows, sizes, ids, span_normalise="xdf")
        with pytest.raises(TypeError):
            ts.divergence_matrix(windoze=[0, 1])
        with pytest.raises(ValueError, match="at least 2"):
            ts.divergence_matrix(
                [0],
                sizes,
                ids,
            )
        with pytest.raises(_tskit.LibraryError, match="BAD_WINDOWS"):
            ts.divergence_matrix([-1, 0, 1], sizes, ids)
        with pytest.raises(ValueError, match="Unrecognised stats mode"):
            ts.divergence_matrix([0, 1], sizes, ids, mode="sdf")
        with pytest.raises(_tskit.LibraryError, match="UNSUPPORTED_STAT_MODE"):
            ts.divergence_matrix([0, 1], sizes, ids, mode="node")

    def test_load_tables_build_indexes(self):
        for ts in self.get_example_tree_sequences():
            tables = _tskit.TableCollection(sequence_length=ts.get_sequence_length())
            ts.dump_tables(tables)
            tables.drop_index()

            # Tables not in tc but rebuilt
            ts2 = _tskit.TreeSequence()
            ts2.load_tables(tables, build_indexes=True)
            tables2 = _tskit.TableCollection(sequence_length=ts.get_sequence_length())
            ts2.dump_tables(tables2)
            assert tables2.has_index()

            # Tables not in tc, not rebuilt so error
            ts3 = _tskit.TreeSequence()
            with pytest.raises(
                _tskit.LibraryError, match="Table collection must be indexed"
            ):
                ts3.load_tables(tables)

            # Tables in tc, not rebuilt
            tables.build_index()
            ts4 = _tskit.TreeSequence()
            ts4.load_tables(tables, build_indexes=False)
            tables4 = _tskit.TableCollection(sequence_length=ts.get_sequence_length())
            ts4.dump_tables(tables4)
            assert tables4.has_index()

    def test_clear_table(self, ts_fixture):
        tables = _tskit.TableCollection(
            sequence_length=ts_fixture.get_sequence_length()
        )
        ts_fixture.ll_tree_sequence.dump_tables(tables)
        tables.clear()
        data_tables = [t for t in tskit.TABLE_NAMES if t != "provenances"]
        for table in data_tables:
            assert getattr(tables, f"{table}").num_rows == 0
            assert len(getattr(tables, f"{table}").metadata_schema) != 0
        assert tables.provenances.num_rows > 0
        assert len(tables.metadata) > 0
        assert len(tables.metadata_schema) > 0

        tables.clear(clear_provenance=True)
        assert tables.provenances.num_rows == 0
        for table in data_tables:
            assert len(getattr(tables, f"{table}").metadata_schema) != 0
        assert len(tables.metadata) > 0
        assert len(tables.metadata_schema) > 0

        tables.clear(clear_metadata_schemas=True)
        for table in data_tables:
            assert len(getattr(tables, f"{table}").metadata_schema) == 0
        assert len(tables.metadata) > 0
        assert len(tables.metadata_schema) > 0

        tables.clear(clear_ts_metadata_and_schema=True)
        assert len(tables.metadata) == 0
        assert len(tables.metadata_schema) == 0

        # Check for attributes that are not cleared
        assert tables.sequence_length == ts_fixture.tables.sequence_length
        assert tables.time_units == ts_fixture.tables.time_units

    def test_discrete_genome(self):
        tables = _tskit.TableCollection(1)
        tables.build_index()
        ts = _tskit.TreeSequence()
        ts.load_tables(tables)
        assert ts.get_discrete_genome() == 1

    def test_discrete_time(self):
        tables = _tskit.TableCollection(1)
        tables.build_index()
        ts = _tskit.TreeSequence()
        ts.load_tables(tables)
        assert ts.get_discrete_time() == 1

    def test_min_time(self):
        ts = self.get_example_tree_sequence(5)
        assert isinstance(ts.get_min_time(), float)

    def test_max_time(self):
        ts = self.get_example_tree_sequence(5)
        assert isinstance(ts.get_max_time(), float)

    def test_split_edges_return_type(self):
        ts = self.get_example_tree_sequence()
        split = ts.split_edges(time=0, flags=0, population=0, metadata=b"")
        assert isinstance(split, _tskit.TreeSequence)

    def test_split_edges_bad_types(self):
        ts = self.get_example_tree_sequence()

        def f(time=0, flags=0, population=0, metadata=b""):
            return ts.split_edges(
                time=time,
                flags=flags,
                population=population,
                metadata=metadata,
            )

        with pytest.raises(TypeError):
            f(time="0")
        with pytest.raises(TypeError):
            f(flags="0")
        with pytest.raises(TypeError):
            f(metadata="0")

    def test_split_edges_bad_population(self):
        ts = self.get_example_tree_sequence()
        with pytest.raises(_tskit.LibraryError, match="POPULATION_OUT_OF_BOUNDS"):
            ts.split_edges(
                time=0,
                flags=0,
                population=ts.get_num_populations(),
                metadata=b"",
            )

    @pytest.mark.parametrize("name", ARRAY_NAMES)
    def test_array_read_only(self, name):
        ts1 = self.get_example_tree_sequence(10)
        with pytest.raises(AttributeError, match="not writable"):
            setattr(ts1, name, None)
        with pytest.raises(AttributeError, match="not writable"):
            delattr(ts1, name)

        a = getattr(ts1, name)
        with pytest.raises(ValueError, match="assignment destination"):
            a[:] = 0
        with pytest.raises(ValueError, match="assignment destination"):
            a[0] = 0
        with pytest.raises(ValueError, match="cannot set WRITEABLE"):
            a.setflags(write=True)

    @pytest.mark.parametrize("name", ARRAY_NAMES)
    def test_array_properties(self, name):
        ts1 = self.get_example_tree_sequence(10)
        a = getattr(ts1, name)
        assert a.base == ts1
        assert not a.flags.writeable
        assert a.flags.aligned
        assert a.flags.c_contiguous
        assert not a.flags.owndata
        b = getattr(ts1, name)
        assert a is not b
        assert np.all(a == b)
        # This checks that the underlying pointer to memory is the same in
        # both arrays.
        assert a.__array_interface__ == b.__array_interface__

    @pytest.mark.parametrize("name", ARRAY_NAMES)
    def test_array_lifetime(self, name):
        ts1 = self.get_example_tree_sequence(10)
        a1 = getattr(ts1, name)
        a2 = a1.copy()
        assert a1 is not a2
        del ts1
        # Do some memory operations
        a3 = np.ones(10**6)
        assert np.all(a1 == a2)
        del a1
        # Just do something to touch memory
        a2[:] = 0
        assert a3 is not a2


class StatsInterfaceMixin:
    """
    Tests for the interface on specific stats.
    """

    def test_mode_errors(self):
        _, f, params = self.get_example()
        for bad_mode in ["", "not a mode", "SITE", "x" * 8192]:
            with pytest.raises(ValueError):
                f(mode=bad_mode, **params)

        for bad_type in [123, {}, None, [[]]]:
            with pytest.raises(TypeError):
                f(mode=bad_type, **params)

    def test_window_errors(self):
        ts, f, params = self.get_example()
        del params["windows"]
        for bad_array in ["asdf", None, [[[[]], [[]]]], np.zeros((10, 3, 4))]:
            with pytest.raises(ValueError):
                f(windows=bad_array, **params)

        for bad_windows in [[], [0]]:
            with pytest.raises(ValueError):
                f(windows=bad_windows, **params)
        L = ts.get_sequence_length()
        bad_windows = [
            [L, 0],
            [0.1, L],
            [-1, L],
            [0, L + 0.1],
            [0, 0.1, 0.1, L],
            [0, -1, L],
            [0, 0.1, 0.05, 0.2, L],
        ]
        for bad_window in bad_windows:
            with pytest.raises(_tskit.LibraryError):
                f(windows=bad_window, **params)

    def test_polarisation(self):
        ts, f, params = self.get_example()
        with pytest.raises(TypeError):
            f(polarised="sdf", **params)
        x1 = f(polarised=False, **params)
        x2 = f(polarised=True, **params)
        # Basic check just to run both code paths
        assert x1.shape == x2.shape

    def test_windows_output(self):
        ts, f, params = self.get_example()
        del params["windows"]
        for num_windows in range(1, 10):
            windows = np.linspace(0, ts.get_sequence_length(), num=num_windows + 1)
            assert windows.shape[0] == num_windows + 1
            sigma = f(windows=windows, **params)
            assert sigma.shape[0] == num_windows


class WeightMixin(StatsInterfaceMixin):
    def get_example(self):
        ts, method = self.get_method()
        params = {
            "weights": np.ones((ts.get_num_samples(), 2)),
            "windows": [0, ts.get_sequence_length()],
        }
        return ts, method, params

    def test_bad_weights(self):
        ts, f, params = self.get_example()
        del params["weights"]
        n = ts.get_num_samples()

        with pytest.raises(_tskit.LibraryError):
            f(weights=np.ones((n, 0)), **params)

        for bad_weight_shape in [(n - 1, 1), (n + 1, 1), (0, 3)]:
            with pytest.raises(ValueError):
                f(weights=np.ones(bad_weight_shape), **params)

    def test_output_dims(self):
        ts, method, params = self.get_example()
        weights = params["weights"]
        nw = weights.shape[1]
        windows = [0, ts.get_sequence_length()]

        for mode in ["site", "branch"]:
            out = method(weights[:, [0]], windows, mode=mode)
            assert out.shape == (1, 1)
            out = method(weights, windows, mode=mode)
            assert out.shape == (1, nw)
            out = method(weights[:, [0, 0, 0]], windows, mode=mode)
            assert out.shape == (1, 3)
        mode = "node"
        N = ts.get_num_nodes()
        out = method(weights[:, [0]], windows, mode=mode)
        assert out.shape == (1, N, 1)
        out = method(weights, windows, mode=mode)
        assert out.shape == (1, N, nw)
        out = method(weights[:, [0, 0, 0]], windows, mode=mode)
        assert out.shape == (1, N, 3)


class WeightCovariateMixin(StatsInterfaceMixin):
    def get_example(self):
        ts, method = self.get_method()
        params = {
            "weights": np.ones((ts.get_num_samples(), 2)),
            "covariates": np.array(
                [np.arange(ts.get_num_samples()), np.arange(ts.get_num_samples()) ** 2]
            ).T,
            "windows": [0, ts.get_sequence_length()],
        }
        return ts, method, params

    def test_output_dims(self):
        ts, method, params = self.get_example()
        weights = params["weights"]
        nw = weights.shape[1]
        windows = [0, ts.get_sequence_length()]
        for covariates in (params["covariates"], params["covariates"][:, :0]):
            for mode in ["site", "branch"]:
                out = method(weights[:, [0]], covariates, windows, mode=mode)
                assert out.shape == (1, 1)
                out = method(weights, covariates, windows, mode=mode)
                assert out.shape == (1, nw)
                out = method(weights[:, [0, 0, 0]], covariates, windows, mode=mode)
                assert out.shape == (1, 3)
            mode = "node"
            N = ts.get_num_nodes()
            out = method(weights[:, [0]], covariates, windows, mode=mode)
            assert out.shape == (1, N, 1)
            out = method(weights, covariates, windows, mode=mode)
            assert out.shape == (1, N, nw)
            out = method(weights[:, [0, 0, 0]], covariates, windows, mode=mode)
            assert out.shape == (1, N, 3)


class SampleSetMixin(StatsInterfaceMixin):
    def test_bad_sample_sets(self):
        ts, f, params = self.get_example()
        del params["sample_set_sizes"]
        del params["sample_sets"]

        with pytest.raises(_tskit.LibraryError):
            f(sample_sets=[], sample_set_sizes=[], **params)

        n = ts.get_num_samples()
        samples = ts.get_samples()
        for bad_set_sizes in [[], [1], [n - 1], [n + 1], [n - 3, 1, 1], [1, n - 2]]:
            with pytest.raises(ValueError):
                f(sample_set_sizes=bad_set_sizes, sample_sets=samples, **params)

        N = ts.get_num_nodes()
        for bad_node in [-1, N, N + 1, -N]:
            with pytest.raises(_tskit.LibraryError):
                f(sample_set_sizes=[2], sample_sets=[0, bad_node], **params)

        for bad_sample in [n, n + 1, N - 1]:
            with pytest.raises(_tskit.LibraryError):
                f(sample_set_sizes=[2], sample_sets=[0, bad_sample], **params)


class OneWaySampleStatsMixin(SampleSetMixin):
    """
    Tests for one-way sample stats.
    """

    def get_example(self):
        ts, method = self.get_method()
        params = {
            "sample_set_sizes": [ts.get_num_samples()],
            "sample_sets": ts.get_samples(),
            "windows": [0, ts.get_sequence_length()],
        }
        return ts, method, params

    def test_basic_example(self):
        ts, method = self.get_method()
        result = method(
            [ts.get_num_samples()], ts.get_samples(), [0, ts.get_sequence_length()]
        )
        assert result.shape == (1, 1)
        result = method(
            [ts.get_num_samples()],
            ts.get_samples(),
            [0, ts.get_sequence_length()],
            mode="node",
        )
        assert result.shape == (1, ts.get_num_nodes(), 1)
        result = method(
            [ts.get_num_samples()], ts.get_samples(), ts.get_breakpoints(), mode="node"
        )
        assert result.shape == (ts.get_num_trees(), ts.get_num_nodes(), 1)

    def test_output_dims(self):
        ts, method = self.get_method()
        samples = ts.get_samples()
        windows = [0, ts.get_sequence_length()]
        n = len(samples)

        for mode in ["site", "branch"]:
            pi = method([n], samples, windows, mode=mode)
            assert pi.shape == (1, 1)
            pi = method([2, n - 2], samples, windows, mode=mode)
            assert pi.shape == (1, 2)
            pi = method([2, 2, n - 4], samples, windows, mode=mode)
            assert pi.shape == (1, 3)
            pi = method(np.ones(n).astype(np.uint32), samples, windows, mode=mode)
            assert pi.shape == (1, n)
        mode = "node"
        N = ts.get_num_nodes()
        pi = method([n], samples, windows, mode=mode)
        assert pi.shape == (1, N, 1)
        pi = method([2, n - 2], samples, windows, mode=mode)
        assert pi.shape == (1, N, 2)
        pi = method([2, 2, n - 4], samples, windows, mode=mode)
        assert pi.shape == (1, N, 3)
        pi = method(np.ones(n).astype(np.uint32), samples, windows, mode=mode)
        assert pi.shape == (1, N, n)

    def test_polarised(self):
        # TODO move this to the top level.
        ts, method = self.get_method()
        samples = ts.get_samples()
        n = len(samples)
        windows = [0, ts.get_sequence_length()]
        method([n], samples, windows, polarised=True)
        method([n], samples, windows, polarised=False)


class TestDiversity(LowLevelTestCase, OneWaySampleStatsMixin):
    """
    Tests for the diversity method.
    """

    def get_method(self):
        ts = self.get_example_tree_sequence()
        return ts, ts.diversity


class TestTraitCovariance(LowLevelTestCase, WeightMixin):
    """
    Tests for trait covariance.
    """

    def get_method(self):
        ts = self.get_example_tree_sequence()
        return ts, ts.trait_covariance


class TestTraitCorrelation(LowLevelTestCase, WeightMixin):
    """
    Tests for trait correlation.
    """

    def get_method(self):
        ts = self.get_example_tree_sequence()
        return ts, ts.trait_correlation


class TestTraitLinearModel(LowLevelTestCase, WeightCovariateMixin):
    """
    Tests for trait correlation.
    """

    def get_method(self):
        ts = self.get_example_tree_sequence()
        return ts, ts.trait_linear_model


class TestSegregatingSites(LowLevelTestCase, OneWaySampleStatsMixin):
    """
    Tests for the diversity method.
    """

    def get_method(self):
        ts = self.get_example_tree_sequence()
        return ts, ts.segregating_sites


class TestY1(LowLevelTestCase, OneWaySampleStatsMixin):
    """
    Tests for the diversity method.
    """

    def get_method(self):
        ts = self.get_example_tree_sequence()
        return ts, ts.Y1


class TestAlleleFrequencySpectrum(LowLevelTestCase, OneWaySampleStatsMixin):
    """
    Tests for the diversity method.
    """

    def get_method(self):
        ts = self.get_example_tree_sequence()
        return ts, ts.allele_frequency_spectrum

    def test_basic_example(self):
        ts = self.get_example_tree_sequence()
        n = ts.get_num_samples()
        result = ts.allele_frequency_spectrum(
            [n], ts.get_samples(), [0, ts.get_sequence_length()]
        )
        assert result.shape == (1, n + 1)
        result = ts.allele_frequency_spectrum(
            [n], ts.get_samples(), [0, ts.get_sequence_length()], polarised=True
        )
        assert result.shape == (1, n + 1)

    def test_output_dims(self):
        ts = self.get_example_tree_sequence()
        samples = ts.get_samples()
        L = ts.get_sequence_length()
        n = len(samples)

        for mode in ["site", "branch"]:
            for s in [[n], [n - 2, 2], [n - 4, 2, 2], [1] * n]:
                s = np.array(s, dtype=np.uint32)
                windows = [0, L]
                for windows in [[0, L], [0, L / 2, L], np.linspace(0, L, num=10)]:
                    jafs = ts.allele_frequency_spectrum(
                        s, samples, windows, mode=mode, polarised=True
                    )
                    assert jafs.shape == tuple([len(windows) - 1] + list(s + 1))
                    jafs = ts.allele_frequency_spectrum(
                        s, samples, windows, mode=mode, polarised=False
                    )
                    assert jafs.shape == tuple([len(windows) - 1] + list(s + 1))

    def test_node_mode_not_supported(self):
        ts = self.get_example_tree_sequence()
        with pytest.raises(_tskit.LibraryError):
            ts.allele_frequency_spectrum(
                [ts.get_num_samples()],
                ts.get_samples(),
                [0, ts.get_sequence_length()],
                mode="node",
            )


class TwoWaySampleStatsMixin(SampleSetMixin):
    """
    Tests for the two way sample stats.
    """

    def get_example(self):
        ts, method = self.get_method()
        params = {
            "sample_set_sizes": [2, ts.get_num_samples() - 2],
            "sample_sets": ts.get_samples(),
            "indexes": [[0, 1]],
            "windows": [0, ts.get_sequence_length()],
        }
        return ts, method, params

    def test_basic_example(self):
        ts, method = self.get_method()
        div = method(
            [2, ts.get_num_samples() - 2],
            ts.get_samples(),
            [[0, 1]],
            windows=[0, ts.get_sequence_length()],
        )
        assert div.shape == (1, 1)

    def test_output_dims(self):
        ts, method = self.get_method()
        samples = ts.get_samples()
        windows = [0, ts.get_sequence_length()]
        n = len(samples)
        for mode in ["site", "branch"]:
            div = method([2, 2, n - 4], samples, [[0, 1]], windows, mode=mode)
            assert div.shape == (1, 1)
            div = method([2, 2, n - 4], samples, [[0, 1], [1, 2]], windows, mode=mode)
            assert div.shape == (1, 2)
            div = method(
                [2, 2, n - 4], samples, [[0, 1], [1, 2], [0, 1]], windows, mode=mode
            )
            assert div.shape == (1, 3)

        N = ts.get_num_nodes()
        mode = "node"
        div = method([2, 2, n - 4], samples, [[0, 1]], windows, mode=mode)
        assert div.shape == (1, N, 1)
        div = method([2, 2, n - 4], samples, [[0, 1], [1, 2]], windows, mode=mode)
        assert div.shape == (1, N, 2)
        div = method(
            [2, 2, n - 4], samples, [[0, 1], [1, 2], [0, 1]], windows, mode=mode
        )
        assert div.shape == (1, N, 3)

    def test_set_index_errors(self):
        ts, method = self.get_method()
        samples = ts.get_samples()
        windows = [0, ts.get_sequence_length()]
        n = len(samples)

        def f(indexes):
            method([2, 2, n - 4], samples, indexes, windows)

        for bad_array in ["wer", {}, [[[], []], [[], []]]]:
            with pytest.raises(ValueError):
                f(bad_array)
        for bad_dim in [[[]], [[1], [1]]]:
            with pytest.raises(ValueError):
                f(bad_dim)


class TwoWayWeightedStatsMixin(StatsInterfaceMixin):
    """
    Tests for the weighted two way sample stats.
    """

    def get_example(self):
        ts, method = self.get_method()
        params = {
            "weights": np.zeros((ts.get_num_samples(), 2)) + 0.5,
            "indexes": [[0, 1]],
            "windows": [0, ts.get_sequence_length()],
        }
        return ts, method, params

    def test_basic_example(self):
        ts, method = self.get_method()
        div = method(
            np.zeros((ts.get_num_samples(), 1)) + 0.5,
            [[0, 1]],
            windows=[0, ts.get_sequence_length()],
        )
        assert div.shape == (1, 1)

    def test_bad_weights(self):
        ts, f, params = self.get_example()
        del params["weights"]
        n = ts.get_num_samples()

        for bad_weight_type in [None, [None, None]]:
            with pytest.raises(ValueError, match="object of too small depth"):
                f(weights=bad_weight_type, **params)

        for bad_weight_shape in [(n - 1, 1), (n + 1, 1), (0, 3)]:
            with pytest.raises(ValueError, match="First dimension must be num_samples"):
                f(weights=np.ones(bad_weight_shape), **params)

    def test_output_dims(self):
        ts, method, params = self.get_example()
        weights = params.pop("weights")
        params["windows"] = [0, ts.get_sequence_length()]

        for mode in ["site", "branch"]:
            out = method(weights[:, [0]], mode=mode, **params)
            assert out.shape == (1, 1)
            out = method(weights, mode=mode, **params)
            assert out.shape == (1, 1)
            out = method(weights[:, [0, 0, 0]], mode=mode, **params)
            assert out.shape == (1, 1)
        mode = "node"
        N = ts.get_num_nodes()
        out = method(weights[:, [0]], mode=mode, **params)
        assert out.shape == (1, N, 1)
        out = method(weights, mode=mode, **params)
        assert out.shape == (1, N, 1)
        out = method(weights[:, [0, 0, 0]], mode=mode, **params)
        assert out.shape == (1, N, 1)

    def test_set_index_errors(self):
        ts, method, params = self.get_example()
        del params["indexes"]

        def f(indexes):
            method(indexes=indexes, **params)

        for bad_array in ["wer", {}, [[[], []], [[], []]]]:
            with pytest.raises(ValueError):
                f(bad_array)
        for bad_dim in [[[]], [[1], [1]]]:
            with pytest.raises(ValueError):
                f(bad_dim)


class ThreeWaySampleStatsMixin(SampleSetMixin):
    """
    Tests for the two way sample stats.
    """

    def get_example(self):
        ts, method = self.get_method()
        params = {
            "sample_set_sizes": [1, 1, ts.get_num_samples() - 2],
            "sample_sets": ts.get_samples(),
            "indexes": [[0, 1, 2]],
            "windows": [0, ts.get_sequence_length()],
        }
        return ts, method, params

    def test_basic_example(self):
        ts, method = self.get_method()
        div = method(
            [1, 1, ts.get_num_samples() - 2],
            ts.get_samples(),
            [[0, 1, 2]],
            windows=[0, ts.get_sequence_length()],
        )
        assert div.shape == (1, 1)

    def test_output_dims(self):
        ts, method = self.get_method()
        samples = ts.get_samples()
        windows = [0, ts.get_sequence_length()]
        n = len(samples)
        for mode in ["site", "branch"]:
            div = method([2, 2, n - 4], samples, [[0, 1, 2]], windows, mode=mode)
            assert div.shape == (1, 1)
            div = method(
                [1, 1, 2, n - 4], samples, [[0, 1, 2], [1, 2, 3]], windows, mode=mode
            )
            assert div.shape == (1, 2)
            div = method(
                [1, 1, 2, n - 4],
                samples,
                [[0, 1, 2], [1, 2, 3], [0, 1, 2]],
                windows,
                mode=mode,
            )
            assert div.shape == (1, 3)

        N = ts.get_num_nodes()
        mode = "node"
        div = method([2, 2, n - 4], samples, [[0, 1, 2]], windows, mode=mode)
        assert div.shape == (1, N, 1)
        div = method(
            [1, 1, 2, n - 4], samples, [[0, 1, 2], [1, 2, 3]], windows, mode=mode
        )
        assert div.shape == (1, N, 2)
        div = method(
            [1, 1, 2, n - 4],
            samples,
            [[0, 1, 2], [1, 2, 3], [0, 1, 2]],
            windows,
            mode=mode,
        )
        assert div.shape == (1, N, 3)

    def test_set_index_errors(self):
        ts, method = self.get_method()
        samples = ts.get_samples()
        windows = [0, ts.get_sequence_length()]
        n = len(samples)

        def f(indexes):
            method([2, 2, n - 4], samples, indexes, windows)

        for bad_array in ["wer", {}, [[[], []], [[], []]]]:
            with pytest.raises(ValueError):
                f(bad_array)
        for bad_dim in [[[]], [[1], [1]], [(0, 1)], [(0, 1, 2, 3)]]:
            with pytest.raises(ValueError):
                f(bad_dim)


class FourWaySampleStatsMixin(SampleSetMixin):
    """
    Tests for the four way sample stats.
    """

    def get_example(self):
        ts, method = self.get_method()
        params = {
            "sample_set_sizes": [1, 1, 1, ts.get_num_samples() - 3],
            "sample_sets": ts.get_samples(),
            "indexes": [[0, 1, 2, 3]],
            "windows": [0, ts.get_sequence_length()],
        }
        return ts, method, params

    def test_basic_example(self):
        ts, method = self.get_method()
        div = method(
            [1, 1, 1, ts.get_num_samples() - 3],
            ts.get_samples(),
            [[0, 1, 2, 3]],
            windows=[0, ts.get_sequence_length()],
        )
        assert div.shape == (1, 1)

    def test_output_dims(self):
        ts, method = self.get_method()
        samples = ts.get_samples()
        windows = [0, ts.get_sequence_length()]
        n = len(samples)
        for mode in ["site", "branch"]:
            div = method([2, 1, 1, n - 4], samples, [[0, 1, 2, 3]], windows, mode=mode)
            assert div.shape == (1, 1)
            div = method(
                [1, 1, 1, 1, n - 4],
                samples,
                [[0, 1, 2, 3], [1, 2, 3, 4]],
                windows,
                mode=mode,
            )
            assert div.shape == (1, 2)
            div = method(
                [1, 1, 1, 1, n - 4],
                samples,
                [[0, 1, 2, 3], [1, 2, 3, 4], [0, 1, 2, 4]],
                windows,
                mode=mode,
            )
            assert div.shape == (1, 3)

        N = ts.get_num_nodes()
        mode = "node"
        div = method([2, 1, 1, n - 4], samples, [[0, 1, 2, 3]], windows, mode=mode)
        assert div.shape == (1, N, 1)
        div = method(
            [1, 1, 1, 1, n - 4],
            samples,
            [[0, 1, 2, 3], [1, 2, 3, 4]],
            windows,
            mode=mode,
        )
        assert div.shape == (1, N, 2)
        div = method(
            [1, 1, 1, 1, n - 4],
            samples,
            [[0, 1, 2, 3], [1, 2, 3, 4], [0, 1, 2, 4]],
            windows,
            mode=mode,
        )
        assert div.shape == (1, N, 3)

    def test_set_index_errors(self):
        ts, method = self.get_method()
        samples = ts.get_samples()
        windows = [0, ts.get_sequence_length()]
        n = len(samples)

        def f(indexes):
            method([2, 1, 1, n - 4], samples, indexes, windows)

        for bad_array in ["wer", {}, [[[], []], [[], []]]]:
            with pytest.raises(ValueError):
                f(bad_array)
        for bad_dim in [[[]], [[1], [1]], [(0, 1)], [(0, 1, 2, 3, 4)]]:
            with pytest.raises(ValueError):
                f(bad_dim)


class TestDivergence(LowLevelTestCase, TwoWaySampleStatsMixin):
    def get_method(self):
        ts = self.get_example_tree_sequence()
        return ts, ts.divergence


class TestY2(LowLevelTestCase, TwoWaySampleStatsMixin):
    def get_method(self):
        ts = self.get_example_tree_sequence()
        return ts, ts.Y2


class Testf2(LowLevelTestCase, TwoWaySampleStatsMixin):
    def get_method(self):
        ts = self.get_example_tree_sequence()
        return ts, ts.f2


class TestGeneticRelatedness(LowLevelTestCase, TwoWaySampleStatsMixin):
    def get_method(self):
        ts = self.get_example_tree_sequence()
        return ts, ts.genetic_relatedness

    def test_options(self):
        ts, _, params = self.get_example()
        x = ts.genetic_relatedness(**params)
        new_params = params.copy()
        new_params["centre"] = False
        y = ts.genetic_relatedness(**new_params)
        assert x.shape == y.shape
        new_params["polarised"] = False
        y = ts.genetic_relatedness(**new_params)
        assert x.shape == y.shape
        del new_params["centre"]
        y = ts.genetic_relatedness(**new_params)
        assert x.shape == y.shape

        del new_params["indexes"]
        with pytest.raises(ValueError, match="object of too small depth"):
            ts.genetic_relatedness(**new_params, indexes="foo")


class TestY3(LowLevelTestCase, ThreeWaySampleStatsMixin):
    def get_method(self):
        ts = self.get_example_tree_sequence()
        return ts, ts.Y3


class Testf3(LowLevelTestCase, ThreeWaySampleStatsMixin):
    def get_method(self):
        ts = self.get_example_tree_sequence()
        return ts, ts.f3


class Testf4(LowLevelTestCase, FourWaySampleStatsMixin):
    def get_method(self):
        ts = self.get_example_tree_sequence()
        return ts, ts.f4


class TestWeightedGeneticRelatedness(LowLevelTestCase, TwoWayWeightedStatsMixin):
    def get_method(self):
        ts = self.get_example_tree_sequence()
        return ts, ts.genetic_relatedness_weighted

    def test_options(self):
        ts, _, params = self.get_example()
        x = ts.genetic_relatedness_weighted(**params)

        new_params = params.copy()
        new_params["centre"] = False
        y = ts.genetic_relatedness_weighted(**new_params)
        assert x.shape == y.shape
        new_params["polarised"] = False
        y = ts.genetic_relatedness_weighted(**new_params)
        assert x.shape == y.shape
        del new_params["centre"]
        y = ts.genetic_relatedness_weighted(**new_params)
        assert x.shape == y.shape

        del new_params["weights"]
        with pytest.raises(ValueError, match="First dimension"):
            ts.genetic_relatedness_weighted(
                **new_params, weights=np.ones((ts.get_num_samples() + 2, 1))
            )


class TestGeneticRelatednessVector(LowLevelTestCase):
    def get_example(self, num_weights=2):
        ts = self.get_example_tree_sequence()
        num_samples = ts.get_num_samples()
        params = {
            "weights": np.linspace(0, 1, num_weights * num_samples).reshape(
                (num_samples, num_weights)
            ),
            "windows": [0, ts.get_sequence_length()],
            "nodes": list(ts.get_samples()),
        }
        return ts, params

    @pytest.mark.parametrize("mode", ["branch"])
    @pytest.mark.parametrize("num_weights", [1, 3])
    def test_basic_example(self, mode, num_weights):
        ts, params = self.get_example(num_weights)
        ns = ts.get_num_samples()
        params["mode"] = mode
        for a, b in ([True, True], [True, False], [False, True]):
            params["span_normalise"] = a
            params["centre"] = b
            result = ts.genetic_relatedness_vector(**params)
            assert result.shape == (1, ns, num_weights)

    @pytest.mark.parametrize("mode", ["branch"])
    def test_good_nodes(self, mode):
        num_weights = 2
        ts, params = self.get_example(num_weights)
        params["mode"] = mode
        for nodes in [
            list(ts.get_samples())[:3],
            list(ts.get_samples())[:1],
            [0, ts.get_num_nodes() - 1],
        ]:
            params["nodes"] = nodes
            result = ts.genetic_relatedness_vector(**params)
            assert result.shape == (1, len(nodes), num_weights)

    def test_bad_nodes(self):
        ts, params = self.get_example()
        params["mode"] = "branch"
        for nodes in ["abc", [[1, 2]]]:
            params["nodes"] = nodes
            with pytest.raises(ValueError, match="desired array"):
                ts.genetic_relatedness_vector(**params)
        for nodes in [[-1, 3], [3, 2 * ts.get_num_nodes()]]:
            params["nodes"] = nodes
            with pytest.raises(_tskit.LibraryError, match="TSK_ERR_NODE_OUT_OF_BOUNDS"):
                ts.genetic_relatedness_vector(**params)

    def test_bad_args(self):
        ts, params = self.get_example()
        for mode in ("", "abc"):
            with pytest.raises(ValueError, match="stats mode"):
                ts.genetic_relatedness_vector(
                    params["weights"], params["windows"], mode, True, True
                )
        for mode in (None, []):
            with pytest.raises(TypeError):
                ts.genetic_relatedness_vector(
                    params["weights"], params["windows"], mode, True, True
                )
        with pytest.raises(TypeError):
            ts.genetic_relatedness_vector(
                params["weights"], params["windows"], "branch", "yes", True
            )
        with pytest.raises(TypeError):
            ts.genetic_relatedness_vector(
                params["weights"], params["windows"], "branch", True, "no"
            )

    @pytest.mark.parametrize("mode", ["site", "node"])
    def test_modes_not_supported(self, mode):
        ts, params = self.get_example()
        params["mode"] = mode
        with pytest.raises(_tskit.LibraryError):
            ts.genetic_relatedness_vector(**params)

    @pytest.mark.parametrize("mode", ["branch"])
    def test_bad_weights(self, mode):
        ts, params = self.get_example()
        del params["weights"]
        ns = ts.get_num_samples()
        for bad_weight_type in [None, [None, None]]:
            with pytest.raises(ValueError, match="object of too small depth"):
                ts.genetic_relatedness_vector(
                    weights=bad_weight_type, mode=mode, **params
                )
        for bad_weight_shape in [(ns - 1, 1), (ns + 1, 1), (0, 3)]:
            with pytest.raises(ValueError, match="First dimension must be num_samples"):
                ts.genetic_relatedness_vector(
                    weights=np.ones(bad_weight_shape), mode=mode, **params
                )

    def test_window_errors(self):
        ts, params = self.get_example()
        del params["windows"]
        for bad_array in ["asdf", None, [[[[]], [[]]]], np.zeros((10, 3, 4))]:
            with pytest.raises(ValueError):
                ts.genetic_relatedness_vector(
                    windows=bad_array, mode="branch", **params
                )

        for bad_windows in [[], [0]]:
            with pytest.raises(ValueError):
                ts.genetic_relatedness_vector(
                    windows=bad_windows, mode="branch", **params
                )
        L = ts.get_sequence_length()
        bad_windows = [
            [L, 0],
            [-1, L],
            [0, L + 0.1],
            [0, 0.1, 0.1, L],
            [0, -1, L],
        ]
        for bad_window in bad_windows:
            with pytest.raises(_tskit.LibraryError):
                ts.genetic_relatedness_vector(
                    windows=bad_window, mode="branch", **params
                )


class TestGeneralStatsInterface(LowLevelTestCase, StatsInterfaceMixin):
    """
    Tests for the general stats interface.
    """

    def get_example(self):
        ts = self.get_example_tree_sequence()
        W = np.zeros((ts.get_num_samples(), 1))
        params = {
            "weights": W,
            "summary_func": lambda x: np.cumsum(x),
            "output_dim": 1,
            "windows": ts.get_breakpoints(),
        }
        return ts, ts.general_stat, params

    def test_basic_example(self):
        ts = self.get_example_tree_sequence()
        W = np.zeros((ts.get_num_samples(), 1))
        sigma = ts.general_stat(
            W, lambda x: np.cumsum(x), 1, ts.get_breakpoints(), mode="branch"
        )
        assert sigma.shape == (ts.get_num_trees(), 1)

    def test_non_numpy_return(self):
        ts = self.get_example_tree_sequence()
        W = np.ones((ts.get_num_samples(), 3))
        sigma = ts.general_stat(
            W, lambda x: [sum(x)], 1, ts.get_breakpoints(), mode="branch"
        )
        assert sigma.shape == (ts.get_num_trees(), 1)
        sigma = ts.general_stat(
            W, lambda x: [2, 2], 2, ts.get_breakpoints(), mode="branch"
        )
        assert sigma.shape == (ts.get_num_trees(), 2)

    def test_complicated_numpy_function(self):
        ts = self.get_example_tree_sequence(sample_size=20, length=30, random_seed=325)
        W = np.zeros((ts.get_num_samples(), 4))

        def f(x):
            y = np.sum(x * x), np.prod(x + np.arange(x.shape[0]))
            return y

        sigma = ts.general_stat(W, f, 2, ts.get_breakpoints(), mode="branch")
        assert sigma.shape == (ts.get_num_trees(), 2)

    def test_input_dims(self):
        ts = self.get_example_tree_sequence()
        for k in range(1, 20):
            W = np.zeros((ts.get_num_samples(), k))
            sigma = ts.general_stat(
                W, lambda x: np.cumsum(x), k, ts.get_breakpoints(), mode="branch"
            )
            assert sigma.shape == (ts.get_num_trees(), k)
            sigma = ts.general_stat(
                W, lambda x: [np.sum(x)], 1, ts.get_breakpoints(), mode="branch"
            )
            assert sigma.shape == (ts.get_num_trees(), 1)

    def test_W_errors(self):
        ts = self.get_example_tree_sequence()
        n = ts.get_num_samples()
        for bad_array in [[], [0, 1], [[[[]], [[]]]], np.zeros((10, 3, 4))]:
            with pytest.raises(ValueError):
                ts.general_stat(bad_array, lambda x: x, 1, ts.get_breakpoints())

        for bad_size in [n - 1, n + 1, 0]:
            W = np.zeros((bad_size, 1))
            with pytest.raises(ValueError):
                ts.general_stat(W, lambda x: x, 1, ts.get_breakpoints())

    def test_summary_func_errors(self):
        ts = self.get_example_tree_sequence()
        W = np.zeros((ts.get_num_samples(), 1))
        for bad_type in ["sdf", 1, {}]:
            with pytest.raises(TypeError):
                ts.general_stat(W, bad_type, 1, ts.get_breakpoints())

        # Wrong numbers of arguments to f
        with pytest.raises(TypeError):
            ts.general_stat(W, lambda: 0, 1, ts.get_breakpoints())
        with pytest.raises(TypeError):
            ts.general_stat(W, lambda x, y: None, 1, ts.get_breakpoints())

        # Exceptions within f are correctly raised.
        for exception in [ValueError, TypeError]:

            def f(x):
                raise exception("test")  # noqa: B023

            with pytest.raises(exception):
                ts.general_stat(W, f, 1, ts.get_breakpoints())

        # Wrong output dimensions
        for bad_array in [[1, 1], range(10)]:
            with pytest.raises(ValueError):
                ts.general_stat(
                    W, lambda x: bad_array, 1, ts.get_breakpoints()  # noqa:B023
                )
        with pytest.raises(ValueError):
            ts.general_stat(W, lambda x: [1], 2, ts.get_breakpoints())

        # Bad arrays returned from f
        for bad_array in [["sdf"], 0, "w4", None]:
            with pytest.raises(ValueError):
                ts.general_stat(
                    W, lambda x: bad_array, 1, ts.get_breakpoints()  # noqa:B023
                )


class TestVariant(LowLevelTestCase):
    """
    Tests for the Variant class.
    """

    def test_uninitialised_tree_sequence(self):
        ts = _tskit.TreeSequence()
        with pytest.raises(ValueError):
            _tskit.Variant(ts)

    def test_constructor(self):
        with pytest.raises(TypeError):
            _tskit.Variant()
        with pytest.raises(TypeError):
            _tskit.Variant(None)
        ts = self.get_example_tree_sequence()
        with pytest.raises(ValueError):
            _tskit.Variant(ts, samples={})
        with pytest.raises(TypeError):
            _tskit.Variant(ts, isolated_as_missing=None)
        with pytest.raises(_tskit.LibraryError):
            _tskit.Variant(ts, samples=[-1, 2])
        with pytest.raises(TypeError):
            _tskit.Variant(ts, alleles=1234)

    def test_bad_decode(self):
        ts = self.get_example_tree_sequence()
        variant = _tskit.Variant(ts)
        with pytest.raises(tskit.LibraryError, match="Site out of bounds"):
            variant.decode(-1)
        with pytest.raises(TypeError):
            variant.decode("42")
        with pytest.raises(TypeError):
            variant.decode({})
        with pytest.raises(TypeError):
            variant.decode()

    def test_alleles(self):
        ts = self.get_example_tree_sequence()
        for bad_type in [["a", "b"], "sdf", 234]:
            with pytest.raises(TypeError):
                _tskit.Variant(ts, samples=[1, 2], alleles=bad_type)
        with pytest.raises(ValueError):
            _tskit.Variant(ts, samples=[1, 2], alleles=tuple())

        for bad_allele_type in [None, 0, b"x", []]:
            with pytest.raises(TypeError):
                _tskit.Variant(ts, samples=[1, 2], alleles=(bad_allele_type,))

    def test_samples(self):
        ts = self.get_example_tree_sequence()
        v = _tskit.Variant(ts, samples=None, alleles=None)
        assert np.array_equal(
            v.samples, np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=np.int32)
        )
        v = _tskit.Variant(ts, samples=[4, 2], alleles=None)
        assert np.array_equal(v.samples, np.array([4, 2], dtype=np.int32))
        v = _tskit.Variant(ts, samples=[], alleles=None)
        assert np.array_equal(v.samples, np.array([], dtype=np.int32))
        with pytest.raises(AttributeError):
            v.samples = [1]

    def test_isolated_as_missing(self):
        ts = self.get_example_tree_sequence()
        v = _tskit.Variant(ts)
        assert v.isolated_as_missing
        v = _tskit.Variant(ts, isolated_as_missing=True)
        assert v.isolated_as_missing
        v = _tskit.Variant(ts, isolated_as_missing=False)
        assert not v.isolated_as_missing

    def test_undecoded(self):
        tables = _tskit.TableCollection(1)
        tables.build_index()
        ts = _tskit.TreeSequence(0)
        ts.load_tables(tables)
        variant = _tskit.Variant(ts)
        assert variant.site_id == tskit.NULL
        assert np.array_equal(variant.genotypes, [])
        assert variant.alleles == ()
        assert np.array_equal(variant.samples, ())
        assert variant.isolated_as_missing

    def test_properties_unwritable(self):
        ts = self.get_example_tree_sequence()
        variant = _tskit.Variant(ts)
        with pytest.raises(AttributeError):
            variant.site_id = 1
        with pytest.raises(AttributeError):
            variant.genotypes = [1]
        with pytest.raises(AttributeError):
            variant.samples = [1]
        with pytest.raises(AttributeError):
            variant.isolated_as_missing = False
        with pytest.raises(AttributeError):
            variant.alleles = "A"

    def test_missing_data(self):
        tables = _tskit.TableCollection(1)
        tables.nodes.add_row(flags=1, time=0)
        tables.nodes.add_row(flags=1, time=0)
        tables.sites.add_row(0.1, "A")
        tables.build_index()
        ts = _tskit.TreeSequence(0)
        ts.load_tables(tables)
        variant = _tskit.Variant(ts)
        variant.decode(0)
        assert variant.site_id == 0
        assert np.array_equal(variant.genotypes, [-1, -1])
        assert variant.alleles == ("A", None)

    def test_variants_lifecycle(self):
        ts = self.get_example_tree_sequence(random_seed=42)
        variant = _tskit.Variant(ts)
        variant.decode(0)
        genotypes = variant.genotypes
        expected = [1, 0, 0, 1, 0, 0, 0, 0, 1, 1]
        assert np.array_equal(genotypes, expected)
        del variant
        assert np.array_equal(genotypes, expected)
        variant = _tskit.Variant(ts)
        del ts
        variant.decode(0)
        del variant
        assert np.array_equal(genotypes, expected)

    @pytest.mark.parametrize("isolated_as_missing", [True, False])
    @pytest.mark.parametrize("samples", [None, [], (0,)])
    @pytest.mark.parametrize("alleles", [None, ("1", "0")])
    def test_copy(self, isolated_as_missing, samples, alleles):
        ts = self.get_example_tree_sequence(random_seed=42)
        variant = _tskit.Variant(
            ts,
            isolated_as_missing=isolated_as_missing,
            samples=samples,
            alleles=alleles,
        )

        # Test taking a copy before decode
        variant2 = variant.restricted_copy()
        assert variant.site_id == variant2.site_id
        assert variant.alleles == variant2.alleles
        assert np.array_equal(variant.genotypes, variant2.genotypes)
        assert np.array_equal(variant.samples, variant2.samples)
        assert variant.isolated_as_missing == variant2.isolated_as_missing

        variant.decode(0)
        # Everything below should work even if the Python ts is free'd
        del ts
        variant2 = variant.restricted_copy()
        assert variant.site_id == variant2.site_id
        assert variant.alleles == variant2.alleles
        assert np.array_equal(variant.genotypes, variant2.genotypes)
        assert np.array_equal(variant.samples, variant2.samples)
        assert variant.isolated_as_missing == variant2.isolated_as_missing

        # Take a copy for comparison, then move the variant to check the copy
        # doesn't move too
        genotypes = variant.genotypes
        genotypes_copy = np.array(variant.genotypes)
        alleles = variant.alleles
        site_id = variant.site_id
        variant.decode(1)
        with pytest.raises(
            tskit.LibraryError, match="Can't decode a copy of a variant"
        ):
            variant2.decode(1)
        assert site_id == variant2.site_id
        assert alleles == variant2.alleles
        # Other properties shouldn't have changed
        assert np.array_equal(variant.samples, variant2.samples)
        assert variant.isolated_as_missing == variant2.isolated_as_missing

        # Variant should be equal to the copy we took earlier
        assert np.array_equal(genotypes_copy, variant2.genotypes)
        # But not equal to the un-copies genotypes anymore as they
        # have decoded a new site as a side effect of reusing the
        # array when decoding
        assert len(variant.samples) == 0 or not np.array_equal(
            genotypes, variant2.genotypes
        )

        # Check the lifecycle of copies and copies of copies
        del variant
        variant3 = variant2.restricted_copy()
        del variant2
        assert np.array_equal(genotypes_copy, variant3.genotypes)
        genotypes3 = variant3.genotypes
        del variant3
        assert np.array_equal(genotypes_copy, genotypes3)


class TestLdCalculator(LowLevelTestCase):
    """
    Tests for the LdCalculator class.
    """

    def test_uninitialised_tree_sequence(self):
        ts = _tskit.TreeSequence()
        with pytest.raises(ValueError):
            _tskit.LdCalculator(ts)

    def test_constructor(self):
        with pytest.raises(TypeError):
            _tskit.LdCalculator()
        with pytest.raises(TypeError):
            _tskit.LdCalculator(None)

    def test_get_r2(self):
        ts = self.get_example_tree_sequence()
        calc = _tskit.LdCalculator(ts)
        n = ts.get_num_sites()
        for bad_id in [-1, n, n + 1]:
            with pytest.raises(_tskit.LibraryError):
                calc.get_r2(0, bad_id)
            with pytest.raises(_tskit.LibraryError):
                calc.get_r2(bad_id, 0)

    def test_get_r2_array(self):
        ts = self.get_example_tree_sequence()
        calc = _tskit.LdCalculator(ts)

        n = ts.get_num_sites()
        assert n > 2

        with pytest.raises(ValueError):
            calc.get_r2_array(0, max_distance=-1)
        with pytest.raises(ValueError):
            calc.get_r2_array(0, direction=1000)

        for bad_max_sites in [-2, -3]:
            with pytest.raises(ValueError, match="cannot be negative"):
                calc.get_r2_array(0, max_sites=bad_max_sites)
        for bad_start_pos in [-1, n, n + 1]:
            with pytest.raises(_tskit.LibraryError):
                calc.get_r2_array(bad_start_pos)

    def test_r2_array_properties(self):
        ts = self.get_example_tree_sequence()
        calc = _tskit.LdCalculator(ts)
        n = ts.get_num_sites()
        a = calc.get_r2_array(0)
        assert a.shape == (n - 1,)
        assert a.dtype == np.float64
        assert not a.flags.writeable
        assert a.flags.aligned
        assert a.flags.c_contiguous
        assert a.flags.owndata

    def test_r2_array_lifetime(self):
        ts = self.get_example_tree_sequence()
        calc = _tskit.LdCalculator(ts)
        n = ts.get_num_sites()

        a1 = calc.get_r2_array(0)
        assert a1.shape[0] == n - 1
        a2 = a1.copy()
        assert a1 is not a2
        del calc
        # Do some memory operations
        a3 = np.ones(10**6)
        assert np.all(a1 == a2)
        del ts
        assert np.all(a1 == a2)
        del a1
        # Just do something to touch memory
        a2[:] = 0
        assert a3 is not a2


class TestLsHmm(LowLevelTestCase):
    """
    Tests for the LsHmm class.
    """

    def test_uninitialised_tree_sequence(self):
        ts = _tskit.TreeSequence()
        with pytest.raises(ValueError):
            _tskit.LsHmm(ts, None, None)

    def test_constructor(self):
        ts = self.get_example_tree_sequence()
        with pytest.raises(TypeError):
            _tskit.LsHmm()
        with pytest.raises(TypeError):
            _tskit.LsHmm(None)
        values = np.zeros(ts.get_num_sites())
        for bad_array in ["asdf", [[], []], None]:
            with pytest.raises(ValueError):
                _tskit.LsHmm(ts, bad_array, values)
            with pytest.raises(ValueError):
                _tskit.LsHmm(ts, values, bad_array)

    def test_bad_rate_arrays(self):
        ts = self.get_example_tree_sequence()
        m = ts.get_num_sites()
        assert m > 0
        values = np.zeros(m)
        for bad_size in [0, m - 1, m + 1, m + 2]:
            bad_array = np.zeros(bad_size)
            with pytest.raises(ValueError):
                _tskit.LsHmm(ts, bad_array, values)
            with pytest.raises(ValueError):
                _tskit.LsHmm(ts, values, bad_array)

    def test_haplotype_input(self):
        ts = self.get_example_tree_sequence()
        m = ts.get_num_sites()
        fm = _tskit.CompressedMatrix(ts)
        vm = _tskit.ViterbiMatrix(ts)
        norm = np.ones(m)
        ls_hmm = _tskit.LsHmm(ts, np.zeros(m), np.zeros(m))
        for bad_size in [0, m - 1, m + 1, m + 2]:
            bad_array = np.zeros(bad_size, dtype=np.int8)
            with pytest.raises(ValueError, match="haplotype array"):
                ls_hmm.forward_matrix(bad_array, fm)
            with pytest.raises(ValueError, match="haplotype array"):
                ls_hmm.backward_matrix(bad_array, norm, fm)
            with pytest.raises(ValueError, match="haplotype array"):
                ls_hmm.viterbi_matrix(bad_array, vm)
        for bad_array in [[0.002], [[], []], None]:
            with pytest.raises(ValueError):
                ls_hmm.forward_matrix(bad_array, fm)
            with pytest.raises(ValueError):
                ls_hmm.viterbi_matrix(bad_array, vm)
            with pytest.raises(ValueError):
                ls_hmm.backward_matrix(bad_array, norm, fm)

    def test_norm_input(self):
        ts = self.get_example_tree_sequence()
        m = ts.get_num_sites()
        cm = _tskit.CompressedMatrix(ts)
        h = np.zeros(m, dtype=np.int32)
        ls_hmm = _tskit.LsHmm(ts, np.zeros(m), np.zeros(m))
        for bad_size in [0, m - 1, m + 1, m + 2]:
            bad_array = np.zeros(bad_size)
            with pytest.raises(ValueError, match="forward_norm array"):
                ls_hmm.backward_matrix(h, bad_array, cm)

        for bad_array in [[0.002], [[], []], None]:
            with pytest.raises(ValueError):
                ls_hmm.backward_matrix(h, bad_array, cm)

    def test_output_type_errors(self):
        ts = self.get_example_tree_sequence()
        m = ts.get_num_sites()
        h = np.zeros(m, dtype=np.int8)
        norm = np.ones(m)
        ls_hmm = _tskit.LsHmm(ts, np.zeros(m), np.zeros(m))
        for bad_type in [ls_hmm, None, m, []]:
            with pytest.raises(TypeError):
                ls_hmm.forward_matrix(h, bad_type)
            with pytest.raises(TypeError):
                ls_hmm.viterbi_matrix(h, bad_type)
            with pytest.raises(TypeError):
                ls_hmm.backward_matrix(h, norm, bad_type)

        other_ts = self.get_example_tree_sequence()
        output = _tskit.CompressedMatrix(other_ts)
        with pytest.raises(_tskit.LibraryError):
            ls_hmm.forward_matrix(h, output)
        with pytest.raises(_tskit.LibraryError):
            ls_hmm.backward_matrix(h, norm, output)
        output = _tskit.ViterbiMatrix(other_ts)
        with pytest.raises(_tskit.LibraryError):
            ls_hmm.viterbi_matrix(h, output)

    def test_empty_forward_matrix(self):
        for mu in [0, 1]:
            ts = self.get_example_tree_sequence(mutation_rate=mu)
            m = ts.get_num_sites()
            fm = _tskit.CompressedMatrix(ts)
            assert fm.num_sites == m
            assert np.array_equal(np.zeros(m), fm.normalisation_factor)
            assert np.array_equal(np.zeros(m, dtype=np.uint32), fm.num_transitions)
            F = fm.decode()
            assert np.all(F >= 0)
            for j in range(m):
                assert fm.get_site(j) == []

    def test_empty_viterbi_matrix(self):
        for mu in [0, 1]:
            ts = self.get_example_tree_sequence(mutation_rate=mu)
            m = ts.get_num_sites()
            vm = _tskit.ViterbiMatrix(ts)
            assert vm.num_sites == m
            # TODO we should have the same semantics for 0 sites
            if m == 0:
                h = vm.traceback()
                assert len(h) == 0
            else:
                with pytest.raises(_tskit.LibraryError):
                    vm.traceback()

    def verify_compressed_matrix(self, ts, output):
        S = output.normalisation_factor
        N = output.num_transitions
        assert np.all(0 < S)
        assert np.all(S < 1)
        assert np.all(N > 0)
        F = output.decode()
        assert F.shape == (ts.get_num_sites(), ts.get_num_samples())
        assert np.all(F >= 0)
        m = ts.get_num_sites()
        for j in range(m):
            site_list = output.get_site(j)
            assert len(site_list) == N[j]
            for item in site_list:
                assert len(item) == 2
                node, value = item
                assert 0 <= node < ts.get_num_nodes()
                assert value >= 0
        for site in [m, m + 1, 2 * m]:
            with pytest.raises(ValueError):
                output.get_site(site)

    def test_forward_matrix(self):
        ts = self.get_example_tree_sequence()
        m = ts.get_num_sites()
        output = _tskit.CompressedMatrix(ts)
        ls_hmm = _tskit.LsHmm(ts, np.zeros(m) + 0.1, np.zeros(m) + 0.1)
        rv = ls_hmm.forward_matrix([0 for _ in range(m)], output)
        assert rv is None
        self.verify_compressed_matrix(ts, output)

    def test_backward_matrix(self):
        ts = self.get_example_tree_sequence()
        m = ts.get_num_sites()
        fm = _tskit.CompressedMatrix(ts)
        bm = _tskit.CompressedMatrix(ts)
        h = np.zeros(m, dtype=np.int32)
        ls_hmm = _tskit.LsHmm(ts, np.zeros(m) + 0.1, np.zeros(m) + 0.1)
        ls_hmm.forward_matrix(h, fm)
        ls_hmm.backward_matrix(h, fm.normalisation_factor, bm)
        self.verify_compressed_matrix(ts, bm)

    def test_viterbi_matrix(self):
        ts = self.get_example_tree_sequence()
        m = ts.get_num_sites()
        output = _tskit.ViterbiMatrix(ts)
        ls_hmm = _tskit.LsHmm(ts, np.zeros(m) + 0.1, np.zeros(m) + 0.1)
        rv = ls_hmm.viterbi_matrix([0 for _ in range(m)], output)
        assert rv is None
        self.verify_compressed_matrix(ts, output)
        h = output.traceback()
        assert isinstance(h, np.ndarray)


class TestTree(LowLevelTestCase):
    """
    Tests on the low-level tree interface.
    """

    ARRAY_NAMES = [
        "parent",
        "left_child",
        "right_child",
        "left_sib",
        "right_sib",
        "num_children",
        "edge",
    ]

    def test_options(self):
        ts = self.get_example_tree_sequence()
        st = _tskit.Tree(ts)
        assert st.get_options() == 0
        all_options = [
            0,
            _tskit.NO_SAMPLE_COUNTS,
            _tskit.SAMPLE_LISTS,
            _tskit.NO_SAMPLE_COUNTS | _tskit.SAMPLE_LISTS,
        ]
        for options in all_options:
            tree = _tskit.Tree(ts, options=options)
            copy = tree.copy()
            for st in [tree, copy]:
                assert st.get_options() == options
                assert st.get_num_samples(0) == 1
                if options & _tskit.NO_SAMPLE_COUNTS:
                    # We should still be able to count the samples, just inefficiently.
                    assert st.get_num_samples(0) == 1
                    with pytest.raises(_tskit.LibraryError):
                        st.get_num_tracked_samples(0)
                else:
                    assert st.get_num_tracked_samples(0) == 0
                if options & _tskit.SAMPLE_LISTS:
                    assert 0 == st.get_left_sample(0)
                    assert 0 == st.get_right_sample(0)
                else:
                    with pytest.raises(ValueError):
                        st.get_left_sample(0)
                    with pytest.raises(ValueError):
                        st.get_right_sample(0)
                    with pytest.raises(ValueError):
                        st.get_next_sample(0)

    def test_site_errors(self):
        ts = self.get_example_tree_sequence()
        for bad_index in [-1, ts.get_num_sites(), ts.get_num_sites() + 1]:
            with pytest.raises(IndexError):
                ts.get_site(bad_index)

    def test_mutation_errors(self):
        ts = self.get_example_tree_sequence()
        for bad_index in [-1, ts.get_num_mutations(), ts.get_num_mutations() + 1]:
            with pytest.raises(IndexError):
                ts.get_mutation(bad_index)

    def test_individual_errors(self):
        ts = self.get_example_tree_sequence()
        for bad_index in [-1, ts.get_num_individuals(), ts.get_num_individuals() + 1]:
            with pytest.raises(IndexError):
                ts.get_individual(bad_index)

    def test_population_errors(self):
        ts = self.get_example_tree_sequence()
        for bad_index in [-1, ts.get_num_populations(), ts.get_num_populations() + 1]:
            with pytest.raises(IndexError):
                ts.get_population(bad_index)

    def test_provenance_errors(self):
        ts = self.get_example_tree_sequence()
        for bad_index in [-1, ts.get_num_provenances(), ts.get_num_provenances() + 1]:
            with pytest.raises(IndexError):
                ts.get_provenance(bad_index)

    def test_sites(self):
        for ts in self.get_example_tree_sequences():
            st = _tskit.Tree(ts)
            all_sites = [ts.get_site(j) for j in range(ts.get_num_sites())]
            all_tree_sites = []
            j = 0
            mutation_id = 0
            while st.next():
                tree_sites = st.get_sites()
                assert st.get_num_sites() == len(tree_sites)
                all_tree_sites.extend(tree_sites)
                for (
                    position,
                    ancestral_state,
                    mutations,
                    index,
                    metadata,
                ) in tree_sites:
                    assert st.get_left() <= position < st.get_right()
                    assert isinstance(ancestral_state, str)
                    assert index == j
                    assert metadata == b""
                    for mut_id in mutations:
                        (
                            site,
                            node,
                            derived_state,
                            parent,
                            metadata,
                            time,
                            edge,
                        ) = ts.get_mutation(mut_id)
                        assert site == index
                        assert mutation_id == mut_id
                        assert st.get_parent(node) != _tskit.NULL
                        assert metadata == b""
                        assert edge != _tskit.NULL
                        mutation_id += 1
                    j += 1
            assert all_tree_sites == all_sites

    def test_root_threshold_errors(self):
        ts = self.get_example_tree_sequence()
        tree = _tskit.Tree(ts)
        for bad_type in ["", "x", {}]:
            with pytest.raises(TypeError):
                tree.set_root_threshold(bad_type)

        with pytest.raises(_tskit.LibraryError):
            tree.set_root_threshold(0)
        tree.set_root_threshold(2)
        # Setting when not in the null state raises an error
        tree.next()
        with pytest.raises(_tskit.LibraryError):
            tree.set_root_threshold(2)

    def test_seek_errors(self):
        ts = self.get_example_tree_sequence()
        tree = _tskit.Tree(ts)
        for bad_type in ["", "x", {}]:
            with pytest.raises(TypeError):
                tree.seek(bad_type)
        for bad_pos in [-1, 1e6]:
            with pytest.raises(_tskit.LibraryError):
                tree.seek(bad_pos)

    def test_seek_index_errors(self):
        ts = self.get_example_tree_sequence()
        tree = _tskit.Tree(ts)
        for bad_type in ["", "x", {}]:
            with pytest.raises(TypeError):
                tree.seek_index(bad_type)
        for bad_index in [-1, 10**6]:
            with pytest.raises(_tskit.LibraryError):
                tree.seek_index(bad_index)

    def test_root_threshold(self):
        for ts in self.get_example_tree_sequences():
            tree = _tskit.Tree(ts)
            for root_threshold in [1, 2, ts.get_num_samples() * 2]:
                tree.set_root_threshold(root_threshold)
                assert tree.get_root_threshold() == root_threshold
                while tree.next():
                    assert tree.get_root_threshold() == root_threshold
                    with pytest.raises(_tskit.LibraryError):
                        tree.set_root_threshold(2)
                assert tree.get_root_threshold() == root_threshold

    def test_constructor(self):
        with pytest.raises(TypeError):
            _tskit.Tree()
        for bad_type in ["", {}, [], None, 0]:
            with pytest.raises(TypeError):
                _tskit.Tree(bad_type)
        ts = self.get_example_tree_sequence()
        for bad_type in ["", {}, True, 1, None]:
            with pytest.raises(TypeError):
                _tskit.Tree(ts, tracked_samples=bad_type)
        for bad_type in ["", {}, None, []]:
            with pytest.raises(TypeError):
                _tskit.Tree(ts, options=bad_type)
        for ts in self.get_example_tree_sequences():
            st = _tskit.Tree(ts)
            # An uninitialised tree should always be zero.
            samples = ts.get_samples()
            assert st.get_left_child(st.get_virtual_root()) == samples[0]
            assert st.get_right_child(st.get_virtual_root()) == samples[-1]
            assert st.get_left() == 0
            assert st.get_right() == 0
            for j in range(ts.get_num_samples()):
                assert st.get_parent(j) == _tskit.NULL
                assert st.get_children(j) == tuple()
                assert st.get_time(j) == 0

    def test_bad_tracked_samples(self):
        ts = self.get_example_tree_sequence()
        options = 0
        for bad_type in ["", {}, [], None]:
            with pytest.raises(TypeError):
                _tskit.Tree(ts, options=options, tracked_samples=[bad_type])
            with pytest.raises(TypeError):
                _tskit.Tree(
                    ts,
                    options=options,
                    tracked_samples=[1, bad_type],
                )
        for bad_sample in [10**6, -1e6]:
            with pytest.raises(ValueError):
                _tskit.Tree(
                    ts,
                    options=options,
                    tracked_samples=[bad_sample],
                )
            with pytest.raises(ValueError):
                _tskit.Tree(
                    ts,
                    options=options,
                    tracked_samples=[1, bad_sample],
                )
            with pytest.raises(ValueError):
                _tskit.Tree(ts, tracked_samples=[1, bad_sample, 1])

    def test_while_loop_semantics(self):
        for ts in self.get_example_tree_sequences():
            tree = _tskit.Tree(ts)
            # Any mixture of prev and next is OK and gives a valid iteration.
            for _ in range(2):
                j = 0
                while tree.next():
                    assert tree.get_index() == j
                    j += 1
                assert j == ts.get_num_trees()
            for _ in range(2):
                j = ts.get_num_trees()
                while tree.prev():
                    assert tree.get_index() == j - 1
                    j -= 1
                assert j == 0
            j = 0
            while tree.next():
                assert tree.get_index() == j
                j += 1
            assert j == ts.get_num_trees()

    def test_count_all_samples(self):
        for ts in self.get_example_tree_sequences():
            st = _tskit.Tree(ts)
            # Without initialisation we should be 0 samples for every node
            # that is not a sample.
            for j in range(ts.get_num_nodes()):
                count = 1 if j < ts.get_num_samples() else 0
                assert st.get_num_samples(j) == count
                assert st.get_num_tracked_samples(j) == 0
            while st.next():
                nu = get_sample_counts(ts, st)
                nu_prime = [st.get_num_samples(j) for j in range(ts.get_num_nodes())]
                assert nu == nu_prime
                # For tracked samples, this should be all zeros.
                nu = [st.get_num_tracked_samples(j) for j in range(ts.get_num_nodes())]
                assert nu == list(0 for _ in nu)

    def test_count_tracked_samples(self):
        # Ensure that there are some non-binary nodes.
        non_binary = False
        for ts in self.get_example_tree_sequences():
            st = _tskit.Tree(ts)
            while st.next():
                for u in range(ts.get_num_nodes()):
                    if len(st.get_children(u)) > 1:
                        non_binary = True
            samples = [j for j in range(ts.get_num_samples())]
            powerset = itertools.chain.from_iterable(
                itertools.combinations(samples, r) for r in range(len(samples) + 1)
            )
            max_sets = 100
            for _, subset in zip(range(max_sets), map(list, powerset)):
                # Ordering shouldn't make any difference.
                random.shuffle(subset)
                st = _tskit.Tree(ts, tracked_samples=subset)
                while st.next():
                    nu = get_tracked_sample_counts(ts, st, subset)
                    nu_prime = [
                        st.get_num_tracked_samples(j) for j in range(ts.get_num_nodes())
                    ]
                    assert nu == nu_prime
            # Passing duplicated values should raise an error
            sample = 1
            for j in range(2, 20):
                tracked_samples = [sample for _ in range(j)]
                with pytest.raises(_tskit.LibraryError):
                    _tskit.Tree(
                        ts,
                        tracked_samples=tracked_samples,
                    )
        assert non_binary

    def test_bounds_checking(self):
        for ts in self.get_example_tree_sequences():
            n = ts.get_num_nodes()
            st = _tskit.Tree(ts, options=_tskit.SAMPLE_LISTS)
            for v in [-100, -1, n + 1, n + 100, n * 100]:
                with pytest.raises(ValueError):
                    st.get_parent(v)
                with pytest.raises(ValueError):
                    st.get_children(v)
                with pytest.raises(ValueError):
                    st.get_time(v)
                with pytest.raises(ValueError):
                    st.get_left_sample(v)
                with pytest.raises(ValueError):
                    st.get_right_sample(v)
                with pytest.raises(ValueError):
                    st.is_descendant(v, 0)
                with pytest.raises(ValueError):
                    st.is_descendant(0, v)
                with pytest.raises(ValueError):
                    st.depth(v)
            n = ts.get_num_samples()
            for v in [-100, -1, n + 1, n + 100, n * 100]:
                with pytest.raises(ValueError):
                    st.get_next_sample(v)

    def test_mrca_interface(self):
        for ts in self.get_example_tree_sequences():
            num_nodes = ts.get_num_nodes()
            st = _tskit.Tree(ts)
            for v in [num_nodes + 1, 10**6, _tskit.NULL]:
                with pytest.raises(ValueError):
                    st.get_mrca(v, v)
                with pytest.raises(ValueError):
                    st.get_mrca(v, 1)
                with pytest.raises(ValueError):
                    st.get_mrca(1, v)
            # All the mrcas for an uninitialised tree should be _tskit.NULL
            for u, v in itertools.combinations(range(num_nodes), 2):
                assert st.get_mrca(u, v) == _tskit.NULL

    def test_newick_precision(self):
        ts = self.get_example_tree_sequence()
        st = _tskit.Tree(ts)
        assert st.next()
        with pytest.raises(ValueError):
            st.get_newick(root=0, precision=-1)
        with pytest.raises(ValueError):
            st.get_newick(root=0, precision=18)
        with pytest.raises(ValueError):
            st.get_newick(root=0, precision=100)

    def test_newick_legacy_ms(self):
        ts = self.get_example_tree_sequence()
        st = _tskit.Tree(ts)
        assert st.next()
        root = st.get_left_child(st.get_virtual_root())
        ns = st.get_newick(root)
        assert "n0" in ns
        assert ns == st.get_newick(root, legacy_ms_labels=False)
        assert ns != st.get_newick(root, legacy_ms_labels=True)

    def test_cleared_tree(self):
        ts = self.get_example_tree_sequence()
        samples = ts.get_samples()

        def check_tree(tree):
            assert tree.get_index() == -1
            assert tree.get_left_child(tree.get_virtual_root()) == samples[0]
            assert tree.get_num_edges() == 0
            assert tree.get_mrca(0, 1) == _tskit.NULL
            for u in range(ts.get_num_nodes()):
                assert tree.get_parent(u) == _tskit.NULL
                assert tree.get_left_child(u) == _tskit.NULL
                assert tree.get_right_child(u) == _tskit.NULL
                assert tree.get_num_children(u) == 0
                assert tree.get_edge(u) == _tskit.NULL

        tree = _tskit.Tree(ts)
        check_tree(tree)
        while tree.next():
            pass
        check_tree(tree)
        while tree.prev():
            pass
        check_tree(tree)

    def test_newick_interface(self):
        ts = self.get_example_tree_sequence()
        st = _tskit.Tree(ts)
        # TODO this will break when we correctly handle multiple roots.
        assert st.get_newick(0) == "n0;"
        for bad_type in [None, "", [], {}]:
            with pytest.raises(TypeError):
                st.get_newick(0, precision=bad_type)
            with pytest.raises(TypeError):
                st.get_newick(0, buffer_size=bad_type)
            with pytest.raises(TypeError):
                st.get_newick(0, legacy_ms_labels=bad_type)

    def test_newick_buffer_size(self):
        ts = self.get_example_tree_sequence()
        st = _tskit.Tree(ts)
        assert st.next
        u = st.get_left_child(st.get_virtual_root())
        newick = st.get_newick(u)
        assert newick.endswith(";")
        with pytest.raises(ValueError):
            st.get_newick(u, buffer_size=-1)
        with pytest.raises(_tskit.LibraryError):
            st.get_newick(u, buffer_size=1)
        newick2 = st.get_newick(u, len(newick))
        assert newick2 == newick
        with pytest.raises(_tskit.LibraryError):
            st.get_newick(u, buffer_size=len(newick) - 1)

    def test_index(self):
        for ts in self.get_example_tree_sequences():
            st = _tskit.Tree(ts)
            index = 0
            while st.next():
                assert index == st.get_index()
                index += 1

    def test_bad_mutations(self):
        ts = self.get_example_tree_sequence()
        tables = _tskit.TableCollection()
        ts.dump_tables(tables)

        def f(mutations):
            position = []
            node = []
            site = []
            ancestral_state = []
            ancestral_state_offset = [0]
            derived_state = []
            derived_state_offset = [0]
            for j, (p, n) in enumerate(mutations):
                site.append(j)
                position.append(p)
                ancestral_state.append("0")
                ancestral_state_offset.append(ancestral_state_offset[-1] + 1)
                derived_state.append("1")
                derived_state_offset.append(derived_state_offset[-1] + 1)
                node.append(n)
            tables.sites.set_columns(
                dict(
                    position=position,
                    ancestral_state=ancestral_state,
                    ancestral_state_offset=ancestral_state_offset,
                    metadata=None,
                    metadata_offset=None,
                )
            )
            tables.mutations.set_columns(
                dict(
                    site=site,
                    node=node,
                    derived_state=derived_state,
                    derived_state_offset=derived_state_offset,
                    parent=None,
                    metadata=None,
                    metadata_offset=None,
                )
            )
            ts2 = _tskit.TreeSequence()
            ts2.load_tables(tables)

        with pytest.raises(_tskit.LibraryError):
            f([(0.1, -1)])
        length = ts.get_sequence_length()
        u = ts.get_num_nodes()
        for bad_node in [u, u + 1, 2 * u]:
            with pytest.raises(_tskit.LibraryError):
                f([(0.1, bad_node)])
        for bad_pos in [-1, length, length + 1]:
            with pytest.raises(_tskit.LibraryError):
                f([(bad_pos, 0)])

    def test_sample_list(self):
        options = _tskit.SAMPLE_LISTS
        # Note: we're assuming that samples are 0-n here.
        for ts in self.get_example_tree_sequences():
            t = _tskit.Tree(ts, options=options)
            while t.next():
                # All sample nodes should have themselves.
                for j in range(ts.get_num_samples()):
                    assert t.get_left_sample(j) == j
                    assert t.get_right_sample(j) == j

                # All non-tree nodes should have 0
                for j in range(ts.get_num_nodes()):
                    if (
                        t.get_parent(j) == _tskit.NULL
                        and t.get_left_child(j) == _tskit.NULL
                    ):
                        assert t.get_left_sample(j) == _tskit.NULL
                        assert t.get_right_sample(j) == _tskit.NULL
                # The roots should have all samples.
                u = t.get_left_child(t.get_virtual_root())
                samples = []
                while u != _tskit.NULL:
                    sample = t.get_left_sample(u)
                    end = t.get_right_sample(u)
                    while True:
                        samples.append(sample)
                        if sample == end:
                            break
                        sample = t.get_next_sample(sample)
                    u = t.get_right_sib(u)
                assert sorted(samples) == list(range(ts.get_num_samples()))

    def test_equality(self):
        last_ts = None
        for ts in self.get_example_tree_sequences():
            t1 = _tskit.Tree(ts)
            t2 = _tskit.Tree(ts)
            assert t1.equals(t2)
            assert t2.equals(t1)
            while True:
                assert t1.equals(t2)
                assert t2.equals(t1)
                n1 = t1.next()
                assert not t1.equals(t2)
                assert not t2.equals(t1)
                n2 = t2.next()
                assert n1 == n2
                if not n1:
                    break
            if last_ts is not None:
                t2 = _tskit.Tree(last_ts)
                assert not t1.equals(t2)
                assert not t2.equals(t1)
            last_ts = ts

    def test_b2_errors(self):
        ts1 = self.get_example_tree_sequence(10)
        t1 = _tskit.Tree(ts1)
        t1.first()
        with pytest.raises(TypeError):
            t1.get_b2_index()
        with pytest.raises(TypeError):
            t1.get_b2_index("asdf")

    def test_b2(self):
        ts1 = self.get_example_tree_sequence(10)
        t1 = _tskit.Tree(ts1)
        t1.first()
        assert t1.get_b2_index(10) > 0

    def test_num_lineages_errors(self):
        ts1 = self.get_example_tree_sequence(10)
        t1 = _tskit.Tree(ts1)
        t1.first()
        with pytest.raises(TypeError):
            t1.get_num_lineages()
        with pytest.raises(TypeError):
            t1.get_num_lineages("asdf")
        with pytest.raises(_tskit.LibraryError, match="TIME_NONFINITE"):
            t1.get_num_lineages(np.inf)

    def test_num_lineages(self):
        ts1 = self.get_example_tree_sequence(10)
        t1 = _tskit.Tree(ts1)
        t1.first()
        assert t1.get_num_lineages(0) == 10

    def test_kc_distance_errors(self):
        ts1 = self.get_example_tree_sequence(10)
        t1 = _tskit.Tree(ts1, options=_tskit.SAMPLE_LISTS)
        t1.first()
        with pytest.raises(TypeError):
            t1.get_kc_distance()
        with pytest.raises(TypeError):
            t1.get_kc_distance(t1)
        for bad_tree in [None, "tree", 0]:
            with pytest.raises(TypeError):
                t1.get_kc_distance(bad_tree, lambda_=0)
        for bad_value in ["tree", [], None]:
            with pytest.raises(TypeError):
                t1.get_kc_distance(t1, lambda_=bad_value)

        t2 = _tskit.Tree(ts1, options=_tskit.SAMPLE_LISTS)
        # If we don't seek to a specific tree, it has multiple roots (i.e., it's
        # in the null state). This fails because we don't accept multiple roots.
        self.verify_kc_library_error(t2, t2)

        # Different numbers of samples fail.
        ts2 = self.get_example_tree_sequence(11)
        t2 = _tskit.Tree(ts2, options=_tskit.SAMPLE_LISTS)
        t2.first()
        self.verify_kc_library_error(t1, t2)

        # Error when tree not initialized with sample lists
        ts2 = self.get_example_tree_sequence(10)
        t2 = _tskit.Tree(ts2)
        t2.first()
        with pytest.raises(
            _tskit.NoSampleListsError, match="requires that sample lists are stored"
        ):
            self.verify_kc_library_error(t1, t2)

        # Unary nodes cause errors.
        tables = _tskit.TableCollection(1.0)
        tables.nodes.add_row(flags=1)
        tables.nodes.add_row(flags=1, time=1)
        tables.edges.add_row(0, 1, 1, 0)
        tables.build_index()
        ts = _tskit.TreeSequence()
        ts.load_tables(tables)
        t1 = _tskit.Tree(ts, options=_tskit.SAMPLE_LISTS)
        t1.first()
        self.verify_kc_library_error(t1, t1)

    def verify_kc_library_error(self, t1, t2):
        with pytest.raises(_tskit.LibraryError):
            t1.get_kc_distance(t2, 0)

    def test_kc_distance(self):
        ts1 = self.get_example_tree_sequence(10, random_seed=123456)
        t1 = _tskit.Tree(ts1, options=_tskit.SAMPLE_LISTS)
        t1.first()
        ts2 = self.get_example_tree_sequence(10, random_seed=1234)
        t2 = _tskit.Tree(ts2, options=_tskit.SAMPLE_LISTS)
        t2.first()
        for lambda_ in [-1, 0, 1, 1000, -1e300]:
            x1 = t1.get_kc_distance(t2, lambda_)
            x2 = t2.get_kc_distance(t1, lambda_)
            assert x1 == x2

    def test_copy(self):
        for ts in self.get_example_tree_sequences():
            t1 = _tskit.Tree(ts)
            t2 = t1.copy()
            assert t1.get_index() == t2.get_index()
            assert t1 is not t2
            while t1.next():
                t2 = t1.copy()
                assert t1.get_index() == t2.get_index()

    def test_map_mutations_null(self):
        ts = self.get_example_tree_sequence()
        tree = _tskit.Tree(ts)
        n = ts.get_num_samples()
        genotypes = np.zeros(n, dtype=np.int8)
        ancestral_state, transitions = tree.map_mutations(genotypes)
        assert ancestral_state == 0
        assert len(transitions) == 0

        genotypes = np.arange(n, dtype=np.int8)
        ancestral_state, transitions = tree.map_mutations(genotypes)
        assert ancestral_state == 0
        assert len(transitions) == n - 1
        for j in range(n - 1):
            x = n - j - 1
            assert transitions[j][0] == x
            assert transitions[j][1] == -1
            assert transitions[j][2] == x

    def test_map_mutations(self):
        ts = self.get_example_tree_sequence()
        tree = _tskit.Tree(ts)
        tree.next()
        n = ts.get_num_samples()
        genotypes = np.zeros(n, dtype=np.int8)
        ancestral_state, transitions = tree.map_mutations(genotypes)
        assert ancestral_state == 0
        assert len(transitions) == 0

    def test_map_mutations_fixed_ancestral_state(self):
        ts = self.get_example_tree_sequence()
        tree = _tskit.Tree(ts)
        tree.next()
        n = ts.get_num_samples()
        genotypes = np.ones(n, dtype=np.int8)
        ancestral_state, transitions = tree.map_mutations(genotypes, 0)
        assert ancestral_state == 0
        assert len(transitions) == 1

    def test_map_mutations_errors(self):
        ts = self.get_example_tree_sequence()
        tree = _tskit.Tree(ts)
        n = ts.get_num_samples()
        with pytest.raises(TypeError):
            tree.map_mutations()
        for bad_size in [0, 1, n - 1, n + 1]:
            with pytest.raises(ValueError):
                tree.map_mutations(np.zeros(bad_size, dtype=np.int8))
        for bad_type in [None, {}, set()]:
            with pytest.raises(TypeError):
                tree.map_mutations([bad_type] * n)
        for bad_type in [np.uint32, np.uint64, np.float32]:
            with pytest.raises(TypeError):
                tree.map_mutations(np.zeros(n, dtype=bad_type))
        genotypes = np.zeros(n, dtype=np.int8)
        tree.map_mutations(genotypes)
        for bad_value in [64, 65, 127, -2]:
            genotypes[0] = bad_value
            with pytest.raises(_tskit.LibraryError):
                tree.map_mutations(genotypes)

        genotypes = np.zeros(n, dtype=np.int8)
        tree.map_mutations(genotypes)
        for bad_type in ["d", []]:
            with pytest.raises(TypeError):
                tree.map_mutations(genotypes, bad_type)
        for bad_state in [-2, -1, 127, 255]:
            with pytest.raises(_tskit.LibraryError, match="Bad ancestral"):
                tree.map_mutations(genotypes, bad_state)

    @pytest.mark.parametrize("array", ARRAY_NAMES)
    def test_array_read_only(self, array):
        name = array + "_array"
        ts1 = self.get_example_tree_sequence(10)
        t1 = _tskit.Tree(ts1)
        t1.first()
        with pytest.raises(AttributeError, match="not writable"):
            setattr(t1, name, None)
        with pytest.raises(AttributeError, match="not writable"):
            delattr(t1, name)

        a = getattr(t1, name)
        with pytest.raises(ValueError, match="assignment destination"):
            a[:] = 0
        with pytest.raises(ValueError, match="assignment destination"):
            a[0] = 0
        with pytest.raises(ValueError, match="cannot set WRITEABLE"):
            a.setflags(write=True)

    @pytest.mark.parametrize("array", ARRAY_NAMES)
    def test_array_properties(self, array):
        ts1 = self.get_example_tree_sequence(10)
        t1 = _tskit.Tree(ts1)
        a = getattr(t1, array + "_array")
        t1.first()
        a = getattr(t1, array + "_array")
        assert a.dtype == np.int32
        assert a.shape == (ts1.get_num_nodes() + 1,)
        assert a.base == t1
        assert not a.flags.writeable
        assert a.flags.aligned
        assert a.flags.c_contiguous
        assert not a.flags.owndata
        b = getattr(t1, array + "_array")
        assert a is not b
        assert np.all(a == b)
        a_copy = a.copy()
        # This checks that the underlying pointer to memory is the same in
        # both arrays.
        assert a.__array_interface__ == b.__array_interface__
        t1.next()
        # NB! Because we are pointing to the underlying memory, the arrays
        # will change as we iterate along the trees! This is a gotcha, but
        # it's just something we have to document as it's a consequence of the
        # zero copy semantics.
        b = getattr(t1, array + "_array")
        assert np.all(a == b)
        assert np.any(a_copy != b)

    @pytest.mark.parametrize("array", ARRAY_NAMES)
    def test_array_lifetime(self, array):
        ts1 = self.get_example_tree_sequence(10)
        t1 = _tskit.Tree(ts1)
        t1.first()
        a1 = getattr(t1, array + "_array")
        a2 = a1.copy()
        assert a1 is not a2
        del t1
        # Do some memory operations
        a3 = np.ones(10**6)
        assert np.all(a1 == a2)
        del ts1
        assert np.all(a1 == a2)
        del a1
        # Just do something to touch memory
        a2[:] = 0
        assert a3 is not a2

    @pytest.mark.parametrize("ordering", ["preorder", "postorder"])
    def test_traversal_arrays(self, ordering):
        ts = self.get_example_tree_sequence(10)
        tree = _tskit.Tree(ts)
        tree.first()
        method = getattr(tree, "get_" + ordering)
        for bad_type in [None, {}]:
            with pytest.raises(TypeError):
                method(bad_type)
        for bad_node in [-2, 10**6]:
            with pytest.raises(_tskit.LibraryError, match="out of bounds"):
                method(bad_node)
        a = method(tree.get_virtual_root())
        assert a.dtype == np.int32
        assert not a.flags.writeable
        assert a.flags.aligned
        assert a.flags.c_contiguous
        assert a.flags.owndata


class TestTableMetadataSchema(MetadataTestMixin):
    def test_metadata_schema_attribute(self):
        tables = _tskit.TableCollection(1.0)
        for table in self.metadata_tables:
            table = getattr(tables, f"{table}s")
            # Check default value
            assert table.metadata_schema == ""
            # Set and read back
            example = "An example of metadata schema with unicode 🎄🌳🌴🌲🎋"
            table.metadata_schema = example
            assert table.metadata_schema == example
            # Can't del, or set to None
            with pytest.raises(AttributeError):
                del table.metadata_schema
            with pytest.raises(TypeError):
                table.metadata_schema = None
            # Del or None had no effect
            assert table.metadata_schema == example
            # Clear and read back
            table.metadata_schema = ""
            assert table.metadata_schema == ""


class TestMetadataSchemaNamedTuple(MetadataTestMixin):
    def test_named_tuple_init(self):
        # Test init errors
        with pytest.raises(TypeError):
            metadata_schemas = _tskit.MetadataSchemas()
        with pytest.raises(TypeError):
            metadata_schemas = _tskit.MetadataSchemas([])
        with pytest.raises(TypeError):
            metadata_schemas = _tskit.MetadataSchemas(["test_schema"])
        # Set and read back
        metadata_schemas = _tskit.MetadataSchemas(
            f"{table}_test_schema" for table in self.metadata_tables
        )
        assert metadata_schemas == tuple(
            f"{table}_test_schema" for table in self.metadata_tables
        )
        for i, table in enumerate(self.metadata_tables):
            # Read back via attr, index
            assert getattr(metadata_schemas, table) == f"{table}_test_schema"
            assert metadata_schemas[i] == f"{table}_test_schema"
            # Check read-only
            with pytest.raises(AttributeError):
                setattr(metadata_schemas, table, "")
            with pytest.raises(TypeError):
                metadata_schemas[i] = ""
        # Equality
        metadata_schemas2 = _tskit.MetadataSchemas(
            f"{table}_test_schema" for table in self.metadata_tables
        )
        assert metadata_schemas == metadata_schemas2
        metadata_schemas3 = _tskit.MetadataSchemas(
            f"{table}_test_schema_diff" for table in self.metadata_tables
        )
        assert metadata_schemas != metadata_schemas3


class TestReferenceSequenceInputErrors:
    @pytest.mark.parametrize("bad_type", [1234, b"bytes", None, {}])
    @pytest.mark.parametrize("attr", ["data", "url", "metadata_schema"])
    def test_string_bad_type(self, attr, bad_type):
        refseq = _tskit.TableCollection().reference_sequence
        with pytest.raises(TypeError, match=f"{attr} must be a string"):
            setattr(refseq, attr, bad_type)

    @pytest.mark.parametrize("bad_type", [1234, "unicode", None, {}])
    def test_metadata_bad_type(self, bad_type):
        refseq = _tskit.TableCollection().reference_sequence
        with pytest.raises(TypeError):
            refseq.metadata = bad_type

    @pytest.mark.parametrize("attr", ["data", "url", "metadata_schema"])
    def test_unicode_error(self, attr):
        refseq = _tskit.TableCollection().reference_sequence
        with pytest.raises(UnicodeEncodeError):
            setattr(refseq, attr, NON_UTF8_STRING)

    @pytest.mark.parametrize("attr", ["data", "url", "metadata", "metadata_schema"])
    def test_del_attr(self, attr):
        refseq = _tskit.TableCollection().reference_sequence
        with pytest.raises(AttributeError, match=f"Cannot del {attr}"):
            delattr(refseq, attr)


class TestReferenceSequenceUpdates:
    @pytest.mark.parametrize("value", ["abc", "🎄🌳🌴🌲🎋"])
    @pytest.mark.parametrize("attr", ["data", "url", "metadata_schema"])
    def test_set_string(self, attr, value):
        refseq = _tskit.TableCollection().reference_sequence
        assert refseq.is_null()
        setattr(refseq, attr, value)
        assert getattr(refseq, attr) == value
        assert not refseq.is_null()

    @pytest.mark.parametrize("attr", ["data", "url", "metadata_schema"])
    def test_set_string_null_none(self, attr):
        refseq = _tskit.TableCollection().reference_sequence
        assert refseq.is_null()
        setattr(refseq, attr, "a")
        assert not refseq.is_null()
        setattr(refseq, attr, "")
        assert refseq.is_null()

    @pytest.mark.parametrize("value", [b"x", b"{}", b"abc\0defg"])
    def test_set_metadata(self, value):
        refseq = _tskit.TableCollection().reference_sequence
        assert refseq.is_null()
        refseq.metadata = value
        assert not refseq.is_null()
        refseq.metadata = b""
        assert refseq.is_null()


class TestReferenceSequenceTableCollection:
    def test_references(self):
        tables = _tskit.TableCollection()
        refseq = tables.reference_sequence
        assert refseq is not tables.reference_sequence

    def test_state(self):
        tables = _tskit.TableCollection()
        refseq = tables.reference_sequence
        assert refseq.is_null()
        assert not tables.has_reference_sequence()
        # Setting any non empty string changes the state to "non-null"
        refseq.data = "x"
        assert tables.has_reference_sequence()
        assert not refseq.is_null()

    @pytest.mark.parametrize("ref_data", ["abc", "A" * 10, "🎄🌳🌴🌲🎋"])
    def test_data(self, ref_data):
        tables = _tskit.TableCollection()
        refseq = tables.reference_sequence
        assert refseq.data == ""
        refseq.data = ref_data
        assert refseq.data == ref_data
        assert tables.reference_sequence.data == ref_data

    @pytest.mark.parametrize("url", ["", "abc", "A" * 10, "🎄🌳🌴🌲🎋"])
    def test_url(self, url):
        tables = _tskit.TableCollection()
        refseq = tables.reference_sequence
        assert refseq.url == ""
        refseq.url = url
        assert refseq.url == url
        assert tables.reference_sequence.url == url

    def test_metadata_default_none(self):
        tables = _tskit.TableCollection()
        assert tables.reference_sequence.metadata_schema == ""
        assert tables.reference_sequence.metadata == b""

    # we don't actually check the form here, just pass in and out strings
    @pytest.mark.parametrize("schema", ["", "{}", "abcdefg"])
    def test_metadata_schema(self, schema):
        tables = _tskit.TableCollection()
        tables.reference_sequence.metadata_schema = schema
        assert tables.has_reference_sequence
        assert tables.reference_sequence.metadata_schema == schema

    @pytest.mark.parametrize("metadata", [b"", b"{}", b"abcdefg"])
    def test_metadata(self, metadata):
        tables = _tskit.TableCollection()
        tables.reference_sequence.metadata = metadata
        assert tables.has_reference_sequence
        assert tables.reference_sequence.metadata == metadata


class TestReferenceSequenceTreeSequence:
    def test_references(self):
        tc = _tskit.TableCollection()
        tc.sequence_length = 1
        ts = _tskit.TreeSequence()
        ts.load_tables(tc, build_indexes=True)
        refseq = ts.reference_sequence
        assert refseq is not ts.reference_sequence
        assert refseq is not tc

    def test_state(self):
        tc = _tskit.TableCollection()
        tc.sequence_length = 1
        ts = _tskit.TreeSequence()
        ts.load_tables(tc, build_indexes=True)
        assert not ts.has_reference_sequence()

    def test_write(self):
        tc = _tskit.TableCollection()
        tc.sequence_length = 1
        ts = _tskit.TreeSequence()
        ts.load_tables(tc, build_indexes=True)
        refseq = ts.reference_sequence
        with pytest.raises(AttributeError, match="read-only"):
            refseq.data = "asdf"
        with pytest.raises(AttributeError, match="read-only"):
            refseq.url = "asdf"
        with pytest.raises(AttributeError, match="read-only"):
            refseq.metadata_schema = "asdf"
        with pytest.raises(AttributeError, match="read-only"):
            refseq.metadata = "asdf"

    @pytest.mark.parametrize("ref_data", ["", "ACTG" * 10, "🎄🌳🌴🌲🎋"])
    def test_data(self, ref_data):
        tc = _tskit.TableCollection()
        tc.sequence_length = 1
        tc.reference_sequence.data = ref_data
        ts = _tskit.TreeSequence()
        ts.load_tables(tc, build_indexes=True)
        assert ts.reference_sequence.data == ref_data

    @pytest.mark.parametrize("url", ["", "ACTG" * 10, "🎄🌳🌴🌲🎋"])
    def test_url(self, url):
        tc = _tskit.TableCollection()
        tc.sequence_length = 1
        tc.reference_sequence.url = url
        ts = _tskit.TreeSequence()
        ts.load_tables(tc, build_indexes=True)
        assert ts.reference_sequence.url == url

    # we don't actually check the form here, just pass in and out strings
    @pytest.mark.parametrize("schema", ["", "{}", "abcdefg"])
    def test_metadata_schema(self, schema):
        tc = _tskit.TableCollection()
        tc.sequence_length = 1
        tc.reference_sequence.metadata_schema = schema
        ts = _tskit.TreeSequence()
        ts.load_tables(tc, build_indexes=True)
        assert ts.has_reference_sequence
        assert ts.reference_sequence.metadata_schema == schema

    @pytest.mark.parametrize("metadata", [b"", b"{}", b"abcdefg"])
    def test_metadata(self, metadata):
        tc = _tskit.TableCollection()
        tc.sequence_length = 1
        tc.reference_sequence.metadata = metadata
        ts = _tskit.TreeSequence()
        ts.load_tables(tc, build_indexes=True)
        assert ts.has_reference_sequence
        assert ts.reference_sequence.metadata == metadata


class TestModuleFunctions:
    """
    Tests for the module level functions.
    """

    def test_kastore_version(self):
        version = _tskit.get_kastore_version()
        assert version == (2, 1, 1)

    def test_tskit_version(self):
        version = _tskit.get_tskit_version()
        assert version == (1, 1, 3)

    def test_tskit_version_file(self):
        maj, min_, patch = _tskit.get_tskit_version()
        with open(f"{tskit.__path__[0]}/../../c/VERSION.txt") as f:
            assert f.read() == f"{maj}.{min_}.{patch}"


def test_uninitialised():
    # These methods work from an instance that has a NULL ref so don't check
    skip_list = [
        "TableCollection_load",
        "TreeSequence_load",
        "TreeSequence_load_tables",
    ]
    for cls_name, cls in inspect.getmembers(_tskit):
        if (
            isinstance(cls, type)
            and not issubclass(cls, Exception)
            and not issubclass(cls, tuple)
        ):
            methods = []
            attributes = []
            for name, value in inspect.getmembers(cls):
                if not name.startswith("__") and f"{cls_name}_{name}" not in skip_list:
                    if inspect.isdatadescriptor(value):
                        attributes.append(name)
                    else:
                        methods.append(name)
            uninitialised = cls.__new__(cls)
            for attr in attributes:
                with pytest.raises((SystemError, ValueError)):
                    getattr(uninitialised, attr)
                with pytest.raises((SystemError, ValueError, AttributeError)):
                    setattr(uninitialised, attr, None)
            for method_name in methods:
                method = getattr(uninitialised, method_name)
                with pytest.raises((SystemError, ValueError)):
                    method()


def test_constants():
    assert _tskit.TIME_UNITS_UNKNOWN == "unknown"
    assert _tskit.TIME_UNITS_UNCALIBRATED == "uncalibrated"


class TestPairCoalescenceCountsErrors:
    def example_ts(self, sample_size=10):
        ts = msprime.sim_ancestry(
            sample_size,
            sequence_length=1e4,
            recombination_rate=1e-8,
            random_seed=1,
            population_size=1e4,
        )
        return ts.ll_tree_sequence

    @staticmethod
    def pair_coalescence_counts(
        ts,
        sample_sets=None,
        sample_set_sizes=None,
        indexes=None,
        windows=None,
        node_bin_map=None,
        span_normalise=False,
        pair_normalise=False,
    ):
        n = ts.get_num_samples()
        N = ts.get_num_nodes()
        if sample_sets is None:
            sample_sets = np.arange(n, dtype=np.int32)
        if sample_set_sizes is None:
            sample_set_sizes = [n // 2, n - n // 2]
        if indexes is None:
            pairs = itertools.combinations_with_replacement(
                range(len(sample_set_sizes)), 2
            )
            indexes = [(i, j) for i, j in pairs]
        if windows is None:
            windows = np.array([0, 0.5, 1.0]) * ts.get_sequence_length()
        if node_bin_map is None:
            node_bin_map = np.arange(N, dtype=np.int32)
        return ts.pair_coalescence_counts(
            sample_sets=sample_sets,
            sample_set_sizes=sample_set_sizes,
            windows=windows,
            indexes=indexes,
            node_bin_map=node_bin_map,
            span_normalise=span_normalise,
            pair_normalise=pair_normalise,
        )

    def test_output_dims(self):
        ts = self.example_ts()
        coal = self.pair_coalescence_counts(ts)
        dim = (2, 3, ts.get_num_nodes())
        assert coal.shape == dim
        coal = self.pair_coalescence_counts(ts, span_normalise=True)
        assert coal.shape == dim
        coal = self.pair_coalescence_counts(ts, pair_normalise=True)
        assert coal.shape == dim

    def test_node_shuffle(self):
        rng = np.random.default_rng(1024)
        ts = self.example_ts()
        coal = self.pair_coalescence_counts(ts)
        node_bin_map = np.arange(ts.get_num_nodes(), dtype=np.int32)
        rng.shuffle(node_bin_map)
        coal_shuffle = self.pair_coalescence_counts(ts, node_bin_map=node_bin_map)
        np.testing.assert_allclose(coal_shuffle[..., node_bin_map], coal)

    @pytest.mark.parametrize("bad_node", [-1, -2, 1000])
    def test_c_tsk_err_node_out_of_bounds(self, bad_node):
        ts = self.example_ts()
        ids = np.arange(ts.get_num_samples(), dtype=np.int32)
        with pytest.raises(_tskit.LibraryError, match="TSK_ERR_NODE_OUT_OF_BOUNDS"):
            self.pair_coalescence_counts(
                ts, sample_sets=np.append(ids[:-1], bad_node).astype(np.int32)
            )

    def test_c_tsk_err_bad_windows(self):
        ts = self.example_ts()
        L = ts.get_sequence_length()
        with pytest.raises(_tskit.LibraryError, match="TSK_ERR_BAD_WINDOWS"):
            self.pair_coalescence_counts(ts, windows=[1.0, L])

    def test_c_tsk_err_bad_node_bin_map(self):
        ts = self.example_ts()
        node_bin_map = np.arange(ts.get_num_nodes(), dtype=np.int32)
        node_bin_map[0] = -10
        with pytest.raises(_tskit.LibraryError, match="BAD_NODE_BIN_MAP"):
            self.pair_coalescence_counts(ts, node_bin_map=node_bin_map)

    @pytest.mark.parametrize("bad_index", [-1, 10])
    def test_c_tsk_err_bad_sample_set_index(self, bad_index):
        ts = self.example_ts()
        with pytest.raises(_tskit.LibraryError, match="BAD_SAMPLE_SET_INDEX"):
            self.pair_coalescence_counts(ts, indexes=[(0, bad_index)])

    @pytest.mark.parametrize("bad_ss_size", [-1, 1000])
    def test_cpy_bad_sample_sets(self, bad_ss_size):
        ts = self.example_ts()
        with pytest.raises(
            (ValueError, OverflowError),
            match="Sum of sample_set_sizes|Overflow|out of bounds",
        ):
            self.pair_coalescence_counts(
                ts, sample_set_sizes=[bad_ss_size, ts.get_num_samples()]
            )

    def test_cpy_bad_parse_inputs(self):
        ts = self.example_ts()
        with pytest.raises(TypeError, match="str"):
            self.pair_coalescence_counts(ts, span_normalise="foo")

    def test_cpy_bad_windows(self):
        ts = self.example_ts()
        with pytest.raises(ValueError, match="at least 2"):
            self.pair_coalescence_counts(ts, windows=[0.0])

    @pytest.mark.parametrize("indexes", [[(0, 0, 0)], np.zeros((0, 2), dtype=np.int32)])
    def test_cpy_bad_indexes(self, indexes):
        ts = self.example_ts()
        with pytest.raises(ValueError, match="k x 2 array"):
            self.pair_coalescence_counts(ts, indexes=indexes)
        with pytest.raises(ValueError, match="too small depth"):
            self.pair_coalescence_counts(ts, indexes=np.ravel(indexes))

    def test_cpy_bad_node_bin_map(self):
        ts = self.example_ts()
        num_nodes = ts.get_num_nodes()
        node_bin_map = np.full(num_nodes, tskit.NULL, dtype=np.int32)
        with pytest.raises(ValueError, match="null values for all nodes"):
            self.pair_coalescence_counts(ts, node_bin_map=node_bin_map)
        with pytest.raises(ValueError, match="a value per node"):
            self.pair_coalescence_counts(ts, node_bin_map=node_bin_map[:-1])
        with pytest.raises(TypeError, match="cast array data"):
            self.pair_coalescence_counts(ts, node_bin_map=np.zeros(num_nodes))


class TestPairCoalescenceQuantilesErrors:
    def example_ts(self, sample_size=10):
        ts = msprime.sim_ancestry(
            sample_size,
            sequence_length=1e4,
            recombination_rate=1e-8,
            random_seed=1,
            population_size=1e4,
        )
        return ts.ll_tree_sequence

    @staticmethod
    def pair_coalescence_quantiles(
        ts,
        quantiles=None,
        sample_sets=None,
        sample_set_sizes=None,
        indexes=None,
        windows=None,
        node_bin_map=None,
    ):
        n = ts.get_num_samples()
        if quantiles is None:
            quantiles = np.linspace(0, 1, 4)
        if sample_sets is None:
            sample_sets = np.arange(n, dtype=np.int32)
        if sample_set_sizes is None:
            sample_set_sizes = [n // 2, n - n // 2]
        if indexes is None:
            pairs = itertools.combinations_with_replacement(
                range(len(sample_set_sizes)), 2
            )
            indexes = [(i, j) for i, j in pairs]
        if windows is None:
            windows = np.array([0, 0.5, 1.0]) * ts.get_sequence_length()
        if node_bin_map is None:
            _, node_bin_map = np.unique(ts.nodes_time, return_inverse=True)
            node_bin_map = node_bin_map.astype(np.int32)
        return ts.pair_coalescence_quantiles(
            sample_sets=sample_sets,
            sample_set_sizes=sample_set_sizes,
            windows=windows,
            indexes=indexes,
            node_bin_map=node_bin_map,
            quantiles=quantiles,
        )

    def test_output_dims(self):
        ts = self.example_ts()
        coal = self.pair_coalescence_quantiles(ts)
        dim = (2, 3, 4)
        assert coal.shape == dim

    @pytest.mark.parametrize("quantiles", [[1.0, 0.0], [-1.0], [2.0]])
    def test_c_tsk_err_bad_quantiles(self, quantiles):
        ts = self.example_ts()
        with pytest.raises(_tskit.LibraryError, match="TSK_ERR_BAD_QUANTILES"):
            self.pair_coalescence_quantiles(ts, quantiles=quantiles)

    def test_c_tsk_err_unsorted_times(self):
        ts = self.example_ts()
        _, node_bin_map = np.unique(ts.nodes_time, return_inverse=True)
        node_bin_map = node_bin_map[::-1]
        node_bin_map = node_bin_map.astype(np.int32)
        with pytest.raises(_tskit.LibraryError, match="TSK_ERR_UNSORTED_TIMES"):
            self.pair_coalescence_quantiles(ts, node_bin_map=node_bin_map)

    def test_c_tsk_err_bad_windows(self):
        ts = self.example_ts()
        with pytest.raises(_tskit.LibraryError, match="TSK_ERR_BAD_WINDOWS"):
            self.pair_coalescence_quantiles(ts, windows=[1.0, ts.get_sequence_length()])

    @pytest.mark.parametrize("bad_ss_size", [-1, 1000])
    def test_cpy_bad_sample_sets(self, bad_ss_size):
        ts = self.example_ts()
        with pytest.raises(
            (ValueError, OverflowError),
            match="Sum of sample_set_sizes|Overflow|out of bounds",
        ):
            self.pair_coalescence_quantiles(
                ts, sample_set_sizes=[bad_ss_size, ts.get_num_samples()]
            )

    def test_cpy_bad_windows(self):
        ts = self.example_ts()
        with pytest.raises(ValueError, match="at least 2"):
            self.pair_coalescence_quantiles(ts, windows=[0.0])

    @pytest.mark.parametrize("indexes", [[(0, 0, 0)], np.zeros((0, 2), dtype=np.int32)])
    def test_cpy_bad_indexes(self, indexes):
        ts = self.example_ts()
        with pytest.raises(ValueError, match="k x 2 array"):
            self.pair_coalescence_quantiles(ts, indexes=indexes)
        with pytest.raises(ValueError, match="too small depth"):
            self.pair_coalescence_quantiles(ts, indexes=np.ravel(indexes))

    def test_cpy_bad_node_bin_map(self):
        ts = self.example_ts()
        num_nodes = ts.get_num_nodes()
        node_bin_map = np.full(num_nodes, tskit.NULL, dtype=np.int32)
        with pytest.raises(ValueError, match="null values for all nodes"):
            self.pair_coalescence_quantiles(ts, node_bin_map=node_bin_map)
        with pytest.raises(ValueError, match="a value per node"):
            self.pair_coalescence_quantiles(ts, node_bin_map=node_bin_map[:-1])
        with pytest.raises(TypeError, match="cast array data"):
            self.pair_coalescence_quantiles(ts, node_bin_map=np.zeros(num_nodes))

    def test_cpy_bad_quantiles(self):
        ts = self.example_ts()
        quantiles = np.zeros(0)
        with pytest.raises(ValueError, match="at least one quantile"):
            self.pair_coalescence_quantiles(ts, quantiles=quantiles)
        quantiles = np.zeros((3, 3))
        with pytest.raises(ValueError, match="object too deep"):
            self.pair_coalescence_quantiles(ts, quantiles=quantiles)

    def test_cpy_bad_inputs(self):
        ts = self.example_ts()
        with pytest.raises(TypeError, match="at most 6 keyword"):
            ts.pair_coalescence_quantiles(
                sample_sets=None,
                sample_set_sizes=None,
                windows=None,
                quantiles=None,
                indexes=None,
                node_bin_map=None,
                foo="bar",
            )


class TestPairCoalescenceRatesErrors:
    def example_ts(self, sample_size=10):
        ts = msprime.sim_ancestry(
            sample_size,
            sequence_length=1e4,
            recombination_rate=1e-8,
            random_seed=1,
            population_size=1e4,
        )
        return ts.ll_tree_sequence

    @staticmethod
    def pair_coalescence_rates(
        ts,
        time_windows=None,
        sample_sets=None,
        sample_set_sizes=None,
        indexes=None,
        windows=None,
        node_bin_map=None,
    ):
        n = ts.get_num_samples()
        if time_windows is None:
            time_windows = np.array([0.0, np.mean(ts.nodes_time), np.inf])
        if sample_sets is None:
            sample_sets = np.arange(n, dtype=np.int32)
        if sample_set_sizes is None:
            sample_set_sizes = [n // 2, n - n // 2]
        if indexes is None:
            pairs = itertools.combinations_with_replacement(
                range(len(sample_set_sizes)), 2
            )
            indexes = [(i, j) for i, j in pairs]
        if windows is None:
            windows = np.array([0, 0.5, 1.0]) * ts.get_sequence_length()
        if node_bin_map is None:
            node_bin_map = np.digitize(ts.nodes_time, time_windows) - 1
            node_bin_map[node_bin_map == time_windows.size - 1] = tskit.NULL
            node_bin_map = node_bin_map.astype(np.int32)
        return ts.pair_coalescence_rates(
            sample_sets=sample_sets,
            sample_set_sizes=sample_set_sizes,
            windows=windows,
            indexes=indexes,
            node_bin_map=node_bin_map,
            time_windows=time_windows,
        )

    def test_output_dims(self):
        ts = self.example_ts()
        coal = self.pair_coalescence_rates(ts)
        dim = (2, 3, 2)
        assert coal.shape == dim

    def test_c_tsk_err_bad_time_windows(self):
        ts = self.example_ts()
        with pytest.raises(_tskit.LibraryError, match="TSK_ERR_BAD_TIME_WINDOWS"):
            self.pair_coalescence_rates(ts, time_windows=np.array([np.inf, 0.0]))

    def test_c_tsk_err_bad_node_time_window(self):
        ts = self.example_ts()
        node_bin_map = np.zeros(ts.nodes_time.size, dtype=np.int32)
        with pytest.raises(_tskit.LibraryError, match="TSK_ERR_BAD_NODE_TIME_WINDOW"):
            self.pair_coalescence_rates(ts, node_bin_map=node_bin_map)

    def test_c_tsk_err_bad_sample_pair_times(self):
        ts = self.example_ts()
        with pytest.raises(_tskit.LibraryError, match="TSK_ERR_BAD_SAMPLE_PAIR_TIMES"):
            self.pair_coalescence_rates(ts, time_windows=np.array([-1.0, np.inf]))

    def test_c_tsk_err_bad_windows(self):
        ts = self.example_ts()
        with pytest.raises(_tskit.LibraryError, match="TSK_ERR_BAD_WINDOWS"):
            self.pair_coalescence_rates(ts, windows=[1.0, ts.get_sequence_length()])

    @pytest.mark.parametrize("bad_ss_size", [-1, 1000])
    def test_cpy_bad_sample_sets(self, bad_ss_size):
        ts = self.example_ts()
        with pytest.raises(
            (ValueError, OverflowError),
            match="Sum of sample_set_sizes|Overflow|out of bounds",
        ):
            self.pair_coalescence_rates(
                ts, sample_set_sizes=[bad_ss_size, ts.get_num_samples()]
            )

    @pytest.mark.parametrize("indexes", [[(0, 0, 0)], np.zeros((0, 2), dtype=np.int32)])
    def test_cpy_bad_indexes(self, indexes):
        ts = self.example_ts()
        with pytest.raises(ValueError, match="k x 2 array"):
            self.pair_coalescence_rates(ts, indexes=indexes)
        with pytest.raises(ValueError, match="too small depth"):
            self.pair_coalescence_rates(ts, indexes=np.ravel(indexes))

    def test_cpy_bad_node_bin_map(self):
        ts = self.example_ts()
        num_nodes = ts.get_num_nodes()
        node_bin_map = np.full(num_nodes, tskit.NULL, dtype=np.int32)
        with pytest.raises(ValueError, match="null values for all nodes"):
            self.pair_coalescence_rates(ts, node_bin_map=node_bin_map)
        with pytest.raises(ValueError, match="a value per node"):
            self.pair_coalescence_rates(ts, node_bin_map=node_bin_map[:-1])
        with pytest.raises(TypeError, match="cast array data"):
            self.pair_coalescence_rates(ts, node_bin_map=np.zeros(num_nodes))

    def test_cpy_bad_time_windows(self):
        ts = self.example_ts()
        time_windows = np.zeros(1)
        node_bin_map = np.zeros(ts.nodes_time.size, dtype=np.int32)
        with pytest.raises(ValueError, match="at least two breakpoints"):
            self.pair_coalescence_rates(
                ts, time_windows=time_windows, node_bin_map=node_bin_map
            )
        time_windows = np.zeros((3, 3))
        with pytest.raises(ValueError, match="object too deep"):
            self.pair_coalescence_rates(
                ts, time_windows=time_windows, node_bin_map=node_bin_map
            )

    def test_cpy_bad_inputs(self):
        ts = self.example_ts()
        with pytest.raises(TypeError, match="at most 6 keyword"):
            ts.pair_coalescence_rates(
                sample_sets=None,
                sample_set_sizes=None,
                windows=None,
                quantiles=None,
                indexes=None,
                node_bin_map=None,
                foo="bar",
            )


--- ../../tskit/python/tests/test_parsimony.py ---


"""
Tests for the tree parsimony methods.
"""
import dataclasses
import io
import itertools

import Bio.Phylo.TreeConstruction
import msprime
import numpy as np
import pytest

import tests.tsutil as tsutil
import tskit


INF = np.inf


def bp_sankoff_score(tree, genotypes, cost_matrix):
    """
    Returns the sankoff score matrix computed by BioPython.
    """
    ts = tree.tree_sequence
    bp_tree = Bio.Phylo.read(io.StringIO(tree.as_newick()), "newick")
    records = [
        Bio.SeqRecord.SeqRecord(Bio.Seq.Seq(str(genotypes[j])), id=f"n{u}")
        for j, u in enumerate(ts.samples())
    ]
    alignment = Bio.Align.MultipleSeqAlignment(records)
    lower_triangular = []
    for j in range(cost_matrix.shape[0]):
        lower_triangular.append(list(cost_matrix[j, : j + 1]))
    bp_matrix = Bio.Phylo.TreeConstruction._Matrix(
        list(map(str, range(cost_matrix.shape[0]))), lower_triangular
    )
    ps = Bio.Phylo.TreeConstruction.ParsimonyScorer(bp_matrix)
    return ps.get_score(bp_tree, alignment)


def bp_fitch_score(tree, genotypes):
    """
    Returns the Fitch parsimony score computed by BioPython.
    """
    ts = tree.tree_sequence
    bp_tree = Bio.Phylo.read(io.StringIO(tree.as_newick()), "newick")
    records = [
        Bio.SeqRecord.SeqRecord(Bio.Seq.Seq(str(genotypes[j])), id=f"n{u}")
        for j, u in enumerate(ts.samples())
    ]
    alignment = Bio.Align.MultipleSeqAlignment(records)
    ps = Bio.Phylo.TreeConstruction.ParsimonyScorer()
    return ps.get_score(bp_tree, alignment)


def sankoff_score(tree, genotypes, cost_matrix):
    """
    Returns a num_nodes * num_alleles numpy array giving the minimum cost
    scores for the specified genotypes on the specified tree. If a cost
    matrix is provided, it must be a num_alleles * num_alleles array giving
    the cost of transitioning from each allele to every other allele.
    """
    num_alleles = cost_matrix.shape[0]
    S = np.zeros((tree.tree_sequence.num_nodes, num_alleles))
    for allele, u in zip(genotypes, tree.tree_sequence.samples()):
        S[u, :] = INF
        S[u, allele] = 0
    for parent in tree.nodes(order="postorder"):
        for child in tree.children(parent):
            for j in range(num_alleles):
                S[parent, j] += np.min(cost_matrix[:, j] + S[child])
    return S


def fitch_score(tree, genotypes):
    """
    Returns the Fitch parsimony score for the specified set of genotypes.
    """
    # Use the simplest set operation encoding of the set operations.
    A = {}
    for allele, u in zip(genotypes, tree.tree_sequence.samples()):
        A[u] = {allele}
    score = 0
    for u in tree.nodes(order="postorder"):
        if tree.is_internal(u):
            A[u] = set.intersection(*[A[v] for v in tree.children(u)])
            if len(A[u]) == 0:
                A[u] = set.union(*[A[v] for v in tree.children(u)])
                score += 1
    return score


def fitch_map_mutations(tree, genotypes, alleles):
    """
    Returns the Fitch parsimony reconstruction for the specified set of genotypes.
    The reconstruction is specified by returning the ancestral state and a
    list of mutations on the tree. Each mutation is a (node, parent, state)
    triple, where node is the node over which the transition occurs, the
    parent is the index of the parent transition above it on the tree (or -1
    if there is none) and state is the new state.
    """
    genotypes = np.array(genotypes)
    # Encode the set operations using a numpy array.
    not_missing = genotypes != -1
    if np.sum(not_missing) == 0:
        raise ValueError("Must have at least one non-missing genotype")
    num_alleles = np.max(genotypes[not_missing]) + 1
    A = np.zeros((tree.tree_sequence.num_nodes, num_alleles), dtype=np.int8)
    for allele, u in zip(genotypes, tree.tree_sequence.samples()):
        if allele != -1:
            A[u, allele] = 1
        else:
            A[u] = 1
    for u in tree.nodes(order="postorder"):
        if tree.num_children(u) > 2:
            raise ValueError("Fitch parsimony is for binary trees only")
        if not tree.is_sample(u):
            A[u] = 1
            for v in tree.children(u):
                A[u] = np.logical_and(A[u], A[v])
            if np.sum(A[u]) == 0:
                for v in tree.children(u):
                    A[u] = np.logical_or(A[u], A[v])

    root_states = np.zeros_like(A[0])
    for root in tree.roots:
        root_states = np.logical_or(root_states, A[root])
    ancestral_state = np.where(root_states == 1)[0][0]

    mutations = []
    state = {}
    for root in tree.roots:
        state[root] = ancestral_state
        parent = tskit.NULL
        if A[root, ancestral_state] != 1:
            state[root] = np.where(A[root] == 1)[0][0]
            mutations.append(
                tskit.Mutation(
                    node=root, parent=tskit.NULL, derived_state=alleles[state[root]]
                )
            )
            parent = len(mutations) - 1
        stack = [(root, parent)]
        while len(stack) > 0:
            u, parent_mutation = stack.pop()
            for v in tree.children(u):
                state[v] = state[u]
                if A[v, state[u]] != 1:
                    state[v] = np.where(A[v] == 1)[0][0]
                    mutations.append(
                        tskit.Mutation(
                            node=v,
                            parent=parent_mutation,
                            derived_state=alleles[state[v]],
                        )
                    )
                    stack.append((v, len(mutations) - 1))
                else:
                    stack.append((v, parent_mutation))
    return alleles[ancestral_state], mutations


def hartigan_map_mutations(tree, genotypes, alleles, ancestral_state=None):
    """
    Returns a Hartigan parsimony reconstruction for the specified set of genotypes.
    The reconstruction is specified by returning the ancestral state and a
    list of mutations on the tree. Each mutation is a (node, parent, state)
    triple, where node is the node over which the transition occurs, the
    parent is the index of the parent transition above it on the tree (or -1
    if there is none) and state is the new state.
    """
    # The python version of map_mutations allows the ancestral_state to be a string
    # from the alleles list, so we implement this at the top of this function although
    # it doesn't need to be in the C equivalent of this function
    if isinstance(ancestral_state, str):
        ancestral_state = alleles.index(ancestral_state)

    # equivalent C implementation can start here
    genotypes = np.array(genotypes)
    not_missing = genotypes != -1
    if np.sum(not_missing) == 0:
        raise ValueError("Must have at least one non-missing genotype")
    num_alleles = np.max(genotypes[not_missing]) + 1
    if ancestral_state is not None:
        if ancestral_state < 0 or ancestral_state >= len(alleles):
            raise ValueError("ancestral_state must be a number from 0..(num_alleles-1)")
        if ancestral_state >= num_alleles:
            num_alleles = ancestral_state + 1
    num_nodes = tree.tree_sequence.num_nodes

    # use a numpy array of 0/1 values to represent the set of states
    # to make the code as similar as possible to the C implementation.
    optimal_set = np.zeros((num_nodes + 1, num_alleles), dtype=np.int8)
    for allele, u in zip(genotypes, tree.tree_sequence.samples()):
        if allele != -1:
            optimal_set[u, allele] = 1
        else:
            optimal_set[u] = 1

    allele_count = np.zeros(num_alleles, dtype=int)
    for u in tree.nodes(tree.virtual_root, order="postorder"):
        allele_count[:] = 0
        for v in tree.children(u):
            for j in range(num_alleles):
                allele_count[j] += optimal_set[v, j]
        if not tree.is_sample(u):
            max_allele_count = np.max(allele_count)
            optimal_set[u, allele_count == max_allele_count] = 1

    if ancestral_state is None:
        ancestral_state = np.argmax(optimal_set[tree.virtual_root])
    else:
        optimal_set[tree.virtual_root] = 1

    @dataclasses.dataclass
    class StackElement:
        node: int
        state: int
        mutation_parent: int

    mutations = []
    stack = [StackElement(tree.virtual_root, ancestral_state, -1)]
    while len(stack) > 0:
        s = stack.pop()
        if optimal_set[s.node, s.state] == 0:
            s.state = np.argmax(optimal_set[s.node])
            mutation = tskit.Mutation(
                node=s.node,
                derived_state=alleles[s.state],
                parent=s.mutation_parent,
            )
            s.mutation_parent = len(mutations)
            mutations.append(mutation)
        for v in tree.children(s.node):
            stack.append(StackElement(v, s.state, s.mutation_parent))
    return alleles[ancestral_state], mutations


def reconstruct_states(tree, genotypes, S, cost_matrix):
    """
    Given the specified observations for the samples and tree score
    matrix computed by sankoff_score and the transition cost matrix,
    return the ancestral_state and state transitions on the tree.
    """
    root_cost = np.zeros_like(S[0])
    for root in tree.roots:
        for j in range(S.shape[1]):
            root_cost[j] += np.min(cost_matrix[:, j] + S[root])
    ancestral_state = np.argmin(root_cost)

    transitions = {}
    A = {}
    for root in tree.roots:
        A[root] = ancestral_state
        for u in tree.nodes(order="preorder"):
            for v in tree.children(u):
                cost = cost_matrix[A[u]] + S[v]
                A[v] = np.argmin(cost)
                if A[u] != A[v]:
                    transitions[v] = A[v]

    return ancestral_state, transitions


def sankoff_map_mutations(tree, genotypes, cost_matrix=None):
    """
    Returns the recontructed minimal state transitions for the specified set of
    genotypes on the specified (optional) cost matrix.

    NOTE: we don't consider complications of multiple roots and internal samples
    here.

    TODO: update this to take the alleles as input like the other methods.
    """
    if cost_matrix is None:
        num_alleles = np.max(genotypes) + 1
        cost_matrix = np.ones((num_alleles, num_alleles))
        np.fill_diagonal(cost_matrix, 0)
    S = sankoff_score(tree, genotypes, cost_matrix)
    return reconstruct_states(tree, genotypes, S, cost_matrix)


def felsenstein_tables():
    """
    Return tables for the example tree.
    """
    #
    #     8
    #   ┏━┻━━┓
    #   ┃    7
    #   ┃   ┏┻┓
    #   6   ┃ ┃
    # ┏━┻┓  ┃ ┃
    # ┃  5  ┃ ┃
    # ┃ ┏┻┓ ┃ ┃
    # 2 3 4 0 1
    #
    tables = tskit.TableCollection(1)
    for _ in range(5):
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
    for j in range(4):
        tables.nodes.add_row(flags=0, time=j + 1)
    tables.edges.add_row(0, 1, 7, 0)
    tables.edges.add_row(0, 1, 7, 1)
    tables.edges.add_row(0, 1, 6, 2)
    tables.edges.add_row(0, 1, 5, 3)
    tables.edges.add_row(0, 1, 5, 4)
    tables.edges.add_row(0, 1, 6, 5)
    tables.edges.add_row(0, 1, 8, 6)
    tables.edges.add_row(0, 1, 8, 7)
    tables.sort()
    return tables


def felsenstein_example():
    """
    Returns the tree used in Felsenstein's book, pg.15.
    """
    ts = felsenstein_tables().tree_sequence()
    return ts.first()


class TestSankoff:
    """
    Tests for the Sankoff algorithm.
    """

    def test_felsenstein_example_score(self):
        tree = felsenstein_example()
        genotypes = [1, 0, 1, 0, 2]
        cost_matrix = np.array(
            [[0, 2.5, 1, 2.5], [2.5, 0, 2.5, 1], [1, 2.5, 0, 2.5], [2.5, 1, 2.5, 0]]
        )
        S = sankoff_score(tree, genotypes, cost_matrix)
        S2 = [
            [INF, 0, INF, INF],
            [0, INF, INF, INF],
            [INF, 0, INF, INF],
            [0, INF, INF, INF],
            [INF, INF, 0, INF],
            [1, 5, 1, 5],
            [3.5, 3.5, 3.5, 4.5],
            [2.5, 2.5, 3.5, 3.5],
            [6, 6, 7, 8],
        ]
        assert np.array_equal(S, np.array(S2))

    def test_felsenstein_example_reconstruct(self):
        tree = felsenstein_example()
        genotypes = [1, 0, 1, 0, 2]
        cost_matrix = np.array(
            [[0, 2.5, 1, 2.5], [2.5, 0, 2.5, 1], [1, 2.5, 0, 2.5], [2.5, 1, 2.5, 0]]
        )
        S = sankoff_score(tree, genotypes, cost_matrix)
        ancestral_state, transitions = reconstruct_states(
            tree, genotypes, S, cost_matrix
        )
        assert {2: 1, 4: 2, 0: 1} == transitions
        assert 0 == ancestral_state

    def verify_infinite_sites(self, ts):
        assert ts.num_trees == 1
        assert ts.num_sites > 5
        tree = ts.first()
        for variant in ts.variants():
            ancestral_state, transitions = sankoff_map_mutations(
                tree, variant.genotypes
            )
            assert len(transitions) == 1
            assert ancestral_state == 0
            assert transitions[variant.site.mutations[0].node] == 1

    def test_infinite_sites_binary_n2(self):
        ts = msprime.simulate(2, mutation_rate=10, random_seed=1)
        self.verify_infinite_sites(ts)

    def test_infinite_sites_binary_n50(self):
        ts = msprime.simulate(50, mutation_rate=2, random_seed=1)
        self.verify_infinite_sites(ts)

    def test_infinite_sites_acgt_n2(self):
        ts = msprime.simulate(2, random_seed=1)
        ts = msprime.mutate(
            ts, rate=3, model=msprime.InfiniteSites(msprime.NUCLEOTIDES), random_seed=1
        )
        self.verify_infinite_sites(ts)

    def test_infinite_sites_acgt_n15(self):
        ts = msprime.simulate(2, random_seed=1)
        ts = msprime.mutate(
            ts, rate=3, model=msprime.InfiniteSites(msprime.NUCLEOTIDES), random_seed=1
        )
        self.verify_infinite_sites(ts)

    def verify_jukes_cantor(self, ts, cost_matrix):
        assert ts.num_trees == 1
        assert ts.num_mutations > ts.num_sites
        tree = ts.first()
        for variant in ts.variants():
            single_score = bp_sankoff_score(tree, variant.genotypes, cost_matrix)
            score_matrix = sankoff_score(tree, variant.genotypes, cost_matrix)
            score = np.min(score_matrix[tree.root])
            assert single_score == score

    def test_jukes_cantor_n2_simple_matrix(self):
        cost_matrix = np.ones((4, 4))
        np.fill_diagonal(cost_matrix, 0)
        ts = msprime.simulate(2, random_seed=1)
        ts = tsutil.jukes_cantor(ts, 5, 2, seed=1)
        self.verify_jukes_cantor(ts, cost_matrix)

    def test_jukes_cantor_n20_simple_matrix(self):
        cost_matrix = np.ones((4, 4))
        np.fill_diagonal(cost_matrix, 0)
        ts = msprime.simulate(20, random_seed=1)
        ts = tsutil.jukes_cantor(ts, 5, 2, seed=1)
        self.verify_jukes_cantor(ts, cost_matrix)

    def test_jukes_cantor_n2_felsenstein_matrix(self):
        cost_matrix = np.array(
            [[0, 2.5, 1, 2.5], [2.5, 0, 2.5, 1], [1, 2.5, 0, 2.5], [2.5, 1, 2.5, 0]]
        )
        ts = msprime.simulate(2, random_seed=1)
        ts = tsutil.jukes_cantor(ts, 5, 2, seed=1)
        self.verify_jukes_cantor(ts, cost_matrix)

    def test_jukes_cantor_n20_felsenstein_matrix(self):
        cost_matrix = np.array(
            [[0, 2.5, 1, 2.5], [2.5, 0, 2.5, 1], [1, 2.5, 0, 2.5], [2.5, 1, 2.5, 0]]
        )
        ts = msprime.simulate(20, random_seed=1)
        ts = tsutil.jukes_cantor(ts, 5, 2, seed=1)
        self.verify_jukes_cantor(ts, cost_matrix)


class TestFitchParsimonyDistance:
    """
    Tests for the Fitch parsimony algorithm.
    """

    def verify(self, ts):
        assert ts.num_trees == 1
        assert ts.num_sites > 3
        tree = ts.first()
        for variant in ts.variants(isolated_as_missing=False):
            score = fitch_score(tree, variant.genotypes)
            bp_score = bp_fitch_score(tree, variant.genotypes)
            assert bp_score == score
            ancestral_state1, transitions1 = fitch_map_mutations(
                tree, variant.genotypes, variant.alleles
            )
            ancestral_state2, transitions2 = tree.map_mutations(
                variant.genotypes, variant.alleles
            )
            assert ancestral_state1 == ancestral_state2
            assert len(transitions1) == len(transitions2)
            # The Sankoff algorithm doesn't recontruct the state in the same way.
            # Just a limitation of the implementation.
            ancestral_state3, transitions3 = sankoff_map_mutations(
                tree, variant.genotypes
            )
            assert ancestral_state1 == variant.alleles[ancestral_state3]
            # The algorithms will make slightly different choices on where to put
            # the transitions, but they are equally parsimonious.
            assert len(transitions1) == len(transitions3)

    def test_infinite_sites_binary_n2(self):
        ts = msprime.simulate(2, mutation_rate=10, random_seed=1)
        self.verify(ts)

    def test_infinite_sites_binary_n50(self):
        ts = msprime.simulate(50, mutation_rate=2, random_seed=1)
        self.verify(ts)

    def test_infinite_sites_acgt_n2(self):
        ts = msprime.simulate(2, random_seed=1)
        ts = msprime.mutate(
            ts, rate=3, model=msprime.InfiniteSites(msprime.NUCLEOTIDES), random_seed=1
        )
        self.verify(ts)

    def test_jukes_cantor_n2(self):
        ts = msprime.simulate(2, random_seed=1)
        ts = tsutil.jukes_cantor(ts, 5, 2, seed=1)
        self.verify(ts)

    def test_jukes_cantor_n5(self):
        ts = msprime.simulate(5, random_seed=2)
        ts = tsutil.jukes_cantor(ts, 5, 1, seed=0)
        self.verify(ts)

    def test_jukes_cantor_n20(self):
        ts = msprime.simulate(20, random_seed=2)
        ts = tsutil.jukes_cantor(ts, 5, 2, seed=11)
        self.verify(ts)

    def test_jukes_cantor_n50(self):
        ts = msprime.simulate(50, random_seed=2)
        ts = tsutil.jukes_cantor(ts, 5, 2, seed=1)
        self.verify(ts)


class TestParsimonyBase:
    """
    Base class for tests of the map_mutations parsimony method.
    """

    def do_map_mutations(
        self, tree, genotypes, alleles=None, ancestral_state=None, compare_lib=True
    ):
        if alleles is None:
            alleles = [str(j) for j in range(max(genotypes) + 1)]
        ancestral_state1, transitions1 = tree.map_mutations(
            genotypes, alleles, ancestral_state
        )
        if compare_lib:
            ancestral_state2, transitions2 = hartigan_map_mutations(
                tree, genotypes, alleles, ancestral_state
            )
            assert ancestral_state1 == ancestral_state2
            assert len(transitions1) == len(transitions2)
            sorted_t1 = sorted((m.node, m.derived_state) for m in transitions1)
            sorted_t2 = sorted((m.node, m.derived_state) for m in transitions2)
            assert sorted_t1 == sorted_t2
            assert transitions1 == transitions2
        return ancestral_state1, transitions1


class TestParsimonyBadAlleles(TestParsimonyBase):
    tree = tskit.Tree.generate_comb(3)

    def test_too_many_alleles(self):
        genotypes = [0, 0, 64]
        alleles = [str(j) for j in range(max(genotypes) + 1)]
        with pytest.raises(ValueError, match="maximum of 64"):
            # Only a limitation in the C version of map_mutations
            self.tree.map_mutations(genotypes, alleles)

    def test_ancestral_state_too_big(self):
        genotypes = [0, 0, 1]
        alleles = [str(x) for x in range(2**8)]  # exceeds HARTIGAN_MAX_ALLELES
        with pytest.raises(ValueError, match="maximum of 64"):
            # Only a limitation in the C version of map_mutations
            self.tree.map_mutations(
                genotypes, alleles=alleles, ancestral_state=alleles[-1]
            )


class TestParsimonyRoundTrip(TestParsimonyBase):
    """
    Tests that we can reproduce the genotypes for set of tree sequences by
    inferring the locations of mutations.
    """

    def verify(self, ts):
        G = ts.genotype_matrix(isolated_as_missing=False)
        alleles = [v.alleles for v in ts.variants()]
        for randomize_ancestral_states in [False, True]:
            tables = ts.dump_tables()
            tables.sites.clear()
            tables.mutations.clear()
            fixed_anc_state = None
            for tree in ts.trees():
                for site in tree.sites():
                    if randomize_ancestral_states:
                        num_alleles = len(alleles[site.id])
                        if alleles[site.id][-1] is None:
                            num_alleles -= 1
                        fixed_anc_state = np.random.randint(num_alleles)
                    ancestral_state, mutations = self.do_map_mutations(
                        tree,
                        G[site.id],
                        alleles[site.id],
                        ancestral_state=fixed_anc_state,
                    )
                    site_id = tables.sites.append(
                        site.replace(ancestral_state=ancestral_state)
                    )
                    parent_offset = len(tables.mutations)
                    for mutation in mutations:
                        parent = mutation.parent
                        if parent != tskit.NULL:
                            parent += parent_offset
                        tables.mutations.append(
                            mutation.replace(site=site_id, parent=parent)
                        )
            other_ts = tables.tree_sequence()
            for h1, h2 in zip(
                ts.haplotypes(isolated_as_missing=False),
                other_ts.haplotypes(isolated_as_missing=False),
            ):
                assert h1 == h2

            # Make sure we're computing the parent correctly.
            tables2 = tables.copy()
            nulled = np.zeros_like(tables.mutations.parent) - 1
            tables2.mutations.parent = nulled
            assert np.array_equal(tables.mutations.parent, tables.mutations.parent)

    def test_infinite_sites_n3(self):
        ts = msprime.simulate(3, mutation_rate=3, random_seed=3)
        self.verify(ts)

    def test_infinite_sites_n20(self):
        ts = msprime.simulate(20, mutation_rate=3, random_seed=3)
        self.verify(ts)

    def test_infinite_sites_n20_recombination(self):
        ts = msprime.simulate(20, mutation_rate=3, recombination_rate=2, random_seed=3)
        assert ts.num_trees > 2
        self.verify(ts)

    def test_infinite_sites_n5_internal_samples(self):
        ts = msprime.simulate(5, mutation_rate=3, random_seed=3)
        self.verify(tsutil.jiggle_samples(ts))

    def test_infinite_sites_n20_internal_samples(self):
        ts = msprime.simulate(20, mutation_rate=3, random_seed=3)
        self.verify(tsutil.jiggle_samples(ts))

    def test_jukes_cantor_n5(self):
        ts = msprime.simulate(5, random_seed=1)
        ts = tsutil.jukes_cantor(ts, 5, 1, seed=1)
        self.verify(ts)

    def test_jukes_cantor_n20(self):
        ts = msprime.simulate(20, random_seed=1)
        ts = tsutil.jukes_cantor(ts, 5, 2, seed=1)
        self.verify(ts)

    def test_jukes_cantor_n50(self):
        ts = msprime.simulate(50, random_seed=1)
        ts = tsutil.jukes_cantor(ts, 5, 2, seed=2)
        self.verify(ts)

    def test_jukes_cantor_n5_internal_samples(self):
        ts = msprime.simulate(5, random_seed=1)
        ts = tsutil.jukes_cantor(ts, 1, 1, seed=1)
        ts = tsutil.jiggle_samples(ts)
        self.verify(ts)

    def test_jukes_cantor_n20_internal_samples(self):
        ts = msprime.simulate(20, random_seed=1)
        ts = tsutil.jukes_cantor(ts, 5, 2, seed=1)
        self.verify(tsutil.jiggle_samples(ts))

    def test_jukes_cantor_n50_internal_samples(self):
        ts = msprime.simulate(50, random_seed=1)
        ts = tsutil.jukes_cantor(ts, 5, 2, seed=2)
        self.verify(tsutil.jiggle_samples(ts))

    def test_jukes_cantor_balanced_ternary_internal_samples(self):
        tree = tskit.Tree.generate_balanced(27, arity=3)
        ts = tsutil.jukes_cantor(tree.tree_sequence, 5, 2, seed=1)
        assert ts.num_sites > 1
        self.verify(tsutil.jiggle_samples(ts))

    def test_infinite_sites_n20_multiroot(self):
        ts = msprime.simulate(20, mutation_rate=3, random_seed=3)
        self.verify(ts.decapitate(np.max(ts.tables.nodes.time) / 2))

    def test_jukes_cantor_n15_multiroot(self):
        ts = msprime.simulate(15, random_seed=1)
        ts = ts.decapitate(np.max(ts.tables.nodes.time) / 5)
        ts = tsutil.jukes_cantor(ts, 15, 2, seed=3)
        self.verify(ts)

    def test_jukes_cantor_balanced_ternary_multiroot(self):
        ts = tskit.Tree.generate_balanced(50, arity=3).tree_sequence
        ts = ts.decapitate(np.max(ts.tables.nodes.time) / 3)
        ts = tsutil.jukes_cantor(ts, 15, 2, seed=3)
        self.verify(ts)
        assert ts.num_sites > 1
        self.verify(tsutil.jiggle_samples(ts))

    def test_jukes_cantor_n50_multiroot(self):
        ts = msprime.simulate(50, random_seed=1)
        ts = ts.decapitate(np.max(ts.tables.nodes.time) / 2)
        ts = tsutil.jukes_cantor(ts, 5, 2, seed=2)
        self.verify(ts)

    def test_jukes_cantor_root_polytomy_n5(self):
        tree = tskit.Tree.unrank(5, (1, 0))
        ts = tsutil.jukes_cantor(tree.tree_sequence, 5, 2, seed=1)
        assert ts.num_sites > 2
        self.verify(ts)

    def test_jukes_cantor_leaf_polytomy_n5(self):
        tree = tskit.Tree.unrank(5, (7, 0))
        ts = tsutil.jukes_cantor(tree.tree_sequence, 5, 2, seed=1)
        assert ts.num_sites > 2
        self.verify(ts)

    @pytest.mark.parametrize(
        "tree_builder", [tskit.Tree.generate_balanced, tskit.Tree.generate_comb]
    )
    @pytest.mark.parametrize("n", [2, 5, 10])
    def test_many_states_binary(self, tree_builder, n):
        tree = tree_builder(n)
        tables = tree.tree_sequence.dump_tables()
        tables.sites.add_row(0.5, "0")
        for j in range(1, n):
            tables.mutations.add_row(0, derived_state=str(j), node=j)
        ts = tables.tree_sequence()
        assert np.array_equal(ts.genotype_matrix(), [np.arange(n, dtype=np.int8)])
        self.verify(tables.tree_sequence())

    @pytest.mark.parametrize("arity", [2, 3, 4])
    @pytest.mark.parametrize("n", [2, 5, 10])
    def test_many_states_arity(self, n, arity):
        tree = tskit.Tree.generate_balanced(n, arity=arity)
        tables = tree.tree_sequence.dump_tables()
        tables.sites.add_row(0.5, "0")
        for j in range(1, n):
            tables.mutations.add_row(0, derived_state=str(j), node=j)
        ts = tables.tree_sequence()
        assert np.array_equal(ts.genotype_matrix(), [np.arange(n, dtype=np.int8)])
        self.verify(tables.tree_sequence())


class TestParsimonyRoundTripMissingData(TestParsimonyRoundTrip):
    """
    Tests that we can reproduce the genotypes for set of tree sequences by
    inferring the locations of mutations.
    """

    def verify(self, ts):
        tables = ts.dump_tables()
        tables.sites.clear()
        tables.mutations.clear()
        G = ts.genotype_matrix(isolated_as_missing=False)
        # Set the first sample to missing data everywhere
        G[:, 0] = -1
        alleles = [v.alleles for v in ts.variants()]
        for tree in ts.trees():
            for site in tree.sites():
                ancestral_state, mutations = self.do_map_mutations(
                    tree, G[site.id], alleles[site.id]
                )
                site_id = tables.sites.append(
                    site.replace(ancestral_state=ancestral_state)
                )
                parent_offset = len(tables.mutations)
                for m in mutations:
                    parent = m.parent
                    if m.parent != tskit.NULL:
                        parent = m.parent + parent_offset
                    tables.mutations.append(m.replace(site=site_id, parent=parent))
        other_ts = tables.tree_sequence()
        assert ts.num_samples == other_ts.num_samples
        H1 = list(ts.haplotypes(isolated_as_missing=False))
        H2 = list(other_ts.haplotypes(isolated_as_missing=False))
        # All samples except 0 should be reproduced exactly.
        assert H1[1:] == H2[1:]


class TestParsimonyMissingData(TestParsimonyBase):
    """
    Tests that we correctly map_mutations when we have missing data.
    """

    @pytest.mark.parametrize("n", range(2, 10))
    def test_all_missing(self, n):
        ts = msprime.simulate(n, random_seed=2)
        tree = ts.first()
        genotypes = np.zeros(n, dtype=np.int8) - 1
        alleles = ["0", "1"]
        with pytest.raises(ValueError):
            fitch_map_mutations(tree, genotypes, alleles)
        with pytest.raises(ValueError):
            hartigan_map_mutations(tree, genotypes, alleles)
        with pytest.raises(tskit.LibraryError):
            tree.map_mutations(genotypes, alleles)

    @pytest.mark.parametrize("n", range(2, 10))
    def test_one_non_missing(self, n):
        ts = msprime.simulate(n, random_seed=2)
        tree = ts.first()
        for j in range(n):
            genotypes = np.zeros(n, dtype=np.int8) - 1
            genotypes[j] = 0
            ancestral_state, transitions = self.do_map_mutations(
                tree, genotypes, ["0", "1"]
            )
            assert ancestral_state == "0"
            assert len(transitions) == 0

    @pytest.mark.parametrize("arity", range(2, 10))
    def test_one_non_missing_balanced(self, arity):
        n = 40
        tree = tskit.Tree.generate_balanced(n, arity=arity)
        for j in range(n):
            genotypes = np.zeros(n, dtype=np.int8) - 1
            genotypes[j] = 0
            ancestral_state, transitions = self.do_map_mutations(
                tree, genotypes, ["0", "1"]
            )
            assert ancestral_state == "0"
            assert len(transitions) == 0

    @pytest.mark.parametrize("n", range(2, 10))
    def test_many_states_half_missing(self, n):
        ts = msprime.simulate(n, random_seed=2)
        tree = ts.first()
        genotypes = np.zeros(n, dtype=np.int8) - 1
        genotypes[0 : n // 2] = np.arange(n // 2, dtype=int)
        alleles = [str(j) for j in range(n)]
        ancestral_state, transitions = self.do_map_mutations(tree, genotypes, alleles)
        assert ancestral_state == "0"
        assert len(transitions) == max(0, n // 2 - 1)

    @pytest.mark.parametrize("n", range(2, 10))
    def test_one_missing(self, n):
        ts = msprime.simulate(n, random_seed=2)
        tree = ts.first()
        alleles = [str(j) for j in range(2)]
        for j in range(n):
            genotypes = np.zeros(n, dtype=np.int8) - 1
            genotypes[j] = 0
            ancestral_state, transitions = self.do_map_mutations(
                tree, genotypes, alleles
            )
            assert ancestral_state == "0"
            assert len(transitions) == 0

    @pytest.mark.parametrize("arity", range(2, 10))
    def test_one_missing_balanced(self, arity):
        n = 40
        tree = tskit.Tree.generate_balanced(n, arity=arity)
        alleles = [str(j) for j in range(2)]
        for j in range(n):
            genotypes = np.zeros(n, dtype=np.int8) - 1
            genotypes[j] = 0
            ancestral_state, transitions = self.do_map_mutations(
                tree, genotypes, alleles
            )
            assert ancestral_state == "0"
            assert len(transitions) == 0

    def test_one_missing_derived_state(self):
        tables = felsenstein_tables()
        ts = tables.tree_sequence()
        genotypes = np.zeros(5, dtype=np.int8)
        genotypes[0] = -1
        genotypes[1] = 1
        alleles = [str(j) for j in range(2)]
        ancestral_state, transitions = self.do_map_mutations(
            ts.first(), genotypes, alleles
        )
        assert ancestral_state == "0"
        assert len(transitions) == 1
        assert transitions[0].node == 7
        assert transitions[0].parent == -1
        assert transitions[0].derived_state == "1"


class TestParsimonyExamples(TestParsimonyBase):
    """
    Some examples on a given tree.
    """

    #
    #          8
    #         / \
    #        /   \
    #       /     \
    #      7       \
    #     / \       6
    #    /   5     / \
    #   /   / \   /   \
    #  4   0   1 2     3
    small_tree_ex_nodes = """\
    id      is_sample   population      time
    0       1       0               0.00000000000000
    1       1       0               0.00000000000000
    2       1       0               0.00000000000000
    3       1       0               0.00000000000000
    4       1       0               0.00000000000000
    5       0       0               0.14567111023387
    6       0       0               0.21385545626353
    7       0       0               0.43508024345063
    8       0       0               1.60156352971203
    """
    small_tree_ex_edges = """\
    id      left            right           parent  child
    0       0.00000000      1.00000000      5       0,1
    1       0.00000000      1.00000000      6       2,3
    2       0.00000000      1.00000000      7       4,5
    3       0.00000000      1.00000000      8       6,7
    """
    tree = tskit.load_text(
        nodes=io.StringIO(small_tree_ex_nodes),
        edges=io.StringIO(small_tree_ex_edges),
        strict=False,
    ).first()

    def test_mutation_over_0(self):
        genotypes = [1, 0, 0, 0, 0]
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 1
        assert transitions[0] == tskit.Mutation(node=0, parent=-1, derived_state="1")

    def test_mutation_over_5(self):
        genotypes = [1, 1, 0, 0, 0]
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 1
        assert transitions[0] == tskit.Mutation(node=5, parent=-1, derived_state="1")

    def test_mutation_over_7(self):
        genotypes = [1, 1, 0, 0, 1]
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 1
        assert transitions[0] == tskit.Mutation(node=7, parent=-1, derived_state="1")

    def test_mutation_over_7_0(self):
        genotypes = [2, 1, 0, 0, 1]
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 2
        assert transitions[0] == tskit.Mutation(node=7, parent=-1, derived_state="1")
        assert transitions[1] == tskit.Mutation(node=0, parent=0, derived_state="2")

    def test_mutation_over_7_0_alleles(self):
        genotypes = [2, 1, 0, 0, 1]
        alleles = ["ANC", "ONE", "TWO"]
        ancestral_state, transitions = self.do_map_mutations(
            self.tree, genotypes, alleles
        )
        assert ancestral_state == "ANC"
        assert len(transitions) == 2
        assert transitions[0] == tskit.Mutation(node=7, parent=-1, derived_state="ONE")
        assert transitions[1] == tskit.Mutation(node=0, parent=0, derived_state="TWO")

    def test_mutation_over_7_missing_data_0(self):
        genotypes = [-1, 1, 0, 0, 1]
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 1
        assert transitions[0] == tskit.Mutation(node=7, parent=-1, derived_state="1")

    def test_mutation_over_leaf_sibling_missing(self):
        genotypes = [0, 0, 1, -1, 0]
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 1
        # We assume that the mutation is over the parent of 2 and the missing data
        # so we impute that 3 also has allele 1. This suprising behaviour to me:
        # I would have thought it was more parsimonious to assume that the missing
        # data had the ancestral state. However, the number of *state changes*
        # is the same, which is what the algorithm is minimising.
        assert transitions[0] == tskit.Mutation(node=6, parent=-1, derived_state="1")

        # Reverse is the same
        genotypes = [0, 0, -1, 1, 0]
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 1
        assert transitions[0] == tskit.Mutation(node=6, parent=-1, derived_state="1")

    def test_mutation_over_6_missing_data_0(self):
        genotypes = [-1, 0, 1, 1, 0]
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 1
        assert transitions[0] == tskit.Mutation(node=6, parent=-1, derived_state="1")

    def test_mutation_over_0_missing_data_4(self):
        genotypes = [1, 0, 0, 0, -1]
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 1
        assert transitions[0] == tskit.Mutation(node=0, parent=-1, derived_state="1")

    def test_multi_mutation_missing_data(self):
        genotypes = [1, 2, -1, 0, 0]
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 2
        assert transitions[0] == tskit.Mutation(node=5, parent=-1, derived_state="1")
        assert transitions[1] == tskit.Mutation(node=1, parent=0, derived_state="2")


class TestParsimonyExamplesPolytomy(TestParsimonyBase):
    """
    Some examples on a given non-binary tree.
    """

    #         9
    #       ┏━┻━━┓
    #       7    8
    #     ┏━┻━┓ ┏┻┓
    #     6   ┃ ┃ ┃
    #   ┏━╋━┓ ┃ ┃ ┃
    #   0 2 4 5 1 3

    nodes = io.StringIO(
        """\
    id      is_sample   time
    0       1           0
    1       1           0
    2       1           0
    3       1           0
    4       1           0
    5       1           0
    6       0           1
    7       0           2
    8       0           2
    9       0           3
    """
    )
    edges = io.StringIO(
        """\
    left    right   parent  child
    0       1       6       0,2,4
    0       1       7       6,5
    0       1       8       1,3
    0       1       9       7,8
    """
    )

    tree = tskit.load_text(
        nodes=nodes,
        edges=edges,
        strict=False,
    ).first()

    def test_all_zeros(self):
        genotypes = [0, 0, 0, 0, 0, 0]
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 0

    def test_mutation_over_8(self):
        genotypes = [0, 1, 0, 1, 0, 0]
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 1
        assert transitions[0] == tskit.Mutation(node=8, derived_state="1")

    def test_mutation_over_6(self):
        genotypes = [1, 0, 1, 0, 1, 0]
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 1
        assert transitions[0] == tskit.Mutation(node=6, derived_state="1")

    def test_mutation_over_0_5(self):
        # Bug reported in https://github.com/tskit-dev/tskit/issues/987
        genotypes = [1, 0, 0, 0, 0, 1]
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 2
        assert transitions[0] == tskit.Mutation(node=0, derived_state="1")
        assert transitions[1] == tskit.Mutation(node=5, derived_state="1")

    def test_mutation_over_7_back_mutation_4(self):
        genotypes = [1, 0, 1, 0, 0, 1]
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 2
        assert transitions[0] == tskit.Mutation(node=7, derived_state="1")
        assert transitions[1] == tskit.Mutation(node=4, derived_state="0", parent=0)


class TestParsimonyExamplesStar(TestParsimonyBase):
    """
    Some examples on star topologies.
    """

    @pytest.mark.parametrize("n", range(3, 8))
    def test_two_states_freq_n_minus_1(self, n):
        tree = tskit.Tree.generate_star(n)
        genotypes = np.zeros(n, dtype=np.int8)
        genotypes[0] = 1
        ancestral_state, transitions = self.do_map_mutations(tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 1
        assert transitions[0] == tskit.Mutation(node=0, derived_state="1")

        genotypes[:] = 1
        genotypes[0] = 0
        ancestral_state, transitions = self.do_map_mutations(tree, genotypes)
        assert ancestral_state == "1"
        assert len(transitions) == 1
        assert transitions[0] == tskit.Mutation(node=0, derived_state="0")

    @pytest.mark.parametrize("n", range(5, 10))
    def test_two_states_freq_n_minus_2(self, n):
        tree = tskit.Tree.generate_star(n)
        genotypes = np.zeros(n, dtype=np.int8)
        genotypes[0:2] = 1
        ancestral_state, transitions = self.do_map_mutations(tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 2
        assert transitions[0] == tskit.Mutation(node=1, derived_state="1")
        assert transitions[1] == tskit.Mutation(node=0, derived_state="1")

        genotypes[:] = 1
        genotypes[0:2] = 0
        ancestral_state, transitions = self.do_map_mutations(tree, genotypes)
        assert ancestral_state == "1"
        assert len(transitions) == 2
        assert transitions[0] == tskit.Mutation(node=1, derived_state="0")
        assert transitions[1] == tskit.Mutation(node=0, derived_state="0")

    @pytest.mark.parametrize("n", range(5, 10))
    def test_three_states_freq_n_minus_2(self, n):
        tree = tskit.Tree.generate_star(n)
        genotypes = np.zeros(n, dtype=np.int8)
        genotypes[0] = 1
        genotypes[1] = 2
        ancestral_state, transitions = self.do_map_mutations(tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 2
        assert transitions[0] == tskit.Mutation(node=1, derived_state="2")
        assert transitions[1] == tskit.Mutation(node=0, derived_state="1")

    @pytest.mark.parametrize("n", range(2, 10))
    def test_n_states(self, n):
        tree = tskit.Tree.generate_star(n)
        genotypes = np.arange(n, dtype=np.int8)
        ancestral_state, transitions = self.do_map_mutations(tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == n - 1

    @pytest.mark.parametrize("n", range(3, 10))
    def test_missing_data(self, n):
        tree = tskit.Tree.generate_star(n)
        genotypes = np.zeros(n, dtype=np.int8)
        genotypes[0] = tskit.MISSING_DATA
        genotypes[1] = 1
        ancestral_state, transitions = self.do_map_mutations(tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 1
        assert transitions[0] == tskit.Mutation(node=1, derived_state="1")


class TestParsimonyExamplesBalancedTernary(TestParsimonyBase):
    """
    Some examples on a given non-binary tree.
    """

    tree = tskit.Tree.generate_balanced(27, arity=3)
    #                                39
    #         ┏━━━━━━━━━━━━━━━━━━━━━┳━┻━━━━━━━━━━━━━━━━━━━━━━━━┓
    #        30                    34                         38
    #   ┏━━━━━╋━━━━━┓      ┏━━━━━━━━╋━━━━━━━━┓        ┏━━━━━━━━╋━━━━━━━━┓
    #  27    28    29     31       32       33       35       36       37
    # ┏━╋━┓ ┏━╋━┓ ┏━╋━┓ ┏━━╋━━┓  ┏━━╋━━┓  ┏━━╋━━┓  ┏━━╋━━┓  ┏━━╋━━┓  ┏━━╋━━┓
    # 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26

    def test_mutation_over_27_29(self):
        genotypes = np.zeros(27, dtype=int)
        genotypes[0:3] = 1
        genotypes[6:9] = 1
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 2
        # the algorithm chooses a back mutation instead
        assert transitions[0] == tskit.Mutation(node=30, derived_state="1")
        assert transitions[1] == tskit.Mutation(node=28, derived_state="0", parent=0)

    def test_three_clades(self):
        genotypes = np.zeros(27, dtype=int)
        genotypes[9:18] = 1
        genotypes[18:27] = 2
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 2
        assert transitions[0] == tskit.Mutation(node=38, derived_state="2")
        assert transitions[1] == tskit.Mutation(node=34, derived_state="1")

    def test_nonzero_ancestral_state(self):
        genotypes = np.ones(27, dtype=int)
        genotypes[0] = 0
        genotypes[26] = 0
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "1"
        assert len(transitions) == 2
        assert transitions[0] == tskit.Mutation(node=26, derived_state="0")
        assert transitions[1] == tskit.Mutation(node=0, derived_state="0")

    def test_many_states(self):
        genotypes = np.arange(27, dtype=int)
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 26

    def test_least_parsimonious(self):
        genotypes = [0, 1, 2] * 9
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 18


class TestParsimonyExamplesUnary(TestParsimonyBase):
    """
    Some examples on a tree with unary nodes. The mutation should be placed
    on the highest node along the lineage compatible with the parsimonious placement
    """

    #        9
    #      ┏━┻━┓
    #      8   ┃
    #    ┏━┻━┓ ┃
    #    6   7 ┃
    #    ┃   ┃ ┃
    #    5   ┃ ┃
    #  ┏━╋━┓ ┃ ┃
    #  0 2 3 1 4

    nodes = io.StringIO(
        """\
    id      is_sample   time
    0       1           0
    1       1           0
    2       1           0
    3       1           0
    4       1           0
    5       0           1
    6       0           2
    7       0           2
    8       0           3
    9       0           4
    """
    )
    edges = io.StringIO(
        """\
    left    right   parent  child
    0       1       5       0,2,3
    0       1       6       5
    0       1       7       1
    0       1       8       6
    0       1       8       7
    0       1       9       8
    0       1       9       4
    """
    )

    tree = tskit.load_text(
        nodes=nodes,
        edges=edges,
        strict=False,
    ).first()

    def test_all_zeros(self):
        genotypes = [0, 0, 0, 0, 0]
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 0

    def test_mutation_over_6(self):
        genotypes = [1, 0, 1, 1, 0]
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 1
        assert transitions[0] == tskit.Mutation(node=6, derived_state="1")

    def test_mutation_over_7(self):
        genotypes = [0, 1, 0, 0, 0]
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "0"
        assert len(transitions) == 1
        assert transitions[0] == tskit.Mutation(node=7, derived_state="1")

    def test_reversed_mutation_over_7(self):
        genotypes = [1, 0, 1, 1, 1]
        ancestral_state, transitions = self.do_map_mutations(self.tree, genotypes)
        assert ancestral_state == "1"
        assert len(transitions) == 1
        assert transitions[0] == tskit.Mutation(node=7, derived_state="0")


class TestParsimonyExamplesAncestralState(TestParsimonyBase):
    """
    Test fixing the ancestral state. Note that a mutation can occur above node 10
    to switch the ancestral state
    """

    #     10
    #    ┏━┻━┓
    #    ┃   9
    #    ┃ ┏━┻━┓
    #    ┃ ┃   8
    #    ┃ ┃ ┏━┻━┓
    #    ┃ ┃ ┃   7
    #    ┃ ┃ ┃ ┏━┻┓
    #    ┃ ┃ ┃ ┃  6
    #    ┃ ┃ ┃ ┃ ┏┻┓
    #    0 1 2 3 4 5
    tree = tskit.Tree.generate_comb(6)

    def test_mutation_over_0(self):
        genotypes = [1, 0, 0, 0, 0, 0]
        ancestral_state, transitions = self.do_map_mutations(
            self.tree, genotypes, ancestral_state=0
        )
        assert ancestral_state == "0"
        assert len(transitions) == 1
        assert transitions[0] == tskit.Mutation(node=0, parent=-1, derived_state="1")

        ancestral_state, transitions = self.do_map_mutations(
            self.tree, genotypes, ancestral_state=1
        )
        assert ancestral_state == "1"
        assert len(transitions) == 1
        assert transitions[0] == tskit.Mutation(node=9, parent=-1, derived_state="0")

    def test_mutation_over_3(self):
        genotypes = [0, 0, 0, 1, 0, 0]
        ancestral_state, transitions = self.do_map_mutations(
            self.tree, genotypes, ancestral_state=None
        )
        assert ancestral_state == "0"
        assert len(transitions) == 1
        assert transitions[0] == tskit.Mutation(node=3, parent=-1, derived_state="1")

        ancestral_state, transitions = self.do_map_mutations(
            self.tree, genotypes, ancestral_state=0
        )
        assert ancestral_state == "0"
        assert len(transitions) == 1
        assert transitions[0] == tskit.Mutation(node=3, parent=-1, derived_state="1")
        ancestral_state, transitions = self.do_map_mutations(
            self.tree, genotypes, ancestral_state=1
        )
        assert ancestral_state == "1"
        assert len(transitions) == 2
        assert transitions[0] == tskit.Mutation(node=10, parent=-1, derived_state="0")
        assert transitions[1] == tskit.Mutation(node=3, parent=0, derived_state="1")

    def test_novel_ancestral_state(self):
        # should put a single mutation above the root
        genotypes = [0, 0, 0, 0, 0, 0]
        for alleles in (["0", "1", "2", "3"], ["0", "1", "2", "3", None]):
            ancestral_state, transitions = self.do_map_mutations(
                self.tree, genotypes, alleles=alleles, ancestral_state=3
            )
            assert len(transitions) == 1
            assert transitions[0] == tskit.Mutation(node=10, derived_state="0")

    def test_mutations_over_root(self):
        tree = tskit.Tree.generate_star(6)
        # Mutations on root children
        genotypes = [0, 0, 0, 1, 1, 1]
        ancestral_state, transitions = self.do_map_mutations(
            tree, genotypes, ancestral_state=1
        )
        assert ancestral_state == "1"
        assert len(transitions) == 3
        assert all(m.derived_state == "0" for m in transitions)
        assert set(range(3)) == {m.node for m in transitions}

        # Should now switch to a mutation over the root
        genotypes = [0, 0, 0, 0, 1, 1]
        ancestral_state, transitions = self.do_map_mutations(
            tree, genotypes, ancestral_state=1
        )
        assert ancestral_state == "1"
        assert len(transitions) == 3
        assert transitions[0] == tskit.Mutation(node=tree.root, derived_state="0")
        assert all(m.derived_state == "1" for m in transitions[1:])
        assert all(m.parent == 0 for m in transitions[1:])
        assert {4, 5} == {m.node for m in transitions[1:]}

    def test_all_isolated_different_from_ancestral(self):
        ts = tskit.Tree.generate_star(6).tree_sequence
        ts = ts.decapitate(0)
        tree = ts.first()
        genotypes = [0, 0, 0, 1, 1, 1]
        ancestral_state, transitions = self.do_map_mutations(
            tree, genotypes, alleles=["A", "T", "G", "C"], ancestral_state=2
        )
        assert len(transitions) == 6
        assert all(m.parent == -1 for m in transitions)
        derived_states = [m.derived_state for m in transitions]
        assert derived_states.count("A") == 3
        assert derived_states.count("T") == 3
        assert {m.node for m in transitions if m.derived_state == "A"} == {0, 1, 2}
        assert {m.node for m in transitions if m.derived_state == "T"} == {3, 4, 5}

    def test_ancestral_as_string(self):
        genotypes = [1, 0, 0, 0, 0, 0]
        ancestral_state, transitions = self.do_map_mutations(
            self.tree, genotypes, alleles=["A", "T", "G", "C"], ancestral_state="A"
        )
        assert ancestral_state == "A"
        assert len(transitions) == 1
        assert transitions[0] == tskit.Mutation(node=0, parent=-1, derived_state="T")

    def test_bad_ancestral_state(self):
        genotypes = [0, 0, 0, 1, 0, 0]
        alleles = [str(j) for j in range(max(genotypes) + 1)]
        for bad, err in {
            2: "ancestral_state",
            -1: "ancestral_state",
            "A": "not in list",
        }.items():
            with pytest.raises(ValueError, match=err):
                hartigan_map_mutations(
                    self.tree, genotypes, alleles=alleles, ancestral_state=bad
                )
            with pytest.raises(ValueError, match=err):
                self.tree.map_mutations(genotypes, alleles=alleles, ancestral_state=bad)


class TestReconstructAllTuples:
    """
    Tests that the parsimony algorithm correctly round-trips all possible
    states.
    """

    def verify(self, ts, k):
        tables = ts.dump_tables()
        assert ts.num_trees == 1
        tree = ts.first()
        n = ts.num_samples
        m = k**n
        tables.sequence_length = m + 1
        tables.edges.set_columns(
            left=tables.edges.left,
            right=np.zeros_like(tables.edges.right) + tables.sequence_length,
            parent=tables.edges.parent,
            child=tables.edges.child,
        )
        G1 = np.zeros((m, n), dtype=np.int8)
        alleles = [str(j) for j in range(k)]
        for j, genotypes in enumerate(itertools.product(range(k), repeat=n)):
            G1[j] = genotypes
            ancestral_state, mutations = tree.map_mutations(G1[j], alleles)
            tables.sites.add_row(j, ancestral_state=ancestral_state)
            parent_offset = len(tables.mutations)
            for mutation in mutations:
                parent = mutation.parent
                if parent != tskit.NULL:
                    parent += parent_offset
                tables.mutations.append(mutation.replace(site=j, parent=parent))

        ts2 = tables.tree_sequence()
        G2 = np.zeros((m, n), dtype=np.int8)
        for j, variant in enumerate(ts2.variants()):
            alleles = np.array(list(map(int, variant.alleles)), dtype=np.int8)
            G2[j] = alleles[variant.genotypes]
        assert np.array_equal(G1, G2)

    def test_simple_n3_k2(self):
        ts = msprime.simulate(3, random_seed=4)
        self.verify(ts, 2)

    def test_simple_n3_k4(self):
        ts = msprime.simulate(3, random_seed=4)
        self.verify(ts, 4)

    def test_simple_n4_k2(self):
        ts = msprime.simulate(4, random_seed=4)
        self.verify(ts, 2)

    def test_simple_n4_k4(self):
        ts = msprime.simulate(4, random_seed=4)
        self.verify(ts, 4)

    def test_simple_n4_k5(self):
        ts = msprime.simulate(4, random_seed=4)
        self.verify(ts, 5)

    def test_simple_n5_k4(self):
        ts = msprime.simulate(5, random_seed=4)
        self.verify(ts, 4)

    def test_simple_n6_k3(self):
        ts = msprime.simulate(6, random_seed=4)
        self.verify(ts, 3)

    def test_root_polytomy_n5_k4(self):
        tree = tskit.Tree.unrank(5, (1, 0))
        self.verify(tree.tree_sequence, 4)

    def test_leaf_polytomy_n5_k4(self):
        tree = tskit.Tree.unrank(5, (7, 0))
        self.verify(tree.tree_sequence, 4)

    def test_leaf_polytomy_n5_k5(self):
        tree = tskit.Tree.unrank(5, (7, 0))
        self.verify(tree.tree_sequence, 5)


--- ../../tskit/python/tests/test_dict_encoding.py ---


"""
Test cases for the low-level dictionary encoding used to move
data around in C.
"""
import pathlib
import pickle

import _tskit
import lwt_interface.dict_encoding_testlib
import tskit


lwt_interface.dict_encoding_testlib.lwt_module = _tskit
# Bring the tests defined in dict_encoding_testlib into the current namespace
# so pytest will find and execute them.
from lwt_interface.dict_encoding_testlib import *  # noqa


def test_pickled_examples():
    seen_msprime = False
    test_dir = pathlib.Path(__file__).parent / "data/dict-encodings"
    for filename in test_dir.glob("*.pkl"):
        if "msprime" in str(filename):
            seen_msprime = True
        with open(test_dir / filename, "rb") as f:
            d = pickle.load(f)
            lwt = _tskit.LightweightTableCollection()
            lwt.fromdict(d)
            tskit.TableCollection.fromdict(d)
    # Check we've done something
    assert seen_msprime


--- ../../tskit/python/tests/test_topology.py ---


"""
Test cases for the supported topological variations and operations.
"""
import functools
import io
import itertools
import json
import random
import sys
import unittest

import msprime
import numpy as np
import pytest

import _tskit
import tests as tests
import tests.test_wright_fisher as wf
import tests.tsutil as tsutil
import tskit
import tskit.provenance as provenance


def simple_keep_intervals(tables, intervals, simplify=True, record_provenance=True):
    """
    Simple Python implementation of keep_intervals.
    """
    ts = tables.tree_sequence()
    last_stop = 0
    for start, stop in intervals:
        if start < 0 or stop > ts.sequence_length:
            raise ValueError("Slice bounds must be within the existing tree sequence")
        if start >= stop:
            raise ValueError("Interval error: start must be < stop")
        if start < last_stop:
            raise ValueError("Intervals must be disjoint")
        last_stop = stop
    tables.edges.clear()
    tables.sites.clear()
    tables.mutations.clear()
    for edge in ts.edges():
        for interval_left, interval_right in intervals:
            if not (edge.right <= interval_left or edge.left >= interval_right):
                left = max(interval_left, edge.left)
                right = min(interval_right, edge.right)
                tables.edges.append(edge.replace(left=left, right=right))
    for site in ts.sites():
        for interval_left, interval_right in intervals:
            if interval_left <= site.position < interval_right:
                site_id = tables.sites.append(site)
                for m in site.mutations:
                    tables.mutations.append(m.replace(site=site_id, parent=tskit.NULL))
    tables.build_index()
    tables.compute_mutation_parents()
    tables.sort()
    if simplify:
        tables.simplify(record_provenance=False)
    if record_provenance:
        parameters = {"command": "keep_intervals", "TODO": "add parameters"}
        tables.provenances.add_row(
            record=json.dumps(provenance.get_provenance_dict(parameters))
        )


def generate_segments(n, sequence_length=100, seed=None):
    rng = random.Random(seed)
    segs = []
    for j in range(n):
        left = rng.randint(0, sequence_length - 1)
        right = rng.randint(left + 1, sequence_length)
        assert left < right
        segs.append(tests.Segment(left, right, j))
    return segs


class ExampleTopologyMixin:
    """
    Some example topologies for tests cases.
    """

    def test_single_coalescent_tree(self):
        ts = msprime.simulate(10, random_seed=1, length=10)
        self.verify(ts)

    def test_coalescent_trees(self):
        ts = msprime.simulate(8, recombination_rate=5, random_seed=1, length=2)
        assert ts.num_trees > 2
        self.verify(ts)

    def test_coalescent_trees_internal_samples(self):
        ts = msprime.simulate(8, recombination_rate=5, random_seed=10, length=2)
        assert ts.num_trees > 2
        self.verify(tsutil.jiggle_samples(ts))

    def test_coalescent_trees_all_samples(self):
        ts = msprime.simulate(8, recombination_rate=5, random_seed=10, length=2)
        assert ts.num_trees > 2
        tables = ts.dump_tables()
        flags = np.zeros_like(tables.nodes.flags) + tskit.NODE_IS_SAMPLE
        tables.nodes.flags = flags
        self.verify(tables.tree_sequence())

    def test_wright_fisher_trees_unsimplified(self):
        tables = wf.wf_sim(10, 5, deep_history=False, seed=2)
        tables.sort()
        ts = tables.tree_sequence()
        self.verify(ts)

    def test_wright_fisher_trees_simplified(self):
        tables = wf.wf_sim(10, 5, deep_history=False, seed=1)
        tables.sort()
        ts = tables.tree_sequence()
        ts = ts.simplify()
        self.verify(ts)

    def test_wright_fisher_trees_simplified_one_gen(self):
        tables = wf.wf_sim(10, 1, deep_history=False, seed=1)
        tables.sort()
        ts = tables.tree_sequence()
        ts = ts.simplify()
        self.verify(ts)

    def test_nonbinary_trees(self):
        demographic_events = [
            msprime.SimpleBottleneck(time=1.0, population=0, proportion=0.95)
        ]
        ts = msprime.simulate(
            20,
            recombination_rate=10,
            mutation_rate=5,
            demographic_events=demographic_events,
            random_seed=7,
        )
        found = False
        for e in ts.edgesets():
            if len(e.children) > 2:
                found = True
        assert found
        self.verify(ts)

    def test_many_multiroot_trees(self):
        ts = msprime.simulate(7, recombination_rate=1, random_seed=10)
        assert ts.num_trees > 3
        ts = ts.decapitate(np.max(ts.tables.nodes.time) / 2)
        self.verify(ts)

    def test_multiroot_tree(self):
        ts = msprime.simulate(15, random_seed=10)
        ts = ts.decapitate(np.max(ts.tables.nodes.time) / 2)
        self.verify(ts)

    def test_all_missing_data(self):
        tables = tskit.TableCollection(1)
        for _ in range(10):
            tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        self.verify(tables.tree_sequence())


class TestOverlappingSegments:
    """
    Tests for the overlapping segments algorithm required for simplify.
    This test probably belongs somewhere else.
    """

    def test_random(self):
        segs = generate_segments(10, 20, 1)
        for left, right, X in tests.overlapping_segments(segs):
            assert right > left
            assert len(X) > 0

    def test_empty(self):
        ret = list(tests.overlapping_segments([]))
        assert len(ret) == 0

    def test_single_interval(self):
        for j in range(1, 10):
            segs = [tests.Segment(0, 1, j) for _ in range(j)]
            ret = list(tests.overlapping_segments(segs))
            assert len(ret) == 1
            left, right, X = ret[0]
            assert left == 0
            assert right == 1
            assert sorted(segs) == sorted(X)

    def test_stairs_down(self):
        segs = [tests.Segment(0, 1, 0), tests.Segment(0, 2, 1), tests.Segment(0, 3, 2)]
        ret = list(tests.overlapping_segments(segs))
        assert len(ret) == 3

        left, right, X = ret[0]
        assert left == 0
        assert right == 1
        assert sorted(X) == sorted(segs)

        left, right, X = ret[1]
        assert left == 1
        assert right == 2
        assert sorted(X) == sorted(segs[1:])

        left, right, X = ret[2]
        assert left == 2
        assert right == 3
        assert sorted(X) == sorted(segs[2:])

    def test_stairs_up(self):
        segs = [tests.Segment(0, 3, 0), tests.Segment(1, 3, 1), tests.Segment(2, 3, 2)]
        ret = list(tests.overlapping_segments(segs))
        assert len(ret) == 3

        left, right, X = ret[0]
        assert left == 0
        assert right == 1
        assert X == segs[:1]

        left, right, X = ret[1]
        assert left == 1
        assert right == 2
        assert sorted(X) == sorted(segs[:2])

        left, right, X = ret[2]
        assert left == 2
        assert right == 3
        assert sorted(X) == sorted(segs)

    def test_pyramid(self):
        segs = [tests.Segment(0, 5, 0), tests.Segment(1, 4, 1), tests.Segment(2, 3, 2)]
        ret = list(tests.overlapping_segments(segs))
        assert len(ret) == 5

        left, right, X = ret[0]
        assert left == 0
        assert right == 1
        assert X == segs[:1]

        left, right, X = ret[1]
        assert left == 1
        assert right == 2
        assert sorted(X) == sorted(segs[:2])

        left, right, X = ret[2]
        assert left == 2
        assert right == 3
        assert sorted(X) == sorted(segs)

        left, right, X = ret[3]
        assert left == 3
        assert right == 4
        assert sorted(X) == sorted(segs[:2])

        left, right, X = ret[4]
        assert left == 4
        assert right == 5
        assert sorted(X) == sorted(segs[:1])

    def test_gap(self):
        segs = [tests.Segment(0, 2, 0), tests.Segment(3, 4, 1)]
        ret = list(tests.overlapping_segments(segs))
        assert len(ret) == 2

        left, right, X = ret[0]
        assert left == 0
        assert right == 2
        assert X == segs[:1]

        left, right, X = ret[1]
        assert left == 3
        assert right == 4
        assert X == segs[1:]


class TopologyTestCase:
    """
    Superclass of test cases containing common utilities.
    """

    random_seed = 123456

    def assert_haplotypes_equal(self, ts1, ts2):
        h1 = list(ts1.haplotypes())
        h2 = list(ts2.haplotypes())
        assert h1 == h2

    def assert_variants_equal(self, ts1, ts2):
        for v1, v2 in zip(
            ts1.variants(copy=False),
            ts2.variants(copy=False),
        ):
            assert v1.alleles == v2.alleles
            assert np.array_equal(v1.genotypes, v2.genotypes)

    def check_num_samples(self, ts, x):
        """
        Compare against x, a list of tuples of the form
        `(tree number, parent, number of samples)`.
        """
        k = 0
        tss = ts.trees()
        t = next(tss)
        for j, node, nl in x:
            while k < j:
                t = next(tss)
                k += 1
            assert nl == t.num_samples(node)

    def check_num_tracked_samples(self, ts, tracked_samples, x):
        k = 0
        tss = ts.trees(tracked_samples=tracked_samples)
        t = next(tss)
        for j, node, nl in x:
            while k < j:
                t = next(tss)
                k += 1
            assert nl == t.num_tracked_samples(node)

    def check_sample_iterator(self, ts, x):
        """
        Compare against x, a list of tuples of the form
        `(tree number, node, sample ID list)`.
        """
        k = 0
        tss = ts.trees(sample_lists=True)
        t = next(tss)
        for j, node, samples in x:
            while k < j:
                t = next(tss)
                k += 1
            for u, v in zip(samples, t.samples(node)):
                assert u == v


class TestZeroRoots:
    """
    Tests that for the case in which we have zero samples and therefore
    zero roots in our trees.
    """

    def remove_samples(self, ts):
        tables = ts.dump_tables()
        tables.nodes.flags = np.zeros_like(tables.nodes.flags)
        return tables.tree_sequence()

    def verify(self, ts, no_root_ts):
        assert ts.num_trees == no_root_ts.num_trees
        for tree, no_root in zip(ts.trees(), no_root_ts.trees()):
            assert no_root.num_roots == 0
            assert no_root.left_root == tskit.NULL
            assert no_root.roots == []
            assert tree.parent_dict == no_root.parent_dict

    def test_single_tree(self):
        ts = msprime.simulate(10, random_seed=1)
        no_root_ts = self.remove_samples(ts)
        assert ts.num_trees == 1
        self.verify(ts, no_root_ts)

    def test_multiple_trees(self):
        ts = msprime.simulate(10, recombination_rate=2, random_seed=1)
        no_root_ts = self.remove_samples(ts)
        assert ts.num_trees > 1
        self.verify(ts, no_root_ts)


class TestEmptyTreeSequences(TopologyTestCase):
    """
    Tests covering tree sequences that have zero edges.
    """

    def test_zero_nodes(self):
        tables = tskit.TableCollection(1)
        ts = tables.tree_sequence()
        assert ts.sequence_length == 1
        assert ts.num_trees == 1
        assert ts.num_nodes == 0
        assert ts.num_edges == 0
        t = next(ts.trees())
        assert t.index == 0
        assert t.left_root == tskit.NULL
        assert t.interval == (0, 1)
        assert t.roots == []
        assert t.root == tskit.NULL
        assert t.parent_dict == {}
        assert t.virtual_root == 0
        assert t.left_child(t.virtual_root) == -1
        assert t.right_child(t.virtual_root) == -1
        assert list(t.nodes()) == []
        assert list(ts.haplotypes()) == []
        assert list(ts.variants()) == []
        methods = [
            t.parent,
            t.left_child,
            t.right_child,
            t.left_sib,
            t.right_sib,
            t.num_children,
        ]
        for method in methods:
            for u in [-1, 1, 100]:
                with pytest.raises(ValueError):
                    method(u)
        tsp = ts.simplify()
        assert tsp.num_nodes == 0
        assert tsp.num_edges == 0

    def test_one_node_zero_samples(self):
        tables = tskit.TableCollection(sequence_length=1)
        tables.nodes.add_row(time=0, flags=0)
        # Without a sequence length this should fail.
        ts = tables.tree_sequence()
        assert ts.sequence_length == 1
        assert ts.num_trees == 1
        assert ts.num_nodes == 1
        assert ts.sample_size == 0
        assert ts.num_edges == 0
        assert ts.num_sites == 0
        assert ts.num_mutations == 0
        t = next(ts.trees())
        assert t.index == 0
        assert t.left_root == tskit.NULL
        assert t.interval == (0, 1)
        assert t.roots == []
        assert t.root == tskit.NULL
        assert t.virtual_root == 1
        assert t.parent_dict == {}
        assert list(t.nodes()) == []
        assert list(ts.haplotypes()) == []
        assert list(ts.variants()) == []
        methods = [
            t.parent,
            t.left_child,
            t.right_child,
            t.left_sib,
            t.right_sib,
            t.num_children,
        ]
        for method in methods:
            expected = tskit.NULL if method != t.num_children else 0
            assert method(0) == expected
            for u in [-1, 2, 100]:
                with pytest.raises(ValueError):
                    method(u)

    def test_one_node_zero_samples_sites(self):
        tables = tskit.TableCollection(sequence_length=1)
        tables.nodes.add_row(time=0, flags=0)
        tables.sites.add_row(position=0.5, ancestral_state="0")
        tables.mutations.add_row(site=0, derived_state="1", node=0)
        ts = tables.tree_sequence()
        assert ts.sequence_length == 1
        assert ts.num_trees == 1
        assert ts.num_nodes == 1
        assert ts.sample_size == 0
        assert ts.num_edges == 0
        assert ts.num_sites == 1
        assert ts.num_mutations == 1
        t = next(ts.trees())
        assert t.index == 0
        assert t.left_root == tskit.NULL
        assert t.interval == (0, 1)
        assert t.roots == []
        assert t.root == tskit.NULL
        assert t.parent_dict == {}
        assert len(list(t.sites())) == 1
        assert list(t.nodes()) == []
        assert list(ts.haplotypes()) == []
        assert len(list(ts.variants())) == 1
        tsp = ts.simplify()
        assert tsp.num_nodes == 0
        assert tsp.num_edges == 0

    def test_one_node_one_sample(self):
        tables = tskit.TableCollection(sequence_length=1)
        tables.nodes.add_row(time=0, flags=tskit.NODE_IS_SAMPLE)
        ts = tables.tree_sequence()
        assert ts.sequence_length == 1
        assert ts.num_trees == 1
        assert ts.num_nodes == 1
        assert ts.sample_size == 1
        assert ts.num_edges == 0
        t = next(ts.trees())
        assert t.index == 0
        assert t.left_root == 0
        assert t.interval == (0, 1)
        assert t.roots == [0]
        assert t.root == 0
        assert t.virtual_root == 1
        assert t.parent_dict == {}
        assert list(t.nodes()) == [0]
        assert list(ts.haplotypes(isolated_as_missing=False)) == [""]
        assert list(ts.variants()) == []
        methods = [
            t.parent,
            t.left_child,
            t.right_child,
            t.left_sib,
            t.right_sib,
            t.num_children,
        ]
        for method in methods:
            expected = tskit.NULL if method != t.num_children else 0
            assert method(0) == expected
            for u in [-1, 2, 100]:
                with pytest.raises(ValueError):
                    method(u)
        tsp = ts.simplify()
        assert tsp.num_nodes == 1
        assert tsp.num_edges == 0

    def test_one_node_one_sample_sites(self):
        tables = tskit.TableCollection(sequence_length=1)
        tables.nodes.add_row(time=0, flags=tskit.NODE_IS_SAMPLE)
        tables.sites.add_row(position=0.5, ancestral_state="0")
        tables.mutations.add_row(site=0, derived_state="1", node=0)
        ts = tables.tree_sequence()
        assert ts.sequence_length == 1
        assert ts.num_trees == 1
        assert ts.num_nodes == 1
        assert ts.sample_size == 1
        assert ts.num_edges == 0
        assert ts.num_sites == 1
        assert ts.num_mutations == 1
        t = next(ts.trees())
        assert t.index == 0
        assert t.left_root == 0
        assert t.interval == (0, 1)
        assert t.roots == [0]
        assert t.root == 0
        assert t.virtual_root == 1
        assert t.parent_dict == {}
        assert list(t.nodes()) == [0]
        assert list(ts.haplotypes(isolated_as_missing=False)) == ["1"]
        assert len(list(ts.variants())) == 1
        methods = [
            t.parent,
            t.left_child,
            t.right_child,
            t.left_sib,
            t.right_sib,
            t.num_children,
        ]
        for method in methods:
            expected = tskit.NULL if method != t.num_children else 0
            assert method(0) == expected
            for u in [-1, 2, 100]:
                with pytest.raises(ValueError):
                    method(u)
        tsp = ts.simplify(filter_sites=False)
        assert tsp.num_nodes == 1
        assert tsp.num_edges == 0
        assert tsp.num_sites == 1


class TestHoleyTreeSequences(TopologyTestCase):
    """
    Tests for tree sequences in which we have partial (or no) trees defined
    over some of the sequence.
    """

    def verify_trees(self, ts, expected):
        observed = []
        for t in ts.trees():
            observed.append((t.interval, t.parent_dict))
        assert expected == observed
        # Test simple algorithm also.
        observed = []
        for interval, parent in tsutil.algorithm_T(ts):
            parent_dict = {j: parent[j] for j in range(ts.num_nodes) if parent[j] >= 0}
            observed.append((interval, parent_dict))
        assert expected == observed

    def verify_zero_roots(self, ts):
        for tree in ts.trees():
            assert tree.num_roots == 0
            assert tree.left_root == tskit.NULL
            assert tree.roots == []

    def test_simple_hole(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time
        0   1           0
        1   1           0
        2   0           1
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       2       0
        2       3       2       0
        0       1       2       1
        2       3       2       1
        """
        )
        ts = tskit.load_text(nodes, edges, strict=False)
        expected = [((0, 1), {0: 2, 1: 2}), ((1, 2), {}), ((2, 3), {0: 2, 1: 2})]
        self.verify_trees(ts, expected)

    def test_simple_hole_zero_roots(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time
        0   0           0
        1   0           0
        2   0           1
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       2       0
        2       3       2       0
        0       1       2       1
        2       3       2       1
        """
        )
        ts = tskit.load_text(nodes, edges, strict=False)
        expected = [((0, 1), {0: 2, 1: 2}), ((1, 2), {}), ((2, 3), {0: 2, 1: 2})]
        self.verify_trees(ts, expected)
        self.verify_zero_roots(ts)

    def test_initial_gap(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time
        0   1           0
        1   1           0
        2   0           1
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        1       2       2       0,1
        """
        )
        ts = tskit.load_text(nodes, edges, strict=False)
        expected = [((0, 1), {}), ((1, 2), {0: 2, 1: 2})]
        self.verify_trees(ts, expected)

    def test_initial_gap_zero_roots(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time
        0   0           0
        1   0           0
        2   0           1
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        1       2       2       0,1
        """
        )
        ts = tskit.load_text(nodes, edges, strict=False)
        expected = [((0, 1), {}), ((1, 2), {0: 2, 1: 2})]
        self.verify_trees(ts, expected)
        self.verify_zero_roots(ts)

    def test_final_gap(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time
        0   1           0
        1   1           0
        2   0           1
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       2       2       0,1
        """
        )
        ts = tskit.load_text(nodes, edges, sequence_length=3, strict=False)
        expected = [((0, 2), {0: 2, 1: 2}), ((2, 3), {})]
        self.verify_trees(ts, expected)

    def test_final_gap_zero_roots(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time
        0   0           0
        1   0           0
        2   0           1
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       2       2       0,1
        """
        )
        ts = tskit.load_text(nodes, edges, sequence_length=3, strict=False)
        expected = [((0, 2), {0: 2, 1: 2}), ((2, 3), {})]
        self.verify_trees(ts, expected)
        self.verify_zero_roots(ts)

    def test_initial_and_final_gap(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time
        0   1           0
        1   1           0
        2   0           1
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        1       2       2       0,1
        """
        )
        ts = tskit.load_text(nodes, edges, sequence_length=3, strict=False)
        expected = [((0, 1), {}), ((1, 2), {0: 2, 1: 2}), ((2, 3), {})]
        self.verify_trees(ts, expected)

    def test_initial_and_final_gap_zero_roots(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time
        0   0           0
        1   0           0
        2   0           1
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        1       2       2       0,1
        """
        )
        ts = tskit.load_text(nodes, edges, sequence_length=3, strict=False)
        expected = [((0, 1), {}), ((1, 2), {0: 2, 1: 2}), ((2, 3), {})]
        self.verify_trees(ts, expected)
        self.verify_zero_roots(ts)


class TestTsinferExamples(TopologyTestCase):
    """
    Test cases on troublesome topology examples that arose from tsinfer.
    """

    def test_no_last_tree(self):
        # The last tree was not being generated here because of a bug in
        # the low-level tree generation code.
        nodes = io.StringIO(
            """\
        id      is_sample   population      time
        0       1       -1              3.00000000000000
        1       1       -1              2.00000000000000
        2       1       -1              2.00000000000000
        3       1       -1              2.00000000000000
        4       1       -1              2.00000000000000
        5       1       -1              1.00000000000000
        6       1       -1              1.00000000000000
        7       1       -1              1.00000000000000
        8       1       -1              1.00000000000000
        9       1       -1              1.00000000000000
        10      1       -1              1.00000000000000
        """
        )
        edges = io.StringIO(
            """\
        id      left            right           parent  child
        0       62291.41659631  79679.17408763  1       5
        1       62291.41659631  62374.60889677  1       6
        2       122179.36037089 138345.43104411 1       7
        3       67608.32330402  79679.17408763  1       8
        4       122179.36037089 138345.43104411 1       8
        5       62291.41659631  79679.17408763  1       9
        6       126684.47550333 138345.43104411 1       10
        7       23972.05905068  62291.41659631  2       5
        8       79679.17408763  82278.53390076  2       5
        9       23972.05905068  62291.41659631  2       6
        10      79679.17408763  110914.43816806 2       7
        11      145458.28890561 189765.31932273 2       7
        12      79679.17408763  110914.43816806 2       8
        13      145458.28890561 200000.00000000 2       8
        14      23972.05905068  62291.41659631  2       9
        15      79679.17408763  110914.43816806 2       9
        16      145458.28890561 145581.18329797 2       10
        17      4331.62138785   23972.05905068  3       6
        18      4331.62138785   23972.05905068  3       9
        19      110914.43816806 122179.36037089 4       7
        20      138345.43104411 145458.28890561 4       7
        21      110914.43816806 122179.36037089 4       8
        22      138345.43104411 145458.28890561 4       8
        23      110914.43816806 112039.30503475 4       9
        24      138345.43104411 145458.28890561 4       10
        25      0.00000000      200000.00000000 0       1
        26      0.00000000      200000.00000000 0       2
        27      0.00000000      200000.00000000 0       3
        28      0.00000000      200000.00000000 0       4
        """
        )
        ts = tskit.load_text(nodes, edges, sequence_length=200000, strict=False)
        pts = tests.PythonTreeSequence(ts)
        num_trees = 0
        for _ in pts.trees():
            num_trees += 1
        assert num_trees == ts.num_trees
        n = 0
        for pt, t in zip(pts.trees(), ts.trees()):
            assert (pt.left, pt.right) == t.interval
            for j in range(ts.num_nodes):
                assert pt.parent[j] == t.parent(j)
                assert pt.left_child[j] == t.left_child(j)
                assert pt.right_child[j] == t.right_child(j)
                assert pt.left_sib[j] == t.left_sib(j)
                assert pt.right_sib[j] == t.right_sib(j)
                assert pt.num_children[j] == t.num_children(j)
            n += 1
        assert n == num_trees
        intervals = [t.interval for t in ts.trees()]
        assert intervals[0][0] == 0
        assert intervals[-1][-1] == ts.sequence_length


class TestRecordSquashing(TopologyTestCase):
    """
    Tests that we correctly squash adjacent equal records together.
    """

    def test_single_record(self):
        nodes = io.StringIO(
            """\
        id  is_sample   time
        0   1           0
        1   1           1
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       1       0
        1       2       1       0
        """
        )
        ts = tskit.load_text(nodes, edges, strict=False)
        tss, node_map = ts.simplify(map_nodes=True)
        assert list(node_map) == [0, 1]
        assert tss.dump_tables().nodes == ts.dump_tables().nodes
        simplified_edges = list(tss.edges())
        assert len(simplified_edges) == 1
        e = simplified_edges[0]
        assert e.left == 0
        assert e.right == 2

    def test_single_tree(self):
        ts = msprime.simulate(10, random_seed=self.random_seed)
        ts_redundant = tsutil.insert_redundant_breakpoints(ts)
        tss = ts_redundant.simplify()
        assert tss.dump_tables().nodes == ts.dump_tables().nodes
        assert tss.dump_tables().edges == ts.dump_tables().edges

    def test_many_trees(self):
        ts = msprime.simulate(20, recombination_rate=5, random_seed=self.random_seed)
        assert ts.num_trees > 2
        ts_redundant = tsutil.insert_redundant_breakpoints(ts)
        tss = ts_redundant.simplify()
        assert tss.dump_tables().nodes == ts.dump_tables().nodes
        assert tss.dump_tables().edges == ts.dump_tables().edges


class TestRedundantBreakpoints(TopologyTestCase):
    """
    Tests for dealing with redundant breakpoints within the tree sequence.
    These are records that may be squashed together into a single record.
    """

    def test_single_tree(self):
        ts = msprime.simulate(10, random_seed=self.random_seed)
        ts_redundant = tsutil.insert_redundant_breakpoints(ts)
        assert ts.sample_size == ts_redundant.sample_size
        assert ts.sequence_length == ts_redundant.sequence_length
        assert ts_redundant.num_trees == 2
        trees = [t.parent_dict for t in ts_redundant.trees()]
        assert len(trees) == 2
        assert trees[0] == trees[1]
        assert [t.parent_dict for t in ts.trees()][0] == trees[0]

    def test_many_trees(self):
        ts = msprime.simulate(20, recombination_rate=5, random_seed=self.random_seed)
        assert ts.num_trees > 2
        ts_redundant = tsutil.insert_redundant_breakpoints(ts)
        assert ts.sample_size == ts_redundant.sample_size
        assert ts.sequence_length == ts_redundant.sequence_length
        assert ts_redundant.num_trees > ts.num_trees
        assert ts_redundant.num_edges > ts.num_edges
        redundant_trees = ts_redundant.trees()
        redundant_t = next(redundant_trees)
        comparisons = 0
        for t in ts.trees():
            while (
                redundant_t is not None
                and redundant_t.interval.right <= t.interval.right
            ):
                assert t.parent_dict == redundant_t.parent_dict
                comparisons += 1
                redundant_t = next(redundant_trees, None)
        assert comparisons == ts_redundant.num_trees


class TestUnaryNodes(TopologyTestCase):
    """
    Tests for situations in which we have unary nodes in the tree sequence.
    """

    def test_simple_case(self):
        # Simple case where we have n = 2 and some unary nodes.
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           1
        3       0           1
        4       0           2
        5       0           3
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       2       0
        0       1       3       1
        0       1       4       2,3
        0       1       5       4
        """
        )
        sites = "position    ancestral_state\n"
        mutations = "site    node    derived_state\n"
        for j in range(5):
            position = j * 1 / 5
            sites += f"{position} 0\n"
            mutations += f"{j} {j} 1\n"
        ts = tskit.load_text(
            nodes=nodes,
            edges=edges,
            sites=io.StringIO(sites),
            mutations=io.StringIO(mutations),
            strict=False,
        )

        assert ts.sample_size == 2
        assert ts.num_nodes == 6
        assert ts.num_trees == 1
        assert ts.num_sites == 5
        assert ts.num_mutations == 5
        assert len(list(ts.edge_diffs())) == ts.num_trees
        t = next(ts.trees())
        assert t.parent_dict == {0: 2, 1: 3, 2: 4, 3: 4, 4: 5}
        assert t.mrca(0, 1) == 4
        assert t.mrca(0, 2) == 2
        assert t.mrca(0, 4) == 4
        assert t.mrca(0, 5) == 5
        assert t.mrca(0, 3) == 4
        H = list(ts.haplotypes())
        assert H[0] == "10101"
        assert H[1] == "01011"

    def test_ladder_tree(self):
        # We have a single tree with a long ladder of unary nodes along a path
        num_unary_nodes = 30
        n = 2
        nodes = """\
            is_sample   time
            1           0
            1           0
        """
        edges = """\
            left right parent child
            0    1     2      0
        """
        for j in range(num_unary_nodes + 2):
            nodes += f"0 {j + 2}\n"
        for j in range(num_unary_nodes):
            edges += f"0 1 {n + j + 1} {n + j}\n"
        root = num_unary_nodes + 3
        root_time = num_unary_nodes + 3
        edges += f"0    1     {root}      1,{num_unary_nodes + 2}\n"
        ts = tskit.load_text(io.StringIO(nodes), io.StringIO(edges), strict=False)
        t = ts.first()
        assert t.mrca(0, 1) == root
        assert t.tmrca(0, 1) == root_time
        ts_simplified, node_map = ts.simplify(map_nodes=True)
        test_map = [tskit.NULL for _ in range(ts.num_nodes)]
        test_map[0] = 0
        test_map[1] = 1
        test_map[root] = 2
        assert list(node_map) == test_map
        assert ts_simplified.num_edges == 2
        t = ts_simplified.first()
        assert t.mrca(0, 1) == 2
        assert t.tmrca(0, 1) == root_time
        ts_simplified = ts.simplify(keep_unary=True, record_provenance=False)
        assert ts_simplified.tables == ts.tables

    def verify_unary_tree_sequence(self, ts):
        """
        Take the specified tree sequence and produce an equivalent in which
        unary records have been interspersed, every other with an associated individual
        """
        assert ts.num_trees > 2
        assert ts.num_mutations > 2
        tables = ts.dump_tables()
        next_node = ts.num_nodes
        node_times = {j: node.time for j, node in enumerate(ts.nodes())}
        edges = []
        for i, e in enumerate(ts.edges()):
            node = ts.node(e.parent)
            t = node.time - 1e-14  # Arbitrary small value.
            next_node = len(tables.nodes)
            indiv = tables.individuals.add_row() if i % 2 == 0 else tskit.NULL
            tables.nodes.add_row(time=t, population=node.population, individual=indiv)
            edges.append(
                tskit.Edge(left=e.left, right=e.right, parent=next_node, child=e.child)
            )
            node_times[next_node] = t
            edges.append(
                tskit.Edge(left=e.left, right=e.right, parent=e.parent, child=next_node)
            )
        edges.sort(key=lambda e: node_times[e.parent])
        tables.edges.reset()
        for e in edges:
            tables.edges.append(e)
        ts_new = tables.tree_sequence()
        assert ts_new.num_edges > ts.num_edges
        self.assert_haplotypes_equal(ts, ts_new)
        self.assert_variants_equal(ts, ts_new)
        ts_simplified = ts_new.simplify()
        assert list(ts_simplified.records()) == list(ts.records())
        self.assert_haplotypes_equal(ts, ts_simplified)
        self.assert_variants_equal(ts, ts_simplified)
        assert len(list(ts.edge_diffs())) == ts.num_trees
        assert 0 < ts_new.num_individuals < ts_new.num_nodes

        for params in [
            {"keep_unary": False, "keep_unary_in_individuals": False},
            {"keep_unary": True, "keep_unary_in_individuals": False},
            {"keep_unary": False, "keep_unary_in_individuals": True},
        ]:
            s = tests.Simplifier(ts_new, ts_new.samples(), **params)
            py_ts, py_node_map = s.simplify()
            lib_ts, lib_node_map = ts_new.simplify(map_nodes=True, **params)
            py_tables = py_ts.dump_tables()
            lib_tables = lib_ts.dump_tables()
            lib_tables.assert_equals(py_tables, ignore_provenance=True)
            assert np.all(lib_node_map == py_node_map)

    def test_binary_tree_sequence_unary_nodes(self):
        ts = msprime.simulate(
            20, recombination_rate=5, mutation_rate=5, random_seed=self.random_seed
        )
        self.verify_unary_tree_sequence(ts)

    def test_nonbinary_tree_sequence_unary_nodes(self):
        demographic_events = [
            msprime.SimpleBottleneck(time=1.0, population=0, proportion=0.95)
        ]
        ts = msprime.simulate(
            20,
            recombination_rate=10,
            mutation_rate=5,
            demographic_events=demographic_events,
            random_seed=self.random_seed,
        )
        found = False
        for r in ts.edgesets():
            if len(r.children) > 2:
                found = True
        assert found
        self.verify_unary_tree_sequence(ts)


class TestGeneralSamples(TopologyTestCase):
    """
    Test cases in which we have samples at arbitrary nodes (i.e., not at
    {0,...,n - 1}).
    """

    def test_simple_case(self):
        # Simple case where we have n = 3 and samples starting at n.
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       0           2
        1       0           1
        2       1           0
        3       1           0
        4       1           0
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       1       2,3
        0       1       0       1,4
        """
        )
        sites = io.StringIO(
            """\
        position    ancestral_state
        0.1     0
        0.2     0
        0.3     0
        0.4     0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    derived_state
        0       2       1
        1       3       1
        2       4       1
        3       1       1
        """
        )
        ts = tskit.load_text(
            nodes=nodes, edges=edges, sites=sites, mutations=mutations, strict=False
        )

        assert ts.sample_size == 3
        assert list(ts.samples()) == [2, 3, 4]
        assert ts.num_nodes == 5
        assert ts.num_nodes == 5
        assert ts.num_sites == 4
        assert ts.num_mutations == 4
        assert len(list(ts.edge_diffs())) == ts.num_trees
        t = next(ts.trees())
        assert t.root == 0
        assert t.parent_dict == {1: 0, 2: 1, 3: 1, 4: 0}
        H = list(ts.haplotypes())
        assert H[0] == "1001"
        assert H[1] == "0101"
        assert H[2] == "0010"

        tss, node_map = ts.simplify(map_nodes=True)
        assert list(node_map) == [4, 3, 0, 1, 2]
        # We should have the same tree sequence just with canonicalised nodes.
        assert tss.sample_size == 3
        assert list(tss.samples()) == [0, 1, 2]
        assert tss.num_nodes == 5
        assert tss.num_trees == 1
        assert tss.num_sites == 4
        assert tss.num_mutations == 4
        assert len(list(ts.edge_diffs())) == ts.num_trees
        t = next(tss.trees())
        assert t.root == 4
        assert t.parent_dict == {0: 3, 1: 3, 2: 4, 3: 4}
        H = list(tss.haplotypes())
        assert H[0] == "1001"
        assert H[1] == "0101"
        assert H[2] == "0010"

    def verify_permuted_nodes(self, ts):
        """
        Take the specified tree sequence and permute the nodes, verifying that we
        get back a tree sequence with the correct properties.
        """
        # Mapping from the original nodes into nodes in the new tree sequence.
        node_map = list(range(ts.num_nodes))
        random.shuffle(node_map)
        # Change the permutation so that the relative order of samples is maintained.
        # Then, we should get back exactly the same tree sequence after simplify
        # and haplotypes and variants are also equal.
        samples = sorted(node_map[: ts.sample_size])
        node_map = samples + node_map[ts.sample_size :]
        permuted = tsutil.permute_nodes(ts, node_map)
        assert ts.sequence_length == permuted.sequence_length
        assert list(permuted.samples()) == samples
        assert list(permuted.haplotypes()) == list(ts.haplotypes())
        for v1, v2 in zip(
            permuted.variants(copy=False),
            ts.variants(copy=False),
        ):
            assert np.array_equal(v1.genotypes, v2.genotypes)

        assert ts.num_trees == permuted.num_trees
        j = 0
        for t1, t2 in zip(ts.trees(), permuted.trees()):
            t1_dict = {node_map[k]: node_map[v] for k, v in t1.parent_dict.items()}
            assert node_map[t1.root] == t2.root
            assert t1_dict == t2.parent_dict
            for u1 in t1.nodes():
                u2 = node_map[u1]
                assert sorted(node_map[v] for v in t1.samples(u1)) == sorted(
                    list(t2.samples(u2))
                )
            j += 1
        assert j == ts.num_trees

        # The simplified version of the permuted tree sequence should be in canonical
        # form, and identical to the original.
        simplified, s_node_map = permuted.simplify(map_nodes=True)

        for u, v in enumerate(node_map):
            assert s_node_map[v] == u
        ts.tables.assert_equals(simplified.tables, ignore_provenance=True)

    def test_single_tree_permuted_nodes(self):
        ts = msprime.simulate(10, mutation_rate=5, random_seed=self.random_seed)
        self.verify_permuted_nodes(ts)

    def test_binary_tree_sequence_permuted_nodes(self):
        ts = msprime.simulate(
            20, recombination_rate=5, mutation_rate=5, random_seed=self.random_seed
        )
        self.verify_permuted_nodes(ts)

    def test_nonbinary_tree_sequence_permuted_nodes(self):
        demographic_events = [
            msprime.SimpleBottleneck(time=1.0, population=0, proportion=0.95)
        ]
        ts = msprime.simulate(
            20,
            recombination_rate=10,
            mutation_rate=5,
            demographic_events=demographic_events,
            random_seed=self.random_seed,
        )
        found = False
        for e in ts.edgesets():
            if len(e.children) > 2:
                found = True
        assert found
        self.verify_permuted_nodes(ts)


class TestSimplifyExamples(TopologyTestCase):
    """
    Tests for simplify where we write out the input and expected output
    or we detect expected errors.
    """

    def verify_simplify(
        self,
        samples,
        *,
        filter_sites=True,
        keep_input_roots=False,
        filter_nodes=True,
        nodes_before=None,
        edges_before=None,
        sites_before=None,
        mutations_before=None,
        nodes_after=None,
        edges_after=None,
        sites_after=None,
        mutations_after=None,
        debug=False,
    ):
        """
        Verifies that if we run simplify on the specified input we get the
        required output.
        """
        before = tskit.load_text(
            nodes=io.StringIO(nodes_before),
            edges=io.StringIO(edges_before),
            sites=io.StringIO(sites_before) if sites_before is not None else None,
            mutations=(
                io.StringIO(mutations_before) if mutations_before is not None else None
            ),
            strict=False,
        )

        after = tskit.load_text(
            nodes=io.StringIO(nodes_after),
            edges=io.StringIO(edges_after),
            sites=io.StringIO(sites_after) if sites_after is not None else None,
            mutations=(
                io.StringIO(mutations_after) if mutations_after is not None else None
            ),
            strict=False,
            sequence_length=before.sequence_length,
        )

        result, _ = do_simplify(
            before,
            samples=samples,
            filter_sites=filter_sites,
            keep_input_roots=keep_input_roots,
            filter_nodes=filter_nodes,
            compare_lib=True,
        )
        if debug:
            print("before")
            print(before)
            print(before.draw_text())
            print("after")
            print(after)
            print(after.draw_text())
            print("result")
            print(result)
            print(result.draw_text())
        after.tables.assert_equals(result.tables)

    def test_unsorted_edges(self):
        # We have two nodes at the same time and interleave edges for
        # these nodes together. This is an error because all edges for
        # a given parent must be contigous.
        nodes_before = """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           1
        3       0           1
        """
        edges_before = """\
        left    right   parent  child
        0       1       2       0,1
        0       1       3       0,1
        1       2       2       0,1
        1       2       3       0,1
        """
        nodes = tskit.parse_nodes(io.StringIO(nodes_before), strict=False)
        edges = tskit.parse_edges(io.StringIO(edges_before), strict=False)
        # Cannot use load_text here because it calls sort()
        tables = tskit.TableCollection(sequence_length=2)
        tables.nodes.set_columns(**nodes.asdict())
        tables.edges.set_columns(**edges.asdict())
        with pytest.raises(_tskit.LibraryError):
            tables.simplify(samples=[0, 1])

    def test_single_binary_tree(self):
        #
        # 2        4
        #         / \
        # 1      3   \
        #       / \   \
        # 0   (0)(1)  (2)
        nodes_before = """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       0           1
        4       0           2
        """
        edges_before = """\
        left    right   parent  child
        0       1       3       0,1
        0       1       4       2,3
        """
        # We sample 0 and 2, so we get
        nodes_after = """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           2
        """
        edges_after = """\
        left    right   parent  child
        0       1       2       0,1
        """
        self.verify_simplify(
            samples=[0, 2],
            nodes_before=nodes_before,
            edges_before=edges_before,
            nodes_after=nodes_after,
            edges_after=edges_after,
        )

    def test_single_binary_tree_no_sample_nodes(self):
        #
        # 2        4
        #         / \
        # 1      3   \
        #       / \   \
        # 0   (0)(1)  (2)
        nodes_before = """\
        id      is_sample   time
        0       0           0
        1       0           0
        2       0           0
        3       0           1
        4       0           2
        """
        edges_before = """\
        left    right   parent  child
        0       1       3       0,1
        0       1       4       2,3
        """
        # We sample 0 and 2, so we get
        nodes_after = """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           2
        """
        edges_after = """\
        left    right   parent  child
        0       1       2       0,1
        """
        self.verify_simplify(
            samples=[0, 2],
            nodes_before=nodes_before,
            edges_before=edges_before,
            nodes_after=nodes_after,
            edges_after=edges_after,
        )

    def test_single_binary_tree_keep_input_root(self):
        #
        # 2        4
        #         / \
        # 1      3   \
        #       / \   \
        # 0   (0)(1)  (2)
        nodes_before = """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       0           1
        4       0           2
        """
        edges_before = """\
        left    right   parent  child
        0       1       3       0,1
        0       1       4       2,3
        """
        nodes_after = """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           1
        3       0           2
        """
        edges_after = """\
        left    right   parent  child
        0       1       2       0,1
        0       1       3       2
        """
        self.verify_simplify(
            samples=[0, 1],
            nodes_before=nodes_before,
            edges_before=edges_before,
            nodes_after=nodes_after,
            edges_after=edges_after,
            keep_input_roots=True,
        )

    def test_single_binary_tree_internal_sample(self):
        #
        # 2        4
        #         / \
        # 1     (3)  \
        #       / \   \
        # 0   (0)  1  (2)
        nodes_before = """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           0
        3       1           1
        4       0           2
        """
        edges_before = """\
        left    right   parent  child
        0       1       3       0,1
        0       1       4       2,3
        """
        # We sample 0 and 3, so we get
        nodes_after = """\
        id      is_sample   time
        0       1           0
        1       1           1
        """
        edges_after = """\
        left    right   parent  child
        0       1       1       0
        """
        self.verify_simplify(
            samples=[0, 3],
            nodes_before=nodes_before,
            edges_before=edges_before,
            nodes_after=nodes_after,
            edges_after=edges_after,
        )

    def test_single_binary_tree_internal_sample_meet_at_root(self):
        # 3          5
        #           / \
        # 2        4  (6)
        #         / \
        # 1     (3)  \
        #       / \   \
        # 0   (0)  1   2
        nodes_before = """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           0
        3       1           1
        4       0           2
        5       0           3
        6       1           2
        """
        edges_before = """\
        left    right   parent  child
        0       1       3       0,1
        0       1       4       2,3
        0       1       5       4,6
        """
        # We sample 0 and 3 and 6, so we get
        nodes_after = """\
        id      is_sample   time
        0       1           0
        1       1           1
        2       1           2
        3       0           3
        """
        edges_after = """\
        left    right   parent  child
        0       1       1       0
        0       1       3       1,2
        """
        self.verify_simplify(
            samples=[0, 3, 6],
            nodes_before=nodes_before,
            edges_before=edges_before,
            nodes_after=nodes_after,
            edges_after=edges_after,
        )

    def test_single_binary_tree_simple_mutations(self):
        # 3          5
        #           / \
        # 2        4   \
        #         / \   s0
        # 1      3   s1  \
        #       / \   \   \
        # 0   (0) (1)  2  (6)
        nodes_before = """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           0
        3       0           1
        4       0           2
        5       0           3
        6       1           0
        """
        edges_before = """\
        left    right   parent  child
        0       1       3       0,1
        0       1       4       2,3
        0       1       5       4,6
        """
        sites_before = """\
        id  position    ancestral_state
        0   0.1         0
        1   0.2         0
        """
        mutations_before = """\
        site    node    derived_state
        0       6       1
        1       2       1
        """

        # We sample 0 and 2 and 6, so we get
        nodes_after = """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       0           1
        3       0           3
        """
        edges_after = """\
        left    right   parent  child
        0       1       3       0,1
        0       1       4       2,3
        """
        sites_after = """\
        id  position    ancestral_state
        0   0.1         0
        """
        mutations_after = """\
        site    node    derived_state
        0       2       1
        """
        self.verify_simplify(
            samples=[0, 1, 6],
            nodes_before=nodes_before,
            edges_before=edges_before,
            sites_before=sites_before,
            mutations_before=mutations_before,
            nodes_after=nodes_after,
            edges_after=edges_after,
            sites_after=sites_after,
            mutations_after=mutations_after,
        )
        # If we don't filter the fixed sites, we should get the same
        # mutations and the original sites table back.
        self.verify_simplify(
            samples=[0, 1, 6],
            filter_sites=False,
            nodes_before=nodes_before,
            edges_before=edges_before,
            sites_before=sites_before,
            mutations_before=mutations_before,
            nodes_after=nodes_after,
            edges_after=edges_after,
            sites_after=sites_before,
            mutations_after=mutations_after,
        )

    def test_single_binary_tree_keep_roots_mutations(self):
        # 3          5
        #        m0 / \
        # 2        4   \
        #      m1 / \   \
        # 1      3   \   \
        #       / \   \   \
        # 0   (0) (1)  2   6
        nodes_before = """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           0
        3       0           1
        4       0           2
        5       0           3
        6       0           0
        """
        edges_before = """\
        left    right   parent  child
        0       1       3       0,1
        0       1       4       2,3
        0       1       5       4,6
        """
        sites_before = """\
        id  position    ancestral_state
        0   0.1         0
        """
        mutations_before = """\
        site    node    derived_state
        0       3       2
        0       4       1
        """

        # We sample 0 and 2
        nodes_after = """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           1
        3       0           3
        """
        edges_after = """\
        left    right   parent  child
        0       1       2       0,1
        0       1       3       2
        """
        sites_after = """\
        id  position    ancestral_state
        0   0.1         0
        """
        mutations_after = """\
        site    node    derived_state
        0       2       2
        0       2       1
        """
        self.verify_simplify(
            samples=[0, 1],
            nodes_before=nodes_before,
            edges_before=edges_before,
            sites_before=sites_before,
            mutations_before=mutations_before,
            nodes_after=nodes_after,
            edges_after=edges_after,
            sites_after=sites_after,
            mutations_after=mutations_after,
            keep_input_roots=True,
        )

    def test_place_mutations_with_and_without_roots(self):
        nodes_before = """\
        id      is_sample   time
        0       1           0
        1       0           1
        2       0           2
        """
        edges_before = """\
        left    right   parent  child
        0       2       1       0
        0       2       2       1
        """
        sites = """\
        id  position    ancestral_state
        0   1.0         0
        """
        mutations_before = """\
        site    node    derived_state time
        0       0       2             0
        0       1       1             1
        0       2       3             2
        """
        # expected result without keep_input_roots
        nodes_after = """\
        id      is_sample   time
        0       1           0
        """
        edges_after = """\
        left    right   parent  child
        """
        mutations_after = """\
        site    node    derived_state time
        0       0       2             0
        0       0       1             1
        0       0       3             2
        """
        # expected result with keep_input_roots
        nodes_after_keep = """\
        id      is_sample   time
        0       1           0
        1       0           2
        """
        edges_after_keep = """\
        left    right   parent  child
        0       2       1       0
        """
        mutations_after_keep = """\
        site    node    derived_state time
        0       0       2             0
        0       0       1             1
        0       1       3             2
        """
        self.verify_simplify(
            samples=[0],
            nodes_before=nodes_before,
            edges_before=edges_before,
            sites_before=sites,
            mutations_before=mutations_before,
            nodes_after=nodes_after,
            edges_after=edges_after,
            sites_after=sites,
            mutations_after=mutations_after,
            keep_input_roots=False,
        )
        self.verify_simplify(
            samples=[0],
            nodes_before=nodes_before,
            edges_before=edges_before,
            sites_before=sites,
            mutations_before=mutations_before,
            nodes_after=nodes_after_keep,
            edges_after=edges_after_keep,
            sites_after=sites,
            mutations_after=mutations_after_keep,
            keep_input_roots=True,
        )

    def test_overlapping_edges(self):
        nodes = """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           1
        """
        edges_before = """\
        left    right   parent  child
        0       2       2       0
        1       3       2       1
        """
        # We resolve the overlapping edges here. Since the flanking regions
        # have no interesting edges, these are left out of the output.
        edges_after = """\
        left    right   parent  child
        1       2       2       0,1
        """
        self.verify_simplify(
            samples=[0, 1],
            nodes_before=nodes,
            edges_before=edges_before,
            nodes_after=nodes,
            edges_after=edges_after,
        )

    def test_overlapping_edges_internal_samples(self):
        nodes = """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           1
        """
        edges = """\
        left    right   parent  child
        0       2       2       0
        1       3       2       1
        """
        self.verify_simplify(
            samples=[0, 1, 2],
            nodes_before=nodes,
            edges_before=edges,
            nodes_after=nodes,
            edges_after=edges,
        )

    def test_unary_edges_no_overlap(self):
        nodes_before = """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           1
        """
        edges_before = """\
        left    right   parent  child
        0       2       2       0
        2       3       2       1
        """
        # Because there is no overlap between the samples, we just get an
        # empty set of output edges.
        nodes_after = """\
        id      is_sample   time
        0       1           0
        1       1           0
        """
        edges_after = """\
        left    right   parent  child
        """
        self.verify_simplify(
            samples=[0, 1],
            nodes_before=nodes_before,
            edges_before=edges_before,
            nodes_after=nodes_after,
            edges_after=edges_after,
        )

    def test_unary_edges_no_overlap_internal_sample(self):
        nodes_before = """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           1
        """
        edges_before = """\
        left    right   parent  child
        0       1       2       0
        1       2       2       1
        """
        self.verify_simplify(
            samples=[0, 1, 2],
            nodes_before=nodes_before,
            edges_before=edges_before,
            nodes_after=nodes_before,
            edges_after=edges_before,
        )

    def test_keep_nodes(self):
        nodes_before = """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           1
        3       0           2
        4       0           3
        """
        edges_before = """\
        left    right   parent  child
        0       1       2       0
        0       1       2       1
        0       1       3       2
        0       1       4       3
        """
        edges_after = """\
        left    right   parent  child
        0       1       2       0
        0       1       2       1
        0       1       4       2
        """
        self.verify_simplify(
            samples=[0, 1],
            nodes_before=nodes_before,
            edges_before=edges_before,
            nodes_after=nodes_before,
            edges_after=edges_after,
            filter_nodes=False,
            keep_input_roots=True,
        )


class TestNonSampleExternalNodes(TopologyTestCase):
    """
    Tests for situations in which we have tips that are not samples.
    """

    def test_simple_case(self):
        # Simplest case where we have n = 2 and external non-sample nodes.
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           1
        3       0           0
        4       0           0
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       2       0,1,3,4
        """
        )
        sites = io.StringIO(
            """\
        id  position    ancestral_state
        0   0.1         0
        1   0.2         0
        2   0.3         0
        3   0.4         0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    derived_state
        0       0       1
        1       1       1
        2       3       1
        3       4       1
        """
        )
        ts = tskit.load_text(
            nodes=nodes, edges=edges, sites=sites, mutations=mutations, strict=False
        )
        assert ts.sample_size == 2
        assert ts.num_trees == 1
        assert ts.num_nodes == 5
        assert ts.num_sites == 4
        assert ts.num_mutations == 4
        t = next(ts.trees())
        assert t.parent_dict == {0: 2, 1: 2, 3: 2, 4: 2}
        assert t.root == 2
        ts_simplified, node_map = ts.simplify(map_nodes=True)
        assert list(node_map) == [0, 1, 2, -1, -1]
        assert ts_simplified.num_nodes == 3
        assert ts_simplified.num_trees == 1
        t = next(ts_simplified.trees())
        assert t.parent_dict == {0: 2, 1: 2}
        assert t.root == 2
        # We should have removed the two non-sample mutations.
        assert [s.position for s in t.sites()] == [0.1, 0.2]

    def test_unary_non_sample_external_nodes(self):
        # Take an ordinary tree sequence and put a bunch of external non
        # sample nodes on it.
        ts = msprime.simulate(
            15, recombination_rate=5, random_seed=self.random_seed, mutation_rate=5
        )
        assert ts.num_trees > 2
        assert ts.num_mutations > 2
        tables = ts.dump_tables()
        next_node = ts.num_nodes
        tables.edges.reset()
        for e in ts.edges():
            tables.edges.append(e)
            tables.edges.append(e.replace(child=next_node))
            tables.nodes.add_row(time=0)
            next_node += 1
        tables.sort()
        ts_new = tables.tree_sequence()
        assert ts_new.num_nodes == next_node
        assert ts_new.sample_size == ts.sample_size
        self.assert_haplotypes_equal(ts, ts_new)
        self.assert_variants_equal(ts, ts_new)
        ts_simplified = ts_new.simplify()
        assert ts_simplified.num_nodes == ts.num_nodes
        assert ts_simplified.sample_size == ts.sample_size
        assert list(ts_simplified.records()) == list(ts.records())
        self.assert_haplotypes_equal(ts, ts_simplified)
        self.assert_variants_equal(ts, ts_simplified)


class TestMultipleRoots(TopologyTestCase):
    """
    Tests for situations where we have multiple roots for the samples.
    """

    def test_simplest_degenerate_case(self):
        # Simplest case where we have n = 2 and no edges.
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        """
        )
        sites = io.StringIO(
            """\
        id  position    ancestral_state
        0   0.1         0
        1   0.2         0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    derived_state
        0       0         1
        1       1         1
        """
        )
        ts = tskit.load_text(
            nodes=nodes,
            edges=edges,
            sites=sites,
            mutations=mutations,
            sequence_length=1,
            strict=False,
        )
        assert ts.num_nodes == 2
        assert ts.num_trees == 1
        assert ts.num_sites == 2
        assert ts.num_mutations == 2
        t = next(ts.trees())
        assert t.parent_dict == {}
        assert sorted(t.roots) == [0, 1]
        assert list(ts.haplotypes(isolated_as_missing=False)) == ["10", "01"]
        assert np.array_equal(
            np.stack([v.genotypes for v in ts.variants(isolated_as_missing=False)]),
            [[1, 0], [0, 1]],
        )
        simplified = ts.simplify()
        t1 = ts.dump_tables()
        t2 = simplified.dump_tables()
        assert t1.nodes == t2.nodes
        assert t1.edges == t2.edges

    def test_simplest_non_degenerate_case(self):
        # Simplest case where we have n = 4 and two trees.
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       1           0
        4       0           1
        5       0           2
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       4       0,1
        0       1       5       2,3
        """
        )
        sites = io.StringIO(
            """\
        id  position    ancestral_state
        0   0.1         0
        1   0.2         0
        2   0.3         0
        3   0.4         0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    derived_state
        0       0       1
        1       1       1
        2       2       1
        3       3       1
        """
        )
        ts = tskit.load_text(
            nodes=nodes, edges=edges, sites=sites, mutations=mutations, strict=False
        )
        assert ts.num_nodes == 6
        assert ts.num_trees == 1
        assert ts.num_sites == 4
        assert ts.num_mutations == 4
        t = next(ts.trees())
        assert t.parent_dict == {0: 4, 1: 4, 2: 5, 3: 5}
        assert list(ts.haplotypes()) == ["1000", "0100", "0010", "0001"]
        assert np.array_equal(
            np.stack([v.genotypes for v in ts.variants()]),
            [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]],
        )
        assert t.mrca(0, 1) == 4
        assert t.mrca(0, 4) == 4
        assert t.mrca(2, 3) == 5
        assert t.mrca(0, 2) == tskit.NULL
        assert t.mrca(0, 3) == tskit.NULL
        assert t.mrca(2, 4) == tskit.NULL
        ts_simplified, node_map = ts.simplify(map_nodes=True)
        for j in range(4):
            assert node_map[j] == j
        assert ts_simplified.num_nodes == 6
        assert ts_simplified.num_trees == 1
        assert ts_simplified.num_sites == 4
        assert ts_simplified.num_mutations == 4
        t = next(ts_simplified.trees())
        assert t.parent_dict == {0: 4, 1: 4, 2: 5, 3: 5}

    def test_two_reducible_trees(self):
        # We have n = 4 and two trees, with some unary nodes and non-sample leaves
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       1           0
        4       0           1
        5       0           1
        6       0           2
        7       0           3
        8       0           0   # Non sample leaf
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1      4         0
        0       1      5         1
        0       1      6         4,5
        0       1      7         2,3,8
        """
        )
        sites = io.StringIO(
            """\
        id  position    ancestral_state
        0   0.1         0
        1   0.2         0
        2   0.3         0
        3   0.4         0
        4   0.5         0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    derived_state
        0       0       1
        1       1       1
        2       2       1
        3       3       1
        4       8       1
        """
        )
        ts = tskit.load_text(
            nodes=nodes, edges=edges, sites=sites, mutations=mutations, strict=False
        )
        assert ts.num_nodes == 9
        assert ts.num_trees == 1
        assert ts.num_sites == 5
        assert ts.num_mutations == 5
        t = next(ts.trees())
        assert t.parent_dict == {0: 4, 1: 5, 2: 7, 3: 7, 4: 6, 5: 6, 8: 7}
        assert list(ts.haplotypes()) == ["10000", "01000", "00100", "00010"]
        assert np.array_equal(
            np.stack([v.genotypes for v in ts.variants()]),
            [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1], [0, 0, 0, 0]],
        )
        assert t.mrca(0, 1) == 6
        assert t.mrca(2, 3) == 7
        assert t.mrca(2, 8) == 7
        assert t.mrca(0, 2) == tskit.NULL
        assert t.mrca(0, 3) == tskit.NULL
        assert t.mrca(0, 8) == tskit.NULL
        ts_simplified, node_map = ts.simplify(map_nodes=True)
        for j in range(4):
            assert node_map[j] == j
        assert ts_simplified.num_nodes == 6
        assert ts_simplified.num_trees == 1
        t = next(ts_simplified.trees())
        assert list(ts_simplified.haplotypes()) == ["1000", "0100", "0010", "0001"]
        assert np.array_equal(
            np.stack([v.genotypes for v in ts_simplified.variants()]),
            [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]],
        )
        # The site over the non-sample external node should have been discarded.
        sites = list(t.sites())
        assert sites[-1].position == 0.4
        assert t.parent_dict == {0: 4, 1: 4, 2: 5, 3: 5}

    def test_one_reducible_tree(self):
        # We have n = 4 and two trees. One tree is reducible and the other isn't.
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       1           0
        4       0           1
        5       0           1
        6       0           2
        7       0           3
        8       0           0   # Non sample leaf
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1      4         0
        0       1      5         1
        0       1      6         4,5
        0       1      7         2,3,8
        """
        )
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        assert ts.num_nodes == 9
        assert ts.num_trees == 1
        t = next(ts.trees())
        assert t.parent_dict == {0: 4, 1: 5, 2: 7, 3: 7, 4: 6, 5: 6, 8: 7}
        assert t.mrca(0, 1) == 6
        assert t.mrca(2, 3) == 7
        assert t.mrca(2, 8) == 7
        assert t.mrca(0, 2) == tskit.NULL
        assert t.mrca(0, 3) == tskit.NULL
        assert t.mrca(0, 8) == tskit.NULL
        ts_simplified = ts.simplify()
        assert ts_simplified.num_nodes == 6
        assert ts_simplified.num_trees == 1
        t = next(ts_simplified.trees())
        assert t.parent_dict == {0: 4, 1: 4, 2: 5, 3: 5}

    # NOTE: This test has not been checked since updating to the text representation
    # so there might be other problems with it.
    def test_mutations_over_roots(self):
        # Mutations over root nodes should be ok when we have multiple roots.
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       0           1
        4       0           2
        5       0           2
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       3       0,1
        0       1       4       3
        0       1       5       2
        """
        )
        sites = io.StringIO(
            """\
        id  position    ancestral_state
        0   0.1         0
        1   0.2         0
        2   0.3         0
        3   0.4         0
        4   0.5         0
        5   0.6         0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    derived_state
        0       0       1
        1       1       1
        2       3       1
        3       4       1
        4       2       1
        5       5       1
        """
        )
        ts = tskit.load_text(
            nodes=nodes, edges=edges, sites=sites, mutations=mutations, strict=False
        )
        assert ts.num_nodes == 6
        assert ts.num_trees == 1
        assert ts.num_sites == 6
        assert ts.num_mutations == 6
        t = next(ts.trees())
        assert len(list(t.sites())) == 6
        haplotypes = ["101100", "011100", "000011"]
        variants = [[1, 0, 0], [0, 1, 0], [1, 1, 0], [1, 1, 0], [0, 0, 1], [0, 0, 1]]
        assert list(ts.haplotypes()) == haplotypes
        assert np.array_equal(np.stack([v.genotypes for v in ts.variants()]), variants)
        ts_simplified = ts.simplify(filter_sites=False)
        assert list(ts_simplified.haplotypes(isolated_as_missing=False)) == haplotypes
        assert np.array_equal(
            np.stack(
                [v.genotypes for v in ts_simplified.variants(isolated_as_missing=False)]
            ),
            variants,
        )

    def test_break_single_tree(self):
        # Take a single largish tree from tskit, and remove the oldest record.
        # This breaks it into two subtrees.
        ts = msprime.simulate(20, random_seed=self.random_seed, mutation_rate=4)
        assert ts.num_mutations > 5
        tables = ts.dump_tables()
        tables.edges.set_columns(
            left=tables.edges.left[:-1],
            right=tables.edges.right[:-1],
            parent=tables.edges.parent[:-1],
            child=tables.edges.child[:-1],
        )
        ts_new = tables.tree_sequence()
        assert ts.sample_size == ts_new.sample_size
        assert ts.num_edges == ts_new.num_edges + 1
        assert ts.num_trees == ts_new.num_trees
        self.assert_haplotypes_equal(ts, ts_new)
        self.assert_variants_equal(ts, ts_new)
        roots = set()
        t_new = next(ts_new.trees())
        for u in ts_new.samples():
            while t_new.parent(u) != tskit.NULL:
                u = t_new.parent(u)
            roots.add(u)
        assert len(roots) == 2
        assert sorted(roots) == sorted(t_new.roots)


class TestWithVisuals(TopologyTestCase):
    """
    Some pedantic tests with ascii depictions of what's supposed to happen.
    """

    def verify_simplify_topology(self, ts, sample, haplotypes=False):
        # copies from test_highlevel.py
        new_ts, node_map = ts.simplify(sample, map_nodes=True)
        old_trees = ts.trees()
        old_tree = next(old_trees)
        assert ts.get_num_trees() >= new_ts.get_num_trees()
        for new_tree in new_ts.trees():
            new_left, new_right = new_tree.get_interval()
            old_left, old_right = old_tree.get_interval()
            # Skip ahead on the old tree until new_left is within its interval
            while old_right <= new_left:
                old_tree = next(old_trees)
                old_left, old_right = old_tree.get_interval()
            # If the TMRCA of all pairs of samples is the same, then we have the
            # same information. We limit this to at most 500 pairs
            pairs = itertools.islice(itertools.combinations(sample, 2), 500)
            for pair in pairs:
                mapped_pair = [node_map[u] for u in pair]
                mrca1 = old_tree.get_mrca(*pair)
                mrca2 = new_tree.get_mrca(*mapped_pair)
                assert mrca2 == node_map[mrca1]
        if haplotypes:
            orig_haps = list(ts.haplotypes())
            simp_haps = list(new_ts.haplotypes())
            for i, j in enumerate(sample):
                assert orig_haps[j] == simp_haps[i]

    def test_partial_non_sample_external_nodes(self):
        # A somewhat more complicated test case with a partially specified,
        # non-sampled tip.
        #
        # Here is the situation:
        #
        # 1.0             7
        # 0.7            / \                                            6
        #               /   \                                          / \
        # 0.5          /     5                      5                 /   5
        #             /     / \                    / \               /   / \
        # 0.4        /     /   4                  /   4             /   /   4
        #           /     /   / \                /   / \           /   /   / \
        #          /     /   3   \              /   /   \         /   /   3   \
        #         /     /         \            /   /     \       /   /         \
        # 0.0    0     1           2          1   0       2     0   1           2
        #
        #          (0.0, 0.2),                 (0.2, 0.8),         (0.8, 1.0)

        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       0           0.2  # Non sample leaf
        4       0           0.4
        5       0           0.5
        6       0           0.7
        7       0           1.0
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.0     0.2     4       2,3
        0.2     0.8     4       0,2
        0.8     1.0     4       2,3
        0.0     1.0     5       1,4
        0.8     1.0     6       0,5
        0.0     0.2     7       0,5
        """
        )
        true_trees = [
            {0: 7, 1: 5, 2: 4, 3: 4, 4: 5, 5: 7, 6: -1, 7: -1},
            {0: 4, 1: 5, 2: 4, 3: -1, 4: 5, 5: -1, 6: -1, 7: -1},
            {0: 6, 1: 5, 2: 4, 3: 4, 4: 5, 5: 6, 6: -1, 7: -1},
        ]
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        tree_dicts = [t.parent_dict for t in ts.trees()]
        assert ts.sample_size == 3
        assert ts.num_trees == 3
        assert ts.num_nodes == 8
        # check topologies agree:
        for a, t in zip(true_trees, tree_dicts):
            for k in a.keys():
                if k in t.keys():
                    assert t[k] == a[k]
                else:
                    assert a[k] == tskit.NULL
        # check .simplify() works here
        self.verify_simplify_topology(ts, [0, 1, 2])

    def test_partial_non_sample_external_nodes_2(self):
        # The same situation as above, but partial tip is labeled '7' not '3':
        #
        # 1.0          6
        # 0.7         / \                                       5
        #            /   \                                     / \
        # 0.5       /     4                 4                 /   4
        #          /     / \               / \               /   / \
        # 0.4     /     /   3             /   3             /   /   3
        #        /     /   / \           /   / \           /   /   / \
        #       /     /   7   \         /   /   \         /   /   7   \
        #      /     /         \       /   /     \       /   /         \
        # 0.0 0     1           2     1   0       2     0   1           2
        #
        #          (0.0, 0.2),         (0.2, 0.8),         (0.8, 1.0)
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       0           0.4
        4       0           0.5
        5       0           0.7
        6       0           1.0
        7       0           0    # Non sample leaf
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.0     0.2     3       2,7
        0.2     0.8     3       0,2
        0.8     1.0     3       2,7
        0.0     0.2     4       1,3
        0.2     0.8     4       1,3
        0.8     1.0     4       1,3
        0.8     1.0     5       0,4
        0.0     0.2     6       0,4
        """
        )
        true_trees = [
            {0: 6, 1: 4, 2: 3, 3: 4, 4: 6, 5: -1, 6: -1, 7: 3},
            {0: 3, 1: 4, 2: 3, 3: 4, 4: -1, 5: -1, 6: -1, 7: -1},
            {0: 5, 1: 4, 2: 3, 3: 4, 4: 5, 5: -1, 6: -1, 7: 3},
        ]
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        tree_dicts = [t.parent_dict for t in ts.trees()]
        # sample size check works here since 7 > 3
        assert ts.sample_size == 3
        assert ts.num_trees == 3
        assert ts.num_nodes == 8
        # check topologies agree:
        for a, t in zip(true_trees, tree_dicts):
            for k in a.keys():
                if k in t.keys():
                    assert t[k] == a[k]
                else:
                    assert a[k] == tskit.NULL
        self.verify_simplify_topology(ts, [0, 1, 2])

    def test_single_offspring_records(self):
        # Here we have inserted a single-offspring record
        # (for 6 on the left segment):
        #
        # 1.0             7
        # 0.7            / 6                                                  6
        #               /   \                                                / \
        # 0.5          /     5                       5                      /   5
        #             /     / \                     / \                    /   / \
        # 0.4        /     /   4                   /   4                  /   /   4
        # 0.3       /     /   / \                 /   / \                /   /   / \
        #          /     /   3   \               /   /   \              /   /   3   \
        #         /     /         \             /   /     \            /   /         \
        # 0.0    0     1           2           1   0       2          0   1           2
        #
        #          (0.0, 0.2),               (0.2, 0.8),              (0.8, 1.0)
        nodes = io.StringIO(
            """\
        id  is_sample   time
        0   1           0
        1   1           0
        2   1           0
        3   0           0       # Non sample leaf
        4   0           0.4
        5   0           0.5
        6   0           0.7
        7   0           1.0
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.0     0.2     4       2,3
        0.2     0.8     4       0,2
        0.8     1.0     4       2,3
        0.0     1.0     5       1,4
        0.8     1.0     6       0,5
        0.0     0.2     6       5
        0.0     0.2     7       0,6
        """
        )
        ts = tskit.load_text(nodes, edges, strict=False)
        true_trees = [
            {0: 7, 1: 5, 2: 4, 3: 4, 4: 5, 5: 6, 6: 7, 7: -1},
            {0: 4, 1: 5, 2: 4, 3: -1, 4: 5, 5: -1, 6: -1, 7: -1},
            {0: 6, 1: 5, 2: 4, 3: 4, 4: 5, 5: 6, 6: -1, 7: -1},
        ]
        tree_dicts = [t.parent_dict for t in ts.trees()]
        assert ts.sample_size == 3
        assert ts.num_trees == 3
        assert ts.num_nodes == 8
        # check topologies agree:
        for a, t in zip(true_trees, tree_dicts):
            for k in a.keys():
                if k in t.keys():
                    assert t[k] == a[k]
                else:
                    assert a[k] == tskit.NULL
        self.verify_simplify_topology(ts, [0, 1, 2])

    def test_many_single_offspring(self):
        # a more complex test with single offspring
        # With `(i,j,x)->k` denoting that individual `k` inherits from `i` on `[0,x)`
        #    and from `j` on `[x,1)`:
        # 1. Begin with an individual `3` (and another anonymous one) at `t=0`.
        # 2. `(3,?,1.0)->4` and `(3,?,1.0)->5` at `t=1`
        # 3. `(4,3,0.9)->6` and `(3,5,0.1)->7` and then `3` dies at `t=2`
        # 4. `(6,7,0.7)->8` at `t=3`
        # 5. `(8,6,0.8)->9` and `(7,8,0.2)->10` at `t=4`.
        # 6. `(3,9,0.6)->0` and `(9,10,0.5)->1` and `(10,4,0.4)->2` at `t=5`.
        # 7. We sample `0`, `1`, and `2`.
        # Here are the trees:
        # t                  |              |              |             |
        #
        # 0       --3--      |     --3--    |     --3--    |    --3--    |    --3--
        #        /  |  \     |    /  |  \   |    /     \   |   /     \   |   /     \
        # 1     4   |   5    |   4   *   5  |   4       5  |  4       5  |  4       5
        #       |\ / \ /|    |   |\   \     |   |\     /   |  |\     /   |  |\     /|
        # 2     | 6   7 |    |   | 6   7    |   | 6   7    |  | 6   7    |  | 6   7 |
        #       | |\ /| |    |   |  \  *    |   |  \  |    |  |  *       |  |  *    | ...
        # 3     | | 8 | |    |   |   8 |    |   *   8 *    |  |   8      |  |   8   |
        #       | |/ \| |    |   |  /  |    |   |  /  |    |  |  * *     |  |  / \  |
        # 4     | 9  10 |    |   | 9  10    |   | 9  10    |  | 9  10    |  | 9  10 |
        #       |/ \ / \|    |   |  \   *   |   |  \   \   |  |  \   *   |  |  \    |
        # 5     0   1   2    |   0   1   2  |   0   1   2  |  0   1   2  |  0   1   2
        #
        #                    |   0.0 - 0.1  |   0.1 - 0.2  |  0.2 - 0.4  |  0.4 - 0.5
        # ... continued:
        # t                  |             |             |             |
        #
        # 0         --3--    |    --3--    |    --3--    |    --3--    |    --3--
        #          /     \   |   /     \   |   /     \   |   /     \   |   /  |  \
        # 1       4       5  |  4       5  |  4       5  |  4       5  |  4   |   5
        #         |\     /|  |   \     /|  |   \     /|  |   \     /|  |     /   /|
        # 2       | 6   7 |  |    6   7 |  |    6   7 |  |    6   7 |  |    6   7 |
        #         |  \    |  |     \    |  |       /  |  |    |  /  |  |    |  /  |
        # 3  ...  |   8   |  |      8   |  |      8   |  |    | 8   |  |    | 8   |
        #         |  / \  |  |     / \  |  |     / \  |  |    |  \  |  |    |  \  |
        # 4       | 9  10 |  |    9  10 |  |    9  10 |  |    9  10 |  |    9  10 |
        #         |    /  |  |   /   /  |  |   /   /  |  |   /   /  |  |   /   /  |
        # 5       0   1   2  |  0   1   2  |  0   1   2  |  0   1   2  |  0   1   2
        #
        #         0.5 - 0.6  |  0.6 - 0.7  |  0.7 - 0.8  |  0.8 - 0.9  |  0.9 - 1.0

        true_trees = [
            {0: 4, 1: 9, 2: 10, 3: -1, 4: 3, 5: 3, 6: 4, 7: 3, 8: 6, 9: 8, 10: 7},
            {0: 4, 1: 9, 2: 10, 3: -1, 4: 3, 5: 3, 6: 4, 7: 5, 8: 6, 9: 8, 10: 7},
            {0: 4, 1: 9, 2: 10, 3: -1, 4: 3, 5: 3, 6: 4, 7: 5, 8: 6, 9: 8, 10: 8},
            {0: 4, 1: 9, 2: 5, 3: -1, 4: 3, 5: 3, 6: 4, 7: 5, 8: 6, 9: 8, 10: 8},
            {0: 4, 1: 10, 2: 5, 3: -1, 4: 3, 5: 3, 6: 4, 7: 5, 8: 6, 9: 8, 10: 8},
            {0: 9, 1: 10, 2: 5, 3: -1, 4: 3, 5: 3, 6: 4, 7: 5, 8: 6, 9: 8, 10: 8},
            {0: 9, 1: 10, 2: 5, 3: -1, 4: 3, 5: 3, 6: 4, 7: 5, 8: 7, 9: 8, 10: 8},
            {0: 9, 1: 10, 2: 5, 3: -1, 4: 3, 5: 3, 6: 4, 7: 5, 8: 7, 9: 6, 10: 8},
            {0: 9, 1: 10, 2: 5, 3: -1, 4: 3, 5: 3, 6: 3, 7: 5, 8: 7, 9: 6, 10: 8},
        ]
        true_haplotypes = ["0100", "0001", "1110"]
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       0           5
        4       0           4
        5       0           4
        6       0           3
        7       0           3
        8       0           2
        9       0           1
        10      0           1
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.5     1.0     10      1
        0.0     0.4     10      2
        0.6     1.0     9       0
        0.0     0.5     9       1
        0.8     1.0     8       10
        0.2     0.8     8       9,10
        0.0     0.2     8       9
        0.7     1.0     7       8
        0.0     0.2     7       10
        0.8     1.0     6       9
        0.0     0.7     6       8
        0.4     1.0     5       2,7
        0.1     0.4     5       7
        0.6     0.9     4       6
        0.0     0.6     4       0,6
        0.9     1.0     3       4,5,6
        0.1     0.9     3       4,5
        0.0     0.1     3       4,5,7
        """
        )
        sites = io.StringIO(
            """\
        position    ancestral_state
        0.05        0
        0.15        0
        0.25        0
        0.4         0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    derived_state   parent
        0       7       1               -1
        0      10       0               0
        0       2       1               1
        1       0       1               -1
        1      10       1               -1
        2       8       1               -1
        2       9       0               5
        2      10       0               5
        2       2       1               7
        3       8       1               -1
        """
        )
        ts = tskit.load_text(nodes, edges, sites, mutations, strict=False)
        tree_dicts = [t.parent_dict for t in ts.trees()]
        assert ts.sample_size == 3
        assert ts.num_trees == len(true_trees)
        assert ts.num_nodes == 11
        assert len(list(ts.edge_diffs())) == ts.num_trees
        # check topologies agree:
        for a, t in zip(true_trees, tree_dicts):
            for k in a.keys():
                if k in t.keys():
                    assert t[k] == a[k]
                else:
                    assert a[k] == tskit.NULL
        for j, x in enumerate(ts.haplotypes()):
            assert x == true_haplotypes[j]
        self.verify_simplify_topology(ts, [0, 1, 2], haplotypes=True)
        self.verify_simplify_topology(ts, [1, 0, 2], haplotypes=True)
        self.verify_simplify_topology(ts, [0, 1], haplotypes=False)
        self.verify_simplify_topology(ts, [1, 2], haplotypes=False)
        self.verify_simplify_topology(ts, [2, 0], haplotypes=False)

    def test_tricky_switches(self):
        # suppose the topology has:
        # left right parent child
        #  0.0   0.5      6      0,1
        #  0.5   1.0      6      4,5
        #  0.0   0.4      7      2,3
        #
        # --------------------------
        #
        #        12         .        12         .        12         .
        #       /  \        .       /  \        .       /  \        .
        #     11    \       .      /    \       .      /    \       .
        #     / \    \      .     /     10      .     /     10      .
        #    /   \    \     .    /     /  \     .    /     /  \     .
        #   6     7    8    .   6     9    8    .   6     9    8    .
        #  / \   / \   /\   .  / \   / \   /\   .  / \   / \   /\   .
        # 0   1 2   3 4  5  . 0   1 2   3 4  5  . 4   5 2   3 0  1  .
        #                   .                   .                   .
        # 0.0              0.4                 0.5                 1.0
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       1           0
        4       1           0
        5       1           0
        6       0           1
        7       0           1
        8       0           1
        9       0           1
        10      0           2
        11      0           3
        12      0           4
        """
        )
        edges = io.StringIO(
            """\
        left right parent child
        0.0  0.5   6      0
        0.0  0.5   6      1
        0.5  1.0   6      4
        0.5  1.0   6      5
        0.0  0.4   7      2,3
        0.5  1.0   8      0
        0.5  1.0   8      1
        0.0  0.5   8      4
        0.0  0.5   8      5
        0.4  1.0   9      2,3
        0.4  1.0   10     8,9
        0.0  0.4   11     6,7
        0.4  1.0   12     6
        0.0  0.4   12     8
        0.4  1.0   12     10
        0.0  0.4   12     11
        """
        )
        true_trees = [
            {
                0: 6,
                1: 6,
                2: 7,
                3: 7,
                4: 8,
                5: 8,
                6: 11,
                7: 11,
                8: 12,
                9: -1,
                10: -1,
                11: 12,
                12: -1,
            },
            {
                0: 6,
                1: 6,
                2: 9,
                3: 9,
                4: 8,
                5: 8,
                6: 12,
                7: -1,
                8: 10,
                9: 10,
                10: 12,
                11: -1,
                12: -1,
            },
            {
                0: 8,
                1: 8,
                2: 9,
                3: 9,
                4: 6,
                5: 6,
                6: 12,
                7: -1,
                8: 10,
                9: 10,
                10: 12,
                11: -1,
                12: -1,
            },
        ]
        ts = tskit.load_text(nodes, edges, strict=False)
        tree_dicts = [t.parent_dict for t in ts.trees()]
        assert ts.sample_size == 6
        assert ts.num_trees == len(true_trees)
        assert ts.num_nodes == 13
        assert len(list(ts.edge_diffs())) == ts.num_trees
        # check topologies agree:
        for a, t in zip(true_trees, tree_dicts):
            for k in a.keys():
                if k in t.keys():
                    assert t[k] == a[k]
                else:
                    assert a[k] == tskit.NULL
        self.verify_simplify_topology(ts, [0, 2])
        self.verify_simplify_topology(ts, [0, 4])
        self.verify_simplify_topology(ts, [2, 4])

    def test_tricky_simplify(self):
        # Continue as above but invoke simplfy:
        #
        #         12         .          12         .
        #        /  \        .         /  \        .
        #      11    \       .       11    \       .
        #      / \    \      .       / \    \      .
        #    13   \    \     .      /  15    \     .
        #    / \   \    \    .     /   / \    \    .
        #   6  14   7    8   .    6  14   7    8   .
        #  / \     / \   /\  .   / \     / \   /\  .
        # 0   1   2   3 4  5 .  0   1   2   3 4  5 .
        #                    .                     .
        # 0.0               0.1                   0.4
        #
        #  .        12         .        12         .
        #  .       /  \        .       /  \        .
        #  .      /    \       .      /    \       .
        #  .     /     10      .     /     10      .
        #  .    /     /  \     .    /     /  \     .
        #  .   6     9    8    .   6     9    8    .
        #  .  / \   / \   /\   .  / \   / \   /\   .
        #  . 0   1 2   3 4  5  . 4   5 2   3 0  1  .
        #  .                   .                   .
        # 0.4                 0.5                 1.0
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       1           0
        4       1           0
        5       1           0
        6       0           1
        7       0           1
        8       0           1
        9       0           1
        10      0           2
        11      0           3
        12      0           4
        13      0           2
        14      0           1
        15      0           2
        """
        )
        edges = io.StringIO(
            """\
        left right parent child
        0.0  0.5   6      0,1
        0.5  1.0   6      4,5
        0.0  0.4   7      2,3
        0.0  0.5   8      4,5
        0.5  1.0   8      0,1
        0.4  1.0   9      2,3
        0.4  1.0   10     8,9
        0.0  0.1   13     6,14
        0.1  0.4   15     7,14
        0.0  0.1   11     7,13
        0.1  0.4   11     6,15
        0.0  0.4   12     8,11
        0.4  1.0   12     6,10
        """
        )
        true_trees = [
            {
                0: 6,
                1: 6,
                2: 7,
                3: 7,
                4: 8,
                5: 8,
                6: 11,
                7: 11,
                8: 12,
                9: -1,
                10: -1,
                11: 12,
                12: -1,
            },
            {
                0: 6,
                1: 6,
                2: 9,
                3: 9,
                4: 8,
                5: 8,
                6: 12,
                7: -1,
                8: 10,
                9: 10,
                10: 12,
                11: -1,
                12: -1,
            },
            {
                0: 8,
                1: 8,
                2: 9,
                3: 9,
                4: 6,
                5: 6,
                6: 12,
                7: -1,
                8: 10,
                9: 10,
                10: 12,
                11: -1,
                12: -1,
            },
        ]
        big_ts = tskit.load_text(nodes, edges, strict=False)
        assert big_ts.num_trees == 1 + len(true_trees)
        assert big_ts.num_nodes == 16
        ts, node_map = big_ts.simplify(map_nodes=True)
        assert list(node_map[:6]) == list(range(6))
        assert ts.sample_size == 6
        assert ts.num_nodes == 13

    def test_ancestral_samples(self):
        # Check that specifying samples to be not at time 0.0 works.
        #
        # 1.0             7
        # 0.7            / \                      8                     6
        #               /   \                    / \                   / \
        # 0.5          /     5                  /   5                 /   5
        #             /     / \                /   / \               /   / \
        # 0.4        /     /   4              /   /   4             /   /   4
        #           /     /   / \            /   /   / \           /   /   / \
        # 0.2      /     /   3   \          3   /   /   \         /   /   3   \
        #         /     /    *    \         *  /   /     \       /   /    *    \
        # 0.0    0     1           2          1   0       2     0   1           2
        #              *           *          *           *         *           *
        #          (0.0, 0.2),                 (0.2, 0.8),         (0.8, 1.0)
        #
        # Simplified, keeping [1,2,3]
        #
        # 1.0
        # 0.7                                     5
        #                                        / \
        # 0.5                4                  /   4                     4
        #                   / \                /   / \                   / \
        # 0.4              /   3              /   /   3                 /   3
        #                 /   / \            /   /     \               /   / \
        # 0.2            /   2   \          2   /       \             /   2   \
        #               /    *    \         *  /         \           /    *    \
        # 0.0          0           1          0           1         0           1
        #              *           *          *           *         *           *
        #          (0.0, 0.2),                 (0.2, 0.8),         (0.8, 1.0)

        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       0           0
        1       1           0
        2       1           0
        3       1           0.2
        4       0           0.4
        5       0           0.5
        6       0           0.7
        7       0           1.0
        8       0           0.8
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.0     0.2     4       2,3
        0.2     0.8     4       0,2
        0.8     1.0     4       2,3
        0.0     1.0     5       1,4
        0.8     1.0     6       0,5
        0.2     0.8     8       3,5
        0.0     0.2     7       0,5
        """
        )
        first_ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        ts, node_map = first_ts.simplify(map_nodes=True)
        true_trees = [
            {0: 7, 1: 5, 2: 4, 3: 4, 4: 5, 5: 7, 6: -1, 7: -1},
            {0: 4, 1: 5, 2: 4, 3: 8, 4: 5, 5: 8, 6: -1, 7: -1},
            {0: 6, 1: 5, 2: 4, 3: 4, 4: 5, 5: 6, 6: -1, 7: -1},
        ]
        # maps [1,2,3] -> [0,1,2]
        assert node_map[1] == 0
        assert node_map[2] == 1
        assert node_map[3] == 2
        true_simplified_trees = [
            {0: 4, 1: 3, 2: 3, 3: 4},
            {0: 4, 1: 4, 2: 5, 4: 5},
            {0: 4, 1: 3, 2: 3, 3: 4},
        ]
        assert first_ts.sample_size == 3
        assert ts.sample_size == 3
        assert first_ts.num_trees == 3
        assert ts.num_trees == 3
        assert first_ts.num_nodes == 9
        assert ts.num_nodes == 6
        assert first_ts.node(3).time == 0.2
        assert ts.node(2).time == 0.2
        # check topologies agree:
        tree_dicts = [t.parent_dict for t in first_ts.trees()]
        for a, t in zip(true_trees, tree_dicts):
            for k in a.keys():
                if k in t.keys():
                    assert t[k] == a[k]
                else:
                    assert a[k] == tskit.NULL
        tree_simplified_dicts = [t.parent_dict for t in ts.trees()]
        for a, t in zip(true_simplified_trees, tree_simplified_dicts):
            for k in a.keys():
                if k in t.keys():
                    assert t[k] == a[k]
                else:
                    assert a[k] == tskit.NULL
        # check .simplify() works here
        self.verify_simplify_topology(first_ts, [1, 2, 3])

    def test_all_ancestral_samples(self):
        # Check that specifying samples all to be not at time 0.0 works.
        #
        # 1.0             7
        # 0.7            / \                      8                     6
        #               /   \                    / \                   / \
        # 0.5          /     5                  /   5                 /   5
        #             /     / \                /   / \               /   / \
        # 0.4        /     /   4              /   /   4             /   /   4
        #           /     /   / \            /   /   / \           /   /   / \
        # 0.2      /     /   3   \          3   /   /   \         /   /   3   \
        #         /     1    *    2         *  1   /     2       /   1    *    2
        # 0.0    0      *         *            *  0      *      0    *         *
        #
        #          (0.0, 0.2),                 (0.2, 0.8),         (0.8, 1.0)

        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       0           0
        1       1           0.1
        2       1           0.1
        3       1           0.2
        4       0           0.4
        5       0           0.5
        6       0           0.7
        7       0           1.0
        8       0           0.8
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.0     0.2     4       2,3
        0.2     0.8     4       0,2
        0.8     1.0     4       2,3
        0.0     1.0     5       1,4
        0.8     1.0     6       0,5
        0.2     0.8     8       3,5
        0.0     0.2     7       0,5
        """
        )
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        true_trees = [
            {0: 7, 1: 5, 2: 4, 3: 4, 4: 5, 5: 7, 6: -1, 7: -1},
            {0: 4, 1: 5, 2: 4, 3: 8, 4: 5, 5: 8, 6: -1, 7: -1},
            {0: 6, 1: 5, 2: 4, 3: 4, 4: 5, 5: 6, 6: -1, 7: -1},
        ]
        assert ts.sample_size == 3
        assert ts.num_trees == 3
        assert ts.num_nodes == 9
        assert ts.node(0).time == 0.0
        assert ts.node(1).time == 0.1
        assert ts.node(2).time == 0.1
        assert ts.node(3).time == 0.2
        # check topologies agree:
        tree_dicts = [t.parent_dict for t in ts.trees()]
        for a, t in zip(true_trees, tree_dicts):
            for k in a.keys():
                if k in t.keys():
                    assert t[k] == a[k]
                else:
                    assert a[k] == tskit.NULL
        # check .simplify() works here
        self.verify_simplify_topology(ts, [1, 2, 3])

    def test_internal_sampled_node(self):
        # 1.0             7
        # 0.7            / \                      8                     6
        #               /   \                    / \                   / \
        # 0.5          /     5                  /   5                 /   5
        #             /     /*\                /   /*\               /   /*\
        # 0.4        /     /   4              /   /   4             /   /   4
        #           /     /   / \            /   /   / \           /   /   / \
        # 0.2      /     /   3   \          3   /   /   \         /   /   3   \
        #         /     1    *    2         *  1   /     2       /   1    *    2
        # 0.0    0      *         *            *  0      *      0    *         *
        #
        #          (0.0, 0.2),                 (0.2, 0.8),         (0.8, 1.0)
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       0           0
        1       1           0.1
        2       1           0.1
        3       1           0.2
        4       0           0.4
        5       1           0.5
        6       0           0.7
        7       0           1.0
        8       0           0.8
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.0     0.2     4       2,3
        0.2     0.8     4       0,2
        0.8     1.0     4       2,3
        0.0     1.0     5       1,4
        0.8     1.0     6       0,5
        0.2     0.8     8       3,5
        0.0     0.2     7       0,5
        """
        )
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        true_trees = [
            {0: 7, 1: 5, 2: 4, 3: 4, 4: 5, 5: 7, 6: -1, 7: -1},
            {0: 4, 1: 5, 2: 4, 3: 8, 4: 5, 5: 8, 6: -1, 7: -1},
            {0: 6, 1: 5, 2: 4, 3: 4, 4: 5, 5: 6, 6: -1, 7: -1},
        ]
        assert ts.sample_size == 4
        assert ts.num_trees == 3
        assert ts.num_nodes == 9
        assert ts.node(0).time == 0.0
        assert ts.node(1).time == 0.1
        assert ts.node(2).time == 0.1
        assert ts.node(3).time == 0.2
        # check topologies agree:
        tree_dicts = [t.parent_dict for t in ts.trees()]
        for a, t in zip(true_trees, tree_dicts):
            for k in a.keys():
                if k in t.keys():
                    assert t[k] == a[k]
                else:
                    assert a[k] == tskit.NULL
        # check .simplify() works here
        self.verify_simplify_topology(ts, [1, 2, 3])
        self.check_num_samples(
            ts,
            [
                (0, 5, 4),
                (0, 2, 1),
                (0, 7, 4),
                (0, 4, 2),
                (1, 4, 1),
                (1, 5, 3),
                (1, 8, 4),
                (1, 0, 0),
                (2, 5, 4),
                (2, 1, 1),
            ],
        )
        self.check_num_tracked_samples(
            ts,
            [1, 2, 5],
            [
                (0, 5, 3),
                (0, 2, 1),
                (0, 7, 3),
                (0, 4, 1),
                (1, 4, 1),
                (1, 5, 3),
                (1, 8, 3),
                (1, 0, 0),
                (2, 5, 3),
                (2, 1, 1),
            ],
        )
        self.check_sample_iterator(
            ts,
            [
                (0, 0, []),
                (0, 5, [5, 1, 2, 3]),
                (0, 4, [2, 3]),
                (1, 5, [5, 1, 2]),
                (2, 4, [2, 3]),
            ],
        )
        # pedantically check the Tree methods on the second tree
        tst = ts.trees()
        t = next(tst)
        t = next(tst)
        assert t.branch_length(1) == 0.4
        assert not t.is_internal(0)
        assert t.is_leaf(0)
        assert not t.is_sample(0)
        assert not t.is_internal(1)
        assert t.is_leaf(1)
        assert t.is_sample(1)
        assert t.is_internal(5)
        assert not t.is_leaf(5)
        assert t.is_sample(5)
        assert t.is_internal(4)
        assert not t.is_leaf(4)
        assert not t.is_sample(4)
        assert t.root == 8
        assert t.mrca(0, 1) == 5
        assert t.sample_size == 4


class TestBadTrees:
    """
    Tests for bad tree sequence topologies that can only be detected when we
    try to create trees.
    """

    def test_simplest_contradictory_children(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           1
        3       0           2
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.0     1.0     2       0
        0.0     1.0     3       0
        """
        )
        with pytest.raises(_tskit.LibraryError):
            tskit.load_text(nodes=nodes, edges=edges, strict=False)

    def test_partial_overlap_contradictory_children(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           1
        3       0           2
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.0     1.0     2       0,1
        0.5     1.0     3       0
        """
        )
        with pytest.raises(_tskit.LibraryError):
            tskit.load_text(nodes=nodes, edges=edges, strict=False)


class TestCoiteration:
    """
    Test ability to iterate over multiple (currently 2) tree sequences simultaneously
    """

    def test_identical_ts(self):
        ts = msprime.simulate(4, recombination_rate=1, random_seed=123)
        assert ts.num_trees > 1
        total_iterations = 0
        for tree, (_, t1, t2) in zip(ts.trees(), ts.coiterate(ts)):
            total_iterations += 1
            assert tree == t1 == t2
        assert ts.num_trees == total_iterations

    def test_intervals(self):
        ts1 = msprime.simulate(4, recombination_rate=1, random_seed=1)
        assert ts1.num_trees > 1
        one_tree_ts = msprime.simulate(5, random_seed=2)
        multi_tree_ts = msprime.simulate(5, recombination_rate=1, random_seed=2)
        assert multi_tree_ts.num_trees > 1
        for ts2 in (one_tree_ts, multi_tree_ts):
            bp1 = set(ts1.breakpoints())
            bp2 = set(ts2.breakpoints())
            assert bp1 != bp2
            breaks = set()
            for interval, t1, t2 in ts1.coiterate(ts2):
                assert set(interval) <= set(t1.interval) | set(t2.interval)
                breaks.add(interval.left)
                breaks.add(interval.right)
                assert t1.tree_sequence == ts1
                assert t2.tree_sequence == ts2
            assert breaks == bp1 | bp2

    def test_simple_ts(self):
        nodes = """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       0           1
        4       0           2
        """
        edges1 = """\
        left    right   parent  child
        0       0.2       3       0,1
        0       0.2       4       2,3
        0.2     1         3       2,1
        0.2     1         4       0,3
        """
        edges2 = """\
        left    right   parent  child
        0       0.8       3       2,1
        0       0.8       4       0,3
        0.8     1         3       0,1
        0.8     1         4       2,3
        """
        ts1 = tskit.load_text(io.StringIO(nodes), io.StringIO(edges1), strict=False)
        ts2 = tskit.load_text(io.StringIO(nodes), io.StringIO(edges2), strict=False)
        coiterator = ts1.coiterate(ts2)
        interval, tree1, tree2 = next(coiterator)
        assert interval.left == 0
        assert interval.right == 0.2
        assert tree1 == ts1.at_index(0)
        assert tree2 == ts2.at_index(0)
        interval, tree1, tree2 = next(coiterator)
        assert interval.left == 0.2
        assert interval.right == 0.8
        assert tree1 == ts1.at_index(1)
        assert tree2 == ts2.at_index(0)
        interval, tree1, tree2 = next(coiterator)
        assert interval.left == 0.8
        assert interval.right == 1
        assert tree1 == ts1.at_index(1)
        assert tree2 == ts2.at_index(1)

    def test_nonequal_lengths(self):
        ts1 = msprime.simulate(4, random_seed=1, length=2)
        ts2 = msprime.simulate(4, random_seed=1)
        with pytest.raises(ValueError, match="equal sequence length"):
            next(ts1.coiterate(ts2))

    def test_kwargs(self):
        ts = msprime.simulate(4, recombination_rate=1, random_seed=123)
        for _, t1, t2 in ts.coiterate(ts):
            assert t1.num_tracked_samples() == t2.num_tracked_samples() == 0
        for _, t1, t2 in ts.coiterate(ts, tracked_samples=ts.samples()):
            assert t1.num_tracked_samples() == t2.num_tracked_samples() == 4


def do_simplify(
    ts,
    samples=None,
    compare_lib=True,
    filter_sites=True,
    filter_populations=True,
    filter_individuals=True,
    filter_nodes=True,
    keep_unary=False,
    keep_input_roots=False,
    update_sample_flags=True,
):
    """
    Runs the Python test implementation of simplify.
    """
    if samples is None:
        samples = ts.samples()
    s = tests.Simplifier(
        ts,
        samples,
        filter_sites=filter_sites,
        filter_populations=filter_populations,
        filter_individuals=filter_individuals,
        filter_nodes=filter_nodes,
        keep_unary=keep_unary,
        keep_input_roots=keep_input_roots,
        update_sample_flags=update_sample_flags,
    )
    new_ts, node_map = s.simplify()
    if compare_lib:
        sts, lib_node_map1 = ts.simplify(
            samples,
            filter_sites=filter_sites,
            filter_individuals=filter_individuals,
            filter_populations=filter_populations,
            filter_nodes=filter_nodes,
            update_sample_flags=update_sample_flags,
            keep_unary=keep_unary,
            keep_input_roots=keep_input_roots,
            map_nodes=True,
        )
        lib_tables1 = sts.dump_tables()

        py_tables = new_ts.dump_tables()
        py_tables.assert_equals(lib_tables1, ignore_provenance=True)
        assert all(node_map == lib_node_map1)
    return new_ts, node_map


class SimplifyTestBase:
    """
    Base class for simplify tests.
    """


class TestSimplify(SimplifyTestBase):
    """
    Tests that the implementations of simplify() do what they are supposed to.
    """

    random_seed = 23
    #
    #          8
    #         / \
    #        /   \
    #       /     \
    #      7       \
    #     / \       6
    #    /   5     / \
    #   /   / \   /   \
    #  4   0   1 2     3
    small_tree_ex_nodes = """\
    id      is_sample   population      time
    0       1       0               0.00000000000000
    1       1       0               0.00000000000000
    2       1       0               0.00000000000000
    3       1       0               0.00000000000000
    4       1       0               0.00000000000000
    5       0       0               0.14567111023387
    6       0       0               0.21385545626353
    7       0       0               0.43508024345063
    8       0       0               1.60156352971203
    """
    small_tree_ex_edges = """\
    id      left            right           parent  child
    0       0.00000000      1.00000000      5       0,1
    1       0.00000000      1.00000000      6       2,3
    2       0.00000000      1.00000000      7       4,5
    3       0.00000000      1.00000000      8       6,7
    """

    def verify_no_samples(self, ts, keep_unary=False):
        """
        Zero out the flags column and verify that we get back the correct
        tree sequence when we run simplify.
        """
        t1 = ts.dump_tables()
        t1.nodes.flags = np.zeros_like(t1.nodes.flags)
        ts1, node_map1 = do_simplify(ts, samples=ts.samples(), keep_unary=keep_unary)
        t1 = ts1.dump_tables()
        ts2, node_map2 = do_simplify(ts, keep_unary=keep_unary)
        t2 = ts2.dump_tables()
        t1.assert_equals(t2)

    def verify_single_childified(self, ts, keep_unary=False):
        """
        Modify the specified tree sequence so that it has lots of unary
        nodes. Run simplify and verify we get the same tree sequence back
        if keep_unary is False. If keep_unary is True, the simplication
        won't do anything to the original treeSequence.
        """
        ts_single = tsutil.single_childify(ts)

        tss, node_map = do_simplify(ts_single, keep_unary=keep_unary)
        # All original nodes should still be present.
        for u in range(ts.num_samples):
            assert u == node_map[u]
        # All introduced nodes should be mapped to null.
        for u in range(ts.num_samples, ts_single.num_samples):
            assert node_map[u] == tskit.NULL
        t1 = ts.dump_tables()
        t2 = tss.dump_tables()
        t3 = ts_single.dump_tables()
        if keep_unary:
            assert set(t3.nodes.time) == set(t2.nodes.time)
            assert len(t3.edges) == len(t2.edges)
            assert t3.sites == t2.sites
            assert len(t3.mutations) == len(t2.mutations)
        else:
            assert t1.nodes == t2.nodes
            assert t1.edges == t2.edges
            assert t1.sites == t2.sites
            assert t1.mutations == t2.mutations

    def verify_multiroot_internal_samples(self, ts, keep_unary=False):
        ts_multiroot = ts.decapitate(np.max(ts.tables.nodes.time) / 2)
        ts1 = tsutil.jiggle_samples(ts_multiroot)
        ts2, node_map = do_simplify(ts1, keep_unary=keep_unary)
        assert ts1.num_trees >= ts2.num_trees
        trees2 = ts2.trees()
        t2 = next(trees2)
        for t1 in ts1.trees():
            assert t2.interval.left <= t1.interval.left
            assert t2.interval.right >= t1.interval.right
            pairs = itertools.combinations(ts1.samples(), 2)
            for pair in pairs:
                mapped_pair = [node_map[u] for u in pair]
                mrca1 = t1.get_mrca(*pair)
                mrca2 = t2.get_mrca(*mapped_pair)
                if mrca1 == tskit.NULL:
                    assert mrca2 == tskit.NULL
                else:
                    assert node_map[mrca1] == mrca2
            if t2.interval.right == t1.interval.right:
                t2 = next(trees2, None)

    def test_single_tree(self):
        ts = msprime.simulate(10, random_seed=self.random_seed)
        self.verify_no_samples(ts)
        self.verify_single_childified(ts)
        self.verify_multiroot_internal_samples(ts)
        # Now with keep_unary=True.
        self.verify_no_samples(ts, keep_unary=True)
        self.verify_single_childified(ts, keep_unary=True)
        self.verify_multiroot_internal_samples(ts, keep_unary=True)

    def test_single_tree_mutations(self):
        ts = msprime.simulate(10, mutation_rate=1, random_seed=self.random_seed)
        assert ts.num_sites > 1
        do_simplify(ts)
        self.verify_single_childified(ts)
        # Also with keep_unary == True.
        do_simplify(ts, keep_unary=True)
        self.verify_single_childified(ts, keep_unary=True)

    def test_many_trees_mutations(self):
        ts = msprime.simulate(
            10, recombination_rate=1, mutation_rate=10, random_seed=self.random_seed
        )
        assert ts.num_trees > 2
        assert ts.num_sites > 2
        self.verify_no_samples(ts)
        do_simplify(ts)
        self.verify_single_childified(ts)
        # Also with keep_unary == True.
        do_simplify(ts, keep_unary=True)
        self.verify_single_childified(ts, keep_unary=True)

    def test_many_trees(self):
        ts = msprime.simulate(5, recombination_rate=4, random_seed=self.random_seed)
        assert ts.num_trees > 2
        self.verify_no_samples(ts)
        self.verify_single_childified(ts)
        self.verify_multiroot_internal_samples(ts)
        # Also with keep_unary == True.
        self.verify_no_samples(ts, keep_unary=True)
        self.verify_single_childified(ts, keep_unary=True)
        self.verify_multiroot_internal_samples(ts, keep_unary=True)

    def test_small_tree_internal_samples(self):
        ts = tskit.load_text(
            nodes=io.StringIO(self.small_tree_ex_nodes),
            edges=io.StringIO(self.small_tree_ex_edges),
            strict=False,
        )
        tables = ts.dump_tables()
        nodes = tables.nodes
        flags = nodes.flags
        # The parent of samples 0 and 1 is 5. Change this to an internal sample
        # and set 0 and 1 to be unsampled.
        flags[0] = 0
        flags[0] = 0
        flags[5] = tskit.NODE_IS_SAMPLE
        nodes.flags = flags
        ts = tables.tree_sequence()
        assert ts.sample_size == 5
        tss, node_map = do_simplify(ts, [3, 5])
        assert node_map[3] == 0
        assert node_map[5] == 1
        assert tss.num_nodes == 3
        assert tss.num_edges == 2
        self.verify_no_samples(ts)
        # with keep_unary == True
        tss, node_map = do_simplify(ts, [3, 5], keep_unary=True)
        assert node_map[3] == 0
        assert node_map[5] == 1
        assert tss.num_nodes == 5
        assert tss.num_edges == 4
        self.verify_no_samples(ts, keep_unary=True)

    def test_small_tree_linear_samples(self):
        ts = tskit.load_text(
            nodes=io.StringIO(self.small_tree_ex_nodes),
            edges=io.StringIO(self.small_tree_ex_edges),
            strict=False,
        )
        tables = ts.dump_tables()
        nodes = tables.nodes
        flags = nodes.flags
        # 7 is above 0. These are the only two samples
        flags[:] = 0
        flags[0] = tskit.NODE_IS_SAMPLE
        flags[7] = tskit.NODE_IS_SAMPLE
        nodes.flags = flags
        ts = tables.tree_sequence()
        assert ts.sample_size == 2
        tss, node_map = do_simplify(ts, [0, 7])
        assert node_map[0] == 0
        assert node_map[7] == 1
        assert tss.num_nodes == 2
        assert tss.num_edges == 1
        t = next(tss.trees())
        assert t.parent_dict == {0: 1}
        # with keep_unary == True
        tss, node_map = do_simplify(ts, [0, 7], keep_unary=True)
        assert node_map[0] == 0
        assert node_map[7] == 1
        assert tss.num_nodes == 4
        assert tss.num_edges == 3
        t = next(tss.trees())

    def test_small_tree_internal_and_external_samples(self):
        ts = tskit.load_text(
            nodes=io.StringIO(self.small_tree_ex_nodes),
            edges=io.StringIO(self.small_tree_ex_edges),
            strict=False,
        )
        tables = ts.dump_tables()
        nodes = tables.nodes
        flags = nodes.flags
        # 7 is above 0 and 1.
        flags[:] = 0
        flags[0] = tskit.NODE_IS_SAMPLE
        flags[1] = tskit.NODE_IS_SAMPLE
        flags[7] = tskit.NODE_IS_SAMPLE
        nodes.flags = flags
        ts = tables.tree_sequence()
        assert ts.sample_size == 3
        tss, node_map = do_simplify(ts, [0, 1, 7])
        assert node_map[0] == 0
        assert node_map[1] == 1
        assert node_map[7] == 2
        assert tss.num_nodes == 4
        assert tss.num_edges == 3
        t = next(tss.trees())
        assert t.parent_dict == {0: 3, 1: 3, 3: 2}
        # with keep_unary == True
        tss, node_map = do_simplify(ts, [0, 1, 7], keep_unary=True)
        assert node_map[0] == 0
        assert node_map[1] == 1
        assert node_map[7] == 2
        assert tss.num_nodes == 5
        assert tss.num_edges == 4
        t = next(tss.trees())
        assert t.parent_dict == {0: 3, 1: 3, 3: 2, 2: 4}

    def test_small_tree_mutations(self):
        ts = tskit.load_text(
            nodes=io.StringIO(self.small_tree_ex_nodes),
            edges=io.StringIO(self.small_tree_ex_edges),
            strict=False,
        )
        tables = ts.dump_tables()
        # Add some simple mutations here above the nodes we're keeping.
        tables.sites.add_row(position=0.25, ancestral_state="0")
        tables.sites.add_row(position=0.5, ancestral_state="0")
        tables.sites.add_row(position=0.75, ancestral_state="0")
        tables.sites.add_row(position=0.8, ancestral_state="0")
        tables.mutations.add_row(site=0, node=0, derived_state="1")
        tables.mutations.add_row(site=1, node=2, derived_state="1")
        tables.mutations.add_row(site=2, node=7, derived_state="1")
        tables.mutations.add_row(site=3, node=0, derived_state="1")
        ts = tables.tree_sequence()
        assert ts.num_sites == 4
        assert ts.num_mutations == 4
        for keep in [True, False]:
            tss = do_simplify(ts, [0, 2], keep_unary=keep)[0]
            assert tss.sample_size == 2
            assert tss.num_mutations == 4
            assert list(tss.haplotypes()) == ["1011", "0100"]

    def test_small_tree_filter_zero_mutations(self):
        ts = tskit.load_text(
            nodes=io.StringIO(self.small_tree_ex_nodes),
            edges=io.StringIO(self.small_tree_ex_edges),
            strict=False,
        )
        ts = tsutil.insert_branch_sites(ts)
        assert ts.num_sites == 8
        assert ts.num_mutations == 8
        for keep in [True, False]:
            tss, _ = do_simplify(ts, [4, 0, 1], filter_sites=True, keep_unary=keep)
            assert tss.num_sites == 5
            assert tss.num_mutations == 5
            tss, _ = do_simplify(ts, [4, 0, 1], filter_sites=False, keep_unary=keep)
            assert tss.num_sites == 8
            assert tss.num_mutations == 5

    def test_small_tree_fixed_sites(self):
        ts = tskit.load_text(
            nodes=io.StringIO(self.small_tree_ex_nodes),
            edges=io.StringIO(self.small_tree_ex_edges),
            strict=False,
        )
        tables = ts.dump_tables()
        # Add some simple mutations that will be fixed after simplify
        tables.sites.add_row(position=0.25, ancestral_state="0")
        tables.sites.add_row(position=0.5, ancestral_state="0")
        tables.sites.add_row(position=0.75, ancestral_state="0")
        tables.mutations.add_row(site=0, node=2, derived_state="1")
        tables.mutations.add_row(site=1, node=3, derived_state="1")
        tables.mutations.add_row(site=2, node=6, derived_state="1")
        ts = tables.tree_sequence()
        assert ts.num_sites == 3
        assert ts.num_mutations == 3
        for keep in [True, False]:
            tss, _ = do_simplify(ts, [4, 1], keep_unary=keep)
            assert tss.sample_size == 2
            assert tss.num_mutations == 0
            assert list(tss.haplotypes()) == ["", ""]

    def test_small_tree_mutations_over_root(self):
        ts = tskit.load_text(
            nodes=io.StringIO(self.small_tree_ex_nodes),
            edges=io.StringIO(self.small_tree_ex_edges),
            strict=False,
        )
        tables = ts.dump_tables()
        tables.sites.add_row(position=0.25, ancestral_state="0")
        tables.mutations.add_row(site=0, node=8, derived_state="1")
        ts = tables.tree_sequence()
        assert ts.num_sites == 1
        assert ts.num_mutations == 1
        for keep_unary, filter_sites in itertools.product([True, False], repeat=2):
            tss, _ = do_simplify(
                ts, [0, 1], filter_sites=filter_sites, keep_unary=keep_unary
            )
            assert tss.num_sites == 1
            assert tss.num_mutations == 1

    def test_small_tree_recurrent_mutations(self):
        ts = tskit.load_text(
            nodes=io.StringIO(self.small_tree_ex_nodes),
            edges=io.StringIO(self.small_tree_ex_edges),
            strict=False,
        )
        tables = ts.dump_tables()
        # Add recurrent mutation on the root branches
        tables.sites.add_row(position=0.25, ancestral_state="0")
        tables.mutations.add_row(site=0, node=6, derived_state="1")
        tables.mutations.add_row(site=0, node=7, derived_state="1")
        ts = tables.tree_sequence()
        assert ts.num_sites == 1
        assert ts.num_mutations == 2
        for keep in [True, False]:
            tss = do_simplify(ts, [4, 3], keep_unary=keep)[0]
            assert tss.sample_size == 2
            assert tss.num_sites == 1
            assert tss.num_mutations == 2
            assert list(tss.haplotypes()) == ["1", "1"]

    def test_small_tree_back_mutations(self):
        ts = tskit.load_text(
            nodes=io.StringIO(self.small_tree_ex_nodes),
            edges=io.StringIO(self.small_tree_ex_edges),
            strict=False,
        )
        tables = ts.dump_tables()
        # Add a chain of mutations
        tables.sites.add_row(position=0.25, ancestral_state="0")
        tables.mutations.add_row(site=0, node=7, derived_state="1")
        tables.mutations.add_row(site=0, node=5, derived_state="0")
        tables.mutations.add_row(site=0, node=1, derived_state="1")
        ts = tables.tree_sequence()
        assert ts.num_sites == 1
        assert ts.num_mutations == 3
        assert list(ts.haplotypes()) == ["0", "1", "0", "0", "1"]
        # First check if we simplify for all samples and keep original state.
        for keep in [True, False]:
            tss = do_simplify(ts, [0, 1, 2, 3, 4], keep_unary=keep)[0]
            assert tss.sample_size == 5
            assert tss.num_sites == 1
            assert tss.num_mutations == 3
            assert list(tss.haplotypes()) == ["0", "1", "0", "0", "1"]

        # The ancestral state above 5 should be 0.
        for keep in [True, False]:
            tss = do_simplify(ts, [0, 1], keep_unary=keep)[0]
            assert tss.sample_size == 2
            assert tss.num_sites == 1
            assert tss.num_mutations == 3
            assert list(tss.haplotypes()) == ["0", "1"]

        # The ancestral state above 7 should be 1.
        for keep in [True, False]:
            tss = do_simplify(ts, [4, 0, 1], keep_unary=keep)[0]
            assert tss.sample_size == 3
            assert tss.num_sites == 1
            assert tss.num_mutations == 3
            assert list(tss.haplotypes()) == ["1", "0", "1"]

    def test_overlapping_unary_edges(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           1
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       2       2       0
        1       3       2       1
        """
        )
        ts = tskit.load_text(nodes, edges, strict=False)
        assert ts.sample_size == 2
        assert ts.num_trees == 3
        assert ts.sequence_length == 3
        for keep in [True, False]:
            tss, node_map = do_simplify(ts, samples=[0, 1, 2], keep_unary=keep)
            assert list(node_map) == [0, 1, 2]
            trees = [{0: 2}, {0: 2, 1: 2}, {1: 2}]
            for t in tss.trees():
                assert t.parent_dict == trees[t.index]

    def test_overlapping_unary_edges_internal_samples(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           1
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       2       2       0
        1       3       2       1
        """
        )
        ts = tskit.load_text(nodes, edges, strict=False)
        assert ts.sample_size == 3
        assert ts.num_trees == 3
        trees = [{0: 2}, {0: 2, 1: 2}, {1: 2}]
        for t in ts.trees():
            assert t.parent_dict == trees[t.index]
        tss, node_map = do_simplify(ts)
        assert list(node_map) == [0, 1, 2]

    def test_isolated_samples(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           1
        2       1           2
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        """
        )
        ts = tskit.load_text(nodes, edges, sequence_length=1, strict=False)
        assert ts.num_samples == 3
        assert ts.num_trees == 1
        assert ts.num_nodes == 3
        for keep in [True, False]:
            tss, node_map = do_simplify(ts, keep_unary=keep)
            assert ts.tables.nodes == tss.tables.nodes
            assert ts.tables.edges == tss.tables.edges
            assert list(node_map) == [0, 1, 2]

    def test_internal_samples(self):
        nodes = io.StringIO(
            """\
        id      is_sample   population      time
        0       1       -1              1.00000000000000
        1       0       -1              1.00000000000000
        2       1       -1              1.00000000000000
        3       0       -1              1.31203521181726
        4       0       -1              2.26776380586006
        5       1       -1              0.00000000000000
        6       0       -1              0.50000000000000
        7       0       -1              1.50000000000000

        """
        )
        edges = io.StringIO(
            """\
        id      left            right           parent  child
        0       0.62185118      1.00000000      1       6
        1       0.00000000      0.62185118      2       6
        2       0.00000000      1.00000000      3       0,2
        3       0.00000000      1.00000000      4       7,3
        4       0.00000000      1.00000000      6       5
        5       0.00000000      1.00000000      7       1
        """
        )

        ts = tskit.load_text(nodes, edges, strict=False)
        tss, node_map = do_simplify(ts, [5, 2, 0])
        assert node_map[0] == 2
        assert node_map[1] == -1
        assert node_map[2] == 1
        assert node_map[3] == 3
        assert node_map[4] == 4
        assert node_map[5] == 0
        assert node_map[6] == -1
        assert node_map[7] == -1
        assert tss.sample_size == 3
        assert tss.num_trees == 2
        trees = [{0: 1, 1: 3, 2: 3}, {0: 4, 1: 3, 2: 3, 3: 4}]
        for t in tss.trees():
            assert t.parent_dict == trees[t.index]
        # with keep_unary == True
        tss, node_map = do_simplify(ts, [5, 2, 0], keep_unary=True)
        assert node_map[0] == 2
        assert node_map[1] == 4
        assert node_map[2] == 1
        assert node_map[3] == 5
        assert node_map[4] == 7
        assert node_map[5] == 0
        assert node_map[6] == 3
        assert node_map[7] == 6
        assert tss.sample_size == 3
        assert tss.num_trees == 2
        trees = [
            {0: 3, 1: 5, 2: 5, 3: 1, 5: 7},
            {0: 3, 1: 5, 2: 5, 3: 4, 4: 6, 5: 7, 6: 7},
        ]
        for t in tss.trees():
            assert t.parent_dict == trees[t.index]

    def test_many_mutations_over_single_sample_ancestral_state(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       0           1
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       1       0
        """
        )
        sites = io.StringIO(
            """\
        position    ancestral_state
        0           0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    derived_state   parent
        0       0       1               -1
        0       0       0               0
        """
        )
        ts = tskit.load_text(
            nodes, edges, sites=sites, mutations=mutations, strict=False
        )
        assert ts.sample_size == 1
        assert ts.num_trees == 1
        assert ts.num_sites == 1
        assert ts.num_mutations == 2
        for keep in [True, False]:
            tss, node_map = do_simplify(ts, keep_unary=keep)
            assert tss.num_sites == 1
            assert tss.num_mutations == 2
            assert list(tss.haplotypes(isolated_as_missing=False)) == ["0"]

    def test_many_mutations_over_single_sample_derived_state(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       0           1
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0       1       1       0
        """
        )
        sites = io.StringIO(
            """\
        position    ancestral_state
        0           0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    derived_state   parent
        0       0       1               -1
        0       0       0               0
        0       0       1               1
        """
        )
        ts = tskit.load_text(
            nodes, edges, sites=sites, mutations=mutations, strict=False
        )
        assert ts.sample_size == 1
        assert ts.num_trees == 1
        assert ts.num_sites == 1
        assert ts.num_mutations == 3
        for keep in [True, False]:
            tss, node_map = do_simplify(ts, keep_unary=keep)
            assert tss.num_sites == 1
            assert tss.num_mutations == 3
            assert list(tss.haplotypes(isolated_as_missing=False)) == ["1"]

    def test_many_trees_filter_zero_mutations(self):
        ts = msprime.simulate(5, recombination_rate=1, random_seed=10)
        assert ts.num_trees > 3
        ts = tsutil.insert_branch_sites(ts)
        assert ts.num_sites == ts.num_mutations
        assert ts.num_sites > ts.num_trees
        for keep in [True, False]:
            for filter_sites in [True, False]:
                tss, _ = do_simplify(
                    ts, samples=None, filter_sites=filter_sites, keep_unary=keep
                )
                assert ts.num_sites == tss.num_sites
                assert ts.num_mutations == tss.num_mutations

    def test_many_trees_filter_zero_multichar_mutations(self):
        ts = msprime.simulate(5, recombination_rate=1, random_seed=10)
        assert ts.num_trees > 3
        ts = tsutil.insert_multichar_mutations(ts)
        assert ts.num_sites == ts.num_trees
        assert ts.num_mutations == ts.num_trees
        for keep in [True, False]:
            for filter_sites in [True, False]:
                tss, _ = do_simplify(
                    ts, samples=None, filter_sites=filter_sites, keep_unary=keep
                )
                assert ts.num_sites == tss.num_sites
                assert ts.num_mutations == tss.num_mutations

    def test_simple_population_filter(self):
        ts = msprime.simulate(10, random_seed=2)
        tables = ts.dump_tables()
        tables.populations.add_row(metadata=b"unreferenced")
        assert len(tables.populations) == 2
        for keep in [True, False]:
            tss, _ = do_simplify(
                tables.tree_sequence(), filter_populations=True, keep_unary=keep
            )
            assert tss.num_populations == 1
            tss, _ = do_simplify(
                tables.tree_sequence(), filter_populations=False, keep_unary=keep
            )
            assert tss.num_populations == 2

    def test_interleaved_populations_filter(self):
        ts = msprime.simulate(
            population_configurations=[
                msprime.PopulationConfiguration(),
                msprime.PopulationConfiguration(10),
                msprime.PopulationConfiguration(),
                msprime.PopulationConfiguration(),
            ],
            random_seed=2,
        )
        assert ts.num_populations == 4
        tables = ts.dump_tables()
        # Edit the populations so we can identify the rows.
        tables.populations.clear()
        for j in range(4):
            tables.populations.add_row(metadata=bytes([j]))
        ts = tables.tree_sequence()
        id_map = np.array([-1, 0, -1, -1], dtype=np.int32)
        for keep in [True, False]:
            tss, _ = do_simplify(ts, filter_populations=True, keep_unary=keep)
            assert tss.num_populations == 1
            population = tss.population(0)
            assert population.metadata == bytes([1])
            assert np.array_equal(
                id_map[ts.tables.nodes.population], tss.tables.nodes.population
            )
            tss, _ = do_simplify(ts, filter_populations=False, keep_unary=keep)
            assert tss.num_populations == 4

    def test_removed_node_population_filter(self):
        tables = tskit.TableCollection(1)
        tables.populations.add_row(metadata=bytes(0))
        tables.populations.add_row(metadata=bytes(1))
        tables.populations.add_row(metadata=bytes(2))
        tables.nodes.add_row(flags=1, population=0)
        # Because flags=0 here, this node will be simplified out and the node
        # will disappear.
        tables.nodes.add_row(flags=0, population=1)
        tables.nodes.add_row(flags=1, population=2)
        for keep in [True, False]:
            tss, _ = do_simplify(
                tables.tree_sequence(), filter_populations=True, keep_unary=keep
            )
            assert tss.num_nodes == 2
            assert tss.num_populations == 2
            assert tss.population(0).metadata == bytes(0)
            assert tss.population(1).metadata == bytes(2)
            assert tss.node(0).population == 0
            assert tss.node(1).population == 1

            tss, _ = do_simplify(
                tables.tree_sequence(), filter_populations=False, keep_unary=keep
            )
            assert tss.tables.populations == tables.populations

    def test_simple_individual_filter(self):
        tables = tskit.TableCollection(1)
        tables.individuals.add_row(flags=0)
        tables.individuals.add_row(flags=1)
        tables.nodes.add_row(flags=1, individual=0)
        tables.nodes.add_row(flags=1, individual=0)
        for keep in [True, False]:
            tss, _ = do_simplify(
                tables.tree_sequence(), filter_individuals=True, keep_unary=keep
            )
            assert tss.num_nodes == 2
            assert tss.num_individuals == 1
            assert tss.individual(0).flags == 0

        tss, _ = do_simplify(tables.tree_sequence(), filter_individuals=False)
        assert tss.tables.individuals == tables.individuals

    def test_interleaved_individual_filter(self):
        tables = tskit.TableCollection(1)
        tables.individuals.add_row(flags=0)
        tables.individuals.add_row(flags=1)
        tables.individuals.add_row(flags=2)
        tables.nodes.add_row(flags=1, individual=1)
        tables.nodes.add_row(flags=1, individual=-1)
        tables.nodes.add_row(flags=1, individual=1)
        for keep in [True, False]:
            tss, _ = do_simplify(
                tables.tree_sequence(), filter_individuals=True, keep_unary=keep
            )
            assert tss.num_nodes == 3
            assert tss.num_individuals == 1
            assert tss.individual(0).flags == 1

            tss, _ = do_simplify(
                tables.tree_sequence(), filter_individuals=False, keep_unary=keep
            )
            assert tss.tables.individuals == tables.individuals

    def test_removed_node_individual_filter(self):
        tables = tskit.TableCollection(1)
        tables.individuals.add_row(flags=0)
        tables.individuals.add_row(flags=1)
        tables.individuals.add_row(flags=2)
        tables.nodes.add_row(flags=1, individual=0)
        # Because flags=0 here, this node will be simplified out and the node
        # will disappear.
        tables.nodes.add_row(flags=0, individual=1)
        tables.nodes.add_row(flags=1, individual=2)
        for keep in [True, False]:
            tss, _ = do_simplify(
                tables.tree_sequence(), filter_individuals=True, keep_unary=keep
            )
            assert tss.num_nodes == 2
            assert tss.num_individuals == 2
            assert tss.individual(0).flags == 0
            assert tss.individual(1).flags == 2
            assert tss.node(0).individual == 0
            assert tss.node(1).individual == 1

            tss, _ = do_simplify(
                tables.tree_sequence(), filter_individuals=False, keep_unary=keep
            )
            assert tss.tables.individuals == tables.individuals

    def verify_simplify_haplotypes(self, ts, samples, keep_unary=False):
        sub_ts, node_map = do_simplify(
            ts, samples, filter_sites=False, keep_unary=keep_unary
        )
        assert ts.num_sites == sub_ts.num_sites
        sub_haplotypes = list(sub_ts.haplotypes(isolated_as_missing=False))
        all_samples = list(ts.samples())
        k = 0
        for j, h in enumerate(ts.haplotypes(isolated_as_missing=False)):
            if k == len(samples):
                break
            if samples[k] == all_samples[j]:
                assert h == sub_haplotypes[k]
                k += 1

    def test_single_tree_recurrent_mutations(self):
        ts = msprime.simulate(6, random_seed=10)
        for mutations_per_branch in [1, 2, 3]:
            ts = tsutil.insert_branch_mutations(ts, mutations_per_branch)
            for num_samples in range(1, ts.num_samples):
                for samples in itertools.combinations(ts.samples(), num_samples):
                    for keep in [True, False]:
                        self.verify_simplify_haplotypes(ts, samples, keep_unary=keep)

    def test_many_trees_recurrent_mutations(self):
        ts = msprime.simulate(5, recombination_rate=1, random_seed=10)
        assert ts.num_trees > 3
        for mutations_per_branch in [1, 2, 3]:
            ts = tsutil.insert_branch_mutations(ts, mutations_per_branch)
            for num_samples in range(1, ts.num_samples):
                for samples in itertools.combinations(ts.samples(), num_samples):
                    for keep in [True, False]:
                        self.verify_simplify_haplotypes(ts, samples, keep_unary=keep)

    def test_single_multiroot_tree_recurrent_mutations(self):
        ts = msprime.simulate(6, random_seed=10)
        ts = ts.decapitate(np.max(ts.tables.nodes.time) / 2)
        for mutations_per_branch in [1, 2, 3]:
            ts = tsutil.insert_branch_mutations(ts, mutations_per_branch)
            for num_samples in range(1, ts.num_samples):
                for samples in itertools.combinations(ts.samples(), num_samples):
                    for keep in [True, False]:
                        self.verify_simplify_haplotypes(ts, samples, keep_unary=keep)

    def test_many_multiroot_trees_recurrent_mutations(self):
        ts = msprime.simulate(7, recombination_rate=1, random_seed=10)
        assert ts.num_trees > 3
        ts = ts.decapitate(np.max(ts.tables.nodes.time) / 2)
        for mutations_per_branch in [1, 2, 3]:
            ts = tsutil.insert_branch_mutations(ts, mutations_per_branch)
            for num_samples in range(1, ts.num_samples):
                for samples in itertools.combinations(ts.samples(), num_samples):
                    for keep in [True, False]:
                        self.verify_simplify_haplotypes(ts, samples, keep_unary=keep)

    def test_single_tree_recurrent_mutations_internal_samples(self):
        ts = msprime.simulate(6, random_seed=10)
        ts = tsutil.jiggle_samples(ts)
        for mutations_per_branch in [1, 2, 3]:
            ts = tsutil.insert_branch_mutations(ts, mutations_per_branch)
            for num_samples in range(1, ts.num_samples):
                for samples in itertools.combinations(ts.samples(), num_samples):
                    for keep in [True, False]:
                        self.verify_simplify_haplotypes(ts, samples, keep_unary=keep)

    def test_many_trees_recurrent_mutations_internal_samples(self):
        ts = msprime.simulate(5, recombination_rate=1, random_seed=10)
        ts = tsutil.jiggle_samples(ts)
        assert ts.num_trees > 3
        for mutations_per_branch in [1, 2, 3]:
            ts = tsutil.insert_branch_mutations(ts, mutations_per_branch)
            for num_samples in range(1, ts.num_samples):
                for samples in itertools.combinations(ts.samples(), num_samples):
                    for keep in [True, False]:
                        self.verify_simplify_haplotypes(ts, samples, keep_unary=keep)


class TestSimplifyUnreferencedPopulations:
    def example(self):
        tables = tskit.TableCollection(1)
        tables.populations.add_row()
        tables.populations.add_row()
        # No references to population 0
        tables.nodes.add_row(time=0, population=1, flags=1)
        tables.nodes.add_row(time=0, population=1, flags=1)
        tables.nodes.add_row(time=1, population=1, flags=0)
        # Unreference node
        tables.nodes.add_row(time=1, population=1, flags=0)
        tables.edges.add_row(0, 1, parent=2, child=0)
        tables.edges.add_row(0, 1, parent=2, child=1)
        tables.sort()
        return tables

    def test_no_filter_populations(self):
        tables = self.example()
        tables.simplify(filter_populations=False)
        assert len(tables.populations) == 2
        assert len(tables.nodes) == 3
        assert np.all(tables.nodes.population == 1)

    def test_no_filter_populations_nodes(self):
        tables = self.example()
        tables.simplify(filter_populations=False, filter_nodes=False)
        assert len(tables.populations) == 2
        assert len(tables.nodes) == 4
        assert np.all(tables.nodes.population == 1)

    def test_filter_populations_no_filter_nodes(self):
        tables = self.example()
        tables.simplify(filter_populations=True, filter_nodes=False)
        assert len(tables.populations) == 1
        assert len(tables.nodes) == 4
        assert np.all(tables.nodes.population == 0)

    def test_remapped_default(self):
        tables = self.example()
        tables.simplify()
        assert len(tables.populations) == 1
        assert len(tables.nodes) == 3
        assert np.all(tables.nodes.population == 0)


class TestSimplifyUnreferencedIndividuals:
    def example(self):
        tables = tskit.TableCollection(1)
        tables.individuals.add_row()
        tables.individuals.add_row()
        # No references to individual 0
        tables.nodes.add_row(time=0, individual=1, flags=1)
        tables.nodes.add_row(time=0, individual=1, flags=1)
        tables.nodes.add_row(time=1, individual=1, flags=0)
        # Unreference node
        tables.nodes.add_row(time=1, individual=1, flags=0)
        tables.edges.add_row(0, 1, parent=2, child=0)
        tables.edges.add_row(0, 1, parent=2, child=1)
        tables.sort()
        return tables

    def test_no_filter_individuals(self):
        tables = self.example()
        tables.simplify(filter_individuals=False)
        assert len(tables.individuals) == 2
        assert len(tables.nodes) == 3
        assert np.all(tables.nodes.individual == 1)

    def test_no_filter_individuals_nodes(self):
        tables = self.example()
        tables.simplify(filter_individuals=False, filter_nodes=False)
        assert len(tables.individuals) == 2
        assert len(tables.nodes) == 4
        assert np.all(tables.nodes.individual == 1)

    def test_filter_individuals_no_filter_nodes(self):
        tables = self.example()
        tables.simplify(filter_individuals=True, filter_nodes=False)
        assert len(tables.individuals) == 1
        assert len(tables.nodes) == 4
        assert np.all(tables.nodes.individual == 0)

    def test_remapped_default(self):
        tables = self.example()
        tables.simplify()
        assert len(tables.individuals) == 1
        assert len(tables.nodes) == 3
        assert np.all(tables.nodes.individual == 0)


class TestSimplifyKeepInputRoots(SimplifyTestBase, ExampleTopologyMixin):
    """
    Tests for the keep_input_roots option to simplify.
    """

    def verify(self, ts):
        # Called by the examples in ExampleTopologyMixin
        samples = ts.samples()
        self.verify_keep_input_roots(ts, samples[:2])
        self.verify_keep_input_roots(ts, samples[:3])
        self.verify_keep_input_roots(ts, samples[:-1])
        self.verify_keep_input_roots(ts, samples)

    def verify_keep_input_roots(self, ts, samples):
        ts = tsutil.insert_unique_metadata(ts, ["individuals"])
        ts_with_roots, node_map = do_simplify(
            ts, samples, keep_input_roots=True, filter_sites=False, compare_lib=True
        )
        new_to_input_map = {
            value: key for key, value in enumerate(node_map) if value != tskit.NULL
        }
        for (left, right), input_tree, tree_with_roots in ts.coiterate(ts_with_roots):
            input_roots = input_tree.roots
            assert len(tree_with_roots.roots) > 0
            for root in tree_with_roots.roots:
                # Check that the roots in the current
                input_root = new_to_input_map[root]
                assert input_root in input_roots
                input_node = ts.node(input_root)
                new_node = ts_with_roots.node(root)
                assert new_node.time == input_node.time
                assert new_node.population == input_node.population
                if new_node.individual == tskit.NULL:
                    assert new_node.individual == input_node.individual
                else:
                    assert (
                        ts_with_roots.individual(new_node.individual).metadata
                        == ts.individual(input_node.individual).metadata
                    )
                assert new_node.metadata == input_node.metadata
                # This should only be marked as a sample if it's an
                # element of the samples list.
                assert new_node.is_sample() == (input_root in samples)
                # Find the MRCA of the samples below this root.
                root_samples = list(tree_with_roots.samples(root))
                mrca = functools.reduce(tree_with_roots.mrca, root_samples)
                if mrca != root:
                    # If the MRCA is not equal to the root, then there should
                    # be a unary branch joining them.
                    assert tree_with_roots.parent(mrca) == root
                    assert tree_with_roots.children(root) == (mrca,)

                    # Any mutations that were on the path from the old MRCA
                    # to the root should be mapped to this node, and any mutations
                    # above the root should still be there.
                    u = new_to_input_map[mrca]
                    root_path = []
                    while u != tskit.NULL:
                        root_path.append(u)
                        u = input_tree.parent(u)
                    input_sites = {
                        site.position: site
                        for site in input_tree.sites()
                        if site.position >= left and site.position < right
                    }
                    new_sites = {
                        site.position: site
                        for site in tree_with_roots.sites()
                        if site.position >= left and site.position < right
                    }
                    assert set(input_sites.keys()) == set(new_sites.keys())
                    positions = input_sites.keys()
                    for position in positions:
                        assert left <= position < right
                        new_site = new_sites[position]
                        # We assume the metadata contains a unique key for each mutation.
                        new_mutations = {
                            mut.metadata: mut for mut in new_site.mutations
                        }
                        # Just make sure the metadata is actually unique.
                        assert len(new_mutations) == len(new_site.mutations)
                        input_site = input_sites[position]
                        for input_mutation in input_site.mutations:
                            if input_mutation.node in root_path:
                                new_node = (
                                    mrca if input_mutation.node != input_root else root
                                )
                                # The same mutation should exist and be mapped to
                                # new_node
                                new_mutation = new_mutations[input_mutation.metadata]
                                # We have turned filter sites off, so sites should
                                # be comparable
                                assert new_mutation.site == input_mutation.site
                                assert (
                                    new_mutation.derived_state
                                    == input_mutation.derived_state
                                )
                                assert new_mutation.node == new_node

        return ts_with_roots

    def test_many_trees(self):
        ts = msprime.simulate(5, recombination_rate=1, random_seed=10)
        assert ts.num_trees > 3
        for num_samples in range(1, ts.num_samples):
            for samples in itertools.combinations(ts.samples(), num_samples):
                self.verify_keep_input_roots(ts, samples)

    def test_many_trees_internal_samples(self):
        ts = msprime.simulate(5, recombination_rate=1, random_seed=10)
        ts = tsutil.jiggle_samples(ts)
        assert ts.num_trees > 3
        for num_samples in range(1, ts.num_samples):
            for samples in itertools.combinations(ts.samples(), num_samples):
                self.verify_keep_input_roots(ts, samples)

    def test_many_multiroot_trees(self):
        ts = msprime.simulate(7, recombination_rate=1, random_seed=10)
        assert ts.num_trees > 3
        ts = ts.decapitate(np.max(ts.tables.nodes.time) / 2)
        for num_samples in range(1, ts.num_samples):
            for samples in itertools.combinations(ts.samples(), num_samples):
                self.verify_keep_input_roots(ts, samples)

    def test_wright_fisher_unsimplified(self):
        num_generations = 10
        tables = wf.wf_sim(10, num_generations, deep_history=False, seed=2)
        tables.sort()
        ts = tables.tree_sequence()
        simplified = self.verify_keep_input_roots(ts, ts.samples())
        roots = set()
        for tree in simplified.trees():
            for root in tree.roots:
                roots.add(root)
                assert tree.time(root) == num_generations
        init_nodes = np.where(simplified.tables.nodes.time == num_generations)[0]
        assert set(init_nodes) == roots

    def test_single_tree_recurrent_mutations(self):
        ts = msprime.simulate(6, random_seed=10)
        for mutations_per_branch in [1, 2, 3]:
            ts = tsutil.insert_branch_mutations(ts, mutations_per_branch)
            for num_samples in range(1, ts.num_samples):
                for samples in itertools.combinations(ts.samples(), num_samples):
                    self.verify_keep_input_roots(ts, samples)

    def test_many_trees_recurrent_mutations(self):
        ts = msprime.simulate(5, recombination_rate=1, random_seed=8)
        assert ts.num_trees > 2
        for mutations_per_branch in [1, 2, 3]:
            ts = tsutil.insert_branch_mutations(ts, mutations_per_branch)
            for num_samples in range(1, ts.num_samples):
                for samples in itertools.combinations(ts.samples(), num_samples):
                    self.verify_keep_input_roots(ts, samples)


class TestSimplifyFilterNodes:
    """
    Tests simplify when nodes are kept in the ts with filter_nodes=False
    """

    def reverse_node_indexes(self, ts):
        tables = ts.dump_tables()
        nodes = tables.nodes
        edges = tables.edges
        mutations = tables.mutations
        nodes.replace_with(nodes[::-1])
        edges.parent = ts.num_nodes - edges.parent - 1
        edges.child = ts.num_nodes - edges.child - 1
        mutations.node = ts.num_nodes - mutations.node - 1
        tables.sort()
        return tables.tree_sequence()

    def verify_nodes_unchanged(self, ts_in, resample_size=None, **kwargs):
        if resample_size is None:
            samples = None
        else:
            np.random.seed(42)
            samples = np.sort(
                np.random.choice(ts_in.num_nodes, resample_size, replace=False)
            )

        for ts in (ts_in, self.reverse_node_indexes(ts_in)):
            filtered, n_map = do_simplify(
                ts, samples=samples, filter_nodes=False, compare_lib=True, **kwargs
            )
            assert np.array_equal(n_map, np.arange(ts.num_nodes, dtype=n_map.dtype))
            referenced_nodes = set(filtered.samples())
            referenced_nodes.update(filtered.edges_parent)
            referenced_nodes.update(filtered.edges_child)
            for n1, n2 in zip(ts.nodes(), filtered.nodes()):
                # Ignore the tskit.NODE_IS_SAMPLE flag which can be changed by simplify
                n1 = n1.replace(flags=n1.flags | tskit.NODE_IS_SAMPLE)
                n2 = n2.replace(flags=n2.flags | tskit.NODE_IS_SAMPLE)
                assert n1 == n2

            # Check that edges are identical to the normal simplify(),
            # with the normal "simplify" having altered IDs
            simplified, node_map = ts.simplify(
                samples=samples, map_nodes=True, **kwargs
            )
            simplified_edges = {e for e in simplified.tables.edges}
            filtered_edges = {
                e.replace(parent=node_map[e.parent], child=node_map[e.child])
                for e in filtered.tables.edges
            }
            assert filtered_edges == simplified_edges

    def test_empty(self):
        ts = tskit.TableCollection(1).tree_sequence()
        self.verify_nodes_unchanged(ts)

    def test_all_samples(self):
        ts = tskit.Tree.generate_comb(5).tree_sequence
        tables = ts.dump_tables()
        flags = tables.nodes.flags
        flags |= tskit.NODE_IS_SAMPLE
        tables.nodes.flags = flags
        ts = tables.tree_sequence()
        assert ts.num_samples == ts.num_nodes
        self.verify_nodes_unchanged(ts)

    @pytest.mark.parametrize("resample_size", [None, 4])
    def test_no_topology(self, resample_size):
        ts = tskit.Tree.generate_comb(5).tree_sequence
        ts = ts.keep_intervals([], simplify=False)
        assert ts.num_nodes > 5  # has unreferenced nodes
        self.verify_nodes_unchanged(ts, resample_size=resample_size)

    @pytest.mark.parametrize("resample_size", [None, 2])
    def test_stick_tree(self, resample_size):
        ts = tskit.Tree.generate_comb(2).tree_sequence
        ts = ts.simplify([0], keep_unary=True)
        assert ts.first().parent(0) != tskit.NULL
        self.verify_nodes_unchanged(ts, resample_size=resample_size)

        # switch to an internal sample
        tables = ts.dump_tables()
        flags = tables.nodes.flags
        flags[0] = 0
        flags[1] = tskit.NODE_IS_SAMPLE
        tables.nodes.flags = flags
        self.verify_nodes_unchanged(tables.tree_sequence(), resample_size=resample_size)

    @pytest.mark.parametrize("resample_size", [None, 4])
    def test_internal_samples(self, resample_size):
        ts = tskit.Tree.generate_comb(4).tree_sequence
        tables = ts.dump_tables()
        flags = tables.nodes.flags
        flags ^= tskit.NODE_IS_SAMPLE
        tables.nodes.flags = flags
        ts = tables.tree_sequence()
        assert np.all(ts.samples() >= ts.num_samples)
        self.verify_nodes_unchanged(ts, resample_size=resample_size)

    @pytest.mark.parametrize("resample_size", [None, 4])
    def test_blank_flanks(self, resample_size):
        ts = tskit.Tree.generate_comb(4).tree_sequence
        ts = ts.keep_intervals([[0.25, 0.75]], simplify=False)
        self.verify_nodes_unchanged(ts, resample_size=resample_size)

    @pytest.mark.parametrize("resample_size", [None, 4])
    def test_multiroot(self, resample_size):
        ts = tskit.Tree.generate_balanced(6).tree_sequence
        ts = ts.decapitate(2.5)
        self.verify_nodes_unchanged(ts, resample_size=resample_size)

    @pytest.mark.parametrize("resample_size", [None, 10])
    def test_with_metadata(self, ts_fixture_for_simplify, resample_size):
        assert ts_fixture_for_simplify.num_nodes > 10
        self.verify_nodes_unchanged(
            ts_fixture_for_simplify, resample_size=resample_size
        )

    @pytest.mark.parametrize("resample_size", [None, 7])
    def test_complex_ts_with_unary(self, resample_size):
        ts = msprime.sim_ancestry(
            3,
            sequence_length=10,
            recombination_rate=1,
            record_full_arg=True,
            random_seed=123,
        )
        assert ts.num_trees > 2
        ts = msprime.sim_mutations(ts, rate=1, random_seed=123)
        # Add some unreferenced nodes
        tables = ts.dump_tables()
        tables.nodes.add_row(flags=0)
        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE)
        ts = tables.tree_sequence()
        self.verify_nodes_unchanged(ts, resample_size=resample_size)

    def test_keeping_unary(self):
        # Test interaction with keeping unary nodes
        n_samples = 6
        ts = tskit.Tree.generate_comb(n_samples).tree_sequence
        num_nodes = ts.num_nodes
        reduced_n_samples = [2, n_samples - 1]  # last sample is most deeply nested
        ts_with_unary = ts.simplify(reduced_n_samples, keep_unary=True)
        assert ts_with_unary.num_nodes == num_nodes - n_samples + len(reduced_n_samples)
        tree = ts_with_unary.first()
        assert any([tree.num_children(u) == 1 for u in tree.nodes()])
        self.verify_nodes_unchanged(ts_with_unary, keep_unary=True)
        self.verify_nodes_unchanged(ts_with_unary, keep_unary=False)

    def test_find_unreferenced_nodes(self):
        # Simple test to show we can find unreferenced nodes easily.
        # 2.00┊    6    ┊
        #     ┊  ┏━┻━┓  ┊
        # 1.00┊  4   5  ┊
        #     ┊ ┏┻┓ ┏┻┓ ┊
        # 0.00┊ 0 1 2 3 ┊
        #     0         1
        ts1 = tskit.Tree.generate_balanced(4).tree_sequence
        ts2, node_map = do_simplify(
            ts1,
            [0, 1, 2],
            filter_nodes=False,
        )
        assert np.array_equal(node_map, np.arange(ts1.num_nodes))
        node_references = np.zeros(ts1.num_nodes, dtype=np.int32)
        node_references[ts2.edges_parent] += 1
        node_references[ts2.edges_child] += 1
        # Simplifying for [0, 1, 2] should remove references to node 3 and 5
        assert list(node_references) == [1, 1, 1, 0, 2, 0, 1]

    def test_mutations_on_removed_branches(self):
        # 2.00┊    6    ┊
        #     ┊  ┏━┻━┓  ┊
        # 1.00┊  4   5  ┊
        #     ┊ ┏┻┓ ┏┻┓ ┊
        # 0.00┊ 0 1 2 3 ┊
        #     0         1
        tables = tskit.Tree.generate_balanced(4).tree_sequence.dump_tables()
        # A mutation on a removed branch should get removed
        tables.sites.add_row(0.5, "A")
        tables.mutations.add_row(0, node=3, derived_state="T")
        ts2, node_map = do_simplify(
            tables.tree_sequence(),
            [0, 1, 2],
            filter_nodes=False,
        )
        assert ts2.num_sites == 0
        assert ts2.num_mutations == 0


class TestSimplifyNoUpdateSampleFlags:
    """
    Tests for simplify when we don't update the sample flags.
    """

    def test_simple_case_filter_nodes(self):
        # 2.00┊    6    ┊
        #     ┊  ┏━┻━┓  ┊
        # 1.00┊  4   5  ┊
        #     ┊ ┏┻┓ ┏┻┓ ┊
        # 0.00┊ 0 1 2 3 ┊
        #     0         1
        ts1 = tskit.Tree.generate_balanced(4).tree_sequence
        ts2, node_map = do_simplify(
            ts1,
            [0, 1, 6],
            update_sample_flags=False,
        )
        # Because we don't retain 2 and 3 here, they don't stay as
        # samples. But, we specified 6 as a sample, so it's coming
        # through where it would ordinarily be dropped.

        # 2.00┊  2  ┊
        #     ┊  ┃  ┊
        # 1.00┊  3  ┊
        #     ┊ ┏┻┓ ┊
        # 0.00┊ 0 1 ┊
        #     0     1
        assert list(ts2.nodes_flags) == [1, 1, 0, 0]
        tree = ts2.first()
        assert list(tree.parent_array) == [3, 3, -1, 2, -1]

    def test_simple_case_no_filter_nodes(self):
        # 2.00┊    6    ┊
        #     ┊  ┏━┻━┓  ┊
        # 1.00┊  4   5  ┊
        #     ┊ ┏┻┓ ┏┻┓ ┊
        # 0.00┊ 0 1 2 3 ┊
        #     0         1
        ts1 = tskit.Tree.generate_balanced(4).tree_sequence
        ts2, node_map = do_simplify(
            ts1,
            [0, 1, 6],
            update_sample_flags=False,
            filter_nodes=False,
        )

        # 2.00┊  6      ┊
        #     ┊  ┃      ┊
        # 1.00┊  4      ┊
        #     ┊ ┏┻┓     ┊
        # 0.00┊ 0 1 2 3 ┊
        #     0         1
        assert list(ts2.nodes_flags) == list(ts1.nodes_flags)
        tree = ts2.first()
        assert list(tree.parent_array) == [4, 4, -1, -1, 6, -1, -1, -1]


class TestMapToAncestors:
    """
    Tests the AncestorMap class.
    """

    random_seed = 13
    #
    #          8
    #         / \
    #        /   \
    #       /     \
    #      7       \
    #     / \       6
    #    /   5     / \
    #   /   / \   /   \
    #  4   0   1 2     3
    nodes = """\
    id      is_sample   population      time
    0       1       0               0.00000000000000
    1       1       0               0.00000000000000
    2       1       0               0.00000000000000
    3       1       0               0.00000000000000
    4       1       0               0.00000000000000
    5       0       0               0.14567111023387
    6       0       0               0.21385545626353
    7       0       0               0.43508024345063
    8       0       0               1.60156352971203
    """
    edges = """\
    id      left            right           parent  child
    0       0.00000000      1.00000000      5       0,1
    1       0.00000000      1.00000000      6       2,3
    2       0.00000000      1.00000000      7       4,5
    3       0.00000000      1.00000000      8       6,7
    """
    #
    #          9                        10
    #         / \                      / \
    #        /   \                    /   8
    #       /     \                  /   / \
    #      7       \                /   /   \
    #     / \       6              /   /     6
    #    /   5     / \            /   5     / \
    #   /   / \   /   \          /   / \   /   \
    #  4   0   1 2     3        4   0   1 2     3
    #
    # 0 ------------------ 0.5 ------------------ 1.0
    nodes0 = """\
    id      is_sample   population      time
    0       1       0               0.00000000000000
    1       1       0               0.00000000000000
    2       1       0               0.00000000000000
    3       1       0               0.00000000000000
    4       1       0               0.00000000000000
    5       0       0               0.14567111023387
    6       0       0               0.21385545626353
    7       0       0               0.43508024345063
    8       0       0               0.60156352971203
    9       0       0               0.90000000000000
    10      0       0               1.20000000000000
    """
    edges0 = """\
    id      left            right           parent  child
    0       0.00000000      1.00000000      5       0,1
    1       0.00000000      1.00000000      6       2,3
    2       0.00000000      0.50000000      7       4,5
    3       0.50000000      1.00000000      8       5,6
    4       0.00000000      0.50000000      9       6,7
    5       0.50000000      1.00000000      10      4,8
    """
    nodes1 = """\
    id      is_sample   population      time
    0       0           0           1.0
    1       1           0           0.0
    2       1           0           0.0
    """
    edges1 = """\
    id      left            right           parent  child
    0       0.00000000      1.00000000      0       1,2
    """

    def do_map(self, ts, ancestors, samples=None, compare_lib=True):
        """
        Runs the Python test implementation of link_ancestors.
        """
        if samples is None:
            samples = ts.samples()
        s = tests.AncestorMap(ts, samples, ancestors)
        ancestor_table = s.link_ancestors()
        if compare_lib:
            lib_result = ts.tables.link_ancestors(samples, ancestors)
            assert ancestor_table == lib_result
        return ancestor_table

    def test_deprecated_name(self):
        # copied from test_single_tree_one_ancestor below
        nodes = io.StringIO(self.nodes)
        edges = io.StringIO(self.edges)
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        samples = ts.samples()
        ancestors = [8]
        s = tests.AncestorMap(ts, samples, ancestors)
        tss = s.link_ancestors()
        lib_result = ts.tables.map_ancestors(samples, ancestors)
        assert tss == lib_result
        assert list(tss.parent) == [8, 8, 8, 8, 8]
        assert list(tss.child) == [0, 1, 2, 3, 4]
        assert all(tss.left) == 0
        assert all(tss.right) == 1

    def test_single_tree_one_ancestor(self):
        nodes = io.StringIO(self.nodes)
        edges = io.StringIO(self.edges)
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        tss = self.do_map(ts, ancestors=[8])
        assert list(tss.parent) == [8, 8, 8, 8, 8]
        assert list(tss.child) == [0, 1, 2, 3, 4]
        assert all(tss.left) == 0
        assert all(tss.right) == 1

    def test_single_tree_unordered_nodes(self):
        nodes = io.StringIO(self.nodes1)
        edges = io.StringIO(self.edges1)
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        tss = self.do_map(ts, ancestors=[0])
        assert list(tss.parent) == [0, 0]
        assert list(tss.child) == [1, 2]
        assert all(tss.left) == 0
        assert all(tss.right) == 1

    def test_single_tree_two_ancestors(self):
        nodes = io.StringIO(self.nodes)
        edges = io.StringIO(self.edges)
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        tss = self.do_map(ts, ancestors=[6, 7])
        assert list(tss.parent) == [6, 6, 7, 7, 7]
        assert list(tss.child) == [2, 3, 0, 1, 4]
        assert all(tss.left) == 0
        assert all(tss.right) == 1

    def test_single_tree_no_ancestors(self):
        nodes = io.StringIO(self.nodes)
        edges = io.StringIO(self.edges)
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        tss = self.do_map(ts, samples=[2, 3], ancestors=[7])
        assert tss.num_rows == 0

    def test_single_tree_samples_or_ancestors_not_in_tree(self):
        nodes = io.StringIO(self.nodes)
        edges = io.StringIO(self.edges)
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        with pytest.raises(AssertionError):
            self.do_map(ts, samples=[-1, 3], ancestors=[5])
        with pytest.raises(AssertionError):
            self.do_map(ts, samples=[2, 3], ancestors=[10])

    def test_single_tree_ancestors_descend_from_other_ancestors(self):
        nodes = io.StringIO(self.nodes)
        edges = io.StringIO(self.edges)
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        tss = self.do_map(ts, ancestors=[7, 8])
        assert list(tss.parent) == [7, 7, 7, 8, 8, 8]
        assert list(tss.child) == [0, 1, 4, 2, 3, 7]
        assert all(tss.left) == 0
        assert all(tss.right) == 1

    def test_single_tree_internal_samples(self):
        nodes = io.StringIO(self.nodes)
        edges = io.StringIO(self.edges)
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        tss = self.do_map(ts, samples=[2, 3, 4, 5], ancestors=[7, 8])
        assert list(tss.parent) == [7, 7, 8, 8, 8]
        assert list(tss.child) == [4, 5, 2, 3, 7]
        assert all(tss.left) == 0
        assert all(tss.right) == 1

    def test_single_tree_samples_and_ancestors_overlap(self):
        nodes = io.StringIO(self.nodes)
        edges = io.StringIO(self.edges)
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        tss = self.do_map(ts, samples=[1, 2, 3, 5], ancestors=[5, 6, 7])
        assert list(tss.parent) == [5, 6, 6, 7]
        assert list(tss.child) == [1, 2, 3, 5]
        assert all(tss.left) == 0
        assert all(tss.right) == 1

    def test_single_tree_unary_ancestor(self):
        nodes = io.StringIO(self.nodes)
        edges = io.StringIO(self.edges)
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        tss = self.do_map(ts, samples=[1, 2, 4], ancestors=[5, 7, 8])
        assert list(tss.parent) == [5, 7, 7, 8, 8]
        assert list(tss.child) == [1, 4, 5, 2, 7]
        assert all(tss.left) == 0
        assert all(tss.right) == 1

    def test_single_tree_ancestors_descend_from_samples(self):
        nodes = io.StringIO(self.nodes)
        edges = io.StringIO(self.edges)
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        tss = self.do_map(ts, samples=[1, 7], ancestors=[5, 8])
        assert list(tss.parent) == [5, 7, 8]
        assert list(tss.child) == [1, 5, 7]
        assert all(tss.left) == 0
        assert all(tss.right) == 1

    def test_single_tree_samples_descend_from_samples(self):
        nodes = io.StringIO(self.nodes)
        edges = io.StringIO(self.edges)
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        tss = self.do_map(ts, samples=[3, 6], ancestors=[8])
        assert list(tss.parent) == [6, 8]
        assert list(tss.child) == [3, 6]
        assert all(tss.left) == 0
        assert all(tss.right) == 1

    def test_multiple_trees_to_single_tree(self):
        nodes = io.StringIO(self.nodes0)
        edges = io.StringIO(self.edges0)
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        tss = self.do_map(ts, ancestors=[5, 6])
        assert list(tss.parent) == [5, 5, 6, 6]
        assert list(tss.child) == [0, 1, 2, 3]
        assert all(tss.left) == 0
        assert all(tss.right) == 1

    def test_multiple_trees_one_ancestor(self):
        nodes = io.StringIO(self.nodes0)
        edges = io.StringIO(self.edges0)
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        tss = self.do_map(ts, ancestors=[9, 10])
        assert list(tss.parent) == [9, 9, 9, 9, 9, 10, 10, 10, 10, 10]
        assert list(tss.child) == [0, 1, 2, 3, 4, 0, 1, 2, 3, 4]
        assert all(tss.left) == 0
        assert all(tss.right) == 1

    def verify(self, ts, sample_nodes, ancestral_nodes):
        tss = self.do_map(ts, ancestors=ancestral_nodes, samples=sample_nodes)
        # ancestors = list(set(tss.parent))
        # Loop through the rows of the ancestral branch table.
        current_ancestor = tss.parent[0]
        current_descendants = [tss.child[0]]
        current_left = tss.left[0]
        current_right = tss.right[0]
        for _, row in enumerate(tss):
            if (
                row.parent != current_ancestor
                or row.left != current_left
                or row.right != current_right
            ):
                # Loop through trees.
                for tree in ts.trees():
                    if tree.interval.left >= current_right:
                        break
                    while tree.interval.right <= current_left:
                        tree.next()
                    # Check that the most recent ancestor of the descendants is the
                    # current_ancestor.
                    current_descendants = list(set(current_descendants))
                    for des in current_descendants:
                        par = tree.get_parent(des)
                        while par not in ancestral_nodes and par not in sample_nodes:
                            par = tree.get_parent(par)
                        assert par == current_ancestor
                # Reset the current ancestor and descendants, left and right coords.
                current_ancestor = row.parent
                current_descendants = [row.child]
                current_left = row.left
                current_right = row.right
            else:
                # Collate a list of children corresponding to each ancestral node.
                current_descendants.append(row.child)

    def test_sim_single_coalescent_tree(self):
        ts = msprime.simulate(30, random_seed=1, length=10)
        ancestors = [3 * n for n in np.arange(0, ts.num_nodes // 3)]
        self.verify(ts, ts.samples(), ancestors)
        random_samples = [4 * n for n in np.arange(0, ts.num_nodes // 4)]
        self.verify(ts, random_samples, ancestors)

    def test_sim_coalescent_trees(self):
        ts = msprime.simulate(8, recombination_rate=5, random_seed=1, length=2)
        ancestors = [3 * n for n in np.arange(0, ts.num_nodes // 3)]
        self.verify(ts, ts.samples(), ancestors)
        random_samples = [4 * n for n in np.arange(0, ts.num_nodes // 4)]
        self.verify(ts, random_samples, ancestors)

    def test_sim_coalescent_trees_internal_samples(self):
        ts = msprime.simulate(8, recombination_rate=5, random_seed=10, length=2)
        assert ts.num_trees > 2
        ancestors = [4 * n for n in np.arange(0, ts.num_nodes // 4)]
        self.verify(tsutil.jiggle_samples(ts), ts.samples(), ancestors)
        random_samples = [4 * n for n in np.arange(0, ts.num_nodes // 4)]
        self.verify(tsutil.jiggle_samples(ts), random_samples, ancestors)

    def test_sim_many_multiroot_trees(self):
        ts = msprime.simulate(7, recombination_rate=1, random_seed=10)
        assert ts.num_trees > 3
        ts = ts.decapitate(np.max(ts.tables.nodes.time) / 2)
        ancestors = [4 * n for n in np.arange(0, ts.num_nodes // 4)]
        self.verify(ts, ts.samples(), ancestors)
        random_samples = [4 * n for n in np.arange(0, ts.num_nodes // 4)]
        self.verify(ts, random_samples, ancestors)

    def test_sim_wright_fisher_generations(self):
        number_of_gens = 5
        tables = wf.wf_sim(10, number_of_gens, deep_history=False, seed=2)
        tables.sort()
        ts = tables.tree_sequence()
        ancestors = [4 * n for n in np.arange(0, ts.num_nodes // 4)]
        self.verify(ts, ts.samples(), ancestors)
        for gen in range(1, number_of_gens):
            ancestors = [u.id for u in ts.nodes() if u.time == gen]
            self.verify(ts, ts.samples(), ancestors)

        random_samples = [4 * n for n in np.arange(0, ts.num_nodes // 4)]
        self.verify(ts, random_samples, ancestors)
        for gen in range(1, number_of_gens):
            ancestors = [u.id for u in ts.nodes() if u.time == gen]
            self.verify(ts, random_samples, ancestors)


class TestMutationParent:
    """
    Tests that mutation parent is correctly specified, and that we correctly
    recompute it with compute_mutation_parent.
    """

    seed = 42

    def verify_parents(self, ts):
        parent = tsutil.compute_mutation_parent(ts)
        tables = ts.tables
        assert np.array_equal(parent, tables.mutations.parent)
        tables.mutations.parent = np.zeros_like(tables.mutations.parent) - 1
        assert np.all(tables.mutations.parent == tskit.NULL)
        tables.compute_mutation_parents()
        assert np.array_equal(parent, tables.mutations.parent)

    def test_example(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       0           2.0
        1       0           1.0
        2       0           1.0
        3       1           0
        4       1           0
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.0    0.5   2  3
        0.0    0.8   2  4
        0.5    1.0   1  3
        0.0    1.0   0  1
        0.0    1.0   0  2
        0.8    1.0   0  4
        """
        )
        sites = io.StringIO(
            """\
        position    ancestral_state
        0.1     0
        0.5     0
        0.9     0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    derived_state   parent
        0       1       1               -1
        0       2       1               -1
        0       3       2               1
        1       0       1               -1
        1       1       1               3
        1       3       2               4
        1       2       1               3
        1       4       2               6
        2       0       1               -1
        2       1       1               8
        2       2       1               8
        2       4       1               8
        """
        )
        ts = tskit.load_text(
            nodes=nodes, edges=edges, sites=sites, mutations=mutations, strict=False
        )
        self.verify_parents(ts)

    def test_single_muts(self):
        ts = msprime.simulate(
            10, random_seed=self.seed, mutation_rate=3.0, recombination_rate=1.0
        )
        self.verify_parents(ts)

    def test_with_jukes_cantor(self):
        ts = msprime.simulate(
            10, random_seed=self.seed, mutation_rate=0.0, recombination_rate=1.0
        )
        # make *lots* of recurrent mutations
        mut_ts = tsutil.jukes_cantor(
            ts, num_sites=10, mu=1, multiple_per_node=False, seed=self.seed
        )
        self.verify_parents(mut_ts)

    def test_with_jukes_cantor_multiple_per_node(self):
        ts = msprime.simulate(
            10, random_seed=self.seed, mutation_rate=0.0, recombination_rate=1.0
        )
        # make *lots* of recurrent mutations
        mut_ts = tsutil.jukes_cantor(
            ts, num_sites=10, mu=1, multiple_per_node=True, seed=self.seed
        )
        self.verify_parents(mut_ts)

    def verify_branch_mutations(self, ts, mutations_per_branch):
        ts = tsutil.insert_branch_mutations(ts, mutations_per_branch)
        assert ts.num_mutations > 1
        self.verify_parents(ts)

    def test_single_tree_one_mutation_per_branch(self):
        ts = msprime.simulate(6, random_seed=10)
        self.verify_branch_mutations(ts, 1)

    def test_single_tree_two_mutations_per_branch(self):
        ts = msprime.simulate(10, random_seed=9)
        self.verify_branch_mutations(ts, 2)

    def test_single_tree_three_mutations_per_branch(self):
        ts = msprime.simulate(8, random_seed=9)
        self.verify_branch_mutations(ts, 3)

    def test_single_multiroot_tree_recurrent_mutations(self):
        ts = msprime.simulate(6, random_seed=10)
        ts = ts.decapitate(np.max(ts.tables.nodes.time) / 2)
        for mutations_per_branch in [1, 2, 3]:
            self.verify_branch_mutations(ts, mutations_per_branch)

    def test_many_multiroot_trees_recurrent_mutations(self):
        ts = msprime.simulate(7, recombination_rate=1, random_seed=10)
        assert ts.num_trees > 3
        ts = ts.decapitate(np.max(ts.tables.nodes.time) / 2)
        for mutations_per_branch in [1, 2, 3]:
            self.verify_branch_mutations(ts, mutations_per_branch)


class TestMutationEdge:
    def verify_mutation_edge(self, ts):
        # print(ts.tables)
        for mutation in ts.mutations():
            site = ts.site(mutation.site)
            if mutation.edge == tskit.NULL:
                edges = [
                    edge
                    for edge in ts.edges()
                    if edge.left <= site.position < edge.right
                    and mutation.node == edge.child
                ]
                assert len(edges) == 0
            else:
                edge = ts.edge(mutation.edge)
                assert edge.left <= site.position < edge.right
                assert edge.child == mutation.node

        for tree in ts.trees():
            for site in tree.sites():
                for mutation in site.mutations:
                    assert mutation.edge == ts.mutation(mutation.id).edge
                    if mutation.edge == tskit.NULL:
                        assert tree.parent(mutation.node) == tskit.NULL

    def verify_branch_mutations(self, ts, mutations_per_branch):
        ts = tsutil.insert_branch_mutations(ts, mutations_per_branch)
        assert ts.num_mutations > 1
        self.verify_mutation_edge(ts)

    def test_single_tree_one_mutation_per_branch(self):
        ts = msprime.simulate(6, random_seed=10)
        self.verify_branch_mutations(ts, 1)

    def test_single_tree_two_mutations_per_branch(self):
        ts = msprime.simulate(10, random_seed=9)
        self.verify_branch_mutations(ts, 2)

    def test_single_tree_three_mutations_per_branch(self):
        ts = msprime.simulate(8, random_seed=9)
        self.verify_branch_mutations(ts, 3)

    def test_single_multiroot_tree_recurrent_mutations(self):
        ts = msprime.simulate(6, random_seed=10)
        ts = ts.decapitate(np.max(ts.tables.nodes.time) / 2)
        for mutations_per_branch in [1, 2, 3]:
            self.verify_branch_mutations(ts, mutations_per_branch)

    def test_many_multiroot_trees_recurrent_mutations(self):
        ts = msprime.simulate(7, recombination_rate=1, random_seed=10)
        assert ts.num_trees > 3
        ts = ts.decapitate(np.max(ts.tables.nodes.time) / 2)
        for mutations_per_branch in [1, 2, 3]:
            self.verify_branch_mutations(ts, mutations_per_branch)

    @pytest.mark.parametrize("n", range(2, 5))
    @pytest.mark.parametrize("mutations_per_branch", range(3))
    def test_balanced_binary_tree(self, n, mutations_per_branch):
        ts = tskit.Tree.generate_balanced(4).tree_sequence
        # These trees have a handy property
        assert all(edge.id == edge.child for edge in ts.edges())
        for mutation in ts.mutations():
            assert mutation.edge == mutation.node
        for site in ts.first().sites():
            for mutation in site.mutations:
                assert mutation.edge == mutation.node


class TestMutationTime:
    """
    Tests that mutation time is correctly specified, and that we correctly
    recompute it with compute_mutation_times.
    """

    seed = 42

    def verify_times(self, ts):
        tables = ts.tables
        # Clear out the existing mutations as they come from msprime
        tables.mutations.time = np.full(
            tables.mutations.time.shape, -1, dtype=np.float64
        )
        assert np.all(tables.mutations.time == -1)
        # Compute times with C method and dumb python method
        tables.compute_mutation_times()
        python_time = tsutil.compute_mutation_times(ts)
        assert np.allclose(python_time, tables.mutations.time, rtol=1e-15, atol=1e-15)

    def test_example(self):
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       0           2.0
        1       0           1.0
        2       0           1.0
        3       1           0
        4       1           0
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.0    0.5   2  3
        0.0    0.8   2  4
        0.5    1.0   1  3
        0.0    1.0   0  1
        0.0    1.0   0  2
        0.8    1.0   0  4
        """
        )
        sites = io.StringIO(
            """\
        position    ancestral_state
        0.1     0
        0.5     0
        0.9     0
        """
        )
        mutations = io.StringIO(
            """\
        site	node	time	derived_state	parent
        0       1       1.5     1               -1
        0       2       1.5     1               -1
        0       3       0.5     2               1
        1       0       2.0     1               -1
        1       1       1.5     1               3
        1       3       0.5     2               4
        1       2       1.5     1               3
        1       4       0.5     2               6
        2       0       2.0     1               -1
        2       1       1.5     1               8
        2       2       1.5     1               8
        2       4       1.0     1               8
        """
        )
        ts = tskit.load_text(
            nodes=nodes, edges=edges, sites=sites, mutations=mutations, strict=False
        )
        # ts.dump_text(mutations=sys.stdout)
        # self.assertFalse(True)
        tables = ts.tables
        python_time = tsutil.compute_mutation_times(ts)
        assert np.allclose(python_time, tables.mutations.time, rtol=1e-15, atol=1e-15)
        tables.mutations.time = np.full(
            tables.mutations.time.shape, -1, dtype=np.float64
        )
        assert np.all(tables.mutations.time == -1)
        tables.compute_mutation_times()
        assert np.allclose(python_time, tables.mutations.time, rtol=1e-15, atol=1e-15)

    def test_single_muts(self):
        ts = msprime.simulate(
            10, random_seed=self.seed, mutation_rate=3.0, recombination_rate=1.0
        )
        self.verify_times(ts)

    def test_with_jukes_cantor(self):
        ts = msprime.simulate(
            10, random_seed=self.seed, mutation_rate=0.0, recombination_rate=1.0
        )
        # make *lots* of recurrent mutations
        mut_ts = tsutil.jukes_cantor(
            ts, num_sites=10, mu=1, multiple_per_node=False, seed=self.seed
        )
        self.verify_times(mut_ts)

    def test_with_jukes_cantor_multiple_per_node(self):
        ts = msprime.simulate(
            10, random_seed=self.seed, mutation_rate=0.0, recombination_rate=1.0
        )
        # make *lots* of recurrent mutations
        mut_ts = tsutil.jukes_cantor(
            ts, num_sites=10, mu=1, multiple_per_node=True, seed=self.seed
        )
        self.verify_times(mut_ts)

    def verify_branch_mutations(self, ts, mutations_per_branch):
        ts = tsutil.insert_branch_mutations(ts, mutations_per_branch)
        assert ts.num_mutations > 1
        self.verify_times(ts)

    def test_single_tree_one_mutation_per_branch(self):
        ts = msprime.simulate(6, random_seed=10)
        self.verify_branch_mutations(ts, 1)

    def test_single_tree_two_mutations_per_branch(self):
        ts = msprime.simulate(10, random_seed=9)
        self.verify_branch_mutations(ts, 2)

    def test_single_tree_three_mutations_per_branch(self):
        ts = msprime.simulate(8, random_seed=9)
        self.verify_branch_mutations(ts, 3)

    def test_single_multiroot_tree_recurrent_mutations(self):
        ts = msprime.simulate(6, random_seed=10)
        ts = ts.decapitate(np.max(ts.tables.nodes.time) / 2)
        for mutations_per_branch in [1, 2, 3]:
            self.verify_branch_mutations(ts, mutations_per_branch)

    def test_many_multiroot_trees_recurrent_mutations(self):
        ts = msprime.simulate(7, recombination_rate=1, random_seed=10)
        assert ts.num_trees > 3
        ts = ts.decapitate(np.max(ts.tables.nodes.time) / 2)
        for mutations_per_branch in [1, 2, 3]:
            self.verify_branch_mutations(ts, mutations_per_branch)


class TestSimpleTreeAlgorithm:
    """
    Tests for the direct implementation of Algorithm T in tsutil.py.

    See TestHoleyTreeSequences above for further tests on wacky topologies.
    """

    def test_zero_nodes(self):
        tables = tskit.TableCollection(1)
        ts = tables.tree_sequence()
        assert ts.sequence_length == 1
        assert ts.num_trees == 1
        # Test the simple tree iterator.
        trees = list(tsutil.algorithm_T(ts))
        assert len(trees) == 1
        (left, right), parent = trees[0]
        assert left == 0
        assert right == 1
        assert parent == []

    def test_one_node(self):
        tables = tskit.TableCollection(1)
        tables.nodes.add_row()
        ts = tables.tree_sequence()
        assert ts.sequence_length == 1
        assert ts.num_trees == 1
        # Test the simple tree iterator.
        trees = list(tsutil.algorithm_T(ts))
        assert len(trees) == 1
        (left, right), parent = trees[0]
        assert left == 0
        assert right == 1
        assert parent == [-1]

    def test_single_coalescent_tree(self):
        ts = msprime.simulate(10, random_seed=1, length=10)
        tree = ts.first()
        p1 = [tree.parent(j) for j in range(ts.num_nodes)]
        interval, p2 = next(tsutil.algorithm_T(ts))
        assert interval == tree.interval
        assert p1 == p2

    def test_coalescent_trees(self):
        ts = msprime.simulate(8, recombination_rate=5, random_seed=1, length=2)
        assert ts.num_trees > 2
        new_trees = tsutil.algorithm_T(ts)
        for tree in ts.trees():
            interval, p2 = next(new_trees)
            p1 = [tree.parent(j) for j in range(ts.num_nodes)]
            assert interval == tree.interval
            assert p1 == p2
        with pytest.raises(StopIteration):
            next(new_trees)


class TestVirtualRootAPIs(ExampleTopologyMixin):
    """
    Tests the APIs based on getting roots.
    """

    def verify(self, ts):
        for tree in ts.trees():
            left_child = tree.left_child_array
            right_child = tree.right_child_array
            assert tree.virtual_root == ts.num_nodes
            assert tree.left_root == tree.left_child(tree.virtual_root)
            assert tree.right_root == tree.right_child(tree.virtual_root)
            assert tree.left_root == left_child[-1]
            assert tree.right_root == right_child[-1]
            assert tree.parent(tree.virtual_root) == tskit.NULL
            assert tree.left_sib(tree.virtual_root) == tskit.NULL
            assert tree.right_sib(tree.virtual_root) == tskit.NULL
            assert tree.num_children(tree.virtual_root) == tree.num_roots

            u = tree.left_root
            roots = []
            while u != tskit.NULL:
                roots.append(u)
                u = tree.right_sib(u)
            assert roots == list(tree.roots)

            # The branch_length for roots is defined as 0, and it's consistent
            # to have the same for the virtual root.
            assert tree.branch_length(tree.virtual_root) == 0
            # The virtual root has depth -1 from the root
            assert tree.depth(tree.virtual_root) == -1
            assert tree.num_children(tree.virtual_root) == tree.num_roots
            assert tree.num_samples(tree.virtual_root) == tree.num_samples()
            # We're not using tracked samples here.
            assert tree.num_tracked_samples(tree.virtual_root) == 0
            # The virtual_root is internal because it has children (the roots)
            assert tree.is_internal(tree.virtual_root)
            assert not tree.is_leaf(tree.virtual_root)
            assert not tree.is_sample(tree.virtual_root)
            # The mrca of the virtual_root and anything is itself
            assert tree.mrca(0, tree.virtual_root) == tree.virtual_root
            assert tree.mrca(tree.virtual_root, 0) == tree.virtual_root
            assert tree.mrca(tree.virtual_root, tree.virtual_root) == tree.virtual_root
            # The virtual_root is a descendant of nothing other than itself
            assert not tree.is_descendant(0, tree.virtual_root)
            assert tree.is_descendant(tree.virtual_root, tree.virtual_root)

            assert list(tree.leaves(tree.virtual_root)) == list(tree.leaves())
            assert list(tree.samples(tree.virtual_root)) == list(tree.samples())

            orders = [
                "preorder",
                "inorder",
                "levelorder",
                "breadthfirst",
                "postorder",
                "timeasc",
                "timedesc",
                "minlex_postorder",
            ]
            for order in orders:
                l_vr = list(tree.nodes(tree.virtual_root, order=order))
                l_standard = list(tree.nodes(order=order))
                assert len(l_vr) == 1 + len(l_standard)
                assert tree.virtual_root in l_vr

            # For pre-order, virtual_root should be first node visited:
            assert next(tree.nodes(tree.virtual_root)) == tree.virtual_root

            # Methods that imply looking up tree sequence properties of the
            # node raise an error
            # Some methods don't apply
            for method in [tree.population]:
                with pytest.raises(tskit.LibraryError, match="Node out of bounds"):
                    method(tree.virtual_root)


class TestSampleLists(ExampleTopologyMixin):
    """
    Tests for the sample lists algorithm.
    """

    def verify(self, ts):
        tree1 = tsutil.SampleListTree(ts)
        s = str(tree1)
        assert s is not None
        trees = ts.trees(sample_lists=True)
        for left, right in tree1.sample_lists():
            tree2 = next(trees)
            assert (left, right) == tree2.interval
            for u in tree2.nodes():
                assert tree1.left_sample[u] == tree2.left_sample(u)
                assert tree1.right_sample[u] == tree2.right_sample(u)
            for j in range(ts.num_samples):
                assert tree1.next_sample[j] == tree2.next_sample(j)
        assert right == ts.sequence_length

        tree1 = tsutil.SampleListTree(ts)
        trees = ts.trees(sample_lists=False)
        sample_index_map = ts.samples()
        for _, _ in tree1.sample_lists():
            tree2 = next(trees)
            for u in range(ts.num_nodes):
                samples2 = list(tree2.samples(u))
                samples1 = []
                index = tree1.left_sample[u]
                if index != tskit.NULL:
                    assert sample_index_map[tree1.left_sample[u]] == samples2[0]
                    assert sample_index_map[tree1.right_sample[u]] == samples2[-1]
                    stop = tree1.right_sample[u]
                    while True:
                        assert index != -1
                        samples1.append(sample_index_map[index])
                        if index == stop:
                            break
                        index = tree1.next_sample[index]
                assert samples1 == samples2
            # The python implementation here doesn't maintain roots
            np.testing.assert_array_equal(tree1.parent, tree2.parent_array[:-1])
            np.testing.assert_array_equal(tree1.left_child, tree2.left_child_array[:-1])
            np.testing.assert_array_equal(
                tree1.right_child, tree2.right_child_array[:-1]
            )
        assert right == ts.sequence_length


class TestOneSampleRoot(ExampleTopologyMixin):
    """
    Tests for the standard root threshold of subtending at least
    one sample.
    """

    def verify(self, ts):
        tree2 = tskit.Tree(ts)
        tree2.first()
        for interval, tree1 in tsutil.algorithm_R(ts, root_threshold=1):
            root_reachable_nodes = len(tree2.preorder())
            size_bound = tree1.num_edges + ts.num_samples
            assert size_bound >= root_reachable_nodes
            assert interval == tree2.interval
            assert tree1.roots() == tree2.roots
            # Definition here is the set unique path ends from samples
            roots = set()
            for u in ts.samples():
                while u != tskit.NULL:
                    path_end = u
                    u = tree2.parent(u)
                roots.add(path_end)
            assert set(tree1.roots()) == roots
            np.testing.assert_array_equal(tree1.parent, tree2.parent_array)
            np.testing.assert_array_equal(tree1.left_child, tree2.left_child_array)
            np.testing.assert_array_equal(tree1.right_child, tree2.right_child_array)
            np.testing.assert_array_equal(tree1.left_sib, tree2.left_sib_array)
            np.testing.assert_array_equal(tree1.right_sib, tree2.right_sib_array)
            np.testing.assert_array_equal(tree1.num_children, tree2.num_children_array)
            tree2.next()
        assert tree2.index == -1


class RootThreshold(ExampleTopologyMixin):
    """
    Tests for the root criteria of subtending at least k samples.
    """

    def verify(self, ts):
        k = self.root_threshold
        trees_py = tsutil.algorithm_R(ts, root_threshold=k)
        tree_lib = tskit.Tree(ts, root_threshold=k)
        tree_lib.first()
        tree_leg = tsutil.LegacyRootThresholdTree(ts, root_threshold=k)
        for (interval_py, tree_py), interval_leg in itertools.zip_longest(
            trees_py, tree_leg.iterate()
        ):
            assert interval_py == tree_lib.interval
            assert interval_leg == tree_lib.interval

            root_reachable_nodes = len(tree_lib.preorder())
            size_bound = tree_py.num_edges + ts.num_samples
            assert size_bound >= root_reachable_nodes
            assert tree_py.num_edges == tree_lib.num_edges

            # Definition here is the set unique path ends from samples
            # that subtend at least k samples
            roots = set()
            for u in ts.samples():
                while u != tskit.NULL:
                    path_end = u
                    u = tree_lib.parent(u)
                if tree_lib.num_samples(path_end) >= k:
                    roots.add(path_end)
            assert set(tree_py.roots()) == roots
            assert set(tree_lib.roots) == roots
            assert set(tree_leg.roots()) == roots
            assert len(tree_leg.roots()) == tree_lib.num_roots
            assert tree_py.roots() == tree_lib.roots

            # # The python class has identical behaviour to the lib version
            assert tree_py.left_child[-1] == tree_lib.left_root
            np.testing.assert_array_equal(tree_py.parent, tree_lib.parent_array)
            np.testing.assert_array_equal(tree_py.left_child, tree_lib.left_child_array)
            np.testing.assert_array_equal(
                tree_py.right_child, tree_lib.right_child_array
            )
            np.testing.assert_array_equal(tree_py.left_sib, tree_lib.left_sib_array)
            np.testing.assert_array_equal(tree_py.right_sib, tree_lib.right_sib_array)
            np.testing.assert_array_equal(
                tree_py.num_children, tree_lib.num_children_array
            )

            # NOTE: the legacy left_root value is *not* necessarily the same as the
            # new left_root.
            # assert tree_leg.left_root == tree_py.left_child[-1]

            # The virtual root version is identical to the legacy tree
            # except for the extra node and the details of the sib arrays.
            np.testing.assert_array_equal(tree_py.parent[:-1], tree_leg.parent)
            np.testing.assert_array_equal(tree_py.left_child[:-1], tree_leg.left_child)
            np.testing.assert_array_equal(
                tree_py.right_child[:-1], tree_leg.right_child
            )
            # The sib arrays are identical except for root nodes.
            for u in range(ts.num_nodes):
                if u not in roots:
                    assert tree_py.left_sib[u] == tree_leg.left_sib[u]
                    assert tree_py.right_sib[u] == tree_leg.right_sib[u]

            tree_lib.next()
        assert tree_lib.index == -1


class TestRootThreshold1(RootThreshold):
    root_threshold = 1


class TestRootThreshold2(RootThreshold):
    root_threshold = 2


class TestRootThreshold3(RootThreshold):
    root_threshold = 3


class TestRootThreshold4(RootThreshold):
    root_threshold = 4


class TestRootThreshold10(RootThreshold):
    root_threshold = 10


class TestSquashEdges:
    """
    Tests of the squash_edges function.
    """

    def do_squash(self, ts, compare_lib=True):
        squashed = ts.tables.edges
        squashed.squash()
        if compare_lib:
            squashed_list = squash_edges(ts)
            squashed_py = tskit.EdgeTable()
            for e in squashed_list:
                squashed_py.append(e)
            # Check the Python and C implementations produce the same output.
            assert squashed_py == squashed
        return squashed

    def test_simple_case(self):
        #   2
        #  / \
        # 0   1
        nodes = io.StringIO(
            """\
        id      is_sample   population      time
        0       1       0               0.00000000000000
        1       1       0               0.00000000000000
        2       0       0               1.00000000000000
        """
        )
        edges = io.StringIO(
            """\
        id      left            right           parent  child
        0       0.00000000      0.50000000      2       0
        1       0.00000000      0.50000000      2       1
        2       0.50000000      1.00000000      2       0
        3       0.50000000      1.00000000      2       1
        """
        )
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        edges = self.do_squash(ts)
        assert all(edges.left) == 0
        assert all(edges.right) == 1
        assert list(edges.parent) == [2, 2]
        assert list(edges.child) == [0, 1]

    def test_simple_case_unordered_intervals(self):
        # 1
        # |
        # 0
        nodes = io.StringIO(
            """\
        id      is_sample   population      time
        0       1           0               0.0
        1       0           0               1.0
        """
        )
        edges = io.StringIO(
            """\
        id      left            right           parent  child
        0       0.40            1.0             1       0
        0       0.00            0.40            1       0
        """
        )
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        edges = self.do_squash(ts)
        assert edges.left[0] == 0
        assert edges.right[0] == 1
        assert edges.parent[0] == 1
        assert edges.child[0] == 0

    def test_simple_case_unordered_children(self):
        #   2
        #  / \
        # 0   1
        nodes = io.StringIO(
            """\
        id      is_sample   population      time
        0       1       0               0.00000000000000
        1       1       0               0.00000000000000
        2       0       0               1.00000000000000
        """
        )
        edges = io.StringIO(
            """\
        id      left            right           parent  child
        0       0.50000000      1.00000000      2       1
        1       0.50000000      1.00000000      2       0
        2       0.00000000      0.50000000      2       1
        3       0.00000000      0.50000000      2       0
        """
        )
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        edges = self.do_squash(ts)
        assert all(edges.left) == 0
        assert all(edges.right) == 1
        assert list(edges.parent) == [2, 2]
        assert list(edges.child) == [0, 1]

    def test_simple_case_unordered_children_and_intervals(self):
        #   2
        #  / \
        # 0   1
        nodes = io.StringIO(
            """\
        id      is_sample   population      time
        0       1       0               0.00000000000000
        1       1       0               0.00000000000000
        2       0       0               1.00000000000000
        """
        )
        edges = io.StringIO(
            """\
        id      left            right           parent  child
        0       0.50000000      1.00000000      2       1
        2       0.00000000      0.50000000      2       1
        3       0.00000000      0.50000000      2       0
        1       0.50000000      1.00000000      2       0
        """
        )
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        edges = self.do_squash(ts)
        assert all(edges.left) == 0
        assert all(edges.right) == 1
        assert list(edges.parent) == [2, 2]
        assert list(edges.child) == [0, 1]

    def test_squash_multiple_parents_and_children(self):
        #   4       5
        #  / \     / \
        # 0   1   2   3
        nodes = io.StringIO(
            """\
        id      is_sample   population      time
        0       1       0               0.00000000000000
        1       1       0               0.00000000000000
        2       1       0               0.00000000000000
        3       1       0               0.00000000000000
        4       0       0               1.00000000000000
        5       0       0               1.00000000000000
        """
        )
        edges = io.StringIO(
            """\
        id      left            right           parent  child
        5       0.50000000      1.00000000      5       3
        6       0.50000000      1.00000000      5       2
        7       0.00000000      0.50000000      5       3
        8       0.00000000      0.50000000      5       2
        9       0.40000000      1.00000000      4       1
        10      0.00000000      0.40000000      4       1
        11      0.40000000      1.00000000      4       0
        12      0.00000000      0.40000000      4       0
        """
        )
        ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)
        edges = self.do_squash(ts)
        assert all(edges.left) == 0
        assert all(edges.right) == 1
        assert list(edges.parent) == [4, 4, 5, 5]
        assert list(edges.child) == [0, 1, 2, 3]

    def test_squash_overlapping_intervals(self):
        nodes = io.StringIO(
            """\
        id      is_sample   population      time
        0       1           0               0.0
        1       0           0               1.0
        """
        )
        edges = io.StringIO(
            """\
        id      left            right           parent  child
        0       0.00            0.50            1       0
        1       0.40            0.80            1       0
        2       0.60            1.00            1       0
        """
        )
        with pytest.raises(tskit.LibraryError):
            tskit.load_text(nodes=nodes, edges=edges, strict=False)

    def verify_slice_and_squash(self, ts):
        """
        Slices a tree sequence so that there are edge endpoints at
        all integer locations, then squashes these edges and verifies
        that the resulting edge table is the same as the input edge table.
        """
        sliced_edges = []
        # Create new sliced edge table.
        for e in ts.edges():
            left = e.left
            right = e.right

            if left == np.floor(left):
                r_left = np.ceil(left) + 1
            else:
                r_left = np.ceil(left)
            if right == np.floor(right):
                r_right = np.floor(right)
            else:
                r_right = np.floor(right) + 1

            new_range = [left]
            for r in np.arange(r_left, r_right):
                new_range.append(r)
            new_range.append(right)
            assert len(new_range) > 1

            # Add new edges to the list.
            for r in range(1, len(new_range)):
                new = tskit.Edge(new_range[r - 1], new_range[r], e.parent, e.child)
                sliced_edges.append(new)

        # Shuffle the edges and create a new edge table.
        random.shuffle(sliced_edges)
        sliced_table = tskit.EdgeTable()
        for e in sliced_edges:
            sliced_table.append(e)

        # Squash the edges and check against input table.
        sliced_table.squash()
        assert sliced_table == ts.tables.edges

    def test_sim_single_coalescent_tree(self):
        ts = msprime.simulate(20, random_seed=4, length=10)
        assert ts.num_trees == 1
        self.verify_slice_and_squash(ts)

    def test_sim_big_coalescent_trees(self):
        ts = msprime.simulate(20, recombination_rate=5, random_seed=4, length=10)
        assert ts.num_trees > 2
        self.verify_slice_and_squash(ts)


def squash_edges(ts):
    """
    Returns the edges in the tree sequence squashed.
    """
    t = ts.tables.nodes.time
    edges = list(ts.edges())
    edges.sort(key=lambda e: (t[e.parent], e.parent, e.child, e.left))
    if len(edges) == 0:
        return []

    squashed = []
    last_e = edges[0]
    for e in edges[1:]:
        condition = (
            e.parent != last_e.parent
            or e.child != last_e.child
            or e.left != last_e.right
        )
        if condition:
            squashed.append(last_e)
            last_e = e
        last_e.right = e.right
    squashed.append(last_e)
    return squashed


def reduce_topology(ts):
    """
    Returns a tree sequence with the minimal information required to represent
    the tree topologies at its sites. Uses a left-to-right algorithm.
    """
    tables = ts.dump_tables()
    edge_map = {}

    def add_edge(left, right, parent, child):
        new_edge = tskit.Edge(left, right, parent, child)
        if child not in edge_map:
            edge_map[child] = new_edge
        else:
            edge = edge_map[child]
            if edge.right == left and edge.parent == parent:
                # Squash
                edge.right = right
            else:
                tables.edges.append(edge)
                edge_map[child] = new_edge

    tables.edges.clear()

    edge_buffer = []
    first_site = True
    for tree in ts.trees():
        # print(tree.interval)
        # print(tree.draw(format="unicode"))
        if tree.num_sites > 0:
            sites = list(tree.sites())
            if first_site:
                x = 0
                # print("First site", sites)
                first_site = False
            else:
                x = sites[0].position
            # Flush the edge buffer.
            for left, parent, child in edge_buffer:
                add_edge(left, x, parent, child)
            # Add edges for each node in the tree.
            edge_buffer = []
            for root in tree.roots:
                for u in tree.nodes(root):
                    if u != root:
                        edge_buffer.append((x, tree.parent(u), u))
    # Add the final edges.
    for left, parent, child in edge_buffer:
        add_edge(left, tables.sequence_length, parent, child)
    # Flush the remaining edges to the table
    for edge in edge_map.values():
        tables.edges.append(edge)
    tables.sort()
    ts = tables.tree_sequence()
    # Now simplify to remove redundant nodes.
    return ts.simplify(map_nodes=True, filter_sites=False)


class TestReduceTopology:
    """
    Tests to ensure that reduce topology in simplify is equivalent to the
    reduce_topology function above.
    """

    def verify(self, ts):
        source_tables = ts.tables
        X = source_tables.sites.position
        position_count = {x: 0 for x in X}
        position_count[0] = 0
        position_count[ts.sequence_length] = 0
        mts, node_map = reduce_topology(ts)
        for edge in mts.edges():
            assert edge.left in position_count
            assert edge.right in position_count
            position_count[edge.left] += 1
            position_count[edge.right] += 1
        if ts.num_sites == 0:
            # We should have zero edges output.
            assert mts.num_edges == 0
        elif X[0] != 0:
            # The first site (if it's not zero) should be mapped to zero so
            # this never occurs in edges.
            assert position_count[X[0]] == 0

        minimised_trees = mts.trees()
        minimised_tree = next(minimised_trees)
        minimised_tree_sites = minimised_tree.sites()
        for tree in ts.trees():
            for site in tree.sites():
                minimised_site = next(minimised_tree_sites, None)
                if minimised_site is None:
                    minimised_tree = next(minimised_trees)
                    minimised_tree_sites = minimised_tree.sites()
                    minimised_site = next(minimised_tree_sites)
                assert site.position == minimised_site.position
                assert site.ancestral_state == minimised_site.ancestral_state
                assert site.metadata == minimised_site.metadata
                assert len(site.mutations) == len(minimised_site.mutations)

                for mutation, minimised_mutation in zip(
                    site.mutations, minimised_site.mutations
                ):
                    assert mutation.derived_state == minimised_mutation.derived_state
                    assert mutation.metadata == minimised_mutation.metadata
                    assert mutation.parent == minimised_mutation.parent
                    assert node_map[mutation.node] == minimised_mutation.node
            if tree.num_sites > 0:
                mapped_dict = {
                    node_map[u]: node_map[v] for u, v in tree.parent_dict.items()
                }
                assert mapped_dict == minimised_tree.parent_dict
        assert np.array_equal(ts.genotype_matrix(), mts.genotype_matrix())

        edges = list(mts.edges())
        squashed = squash_edges(mts)
        assert len(edges) == len(squashed)
        assert edges == squashed

        # Verify against simplify implementations.
        s = tests.Simplifier(
            ts, ts.samples(), reduce_to_site_topology=True, filter_sites=False
        )
        sts1, _ = s.simplify()
        sts2 = ts.simplify(reduce_to_site_topology=True, filter_sites=False)
        t1 = mts.tables
        for sts in [sts2, sts2]:
            t2 = sts.tables
            assert t1.nodes == t2.nodes
            assert t1.edges == t2.edges
            assert t1.sites == t2.sites
            assert t1.mutations == t2.mutations
            assert t1.populations == t2.populations
            assert t1.individuals == t2.individuals
        return mts

    def test_no_recombination_one_site(self):
        ts = msprime.simulate(15, random_seed=1)
        tables = ts.dump_tables()
        tables.sites.add_row(position=0.25, ancestral_state="0")
        mts = self.verify(tables.tree_sequence())
        assert mts.num_trees == 1

    def test_simple_recombination_one_site(self):
        ts = msprime.simulate(15, random_seed=1, recombination_rate=2)
        tables = ts.dump_tables()
        tables.sites.add_row(position=0.25, ancestral_state="0")
        mts = self.verify(tables.tree_sequence())
        assert mts.num_trees == 1

    def test_simple_recombination_fixed_sites(self):
        ts = msprime.simulate(5, random_seed=1, recombination_rate=2)
        tables = ts.dump_tables()
        for x in [0.25, 0.5, 0.75]:
            tables.sites.add_row(position=x, ancestral_state="0")
        self.verify(tables.tree_sequence())

    def get_integer_edge_ts(self, n, m):
        recombination_map = msprime.RecombinationMap.uniform_map(m, 1, num_loci=m)
        ts = msprime.simulate(n, random_seed=1, recombination_map=recombination_map)
        assert ts.num_trees > 1
        for edge in ts.edges():
            assert int(edge.left) == edge.left
            assert int(edge.right) == edge.right
        return ts

    def test_integer_edges_one_site(self):
        ts = self.get_integer_edge_ts(5, 10)
        tables = ts.dump_tables()
        tables.sites.add_row(position=1, ancestral_state="0")
        mts = self.verify(tables.tree_sequence())
        assert mts.num_trees == 1

    def test_integer_edges_all_sites(self):
        ts = self.get_integer_edge_ts(5, 10)
        tables = ts.dump_tables()
        for x in range(10):
            tables.sites.add_row(position=x, ancestral_state="0")
        mts = self.verify(tables.tree_sequence())
        assert mts.num_trees == ts.num_trees

    def test_simple_recombination_site_at_zero(self):
        ts = msprime.simulate(5, random_seed=1, recombination_rate=2)
        tables = ts.dump_tables()
        tables.sites.add_row(position=0, ancestral_state="0")
        mts = self.verify(tables.tree_sequence())
        assert mts.num_trees == 1

    def test_simple_recombination(self):
        ts = msprime.simulate(5, random_seed=1, recombination_rate=2, mutation_rate=2)
        self.verify(ts)

    def test_large_recombination(self):
        ts = msprime.simulate(
            25, random_seed=12, recombination_rate=5, mutation_rate=15
        )
        self.verify(ts)

    def test_no_recombination(self):
        ts = msprime.simulate(5, random_seed=1, mutation_rate=2)
        self.verify(ts)

    def test_no_mutation(self):
        ts = msprime.simulate(5, random_seed=1)
        self.verify(ts)

    def test_zero_sites(self):
        ts = msprime.simulate(5, random_seed=2)
        assert ts.num_sites == 0
        mts = ts.simplify(reduce_to_site_topology=True)
        assert mts.num_trees == 1
        assert mts.num_edges == 0

    def test_branch_sites(self):
        ts = msprime.simulate(15, random_seed=12, recombination_rate=2, length=10)
        ts = tsutil.insert_branch_sites(ts)
        self.verify(ts)

    def test_jiggled_samples(self):
        ts = msprime.simulate(8, random_seed=13, recombination_rate=2, length=10)
        ts = tsutil.jiggle_samples(ts)
        self.verify(ts)


def search_sorted(a, v):
    """
    Implementation of searchsorted based on binary search with the same
    semantics as numpy's searchsorted. Used as the basis of the C
    implementation which we use in the simplify algorithm.
    """
    upper = len(a)
    if upper == 0:
        return 0
    lower = 0
    while upper - lower > 1:
        mid = (upper + lower) // 2
        if v >= a[mid]:
            lower = mid
        else:
            upper = mid
    offset = 0
    if a[lower] < v:
        offset = 1
    return lower + offset


class TestSearchSorted:
    """
    Tests for the basic implementation of search_sorted.
    """

    def verify(self, a):
        a = np.array(a)
        start, end = a[0], a[-1]
        # Check random values.
        np.random.seed(43)
        for v in np.random.uniform(start, end, 10):
            assert search_sorted(a, v) == np.searchsorted(a, v)
        # Check equal values.
        for v in a:
            assert search_sorted(a, v) == np.searchsorted(a, v)
        # Check values outside bounds.
        for v in [start - 2, start - 1, end, end + 1, end + 2]:
            assert search_sorted(a, v) == np.searchsorted(a, v)

    def test_range(self):
        for j in range(1, 20):
            self.verify(range(j))

    def test_negative_range(self):
        for j in range(1, 20):
            self.verify(-1 * np.arange(j)[::-1])

    def test_random_unit_interval(self):
        np.random.seed(143)
        for size in range(1, 100):
            a = np.random.random(size=size)
            a.sort()
            self.verify(a)

    def test_random_interval(self):
        np.random.seed(143)
        for _ in range(10):
            interval = np.random.random(2) * 10
            interval.sort()
            a = np.random.uniform(*interval, size=100)
            a.sort()
            self.verify(a)

    def test_random_negative(self):
        np.random.seed(143)
        for _ in range(10):
            interval = np.random.random(2) * 5
            interval.sort()
            a = -1 * np.random.uniform(*interval, size=100)
            a.sort()
            self.verify(a)

    def test_edge_cases(self):
        for v in [0, 1]:
            assert search_sorted([], v) == np.searchsorted([], v)
            assert search_sorted([1], v) == np.searchsorted([1], v)


class TestDeleteSites:
    """
    Tests for the TreeSequence.delete_sites method
    """

    def ts_with_4_sites(self):
        ts = msprime.simulate(8, random_seed=3)
        tables = ts.dump_tables()
        tables.sites.set_columns(np.arange(0, 1, 0.25), *tskit.pack_strings(["G"] * 4))
        tables.mutations.add_row(site=1, node=ts.first().parent(0), derived_state="C")
        tables.mutations.add_row(site=1, node=0, derived_state="T", parent=0)
        tables.mutations.add_row(site=2, node=1, derived_state="A")
        return tables.tree_sequence()

    def test_remove_by_index(self):
        ts = self.ts_with_4_sites().delete_sites([])
        assert ts.num_sites == 4
        assert ts.num_mutations == 3
        ts = ts.delete_sites(2)
        assert ts.num_sites == 3
        assert ts.num_mutations == 2
        ts = ts.delete_sites([1, 2])
        assert ts.num_sites == 1
        assert ts.num_mutations == 0

    def test_remove_all(self):
        ts = self.ts_with_4_sites().delete_sites(range(4))
        assert ts.num_sites == 0
        assert ts.num_mutations == 0
        # should be OK to run on a siteless tree seq as no sites specified
        ts.delete_sites([])

    def test_remove_repeated_sites(self):
        ts = self.ts_with_4_sites()
        t1 = ts.delete_sites([0, 1], record_provenance=False)
        t2 = ts.delete_sites([0, 0, 1], record_provenance=False)
        t3 = ts.delete_sites([0, 0, 0, 1], record_provenance=False)
        assert t1.tables == t2.tables
        assert t1.tables == t3.tables

    def test_remove_different_orders(self):
        ts = self.ts_with_4_sites()
        t1 = ts.delete_sites([0, 1, 3], record_provenance=False)
        t2 = ts.delete_sites([0, 3, 1], record_provenance=False)
        t3 = ts.delete_sites([3, 0, 1], record_provenance=False)
        assert t1.tables == t2.tables
        assert t1.tables == t3.tables

    def test_remove_bad(self):
        ts = self.ts_with_4_sites()
        with pytest.raises(TypeError):
            ts.delete_sites(["1"])
        with pytest.raises(ValueError):
            ts.delete_sites(4)
        with pytest.raises(ValueError):
            ts.delete_sites(-5)

    def verify_removal(self, ts, remove_sites):
        tables = ts.dump_tables()
        tables.delete_sites(remove_sites)

        # Make sure we've computed the mutation parents properly.
        mutation_parent = tables.mutations.parent
        tables.compute_mutation_parents()
        assert np.array_equal(mutation_parent, tables.mutations.parent)

        tsd = tables.tree_sequence()
        assert tsd.num_sites == ts.num_sites - len(remove_sites)
        source_sites = [site for site in ts.sites() if site.id not in remove_sites]
        assert len(source_sites) == tsd.num_sites
        for s1, s2 in zip(source_sites, tsd.sites()):
            assert s1.position == s2.position
            assert s1.ancestral_state == s2.ancestral_state
            assert s1.metadata == s2.metadata
            assert len(s1.mutations) == len(s2.mutations)
            for m1, m2 in zip(s1.mutations, s2.mutations):
                assert m1.node == m2.node
                assert m1.derived_state == m2.derived_state
                assert m1.metadata == m2.metadata

        # Check we get the same genotype_matrix
        G1 = ts.genotype_matrix()
        G2 = tsd.genotype_matrix()
        keep = np.ones(ts.num_sites, dtype=bool)
        keep[remove_sites] = 0
        assert np.array_equal(G1[keep], G2)

    def test_simple_random_metadata(self):
        ts = msprime.simulate(10, mutation_rate=10, random_seed=2)
        ts = tsutil.add_random_metadata(ts)
        assert ts.num_mutations > 5
        self.verify_removal(ts, [1, 3])

    def test_simple_mixed_length_states(self):
        ts = msprime.simulate(10, random_seed=2, length=10)
        tables = ts.dump_tables()
        for j in range(10):
            tables.sites.add_row(j, "X" * j)
            tables.mutations.add_row(site=j, node=j, derived_state="X" * (j + 1))
        ts = tables.tree_sequence()
        self.verify_removal(ts, [9])

    def test_jukes_cantor_random_metadata(self):
        ts = msprime.simulate(10, random_seed=2)
        ts = tsutil.jukes_cantor(ts, 10, 1, seed=2)
        ts = tsutil.add_random_metadata(ts)
        assert ts.num_mutations > 10
        self.verify_removal(ts, [])
        self.verify_removal(ts, [0, 2, 4, 8])
        self.verify_removal(ts, range(5))

    def test_jukes_cantor_many_mutations(self):
        ts = msprime.simulate(2, random_seed=2)
        ts = tsutil.jukes_cantor(ts, 10, mu=10, seed=2)
        assert ts.num_mutations > 100
        self.verify_removal(ts, [1, 3, 5, 7])
        self.verify_removal(ts, [1])
        self.verify_removal(ts, [9])

    def test_jukes_cantor_one_site(self):
        ts = msprime.simulate(5, random_seed=2)
        ts = tsutil.jukes_cantor(ts, 1, mu=10, seed=2)
        assert ts.num_mutations > 10
        self.verify_removal(ts, [])
        self.verify_removal(ts, [0])


class TestKeepSingleInterval(unittest.TestCase):
    """
    Tests for cutting up tree sequences along the genome.
    """

    def test_slice_unchanged(self):
        ts = msprime.simulate(5, random_seed=1, recombination_rate=2, mutation_rate=2)
        tables = ts.dump_tables()
        tables.edges.packset_metadata([b"edge {i}" for i in range(ts.num_edges)])
        ts1 = tables.tree_sequence()
        ts2 = ts1.keep_intervals([[0, 1]], simplify=False, record_provenance=False)
        ts1.tables.assert_equals(ts2.tables)

    def test_slice_by_tree_positions(self):
        ts = msprime.simulate(5, random_seed=1, recombination_rate=2, mutation_rate=2)
        breakpoints = list(ts.breakpoints())

        # Keep the last 3 trees (from 4th last breakpoint onwards)
        ts_sliced = ts.keep_intervals([[breakpoints[-4], ts.sequence_length]])
        assert ts_sliced.num_trees == 4
        assert ts_sliced.num_edges < ts.num_edges
        self.assertAlmostEqual(ts_sliced.sequence_length, 1.0)
        last_3_mutations = 0
        for tree_index in range(-3, 0):
            last_3_mutations += ts.at_index(tree_index).num_mutations
        assert ts_sliced.num_mutations == last_3_mutations

        # Keep the first 3 trees
        ts_sliced = ts.keep_intervals([[0, breakpoints[3]]])
        assert ts_sliced.num_trees == 4
        assert ts_sliced.num_edges < ts.num_edges
        self.assertAlmostEqual(ts_sliced.sequence_length, 1)
        first_3_mutations = 0
        for tree_index in range(0, 3):
            first_3_mutations += ts.at_index(tree_index).num_mutations
        assert ts_sliced.num_mutations == first_3_mutations

        # Slice out the middle
        ts_sliced = ts.keep_intervals([[breakpoints[3], breakpoints[-4]]])
        assert ts_sliced.num_trees == ts.num_trees - 4
        assert ts_sliced.num_edges < ts.num_edges
        self.assertAlmostEqual(ts_sliced.sequence_length, 1.0)
        assert (
            ts_sliced.num_mutations
            == ts.num_mutations - first_3_mutations - last_3_mutations
        )

    def test_slice_by_position(self):
        ts = msprime.simulate(5, random_seed=1, recombination_rate=2, mutation_rate=2)
        ts_sliced = ts.keep_intervals([[0.4, 0.6]])
        positions = ts.tables.sites.position
        assert ts_sliced.num_sites == np.sum((positions >= 0.4) & (positions < 0.6))

    def test_slice_unsimplified(self):
        ts = msprime.simulate(5, random_seed=1, recombination_rate=2, mutation_rate=2)
        ts_sliced = ts.keep_intervals([[0.4, 0.6]], simplify=True)
        assert ts.num_nodes != ts_sliced.num_nodes
        self.assertAlmostEqual(ts_sliced.sequence_length, 1.0)
        ts_sliced = ts.keep_intervals([[0.4, 0.6]], simplify=False)
        assert ts.num_nodes == ts_sliced.num_nodes
        self.assertAlmostEqual(ts_sliced.sequence_length, 1.0)

    def test_slice_coordinates(self):
        ts = msprime.simulate(5, random_seed=1, recombination_rate=2, mutation_rate=2)
        ts_sliced = ts.keep_intervals([[0.4, 0.6]])
        self.assertAlmostEqual(ts_sliced.sequence_length, 1)
        assert ts_sliced.num_trees != ts.num_trees
        assert ts_sliced.at_index(0).total_branch_length == 0
        assert ts_sliced.at(0).total_branch_length == 0
        assert ts_sliced.at(0.399).total_branch_length == 0
        assert ts_sliced.at(0.4).total_branch_length != 0
        assert ts_sliced.at(0.5).total_branch_length != 0
        assert ts_sliced.at(0.599).total_branch_length != 0
        assert ts_sliced.at(0.6).total_branch_length == 0
        assert ts_sliced.at(0.999).total_branch_length == 0
        assert ts_sliced.at_index(-1).total_branch_length == 0

    def test_slice_migrations(self):
        pop_configs = [msprime.PopulationConfiguration(5) for _ in range(2)]
        migration_matrix = [[0, 0.05], [0.05, 0]]
        ts = msprime.simulate(
            population_configurations=pop_configs,
            migration_matrix=migration_matrix,
            record_migrations=True,
            recombination_rate=2,
            random_seed=1,
        )
        tables = ts.dump_tables()
        tables.migrations.packset_metadata(
            [b"migration {i}" for i in range(ts.num_migrations)]
        )
        ts = tables.tree_sequence()

        ts_sliced = ts.keep_intervals([[0, 1]], simplify=False)
        assert ts.tables.migrations == ts_sliced.tables.migrations
        ts_sliced = ts.keep_intervals([[0, 0.5]], simplify=False)
        assert np.max(ts_sliced.tables.migrations.right) <= 0.5
        assert ts.num_migrations > ts_sliced.num_migrations

        ts_sliced = ts.keep_intervals([[0.5, 1]], simplify=False)
        assert np.max(ts_sliced.tables.migrations.left) >= 0.5
        assert ts.num_migrations > ts_sliced.num_migrations

        ts_sliced = ts.keep_intervals([[0.4, 0.6]], simplify=False)
        assert np.max(ts_sliced.tables.migrations.right) <= 0.6
        assert np.max(ts_sliced.tables.migrations.left) >= 0.4
        assert ts.num_migrations > ts_sliced.num_migrations


class TestKeepIntervals(TopologyTestCase):
    """
    Tests for keep_intervals operation, where we slice out multiple disjoint
    intervals concurrently.
    """

    def example_intervals(self, tables):
        L = tables.sequence_length
        yield []
        yield [(0, L)]
        yield [(0, L / 2), (L / 2, L)]
        yield [(0, 0.25 * L), (0.75 * L, L)]
        yield [(0.25 * L, L)]
        yield [(0.25 * L, 0.5 * L)]
        yield [(0.25 * L, 0.5 * L), (0.75 * L, 0.8 * L)]

    def do_keep_intervals(
        self, tables, intervals, simplify=True, record_provenance=True
    ):
        t1 = tables.copy()
        simple_keep_intervals(t1, intervals, simplify, record_provenance)
        t2 = tables.copy()
        t2.keep_intervals(intervals, simplify, record_provenance)
        t1.assert_equals(t2, ignore_timestamps=True)
        return t2

    def test_migration_error(self):
        # keep_intervals should fail if simplify=True (default)
        pop_configs = [msprime.PopulationConfiguration(5) for _ in range(2)]
        migration_matrix = [[0, 0.05], [0.05, 0]]
        ts = msprime.simulate(
            population_configurations=pop_configs,
            migration_matrix=migration_matrix,
            record_migrations=True,
            recombination_rate=2,
            random_seed=1,
        )
        with pytest.raises(tskit.LibraryError):
            ts.tables.keep_intervals([[0, 1]])

    def test_bad_intervals(self):
        tables = tskit.TableCollection(10)
        bad_intervals = [[[1, 1]], [[-1, 0]], [[0, 11]], [[0, 5], [4, 6]]]
        for intervals in bad_intervals:
            with pytest.raises(ValueError):
                tables.keep_intervals(intervals)
            with pytest.raises(ValueError):
                tables.delete_intervals(intervals)

    def test_one_interval(self):
        ts = msprime.simulate(
            10, random_seed=self.random_seed, recombination_rate=2, mutation_rate=2
        )
        tables = ts.tables
        intervals = [(0.3, 0.7)]
        for simplify in (True, False):
            for rec_prov in (True, False):
                self.do_keep_intervals(tables, intervals, simplify, rec_prov)

    def test_two_intervals(self):
        ts = msprime.simulate(
            10, random_seed=self.random_seed, recombination_rate=2, mutation_rate=2
        )
        tables = ts.tables
        intervals = [(0.1, 0.2), (0.8, 0.9)]
        for simplify in (True, False):
            for rec_prov in (True, False):
                self.do_keep_intervals(tables, intervals, simplify, rec_prov)

    def test_ten_intervals(self):
        ts = msprime.simulate(
            10, random_seed=self.random_seed, recombination_rate=2, mutation_rate=2
        )
        tables = ts.tables
        intervals = [(x, x + 0.05) for x in np.arange(0.0, 1.0, 0.1)]
        for simplify in (True, False):
            for rec_prov in (True, False):
                self.do_keep_intervals(tables, intervals, simplify, rec_prov)

    def test_hundred_intervals(self):
        ts = msprime.simulate(
            10, random_seed=self.random_seed, recombination_rate=2, mutation_rate=2
        )
        tables = ts.tables
        intervals = [(x, x + 0.005) for x in np.arange(0.0, 1.0, 0.01)]
        for simplify in (True, False):
            for rec_prov in (True, False):
                self.do_keep_intervals(tables, intervals, simplify, rec_prov)

    def test_regular_intervals(self):
        ts = msprime.simulate(
            3, random_seed=1234, recombination_rate=2, mutation_rate=2
        )
        tables = ts.tables
        eps = 0.0125
        for num_intervals in range(2, 10):
            breaks = np.linspace(0, ts.sequence_length, num=num_intervals)
            intervals = [(x, x + eps) for x in breaks[:-1]]
            self.do_keep_intervals(tables, intervals)

    def test_no_edges_sites(self):
        tables = tskit.TableCollection(1.0)
        tables.sites.add_row(0.1, "A")
        tables.sites.add_row(0.2, "T")
        for intervals in self.example_intervals(tables):
            assert len(tables.sites) == 2
            diced = self.do_keep_intervals(tables, intervals)
            assert diced.sequence_length == 1
            assert len(diced.edges) == 0
            assert len(diced.sites) == 0

    def verify(self, tables):
        for intervals in self.example_intervals(tables):
            for simplify in [True, False]:
                self.do_keep_intervals(tables, intervals, simplify=simplify)

    def test_empty_tables(self):
        tables = tskit.TableCollection(1.0)
        self.verify(tables)

    def test_single_tree_jukes_cantor(self):
        ts = msprime.simulate(6, random_seed=1, mutation_rate=1)
        ts = tsutil.jukes_cantor(ts, 20, 1, seed=10)
        self.verify(ts.tables)

    def test_single_tree_multichar_mutations(self):
        ts = msprime.simulate(6, random_seed=1, mutation_rate=1)
        ts = tsutil.insert_multichar_mutations(ts)
        self.verify(ts.tables)

    def test_many_trees_infinite_sites(self):
        ts = msprime.simulate(6, recombination_rate=2, mutation_rate=2, random_seed=1)
        assert ts.num_sites > 0
        assert ts.num_trees > 2
        self.verify(ts.tables)

    def test_many_trees_sequence_length_infinite_sites(self):
        for L in [0.5, 1.5, 3.3333]:
            ts = msprime.simulate(
                6, length=L, recombination_rate=2, mutation_rate=1, random_seed=1
            )
            self.verify(ts.tables)

    def test_wright_fisher_unsimplified(self):
        tables = wf.wf_sim(
            4,
            5,
            seed=1,
            deep_history=True,
            initial_generation_samples=False,
            num_loci=10,
        )
        tables.sort()
        ts = msprime.mutate(tables.tree_sequence(), rate=0.05, random_seed=234)
        assert ts.num_sites > 0
        self.verify(ts.tables)

    def test_wright_fisher_initial_generation(self):
        tables = wf.wf_sim(
            6, 5, seed=3, deep_history=True, initial_generation_samples=True, num_loci=2
        )
        tables.sort()
        tables.simplify()
        ts = msprime.mutate(tables.tree_sequence(), rate=0.08, random_seed=2)
        assert ts.num_sites > 0
        self.verify(ts.tables)

    def test_wright_fisher_initial_generation_no_deep_history(self):
        tables = wf.wf_sim(
            7,
            15,
            seed=202,
            deep_history=False,
            initial_generation_samples=True,
            num_loci=5,
        )
        tables.sort()
        tables.simplify()
        ts = msprime.mutate(tables.tree_sequence(), rate=0.2, random_seed=2)
        assert ts.num_sites > 0
        self.verify(ts.tables)

    def test_wright_fisher_unsimplified_multiple_roots(self):
        tables = wf.wf_sim(
            8,
            15,
            seed=1,
            deep_history=False,
            initial_generation_samples=False,
            num_loci=20,
        )
        tables.sort()
        ts = msprime.mutate(tables.tree_sequence(), rate=0.006, random_seed=2)
        assert ts.num_sites > 0
        self.verify(ts.tables)

    def test_wright_fisher_simplified(self):
        tables = wf.wf_sim(
            9,
            10,
            seed=1,
            deep_history=True,
            initial_generation_samples=False,
            num_loci=5,
        )
        tables.sort()
        ts = tables.tree_sequence().simplify()
        ts = msprime.mutate(ts, rate=0.2, random_seed=1234)
        assert ts.num_sites > 0
        self.verify(ts.tables)


class TestKeepDeleteIntervalsExamples:
    """
    Simple examples of keep/delete intervals at work.
    """

    def test_tables_single_tree_keep_middle(self):
        ts = msprime.simulate(10, random_seed=2)
        t_keep = ts.dump_tables()
        t_keep.keep_intervals([[0.25, 0.5]], record_provenance=False)
        t_delete = ts.dump_tables()
        t_delete.delete_intervals([[0, 0.25], [0.5, 1.0]], record_provenance=False)
        assert t_keep == t_delete

    def test_tables_single_tree_delete_middle(self):
        ts = msprime.simulate(10, random_seed=2)
        t_keep = ts.dump_tables()
        t_keep.delete_intervals([[0.25, 0.5]], record_provenance=False)
        t_delete = ts.dump_tables()
        t_delete.keep_intervals([[0, 0.25], [0.5, 1.0]], record_provenance=False)
        assert t_keep == t_delete

    def test_ts_single_tree_keep_middle(self):
        ts = msprime.simulate(10, random_seed=2)
        ts_keep = ts.keep_intervals([[0.25, 0.5]], record_provenance=False)
        ts_delete = ts.delete_intervals(
            [[0, 0.25], [0.5, 1.0]], record_provenance=False
        )
        assert ts_keep == ts_delete

    def test_ts_single_tree_delete_middle(self):
        ts = msprime.simulate(10, random_seed=2)
        ts_keep = ts.delete_intervals([[0.25, 0.5]], record_provenance=False)
        ts_delete = ts.keep_intervals([[0, 0.25], [0.5, 1.0]], record_provenance=False)
        assert ts_keep == ts_delete

    def test_ts_migrations(self):
        pop_configs = [msprime.PopulationConfiguration(5) for _ in range(2)]
        migration_matrix = [[0, 0.05], [0.05, 0]]
        ts = msprime.simulate(
            population_configurations=pop_configs,
            migration_matrix=migration_matrix,
            record_migrations=True,
            recombination_rate=2,
            random_seed=1,
        )
        ts_keep = ts.delete_intervals(
            [[0.25, 0.5]], record_provenance=False, simplify=False
        )
        ts_delete = ts.keep_intervals(
            [[0, 0.25], [0.5, 1.0]], record_provenance=False, simplify=False
        )
        assert ts_keep == ts_delete


class TestTrim(unittest.TestCase):
    """
    Test the trimming functionality
    """

    def add_mutations(self, ts, position, ancestral_state, derived_states, nodes):
        """
        Create a site at the specified position and assign mutations to the specified
        nodes (could be sequential mutations)
        """
        tables = ts.dump_tables()
        site = tables.sites.add_row(position, ancestral_state)
        for state, node in zip(derived_states, nodes):
            tables.mutations.add_row(site, node, state)
        tables.sort()
        tables.build_index()
        tables.compute_mutation_parents()
        return tables.tree_sequence()

    def verify_sites(self, source_tree, trimmed_tree, position_offset):
        source_sites = list(source_tree.sites())
        trimmed_sites = list(trimmed_tree.sites())
        assert len(source_sites) == len(trimmed_sites)
        for source_site, trimmed_site in zip(source_sites, trimmed_sites):
            self.assertAlmostEqual(
                source_site.position, position_offset + trimmed_site.position
            )
            assert source_site.ancestral_state == trimmed_site.ancestral_state
            assert source_site.metadata == trimmed_site.metadata
            assert len(source_site.mutations) == len(trimmed_site.mutations)
            for source_mut, trimmed_mut in zip(
                source_site.mutations, trimmed_site.mutations
            ):
                assert source_mut.node == trimmed_mut.node
                assert source_mut.derived_state == trimmed_mut.derived_state
                assert source_mut.metadata == trimmed_mut.metadata
                # mutation.parent id may have changed after deleting redundant mutations
                if source_mut.parent == trimmed_mut.parent == tskit.NULL:
                    pass
                else:
                    assert (
                        source_tree.tree_sequence.mutation(source_mut.parent).node
                        == trimmed_tree.tree_sequence.mutation(trimmed_mut.parent).node
                    )

    def verify_ltrim(self, source_ts, trimmed_ts):
        deleted_span = source_ts.first().span
        self.assertAlmostEqual(
            source_ts.sequence_length, trimmed_ts.sequence_length + deleted_span
        )
        assert source_ts.num_trees == trimmed_ts.num_trees + 1
        for j in range(trimmed_ts.num_trees):
            source_tree = source_ts.at_index(j + 1)
            trimmed_tree = trimmed_ts.at_index(j)
            assert source_tree.parent_dict == trimmed_tree.parent_dict
            self.assertAlmostEqual(source_tree.span, trimmed_tree.span)
            self.assertAlmostEqual(
                source_tree.interval.left, trimmed_tree.interval.left + deleted_span
            )
            self.verify_sites(source_tree, trimmed_tree, deleted_span)

    def verify_rtrim(self, source_ts, trimmed_ts):
        deleted_span = source_ts.last().span
        self.assertAlmostEqual(
            source_ts.sequence_length, trimmed_ts.sequence_length + deleted_span
        )
        assert source_ts.num_trees == trimmed_ts.num_trees + 1
        for j in range(trimmed_ts.num_trees):
            source_tree = source_ts.at_index(j)
            trimmed_tree = trimmed_ts.at_index(j)
            assert source_tree.parent_dict == trimmed_tree.parent_dict
            assert source_tree.interval == trimmed_tree.interval
            self.verify_sites(source_tree, trimmed_tree, 0)

    def clear_left_mutate(self, ts, left, num_sites):
        """
        Clear the edges from a tree sequence left of the specified coordinate
        and add in num_sites regularly spaced sites into the cleared region.
        """
        new_ts = ts.delete_intervals([[0.0, left]])
        for j, x in enumerate(np.linspace(0, left, num_sites, endpoint=False)):
            new_ts = self.add_mutations(new_ts, x, "A" * j, ["T"] * j, range(j + 1))
        return new_ts

    def clear_right_mutate(self, ts, right, num_sites):
        """
        Clear the edges from a tree sequence right of the specified coordinate
        and add in num_sites regularly spaced sites into the cleared region.
        """
        new_ts = ts.delete_intervals([[right, ts.sequence_length]])
        for j, x in enumerate(
            np.linspace(right, ts.sequence_length, num_sites, endpoint=False)
        ):
            new_ts = self.add_mutations(new_ts, x, "A" * j, ["T"] * j, range(j + 1))
        return new_ts

    def clear_left_right_234(self, left, right):
        """
        Clear edges to left and right and add 2 mutations at the same site into the left
        cleared region, 3 at the same site into the untouched region, and 4 into the
        right cleared region.
        """
        assert 0.0 < left < right < 1.0
        ts = msprime.simulate(10, recombination_rate=10, random_seed=2)
        left_pos = np.mean([0.0, left])
        left_root = ts.at(left_pos).root
        mid_pos = np.mean([left, right])
        mid_root = ts.at(mid_pos).root
        right_pos = np.mean([right, ts.sequence_length])
        right_root = ts.at(right_pos).root
        # Clear
        ts = ts.keep_intervals([[left, right]], simplify=False)
        ts = self.add_mutations(ts, left_pos, "A", ["T", "C"], [left_root, 0])
        ts = self.add_mutations(ts, mid_pos, "T", ["A", "C", "G"], [mid_root, 0, 1])
        ts = self.add_mutations(
            ts, right_pos, "X", ["T", "C", "G", "A"], [right_root, 0, 1, 2]
        )
        assert np.min(ts.tables.edges.left) != 0
        assert ts.num_mutations == 9
        assert ts.num_sites == 3
        return ts

    def migration_sim(self):
        pop_configs = [msprime.PopulationConfiguration(5) for _ in range(2)]
        migration_matrix = [[0, 0.05], [0.05, 0]]
        ts = msprime.simulate(
            population_configurations=pop_configs,
            migration_matrix=migration_matrix,
            record_migrations=True,
            recombination_rate=2,
            random_seed=1,
        )
        return ts

    def test_ltrim_single_tree(self):
        ts = msprime.simulate(10, mutation_rate=12, random_seed=2)
        ts = self.clear_left_mutate(ts, 0.5, 10)
        self.verify_ltrim(ts, ts.ltrim())

    def test_ltrim_single_tree_no_mutations(self):
        ts = msprime.simulate(10, random_seed=2)
        ts = self.clear_left_mutate(ts, 0.5, 0)
        self.verify_ltrim(ts, ts.ltrim())

    def test_ltrim_single_tree_tiny_left(self):
        ts = msprime.simulate(10, mutation_rate=12, random_seed=2)
        ts = self.clear_left_mutate(ts, 1e-200, 10)
        self.verify_ltrim(ts, ts.ltrim())

    def test_ltrim_many_trees(self):
        ts = msprime.simulate(
            10, recombination_rate=10, mutation_rate=12, random_seed=2
        )
        ts = self.clear_left_mutate(ts, 0.5, 10)
        self.verify_ltrim(ts, ts.ltrim())

    def test_ltrim_many_trees_left_min(self):
        ts = msprime.simulate(
            10, recombination_rate=10, mutation_rate=12, random_seed=2
        )
        ts = self.clear_left_mutate(ts, sys.float_info.min, 10)
        self.verify_ltrim(ts, ts.ltrim())

    def test_ltrim_many_trees_left_epsilon(self):
        ts = msprime.simulate(
            10, recombination_rate=10, mutation_rate=12, random_seed=2
        )
        ts = self.clear_left_mutate(ts, sys.float_info.epsilon, 0)
        self.verify_ltrim(ts, ts.ltrim())

    def test_ltrim_empty(self):
        ts = msprime.simulate(2, random_seed=2)
        ts = ts.delete_intervals([[0, 1]])
        with pytest.raises(ValueError):
            ts.ltrim()

    def test_ltrim_multiple_mutations(self):
        ts = self.clear_left_right_234(0.1, 0.5)
        trimmed_ts = ts.ltrim()
        self.assertAlmostEqual(trimmed_ts.sequence_length, 0.9)
        assert trimmed_ts.num_sites == 2
        assert trimmed_ts.num_mutations == 7  # We should have deleted 2
        assert np.min(trimmed_ts.tables.edges.left) == 0
        self.verify_ltrim(ts, trimmed_ts)

    def test_ltrim_migrations(self):
        ts = self.migration_sim()
        ts = ts.delete_intervals([[0, 0.1]], simplify=False)
        trimmed_ts = ts.ltrim()
        assert np.array_equal(
            trimmed_ts.tables.migrations.left, ts.tables.migrations.left - 0.1
        )
        assert np.array_equal(
            trimmed_ts.tables.migrations.right, ts.tables.migrations.right - 0.1
        )

    def test_rtrim_single_tree(self):
        ts = msprime.simulate(10, mutation_rate=12, random_seed=2)
        ts = self.clear_right_mutate(ts, 0.5, 10)
        self.verify_rtrim(ts, ts.rtrim())

    def test_rtrim_single_tree_no_mutations(self):
        ts = msprime.simulate(10, random_seed=2)
        ts = self.clear_right_mutate(ts, 0.5, 0)
        self.verify_rtrim(ts, ts.rtrim())

    def test_rtrim_single_tree_tiny_left(self):
        ts = msprime.simulate(10, mutation_rate=12, random_seed=2)
        ts = self.clear_right_mutate(ts, 1e-200, 10)
        self.verify_rtrim(ts, ts.rtrim())

    def test_rtrim_many_trees(self):
        ts = msprime.simulate(
            10, recombination_rate=10, mutation_rate=12, random_seed=2
        )
        ts = self.clear_right_mutate(ts, 0.5, 10)
        self.verify_rtrim(ts, ts.rtrim())

    def test_rtrim_many_trees_left_min(self):
        ts = msprime.simulate(
            10, recombination_rate=10, mutation_rate=12, random_seed=2
        )
        ts = self.clear_right_mutate(ts, sys.float_info.min, 10)
        self.verify_rtrim(ts, ts.rtrim())

    def test_rtrim_many_trees_left_epsilon(self):
        ts = msprime.simulate(
            10, recombination_rate=10, mutation_rate=12, random_seed=2
        )
        ts = self.clear_right_mutate(ts, sys.float_info.epsilon, 0)
        self.verify_rtrim(ts, ts.rtrim())

    def test_rtrim_empty(self):
        ts = msprime.simulate(2, random_seed=2)
        ts = ts.delete_intervals([[0, 1]])
        with pytest.raises(ValueError):
            ts.rtrim()

    def test_rtrim_multiple_mutations(self):
        ts = self.clear_left_right_234(0.1, 0.5)
        trimmed_ts = ts.rtrim()
        self.assertAlmostEqual(trimmed_ts.sequence_length, 0.5)
        assert trimmed_ts.num_sites == 2
        assert trimmed_ts.num_mutations == 5  # We should have deleted 4
        assert (
            np.max(trimmed_ts.tables.edges.right) == trimmed_ts.tables.sequence_length
        )
        self.verify_rtrim(ts, trimmed_ts)

    def test_rtrim_migrations(self):
        ts = self.migration_sim()
        ts = ts.delete_intervals([[0.9, 1]], simplify=False)
        trimmed_ts = ts.rtrim()
        trimmed_rights = trimmed_ts.tables.migrations.right
        assert np.max(trimmed_rights) == 0.9

    def test_trim_multiple_mutations(self):
        ts = self.clear_left_right_234(0.1, 0.5)
        trimmed_ts = ts.trim()
        self.assertAlmostEqual(trimmed_ts.sequence_length, 0.4)
        assert trimmed_ts.num_mutations == 3
        assert trimmed_ts.num_sites == 1
        assert np.min(trimmed_ts.tables.edges.left) == 0
        assert (
            np.max(trimmed_ts.tables.edges.right) == trimmed_ts.tables.sequence_length
        )

    def test_trims_no_effect(self):
        # Deleting from middle should have no effect on any trim function
        ts = msprime.simulate(10, recombination_rate=2, mutation_rate=50, random_seed=2)
        ts = ts.delete_intervals([[0.1, 0.5]])
        trimmed_ts = ts.ltrim(record_provenance=False)
        assert ts == trimmed_ts
        trimmed_ts = ts.rtrim(record_provenance=False)
        assert ts == trimmed_ts
        trimmed_ts = ts.trim(record_provenance=False)
        assert ts == trimmed_ts

    def test_failure_with_migrations(self):
        # All trim functions fail if migrations extend further than rightmost or
        # leftmost edges
        ts = msprime.simulate(10, recombination_rate=2, random_seed=2)
        ts = ts.keep_intervals([[0.1, 0.5]])
        tables = ts.dump_tables()
        tables.migrations.add_row(0, 1, 0, 0, 0, 0)
        ts = tables.tree_sequence()
        with pytest.raises(ValueError):
            ts.ltrim()
        with pytest.raises(ValueError):
            ts.rtrim()
        with pytest.raises(ValueError):
            ts.trim()


class TestMissingData:
    """
    Test various aspects of missing data functionality
    """

    # TODO tests for missing data currently sparse: more tests should go here

    def ts_missing_middle(self):
        # Simple ts with sample 0 missing a middle section
        ts = msprime.simulate(4, mutation_rate=1, recombination_rate=4, random_seed=2)
        tables = ts.dump_tables()
        tables.edges.clear()
        # mark the middle as missing
        for e in ts.tables.edges:
            if e.child == 0:
                if e.left == 0.0:
                    missing_from = e.right
                elif e.right == 1.0:
                    missing_to = e.left
                else:
                    continue  # omit this edge => node is isolated
            tables.edges.append(e)
        # Check we have non-missing to L & R
        assert 0.0 < missing_from < 1.0
        assert 0.0 < missing_to < 1.0
        return tables.tree_sequence(), missing_from, missing_to

    def test_is_isolated(self):
        ts, missing_from, missing_to = self.ts_missing_middle()
        for tree in ts.trees():
            if tree.interval.right > missing_from and tree.interval.left < missing_to:
                assert tree.is_isolated(0)
                assert not tree.is_isolated(1)
            else:
                assert not tree.is_isolated(0)
                assert not tree.is_isolated(1)
            # A non-sample node is isolated if not in the tree
            tree_nodes = set(tree.nodes())
            for nonsample_node in np.setdiff1d(np.arange(ts.num_nodes), ts.samples()):
                if nonsample_node in tree_nodes:
                    assert not tree.is_isolated(nonsample_node)
                else:
                    assert tree.is_isolated(nonsample_node)

    def test_is_isolated_bad(self):
        ts, missing_from, missing_to = self.ts_missing_middle()
        for tree in ts.trees():
            with pytest.raises(ValueError):
                tree.is_isolated(tskit.NULL)
            with pytest.raises(ValueError):
                tree.is_isolated(ts.num_nodes + 1)
            with pytest.raises(ValueError):
                tree.is_isolated(-2)
            with pytest.raises(TypeError):
                tree.is_isolated(None)
            with pytest.raises(TypeError):
                tree.is_isolated("abc")
            with pytest.raises(TypeError):
                tree.is_isolated(1.1)


--- ../../tskit/python/tests/test_ld_matrix.py ---


"""
Test cases for two-locus statistics
"""
import contextlib
import io
from dataclasses import dataclass
from itertools import combinations_with_replacement
from itertools import permutations
from itertools import product
from typing import Any
from typing import Callable
from typing import Dict
from typing import Generator
from typing import List
from typing import Tuple

import msprime
import numpy as np
import pytest

import tskit
from tests import tsutil
from tests.test_highlevel import get_example_tree_sequences


@contextlib.contextmanager
def suppress_overflow_div0_warning():
    with np.errstate(over="ignore", invalid="ignore", divide="ignore"):
        yield


class BitSet:
    """BitSet object, which stores values in arrays of unsigned integers.
    The rows represent all possible values a bit can take, and the rows
    represent each item that can be stored in the array.

    :param num_bits: The number of values that a single row can contain.
    :param length: The number of rows.
    """

    DTYPE = np.uint32  # Data type to be stored in the bitset
    CHUNK_SIZE = DTYPE(32)  # Size of integer field to store the data in

    def __init__(self: "BitSet", num_bits: int, length: int) -> None:
        self.row_len = num_bits // self.CHUNK_SIZE
        self.row_len += 1 if num_bits % self.CHUNK_SIZE else 0
        self.row_len = int(self.row_len)
        self.data = np.zeros(self.row_len * length, dtype=self.DTYPE)

    def intersect(
        self: "BitSet", self_row: int, other: "BitSet", other_row: int, out: "BitSet"
    ) -> None:
        """Intersect a row from the current array instance with a row from
        another BitSet and store it in an output bit array of length 1.

        NB: we don't specify the row in the output array, it is expected
        to be length 1.

        :param self_row: Row from the current array instance to be intersected.
        :param other: Other BitSet to intersect with.
        :param other_row: Row from the other BitSet instance.
        :param out: BitArray to store the result.
        """
        self_offset = self_row * self.row_len
        other_offset = other_row * self.row_len

        for i in range(self.row_len):
            out.data[i] = self.data[i + self_offset] & other.data[i + other_offset]

    def difference(
        self: "BitSet", self_row: int, other: "BitSet", other_row: int
    ) -> None:
        """Take the difference between the current array instance and another
        array instance. Store the result in the specified row of the current
        instance.

        :param self_row: Row from the current array from which to subtract.
        :param other: Other BitSet to subtract from the current instance.
        :param other_row: Row from the other BitSet instance.
        """
        self_offset = self_row * self.row_len
        other_offset = other_row * self.row_len

        for i in range(self.row_len):
            self.data[i + self_offset] &= ~(other.data[i + other_offset])

    def union(self: "BitSet", self_row: int, other: "BitSet", other_row: int) -> None:
        """Take the union between the current array instance and another
        array instance. Store the result in the specified row of the current
        instance.

        :param self_row: Row from the current array with which to union.
        :param other: Other BitSet to union with the current instance.
        :param other_row: Row from the other BitSet instance.
        """
        self_offset = self_row * self.row_len
        other_offset = other_row * self.row_len

        for i in range(self.row_len):
            self.data[i + self_offset] |= other.data[i + other_offset]

    def add(self: "BitSet", row: int, bit: int) -> None:
        """Add a single bit to the row of a bit array

        :param row: Row to be modified.
        :param bit: Bit to be added.
        """
        offset = row * self.row_len
        i = bit // self.CHUNK_SIZE
        self.data[i + offset] |= self.DTYPE(1) << (bit - (self.CHUNK_SIZE * i))

    def get_items(self: "BitSet", row: int) -> Generator[int, None, None]:
        """Get the items stored in the row of a bitset

        :param row: Row from the array to list from.
        :returns: A generator of integers stored in the array.
        """
        offset = row * self.row_len
        for i in range(self.row_len):
            for item in range(self.CHUNK_SIZE):
                if self.data[i + offset] & (self.DTYPE(1) << item):
                    yield item + (i * self.CHUNK_SIZE)

    def contains(self: "BitSet", row: int, bit: int) -> bool:
        """Test if a bit is contained within a bit array row

        :param row: Row to test.
        :param bit: Bit to check.
        :returns: True if the bit is set in the row, else false.
        """
        i = bit // self.CHUNK_SIZE
        offset = row * self.row_len
        return bool(
            self.data[i + offset] & (self.DTYPE(1) << (bit - (self.CHUNK_SIZE * i)))
        )

    def count(self: "BitSet", row: int) -> int:
        """Count all of the set bits in a specified row. Uses a SWAR
        algorithm to count in parallel with a constant number (12) of operations.

        NB: we have to cast all values to our unsigned dtype to avoid type promotion

        Details here:
        # https://graphics.stanford.edu/~seander/bithacks.html#CountBitsSetParallel

        :param row: Row to count.
        :returns: Count of all of the set bits.
        """
        count = 0
        offset = row * self.row_len
        D = self.DTYPE

        for i in range(offset, offset + self.row_len):
            v = self.data[i]
            v = v - ((v >> D(1)) & D(0x55555555))
            v = (v & D(0x33333333)) + ((v >> D(2)) & D(0x33333333))
            # this operation relies on integer overflow
            with np.errstate(over="ignore"):
                count += ((v + (v >> D(4)) & D(0xF0F0F0F)) * D(0x1010101)) >> D(24)

        return count

    def count_naive(self: "BitSet", row: int) -> int:
        """Naive counting algorithm implementing the same functionality as the count
        method. Useful for testing correctness, uses the same number of operations
        as set bits.

        Details here:
        # https://graphics.stanford.edu/~seander/bithacks.html#CountBitsSetNaive

        :param row: Row to count.
        :returns: Count of all of the set bits.
        """
        count = 0
        offset = row * self.row_len

        for i in range(offset, offset + self.row_len):
            v = self.data[i]
            while v:
                v &= v - self.DTYPE(1)
                count += self.DTYPE(1)
        return count


def norm_hap_weighted(
    state_dim: int,
    hap_weights: np.ndarray,
    n_a: int,
    n_b: int,
    result: np.ndarray,
    params: Dict[str, Any],
) -> None:
    """Create a vector of normalizing coefficients, length of the number of
    sample sets. In this normalization strategy, we weight each allele's
    statistic by the proportion of the haplotype present.

    :param state_dim: Number of sample sets.
    :param hap_weights: Proportion of each two-locus haplotype.
    :param n_a: Number of alleles at the A locus.
    :param n_b: Number of alleles at the B locus.
    :param result: Result vector to store the normalizing coefficients in.
    :param params: Params of summary function.
    """
    del n_a, n_b  # handle unused params
    sample_set_sizes = params["sample_set_sizes"]
    for k in range(state_dim):
        n = sample_set_sizes[k]
        result[k] = hap_weights[0, k] / n


def norm_total_weighted(
    state_dim: int,
    hap_weights: np.ndarray,
    n_a: int,
    n_b: int,
    result: np.ndarray,
    params: Dict[str, Any],
) -> None:
    """Create a vector of normalizing coefficients, length of the number of
    sample sets. In this normalization strategy, we weight each allele's
    statistic by the product of the allele frequencies

    :param state_dim: Number of sample sets.
    :param hap_weights: Proportion of each two-locus haplotype.
    :param n_a: Number of alleles at the A locus.
    :param n_b: Number of alleles at the B locus.
    :param result: Result vector to store the normalizing coefficients in.
    :param params: Params of summary function.
    """
    del hap_weights, params  # handle unused params
    for k in range(state_dim):
        result[k] = 1 / (n_a * n_b)


def check_order_bounds_dups(values, max_value):
    """Validate the specified values.

    We require that values are:

    1) Within the boundaries of the max value in the tree sequence
    2) Sorted
    3) Non-repeating

    Raises an exception if any error is found.

    :param sites: 1d array of values to validate.
    :param max_sites: The upper bound for the provided values.
    """
    if len(values) == 0:
        return
    i = 0
    for i in range(len(values) - 1):
        if values[i] < 0 or values[i] >= max_value:
            raise ValueError(f"Value out of bounds: {values[i]}")
        if values[i] >= values[i + 1]:
            raise ValueError(f"Value not sorted: {values[i], values[i + 1]}")
    if values[-1] < 0 or values[-1] >= max_value:
        raise ValueError(f"Value out of bounds: {values[i + 1]}")


def get_site_row_col_indices(
    row_sites: np.ndarray, col_sites: np.ndarray
) -> Tuple[List[int], List[int], List[int]]:
    """Co-iterate over the row and column sites, keeping a sorted union of
    site values and an index into the unique list of sites for both the row
    and column sites. This function produces a list of sites of interest and
    row and column indexes into this list of sites.

    NB: This routine requires that the site lists are sorted and deduplicated.

    :param row_sites: List of sites that will be represented in the output
                      matrix rows.
    :param col_sites: List of sites that will be represented in the output
                      matrix columns.
    :returns: Tuple of lists of sites, row, and column indices.
    """
    r = 0
    c = 0
    s = 0
    sites = []
    col_idx = []
    row_idx = []

    while r < len(row_sites) and c < len(col_sites):
        if row_sites[r] < col_sites[c]:
            sites.append(row_sites[r])
            row_idx.append(s)
            s += 1
            r += 1
        elif row_sites[r] > col_sites[c]:
            sites.append(col_sites[c])
            col_idx.append(s)
            s += 1
            c += 1
        else:
            sites.append(row_sites[r])
            row_idx.append(s)
            col_idx.append(s)
            s += 1
            r += 1
            c += 1
    while r < len(row_sites):
        sites.append(row_sites[r])
        row_idx.append(s)
        s += 1
        r += 1
    while c < len(col_sites):
        sites.append(col_sites[c])
        col_idx.append(s)
        s += 1
        c += 1

    return sites, row_idx, col_idx


def get_all_samples_bits(num_samples: int) -> BitSet:
    """Get the bits for all samples in the tree sequence. This is achieved
    by creating a length 1 bitset and adding every sample's bit to it.

    :param num_samples: Number of samples contained in the tree sequence.
    :returns: Length 1 BitSet containing all samples in the tree sequence.
    """
    all_samples = BitSet(num_samples, 1)
    for i in range(num_samples):
        all_samples.add(0, i)
    return all_samples


def get_allele_samples(
    site: tskit.Site, site_offset: int, mut_samples: BitSet, allele_samples: BitSet
) -> int:
    """Given a BitSet that has been arranged so that we have every sample under
    a given mutation's node, create the final output where we know which samples
    should belong under each mutation, considering the mutation's parentage,
    back mutations, and ancestral state.

    To this end, we iterate over each mutation and store the samples under the
    focal mutation in the output BitSet (allele_samples). Then, we check the
    parent of the focal mutation (either a mutation or the ancestral allele),
    and we subtract the samples in the focal mutation from the parent allele's
    samples.

    :param site: Focal site for which to adjust mutation data.
    :param site_offset: Offset into allele_samples for our focal site.
    :param mut_samples: BitSet containing the samples under each mutation in the
                        focal site.
    :param allele_samples: Output BitSet, initially passed in with all of the
                           tree sequence samples set in the ancestral allele
                           state.
    :returns: number of alleles actually encountered (adjusting for back-mutation).
    """
    alleles = []
    num_alleles = 1
    alleles.append(site.ancestral_state)

    for m, mut in enumerate(site.mutations):
        try:
            allele = alleles.index(mut.derived_state)
        except ValueError:
            allele = len(alleles)
            alleles.append(mut.derived_state)
            num_alleles += 1
        allele_samples.union(allele + site_offset, mut_samples, m)
        # now to find the parent allele from which we must subtract
        alt_allele_state = site.ancestral_state
        if mut.parent != tskit.NULL:
            parent_mut = site.mutations[mut.parent - site.mutations[0].id]
            alt_allele_state = parent_mut.derived_state
        alt_allele = alleles.index(alt_allele_state)
        # subtract focal allele's samples from the alt allele
        allele_samples.difference(
            alt_allele + site_offset, allele_samples, allele + site_offset
        )

    return num_alleles


def get_mutation_samples(
    ts: tskit.TreeSequence, sites: List[int], sample_index_map: np.ndarray
) -> Tuple[np.ndarray, np.ndarray, BitSet]:
    """For a given set of sites, generate a BitSet of all samples posessing
    each allelic state for each site. This includes the ancestral state, along
    with any mutations contained in the site.

    We achieve this goal by starting at the tree containing the first site in
    our list, then we walk along each tree until we've encountered the last
    tree containing the last site in our list. Along the way, we perform a
    preorder traversal from the node of each mutation in a given site, storing
    the samples under that particular node. After we've stored all of the samples
    for each allele at a site, we adjust each allele's samples by removing
    samples that have a different allele at a child mutation down the tree (see
    get_allele_samples for more details).

    We also gather some ancillary data while we iterate over the sites: the
    number of alleles for each site, and the offset of each site. The number of
    alleles at each site includes the count of mutations + the ancestral allele.
    The offeset for each site indicates how many array entries we must skip (ie
    how many alleles exist before a specific site's entry) in order to address
    the data for a given site.

    :param ts: Tree sequence to gather data from.
    :param sites: Subset of sites to consider when gathering data.
    :returns: Tuple of the number of alleles per site, site offsets, and the
              BitSet of all samples in each allelic state.
    """
    num_alleles = np.zeros(len(sites), dtype=np.uint64)
    site_offsets = np.zeros(len(sites), dtype=np.uint64)
    all_samples = get_all_samples_bits(ts.num_samples)
    allele_samples = BitSet(
        ts.num_samples, sum(len(ts.site(i).mutations) + 1 for i in sites)
    )

    site_offset = 0
    site_idx = 0
    for site_idx, site_id in enumerate(sites):
        site = ts.site(site_id)
        tree = ts.at(site.position)
        # initialize the ancestral allele with all samples
        allele_samples.union(site_offset, all_samples, 0)
        # store samples for each mutation in mut_samples
        mut_samples = BitSet(ts.num_samples, len(site.mutations))
        for m, mut in enumerate(site.mutations):
            for node in tree.preorder(mut.node):
                if ts.node(node).is_sample():
                    mut_samples.add(m, sample_index_map[node])
        # account for mutation parentage, subtract samples from mutation parents
        num_alleles[site_idx] = get_allele_samples(
            site, site_offset, mut_samples, allele_samples
        )
        # increment the offset for ancestral + mutation alleles
        site_offsets[site_idx] = site_offset
        site_offset += len(site.mutations) + 1

    return num_alleles, site_offsets, allele_samples


def compute_general_two_site_stat_result(
    row_site_offset: int,
    col_site_offset: int,
    num_row_alleles: int,
    num_col_alleles: int,
    num_samples: int,
    allele_samples: BitSet,
    state_dim: int,
    sample_sets: BitSet,
    func: Callable[[int, np.ndarray, np.ndarray, Dict[str, Any]], None],
    norm_func: Callable[[int, np.ndarray, int, int, np.ndarray, Dict[str, Any]], None],
    params: Dict[str, Any],
    polarised: bool,
    result: np.ndarray,
) -> None:
    """For a given pair of sites, compute the summary statistic for the allele
    frequencies for each allelic state of the two pairs.

    :param row_site_offset: Offset of the row site's data in the allele_samples.
    :param row_site_offset: Offset of the col site's data in the allele_samples.
    :param num_row_alleles: Number of alleles in the row site.
    :param num_col_alleles: Number of alleles in the col site.
    :param num_samples: Number of samples in tree sequence.
    :param allele_samples: BitSet containing the samples with each allelic state
                           for each site of interest.
    :param state_dim: Number of sample sets.
    :param sample_sets: BitSet of sample sets to be intersected with the samples
                        contained within each allele.
    :param func: Summary function used to compute each two-locus statistic.
    :param norm_func: Function used to generate the normalization coefficients
                      for each statistic.
    :param params: Parameters to pass to the norm and summary function.
    :param polarised: If true, skip the computation of the statistic for the
                      ancestral state.
    :param result: Vector of the results matrix to populate. We will produce one
                   value per sample set, hence the vector of length state_dim.
    """
    ss_A_samples = BitSet(num_samples, 1)
    ss_B_samples = BitSet(num_samples, 1)
    ss_AB_samples = BitSet(num_samples, 1)
    AB_samples = BitSet(num_samples, 1)
    weights = np.zeros((3, state_dim), np.float64)
    norm = np.zeros(state_dim, np.float64)
    result_tmp = np.zeros(state_dim, np.float64)

    polarised_val = 1 if polarised else 0

    for mut_a in range(polarised_val, num_row_alleles):
        a = int(mut_a + row_site_offset)
        for mut_b in range(polarised_val, num_col_alleles):
            b = int(mut_b + col_site_offset)
            allele_samples.intersect(a, allele_samples, b, AB_samples)
            for k in range(state_dim):
                allele_samples.intersect(a, sample_sets, k, ss_A_samples)
                allele_samples.intersect(b, sample_sets, k, ss_B_samples)
                AB_samples.intersect(0, sample_sets, k, ss_AB_samples)

                w_AB = ss_AB_samples.count(0)
                w_A = ss_A_samples.count(0)
                w_B = ss_B_samples.count(0)

                weights[0, k] = w_AB
                weights[1, k] = w_A - w_AB  # w_Ab
                weights[2, k] = w_B - w_AB  # w_aB

            func(state_dim, weights, result_tmp, params)

            norm_func(
                state_dim,
                weights,
                num_row_alleles - polarised_val,
                num_col_alleles - polarised_val,
                norm,
                params,
            )

            for k in range(state_dim):
                result[k] += result_tmp[k] * norm[k]


def two_site_count_stat(
    ts: tskit.TreeSequence,
    func: Callable[[int, np.ndarray, np.ndarray, Dict[str, Any]], None],
    norm_func: Callable[[int, np.ndarray, int, int, np.ndarray, Dict[str, Any]], None],
    num_sample_sets: int,
    sample_set_sizes: np.ndarray,
    sample_sets: BitSet,
    sample_index_map: np.ndarray,
    row_sites: np.ndarray,
    col_sites: np.ndarray,
    polarised: bool,
) -> np.ndarray:
    """Outer function that generates the high-level intermediates used in the
    computation of our two-locus statistics. First, we compute the row and
    column indices for our unique list of sites, then we get each sample for
    each allele in our list of specified sites.

    With those intermediates in hand, we iterate over the row and column indices
    to compute comparisons between each of the specified lists of sites. We pass
    a vector of results to the computation, which will compute a single result
    for each sample set, inserting that into our result matrix.

    :param ts: Tree sequence to gather data from.
    :param func: Function used to compute each two-locus statistic.
    :param norm_func: Function used to generate the normalization coefficients
                      for each statistic.
    :param num_sample_sets: Number of sample sets that we will consider.
    :param sample_set_sizes: Number of samples in each sample set.
    :param sample_sets: BitSet of samples to compute stats for. We will only
                        consider these samples in our computations, resulting
                        in stats that are computed on subsets of the samples
                        on the tree sequence.
    :param row_sites: Sites contained in the rows of the output matrix.
    :param col_sites: Sites contained in the columns of the output matrix.
    :param polarised: If true, skip the computation of the statistic for the
                      ancestral state.
    :returns: 3D array of results, dimensions (sample_sets, row_sites, col_sites).
    """
    params = {"sample_set_sizes": sample_set_sizes}
    result = np.zeros(
        (num_sample_sets, len(row_sites), len(col_sites)), dtype=np.float64
    )

    sites, row_idx, col_idx = get_site_row_col_indices(row_sites, col_sites)
    num_alleles, site_offsets, allele_samples = get_mutation_samples(
        ts, sites, sample_index_map
    )

    for row, row_site in enumerate(row_idx):
        for col, col_site in enumerate(col_idx):
            compute_general_two_site_stat_result(
                site_offsets[row_site],
                site_offsets[col_site],
                num_alleles[row_site],
                num_alleles[col_site],
                ts.num_samples,
                allele_samples,
                num_sample_sets,
                sample_sets,
                func,
                norm_func,
                params,
                polarised,
                result[:, row, col],
            )

    return result


def get_index_repeats(indices):
    """In a list of indices, find the repeat values. The first value
    is offset by the first index and ranges to the last index.
    For instance, [4, 4, 5, 6, 8] becomes [2, 1, 1, 0, 1].
    The list must be sorted and ordered.

    :param indices: List of indices to count
    :returns: Counts of index repeats
    """
    counts = np.zeros(indices[-1] - indices[0] + 1, dtype=np.int32)
    idx = indices[0]
    count = 1
    for i in range(1, len(indices)):
        if indices[i] == indices[i - 1]:
            count += 1
        else:
            counts[idx - indices[0]] = count
            count = 1
            idx = indices[i]
    counts[idx - indices[0]] = count
    return counts


def two_branch_count_stat(
    ts: tskit.TreeSequence,
    func: Callable[[int, np.ndarray, np.ndarray, Dict[str, Any]], None],
    norm_func,
    num_sample_sets: int,
    sample_set_sizes: np.ndarray,
    sample_sets: BitSet,
    sample_index_map: np.ndarray,
    row_trees: np.ndarray,
    col_trees: np.ndarray,
    polarised: bool,
) -> np.ndarray:
    """
    Compute a tree X tree LD matrix by walking along the tree sequence and
    computing haplotype counts. This method incrementally adds and removes
    branches from a tree sequence and updates the stat based on sample additions
    and removals. We bifurcate the tree with a given branch on each locus and
    intersect the samples under each branch to produce haplotype counts. It
    is possible to subset the output matrix with genomic positions. Positions
    lying on the same tree will receive the same LD value in the output matrix.

    :param ts: Tree sequence to gather data from.
    :param func: Function used to compute each two-locus statistic.
    :param norm_func: Not (YET) applicable for branch stats: TODO?
    :param num_sample_sets: Number of sample sets that we will consider.
    :param sample_set_sizes: Number of samples in each sample set.
    :param sample_sets: BitSet of samples to compute stats for. We will only
                        consider these samples in our computations, resulting
                        in stats that are computed on subsets of the samples
                        on the tree sequence.
    :param row_trees: Trees contained in the rows of the output matrix (repeats ok)
    :param col_trees: Trees contained in the rows of the output matrix (repeats ok)
    :param polarised: If true, skip the computation of the statistic for the
                      ancestral state.
    :returns: 3D array of results, dimensions (sample_sets, row_sites, col_sites).
    """
    params = {"sample_set_sizes": sample_set_sizes}
    result = np.zeros(
        (num_sample_sets, len(row_trees), len(col_trees)), dtype=np.float64
    )
    row_repeats = get_index_repeats(row_trees)
    col_repeats = get_index_repeats(col_trees)

    stat = np.zeros(num_sample_sets, dtype=np.float64)
    # State is initialized at tree -1
    l_state = TreeState(ts, sample_sets, num_sample_sets, sample_index_map)
    r_state = TreeState(ts, sample_sets, num_sample_sets, sample_index_map)

    # Even if we're skipping trees, we must iterate over the range to keep the
    # running total of the statistic consistent.
    row = 0
    for r in range(row_trees[-1] + 1 - row_trees[0]):
        # zero out stat and r_state at the beginning of each row
        stat = np.zeros_like(stat)
        r_state = TreeState(ts, sample_sets, num_sample_sets, sample_index_map)
        l_state.advance(r + row_trees[0])
        # use null TreeState to advance l_state, conveniently we just zerod r_state
        _, l_state = compute_branch_stat(
            ts, func, stat, params, num_sample_sets, r_state, l_state
        )
        col = 0
        for c in range(col_trees[-1] + 1 - col_trees[0]):
            r_state.advance(c + col_trees[0])
            stat, r_state = compute_branch_stat(
                ts, func, stat, params, num_sample_sets, l_state, r_state
            )
            # Fill in repeated values for all sample sets
            for i in range(row_repeats[r]):
                for j in range(col_repeats[c]):
                    result[:, i + row, j + col] = stat
            col += col_repeats[c]
        row += row_repeats[r]
    return result


def sample_sets_to_bit_array(
    ts: tskit.TreeSequence, sample_sets: List[List[int]]
) -> Tuple[np.ndarray, np.ndarray, BitSet]:
    """Convert the list of sample ids to a bit array. This function takes
    sample identifiers and maps them to their enumerated integer values, then
    stores these values in a bit array. We produce a BitArray and a numpy
    array of integers that specify how many samples there are in each sample set.

    NB: this function's type signature is of type integer, but I believe this
        could be expanded to Any, currently untested so the integer
        specification remains.

    :param ts: Tree sequence to gather data from.
    :param sample_sets: List of sample identifiers to store in bit array.
    :returns: Tuple containing numpy array of sample set sizes and the sample
              set BitSet.
    """
    sample_sets_bits = BitSet(ts.num_samples, len(sample_sets))
    sample_index_map = -np.ones(ts.num_nodes, dtype=np.int32)
    sample_set_sizes = np.zeros(len(sample_sets), dtype=np.uint64)

    sample_count = 0
    for node in ts.nodes():
        if node.flags & tskit.NODE_IS_SAMPLE:
            sample_index_map[node.id] = sample_count
            sample_count += 1

    for k, sample_set in enumerate(sample_sets):
        sample_set_sizes[k] = len(sample_set)
        for sample in sample_set:
            sample_index = sample_index_map[sample]
            if sample_index == tskit.NULL:
                raise ValueError(f"Sample out of bounds: {sample}")
            if sample_sets_bits.contains(k, sample_index):
                raise ValueError(f"Duplicate sample detected: {sample}")
            sample_sets_bits.add(k, sample_index)

    return sample_index_map, sample_set_sizes, sample_sets_bits


def positions_to_tree_indices(bp, positions):
    """Given a set of breakpoints and positions, provide an array of tree
    indices that correspond with positions. We have already validated that the
    bounds of the positions are correct and that they are sorted, and
    deduplicated.

    :param bp: Breakpoints of the tree sequence
    :param positions: Positions to search over
    :returns: Array of tree indices
    """
    tree_idx = 0
    tree_indices = -np.ones_like(positions, dtype=np.int32)

    for i in range(len(positions)):
        while bp[tree_idx + 1] <= positions[i]:
            tree_idx += 1
        tree_indices[i] = tree_idx

    return tree_indices


def two_locus_count_stat(
    ts,
    summary_func,
    norm_func,
    polarised,
    mode,
    sites=None,
    positions=None,
    sample_sets=None,
):
    """Outer wrapper for two site general stat functionality. Perform some input
    validation, get the site index and allele state, then compute the LD matrix.

    :param ts: Tree sequence to gather data from.
    :param summary_func: Function used to compute each two-locus statistic.
    :param norm_func: Function used to generate the normalization coefficients
                      for each statistic.
    :param polarised: If true, skip the computation of the statistic for the
                      ancestral state.
    :param sites: List of two lists containing [row_sites, column_sites].
    :param positions: List of two lists containing [row_positions, col_positions],
                      which are genomic positions to compute LD on.
    :param sample_sets: List of lists of samples to compute stats for. We will
                        only consider these samples in our computations,
                        resulting in stats that are computed on subsets of the
                        samples on the tree sequence.
    :returns: 3d numpy array containing LD for (sample_set,row_site,column_site)
              unless one or no sample sets are specified, then 2d array
              containing LD for (row_site,column_site).
    """
    if sample_sets is None:
        sample_sets = [ts.samples()]

    sample_index_map, ss_sizes, ss_bits = sample_sets_to_bit_array(ts, sample_sets)

    if mode == "site":
        if positions is not None:
            raise ValueError("Cannot specify positions in site mode")
        if sites is None:
            row_sites = np.arange(ts.num_sites, dtype=np.int32)
            col_sites = np.arange(ts.num_sites, dtype=np.int32)
        elif len(sites) == 2:
            row_sites = np.asarray(sites[0])
            col_sites = np.asarray(sites[1])
        elif len(sites) == 1:
            row_sites = np.asarray(sites[0])
            col_sites = row_sites
        else:
            raise ValueError(
                f"Sites must be a length 1 or 2 list, got a length {len(sites)} list"
            )
        check_order_bounds_dups(row_sites, ts.num_sites)
        check_order_bounds_dups(col_sites, ts.num_sites)
        result = two_site_count_stat(
            ts,
            summary_func,
            norm_func,
            len(ss_sizes),
            ss_sizes,
            ss_bits,
            sample_index_map,
            row_sites,
            col_sites,
            polarised,
        )
    elif mode == "branch":
        if sites is not None:
            raise ValueError("Cannot specify sites in branch mode")
        if positions is None:
            row_trees = np.arange(ts.num_trees, dtype=np.int32)
            col_trees = np.arange(ts.num_trees, dtype=np.int32)
        elif len(positions) == 2:
            breakpoints = ts.breakpoints(as_array=True)
            row_positions = np.asarray(positions[0])
            col_positions = np.asarray(positions[1])
            check_order_bounds_dups(row_positions, breakpoints[-1])
            check_order_bounds_dups(col_positions, breakpoints[-1])
            row_trees = positions_to_tree_indices(breakpoints, row_positions)
            col_trees = positions_to_tree_indices(breakpoints, col_positions)
        elif len(positions) == 1:
            breakpoints = ts.breakpoints(as_array=True)
            row_positions = np.asarray(positions[0])
            col_positions = row_positions
            check_order_bounds_dups(row_positions, breakpoints[-1])
            row_trees = positions_to_tree_indices(breakpoints, row_positions)
            col_trees = row_trees
        else:
            raise ValueError(
                "Positions must be a length 1 or 2 list, "
                f"got a length {len(positions)} list"
            )
        result = two_branch_count_stat(
            ts,
            summary_func,
            None,
            len(ss_sizes),
            ss_sizes,
            ss_bits,
            sample_index_map,
            row_trees,
            col_trees,
            False,
        )
    else:
        raise ValueError(f"Unknown mode: {mode}")

    # If there is one sample set, return a 2d numpy array of row/site LD
    if len(sample_sets) == 1:
        return result.reshape(result.shape[1:3])
    return result


def r2_summary_func(
    state_dim: int, state: np.ndarray, result: np.ndarray, params: Dict[str, Any]
) -> None:
    """Summary function for the r2 statistic. We first compute the proportion of
    AB, A, and B haplotypes, then we compute the r2 statistic, storing the outputs
    in the result vector, one entry per sample set.

    :param state_dim: Number of sample sets.
    :param state: Counts of 3 haplotype configurations for each sample set.
    :param result: Vector of length state_dim to store the results in.
    :param params: Parameters for the summary function.
    """
    sample_set_sizes = params["sample_set_sizes"]
    for k in range(state_dim):
        n = sample_set_sizes[k]
        p_AB = state[0, k] / n
        p_Ab = state[1, k] / n
        p_aB = state[2, k] / n

        p_A = p_AB + p_Ab
        p_B = p_AB + p_aB

        D = p_AB - (p_A * p_B)
        denom = p_A * p_B * (1 - p_A) * (1 - p_B)

        with suppress_overflow_div0_warning():
            result[k] = (D * D) / denom


def D_summary_func(
    state_dim: int, state: np.ndarray, result: np.ndarray, params: Dict[str, Any]
) -> None:
    sample_set_sizes = params["sample_set_sizes"]
    for k in range(state_dim):
        n = sample_set_sizes[k]
        p_AB = state[0, k] / float(n)
        p_Ab = state[1, k] / float(n)
        p_aB = state[2, k] / float(n)

        p_A = p_AB + p_Ab
        p_B = p_AB + p_aB

        result[k] = p_AB - (p_A * p_B)


def D2_summary_func(
    state_dim: int, state: np.ndarray, result: np.ndarray, params: Dict[str, Any]
) -> None:
    sample_set_sizes = params["sample_set_sizes"]
    for k in range(state_dim):
        n = sample_set_sizes[k]
        p_AB = state[0, k] / float(n)
        p_Ab = state[1, k] / float(n)
        p_aB = state[2, k] / float(n)

        p_A = p_AB + p_Ab
        p_B = p_AB + p_aB

        result[k] = p_AB - (p_A * p_B)
        result[k] = result[k] * result[k]


def D_prime_summary_func(
    state_dim: int, state: np.ndarray, result: np.ndarray, params: Dict[str, Any]
) -> None:
    sample_set_sizes = params["sample_set_sizes"]
    for k in range(state_dim):
        n = sample_set_sizes[k]
        p_AB = state[0, k] / float(n)
        p_Ab = state[1, k] / float(n)
        p_aB = state[2, k] / float(n)

        p_A = p_AB + p_Ab
        p_B = p_AB + p_aB

        D = p_AB - (p_A * p_B)
        with suppress_overflow_div0_warning():
            if D >= 0:
                result[k] = D / min(p_A * (1 - p_B), (1 - p_A) * p_B)
            else:
                result[k] = D / min(p_A * p_B, (1 - p_A) * (1 - p_B))


def r_summary_func(
    state_dim: int, state: np.ndarray, result: np.ndarray, params: Dict[str, Any]
) -> None:
    sample_set_sizes = params["sample_set_sizes"]
    for k in range(state_dim):
        n = sample_set_sizes[k]
        p_AB = state[0, k] / n
        p_Ab = state[1, k] / n
        p_aB = state[2, k] / n

        p_A = p_AB + p_Ab
        p_B = p_AB + p_aB

        D = p_AB - (p_A * p_B)
        denom = p_A * p_B * (1 - p_A) * (1 - p_B)

        with suppress_overflow_div0_warning():
            result[k] = D / np.sqrt(denom)


def Dz_summary_func(
    state_dim: int, state: np.ndarray, result: np.ndarray, params: Dict[str, Any]
) -> None:
    sample_set_sizes = params["sample_set_sizes"]
    for k in range(state_dim):
        n = sample_set_sizes[k]
        p_AB = state[0, k] / n
        p_Ab = state[1, k] / n
        p_aB = state[2, k] / n

        p_A = p_AB + p_Ab
        p_B = p_AB + p_aB

        D = p_AB - (p_A * p_B)

        result[k] = D * (1 - 2 * p_A) * (1 - 2 * p_B)


def pi2_summary_func(
    state_dim: int, state: np.ndarray, result: np.ndarray, params: Dict[str, Any]
) -> None:
    sample_set_sizes = params["sample_set_sizes"]
    for k in range(state_dim):
        n = sample_set_sizes[k]
        p_AB = state[0, k] / n
        p_Ab = state[1, k] / n
        p_aB = state[2, k] / n

        p_A = p_AB + p_Ab
        p_B = p_AB + p_aB

        result[k] = p_A * (1 - p_A) * p_B * (1 - p_B)


# Unbiased estimators of pi2, dz, and d2. These are derived in Ragsdale 2019
# (https://doi.org/10.1093/molbev/msz265) and can be used in place of the method
# outlined by McVean 2002. The reason for using haplotype counts in the branch
# methods is that we can compute statistics that cannot be represented by tMRCA
# covariance. With these unbiased estimators, we still reproduce the values
# estimated with tMRCA covariance.

# TODO: update these summary functions to have the same function signature as
#       the summary functions defined above.


def pi2_unbiased(
    state_dim: int, state: np.ndarray, result: np.ndarray, params: Dict[str, Any]
):
    sample_set_sizes = params["sample_set_sizes"]
    for k in range(state_dim):
        n = sample_set_sizes[k]
        w_AB = state[0, k]
        w_Ab = state[1, k]
        w_aB = state[2, k]
        w_ab = n - (w_AB + w_Ab + w_aB)
        with suppress_overflow_div0_warning():
            result[k] = (1 / (n * (n - 1) * (n - 2) * (n - 3))) * (
                ((w_AB + w_Ab) * (w_aB + w_ab) * (w_AB + w_aB) * (w_Ab + w_ab))
                - ((w_AB * w_ab) * (w_AB + w_ab + (3 * w_Ab) + (3 * w_aB) - 1))
                - ((w_Ab * w_aB) * (w_Ab + w_aB + (3 * w_AB) + (3 * w_ab) - 1))
            )


def dz_unbiased(
    state_dim: int, state: np.ndarray, result: np.ndarray, params: Dict[str, Any]
):
    sample_set_sizes = params["sample_set_sizes"]
    for k in range(state_dim):
        n = sample_set_sizes[k]
        w_AB = state[0, k]
        w_Ab = state[1, k]
        w_aB = state[2, k]
        w_ab = n - (w_AB + w_Ab + w_aB)
        with suppress_overflow_div0_warning():
            result[k] = (1 / (n * (n - 1) * (n - 2) * (n - 3))) * (
                (
                    ((w_AB * w_ab) - (w_Ab * w_aB))
                    * (w_aB + w_ab - w_AB - w_Ab)
                    * (w_Ab + w_ab - w_AB - w_aB)
                )
                - ((w_AB * w_ab) * (w_AB + w_ab - w_Ab - w_aB - 2))
                - ((w_Ab * w_aB) * (w_Ab + w_aB - w_AB - w_ab - 2))
            )


def d2_unbiased(
    state_dim: int, state: np.ndarray, result: np.ndarray, params: Dict[str, Any]
):
    sample_set_sizes = params["sample_set_sizes"]
    for k in range(state_dim):
        n = sample_set_sizes[k]
        w_AB = state[0, k]
        w_Ab = state[1, k]
        w_aB = state[2, k]
        w_ab = n - (w_AB + w_Ab + w_aB)
        with suppress_overflow_div0_warning():
            result[k] = (1 / (n * (n - 1) * (n - 2) * (n - 3))) * (
                ((w_aB**2) * (w_Ab - 1) * w_Ab)
                + ((w_ab - 1) * w_ab * (w_AB - 1) * w_AB)
                - (w_aB * w_Ab * (w_Ab + (2 * w_ab * w_AB) - 1))
            )


SUMMARY_FUNCS = {
    "r": r_summary_func,
    "r2": r2_summary_func,
    "D": D_summary_func,
    "D2": D2_summary_func,
    "D_prime": D_prime_summary_func,
    "pi2": pi2_summary_func,
    "Dz": Dz_summary_func,
    "D2_unbiased": d2_unbiased,
    "Dz_unbiased": dz_unbiased,
    "pi2_unbiased": pi2_unbiased,
}

NORM_METHOD = {
    D_summary_func: norm_total_weighted,
    D_prime_summary_func: norm_hap_weighted,
    D2_summary_func: norm_total_weighted,
    Dz_summary_func: norm_total_weighted,
    pi2_summary_func: norm_total_weighted,
    r_summary_func: norm_total_weighted,
    r2_summary_func: norm_hap_weighted,
    d2_unbiased: norm_total_weighted,
    dz_unbiased: norm_total_weighted,
    pi2_unbiased: norm_total_weighted,
}

POLARIZATION = {
    D_summary_func: True,
    D_prime_summary_func: True,
    D2_summary_func: False,
    Dz_summary_func: False,
    pi2_summary_func: False,
    r_summary_func: True,
    r2_summary_func: False,
    d2_unbiased: False,
    dz_unbiased: False,
    pi2_unbiased: False,
}


def ld_matrix(ts, sample_sets=None, sites=None, positions=None, stat="r2", mode="site"):
    summary_func = SUMMARY_FUNCS[stat]
    return two_locus_count_stat(
        ts,
        summary_func,
        NORM_METHOD[summary_func],
        POLARIZATION[summary_func],
        mode,
        sites=sites,
        positions=positions,
        sample_sets=sample_sets,
    )


def get_paper_ex_ts():
    """Generate the tree sequence example from the tskit paper

    Data taken from the tests:
    https://github.com/tskit-dev/tskit/blob/61a844a/c/tests/testlib.c#L55-L96

    :returns: Tree sequence
    """
    nodes = """\
    is_sample time population individual
    1  0       -1   0
    1  0       -1   0
    1  0       -1   1
    1  0       -1   1
    0  0.071   -1   -1
    0  0.090   -1   -1
    0  0.170   -1   -1
    0  0.202   -1   -1
    0  0.253   -1   -1
    """

    edges = """\
    left   right   parent  child
    2 10 4 2
    2 10 4 3
    0 10 5 1
    0 2  5 3
    2 10 5 4
    0 7  6 0,5
    7 10 7 0,5
    0 2  8 2,6
    """

    sites = """\
    position ancestral_state
    1      0
    4.5    0
    8.5    0
    """

    mutations = """\
    site node derived_state
    0      2   1
    1      0   1
    2      5   1
    """

    individuals = """\
    flags  location   parents
    0      0.2,1.5    -1,-1
    0      0.0,0.0    -1,-1
    """

    return tskit.load_text(
        nodes=io.StringIO(nodes),
        edges=io.StringIO(edges),
        sites=io.StringIO(sites),
        individuals=io.StringIO(individuals),
        mutations=io.StringIO(mutations),
        strict=False,
    )


# fmt:off
# true r2 values for the tree sequence from the tskit paper
PAPER_EX_TRUTH_MATRIX = np.array(
    [[1.0,        0.11111111, 0.11111111],  # noqa: E241
     [0.11111111, 1.0,        1.0],  # noqa: E241
     [0.11111111, 1.0,        1.0]]  # noqa: E241
)
PAPER_EX_BRANCH_TRUTH_MATRIX = np.array(
    [[ 1.06666667e-03, -1.26666667e-04, -1.26666667e-04],  # noqa: E241,E201
     [-1.26666667e-04,  6.01666667e-05,  6.01666667e-05],  # noqa: E241
     [-1.26666667e-04,  6.01666667e-05,  6.01666667e-05]]  # noqa: E241
)
# fmt:on


def get_matrix_partitions(n):
    """Generate all partitions for square matricies, then combine with replacement
    and return all possible pairs of all partitions.

    TODO: only works for square matricies, would need to generate two lists of
    partitions to get around this

    :param n: length of one dimension of the !square! matrix.
    :returns: combinations of partitions.
    """
    parts = []
    for part in tskit.combinatorics.rule_asc(n):
        for g in set(permutations(part, len(part))):
            p = []
            i = iter(range(n))
            for item in g:
                p.append([next(i) for _ in range(item)])
            parts.append(p)
    combos = []
    for a, b in combinations_with_replacement({tuple(j) for i in parts for j in i}, 2):
        combos.append((a, b))
        combos.append((b, a))
    combos = [[list(a), list(b)] for a, b in set(combos)]
    return combos


# Generate all partitions of the LD matrix, then pass into test_subset
@pytest.mark.parametrize("partition", get_matrix_partitions(len(PAPER_EX_TRUTH_MATRIX)))
def test_subset_sites(partition):
    """Given a partition of the truth matrix, check that we can successfully
    compute the LD matrix for that given partition, effectively ensuring that
    our handling of site subsets is correct.

    :param partition: length 2 list of [row_sites, column_sites].
    """
    a, b = partition
    ts = get_paper_ex_ts()
    np.testing.assert_allclose(
        ld_matrix(ts, sites=partition),
        PAPER_EX_TRUTH_MATRIX[a[0] : a[-1] + 1, b[0] : b[-1] + 1],
    )
    np.testing.assert_equal(
        ld_matrix(ts, sites=partition), ts.ld_matrix(sites=partition)
    )


@pytest.mark.parametrize(
    "partition", get_matrix_partitions(len(PAPER_EX_BRANCH_TRUTH_MATRIX))
)
def test_subset_positions(partition):
    """Given a partition of the truth matrix, check that we can successfully
    compute the LD matrix for that given partition, effectively ensuring that
    our handling of positions is correct. We use the midpoint inside of the
    tree interval as the position for a particular tree.

    :param partition: length 2 list of [row_positions, column_positions].
    """
    a, b = partition
    ts = get_paper_ex_ts()
    bp = ts.breakpoints(as_array=True)
    mid = (bp[1:] + bp[:-1]) / 2
    np.testing.assert_allclose(
        ld_matrix(ts, mode="branch", stat="D2_unbiased", positions=[mid[a], mid[b]]),
        PAPER_EX_BRANCH_TRUTH_MATRIX[a[0] : a[-1] + 1, b[0] : b[-1] + 1],
    )
    np.testing.assert_allclose(
        ts.ld_matrix(mode="branch", stat="D2_unbiased", positions=[mid[a], mid[b]]),
        PAPER_EX_BRANCH_TRUTH_MATRIX[a[0] : a[-1] + 1, b[0] : b[-1] + 1],
    )


@pytest.mark.parametrize(
    "positions,truth",
    [
        ([0, 1, 2, 3, 4, 5, 6, 7, 8], [0, 1, 2, 3, 4, 5, 6, 7, 8]),
        ([0], [0]),
        ([8], [8]),
        ([1], [1]),
        ([1, 2, 3], [1, 2, 3]),
        ([], []),
    ],
)
def test_positions_to_tree_indices(positions, truth):
    breakpoints = np.arange(10, dtype=np.float64)
    np.testing.assert_equal(positions_to_tree_indices(breakpoints, positions), truth)


def test_bad_positions():
    with pytest.raises(IndexError, match="out of bounds"):
        breakpoints = np.arange(10, dtype=np.float64)
        positions_to_tree_indices(breakpoints, breakpoints)


@pytest.mark.parametrize("sites", [[0, 1, 2], [1, 2], [0, 1], [0], [1]])
def test_subset_sites_one_list(sites):
    """Test the case where we only pass only one list of sites to compute. This
    should return a square matrix comparing the sites to themselves.
    """
    ts = get_paper_ex_ts()
    np.testing.assert_equal(ld_matrix(ts, sites=[sites]), ts.ld_matrix(sites=[sites]))


@pytest.mark.parametrize("tree_index", [[0, 1, 2], [1, 2], [0, 1], [0], [1]])
def test_subset_positions_one_list(tree_index):
    """Test the case where we only pass only one list of positions to compute. This
    should return a square matrix comparing the positions to themselves.
    """
    ts = get_paper_ex_ts()
    bp = ts.breakpoints(as_array=True)
    mid = (bp[1:] + bp[:-1]) / 2
    np.testing.assert_allclose(
        ld_matrix(ts, mode="branch", stat="D2_unbiased", positions=[mid[tree_index]]),
        PAPER_EX_BRANCH_TRUTH_MATRIX[
            tree_index[0] : tree_index[-1] + 1, tree_index[0] : tree_index[-1] + 1
        ],
    )
    np.testing.assert_allclose(
        ts.ld_matrix(mode="branch", stat="D2_unbiased", positions=[mid[tree_index]]),
        PAPER_EX_BRANCH_TRUTH_MATRIX[
            tree_index[0] : tree_index[-1] + 1, tree_index[0] : tree_index[-1] + 1
        ],
    )


@pytest.mark.parametrize(
    "tree_index",
    [
        ([0, 0, 1, 2], [1, 2]),
        ([0, 0, 0, 2], [0, 2]),
        ([1, 1, 1], [1]),
        ([2, 2], [1]),
        ([0, 2, 2, 2], [0, 0, 0]),
        ([0, 0, 0, 1, 1, 1, 2, 2, 2], [0, 0, 0, 1, 1, 1, 2, 2, 2]),
    ],
)
def test_repeated_position_elements(tree_index):
    """Test that we repeat positions in the LD matrix when we have multiple positions
    that overlap the same tree when specifying positions in branch mode.
    """
    ts = get_paper_ex_ts()
    l, r = tree_index
    bp = ts.breakpoints(as_array=True)
    val, count = np.unique(l, return_counts=True)
    l_pos = np.hstack(
        [np.linspace(bp[v], bp[v + 1], count[i] + 2)[1:-1] for i, v in enumerate(val)]
    )
    val, count = np.unique(r, return_counts=True)
    r_pos = np.hstack(
        [np.linspace(bp[v], bp[v + 1], count[i] + 2)[1:-1] for i, v in enumerate(val)]
    )
    assert (positions_to_tree_indices(bp, l_pos) == l).all()
    assert (positions_to_tree_indices(bp, r_pos) == r).all()

    truth = PAPER_EX_BRANCH_TRUTH_MATRIX[
        [i for i, _ in product(l, r)], [i for _, i in product(l, r)]
    ].reshape(len(l), len(r))

    np.testing.assert_allclose(
        truth,
        ld_matrix(ts, mode="branch", stat="D2_unbiased", positions=[l_pos, r_pos]),
    )
    np.testing.assert_allclose(
        truth,
        ts.ld_matrix(mode="branch", stat="D2_unbiased", positions=[l_pos, r_pos]),
    )


# Generate all partitions of the samples, producing pairs of sample sets
@pytest.mark.parametrize(
    "partition", get_matrix_partitions(get_paper_ex_ts().num_samples)
)
def test_sample_sets(partition):
    """Test all partitions of sample sets, ensuring that we are correctly
    computing stats for various subsets of the samples in a given tree.

    :param partition: length 2 list of [ss_1, ss_2].
    """
    ts = get_paper_ex_ts()
    np.testing.assert_allclose(
        ld_matrix(ts, sample_sets=partition), ts.ld_matrix(sample_sets=partition)
    )


def test_compare_to_ld_calculator():
    ts = msprime.sim_ancestry(
        samples=4, recombination_rate=0.2, sequence_length=10, random_seed=1
    )
    ts = msprime.sim_mutations(ts, rate=0.5, random_seed=1, discrete_genome=False)
    ld_calc = tskit.LdCalculator(ts)
    np.testing.assert_array_almost_equal(ld_calc.get_r2_matrix(), ts.ld_matrix())


@pytest.mark.parametrize(
    "stat",
    sorted(SUMMARY_FUNCS.keys()),
)
def test_multiallelic_with_back_mutation(stat):
    ts = msprime.sim_ancestry(
        samples=4, recombination_rate=0.2, sequence_length=10, random_seed=1
    )
    ts = msprime.sim_mutations(ts, rate=0.5, random_seed=1)
    np.testing.assert_array_almost_equal(
        ld_matrix(ts, stat=stat), ts.ld_matrix(stat=stat)
    )


@pytest.mark.parametrize(
    "ts",
    [
        ts
        for ts in get_example_tree_sequences()
        if ts.id not in {"no_samples", "empty_ts"}
    ],
)
# TODO: port unbiased summary functions
@pytest.mark.parametrize(
    "stat",
    sorted(SUMMARY_FUNCS.keys()),
)
def test_ld_matrix(ts, stat):
    np.testing.assert_array_almost_equal(
        ld_matrix(ts, stat=stat), ts.ld_matrix(stat=stat)
    )


@pytest.mark.parametrize(
    "ts",
    [ts for ts in get_example_tree_sequences() if ts.id in {"no_samples", "empty_ts"}],
)
def test_ld_empty_examples(ts):
    with pytest.raises(ValueError, match="at least one element"):
        ts.ld_matrix()
    with pytest.raises(ValueError, match="at least one element"):
        ts.ld_matrix(mode="branch")


def test_input_validation():
    ts = get_paper_ex_ts()
    with pytest.raises(ValueError, match="Unknown two-locus statistic"):
        ts.ld_matrix(stat="bad_stat")

    with pytest.raises(ValueError, match="must be a list of"):
        ts.ld_matrix(sites=["abc"])
    with pytest.raises(ValueError, match="must be a list of"):
        ts.ld_matrix(sites=[1, 2, 3])
    with pytest.raises(ValueError, match="must be a length 1 or 2 list"):
        ts.ld_matrix(sites=[[1, 2], [2, 3], [3, 4]])
    with pytest.raises(ValueError, match="must be a length 1 or 2 list"):
        ts.ld_matrix(sites=[[1, 2], [2, 3], [3, 4]])
    with pytest.raises(ValueError, match="must be a length 1 or 2 list"):
        ts.ld_matrix(sites=[])

    with pytest.raises(ValueError, match="must be a list of"):
        ts.ld_matrix(positions=["abc"], mode="branch")
    with pytest.raises(ValueError, match="must be a list of"):
        ts.ld_matrix(positions=[1.0, 2.0, 3.0], mode="branch")
    with pytest.raises(ValueError, match="must be a length 1 or 2 list"):
        ts.ld_matrix(positions=[[1.0, 2.0], [2.0, 3.0], [3.0, 4.0]], mode="branch")
    with pytest.raises(ValueError, match="must be a length 1 or 2 list"):
        ts.ld_matrix(positions=[], mode="branch")


@dataclass
class TreeState:
    """
    Class for storing tree state from one iteration to the next. This object
    enables easy copying of the state for computing a matrix.
    """

    pos: tsutil.TreePosition  # current position in the tree sequence
    parent: np.ndarray  # parent node of a given node (connected by an edge)
    branch_len: np.ndarray  # length of the branch above a particular child node
    node_samples: BitSet  # samples that exist under a given node, this is a
    # bitset with a row for each node and sample set. Rows are grouped by node,
    # for example:
    # node sample_set
    # 0    0
    # 0    1
    # 1    0
    # 1    1
    edges_out: List[int]  # list of edges removed during iteration
    edges_in: List[int]  # list of edges added during iteration

    def __init__(self, ts, sample_sets, num_sample_sets, sample_index_map):
        self.pos = tsutil.TreePosition(ts)
        self.parent = -np.ones(ts.num_nodes, dtype=np.int64)
        self.branch_len = np.zeros(ts.num_nodes, dtype=np.float64)
        self.node_samples = BitSet(ts.num_samples, ts.num_nodes * num_sample_sets)
        # Create a bit array to store all samples under each node for each sample set.
        # We initialize with the samples under the sample nodes.
        for n in range(ts.num_nodes):
            for k in range(num_sample_sets):
                if sample_sets.contains(k, sample_index_map[n]):
                    self.node_samples.add(
                        (num_sample_sets * n) + k, sample_index_map[n]
                    )
        # these are empty for the uninitialized state (index = -1)
        self.edges_in = []
        self.edges_out = []

    def advance(self, index):
        """
        Advance tree to next tree position. If the tree is still uninitialized,
        seeks may be performed to an arbitrary position. Since we need to
        compute stats over contiguous ranges of trees, once we've seeked to a
        position, we step forward by one tree. Finally, we set `edges_in` and
        `edges_out` to be consumed by the downstream stats function.

        :param index: Tree index to advance to
        """

        # if initialized or seeking to the first position from the beginning, jump
        # forward one tree
        if self.pos.index != tskit.NULL or index == 0:
            if index != 0:
                assert index == self.pos.index + 1, "only one step allowed"
            assert self.pos.next(), "out of bounds"
            edges_out = [
                self.pos.out_range.order[j]
                for j in range(self.pos.out_range.start, self.pos.out_range.stop)
            ]
            edges_in = [
                self.pos.in_range.order[j]
                for j in range(self.pos.in_range.start, self.pos.in_range.stop)
            ]
            self.edges_out = edges_out
            self.edges_in = edges_in
            return

        # if uninitialized (no current position), and seeking to an arbitrary point
        # in the tree, use seek_forward
        edges_out, edges_in = [], []
        self.pos.seek_forward(index)
        left = self.pos.interval.left
        # since we're starting from an uninitialized tree, we only add edges
        for j in range(self.pos.in_range.start, self.pos.in_range.stop):
            e = self.pos.in_range.order[j]
            # skip over edges that are not in the current tree
            if self.pos.ts.edges_left[e] <= left < self.pos.ts.edges_right[e]:
                edges_in.append(e)

        self.edges_out = edges_out
        self.edges_in = edges_in
        return


def compute_branch_stat_update(
    c,
    child_samples,
    A_state,
    B_state,
    state_dim,
    sign,
    stat_func,
    num_samples,
    result,
    params,
):
    """Compute an update to the two-locus statistic for a single subset of the
    tree being modified, relative to all subsets of the fixed tree. We perform
    this operation for all samples edge being modified. For subsequent parent
    nodes, we update the statistic by removing the existing contribution after
    adding in the update contribution.

    i.e. if we're adding two samples ({3, 4}) to a node, if the parent node
    contains {1, 2}, we first add the statistic for {1, 2, 3, 4}, then
    subtract the stat for {1, 2}.

    :param c: Child node of the edge we're modifying
    :param child_samples: Samples under the edge being added/removed
    :param A_state: State for the tree contributing to the A samples (fixed)
    :param A_state: State for the tree contributing to the B samples (modified)
    :param state_dim: Number of sample sets.
    :param sign: The sign of the update
    :param stat_func: Function used to compute the two-locus statistic
    :param num_samples: Number of samples in the tree sequence
    :param result: Vector of LD results, length of number of sample sets
    :param params: Params of summary function.
    """
    b_len = B_state.branch_len[c] * sign
    if b_len == 0:
        return result

    AB_samples = BitSet(num_samples, 1)
    node_samples_tmp = BitSet(num_samples, 1)
    weights = np.zeros((3, state_dim), dtype=np.int64)
    result_tmp = np.zeros(state_dim, np.float64)

    for n in np.where(A_state.branch_len > 0)[0]:
        a_len = A_state.branch_len[n]
        for k in range(state_dim):
            row = (state_dim * n) + k
            c_row = (state_dim * c) + k
            # Samples under the modified edge and the current fixed tree node are AB
            A_state.node_samples.intersect(row, B_state.node_samples, c_row, AB_samples)

            w_AB = AB_samples.count(0)
            w_A = A_state.node_samples.count(row)
            w_B = B_state.node_samples.count(c_row)

            weights[0, k] = w_AB
            weights[1, k] = w_A - w_AB  # w_Ab
            weights[2, k] = w_B - w_AB  # w_aB

        stat_func(state_dim, weights, result_tmp, params)
        for k in range(state_dim):
            result[k] += result_tmp[k] * a_len * b_len

        # If we've begun our walk up the parents of the current edge removal, we
        # must adjust the statistic for samples that were already present before
        # addition or that remain after removal.
        if child_samples is not None:
            for k in range(state_dim):
                row = (state_dim * n) + k
                c_row = (state_dim * c) + k
                node_samples_tmp.union(0, B_state.node_samples, c_row)
                node_samples_tmp.difference(0, child_samples, k)
                AB_samples.data[:] = 0  # Zero out the bitset so that we can reuse it
                A_state.node_samples.intersect(row, node_samples_tmp, 0, AB_samples)

                w_AB = AB_samples.count(0)
                w_A = A_state.node_samples.count(row)
                w_B = node_samples_tmp.count(0)

                weights[0, k] = w_AB
                weights[1, k] = w_A - w_AB  # w_Ab
                weights[2, k] = w_B - w_AB  # w_aB

            stat_func(state_dim, weights, result_tmp, params)
            for k in range(state_dim):
                result[k] -= result_tmp[k] * a_len * b_len


def compute_branch_stat(ts, stat_func, stat, params, state_dim, l_state, r_state):
    """Step between trees in a tree sequence, updating our two-locus statistic
    as we add or remove edges. Since we're computing statistics for two loci, we
    have a focal tree that remains constant, and a tree that is updated to
    represent the tree we're comparing to. The lefthand tree is held constant
    and the righthand tree is modified. The statistic is updated as we add and
    remove branches, and when we reach the point where the righthand tree is
    fully updated, the statistic will have been updated to the two-locus
    statistic between both trees.

    For instance, if we pass in the l_state for tree 0 and the r_state for tree
    0, we will update the r_state until r_state contains the information for
    tree 1. Then, the statistic will represent the LD between tree 1 and tree 2.

    Currenty, iteration happens in the forward direction.

    :param ts: The underlying tree sequence object that we're iterating across.
    :param stat_func: A function that computes the two locus statistic, given
                      haplotype counts.
    :param stat: The two-locus statistic computed between two trees.
    :param params: Params of summary function.
    :param state_dim: Number of sample sets.
    :param l_state: The lefthand constant state
    :param r_state: The righthand state to be updated
    :returns: A tuple containing the statistic between the two trees after
              branch updates and the righthand tree state.
    """
    time = ts.tables.nodes.time

    child_samples = BitSet(ts.num_samples, state_dim)
    for e in r_state.edges_out:
        p = ts.edges_parent[e]
        c = ts.edges_child[e]
        child_samples.data[:] = 0
        for k in range(state_dim):
            c_row = (state_dim * c) + k
            child_samples.union(k, r_state.node_samples, c_row)

        # Remove the LD contributed by the samples under removed edges. When
        # we walk up the tree to propagate these changes to parents of the
        # removed edge, we need to add back in the LD contributed by samples
        # that aren't removed. We remove samples from the parents of the removed
        # branch as we propagate changes upward
        in_parent = None
        while p != tskit.NULL:
            compute_branch_stat_update(
                c,
                in_parent,
                l_state,
                r_state,
                state_dim,
                -1,
                stat_func,
                ts.num_samples,
                stat,
                params,
            )
            if in_parent is not None:
                # remove samples from the parents of the branch being removed
                # we remove the child node after the first iteration
                for k in range(state_dim):
                    c_row = (state_dim * c) + k
                    r_state.node_samples.difference(c_row, child_samples, k)
            in_parent = child_samples
            c = p
            p = r_state.parent[p]
        for k in range(state_dim):
            c_row = (state_dim * c) + k
            r_state.node_samples.difference(c_row, child_samples, k)

        # reset to the child of the edge being removed.
        c = ts.edges_child[e]
        r_state.branch_len[c] = 0
        r_state.parent[c] = tskit.NULL

    for e in r_state.edges_in:
        p = ts.edges_parent[e]
        c = ts.edges_child[e]
        child_samples.data[:] = 0
        for k in range(state_dim):
            c_row = (state_dim * c) + k
            child_samples.union(k, r_state.node_samples, c_row)
        r_state.branch_len[c] = time[p] - time[c]
        r_state.parent[c] = p

        # Add the LD contributed by the samples under added edges. When we walk
        # up the tree to propagate these changes to parents of the removed edge,
        # we need to remove the LD contributed by samples that were already
        # there
        in_parent = None
        while p != tskit.NULL:
            for k in range(state_dim):
                p_row = (state_dim * p) + k
                r_state.node_samples.union(p_row, child_samples, k)
            compute_branch_stat_update(
                c,
                in_parent,
                l_state,
                r_state,
                state_dim,
                +1,
                stat_func,
                ts.num_samples,
                stat,
                params,
            )
            in_parent = child_samples
            c = p
            p = r_state.parent[p]

    return stat, r_state


@pytest.mark.parametrize(
    "ts",
    [
        ts
        for ts in get_example_tree_sequences()
        if ts.id
        not in {
            "no_samples",
            "empty_ts",
            # We must skip these cases so that tests run in a reasonable
            # amount of time. To get more complete testing, these filters
            # can be commented out. (runtime ~1hr)
            "gap_0",
            "gap_0.1",
            "gap_0.5",
            "gap_0.75",
            "n=2_m=32_rho=0",
            "n=10_m=1_rho=0",
            "n=10_m=1_rho=0.1",
            "n=10_m=2_rho=0",
            "n=10_m=2_rho=0.1",
            "n=10_m=32_rho=0",
            "n=10_m=32_rho=0.1",
            "n=10_m=32_rho=0.5",
            # we keep one n=100 case to ensure bit arrays are working
            "n=100_m=1_rho=0.1",
            "n=100_m=1_rho=0.5",
            "n=100_m=2_rho=0",
            "n=100_m=2_rho=0.1",
            "n=100_m=2_rho=0.5",
            "n=100_m=32_rho=0",
            "n=100_m=32_rho=0.1",
            "n=100_m=32_rho=0.5",
            "all_fields",
            "back_mutations",
            "multichar",
            "multichar_no_metadata",
            "bottleneck_n=100_mutated",
        }
    ],
)
@pytest.mark.parametrize("stat", sorted(SUMMARY_FUNCS.keys()))
def test_branch_ld_matrix(ts, stat):
    np.testing.assert_array_almost_equal(
        ts.ld_matrix(stat=stat, mode="branch"), ld_matrix(ts, stat=stat, mode="branch")
    )


def get_test_branch_sample_set_test_cases():
    p_dict = {ps.id: ps for ps in get_example_tree_sequences()}
    return [
        pytest.param(
            p_dict["n=100_m=1_rho=0"].values[0],
            [[51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67]],
            id="n=100_m=1_rho=0",
        ),
        pytest.param(
            p_dict["all_nodes_samples"].values[0],
            [[2, 4, 5, 6]],
            id="all_nodes_samples",
        ),
        pytest.param(
            p_dict["bottleneck_n=10_mutated"].values[0],
            [[1, 2, 4, 9]],
            id="bottleneck_n=10_mutated",
        ),
        pytest.param(p_dict["gap_at_end"].values[0], [[1, 3, 5, 8]], id="gap_at_end"),
    ]


@pytest.mark.parametrize("ts,sample_set", get_test_branch_sample_set_test_cases())
@pytest.mark.parametrize("stat", sorted(SUMMARY_FUNCS.keys()))
def test_branch_ld_matrix_sample_sets(ts, sample_set, stat):
    np.testing.assert_array_almost_equal(
        np.expand_dims(
            ld_matrix(ts, stat=stat, mode="branch", sample_sets=sample_set), axis=0
        ),
        ts.ld_matrix(stat=stat, mode="branch", sample_sets=sample_set),
    )


--- ../../tskit/python/tests/data/dict-encodings/generate_msprime.py ---

import pathlib
import pickle

import _msprime
import msprime

pop_configs = [msprime.PopulationConfiguration(5) for _ in range(2)]
migration_matrix = [[0, 1], [1, 0]]
ts = msprime.simulate(
    population_configurations=pop_configs,
    migration_matrix=migration_matrix,
    mutation_rate=1,
    record_migrations=True,
    random_seed=1,
)
lwt = _msprime.LightweightTableCollection()
lwt.fromdict(ts.tables.asdict())

test_dir = pathlib.Path(__file__).parent
with open(test_dir / f"msprime-{msprime.__version__}.pkl", "wb") as f:
    pickle.dump(lwt.asdict(), f)


--- ../../tskit/python/tskit/combinatorics.py ---


"""
Module for ranking and unranking trees. Trees are considered only
leaf-labelled and unordered, so order of children does not influence equality.
"""
import collections
import functools
import heapq
import itertools
import json
import random
from typing import NamedTuple

import attr
import numpy as np

import tskit


class Rank(NamedTuple):
    """
    A tuple of 2 numbers, ``(shape, label)``, together defining a unique
    topology for a labeled tree. See :ref:`sec_combinatorics`.
    """

    shape: int
    """
    A non-negative integer representing the (unlabelled) topology of a tree with a
    defined number of tips.
    """
    label: int
    """
    A non-negative integer representing the order of labels for a given tree topology.
    """


def equal_chunks(lst, k):
    """
    Yield k successive equally sized chunks from lst of size n.

    If k >= n, we return n chunks of size 1.

    Otherwise, we always return k chunks. The first k - 1 chunks will
    contain exactly n // k items, and the last chunk the remainder.
    """
    n = len(lst)
    if k <= 0 or int(k) != k:
        raise ValueError("Number of chunks must be a positive integer")

    if n > 0:
        chunk_size = max(1, n // k)
        offset = 0
        j = 0
        while offset < n - chunk_size and j < k - 1:
            yield lst[offset : offset + chunk_size]
            offset += chunk_size
            j += 1
        yield lst[offset:]


@attr.s(eq=False)
class TreeNode:
    """
    Simple linked tree class used to generate tree topologies.
    """

    parent = attr.ib(default=None)
    children = attr.ib(factory=list)
    label = attr.ib(default=None)

    def as_tables(self, *, num_leaves, span, branch_length):
        """
        Convert the tree rooted at this node into an equivalent
        TableCollection. Internal nodes are allocated in postorder.
        """
        tables = tskit.TableCollection(span)
        for _ in range(num_leaves):
            tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)

        def assign_internal_labels(node):
            if len(node.children) == 0:
                node.time = 0
            else:
                max_child_time = 0
                for child in node.children:
                    assign_internal_labels(child)
                    max_child_time = max(max_child_time, child.time)
                node.time = max_child_time + branch_length
                node.label = tables.nodes.add_row(time=node.time)
                for child in node.children:
                    tables.edges.add_row(0, span, node.label, child.label)

        # Do a postorder traversal to assign the internal node labels and times.
        assign_internal_labels(self)
        tables.sort()
        return tables

    @staticmethod
    def random_binary_tree(leaf_labels, rng):
        """
        Returns a random binary tree where the leaves have the specified
        labels using the specified random.Random instance. The root node
        of this tree is returned.

        Based on the description of Remy's method of generating "decorated"
        random binary trees in TAOCP 7.2.1.6. This is not a direct
        implementation of Algorithm R, because we are interested in
        the leaf node labellings.

        The pre-fascicle text is available here, page 16:
        http://www.cs.utsa.edu/~wagner/knuth/fasc4a.pdf
        """
        nodes = [TreeNode(label=leaf_labels[0])]
        for label in leaf_labels[1:]:
            # Choose a node x randomly and insert a new internal node above
            # it with the (n + 1)th labelled leaf as its sibling.
            x = rng.choice(nodes)
            new_leaf = TreeNode(label=label)
            new_internal = TreeNode(parent=x.parent, children=[x, new_leaf])
            if x.parent is not None:
                index = x.parent.children.index(x)
                x.parent.children[index] = new_internal
            rng.shuffle(new_internal.children)
            x.parent = new_internal
            new_leaf.parent = new_internal
            nodes.extend([new_leaf, new_internal])

        root = nodes[0]
        while root.parent is not None:
            root = root.parent

        # Canonicalise the order of the children within a node. This
        # is given by (num_leaves, min_label). See also the
        # RankTree.canonical_order function for the definition of
        # how these are ordered during rank/unrank.

        def reorder_children(node):
            if len(node.children) == 0:
                return 1, node.label
            keys = [reorder_children(child) for child in node.children]
            if keys[0] > keys[1]:
                node.children = node.children[::-1]
            return (
                sum(leaf_count for leaf_count, _ in keys),
                min(min_label for _, min_label in keys),
            )

        reorder_children(root)
        return root

    @classmethod
    def balanced_tree(cls, leaf_labels, arity):
        """
        Returns a balanced tree of the specified arity. At each node the
        leaf labels are split equally among the arity children using the
        equal_chunks method.
        """
        assert len(leaf_labels) > 0
        if len(leaf_labels) == 1:
            root = cls(label=leaf_labels[0])
        else:
            children = [
                cls.balanced_tree(chunk, arity)
                for chunk in equal_chunks(leaf_labels, arity)
            ]
            root = cls(children=children)
            for child in children:
                child.parent = root
        return root


def generate_star(num_leaves, *, span, branch_length, record_provenance, **kwargs):
    """
    Generate a star tree for the specified number of leaves.

    See the documentation for :meth:`Tree.generate_star` for more details.
    """
    if num_leaves < 2:
        raise ValueError("The number of leaves must be 2 or greater")
    tables = tskit.TableCollection(sequence_length=span)
    tables.nodes.set_columns(
        flags=np.full(num_leaves, tskit.NODE_IS_SAMPLE, dtype=np.uint32),
        time=np.zeros(num_leaves),
    )
    root = tables.nodes.add_row(time=branch_length)
    tables.edges.set_columns(
        left=np.full(num_leaves, 0),
        right=np.full(num_leaves, span),
        parent=np.full(num_leaves, root, dtype=np.int32),
        child=np.arange(num_leaves, dtype=np.int32),
    )
    if record_provenance:
        # TODO replace with a version of https://github.com/tskit-dev/tskit/pull/243
        # TODO also make sure we convert all the arguments so that they are
        # definitely JSON encodable.
        parameters = {"command": "generate_star", "TODO": "add parameters"}
        tables.provenances.add_row(
            record=json.dumps(tskit.provenance.get_provenance_dict(parameters))
        )
    return tables.tree_sequence().first(**kwargs)


def generate_comb(num_leaves, *, span, branch_length, record_provenance, **kwargs):
    """
    Generate a comb tree for the specified number of leaves.

    See the documentation for :meth:`Tree.generate_comb` for more details.
    """
    if num_leaves < 2:
        raise ValueError("The number of leaves must be 2 or greater")
    tables = tskit.TableCollection(sequence_length=span)
    tables.nodes.set_columns(
        flags=np.full(num_leaves, tskit.NODE_IS_SAMPLE, dtype=np.uint32),
        time=np.zeros(num_leaves),
    )
    right_child = num_leaves - 1
    time = branch_length
    for left_child in range(num_leaves - 2, -1, -1):
        parent = tables.nodes.add_row(time=time)
        time += branch_length
        tables.edges.add_row(0, span, parent, left_child)
        tables.edges.add_row(0, span, parent, right_child)
        right_child = parent

    if record_provenance:
        # TODO replace with a version of https://github.com/tskit-dev/tskit/pull/243
        # TODO also make sure we convert all the arguments so that they are
        # definitely JSON encodable.
        parameters = {"command": "generate_comb", "TODO": "add parameters"}
        tables.provenances.add_row(
            record=json.dumps(tskit.provenance.get_provenance_dict(parameters))
        )
    return tables.tree_sequence().first(**kwargs)


def generate_balanced(
    num_leaves, *, arity, span, branch_length, record_provenance, **kwargs
):
    """
    Generate a balanced tree for the specified number of leaves.

    See the documentation for :meth:`Tree.generate_balanced` for more details.
    """
    if num_leaves < 1:
        raise ValueError("The number of leaves must be at least 1")
    if arity < 2:
        raise ValueError("The arity must be at least 2")

    root = TreeNode.balanced_tree(range(num_leaves), arity)
    tables = root.as_tables(
        num_leaves=num_leaves, span=span, branch_length=branch_length
    )

    if record_provenance:
        # TODO replace with a version of https://github.com/tskit-dev/tskit/pull/243
        # TODO also make sure we convert all the arguments so that they are
        # definitely JSON encodable.
        parameters = {"command": "generate_balanced", "TODO": "add parameters"}
        tables.provenances.add_row(
            record=json.dumps(tskit.provenance.get_provenance_dict(parameters))
        )

    return tables.tree_sequence().first(**kwargs)


def generate_random_binary(
    num_leaves, *, span, branch_length, random_seed, record_provenance, **kwargs
):
    """
    Sample a leaf-labelled binary tree uniformly.

    See the documentation for :meth:`Tree.generate_random_binary` for more details.
    """
    if num_leaves < 1:
        raise ValueError("The number of leaves must be at least 1")

    rng = random.Random(random_seed)
    root = TreeNode.random_binary_tree(range(num_leaves), rng)
    tables = root.as_tables(
        num_leaves=num_leaves, span=span, branch_length=branch_length
    )

    if record_provenance:
        # TODO replace with a version of https://github.com/tskit-dev/tskit/pull/243
        # TODO also make sure we convert all the arguments so that they are
        # definitely JSON encodable.
        parameters = {"command": "generate_random_binary", "TODO": "add parameters"}
        tables.provenances.add_row(
            record=json.dumps(tskit.provenance.get_provenance_dict(parameters))
        )
    ts = tables.tree_sequence()
    return ts.first(**kwargs)


def split_polytomies(
    tree,
    *,
    epsilon=None,
    method=None,
    record_provenance=True,
    random_seed=None,
    **kwargs,
):
    """
    Return a new tree where extra nodes and edges have been inserted
    so that any any node with more than two children is resolved into
    a binary tree.

    See the documentation for :meth:`Tree.split_polytomies` for more details.
    """
    allowed_methods = ["random"]
    if method is None:
        method = "random"
    if method not in allowed_methods:
        raise ValueError(f"Method must be chosen from {allowed_methods}")

    tables = tree.tree_sequence.dump_tables()
    tables.keep_intervals([tree.interval], simplify=False)
    tables.edges.clear()
    rng = random.Random(random_seed)

    for u in tree.nodes():
        if tree.num_children(u) > 2:
            root = TreeNode.random_binary_tree(tree.children(u), rng)
            root.label = u
            root_time = tree.time(u)
            stack = [(child, root_time) for child in root.children]
            while len(stack) > 0:
                node, parent_time = stack.pop()
                if node.label is None:
                    if epsilon is None:
                        child_time = np.nextafter(parent_time, -np.inf)
                    else:
                        child_time = parent_time - epsilon
                    node.label = tables.nodes.add_row(time=child_time)
                else:
                    assert len(node.children) == 0
                    # This is a leaf node connecting back into the original tree
                    child_time = tree.time(node.label)
                if parent_time <= child_time:
                    u = root.label
                    min_child_time = min(tree.time(v) for v in tree.children(u))
                    min_time = root_time - min_child_time
                    message = (
                        f"Cannot resolve the degree {tree.num_children(u)} "
                        f"polytomy rooted at node {u} with minimum time difference "
                        f"of {min_time} to the resolved leaves."
                    )
                    if epsilon is None:
                        message += (
                            " The time difference between nodes is so small that "
                            "more nodes cannot be inserted between within the limits "
                            "of floating point precision."
                        )
                    else:
                        # We can also have parent_time == child_time if epsilon is
                        # chosen such that we exactly divide up the branch in the
                        # original tree. We avoid saying this is caused by a
                        # too-small epsilon by noting it can only happen when we
                        # are at leaf node in the randomly generated tree.
                        if parent_time == child_time and len(node.children) > 0:
                            message += (
                                f" The fixed epsilon value of {epsilon} is too small, "
                                "resulting in the parent and child times being equal "
                                "within the limits of numerical precision."
                            )
                        else:
                            message += (
                                f" The fixed epsilon value of {epsilon} is too large, "
                                "resulting in the parent time being less than the child "
                                "time."
                            )
                    raise tskit.LibraryError(message)
                tables.edges.add_row(*tree.interval, node.parent.label, node.label)
                for child in node.children:
                    stack.append((child, child_time))
        else:
            for v in tree.children(u):
                tables.edges.add_row(*tree.interval, u, v)

    if record_provenance:
        parameters = {"command": "split_polytomies"}
        tables.provenances.add_row(
            record=json.dumps(tskit.provenance.get_provenance_dict(parameters))
        )
    try:
        tables.sort()
        ts = tables.tree_sequence()
    except tskit.LibraryError as e:
        msg = str(e)
        # We should have caught all topology time travel above.
        assert not msg.startswith("time[parent] must be greater than time[child]")
        if msg.startswith(
            "A mutation's time must be < the parent node of the edge on which it occurs"
        ):
            if epsilon is not None:
                msg = (
                    f"epsilon={epsilon} not small enough to create new nodes below a "
                    "polytomy, due to the time of a mutation above a child of the "
                    "polytomy."
                )
            else:
                msg = (
                    "Cannot split polytomy: mutation with numerical precision "
                    "of the parent time."
                )
            e.args += (msg,)
        raise e
    return ts.at(tree.interval.left, **kwargs)


def treeseq_count_topologies(ts, sample_sets):
    topology_counter = np.full(ts.num_nodes, None, dtype=object)
    parent = np.full(ts.num_nodes, -1)

    def update_state(tree, u):
        stack = [u]
        while len(stack) > 0:
            v = stack.pop()
            children = []
            for c in tree.children(v):
                if topology_counter[c] is not None:
                    children.append(topology_counter[c])
            if len(children) > 0:
                topology_counter[v] = combine_child_topologies(children)
            else:
                topology_counter[v] = None
            p = parent[v]
            if p != -1:
                stack.append(p)

    for sample_set_index, sample_set in enumerate(sample_sets):
        for u in sample_set:
            if not ts.node(u).is_sample():
                raise ValueError(f"Node {u} in sample_sets is not a sample.")
            topology_counter[u] = TopologyCounter.from_sample(sample_set_index)

    for tree, (_, edges_out, edges_in) in zip(ts.trees(), ts.edge_diffs()):
        # Avoid recomputing anything for the parent until all child edges
        # for that parent are inserted/removed
        for p, sibling_edges in itertools.groupby(edges_out, key=lambda e: e.parent):
            for e in sibling_edges:
                parent[e.child] = -1
            update_state(tree, p)
        for p, sibling_edges in itertools.groupby(edges_in, key=lambda e: e.parent):
            if tree.is_sample(p):
                raise ValueError("Internal samples not supported.")
            for e in sibling_edges:
                parent[e.child] = p
            update_state(tree, p)

        counters = []
        for root in tree.roots:
            if topology_counter[root] is not None:
                counters.append(topology_counter[root])
        yield TopologyCounter.merge(counters)


def tree_count_topologies(tree, sample_sets):
    for u in tree.samples():
        if not tree.is_leaf(u):
            raise ValueError("Internal samples not supported.")

    topology_counter = np.full(tree.tree_sequence.num_nodes, None, dtype=object)
    for sample_set_index, sample_set in enumerate(sample_sets):
        for u in sample_set:
            if not tree.is_sample(u):
                raise ValueError(f"Node {u} in sample_sets is not a sample.")
            topology_counter[u] = TopologyCounter.from_sample(sample_set_index)

    for u in tree.nodes(order="postorder"):
        children = []
        for v in tree.children(u):
            if topology_counter[v] is not None:
                children.append(topology_counter[v])
        if len(children) > 0:
            topology_counter[u] = combine_child_topologies(children)

    counters = []
    for root in tree.roots:
        if topology_counter[root] is not None:
            counters.append(topology_counter[root])
    return TopologyCounter.merge(counters)


def combine_child_topologies(topology_counters):
    """
    Select all combinations of topologies from different
    counters in ``topology_counters`` that are capable of
    being combined into a single topology. This includes
    any combination of at least two topologies, all from
    different children, where no topologies share a
    sample set index.
    """
    partial_topologies = PartialTopologyCounter()
    for tc in topology_counters:
        partial_topologies.add_sibling_topologies(tc)

    return partial_topologies.join_all_combinations()


class TopologyCounter:
    """
    Contains the distributions of embedded topologies for every combination
    of the sample sets used to generate the ``TopologyCounter``. It is
    indexable by a combination of sample set indexes and returns a
    ``collections.Counter`` whose keys are topology ranks
    (see :ref:`sec_tree_ranks`). See :meth:`Tree.count_topologies` for more
    detail on how this structure is used.
    """

    def __init__(self):
        self.topologies = collections.defaultdict(collections.Counter)

    def __getitem__(self, sample_set_indexes):
        k = TopologyCounter._to_key(sample_set_indexes)
        return self.topologies[k]

    def __setitem__(self, sample_set_indexes, counter):
        k = TopologyCounter._to_key(sample_set_indexes)
        self.topologies[k] = counter

    @staticmethod
    def _to_key(sample_set_indexes):
        if not isinstance(sample_set_indexes, collections.abc.Iterable):
            sample_set_indexes = (sample_set_indexes,)
        return tuple(sorted(sample_set_indexes))

    def __eq__(self, other):
        return self.__class__ == other.__class__ and self.topologies == other.topologies

    @staticmethod
    def merge(topology_counters):
        """
        Union together independent topology counters into one.
        """
        total = TopologyCounter()
        for tc in topology_counters:
            for k, v in tc.topologies.items():
                total.topologies[k] += v

        return total

    @staticmethod
    def from_sample(sample_set_index):
        """
        Generate the topologies covered by a single sample. This
        is the single-leaf topology representing the single sample
        set.
        """
        rank_tree = RankTree(children=[], label=sample_set_index)
        tc = TopologyCounter()
        tc[sample_set_index][rank_tree.rank()] = 1
        return tc


class PartialTopologyCounter:
    """
    Represents the possible combinations of children under a node in a tree
    and the combinations of embedded topologies that are rooted at the node.
    This allows an efficient way of calculating which unique embedded
    topologies arise by only every storing a given pairing of sibling topologies
    once.
    ``partials`` is a dictionary where a key is a tuple of sample set indexes,
    and the value is a ``collections.Counter`` that counts combinations of
    sibling topologies whose tips represent the sample sets in the key.
    Each element of the counter is a homogeneous tuple where each element represents
    a topology. The topology is itself a tuple of the sample set indexes in that
    topology and the rank.
    """

    def __init__(self):
        self.partials = collections.defaultdict(collections.Counter)

    def add_sibling_topologies(self, topology_counter):
        """
        Combine each topology in the given TopologyCounter with every existing
        combination of topologies whose sample set indexes are disjoint from the
        topology from the counter. This also includes adding the topologies from
        the counter without joining them to any existing combinations.
        """
        merged = collections.defaultdict(collections.Counter)
        for sample_set_indexes, topologies in topology_counter.topologies.items():
            for rank, count in topologies.items():
                topology = ((sample_set_indexes, rank),)
                # Cross with existing topology combinations
                for sibling_sample_set_indexes, siblings in self.partials.items():
                    if isdisjoint(sample_set_indexes, sibling_sample_set_indexes):
                        for sib_topologies, sib_count in siblings.items():
                            merged_topologies = merge_tuple(sib_topologies, topology)
                            merged_sample_set_indexes = merge_tuple(
                                sibling_sample_set_indexes, sample_set_indexes
                            )
                            merged[merged_sample_set_indexes][merged_topologies] += (
                                count * sib_count
                            )
                # Propagate without combining
                merged[sample_set_indexes][topology] += count

        for sample_set_indexes, counter in merged.items():
            self.partials[sample_set_indexes] += counter

    def join_all_combinations(self):
        """
        For each pairing of child topologies, join them together into a new
        tree and count the resulting topologies.
        """
        topology_counter = TopologyCounter()
        for sample_set_indexes, sibling_topologies in self.partials.items():
            for topologies, count in sibling_topologies.items():
                # A node must have at least two children
                if len(topologies) >= 2:
                    rank = PartialTopologyCounter.join_topologies(topologies)
                    topology_counter[sample_set_indexes][rank] += count
                else:
                    # Pass on the single tree without adding a parent
                    for _, rank in topologies:
                        topology_counter[sample_set_indexes][rank] += count

        return topology_counter

    @staticmethod
    def join_topologies(child_topologies):
        children = []
        for sample_set_indexes, rank in child_topologies:
            n = len(sample_set_indexes)
            t = RankTree.unrank(n, rank, list(sample_set_indexes))
            children.append(t)
        children.sort(key=RankTree.canonical_order)
        return RankTree(children).rank()


def all_trees(num_leaves, span=1):
    """
    Generates all unique leaf-labelled trees with ``num_leaves``
    leaves. See :ref:`sec_combinatorics` on the details of this
    enumeration. The leaf labels are selected from the set
    ``[0, num_leaves)``. The times and labels on internal nodes are
    chosen arbitrarily.

    :param int num_leaves: The number of leaves of the tree to generate.
    :param float span: The genomic span of each returned tree.
    :rtype: tskit.Tree
    """
    for rank_tree in RankTree.all_labelled_trees(num_leaves):
        yield rank_tree.to_tsk_tree(span=span)


def all_tree_shapes(num_leaves, span=1):
    """
    Generates all unique shapes of trees with ``num_leaves`` leaves.

    :param int num_leaves: The number of leaves of the tree to generate.
    :param float span: The genomic span of each returned tree.
    :rtype: tskit.Tree
    """
    for rank_tree in RankTree.all_unlabelled_trees(num_leaves):
        default_labelling = rank_tree.label_unrank(0)
        yield default_labelling.to_tsk_tree(span=span)


def all_tree_labellings(tree, span=1):
    """
    Generates all unique labellings of the leaves of a
    :class:`tskit.Tree`. Leaves are labelled from the set
    ``[0, n)`` where ``n`` is the number of leaves of ``tree``.

    :param tskit.Tree tree: The tree used to generate
        labelled trees of the same shape.
    :param float span: The genomic span of each returned tree.
    :rtype: tskit.Tree
    """
    rank_tree = RankTree.from_tsk_tree(tree)
    for labelling in RankTree.all_labellings(rank_tree):
        yield labelling.to_tsk_tree(span=span)


class RankTree:
    """
    A tree class that maintains the topological ranks of each node in the tree.
    This structure can be used to efficiently compute the rank of a tree of
    n labelled leaves and produce a tree given a rank.
    """

    def __init__(self, children, label=None):
        # Children are assumed to be sorted by RankTree.canonical_order
        self.children = children
        if len(children) == 0:
            self.num_leaves = 1
            self.labels = [label]
        else:
            self.num_leaves = sum(c.num_leaves for c in children)
            self.labels = list(heapq.merge(*(c.labels for c in children)))

        self._shape_rank = None
        self._label_rank = None

    def compute_shape_rank(self):
        """
        Mirroring the way in which unlabelled trees are enumerated, we must
        first calculate the number of trees whose partitions of number of leaves
        rank lesser than this tree's partition.

        Once we reach the partition of leaves in this tree, we examine the
        groups of child subtrees assigned to subsequences of the partition.
        For each group of children with the same number of leaves, k, the trees
        in that group were selected according to a combination with replacement
        of those trees from S(k). By finding the rank of that combination,
        we find how many combinations preceded the current one in that group.
        That rank is then multiplied by the total number of arrangements that
        could be made in the following groups, added to the total rank,
        and then we recur on the rest of the group and groups.
        """
        part = self.leaf_partition()
        total = 0
        for prev_part in partitions(self.num_leaves):
            if prev_part == part:
                break
            total += num_tree_pairings(prev_part)

        child_groups = self.group_children_by_num_leaves()
        next_child_idx = 0
        for g in child_groups:
            next_child_idx += len(g)
            k = g[0].num_leaves
            S_k = num_shapes(k)

            child_ranks = [c.shape_rank() for c in g]
            g_rank = Combination.with_replacement_rank(child_ranks, S_k)

            # TODO precompute vector before loop
            rest_part = part[next_child_idx:]
            total_rest = num_tree_pairings(rest_part)

            total += g_rank * total_rest

        return total

    def compute_label_rank(self):
        """
        Again mirroring how we've labeled a particular tree, T, we can rank the
        labelling on T.

        We group the children into symmetric groups. In the context of labelling,
        symmetric groups contain child trees that are of the same shape. Each
        group contains a combination of labels selected from all the labels
        available to T.

        The different variables to consider are:
        1. How to assign a combination of labels to the first group.
        2. Given a combination of labels assigned to the group, how can we
            distribute those labels to each tree in the group.
        3. Given an assignment of the labels to each tree in the group, how many
            distinct ways could all the trees in the group be labelled.

        These steps for generating labelled trees break down the stages of
        ranking them.
        For each group G, we can find the rank of the combination of labels
        assigned to G. This rank times the number of ways the trees in G
        could be labelled, times the number of possible labellings of the
        rest of the trees, gives the number of labellings that precede those with
        the given combination of labels assigned to G. This process repeats and
        breaks down to give the rank of the assignment of labels to trees in G,
        and the label ranks of the trees themselves in G.
        """
        all_labels = self.labels
        child_groups = self.group_children_by_shape()
        total = 0
        for i, g in enumerate(child_groups):
            rest_groups = child_groups[i + 1 :]
            g_labels = list(heapq.merge(*(t.labels for t in g)))
            num_rest_labellings = num_list_of_group_labellings(rest_groups)

            # Preceded by all of the ways to label all the groups
            # with a lower ranking combination given to g.
            comb_rank = Combination.rank(g_labels, all_labels)
            num_g_labellings = num_group_labellings(g)
            preceding_comb = comb_rank * num_g_labellings * num_rest_labellings

            # Preceded then by all the configurations of g ranking less than
            # the current one
            rank_from_g = group_rank(g) * num_rest_labellings

            total += preceding_comb + rank_from_g
            all_labels = set_minus(all_labels, g_labels)

        return total

    # TODO I think this would boost performance if it were a field and not
    # recomputed.
    def num_labellings(self):
        child_groups = self.group_children_by_shape()
        return num_list_of_group_labellings(child_groups)

    def rank(self):
        return Rank(self.shape_rank(), self.label_rank())

    def shape_rank(self):
        if self._shape_rank is None:
            self._shape_rank = self.compute_shape_rank()
        return self._shape_rank

    def label_rank(self):
        if self._label_rank is None:
            assert self.shape_rank() is not None
            self._label_rank = self.compute_label_rank()
        return self._label_rank

    @staticmethod
    def unrank(num_leaves, rank, labels=None):
        """
        Produce a ``RankTree`` of the given ``rank`` with ``num_leaves`` leaves,
        labelled with ``labels``. Labels must be sorted, and if ``None`` default
        to ``[0, num_leaves)``.
        """
        shape_rank, label_rank = rank
        if shape_rank < 0 or label_rank < 0:
            raise ValueError("Rank is out of bounds.")
        unlabelled = RankTree.shape_unrank(num_leaves, shape_rank)
        return unlabelled.label_unrank(label_rank, labels)

    @staticmethod
    def shape_unrank(n, shape_rank):
        """
        Generate an unlabelled tree with n leaves with a shape corresponding to
        the `shape_rank`.
        """
        part, child_shape_ranks = children_shape_ranks(shape_rank, n)
        children = [
            RankTree.shape_unrank(k, rk) for k, rk in zip(part, child_shape_ranks)
        ]

        t = RankTree(children=children)
        t._shape_rank = shape_rank
        return t

    def label_unrank(self, label_rank, labels=None):
        """
        Generate a tree with the same shape, whose leaves are labelled
        from ``labels`` with the labelling corresponding to ``label_rank``.
        """
        if labels is None:
            labels = list(range(self.num_leaves))

        if self.is_leaf():
            if label_rank != 0:
                raise ValueError("Rank is out of bounds.")
            return RankTree(children=[], label=labels[0])

        child_groups = self.group_children_by_shape()
        child_labels, child_label_ranks = children_label_ranks(
            child_groups, label_rank, labels
        )

        children = self.children
        labelled_children = [
            RankTree.label_unrank(c, c_rank, c_labels)
            for c, c_rank, c_labels in zip(children, child_label_ranks, child_labels)
        ]

        t = RankTree(children=labelled_children)
        t._shape_rank = self.shape_rank()
        t._label_rank = label_rank
        return t

    @staticmethod
    def canonical_order(c):
        """
        Defines the canonical ordering of sibling subtrees.
        """
        return c.num_leaves, c.shape_rank(), c.min_label()

    @staticmethod
    def from_tsk_tree_node(tree, u):
        if tree.is_leaf(u):
            return RankTree(children=[], label=u)

        if tree.num_children(u) == 1:
            raise ValueError("Cannot rank trees with unary nodes")

        children = list(
            sorted(
                (RankTree.from_tsk_tree_node(tree, c) for c in tree.children(u)),
                key=RankTree.canonical_order,
            )
        )
        return RankTree(children=children)

    @staticmethod
    def from_tsk_tree(tree):
        if tree.num_roots != 1:
            raise ValueError("Cannot rank trees with multiple roots")

        return RankTree.from_tsk_tree_node(tree, tree.root)

    def to_tsk_tree(self, span=1, branch_length=1):
        """
        Convert a ``RankTree`` into the only tree in a new tree sequence. Internal
        nodes and their times are assigned via a postorder traversal of the tree.

        :param float span: The genomic span of the returned tree. The tree will cover
            the interval :math:`[0, span)` and the :attr:`~Tree.tree_sequence` from which
            the tree is taken will have its :attr:`~tskit.TreeSequence.sequence_length`
            equal to ``span``.
        :param float branch_length: The minimum length of a branch in the returned
            tree.
        """
        if set(self.labels) != set(range(self.num_leaves)):
            raise ValueError("Labels set must be equivalent to [0, num_leaves)")

        tables = tskit.TableCollection(span)

        def add_node(node):
            if node.is_leaf():
                assert node.label is not None
                return node.label

            child_ids = [add_node(child) for child in node.children]
            max_child_time = max(tables.nodes.time[c] for c in child_ids)
            parent_id = tables.nodes.add_row(time=max_child_time + branch_length)
            for child_id in child_ids:
                tables.edges.add_row(0, span, parent_id, child_id)

            return parent_id

        for _ in range(self.num_leaves):
            tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
        add_node(self)

        # The way in which we're inserting nodes doesn't necessarily
        # adhere to the ordering constraint on edges, so we have
        # to sort.
        tables.sort()
        return tables.tree_sequence().first()

    @staticmethod
    def all_labelled_trees(n):
        """
        Generate all unordered, leaf-labelled trees with n leaves.
        """
        for tree in RankTree.all_unlabelled_trees(n):
            yield from RankTree.all_labellings(tree)

    @staticmethod
    def all_unlabelled_trees(n):
        """
        Generate all tree shapes with n leaves. See :ref:`sec_combinatorics`
        for how tree shapes are enumerated.
        """
        if n == 1:
            yield RankTree(children=[])
        else:
            for part in partitions(n):
                for subtree_pairing in RankTree.all_subtree_pairings(
                    group_partition(part)
                ):
                    yield RankTree(children=subtree_pairing)

    @staticmethod
    def all_subtree_pairings(grouped_part):
        if len(grouped_part) == 0:
            yield []
        else:
            g = grouped_part[0]
            k = g[0]
            all_k_leaf_trees = RankTree.all_unlabelled_trees(k)
            num_k_leaf_trees = len(g)
            g_trees = itertools.combinations_with_replacement(
                all_k_leaf_trees, num_k_leaf_trees
            )
            for first_trees in g_trees:
                for rest in RankTree.all_subtree_pairings(grouped_part[1:]):
                    yield list(first_trees) + rest

    @staticmethod
    def all_labellings(tree, labels=None):
        """
        Given a tree, generate all the unique labellings of that tree.
        See :ref:`sec_combinatorics` for how labellings of a tree are
        enumerated.
        """
        if labels is None:
            labels = list(range(tree.num_leaves))

        if tree.is_leaf():
            assert len(labels) == 1
            yield RankTree(children=[], label=labels[0])
        else:
            groups = tree.group_children_by_shape()
            for labeled_children in RankTree.label_all_groups(groups, labels):
                yield RankTree(children=labeled_children)

    @staticmethod
    def label_all_groups(groups, labels):
        if len(groups) == 0:
            yield []
        else:
            g, rest = groups[0], groups[1:]
            x = len(g)
            k = g[0].num_leaves
            for g_labels in itertools.combinations(labels, x * k):
                rest_labels = set_minus(labels, g_labels)
                for labeled_g in RankTree.label_tree_group(g, g_labels):
                    for labeled_rest in RankTree.label_all_groups(rest, rest_labels):
                        yield labeled_g + labeled_rest

    @staticmethod
    def label_tree_group(trees, labels):
        if len(trees) == 0:
            assert len(labels) == 0
            yield []
        else:
            first, rest = trees[0], trees[1:]
            k = first.num_leaves
            min_label = labels[0]
            for first_other_labels in itertools.combinations(labels[1:], k - 1):
                first_labels = [min_label] + list(first_other_labels)
                rest_labels = set_minus(labels, first_labels)
                for labeled_first in RankTree.all_labellings(first, first_labels):
                    for labeled_rest in RankTree.label_tree_group(rest, rest_labels):
                        yield [labeled_first] + labeled_rest

    def _newick(self):
        if self.is_leaf():
            return str(self.label) if self.labelled() else ""
        return "(" + ",".join(c._newick() for c in self.children) + ")"

    def newick(self):
        return self._newick() + ";"

    @property
    def label(self):
        return self.labels[0]

    def labelled(self):
        return all(label is not None for label in self.labels)

    def min_label(self):
        return self.labels[0]

    def is_leaf(self):
        return len(self.children) == 0

    def leaf_partition(self):
        return [c.num_leaves for c in self.children]

    def group_children_by_num_leaves(self):
        def same_num_leaves(c1, c2):
            return c1.num_leaves == c2.num_leaves

        return group_by(self.children, same_num_leaves)

    def group_children_by_shape(self):
        def same_shape(c1, c2):
            return c1.num_leaves == c2.num_leaves and c1.shape_rank() == c2.shape_rank()

        return group_by(self.children, same_shape)

    def __eq__(self, other):
        if self.__class__ != other.__class__:
            return False

        if self.is_leaf() and other.is_leaf():
            return self.label == other.label

        if len(self.children) != len(other.children):
            return False

        return all(c1 == c2 for c1, c2 in zip(self.children, other.children))

    def __ne__(self, other):
        return not self.__eq__(other)

    def shape_equal(self, other):
        if self.is_leaf() and other.is_leaf():
            return True

        if len(self.children) != len(other.children):
            return False

        return all(c1.shape_equal(c2) for c1, c2 in zip(self.children, other.children))

    def is_canonical(self):
        if self.is_leaf():
            return True

        children = self.children
        for c1, c2 in zip(children, children[1:]):
            if RankTree.canonical_order(c1) > RankTree.canonical_order(c2):
                return False
        return all(c.is_canonical() for c in children)

    def is_symmetrical(self):
        if self.is_leaf():
            return True

        even_split_leaves = len(set(self.leaf_partition())) == 1
        all_same_rank = len({c.shape_rank() for c in self.children}) == 1

        return even_split_leaves and all_same_rank


# TODO This is called repeatedly in ranking and unranking and has a perfect
# subtructure for DP. It's only every called on n in [0, num_leaves]
# so we should compute a vector of those results up front instead of using
# repeated calls to this function.
# Put an lru_cache on for now as a quick replacement (cuts test time down by 80%)
@functools.lru_cache
def num_shapes(n):
    """
    The cardinality of the set of unlabelled trees with n leaves,
    up to isomorphism.
    """
    if n <= 1:
        return n
    return sum(num_tree_pairings(part) for part in partitions(n))


def num_tree_pairings(part):
    """
    The number of unique tree shapes that could be assembled from
    a given partition of leaves. If we group the elements of the partition
    by number of leaves, each group can be independently enumerated and the
    cardinalities of each group's pairings can be multiplied. Within a group,
    subsequent trees must have equivalent or greater rank, so the number of
    ways to select trees follows combinations with replacement from the set
    of all possible trees for that group.
    """
    total = 1
    for g in group_partition(part):
        k = g[0]
        total *= Combination.comb_with_replacement(num_shapes(k), len(g))
    return total


def num_labellings(n, shape_rk):
    return RankTree.shape_unrank(n, shape_rk).num_labellings()


def children_shape_ranks(rank, n):
    """
    Return the partition of leaves associated
    with the children of the tree of rank `rank`, and
    the ranks of each child tree.
    """
    part = []
    for prev_part in partitions(n):
        num_trees_with_part = num_tree_pairings(prev_part)
        if rank < num_trees_with_part:
            part = prev_part
            break
        rank -= num_trees_with_part
    else:
        if n != 1:
            raise ValueError("Rank is out of bounds.")

    grouped_part = group_partition(part)
    child_ranks = []
    next_child = 0
    for g in grouped_part:
        next_child += len(g)
        k = g[0]

        # TODO precompute vector up front
        rest_children = part[next_child:]
        rest_num_pairings = num_tree_pairings(rest_children)

        shapes_comb_rank = rank // rest_num_pairings
        g_shape_ranks = Combination.with_replacement_unrank(
            shapes_comb_rank, num_shapes(k), len(g)
        )
        child_ranks += g_shape_ranks
        rank %= rest_num_pairings

    return part, child_ranks


def children_label_ranks(child_groups, rank, labels):
    """
    Produces the subsets of labels assigned to each child
    and the associated label rank of each child.
    """
    child_labels = []
    child_label_ranks = []

    for i, g in enumerate(child_groups):
        k = g[0].num_leaves
        g_num_leaves = k * len(g)
        num_g_labellings = num_group_labellings(g)
        # TODO precompute vector of partial products outside of loop
        rest_groups = child_groups[i + 1 :]
        num_rest_labellings = num_list_of_group_labellings(rest_groups)

        num_labellings_per_label_comb = num_g_labellings * num_rest_labellings
        comb_rank = rank // num_labellings_per_label_comb
        rank_given_label_comb = rank % num_labellings_per_label_comb
        g_rank = rank_given_label_comb // num_rest_labellings

        g_labels = Combination.unrank(comb_rank, labels, g_num_leaves)

        g_child_labels, g_child_ranks = group_label_ranks(g_rank, g, g_labels)
        child_labels += g_child_labels
        child_label_ranks += g_child_ranks

        labels = set_minus(labels, g_labels)
        rank %= num_rest_labellings

    return child_labels, child_label_ranks


def group_rank(g):
    k = g[0].num_leaves
    n = len(g) * k
    # Num ways to label a single one of the trees
    # We can do this once because all the trees in the group
    # are of the same shape rank
    y = g[0].num_labellings()
    all_labels = list(heapq.merge(*(t.labels for t in g)))
    rank = 0
    for i, t in enumerate(g):
        u_labels = t.labels
        curr_trees = len(g) - i
        # Kind of cheating here leaving the selection of min labels implicit
        # because the rank of the comb without min labels is the same
        comb_rank = Combination.rank(u_labels, all_labels)

        # number of ways to distribute labels to rest leaves
        num_rest_combs = 1
        remaining_leaves = n - (i + 1) * k
        for j in range(curr_trees - 1):
            num_rest_combs *= Combination.comb(remaining_leaves - j * k - 1, k - 1)

        preceding_combs = comb_rank * num_rest_combs * (y**curr_trees)
        curr_comb = t.label_rank() * num_rest_combs * (y ** (curr_trees - 1))
        rank += preceding_combs + curr_comb
        all_labels = set_minus(all_labels, u_labels)
    return rank


# TODO This is only used in a few cases and mostly in a n^2 way. Would
# be easy and useful to do this DP and produce a list of partial products
def num_list_of_group_labellings(groups):
    """
    Given a set of labels and a list of groups, how many unique ways are there
    to assign subsets of labels to each group in the list and subsequently
    label all the trees in all the groups.
    """
    remaining_leaves = sum(len(g) * g[0].num_leaves for g in groups)
    total = 1
    for g in groups:
        k = g[0].num_leaves
        x = len(g)
        num_label_choices = Combination.comb(remaining_leaves, x * k)
        total *= num_label_choices * num_group_labellings(g)
        remaining_leaves -= x * k

    return total


def num_group_labellings(g):
    """
    Given a particular set of labels, how many unique ways are there
    to assign subsets of labels to each tree in the group and subsequently
    label those trees.
    """
    # Shortcut because all the trees are identical and can therefore
    # be labelled in the same ways
    num_tree_labelings = g[0].num_labellings() ** len(g)
    return num_assignments_in_group(g) * num_tree_labelings


def num_assignments_in_group(g):
    """
    Given this group of identical trees, how many unique ways
    are there to divide up a set of n labels?
    """
    n = sum(t.num_leaves for t in g)
    total = 1
    for t in g:
        k = t.num_leaves
        # Choose k - 1 from n - 1 because the minimum label must be
        # assigned to the first tree for a canonical labelling.
        total *= Combination.comb(n - 1, k - 1)
        n -= k
    return total


def group_label_ranks(rank, child_group, labels):
    """
    Given a group of trees of the same shape, a label rank and list of labels,
    produce assignment of label subsets to each tree in the group and the
    label rank of each tree.
    """
    child_labels = []
    child_label_ranks = []

    for i, rank_tree in enumerate(child_group):
        k = rank_tree.num_leaves
        num_t_labellings = rank_tree.num_labellings()
        rest_trees = child_group[i + 1 :]
        num_rest_assignments = num_assignments_in_group(rest_trees)
        num_rest_labellings = num_rest_assignments * (
            num_t_labellings ** len(rest_trees)
        )
        num_labellings_per_label_comb = num_t_labellings * num_rest_labellings

        comb_rank = rank // num_labellings_per_label_comb
        rank_given_comb = rank % num_labellings_per_label_comb
        t_rank = rank_given_comb // num_rest_labellings
        rank %= num_rest_labellings

        min_label = labels[0]
        t_labels = [min_label] + Combination.unrank(comb_rank, labels[1:], k - 1)
        labels = set_minus(labels, t_labels)

        child_labels.append(t_labels)
        child_label_ranks.append(t_rank)

    return child_labels, child_label_ranks


class Combination:
    @staticmethod
    def comb(n, k):
        """
        The number of times you can select k items from
        n items without order and without replacement.

        FIXME: This function will be available in `math` in Python 3.8
        and should be replaced eventually.
        """
        k = min(k, n - k)
        res = 1
        for i in range(1, k + 1):
            res *= n - k + i
            res //= i

        return res

    @staticmethod
    def comb_with_replacement(n, k):
        """
        Also called multichoose, the number of times you can select
        k items from n items without order but *with* replacement.
        """
        return Combination.comb(n + k - 1, k)

    @staticmethod
    def rank(combination, elements):
        """
        Find the combination of k elements from the given set of elements
        with the given rank in a lexicographic ordering.
        """
        indices = [elements.index(x) for x in combination]
        return Combination.from_range_rank(indices, len(elements))

    @staticmethod
    def from_range_rank(combination, n):
        """
        Find the combination of k integers from [0, n)
        with the given rank in a lexicographic ordering.
        """
        k = len(combination)
        if k == 0 or k == n:
            return 0

        j = combination[0]
        combination = [x - 1 for x in combination]
        if j == 0:
            return Combination.from_range_rank(combination[1:], n - 1)

        first_rank = Combination.comb(n - 1, k - 1)
        rest_rank = Combination.from_range_rank(combination, n - 1)
        return first_rank + rest_rank

    @staticmethod
    def unrank(rank, elements, k):
        n = len(elements)
        if k == 0:
            return []
        if len(elements) == 0:
            raise ValueError("Rank is out of bounds.")

        n_rest_combs = Combination.comb(n - 1, k - 1)
        if rank < n_rest_combs:
            return elements[:1] + Combination.unrank(rank, elements[1:], k - 1)

        return Combination.unrank(rank - n_rest_combs, elements[1:], k)

    @staticmethod
    def with_replacement_rank(combination, n):
        """
        Find the rank of ``combination`` in the lexicographic ordering of
        combinations with replacement of integers from [0, n).
        """
        k = len(combination)
        if k == 0:
            return 0
        j = combination[0]
        if k == 1:
            return j

        if j == 0:
            return Combination.with_replacement_rank(combination[1:], n)

        rest = [x - j for x in combination[1:]]
        preceding = 0
        for i in range(j):
            preceding += Combination.comb_with_replacement(n - i, k - 1)
        return preceding + Combination.with_replacement_rank(rest, n - j)

    @staticmethod
    def with_replacement_unrank(rank, n, k):
        """
        Find the combination with replacement of k integers from [0, n)
        with the given rank in a lexicographic ordering.
        """
        if k == 0:
            return []

        i = 0
        preceding = Combination.comb_with_replacement(n, k - 1)
        while rank >= preceding:
            rank -= preceding
            i += 1
            preceding = Combination.comb_with_replacement(n - i, k - 1)

        rest = Combination.with_replacement_unrank(rank, n - i, k - 1)
        return [i] + [x + i for x in rest]


def set_minus(arr, subset):
    return [x for x in arr if x not in set(subset)]


# TODO I think we can use part-count form everywhere. Right now
# there's a janky work-around of grouping the partition when
# we needed in part-count form but it doesn't look like there's any
# place that can't just accept it from the start.
def partitions(n):
    """
    Ascending integer partitions of n, excluding the partition [n].
    Since trees with unary nodes are uncountable, the partition of
    leaves must be at least size two.
    """
    if n > 0:
        # last partition is guaranteed to be length 1.
        yield from itertools.takewhile(lambda a: len(a) > 1, rule_asc(n))


def rule_asc(n):
    """
    Produce the integer partitions of n as ascending compositions.
    See: http://jeromekelleher.net/generating-integer-partitions.html
    """
    a = [0 for _ in range(n + 1)]
    k = 1
    a[1] = n
    while k != 0:
        x = a[k - 1] + 1
        y = a[k] - 1
        k -= 1
        while x <= y:
            a[k] = x
            y -= x
            k += 1
        a[k] = x + y
        yield a[: k + 1]


def group_by(values, equal):
    groups = []
    curr_group = []
    for x in values:
        if len(curr_group) == 0 or equal(x, curr_group[0]):
            curr_group.append(x)
        else:
            groups.append(curr_group)
            curr_group = [x]

    if len(curr_group) != 0:
        groups.append(curr_group)
    return groups


def group_partition(part):
    return group_by(part, lambda x, y: x == y)


def merge_tuple(tup1, tup2):
    return tuple(heapq.merge(tup1, tup2))


def isdisjoint(iterable1, iterable2):
    return set(iterable1).isdisjoint(iterable2)


--- ../../tskit/python/tskit/cli.py ---


"""
Command line utilities for tskit.
"""
import argparse
import json
import os
import signal
import sys

import tskit


def set_sigpipe_handler():
    if os.name == "posix":
        # Set signal handler for SIGPIPE to quietly kill the program.
        signal.signal(signal.SIGPIPE, signal.SIG_DFL)


def sys_exit(message):
    sys.exit(message)


def load_tree_sequence(path):
    try:
        return tskit.load(path)
    except OSError as e:
        sys_exit(f"Load error: {e}")


def run_info(args):
    print(load_tree_sequence(args.tree_sequence))


def run_trees(args):
    ts = load_tree_sequence(args.tree_sequence)
    for tree in ts.trees():
        print(f"tree {tree.index}:")
        print(f"  num_sites: {tree.num_sites}")
        print(
            "  interval:  ({0:.{2}f}, {1:.{2}f})".format(
                tree.interval.left, tree.interval.right, args.precision
            )
        )
        if args.draw:
            print(tree.draw(format="unicode"))


def run_upgrade(args):
    try:
        tree_sequence = tskit.load_legacy(args.source, args.remove_duplicate_positions)
        tree_sequence.dump(args.destination)
    except tskit.DuplicatePositionsError:
        sys_exit(
            "Error: Duplicate mutation positions in the source file detected.\n\n"
            'This is not supported in the current file format. Running "upgrade -d" '
            "will remove these duplicate positions. However, this will result in loss "
            "of data from the original file!"
        )


def run_individuals(args):
    tree_sequence = load_tree_sequence(args.tree_sequence)
    tree_sequence.dump_text(individuals=sys.stdout, precision=args.precision)


def run_nodes(args):
    tree_sequence = load_tree_sequence(args.tree_sequence)
    tree_sequence.dump_text(nodes=sys.stdout, precision=args.precision)


def run_edges(args):
    tree_sequence = load_tree_sequence(args.tree_sequence)
    tree_sequence.dump_text(edges=sys.stdout, precision=args.precision)


def run_sites(args):
    tree_sequence = load_tree_sequence(args.tree_sequence)
    tree_sequence.dump_text(sites=sys.stdout, precision=args.precision)


def run_mutations(args):
    tree_sequence = load_tree_sequence(args.tree_sequence)
    tree_sequence.dump_text(mutations=sys.stdout, precision=args.precision)


def run_populations(args):
    tree_sequence = load_tree_sequence(args.tree_sequence)
    tree_sequence.dump_text(populations=sys.stdout)


def run_migrations(args):
    tree_sequence = load_tree_sequence(args.tree_sequence)
    tree_sequence.dump_text(migrations=sys.stdout, precision=args.precision)


def run_provenances(args):
    tree_sequence = load_tree_sequence(args.tree_sequence)
    if args.human:
        for provenance in tree_sequence.provenances():
            d = json.loads(provenance.record)
            print(
                "id={}, timestamp={}, record={}".format(
                    provenance.id, provenance.timestamp, json.dumps(d, indent=4)
                )
            )
    else:
        tree_sequence.dump_text(provenances=sys.stdout)


def run_fasta(args):
    tree_sequence = load_tree_sequence(args.tree_sequence)
    tree_sequence.write_fasta(sys.stdout, wrap_width=args.wrap)


def run_vcf(args):
    tree_sequence = load_tree_sequence(args.tree_sequence)
    tree_sequence.write_vcf(
        sys.stdout,
        ploidy=args.ploidy,
        contig_id=args.contig_id,
        allow_position_zero=args.allow_position_zero,
    )


def add_tree_sequence_argument(parser):
    parser.add_argument("tree_sequence", help="The tskit tree sequence file")


def add_precision_argument(parser):
    parser.add_argument(
        "--precision",
        "-p",
        type=int,
        default=6,
        help="The number of decimal places to print in records",
    )


def get_tskit_parser():
    top_parser = argparse.ArgumentParser(
        prog="python3 -m tskit", description="Command line interface for tskit."
    )
    top_parser.add_argument(
        "-V", "--version", action="version", version=f"%(prog)s {tskit.__version__}"
    )
    subparsers = top_parser.add_subparsers(dest="subcommand")
    subparsers.required = True

    parser = subparsers.add_parser(
        "info", help="Print summary information about a tree sequence."
    )
    add_tree_sequence_argument(parser)
    parser.set_defaults(runner=run_info)

    parser = subparsers.add_parser("trees", help="Print information about trees.")
    add_tree_sequence_argument(parser)
    add_precision_argument(parser)
    parser.add_argument(
        "--draw", "-d", action="store_true", default=False, help="Draw the trees"
    )
    parser.set_defaults(runner=run_trees)

    parser = subparsers.add_parser(
        "upgrade", help="Upgrade legacy tree sequence files."
    )
    parser.add_argument(
        "source", help="The source tskit tree sequence file in legacy format"
    )
    parser.add_argument("destination", help="The filename of the upgraded copy.")
    parser.add_argument(
        "--remove-duplicate-positions",
        "-d",
        action="store_true",
        default=False,
        help="Remove any duplicated mutation positions in the source file. ",
    )
    parser.set_defaults(runner=run_upgrade)
    # suppress fasta visibility until we have a reference sequence
    # See https://github.com/tskit-dev/tskit/issues/1888
    # parser = subparsers.add_parser(
    #    "fasta",
    #     help="Convert the tree sequence haplotypes to fasta format")
    # add_tree_sequence_argument(parser)
    # parser.add_argument(
    #     "--wrap", "-w", type=int, default=60,
    #     help=("line-wrapping width for printed sequences"))
    # parser.set_defaults(runner=run_fasta)
    parser = subparsers.add_parser(
        "vcf", help="Convert the tree sequence genotypes to VCF format."
    )
    add_tree_sequence_argument(parser)
    parser.add_argument(
        "--ploidy",
        "-P",
        type=int,
        default=None,
        help=(
            "If the tree sequence does not contain information about "
            "individuals, create them by combining adjacent samples nodes "
            "into individuals of the specified ploidy. It is an error "
            "to provide this argument if the tree sequence does contain "
            "individuals"
        ),
    )
    parser.add_argument(
        "--contig-id", "-c", type=str, default="1", help="Specify the contig id"
    )
    parser.add_argument(
        "--allow-position-zero",
        "-0",
        action="store_true",
        default=False,
        help="Allow position 0 sites",
    )
    parser.set_defaults(runner=run_vcf)

    parser = subparsers.add_parser(
        "individuals", help="Output individuals in tabular format."
    )
    add_tree_sequence_argument(parser)
    add_precision_argument(parser)
    parser.set_defaults(runner=run_individuals)

    parser = subparsers.add_parser("nodes", help="Output nodes in tabular format.")
    add_tree_sequence_argument(parser)
    add_precision_argument(parser)
    parser.set_defaults(runner=run_nodes)

    parser = subparsers.add_parser("edges", help="Output edges in tabular format.")
    add_tree_sequence_argument(parser)
    add_precision_argument(parser)
    parser.set_defaults(runner=run_edges)

    parser = subparsers.add_parser("sites", help="Output sites in tabular format.")
    add_tree_sequence_argument(parser)
    add_precision_argument(parser)
    parser.set_defaults(runner=run_sites)

    parser = subparsers.add_parser(
        "mutations", help="Output mutations in tabular format."
    )
    add_tree_sequence_argument(parser)
    add_precision_argument(parser)
    parser.set_defaults(runner=run_mutations)

    parser = subparsers.add_parser(
        "populations", help="Output population information in tabular format."
    )
    add_tree_sequence_argument(parser)
    parser.set_defaults(runner=run_populations)

    parser = subparsers.add_parser(
        "migrations", help="Output migration information in tabular format."
    )
    add_tree_sequence_argument(parser)
    add_precision_argument(parser)
    parser.set_defaults(runner=run_migrations)

    parser = subparsers.add_parser(
        "provenances", help="Output provenance information in tabular format."
    )
    add_tree_sequence_argument(parser)
    parser.add_argument(
        "-H",
        "--human",
        action="store_true",
        help="Print out the provenances in a human readable format",
    )
    parser.set_defaults(runner=run_provenances)

    return top_parser


def tskit_main(arg_list=None):
    set_sigpipe_handler()
    parser = get_tskit_parser()
    args = parser.parse_args(arg_list)
    args.runner(args)


--- ../../tskit/python/tskit/provenance.py ---


"""
Common provenance methods used to determine the state and versions
of various dependencies and the OS.
"""
import json
import os.path
import platform
import sys
import time

try:
    import resource
except ImportError:
    resource = None  # resource.getrusage absent on windows

import jsonschema

import _tskit
import tskit.exceptions as exceptions
from . import _version

__version__ = _version.tskit_version


# NOTE: the APIs here are all preliminary. We should have a class that encapsulates
# all of the required functionality, including parsing and printing out provenance
# records. This will replace the current functions.


def get_environment(extra_libs=None, include_tskit=True):
    """
    Returns a dictionary describing the environment in which tskit
    is currently running.

    This API is tentative and will change in the future when a more
    comprehensive provenance API is implemented.
    """
    env = {
        "os": {
            "system": platform.system(),
            "node": platform.node(),
            "release": platform.release(),
            "version": platform.version(),
            "machine": platform.machine(),
        },
        "python": {
            "implementation": platform.python_implementation(),
            "version": platform.python_version(),
        },
    }
    libs = {"kastore": {"version": ".".join(map(str, _tskit.get_kastore_version()))}}
    if include_tskit:
        libs["tskit"] = {"version": __version__}
    if extra_libs is not None:
        libs.update(extra_libs)
    env["libraries"] = libs
    return env


def get_resources(start_time):
    # Returns a dict describing the resources used by the current process
    times = os.times()
    ret = {
        "elapsed_time": time.time() - start_time,
        "user_time": times.user + times.children_user,
        "sys_time": times.system + times.children_system,
    }
    if resource is not None:
        # Don't report max memory on Windows, we would need an external dep like psutil
        ret["max_memory"] = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
        if sys.platform != "darwin":
            ret["max_memory"] *= 1024  # Linux, freeBSD et al reports in KiB, not bytes

    return ret


def get_provenance_dict(parameters=None):
    """
    Returns a dictionary encoding an execution of tskit conforming to the
    provenance schema.
    """
    document = {
        "schema_version": "1.0.0",
        "software": {"name": "tskit", "version": __version__},
        "parameters": parameters,
        "environment": get_environment(include_tskit=False),
    }
    return document


# Cache the schema
_schema = None


def get_schema():
    """
    Returns the tskit provenance :ref:`provenance schema <sec_provenance>` as
    a dict.

    :return: The provenance schema.
    :rtype: dict
    """
    global _schema
    if _schema is None:
        base = os.path.dirname(__file__)
        schema_file = os.path.join(base, "provenance.schema.json")
        with open(schema_file) as f:
            _schema = json.load(f)
    # Return a copy to avoid issues with modifying the cached schema
    return dict(_schema)


def validate_provenance(provenance):
    """
    Validates the specified dict-like object against the tskit
    :ref:`provenance schema <sec_provenance>`. If the input does
    not represent a valid instance of the schema an exception is
    raised.

    :param dict provenance: The dictionary representing a JSON document
        to be validated against the schema.
    :raises ProvenanceValidationError: if the schema is not valid.
    """
    schema = get_schema()
    try:
        jsonschema.validate(provenance, schema)
    except jsonschema.exceptions.ValidationError as ve:
        raise exceptions.ProvenanceValidationError from ve


--- ../../tskit/python/tskit/text_formats.py ---


"""
Module responsible for working with text format data.
"""
import base64

import numpy as np

import tskit
from tskit import util


def parse_fam(fam_file):
    """
    Parse PLINK .fam file and convert to tskit IndividualTable.

    Assumes fam file contains five columns: FID, IID, PAT, MAT, SEX

    :param fam_file: PLINK .fam file object
    :param tskit.TableCollection tc: TableCollection with IndividualTable to
        which the individuals will be added
    """
    individuals = np.loadtxt(
        fname=fam_file,
        dtype=str,
        ndmin=2,  # read file as 2-D table
        usecols=(0, 1, 2, 3, 4),  # only keep FID, IID, PAT, MAT, SEX columns
    )  # requires same number of columns in each row, i.e. not ragged

    id_map = {}  # dict for translating PLINK ID to tskit IndividualTable ID
    for tskit_id, (plink_fid, plink_iid, _pat, _mat, _sex) in enumerate(individuals):
        # include space between strings to ensure uniqueness
        plink_id = f"{plink_fid} {plink_iid}"
        if plink_id in id_map:
            raise ValueError("Duplicate PLINK ID: {plink_id}")
        id_map[plink_id] = tskit_id
    id_map["0"] = -1  # -1 is used in tskit to denote "missing"

    tc = tskit.TableCollection(1)
    tb = tc.individuals
    tb.metadata_schema = tskit.MetadataSchema(
        {
            "codec": "json",
            "type": "object",
            "properties": {
                "plink_fid": {"type": "string"},
                "plink_iid": {"type": "string"},
                "sex": {"type": "integer"},
            },
            "required": ["plink_fid", "plink_iid", "sex"],
            "additionalProperties": True,
        }
    )
    for plink_fid, plink_iid, pat, mat, sex in individuals:
        sex = int(sex)
        if not (sex in range(3)):
            raise ValueError(
                "Sex must be one of the following: 0 (unknown), 1 (male), 2 (female)"
            )
        metadata_dict = {"plink_fid": plink_fid, "plink_iid": plink_iid, "sex": sex}
        pat_id = f"{plink_fid} {pat}" if pat != "0" else pat
        mat_id = f"{plink_fid} {mat}" if mat != "0" else mat
        tb.add_row(
            parents=[
                id_map[pat_id],
                id_map[mat_id],
            ],
            metadata=metadata_dict,
        )
    tc.sort()

    return tb


def flexible_file_output(ts_export_func):
    """
    Decorator to support writing to either an open file-like object
    or to a path. Assumes the second argument is the output.
    """

    def f(ts, file_or_path, **kwargs):
        file, local_file = util.convert_file_like_to_open_file(file_or_path, "w")
        try:
            ts_export_func(ts, file, **kwargs)
        finally:
            if local_file:
                file.close()

    return f


@flexible_file_output
def write_nexus(
    ts,
    out,
    *,
    precision,
    include_trees,
    include_alignments,
    reference_sequence,
    missing_data_character,
):
    # See TreeSequence.write_nexus for documentation on parameters.
    if precision is None:
        pos_precision = 0 if ts.discrete_genome else 17
        time_precision = None
    else:
        pos_precision = precision
        time_precision = precision

    indent = "  "
    print("#NEXUS", file=out)
    print("BEGIN TAXA;", file=out)
    print("", f"DIMENSIONS NTAX={ts.num_samples};", sep=indent, file=out)
    taxlabels = " ".join(f"n{u}" for u in ts.samples())
    print("", f"TAXLABELS {taxlabels};", sep=indent, file=out)
    print("END;", file=out)

    if include_alignments is None:
        include_alignments = ts.discrete_genome and ts.num_sites > 0
    if include_alignments:
        missing_data_character = (
            "?" if missing_data_character is None else missing_data_character
        )
        print("BEGIN DATA;", file=out)
        print("", f"DIMENSIONS NCHAR={int(ts.sequence_length)};", sep=indent, file=out)
        print(
            "",
            f"FORMAT DATATYPE=DNA MISSING={missing_data_character};",
            sep=indent,
            file=out,
        )
        print("", "MATRIX", file=out, sep=indent)
        alignments = ts.alignments(
            reference_sequence=reference_sequence,
            missing_data_character=missing_data_character,
        )
        for u, alignment in zip(ts.samples(), alignments):
            print(2 * indent, f"n{u}", " ", alignment, sep="", file=out)
        print("", ";", sep=indent, file=out)
        print("END;", file=out)

    include_trees = True if include_trees is None else include_trees
    if include_trees:
        print("BEGIN TREES;", file=out)
        for tree in ts.trees():
            start_interval = "{0:.{1}f}".format(tree.interval.left, pos_precision)
            end_interval = "{0:.{1}f}".format(tree.interval.right, pos_precision)
            tree_label = f"t{start_interval}^{end_interval}"
            newick = tree.as_newick(precision=time_precision)
            print("", f"TREE {tree_label} = [&R] {newick}", sep=indent, file=out)
        print("END;", file=out)


def wrap_text(text, width):
    """
    Return an iterator over the lines in the specified string of at most the
    specified width. (We could use textwrap.wrap for this, but it uses a
    more complicated algorithm appropriate for blocks of words.)
    """
    width = len(text) if width == 0 else width
    N = len(text) // width
    offset = 0
    for _ in range(N):
        yield text[offset : offset + width]
        offset += width
    if offset != len(text):
        yield text[offset:]


@flexible_file_output
def write_fasta(
    ts,
    output,
    *,
    wrap_width,
    reference_sequence,
    missing_data_character,
):
    # See TreeSequence.write_fasta for documentation
    if wrap_width < 0 or int(wrap_width) != wrap_width:
        raise ValueError(
            "wrap_width must be a non-negative integer. "
            "You may specify `wrap_width=0` "
            "if you do not want any wrapping."
        )
    wrap_width = int(wrap_width)
    alignments = ts.alignments(
        reference_sequence=reference_sequence,
        missing_data_character=missing_data_character,
    )
    for u, alignment in zip(ts.samples(), alignments):
        print(">", f"n{u}", sep="", file=output)
        for line in wrap_text(alignment, wrap_width):
            print(line, file=output)


def _build_newick(tree, *, node, precision, node_labels, include_branch_lengths):
    label = node_labels.get(node, "")
    if tree.is_leaf(node):
        s = f"{label}"
    else:
        s = "("
        for child in tree.children(node):
            branch_length = tree.branch_length(child)
            subtree = _build_newick(
                tree,
                node=child,
                precision=precision,
                node_labels=node_labels,
                include_branch_lengths=include_branch_lengths,
            )
            if include_branch_lengths:
                subtree += ":{0:.{1}f}".format(branch_length, precision)
            s += subtree + ","
        s = s[:-1] + f"){label}"
    return s


def build_newick(tree, *, root, precision, node_labels, include_branch_lengths):
    """
    Simple recursive version of the newick generator used when non-default
    node labels are needed, or when branch lengths are omitted
    """
    s = _build_newick(
        tree,
        node=root,
        precision=precision,
        node_labels=node_labels,
        include_branch_lengths=include_branch_lengths,
    )
    return s + ";"


def dump_text(
    ts,
    *,
    nodes,
    edges,
    sites,
    mutations,
    individuals,
    populations,
    migrations,
    provenances,
    precision,
    encoding,
    base64_metadata,
):
    if nodes is not None:
        print(
            "id",
            "is_sample",
            "time",
            "population",
            "individual",
            "metadata",
            sep="\t",
            file=nodes,
        )
        for node in ts.nodes():
            metadata = text_metadata(base64_metadata, encoding, node)
            row = (
                "{id:d}\t"
                "{is_sample:d}\t"
                "{time:.{precision}f}\t"
                "{population:d}\t"
                "{individual:d}\t"
                "{metadata}"
            ).format(
                precision=precision,
                id=node.id,
                is_sample=node.is_sample(),
                time=node.time,
                population=node.population,
                individual=node.individual,
                metadata=metadata,
            )
            print(row, file=nodes)

    if edges is not None:
        print("left", "right", "parent", "child", "metadata", sep="\t", file=edges)
        for edge in ts.edges():
            metadata = text_metadata(base64_metadata, encoding, edge)
            row = (
                "{left:.{precision}f}\t"
                "{right:.{precision}f}\t"
                "{parent:d}\t"
                "{child:d}\t"
                "{metadata}"
            ).format(
                precision=precision,
                left=edge.left,
                right=edge.right,
                parent=edge.parent,
                child=edge.child,
                metadata=metadata,
            )
            print(row, file=edges)

    if sites is not None:
        print("position", "ancestral_state", "metadata", sep="\t", file=sites)
        for site in ts.sites():
            metadata = text_metadata(base64_metadata, encoding, site)
            row = (
                "{position:.{precision}f}\t" "{ancestral_state}\t" "{metadata}"
            ).format(
                precision=precision,
                position=site.position,
                ancestral_state=site.ancestral_state,
                metadata=metadata,
            )
            print(row, file=sites)

    if mutations is not None:
        print(
            "site",
            "node",
            "time",
            "derived_state",
            "parent",
            "metadata",
            sep="\t",
            file=mutations,
        )
        for site in ts.sites():
            for mutation in site.mutations:
                metadata = text_metadata(base64_metadata, encoding, mutation)
                row = (
                    "{site}\t"
                    "{node}\t"
                    "{time}\t"
                    "{derived_state}\t"
                    "{parent}\t"
                    "{metadata}"
                ).format(
                    site=mutation.site,
                    node=mutation.node,
                    time=(
                        "unknown"
                        if util.is_unknown_time(mutation.time)
                        else mutation.time
                    ),
                    derived_state=mutation.derived_state,
                    parent=mutation.parent,
                    metadata=metadata,
                )
                print(row, file=mutations)

    if individuals is not None:
        print(
            "id",
            "flags",
            "location",
            "parents",
            "metadata",
            sep="\t",
            file=individuals,
        )
        for individual in ts.individuals():
            metadata = text_metadata(base64_metadata, encoding, individual)
            location = ",".join(map(str, individual.location))
            parents = ",".join(map(str, individual.parents))
            row = (
                "{id}\t" "{flags}\t" "{location}\t" "{parents}\t" "{metadata}"
            ).format(
                id=individual.id,
                flags=individual.flags,
                location=location,
                parents=parents,
                metadata=metadata,
            )
            print(row, file=individuals)

    if populations is not None:
        print("id", "metadata", sep="\t", file=populations)
        for population in ts.populations():
            metadata = text_metadata(base64_metadata, encoding, population)
            row = ("{id}\t" "{metadata}").format(id=population.id, metadata=metadata)
            print(row, file=populations)

    if migrations is not None:
        print(
            "left",
            "right",
            "node",
            "source",
            "dest",
            "time",
            "metadata",
            sep="\t",
            file=migrations,
        )
        for migration in ts.migrations():
            metadata = text_metadata(base64_metadata, encoding, migration)
            row = (
                "{left}\t"
                "{right}\t"
                "{node}\t"
                "{source}\t"
                "{dest}\t"
                "{time}\t"
                "{metadata}\t"
            ).format(
                left=migration.left,
                right=migration.right,
                node=migration.node,
                source=migration.source,
                dest=migration.dest,
                time=migration.time,
                metadata=metadata,
            )
            print(row, file=migrations)

    if provenances is not None:
        print("id", "timestamp", "record", sep="\t", file=provenances)
        for provenance in ts.provenances():
            row = ("{id}\t" "{timestamp}\t" "{record}\t").format(
                id=provenance.id,
                timestamp=provenance.timestamp,
                record=provenance.record,
            )
            print(row, file=provenances)


def text_metadata(base64_metadata, encoding, node):
    metadata = node.metadata
    if isinstance(metadata, bytes) and base64_metadata:
        metadata = base64.b64encode(metadata).decode(encoding)
    else:
        metadata = repr(metadata)
    return metadata


--- ../../tskit/python/tskit/vcf.py ---


"""
Convert tree sequences to VCF.
"""
import numpy as np

import tskit
from . import provenance


def legacy_position_transform(positions):
    """
    Transforms positions in the tree sequence into VCF coordinates under
    the pre 0.2.0 legacy rule.
    """
    last_pos = 0
    transformed = []
    for pos in positions:
        pos = int(round(pos))
        if pos <= last_pos:
            pos = last_pos + 1
        transformed.append(pos)
        last_pos = pos
    return transformed


class VcfWriter:
    """
    Writes a VCF representation of the genotypes tree sequence to a
    file-like object.
    """

    def __init__(
        self,
        tree_sequence,
        *,
        ploidy,
        contig_id,
        individuals,
        individual_names,
        position_transform,
        site_mask,
        sample_mask,
        isolated_as_missing,
        allow_position_zero,
    ):
        self.tree_sequence = tree_sequence
        self.contig_id = contig_id
        self.isolated_as_missing = isolated_as_missing

        self.__make_sample_mapping(ploidy, individuals)
        if individual_names is None:
            individual_names = [f"tsk_{j}" for j in range(self.num_individuals)]
        self.individual_names = individual_names
        if len(self.individual_names) != self.num_individuals:
            raise ValueError(
                "individual_names must have length equal to the number of individuals"
            )

        # Transform coordinates for VCF
        if position_transform is None:
            position_transform = np.round
        elif position_transform == "legacy":
            position_transform = legacy_position_transform
        self.transformed_positions = np.array(
            position_transform(tree_sequence.tables.sites.position), dtype=int
        )
        if self.transformed_positions.shape != (tree_sequence.num_sites,):
            raise ValueError(
                "Position transform must return an array of the same length"
            )
        self.contig_length = max(
            1, int(position_transform([tree_sequence.sequence_length])[0])
        )
        if len(self.transformed_positions) > 0:
            # Arguably this should be last_pos + 1, but if we hit this
            # condition the coordinate systems are all muddled up anyway
            # so it's simpler to stay with this rule that was inherited
            # from the legacy VCF output code.
            self.contig_length = max(self.transformed_positions[-1], self.contig_length)

        if site_mask is None:
            site_mask = np.zeros(tree_sequence.num_sites, dtype=bool)
        self.site_mask = np.array(site_mask, dtype=bool)
        if self.site_mask.shape != (tree_sequence.num_sites,):
            raise ValueError("Site mask must be 1D a boolean array of length num_sites")

        self.sample_mask = sample_mask
        if sample_mask is not None:
            if not callable(sample_mask):
                sample_mask = np.array(sample_mask, dtype=bool)
                self.sample_mask = lambda _: sample_mask

        # The VCF spec does not allow for positions to be 0, so we error if one of the
        # transformed positions is 0 and allow_position_zero is False.
        if not allow_position_zero and np.any(
            self.transformed_positions[~site_mask] == 0
        ):
            raise ValueError(
                "A variant position of 0 was found in the VCF output, this is not "
                "fully compliant with the VCF spec. If you still wish to write the VCF "
                'please use the "allow_position_zero" argument to write_vcf. '
                "Alternatively, you can increment all the positions by one using "
                '"position_transform = lambda x: 1 + x" or coerce the zero to one with '
                '"position_transform = lambda x: np.fmax(1, x)"'
            )

    def __make_sample_mapping(self, ploidy, individuals):
        """
        Compute the sample IDs for each VCF individual and the template for
        writing out genotypes.
        """
        ts = self.tree_sequence
        self.samples = None
        self.individual_ploidies = []

        # Cannot use "ploidy" when *any* individuals are present.
        if ts.num_individuals > 0 and ploidy is not None:
            raise ValueError(
                "Cannot specify ploidy when individuals are present in tables "
            )

        if individuals is None:
            # Find all sample nodes that reference individuals
            individuals = np.unique(ts.nodes_individual[ts.samples()])
            if len(individuals) == 1 and individuals[0] == tskit.NULL:
                # No samples refer to individuals
                individuals = None
            else:
                # np.unique sorts the argument, so if NULL (-1) is present it
                # will be the first value.
                if individuals[0] == tskit.NULL:
                    raise ValueError(
                        "Sample nodes must either all be associated with individuals "
                        "or not associated with any individuals"
                    )
        else:
            individuals = np.array(individuals, dtype=np.int32)
            if len(individuals) == 0:
                raise ValueError("List of sample individuals empty")

        if individuals is not None:
            self.samples = []
            # FIXME this could probably be done more efficiently.
            for i in individuals:
                if i < 0 or i >= self.tree_sequence.num_individuals:
                    raise ValueError("Invalid individual IDs provided.")
                ind = self.tree_sequence.individual(i)
                if len(ind.nodes) == 0:
                    raise ValueError(f"Individual {i} not associated with a node")
                is_sample = {ts.node(u).is_sample() for u in ind.nodes}
                if len(is_sample) != 1:
                    raise ValueError(
                        f"Individual {ind.id} has nodes that are sample and "
                        "non-samples"
                    )
                self.samples.extend(ind.nodes)
                self.individual_ploidies.append(len(ind.nodes))
        else:
            if ploidy is None:
                ploidy = 1
            if ploidy < 1:
                raise ValueError("Ploidy must be >= 1")
            if ts.num_samples % ploidy != 0:
                raise ValueError("Sample size must be divisible by ploidy")
            self.individual_ploidies = np.full(
                ts.sample_size // ploidy, ploidy, dtype=np.int32
            )
        self.num_individuals = len(self.individual_ploidies)

    def __write_header(self, output):
        print("##fileformat=VCFv4.2", file=output)
        print(f"##source=tskit {provenance.__version__}", file=output)
        print('##FILTER=<ID=PASS,Description="All filters passed">', file=output)
        print(
            f"##contig=<ID={self.contig_id},length={self.contig_length}>", file=output
        )
        print(
            '##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">', file=output
        )
        vcf_samples = "\t".join(self.individual_names)
        print(
            "#CHROM",
            "POS",
            "ID",
            "REF",
            "ALT",
            "QUAL",
            "FILTER",
            "INFO",
            "FORMAT",
            vcf_samples,
            sep="\t",
            file=output,
        )

    def write(self, output):
        self.__write_header(output)

        # Build the array for hold the text genotype VCF data and the indexes into
        # this array for when we're updating it.
        gt_array = []
        indexes = []
        for ploidy in self.individual_ploidies:
            for _ in range(ploidy):
                indexes.append(len(gt_array))
                # First element here is a placeholder that we'll write the actual
                # genotypes into when for each variant.
                gt_array.extend([0, ord("|")])
            gt_array[-1] = ord("\t")
        gt_array[-1] = ord("\n")
        gt_array = np.array(gt_array, dtype=np.int8)
        # TODO Unclear here whether using int64 or int32 will be faster for this index
        # array. Test it out.
        indexes = np.array(indexes, dtype=int)

        for variant in self.tree_sequence.variants(
            samples=self.samples, isolated_as_missing=self.isolated_as_missing
        ):
            site_id = variant.site.id
            # We check the mask before we do any checks so we can use this as a
            # way of skipping problematic sites.
            if self.site_mask[site_id]:
                continue

            if variant.num_alleles > 9:
                raise ValueError(
                    "More than 9 alleles not currently supported. Please open an issue "
                    "on GitHub if this limitation affects you."
                )
            pos = self.transformed_positions[variant.index]
            ref = variant.alleles[0]
            alt = "."
            if variant.num_alleles > 1:
                alt = ",".join(variant.alleles[1 : variant.num_alleles])
            print(
                self.contig_id,
                pos,
                site_id,
                ref,
                alt,
                ".",
                "PASS",
                ".",
                "GT",
                sep="\t",
                end="\t",
                file=output,
            )
            genotypes = variant.genotypes
            gt_array[indexes] = genotypes + ord("0")
            if self.sample_mask is not None:
                genotypes = genotypes.copy()
                sample_mask = np.array(self.sample_mask(variant), dtype=bool)
                if sample_mask.shape != genotypes.shape:
                    raise ValueError(
                        "Sample mask must be a numpy array of size num_samples"
                    )
                genotypes[sample_mask] = -1
            if self.sample_mask is not None or variant.has_missing_data:
                missing = genotypes == -1
                gt_array[indexes[missing]] = ord(".")
            g_bytes = memoryview(gt_array).tobytes()
            g_str = g_bytes.decode()
            print(g_str, end="", file=output)


--- ../../tskit/python/tskit/exceptions.py ---


"""
Exceptions defined in tskit.
"""
from _tskit import FileFormatError  # noqa: F401
from _tskit import IdentityPairsNotStoredError  # noqa: F401
from _tskit import IdentitySegmentsNotStoredError  # noqa: F401
from _tskit import LibraryError  # noqa: F401
from _tskit import TskitException  # noqa: F401
from _tskit import VersionTooNewError  # noqa: F401
from _tskit import VersionTooOldError  # noqa: F401


class DuplicatePositionsError(TskitException):
    """
    Duplicate positions in the list of sites.
    """


class ProvenanceValidationError(TskitException):
    """
    A JSON document did not validate against the provenance schema.
    """


class MetadataValidationError(TskitException):
    """
    A metadata object did not validate against the provenance schema.
    """


class MetadataSchemaValidationError(TskitException):
    """
    A metadata schema object did not validate against the metaschema.
    """


class MetadataEncodingError(TskitException):
    """
    A metadata object was of a type that could not be encoded
    """


--- ../../tskit/python/tskit/tables.py ---


"""
Tree sequence IO via the tables API.
"""
import collections.abc
import dataclasses
import datetime
import json
import numbers
import warnings
from collections.abc import Mapping
from dataclasses import dataclass
from functools import reduce
from typing import Dict
from typing import Optional
from typing import Union

import numpy as np

import _tskit
import tskit
import tskit.metadata as metadata
import tskit.provenance as provenance
import tskit.util as util
from tskit import UNKNOWN_TIME

dataclass_options = {"frozen": True}


# Needed for cases where `None` can be an appropriate kwarg value,
# we override the meta so that it looks good in the docs.
class NotSetMeta(type):
    def __repr__(cls):
        return "Not set"


class NOTSET(metaclass=NotSetMeta):
    pass


@metadata.lazy_decode()
@dataclass(**dataclass_options)
class IndividualTableRow(util.Dataclass):
    """
    A row in an :class:`IndividualTable`.
    """

    __slots__ = ["flags", "location", "parents", "metadata"]
    flags: int
    """
    See :attr:`Individual.flags`
    """
    location: np.ndarray
    """
    See :attr:`Individual.location`
    """
    parents: np.ndarray
    """
    See :attr:`Individual.parents`
    """
    metadata: Optional[Union[bytes, dict]]
    """
    See :attr:`Individual.metadata`
    """

    # We need a custom eq for the numpy arrays
    def __eq__(self, other):
        return (
            isinstance(other, IndividualTableRow)
            and self.flags == other.flags
            and np.array_equal(self.location, other.location)
            and np.array_equal(self.parents, other.parents)
            and self.metadata == other.metadata
        )


@metadata.lazy_decode()
@dataclass(**dataclass_options)
class NodeTableRow(util.Dataclass):
    """
    A row in a :class:`NodeTable`.
    """

    __slots__ = ["flags", "time", "population", "individual", "metadata"]
    flags: int
    """
    See :attr:`Node.flags`
    """
    time: float
    """
    See :attr:`Node.time`
    """
    population: int
    """
    See :attr:`Node.population`
    """
    individual: int
    """
    See :attr:`Node.individual`
    """
    metadata: Optional[Union[bytes, dict]]
    """
    See :attr:`Node.metadata`
    """


@metadata.lazy_decode()
@dataclass(**dataclass_options)
class EdgeTableRow(util.Dataclass):
    """
    A row in an :class:`EdgeTable`.
    """

    __slots__ = ["left", "right", "parent", "child", "metadata"]
    left: float
    """
    See :attr:`Edge.left`
    """
    right: float
    """
    See :attr:`Edge.right`
    """
    parent: int
    """
    See :attr:`Edge.parent`
    """
    child: int
    """
    See :attr:`Edge.child`
    """
    metadata: Optional[Union[bytes, dict]]
    """
    See :attr:`Edge.metadata`
    """


@metadata.lazy_decode()
@dataclass(**dataclass_options)
class MigrationTableRow(util.Dataclass):
    """
    A row in a :class:`MigrationTable`.
    """

    __slots__ = ["left", "right", "node", "source", "dest", "time", "metadata"]
    left: float
    """
    See :attr:`Migration.left`
    """
    right: float
    """
    See :attr:`Migration.right`
    """
    node: int
    """
    See :attr:`Migration.node`
    """
    source: int
    """
    See :attr:`Migration.source`
    """
    dest: int
    """
    See :attr:`Migration.dest`
    """
    time: float
    """
    See :attr:`Migration.time`
    """
    metadata: Optional[Union[bytes, dict]]
    """
    See :attr:`Migration.metadata`
    """


@metadata.lazy_decode()
@dataclass(**dataclass_options)
class SiteTableRow(util.Dataclass):
    """
    A row in a :class:`SiteTable`.
    """

    __slots__ = ["position", "ancestral_state", "metadata"]
    position: float
    """
    See :attr:`Site.position`
    """
    ancestral_state: str
    """
    See :attr:`Site.ancestral_state`
    """
    metadata: Optional[Union[bytes, dict]]
    """
    See :attr:`Site.metadata`
    """


@metadata.lazy_decode()
@dataclass(**dataclass_options)
class MutationTableRow(util.Dataclass):
    """
    A row in a :class:`MutationTable`.
    """

    __slots__ = ["site", "node", "derived_state", "parent", "metadata", "time"]
    site: int
    """
    See :attr:`Mutation.site`
    """
    node: int
    """
    See :attr:`Mutation.node`
    """
    derived_state: str
    """
    See :attr:`Mutation.derived_state`
    """
    parent: int
    """
    See :attr:`Mutation.parent`
    """
    metadata: Optional[Union[bytes, dict]]
    """
    See :attr:`Mutation.metadata`
    """
    time: float
    """
    See :attr:`Mutation.time`
    """

    # We need a custom eq here as we have unknown times (nans) to check
    def __eq__(self, other):
        return (
            isinstance(other, MutationTableRow)
            and self.site == other.site
            and self.node == other.node
            and self.derived_state == other.derived_state
            and self.parent == other.parent
            and self.metadata == other.metadata
            and (
                self.time == other.time
                or (
                    util.is_unknown_time(self.time) and util.is_unknown_time(other.time)
                )
            )
        )


@metadata.lazy_decode()
@dataclass(**dataclass_options)
class PopulationTableRow(util.Dataclass):
    """
    A row in a :class:`PopulationTable`.
    """

    __slots__ = ["metadata"]
    metadata: Optional[Union[bytes, dict]]
    """
    See :attr:`Population.metadata`
    """


@dataclass(**dataclass_options)
class ProvenanceTableRow(util.Dataclass):
    """
    A row in a :class:`ProvenanceTable`.
    """

    __slots__ = ["timestamp", "record"]
    timestamp: str
    """
    See :attr:`Provenance.timestamp`
    """
    record: str
    """
    See :attr:`Provenance.record`
    """


@dataclass(**dataclass_options)
class TableCollectionIndexes(util.Dataclass):
    """
    A class encapsulating the indexes of a :class:`TableCollection`
    """

    edge_insertion_order: np.ndarray = None
    edge_removal_order: np.ndarray = None

    def asdict(self):
        return {k: v for k, v in dataclasses.asdict(self).items() if v is not None}

    @property
    def nbytes(self) -> int:
        """
        The number of bytes taken by the indexes
        """
        total = 0
        if self.edge_removal_order is not None:
            total += self.edge_removal_order.nbytes
        if self.edge_insertion_order is not None:
            total += self.edge_insertion_order.nbytes
        return total


def keep_with_offset(keep, data, offset):
    """
    Used when filtering _offset columns in tables
    """
    # We need the astype here for 32 bit machines
    lens = np.diff(offset).astype(np.int32)
    return (
        data[np.repeat(keep, lens)],
        np.concatenate(
            [
                np.array([0], dtype=offset.dtype),
                np.cumsum(lens[keep], dtype=offset.dtype),
            ]
        ),
    )


class BaseTable:
    """
    Superclass of high-level tables. Not intended for direct instantiation.
    """

    # The list of columns in the table. Must be set by subclasses.
    column_names = []

    def __init__(self, ll_table, row_class):
        self.ll_table = ll_table
        self.row_class = row_class

    def _check_required_args(self, **kwargs):
        for k, v in kwargs.items():
            if v is None:
                raise TypeError(f"{k} is required")

    @property
    def num_rows(self) -> int:
        return self.ll_table.num_rows

    @property
    def max_rows(self) -> int:
        return self.ll_table.max_rows

    @property
    def max_rows_increment(self) -> int:
        return self.ll_table.max_rows_increment

    @property
    def nbytes(self) -> int:
        """
        Returns the total number of bytes required to store the data
        in this table. Note that this may not be equal to
        the actual memory footprint.
        """
        # It's not ideal that we run asdict() here to do this as we're
        # currently creating copies of the column arrays, so it would
        # be more efficient to have dedicated low-level methods. However,
        # if we do have read-only views on the underlying memory for the
        # column arrays then this will be a perfectly good way of
        # computing the nbytes values and the overhead minimal.
        d = self.asdict()
        nbytes = 0
        # Some tables don't have a metadata_schema
        metadata_schema = d.pop("metadata_schema", None)
        if metadata_schema is not None:
            nbytes += len(metadata_schema.encode())
        nbytes += sum(col.nbytes for col in d.values())
        return nbytes

    def equals(self, other, ignore_metadata=False):
        """
        Returns True if  `self` and `other` are equal. By default, two tables
        are considered equal if their columns and metadata schemas are
        byte-for-byte identical.

        :param other: Another table instance
        :param bool ignore_metadata: If True exclude metadata and metadata schemas
            from the comparison.
        :return: True if other is equal to this table; False otherwise.
        :rtype: bool
        """
        # Note: most tables support ignore_metadata, we can override for those that don't
        ret = False
        if type(other) is type(self):
            ret = bool(
                self.ll_table.equals(other.ll_table, ignore_metadata=ignore_metadata)
            )
        return ret

    def assert_equals(self, other, *, ignore_metadata=False):
        """
        Raise an AssertionError for the first found difference between
        this and another table of the same type.

        :param other: Another table instance
        :param bool ignore_metadata: If True exclude metadata and metadata schemas
            from the comparison.
        """
        if type(other) is not type(self):
            raise AssertionError(f"Types differ: self={type(self)} other={type(other)}")

        # Check using the low-level method to avoid slowly going through everything
        if self.equals(other, ignore_metadata=ignore_metadata):
            return

        if not ignore_metadata and self.metadata_schema != other.metadata_schema:
            raise AssertionError(
                f"{type(self).__name__} metadata schemas differ: "
                f"self={self.metadata_schema} "
                f"other={other.metadata_schema}"
            )

        for n, (row_self, row_other) in enumerate(zip(self, other)):
            if ignore_metadata:
                row_self = dataclasses.replace(row_self, metadata=None)
                row_other = dataclasses.replace(row_other, metadata=None)
            if row_self != row_other:
                self_dict = dataclasses.asdict(self[n])
                other_dict = dataclasses.asdict(other[n])
                diff_string = []
                for col in self_dict.keys():
                    if isinstance(self_dict[col], np.ndarray):
                        equal = np.array_equal(self_dict[col], other_dict[col])
                    else:
                        equal = self_dict[col] == other_dict[col]
                    if not equal:
                        diff_string.append(
                            f"self.{col}={self_dict[col]} other.{col}={other_dict[col]}"
                        )
                diff_string = "\n".join(diff_string)
                raise AssertionError(
                    f"{type(self).__name__} row {n} differs:\n{diff_string}"
                )

        if self.num_rows != other.num_rows:
            raise AssertionError(
                f"{type(self).__name__} number of rows differ: self={self.num_rows} "
                f"other={other.num_rows}"
            )

        raise AssertionError(
            "Tables differ in an undetected way - "
            "this is a bug, please report an issue on gitub"
        )  # pragma: no cover

    def __eq__(self, other):
        return self.equals(other)

    def __len__(self):
        return self.num_rows

    def __getattr__(self, name):
        if name in self.column_names:
            return getattr(self.ll_table, name)
        else:
            raise AttributeError(
                f"{self.__class__.__name__} object has no attribute {name}"
            )

    def __setattr__(self, name, value):
        if name in self.column_names:
            d = self.asdict()
            d[name] = value
            self.set_columns(**d)
        else:
            object.__setattr__(self, name, value)

    def _make_row(self, *args):
        return self.row_class(*args)

    def __getitem__(self, index):
        """
        If passed an integer, return the specified row of this table, decoding metadata
        if it is present. Supports negative indexing, e.g. ``table[-5]``.
        If passed a slice, iterable or array return a new table containing the specified
        rows. Similar to numpy fancy indexing, if the array or iterables contains
        booleans then the index acts as a mask, returning those rows for which the mask
        is True. Note that as the result is a new table, the row ids will change as tskit
        row ids are row indexes.

        :param index: the index of a desired row, a slice of the desired rows, an
            iterable or array of the desired row numbers, or a boolean array to use as
            a mask.
        """

        if isinstance(index, numbers.Integral):
            # Single row by integer
            if index < 0:
                index += len(self)
            if index < 0 or index >= len(self):
                raise IndexError("Index out of bounds")
            return self._make_row(*self.ll_table.get_row(index))
        elif isinstance(index, numbers.Number):
            raise TypeError("Index must be integer, slice or iterable")
        elif isinstance(index, slice):
            index = range(*index.indices(len(self)))
        else:
            index = np.asarray(index)
            if index.dtype == np.bool_:
                if len(index) != len(self):
                    raise IndexError("Boolean index must be same length as table")
                index = np.flatnonzero(index)
            index = util.safe_np_int_cast(index, np.int32)

        ret = self.__class__()
        ret.metadata_schema = self.metadata_schema
        ret.ll_table.extend(self.ll_table, row_indexes=index)

        return ret

    def __setitem__(self, index, new_row):
        """
        Replaces a row of this table at the specified index with information from a
        row-like object. Metadata, will be validated and encoded according to the table's
        :attr:`metadata_schema<tskit.IndividualTable.metadata_schema>`.

        :param index: the index of the row to change
        :param row-like new_row: An object that has attributes corresponding to the
            properties of the new row. Both the objects returned from ``table[i]`` and
            e.g. ``ts.individual(i)`` work for this purpose, along with any other
            object with the correct attributes.
        """
        if isinstance(index, numbers.Integral):
            # Single row by integer
            if index < 0:
                index += len(self)
            if index < 0 or index >= len(self):
                raise IndexError("Index out of bounds")
        else:
            raise TypeError("Index must be integer")

        row_data = {
            column: getattr(new_row, column)
            for column in self.column_names
            if "_offset" not in column
        }

        # Encode the metadata - note that if this becomes a perf bottleneck it is
        # possible to use the cached, encoded metadata in the row object, rather than
        # decode and reencode
        if "metadata" in row_data:
            row_data["metadata"] = self.metadata_schema.validate_and_encode_row(
                row_data["metadata"]
            )

        self.ll_table.update_row(row_index=index, **row_data)

    def append(self, row):
        """
        Adds a new row to this table and returns the ID of the new row. Metadata, if
        specified, will be validated and encoded according to the table's
        :attr:`metadata_schema<tskit.IndividualTable.metadata_schema>`.

        :param row-like row: An object that has attributes corresponding to the
            properties of the new row. Both the objects returned from ``table[i]`` and
            e.g. ``ts.individual(i)`` work for this purpose, along with any other
            object with the correct attributes.
        :return: The index of the newly added row.
        :rtype: int
        """
        return self.add_row(
            **{
                column: getattr(row, column)
                for column in self.column_names
                if "_offset" not in column
            }
        )

    def replace_with(self, other):
        # Overwrite the contents of this table with a copy of the other table
        self.set_columns(**other.asdict())

    def clear(self):
        """
        Deletes all rows in this table.
        """
        self.ll_table.clear()

    def reset(self):
        # Deprecated alias for clear
        self.clear()

    def truncate(self, num_rows):
        """
        Truncates this table so that the only the first ``num_rows`` are retained.

        :param int num_rows: The number of rows to retain in this table.
        """
        return self.ll_table.truncate(num_rows)

    def keep_rows(self, keep):
        """
        .. include:: substitutions/table_keep_rows_main.rst

        :param array-like keep: The rows to keep as a boolean array. Must
            be the same length as the table, and convertible to a numpy
            array of dtype bool.
        :return: The mapping between old and new row IDs as a numpy
            array (dtype int32).
        :rtype: numpy.ndarray (dtype=np.int32)
        """
        # We do this check here rather than in the C code because calling
        # len() on the input will cause a more readable exception to be
        # raised than the inscrutable errors we get from numpy when
        # converting arguments of the wrong type.
        if len(keep) != len(self):
            msg = (
                "Argument for keep_rows must be a boolean array of "
                "the same length as the table. "
                f"(need:{len(self)}, got:{len(keep)})"
            )
            raise ValueError(msg)
        return self.ll_table.keep_rows(keep)

    # Pickle support
    def __getstate__(self):
        return self.asdict()

    # Unpickle support
    def __setstate__(self, state):
        self.__init__()
        self.set_columns(**state)

    def copy(self):
        """
        Returns a deep copy of this table
        """
        copy = self.__class__()
        copy.set_columns(**self.asdict())
        return copy

    def asdict(self):
        """
        Returns a dictionary mapping the names of the columns in this table
        to the corresponding numpy arrays.
        """
        ret = {col: getattr(self, col) for col in self.column_names}
        # Not all tables have metadata
        try:
            ret["metadata_schema"] = repr(self.metadata_schema)
        except AttributeError:
            pass
        return ret

    def set_columns(self, **kwargs):
        """
        Sets the values for each column in this :class:`Table` using values
        provided in numpy arrays. Overwrites existing data in all the table columns.
        """
        raise NotImplementedError()

    def __str__(self):
        headers, rows = self._text_header_and_rows(
            limit=tskit._print_options["max_lines"]
        )
        return util.unicode_table(rows, header=headers, row_separator=False)

    def _repr_html_(self):
        """
        Called e.g. by jupyter notebooks to render tables
        """
        headers, rows = self._text_header_and_rows(
            limit=tskit._print_options["max_lines"]
        )
        return util.html_table(rows, header=headers)

    def _columns_all_integer(self, *colnames):
        # For displaying floating point values without loads of decimal places
        return all(
            np.all(getattr(self, col) == np.floor(getattr(self, col)))
            for col in colnames
        )


class MetadataTable(BaseTable):
    """
    Base class for tables that have a metadata column.
    """

    # TODO this class has some overlap with the MetadataProvider base class
    # and also the TreeSequence class. These all have methods to deal with
    # schemas and essentially do the same thing (provide a facade for the
    # low-level get/set metadata schemas functionality). We should refactor
    # this so we're only doing it in one place.
    # https://github.com/tskit-dev/tskit/issues/1957
    def __init__(self, ll_table, row_class):
        super().__init__(ll_table, row_class)

    def _make_row(self, *args):
        return self.row_class(*args, metadata_decoder=self.metadata_schema.decode_row)

    def packset_metadata(self, metadatas):
        """
        Packs the specified list of metadata values and updates the ``metadata``
        and ``metadata_offset`` columns. The length of the metadatas array
        must be equal to the number of rows in the table.

        :param list metadatas: A list of metadata bytes values.
        """
        packed, offset = util.pack_bytes(metadatas)
        d = self.asdict()
        d["metadata"] = packed
        d["metadata_offset"] = offset
        self.set_columns(**d)

    @property
    def metadata_schema(self) -> metadata.MetadataSchema:
        """
        The :class:`tskit.MetadataSchema` for this table.
        """
        # This isn't as inefficient as it looks because we're using an LRU cache on
        # the parse_metadata_schema function. Thus, we're really only incurring the
        # cost of creating the unicode string from the low-level schema and looking
        # up the functools cache.
        return metadata.parse_metadata_schema(self.ll_table.metadata_schema)

    @metadata_schema.setter
    def metadata_schema(self, schema: metadata.MetadataSchema) -> None:
        if not isinstance(schema, metadata.MetadataSchema):
            raise TypeError(
                "Only instances of tskit.MetadataSchema can be assigned to "
                f"metadata_schema, not {type(schema)}"
            )
        self.ll_table.metadata_schema = repr(schema)

    def metadata_vector(self, key, *, dtype=None, default_value=NOTSET):
        """
        Returns a numpy array of metadata values obtained by extracting ``key``
        from each metadata entry, and using ``default_value`` if the key is
        not present. ``key`` may be a list, in which case nested values are returned.
        For instance, ``key = ["a", "x"]`` will return an array of
        ``row.metadata["a"]["x"]`` values, iterated over rows in this table.

        :param str key: The name, or a list of names, of metadata entries.
        :param str dtype: The dtype of the result (can usually be omitted).
        :param object default_value: The value to be inserted if the metadata key
            is not present. Note that for numeric columns, a default value of None
            will result in a non-numeric array. The default behaviour is to raise
            ``KeyError`` on missing entries.
        """

        if default_value == NOTSET:

            def getter(d, k):
                return d[k]

        else:

            def getter(d, k):
                return (
                    d.get(k, default_value) if isinstance(d, Mapping) else default_value
                )

        if isinstance(key, list):
            out = np.array(
                [
                    reduce(
                        getter,
                        key,
                        row.metadata,
                    )
                    for row in self
                ],
                dtype=dtype,
            )
        else:
            out = np.array(
                [getter(row.metadata, key) for row in self],
                dtype=dtype,
            )
        return out

    def drop_metadata(self, *, keep_schema=False):
        """
        Drops all metadata in this table. By default, the schema is also cleared,
        except if ``keep_schema`` is True.

        :param bool keep_schema: True if the current schema should be kept intact.
        """
        data = self.asdict()
        data["metadata"] = []
        data["metadata_offset"][:] = 0
        self.set_columns(**data)
        if not keep_schema:
            self.metadata_schema = metadata.MetadataSchema.null()


class IndividualTable(MetadataTable):
    """
    A table defining the individuals in a tree sequence. Note that although
    each Individual has associated nodes, reference to these is not stored in
    the individual table, but rather reference to the individual is stored for
    each node in the :class:`NodeTable`.  This is similar to the way in which
    the relationship between sites and mutations is modelled.

    .. include:: substitutions/table_edit_warning.rst

    :ivar flags: The array of flags values.
    :vartype flags: numpy.ndarray, dtype=np.uint32
    :ivar location: The flattened array of floating point location values. See
        :ref:`sec_encoding_ragged_columns` for more details.
    :vartype location: numpy.ndarray, dtype=np.float64
    :ivar location_offset: The array of offsets into the location column. See
        :ref:`sec_encoding_ragged_columns` for more details.
    :vartype location_offset: numpy.ndarray, dtype=np.uint32
    :ivar parents: The flattened array of parent individual ids. See
        :ref:`sec_encoding_ragged_columns` for more details.
    :vartype parents: numpy.ndarray, dtype=np.int32
    :ivar parents_offset: The array of offsets into the parents column. See
        :ref:`sec_encoding_ragged_columns` for more details.
    :vartype parents_offset: numpy.ndarray, dtype=np.uint32
    :ivar metadata: The flattened array of binary metadata values. See
        :ref:`sec_tables_api_binary_columns` for more details.
    :vartype metadata: numpy.ndarray, dtype=np.int8
    :ivar metadata_offset: The array of offsets into the metadata column. See
        :ref:`sec_tables_api_binary_columns` for more details.
    :vartype metadata_offset: numpy.ndarray, dtype=np.uint32
    :ivar metadata_schema: The metadata schema for this table's metadata column
    :vartype metadata_schema: tskit.MetadataSchema
    """

    column_names = [
        "flags",
        "location",
        "location_offset",
        "parents",
        "parents_offset",
        "metadata",
        "metadata_offset",
    ]

    def __init__(self, max_rows_increment=0, ll_table=None):
        if ll_table is None:
            ll_table = _tskit.IndividualTable(max_rows_increment=max_rows_increment)
        super().__init__(ll_table, IndividualTableRow)

    def _text_header_and_rows(self, limit=None):
        headers = ("id", "flags", "location", "parents", "metadata")
        rows = []
        row_indexes = util.truncate_rows(self.num_rows, limit)
        for j in row_indexes:
            if j == -1:
                rows.append(f"__skipped__{self.num_rows - limit}")
            else:
                row = self[j]
                location_str = ", ".join(map(str, row.location))
                parents_str = ", ".join([f"{p:,}" for p in row.parents])
                rows.append(
                    "{:,}\t{}\t{}\t{}\t{}".format(
                        j,
                        row.flags,
                        location_str,
                        parents_str,
                        util.render_metadata(row.metadata),
                    ).split("\t")
                )
        return headers, rows

    def add_row(self, flags=0, location=None, parents=None, metadata=None):
        """
        Adds a new row to this :class:`IndividualTable` and returns the ID of the
        corresponding individual. Metadata, if specified, will be validated and encoded
        according to the table's
        :attr:`metadata_schema<tskit.IndividualTable.metadata_schema>`.

        :param int flags: The bitwise flags for the new node.
        :param array-like location: A list of numeric values or one-dimensional numpy
            array describing the location of this individual. If not specified
            or None, a zero-dimensional location is stored.
        :param array-like parents: A list or array of ids of parent individuals. If not
            specified an empty array is stored.
        :param object metadata: Any object that is valid metadata for the table's schema.
            Defaults to the default metadata value for the table's schema. This is
            typically ``{}``. For no schema, ``None``.
        :return: The ID of the newly added individual.
        :rtype: int
        """
        if metadata is None:
            metadata = self.metadata_schema.empty_value
        metadata = self.metadata_schema.validate_and_encode_row(metadata)
        return self.ll_table.add_row(
            flags=flags, location=location, parents=parents, metadata=metadata
        )

    def set_columns(
        self,
        flags=None,
        location=None,
        location_offset=None,
        parents=None,
        parents_offset=None,
        metadata=None,
        metadata_offset=None,
        metadata_schema=None,
    ):
        """
        Sets the values for each column in this :class:`IndividualTable` using the
        values in the specified arrays. Overwrites existing data in all the table
        columns.

        The ``flags`` array is mandatory and defines the number of individuals
        the table will contain.
        The ``location`` and ``location_offset`` parameters must be supplied
        together, and meet the requirements for :ref:`sec_encoding_ragged_columns`.
        The ``parents`` and ``parents_offset`` parameters must be supplied
        together, and meet the requirements for :ref:`sec_encoding_ragged_columns`.
        The ``metadata`` and ``metadata_offset`` parameters must be supplied
        together, and meet the requirements for :ref:`sec_encoding_ragged_columns`.
        See :ref:`sec_tables_api_binary_columns` for more information and
        :ref:`sec_tutorial_metadata_bulk` for an example of how to prepare metadata.

        :param flags: The bitwise flags for each individual. Required.
        :type flags: numpy.ndarray, dtype=np.uint32
        :param location: The flattened location array. Must be specified along
            with ``location_offset``. If not specified or None, an empty location
            value is stored for each individual.
        :type location: numpy.ndarray, dtype=np.float64
        :param location_offset: The offsets into the ``location`` array.
        :type location_offset: numpy.ndarray, dtype=np.uint32.
        :param parents: The flattened parents array. Must be specified along
            with ``parents_offset``. If not specified or None, an empty parents array
            is stored for each individual.
        :type parents: numpy.ndarray, dtype=np.int32
        :param parents_offset: The offsets into the ``parents`` array.
        :type parents_offset: numpy.ndarray, dtype=np.uint32.
        :param metadata: The flattened metadata array. Must be specified along
            with ``metadata_offset``. If not specified or None, an empty metadata
            value is stored for each individual.
        :type metadata: numpy.ndarray, dtype=np.int8
        :param metadata_offset: The offsets into the ``metadata`` array.
        :type metadata_offset: numpy.ndarray, dtype=np.uint32.
        :param metadata_schema: The encoded metadata schema. If None (default)
            do not overwrite the exising schema. Note that a schema will need to be
            encoded as a string, e.g. via ``repr(new_metadata_schema)``.
        :type metadata_schema: str

        """
        self._check_required_args(flags=flags)
        self.ll_table.set_columns(
            dict(
                flags=flags,
                location=location,
                location_offset=location_offset,
                parents=parents,
                parents_offset=parents_offset,
                metadata=metadata,
                metadata_offset=metadata_offset,
                metadata_schema=metadata_schema,
            )
        )

    def append_columns(
        self,
        flags=None,
        location=None,
        location_offset=None,
        parents=None,
        parents_offset=None,
        metadata=None,
        metadata_offset=None,
    ):
        """
        Appends the specified arrays to the end of the columns in this
        :class:`IndividualTable`. This allows many new rows to be added at once.

        The ``flags`` array is mandatory and defines the number of
        extra individuals to add to the table.
        The ``parents`` and ``parents_offset`` parameters must be supplied
        together, and meet the requirements for :ref:`sec_encoding_ragged_columns`.
        The ``location`` and ``location_offset`` parameters must be supplied
        together, and meet the requirements for :ref:`sec_encoding_ragged_columns`.
        The ``metadata`` and ``metadata_offset`` parameters must be supplied
        together, and meet the requirements for :ref:`sec_encoding_ragged_columns`.
        See :ref:`sec_tables_api_binary_columns` for more information and
        :ref:`sec_tutorial_metadata_bulk` for an example of how to prepare metadata.

        :param flags: The bitwise flags for each individual. Required.
        :type flags: numpy.ndarray, dtype=np.uint32
        :param location: The flattened location array. Must be specified along
            with ``location_offset``. If not specified or None, an empty location
            value is stored for each individual.
        :type location: numpy.ndarray, dtype=np.float64
        :param location_offset: The offsets into the ``location`` array.
        :type location_offset: numpy.ndarray, dtype=np.uint32.
        :param metadata: The flattened metadata array. Must be specified along
            with ``metadata_offset``. If not specified or None, an empty metadata
            value is stored for each individual.
        :param parents: The flattened parents array. Must be specified along
            with ``parents_offset``. If not specified or None, an empty parents array
            is stored for each individual.
        :type parents: numpy.ndarray, dtype=np.int32
        :param parents_offset: The offsets into the ``parents`` array.
        :type metadata: numpy.ndarray, dtype=np.int8
        :param metadata_offset: The offsets into the ``metadata`` array.
        :type metadata_offset: numpy.ndarray, dtype=np.uint32.
        """
        self._check_required_args(flags=flags)
        self.ll_table.append_columns(
            dict(
                flags=flags,
                location=location,
                location_offset=location_offset,
                parents=parents,
                parents_offset=parents_offset,
                metadata=metadata,
                metadata_offset=metadata_offset,
            )
        )

    def packset_location(self, locations):
        """
        Packs the specified list of location values and updates the ``location``
        and ``location_offset`` columns. The length of the locations array
        must be equal to the number of rows in the table.

        :param list locations: A list of locations interpreted as numpy float64
            arrays.
        """
        packed, offset = util.pack_arrays(locations)
        d = self.asdict()
        d["location"] = packed
        d["location_offset"] = offset
        self.set_columns(**d)

    def packset_parents(self, parents):
        """
        Packs the specified list of parent values and updates the ``parent``
        and ``parent_offset`` columns. The length of the parents array
        must be equal to the number of rows in the table.

        :param list parents: A list of list of parent ids, interpreted as numpy int32
            arrays.
        """
        packed, offset = util.pack_arrays(parents, np.int32)
        d = self.asdict()
        d["parents"] = packed
        d["parents_offset"] = offset
        self.set_columns(**d)

    def keep_rows(self, keep):
        """
        .. include:: substitutions/table_keep_rows_main.rst

        The values in the ``parents`` column are updated according to this
        map, so that reference integrity within the table is maintained.
        As a consequence of this, the values in the ``parents`` column
        for kept rows are bounds-checked and an error raised if they
        are not valid. Rows that are deleted are not checked for
        parent ID integrity.

        If an attempt is made to delete rows that are referred to by
        the ``parents`` column of rows that are retained, an error
        is raised.

        These error conditions are checked before any alterations to
        the table are made.

        :param array-like keep: The rows to keep as a boolean array. Must
            be the same length as the table, and convertible to a numpy
            array of dtype bool.
        :return: The mapping between old and new row IDs as a numpy
            array (dtype int32).
        :rtype: numpy.ndarray (dtype=np.int32)
        """
        return super().keep_rows(keep)


class NodeTable(MetadataTable):
    """
    A table defining the nodes in a tree sequence. See the
    :ref:`definitions <sec_node_table_definition>` for details on the columns
    in this table and the
    :ref:`tree sequence requirements <sec_valid_tree_sequence_requirements>` section
    for the properties needed for a node table to be a part of a valid tree sequence.

    .. include:: substitutions/table_edit_warning.rst

    :ivar time: The array of time values.
    :vartype time: numpy.ndarray, dtype=np.float64
    :ivar flags: The array of flags values.
    :vartype flags: numpy.ndarray, dtype=np.uint32
    :ivar population: The array of population IDs.
    :vartype population: numpy.ndarray, dtype=np.int32
    :ivar individual: The array of individual IDs that each node belongs to.
    :vartype individual: numpy.ndarray, dtype=np.int32
    :ivar metadata: The flattened array of binary metadata values. See
        :ref:`sec_tables_api_binary_columns` for more details.
    :vartype metadata: numpy.ndarray, dtype=np.int8
    :ivar metadata_offset: The array of offsets into the metadata column. See
        :ref:`sec_tables_api_binary_columns` for more details.
    :vartype metadata_offset: numpy.ndarray, dtype=np.uint32
    :ivar metadata_schema: The metadata schema for this table's metadata column
    :vartype metadata_schema: tskit.MetadataSchema
    """

    column_names = [
        "time",
        "flags",
        "population",
        "individual",
        "metadata",
        "metadata_offset",
    ]

    def __init__(self, max_rows_increment=0, ll_table=None):
        if ll_table is None:
            ll_table = _tskit.NodeTable(max_rows_increment=max_rows_increment)
        super().__init__(ll_table, NodeTableRow)

    def _text_header_and_rows(self, limit=None):
        headers = ("id", "flags", "population", "individual", "time", "metadata")
        rows = []
        row_indexes = util.truncate_rows(self.num_rows, limit)
        decimal_places_times = 0 if self._columns_all_integer("time") else 8
        for j in row_indexes:
            row = self[j]
            if j == -1:
                rows.append(f"__skipped__{self.num_rows - limit}")
            else:
                rows.append(
                    "{:,}\t{}\t{:,}\t{:,}\t{:,.{dp}f}\t{}".format(
                        j,
                        row.flags,
                        row.population,
                        row.individual,
                        row.time,
                        util.render_metadata(row.metadata),
                        dp=decimal_places_times,
                    ).split("\t")
                )
        return headers, rows

    def add_row(self, flags=0, time=0, population=-1, individual=-1, metadata=None):
        """
        Adds a new row to this :class:`NodeTable` and returns the ID of the
        corresponding node. Metadata, if specified, will be validated and encoded
        according to the table's
        :attr:`metadata_schema<tskit.NodeTable.metadata_schema>`.

        :param int flags: The bitwise flags for the new node.
        :param float time: The birth time for the new node.
        :param int population: The ID of the population in which the new node was born.
            Defaults to :data:`tskit.NULL`.
        :param int individual: The ID of the individual in which the new node was born.
            Defaults to :data:`tskit.NULL`.
        :param object metadata: Any object that is valid metadata for the table's schema.
            Defaults to the default metadata value for the table's schema. This is
            typically ``{}``. For no schema, ``None``.
        :return: The ID of the newly added node.
        :rtype: int
        """
        if metadata is None:
            metadata = self.metadata_schema.empty_value
        metadata = self.metadata_schema.validate_and_encode_row(metadata)
        return self.ll_table.add_row(flags, time, population, individual, metadata)

    def set_columns(
        self,
        flags=None,
        time=None,
        population=None,
        individual=None,
        metadata=None,
        metadata_offset=None,
        metadata_schema=None,
    ):
        """
        Sets the values for each column in this :class:`NodeTable` using the values in
        the specified arrays. Overwrites existing data in all the table columns.

        The ``flags``, ``time`` and ``population`` arrays must all be of the same length,
        which is equal to the number of nodes the table will contain. The
        ``metadata`` and ``metadata_offset`` parameters must be supplied together, and
        meet the requirements for :ref:`sec_encoding_ragged_columns`.
        See :ref:`sec_tables_api_binary_columns` for more information and
        :ref:`sec_tutorial_metadata_bulk` for an example of how to prepare metadata.

        :param flags: The bitwise flags for each node. Required.
        :type flags: numpy.ndarray, dtype=np.uint32
        :param time: The time values for each node. Required.
        :type time: numpy.ndarray, dtype=np.float64
        :param population: The population values for each node. If not specified
            or None, the :data:`tskit.NULL` value is stored for each node.
        :type population: numpy.ndarray, dtype=np.int32
        :param individual: The individual values for each node. If not specified
            or None, the :data:`tskit.NULL` value is stored for each node.
        :type individual: numpy.ndarray, dtype=np.int32
        :param metadata: The flattened metadata array. Must be specified along
            with ``metadata_offset``. If not specified or None, an empty metadata
            value is stored for each node.
        :type metadata: numpy.ndarray, dtype=np.int8
        :param metadata_offset: The offsets into the ``metadata`` array.
        :type metadata_offset: numpy.ndarray, dtype=np.uint32.
        :param metadata_schema: The encoded metadata schema. If None (default)
            do not overwrite the exising schema. Note that a schema will need to be
            encoded as a string, e.g. via ``repr(new_metadata_schema)``.
        :type metadata_schema: str
        """
        self._check_required_args(flags=flags, time=time)
        self.ll_table.set_columns(
            dict(
                flags=flags,
                time=time,
                population=population,
                individual=individual,
                metadata=metadata,
                metadata_offset=metadata_offset,
                metadata_schema=metadata_schema,
            )
        )

    def append_columns(
        self,
        flags=None,
        time=None,
        population=None,
        individual=None,
        metadata=None,
        metadata_offset=None,
    ):
        """
        Appends the specified arrays to the end of the columns in this
        :class:`NodeTable`. This allows many new rows to be added at once.

        The ``flags``, ``time`` and ``population`` arrays must all be of the same length,
        which is equal to the number of nodes that will be added to the table. The
        ``metadata`` and ``metadata_offset`` parameters must be supplied together, and
        meet the requirements for :ref:`sec_encoding_ragged_columns`.
        See :ref:`sec_tables_api_binary_columns` for more information and
        :ref:`sec_tutorial_metadata_bulk` for an example of how to prepare metadata.

        :param flags: The bitwise flags for each node. Required.
        :type flags: numpy.ndarray, dtype=np.uint32
        :param time: The time values for each node. Required.
        :type time: numpy.ndarray, dtype=np.float64
        :param population: The population values for each node. If not specified
            or None, the :data:`tskit.NULL` value is stored for each node.
        :type population: numpy.ndarray, dtype=np.int32
        :param individual: The individual values for each node. If not specified
            or None, the :data:`tskit.NULL` value is stored for each node.
        :type individual: numpy.ndarray, dtype=np.int32
        :param metadata: The flattened metadata array. Must be specified along
            with ``metadata_offset``. If not specified or None, an empty metadata
            value is stored for each node.
        :type metadata: numpy.ndarray, dtype=np.int8
        :param metadata_offset: The offsets into the ``metadata`` array.
        :type metadata_offset: numpy.ndarray, dtype=np.uint32.
        """
        self._check_required_args(flags=flags, time=time)
        self.ll_table.append_columns(
            dict(
                flags=flags,
                time=time,
                population=population,
                individual=individual,
                metadata=metadata,
                metadata_offset=metadata_offset,
            )
        )


class EdgeTable(MetadataTable):
    """
    A table defining the edges in a tree sequence. See the
    :ref:`definitions <sec_edge_table_definition>` for details on the columns
    in this table and the
    :ref:`tree sequence requirements <sec_valid_tree_sequence_requirements>` section
    for the properties needed for an edge table to be a part of a valid tree sequence.

    .. include:: substitutions/table_edit_warning.rst

    :ivar left: The array of left coordinates.
    :vartype left: numpy.ndarray, dtype=np.float64
    :ivar right: The array of right coordinates.
    :vartype right: numpy.ndarray, dtype=np.float64
    :ivar parent: The array of parent node IDs.
    :vartype parent: numpy.ndarray, dtype=np.int32
    :ivar child: The array of child node IDs.
    :vartype child: numpy.ndarray, dtype=np.int32
    :ivar metadata: The flattened array of binary metadata values. See
        :ref:`sec_tables_api_binary_columns` for more details.
    :vartype metadata: numpy.ndarray, dtype=np.int8
    :ivar metadata_offset: The array of offsets into the metadata column. See
        :ref:`sec_tables_api_binary_columns` for more details.
    :vartype metadata_offset: numpy.ndarray, dtype=np.uint32
    :ivar metadata_schema: The metadata schema for this table's metadata column
    :vartype metadata_schema: tskit.MetadataSchema
    """

    column_names = [
        "left",
        "right",
        "parent",
        "child",
        "metadata",
        "metadata_offset",
    ]

    def __init__(self, max_rows_increment=0, ll_table=None):
        if ll_table is None:
            ll_table = _tskit.EdgeTable(max_rows_increment=max_rows_increment)
        super().__init__(ll_table, EdgeTableRow)

    def _text_header_and_rows(self, limit=None):
        headers = ("id", "left", "right", "parent", "child", "metadata")
        rows = []
        row_indexes = util.truncate_rows(self.num_rows, limit)
        decimal_places = 0 if self._columns_all_integer("left", "right") else 8
        for j in row_indexes:
            if j == -1:
                rows.append(f"__skipped__{self.num_rows - limit}")
            else:
                row = self[j]
                rows.append(
                    "{:,}\t{:,.{dp}f}\t{:,.{dp}f}\t{:,}\t{:,}\t{}".format(
                        j,
                        row.left,
                        row.right,
                        row.parent,
                        row.child,
                        util.render_metadata(row.metadata),
                        dp=decimal_places,
                    ).split("\t")
                )
        return headers, rows

    def add_row(self, left, right, parent, child, metadata=None):
        """
        Adds a new row to this :class:`EdgeTable` and returns the ID of the
        corresponding edge. Metadata, if specified, will be validated and encoded
        according to the table's
        :attr:`metadata_schema<tskit.EdgeTable.metadata_schema>`.

        :param float left: The left coordinate (inclusive).
        :param float right: The right coordinate (exclusive).
        :param int parent: The ID of parent node.
        :param int child: The ID of child node.
        :param object metadata: Any object that is valid metadata for the table's schema.
            Defaults to the default metadata value for the table's schema. This is
            typically ``{}``. For no schema, ``None``.
        :return: The ID of the newly added edge.
        :rtype: int
        """
        if metadata is None:
            metadata = self.metadata_schema.empty_value
        metadata = self.metadata_schema.validate_and_encode_row(metadata)
        return self.ll_table.add_row(left, right, parent, child, metadata)

    def set_columns(
        self,
        left=None,
        right=None,
        parent=None,
        child=None,
        metadata=None,
        metadata_offset=None,
        metadata_schema=None,
    ):
        """
        Sets the values for each column in this :class:`EdgeTable` using the values
        in the specified arrays. Overwrites existing data in all the table columns.

        The ``left``, ``right``, ``parent`` and ``child`` parameters are mandatory,
        and must be numpy arrays of the same length (which is equal to the number of
        edges the table will contain).
        The ``metadata`` and ``metadata_offset`` parameters must be supplied together,
        and meet the requirements for :ref:`sec_encoding_ragged_columns`.
        See :ref:`sec_tables_api_binary_columns` for more information and
        :ref:`sec_tutorial_metadata_bulk` for an example of how to prepare metadata.


        :param left: The left coordinates (inclusive).
        :type left: numpy.ndarray, dtype=np.float64
        :param right: The right coordinates (exclusive).
        :type right: numpy.ndarray, dtype=np.float64
        :param parent: The parent node IDs.
        :type parent: numpy.ndarray, dtype=np.int32
        :param child: The child node IDs.
        :type child: numpy.ndarray, dtype=np.int32
        :param metadata: The flattened metadata array. Must be specified along
            with ``metadata_offset``. If not specified or None, an empty metadata
            value is stored for each node.
        :type metadata: numpy.ndarray, dtype=np.int8
        :param metadata_offset: The offsets into the ``metadata`` array.
        :type metadata_offset: numpy.ndarray, dtype=np.uint32.
        :param metadata_schema: The encoded metadata schema. If None (default)
            do not overwrite the exising schema. Note that a schema will need to be
            encoded as a string, e.g. via ``repr(new_metadata_schema)``.
        :type metadata_schema: str
        """
        self._check_required_args(left=left, right=right, parent=parent, child=child)
        self.ll_table.set_columns(
            dict(
                left=left,
                right=right,
                parent=parent,
                child=child,
                metadata=metadata,
                metadata_offset=metadata_offset,
                metadata_schema=metadata_schema,
            )
        )

    def append_columns(
        self, left, right, parent, child, metadata=None, metadata_offset=None
    ):
        """
        Appends the specified arrays to the end of the columns of this
        :class:`EdgeTable`. This allows many new rows to be added at once.

        The ``left``, ``right``, ``parent`` and ``child`` parameters are mandatory,
        and must be numpy arrays of the same length (which is equal to the number of
        additional edges to add to the table). The ``metadata`` and
        ``metadata_offset`` parameters must be supplied together, and
        meet the requirements for :ref:`sec_encoding_ragged_columns`.
        See :ref:`sec_tables_api_binary_columns` for more information and
        :ref:`sec_tutorial_metadata_bulk` for an example of how to prepare metadata.


        :param left: The left coordinates (inclusive).
        :type left: numpy.ndarray, dtype=np.float64
        :param right: The right coordinates (exclusive).
        :type right: numpy.ndarray, dtype=np.float64
        :param parent: The parent node IDs.
        :type parent: numpy.ndarray, dtype=np.int32
        :param child: The child node IDs.
        :type child: numpy.ndarray, dtype=np.int32
        :param metadata: The flattened metadata array. Must be specified along
            with ``metadata_offset``. If not specified or None, an empty metadata
            value is stored for each node.
        :type metadata: numpy.ndarray, dtype=np.int8
        :param metadata_offset: The offsets into the ``metadata`` array.
        :type metadata_offset: numpy.ndarray, dtype=np.uint32.
        """
        self.ll_table.append_columns(
            dict(
                left=left,
                right=right,
                parent=parent,
                child=child,
                metadata=metadata,
                metadata_offset=metadata_offset,
            )
        )

    def squash(self):
        """
        Sorts, then condenses the table into the smallest possible number of rows by
        combining any adjacent edges.
        A pair of edges is said to be `adjacent` if they have the same parent and child
        nodes, and if the left coordinate of one of the edges is equal to the right
        coordinate of the other edge.
        The ``squash`` method modifies an :class:`EdgeTable` in place so that any set of
        adjacent edges is replaced by a single edge.
        The new edge will have the same parent and child node, a left coordinate
        equal to the smallest left coordinate in the set, and a right coordinate
        equal to the largest right coordinate in the set.
        The new edge table will be sorted in the order (P, C, L, R): if the node table
        is ordered by increasing node time, as is common, this order will meet the
        :ref:`sec_edge_requirements` for a valid tree sequence, otherwise you will need
        to call :meth:`.sort` on the entire :class:`TableCollection`.

        .. note::
            Note that this method will fail if any edges have non-empty metadata.

        """
        self.ll_table.squash()


class MigrationTable(MetadataTable):
    """
    A table defining the migrations in a tree sequence. See the
    :ref:`definitions <sec_migration_table_definition>` for details on the columns
    in this table and the
    :ref:`tree sequence requirements <sec_valid_tree_sequence_requirements>` section
    for the properties needed for a migration table to be a part of a valid tree
    sequence.

    .. include:: substitutions/table_edit_warning.rst

    :ivar left: The array of left coordinates.
    :vartype left: numpy.ndarray, dtype=np.float64
    :ivar right: The array of right coordinates.
    :vartype right: numpy.ndarray, dtype=np.float64
    :ivar node: The array of node IDs.
    :vartype node: numpy.ndarray, dtype=np.int32
    :ivar source: The array of source population IDs.
    :vartype source: numpy.ndarray, dtype=np.int32
    :ivar dest: The array of destination population IDs.
    :vartype dest: numpy.ndarray, dtype=np.int32
    :ivar time: The array of time values.
    :vartype time: numpy.ndarray, dtype=np.float64
    :ivar metadata: The flattened array of binary metadata values. See
        :ref:`sec_tables_api_binary_columns` for more details.
    :vartype metadata: numpy.ndarray, dtype=np.int8
    :ivar metadata_offset: The array of offsets into the metadata column. See
        :ref:`sec_tables_api_binary_columns` for more details.
    :vartype metadata_offset: numpy.ndarray, dtype=np.uint32
    :ivar metadata_schema: The metadata schema for this table's metadata column
    :vartype metadata_schema: tskit.MetadataSchema
    """

    column_names = [
        "left",
        "right",
        "node",
        "source",
        "dest",
        "time",
        "metadata",
        "metadata_offset",
    ]

    def __init__(self, max_rows_increment=0, ll_table=None):
        if ll_table is None:
            ll_table = _tskit.MigrationTable(max_rows_increment=max_rows_increment)
        super().__init__(ll_table, MigrationTableRow)

    def _text_header_and_rows(self, limit=None):
        headers = ("id", "left", "right", "node", "source", "dest", "time", "metadata")
        rows = []
        row_indexes = util.truncate_rows(self.num_rows, limit)
        decimal_places_coords = 0 if self._columns_all_integer("left", "right") else 8
        decimal_places_times = 0 if self._columns_all_integer("time") else 8
        for j in row_indexes:
            if j == -1:
                rows.append(f"__skipped__{self.num_rows - limit}")
            else:
                row = self[j]
                rows.append(
                    "{:,}\t{:,.{dp_c}f}\t{:,.{dp_c}f}\t{:,}\t{:,}\t{:,}\t{:,.{dp_t}f}\t{}".format(
                        j,
                        row.left,
                        row.right,
                        row.node,
                        row.source,
                        row.dest,
                        row.time,
                        util.render_metadata(row.metadata),
                        dp_c=decimal_places_coords,
                        dp_t=decimal_places_times,
                    ).split(
                        "\t"
                    )
                )
        return headers, rows

    def add_row(self, left, right, node, source, dest, time, metadata=None):
        """
        Adds a new row to this :class:`MigrationTable` and returns the ID of the
        corresponding migration. Metadata, if specified, will be validated and encoded
        according to the table's
        :attr:`metadata_schema<tskit.MigrationTable.metadata_schema>`.

        :param float left: The left coordinate (inclusive).
        :param float right: The right coordinate (exclusive).
        :param int node: The node ID.
        :param int source: The ID of the source population.
        :param int dest: The ID of the destination population.
        :param float time: The time of the migration event.
        :param object metadata: Any object that is valid metadata for the table's schema.
            Defaults to the default metadata value for the table's schema. This is
            typically ``{}``. For no schema, ``None``.
        :return: The ID of the newly added migration.
        :rtype: int
        """
        if metadata is None:
            metadata = self.metadata_schema.empty_value
        metadata = self.metadata_schema.validate_and_encode_row(metadata)
        return self.ll_table.add_row(left, right, node, source, dest, time, metadata)

    def set_columns(
        self,
        left=None,
        right=None,
        node=None,
        source=None,
        dest=None,
        time=None,
        metadata=None,
        metadata_offset=None,
        metadata_schema=None,
    ):
        """
        Sets the values for each column in this :class:`MigrationTable` using the values
        in the specified arrays. Overwrites existing data in all the table columns.

        All parameters except ``metadata`` and ``metadata_offset`` and are mandatory,
        and must be numpy arrays of the same length (which is equal to the number of
        migrations the table will contain).
        The ``metadata`` and ``metadata_offset`` parameters must be supplied together,
        and meet the requirements for :ref:`sec_encoding_ragged_columns`.
        See :ref:`sec_tables_api_binary_columns` for more information and
        :ref:`sec_tutorial_metadata_bulk` for an example of how to prepare metadata.

        :param left: The left coordinates (inclusive).
        :type left: numpy.ndarray, dtype=np.float64
        :param right: The right coordinates (exclusive).
        :type right: numpy.ndarray, dtype=np.float64
        :param node: The node IDs.
        :type node: numpy.ndarray, dtype=np.int32
        :param source: The source population IDs.
        :type source: numpy.ndarray, dtype=np.int32
        :param dest: The destination population IDs.
        :type dest: numpy.ndarray, dtype=np.int32
        :param time: The time of each migration.
        :type time: numpy.ndarray, dtype=np.int64
        :param metadata: The flattened metadata array. Must be specified along
            with ``metadata_offset``. If not specified or None, an empty metadata
            value is stored for each migration.
        :type metadata: numpy.ndarray, dtype=np.int8
        :param metadata_offset: The offsets into the ``metadata`` array.
        :type metadata_offset: numpy.ndarray, dtype=np.uint32.
        :param metadata_schema: The encoded metadata schema. If None (default)
            do not overwrite the exising schema. Note that a schema will need to be
            encoded as a string, e.g. via ``repr(new_metadata_schema)``.
        :type metadata_schema: str
        """
        self._check_required_args(
            left=left, right=right, node=node, source=source, dest=dest, time=time
        )
        self.ll_table.set_columns(
            dict(
                left=left,
                right=right,
                node=node,
                source=source,
                dest=dest,
                time=time,
                metadata=metadata,
                metadata_offset=metadata_offset,
                metadata_schema=metadata_schema,
            )
        )

    def append_columns(
        self,
        left,
        right,
        node,
        source,
        dest,
        time,
        metadata=None,
        metadata_offset=None,
    ):
        """
        Appends the specified arrays to the end of the columns of this
        :class:`MigrationTable`. This allows many new rows to be added at once.

        All parameters except ``metadata`` and ``metadata_offset`` and are mandatory,
        and must be numpy arrays of the same length (which is equal to the number of
        additional migrations to add to the table). The ``metadata`` and
        ``metadata_offset`` parameters must be supplied together, and
        meet the requirements for :ref:`sec_encoding_ragged_columns`.
        See :ref:`sec_tables_api_binary_columns` for more information and
        :ref:`sec_tutorial_metadata_bulk` for an example of how to prepare metadata.

        :param left: The left coordinates (inclusive).
        :type left: numpy.ndarray, dtype=np.float64
        :param right: The right coordinates (exclusive).
        :type right: numpy.ndarray, dtype=np.float64
        :param node: The node IDs.
        :type node: numpy.ndarray, dtype=np.int32
        :param source: The source population IDs.
        :type source: numpy.ndarray, dtype=np.int32
        :param dest: The destination population IDs.
        :type dest: numpy.ndarray, dtype=np.int32
        :param time: The time of each migration.
        :type time: numpy.ndarray, dtype=np.int64
        :param metadata: The flattened metadata array. Must be specified along
            with ``metadata_offset``. If not specified or None, an empty metadata
            value is stored for each migration.
        :type metadata: numpy.ndarray, dtype=np.int8
        :param metadata_offset: The offsets into the ``metadata`` array.
        :type metadata_offset: numpy.ndarray, dtype=np.uint32.
        """
        self.ll_table.append_columns(
            dict(
                left=left,
                right=right,
                node=node,
                source=source,
                dest=dest,
                time=time,
                metadata=metadata,
                metadata_offset=metadata_offset,
            )
        )


class SiteTable(MetadataTable):
    """
    A table defining the sites in a tree sequence. See the
    :ref:`definitions <sec_site_table_definition>` for details on the columns
    in this table and the
    :ref:`tree sequence requirements <sec_valid_tree_sequence_requirements>` section
    for the properties needed for a site table to be a part of a valid tree
    sequence.

    .. include:: substitutions/table_edit_warning.rst

    :ivar position: The array of site position coordinates.
    :vartype position: numpy.ndarray, dtype=np.float64
    :ivar ancestral_state: The flattened array of ancestral state strings.
        See :ref:`sec_tables_api_text_columns` for more details.
    :vartype ancestral_state: numpy.ndarray, dtype=np.int8
    :ivar ancestral_state_offset: The offsets of rows in the ancestral_state
        array. See :ref:`sec_tables_api_text_columns` for more details.
    :vartype ancestral_state_offset: numpy.ndarray, dtype=np.uint32
    :ivar metadata: The flattened array of binary metadata values. See
        :ref:`sec_tables_api_binary_columns` for more details.
    :vartype metadata: numpy.ndarray, dtype=np.int8
    :ivar metadata_offset: The array of offsets into the metadata column. See
        :ref:`sec_tables_api_binary_columns` for more details.
    :vartype metadata_offset: numpy.ndarray, dtype=np.uint32
    :ivar metadata_schema: The metadata schema for this table's metadata column
    :vartype metadata_schema: tskit.MetadataSchema
    """

    column_names = [
        "position",
        "ancestral_state",
        "ancestral_state_offset",
        "metadata",
        "metadata_offset",
    ]

    def __init__(self, max_rows_increment=0, ll_table=None):
        if ll_table is None:
            ll_table = _tskit.SiteTable(max_rows_increment=max_rows_increment)
        super().__init__(ll_table, SiteTableRow)

    def _text_header_and_rows(self, limit=None):
        headers = ("id", "position", "ancestral_state", "metadata")
        rows = []
        row_indexes = util.truncate_rows(self.num_rows, limit)
        decimal_places = 0 if self._columns_all_integer("position") else 8
        for j in row_indexes:
            if j == -1:
                rows.append(f"__skipped__{self.num_rows - limit}")
            else:
                row = self[j]
                rows.append(
                    "{:,}\t{:,.{dp}f}\t{}\t{}".format(
                        j,
                        row.position,
                        row.ancestral_state,
                        util.render_metadata(row.metadata),
                        dp=decimal_places,
                    ).split("\t")
                )
        return headers, rows

    def add_row(self, position, ancestral_state, metadata=None):
        """
        Adds a new row to this :class:`SiteTable` and returns the ID of the
        corresponding site. Metadata, if specified, will be validated and encoded
        according to the table's
        :attr:`metadata_schema<tskit.SiteTable.metadata_schema>`.

        :param float position: The position of this site in genome coordinates.
        :param str ancestral_state: The state of this site at the root of the tree.
        :param object metadata: Any object that is valid metadata for the table's schema.
            Defaults to the default metadata value for the table's schema. This is
            typically ``{}``. For no schema, ``None``.
        :return: The ID of the newly added site.
        :rtype: int
        """
        if metadata is None:
            metadata = self.metadata_schema.empty_value
        metadata = self.metadata_schema.validate_and_encode_row(metadata)
        return self.ll_table.add_row(position, ancestral_state, metadata)

    def set_columns(
        self,
        position=None,
        ancestral_state=None,
        ancestral_state_offset=None,
        metadata=None,
        metadata_offset=None,
        metadata_schema=None,
    ):
        """
        Sets the values for each column in this :class:`SiteTable` using the values
        in the specified arrays. Overwrites existing data in all the table columns.

        The ``position``, ``ancestral_state`` and ``ancestral_state_offset``
        parameters are mandatory, and must be 1D numpy arrays. The length
        of the ``position`` array determines the number of rows in table.
        The ``ancestral_state`` and ``ancestral_state_offset`` parameters must
        be supplied together, and meet the requirements for
        :ref:`sec_encoding_ragged_columns` (see
        :ref:`sec_tables_api_text_columns` for more information). The
        ``metadata`` and ``metadata_offset`` parameters must be supplied
        together, and meet the requirements for
        :ref:`sec_encoding_ragged_columns` (see
        :ref:`sec_tables_api_binary_columns` for more information) and
        :ref:`sec_tutorial_metadata_bulk` for an example of how to prepare metadata.

        :param position: The position of each site in genome coordinates.
        :type position: numpy.ndarray, dtype=np.float64
        :param ancestral_state: The flattened ancestral_state array. Required.
        :type ancestral_state: numpy.ndarray, dtype=np.int8
        :param ancestral_state_offset: The offsets into the ``ancestral_state`` array.
        :type ancestral_state_offset: numpy.ndarray, dtype=np.uint32.
        :param metadata: The flattened metadata array. Must be specified along
            with ``metadata_offset``. If not specified or None, an empty metadata
            value is stored for each node.
        :type metadata: numpy.ndarray, dtype=np.int8
        :param metadata_offset: The offsets into the ``metadata`` array.
        :type metadata_offset: numpy.ndarray, dtype=np.uint32.
        :param metadata_schema: The encoded metadata schema. If None (default)
            do not overwrite the exising schema. Note that a schema will need to be
            encoded as a string, e.g. via ``repr(new_metadata_schema)``.
        :type metadata_schema: str
        """
        self._check_required_args(
            position=position,
            ancestral_state=ancestral_state,
            ancestral_state_offset=ancestral_state_offset,
        )
        self.ll_table.set_columns(
            dict(
                position=position,
                ancestral_state=ancestral_state,
                ancestral_state_offset=ancestral_state_offset,
                metadata=metadata,
                metadata_offset=metadata_offset,
                metadata_schema=metadata_schema,
            )
        )

    def append_columns(
        self,
        position,
        ancestral_state,
        ancestral_state_offset,
        metadata=None,
        metadata_offset=None,
    ):
        """
        Appends the specified arrays to the end of the columns of this
        :class:`SiteTable`. This allows many new rows to be added at once.

        The ``position``, ``ancestral_state`` and ``ancestral_state_offset``
        parameters are mandatory, and must be 1D numpy arrays. The length
        of the ``position`` array determines the number of additional rows
        to add the table.
        The ``ancestral_state`` and ``ancestral_state_offset`` parameters must
        be supplied together, and meet the requirements for
        :ref:`sec_encoding_ragged_columns` (see
        :ref:`sec_tables_api_text_columns` for more information). The
        ``metadata`` and ``metadata_offset`` parameters must be supplied
        together, and meet the requirements for
        :ref:`sec_encoding_ragged_columns` (see
        :ref:`sec_tables_api_binary_columns` for more information) and
        :ref:`sec_tutorial_metadata_bulk` for an example of how to prepare metadata.

        :param position: The position of each site in genome coordinates.
        :type position: numpy.ndarray, dtype=np.float64
        :param ancestral_state: The flattened ancestral_state array. Required.
        :type ancestral_state: numpy.ndarray, dtype=np.int8
        :param ancestral_state_offset: The offsets into the ``ancestral_state`` array.
        :type ancestral_state_offset: numpy.ndarray, dtype=np.uint32.
        :param metadata: The flattened metadata array. Must be specified along
            with ``metadata_offset``. If not specified or None, an empty metadata
            value is stored for each node.
        :type metadata: numpy.ndarray, dtype=np.int8
        :param metadata_offset: The offsets into the ``metadata`` array.
        :type metadata_offset: numpy.ndarray, dtype=np.uint32.
        """
        self.ll_table.append_columns(
            dict(
                position=position,
                ancestral_state=ancestral_state,
                ancestral_state_offset=ancestral_state_offset,
                metadata=metadata,
                metadata_offset=metadata_offset,
            )
        )

    def packset_ancestral_state(self, ancestral_states):
        """
        Packs the specified list of ancestral_state values and updates the
        ``ancestral_state`` and ``ancestral_state_offset`` columns. The length
        of the ancestral_states array must be equal to the number of rows in
        the table.

        :param list(str) ancestral_states: A list of string ancestral state values.
        """
        packed, offset = util.pack_strings(ancestral_states)
        d = self.asdict()
        d["ancestral_state"] = packed
        d["ancestral_state_offset"] = offset
        self.set_columns(**d)


class MutationTable(MetadataTable):
    """
    A table defining the mutations in a tree sequence. See the
    :ref:`definitions <sec_mutation_table_definition>` for details on the columns
    in this table and the
    :ref:`tree sequence requirements <sec_valid_tree_sequence_requirements>` section
    for the properties needed for a mutation table to be a part of a valid tree
    sequence.

    .. include:: substitutions/table_edit_warning.rst

    :ivar site: The array of site IDs.
    :vartype site: numpy.ndarray, dtype=np.int32
    :ivar node: The array of node IDs.
    :vartype node: numpy.ndarray, dtype=np.int32
    :ivar time: The array of time values.
    :vartype time: numpy.ndarray, dtype=np.float64
    :ivar derived_state: The flattened array of derived state strings.
        See :ref:`sec_tables_api_text_columns` for more details.
    :vartype derived_state: numpy.ndarray, dtype=np.int8
    :ivar derived_state_offset: The offsets of rows in the derived_state
        array. See :ref:`sec_tables_api_text_columns` for more details.
    :vartype derived_state_offset: numpy.ndarray, dtype=np.uint32
    :ivar parent: The array of parent mutation IDs.
    :vartype parent: numpy.ndarray, dtype=np.int32
    :ivar metadata: The flattened array of binary metadata values. See
        :ref:`sec_tables_api_binary_columns` for more details.
    :vartype metadata: numpy.ndarray, dtype=np.int8
    :ivar metadata_offset: The array of offsets into the metadata column. See
        :ref:`sec_tables_api_binary_columns` for more details.
    :vartype metadata_offset: numpy.ndarray, dtype=np.uint32
    :ivar metadata_schema: The metadata schema for this table's metadata column
    :vartype metadata_schema: tskit.MetadataSchema
    """

    column_names = [
        "site",
        "node",
        "time",
        "derived_state",
        "derived_state_offset",
        "parent",
        "metadata",
        "metadata_offset",
    ]

    def __init__(self, max_rows_increment=0, ll_table=None):
        if ll_table is None:
            ll_table = _tskit.MutationTable(max_rows_increment=max_rows_increment)
        super().__init__(ll_table, MutationTableRow)

    def _text_header_and_rows(self, limit=None):
        headers = ("id", "site", "node", "time", "derived_state", "parent", "metadata")
        rows = []
        row_indexes = util.truncate_rows(self.num_rows, limit)
        # Currently mutations do not have discretised times: this for consistency
        decimal_places_times = 0 if self._columns_all_integer("time") else 8
        for j in row_indexes:
            if j == -1:
                rows.append(f"__skipped__{self.num_rows - limit}")
            else:
                row = self[j]
                rows.append(
                    "{:,}\t{:,}\t{:,}\t{:,.{dp}f}\t{}\t{:,}\t{}".format(
                        j,
                        row.site,
                        row.node,
                        row.time,
                        row.derived_state,
                        row.parent,
                        util.render_metadata(row.metadata),
                        dp=decimal_places_times,
                    ).split("\t")
                )
        return headers, rows

    def add_row(self, site, node, derived_state, parent=-1, metadata=None, time=None):
        """
        Adds a new row to this :class:`MutationTable` and returns the ID of the
        corresponding mutation. Metadata, if specified, will be validated and encoded
        according to the table's
        :attr:`metadata_schema<tskit.MutationTable.metadata_schema>`.

        :param int site: The ID of the site that this mutation occurs at.
        :param int node: The ID of the first node inheriting this mutation.
        :param str derived_state: The state of the site at this mutation's node.
        :param int parent: The ID of the parent mutation. If not specified,
            defaults to :attr:`NULL`.
        :param object metadata: Any object that is valid metadata for the table's schema.
            Defaults to the default metadata value for the table's schema. This is
            typically ``{}``. For no schema, ``None``.
        :return: The ID of the newly added mutation.
        :param float time: The occurrence time for the new mutation. If not specified,
            defaults to ``UNKNOWN_TIME``, indicating the time is unknown.
        :rtype: int
        """
        if metadata is None:
            metadata = self.metadata_schema.empty_value
        metadata = self.metadata_schema.validate_and_encode_row(metadata)
        return self.ll_table.add_row(
            site,
            node,
            derived_state,
            parent,
            metadata,
            UNKNOWN_TIME if time is None else time,
        )

    def set_columns(
        self,
        site=None,
        node=None,
        time=None,
        derived_state=None,
        derived_state_offset=None,
        parent=None,
        metadata=None,
        metadata_offset=None,
        metadata_schema=None,
    ):
        """
        Sets the values for each column in this :class:`MutationTable` using the values
        in the specified arrays. Overwrites existing data in all the the table columns.

        The ``site``, ``node``, ``derived_state`` and ``derived_state_offset``
        parameters are mandatory, and must be 1D numpy arrays. The
        ``site`` and ``node`` (also ``parent`` and ``time``, if supplied) arrays
        must be of equal length, and determine the number of rows in the table.
        The ``derived_state`` and ``derived_state_offset`` parameters must
        be supplied together, and meet the requirements for
        :ref:`sec_encoding_ragged_columns` (see
        :ref:`sec_tables_api_text_columns` for more information). The
        ``metadata`` and ``metadata_offset`` parameters must be supplied
        together, and meet the requirements for
        :ref:`sec_encoding_ragged_columns` (see
        :ref:`sec_tables_api_binary_columns` for more information) and
        :ref:`sec_tutorial_metadata_bulk` for an example of how to prepare metadata.

        :param site: The ID of the site each mutation occurs at.
        :type site: numpy.ndarray, dtype=np.int32
        :param node: The ID of the node each mutation is associated with.
        :type node: numpy.ndarray, dtype=np.int32
        :param time: The time values for each mutation.
        :type time: numpy.ndarray, dtype=np.float64
        :param derived_state: The flattened derived_state array. Required.
        :type derived_state: numpy.ndarray, dtype=np.int8
        :param derived_state_offset: The offsets into the ``derived_state`` array.
        :type derived_state_offset: numpy.ndarray, dtype=np.uint32.
        :param parent: The ID of the parent mutation for each mutation.
        :type parent: numpy.ndarray, dtype=np.int32
        :param metadata: The flattened metadata array. Must be specified along
            with ``metadata_offset``. If not specified or None, an empty metadata
            value is stored for each node.
        :type metadata: numpy.ndarray, dtype=np.int8
        :param metadata_offset: The offsets into the ``metadata`` array.
        :type metadata_offset: numpy.ndarray, dtype=np.uint32.
        :param metadata_schema: The encoded metadata schema. If None (default)
            do not overwrite the exising schema. Note that a schema will need to be
            encoded as a string, e.g. via ``repr(new_metadata_schema)``.
        :type metadata_schema: str
        """
        self._check_required_args(
            site=site,
            node=node,
            derived_state=derived_state,
            derived_state_offset=derived_state_offset,
        )
        self.ll_table.set_columns(
            dict(
                site=site,
                node=node,
                parent=parent,
                time=time,
                derived_state=derived_state,
                derived_state_offset=derived_state_offset,
                metadata=metadata,
                metadata_offset=metadata_offset,
                metadata_schema=metadata_schema,
            )
        )

    def append_columns(
        self,
        site,
        node,
        derived_state,
        derived_state_offset,
        parent=None,
        time=None,
        metadata=None,
        metadata_offset=None,
    ):
        """
        Appends the specified arrays to the end of the columns of this
        :class:`MutationTable`. This allows many new rows to be added at once.

        The ``site``, ``node``, ``derived_state`` and ``derived_state_offset``
        parameters are mandatory, and must be 1D numpy arrays. The
        ``site`` and ``node`` (also ``time`` and ``parent``, if supplied) arrays
        must be of equal length, and determine the number of additional
        rows to add to the table.
        The ``derived_state`` and ``derived_state_offset`` parameters must
        be supplied together, and meet the requirements for
        :ref:`sec_encoding_ragged_columns` (see
        :ref:`sec_tables_api_text_columns` for more information). The
        ``metadata`` and ``metadata_offset`` parameters must be supplied
        together, and meet the requirements for
        :ref:`sec_encoding_ragged_columns` (see
        :ref:`sec_tables_api_binary_columns` for more information) and
        :ref:`sec_tutorial_metadata_bulk` for an example of how to prepare metadata.

        :param site: The ID of the site each mutation occurs at.
        :type site: numpy.ndarray, dtype=np.int32
        :param node: The ID of the node each mutation is associated with.
        :type node: numpy.ndarray, dtype=np.int32
        :param time: The time values for each mutation.
        :type time: numpy.ndarray, dtype=np.float64
        :param derived_state: The flattened derived_state array. Required.
        :type derived_state: numpy.ndarray, dtype=np.int8
        :param derived_state_offset: The offsets into the ``derived_state`` array.
        :type derived_state_offset: numpy.ndarray, dtype=np.uint32.
        :param parent: The ID of the parent mutation for each mutation.
        :type parent: numpy.ndarray, dtype=np.int32
        :param metadata: The flattened metadata array. Must be specified along
            with ``metadata_offset``. If not specified or None, an empty metadata
            value is stored for each node.
        :type metadata: numpy.ndarray, dtype=np.int8
        :param metadata_offset: The offsets into the ``metadata`` array.
        :type metadata_offset: numpy.ndarray, dtype=np.uint32.
        """
        self.ll_table.append_columns(
            dict(
                site=site,
                node=node,
                time=time,
                parent=parent,
                derived_state=derived_state,
                derived_state_offset=derived_state_offset,
                metadata=metadata,
                metadata_offset=metadata_offset,
            )
        )

    def packset_derived_state(self, derived_states):
        """
        Packs the specified list of derived_state values and updates the
        ``derived_state`` and ``derived_state_offset`` columns. The length
        of the derived_states array must be equal to the number of rows in
        the table.

        :param list(str) derived_states: A list of string derived state values.
        """
        packed, offset = util.pack_strings(derived_states)
        d = self.asdict()
        d["derived_state"] = packed
        d["derived_state_offset"] = offset
        self.set_columns(**d)

    def keep_rows(self, keep):
        """
        .. include:: substitutions/table_keep_rows_main.rst

        The values in the ``parent`` column are updated according to this
        map, so that reference integrity within the table is maintained.
        As a consequence of this, the values in the ``parent`` column
        for kept rows are bounds-checked and an error raised if they
        are not valid. Rows that are deleted are not checked for
        parent ID integrity.

        If an attempt is made to delete rows that are referred to by
        the ``parent`` column of rows that are retained, an error
        is raised.

        These error conditions are checked before any alterations to
        the table are made.

        :param array-like keep: The rows to keep as a boolean array. Must
            be the same length as the table, and convertible to a numpy
            array of dtype bool.
        :return: The mapping between old and new row IDs as a numpy
            array (dtype int32).
        :rtype: numpy.ndarray (dtype=np.int32)
        """
        return super().keep_rows(keep)


class PopulationTable(MetadataTable):
    """
    A table defining the populations referred to in a tree sequence.
    The PopulationTable stores metadata for populations that may be referred to
    in the NodeTable and MigrationTable".  Note that although nodes
    may be associated with populations, this association is stored in
    the :class:`NodeTable`: only metadata on each population is stored
    in the population table.

    .. include:: substitutions/table_edit_warning.rst

    :ivar metadata: The flattened array of binary metadata values. See
        :ref:`sec_tables_api_binary_columns` for more details.
    :vartype metadata: numpy.ndarray, dtype=np.int8
    :ivar metadata_offset: The array of offsets into the metadata column. See
        :ref:`sec_tables_api_binary_columns` for more details.
    :vartype metadata_offset: numpy.ndarray, dtype=np.uint32
    :ivar metadata_schema: The metadata schema for this table's metadata column
    :vartype metadata_schema: tskit.MetadataSchema
    """

    column_names = ["metadata", "metadata_offset"]

    def __init__(self, max_rows_increment=0, ll_table=None):
        if ll_table is None:
            ll_table = _tskit.PopulationTable(max_rows_increment=max_rows_increment)
        super().__init__(ll_table, PopulationTableRow)

    def add_row(self, metadata=None):
        """
        Adds a new row to this :class:`PopulationTable` and returns the ID of the
        corresponding population. Metadata, if specified, will be validated and encoded
        according to the table's
        :attr:`metadata_schema<tskit.PopulationTable.metadata_schema>`.

        :param object metadata: Any object that is valid metadata for the table's schema.
            Defaults to the default metadata value for the table's schema. This is
            typically ``{}``. For no schema, ``None``.
        :return: The ID of the newly added population.
        :rtype: int
        """
        if metadata is None:
            metadata = self.metadata_schema.empty_value
        metadata = self.metadata_schema.validate_and_encode_row(metadata)
        return self.ll_table.add_row(metadata=metadata)

    def _text_header_and_rows(self, limit=None):
        headers = ("id", "metadata")
        rows = []
        row_indexes = util.truncate_rows(self.num_rows, limit)
        for j in row_indexes:
            if j == -1:
                rows.append(f"__skipped__{self.num_rows - limit}")
            else:
                rows.append((str(j), util.render_metadata(self[j].metadata, length=70)))
        return headers, rows

    def set_columns(self, metadata=None, metadata_offset=None, metadata_schema=None):
        """
        Sets the values for each column in this :class:`PopulationTable` using the
        values in the specified arrays. Overwrites existing data in all the table
        columns.

        The ``metadata`` and ``metadata_offset`` parameters must be supplied
        together, and meet the requirements for
        :ref:`sec_encoding_ragged_columns` (see
        :ref:`sec_tables_api_binary_columns` for more information) and
        :ref:`sec_tutorial_metadata_bulk` for an example of how to prepare metadata.

        :param metadata: The flattened metadata array. Must be specified along
            with ``metadata_offset``. If not specified or None, an empty metadata
            value is stored for each node.
        :type metadata: numpy.ndarray, dtype=np.int8
        :param metadata_offset: The offsets into the ``metadata`` array.
        :type metadata_offset: numpy.ndarray, dtype=np.uint32.
        :param metadata_schema: The encoded metadata schema. If None (default)
            do not overwrite the exising schema. Note that a schema will need to be
            encoded as a string, e.g. via ``repr(new_metadata_schema)``.
        :type metadata_schema: str
        """
        self.ll_table.set_columns(
            dict(
                metadata=metadata,
                metadata_offset=metadata_offset,
                metadata_schema=metadata_schema,
            )
        )

    def append_columns(self, metadata=None, metadata_offset=None):
        """
        Appends the specified arrays to the end of the columns of this
        :class:`PopulationTable`. This allows many new rows to be added at once.

        The ``metadata`` and ``metadata_offset`` parameters must be supplied
        together, and meet the requirements for
        :ref:`sec_encoding_ragged_columns` (see
        :ref:`sec_tables_api_binary_columns` for more information) and
        :ref:`sec_tutorial_metadata_bulk` for an example of how to prepare metadata.

        :param metadata: The flattened metadata array. Must be specified along
            with ``metadata_offset``. If not specified or None, an empty metadata
            value is stored for each node.
        :type metadata: numpy.ndarray, dtype=np.int8
        :param metadata_offset: The offsets into the ``metadata`` array.
        :type metadata_offset: numpy.ndarray, dtype=np.uint32.
        """
        self.ll_table.append_columns(
            dict(metadata=metadata, metadata_offset=metadata_offset)
        )


class ProvenanceTable(BaseTable):
    """
    A table recording the provenance (i.e., history) of this table, so that the
    origin of the underlying data and sequence of subsequent operations can be
    traced. Each row contains a "record" string (recommended format: JSON) and
    a timestamp.

    .. todo::
        The format of the `record` field will be more precisely specified in
        the future.

    :ivar record: The flattened array containing the record strings.
        :ref:`sec_tables_api_text_columns` for more details.
    :vartype record: numpy.ndarray, dtype=np.int8
    :ivar record_offset: The array of offsets into the record column. See
        :ref:`sec_tables_api_text_columns` for more details.
    :vartype record_offset: numpy.ndarray, dtype=np.uint32
    :ivar timestamp: The flattened array containing the timestamp strings.
        :ref:`sec_tables_api_text_columns` for more details.
    :vartype timestamp: numpy.ndarray, dtype=np.int8
    :ivar timestamp_offset: The array of offsets into the timestamp column. See
        :ref:`sec_tables_api_text_columns` for more details.
    :vartype timestamp_offset: numpy.ndarray, dtype=np.uint32
    """

    column_names = ["record", "record_offset", "timestamp", "timestamp_offset"]

    def __init__(self, max_rows_increment=0, ll_table=None):
        if ll_table is None:
            ll_table = _tskit.ProvenanceTable(max_rows_increment=max_rows_increment)
        super().__init__(ll_table, ProvenanceTableRow)

    def equals(self, other, ignore_timestamps=False):
        """
        Returns True if  `self` and `other` are equal. By default, two provenance
        tables are considered equal if their columns are byte-for-byte identical.

        :param other: Another provenance table instance
        :param bool ignore_timestamps: If True exclude the timestamp column
            from the comparison.
        :return: True if other is equal to this provenance table; False otherwise.
        :rtype: bool
        """
        ret = False
        if type(other) is type(self):
            ret = bool(
                self.ll_table.equals(
                    other.ll_table, ignore_timestamps=ignore_timestamps
                )
            )
        return ret

    def assert_equals(self, other, *, ignore_timestamps=False):
        """
        Raise an AssertionError for the first found difference between
        this and another provenance table.

        :param other: Another provenance table instance
        :param bool ignore_timestamps: If True exclude the timestamp column
            from the comparison.
        """
        if type(other) is not type(self):
            raise AssertionError(f"Types differ: self={type(self)} other={type(other)}")

        # Check using the low-level method to avoid slowly going through everything
        if self.equals(other, ignore_timestamps=ignore_timestamps):
            return

        for n, (row_self, row_other) in enumerate(zip(self, other)):
            if ignore_timestamps:
                row_self = dataclasses.replace(row_self, timestamp=None)
                row_other = dataclasses.replace(row_other, timestamp=None)
            if row_self != row_other:
                self_dict = dataclasses.asdict(self[n])
                other_dict = dataclasses.asdict(other[n])
                diff_string = []
                for col in self_dict.keys():
                    if self_dict[col] != other_dict[col]:
                        diff_string.append(
                            f"self.{col}={self_dict[col]} other.{col}={other_dict[col]}"
                        )
                diff_string = "\n".join(diff_string)
                raise AssertionError(
                    f"{type(self).__name__} row {n} differs:\n{diff_string}"
                )

        if self.num_rows != other.num_rows:
            raise AssertionError(
                f"{type(self).__name__} number of rows differ: self={self.num_rows} "
                f"other={other.num_rows}"
            )

        raise AssertionError(
            "Tables differ in an undetected way - "
            "this is a bug, please report an issue on gitub"
        )  # pragma: no cover

    def add_row(self, record, timestamp=None):
        """
        Adds a new row to this ProvenanceTable consisting of the specified record and
        timestamp. If timestamp is not specified, it is automatically generated from
        the current time.

        :param str record: A provenance record, describing the parameters and
            environment used to generate the current set of tables.
        :param str timestamp: A string timestamp. This should be in ISO8601 form.
        """
        if timestamp is None:
            timestamp = datetime.datetime.now().isoformat()
        # Note that the order of the positional arguments has been reversed
        # from the low-level module, which is a bit confusing. However, we
        # want the default behaviour here to be to add a row to the table at
        # the current time as simply as possible.
        return self.ll_table.add_row(record=record, timestamp=timestamp)

    def set_columns(
        self, timestamp=None, timestamp_offset=None, record=None, record_offset=None
    ):
        """
        Sets the values for each column in this :class:`ProvenanceTable` using the
        values in the specified arrays. Overwrites existing data in all the table
        columns.

        The ``timestamp`` and ``timestamp_offset`` parameters must be supplied
        together, and meet the requirements for
        :ref:`sec_encoding_ragged_columns` (see
        :ref:`sec_tables_api_binary_columns` for more information). Likewise
        for the ``record`` and ``record_offset`` columns

        :param timestamp: The flattened timestamp array. Must be specified along
            with ``timestamp_offset``. If not specified or None, an empty timestamp
            value is stored for each node.
        :type timestamp: numpy.ndarray, dtype=np.int8
        :param timestamp_offset: The offsets into the ``timestamp`` array.
        :type timestamp_offset: numpy.ndarray, dtype=np.uint32.
        :param record: The flattened record array. Must be specified along
            with ``record_offset``. If not specified or None, an empty record
            value is stored for each node.
        :type record: numpy.ndarray, dtype=np.int8
        :param record_offset: The offsets into the ``record`` array.
        :type record_offset: numpy.ndarray, dtype=np.uint32.
        """
        self.ll_table.set_columns(
            dict(
                timestamp=timestamp,
                timestamp_offset=timestamp_offset,
                record=record,
                record_offset=record_offset,
            )
        )

    def append_columns(
        self, timestamp=None, timestamp_offset=None, record=None, record_offset=None
    ):
        """
        Appends the specified arrays to the end of the columns of this
        :class:`ProvenanceTable`. This allows many new rows to be added at once.

        The ``timestamp`` and ``timestamp_offset`` parameters must be supplied
        together, and meet the requirements for
        :ref:`sec_encoding_ragged_columns` (see
        :ref:`sec_tables_api_binary_columns` for more information). Likewise
        for the ``record`` and ``record_offset`` columns

        :param timestamp: The flattened timestamp array. Must be specified along
            with ``timestamp_offset``. If not specified or None, an empty timestamp
            value is stored for each node.
        :type timestamp: numpy.ndarray, dtype=np.int8
        :param timestamp_offset: The offsets into the ``timestamp`` array.
        :type timestamp_offset: numpy.ndarray, dtype=np.uint32.
        :param record: The flattened record array. Must be specified along
            with ``record_offset``. If not specified or None, an empty record
            value is stored for each node.
        :type record: numpy.ndarray, dtype=np.int8
        :param record_offset: The offsets into the ``record`` array.
        :type record_offset: numpy.ndarray, dtype=np.uint32.
        """
        self.ll_table.append_columns(
            dict(
                timestamp=timestamp,
                timestamp_offset=timestamp_offset,
                record=record,
                record_offset=record_offset,
            )
        )

    def _text_header_and_rows(self, limit=None):
        headers = ("id", "timestamp", "record")
        rows = []
        row_indexes = util.truncate_rows(self.num_rows, limit)
        for j in row_indexes:
            if j == -1:
                rows.append(f"__skipped__{self.num_rows - limit}")
            else:
                row = self[j]
                rows.append(
                    (
                        str(j),
                        str(row.timestamp),
                        util.truncate_string_end(str(row.record), length=60),
                    )
                )
        return headers, rows

    def packset_record(self, records):
        """
        Packs the specified list of record values and updates the
        ``record`` and ``record_offset`` columns. The length
        of the records array must be equal to the number of rows in
        the table.

        :param list(str) records: A list of string record values.
        """
        packed, offset = util.pack_strings(records)
        d = self.asdict()
        d["record"] = packed
        d["record_offset"] = offset
        self.set_columns(**d)

    def packset_timestamp(self, timestamps):
        """
        Packs the specified list of timestamp values and updates the
        ``timestamp`` and ``timestamp_offset`` columns. The length
        of the timestamps array must be equal to the number of rows in
        the table.

        :param list(str) timestamps: A list of string timestamp values.
        """
        packed, offset = util.pack_strings(timestamps)
        d = self.asdict()
        d["timestamp"] = packed
        d["timestamp_offset"] = offset
        self.set_columns(**d)


# We define segment ordering by (left, right, node) tuples
@dataclasses.dataclass(eq=True, order=True)
class IdentitySegment:
    """
    A single segment of identity spanning a genomic interval for a
    a specific ancestor node.
    """

    left: float
    """The left genomic coordinate (inclusive)."""
    right: float
    """The right genomic coordinate (exclusive)."""
    node: int
    """The ID of the most recent common ancestor node."""

    @property
    def span(self) -> float:
        """
        The length of the genomic region spanned by this identity segment.
        """
        return self.right - self.left


class IdentitySegmentList(collections.abc.Iterable, collections.abc.Sized):
    """
    A summary of identity segments for some pair of samples in a
    :class:`.IdentitySegments` result. If the ``store_segments`` argument
    has been specified to :meth:`.TreeSequence.ibd_segments`, this class
    can be treated as a sequence of :class:`.IdentitySegment` objects.

    Access to the segment data via numpy arrays is also available via
    the :attr:`.IdentitySegmentList.left`, :attr:`.IdentitySegmentList.right`
    and :attr:`.IdentitySegmentList.node` attributes.

    If ``store_segments`` is False, only the overall summary values
    such as :attr:`.IdentitySegmentList.total_span` and ``len()`` are
    available.

    .. warning:: The order of segments within an IdentitySegmentList is
        arbitrary and may change in the future

    """

    def __init__(self, ll_segment_list):
        self._ll_segment_list = ll_segment_list

    def __iter__(self):
        for left, right, node in zip(self.left, self.right, self.node):
            yield IdentitySegment(float(left), float(right), int(node))

    def __len__(self):
        return self._ll_segment_list.num_segments

    def __str__(self):
        return (
            f"IdentitySegmentList(num_segments={len(self)}, "
            f"total_span={self.total_span})"
        )

    def __repr__(self):
        return f"IdentitySegmentList({repr(list(self))})"

    def __eq__(self, other):
        if not isinstance(other, IdentitySegmentList):
            return False
        return list(self) == list(other)

    @property
    def total_span(self):
        """
        The total genomic span covered by segments in this list. Equal to
        ``sum(seg.span for seg in seglst)``.
        """
        return self._ll_segment_list.total_span

    @property
    def left(self):
        """
        A numpy array (dtype=np.float64) of the ``left`` coordinates of segments.
        """
        return self._ll_segment_list.left

    @property
    def right(self):
        """
        A numpy array (dtype=np.float64) of the ``right`` coordinates of segments.
        """
        return self._ll_segment_list.right

    @property
    def node(self):
        """
        A numpy array (dtype=np.int32) of the MRCA node IDs in segments.
        """
        return self._ll_segment_list.node


class IdentitySegments(collections.abc.Mapping):
    """
    A class summarising and optionally storing the segments of identity
    by state returned by :meth:`.TreeSequence.ibd_segments`. See the
    :ref:`sec_identity` for more information and examples.

    Along with the documented methods and attributes, the class supports
    the Python mapping protocol, and can be regarded as a dictionary
    mapping sample node pair tuples to the corresponding
    :class:`.IdentitySegmentList`.

    .. note:: It is important to note that the facilities available
       for a given instance of this class are determined by the
       ``store_pairs`` and ``store_segments`` arguments provided to
       :meth:`.TreeSequence.ibd_segments`. For example, attempting
       to access per-sample pair information if ``store_pairs``
       is False will result in a (hopefully informative) error being
       raised.

    .. warning:: This class should not be instantiated directly.
    """

    def __init__(self, ll_result, *, max_time, min_span, store_segments, store_pairs):
        self._ll_identity_segments = ll_result
        self.max_time = max_time
        self.min_span = min_span
        self.store_segments = store_segments
        self.store_pairs = store_pairs

    @property
    def num_segments(self):
        """
        The total number of identity segments found.
        """
        return self._ll_identity_segments.num_segments

    @property
    def num_pairs(self):
        """
        The total number of distinct sample pairs for which identity
        segments were found. (Only available when ``store_pairs`` or
        ``store_segments`` is specified).
        """
        return self._ll_identity_segments.num_pairs

    @property
    def total_span(self):
        """
        The total genomic sequence length spanned by all identity
        segments that were found.
        """
        return self._ll_identity_segments.total_span

    @property
    def pairs(self):
        """
        A numpy array with shape ``(segs.num_pairs, 2)`` and dtype=np.int32
        containing the sample pairs for which IBD segments were found.
        """
        return self._ll_identity_segments.get_keys()

    # We have two different versions of repr - one where we list out the segments
    # for debugging, and the other that just shows the standard representation.
    # We could have repr fail if store_segments isn't true, but then printing,
    # e.g., a list of IdentitySegments objects would fail unexpectedly.
    def __repr__(self):
        if self.store_segments:
            return f"IdentitySegments({dict(self)})"
        return super().__repr__()

    def __str__(self):
        # TODO it would be nice to add horizontal lines as
        # table separators to distinguish the two parts of the
        # table like suggested here:
        # https://github.com/tskit-dev/tskit/pull/1902#issuecomment-989943424
        rows = [
            ["Parameters:", ""],
            ["max_time", str(self.max_time)],
            ["min_span", str(self.min_span)],
            ["store_pairs", str(self.store_pairs)],
            ["store_segments", str(self.store_segments)],
            ["Results:", ""],
            ["num_segments", str(self.num_segments)],
            ["total_span", str(self.total_span)],
        ]
        if self.store_pairs:
            rows.append(["num_pairs", str(len(self))])
        return util.unicode_table(rows, title="IdentitySegments", row_separator=False)

    def __getitem__(self, key):
        sample_a, sample_b = key
        return IdentitySegmentList(self._ll_identity_segments.get(sample_a, sample_b))

    def __iter__(self):
        return map(tuple, self._ll_identity_segments.get_keys())

    def __len__(self):
        return self.num_pairs


# TODO move to reference_sequence.py when we start adding more functionality.
class ReferenceSequence(metadata.MetadataProvider):
    """
    The :ref:`reference sequence<sec_data_model_reference_sequence>` associated
    with a given :class:`.TableCollection` or :class:`.TreeSequence`.

    Metadata concerning reference sequences can be described using the
    :attr:`.metadata_schema` and stored in the :attr:`.metadata` attribute.
    See the :ref:`examples<sec_metadata_examples_reference_sequence>` for
    idiomatic usage.

    .. warning:: This API is preliminary and currently only supports accessing
       reference sequence information via the ``.data`` attribute. Future versions
       will also enable transparent fetching of known reference sequences
       from a URL (see https://github.com/tskit-dev/tskit/issues/2022).
    """

    def __init__(self, ll_reference_sequence):
        super().__init__(ll_reference_sequence)
        self._ll_reference_sequence = ll_reference_sequence

    def is_null(self) -> bool:
        """
        Returns True if this :class:`.ReferenceSequence` is null, i.e.,
        all fields are empty.
        """
        return bool(self._ll_reference_sequence.is_null())

    def clear(self):
        self.data = ""
        self.url = ""
        self.metadata_schema = tskit.MetadataSchema(None)
        self.metadata = b""

    # https://github.com/tskit-dev/tskit/issues/1984
    # TODO add a __str__ method
    # TODO add a _repr_html_
    # FIXME This is a shortcut, we want to put the values in explicitly
    # here to get more control over how they are displayed.
    def __repr__(self):
        return f"ReferenceSequence({repr(self.asdict())})"

    @property
    def data(self) -> str:
        """
        The string encoding of the reference sequence such that ``data[j]``
        represents the reference nucleotide at base ``j``. If this reference
        sequence is writable, the value can be assigned, e.g.
        ``tables.reference_sequence.data = "ACGT"``
        """
        return self._ll_reference_sequence.data

    @data.setter
    def data(self, value):
        self._ll_reference_sequence.data = value

    @property
    def url(self) -> str:
        return self._ll_reference_sequence.url

    @url.setter
    def url(self, value):
        self._ll_reference_sequence.url = value

    def asdict(self) -> dict:
        return {
            "metadata_schema": repr(self.metadata_schema),
            "metadata": self.metadata_bytes,
            "data": self.data,
            "url": self.url,
        }

    def __eq__(self, other):
        return self.equals(other)

    def equals(self, other, ignore_metadata=False):
        try:
            self.assert_equals(other, ignore_metadata)
            return True
        except AssertionError:
            return False

    def assert_equals(self, other, ignore_metadata=False):
        if not ignore_metadata:
            super().assert_equals(other)

        if self.data != other.data:
            raise AssertionError(
                f"Reference sequence data differs: self={self.data} "
                f"other={other.data}"
            )
        if self.url != other.url:
            raise AssertionError(
                f"Reference sequence url differs: self={self.url} " f"other={other.url}"
            )

    @property
    def nbytes(self):
        # TODO this will be inefficient when we work with large references.
        # Make a dedicated low-level method for getting the length of data.
        return super().nbytes + len(self.url) + len(self.data)


class TableCollection(metadata.MetadataProvider):
    """
    A collection of mutable tables defining a tree sequence. See the
    :ref:`sec_data_model` section for definition on the various tables
    and how they together define a :class:`TreeSequence`. Arbitrary
    data can be stored in a TableCollection, but there are certain
    :ref:`requirements <sec_valid_tree_sequence_requirements>` that must be
    satisfied for these tables to be interpreted as a tree sequence.

    To obtain an immutable :class:`TreeSequence` instance corresponding to the
    current state of a ``TableCollection``, please use the :meth:`.tree_sequence`
    method.
    """

    set_err_text = (
        "Cannot set tables in a table collection: use table.replace_with() instead."
    )

    def __init__(self, sequence_length=0, *, ll_tables=None):
        self._ll_tables = ll_tables
        if ll_tables is None:
            self._ll_tables = _tskit.TableCollection(sequence_length)
        super().__init__(self._ll_tables)
        self._individuals = IndividualTable(ll_table=self._ll_tables.individuals)
        self._nodes = NodeTable(ll_table=self._ll_tables.nodes)
        self._edges = EdgeTable(ll_table=self._ll_tables.edges)
        self._migrations = MigrationTable(ll_table=self._ll_tables.migrations)
        self._sites = SiteTable(ll_table=self._ll_tables.sites)
        self._mutations = MutationTable(ll_table=self._ll_tables.mutations)
        self._populations = PopulationTable(ll_table=self._ll_tables.populations)
        self._provenances = ProvenanceTable(ll_table=self._ll_tables.provenances)

    @property
    def individuals(self) -> IndividualTable:
        """
        The :ref:`sec_individual_table_definition` in this collection.
        """
        return self._individuals

    @individuals.setter
    def individuals(self, value):
        raise AttributeError(self.set_err_text)

    @property
    def nodes(self) -> NodeTable:
        """
        The :ref:`sec_node_table_definition` in this collection.
        """
        return self._nodes

    @nodes.setter
    def nodes(self, value):
        raise AttributeError(self.set_err_text)

    @property
    def edges(self) -> EdgeTable:
        """
        The :ref:`sec_edge_table_definition` in this collection.
        """
        return self._edges

    @edges.setter
    def edges(self, value):
        raise AttributeError(self.set_err_text)

    @property
    def migrations(self) -> MigrationTable:
        """
        The :ref:`sec_migration_table_definition` in this collection
        """
        return self._migrations

    @migrations.setter
    def migrations(self, value):
        raise AttributeError(self.set_err_text)

    @property
    def sites(self) -> SiteTable:
        """
        The :ref:`sec_site_table_definition` in this collection.
        """
        return self._sites

    @sites.setter
    def sites(self, value):
        raise AttributeError(self.set_err_text)

    @property
    def mutations(self) -> MutationTable:
        """
        The :ref:`sec_mutation_table_definition` in this collection.
        """
        return self._mutations

    @mutations.setter
    def mutations(self, value):
        raise AttributeError(self.set_err_text)

    @property
    def populations(self) -> PopulationTable:
        """
        The :ref:`sec_population_table_definition` in this collection.
        """
        return self._populations

    @populations.setter
    def populations(self, value):
        raise AttributeError(self.set_err_text)

    @property
    def provenances(self) -> ProvenanceTable:
        """
        The :ref:`sec_provenance_table_definition` in this collection.
        """
        return self._provenances

    @provenances.setter
    def provenances(self, value):
        raise AttributeError(self.set_err_text)

    @property
    def indexes(self) -> TableCollectionIndexes:
        """
        The edge insertion and removal indexes.
        """
        indexes = self._ll_tables.indexes
        return TableCollectionIndexes(**indexes)

    @indexes.setter
    def indexes(self, indexes):
        self._ll_tables.indexes = indexes.asdict()

    @property
    def sequence_length(self) -> float:
        """
        The sequence length defining the coordinate space.
        """
        return self._ll_tables.sequence_length

    @sequence_length.setter
    def sequence_length(self, sequence_length):
        self._ll_tables.sequence_length = sequence_length

    @property
    def file_uuid(self) -> str:
        """
        The UUID for the file this TableCollection is derived
        from, or None if not derived from a file.
        """
        return self._ll_tables.file_uuid

    @property
    def time_units(self) -> str:
        """
        The units used for the time dimension of this TableCollection
        """
        return self._ll_tables.time_units

    @time_units.setter
    def time_units(self, time_units: str) -> None:
        self._ll_tables.time_units = time_units

    def has_reference_sequence(self):
        """
        Returns True if this :class:`.TableCollection` has an associated
        :ref:`reference sequence<sec_data_model_reference_sequence>`.
        """
        return bool(self._ll_tables.has_reference_sequence())

    @property
    def reference_sequence(self):
        """
        The :class:`.ReferenceSequence` associated with this :class:`.TableCollection`.

        .. note:: Note that the behaviour of this attribute differs from
            :attr:`.TreeSequence.reference_sequence` in that we return a valid
            instance of :class:`.ReferenceSequence` even when
            :attr:`.TableCollection.has_reference_sequence` is False. This is
            to allow us to update the state of the reference sequence.
        """
        # NOTE: arguably we should cache the reference to this object
        # during init, rather than creating a new instance each time.
        # However, following the pattern of the Table classes for now
        # for consistency.
        return ReferenceSequence(self._ll_tables.reference_sequence)

    @reference_sequence.setter
    def reference_sequence(self, value: ReferenceSequence):
        self.reference_sequence.metadata_schema = value.metadata_schema
        self.reference_sequence.metadata = value.metadata
        self.reference_sequence.data = value.data
        self.reference_sequence.url = value.url

    def asdict(self, force_offset_64=False):
        """
        Returns the nested dictionary representation of this TableCollection
        used for interchange.

        Note: the semantics of this method changed at tskit 0.1.0. Previously a
        map of table names to the tables themselves was returned.

        :param bool force_offset_64: If True, all offset columns will have dtype
            np.uint64. If False (the default) the offset array columns will have
            a dtype of either np.uint32 or np.uint64, depending on the size of the
            corresponding data array.
        :return: The dictionary representation of this table collection.
        :rtype: dict
        """
        return self._ll_tables.asdict(force_offset_64)

    @property
    def table_name_map(self) -> Dict:
        """
        Returns a dictionary mapping table names to the corresponding
        table instances. For example, the returned dictionary will contain the
        key "edges" that maps to an :class:`.EdgeTable` instance.
        """
        return {
            "edges": self.edges,
            "individuals": self.individuals,
            "migrations": self.migrations,
            "mutations": self.mutations,
            "nodes": self.nodes,
            "populations": self.populations,
            "provenances": self.provenances,
            "sites": self.sites,
        }

    @property
    def name_map(self) -> Dict:
        # Deprecated in 0.4.1
        warnings.warn(
            "name_map is deprecated; use table_name_map instead",
            FutureWarning,
            stacklevel=4,
        )
        return self.table_name_map

    @property
    def nbytes(self) -> int:
        """
        Returns the total number of bytes required to store the data
        in this table collection. Note that this may not be equal to
        the actual memory footprint.
        """
        return sum(
            (
                8,  # sequence_length takes 8 bytes
                super().nbytes,  # metadata
                len(self.time_units.encode()),
                self.indexes.nbytes,
                self.reference_sequence.nbytes,
                sum(table.nbytes for table in self.table_name_map.values()),
            )
        )

    def __str__(self):
        """
        Return a plain text summary of this TableCollection
        """
        return "\n".join(
            [
                "TableCollection",
                "",
                f"Sequence Length: {self.sequence_length}",
                f"Time units: {self.time_units}",
                f"Metadata: {self.metadata}",
                "",
                "Individuals",
                str(self.individuals),
                "Nodes",
                str(self.nodes),
                "Edges",
                str(self.edges),
                "Sites",
                str(self.sites),
                "Mutations",
                str(self.mutations),
                "Migrations",
                str(self.migrations),
                "Populations",
                str(self.populations),
                "Provenances",
                str(self.provenances),
            ]
        )

    def equals(
        self,
        other,
        *,
        ignore_metadata=False,
        ignore_ts_metadata=False,
        ignore_provenance=False,
        ignore_timestamps=False,
        ignore_tables=False,
        ignore_reference_sequence=False,
    ):
        """
        Returns True if  `self` and `other` are equal. By default, two table
        collections are considered equal if their

        - ``sequence_length`` properties are identical;
        - top-level tree sequence metadata and metadata schemas are
          byte-wise identical;
        - constituent tables are byte-wise identical.

        Some of the requirements in this definition can be relaxed using the
        parameters, which can be used to remove certain parts of the data model
        from the comparison.

        Table indexes are not considered in the equality comparison.

        :param TableCollection other: Another table collection.
        :param bool ignore_metadata: If True *all* metadata and metadata schemas
            will be excluded from the comparison. This includes the top-level
            tree sequence and constituent table metadata (default=False).
        :param bool ignore_ts_metadata: If True the top-level tree sequence
            metadata and metadata schemas will be excluded from the comparison.
            If ``ignore_metadata`` is True, this parameter has no effect.
        :param bool ignore_provenance: If True the provenance tables are
            not included in the comparison.
        :param bool ignore_timestamps: If True the provenance timestamp column
            is ignored in the comparison. If ``ignore_provenance`` is True, this
            parameter has no effect.
        :param bool ignore_tables: If True no tables are included in the
            comparison, thus comparing only the top-level information.
        :param bool ignore_reference_sequence: If True the reference sequence
            is not included in the comparison.
        :return: True if other is equal to this table collection; False otherwise.
        :rtype: bool
        """
        ret = False
        if type(other) is type(self):
            ret = bool(
                self._ll_tables.equals(
                    other._ll_tables,
                    ignore_metadata=bool(ignore_metadata),
                    ignore_ts_metadata=bool(ignore_ts_metadata),
                    ignore_provenance=bool(ignore_provenance),
                    ignore_timestamps=bool(ignore_timestamps),
                    ignore_tables=bool(ignore_tables),
                    ignore_reference_sequence=bool(ignore_reference_sequence),
                )
            )
        return ret

    def assert_equals(
        self,
        other,
        *,
        ignore_metadata=False,
        ignore_ts_metadata=False,
        ignore_provenance=False,
        ignore_timestamps=False,
        ignore_tables=False,
        ignore_reference_sequence=False,
    ):
        """
        Raise an AssertionError for the first found difference between
        this and another table collection. Note that table indexes are not checked.

        :param TableCollection other: Another table collection.
        :param bool ignore_metadata: If True *all* metadata and metadata schemas
            will be excluded from the comparison. This includes the top-level
            tree sequence and constituent table metadata (default=False).
        :param bool ignore_ts_metadata: If True the top-level tree sequence
            metadata and metadata schemas will be excluded from the comparison.
            If ``ignore_metadata`` is True, this parameter has no effect.
        :param bool ignore_provenance: If True the provenance tables are
            not included in the comparison.
        :param bool ignore_timestamps: If True the provenance timestamp column
            is ignored in the comparison. If ``ignore_provenance`` is True, this
            parameter has no effect.
        :param bool ignore_tables: If True no tables are included in the
            comparison, thus comparing only the top-level information.
        :param bool ignore_reference_sequence: If True the reference sequence
            is not included in the comparison.
        """
        if type(other) is not type(self):
            raise AssertionError(f"Types differ: self={type(self)} other={type(other)}")

        # Check using the low-level method to avoid slowly going through everything
        if self.equals(
            other,
            ignore_metadata=ignore_metadata,
            ignore_ts_metadata=ignore_ts_metadata,
            ignore_provenance=ignore_provenance,
            ignore_timestamps=ignore_timestamps,
            ignore_tables=ignore_tables,
            ignore_reference_sequence=ignore_reference_sequence,
        ):
            return

        if not ignore_metadata or ignore_ts_metadata:
            super().assert_equals(other)

        if not ignore_reference_sequence:
            self.reference_sequence.assert_equals(
                other.reference_sequence, ignore_metadata=ignore_metadata
            )

        if self.time_units != other.time_units:
            raise AssertionError(
                f"Time units differs: self={self.time_units} "
                f"other={other.time_units}"
            )

        if self.sequence_length != other.sequence_length:
            raise AssertionError(
                f"Sequence Length"
                f" differs: self={self.sequence_length} other={other.sequence_length}"
            )

        for table_name, table in self.table_name_map.items():
            if table_name != "provenances":
                table.assert_equals(
                    getattr(other, table_name), ignore_metadata=ignore_metadata
                )

        if not ignore_provenance:
            self.provenances.assert_equals(
                other.provenances, ignore_timestamps=ignore_timestamps
            )

        raise AssertionError(
            "TableCollections differ in an undetected way - "
            "this is a bug, please report an issue on gitub"
        )  # pragma: no cover

    def __eq__(self, other):
        return self.equals(other)

    def __getstate__(self):
        return self.asdict()

    @classmethod
    def load(cls, file_or_path, *, skip_tables=False, skip_reference_sequence=False):
        file, local_file = util.convert_file_like_to_open_file(file_or_path, "rb")
        ll_tc = _tskit.TableCollection()
        try:
            ll_tc.load(
                file,
                skip_tables=skip_tables,
                skip_reference_sequence=skip_reference_sequence,
            )
            return TableCollection(ll_tables=ll_tc)
        except tskit.FileFormatError as e:
            util.raise_known_file_format_errors(file, e)
        finally:
            if local_file:
                file.close()

    def dump(self, file_or_path):
        """
        Writes the table collection to the specified path or file object.

        :param str file_or_path: The file object or path to write the TreeSequence to.
        """
        file, local_file = util.convert_file_like_to_open_file(file_or_path, "wb")
        try:
            self._ll_tables.dump(file)
        finally:
            if local_file:
                file.close()

    # Unpickle support
    def __setstate__(self, state):
        self.__init__()
        self._ll_tables.fromdict(state)

    @classmethod
    def fromdict(self, tables_dict):
        ll_tc = _tskit.TableCollection()
        ll_tc.fromdict(tables_dict)
        return TableCollection(ll_tables=ll_tc)

    def copy(self):
        """
        Returns a deep copy of this TableCollection.

        :return: A deep copy of this TableCollection.
        :rtype: tskit.TableCollection
        """
        return TableCollection.fromdict(self.asdict())

    def tree_sequence(self):
        """
        Returns a :class:`TreeSequence` instance from the tables defined in this
        :class:`TableCollection`, building the required indexes if they have not yet
        been created by :meth:`.build_index`. If the table collection does not meet
        the :ref:`sec_valid_tree_sequence_requirements`, for example if the tables
        are not correctly sorted or if they cannot be interpreted as a tree sequence,
        an exception is raised. Note that in the former case, the :meth:`.sort`
        method may be used to ensure that sorting requirements are met.

        :return: A :class:`TreeSequence` instance reflecting the structures
            defined in this set of tables.
        :rtype: tskit.TreeSequence
        """
        if not self.has_index():
            self.build_index()
        return tskit.TreeSequence.load_tables(self)

    def simplify(
        self,
        samples=None,
        *,
        reduce_to_site_topology=False,
        filter_populations=None,
        filter_individuals=None,
        filter_sites=None,
        filter_nodes=None,
        update_sample_flags=None,
        keep_unary=False,
        keep_unary_in_individuals=None,
        keep_input_roots=False,
        record_provenance=True,
        filter_zero_mutation_sites=None,  # Deprecated alias for filter_sites
    ):
        """
        Simplifies the tables in place to retain only the information necessary
        to reconstruct the tree sequence describing the given ``samples``.
        If ``filter_nodes`` is True (the default), this can change the ID of
        the nodes, so that the node ``samples[k]`` will have ID ``k`` in the
        result, resulting in a NodeTable where only the first ``len(samples)``
        nodes are marked as samples. The mapping from node IDs in the current
        set of tables to their equivalent values in the simplified tables is
        returned as a numpy array. If an array ``a`` is returned by this
        function and ``u`` is the ID of a node in the input table, then
        ``a[u]`` is the ID of this node in the output table. For any node ``u``
        that is not mapped into the output tables, this mapping will equal
        ``tskit.NULL`` (``-1``).

        Tables operated on by this function must: be sorted (see
        :meth:`TableCollection.sort`), have children be born strictly after their
        parents, and the intervals on which any node is a child must be
        disjoint. Other than this the tables need not satisfy remaining
        requirements to specify a valid tree sequence (but the resulting tables
        will).

        .. note::
            To invert the returned ``node_map``, that is, to obtain a reverse
            mapping from the node ID in the output table to the node ID in
            the input table, you can use::

                rev_map = np.zeros_like(node_map, shape=simplified_ts.num_nodes)
                kept = node_map != tskit.NULL
                rev_map[node_map[kept]] = np.arange(len(node_map))[kept]

            In this case, no elements of the ``rev_map`` array will be set to
            ``tskit.NULL``.

        .. seealso::
            This is identical to :meth:`TreeSequence.simplify` but acts *in place* to
            alter the data in this :class:`TableCollection`. Please see the
            :meth:`TreeSequence.simplify` method for a description of the remaining
            parameters.

        :param list[int] samples: A list of node IDs to retain as samples. They
            need not be nodes marked as samples in the original tree sequence, but
            will constitute the entire set of samples in the returned tree sequence.
            If not specified or None, use all nodes marked with the IS_SAMPLE flag.
            The list may be provided as a numpy array (or array-like) object
            (dtype=np.int32).
        :param bool reduce_to_site_topology: Whether to reduce the topology down
            to the trees that are present at sites. (Default: False).
        :param bool filter_populations: If True, remove any populations that are
            not referenced by nodes after simplification; new population IDs are
            allocated sequentially from zero. If False, the population table will
            not be altered in any way. (Default: None, treated as True)
        :param bool filter_individuals: If True, remove any individuals that are
            not referenced by nodes after simplification; new individual IDs are
            allocated sequentially from zero. If False, the individual table will
            not be altered in any way. (Default: None, treated as True)
        :param bool filter_sites: If True, remove any sites that are
            not referenced by mutations after simplification; new site IDs are
            allocated sequentially from zero. If False, the site table will not
            be altered in any way. (Default: None, treated as True)
        :param bool filter_nodes: If True, remove any nodes that are
            not referenced by edges after simplification. If False, the only
            potential change to the node table may be to change the node flags
            (if ``samples`` is specified and different from the existing samples).
            (Default: None, treated as True)
        :param bool update_sample_flags: If True, update node flags to so that
            nodes in the specified list of samples have the NODE_IS_SAMPLE
            flag after simplification, and nodes that are not in this list
            do not. (Default: None, treated as True)
        :param bool keep_unary: If True, preserve unary nodes (i.e. nodes with
            exactly one child) that exist on the path from samples to root.
            (Default: False)
        :param bool keep_unary_in_individuals: If True, preserve unary nodes
            that exist on the path from samples to root, but only if they are
            associated with an individual in the individuals table. Cannot be
            specified at the same time as ``keep_unary``. (Default: ``None``,
            equivalent to False)
        :param bool keep_input_roots: Whether to retain history ancestral to the
            MRCA of the samples. If ``False``, no topology older than the MRCAs of the
            samples will be included. If ``True`` the roots of all trees in the returned
            tree sequence will be the same roots as in the original tree sequence.
            (Default: False)
        :param bool record_provenance: If True, record details of this call to
            simplify in the returned tree sequence's provenance information
            (Default: True).
        :param bool filter_zero_mutation_sites: Deprecated alias for ``filter_sites``.
        :return: A numpy array mapping node IDs in the input tables to their
            corresponding node IDs in the output tables.
        :rtype: numpy.ndarray (dtype=np.int32)
        """
        if filter_zero_mutation_sites is not None:
            # Deprecated in msprime 0.6.1.
            warnings.warn(
                "filter_zero_mutation_sites is deprecated; use filter_sites instead",
                FutureWarning,
                stacklevel=4,
            )
            filter_sites = filter_zero_mutation_sites
        if samples is None:
            flags = self.nodes.flags
            samples = np.where(np.bitwise_and(flags, _tskit.NODE_IS_SAMPLE) != 0)[
                0
            ].astype(np.int32)
        else:
            samples = util.safe_np_int_cast(samples, np.int32)
        if filter_populations is None:
            filter_populations = True
        if filter_individuals is None:
            filter_individuals = True
        if filter_sites is None:
            filter_sites = True
        if filter_nodes is None:
            filter_nodes = True
        if update_sample_flags is None:
            update_sample_flags = True
        if keep_unary_in_individuals is None:
            keep_unary_in_individuals = False

        node_map = self._ll_tables.simplify(
            samples,
            filter_sites=filter_sites,
            filter_individuals=filter_individuals,
            filter_populations=filter_populations,
            filter_nodes=filter_nodes,
            update_sample_flags=update_sample_flags,
            reduce_to_site_topology=reduce_to_site_topology,
            keep_unary=keep_unary,
            keep_unary_in_individuals=keep_unary_in_individuals,
            keep_input_roots=keep_input_roots,
        )
        if record_provenance:
            # TODO replace with a version of https://github.com/tskit-dev/tskit/pull/243
            # TODO also make sure we convert all the arguments so that they are
            # definitely JSON encodable.
            parameters = {"command": "simplify", "TODO": "add simplify parameters"}
            self.provenances.add_row(
                record=json.dumps(provenance.get_provenance_dict(parameters))
            )
        return node_map

    def link_ancestors(self, samples, ancestors):
        """
        Returns an :class:`EdgeTable` instance describing a subset of the genealogical
        relationships between the nodes in ``samples`` and ``ancestors``.

        Each row ``parent, child, left, right`` in the output table indicates that
        ``child`` has inherited the segment ``[left, right)`` from ``parent`` more
        recently than from any other node in these lists.

        In particular, suppose ``samples`` is a list of nodes such that ``time`` is 0
        for each node, and ``ancestors`` is a list of nodes such that ``time`` is
        greater than 0.0 for each node. Then each row of the output table will show
        an interval ``[left, right)`` over which a node in ``samples`` has inherited
        most recently from a node in ``ancestors``, or an interval over which one of
        these ``ancestors`` has inherited most recently from another node in
        ``ancestors``.

        The following table shows which ``parent->child`` pairs will be shown in the
        output of ``link_ancestors``.
        A node is a relevant descendant on a given interval if it also appears somewhere
        in the ``parent`` column of the outputted table.

        ========================  ===============================================
        Type of relationship      Shown in output of ``link_ancestors``
        ------------------------  -----------------------------------------------
        ``ancestor->sample``      Always
        ``ancestor1->ancestor2``  Only if ``ancestor2`` has a relevant descendant
        ``sample1->sample2``      Always
        ``sample->ancestor``      Only if ``ancestor`` has a relevant descendant
        ========================  ===============================================

        The difference between ``samples`` and ``ancestors`` is that information about
        the ancestors of a node in ``ancestors`` will only be retained if it also has a
        relevant descendant, while information about the ancestors of a node in
        ``samples`` will always be retained.
        The node IDs in ``parent`` and ``child`` refer to the IDs in the node table
        of the inputted tree sequence.

        The supplied nodes must be non-empty lists of the node IDs in the tree sequence:
        in particular, they do not have to be *samples* of the tree sequence. The lists
        of ``samples`` and ``ancestors`` may overlap, although adding a node from
        ``samples`` to ``ancestors`` will not change the output. So, setting ``samples``
        and ``ancestors`` to the same list of nodes will find all genealogical
        relationships within this list.

        If none of the nodes in ``ancestors`` or ``samples`` are ancestral to ``samples``
        anywhere in the tree sequence, an empty table will be returned.

        :param list[int] samples: A list of node IDs to retain as samples.
        :param list[int] ancestors: A list of node IDs to use as ancestors.
        :return: An :class:`EdgeTable` instance displaying relationships between
            the `samples` and `ancestors`.
        """
        samples = util.safe_np_int_cast(samples, np.int32)
        ancestors = util.safe_np_int_cast(ancestors, np.int32)
        ll_edge_table = self._ll_tables.link_ancestors(samples, ancestors)
        return EdgeTable(ll_table=ll_edge_table)

    def map_ancestors(self, *args, **kwargs):
        # A deprecated alias for link_ancestors()
        return self.link_ancestors(*args, **kwargs)

    def sort(self, edge_start=0, *, site_start=0, mutation_start=0):
        """
        Sorts the tables in place. This ensures that all tree sequence ordering
        requirements listed in the
        :ref:`sec_valid_tree_sequence_requirements` section are met, as long
        as each site has at most one mutation (see below).

        If the ``edge_start`` parameter is provided, this specifies the index
        in the edge table where sorting should start. Only rows with index
        greater than or equal to ``edge_start`` are sorted; rows before this index
        are not affected. This parameter is provided to allow for efficient sorting
        when the user knows that the edges up to a given index are already sorted.

        If both ``site_start`` and ``mutation_start`` are equal to the number of rows
        in their retrospective tables then neither is sorted. Note that a partial
        non-sorting is not possible, and both or neither must be skipped.

        The node, individual, population and provenance tables are not affected
        by this method.

        Edges are sorted as follows:

        - time of parent, then
        - parent node ID, then
        - child node ID, then
        - left endpoint.

        Note that this sorting order exceeds the
        :ref:`edge sorting requirements <sec_edge_requirements>` for a valid
        tree sequence. For a valid tree sequence, we require that all edges for a
        given parent ID are adjacent, but we do not require that they be listed in
        sorted order.

        Sites are sorted by position, and sites with the same position retain
        their relative ordering.

        Mutations are sorted by site ID, and within the same site are sorted by time.
        Those with equal or unknown time retain their relative ordering. This does not
        currently rearrange tables so that mutations occur after their mutation parents,
        which is a requirement for valid tree sequences.

        Migrations are sorted by ``time``, ``source``, ``dest``, ``left`` and
        ``node`` values. This defines a total sort order, such that any permutation
        of a valid migration table will be sorted into the same output order.
        Note that this sorting order exceeds the
        :ref:`migration sorting requirements <sec_migration_requirements>` for a
        valid tree sequence, which only requires that migrations are sorted by
        time value.

        :param int edge_start: The index in the edge table where sorting starts
            (default=0; must be <= len(edges)).
        :param int site_start: The index in the site table where sorting starts
            (default=0; must be one of [0, len(sites)]).
        :param int mutation_start: The index in the mutation table where sorting starts
            (default=0; must be one of [0, len(mutations)]).
        """
        self._ll_tables.sort(edge_start, site_start, mutation_start)
        # TODO add provenance

    def sort_individuals(self):
        """
        Sorts the individual table in place, so that parents come before children,
        and the parent column is remapped as required. Node references to individuals
        are also updated.
        """
        self._ll_tables.sort_individuals()
        # TODO add provenance

    def canonicalise(self, remove_unreferenced=None):
        """
        This puts the tables in *canonical* form, imposing a stricter order on the
        tables than :ref:`required <sec_valid_tree_sequence_requirements>` for
        a valid tree sequence. In particular, the individual
        and population tables are sorted by the first node that refers to each
        (see :meth:`TreeSequence.subset`). Then, the remaining tables are sorted
        as in :meth:`.sort`, with the modification that mutations are sorted by
        site, then time, then number of descendant mutations (ensuring that
        parent mutations occur before children), then node, then original order
        in the tables. This ensures that any two tables with the same information
        and node order should be identical after canonical sorting (note
        that no canonical order exists for the node table).

        By default, the method removes sites, individuals, and populations that
        are not referenced (by mutations and nodes, respectively). If you wish
        to keep these, pass ``remove_unreferenced=False``, but note that
        unreferenced individuals and populations are put at the end of the tables
        in their original order.

        .. seealso::

            :meth:`.sort` for sorting edges, mutations, and sites, and
            :meth:`.subset` for reordering nodes, individuals, and populations.

        :param bool remove_unreferenced: Whether to remove unreferenced sites,
            individuals, and populations (default=True).
        """
        remove_unreferenced = (
            True if remove_unreferenced is None else remove_unreferenced
        )
        self._ll_tables.canonicalise(remove_unreferenced=remove_unreferenced)
        # TODO add provenance

    def compute_mutation_parents(self):
        """
        Modifies the tables in place, computing the ``parent`` column of the
        mutation table. For this to work, the node and edge tables must be
        valid, and the site and mutation tables must be sorted (see
        :meth:`TableCollection.sort`).  This will produce an error if mutations
        are not sorted (i.e., if a mutation appears before its mutation parent)
        *unless* the two mutations occur on the same branch, in which case
        there is no way to detect the error.

        The ``parent`` of a given mutation is the ID of the next mutation
        encountered traversing the tree upwards from that mutation, or
        ``NULL`` if there is no such mutation.
        """
        self._ll_tables.compute_mutation_parents()
        # TODO add provenance

    def compute_mutation_times(self):
        """
        Modifies the tables in place, computing valid values for the ``time`` column of
        the mutation table. For this to work, the node and edge tables must be
        valid, and the site and mutation tables must be sorted and indexed(see
        :meth:`TableCollection.sort` and :meth:`TableCollection.build_index`).

        For a single mutation on an edge at a site, the ``time`` assigned to a mutation
        by this method is the mid-point between the times of the nodes above and below
        the mutation. In the case where there is more than one mutation on an edge for
        a site, the times are evenly spread along the edge. For mutations that are
        above a root node, the time of the root node is assigned.

        The mutation table will be sorted if the new times mean that the original order
        is no longer valid.

        """
        self._ll_tables.compute_mutation_times()
        # TODO add provenance

    def deduplicate_sites(self):
        """
        Modifies the tables in place, removing entries in the site table with
        duplicate ``position`` (and keeping only the *first* entry for each
        site), and renumbering the ``site`` column of the mutation table
        appropriately.  This requires the site table to be sorted by position.

        .. warning:: This method does not sort the tables afterwards, so
            mutations may no longer be sorted by time.
        """
        self._ll_tables.deduplicate_sites()
        # TODO add provenance

    def delete_sites(self, site_ids, record_provenance=True):
        """
        Remove the specified sites entirely from the sites and mutations tables in this
        collection. This is identical to :meth:`TreeSequence.delete_sites` but acts
        *in place* to alter the data in this :class:`TableCollection`.

        :param list[int] site_ids: A list of site IDs specifying the sites to remove.
        :param bool record_provenance: If ``True``, add details of this operation
            to the provenance table in this TableCollection. (Default: ``True``).
        """
        keep_sites = np.ones(len(self.sites), dtype=bool)
        site_ids = util.safe_np_int_cast(site_ids, np.int32)
        if np.any(site_ids < 0) or np.any(site_ids >= len(self.sites)):
            raise ValueError("Site ID out of bounds")
        keep_sites[site_ids] = 0
        new_as, new_as_offset = keep_with_offset(
            keep_sites, self.sites.ancestral_state, self.sites.ancestral_state_offset
        )
        new_md, new_md_offset = keep_with_offset(
            keep_sites, self.sites.metadata, self.sites.metadata_offset
        )
        self.sites.set_columns(
            position=self.sites.position[keep_sites],
            ancestral_state=new_as,
            ancestral_state_offset=new_as_offset,
            metadata=new_md,
            metadata_offset=new_md_offset,
        )
        # We also need to adjust the mutations table, as it references into sites
        keep_mutations = keep_sites[self.mutations.site]
        new_ds, new_ds_offset = keep_with_offset(
            keep_mutations,
            self.mutations.derived_state,
            self.mutations.derived_state_offset,
        )
        new_md, new_md_offset = keep_with_offset(
            keep_mutations, self.mutations.metadata, self.mutations.metadata_offset
        )
        # Site numbers will have changed
        site_map = np.cumsum(keep_sites, dtype=self.mutations.site.dtype) - 1
        # Mutation numbers will change, so the parent references need altering
        mutation_map = np.cumsum(keep_mutations, dtype=self.mutations.parent.dtype) - 1
        # Map parent == -1 to -1, and check this has worked (assumes tskit.NULL == -1)
        mutation_map = np.append(mutation_map, -1).astype(self.mutations.parent.dtype)
        assert mutation_map[tskit.NULL] == tskit.NULL
        self.mutations.set_columns(
            site=site_map[self.mutations.site[keep_mutations]],
            node=self.mutations.node[keep_mutations],
            time=self.mutations.time[keep_mutations],
            derived_state=new_ds,
            derived_state_offset=new_ds_offset,
            parent=mutation_map[self.mutations.parent[keep_mutations]],
            metadata=new_md,
            metadata_offset=new_md_offset,
        )
        if record_provenance:
            # TODO replace with a version of https://github.com/tskit-dev/tskit/pull/243
            parameters = {"command": "delete_sites", "TODO": "add parameters"}
            self.provenances.add_row(
                record=json.dumps(provenance.get_provenance_dict(parameters))
            )

    def delete_intervals(self, intervals, simplify=True, record_provenance=True):
        """
        Delete all information from this set of tables which lies *within* the
        specified list of genomic intervals. This is identical to
        :meth:`TreeSequence.delete_intervals` but acts *in place* to alter
        the data in this :class:`TableCollection`.

        :param array_like intervals: A list (start, end) pairs describing the
            genomic intervals to delete. Intervals must be non-overlapping and
            in increasing order. The list of intervals must be interpretable as a
            2D numpy array with shape (N, 2), where N is the number of intervals.
        :param bool simplify: If True, run simplify on the tables so that nodes
            no longer used are discarded. (Default: True).
        :param bool record_provenance: If ``True``, add details of this operation
            to the provenance table in this TableCollection. (Default: ``True``).
        """
        self.keep_intervals(
            util.negate_intervals(intervals, 0, self.sequence_length),
            simplify=simplify,
            record_provenance=False,
        )
        if record_provenance:
            parameters = {"command": "delete_intervals", "TODO": "add parameters"}
            self.provenances.add_row(
                record=json.dumps(provenance.get_provenance_dict(parameters))
            )

    def keep_intervals(self, intervals, simplify=True, record_provenance=True):
        """
        Delete all information from this set of tables which lies *outside* the
        specified list of genomic intervals. This is identical to
        :meth:`TreeSequence.keep_intervals` but acts *in place* to alter
        the data in this :class:`TableCollection`.

        :param array_like intervals: A list (start, end) pairs describing the
            genomic intervals to keep. Intervals must be non-overlapping and
            in increasing order. The list of intervals must be interpretable as a
            2D numpy array with shape (N, 2), where N is the number of intervals.
        :param bool simplify: If True, run simplify on the tables so that nodes
            no longer used are discarded. Must be ``False`` if input tree sequence
            includes migrations. (Default: True).
        :param bool record_provenance: If ``True``, add details of this operation
            to the provenance table in this TableCollection. (Default: ``True``).
        """
        intervals = util.intervals_to_np_array(intervals, 0, self.sequence_length)

        edges = self.edges.copy()
        self.edges.clear()
        migrations = self.migrations.copy()
        self.migrations.clear()
        keep_sites = np.repeat(False, self.sites.num_rows)
        for s, e in intervals:
            curr_keep_sites = np.logical_and(
                self.sites.position >= s, self.sites.position < e
            )
            keep_sites = np.logical_or(keep_sites, curr_keep_sites)
            keep_edges = np.logical_not(
                np.logical_or(edges.right <= s, edges.left >= e)
            )
            metadata, metadata_offset = keep_with_offset(
                keep_edges, edges.metadata, edges.metadata_offset
            )
            self.edges.append_columns(
                left=np.fmax(s, edges.left[keep_edges]),
                right=np.fmin(e, edges.right[keep_edges]),
                parent=edges.parent[keep_edges],
                child=edges.child[keep_edges],
                metadata=metadata,
                metadata_offset=metadata_offset,
            )
            keep_migrations = np.logical_not(
                np.logical_or(migrations.right <= s, migrations.left >= e)
            )
            metadata, metadata_offset = keep_with_offset(
                keep_migrations, migrations.metadata, migrations.metadata_offset
            )
            self.migrations.append_columns(
                left=np.fmax(s, migrations.left[keep_migrations]),
                right=np.fmin(e, migrations.right[keep_migrations]),
                node=migrations.node[keep_migrations],
                source=migrations.source[keep_migrations],
                dest=migrations.dest[keep_migrations],
                time=migrations.time[keep_migrations],
                metadata=metadata,
                metadata_offset=metadata_offset,
            )
        self.delete_sites(
            np.where(np.logical_not(keep_sites))[0], record_provenance=False
        )

        self.sort()
        if simplify:
            self.simplify(record_provenance=False)
        if record_provenance:
            parameters = {"command": "keep_intervals", "TODO": "add parameters"}
            self.provenances.add_row(
                record=json.dumps(provenance.get_provenance_dict(parameters))
            )

    def _check_trim_conditions(self):
        if self.migrations.num_rows > 0:
            if (np.min(self.migrations.left) < np.min(self.edges.left)) and (
                np.max(self.migrations.right) > np.max(self.edges.right)
            ):
                raise ValueError(
                    "Cannot trim a tree sequence with migrations which exist to the"
                    "left of the leftmost edge or to the right of the rightmost edge."
                )
        if self.edges.num_rows == 0:
            raise ValueError(
                "Trimming a tree sequence with no edges would reduce the sequence length"
                " to zero, which is not allowed"
            )

    def ltrim(self, record_provenance=True):
        """
        Reset the coordinate system used in these tables, changing the left and right
        genomic positions in the edge table such that the leftmost edge now starts at
        position 0. This is identical to :meth:`TreeSequence.ltrim` but acts *in place*
        to alter the data in this :class:`TableCollection`.

        :param bool record_provenance: If ``True``, add details of this operation
            to the provenance table in this TableCollection. (Default: ``True``).
        """
        self._check_trim_conditions()
        leftmost = np.min(self.edges.left)
        self.delete_sites(
            np.where(self.sites.position < leftmost), record_provenance=False
        )
        self.edges.set_columns(
            left=self.edges.left - leftmost,
            right=self.edges.right - leftmost,
            parent=self.edges.parent,
            child=self.edges.child,
        )
        self.sites.set_columns(
            position=self.sites.position - leftmost,
            ancestral_state=self.sites.ancestral_state,
            ancestral_state_offset=self.sites.ancestral_state_offset,
            metadata=self.sites.metadata,
            metadata_offset=self.sites.metadata_offset,
        )
        self.migrations.set_columns(
            left=self.migrations.left - leftmost,
            right=self.migrations.right - leftmost,
            time=self.migrations.time,
            node=self.migrations.node,
            source=self.migrations.source,
            dest=self.migrations.dest,
        )
        self.sequence_length = self.sequence_length - leftmost
        if record_provenance:
            # TODO replace with a version of https://github.com/tskit-dev/tskit/pull/243
            parameters = {
                "command": "ltrim",
            }
            self.provenances.add_row(
                record=json.dumps(provenance.get_provenance_dict(parameters))
            )

    def rtrim(self, record_provenance=True):
        """
        Reset the ``sequence_length`` property so that the sequence ends at the end of
        the last edge. This is identical to :meth:`TreeSequence.rtrim` but acts
        *in place* to alter the data in this :class:`TableCollection`.

        :param bool record_provenance: If ``True``, add details of this operation
            to the provenance table in this TableCollection. (Default: ``True``).
        """
        self._check_trim_conditions()
        rightmost = np.max(self.edges.right)
        self.delete_sites(
            np.where(self.sites.position >= rightmost), record_provenance=False
        )
        self.sequence_length = rightmost
        if record_provenance:
            # TODO replace with a version of https://github.com/tskit-dev/tskit/pull/243
            parameters = {
                "command": "rtrim",
            }
            self.provenances.add_row(
                record=json.dumps(provenance.get_provenance_dict(parameters))
            )

    def trim(self, record_provenance=True):
        """
        Trim away any empty regions on the right and left of the tree sequence encoded by
        these tables. This is identical to :meth:`TreeSequence.trim` but acts *in place*
        to alter the data in this :class:`TableCollection`.

        :param bool record_provenance: If ``True``, add details of this operation
            to the provenance table in this TableCollection. (Default: ``True``).
        """
        self.rtrim(record_provenance=False)
        self.ltrim(record_provenance=False)
        if record_provenance:
            # TODO replace with a version of https://github.com/tskit-dev/tskit/pull/243
            parameters = {
                "command": "trim",
            }
            self.provenances.add_row(
                record=json.dumps(provenance.get_provenance_dict(parameters))
            )

    def delete_older(self, time):
        """
        Deletes edge, mutation and migration information at least as old as
        the specified time.

        .. seealso:: This method is similar to the higher-level
            :meth:`TreeSequence.decapitate` method, which also splits
            edges that intersect with the given time.
            :meth:`TreeSequence.decapitate`
            is more useful for most purposes, and may be what
            you need instead of this method!

        For the purposes of this method, an edge covers the times from the
        child node up until the *parent* node, so that any any edge with parent
        node time > ``time`` will be removed.

        Any mutation whose time is >= ``time`` will be removed. A mutation's time
        is its associated ``time`` value, or the time of its node if the
        mutation's time was marked as unknown (:data:`UNKNOWN_TIME`).

        Any migration with time >= ``time`` will be removed.

        The node table is not affected by this operation.

        .. note:: This method does not have any specific sorting requirements
            and will maintain mutation parent mappings.

        :param float time: The cutoff time.
        """
        self._ll_tables.delete_older(time)

    def clear(
        self,
        clear_provenance=False,
        clear_metadata_schemas=False,
        clear_ts_metadata_and_schema=False,
    ):
        """
        Remove all rows of the data tables, optionally remove provenance, metadata
        schemas and ts-level metadata.

        :param bool clear_provenance: If ``True``, remove all rows of the provenance
            table. (Default: ``False``).
        :param bool clear_metadata_schemas: If ``True``, clear the table metadata
            schemas. (Default: ``False``).
        :param bool clear_ts_metadata_and_schema: If ``True``, clear the tree-sequence
            level metadata and schema (Default: ``False``).
        """
        self._ll_tables.clear(
            clear_provenance=clear_provenance,
            clear_metadata_schemas=clear_metadata_schemas,
            clear_ts_metadata_and_schema=clear_ts_metadata_and_schema,
        )

    def has_index(self):
        """
        Returns True if this TableCollection is indexed. See :ref:`sec_table_indexes`
        for information on indexes.
        """
        return bool(self._ll_tables.has_index())

    def build_index(self):
        """
        Builds an index on this TableCollection. Any existing indexes are automatically
        dropped.  See :ref:`sec_table_indexes` for information on indexes.
        """
        self._ll_tables.build_index()

    def drop_index(self):
        """
        Drops any indexes present on this table collection. If the tables are not
        currently indexed this method has no effect.  See :ref:`sec_table_indexes`
        for information on indexes.
        """
        self._ll_tables.drop_index()

    def subset(
        self,
        nodes,
        record_provenance=True,
        *,
        reorder_populations=None,
        remove_unreferenced=None,
    ):
        """
        Modifies the tables in place to contain only the entries referring to
        the provided list of node IDs, with nodes reordered according to the
        order they appear in the list. Other tables are :meth:`sorted <sort>`
        to conform to the :ref:`sec_valid_tree_sequence_requirements`, and
        additionally sorted as described in the documentation for the equivalent
        tree sequence method :meth:`TreeSequence.subset`: please see this for more
        detail.

        :param list nodes: The list of nodes for which to retain information. This
            may be a numpy array (or array-like) object (dtype=np.int32).
        :param bool record_provenance: Whether to record a provenance entry
            in the provenance table for this operation.
        :param bool reorder_populations: Whether to reorder the population table
            (default: True).  If False, the population table will not be altered
            in any way.
        :param bool remove_unreferenced: Whether sites, individuals, and populations
            that are not referred to by any retained entries in the tables should
            be removed (default: True). See the description for details.
        """
        reorder_populations = (
            True if reorder_populations is None else reorder_populations
        )
        remove_unreferenced = (
            True if remove_unreferenced is None else remove_unreferenced
        )
        nodes = util.safe_np_int_cast(nodes, np.int32)
        self._ll_tables.subset(
            nodes,
            reorder_populations=reorder_populations,
            remove_unreferenced=remove_unreferenced,
        )
        self.sort()
        if record_provenance:
            parameters = {"command": "subset", "nodes": nodes.tolist()}
            self.provenances.add_row(
                record=json.dumps(provenance.get_provenance_dict(parameters))
            )

    def union(
        self,
        other,
        node_mapping,
        check_shared_equality=True,
        add_populations=True,
        record_provenance=True,
    ):
        """
        Modifies the table collection in place by adding the non-shared
        portions of ``other`` to itself. To perform the node-wise union,
        the method relies on a ``node_mapping`` array, that maps nodes in
        ``other`` to its equivalent node in ``self`` or ``tskit.NULL`` if
        the node is exclusive to ``other``. See :meth:`TreeSequence.union` for a more
        detailed description.

        :param TableCollection other: Another table collection.
        :param list node_mapping: An array of node IDs that relate nodes in
            ``other`` to nodes in ``self``: the k-th element of ``node_mapping``
            should be the index of the equivalent node in ``self``, or
            ``tskit.NULL`` if the node is not present in ``self`` (in which case it
            will be added to self).
        :param bool check_shared_equality: If True, the shared portions of the
            table collections will be checked for equality.
        :param bool add_populations: If True, nodes new to ``self`` will be
            assigned new population IDs.
        :param bool record_provenance: Whether to record a provenance entry
            in the provenance table for this operation.
        """
        node_mapping = util.safe_np_int_cast(node_mapping, np.int32)
        self._ll_tables.union(
            other._ll_tables,
            node_mapping,
            check_shared_equality=check_shared_equality,
            add_populations=add_populations,
        )
        if record_provenance:
            other_records = [prov.record for prov in other.provenances]
            other_timestamps = [prov.timestamp for prov in other.provenances]
            parameters = {
                "command": "union",
                "other": {"timestamp": other_timestamps, "record": other_records},
                "node_mapping": node_mapping.tolist(),
            }
            self.provenances.add_row(
                record=json.dumps(provenance.get_provenance_dict(parameters))
            )

    def ibd_segments(
        self,
        *,
        within=None,
        between=None,
        max_time=None,
        min_span=None,
        store_pairs=None,
        store_segments=None,
    ):
        """
        Equivalent to the :meth:`TreeSequence.ibd_segments` method; please see its
        documentation for more details, and use this method only if you specifically need
        to work with a :class:`TableCollection` object.

        This method has the same data requirements as
        :meth:`TableCollection.simplify`. In particular, the tables in the collection
        have :ref:`required <sec_valid_tree_sequence_requirements>` sorting orders.
        To enforce this, you can call :meth:`TableCollection.sort` before using this
        method. If the edge table contains any edges with identical
        parents and children over adjacent genomic intervals, any IBD intervals
        underneath the edges will also be split across the breakpoint(s). To prevent this
        behaviour in this situation, use :meth:`EdgeTable.squash` beforehand.

        :param list within: As for the :meth:`TreeSequence.ibd_segments` method.
        :param list[list] between: As for the :meth:`TreeSequence.ibd_segments` method.
        :param float max_time: As for the :meth:`TreeSequence.ibd_segments` method.
        :param float min_span: As for the :meth:`TreeSequence.ibd_segments` method.
        :param bool store_pairs: As for the :meth:`TreeSequence.ibd_segments` method.
        :param bool store_segments: As for the :meth:`TreeSequence.ibd_segments` method.
        :return: An :class:`.IdentitySegments` object containing the recorded
            IBD information.
        :rtype: IdentitySegments
        """
        max_time = np.inf if max_time is None else max_time
        min_span = 0 if min_span is None else min_span
        store_pairs = False if store_pairs is None else store_pairs
        store_segments = False if store_segments is None else store_segments
        if within is not None and between is not None:
            raise ValueError(
                "The ``within`` and ``between`` arguments are mutually exclusive"
            )
        if between is not None:
            sample_set_sizes = np.array(
                [len(sample_set) for sample_set in between], dtype=np.uint64
            )
            # hstack has some annoying quirks around its handling of empty
            # lists which we need to work around. In a way it would be more
            # convenient to detect these conditions as errors, but then we
            # end up having to workaround edge cases in the tests and its
            # mathematically neater this way.
            pre_flattened = [lst for lst in between if len(lst) > 0]
            if len(pre_flattened) == 0:
                flattened = []
            else:
                flattened = util.safe_np_int_cast(np.hstack(pre_flattened), np.int32)
            ll_result = self._ll_tables.ibd_segments_between(
                sample_set_sizes=sample_set_sizes,
                sample_sets=flattened,
                max_time=max_time,
                min_span=min_span,
                store_pairs=store_pairs,
                store_segments=store_segments,
            )
        else:
            if within is not None:
                within = util.safe_np_int_cast(within, np.int32)
            ll_result = self._ll_tables.ibd_segments_within(
                samples=within,
                max_time=max_time,
                min_span=min_span,
                store_pairs=store_pairs,
                store_segments=store_segments,
            )
        return IdentitySegments(
            ll_result,
            max_time=max_time,
            min_span=min_span,
            store_pairs=store_pairs,
            store_segments=store_segments,
        )


--- ../../tskit/python/tskit/__init__.py ---

 
import _tskit

#: Special reserved value representing a null ID.
NULL = _tskit.NULL

#: Special value representing missing data in a genotype array
MISSING_DATA = _tskit.MISSING_DATA

#: Node flag value indicating that it is a sample.
NODE_IS_SAMPLE = _tskit.NODE_IS_SAMPLE

#: Constant representing the forward direction of travel (i.e.,
#: increasing genomic coordinate values).
FORWARD = _tskit.FORWARD

#: Constant representing the reverse direction of travel (i.e.,
#: decreasing genomic coordinate values).
REVERSE = _tskit.REVERSE

#: The allele mapping where the strings "0" and "1" map to genotype
#: values 0 and 1.
ALLELES_01 = ("0", "1")

#: The allele mapping where the four nucleotides A, C, G and T map to
#: the genotype integers 0, 1, 2, and 3, respectively.
ALLELES_ACGT = ("A", "C", "G", "T")

#: Special NAN value used to indicate unknown mutation times. Since this is a
#: NAN value, you cannot use `==` to test for it. Use :func:`is_unknown_time` instead.
UNKNOWN_TIME = _tskit.UNKNOWN_TIME

#: Default value of ts.time_units
TIME_UNITS_UNKNOWN = _tskit.TIME_UNITS_UNKNOWN

#: ts.time_units value when dimension is uncalibrated
TIME_UNITS_UNCALIBRATED = _tskit.TIME_UNITS_UNCALIBRATED

#: Options for printing to strings and HTML, modify with tskit.set_print_options.
_print_options = {"max_lines": 40}

TABLE_NAMES = [
    "individuals",
    "nodes",
    "edges",
    "migrations",
    "sites",
    "mutations",
    "populations",
    "provenances",
]


from tskit.provenance import __version__  # NOQA
from tskit.provenance import validate_provenance  # NOQA
from tskit.formats import *  # NOQA
from tskit.trees import *  # NOQA
from tskit.genotypes import Variant  # NOQA
from tskit.tables import *  # NOQA
from tskit.stats import *  # NOQA
from tskit.combinatorics import (  # NOQA
    all_trees,
    all_tree_shapes,
    all_tree_labellings,
    TopologyCounter,
    Rank,
)
from tskit.drawing import SVGString  # NOQA
from tskit.exceptions import *  # NOQA
from tskit.util import *  # NOQA
from tskit.metadata import *  # NOQA
from tskit.text_formats import *  # NOQA
from tskit.intervals import RateMap  # NOQA


--- ../../tskit/python/tskit/drawing.py ---


"""
Module responsible for visualisations.
"""
import collections
import itertools
import logging
import math
import numbers
import operator
import warnings
from dataclasses import dataclass
from typing import List
from typing import Mapping
from typing import Union

import numpy as np
import svgwrite

import tskit
import tskit.util as util
from _tskit import NODE_IS_SAMPLE
from _tskit import NULL

LEFT = "left"
RIGHT = "right"
TOP = "top"
BOTTOM = "bottom"

# constants for whether to plot a tree in a tree sequence
OMIT = 1
LEFT_CLIP = 2
RIGHT_CLIP = 4
OMIT_MIDDLE = 8


@dataclass
class Offsets:
    "Used when x_lim set, and displayed ts has been cut down by keep_intervals"
    tree: int = 0
    site: int = 0
    mutation: int = 0


@dataclass(frozen=True)
class Timescaling:
    "Class used to transform the time axis"
    max_time: float
    min_time: float
    plot_min: float
    plot_range: float
    use_log_transform: bool

    def __post_init__(self):
        if self.plot_range < 0:
            raise ValueError("Image size too small to allow space to plot tree")
        if self.use_log_transform:
            if self.min_time < 0:
                raise ValueError("Cannot use a log scale if there are negative times")
            super().__setattr__("transform", self.log_transform)
        else:
            super().__setattr__("transform", self.linear_transform)

    def log_transform(self, y):
        "Standard log transform but allowing for values of 0 by adding 1"
        delta = 1 if self.min_time == 0 else 0
        log_max = np.log(self.max_time + delta)
        log_min = np.log(self.min_time + delta)
        y_scale = self.plot_range / (log_max - log_min)
        return self.plot_min - (np.log(y + delta) - log_min) * y_scale

    def linear_transform(self, y):
        y_scale = self.plot_range / (self.max_time - self.min_time)
        return self.plot_min - (y - self.min_time) * y_scale


class SVGString(str):
    "A string containing an SVG representation"

    def _repr_svg_(self):
        """
        Simply return the SVG string: called by jupyter notebooks to render trees.
        """
        return self


def check_orientation(orientation):
    if orientation is None:
        orientation = TOP
    else:
        orientation = orientation.lower()
        orientations = [LEFT, RIGHT, TOP, BOTTOM]
        if orientation not in orientations:
            raise ValueError(f"Unknown orientiation: choose from {orientations}")
    return orientation


def check_max_time(max_time, allow_numeric=True):
    if max_time is None:
        max_time = "tree"
    is_numeric = isinstance(max_time, numbers.Real)
    if max_time not in ["tree", "ts"] and not allow_numeric:
        raise ValueError("max_time must be 'tree' or 'ts'")
    if max_time not in ["tree", "ts"] and (allow_numeric and not is_numeric):
        raise ValueError("max_time must be a numeric value or one of 'tree' or 'ts'")
    return max_time


def check_min_time(min_time, allow_numeric=True):
    if min_time is None:
        min_time = "tree"
    if allow_numeric:
        is_numeric = isinstance(min_time, numbers.Real)
        if min_time not in ["tree", "ts"] and not is_numeric:
            raise ValueError(
                "min_time must be a numeric value or one of 'tree' or 'ts'"
            )
    else:
        if min_time not in ["tree", "ts"]:
            raise ValueError("min_time must be 'tree' or 'ts'")
    return min_time


def check_time_scale(time_scale):
    if time_scale is None:
        time_scale = "time"
    if time_scale not in ["time", "log_time", "rank"]:
        raise ValueError("time_scale must be 'time', 'log_time' or 'rank'")
    return time_scale


def check_format(format):  # noqa A002
    if format is None:
        format = "SVG"  # noqa A001
    fmt = format.lower()
    supported_formats = ["svg", "ascii", "unicode"]
    if fmt not in supported_formats:
        raise ValueError(
            "Unknown format '{}'. Supported formats are {}".format(
                format, supported_formats
            )
        )
    return fmt


def check_order(order):
    """
    Checks the specified drawing order is valid and returns the corresponding
    tree traversal order.
    """
    if order is None:
        order = "minlex"
    traversal_orders = {
        "minlex": "minlex_postorder",
        "tree": "postorder",
    }
    # Silently accept a tree traversal order as a valid order, so we can
    # call this check twice if necessary
    if order in traversal_orders.values():
        return order
    if order not in traversal_orders:
        raise ValueError(
            f"Unknown display order '{order}'. "
            f"Supported orders are {list(traversal_orders.keys())}"
        )
    return traversal_orders[order]


def check_x_scale(x_scale):
    """
    Checks the specified x_scale is valid and sets default if None
    """
    if x_scale is None:
        x_scale = "physical"
    x_scales = ["physical", "treewise"]
    if x_scale not in x_scales:
        raise ValueError(
            f"Unknown display x_scale '{x_scale}'. " f"Supported orders are {x_scales}"
        )
    return x_scale


def check_x_lim(x_lim, max_x):
    """
    Checks the specified x_limits are valid and sets default if None.
    """
    if x_lim is None:
        x_lim = (None, None)
    if len(x_lim) != 2:
        raise ValueError("The x_lim parameter must be a list of length 2, or None")
    try:
        if x_lim[0] is not None and x_lim[0] < 0:
            raise ValueError("x_lim[0] cannot be negative")
        if x_lim[1] is not None and x_lim[1] > max_x:
            raise ValueError("x_lim[1] cannot be greater than the sequence length")
        if x_lim[0] is not None and x_lim[1] is not None and x_lim[0] >= x_lim[1]:
            raise ValueError("x_lim[0] must be less than x_lim[1]")
    except TypeError:
        raise TypeError("x_lim parameters must be numeric")
    return x_lim


def create_tick_labels(tick_values, decimal_places=2):
    """
    If tick_values are numeric, round the labels to X decimal_places, but do not print
    decimals if all values are integers
    """
    try:
        integer_ticks = np.all(np.round(tick_values) == tick_values)
    except TypeError:
        return tick_values
    label_precision = 0 if integer_ticks else decimal_places
    return [f"{lab:.{label_precision}f}" for lab in tick_values]


def clip_ts(ts, x_min, x_max, max_num_trees=None):
    """
    Culls the edges of the tree sequence outside the limits of x_min and x_max if
    necessary, and flags internal trees for omission if there are more than
    max_num_trees in the tree sequence

    Returns the new tree sequence using the same genomic scale, and an
    array specifying which trees to actually plot from it. This array contains
    information about whether a plotted tree was clipped, because clipping can
    cause the rightmost and leftmost tree in this new TS to have reduced spans, and
    should be displayed by omitting the appropriate breakpoint.

    If x_min is None, we take it to be 0 if the first tree has edges or sites, or
    ``min(edges.left)`` if the first tree represents an empty region.
    Similarly, if x_max is None we take it to be ``ts.sequence_length`` if the last tree
    has edges or mutations, or ``ts.last().interval.left`` if the last tree represents
    an empty region.

    To plot the full ts, including empty flanking regions, specify x_limits of
    [0, seq_len].

    """
    edges = ts.tables.edges
    sites = ts.tables.sites
    offsets = Offsets()
    if x_min is None:
        if ts.num_edges == 0:
            if ts.num_sites == 0:
                raise ValueError(
                    "To plot an empty tree sequence, specify x_lim=[0, sequence_length]"
                )
            x_min = 0
        else:
            x_min = np.min(edges.left)
            if ts.num_sites > 0 and np.min(sites.position) < x_min:
                x_min = 0  # First region has no edges, but does have sites => keep
    if x_max is None:
        if ts.num_edges == 0:
            if ts.num_sites == 0:
                raise ValueError(
                    "To plot an empty tree sequence, specify x_lim=[0, sequence_length]"
                )
            x_max = ts.sequence_length
        else:
            x_max = np.max(edges.right)
            if ts.num_sites > 0 and np.max(sites.position) > x_max:
                x_max = ts.sequence_length  # Last region has sites but no edges => keep

    if max_num_trees is None:
        max_num_trees = np.inf

    if max_num_trees < 2:
        raise ValueError("Must show at least 2 trees when clipping a tree sequence")

    if (x_min > 0) or (x_max < ts.sequence_length):
        old_breaks = ts.breakpoints(as_array=True)
        offsets.tree = np.searchsorted(old_breaks, x_min, "right") - 2
        offsets.site = np.searchsorted(sites.position, x_min)
        offsets.mutation = np.searchsorted(ts.tables.mutations.site, offsets.site)
        ts = ts.keep_intervals([[x_min, x_max]], simplify=False)
        if ts.num_edges == 0:
            raise ValueError(
                f"Can't limit plotting from {x_min} to {x_max} as whole region is empty"
            )
        edges = ts.tables.edges
        sites = ts.tables.sites
        trees_start = np.min(edges.left)
        trees_end = np.max(edges.right)
        tree_status = np.zeros(ts.num_trees, dtype=np.uint8)
        # Are the leftmost/rightmost regions completely empty - if so, don't plot them
        if 0 < x_min <= trees_start and (
            ts.num_sites == 0 or trees_start <= np.min(sites.position)
        ):
            tree_status[0] = OMIT
        if trees_end <= x_max < ts.sequence_length and (
            ts.num_sites == 0 or trees_end >= np.max(sites.position)
        ):
            tree_status[-1] = OMIT

        # Which breakpoints are new ones, as a result of clipping
        new_breaks = np.logical_not(np.isin(ts.breakpoints(as_array=True), old_breaks))
        tree_status[new_breaks[:-1]] |= LEFT_CLIP
        tree_status[new_breaks[1:]] |= RIGHT_CLIP
    else:
        tree_status = np.zeros(ts.num_trees, dtype=np.uint8)

    first_tree = 1 if tree_status[0] & OMIT else 0
    last_tree = ts.num_trees - 2 if tree_status[-1] & OMIT else ts.num_trees - 1
    num_shown_trees = last_tree - first_tree + 1
    if num_shown_trees > max_num_trees:
        num_start_trees = max_num_trees // 2 + (1 if max_num_trees % 2 else 0)
        num_end_trees = max_num_trees // 2
        assert num_start_trees + num_end_trees == max_num_trees
        tree_status[
            (first_tree + num_start_trees) : (last_tree - num_end_trees + 1)
        ] = (OMIT | OMIT_MIDDLE)

    return ts, tree_status, offsets


def check_y_ticks(ticks: Union[List, Mapping, None]) -> Mapping:
    """
    Later we might want to implement a tick locator function, such that e.g. ticks=5
    selects ~5 nicely spaced tick locations (with sensible behaviour for log scales)
    """
    if ticks is None:
        return {}
    if isinstance(ticks, Mapping):
        return dict(zip(ticks, create_tick_labels(list(ticks.values()))))
    return dict(zip(ticks, create_tick_labels(ticks)))


def rnd(x):
    """
    Round a number so that the output SVG doesn't have unneeded precision
    """
    digits = 6
    if x == 0 or not math.isfinite(x):
        return x
    digits -= math.ceil(math.log10(abs(x)))
    x = round(x, digits)
    if int(x) == x:
        return int(x)
    return x


def bold_integer(number):
    # For simple integers, it's easier to use bold unicode characters
    # than to try to get the SVG to render a bold font for part of a string
    return "".join("𝟎𝟏𝟐𝟑𝟒𝟓𝟔𝟕𝟖𝟗"[int(digit)] for digit in str(number))


def edge_and_sample_nodes(ts, omit_regions=None):
    """
    Return ids of nodes which are mentioned in an edge in this tree sequence or which
    are samples: nodes not connected to an edge are often found if x_lim is specified.
    """
    if omit_regions is None or len(omit_regions) == 0:
        ids = np.concatenate((ts.edges_child, ts.edges_parent))
    else:
        ids = np.array([], dtype=ts.edges_child.dtype)
        edges = ts.tables.edges
        assert omit_regions.shape[1] == 2
        omit_regions = omit_regions.flatten()
        assert np.all(omit_regions == np.unique(omit_regions))  # Check they're in order
        use_regions = np.concatenate(([0.0], omit_regions, [ts.sequence_length]))
        use_regions = use_regions.reshape(-1, 2)
        for left, right in use_regions:
            used_edges = edges[np.logical_and(edges.left >= left, edges.right < right)]
            ids = np.concatenate((ids, used_edges.child, used_edges.parent))
    return np.unique(
        np.concatenate((ids, np.where(ts.nodes_flags & NODE_IS_SAMPLE)[0]))
    )


def _postorder_tracked_node_traversal(tree, root, collapse_tracked, key_dict=None):
    # Postorder traversal that only descends into subtrees if they contain
    # a tracked node. Additionally, if collapse_tracked is not None, it is
    # interpreted as a proportion, so that we do not descend into a subtree if
    # that proportion or greater of the samples in the subtree are tracked.
    # If key_dict is provided, use this to sort the children. This allows
    # us to put e.g. the subtrees containing the most tracked nodes first.
    # Private function, for use only in drawing.postorder_tracked_minlex_traversal()

    # If we deliberately specify the virtual root, it should also be returned
    is_virtual_root = root == tree.virtual_root
    if root == tskit.NULL:
        root = tree.virtual_root
    stack = [(root, False)]
    while stack:
        u, visited = stack.pop()
        if visited:
            if u != tree.virtual_root or is_virtual_root:
                yield u
        else:
            if tree.num_children(u) == 0:
                yield u
            elif tree.num_tracked_samples(u) == 0:
                yield u
            elif (
                collapse_tracked is not None
                and tree.num_children(u) != 1
                and tree.num_tracked_samples(u)
                >= collapse_tracked * tree.num_samples(u)
            ):
                yield u
            else:
                stack.append((u, True))
                if key_dict is None:
                    stack.extend((c, False) for c in tree.children(u))
                else:
                    stack.extend(
                        sorted(
                            ((c, False) for c in tree.children(u)),
                            key=lambda v: key_dict[v[0]],
                            reverse=True,
                        )
                    )


def _postorder_tracked_minlex_traversal(tree, root=None, *, collapse_tracked=None):
    """
    Postorder traversal for drawing purposes that places child nodes with the
    most tracked sample descendants first (then sorts ties by minlex on leaf node ids).
    Additionally, this traversal only descends into subtrees if they contain a tracked
    node, and may not descend into other subtree, if the ``collapse_tracked``
    parameter is set to a numeric value. More specifically, if the proportion of
    tracked samples in the subtree is greater than or equal to ``collapse_tracked``,
    the subtree is not descended into.
    """

    key_dict = {}
    parent_array = tree.parent_array
    prev = tree.virtual_root
    if root is None:
        root = tskit.NULL
    for u in _postorder_tracked_node_traversal(tree, root, collapse_tracked):
        is_tip = parent_array[prev] != u
        prev = u
        if is_tip:
            # Sort by number of tracked samples (desc), then by minlex
            key_dict[u] = (-tree.num_tracked_samples(u), u)
        else:
            min_tip_id = min(key_dict[v][1] for v in tree.children(u) if v in key_dict)
            key_dict[u] = (-tree.num_tracked_samples(u), min_tip_id)

    return _postorder_tracked_node_traversal(
        tree, root, collapse_tracked, key_dict=key_dict
    )


def draw_tree(
    tree,
    width=None,
    height=None,
    node_labels=None,
    node_colours=None,
    mutation_labels=None,
    mutation_colours=None,
    format=None,  # noqa A002
    edge_colours=None,
    time_scale=None,
    tree_height_scale=None,
    max_time=None,
    min_time=None,
    max_tree_height=None,
    order=None,
    omit_sites=None,
):
    if time_scale is None and tree_height_scale is not None:
        time_scale = tree_height_scale
        # Deprecated in 0.3.6
        warnings.warn(
            "tree_height_scale is deprecated; use time_scale instead",
            FutureWarning,
            stacklevel=4,
        )
    if max_time is None and max_tree_height is not None:
        max_time = max_tree_height
        # Deprecated in 0.3.6
        warnings.warn(
            "max_tree_height is deprecated; use max_time instead",
            FutureWarning,
            stacklevel=4,
        )

    # See tree.draw() for documentation on these arguments.
    fmt = check_format(format)
    if fmt == "svg":
        if width is None:
            width = 200
        if height is None:
            height = 200

        def remap_style(original_map, new_key, none_value):
            if original_map is None:
                return None
            new_map = {}
            for key, value in original_map.items():
                if value is None:
                    new_map[key] = {"style": none_value}
                else:
                    new_map[key] = {"style": f"{new_key}:{value};"}
            return new_map

        # Set style rather than fill & stroke directly to override top stylesheet
        # Old semantics were to not draw the node if colour is None.
        # Setting opacity to zero has the same effect.
        node_attrs = remap_style(node_colours, "fill", "fill-opacity:0;")
        edge_attrs = remap_style(edge_colours, "stroke", "stroke-opacity:0;")
        mutation_attrs = remap_style(mutation_colours, "fill", "fill-opacity:0;")

        node_label_attrs = None
        tree = SvgTree(
            tree,
            (width, height),
            node_labels=node_labels,
            mutation_labels=mutation_labels,
            time_scale=time_scale,
            max_time=max_time,
            min_time=min_time,
            node_attrs=node_attrs,
            edge_attrs=edge_attrs,
            node_label_attrs=node_label_attrs,
            mutation_attrs=mutation_attrs,
            order=order,
            omit_sites=omit_sites,
        )
        return SVGString(tree.drawing.tostring())

    else:
        if width is not None:
            raise ValueError("Text trees do not support width")
        if height is not None:
            raise ValueError("Text trees do not support height")
        if mutation_labels is not None:
            raise ValueError("Text trees do not support mutation_labels")
        if mutation_colours is not None:
            raise ValueError("Text trees do not support mutation_colours")
        if node_colours is not None:
            raise ValueError("Text trees do not support node_colours")
        if edge_colours is not None:
            raise ValueError("Text trees do not support edge_colours")
        if time_scale is not None:
            raise ValueError("Text trees do not support time_scale")

        use_ascii = fmt == "ascii"
        text_tree = VerticalTextTree(
            tree,
            node_labels=node_labels,
            max_time=max_time,
            min_time=min_time,
            use_ascii=use_ascii,
            orientation=TOP,
            order=order,
        )
        return str(text_tree)


def add_class(attrs_dict, classes_str):
    """Adds the classes_str to the 'class' key in attrs_dict, or creates it"""
    try:
        attrs_dict["class"] += " " + classes_str
    except KeyError:
        attrs_dict["class"] = classes_str


@dataclass
class Plotbox:
    total_size: list
    pad_top: float = 0
    pad_left: float = 0
    pad_bottom: float = 0
    pad_right: float = 0

    def set_padding(self, top, left, bottom, right):
        self.pad_top = top
        self.pad_left = left
        self.pad_bottom = bottom
        self.pad_right = right
        self._check()

    @property
    def max_x(self):
        return self.total_size[0]

    @property
    def max_y(self):
        return self.total_size[1]

    @property
    def top(self):  # Alias for consistency with top & bottom
        return self.pad_top

    @property
    def left(self):  # Alias for consistency with top & bottom
        return self.pad_left

    @property
    def bottom(self):
        return self.max_y - self.pad_bottom

    @property
    def right(self):
        return self.max_x - self.pad_right

    @property
    def width(self):
        return self.right - self.left

    @property
    def height(self):
        return self.bottom - self.top

    def __post_init__(self):
        self._check()

    def _check(self):
        if self.width < 1 or self.height < 1:
            raise ValueError("Image size too small to fit")

    def draw(self, dwg, add_to, colour="grey"):
        # used for debugging
        add_to.add(
            dwg.rect(
                (0, 0),
                (self.max_x, self.max_y),
                fill="white",
                fill_opacity=0,
                stroke=colour,
                stroke_dasharray="15,15",
                class_="outer_plotbox",
            )
        )
        add_to.add(
            dwg.rect(
                (self.left, self.top),
                (self.width, self.height),
                fill="white",
                fill_opacity=0,
                stroke=colour,
                stroke_dasharray="5,5",
                class_="inner_plotbox",
            )
        )


class SvgPlot:
    """
    The base class for plotting any box to canvas
    """

    text_height = 14  # May want to calculate this based on a font size
    line_height = text_height * 1.2  # allowing padding above and below a line

    def __init__(
        self,
        size,
        svg_class,
        root_svg_attributes=None,
        canvas_size=None,
    ):
        """
        Creates self.drawing, an svgwrite.Drawing object for further use, and populates
        it with a base group. The root_groups will be populated with
        items that can be accessed from the outside, such as the plotbox, axes, etc.
        """

        if root_svg_attributes is None:
            root_svg_attributes = {}
        if canvas_size is None:
            canvas_size = size
        dwg = svgwrite.Drawing(size=canvas_size, debug=True, **root_svg_attributes)

        self.image_size = size
        self.plotbox = Plotbox(size)
        self.root_groups = {}
        self.svg_class = svg_class
        self.timescaling = None
        self.root_svg_attributes = root_svg_attributes
        self.dwg_base = dwg.add(dwg.g(class_=svg_class))
        self.drawing = dwg

    def get_plotbox(self):
        """
        Get the svgwrite plotbox, creating it if necessary.
        """
        if "plotbox" not in self.root_groups:
            dwg = self.drawing
            self.root_groups["plotbox"] = self.dwg_base.add(dwg.g(class_="plotbox"))
        return self.root_groups["plotbox"]

    def add_text_in_group(self, text, add_to, pos, group_class=None, **kwargs):
        """
        Add the text to the elem within a group; allows text rotations to work smoothly,
        otherwise, if x & y parameters are used to position text, rotations applied to
        the text tag occur around the (0,0) point of the containing group
        """
        dwg = self.drawing
        group_attributes = {"transform": f"translate({rnd(pos[0])} {rnd(pos[1])})"}
        if group_class is not None:
            group_attributes["class_"] = group_class
        grp = add_to.add(dwg.g(**group_attributes))
        grp.add(dwg.text(text, **kwargs))


class SvgSkippedPlot(SvgPlot):
    def __init__(
        self,
        size,
        num_skipped,
    ):
        super().__init__(
            size,
            svg_class="skipped",
        )
        container = self.get_plotbox()
        x = self.plotbox.width / 2
        y = self.plotbox.height / 2
        self.add_text_in_group(
            f"{num_skipped} trees",
            container,
            (x, y - self.line_height / 2),
            text_anchor="middle",
        )
        self.add_text_in_group(
            "skipped", container, (x, y + self.line_height / 2), text_anchor="middle"
        )


class SvgAxisPlot(SvgPlot):
    """
    The class used for plotting either a tree or a tree sequence as an SVG file
    """

    standard_style = (
        ".background path {fill: #808080; fill-opacity: 0}"
        ".background path:nth-child(odd) {fill-opacity: .1}"
        ".x-regions rect {fill: yellow; stroke: black; opacity: 0.5}"  # opaque 4 overlap
        ".axes {font-size: 14px}"
        ".x-axis .tick .lab {font-weight: bold; dominant-baseline: hanging}"
        ".axes, .tree {font-size: 14px; text-anchor: middle}"
        ".axes line, .edge {stroke: black; fill: none}"
        ".axes .ax-skip {stroke-dasharray: 4}"
        ".y-axis .grid {stroke: #FAFAFA}"
        ".node > .sym {fill: black; stroke: none}"
        ".site > .sym {stroke: black}"
        ".mut text {fill: red; font-style: italic}"
        ".mut.extra text {fill: hotpink}"
        ".mut line {fill: none; stroke: none}"  # Default hide mut line to expose edges
        ".mut .sym {fill: none; stroke: red}"
        ".mut.extra .sym {stroke: hotpink}"
        ".node .mut .sym {stroke-width: 1.5px}"
        ".tree text, .tree-sequence text {dominant-baseline: central}"
        ".plotbox .lab.lft {text-anchor: end}"
        ".plotbox .lab.rgt {text-anchor: start}"
        ".polytomy line {stroke: black; stroke-dasharray: 1px, 1px}"
        ".polytomy text {paint-order:stroke;stroke-width:0.3em;stroke:white}"
    )

    # TODO: we may want to make some of the constants below into parameters
    root_branch_fraction = 1 / 8  # Rel root branch len, unless it has a timed mutation
    default_tick_length = 5
    default_tick_length_site = 10
    # Placement of the axes lines within the padding - not used unless axis is plotted
    default_x_axis_offset = 20
    default_y_axis_offset = 40

    def __init__(
        self,
        ts,
        size,
        root_svg_attributes,
        style,
        svg_class,
        time_scale,
        x_axis=None,
        y_axis=None,
        x_label=None,
        y_label=None,
        offsets=None,
        debug_box=None,
        omit_sites=None,
        canvas_size=None,
        mutation_titles=None,
    ):
        super().__init__(
            size,
            svg_class,
            root_svg_attributes,
            canvas_size,
        )
        self.ts = ts
        dwg = self.drawing
        # Put all styles in a single stylesheet (required for Inkscape 0.92)
        style = self.standard_style + ("" if style is None else style)
        dwg.defs.add(dwg.style(style))
        self.debug_box = debug_box
        self.time_scale = check_time_scale(time_scale)
        self.y_axis = y_axis
        self.x_axis = x_axis
        if x_label is None and x_axis:
            x_label = "Genome position"
        if y_label is None and y_axis:
            if time_scale == "rank":
                y_label = "Node time"
            else:
                y_label = "Time ago"
            if ts.time_units != tskit.TIME_UNITS_UNKNOWN:
                y_label += f" ({ts.time_units})"
        self.x_label = x_label
        self.y_label = y_label
        self.offsets = Offsets() if offsets is None else offsets
        self.omit_sites = omit_sites
        self.mutation_titles = {} if mutation_titles is None else mutation_titles
        self.mutations_outside_tree = set()  # mutations in here get an additional class

    def set_spacing(self, top=0, left=0, bottom=0, right=0):
        """
        Set edges, but allow space for axes etc
        """
        self.x_axis_offset = self.default_x_axis_offset
        self.y_axis_offset = self.default_y_axis_offset
        if self.x_label:
            self.x_axis_offset += self.line_height
        if self.y_label:
            self.y_axis_offset += self.line_height
        if self.x_axis:
            bottom += self.x_axis_offset
        if self.y_axis:
            left = self.y_axis_offset  # Override user-provided, so y-axis is at x=0
        self.plotbox.set_padding(top, left, bottom, right)
        if self.debug_box:
            self.root_groups["debug"] = self.dwg_base.add(
                self.drawing.g(class_="debug")
            )
            self.plotbox.draw(self.drawing, self.root_groups["debug"])

    def get_axes(self):
        if "axes" not in self.root_groups:
            self.root_groups["axes"] = self.dwg_base.add(self.drawing.g(class_="axes"))
        return self.root_groups["axes"]

    def draw_x_axis(
        self,
        tick_positions=None,  # np.array of ax ticks below (+ above if sites is None)
        tick_labels=None,  # Tick labels below axis. If None, use the position value
        tick_length_lower=default_tick_length,
        tick_length_upper=None,  # If None, use the same as tick_length_lower
        site_muts=None,  # A dict of site id => mutation to plot as ticks on the x axis
        alternate_dash_positions=None,  # Where to alternate the axis from solid to dash
        x_regions=None,  # A dict of (left, right):label items to place in boxes
    ):
        if not self.x_axis:
            return
        if alternate_dash_positions is None:
            alternate_dash_positions = np.array([])
        if x_regions is None:
            x_regions = {}
        dwg = self.drawing
        axes = self.get_axes()
        x_axis = axes.add(dwg.g(class_="x-axis"))
        if self.x_label:
            self.add_text_in_group(
                self.x_label,
                x_axis,
                pos=((self.plotbox.left + self.plotbox.right) / 2, self.plotbox.max_y),
                group_class="title",
                class_="lab",
                transform="translate(0 -11)",
                text_anchor="middle",
            )
        if len(x_regions) > 0:
            regions_group = x_axis.add(dwg.g(class_="x-regions"))
            for i, ((left, right), label) in enumerate(x_regions.items()):
                if not (0 <= left < right <= self.ts.sequence_length):
                    raise ValueError(
                        f"Invalid coordinates ({left} to {right}) for x-axis region"
                    )
                x1 = self.x_transform(left)
                x2 = self.x_transform(right)
                y = self.plotbox.max_y - self.x_axis_offset
                region = regions_group.add(dwg.g(class_=f"r{i}"))
                region.add(
                    dwg.rect((x1, y), (x2 - x1, self.line_height), class_="r{i}")
                )
                self.add_text_in_group(
                    label,
                    region,
                    pos=((x2 + x1) / 2, y + self.line_height / 2),
                    class_="lab",
                    text_anchor="middle",
                )
        if tick_length_upper is None:
            tick_length_upper = tick_length_lower
        y = rnd(self.plotbox.max_y - self.x_axis_offset)
        dash_locs = np.concatenate(
            (
                [self.plotbox.left],
                self.x_transform(alternate_dash_positions),
                [self.plotbox.right],
            )
        )
        for i, (x1, x2) in enumerate(zip(dash_locs[:-1], dash_locs[1:])):
            x_axis.add(
                dwg.line(
                    (rnd(x1), y),
                    (rnd(x2), y),
                    class_="ax-skip" if i % 2 else "ax-line",
                )
            )
        if tick_positions is not None:
            if tick_labels is None or isinstance(tick_labels, np.ndarray):
                if tick_labels is None:
                    tick_labels = tick_positions
                tick_labels = create_tick_labels(tick_labels)  # format integers

            upper_length = -tick_length_upper if site_muts is None else 0
            ticks_group = x_axis.add(dwg.g(class_="ticks"))
            for pos, lab in itertools.zip_longest(tick_positions, tick_labels):
                tick = ticks_group.add(
                    dwg.g(
                        class_="tick",
                        transform=f"translate({rnd(self.x_transform(pos))} {y})",
                    )
                )
                tick.add(dwg.line((0, rnd(upper_length)), (0, rnd(tick_length_lower))))
                self.add_text_in_group(
                    lab,
                    tick,
                    class_="lab",
                    # place origin at the bottom of the tick plus a single px space
                    pos=(0, tick_length_lower + 1),
                )
        if not self.omit_sites and site_muts is not None:
            # Add sites as vertical lines with overlaid mutations as upper chevrons
            for s_id, mutations in site_muts.items():
                s = self.ts.site(s_id)
                x = self.x_transform(s.position)
                site = x_axis.add(
                    dwg.g(
                        class_=f"site s{s.id + self.offsets.site}",
                        transform=f"translate({rnd(x)} {y})",
                    )
                )
                site.add(dwg.line((0, 0), (0, rnd(-tick_length_upper)), class_="sym"))
                for i, m in enumerate(reversed(mutations)):
                    mutation_class = f"mut m{m.id + self.offsets.mutation}"
                    if m.id in self.mutations_outside_tree:
                        mutation_class += " extra"
                    mut = dwg.g(class_=mutation_class)
                    h = -i * 4 - 1.5
                    w = tick_length_upper / 4
                    # Chevron symbol
                    symbol = mut.add(
                        dwg.polyline(
                            [
                                (rnd(w), rnd(h - 2 * w)),
                                (0, rnd(h)),
                                (rnd(-w), rnd(h - 2 * w)),
                            ],
                            class_="sym",
                        )
                    )
                    if m.id in self.mutation_titles:
                        symbol.set_desc(title=self.mutation_titles[m.id])
                    site.add(mut)

    def draw_y_axis(
        self,
        ticks,  # A dict of pos->label
        upper=None,  # In plot coords
        lower=None,  # In plot coords
        tick_length_left=default_tick_length,
        gridlines=None,
    ):
        if not self.y_axis and not self.y_label:
            return
        if upper is None:
            upper = self.plotbox.top
        if lower is None:
            lower = self.plotbox.bottom
        dwg = self.drawing
        x = rnd(self.y_axis_offset)
        axes = self.get_axes()
        y_axis = axes.add(dwg.g(class_="y-axis"))
        if self.y_label:
            self.add_text_in_group(
                self.y_label,
                y_axis,
                pos=(0, (upper + lower) / 2),
                group_class="title",
                class_="lab",
                text_anchor="middle",
                transform="translate(11) rotate(-90)",
            )
        if self.y_axis:
            y_axis.add(dwg.line((x, rnd(lower)), (x, rnd(upper)), class_="ax-line"))
            ticks_group = y_axis.add(dwg.g(class_="ticks"))
            tick_outside_axis = {}
            for y, label in ticks.items():
                y_pos = self.timescaling.transform(y)
                if y_pos > lower or y_pos < upper:  # nb lower > upper in SVG coords
                    tick_outside_axis[y] = label
                tick = ticks_group.add(
                    dwg.g(class_="tick", transform=f"translate({x} {rnd(y_pos)})")
                )
                if gridlines:
                    tick.add(
                        dwg.line(
                            (0, 0), (rnd(self.plotbox.right - x), 0), class_="grid"
                        )
                    )
                tick.add(dwg.line((0, 0), (rnd(-tick_length_left), 0)))
                self.add_text_in_group(
                    # place the origin at the left of the tickmark plus a single px space
                    label,
                    tick,
                    pos=(rnd(-tick_length_left - 1), 0),
                    class_="lab",
                    text_anchor="end",
                )
            if len(tick_outside_axis) > 0:
                logging.warning(
                    f"Ticks {tick_outside_axis} lie outside the plotted axis"
                )

    def shade_background(
        self,
        breaks,
        tick_length_lower,
        tree_width=None,
        bottom_padding=None,
    ):
        if not self.x_axis:
            return
        if tree_width is None:
            tree_width = self.plotbox.width
        if bottom_padding is None:
            bottom_padding = self.plotbox.pad_bottom
        plot_breaks = self.x_transform(np.array(breaks))
        dwg = self.drawing

        # For tree sequences, we need to add on the background shaded regions
        self.root_groups["background"] = self.dwg_base.add(dwg.g(class_="background"))
        y = self.image_size[1] - self.x_axis_offset - self.plotbox.top
        for i in range(1, len(breaks)):
            break_x = plot_breaks[i]
            prev_break_x = plot_breaks[i - 1]
            tree_x = i * tree_width + self.plotbox.left
            prev_tree_x = (i - 1) * tree_width + self.plotbox.left
            # Shift diagonal lines between tree & axis into the treebox a little
            diag_height = y - (self.image_size[1] - bottom_padding) + self.plotbox.top
            self.root_groups["background"].add(
                # NB: the path below draws straight diagonal lines between the tree boxes
                # and the X axis. An alternative implementation using bezier curves could
                # substitute the following for lines 2 and 4 of the path spec string
                # "l0,{box_h:g} c0,{diag_h} {rdiag_x},0 {rdiag_x},{diag_h} "
                # "c0,-{diag_h} {ldiag_x},0 {ldiag_x},-{diag_h} l0,-{box_h:g}z"
                dwg.path(
                    "M{start_x:g},{top:g} l{box_w:g},0 "  # Top left to top right of tree
                    "l0,{box_h:g} l{rdiag_x:g},{diag_h:g} "  # Down to axis
                    "l0,{tick_h:g} l{ax_x:g},0 l0,-{tick_h:g} "  # Between axis ticks
                    "l{ldiag_x:g},-{diag_h:g} l0,-{box_h:g}z".format(  # Up from axis
                        top=rnd(self.plotbox.top),
                        start_x=rnd(prev_tree_x),
                        box_w=rnd(tree_x - prev_tree_x),
                        box_h=rnd(y - diag_height),
                        rdiag_x=rnd(break_x - tree_x),
                        diag_h=rnd(diag_height),
                        tick_h=rnd(tick_length_lower),
                        ax_x=rnd(prev_break_x - break_x),
                        ldiag_x=rnd(rnd(prev_tree_x) - rnd(prev_break_x)),
                    )
                )
            )

    def x_transform(self, x):
        raise NotImplementedError(
            "No transform func defined for genome pos -> plot coords"
        )


class SvgTreeSequence(SvgAxisPlot):
    """
    A class to draw a tree sequence in SVG format.

    See :meth:`TreeSequence.draw_svg` for a description of usage and parameters.
    """

    def __init__(
        self,
        ts,
        size,
        x_scale,
        time_scale,
        node_labels,
        mutation_labels,
        root_svg_attributes,
        style,
        order,
        force_root_branch,
        symbol_size,
        x_axis,
        y_axis,
        x_label,
        y_label,
        y_ticks,
        x_regions=None,
        y_gridlines=None,
        x_lim=None,
        max_time=None,
        min_time=None,
        node_attrs=None,
        mutation_attrs=None,
        edge_attrs=None,
        node_label_attrs=None,
        mutation_label_attrs=None,
        node_titles=None,
        mutation_titles=None,
        tree_height_scale=None,
        max_tree_height=None,
        max_num_trees=None,
        title=None,
        **kwargs,
    ):
        if max_time is None and max_tree_height is not None:
            max_time = max_tree_height
            # Deprecated in 0.3.6
            warnings.warn(
                "max_tree_height is deprecated; use max_time instead",
                FutureWarning,
                stacklevel=4,
            )
        if time_scale is None and tree_height_scale is not None:
            time_scale = tree_height_scale
            # Deprecated in 0.3.6
            warnings.warn(
                "tree_height_scale is deprecated; use time_scale instead",
                FutureWarning,
                stacklevel=4,
            )
        x_lim = check_x_lim(x_lim, max_x=ts.sequence_length)
        ts, self.tree_status, offsets = clip_ts(ts, x_lim[0], x_lim[1], max_num_trees)

        use_tree = self.tree_status & OMIT == 0
        use_skipped = np.append(np.diff(self.tree_status & OMIT_MIDDLE == 0) == 1, 0)
        num_plotboxes = np.sum(np.logical_or(use_tree, use_skipped))
        if size is None:
            size = (200 * int(num_plotboxes), 200)
        if max_time is None:
            max_time = "ts"
        if min_time is None:
            min_time = "ts"
        # X axis shown by default
        if x_axis is None:
            x_axis = True
        super().__init__(
            ts,
            size,
            root_svg_attributes,
            style,
            svg_class="tree-sequence",
            time_scale=time_scale,
            x_axis=x_axis,
            y_axis=y_axis,
            x_label=x_label,
            y_label=y_label,
            offsets=offsets,
            mutation_titles=mutation_titles,
            **kwargs,
        )
        x_scale = check_x_scale(x_scale)
        order = check_order(order)
        if node_labels is None:
            node_labels = {u: str(u) for u in range(ts.num_nodes)}
        if force_root_branch is None:
            force_root_branch = any(
                any(tree.parent(mut.node) == NULL for mut in tree.mutations())
                for tree, use in zip(ts.trees(), use_tree)
                if use
            )

        # TODO add general padding arguments following matplotlib's terminology.
        self.set_spacing(
            top=0 if title is None else self.line_height, left=20, bottom=10, right=20
        )
        subplot_size = (self.plotbox.width / num_plotboxes, self.plotbox.height)
        subplots = []
        for tree, use, summary in zip(ts.trees(), use_tree, use_skipped):
            if use:
                subplots.append(
                    SvgTree(
                        tree,
                        size=subplot_size,
                        time_scale=time_scale,
                        node_labels=node_labels,
                        mutation_labels=mutation_labels,
                        node_titles=node_titles,
                        mutation_titles=mutation_titles,
                        order=order,
                        force_root_branch=force_root_branch,
                        symbol_size=symbol_size,
                        max_time=max_time,
                        min_time=min_time,
                        node_attrs=node_attrs,
                        mutation_attrs=mutation_attrs,
                        edge_attrs=edge_attrs,
                        node_label_attrs=node_label_attrs,
                        mutation_label_attrs=mutation_label_attrs,
                        offsets=offsets,
                        # Do not plot axes on these subplots
                        **kwargs,  # pass though e.g. debug boxes
                    )
                )
                last_used_index = tree.index
            elif summary:
                subplots.append(
                    SvgSkippedPlot(
                        size=subplot_size, num_skipped=tree.index - last_used_index
                    )
                )
        y = self.plotbox.top
        if title is not None:
            self.add_text_in_group(
                title,
                self.drawing,
                pos=(self.plotbox.max_x / 2, 0),
                dominant_baseline="hanging",
                group_class="title",
                text_anchor="middle",
            )
        self.tree_plotbox = subplots[0].plotbox
        tree_is_used, breaks, skipbreaks = self.find_used_trees()
        self.draw_x_axis(
            x_scale,
            tree_is_used,
            breaks,
            skipbreaks,
            tick_length_lower=self.default_tick_length,  # TODO - parameterize
            tick_length_upper=self.default_tick_length_site,  # TODO - parameterize
            x_regions=x_regions,
        )
        y_low = self.tree_plotbox.bottom
        if y_axis is not None:
            tscales = {s.timescaling for s in subplots if s.timescaling}
            if len(tscales) > 1:
                raise ValueError(
                    "Can't draw a tree sequence Y axis if trees vary in timescale"
                )
            self.timescaling = tscales.pop()
            y_low = self.timescaling.transform(self.timescaling.min_time)
            if y_ticks is None:
                used_nodes = edge_and_sample_nodes(ts, breaks[skipbreaks])
                y_ticks = np.unique(ts.nodes_time[used_nodes])
                if self.time_scale == "rank":
                    # Ticks labelled by time not rank
                    y_ticks = dict(enumerate(y_ticks))

        self.draw_y_axis(
            ticks=check_y_ticks(y_ticks),
            upper=self.tree_plotbox.top,
            lower=y_low,
            tick_length_left=self.default_tick_length,
            gridlines=y_gridlines,
        )

        subplot_x = self.plotbox.left
        container = self.get_plotbox()  # Top-level TS plotbox contains all trees
        container["class"] = container["class"] + " trees"
        for subplot in subplots:
            svg_subplot = container.add(
                self.drawing.g(
                    class_=subplot.svg_class,
                    transform=f"translate({rnd(subplot_x)} {y})",
                )
            )
            for svg_items in subplot.root_groups.values():
                svg_subplot.add(svg_items)
            subplot_x += subplot.image_size[0]

    def find_used_trees(self):
        """
        Return a boolean array of which trees are actually plotted,
        a list of which breakpoints are used to transition between plotted trees,
        and a 2 x n array (often n=0) of indexes into these breakpoints delimiting
        the regions that should be plotted as "skipped"
        """
        tree_is_used = (self.tree_status & OMIT) != OMIT
        break_used_as_tree_left = np.append(tree_is_used, False)
        break_used_as_tree_right = np.insert(tree_is_used, 0, False)
        break_used = np.logical_or(break_used_as_tree_left, break_used_as_tree_right)
        all_breaks = self.ts.breakpoints(True)
        used_breaks = all_breaks[break_used]
        mark_skip_transitions = np.concatenate(
            ([False], np.diff(self.tree_status & OMIT_MIDDLE) != 0, [False])
        )
        skipregion_indexes = np.where(mark_skip_transitions[break_used])[0]
        assert len(skipregion_indexes) % 2 == 0  # all skipped regions have start, end
        return tree_is_used, used_breaks, skipregion_indexes.reshape((-1, 2))

    def draw_x_axis(
        self,
        x_scale,
        tree_is_used,
        breaks,
        skipbreaks,
        x_regions,
        tick_length_lower=SvgAxisPlot.default_tick_length,
        tick_length_upper=SvgAxisPlot.default_tick_length_site,
    ):
        """
        Add extra functionality to the original draw_x_axis method in SvgAxisPlot,
        to account for the background shading that is displayed in a tree sequence
        and in case trees are omitted from the middle of the tree sequence
        """
        if not self.x_axis and not self.x_label:
            return
        if x_scale == "physical":
            # In a tree sequence plot, the x_transform is used for the ticks, background
            # shading positions, and sites along the x-axis. Each tree will have its own
            # separate x_transform function for node positions within the tree.

            # For a plot with a break on the x-axis (representing "skipped" trees), the
            # x_transform is a piecewise function. We need to identify the breakpoints
            # where the x-scale transitions from the standard scale to the scale(s) used
            # within a skipped region

            skipregion_plot_width = self.tree_plotbox.width
            skipregion_span = np.diff(breaks[skipbreaks]).T[0]
            std_scale = (
                self.plotbox.width - skipregion_plot_width * len(skipregion_span)
            ) / (breaks[-1] - breaks[0] - np.sum(skipregion_span))
            skipregion_pos = breaks[skipbreaks].flatten()
            genome_pos = np.concatenate(([breaks[0]], skipregion_pos, [breaks[-1]]))
            plot_step = np.full(len(genome_pos) - 1, skipregion_plot_width)
            plot_step[::2] = std_scale * np.diff(genome_pos)[::2]
            plot_pos = np.cumsum(np.insert(plot_step, 0, self.plotbox.left))
            # Convert to slope + intercept form
            slope = np.diff(plot_pos) / np.diff(genome_pos)
            intercept = plot_pos[1:] - slope * genome_pos[1:]
            self.x_transform = lambda y: (
                y * slope[np.searchsorted(skipregion_pos, y)]
                + intercept[np.searchsorted(skipregion_pos, y)]
            )
            tick_positions = breaks
            site_muts = {
                s.id: s.mutations
                for tree, use in zip(self.ts.trees(), tree_is_used)
                for s in tree.sites()
                if use
            }

            self.shade_background(
                breaks,
                tick_length_lower,
                self.tree_plotbox.max_x,
                self.plotbox.pad_bottom + self.tree_plotbox.pad_bottom,
            )
        else:
            # For a treewise plot, the only time the x_transform is used is to apply
            # to tick positions, so simply use positions 0..num_used_breaks for the
            # positions, and a simple transform
            self.x_transform = (
                lambda x: self.plotbox.left + x / (len(breaks) - 1) * self.plotbox.width
            )
            tick_positions = np.arange(len(breaks))

            site_muts = None  # It doesn't make sense to plot sites for "treewise" plots
            tick_length_upper = None  # No sites plotted, so use the default upper tick
            if x_regions is not None and len(x_regions) > 0:
                raise ValueError("x_regions are not supported for treewise plots")

            # NB: no background shading needed if x_scale is "treewise"

            skipregion_pos = skipbreaks.flatten()

        first_tick = 1 if np.any(self.tree_status[tree_is_used] & LEFT_CLIP) else 0
        last_tick = -1 if np.any(self.tree_status[tree_is_used] & RIGHT_CLIP) else None

        super().draw_x_axis(
            tick_positions=tick_positions[first_tick:last_tick],
            tick_labels=breaks[first_tick:last_tick],
            tick_length_lower=tick_length_lower,
            tick_length_upper=tick_length_upper,
            site_muts=site_muts,
            alternate_dash_positions=skipregion_pos,
            x_regions=x_regions,
        )


class SvgTree(SvgAxisPlot):
    """
    A class to draw a tree in SVG format.

    See :meth:`Tree.draw_svg` for a description of usage and frequently used parameters.
    """

    PolytomyLine = collections.namedtuple(
        "PolytomyLine", "num_branches, num_samples, line_pos"
    )

    def __init__(
        self,
        tree,
        size=None,
        max_time=None,
        min_time=None,
        max_tree_height=None,
        node_labels=None,
        mutation_labels=None,
        node_titles=None,
        mutation_titles=None,
        root_svg_attributes=None,
        style=None,
        order=None,
        force_root_branch=None,
        symbol_size=None,
        x_axis=None,
        y_axis=None,
        x_label=None,
        y_label=None,
        title=None,
        x_regions=None,
        y_ticks=None,
        y_gridlines=None,
        all_edge_mutations=None,
        time_scale=None,
        tree_height_scale=None,
        node_attrs=None,
        mutation_attrs=None,
        edge_attrs=None,
        node_label_attrs=None,
        mutation_label_attrs=None,
        offsets=None,
        omit_sites=None,
        pack_untracked_polytomies=None,
        **kwargs,
    ):
        if max_time is None and max_tree_height is not None:
            max_time = max_tree_height
            # Deprecated in 0.3.6
            warnings.warn(
                "max_tree_height is deprecated; use max_time instead",
                FutureWarning,
                stacklevel=4,
            )
        if time_scale is None and tree_height_scale is not None:
            time_scale = tree_height_scale
            # Deprecated in 0.3.6
            warnings.warn(
                "tree_height_scale is deprecated; use time_scale instead",
                FutureWarning,
                stacklevel=4,
            )
        if size is None:
            size = (200, 200)
        if symbol_size is None:
            symbol_size = 6
        self.symbol_size = symbol_size
        self.pack_untracked_polytomies = pack_untracked_polytomies
        ts = tree.tree_sequence
        tree_index = tree.index
        if offsets is not None:
            tree_index += offsets.tree
        super().__init__(
            ts,
            size,
            root_svg_attributes,
            style,
            svg_class=f"tree t{tree_index}",
            time_scale=time_scale,
            x_axis=x_axis,
            y_axis=y_axis,
            x_label=x_label,
            y_label=y_label,
            offsets=offsets,
            omit_sites=omit_sites,
            **kwargs,
        )
        self.tree = tree
        if order is None or isinstance(order, str):
            # Can't use the Tree.postorder array as we need minlex
            self.postorder_nodes = list(tree.nodes(order=check_order(order)))
        else:
            # Currently undocumented feature: we can pass a (postorder) list
            # of nodes to plot, which allows us to draw a subset of nodes, or
            # stop traversing certain subtrees
            self.postorder_nodes = order

        # Create some instance variables for later use in plotting
        self.node_mutations = collections.defaultdict(list)
        self.edge_attrs = {}
        self.node_attrs = {}
        self.node_label_attrs = {}
        self.mutation_attrs = {}
        self.mutation_label_attrs = {}
        self.node_titles = {} if node_titles is None else node_titles
        self.mutation_titles = {} if mutation_titles is None else mutation_titles
        self.mutations_over_roots = False
        # mutations collected per node
        nodes = set(tree.nodes())
        unplotted = []
        if not omit_sites:
            for site in tree.sites():
                for mutation in site.mutations:
                    if mutation.node in nodes:
                        self.node_mutations[mutation.node].append(mutation)
                        if tree.parent(mutation.node) == NULL:
                            self.mutations_over_roots = True
                    else:
                        unplotted.append(mutation.id + self.offsets.mutation)
        if len(unplotted) > 0:
            warnings.warn(
                f"Mutations {unplotted} are above nodes which are not present in the "
                "displayed tree, so are not plotted on the topology.",
                UserWarning,
                stacklevel=2,
            )
        self.left_extent = tree.interval.left
        self.right_extent = tree.interval.right
        if not omit_sites and all_edge_mutations:
            tree_left = tree.interval.left
            tree_right = tree.interval.right
            edge_left = ts.tables.edges.left
            edge_right = ts.tables.edges.right
            node_edges = tree.edge_array
            # whittle mutations down so we only need look at those above the tree nodes
            mut_t = ts.tables.mutations
            focal_mutations = np.isin(mut_t.node, np.fromiter(nodes, mut_t.node.dtype))
            mutation_nodes = mut_t.node[focal_mutations]
            mutation_positions = ts.tables.sites.position[mut_t.site][focal_mutations]
            mutation_ids = np.arange(ts.num_mutations, dtype=int)[focal_mutations]
            for m_id, node, pos in zip(
                mutation_ids, mutation_nodes, mutation_positions
            ):
                curr_edge = node_edges[node]
                if curr_edge >= 0:
                    if (
                        edge_left[curr_edge] <= pos < tree_left
                    ):  # Mutation on this edge but to left of plotted tree
                        self.node_mutations[node].append(ts.mutation(m_id))
                        self.mutations_outside_tree.add(m_id)
                        self.left_extent = min(self.left_extent, pos)
                    elif (
                        tree_right <= pos < edge_right[curr_edge]
                    ):  # Mutation on this edge but to right of plotted tree
                        self.node_mutations[node].append(ts.mutation(m_id))
                        self.mutations_outside_tree.add(m_id)
                        self.right_extent = max(self.right_extent, pos)
            if self.right_extent != tree.interval.right:
                # Use nextafter so extent of plotting incorporates the mutation
                self.right_extent = np.nextafter(
                    self.right_extent, self.right_extent + 1
                )
        # attributes for symbols
        half_symbol_size = f"{rnd(symbol_size / 2):g}"
        symbol_size = f"{rnd(symbol_size):g}"
        for u in tree.nodes():
            self.edge_attrs[u] = {}
            if edge_attrs is not None and u in edge_attrs:
                self.edge_attrs[u].update(edge_attrs[u])
            if tree.is_sample(u):
                # a square: set bespoke svgwrite params
                self.node_attrs[u] = {
                    "size": (symbol_size,) * 2,
                    "insert": ("-" + half_symbol_size,) * 2,
                }
            else:
                # a circle: set bespoke svgwrite param `centre` and default radius
                self.node_attrs[u] = {"center": (0, 0), "r": half_symbol_size}
            if node_attrs is not None and u in node_attrs:
                self.node_attrs[u].update(node_attrs[u])
            add_class(self.node_attrs[u], "sym")  # class 'sym' for symbol
            label = ""
            if node_labels is None:
                label = str(u)
            elif u in node_labels:
                label = str(node_labels[u])
            self.node_label_attrs[u] = {"text": label}
            add_class(self.node_label_attrs[u], "lab")  # class 'lab' for label
            if node_label_attrs is not None and u in node_label_attrs:
                self.node_label_attrs[u].update(node_label_attrs[u])
        for _, mutations in self.node_mutations.items():
            for mutation in mutations:
                m = mutation.id + self.offsets.mutation
                # We need to offset the mutation symbol so that it's centred
                self.mutation_attrs[m] = {
                    "d": "M -{0},-{0} l {1},{1} M -{0},{0} l {1},-{1}".format(
                        half_symbol_size, symbol_size
                    )
                }
                if mutation_attrs is not None and m in mutation_attrs:
                    self.mutation_attrs[m].update(mutation_attrs[m])
                add_class(self.mutation_attrs[m], "sym")  # class 'sym' for symbol
                label = ""
                if mutation_labels is None:
                    label = str(m)
                elif m in mutation_labels:
                    label = str(mutation_labels[m])
                self.mutation_label_attrs[m] = {"text": label}
                if mutation_label_attrs is not None and m in mutation_label_attrs:
                    self.mutation_label_attrs[m].update(mutation_label_attrs[m])
                add_class(self.mutation_label_attrs[m], "lab")

        self.set_spacing(
            top=10 if title is None else 10 + self.line_height,
            left=20,
            bottom=15,
            right=20,
        )
        if title is not None:
            self.add_text_in_group(
                title,
                self.drawing,
                pos=(self.plotbox.max_x / 2, 0),
                dominant_baseline="hanging",
                group_class="title",
                text_anchor="middle",
            )

        self.assign_x_coordinates()
        self.assign_y_coordinates(max_time, min_time, force_root_branch)
        tick_length_lower = self.default_tick_length  # TODO - parameterize
        tick_length_upper = self.default_tick_length_site  # TODO - parameterize
        if all_edge_mutations:
            self.shade_background(tree.interval, tick_length_lower)

        first_site, last_site = np.searchsorted(
            self.ts.tables.sites.position, [self.left_extent, self.right_extent]
        )
        site_muts = {site_id: [] for site_id in range(first_site, last_site)}
        # Only use mutations plotted on the tree (not necessarily all at the site)
        for muts in self.node_mutations.values():
            for mut in muts:
                site_muts[mut.site].append(mut)

        self.draw_x_axis(
            tick_positions=np.array(tree.interval),
            tick_length_lower=tick_length_lower,
            tick_length_upper=tick_length_upper,
            site_muts=site_muts,
            x_regions=x_regions,
        )
        if y_ticks is None:
            y_ticks = {h: ts.node(u).time for u, h in sorted(self.node_height.items())}

        self.draw_y_axis(
            ticks=check_y_ticks(y_ticks),
            lower=self.timescaling.transform(self.timescaling.min_time),
            tick_length_left=self.default_tick_length,
            gridlines=y_gridlines,
        )
        self.draw_tree()

    def process_mutations_over_node(self, u, low_bound, high_bound, ignore_times=False):
        """
        Sort the self.node_mutations array for a given node ``u`` in reverse time order.
        The main complication is with UNKNOWN_TIME values: we replace these with times
        spaced between the low & high bounds (this is always done if ignore_times=True).
        We do not currently allow a mix of known & unknown mutation times in a tree
        sequence, which makes the logic easy. If we were to allow it, more complex
        logic can be neatly encapsulated in this method.
        """
        mutations = self.node_mutations[u]
        time_unknown = [util.is_unknown_time(m.time) for m in mutations]
        if all(time_unknown) or ignore_times is True:
            # sort by site then within site by parent: will end up with oldest first
            mutations.sort(key=operator.attrgetter("site", "parent"))
            diff = high_bound - low_bound
            for i in range(len(mutations)):
                mutations[i].time = high_bound - diff * (i + 1) / (len(mutations) + 1)
        else:
            assert not any(time_unknown)
            mutations.sort(key=operator.attrgetter("time"), reverse=True)

    def assign_y_coordinates(
        self,
        max_time,
        min_time,
        force_root_branch,
        bottom_space=SvgAxisPlot.line_height,
        top_space=SvgAxisPlot.line_height,
    ):
        """
        Create a self.node_height dict, a self.timescaling instance and
        self.min_root_branch_plot_length for use in plotting. Allow extra space within
        the plotbox, at the bottom for leaf labels, and  (potentially, if no root
        branches are plotted) above the topmost root node for root labels.
        """
        max_time = check_max_time(max_time, self.time_scale != "rank")
        min_time = check_min_time(min_time, self.time_scale != "rank")
        node_time = self.ts.nodes_time
        mut_time = self.ts.mutations_time
        root_branch_len = 0
        if self.time_scale == "rank":
            t = np.zeros_like(node_time)
            if max_time == "tree":
                # We only rank the times within the tree in this case.
                for u in self.node_x_coord.keys():
                    t[u] = node_time[u]
            else:
                # only rank the nodes that are actually referenced in the edge table
                # (non-referenced nodes could occur if the user specifies x_lim values)
                # However, we do include nodes in trees that have been skipped
                use_time = edge_and_sample_nodes(self.ts)
                t[use_time] = node_time[use_time]
            node_time = t
            times = np.unique(node_time[node_time <= self.ts.max_root_time])
            max_node_height = len(times)
            depth = {t: j for j, t in enumerate(times)}
            if self.mutations_over_roots or force_root_branch:
                root_branch_len = 1  # Will get scaled later
            max_time = max(depth.values()) + root_branch_len
            if min_time in (None, "tree", "ts"):
                assert min(depth.values()) == 0
                min_time = 0
            # In pathological cases, all the nodes are at the same time
            if max_time == min_time:
                max_time = min_time + 1
            self.node_height = {
                u: depth[node_time[u]] for u in self.node_x_coord.keys()
            }
            for u in self.node_mutations.keys():
                if u in self.node_height:
                    parent = self.tree.parent(u)
                    if parent == NULL:
                        top = self.node_height[u] + root_branch_len
                    else:
                        top = depth[node_time[parent]]
                    self.process_mutations_over_node(
                        u, self.node_height[u], top, ignore_times=True
                    )
        else:
            assert self.time_scale in ["time", "log_time"]
            self.node_height = {u: node_time[u] for u in self.node_x_coord.keys()}
            if max_time == "tree":
                max_node_height = max(self.node_height.values())
                max_mut_height = np.nanmax(
                    [0] + [mut.time for m in self.node_mutations.values() for mut in m]
                )
                max_time = max(max_node_height, max_mut_height)  # Reuse variable
            elif max_time == "ts":
                max_node_height = self.ts.max_root_time
                max_mut_height = np.nanmax(np.append(mut_time, 0))
                max_time = max(max_node_height, max_mut_height)  # Reuse variable
            if min_time == "tree":
                min_time = min(self.node_height.values())
                # don't need to check mutation times, as they must be above a node
            elif min_time == "ts":
                min_time = np.min(self.ts.nodes_time[edge_and_sample_nodes(self.ts)])
            # In pathological cases, all the nodes are at the same time
            if min_time == max_time:
                max_time = min_time + 1
            if self.mutations_over_roots or force_root_branch:
                # Define a minimum root branch length, after transformation if necessary
                if self.time_scale != "log_time":
                    root_branch_len = (max_time - min_time) * self.root_branch_fraction
                else:
                    max_plot_y = np.log(max_time + 1)
                    diff_plot_y = max_plot_y - np.log(min_time + 1)
                    root_plot_y = max_plot_y + diff_plot_y * self.root_branch_fraction
                    root_branch_len = np.exp(root_plot_y) - 1 - max_time
                # If necessary, allow for this extra branch in max_time
                if max_node_height + root_branch_len > max_time:
                    max_time = max_node_height + root_branch_len
            for u in self.node_mutations.keys():
                if u in self.node_height:
                    parent = self.tree.parent(u)
                    if parent == NULL:
                        # This is a root: if muts have no times we specify an upper time
                        top = self.node_height[u] + root_branch_len
                    else:
                        top = node_time[parent]
                    self.process_mutations_over_node(u, self.node_height[u], top)

        assert float(max_time) == max_time
        assert float(min_time) == min_time
        # Add extra space above the top and below the bottom of the tree to keep the
        # node labels within the plotbox (but top label space not needed if the
        # existence of a root branch pushes the whole tree + labels downwards anyway)
        top_space = 0 if root_branch_len > 0 else top_space
        self.timescaling = Timescaling(
            max_time=max_time,
            min_time=min_time,
            plot_min=self.plotbox.height + self.plotbox.top - bottom_space,
            plot_range=self.plotbox.height - top_space - bottom_space,
            use_log_transform=(self.time_scale == "log_time"),
        )

        # Calculate default root branch length to use (in plot coords). This is a
        # minimum, as branches with deep root mutations could be longer
        self.min_root_branch_plot_length = self.timescaling.transform(
            self.timescaling.max_time
        ) - self.timescaling.transform(self.timescaling.max_time + root_branch_len)

    def assign_x_coordinates(self):
        # Set up transformation for genome positions
        self.x_transform = lambda x: (
            (x - self.left_extent)
            / (self.right_extent - self.left_extent)
            * self.plotbox.width
            + self.plotbox.left
        )
        # Set up x positions for nodes
        node_xpos = {}
        untracked_children = collections.defaultdict(list)
        self.extra_line = {}  # To store a dotted line to represent polytomies
        leaf_x = 0  # First leaf starts at x=1, to give some space between Y axis & leaf
        tree = self.tree
        prev = tree.virtual_root
        for u in self.postorder_nodes:
            parent = tree.parent(u)
            if parent == prev:
                raise ValueError("Nodes must be passed in postorder to Tree.draw_svg()")
            is_tip = tree.parent(prev) != u
            if is_tip:
                if self.pack_untracked_polytomies and tree.num_tracked_samples(u) == 0:
                    untracked_children[parent].append(u)
                else:
                    leaf_x += 1
                    node_xpos[u] = leaf_x
            else:
                # Concatenate all the untracked children
                num_untracked_children = len(untracked_children[u])
                child_x = [node_xpos[c] for c in tree.children(u) if c in node_xpos]
                if num_untracked_children > 0:
                    if num_untracked_children <= 1:
                        # If only a single non-focal lineage, we might as well show it
                        for child in untracked_children[u]:
                            leaf_x += 1
                            node_xpos[child] = leaf_x
                            child_x.append(leaf_x)
                    else:
                        # Otherwise show a horizontal line with the number of lineages
                        # Extra length of line is equal to log of the polytomy size
                        self.extra_line[u] = self.PolytomyLine(
                            num_untracked_children,
                            sum(tree.num_samples(v) for v in untracked_children[u]),
                            [leaf_x, leaf_x + 1 + np.log(num_untracked_children)],
                        )
                        child_x.append(leaf_x + 1)
                        leaf_x = self.extra_line[u].line_pos[1]
                assert len(child_x) != 0  # Must have prev hit somethng defined as a tip
                if len(child_x) == 1:
                    node_xpos[u] = child_x[0]
                else:
                    a = min(child_x)
                    b = max(child_x)
                    node_xpos[u] = a + (b - a) / 2
            prev = u
        # Now rescale to the plot width: leaf_x is the maximum value of the last leaf
        if len(node_xpos) > 0:
            scale = self.plotbox.width / leaf_x
            lft = self.plotbox.left - scale / 2
        self.node_x_coord = {k: lft + v * scale for k, v in node_xpos.items()}
        for v in self.extra_line.values():
            for i in range(len(v.line_pos)):
                v.line_pos[i] = lft + v.line_pos[i] * scale

    def info_classes(self, focal_node_id):
        """
        For a focal node id, return a set of classes that encode this useful information:
            "a<X>" or "root": where <X> == id of immediate ancestor (parent) node
            "i<I>":           where <I> == individual id
            "p<P>":           where <P> == population id
            "n<Y>":           where <Y> == focal node id
            "m<A>":           where <A> == mutation id
            "s<B>":           where <B> == site id of all mutations
            "c<N>" or "leaf": where <N> == number of direct children of this node
        """
        # Add a new group for each node, and give it classes for css targetting
        focal_node = self.ts.node(focal_node_id)
        classes = set()
        classes.add(f"node n{focal_node_id}")
        if focal_node.individual != NULL:
            classes.add(f"i{focal_node.individual}")
        if focal_node.population != NULL:
            classes.add(f"p{focal_node.population}")
        v = self.tree.parent(focal_node_id)
        if v == NULL:
            classes.add("root")
        else:
            classes.add(f"a{v}")
        if self.tree.is_sample(focal_node_id):
            classes.add("sample")
        if self.tree.is_leaf(focal_node_id):
            classes.add("leaf")
        else:
            classes.add(f"c{self.tree.num_children(focal_node_id)}")
        for mutation in self.node_mutations[focal_node_id]:
            # Adding mutations and sites above this node allows identification
            # of the tree under any specific mutation
            classes.add(f"m{mutation.id + self.offsets.mutation}")
            classes.add(f"s{mutation.site+ self.offsets.site}")
        return sorted(classes)

    def text_transform(self, position, dy=0):
        line_h = self.text_height
        sym_sz = self.symbol_size
        transforms = {
            "below": f"translate(0 {rnd(line_h - sym_sz / 2 + dy)})",
            "above": f"translate(0 {rnd(-(line_h - sym_sz / 2) + dy)})",
            "above_left": f"translate({rnd(-sym_sz / 2)} {rnd(-line_h / 2 + dy)})",
            "above_right": f"translate({rnd(sym_sz / 2)} {-rnd(line_h / 2 + dy)})",
            "left": f"translate({-rnd(2 + sym_sz / 2)} {rnd(dy)})",
            "right": f"translate({rnd(2 + sym_sz / 2)} {rnd(dy)})",
        }
        return transforms[position]

    def draw_tree(self):
        # Note: the displayed tree may not be the same as self.tree, e.g. if the nodes
        # have been collapsed, or a subtree is being displayed. The node_x_coord
        # dictionary keys gives the nodes of the displayed tree, in postorder.
        NodeDrawInfo = collections.namedtuple("NodeDrawInfo", ["pos", "is_tip"])
        dwg = self.drawing
        tree = self.tree
        left_child = get_left_child(tree, self.postorder_nodes)
        parent_array = tree.parent_array

        node_info = {}
        roots = []  # Roots of the displated tree
        prev = tree.virtual_root
        for u, x in self.node_x_coord.items():  # Node ids `u` returned in postorder
            node_info[u] = NodeDrawInfo(
                pos=np.array([x, self.timescaling.transform(self.node_height[u])]),
                # Detect if this is a "tip" in the displayed tree, even if
                # it is not a leaf in the original tree, by looking at the prev parent
                is_tip=(parent_array[prev] != u),
            )
            prev = u
            if parent_array[u] not in self.node_x_coord:
                roots.append(u)
        # Iterate over displayed nodes, adding groups to reflect the tree hierarchy
        stack = []
        for u in roots:
            x, y = node_info[u].pos
            grp = dwg.g(
                class_=" ".join(self.info_classes(u)),
                transform=f"translate({rnd(x)} {rnd(y)})",
            )
            stack.append((u, self.get_plotbox().add(grp)))

        # Preorder traversal, so we can create nested groups
        while len(stack) > 0:
            u, curr_svg_group = stack.pop()
            pu, is_tip = node_info[u]
            for focal in tree.children(u):
                if focal not in node_info:
                    continue
                fx, fy = node_info[focal].pos - pu
                new_svg_group = curr_svg_group.add(
                    dwg.g(
                        class_=" ".join(self.info_classes(focal)),
                        transform=f"translate({rnd(fx)} {rnd(fy)})",
                    )
                )
                stack.append((focal, new_svg_group))

            o = (0, 0)
            v = parent_array[u]

            # Add polytomy line if necessary
            if u in self.extra_line:
                info = self.extra_line[u]
                x2 = info.line_pos[1] - pu[0]
                poly = dwg.g(class_="polytomy")
                poly.add(
                    dwg.line(
                        start=(0, 0),
                        end=(x2, 0),
                    )
                )
                poly.add(
                    dwg.text(
                        f"+{info.num_samples}/{bold_integer(info.num_branches)}",
                        font_style="italic",
                        x=[rnd(x2)],
                        dy=[rnd(-self.text_height / 10)],  # make the plus sign line up
                        text_anchor="end",
                    )
                )
                curr_svg_group.add(poly)

            # Add edge above node first => on layer underneath anything else
            draw_edge_above_node = False
            try:
                dx, dy = node_info[v].pos - pu
                draw_edge_above_node = True
            except KeyError:
                # Must be a root
                root_branch_l = self.min_root_branch_plot_length
                if root_branch_l > 0:
                    if len(self.node_mutations[u]) > 0:
                        mtop = self.timescaling.transform(
                            self.node_mutations[u][0].time
                        )
                        root_branch_l = max(root_branch_l, pu[1] - mtop)
                    dx, dy = 0, -root_branch_l
                    draw_edge_above_node = True
            if draw_edge_above_node:
                add_class(self.edge_attrs[u], "edge")
                path = dwg.path(
                    [("M", o), ("V", rnd(dy)), ("H", rnd(dx))], **self.edge_attrs[u]
                )
                curr_svg_group.add(path)

            # Add mutation symbols + labels
            for mutation in self.node_mutations[u]:
                # TODO get rid of these manual positioning tweaks and add them
                # as offsets the user can access via a transform or something.
                dy = self.timescaling.transform(mutation.time) - pu[1]
                mutation_id = mutation.id + self.offsets.mutation
                mutation_class = (
                    f"mut m{mutation_id} " f"s{mutation.site+ self.offsets.site}"
                )
                # Use the real mutation ID here, since we are referencing into the ts
                if util.is_unknown_time(self.ts.mutation(mutation.id).time):
                    mutation_class += " unknown_time"
                if mutation_id in self.mutations_outside_tree:
                    mutation_class += " extra"
                mut_group = curr_svg_group.add(
                    dwg.g(class_=mutation_class, transform=f"translate(0 {rnd(dy)})")
                )
                # A line from the mutation to the node below, normally hidden, but
                # revealable if we want to flag the path below a mutation
                mut_group.add(dwg.line(end=(0, -rnd(dy))))
                # Symbols
                symbol = mut_group.add(dwg.path(**self.mutation_attrs[mutation_id]))
                if mutation_id in self.mutation_titles:
                    symbol.set_desc(title=self.mutation_titles[mutation_id])
                # Labels
                if u == left_child[parent_array[u]]:
                    mut_label_class = "lft"
                    transform = self.text_transform("left")
                else:
                    mut_label_class = "rgt"
                    transform = self.text_transform("right")
                add_class(self.mutation_label_attrs[mutation_id], mut_label_class)
                self.mutation_label_attrs[mutation_id]["transform"] = transform
                mut_group.add(dwg.text(**self.mutation_label_attrs[mutation_id]))

            # Add node symbol + label (visually above the edge subtending this node)
            # -> symbols
            if tree.is_sample(u):
                symbol = curr_svg_group.add(dwg.rect(**self.node_attrs[u]))
            else:
                symbol = curr_svg_group.add(dwg.circle(**self.node_attrs[u]))
            multi_samples = None
            if (
                is_tip and tree.num_samples(u) > 1
            ):  # Multi-sample tip => trapezium shape
                multi_samples = tree.num_samples(u)
                trapezium_attrs = self.node_attrs[u].copy()
                # Remove the shape-styling attributes
                for unwanted_attr in ("size", "insert", "center", "r"):
                    trapezium_attrs.pop(unwanted_attr, None)
                trapezium_attrs["points"] = [  # add a trapezium shape below the symbol
                    (self.symbol_size / 2, 0),
                    (self.symbol_size, self.symbol_size),
                    (-self.symbol_size, self.symbol_size),
                    (-self.symbol_size / 2, 0),
                ]
                add_class(trapezium_attrs, "multi")
                curr_svg_group.add(dwg.polygon(**trapezium_attrs))
            if u in self.node_titles:
                symbol.set_desc(title=self.node_titles[u])
            # -> labels
            node_lab_attr = self.node_label_attrs[u]
            if is_tip and multi_samples is None:
                node_lab_attr["transform"] = self.text_transform("below")
            elif u in roots and self.min_root_branch_plot_length == 0:
                node_lab_attr["transform"] = self.text_transform("above")
            else:
                if multi_samples is not None:
                    curr_svg_group.add(
                        dwg.text(
                            text=f"+{multi_samples}",
                            transform=self.text_transform("below", dy=1),
                            font_style="italic",
                            class_="lab summary",
                        )
                    )
                if u == left_child[tree.parent(u)]:
                    add_class(node_lab_attr, "lft")
                    node_lab_attr["transform"] = self.text_transform("above_left")
                else:
                    add_class(node_lab_attr, "rgt")
                    node_lab_attr["transform"] = self.text_transform("above_right")
            curr_svg_group.add(dwg.text(**node_lab_attr))


class TextTreeSequence:
    """
    Draw a tree sequence as horizontal line of trees.
    """

    def __init__(
        self,
        ts,
        node_labels=None,
        use_ascii=False,
        time_label_format=None,
        position_label_format=None,
        order=None,
    ):
        self.ts = ts

        time_label_format = "{:.2f}" if time_label_format is None else time_label_format
        tick_labels = ts.breakpoints(as_array=True)
        if position_label_format is None:
            position_scale_labels = create_tick_labels(tick_labels)
        else:
            position_scale_labels = [
                position_label_format.format(x) for x in tick_labels
            ]

        time = ts.tables.nodes.time
        time_scale_labels = [
            time_label_format.format(time[u]) for u in range(ts.num_nodes)
        ]

        trees = [
            VerticalTextTree(
                tree,
                max_time="ts",
                node_labels=node_labels,
                use_ascii=use_ascii,
                order=order,
            )
            for tree in self.ts.trees()
        ]

        self.height = 1 + max(tree.height for tree in trees)
        self.width = sum(tree.width + 2 for tree in trees) - 1
        max_time_scale_label_len = max(map(len, time_scale_labels))
        self.width += 3 + max_time_scale_label_len + len(position_scale_labels[-1]) // 2

        self.canvas = np.zeros((self.height, self.width), dtype=str)
        self.canvas[:] = " "

        vertical_sep = "|" if use_ascii else "┊"
        x = 0
        time_position = trees[0].time_position
        for u, label in enumerate(map(to_np_unicode, time_scale_labels)):
            y = time_position[u]
            self.canvas[y, 0 : label.shape[0]] = label
        self.canvas[:, max_time_scale_label_len] = vertical_sep
        x = 2 + max_time_scale_label_len

        for j, tree in enumerate(trees):
            pos_label = to_np_unicode(position_scale_labels[j])
            k = len(pos_label)
            label_x = max(x - k // 2 - 2, 0)
            self.canvas[-1, label_x : label_x + k] = pos_label
            h, w = tree.canvas.shape
            self.canvas[-h - 1 : -1, x : x + w - 1] = tree.canvas[:, :-1]
            x += w
            self.canvas[:, x] = vertical_sep
            x += 2

        pos_label = to_np_unicode(position_scale_labels[-1])
        k = len(pos_label)
        label_x = max(x - k // 2 - 2, 0)
        self.canvas[-1, label_x : label_x + k] = pos_label
        self.canvas[:, -1] = "\n"

    def __str__(self):
        return "".join(self.canvas.reshape(self.width * self.height))


def to_np_unicode(string):
    """
    Converts the specified string to a numpy unicode array.
    """
    # TODO: what's the clean of doing this with numpy?
    # It really wants to create a zero-d Un array here
    # which breaks the assignment below and we end up
    # with n copies of the first char.
    n = len(string)
    np_string = np.zeros(n, dtype="U")
    for j in range(n):
        np_string[j] = string[j]
    return np_string


def get_left_neighbour(tree, traversal_order):
    """
    Returns the left-most neighbour of each node in the tree according to the
    specified traversal order. The left neighbour is the closest node in terms
    of path distance to the left of a given node.
    """
    # The traversal order will define the order of children and roots.
    # Root order is defined by this traversal, and the roots are
    # the children of -1
    children = collections.defaultdict(list)
    for u in tree.nodes(order=traversal_order):
        children[tree.parent(u)].append(u)

    left_neighbour = np.full(tree.tree_sequence.num_nodes + 1, NULL, dtype=int)

    def find_neighbours(u, neighbour):
        left_neighbour[u] = neighbour
        for v in children[u]:
            find_neighbours(v, neighbour)
            neighbour = v

    # The children of -1 are the roots and the neighbour of all left-most
    # nodes in the tree is also -1 (NULL)
    find_neighbours(-1, -1)

    return left_neighbour[:-1]


def get_left_child(tree, postorder_nodes):
    """
    Returns the left-most child of each node in the tree according to the
    traversal order listed in postorder_nodes. If a node has no children or
    NULL is passed in, return NULL.
    """
    left_child = np.full(tree.tree_sequence.num_nodes + 1, NULL, dtype=int)
    for u in postorder_nodes:
        parent = tree.parent(u)
        if parent != NULL and left_child[parent] == NULL:
            left_child[parent] = u
    return left_child


def node_time_depth(tree, min_branch_length=None, max_time="tree"):
    """
    Returns a dictionary mapping nodes in the specified tree to their depth
    in the specified tree (from the root direction). If min_branch_len is
    provided, it specifies the minimum length of each branch. If not specified,
    default to 1.
    """
    if min_branch_length is None:
        min_branch_length = {u: 1 for u in range(tree.tree_sequence.num_nodes)}
    time_node_map = collections.defaultdict(list)
    current_depth = 0
    depth = {}
    # TODO this is basically the same code for the two cases. Refactor so that
    # we use the same code.
    if max_time == "tree":
        for u in tree.nodes():
            time_node_map[tree.time(u)].append(u)
        for t in sorted(time_node_map.keys()):
            for u in time_node_map[t]:
                for v in tree.children(u):
                    current_depth = max(current_depth, depth[v] + min_branch_length[v])
            for u in time_node_map[t]:
                depth[u] = current_depth
            current_depth += 2
        for root in tree.roots:
            current_depth = max(current_depth, depth[root] + min_branch_length[root])
    else:
        assert max_time == "ts"
        ts = tree.tree_sequence
        for node in ts.nodes():
            time_node_map[node.time].append(node.id)
        node_edges = collections.defaultdict(list)
        for edge in ts.edges():
            node_edges[edge.parent].append(edge)

        for t in sorted(time_node_map.keys()):
            for u in time_node_map[t]:
                for edge in node_edges[u]:
                    v = edge.child
                    current_depth = max(current_depth, depth[v] + min_branch_length[v])
            for u in time_node_map[t]:
                depth[u] = current_depth
            current_depth += 2

    return depth, current_depth


class TextTree:
    """
    Draws a reprentation of a tree using unicode drawing characters written
    to a 2D array.
    """

    def __init__(
        self,
        tree,
        node_labels=None,
        max_time=None,
        min_time=None,
        use_ascii=False,
        orientation=None,
        order=None,
    ):
        self.tree = tree
        self.traversal_order = check_order(order)
        self.max_time = check_max_time(max_time, allow_numeric=False)
        self.min_time = check_min_time(min_time, allow_numeric=False)
        self.use_ascii = use_ascii
        self.orientation = check_orientation(orientation)
        self.horizontal_line_char = "━"
        self.vertical_line_char = "┃"
        if use_ascii:
            self.horizontal_line_char = "-"
            self.vertical_line_char = "|"
        # These are set below by the placement algorithms.
        self.width = None
        self.height = None
        self.canvas = None
        # Placement of nodes in the 2D space. Nodes are positioned in one
        # dimension based on traversal ordering and by their time in the
        # other dimension. These are mapped to x and y coordinates according
        # to the orientation.
        self.traversal_position = {}  # Position of nodes in traversal space
        self.time_position = {}
        # Labels for nodes
        self.node_labels = {}

        # Set the node labels
        for u in tree.nodes():
            if node_labels is None:
                # If we don't specify node_labels, default to node ID
                self.node_labels[u] = str(u)
            else:
                # If we do specify node_labels, default to an empty line
                self.node_labels[u] = self.default_node_label
        if node_labels is not None:
            for node, label in node_labels.items():
                self.node_labels[node] = label

        self._assign_time_positions()
        self._assign_traversal_positions()
        self.canvas = np.zeros((self.height, self.width), dtype=str)
        self.canvas[:] = " "
        self._draw()
        self.canvas[:, -1] = "\n"

    def __str__(self):
        return "".join(self.canvas.reshape(self.width * self.height))


class VerticalTextTree(TextTree):
    """
    Text tree rendering where root nodes are at the top and time goes downwards
    into the present.
    """

    @property
    def default_node_label(self):
        return self.vertical_line_char

    def _assign_time_positions(self):
        tree = self.tree
        # TODO when we add mutations to the text tree we'll need to take it into
        # account here. Presumably we need to get the maximum number of mutations
        # per branch.
        self.time_position, total_depth = node_time_depth(tree, max_time=self.max_time)
        self.height = total_depth - 1

    def _assign_traversal_positions(self):
        self.label_x = {}
        left_neighbour = get_left_neighbour(self.tree, self.traversal_order)
        x = 0
        for u in self.tree.nodes(order=self.traversal_order):
            label_size = len(self.node_labels[u])
            if self.tree.is_leaf(u):
                self.traversal_position[u] = x + label_size // 2
                self.label_x[u] = x
                x += label_size + 1
            else:
                coords = [self.traversal_position[c] for c in self.tree.children(u)]
                if len(coords) == 1:
                    self.traversal_position[u] = coords[0]
                else:
                    a = min(coords)
                    b = max(coords)
                    child_mid = int(round(a + (b - a) / 2))
                    self.traversal_position[u] = child_mid
                self.label_x[u] = self.traversal_position[u] - label_size // 2
                neighbour_x = -1
                neighbour = left_neighbour[u]
                if neighbour != NULL:
                    neighbour_x = self.traversal_position[neighbour]
                self.label_x[u] = max(neighbour_x + 1, self.label_x[u])
                x = max(x, self.label_x[u] + label_size + 1)
            assert self.label_x[u] >= 0
        self.width = x

    def _draw(self):
        if self.use_ascii:
            left_child = "+"
            right_child = "+"
            mid_parent = "+"
            mid_parent_child = "+"
            mid_child = "+"
        elif self.orientation == TOP:
            left_child = "┏"
            right_child = "┓"
            mid_parent = "┻"
            mid_parent_child = "╋"
            mid_child = "┳"
        else:
            left_child = "┗"
            right_child = "┛"
            mid_parent = "┳"
            mid_parent_child = "╋"
            mid_child = "┻"

        for u in self.tree.nodes():
            xu = self.traversal_position[u]
            yu = self.time_position[u]
            label = to_np_unicode(self.node_labels[u])
            label_len = label.shape[0]
            label_x = self.label_x[u]
            assert label_x >= 0
            self.canvas[yu, label_x : label_x + label_len] = label
            children = self.tree.children(u)
            if len(children) > 0:
                if len(children) == 1:
                    yv = self.time_position[children[0]]
                    self.canvas[yv:yu, xu] = self.vertical_line_char
                else:
                    left = min(self.traversal_position[v] for v in children)
                    right = max(self.traversal_position[v] for v in children)
                    y = yu - 1
                    self.canvas[y, left + 1 : right] = self.horizontal_line_char
                    self.canvas[y, xu] = mid_parent
                    for v in children:
                        xv = self.traversal_position[v]
                        yv = self.time_position[v]
                        self.canvas[yv:yu, xv] = self.vertical_line_char
                        mid_char = mid_parent_child if xv == xu else mid_child
                        self.canvas[y, xv] = mid_char
                    self.canvas[y, left] = left_child
                    self.canvas[y, right] = right_child
        if self.orientation == TOP:
            self.canvas = np.flip(self.canvas, axis=0)
            # Reverse the time positions so that we can use them in the tree
            # sequence drawing as well.
            flipped_time_position = {
                u: self.height - y - 1 for u, y in self.time_position.items()
            }
            self.time_position = flipped_time_position


class HorizontalTextTree(TextTree):
    """
    Text tree rendering where root nodes are at the left and time goes
    rightwards into the present.
    """

    @property
    def default_node_label(self):
        return self.horizontal_line_char

    def _assign_time_positions(self):
        # TODO when we add mutations to the text tree we'll need to take it into
        # account here. Presumably we need to get the maximum number of mutations
        # per branch.
        self.time_position, total_depth = node_time_depth(
            self.tree, {u: 1 + len(self.node_labels[u]) for u in self.tree.nodes()}
        )
        self.width = total_depth

    def _assign_traversal_positions(self):
        y = 0
        for root in self.tree.roots:
            for u in self.tree.nodes(root, order=self.traversal_order):
                if self.tree.is_leaf(u):
                    self.traversal_position[u] = y
                    y += 2
                else:
                    coords = [self.traversal_position[c] for c in self.tree.children(u)]
                    if len(coords) == 1:
                        self.traversal_position[u] = coords[0]
                    else:
                        a = min(coords)
                        b = max(coords)
                        child_mid = int(round(a + (b - a) / 2))
                        self.traversal_position[u] = child_mid
            y += 1
        self.height = y - 2

    def _draw(self):
        if self.use_ascii:
            top_across = "+"
            bot_across = "+"
            mid_parent = "+"
            mid_parent_child = "+"
            mid_child = "+"
        elif self.orientation == LEFT:
            top_across = "┏"
            bot_across = "┗"
            mid_parent = "┫"
            mid_parent_child = "╋"
            mid_child = "┣"
        else:
            top_across = "┓"
            bot_across = "┛"
            mid_parent = "┣"
            mid_parent_child = "╋"
            mid_child = "┫"

        # Draw in root-right mode as the coordinates go in the expected direction.
        for u in self.tree.nodes():
            yu = self.traversal_position[u]
            xu = self.time_position[u]
            label = to_np_unicode(self.node_labels[u])
            if self.orientation == LEFT:
                # We flip the array at the end so need to reverse the label.
                label = label[::-1]
            label_len = label.shape[0]
            self.canvas[yu, xu : xu + label_len] = label
            children = self.tree.children(u)
            if len(children) > 0:
                if len(children) == 1:
                    xv = self.time_position[children[0]]
                    self.canvas[yu, xv:xu] = self.horizontal_line_char
                else:
                    bot = min(self.traversal_position[v] for v in children)
                    top = max(self.traversal_position[v] for v in children)
                    x = xu - 1
                    self.canvas[bot + 1 : top, x] = self.vertical_line_char
                    self.canvas[yu, x] = mid_parent
                    for v in children:
                        yv = self.traversal_position[v]
                        xv = self.time_position[v]
                        self.canvas[yv, xv:x] = self.horizontal_line_char
                        mid_char = mid_parent_child if yv == yu else mid_child
                        self.canvas[yv, x] = mid_char
                    self.canvas[bot, x] = top_across
                    self.canvas[top, x] = bot_across
        if self.orientation == LEFT:
            self.canvas = np.flip(self.canvas, axis=1)
            # Move the padding to the left.
            self.canvas[:, :-1] = self.canvas[:, 1:]
            self.canvas[:, -1] = " "


--- ../../tskit/python/tskit/stats.py ---

 
"""
Module responsible for computing various statistics on tree sequences.
"""
import sys
import threading

import numpy as np

import _tskit


class LdCalculator:
    """
    Class for calculating `linkage disequilibrium
    <https://en.wikipedia.org/wiki/Linkage_disequilibrium>`_ coefficients
    between pairs of sites in a :class:`TreeSequence`.

    .. note:: This interface is deprecated and a replacement is planned.
        Please see https://github.com/tskit-dev/tskit/issues/1900 for
        more information. Note also that the current implementation is
        quite limited (see warning below).

    .. warning:: This class does not currently support sites that have more than one
        mutation. Using it on such a tree sequence will raise a LibraryError with
        an "Only infinite sites mutations supported" message.

        Silent mutations are also not supported and will result in a LibraryError.

    :param TreeSequence tree_sequence: The tree sequence of interest.
    """

    def __init__(self, tree_sequence):
        self._tree_sequence = tree_sequence
        self._ll_ld_calculator = _tskit.LdCalculator(
            tree_sequence.get_ll_tree_sequence()
        )
        # To protect low-level C code, only one method may execute on the
        # low-level objects at one time.
        self._instance_lock = threading.Lock()

    def get_r2(self, a, b):
        # Deprecated alias for r2(a, b)
        return self.r2(a, b)

    def r2(self, a, b):
        """
        Returns the value of the :math:`r^2` statistic between the pair of
        sites at the specified indexes. This method is *not* an efficient
        method for computing large numbers of pairwise LD values; please use either
        :meth:`.r2_array` or :meth:`.r2_matrix` for this purpose.

        :param int a: The index of the first site.
        :param int b: The index of the second site.
        :return: The value of :math:`r^2` between the sites at indexes
            ``a`` and ``b``.
        :rtype: float
        """
        with self._instance_lock:
            return self._ll_ld_calculator.get_r2(a, b)

    def get_r2_array(self, a, direction=1, max_mutations=None, max_distance=None):
        # Deprecated alias for r2_array
        return self.r2_array(
            a,
            direction=direction,
            max_mutations=max_mutations,
            max_distance=max_distance,
        )

    def r2_array(
        self, a, direction=1, max_mutations=None, max_distance=None, max_sites=None
    ):
        """
        Returns the value of the :math:`r^2` statistic between the focal
        site at index :math:`a` and a set of other sites. The method
        operates by starting at the focal site and iterating over adjacent
        sites (in either the forward or backwards direction) until either a
        maximum number of other sites have been considered (using the
        ``max_sites`` parameter), a maximum distance in sequence
        coordinates has been reached (using the ``max_distance`` parameter) or
        the start/end of the sequence has been reached. For every site
        :math:`b` considered, we then insert the value of :math:`r^2` between
        :math:`a` and :math:`b` at the corresponding index in an array, and
        return the entire array. If the returned array is :math:`x` and
        ``direction`` is :data:`tskit.FORWARD` then :math:`x[0]` is the
        value of the statistic for :math:`a` and :math:`a + 1`, :math:`x[1]`
        the value for :math:`a` and :math:`a + 2`, etc. Similarly, if
        ``direction`` is :data:`tskit.REVERSE` then :math:`x[0]` is the
        value of the statistic for :math:`a` and :math:`a - 1`, :math:`x[1]`
        the value for :math:`a` and :math:`a - 2`, etc.

        :param int a: The index of the focal sites.
        :param int direction: The direction in which to travel when
            examining other sites. Must be either
            :data:`tskit.FORWARD` or :data:`tskit.REVERSE`. Defaults
            to :data:`tskit.FORWARD`.
        :param int max_sites: The maximum number of sites to return
            :math:`r^2` values for. Defaults to as many sites as
            possible.
        :param int max_mutations: Deprecated synonym for max_sites.
        :param float max_distance: The maximum absolute distance between
            the focal sites and those for which :math:`r^2` values
            are returned.
        :return: An array of double precision floating point values
            representing the :math:`r^2` values for sites in the
            specified direction.
        :rtype: numpy.ndarray
        """
        if max_mutations is not None and max_sites is not None:
            raise ValueError("max_mutations is a deprecated synonym for max_sites")
        if max_mutations is not None:
            max_sites = max_mutations
        max_sites = -1 if max_sites is None else max_sites
        if max_distance is None:
            max_distance = sys.float_info.max
        with self._instance_lock:
            return self._ll_ld_calculator.get_r2_array(
                a,
                direction=direction,
                max_sites=max_sites,
                max_distance=max_distance,
            )

    def get_r2_matrix(self):
        # Deprecated alias for r2_matrix
        return self.r2_matrix()

    def r2_matrix(self):
        """
        Returns the complete :math:`m \\times m` matrix of pairwise
        :math:`r^2` values in a tree sequence with :math:`m` sites.

        :return: An 2 dimensional square array of double precision
            floating point values representing the :math:`r^2` values for
            all pairs of sites.
        :rtype: numpy.ndarray
        """
        m = self._tree_sequence.num_sites
        A = np.ones((m, m), dtype=float)
        for j in range(m - 1):
            a = self.get_r2_array(j)
            A[j, j + 1 :] = a
            A[j + 1 :, j] = a
        return A


--- ../../tskit/python/tskit/intervals.py ---


"""
Utilities for working with intervals and interval maps.
"""
from __future__ import annotations

import collections.abc
import numbers

import numpy as np

import tskit
import tskit.util as util


class RateMap(collections.abc.Mapping):
    """
    A class mapping a non-negative rate value to a set of non-overlapping intervals
    along the genome. Intervals for which the rate is unknown (i.e., missing data)
    are encoded by NaN values in the ``rate`` array.

    :param list position: A list of :math:`n+1` positions, starting at 0, and ending
        in the sequence length over which the RateMap will apply.
    :param list rate: A list of :math:`n` positive rates that apply between each
        position. Intervals with missing data are encoded by NaN values.
    """

    # The args are marked keyword only to give us some flexibility in how we
    # create class this in the future.
    def __init__(
        self,
        *,
        position,
        rate,
    ):
        # Making the arrays read-only guarantees rate and cumulative mass stay in sync
        # We prevent the arrays themselves being overwritten by making self.position,
        # etc properties.

        # TODO we always coerce the position type to float here, but we may not
        # want to do this. int32 is a perfectly good choice a lot of the time.
        self._position = np.array(position, dtype=float)
        self._position.flags.writeable = False
        self._rate = np.array(rate, dtype=float)
        self._rate.flags.writeable = False
        size = len(self._position)
        if size < 2:
            raise ValueError("Must have at least two positions")
        if len(self._rate) != size - 1:
            raise ValueError(
                "Rate array must have one less entry than the position array"
            )
        if self._position[0] != 0:
            raise ValueError("First position must be zero")

        span = self.span
        if np.any(span <= 0):
            bad_pos = np.where(span <= 0)[0] + 1
            raise ValueError(
                f"Position values not strictly increasing at indexes {bad_pos}"
            )
        if np.any(self._rate < 0):
            bad_rates = np.where(self._rate < 0)[0]
            raise ValueError(f"Rate values negative at indexes {bad_rates}")
        self._missing = np.isnan(self.rate)
        self._num_missing_intervals = np.sum(self._missing)
        if self._num_missing_intervals == len(self.rate):
            raise ValueError("All intervals are missing data")
        # We don't expose the cumulative mass array as a part of the array
        # API is it's not quite as obvious how it lines up for each interval.
        # It's really the sum of the mass up to but not including the current
        # interval, which is a bit confusing. Probably best to just leave
        # it as a function, so that people can sample at regular positions
        # along the genome anyway, emphasising that it's a continuous function,
        # not a step function like the other interval attributes.
        self._cumulative_mass = np.insert(np.nancumsum(self.mass), 0, 0)
        assert self._cumulative_mass[0] == 0
        self._cumulative_mass.flags.writeable = False

    @property
    def left(self):
        """
        The left position of each interval (inclusive).
        """
        return self._position[:-1]

    @property
    def right(self):
        """
        The right position of each interval (exclusive).
        """
        return self._position[1:]

    @property
    def mid(self):
        """
        Returns the midpoint of each interval.
        """
        mid = self.left + self.span / 2
        mid.flags.writeable = False
        return mid

    @property
    def span(self):
        """
        Returns the span (i.e., ``right - left``) of each of the intervals.
        """
        span = self.right - self.left
        span.flags.writeable = False
        return span

    @property
    def position(self):
        """
        The breakpoint positions between intervals. This is equal to the
        :attr:`~.RateMap.left` array with the :attr:`sequence_length`
        appended.
        """
        return self._position

    @property
    def rate(self):
        """
        The rate associated with each interval. Missing data is encoded
        by NaN values.
        """
        return self._rate

    @property
    def mass(self):
        r"""
        The "mass" of each interval, defined as the :attr:`~.RateMap.rate`
        :math:`\times` :attr:`~.RateMap.span`. This is NaN for intervals
        containing missing data.
        """
        return self._rate * self.span

    @property
    def missing(self):
        """
        A boolean array encoding whether each interval contains missing data.
        Equivalent to ``np.isnan(rate_map.rate)``
        """
        return self._missing

    @property
    def non_missing(self):
        """
        A boolean array encoding whether each interval contains non-missing data.
        Equivalent to ``np.logical_not(np.isnan(rate_map.rate))``
        """
        return ~self._missing

    #
    # Interval counts
    #

    @property
    def num_intervals(self) -> int:
        """
        The total number of intervals in this map. Equal to
        :attr:`~.RateMap.num_missing_intervals` +
        :attr:`~.RateMap.num_non_missing_intervals`.
        """
        return len(self._rate)

    @property
    def num_missing_intervals(self) -> int:
        """
        Returns the number of missing intervals, i.e., those in which the
        :attr:`~.RateMap.rate` value is a NaN.
        """
        return self._num_missing_intervals

    @property
    def num_non_missing_intervals(self) -> int:
        """
        The number of non missing intervals, i.e., those in which the
        :attr:`~.RateMap.rate` value is not a NaN.
        """
        return self.num_intervals - self.num_missing_intervals

    @property
    def sequence_length(self):
        """
        The sequence length covered by this map
        """
        return self.position[-1]

    @property
    def total_mass(self):
        """
        The cumulative total mass over the entire map.
        """
        return self._cumulative_mass[-1]

    @property
    def mean_rate(self):
        """
        The mean rate over this map weighted by the span covered by each rate.
        Unknown intervals are excluded.
        """
        total_span = np.sum(self.span[self.non_missing])
        return self.total_mass / total_span

    def get_rate(self, x):
        """
        Return the rate at the specified list of positions.

        .. note:: This function will return a NaN value for any positions
            that contain missing data.

        :param numpy.ndarray x: The positions for which to return values.
        :return: An array of rates, the same length as ``x``.
        :rtype: numpy.ndarray
        """
        loc = np.searchsorted(self.position, x, side="right") - 1
        if np.any(loc < 0) or np.any(loc >= len(self.rate)):
            raise ValueError("position out of bounds")
        return self.rate[loc]

    def get_cumulative_mass(self, x):
        """
        Return the cumulative mass of the map up to (but not including) a
        given point for a list of positions along the map. This is equal to
        the integral of the rate from 0 to the point.

        :param numpy.ndarray x: The positions for which to return values.

        :return: An array of cumulative mass values, the same length as ``x``
        :rtype: numpy.ndarray
        """
        x = np.array(x)
        if np.any(x < 0) or np.any(x > self.sequence_length):
            raise ValueError(f"Cannot have positions < 0 or > {self.sequence_length}")
        return np.interp(x, self.position, self._cumulative_mass)

    def find_index(self, x: float) -> int:
        """
        Returns the index of the interval that the specified position falls within,
        such that ``rate_map.left[index] <= x < self.rate_map.right[index]``.

        :param float x: The position to search.
        :return: The index of the interval containing this point.
        :rtype: int
        :raises KeyError: if the position is not contained in any of the intervals.
        """
        if x < 0 or x >= self.sequence_length:
            raise KeyError(f"Position {x} out of bounds")
        index = np.searchsorted(self.position, x, side="left")
        if x < self.position[index]:
            index -= 1
        assert self.left[index] <= x < self.right[index]
        return index

    def missing_intervals(self):
        """
        Returns the left and right coordinates of the intervals containing
        missing data in this map as a 2D numpy array
        with shape (:attr:`~.RateMap.num_missing_intervals`, 2). Each row
        of this returned array is therefore a ``left``, ``right`` tuple
        corresponding to the coordinates of the missing intervals.

        :return: A numpy array of the coordinates of intervals containing
            missing data.
        :rtype: numpy.ndarray
        """
        out = np.empty((self.num_missing_intervals, 2))
        out[:, 0] = self.left[self.missing]
        out[:, 1] = self.right[self.missing]
        return out

    def asdict(self):
        return {"position": self.position, "rate": self.rate}

    #
    # Dunder methods. We implement the Mapping protocol via __iter__, __len__
    # and __getitem__. We have some extra semantics for __getitem__, providing
    # slice notation.
    #

    def __iter__(self):
        # The clinching argument for using mid here is that if we used
        # left instead we would have
        #   RateMap([0, 1], [0.1]) == RateMap([0, 100], [0.1])
        # by the inherited definition of equality since the dictionary items
        # would be equal.
        # Similarly, we only return the midpoints of known intervals
        # because NaN values are not equal, and we would need to do
        # something to work around this. It seems reasonable that
        # this high-level operation returns the *known* values only
        # anyway.
        yield from self.mid[self.non_missing]

    def __len__(self):
        return np.sum(self.non_missing)

    def __getitem__(self, key):
        if isinstance(key, slice):
            if key.step is not None:
                raise TypeError("Only interval slicing is supported")
            return self.slice(key.start, key.stop)
        if isinstance(key, numbers.Number):
            index = self.find_index(key)
            if np.isnan(self.rate[index]):
                # To be consistent with the __iter__ definition above we
                # don't consider these missing positions to be "in" the map.
                raise KeyError(f"Position {key} is within a missing interval")
            return self.rate[index]
        # TODO we could implement numpy array indexing here and call
        # to get_rate. Note we'd need to take care that we return a keyerror
        # if the returned array contains any nans though.
        raise KeyError("Key {key} not in map")

    def _text_header_and_rows(self, limit=None):
        headers = ("left", "right", "mid", "span", "rate")
        num_rows = len(self.left)
        rows = []
        row_indexes = util.truncate_rows(num_rows, limit)
        for j in row_indexes:
            if j == -1:
                rows.append(f"__skipped__{num_rows-limit}")
            else:
                rows.append(
                    [
                        f"{self.left[j]:.10g}",
                        f"{self.right[j]:.10g}",
                        f"{self.mid[j]:.10g}",
                        f"{self.span[j]:.10g}",
                        f"{self.rate[j]:.2g}",
                    ]
                )
        return headers, rows

    def __str__(self):
        header, rows = self._text_header_and_rows(
            limit=tskit._print_options["max_lines"]
        )
        table = util.unicode_table(
            rows=rows,
            header=header,
            column_alignments="<<>>>",
        )
        return table

    def _repr_html_(self):
        header, rows = self._text_header_and_rows(
            limit=tskit._print_options["max_lines"]
        )
        return util.html_table(rows, header=header)

    def __repr__(self):
        return f"RateMap(position={repr(self.position)}, rate={repr(self.rate)})"

    #
    # Methods for building rate maps.
    #

    def copy(self) -> RateMap:
        """
        Returns a deep copy of this RateMap.
        """
        # We take read-only copies of the arrays in the constructor anyway, so
        # no need for copying.
        return RateMap(position=self.position, rate=self.rate)

    def slice(self, left=None, right=None, *, trim=False) -> RateMap:  # noqa: A003
        """
        Returns a subset of this rate map in the specified interval.

        :param float left: The left coordinate (inclusive) of the region to keep.
            If ``None``, defaults to 0.
        :param float right: The right coordinate (exclusive) of the region to keep.
            If ``None``, defaults to the sequence length.
        :param bool trim: If True, remove the flanking regions such that the
            sequence length of the new rate map is ``right`` - ``left``. If ``False``
            (default), do not change the coordinate system and mark the flanking
            regions as "unknown".
        :return: A new RateMap instance
        :rtype: RateMap
        """
        left = 0 if left is None else left
        right = self.sequence_length if right is None else right
        if not (0 <= left < right <= self.sequence_length):
            raise KeyError(f"Invalid slice: left={left}, right={right}")

        i = self.find_index(left)
        j = i + np.searchsorted(self.position[i:], right, side="right")
        if right > self.position[j - 1]:
            j += 1

        position = self.position[i:j].copy()
        rate = self.rate[i : j - 1].copy()
        position[0] = left
        position[-1] = right

        if trim:
            # Return trimmed map with changed coords
            return RateMap(position=position - left, rate=rate)

        # Need to check regions before & after sliced region are filled out:
        if left != 0:
            if np.isnan(rate[0]):
                position[0] = 0  # Extend
            else:
                rate = np.insert(rate, 0, np.nan)  # Prepend
                position = np.insert(position, 0, 0)
        if right != self.position[-1]:
            if np.isnan(rate[-1]):
                position[-1] = self.sequence_length  # Extend
            else:
                rate = np.append(rate, np.nan)  # Append
                position = np.append(position, self.position[-1])
        return RateMap(position=position, rate=rate)

    @staticmethod
    def uniform(sequence_length, rate) -> RateMap:
        """
        Create a uniform rate map
        """
        return RateMap(position=[0, sequence_length], rate=[rate])

    @staticmethod
    def read_hapmap(
        fileobj,
        sequence_length=None,
        *,
        has_header=True,
        position_col=None,
        rate_col=None,
        map_col=None,
    ):
        # Black barfs with an INTERNAL_ERROR trying to reformat this docstring,
        # so we explicitly disable reformatting here.
        # fmt: off
        """
        Parses the specified file in HapMap format and returns a :class:`.RateMap`.
        HapMap files must white-space-delimited, and by default are assumed to
        contain a single header line (which is ignored). Each subsequent line
        then contains a physical position (in base pairs) and either a genetic
        map position (in centiMorgans) or a recombination rate (in centiMorgans
        per megabase). The value in the rate column in a given line gives the
        constant rate between the physical position in that line (inclusive) and the
        physical position on the next line (exclusive).
        By default, the second column of the file is taken
        as the physical position and the fourth column is taken as the genetic
        position, as seen in the following sample of the format::

            Chromosome	Position(bp)  Rate(cM/Mb)  Map(cM)
            chr10       48232         0.1614       0.002664
            chr10       48486         0.1589       0.002705
            chr10       50009         0.159        0.002947
            chr10       52147         0.1574       0.003287
            ...
            chr10	133762002     3.358        181.129345
            chr10	133766368     0.000        181.144008

        In the example above, the first row has a nonzero genetic map position
        (last column, cM), implying a nonzero recombination rate before that
        position, that is assumed to extend to the start of the chromosome
        (at position 0 bp). However, if the first line has a nonzero bp position
        (second column) and a zero genetic map position (last column, cM),
        then the recombination rate before that position is *unknown*, producing
        :ref:`missing data <sec_rate_maps_missing>`.

        .. note::
            The rows are all assumed to come from the same contig, and the
            first column is currently ignored. Therefore if you have a single
            file containing several contigs or chromosomes, you must must split
            it up into multiple files, and pass each one separately to this
            function.

        :param str fileobj: Filename or file to read. This is passed directly
            to :func:`numpy.loadtxt`, so if the filename extension is .gz or .bz2,
            the file is decompressed first
        :param float sequence_length: The total length of the map. If ``None``,
            then assume it is the last physical position listed in the file.
            Otherwise it must be greater then or equal to the last physical
            position in the file, and the region between the last physical position
            and the sequence_length is padded with a rate of zero.
        :param bool has_header: If True (default), assume the file has a header row
            and ignore the first line of the file.
        :param int position_col: The zero-based index of the column in the file
            specifying the physical position in base pairs. If ``None`` (default)
            assume an index of 1 (i.e. the second column).
        :param int rate_col: The zero-based index of the column in the file
            specifying the rate in cM/Mb. If ``None`` (default) do not use the rate
            column, but calculate rates using the genetic map positions, as
            specified in ``map_col``. If the rate column is used, the
            interval from 0 to first physical position in the file is marked as
            unknown, and the last value in the rate column must be zero.
        :param int map_col: The zero-based index of the column in the file
            specifying the genetic map position in centiMorgans. If ``None``
            (default), assume an index of 3 (i.e. the fourth column). If the first
            genetic position is 0 the interval from position 0 to the first
            physical position in the file is marked as unknown. Otherwise, act
            as if an additional row, specifying physical position 0 and genetic
            position 0, exists at the start of the file.
        :return: A RateMap object.
        :rtype: RateMap
        """
        # fmt: on
        column_defs = {}  # column definitions passed to np.loadtxt
        if rate_col is None and map_col is None:
            # Default to map_col
            map_col = 3
        elif rate_col is not None and map_col is not None:
            raise ValueError("Cannot specify both rate_col and map_col")
        if map_col is not None:
            column_defs[map_col] = ("map", float)
        else:
            column_defs[rate_col] = ("rate", float)
        position_col = 1 if position_col is None else position_col
        if position_col in column_defs:
            raise ValueError(
                "Cannot specify the same columns for position_col and "
                "rate_col or map_col"
            )
        column_defs[position_col] = ("pos", int)

        column_names = [c[0] for c in column_defs.values()]
        column_data = np.loadtxt(
            fileobj,
            skiprows=1 if has_header else 0,
            dtype=list(column_defs.values()),
            usecols=list(column_defs.keys()),
            unpack=True,
        )
        data = dict(zip(column_names, column_data))

        if "map" not in data:
            assert "rate" in data
            if data["rate"][-1] != 0:
                raise ValueError("The last entry in the 'rate' column must be zero")
            pos_Mb = data["pos"] / 1e6
            map_pos = np.cumsum(data["rate"][:-1] * np.diff(pos_Mb))
            data["map"] = np.insert(map_pos, 0, 0) / 100
        else:
            data["map"] /= 100  # Convert centiMorgans to Morgans
        if len(data["map"]) == 0:
            raise ValueError("Empty hapmap file")

        # TO DO: read in chrom name from col 0 and poss set as .name
        # attribute on the RateMap

        physical_positions = data["pos"]
        genetic_positions = data["map"]
        start = physical_positions[0]
        end = physical_positions[-1]

        if genetic_positions[0] > 0 and start == 0:
            raise ValueError(
                "The map distance at the start of the chromosome must be zero"
            )
        if start > 0:
            physical_positions = np.insert(physical_positions, 0, 0)
            if genetic_positions[0] > 0:
                # Exception for a map that starts > 0cM: include the start rate
                # in the mean
                start = 0
            genetic_positions = np.insert(genetic_positions, 0, 0)

        if sequence_length is not None:
            if sequence_length < end:
                raise ValueError(
                    "The sequence_length cannot be less than the last physical position "
                    f" ({physical_positions[-1]})"
                )
            if sequence_length > end:
                physical_positions = np.append(physical_positions, sequence_length)
                genetic_positions = np.append(genetic_positions, genetic_positions[-1])

        assert genetic_positions[0] == 0
        rate = np.diff(genetic_positions) / np.diff(physical_positions)
        if start != 0:
            rate[0] = np.nan
        if end != physical_positions[-1]:
            rate[-1] = np.nan
        return RateMap(position=physical_positions, rate=rate)


--- ../../tskit/python/tskit/genotypes.py ---


from __future__ import annotations

import collections
import logging
import typing

import numpy as np

import _tskit
import tskit
import tskit.trees as trees
import tskit.util as util


class Variant:
    """
    A variant in a tree sequence, describing the observed genetic variation
    among samples for a given site. A variant consists of (a) a tuple of
    **alleles** listing the potential allelic states which samples at this site
    can possess; (b) an array of **genotypes** mapping sample IDs to the observed
    alleles (c) a reference to the :class:`Site` at which the Variant has been decoded
    and (d) an array of **samples** giving the node id to which the each element of
    the genotypes array corresponds.

    After creation a Variant is not yet decoded, and has no genotypes.
    To decode a Variant, call the :meth:`decode` method. The Variant class will then
    use a Tree, internal to the Variant, to seek to the position of the site and
    decode the genotypes at that site. It is therefore much more efficient to visit
    sites in sequential genomic order, either in a forwards or backwards direction,
    than to do so randomly.

    Each element in the ``alleles`` tuple is a string, representing an
    observed allelic state that may be seen at this site. The ``alleles`` tuple,
    which is guaranteed not to contain any duplicates, is generated in one of two
    ways. The first (and default) way is for ``tskit`` to generate the encoding on
    the fly while generating genotypes. In this case, the first element of this
    tuple is guaranteed to be the same as the site's ``ancestral_state`` value.
    Note that allelic values may be listed that are not referred to by any
    samples. For example, if we have a site that is fixed for the derived state
    (i.e., we have a mutation over the tree root), all genotypes will be 1, but
    the alleles list will be equal to ``('0', '1')``. Other than the
    ancestral state being the first allele, the alleles are listed in
    no particular order, and the ordering should not be relied upon
    (but see the notes on missing data below).

    The second way is for the user to define the mapping between
    genotype values and allelic state strings using the
    ``alleles`` parameter to the :meth:`TreeSequence.variants` method.
    In this case, there is no indication of which allele is the ancestral state,
    as the ordering is determined by the user.

    The ``genotypes`` represent the observed allelic states for each sample,
    such that ``var.alleles[var.genotypes[j]]`` gives the string allele
    for sample ID ``j``. Thus, the elements of the genotypes array are
    indexes into the ``alleles`` list. The genotypes are provided in this
    way via a numpy numeric array to enable efficient calculations. To obtain a
    (less efficient) array of allele strings for each sample, you can use e.g.
    ``np.asarray(variant.alleles)[variant.genotypes]``.

    When :ref:`missing data<sec_data_model_missing_data>` is present at a given
    site, the property ``has_missing_data`` will be True, at least one element
    of the ``genotypes`` array will be equal to ``tskit.MISSING_DATA``, and the
    last element of the ``alleles`` array will be ``None``. Note that in this
    case ``variant.num_alleles`` will **not** be equal to
    ``len(variant.alleles)``. The rationale for adding ``None`` to the end of
    the ``alleles`` list is to help code that does not handle missing data
    correctly fail early rather than introducing subtle and hard-to-find bugs.
    As ``tskit.MISSING_DATA`` is equal to -1, code that decodes genotypes into
    allelic values without taking missing data into account would otherwise
    incorrectly output the last allele in the list.

    :param TreeSequence tree_sequence: The tree sequence to which this variant
        belongs.
    :param array_like samples: An array of node IDs for which to generate
        genotypes, or None for all sample nodes. Default: None.
    :param bool isolated_as_missing: If True, the genotype value assigned to
        missing samples (i.e., isolated samples without mutations) is
        :data:`.MISSING_DATA` (-1). If False, missing samples will be
        assigned the allele index for the ancestral state.
        Default: True.
    :param tuple alleles: A tuple of strings defining the encoding of
        alleles as integer genotype values. At least one allele must be provided.
        If duplicate alleles are provided, output genotypes will always be
        encoded as the first occurrence of the allele. If None (the default),
        the alleles are encoded as they are encountered during genotype
        generation.

    """

    def __init__(
        self, tree_sequence, samples=None, isolated_as_missing=None, alleles=None
    ):
        if isolated_as_missing is None:
            isolated_as_missing = True
        self.tree_sequence = tree_sequence
        if samples is not None:
            samples = util.safe_np_int_cast(samples, np.int32)
        self._ll_variant = _tskit.Variant(
            tree_sequence._ll_tree_sequence,
            samples=samples,
            isolated_as_missing=isolated_as_missing,
            alleles=alleles,
        )

    def _check_decoded(self):
        if self._ll_variant.site_id == tskit.NULL:
            raise ValueError(
                "This variant has not yet been decoded at a specific site, "
                "call Variant.decode to set the site."
            )

    @property
    def site(self) -> trees.Site:
        """
        The Site object for the site at which this variant has been decoded.
        """
        self._check_decoded()
        return self.tree_sequence.site(self._ll_variant.site_id)

    @property
    def alleles(self) -> tuple[str | None, ...]:
        """
        A tuple of the allelic values which samples can possess at the current
        site. Unless an encoding of alleles is specified when creating this
        variant instance, the first element of this tuple is always the site's
        ancestral state.
        """
        return self._ll_variant.alleles

    @property
    def samples(self) -> np.ndarray:
        """
        A numpy array of the node ids whose genotypes will be returned
        by the :meth:`genotypes` method.
        """
        return self._ll_variant.samples

    @property
    def genotypes(self) -> np.ndarray:
        """
        An array of indexes into the list ``alleles``, giving the
        state of each sample at the current site.
        """
        self._check_decoded()
        return self._ll_variant.genotypes

    @property
    def isolated_as_missing(self) -> bool:
        """
        True if isolated samples are decoded to missing data. If False, isolated
        samples are decoded to the ancestral state.
        """
        return self._ll_variant.isolated_as_missing

    @property
    def has_missing_data(self) -> bool:
        """
        True if there is missing data for any of the
        samples at the current site.
        """
        alleles = self._ll_variant.alleles
        return len(alleles) > 0 and alleles[-1] is None

    @property
    def num_missing(self) -> int:
        """
        The number of samples with missing data at this site.
        """
        return np.sum(self.genotypes == tskit.NULL)

    @property
    def num_alleles(self) -> int:
        """
        The number of distinct alleles at this site. Note that this may
        not be the same as the number of distinct values in the genotypes
        array: firstly missing data is not counted as an allele, and secondly,
        the site may contain mutations to alternative allele states (which are
        counted in the number of alleles) without the mutation being inherited
        by any of the samples.
        """
        return len(self.alleles) - self.has_missing_data

    # Deprecated alias to avoid breaking existing code.
    @property
    def position(self) -> float:
        return self.site.position

    # Deprecated alias to avoid breaking existing code.
    @property
    def index(self) -> int:
        return self._ll_variant.site_id

    # We need a custom eq for the numpy array
    def __eq__(self, other) -> bool:
        return (
            isinstance(other, Variant)
            and self.tree_sequence == other.tree_sequence
            and self._ll_variant.site_id == other._ll_variant.site_id
            and self._ll_variant.site_id != tskit.NULL
            and self._ll_variant.alleles == other._ll_variant.alleles
            and np.array_equal(self._ll_variant.genotypes, other._ll_variant.genotypes)
        )

    def decode(self, site_id) -> None:
        """
        Decode the variant at the given site, setting the site ID, genotypes and
        alleles to those of the site and samples of this Variant.

        :param int site_id: The ID of the site to decode. This must be a valid site ID.
        """
        self._ll_variant.decode(site_id)

    def copy(self) -> Variant:
        """
        Create a copy of this Variant. Note that calling :meth:`decode` on the
        copy will fail as it does not take a copy of the internal tree.

        :return: The copy of this Variant.
        """
        variant_copy = Variant.__new__(Variant)
        variant_copy.tree_sequence = self.tree_sequence
        variant_copy._ll_variant = self._ll_variant.restricted_copy()
        return variant_copy

    def states(self, missing_data_string=None) -> np.ndarray:
        """
        Returns the allelic states at this site as an array of strings.

        .. warning::
            Using this method is inefficient compared to working with the
            underlying integer representation of genotypes as returned by
            the :attr:`~Variant.genotypes` property.

        :param str missing_data_string: A string that will be used to represent missing
            data. If any normal allele contains this character, an error is raised.
            Default: `None`, treated as `'N'`.
        :return: An numpy array of strings of length ``num_sites``.
        """
        if missing_data_string is None:
            missing_data_string = "N"
        elif not isinstance(missing_data_string, str):
            # Must explicitly test here, otherwise we output a numpy object array
            raise ValueError("Missing data string is not a string")
        alleles = self.alleles
        if alleles[-1] is None:
            if missing_data_string in alleles:
                raise ValueError(
                    "An existing allele is equal to the "
                    f"missing data string '{missing_data_string}'"
                )
            alleles = alleles[:-1] + (missing_data_string,)
        return np.array(alleles)[self.genotypes]

    def counts(self) -> typing.Counter[str | None]:
        """
        Returns a :class:`python:collections.Counter` object providing counts for each
        possible :attr:`allele <Variant.alleles>` at this site: i.e. the number of
        samples possessing that allele among the set of samples specified when creating
        this Variant (by default, this is all the sample nodes in the tree sequence).
        Missing data is represented by an allelic state of ``None``.

        :return: A counter of the number of samples associated with each allele.
        """
        counts = collections.Counter()
        if self.alleles[-1] is None:
            # we have to treat the last element of the genotypes array as special
            counts[None] = np.sum(self.genotypes == tskit.MISSING_DATA)
            for i, allele in enumerate(self.alleles[:-1]):
                counts[allele] = np.sum(self.genotypes == i)
        else:
            bincounts = np.bincount(self.genotypes, minlength=self.num_alleles)
            for i, allele in enumerate(self.alleles):
                counts[allele] = bincounts[i]
        return counts

    def frequencies(self, remove_missing=None) -> dict[str, float]:
        """
        Return a dictionary mapping each possible :attr:`allele <Variant.alleles>`
        at this site to the frequency of that allele: i.e. the number of samples
        with that allele divided by the total number of samples, among the set of
        samples specified when creating this Variant (by default, this is all the
        sample nodes in the tree sequence). Note, therefore, that if a restricted set
        of samples was specified on creation, the allele frequencies returned here
        will *not* be the global allele frequencies in the whole tree sequence.

        :param bool remove_missing: If True, only samples with non-missing data will
            be counted in the total number of samples used to calculate the frequency,
            and no information on the frequency of missing data is returned. Otherwise
            (default), samples with missing data are included when calculating
            frequencies.
        :return: A dictionary mapping allelic states to the frequency of each allele
            among the samples
        """
        if remove_missing is None:
            remove_missing = False
        total = len(self.samples)
        if remove_missing:
            total -= self.num_missing
        if total == 0:
            logging.warning(
                "No non-missing samples at this site, frequencies undefined"
            )
        return {
            allele: count / total if total > 0 else np.nan
            for allele, count in self.counts().items()
            if not (allele is None and remove_missing)
        }

    def __str__(self) -> str:
        """
        Return a plain text summary of the contents of a variant.
        """
        try:
            site_id = self.site.id
            site_position = self.site.position
            counts = self.counts()
            freqs = self.frequencies()
            rows = (
                [
                    ["Site id", f"{site_id:,}"],
                    ["Site position", f"{site_position:,}"],
                    ["Number of samples", f"{len(self.samples):,}"],
                    ["Number of alleles", f"{self.num_alleles:,}"],
                ]
                + [
                    [
                        f"Samples with allele "
                        f"""{'missing' if k is None else "'" + k + "'"}""",
                        f"{counts[k]:,} ({freqs[k] * 100:.2g}%)",
                    ]
                    for k in self.alleles
                ]
                + [
                    ["Has missing data", str(self.has_missing_data)],
                    ["Isolated as missing", str(bool(self.isolated_as_missing))],
                ]
            )
        except ValueError as err:
            rows = [[str(err), ""]]
        return util.unicode_table(rows, title="Variant")

    def _repr_html_(self) -> str:
        """
        Return an html summary of a variant. Called by Jupyter notebooks
        to render a Variant.
        """
        return util.variant_html(self)

    def __repr__(self):
        d = {
            "site": self.site,
            "samples": self.samples,
            "alleles": self.alleles,
            "genotypes": self.genotypes,
            "has_missing_data": self.has_missing_data,
            "isolated_as_missing": self.isolated_as_missing,
        }
        return f"Variant({repr(d)})"


#
# Miscellaneous auxiliary methods.
#
def allele_remap(alleles_from, alleles_to):
    # Returns an index map from the elements in one list (alleles_from)
    # to the elements of another list (alleles_to).
    #
    # If some elements in alleles_from are not in alleles_to,
    # then indices outside of alleles_to are used.
    alleles_to = np.array(alleles_to, dtype="U")
    alleles_from = np.array(alleles_from, dtype="U")
    allele_map = np.empty_like(alleles_from, dtype="uint32")
    overflow = len(alleles_to)
    for i, allele in enumerate(alleles_from):
        try:
            # Use the index of the first matching element.
            allele_map[i] = np.where(alleles_to == allele)[0][0]
        except IndexError:
            allele_map[i] = overflow
            overflow += 1
    return allele_map


--- ../../tskit/python/tskit/__main__.py ---

from . import cli

if __name__ == "__main__":
    cli.tskit_main()


--- ../../tskit/python/tskit/trees.py ---


"""
Module responsible for managing trees and tree sequences.
"""
from __future__ import annotations

import base64
import builtins
import collections
import concurrent.futures
import functools
import io
import itertools
import math
import numbers
import warnings
from dataclasses import dataclass
from typing import Any
from typing import NamedTuple

import numpy as np

import _tskit
import tskit
import tskit.combinatorics as combinatorics
import tskit.drawing as drawing
import tskit.metadata as metadata_module
import tskit.tables as tables
import tskit.text_formats as text_formats
import tskit.util as util
import tskit.vcf as vcf
from tskit import NODE_IS_SAMPLE
from tskit import NULL
from tskit import UNKNOWN_TIME

LEGACY_MS_LABELS = "legacy_ms"


class CoalescenceRecord(NamedTuple):
    left: float
    right: float
    node: int
    children: np.ndarray
    time: float
    population: int


class Interval(NamedTuple):
    """
    A tuple of 2 numbers, ``[left, right)``, defining an interval over the genome.
    """

    left: float | int
    """
    The left hand end of the interval. By convention this value is included
    in the interval
    """
    right: float | int
    """
    The right hand end of the interval. By convention this value is *not*
    included in the interval, i.e., the interval is half-open.
    """

    @property
    def span(self) -> float | int:
        """
        The span of the genome covered by this interval, simply ``right-left``.
        """
        return self.right - self.left

    @property
    def mid(self) -> float | int:
        """
        The middle point of this interval, simply ``left+(right-left)/2``.
        """
        return self.left + (self.right - self.left) / 2


class EdgeDiff(NamedTuple):
    interval: Interval
    edges_out: list
    edges_in: list


def store_tree_sequence(cls):
    wrapped_init = cls.__init__

    # Intercept the init to record the tree_sequence
    def new_init(self, *args, tree_sequence=None, **kwargs):
        builtins.object.__setattr__(self, "_tree_sequence", tree_sequence)
        wrapped_init(self, *args, **kwargs)

    cls.__init__ = new_init
    return cls


@store_tree_sequence
@metadata_module.lazy_decode()
@dataclass
class Individual(util.Dataclass):
    """
    An :ref:`individual <sec_individual_table_definition>` in a tree sequence.
    Since nodes correspond to genomes, individuals are associated with a collection
    of nodes (e.g., two nodes per diploid). See :ref:`sec_nodes_or_individuals`
    for more discussion of this distinction.

    Modifying the attributes in this class will have **no effect** on the
    underlying tree sequence data.
    """

    __slots__ = [
        "id",
        "flags",
        "location",
        "parents",
        "nodes",
        "metadata",
        "_tree_sequence",
    ]
    id: int  # noqa A003
    """
    The integer ID of this individual. Varies from 0 to
    :attr:`TreeSequence.num_individuals` - 1."""
    flags: int
    """
    The bitwise flags for this individual.
    """
    location: np.ndarray
    """
    The spatial location of this individual as a numpy array. The location is an empty
    array if no spatial location is defined.
    """
    parents: np.ndarray
    """
    The parent individual ids of this individual as a numpy array. The parents is an
    empty array if no parents are defined.
    """
    nodes: np.ndarray
    """
    The IDs of the nodes that are associated with this individual as
    a numpy array (dtype=np.int32). If no nodes are associated with the
    individual this array will be empty.
    """
    metadata: bytes | dict | None
    """
    The :ref:`metadata <sec_metadata_definition>`
    for this individual, decoded if a schema applies.
    """

    @property
    def population(self) -> int:
        populations = {self._tree_sequence.node(n).population for n in self.nodes}
        if len(populations) > 1:
            raise ValueError("Individual has nodes with mis-matched populations")
        if len(populations) == 0:
            return tskit.NULL
        return populations.pop()

    @property
    def time(self) -> int:
        times = {self._tree_sequence.node(n).time for n in self.nodes}
        if len(times) > 1:
            raise ValueError("Individual has nodes with mis-matched times")
        if len(times) == 0:
            return tskit.UNKNOWN_TIME
        return times.pop()

    # Custom eq for the numpy arrays
    def __eq__(self, other):
        return (
            self.id == other.id
            and self.flags == other.flags
            and np.array_equal(self.location, other.location)
            and np.array_equal(self.parents, other.parents)
            and np.array_equal(self.nodes, other.nodes)
            and self.metadata == other.metadata
        )


@metadata_module.lazy_decode()
@dataclass
class Node(util.Dataclass):
    """
    A :ref:`node <sec_node_table_definition>` in a tree sequence, corresponding
    to a single genome. The ``time`` and ``population`` are attributes of the
    ``Node``, rather than the ``Individual``, as discussed in
    :ref:`sec_nodes_or_individuals`.

    Modifying the attributes in this class will have **no effect** on the
    underlying tree sequence data.
    """

    __slots__ = ["id", "flags", "time", "population", "individual", "metadata"]
    id: int  # noqa A003
    """
    The integer ID of this node. Varies from 0 to :attr:`TreeSequence.num_nodes` - 1.
    """
    flags: int
    """
    The bitwise flags for this node.
    """
    time: float
    """
    The birth time of this node.
    """
    population: int
    """
    The integer ID of the population that this node was born in.
    """
    individual: int
    """
    The integer ID of the individual that this node was a part of.
    """
    metadata: bytes | dict | None
    """
    The :ref:`metadata <sec_metadata_definition>` for this node, decoded if a schema
    applies.
    """

    def is_sample(self):
        """
        Returns True if this node is a sample. This value is derived from the
        ``flag`` variable.

        :rtype: bool
        """
        return self.flags & NODE_IS_SAMPLE


@metadata_module.lazy_decode(own_init=True)
@dataclass
class Edge(util.Dataclass):
    """
    An :ref:`edge <sec_edge_table_definition>` in a tree sequence.

    Modifying the attributes in this class will have **no effect** on the
    underlying tree sequence data.
    """

    __slots__ = ["left", "right", "parent", "child", "metadata", "id"]
    left: float
    """
    The left coordinate of this edge.
    """
    right: float
    """
    The right coordinate of this edge.
    """
    parent: int
    """
    The integer ID of the parent node for this edge.
    To obtain further information about a node with a given ID, use
    :meth:`TreeSequence.node`.
    """
    child: int
    """
    The integer ID of the child node for this edge.
    To obtain further information about a node with a given ID, use
    :meth:`TreeSequence.node`.
    """
    metadata: bytes | dict | None
    """
    The :ref:`metadata <sec_metadata_definition>` for this edge, decoded if a schema
    applies.
    """
    id: int  # noqa A003
    """
    The integer ID of this edge. Varies from 0 to
    :attr:`TreeSequence.num_edges` - 1.
    """

    # Custom init to define default values with slots
    def __init__(
        self,
        left,
        right,
        parent,
        child,
        metadata=b"",
        id=None,  # noqa A002
        metadata_decoder=None,
    ):
        self.id = id
        self.left = left
        self.right = right
        self.parent = parent
        self.child = child
        self.metadata = metadata
        self._metadata_decoder = metadata_decoder

    @property
    def span(self):
        """
        Returns the span of this edge, i.e., the right position minus the left position

        :return: The span of this edge.
        :rtype: float
        """
        return self.right - self.left

    @property
    def interval(self):
        """
        Returns the left and right positions of this edge as an :class:`Interval`

        :return: The interval covered by this edge.
        :rtype: :class:`Interval`
        """
        return Interval(self.left, self.right)


@metadata_module.lazy_decode()
@dataclass
class Site(util.Dataclass):
    """
    A :ref:`site <sec_site_table_definition>` in a tree sequence.

    Modifying the attributes in this class will have **no effect** on the
    underlying tree sequence data.
    """

    __slots__ = ["id", "position", "ancestral_state", "mutations", "metadata"]
    id: int  # noqa A003
    """
    The integer ID of this site. Varies from 0 to :attr:`TreeSequence.num_sites` - 1.
    """
    position: float
    """
    The floating point location of this site in genome coordinates.
    Ranges from 0 (inclusive) to :attr:`TreeSequence.sequence_length` (exclusive).
    """
    ancestral_state: str
    """
    The ancestral state at this site (i.e., the state inherited by nodes, unless
    mutations occur).
    """
    mutations: np.ndarray
    """
    The list of mutations at this site. Mutations within a site are returned in the
    order they are specified in the underlying :class:`MutationTable`.
    """
    metadata: bytes | dict | None
    """
    The :ref:`metadata <sec_metadata_definition>` for this site, decoded if a schema
    applies.
    """

    # We need a custom eq for the numpy arrays
    def __eq__(self, other):
        return (
            isinstance(other, Site)
            and self.id == other.id
            and self.position == other.position
            and self.ancestral_state == other.ancestral_state
            and np.array_equal(self.mutations, other.mutations)
            and self.metadata == other.metadata
        )

    @property
    def alleles(self) -> set[str]:
        """
        Return the set of all the alleles defined at this site

        .. note::
            This deliberately returns an (unordered) *set* of the possible allelic
            states (as defined by the site's ancestral allele and its associated
            mutations). If you wish to obtain an (ordered) *list* of alleles, for
            example to translate the numeric genotypes at a site into allelic states,
            you should instead use ``.alleles`` attribute of the :class:`Variant` class,
            which unlike this attribute includes ``None`` as a state when there is
            missing data at a site.
        """
        return {self.ancestral_state} | {m.derived_state for m in self.mutations}


@metadata_module.lazy_decode()
@dataclass
class Mutation(util.Dataclass):
    """
    A :ref:`mutation <sec_mutation_table_definition>` in a tree sequence.

    Modifying the attributes in this class will have **no effect** on the
    underlying tree sequence data.
    """

    __slots__ = [
        "id",
        "site",
        "node",
        "derived_state",
        "parent",
        "metadata",
        "time",
        "edge",
    ]
    id: int  # noqa A003
    """
    The integer ID of this mutation. Varies from 0 to
    :attr:`TreeSequence.num_mutations` - 1.

    Modifying the attributes in this class will have **no effect** on the
    underlying tree sequence data.
    """
    site: int
    """
    The integer ID of the site that this mutation occurs at. To obtain
    further information about a site with a given ID use :meth:`TreeSequence.site`.
    """
    node: int
    """
    The integer ID of the first node that inherits this mutation.
    To obtain further information about a node with a given ID, use
    :meth:`TreeSequence.node`.
    """
    derived_state: str
    """
    The derived state for this mutation. This is the state
    inherited by nodes in the subtree rooted at this mutation's node, unless
    another mutation occurs.
    """
    parent: int
    """
    The integer ID of this mutation's parent mutation. When multiple
    mutations occur at a site along a path in the tree, mutations must
    record the mutation that is immediately above them. If the mutation does
    not have a parent, this is equal to the :data:`NULL` (-1).
    To obtain further information about a mutation with a given ID, use
    :meth:`TreeSequence.mutation`.
    """
    metadata: bytes | dict | None
    """
    The :ref:`metadata <sec_metadata_definition>` for this mutation, decoded if a schema
    applies.
    """
    time: float
    """
    The occurrence time of this mutation.
    """
    edge: int
    """
    The ID of the edge that this mutation is on.
    """

    # To get default values on slots we define a custom init
    def __init__(
        self,
        id=NULL,  # noqa A003
        site=NULL,
        node=NULL,
        time=UNKNOWN_TIME,
        derived_state=None,
        parent=NULL,
        metadata=b"",
        edge=NULL,
    ):
        self.id = id
        self.site = site
        self.node = node
        self.time = time
        self.derived_state = derived_state
        self.parent = parent
        self.metadata = metadata
        self.edge = edge

    # We need a custom eq to compare unknown times.
    def __eq__(self, other):
        return (
            isinstance(other, Mutation)
            and self.id == other.id
            and self.site == other.site
            and self.node == other.node
            and self.derived_state == other.derived_state
            and self.parent == other.parent
            and self.edge == other.edge
            and self.metadata == other.metadata
            and (
                self.time == other.time
                or (
                    util.is_unknown_time(self.time) and util.is_unknown_time(other.time)
                )
            )
        )


@metadata_module.lazy_decode()
@dataclass
class Migration(util.Dataclass):
    """
    A :ref:`migration <sec_migration_table_definition>` in a tree sequence.

    Modifying the attributes in this class will have **no effect** on the
    underlying tree sequence data.
    """

    __slots__ = ["left", "right", "node", "source", "dest", "time", "metadata", "id"]
    left: float
    """
    The left end of the genomic interval covered by this
    migration (inclusive).
    """
    right: float
    """
    The right end of the genomic interval covered by this migration
    (exclusive).
    """
    node: int
    """
    The integer ID of the node involved in this migration event.
    To obtain further information about a node with a given ID, use
    :meth:`TreeSequence.node`.
    """
    source: int
    """
    The source population ID.
    """
    dest: int
    """
    The destination population ID.
    """
    time: float
    """
    The time at which this migration occurred at.
    """
    metadata: bytes | dict | None
    """
    The :ref:`metadata <sec_metadata_definition>` for this migration, decoded if a schema
    applies.
    """
    id: int  # noqa A003
    """
    The integer ID of this mutation. Varies from 0 to
    :attr:`TreeSequence.num_mutations` - 1.
    """


@metadata_module.lazy_decode()
@dataclass
class Population(util.Dataclass):
    """
    A :ref:`population <sec_population_table_definition>` in a tree sequence.

    Modifying the attributes in this class will have **no effect** on the
    underlying tree sequence data.
    """

    __slots__ = ["id", "metadata"]
    id: int  # noqa A003
    """
    The integer ID of this population. Varies from 0 to
    :attr:`TreeSequence.num_populations` - 1.
    """
    metadata: bytes | dict | None
    """
    The :ref:`metadata <sec_metadata_definition>` for this population, decoded if a
    schema applies.
    """


@dataclass
class Edgeset(util.Dataclass):
    __slots__ = ["left", "right", "parent", "children"]
    left: int
    right: int
    parent: int
    children: np.ndarray

    # We need a custom eq for the numpy array
    def __eq__(self, other):
        return (
            isinstance(other, Edgeset)
            and self.left == other.left
            and self.right == other.right
            and self.parent == other.parent
            and np.array_equal(self.children, other.children)
        )


@dataclass
class Provenance(util.Dataclass):
    """
    A provenance entry in a tree sequence, detailing how this tree
    sequence was generated, or subsequent operations on it (see :ref:`sec_provenance`).
    """

    __slots__ = ["id", "timestamp", "record"]
    id: int  # noqa A003
    timestamp: str
    """
    The time that this entry was made
    """
    record: str
    """
    A JSON string giving details of the provenance (see :ref:`sec_provenance_example`
    for an example JSON string)
    """


class Tree:
    """
    A single tree in a :class:`TreeSequence`. Please see the
    :ref:`tutorials:sec_processing_trees` section for information
    on how efficiently access trees sequentially or obtain a list
    of individual trees in a tree sequence.

    The ``sample_lists`` parameter controls the features that are enabled
    for this tree. If ``sample_lists`` is True a more efficient algorithm is
    used in the :meth:`Tree.samples` method.

    The ``tracked_samples`` parameter can be used to efficiently count the
    number of samples in a given set that exist in a particular subtree
    using the :meth:`Tree.num_tracked_samples` method.

    The :class:`Tree` class is a state-machine which has a state
    corresponding to each of the trees in the parent tree sequence. We
    transition between these states by using the seek functions like
    :meth:`Tree.first`, :meth:`Tree.last`, :meth:`Tree.seek` and
    :meth:`Tree.seek_index`. There is one more state, the so-called "null"
    or "cleared" state. This is the state that a :class:`Tree` is in
    immediately after initialisation;  it has an index of -1, and no edges. We
    can also enter the null state by calling :meth:`Tree.next` on the last
    tree in a sequence, calling :meth:`Tree.prev` on the first tree in a
    sequence or calling calling the :meth:`Tree.clear` method at any time.

    The high-level TreeSequence seeking and iterations methods (e.g,
    :meth:`TreeSequence.trees`) are built on these low-level state-machine
    seek operations. We recommend these higher level operations for most
    users.

    :param TreeSequence tree_sequence: The parent tree sequence.
    :param list tracked_samples: The list of samples to be tracked and
        counted using the :meth:`Tree.num_tracked_samples` method.
    :param bool sample_lists: If True, provide more efficient access
        to the samples beneath a given node using the
        :meth:`Tree.samples` method.
    :param int root_threshold: The minimum number of samples that a node
        must be ancestral to for it to be in the list of roots. By default
        this is 1, so that isolated samples (representing missing data)
        are roots. To efficiently restrict the roots of the tree to
        those subtending meaningful topology, set this to 2. This value
        is only relevant when trees have multiple roots.
    :param bool sample_counts: Deprecated since 0.2.4.
    """

    def __init__(
        self,
        tree_sequence,
        tracked_samples=None,
        *,
        sample_lists=False,
        root_threshold=1,
        sample_counts=None,
    ):
        options = 0
        if sample_counts is not None:
            warnings.warn(
                "The sample_counts option is not supported since 0.2.4 "
                "and is ignored",
                RuntimeWarning,
                stacklevel=4,
            )
        if sample_lists:
            options |= _tskit.SAMPLE_LISTS
        kwargs = {"options": options}
        if root_threshold <= 0:
            raise ValueError("Root threshold must be greater than 0")
        if tracked_samples is not None:
            # TODO remove this when we allow numpy arrays in the low-level API.
            kwargs["tracked_samples"] = list(tracked_samples)

        self._tree_sequence = tree_sequence
        self._ll_tree = _tskit.Tree(tree_sequence.ll_tree_sequence, **kwargs)
        self._ll_tree.set_root_threshold(root_threshold)
        self._make_arrays()

    def copy(self):
        """
        Returns a deep copy of this tree. The returned tree will have identical state
        to this tree.

        :return: A copy of this tree.
        :rtype: Tree
        """
        copy = type(self).__new__(type(self))
        copy._tree_sequence = self._tree_sequence
        copy._ll_tree = self._ll_tree.copy()
        copy._make_arrays()
        return copy

    # TODO make this method public and document it.
    # Note that this probably does not cover all corner cases correctly
    # https://github.com/tskit-dev/tskit/issues/1908
    def _has_isolated_samples(self):
        # TODO Is this definition correct for a single-node tree sequence?
        for root in self.roots:
            # If the root has no children then it must be a sample
            if self.left_child(root) == NULL:
                return True
        return False

    def _make_arrays(self):
        # Store the low-level arrays for efficiency. There's no real overhead
        # in this, because the refer to the same underlying memory as the
        # tree object.
        self._parent_array = self._ll_tree.parent_array
        self._left_child_array = self._ll_tree.left_child_array
        self._right_child_array = self._ll_tree.right_child_array
        self._left_sib_array = self._ll_tree.left_sib_array
        self._right_sib_array = self._ll_tree.right_sib_array
        self._num_children_array = self._ll_tree.num_children_array
        self._edge_array = self._ll_tree.edge_array

    @property
    def tree_sequence(self):
        """
        Returns the tree sequence that this tree is from.

        :return: The parent tree sequence for this tree.
        :rtype: :class:`TreeSequence`
        """
        return self._tree_sequence

    @property
    def root_threshold(self):
        """
        Returns the minimum number of samples that a node must be an ancestor
        of to be considered a potential root. This can be set, for example, when
        calling the :meth:`TreeSequence.trees` iterator.

        :return: The root threshold.
        :rtype: :class:`TreeSequence`
        """
        return self._ll_tree.get_root_threshold()

    def __eq__(self, other):
        ret = False
        if type(other) is type(self):
            ret = bool(self._ll_tree.equals(other._ll_tree))
        return ret

    def __ne__(self, other):
        return not self.__eq__(other)

    def first(self):
        """
        Seeks to the first tree in the sequence. This can be called whether
        the tree is in the null state or not.
        """
        self._ll_tree.first()

    def last(self):
        """
        Seeks to the last tree in the sequence. This can be called whether
        the tree is in the null state or not.
        """
        self._ll_tree.last()

    def next(self):  # noqa A002
        """
        Seeks to the next tree in the sequence. If the tree is in the initial
        null state we seek to the first tree (equivalent to calling :meth:`~Tree.first`).
        Calling ``next`` on the last tree in the sequence results in the tree
        being cleared back into the null initial state (equivalent to calling
        :meth:`~Tree.clear`). The return value of the function indicates whether the
        tree is in a non-null state, and can be used to loop over the trees::

            # Iterate over the trees from left-to-right
            tree = tskit.Tree(tree_sequence)
            while tree.next()
                # Do something with the tree.
                print(tree.index)
            # tree is now back in the null state.

        :return: True if the tree has been transformed into one of the trees
            in the sequence; False if the tree has been transformed into the
            null state.
        :rtype: bool
        """
        return bool(self._ll_tree.next())

    def prev(self):
        """
        Seeks to the previous tree in the sequence. If the tree is in the initial
        null state we seek to the last tree (equivalent to calling :meth:`~Tree.last`).
        Calling ``prev`` on the first tree in the sequence results in the tree
        being cleared back into the null initial state (equivalent to calling
        :meth:`~Tree.clear`). The return value of the function indicates whether the
        tree is in a non-null state, and can be used to loop over the trees::

            # Iterate over the trees from right-to-left
            tree = tskit.Tree(tree_sequence)
            while tree.prev()
                # Do something with the tree.
                print(tree.index)
            # tree is now back in the null state.

        :return: True if the tree has been transformed into one of the trees
            in the sequence; False if the tree has been transformed into the
            null state.
        :rtype: bool
        """
        return bool(self._ll_tree.prev())

    def clear(self):
        """
        Resets this tree back to the initial null state. Calling this method
        on a tree already in the null state has no effect.
        """
        self._ll_tree.clear()

    def seek_index(self, index):
        """
        Sets the state to represent the tree at the specified
        index in the parent tree sequence. Negative indexes following the
        standard Python conventions are allowed, i.e., ``index=-1`` will
        seek to the last tree in the sequence.

        .. include:: substitutions/linear_traversal_warning.rst

        :param int index: The tree index to seek to.
        :raises IndexError: If an index outside the acceptable range is provided.
        """
        num_trees = self.tree_sequence.num_trees
        if index < 0:
            index += num_trees
        if index < 0 or index >= num_trees:
            raise IndexError("Index out of bounds")
        self._ll_tree.seek_index(index)

    def seek(self, position):
        """
        Sets the state to represent the tree that covers the specified
        position in the parent tree sequence. After a successful return
        of this method we have ``tree.interval.left`` <= ``position``
        < ``tree.interval.right``.

        .. include:: substitutions/linear_traversal_warning.rst

        :param float position: The position along the sequence length to
            seek to.
        :raises ValueError: If 0 < position or position >=
            :attr:`TreeSequence.sequence_length`.
        """
        if position < 0 or position >= self.tree_sequence.sequence_length:
            raise ValueError("Position out of bounds")
        self._ll_tree.seek(position)

    def rank(self) -> tskit.Rank:
        """
        Produce the rank of this tree in the enumeration of all leaf-labelled
        trees of n leaves. See the :ref:`sec_tree_ranks` section for
        details on ranking and unranking trees.

        :raises ValueError: If the tree has multiple roots.
        """
        return combinatorics.RankTree.from_tsk_tree(self).rank()

    @staticmethod
    def unrank(num_leaves, rank, *, span=1, branch_length=1) -> Tree:
        """
        Reconstruct the tree of the given ``rank``
        (see :meth:`tskit.Tree.rank`) with ``num_leaves`` leaves.
        The labels and times of internal nodes are assigned by a postorder
        traversal of the nodes, such that the time of each internal node
        is the maximum time of its children plus the specified ``branch_length``.
        The time of each leaf is 0.

        See the :ref:`sec_tree_ranks` section for details on ranking and
        unranking trees and what constitutes valid ranks.

        :param int num_leaves: The number of leaves of the tree to generate.
        :param tuple(int) rank: The rank of the tree to generate.
        :param float span: The genomic span of the returned tree. The tree will cover
            the interval :math:`[0, \\text{span})` and the :attr:`~Tree.tree_sequence`
            from which the tree is taken will have its
            :attr:`~tskit.TreeSequence.sequence_length` equal to ``span``.
        :param: float branch_length: The minimum length of a branch in this tree.
        :raises ValueError: If the given rank is out of bounds for trees
            with ``num_leaves`` leaves.
        """
        rank_tree = combinatorics.RankTree.unrank(num_leaves, rank)
        return rank_tree.to_tsk_tree(span=span, branch_length=branch_length)

    def count_topologies(self, sample_sets=None) -> tskit.TopologyCounter:
        """
        Calculates the distribution of embedded topologies for every combination
        of the sample sets in ``sample_sets``. ``sample_sets`` defaults to all
        samples in the tree grouped by population.

        ``sample_sets`` need not include all samples but must be pairwise disjoint.

        The returned object is a :class:`tskit.TopologyCounter` that contains
        counts of topologies per combination of sample sets. For example,

        >>> topology_counter = tree.count_topologies()
        >>> rank, count = topology_counter[0, 1, 2].most_common(1)[0]

        produces the most common tree topology, with populations 0, 1
        and 2 as its tips, according to the genealogies of those
        populations' samples in this tree.

        The counts for each topology in the :class:`tskit.TopologyCounter`
        are absolute counts that we would get if we were to select all
        combinations of samples from the relevant sample sets.
        For sample sets :math:`[s_0, s_1, ..., s_n]`, the total number of
        topologies for those sample sets is equal to
        :math:`|s_0| * |s_1| * ... * |s_n|`, so the counts in the counter
        ``topology_counter[0, 1, ..., n]`` should sum to
        :math:`|s_0| * |s_1| * ... * |s_n|`.

        To convert the topology counts to probabilities, divide by the total
        possible number of sample combinations from the sample sets in question::

            >>> set_sizes = [len(sample_set) for sample_set in sample_sets]
            >>> p = count / (set_sizes[0] * set_sizes[1] * set_sizes[2])

        .. warning:: The interface for this method is preliminary and may be subject to
            backwards incompatible changes in the near future.

        :param list sample_sets: A list of lists of Node IDs, specifying the
            groups of nodes to compute the statistic with.
            Defaults to all samples grouped by population.
        :raises ValueError: If nodes in ``sample_sets`` are invalid or are
            internal samples.
        """
        if sample_sets is None:
            sample_sets = [
                self.tree_sequence.samples(population=pop.id)
                for pop in self.tree_sequence.populations()
            ]

        return combinatorics.tree_count_topologies(self, sample_sets)

    def get_branch_length(self, u):
        # Deprecated alias for branch_length
        return self.branch_length(u)

    def branch_length(self, u):
        """
        Returns the length of the branch (in units of time) joining the
        specified node to its parent. This is equivalent to

        >>> tree.time(tree.parent(u)) - tree.time(u)

        The branch length for a node that has no parent (e.g., a root) is
        defined as zero.

        Note that this is not related to the property `.length` which
        is a deprecated alias for the genomic :attr:`~Tree.span` covered by a tree.

        :param int u: The node of interest.
        :return: The branch length from u to its parent.
        :rtype: float
        """
        ret = 0
        parent = self.parent(u)
        if parent != NULL:
            ret = self.time(parent) - self.time(u)
        return ret

    def get_total_branch_length(self):
        # Deprecated alias for total_branch_length
        return self.total_branch_length

    @property
    def total_branch_length(self):
        """
        Returns the sum of all the branch lengths in this tree (in
        units of time). This is equivalent to

        >>> sum(tree.branch_length(u) for u in tree.nodes())

        Note that the branch lengths for root nodes are defined as zero.

        As this is defined by a traversal of the tree, technically we
        return the sum of all branch lengths that are reachable from
        roots. Thus, this is the total length of all branches that are connected
        to at least one sample. This distinction is only important
        in tree sequences that contain 'dead branches', i.e., those
        that define topology that is not connected to a tree root
        (see :ref:`sec_data_model_tree_dead_leaves_and_branches`)

        :return: The sum of lengths of branches in this tree.
        :rtype: float
        """
        return self._ll_tree.get_total_branch_length()

    def get_mrca(self, u, v):
        # Deprecated alias for mrca
        return self.mrca(u, v)

    def mrca(self, *args):
        """
        Returns the most recent common ancestor of the specified nodes.

        :param int `*args`: input node IDs, at least 2 arguments are required.
        :return: The node ID of the most recent common ancestor of the
            input nodes, or :data:`tskit.NULL` if the nodes do not share
            a common ancestor in the tree.
        :rtype: int
        """
        if len(args) < 2:
            raise ValueError("Must supply at least two arguments")
        mrca = args[0]
        for node in args[1:]:
            mrca = self._ll_tree.get_mrca(mrca, node)
            if mrca == tskit.NULL:
                break
        return mrca

    def get_tmrca(self, u, v):
        # Deprecated alias for tmrca
        return self.tmrca(u, v)

    def tmrca(self, *args):
        """
        Returns the time of the most recent common ancestor of the specified
        nodes. This is equivalent to::

            >>> tree.time(tree.mrca(*args))

        .. note::
            If you are using this method to calculate average tmrca values along the
            genome between pairs of sample nodes, for efficiency reasons you should
            instead consider the ``mode="branch"`` option of the
            :meth:`TreeSequence.divergence` or :meth:`TreeSequence.diversity` methods.
            Since these calculate the average branch length between pairs of sample
            nodes, for samples at time 0 the resulting statistics will be exactly
            twice the tmrca value.

        :param `*args`: input node IDs, at least 2 arguments are required.
        :return: The time of the most recent common ancestor of all the nodes.
        :rtype: float
        :raises ValueError: If the nodes do not share a single common ancestor in this
            tree (i.e., if ``tree.mrca(*args) == tskit.NULL``)
        """
        mrca = self.mrca(*args)
        if mrca == tskit.NULL:
            raise ValueError(f"Nodes {args} do not share a common ancestor in the tree")
        return self.get_time(mrca)

    def get_parent(self, u):
        # Deprecated alias for parent
        return self.parent(u)

    def parent(self, u):
        """
        Returns the parent of the specified node. Returns
        :data:`tskit.NULL` if u is a root or is not a node in
        the current tree.

        :param int u: The node of interest.
        :return: The parent of u.
        :rtype: int
        """
        return self._ll_tree.get_parent(u)

    @property
    def parent_array(self):
        """
        A numpy array (dtype=np.int32) encoding the parent of each node
        in this tree, such that ``tree.parent_array[u] == tree.parent(u)``
        for all ``0 <= u <= ts.num_nodes``. See the :meth:`~.parent`
        method for details on the semantics of tree parents and the
        :ref:`sec_data_model_tree_structure` section for information on the
        quintuply linked tree encoding.

        .. include:: substitutions/virtual_root_array_note.rst

        .. include:: substitutions/tree_array_warning.rst
        """
        return self._parent_array

    def ancestors(self, u):
        """
        Returns an iterator over the ancestors of node ``u`` in this tree
        (i.e. the chain of parents from ``u`` to the root).
        """
        u = self.parent(u)
        while u != -1:
            yield u
            u = self.parent(u)

    # Quintuply linked tree structure.

    def left_child(self, u):
        """
        Returns the leftmost child of the specified node. Returns
        :data:`tskit.NULL` if u is a leaf or is not a node in
        the current tree. The left-to-right ordering of children
        is arbitrary and should not be depended on; see the
        :ref:`data model <sec_data_model_tree_structure>` section
        for details.

        This is a low-level method giving access to the quintuply linked
        tree structure in memory; the :meth:`.children` method is a more
        convenient way to obtain the children of a given node.

        :param int u: The node of interest.
        :return: The leftmost child of u.
        :rtype: int
        """
        return self._ll_tree.get_left_child(u)

    @property
    def left_child_array(self):
        """
        A numpy array (dtype=np.int32) encoding the left child of each node
        in this tree, such that ``tree.left_child_array[u] == tree.left_child(u)``
        for all ``0 <= u <= ts.num_nodes``. See the :meth:`~.left_child`
        method for details on the semantics of tree left_child and the
        :ref:`sec_data_model_tree_structure` section for information on the
        quintuply linked tree encoding.

        .. include:: substitutions/virtual_root_array_note.rst

        .. include:: substitutions/tree_array_warning.rst
        """
        return self._left_child_array

    def right_child(self, u):
        """
        Returns the rightmost child of the specified node. Returns
        :data:`tskit.NULL` if u is a leaf or is not a node in
        the current tree. The left-to-right ordering of children
        is arbitrary and should not be depended on; see the
        :ref:`data model <sec_data_model_tree_structure>` section
        for details.

        This is a low-level method giving access to the quintuply linked
        tree structure in memory; the :meth:`.children` method is a more
        convenient way to obtain the children of a given node.

        :param int u: The node of interest.
        :return: The rightmost child of u.
        :rtype: int
        """
        return self._ll_tree.get_right_child(u)

    @property
    def right_child_array(self):
        """
        A numpy array (dtype=np.int32) encoding the right child of each node
        in this tree, such that ``tree.right_child_array[u] == tree.right_child(u)``
        for all ``0 <= u <= ts.num_nodes``. See the :meth:`~.right_child`
        method for details on the semantics of tree right_child and the
        :ref:`sec_data_model_tree_structure` section for information on the
        quintuply linked tree encoding.

        .. include:: substitutions/virtual_root_array_note.rst

        .. include:: substitutions/tree_array_warning.rst
        """
        return self._right_child_array

    def left_sib(self, u):
        """
        Returns the sibling node to the left of u, or :data:`tskit.NULL`
        if u does not have a left sibling.
        The left-to-right ordering of children
        is arbitrary and should not be depended on; see the
        :ref:`data model <sec_data_model_tree_structure>` section
        for details.

        :param int u: The node of interest.
        :return: The sibling node to the left of u.
        :rtype: int
        """
        return self._ll_tree.get_left_sib(u)

    @property
    def left_sib_array(self):
        """
        A numpy array (dtype=np.int32) encoding the left sib of each node
        in this tree, such that ``tree.left_sib_array[u] == tree.left_sib(u)``
        for all ``0 <= u <= ts.num_nodes``. See the :meth:`~.left_sib`
        method for details on the semantics of tree left_sib and the
        :ref:`sec_data_model_tree_structure` section for information on the
        quintuply linked tree encoding.

        .. include:: substitutions/virtual_root_array_note.rst

        .. include:: substitutions/tree_array_warning.rst
        """
        return self._left_sib_array

    def right_sib(self, u):
        """
        Returns the sibling node to the right of u, or :data:`tskit.NULL`
        if u does not have a right sibling.
        The left-to-right ordering of children
        is arbitrary and should not be depended on; see the
        :ref:`data model <sec_data_model_tree_structure>` section
        for details.

        :param int u: The node of interest.
        :return: The sibling node to the right of u.
        :rtype: int
        """
        return self._ll_tree.get_right_sib(u)

    @property
    def right_sib_array(self):
        """
        A numpy array (dtype=np.int32) encoding the right sib of each node
        in this tree, such that ``tree.right_sib_array[u] == tree.right_sib(u)``
        for all ``0 <= u <= ts.num_nodes``. See the :meth:`~.right_sib`
        method for details on the semantics of tree right_sib and the
        :ref:`sec_data_model_tree_structure` section for information on the
        quintuply linked tree encoding.

        .. include:: substitutions/virtual_root_array_note.rst

        .. include:: substitutions/tree_array_warning.rst
        """
        return self._right_sib_array

    def siblings(self, u):
        """
        Returns the sibling(s) of the specified node ``u`` as a tuple of integer
        node IDs. If ``u`` has no siblings or is not a node in the current tree,
        returns an empty tuple. If ``u`` is the root of a single-root tree,
        returns an empty tuple; if ``u`` is the root of a multi-root tree,
        returns the other roots (note all the roots are related by the virtual root).
        If ``u`` is the virtual root (which has no siblings), returns an empty tuple.
        If ``u`` is an isolated node, whether it has siblings or not depends on
        whether it is a sample or non-sample node; if it is a sample node,
        returns the root(s) of the tree, otherwise, returns an empty tuple.
        The ordering of siblings  is arbitrary and should not be depended on;
        see the :ref:`data model <sec_data_model_tree_structure>` section for details.

        :param int u: The node of interest.
        :return: The siblings of ``u``.
        :rtype: tuple(int)
        """
        if u == self.virtual_root:
            return tuple()
        parent = self.parent(u)
        if self.is_root(u):
            parent = self.virtual_root
        if parent != tskit.NULL:
            return tuple(v for v in self.children(parent) if u != v)
        return tuple()

    @property
    def num_children_array(self):
        """
        A numpy array (dtype=np.int32) encoding the number of children of
        each node in this tree, such that
        ``tree.num_children_array[u] == tree.num_children(u)`` for all
        ``0 <= u <= ts.num_nodes``. See the :meth:`~.num_children`
        method for details on the semantics of tree num_children and the
        :ref:`sec_data_model_tree_structure` section for information on the
        quintuply linked tree encoding.

        .. include:: substitutions/virtual_root_array_note.rst

        .. include:: substitutions/tree_array_warning.rst
        """
        return self._num_children_array

    def edge(self, u):
        """
        Returns the id of the edge encoding the relationship between ``u``
        and its parent, or :data:`tskit.NULL` if ``u`` is a root, virtual root
        or is not a node in the current tree.

        :param int u: The node of interest.
        :return: Id of edge connecting u to its parent.
        :rtype: int
        """
        return self._ll_tree.get_edge(u)

    @property
    def edge_array(self):
        """
        A numpy array (dtype=np.int32) of edge ids encoding the relationship
        between the child node ``u`` and its parent, such that
        ``tree.edge_array[u] == tree.edge(u)`` for all
        ``0 <= u <= ts.num_nodes``. See the :meth:`~.edge`
        method for details on the semantics of tree edge and the
        :ref:`sec_data_model_tree_structure` section for information on the
        quintuply linked tree encoding.

        .. include:: substitutions/virtual_root_array_note.rst

        .. include:: substitutions/tree_array_warning.rst
        """
        return self._edge_array

    # Sample list.

    def left_sample(self, u):
        return self._ll_tree.get_left_sample(u)

    def right_sample(self, u):
        return self._ll_tree.get_right_sample(u)

    def next_sample(self, u):
        return self._ll_tree.get_next_sample(u)

    @property
    def virtual_root(self):
        """
        The ID of the virtual root in this tree. This is equal to
        :attr:`TreeSequence.num_nodes`.

        Please see the :ref:`tree roots <sec_data_model_tree_roots>`
        section for more details.
        """
        return self._ll_tree.get_virtual_root()

    @property
    def num_edges(self):
        """
        The total number of edges in this tree. This is equal to the
        number of tree sequence edges that intersect with this tree's
        genomic interval.

        Note that this may be greater than the number of branches that
        are reachable from the tree's roots, since we can have topology
        that is not associated with any samples.
        """
        return self._ll_tree.get_num_edges()

    @property
    def left_root(self):
        """
        The leftmost root in this tree. If there are multiple roots
        in this tree, they are siblings of this node, and so we can
        use :meth:`.right_sib` to iterate over all roots:

        .. code-block:: python

            u = tree.left_root
            while u != tskit.NULL:
                print("Root:", u)
                u = tree.right_sib(u)

        The left-to-right ordering of roots is arbitrary and should
        not be depended on; see the
        :ref:`data model <sec_data_model_tree_structure>`
        section for details.

        This is a low-level method giving access to the quintuply linked
        tree structure in memory; the :attr:`~Tree.roots` attribute is a more
        convenient way to obtain the roots of a tree. If you are assuming
        that there is a single root in the tree you should use the
        :attr:`~Tree.root` property.

        .. warning:: Do not use this property if you are assuming that there
            is a single root in trees that are being processed. The
            :attr:`~Tree.root` property should be used in this case, as it will
            raise an error when multiple roots exists.

        :rtype: int
        """
        return self.left_child(self.virtual_root)

    @property
    def right_root(self):
        return self.right_child(self.virtual_root)

    def get_children(self, u):
        # Deprecated alias for self.children
        return self.children(u)

    def children(self, u):
        """
        Returns the children of the specified node ``u`` as a tuple of integer node IDs.
        If ``u`` is a leaf, return the empty tuple. The ordering of children
        is arbitrary and should not be depended on; see the
        :ref:`data model <sec_data_model_tree_structure>` section
        for details.

        :param int u: The node of interest.
        :return: The children of ``u`` as a tuple of integers
        :rtype: tuple(int)
        """
        return self._ll_tree.get_children(u)

    def get_time(self, u):
        # Deprecated alias for self.time
        return self.time(u)

    def time(self, u):
        """
        Returns the time of the specified node. This is equivalently
        to ``tree.tree_sequence.node(u).time`` except for the special
        case of the tree's :ref:`virtual root <sec_data_model_tree_roots>`,
        which is defined as positive infinity.

        :param int u: The node of interest.
        :return: The time of u.
        :rtype: float
        """
        return self._ll_tree.get_time(u)

    def depth(self, u):
        """
        Returns the number of nodes on the path from ``u`` to a
        root, not including ``u``. Thus, the depth of a root is
        zero.

        As a special case, the depth of the :ref:`virtual root
        <sec_data_model_tree_roots>` is defined as -1.

        :param int u: The node of interest.
        :return: The depth of u.
        :rtype: int
        """
        return self._ll_tree.depth(u)

    def get_population(self, u):
        # Deprecated alias for self.population
        return self.population(u)

    def population(self, u):
        """
        Returns the population associated with the specified node.
        Equivalent to ``tree.tree_sequence.node(u).population``.

        :param int u: The node of interest.
        :return: The ID of the population associated with node u.
        :rtype: int
        """
        return self._ll_tree.get_population(u)

    def is_internal(self, u):
        """
        Returns True if the specified node is not a leaf. A node is internal
        if it has one or more children in the current tree.

        :param int u: The node of interest.
        :return: True if u is not a leaf node.
        :rtype: bool
        """
        return not self.is_leaf(u)

    def is_leaf(self, u):
        """
        Returns True if the specified node is a leaf. A node :math:`u` is a
        leaf if it has zero children.

        .. note::
            :math:`u` can be any node in the entire tree sequence, including ones
            which are not connected via branches to a root node of the tree (and which
            are therefore not conventionally considered part of the tree). Indeed, if
            there are many trees in the tree sequence, it is common for the majority of
            non-sample nodes to be :meth:`isolated<is_isolated>` in any one
            tree. By the definition above, this method will return ``True`` for such
            a tree when a node of this sort is specified. Such nodes can be thought of
            as "dead leaves", see :ref:`sec_data_model_tree_dead_leaves_and_branches`.

        :param int u: The node of interest.
        :return: True if u is a leaf node.
        :rtype: bool
        """
        return len(self.children(u)) == 0

    def is_isolated(self, u):
        """
        Returns True if the specified node is isolated in this tree: that is
        it has no parents and no children (note that all isolated nodes in the tree
        are therefore also :meth:`leaves<Tree.is_leaf>`). Sample nodes that are isolated
        and have no mutations above them are used to represent
        :ref:`missing data<sec_data_model_missing_data>`.

        :param int u: The node of interest.
        :return: True if u is an isolated node.
        :rtype: bool
        """
        return self.num_children(u) == 0 and self.parent(u) == NULL

    def is_sample(self, u):
        """
        Returns True if the specified node is a sample. A node :math:`u` is a
        sample if it has been marked as a sample in the parent tree sequence.

        :param int u: The node of interest.
        :return: True if u is a sample.
        :rtype: bool
        """
        return bool(self._ll_tree.is_sample(u))

    def is_descendant(self, u, v):
        """
        Returns True if the specified node u is a descendant of node v and False
        otherwise. A node :math:`u` is a descendant of another node :math:`v` if
        :math:`v` is on the path from :math:`u` to root. A node is considered
        to be a descendant of itself, so ``tree.is_descendant(u, u)`` will be
        True for any valid node.

        :param int u: The descendant node.
        :param int v: The ancestral node.
        :return: True if u is a descendant of v.
        :rtype: bool
        :raises ValueError: If u or v are not valid node IDs.
        """
        return bool(self._ll_tree.is_descendant(u, v))

    @property
    def num_nodes(self):
        """
        Returns the number of nodes in the :class:`TreeSequence` this tree is in.
        Equivalent to ``tree.tree_sequence.num_nodes``.

        .. deprecated:: 0.4
            Use :attr:`Tree.tree_sequence.num_nodes<TreeSequence.num_nodes>` if you want
            the number of nodes in the entire tree sequence, or
            ``len(tree.preorder())`` to find the number of nodes that are
            reachable from all roots in this tree.

        :rtype: int

        """
        warnings.warn(
            "This property is a deprecated alias for Tree.tree_sequence.num_nodes "
            "and will be removed in the future. To obtain the number of nodes "
            "in the topology of the current tree (i.e. reachable from the roots) "
            "use len(tree.preorder()).",
            FutureWarning,
            stacklevel=4,
        )
        return self.tree_sequence.num_nodes

    @property
    def num_roots(self):
        """
        The number of roots in this tree, as defined in the :attr:`~Tree.roots`
        attribute.

        Only requires O(number of roots) time.

        :rtype: int
        """
        return self._ll_tree.get_num_roots()

    @property
    def has_single_root(self):
        """
        ``True`` if this tree has a single root, ``False`` otherwise.
        Equivalent to tree.num_roots == 1. This is a O(1) operation.

        :rtype: bool
        """
        root = self.left_root
        if root != NULL and self.right_sib(root) == NULL:
            return True
        return False

    @property
    def has_multiple_roots(self):
        """
        ``True`` if this tree has more than one root, ``False`` otherwise.
        Equivalent to tree.num_roots > 1. This is a O(1) operation.

        :rtype: bool
        """
        root = self.left_root
        if root != NULL and self.right_sib(root) != NULL:
            return True
        return False

    @property
    def roots(self):
        """
        The list of roots in this tree. A root is defined as a unique endpoint of the
        paths starting at samples, subject to the condition that it is connected to at
        least :attr:`root_threshold` samples. We can define the set of roots as follows:

        .. code-block:: python

            roots = set()
            for u in tree_sequence.samples():
                while tree.parent(u) != tskit.NULL:
                    u = tree.parent(u)
                if tree.num_samples(u) >= tree.root_threshold:
                    roots.add(u)
            # roots is now the set of all roots in this tree.
            assert sorted(roots) == sorted(tree.roots)

        The roots of the tree are returned in a list, in no particular order.

        Only requires O(number of roots) time.

        .. note::
            In trees with large amounts of :ref:`sec_data_model_missing_data`,
            for example where a region of the genome lacks any ancestral information,
            there can be a very large number of roots, potentially all the samples
            in the tree sequence.

        :return: The list of roots in this tree.
        :rtype: list
        """
        roots = []
        u = self.left_root
        while u != NULL:
            roots.append(u)
            u = self.right_sib(u)
        return roots

    def get_root(self):
        # Deprecated alias for self.root
        return self.root

    @property
    def root(self):
        """
        The root of this tree. If the tree contains multiple roots, a ValueError is
        raised indicating that the :attr:`~Tree.roots` attribute should be used instead.

        :return: The root node.
        :rtype: int
        :raises ValueError: if this tree contains more than one root.
        """
        if self.has_multiple_roots:
            raise ValueError("More than one root exists. Use tree.roots instead")
        return self.left_root

    def is_root(self, u) -> bool:
        """
        Returns ``True`` if the specified node is a root in this tree (see
        :attr:`~Tree.roots` for the definition of a root). This is exactly equivalent to
        finding the node ID in :attr:`~Tree.roots`, but is more efficient for trees
        with large numbers of roots, such as in regions with extensive
        :ref:`sec_data_model_missing_data`.  Note that ``False`` is returned for all
        other nodes, including :ref:`isolated<sec_data_model_tree_isolated_nodes>`
        non-sample nodes which are not found in the topology of the current tree.

        :param int u: The node of interest.
        :return: ``True`` if u is a root.
        """
        return (
            self.num_samples(u) >= self.root_threshold and self.parent(u) == tskit.NULL
        )

    def get_index(self):
        # Deprecated alias for self.index
        return self.index

    @property
    def index(self):
        """
        Returns the index this tree occupies in the parent tree sequence.
        This index is zero based, so the first tree in the sequence has index 0.

        :return: The index of this tree.
        :rtype: int
        """
        return self._ll_tree.get_index()

    def get_interval(self):
        # Deprecated alias for self.interval
        return self.interval

    @property
    def interval(self):
        """
        Returns the coordinates of the genomic interval that this tree
        represents the history of. The interval is returned as a tuple
        :math:`(l, r)` and is a half-open interval such that the left
        coordinate is inclusive and the right coordinate is exclusive. This
        tree therefore applies to all genomic locations :math:`x` such that
        :math:`l \\leq x < r`.

        :return: A named tuple (l, r) representing the left-most (inclusive)
            and right-most (exclusive) coordinates of the genomic region
            covered by this tree. The coordinates can be accessed by index
            (``0`` or ``1``) or equivalently by name (``.left`` or ``.right``)
        :rtype: Interval
        """
        return Interval(self._ll_tree.get_left(), self._ll_tree.get_right())

    def get_length(self):
        # Deprecated alias for self.span
        return self.length

    @property
    def length(self):
        # Deprecated alias for self.span
        return self.span

    @property
    def span(self):
        """
        Returns the genomic distance that this tree spans.
        This is defined as :math:`r - l`, where :math:`(l, r)` is the genomic
        interval returned by :attr:`~Tree.interval`.

        :return: The genomic distance covered by this tree.
        :rtype: float
        """
        return self.interval.span

    @property
    def mid(self):
        """
        Returns the midpoint of the genomic interval that this tree represents
        the history of. This is defined as :math:`(l + (r - l) / 2)`, where
        :math:`(l, r)` is the genomic interval returned by
        :attr:`~Tree.interval`.

        :return: The genomic distance covered by this tree.
        :rtype: float
        """
        return self.interval.mid

    def get_sample_size(self):
        # Deprecated alias for self.sample_size
        return self.sample_size

    @property
    def sample_size(self):
        # Deliberately undocumented but kept for backwards compatibility.
        # The proper way to access this is via tree.tree_sequence.num_samples
        return self._ll_tree.get_sample_size()

    def draw_text(
        self,
        orientation=None,
        *,
        node_labels=None,
        max_time=None,
        use_ascii=False,
        order=None,
    ):
        """
        Create a text representation of a tree.

        :param str orientation: one of ``"top"``, ``"left"``, ``"bottom"``, or
            ``"right"``, specifying the margin on which the root is placed. Specifying
            ``"left"`` or ``"right"`` will lead to time being shown on the x axis (i.e.
            a "horizontal" tree. If ``None`` (default) use the standard coalescent
            arrangement of a vertical tree with recent nodes at the bottom of the plot
            and older nodes above.
        :param dict node_labels: If specified, show custom labels for the nodes
            that are present in the map. Any nodes not specified in the map will
            not have a node label.
        :param str max_time: If equal to ``"tree"`` (the default), the maximum time
            is set to be that of the oldest root in the tree. If equal to ``"ts"`` the
            maximum time is set to be the time of the oldest root in the tree
            sequence; this is useful when drawing trees from the same tree sequence as it
            ensures that node heights are consistent.
        :param bool use_ascii: If ``False`` (default) then use unicode
            `box drawing characters \
<https://en.wikipedia.org/wiki/Box-drawing_character>`_
            to render the tree. If ``True``, use plain ascii characters, which look
            cruder but are less susceptible to misalignment or font substitution.
            Alternatively, if you are having alignment problems with Unicode, you can try
            out the solution documented `here \
<https://github.com/tskit-dev/tskit/issues/189#issuecomment-499114811>`_.
        :param str order: The left-to-right ordering of child nodes in the drawn tree.
            This can be either: ``"minlex"``, which minimises the differences
            between adjacent trees (see also the ``"minlex_postorder"`` traversal
            order for the :meth:`.nodes` method); or ``"tree"`` which draws trees
            in the left-to-right order defined by the
            :ref:`quintuply linked tree structure <sec_data_model_tree_structure>`.
            If not specified or None, this defaults to ``"minlex"``.

        :return: A text representation of a tree.
        :rtype: str
        """
        orientation = drawing.check_orientation(orientation)
        if orientation in (drawing.LEFT, drawing.RIGHT):
            text_tree = drawing.HorizontalTextTree(
                self,
                orientation=orientation,
                node_labels=node_labels,
                max_time=max_time,
                use_ascii=use_ascii,
                order=order,
            )
        else:
            text_tree = drawing.VerticalTextTree(
                self,
                orientation=orientation,
                node_labels=node_labels,
                max_time=max_time,
                use_ascii=use_ascii,
                order=order,
            )
        return str(text_tree)

    def draw_svg(
        self,
        path=None,
        *,
        size=None,
        time_scale=None,
        tree_height_scale=None,
        title=None,
        max_time=None,
        min_time=None,
        max_tree_height=None,
        node_labels=None,
        mutation_labels=None,
        node_titles=None,
        mutation_titles=None,
        root_svg_attributes=None,
        style=None,
        order=None,
        force_root_branch=None,
        symbol_size=None,
        x_axis=None,
        x_label=None,
        x_regions=None,
        y_axis=None,
        y_label=None,
        y_ticks=None,
        y_gridlines=None,
        all_edge_mutations=None,
        omit_sites=None,
        canvas_size=None,
        **kwargs,
    ):
        """
        Return an SVG representation of a single tree. By default, numeric
        labels are drawn beside nodes and mutations: these can be altered using the
        ``node_labels`` and ``mutation_labels`` parameters. See the
        :ref:`visualization tutorial<tutorials:sec_tskit_viz>` for more details.

        :param str path: The path to the file to write the output. If None, do not
            write to file.
        :param tuple(int, int) size: A tuple of (width, height) specifying a target
            drawing size in abstract user units (usually interpreted as pixels on
            initial display). Components of the drawing will be scaled so that the total
            plot including labels etc. normally fits onto a canvas of this size (see
            ``canvas_size`` below). If ``None``, pick a size appropriate for a tree
            with a reasonably small number (i.e. tens) of samples. Default: ``None``
        :type size:
        :param str time_scale: Control how height values for nodes are computed.
            If this is equal to ``"time"`` (the default), node heights are proportional
            to their time values. If this is equal to ``"log_time"``, node heights are
            proportional to their log(time) values. If it is equal to ``"rank"``, node
            heights are spaced equally according to their ranked times.
        :param str tree_height_scale: Deprecated alias for time_scale. (Deprecated in
                0.3.6)
        :param str title: A title string to be included in the SVG output. If ``None``
            (default) no title is shown, which gives more vertical space for the tree.
        :param str,float max_time: The maximum plotted time value in the current
            scaling system (see ``time_scale``). Can be either a string or a
            numeric value. If equal to ``"tree"`` (the default), the maximum time
            is set to be that of the oldest root in the tree. If equal to ``"ts"`` the
            maximum time is set to be the time of the oldest root in the tree
            sequence; this is useful when drawing trees from the same tree sequence as it
            ensures that node heights are consistent. If a numeric value, this is used as
            the maximum plotted time by which to scale other nodes.
        :param str,float min_time: The minimum plotted time value in the current
            scaling system (see ``time_scale``). Can be either a string or a
            numeric value. If equal to ``"tree"`` (the default), the minimum time
            is set to be that of the youngest node in the tree. If equal to ``"ts"`` the
            minimum time is set to be the time of the youngest node in the tree
            sequence; this is useful when drawing trees from the same tree sequence as it
            ensures that node heights are consistent. If a numeric value, this is used as
            the minimum plotted time.
        :param str,float max_tree_height: Deprecated alias for max_time. (Deprecated in
            0.3.6)
        :param node_labels: If specified, show custom labels for the nodes
            (specified by ID) that are present in this map; any nodes not present will
            not have a label.
        :type node_labels: dict(int, str)
        :param mutation_labels: If specified, show custom labels for the
            mutations (specified by ID) that are present in the map; any mutations
            not present will not have a label.
        :type mutation_labels: dict(int, str)
        :param dict(int, str) node_titles: If specified, add a ``<title>`` string to
            symbols for each node (specified by ID) present in this map. SVG visualizers
            such as web browsers will commonly display this string on mousing over the
            node symbol.
        :param dict(int, str) mutation_titles: If specified, add a ``<title>`` string to
            symbols for each mutation (specified by ID) present in this map. SVG
            visualizers such as web browsers will commonly display this string on
            mousing over the mutation symbol in the tree and (if show) on the x axis.
        :param dict root_svg_attributes: Additional attributes, such as an id, that will
            be embedded in the root ``<svg>`` tag of the generated drawing.
        :param str style: A
            `css style string <https://www.w3.org/TR/CSS22/syndata.html>`_ that will be
            included in the ``<style>`` tag of the generated svg.
        :param str order: The left-to-right ordering of child nodes in the drawn tree.
            This can be either: ``"minlex"``, which minimises the differences
            between adjacent trees (see also the ``"minlex_postorder"`` traversal
            order for the :meth:`.nodes` method); or ``"tree"`` which draws trees
            in the left-to-right order defined by the
            :ref:`quintuply linked tree structure <sec_data_model_tree_structure>`.
            If not specified or None, this defaults to ``"minlex"``.
        :param bool force_root_branch: If ``True`` always plot a branch (edge) above the
            root(s) in the tree. If ``None`` (default) then only plot such root branches
            if there is a mutation above a root of the tree.
        :param float symbol_size: Change the default size of the node and mutation
            plotting symbols. If ``None`` (default) use a standard size.
        :param bool x_axis: Should the plot have an X axis line, showing the start and
            end position of this tree along the genome. If ``None`` (default) do not
            plot an X axis.
        :param str x_label: Place a label under the plot. If ``None`` (default) and
            there is an X axis, create and place an appropriate label.
        :param dict x_regions: A dictionary mapping (left, right) tuples to names. This
            draws a box, labelled with the name, on the X axis between the left and
            right positions, and can be used for annotating genomic regions (e.g.
            genes) on the X axis. If ``None`` (default) do not plot any regions.
        :param bool y_axis: Should the plot have an Y axis line, showing time (or
            ranked node time if ``time_scale="rank"``). If ``None`` (default)
            do not plot a Y axis.
        :param str y_label: Place a label to the left of the plot. If ``None`` (default)
            and there is a Y axis,  create and place an appropriate label.
        :param Union[list, dict] y_ticks: A list of Y values at which to plot
            tickmarks, or a dictionary mapping Y values to labels (``[]`` gives no
            tickmarks). If ``None`` (default), plot one tickmark for each unique node
            value. Note that if ``time_scale="rank"``, the Y values refer to the
            zero-based rank of the plotted nodes, rather than the node time itself.
        :param bool y_gridlines: Whether to plot horizontal lines behind the tree
            at each y tickmark.
        :param bool all_edge_mutations: The edge on which a mutation occurs may span
            multiple trees. If ``False`` or ``None`` (default) mutations are only drawn
            on an edge if their site position exists within the genomic interval covered
            by this tree. If ``True``, all mutations on each edge of the tree are drawn,
            even if their genomic position is to the left or right of the tree
            itself. Note that this means that independent drawings of different trees
            from the same tree sequence may share some plotted mutations.
        :param bool omit_sites: If True, omit sites and mutations from the drawing.
            Default: False
        :param tuple(int, int) canvas_size: The (width, height) of the SVG canvas.
            This will change the SVG width and height without rescaling graphical
            elements, allowing extra room e.g. for unusually long labels. If ``None``
            take the canvas size to be the same as the target drawing size (see
            ``size``, above). Default: None

        :return: An SVG representation of a tree.
        :rtype: SVGString
        """
        draw = drawing.SvgTree(
            self,
            size,
            time_scale=time_scale,
            tree_height_scale=tree_height_scale,
            title=title,
            max_time=max_time,
            min_time=min_time,
            max_tree_height=max_tree_height,
            node_labels=node_labels,
            mutation_labels=mutation_labels,
            node_titles=node_titles,
            mutation_titles=mutation_titles,
            root_svg_attributes=root_svg_attributes,
            style=style,
            order=order,  # NB undocumented: Tree.draw_svg can also take an iterable here
            force_root_branch=force_root_branch,
            symbol_size=symbol_size,
            x_axis=x_axis,
            x_label=x_label,
            x_regions=x_regions,
            y_axis=y_axis,
            y_label=y_label,
            y_ticks=y_ticks,
            y_gridlines=y_gridlines,
            all_edge_mutations=all_edge_mutations,
            omit_sites=omit_sites,
            canvas_size=canvas_size,
            **kwargs,
        )
        output = draw.drawing.tostring()
        if path is not None:
            # TODO: removed the pretty here when this is stable.
            draw.drawing.saveas(path, pretty=True)
        return drawing.SVGString(output)

    def draw(
        self,
        path=None,
        width=None,
        height=None,
        node_labels=None,
        node_colours=None,
        mutation_labels=None,
        mutation_colours=None,
        format=None,  # noqa A002
        edge_colours=None,
        time_scale=None,
        tree_height_scale=None,
        max_time=None,
        min_time=None,
        max_tree_height=None,
        order=None,
        omit_sites=None,
    ):
        """
        Returns a drawing of this tree.

        When working in a Jupyter notebook, use the ``IPython.display.SVG``
        function to display the SVG output from this function inline in the notebook::

            >>> SVG(tree.draw())

        The unicode format uses unicode `box drawing characters
        <https://en.wikipedia.org/wiki/Box-drawing_character>`_ to render the tree.
        This allows rendered trees to be printed out to the terminal::

            >>> print(tree.draw(format="unicode"))
              6
            ┏━┻━┓
            ┃   5
            ┃ ┏━┻┓
            ┃ ┃  4
            ┃ ┃ ┏┻┓
            3 0 1 2

        The ``node_labels`` argument allows the user to specify custom labels
        for nodes, or no labels at all::

            >>> print(tree.draw(format="unicode", node_labels={}))
              ┃
            ┏━┻━┓
            ┃   ┃
            ┃ ┏━┻┓
            ┃ ┃  ┃
            ┃ ┃ ┏┻┓
            ┃ ┃ ┃ ┃

        Note: in some environments such as Jupyter notebooks with Windows or Mac,
        users have observed that the Unicode box drawings can be misaligned. In
        these cases, we recommend using the SVG or ASCII display formats instead.
        If you have a strong preference for aligned Unicode, you can try out the
        solution documented
        `here <https://github.com/tskit-dev/tskit/issues/189#issuecomment-499114811>`_.

        :param str path: The path to the file to write the output. If None, do not
            write to file.
        :param int width: The width of the image in pixels. If not specified, either
            defaults to the minimum size required to depict the tree (text formats)
            or 200 pixels.
        :param int height: The height of the image in pixels. If not specified, either
            defaults to the minimum size required to depict the tree (text formats)
            or 200 pixels.
        :param dict node_labels: If specified, show custom labels for the nodes
            that are present in the map. Any nodes not specified in the map will
            not have a node label.
        :param dict node_colours: If specified, show custom colours for the nodes
            given in the map. Any nodes not specified in the map will take the default
            colour; a value of ``None`` is treated as transparent and hence the node
            symbol is not plotted. (Only supported in the SVG format.)
        :param dict mutation_labels: If specified, show custom labels for the mutations
            (specified by ID) that are present in the map. Any mutations not in the map
            will not have a label. (Showing mutations is currently only supported in the
            SVG format)
        :param dict mutation_colours: If specified, show custom colours for the mutations
            given in the map (specified by ID). As for ``node_colours``, mutations not
            present in the map take the default colour, and those mapping to ``None``
            are not drawn. (Only supported in the SVG format.)
        :param str format: The format of the returned image. Currently supported
            are 'svg', 'ascii' and 'unicode'. Note that the :meth:`Tree.draw_svg`
            method provides more comprehensive functionality for creating SVGs.
        :param dict edge_colours: If specified, show custom colours for the edge
            joining each node in the map to its parent. As for ``node_colours``,
            unspecified edges take the default colour, and ``None`` values result in the
            edge being omitted. (Only supported in the SVG format.)
        :param str time_scale: Control how height values for nodes are computed.
            If this is equal to ``"time"``, node heights are proportional to their time
            values. If this is equal to ``"log_time"``, node heights are proportional to
            their log(time) values. If it is equal to ``"rank"``, node heights are spaced
            equally according to their ranked times. For SVG output the default is
            'time'-scale whereas for text output the default is 'rank'-scale.
            Time scaling is not currently supported for text output.
        :param str tree_height_scale: Deprecated alias for time_scale. (Deprecated in
                0.3.6)
        :param str,float max_time: The maximum time value in the current
            scaling system (see ``time_scale``). Can be either a string or a
            numeric value. If equal to ``"tree"``, the maximum time is set to be
            that of the oldest root in the tree. If equal to ``"ts"`` the maximum
            time is set to be the time of the oldest root in the tree sequence;
            this is useful when drawing trees from the same tree sequence as it ensures
            that node heights are consistent. If a numeric value, this is used as the
            maximum time by which to scale other nodes. This parameter
            is not currently supported for text output.
        :param str,float min_time: The minimum time value in the current
            scaling system (see ``time_scale``). Can be either a string or a
            numeric value. If equal to ``"tree"``, the minimum time is set to be
            that of the youngest node in the tree. If equal to ``"ts"`` the minimum
            time is set to be the time of the youngest node in the tree sequence;
            this is useful when drawing trees from the same tree sequence as it ensures
            that node heights are consistent. If a numeric value, this is used as the
            minimum time to display. This parameter is not currently supported for text
            output.
        :param str max_tree_height: Deprecated alias for max_time. (Deprecated in
                0.3.6)
        :param str order: The left-to-right ordering of child nodes in the drawn tree.
            This can be either: ``"minlex"``, which minimises the differences
            between adjacent trees (see also the ``"minlex_postorder"`` traversal
            order for the :meth:`.nodes` method); or ``"tree"`` which draws trees
            in the left-to-right order defined by the
            :ref:`quintuply linked tree structure <sec_data_model_tree_structure>`.
            If not specified or None, this defaults to ``"minlex"``.
        :param bool omit_sites: If True, omit sites and mutations from the drawing
            (only relevant to the SVG format). Default: False
        :return: A representation of this tree in the requested format.
        :rtype: str
        """
        output = drawing.draw_tree(
            self,
            format=format,
            width=width,
            height=height,
            node_labels=node_labels,
            node_colours=node_colours,
            mutation_labels=mutation_labels,
            mutation_colours=mutation_colours,
            edge_colours=edge_colours,
            time_scale=time_scale,
            tree_height_scale=tree_height_scale,
            max_time=max_time,
            min_time=min_time,
            max_tree_height=max_tree_height,
            order=order,
            omit_sites=omit_sites,
        )
        if path is not None:
            with open(path, "w") as f:
                f.write(output)
        return output

    def get_num_mutations(self):
        return self.num_mutations

    @property
    def num_mutations(self):
        """
        Returns the total number of mutations across all sites on this tree.

        :return: The total number of mutations over all sites on this tree.
        :rtype: int
        """
        return sum(len(site.mutations) for site in self.sites())

    @property
    def num_sites(self):
        """
        Returns the number of sites on this tree.

        :return: The number of sites on this tree.
        :rtype: int
        """
        return self._ll_tree.get_num_sites()

    def sites(self):
        """
        Returns an iterator over all the :ref:`sites <sec_site_table_definition>`
        in this tree. Sites are returned in order of increasing ID
        (and also position). See the :class:`Site` class for details on
        the available fields for each site.

        :return: An iterator over all sites in this tree.
        """
        # TODO change the low-level API to just return the IDs of the sites.
        for ll_site in self._ll_tree.get_sites():
            _, _, _, id_, _ = ll_site
            yield self.tree_sequence.site(id_)

    def mutations(self):
        """
        Returns an iterator over all the
        :ref:`mutations <sec_mutation_table_definition>` in this tree.
        Mutations are returned in their
        :ref:`order in the mutations table<sec_mutation_requirements>`,
        that is, by nondecreasing site ID, and within a site, by decreasing
        mutation time with parent mutations before their children.
        See the :class:`Mutation` class for details on the available fields for
        each mutation.

        The returned iterator is equivalent to iterating over all sites
        and all mutations in each site, i.e.::

            >>> for site in tree.sites():
            >>>     for mutation in site.mutations:
            >>>         yield mutation

        :return: An iterator over all :class:`Mutation` objects in this tree.
        :rtype: iter(:class:`Mutation`)
        """
        for site in self.sites():
            yield from site.mutations

    def get_leaves(self, u):
        # Deprecated alias for samples. See the discussion in the get_num_leaves
        # method for why this method is here and why it is semantically incorrect.
        # The 'leaves' iterator below correctly returns the leaves below a given
        # node.
        return self.samples(u)

    def leaves(self, u=None):
        """
        Returns an iterator over all the leaves in this tree that descend from
        the specified node. If :math:`u`  is not specified, return all leaves on
        the tree (i.e. all leaves reachable from the tree root(s), see note below).

        .. note::
            :math:`u` can be any node in the entire tree sequence, including ones
            which are not connected via branches to a root node of the tree. If
            called on such a node, the iterator will return "dead" leaves
            (see :ref:`sec_data_model_tree_dead_leaves_and_branches`) which cannot
            be reached from a root of this tree. However, dead leaves will never be
            returned if :math:`u` is left unspecified.

        :param int u: The node of interest.
        :return: An iterator over all leaves in the subtree rooted at u.
        :rtype: collections.abc.Iterable
        """
        roots = [u]
        if u is None:
            roots = self.roots
        for root in roots:
            for v in self.nodes(root):
                if self.is_leaf(v):
                    yield v

    def _sample_generator(self, u):
        if self._ll_tree.get_options() & _tskit.SAMPLE_LISTS:
            samples = self.tree_sequence.samples()
            index = self.left_sample(u)
            if index != NULL:
                stop = self.right_sample(u)
                while True:
                    yield samples[index]
                    if index == stop:
                        break
                    index = self.next_sample(index)
        else:
            # Fall back on iterating over all nodes in the tree, yielding
            # samples as we see them.
            for v in self.nodes(u):
                if self.is_sample(v):
                    yield v

    def samples(self, u=None):
        """
        Returns an iterator over the numerical IDs of all the sample nodes in
        this tree that are underneath the node with ID ``u``. If ``u`` is a sample,
        it is included in the returned iterator. If ``u`` is not a sample, it is
        possible for the returned iterator to be empty, for example if ``u`` is an
        :meth:`isolated<Tree.is_isolated>` node that is not part of the the current
        topology. If u is not specified, return all sample node IDs in the tree
        (equivalent to all the sample node IDs in the tree sequence).

        If the :meth:`TreeSequence.trees` method is called with
        ``sample_lists=True``, this method uses an efficient algorithm to find
        the sample nodes. If not, a simple traversal based method is used.

        .. note::

            The iterator is *not* guaranteed to return the sample node IDs in
            numerical or any other particular order.

        :param int u: The node of interest.
        :return: An iterator over all sample node IDs in the subtree rooted at u.
        :rtype: collections.abc.Iterable
        """
        roots = [u]
        if u is None:
            roots = self.roots
        for root in roots:
            yield from self._sample_generator(root)

    def num_children(self, u):
        """
        Returns the number of children of the specified
        node (i.e., ``len(tree.children(u))``)

        :param int u: The node of interest.
        :return: The number of immediate children of the node u in this tree.
        :rtype: int
        """
        return self._ll_tree.get_num_children(u)

    def get_num_leaves(self, u):
        # Deprecated alias for num_samples. The method name is inaccurate
        # as this will count the number of tracked _samples_. This is only provided to
        # avoid breaking existing code and should not be used in new code. We could
        # change this method to be semantically correct and just count the
        # number of leaves we hit in the leaves() iterator. However, this would
        # have the undesirable effect of making code that depends on the constant
        # time performance of get_num_leaves many times slower. So, the best option
        # is to leave this method as is, and to slowly deprecate it out. Once this
        # has been removed, we might add in a ``num_leaves`` method that returns the
        # length of the leaves() iterator as one would expect.
        return self.num_samples(u)

    def get_num_samples(self, u=None):
        # Deprecated alias for num_samples.
        return self.num_samples(u)

    def num_samples(self, u=None):
        """
        Returns the number of sample nodes in this tree underneath the specified
        node (including the node itself). If u is not specified return
        the total number of samples in the tree.

        This is a constant time operation.

        :param int u: The node of interest.
        :return: The number of samples in the subtree rooted at u.
        :rtype: int
        """
        u = self.virtual_root if u is None else u
        return self._ll_tree.get_num_samples(u)

    def get_num_tracked_leaves(self, u):
        # Deprecated alias for num_tracked_samples. The method name is inaccurate
        # as this will count the number of tracked _samples_. This is only provided to
        # avoid breaking existing code and should not be used in new code.
        return self.num_tracked_samples(u)

    def get_num_tracked_samples(self, u=None):
        # Deprecated alias for num_tracked_samples
        return self.num_tracked_samples(u)

    def num_tracked_samples(self, u=None):
        """
        Returns the number of samples in the set specified in the
        ``tracked_samples`` parameter of the :meth:`TreeSequence.trees` method
        underneath the specified node. If the input node is not specified,
        return the total number of tracked samples in the tree.

        This is a constant time operation.

        :param int u: The node of interest.
        :return: The number of samples within the set of tracked samples in
            the subtree rooted at u.
        :rtype: int
        """
        u = self.virtual_root if u is None else u
        return self._ll_tree.get_num_tracked_samples(u)

    def preorder(self, u=NULL):
        """
        Returns a numpy array of node ids in `preorder
        <https://en.wikipedia.org/wiki/Tree_traversal#Pre-order_(NLR)>`_. If the node u
        is specified the traversal is rooted at this node (and it will be the first
        element in the returned array). Otherwise, all nodes reachable from the tree
        roots will be returned. See :ref:`tutorials:sec_analysing_trees_traversals` for
        examples.

        :param int u: If specified, return all nodes in the subtree rooted at u
            (including u) in traversal order.
        :return: Array of node ids
        :rtype: numpy.ndarray (dtype=np.int32)
        """
        return self._ll_tree.get_preorder(u)

    def postorder(self, u=NULL):
        """
        Returns a numpy array of node ids in `postorder
        <https://en.wikipedia.org/wiki/Tree_traversal##Post-order_(LRN)>`_. If the node u
        is specified the traversal is rooted at this node (and it will be the last
        element in the returned array). Otherwise, all nodes reachable from the tree
        roots will be returned. See :ref:`tutorials:sec_analysing_trees_traversals` for
        examples.

        :param int u: If specified, return all nodes in the subtree rooted at u
            (including u) in traversal order.
        :return: Array of node ids
        :rtype: numpy.ndarray (dtype=np.int32)
        """
        return self._ll_tree.get_postorder(u)

    def timeasc(self, u=NULL):
        """
        Returns a numpy array of node ids. Starting at `u`, returns the reachable
        descendant nodes in order of increasing time (most recent first), falling back
        to increasing ID if times are equal. Also see
        :ref:`tutorials:sec_analysing_trees_traversals` for examples of how to use
        traversals.

        :param int u: If specified, return all nodes in the subtree rooted at u
            (including u) in traversal order.
        :return: Array of node ids
        :rtype: numpy.ndarray (dtype=np.int32)
        """
        nodes = self.preorder(u)
        is_virtual_root = u == self.virtual_root
        time = self.tree_sequence.nodes_time
        if is_virtual_root:
            # We could avoid creating this array if we wanted to, but
            # it's not that often people will be using this with the
            # virtual_root as an argument, so doesn't seem worth
            # the complexity
            time = np.append(time, [np.inf])
        order = np.lexsort([nodes, time[nodes]])
        return nodes[order]

    def timedesc(self, u=NULL):
        """
        Returns a numpy array of node ids. Starting at `u`, returns the reachable
        descendant nodes in order of decreasing time (least recent first), falling back
        to decreasing ID if times are equal. Also see
        :ref:`tutorials:sec_analysing_trees_traversals` for examples of how to use
        traversals.

        :param int u: If specified, return all nodes in the subtree rooted at u
            (including u) in traversal order.
        :return: Array of node ids
        :rtype: numpy.ndarray (dtype=np.int32)
        """
        return self.timeasc(u)[::-1]

    def _preorder_traversal(self, root):
        # Return Python integers for compatibility
        return map(int, self.preorder(root))

    def _postorder_traversal(self, root):
        # Return Python integers for compatibility
        return map(int, self.postorder(root))

    def _inorder_traversal(self, root):
        # TODO add a nonrecursive version of the inorder traversal.

        def traverse(u):
            children = self.get_children(u)
            mid = len(children) // 2
            for c in children[:mid]:
                yield from traverse(c)
            yield u
            for c in children[mid:]:
                yield from traverse(c)

        roots = self.roots if root == NULL else [root]
        for root in roots:
            yield from traverse(root)

    def _levelorder_traversal(self, root):
        roots = self.roots if root == NULL else [root]
        queue = collections.deque(roots)
        # For perf we store these to avoid lookups in the tight loop
        pop = queue.popleft
        extend = queue.extend
        children = self.children
        # Note: the usual style is to be explicit about what we're testing
        # and use while len(queue) > 0, but this form is slightly faster.
        while queue:
            v = pop()
            extend(children(v))
            yield v

    def _timeasc_traversal(self, root):
        """
        Sorts by increasing time but falls back to increasing ID for equal times.
        """
        return map(int, self.timeasc(root))

    def _timedesc_traversal(self, root):
        """
        The reverse of timeasc.
        """
        return map(int, self.timedesc(root))

    def _minlex_postorder_traversal(self, root):
        """
        Postorder traversal that visits leaves in minimum lexicographic order.

        Minlex stands for minimum lexicographic. We wish to visit a tree in such
        a way that the leaves visited, when their IDs are listed out, have
        minimum lexicographic order. This is a useful ordering for drawing
        multiple Trees of a TreeSequence, as it leads to more consistency
        between adjacent Trees.
        """

        # We compute a dictionary mapping from internal node ID to min leaf ID
        # under the node, using a first postorder traversal
        min_leaf = {}
        for u in self.nodes(root, order="postorder"):
            if self.is_leaf(u):
                min_leaf[u] = u
            else:
                min_leaf[u] = min(min_leaf[v] for v in self.children(u))

        # If we deliberately specify the virtual root, it should also be returned
        is_virtual_root = root == self.virtual_root
        if root == -1:
            root = self.virtual_root

        stack = [(root, False)]
        while len(stack) > 0:
            u, visited = stack.pop()
            if visited:
                if u != self.virtual_root or is_virtual_root:
                    yield u
            else:
                stack.append((u, True))  # Reappend, marking visited
                stack.extend(
                    sorted(
                        ((c, False) for c in self.children(u)),
                        key=lambda v: min_leaf[v[0]],
                        reverse=True,
                    )
                )

    def nodes(self, root=None, order="preorder"):
        """
        Returns an iterator over the node IDs reachable from the specified node in this
        tree in the specified traversal order.

        .. note::
            Unlike the :meth:`TreeSequence.nodes` method, this iterator produces
            integer node IDs, not :class:`Node` objects.

        If the ``root`` parameter is not provided or ``None``, iterate over all
        nodes reachable from the roots (see :attr:`Tree.roots` for details
        on which nodes are considered roots). If the ``root`` parameter
        is provided, only the nodes in the subtree rooted at this node
        (including the specified node) will be iterated over. If the
        :attr:`.virtual_root` is specified as the traversal root, it will
        be included in the traversed nodes in the appropriate position
        for the given ordering. (See the
        :ref:`tree roots <sec_data_model_tree_virtual_root>` section for more
        information on the virtual root.)

        The ``order`` parameter defines the order in which tree nodes are visited
        in the iteration (also see the :ref:`sec_analysing_trees_traversals` section
        in the `tutorials <https://tskit.dev/tutorials>`__). The available orders are:

        - 'preorder': starting at root, yield the current node, then recurse
          and do a preorder on each child of the current node. See also `Wikipedia
          <https://en.wikipedia.org/wiki/Tree_traversal#Pre-order_(NLR)>`__.
        - 'inorder': starting at root, assuming binary trees, recurse and do
          an inorder on the first child, then yield the current node, then
          recurse and do an inorder on the second child. In the case of ``n``
          child nodes (not necessarily 2), the first ``n // 2`` children are
          visited in the first stage, and the remaining ``n - n // 2`` children
          are visited in the second stage. See also `Wikipedia
          <https://en.wikipedia.org/wiki/Tree_traversal#In-order_(LNR)>`__.
        - 'postorder': starting at root, recurse and do a postorder on each
          child of the current node, then yield the current node. See also
          `Wikipedia
          <https://en.wikipedia.org/wiki/Tree_traversal#Post-order_(LRN)>`__.
        - 'levelorder' ('breadthfirst'): visit the nodes under root (including
          the root) in increasing order of their depth from root. See also
          `Wikipedia
          <https://en.wikipedia.org/wiki/Tree_traversal\
#Breadth-first_search_/_level_order>`__.
        - 'timeasc': visits the nodes in order of increasing time, falling back to
          increasing ID if times are equal.
        - 'timedesc': visits the nodes in order of decreasing time, falling back to
          decreasing ID if times are equal.
        - 'minlex_postorder': a usual postorder has ambiguity in the order in
          which children of a node are visited. We constrain this by outputting
          a postorder such that the leaves visited, when their IDs are
          listed out, have minimum `lexicographic order
          <https://en.wikipedia.org/wiki/Lexicographical_order>`__ out of all valid
          traversals. This traversal is useful for drawing multiple trees of
          a ``TreeSequence``, as it leads to more consistency between adjacent
          trees. Note that internal non-leaf nodes are not counted in
          assessing the lexicographic order.

        :param int root: The root of the subtree we are traversing.
        :param str order: The traversal ordering. Currently 'preorder',
            'inorder', 'postorder', 'levelorder' ('breadthfirst'), 'timeasc' and
            'timedesc' and 'minlex_postorder' are supported.
        :return: An iterator over the node IDs in the tree in some traversal order.
        :rtype: collections.abc.Iterable, int
        """
        methods = {
            "preorder": self._preorder_traversal,
            "inorder": self._inorder_traversal,
            "postorder": self._postorder_traversal,
            "levelorder": self._levelorder_traversal,
            "breadthfirst": self._levelorder_traversal,
            "timeasc": self._timeasc_traversal,
            "timedesc": self._timedesc_traversal,
            "minlex_postorder": self._minlex_postorder_traversal,
        }
        try:
            iterator = methods[order]
        except KeyError:
            raise ValueError(f"Traversal ordering '{order}' not supported")

        root = -1 if root is None else root
        return iterator(root)

    def _as_newick_fast(self, *, root, precision, legacy_ms_labels):
        """
        Call into the fast but limited C implementation of the newick conversion.
        """
        root_time = max(1, self.time(root))
        max_label_size = math.ceil(math.log10(self.tree_sequence.num_nodes))
        single_node_size = (
            5 + max_label_size + math.ceil(math.log10(root_time)) + precision
        )
        buffer_size = 1 + single_node_size * self.tree_sequence.num_nodes
        return self._ll_tree.get_newick(
            precision=precision,
            root=root,
            buffer_size=buffer_size,
            legacy_ms_labels=legacy_ms_labels,
        )

    def as_newick(
        self,
        *,
        root=None,
        precision=None,
        node_labels=None,
        include_branch_lengths=None,
    ):
        """
        Returns a `newick encoding
        <https://en.wikipedia.org/wiki/Newick_format>`_ of this tree.
        For example, a binary tree with 3 leaves generated by
        :meth:`Tree.generate_balanced(3)<Tree.generate_balanced>`
        encodes as::

            (n0:2,(n1:1,n2:1):1);

        By default :ref:`sample nodes<sec_data_model_definitions>` are
        labelled using the form ``f"n{node_id}"``, i.e. the sample node's
        ID prefixed with the string ``"n"``. Node labels can be specified
        explicitly using the ``node_labels`` argument, which is a mapping from
        integer node IDs to the corresponding string label. If a node is not
        present in the mapping, no label is associated with the node in
        output.

        .. warning:: Node labels are **not** Newick escaped, so care must be taken
            to provide labels that will not break the encoding.

        .. note:: Specifying a ``node_labels`` dictionary or setting
            ``include_branch_lengths=False`` results in a less efficient
            method being used to generate the newick output. The performance
            difference can be substantial for large trees.

        By default, branch lengths are printed out with sufficient precision
        for them to be recovered exactly in double precision (although note
        that this does not necessarily mean that we can precisely recover the
        corresponding node times, since branch lengths are obtained by
        subtraction). If all times on the tree sequence are discrete, then
        branch lengths are printed as integers. Otherwise, branch lengths are
        printed with 17 digits of precision (i.e., ``"%.17f"`` in
        printf-notation).

        The precision for branch lengths can be specified using the ``precision``
        argument. Branch lengths can be omitted entirely by setting
        ``include_branch_lengths=False``.

        If the ``root`` argument is specified, we return the newick encoding of
        the specified subtree, otherwise the full tree is returned. If the tree
        has :ref:`multiple roots <sec_data_model_tree_roots>` and a root node
        is not explicitly specified, we raise a ``ValueError``. This is because
        most libraries and downstream software consider a newick string that
        contains multiple disconnected subtrees an error, and it is therefore
        best to consider how such topologies should be interchanged on a
        case-by-base basis. A list of the newick strings for each root can be
        obtained by ``[tree.as_newick(root=root) for root in tree.roots]``.

        :param int precision: The numerical precision with which branch lengths are
            printed. If not specified or None default to 0 if the tree sequence
            contains only integer node times, or 17 otherwise.
        :param int root: If specified, return the tree rooted at this node.
        :param dict node_labels: If specified, show custom labels for the nodes
            that are present in the map. Any nodes not specified in the map will
            not have a node label.
        :param include_branch_lengths: If True (default), output branch lengths in the
            Newick string. If False, only output the topology, without branch lengths.
        :return: A newick representation of this tree.
        :rtype: str
        """
        if root is None:
            if not self.has_single_root:
                raise ValueError(
                    "Cannot get newick unless a tree has a single root. Try "
                    "[t.as_newick(root) for root in t.roots] to get a list of "
                    "newick trees, one for each root."
                )
            root = self.root

        if precision is None:
            # 17 decimal digits provides the full precision of an IEEE double,
            # as defined by DBL_DECIMAL_DIG macro. If we have discrete time
            # then write out integer branch lengths.
            precision = 0 if self.tree_sequence.discrete_time else 17
        include_branch_lengths = (
            True if include_branch_lengths is None else include_branch_lengths
        )
        # Can we run this through the fast path?
        if include_branch_lengths and node_labels in [LEGACY_MS_LABELS, None]:
            # Note the LEGACY_MS_LABELS code path is not part of the documented
            # interface and should not be depended on by client code.
            return self._as_newick_fast(
                root=root,
                precision=precision,
                legacy_ms_labels=node_labels == LEGACY_MS_LABELS,
            )

        # No, we have to use the slower Python code.
        if node_labels is None:
            node_labels = {u: f"n{u}" for u in self.tree_sequence.samples()}
        elif node_labels == LEGACY_MS_LABELS:
            # NOTE in the ms format it's the *leaf* nodes we label not
            # necessarily the samples. We keep this behaviour to avoid
            # breaking legacy code that may depend on it.
            node_labels = {u: f"{u + 1}" for u in self.leaves()}
        return text_formats.build_newick(
            self,
            root=root,
            precision=precision,
            node_labels=node_labels,
            include_branch_lengths=include_branch_lengths,
        )

    def newick(
        self,
        precision=14,
        *,
        root=None,
        node_labels=None,
        include_branch_lengths=True,
    ):
        """
        .. warning:: This method is deprecated and may be removed in future
            versions of tskit. Please use the :meth:`.as_newick` method
            in new code.

        This method is a deprecated version of the :meth:`.as_newick` method.
        Functionality is equivalent, except for the default node labels.

        By default, *leaf* nodes are labelled with their numerical ID + 1,
        and internal nodes are not labelled. This default strategy was originally
        used to mimic the output of the ``ms`` simulator. However, the choice
        of labelling leaf nodes rather than samples is problematic, and this
        behaviour is only retained to avoid breaking existing code which may
        rely on it.

        Other parameters behave as documented in the :meth:`.as_newick` method.

        :param int precision: The numerical precision with which branch lengths are
            printed. Defaults to 14.
        :param int root: If specified, return the tree rooted at this node.
        :param dict node_labels: If specified, show custom labels for the nodes
            that are present in the map. Any nodes not specified in the map will
            not have a node label.
        :param include_branch_lengths: If True (default), output branch lengths in the
            Newick string. If False, only output the topology, without branch lengths.
        :return: A newick representation of this tree.
        :rtype: str
        """
        node_labels = LEGACY_MS_LABELS if node_labels is None else node_labels
        return self.as_newick(
            root=root,
            precision=precision,
            node_labels=node_labels,
            include_branch_lengths=include_branch_lengths,
        )

    def as_dict_of_dicts(self):
        """
        Convert tree to dict of dicts for conversion to a
        `networkx graph <https://networkx.github.io/documentation/stable/
        reference/classes/digraph.html>`_.

        For example::

            >>> import networkx as nx
            >>> nx.DiGraph(tree.as_dict_of_dicts())
            >>> # undirected graphs work as well
            >>> nx.Graph(tree.as_dict_of_dicts())

        :return: Dictionary of dictionaries of dictionaries where the first key
            is the source, the second key is the target of an edge, and the
            third key is an edge annotation. At this point the only annotation
            is "branch_length", the length of the branch (in units of time).
        """
        dod = {}
        for parent in self.nodes():
            dod[parent] = {}
            for child in self.children(parent):
                dod[parent][child] = {"branch_length": self.branch_length(child)}
        return dod

    @property
    def parent_dict(self):
        return self.get_parent_dict()

    def get_parent_dict(self):
        pi = {
            u: self.parent(u)
            for u in range(self.tree_sequence.num_nodes)
            if self.parent(u) != NULL
        }
        return pi

    def __str__(self):
        """
        Return a plain text summary of a tree in a tree sequence
        """
        tree_rows = [
            ["Index", f"{self.index:,}"],
            [
                "Interval",
                f"{self.interval.left:,.8g}-{self.interval.right:,.8g}"
                f"({self.span:,.8g})",
            ],
            ["Roots", f"{self.num_roots:,}"],
            ["Nodes", f"{len(self.preorder()):,}"],
            ["Sites", f"{self.num_sites:,}"],
            ["Mutations", f"{self.num_mutations:,}"],
            ["Total Branch Length", f"{self.total_branch_length:,.8g}"],
        ]
        return util.unicode_table(tree_rows, title="Tree")

    def _repr_html_(self):
        """
        Return an html summary of a tree in a tree sequence. Called by jupyter
        notebooks to render trees
        """
        return util.tree_html(self)

    def map_mutations(self, genotypes, alleles, ancestral_state=None):
        """
        Given observations for the samples in this tree described by the specified
        set of genotypes and alleles, return a parsimonious set of state transitions
        explaining these observations. The genotypes array is interpreted as indexes
        into the alleles list in the same manner as described in the
        :meth:`TreeSequence.variants` method. Thus, if sample ``j`` carries the
        allele at index ``k``, then we have ``genotypes[j] = k``.
        Missing observations can be specified for a sample using the value
        ``tskit.MISSING_DATA`` (-1), in which case the state at this sample does not
        influence the ancestral state or the position of mutations returned. At least
        one non-missing observation must be provided. A maximum of 64 alleles are
        supported.

        The current implementation uses the Hartigan parsimony algorithm to determine
        the minimum number of state transitions required to explain the data. In this
        model, transitions between any of the non-missing states is equally likely.

        The returned values correspond directly to the data model for describing
        variation at sites using mutations. See the :ref:`sec_site_table_definition`
        and :ref:`sec_mutation_table_definition` definitions for details and background.

        The state reconstruction is returned as two-tuple, ``(ancestral_state,
        mutations)``, where ``ancestral_state`` is the allele assigned to the
        tree root(s) and ``mutations`` is a list of :class:`Mutation` objects,
        ordered as :ref:`required in a mutation table<sec_mutation_requirements>`.
        For each mutation, ``derived_state`` is the new state after this mutation and
        ``node`` is the tree node immediately beneath the mutation (if there are unary
        nodes between two branch points, hence multiple nodes above which the
        mutation could be parsimoniously placed, the oldest node is used). The
        ``parent`` property contains the index in the returned list of the previous
        mutation on the path to root, or ``tskit.NULL``
        if there are no previous mutations (see the :ref:`sec_mutation_table_definition`
        for more information on the concept of mutation parents). All other attributes
        of the :class:`Mutation` object are undefined and should not be used.

        .. note::
            Sample states observed as missing in the input ``genotypes`` need
            not correspond to samples whose nodes are actually "missing" (i.e.,
            :ref:`isolated<sec_data_model_missing_data>`) in the tree. In this
            case, mapping the mutations returned by this method onto the tree
            will result in these missing observations being imputed to the
            most parsimonious state.

        Because the ``parent`` in the returned list of mutations refers to the index
        in that list, if you are adding mutations to an existing tree sequence, you
        will need to maintain a map of list IDs to the newly added mutations, for
        instance::

            last_tree = ts.last()
            anc_state, parsimonious_muts = last_tree.map_mutations([0, 1, 0], ("A", "T"))
            # Edit the tree sequence, see the "Tables and Editing" tutorial
            tables = ts.dump_tables()
            # add a new site at the end of ts, assumes there isn't one there already
            site_id = tables.sites.add_row(ts.sequence_length - 1, anc_state)

            mut_id_map = {tskit.NULL: tskit.NULL}  # don't change if parent id is -1
            for list_id, mutation in enumerate(parsimonious_muts):
                mut_id_map[list_id] = tables.mutations.append(
                    mutation.replace(site=site_id, parent=mut_id_map[mutation.parent]))
            tables.sort()  # Redundant here, but needed if the site is not the last one
            new_ts = tables.tree_sequence()

        See the :ref:`tutorials:sec_analysing_trees_parsimony` section in the tutorial
        for further examples of how to use this method.

        :param array_like genotypes: The input observations for the samples in this tree.
        :param tuple(str) alleles: The alleles for the specified ``genotypes``. Each
            positive value in the ``genotypes`` array is treated as an index into this
            list of alleles.
        :param ancestral_state: A fixed ancestral state, specified either as a
            non-negative integer less than the number of alleles, or a string which
            must be one of the ``alleles`` provided above. If ``None`` (default) then
            an ancestral state is chosen arbitrarily from among those that provide
            the most parsimonious placement of mutations. Note that if the ancestral
            state is specified, the placement of mutations may not be as parsimonious
            as that which could be achieved by leaving the ancestral state unspecified;
            additionally it may lead to mutations being placed above the root node(s) of
            the tree (for example if all the samples have a genotype of 1 but the
            ancestral state is fixed to be 0).
        :type ancestral_state: Union[int, str]
        :return: The inferred ancestral state and list of mutations on this tree
            that encode the specified observations.
        :rtype: (str, list(tskit.Mutation))
        """
        genotypes = util.safe_np_int_cast(genotypes, np.int8)
        max_alleles = np.max(genotypes)
        if ancestral_state is not None:
            if isinstance(ancestral_state, str):
                # Will raise a ValueError if not in the list
                ancestral_state = alleles.index(ancestral_state)
            if ancestral_state < 0 or ancestral_state >= len(alleles):
                raise ValueError("ancestral_state not between 0 and (num_alleles-1)")
            max_alleles = max(ancestral_state, max_alleles)
        if max_alleles >= 64:
            raise ValueError("A maximum of 64 states is supported")
        ancestral_state, transitions = self._ll_tree.map_mutations(
            genotypes, ancestral_state
        )
        # Translate back into string alleles
        ancestral_state = alleles[ancestral_state]
        mutations = [
            Mutation(
                node=node,
                derived_state=alleles[derived_state],
                parent=parent,
                metadata=self.tree_sequence.table_metadata_schemas.mutation.empty_value,
            )
            for node, parent, derived_state in transitions
        ]
        return ancestral_state, mutations

    def kc_distance(self, other, lambda_=0.0):
        """
        Returns the Kendall-Colijn distance between the specified pair of trees.
        The ``lambda_`` parameter  determines the relative weight of topology
        vs branch lengths in calculating the distance. If ``lambda_`` is 0
        (the default) we only consider topology, and if it is 1 we only
        consider branch lengths. See `Kendall & Colijn (2016)
        <https://academic.oup.com/mbe/article/33/10/2735/2925548>`_ for details.

        The trees we are comparing to must have identical lists of sample
        nodes (i.e., the same IDs in the same order). The metric operates on
        samples, not leaves, so internal samples are treated identically to
        sample tips. Subtrees with no samples do not contribute to the metric.

        :param Tree other: The other tree to compare to.
        :param float lambda_: The KC metric lambda parameter determining the
            relative weight of topology and branch length.
        :return: The computed KC distance between this tree and other.
        :rtype: float
        """
        return self._ll_tree.get_kc_distance(other._ll_tree, lambda_)

    def _get_sample_sets(self):
        ret = {}
        for u in self.nodes(order="postorder"):
            u_sample_set = set()
            if self.is_sample(u):
                u_sample_set.add(u)
            for v in self.children(u):
                u_sample_set |= ret[v]
            ret[u] = frozenset(u_sample_set)
        return ret

    def rf_distance(self, other):
        """
        Returns the (unweighted) Robinson-Foulds distance between the specified pair
        of trees, where corresponding samples between the two trees are identified by
        node ID. The Robinson-Foulds distance (also known as the symmetric difference)
        is defined as the number of bipartitions that are present in one tree but
        not the other (see
        `Robinson & Foulds (1981) <https://doi.org/10.1016/0025-5564(81)90043-2>`_).
        This method returns the unnormalised RF distance: if the
        trees are strictly bifurcating, i.e. binary, the value can be
        normalised by dividing by the maximum, which is $2n-4$ for two rooted
        trees of $n$ samples (however, if the trees contain polytomies, the maximum
        RF distance is less easily defined).

        .. note::
            The RF distance can be sensitive to small changes in topology: in some
            cases, changing the position of a single leaf can result in the maximum
            RF distance. Therefore even if adjacent trees in a tree sequence differ
            by a single subtree-prune-and-regraft operation, the RF distance
            between them can be large.

        :param Tree other: The other tree to compare to. Trees are treated as rooted.
        :return: The unweighted Robinson-Foulds distance between this tree and ``other``.
        :rtype: int
        :raises ValueError: If either tree has multiple roots, or the trees have
            different sample nodes.
        """
        if self.num_roots != 1 or other.num_roots != 1:
            raise ValueError("Trees must have a single root")

        s1 = set(self._get_sample_sets().values())
        s2 = set(other._get_sample_sets().values())

        return len(s1.symmetric_difference(s2))

    def path_length(self, u, v):
        """
        Returns the number of edges on the path in this tree between the two nodes.
        If the two nodes have a most recent common ancestor, then this is defined as
        ``tree.depth(u) + tree.depth(v) - 2 * tree.depth(tree.mrca(u, v))``. If the nodes
        do not have an MRCA (i.e., they are in disconnected subtrees) the path length
        is infinity.

        .. note:: This counts the number of "hops" between two nodes. To find the branch
            length distance between them, in units of time (i.e. the sum of edge lengths
            that separate two nodes) use the :meth:`.distance_between` method instead.

        .. seealso:: See also the :meth:`.depth` method

        :param int u: The first node for path length computation.
        :param int v: The second node for path length computation.
        :return: The number of edges between the two nodes.
        :rtype: int
        """
        mrca = self.mrca(u, v)
        if mrca == -1:
            return math.inf
        return self.depth(u) + self.depth(v) - 2 * self.depth(mrca)

    def distance_between(self, u, v):
        """
        Returns the total distance between two nodes in the tree, expressed as
        the sum of "branch lengths" from both nodes to their most recent common ancestor.

        :param int u: The first node for path length computation.
        :param int v: The second node for path length computation.
        :return: The distance between the two nodes, the sum of "branch lengths" .
        :rtype: float
        """
        tmrca = self.tmrca(u, v)
        return tmrca - self.time(u) + tmrca - self.time(v)

    def b1_index(self):
        """
        Returns the
        `B1 balance index <https://treebalance.wordpress.com/b₁-index/>`_
        for this tree. This is defined as the inverse of the sum of all
        longest paths to leaves for each node besides roots.

        .. seealso:: See `Shao and Sokal (1990)
            <https://www.jstor.org/stable/2992186>`_ for details.

        :return: The B1 balance index.
        :rtype: float
        """
        return self._ll_tree.get_b1_index()

    def b2_index(self, base=10):
        """
        Returns the
        `B2 balance index <https://treebalance.wordpress.com/b₂-index/>`_
        this tree.
        This is defined as the Shannon entropy of the probability
        distribution to reach leaves assuming a random walk
        from a root. The default base is 10, following Shao and Sokal (1990).

        .. seealso:: See `Shao and Sokal (1990)
            <https://www.jstor.org/stable/2992186>`_ for details.

        :param int base: The base used for the logarithm in the
            Shannon entropy computation.
        :return: The B2 balance index.
        :rtype: float
        """
        # Let Python decide if the base is acceptable
        math.log(10, base)
        return self._ll_tree.get_b2_index(base)

    def colless_index(self):
        """
        Returns the
        `Colless imbalance index <https://treebalance.wordpress.com/colless-index/>`_
        for this tree. This is defined as the sum of all differences between
        number of leaves subtended by the left and right child of each node.
        The Colless index is undefined for non-binary trees and trees with
        multiple roots. This method will raise a LibraryError if the tree is
        not singly-rooted and binary.

        .. seealso:: See `Shao and Sokal (1990)
            <https://www.jstor.org/stable/2992186>`_ for details.

        :return: The Colless imbalance index.
        :rtype: int
        """
        return self._ll_tree.get_colless_index()

    def sackin_index(self):
        """
        Returns the
        `Sackin imbalance index <https://treebalance.wordpress.com/sackin-index/>`_
        for this tree. This is defined as the sum of the depths of all leaves
        in the tree. Equivalent to ``sum(tree.depth(u) for u in
        tree.leaves())``

        .. seealso:: See `Shao and Sokal (1990)
            <https://www.jstor.org/stable/2992186>`_ for details.

        :return: The Sackin imbalance index.
        :rtype: int
        """
        return self._ll_tree.get_sackin_index()

    def num_lineages(self, t):
        """
        Returns the number of lineages present in this tree at time ``t``. This
        is defined as the number of branches in this tree (reachable from the
        samples) that intersect with ``t``. Thus, ``tree.num_lineages(t)``
        is equal to 0 for any ``t`` greater than or equal to the time of
        the root in a singly-rooted tree.

        .. note:: Note that this definition means that if a (non root) node
            with three children has time ``t``, then it will count as one lineage,
            not three.

        :param int t: The time to count lineages at.
        :return: The number of lineages in the tree at time t.
        :rtype: int
        """
        return self._ll_tree.get_num_lineages(t)

    def split_polytomies(
        self,
        *,
        epsilon=None,
        method=None,
        record_provenance=True,
        random_seed=None,
        **kwargs,
    ):
        """
        Return a new :class:`.Tree` where extra nodes and edges have been inserted
        so that any any node ``u`` with greater than 2 children --- a multifurcation
        or "polytomy" --- is resolved into successive bifurcations. New nodes are
        inserted at times fractionally less than than the time of node ``u``.
        Times are allocated to different levels of the tree, such that any newly
        inserted sibling nodes will have the same time.

        By default, the times of the newly generated children of a particular
        node are the minimum representable distance in floating point arithmetic
        from their parents (using the `nextafter
        <https://numpy.org/doc/stable/reference/generated/numpy.nextafter.html>`_
        function). Thus, the generated branches have the shortest possible nonzero
        length. A fixed branch length between inserted nodes and their parents
        can also be specified by using the ``epsilon`` parameter.

        .. note::
            A tree sequence :ref:`requires<sec_valid_tree_sequence_requirements>` that
            parents be older than children and that mutations are younger than the
            parent of the edge on which they lie. If a fixed ``epsilon`` is specifed
            and is not small enough compared to the distance between a polytomy and
            its oldest child (or oldest child mutation) these requirements may not
            be met. In this case an error will be raised.

        If the ``method`` is ``"random"`` (currently the only option, and the default
        when no method is specified), then for a node with :math:`n` children, the
        :math:`(2n - 3)! / (2^(n - 2) (n - 2!))` possible binary trees with equal
        probability.

        The returned :class:`.Tree` will have the same genomic span as this tree,
        and node IDs will be conserved (that is, node ``u`` in this tree will
        be the same node in the returned tree). The returned tree is derived from a
        tree sequence that contains only one non-degenerate tree, that is, where
        edges cover only the interval spanned by this tree.

        :param epsilon: If specified, the fixed branch length between inserted
            nodes and their parents. If None (the default), the minimal possible
            nonzero branch length is generated for each node.
        :param str method: The method used to break polytomies. Currently only "random"
            is supported, which can also be specified by ``method=None``
            (Default: ``None``).
        :param bool record_provenance: If True, add details of this operation to the
            provenance information of the returned tree sequence. (Default: True).
        :param int random_seed: The random seed. If this is None, a random seed will
            be automatically generated. Valid random seeds must be between 1 and
            :math:`2^32 − 1`.
        :param \\**kwargs: Further arguments used as parameters when constructing the
            returned :class:`Tree`. For example
            ``tree.split_polytomies(sample_lists=True)`` will
            return a :class:`Tree` created with ``sample_lists=True``.
        :return: A new tree with polytomies split into random bifurcations.
        :rtype: tskit.Tree
        """
        return combinatorics.split_polytomies(
            self,
            epsilon=epsilon,
            method=method,
            record_provenance=record_provenance,
            random_seed=random_seed,
            **kwargs,
        )

    @staticmethod
    def generate_star(
        num_leaves, *, span=1, branch_length=1, record_provenance=True, **kwargs
    ):
        """
        Generate a :class:`Tree` whose leaf nodes all have the same parent (i.e.,
        a "star" tree). The leaf nodes are all at time 0 and are marked as sample nodes.

        The tree produced by this method is identical to
        ``tskit.Tree.unrank(n, (0, 0))``, but generated more efficiently for large ``n``.

        :param int num_leaves: The number of leaf nodes in the returned tree (must be
            2 or greater).
        :param float span: The span of the tree, and therefore the
            :attr:`~TreeSequence.sequence_length` of the :attr:`.tree_sequence`
            property of the returned :class:`Tree`.
        :param float branch_length: The length of every branch in the tree (equivalent
            to the time of the root node).
        :param bool record_provenance: If True, add details of this operation to the
            provenance information of the returned tree sequence. (Default: True).
        :param \\**kwargs: Further arguments used as parameters when constructing the
            returned :class:`Tree`. For example
            ``tskit.Tree.generate_star(sample_lists=True)`` will
            return a :class:`Tree` created with ``sample_lists=True``.
        :return: A star-shaped tree. Its corresponding :class:`TreeSequence` is available
            via the :attr:`.tree_sequence` attribute.
        :rtype: Tree
        """
        return combinatorics.generate_star(
            num_leaves,
            span=span,
            branch_length=branch_length,
            record_provenance=record_provenance,
            **kwargs,
        )

    @staticmethod
    def generate_balanced(
        num_leaves,
        *,
        arity=2,
        span=1,
        branch_length=1,
        record_provenance=True,
        **kwargs,
    ):
        """
        Generate a :class:`Tree` with the specified number of leaves that is maximally
        balanced. By default, the tree returned is binary, such that for each
        node that subtends :math:`n` leaves, the left child will subtend
        :math:`\\lfloor{n / 2}\\rfloor` leaves and the right child the
        remainder. Balanced trees with higher arity can also generated using the
        ``arity`` parameter, where the leaves subtending a node are distributed
        among its children analogously.

        In the returned tree, the leaf nodes are all at time 0, marked as samples,
        and labelled 0 to n from left-to-right. Internal node IDs are assigned
        sequentially from n in a postorder traversal, and the time of an internal
        node is the maximum time of its children plus the specified ``branch_length``.

        :param int num_leaves: The number of leaf nodes in the returned tree (must be
            be 2 or greater).
        :param int arity: The maximum number of children a node can have in the returned
            tree.
        :param float span: The span of the tree, and therefore the
            :attr:`~TreeSequence.sequence_length` of the :attr:`.tree_sequence`
            property of the returned :class:`Tree`.
        :param float branch_length: The minimum length of a branch in the tree (see
            above for details on how internal node times are assigned).
        :param bool record_provenance: If True, add details of this operation to the
            provenance information of the returned tree sequence. (Default: True).
        :param \\**kwargs: Further arguments used as parameters when constructing the
            returned :class:`Tree`. For example
            ``tskit.Tree.generate_balanced(sample_lists=True)`` will
            return a :class:`Tree` created with ``sample_lists=True``.
        :return: A balanced tree. Its corresponding :class:`TreeSequence` is available
            via the :attr:`.tree_sequence` attribute.
        :rtype: Tree
        """
        return combinatorics.generate_balanced(
            num_leaves,
            arity=arity,
            span=span,
            branch_length=branch_length,
            record_provenance=record_provenance,
            **kwargs,
        )

    @staticmethod
    def generate_comb(
        num_leaves, *, span=1, branch_length=1, record_provenance=True, **kwargs
    ):
        """
        Generate a :class:`Tree` in which all internal nodes have two children
        and the left child is a leaf. This is a "comb", "ladder" or "pectinate"
        phylogeny, and also known as a `caterpillar tree
        <https://en.wikipedia.org/wiki/Caterpillar_tree>`_.

        The leaf nodes are all at time 0, marked as samples,
        and labelled 0 to n from left-to-right. Internal node IDs are assigned
        sequentially from n as we ascend the tree, and the time of an internal
        node is the maximum time of its children plus the specified ``branch_length``.

        :param int num_leaves: The number of leaf nodes in the returned tree (must be
            2 or greater).
        :param float span: The span of the tree, and therefore the
            :attr:`~TreeSequence.sequence_length` of the :attr:`.tree_sequence`
            property of the returned :class:`Tree`.
        :param float branch_length: The branch length between each internal node; the
            root node is therefore placed at time ``branch_length * (num_leaves - 1)``.
        :param bool record_provenance: If True, add details of this operation to the
            provenance information of the returned tree sequence. (Default: True).
        :param \\**kwargs: Further arguments used as parameters when constructing the
            returned :class:`Tree`. For example
            ``tskit.Tree.generate_comb(sample_lists=True)`` will
            return a :class:`Tree` created with ``sample_lists=True``.
        :return: A comb-shaped bifurcating tree. Its corresponding :class:`TreeSequence`
            is available via the :attr:`.tree_sequence` attribute.
        :rtype: Tree
        """
        return combinatorics.generate_comb(
            num_leaves,
            span=span,
            branch_length=branch_length,
            record_provenance=record_provenance,
            **kwargs,
        )

    @staticmethod
    def generate_random_binary(
        num_leaves,
        *,
        span=1,
        branch_length=1,
        random_seed=None,
        record_provenance=True,
        **kwargs,
    ):
        """
        Generate a random binary :class:`Tree` with :math:`n` = ``num_leaves``
        leaves with an equal probability of returning any topology and
        leaf label permutation among the :math:`(2n - 3)! / (2^{n - 2} (n - 2)!)`
        leaf-labelled binary trees.

        The leaf nodes are marked as samples, labelled 0 to n, and placed at
        time 0. Internal node IDs are assigned sequentially from n as we ascend
        the tree, and the time of an internal node is the maximum time of its
        children plus the specified ``branch_length``.

        .. note::
            The returned tree has not been created under any explicit model of
            evolution. In order to simulate such trees, additional software
            such as `msprime <https://github.com/tskit-dev/msprime>`` is required.

        :param int num_leaves: The number of leaf nodes in the returned tree (must
            be 2 or greater).
        :param float span: The span of the tree, and therefore the
            :attr:`~TreeSequence.sequence_length` of the :attr:`.tree_sequence`
            property of the returned :class:`Tree`.
        :param float branch_length: The minimum time between parent and child nodes.
        :param int random_seed: The random seed. If this is None, a random seed will
            be automatically generated. Valid random seeds must be between 1 and
            :math:`2^32 − 1`.
        :param bool record_provenance: If True, add details of this operation to the
            provenance information of the returned tree sequence. (Default: True).
        :param \\**kwargs: Further arguments used as parameters when constructing the
            returned :class:`Tree`. For example
            ``tskit.Tree.generate_comb(sample_lists=True)`` will
            return a :class:`Tree` created with ``sample_lists=True``.
        :return: A random binary tree. Its corresponding :class:`TreeSequence` is
            available via the :attr:`.tree_sequence` attribute.
        :rtype: Tree
        """
        return combinatorics.generate_random_binary(
            num_leaves,
            span=span,
            branch_length=branch_length,
            random_seed=random_seed,
            record_provenance=record_provenance,
            **kwargs,
        )


def load(file, *, skip_tables=False, skip_reference_sequence=False):
    """
    Return a :class:`TreeSequence` instance loaded from the specified file object or
    path. The file must be in the
    :ref:`tree sequence file format <sec_tree_sequence_file_format>`
    produced by the :meth:`TreeSequence.dump` method.

    .. warning:: With any of the ``skip_tables`` or ``skip_reference_sequence``
        options set, it is not possible to load data from a non-seekable stream
        (e.g. a socket or STDIN) of multiple tree sequences using consecutive
        calls to :meth:`tskit.load`.

    :param str file: The file object or path of the ``.trees`` file containing the
        tree sequence we wish to load.
    :param bool skip_tables: If True, no tables are read from the ``.trees``
        file and only the top-level information is populated in the tree
        sequence object.
    :param bool skip_reference_sequence: If True, the tree sequence is read
        without loading its reference sequence.
    :return: The tree sequence object containing the information
        stored in the specified file path.
    :rtype: :class:`tskit.TreeSequence`
    """
    return TreeSequence.load(
        file, skip_tables=skip_tables, skip_reference_sequence=skip_reference_sequence
    )


def parse_individuals(
    source, strict=True, encoding="utf8", base64_metadata=True, table=None
):
    """
    Parse the specified file-like object containing a whitespace delimited
    description of an individual table and returns the corresponding
    :class:`IndividualTable` instance. See the :ref:`individual text format
    <sec_individual_text_format>` section for the details of the required
    format and the :ref:`individual table definition
    <sec_individual_table_definition>` section for the required properties of
    the contents.

    See :func:`tskit.load_text` for a detailed explanation of the ``strict``
    parameter.

    :param io.TextIOBase source: The file-like object containing the text.
    :param bool strict: If True, require strict tab delimiting (default). If
        False, a relaxed whitespace splitting algorithm is used.
    :param str encoding: Encoding used for text representation.
    :param bool base64_metadata: If True, metadata is encoded using Base64
        encoding; otherwise, as plain text.
    :param IndividualTable table: If specified write into this table. If not,
        create a new :class:`IndividualTable` instance.
    """
    sep = None
    if strict:
        sep = "\t"
    if table is None:
        table = tables.IndividualTable()
    # Read the header and find the indexes of the required fields.
    header = source.readline().rstrip("\n").split(sep)
    flags_index = header.index("flags")
    location_index = None
    parents_index = None
    metadata_index = None
    try:
        location_index = header.index("location")
    except ValueError:
        pass
    try:
        parents_index = header.index("parents")
    except ValueError:
        pass
    try:
        metadata_index = header.index("metadata")
    except ValueError:
        pass
    for line in source:
        tokens = line.rstrip("\n").split(sep)
        if len(tokens) >= 1:
            flags = int(tokens[flags_index])
            location = ()
            if location_index is not None:
                location_string = tokens[location_index]
                if len(location_string) > 0:
                    location = tuple(map(float, location_string.split(",")))
            parents = ()
            if parents_index is not None:
                parents_string = tokens[parents_index]
                if len(parents_string) > 0:
                    parents = tuple(map(int, parents_string.split(",")))
            metadata = b""
            if metadata_index is not None and metadata_index < len(tokens):
                metadata = tokens[metadata_index].encode(encoding)
                if base64_metadata:
                    metadata = base64.b64decode(metadata)
            table.add_row(
                flags=flags, location=location, parents=parents, metadata=metadata
            )
    return table


def parse_nodes(source, strict=True, encoding="utf8", base64_metadata=True, table=None):
    """
    Parse the specified file-like object containing a whitespace delimited
    description of a node table and returns the corresponding :class:`NodeTable`
    instance. See the :ref:`node text format <sec_node_text_format>` section
    for the details of the required format and the
    :ref:`node table definition <sec_node_table_definition>` section for the
    required properties of the contents.

    See :func:`tskit.load_text` for a detailed explanation of the ``strict``
    parameter.

    :param io.TextIOBase source: The file-like object containing the text.
    :param bool strict: If True, require strict tab delimiting (default). If
        False, a relaxed whitespace splitting algorithm is used.
    :param str encoding: Encoding used for text representation.
    :param bool base64_metadata: If True, metadata is encoded using Base64
        encoding; otherwise, as plain text.
    :param NodeTable table: If specified write into this table. If not,
        create a new :class:`NodeTable` instance.
    """
    sep = None
    if strict:
        sep = "\t"
    if table is None:
        table = tables.NodeTable()
    # Read the header and find the indexes of the required fields.
    header = source.readline().rstrip("\n").split(sep)
    is_sample_index = header.index("is_sample")
    time_index = header.index("time")
    population_index = None
    individual_index = None
    metadata_index = None
    try:
        population_index = header.index("population")
    except ValueError:
        pass
    try:
        individual_index = header.index("individual")
    except ValueError:
        pass
    try:
        metadata_index = header.index("metadata")
    except ValueError:
        pass
    for line in source:
        tokens = line.rstrip("\n").split(sep)
        if len(tokens) >= 2:
            is_sample = int(tokens[is_sample_index])
            time = float(tokens[time_index])
            flags = 0
            if is_sample != 0:
                flags |= NODE_IS_SAMPLE
            population = NULL
            if population_index is not None:
                population = int(tokens[population_index])
            individual = NULL
            if individual_index is not None:
                individual = int(tokens[individual_index])
            metadata = b""
            if metadata_index is not None and metadata_index < len(tokens):
                metadata = tokens[metadata_index].encode(encoding)
                if base64_metadata:
                    metadata = base64.b64decode(metadata)
            table.add_row(
                flags=flags,
                time=time,
                population=population,
                individual=individual,
                metadata=metadata,
            )
    return table


def parse_edges(source, strict=True, table=None):
    """
    Parse the specified file-like object containing a whitespace delimited
    description of a edge table and returns the corresponding :class:`EdgeTable`
    instance. See the :ref:`edge text format <sec_edge_text_format>` section
    for the details of the required format and the
    :ref:`edge table definition <sec_edge_table_definition>` section for the
    required properties of the contents.

    See :func:`tskit.load_text` for a detailed explanation of the ``strict`` parameter.

    :param io.TextIOBase source: The file-like object containing the text.
    :param bool strict: If True, require strict tab delimiting (default). If
        False, a relaxed whitespace splitting algorithm is used.
    :param EdgeTable table: If specified, write the edges into this table. If
        not, create a new :class:`EdgeTable` instance and return.
    """
    sep = None
    if strict:
        sep = "\t"
    if table is None:
        table = tables.EdgeTable()
    header = source.readline().rstrip("\n").split(sep)
    left_index = header.index("left")
    right_index = header.index("right")
    parent_index = header.index("parent")
    children_index = header.index("child")
    for line in source:
        tokens = line.rstrip("\n").split(sep)
        if len(tokens) >= 4:
            left = float(tokens[left_index])
            right = float(tokens[right_index])
            parent = int(tokens[parent_index])
            children = tuple(map(int, tokens[children_index].split(",")))
            for child in children:
                table.add_row(left=left, right=right, parent=parent, child=child)
    return table


def parse_sites(source, strict=True, encoding="utf8", base64_metadata=True, table=None):
    """
    Parse the specified file-like object containing a whitespace delimited
    description of a site table and returns the corresponding :class:`SiteTable`
    instance. See the :ref:`site text format <sec_site_text_format>` section
    for the details of the required format and the
    :ref:`site table definition <sec_site_table_definition>` section for the
    required properties of the contents.

    See :func:`tskit.load_text` for a detailed explanation of the ``strict``
    parameter.

    :param io.TextIOBase source: The file-like object containing the text.
    :param bool strict: If True, require strict tab delimiting (default). If
        False, a relaxed whitespace splitting algorithm is used.
    :param str encoding: Encoding used for text representation.
    :param bool base64_metadata: If True, metadata is encoded using Base64
        encoding; otherwise, as plain text.
    :param SiteTable table: If specified write site into this table. If not,
        create a new :class:`SiteTable` instance.
    """
    sep = None
    if strict:
        sep = "\t"
    if table is None:
        table = tables.SiteTable()
    header = source.readline().rstrip("\n").split(sep)
    position_index = header.index("position")
    ancestral_state_index = header.index("ancestral_state")
    metadata_index = None
    try:
        metadata_index = header.index("metadata")
    except ValueError:
        pass
    for line in source:
        tokens = line.rstrip("\n").split(sep)
        if len(tokens) >= 2:
            position = float(tokens[position_index])
            ancestral_state = tokens[ancestral_state_index]
            metadata = b""
            if metadata_index is not None and metadata_index < len(tokens):
                metadata = tokens[metadata_index].encode(encoding)
                if base64_metadata:
                    metadata = base64.b64decode(metadata)
            table.add_row(
                position=position, ancestral_state=ancestral_state, metadata=metadata
            )
    return table


def parse_mutations(
    source, strict=True, encoding="utf8", base64_metadata=True, table=None
):
    """
    Parse the specified file-like object containing a whitespace delimited
    description of a mutation table and returns the corresponding :class:`MutationTable`
    instance. See the :ref:`mutation text format <sec_mutation_text_format>` section
    for the details of the required format and the
    :ref:`mutation table definition <sec_mutation_table_definition>` section for the
    required properties of the contents. Note that if the ``time`` column is missing its
    entries are filled with ``UNKNOWN_TIME``.

    See :func:`tskit.load_text` for a detailed explanation of the ``strict``
    parameter.

    :param io.TextIOBase source: The file-like object containing the text.
    :param bool strict: If True, require strict tab delimiting (default). If
        False, a relaxed whitespace splitting algorithm is used.
    :param str encoding: Encoding used for text representation.
    :param bool base64_metadata: If True, metadata is encoded using Base64
        encoding; otherwise, as plain text.
    :param MutationTable table: If specified, write mutations into this table.
        If not, create a new :class:`MutationTable` instance.
    """
    sep = None
    if strict:
        sep = "\t"
    if table is None:
        table = tables.MutationTable()
    header = source.readline().rstrip("\n").split(sep)
    site_index = header.index("site")
    node_index = header.index("node")
    try:
        time_index = header.index("time")
    except ValueError:
        time_index = None
    derived_state_index = header.index("derived_state")
    parent_index = None
    parent = NULL
    try:
        parent_index = header.index("parent")
    except ValueError:
        pass
    metadata_index = None
    try:
        metadata_index = header.index("metadata")
    except ValueError:
        pass
    for line in source:
        tokens = line.rstrip("\n").split(sep)
        if len(tokens) >= 3:
            site = int(tokens[site_index])
            node = int(tokens[node_index])
            if time_index is None or tokens[time_index] == tskit.TIME_UNITS_UNKNOWN:
                time = UNKNOWN_TIME
            else:
                time = float(tokens[time_index])
            derived_state = tokens[derived_state_index]
            if parent_index is not None:
                parent = int(tokens[parent_index])
            metadata = b""
            if metadata_index is not None and metadata_index < len(tokens):
                metadata = tokens[metadata_index].encode(encoding)
                if base64_metadata:
                    metadata = base64.b64decode(metadata)
            table.add_row(
                site=site,
                node=node,
                time=time,
                derived_state=derived_state,
                parent=parent,
                metadata=metadata,
            )
    return table


def parse_populations(
    source, strict=True, encoding="utf8", base64_metadata=True, table=None
):
    """
    Parse the specified file-like object containing a whitespace delimited
    description of a population table and returns the corresponding
    :class:`PopulationTable` instance. See the :ref:`population text format
    <sec_population_text_format>` section for the details of the required
    format and the :ref:`population table definition
    <sec_population_table_definition>` section for the required properties of
    the contents.

    See :func:`tskit.load_text` for a detailed explanation of the ``strict``
    parameter.

    :param io.TextIOBase source: The file-like object containing the text.
    :param bool strict: If True, require strict tab delimiting (default). If
        False, a relaxed whitespace splitting algorithm is used.
    :param str encoding: Encoding used for text representation.
    :param bool base64_metadata: If True, metadata is encoded using Base64
        encoding; otherwise, as plain text.
    :param PopulationTable table: If specified write into this table. If not,
        create a new :class:`PopulationTable` instance.
    """
    sep = None
    if strict:
        sep = "\t"
    if table is None:
        table = tables.PopulationTable()
    # Read the header and find the indexes of the required fields.
    header = source.readline().rstrip("\n").split(sep)
    metadata_index = header.index("metadata")
    for line in source:
        tokens = line.rstrip("\n").split(sep)
        if len(tokens) >= 1:
            metadata = tokens[metadata_index].encode(encoding)
            if base64_metadata:
                metadata = base64.b64decode(metadata)
            table.add_row(metadata=metadata)
    return table


def parse_migrations(
    source, strict=True, encoding="utf8", base64_metadata=True, table=None
):
    """
    Parse the specified file-like object containing a whitespace delimited
    description of a migration table and returns the corresponding
    :class:`MigrationTable` instance.

    See the :ref:`migration text format <sec_migration_text_format>` section
    for the details of the required format and the
    :ref:`migration table definition <sec_migration_table_definition>` section
    for the required properties of the contents. Note that if the ``time`` column
    is missing its entries are filled with :data:`UNKNOWN_TIME`.

    See :func:`tskit.load_text` for a detailed explanation of the ``strict``
    parameter.

    :param io.TextIOBase source: The file-like object containing the text.
    :param bool strict: If True, require strict tab delimiting (default). If
        False, a relaxed whitespace splitting algorithm is used.
    :param str encoding: Encoding used for text representation.
    :param bool base64_metadata: If True, metadata is encoded using Base64
        encoding; otherwise, as plain text.
    :param MigrationTable table: If specified, write migrations into this table.
        If not, create a new :class:`MigrationTable` instance.
    """
    sep = None
    if strict:
        sep = "\t"
    if table is None:
        table = tables.MigrationTable()
    header = source.readline().rstrip("\n").split(sep)
    left_index = header.index("left")
    right_index = header.index("right")
    node_index = header.index("node")
    source_index = header.index("source")
    dest_index = header.index("dest")
    time_index = header.index("time")
    metadata_index = None
    try:
        metadata_index = header.index("metadata")
    except ValueError:
        pass
    for line in source:
        tokens = line.rstrip("\n").split(sep)
        if len(tokens) >= 6:
            left = float(tokens[left_index])
            right = float(tokens[right_index])
            node = int(tokens[node_index])
            source = int(tokens[source_index])
            dest = int(tokens[dest_index])
            time = float(tokens[time_index])
            metadata = b""
            if metadata_index is not None and metadata_index < len(tokens):
                metadata = tokens[metadata_index].encode(encoding)
                if base64_metadata:
                    metadata = base64.b64decode(metadata)
            table.add_row(
                left=left,
                right=right,
                node=node,
                source=source,
                dest=dest,
                time=time,
                metadata=metadata,
            )
    return table


def load_text(
    nodes,
    edges,
    sites=None,
    mutations=None,
    individuals=None,
    populations=None,
    migrations=None,
    sequence_length=0,
    strict=True,
    encoding="utf8",
    base64_metadata=True,
):
    """
    Return a :class:`TreeSequence` instance parsed from tabulated text data
    contained in the specified file-like objects. The format
    for these files is documented in the :ref:`sec_text_file_format` section,
    and is produced by the :meth:`TreeSequence.dump_text` method. Further
    properties required for an input tree sequence are described in the
    :ref:`sec_valid_tree_sequence_requirements` section. This method is intended as a
    convenient interface for importing external data into tskit; the binary
    file format using by :meth:`tskit.load` is many times more efficient than
    this text format.

    The ``nodes`` and ``edges`` parameters are mandatory and must be file-like
    objects containing text with whitespace delimited columns,  parsable by
    :func:`parse_nodes` and :func:`parse_edges`, respectively. ``sites``,
    ``individuals``, ``populations``, ``mutations``, and ``migrations`` are optional,
    and must be parsable by :func:`parse_sites`, :func:`parse_individuals`,
    :func:`parse_populations`, :func:`parse_mutations`, and :func:`parse_migrations`,
    respectively. For convenience, if the node table refers to populations,
    but the ``populations`` parameter is not provided, a minimal set of rows are
    added to the population table, so that a valid tree sequence can be returned.

    The ``sequence_length`` parameter determines the
    :attr:`TreeSequence.sequence_length` of the returned tree sequence. If it
    is 0 or not specified, the value is taken to be the maximum right
    coordinate of the input edges. This parameter is useful in degenerate
    situations (such as when there are zero edges), but can usually be ignored.

    The ``strict`` parameter controls the field delimiting algorithm that
    is used. If ``strict`` is True (the default), we require exactly one
    tab character separating each field. If ``strict`` is False, a more relaxed
    whitespace delimiting algorithm is used, such that any run of whitespace
    is regarded as a field separator. In most situations, ``strict=False``
    is more convenient, but it can lead to error in certain situations. For
    example, if a deletion is encoded in the mutation table this will not
    be parseable when ``strict=False``.

    After parsing the tables, :meth:`TableCollection.sort` is called to ensure that
    the loaded tables satisfy the tree sequence :ref:`ordering requirements
    <sec_valid_tree_sequence_requirements>`. Note that this may result in the
    IDs of various entities changing from their positions in the input file.

    :param io.TextIOBase nodes: The file-like object containing text describing a
        :class:`NodeTable`.
    :param io.TextIOBase edges: The file-like object containing text
        describing an :class:`EdgeTable`.
    :param io.TextIOBase sites: The file-like object containing text describing a
        :class:`SiteTable`.
    :param io.TextIOBase mutations: The file-like object containing text
        describing a :class:`MutationTable`.
    :param io.TextIOBase individuals: The file-like object containing text
        describing a :class:`IndividualTable`.
    :param io.TextIOBase populations: The file-like object containing text
        describing a :class:`PopulationTable`.
    :param io.TextIOBase migrations: The file-like object containing text
        describing a :class:`MigrationTable`.
    :param float sequence_length: The sequence length of the returned tree sequence. If
        not supplied or zero this will be inferred from the set of edges.
    :param bool strict: If True, require strict tab delimiting (default). If
        False, a relaxed whitespace splitting algorithm is used.
    :param str encoding: Encoding used for text representation.
    :param bool base64_metadata: If True, metadata is encoded using Base64
        encoding; otherwise, as plain text.
    :return: The tree sequence object containing the information
        stored in the specified file paths.
    :rtype: :class:`tskit.TreeSequence`
    """
    # We need to parse the edges so we can figure out the sequence length, and
    # TableCollection.sequence_length is immutable so we need to create a temporary
    # edge table.
    edge_table = parse_edges(edges, strict=strict)
    if sequence_length == 0 and len(edge_table) > 0:
        sequence_length = edge_table.right.max()
    tc = tables.TableCollection(sequence_length)
    tc.edges.set_columns(
        left=edge_table.left,
        right=edge_table.right,
        parent=edge_table.parent,
        child=edge_table.child,
    )
    parse_nodes(
        nodes,
        strict=strict,
        encoding=encoding,
        base64_metadata=base64_metadata,
        table=tc.nodes,
    )
    if sites is not None:
        parse_sites(
            sites,
            strict=strict,
            encoding=encoding,
            base64_metadata=base64_metadata,
            table=tc.sites,
        )
    if mutations is not None:
        parse_mutations(
            mutations,
            strict=strict,
            encoding=encoding,
            base64_metadata=base64_metadata,
            table=tc.mutations,
        )
    if individuals is not None:
        parse_individuals(
            individuals,
            strict=strict,
            encoding=encoding,
            base64_metadata=base64_metadata,
            table=tc.individuals,
        )
    if populations is None:
        # As a convenience we add any populations referenced in the node table.
        if len(tc.nodes) > 0:
            max_population = tc.nodes.population.max()
            if max_population != NULL:
                for _ in range(max_population + 1):
                    tc.populations.add_row()
    else:
        parse_populations(
            populations,
            strict=strict,
            encoding=encoding,
            base64_metadata=base64_metadata,
            table=tc.populations,
        )
    if migrations is not None:
        parse_migrations(
            migrations,
            strict=strict,
            encoding=encoding,
            base64_metadata=base64_metadata,
            table=tc.migrations,
        )
    tc.sort()
    return tc.tree_sequence()


class TreeIterator:
    """
    Simple class providing forward and backward iteration over a tree sequence.
    """

    def __init__(self, tree):
        self.tree = tree
        self.more_trees = True
        self.forward = True

    def __iter__(self):
        return self

    def __reversed__(self):
        self.forward = False
        return self

    def __next__(self):
        if self.forward:
            self.more_trees = self.more_trees and self.tree.next()
        else:
            self.more_trees = self.more_trees and self.tree.prev()
        if not self.more_trees:
            raise StopIteration()
        return self.tree

    def __len__(self):
        return self.tree.tree_sequence.num_trees


class SimpleContainerSequence:
    """
    Simple wrapper to allow arrays of SimpleContainers (e.g. edges, nodes) that have a
    function allowing access by index (e.g. ts.edge(i), ts.node(i)) to be treated as a
    python sequence, allowing forward and reverse iteration.

    To generate a sequence of items in a different order, the ``order`` parameter allows
    an array of indexes to be passed in, such as returned from np.argsort or np.lexsort.
    """

    def __init__(self, getter, length, order=None):
        if order is None:
            self.getter = getter
        else:
            self.getter = lambda index: getter(order[index])
        self.length = length

    def __len__(self):
        return self.length

    def __getitem__(self, index):
        return self.getter(index)


@dataclass(frozen=True)
class TableMetadataSchemas:
    """
    Convenience class for returning the schemas of all the tables in a tree sequence.
    """

    node: metadata_module.MetadataSchema = None
    """
    The metadata schema of the node table.
    """

    edge: metadata_module.MetadataSchema = None
    """
    The metadata schema of the edge table.
    """

    site: metadata_module.MetadataSchema = None
    """
    The metadata schema of the site table.
    """

    mutation: metadata_module.MetadataSchema = None
    """
    The metadata schema of the mutation table.
    """

    migration: metadata_module.MetadataSchema = None
    """
    The metadata schema of the migration table.
    """

    individual: metadata_module.MetadataSchema = None
    """
    The metadata schema of the individual table.
    """

    population: metadata_module.MetadataSchema = None
    """
    The metadata schema of the population table.
    """


class TreeSequence:
    """
    A single tree sequence, as defined by the :ref:`data model <sec_data_model>`.
    A TreeSequence instance can be created from a set of
    :ref:`tables <sec_table_definitions>` using
    :meth:`TableCollection.tree_sequence`, or loaded from a set of text files
    using :func:`tskit.load_text`, or loaded from a native binary file using
    :func:`tskit.load`.

    TreeSequences are immutable. To change the data held in a particular
    tree sequence, first get the table information as a :class:`TableCollection`
    instance (using :meth:`.dump_tables`), edit those tables using the
    :ref:`tables api <sec_tables_api>`, and create a new tree sequence using
    :meth:`TableCollection.tree_sequence`.

    The :meth:`.trees` method iterates over all trees in a tree sequence, and
    the :meth:`.variants` method iterates over all sites and their genotypes.
    """

    def __init__(self, ll_tree_sequence):
        self._ll_tree_sequence = ll_tree_sequence
        metadata_schema_strings = self._ll_tree_sequence.get_table_metadata_schemas()
        metadata_schema_instances = {
            name: metadata_module.parse_metadata_schema(
                getattr(metadata_schema_strings, name)
            )
            for name in vars(TableMetadataSchemas)
            if not name.startswith("_")
        }
        self._table_metadata_schemas = TableMetadataSchemas(**metadata_schema_instances)
        self._individuals_time = None
        self._individuals_population = None
        self._individuals_location = None
        # NOTE: when we've implemented read-only access via the underlying
        # tables we can replace these arrays with reference to the read-only
        # tables here (and remove the low-level boilerplate).
        llts = self._ll_tree_sequence
        self._individuals_flags = llts.individuals_flags
        self._nodes_time = llts.nodes_time
        self._nodes_flags = llts.nodes_flags
        self._nodes_population = llts.nodes_population
        self._nodes_individual = llts.nodes_individual
        self._edges_left = llts.edges_left
        self._edges_right = llts.edges_right
        self._edges_parent = llts.edges_parent
        self._edges_child = llts.edges_child
        self._sites_position = llts.sites_position
        self._mutations_site = llts.mutations_site
        self._mutations_node = llts.mutations_node
        self._mutations_parent = llts.mutations_parent
        self._mutations_time = llts.mutations_time
        self._migrations_left = llts.migrations_left
        self._migrations_right = llts.migrations_right
        self._migrations_right = llts.migrations_right
        self._migrations_node = llts.migrations_node
        self._migrations_source = llts.migrations_source
        self._migrations_dest = llts.migrations_dest
        self._migrations_time = llts.migrations_time
        self._indexes_edge_insertion_order = llts.indexes_edge_insertion_order
        self._indexes_edge_removal_order = llts.indexes_edge_removal_order

    # Implement the pickle protocol for TreeSequence
    def __getstate__(self):
        return self.dump_tables()

    def __setstate__(self, tc):
        self.__init__(tc.tree_sequence().ll_tree_sequence)

    def __eq__(self, other):
        return self.tables == other.tables

    def equals(
        self,
        other,
        *,
        ignore_metadata=False,
        ignore_ts_metadata=False,
        ignore_provenance=False,
        ignore_timestamps=False,
        ignore_tables=False,
        ignore_reference_sequence=False,
    ):
        """
        Returns True if  `self` and `other` are equal. Uses the underlying table
        equality, see :meth:`TableCollection.equals` for details and options.
        """
        return self.tables.equals(
            other.tables,
            ignore_metadata=ignore_metadata,
            ignore_ts_metadata=ignore_ts_metadata,
            ignore_provenance=ignore_provenance,
            ignore_timestamps=ignore_timestamps,
            ignore_tables=ignore_tables,
            ignore_reference_sequence=ignore_reference_sequence,
        )

    @property
    def ll_tree_sequence(self):
        return self.get_ll_tree_sequence()

    def get_ll_tree_sequence(self):
        return self._ll_tree_sequence

    def aslist(self, **kwargs):
        """
        Returns the trees in this tree sequence as a list. Each tree is
        represented by a different instance of :class:`Tree`. As such, this
        method is inefficient and may use a large amount of memory, and should
        not be used when performance is a consideration. The :meth:`.trees`
        method is the recommended way to efficiently iterate over the trees
        in a tree sequence.

        :param \\**kwargs: Further arguments used as parameters when constructing the
            returned trees. For example ``ts.aslist(sample_lists=True)`` will result
            in a list of :class:`Tree` instances created with ``sample_lists=True``.
        :return: A list of the trees in this tree sequence.
        :rtype: list
        """
        return [tree.copy() for tree in self.trees(**kwargs)]

    @classmethod
    def load(cls, file_or_path, *, skip_tables=False, skip_reference_sequence=False):
        file, local_file = util.convert_file_like_to_open_file(file_or_path, "rb")
        try:
            ts = _tskit.TreeSequence()
            ts.load(
                file,
                skip_tables=skip_tables,
                skip_reference_sequence=skip_reference_sequence,
            )
            return TreeSequence(ts)
        except tskit.FileFormatError as e:
            util.raise_known_file_format_errors(file, e)
        finally:
            if local_file:
                file.close()

    @classmethod
    def load_tables(cls, tables, *, build_indexes=False):
        ts = _tskit.TreeSequence()
        ts.load_tables(tables._ll_tables, build_indexes=build_indexes)
        return TreeSequence(ts)

    def dump(self, file_or_path, zlib_compression=False):
        """
        Writes the tree sequence to the specified path or file object.

        :param str file_or_path: The file object or path to write the TreeSequence to.
        :param bool zlib_compression: This parameter is deprecated and ignored.
        """
        if zlib_compression:
            # Note: the msprime CLI before version 1.0 uses this option, so we need
            # to keep it indefinitely.
            warnings.warn(
                "The zlib_compression option is no longer supported and is ignored",
                RuntimeWarning,
                stacklevel=4,
            )
        file, local_file = util.convert_file_like_to_open_file(file_or_path, "wb")
        try:
            self._ll_tree_sequence.dump(file)
        finally:
            if local_file:
                file.close()

    @property
    def reference_sequence(self):
        """
        The :class:`.ReferenceSequence` associated with this :class:`.TreeSequence`
        if one is defined (see :meth:`.TreeSequence.has_reference_sequence`),
        or None otherwise.
        """
        if self.has_reference_sequence():
            return tables.ReferenceSequence(self._ll_tree_sequence.reference_sequence)
        return None

    def has_reference_sequence(self):
        """
        Returns True if this :class:`.TreeSequence` has an associated
        :ref:`reference sequence<sec_data_model_reference_sequence>`.
        """
        return bool(self._ll_tree_sequence.has_reference_sequence())

    @property
    def tables_dict(self):
        """
        Returns a dictionary mapping names to tables in the
        underlying :class:`.TableCollection`. Equivalent to calling
        ``ts.tables.table_name_map``.
        """
        return self.tables.table_name_map

    @property
    def tables(self):
        """
        Returns the :class:`tables<TableCollection>` underlying this tree
        sequence, intended for read-only access. See :meth:`.dump_tables` if you wish
        to modify the tables.

        .. warning:: This property currently returns a copy of the tables
            underlying a tree sequence but it may return a read-only
            **view** in the future. Thus, if the tables will subsequently be
            updated, please use the :meth:`.dump_tables` method instead as
            this will always return a new copy of the TableCollection.

        :return: A :class:`TableCollection` containing all a copy of the
            tables underlying this tree sequence.
        :rtype: TableCollection
        """
        return self.dump_tables()

    @property
    def nbytes(self):
        """
        Returns the total number of bytes required to store the data
        in this tree sequence. Note that this may not be equal to
        the actual memory footprint.
        """
        return self.tables.nbytes

    def dump_tables(self):
        """
        Returns a modifiable copy of the :class:`tables<TableCollection>` defining
        this tree sequence.

        :return: A :class:`TableCollection` containing all tables underlying
            the tree sequence.
        :rtype: TableCollection
        """
        ll_tables = _tskit.TableCollection(self.sequence_length)
        self._ll_tree_sequence.dump_tables(ll_tables)
        return tables.TableCollection(ll_tables=ll_tables)

    def dump_text(
        self,
        nodes=None,
        edges=None,
        sites=None,
        mutations=None,
        individuals=None,
        populations=None,
        migrations=None,
        provenances=None,
        precision=6,
        encoding="utf8",
        base64_metadata=True,
    ):
        """
        Writes a text representation of the tables underlying the tree sequence
        to the specified connections.

        If Base64 encoding is not used, then metadata will be saved directly, possibly
        resulting in errors reading the tables back in if metadata includes whitespace.

        :param io.TextIOBase nodes: The file-like object (having a .write() method) to
            write the NodeTable to.
        :param io.TextIOBase edges: The file-like object to write the EdgeTable to.
        :param io.TextIOBase sites: The file-like object to write the SiteTable to.
        :param io.TextIOBase mutations: The file-like object to write the
            MutationTable to.
        :param io.TextIOBase individuals: The file-like object to write the
            IndividualTable to.
        :param io.TextIOBase populations: The file-like object to write the
            PopulationTable to.
        :param io.TextIOBase migrations: The file-like object to write the
            MigrationTable to.
        :param io.TextIOBase provenances: The file-like object to write the
            ProvenanceTable to.
        :param int precision: The number of digits of precision.
        :param str encoding: Encoding used for text representation.
        :param bool base64_metadata: Only used if a schema is not present on each table
            being dumped. If True, metadata is encoded using Base64
            encoding; otherwise, as plain text.
        """
        text_formats.dump_text(
            self,
            nodes=nodes,
            edges=edges,
            sites=sites,
            mutations=mutations,
            individuals=individuals,
            populations=populations,
            migrations=migrations,
            provenances=provenances,
            precision=precision,
            encoding=encoding,
            base64_metadata=base64_metadata,
        )

    def __str__(self):
        """
        Return a plain text summary of the contents of a tree sequence
        """
        ts_rows = [
            ["Trees", str(self.num_trees)],
            [
                "Sequence Length",
                str(
                    int(self.sequence_length)
                    if self.discrete_genome
                    else self.sequence_length
                ),
            ],
            ["Time Units", self.time_units],
            ["Sample Nodes", str(self.num_samples)],
            ["Total Size", util.naturalsize(self.nbytes)],
        ]
        header = ["Table", "Rows", "Size", "Has Metadata"]
        table_rows = []
        for name, table in self.tables.table_name_map.items():
            table_rows.append(
                [
                    name.capitalize(),
                    f"{table.num_rows:,}",
                    util.naturalsize(table.nbytes),
                    (
                        "Yes"
                        if hasattr(table, "metadata") and len(table.metadata) > 0
                        else "No"
                    ),
                ]
            )
        return util.unicode_table(ts_rows, title="TreeSequence") + util.unicode_table(
            table_rows, header=header
        )

    def _repr_html_(self):
        """
        Return an html summary of a tree sequence. Called by jupyter notebooks
        to render a TreeSequence.
        """
        return util.tree_sequence_html(self)

    # num_samples was originally called sample_size, and so we must keep sample_size
    # around as a deprecated alias.
    @property
    def num_samples(self):
        """
        Returns the number of sample nodes in this tree sequence. This is also the
        number of sample nodes in each tree.

        :return: The number of sample nodes in this tree sequence.
        :rtype: int
        """
        return self._ll_tree_sequence.get_num_samples()

    @property
    def table_metadata_schemas(self) -> TableMetadataSchemas:
        """
        The set of metadata schemas for the tables in this tree sequence.
        """
        return self._table_metadata_schemas

    @property
    def sample_size(self):
        # Deprecated alias for num_samples
        return self.num_samples

    def get_sample_size(self):
        # Deprecated alias for num_samples
        return self.num_samples

    @property
    def file_uuid(self):
        return self._ll_tree_sequence.get_file_uuid()

    @property
    def discrete_genome(self):
        """
        Returns True if all genome coordinates in this TreeSequence are
        discrete integer values. This is true iff all the following are true:

        - The sequence length is discrete
        - All site positions are discrete
        - All left and right edge coordinates are discrete
        - All migration left and right coordinates are discrete

        :return: True if this TreeSequence uses discrete genome coordinates.
        :rtype: bool
        """
        return bool(self._ll_tree_sequence.get_discrete_genome())

    @property
    def discrete_time(self):
        """
        Returns True if all time coordinates in this TreeSequence are
        discrete integer values. This is true iff all the following are true:

        - All node times are discrete
        - All mutation times are discrete
        - All migration times are discrete

        Note that ``tskit.UNKNOWN_TIME`` counts as discrete.

        :return: True if this TreeSequence uses discrete time coordinates.
        :rtype: bool
        """
        return bool(self._ll_tree_sequence.get_discrete_time())

    @property
    def min_time(self):
        """
        Returns the min time in this tree sequence. This is the minimum
        of the node times and mutation times.

        Note that mutation times with the value ``tskit.UNKNOWN_TIME``
        are ignored.

        :return: The min time of the nodes and mutations in this tree sequence.
        :rtype: float
        """
        return self._ll_tree_sequence.get_min_time()

    @property
    def max_time(self):
        """
        Returns the max time in this tree sequence. This is the maximum
        of the node times and mutation times.

        Note that mutation times with the value ``tskit.UNKNOWN_TIME``
        are ignored.

        :return: The max time of the nodes and mutations in this tree sequence.
        :rtype: float
        """
        return self._ll_tree_sequence.get_max_time()

    @property
    def sequence_length(self):
        """
        Returns the sequence length in this tree sequence. This defines the
        genomic scale over which tree coordinates are defined. Given a
        tree sequence with a sequence length :math:`L`, the constituent
        trees will be defined over the half-closed interval
        :math:`[0, L)`. Each tree then covers some subset of this
        interval --- see :attr:`tskit.Tree.interval` for details.

        :return: The length of the sequence in this tree sequence in bases.
        :rtype: float
        """
        return self.get_sequence_length()

    def get_sequence_length(self):
        return self._ll_tree_sequence.get_sequence_length()

    @property
    def metadata(self) -> Any:
        """
        The decoded metadata for this TreeSequence.
        """
        return self.metadata_schema.decode_row(self._ll_tree_sequence.get_metadata())

    @property
    def metadata_schema(self) -> metadata_module.MetadataSchema:
        """
        The :class:`tskit.MetadataSchema` for this TreeSequence.
        """
        return metadata_module.parse_metadata_schema(
            self._ll_tree_sequence.get_metadata_schema()
        )

    @property
    def time_units(self) -> str:
        """
        String describing the units of the time dimension for this TreeSequence.
        """
        return self._ll_tree_sequence.get_time_units()

    @property
    def num_edges(self):
        """
        Returns the number of :ref:`edges <sec_edge_table_definition>` in this
        tree sequence.

        :return: The number of edges in this tree sequence.
        :rtype: int
        """
        return self._ll_tree_sequence.get_num_edges()

    def get_num_trees(self):
        # Deprecated alias for self.num_trees
        return self.num_trees

    @property
    def num_trees(self):
        """
        Returns the number of distinct trees in this tree sequence. This
        is equal to the number of trees returned by the :meth:`.trees`
        method.

        :return: The number of trees in this tree sequence.
        :rtype: int
        """
        return self._ll_tree_sequence.get_num_trees()

    def get_num_sites(self):
        # Deprecated alias for self.num_sites
        return self._ll_tree_sequence.get_num_sites()

    @property
    def num_sites(self):
        """
        Returns the number of :ref:`sites <sec_site_table_definition>` in
        this tree sequence.

        :return: The number of sites in this tree sequence.
        :rtype: int
        """
        return self.get_num_sites()

    def get_num_mutations(self):
        # Deprecated alias for self.num_mutations
        return self.num_mutations

    @property
    def num_mutations(self):
        """
        Returns the number of :ref:`mutations <sec_mutation_table_definition>`
        in this tree sequence.

        :return: The number of mutations in this tree sequence.
        :rtype: int
        """
        return self._ll_tree_sequence.get_num_mutations()

    def get_num_nodes(self):
        # Deprecated alias for self.num_nodes
        return self.num_nodes

    @property
    def num_individuals(self):
        """
        Returns the number of :ref:`individuals <sec_individual_table_definition>` in
        this tree sequence.

        :return: The number of individuals in this tree sequence.
        :rtype: int
        """
        return self._ll_tree_sequence.get_num_individuals()

    @property
    def num_nodes(self):
        """
        Returns the number of :ref:`nodes <sec_node_table_definition>` in
        this tree sequence.

        :return: The number of nodes in this tree sequence.
        :rtype: int
        """
        return self._ll_tree_sequence.get_num_nodes()

    @property
    def num_provenances(self):
        """
        Returns the number of :ref:`provenances <sec_provenance_table_definition>`
        in this tree sequence.

        :return: The number of provenances in this tree sequence.
        :rtype: int
        """
        return self._ll_tree_sequence.get_num_provenances()

    @property
    def num_populations(self):
        """
        Returns the number of :ref:`populations <sec_population_table_definition>`
        in this tree sequence.

        :return: The number of populations in this tree sequence.
        :rtype: int
        """
        return self._ll_tree_sequence.get_num_populations()

    @property
    def num_migrations(self):
        """
        Returns the number of :ref:`migrations <sec_migration_table_definition>`
        in this tree sequence.

        :return: The number of migrations in this tree sequence.
        :rtype: int
        """
        return self._ll_tree_sequence.get_num_migrations()

    @property
    def max_root_time(self):
        """
        Returns the time of the oldest root in any of the trees in this tree sequence.
        This is usually equal to ``np.max(ts.tables.nodes.time)`` but may not be
        since there can be non-sample nodes that are not present in any tree. Note that
        isolated samples are also defined as roots (so there can be a max_root_time
        even in a tree sequence with no edges).

        :return: The maximum time of a root in this tree sequence.
        :rtype: float
        :raises ValueError: If there are no samples in the tree, and hence no roots (as
            roots are defined by the ends of the upward paths from the set of samples).
        """
        if self.num_samples == 0:
            raise ValueError(
                "max_root_time is not defined in a tree sequence with 0 samples"
            )
        ret = max(self.nodes_time[u] for u in self.samples())
        if self.num_edges > 0:
            # Edges are guaranteed to be listed in parent-time order, so we can get the
            # last one to get the oldest root
            edge = self.edge(self.num_edges - 1)
            # However, we can have situations where there is a sample older than a
            # 'proper' root
            ret = max(ret, self.nodes_time[edge.parent])
        return ret

    def migrations(self):
        """
        Returns an iterable sequence of all the
        :ref:`migrations <sec_migration_table_definition>` in this tree sequence.

        Migrations are returned in nondecreasing order of the ``time`` value.

        :return: An iterable sequence of all migrations.
        :rtype: Sequence(:class:`.Migration`)
        """
        return SimpleContainerSequence(self.migration, self.num_migrations)

    def individuals(self):
        """
        Returns an iterable sequence of all the
        :ref:`individuals <sec_individual_table_definition>` in this tree sequence.

        :return: An iterable sequence of all individuals.
        :rtype: Sequence(:class:`.Individual`)
        """
        return SimpleContainerSequence(self.individual, self.num_individuals)

    def nodes(self, *, order=None):
        """
        Returns an iterable sequence of all the :ref:`nodes <sec_node_table_definition>`
        in this tree sequence.

        .. note:: Although node ids are commonly ordered by node time, this is not a
            formal tree sequence requirement. If you wish to iterate over nodes in
            time order, you should therefore use ``order="timeasc"`` (and wrap the
            resulting sequence in the standard Python :func:`python:reversed` function
            if you wish to iterate over older nodes before younger ones)

        :param str order: The order in which the nodes should be returned: must be
            one of "id" (default) or "timeasc" (ascending order of time, then by
            ascending node id, matching the first two ordering requirements of
            parent nodes in a :meth:`sorted <TableCollection.sort>` edge table).
        :return: An iterable sequence of all nodes.
        :rtype: Sequence(:class:`.Node`)
        """
        order = "id" if order is None else order
        if order not in ["id", "timeasc"]:
            raise ValueError('order must be "id" or "timeasc"')
        odr = None
        if order == "timeasc":
            odr = np.lexsort((np.arange(self.num_nodes), self.nodes_time))
        return SimpleContainerSequence(self.node, self.num_nodes, order=odr)

    def edges(self):
        """
        Returns an iterable sequence of all the :ref:`edges <sec_edge_table_definition>`
        in this tree sequence. Edges are returned in the order required
        for a :ref:`valid tree sequence <sec_valid_tree_sequence_requirements>`. So,
        edges are guaranteed to be ordered such that (a) all parents with a
        given ID are contiguous; (b) edges are returned in non-decreasing
        order of parent time ago; (c) within the edges for a given parent, edges
        are sorted first by child ID and then by left coordinate.

        :return: An iterable sequence of all edges.
        :rtype: Sequence(:class:`.Edge`)
        """
        return SimpleContainerSequence(self.edge, self.num_edges)

    def edgesets(self):
        # TODO the order that these records are returned in is not well specified.
        # Hopefully this does not matter, and we can just state that the ordering
        # should not be depended on.
        children = collections.defaultdict(set)
        active_edgesets = {}
        for (left, right), edges_out, edges_in in self.edge_diffs():
            # Complete and return any edgesets that are affected by this tree
            # transition
            parents = iter(edge.parent for edge in itertools.chain(edges_out, edges_in))
            for parent in parents:
                if parent in active_edgesets:
                    edgeset = active_edgesets.pop(parent)
                    edgeset.right = left
                    edgeset.children = sorted(children[parent])
                    yield edgeset
            for edge in edges_out:
                children[edge.parent].remove(edge.child)
            for edge in edges_in:
                children[edge.parent].add(edge.child)
            # Update the active edgesets
            for edge in itertools.chain(edges_out, edges_in):
                if (
                    len(children[edge.parent]) > 0
                    and edge.parent not in active_edgesets
                ):
                    active_edgesets[edge.parent] = Edgeset(left, right, edge.parent, [])

        for parent in active_edgesets.keys():
            edgeset = active_edgesets[parent]
            edgeset.right = self.sequence_length
            edgeset.children = sorted(children[edgeset.parent])
            yield edgeset

    def _edge_diffs_forward(self, include_terminal=False):
        metadata_decoder = self.table_metadata_schemas.edge.decode_row
        edge_left = self.edges_left
        edge_right = self.edges_right
        sequence_length = self.sequence_length
        in_order = self.indexes_edge_insertion_order
        out_order = self.indexes_edge_removal_order
        M = self.num_edges
        j = 0
        k = 0
        left = 0.0
        while j < M or left < sequence_length:
            edges_out = []
            edges_in = []
            while k < M and edge_right[out_order[k]] == left:
                edges_out.append(
                    Edge(
                        *self._ll_tree_sequence.get_edge(out_order[k]),
                        id=out_order[k],
                        metadata_decoder=metadata_decoder,
                    )
                )
                k += 1
            while j < M and edge_left[in_order[j]] == left:
                edges_in.append(
                    Edge(
                        *self._ll_tree_sequence.get_edge(in_order[j]),
                        id=in_order[j],
                        metadata_decoder=metadata_decoder,
                    )
                )
                j += 1
            right = sequence_length
            if j < M:
                right = min(right, edge_left[in_order[j]])
            if k < M:
                right = min(right, edge_right[out_order[k]])
            yield EdgeDiff(Interval(left, right), edges_out, edges_in)
            left = right

        if include_terminal:
            edges_out = []
            while k < M:
                edges_out.append(
                    Edge(
                        *self._ll_tree_sequence.get_edge(out_order[k]),
                        id=out_order[k],
                        metadata_decoder=metadata_decoder,
                    )
                )
                k += 1
            yield EdgeDiff(Interval(left, right), edges_out, [])

    def _edge_diffs_reverse(self, include_terminal=False):
        metadata_decoder = self.table_metadata_schemas.edge.decode_row
        edge_left = self.edges_left
        edge_right = self.edges_right
        sequence_length = self.sequence_length
        in_order = self.indexes_edge_removal_order
        out_order = self.indexes_edge_insertion_order
        M = self.num_edges
        j = M - 1
        k = M - 1
        right = sequence_length
        while j >= 0 or right > 0:
            edges_out = []
            edges_in = []
            while k >= 0 and edge_left[out_order[k]] == right:
                edges_out.append(
                    Edge(
                        *self._ll_tree_sequence.get_edge(out_order[k]),
                        id=out_order[k],
                        metadata_decoder=metadata_decoder,
                    )
                )
                k -= 1
            while j >= 0 and edge_right[in_order[j]] == right:
                edges_in.append(
                    Edge(
                        *self._ll_tree_sequence.get_edge(in_order[j]),
                        id=in_order[j],
                        metadata_decoder=metadata_decoder,
                    )
                )
                j -= 1
            left = 0
            if j >= 0:
                left = max(left, edge_right[in_order[j]])
            if k >= 0:
                left = max(left, edge_left[out_order[k]])
            yield EdgeDiff(Interval(left, right), edges_out, edges_in)
            right = left

        if include_terminal:
            edges_out = []
            while k >= 0:
                edges_out.append(
                    Edge(
                        *self._ll_tree_sequence.get_edge(out_order[k]),
                        id=out_order[k],
                        metadata_decoder=metadata_decoder,
                    )
                )
                k -= 1
            yield EdgeDiff(Interval(left, right), edges_out, [])

    def edge_diffs(self, include_terminal=False, *, direction=tskit.FORWARD):
        """
        Returns an iterator over all the :ref:`edges <sec_edge_table_definition>` that
        are inserted and removed to build the trees as we move from left-to-right along
        the tree sequence. Each iteration yields a named tuple consisting of 3 values,
        ``(interval, edges_out, edges_in)``. The first value, ``interval``, is the
        genomic interval ``(left, right)`` covered by the incoming tree
        (see :attr:`Tree.interval`). The second, ``edges_out`` is a list of the edges
        that were just-removed to create the tree covering the interval
        (hence ``edges_out`` will always be empty for the first tree). The last value,
        ``edges_in``, is a list of edges that were just
        inserted to construct the tree covering the current interval.

        The edges returned within each ``edges_in`` list are ordered by ascending
        time of the parent node, then ascending parent id, then ascending child id.
        The edges within each ``edges_out`` list are the reverse order (e.g.
        descending parent time, parent id, then child_id). This means that within
        each list, edges with the same parent appear consecutively.

        The ``direction`` argument can be used to control whether diffs are produced
        in the forward (left-to-right, increasing genome coordinate value)
        or reverse (right-to-left, decreasing genome coordinate value) direction.

        :param bool include_terminal: If False (default), the iterator terminates
            after the final interval in the tree sequence (i.e., it does not
            report a final removal of all remaining edges), and the number
            of iterations will be equal to the number of trees in the tree
            sequence. If True, an additional iteration takes place, with the last
            ``edges_out`` value reporting all the edges contained in the final
            tree (with both ``left`` and ``right`` equal to the sequence length).
        :param int direction: The direction of travel along the sequence for
            diffs. Must be one of :data:`.FORWARD` or :data:`.REVERSE`.
            (Default: :data:`.FORWARD`).
        :return: An iterator over the (interval, edges_out, edges_in) tuples. This
            is a named tuple, so the 3 values can be accessed by position
            (e.g. ``returned_tuple[0]``) or name (e.g. ``returned_tuple.interval``).
        :rtype: :class:`collections.abc.Iterable`
        """
        if direction == _tskit.FORWARD:
            return self._edge_diffs_forward(include_terminal=include_terminal)
        elif direction == _tskit.REVERSE:
            return self._edge_diffs_reverse(include_terminal=include_terminal)
        else:
            raise ValueError("direction must be either tskit.FORWARD or tskit.REVERSE")

    def sites(self):
        """
        Returns an iterable sequence of all the :ref:`sites <sec_site_table_definition>`
        in this tree sequence. Sites are returned in order of increasing ID
        (and also position). See the :class:`Site` class for details on
        the available fields for each site.

        :return: An iterable sequence of all sites.
        :rtype: Sequence(:class:`.Site`)
        """
        return SimpleContainerSequence(self.site, self.num_sites)

    def mutations(self):
        """
        Returns an iterator over all the
        :ref:`mutations <sec_mutation_table_definition>` in this tree sequence.
        Mutations are returned in order of nondecreasing site ID.
        See the :class:`Mutation` class for details on the available fields for
        each mutation.

        The returned iterator is equivalent to iterating over all sites
        and all mutations in each site, i.e.::

            >>> for site in tree_sequence.sites():
            >>>     for mutation in site.mutations:
            >>>         yield mutation

        :return: An iterator over all mutations in this tree sequence.
        :rtype: iter(:class:`Mutation`)
        """
        for site in self.sites():
            yield from site.mutations

    def populations(self):
        """
        Returns an iterable sequence of all the
        :ref:`populations <sec_population_table_definition>` in this tree sequence.

        :return: An iterable sequence of all populations.
        :rtype: Sequence(:class:`.Population`)
        """
        return SimpleContainerSequence(self.population, self.num_populations)

    def provenances(self):
        """
        Returns an iterable sequence of all the
        :ref:`provenances <sec_provenance_table_definition>` in this tree sequence.

        :return: An iterable sequence of all provenances.
        :rtype: Sequence(:class:`.Provenance`)
        """
        return SimpleContainerSequence(self.provenance, self.num_provenances)

    def breakpoints(self, as_array=False):
        """
        Returns the breakpoints that separate trees along the chromosome, including the
        two extreme points 0 and L. This is equivalent to

        >>> iter([0] + [t.interval.right for t in self.trees()])

        By default we return an iterator over the breakpoints as Python float objects;
        if ``as_array`` is True we return them as a numpy array.

        Note that the ``as_array`` form will be more efficient and convenient in most
        cases; the default iterator behaviour is mainly kept to ensure compatibility
        with existing code.

        :param bool as_array: If True, return the breakpoints as a numpy array.
        :return: The breakpoints defined by the tree intervals along the sequence.
        :rtype: collections.abc.Iterable or numpy.ndarray
        """
        breakpoints = self.ll_tree_sequence.get_breakpoints()
        if not as_array:
            # Convert to Python floats for backward compatibility.
            breakpoints = map(float, breakpoints)
        return breakpoints

    def at(self, position, **kwargs):
        """
        Returns the tree covering the specified genomic location. The returned tree
        will have ``tree.interval.left`` <= ``position`` < ``tree.interval.right``.
        See also :meth:`Tree.seek`.

        .. include:: substitutions/linear_traversal_warning.rst

        :param float position: A genomic location.
        :param \\**kwargs: Further arguments used as parameters when constructing the
            returned :class:`Tree`. For example ``ts.at(2.5, sample_lists=True)`` will
            result in a :class:`Tree` created with ``sample_lists=True``.
        :return: A new instance of :class:`Tree` positioned to cover the specified
            genomic location.
        :rtype: Tree
        """
        tree = Tree(self, **kwargs)
        tree.seek(position)
        return tree

    def at_index(self, index, **kwargs):
        """
        Returns the tree at the specified index. See also :meth:`Tree.seek_index`.

        .. include:: substitutions/linear_traversal_warning.rst

        :param int index: The index of the required tree.
        :param \\**kwargs: Further arguments used as parameters when constructing the
            returned :class:`Tree`. For example ``ts.at_index(4, sample_lists=True)``
            will result in a :class:`Tree` created with ``sample_lists=True``.
        :return: A new instance of :class:`Tree` positioned at the specified index.
        :rtype: Tree
        """
        tree = Tree(self, **kwargs)
        tree.seek_index(index)
        return tree

    def first(self, **kwargs):
        """
        Returns the first tree in this :class:`TreeSequence`. To iterate over all
        trees in the sequence, use the :meth:`.trees` method.

        :param \\**kwargs: Further arguments used as parameters when constructing the
            returned :class:`Tree`. For example ``ts.first(sample_lists=True)`` will
            result in a :class:`Tree` created with ``sample_lists=True``.
        :return: The first tree in this tree sequence.
        :rtype: :class:`Tree`.
        """
        tree = Tree(self, **kwargs)
        tree.first()
        return tree

    def last(self, **kwargs):
        """
        Returns the last tree in this :class:`TreeSequence`. To iterate over all
        trees in the sequence, use the :meth:`.trees` method.

        :param \\**kwargs: Further arguments used as parameters when constructing the
            returned :class:`Tree`. For example ``ts.first(sample_lists=True)`` will
            result in a :class:`Tree` created with ``sample_lists=True``.
        :return: The last tree in this tree sequence.
        :rtype: :class:`Tree`.
        """
        tree = Tree(self, **kwargs)
        tree.last()
        return tree

    def trees(
        self,
        tracked_samples=None,
        *,
        sample_lists=False,
        root_threshold=1,
        sample_counts=None,
        tracked_leaves=None,
        leaf_counts=None,
        leaf_lists=None,
    ):
        """
        Returns an iterator over the trees in this tree sequence. Each value
        returned in this iterator is an instance of :class:`Tree`. Upon
        successful termination of the iterator, the tree will be in the
        "cleared" null state.

        The ``sample_lists`` and ``tracked_samples`` parameters are passed
        to the :class:`Tree` constructor, and control
        the options that are set in the returned tree instance.

        .. warning:: Do not store the results of this iterator in a list!
           For performance reasons, the same underlying object is used
           for every tree returned which will most likely lead to unexpected
           behaviour. If you wish to obtain a list of trees in a tree sequence
           please use ``ts.aslist()`` instead.

        :param list tracked_samples: The list of samples to be tracked and
            counted using the :meth:`Tree.num_tracked_samples` method.
        :param bool sample_lists: If True, provide more efficient access
            to the samples beneath a given node using the
            :meth:`Tree.samples` method.
        :param int root_threshold: The minimum number of samples that a node
            must be ancestral to for it to be in the list of roots. By default
            this is 1, so that isolated samples (representing missing data)
            are roots. To efficiently restrict the roots of the tree to
            those subtending meaningful topology, set this to 2. This value
            is only relevant when trees have multiple roots.
        :param bool sample_counts: Deprecated since 0.2.4.
        :return: An iterator over the Trees in this tree sequence.
        :rtype: collections.abc.Iterable, :class:`Tree`
        """
        # tracked_leaves, leaf_counts and leaf_lists are deprecated aliases
        # for tracked_samples, sample_counts and sample_lists respectively.
        # These are left over from an older version of the API when leaves
        # and samples were synonymous.
        if tracked_leaves is not None:
            tracked_samples = tracked_leaves
        if leaf_counts is not None:
            sample_counts = leaf_counts
        if leaf_lists is not None:
            sample_lists = leaf_lists
        tree = Tree(
            self,
            tracked_samples=tracked_samples,
            sample_lists=sample_lists,
            root_threshold=root_threshold,
            sample_counts=sample_counts,
        )
        return TreeIterator(tree)

    def coiterate(self, other, **kwargs):
        """
        Returns an iterator over the pairs of trees for each distinct
        interval in the specified pair of tree sequences.

        :param TreeSequence other: The other tree sequence from which to take trees. The
            sequence length must be the same as the current tree sequence.
        :param \\**kwargs: Further named arguments that will be passed to the
            :meth:`.trees` method when constructing the returned trees.

        :return: An iterator returning successive tuples of the form
            ``(interval, tree_self, tree_other)``. For example, the first item returned
            will consist of an tuple of the initial interval, the first tree of the
            current tree sequence, and the first tree of the ``other`` tree sequence;
            the ``.left`` attribute of the initial interval will be 0 and the ``.right``
            attribute will be the smallest non-zero breakpoint of the 2 tree sequences.
        :rtype: iter(:class:`Interval`, :class:`Tree`, :class:`Tree`)

        """
        if self.sequence_length != other.sequence_length:
            raise ValueError("Tree sequences must be of equal sequence length.")
        L = self.sequence_length
        trees1 = self.trees(**kwargs)
        trees2 = other.trees(**kwargs)
        tree1 = next(trees1)
        tree2 = next(trees2)
        right = 0
        while right != L:
            left = right
            right = min(tree1.interval.right, tree2.interval.right)
            yield Interval(left, right), tree1, tree2
            # Advance
            if tree1.interval.right == right:
                tree1 = next(trees1, None)
            if tree2.interval.right == right:
                tree2 = next(trees2, None)

    def _check_genomic_range(self, left, right, ensure_integer=False):
        if left is None:
            left = 0
        if right is None:
            right = self.sequence_length
        if np.isnan(left) or left < 0 or left >= self.sequence_length:
            raise ValueError(
                "`left` not between zero (inclusive) and sequence length (exclusive)"
            )
        if np.isnan(right) or right <= 0 or right > self.sequence_length:
            raise ValueError(
                "`right` not between zero (exclusive) and sequence length (inclusive)"
            )
        if left >= right:
            raise ValueError("`left` must be less than `right`")
        if ensure_integer:
            if left != int(left) or right != int(right):
                raise ValueError("`left` and `right` must be integers")
            return Interval(int(left), int(right))
        return Interval(left, right)

    def _haplotypes_array(
        self,
        *,
        interval,
        isolated_as_missing=None,
        missing_data_character=None,
        samples=None,
    ):
        # return an array of haplotypes and the first and last site positions
        if missing_data_character is None:
            missing_data_character = "N"

        start_site, stop_site = np.searchsorted(self.sites_position, interval)
        H = np.empty(
            (
                self.num_samples if samples is None else len(samples),
                stop_site - start_site,
            ),
            dtype=np.int8,
        )
        missing_int8 = ord(missing_data_character.encode("ascii"))
        for var in self.variants(
            samples=samples,
            isolated_as_missing=isolated_as_missing,
            left=interval.left,
            right=interval.right,
        ):
            alleles = np.full(len(var.alleles), missing_int8, dtype=np.int8)
            for i, allele in enumerate(var.alleles):
                if allele is not None:
                    if len(allele) != 1:
                        raise TypeError(
                            "Multi-letter allele or deletion detected at site {}".format(
                                var.site.id
                            )
                        )
                    try:
                        ascii_allele = allele.encode("ascii")
                    except UnicodeEncodeError:
                        raise TypeError(
                            "Non-ascii character in allele at site {}".format(
                                var.site.id
                            )
                        )
                    allele_int8 = ord(ascii_allele)
                    if allele_int8 == missing_int8:
                        raise ValueError(
                            "The missing data character '{}' clashes with an "
                            "existing allele at site {}".format(
                                missing_data_character, var.site.id
                            )
                        )
                    alleles[i] = allele_int8
            H[:, var.site.id - start_site] = alleles[var.genotypes]
        return H, (start_site, stop_site - 1)

    def haplotypes(
        self,
        *,
        isolated_as_missing=None,
        missing_data_character=None,
        samples=None,
        left=None,
        right=None,
        impute_missing_data=None,
    ):
        """
        Returns an iterator over the strings of haplotypes that result from
        the trees and mutations in this tree sequence. Each haplotype string
        is guaranteed to be of the same length. A tree sequence with
        :math:`n` samples and with :math:`s` sites lying between ``left`` and
        ``right`` will return a total of :math:`n`
        strings of :math:`s` alleles concatenated together, where an allele
        consists of a single ascii character (tree sequences that include alleles
        which are not a single character in length, or where the character is
        non-ascii, will raise an error). The first string returned is the
        haplotype for the first requested sample, and so on.

        The alleles at each site must be represented by single byte characters,
        (i.e., variants must be single nucleotide polymorphisms, or SNPs), hence
        the strings returned will all be of length :math:`s`. If the ``left``
        position is less than or equal to the position of the first site, for a
        haplotype ``h``, the value of ``h[j]`` will therefore be the observed
        allelic state at site ``j``.

        If ``isolated_as_missing`` is True (the default), isolated samples without
        mutations directly above them will be treated as
        :ref:`missing data<sec_data_model_missing_data>` and will be
        represented in the string by the ``missing_data_character``. If
        instead it is set to False, missing data will be assigned the ancestral state
        (unless they have mutations directly above them, in which case they will take
        the most recent derived mutational state for that node). This was the default
        behaviour in versions prior to 0.2.0. Prior to 0.3.0 the `impute_missing_data`
        argument controlled this behaviour.

        See also the :meth:`.variants` iterator for site-centric access
        to sample genotypes.

        .. warning::
            For large datasets, this method can consume a **very large** amount of
            memory! To output all the sample data, it is more efficient to iterate
            over sites rather than over samples.

        :return: An iterator over the haplotype strings for the samples in
            this tree sequence.
        :param bool isolated_as_missing: If True, the allele assigned to
            missing samples (i.e., isolated samples without mutations) is
            the ``missing_data_character``. If False,
            missing samples will be assigned the ancestral state.
            Default: True.
        :param str missing_data_character: A single ascii character that will
            be used to represent missing data.
            If any normal allele contains this character, an error is raised.
            Default: 'N'.
        :param list[int] samples: The samples for which to output haplotypes. If
            ``None`` (default), return haplotypes for all the samples in the tree
            sequence, in the order given by the :meth:`.samples` method.
        :param int left: Haplotype strings will start with the first site at or after
            this genomic position. If ``None`` (default) start at the first site.
        :param int right: Haplotype strings will end with the last site before this
            position. If ``None`` (default) assume ``right`` is the sequence length
            (i.e. the last character in the string will be the last site in the tree
            sequence).
        :param bool impute_missing_data:
            *Deprecated in 0.3.0. Use ``isolated_as_missing``, but inverting value.
            Will be removed in a future version*
        :rtype: collections.abc.Iterable
        :raises TypeError: if the ``missing_data_character`` or any of the alleles
            at a site are not a single ascii character.
        :raises ValueError: if the ``missing_data_character`` exists in one of the
            alleles
        """
        if impute_missing_data is not None:
            warnings.warn(
                "The impute_missing_data parameter was deprecated in 0.3.0 and will"
                " be removed. Use ``isolated_as_missing=False`` instead of"
                "``impute_missing_data=True``.",
                FutureWarning,
                stacklevel=4,
            )
        # Only use impute_missing_data if isolated_as_missing has the default value
        if isolated_as_missing is None:
            isolated_as_missing = not impute_missing_data
        interval = self._check_genomic_range(left, right)
        H, _ = self._haplotypes_array(
            interval=interval,
            isolated_as_missing=isolated_as_missing,
            missing_data_character=missing_data_character,
            samples=samples,
        )
        for h in H:
            yield h.tobytes().decode("ascii")

    def variants(
        self,
        *,
        samples=None,
        isolated_as_missing=None,
        alleles=None,
        impute_missing_data=None,
        copy=None,
        left=None,
        right=None,
    ):
        """
        Returns an iterator over the variants between the ``left`` (inclusive)
        and ``right`` (exclusive) genomic positions in this tree sequence. Each
        returned :class:`Variant` object has a site, a list of possible allelic
        states and an array of genotypes for the specified ``samples``. The
        ``genotypes`` value is a numpy array containing indexes into the
        ``alleles`` list. By default, this list is generated automatically for
        each site such that the first entry, ``alleles[0]``, is the ancestral
        state and subsequent alleles are listed in no
        particular order. This means that the encoding of alleles in
        terms of genotype values can vary from site-to-site, which is
        sometimes inconvenient. It is possible to specify a fixed mapping
        from allele strings to genotype values using the ``alleles``
        parameter. For example, if we set ``alleles=("A", "C", "G", "T")``,
        this will map allele "A" to 0, "C" to 1 and so on (the
        :data:`ALLELES_ACGT` constant provides a shortcut for this
        common mapping).

        By default, genotypes are generated for all samples. The ``samples``
        parameter allows us to specify the nodes for which genotypes are
        generated; output order of genotypes in the returned variants
        corresponds to the order of the samples in this list. It is also
        possible to provide **non-sample** nodes as an argument here, if you
        wish to generate genotypes for (e.g.) internal nodes. However,
        ``isolated_as_missing`` must be False in this case, as it is not
        possible to detect missing data for non-sample nodes.

        If isolated samples are present at a given site without mutations above them,
        they are interpreted by default as
        :ref:`missing data<sec_data_model_missing_data>`, and the genotypes array
        will contain a special value :data:`MISSING_DATA` (-1) to identify them
        while the ``alleles`` tuple will end with the value ``None`` (note that this
        will be the case whether or not we specify a fixed mapping using the
        ``alleles`` parameter; see the :class:`Variant` class for more details).
        Alternatively, if ``isolated_as_missing`` is set to to False, such isolated
        samples will not be treated as missing, and instead assigned the ancestral
        state (this was the default behaviour in versions prior to 0.2.0). Prior to
        0.3.0 the `impute_missing_data` argument controlled this behaviour.

        :param array_like samples: An array of node IDs for which to generate
            genotypes, or None for all sample nodes. Default: None.
        :param bool isolated_as_missing: If True, the genotype value assigned to
            missing samples (i.e., isolated samples without mutations) is
            :data:`.MISSING_DATA` (-1). If False, missing samples will be
            assigned the allele index for the ancestral state.
            Default: True.
        :param tuple alleles: A tuple of strings defining the encoding of
            alleles as integer genotype values. At least one allele must be provided.
            If duplicate alleles are provided, output genotypes will always be
            encoded as the first occurrence of the allele. If None (the default),
            the alleles are encoded as they are encountered during genotype
            generation.
        :param bool impute_missing_data:
            *Deprecated in 0.3.0. Use ``isolated_as_missing``, but inverting value.
            Will be removed in a future version*
        :param bool copy:
            If False re-use the same Variant object for each site such that any
            references held to it are overwritten when the next site is visited.
            If True return a fresh :class:`Variant` for each site. Default: True.
        :param int left: Start with the first site at or after
            this genomic position. If ``None`` (default) start at the first site.
        :param int right: End with the last site before this position. If ``None``
            (default) assume ``right`` is the sequence length, so that the last
            variant corresponds to the last site in the tree sequence.
        :return: An iterator over all variants in this tree sequence.
        :rtype: iter(:class:`Variant`)
        """
        interval = self._check_genomic_range(left, right)
        if impute_missing_data is not None:
            warnings.warn(
                "The impute_missing_data parameter was deprecated in 0.3.0 and will"
                " be removed. Use ``isolated_as_missing=False`` instead of"
                "``impute_missing_data=True``.",
                FutureWarning,
                stacklevel=4,
            )
        # Only use impute_missing_data if isolated_as_missing has the default value
        if isolated_as_missing is None:
            isolated_as_missing = not impute_missing_data
        if copy is None:
            copy = True
        # See comments for the Variant type for discussion on why the
        # present form was chosen.
        variant = tskit.Variant(
            self,
            samples=samples,
            isolated_as_missing=isolated_as_missing,
            alleles=alleles,
        )
        if left == 0 and right == self.sequence_length:
            start = 0
            stop = self.num_sites
        else:
            start, stop = np.searchsorted(self.sites_position, interval)

        if copy:
            for site_id in range(start, stop):
                variant.decode(site_id)
                yield variant.copy()
        else:
            for site_id in range(start, stop):
                variant.decode(site_id)
                yield variant

    def genotype_matrix(
        self,
        *,
        samples=None,
        isolated_as_missing=None,
        alleles=None,
        impute_missing_data=None,
    ):
        """
        Returns an :math:`m \\times n` numpy array of the genotypes in this
        tree sequence, where :math:`m` is the number of sites and :math:`n`
        the number of samples. The genotypes are the indexes into the array
        of ``alleles``, as described for the :class:`Variant` class.

        If isolated samples are present at a given site without mutations above them,
        they will be interpreted as :ref:`missing data<sec_data_model_missing_data>`
        the genotypes array will contain a special value :data:`MISSING_DATA`
        (-1) to identify these missing samples.

        Such samples are treated as missing data by default, but if
        ``isolated_as_missing`` is set to to False, they will not be treated as missing,
        and so assigned the ancestral state. This was the default behaviour in
        versions prior to 0.2.0. Prior to 0.3.0 the `impute_missing_data`
        argument controlled this behaviour.

        .. warning::
            This method can consume a **very large** amount of memory! If
            all genotypes are not needed at once, it is usually better to
            access them sequentially using the :meth:`.variants` iterator.

        :param array_like samples: An array of node IDs for which to generate
            genotypes, or None for all sample nodes. Default: None.
        :param bool isolated_as_missing: If True, the genotype value assigned to
            missing samples (i.e., isolated samples without mutations) is
            :data:`.MISSING_DATA` (-1). If False, missing samples will be
            assigned the allele index for the ancestral state.
            Default: True.
        :param tuple alleles: A tuple of strings describing the encoding of
            alleles to genotype values. At least one allele must be provided.
            If duplicate alleles are provided, output genotypes will always be
            encoded as the first occurrence of the allele. If None (the default),
            the alleles are encoded as they are encountered during genotype
            generation.
        :param bool impute_missing_data:
            *Deprecated in 0.3.0. Use ``isolated_as_missing``, but inverting value.
            Will be removed in a future version*

        :return: The full matrix of genotypes.
        :rtype: numpy.ndarray (dtype=np.int32)
        """
        if impute_missing_data is not None:
            warnings.warn(
                "The impute_missing_data parameter was deprecated in 0.3.0 and will"
                " be removed. Use ``isolated_as_missing=False`` instead of"
                "``impute_missing_data=True``.",
                FutureWarning,
                stacklevel=4,
            )
        # Only use impute_missing_data if isolated_as_missing has the default value
        if isolated_as_missing is None:
            isolated_as_missing = not impute_missing_data

        variant = tskit.Variant(
            self,
            samples=samples,
            isolated_as_missing=isolated_as_missing,
            alleles=alleles,
        )

        num_samples = self.num_samples if samples is None else len(samples)
        ret = np.zeros(shape=(self.num_sites, num_samples), dtype=np.int32)

        for site_id in range(self.num_sites):
            variant.decode(site_id)
            ret[site_id, :] = variant.genotypes

        return ret

    def alignments(
        self,
        *,
        reference_sequence=None,
        missing_data_character=None,
        samples=None,
        left=None,
        right=None,
    ):
        """
        Returns an iterator over the full sequence alignments for the defined samples
        in this tree sequence. Each alignment ``a`` is a string of length ``L`` where
        the first character is the genomic sequence at the ``start`` position in the
        genome (defaulting to 0) and the last character is the genomic sequence one
        position before the ``stop`` value (defaulting to the :attr:`.sequence_length`
        of this tree sequence, which must have :attr:`.discrete_genome` equal to True).
        By default ``L`` is therefore equal to the :attr:`.sequence_length`,
        and ``a[j]`` is the nucleotide value at genomic position ``j``.

        .. note:: This is inherently a **zero-based** representation of the sequence
            coordinate space. Care will be needed when interacting with other
            libraries and upstream coordinate spaces.


        The :ref:`sites<sec_data_model_definitions_site>` in a tree sequence will
        usually only define the variation for a subset of the ``L`` nucleotide
        positions along the genome, and the remaining positions are filled using
        a :ref:`reference sequence <sec_data_model_reference_sequence>`.
        The reference sequence data is defined either via the
        ``reference_sequence`` parameter to this method, or embedded within
        with the tree sequence itself via the :attr:`.TreeSequence.reference_sequence`.

        Site information from the tree sequence takes precedence over the reference
        sequence so that, for example, at a site with no mutations all samples
        will have the site's ancestral state.

        The reference sequence bases are determined in the following way:

        - If the ``reference_sequence`` parameter is supplied this will be
          used, regardless of whether the tree sequence has an embedded
          reference sequence.
        - Otherwise, if the tree sequence has an embedded reference sequence,
          this will be used.
        - If the ``reference_sequence`` parameter is not specified and
          there is no embedded reference sequence, ``L`` copies of the
          ``missing_data_character`` (which defaults to 'N') are used
          instead.

        .. warning:: The :class:`.ReferenceSequence` API is preliminary and
           some behaviours may change in the future. In particular, a
           tree sequence is currently regarded as having an embedded reference
           sequence even if it only has some metadata defined. In this case
           the ``reference_sequence`` parameter will need to be explicitly set.

        .. note::
            Two common options for setting a reference sequence are:

            - Mark them as missing data, by setting
              ``reference_sequence="N" * int(ts.sequence_length)``
            - Fill the gaps with random nucleotides, by setting
              ``reference_sequence=tskit.random_nucleotides(ts.sequence_length)``.
              See the :func:`.random_nucleotides` function for more information.

        .. warning:: Insertions and deletions are not currently supported and
           the alleles at each site must be represented by
           single byte characters, (i.e., variants must be single nucleotide
           polymorphisms, or SNPs).

        .. warning:: :ref:`Missing data<sec_data_model_missing_data>` is not
           currently supported by this method and it will raise a ValueError
           if called on tree sequences containing isolated samples.
           See https://github.com/tskit-dev/tskit/issues/1896 for more
           information.

        See also the :meth:`.variants` iterator for site-centric access
        to sample genotypes and :meth:`.haplotypes` for access to sample sequences
        at just the sites in the tree sequence.

        :param str reference_sequence: The reference sequence to fill in
            gaps between sites in the alignments.
        :param str missing_data_character: A single ascii character that will
            be used to represent missing data.
            If any normal allele contains this character, an error is raised.
            Default: 'N'.
        :param list[int] samples: The samples for which to output alignments. If
            ``None`` (default), return alignments for all the samples in the tree
            sequence, in the order given by the :meth:`.samples` method.
        :param int left: Alignments will start at this genomic position. If ``None``
            (default) alignments start at 0.
        :param int right: Alignments will stop before this genomic position. If ``None``
            (default) alignments will continue until the end of the tree sequence.
        :return: An iterator over the alignment strings for specified samples in
            this tree sequence, in the order given in ``samples``.
        :rtype: collections.abc.Iterable
        :raises ValueError: if any genome coordinate in this tree sequence is not
            discrete, or if the ``reference_sequence`` is not of the correct length.
        :raises TypeError: if any of the alleles at a site are not a
            single ascii character.
        """
        if not self.discrete_genome:
            raise ValueError("sequence alignments only defined for discrete genomes")
        interval = self._check_genomic_range(left, right, ensure_integer=True)
        missing_data_character = (
            "N" if missing_data_character is None else missing_data_character
        )

        L = interval.span
        a = np.empty(L, dtype=np.int8)
        if reference_sequence is None:
            if self.has_reference_sequence():
                # This may be inefficient - see #1989. However, since we're
                # n copies of the reference sequence anyway, this is a relatively
                # minor tweak. We may also want to recode the below not to use direct
                # access to the .data attribute, e.g. if we allow reference sequences
                # to start at non-zero positions
                reference_sequence = self.reference_sequence.data[
                    interval.left : interval.right
                ]
            else:
                reference_sequence = missing_data_character * L

        if len(reference_sequence) != L:
            if interval.right == int(self.sequence_length):
                raise ValueError(
                    "The reference sequence is shorter than the tree sequence length"
                )
            else:
                raise ValueError(
                    "The reference sequence ends before the requested stop position"
                )
        ref_bytes = reference_sequence.encode("ascii")
        a[:] = np.frombuffer(ref_bytes, dtype=np.int8)

        # To do this properly we'll have to detect the missing data as
        # part of a full implementation of alignments in C. The current
        # definition might not be calling some degenerate cases correctly;
        # see https://github.com/tskit-dev/tskit/issues/1908
        #
        # Note also that this will call the presence of missing data
        # incorrectly if have a sample isolated over the region (a, b],
        # and if we have sites at each position from a to b, and at
        # each site there is a mutation over the isolated sample.
        if any(tree._has_isolated_samples() for tree in self.trees()):
            raise ValueError(
                "Missing data not currently supported in alignments; see "
                "https://github.com/tskit-dev/tskit/issues/1896 for details."
                "The current implementation may also incorrectly identify an "
                "input tree sequence has having missing data."
            )
        H, (first_site_id, last_site_id) = self._haplotypes_array(
            interval=interval,
            missing_data_character=missing_data_character,
            samples=samples,
        )
        site_pos = self.sites_position.astype(np.int64)[
            first_site_id : last_site_id + 1
        ]
        for h in H:
            a[site_pos - interval.left] = h
            yield a.tobytes().decode("ascii")

    @property
    def individuals_population(self):
        """
        Returns the length-``num_individuals`` array containing, for each
        individual, the ``population`` attribute of their nodes, or
        ``tskit.NULL`` for individuals with no nodes. Errors if any individual
        has nodes with inconsistent non-NULL populations.
        """
        if self._individuals_population is None:
            self._individuals_population = (
                self._ll_tree_sequence.get_individuals_population()
            )
        return self._individuals_population

    @property
    def individual_populations(self):
        # Undocumented alias for individuals_population to avoid breaking
        # pre-1.0 pyslim code
        return self.individuals_population

    @property
    def individuals_time(self):
        """
        Returns the length-``num_individuals`` array containing, for each
        individual, the ``time`` attribute of their nodes or ``np.nan`` for
        individuals with no nodes. Errors if any individual has nodes with
        inconsistent times.
        """
        if self._individuals_time is None:
            self._individuals_time = self._ll_tree_sequence.get_individuals_time()
        return self._individuals_time

    @property
    def individual_times(self):
        # Undocumented alias for individuals_time to avoid breaking
        # pre-1.0 pyslim code
        return self.individuals_time

    @property
    def individuals_location(self):
        """
        Convenience method returning the ``num_individuals x n`` array
        whose row k-th row contains the ``location`` property of the k-th
        individual. The method only works if all individuals' locations
        have the same length (which is ``n``), and errors otherwise.
        """
        if self._individuals_location is None:
            individuals = self.tables.individuals
            n = 0
            lens = np.unique(np.diff(individuals.location_offset))
            if len(lens) > 1:
                raise ValueError("Individual locations are not all the same length.")
            if len(lens) > 0:
                n = lens[0]
            self._individuals_location = individuals.location.reshape(
                (self.num_individuals, n)
            )
        return self._individuals_location

    @property
    def individual_locations(self):
        # Undocumented alias for individuals_time to avoid breaking
        # pre-1.0 pyslim code
        return self.individuals_location

    @property
    def individuals_flags(self):
        """
        Efficient access to the bitwise ``flags`` column in the
        :ref:`sec_individual_table_definition` as a numpy array (dtype=np.uint32).
        Equivalent to ``ts.tables.individuals.flags`` (but avoiding the full copy
        of the table data that accessing ``ts.tables`` currently entails).
        """
        return self._individuals_flags

    @property
    def nodes_time(self):
        """
        Efficient access to the ``time`` column in the
        :ref:`sec_node_table_definition` as a numpy array (dtype=np.float64).
        Equivalent to ``ts.tables.nodes.time`` (but avoiding the full copy
        of the table data that accessing ``ts.tables`` currently entails).
        """
        return self._nodes_time

    @property
    def nodes_flags(self):
        """
        Efficient access to the bitwise ``flags`` column in the
        :ref:`sec_node_table_definition` as a numpy array (dtype=np.uint32).
        Equivalent to ``ts.tables.nodes.flags`` (but avoiding the full copy
        of the table data that accessing ``ts.tables`` currently entails).
        """
        return self._nodes_flags

    @property
    def nodes_population(self):
        """
        Efficient access to the ``population`` column in the
        :ref:`sec_node_table_definition` as a numpy array (dtype=np.int32).
        Equivalent to ``ts.tables.nodes.population`` (but avoiding the full copy
        of the table data that accessing ``ts.tables`` currently entails).
        """
        return self._nodes_population

    @property
    def nodes_individual(self):
        """
        Efficient access to the ``individual`` column in the
        :ref:`sec_node_table_definition` as a numpy array (dtype=np.int32).
        Equivalent to ``ts.tables.nodes.individual`` (but avoiding the full copy
        of the table data that accessing ``ts.tables`` currently entails).
        """
        return self._nodes_individual

    @property
    def edges_left(self):
        """
        Efficient access to the ``left`` column in the
        :ref:`sec_edge_table_definition` as a numpy array (dtype=np.float64).
        Equivalent to ``ts.tables.edges.left`` (but avoiding the full copy
        of the table data that accessing ``ts.tables`` currently entails).
        """
        return self._edges_left

    @property
    def edges_right(self):
        """
        Efficient access to the ``right`` column in the
        :ref:`sec_edge_table_definition` as a numpy array (dtype=np.float64).
        Equivalent to ``ts.tables.edges.right`` (but avoiding the full copy
        of the table data that accessing ``ts.tables`` currently entails).
        """
        return self._edges_right

    @property
    def edges_parent(self):
        """
        Efficient access to the ``parent`` column in the
        :ref:`sec_edge_table_definition` as a numpy array (dtype=np.int32).
        Equivalent to ``ts.tables.edges.parent`` (but avoiding the full copy
        of the table data that accessing ``ts.tables`` currently entails).
        """
        return self._edges_parent

    @property
    def edges_child(self):
        """
        Efficient access to the ``child`` column in the
        :ref:`sec_edge_table_definition` as a numpy array (dtype=np.int32).
        Equivalent to ``ts.tables.edges.child`` (but avoiding the full copy
        of the table data that accessing ``ts.tables`` currently entails).
        """
        return self._edges_child

    @property
    def sites_position(self):
        """
        Efficient access to the ``position`` column in the
        :ref:`sec_site_table_definition` as a numpy array (dtype=np.float64).
        Equivalent to ``ts.tables.sites.position`` (but avoiding the full copy
        of the table data that accessing ``ts.tables`` currently entails).
        """
        return self._sites_position

    @property
    def mutations_site(self):
        """
        Efficient access to the ``site`` column in the
        :ref:`sec_mutation_table_definition` as a numpy array (dtype=np.int32).
        Equivalent to ``ts.tables.mutations.site`` (but avoiding the full copy
        of the table data that accessing ``ts.tables`` currently entails).
        """
        return self._mutations_site

    @property
    def mutations_node(self):
        """
        Efficient access to the ``node`` column in the
        :ref:`sec_mutation_table_definition` as a numpy array (dtype=np.int32).
        Equivalent to ``ts.tables.mutations.node`` (but avoiding the full copy
        of the table data that accessing ``ts.tables`` currently entails).
        """
        return self._mutations_node

    @property
    def mutations_parent(self):
        """
        Efficient access to the ``parent`` column in the
        :ref:`sec_mutation_table_definition` as a numpy array (dtype=np.int32).
        Equivalent to ``ts.tables.mutations.parent`` (but avoiding the full copy
        of the table data that accessing ``ts.tables`` currently entails).
        """
        return self._mutations_parent

    @property
    def mutations_time(self):
        """
        Efficient access to the ``time`` column in the
        :ref:`sec_mutation_table_definition` as a numpy array (dtype=np.float64).
        Equivalent to ``ts.tables.mutations.time`` (but avoiding the full copy
        of the table data that accessing ``ts.tables`` currently entails).
        """
        return self._mutations_time

    @property
    def migrations_left(self):
        """
        Efficient access to the ``left`` column in the
        :ref:`sec_migration_table_definition` as a numpy array (dtype=np.float64).
        Equivalent to ``ts.tables.migrations.left`` (but avoiding the full copy
        of the table data that accessing ``ts.tables`` currently entails).
        """
        return self._migrations_left

    @property
    def migrations_right(self):
        """
        Efficient access to the ``right`` column in the
        :ref:`sec_migration_table_definition` as a numpy array (dtype=np.float64).
        Equivalent to ``ts.tables.migrations.right`` (but avoiding the full copy
        of the table data that accessing ``ts.tables`` currently entails).
        """
        return self._migrations_right

    @property
    def migrations_node(self):
        """
        Efficient access to the ``node`` column in the
        :ref:`sec_migration_table_definition` as a numpy array (dtype=np.int32).
        Equivalent to ``ts.tables.migrations.node`` (but avoiding the full copy
        of the table data that accessing ``ts.tables`` currently entails).
        """
        return self._migrations_node

    @property
    def migrations_source(self):
        """
        Efficient access to the ``source`` column in the
        :ref:`sec_migration_table_definition` as a numpy array (dtype=np.int32).
        Equivalent to ``ts.tables.migrations.source`` (but avoiding the full copy
        of the table data that accessing ``ts.tables`` currently entails).
        """
        return self._migrations_source

    @property
    def migrations_dest(self):
        """
        Efficient access to the ``dest`` column in the
        :ref:`sec_migration_table_definition` as a numpy array (dtype=np.int32).
        Equivalent to ``ts.tables.migrations.dest`` (but avoiding the full copy
        of the table data that accessing ``ts.tables`` currently entails).
        """
        return self._migrations_dest

    @property
    def migrations_time(self):
        """
        Efficient access to the ``time`` column in the
        :ref:`sec_migration_table_definition` as a numpy array (dtype=np.float64).
        Equivalent to ``ts.tables.migrations.time`` (but avoiding the full copy
        of the table data that accessing ``ts.tables`` currently entails).
        """
        return self._migrations_time

    @property
    def indexes_edge_insertion_order(self):
        """
        Efficient access to the ``edge_insertion_order`` column in the
        :ref:`sec_table_indexes` as a numpy array (dtype=np.int32).
        Equivalent to ``ts.tables.indexes.edge_insertion_order`` (but avoiding
        the full copy of the table data that accessing ``ts.tables``
        currently entails).
        """
        return self._indexes_edge_insertion_order

    @property
    def indexes_edge_removal_order(self):
        """
        Efficient access to the ``edge_removal_order`` column in the
        :ref:`sec_table_indexes` as a numpy array (dtype=np.int32).
        Equivalent to ``ts.tables.indexes.edge_removal_order`` (but avoiding
        the full copy of the table data that accessing ``ts.tables``
        currently entails).
        """
        return self._indexes_edge_removal_order

    def individual(self, id_):
        """
        Returns the :ref:`individual <sec_individual_table_definition>`
        in this tree sequence with the specified ID.  As with python lists, negative
        IDs can be used to index backwards from the last individual.

        :rtype: :class:`Individual`
        """
        id_ = self.check_index(id_, self.num_individuals)
        (
            flags,
            location,
            parents,
            metadata,
            nodes,
        ) = self._ll_tree_sequence.get_individual(id_)
        ind = Individual(
            id=id_,
            flags=flags,
            location=location,
            parents=parents,
            metadata=metadata,
            nodes=nodes,
            metadata_decoder=self.table_metadata_schemas.individual.decode_row,
            tree_sequence=self,
        )
        return ind

    def node(self, id_):
        """
        Returns the :ref:`node <sec_node_table_definition>` in this tree sequence
        with the specified ID. As with python lists, negative IDs can be used to
        index backwards from the last node.

        :rtype: :class:`Node`
        """
        id_ = self.check_index(id_, self.num_nodes)
        (
            flags,
            time,
            population,
            individual,
            metadata,
        ) = self._ll_tree_sequence.get_node(id_)
        return Node(
            id=id_,
            flags=flags,
            time=time,
            population=population,
            individual=individual,
            metadata=metadata,
            metadata_decoder=self.table_metadata_schemas.node.decode_row,
        )

    @staticmethod
    def check_index(index, length):
        if not isinstance(index, numbers.Integral):
            raise TypeError(
                f"Index must be of integer type, not '{type(index).__name__}'"
            )
        if index < 0:
            index += length
        if index < 0 or index >= length:
            raise IndexError("Index out of bounds")
        return index

    def edge(self, id_):
        """
        Returns the :ref:`edge <sec_edge_table_definition>` in this tree sequence
        with the specified ID. As with python lists, negative IDs can be used to
        index backwards from the last edge.

        :rtype: :class:`Edge`
        """
        id_ = self.check_index(id_, self.num_edges)
        left, right, parent, child, metadata = self._ll_tree_sequence.get_edge(id_)
        return Edge(
            id=id_,
            left=left,
            right=right,
            parent=parent,
            child=child,
            metadata=metadata,
            metadata_decoder=self.table_metadata_schemas.edge.decode_row,
        )

    def migration(self, id_):
        """
        Returns the :ref:`migration <sec_migration_table_definition>` in this tree
        sequence with the specified ID. As with python lists, negative IDs can be
        used to index backwards from the last migration.

        :rtype: :class:`.Migration`
        """
        id_ = self.check_index(id_, self.num_migrations)
        (
            left,
            right,
            node,
            source,
            dest,
            time,
            metadata,
        ) = self._ll_tree_sequence.get_migration(id_)
        return Migration(
            id=id_,
            left=left,
            right=right,
            node=node,
            source=source,
            dest=dest,
            time=time,
            metadata=metadata,
            metadata_decoder=self.table_metadata_schemas.migration.decode_row,
        )

    def mutation(self, id_):
        """
        Returns the :ref:`mutation <sec_mutation_table_definition>` in this tree sequence
        with the specified ID. As with python lists, negative IDs can be used to
        index backwards from the last mutation.

        :rtype: :class:`Mutation`
        """
        id_ = self.check_index(id_, self.num_mutations)
        (
            site,
            node,
            derived_state,
            parent,
            metadata,
            time,
            edge,
        ) = self._ll_tree_sequence.get_mutation(id_)
        return Mutation(
            id=id_,
            site=site,
            node=node,
            derived_state=derived_state,
            parent=parent,
            metadata=metadata,
            time=time,
            edge=edge,
            metadata_decoder=self.table_metadata_schemas.mutation.decode_row,
        )

    def site(self, id_=None, *, position=None):
        """
        Returns the :ref:`site <sec_site_table_definition>` in this tree sequence
        with either the specified ID or position. As with python lists, negative IDs
        can be used to index backwards from the last site.

        When position is specified instead of site ID, a binary search is done
        on the list of positions of the sites to try to find a site
        with the user-specified position.

        :rtype: :class:`Site`
        """
        if id_ is None and position is None:
            raise TypeError("Site id or position must be provided.")
        elif id_ is not None and position is not None:
            raise TypeError("Only one of site id or position needs to be provided.")
        elif id_ is None:
            position = np.array(position)
            if len(position.shape) > 0:
                raise ValueError("Position must be provided as a scalar value.")
            if position < 0 or position >= self.sequence_length:
                raise ValueError(
                    "Position is beyond the coordinates defined by sequence length."
                )
            site_pos = self.sites_position
            id_ = site_pos.searchsorted(position)
            if id_ >= len(site_pos) or site_pos[id_] != position:
                raise ValueError(f"There is no site at position {position}.")
        else:
            id_ = self.check_index(id_, self.num_sites)
        ll_site = self._ll_tree_sequence.get_site(id_)
        pos, ancestral_state, ll_mutations, _, metadata = ll_site
        mutations = [self.mutation(mut_id) for mut_id in ll_mutations]
        return Site(
            id=id_,
            position=pos,
            ancestral_state=ancestral_state,
            mutations=mutations,
            metadata=metadata,
            metadata_decoder=self.table_metadata_schemas.site.decode_row,
        )

    def population(self, id_):
        """
        Returns the :ref:`population <sec_population_table_definition>`
        in this tree sequence with the specified ID.  As with python lists, negative
        IDs can be used to index backwards from the last population.

        :rtype: :class:`Population`
        """
        id_ = self.check_index(id_, self.num_populations)
        (metadata,) = self._ll_tree_sequence.get_population(id_)
        return Population(
            id=id_,
            metadata=metadata,
            metadata_decoder=self.table_metadata_schemas.population.decode_row,
        )

    def provenance(self, id_):
        """
        Returns the :ref:`provenance <sec_provenance_table_definition>`
        in this tree sequence with the specified ID.  As with python lists,
        negative IDs can be used to index backwards from the last provenance.
        """
        id_ = self.check_index(id_, self.num_provenances)
        timestamp, record = self._ll_tree_sequence.get_provenance(id_)
        return Provenance(id=id_, timestamp=timestamp, record=record)

    def get_samples(self, population_id=None):
        # Deprecated alias for samples()
        return self.samples(population_id)

    def samples(self, population=None, *, population_id=None, time=None):
        """
        Returns an array of the sample node IDs in this tree sequence. If
        `population` is specified, only return sample IDs from that population.
        It is also possible to restrict samples by time using the parameter
        `time`. If `time` is a numeric value, only return sample IDs whose node
        time is approximately equal to the specified time. If `time` is a pair
        of values of the form `(min_time, max_time)`, only return sample IDs
        whose node time `t` is in this interval such that `min_time <= t < max_time`.

        :param int population: The population of interest. If None, do not
            filter samples by population.
        :param int population_id: Deprecated alias for ``population``.
        :param float,tuple time: The time or time interval of interest. If
            None, do not filter samples by time.
        :return: A numpy array of the node IDs for the samples of interest,
            listed in numerical order.
        :rtype: numpy.ndarray (dtype=np.int32)
        """
        if population is not None and population_id is not None:
            raise ValueError(
                "population_id and population are aliases. Cannot specify both"
            )
        if population_id is not None:
            population = population_id
        samples = self._ll_tree_sequence.get_samples()
        keep = np.full(shape=samples.shape, fill_value=True)
        if population is not None:
            sample_population = self.nodes_population[samples]
            keep = np.logical_and(keep, sample_population == population)
        if time is not None:
            # ndmin is set so that scalars are converted into 1d arrays
            time = np.array(time, ndmin=1, dtype=float)
            sample_times = self.nodes_time[samples]
            if time.shape == (1,):
                keep = np.logical_and(keep, np.isclose(sample_times, time))
            elif time.shape == (2,):
                if time[1] <= time[0]:
                    raise ValueError("time_interval max is less than or equal to min.")
                keep = np.logical_and(keep, sample_times >= time[0])
                keep = np.logical_and(keep, sample_times < time[1])
            else:
                raise ValueError(
                    "time must be either a single value or a pair of values "
                    "(min_time, max_time)."
                )
        return samples[keep]

    def as_vcf(self, *args, **kwargs):
        """
        Return the result of :meth:`.write_vcf` as a string.
        Keyword parameters are as defined in :meth:`.write_vcf`.

        :return: A VCF encoding of the variants in this tree sequence as a string.
        :rtype: str
        """
        buff = io.StringIO()
        self.write_vcf(buff, *args, **kwargs)
        return buff.getvalue()

    def write_vcf(
        self,
        output,
        ploidy=None,
        *,
        contig_id="1",
        individuals=None,
        individual_names=None,
        position_transform=None,
        site_mask=None,
        sample_mask=None,
        isolated_as_missing=None,
        allow_position_zero=None,
    ):
        """
        Convert the genetic variation data in this tree sequence to Variant
        Call Format and write to the specified file-like object.

        .. seealso: See the :ref:`sec_export_vcf` section for examples
            and explanations of how we map VCF to the tskit data model.

        Multiploid samples in the output VCF are generated either using
        individual information in the data model (see
        :ref:`sec_individual_table_definition`), or by combining genotypes for
        adjacent sample nodes using the ``ploidy`` argument. See the
        :ref:`sec_export_vcf_constructing_gt` section for more details
        and examples.

        If individuals that are associated with sample nodes are defined in the
        data model (see :ref:`sec_individual_table_definition`), the genotypes
        for each of the individual's samples are combined into a phased
        multiploid values at each site. By default, all individuals associated
        with sample nodes are included in increasing order of individual ID.

        Subsets or permutations of the sample individuals may be specified
        using the ``individuals`` argument. It is an error to specify any
        individuals that are not associated with any nodes, or whose
        nodes are not all samples.

        Mixed-sample individuals (e.g., those associated with one node
        that is a sample and another that is not) in the data model will
        result in an error by default. However, such individuals can be
        excluded using the ``individuals`` argument.

        If there are no individuals in the tree sequence,
        synthetic individuals are created by combining adjacent samples, and
        the number of samples combined is equal to the ``ploidy`` value (1 by
        default). For example, if we have a ``ploidy`` of 2 and 6 sample nodes,
        then we will have 3 diploid samples in the VCF, consisting of the
        combined genotypes for samples [0, 1], [2, 3] and [4, 5]. If we had
        genotypes 011110 at a particular variant, then we would output the
        diploid genotypes 0|1, 1|1 and 1|0 in VCF.

        Each individual in the output is identified by a string; these are the
        VCF "sample" names. By default, these are of the form ``tsk_0``,
        ``tsk_1`` etc, up to the number of individuals, but can be manually
        specified using the ``individual_names`` argument. We do not check
        for duplicates in this array, or perform any checks to ensure that
        the output VCF is well-formed.

        .. note::
            The default individual names (VCF sample IDs) are always of
            the form ``tsk_0``, ``tsk_1``, ..., ``tsk_{N - 1}``, where
            N is the number of individuals we output. These numbers
            are **not** necessarily the individual IDs.

        The REF value in the output VCF is the ancestral allele for a site
        and ALT values are the remaining alleles. It is important to note,
        therefore, that for real data this means that the REF value for a given
        site **may not** be equal to the reference allele. We also do not
        check that the alleles result in a valid VCF---for example, it is possible
        to use the tab character as an allele, leading to a broken VCF.

        The ID value in the output VCF file is the integer ID of the
        corresponding :ref:`site <sec_site_table_definition>` (``site.id``).
        These ID values can be utilized to match the contents of the VCF file
        to the sites in the tree sequence object.

        .. note::
           Older code often uses the ``ploidy=2`` argument, because old
           versions of msprime did not output individual data. Specifying
           individuals in the tree sequence is more robust, and since tree
           sequences now  typically contain individuals (e.g., as produced by
           ``msprime.sim_ancestry( )``), this is not necessary, and the
           ``ploidy`` argument can safely be removed as part of the process
           of updating from the msprime 0.x legacy API.

        :param io.IOBase output: The file-like object to write the VCF output.
        :param int ploidy: The ploidy of the individuals to be written to
            VCF. This sample size must be evenly divisible by ploidy. Cannot be
            used if there is individual data in the tree sequence.
        :param str contig_id: The value of the CHROM column in the output VCF.
        :param list(int) individuals: A list containing the individual IDs to
            corresponding to the VCF samples. Defaults to all individuals
            associated with sample nodes in the tree sequence.
            See the {ref}`sec_export_vcf_constructing_gt` section for more
            details and examples.
        :param list(str) individual_names: A list of string names to identify
            individual columns in the VCF. In VCF nomenclature, these are the
            sample IDs. If specified, this must be a list of strings of
            length equal to the number of individuals to be output. Note that
            we do not check the form of these strings in any way, so that is
            is possible to output malformed VCF (for example, by embedding a
            tab character within on of the names). The default is to output
            ``tsk_j`` for the jth individual.
            See the :ref:`sec_export_vcf_individual_names` for examples
            and more information.
        :param position_transform: A callable that transforms the
            site position values into integer valued coordinates suitable for
            VCF. The function takes a single positional parameter x and must
            return an integer numpy array the same dimension as x. By default,
            this is set to ``numpy.round()`` which will round values to the
            nearest integer. If the string "legacy" is provided here, the
            pre 0.2.0 legacy behaviour of rounding values to the nearest integer
            (starting from 1) and avoiding the output of identical positions
            by incrementing is used.
            See the :ref:`sec_export_vcf_modifying_coordinates` for examples
            and more information.
        :param site_mask: A numpy boolean array (or something convertable to
            a numpy boolean array) with num_sites elements, used to mask out
            sites in the output. If  ``site_mask[j]`` is True, then this
            site (i.e., the line in the VCF file) will be omitted.
            See the :ref:`sec_export_vcf_masking_output` for examples
            and more information.
        :param sample_mask: A numpy boolean array (or something convertable to
            a numpy boolean array) with num_samples elements, or a callable
            that returns such an array, such that if
            ``sample_mask[j]`` is True, then the genotype for sample ``j``
            will be marked as missing using a ".". If ``sample_mask`` is a
            callable, it must take a single argument and return a boolean
            numpy array. This function will be called for each (unmasked) site
            with the corresponding :class:`.Variant` object, allowing
            for dynamic masks to be generated.
            See the :ref:`sec_export_vcf_masking_output` for examples
            and more information.
        :param bool isolated_as_missing: If True, the genotype value assigned to
            missing samples (i.e., isolated samples without mutations) is "."
            If False, missing samples will be assigned the ancestral allele.
            See :meth:`.variants` for more information. Default: True.
        :param bool allow_position_zero: If True allow sites with position zero to be
            output to the VCF, otherwise if one is present an error will be raised.
            The VCF spec does not allow for sites at position 0. However, in practise
            many tools will be fine with this. Default: False.
        """
        if allow_position_zero is None:
            allow_position_zero = False
        writer = vcf.VcfWriter(
            self,
            ploidy=ploidy,
            contig_id=contig_id,
            individuals=individuals,
            individual_names=individual_names,
            position_transform=position_transform,
            site_mask=site_mask,
            sample_mask=sample_mask,
            isolated_as_missing=isolated_as_missing,
            allow_position_zero=allow_position_zero,
        )
        writer.write(output)

    def write_fasta(
        self,
        file_or_path,
        *,
        wrap_width=60,
        reference_sequence=None,
        missing_data_character=None,
    ):
        """
        Writes the :meth:`.alignments` for this tree sequence to file in
        `FASTA <https://en.wikipedia.org/wiki/FASTA_format>`__ format.
        Please see the :meth:`.alignments` method for details on how
        reference sequences are handled.

        Alignments are returned for the
        :ref:`sample nodes<sec_data_model_definitions>` in this tree
        sequence, and a sample with node id ``u`` is given the label
        ``f"n{u}"``, following the same convention as the
        :meth:`.write_nexus` and :meth:`Tree.as_newick` methods.

        The ``wrap_width`` parameter controls the maximum width of lines
        of sequence data in the output. By default this is 60
        characters in accordance with fasta standard outputs. To turn off
        line-wrapping of sequences, set ``wrap_width`` = 0.

        Example usage:

        .. code-block:: python

            ts.write_fasta("output.fa")

        .. warning:: :ref:`Missing data<sec_data_model_missing_data>` is not
            currently supported by this method and it will raise a ValueError
            if called on tree sequences containing isolated samples.
            See https://github.com/tskit-dev/tskit/issues/1896 for more
            information.

        :param file_or_path: The file object or path to write the output.
            Paths can be either strings or :class:`python:pathlib.Path` objects.
        :param int wrap_width: The number of sequence
            characters to include on each line in the fasta file, before wrapping
            to the next line for each sequence, or 0 to turn off line wrapping.
            (Default=60).
        :param str reference_sequence: As for the :meth:`.alignments` method.
        :param str missing_data_character: As for the :meth:`.alignments` method.
        """
        text_formats.write_fasta(
            self,
            file_or_path,
            wrap_width=wrap_width,
            reference_sequence=reference_sequence,
            missing_data_character=missing_data_character,
        )

    def as_fasta(self, **kwargs):
        """
        Return the result of :meth:`.write_fasta` as a string.
        Keyword parameters are as defined in :meth:`.write_fasta`.

        :return: A FASTA encoding of the alignments in this tree sequence as a string.
        :rtype: str
        """
        buff = io.StringIO()
        self.write_fasta(buff, **kwargs)
        return buff.getvalue()

    def write_nexus(
        self,
        file_or_path,
        *,
        precision=None,
        include_trees=None,
        include_alignments=None,
        reference_sequence=None,
        missing_data_character=None,
    ):
        """
        Returns a `nexus encoding <https://en.wikipedia.org/wiki/Nexus_file>`_
        of this tree sequence. By default, tree topologies are included
        in the output, and sequence data alignments are included by default
        if this tree sequence has discrete genome coordinates and one or
        more sites. Inclusion of these sections can be controlled manually
        using the ``include_trees`` and ``include_alignments`` parameters.

        Tree topologies and branch lengths are listed
        sequentially in the TREES block and the spatial location of each tree
        encoded within the tree name labels. Specifically, a tree spanning
        the interval :math:`[x, y)`` is given the name ``f"t{x}^{y}"``
        (See below for a description of the precision at which these spatial
        coordinates are printed out).

        The :ref:`sample nodes<sec_data_model_definitions>` in this tree
        sequence are regarded as taxa, and a sample with node id ``u``
        is given the label ``f"n{u}"``, following the same convention
        as the :meth:`Tree.as_newick` method.

        By default, genome positions are printed out with with sufficient
        precision for them to be recovered exactly in double precision.
        If the tree sequence is defined on a :attr:`.discrete_genome`,
        then positions are written out as integers. Otherwise, 17 digits
        of precision is used. Branch length precision defaults are handled
        in the same way as :meth:`.Tree.as_newick`.

        If the ``precision`` argument is provided, genome positions and
        branch lengths are printed out with this many digits of precision.

        For example, here is the nexus encoding of a simple tree sequence
        with integer times and genome coordinates with three samples
        and two trees::

            #NEXUS
            BEGIN TAXA;
              DIMENSIONS NTAX=3;
              TAXLABELS n0 n1 n2;
            END;
            BEGIN TREES;
              TREE t0^2 = [&R] (n0:3,(n1:2,n2:2):1);
              TREE t2^10 = [&R] (n1:2,(n0:1,n2:1):1);
            END;

        If sequence data :meth:`.alignments` are defined for this tree sequence
        and there is at least one site present, sequence alignment data will also
        be included by default (this can be suppressed by setting
        ``include_alignments=False``). For example, this tree sequence has
        a sequence length of 10, two variable sites and no
        :ref:`reference sequence<sec_data_model_reference_sequence>`::

            #NEXUS
            BEGIN TAXA;
              DIMENSIONS NTAX=3;
              TAXLABELS n0 n1 n2;
            END;
            BEGIN DATA;
              DIMENSIONS NCHAR=10;
              FORMAT DATATYPE=DNA MISSING=?;
              MATRIX
                n0 ??G??????T
                n1 ??A??????C
                n2 ??A??????C
              ;
            END;
            BEGIN TREES;
              TREE t0^10 = [&R] (n0:2,(n1:1,n2:1):1);
            END;

        Please see the :meth:`.alignments` method for details on how
        reference sequences are handled.

        .. note:: Note the default ``missing_data_character`` for this method
            is "?" rather then "N", in keeping with common conventions for
            nexus data. This can be changed using the ``missing_data_character``
            parameter.

        .. warning:: :ref:`Missing data<sec_data_model_missing_data>`
            is not supported for encoding tree topology information
            as our convention of using trees with multiple roots
            is not often supported by newick parsers. Thus, the method
            will raise a ValueError if we try to output trees with
            multiple roots. Additionally, missing data
            is not currently supported for alignment data.
            See https://github.com/tskit-dev/tskit/issues/1896 for more
            information.

        .. seealso: See also the :meth:`.as_nexus` method which will
            return this nexus representation as a string.

        :param int precision: The numerical precision with which branch lengths
            and tree positions are printed.
        :param bool include_trees: True if the tree topology information should
            be included; False otherwise (default=True).
        :param bool include_alignments: True if the sequence data alignment information
            should be included; False otherwise (default=True if sequence alignments
            are well-defined and the tree sequence contains at least one site).
        :param str reference_sequence: As for the :meth:`.alignments` method.
        :param str missing_data_character: As for the :meth:`.alignments` method,
            but defaults to "?".
        :return: A nexus representation of this :class:`TreeSequence`
        :rtype: str
        """
        text_formats.write_nexus(
            self,
            file_or_path,
            precision=precision,
            include_trees=include_trees,
            include_alignments=include_alignments,
            reference_sequence=reference_sequence,
            missing_data_character=missing_data_character,
        )

    def as_nexus(self, **kwargs):
        """
        Return the result of :meth:`.write_nexus` as a string.
        Keyword parameters are as defined in :meth:`.write_nexus`.

        :return: A nexus encoding of the alignments in this tree sequence as a string.
        :rtype: str
        """
        buff = io.StringIO()
        self.write_nexus(buff, **kwargs)
        return buff.getvalue()

    # TODO
    # (1) Move the definition to text_formats.py
    # (2) Rename to as_macs and keep to_macs as a deprecated synonym
    def to_macs(self):
        """
        Return a `macs encoding <https://github.com/gchen98/macs>`_
        of this tree sequence.

        :return: The macs representation of this TreeSequence as a string.
        :rtype: str
        """
        n = self.get_sample_size()
        m = self.get_sequence_length()
        output = [f"COMMAND:\tnot_macs {n} {m}"]
        output.append("SEED:\tASEED")
        for variant in self.variants(copy=False):
            if any(len(allele) > 1 for allele in variant.alleles):
                raise ValueError("macs output only supports single letter alleles")
            bytes_genotypes = np.empty(self.num_samples, dtype=np.uint8)
            lookup = np.array([ord(a[0]) for a in variant.alleles], dtype=np.uint8)
            bytes_genotypes[:] = lookup[variant.genotypes]
            genotypes = bytes_genotypes.tobytes().decode()
            output.append(
                f"SITE:\t{variant.index}\t{variant.position / m}\t0.0\t" f"{genotypes}"
            )
        return "\n".join(output) + "\n"

    def simplify(
        self,
        samples=None,
        *,
        map_nodes=False,
        reduce_to_site_topology=False,
        filter_populations=None,
        filter_individuals=None,
        filter_sites=None,
        filter_nodes=None,
        update_sample_flags=None,
        keep_unary=False,
        keep_unary_in_individuals=None,
        keep_input_roots=False,
        record_provenance=True,
        filter_zero_mutation_sites=None,  # Deprecated alias for filter_sites
    ):
        """
        Returns a simplified tree sequence that retains only the history of
        the nodes given in the list ``samples``. If ``map_nodes`` is true,
        also return a numpy array whose ``u``-th element is the ID of the node
        in the simplified tree sequence that corresponds to node ``u`` in the
        original tree sequence, or :data:`tskit.NULL` (-1) if ``u`` is no longer
        present in the simplified tree sequence.

        .. note::
            If you wish to simplify a set of tables that do not satisfy all
            requirements for building a TreeSequence, then use
            :meth:`TableCollection.simplify`.

        If the ``reduce_to_site_topology`` parameter is True, the returned tree
        sequence will contain only topological information that is necessary to
        represent the trees that contain sites. If there are zero sites in this
        tree sequence, this will result in an output tree sequence with zero edges.
        When the number of sites is greater than zero, every tree in the output
        tree sequence will contain at least one site. For a given site, the
        topology of the tree containing that site will be identical
        (up to node ID remapping) to the topology of the corresponding tree
        in the input tree sequence.

        If ``filter_populations``, ``filter_individuals``, ``filter_sites``, or
        ``filter_nodes`` is True, any of the corresponding objects that are not
        referenced elsewhere are filtered out. As this is the default behaviour,
        it is important to realise IDs for these objects may change through
        simplification. By setting these parameters to False, however, the
        corresponding tables can be preserved without changes.

        If ``filter_nodes`` is False, then the output node table will be
        unchanged except for updating the sample status of nodes and any ID
        remappings caused by filtering individuals and populations (if the
        ``filter_individuals`` and ``filter_populations`` options are enabled).
        Nodes that are in the specified list of ``samples`` will be marked as
        samples in the output, and nodes that are currently marked as samples
        in the node table but not in the specified list of ``samples`` will
        have their :data:`tskit.NODE_IS_SAMPLE` flag cleared. Note also that
        the order of the ``samples`` list is not meaningful when
        ``filter_nodes`` is False. In this case, the returned node mapping is
        always the identity mapping, such that ``a[u] == u`` for all nodes.

        Setting the ``update_sample_flags`` parameter to False disables the
        automatic sample status update of nodes (described above) from
        occuring, making it the responsibility of calling code to keep track of
        the ultimate sample status of nodes. This is an advanced option, mostly
        of use when combined with the ``filter_nodes=False``,
        ``filter_populations=False`` and ``filter_individuals=False`` options,
        which then guarantees that the node table will not be altered by
        simplification.

        :param list[int] samples: A list of node IDs to retain as samples. They
            need not be nodes marked as samples in the original tree sequence, but
            will constitute the entire set of samples in the returned tree sequence.
            If not specified or None, use all nodes marked with the IS_SAMPLE flag.
            The list may be provided as a numpy array (or array-like) object
            (dtype=np.int32).
        :param bool map_nodes: If True, return a tuple containing the resulting
            tree sequence and a numpy array mapping node IDs in the current tree
            sequence to their corresponding node IDs in the returned tree sequence.
            If False (the default), return only the tree sequence object itself.
        :param bool reduce_to_site_topology: Whether to reduce the topology down
            to the trees that are present at sites. (Default: False)
        :param bool filter_populations: If True, remove any populations that are
            not referenced by nodes after simplification; new population IDs are
            allocated sequentially from zero. If False, the population table will
            not be altered in any way. (Default: None, treated as True)
        :param bool filter_individuals: If True, remove any individuals that are
            not referenced by nodes after simplification; new individual IDs are
            allocated sequentially from zero. If False, the individual table will
            not be altered in any way. (Default: None, treated as True)
        :param bool filter_sites: If True, remove any sites that are
            not referenced by mutations after simplification; new site IDs are
            allocated sequentially from zero. If False, the site table will not
            be altered in any way. (Default: None, treated as True)
        :param bool filter_nodes: If True, remove any nodes that are
            not referenced by edges after simplification. If False, the only
            potential change to the node table may be to change the node flags
            (if ``samples`` is specified and different from the existing samples).
            (Default: None, treated as True)
        :param bool update_sample_flags: If True, update node flags to so that
            nodes in the specified list of samples have the NODE_IS_SAMPLE
            flag after simplification, and nodes that are not in this list
            do not. (Default: None, treated as True)
        :param bool keep_unary: If True, preserve unary nodes (i.e., nodes with
            exactly one child) that exist on the path from samples to root.
            (Default: False)
        :param bool keep_unary_in_individuals: If True, preserve unary nodes
            that exist on the path from samples to root, but only if they are
            associated with an individual in the individuals table. Cannot be
            specified at the same time as ``keep_unary``. (Default: ``None``,
            equivalent to False)
        :param bool keep_input_roots: Whether to retain history ancestral to the
            MRCA of the samples. If ``False``, no topology older than the MRCAs of the
            samples will be included. If ``True`` the roots of all trees in the returned
            tree sequence will be the same roots as in the original tree sequence.
            (Default: False)
        :param bool record_provenance: If True, record details of this call to
            simplify in the returned tree sequence's provenance information
            (Default: True).
        :param bool filter_zero_mutation_sites: Deprecated alias for ``filter_sites``.
        :return: The simplified tree sequence, or (if ``map_nodes`` is True)
            a tuple consisting of the simplified tree sequence and a numpy array
            mapping source node IDs to their corresponding IDs in the new tree
            sequence.
        :rtype: tskit.TreeSequence or (tskit.TreeSequence, numpy.ndarray)
        """
        tables = self.dump_tables()
        assert tables.sequence_length == self.sequence_length
        node_map = tables.simplify(
            samples=samples,
            reduce_to_site_topology=reduce_to_site_topology,
            filter_populations=filter_populations,
            filter_individuals=filter_individuals,
            filter_sites=filter_sites,
            filter_nodes=filter_nodes,
            update_sample_flags=update_sample_flags,
            keep_unary=keep_unary,
            keep_unary_in_individuals=keep_unary_in_individuals,
            keep_input_roots=keep_input_roots,
            record_provenance=record_provenance,
            filter_zero_mutation_sites=filter_zero_mutation_sites,
        )
        new_ts = tables.tree_sequence()
        assert new_ts.sequence_length == self.sequence_length
        if map_nodes:
            return new_ts, node_map
        else:
            return new_ts

    def delete_sites(self, site_ids, record_provenance=True):
        """
        Returns a copy of this tree sequence with the specified sites (and their
        associated mutations) entirely removed. The site IDs do not need to be in any
        particular order, and specifying the same ID multiple times does not have any
        effect (i.e., calling ``tree_sequence.delete_sites([0, 1, 1])`` has the same
        effect as calling ``tree_sequence.delete_sites([0, 1])``.

        :param list[int] site_ids: A list of site IDs specifying the sites to remove.
        :param bool record_provenance: If ``True``, add details of this operation to the
            provenance information of the returned tree sequence. (Default: ``True``).
        """
        tables = self.dump_tables()
        tables.delete_sites(site_ids, record_provenance)
        return tables.tree_sequence()

    def delete_intervals(self, intervals, simplify=True, record_provenance=True):
        """
        Returns a copy of this tree sequence for which information in the
        specified list of genomic intervals has been deleted. Edges spanning these
        intervals are truncated or deleted, and sites and mutations falling within
        them are discarded. Note that it is the information in the intervals that
        is deleted, not the intervals themselves, so in particular, all samples
        will be isolated in the deleted intervals.

        Note that node IDs may change as a result of this operation,
        as by default :meth:`.simplify` is called on the returned tree sequence
        to remove redundant nodes. If you wish to map node IDs onto the same
        nodes before and after this method has been called, specify ``simplify=False``.

        See also :meth:`.keep_intervals`, :meth:`.ltrim`, :meth:`.rtrim`, and
        :ref:`missing data<sec_data_model_missing_data>`.

        :param array_like intervals: A list (start, end) pairs describing the
            genomic intervals to delete. Intervals must be non-overlapping and
            in increasing order. The list of intervals must be interpretable as a
            2D numpy array with shape (N, 2), where N is the number of intervals.
        :param bool simplify: If True, return a simplified tree sequence where nodes
            no longer used are discarded. (Default: True).
        :param bool record_provenance: If ``True``, add details of this operation to the
            provenance information of the returned tree sequence. (Default: ``True``).
        :rtype: tskit.TreeSequence
        """
        tables = self.dump_tables()
        tables.delete_intervals(intervals, simplify, record_provenance)
        return tables.tree_sequence()

    def keep_intervals(self, intervals, simplify=True, record_provenance=True):
        """
        Returns a copy of this tree sequence which includes only information in
        the specified list of genomic intervals. Edges are truncated to lie within
        these intervals, and sites and mutations falling outside these intervals
        are discarded.  Note that it is the information outside the intervals that
        is deleted, not the intervals themselves, so in particular, all samples
        will be isolated outside of the retained intervals.

        Note that node IDs may change as a result of this operation,
        as by default :meth:`.simplify` is called on the returned tree sequence
        to remove redundant nodes. If you wish to map node IDs onto the same
        nodes before and after this method has been called, specify ``simplify=False``.

        See also :meth:`.keep_intervals`, :meth:`.ltrim`, :meth:`.rtrim`, and
        :ref:`missing data<sec_data_model_missing_data>`.

        :param array_like intervals: A list (start, end) pairs describing the
            genomic intervals to keep. Intervals must be non-overlapping and
            in increasing order. The list of intervals must be interpretable as a
            2D numpy array with shape (N, 2), where N is the number of intervals.
        :param bool simplify: If True, return a simplified tree sequence where nodes
            no longer used are discarded. (Default: True).
        :param bool record_provenance: If True, add details of this operation to the
            provenance information of the returned tree sequence.
            (Default: True).
        :rtype: tskit.TreeSequence
        """
        tables = self.dump_tables()
        tables.keep_intervals(intervals, simplify, record_provenance)
        return tables.tree_sequence()

    def ltrim(self, record_provenance=True):
        """
        Returns a copy of this tree sequence with a potentially changed coordinate
        system, such that empty regions (i.e., those not covered by any edge) at the
        start of the tree sequence are trimmed away, and the leftmost edge starts at
        position 0. This affects the reported position of sites and
        edges. Additionally, sites and their associated mutations to the left of
        the new zero point are thrown away.

        :param bool record_provenance: If True, add details of this operation to the
            provenance information of the returned tree sequence. (Default: True).
        """
        tables = self.dump_tables()
        tables.ltrim(record_provenance)
        return tables.tree_sequence()

    def rtrim(self, record_provenance=True):
        """
        Returns a copy of this tree sequence with the ``sequence_length`` property reset
        so that the sequence ends at the end of the rightmost edge. Additionally, sites
        and their associated mutations at positions greater than the new
        ``sequence_length`` are thrown away.

        :param bool record_provenance: If True, add details of this operation to the
            provenance information of the returned tree sequence. (Default: True).
        """
        tables = self.dump_tables()
        tables.rtrim(record_provenance)
        return tables.tree_sequence()

    def trim(self, record_provenance=True):
        """
        Returns a copy of this tree sequence with any empty regions (i.e., those not
        covered by any edge) on the right and left trimmed away. This may reset both the
        coordinate system and the ``sequence_length`` property. It is functionally
        equivalent to :meth:`.rtrim` followed by :meth:`.ltrim`. Sites and their
        associated mutations in the empty regions are thrown away.

        :param bool record_provenance: If True, add details of this operation to the
            provenance information of the returned tree sequence. (Default: True).
        """
        tables = self.dump_tables()
        tables.trim(record_provenance)
        return tables.tree_sequence()

    def split_edges(self, time, *, flags=None, population=None, metadata=None):
        """
        Returns a copy of this tree sequence in which we replace any
        edge ``(left, right, parent, child)`` in which
        ``node_time[child] < time < node_time[parent]`` with two edges
        ``(left, right, parent, u)`` and ``(left, right, u, child)``,
        where ``u`` is a newly added node for each intersecting edge.

        If ``metadata``, ``flags``, or ``population`` are specified, newly
        added nodes will be assigned these values. Otherwise, default values
        will be used. The default metadata is an empty dictionary if a metadata
        schema is defined for the node table, and is an empty byte string
        otherwise. The default population for the new node is
        :data:`tskit.NULL`. Newly added have a default ``flags`` value of 0.

        Any metadata associated with a split edge will be copied to the new edge.

        .. warning:: This method currently does not support migrations
            and a error will be raised if the migration table is not
            empty. Future versions may take migrations that intersect with the
            edge into account when determining the default population
            assignments for new nodes.

        Any mutations lying on the edge whose time is >= ``time`` will have
        their node value set to ``u``. Note that the time of the mutation is
        defined as the time of the child node if the mutation's time is
        unknown.

        :param float time: The cutoff time.
        :param int flags: The flags value for newly-inserted nodes. (Default = 0)
        :param int population: The population value for newly inserted nodes.
            Defaults to ``tskit.NULL`` if not specified.
        :param metadata: The metadata for any newly inserted nodes. See
            :meth:`.NodeTable.add_row` for details on how default metadata
            is produced for a given schema (or none).
        :return: A copy of this tree sequence with edges split at the specified time.
        :rtype: tskit.TreeSequence
        """
        population = tskit.NULL if population is None else population
        flags = 0 if flags is None else flags
        schema = self.table_metadata_schemas.node
        if metadata is None:
            metadata = schema.empty_value
        metadata = schema.validate_and_encode_row(metadata)
        ll_ts = self._ll_tree_sequence.split_edges(
            time=time,
            flags=flags,
            population=population,
            metadata=metadata,
        )
        return TreeSequence(ll_ts)

    def decapitate(self, time, *, flags=None, population=None, metadata=None):
        """
        Delete all edge topology and mutational information at least as old
        as the specified time from this tree sequence.

        Removes all edges in which the time of the child is >= the specified
        time ``t``, and breaks edges that intersect with ``t``. For each edge
        intersecting with ``t`` we create a new node with time equal to ``t``,
        and set the parent of the edge to this new node. The node table
        is not altered in any other way. Newly added nodes have values
        for ``flags``, ``population`` and ``metadata`` controlled by parameters
        to this function in the same way as :meth:`.split_edges`.

        .. note::
            Note that each edge is treated independently, so that even if two
            edges that are broken by this operation share the same parent and
            child nodes, there will be two different new parent nodes inserted.

        Any mutation whose time is >= ``t`` will be removed. A mutation's time
        is its associated ``time`` value, or the time of its node if the
        mutation's time was marked as unknown (:data:`UNKNOWN_TIME`).

        Migrations are not supported, and a LibraryError will be raise if
        called on a tree sequence containing migration information.

        .. seealso:: This method is implemented using the :meth:`.split_edges`
            and :meth:`TableCollection.delete_older` functions.

        :param float time: The cutoff time.
        :param int flags: The flags value for newly-inserted nodes. (Default = 0)
        :param int population: The population value for newly inserted nodes.
            Defaults to ``tskit.NULL`` if not specified.
        :param metadata: The metadata for any newly inserted nodes. See
            :meth:`.NodeTable.add_row` for details on how default metadata
            is produced for a given schema (or none).
        :return: A copy of this tree sequence with edges split at the specified time.
        :rtype: tskit.TreeSequence
        """
        split_ts = self.split_edges(
            time, flags=flags, population=population, metadata=metadata
        )
        tables = split_ts.dump_tables()
        del split_ts
        tables.delete_older(time)
        return tables.tree_sequence()

    def extend_haplotypes(self, max_iter=10):
        """
        Returns a new tree sequence in which the span covered by ancestral nodes
        is "extended" to regions of the genome according to the following rule:
        If an ancestral segment corresponding to node `n` has ancestor `p` and
        descendant `c` on some portion of the genome, and on an adjacent segment of
        genome `p` is still an ancestor of `c`, then `n` is inserted into the
        path from `p` to `c`. For instance, if `p` is the parent of `n` and `n`
        is the parent of `c`, then the span of the edges from `p` to `n` and
        `n` to `c` are extended, and the span of the edge from `p` to `c` is
        reduced. Thus, the ancestral haplotype represented by `n` is extended
        to a longer span of the genome. However, any edges whose child node is
        a sample are not modified.

        Since some edges may be removed entirely, this process usually reduces
        the number of edges in the tree sequence.

        The method works by iterating over the genome to look for paths that can
        be extended in this way; the maximum number of such iterations is
        controlled by ``max_iter``.

        The rationale is that we know that `n` carries a portion of the segment
        of ancestral genome inherited by `c` from `p`, and so likely carries
        the *entire* inherited segment (since the implication otherwise would
        be that distinct recombined segments were passed down separately from
        `p` to `c`).

        In the example above, if there was a mutation on the node above `c`
        older than the time of `n` in the span into which `n` was extended,
        then the mutation will now occur above `n`. So, this operation may change
        mutations' nodes (but will not affect genotypes).  This is only
        unambiguous if the mutation's time is known, so the method requires
        known mutation times.  See :meth:`.impute_unknown_mutations_time` if
        mutation times are not known.

        The method will not affect the marginal trees (so, if the original tree
        sequence was simplified, then following up with `simplify` will recover
        the original tree sequence, possibly with edges in a different order).
        It will also not affect the genotype matrix, or any of the tables other
        than the edge table or the node column in the mutation table.

        :param int max_iters: The maximum number of iterations over the tree
            sequence. Defaults to 10.

        :return: A new tree sequence with unary nodes extended.
        :rtype: tskit.TreeSequence
        """
        max_iter = int(max_iter)
        ll_ts = self._ll_tree_sequence.extend_haplotypes(max_iter)
        return TreeSequence(ll_ts)

    def subset(
        self,
        nodes,
        record_provenance=True,
        reorder_populations=True,
        remove_unreferenced=True,
    ):
        """
        Returns a tree sequence containing only information directly
        referencing the provided list of nodes to retain.  The result will
        retain only the nodes whose IDs are listed in ``nodes``, only edges for
        which both parent and child are in ``nodes```, only mutations whose
        node is in ``nodes``, and only individuals that are referred to by one
        of the retained nodes.  Note that this does *not* retain
        the ancestry of these nodes - for that, see :meth:`.simplify`.

        This has the side effect that it may change the order of the nodes,
        individuals, populations, and migrations in the tree sequence: the nodes
        in the new tree sequence will be in the order provided in ``nodes``, and
        both individuals and populations will be ordered by the earliest retained
        node that refers to them. (However, ``reorder_populations`` may be set to
        False to keep the population table unchanged.)

        By default, the method removes all individuals and populations not
        referenced by any nodes, and all sites not referenced by any mutations.
        To retain these unreferenced individuals, populations, and sites, pass
        ``remove_unreferenced=False``. If this is done, the site table will
        remain unchanged, unreferenced individuals will appear at the end of
        the individuals table (and in their original order), and unreferenced
        populations will appear at the end of the population table (unless
        ``reorder_populations=False``).

        .. seealso::

            :meth:`.keep_intervals` for subsetting a given portion of the genome;
            :meth:`.simplify` for retaining the ancestry of a subset of nodes.

        :param list nodes: The list of nodes for which to retain information. This
            may be a numpy array (or array-like) object (dtype=np.int32).
        :param bool record_provenance: Whether to record a provenance entry
            in the provenance table for this operation.
        :param bool reorder_populations: Whether to reorder populations
            (default: True).  If False, the population table will not be altered in
            any way.
        :param bool remove_unreferenced: Whether sites, individuals, and populations
            that are not referred to by any retained entries in the tables should
            be removed (default: True). See the description for details.
        :rtype: tskit.TreeSequence
        """
        tables = self.dump_tables()
        tables.subset(
            nodes,
            record_provenance=record_provenance,
            reorder_populations=reorder_populations,
            remove_unreferenced=remove_unreferenced,
        )
        return tables.tree_sequence()

    def union(
        self,
        other,
        node_mapping,
        check_shared_equality=True,
        add_populations=True,
        record_provenance=True,
    ):
        """
        Returns an expanded tree sequence which contains the node-wise union of
        ``self`` and ``other``, obtained by adding the non-shared portions of
        ``other`` onto ``self``. The "shared" portions are specified using a
        map that specifies which nodes in ``other`` are equivalent to those in
        ``self``: the ``node_mapping`` argument should be an array of length
        equal to the number of nodes in ``other`` and whose entries are the ID
        of the matching node in ``self``, or ``tskit.NULL`` if there is no
        matching node. Those nodes in ``other`` that map to ``tskit.NULL`` will
        be added to ``self``, along with:

        1. Individuals whose nodes are new to ``self``.
        2. Edges whose parent or child are new to ``self``.
        3. Mutations whose nodes are new to ``self``.
        4. Sites which were not present in ``self``, if the site contains a newly
           added mutation.

        By default, populations of newly added nodes are assumed to be new
        populations, and added to the population table as well.

        Note that this operation also sorts the resulting tables, so the
        resulting tree sequence may not be equal to ``self`` even if nothing
        new was added (although it would differ only in ordering of the tables).

        :param TableCollection other: Another table collection.
        :param list node_mapping: An array of node IDs that relate nodes in
            ``other`` to nodes in ``self``.
        :param bool check_shared_equality: If True, the shared portions of the
            tree sequences will be checked for equality. It does so by
            subsetting both ``self`` and ``other`` on the equivalent nodes
            specified in ``node_mapping``, and then checking for equality of
            the subsets.
        :param bool add_populations: If True, nodes new to ``self`` will be
            assigned new population IDs.
        :param bool record_provenance: Whether to record a provenance entry
            in the provenance table for this operation.
        """
        tables = self.dump_tables()
        other_tables = other.dump_tables()
        tables.union(
            other_tables,
            node_mapping,
            check_shared_equality=check_shared_equality,
            add_populations=add_populations,
            record_provenance=record_provenance,
        )
        return tables.tree_sequence()

    def draw_svg(
        self,
        path=None,
        *,
        size=None,
        x_scale=None,
        time_scale=None,
        tree_height_scale=None,
        title=None,
        node_labels=None,
        mutation_labels=None,
        node_titles=None,
        mutation_titles=None,
        root_svg_attributes=None,
        style=None,
        order=None,
        force_root_branch=None,
        symbol_size=None,
        x_axis=None,
        x_label=None,
        x_lim=None,
        x_regions=None,
        y_axis=None,
        y_label=None,
        y_ticks=None,
        y_gridlines=None,
        omit_sites=None,
        canvas_size=None,
        max_num_trees=None,
        **kwargs,
    ):
        """
        Return an SVG representation of a tree sequence. See the
        :ref:`visualization tutorial<tutorials:sec_tskit_viz>` for more details.

        :param str path: The path to the file to write the output. If None, do not write
            to file.
        :param tuple(int, int) size: A tuple of (width, height) specifying a target
            drawing size in abstract user units (usually interpreted as pixels on
            initial display). Components of the drawing will be scaled so that the total
            plot including labels etc. normally fits onto a canvas of this size (see
            ``canvas_size`` below). If ``None``, chose values such that each tree is
            drawn at a size appropriate for a reasonably small set of samples (this will
            nevertheless result in a very wide drawing if there are many trees to
            display). Default: ``None``
        :param str x_scale: Control how the X axis is drawn. If "physical" (the default)
            the axis scales linearly with physical distance along the sequence,
            background shading is used to indicate the position of the trees along the
            X axis, and sites (with associated mutations) are marked at the
            appropriate physical position on axis line. If "treewise", each axis tick
            corresponds to a tree boundary, which are positioned evenly along the axis,
            so that the X axis is of variable scale, no background scaling is required,
            and site positions are not marked on the axis.
        :param str time_scale: Control how height values for nodes are computed.
            If this is equal to ``"time"``, node heights are proportional to their time
            values (this is the default). If this is equal to ``"log_time"``, node
            heights are proportional to their log(time) values. If it is equal to
            ``"rank"``, node heights are spaced equally according to their ranked times.
        :param str tree_height_scale: Deprecated alias for time_scale. (Deprecated in
            0.3.6)
        :param str title: A title string to be included in the SVG output. If ``None``
            (default) no title is shown, which gives more vertical space for the tree.
        :param node_labels: If specified, show custom labels for the nodes
            (specified by ID) that are present in this map; any nodes not present will
            not have a label.
        :type node_labels: dict(int, str)
        :param mutation_labels: If specified, show custom labels for the
            mutations (specified by ID) that are present in the map; any mutations
            not present will not have a label.
        :param dict(int, str) node_titles: If specified, add a ``<title>`` string to
            symbols for each node (specified by ID) present in this map. SVG visualizers
            such as web browsers will commonly display this string on mousing over
            node symbol.
        :param dict(int, str) mutation_titles: If specified, add a ``<title>`` string to
            symbols for each mutation (specified by ID) present in this map. SVG
            visualizers such as web browsers will commonly display this string on
            mousing over the mutation symbol in the tree and (if show) on the x axis.
        :param dict root_svg_attributes: Additional attributes, such as an id, that will
            be embedded in the root ``<svg>`` tag of the generated drawing.
        :param str style: A `css string <https://www.w3.org/TR/CSS21/syndata.htm>`_
            that will be included in the ``<style>`` tag of the generated svg.
        :param str order: The left-to-right ordering of child nodes in each drawn tree.
            This can be either: ``"minlex"``, which minimises the differences
            between adjacent trees (see also the ``"minlex_postorder"`` traversal
            order for the :meth:`.Tree.nodes` method); or ``"tree"`` which draws trees
            in the left-to-right order defined by the
            :ref:`quintuply linked tree structure <sec_data_model_tree_structure>`.
            If not specified or None, this defaults to ``"minlex"``.
        :param bool force_root_branch: If ``True`` plot a branch (edge) above every tree
            root in the tree sequence. If ``None`` (default) then only plot such
            root branches if any root in the tree sequence has a mutation above it.
        :param float symbol_size: Change the default size of the node and mutation
            plotting symbols. If ``None`` (default) use a standard size.
        :param bool x_axis: Should the plot have an X axis line, showing the positions
            of trees along the genome. The scale used is determined by the ``x_scale``
            parameter. If ``None`` (default) plot an X axis.
        :param str x_label: Place a label under the plot. If ``None`` (default) and
            there is an X axis, create and place an appropriate label.
        :param list x_lim: A list of size two giving the genomic positions between which
            trees should be plotted. If the first is ``None``, then plot from the first
            non-empty region of the tree sequence. If the second is ``None``, then plot
            up to the end of the last non-empty region of the tree sequence. The default
            value ``x_lim=None`` is shorthand for the list [``None``, ``None``]. If
            numerical values are given, then regions outside the interval have all
            information discarded: this means that mutations outside the interval will
            not be shown. To force display of the entire tree sequence, including empty
            flanking regions, specify ``x_lim=[0, ts.sequence_length]``.
        :param dict x_regions: A dictionary mapping (left, right) tuples to names. This
            draws a box, labelled with the name, on the X axis between the left and
            right positions, and can be used for annotating genomic regions (e.g.
            genes) on the X axis. If ``None`` (default) do not plot any regions.
        :param bool y_axis: Should the plot have an Y axis line, showing time (or
            ranked node time if ``time_scale="rank"``. If ``None`` (default)
            do not plot a Y axis.
        :param str y_label: Place a label to the left of the plot. If ``None`` (default)
            and there is a Y axis, create and place an appropriate label.
        :param Union[list, dict] y_ticks: A list of Y values at which to plot
            tickmarks, or a dictionary mapping Y values to labels (``[]`` gives no
            tickmarks). If ``None`` (default), plot one tickmark for each unique node
            value. Note that if ``time_scale="rank"``, the Y values refer to the
            zero-based rank of the plotted nodes, rather than the node time itself.
        :param bool y_gridlines: Whether to plot horizontal lines behind the tree
            at each y tickmark.
        :param bool omit_sites: If True, omit sites and mutations from the drawing.
            Default: False
        :param tuple(int, int) canvas_size: The (width, height) of the SVG canvas.
            This will change the SVG width and height without rescaling graphical
            elements, allowing extra room e.g. for unusually long labels. If ``None``
            take the canvas size to be the same as the target drawing size (see
            ``size``, above). Default: None
        :param int max_num_trees: The maximum number of trees to plot. If there are
            more trees than this in the tree sequence, the middle trees will be skipped
            from the plot and a message "XX trees skipped" displayed in their place.
            If ``None``, all the trees will be plotted: this can produce a very wide
            plot if there are many trees in the tree sequence. Default: None

        :return: An SVG representation of a tree sequence.
        :rtype: SVGString

        .. note::
            Technically, x_lim[0] specifies a *minimum* value for the start of the X
            axis, and x_lim[1] specifies a *maximum* value for the end. This is only
            relevant if the tree sequence contains "empty" regions with no edges or
            mutations. In this case if x_lim[0] lies strictly within an empty region
            (i.e., ``empty_tree.interval.left < x_lim[0] < empty_tree.interval.right``)
            then that tree will not be plotted on the left hand side, and the X axis
            will start at ``empty_tree.interval.right``. Similarly, if x_lim[1] lies
            strictly within an empty region then that tree will not be plotted on the
            right hand side, and the X axis will end at ``empty_tree.interval.left``
        """
        draw = drawing.SvgTreeSequence(
            self,
            size,
            x_scale=x_scale,
            time_scale=time_scale,
            tree_height_scale=tree_height_scale,
            title=title,
            node_labels=node_labels,
            mutation_labels=mutation_labels,
            node_titles=node_titles,
            mutation_titles=mutation_titles,
            root_svg_attributes=root_svg_attributes,
            style=style,
            order=order,
            force_root_branch=force_root_branch,
            symbol_size=symbol_size,
            x_axis=x_axis,
            x_label=x_label,
            x_lim=x_lim,
            x_regions=x_regions,
            y_axis=y_axis,
            y_label=y_label,
            y_ticks=y_ticks,
            y_gridlines=y_gridlines,
            omit_sites=omit_sites,
            canvas_size=canvas_size,
            max_num_trees=max_num_trees,
            **kwargs,
        )
        output = draw.drawing.tostring()
        if path is not None:
            # TODO remove the 'pretty' when we are done debugging this.
            draw.drawing.saveas(path, pretty=True)
        return drawing.SVGString(output)

    def draw_text(
        self,
        *,
        node_labels=None,
        use_ascii=False,
        time_label_format=None,
        position_label_format=None,
        order=None,
        **kwargs,
    ):
        """
        Create a text representation of a tree sequence.

        :param dict node_labels: If specified, show custom labels for the nodes
            that are present in the map. Any nodes not specified in the map will
            not have a node label.
        :param bool use_ascii: If ``False`` (default) then use unicode
            `box drawing characters \
<https://en.wikipedia.org/wiki/Box-drawing_character>`_
            to render the tree. If ``True``, use plain ascii characters, which look
            cruder but are less susceptible to misalignment or font substitution.
            Alternatively, if you are having alignment problems with Unicode, you can try
            out the solution documented `here \
<https://github.com/tskit-dev/tskit/issues/189#issuecomment-499114811>`_.
        :param str time_label_format: A python format string specifying the format (e.g.
            number of decimal places or significant figures) used to print the numerical
            time values on the time axis. If ``None``, this defaults to ``"{:.2f}"``.
        :param str position_label_format: A python format string specifying the format
            (e.g. number of decimal places or significant figures) used to print genomic
            positions. If ``None``, this defaults to ``"{:.2f}"``.
        :param str order: The left-to-right ordering of child nodes in the drawn tree.
            This can be either: ``"minlex"``, which minimises the differences
            between adjacent trees (see also the ``"minlex_postorder"`` traversal
            order for the :meth:`.Tree.nodes` method); or ``"tree"`` which draws trees
            in the left-to-right order defined by the
            :ref:`quintuply linked tree structure <sec_data_model_tree_structure>`.
            If not specified or None, this defaults to ``"minlex"``.

        :return: A text representation of a tree sequence.
        :rtype: str
        """
        return str(
            drawing.TextTreeSequence(
                self,
                node_labels=node_labels,
                use_ascii=use_ascii,
                time_label_format=time_label_format,
                position_label_format=position_label_format,
                order=order,
            )
        )

    ############################################
    #
    # Statistics computation
    #
    ############################################

    def general_stat(
        self,
        W,
        f,
        output_dim,
        windows=None,
        polarised=False,
        mode=None,
        span_normalise=True,
        strict=True,
    ):
        """
        Compute a windowed statistic from weights and a summary function.
        See the :ref:`statistics interface <sec_stats_interface>` section for details on
        :ref:`windows <sec_stats_windows>`,
        :ref:`mode <sec_stats_mode>`,
        :ref:`span normalise <sec_stats_span_normalise>`,
        and :ref:`return value <sec_stats_output_format>`.
        On each tree, this
        propagates the weights ``W`` up the tree, so that the "weight" of each
        node is the sum of the weights of all samples at or below the node.
        Then the summary function ``f`` is applied to the weights, giving a
        summary for each node in each tree. How this is then aggregated depends
        on ``mode``:

        "site"
            Adds together the total summary value across all alleles in each window.

        "branch"
            Adds together the summary value for each node, multiplied by the
            length of the branch above the node and the span of the tree.

        "node"
            Returns each node's summary value added across trees and multiplied
            by the span of the tree.

        Both the weights and the summary can be multidimensional: if ``W`` has ``k``
        columns, and ``f`` takes a ``k``-vector and returns an ``m``-vector,
        then the output will be ``m``-dimensional for each node or window (depending
        on "mode").

        .. note::
            The summary function ``f`` should return zero when given both 0 and
            the total weight (i.e., ``f(0) = 0`` and ``f(np.sum(W, axis=0)) = 0``),
            unless ``strict=False``.  This is necessary for the statistic to be
            unaffected by parts of the tree sequence ancestral to none or all
            of the samples, respectively.

        :param numpy.ndarray W: An array of values with one row for each sample and one
            column for each weight.
        :param f: A function that takes a one-dimensional array of length
            equal to the number of columns of ``W`` and returns a one-dimensional
            array.
        :param int output_dim: The length of ``f``'s return value.
        :param list windows: An increasing list of breakpoints between the windows
            to compute the statistic in.
        :param bool polarised: Whether to leave the ancestral state out of computations:
            see :ref:`sec_stats` for more details.
        :param str mode: A string giving the "type" of the statistic to be computed
            (defaults to "site").
        :param bool span_normalise: Whether to divide the result by the span of the
            window (defaults to True).
        :param bool strict: Whether to check that f(0) and f(total weight) are zero.
        :return: A ndarray with shape equal to (num windows, num statistics).
        """
        if mode is None:
            mode = "site"
        if strict:
            total_weights = np.sum(W, axis=0)
            for x in [total_weights, total_weights * 0.0]:
                with np.errstate(invalid="ignore", divide="ignore"):
                    fx = np.array(f(x))
                fx[np.isnan(fx)] = 0.0
                if not np.allclose(fx, np.zeros((output_dim,))):
                    raise ValueError(
                        "Summary function does not return zero for both "
                        "zero weight and total weight."
                    )
        return self.__run_windowed_stat(
            windows,
            self.ll_tree_sequence.general_stat,
            W,
            f,
            output_dim,
            polarised=polarised,
            span_normalise=span_normalise,
            mode=mode,
        )

    def sample_count_stat(
        self,
        sample_sets,
        f,
        output_dim,
        windows=None,
        polarised=False,
        mode=None,
        span_normalise=True,
        strict=True,
    ):
        """
        Compute a windowed statistic from sample counts and a summary function.
        This is a wrapper around :meth:`.general_stat` for the common case in
        which the weights are all either 1 or 0, i.e., functions of the joint
        allele frequency spectrum.
        See the :ref:`statistics interface <sec_stats_interface>` section for details on
        :ref:`sample sets <sec_stats_sample_sets>`,
        :ref:`windows <sec_stats_windows>`,
        :ref:`mode <sec_stats_mode>`,
        :ref:`span normalise <sec_stats_span_normalise>`,
        and :ref:`return value <sec_stats_output_format>`.
        If ``sample_sets`` is a list of ``k`` sets of samples, then
        ``f`` should be a function that takes an argument of length ``k`` and
        returns a one-dimensional array. The ``j``-th element of the argument
        to ``f`` will be the number of samples in ``sample_sets[j]`` that lie
        below the node that ``f`` is being evaluated for. See
        :meth:`.general_stat`  for more details.

        Here is a contrived example: suppose that ``A`` and ``B`` are two sets
        of samples with ``nA`` and ``nB`` elements, respectively. Passing these
        as sample sets will give ``f`` an argument of length two, giving the number
        of samples in ``A`` and ``B`` below the node in question. So, if we define


        .. code-block:: python

            def f(x):
                pA = x[0] / nA
                pB = x[1] / nB
                return np.array([pA * pB])

        then if all sites are biallelic,

        .. code-block:: python

            ts.sample_count_stat([A, B], f, 1, windows="sites", polarised=False, mode="site")

        would compute, for each site, the product of the derived allele
        frequencies in the two sample sets, in a (num sites, 1) array.  If
        instead ``f`` returns ``np.array([pA, pB, pA * pB])``, then the
        output would be a (num sites, 3) array, with the first two columns
        giving the allele frequencies in ``A`` and ``B``, respectively.

        .. note::
            The summary function ``f`` should return zero when given both 0 and
            the sample size (i.e., ``f(0) = 0`` and
            ``f(np.array([len(x) for x in sample_sets])) = 0``).  This is
            necessary for the statistic to be unaffected by parts of the tree
            sequence ancestral to none or all of the samples, respectively.

        :param list sample_sets: A list of lists of Node IDs, specifying the
            groups of nodes to compute the statistic with.
        :param f: A function that takes a one-dimensional array of length
            equal to the number of sample sets and returns a one-dimensional array.
        :param int output_dim: The length of ``f``'s return value.
        :param list windows: An increasing list of breakpoints between the windows
            to compute the statistic in.
        :param bool polarised: Whether to leave the ancestral state out of computations:
            see :ref:`sec_stats` for more details.
        :param str mode: A string giving the "type" of the statistic to be computed
            (defaults to "site").
        :param bool span_normalise: Whether to divide the result by the span of the
            window (defaults to True).
        :param bool strict: Whether to check that f(0) and f(total weight) are zero.
        :return: A ndarray with shape equal to (num windows, num statistics).
        """  # noqa: B950
        # helper function for common case where weights are indicators of sample sets
        for U in sample_sets:
            if len(U) != len(set(U)):
                raise ValueError(
                    "Elements of sample_sets must be lists without repeated elements."
                )
            if len(U) == 0:
                raise ValueError("Elements of sample_sets cannot be empty.")
            for u in U:
                if not self.node(u).is_sample():
                    raise ValueError("Not all elements of sample_sets are samples.")

        W = np.array([[float(u in A) for A in sample_sets] for u in self.samples()])
        return self.general_stat(
            W,
            f,
            output_dim,
            windows=windows,
            polarised=polarised,
            mode=mode,
            span_normalise=span_normalise,
            strict=strict,
        )

    def parse_windows(self, windows):
        # Note: need to make sure windows is a string or we try to compare the
        # target with a numpy array elementwise.
        if windows is None:
            windows = [0.0, self.sequence_length]
        elif isinstance(windows, str):
            if windows == "trees":
                windows = self.breakpoints(as_array=True)
            elif windows == "sites":
                # breakpoints are at 0.0 and at the sites and at the end
                windows = np.concatenate(
                    [
                        [] if self.num_sites > 0 else [0.0],
                        self.sites_position,
                        [self.sequence_length],
                    ]
                )
                windows[0] = 0.0
            else:
                raise ValueError(
                    f"Unrecognized window specification {windows}:",
                    "the only allowed strings are 'sites' or 'trees'",
                )
        return np.array(windows)

    def __run_windowed_stat(self, windows, method, *args, **kwargs):
        strip_dim = windows is None
        windows = self.parse_windows(windows)
        stat = method(*args, **kwargs, windows=windows)
        if strip_dim:
            stat = stat[0]
        return stat

    def __one_way_sample_set_stat(
        self,
        ll_method,
        sample_sets,
        windows=None,
        mode=None,
        span_normalise=True,
        polarised=False,
    ):
        if sample_sets is None:
            sample_sets = self.samples()

        # First try to convert to a 1D numpy array. If it is, then we strip off
        # the corresponding dimension from the output.
        drop_dimension = False
        try:
            sample_sets = np.array(sample_sets, dtype=np.uint64)
        except ValueError:
            pass
        else:
            # If we've successfully converted sample_sets to a 1D numpy array
            # of integers then drop the dimension
            if len(sample_sets.shape) == 1:
                sample_sets = [sample_sets]
                drop_dimension = True

        sample_set_sizes = np.array(
            [len(sample_set) for sample_set in sample_sets], dtype=np.uint32
        )
        if np.any(sample_set_sizes == 0):
            raise ValueError("Sample sets must contain at least one element")

        flattened = util.safe_np_int_cast(np.hstack(sample_sets), np.int32)
        stat = self.__run_windowed_stat(
            windows,
            ll_method,
            sample_set_sizes,
            flattened,
            mode=mode,
            span_normalise=span_normalise,
            polarised=polarised,
        )
        if drop_dimension:
            stat = stat.reshape(stat.shape[:-1])
            if stat.shape == () and windows is None:
                stat = stat[()]
        return stat

    def parse_sites(self, sites):
        row_sites, col_sites = None, None
        if sites is not None:
            if any(not hasattr(a, "__getitem__") or isinstance(a, str) for a in sites):
                raise ValueError("Sites must be a list of lists, tuples, or ndarrays")
            if len(sites) == 2:
                row_sites, col_sites = sites
            elif len(sites) == 1:
                row_sites = col_sites = sites[0]
            else:
                raise ValueError(
                    f"Sites must be a length 1 or 2 list, got a length {len(sites)} list"
                )
        return row_sites, col_sites

    def parse_positions(self, positions):
        row_positions, col_positions = None, None
        if positions is not None:
            if any(
                not hasattr(a, "__getitem__") or isinstance(a, str) for a in positions
            ):
                raise ValueError(
                    "Positions must be a list of lists, tuples, or ndarrays"
                )
            if len(positions) == 2:
                row_positions, col_positions = positions
            elif len(positions) == 1:
                row_positions = col_positions = positions[0]
            else:
                raise ValueError(
                    "Positions must be a length 1 or 2 list, "
                    f"got a length {len(positions)} list"
                )
        return row_positions, col_positions

    def __two_locus_sample_set_stat(
        self,
        ll_method,
        sample_sets,
        sites=None,
        positions=None,
        mode=None,
    ):
        if sample_sets is None:
            sample_sets = self.samples()
        row_sites, col_sites = self.parse_sites(sites)
        row_positions, col_positions = self.parse_positions(positions)

        # First try to convert to a 1D numpy array. If we succeed, then we strip off
        # the corresponding dimension from the output.
        drop_dimension = False
        try:
            sample_sets = np.array(sample_sets, dtype=np.uint64)
        except ValueError:
            pass
        else:
            # If we've successfully converted sample_sets to a 1D numpy array
            # of integers then drop the dimension
            if len(sample_sets.shape) == 1:
                sample_sets = [sample_sets]
                drop_dimension = True

        sample_set_sizes = np.array(
            [len(sample_set) for sample_set in sample_sets], dtype=np.uint32
        )
        if np.any(sample_set_sizes == 0):
            raise ValueError("Sample sets must contain at least one element")

        flattened = util.safe_np_int_cast(np.hstack(sample_sets), np.int32)

        result = ll_method(
            sample_set_sizes,
            flattened,
            row_sites,
            col_sites,
            row_positions,
            col_positions,
            mode,
        )

        if drop_dimension:
            result = result.reshape(result.shape[:2])
        else:
            # Orient the data so that the first dimension is the sample set.
            # With this orientation, we get one LD matrix per sample set.
            result = result.swapaxes(0, 2).swapaxes(1, 2)

        return result

    def __k_way_sample_set_stat(
        self,
        ll_method,
        k,
        sample_sets,
        indexes=None,
        windows=None,
        mode=None,
        span_normalise=True,
        polarised=False,
        centre=True,
    ):
        sample_set_sizes = np.array(
            [len(sample_set) for sample_set in sample_sets], dtype=np.uint32
        )
        if np.any(sample_set_sizes == 0):
            raise ValueError("Sample sets must contain at least one element")
        flattened = util.safe_np_int_cast(np.hstack(sample_sets), np.int32)
        drop_based_on_index = False
        if indexes is None:
            drop_based_on_index = True
            if len(sample_sets) != k:
                raise ValueError(
                    "Must specify indexes if there are not exactly {} sample "
                    "sets.".format(k)
                )
            indexes = np.arange(k, dtype=np.int32)
        drop_dimension = False
        indexes = util.safe_np_int_cast(indexes, np.int32)
        if len(indexes.shape) == 1:
            indexes = indexes.reshape((1, indexes.shape[0]))
            drop_dimension = True
        if len(indexes.shape) != 2 or indexes.shape[1] != k:
            raise ValueError(
                "Indexes must be convertable to a 2D numpy array with {} "
                "columns".format(k)
            )
        stat = self.__run_windowed_stat(
            windows,
            ll_method,
            sample_set_sizes,
            flattened,
            indexes,
            mode=mode,
            span_normalise=span_normalise,
            polarised=polarised,
            centre=centre,
        )
        if drop_dimension:
            stat = stat.reshape(stat.shape[:-1])
            if stat.shape == () and windows is None and drop_based_on_index:
                stat = stat[()]
        return stat

    def __k_way_weighted_stat(
        self,
        ll_method,
        k,
        W,
        indexes=None,
        windows=None,
        mode=None,
        span_normalise=True,
        polarised=False,
        centre=True,
    ):
        W = np.asarray(W)
        if indexes is None:
            if W.shape[1] != k:
                raise ValueError(
                    "Must specify indexes if there are not exactly {} columns "
                    "in W.".format(k)
                )
            indexes = np.arange(k, dtype=np.int32)
        drop_dimension = False
        indexes = util.safe_np_int_cast(indexes, np.int32)
        if len(indexes.shape) == 1:
            indexes = indexes.reshape((1, indexes.shape[0]))
            drop_dimension = True
        if len(indexes.shape) != 2 or indexes.shape[1] != k:
            raise ValueError(
                "Indexes must be convertable to a 2D numpy array with {} "
                "columns".format(k)
            )
        stat = self.__run_windowed_stat(
            windows,
            ll_method,
            W,
            indexes,
            mode=mode,
            span_normalise=span_normalise,
            polarised=polarised,
            centre=centre,
        )
        if drop_dimension:
            stat = stat.reshape(stat.shape[:-1])
        return stat

    def __weighted_vector_stat(
        self,
        ll_method,
        W,
        windows=None,
        mode=None,
        span_normalise=True,
        centre=True,
        nodes=None,
    ):
        W = np.asarray(W)
        if len(W.shape) == 1:
            W = W.reshape(W.shape[0], 1)
        if nodes is None:
            nodes = list(self.samples())
        else:
            if centre:
                raise ValueError("If `nodes` is provided, must have centre=False.")
        try:
            nodes = util.safe_np_int_cast(nodes, np.int32)
        except Exception:
            raise ValueError("Could not interpret `nodes` as a list of node IDs.")
        stat = self.__run_windowed_stat(
            windows,
            ll_method,
            W,
            mode=mode,
            span_normalise=span_normalise,
            centre=centre,
            nodes=nodes,
        )
        return stat

    ############################################
    # Statistics definitions
    ############################################

    def diversity(
        self, sample_sets=None, windows=None, mode="site", span_normalise=True
    ):
        """
        Computes mean genetic diversity (also known as "pi") in each of the
        sets of nodes from ``sample_sets``.  The statistic is also known as
        "sample heterozygosity"; a common citation for the definition is
        `Nei and Li (1979) <https://doi.org/10.1073/pnas.76.10.5269>`_
        (equation 22), so it is sometimes called called "Nei's pi"
        (but also sometimes "Tajima's pi").

        Please see the :ref:`one-way statistics <sec_stats_sample_sets_one_way>`
        section for details on how the ``sample_sets`` argument is interpreted
        and how it interacts with the dimensions of the output array.
        See the :ref:`statistics interface <sec_stats_interface>` section for details on
        :ref:`windows <sec_stats_windows>`,
        :ref:`mode <sec_stats_mode>`,
        :ref:`span normalise <sec_stats_span_normalise>`,
        and :ref:`return value <sec_stats_output_format>`.

        Note that this quantity can also be computed by the
        :meth:`divergence <.TreeSequence.divergence>` method.

        What is computed depends on ``mode``:

        "site"
            Mean pairwise genetic diversity: the average over all n choose 2 pairs of
            sample nodes, of the density of sites at
            which the two carry different alleles, per unit of chromosome length.

        "branch"
            Mean distance in the tree: the average across over all n choose 2 pairs of
            sample nodes and locations in the window, of the mean distance in
            the tree between the two samples (in units of time).

        "node"
            For each node, the proportion of genome on which the node is an ancestor to
            only one of a pair of sample nodes from the sample set, averaged
            over over all n choose 2 pairs of sample nodes.

        :param list sample_sets: A list of lists of Node IDs, specifying the
            groups of nodes for which the statistic is computed. If any of the
            sample sets contain only a single node, the returned diversity will be
            NaN. If ``None`` (default), average over all n choose 2 pairs of distinct
            sample nodes in the tree sequence.
        :param list windows: An increasing list of breakpoints between the windows
            to compute the statistic in.
        :param str mode: A string giving the "type" of the statistic to be computed
            (defaults to "site").
        :param bool span_normalise: Whether to divide the result by the span of the
            window (defaults to True).
        :return: A numpy array whose length is equal to the number of sample sets.
            If there is one sample set and windows=None, a numpy scalar is returned.
        """
        return self.__one_way_sample_set_stat(
            self._ll_tree_sequence.diversity,
            sample_sets,
            windows=windows,
            mode=mode,
            span_normalise=span_normalise,
        )

    def divergence(
        self, sample_sets, indexes=None, windows=None, mode="site", span_normalise=True
    ):
        r"""
        Computes mean genetic divergence between (and within) pairs of
        sets of nodes from ``sample_sets``.
        This is the "average number of differences", usually referred to as "dxy";
        a common citation for this definition is Nei and Li (1979), who called it
        :math:`\pi_{XY}`. Note that the mean pairwise nucleotide diversity of a
        sample set to itself (computed by passing an index of the form (j,j))
        is its :meth:`diversity <.TreeSequence.diversity>` (see the note below).

        Operates on ``k = 2`` sample sets at a time; please see the
        :ref:`multi-way statistics <sec_stats_sample_sets_multi_way>`
        section for details on how the ``sample_sets`` and ``indexes`` arguments are
        interpreted and how they interact with the dimensions of the output array.
        See the :ref:`statistics interface <sec_stats_interface>` section for details on
        :ref:`windows <sec_stats_windows>`,
        :ref:`mode <sec_stats_mode>`,
        :ref:`span normalise <sec_stats_span_normalise>`,
        and :ref:`return value <sec_stats_output_format>`.

        .. note::
            To avoid unexpected results, sample sets should be nonoverlapping,
            since comparisons of individuals to themselves are not removed when computing
            divergence between distinct sample sets. (However, specifying an index
            ``(j, j)`` computes the :meth:`diversity <.TreeSequence.diversity>`
            of ``sample_set[j]``, which removes self comparisons to provide
            an unbiased estimate.)

        What is computed depends on ``mode``:

        "site"
            Mean pairwise genetic divergence: the average across every possible pair of
            chromosomes (one from each sample set), of the density of sites at which
            the two carry different alleles, per unit of chromosome length.

        "branch"
            Mean distance in the tree: the average across every possible pair of
            chromosomes (one from each sample set) and locations in the window, of
            the mean distance in the tree between the two samples (in units of time).

        "node"
            For each node, the proportion of genome on which the node is an ancestor to
            only one of a pair of chromosomes from the sample set, averaged
            over all possible pairs.

        :param list sample_sets: A list of lists of Node IDs, specifying the
            groups of nodes to compute the statistic with.
        :param list indexes: A list of 2-tuples, or None.
        :param list windows: An increasing list of breakpoints between the windows
            to compute the statistic in.
        :param str mode: A string giving the "type" of the statistic to be computed
            (defaults to "site").
        :param bool span_normalise: Whether to divide the result by the span of the
            window (defaults to True).
        :return: A ndarray with shape equal to (num windows, num statistics).
            If there is one pair of sample sets and windows=None, a numpy scalar is
            returned.

        """
        return self.__k_way_sample_set_stat(
            self._ll_tree_sequence.divergence,
            2,
            sample_sets,
            indexes=indexes,
            windows=windows,
            mode=mode,
            span_normalise=span_normalise,
        )

    ############################################
    # Pairwise sample x sample statistics
    ############################################

    def _chunk_sequence_by_tree(self, num_chunks):
        """
        Return list of (left, right) genome interval tuples that contain
        approximately equal numbers of trees as a 2D numpy array. A
        maximum of self.num_trees single-tree intervals can be returned.
        """
        if num_chunks <= 0 or int(num_chunks) != num_chunks:
            raise ValueError("Number of chunks must be an integer > 0")
        num_chunks = min(self.num_trees, num_chunks)
        breakpoints = self.breakpoints(as_array=True)[:-1]
        splits = np.array_split(breakpoints, num_chunks)
        chunks = []
        for j in range(num_chunks - 1):
            chunks.append((splits[j][0], splits[j + 1][0]))
        chunks.append((splits[-1][0], self.sequence_length))
        return chunks

    @staticmethod
    def _chunk_windows(windows, num_chunks):
        """
        Returns a list of (at most) num_chunks windows, which represent splitting
        up the specified list of windows into roughly equal work.

        Currently this is implemented by just splitting up into roughly equal
        numbers of windows in each chunk.
        """
        if num_chunks <= 0 or int(num_chunks) != num_chunks:
            raise ValueError("Number of chunks must be an integer > 0")
        num_chunks = min(len(windows) - 1, num_chunks)
        splits = np.array_split(windows[:-1], num_chunks)
        chunks = []
        for j in range(num_chunks - 1):
            chunk = np.append(splits[j], splits[j + 1][0])
            chunks.append(chunk)
        chunk = np.append(splits[-1], windows[-1])
        chunks.append(chunk)
        return chunks

    def _parallelise_divmat_by_tree(self, num_threads, span_normalise, **kwargs):
        """
        No windows were specified, so we can chunk up the whole genome by
        tree, and do a simple sum of the results. This means that we have to
        handle span_normalise specially, though.
        """

        def worker(interval):
            return self._ll_tree_sequence.divergence_matrix(interval, **kwargs)

        work = self._chunk_sequence_by_tree(num_threads)
        with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as pool:
            results = pool.map(worker, work)
        total = sum(results)
        if span_normalise:
            total /= self.sequence_length
        return total

    def _parallelise_divmat_by_window(self, windows, num_threads, **kwargs):
        """
        We assume we have a number of windows that's >= to the number
        of threads available, and let each thread have a chunk of the
        windows. There will definitely cases where this leads to
        pathological behaviour, so we may need a more sophisticated
        strategy at some point.
        """

        def worker(sub_windows):
            return self._ll_tree_sequence.divergence_matrix(sub_windows, **kwargs)

        work = self._chunk_windows(windows, num_threads)
        with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as executor:
            futures = [executor.submit(worker, sub_windows) for sub_windows in work]
            concurrent.futures.wait(futures)
        return np.vstack([future.result() for future in futures])

    @staticmethod
    def _parse_stat_matrix_sample_sets(ids):
        """
        Returns a flattened list of sets of IDs. If ids is a 1D list,
        interpret as n one-element sets. Otherwise, it must be a sequence
        of ID lists.
        """
        id_dtype = np.int32
        size_dtype = np.uint64
        # Exclude some types that could be specified accidentally, and
        # we may want to reserve for future use.
        if isinstance(ids, (str, bytes, collections.abc.Mapping, numbers.Number)):
            raise TypeError(f"ID specification cannot be a {type(ids)}")
        if len(ids) == 0:
            return np.array([], dtype=id_dtype), np.array([], dtype=size_dtype)
        if isinstance(ids[0], numbers.Number):
            # Interpret as a 1D array
            flat = util.safe_np_int_cast(ids, id_dtype)
            sizes = np.ones(len(flat), dtype=size_dtype)
        else:
            set_lists = []
            sizes = []
            for id_list in ids:
                a = util.safe_np_int_cast(id_list, id_dtype)
                if len(a.shape) != 1:
                    raise ValueError("ID sets must be 1D integer arrays")
                set_lists.append(a)
                sizes.append(len(a))
            flat = np.hstack(set_lists)
            sizes = np.array(sizes, dtype=size_dtype)
        return flat, sizes

    # def divergence_matrix(self, sample_sets, windows=None, mode="site"):
    #     """
    #     Finds the mean divergence  between pairs of samples from each set of
    #     samples and in each window. Returns a numpy array indexed by (window,
    #     sample_set, sample_set).  Diagonal entries are corrected so that the
    #     value gives the mean divergence for *distinct* samples, but it is not
    #     checked whether the sample_sets are disjoint (so offdiagonals are not
    #     corrected).  For this reason, if an element of `sample_sets` has only
    #     one element, the corresponding diagonal will be NaN.

    #     The mean divergence between two samples is defined to be the mean: (as
    #     a TreeStat) length of all edges separating them in the tree, or (as a
    #     SiteStat) density of segregating sites, at a uniformly chosen position
    #     on the genome.

    #     :param list sample_sets: A list of sets of IDs of samples.
    #     :param iterable windows: The breakpoints of the windows (including start
    #         and end, so has one more entry than number of windows).
    #     :return: A list of the upper triangle of mean TMRCA values in row-major
    #         order, including the diagonal.
    #     """
    #     ns = len(sample_sets)
    #     indexes = [(i, j) for i in range(ns) for j in range(i, ns)]
    #     x = self.divergence(sample_sets, indexes, windows, mode=mode)
    #     nw = len(windows) - 1
    #     A = np.ones((nw, ns, ns), dtype=float)
    #     for w in range(nw):
    #         k = 0
    #         for i in range(ns):
    #             for j in range(i, ns):
    #                 A[w, i, j] = A[w, j, i] = x[w][k]
    #                 k += 1
    #     return A
    # NOTE: see older definition of divmat here, which may be useful when documenting
    # this function. See https://github.com/tskit-dev/tskit/issues/2781

    # NOTE for documentation of sample_sets. We *must* use samples currently because
    # the normalisation for non-sample nodes is tricky. Do we normalise by the
    # total span of the ts where the node is 'present' in the tree? We avoid this
    # by insisting on sample nodes.

    # NOTE for documentation of num_threads. Need to explain that the
    # its best to think of as the number of background *worker* threads.
    # default is to run without any worker threads. If you want to run
    # with all the cores on the machine, use num_threads=os.cpu_count().

    def divergence_matrix(
        self,
        sample_sets=None,
        *,
        windows=None,
        num_threads=0,
        mode=None,
        span_normalise=True,
    ):
        windows_specified = windows is not None
        windows = self.parse_windows(windows)
        mode = "site" if mode is None else mode

        if sample_sets is None:
            sample_sets = self.samples()
            flattened_samples = self.samples()
            sample_set_sizes = np.ones(len(sample_sets), dtype=np.uint32)
        else:
            flattened_samples, sample_set_sizes = self._parse_stat_matrix_sample_sets(
                sample_sets
            )

        # FIXME this logic should be merged into __run_windowed_stat if
        # we generalise the num_threads argument to all stats.
        if num_threads <= 0:
            D = self._ll_tree_sequence.divergence_matrix(
                windows,
                sample_sets=flattened_samples,
                sample_set_sizes=sample_set_sizes,
                mode=mode,
                span_normalise=span_normalise,
            )
        else:
            if windows_specified:
                D = self._parallelise_divmat_by_window(
                    windows,
                    num_threads,
                    sample_sets=flattened_samples,
                    sample_set_sizes=sample_set_sizes,
                    mode=mode,
                    span_normalise=span_normalise,
                )
            else:
                D = self._parallelise_divmat_by_tree(
                    num_threads,
                    span_normalise=span_normalise,
                    sample_sets=flattened_samples,
                    sample_set_sizes=sample_set_sizes,
                    mode=mode,
                )

        if not windows_specified:
            # Drop the windows dimension
            D = D[0]
        return D

    def genetic_relatedness(
        self,
        sample_sets,
        indexes=None,
        windows=None,
        mode="site",
        span_normalise=True,
        polarised=True,
        proportion=True,
        centre=True,
    ):
        """
        Computes genetic relatedness between (and within) pairs of
        sets of nodes from ``sample_sets``.
        Operates on ``k = 2`` sample sets at a time; please see the
        :ref:`multi-way statistics <sec_stats_sample_sets_multi_way>`
        section for details on how the ``sample_sets`` and ``indexes`` arguments are
        interpreted and how they interact with the dimensions of the output array.
        See the :ref:`statistics interface <sec_stats_interface>` section for details on
        :ref:`windows <sec_stats_windows>`,
        :ref:`mode <sec_stats_mode>`,
        :ref:`span normalise <sec_stats_span_normalise>`,
        :ref:`polarised <sec_stats_polarisation>`,
        and :ref:`return value <sec_stats_output_format>`.

        What is computed depends on ``mode``:

        "site"
            Frequency of pairwise allelic matches in the window between two
            sample sets relative to the rest of the sample sets. To be precise,
            let `m(u,v)` denote the total number of alleles shared between
            nodes `u` and `v`, and let `m(I,J)` be the average of `m(u,v)` over
            all nodes `u` in sample set `I` and `v` in sample set `J`. Let `S`
            and `T` be independently chosen sample sets. Then, for sample sets
            `I` and `J`, this computes `E[m(I,J) - m(I,S) - m(J,T) + m(S,T)]`
            if centre=True (the default), or `E[m(I,J)]` if centre=False.
            This can also be seen as the covariance of a quantitative trait
            determined by additive contributions from the genomes in each
            sample set. Let each derived allele be associated with an effect
            drawn from a `N(0,1)` distribution, and let the trait value of a
            sample be the sum of its allele effects. Then, this computes
            the covariance between the average trait values of two sample sets.
            For example, to compute covariance between the traits of diploid
            individuals, each sample set would be the pair of genomes of each
            individual, with the trait being the average of the two genomes.
            If ``proportion=True``, this then corresponds to :math:`K_{c0}` in
            `Speed & Balding (2014) <https://www.nature.com/articles/nrg3821>`_,
            multiplied by four (see below).

        "branch"
            Average area of branches in the window ancestral to pairs of samples
            in two sample sets relative to the rest of the sample sets. To be
            precise, let `B(u,v)` denote the total area of all branches
            ancestral to nodes `u` and `v`, and let `B(I,J)` be the average of
            `B(u,v)` over all nodes `u` in sample set `I` and `v` in sample set
            `J`. Let `S` and `T` be two independently chosen sample sets. Then
            for sample sets `I` and `J`, this computes
            `E[B(I,J) - B(I,S) - B(J,T) + B(S,T)]` if centre=True (the default),
            or `E[B(I,J)]` if centre=False.

        "node"
            For each node, the proportion of the window over which pairs of
            samples in two sample sets are descendants, relative to the rest of
            the sample sets. To be precise, for each node `n`, let `N(u,v)`
            denote the proportion of the window over which samples `u` and `v`
            are descendants of `n`, and let and let `N(I,J)` be the average of
            `N(u,v)` over all nodes `u` in sample set `I` and `v` in sample set
            `J`. Let `S` and `T` be two independently chosen sample sets. Then
            for sample sets `I` and `J`, this computes
            `E[N(I,J) - N(I,S) - N(J,T) + N(S,T)]` if centre=True (the default),
            or `E[N(I,J)]` if centre=False.

        *Note:* The default for this statistic - unlike most other statistics - is
        ``polarised=True``. Using the default value ``centre=True``, setting
        ``polarised=False`` will only multiply the result by a factor of two
        for branch-mode, or site-mode if all sites are biallelic. (With
        multiallelic sites the difference is more complicated.) The uncentred
        and unpolarised value is probably not what you are looking for: for
        instance, the unpolarised, uncentred site statistic between two samples
        counts the number of alleles inherited by both *and* the number of
        alleles inherited by neither of the two samples.

        *Note:* Some authors
        (see `Speed & Balding (2014) <https://www.nature.com/articles/nrg3821>`_)
        compute relatedness between `I` and `J` as the total number of all pairwise
        allelic matches between `I` and `J`, rather than the frequency,
        which would define `m(I,J)` as the sum of `m(u,v)` rather than the average
        in the definition of "site" relatedness above. If every sample set is the
        samples of a :math:`k`-ploid individual, this would simply multiply the
        result by :math:`k^2`. However, this definition would make the result not
        useful as a summary statistic of typical relatedness for larger sample
        sets.

        :param list sample_sets: A list of lists of Node IDs, specifying the
            groups of nodes to compute the statistic with.
        :param list indexes: A list of 2-tuples, or None.
        :param list windows: An increasing list of breakpoints between the windows
            to compute the statistic in.
        :param str mode: A string giving the "type" of the statistic to be computed
            (defaults to "site").
        :param bool span_normalise: Whether to divide the result by the span of the
            window (defaults to True). Has no effect if ``proportion`` is True.
        :param bool proportion: Defaults to True.  Whether to divide the result by
            :meth:`.segregating_sites`, called with the same ``windows``,
            ``mode``, and ``span_normalise``. Note that this counts sites
            that are segregating between *any* of the samples of *any* of the
            sample sets (rather than segregating between all of the samples of
            the tree sequence).
        :param bool polarised: Whether to leave the ancestral state out of computations:
            see :ref:`sec_stats` for more details. Defaults to True.
        :param bool centre: Defaults to True. Whether to 'centre' the result, as
            described above (the usual definition is centred).
        :return: A ndarray with shape equal to (num windows, num statistics).
            If there is one pair of sample sets and windows=None, a numpy scalar is
            returned.
        """

        out = self.__k_way_sample_set_stat(
            self._ll_tree_sequence.genetic_relatedness,
            2,
            sample_sets,
            indexes=indexes,
            windows=windows,
            mode=mode,
            span_normalise=span_normalise,
            polarised=polarised,
            centre=centre,
        )
        if proportion:
            # TODO this should be done in C also
            all_samples = np.array(list({u for s in sample_sets for u in s}))
            denominator = self.segregating_sites(
                sample_sets=all_samples,
                windows=windows,
                mode=mode,
                span_normalise=span_normalise,
            )
            # the shapes of out and denominator should be the same except that
            # out may have an extra dimension if indexes is not None
            if indexes is not None and not isinstance(denominator, float):
                oshape = list(out.shape)
                oshape[-1] = 1
                denominator = denominator.reshape(oshape)
            with np.errstate(divide="ignore", invalid="ignore"):
                out /= denominator

        return out

    def genetic_relatedness_matrix(
        self,
        sample_sets=None,
        *,
        windows=None,
        num_threads=0,
        mode=None,
        span_normalise=True,
    ):
        """
        Computes the full matrix of pairwise genetic relatedness values
        between (and within) pairs of sets of nodes from ``sample_sets``.
        *Warning:* this does not compute exactly the same thing as
        :meth:`.genetic_relatedness`: see below for more details.

        If `mode="branch"`, then the value obtained is the same as that from
        :meth:`.genetic_relatedness`, using the options `centre=True` and
        `proportion=False`. The same is true if `mode="site"` and all sites have
        at most one mutation.

        However, if some sites have more than one mutation, the value may differ.
        The reason is that this function (for efficiency) computes relatedness
        using :meth:`.divergence` and the following relationship.
        "Relatedness" measures the number of *shared* alleles (or branches),
        while "divergence" measures the number of *non-shared* alleles (or branches).
        Let :math:`T_i` be the total distance from sample :math:`i` up to the root;
        then if :math:`D_{ij}` is the divergence between :math:`i` and :math:`j`
        and :math:`R_{ij}` is the relatedness between :math:`i` and :math:`j`, then
        :math:`T_i + T_j = D_{ij} + 2 R_{ij}.`
        So, for any samples :math:`I`, :math:`J`, :math:`S`, :math:`T`
        (that may now be random choices),
        :math:`R_{IJ}-R_{IS}-R_{JT}+R_{ST} = (D_{IJ}-D_{IS}-D_{JT}+D_{ST})/ (-2)`.
        Note, however, that this relationship only holds for `mode="site"`
        if we can treat "number of differing alleles" as distances on the tree;
        this is not necessarily the case in the presence of multiple mutations.

        Another caveat in the above relationship between :math:`R` and :math:`D`
        is that :meth:`.divergence` of a sample set to itself does not include
        the "self" comparisons (so as to provide an unbiased estimator of a
        population quantity), while the usual definition of genetic relatedness
        *does* include such comparisons (to provide, for instance, an appropriate
        value for prospective results beginning with only a given set of
        individuals).

        :param list sample_sets: A list of lists of Node IDs, specifying the
            groups of nodes to compute the statistic with.
        :param list windows: An increasing list of breakpoints between the windows
            to compute the statistic in.
        :param str mode: A string giving the "type" of the statistic to be computed
            (defaults to "site").
        :param bool span_normalise: Whether to divide the result by the span of the
            window (defaults to True). Has no effect if ``proportion`` is True.
        :return: A ndarray with shape equal to (num windows, num statistics).
            If there is one pair of sample sets and windows=None, a numpy scalar is
            returned.
        """
        D = self.divergence_matrix(
            sample_sets,
            windows=windows,
            num_threads=num_threads,
            mode=mode,
            span_normalise=span_normalise,
        )

        if sample_sets is None:
            n = np.ones(self.num_samples)
        else:
            n = np.array([len(x) for x in sample_sets])

        def _normalise(B):
            if len(B) == 0:
                return B
            # correct for lack of self comparisons in divergence
            np.fill_diagonal(B, np.diag(B) * (n - 1) / n)
            K = B + np.mean(B)
            y = np.mean(B, axis=0)
            X = y[:, np.newaxis] + y[np.newaxis, :]
            K -= X
            return K / -2

        if windows is None:
            return _normalise(D)
        else:
            for j in range(D.shape[0]):
                D[j] = _normalise(D[j])
        return D

    def genetic_relatedness_weighted(
        self,
        W,
        indexes=None,
        windows=None,
        mode="site",
        span_normalise=True,
        polarised=False,
        centre=True,
    ):
        r"""
        Computes weighted genetic relatedness. If the :math:`k` th pair of indices
        is (i, j) then the :math:`k` th column of output will be
        :math:`\sum_{a,b} W_{ai} W_{bj} C_{ab}`,
        where :math:`W` is the matrix of weights, and :math:`C_{ab}` is the
        :meth:`genetic_relatedness <.TreeSequence.genetic_relatedness>` between sample
        a and sample b, summing over all pairs of samples in the tree sequence.

        *Note:* the genetic relatedness matrix :math:`C` here is as returned by
        :meth:`.genetic_relatedness`, rather than by :meth:`.genetic_relatedness_matrix`
        (see the latter's documentation for the difference).

        :param numpy.ndarray W: An array of values with one row for each sample node and
            one column for each set of weights.
        :param list indexes: A list of 2-tuples, or None (default). Note that if
            indexes = None, then W must have exactly two columns and this is equivalent
            to indexes = [(0,1)].
        :param list windows: An increasing list of breakpoints between the windows
            to compute the statistic in.
        :param str mode: A string giving the "type" of the statistic to be computed
            (defaults to "site").
        :param bool span_normalise: Whether to divide the result by the span of the
            window (defaults to True).
        :param bool polarised: Whether to leave the ancestral state out of computations:
            see :ref:`sec_stats` for more details. Defaults to True.
        :param bool centre: Defaults to True. Whether to 'centre' the result, as
            described above (the usual definition is centred).
        :return: A ndarray with shape equal to (num windows, num statistics).
        """
        if len(W) != self.num_samples:
            raise ValueError(
                "First trait dimension must be equal to number of samples."
            )
        return self.__k_way_weighted_stat(
            self._ll_tree_sequence.genetic_relatedness_weighted,
            2,
            W,
            indexes=indexes,
            windows=windows,
            mode=mode,
            span_normalise=span_normalise,
            polarised=polarised,
            centre=centre,
        )

    def genetic_relatedness_vector(
        self,
        W,
        windows=None,
        mode="site",
        span_normalise=True,
        centre=True,
        nodes=None,
    ):
        r"""
        Computes the product of the genetic relatedness matrix and a vector of weights
        (one per sample). The output is a (num windows) x (num samples) x (num weights)
        array whose :math:`(w,i,j)`-th element is :math:`\sum_{b} W_{bj} C_{ib}`,
        where :math:`W` is the matrix of weights, and :math:`C_{ab}` is the
        :meth:`genetic_relatedness <.TreeSequence.genetic_relatedness>` between sample
        `a` and sample `b` in window `w`, and the sum is over all samples in the tree
        sequence.  Like other statistics, if windows is None, the first dimension in
        the output is dropped.

        The relatedness used here corresponds to `polarised=True`; no unpolarised option
        is available for this method.

        Optionally, you may provide a list of focal nodes that modifies the behavior
        as follows. If `nodes` is a list of `n` node IDs (that do not need to be
        samples), then the output will have dimension (num windows) x n x (num weights),
        and the matrix :math:`C` used in the definition above is the rectangular matrix
        with :math:`C_{ij}` the relatedness between `nodes[i]` and `samples[j]`. This
        can only be used with `centre=False`; if relatedness between uncentred nodes
        and centred samples is desired, then simply subtract column means from `W` first.
        The default is `nodes=None`, which is equivalent to setting `nodes` equal to
        `ts.samples()`.

        :param numpy.ndarray W: An array of values with one row for each sample node and
            one column for each set of weights.
        :param list windows: An increasing list of breakpoints between the windows
            to compute the statistic in.
        :param str mode: A string giving the "type" of the statistic to be computed
            (defaults to "site").
        :param bool span_normalise: Whether to divide the result by the span of the
            window (defaults to True).
        :param bool centre: Whether to use the *centred* relatedness matrix or not:
            see :meth:`genetic_relatedness <.TreeSequence.genetic_relatedness>`.
        :param list nodes: Optionally, a list of focal nodes as described above
            (default: None).
        :return: A ndarray with shape equal to (num windows, num samples, num weights),
            or (num samples, num weights) if windows is None.
        """
        if (not hasattr(W, "__len__")) or (len(W) != self.num_samples):
            raise ValueError(
                "First weight dimension must be equal to number of samples."
            )

        out = self.__weighted_vector_stat(
            self._ll_tree_sequence.genetic_relatedness_vector,
            W,
            windows=windows,
            mode=mode,
            span_normalise=span_normalise,
            centre=centre,
            nodes=nodes,
        )
        return out

    def trait_covariance(self, W, windows=None, mode="site", span_normalise=True):
        """
        Computes the mean squared covariances between each of the columns of ``W``
        (the "phenotypes") and inheritance along the tree sequence.
        See the :ref:`statistics interface <sec_stats_interface>` section for details on
        :ref:`windows <sec_stats_windows>`,
        :ref:`mode <sec_stats_mode>`,
        :ref:`span normalise <sec_stats_span_normalise>`,
        and :ref:`return value <sec_stats_output_format>`.
        Operates on all samples in the tree sequence.

        Concretely, if `g` is a binary vector that indicates inheritance from an allele,
        branch, or node and `w` is a column of W, normalised to have mean zero,
        then the covariance of `g` and `w` is :math:`\\sum_i g_i w_i`, the sum of the
        weights corresponding to entries of `g` that are `1`. Since weights sum to
        zero, this is also equal to the sum of weights whose entries of `g` are 0.
        So, :math:`cov(g,w)^2 = ((\\sum_i g_i w_i)^2 + (\\sum_i (1-g_i) w_i)^2)/2`.

        What is computed depends on ``mode``:

        "site"
            The sum of squared covariances between presence/absence of each allele and
            phenotypes, divided by length of the window (if ``span_normalise=True``).
            This is computed as sum_a (sum(w[a])^2 / 2), where
            w is a column of W with the average subtracted off,
            and w[a] is the sum of all entries of w corresponding to samples
            carrying allele "a", and the first sum is over all alleles.

        "branch"
            The sum of squared covariances between the split induced by each branch and
            phenotypes, multiplied by branch length, averaged across trees in
            the window. This is computed as above: a branch with total weight
            w[b] below b contributes (branch length) * w[b]^2 to the total
            value for a tree. (Since the sum of w is zero, the total weight
            below b and not below b are equal, canceling the factor of 2
            above.)

        "node"
            For each node, the squared covariance between the property of
            inheriting from this node and phenotypes, computed as in "branch".

        :param numpy.ndarray W: An array of values with one row for each sample and one
            column for each "phenotype".
        :param list windows: An increasing list of breakpoints between the windows
            to compute the statistic in.
        :param str mode: A string giving the "type" of the statistic to be computed
            (defaults to "site").
        :param bool span_normalise: Whether to divide the result by the span of the
            window (defaults to True).
        :return: A ndarray with shape equal to (num windows, num statistics).
            If windows=None and W is a single column, a numpy scalar is returned.
        """
        if W.shape[0] != self.num_samples:
            raise ValueError(
                "First trait dimension must be equal to number of samples."
            )
        return self.__run_windowed_stat(
            windows,
            self._ll_tree_sequence.trait_covariance,
            W,
            mode=mode,
            span_normalise=span_normalise,
        )

    def trait_correlation(self, W, windows=None, mode="site", span_normalise=True):
        """
        Computes the mean squared correlations between each of the columns of ``W``
        (the "phenotypes") and inheritance along the tree sequence.
        See the :ref:`statistics interface <sec_stats_interface>` section for details on
        :ref:`windows <sec_stats_windows>`,
        :ref:`mode <sec_stats_mode>`,
        :ref:`span normalise <sec_stats_span_normalise>`,
        and :ref:`return value <sec_stats_output_format>`.
        Operates on all samples in the tree sequence.

        This is computed as squared covariance in
        :meth:`trait_covariance <.TreeSequence.trait_covariance>`,
        but divided by :math:`p (1-p)`, where `p` is the proportion of samples
        inheriting from the allele, branch, or node in question.

        What is computed depends on ``mode``:

        "site"
            The sum of squared correlations between presence/absence of each allele and
            phenotypes, divided by length of the window (if ``span_normalise=True``).
            This is computed as the
            :meth:`trait_covariance <.TreeSequence.trait_covariance>`
            divided by the variance of the relevant column of W
            and by :math:`p * (1 - p)`, where :math:`p` is the allele frequency.

        "branch"
            The sum of squared correlations between the split induced by each branch and
            phenotypes, multiplied by branch length, averaged across trees in
            the window. This is computed as the
            :meth:`trait_covariance <.TreeSequence.trait_covariance>`,
            divided by the variance of the column of w
            and by :math:`p * (1 - p)`, where :math:`p` is the proportion of
            the samples lying below the branch.

        "node"
            For each node, the squared correlation between the property of
            inheriting from this node and phenotypes, computed as in "branch".

        Note that above we divide by the **sample** variance, which for a
        vector x of length n is ``np.var(x) * n / (n-1)``.

        :param numpy.ndarray W: An array of values with one row for each sample and one
            column for each "phenotype". Each column must have positive standard
            deviation.
        :param list windows: An increasing list of breakpoints between the windows
            to compute the statistic in.
        :param str mode: A string giving the "type" of the statistic to be computed
            (defaults to "site").
        :param bool span_normalise: Whether to divide the result by the span of the
            window (defaults to True).
        :return: A ndarray with shape equal to (num windows, num statistics).
            If windows=None and W is a single column, a numpy scalar is returned.
        """
        if W.shape[0] != self.num_samples:
            raise ValueError(
                "First trait dimension must be equal to number of samples."
            )
        sds = np.std(W, axis=0)
        if np.any(sds == 0):
            raise ValueError(
                "Weight columns must have positive variance", "to compute correlation."
            )
        return self.__run_windowed_stat(
            windows,
            self._ll_tree_sequence.trait_correlation,
            W,
            mode=mode,
            span_normalise=span_normalise,
        )

    def trait_regression(self, *args, **kwargs):
        """
        Deprecated synonym for
        :meth:`trait_linear_model <.TreeSequence.trait_linear_model>`.
        """
        warnings.warn(
            "This is deprecated: please use trait_linear_model( ) instead.",
            FutureWarning,
            stacklevel=4,
        )
        return self.trait_linear_model(*args, **kwargs)

    def trait_linear_model(
        self, W, Z=None, windows=None, mode="site", span_normalise=True
    ):
        """
        Finds the relationship between trait and genotype after accounting for
        covariates.  Concretely, for each trait w (i.e., each column of W),
        this does a least-squares fit of the linear model :math:`w \\sim g + Z`,
        where :math:`g` is inheritance in the tree sequence (e.g., genotype)
        and the columns of :math:`Z` are covariates, and returns the squared
        coefficient of :math:`g` in this linear model.
        See the :ref:`statistics interface <sec_stats_interface>` section for details on
        :ref:`windows <sec_stats_windows>`,
        :ref:`mode <sec_stats_mode>`,
        :ref:`span normalise <sec_stats_span_normalise>`,
        and :ref:`return value <sec_stats_output_format>`.
        Operates on all samples in the tree sequence.

        To do this, if `g` is a binary vector that indicates inheritance from an allele,
        branch, or node and `w` is a column of W, there are :math:`k` columns of
        :math:`Z`, and the :math:`k+2`-vector :math:`b` minimises
        :math:`\\sum_i (w_i - b_0 - b_1 g_i - b_2 z_{2,i} - ... b_{k+2} z_{k+2,i})^2`
        then this returns the number :math:`b_1^2`. If :math:`g` lies in the linear span
        of the columns of :math:`Z`, then :math:`b_1` is set to 0. To fit the
        linear model without covariates (only the intercept), set `Z = None`.

        What is computed depends on ``mode``:

        "site"
            Computes the sum of :math:`b_1^2/2` for each allele in the window,
            as above with :math:`g` indicating presence/absence of the allele,
            then divided by the length of the window if ``span_normalise=True``.
            (For biallelic loci, this number is the same for both alleles, and so summing
            over each cancels the factor of two.)

        "branch"
            The squared coefficient :math:`b_1^2`, computed for the split induced by each
            branch (i.e., with :math:`g` indicating inheritance from that branch),
            multiplied by branch length and tree span, summed over all trees
            in the window, and divided by the length of the window if
            ``span_normalise=True``.

        "node"
            For each node, the squared coefficient :math:`b_1^2`, computed for
            the property of inheriting from this node, as in "branch".

        :param numpy.ndarray W: An array of values with one row for each sample and one
            column for each "phenotype".
        :param numpy.ndarray Z: An array of values with one row for each sample and one
            column for each "covariate", or `None`. Columns of `Z` must be linearly
            independent.
        :param list windows: An increasing list of breakpoints between the windows
            to compute the statistic in.
        :param str mode: A string giving the "type" of the statistic to be computed
            (defaults to "site").
        :param bool span_normalise: Whether to divide the result by the span of the
            window (defaults to True).
        :return: A ndarray with shape equal to (num windows, num statistics).
            If windows=None and W is a single column, a numpy scalar is returned.
        """
        if W.shape[0] != self.num_samples:
            raise ValueError(
                "First trait dimension must be equal to number of samples."
            )
        if Z is None:
            Z = np.ones((self.num_samples, 1))
        else:
            tZ = np.column_stack([Z, np.ones((Z.shape[0], 1))])
            if np.linalg.matrix_rank(tZ) == tZ.shape[1]:
                Z = tZ
        if Z.shape[0] != self.num_samples:
            raise ValueError("First dimension of Z must equal the number of samples.")
        if np.linalg.matrix_rank(Z) < Z.shape[1]:
            raise ValueError("Matrix of covariates is computationally singular.")
        # numpy returns a lower-triangular cholesky
        K = np.linalg.cholesky(np.matmul(Z.T, Z)).T
        Z = np.matmul(Z, np.linalg.inv(K))
        return self.__run_windowed_stat(
            windows,
            self._ll_tree_sequence.trait_linear_model,
            W,
            Z,
            mode=mode,
            span_normalise=span_normalise,
        )

    def segregating_sites(
        self, sample_sets=None, windows=None, mode="site", span_normalise=True
    ):
        """
        Computes the density of segregating sites for each of the sets of nodes
        from ``sample_sets``, and related quantities.
        Please see the :ref:`one-way statistics <sec_stats_sample_sets_one_way>`
        section for details on how the ``sample_sets`` argument is interpreted
        and how it interacts with the dimensions of the output array.
        See the :ref:`statistics interface <sec_stats_interface>` section for details on
        :ref:`windows <sec_stats_windows>`, :ref:`mode <sec_stats_mode>`,
        :ref:`span normalise <sec_stats_span_normalise>`,
        and :ref:`return value <sec_stats_output_format>`.

        What is computed depends on ``mode``. For a sample set ``A``, computes:

        "site"
            The sum over sites of the number of alleles found in ``A`` at each site
            minus one, per unit of chromosome length.
            If all sites have at most two alleles, this is the density of sites
            that are polymorphic in ``A``. To get the **number** of segregating minor
            alleles per window, pass ``span_normalise=False``.

        "branch"
            The total length of all branches in the tree subtended by the samples in
            ``A``, averaged across the window.

        "node"
            The proportion of the window on which the node is ancestral to some,
            but not all, of the samples in ``A``.

        :param list sample_sets: A list of lists of Node IDs, specifying the
            groups of nodes to compute the statistic with.
        :param list windows: An increasing list of breakpoints between the windows
            to compute the statistic in.
        :param str mode: A string giving the "type" of the statistic to be computed
            (defaults to "site").
        :param bool span_normalise: Whether to divide the result by the span of the
            window (defaults to True).
        :return: A ndarray with shape equal to (num windows, num statistics).
            If there is one sample set and windows=None, a numpy scalar is returned.
        """
        return self.__one_way_sample_set_stat(
            self._ll_tree_sequence.segregating_sites,
            sample_sets,
            windows=windows,
            mode=mode,
            span_normalise=span_normalise,
        )

    def allele_frequency_spectrum(
        self,
        sample_sets=None,
        windows=None,
        mode="site",
        span_normalise=True,
        polarised=False,
    ):
        """
        Computes the allele frequency spectrum (AFS) in windows across the genome for
        with respect to the specified ``sample_sets``.
        See the :ref:`statistics interface <sec_stats_interface>` section for details on
        :ref:`sample sets <sec_stats_sample_sets>`,
        :ref:`windows <sec_stats_windows>`,
        :ref:`mode <sec_stats_mode>`,
        :ref:`span normalise <sec_stats_span_normalise>`,
        :ref:`polarised <sec_stats_polarisation>`,
        and :ref:`return value <sec_stats_output_format>`.
        and see :ref:`sec_tutorial_afs` for examples of how to use this method.

        Similar to other windowed stats, the first dimension in the returned array
        corresponds to windows, such that ``result[i]`` is the AFS in the ith
        window. The AFS in each window is a k-dimensional numpy array, where k is
        the number of input sample sets, such that ``result[i, j0, j1, ...]`` is the
        value associated with frequency ``j0`` in ``sample_sets[0]``, ``j1`` in
        ``sample_sets[1]``, etc, in window ``i``. From here, we will assume that
        ``afs`` corresponds to the result in a single window, i.e.,
        ``afs = result[i]``.

        If a single sample set is specified, the allele frequency spectrum within
        this set is returned, such that ``afs[j]`` is the value associated with
        frequency ``j``. Thus, singletons are counted in ``afs[1]``, doubletons in
        ``afs[2]``, and so on. The zeroth entry counts alleles or branches not
        seen in the samples but that are polymorphic among the rest of the samples
        of the tree sequence; likewise, the last entry counts alleles fixed in
        the sample set but polymorphic in the entire set of samples. Please see
        the :ref:`sec_tutorial_afs_zeroth_entry` for an illustration.

        .. warning:: Please note that singletons are **not** counted in the initial
            entry in each AFS array (i.e., ``afs[0]``), but in ``afs[1]``.

        If ``sample_sets`` is None (the default), the allele frequency spectrum
        for all samples in the tree sequence is returned.

        If more than one sample set is specified, the **joint** allele frequency
        spectrum within windows is returned. For example, if we set
        ``sample_sets = [S0, S1]``, then afs[1, 2] counts the number of sites that
        are singletons within S0 and doubletons in S1. The dimensions of the
        output array will be ``[num_windows] + [1 + len(S) for S in sample_sets]``.

        If ``polarised`` is False (the default) the AFS will be *folded*, so that
        the counts do not depend on knowing which allele is ancestral. If folded,
        the frequency spectrum for a single sample set ``S`` has ``afs[j] = 0`` for
        all ``j > len(S) / 2``, so that alleles at frequency ``j`` and ``len(S) - j``
        both add to the same entry. If there is more than one sample set, the
        returned array is "lower triangular" in a similar way. For more details,
        especially about handling of multiallelic sites, see :ref:`sec_stats_notes_afs`.

        What is computed depends on ``mode``:

        "site"
            The number of alleles at a given frequency within the specified sample
            sets for each window, per unit of sequence length. To obtain the total
            number of alleles, set ``span_normalise`` to False.

        "branch"
            The total length of branches in the trees subtended by subsets of the
            specified sample sets, per unit of sequence length. To obtain the
            total, set ``span_normalise`` to False.

        "node"
            Not supported for this method (raises a ValueError).

        For example, suppose that `S0` is a list of 5 sample IDs, and `S1` is
        a list of 3 other sample IDs. Then `afs = ts.allele_frequency_spectrum([S0, S1],
        mode="site", span_normalise=False)` will be a 5x3 numpy array, and if
        there are six alleles that are present in only one sample of `S0` but
        two samples of `S1`, then `afs[1,2]` will be equal to 6.  Similarly,
        `branch_afs = ts.allele_frequency_spectrum([S0, S1], mode="branch",
        span_normalise=False)` will also be a 5x3 array, and `branch_afs[1,2]`
        will be the total area (i.e., length times span) of all branches that
        are above exactly one sample of `S0` and two samples of `S1`.

        :param list sample_sets: A list of lists of Node IDs, specifying the
            groups of samples to compute the joint allele frequency
        :param list windows: An increasing list of breakpoints between windows
            along the genome.
        :param str mode: A string giving the "type" of the statistic to be computed
            (defaults to "site").
        :param bool span_normalise: Whether to divide the result by the span of the
            window (defaults to True).
        :return: A (k + 1) dimensional numpy array, where k is the number of sample
            sets specified.
            If there is one sample set and windows=None, a 1 dimensional array is
            returned.
        """
        if sample_sets is None:
            sample_sets = [self.samples()]
        return self.__one_way_sample_set_stat(
            self._ll_tree_sequence.allele_frequency_spectrum,
            sample_sets,
            windows=windows,
            mode=mode,
            span_normalise=span_normalise,
            polarised=polarised,
        )

    def Tajimas_D(self, sample_sets=None, windows=None, mode="site"):
        """
        Computes Tajima's D of sets of nodes from ``sample_sets`` in windows.
        Please see the :ref:`one-way statistics <sec_stats_sample_sets_one_way>`
        section for details on how the ``sample_sets`` argument is interpreted
        and how it interacts with the dimensions of the output array.
        See the :ref:`statistics interface <sec_stats_interface>` section for details on
        :ref:`windows <sec_stats_windows>`, :ref:`mode <sec_stats_mode>`,
        and :ref:`return value <sec_stats_output_format>`.
        Operates on ``k = 1`` sample sets at a
        time. For a sample set ``X`` of ``n`` nodes, if and ``T`` is the mean
        number of pairwise differing sites in ``X`` and ``S`` is the number of
        sites segregating in ``X`` (computed with :meth:`diversity
        <.TreeSequence.diversity>` and :meth:`segregating sites
        <.TreeSequence.segregating_sites>`, respectively, both not span
        normalised), then Tajima's D is

        .. code-block:: python

            D = (T - S / h) / sqrt(a * S + (b / c) * S * (S - 1))
            h = 1 + 1 / 2 + ... + 1 / (n - 1)
            g = 1 + 1 / 2**2 + ... + 1 / (n - 1) ** 2
            a = (n + 1) / (3 * (n - 1) * h) - 1 / h**2
            b = 2 * (n**2 + n + 3) / (9 * n * (n - 1)) - (n + 2) / (h * n) + g / h**2
            c = h**2 + g

        What is computed for diversity and divergence depends on ``mode``;
        see those functions for more details.

        :param list sample_sets: A list of lists of Node IDs, specifying the
            groups of nodes to compute the statistic with.
        :param list indexes: A list of 2-tuples, or None.
        :param list windows: An increasing list of breakpoints between the windows
            to compute the statistic in.
        :param str mode: A string giving the "type" of the statistic to be computed
            (defaults to "site").
        :return: A ndarray with shape equal to (num windows, num statistics).
            If there is one sample set and windows=None, a numpy scalar is returned.
        """

        # TODO this should be done in C as we'll want to support this method there.
        def tjd_func(sample_set_sizes, flattened, **kwargs):
            n = sample_set_sizes
            T = self.ll_tree_sequence.diversity(n, flattened, **kwargs)
            S = self.ll_tree_sequence.segregating_sites(n, flattened, **kwargs)
            h = np.array([np.sum(1 / np.arange(1, nn)) for nn in n])
            g = np.array([np.sum(1 / np.arange(1, nn) ** 2) for nn in n])
            with np.errstate(invalid="ignore", divide="ignore"):
                a = (n + 1) / (3 * (n - 1) * h) - 1 / h**2
                b = (
                    2 * (n**2 + n + 3) / (9 * n * (n - 1))
                    - (n + 2) / (h * n)
                    + g / h**2
                )
                D = (T - S / h) / np.sqrt(a * S + (b / (h**2 + g)) * S * (S - 1))
            return D

        return self.__one_way_sample_set_stat(
            tjd_func, sample_sets, windows=windows, mode=mode, span_normalise=False
        )

    def Fst(
        self, sample_sets, indexes=None, windows=None, mode="site", span_normalise=True
    ):
        """
        Computes "windowed" Fst between pairs of sets of nodes from ``sample_sets``.
        Operates on ``k = 2`` sample sets at a time; please see the
        :ref:`multi-way statistics <sec_stats_sample_sets_multi_way>`
        section for details on how the ``sample_sets`` and ``indexes`` arguments are
        interpreted and how they interact with the dimensions of the output array.
        See the :ref:`statistics interface <sec_stats_interface>` section for details on
        :ref:`windows <sec_stats_windows>`,
        :ref:`mode <sec_stats_mode>`,
        :ref:`span normalise <sec_stats_span_normalise>`,
        and :ref:`return value <sec_stats_output_format>`.

        For sample sets ``X`` and ``Y``, if ``d(X, Y)`` is the
        :meth:`divergence <.TreeSequence.divergence>`
        between ``X`` and ``Y``, and ``d(X)`` is the
        :meth:`diversity <.TreeSequence.diversity>` of ``X``, then what is
        computed is

        .. code-block:: python

            Fst = 1 - 2 * (d(X) + d(Y)) / (d(X) + 2 * d(X, Y) + d(Y))

        What is computed for diversity and divergence depends on ``mode``;
        see those functions for more details.

        :param list sample_sets: A list of lists of Node IDs, specifying the
            groups of nodes to compute the statistic with.
        :param list indexes: A list of 2-tuples.
        :param list windows: An increasing list of breakpoints between the windows
            to compute the statistic in.
        :param str mode: A string giving the "type" of the statistic to be computed
            (defaults to "site").
        :param bool span_normalise: Whether to divide the result by the span of the
            window (defaults to True).
        :return: A ndarray with shape equal to (num windows, num statistics).
            If there is one pair of sample sets and windows=None, a numpy scalar is
            returned.
        """
        # TODO this should really be implemented in C (presumably C programmers will want
        # to compute Fst too), but in the mean time implementing using the low-level
        # calls has two advantages: (a) we automatically change dimensions like the other
        # two-way stats and (b) it's a bit more efficient because we're not messing
        # around with indexes and samples sets twice.

        def fst_func(
            sample_set_sizes,
            flattened,
            indexes,
            windows,
            mode,
            span_normalise,
            polarised,
            centre,
        ):
            # note: this is kinda hacky - polarised and centre are not used here -
            # but this seems necessary to use our __k_way_sample_set_stat framework
            divergences = self._ll_tree_sequence.divergence(
                sample_set_sizes,
                flattened,
                indexes=indexes,
                windows=windows,
                mode=mode,
                span_normalise=span_normalise,
                polarised=polarised,
                centre=centre,
            )
            diversities = self._ll_tree_sequence.diversity(
                sample_set_sizes,
                flattened,
                windows=windows,
                mode=mode,
                span_normalise=span_normalise,
                polarised=polarised,
            )

            orig_shape = divergences.shape
            # "node" statistics produce a 3D array
            if len(divergences.shape) == 2:
                divergences.shape = (divergences.shape[0], 1, divergences.shape[1])
                diversities.shape = (diversities.shape[0], 1, diversities.shape[1])

            fst = np.repeat(1.0, np.prod(divergences.shape))
            fst.shape = divergences.shape
            for i, (u, v) in enumerate(indexes):
                denom = (
                    diversities[:, :, u]
                    + diversities[:, :, v]
                    + 2 * divergences[:, :, i]
                )
                with np.errstate(divide="ignore", invalid="ignore"):
                    fst[:, :, i] -= (
                        2 * (diversities[:, :, u] + diversities[:, :, v]) / denom
                    )
            fst.shape = orig_shape
            return fst

        return self.__k_way_sample_set_stat(
            fst_func,
            2,
            sample_sets,
            indexes=indexes,
            windows=windows,
            mode=mode,
            span_normalise=span_normalise,
        )

    def Y3(
        self, sample_sets, indexes=None, windows=None, mode="site", span_normalise=True
    ):
        """
        Computes the 'Y' statistic between triples of sets of nodes from ``sample_sets``.
        Operates on ``k = 3`` sample sets at a time; please see the
        :ref:`multi-way statistics <sec_stats_sample_sets_multi_way>`
        section for details on how the ``sample_sets`` and ``indexes`` arguments are
        interpreted and how they interact with the dimensions of the output array.
        See the :ref:`statistics interface <sec_stats_interface>` section for details on
        :ref:`windows <sec_stats_windows>`,
        :ref:`mode <sec_stats_mode>`,
        :ref:`span normalise <sec_stats_span_normalise>`,
        and :ref:`return value <sec_stats_output_format>`.

        What is computed depends on ``mode``. Each is an average across every
        combination of trios of samples ``(a, b, c)``, one chosen from each sample set:

        "site"
            The average density of sites at which ``a`` differs from ``b`` and
            ``c``, per unit of chromosome length.

        "branch"
            The average length of all branches that separate ``a`` from ``b``
            and ``c`` (in units of time).

        "node"
            For each node, the average proportion of the window on which ``a``
            inherits from that node but ``b`` and ``c`` do not, or vice-versa.

        :param list sample_sets: A list of lists of Node IDs, specifying the
            groups of nodes to compute the statistic with.
        :param list indexes: A list of 3-tuples, or None.
        :param list windows: An increasing list of breakpoints between the windows
            to compute the statistic in.
        :param str mode: A string giving the "type" of the statistic to be computed
            (defaults to "site").
        :param bool span_normalise: Whether to divide the result by the span of the
            window (defaults to True).
        :return: A ndarray with shape equal to (num windows, num statistics).
            If there is one triple of sample sets and windows=None, a numpy scalar is
            returned.
        """
        return self.__k_way_sample_set_stat(
            self._ll_tree_sequence.Y3,
            3,
            sample_sets,
            indexes=indexes,
            windows=windows,
            mode=mode,
            span_normalise=span_normalise,
        )

    def Y2(
        self, sample_sets, indexes=None, windows=None, mode="site", span_normalise=True
    ):
        """
        Computes the 'Y2' statistic between pairs of sets of nodes from ``sample_sets``.
        Operates on ``k = 2`` sample sets at a time; please see the
        :ref:`multi-way statistics <sec_stats_sample_sets_multi_way>`
        section for details on how the ``sample_sets`` and ``indexes`` arguments are
        interpreted and how they interact with the dimensions of the output array.
        See the :ref:`statistics interface <sec_stats_interface>` section for details on
        :ref:`windows <sec_stats_windows>`,
        :ref:`mode <sec_stats_mode>`,
        :ref:`span normalise <sec_stats_span_normalise>`,
        and :ref:`return value <sec_stats_output_format>`.

        What is computed depends on ``mode``. Each is computed exactly as
        ``Y3``, except that the average is across every possible trio of samples
        ``(a, b1, b2)``, where ``a`` is chosen from the first sample set, and
        ``b1, b2`` are chosen (without replacement) from the second sample set.
        See :meth:`Y3 <.TreeSequence.Y3>` for more details.

        :param list sample_sets: A list of lists of Node IDs, specifying the
            groups of nodes to compute the statistic with.
        :param list indexes: A list of 2-tuples, or None.
        :param list windows: An increasing list of breakpoints between the windows
            to compute the statistic in.
        :param str mode: A string giving the "type" of the statistic to be computed
            (defaults to "site").
        :param bool span_normalise: Whether to divide the result by the span of the
            window (defaults to True).
        :return: A ndarray with shape equal to (num windows, num statistics).
            If there is one pair of sample sets and windows=None, a numpy scalar is
            returned.
        """
        return self.__k_way_sample_set_stat(
            self._ll_tree_sequence.Y2,
            2,
            sample_sets,
            indexes=indexes,
            windows=windows,
            mode=mode,
            span_normalise=span_normalise,
        )

    def Y1(self, sample_sets, windows=None, mode="site", span_normalise=True):
        """
        Computes the 'Y1' statistic within each of the sets of nodes given by
        ``sample_sets``.
        Please see the :ref:`one-way statistics <sec_stats_sample_sets_one_way>`
        section for details on how the ``sample_sets`` argument is interpreted
        and how it interacts with the dimensions of the output array.
        See the :ref:`statistics interface <sec_stats_interface>` section for details on
        :ref:`windows <sec_stats_windows>`, :ref:`mode <sec_stats_mode>`,
        :ref:`span normalise <sec_stats_span_normalise>`,
        and :ref:`return value <sec_stats_output_format>`.
        Operates on ``k = 1`` sample set at a time.

        What is computed depends on ``mode``. Each is computed exactly as
        ``Y3``, except that the average is across every possible trio of samples
        samples ``(a1, a2, a3)`` all chosen without replacement from the same
        sample set. See :meth:`Y3 <.TreeSequence.Y3>` for more details.

        :param list sample_sets: A list of lists of Node IDs, specifying the
            groups of nodes to compute the statistic with.
        :param list windows: An increasing list of breakpoints between the windows
            to compute the statistic in.
        :param str mode: A string giving the "type" of the statistic to be computed
            (defaults to "site").
        :param bool span_normalise: Whether to divide the result by the span of the
            window (defaults to True).
        :return: A ndarray with shape equal to (num windows, num statistics).
            If there is one sample set and windows=None, a numpy scalar is returned.
        """
        return self.__one_way_sample_set_stat(
            self._ll_tree_sequence.Y1,
            sample_sets,
            windows=windows,
            mode=mode,
            span_normalise=span_normalise,
        )

    def f4(
        self, sample_sets, indexes=None, windows=None, mode="site", span_normalise=True
    ):
        """
        Computes Patterson's f4 statistic between four groups of nodes from
        ``sample_sets``.
        Operates on ``k = 4`` sample sets at a time; please see the
        :ref:`multi-way statistics <sec_stats_sample_sets_multi_way>`
        section for details on how the ``sample_sets`` and ``indexes`` arguments are
        interpreted and how they interact with the dimensions of the output array.
        See the :ref:`statistics interface <sec_stats_interface>` section for details on
        :ref:`windows <sec_stats_windows>`,
        :ref:`mode <sec_stats_mode>`,
        :ref:`span normalise <sec_stats_span_normalise>`,
        and :ref:`return value <sec_stats_output_format>`.

        What is computed depends on ``mode``. Each is an average across every possible
        combination of four samples ``(a, b; c, d)``, one chosen from each sample set:

        "site"
            The average density of sites at which ``a`` and ``c`` agree but
            differs from ``b`` and ``d``, minus the average density of sites at
            which ``a`` and ``d`` agree but differs from ``b`` and ``c``, per
            unit of chromosome length.

        "branch"
            The average length of all branches that separate ``a`` and ``c``
            from ``b`` and ``d``, minus the average length of all branches that
            separate ``a`` and ``d`` from ``b`` and ``c`` (in units of time).

        "node"
            For each node, the average proportion of the window on which ``a`` and ``c``
            inherit from that node but ``b`` and ``d`` do not, or vice-versa,
            minus the average proportion of the window on which ``a`` and ``d``
            inherit from that node but ``b`` and ``c`` do not, or vice-versa.

        :param list sample_sets: A list of lists of Node IDs, specifying the
            groups of nodes to compute the statistic with.
        :param list indexes: A list of 4-tuples, or None.
        :param list windows: An increasing list of breakpoints between the windows
            to compute the statistic in.
        :param str mode: A string giving the "type" of the statistic to be computed
            (defaults to "site").
        :param bool span_normalise: Whether to divide the result by the span of the
            window (defaults to True).
        :return: A ndarray with shape equal to (num windows, num statistics).
            If there are four sample sets and windows=None, a numpy scalar is returned.
        """
        return self.__k_way_sample_set_stat(
            self._ll_tree_sequence.f4,
            4,
            sample_sets,
            indexes=indexes,
            windows=windows,
            mode=mode,
            span_normalise=span_normalise,
        )

    def f3(
        self, sample_sets, indexes=None, windows=None, mode="site", span_normalise=True
    ):
        r"""
        Computes Patterson's f3 statistic between three groups of nodes from
        ``sample_sets``.
        Note that the order of the arguments of f3 differs across the literature:
        here, ``f3([A, B, C])`` for sample sets ``A``, ``B``, and ``C``
        will estimate
        :math:`f_3(A; B, C) = \mathbb{E}[(p_A - p_B) (p_A - p_C)]`,
        where :math:`p_A` is the allele frequency in ``A``.
        When used as a test for admixture, the putatively admixed population
        is usually placed as population ``A`` (see
        `Peter (2016) <https://doi.org/10.1534/genetics.115.183913>`_
        for more discussion).

        Operates on ``k = 3`` sample sets at a time; please see the
        :ref:`multi-way statistics <sec_stats_sample_sets_multi_way>`
        section for details on how the ``sample_sets`` and ``indexes`` arguments are
        interpreted and how they interact with the dimensions of the output array.
        See the :ref:`statistics interface <sec_stats_interface>` section for details on
        :ref:`windows <sec_stats_windows>`,
        :ref:`mode <sec_stats_mode>`,
        :ref:`span normalise <sec_stats_span_normalise>`,
        and :ref:`return value <sec_stats_output_format>`.

        What is computed depends on ``mode``. Each works exactly as
        :meth:`f4 <.TreeSequence.f4>`, except the average is across every possible
        combination of four samples ``(a1, b; a2, c)`` where `a1` and `a2` have both
        been chosen (without replacement) from the first sample set. See
        :meth:`f4 <.TreeSequence.f4>` for more details.

        :param list sample_sets: A list of lists of Node IDs, specifying the
            groups of nodes to compute the statistic with.
        :param list indexes: A list of 3-tuples, or None.
        :param list windows: An increasing list of breakpoints between the windows
            to compute the statistic in.
        :param str mode: A string giving the "type" of the statistic to be computed
            (defaults to "site").
        :param bool span_normalise: Whether to divide the result by the span of the
            window (defaults to True).
        :return: A ndarray with shape equal to (num windows, num statistics).
            If there are three sample sets and windows=None, a numpy scalar is returned.
        """
        return self.__k_way_sample_set_stat(
            self._ll_tree_sequence.f3,
            3,
            sample_sets,
            indexes=indexes,
            windows=windows,
            mode=mode,
            span_normalise=span_normalise,
        )

    def f2(
        self, sample_sets, indexes=None, windows=None, mode="site", span_normalise=True
    ):
        """
        Computes Patterson's f2 statistic between two groups of nodes from
        ``sample_sets``.
        Operates on ``k = 2`` sample sets at a time; please see the
        :ref:`multi-way statistics <sec_stats_sample_sets_multi_way>`
        section for details on how the ``sample_sets`` and ``indexes`` arguments are
        interpreted and how they interact with the dimensions of the output array.
        See the :ref:`statistics interface <sec_stats_interface>` section for details on
        :ref:`windows <sec_stats_windows>`,
        :ref:`mode <sec_stats_mode>`,
        :ref:`span normalise <sec_stats_span_normalise>`,
        and :ref:`return value <sec_stats_output_format>`.

        What is computed depends on ``mode``. Each works exactly as
        :meth:`f4 <.TreeSequence.f4>`, except the average is across every possible
        combination of four samples ``(a1, b1; a2, b2)`` where `a1` and `a2` have
        both been chosen (without replacement) from the first sample set, and ``b1``
        and ``b2`` have both been chosen (without replacement) from the second
        sample set. See :meth:`f4 <.TreeSequence.f4>` for more details.

        :param list sample_sets: A list of lists of Node IDs, specifying the
            groups of nodes to compute the statistic with.
        :param list indexes: A list of 2-tuples, or None.
        :param list windows: An increasing list of breakpoints between the windows
            to compute the statistic in.
        :param str mode: A string giving the "type" of the statistic to be computed
            (defaults to "site").
        :param bool span_normalise: Whether to divide the result by the span of the
            window (defaults to True).
        :return: A ndarray with shape equal to (num windows, num statistics).
            If there is one pair of sample sets and windows=None, a numpy scalar is
            returned.
        """
        return self.__k_way_sample_set_stat(
            self._ll_tree_sequence.f2,
            2,
            sample_sets,
            indexes=indexes,
            windows=windows,
            mode=mode,
            span_normalise=span_normalise,
        )

    def mean_descendants(self, sample_sets):
        """
        Computes for every node the mean number of samples in each of the
        `sample_sets` that descend from that node, averaged over the
        portions of the genome for which the node is ancestral to *any* sample.
        The output is an array, `C[node, j]`, which reports the total span of
        all genomes in `sample_sets[j]` that inherit from `node`, divided by
        the total span of the genome on which `node` is an ancestor to any
        sample in the tree sequence.

        .. warning:: The interface for this method is preliminary and may be subject to
            backwards incompatible changes in the near future. The long-term stable
            API for this method will be consistent with other :ref:`sec_stats`.
            In particular, the normalization by proportion of the genome that `node`
            is an ancestor to anyone may not be the default behaviour in the future.

        :param list sample_sets: A list of lists of node IDs.
        :return: An array with dimensions (number of nodes in the tree sequence,
            number of reference sets)
        """
        return self._ll_tree_sequence.mean_descendants(sample_sets)

    def genealogical_nearest_neighbours(self, focal, sample_sets, num_threads=0):
        """
        Return the genealogical nearest neighbours (GNN) proportions for the given
        focal nodes, with reference to two or more sets of interest, averaged over all
        trees in the tree sequence.

        The GNN proportions for a focal node in a single tree are given by first finding
        the most recent common ancestral node :math:`a` between the focal node and any
        other node present in the reference sets. The GNN proportion for a specific
        reference set, :math:`S` is the number of nodes in :math:`S` that descend from
        :math:`a`, as a proportion of the total number of descendant nodes in any of the
        reference sets.

        For example, consider a case with 2 sample sets, :math:`S_1` and :math:`S_2`.
        For a given tree, :math:`a` is the node that includes at least one descendant in
        :math:`S_1` or :math:`S_2` (not including the focal node). If the descendants of
        :math:`a` include some nodes in :math:`S_1` but no nodes in :math:`S_2`, then the
        GNN proportions for that tree will be 100% :math:`S_1` and 0% :math:`S_2`, or
        :math:`[1.0, 0.0]`.

        For a given focal node, the GNN proportions returned by this function are an
        average of the GNNs for each tree, weighted by the genomic distance spanned by
        that tree.

        For an precise mathematical definition of GNN, see https://doi.org/10.1101/458067

        .. note:: The reference sets need not include all the samples, hence the most
            recent common ancestral node of the reference sets, :math:`a`, need not be
            the immediate ancestor of the focal node. If the reference sets only comprise
            sequences from relatively distant individuals, the GNN statistic may end up
            as a measure of comparatively distant ancestry, even for tree sequences that
            contain many closely related individuals.

        .. warning:: The interface for this method is preliminary and may be subject to
            backwards incompatible changes in the near future. The long-term stable
            API for this method will be consistent with other :ref:`sec_stats`.

        :param list focal: A list of :math:`n` nodes whose GNNs should be calculated.
        :param list sample_sets: A list of :math:`m` lists of node IDs.
        :return: An :math:`n`  by :math:`m` array of focal nodes by GNN proportions.
            Every focal node corresponds to a row. The numbers in each
            row corresponding to the GNN proportion for each of the passed-in reference
            sets. Rows therefore sum to one.
        :rtype: numpy.ndarray
        """
        # TODO add windows=None option: https://github.com/tskit-dev/tskit/issues/193
        if num_threads <= 0:
            return self._ll_tree_sequence.genealogical_nearest_neighbours(
                focal, sample_sets
            )
        else:
            worker = functools.partial(
                self._ll_tree_sequence.genealogical_nearest_neighbours,
                reference_sets=sample_sets,
            )
            focal = util.safe_np_int_cast(focal, np.int32)
            splits = np.array_split(focal, num_threads)
            with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as pool:
                arrays = pool.map(worker, splits)
            return np.vstack(list(arrays))

    def kc_distance(self, other, lambda_=0.0):
        """
        Returns the average :meth:`Tree.kc_distance` between pairs of trees along
        the sequence whose intervals overlap. The average is weighted by the
        fraction of the sequence on which each pair of trees overlap.

        :param TreeSequence other: The other tree sequence to compare to.
        :param float lambda_: The KC metric lambda parameter determining the
            relative weight of topology and branch length.
        :return: The computed KC distance between this tree sequence and other.
        :rtype: float
        """
        return self._ll_tree_sequence.get_kc_distance(other._ll_tree_sequence, lambda_)

    def count_topologies(self, sample_sets=None):
        """
        Returns a generator that produces the same distribution of topologies as
        :meth:`Tree.count_topologies` but sequentially for every tree in a tree
        sequence. For use on a tree sequence this method is much faster than
        computing the result independently per tree.

        .. warning:: The interface for this method is preliminary and may be subject to
            backwards incompatible changes in the near future.

        :param list sample_sets: A list of lists of Node IDs, specifying the
            groups of nodes to compute the statistic with.
        :rtype: iter(:class:`tskit.TopologyCounter`)
        :raises ValueError: If nodes in ``sample_sets`` are invalid or are
            internal samples.
        """
        if sample_sets is None:
            sample_sets = [
                self.samples(population=pop.id) for pop in self.populations()
            ]

        yield from combinatorics.treeseq_count_topologies(self, sample_sets)

    def ibd_segments(
        self,
        *,
        within=None,
        between=None,
        max_time=None,
        min_span=None,
        store_pairs=None,
        store_segments=None,
    ):
        """
        Finds pairs of samples that are identical by descent (IBD) and returns
        the result as an :class:`.IdentitySegments` instance. The information
        stored in this object is controlled by the ``store_pairs`` and
        ``store_segments`` parameters. By default only total counts and other
        statistics of the IBD segments are stored (i.e.,
        ``store_pairs=False``), since storing pairs and segments has a
        substantial CPU and memory overhead. Please see the
        :ref:`sec_identity` section for more details on how to access the
        information stored in the :class:`.IdentitySegments`.

        If ``within`` is specified, only IBD segments for pairs of nodes within
        that set will be recorded. If ``between`` is specified, only IBD
        segments from pairs that are in one or other of the specified sample
        sets will be reported. Note that ``within`` and ``between`` are
        mutually exclusive.

        A pair of nodes ``(u, v)`` has an IBD segment with a left and right
        coordinate ``[left, right)`` and ancestral node ``a`` iff the most
        recent common ancestor of the segment ``[left, right)`` in nodes ``u``
        and ``v`` is ``a``, and the segment has been inherited along the same
        genealogical path (ie. it has not been broken by recombination). The
        segments returned are the longest possible ones.

        Note that this definition is purely genealogical --- allelic states
        *are not* considered here. If used without time or length thresholds, the
        segments returned for a given pair will partition the span of the contig
        represented by the tree sequence.

        :param list within: A list of node IDs defining set of nodes that
            we finding IBD segments for. If not specified, this defaults to
            all samples in the tree sequence.
        :param list[list] between: A list of lists of sample node IDs. Given
            two sample sets A and B, only IBD segments will be returned such
            that one of the samples is an element of A and the other is
            an element of B. Cannot be specified with ``within``.
        :param float max_time: Only segments inherited from common
            ancestors whose node times are more recent than the specified time
            will be returned. Specifying a maximum time is strongly recommended when
            working with large tree sequences.
        :param float min_span: Only segments in which the difference between
            the right and left genome coordinates (i.e., the span of the
            segment) is greater than this value will be included. (Default=0)
        :param bool store_pairs: If True store information separately for each
            pair of samples ``(a, b)`` that are found to be IBD. Otherwise
            store summary information about all sample apirs. (Default=False)
        :param bool store_segments: If True store each IBD segment
            ``(left, right, c)`` and associate it with the corresponding
            sample pair ``(a, b)``. If True, implies ``store_pairs``.
            (Default=False).
        :return: An :class:`.IdentitySegments` object containing the recorded
            IBD information.
        :rtype: IdentitySegments
        """
        return self.tables.ibd_segments(
            within=within,
            between=between,
            max_time=max_time,
            min_span=min_span,
            store_segments=store_segments,
            store_pairs=store_pairs,
        )

    def pair_coalescence_counts(
        self,
        sample_sets=None,
        indexes=None,
        windows=None,
        span_normalise=True,
        pair_normalise=False,
        time_windows="nodes",
    ):
        """
        Calculate the number of coalescing sample pairs per node, summed over
        trees and weighted by tree span.

        The number of coalescing pairs may be calculated within or between the
        non-overlapping lists of samples contained in `sample_sets`. In the
        latter case, pairs are counted if they have exactly one member in each
        of two sample sets. If `sample_sets` is omitted, a single set
        containing all samples is assumed.

        The argument `indexes` may be used to specify which pairs of sample
        sets to compute the statistic between, and in what order. If
        `indexes=None`, then `indexes` is assumed to equal `[(0,0)]` for a
        single sample set and `[(0,1)]` for two sample sets. For more than two
        sample sets, `indexes` must be explicitly passed.

        The argument `time_windows` may be used to count coalescence
        events within time intervals (if an array of breakpoints is supplied)
        rather than for individual nodes (the default).

        The output array has dimension `(windows, indexes, nodes)` with
        dimensions dropped when the corresponding argument is set to None.

        :param list sample_sets: A list of lists of Node IDs, specifying the
            groups of nodes to compute the statistic with, or None.
        :param list indexes: A list of 2-tuples, or None.
        :param list windows: An increasing list of breakpoints between the
            sequence windows to compute the statistic in, or None.
        :param bool span_normalise: Whether to divide the result by the span of
            the window (defaults to True).
        :param bool pair_normalise: Whether to divide the result by the total
            number of pairs for a given index (defaults to False).
        :param time_windows: Either a string "nodes" or an increasing
            list of breakpoints between time intervals.
        """

        if sample_sets is None:
            sample_sets = [list(self.samples())]

        drop_middle_dimension = False
        if indexes is None:
            drop_middle_dimension = True
            if len(sample_sets) == 1:
                indexes = [(0, 0)]
            elif len(sample_sets) == 2:
                indexes = [(0, 1)]
            else:
                raise ValueError(
                    "Must specify indexes if there are more than two sample sets"
                )

        drop_left_dimension = False
        if windows is None:
            drop_left_dimension = True
            windows = np.array([0.0, self.sequence_length])

        if isinstance(time_windows, str) and time_windows == "nodes":
            node_bin_map = np.arange(self.num_nodes, dtype=np.int32)
        else:
            if self.time_units == tskit.TIME_UNITS_UNCALIBRATED:
                raise ValueError("Time windows require calibrated node times")
            node_bin_map = np.digitize(self.nodes_time, time_windows) - 1
            node_bin_map[node_bin_map == time_windows.size - 1] = tskit.NULL
            node_bin_map = node_bin_map.astype(np.int32)

        sample_set_sizes = np.array([len(s) for s in sample_sets], dtype=np.uint32)
        sample_sets = util.safe_np_int_cast(np.hstack(sample_sets), np.int32)

        coalescing_pairs = self.ll_tree_sequence.pair_coalescence_counts(
            sample_sets=sample_sets,
            sample_set_sizes=sample_set_sizes,
            windows=windows,
            indexes=indexes,
            node_bin_map=node_bin_map,
            span_normalise=span_normalise,
            pair_normalise=pair_normalise,
        )

        if drop_middle_dimension:
            coalescing_pairs = np.squeeze(coalescing_pairs, axis=1)
        if drop_left_dimension:
            coalescing_pairs = np.squeeze(coalescing_pairs, axis=0)

        return coalescing_pairs

    def pair_coalescence_quantiles(
        self,
        quantiles,
        sample_sets=None,
        indexes=None,
        windows=None,
    ):
        """
        Estimate quantiles of pair coalescence times by inverting the empirical
        CDF. This is equivalent to the "inverted_cdf" method of
        `numpy.quantile` applied to node times, with weights proportional to
        the number of coalescing pairs per node (averaged over trees, see
        `TreeSequence.pair_coalescence_counts`).

        Quantiles of pair coalescence times may be calculated within or
        between the non-overlapping lists of samples contained in `sample_sets`. In
        the latter case, pairs are counted if they have exactly one member in each
        of two sample sets. If `sample_sets` is omitted, a single set containing
        all samples is assumed.

        The argument `indexes` may be used to specify which pairs of sample sets to
        compute coalescences between, and in what order. If `indexes=None`, then
        `indexes` is assumed to equal `[(0,0)]` for a single sample set and
        `[(0,1)]` for two sample sets. For more than two sample sets, `indexes`
        must be explicitly passed.

        The output array has dimension `(windows, indexes, quantiles)` with
        dimensions dropped when the corresponding argument is set to None.

        :param quantiles: A list of increasing breakpoints between [0, 1].
        :param list sample_sets: A list of lists of Node IDs, specifying the
            groups of nodes to compute the statistic with, or None.
        :param list indexes: A list of 2-tuples, or None.
        :param list windows: An increasing list of breakpoints between the
            sequence windows to compute the statistic in, or None.
        """

        if sample_sets is None:
            sample_sets = [list(self.samples())]

        drop_middle_dimension = False
        if indexes is None:
            drop_middle_dimension = True
            if len(sample_sets) == 1:
                indexes = [(0, 0)]
            elif len(sample_sets) == 2:
                indexes = [(0, 1)]
            else:
                raise ValueError(
                    "Must specify indexes if there are more than two sample sets"
                )

        if self.time_units == tskit.TIME_UNITS_UNCALIBRATED:
            raise ValueError("Pair coalescence quantiles require calibrated node times")

        drop_left_dimension = False
        if windows is None:
            drop_left_dimension = True
            windows = np.array([0.0, self.sequence_length])

        sample_set_sizes = np.array([len(s) for s in sample_sets], dtype=np.uint32)
        sample_sets = util.safe_np_int_cast(np.hstack(sample_sets), np.int32)
        _, node_bin_map = np.unique(self.nodes_time, return_inverse=True)
        node_bin_map = util.safe_np_int_cast(node_bin_map, np.int32)

        coalescence_times = self.ll_tree_sequence.pair_coalescence_quantiles(
            sample_sets=sample_sets,
            sample_set_sizes=sample_set_sizes,
            windows=windows,
            indexes=indexes,
            node_bin_map=node_bin_map,
            quantiles=quantiles,
        )

        if drop_middle_dimension:
            coalescence_times = np.squeeze(coalescence_times, axis=1)
        if drop_left_dimension:
            coalescence_times = np.squeeze(coalescence_times, axis=0)

        return coalescence_times

    def pair_coalescence_rates(
        self,
        time_windows,
        sample_sets=None,
        indexes=None,
        windows=None,
    ):
        """
        Estimate the rate at which pairs of samples coalesce within time
        windows, using the empirical cumulative distribution function (ecdf) of
        pair coalescence times.  Assuming that pair coalescence events follow a
        nonhomogeneous Poisson process, the empirical rate for a time window
        :math:`[a, b)` where :math:`ecdf(b) < 1` is,

        ..math:

            log(1 - \\frac{ecdf(b) - ecdf(a)}{1 - ecdf(a)}) / (a - b)

        If the last coalescence event is within :math:`[a, b)`, so that
        :math:`ecdf(b) = 1`, then an estimate of the empirical rate is

        ..math:

            (\\mathbb{E}[t | t > a] - a)^{-1}

        where :math:`\\mathbb{E}[t | t < a]` is the average pair coalescence time
        conditional on coalescence after the start of the last epoch.

        The first breakpoint in `time_windows` must start at the age of the
        samples, and the last must end at infinity. In the output array, any
        time windows that do not contain coalescence events will have `NaN`
        values.

        Pair coalescence rates may be calculated within or between the
        non-overlapping lists of samples contained in `sample_sets`. In the
        latter case, pairs are counted if they have exactly one member in each
        of two sample sets. If `sample_sets` is omitted, a single group
        containing all samples is assumed.

        The argument `indexes` may be used to specify which pairs of sample
        sets to compute the statistic between, and in what order. If
        `indexes=None`, then `indexes` is assumed to equal `[(0,0)]` for a
        single sample set and `[(0,1)]` for two sample sets. For more than two
        sample sets, `indexes` must be explicitly passed.

        The output array has dimension `(windows, indexes, time_windows)` with
        dimensions dropped when the corresponding argument is set to None.

        :param time_windows: An increasing list of breakpoints between time
            intervals, starting at the age of the samples and ending at
            infinity.
        :param list sample_sets: A list of lists of Node IDs, specifying the
            groups of nodes to compute the statistic with, or None.
        :param list indexes: A list of 2-tuples, or None.
        :param list windows: An increasing list of breakpoints between the
            sequence windows to compute the statistic in, or None.
        """

        if sample_sets is None:
            sample_sets = [list(self.samples())]

        drop_middle_dimension = False
        if indexes is None:
            drop_middle_dimension = True
            if len(sample_sets) == 1:
                indexes = [(0, 0)]
            elif len(sample_sets) == 2:
                indexes = [(0, 1)]
            else:
                raise ValueError(
                    "Must specify indexes if there are more than two sample sets"
                )

        if self.time_units == tskit.TIME_UNITS_UNCALIBRATED:
            raise ValueError("Pair coalescence rates require calibrated node times")

        drop_left_dimension = False
        if windows is None:
            drop_left_dimension = True
            windows = np.array([0.0, self.sequence_length])

        sample_set_sizes = np.array([len(s) for s in sample_sets], dtype=np.uint32)
        sample_sets = util.safe_np_int_cast(np.hstack(sample_sets), np.int32)
        node_bin_map = np.digitize(self.nodes_time, time_windows) - 1
        node_bin_map[node_bin_map == time_windows.size - 1] = tskit.NULL
        node_bin_map = node_bin_map.astype(np.int32)

        coalescence_rates = self.ll_tree_sequence.pair_coalescence_rates(
            sample_sets=sample_sets,
            sample_set_sizes=sample_set_sizes,
            windows=windows,
            indexes=indexes,
            node_bin_map=node_bin_map,
            time_windows=time_windows,
        )

        if drop_middle_dimension:
            coalescence_rates = np.squeeze(coalescence_rates, axis=1)
        if drop_left_dimension:
            coalescence_rates = np.squeeze(coalescence_rates, axis=0)

        return coalescence_rates

    def impute_unknown_mutations_time(
        self,
        method=None,
    ):
        """
        Returns an array of mutation times, where any unknown times are
        imputed from the times of associated nodes. Not to be confused with
        :meth:`TableCollection.compute_mutation_times`, which modifies the
        ``time`` column of the mutations table in place.

        :param str method: The method used to impute the unknown mutation times.
            Currently only "min" is supported, which uses the time of the node
            below the mutation as the mutation time. The "min" method can also
            be specified by ``method=None`` (Default: ``None``).
        :return: An array of length equal to the number of mutations in the
            tree sequence.
        """
        allowed_methods = ["min"]
        if method is None:
            method = "min"
        if method not in allowed_methods:
            raise ValueError(
                f"Mutations time imputation method must be chosen from {allowed_methods}"
            )
        if method == "min":
            mutations_time = self.mutations_time.copy()
            unknown = tskit.is_unknown_time(mutations_time)
            mutations_time[unknown] = self.nodes_time[self.mutations_node[unknown]]
            return mutations_time

    def ld_matrix(
        self, sample_sets=None, sites=None, positions=None, mode="site", stat="r2"
    ):
        stats = {
            "D": self._ll_tree_sequence.D_matrix,
            "D2": self._ll_tree_sequence.D2_matrix,
            "r2": self._ll_tree_sequence.r2_matrix,
            "D_prime": self._ll_tree_sequence.D_prime_matrix,
            "r": self._ll_tree_sequence.r_matrix,
            "Dz": self._ll_tree_sequence.Dz_matrix,
            "pi2": self._ll_tree_sequence.pi2_matrix,
            "Dz_unbiased": self._ll_tree_sequence.Dz_unbiased_matrix,
            "D2_unbiased": self._ll_tree_sequence.D2_unbiased_matrix,
            "pi2_unbiased": self._ll_tree_sequence.pi2_unbiased_matrix,
        }

        try:
            two_locus_stat = stats[stat]
        except KeyError:
            raise ValueError(
                f"Unknown two-locus statistic '{stat}', we support: {list(stats.keys())}"
            )

        return self.__two_locus_sample_set_stat(
            two_locus_stat,
            sample_sets,
            sites=sites,
            positions=positions,
            mode=mode,
        )

    ############################################
    #
    # Deprecated APIs. These are either already unsupported, or will be unsupported in a
    # later release.
    #
    ############################################

    def get_pairwise_diversity(self, samples=None):
        # Deprecated alias for self.pairwise_diversity
        return self.pairwise_diversity(samples)

    def pairwise_diversity(self, samples=None):
        """
        Returns the pairwise nucleotide site diversity, the average number of sites
        that differ between a every possible pair of distinct samples.  If `samples` is
        specified, calculate the diversity within this set.

         .. deprecated:: 0.2.0
             please use :meth:`.diversity` instead. Since version 0.2.0 the error
             semantics have also changed slightly. It is no longer an error
             when there is one sample and a tskit.LibraryError is raised
             when non-sample IDs are provided rather than a ValueError. It is
             also no longer an error to compute pairwise diversity at sites
             with multiple mutations.

        :param list samples: The set of samples within which we calculate
            the diversity. If None, calculate diversity within the entire sample.
        :return: The pairwise nucleotide site diversity.
        :rtype: float
        """
        if samples is None:
            samples = self.samples()
        return float(
            self.diversity(
                [samples], windows=[0, self.sequence_length], span_normalise=False
            )[0][0]
        )

    def get_time(self, u):
        # Deprecated. Use ts.node(u).time
        if u < 0 or u >= self.get_num_nodes():
            raise ValueError("ID out of bounds")
        node = self.node(u)
        return node.time

    def get_population(self, u):
        # Deprecated. Use ts.node(u).population
        if u < 0 or u >= self.get_num_nodes():
            raise ValueError("ID out of bounds")
        node = self.node(u)
        return node.population

    def records(self):
        # Deprecated. Use either ts.edges() or ts.edgesets().
        t = [node.time for node in self.nodes()]
        pop = [node.population for node in self.nodes()]
        for e in self.edgesets():
            yield CoalescenceRecord(
                e.left, e.right, e.parent, e.children, t[e.parent], pop[e.parent]
            )

    # Unsupported old methods.

    def get_num_records(self):
        raise NotImplementedError(
            "This method is no longer supported. Please use the "
            "TreeSequence.num_edges if possible to work with edges rather "
            "than coalescence records. If not, please use len(list(ts.edgesets())) "
            "which should return the number of coalescence records, as previously "
            "defined. Please open an issue on GitHub if this is "
            "important for your workflow."
        )

    def diffs(self):
        raise NotImplementedError(
            "This method is no longer supported. Please use the "
            "TreeSequence.edge_diffs() method instead"
        )

    def newick_trees(self, precision=3, breakpoints=None, Ne=1):
        raise NotImplementedError(
            "This method is no longer supported. Please use the Tree.newick"
            " method instead"
        )

    def to_nexus(self, precision=14):
        raise NotImplementedError(
            "This method is no longer supported since 0.4.0. Please use the as_nexus "
            "or write_nexus methods instead"
        )


# TODO move to "text_formats.py"
def write_ms(
    tree_sequence,
    output,
    print_trees=False,
    precision=4,
    num_replicates=1,
    write_header=True,
):
    """
    Write ``ms`` formatted output from the genotypes of a tree sequence
    or an iterator over tree sequences. Usage:

    .. code-block:: python

        import tskit as ts

        tree_sequence = msprime.simulate(
            sample_size=sample_size,
            Ne=Ne,
            length=length,
            mutation_rate=mutation_rate,
            recombination_rate=recombination_rate,
            random_seed=random_seed,
            num_replicates=num_replicates,
        )
        with open("output.ms", "w") as ms_file:
            ts.write_ms(tree_sequence, ms_file)

    :param ts tree_sequence: The tree sequence (or iterator over tree sequences) to
        write to ms file
    :param io.IOBase output: The file-like object to write the ms-style output
    :param bool print_trees: Boolean parameter to write out newick format trees
        to output [optional]
    :param int precision: Numerical precision with which to write the ms
        output [optional]
    :param bool write_header: Boolean parameter to write out the header. [optional]
    :param int num_replicates: Number of replicates simulated [required if
        num_replicates used in simulation]

    The first line of this ms-style output file written has two arguments which
    are sample size and number of replicates. The second line has a 0 as a substitute
    for the random seed.
    """
    if not isinstance(tree_sequence, collections.abc.Iterable):
        tree_sequence = [tree_sequence]

    i = 0
    for tree_seq in tree_sequence:
        if i > 0:
            write_header = False
        i = i + 1

        if write_header is True:
            print(
                f"ms {tree_seq.sample_size} {num_replicates}",
                file=output,
            )
            print("0", file=output)

        print(file=output)
        print("//", file=output)
        if print_trees is True:
            """
            Print out the trees in ms-format from the specified tree sequence.
            """
            if len(tree_seq.trees()) == 1:
                tree = next(tree_seq.trees())
                newick = tree.newick(precision=precision)
                print(newick, file=output)
            else:
                for tree in tree_seq.trees():
                    newick = tree.newick(precision=precision)
                    print(f"[{tree.span:.{precision}f}]", newick, file=output)

        else:
            s = tree_seq.get_num_sites()
            print("segsites:", s, file=output)
            if s != 0:
                print("positions: ", end="", file=output)
                positions = [
                    variant.position / (tree_seq.sequence_length)
                    for variant in tree_seq.variants()
                ]
                for position in positions:
                    print(
                        f"{position:.{precision}f}",
                        end=" ",
                        file=output,
                    )
                print(file=output)

                genotypes = tree_seq.genotype_matrix()
                for k in range(tree_seq.num_samples):
                    tmp_str = "".join(map(str, genotypes[:, k]))
                    if set(tmp_str).issubset({"0", "1", "-"}):
                        print(tmp_str, file=output)
                    else:
                        raise ValueError(
                            "This tree sequence contains non-biallelic"
                            "SNPs and is incompatible with the ms format!"
                        )
            else:
                print(file=output)


--- ../../tskit/python/tskit/formats.py ---


"""
Module responsible for converting tree sequence files from older
formats.
"""
import datetime
import json
import logging

import numpy as np

import tskit
import tskit.exceptions as exceptions
import tskit.provenance as provenance


def _get_v2_provenance(command, attrs):
    """
    Returns the V2 tree provenance attributes reformatted as a provenance record.
    """
    environment = {}
    parameters = {}
    # Try to get the provenance strings. Malformed JSON should not prevent us
    # from finishing the conversion.
    try:
        environment = json.loads(str(attrs["environment"]))
    except ValueError:
        logging.warning("Failed to convert environment provenance")
    try:
        parameters = json.loads(str(attrs["parameters"]))
    except ValueError:
        logging.warning("Failed to convert parameters provenance")
    parameters["command"] = command
    provenance_dict = provenance.get_provenance_dict(parameters)
    provenance_dict["version"] = environment.get("msprime_version", "Unknown_version")
    provenance_dict["environment"] = environment
    return json.dumps(provenance_dict).encode()


def _get_upgrade_provenance(root):
    """
    Returns the provenance string from upgrading the specified HDF5 file.
    """
    # TODO add more parameters here like filename, etc.
    parameters = {
        "command": "upgrade",
        "source_version": list(map(int, root.attrs["format_version"])),
    }
    s = json.dumps(provenance.get_provenance_dict(parameters))
    return s.encode()


def _convert_hdf5_mutations(
    mutations_group, sites, mutations, remove_duplicate_positions
):
    """
    Loads the v2/v3 into the specified tables.
    """
    position = np.array(mutations_group["position"])
    node = np.array(mutations_group["node"], dtype=np.int32)
    unique_position, index = np.unique(position, return_index=True)
    if unique_position.shape != position.shape:
        if remove_duplicate_positions:
            position = position[index]
            node = node[index]
        else:
            # TODO add the number of duplicates so that we can improve the
            # error message.
            raise exceptions.DuplicatePositionsError()
    num_mutations = position.shape[0]
    sites.set_columns(
        position=position,
        ancestral_state=ord("0") * np.ones(num_mutations, dtype=np.int8),
        ancestral_state_offset=np.arange(num_mutations + 1, dtype=np.uint32),
    )
    mutations.set_columns(
        node=node,
        site=np.arange(num_mutations, dtype=np.int32),
        derived_state=ord("1") * np.ones(num_mutations, dtype=np.int8),
        derived_state_offset=np.arange(num_mutations + 1, dtype=np.uint32),
    )


def _set_populations(tables):
    """
    Updates PopulationTable suitable to represent the populations referred to
    in the node table.
    """
    if len(tables.nodes) > 0:
        for _ in range(np.max(tables.nodes.population) + 1):
            tables.populations.add_row()


def _load_legacy_hdf5_v2(root, remove_duplicate_positions):
    # Get the coalescence records
    trees_group = root["trees"]
    old_timestamp = datetime.datetime.min.isoformat()
    provenances = tskit.ProvenanceTable()
    provenances.add_row(
        timestamp=old_timestamp,
        record=_get_v2_provenance("generate_trees", trees_group.attrs),
    )
    num_rows = trees_group["node"].shape[0]
    index = np.arange(num_rows, dtype=int)
    parent = np.zeros(2 * num_rows, dtype=np.int32)
    parent[2 * index] = trees_group["node"]
    parent[2 * index + 1] = trees_group["node"]
    left = np.zeros(2 * num_rows, dtype=np.float64)
    left[2 * index] = trees_group["left"]
    left[2 * index + 1] = trees_group["left"]
    right = np.zeros(2 * num_rows, dtype=np.float64)
    right[2 * index] = trees_group["right"]
    right[2 * index + 1] = trees_group["right"]
    child = np.array(trees_group["children"], dtype=np.int32).flatten()

    tables = tskit.TableCollection(np.max(right))
    tables.edges.set_columns(left=left, right=right, parent=parent, child=child)

    cr_node = np.array(trees_group["node"], dtype=np.int32)
    num_nodes = max(np.max(child), np.max(cr_node)) + 1
    sample_size = np.min(cr_node)
    flags = np.zeros(num_nodes, dtype=np.uint32)
    population = np.zeros(num_nodes, dtype=np.int32)
    time = np.zeros(num_nodes, dtype=np.float64)
    flags[:sample_size] = tskit.NODE_IS_SAMPLE
    cr_population = np.array(trees_group["population"], dtype=np.int32)
    cr_time = np.array(trees_group["time"])
    time[cr_node] = cr_time
    population[cr_node] = cr_population
    if "samples" in root:
        samples_group = root["samples"]
        population[:sample_size] = np.array(samples_group["population"], copy=True)
        if "time" in samples_group:
            time[:sample_size] = np.array(samples_group["time"], copy=True)
    tables.nodes.set_columns(flags=flags, population=population, time=time)
    _set_populations(tables)

    if "mutations" in root:
        mutations_group = root["mutations"]
        _convert_hdf5_mutations(
            mutations_group, tables.sites, tables.mutations, remove_duplicate_positions
        )
        provenances.add_row(
            timestamp=old_timestamp,
            record=_get_v2_provenance("generate_mutations", mutations_group.attrs),
        )
    tables.provenances.add_row(_get_upgrade_provenance(root))
    tables.sort()
    return tables.tree_sequence()


def _load_legacy_hdf5_v3(root, remove_duplicate_positions):
    # get the trees group for the records and samples
    trees_group = root["trees"]
    nodes_group = trees_group["nodes"]
    time = np.array(nodes_group["time"])

    breakpoints = np.array(trees_group["breakpoints"])
    records_group = trees_group["records"]
    left_indexes = np.array(records_group["left"])
    right_indexes = np.array(records_group["right"])
    record_node = np.array(records_group["node"], dtype=np.int32)
    num_nodes = time.shape[0]
    sample_size = np.min(record_node)
    flags = np.zeros(num_nodes, dtype=np.uint32)
    flags[:sample_size] = tskit.NODE_IS_SAMPLE

    children_length = np.array(records_group["num_children"], dtype=np.uint32)
    total_rows = np.sum(children_length)
    left = np.zeros(total_rows, dtype=np.float64)
    right = np.zeros(total_rows, dtype=np.float64)
    parent = np.zeros(total_rows, dtype=np.int32)
    record_left = breakpoints[left_indexes]
    record_right = breakpoints[right_indexes]
    k = 0
    for j in range(left_indexes.shape[0]):
        for _ in range(children_length[j]):
            left[k] = record_left[j]
            right[k] = record_right[j]
            parent[k] = record_node[j]
            k += 1
    tables = tskit.TableCollection(np.max(right))
    tables.nodes.set_columns(
        flags=flags, time=nodes_group["time"], population=nodes_group["population"]
    )
    _set_populations(tables)
    tables.edges.set_columns(
        left=left, right=right, parent=parent, child=records_group["children"]
    )
    if "mutations" in root:
        _convert_hdf5_mutations(
            root["mutations"],
            tables.sites,
            tables.mutations,
            remove_duplicate_positions,
        )
    old_timestamp = datetime.datetime.min.isoformat()
    if "provenance" in root:
        for record in root["provenance"]:
            tables.provenances.add_row(timestamp=old_timestamp, record=record)
    tables.provenances.add_row(_get_upgrade_provenance(root))
    tables.sort()
    return tables.tree_sequence()


def get_h5py():
    try:
        import h5py
    except ImportError:
        raise ImportError(
            "Legacy formats require h5py. Install via `pip install h5py`"
            " or `conda install h5py`"
        )
    return h5py


def load_legacy(filename, remove_duplicate_positions=False):
    """
    Reads the specified msprime HDF5 file and returns a tree sequence. This
    method is only intended to be used to read old format HDF5 files.

    If remove_duplicate_positions is True, remove all sites (except the
    first) that contain duplicate positions. If this is False, any input
    files that contain duplicate positions will raise an DuplicatePositionsError.
    """
    loaders = {
        2: _load_legacy_hdf5_v2,
        3: _load_legacy_hdf5_v3,
        10: _load_legacy_hdf5_v10,
    }
    h5py = get_h5py()
    root = h5py.File(filename, "r")
    if "format_version" not in root.attrs:
        raise ValueError("HDF5 file not in msprime format")
    format_version = root.attrs["format_version"]
    if format_version[0] not in loaders:
        raise ValueError(f"Version {format_version} not supported for loading")
    try:
        ts = loaders[format_version[0]](root, remove_duplicate_positions)
    finally:
        root.close()
    return ts


def _dump_legacy_hdf5_v2(tree_sequence, root):
    root.attrs["format_version"] = (2, 999)
    root.attrs["sample_size"] = tree_sequence.get_sample_size()
    root.attrs["sequence_length"] = (tree_sequence.get_sequence_length(),)
    left = []
    right = []
    node = []
    children = []
    time = []
    population = []
    for record in tree_sequence.records():
        left.append(record.left)
        right.append(record.right)
        node.append(record.node)
        if len(record.children) != 2:
            raise ValueError("V2 files only support binary records")
        children.append(record.children)
        time.append(record.time)
        population.append(record.population)
    length = len(time)
    trees = root.create_group("trees")
    trees.attrs["environment"] = json.dumps({"msprime_version": 0})
    trees.attrs["parameters"] = "{}"
    trees.create_dataset("left", (length,), data=left, dtype=float)
    trees.create_dataset("right", (length,), data=right, dtype=float)
    trees.create_dataset("time", (length,), data=time, dtype=float)
    trees.create_dataset("node", (length,), data=node, dtype="u4")
    trees.create_dataset("population", (length,), data=population, dtype="u1")
    trees.create_dataset("children", (length, 2), data=children, dtype="u4")
    samples = root.create_group("samples")
    population = []
    time = []
    length = tree_sequence.get_sample_size()
    for u in range(length):
        time.append(tree_sequence.get_time(u))
        population.append(tree_sequence.get_population(u))
    samples.create_dataset("time", (length,), data=time, dtype=float)
    samples.create_dataset("population", (length,), data=population, dtype="u1")
    if tree_sequence.get_num_mutations() > 0:
        node = []
        position = []
        for site in tree_sequence.sites():
            if len(site.mutations) != 1:
                raise ValueError("v2 does not support recurrent mutations")
            if site.ancestral_state != "0" or site.mutations[0].derived_state != "1":
                raise ValueError("v2 does not support non-binary mutations")
            position.append(site.position)
            node.append(site.mutations[0].node)
        length = len(node)
        mutations = root.create_group("mutations")
        mutations.attrs["environment"] = json.dumps({"msprime_version": 0})
        mutations.attrs["parameters"] = "{}"
        mutations.create_dataset("position", (length,), data=position, dtype=float)
        mutations.create_dataset("node", (length,), data=node, dtype="u4")


def _dump_legacy_hdf5_v3(tree_sequence, root):
    root.attrs["format_version"] = (3, 999)
    root.attrs["sample_size"] = (0,)
    root.attrs["sequence_length"] = (0,)
    trees = root.create_group("trees")
    # Get the breakpoints from the records.
    left = [cr.left for cr in tree_sequence.records()]
    breakpoints = np.unique(left + [tree_sequence.sequence_length])
    trees.create_dataset(
        "breakpoints", (len(breakpoints),), data=breakpoints, dtype=float
    )

    left = []
    right = []
    node = []
    children = []
    num_children = []
    time = []
    for cr in tree_sequence.records():
        node.append(cr.node)
        left.append(np.searchsorted(breakpoints, cr.left))
        right.append(np.searchsorted(breakpoints, cr.right))
        children.extend(cr.children)
        num_children.append(len(cr.children))
        time.append(cr.time)
    records_group = trees.create_group("records")
    length = len(num_children)
    records_group.create_dataset("left", (length,), data=left, dtype="u4")
    records_group.create_dataset("right", (length,), data=right, dtype="u4")
    records_group.create_dataset("node", (length,), data=node, dtype="u4")
    records_group.create_dataset(
        "num_children", (length,), data=num_children, dtype="u4"
    )
    records_group.create_dataset(
        "children", (len(children),), data=children, dtype="u4"
    )

    indexes_group = trees.create_group("indexes")
    left_index = sorted(range(length), key=lambda j: (left[j], time[j]))
    right_index = sorted(range(length), key=lambda j: (right[j], -time[j]))
    indexes_group.create_dataset(
        "insertion_order", (length,), data=left_index, dtype="u4"
    )
    indexes_group.create_dataset(
        "removal_order", (length,), data=right_index, dtype="u4"
    )

    nodes_group = trees.create_group("nodes")
    population = np.zeros(tree_sequence.num_nodes, dtype="u4")
    time = np.zeros(tree_sequence.num_nodes, dtype=float)
    tree = next(tree_sequence.trees())
    for u in range(tree_sequence.sample_size):
        population[u] = tree.population(u)
        time[u] = tree.time(u)
    for cr in tree_sequence.records():
        population[cr.node] = cr.population
        time[cr.node] = cr.time
    length = tree_sequence.num_nodes
    nodes_group.create_dataset("time", (length,), data=time, dtype=float)
    nodes_group.create_dataset("population", (length,), data=population, dtype="u4")

    node = []
    position = []
    for site in tree_sequence.sites():
        if len(site.mutations) != 1:
            raise ValueError("v3 does not support recurrent mutations")
        if site.ancestral_state != "0" or site.mutations[0].derived_state != "1":
            raise ValueError("v3 does not support non-binary mutations")
        position.append(site.position)
        node.append(site.mutations[0].node)
    length = len(position)
    if length > 0:
        mutations = root.create_group("mutations")
        mutations.create_dataset("position", (length,), data=position, dtype=float)
        mutations.create_dataset("node", (length,), data=node, dtype="u4")


def _add_dataset(group, name, data):
    # In the HDF5 format any zero-d arrays must be excluded.
    if data.shape[0] > 0:
        group.create_dataset(name, data=data)


def _dump_legacy_hdf5_v10(tree_sequence, root):
    root.attrs["format_version"] = (10, 999)
    root.attrs["sample_size"] = (0,)
    root.attrs["sequence_length"] = (tree_sequence.sequence_length,)
    tables = tree_sequence.dump_tables()

    nodes = root.create_group("nodes")
    _add_dataset(nodes, "time", tables.nodes.time)
    _add_dataset(nodes, "flags", tables.nodes.flags)
    _add_dataset(nodes, "population", tables.nodes.population)
    _add_dataset(nodes, "metadata", tables.nodes.metadata)
    _add_dataset(nodes, "metadata_offset", tables.nodes.metadata_offset)

    edges = root.create_group("edges")
    if len(tables.edges) > 0:
        edges.create_dataset("left", data=tables.edges.left)
        edges.create_dataset("right", data=tables.edges.right)
        edges.create_dataset("parent", data=tables.edges.parent)
        edges.create_dataset("child", data=tables.edges.child)

        left = tables.edges.left
        right = tables.edges.right
        time = tables.nodes.time[tables.edges.parent]
        # We can do this more efficiently if we ever need to do it for anything
        # other than testing.
        indexes_group = edges.create_group("indexes")
        length = len(tables.edges)
        left_index = sorted(range(length), key=lambda j: (left[j], time[j]))
        right_index = sorted(range(length), key=lambda j: (right[j], -time[j]))
        indexes_group.create_dataset("insertion_order", data=left_index, dtype="u4")
        indexes_group.create_dataset("removal_order", data=right_index, dtype="u4")

    migrations = root.create_group("migrations")
    if len(tables.migrations) > 0:
        migrations.create_dataset("left", data=tables.migrations.left)
        migrations.create_dataset("right", data=tables.migrations.right)
        migrations.create_dataset("node", data=tables.migrations.node)
        migrations.create_dataset("source", data=tables.migrations.source)
        migrations.create_dataset("dest", data=tables.migrations.dest)
        migrations.create_dataset("time", data=tables.migrations.time)

    sites = root.create_group("sites")
    _add_dataset(sites, "position", tables.sites.position)
    _add_dataset(sites, "ancestral_state", tables.sites.ancestral_state)
    _add_dataset(sites, "ancestral_state_offset", tables.sites.ancestral_state_offset)
    _add_dataset(sites, "metadata", tables.sites.metadata)
    _add_dataset(sites, "metadata_offset", tables.sites.metadata_offset)

    mutations = root.create_group("mutations")
    _add_dataset(mutations, "site", tables.mutations.site)
    _add_dataset(mutations, "node", tables.mutations.node)
    _add_dataset(mutations, "parent", tables.mutations.parent)
    _add_dataset(mutations, "derived_state", tables.mutations.derived_state)
    _add_dataset(
        mutations, "derived_state_offset", tables.mutations.derived_state_offset
    )
    _add_dataset(mutations, "metadata", tables.mutations.metadata)
    _add_dataset(mutations, "metadata_offset", tables.mutations.metadata_offset)

    provenances = root.create_group("provenances")
    _add_dataset(provenances, "timestamp", tables.provenances.timestamp)
    _add_dataset(provenances, "timestamp_offset", tables.provenances.timestamp_offset)
    _add_dataset(provenances, "record", tables.provenances.record)
    _add_dataset(provenances, "record_offset", tables.provenances.record_offset)


def _load_legacy_hdf5_v10(root, remove_duplicate_positions=False):
    # We cannot have duplicate positions in v10, so this parameter is ignored
    sequence_length = root.attrs["sequence_length"]
    try:
        sequence_length = sequence_length[0]
    except TypeError:
        pass
    tables = tskit.TableCollection(sequence_length)

    nodes_group = root["nodes"]
    metadata = None
    metadata_offset = None
    if "metadata" in nodes_group:
        metadata = nodes_group["metadata"]
        metadata_offset = nodes_group["metadata_offset"]
    if "flags" in nodes_group:
        tables.nodes.set_columns(
            flags=nodes_group["flags"],
            population=nodes_group["population"],
            time=nodes_group["time"],
            metadata=metadata,
            metadata_offset=metadata_offset,
        )

    edges_group = root["edges"]
    if "left" in edges_group:
        tables.edges.set_columns(
            left=edges_group["left"],
            right=edges_group["right"],
            parent=edges_group["parent"],
            child=edges_group["child"],
        )

    migrations_group = root["migrations"]
    if "left" in migrations_group:
        tables.migrations.set_columns(
            left=migrations_group["left"],
            right=migrations_group["right"],
            node=migrations_group["node"],
            source=migrations_group["source"],
            dest=migrations_group["dest"],
            time=migrations_group["time"],
        )

    sites_group = root["sites"]
    if "position" in sites_group:
        metadata = None
        metadata_offset = None
        if "metadata" in sites_group:
            metadata = sites_group["metadata"]
            metadata_offset = sites_group["metadata_offset"]
        tables.sites.set_columns(
            position=sites_group["position"],
            ancestral_state=sites_group["ancestral_state"],
            ancestral_state_offset=sites_group["ancestral_state_offset"],
            metadata=metadata,
            metadata_offset=metadata_offset,
        )

    mutations_group = root["mutations"]
    if "site" in mutations_group:
        metadata = None
        metadata_offset = None
        if "metadata" in mutations_group:
            metadata = mutations_group["metadata"]
            metadata_offset = mutations_group["metadata_offset"]
        tables.mutations.set_columns(
            site=mutations_group["site"],
            node=mutations_group["node"],
            parent=mutations_group["parent"],
            derived_state=mutations_group["derived_state"],
            derived_state_offset=mutations_group["derived_state_offset"],
            metadata=metadata,
            metadata_offset=metadata_offset,
        )

    provenances_group = root["provenances"]
    if "timestamp" in provenances_group:
        timestamp = provenances_group["timestamp"]
        timestamp_offset = provenances_group["timestamp_offset"]
        record = provenances_group["record"]
        record_offset = provenances_group["record_offset"]
        tables.provenances.set_columns(
            timestamp=timestamp,
            timestamp_offset=timestamp_offset,
            record=record,
            record_offset=record_offset,
        )
    tables.provenances.add_row(_get_upgrade_provenance(root))
    _set_populations(tables)
    return tables.tree_sequence()


def dump_legacy(tree_sequence, filename, version=3):
    """
    Writes the specified tree sequence to a HDF5 file in the specified
    legacy file format version.
    """
    dumpers = {
        2: _dump_legacy_hdf5_v2,
        3: _dump_legacy_hdf5_v3,
        10: _dump_legacy_hdf5_v10,
    }
    if version not in dumpers:
        raise ValueError(f"Version {version} file format is supported")
    h5py = get_h5py()
    root = h5py.File(filename, "w")
    try:
        dumpers[version](tree_sequence, root)
    finally:
        root.close()


--- ../../tskit/python/tskit/util.py ---

 
"""
Module responsible for various utility functions used in other modules.
"""
import dataclasses
import datetime
import html
import io
import itertools
import json
import numbers
import os
import textwrap
from typing import Union

import numpy as np

import tskit
from tskit import UNKNOWN_TIME


# Extra methods for dataclasses
class Dataclass:
    def replace(self, **kwargs):
        """
        Return a new instance of this dataclass, with the specified attributes
        overwritten by new values.

        :return: A new instance of the same type
        """
        return dataclasses.replace(self, **kwargs)

    def asdict(self, **kwargs):
        """
        Return a new dict which maps field names to their corresponding values
        in this dataclass.
        """
        return dataclasses.asdict(self, **kwargs)


def canonical_json(obj):
    """
    Returns string of encoded JSON with keys sorted and whitespace removed to enable
    byte-level comparison of encoded data.

    :param Any obj: Python object to encode
    :return: The encoded string
    :rtype: str
    """
    return json.dumps(obj, sort_keys=True, separators=(",", ":"))


def is_unknown_time(time):
    """
    As the default unknown mutation time (:const:`UNKNOWN_TIME`) is a specific NAN value,
    equality always fails (A NAN value is not equal to itself by definition).
    This method compares the bitfield such that unknown times can be detected. Either
    single floats can be passed or lists/arrays.

    Note that NANs are a set of floating-point values. `tskit.UNKNOWN_TIME` is a specific
    value in this set. `np.nan` is a differing value, but both are NAN.
    See https://en.wikipedia.org/wiki/NaN

    This function only returns true for ``tskit.is_unknown_time(tskit.UNKNOWN_TIME)``
    and will return false for ``tskit.is_unknown_time(np.nan)`` or any other NAN or
    non-NAN value.

    :param time: Value or array to check.
    :type time: Union[float, array-like]
    :return: A single boolean or array of booleans the same shape as ``time``.
    :rtype: Union[bool, numpy.ndarray[bool]]
    """
    return np.asarray(time, dtype=np.float64).view(np.uint64) == np.float64(
        UNKNOWN_TIME
    ).view(np.uint64)


def safe_np_int_cast(int_array, dtype, copy=False):
    """
    A few functions require arrays of certain dtypes (e.g. node indices are np.int32,
    genotypes are np.int8, etc. Standard numpy integer arrays are of (dtype=np.int64),
    so need casting. This function casts but checks bounds to avoid wrap-around
    conversion errors. Strangely, numpy seems not to have this functionality built-in.

    If copy=False, and the original array is a numpy array of exactly the same dtype
    required, simply return the original rather than making a copy (same as the numpy
    .astype(copy=...) function)
    """
    if not isinstance(int_array, np.ndarray):
        int_array = np.array(int_array)
        # Since this is a new numpy array anyway, it's always a copy, so economize by
        # setting copy=False
        copy = False
    if int_array.size == 0:
        return int_array.astype(dtype, copy=copy)  # Allow empty arrays of any type
    try:
        return int_array.astype(dtype, casting="safe", copy=copy)
    except TypeError:
        if int_array.dtype == np.dtype("O"):
            # this occurs e.g. if we're passed a list of lists of different lengths
            raise TypeError("Cannot convert to a rectangular array.")
        bounds = np.iinfo(dtype)
        if np.any(int_array < bounds.min) or np.any(int_array > bounds.max):
            raise OverflowError(f"Cannot convert safely to {dtype} type")
        if int_array.dtype.kind == "i" and np.dtype(dtype).kind == "u":
            # Allow casting from int to unsigned int, since we have checked bounds
            casting = "unsafe"
        else:
            # Raise a TypeError when we try to convert from, e.g., a float.
            casting = "same_kind"
        return int_array.astype(dtype, casting=casting, copy=copy)


#
# Pack/unpack lists of data into flattened numpy arrays.
#


def pack_bytes(data):
    """
    Packs the specified list of bytes into a flattened numpy array of 8 bit integers
    and corresponding offsets. See :ref:`sec_encoding_ragged_columns` for details
    of this encoding.

    :param list[bytes] data: The list of bytes values to encode.
    :return: The tuple (packed, offset) of numpy arrays representing the flattened
        input data and offsets.
    :rtype: numpy.ndarray (dtype=np.int8), numpy.ndarray (dtype=np.uint32)
    """
    n = len(data)
    offsets = np.zeros(n + 1, dtype=np.uint32)
    for j in range(n):
        offsets[j + 1] = offsets[j] + len(data[j])
    column = np.zeros(offsets[-1], dtype=np.int8)
    for j, value in enumerate(data):
        column[offsets[j] : offsets[j + 1]] = bytearray(value)
    return column, offsets


def unpack_bytes(packed, offset):
    """
    Unpacks a list of bytes from the specified numpy arrays of packed byte
    data and corresponding offsets. See :ref:`sec_encoding_ragged_columns` for details
    of this encoding.

    :param numpy.ndarray packed: The flattened array of byte values.
    :param numpy.ndarray offset: The array of offsets into the ``packed`` array.
    :return: The list of bytes values unpacked from the parameter arrays.
    :rtype: list[bytes]
    """
    # This could be done a lot more efficiently...
    ret = []
    for j in range(offset.shape[0] - 1):
        raw = packed[offset[j] : offset[j + 1]].tobytes()
        ret.append(raw)
    return ret


def pack_strings(strings, encoding="utf8"):
    """
    Packs the specified list of strings into a flattened numpy array of 8 bit integers
    and corresponding offsets using the specified text encoding.
    See :ref:`sec_encoding_ragged_columns` for details of this encoding of
    columns of variable length data.

    :param list[str] data: The list of strings to encode.
    :param str encoding: The text encoding to use when converting string data
        to bytes. See the :mod:`codecs` module for information on available
        string encodings.
    :return: The tuple (packed, offset) of numpy arrays representing the flattened
        input data and offsets.
    :rtype: numpy.ndarray (dtype=np.int8), numpy.ndarray (dtype=np.uint32)
    """
    return pack_bytes([bytearray(s.encode(encoding)) for s in strings])


def unpack_strings(packed, offset, encoding="utf8"):
    """
    Unpacks a list of strings from the specified numpy arrays of packed byte
    data and corresponding offsets using the specified text encoding.
    See :ref:`sec_encoding_ragged_columns` for details of this encoding of
    columns of variable length data.

    :param numpy.ndarray packed: The flattened array of byte values.
    :param numpy.ndarray offset: The array of offsets into the ``packed`` array.
    :param str encoding: The text encoding to use when converting string data
        to bytes. See the :mod:`codecs` module for information on available
        string encodings.
    :return: The list of strings unpacked from the parameter arrays.
    :rtype: list[str]
    """
    return [b.decode(encoding) for b in unpack_bytes(packed, offset)]


def pack_arrays(list_of_lists, dtype=np.float64):
    """
    Packs the specified list of numeric lists into a flattened numpy array
    of the specified dtype with corresponding offsets. See
    :ref:`sec_encoding_ragged_columns` for details of this encoding of columns
    of variable length data.

    :param list[list] list_of_lists: The list of numeric lists to encode.
    :param dtype: The dtype for the packed array, defaults to float64
    :return: The tuple (packed, offset) of numpy arrays representing the flattened
        input data and offsets.
    :rtype: numpy.array (dtype=dtype), numpy.array (dtype=np.uint32)
    """
    # TODO must be possible to do this more efficiently with numpy
    n = len(list_of_lists)
    offset = np.zeros(n + 1, dtype=np.uint32)
    for j in range(n):
        offset[j + 1] = offset[j] + len(list_of_lists[j])
    data = np.empty(offset[-1], dtype=dtype)
    for j in range(n):
        data[offset[j] : offset[j + 1]] = list_of_lists[j]
    return data, offset


def unpack_arrays(packed, offset):
    """
    Unpacks a list of arrays from the specified numpy array of packed
    data and its associated offset array. See
    :ref:`sec_encoding_ragged_columns` for details of how columns
    of variable length data are encoded in this way.

    :param numpy.ndarray packed: The flattened array of data.
    :param numpy.ndarray offset: The array of offsets into the ``packed`` array.
    :return: A list of numpy arrays unpacked from the flattened ``packed`` array.
    :rtype: list[numpy.ndarray]
    """
    ret = []
    for j in range(offset.shape[0] - 1):
        ret.append(packed[offset[j] : offset[j + 1]])
    return ret


#
# Interval utilities
#


def intervals_to_np_array(intervals, start, end):
    """
    Converts the specified intervals to a numpy array and checks for
    errors.
    """
    intervals = np.array(intervals, dtype=np.float64)
    # Special case the empty list of intervals
    if len(intervals) == 0:
        intervals = np.zeros((0, 2), dtype=np.float64)
    if len(intervals.shape) != 2:
        raise ValueError("Intervals must be a 2D numpy array")
    if intervals.shape[1] != 2:
        raise ValueError("Intervals array shape must be (N, 2)")
    # TODO do this with numpy operations.
    last_right = start
    for left, right in intervals:
        if left < start or right > end:
            raise ValueError(f"Intervals must be within {start} and {end}")
        if right <= left:
            raise ValueError("Bad interval: right <= left")
        if left < last_right:
            raise ValueError("Intervals must be disjoint.")
        last_right = right
    return intervals


def negate_intervals(intervals, start, end):
    """
    Returns the set of intervals *not* covered by the specified set of
    disjoint intervals in the specified range.
    """
    intervals = intervals_to_np_array(intervals, start, end)
    other_intervals = []
    last_right = start
    for left, right in intervals:
        if left != last_right:
            other_intervals.append((last_right, left))
        last_right = right
    if last_right != end:
        other_intervals.append((last_right, end))
    return np.array(other_intervals)


def naturalsize(value):
    """
    Format a number of bytes like a human readable filesize (e.g. 10 kiB)
    """
    # Taken from https://github.com/jmoiron/humanize
    suffix = ("KiB", "MiB", "GiB", "TiB", "PiB", "EiB", "ZiB", "YiB")
    base = 1024
    format_ = "%.1f"

    bytes_ = float(value)
    abs_bytes = abs(bytes_)

    if abs_bytes == 1:
        return "%d Byte" % bytes_
    elif abs_bytes < base:
        return "%d Bytes" % bytes_

    for i, s in enumerate(suffix):
        unit = base ** (i + 2)
        if abs_bytes < unit:
            return (format_ + " %s") % ((base * bytes_ / unit), s)
    return (format_ + " %s") % ((base * bytes_ / unit), s)


def obj_to_collapsed_html(d, name=None, open_depth=0, max_items=30, max_item_len=100):
    """
    Recursively make an HTML representation of python objects.

    :param str name: Name for this object
    :param int open_depth: By default sub-sections are collapsed. If this number is
        non-zero the first layers up to open_depth will be opened.
    :param int max_items: Maximum number of items to display per collection
    :return: The HTML as a string
    :rtype: str
    """
    opened = "open" if open_depth > 0 else ""
    open_depth -= 1
    name = f"{str(name)}:" if name is not None else ""
    if isinstance(d, dict):
        items = list(d.items())
        more = len(items) - max_items
        display_items = items[:max_items] if more > 0 else items
        inner_html = "".join(
            f"{obj_to_collapsed_html(val, key, open_depth, max_items)}<br/>"
            for key, val in display_items
        )
        if more > 0:
            inner_html += f"... and {more} more"
        return f"""
            <div>
                <span class="tskit-details-label">{name}</span>
                <details {opened}>
                    <summary>dict</summary>
                    {inner_html}
                </details>
            </div>
            """
    elif isinstance(d, list):
        items = d
        more = len(items) - max_items
        display_items = items[:max_items] if more > 0 else items
        inner_html = "".join(
            f"{obj_to_collapsed_html(val, None, open_depth, max_items)}<br/>"
            for val in display_items
        )
        if more > 0:
            inner_html += f"... and {more} more"
        return f"""
            <div>
                <span class="tskit-details-label">{name}</span>
                <details {opened}>
                    <summary>list</summary>
                    {inner_html}
                </details>
            </div>
            """
    else:
        d_str = str(d)
        if len(d_str) > max_item_len:
            d_str = d_str[:max_item_len] + "..."
        d_str = textwrap.fill(d_str, width=30)
        d_str = f"{name} {html.escape(str(d_str))}"
        d_str = d_str.replace("\n", "<br/>")
        return d_str


def truncate_string_end(string, length):
    """
    If a string is longer than "length" then snip out the middle and replace with an
    ellipsis.
    """
    if len(string) <= length:
        return string
    return f"{string[:length - 3]}..."


def render_metadata(md, length=40):
    if md == b"":
        return ""
    return truncate_string_end(str(md), length)


def unicode_table(
    rows, *, title=None, header=None, row_separator=True, column_alignments=None
):
    """
    Convert a table (list of lists) of strings to a unicode table. If a row contains
    the string "__skipped__NNN" then "skipped N rows" is displayed.

    :param list[list[str]] rows: List of rows, each of which is a list of strings for
        each cell. Each row must have the same number of cells.
    :param str title: If specified the first output row will be a single cell
        containing this string, left-justified. [optional]
    :param list[str] header: Specifies a row above the main rows which will be in double
        lined borders and left justified. Must be same length as each row. [optional]
    :param boolean row_separator: If True add lines between each row. [Default: True]
    :param column_alignments str: A string of the same length as the number of cells in
        a row (i.e. columns) where each character specifies an alignment such as ``<``,
        ``>`` or ``^`` as used in Python's string formatting mini-language. If ``None``,
        set the first column to be left justified and the remaining columns to be right
        justified [Default: ``None``]
    :return: The table as a string
    :rtype: str
    """
    if header is not None:
        all_rows = [header] + rows
    else:
        all_rows = rows
    widths = [
        max(len(row[i_col]) for row in all_rows) for i_col in range(len(all_rows[0]))
    ]
    if column_alignments is None:
        column_alignments = "<" + ">" * (len(widths) - 1)
    out = []
    inner_width = sum(widths) + len(header or rows[0]) - 1
    if title is not None:
        out += [
            f"╔{'═' * inner_width}╗\n" f"║{title.ljust(inner_width)}║\n",
            f"╠{'╤'.join('═' * w for w in widths)}╣\n",
        ]
    if header is not None:
        out += [
            f"╔{'╤'.join('═' * w for w in widths)}╗\n",
            f"║{'│'.join(cell.ljust(w) for cell, w in zip(header, widths))}║\n",
            f"╠{'╪'.join('═' * w for w in widths)}╣\n",
        ]
    last_skipped = False
    for i, row in enumerate(rows):
        if "__skipped__" in row:
            msg = f"{row[11:]} rows skipped (tskit.set_print_options)"[
                :inner_width
            ].center(inner_width)
            row_str = f"║{msg}║\n"
            if row_separator:
                out += [
                    f"╟{'┴'.join('─' * w for w in widths)}╢\n" + row_str,
                    f"╟{'┬'.join('─' * w for w in widths)}╢\n",
                ]
            else:
                out.append(row_str)
            last_skipped = True
        else:
            if i != 0 and not last_skipped and row_separator:
                out.append(f"╟{'┼'.join('─' * w for w in widths)}╢\n")

            out.append(
                "║"
                + "│".join(
                    f"{r:{a}{w}}" for r, w, a in zip(row, widths, column_alignments)
                )
                + "║\n"
            )
            last_skipped = False

    out.append(f"╚{'╧'.join('═' * w for w in widths)}╝\n")
    return "".join(out)


def html_table(rows, *, header):
    headers = "".join(f"<th>{h}</th>" for h in header)
    rows = (
        (
            f'<td style="text-align: center;" colspan="{len(headers)}"><em>{row[11:]}'
            f" rows skipped (tskit.set_print_options)</em></td>"
            if "__skipped__" in row
            else "".join(f"<td>{cell}</td>" for cell in row)
        )
        for row in rows
    )
    rows = "".join(f"<tr>{row}</tr>\n" for row in rows)
    return f"""
        <div>
            <style scoped="">
                .tskit-table tbody tr th:only-of-type {{vertical-align: middle;}}
                .tskit-table tbody tr th {{vertical-align: top;}}
                .tskit-table tbody td {{text-align: right;padding: 0.5em 0.5em;}}
                .tskit-table tbody th {{padding: 0.5em 0.5em;}}
            </style>
            <table border="1" class="tskit-table">
                <thead>
                    <tr>
                        {headers}
                    </tr>
                </thead>
                <tbody>
                    {rows}
                </tbody>
            </table>
        </div>
    """


def tree_sequence_html(ts):
    table_rows = "".join(
        f"""
            <tr>
                <td>{name.capitalize()}</td>
                <td>{table.num_rows:,}</td>
                <td>{naturalsize(table.nbytes)}</td>
                <td style="text-align: center;">
                    {'✅' if hasattr(table, "metadata") and len(table.metadata) > 0
                     else ''}
                </td>
            </tr>
        """
        for name, table in ts.tables.table_name_map.items()
    )

    provenance_rows = ""
    provenances = list(ts.provenances())
    # Detail the most recent 10 provenances, and collapse the rest
    display_provenances = provenances[-10:]
    extra_provenances = provenances[0:-10]
    for prov in reversed(display_provenances):
        try:
            timestamp = datetime.datetime.fromisoformat(prov.timestamp).strftime(
                "%d %B, %Y at %I:%M:%S %p"
            )
            record = json.loads(prov.record)
            software_name = record.get("software", {}).get("name", "Unknown")
            software_version = record.get("software", {}).get("version", "Unknown")
            command = record.get("parameters", {}).get("command", "Unknown")
            details = obj_to_collapsed_html(record, None, 0)
            provenance_rows += f"""
                <tr>
                    <td>{timestamp}</td>
                    <td>{software_name}</td>
                    <td>{software_version}</td>
                    <td>{command}</td>
                    <td>
                        <details>
                            <summary>Details</summary>
                            {details}
                        </details>
                    </td>
                </tr>
            """
        except Exception as e:
            provenance_rows += (
                f"""Could not parse provenance record: """
                f"""{e.__class__.__name__} {str(e)}"""
            )
    if len(extra_provenances) > 0:
        provenance_rows += f"""
            <tr>
                <td colspan="5"><i>... {len(extra_provenances)} more</i></td>
            </tr>
        """
    md = (
        obj_to_collapsed_html(ts.metadata, None, 1)
        if len(ts.tables.metadata_bytes) > 0
        else "No Metadata"
    )
    return f"""
        <div>
            <style>
                .tskit-table thead tr th {{text-align: left;padding: 0.5em 0.5em;}}
                .tskit-table tbody tr td {{padding: 0.5em 0.5em;}}
                .tskit-table tbody tr td:first-of-type {{text-align: left;}}
                .tskit-details-label {{vertical-align: top; padding-right:5px;}}
                .tskit-table-set {{display: inline-flex;flex-wrap: wrap;margin: -12px 0 0 -12px;width: calc(100% + 12px);}}
                .tskit-table-set-table {{margin: 12px 0 0 12px;}}
                details {{display: inline-block;}}
                summary {{cursor: pointer; outline: 0; display: list-item;}}
            </style>
            <div class="tskit-table-set">
                <div class="tskit-table-set-table">
                    <table class="tskit-table">
                        <thead>
                            <tr>
                                <th style="padding:0;line-height:21px;">
                                    <img style="height: 32px;display: inline-block;padding: 3px 5px 3px 0;" src="https://raw.githubusercontent.com/tskit-dev/administrative/main/tskit_logo.svg"/>
                                    <a target="_blank" href="https://tskit.dev/tskit/docs/latest/python-api.html#the-treesequence-class"> Tree Sequence </a>
                                </th>
                            </tr>
                        </thead>
                        <tbody>
                            <tr><td>Trees</td><td>{ts.num_trees:,}</td></tr>
                            <tr><td>Sequence Length</td><td>{ts.sequence_length:,}</td></tr>
                            <tr><td>Time Units</td><td>{ts.time_units}</td></tr>
                            <tr><td>Sample Nodes</td><td>{ts.num_samples:,}</td></tr>
                            <tr><td>Total Size</td><td>{naturalsize(ts.nbytes)}</td></tr>
                            <tr>
                                <td>Metadata</td><td style="text-align: left;">{md}</td>
                            </tr>
                        </tbody>
                    </table>
                </div>
                <div class="tskit-table-set-table">
                    <table class="tskit-table">
                        <thead>
                            <tr>
                                <th style="line-height:21px;">Table</th>
                                <th>Rows</th>
                                <th>Size</th>
                                <th>Has Metadata</th>
                            </tr>
                        </thead>
                        <tbody>
                            {table_rows}
                        </tbody>
                    </table>
                </div>
                <div class="tskit-table-set-table">
                    <table class="tskit-table">
                        <thead>
                            <tr>
                                <th>Provenance Timestamp</th>
                                <th>Software Name</th>
                                <th>Version</th>
                                <th>Command</th>
                                <th>Full record</th>
                            </tr>
                        </thead>
                        <tbody>
                            {provenance_rows}
                        </tbody>
                    </table>
                </div>
            </div>
        </div>
    """  # noqa: B950


def tree_html(tree):
    return f"""
            <div>
              <style>
                .tskit-table thead tr th {{text-align: left;padding: 0.5em 0.5em;}}
                .tskit-table tbody tr td {{padding: 0.5em 0.5em;}}
                .tskit-table tbody tr td:first-of-type {{text-align: left;}}
                .tskit-details-label {{vertical-align: top; padding-right:5px;}}
                .tskit-table-set {{display: inline-flex;flex-wrap: wrap;margin: -12px 0 0 -12px;width: calc(100% + 12px);}}
                .tskit-table-set-table {{margin: 12px 0 0 12px;}}
                details {{display: inline-block;}}
                summary {{cursor: pointer; outline: 0; display: list-item;}}
              </style>
              <div class="tskit-table-set">
                <div class="tskit-table-set-table">
                  <table class="tskit-table">
                    <thead>
                      <tr>
                        <th style="padding:0;line-height:21px;">
                          <img style="height: 32px;display: inline-block;padding: 3px 5px 3px 0;" src="https://raw.githubusercontent.com/tskit-dev/administrative/main/tskit_logo.svg"/>
                          <a target="_blank" href="https://tskit.dev/tskit/docs/latest/python-api.html#the-tree-class"> Tree </a>
                        </th>
                      </tr>
                    </thead>
                    <tbody>
                      <tr><td>Index</td><td>{tree.index:,}</td></tr>
                      <tr><td>Interval</td><td>{tree.interval.left:,.8g}-{tree.interval.right:,.8g} ({tree.span:,.8g})</td></tr>
                      <tr><td>Roots</td><td>{tree.num_roots:,}</td></tr>
                      <tr><td>Nodes</td><td>{len(tree.preorder()):,}</td></tr>
                      <tr><td>Sites</td><td>{tree.num_sites:,}</td></tr>
                      <tr><td>Mutations</td><td>{tree.num_mutations:,}</td></tr>
                      <tr><td>Total Branch Length</td><td>{tree.total_branch_length:,.8g}</td></tr>
                    </tbody>
                  </table>
                </div>
              </div>
            </div>
            """  # noqa: B950


def variant_html(variant):
    class_type = "Variant"

    url_tskit_logo = (
        "https://raw.githubusercontent.com/tskit-dev/administrative/main/tskit_logo.svg"
    )
    url_variant_class_doc = (
        "https://tskit.dev/tskit/docs/latest/python-api.html#the-variant-class"
    )

    html_body_head = f"""
        <div>
            <style>
                .tskit-table thead tr th {{text-align: left;padding: 0.5em 0.5em;}}
                .tskit-table tbody tr td {{padding: 0.5em 0.5em;}}
                .tskit-table tbody tr td:first-of-type {{text-align: left;}}
                .tskit-details-label {{vertical-align: top; padding-right:5px;}}
                .tskit-table-set {{display: inline-flex;flex-wrap: wrap;margin: -12px 0 0 -12px;width: calc(100% + 12px);}}
                .tskit-table-set-table {{margin: 12px 0 0 12px;}}
                details {{display: inline-block;}}
                summary {{cursor: pointer; outline: 0; display: list-item;}}
            </style>
            <div class="tskit-table-set">
                <div class="tskit-table-set-table">
                <table class="tskit-table">
                    <thead>
                    <tr>
                        <th style="padding:0;line-height:21px;">
                        <img style="height: 32px;display: inline-block;padding: 3px 5px 3px 0;" src="{url_tskit_logo}"/>
                        <a target="_blank" href="{url_variant_class_doc}"> {class_type} </a>
                        </th>
                    </tr>
                    </thead>
                    <tbody>
        """  # noqa: B950

    html_body_tail = """
                    </tbody>
            </table>
            </div>
        </div>
        </div>
        """

    try:
        variant.site

        site_id = variant.site.id
        site_position = variant.site.position
        num_samples = len(variant.samples)
        num_alleles = variant.num_alleles
        has_missing_data = str(variant.has_missing_data)
        isolated_as_missing = str(bool(variant.isolated_as_missing))

        counts = variant.counts()
        freqs = variant.frequencies()

        return (
            html_body_head
            + f"""
                <tr><td>Site Id</td><td>{site_id:,}</td></tr>
                <tr><td>Site Position</td><td>{site_position:,.8g}</td></tr>
                <tr><td>Number of Samples</td><td>{num_samples:,}</td></tr>
                <tr><td>Number of Alleles</td><td>{num_alleles:,}</td></tr>
            """
            + "\n".join(
                [
                    f"""<tr><td>Samples with Allele {'missing' if k is None
                                                     else "'" + k + "'"}</td><td>"""
                    + f"{counts[k]:,}"
                    + " "
                    + f"({freqs[k] * 100:,.2g}%)"
                    + "</td></tr>"
                    for k in variant.alleles
                ]
            )
            + f"""
                <tr><td>Has Missing Data</td><td>{has_missing_data}</td></tr>
                <tr><td>Isolated as Missing</td><td>{isolated_as_missing}</td></tr>
            """
            + html_body_tail
        )
    except ValueError as err:
        return (
            html_body_head
            + f"""
                        <tr><td>Error</td><td>{str(err)}</td></tr>
            """
            + html_body_tail
        )


def convert_file_like_to_open_file(file_like, mode):
    # Get ourselves a local version of the file. The semantics here are complex
    # because need to support a range of inputs and the free behaviour is
    # slightly different on each.
    _file = None
    local_file = True
    try:
        # First, see if we can interpret the argument as a pathlike object.
        path = os.fspath(file_like)
        _file = open(path, mode)
    except TypeError:
        pass
    if _file is None:
        # Now we try to open file. If it's not a pathlike object, it could be
        # an integer fd or object with a fileno method. In this case we
        # must make sure that close is **not** called on the fd.
        try:
            _file = open(file_like, mode, closefd=False, buffering=0)
        except TypeError:
            pass
    if _file is None:
        # Assume that this is a file **but** we haven't opened it, so we must
        # not close it.
        if mode == "wb" and not hasattr(file_like, "write"):
            raise TypeError("file object must have a write method")
        _file = file_like
        local_file = False
    return _file, local_file


def set_print_options(*, max_lines=40):
    """
    Set the options for printing to strings and HTML

    :param integer max_lines: The maximum number of lines to print from a table, beyond
    this number the middle of the table will be skipped.
    """
    tskit._print_options = {"max_lines": max_lines}


def truncate_rows(num_rows, limit=None):
    """
    Return a list of indexes into a set of rows, but if a ``limit`` is set, truncate the
    number of rows and place a single ``-1`` entry, instead of the intermediate indexes
    """
    if limit is None or num_rows <= limit:
        return range(num_rows)
    return itertools.chain(
        range(limit // 2),
        [-1],
        range(num_rows - (limit - (limit // 2)), num_rows),
    )


def random_nucleotides(length: numbers.Number, *, seed: Union[int, None] = None) -> str:
    """
    Returns a random string of nucleotides of the specified length. Characters
    are drawn uniformly from the alphabet "ACTG".

    :param int length: The length of the random sequence.
    :return: A string of the specified length consisting of random nucleotide
       characters.
    :rtype: str
    """
    if int(length) != length:
        raise ValueError("length must be an integer")
    rng = np.random.RandomState(seed)
    encoded_nucleotides = np.array(list(map(ord, "ACTG")), dtype=np.int8)
    a = rng.choice(encoded_nucleotides, size=int(length))
    return a.tobytes().decode("ascii")


def raise_known_file_format_errors(open_file, existing_exception):
    """
    Sniffs the file for pk-zip or hdf header bytes, then raises an exception
    if these are detected, if not raises the existing exception.
    """
    # Check for HDF5 header bytes
    try:
        open_file.seek(0)
        header = open_file.read(4)
    except io.UnsupportedOperation:
        # If we can't seek, we can't sniff the file.
        raise existing_exception
    if header == b"\x89HDF":
        raise tskit.FileFormatError(
            "The specified file appears to be in HDF5 format. This file "
            "may have been generated by msprime < 0.6.0 (June 2018) which "
            "can no longer be read directly. Please convert to the new "
            "kastore format using the ``tskit upgrade`` command."
        ) from existing_exception
    if header[:2] == b"\x50\x4b":
        raise tskit.FileFormatError(
            "The specified file appears to be in zip format, so may be a compressed "
            "tree sequence. Try using the tszip module to decompress this file before "
            "loading. `pip install tszip; tsunzip <filename>` or use "
            "`tszip.decompress` in Python code."
        ) from existing_exception
    raise existing_exception


--- ../../tskit/python/tskit/metadata.py ---


"""
Classes for metadata decoding, encoding and validation
"""
from __future__ import annotations

import abc
import builtins
import collections
import copy
import functools
import json
import pprint
import struct
import types
from itertools import islice
from typing import Any
from typing import Mapping

import jsonschema

import tskit
import tskit.exceptions as exceptions

__builtins__object__setattr__ = builtins.object.__setattr__


def replace_root_refs(obj):
    if type(obj) is list:
        return [replace_root_refs(j) for j in obj]
    elif type(obj) is dict:
        ret = {k: replace_root_refs(v) for k, v in obj.items()}
        if ret.get("$ref") == "#":
            ret["$ref"] = "#/definitions/root"
        return ret
    else:
        return obj


# Our schema is the Draft7Validator schema with added codec information.
TSKITMetadataSchemaValidator = jsonschema.validators.extend(
    jsonschema.validators.Draft7Validator
)
deref_meta_schema: Mapping[str, Any] = copy.deepcopy(
    TSKITMetadataSchemaValidator.META_SCHEMA
)
# We need a top-level only required property so we need to rewrite any reference
# to the top-level schema to a copy in a definition.
deref_meta_schema = replace_root_refs(deref_meta_schema)
deref_meta_schema["definitions"]["root"] = copy.deepcopy(deref_meta_schema)
deref_meta_schema["codec"] = {"type": "string"}
deref_meta_schema["required"] = ["codec"]
# For interoperability reasons, force the top-level to be an object or union
# of object and null
deref_meta_schema["properties"]["type"] = {"enum": ["object", ["object", "null"]]}
# Change the schema URL to avoid jsonschema's cache
deref_meta_schema["$schema"] = "http://json-schema.org/draft-o=07/schema#tskit"
TSKITMetadataSchemaValidator.META_SCHEMA = deref_meta_schema


class AbstractMetadataCodec(metaclass=abc.ABCMeta):
    """
    Superclass of all MetadataCodecs.
    """

    def __init__(self, schema: Mapping[str, Any]) -> None:
        raise NotImplementedError  # pragma: no cover

    @classmethod
    def modify_schema(self, schema: Mapping) -> Mapping:
        return schema

    @classmethod
    def is_schema_trivial(self, schema: Mapping) -> bool:
        return False

    @abc.abstractmethod
    def encode(self, obj: Any) -> bytes:
        raise NotImplementedError  # pragma: no cover

    @abc.abstractmethod
    def decode(self, encoded: bytes) -> Any:
        raise NotImplementedError  # pragma: no cover


codec_registry = {}


def register_metadata_codec(
    codec_cls: type[AbstractMetadataCodec], codec_id: str
) -> None:
    """
    Register a metadata codec class.
    This function maintains a mapping from metadata codec identifiers used in schemas
    to codec classes. When a codec class is registered, it will replace any class
    previously registered under the same codec identifier, if present.

    :param str codec_id: String to use to refer to the codec in the schema.
    """
    codec_registry[codec_id] = codec_cls


class JSONCodec(AbstractMetadataCodec):
    def default_validator(validator, types, instance, schema):
        # For json codec defaults must be at the top level
        if validator.is_type(instance, "object"):
            for v in instance.get("properties", {}).values():
                for v2 in v.get("properties", {}).values():
                    if "default" in v2:
                        yield jsonschema.ValidationError(
                            "Defaults can only be specified at the top level"
                            " for JSON codec"
                        )

    schema_validator = jsonschema.validators.extend(
        TSKITMetadataSchemaValidator, {"default": default_validator}
    )

    @classmethod
    def is_schema_trivial(self, schema: Mapping) -> bool:
        return len(schema.get("properties", {})) == 0

    def __init__(self, schema: Mapping[str, Any]) -> None:
        try:
            self.schema_validator.check_schema(schema)
        except jsonschema.exceptions.SchemaError as ve:
            raise exceptions.MetadataSchemaValidationError(str(ve)) from ve

        # Find default values to fill in on decode, top level only
        self.defaults = {
            key: prop["default"]
            for key, prop in schema.get("properties", {}).items()
            if "default" in prop
        }

    def encode(self, obj: Any) -> bytes:
        try:
            return tskit.canonical_json(obj).encode()
        except TypeError as e:
            raise exceptions.MetadataEncodingError(
                f"Could not encode metadata of type {str(e).split()[3]}"
            )

    def decode(self, encoded: bytes) -> Any:
        if len(encoded) == 0:
            result = {}
        else:
            result = json.loads(encoded.decode())

        # Assign default values
        if isinstance(result, dict):
            return dict(self.defaults, **result)
        else:
            return result


register_metadata_codec(JSONCodec, "json")


class NOOPCodec(AbstractMetadataCodec):
    def __init__(self, schema: Mapping[str, Any]) -> None:
        pass

    def encode(self, data: bytes) -> bytes:
        return data

    def decode(self, data: bytes) -> bytes:
        return data


def binary_format_validator(validator, types, instance, schema):
    # We're hooking into jsonschemas validation code here, which works by creating
    # generators of exceptions, hence the yielding

    # Make sure the normal type validation gets done
    try:
        yield from jsonschema._validators.type(validator, types, instance, schema)
    except AttributeError:
        # Needed since jsonschema==4.19.1
        yield from jsonschema._keywords.type(validator, types, instance, schema)

    # Non-composite types must have a binaryFormat
    if validator.is_type(instance, "object"):
        for v in instance.values():
            if (
                isinstance(v, dict)
                and v.get("type")
                not in (None, "object", "array", "null", ["object", "null"])
                and "binaryFormat" not in v
            ):
                yield jsonschema.ValidationError(
                    f"{v['type']} type must have binaryFormat set"
                )
    # null type must be padding
    if (
        validator.is_type(instance, "object")
        and "null" in instance
        and instance["null"].get("type") == "null"
        and "binaryFormat" in instance["null"]
        and instance["null"]["binaryFormat"][-1] != "x"
    ):
        yield jsonschema.ValidationError(
            'null type binaryFormat must be padding ("x") if set'
        )


def required_validator(validator, required, instance, schema):
    # Do the normal validation
    try:
        yield from jsonschema._validators.required(
            validator, required, instance, schema
        )
    except AttributeError:
        # Needed since jsonschema==4.19.1
        yield from jsonschema._keywords.required(validator, required, instance, schema)

    # For struct codec if a property is not required, then it must have a default
    for prop, sub_schema in instance["properties"].items():
        if prop not in instance["required"] and "default" not in sub_schema:
            yield jsonschema.ValidationError(
                f"Optional property '{prop}' must have" f" a default value"
            )


StructCodecSchemaValidator = jsonschema.validators.extend(
    TSKITMetadataSchemaValidator,
    {"type": binary_format_validator, "required": required_validator},
)
struct_meta_schema: Mapping[str, Any] = copy.deepcopy(
    StructCodecSchemaValidator.META_SCHEMA
)
# No union types
struct_meta_schema["definitions"]["root"]["properties"]["type"] = {
    "$ref": "#/definitions/simpleTypes"
}
# No hetrogeneous arrays
struct_meta_schema["properties"]["items"] = {"$ref": "#/definitions/root"}
struct_meta_schema["definitions"]["root"]["properties"]["items"] = struct_meta_schema[
    "properties"
]["items"]
# binaryFormat matches regex
struct_meta_schema["properties"]["binaryFormat"] = {
    "type": "string",
    "pattern": r"^([cbB\?hHiIlLqQfd]|\d*[spx])$",
}
struct_meta_schema["definitions"]["root"]["properties"]["binaryFormat"] = (
    struct_meta_schema["properties"]["binaryFormat"]
)
# arrayLengthFormat matches regex and has default
struct_meta_schema["properties"]["arrayLengthFormat"] = {
    "type": "string",
    "pattern": r"^[BHILQ]$",
    "default": "L",
}
struct_meta_schema["definitions"]["root"]["properties"]["arrayLengthFormat"] = (
    struct_meta_schema["properties"]["arrayLengthFormat"]
)
# index is numeric
struct_meta_schema["properties"]["index"] = {"type": "number"}
struct_meta_schema["definitions"]["root"]["properties"]["index"] = struct_meta_schema[
    "properties"
]["index"]
# stringEncoding is string and has default
struct_meta_schema["properties"]["stringEncoding"] = {
    "type": "string",
    "default": "utf-8",
}
struct_meta_schema["definitions"]["root"]["properties"]["stringEncoding"] = (
    struct_meta_schema["properties"]["stringEncoding"]
)
# nullTerminated is a boolean
struct_meta_schema["properties"]["nullTerminated"] = {"type": "boolean"}
struct_meta_schema["definitions"]["root"]["properties"]["nullTerminated"] = (
    struct_meta_schema["properties"]["nullTerminated"]
)
# noLengthEncodingExhaustBuffer is a boolean
struct_meta_schema["properties"]["noLengthEncodingExhaustBuffer"] = {"type": "boolean"}
struct_meta_schema["definitions"]["root"]["properties"][
    "noLengthEncodingExhaustBuffer"
] = struct_meta_schema["properties"]["noLengthEncodingExhaustBuffer"]
StructCodecSchemaValidator.META_SCHEMA = struct_meta_schema


class StructCodec(AbstractMetadataCodec):
    """
    Codec that encodes data using struct. Note that this codec has extra restrictions
    Namely that object keys must be fixed (all present and no extra); each entry should
    have a binaryFormat; that arrays are homogeneous and that types are not unions.
    """

    @classmethod
    def order_by_index(cls, obj, do_sort=False):
        """
        Take a schema and recursively convert any dict that is under the key
        name ``properties`` to an OrderedDict.
        """
        if isinstance(obj, collections.abc.Mapping):
            items = obj.items()
            if do_sort:
                # Python sort is stable so we can do the sorts in reverse priority
                items = sorted(items, key=lambda k_v: k_v[0])
                items = sorted(items, key=lambda k_v: k_v[1].get("index", 0))
            items = ((k, cls.order_by_index(v, k == "properties")) for k, v in items)
            if do_sort:
                return collections.OrderedDict(items)
            else:
                return dict(items)
        elif isinstance(obj, list) or isinstance(obj, tuple):
            return [cls.order_by_index(v, False) for v in obj]
        else:
            return obj

    @classmethod
    def make_decode(cls, sub_schema):
        """
        Create a function that can decode objects of this schema
        """
        if set(sub_schema["type"]) == {"object", "null"}:
            return StructCodec.make_object_or_null_decode(sub_schema)
        else:
            return {
                "array": StructCodec.make_array_decode,
                "object": StructCodec.make_object_decode,
                "string": StructCodec.make_string_decode,
                "null": StructCodec.make_null_decode,
                "number": StructCodec.make_numeric_decode,
                "integer": StructCodec.make_numeric_decode,
                "boolean": StructCodec.make_numeric_decode,
            }[sub_schema["type"]](sub_schema)

    @classmethod
    def make_array_decode(cls, sub_schema):
        element_decoder = StructCodec.make_decode(sub_schema["items"])
        array_length_f = "<" + sub_schema.get("arrayLengthFormat", "L")
        array_length_size = struct.calcsize(array_length_f)
        exhaust_buffer = sub_schema.get("noLengthEncodingExhaustBuffer", False)

        def array_decode(buffer):
            array_length = struct.unpack(
                array_length_f, bytes(islice(buffer, array_length_size))
            )[0]
            return [element_decoder(buffer) for _ in range(array_length)]

        def array_decode_exhaust(buffer):
            ret = []
            while True:
                try:
                    ret.append(element_decoder(buffer))
                except struct.error as e:
                    if "unpack requires a buffer" in str(e):
                        break
                    else:
                        raise e
            return ret

        if exhaust_buffer:
            return array_decode_exhaust
        else:
            return array_decode

    @classmethod
    def make_object_decode(cls, sub_schema):
        sub_decoders = {
            key: StructCodec.make_decode(prop)
            for key, prop in sub_schema["properties"].items()
        }
        return lambda buffer: {
            key: sub_decoder(buffer) for key, sub_decoder in sub_decoders.items()
        }

    @classmethod
    def make_object_or_null_decode(cls, sub_schema):
        sub_decoders = {
            key: StructCodec.make_decode(prop)
            for key, prop in sub_schema["properties"].items()
        }

        def decode_object_or_null(buffer):
            # We have to check the buffer length for null, as the islices in
            # sub-decoders won't raise StopIteration
            buffer = list(buffer)
            if len(buffer) == 0:
                return None
            else:
                buffer = iter(buffer)
                return {
                    key: sub_decoder(buffer)
                    for key, sub_decoder in sub_decoders.items()
                }

        return decode_object_or_null

    @classmethod
    def make_string_decode(cls, sub_schema):
        f = "<" + sub_schema["binaryFormat"]
        size = struct.calcsize(f)
        encoding = sub_schema.get("stringEncoding", "utf-8")
        null_terminated = sub_schema.get("nullTerminated", False)
        if not null_terminated:
            return lambda buffer: struct.unpack(f, bytes(islice(buffer, size)))[
                0
            ].decode(encoding)
        else:

            def decode_string(buffer):
                s = struct.unpack(f, bytes(islice(buffer, size)))[0].decode(encoding)
                i = s.find("\x00")
                if i == -1:
                    return s
                return s[:i]

            return decode_string

    @classmethod
    def make_null_decode(cls, sub_schema):
        if sub_schema.get("binaryFormat") is not None:
            f = sub_schema["binaryFormat"]
            size = struct.calcsize(f)

            def padding_decode(buffer):
                struct.unpack(f, bytes(islice(buffer, size)))

            return padding_decode
        else:
            return lambda _: None

    @classmethod
    def make_numeric_decode(cls, sub_schema):
        f = "<" + sub_schema["binaryFormat"]
        size = struct.calcsize(f)
        return lambda buffer: struct.unpack(f, bytes(islice(buffer, size)))[0]

    @classmethod
    def make_encode(cls, sub_schema):
        """
        Create a function that can encode objects of this schema
        """
        if set(sub_schema["type"]) == {"object", "null"}:
            return StructCodec.make_object_or_null_encode(sub_schema)
        else:
            return {
                "array": StructCodec.make_array_encode,
                "object": StructCodec.make_object_encode,
                "string": StructCodec.make_string_encode,
                "null": StructCodec.make_null_encode,
                "number": StructCodec.make_numeric_encode,
                "integer": StructCodec.make_numeric_encode,
                "boolean": StructCodec.make_numeric_encode,
            }[sub_schema["type"]](sub_schema)

    @classmethod
    def make_array_encode(cls, sub_schema):
        array_length_f = "<" + sub_schema.get("arrayLengthFormat", "L")
        element_encoder = StructCodec.make_encode(sub_schema["items"])
        exhaust_buffer = sub_schema.get("noLengthEncodingExhaustBuffer", False)
        if exhaust_buffer:
            return lambda array: b"".join(element_encoder(ele) for ele in array)
        else:

            def array_encode_with_length(array):
                try:
                    packed_length = struct.pack(array_length_f, len(array))
                except struct.error:
                    raise ValueError(
                        "Couldn't pack array size - it is likely too long"
                        " for the specified arrayLengthFormat"
                    )
                return packed_length + b"".join(element_encoder(ele) for ele in array)

            return array_encode_with_length

    @classmethod
    def make_object_encode(cls, sub_schema):
        sub_encoders = {
            key: StructCodec.make_encode(prop)
            for key, prop in sub_schema["properties"].items()
        }
        defaults = {
            key: prop["default"]
            for key, prop in sub_schema["properties"].items()
            if "default" in prop
        }

        def object_encode(obj):
            values = []
            for key, sub_encoder in sub_encoders.items():
                try:
                    values.append(sub_encoder(obj[key]))
                except KeyError:
                    values.append(sub_encoder(defaults[key]))
            return b"".join(values)

        return object_encode

    @classmethod
    def make_object_or_null_encode(cls, sub_schema):
        sub_encoders = {
            key: StructCodec.make_encode(prop)
            for key, prop in sub_schema["properties"].items()
        }
        defaults = {
            key: prop["default"]
            for key, prop in sub_schema["properties"].items()
            if "default" in prop
        }

        def object_encode(obj):
            values = []
            if obj is not None:
                for key, sub_encoder in sub_encoders.items():
                    try:
                        values.append(sub_encoder(obj[key]))
                    except KeyError:
                        values.append(sub_encoder(defaults[key]))
            return b"".join(values)

        return object_encode

    @classmethod
    def make_string_encode(cls, sub_schema):
        encoding = sub_schema.get("stringEncoding", "utf-8")
        return lambda string: struct.pack(
            "<" + sub_schema["binaryFormat"], string.encode(encoding)
        )

    @classmethod
    def make_null_encode(cls, sub_schema):
        return lambda _: struct.pack(sub_schema.get("binaryFormat", "0x"))

    @classmethod
    def make_numeric_encode(cls, sub_schema):
        return struct.Struct("<" + sub_schema["binaryFormat"]).pack

    @classmethod
    def modify_schema(cls, schema: Mapping) -> Mapping:
        # This codec requires that additional properties are
        # not allowed. Rather than get schema authors to repeat that everywhere
        # we add it here, sadly we can't do this in the metaschema as "default" isn't
        # used by the validator.
        def enforce_fixed_properties(obj):
            if type(obj) is list:
                return [enforce_fixed_properties(j) for j in obj]
            elif type(obj) is dict:
                ret = {k: enforce_fixed_properties(v) for k, v in obj.items()}
                if "object" in ret.get("type", []):
                    if ret.get("additional_properties"):
                        raise ValueError(
                            "Struct codec does not support additional_properties"
                        )
                    # To prevent authors having to list required properties the default
                    # is that all without a default are required.
                    if "required" not in ret:
                        ret["required"] = [
                            prop
                            for prop, sub_schema in ret.get("properties", {}).items()
                            if "default" not in sub_schema
                        ]
                    ret["additionalProperties"] = False
                return ret
            else:
                return obj

        schema = enforce_fixed_properties(schema)

        # We also give the schema an explicit ordering
        return StructCodec.order_by_index(schema)

    def __init__(self, schema: Mapping[str, Any]) -> None:
        try:
            StructCodecSchemaValidator.check_schema(schema)
        except jsonschema.exceptions.SchemaError as ve:
            raise exceptions.MetadataSchemaValidationError(str(ve)) from ve

        self.encode = StructCodec.make_encode(schema)
        decoder = StructCodec.make_decode(schema)
        self.decode = lambda buffer: decoder(iter(buffer))

    def encode(self, obj: Any) -> bytes:
        # Set by __init__
        pass  # pragma: nocover

    def decode(self, encoded: bytes) -> Any:
        # Set by __init__
        pass  # pragma: nocover


register_metadata_codec(StructCodec, "struct")


def validate_bytes(data: bytes | None) -> None:
    if data is not None and not isinstance(data, bytes):
        raise TypeError(
            f"If no encoding is set metadata should be bytes, found {type(data)}"
        )


class MetadataSchema:
    """
    Class for validating, encoding and decoding metadata.

    :param dict schema: A dict containing a valid JSONSchema object.
    """

    def __init__(self, schema: Mapping[str, Any] | None) -> None:
        self._schema = schema
        self._bypass_validation = False

        if schema is None:
            self._string = ""
            self._validate_row = validate_bytes
            self.encode_row = NOOPCodec({}).encode
            self.decode_row = NOOPCodec({}).decode
            self.empty_value = b""
        else:
            try:
                TSKITMetadataSchemaValidator.check_schema(schema)
            except jsonschema.exceptions.SchemaError as ve:
                raise exceptions.MetadataSchemaValidationError(str(ve)) from ve
            try:
                codec_cls = codec_registry[schema["codec"]]
            except KeyError:
                raise exceptions.MetadataSchemaValidationError(
                    f"Unrecognised metadata codec '{schema['codec']}'. "
                    f"Valid options are {str(list(codec_registry.keys()))}."
                )
            # Codecs can modify the schema, for example to set defaults as the validator
            # does not.
            schema = codec_cls.modify_schema(schema)
            codec_instance = codec_cls(schema)
            self._string = tskit.canonical_json(schema)
            self._validate_row = TSKITMetadataSchemaValidator(schema).validate
            self._bypass_validation = codec_cls.is_schema_trivial(schema)
            self.encode_row = codec_instance.encode
            self.decode_row = codec_instance.decode

            # If None is allowed by the schema as the top-level type, it gets used even
            # in the presence of default and required values.
            if "type" in schema and "null" in schema["type"]:
                self.empty_value = None
            else:
                self.empty_value = {}

    def __repr__(self) -> str:
        return self._string

    def __str__(self) -> str:
        if isinstance(self._schema, collections.OrderedDict):
            s = pprint.pformat(dict(self._schema))
        else:
            s = pprint.pformat(self._schema)
        if "\n" in s:
            return f"tskit.MetadataSchema(\n{s}\n)"
        else:
            return f"tskit.MetadataSchema({s})"

    def __eq__(self, other) -> bool:
        return self._string == other._string

    @property
    def schema(self) -> Mapping[str, Any] | None:
        # Return a copy to avoid unintentional mutation
        return copy.deepcopy(self._schema)

    def asdict(self) -> Mapping[str, Any] | None:
        """
        Returns a dict representation of this schema. One possible use of this is to
        modify this dict and then pass it to the ``MetadataSchema`` constructor to create
        a similar schema.
        """
        return self.schema

    def validate_and_encode_row(self, row: Any) -> bytes:
        """
        Validate a row (dict) of metadata against this schema and return the encoded
        representation (bytes) using the codec specified in the schema.
        """
        # If the schema is permissive then validation can't fail
        if not self._bypass_validation:
            try:
                self._validate_row(row)
            except jsonschema.exceptions.ValidationError as ve:
                raise exceptions.MetadataValidationError(str(ve)) from ve
        return self.encode_row(row)

    def decode_row(self, row: bytes) -> Any:
        """
        Decode an encoded row (bytes) of metadata, using the codec specifed in the schema
        and return a python dict. Note that no validation of the metadata against the
        schema is performed.
        """
        # Set by __init__
        pass  # pragma: no cover

    def encode_row(self, row: Any) -> bytes:
        """
        Encode a row (dict) of metadata to its binary representation (bytes)
        using the codec specified in the schema. Note that unlike
        :meth:`validate_and_encode_row` no validation against the schema is performed.
        This should only be used for performance if a validation check is not needed.
        """
        # Set by __init__
        pass  # pragma: no cover

    @staticmethod
    def permissive_json():
        """
        The simplest, permissive JSON schema. Only specifies the JSON codec and has
        no constraints on the properties.
        """
        return MetadataSchema({"codec": "json"})

    @staticmethod
    def null():
        """
        The null schema which defines no properties and results in raw bytes
        being returned on accessing metadata column.
        """
        return MetadataSchema(None)


# Often many replicate tree sequences are processed with identical schemas, so cache them
@functools.lru_cache(maxsize=128)
def parse_metadata_schema(encoded_schema: str) -> MetadataSchema:
    """
    Create a schema object from its string encoding. The exact class returned is
    determined by the ``encoding`` specification in the string.

    :param str encoded_schema: The string encoded schema.
    :return: A subclass of AbstractMetadataSchema.
    """
    if encoded_schema == "":
        return MetadataSchema.null()
    else:
        try:
            decoded = json.loads(
                encoded_schema, object_pairs_hook=collections.OrderedDict
            )
        except json.decoder.JSONDecodeError:
            raise ValueError(f"Metadata schema is not JSON, found {encoded_schema}")
        return MetadataSchema(decoded)


class _CachedMetadata:
    """
    Descriptor for lazy decoding of metadata on attribute access.
    """

    def __get__(self, row, owner):
        if row._metadata_decoder is not None:
            # Some classes that use this are frozen so we need to directly setattr.
            __builtins__object__setattr__(
                row, "_metadata", row._metadata_decoder(row._metadata)
            )
            # Decoder being None indicates that metadata is decoded
            __builtins__object__setattr__(row, "_metadata_decoder", None)
        return row._metadata

    def __set__(self, row, value):
        __builtins__object__setattr__(row, "_metadata", value)


def lazy_decode(own_init=False):
    def _lazy_decode(cls):
        """
        Modifies a dataclass such that it lazily decodes metadata, if it is encoded.
        If the metadata passed to the constructor is encoded a `metadata_decoder`
        parameter must be also be passed.
        """
        if not own_init:
            wrapped_init = cls.__init__

            # Intercept the init to record the decoder
            def new_init(self, *args, metadata_decoder=None, **kwargs):
                __builtins__object__setattr__(
                    self, "_metadata_decoder", metadata_decoder
                )
                wrapped_init(self, *args, **kwargs)

            cls.__init__ = new_init

        # Add a descriptor to the class to decode and cache metadata
        cls.metadata = _CachedMetadata()

        # Add slots needed to the class
        slots = cls.__slots__
        slots.extend(["_metadata", "_metadata_decoder"])
        dict_ = dict()
        sloted_members = dict()
        for k, v in cls.__dict__.items():
            if k not in slots:
                dict_[k] = v
            elif not isinstance(v, types.MemberDescriptorType):
                sloted_members[k] = v
        new_cls = type(cls.__name__, cls.__bases__, dict_)
        for k, v in sloted_members.items():
            setattr(new_cls, k, v)
        return new_cls

    return _lazy_decode


class MetadataProvider:
    """
    Abstract superclass of container objects that provide metadata.
    """

    def __init__(self, ll_object):
        self._ll_object = ll_object

    @property
    def metadata_schema(self) -> MetadataSchema:
        """
        The :class:`tskit.MetadataSchema` for this object.
        """
        return parse_metadata_schema(self._ll_object.metadata_schema)

    @metadata_schema.setter
    def metadata_schema(self, schema: MetadataSchema) -> None:
        # Check the schema is a valid schema instance by roundtripping it.
        text_version = repr(schema)
        parse_metadata_schema(text_version)
        self._ll_object.metadata_schema = text_version

    @property
    def metadata(self) -> Any:
        """
        The decoded metadata for this object.
        """
        return self.metadata_schema.decode_row(self.metadata_bytes)

    @metadata.setter
    def metadata(self, metadata: bytes | dict | None) -> None:
        encoded = self.metadata_schema.validate_and_encode_row(metadata)
        self._ll_object.metadata = encoded

    @property
    def metadata_bytes(self) -> Any:
        """
        The raw bytes of metadata for this TableCollection
        """
        return self._ll_object.metadata

    @property
    def nbytes(self) -> int:
        return len(self._ll_object.metadata) + len(self._ll_object.metadata_schema)

    def assert_equals(self, other: MetadataProvider):
        if self.metadata_schema != other.metadata_schema:
            raise AssertionError(
                f"Metadata schemas differ: self={self.metadata_schema} "
                f"other={other.metadata_schema}"
            )
        if self.metadata != other.metadata:
            raise AssertionError(
                f"Metadata differs: self={self.metadata} " f"other={other.metadata}"
            )


--- ../../tskit/python/tskit/_version.py ---

# Definitive location for the version number.
# During development, should be x.y.z.devN
# For beta should be x.y.zbN
tskit_version = "0.6.1.dev0"


--- ../../tskit/python/lwt_interface/setup.py ---

import os.path
import platform

from setuptools import Extension
from setuptools import setup
from setuptools.command.build_ext import build_ext


IS_WINDOWS = platform.system() == "Windows"


# Obscure magic required to allow numpy be used as a 'setup_requires'.
# Based on https://stackoverflow.com/questions/19919905
class local_build_ext(build_ext):
    def finalize_options(self):
        build_ext.finalize_options(self)
        import builtins

        # Prevent numpy from thinking it is still in its setup process:
        builtins.__NUMPY_SETUP__ = False
        import numpy

        self.include_dirs.append(numpy.get_include())


libdir = "../lib"
kastore_dir = os.path.join(libdir, "subprojects", "kastore")
# TODO pathlib glob this.
tsk_source_files = [
    "core.c",
    "tables.c",
    "trees.c",
    "genotypes.c",
    "stats.c",
    "convert.c",
    "haplotype_matching.c",
]
sources = (
    ["example_c_module.c"]
    + [os.path.join(libdir, "tskit", f) for f in tsk_source_files]
    + [os.path.join(kastore_dir, "kastore.c")]
)

defines = []
libraries = []
if IS_WINDOWS:
    # Needed for generating UUIDs in tskit
    libraries.append("Advapi32")
    defines.append(("WIN32", None))

extension_module = Extension(
    "example_c_module",
    sources=sources,
    extra_compile_args=["-std=c99"],
    libraries=libraries,
    define_macros=defines,
    include_dirs=[libdir, kastore_dir],
)

numpy_ver = "numpy>=1.7"

setup(
    name="example_c_module",
    description="Example usage of the LightweightTableCollection tskit interface",
    ext_modules=[extension_module],
    setup_requires=[numpy_ver],
    cmdclass={"build_ext": local_build_ext},
    license="MIT",
    platforms=["POSIX", "Windows", "MacOS X"],
)


--- ../../tskit/python/lwt_interface/dict_encoding_testlib.py ---

 
"""
Test definitions for the low-level LightweightTableCollection class
defined here. These tests are not intended to be executed directly,
but should be imported into another test module that imports a
compiled module exporting the LightweightTableCollection class.
See the test_example_c_module file for an example.
"""
import copy

import kastore
import msprime
import numpy as np
import pytest

import tskit
import tskit.util as util

lwt_module = None

NON_UTF8_STRING = "\ud861\udd37"


@pytest.fixture(scope="session")
def full_ts():
    """
    A tree sequence with data in all fields - duplicated from tskit's conftest.py
    as other test suites using this file will not have that fixture defined.
    """
    demography = msprime.Demography()
    demography.add_population(initial_size=100, name="A")
    demography.add_population(initial_size=100, name="B")
    demography.add_population(initial_size=100, name="C")
    demography.add_population_split(time=10, ancestral="C", derived=["A", "B"])

    ts = msprime.sim_ancestry(
        {"A": 5, "B": 5},
        demography=demography,
        random_seed=1,
        sequence_length=10,
        record_migrations=True,
    )
    assert ts.num_migrations > 0
    assert ts.num_individuals > 0
    ts = msprime.sim_mutations(ts, rate=0.1, random_seed=2)
    assert ts.num_mutations > 0
    tables = ts.dump_tables()
    tables.individuals.clear()

    for ind in ts.individuals():
        tables.individuals.add_row(flags=0, location=[ind.id, ind.id], parents=[-1, -1])

    for name, table in tables.table_name_map.items():
        if name != "provenances":
            table.metadata_schema = tskit.MetadataSchema({"codec": "json"})
            metadatas = [f"n_{name}_{u}" for u in range(len(table))]
            metadata, metadata_offset = tskit.pack_strings(metadatas)
            table.set_columns(
                **{
                    **table.asdict(),
                    "metadata": metadata,
                    "metadata_offset": metadata_offset,
                }
            )
    tables.metadata_schema = tskit.MetadataSchema({"codec": "json"})
    tables.metadata = {"A": "Test metadata"}

    tables.reference_sequence.data = "A" * int(tables.sequence_length)
    tables.reference_sequence.url = "https://example.com/sequence"
    tables.reference_sequence.metadata_schema = tskit.MetadataSchema.permissive_json()
    tables.reference_sequence.metadata = {"A": "Test metadata"}

    # Add some more provenance so we have enough rows for the offset deletion test.
    for j in range(10):
        tables.provenances.add_row(timestamp="x" * j, record="y" * j)
    return tables.tree_sequence()


# The ts above is used for the whole test session, but our tests need fresh tables to
# modify
@pytest.fixture
def tables(full_ts):
    return full_ts.dump_tables()


def test_check_ts_full(tmp_path, full_ts):
    """
    Check that the example ts has data in all fields
    """
    full_ts.dump(tmp_path / "tables")
    store = kastore.load(tmp_path / "tables")
    for v in store.values():
        assert v.nbytes > 0


class TestEncodingVersion:
    def test_version(self):
        lwt = lwt_module.LightweightTableCollection()
        assert lwt.asdict()["encoding_version"] == (1, 6)


class TestRoundTrip:
    """
    Tests if we can do a simple round trip on simulated data.
    """

    def verify(self, tables):
        lwt = lwt_module.LightweightTableCollection()
        lwt.fromdict(tables.asdict())
        other_tables = tskit.TableCollection.fromdict(lwt.asdict())
        tables.assert_equals(other_tables)

    def test_simple(self):
        ts = msprime.simulate(10, mutation_rate=1, random_seed=2)
        self.verify(ts.tables)

    def test_empty(self):
        tables = tskit.TableCollection(sequence_length=1)
        self.verify(tables)

    def test_individuals(self):
        n = 10
        ts = msprime.simulate(n, mutation_rate=1, random_seed=2)
        tables = ts.dump_tables()
        for j in range(n):
            tables.individuals.add_row(
                flags=j, location=(j, j), parents=(j, j), metadata=b"x" * j
            )
        self.verify(tables)

    def test_sequence_length(self):
        ts = msprime.simulate(
            10, recombination_rate=0.1, mutation_rate=1, length=0.99, random_seed=2
        )
        self.verify(ts.tables)

    def test_migration(self):
        pop_configs = [msprime.PopulationConfiguration(5) for _ in range(2)]
        migration_matrix = [[0, 1], [1, 0]]
        ts = msprime.simulate(
            population_configurations=pop_configs,
            migration_matrix=migration_matrix,
            mutation_rate=1,
            record_migrations=True,
            random_seed=1,
        )
        self.verify(ts.tables)

    def test_example(self, tables):
        tables.metadata_schema = tskit.MetadataSchema(
            {
                "codec": "struct",
                "type": "object",
                "properties": {"top-level": {"type": "string", "binaryFormat": "50p"}},
            }
        )
        tables.metadata = {"top-level": "top-level-metadata"}
        for table in tskit.TABLE_NAMES:
            t = getattr(tables, table)
            if hasattr(t, "metadata_schema"):
                t.packset_metadata([f"{table}-{i}".encode() for i in range(t.num_rows)])
                t.metadata_schema = tskit.MetadataSchema(
                    {
                        "codec": "struct",
                        "type": "object",
                        "properties": {
                            table: {"type": "string", "binaryFormat": "50p"}
                        },
                    }
                )

        self.verify(tables)


class TestMissingData:
    """
    Tests what happens when we have missing data in the encoded dict.
    """

    def test_missing_sequence_length(self, tables):
        d = tables.asdict()
        del d["sequence_length"]
        lwt = lwt_module.LightweightTableCollection()
        with pytest.raises(TypeError):
            lwt.fromdict(d)

    def test_missing_time_units(self, tables):
        assert tables.time_units != ""
        d = tables.asdict()
        del d["time_units"]
        lwt = lwt_module.LightweightTableCollection()
        lwt.fromdict(d)
        tables = tskit.TableCollection.fromdict(lwt.asdict())
        assert tables.time_units == tskit.TIME_UNITS_UNKNOWN

    def test_missing_metadata(self, tables):
        assert tables.metadata != b""
        d = tables.asdict()
        del d["metadata"]
        lwt = lwt_module.LightweightTableCollection()
        lwt.fromdict(d)
        tables = tskit.TableCollection.fromdict(lwt.asdict())
        # Empty byte field still gets interpreted by schema
        assert tables.metadata == {}

    def test_missing_metadata_schema(self, tables):
        assert repr(tables.metadata_schema) != ""
        d = tables.asdict()
        del d["metadata_schema"]
        lwt = lwt_module.LightweightTableCollection()
        lwt.fromdict(d)
        tables = tskit.TableCollection.fromdict(lwt.asdict())
        assert repr(tables.metadata_schema) == ""

    def test_missing_tables(self, tables):
        d = tables.asdict()
        table_names = d.keys() - {
            "sequence_length",
            "time_units",
            "metadata",
            "metadata_schema",
            "encoding_version",
            "indexes",
            "reference_sequence",
        }
        for table_name in table_names:
            d = tables.asdict()
            del d[table_name]
            lwt = lwt_module.LightweightTableCollection()
            with pytest.raises(TypeError):
                lwt.fromdict(d)


class TestBadTypes:
    """
    Tests for setting each column to a type that can't be converted to 1D numpy array.
    """

    def verify_columns(self, value, tables):
        d = tables.asdict()
        table_names = set(d.keys()) - {
            "sequence_length",
            "time_units",
            "metadata",
            "metadata_schema",
            "encoding_version",
            "indexes",
            "reference_sequence",
        }
        for table_name in table_names:
            table_dict = d[table_name]
            for colname in set(table_dict.keys()) - {"metadata_schema"}:
                d_copy = dict(table_dict)
                d_copy[colname] = value
                lwt = lwt_module.LightweightTableCollection()
                d = tables.asdict()
                d[table_name] = d_copy
                with pytest.raises(ValueError):
                    lwt.fromdict(d)

    def test_2d_array(self, tables):
        self.verify_columns([[1, 2], [3, 4]], tables)

    def test_str(self, tables):
        self.verify_columns("aserg", tables)

    def test_bad_top_level_types(self, tables):
        d = tables.asdict()
        for key in set(d.keys()) - {"encoding_version", "indexes"}:
            bad_type_dict = tables.asdict()
            # A list should be a ValueError for both the tables and sequence_length
            bad_type_dict[key] = ["12345"]
            lwt = lwt_module.LightweightTableCollection()
            with pytest.raises(TypeError):
                lwt.fromdict(bad_type_dict)


class TestBadLengths:
    """
    Tests for setting each column to a length incompatible with the table.
    """

    def verify(self, num_rows, tables):
        d = tables.asdict()
        table_names = set(d.keys()) - {
            "sequence_length",
            "time_units",
            "metadata",
            "metadata_schema",
            "encoding_version",
            "indexes",
            "reference_sequence",
        }
        for table_name in sorted(table_names):
            table_dict = d[table_name]
            for colname in set(table_dict.keys()) - {"metadata_schema"}:
                d_copy = dict(table_dict)
                d_copy[colname] = table_dict[colname][:num_rows].copy()
                lwt = lwt_module.LightweightTableCollection()
                d = tables.asdict()
                d[table_name] = d_copy
                with pytest.raises(ValueError):
                    lwt.fromdict(d)

    def test_two_rows(self, tables):
        self.verify(2, tables)

    def test_zero_rows(self, tables):
        self.verify(0, tables)

    def test_bad_index_length(self, tables):
        for col in ("insertion", "removal"):
            d = tables.asdict()
            d["indexes"][f"edge_{col}_order"] = d["indexes"][f"edge_{col}_order"][:-1]
            lwt = lwt_module.LightweightTableCollection()
            with pytest.raises(
                ValueError,
                match="^edge_insertion_order and"
                " edge_removal_order must be the same"
                " length$",
            ):
                lwt.fromdict(d)
        d = tables.asdict()
        for col in ("insertion", "removal"):
            d["indexes"][f"edge_{col}_order"] = d["indexes"][f"edge_{col}_order"][:-1]
        lwt = lwt_module.LightweightTableCollection()
        with pytest.raises(
            ValueError,
            match="^edge_insertion_order and edge_removal_order must be"
            " the same length as the number of edges$",
        ):
            lwt.fromdict(d)


class TestParsingUtilities:
    def test_missing_required(self, tables):
        d = tables.asdict()
        del d["sequence_length"]
        lwt = lwt_module.LightweightTableCollection()
        with pytest.raises(TypeError, match="'sequence_length' is required"):
            lwt.fromdict(d)

    def test_string_bad_type(self, tables):
        d = tables.asdict()
        d["time_units"] = b"sdf"
        lwt = lwt_module.LightweightTableCollection()
        with pytest.raises(TypeError, match="'time_units' is not a string"):
            lwt.fromdict(d)

    def test_bytes_bad_type(self, tables):
        d = tables.asdict()
        d["metadata"] = 1234
        lwt = lwt_module.LightweightTableCollection()
        with pytest.raises(TypeError, match="'metadata' is not bytes"):
            lwt.fromdict(d)

    def test_dict_bad_type(self, tables):
        d = tables.asdict()
        d["nodes"] = b"sdf"
        lwt = lwt_module.LightweightTableCollection()
        with pytest.raises(TypeError, match="'nodes' is not a dict"):
            lwt.fromdict(d)

    def test_bad_strings(self, tables):
        def verify_unicode_error(d):
            lwt = lwt_module.LightweightTableCollection()
            with pytest.raises(UnicodeEncodeError):
                lwt.fromdict(d)

        def verify_bad_string_type(d):
            lwt = lwt_module.LightweightTableCollection()
            with pytest.raises(TypeError):
                lwt.fromdict(d)

        d = tables.asdict()
        for k, v in d.items():
            if isinstance(v, str):
                d_copy = copy.deepcopy(d)
                d_copy[k] = NON_UTF8_STRING
                verify_unicode_error(d_copy)
                d_copy[k] = 12345
                verify_bad_string_type(d_copy)
            if isinstance(v, dict):
                for kp, vp in v.items():
                    if isinstance(vp, str):
                        d_copy = copy.deepcopy(d)
                        d_copy[k][kp] = NON_UTF8_STRING
                        verify_unicode_error(d_copy)
                        d_copy[k][kp] = 12345
                        verify_bad_string_type(d_copy)


class TestRequiredAndOptionalColumns:
    """
    Tests that specifying None for some columns will give the intended
    outcome.
    """

    def verify_required_columns(self, tables, table_name, required_cols):
        d = tables.asdict()
        table_dict = {col: None for col in d[table_name].keys()}
        for col in required_cols:
            table_dict[col] = d[table_name][col]
        lwt = lwt_module.LightweightTableCollection()
        d[table_name] = table_dict
        lwt.fromdict(d)
        other = lwt.asdict()
        for col in required_cols:
            assert np.array_equal(other[table_name][col], table_dict[col])

        # Any one of these required columns as None gives an error.
        for col in required_cols:
            d = tables.asdict()
            d_copy = copy.deepcopy(table_dict)
            d_copy[col] = None
            d[table_name] = d_copy
            lwt = lwt_module.LightweightTableCollection()
            with pytest.raises(TypeError):
                lwt.fromdict(d)

        # Removing any one of these required columns gives an error.
        for col in required_cols:
            d = tables.asdict()
            d_copy = copy.deepcopy(table_dict)
            del d_copy[col]
            d[table_name] = d_copy
            lwt = lwt_module.LightweightTableCollection()
            with pytest.raises(TypeError):
                lwt.fromdict(d)

    def verify_optional_column(self, tables, table_len, table_name, col_name):
        d = tables.asdict()
        table_dict = d[table_name]
        table_dict[col_name] = None
        lwt = lwt_module.LightweightTableCollection()
        lwt.fromdict(d)
        out = lwt.asdict()
        assert np.array_equal(
            out[table_name][col_name], np.zeros(table_len, dtype=np.int32) - 1
        )

    def verify_offset_pair(
        self, tables, table_len, table_name, col_name, required=False
    ):
        offset_col = col_name + "_offset"

        if not required:
            d = tables.asdict()
            table_dict = d[table_name]
            table_dict[col_name] = None
            table_dict[offset_col] = None
            lwt = lwt_module.LightweightTableCollection()
            lwt.fromdict(d)
            out = lwt.asdict()
            assert out[table_name][col_name].shape == (0,)
            assert np.array_equal(
                out[table_name][offset_col],
                np.zeros(table_len + 1, dtype=np.uint32),
            )
            d = tables.asdict()
            table_dict = d[table_name]
            del table_dict[col_name]
            del table_dict[offset_col]
            lwt = lwt_module.LightweightTableCollection()
            lwt.fromdict(d)
            out = lwt.asdict()
            assert out[table_name][col_name].shape == (0,)
            assert np.array_equal(
                out[table_name][offset_col],
                np.zeros(table_len + 1, dtype=np.uint32),
            )

        # Setting one or the other raises a TypeError
        d = tables.asdict()
        table_dict = d[table_name]
        table_dict[col_name] = None
        lwt = lwt_module.LightweightTableCollection()
        with pytest.raises(TypeError):
            lwt.fromdict(d)

        d = tables.asdict()
        table_dict = d[table_name]
        del table_dict[col_name]
        lwt = lwt_module.LightweightTableCollection()
        with pytest.raises(TypeError):
            lwt.fromdict(d)

        d = tables.asdict()
        table_dict = d[table_name]
        table_dict[offset_col] = None
        lwt = lwt_module.LightweightTableCollection()
        with pytest.raises(TypeError):
            lwt.fromdict(d)

        d = tables.asdict()
        table_dict = d[table_name]
        del table_dict[offset_col]
        lwt = lwt_module.LightweightTableCollection()
        with pytest.raises(TypeError):
            lwt.fromdict(d)

        d = tables.asdict()
        table_dict = d[table_name]
        bad_offset = np.zeros_like(table_dict[offset_col])
        bad_offset[:-1] = table_dict[offset_col][:-1][::-1]
        bad_offset[-1] = table_dict[offset_col][-1]
        table_dict[offset_col] = bad_offset
        lwt = lwt_module.LightweightTableCollection()
        with pytest.raises(ValueError):
            lwt.fromdict(d)

    def verify_metadata_schema(self, tables, table_name):
        d = tables.asdict()
        d[table_name]["metadata_schema"] = None
        lwt = lwt_module.LightweightTableCollection()
        lwt.fromdict(d)
        out = lwt.asdict()
        assert "metadata_schema" not in out[table_name]
        tables = tskit.TableCollection.fromdict(out)
        assert repr(getattr(tables, table_name).metadata_schema) == ""

    def test_individuals(self, tables):
        self.verify_required_columns(tables, "individuals", ["flags"])
        self.verify_offset_pair(
            tables, len(tables.individuals), "individuals", "location"
        )
        self.verify_offset_pair(
            tables, len(tables.individuals), "individuals", "parents"
        )
        self.verify_offset_pair(
            tables, len(tables.individuals), "individuals", "metadata"
        )
        self.verify_metadata_schema(tables, "individuals")
        # Verify optional parents column
        d = tables.asdict()
        d["individuals"]["parents"] = None
        d["individuals"]["parents_offset"] = None
        lwt = lwt_module.LightweightTableCollection()
        lwt.fromdict(d)
        out = lwt.asdict()
        assert all(val == [] for val in out["individuals"]["parents"])

    def test_nodes(self, tables):
        self.verify_offset_pair(tables, len(tables.nodes), "nodes", "metadata")
        self.verify_optional_column(tables, len(tables.nodes), "nodes", "population")
        self.verify_optional_column(tables, len(tables.nodes), "nodes", "individual")
        self.verify_required_columns(tables, "nodes", ["flags", "time"])
        self.verify_metadata_schema(tables, "nodes")

    def test_edges(self, tables):
        self.verify_required_columns(
            tables, "edges", ["left", "right", "parent", "child"]
        )
        self.verify_offset_pair(tables, len(tables.edges), "edges", "metadata")
        self.verify_metadata_schema(tables, "edges")

    def test_migrations(self, tables):
        self.verify_required_columns(
            tables, "migrations", ["left", "right", "node", "source", "dest", "time"]
        )
        self.verify_offset_pair(
            tables, len(tables.migrations), "migrations", "metadata"
        )
        self.verify_optional_column(tables, len(tables.nodes), "nodes", "individual")
        self.verify_metadata_schema(tables, "migrations")

    def test_sites(self, tables):
        self.verify_required_columns(
            tables, "sites", ["position", "ancestral_state", "ancestral_state_offset"]
        )
        self.verify_offset_pair(tables, len(tables.sites), "sites", "metadata")
        self.verify_metadata_schema(tables, "sites")

    def test_mutations(self, tables):
        self.verify_required_columns(
            tables,
            "mutations",
            ["site", "node", "derived_state", "derived_state_offset"],
        )
        self.verify_offset_pair(tables, len(tables.mutations), "mutations", "metadata")
        self.verify_metadata_schema(tables, "mutations")
        # Verify optional time column
        d = tables.asdict()
        d["mutations"]["time"] = None
        lwt = lwt_module.LightweightTableCollection()
        lwt.fromdict(d)
        out = lwt.asdict()
        assert all(util.is_unknown_time(val) for val in out["mutations"]["time"])

    def test_populations(self, tables):
        self.verify_required_columns(
            tables, "populations", ["metadata", "metadata_offset"]
        )
        self.verify_metadata_schema(tables, "populations")
        self.verify_offset_pair(tables, len(tables.nodes), "nodes", "metadata", True)

    def test_provenances(self, tables):
        self.verify_required_columns(
            tables,
            "provenances",
            ["record", "record_offset", "timestamp", "timestamp_offset"],
        )

    def test_index(self, tables):
        d = tables.asdict()
        lwt = lwt_module.LightweightTableCollection()
        lwt.fromdict(d)
        other = lwt.asdict()
        assert np.array_equal(
            d["indexes"]["edge_insertion_order"],
            other["indexes"]["edge_insertion_order"],
        )
        assert np.array_equal(
            d["indexes"]["edge_removal_order"], other["indexes"]["edge_removal_order"]
        )

        # index is optional
        d = tables.asdict()
        del d["indexes"]
        lwt = lwt_module.LightweightTableCollection()
        lwt.fromdict(d)
        # and a tc without indexes has empty dict
        assert lwt.asdict()["indexes"] == {}

        # Both columns must be provided, if one is
        for col in ("insertion", "removal"):
            d = tables.asdict()
            del d["indexes"][f"edge_{col}_order"]
            lwt = lwt_module.LightweightTableCollection()
            with pytest.raises(
                TypeError,
                match="^edge_insertion_order and "
                "edge_removal_order must be specified "
                "together$",
            ):
                lwt.fromdict(d)

    def test_index_bad_type(self, tables):
        d = tables.asdict()
        lwt = lwt_module.LightweightTableCollection()
        d["indexes"] = "asdf"
        with pytest.raises(TypeError):
            lwt.fromdict(d)

    def test_reference_sequence(self, tables):
        self.verify_metadata_schema(tables, "reference_sequence")

        def get_refseq(d):
            tables = tskit.TableCollection.fromdict(d)
            return tables.reference_sequence

        d = tables.asdict()
        refseq_dict = d.pop("reference_sequence")
        assert get_refseq(d).is_null()

        # All empty strings is the same thing
        d["reference_sequence"] = dict(
            data="", url="", metadata_schema="", metadata=b""
        )
        assert get_refseq(d).is_null()

        del refseq_dict["metadata_schema"]  # handled above
        for key, value in refseq_dict.items():
            d["reference_sequence"] = {key: value}
            refseq = get_refseq(d)
            assert not refseq.is_null()
            assert getattr(refseq, key) == value

    def test_top_level_time_units(self, tables):
        d = tables.asdict()
        # None should give default value
        d["time_units"] = None
        lwt = lwt_module.LightweightTableCollection()
        lwt.fromdict(d)
        out = lwt.asdict()
        tables = tskit.TableCollection.fromdict(out)
        assert tables.time_units == tskit.TIME_UNITS_UNKNOWN
        # Missing is tested in TestMissingData above
        d = tables.asdict()
        d["time_units"] = NON_UTF8_STRING
        lwt = lwt_module.LightweightTableCollection()
        with pytest.raises(UnicodeEncodeError):
            lwt.fromdict(d)

    def test_top_level_metadata(self, tables):
        d = tables.asdict()
        # None should give default value
        d["metadata"] = None
        lwt = lwt_module.LightweightTableCollection()
        lwt.fromdict(d)
        out = lwt.asdict()
        assert "metadata" not in out
        tables = tskit.TableCollection.fromdict(out)
        assert tables.metadata == {}
        # Missing is tested in TestMissingData above

    def test_top_level_metadata_schema(self, tables):
        d = tables.asdict()
        # None should give default value
        d["metadata_schema"] = None
        lwt = lwt_module.LightweightTableCollection()
        lwt.fromdict(d)
        out = lwt.asdict()
        assert "metadata_schema" not in out
        tables = tskit.TableCollection.fromdict(out)
        assert repr(tables.metadata_schema) == ""
        # Missing is tested in TestMissingData above


class TestLifecycle:
    def test_unassigned_empty(self):
        lwt_dict = lwt_module.LightweightTableCollection().asdict()
        assert tskit.TableCollection.fromdict(lwt_dict) == tskit.TableCollection(-1)

    def test_del_empty(self):
        lwt = lwt_module.LightweightTableCollection()
        lwt_dict = lwt.asdict()
        del lwt
        assert tskit.TableCollection.fromdict(lwt_dict) == tskit.TableCollection(-1)

    def test_del_full(self, tables):
        lwt = lwt_module.LightweightTableCollection()
        lwt.fromdict(tables.asdict())
        lwt_dict = lwt.asdict()
        del lwt
        assert tskit.TableCollection.fromdict(lwt_dict) == tables

    def test_del_lwt_and_tables(self, tables):
        lwt = lwt_module.LightweightTableCollection()
        lwt.fromdict(tables.asdict())
        lwt_dict = lwt.asdict()
        del lwt
        tables2 = tables.copy()
        del tables
        assert tskit.TableCollection.fromdict(lwt_dict) == tables2


class TestForceOffset64:
    def get_offset_columns(self, dict_encoding):
        for table_name, table in dict_encoding.items():
            if isinstance(table, dict):
                for name, array in table.items():
                    if name.endswith("_offset"):
                        yield f"{table_name}/{name}", array

    def test_bad_args(self, tables):
        lwt = lwt_module.LightweightTableCollection()
        lwt.fromdict(tables.asdict())
        for bad_type in [None, {}, "sdf"]:
            with pytest.raises(TypeError):
                lwt.asdict(bad_type)

    def test_off_by_default(self, tables):
        lwt = lwt_module.LightweightTableCollection()
        lwt.fromdict(tables.asdict())
        d = lwt.asdict()
        for _, array in self.get_offset_columns(d):
            assert array.dtype == np.uint32

    def test_types_64(self, tables):
        lwt = lwt_module.LightweightTableCollection()
        lwt.fromdict(tables.asdict())
        d = lwt.asdict(force_offset_64=True)
        for _, array in self.get_offset_columns(d):
            assert array.dtype == np.uint64

    def test_types_32(self, tables):
        lwt = lwt_module.LightweightTableCollection()
        lwt.fromdict(tables.asdict())
        d = lwt.asdict(force_offset_64=False)
        for _, array in self.get_offset_columns(d):
            assert array.dtype == np.uint32

    def test_values_equal(self, tables):
        lwt = lwt_module.LightweightTableCollection()
        lwt.fromdict(tables.asdict())
        d64 = lwt.asdict(force_offset_64=True)
        d32 = lwt.asdict(force_offset_64=False)
        offsets_64 = dict(self.get_offset_columns(d64))
        offsets_32 = dict(self.get_offset_columns(d32))
        for col_name, col_32 in offsets_32.items():
            col_64 = offsets_64[col_name]
            assert col_64.shape == col_32.shape
            assert np.all(col_64 == col_32)


@pytest.mark.parametrize("bad_type", [None, "", []])
def test_fromdict_bad_type(bad_type):
    lwt = lwt_module.LightweightTableCollection()
    with pytest.raises(TypeError):
        lwt.fromdict(bad_type)


--- ../../tskit/python/lwt_interface/test_example_c_module.py ---

# flake8: noqa
import os
import sys

import pytest

# Make sure we use the local tskit version.

sys.path.insert(0, os.path.abspath("../"))

# An example of how to run the tests defined in the dict_encoding_testlib.py
# file for a given compiled version of the code.
import dict_encoding_testlib
import example_c_module
import tskit

# The test cases defined in dict_encoding_testlib all use the form
# lwt_module.LightweightTableCollection() to create an instance
# of LightweightTableCollection. So, by setting this variable in
# the module here, we can control which definition of the
# LightweightTableCollection gets used.
dict_encoding_testlib.lwt_module = example_c_module

from dict_encoding_testlib import *


def test_example_receiving():
    # The example_receiving function returns true if the first tree
    # has more than one root
    lwt = example_c_module.LightweightTableCollection()
    tables = tskit.TableCollection(1)
    lwt.fromdict(tables.asdict())
    # Our example function throws an error for an empty table collection
    with pytest.raises(ValueError, match="Table collection must be indexed"):
        example_c_module.example_receiving(lwt)

    # This tree sequence has one root so we get false
    tables = msprime.simulate(10).tables
    lwt.fromdict(tables.asdict())
    assert not example_c_module.example_receiving(lwt)

    # Add a root and we get true
    tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE)
    lwt.fromdict(tables.asdict())
    assert example_c_module.example_receiving(lwt)


def test_example_modifying():
    lwt = example_c_module.LightweightTableCollection()
    # The example_modifying function clears out the table and adds two rows
    tables = msprime.simulate(10, random_seed=42).tables
    assert tables.edges.num_rows == 18
    assert tables.nodes.num_rows == 19
    lwt.fromdict(tables.asdict())
    example_c_module.example_modifying(lwt)
    modified_tables = tskit.TableCollection.fromdict(lwt.asdict())
    assert modified_tables.edges.num_rows == 0
    assert modified_tables.nodes.num_rows == 2


--- ../../tskit/python/lwt_interface/README.md ---

# LightweightTableCollection interface

The files in this directory define the LightweightTableCollection
interface used to safely interchange table collection data between 
different compiled instances of the tskit C library. This is a 
*very* specialised use-case, and unless you are using the tskit
C API in your own compiled Python module (either via Cython
or the Python C API), you almost certainly don't need to use
this code.

## Overview

To allow a tskit table collection to be transferred from one compiled Python
extension module to another the table collection is converted to a `dict` of
basic python types and numpy arrays. This is then converted back in the receiving
module. `tskit_lwt_interface.h` provides a function `register_lwt_class` that 
defines a Python class `LightweightTableCollection` that performs these conversions
with methods `asdict` and `fromdict`. These methods mirror the `asdict` and `fromdict`
methods on `tskit.TableCollection`.

## Usage
An example C module skeleton `example_c_module.c` is provided, which shows passing tables
to the C module. See `test_example_c_module.py` for the python example usage
of the example module.

To add the 
`LightweightTableCollection` type to your module you include `tskit_lwt_interface.h`
and then call `register_lwt_class` on your C Python module object. You can then convert
to and from the lightweight table collection in Python, for example to convert a tskit
`TableCollection` to a `LightweightTableCollection`:
```python
tables = tskit.TableCollection(1)
lwt = example_c_module.LightweightTableCollection()
lwt.fromdict(tables.asdict())
```
and vice-versa:
```python
tc = tskit.TableCollection(lwt.asdict())
```
In C you can access the tables in a `LightweightTableCollection` instance that is passed 
to your function, as shown in the `example_receiving` function in `example_c_module.c`. 
Note the requirement to check for errors from tskit functions and to call
`handle_tskit_error` to set a Python error, returning `NULL` to Python to indicate error.

Tables can also be modified in the extension code as in `example_modifying`. We recommend
creating table collections in Python then passing them to C for modification rather than
creating them in C and returning them. This avoids complex managing of object lifecycles
in C code.




    





--- ../../tskit/python/benchmark/run-for-all-releases.py ---

import json
import subprocess
from urllib.request import urlopen

import tqdm
from distutils.version import StrictVersion


def versions(package_name):
    url = f"https://pypi.org/pypi/{package_name}/json"
    data = json.load(urlopen(url))
    return sorted(data["releases"].keys(), key=StrictVersion)


def sh(command):
    subprocess.run(command, check=True, shell=True)


if __name__ == "__main__":
    try:
        sh("python -m venv _bench-temp-venv")
        sh("_bench-temp-venv/bin/pip install -r ../requirements/development.txt")
        versions = [
            v
            for v in versions("tskit")
            # We don't want alphas, betas or two broken versions:
            if "a" not in v and "b" not in v and v not in ("0.0.0", "0.1.0")
        ]
        for v in tqdm.tqdm(versions):
            sh(f"_bench-temp-venv/bin/pip install tskit=={v}")
            sh("_bench-temp-venv/bin/python run.py")
    finally:
        sh("rm -rf _bench-temp-venv")


--- ../../tskit/python/benchmark/run.py ---

import json
import os.path
import platform
import sys
import timeit
from pathlib import Path

import click
import psutil
import tqdm
import yaml
from matplotlib.colors import LinearSegmentedColormap
from si_prefix import si_format

tskit_dir = Path(__file__).parent.parent
sys.path.append(str(tskit_dir))
import tskit  # noqa: E402
import msprime  # noqa: E402

with open("config.yaml") as f:
    config = yaml.load(f, Loader=yaml.FullLoader)


def system_info():
    ret = {}
    uname = platform.uname()
    for attr in ["system", "node", "release", "version", "machine", "processor"]:
        ret[attr] = getattr(uname, attr)
    ret["python_version"] = sys.version
    cpufreq = psutil.cpu_freq()
    ret["physical_cores"] = psutil.cpu_count(logical=False)
    ret["total_cores"] = psutil.cpu_count(logical=True)
    ret["max_frequency"] = cpufreq.max
    ret["min_frequency"] = cpufreq.min
    ret["current_frequency"] = cpufreq.current
    ret["cpu_usage_per_core"] = [
        percentage for percentage in psutil.cpu_percent(percpu=True, interval=1)
    ]
    ret["total_cpu_usage"] = psutil.cpu_percent()
    return ret


def make_file():
    benchmark_trees = tskit_dir / "benchmark" / "bench.trees"
    if not os.path.exists(benchmark_trees):
        print("Generating benchmark trees...")
        demography = msprime.Demography()
        demography.add_population(name="A", initial_size=10_000)
        demography.add_population(name="B", initial_size=5_000)
        demography.add_population(name="C", initial_size=1_000)
        demography.add_population_split(time=1000, derived=["A", "B"], ancestral="C")
        ts = msprime.sim_ancestry(
            samples={"A": 25000, "B": 25000},
            demography=demography,
            sequence_length=1_000_000,
            random_seed=42,
            recombination_rate=0.0000001,
            record_migrations=True,
            record_provenance=True,
        )
        ts = msprime.sim_mutations(ts, rate=0.000001, random_seed=42)
        ts.dump(benchmark_trees)
        ts = msprime.sim_ancestry(
            samples={"A": 1, "B": 1},
            demography=demography,
            sequence_length=1,
            random_seed=42,
            recombination_rate=0,
            record_migrations=True,
            record_provenance=True,
        )
        ts = msprime.sim_mutations(ts, rate=0.001, random_seed=42)
        ts.dump(tskit_dir / "benchmark" / "tiny.trees")


def autotime(setup, code):
    t = timeit.Timer(setup=setup, stmt=code)
    try:
        one_run = t.timeit(number=1)
    except Exception as e:
        print(f"{code}: Error running benchmark: {e}")
        return None
    num_trials = int(max(1, 2 / one_run))
    return one_run, num_trials, t.timeit(number=num_trials) / num_trials


def run_benchmarks(keyword_filter):
    results = {}
    for benchmark in tqdm.tqdm(config["benchmarks"]):
        bench_name = benchmark.get("name", benchmark["code"])
        if keyword_filter not in bench_name:
            continue
        params = benchmark.get("parameters", {"noop": [None]})

        # Expand the parameters
        def sub_expand(context, name, d):
            if isinstance(d, dict):
                ret = []
                for k, v in d.items():
                    new_context = {**{k: v for k, v in context.items()}, name: k}
                    for k2, v2 in v.items():
                        ret += sub_expand(new_context, k2, v2)
                return ret
            elif isinstance(d, list):
                return [
                    {**{k: v for k, v in context.items()}, name: value} for value in d
                ]
            else:
                raise ValueError(f"Invalid parameter type: {type(d)}-{d}")

        expanded_params = []
        for k, v in params.items():
            expanded_params += sub_expand({}, k, v)

        for values in expanded_params:
            setup = (
                f"import sys;sys.path.append('{tskit_dir}');"
                + config["setup"].replace("\n", "\n")
                + benchmark.get("setup", "").replace("\n", "\n").format(**values)
            )
            code = benchmark["code"].replace("\n", "\n").format(**values)
            result = autotime(setup, code)
            if result is not None:
                one_run, num_trials, avg = result
                results.setdefault(bench_name, {})[code] = {
                    "one_run": one_run,
                    "num_trials": num_trials,
                    "avg": avg,
                }

    return results


def generate_report(all_versions_results):
    all_benchmarks = {}
    for _version, results in all_versions_results.items():
        for benchmark, values in results["tskit_benchmarks"].items():
            for code in values.keys():
                all_benchmarks.setdefault(benchmark, set()).add(code)

    all_versions = sorted(all_versions_results.keys())

    cmap = LinearSegmentedColormap.from_list("rg", ["g", "w", "r"], N=256)

    with open(tskit_dir / "benchmark" / "bench-results.html", "w") as f:
        f.write("<html><body>\n")
        f.write("<h1>tskit benchmark results</h1>\n")
        f.write("<table>\n")
        f.write("<tr><th></th>")
        for version in all_versions:
            f.write(f"<th>{version}</th>")
        f.write("</tr>\n")
        for benchmark in sorted(all_benchmarks.keys()):
            values = all_benchmarks[benchmark]
            indent = False
            if len(values) > 1:
                indent = True
                f.write(
                    f"<tr>"
                    f"  <td style='font-family: monospace'>"
                    f"    {benchmark}"
                    f"  </td>"
                    f"</tr>\n"
                )
            for code in sorted(values):
                f.write(
                    f"<tr><td style='font-family: monospace;"
                    f"padding-left: {'10px' if indent else 'inherit'}'>{code}</td>"
                )
                last_avg = None
                for version in all_versions:
                    try:
                        avg = all_versions_results[version]["tskit_benchmarks"][
                            benchmark
                        ][code]["avg"]
                        if last_avg is not None:
                            percent_change = 100 * ((avg - last_avg) / last_avg)
                            col = cmap(int(((percent_change / 100) * 128) + 128))
                            f.write(
                                f"<td style='background-color: rgba({col[0]*255},"
                                f" {col[1]*255}, {col[2]*255}, 1)'>"
                            )

                            f.write(f"{si_format(avg)} ({percent_change:.1f}%)")
                        else:
                            f.write(f"<td>{si_format(avg)}</td>")
                        last_avg = avg
                    except KeyError:
                        f.write("<td>N/A</td>")

                f.write("</tr>\n")
        f.write("</table>\n")


def print_result(results):
    max_name_length = max(len(name) for name in results.keys()) + 1
    for _bench, param_results in results.items():
        for name, data in param_results.items():
            print(name.ljust(max_name_length), si_format(data["avg"]))


@click.command()
@click.option(
    "--keyword_filter",
    "-k",
    type=str,
    default="",
    help="Only benchmarks with a name containing this string will be run",
)
@click.option("--print_results", "-p", is_flag=True, help="Print results to STDOUT")
def run_benchmark_and_save(keyword_filter, print_results):
    print("Benchmarking tskit version:", tskit._version.tskit_version)
    make_file()
    results = {}
    results["system"] = system_info()
    results["tskit_benchmarks"] = run_benchmarks(keyword_filter)

    if print_results:
        print_result(results["tskit_benchmarks"])

    all_versions_results = {}
    results_json = tskit_dir / "benchmark" / "bench-results.json"
    if os.path.exists(results_json):
        with open(results_json) as f:
            all_versions_results = json.load(f)

    all_versions_results[tskit._version.tskit_version] = results
    with open(results_json, "w") as f:
        json.dump(all_versions_results, f, indent=2)
    generate_report(all_versions_results)

    sys.exit(0)


if __name__ == "__main__":
    run_benchmark_and_save()


--- ../../tskit/docs/cli.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

```{currentmodule} tskit
```

(sec_cli)=

# Command line interface

```{eval-rst}
.. argparse::
    :module: tskit.cli
    :func: get_tskit_parser
    :prog: python3 -m tskit
```

--- ../../tskit/docs/topological-analysis.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

```{currentmodule} tskit
```

(sec_topological_analysis)=

# Topological analysis

The branching structure of a tree is known as the tree *topology*. Dealing with the
topological structure of trees is of key importance in dealing with genealogies
such as those encoded in the tree sequence structure.


(sec_topological_analysis_traversal)=

## Visiting nodes


(sec_topological_analysis_descending)=

### Descendant traversal

A basic thing to want to do is to visit all the nodes in a tree under a focal node
(i.e. the focal node and all its descendants). This can be done in various
*traversal orders*. The `tskit` library provides several methods to do this, such
as {meth}`Tree.nodes` in the {ref}`sec_python_api`. By default, this method
iterates over all the descendant nodes of the root(s) of the tree, and hence
visits all the tree's nodes:

```{code-cell} ipython3
import tskit
from IPython.display import SVG, display

tree = tskit.Tree.generate_balanced(10, arity=3)
display(SVG(tree.draw_svg()))

for order in ["preorder", "inorder", "postorder"]:
    print(f"{order}:\t", list(tree.nodes(order=order)))
```

Providing a focal node ID allows us to traverse through that node's descendants.
For instance, here we visit node 13 and its descendants:

```{code-cell} ipython3
print("Nodes under node 13:", [u for u in tree.nodes(13)])
```

The node IDs returned by traversal methods allow us to to access node information.
Below, for example, we use the {meth}`Tree.num_children` method to find the number of
children (the "arity") of every non-leaf node, and take the average. Since the
specific ordering of descendant nodes is not important in this case, we can leave
it out (defaulting to preorder traversal, the most efficient order):

```{code-cell} ipython3
import numpy as np
av_arity = np.mean([tree.num_children(u) for u in tree.nodes() if not tree.is_leaf(u)])
print(f"Average arity of internal nodes: {av_arity}")
```

:::{note}
In `tskit`, a tree can have multiple {ref}`sec_data_model_tree_roots`, each with a
descendant topology.  However, for some algorithms, instead of traversing through
the descendants of each root of a tree in turn, it can be helpful to start at a
single node and traverse downwards through the entire genealogy. The
{ref}`virtual root<sec_data_model_tree_virtual_root>` is provided for this purpose.
:::

#### Array methods

The {meth}`Tree.nodes` iterator provides a convenient way of looping over descendant
node IDs, but it can be more efficient to deal with all the IDs at once, as a
single array of values. This can be combined with
{ref}`direct memory access<sec_python_api_trees_node_measures_array>` resulting in a
high performance approach. Here, for example, is an equivalent
array-based method to find the average arity of internal nodes, by counting
how many times a node is referenced as a parent:

```{code-cell} ipython3
parent_id, count = np.unique(tree.parent_array[tree.preorder()], return_counts=True)
print(f"Average arity is {count[parent_id != tskit.NULL].mean()}")
```

:::{seealso}
The {ref}`sec_analysing_trees` tutorial provides a number of additional examples
of tree traversal techniques, with different performance characteristics.
:::


(sec_topological_analysis_ascending)=

### Ascending traversal

For many applications it is useful to be able to traverse upwards from a node or set
of nodes, such as the leaves. We can do this by iterating over parents. Here, for
example, we traverse upwards from each of the samples in the tree:

```{code-cell} ipython3
for u in tree.samples():
    path = []
    v = u
    while v != tskit.NULL:
        path.append(v)
        v = tree.parent(v)
    print(u, "->", path)
```

:::{todo}
Indicate that this can be made performant using `numba`, and link out to a tutorial
on high performance methods including the `numba` approach.
:::


(sec_combinatorics)=

## Identifying and counting topologies

tskit uses a combinatorial approach to identify unique topologies of
rooted, leaf-labelled trees. It provides methods
for enumerating all possible tree topologies, as well as converting
back and forth between a tree and its position, or rank, in the
enumeration of all possible topologies.
These methods do not only apply to binary trees;
rather, they cover general, rooted trees without unary nodes.

```{list-table} 
* - {meth}`Tree.rank`
  - Return the rank of this tree.
* - {meth}`Tree.unrank`
  - Return a Tree given its rank and a number of leaves.
* - {func}`tskit.all_trees`
  - Return a generator over all leaf-labelled trees of n leaves.
* - {func}`tskit.all_tree_shapes`
  - Return a generator over all tree shapes of n leaves.
* - {func}`tskit.all_tree_labellings`
  - Return a generator over all labellings of the given tree's shape.
```

:::{note}
As the number of nodes increases, the number of different topologies rises
extremely rapidly (see its entry in the
[On-Line Encyclopedia of Integer Sequences](https://oeis.org/A000311)). This
combinatorial explosion is a major limitation in any analysis that attempts to
explore possible topologies. For example, although the {func}`tskit.all_trees`
function above will happily start generating topologies for (say) a tree of 50
leaves, the total number of possible topologies is over $6^81$, which is of
the same order as the number of atoms in the observable universe. Generating
all the topologies of a tree with anything much more than 10 tips is likely
to be impracticable.
:::


(sec_tree_ranks)=

### Interpreting Tree Ranks

To understand tree ranks we must look at how leaf-labelled tree topologies
are enumerated. For example, we can use {func}`tskit.all_trees`
to generate all possible topologies of three leaves:

```{code-cell} ipython3
import tskit
from IPython.display import display, SVG

for t in tskit.all_trees(num_leaves=3):
    display(SVG(t.draw_svg(node_labels={0: 0, 1: 1, 2: 2}, order="tree", size=(120, 120))))
```

In this sequence, there exist two distinct tree shapes and each shape
can be labelled in at least one unique way. Given that topologies are
ordered first by their shape and then by their labelling, a tree
topology can be uniquely identified by

1. The shape of the tree
2. The labelling of the tree's shape

We can refer to the first tree in the above enumeration as the
first labelling of the first shape of trees with three leaves, or tree
$(0, 0)$. The second tree can be identified as the first labelling
of the second shape, or $(1, 0)$, and so on.
This pair of indexes for the shape and labelling of a tree is referred
to as the rank of the tree, and can be computed using the
{meth}`Tree.rank` method.

```{code-cell} ipython3
ranks = [t.rank() for t in tskit.all_trees(num_leaves=3)]
print("Ranks of 3-leaf trees:", ranks)
```

```{note}
Ranks in combinatorics are typically natural numbers. However,
we refer to this tuple of shape and label rank as a rank because
it serves the same purpose of indexing trees in an enumeration.
```

For details on how shapes and labellings are ordered, see
{ref}`sec_enumerating_topologies`.

We can also reconstruct a leaf-labelled tree given its rank. This process
is known as unranking, and can be performed using the {meth}`Tree.unrank`
method.

```{code-cell} ipython3
for rank in [(0, 0), (1, 0), (1, 1), (1, 2)]:
    t = tskit.Tree.unrank(num_leaves=3, rank=rank)
    display(SVG(t.draw_svg(node_labels={0: 0, 1: 1, 2: 2}, order="tree", size=(120, 120))))
```

#### Examples

One application of tree ranks is to count the different
leaf-labelled topologies in a tree sequence. Since the ranks
are just tuples, we can use a Python ``Counter`` to track them.
Here, we count and unrank the most frequently seen
topology in a tree sequence. For brevity, this example assumes
samples are synonymous with leaves.

```{code-cell} ipython3
import collections
import msprime
# Simulate a tree sequence with 2 diploid individuals (i.e. 4 samples)
ts = msprime.sim_ancestry(2, sequence_length=1e8, recombination_rate=1e-7, random_seed=1)
rank_counts = collections.Counter(t.rank() for t in ts.trees())
most_freq_rank, count = rank_counts.most_common(1)[0]
most_freq_topology = tskit.Tree.unrank(ts.num_samples, most_freq_rank)
print("Most frequent topology")
display(SVG(most_freq_topology.draw_svg(node_labels={0: 0, 1: 1, 2: 2, 3: 3})))
```

(sec_enumerating_topologies)=

### Enumerating Topologies

This section expands briefly on the approach used to enumerate
tree topologies that serves as the basis for {meth}`Tree.rank`
and {meth}`Tree.unrank`.
To enumerate all rooted, leaf-labelled tree topologies, we first
formulate a system of ordering and enumerating tree shapes. Then
we define an enumeration of labellings given an arbitrary tree shape.

#### Enumerating Tree Shapes

Starting with $n = 1$, we see that the only shape for a tree
with a single leaf is a single root leaf. A tree with $n > 1$
leaves can be obtained by joining at least two trees whose number of
leaves sum to $n$.
This maps very closely to the concept of integer partitions.
Each tree shape of $n$ leaves can be represented by taking a
nondecreasing integer partition of $n$ (elements of the partition
are sorted in nondecreasing order) and recursively partitioning its
elements. The order in which we select partitions of $n$ is
determined by the efficient
[rule_asc](http://jeromekelleher.net/generating-integer-partitions.html)
algorithm for generating them.

All tree shapes with four leaves, and the partitions that generate
them, are:

```{image} _static/four_leaf_tree_shapes.png
:alt: All four-leaf tree shapes and their generating partitions
```

Note that the middle column reflects all tree shapes of three leaves
in the right subtree!

`*` This excludes the partition $[n]$, since this would create a unary node
and trees with unary nodes are inumerable (and potentially infinite).

```{note}
Using nondecreasing integer partitions enforces a
*canonical orientation* on the tree shapes, where children under a node are
ordered by the number of leaves below them.
This is important because it prevents us from repeating trees that are
topologically the same but whose children are ordered differently.
```

#### Labelling Tree Shapes

Tree shapes are useful in and of themselves, but we can use the enumeration
formulated above to go further and assign labels to the leaves of each shape.

Say we are given a tree $T$ with $n$ leaves, whose left-most
subtree, $T_l$, has `k` leaves. For each of the $n \choose k$
ways to select labels to assign to $T_l$, we produce a unique labelling
of $T$. This process of choosing labels is repeated for the other
children of $T$ and then recursively for the subtrees.

Looking back to the example from {ref}`sec_tree_ranks`, we can see
the second tree shape, where the tree is a strictly bifucating tree of
three leaves, can be labelled in 3 different unique ways.

```{code-cell} ipython3
second_tree = tskit.Tree.unrank(num_leaves=3, rank=(1, 0))
for t in tskit.all_tree_labellings(second_tree):
    display(SVG(t.draw_svg(node_labels={0: 0, 1: 1, 2: 2}, order="tree", size=(120, 120))))
```

The order of the tree labellings is a direct result of the way in which
combinations of labels are chosen. The implementation in tskit uses a
standard lexicographic ordering to choose labels. See how the trees
are sorted by the order in which the left leaf's label was chosen.

```{note}
There is a caveat here regarding symmetry, similar to that of repeating
tree shapes. Symmetrical trees run the risk of creating redundant labellings
if all combinations of labels were exhausted. To prevent redundant labellings
we impose a *canonical labelling*. In the case of two symmetrical subtrees,
the left subtree must receive the minimum label from the label set. Notice
how this is the case in the right subtrees above.
```

These two enumerations create a complete ordering of topologies where trees are
ordered first by size (number of leaves), then by shape, then by their minimum
label. It is this canonical order that enables efficient ranking and unranking
of topologies.



--- ../../tskit/docs/terminology_and_concepts.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

```{currentmodule} tskit
```

```{code-cell} ipython3
:tags: [remove-cell]
import msprime
import tskit
import numpy as np

def basics():
    # Use "record full ARG" so we can get a nice SPR animation if we want it at the end
    # random_seed picked to get 2nd SPR involving 2 edges only (at the root)
    arg = msprime.sim_ancestry(
        3,
        population_size=100,
        random_seed=1440,
        sequence_length=1000,
        recombination_rate=1e-6,
        record_full_arg=True)
    arg = msprime.sim_mutations(arg, rate=1e-6, random_seed=23)
    ts = arg.simplify()
    tables = ts.dump_tables()
    # populations already has a schema
    populations_md = [p.metadata for p in ts.populations()]
    populations_md[0]["name"] = "GBR"
    populations_md[0]["description"] = "British from England and Scotland"
    
    tables.populations.packset_metadata(
        [tables.populations.metadata_schema.validate_and_encode_row(r) for r in populations_md]
    )
    individuals_md = [{"name": "Ada"}, {"name": "Bob"}, {"name": "Cat"}]
    tables.individuals.metadata_schema = tskit.MetadataSchema({'codec': 'json'})
    tables.individuals.packset_metadata(
        [tables.individuals.metadata_schema.validate_and_encode_row(r) for r in individuals_md]
    )
    tables.mutations.time = np.full_like(tables.mutations.time, tskit.UNKNOWN_TIME)
    tables.tree_sequence().dump("data/basics.trees")
    


def create_notebook_data():
    basics()

# create_notebook_data()  # uncomment to recreate the tree seqs used in this notebook
```

(sec_terminology_and_concepts)=

# Terminology & concepts

If you have read "{ref}`sec_what_is`" we hope to have piqued your interest in tree
sequences. This tutorial serves as an introduction to the terminology and some of the
concepts behind {program}`tskit`, the tree sequence toolkit.

(sec_terminology)=

## Terminology

::::{margin}
:::{note}
See {ref}`sec_intro_downloading_datafiles` to run this tutorial on your own computer
:::
::::

A tree sequence is a data structure which describes a set of correlated
evolutionary trees, together with some associated data that specifies, for example,
the location of mutations in the tree sequence. More technically, a tree sequence
stores a biological structure known as an "Ancestral Recombination Graph", or ARG.

Below are the most important {ref}`terms and concepts <tskit:sec_data_model_definitions>`
that you'll encounter in these tutorials,  but first we'll {func}`~tskit.load` a tree
sequence from a `.trees` file using the
{ref}`tskit:sec_python_api` (which will be used in the rest of this tutorial):

```{code-cell} ipython3
import tskit
# We will often store the python tree sequence object in a variable named "ts"
ts = tskit.load("data/basics.trees")
```

### Genomes

::::{margin}
:::{note}
{ref}`Workarounds<msprime:sec_ancestry_multiple_chromosomes>` exist
to represent a multi-chromosome genome as a tree
sequence, but are not covered here
:::
::::

A tree sequence covers a contiguous section of a single chromosome.
For clarity of explanation, we use the word "genome" to refer to the section spanned by
the whole tree sequence (if you find it easier, imagine our examples describe an
organism with a single, short chromosome).

In the tree sequence we have just loaded,
the genome is of length 1000, conventionally interpreted as the number of base
pairs. This is accessible using the {attr}`~TreeSequence.sequence_length` attribute.

```{code-cell} ipython3
ts.sequence_length  # a floating point number: positions need not be integers
```

It can be helpful to visualize the tree sequence as a set of local trees along the genome:

```{code-cell} ipython3
:"tags": ["hide-input"]
ts.draw_svg(
    y_axis=True,
    y_gridlines=True,
    time_scale="log_time",
    y_ticks=[0, 3, 10, 30, 100, 300,1000],
)
```

Each tree records the lines of descent along which a piece of DNA has been
inherited (ignore for the moment the red symbols, which represent a mutation).
For example, the first tree tells us that DNA from ancestral genome 7 duplicated
to produce two lineages, which ended up in genomes 1 and 4, both of which exist in the
current population. In fact, since this pattern is seen in all trees, these particular
lines of inheritance were taken by all the DNA in this 1000 base pair genome.


(sec_terminology_nodes)=

### Nodes

Trees in a tree sequence are represented in a slightly unusual way.
In the majority of libraries dealing with trees, each node is represented
as an object in memory and the relationship between nodes as pointers between these
objects. In tree sequences, however, all nodes, both internal and terminal,
are represented by an **integer ID**, unique over the entire tree sequence, and which exists
at a specific point in time. A branch point in any of the trees is associated with
an *internal node*, representing an ancestor in which a single DNA
sequence was duplicated (in forwards-time terminology) or in which multiple sequences
coalesced (in backwards-time terminology). 


(sec_terminology_nodes_samples)=

#### Sample nodes

The tip (leaf) nodes of the tree are usually marked as *sample nodes*, plotted as square
symbols in the visualization. Sample nodes are the focal nodes in a tree sequence; they
normally represent genomes from which we have obtained data (i.e. which we have "sampled"),
and which therefore have a fully known genetic sequence. There are 6 sample nodes,
labelled $0..5$, and also 6 non-sample nodes, labelled $6..11$, in the tree sequence above:

```{code-cell} ipython3
print("There are", ts.num_nodes, "nodes, of which", ts.num_samples, "are sample nodes")
```


(sec_terminology_edges)=

### Edges

In an evolutionary tree, the connections between nodes are conventionally referred to as
"branches". In a tree sequence, we instead refer to directional connections between nodes
as _edges_, emphasizing that they usually persist across multiple trees. Thus each of the
three trees in the example above has a branch from node 7 to node 1, but those three
branches represent just a single edge.

Each edge is associated with a parent node ID and a child node ID. The time of the parent
node must be
strictly greater than the time of the child node, and the difference in these times is
sometimes referred to as the "length" of the edge. Since trees in a tree sequence are
usually taken to represent marginal trees along a genome, as well as the time dimension
each edge also has a genomic _span_, defined by a *left* and a *right* position
along the genome. There are 15 edges in the tree sequence above. Here's an example of
one of them:

```{code-cell} ipython3
print("There are", ts.num_edges, "edges, one of which is:\n ", ts.edge(10))
```

(the `metadata` value is discussed below)

:::{note}
The *number of edges* largely determines the entire size of a tree sequence file. Edges
essentially encode ancestral information, and the number in a tree sequence can
be used as a rough measure of many things such as the compressibility of a simulation,
the speed with which analyses are likely to run, or even the effectiveness of ancestral
inference. To give some sense of scale, a basic inferred tree sequence of human
chromosome 20 in the UK Biobank dataset of a million genomes consists of 62 million edges
which define 16 thousand trees; the uncompressed tree sequence takes 2.1 gigabytes of
disk space, and it takes of the order of a few seconds to iterate over all the trees.
:::

(sec_terminology_trees_and_tree_sequences)=

### Trees and tree sequences

The combination of nodes and edges results in a tree sequence. This can be visualized,
as in the drawing above, as a set of correlated trees. The
sample nodes, $0..5$ in the drawing, are present in all the trees (since we have their
full genomes), but the other nodes, such as node $9$, need not be: indeed in larger
tree sequences they are rarely so.

In tree sequence terminology, we don't explictly keep track of where nodes
start and end. Only edges (not nodes) possess a genomic span. So for example, this tree
sequence is defined using edges like $(7\rightarrow1)$ which span the entire genome,
edges like $(11\rightarrow10)$ which only span the leftmost section of the genome, and
edges like $(11\rightarrow8)$ which span the rightmost part of the genome. Every time an
edge changes it creates a new tree; the location of the change is known as a _breakpoint_:

```{code-cell} ipython3
breakpoints = ts.breakpoints(as_array=True)
print("There are", ts.num_trees, "trees, associated with breakpoints", breakpoints)
```

::::{margin}
More details about these and related methods are given in the {ref}`sec_analysing_trees`
and {ref}`Getting started<sec_tskit_getting_started>` tutorials.
::::

Note that in keeping with the reliance on node numbers (rather than objects), information
about the relationships between nodes in a tree is obtained using their **integer IDs**.
In the Python API, for instance, the {meth}`TreeSequence.first`
method will return the first tree, which then provides basic methods to obtain
information about (say) node $7$, including the IDs of its parent and child nodes:

```{code-cell} ipython3
first_tree = ts.first()
parent_of_7 = first_tree.parent(7)
children_of_7 = first_tree.children(7)
print("Node 7's parent is", parent_of_7, "and childen are", children_of_7, "in the first tree")
```




(sec_terminology_individuals_and_populations)=

### Individuals and populations

Sometimes we know that one or more nodes belong to a particular
*{ref}`individual<tskit:sec_nodes_or_individuals>`*. For example,
most organisms are diploid, so if we sequence the DNA from an individual we are likely to
obtain two copies of each autosomal chromosome. The tree with six sample nodes above
could therefore represent the result of sampling three diploid individuals from a larger
population. The tree sequence can keep track of the individuals in which nodes reside,
and store specific information about them (such as the individuals' spatial location)
as well as arbitrary {ref}`metadata <sec_metadata>` (such as a name). In this particular
tree sequence the sample nodes are indeed associated with three named diploid
individuals: ``Ada``, ``Bob`` and ``Cat``.

```{code-cell} ipython3
print(f"There are {ts.num_individuals} individuals defined in the tree sequence:")
for individual in ts.individuals():
    print("*", f"Individual {individual.id}: metadata={individual.metadata}")
```

Adding these individual names to the sample node labels show that indeed, each of the
three individuals has 2 sample nodes (note that these 2 nodes do not always cluster
together in the trees, as the original simulation was of a randomly mating population)

```{code-cell} ipython3
:"tags": ["hide-input"]
node_labels = {node.id: str(node.id) for node in ts.nodes()}
for n in ts.samples():
    node_labels[n] = f"{n} ({ts.individual(ts.node(n).individual).metadata['name']})"
styles = (
    ".node > .lab {font-size: 70%}"
    ".leaf > .lab {text-anchor: start; transform: rotate(90deg) translate(6px)}"
)
ts.draw_svg(
    node_labels=node_labels,
    style=styles,
    y_axis=True,
    y_gridlines=True,
    time_scale="log_time",
    y_ticks=[0, 3, 10, 30, 100, 300,1000],
)
```

In the same way that nodes can be associated with a specific individual, nodes can also
be associated with a larger _population_. In this example, there is only one population
(to which all the nodes belong):

```{code-cell} ipython3
print(f"{ts.num_populations} population(s) defined:")
for population in ts.populations():
    print("*", population)
```


(sec_terminology_mutations_and_sites)=

### Mutations and sites

In {program}`tskit`, _mutations_ occur at specific, defined _sites_. A mutation is
positioned above a particular node (specified by the node ID), occurs at a particular
site (specified by a site ID), and involves a change of allelic state to a defined
{attr}`~Mutation.derived_state`. A mutation therefore can be specified something like
this:

```{code-cell} ipython3
print(f"{ts.num_mutations} mutation(s) in the tree sequence:")
for mutation in ts.mutations():
    print("*", mutation)
```

The mutation can have a {attr}`~Mutation.time` or if, as in this case, the times of
mutations in the tree sequence are unknown, all mutations can have the special NaN value
{data}`tskit.UNKNOWN_TIME`. Notice that the genomic position of the mutation is not
included. Instead, that is a property of the _site_ to which the mutation refers, in
this case, site ID 0 (which happens to be at position 751):

```{code-cell} ipython3
print(ts.site(0))  # For convenience, the Python API also returns the mutations at the site
```

In the plot above, since the the only mutation is above node 8 in the last tree, and has
a {attr}`~Mutation.derived_state` of "G", we know that the samples descending from node
8 in the last tree (sample genomes 2, 3, and 5) have a "G" at {attr}`~Site.position` 751,
while the others have the {attr}`~Site.ancestral_state` of "T". This means that Ada is
homozygous for "T", Bob is homozygous for "G", and Cat is heterozygous "T/G".
In other words the ancestral state and the details of any mutations at that site,
when coupled with the tree topology at the site {attr}`~Site.position`, is sufficient to
define the allelic state possessed by each sample.

Note that even though the genome is 1000 base pairs long, the tree sequence only contains
a single site, because we usually only bother defining *variable* sites in a tree
sequence (e.g. positions seen in studies to have samples possessing different alleles at
that genomic location). It is perfectly possible to have a site with no mutations
(or silent mutations) --- i.e. a "monomorphic" site --- but such sites are not normally
used in further analysis.


(sec_terminology_provenance)=

### Provenance

The history of the tree sequence is recorded in a set of provenance entries. That
includes what produced it, and what operations have been done to it and when. For
instance, if it was produced by {ref}`simplifying<sec_simplification>` a tree
sequence produced by the msprime simulator, the first provenance entry would record the
call to msprime that produced it, and the second the call to
{meth}`~TreeSequence.simplify` that was done on the result. Ideally, the list of
provenance entries are sufficient to exactly recreate the tree sequence, but this
is not always possible.


(sec_concepts)=

## Concepts

There are some basic concepts which can be helpful with thinking about tree sequences.
For reference, here is the tree sequence topology that we have been using:

```{code-cell} ipython3
:"tags": ["hide-input"]
ts.draw_svg(
    y_axis=True,
    y_gridlines=True,
    time_scale="log_time",
    y_ticks=[0, 3, 10, 30, 100, 300,1000],
    style=".mut, .site {display: none}",
)
```

Note that all these trees show strictly bifurcating splits, but this does not need to be
the case for tree sequences in general. In particular, where tree sequences have been
created by the multiple-merger coalecent process, or in tree
sequences that have been inferred from real data, it is possible to have a parent node
with 3 or more children in a particular tree (these are known as *polytomies*).


(sec_concepts_sprs)=

### Tree changes, ancestral recombinations, and SPRs

The process of recombination usually results in trees along a genome where adjacent
trees differ by only a few "tree edit" or SPR (subtree-prune-and-regraft) operations.
The result is a tree sequence in which very few edges
{ref}`change from tree to tree<fig_what_is_edge_diffs>`.
This is the underlying reason that `tskit` is so
efficient, and is well illustrated in the example tree sequence above.

In this (simulated) tree sequence, each tree differs from the next by a single SPR.
The subtree defined by node 7 in the first tree has been pruned and regrafted onto the
branch between 0 and 10, to create the second tree. The second and third trees have the
same topology, but differ because their ultimate coalesence happened in a different
ancestor (easy to spot in a simulation, but hard to detect in real data). This is also
caused by a single SPR: looking at the second tree, either the subtree below node 8 or
the subtree below node 9 must have been pruned and regrafted higher up on the same
lineage to create the third tree. Because this is a fully {ref}`simplified<sec_simplification>`
tree sequence, it is impossible to tell which of these two possible SPRs occurred. To
know this, we would need to have a tree sequence with the exact history of recombinations
recorded (see below).

In general, each detectable recombination occurring in ancestral history results in a
single SPR in the tree sequence. If recombination breakpoints occurs at unique
positions (an "infinite sites" model of breakpoints), then the number of trees in a tree
sequence equals the number of ancestral recombination events plus one. If recombinations
can occur at the same physical position (e.g. if the genome is treated as a set of
discrete integer positions, as in the simulation that created this tree sequence) then
moving from  one tree to the next in a tree sequence might require multiple SPRs if
there are multiple, overlaid ancestral recombination events.

(sec_concepts_args)=

### Tree sequences and ARGs

::::{margin}
:::{note}
There is a subtle distinction between common ancestry and coalescence. In particular, all coalescent nodes are common ancestor events, but not all common ancestor events in an ARG result in coalescence in a local tree.
:::
::::

The term "Ancestral Recombination Graph", or ARG, is commonly used to describe a genetic
genealogy. In particular, many (but not all) authors use it to mean a genetic
genealogy in which details of the position and potentially the timing of all
recombination and common ancestor events are explictly stored. For clarity
we refer to this sort of genetic genealogy as a "full ARG". Succinct tree sequences can
represent many different sorts of ARGs, including "full ARGs", by incorporating extra
non-coalescent nodes (see the {ref}`sec_args` tutorial). However, tree sequences are
often shown and stored in {ref}`fully simplified<sec_simplification>` form,
which omits these extra nodes. This is for two main reasons:

1. Many recombination events are undetectable from sequence data, and even if they are
   detectable, they can be logically impossible to place in the genealogy (as in the
   second SPR example above).
2. The number of recombination and non-coalescing common ancestor events in the genealogy
   quickly grows to dominate the total number of nodes in the tree sequence,
   without actually contributing to the mutations inherited by the samples.
   In other words, these nodes are redundant to the storing of genome data.

Therefore, compared to a full ARG, you can think of a simplified tree sequence as
storing the trees *created by* recombination events, rather than attempting to record the
recombination events themselves. The actual recombination events can be sometimes be
inferred from these trees but, as we have seen, it's not always possible. Here's another
way to put it:

> "an ARG encodes the events that occurred in the history of a sample,
> whereas a [simplified] tree sequence encodes the outcome of those events"
> ([Kelleher _et al._, 2019](https://doi.org/10.1534/genetics.120.303253))


### Tables

The underlying data in a tree sequence are stored in a collection of tables. Details of
these tables, as well as some basic properties of the tree sequence which they encode,
are shown when printing a tree sequence to the screen:

```{code-cell} ipython3
ts  # or use `print(ts)` for the plain text representation
```

There is a separate API for dealing with a tree sequence as a collection of tables,
which can be useful for dealing with large datasets, and is needed if you want to
{ref}`edit<sec_tables_editing>` an existing tree sequence. This is covered in
the {ref}`sec_tables` tutorial.




--- ../../tskit/docs/installation.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

```{currentmodule} tskit
```

(sec_installation)=


# Installation

There are two basic options for installing `tskit`: either through
pre-built binary packages using {ref}`sec_installation_conda` or
by compiling locally using {ref}`sec_installation_pip`. We recommend using `conda`
for most users, although `pip` can be more convenient in certain cases.
Tskit is installed to provide succinct tree sequence functionality
to other software (such as [msprime](https://github.com/tskit-dev/msprime)),
so it may already be installed if you use such software.

(sec_installation_requirements)=


## Requirements

Tskit requires Python 3.8+. There are no external C library dependencies. Python
dependencies are installed automatically by `pip` or `conda`.

(sec_installation_conda)=


## Conda

Pre-built binary packages for `tskit` are available through
[conda](https://conda.io/docs/), and built using [conda-forge](https://conda-forge.org/).
Packages for recent version of Python are available for Linux, OSX and Windows. Install
using:

```bash
$ conda install -c conda-forge tskit
```

### Quick Start

1. Install `conda` using [miniconda ](https://conda.io/miniconda.html).
   Make sure you follow the instructions to fully activate your `conda`
   installation!
2. Set up the [conda-forge channel ](https://conda-forge.org/) using
   `conda config --add channels conda-forge`.
3. Install tskit: `conda install tskit`.
4. Try it out: `tskit --version`.


There are several different ways to obtain `conda`. Please see the
[anaconda installation documentation](https://docs.anaconda.com/anaconda/install/)
for full details.

(sec_installation_pip)=


## Pip

Installing using `pip` is somewhat more flexible than `conda` and
may result in code that is (slightly) faster on your specific hardware.
`Pip` is the recommended method when using the system provided Python
installations. Installation is straightforward:

```bash
$ python3 -m pip install tskit
```

(sec_installation_development_versions)=


## Development versions

For general use, we do not recommend installing development versions.
Occasionally pre-release versions are made available, which can be
installed using `python3 -m pip install --pre tskit`. If you really need to install a
bleeding-edge version, see {ref}`sec_development_installing`.


--- ../../tskit/docs/convert_changelog.py ---

import re
import sys

SUBS = [
    (r":user:`([A-Za-z0-9-]*)`", r"[@\1](https://github.com/\1)"),
    (r":pr:`([0-9]*)`", r"[#\1](https://github.com/tskit-dev/tskit/issues/\1)"),
    (r":issue:`([0-9]*)`", r"[#\1](https://github.com/tskit-dev/tskit/issues/\1)"),
]


def process_log(log):
    delimiters_seen = 0
    for line in log:
        if line.startswith("-------"):
            delimiters_seen += 1
            continue
        if delimiters_seen == 3:
            return
        if delimiters_seen % 2 == 0:
            for pattern, replace in SUBS:
                line = re.sub(pattern, replace, line)
            yield line


with open(sys.argv[1]) as f:
    print("".join(process_log(f.readlines())))


--- ../../tskit/docs/data-model.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

:::{currentmodule} tskit
:::


(sec_data_model)=

# Data model

The `tskit` library deals with sets of sampled genome sequences through storage
and analysis of their shared genetic ancestry. This genealogical ancestry (sometimes
known as an Ancestral Recombination Graph) is stored concisely in `tskit` in the 
"succinct tree sequence" format, which comprises a collection of easy-to-understand
tables. This page documents the structure of the tables and encoding of table data,
as well as the encoding of the correlated genetic trees that can be extracted from
a `tskit` tree sequence.

We begin by defining the the structure of the tables in the {ref}`sec_table_definitions`
section. The {ref}`sec_data_model_data_encoding` section then describe how data is
stored in those tables (also see the {ref}`sec_file_formats`
chapter). The {ref}`sec_data_model_tree_structure` section then
describes the encoding of the trees that are generated from the {class}`NodeTable`
and {class}`EdgeTable`. Finally, we describe how genotype data arises from tree
structure, especially how we can incorporate the idea of missing data.

(sec_table_definitions)=

## Table definitions

(sec_table_types_definitions)=

### Table types

A tree sequence can be stored in a collection of eight tables:
{ref}`Node <sec_node_table_definition>`,
{ref}`Edge <sec_edge_table_definition>`,
{ref}`Individual <sec_individual_table_definition>`,
{ref}`Site <sec_site_table_definition>`,
{ref}`Mutation <sec_mutation_table_definition>`,
{ref}`Migration <sec_migration_table_definition>`,
{ref}`Population <sec_population_table_definition>`, and
{ref}`Provenance <sec_provenance_table_definition>`.
The Node and Edge tables store the genealogical
relationships that define the trees, and the Individual table
describes how multiple genomes are grouped within individuals;
the Site and Mutation tables describe where mutations fall
on the trees; the Migration table describes how lineages move across space;
and the Provenance table contains information on where the data came from.
Only Node and Edge tables are necessary to encode the genealogical trees;
Sites and Mutations are optional but necessary to encode polymorphism
(sequence) data; the remainder are optional.
In the following sections we define these components of a tree sequence in
more detail.

(sec_node_table_definition)=

#### Node Table

A **node** defines a monoploid set of chromosomes (a "genome") of a specific
individual that was born at some time in the past: the set of
chromosomes inherited from a particular one of the individual's parents.
(See {ref}`sec_nodes_or_individuals` for more discussion.)
Every vertex in the marginal trees of a tree sequence corresponds
to exactly one node, and a node may be present in many trees. The
node table contains five columns, of which `flags` and `time` are
mandatory:


| Column        |  Type       |  Description                           |
| :------------ | ----------- | -------------------------------------: |
| flags         |  uint32     |  Bitwise flags.                        |
| time          |  double     |  Birth time of node.                   |
| population    |  int32      |  Birth population of node.             |
| individual    |  int32      |  The individual the node belongs to.   |
| metadata      |  binary     |  Node {ref}`sec_metadata_definition`.  |

The `time` column records the birth time of the individual in question,
and is a floating point value. Similarly,
the `population` column records the ID of the population where this
individual was born. If not provided, `population` defaults to the
null ID (-1). Otherwise, the population ID must refer to a row in the
{ref}`sec_population_table_definition`.
The `individual` column records the ID of the
{ref}`Individual <sec_individual_table_definition>`
individual that this node belongs to. If specified, the ID must refer
to a valid individual. If not provided, `individual`
defaults to the null ID (-1).

The `flags` column stores information about a particular node, and
is composed of 32 bitwise boolean values. Currently, the only flag defined
is `NODE_IS_SAMPLE = 1`, which defines the *sample* status of nodes. Marking
a particular node as a "sample" means, for example, that the mutational state
of the node will be included in the genotypes produced by
{meth}`TreeSequence.variants`.

Bits 0-15 (inclusive) of the `flags` column are reserved for internal use by
`tskit` and should not be used by applications for anything other
than the purposes documented here. Bits 16-31 (inclusive) are free for applications
to use for any purpose and will not be altered or interpreteted by
`tskit`.

See the {ref}`sec_node_requirements` section for details on the properties
required for a valid set of nodes.

For convenience, the {ref}`text format <sec_text_file_format>` for nodes
decomposes the `flags` value into its separate values. Thus, in the
text format we have a column for `is_sample`, which corresponds to the
`flags` column in the underlying table. As more flags values are
defined, these will be added to the text file format.

The `metadata` column provides a location for client code to store
information about each node. See the {ref}`sec_metadata_definition` section for
more details on how metadata columns should be used.

:::{note}
The distinction between `flags` and `metadata` is that flags
holds information about a node that the library understands, whereas
metadata holds information about a node that the library *does not*
understand. Metadata is for storing auxiliarly information that is
not necessary for the core tree sequence algorithms.
:::


(sec_individual_table_definition)=

#### Individual Table

An **individual** defines how nodes (which can be seen
as representing single chromosomes) group together in a polyploid individual.
The individual table contains three columns, of which only `flags` is mandatory.

| Column        | Type       | Description                                |
| :------------ | ---------- | -----------------------------------------: |
| flags         | uint32     | Bitwise flags.                             |
| location      | double     | Location in arbitrary dimensions.          |
| parents       | int32      | Ids of parent individuals.                 |
| metadata      | binary     | Individual {ref}`sec_metadata_definition`. |

See the {ref}`sec_individual_requirements` section for details on the properties
required for a valid set of individuals.

The `flags` column stores information about a particular individual, and
is composed of 32 bitwise boolean values. Currently, no flags are
defined.

Bits 0-15 (inclusive) of the `flags` column are reserved for internal use by
`tskit` and should not be used by applications for anything other
than the purposes documented here. Bits 16-31 (inclusive) are free for applications
to use for any purpose and will not be altered or interpreteted by
`tskit`.

The `location` column stores the location of an individual in arbitrary
dimensions. This column is {ref}`ragged <sec_encoding_ragged_columns>`, and
so different individuals can have locations with different dimensions (i.e.,
one individual may have location `[]` and another `[0, 1, 0]`. This could
therefore be used to store other quantities (e.g., phenotype).

The `parents` column stores the ids of other individuals that are the parents of
an individual. This can be used to store pedigree information for individuals.
This column is {ref}`ragged <sec_encoding_ragged_columns>` such that an
individual can have any number of parents.

The `metadata` column provides a location for client code to store
information about each individual. See the {ref}`sec_metadata_definition` section for
more details on how metadata columns should be used.

:::{note}
The distinction between `flags` and `metadata` is that flags
holds information about a individual that the library understands, whereas
metadata holds information about a individual that the library *does not*
understand. Metadata is for storing auxiliarly information that is
not necessary for the core tree sequence algorithms.
:::


(sec_edge_table_definition)=

#### Edge Table

An **edge** defines a parent-child relationship between a pair of nodes
over a specific sequence interval. The edge table contains five columns,
all of which are mandatory except `metadata`:

| Column        | Type       | Description                                |
| :------------ | ---------- | -----------------------------------------: |
| left          | double     | Left coordinate of the edge (inclusive).   |
| right         | double     | Right coordinate of the edge (exclusive).  |
| parent        | int32      | Parent node ID.                            |
| child         | int32      | Child node ID.                             |
| metadata      | binary     | Node {ref}`sec_metadata_definition`.       |

Each row in an edge table describes a half-open genomic interval `[left, right)`
over which the `child` inherited from the given `parent`.
The `left` and `right` columns are defined using double precision
floating point values. The `parent` and `child`
columns specify integer IDs in the associated {ref}`sec_node_table_definition`.

The `metadata` column provides a location for client code to store
information about each edge. See the {ref}`sec_metadata_definition` section for
more details on how metadata columns should be used.

See the {ref}`sec_edge_requirements` section for details on the properties
required for a valid set of edges.


(sec_site_table_definition)=

#### Site Table

A **site** defines a particular location along the genome in which
we are interested in observing the allelic state. The site table
contains three columns, of which `position` and `ancestral_state`
are mandatory.

| Column          | Type       | Description                                |
| :-------------- | ---------- | -----------------------------------------: |
| position        | double     | Position of site in genome coordinates.    |
| ancestral_state | text       | The state at the root of the tree.         |
| metadata        | binary     | Site {ref}`sec_metadata_definition`.       |

The `position` column is a floating point value defining the location
of the site in question along the genome.

The `ancestral_state` column specifies the allelic state at the root
of the tree, thus defining the state that nodes inherit if no mutations
intervene. The column stores text character data of arbitrary length.

The `metadata` column provides a location for client code to store
information about each site. See the {ref}`sec_metadata_definition` section for
more details on how metadata columns should be used.

See the {ref}`sec_site_requirements` section for details on the properties
required for a valid set of sites.


(sec_mutation_table_definition)=

#### Mutation Table

A **mutation** defines a change of allelic state on a tree at a particular site.
The mutation table contains five columns, of which `site`, `node` and
`derived_state` are mandatory.

| Column          | Type       | Description                                    |
| :-------------- | ---------- | ---------------------------------------------: |
| site            | int32      | The ID of the site the mutation occurs at.     |
| node            | int32      | The node this mutation occurs at.              |
| parent          | int32      | The ID of the parent mutation.                 |
| time            | double     | Time at which the mutation occurred.           |
| derived_state   | char       | The allelic state resulting from the mutation. |
| metadata        | binary     | Mutation {ref}`sec_metadata_definition`.       |

The `site` column is an integer value defining the ID of the
{ref}`site <sec_site_table_definition>` at which this mutation occurred.

The `node` column is an integer value defining the ID of the
first {ref}`node <sec_node_table_definition>` in the tree below this mutation.

The  `time` column is a double precision floating point value recording how long ago
the mutation happened.

The `derived_state` column specifies the allelic state resulting from the mutation,
thus defining the state that the `node` and any descendant nodes in the
subtree inherit unless further mutations occur. The column stores text
character data of arbitrary length.

The `parent` column is an integer value defining the ID of the mutation whose
allelic state this mutation replaced. If there is no mutation at the
site in question on the path back to root, then this field is set to the
null ID (-1). (The `parent` column is only required in situations
where there are multiple mutations at a given site. For
"infinite sites" mutations, it can be ignored.)

The `metadata` column provides a location for client code to store
information about each site. See the {ref}`sec_metadata_definition` section for
more details on how metadata columns should be used.

See the {ref}`sec_mutation_requirements` section for details on the properties
required for a valid set of mutations.


(sec_migration_table_definition)=

#### Migration Table

In simulations, trees can be thought of as spread across space, and it is
helpful for inferring demographic history to record this history.
Migrations are performed by individual ancestors, but most likely not by an
individual whose genome is tracked as a `node` (as in a discrete-deme model they are
unlikely to be both a migrant and a most recent common ancestor).  So,
`tskit` records when a segment of ancestry has moved between
populations. This table is not required, even if different nodes come from
different populations.

| Column     | Type     | Description                                            |
| :--------- | -------- | -----------------------------------------------------: |
| left       | double   | Left coordinate of the migrating segment (inclusive).  |
| right      | double   | Right coordinate of the migrating segment (exclusive). |
| node       | int32    | Node ID.                                               |
| source     | int32    | Source population ID.                                  |
| dest       | int32    | Destination population ID.                             |
| time       | double   | Time of migration event.                               |
| metadata   | binary   | Migration {ref}`sec_metadata_definition`.              |


The `left` and `right` columns are floating point values defining the
half-open segment of genome affected. The `source` and `dest` columns
record the IDs of the respective populations. The `node` column records the
ID of the node that was associated with the ancestry segment in question
at the time of the migration event. The `time` column is holds floating
point values recording the time of the event.

The `metadata` column provides a location for client code to store
information about each migration. See the {ref}`sec_metadata_definition` section for
more details on how metadata columns should be used.

See the {ref}`sec_migration_requirements` section for details on the properties
required for a valid set of mutations.


(sec_population_table_definition)=

#### Population Table

A **population** defines a grouping of individuals that a node can
be said to belong to.

The population table contains one column, `metadata`.

| Column     | Type     | Description                                |
| :--------- | -------- | -----------------------------------------: |
| metadata   | binary   | Population {ref}`sec_metadata_definition`. |


The `metadata` column provides a location for client code to store
information about each population. See the {ref}`sec_metadata_definition` section for
more details on how metadata columns should be used.

See the {ref}`sec_population_requirements` section for details on the properties
required for a valid set of populations.


(sec_provenance_table_definition)=

#### Provenance Table

:::{todo}
Document the provenance table.
:::

| Column    | Type  | Description                                                             |
| :-------- | ----- | ----------------------------------------------------------------------: |
| timestamp | char  | Timestamp in [ISO-8601](https://en.wikipedia.org/wiki/ISO_8601) format. |
| record    | char  | Provenance record.                                                      |


(sec_metadata_definition)=

### Metadata

Each table (excluding provenance) has a metadata column for storing and passing along
information that tskit does not use or interpret. See {ref}`sec_metadata` for details.
The metadata columns are {ref}`binary columns <sec_tables_api_binary_columns>`.

When using the {ref}`sec_text_file_format`, to ensure that metadata can be safely
interchanged, each row is [base 64 encoded](https://en.wikipedia.org/wiki/Base64).
Thus, binary information can be safely printed and exchanged, but may not be
human readable.

The tree sequence itself also has metadata stored as a byte array.


(sec_valid_tree_sequence_requirements)=

### Valid tree sequence requirements

Arbitrary data can be stored in tables using the classes in the
{ref}`sec_tables_api`. However, only a {class}`TableCollection`
that fulfils a set of requirements represents
a valid {class}`TreeSequence` object which can be obtained
using the {meth}`TableCollection.tree_sequence` method. In this
section we list these requirements, and explain their rationale.
Violations of most of these requirements are detected when the
user attempts to load a tree sequence via {func}`tskit.load` or
{meth}`TableCollection.tree_sequence`, raising an informative
error message. Some more complex requirements may not be detectable at load-time,
and errors may not occur until certain operations are attempted.
These are documented below.

The Python API also provides tools that can transform a collection of
tables into a valid collection of tables, so long as they are logically
consistent, see {ref}`sec_tables_api_creating_valid_tree_sequence`.


(sec_individual_requirements)=

#### Individual requirements

Individuals are a basic type in a tree sequence and are not defined with
respect to any other tables. Individuals can have a reference to their parent
individuals, if present these references must be valid or null (-1).

A valid tree sequence does not require individuals to be sorted in any
particular order, and sorting a set of tables using {meth}`TableCollection.sort`
has no effect on individuals. However, individuals can be optionally sorted
using {meth}`TableCollection.sort_individuals`.

(sec_node_requirements)=

#### Node requirements

Given a valid set of individuals and populations, the requirements for
each node are:

- `population` must either be null (-1) or refer to a valid population ID;
- `individual` must either be null (-1) or refer to a valid individual ID.

An ID refers to a zero-indexed row number in the relevant table,
and so is "valid" if is between 0 and one less than the number of rows in the relevant table.

There are no requirements regarding the ordering of nodes with respect to time.

Sorting a set of tables using {meth}`TableCollection.sort`
has no effect on nodes.


(sec_edge_requirements)=

#### Edge requirements

Given a valid set of nodes and a sequence length {math}`L`, the simple
requirements for each edge are:

- We must have {math}`0 \leq` `left` {math}`<` `right` {math}`\leq L`;
- `parent` and `child` must be valid node IDs;
- `time[parent]` > `time[child]`;
- edges must be unique (i.e., no duplicate edges are allowed).

The first requirement simply ensures that the interval makes sense. The
third requirement ensures that we cannot have loops, since time is
always increasing as we ascend the tree.

To ensure a valid tree sequence there is one further requirement:

- The set of intervals on which each node is a child must be disjoint.

This guarantees that we cannot have contradictory edges (i.e.,
where a node `a` is a child of both `b` and `c`), and ensures that
at each point on the sequence we have a well-formed forest of trees.

In the interest of algorithmic efficiency, edges must have the following
sortedness properties:

- All edges for a given parent must be contiguous;
- Edges must be listed in nondecreasing order of `parent` time;
- Within the edges for a given `parent`, edges must be sorted
  first by `child` ID and then by `left` coordinate.

Violations of these requirements are detected at load time.
The {meth}`TableCollection.sort` method will ensure that these sortedness
properties are fulfilled.


(sec_site_requirements)=

#### Site requirements

Given a valid set of nodes and a sequence length {math}`L`, the simple
requirements for a valid set of sites are:

- We must have {math}`0 \leq` `position` {math}`< L`;
- `position` values must be unique.

For simplicity and algorithmic efficiency, sites must also:

- Be sorted in increasing order of `position`.

Violations of these requirements are detected at load time.
The {meth}`TableCollection.sort` method ensures that sites are sorted
according to these criteria.


(sec_mutation_requirements)=

#### Mutation requirements

Given a valid set of nodes, edges and sites, the
requirements for a valid set of mutations are:

- `site` must refer to a valid site ID;
- `node` must refer to a valid node ID;
- `time` must either be `UNKNOWN_TIME` (a NAN value which indicates
  the time is unknown) or be a finite value which is greater or equal to the
  mutation `node`'s `time`, less than the `node` above the mutation's
  `time` and equal to or less than the `time` of the `parent` mutation
  if this mutation has one. If one mutation on a site has `UNKNOWN_TIME` then
  all mutations at that site must, as a mixture of known and unknown is not valid.
- `parent` must either be the null ID (-1) or a valid mutation ID within the
  current table

Furthermore,

- If another mutation occurs on the tree above the mutation in
  question, its ID must be listed as the `parent`.

For simplicity and algorithmic efficiency, mutations must also:

- be sorted by site ID;
- when there are multiple mutations per site, mutations should be ordered by
  decreasing time, if known, and parent mutations must occur
  **before** their children (i.e. if a mutation with ID {math}`x` has
  `parent` with ID {math}`y`, then we must have {math}`y < x`).

Violations of these sorting requirements are detected at load time.
The {meth}`TableCollection.sort` method ensures that mutations are sorted
according site ID, but does not at present enforce that mutations occur
after their parent mutations.

Silent mutations (i.e., mutations for which the ancestral and derived
states are the same) are allowed.
For example, if we have a site with ancestral state
of "A" and a single mutation with derived state "A", then this
mutation does not result in any change of state.
(This addition was made in release C_0.99.11.)

:::{note}
As `tskit.UNKNOWN_TIME` is implemented as a `NaN` value, tests for equality
will always fail. Use `tskit.is_unknown_time` to detect unknown values.
:::


(sec_migration_requirements)=

#### Migration requirements

Given a valid set of nodes and edges, the requirements for a value set of
migrations are:

- `left` and `right` must lie within the tree sequence coordinate space (i.e.,
  from 0 to `sequence_length`).
- `time` must be strictly between the time of its `node` and the time of any
  ancestral node from which that node inherits on the segment `[left, right)`.
- The `population` of any such ancestor matching `source`, if another
  `migration` does not intervene.

To enable efficient processing, migrations must also be:

- Sorted by nondecreasing `time` value.

Note in particular that there is no requirement that adjacent migration records
should be "squashed". That is, we can have two records `m1` and `m2`
such that `m1.right` = `m2.left` and with the `node`, `source`,
`dest` and `time` fields equal. This is because such records will usually
represent two independent ancestral segments migrating at the same time, and
as such squashing them into a single record would result in a loss of information.


(sec_population_requirements)=

#### Population requirements

There are no requirements on a population table.


(sec_provenance_requirements)=

#### Provenance requirements

The `timestamp` column of a provenance table should be in
[ISO-8601](https://en.wikipedia.org/wiki/ISO_8601) format.

The `record` should be valid JSON with structure defined in the Provenance
Schema section (TODO).


(sec_table_indexes)=

### Table indexes

To efficiently iterate over the trees in a tree sequence, `tskit` uses
indexes built on the edges. To create a tree sequence from a table collection
the tables must be indexed; the {meth}`TableCollection.build_index` method
can be used to create an index on a table collection if necessary.

:::{todo}
Add more details on what the indexes actually are.
:::

(sec_data_model_data_encoding)=

## Data encoding

In this section we describe the high-level details of how data is encoded in
tables. Tables store data in a **columnar** manner. In memory, each
table is organised as a number of blocks of contiguous storage, one for
each column. There are many advantages to this approach, but the key
property for us is that allows for very efficient transfer of data
in and out of tables. Rather than inserting data into tables row-by-row
(which can be done in Python
{ref}`using the add_row methods<sec_tables_api_accessing_table_data>`), it is much
more efficient to add many rows at the same time by providing pointers to blocks of
contiguous memory. By taking this approach, we can work with tables containing
gigabytes of data very efficiently.

For instance, in the {ref}`sec_python_api` we can use the
[numpy Array API](https://docs.scipy.org/doc/numpy-1.13.0/reference/arrays.html)
to allow us to define and work with numeric arrays of the required types.
Node IDs, for example, are defined using 32 bit integers. Thus, the
`parent` column of an {ref}`sec_edge_table_definition`'s with `n` rows
is a block `4n` bytes.

This approach is very straightforward for columns in which each row contains
a fixed number of values. However, dealing with columns containing a
**variable** number of values is more problematic.


(sec_encoding_ragged_columns)=

### Encoding ragged columns

A **ragged** column is a column in which the rows are not of a fixed length.
For example, {ref}`sec_metadata_definition` columns contain binary of data of arbitrary
length. To encode such columns in the tables API, we store **two** columns:
one contains the flattened array of data and another stores the **offsets**
of each row into this flattened array. Consider an example:

```{code-cell} ipython3
import tskit

s = tskit.SiteTable()
s.add_row(0, "A")
s.add_row(0, "")
s.add_row(0, "TTT")
s.add_row(0, "G")
s
```

In this example we create a {ref}`sec_site_table_definition` with four rows,
and then display this table. We can see that the second row has the
empty string as its `ancestral_state`, and the third row's
`ancestral_state` is `TTT`. Now let's print out the columns:

```{code-cell} ipython3
print("Ancestral state (numerical): ", s.ancestral_state)
print("Ancestral state (as bytes): ", s.ancestral_state.tobytes())
print("Ancestral state offsets: ", s.ancestral_state_offset)
```

When we print out the tables `ancestral_state`
column, we see that its a numpy array of length 5: this is the
flattened array of [ASCII encoded](https://en.wikipedia.org/wiki/ASCII)
values for these rows. When we decode these bytes using the
numpy {meth}`tobytes<numpy:numpy.ndarray.tobytes>` method, we get the string 'ATTTG'.
This flattened array can now be transferred efficiently in memory like any other column
We then use the `ancestral_state_offset` column to allow us find the individual rows.
For a row `j`:

    ancestral_state[ancestral_state_offset[j]: ancestral_state_offset[j + 1]]

gives us the array of bytes for the ancestral state in that row. For example, here is
row 2:

```{code-cell} ipython3
s.ancestral_state[s.ancestral_state_offset[2]: s.ancestral_state_offset[3]].tobytes()
```

For a table with `n` rows, any offset column must have `n + 1`
values, the first of which is always `0`. The values in this column must be
nondecreasing, and cannot exceed the length of the ragged column in question.


(sec_data_model_reference_sequence)=

## Reference sequence

Along with the topology and site information stored in the tskit tree
sequence, we can also optionally store an associated reference sequence.
Reference sequences are flexible, and can consist simply of some
metadata recording which assembly build a tree sequence uses, or
storing the entire sequence itself.

:::{warning}
Reference sequence support in tskit is preliminary. Reference sequence
data can be stored and accessed via the C API. Support in the Python
API is limited to usage in {meth}`.TreeSequence.alignments` and
related methods, where it provides the default values for nucleotide
positions between {ref}`sites<sec_data_model_definitions_site>`.
:::

(sec_data_model_tree_structure)=

## Tree structure

(sec_data_model_quintuply_linked_trees)=

### Quintuply linked trees

Tree structure in `tskit` is encoded internally as a "quintuply
linked tree", a generalisation of the triply linked tree encoding
used by Knuth and others. Nodes are represented by their integer
IDs, and their relationships to other nodes are recorded in the
`parent`, `left_child`, `right_child`, `left_sib` and
`right_sib` arrays. For example, consider the following tree
and its associated arrays:

```{code-cell} ipython3
:tags: ["hide-input"]
import io

import tskit
from IPython.display import SVG

nodes = """\
id      is_sample   time
0       1           0
1       1           0
2       1           0
3       1           0
4       1           0
5       0           1
6       0           2
7       0           3
"""
edges = """\
left    right   parent  child
0       60      5       4,3
0       40      6       2
0       60      6       1,0
20      40      6       5
0       20      7       5
40      60      7       5
0       60      7       6
40      60      7       2
"""
ts = tskit.load_text(
    nodes=io.StringIO(nodes), edges=io.StringIO(edges), strict=False
)

SVG(ts.first().draw_svg(time_scale="rank"))
```

```{code-cell} ipython3
:tags: ["hide-input"]
from IPython.display import HTML

def html_quintuple_table(ts, show_virtual_root=False, show_convenience_arrays=False):
    tree = ts.first()
    columns = ["node", "parent", "left_child", "right_child", "left_sib", "right_sib"]
    convenience_arrays = ["num_children", "edge"]
    if show_convenience_arrays:
        columns += convenience_arrays
    data = {k:[] for k in columns}
    for u in sorted(tree.nodes(tree.virtual_root if show_virtual_root else None)):
        for colname in columns:
            data[colname].append(u if colname == "node" else getattr(tree, colname)(u))
    html = "<tr>"
    for colname in columns:
        html += f"<th>{colname}</th>"
    html += "</tr>"
    for u in range(len(data["node"])):
        html += "<tr>" if u < ts.num_nodes else "<tr style='font-style: italic; color:red'>"
        for colname in columns:
            html += f"<td>{data[colname][u]}</td>"
        html += "</tr>"
    return "<table>" + html + "</table>"

HTML(html_quintuple_table(ts))
```

Each node in the tree corresponds to a row in this table, and
the columns are the individual arrays recording the quintuply linked
structure. Thus, we can see that the parent of nodes `0`, `1`, and `2`
is `6`. Similarly, the left child of `6` is `0` and the
right child of `6` is `2`. The `left_sib` and `right_sib` arrays
then record each nodes sibling on its left or right, respectively;
hence the right sib of `0` is `1`, and the right sib of `1` is `2`.
Thus, sibling information allows us to efficiently support trees
with arbitrary numbers of children. In each of the five pointer arrays,
the null node (-1) is used to indicate the end of a path; thus,
for example, the parent of `7` and left sib of `0` are null.

Please see this {ref}`example <sec_c_api_examples_tree_traversals>` for
details of how to use the quintuply linked structure in the C API.

:::{note}
For many applications we do not need the quintuply linked trees,
and (for example) the `left_sib` and `right_child` arrays can be
ignored. The reason for using a quintuply instead of triply linked
encoding is that it is not possible to efficiently update the trees
as we move along the sequence without the quintuply linked structure.
:::

:::{warning}
The left-to-right ordering of nodes is determined by the order
in which edges are inserted into the tree during iteration along the sequence.
Thus, if we arrive at the same tree by iterating from different directions,
the left-to-right ordering of nodes may be different! The specific
ordering of the children of a node should therefore not be depended on.
:::

### Convenience arrays
 
Similar to the five arrays representing the {ref}`quintuply linked tree<sec_data_model_quintuply_linked_trees>`,
convenience arrays track information on each node in the tree. These arrays are not essential to 
represent the trees within a tree sequence. However, they can be useful for
specific algorithms (e.g. when computing tree (im)balance metrics). Two 
convenience arrays have been implemented so far: 
{attr}`Tree.num_children_array` and {attr}`Tree.edge_array`.

Here is the table above with the convenience arrays also shown:

```{code-cell} ipython3
:tags: ["hide-input"]

HTML(html_quintuple_table(ts, show_convenience_arrays=True))
```

(sec_data_model_tree_roots)=

### Roots

In the `tskit` {class}`trees <Tree>` we have shown so far, all the sample nodes have
been connected to each other. This means each tree has only a single {attr}`~Tree.root`
(i.e. the oldest node found when tracing a path backwards in time from any sample).
However, a tree can contain {ref}`sec_data_model_tree_isolated_sample_nodes`
or unconnected topologies, and can therefore have *multiple* {attr}`~Tree.roots`.
Here's an example, created by deleting the edge joining `6` and `7` in the tree sequence
used above:

```{code-cell} ipython3
:tags: ["hide-input"]
tables = ts.dump_tables()
tables.edges.truncate(ts.num_edges - 1)
ts_multiroot = tables.tree_sequence()
SVG(ts_multiroot.first().draw_svg(time_scale="rank"))
```

In `tskit` terminology, this should *not* be thought
of as two separate trees, but as a single multi-root "tree", comprising
two unlinked topologies. This fits with the definition of a tree
in a tree sequence: a tree describes the ancestry of the same
fixed set of sample nodes at a single position in the genome. In the
picture above, *both* the left and right hand topologies are required
to describe the genealogy of samples 0..4 at this position.

Here's what the entire tree sequence now looks like:

```{code-cell} ipython3
:tags: ["hide-input"]
SVG(ts_multiroot.draw_svg(time_scale="rank"))
```

From the terminology above, it can be seen that this tree sequence consists of only
three trees (not five). The first tree, which applies from position 0 to 20, is the one
used in our example. As we saw, removing the edge connecting node 6 to node 7 has
created a tree with 2 roots (and thus 2 unconnected topologies in a single tree).
In contrast, the second tree, from position 20 to 40, has a single root. Finally the
third tree, from position 40 to 60, again has two roots.

(sec_data_model_tree_root_threshold)=

#### The root threshold

The roots of a tree are defined by reference to the
{ref}`sample nodes<sec_data_model_definitions_sample>`. By default, roots are the unique
endpoints of the paths traced upwards from the sample nodes; equivalently, each root
counts one or more samples among its descendants (or is itself a sample node). This is
the case when the {attr}`~Tree.root_threshold` property of a tree is left at its default
value of `1`. If, however, the `root_threshold` is (say) `2`, then a node is
considered a root only if it counts at least two samples among its descendants. Setting
an alternative `root_threshold` value can be used to avoid visiting
{ref}`sec_data_model_tree_isolated_sample_nodes`, for example when dealing with trees
containing {ref}`sec_data_model_missing_data`.

(sec_data_model_tree_virtual_root)=

#### The virtual root

To access all the {attr}`~Tree.roots` in a tree, tskit uses a special additional node
called the **virtual root**. This is primarily a bookkeeping device, and
can normally be ignored: it is not plotted in any visualizations and
does not exist as an independent node in the node table.
However, the virtual root can be useful in certain algorithms because its
children are defined as all the "real" roots in a tree. Hence by
descending downwards from the virtual root, it is possible
to access the entire genealogy at a given site, even in a multi-root
tree. In the quintuply linked tree encoding, the virtual root appears as an
extra element at the end of each of the tree arrays. Here's the same table
as before but with the virtual root also shown, using red italics to
emphasise that it is not a "real" node:

```{code-cell} ipython3
:tags: ["hide-input"]
HTML(html_quintuple_table(ts_multiroot, show_virtual_root=True))
```

You can see that the virtual root (node 8) has 6 as its left child and 7
as its right child. Importantly, though, this is an asymmetric
relationship: the parent of the "real" roots 6 and 7 is null
(-1) and *not* the virtual root. Hence when we ascend up the tree from the
sample nodes to their parents, we stop at the "real" roots, and never
encounter the virtual root.

Because the virtual root can be useful in some algorithms, it can
optionally be returned in traversal orders (see {meth}`.Tree.nodes`).
The following properties apply:

- All trees in a tree sequence share the same virtual root.
- The virtual root's ID is always equal to the number of nodes in the tree sequence
  (i.e. the length of the node table). However, there is **no corresponding row**
  in the node table, and any attempts to access information about the
  virtual root via either the tree sequence or tables APIs will fail with
  an out-of-bounds error.
- The parent and siblings of the virtual root are null.
- The time of the virtual root is defined as positive infinity (if
  accessed via {meth}`.Tree.time`). This is useful in defining the
  time-based node traversal orderings.
- The virtual root is the parent of no other node---roots do **not**
  have parent pointers to the virtual root.


(sec_data_model_tree_isolated_nodes)=

### Isolated nodes

In a tree, it is possible for a node to have no children and no parent. Such a node is
said to be *isolated*, meaning that we don't know anything about its relationships
over a specific genomic interval. This is commonly true for ancestral genomes, which
often have large regions that have
not been inherited by any of the {ref}`sample nodes<sec_data_model_definitions_sample>`
in the tree sequence, and therefore regions about which we know nothing. This is true,
for example, of node 7 in the middle tree of our previous example, which is why it is
not plotted on that tree:

```{code-cell} ipython3
display(SVG(ts_multiroot.draw_svg(time_scale="rank")))
for tree in ts_multiroot.trees():
    print(
        "Node 7",
        "is" if tree.is_isolated(7) else "is not",
        "isolated from position",
        tree.interval.left,
        "to",
        tree.interval.right,
    )
```


(sec_data_model_tree_isolated_sample_nodes)=

#### Isolated sample nodes

It is also possible for a {ref}`sample node<sec_data_model_definitions_sample>`
to be isolated. As long as the {ref}`root threshold<sec_data_model_tree_root_threshold>`
is set to its default value, an isolated *sample* node will count as a root, and
therefore be considered as being present on the tree (meaning it will be
returned by the {meth}`Tree.nodes`
and {meth}`Tree.samples` methods). When displaying a tree, isolated samples are shown
unconnected to other nodes. To illustrate, we can remove the edge from node 2 to node 7:

```{code-cell} ipython3
:tags: ["hide-input"]
tables = ts_multiroot.dump_tables()
tables.edges.set_columns(
    **tables.edges[(tables.edges.parent != 7) | (tables.edges.child != 2)].asdict())
ts_isolated = tables.tree_sequence()
SVG(ts_isolated.draw_svg(time_scale="rank"))
```

The rightmost tree now contains an isolated sample node (node 2), which counts as
one of the {ref}`sec_data_model_tree_roots` of the tree. This tree therefore has three
roots, one of which is node 2:

```{code-cell} ipython3
rightmost_tree = ts_isolated.at_index(-1)
print(rightmost_tree.num_roots, "roots in the rightmost tree, with IDs", rightmost_tree.roots)
print(
    "IDs of isolated samples in this tree:",
    [u for u in rightmost_tree.samples() if rightmost_tree.is_isolated(u)],
)
```

In `tskit`, isolated sample nodes are closely associated with the encoding of
{ref}`sec_data_model_missing_data`.


(sec_data_model_tree_dead_leaves_and_branches)=

### Dead leaves and branches

In a `tskit` tree, a *leaf node* is defined as a node without any children. The
implications of this turn out to be slighly unintuitive, and so are worth briefly
documenting here. Firstly, the same node can be a leaf in one tree, and not a leaf
in the next tree along the tree sequence. Secondly all isolated nodes must be leaves
(as by definition they have no children). Thirdly sample nodes need not be leaves
(they could be "internal samples"); likewise leaf nodes need not be samples.

Node 7 in the example above provides a good case study. Note that it is a root node with
at least one child (i.e. not a leaf) in trees 0 and 2; in contrast in tree 1 it is
isolated. Strictly, because it is isolated in tree 1, it is also a leaf node there,
although it is not attached to a root, not a sample, and is therefore not plotted. In
this case, in that tree we can think of node 7 as a "dead leaf" (and we don't normally
plot dead leaves). In fact, in a large tree sequence of many trees, most ancestral nodes
will be isolated in any given tree, and therefore most nodes in such a tree will be of
this sort. However, these dead leaves are excluded from most calculations on trees,
because algorithms usually traverse the tree by starting at a root and working down,
or by starting at a sample and working up. Hence when we refer to the leaves of a tree,
it is usually shorthand for the leaves **on** the tree (that is, attached via branches,
to one of the the tree roots). Dead leaves are excluded from this definition.

Note that it is also possible to have trees in which there are "dead branches": that is
sections of topology which are not accessible from a root, and whose tips are all
dead leaves. Although valid, this is a relatively unusual state of affairs, and such
branches are not plotted by the standard {ref}`sec_tskit_viz` methods. The
{meth}`Tree.nodes` method will not, by default, traverse through dead branches, although
it can be made to do so by specifying the ID of a dead node as the root for traversal.


(sec_data_model_genetic_data)=

## Encoding genetic variation

Genetic variation is incorporated into a tree sequence by placing
{ref}`mutations<sec_mutation_table_definition>` at
{ref}`sites<sec_mutation_table_definition>` along the genome.
The genotypes of the different samples at each site
can be found by using the tree to calculate which mutations are inherited by
the different samples. This is the fundamental basis of how tree sequences
efficiently encode DNA sequences, and is explained in depth elsewhere
(e.g. {ref}`in the tutorials<sec_what_is_dna_data>`).

Below, we discuss some implications of this encoding in more detail, in particular
the way in which it can be used to model missing data.

(sec_data_model_missing_data)=

### Missing data

If, at a particular genomic position, a node is
{ref}`isolated<sec_data_model_tree_isolated_nodes>` *and* additionally has
no mutations directly above it, its genotype at that position is considered to be
unknown (however, if there is a mutation above an isolated node, it
can be thought of as saying directly what the genotype is, and so renders the
genotype at that position not missing).

By way of illustration, we'll use the {meth}`~TableCollection.delete_intervals` method
to remove all knowledge of the ancestry in the
middle portion of the previous example (say from position 15 to 45) sprinkle
on some mutations, and make sure there are sites at every position:

```{code-cell} ipython3
:tags: ["hide-input"]

import numpy as np
import msprime

tables = msprime.sim_mutations(ts_isolated, rate=0.1, random_seed=123).dump_tables()
tables.delete_intervals([[15, 45]], simplify=False)
missing_sites = np.setdiff1d(np.arange(tables.sequence_length), tables.sites.position)
for pos in missing_sites:
    tables.sites.add_row(position=pos, ancestral_state="A")  # Add sites at every pos
tables.sort()
missing_ts = tables.tree_sequence()
SVG(missing_ts.draw_svg())
```


The middle section of the genome now has no ancestry at all, and therefore for any site
that is in this region, the genotypic state that it is assigned is a special value
`tskit.MISSING_DATA`, or `-1`. The {meth}`~TreeSequence.haplotypes()` method, which
outputs the actual allelic state for each sample, defaults to outputting an `N` at
these sites. Therefore where any sample node is isolated, the haplotype will show
an `N`, indicating the DNA sequence is unknown. This will be so not only in the
middle of all of the sample genomes, but also at the right hand end of the genome of
sample 2, as it is an isolated sample node in the rightmost tree:

```{code-cell} ipython3
for i, h in enumerate(missing_ts.haplotypes()):
    print(f"Sample {i}: {h}")
```

See the {meth}`TreeSequence.variants` method and {class}`Variant` class for
more information on how missing data is represented in variant data.


(sec_gotchas)=

## Possibly surprising consequences of the data model

This is a section of miscellaneous issues that might trip even an experienced user up,
also known as "gotchas".
The current examples are quite uncommon, so can be ignored for most purposes,
but the list may be expanded in the future.

### Unrelated material

Usually, all parts of a tree sequence are ancestral to at least one sample,
since that's essentially the definition of a sample: the genomes that
we're describing the ancestry of.
However, in some cases there will be portions of the tree sequence from which
no samples inherit - notably, the result of a forwards simulation that has
not been simplified.
In fact, if the simulation has not coalesced,
one can have entire portions of some marginal tree that are
unrelated to any of the samples
(for instance, an individual in the initial generation of the simulation
that had no offspring).
This can lead to a gotcha:
the *roots* of a tree are defined to be only those roots *reachable from the samples*
(and, furthermore, reachable from at least `root_threshold` samples;
see {meth}`TreeSequence.trees`).
So, our unlucky ancestor would not appear in the list of `roots`, even though
if we drew all the relationships provided by the tree sequence,
they'd definitely be a root.
Furthermore, only nodes *reachable from a root* are included in the
{meth}`Tree.nodes`. So, if you iterate over all the nodes in each marginal tree,
you won't see those parts of the tree sequence that are unrelated to the samples.
If you need to get those, too, you could either
work with the {meth}`TreeSequence.edge_diffs` directly,
or iterate over all nodes (instead of over {meth}`Tree.nodes`).



--- ../../tskit/docs/analysing_tree_sequences.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

```{currentmodule} tskit
```

```{code-cell} ipython3
:tags: [remove-cell]
import msprime
import numpy as np
import tskit

def computing_statistics():
    ts = msprime.simulate(
        10**4, Ne=10**4, recombination_rate=1e-8, mutation_rate=1e-8, length=10**7, random_seed=42)
    ts.dump("data/computing_statistics.trees")
    
def afs():
    ts = msprime.simulate(6, mutation_rate=1, random_seed=47)
    # remove the mutation times so the plot is nicer
    tables = ts.dump_tables()
    tables.mutations.time = np.full_like(tables.mutations.time, tskit.UNKNOWN_TIME)
    ts = tables.tree_sequence()
    ts.dump("data/afs.trees")


def create_notebook_data():
    computing_statistics()
    afs()

# create_notebook_data()  # uncomment to recreate the tree seqs used in this notebook
```

(sec_analysing_tree_sequences)=

# _Analysing tree sequences_
% remove underscores in title when tutorial is complete or near-complete

:::{note}
This tutorial is a work in progress.
:::


(sec_tutorial_stats)=

## Computing statistics

Tskit provides an extensive and flexible interface for computing population
genetic statistics, which is documented in detail in the
{ref}`general statistics <sec_stats>` section of the offical documentation.
This tutorial aims to give a quick overview of how the APIs work and how to use
them effectively.

First, let's load a tree sequence to work with which has roughly human
parameters for 10 thousand samples and 10Mb chromosomes:

```{code-cell} ipython3
ts = tskit.load("data/computing_statistics.trees")
ts
```

This tree sequence has ~36.6 thousand trees & ~39 thousand segregating sites.
We'd now like to compute some statistics on this dataset.

### One-way statistics

We refer to statistics that are defined with respect to a single set of
samples as "one-way". An example of such a statistic is diversity, which
is computed using the {meth}`TreeSequence.diversity` method:

```{code-cell} ipython3
d = ts.diversity()
print("Average diversity per unit sequence length = {d:.3G}")
```

This tells the average diversity across the whole sequence and returns a single
number. We'll usually want to compute statistics in
{ref}`windows <tskit:sec_stats_windows>` along the genome and we
use the ``windows`` argument to do this:

```{code-cell} ipython3
windows = np.linspace(0, ts.sequence_length, num=5)
d = ts.diversity(windows=windows)
print(windows, d, sep="\n")
```

The ``windows`` argument takes a numpy array specifying the breakpoints
along the genome. Here, we use numpy to create four equally spaced windows
of size 2.5 megabases (the windows array contains k + 1 elements to define
k windows). Because we have asked for values in windows, tskit now returns
a numpy array rather than a single value. (See
{ref}`sec_stats_output_dimensions` for a full description of how the output
dimensions of statistics are determined by the ``windows`` argument.)

Suppose we wanted to compute diversity within a specific subset of samples.
We can do this using the ``sample_sets`` argument:

```{code-cell} ipython3
A = ts.samples()[:100]
d = ts.diversity(sample_sets=A)
print(d)
```

Here, we've computed the average diversity within the first hundred samples across
the whole genome. As we've not specified any windows, this is again a single value.

We can also compute diversity in *multiple* sample sets at the same time by providing
a list of sample sets as an argument:

```{code-cell} ipython3
A = ts.samples()[:100]
B = ts.samples()[100:200]
C = ts.samples()[200:300]
d = ts.diversity(sample_sets=[A, B, C])
print(d)
```

Because we've computed multiple statistics concurrently, tskit returns a numpy array
of these statistics. We have asked for diversity within three different sample sets,
and tskit therefore returns an array with three values. (In general, the
dimensions of the input determines the dimensions of the output: see
{ref}`tskit:sec_stats_output_dimensions` for a detailed description of the rules.)

We can also compute multiple statistics in multiple windows:

```{code-cell} ipython3
d = ts.diversity(sample_sets=[A, B, C], windows=windows)
print("shape = ", d.shape, "\n", d)
```

We have computed diversity within three different sample sets across four
genomic windows, and our output is therefore a 2D numpy array with four
rows and three columns: each row contains the diversity values within
A, B and C for a particular window.

### Multi-way statistics

Many population genetic statistics compare multiple sets of samples to
each other. For example, the {meth}`TreeSequence.divergence` method computes
the divergence between two subsets of samples:

```{code-cell} ipython3
A = ts.samples()[:100]
B = ts.samples()[:100]
d = ts.divergence([A, B])
print(d)
```

The divergence between two sets of samples A and B is a single number,
and we we again return a single floating point value as the result. We can also
compute this in windows along the genome, as before:

```{code-cell} ipython3
d = ts.divergence([A, B], windows=windows)
print(d)
```

Again, as we have defined four genomic windows along the sequence, the result is
numpy array with four values.

A powerful feature of tskit's stats API is that we can compute the divergences
between multiple sets of samples simultaneously using the ``indexes`` argument:


```{code-cell} ipython3
d = ts.divergence([A, B, C], indexes=[(0, 1), (0, 2)])
print(d)
```

Here, we've specified three sample sets A, B and C and we've computed the
divergences between A and B,  and between A and C. The ``indexes`` argument is used
to specify which pairs of sets we are interested in. In this example
we've computed two different divergence values and the output is therefore
a numpy array of length 2.

As before, we can combine computing multiple statistics in multiple windows
to return a 2D numpy array:

```{code-cell} ipython3
windows = np.linspace(0, ts.sequence_length, num=5)
d = ts.divergence([A, B, C], indexes=[(0, 1), (0, 2)], windows=windows)
print(d)
```

Each row again corresponds to a window, which contains the average divergence
values between the chosen sets.

If the ``indexes`` parameter is 1D array, we interpret this as specifying
a single statistic and remove the empty outer dimension:

```{code-cell} ipython3
d = ts.divergence([A, B, C], indexes=(0, 1))
print(d)
```

It's important to note that we don't **have** to remove empty dimensions: tskit
will only do this if you explicitly ask it to. Here, for example, we can keep the
output as an array with one value if we wish:

```
d = ts.divergence([A, B, C], indexes=[(0, 1)])
print(d)
```

Please see {ref}`tskit:sec_stats_sample_sets` for a
full description of the ``sample_sets`` and ``indexes`` arguments.

(sec_tutorial_afs)=

## Allele frequency spectra

The allele frequency spectrum is a fundamental tool in population genetics, and
tskit provides a flexible and powerful approach to computing such spectra.
Suppose we have simulated the following tree sequence:

```{code-cell} ipython3
from IPython.display import display
ts = tskit.load("data/afs.trees")
tree = ts.first()
display(tree.draw_svg())
ts.tables.sites
```

Computing the allele frequency spectrum is then easy:

```{code-cell} ipython3
afs = ts.allele_frequency_spectrum(polarised=True, span_normalise=False)
print(afs)
```

This tells us that we have two singletons, six doubletons and one 3-ton and
one 4-ton. Note that the first element of the returned AFS array does *not* correspond
to the singletons (see below for why). Because we have simulated these mutations,
we know the ancestral and derived states we have set ``polarised`` to True. We
can get the "folded" AFS by setting polarised to False. Because we want simple
counts here and not averaged values, we set ``span_normalise=False``: by
default, windowed statistics are divided by the sequence length, so they are
comparable between windows.

The returned value here is actually a 2D array, and this is because we can
also perform these computations in windows along the genome:

```{code-cell} ipython3
afs = ts.allele_frequency_spectrum(windows=[0, 0.5, 1], span_normalise=False, polarised=True)
print(afs)
```

This time, we've asked for the number of sites at each frequency in two
equal windows. Now we can see that in the first half of the sequence we
have three sites (compare with the site table above): one singleton,
one doubleton and one tripleton.

### Joint spectra

We can also compute allele frequencies within multiple sets of samples,
the *joint allele frequency spectra*.

```{code-cell} ipython3
node_colours = {0: "blue", 2: "blue", 3: "blue", 1: "green", 4: "green", 5: "green"}
styles = [f".n{k} > .sym {{fill: {v}}}" for k, v in node_colours.items()]
tree.draw_svg(style = "".join(styles))
```

Here we've marked the samples as either blue or green (we can imagine
these belonging to different populations, for example). We can then compute
the joint AFS based on these two sets:

```{code-cell} ipython3
afs = ts.allele_frequency_spectrum([[0, 2, 3], [1, 4, 5]], polarised=True)
print(afs)
```

Now, each window in our AFS is a 2D numpy array, where each dimension
corresponds to frequencies within the different sets. So, we see for example
that there are six sites that are singletons in both sets, 1 site
that is a doubleton in both sets, and 2 sites that singletons in $[1, 4, 5]$
and not present in the other sample set.

### Branch length spectra

Up to now we've used the {meth}`~TreeSequence.allele_frequency_spectrum` method
to summarise the number of sites that occur at different frequencies. We can also
use this approach to compute the total branch lengths subtending a given
number of samples by setting ``mode="branch"``:

```{code-cell} ipython3
afs = ts.allele_frequency_spectrum(mode="branch", polarised=True, span_normalise=False)
print(afs)
```

Thus, the total branch length over example one sample is 4.86, over two is
5.39, and so on.


(sec_tutorial_afs_zeroth_entry)=

### Zeroth and final entries in the AFS

The zeroth element of the AFS is significant when we are working with
sample sets that are a subset of all samples in the tree sequence.
For example, in the following we compute the AFS within the sample set
[0, 1, 2]:

```{code-cell} ipython3
afs = ts.allele_frequency_spectrum([[0, 1, 2]], mode="branch", polarised=True)
print(afs)
```

Thus, the total branch length over 0, 1 and 2 is 5.3, and over pairs from this set
is 5.25. What does the zeroth value of 4.33 signify? This is the total branch length
over all samples that are **not** in this sample set. By including this value, we
maintain the property that for each tree, the sum of the AFS for any sample set
is always equal to the total branch length. For example, here we compute:

```{code-cell} ipython3
print("sum afs          = ", np.sum(afs))
print("total branch len = ", tree.total_branch_length)
```

The final entry of the AFS is similar: it counts alleles (for mode="site") or
branches (for mode="branch") that are ancestral to all of the given sample set,
but are still polymorphic in the entire set of samples of the tree sequence.
Note, however, that alleles fixed among all the samples, e.g., ones above
the root of the tree, will not be included.


--- ../../tskit/docs/glossary.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

:::{currentmodule} tskit
:::


(sec_glossary)=

# Glossary

(sec_data_model_definitions)=

## Definitions

Here are some definitions of some key ideas encountered in this documentation.

(sec_data_model_definitions_tree)=

tree
: A "gene tree", i.e., the genealogical tree describing how a collection of
  genomes (usually at the tips of the tree) are related to each other at some
  chromosomal location. See {ref}`sec_nodes_or_individuals` for discussion
  of what a "genome" is.

(sec_data_model_definitions_tree_sequence)=

tree sequence
: A "succinct tree sequence" (or tree sequence, for brevity) is an efficient
  encoding of a sequence of correlated trees, such as one encounters looking
  at the gene trees along a genome. A tree sequence efficiently captures the
  structure shared by adjacent trees, (essentially) storing only what differs
  between them.

(sec_data_model_definitions_node)=

node
: Each branching point in each tree is associated with a particular genome
  in a particular ancestor, called a "node".  Since each node represents a
  specific genome it has a unique `time`, thought of as its birth time,
  which determines the height of any branching points it is associated with.
  See {ref}`sec_nodes_or_individuals` for discussion of what a "node" is.

(sec_data_model_definitions_individual)=

individual
: In certain situations we are interested in how nodes (representing
  individual homologous genomes) are grouped together into individuals
  (e.g. two nodes per diploid individual). For example, when we are working
  with polyploid samples it is useful to associate metadata with a specific
  individual rather than duplicate this information on the constituent nodes.
  See {ref}`sec_nodes_or_individuals` for more discussion on this point.

(sec_data_model_definitions_sample)=

sample
: The focal nodes of a tree sequence, usually thought of as those from which
  we have obtained data. The specification of these affects various
  methods: (1) {meth}`TreeSequence.variants` and
  {meth}`TreeSequence.haplotypes` will output the genotypes of the samples,
  and {attr}`Tree.roots` only return roots ancestral to at least one
  sample.
  (This can be checked with {meth}`~Node.is_sample`;
  see the {ref}`node table definitions <sec_node_table_definition>`
  for information on how the sample
  status a node is encoded in the `flags` column.)

(sec_data_model_definitions_edge)=

edge
: The topology of a tree sequence is defined by a set of **edges**. Each
  edge is a tuple `(left, right, parent, child)`, which records a
  parent-child relationship among a pair of nodes on the
  on the half-open interval of chromosome `[left, right)`.

(sec_data_model_definitions_site)=

site
: Tree sequences can define the mutational state of nodes as well as their
  topological relationships. A **site** is thought of as some position along
  the genome at which variation occurs. Each site is associated with
  a unique position and ancestral state.

(sec_data_model_definitions_mutation)=

mutation
: A mutation records the change of state at a particular site 'above'
  a particular node (more precisely, along the branch between the node
  in question and its parent). Each mutation is associated with a specific
  site (which defines the position along the genome), a node (which defines
  where it occurs within the tree at this position), and a derived state
  (which defines the mutational state inherited by all nodes in the subtree
  rooted at the focal node). In more complex situations in which we have
  back or recurrent mutations, a mutation must also specify its 'parent'
  mutation.

(sec_data_model_definitions_migration)=

migration
: An event at which a parent and child node were born in different populations.

(sec_data_model_definitions_population)=

population
: A grouping of nodes, e.g., by sampling location.

(sec_data_model_definitions_provenance)=

provenance
: An entry recording the origin and history of the data encoded in a tree sequence.

(sec_data_model_definitions_ID)=

ID
: In the set of interconnected tables that we define here, we refer
  throughout to the IDs of particular entities. The ID of an
  entity (e.g., a node) is defined by the position of the corresponding
  row in the table. These positions are zero indexed. For example, if we
  refer to node with ID zero, this corresponds to the node defined by the
  first row in the node table.

(sec_data_model_definitions_sequence_length)=

sequence length
: This value defines the coordinate space in which the edges and site positions
  are defined. This is most often assumed to be equal to the largest
  `right` coordinate in the edge table, but there are situations in which
  we might wish to specify the sequence length explicitly.

## Further discussion

(sec_nodes_or_individuals)=

### Nodes, Genomes, or Individuals?

The natural unit of biological analysis is (usually) the *individual*. However,
many organisms we study are diploid, and so each individual contains *two*
homologous copies of the entire genome, separately inherited from the two
parental individuals. Since each monoploid copy of the genome is inherited separately,
each diploid individual lies at the end of two distinct lineages, and so will
be represented by *two* places in any given genealogical tree. This makes it
difficult to precisely discuss tree sequences for diploids, as we have no
simple way to refer to the bundle of chromosomes that make up the "copy of the
genome inherited from one particular parent". For this reason, in this
documentation we use the non-descriptive term "node" to refer to this concept
-- and so, a diploid individual is composed of two nodes -- although we use the
term "genome" at times, for concreteness.

Several properties naturally associated with individuals are in fact assigned
to nodes in what follows: birth time and population. This is for two reasons:
First, since coalescent simulations naturally lack a notion of polyploidy, earlier
versions of `tskit` lacked the notion of an individual. Second, ancestral
nodes are not naturally grouped together into individuals -- we know they must have
existed, but have no way of inferring this grouping, so in fact many nodes in
an empirically-derived tree sequence will not be associated with individuals,
even though their birth times might be inferred.



--- ../../tskit/docs/python-api.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

```{currentmodule} tskit
```

(sec_python_api)=

# Python API

This page documents the full tskit Python API. Brief thematic summaries of common
classes and methods are presented first. The {ref}`sec_python_api_reference` section
at the end then contains full details which aim to be concise, precise and exhaustive.
Note that this may not therefore be the best place to start if you are new
to a particular piece of functionality.

(sec_python_api_trees_and_tree_sequences)=

## Trees and tree sequences

The {class}`TreeSequence` class represents a sequence of correlated
evolutionary trees along a genome. The {class}`Tree` class represents a
single tree in this sequence. These classes are the interfaces used to
interact with the trees and mutational information stored in a tree sequence,
for example as returned from a simulation or inferred from a set of DNA
sequences.


(sec_python_api_tree_sequences)=

### {class}`TreeSequence` API


(sec_python_api_tree_sequences_properties)=

#### General properties

```{eval-rst}
.. autosummary::
  TreeSequence.time_units
  TreeSequence.nbytes
  TreeSequence.sequence_length
  TreeSequence.max_root_time
  TreeSequence.discrete_genome
  TreeSequence.discrete_time
  TreeSequence.metadata
  TreeSequence.metadata_schema
  TreeSequence.reference_sequence
```

#### Efficient table column access

The {class}`.TreeSequence` class provides access to underlying numerical
data defined in the {ref}`data model<sec_data_model>` in two ways:

1. Via the {attr}`.TreeSequence.tables` property and the
    {ref}`Tables API<sec_tables_api_accessing_table_data>`
2. Via a set of properties on the ``TreeSequence`` class that provide
   direct and efficient access to the underlying memory.

:::{warning}
Accessing table data via {attr}`.TreeSequence.tables` can be very inefficient
at the moment because accessing the `.tables` property incurs a **full copy**
of the data model. While we intend to implement this as a read-only view
in the future, the engineering involved is nontrivial, and so we recommend
using the properties listed here like ``ts.nodes_time`` in favour of
``ts.tables.nodes.time``.
Please see [issue #760](https://github.com/tskit-dev/tskit/issues/760)
for more information.
:::


```{eval-rst}
.. autosummary::
  TreeSequence.individuals_flags
  TreeSequence.nodes_time
  TreeSequence.nodes_flags
  TreeSequence.nodes_population
  TreeSequence.nodes_individual
  TreeSequence.edges_left
  TreeSequence.edges_right
  TreeSequence.edges_parent
  TreeSequence.edges_child
  TreeSequence.sites_position
  TreeSequence.mutations_site
  TreeSequence.mutations_node
  TreeSequence.mutations_parent
  TreeSequence.mutations_time
  TreeSequence.migrations_left
  TreeSequence.migrations_right
  TreeSequence.migrations_right
  TreeSequence.migrations_node
  TreeSequence.migrations_source
  TreeSequence.migrations_dest
  TreeSequence.migrations_time
  TreeSequence.indexes_edge_insertion_order
  TreeSequence.indexes_edge_removal_order
```

(sec_python_api_tree_sequences_loading_and_saving)=

#### Loading and saving

There are several methods for loading data into a {class}`TreeSequence`
instance. The simplest and most convenient is the use the {func}`tskit.load`
function to load a {ref}`tree sequence file <sec_tree_sequence_file_format>`. For small
scale data and debugging, it is often convenient to use the {func}`tskit.load_text`
function to read data in the {ref}`text file format<sec_text_file_format>`.
The {meth}`TableCollection.tree_sequence` function
efficiently creates a {class}`TreeSequence` object from a
{class}`collection of tables<TableCollection>`
using the {ref}`Tables API <sec_tables_api>`.

```{eval-rst}
Load a tree sequence
    .. autosummary::
      load
      load_text
      TableCollection.tree_sequence

Save a tree sequence
    .. autosummary::
      TreeSequence.dump
```

:::{seealso}
Tree sequences with a single simple topology can also be created from scratch by
{ref}`generating<sec_python_api_trees_creating>` a {class}`Tree` and accessing its
{attr}`~Tree.tree_sequence` property.
:::

(sec_python_api_tree_sequences_obtaining_trees)=

#### Obtaining trees

The following properties and methods return information about the
{class}`trees<Tree>` that are generated along a tree sequence.

```{eval-rst}
.. autosummary::

  TreeSequence.num_trees
  TreeSequence.trees
  TreeSequence.breakpoints
  TreeSequence.coiterate
  TreeSequence.first
  TreeSequence.last
  TreeSequence.aslist
  TreeSequence.at
  TreeSequence.at_index
```

#### Obtaining other objects

(sec_python_api_tree_sequences_obtaining_other_objects)=

Various components make up a tree sequence, such as nodes and edges, sites and
mutations, and populations and individuals. These can be counted or converted into
Python objects using the following classes, properties, and methods.

##### Tree topology

```{eval-rst}
Nodes
    .. autosummary::
      Node
      TreeSequence.num_nodes
      TreeSequence.nodes
      TreeSequence.node
      TreeSequence.num_samples
      TreeSequence.samples

Edges
    .. autosummary::
      Edge
      TreeSequence.num_edges
      TreeSequence.edges
      TreeSequence.edge
```

##### Genetic variation

```{eval-rst}
Sites
    .. autosummary::
      Site
      TreeSequence.num_sites
      TreeSequence.sites
      TreeSequence.site
      Variant
      TreeSequence.variants
      TreeSequence.genotype_matrix
      TreeSequence.haplotypes
      TreeSequence.alignments

Mutations
    .. autosummary::
      Mutation
      TreeSequence.num_mutations
      TreeSequence.mutations
      TreeSequence.mutation
```

##### Demography

```{eval-rst}
Populations
    .. autosummary::
      Population
      TreeSequence.num_populations
      TreeSequence.populations
      TreeSequence.population

Migrations
    .. autosummary::
      Migration
      TreeSequence.num_migrations
      TreeSequence.migrations
      TreeSequence.migration
```

##### Other

```{eval-rst}
Individuals
    .. autosummary::
      Individual
      TreeSequence.num_individuals
      TreeSequence.individuals
      TreeSequence.individual


Provenance entries (also see :ref:`sec_python_api_provenance`)
    .. autosummary::
      Provenance
      TreeSequence.num_provenances
      TreeSequence.provenances
      TreeSequence.provenance
```

(sec_python_api_tree_sequences_modification)=

#### Tree sequence modification

Although tree sequences are immutable, several methods will taken an existing tree
sequence and return a modifed version. These are thin wrappers around the
{ref}`identically named methods of a TableCollection<sec_tables_api_modification>`,
which perform the same actions but modify the {class}`TableCollection` in place.

```{eval-rst}
.. autosummary::
  TreeSequence.simplify
  TreeSequence.subset
  TreeSequence.union
  TreeSequence.keep_intervals
  TreeSequence.delete_intervals
  TreeSequence.delete_sites
  TreeSequence.trim
  TreeSequence.split_edges
  TreeSequence.decapitate
  TreeSequence.extend_haplotypes
```

(sec_python_api_tree_sequences_ibd)=

#### Identity by descent

The {meth}`.TreeSequence.ibd_segments` method allows us to compute
identity relationships between pairs of samples. See the
{ref}`sec_identity` section for more details and examples
and the {ref}`sec_python_api_reference_identity` section for
API documentation on the associated classes.

```{eval-rst}
.. autosummary::
  TreeSequence.ibd_segments
```

(sec_python_api_tree_sequences_tables)=

#### Tables

The underlying data in a tree sequence is stored in a
{ref}`collection of tables<sec_tables_api>`. The following methods give access
to tables and associated functionality. Since tables can be modified, this
allows tree sequences to be edited: see the {ref}`sec_tables` tutorial for
an introduction.

```{eval-rst}
.. autosummary::
  TreeSequence.tables
  TreeSequence.dump_tables
  TreeSequence.table_metadata_schemas
  TreeSequence.tables_dict
```


(sec_python_api_tree_sequences_statistics)=

#### Statistics

```{eval-rst}

Single site
    .. autosummary::
      TreeSequence.allele_frequency_spectrum
      TreeSequence.divergence
      TreeSequence.diversity
      TreeSequence.f2
      TreeSequence.f3
      TreeSequence.f4
      TreeSequence.Fst
      TreeSequence.genealogical_nearest_neighbours
      TreeSequence.genetic_relatedness
      TreeSequence.genetic_relatedness_weighted
      TreeSequence.genetic_relatedness_vector
      TreeSequence.general_stat
      TreeSequence.segregating_sites
      TreeSequence.sample_count_stat
      TreeSequence.mean_descendants
      TreeSequence.Tajimas_D
      TreeSequence.trait_correlation
      TreeSequence.trait_covariance
      TreeSequence.trait_linear_model
      TreeSequence.Y2
      TreeSequence.Y3

Comparative
    .. autosummary::
      TreeSequence.kc_distance
```

(sec_python_api_tree_sequences_topological_analysis)=

#### Topological analysis

The topology of a tree in a tree sequence refers to the relationship among
samples ignoring branch lengths. Functionality as described in
{ref}`sec_topological_analysis` is mainly provided via
{ref}`methods on trees<sec_python_api_trees_topological_analysis>`, but more
efficient methods sometimes exist for entire tree sequences:

```{eval-rst}
.. autosummary::
  TreeSequence.count_topologies
```

(sec_python_api_tree_sequences_display)=

#### Display

```{eval-rst}
.. autosummary::
  TreeSequence.draw_svg
  TreeSequence.draw_text
  TreeSequence.__str__
  TreeSequence._repr_html_
```


(sec_python_api_tree_sequences_export)=

#### Export
```{eval-rst}
.. autosummary::
  TreeSequence.as_fasta
  TreeSequence.as_nexus
  TreeSequence.dump_text
  TreeSequence.to_macs
  TreeSequence.write_fasta
  TreeSequence.write_nexus
  TreeSequence.write_vcf
```


(sec_python_api_trees)=

### {class}`Tree<Tree>` API

A tree is an instance of the {class}`Tree` class. These trees cannot exist
independently of the {class}`TreeSequence` from which they are generated.
Usually, therefore, a {class}`Tree` instance is created by
{ref}`sec_python_api_tree_sequences_obtaining_trees` from an existing tree
sequence (although it is also possible to generate a new instance of a
{class}`Tree` belonging to the same tree sequence using {meth}`Tree.copy`).

:::{note}
For efficiency, each instance of a {class}`Tree` is a state-machine
whose internal state corresponds to one of the trees in the parent tree sequence:
{ref}`sec_python_api_trees_moving_to` in the tree sequence does not require a
new instance to be created, but simply the internal state to be changed.
:::

(sec_python_api_trees_general_properties)=

#### General properties


```{eval-rst}
.. autosummary::
  Tree.tree_sequence
  Tree.total_branch_length
  Tree.root_threshold
  Tree.virtual_root
  Tree.num_edges
  Tree.num_roots
  Tree.has_single_root
  Tree.has_multiple_roots
  Tree.root
  Tree.roots
  Tree.index
  Tree.interval
  Tree.span
```


(sec_python_api_trees_creating)=

#### Creating new trees

It is sometimes useful to create an entirely new tree sequence consisting
of just a single tree (a "one-tree sequence"). The follow methods create such an
object and return a {class}`Tree` instance corresponding to that tree.
The new tree sequence to which the tree belongs is available through the
{attr}`~Tree.tree_sequence` property.

```{eval-rst}
Creating a new tree
    .. autosummary::
      Tree.generate_balanced
      Tree.generate_comb
      Tree.generate_random_binary
      Tree.generate_star

Creating a new tree from an existing tree
    .. autosummary::
      Tree.split_polytomies
```

:::{seealso}
{meth}`Tree.unrank` for creating a new one-tree sequence from its
{ref}`topological rank<sec_python_api_trees_topological_analysis>`.
:::

:::{note}
Several of these methods are {func}`static<python:staticmethod>`, so should
be called e.g. as `tskit.Tree.generate_balanced(4)` rather than used on
a specific {class}`Tree` instance.
:::


(sec_python_api_trees_node_measures)=

#### Node measures

Often it is useful to access information pertinant to a specific node or set of nodes
but which might also change from tree to tree in the tree sequence. Examples include
the encoding of the tree via `parent`, `left_child`, etc.
(see {ref}`sec_data_model_tree_structure`), the number of samples under a node,
or the most recent common ancestor (MRCA) of two nodes. This sort of information is
available via simple and high performance {class}`Tree` methods


(sec_python_api_trees_node_measures_simple)=

##### Simple measures

These return a simple number, or (usually) short list of numbers relevant to a specific
node or limited set of nodes.

```{eval-rst}
Node information
    .. autosummary::
      Tree.is_sample
      Tree.is_isolated
      Tree.is_leaf
      Tree.is_internal
      Tree.parent
      Tree.num_children
      Tree.time
      Tree.branch_length
      Tree.depth
      Tree.population
      Tree.right_sib
      Tree.left_sib
      Tree.right_child
      Tree.left_child
      Tree.children
      Tree.edge

Descendant nodes
    .. autosummary::
      Tree.leaves
      Tree.samples
      Tree.num_samples
      Tree.num_tracked_samples

Multiple nodes
    .. autosummary::
      Tree.is_descendant
      Tree.mrca
      Tree.tmrca
```


(sec_python_api_trees_node_measures_array)=

##### Array access

These all return a numpy array whose length corresponds to
the total number of nodes in the tree sequence. They provide direct access
to the underlying memory structures, and are thus very efficient, providing a
high performance interface which can be used in conjunction with the equivalent
{ref}`traversal methods<sec_python_api_trees_traversal>`.

```{eval-rst}
.. autosummary::
  Tree.parent_array
  Tree.left_child_array
  Tree.right_child_array
  Tree.left_sib_array
  Tree.right_sib_array
  Tree.num_children_array
  Tree.edge_array
```


(sec_python_api_trees_traversal)=

#### Tree traversal

Moving around within a tree usually involves visiting the tree nodes in some sort of
order. Often, given a particular order, it is convenient to iterate over each node
using the {meth}`Tree.nodes` method. However, for high performance algorithms, it
may be more convenient to access the node indices for a particular order as
an array, and use this, for example, to index into one of the node arrays (see
{ref}`sec_topological_analysis_traversal`).

```{eval-rst}
Iterator access
    .. autosummary::

      Tree.nodes
      Tree.ancestors

Array access
    .. autosummary::

      Tree.postorder
      Tree.preorder
      Tree.timeasc
      Tree.timedesc
```


(sec_python_api_trees_topological_analysis)=

#### Topological analysis

The topology of a tree refers to the simple relationship among samples
(i.e. ignoring branch lengths), see {ref}`sec_combinatorics` for more details. These
methods provide ways to enumerate and count tree topologies.

Briefly, the position of a tree in the enumeration `all_trees` can be obtained using
the tree's {meth}`~Tree.rank` method. Inversely, a {class}`Tree` can be constructed
from a position in the enumeration with {meth}`Tree.unrank`.


```{eval-rst}
Methods of a tree
    .. autosummary::
      Tree.rank
      Tree.count_topologies

Functions and static methods
    .. autosummary::
      Tree.unrank
      all_tree_shapes
      all_tree_labellings
      all_trees
```


(sec_python_api_trees_comparing)=

#### Comparing trees

```{eval-rst}
.. autosummary::
  Tree.kc_distance
  Tree.rf_distance
```

(sec_python_api_trees_balance)=

#### Balance/imbalance indices

```{eval-rst}
.. autosummary::
  Tree.colless_index
  Tree.sackin_index
  Tree.b1_index
  Tree.b2_index
```

(sec_python_api_trees_sites_mutations)=

#### Sites and mutations

```{eval-rst}
.. autosummary::
  Tree.sites
  Tree.num_sites
  Tree.mutations
  Tree.num_mutations
  Tree.map_mutations
```


(sec_python_api_trees_moving_to)=

#### Moving to other trees

```{eval-rst}
.. autosummary::

  Tree.next
  Tree.prev
  Tree.first
  Tree.last
  Tree.seek
  Tree.seek_index
  Tree.clear
```

#### Display

```{eval-rst}
.. autosummary::

  Tree.draw_svg
  Tree.draw_text
  Tree.__str__
  Tree._repr_html_
```

#### Export

```{eval-rst}
.. autosummary::

  Tree.as_dict_of_dicts
  Tree.as_newick
```


(sec_tables_api)=

## Tables and Table Collections

The information required to construct a tree sequence is stored in a collection
of *tables*, each defining a different aspect of the structure of a tree
sequence. These tables are described individually in
{ref}`the next section<sec_tables_api_table>`. However, these are interrelated,
and so many operations work
on the entire collection of tables, known as a *table collection*.

(sec_tables_api_table_collection)=

### `TableCollection` API

The {class}`TableCollection` and {class}`TreeSequence` classes are
deeply related. A `TreeSequence` instance is based on the information
encoded in a `TableCollection`. Tree sequences are **immutable**, and
provide methods for obtaining trees from the sequence. A `TableCollection`
is **mutable**, and does not have any methods for obtaining trees.
The `TableCollection` class thus allows creation and modification of
tree sequences (see the {ref}`sec_tables` tutorial).


#### General properties

Specific {ref}`tables<sec_tables_api_table>` in the {class}`TableCollection`
are be accessed using the plural version of their name, so that, for instance, the
individual table can be accessed using `table_collection.individuals`. A table
collection also has other properties containing, for example, number of bytes taken
to store it and the top-level metadata associated with the tree sequence as a whole.

```{eval-rst}
Table access
    .. autosummary::
      TableCollection.individuals
      TableCollection.nodes
      TableCollection.edges
      TableCollection.migrations
      TableCollection.sites
      TableCollection.mutations
      TableCollection.populations
      TableCollection.provenances

Other properties
    .. autosummary::
      TableCollection.file_uuid
      TableCollection.indexes
      TableCollection.nbytes
      TableCollection.table_name_map
      TableCollection.metadata
      TableCollection.metadata_bytes
      TableCollection.metadata_schema
      TableCollection.sequence_length
      TableCollection.time_units
```


(sec_tables_api_transformation)=

#### Transformation

These methods act in-place to transform the contents of a {class}`TableCollection`,
either by modifying the underlying tables (removing, editing, or adding to them) or
by adjusting the table collection so that it meets the
{ref}`sec_valid_tree_sequence_requirements`.


(sec_tables_api_modification)=

##### Modification

These methods modify the data stored in a {class}`TableCollection`. They also have
{ref}`equivalant TreeSequence versions<sec_python_api_tree_sequences_modification>`
(unlike the methods described below those do *not* operate in place, but rather act in
a functional way, returning a new tree sequence while leaving the original unchanged).

```{eval-rst}
.. autosummary::
  TableCollection.clear
  TableCollection.simplify
  TableCollection.subset
  TableCollection.delete_intervals
  TableCollection.keep_intervals
  TableCollection.delete_sites
  TableCollection.trim
  TableCollection.union
  TableCollection.delete_older
```

(sec_tables_api_creating_valid_tree_sequence)=

##### Creating a valid tree sequence

These methods can be used to help reorganise or rationalise the
{class}`TableCollection` so that it is in the form
{ref}`required<sec_valid_tree_sequence_requirements>` for
it to be {meth}`converted<TableCollection.tree_sequence>`
into a {class}`TreeSequence`. This may require sorting the tables,
ensuring they are logically consistent, and adding {ref}`sec_table_indexes`.

:::{note}
These methods are not guaranteed to make valid a {class}`TableCollection` which is
logically inconsistent, for example if multiple edges have the same child at a
given position on the genome or if non-existent node IDs are referenced.
:::

```{eval-rst}
Sorting
    .. autosummary::
      TableCollection.sort
      TableCollection.sort_individuals
      TableCollection.canonicalise

Logical consistency
    .. autosummary::
      TableCollection.compute_mutation_parents
      TableCollection.compute_mutation_times
      TableCollection.deduplicate_sites

Indexing
    .. autosummary::
      TableCollection.has_index
      TableCollection.build_index
      TableCollection.drop_index
```

#### Miscellaneous methods

```{eval-rst}
.. autosummary::
  TableCollection.copy
  TableCollection.equals
  TableCollection.link_ancestors
```

#### Export
```{eval-rst}
.. autosummary::
  TableCollection.tree_sequence
  TableCollection.dump
```

(sec_tables_api_table)=

### Table APIs

Here we outline the table classes and the common methods and variables available for
each. For description and definition of each table's meaning
and use, see {ref}`the table definitions <sec_table_definitions>`.

```{eval-rst}
.. autosummary::

  IndividualTable
  NodeTable
  EdgeTable
  MigrationTable
  SiteTable
  MutationTable
  PopulationTable
  ProvenanceTable
```

(sec_tables_api_accessing_table_data)=

#### Accessing table data

The tables API provides an efficient way of working
with and interchanging {ref}`tree sequence data <sec_data_model>`. Each table class
(e.g, {class}`NodeTable`, {class}`EdgeTable`, {class}`SiteTable`) has a specific set
of columns with fixed types, and a set of methods for setting and getting the data
in these columns. The number of rows in the table `t` is given by `len(t)`.

```{code-cell} ipython3
import tskit
t = tskit.EdgeTable()
t.add_row(left=0, right=1, parent=10, child=11)
t.add_row(left=1, right=2, parent=9, child=11)
print("The table contains", len(t), "rows")
print(t)
```

Each table supports accessing the data either by row or column. To access the data in
a *column*, we can use standard attribute access which will
return a copy of the column data as a numpy array:

```{code-cell} ipython3
t.left
```

```{code-cell} ipython3
t.parent
```

To access the data in a *row*, say row number `j` in table `t`, simply use `t[j]`:

```{code-cell} ipython3
t[0]
```

This also works as expected with negative `j`, counting rows from the end of the table

```{code-cell} ipython3
t[-1]
```

The returned row has attributes allowing contents to be accessed by name, e.g.
`site_table[0].position`, `site_table[0].ancestral_state`, `site_table[0].metadata`
etc.:

```{code-cell} ipython3
t[-1].right
```

Row attributes cannot be modified directly. Instead, the `replace` method of a row
object can be used to create a new row with one or more changed column
values, which can then be used to replace the original. For example:

```{code-cell} ipython3
t[-1] = t[-1].replace(child=4, right=3)
print(t)
```

Tables also support the {mod}`pickle` protocol, and so can be easily serialised and
deserialised. This can be useful, for example, when performing parallel computations
using the {mod}`multiprocessing` module (however, pickling will not be as efficient
as storing tables in the native {ref}`format <sec_tree_sequence_file_format>`).

```{code-cell} ipython3
import pickle
serialised = pickle.dumps(t)
t2 = pickle.loads(serialised)
print(t2)
```

Tables support the equality operator `==` based on the data
held in the columns:

```{code-cell} ipython3
t == t2
```

```{code-cell} ipython3
t is t2
```

```{code-cell} ipython3
t2.add_row(0, 1, 2, 3)
print(t2)
t == t2
```

:::{todo}
Move some or all of these examples into a suitable alternative chapter.
:::


(sec_tables_api_text_columns)=

##### Text columns

As described in the {ref}`sec_encoding_ragged_columns`, working with
variable length columns is somewhat more involved. Columns
encoding text data store the **encoded bytes** of the flattened
strings, and the offsets into this column in two separate
arrays.

Consider the following example:

```{code-cell} ipython3
t = tskit.SiteTable()
t.add_row(0, "A")
t.add_row(1, "BB")
t.add_row(2, "")
t.add_row(3, "CCC")
print(t)
print(t[0])
print(t[1])
print(t[2])
print(t[3])
```

Here we create a {class}`SiteTable` and add four rows, each with a different
`ancestral_state`. We can then access this information from each
row in a straightforward manner. Working with columns of text data
is a little trickier, however:

```{code-cell} ipython3
print(t.ancestral_state)
print(t.ancestral_state_offset)
```

```{code-cell} ipython3
tskit.unpack_strings(t.ancestral_state, t.ancestral_state_offset)
```

Here, the `ancestral_state` array is the UTF8 encoded bytes of the flattened
strings, and the `ancestral_state_offset` is the offset into this array
for each row. The {func}`tskit.unpack_strings` function, however, is a convient
way to recover the original strings from this encoding. We can also use the
{func}`tskit.pack_strings` to insert data using this approach:

```{code-cell} ipython3
a, off = tskit.pack_strings(["0", "12", ""])
t.set_columns(position=[0, 1, 2], ancestral_state=a, ancestral_state_offset=off)
print(t)
```

When inserting many rows with standard infinite sites mutations (i.e.,
ancestral state is "0"), it is more efficient to construct the
numpy arrays directly than to create a list of strings and use
{func}`pack_strings`. When doing this, it is important to note that
it is the **encoded** byte values that are stored; by default, we
use UTF8 (which corresponds to ASCII for simple printable characters).:

```{code-cell} ipython3
import numpy as np
t_s = tskit.SiteTable()
m = 10
a = ord("0") + np.zeros(m, dtype=np.int8)
off = np.arange(m + 1, dtype=np.uint32)
t_s.set_columns(position=np.arange(m), ancestral_state=a, ancestral_state_offset=off)
print(t_s)
print("ancestral state data", t_s.ancestral_state)
print("ancestral state offsets", t_s.ancestral_state_offset)
```


In the mutation table, the derived state of each mutation can be handled similarly:

```{code-cell} ipython3
t_m = tskit.MutationTable()
site = np.arange(m, dtype=np.int32)
d, off = tskit.pack_strings(["1"] * m)
node = np.zeros(m, dtype=np.int32)
t_m.set_columns(site=site, node=node, derived_state=d, derived_state_offset=off)
print(t_m)
```

:::{todo}
Move some or all of these examples into a suitable alternative chapter.
:::


(sec_tables_api_binary_columns)=

##### Binary columns

Columns storing binary data take the same approach as
{ref}`sec_tables_api_text_columns` to encoding
{ref}`variable length data <sec_encoding_ragged_columns>`.
The difference between the two is only raw {class}`bytes` values are accepted: no
character encoding or decoding is done on the data. Consider the following example
where a table has no `metadata_schema` such that arbitrary bytes can be stored and
no automatic encoding or decoding of objects is performed by the Python API and we can
store and retrieve raw `bytes`. (See {ref}`sec_metadata` for details):

Below, we add two rows to a {class}`NodeTable`, with different
{ref}`metadata <sec_metadata_definition>`. The first row contains a simple
byte string, and the second contains a Python dictionary serialised using
{mod}`pickle`.

```{code-cell} ipython3
t = tskit.NodeTable()
t.add_row(metadata=b"these are raw bytes")
t.add_row(metadata=pickle.dumps({"x": 1.1}))
print(t)
```

Note that the pickled dictionary is encoded in 24 bytes containing unprintable
characters. It appears to be unrelated to the original contents, because the binary
data is [base64 encoded](https://en.wikipedia.org/wiki/Base64) to ensure that it is
print-safe (and doesn't break your terminal). (See the
{ref}`sec_metadata_definition` section for more information on the
use of base64 encoding.).

We can access the metadata in a row (e.g., `t[0].metadata`) which returns a Python
bytes object containing precisely the bytes that were inserted.

```{code-cell} ipython3
print(t[0].metadata)
print(t[1].metadata)
```

The metadata containing the pickled dictionary can be unpickled using
{func}`pickle.loads`:

```{code-cell} ipython3
print(pickle.loads(t[1].metadata))
```

As previously, the `replace` method can be used to change the metadata,
by overwriting an existing row with an updated one:

```{code-cell} ipython3
t[0] = t[0].replace(metadata=b"different raw bytes")
print(t)
```

Finally, when we print the `metadata` column, we see the raw byte values
encoded as signed integers. As for {ref}`sec_tables_api_text_columns`,
the `metadata_offset` column encodes the offsets into this array. So, we
see that the first metadata value is 9 bytes long and the second is 24.

```{code-cell} ipython3
print(t.metadata)
print(t.metadata_offset)
```

The {func}`tskit.pack_bytes` and {func}`tskit.unpack_bytes` functions are
also useful for encoding data in these columns.

:::{todo}
Move some or all of these examples into a suitable alternative chapter.
:::



#### Table functions

```{eval-rst}
.. autosummary::

  parse_nodes
  parse_edges
  parse_sites
  parse_mutations
  parse_individuals
  parse_populations
  parse_migrations
  pack_strings
  unpack_strings
  pack_bytes
  unpack_bytes
```


(sec_python_api_metadata)=

## Metadata API

The `metadata` module provides validation, encoding and decoding of metadata
using a schema. See {ref}`sec_metadata`, {ref}`sec_metadata_api_overview` and
{ref}`sec_tutorial_metadata`.

```{eval-rst}
.. autosummary::
  MetadataSchema
  register_metadata_codec
```

:::{seealso}
Refer to the top level metadata-related properties of TreeSequences and TableCollections,
such as {attr}`TreeSequence.metadata` and {attr}`TreeSequence.metadata_schema`. Also the
metadata fields of
{ref}`objects accessed<sec_python_api_tree_sequences_obtaining_other_objects>` through
the {class}`TreeSequence` API.
:::


(sec_python_api_provenance)=

## Provenance

We provide some preliminary support for validating JSON documents against the
{ref}`provenance schema <sec_provenance>`. Programmatic access to provenance
information is planned for future versions.


```{eval-rst}
.. autosummary::
  validate_provenance
```

(sec_utility_api)=

## Utility functions

Miscellaneous top-level utility functions.

```{eval-rst}
.. autosummary::
  is_unknown_time
  random_nucleotides
```


(sec_python_api_reference)=

## Reference documentation

(sec_python_api_constants)=

### Constants

The following constants are used throughout the `tskit` API.

```{eval-rst}
.. automodule:: tskit
   :members:
```

(sec_python_api_exceptions)=

### Exceptions

```{eval-rst}
.. autoexception:: DuplicatePositionsError
.. autoexception:: MetadataEncodingError
.. autoexception:: MetadataSchemaValidationError
.. autoexception:: MetadataValidationError
.. autoexception:: ProvenanceValidationError
```

(sec_python_api_functions)=

### Top-level functions

```{eval-rst}
.. autofunction:: all_trees
.. autofunction:: all_tree_shapes
.. autofunction:: all_tree_labellings
.. autofunction:: is_unknown_time
.. autofunction:: load
.. autofunction:: load_text
.. autofunction:: pack_bytes
.. autofunction:: pack_strings
.. autofunction:: parse_edges
.. autofunction:: parse_individuals
.. autofunction:: parse_mutations
.. autofunction:: parse_nodes
.. autofunction:: parse_populations
.. autofunction:: parse_migrations
.. autofunction:: parse_sites
.. autofunction:: random_nucleotides
.. autofunction:: register_metadata_codec
.. autofunction:: validate_provenance
.. autofunction:: unpack_bytes
.. autofunction:: unpack_strings

```

### Tree and tree sequence classes

#### The {class}`Tree` class

Also see the {ref}`sec_python_api_trees` summary.

```{eval-rst}
.. autoclass:: Tree()
    :members:
    :special-members: __str__
    :private-members: _repr_html_
```

#### The {class}`TreeSequence` class

Also see the {ref}`sec_python_api_tree_sequences` summary.

```{eval-rst}
.. autoclass:: TreeSequence()
    :members:
    :special-members: __str__
    :private-members: _repr_html_
```

### Simple container classes

#### The {class}`Individual` class

```{eval-rst}
.. autoclass:: Individual()
    :members:
```

#### The {class}`Node` class

```{eval-rst}
.. autoclass:: Node()
    :members:
```

#### The {class}`Edge` class

```{eval-rst}
.. autoclass:: Edge()
    :members:
```

#### The {class}`Site` class

```{eval-rst}
.. autoclass:: Site()
    :members:
```

#### The {class}`Mutation` class

```{eval-rst}
.. autoclass:: Mutation()
    :members:
```

#### The {class}`Variant` class

```{eval-rst}
.. autoclass:: Variant()
    :members:
```

#### The {class}`Migration` class

```{eval-rst}
.. autoclass:: Migration()
    :members:
```

#### The {class}`Population` class

```{eval-rst}
.. autoclass:: Population()
    :members:
```

#### The {class}`Provenance` class

```{eval-rst}
.. autoclass:: Provenance()
    :members:
```

#### The {class}`Interval` class

```{eval-rst}
.. autoclass:: Interval()
    :members:
```

#### The {class}`Rank` class

```{eval-rst}
.. autoclass:: Rank()
    :members:
```

### TableCollection and Table classes

#### The {class}`TableCollection` class

Also see the {ref}`sec_tables_api_table_collection` summary.

```{eval-rst}
.. autoclass:: TableCollection
    :inherited-members:
    :members:
```

% Overriding the default signatures for the tables here as they will be
% confusing to most users.


#### {class}`IndividualTable` classes

```{eval-rst}
.. autoclass:: IndividualTable()
    :members:
    :inherited-members:
    :special-members: __getitem__
```

##### Associated row class

A row returned from an {class}`IndividualTable` is an instance of the following
basic class, where each attribute matches an identically named attribute in the
{class}`Individual` class.

```{eval-rst}
.. autoclass:: IndividualTableRow()
    :members:
    :inherited-members:
```


#### {class}`NodeTable` classes

```{eval-rst}
.. autoclass:: NodeTable()
    :members:
    :inherited-members:
    :special-members: __getitem__
```

##### Associated row class

A row returned from a {class}`NodeTable` is an instance of the following
basic class, where each attribute matches an identically named attribute in the
{class}`Node` class.

```{eval-rst}
.. autoclass:: NodeTableRow()
    :members:
    :inherited-members:
```


#### {class}`EdgeTable` classes

```{eval-rst}
.. autoclass:: EdgeTable()
    :members:
    :inherited-members:
    :special-members: __getitem__
```

##### Associated row class

A row returned from an {class}`EdgeTable` is an instance of the following
basic class, where each attribute matches an identically named attribute in the
{class}`Edge` class.

```{eval-rst}
.. autoclass:: EdgeTableRow()
    :members:
    :inherited-members:
```


#### {class}`MigrationTable` classes

```{eval-rst}
.. autoclass:: MigrationTable()
    :members:
    :inherited-members:
    :special-members: __getitem__
```

##### Associated row class

A row returned from a {class}`MigrationTable` is an instance of the following
basic class, where each attribute matches an identically named attribute in the
{class}`Migration` class.

```{eval-rst}
.. autoclass:: MigrationTableRow()
    :members:
    :inherited-members:
```


#### {class}`SiteTable` classes

```{eval-rst}
.. autoclass:: SiteTable()
    :members:
    :inherited-members:
    :special-members: __getitem__
```

##### Associated row class

A row returned from a {class}`SiteTable` is an instance of the following
basic class, where each attribute matches an identically named attribute in the
{class}`Site` class.

```{eval-rst}
.. autoclass:: SiteTableRow()
    :members:
    :inherited-members:
```


#### {class}`MutationTable` classes

```{eval-rst}
.. autoclass:: MutationTable()
    :members:
    :inherited-members:
    :special-members: __getitem__
```

##### Associated row class

A row returned from a {class}`MutationTable` is an instance of the following
basic class, where each attribute matches an identically named attribute in the
{class}`Mutation` class.

```{eval-rst}
.. autoclass:: MutationTableRow()
    :members:
    :inherited-members:
```


#### {class}`PopulationTable` classes

```{eval-rst}
.. autoclass:: PopulationTable()
    :members:
    :inherited-members:
    :special-members: __getitem__
```

##### Associated row class

A row returned from a {class}`PopulationTable` is an instance of the following
basic class, where each attribute matches an identically named attribute in the
{class}`Population` class.

```{eval-rst}
.. autoclass:: PopulationTableRow()
    :members:
    :inherited-members:
```


#### {class}`ProvenanceTable` classes

Also see the {ref}`sec_provenance` and
{ref}`provenance API methods<sec_python_api_provenance>`.

```{eval-rst}
.. autoclass:: ProvenanceTable()
    :members:
    :inherited-members:
```

##### Associated row class

A row returned from a {class}`ProvenanceTable` is an instance of the following
basic class, where each attribute matches an identically named attribute in the
{class}`Provenance` class.

```{eval-rst}
.. autoclass:: ProvenanceTableRow()
    :members:
    :inherited-members:
```

(sec_python_api_reference_identity)=

### Identity classes

The classes documented in this section are associated with summarising
identity relationships between pairs of samples. See the {ref}`sec_identity`
section for more details and examples.

#### The {class}`IdentitySegments` class

```{eval-rst}
.. autoclass:: IdentitySegments()
    :members:
```

#### The {class}`IdentitySegmentList` class

```{eval-rst}
.. autoclass:: IdentitySegmentList()
    :members:
```

#### The {class}`IdentitySegment` class

```{eval-rst}
.. autoclass:: IdentitySegment()
    :members:
```

### Miscellaneous classes

#### The {class}`ReferenceSequence` class

```{eval-rst}
.. todo:: Add a top-level summary section that we can link to from here.
```

```{eval-rst}
.. autoclass:: ReferenceSequence()
    :members:
    :inherited-members:
```

#### The {class}`MetadataSchema` class

Also see the {ref}`sec_python_api_metadata` summary.

```{eval-rst}
.. autoclass:: MetadataSchema
    :members:
    :inherited-members:
```

#### The {class}`TableMetadataSchemas` class

```{eval-rst}
.. autoclass:: TableMetadataSchemas
    :members:
```

#### The {class}`TopologyCounter` class

```{eval-rst}
.. autoclass:: TopologyCounter
```

#### The {class}`LdCalculator` class

```{eval-rst}
.. autoclass:: LdCalculator
    :members:
```

#### The {class}`TableCollectionIndexes` class

```{eval-rst}
.. autoclass:: TableCollectionIndexes
    :members:
```

#### The {class}`SVGString` class

```{eval-rst}
.. autoclass:: SVGString
    :members:
    :private-members: _repr_svg_
```


--- ../../tskit/docs/tables_and_editing.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

```{currentmodule} tskit
```

```{code-cell} ipython3
:tags: [remove-cell]
import io
import pickle

import msprime
import tskit
import numpy as np

def tables_examples():
    nodes = """\
    is_sample   time    population    individual
    1           0       0             0
    1           0       0             0
    1           0       0             1
    0           0.15    0             2
    0           0.6     -1            -1
    0           0.8     -1            -1
    0           1.0     -1            -1
    """
    edges = """\
    left    right   parent  child
    20      80      3       0
    20      80      3       2
    0       100     4       1
    0       20      4       2
    80      100     4       2
    20      80      4       3
    80      100     5       0
    80      100     5       4
    0       20      6       0
    0       20      6       4
    """
    sites = """\
    id    position    ancestral_state
    0     15          A
    1     42          G
    2     60          T
    """
    mutations = """\
    site    node    derived_state    parent    time
    0       4       G                -1        0.9
    1       1       A                -1        0.4
    2       3       C                -1        0.55
    2       0       T                2         0.1
    """
    individuals = """\
    flags    parents    metadata
    0        2,-1       {"name":"Alice"}
    0        2,-1       {"name":"Bob"}
    0        -1,-1      {"name":"Carol"}
    """
    populations = """\
    metadata
    {"name":"GBR","description":"British_in_England_and_Scotland"}
    """  # (NB: can't use spaces in load_text with strict==False)
    ts = tskit.load_text(
        nodes=io.StringIO(nodes),
        edges=io.StringIO(edges),
        individuals=io.StringIO(individuals),
        populations=io.StringIO(populations),
        strict=False,
        base64_metadata=False,
    )
    ts.dump("data/tables_example.trees")
    tables = tskit.load_text(
        nodes=io.StringIO(nodes),
        edges=io.StringIO(edges),
        sites=io.StringIO(sites),
        mutations=io.StringIO(mutations),
        individuals=io.StringIO(individuals),
        populations=io.StringIO(populations),
        strict=False,
        base64_metadata=False,
    ).dump_tables()
    tables.individuals.metadata_schema = tskit.MetadataSchema({'codec': 'json'})
    tables.populations.metadata = np.where(  # Hack to replace underscores with spaces
        tables.populations.metadata == ord("_"), ord(" "), tables.populations.metadata)
    tables.populations.metadata_schema = tskit.MetadataSchema({'codec': 'json'})
    tables.tree_sequence().dump("data/tables_example_muts.trees")

def construction_example():
    nodes = """\
    id      is_sample   time
    0       1           0
    1       1           0
    2       0           10
    3       1           0
    4       0           20
    """
    edges = """\
    left    right   parent  child
    0       1000    2       0
    0       1000    2       1
    0       1000    4       2
    0       1000    4       3
    """
    sites = """\
    id    position    ancestral_state
    0     500          A
    """
    mutations = """\
    site    node    derived_state    parent
    0       3       G                -1
    """
    ts = tskit.load_text(
        nodes=io.StringIO(nodes),
        edges=io.StringIO(edges),
        sites=io.StringIO(sites),
        mutations=io.StringIO(mutations),
        strict=False,
    )
    ts.dump("data/construction_example.trees")
    

def create_notebook_data():
    tables_examples()
    construction_example()

# create_notebook_data()  # uncomment to recreate the tree seqs used in this notebook
```


(sec_tables)=

# Tables and editing

Internally, a tree sequence can be thought of as a set of tables, and {program}`tskit`
provides an interface for dealing with these tables directly. This is particularly
relevant when you wish to {ref}`edit or otherwise modify<sec_tables_editing>`
a tree sequence, although tables are also useful for bulk access to data contained in
the tree sequence, such as the times of all nodes.

There are eight tables that together define a tree sequence, although some may be empty,
and together they form a {class}`TableCollection`.
The tables are defined in the official {program}`tskit` documentation for
{ref}`Table Definitions <tskit:sec_table_definitions>`, and the
{ref}`Tables API <tskit:sec_tables_api>` section in the docs describes how to work with
them. In this tutorial we give some pointers about what you can and cannot do with them.

(sec_tables_correspondence)=

## Correspondence between tables and trees

Consider the following sequence of trees:

```{code-cell} ipython3
ts = tskit.load("data/tables_example.trees")
```

```{code-cell} ipython3
:"tags": ["hide-input"]
ts.draw_svg(y_axis=True)
```

Ancestral recombination events have produced three different trees
that relate the three sampled genomes ``0``, ``1``, and ``2`` to each other
along the chromosome of length 100.

(sec_tables_correspondence_nodes_and_edges)=

### Node and edge tables

Each node in each of the above trees represents a particular ancestral genome
(a *haploid* genome; diploid individuals would be represented by two nodes).
Details about each node, including the time it lived, are stored in a
{class}`NodeTable` (details {ref}`here<tskit:sec_mutation_table_definition>`)

```{code-cell} ipython3
ts.tables.nodes
```

Importantly, the first column, `id`, is not actually recorded, and is
only shown when printing out node tables (as here) for convenience.
The second column, `flags` records a `1` for the nodes that are *samples*,
i.e., whose entire genealogical history is recorded by these trees.
You can see that the sample nodes are present in all the trees, but other nodes need
not be: for example node 3 is only seen in the middle tree. The `individual` and
`population` columns store an index into the rows of the
{ref}`sec_tables_overview_individuals_and_populations`, or `tskit.NULL` (`-1`)
if the node is not associated with an individual or population. The ``time`` column
records the historical age of the node, with contemporary nodes having a ``time`` of 0.

The way the nodes are related to each other (i.e. the tree topology) is stored
in an {class}`EdgeTable` (details {ref}`here<tskit:sec_edge_table_definition>`).
Since some edges are present in more than one tree (e.g., node 1 inherits from node 4
across the entire sequence), each edge contains not only the IDs of a parent and a child
node, but also the left and right positions which define the genomic region for which it
appears in the trees:

```{code-cell} ipython3
ts.tables.edges
```

Since node 3 is most recent, the edge that says that nodes 0 and 2 inherit
from node 3 on the interval between 0.2 and 0.8 comes first.  Next are the
edges from node 4: there are four of these, as the edge from node 4 to node
1 is shared across the entire sequence, and for each of the three
genomic intervals there is an additional child node. At this
point, we know the full tree on the middle interval.  Finally, edges
specifying the common ancestor of 0 and 4 on the remaining intervals (parents 6
and 5 respectively) allow us to construct all trees across the entire interval.


(sec_tables_correspondence_sites_and_mutations)=

### Site and mutation tables

Most tree sequences have DNA variation data associated with them,
{ref}`stored as mutations overlaid on the trees<sec_what_is_dna_data>`:

```{code-cell} ipython3
ts = tskit.load("data/tables_example_muts.trees")
```

```{code-cell} ipython3
:"tags": ["hide-input"]
mut_labels = {}  # An array of labels for the mutations, listing position & allele change
for mut in ts.mutations():  # This entire loop is just to make pretty labels
    site = ts.site(mut.site)
    older_mut = mut.parent >= 0  # is there an older mutation at the same position?
    prev = ts.mutation(mut.parent).derived_state if older_mut else site.ancestral_state
    mut_labels[mut.id] = "{}→{} @{:g}".format(prev, mut.derived_state, site.position)
ts.draw_svg(y_axis=True, mutation_labels=mut_labels)
```

There are four mutations in the depiction above,
marked by red crosses: one above node ``4`` on the first tree which records an A to G
transition at position 15, another above node ``1`` in the second tree which records a G
to A transition at position 42, and the final two above nodes ``0`` and ``3`` recording
transitions, both at position 60, on the second tree. The positions are recorded in the
{class}`SiteTable` (details {ref}`here<tskit:sec_site_table_definition>`):

```{code-cell} ipython3
ts.tables.sites
```

As with node tables, the ``id`` column is **not** actually recorded, but is
implied by the position in the table.  The mutations themselves are recorded in the
{class}`MutationTable` (details {ref}`here<tskit:sec_mutation_table_definition>`).
This associates each mutation with a site ID, the ID of the node above which the
mutation occurs, the derived state to which the allele has mutated, and (optionally) a
time at which the mutation occured.


```{code-cell} ipython3
ts.tables.mutations
```

Where there are multiple mutations at the same site, there can be mutations
stacked on top of each other. The "parent" column therefore contains the ID of the
mutation immediately above the current one at the same site, or -1 if there is no parent.

These sites and mutations allow us to calculate the DNA sequence, or haplotype, for each
of the sample nodes:

```{code-cell} ipython3
for sample, h in enumerate(ts.haplotypes()):
    print(f"Sample {sample}: {h}")
```

(sec_tables_overview)=

## Tables overview

Now that we've introduced the most important tables that make up a tree sequence, it's
worth seeing how they come together to form a {class}`TableCollection`.


::::{margin}
:::{note}
For simplicity, this diagram omits 2 extra tables which form part of a
{class}`TableCollection`: the {class}`MigrationTable`
(recording details of migrations between populations), and the {class}`ProvenanceTable`
(recording the {ref}`provenance<tskit:sec_provenance>` of the data in a tree sequence)
:::
::::

The following visualization shows how they relate to one another, as well as introducing
two additional tables: the {class}`IndividualTable` and {class}`PopulationTable`
(which record which
{ref}`individuals and populations<sec_terminology_individuals_and_populations>` a node
is contained in).

![The tables in a tree sequence, and their interdependencies](_static/tables_uml.svg)

The lines in this diagram below show where one table refers to the IDs
of another table, that is, they indicate dependencies between tables.


(sec_tables_overview_individuals_and_populations)=

### Individual and population tables

Previously we noted that a row in the node table could contain an index into the
{class}`IndividualTable` and the {class}`PopulationTable`. In fact, in our example
tree sequence, the {ref}`node table above<sec_tables_correspondence_nodes_and_edges>`
indicates that nodes 0 to 3 belong to the population with ID 0 in the population table.
Here's what that table looks like in this tree sequence:

```{code-cell} ipython3
ts.tables.populations
```

The data in this table, and indeed in the individual table, is not visible in the
default tree sequence visualization, but can be very useful for keeping track of
information about the identity and location of nodes in the tree sequence. In
particular, you can see that the populations in this tree sequence have associated
*metadata*, a topic we will discuss in the
{ref}`next section<sec_tables_overview_metadata>`.

The node table also indicated that two nodes (node 0 and 1) point to the individual with
ID 0, one node (node 2) points to individual 1, and another node points to individual 2.
Here's the relevant individual table:

```{code-cell} ipython3
ts.tables.individuals
```

Since individual 0 (Alice) is referenced by two nodes (genomes), and assuming that this
is a diploid species, this indicates that we have sampled both maternal and paternal
genomes for this individual. The other two indivividuals, Bob and Carol, are referenced
by only one node, so either we only managed to sample one of their two genomes,
or alternatively they could represent haploid individuals.

The `location` column of the individual table is intended for geographical location,
if relevant. The `parents` column, if filled out, allows us to construct
parts of the *pedigree* (the genealogical family tree) connecting the individuals
in question. It contains a list of individual IDs (or `tskit.NULL`/`-1` if a
parent is unknown). In this case, the `parents` column indicates that Alice and Bob
share the same parent, individual 2 (i.e. Carol). That means Alice and Bob must be
either siblings or half siblings.


(sec_tables_overview_metadata)=

### Metadata

::::{margin}
:::{note}
For more information on schemas, see the {ref}`sec_metadata` section of the
official {program}`tskit` documentation
:::
::::

As can be seen above, each table row, as well as the entire table collection, can be
associated with (optional) metadata, which can be stored in various formats. The format
is determined by the *schema* of the metadata. In this example, for clarity, the
metadata of the population and individual tables is stored as
[JSON](https://www.json.org/).

Because we have chosen to use JSON for the metadata of the individuals table, it is
easy to access the metadata values:

```{code-cell} ipython3
print("Information from the tree sequence")
for i in ts.individuals():
    print(f"Individual {i.id} is named {i.metadata['name']}")
print("\nEquivalent information from the associated tables")
for row in ts.tables.individuals:
    print(f"Individual name: {row.metadata['name']}")
```

{ref}`Later in this tutorial<sec_tables_editing_minor>` there are some basic
examples of editing the metadata fields; there is also an extensive
{ref}`sec_tutorial_metadata` tutorial which provides much more
detail about how to construct and use `tskit` metadata. 


(sec_tables_accessing)=

## Accessing table data

A simple way to see the contents of a table is to `print` it:

```{code-cell} ipython3
print(ts.tables.nodes)
```

But {program}`tskit` also provides access to columns and rows separately.


(sec_tables_accessing_columns)=

### Column access

An entire column in a table can be extracted as a {program}`numpy` array from the table
object. For instance, if ``n`` is a {class}`NodeTable`, then ``n.time``
will return an array containing the birth times of the individuals whose genomes
are represented by the nodes in the table. 

```{code-cell} ipython3
ts.tables.nodes.time
```

You can also *modify* entire columns, but when doing this to tables from a tree
sequence, you should always modify a *copy* of the original data, using the
{meth}`TreeSequence.dump_tables` method:

```{code-cell} ipython3
new_tables = ts.dump_tables()  # make a copy of the tree sequence tables, for editing
new_tables.nodes.time
```

Moreover, it is important to note that this numpy array itself is copy of the data
stored in `new_tables`. In other words, modifying individual elements of ``n.time`` will
*not* change the node table ``n``. To change the column data, you need to take a copy
of the array, modify it, and assign it back in. For example, here we add 0.25 to
every ``time`` except the first in the node table:

```{code-cell} ipython3
node_table = new_tables.nodes
times = node_table.time
print("Old node times:", times)
times[1:] = times[1:] + 0.25
node_table.time = times
node_table
```

When assigning columns like this, an error will be raised if the column is not of the
expected length:

```{code-cell} ipython3
:tags: [raises-exception]
node_table.time = times[2:]
```

(sec_tables_accessing_rows)=

### Row access

Rows can be accessed using square braces, which will return an object containing the
raw values, accessible via attributes

```{code-cell} ipython3
row = ts.tables.nodes[3]
print(
    "Node 3, belonging to individual", row.individual,
    "in population", row.population,
    "exists at time", row.time,
)
print(row)
```

Additionally, many rows can be extracted into a new table using slices or
{ref}`numpy indexing<numpy:basics.indexing>`

```{code-cell} ipython3
ts.tables.nodes[2:4]
```

As with columns, to change the data in a row of a tree sequence table, you must operate
on the copy that was produced by {meth}`TreeSequence.dump_tables`. Here we'll try to
change the ancestral state value of site 0.

```{code-cell} ipython3
new_tables.sites[0]  # Check the site's ancestral state ("A")
```

Since the row object(s) returned by the calls above are themselves copies,
the attributes enforce read-only access:

```{code-cell} ipython3
:tags: [raises-exception]
new_tables.sites[0].ancestral_state = "C" # This will not work: attributes are read-only
```

So to change row data you need to replace the row completely:

```{code-cell} ipython3
new_tables.sites[0] = tskit.SiteTableRow(position=15, ancestral_state="C", metadata=None)
```

For convenience, each row object has a `replace` method which creates a new table row
from an old one. This is often an easier way to changing one or two attributes:

```{code-cell} ipython3
# Another way to change a single entry: create a new row to replace the old one ...
replacement_row = new_tables.sites[0].replace(ancestral_state="C")
# ... then overwrite the original
new_tables.sites[0] = replacement_row

# Or equivalently, do both in a single step
new_tables.sites[0] = new_tables.sites[0].replace(ancestral_state="C")
```

It is also easy to add more rows:

```{code-cell} ipython3
new_pos = 10
new_site_id = new_tables.sites.add_row(
    position=new_pos, ancestral_state="G", metadata=b"An empty site"
    # NB: For this example table we have to feed a raw byte string as metadata
)
print(f"New empty site allocated at position {new_pos} with ID {new_site_id}")
new_tables.sites
```


(sec_tables_to_tree_sequence)=

## Turning tables into a tree sequence

The {meth}`TableCollection.tree_sequence` method will attempt to turn a table collection
into a tree sequence. This is not guaranteed to work: it's possible you have created a
nonsensical tree sequence where, for example, a child node has multiple parents at
a given position, or where edges reference non-existant nodes. The
{ref}`tskit:sec_valid_tree_sequence_requirements` section of the official tskit docs
lists the requirements for a {class}`TableCollection` to represent a valid
{class}`TreeSequence`.

(sec_tables_to_tree_sequence_sorting)=

### Sorting

Even if the tables represent a valid tree sequence, for efficiency reasons tree
sequences also require several of their constituent tables to be sorted in a specific
order. For instance, edges are {ref}`required<tskit:sec_edge_requirements>` to be sorted
in nondecreasing order of parent time, and sites are
{ref}`required<tskit:sec_site_requirements>` to be sorted in order of position.
We can ensure that a set of tables are correctly sorted by calling the
{meth}`TableCollection.sort` method.

In fact the new empty site that we added to the `new_tables` collection in the previous
section was at the start of the genome, so the table collection no longer has its sites
in the order required for a tree sequence. That means we need to call
`new_tables.sort()` before turning the table collection into a tree sequence:

```{code-cell} ipython3
# altered_ts = new_tables.tree_sequence()  # This won't work
new_tables.sort()
altered_ts = new_tables.tree_sequence()  # Now it will
```

Now that it had been turned into a tree sequence, we can plot it:

```{code-cell} ipython3
# Plot without mutation labels, for clarity 
altered_ts.draw_svg(y_axis=True, y_gridlines=True, mutation_labels={})
```

You can see that the new tree sequence has been modified as expected: there is a new
empty site at position 10 (represented by a tickmark above the X axis with no mutations
on it), and all the nodes except node 0 have had their times increased by 0.25.

(sec_tables_creating)=

## Constructing a tree sequence

With the tools above in mind, we can now see how to construct a tree sequence by hand.
It's unlikely that you would ever need to do this from scratch, but it's helpful to
understand how tables work. We'll build an extremely simple tree sequence, consisting of
a single tree that looks like this:

```{code-cell} ipython3
tskit.load("data/construction_example.trees").draw_svg(y_axis=True)
```

Starting with an empty set of tables, we can fill, say, the node information by using
{meth}`NodeTable.add_row` as follows:

```{code-cell} ipython3
tables = tskit.TableCollection(sequence_length=1e3)
node_table = tables.nodes  # set up an alias, for efficiency
node_table.add_row(flags=tskit.NODE_IS_SAMPLE)  # Node 0: defaults to time 0
node_table.add_row(flags=tskit.NODE_IS_SAMPLE)  # Node 1: defaults to time 0
node_table.add_row(time=10)  # Node 2: not a sample
node_table.add_row(flags=tskit.NODE_IS_SAMPLE)  # Node 3
node_table.add_row(time=20)  # Node 4
node_table
```

The ``.add_row()`` method is natural (and should be reasonably efficient) if
new records appear one-by-one. Alternatively ``.set_columns()`` can be used to
set columns for all the rows at once, by passing in numpy arrays of the appropriate
type. We'll use that for the edges:

```{code-cell} ipython3
import numpy as np
edge_table = tables.edges
edge_table.set_columns(
    left=np.array([0.0, 0.0, 0.0, 0.0]),
    right=np.array([1e3, 1e3, 1e3, 1e3]),
    parent=np.array([2, 2, 4, 4], dtype=np.int32),  # References IDs in the node table
    child=np.array([0, 1, 2, 3], dtype=np.int32),  # References IDs in the node table
)
edge_table
```

Finally we can add a site and a mutation: here we'll use 0/1 mutations rather than
ATGC.

```{code-cell} ipython3
site_id = tables.sites.add_row(position=500.0, ancestral_state='0')
tables.mutations.add_row(site=site_id, node=2, derived_state='1')
tables.sort()  # make sure the edges & sites are in the right order
ts = tables.tree_sequence()
print("A hand-built tree sequence!")
ts.draw_svg(y_axis=True)
```

:::{note}
The ordering requirements for a valid tree sequence
{ref}`do not specify<tskit:sec_node_requirements>` that rows in the
node table have to be sorted in any particular order, e.g. by time. By convention,
sample nodes are often the first ones listed in the node table, and this is the node
order returned by {meth}`~TreeSequence.simplify`, but the example above shows that
sample nodes need not necessarily be those with the IDs $0..n$.
:::

(sec_tables_editing)=

## Editing tree sequences

If you may wish to make modifications to a tree sequence, you will immediately encounter
a barrier: tree sequence objects are **immutable** and so cannot be edited in-place.
Therefore, the way to change a tree sequence is to edit its underlying tables.

Before providing further examples of how to edit tables, it's worth noting that
{program}`tskit` may provide a built-in method to achieve what you want to do, without
you having to think about the underlying tables at all. In particular, several methods
will return a new tree sequence that has been modified in some way. Here are some key
ones:

- {meth}`TreeSequence.delete_sites` returns a tree sequence with certain sites
  deleted
- {meth}`TreeSequence.delete_intervals` and {meth}`TreeSequence.keep_intervals` 
  return a tree sequence whose trees cover a smaller fraction of the genome (and which can
  be combined with {meth}`TreeSequence.trim` to change the coordinate system)
- {meth}`TreeSequence.simplify` returns a new tree sequence with some sample nodes removed
  (see the {ref}`simplification tutorial<sec_simplification>`).
- {meth}`TreeSequence.union` returns a new tree sequence formed by merging two other tree
  sequences together.

For example, to strip singleton sites from a tree sequence, we can simply call the
{meth}`TreeSequence.delete_sites` method:

```{code-cell} ipython3
import msprime
import numpy as np

ts = msprime.sim_ancestry(10, random_seed=123)
ts = msprime.sim_mutations(ts, rate=10, discrete_genome=False, random_seed=123)
print(ts.num_sites, "sites in the simulated tree sequence")

# Identify sites with a singleton allele
sites_with_a_singleton_allele = []
for v in ts.variants():
    non_missing_genotypes = v.genotypes[v.genotypes != tskit.MISSING_DATA]
    if np.any(np.bincount(non_missing_genotypes) == 1):
        sites_with_a_singleton_allele.append(v.site.id)
        
# Strip those sites from the tree sequence
ts_new = ts.delete_sites(sites_with_a_singleton_allele)
print(ts_new.num_sites, "sites after removing singletons")
```

However, you want to do something not covered by a built-in method, you will
need to edit a copy of the underlying tables and then create a new tree sequence from
the modified tables. 


(sec_tables_editing_minor)=

### Minor edits

{ref}`Previously<sec_tables_accessing_rows>` we saw how to change specific rows by
overwriting them. Here, for example, we use {meth}`IndividualTableRow.replace` to change
the `metadata` associated with an individual in one of our example tree sequences:

```{code-cell} ipython3
tables = altered_ts.dump_tables()
# Change the name "Bob" to "Robert"
tables.individuals[1] = tables.individuals[1].replace(metadata={"name": "Robert"})
edited_ts = tables.tree_sequence()
edited_ts.tables.individuals
```

If you are overwriting an entire column, especially something like metadata which is a
{ref}`ragged column<sec_encoding_ragged_columns>`, looping through each row and changing
it as above will not be very performant. Instead, you should overwrite the entire column
(or multiple columns) at once. We {ref}`previously saw<sec_tables_accessing_columns>`
how to overwrite an entire column of numbers, and for most columns that should suffice.
However, overwriting ragged columns like `metadata` or `parents` is more involved,
as you cannot set them to a single numpy array. Instead, for these columns, `tskit`
provides a number of `packset_` routines, which set the appropriate column from a Python
list of values.

Here's an example of changing an entire column of metadata in one pass (in this case,
note the extra step of validating and encoding before calling
{meth}`IndividualTable.packset_metadata`, as
{ref}`described<sec_tutorial_metadata_bulk>` in the
{ref}`metadata tutorial<sec_tutorial_metadata>`). 

```{code-cell} ipython3
# Add a surname to all individuals
new_metadata = [{"name": i.metadata["name"] + " Smith"} for i in tables.individuals]
validated_metadata = [
    tables.individuals.metadata_schema.validate_and_encode_row(row) for row in new_metadata
]
tables.individuals.packset_metadata(validated_metadata)
edited_ts = tables.tree_sequence()
edited_ts.tables.individuals
```

(sec_tables_editing_major)=

### Major edits

For larger changes, it is often worth clearing the entire table and filling it. Here,
for example, we clear all mutations from the mutations table, then fill it up again,
if possible using a more parsimonious placement of mutations (i.e. a placement that
would better explain the observed variation among samples - see the
{ref}`sec_analysing_trees_parsimony` tutorial section).
This also involves changing the ancestral state at the site if necessary.

```{code-cell} ipython3
# Make a tree sequence with loads of recurrent/back mutations, for testing only
ts = msprime.sim_ancestry(10, sequence_length=100, recombination_rate=0.1, random_seed=123)
ts = msprime.sim_mutations(ts, rate=0.1, random_seed=456)
print("Original ts has", ts.num_mutations, "mutations on", ts.num_trees, "trees")

tables = ts.dump_tables()  # Copy the table collection associated with this tree sequence
tables.mutations.clear()  # Clear all mutations in the table collection copy
variant = tskit.Variant(ts)  # Reuse the same Variant object
for tree in ts.trees():
    for site in tree.sites():
        variant.decode(site.id)  # Efficient if ids are sequential
        anc_state, mutations = tree.map_mutations(variant.genotypes, variant.alleles)
        if len(mutations) <  len(site.mutations):
            # Genotypes can be explained with a more parsimonious distribution of mutations
            if anc_state != site.ancestral_state:
                tables.sites[site.id] = tables.sites[site.id].replace(ancestral_state=anc_state)
            for mut in mutations:
                tables.mutations.append(mut.replace(site=site.id))            
        else:
            # Add the original mutations back in (note that parent ids may have changed,
            # so we will need to call compute_mutation_parents at the end)
            for mut in site.mutations:
                tables.mutations.append(mut)
tables.compute_mutation_parents()
tables.compute_mutation_times()
tables.sort()
parsimonious_ts = tables.tree_sequence()

print("Parsimonious ts has", ts.num_mutations - parsimonious_ts.num_mutations, "fewer muts")

# Optional checks
assert ts.num_trees == parsimonious_ts.num_trees
assert ts.num_sites == parsimonious_ts.num_sites
for v1, v2 in zip(ts.variants(), parsimonious_ts.variants()):
    # convert alleles to numpy arrays so we can index into them easily
    alleles1, alleles2 = np.array(v1.alleles), np.array(v2.alleles)
    # Verify the distribution of variation is the same, by comparing the allelic state
    assert np.all(alleles1[v1.genotypes] == alleles2[v2.genotypes])
```

:::{note}
In the code above we iterate over the trees in the original tree sequence, but make
changes to a *copy* of the tables associated with those trees, modifying the table
collection as we go. This is a common idiom.
:::


--- ../../tskit/docs/getting_started.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

```{currentmodule} tskit
```

(sec_tskit_getting_started)=
# Getting started with {program}`tskit`

You've run some simulations or inference methods, and you now have a 
{class}`TreeSequence` object; what now? This tutorial is aimed 
users who are new to {program}`tskit` and would like to get some
basic tasks completed. We'll look at five fundamental things you might
need to do:
{ref}`process trees<sec_processing_trees>`,
{ref}`sites & mutations<sec_processing_sites_and_mutations>`, and
{ref}`genotypes<sec_processing_genotypes>`,
{ref}`compute statistics<sec_tskit_getting_started_compute_statistics>`, and
{ref}`save or export data<sec_tskit_getting_started_exporting_data>`.
Throughout, we'll also provide pointers to where you can learn more.

:::{note}
The examples in this
tutorial are all written using the {ref}`sec_python_api`, but it's also possible to
{ref}`use R <sec_tskit_r>`, or access the API in other languages, notably
{ref}`C<sec_c_api>` and [Rust](https://github.com/tskit-dev/tskit-rust).
:::

A number of different software programs can generate tree sequences. For the purposes
of this tutorial we'll use [{program}`msprime`](https://tskit.dev/msprime) to create
an example tree sequence representing the genetic genealogy of a 10Mb chromosome in
twenty diploid individuals. To make it a bit more interesting, we'll simulate the effects
of a {ref}`selective sweep<msprime:sec_ancestry_models_selective_sweeps>` in the middle
of the chromosome, then throw some neutral mutations onto the resulting tree sequence.


```{code-cell} ipython3
import msprime

pop_size=10_000
seq_length=10_000_000

sweep_model = msprime.SweepGenicSelection(
    position=seq_length/2, start_frequency=0.0001, end_frequency=0.9999, s=0.25, dt=1e-6)

ts = msprime.sim_ancestry(
    20,
    model=[sweep_model, msprime.StandardCoalescent()],
    population_size=pop_size,
    sequence_length=seq_length,
    recombination_rate=1e-8,
    random_seed=1234,  # only needed for repeatabilty
    )
# Optionally add finite-site mutations to the ts using the Jukes & Cantor model, creating SNPs
ts = msprime.sim_mutations(ts, rate=1e-8, random_seed=4321)
ts
```

You can see that there are many thousands of trees in this tree sequence.

:::{note}
Since we simulated the ancestry of 20 *diploid* individuals, our tree sequence
contains 40 *sample nodes*, one for each genome.
:::

(sec_processing_trees)=

## Processing trees

Moving along a tree sequence usually involves iterating over all of its {class}`Tree`
objects. This common idiom underlies many tree sequence algorithms, including those
we'll encounter later in this tutorial for calculating 
{ref}`population genetic statistics<tskit:sec_stats>`.
To iterate over a tree sequence you can use
{meth}`TreeSequence.trees`.

```{code-cell} ipython3
for tree in ts.trees():
    print(f"Tree {tree.index} covers {tree.interval}")
    if tree.index >= 4:
        print("...")
        break
print(f"Tree {ts.last().index} covers {ts.last().interval}")
```
::::{margin}
:::{caution}
For efficiency, {meth}`~TreeSequence.trees` repeatedly returns the
same tree object, updating it internally to reflect the (usually small) changes
between adjacent trees. So take care to treat each tree within the ``for``
loop separately, avoiding e.g. references outside the loop. The following
produces a list of identical "null" trees:

```
# Don't do this!
list(ts.trees())
```

If you need separate instances of each tree (inefficient, and will
eat up your computer memory), use {meth}`~TreeSequence.aslist`.
:::
::::

In this code snippet, as well as the {meth}`~TreeSequence.trees` iterator, we've also
used {meth}`TreeSequence.last` to access the last tree directly; it may not surprise you
that there's a corresponding {meth}`TreeSequence.first` method to return the first tree.

Above, we stopped iterating after Tree 4 to limit the printed output, but iterating
forwards through trees in a tree sequence (or indeed backwards using the standard Python
{func}`~py:reversed` function) is efficient. That means it's quick, for example to check if all
the trees in a tree sequence have fully coalesced (which is to be expected in
reverse-time, coalescent simulations, but not always for tree sequences produced by
{ref}`forward simulation <sec_tskit_forward_simulations>`).

```{code-cell} ipython3
import time
elapsed = time.time()
for tree in ts.trees():
    if tree.has_multiple_roots:
        print("Tree {tree.index} has not coalesced")
        break
else:
    elapsed = time.time() - elapsed
    print(f"All {ts.num_trees} trees coalesced")
    print(f"Checked in {elapsed:.6g} secs")
```

Now that we know all trees have coalesced, we know that at each position in the genome
all the 40 sample nodes must have one most recent common ancestor (MRCA). Below, we
iterate over the trees, finding the IDs of the root (MRCA) node for each tree. The
time of this root node can be found via the {meth}`tskit.TreeSequence.node` method, which
returns a {class}`Node` object whose attributes include the node time:

```{code-cell} ipython3
import matplotlib.pyplot as plt

kb = [0]  # Starting genomic position
mrca_time = []
for tree in ts.trees():
    kb.append(tree.interval.right/1000)  # convert to kb
    mrca = ts.node(tree.root)  # For msprime tree sequences, the root node is the MRCA
    mrca_time.append(mrca.time)
plt.stairs(mrca_time, kb, baseline=None)
plt.xlabel("Genome position (kb)")
plt.ylabel("Time of root (or MRCA) in generations")
plt.yscale("log")
plt.show()
```

It's obvious that there's something unusual about the trees in the middle of this
chromosome, where the selective sweep occurred. 

Although tskit is designed so that is it rapid to pass through trees sequentially,
it is also possible to pull out individual trees from the middle of a tree sequence
via the {meth}`TreeSequence.at` method. Here's how you can use that to extract
the tree at location $5\ 000\ 000$ --- the position of the sweep --- and draw it
using the {meth}`Tree.draw_svg` method.

```{code-cell} ipython3
swept_tree = ts.at(5_000_000)  # or you can get e.g. the nth tree using ts.at_index(n)
intvl = swept_tree.interval
print(f"Tree number {swept_tree.index}, which runs from position {intvl.left} to {intvl.right}:")
# Draw it at a wide size, to make room for all 40 tips
swept_tree.draw_svg(size=(1000, 200))
```
:::{margin}
The {ref}`visualization tutorial <sec_tskit_viz>` gives more drawing possibilities
:::

This tree shows the classic signature of a recent expansion or selection event, with many
long terminal branches, resulting in an excess of singleton mutations.

:::{margin}
The {ref}`Simplification tutorial<sec_simplification>` details many other uses
for {meth}`~TreeSequence.simplify`.
:::

It can often be helpful to slim down a tree sequence so that it represents the genealogy
of a smaller subset of the original samples. This can be done using the powerful
{meth}`TreeSequence.simplify` method.

The {meth}`TreeSequence.draw_svg` method allows us to draw
more than one tree: either the entire tree sequence, or
(by using the ``x_lim`` parameter) a smaller region of the genome:

```{code-cell} ipython3
reduced_ts = ts.simplify([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])  # simplify to the first 10 samples
print("Genealogy of the first 10 samples for the first 5kb of the genome")
reduced_ts.draw_svg(x_lim=(0, 5000))
```

These are much more standard-looking coalescent trees, with far longer branches higher
up in the tree, and therefore many more mutations at higher-frequencies.

:::{margin}
You cannot directly edit a tree sequence; to add e.g. metadata you must edit a
copy of the underlying tables. This is described in the
{ref}`Tables and editing tutorial<sec_tables_editing>`.
:::

:::{note}
In this tutorial we refer to objects, such as sample nodes, by their numerical IDs. These
can change after simplification, and it is often more meaningful to 
{ref}`work with metadata<sec_tutorial_metadata>`, such as sample and population names,
which can be permanently attached to objects in the tree sequence. Such metadata is
often incorporated automatically by the tools generating the tree sequence. 
:::

(sec_processing_sites_and_mutations)=

## Processing sites and mutations

:::{margin}
See the tutorial entitled "{ref}`sec_tskit_no_mutations`" for why you may not need
sites or mutations in your analyses.
:::

For many purposes it may be better to focus on the genealogy of your samples, rather than
the {ref}`sites<sec_data_model_definitions_site>` and
{ref}`mutations<sec_data_model_definitions_mutation>` that
{ref}`define <sec_what_is_dna_data>` the genome sequence itself. Nevertheless,
{program}`tskit` also provides efficient ways to return {class}`Site` object and
{class}`Mutation` objects from a tree sequence.
For instance, under the finite sites model of mutation that we used above, multiple mutations
can occur at some sites, and we can identify them by iterating over the sites using the
{meth}`TreeSequence.sites` method:

```{code-cell} ipython3
import numpy as np
num_muts = np.zeros(ts.num_sites, dtype=int)
for site in ts.sites():
    num_muts[site.id] = len(site.mutations)  # site.mutations is a list of mutations at the site

# Print out some info about mutations per site
for nmuts, count in enumerate(np.bincount(num_muts)):
    info = f"{count} sites"
    if nmuts > 1:
        info += f", with IDs {np.where(num_muts==nmuts)[0]},"
    print(info, f"have {nmuts} mutation" + ("s" if nmuts != 1 else ""))
```

(sec_processing_genotypes)=

## Processing genotypes

At each site, the sample nodes will have a particular allelic state (or be flagged as
{ref}`tskit:sec_data_model_missing_data`). The
{meth}`TreeSequence.variants` method gives access to the
full variation data. For efficiency, the {attr}`~Variant.genotypes`
at a site are returned as a [numpy](https://numpy.org) array of integers:

```{code-cell} ipython3
import numpy as np
np.set_printoptions(linewidth=200)  # print genotypes on a single line

print("Genotypes")
for v in ts.variants():
    print(f"Site {v.site.id}: {v.genotypes}")
    if v.site.id >= 4:  # only print up to site ID 4
        print("...")
        break
```

:::{note}
Tree sequences are optimised to look at all samples at one site, then all samples at an
adjacent site, and so on along the genome. It is much less efficient look at all the
sites for a single sample, then all the sites for the next sample, etc. In other words,
you should generally iterate over sites, not samples. Nevertheless, all the alleles for
a single sample can be obtained via the
{meth}`TreeSequence.haplotypes` method.
:::


To find the actual allelic states at a site, you can refer to the
{attr}`~Variant.alleles` provided for each {class}`Variant`:
the genotype value is an index into this list. Here's one way to print them out; for
clarity this example also prints out the IDs of both the sample nodes (i.e. the genomes)
and the diploid {ref}`individuals <sec_nodes_or_individuals>` in which each sample
node resides.



````{code-cell} ipython3
samp_ids = ts.samples()
print("  ID of diploid individual: ", " ".join([f"{ts.node(s).individual:3}" for s in samp_ids]))
print("       ID of (sample) node: ", " ".join([f"{s:3}" for s in samp_ids]))
for v in ts.variants():
    site = v.site
    alleles = np.array(v.alleles)
    print(f"Site {site.id} (ancestral state '{site.ancestral_state}')",  alleles[v.genotypes])
    if site.id >= 4:  # only print up to site ID 4
        print("...")
        break
````

:::{note}
Since we have used the {class}`msprime.JC69` model of mutations, the alleles are all
either 'A', 'T', 'G', or 'C'. However, more complex mutation models can involve mutations
such as indels, leading to allelic states which need not be one of these 4 letters, nor
even be a single letter.
:::


(sec_tskit_getting_started_compute_statistics)=

## Computing statistics

There are a {ref}`large number of statistics<tskit:sec_stats>` and related calculations
built in to {program}`tskit`. Indeed, many basic population genetic statistics are based
on the allele (or site) frequency spectrum (AFS), which can be obtained from a tree sequence
using the {meth}`TreeSequence.allele_frequency_spectrum`
method:

```{code-cell} ipython3
afs = ts.allele_frequency_spectrum()
plt.bar(np.arange(ts.num_samples + 1), afs)
plt.title("Unpolarised allele frequency spectrum")
plt.show()
```

By default this method returns the "folded" or unpolarized AFS that doesn't
{ref}`take account of the ancestral state<tskit:sec_stats_polarisation>`.
However, since the tree sequence provides the ancestral state, we can plot the polarized
version; additionally we can base our calculations on branch lengths rather than alleles,
which provides an estimate that is not influenced by random mutational "noise".

```{code-cell} ipython3
fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(12, 3))

afs1 = ts.allele_frequency_spectrum(polarised=True, mode="branch")
ax1.bar(np.arange(ts.num_samples+1), afs1)
ax1.set_title("Genome-wide branch-length AFS")

restricted_ts = ts.keep_intervals([[5e6, 5.5e6]])
afs2 = restricted_ts.allele_frequency_spectrum(polarised=True, mode="branch")
ax2.bar(np.arange(restricted_ts.num_samples+1), afs2)
ax2.set_title("Branch-length AFS between 5 and 5.5Mb")

plt.show()
```

On the left is the frequency spectrum averaged over the entire genome, and on the right
is the spectrum for a section of the tree sequence between 5 and 5.5Mb, which we've
created by deleting the regions outside that interval using
{meth}`TreeSequence.keep_intervals`. Unsurprisingly,
as we noted when looking at the trees, there's a far higher proportion of singletons in
the region of the sweep.

(sec_tskit_getting_started_compute_statistics_windowing)=

### Windowing

It is often useful to see how statistics vary in different genomic regions. This is done
by calculating them in {ref}`tskit:sec_stats_windows` along the genome. For this,
let's look at a single statistic, the genetic {meth}`~TreeSequence.diversity` (π). As a
site statistic this measures the average number of genetic differences between two
randomly chosen samples, whereas as a branch length statistic it measures the average
branch length between them. We'll plot how the value of π changes using 10kb windows,
plotting the resulting diversity between positions 4 and 6 Mb:

```{code-cell} ipython3
fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(12, 3))
L = int(ts.sequence_length)
windows = np.linspace(0, L, num=L//10_000)
ax1.stairs(ts.diversity(windows=windows), windows/1_000, baseline=None)  # Default is mode="site"
ax1.set_ylabel("Diversity")
ax1.set_xlabel("Genome position (kb)")
ax1.set_title("Site-based calculation")
ax1.set_xlim(4e3, 6e3)
ax1.set_yscale("log")
ax1.set_ylim(1e-6, 1e-3)
ax2.stairs(ts.diversity(windows=windows, mode="branch"), windows/1_000, baseline=None)
ax2.set_xlabel("Genome position (kb)")
ax2.set_title("Branch-length-based calculation")
ax2.set_xlim(4e3, 6e3)
ax2.set_yscale("log")
plt.show()
```

There's a clear drop in diversity in the region of the selective sweep. And as expected,
the statistic based on branch-lengths gives a less noisy signal.


(sec_tskit_getting_started_exporting_data)=

## Saving and exporting data

Tree sequences can be efficiently saved to file using {meth}`TreeSequence.dump`, and
loaded back again using {func}`tskit.load`. By convention, we use the suffix ``.trees``
for such files:

```{code-cell} ipython3
import tskit

ts.dump("data/my_tree_sequence.trees")
new_ts = tskit.load("data/my_tree_sequence.trees")
```

It's also possible to export tree sequences to different formats. Note, however, that
not only are these usually much larger files, but that analysis is usually much faster
when performed by built-in tskit functions than by exporting and using alternative
software. If you have a large tree sequence, you should *try to avoid exporting
to other formats*.

### Newick and Nexus format

The most common format for interchanging tree data is Newick. 
We can export to a newick format string quite easily. This can be useful for
interoperating with existing tree processing libraries but is very inefficient for
large trees. There is also no support for including sites and mutations in the trees.

```{code-cell} ipython3
small_ts = reduced_ts.keep_intervals([[0, 10000]])
tree = small_ts.first()
print(tree.newick(precision=3))
```

For an entire set of trees, you can use the Nexus file format, which acts as a container
for a list of Newick format trees, one per line:

```{code-cell} ipython3
small_ts = small_ts.trim()  # Must trim off the blank region at the end of cut-down ts
print(small_ts.as_nexus(precision=3, include_alignments=False))
```

### VCF

The standard way of interchanging genetic variation data is the Variant Call Format, 
for which tskit has basic support:

```{code-cell} ipython3
import sys
small_ts.write_vcf(sys.stdout)
```

The write_vcf method takes a file object as a parameter; to get it to write out to the
notebook here we ask it to write to stdout.


(sec_tskit_getting_started_exporting_data_allel)=

### Scikit-allel

Because tskit integrates very closely with numpy, we can interoperate very efficiently
with downstream Python libraries for working with genetic
sequence data, such as [scikit-allel](https://scikit-allel.readthedocs.io/en/stable/).
We can interoperate with {program}`scikit-allel` by exporting the genotype matrix as a
numpy array, which {program}`scikit-allel` can then process in various ways.

```{code-cell} ipython3
import allel
# Export the genotype data to allel. Unfortunately there's a slight mismatch in the 
# terminology here where genotypes and haplotypes mean different things in the two
# libraries.
h = allel.HaplotypeArray(small_ts.genotype_matrix())
print(h.n_variants, h.n_haplotypes)
h
```

Sckit.allel has a wide-ranging and efficient suite of tools for working with genotype
data, so should provide anything that's needed. For example, it gives us an
another way to compute the pairwise diversity statistic (that we calculated
{ref}`above<sec_tskit_getting_started_compute_statistics_windowing>`
using the native {meth}`TreeSequence.diversity` method):

```{code-cell} ipython3
ac = h.count_alleles()
allel.mean_pairwise_difference(ac)
```


(sec_tskit_getting_started_key_points)=

## Key points covered above

Some simple methods and take-home messages from this introduction to the
{program}`tskit` {ref}`sec_python_api`,
in rough order of importance:

* Objects and their attributes
    * In Python, a {class}`TreeSequence` object has a number of basic attributes such as
        {attr}`~TreeSequence.num_trees`, {attr}`~TreeSequence.num_sites`,
        {attr}`~TreeSequence.num_samples`, {attr}`~TreeSequence.sequence_length`, etc.
        Similarly a {class}`Tree` object has e.g. an {attr}`~Tree.interval` attribute, a
        {class}`Site` object has a {attr}`~Site.mutations` attribute, a {class}`Node`
        object has a {attr}`~Node.time` attribute, and so on.
    * {ref}`sec_terminology_nodes` (i.e. genomes) can belong to
        {ref}`individuals<sec_terminology_individuals_and_populations>`. For example,
        sampling a diploid individual results in an {class}`Individual` object which
        possesses two distinct {ref}`sample nodes<sec_terminology_nodes_samples>`.
* Key tree sequence methods
    * {meth}`~TreeSequence.samples()` returns an array of node IDs specifying the
        nodes that are marked as samples
    * {meth}`~TreeSequence.node` returns the node object for a given integer node ID
    * {meth}`~TreeSequence.trees` iterates over all the trees
    * {meth}`~TreeSequence.sites` iterates over all the sites
    * {meth}`~TreeSequence.variants` iterates over all the sites with their genotypes
        and alleles
    * {meth}`~TreeSequence.simplify()` reduces the number of sample nodes in the tree
        sequence to a specified subset
    * {meth}`~TreeSequence.keep_intervals()` (or its complement,
        {meth}`~TreeSequence.delete_intervals()`) removes genetic information from
        specific regions of the genome
    * {meth}`~TreeSequence.draw_svg()` returns an SVG representation of a tree sequence
        (and plots it if in a Jupyter notebook). Similarly, {meth}`Tree.draw_svg()`
        plots individual trees.
    * {meth}`~TreeSequence.at()` returns a tree at a particular genomic position
        (but using {meth}`~TreeSequence.trees` is usually preferable)
    * Various population genetic statistics can be calculated using methods on a tree
        sequence, for example {meth}`~TreeSequence.allele_frequency_spectrum`,
        {meth}`~TreeSequence.diversity`, and {meth}`~TreeSequence.Fst`; these can
        also be calculated in windows along the genome.



--- ../../tskit/docs/development.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

```{currentmodule} tskit
```

(sec_development)=


# Development

If you would like to add some features to `tskit`, this
documentation should help you get set up and contributing.
Please help us to improve the documentation by either
opening an [issue](http://github.com/tskit-dev/tskit/issues) or
[pull request](http://github.com/tskit-dev/tskit/pulls) if you
see any problems.

The tskit-dev team strives to create a welcoming and open environment for
contributors; please see our
[code of conduct](https://github.com/tskit-dev/.github/blob/main/CODE_OF_CONDUCT.md) for
details. We wish our code and documentation to be
[inclusive](https://chromium.googlesource.com/chromium/src/+/master/styleguide/inclusive_code.md)
and in particular to be gender and racially neutral.

(sec_development_structure)=


## Project structure

Tskit is a multi-language project, which is reflected in the directory
structure:

- The `python` directory contains the Python library and command line interface,
  which is what most contributors are likely to be interested in. Please
  see the {ref}`sec_development_python` section for details. The
  low-level {ref}`sec_development_python_c` is also defined here.

- The `c` directory contains the high-performance C library code. Please
  see the {ref}`sec_development_c` for details on how to contribute.

- The `docs` directory contains the source for this documentation,
  which covers both the Python and C APIs. Please see the {ref}`sec_development_documentation`
  for details.

The remaining files in the root directory of the project are for
controlling {ref}`sec_development_continuous_integration` providers
and other administrative purposes.

Please see the {ref}`sec_development_best_practices` section for
an overview of how to contribute a new feature to `tskit`.

(sec_development_getting_started)=


## Getting started

(sec_development_getting_started_requirements)=


### Requirements

To develop the Python code you will need a working C compiler and a
development installation of Python (>= 3.9). On Debian/Ubuntu we can install these
with:

```bash
$ sudo apt install python3-dev python-is-python3 build-essential doxygen
```

Python packages required for development are listed in `python/requirements/development.txt`.
These can be installed using `pip`::

```bash
$ python3 -m pip install -r python/requirements/development.txt
```

You may wish isolate your development environment using a `virtualenv
<https://docs.python-guide.org/dev/virtualenvs/>`_.

A few extra dependencies are required if you wish to work on the
{ref}`C library <sec_development_c_requirements>`.

For OSX and Windows users we recommending using
[conda](https://docs.conda.io/projects/conda/en/latest/)_,
and isolating development in a dedicated environment as follows::

```bash
$ conda env create -f python/requirements/development.yml
$ conda activate tskit-dev
```

On macOS, conda builds are generally done using `clang` packages that are kept up to date:

```bash
$ conda install clang_osx-64  clangxx_osx-64
```

In order to make sure that these compilers work correctly (*e.g.*, so that they can find
other dependencies installed via `conda`), you need to compile `tskit` with this command
on versions of macOS older than "Mojave":

```bash
$ cd python
$ CONDA_BUILD_SYSROOT=/ python3 setup.py build_ext -i
```

On more recent macOS releases, you may omit the `CONDA_BUILD_SYSROOT` prefix.

If you run into issues with the conda compiler, be sure that your command line tools are installed
and up to date (you should also reboot your system after installing CLI tools). Note that you may
also have to install a
[specific version of the Xcode command line tools](https://stackoverflow.com/a/64416852/2752221).

:::{note}
The use of the C toolchain on macOS is a moving target.  The above advice
was updated on 22 June, 2021 and was validated by a few `tskit` contributors.
Caveat emptor, etc..
:::

(sec_development_getting_started_environment)=


### Environment

To get a local git development environment, please follow these steps:

- Make a fork of the tskit repo on [GitHub](http://github.com/tskit-dev/tskit)
- Clone your fork into a local directory:
  ```bash
  $ git clone git@github.com:YOUR_GITHUB_USERNAME/tskit.git
  ```
- Install the {ref}`sec_development_workflow_pre_commit`:
  ```bash
  $ pre-commit install
  ```
  If you later have trouble with these checks, you can skip them with ``git commit --no-verify``.

See the {ref}`sec_development_workflow_git` section for detailed information
on the recommended way to use git and GitHub.

(sec_development_workflow)=


## Workflow

(sec_development_workflow_git)=


### Git workflow

If you would like to make an addition/fix to tskit, then follow the steps below
to get things set up.
If you would just like to review someone else's proposed changes
(either to the code or to the docs), then
skip to {ref}`sec_development_workflow_anothers_commit`.

0.  Open an [issue](http://github.com/tskit-dev/tskit/issues) with your proposed
    functionality/fix. If adding or changing the public API close thought should be given to
    names and signatures of proposed functions. If consensus is reached that your
    proposed addition should be added to the codebase, proceed!

1. Make your own [fork](https://help.github.com/articles/fork-a-repo/)
   of the `tskit` repository on GitHub, and
   [clone](https://help.github.com/articles/cloning-a-repository/)
   a local copy as detailed in {ref}`sec_development_getting_started_environment`.

2. Make sure that your local repository has been configured with an
   [upstream remote](
   https://help.github.com/articles/configuring-a-remote-for-a-fork/):
   ```bash
   $ git remote add upstream https://github.com/tskit-dev/tskit.git
   ```

3. Create a "topic branch" to work on. One reliable way to do it
   is to follow this recipe:
   ```bash
   $ git fetch upstream
   $ git checkout upstream/main
   $ git checkout -b topic_branch_name
   ```

4. Write your code following the outline in {ref}`sec_development_best_practices`.
   As you work on your topic branch you can add commits to it. Once you're
   ready to share this, you can then open a
   [pull request (PR)](https://help.github.com/articles/about-pull-requests/). This can be done at any
   time! You don't have to have code that is completely functional and tested to get
   feedback. Use the drop-down button to create a "draft PR" to indicate that it's not
   done, and explain in the comments what feedback you need and/or what you think needs
   to be done.

5. As you code it is best to
   [rebase](https://stdpopsim.readthedocs.io/en/latest/development.html#rebasing) your
   work onto the `main` branch periodically (e.g. once a week) to keep up with changes.
   If you merge `main` via `git pull upstream main`
   it will create a much more complex rebase when your code is finally ready to be
   incorporated into the main branch, so should be avoided.

6. Once you're done coding add content to the tutorial and other documentation pages if
   appropriate.

7. Update the change logs at `python/CHANGELOG.rst` and `c/CHANGELOG.rst`, taking care
   to document any breaking changes separately in a "breaking changes" section.

8. Push your changes to your topic branch and either open the PR or, if you
   opened a draft PR above change it to a non-draft PR by clicking "Ready to
   Review".

9. The tskit community will review the code, asking you to make changes where appropriate.
   This usually takes at least two rounds of review.

10. Once the review process is complete, squash the commits to the minimal set of changes -
    usually one or two commits. Please follow
    [this guide](https://stdpopsim.readthedocs.io/en/stable/development.html#rebasing) for
    step-by-step instructions on rebasing and squashing commits.

11. Your PR will be merged, time to celebrate! 🎉🍾


(sec_development_workflow_anothers_commit)=


### Checking out someone else's pull request

Sometimes you want to just check out someone else's pull request,
for the purpose of trying it out and giving them feedback.
To do this, you first need your own local version of the git repository,
so you should first do steps 1 and 2 above.
(Strictly speaking, you don't need a fork on github
if you don't plan to edit, but it won't hurt.)
Continuing from there, let's say you want to check out the current
state of the code on [pull request #854](https://github.com/tskit-dev/tskit/pull/854).
(So, below you should replace `854` with the number of the pull request
that you actually want to investigate.)
Then, continuing from above:

3. Fetch the pull request, and store it as a local branch.
   For instance, to name the local branch `my_pr_copy`:
   ```bash
   $ git fetch upstream pull/854/head:my_pr_copy
   ```
   You should probably call the branch something more descriptive,
   though. (Also note that you might need to put `origin` instead
   of `upstream` for the remote repository name: see `git remote -v`
   for a list of possible remotes.)

4. Check out the pull request's local branch:
   ```bash
   $ git checkout my_pr_copy
   ```

Now, your repository will be in exactly the same state as
that of the person who's submitted the pull request.
Great! Now you can test things out.

To view the documentation,
`cd docs && make`, which should build the documentation,
and then navigate your web browser to the `docs/_build/html/`
subdirectory.

To test out changes to the *code*, you can change to the `python/` subdirectory,
and run `make` to compile the C code.
If you then execute `python` from this subdirectory (and only this one!),
it will use the modified version of the package.
(For instance, you might want to
open an interactive `python` shell from the `python/` subdirectory,
or running `python3 -m pytest` from this subdirectory.)

After you're done, you should do:

```bash
$ git checkout main
```

to get your repository back to the "main" branch of development.
If the pull request is changed and you want to do the same thing again,
then first *delete* your local copy (by doing `git branch -d my_pr_copy`)
and repeat the steps again.


(sec_development_workflow_pre_commit)=


### Pre-commit checks

On each commit a [pre-commit hook](https://pre-commit.com/)  will run
that checks for violations of code style
(see the {ref}`sec_development_python_style` section for details)
and other common problems.
Where possible, these hooks will try to fix any problems that they find (including reformatting
your code to conform to the required style). In this case, the commit
will *not complete* and report that "files were modified by this hook".
To include the changes that the hooks made, `git add` any
files that were modified and run `git commit` (or, use `git commit -a`
to commit all changed files.)

If you would like to run the checks without committing, use `pre-commit run`
(but, note that this will *only* check changes that have been *staged*;
do `pre-commit run --all` to check unstaged changes as well).
To bypass the checks (to save or get feedback on work-in-progress) use
`git commit --no-verify`

(sec_development_documentation)=


## Documentation

The documentation for tskit is written using
[Sphinx](http://www.sphinx-doc.org/en/stable/)
and contained in the `docs` directory. The files in this directory are
markdown files that serve as an input to [jupyterbook](https://jupyterbook.org/),
which allows jupyter notebook code, primarily in Python, to be automatically
executed and the output inserted before deployment. The docs are then
deployed automatically to the [tskit.dev website](https://tskit.dev/).
API documentation for both Python and C are generated automatically from
source: documentation embedded in the source code makes use of sphinx and
the [reStructuredText](http://docutils.sourceforge.net/rst.html) format to
alloow formating and cross referencing.
For the C code, a combination of [Doxygen](http://www.doxygen.nl/)
and [breathe](https://breathe.readthedocs.io/en/latest/) is used to
generate API documentation.

Please help us to improve the documentation! You can check on the list of
[documentation issues](https://github.com/tskit-dev/tskit/issues?q=is%3Aissue+is%3Aopen+label%3Adocumentation)
on GitHub, and help us fix any, or add issues for anything that's wrong or missing.


### Small edits

If you see a typo or some other small problem that you'd like to fix,
this is most easily done through the GitHub UI.

If the typo is in a large section of text (like this page), go to the
top of the page and click on the "Edit on GitHub" link at the top
right. This will bring you to the page on GitHub for the RST
source file in question. Then, click on the pencil icon on the
right hand side. This will open a web editor allowing you to
quickly fix the typo and submit a pull request with the changes.
Fix the typo, add a commit message like "Fixed typo" and click
on the green "Propose file change" button. Then follow the dialogues
until you've created a new pull request with your changes,
so that we can incorporate them.

If the change you'd like to make is in the API documentation
for a particular function, then you'll need to find where this
function is defined first. The simplest way to do this is
to click the green "[source]" link next to the function. This
will show you a HTML rendered version of the function, and the
rest of the file that it is in. You can then navigate to this
file on GitHub, and edit it using the same approach as above.


### Significant edits

When making changes more substantial than typo fixes it's best
to check out a local copy.
Follow the steps in the {ref}`sec_development_workflow_git` to
get a fork of tskit, a local clone and newly checked out
feature branch. Then follow the steps in the
{ref}`sec_development_getting_started` section to get a
working development environment.

Once you are ready to make edits to the documentation,
`cd` into the `docs` directory and run `make`.
This should build the HTML
documentation in `docs/_build/html/`, which you can then view in
your browser. As you make changes, run `make` regularly and
view the final result to see if it matches your expectations.

Once you are happy with the changes, commit your updates and
open a pull request on GitHub.


### Tips and resources

- The reStructuredText
  [primer](https://www.sphinx-doc.org/en/1.8/usage/restructuredtext/basics.html)
  is a useful general resource on rst.

- See also the sphinx and rst [cheatsheet
 ](https://docs.typo3.org/m/typo3/docs-how-to-document/master/en-us/WritingReST/CheatSheet.html)

- The Sphinx Python and C
  [domains](https://www.sphinx-doc.org/en/master/usage/restructuredtext/domains.html)
  have extensive options for marking up code.

- Make extensive use of
  [cross referencing](https://docs.typo3.org/m/typo3/docs-how-to-document/master/en-us/WritingReST/Hyperlinks.html).
  When linking to sections in the documentation, use the
  `` :ref:`sec_some_section_label` `` form rather than matching on the section
  title (which is brittle). Use `` :meth:`.Tree.some_method` ``,
  `` :func:`some_function` `` etc to refer to parts of the API.

(sec_development_python)=


## Python library

The Python library is defined in the `python` directory. We assume throughout
this section that you have `cd`'d into this directory.
We also assume that the `tskit` package is built and
run locally *within* this directory. That is, `tskit` is *not* installed
into the Python installation using `pip install -e` or setuptools
[development mode](http://setuptools.readthedocs.io/en/latest/setuptools.html#id23).
Please see the {ref}`sec_development_python_troubleshooting` section for help
if you encounter problems with compiling or running the tests.


### Getting started

After you have installed the basic {ref}`sec_development_getting_started_requirements`
and created a {ref}`development environment <sec_development_getting_started_environment>`,
you will need to compile the low-level {ref}`sec_development_python_c` module.
This is most easily done using `make`:

```bash
$ make
```

If this has completed successfully you should see a file `_tskit.cpython-XXXXXX.so`
in the current directory (the suffix depends on your platform and Python version;
with Python 3.11 on Linux it's `_tskit.cpython-311-x86_64-linux-gnu.so`).

To make sure that your development environment is working, run some
{ref}`tests <sec_development_python_tests>`.


### Layout

Code for the `tskit` module is in the `tskit` directory. The code is split
into a number of modules that are roughly split by function; for example,
code for visualisation is kept in the `tskit/drawing.py`.

Test code is contained in the `tests` directory. Tests are also roughly split
by function, so that tests for the `drawing` module are in the
`tests/test_drawing.py` file. This is not a one-to-one mapping, though.

The `requirements` directory contains descriptions of the requirements
needed for development and on various
{ref}`sec_development_continuous_integration` providers.

(sec_development_python_style)=


### Code style

Python code in tskit is formatted using [Black](https://github.com/psf/black).
Any code submitted as a pull request will be checked to
see if it conforms to this format as part of the
{ref}`sec_development_continuous_integration`. Black is very strict, which seems
unhelpful and nitpicky at first but is actually very useful. This is because it
can also automatically *format* code for you, removing tedious and subjective
choices (and even more tedious and subjective discussions!)
about exactly how a particular code block should be aligned and indented.

In addition to Black autoformatting, code is checked for common problems using
[flake8](https://flake8.pycqa.org/en/latest/)

Black autoformatting and flake8 checks are performed as part of the
{ref}`pre-commit checks <sec_development_workflow_pre_commit>`, which
ensures that your code is always formatted correctly.

Vim users may find the
[black](https://github.com/psf/black),
and
[vim-flake8](https://github.com/nvie/vim-flake8)
plugins useful for automatically formatting code and lint checking
within vim.
There is good support for Black in a number of
[other editors](https://black.readthedocs.io/en/stable/editor_integration.html#editor-integration).


(sec_development_python_tests)=


### Tests

The tests are defined in the `tests` directory, and run using
[pytest](https://docs.pytest.org/en/stable/) from the `python` directory.
If you want to run the tests in a particular module (say, `test_tables.py`), use:

```bash
$ python3 -m pytest tests/test_tables.py
```

To run all the tests in a particular class in this module (say, `TestNodeTable`)
use:

```bash
$ python3 -m pytest tests/test_tables.py::TestNodeTable
```

To run a specific test case in this class (say, `test_copy`) use:

```bash
$ python3 -m pytest tests/test_tables.py::TestNodeTable::test_copy
```

You can also run tests with a keyword expression search. For example this will
run all tests that have `TestNodeTable` but not `copy` in their name:

```bash
$ python3 -m pytest -k "TestNodeTable and not copy"
```

When developing your own tests, it is much quicker to run the specific tests
that you are developing rather than rerunning large sections of the test
suite each time.

To run all of the tests, we can use:

```bash
$ python3 -m pytest
```

By default the tests are run on 4 cores, if you have more you can specify:

```bash
$ python3 -m pytest -n8
```

A few of the tests take most of the time, we can skip the slow tests to get the test run
under 20 seconds on an modern workstation:

```bash
$ python3 -m pytest --skip-slow
```

If you have a lot of failing tests it can be useful to have a shorter summary
of the failing lines:

```bash
$ python3 -m pytest --tb=line
```

If you need to see the output of tests (e.g. `print` statements) then you need to use
these flags to run a single thread and capture output:

```bash
$ python3 -m pytest -n0 -vs
```

All new code must have high test coverage, which will be checked as part of the
{ref}`sec_development_continuous_integration`
tests by [CodeCov](https://codecov.io/gh/tskit-dev/tskit/).
All tests must pass for a PR to be accepted.


### Packaging

The `tskit` Python module follows the current
[best-practices](http://packaging.python.org) advocated by the
[Python Packaging Authority](http://pypa.io/en/latest/). The primary means of
distribution is though [PyPI](http://pypi.python.org/pypi/tskit), which provides the
canonical source for each release.

A package for [conda](http://conda.io/docs/) is also available on
[conda-forge](https://github.com/conda-forge/tskit-feedstock).


### Interfacing with low-level module

Much of the high-level Python code only exists to provide a simpler interface to
the low-level {ref}`_tskit <sec_development_python_c>` module.
As such, many objects (e.g. {class}`.Tree`)
are really just a shallow layer on top of the corresponding low-level object.
The usual convention here is to keep a reference to the low-level object via
a private instance variable such as `self._ll_tree`.


### Command line interface

The command line interface for `tskit` is defined in the `tskit/cli.py` file.
The CLI has a single entry point (e.g. `tskit_main`) which is invoked to run the
program. These entry points are registered with `setuptools` using the
`console_scripts` argument in `setup.py`, which allows them to be deployed as
first-class executable programs in a cross-platform manner.

The CLI can also be run using `python3 -m tskit`. This is the recommended
approach for running the CLI during development.

(sec_development_installing)=

### Installing development versions

We **strongly** recommend that you do not install development versions of
`tskit` and instead use versions released to PyPI and conda-forge.
However, if you really need to be on the bleeding edge, you can use
the following command to install:

```bash
$ python3 -m pip install git+https://github.com/tskit-dev/tskit.git#subdirectory=python
```

(Because the Python package is not defined in the project root directory, using pip to
install directly from  GitHub requires you to specify `subdirectory=python`.)


(sec_development_python_troubleshooting)=

### Troubleshooting

- If `make` is giving you strange errors, or if tests are failing for
  strange reasons, try running `make clean` in the project root
  and then rebuilding.
- Beware of multiple versions of the python library installed by different
  programs (e.g., pip versus installing locally from source)! In python,
  `tskit.__file__` will tell you the location of the package that is being
  used.
- Installation of development version is not supported in Windows. Windows
  users should try using a Linux envronment by using
  [WSL](https://learn.microsoft.com/windows/wsl/), for example.


(sec_development_c)=

## C Library

The Python module uses the high-performance tskit {ref}`sec_c_api`
behind the scenes. All C code and associated development infrastructure
is held in the `c` directory.


(sec_development_c_requirements)=

### Requirements

We use the
[meson](https://mesonbuild.com) build system in conjunction with
[ninja-build](https://ninja-build.org) to compile the C code.
Unit tests use the [CUnit](http://cunit.sourceforge.net) library
and we use [clang-format](https://clang.llvm.org/docs/ClangFormat.html)
to automatically format code.
On Debian/Ubuntu, these can be installed using

```bash
$ sudo apt install libcunit1-dev ninja-build meson clang-format-6.0
```

**Notes:**

1. A more recent version of meson can alternatively be installed using `pip`, if you wish.
2. Recent versions of Debian do not have clang-format-6.0 available;
    if so, you can install it instead with `pip` by running
    `pip3 install clang-format==6.0.1 && ln -s clang-format $(which clang-format)-6.0`.

Conda users can install the basic requirements from `python/requirements/development.txt`.

Unfortunately clang-format is not available on conda, but it is not essential.


(sec_development_c_code_style)=

### Code style

C code is formatted using
[clang-format](https://clang.llvm.org/docs/ClangFormat.html)
with a custom configuration and version 6.0. This is checked as part of the pre-commit
checks. To manually format run:

```bash
$ clang-format-6.0 -i c/tskit/* c/tests/*.c c/tests/*.h
```

Vim users may find the
[vim-clang-format](https://github.com/rhysd/vim-clang-format)
plugin useful for automatically formatting code.


### Building

We use [meson](https://mesonbuild.com) and [ninja-build](https://ninja-build.org) to
compile the C code. Meson keeps all compiled binaries in a build directory (this has many advantages
such as allowing multiple builds with different options to coexist). The build configuration
is defined in `meson.build`. To set up the initial build
directory, run

```bash
$ cd c
$ meson build
```

To compile the code run

```bash
$ ninja -C build
```

All the tests and other artefacts are in the build directory. Individual test
suites can be run, via (e.g.) `./build/test_trees`. To run all of the tests,
run

```bash
$ ninja -C build test
```

For vim users, the [mesonic](https://www.vim.org/scripts/script.php?script_id=5378) plugin
simplifies this process and allows code to be compiled seamlessly within the
editor.


### Unit Tests

The C-library has an extensive suite of unit tests written using
[CUnit](http://cunit.sourceforge.net). These tests aim to establish that the
low-level APIs work correctly over a variety of inputs, and particularly, that
the tests don't result in leaked memory or illegal memory accesses. All tests
are run under valgrind to make sure of this as part of the
{ref}`sec_development_continuous_integration`.

Tests are defined in the `tests/*.c` files. These are roughly split by
the source files, so that the tests for functionality in the `tskit/tables.c` file
will be tested in `tests/test_tables.c`.
To run all the tests
in the `test_tables` suite, run (e.g.) `./build/test_tables`.
To just run a specific test on its own, provide
this test name as a command line argument, e.g.:

```bash
$ ./build/test_tables test_node_table
```

While 100% test coverage is not feasible for C code, we aim to cover all code
that can be reached. (Some classes of error such as malloc failures
and IO errors are difficult to simulate in C.) Code coverage statistics are
automatically tracked using [CodeCov](https://codecov.io/gh/tskit-dev/tskit/).


### Coding conventions

The code is written using the [C99](https://en.wikipedia.org/wiki/C99) standard. All
variable declarations should be done at the start of a function, and functions
kept short and simple where at all possible.

No global or module level variables are used for production code.

Function parameters should be marked as ``const`` where possible.
Parameters that are used as return variables should come last.
The common ``options`` parameter should be the last non-output
parameter.

Please see the {ref}`sec_c_api_overview_structure` section for more information
about how the API is structured.

### Error handling

A critical element of producing reliable C programs is consistent error handling
and checking of return values. All return values **must** be checked! In tskit,
all functions (except the most trivial accessors) return an integer to indicate
success or failure. Any negative value is an error, and must be handled accordingly.
The following pattern is canonical:

```C
   ret = tsk_tree_do_something(self, argument);
    if (ret != 0) {
        goto out;
    }
    // rest of function
out:
    return ret;
```

Here we test the return value of `tsk_tree_do_something` and if it is non-zero,
abort the function and return this same value from the current function. This
is a bit like throwing an exception in higher-level languages, but discipline
is required to ensure that the error codes are propagated back to the original
caller correctly.

Particular care must be taken in functions that allocate memory, because
we must ensure that this memory is freed in all possible success and
failure scenarios. The following pattern is used throughout for this purpose:

```C
    double *x = NULL;

    x = malloc(n * sizeof(double));
    if (x == NULL) {
        ret = TSK_ERR_NO_MEMORY;
        goto out;
    }
    // rest of function
out:
    tsk_safe_free(x);
    return ret;
```

It is vital here that `x` is initialised to `NULL` so that we are guaranteed
correct behaviour in all cases. For this reason, the convention is to declare all
pointer variables on a single line and to initialise them to `NULL` as part
of the declaration.

Error codes are defined in `core.h`, and these can be translated into a
message using `tsk_strerror(err)`.


#### Using assertions

There are two different ways to express assertions in tskit code.
The first is using the custom `tsk_bug_assert` macro, which is used to
make inexpensive checks at key points during execution. These assertions
are always run, regardless of the compiler settings, and should not
contribute significantly to the overall runtime.

More expensive assertions, used, for example, to check pre and post conditions
on performance critical loops should be expressed using the standard
`assert` macro from `assert.h`. These assertions will be checked
during the execution of C unit tests, but will not be enabled when
compiled into the Python C module.


### Type conventions

- `tsk_id_t` is an ID for any entity in a table.
- `tsk_size_t` refers to any size or count values in tskit.
- `size_t` is a standard C type and refers to the size of a memory block.
  This should only be used when computing memory block sizes for functions
  like `malloc` or passing the size of a memory buffer as a parameter.
- Error indicators (the return type of most functions) are `int`.
- `uint32_t` etc should be avoided (any that exist are a leftover from older
  code that didn't use `tsk_size_t` etc.)
- `int64_t` and `uint64_t` are sometimes useful when working with
  bitstrings (e.g. to implement a set).

(sec_development_python_c)=


## Python C Interface


### Overview

The Python C interface is defined in the `python` directory
and written using the [Python C API](https://docs.python.org/3.6/c-api/).
The source code for this interface is in the `_tskitmodule.c` file.
When compiled, this produces the `_tskit` module,
which is imported by the high-level Python code. The low-level Python module is
not intended to be used directly by users and may change arbitrarily over time.

The usual pattern in the low-level Python API is to define a Python class
which corresponds to a given "class" in the C API. For example, we define
a `TreeSequence` class, which is essentially a thin wrapper around the
`tsk_tree_t` type from the C library.

The `_tskitmodule.c` file follows the standard conventions given in the
[Python documentation](https://docs.python.org/3.6/extending/index.html).


### Compiling and debugging

The `setup.py` file describes the requirements for the low-level `_tskit`
module and how it is built from source. The simplest way to compile the
low-level module is to run `make` in the `python` directory:

```bash
$ make
```

If `make` is not available, you can run the same command manually:

```bash
$ python3 setup.py build_ext --inplace
```

It is sometimes useful to specify compiler flags when building the low
level module. For example, to make a debug build you can use:

```bash
$ CFLAGS='-Wall -O0 -g' make
```

If you need to track down a segfault etc, running some code through gdb can
be very useful. For example, to run a particular test case, we can do:


```bash
$ gdb python3
(gdb) run -m pytest tests/test_lowlevel.py


(gdb) run  -m pytest -vs tests/test_tables.py::TestNodeTable::test_copy
Starting program: /usr/bin/python3 run  -m pytest tests/test_tables.py::TestNodeTable::test_copy
[Thread debugging using libthread_db enabled]
Using host libthread_db library "/lib/x86_64-linux-gnu/libthread_db.so.1".
[New Thread 0x7ffff1e48700 (LWP 1503)]
[New Thread 0x7fffef647700 (LWP 1504)]
[New Thread 0x7fffeee46700 (LWP 1505)]
[Thread 0x7fffeee46700 (LWP 1505) exited]
[Thread 0x7fffef647700 (LWP 1504) exited]
[Thread 0x7ffff1e48700 (LWP 1503) exited]
collected 1 item

tests/test_tables.py::TestNodeTable::test_copy PASSED

[Inferior 1 (process 1499) exited normally]
(gdb)
```

Tracing problems in C code is many times more difficult when the Python C API
is involved because of the complexity of Python's memory management. It is
nearly always best to start by making sure that the tskit C API part of your
addition is thoroughly tested with valgrind before resorting to the debugger.


### Testing for memory leaks

The Python C API can be subtle, and it is easy to get the reference counting wrong.
The `stress_lowlevel.py` script makes it easier to track down memory leaks
when they do occur. The script runs the unit tests in a loop, and outputs
memory usage statistics.


(sec_development_continuous_integration)=


## Continuous Integration tests

A number of different continuous integration providers are used, which run different
combinations of tests on different platforms, as well as running various
checks for code quality.

- A [Github action](https://help.github.com/en/actions) runs some code style and
  quality checks along with running the Python test suite on Linux, OSX and Windows. It
  uses conda for those dependencies which are tricky to compile on all systems. An
  additional action builds the docs and posts a link to preview them.

- [CircleCI](https://circleci.com/) runs all Python tests using the apt-get
  infrastructure for system requirements. We also runs C tests, compiled
  using gcc and clang, and check for memory leaks using valgrind.

- [CodeCov](https://codecov.io/gh)_ tracks test coverage in Python and C.

- [PyUp](https://pyup.io/) Runs monthly checks on the Python dependencies listed in the
  requirements files, which are pinned to ensure CI reproducibility. PyUp opens one PR
  a month with updated pins.


(sec_development_best_practices)=


## Best Practices for Development

The following is a rough guide of best practices for contributing a function to the
tskit codebase.

Note that this guide covers the most complex case of adding a new function to both
the C and Python APIs.

1.  Write your function in Python: in `python/tests/` find the test module that
    pertains to the functionality you wish to add. For instance, the kc_distance
    metric was added to
    [test_topology.py](https://github.com/tskit-dev/tskit/blob/main/python/tests/test_topology.py).
    Add a python version of your function here.
2.  Create a new class in this module to write unit tests for your function: in addition
    to making sure that your function is correct, make sure it fails on inappropriate inputs.
    This can often require judgement. For instance, {meth}`Tree.kc_distance` fails on a tree
    with multiple roots, but allows users to input parameter values that are nonsensical,
    as long as they don't break functionality. See the
    [TestKCMetric](https://github.com/tskit-dev/tskit/blob/4e707ea04adca256036669cd852656a08ec45590/python/tests/test_topology.py#L293) for example.
3.  Write your function in C: check out the {ref}`sec_c_api` for guidance. There
    are also many examples in the
    [c directory](https://github.com/tskit-dev/tskit/tree/main/c/tskit).
    Your function will probably go in
    [trees.c](https://github.com/tskit-dev/tskit/blob/main/c/tskit/trees.c).
4.  Write a few tests for your function in C: again, write your tests in
    [tskit/c/tests/test_tree.c](https://github.com/tskit-dev/tskit/blob/main/c/tests/test_trees.c).
    The key here is code coverage, you don't need to worry as much about covering every
    corner case, as we will proceed to link this function to the Python tests you
    wrote earlier.
5.  Create a low-level definition of your function using Python's C API: this will
    go in [_tskitmodule.c](https://github.com/tskit-dev/tskit/blob/main/python/_tskitmodule.c).
6.  Test your low-level implementation in [tskit/python/tests/test_lowlevel.py
   ](https://github.com/tskit-dev/tskit/blob/main/python/tests/test_lowlevel.py):
    again, these tests don't need to be as comprehensive as your first python tests,
    instead, they should focus on the interface, e.g., does the function behave
    correctly on malformed inputs?
7.  Link your C function to the Python API: write a function in tskit's Python API,
    for example the kc_distance function lives in
    [tskit/python/tskit/trees.py](https://github.com/tskit-dev/tskit/blob/main/python/tskit/trees.py).
8.  Modify your Python tests to test the new C-linked function: if you followed
    the example of other tests, you might need to only add a single line of code
    here. In this case, the tests are well factored so that we can easily compare
    the results from both the Python and C versions.
9.  Write a docstring for your function in the Python API: for instance, the kc_distance
    docstring is in
    [tskit/python/tskit/trees.py](https://github.com/tskit-dev/tskit/blob/main/python/tskit/trees.py).
    Ensure that your docstring renders correctly by building the documentation
    (see {ref}`sec_development_documentation`).


## Troubleshooting

### pre-commit is blocking me!

You might be having a hard time committing because of the "pre-commit" checks
(described above). First, consider: the pre-commit hooks are supposed to make your life *easier*,
not add a layer of frustration to contributing.
So, you should feel free to just ask git to skip the pre-commit!
There's no shame in a broken build - you can get it fixed up (and we'll help)
before it's merged into the rest of the project.
To skip, just append `--no-verify` to the `git commit` command.
Below are some more specific situations.


### pre-commit complains about files I didn't edit

For instance, suppose you have *not* edited `util.py` and yet:

```bash
> git commit -a -m 'dev docs'
python/tskit/util.py:117:26: E203 whitespace before ':'
python/tskit/util.py:135:31: E203 whitespace before ':'
python/tskit/util.py:195:23: E203 whitespace before ':'
python/tskit/util.py:213:36: E203 whitespace before ':'
... lots more, gah, what is this ...
```

First, check (with `git status`) that you didn't actually edit `util.py`.
Then, you should **not** try to fix these errors; this is **not your problem**.
You might first try restarting your pre-commit, by running

```bash
pre-commit clean
pre-commit gc
```

You might also check you don't have other pre-commit hook files in `.git/hooks`.
If this doesn't fix the problem,
then you should just *skip* the pre-commit (but alert us to the problem),
by appending `--no-verify`:

```bash
> git commit -a -m 'dev docs' --no-verify
[main 46f3f2e] dev docs
 1 file changed, 43 insertions(+)
```

Now you can go ahead and push your changes!


### pre-commit won't run

For instance:

```bash
> git commit -a -m 'fixed all the things'
/usr/bin/env: ‘python3.8’: No such file or directory
```

What the heck? Why is this even looking for python3.8?
This is because of the "pre-commit hook", mentioned above.
As above, you can proceed by just appending `--no-verify`:

```bash
> git commit -a -m 'fixed all the things' --no-verify
[main 99a01da] fixed all the things
 1 file changed, 10 insertions(+)
```

We'll help you sort it out in the PR.
But, you should fix the problem at some point. In this case,
uninstalling and reinstalling the pre-commit hooks fixed the problem:

```bash
> pre-commit uninstall
pre-commit uninstalled
Restored previous hooks to .git/hooks/pre-commit
> pre-commit install -f
pre-commit installed at .git/hooks/pre-commit
> # do some more edits
> git commit -a -m 'wrote the docs'
[main 79b81ff] fixed all the things
 1 file changed, 42 insertions(+)
```


## Benchmarking

Tskit has a simple benchmarking tool to help keep track of performance.

### Running benchmarks

The benchmark suite can be run with:

```bash
> cd python/benchmark
> python run.py
```

A subset of benchmarks can be run by specifying a string. For example, the following command runs all the benchmarks whose names contain "genotype", e.g. "genotype_matrix".

```bash
> python run.py -k genotype
```

If desired, the results of the benchmarks can be printed to STDOUT.

```bash
> python run.py -k genotype -p
```

Results are written to `bench-results.json` in the same folder. Note that if any version of `tskit`
is installed then that will be used for the benchmarking. To use the local development version of
tskit ensure you have `pip uninstall tskit` before running the benchmarking. The version used is
shown in the header of the report.

### Adding a new benchmark

The benchmarks are specified by the `config.yaml` file in `python/benchmark`. To add a new benchmark 
add an entry to the `benchmarks` dictionary. For example:

```yaml
  - code: do_my_thing({option_name})
    setup: |
      import a_module
    name: my_benchmark #optional, the code is used by default
    parameters:
      option_name:
        - "reticulate_splines"
        - "foobar"
```

Strings are interpreted as Python f-strings, so you can use the `parameters` dictionary to provide
values that will be interpolated into both the `setup` and `code` strings.

The suite can be run for all released versions with the `run-for-all-releases.py` script. 

## Releasing a new version

Tskit maintains separate versioning for the C API and Python package, each has its own
release process.


### C API

To release the C API, the ``TSK_VERSION_*`` macros should be updated,
along with ``VERSION.txt`` and the changelog  updated with the release
date and version. The changelog should also be checked for
completeness. Comparing  ``git log --follow --oneline -- c`` with
`git log --follow --oneline -- c/CHANGELOG.rst` may help here.
After the commit including these changes has been merged, tag a
release on GitHub using the pattern `C_MAJOR.MINOR.PATCH`, with:

```bash
git tag -a C_MAJOR.MINOR.PATCH -m "C API version C_MAJOR.MINOR.PATCH"
git push upstream --tags
```

After a couple of minutes a github action will make a draft release with the changelog
at the [releases page](https://github.com/tskit-dev/tskit/releases). Check it looks
right and publish the release (Click on the little pencil).
After release, start a section in the changelog for new developments and close the
GitHub issue milestone of the release.


### Python

It is worth running the benchmarks as above before release to check for any unexpected
major regressions. To make a release first prepare a pull request that sets the correct
version number in `tskit/_version.py`  following PEP440 format. For a normal release
this should be MAJOR.MINOR.PATCH, for a beta release use MAJOR.MINOR.PATCHbX
e.g. 1.0.0b1. Update the Python CHANGELOG.rst, ensuring that all significant
changes since the last release have been listed. Comparing
`git log --follow --oneline -- python`
with `git log --follow --oneline -- python/CHANGELOG.rst` may help here.
Once this PR is merged, push a tag to github:

```bash
git tag -a MAJOR.MINOR.PATCH -m "Python version MAJOR.MINOR.PATCH"
git push upstream --tags
```

This will trigger a build of the distribution artifacts for Python
on [Github Actions](https://github.com/tskit-dev/tskit/actions). and deploy
them to the [test PyPI](https://test.pypi.org/project/tskit/). Check
the release looks good there, then publish the draft release on the
[releases page](https://github.com/tskit-dev/tskit/releases) (Click on the little pencil).
Publishing this release will cause the github
action to deploy to the [production PyPI](https://pypi.org/project/tskit/).
After release, start a section in the changelog for new developments, close the
GitHub issue milestone of the release and update the python version to a `.dev`.
For a major release the website (github repo tskit-dev/tskit-site) should then
be updated with a notebook of new features. The benchmarks should be run as above
and the `bench-results.html` updated on the website.


--- ../../tskit/docs/provenance.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

```{currentmodule} tskit
```

(sec_provenance)=

# Provenance

Every tree sequence has provenance information associated with it. The purpose of this
information is to improve [reproducibility](https://en.wikipedia.org/wiki/Reproducibility):
given the provenance associated with a given tree sequence, it should be possible to
reproduce it. Provenance is split into three sections: the primary **software** used to
produce a tree sequence; the **parameters** provided to this software; and the computational
**environment** where the software was run.

This documentation serves two distinct purposes:

1. For developers using `tskit` in their own applications, it provides normative documentation
   for how provenance information should be stored.
2. For end-users of `tskit`, it provides documentation to allows them to inspect and interpret
   the provenance information stored in `.trees` files.

Provenance information is encoded using [JSON](https://www.json.org/).
To standardise the provenance information produced by different software and improve
interoperability we define a formal specification using [JSON Schema](http://json-schema.org/).
The full schema is provided {ref}`below <sec_provenance_schema>`, which may be used to
automatically validate input. In the following we describe the intention of the various
sections in more detail.

This document defines specification version 1.0.0. Specification version numbers follow
[SemVer](https://semver.org/) semantics.

(sec_provenance_example)=

## Example

To make things more concrete, let's consider an example:

```json
{
  "schema_version": "1.0.0",
  "software": {
    "name": "msprime",
    "version": "0.6.1.dev123+ga252341.d20180820"
  },
  "parameters": {
    "sample_size": 5,
    "random_seed": 12345,
    "command": "simulate"
  },
  "environment": {
    "libraries": {
      "gsl": {
        "version": "2.1"
      },
      "kastore": {
        "version": "0.1.0"
      }
    },
    "python": {
      "version": "3.5.2",
      "implementation": "CPython"
    },
    "os": {
      "system": "Linux",
      "node": "powderfinger",
      "release": "4.15.0-29-generic",
      "version": "#31~16.04.1-Ubuntu SMP Wed Jul 18 08:54:04 UTC 2018",
      "machine": "x86_64"
    }
  },
  "resources": {
    "elapsed_time": 12.34,
    "user_time": 10.56,
    "sys_time": 1.78,
    "max_memory": 1048576
  }
}
```

This information records the provenance for a very simple msprime simulation. The record is a JSON
object with three mandatory fields ("software", "parameters" and "environment") and one optional
("resources") which we discuss separately in the following sections.

(sec_provenance_software)=


## Software


Every tree sequence is produced by some piece of software. For example, this may be a
coalescent simulation produced by `msprime`, a forwards-time simulation from `SLiM`
or tree sequence inferred from data by `tsinfer`. The software provenance is
intended to capture the details about this primary software.

```{list-table}
:header-rows: 1

* - Field
  - Type
  - Description
* - name
  - string
  - The name of the software.
* - version
  - string
  - The software version.
```

Note that libraries that the primary software links against are considered part of the
{ref}`sec_provenance_environment` and should be recorded there.

(sec_provenance_parameters)=

## Parameters

The parameters section of a provenance document records the input that was used to
produce a particular tree sequence. There are no requirements on what may be stored
within it, but we make some recommendations here on how to encode such information.

As a general principle, sufficient information should be recorded in the parameters
section to allow the output tree sequence to be reproduced exactly. There will be instances,
however, where this is not possible due to missing files, issues with numerical precision
and so on.

### API invocations

Consider an API call like the following simple msprime simulation:

```python
ts = msprime.simulate(sample_size=10, recombination_rate=2)
```

We recommend encoding the parameters provenance as follows (other fields omitted
for clarity):

```json
{
  "parameters": {
    "command": "simulate",
    "sample_size": 10,
    "recombination_rate": 2,
    "random_seed": 123456789,
  }
}
```

Specifically, we encode the name of the function using the `command` key and
the function parameters in the obvious way. Note that we include the `random_seed`
here even though it was automatically generated.


### CLI invocations

Consider the following invocation of a hypothetical command line program:

```bash
$ supersim --sample-size=10 --do-some-stuff -O out.trees
```

We recommend encoding the parameters provenance as follows (other fields omitted
for clarity):

```json
{
  "parameters": {
    "command": "supersim",
    "args": ["--sample-size=10", "--do-some-stuff", "-O", "out.trees"],
    "random_seed": 56789
  }
}
```

Here we encode the name of the program using the `command` key
and its command line arguments as a list of strings in the `args` key. We
also include the automatically generated random seed in the parameters list.

If parameters that affect the output tree sequence are derived from environment
variables these should also be recorded.

(sec_provenance_environment)=

## Environment

The environment section captures details about the computational environment in
which the software was executed. Two optional fields are defined: `os`
and `libraries`. We recommend including any additional relevant platform
information here; for example, if using Python store the interpreter information
as shown in the example above.

### Operating system

The `os` section records details about the operating system on which the
software was executed. This section is optional and has no required internal
structure. We recommend the following structure based on the output of the
POSIX [uname](http://pubs.opengroup.org/onlinepubs/009695399/functions/uname.html)
function:

```json
{
  "environment": {
    "os": {
      "system": "Linux",
      "node": "powderfinger",
      "release": "4.15.0-29-generic",
      "version": "#31~16.04.1-Ubuntu SMP Wed Jul 18 08:54:04 UTC 2018",
      "machine": "x86_64"
    }
}
```

### Libraries

The `libraries` section captures information about important libraries that the
primary software links against. There is no required structure.


## Resources

The resources section captures details about the computational resources used during the execution of the software. This section is optional and has the following fields, each of which is optional and may not be filled depending on os support:


- `elapsed_time`: The total elapsed time in seconds.
- `user_time`: The total user CPU time in seconds.
- `sys_time`: The total system CPU time in seconds.
- `max_memory`: The maximum memory usage in bytes.

Including this information makes it easy for users of tree-sequence producing software to
account for resource usage across pipelines of tools.

(sec_provenance_schema)=

## Full schema

This schema is formally defined using [JSON Schema](http://json-schema.org/) and
given in full here. Developers writing provenance information to `.trees` files
should validate the output JSON against this schema.

```{eval-rst}
.. literalinclude:: ../python/tskit/provenance.schema.json
    :language: json
```

--- ../../tskit/docs/file-formats.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

:::{currentmodule} tskit
:::

(sec_file_formats)=

# File formats


(sec_tree_sequence_file_format)=

## Tree sequence file format

To make tree sequence data as efficient and easy as possible to use, we store the
data on file in a columnar, binary format. The format is based on the
[kastore](https://pypi.org/project/kastore/) package, which is a simple
key-value store for numerical data. There is a one-to-one correspondence
between the tables described above and the arrays stored in these files.

By convention, these files are given the `.trees` suffix (although this
is not enforced in any way), and we will sometimes refer to them as ".trees"
files. We also refer to them as "tree sequence files".

:::{todo}
Link to the documentation for kastore, and describe the arrays that are
stored as well as the top-level metadata.
:::


### Legacy Versions

Tree sequence files written by older versions of tskit are not readable by
newer versions of tskit. For major releases of tskit, `tskit upgrade`
will convert older tree sequence files to the latest version.


(sec_text_file_format)=

## Text file formats

The tree sequence text file format is based on a simple whitespace
delimited approach. Each table corresponds to a single file, and is
composed of a number of whitespace delimited columns. The first
line of each file must be a **header** giving the names of each column.
Subsequent rows must contain data for each of these columns, following
the usual conventions. Each table has a set of mandatory and optional columns which are
described below. The columns can be provided in any order, and extra columns
can be included in the file. Note, in particular, that this means that
an `id` column may be present in any of these files, but it will be
ignored (IDs are always determined by the position of the row in a table).

The {meth}`load_text` method can be used to read tables in text format. This has been
used to create the following very simple tree sequence, with four nodes, two trees,
and three mutations at two sites, both on the first tree:


```{code-cell} ipython3
:tags: ["hide-input"]
# TODO once https://github.com/tskit-dev/tskit/issues/1824 is solved
# change the individual table to include some with blank parents / locations
import io

import tskit
from IPython.display import SVG

individuals = """\
flags       location     parents
0           0.5,1.2      -1,-1
0           1.0,3.4      0,-1
0           3.5,6.3      0,1
0           0.5          -1,-1
0           0.5,0.5      2,3
"""

nodes = """\
is_sample   individual   time
1           0            0.0
1           0            0.0
0           -1           2.0
0           -1           3.0
"""
edges = """\
left   right   parent  child
0.0    7.0     2       0
0.0    7.0     2       1
7.0    10.0    3       0
7.0    10.0    3       1
"""

sites = """\
position      ancestral_state
2.0           AT
4.0           A
"""

mutations = """\
site   node    derived_state    time    parent
0      0       A                0.5     -1
1      0       T                1.5     -1
1      1       A                1.0     1
"""

migrations = """\
left   right   node   source   dest   time
0.0    0.7     5      2        3      1.0
0.8    0.9     8      3        4      3.0
"""

populations = """\
id   metadata
0    cG9wMQ==
1    cG9wMg==
"""

ts = tskit.load_text(
    individuals=io.StringIO(individuals),
    nodes=io.StringIO(nodes),
    edges=io.StringIO(edges),
    sites=io.StringIO(sites),
    mutations=io.StringIO(mutations),
    # migrations=io.StringIO(migrations),  # uncomment when https://github.com/tskit-dev/tskit/issues/19 fixed
    populations=io.StringIO(populations),
    strict=False
)
SVG(ts.draw_svg(y_axis=True))

```

A deletion from AT to A has occurred at position 2 on the branch leading to
node 0, and two mutations have occurred at position 4 on the branch leading to
node 1, first from A to T, then a back mutation to A. The genotypes of our two
samples, nodes 0 and 1, are therefore AA and ATA. Note that this tree sequence
also contains entries in the individual, population,
and migration tables, but this is not shown plot above.


(sec_individual_text_format)=

### Individual text format

The individual text format must contain a `flags` column.
Optionally, there may also be `location`, `parents` and
`metadata` columns. See the
{ref}`individual table definitions<sec_individual_table_definition>`
for details on these columns.

Note that there are currently no globally defined `flags`, but the column
is still required; a value of `0` means that there are no flags set.

The `location` and `parents` columns should be a sequence of comma-separated numeric
values. They do not all have to be the same length.

```{code-cell} python
:tags: ["hide-input", "output-wide-tabs"]
import sys
from IPython.display import display, HTML

display(HTML("An example individual table:"))
ts.dump_text(individuals=sys.stdout)
```

(sec_node_text_format)=

### Node text format

The node text format must contain the columns `is_sample` and
`time`. Optionally, there may also be `population`, `individual`, and
`metadata` columns. See the
{ref}`node table definitions<sec_node_table_definition>` for details on these columns.

Note that we do not have a `flags` column in the text file format, but
instead use `is_sample` (which may be 0 or 1). Currently, `NODE_IS_SAMPLE` is
the only flag value defined for nodes, and as more flags are defined we will
allow for extra columns in the text format.

```{code-cell} ipython3
:tags: ["hide-input", "output-wide-tabs"]
display(HTML("An example node table:"))
ts.dump_text(nodes=sys.stdout)
```


(sec_edge_text_format)=

### Edge text format

The edge text format must contain the columns `left`,
`right`, `parent` and `child`. Optionally, there may also be
a `metadata` column.
See the {ref}`edge table definitions <sec_edge_table_definition>`
for details on these columns.

```{code-cell} ipython3
:tags: ["hide-input", "output-wide-tabs"]
display(HTML("An example edge table:"))
ts.dump_text(edges=sys.stdout)
```

(sec_site_text_format)=

### Site text format

The site text format must contain the columns `position` and
`ancestral_state`. The `metadata` column may also be optionally
present. See the
{ref}`site table definitions <sec_site_table_definition>`
for details on these columns.

```{code-cell} ipython3
:tags: ["hide-input", "output-wide-tabs"]
display(HTML("An example site table:"))
ts.dump_text(sites=sys.stdout)
```


(sec_mutation_text_format)=

### Mutation text format

The mutation text format must contain the columns `site`,
`node` and `derived_state`. The `time`, `parent` and `metadata` columns
may also be optionally present (but `parent` must be specified if
more than one mutation occurs at the same site). If `time` is absent
`UNKNOWN_TIME` will be used to fill the column. See the
{ref}`mutation table definitions <sec_mutation_table_definition>`
for details on these columns.

```{code-cell} ipython3
:tags: ["hide-input", "output-wide-tabs"]
display(HTML("An example mutation table:"))
ts.dump_text(mutations=sys.stdout)
```


(sec_migration_text_format)=

### Migration text format

The migration text format must contain the columns `left`,
`right`, `node`, `source`, `dest` and `time`. The `metadata` column
may also be optionally present. See the
{ref}`migration table definitions <sec_migration_table_definition>`
for details on these columns.

```{code-cell} ipython3
:tags: ["hide-input", "output-wide-tabs"]
display(HTML("An example migration table:"))
print(migrations)  # fixme
# ts.dump_text(migrations=sys.stdout)
```


(sec_population_text_format)=

### Population text format

Population tables only have a `metadata` column, so the text format for
a population table requires there to be a `metadata` column. See the
{ref}`population table definitions <sec_population_table_definition>` for
details.

```{code-cell} ipython3
:tags: ["hide-input", "output-wide-tabs"]
display(HTML("An example population table:"))
ts.dump_text(populations=sys.stdout)
```

The `metadata` contains base64-encoded data (in this case, the strings
`pop1` and `pop1`).



--- ../../tskit/docs/phylogen.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.13
    jupytext_version: 1.10.3
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

```{currentmodule} tskit
```

(sec_phylogen)=

# `Tskit` for phylogenetics

`Tskit`, the tree sequence toolkit, can be used as an efficient library for very large evolutionary trees. `Tskit` makes it easy to deal with trees with millions of
tips, as in the example below:

```{code-cell}
:"tags": ["hide-input"]
import tskit
```

```{code-cell}
%%time

num_tips = 1_000_000
big_tree = tskit.Tree.generate_comb(num_tips);
print("Tree sequence takes up", big_tree.tree_sequence.nbytes / 1024**2, "Mb")
print(f"Generating a 'comb' (pectinate) tree of {num_tips} tips took:")
```

:::{todo}
Display the million tip tree in condensed form, when
https://github.com/tskit-dev/tskit/issues/2372#issuecomment-1298518380 and
https://github.com/tskit-dev/tskit/issues/2628 are solved
:::

Calculations on these huge trees can be very efficient:

```{code-cell}
%%time

traversed_nodes = 0
for u in big_tree.nodes(order="postorder"):
    traversed_nodes += 1
print(f"Postorder traversal through {traversed_nodes} nodes took:")
```

```{code-cell}
%%time

b1_index = big_tree.b1_index()
print(f"B1 balance index is {b1_index}. Calculation took:")
```

We can also read trees efficiently, e.g. in Newick format:

:::{todo}
Add example of fast reading of a newick file, once
https://github.com/tskit-dev/tskit/issues/2187 is solved. E.g. we could say
> For example, we can read in a newick format tree of XXX tips in XXX secs. In practice,
> this means when reading in a large tree you are mainly limited by disk access speeds
:::

```{code-cell}
import tsconvert  # used for reading tree sequences from different formats

# example code reading in a large file, timed

# Or read smaller trees from strings (here we create a tree spanning 1000 genomic units)
ts = tsconvert.from_newick("(A:6,((B:1,C:1):2,(D:2,E:2):1):3);", span=1000)
```

The "succinct tree sequence" format used by `tskit` can also store mutations
(and optionally a reference genome) along with the tree(s). This results in a
single unified representation of large genomic datasets, storing trees,
sequence data and metadata in a single efficient structure. Examples are given
in the section below entitled {ref}`sec_phylogen_unified_structure`.

As the name suggests, a tree sequence can also store and analyse a sequence of
trees along a genome (i.e. a "phylogenetic network"). This is necessary to
account for recombination between lineages, and may be important even when looking at
species-level phylogenies due to the effects of hybridization and incomplete lineage
sorting. An overview, and links to further details are given at the
{ref}`end of this page <sec_phylogen_multiple_trees>`.

## Hints for phylogeneticists

Unlike other phylogenetic libraries, `tskit` is designed to efficiently store not just 
single trees, but sequences of correlated trees along a genome. This means that the
library has some features not found in more standard phylogenetic libraries.
Here we focus on the {ref}`sec_python_api`,
introducing seven `tskit` concepts that may be useful to those with a background in
phylogenetics (each is linked to a separate section below):

1. An evolutionary tree is always contained within a "tree sequence".
   See {ref}`sec_phylogen_tree_in_sequence`
2. The basic elements of a tree are *nodes* and *edges*, referred to by integer IDs.
   See {ref}`sec_phylogen_ids`
3. The `Tree` object in the Python API provides useful phylogenetic methods.
   See {ref}`sec_phylogen_tree_object`
4. Specific nodes in a tree (often the tips) are marked as "samples", meaning they are
   known from data. See {ref}`sec_phylogen_samples`
5. Nodes and edges have additional attributes, with arbitrary information stored in
   *metadata*. See {ref}`sec_phylogen_metadata`
6. All nodes *must* have a valid time, meaning trees are always directional
   (i.e. "rooted"). See {ref}`sec_phylogen_node_time`
7. "Roots" in trees have a specific definition, and a single tree can consist of
   topologically independent clades (a "multiroot" tree).
   See {ref}`sec_phylogen_multiroot`

(sec_phylogen_tree_in_sequence)=
### Trees are always part of a tree sequence

In tskit, all trees are stored as a "tree sequence" of correlated trees. This allows
easy extension of the library to multiple trees produced e.g. by hybridization.
In the simplest case, however, the tree sequence can contain just a single tree. This
can be obtained using the {meth}`~TreeSequence.first()` method.

```{code-cell}
:"tags": ["hide-input"]
# Make sensible node labels. It would be nice to roll this into tree.draw()
node_labels = {
    node_object.id: node_object.metadata["name"]
    for node_object in ts.nodes()
    if "name" in node_object.metadata
}
```

```{code-cell}
tree = ts.first()
tree.draw(node_labels=node_labels)  # or use draw_svg() for more options
```

Often you will have the tree sequence stored in a variable, such as the `ts` variable
used above. However, you can also obtain the tree sequence in which a tree is
contained using the {attr}`Tree.tree_sequence` attribute:

```{code-cell}
tree.tree_sequence  # When output in a notebook, prints a summary of the tree sequence
```

(sec_phylogen_ids)=
### Integer node and edge IDs

The plot above labels nodes by their name, but internally the `tskit` library relies
heavily on integer IDs. Here's the same tree with node IDs plotted instead:

```{code-cell}
tree = ts.first()
tree.draw_svg()
```

#### Nodes

Each {ref}`node<sec_terminology_nodes>` in a tree sequence is allocated an
integer ID starting from 0 to `ts.num_nodes - 1` (IDs can be allocated in any order;
often the tips are labelled starting from 0 but this is not necessarily so, and
is not the case in the example above).

For efficiency reasons, tree traversal routines, as well as many other `tskit` methods,
tend to return integer IDs. You can use this ID to get specific information about the
node and its position in the tree, for example

```{code-cell}
node_id = 4
parent_id = tree.parent(node_id)
child_ids = tree.children(node_id)
print("The parent of", node_id, "is", parent_id, "and its children are", child_ids)
# or e.g. get all parents them as an array (where -1 means there is no parent)
print(f"The parents of nodes 0..{ts.num_nodes-1} are", tree.parent_array)
```

Other methods also exist to
{ref}`examine nodes in a tree<sec_python_api_trees_node_measures>`, e.g.
{meth}`Tree.is_leaf`, {meth}`Tree.mrca` for the most recent common ancestor between
2 or more nodes, etc.

#### Edges

Rather than refer to "branches" of a tree, tskit tends to refer to
{ref}`sec_terminology_edges` (the term "edge" emphasises that these can span
{ref}`sec_phylogen_multiple_trees`, although for tree sequences containing a single
tree, the terms are interchangeable). Like other entities in `tskit`, edges are referred
to by an integer ID. For instance, here is the edge above the internal node 4

```{code-cell}
node_id = 4
edge_id = tree.edge(node_id)
print("The edge above", node_id, "has ID", edge_id)
print(tree.tree_sequence.edge(edge_id))
```

The `left` and `right` attributes of an edge give genomic coordinates, and are
important in tree sequences that contain more than one tree.


(sec_phylogen_tree_object)=
### The `Tree` object

The {class}`Tree` object has {ref}`methods<sec_python_api_trees>` to perform basic operations
on a tree such as traversing the nodes, identifying parents, children, and common
ancestors, etc. {ref}`Several methods<sec_python_api_trees_node_measures_array>`
also return numpy arrays for use in
{ref}`efficient algorithms using numba<sec_trees_numba>`

```{code-cell}
for n_id in tree.nodes(order="postorder"):
    # you can also use "preorder", "levelorder", "timeasc", etc.
    print(n_id)
# Or get all of them as arrays
print("Node IDs in postorder:", tree.postorder())
```

Various phylogenetic statistics are also available on trees, e.g

```{code-cell}
print(f"The colless imbalance index is {tree.colless_index()}")
```

See {ref}`sec_phylogen_methods` for more examples.

(sec_phylogen_samples)=
### Sample nodes

Often we are only have detailed information about specific nodes that we have sampled,
such as genomes A, B, C, D, and E in the example above. These are designated as
*sample nodes*, and are plotted as square nodes. The concept of
{ref}`sample nodes<sec_data_model_definitions_sample>` is integral
to the `tskit` format. They can be identified by using the
{meth}`Node.is_sample` and {meth}`Tree.is_sample` methods, or can be listed using
{meth}`TreeSequence.samples` or {meth}`Tree.samples()` (internally, the `node.flags`
field is used to {ref}`flag up<sec_node_table_definition>` which nodes are samples):

```{code-cell}
for n_id in tree.nodes():
    n_is_sample = tree.is_sample(n_id)
    print(f"Node {n_id} {'is' if n_is_sample else 'is not'} a sample node")

print("Sample nodes are", tree.tree_sequence.samples())
```

Often the sample nodes are the leaves of a tree, but this need not be the case. There
are fast methods for identifying the sample nodes under an internal node in the tree,
etc.


(sec_phylogen_metadata)=
### Attributes and metadata

Given a node ID, you can access more information about the node from a
{class}`node object<tskit.Node>`.
Because nodes are shared across all trees in a tree sequence, you access the
node object via the `tree_sequence` to which this tree belongs:

```{code-cell}
tree.tree_sequence.node(node_id)  # or simply ts.node(node_id)
```

Attributes such as `id`, `flags` and `time` are always present. Arbitrary information,
such a name or e.g. bootstrap values, are stored in *metadata*

```{code-cell}
for n_id in tree.nodes():
    print("Node", n_id, tree.tree_sequence.node(n_id).metadata.get("name", "<no name>"))
```

However, for large datasets, it may be more efficient to access the array of e.g.
times for all nodes, which provides direct memory access into the
{ref}`tables<sec_tables>` that underlie the tree sequence format:

```{code-cell}
tree.tree_sequence.nodes_time
```


(sec_phylogen_node_time)=
### Nodes must have times

Perhaps the most noticable different between a `tskit` tree and the encoding of trees
in other phylogenetic libraries is that `tskit` does not explicitly store branch lengths.
Instead, each node has a *time* associated with it. Branch lengths can therefore be
found by calculating the difference between the time of a node and the time of its
parent node.

Since nodes *must* have a time, `tskit` trees aways have these (implicit) branch
lengths. To represent a tree ("cladogram") in which the branch lengths are not
meaningful, the {attr}`TreeSequence.time_units` of a tree sequence can be
specified as `"uncalibrated"` (see below)

Another implication of storing node times rather than branch lengths is that `tskit`
trees are always directional (i.e. they are "rooted"). The reason that `tskit` stores
times of nodes (rather than e.g. genetic distances between them) is to ensure temporal 
consistency. In particular it makes it impossible for a node to be an ancestor of a
node in one tree, and a descendant of the same node in another tree in the tree sequence.
This is of critical importance when extending the concept of genetic ancestry to
{ref}`sec_phylogen_multiple_trees` along a genome.

The {attr}`TreeSequence.time_units` attribute stores the units in which time is
measured: if not known, this defaults to "unknown":

```{code-cell}
print("Time units are", tree.tree_sequence.time_units)
tree.draw_svg(y_axis=True)
```

Although branch lengths are not stored explicitly, for convenience `tskit` provides a
{meth}`Tree.branch_length` method:

```{code-cell}
print(
    "The edge (i.e. branch) immediately above node",
    node_id,
    "has a 'length' of",
    tree.branch_length(node_id),
)
```
:::{todo}
The branch distance between two samples is also easy to calculate

NB: Turn the following in to a code cell
```
target_node_1 = 5
target_node_2 = 7
print(
    "The branch distance between nodes",
    target_node_1,
    "and",
    target_node_2,
    "is",
    # See https://github.com/tskit-dev/tskit/issues/2627 - what should be call this
    # so as not to get mixed up with tree.path_length which counts the number of edges
    # tree.branch_distance(target_node_1, target_node_2),
)
```

It is worth noting that this distance is the basis for the "genetic divergence"
between two samples in a tree. For this reason, an equivalent way to carry out the
calculation is to use {meth}`TreeSequence.divergence`, part of the the standard `tskit`
{ref}`sec_stats` framework, setting `mode="branch"` and
`windows="trees"`. This is a more flexible approach, as it allows the distance between
multiple sets of samples in {ref}`sec_phylogen_multiple_trees` to be calculated
efficiently:

NB: Turn the following in to a code cell
```
target_node_1 = 5
target_node_2 = 7
print(
    "Branch distance using built-in stats framework:"
    tree.tree_sequence.divergence(([5], [7]), mode="branch", windows="trees")
)
```

:::


(sec_phylogen_multiroot)=
### Roots and multiroot trees

In `tskit`, {ref}`sec_data_model_tree_roots` of trees are defined with respect to the
sample nodes. In particular, if we move back in time along the tree branches from a
sample, the oldest node that we encounter is defined as a root. The ID of a root can be
obtained using {attr}`Tree.root`:

```{code-cell}
print("The root node of the following tree has ID", tree.root)
tree.draw_svg()
```

But in `tskit`, we can also create a single "tree" consisting of multiple unlinked
clades. In our example, we can create one of these phylogenetically unusual objects
if we remove the edge above node 4, by
{ref}`editing the underlying tables<sec_tables_editing>`:

```{code-cell}
# Trees & tree seqs are immutable: to change them, modify a copy of the underlying tables
tables = ts.dump_tables()
keep_edge = tables.edges.child != 4
tables.edges.replace_with(tables.edges[keep_edge])
new_ts = tables.tree_sequence()  # Turn the edited tables back into a tree sequence
new_tree = new_ts.first()
new_tree.draw_svg()
```

Although there are two separate topologies in this plot, in `tskit` terminology, it is
considered a single tree, but with two roots:

```{code-cell}
print("The first tree has", len(new_tree.roots), "roots:", new_tree.roots)
```

This also means that if we have no topology at all (i.e. an "empty tree"), each
sample is its own root.

```{code-cell}
tables.edges.clear()
erased_ts = tables.tree_sequence()
empty_tree = erased_ts.first()
print("This empty tree has", len(empty_tree.roots), "roots:", empty_tree.roots)
empty_tree.draw_svg()
```

The samples here are {ref}`sec_data_model_tree_isolated_nodes`. This may seem like a
strange corner case, but in `tskit`, isolated sample nodes are used to represent
{ref}`sec_data_model_missing_data`. This therefore represents a tree in which
relationships between the samples are not known. This could apply, for instance,
in regions of the genome where no genetic data exists, or where genetic ancestry
has not been simulated.

(sec_phylogen_methods)=
## Phylogenetic methods

:::{todo}
Demo some phylogenetic methods. e.g.
1. Total branch length - demo quick calculation across multiple trees - incremental algorithm used extensively in population genetics. ("bringing tree thinking to popgen").
2. KC distance
3. Balance metrics
4. Topology rankings (see https://github.com/tskit-dev/tutorials/issues/93)
:::


(sec_phylogen_unified_structure)=
## Storing and accessing genetic data

`Tskit` has been designed to capture both evolutionary tree topologies and the genetic
sequences that evolve along the branches of these trees. This is achieved by defining
{ref}`sec_terminology_mutations_and_sites` which are associated with specific positions
along the genome.

```{code-cell}
import msprime  # The `msprime` package can throw mutations onto a tree sequence
mutated_ts = msprime.sim_mutations(ts, rate=3e-3, random_seed=321)
mutated_tree = mutated_ts.first()
print("Variable sites with the following IDs generated")
for site in mutated_tree.sites():
    print(
        f"Site ID {site.id} @ genomic position {site.position:g}:",
        f"{site.ancestral_state} -> {site.mutations[0].derived_state}"
    )
mutated_tree.draw_svg()
```

Mutations occur above nodes in a tree, with all the descendant
nodes inheriting that specific mutation (unless replaced by a subsequent
mutation at the same site). This allows genetic variation to be
{ref}`efficiently represented<sec_what_is_dna_data>` using the tree topology.
To obtain the genetic variation at each site across the entire genome, you can use the
{meth}`TreeSequence.sites` method, or (less efficiently), you can use
{meth}`TreeSequence.alignments` to output the
entire sequences for each sample node:

```{code-cell}
for node_id, alignment in zip(
    mutated_ts.samples(),
    mutated_ts.alignments(missing_data_character="."),
):
    print(f"Node {node_id}: {alignment}")
```


(sec_phylogen_multiple_trees)=
## Multiple trees

Where `tskit` really shines is when the ancestry of your dataset cannot be adequately
represented by a single tree. This is a pervasive issue in genomes (even from different
species) that have undergone recombination in the past. The resulting series of
{ref}`local trees<sec_what_is_local_trees>` along a genome are highly correlated
(see {ref}`sec_concepts`).

Instead of storing each tree along a genome separately, `tskit` records the genomic
coordinates of each edge, which leads to enormous efficiencies in storage and
analysis. As a basic demonstration, we can repeat the edge removal example
{ref}`above <sec_phylogen_multiroot>`, but only remove the ancestral link above node 4
for the first half of the genome.

```{code-cell}
tables = ts.dump_tables()
edge_id_above_node_4 = ts.first().edge(4)
left_coord_for_edges = tables.edges.left
left_coord_for_edges[edge_id_above_node_4] = 50
tables.edges.left = left_coord_for_edges  # reset the right coords
tables.sort()
multi_ts = tables.tree_sequence()

multi_ts.draw_svg()
```

For the left hand side of the genome we lack information about the ancestry of
node 4, but for the right hand side we know this information. The result is to
generate 2 trees in the tree sequence, which differ only in the presence of absence of
a single branch. We do not have to separately store the entire tree on the right: all
the edges that are shared between trees are stored only once.

The rest of the `tskit` tutorials will lead you through the concepts involved with
storing and analysing sequences of many correlated trees. For a simple introduction, you
might want to start with {ref}`sec_what_is`.




--- ../../tskit/docs/popgen.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.13
    jupytext_version: 1.10.3
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

```{currentmodule} tskit
```

```{code-cell}
:tags: [remove-cell]
import urllib.request

import tqdm
import tskit
import tszip

class DownloadProgressBar(tqdm.tqdm):
    def update_to(self, b=1, bsize=1, tsize=None):
        if tsize is not None:
            self.total = tsize
        self.update(b * bsize - self.n)

def download(url, progress=True):
    with DownloadProgressBar(
        unit='B',
        unit_scale=True,
        miniters=1,
        desc=url.split('/')[-1],
        disable=not progress,
    ) as t:
        tmp_fn, _ = urllib.request.urlretrieve(url, reporthook=t.update_to)
        try:
            ts = tskit.load(tmp_fn)
        except tskit.FileFormatError:
            # could be a tsz file
            ts = tszip.decompress(tmp_fn)
        urllib.request.urlcleanup() # Remove tmp_fn
    return ts

def download_unified_genealogy():
    keep_span = [108_000_000, 110_000_000]  # cut down to this genome region
    keep_regions = {"EastAsia", "EAST_ASIA", "AFRICA", "Africa"}

    # Downloads 138 Mb of data - this may take a while
    tables = download(
        "https://zenodo.org/record/5512994/files/"
        "hgdp_tgp_sgdp_high_cov_ancients_chr2_q.dated.trees.tsz"
    ).dump_tables()
    tables.keep_intervals([keep_span])
    tables.populations.metadata_schema = tskit.MetadataSchema.permissive_json()
    tables.sites.metadata_schema = tskit.MetadataSchema.permissive_json()
    ts = tables.tree_sequence()
    ts = ts.simplify([
        u
        for u in ts.samples()
        if (
            ts.population(ts.node(u).population).metadata.get("region") in keep_regions
            or ts.population(ts.node(u).population).metadata.get("name") == "Denisovan"
        )
    ])
    tszip.compress(ts, "data/unified_genealogy_2q_108Mb-110Mb.tsz")

def create_notebook_data():
    download_unified_genealogy()

# create_notebook_data()  # uncomment to recreate the tree seqs used in this notebook
```

(sec_intro_popgen)=

# `Tskit` for population genetics

{ref}`Tskit<tskit:sec_introduction>`, the tree sequence toolkit, brings the power of
evolutionary trees to the field of population genetics. The
{ref}`succinct tree sequence<sec_what_is>` format
is designed to store DNA sequences jointly with their ancestral history (the
"genetic genealogy" or {ref}`ARG<sec_args>`). Storing population genetic data in this
form enables highly efficient computation and analysis.

The core `tskit` library provides methods for storing genetic data, a flexible
analysis framework, and APIs to build your own efficient population genetic algorithms.
Because of its speed and scalability, `tskit` is well-suited to interactive analysis of
large genomic datasets. 

## Population genetic simulation

Several simulation tools output tree sequences. Below we use the
standard library for population genetic simulation models
([stdpopsim](https://popsim-consortium.github.io/)) to generate a model of
*Homo sapiens*, in which African, Eurasian,
and Asian populations combine to generate a mixed American population. We can use the
[demesdraw](https://pypi.org/project/demesdraw/) package to plot a schematic of the
migrations and population size changes that define this model. 


```{code-cell}
import stdpopsim
import demesdraw
from matplotlib import pyplot as plt

species = stdpopsim.get_species("HomSap")
model = species.get_demographic_model("AmericanAdmixture_4B11")

# Plot a schematic of the model
demesdraw.tubes(model.model.to_demes(), ax=plt.gca(), seed=1, log_time=True)
plt.show()
```

Genomic data in tree sequence format can be generated via the widely-used
[msprime](https://tskit.dev/software/msprime.html) simulator. Here we simulate 20
kilobases of genome sequence at the start of human chromosome 1 under this model,
together with its evolutionary history. We generate 16 diploid genomes: 4 from each of
the populations in the model. The DNA sequences and their ancestry are stored in a
succinct tree sequence named `ts`:

```{code-cell}
contig = species.get_contig("chr1", mutation_rate=model.mutation_rate, right=20_000)
samples = {"AFR": 4, "EUR": 4, "ASIA": 4, "ADMIX": 4} # 16 diploid samples
engine = stdpopsim.get_engine("msprime")
ts = engine.simulate(model, contig, samples, seed=9)
print(f"Simulated a tree sequence of {ts.num_samples} haploid genomes:")
print(f"{ts.num_sites} variable sites over {ts.sequence_length} base pairs")
```

We can now inspect alleles and their frequencies at the variable sites we have simulated
along the genome:

```{code-cell}
for v in ts.variants():
    display(v)
    if v.site.id >= 2: #  Only show site 0, 1, and 2, for brevity
        break
```

Or we can display the {meth}`~TreeSequence.haplotypes` (i.e. the variable sites) for
each sample

```{code-cell}
samples = ts.samples()
for sample_id, h in zip(samples, ts.haplotypes(samples=samples)):
    pop = ts.node(sample_id).population
    print(f"Sample {sample_id:<2} ({ts.population(pop).metadata['name']:^5}): {h}")
```

From the tree sequence it is easy to obtain the
{meth}`TreeSequence.allele_frequency_spectrum` for the entire region (or for
{ref}`windowed regions<sec_tskit_getting_started_compute_statistics_windowing>`)

```{code-cell}
afs = ts.allele_frequency_spectrum()
plt.bar(range(ts.num_samples + 1), afs)
plt.title("Allele frequency spectrum")
plt.show()
```

Similarly `tskit` allows fast and easy
{ref}`calculation of statistics<sec_tutorial_stats>` along the genome. Here is
a plot of windowed $F_{st}$ between Africans and admixed Americans over this short
region of chromosome:

```{code-cell}
# Define the samples between which Fst will be calculated
pop_id = {p.metadata["name"]: p.id for p in ts.populations()}
sample_sets=[ts.samples(pop_id["AFR"]), ts.samples(pop_id["ADMIX"])]

# Do the windowed calculation, using windows of 2 kilobases
windows = list(range(0, int(ts.sequence_length + 1), 2_000))
F_st = ts.Fst(sample_sets, windows=windows)

# Plot
plt.stairs(F_st, windows, baseline=None)
plt.ylabel("AFR-ADMIX Fst")
plt.xlabel("Genome position")
plt.show()
```

Extracting the genetic tree at a specific genomic location is easy using `tskit`, which
also provides methods to {ref}`plot<sec_tskit_viz>` these trees. Here we
grab the tree at position 10kb, and colour the different populations by
different colours, as described in the {ref}`viz tutorial<sec_tskit_viz_styling>`:

```{code-cell}
tree = ts.at(10_000)

colours = dict(AFR="yellow", EUR="cyan", ASIA="green", ADMIX="red")
styles = [
    f".leaf.p{pop.id} > .sym {{fill: {colours[pop.metadata['name']]}}}"
    for pop in ts.populations()
]

styles += [ # rotate the population labels, etc
    ".leaf > .lab {text-anchor: start; transform: rotate(90deg) translate(6px)}",
    ".leaf > .sym {stroke: black}"
]

labels = { # Label samples by population
    u: ts.population(ts.node(u).population).metadata["name"].capitalize()
    for u in ts.samples()
}

tree.draw_svg(
    size=(800, 500),
    canvas_size=(800, 520),
    node_labels=labels,
    style="".join(styles),
    y_axis=True,
    y_ticks=range(0, 30_000, 10_000)
)

```

## Population genetic inference

If, instead of simulations, you want to analyse existing genomic data (for example
stored in a VCF file), you will need to infer a tree sequence from it, using e.g.
[tsinfer](https://tskit.dev/tsinfer/docs/stable/). Here we load an illustrative portion
of an [inferred tree sequence](https://zenodo.org/record/5512994)
based on about 7500 public human genomes, including genomes from the
[Thousand Genomes Project](https://www.internationalgenome.org/data-portal/data-collection/grch38) and
[Human Genome Diversity Project](https://www.internationalgenome.org/data-portal/data-collection/hgdp).
The genomic region encoded in this tree sequence has been cut down to
span positions 108Mb-110Mb of human chromosome 2, which spans the
[EDAR](https://en.wikipedia.org/wiki/Ectodysplasin_A_receptor) gene.

Note that tree sequence files are usually imported using {func}`load`,
but because this file has been additionally compressed, we load it via
{func}`tszip:tszip.decompress`:

```{code-cell}
import tszip
ts = tszip.decompress("data/unified_genealogy_2q_108Mb-110Mb.tsz")

# The ts encompasses a region on chr 2 with an interesting SNP (rs3827760) in the EDAR gene
edar_gene_bounds = [108_894_471, 108_989_220]  # In Mb from the start of chromosome 2
focal_variant = [v for v in ts.variants() if v.site.metadata.get("ID") == "rs3827760"].pop()
print("An interesting SNP within the EDAR gene:")
focal_variant
```

For simplicity, this tree sequence has been {ref}`simplified<sec_simplification>` to
include only those samples from the African and East Asian regions. These belong to a
number of populations. The population information, as well as information describing the
variable sites, is stored in tree sequence {ref}`metadata<sec_tutorial_metadata>`:

```{code-cell}
import pandas as pd

print(ts.num_populations, "populations defined in the tree sequence:")

pop_names_regions = [
    [p.metadata.get("name"), p.metadata.get("region")]
    for p in ts.populations()
]
display(pd.DataFrame(pop_names_regions, columns=["population name", "region"]))
```

You can see that there are multiple African and East asian populations, grouped by
region. Here we collect two lists of IDs for the sample
{ref}`nodes<sec_terminology_nodes>` from the African region and from the East asian
region:

```{code-cell}

sample_lists = {}
for n, rgns in {"Africa": {'AFRICA', 'Africa'}, "East asia": {'EAST_ASIA', 'EastAsia'}}.items():
    pop_ids = [p.id for p in ts.populations() if p.metadata.get("region") in rgns]
    sample_lists[n] = [u for p in pop_ids for u in ts.samples(population=p)]
```


With these lists we can calculate different windowed statistics
(here {meth}`genetic diversity<TreeSequence.diversity>` and
{meth}`Tajima's D<TreeSequence.Tajimas_D>`) within each of these regions:

```{code-cell}
edar_ts = ts.trim()  # remove regions with no data (changes the coordinate system)
windows = list(range(0, int(edar_ts.sequence_length)+1, 10_000))
data = {
    "Genetic diversity": {
        region: edar_ts.diversity(samples, windows=windows)
        for region, samples in sample_lists.items()
    },
    "Tajima's D": {
        region: edar_ts.Tajimas_D(samples, windows=windows)
        for region, samples in sample_lists.items()
    },  
}

# Plot the `data`
fig, axes = plt.subplots(ncols=2, figsize=(15, 3))
start = ts.edges_left.min()  # the empty amount at the start of the tree sequence

for (title, plot_data), ax in zip(data.items(), axes):
    ax.set_title(title)
    ax.axvspan(edar_gene_bounds[0], edar_gene_bounds[1], color="lightgray")
    ax.axvline(focal_variant.site.position, ls=":")
    for label, stat in plot_data.items():
        ax.stairs(stat, windows+start, baseline=None, label=label)
    ax.text(edar_gene_bounds[0], 0, "EDAR")
    ax.legend()
plt.show()
```

Other population genetic libraries such as
[scikit-allel](https://scikit-allel.readthedocs.io/en/stable/) (which is
{ref}`interoperable<sec_tskit_getting_started_exporting_data_allel>` with `tskit`)
could also have been used to produce the plot above. In this case, the advantage of
using tree sequences is simply that they allow these sorts of analysis to
{ref}`scale<plot_incremental_calculation>` to datasets of millions of whole genomes.

(sec_popgen_topological)=

### Topological analysis

As this inferred tree sequence stores (an estimate of) the underlying
genealogy, we can also derive statistics based on genealogical relationships. For
example, this tree sequence also contains a sample genome based on an ancient
genome, a [Denisovan](https://en.wikipedia.org/wiki/Denisovan) individual. We can
look at the closeness of relationship between samples from the different geographical
regions and the Denisovan:

:::{todo}
Show an example of looking at topological relationships between the Denisovan and
various East Asian groups, using the {ref}`sec_counting_topologies` functionality.
:::

See {ref}`sec_counting_topologies` for an introduction to topological methods in
`tskit`.

## Further information

This brief introduction is meant as a simple taster. Many other efficient population
genetic {ref}`analyses<sec_analysing_tree_sequences>` are possible when you have
genomic data stored as a tree sequence.

The rest of the {ref}`tutorials<sec_intro>` contain a large number of examples which
are relevant to population genetic analysis and research. You can also visit the
[learning section](https://tskit.dev/learn/) of the [tskit website](https://tskit.dev/).


--- ../../tskit/docs/counting_topologies.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

```{currentmodule} tskit
```

(sec_counting_topologies)=

```{code-cell} ipython3
:tags: [remove-cell]
import msprime
import stdpopsim

def topologies_sim_speciestree():
    newick_species_tree = "((A:100.0,B:100.0):100.0,C:200.0)"
    demography = msprime.Demography.from_species_tree(newick_species_tree, initial_size=100)
    ts = msprime.sim_ancestry({0: 2, 1: 2, 2: 2}, demography=demography, random_seed=321)
    ts.dump("data/topologies_sim_speciestree.trees")

def topologies_sim_stdpopsim():
    species = stdpopsim.get_species("HomSap")
    model = species.get_demographic_model("OutOfAfrica_3G09")
    contig = species.get_contig("chr1", length_multiplier=0.0002, mutation_rate=model.mutation_rate)
    samples = {"YRI": 1000, "CEU": 1000, "CHB": 1000}
    engine = stdpopsim.get_engine("msprime")
    ts = engine.simulate(model, contig, samples, seed=321)
    ts.dump("data/topologies_sim_stdpopsim.trees")


def create_notebook_data():
    topologies_sim_speciestree()
    topologies_sim_stdpopsim()

# create_notebook_data()  # uncomment to recreate the tree seqs used in this notebook
```

# Counting topologies

**Yan Wong**

This tutorial is intended to be a gentle introduction to the combinatorial
treatment of tree topologies in `tskit`. For a more formal introduction,
see the {ref}`sec_combinatorics` section of the
official `tskit` {ref}`documentation<tskit:sec_introduction>`.

The *topology* of a single tree is the term used to describe the branching pattern,
regardless of the lengths of the branches. For example, both trees below have the
same topology, although the branch lengths differ:

```{code-cell}
import tskit
node_labels = {0: "a", 1: "b", 2: "c"}  # avoid confusion by using letters to label tips
tree = tskit.Tree.generate_comb(3)
display(tree.draw_svg(node_labels=node_labels, y_axis=True))

deep_tree = tskit.Tree.generate_comb(10).tree_sequence.simplify([0, 1, 2]).first()
display(deep_tree.draw_svg(node_labels=node_labels, y_axis=True))
```

:::{note}
The treatment of topologies in `tskit` is restricted to trees with a single defined root,
without nodes with a single child (i.e. trees must consist of nodes that are either leaves,
or internal nodes with two or more children).  For convenience in the examples
below, trees are drawn with the tips flagged as samples, although whether a node is a sample or
not does not change the topology of the tree.
:::

## Tree labellings and shapes

The topology of a tree also takes into account the labelling of tips, so that
the trees below, although they have the same *shape*, count as three
different topologies:

```{code-cell}
:tags: [hide-input]
from string import ascii_lowercase
from IPython.display import SVG

def str_none(s, prefix=None):
    if s is not None:
        if prefix is None:
            return str(s)
        else:
            return prefix + " = " + str(s)
    return None

def draw_svg_trees(trees, node_labels={}, x_lab_attr=None, width=100, height=150, space=10):
    w = width + space
    h = height + space
    trees = list(trees)
    s = f'<svg height="{h}" width="{w * len(trees)}" xmlns="http://www.w3.org/2000/svg">'
    s += f'<style>.x-axis {{transform: translateY({space}px)}}</style>'
    for i, tree in enumerate(trees):
        s += tree.draw_svg(
            size=(width, height),
            canvas_size=(w, h),
            root_svg_attributes={"x": i * w},
            node_labels=node_labels,
            x_label=str_none(getattr(tree.rank(), x_lab_attr or "", None), x_lab_attr)
        )
    s += '</svg>'
    return SVG(s)

draw_svg_trees(tskit.all_tree_labellings(tree), node_labels={u: ascii_lowercase[u] for u in tree.samples()})
```

These are, in fact, the only possible three labellings for a three-tip tree of that shape.
There is only one other possible shape for a three-tip tree, and for this shape,
all labelling orders are equivalent (in other words, there is only one
possible labelling):

```{code-cell}
:tags: [hide-input]
tskit.Tree.generate_star(3).draw_svg(node_labels={})
```

A 3-tip tree therefore has only four possible topologies.
These can be generated with the {func}`~tskit.all_trees` function.

```{code-cell}
generated_trees = tskit.all_trees(3)
print("For a three-tip tree there are", len(list(generated_trees)), "labelled topologies.")
```

Here they are, plotted out with their shapes enumerated from zero:

```{code-cell}
:tags: [hide-input]
draw_svg_trees(
    tskit.all_trees(3),
    node_labels={u: ascii_lowercase[u] for u in tree.samples()},
    x_lab_attr="shape"
)
```

### Enumerating shapes and labellings

For a tree with four tips, more topologies and shapes are possible. As before, we can generate the
topologies using {func}`~tskit.all_trees`. Alternatively, if we only want the (unlabelled) shapes,
we can use the {func}`~tskit.all_tree_shapes` function:

```{code-cell}
print("For a four-tip tree there are", len(list(tskit.all_trees(4))), "labelled topologies.")

generated_trees = tskit.all_tree_shapes(4)
print("These can be categorised into", len(list(generated_trees)), "shapes.")
```

Again, we can give each shape a number or *index*, starting from zero:

```{code-cell}
:tags: [hide-input]
draw_svg_trees(tskit.all_tree_shapes(4), x_lab_attr="shape")
```

Each of these shapes will have a separate number of possible labellings, and trees with
these labellings can be created using {func}`~tskit.all_tree_labellings`:

```{code-cell}
for shape_index, tree in enumerate(tskit.all_tree_shapes(4)):
    labellings = tskit.all_tree_labellings(tree)
    num_labellings = len(list(labellings))
    print(
        f"Tree shape {shape_index} for a four-tip tree has "
        f"{num_labellings} labelling{'' if num_labellings==1 else 's'}."
    )
```

Any tree topology for a tree of $N$ tips can therefore be described by a
shape index combined with a labelling index. This is known as the
*rank* of a tree, and it can be obtained using the
{meth}`Tree.rank` method. For instance, here is the rank of a simulated tree
of 10 tips:

```{code-cell}
:tags: [hide-input]
import msprime
num_tips = 10
simulated_ts = msprime.sim_ancestry(10, ploidy=1, random_seed=123)
simulated_tree = simulated_ts.first()
print("The topology of the simulated tree below can be described as", simulated_tree.rank())
ascii_node_labels = {u: ascii_lowercase[u] for u in simulated_tree.samples()}
simulated_tree.draw_svg(node_labels=ascii_node_labels)
```


A tree with the same topology (i.e. the same shape and labelling, but ignoring
the branch lengths) can be generated using the {meth}`Tree.unrank` method, by
specifying the number of tips and the appropriate `(shape, labelling)` tuple:

```{code-cell}
new_tree = tskit.Tree.unrank(num_tips, (1270, 21580))
new_tree.draw_svg(node_labels=ascii_node_labels)
```

Note that this method generates a single tree in a new tree sequence
whose a default sequence length is 1.0.

## Methods for large trees

The number of possible topologies for a tree with $N$ tips
grows very rapidly with $N$. For instance, with 10 tips, there are
282,137,824 possible topologies.

For this reason, the {func}`~tskit.all_trees`, {func}`~tskit.all_tree_shapes` and
{func}`~tskit.all_tree_labellings` methods do not return a list of trees
but an iterator over the trees. This means it is perfectly possible to start
iterating over (say) all tree shapes for a tree of 100 leaves, but
the iterator will not finish before the death of our galaxy.

```{code-cell}
for num_trees, tree in enumerate(tskit.all_tree_shapes(100)):
    shape = tree.rank().shape
    b2 = tree.b2_index()
    print(f"A 100-tip tree with shape index {shape} has a b2 balance index of {b2}")
    if num_trees > 5:
      break  # better not let this run too long!
```

For similar combinatorial reasons, the {meth}`Tree.rank` method can be
inefficient for large trees. To compare the topology of two trees, you are
therefore recommended to use e.g. the {meth}`Tree.kc_distance` method
rather than comparing ranks directly.

```{code-cell}
simulated_tree = simulated_ts.first(sample_lists=True)  # kc_distance requires sample lists
if simulated_ts.first(sample_lists=True).kc_distance(simulated_tree) == 0:
    print("Trees are identical")
    # To compare to the new_tree we need to fix 
    # print("The simulated and topology-constructed trees have the same topology")
```

Despite the combinatorial explosion associated with topologies of
many-tip trees, it is still possible to efficiently count
the number of *embedded topologies* in a large tree.

### Embedded topologies

An embedded topology is a a topology involving a subset of the tips of a tree.
If the tips are classified into (say) three groups, red, green, and blue,
we can efficiently count all the embedded three-tip trees which have
one tip from each group using the {meth}`Tree.count_topologies` method.

```{code-cell}
:tags: [hide-input]
big_tree = tskit.load("data/topologies_sim_speciestree.trees").first()
# Check all observed topologies have the same counts
assert list(big_tree.count_topologies()[0, 1, 2].values()) == [32, 32]
styles = [
    f".node.sample.p{p.id} > .sym " + "{" + f"fill: {colour}" + "}"
    for colour, p in zip(['red', 'green', 'blue'], big_tree.tree_sequence.populations())
]
big_tree.draw_svg(style="".join(styles), node_labels={}, time_scale="rank", x_label="big_tree")
```

In this tree, it is clear that the green and blue tips never cluster together.
The {meth}`Tree.count_topologies` method exhaustively looks at all
combinations of one red, one blue, and one green tip, and confirms that we never see
the topology grouping green and blue. However, as might be expected from
examination of the plot above, a red tip is equally likely to be a sister to a
green tip as to a blue tip:

```{code-cell}
# By default `count_topologies` chooses one tip from each population, like setting
# sample_sets=[ts.samples(p.id) for p in ts.populations() if len(ts.samples(p.id)) > 0]

topology_counter = big_tree.count_topologies()

colours = ['red', 'green', 'blue']
styles = [f".n{u}>.sym {{fill: {c} }}" for u, c in enumerate(colours)]

embedded_counts = topology_counter[0, 1, 2]
for embedded_tree in tskit.all_trees(3):
    rank = embedded_tree.rank()
    number_of_instances = embedded_counts[rank]
    label = f"{number_of_instances} instances embedded in big_tree"
    display(embedded_tree.draw_svg(style="".join(styles), node_labels={}, x_label=label))
```

## Methods over tree sequences

It can be useful to count embedded topologies over an entire tree sequence.
For instance, we might want to know the number of embedded topologies
that support Neanderthals as a sister group to europeans versus africans.
`Tskit` provides the efficient {meth}`TreeSequence.count_topologies` method to
do this [incrementally](sec_incremental), without having to re-count the topologies
independently in each tree.

```{code-cell}
:tags: [hide-input]
from myst_nb import glue
ts = tskit.load("data/topologies_sim_stdpopsim.trees")
print(f"Loaded a stdpopsim of {ts.num_trees} African+European+Chinese trees, each with {ts.num_samples} tips")
glue("seq_len", int(ts.sequence_length/1000), display=False)
```

Although the trees in this tree sequence are very large, counting the embedded topologies is
quite doable (for speed in this demo we are only simulating {glue:}`seq_len` kilobases, but
calculating the average over an entire chromosome simply takes a little longer)

```{code-cell}
from datetime import datetime
names = {"YRI": "African", "CEU": "European", "CHB": "Chinese"}
colours = {"YRI": "yellow", "CEU": "green", "CHB": "blue"}

population_map = {p.metadata["id"]: p.id for p in ts.populations()}
sample_populations = list(sorted({ts.node(u).population for u in ts.samples()}))
topology_span = {tree.rank(): 0 for tree in tskit.all_trees(len(sample_populations))}

start = datetime.now()
total = 0
for topology_counter, tree in zip(ts.count_topologies(), ts.trees()):
    embedded_topologies = topology_counter[sample_populations]
    weight = tree.span / ts.sequence_length
    for rank, count in embedded_topologies.items():
        topology_span[rank] += count * weight
        total += count
print(f"Counted {total} embedded topologies in {datetime.now() - start} seconds")
```

```{code-cell}
:tags: [hide-input]
ntips = len(sample_populations)
styles = ".sample text.lab {baseline-shift: super; font-size: 0.7em;}"
node_labels = {}

for p in range(ntips):
    name = ts.population(sample_populations[p]).metadata["id"]
    node_labels[p] = names[name]
    styles += f".n{p}>.sym {{fill: {colours[name]} }}"

total = sum(topology_span.values())
for rank, weight in topology_span.items():
    label = f"{weight/total *100:.1f}% of genome"
    embedded_tree = tskit.Tree.unrank(ntips, rank)
    display(embedded_tree.draw_svg(size=(160, 150), style="".join(styles), node_labels=node_labels, x_label=label))
```

Perhaps unsurprisingly, the most common topology is the one that groups the non-African
populations together (although there are many trees of the other two topologies,
mostly reflecting genetic divergence prior to the emergence of humans out of Africa).

For an example with real data, see {ref}`sec_popgen_topological`
in the {ref}`sec_intro_popgen` tutorial.

--- ../../tskit/docs/stats.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

```{currentmodule} tskit
```


(sec_stats)=

# Statistics

The `tskit` library provides a large number of functions for calculating population
genetic statistics from tree sequences. Statistics can reflect either the distribution
of genetic variation or the underlying trees that generate it; the
[duality](https://doi.org/10.1534/genetics.120.303253) between
mutations and branch lengths on trees means that statistics based on genetic variation
often have an corresponding version based on branch lengths in trees in a tree sequence.

Note that `tskit` provides a unified interface for computing so-called
"single site statistics" that are summaries across sites in windows of the genome, as
well as a standard method for specifying "multi-way" statistics that are calculated
over many combinations of sets of samples simultaneously.

Please see the {ref}`tutorial <tutorials:sec_tutorial_stats>` for examples of the
statistics API in use.

:::{warning}
{ref}`sec_data_model_missing_data` is not currently
handled correctly by site statistics defined here, as we always
assign missing data to be equal to the ancestral state. Later
versions will add this behaviour as an option and will account
for the presence of missing data by default.
:::

(sec_stats_available)=

## Available statistics

Here are the statistics that can be computed using `tskit`,
grouped by basic classification and type. Single-site statistics are ones that are
averages across sites in windows of the genome, returning numpy arrays whose
dimensions are determined by the parameters (see {ref}`sec_stats_output_dimensions`).

{ref}`sec_stats_sample_sets_one_way` are defined over a single sample set, 
whereas {ref}`sec_stats_sample_sets_multi_way` compare 2 or more sets of samples.

Some of the methods below benefit from a little extra discussion, provided in the
{ref}`sec_stats_notes` section at the end of this chapter: if so, a link to the note
appears beside the listed method.

* Single site
    * One-way
        * {meth}`~TreeSequence.allele_frequency_spectrum` (see {ref}`notes<sec_stats_notes_afs>`)
        * {meth}`~TreeSequence.diversity`
        * {meth}`~TreeSequence.segregating_sites`
        * {meth}`~TreeSequence.trait_covariance`
          {meth}`~TreeSequence.trait_correlation`
          {meth}`~TreeSequence.trait_linear_model`
          (see {ref}`sec_stats_notes_trait`)
        * {meth}`~TreeSequence.Tajimas_D` (see {ref}`notes<sec_stats_notes_derived>`)
    * Multi-way
        * {meth}`~TreeSequence.divergence`
        * {meth}`~TreeSequence.genetic_relatedness`
          {meth}`~TreeSequence.genetic_relatedness_weighted`
        * {meth}`~TreeSequence.f4`
          {meth}`~TreeSequence.f3`
          {meth}`~TreeSequence.f2`
          (see {ref}`sec_stats_notes_f`)
        * {meth}`~TreeSequence.Y3`
          {meth}`~TreeSequence.Y2`
          (see {ref}`sec_stats_notes_y`)
        * {meth}`~TreeSequence.genealogical_nearest_neighbours` (see {ref}`sec_stats_notes_gnn`)
        * {meth}`~TreeSequence.Fst` (see {ref}`sec_stats_notes_derived`)
* Multi site
    * {meth}`~LdCalculator` (note this is soon to be deprecated)

:::{note}
There is a general framework provided for calculating additional single site
statistics (see the {ref}`sec_stats_general_api` section). However, the
pre-implemented statistics in the table above will be faster than rederiving
them using the general framework directly, so the versions above should be preferred.
:::


(sec_stats_examples)=

### Quick examples

```{code-cell} ipython3
:"tags": ["hide-input"]
from IPython.display import Markdown
import msprime
import numpy as np

demography = msprime.Demography()
demography.add_population(name="A", initial_size=10_000)
demography.add_population(name="B", initial_size=10_000)
demography.set_symmetric_migration_rate(["A", "B"], 0.001)
ts = msprime.sim_ancestry(
    samples={"A": 2, "B": 2},
    sequence_length=1000,
    demography=demography,
    recombination_rate=2e-8,
    random_seed=12)
ts = msprime.sim_mutations(ts, rate=2e-8, random_seed=12)
Markdown(
    f"These examples use a tree sequence of {ts.num_samples} samples "
    f"in {ts.num_populations} populations, "
    f"with a sequence length of {int(ts.sequence_length)}. "
    f"There are {ts.num_trees} trees and "
    f"{ts.num_sites} variable sites in the tree sequence."
)
```

#### Basic calling convention

```{code-cell} ipython3
pi = ts.diversity()
print(pi) # Genetic diversity within the sample set
```

#### Restrict to {ref}`sample sets<sec_stats_sample_sets>`

```{code-cell} ipython3
pi_0 = ts.diversity(sample_sets=ts.samples(population=0))
print(pi_0)  # Genetic diversity within population 0
```

#### Summarise in genomic {ref}`windows<sec_stats_windows>`

```{code-cell} ipython3
pi_window = ts.diversity(sample_sets=ts.samples(population=1), windows=[0, 400,  600, 1000])
print(pi_window)  # Genetic diversity within population 1 in three windows along the genome
```

#### Compare {ref}`between<sec_stats_sample_sets_multi_way>` sample sets

```{code-cell} ipython3
dxy = ts.divergence(sample_sets=[ts.samples(population=0), ts.samples(population=1)])
print(dxy)  # Av number of differences per bp between samples in population 0 and 1
```

#### Change the {ref}`mode<sec_stats_mode>`

```{code-cell} ipython3
bl = ts.divergence(
    mode="branch",  # Use branch lengths rather than genetic differences
    sample_sets=[ts.samples(population=0), ts.samples(population=1)],
)
print(bl)  # Av branch length separating samples in population 0 and 1
```

(sec_stats_single_site)=

## Single site statistics


(sec_stats_interface)=

### Interface

Tskit offers a powerful and flexible interface for computing population genetic
statistics. Consequently, the interface is a little complicated and there are a
lot of options. However, we provide sensible defaults for these options and
`tskit` should do what you want in most cases. There are several major options
shared by many statistics, which we describe in detail in the following subsections:

{ref}`sec_stats_mode`
: What are we summarising information about?

{ref}`sec_stats_windows`
: What section(s) of the genome are we interested in?

{ref}`sec_stats_span_normalise`
: Should the statistic calculated for each window be normalised by the span
  (i.e. the sequence length) of that window?

The statistics functions are highly efficient and are based where possible
on numpy arrays. Each of these statistics will return the results as a numpy
array, and the format of this array will depend on the statistic being
computed (see the {ref}`sec_stats_output_format` section for details).
A convenient feature of the statistics API is that the dimensions of the
output array is defined in a simple and intuitive manner by the
parameters provided. The {ref}`sec_stats_output_dimensions` section
defines the rules that are used.

Please see the {ref}`tutorial <sec_tutorial_stats>` for examples of the
statistics API in use.


(sec_stats_mode)=

#### Mode

There are three **modes** of statistic: `site`, `branch`, and `node`,
that each summarize aspects of the tree sequence in different but related ways.
Roughly speaking, these answer the following sorts of question:

site
: How many mutations differentiate these two genomes?

branch
: How long since these genomes' common ancestor?

node
: On how much of the genome is each node an ancestor of only one of these genomes, but not both?

These three examples can all be answered in the same way with the tree sequence:
first, draw all the paths from one genome to the other through the tree sequence
(back up to their common ancestor and back down in each marginal tree).
Then,
(`site`) count the number of mutations falling on the paths,
(`branch`) measure the length of the paths, or
(`node`) count how often the path goes through each node.
There is more discussion of this correspondence in the paper describing these statistics,
and precise definitions are given in each statistic.

Here's an example of using the {meth}`~TreeSequence.diversity` statistic to return the
average branch length between all pairs of samples:

```{code-cell} ipython3
ts.diversity(mode="branch")
```

One important thing to know is that `node` statistics have somewhat different output.
While `site` and `branch` statistics naturally return one number
for each portion of the genome (and thus incorporates information about many nodes: see below),
the `node` statistics return one number **for each node** in the tree sequence (and for each window).
There can be a lot of nodes in the tree sequence, so beware.

Also remember that in a tree sequence the "sites" are usually just the **variant** sites,
e.g., the sites of the SNPs. Although the tree sequence may in principle have monomorphic
sites, those produced by simulation usually don't.


(sec_stats_sample_sets)=

#### Sample sets and indexes

Many standard population genetics statistics
are defined with respect to some number of groups of genomes,
usually called "populations".
(However, it's clear from the correspondence to descriptors of tree shape
that the definitions can usefully describe *something*
even if the groups of samples don't come from "separate populations" in some sense.)
Basically, statistics defined in terms of sample sets can use the frequency of any allele
in each of the sample sets when computing the statistic.
For instance, nucleotide divergence is defined for a *pair* of groups of samples,
so if you wanted to compute pairwise divergences between some groups of samples,
you'd specify these as your `sample_sets`.
Then, if `p[i]` is the derived allele frequency in sample set `i`,
under the hood we (essentially) compute the divergence between sample sets `i` and `j`
by averaging `p[i] * (1 - p[j]) + (1 - p[i]) * p[j]` across the genome.

Concretely, `sample_sets` specifies the IDs of the nodes to compute statistics of.
Importantly, these nodes must be {ref}`samples <sec_data_model_definitions_sample>`.

Here's an example of calculating the average
{meth}`genetic diversity<TreeSequence.diversity>` within a specific population:

```{code-cell} ipython3
ts.diversity(sample_sets=ts.samples(population=0))
```


So, what if you
have samples from each of 10 populations,
and want to compute **all** fourty-five pairwise divergences?
You could call `divergence` fourty-five times, but this would be tedious
and also inefficient, because the allele frequencies for one population
gets used in computing many of those values.
So, statistics that take a `sample_sets` argument also take an `indexes` argument,
which for a statistic that operates on `k` sample sets will be a list of `k`-tuples.
If `indexes` is a length `n` list of `k`-tuples,
then the output will have `n` columns,
and if `indexes[j]` is a tuple `(i0, ..., ik)`,
then the `j`-th column will contain values of the statistic computed on
`(sample_sets[i0], sample_sets[i1], ..., sample_sets[ik])`.


How multiple statistics are handled differs slightly between statistics
that operate on single sample sets and multiple sample sets.


(sec_stats_sample_sets_one_way)=

##### One-way methods

One-way statistics such as {meth}`TreeSequence.diversity` are defined over a single
sample set. For these methods, `sample_sets` is interpreted in the following way:

- If it is a single list of node IDs (e.g., `sample_sets=[0, 1 ,2]`), this is
  interpreted as running the calculation over one sample set and we remove
  the last dimension in the result array as described in the
  {ref}`sec_stats_output_dimensions` section.

- If it is `None` (the default), this is equivalent to `sample_sets=ts.samples()`,
  and we therefore compute the statistic over all samples in the tree sequence. **Note
  that we also drop the outer dimension of the result array in this case**.

- If it is a list of lists of samples we return an array for each window in the output,
  which contains the value of the statistic separately for each of `sample_sets`
  in the order they are given.


(sec_stats_sample_sets_multi_way)=

##### Multi-way methods

Multi-way statistics such as {meth}`TreeSequence.divergence` are defined over a
`k` sample sets. In this case, `sample_sets` must be a list of lists of sample IDs,
and there is no default. For example, this finds the average
{meth}`genetic divergence<TreeSequence.divergence>` between samples in populations
0 and 1

```{code-cell} ipython3
ts.divergence(
    sample_sets=[
        ts.samples(population=0),
        ts.samples(population=1),
    ]
)
```


The `indexes` parameter is interpreted in the following way:

- If it is a single `k`-tuple, this is interpreted as computing a single
  statistic selecting the specified sample sets and we remove the last dimension
  in the result array as described in the {ref}`sec_stats_output_dimensions` section.

- If it is `None` and `sample_sets` contains exactly `k` sample sets,
  this is equivalent to `indexes=range(k)`. **Note
  that we also drop the outer dimension of the result array in this case**.

- If it is a list of `k`-tuples (each consisting of integers
  of integers between `0` and `len(sample_sets) - 1`) of length `n` we
  compute `n` statistics based on these selections of sample sets.


(sec_stats_windows)=

#### Windows

Each statistic has an argument, `windows`,
which defines a collection of contiguous windows spanning the genome.
`windows` should be a list of `n+1` increasing numbers beginning with 0
and ending with the `sequence_length`.
The statistic will be computed separately in each of the `n` windows,
and the `k`-th row of the output will report the values of the statistic
in the `k`-th window, i.e., from (and including) `windows[k]` to
(but not including) `windows[k+1]`. For example, this calculates
{meth}`Tajima's D<TreeSequence.Tajimas_D>` in four evenly spaced windows along the
genome:

```{code-cell} ipython3
num_windows = 4
ts.Tajimas_D(windows=np.linspace(0, ts.sequence_length, num_windows + 1))
```

Most windowed statistics by default return **averages** within each of the windows,
so the values are comparable between windows, even of different spans.
(However, shorter windows may be noisier.)
Suppose for instance  that you compute some statistic with `windows = [0, a, b]`
for some valid positions `0 < a < b`,
and get an output array `S` with two rows.
Then, computing the same statistic with `windows = [0, b]`
would be equivalent to averaging the rows of `S`,
obtaining `((a - 0) * S[0] + (b - a) * S[1]) / (b - 0)`.

There are some shortcuts to other useful options:

`windows = None`
   This is the default and computes statistics in single window over the whole
   sequence. As the first returned array contains only a single
   value, we drop this dimension as described in the
   {ref}`output dimensions <sec_stats_output_dimensions>` section. **Note:** if you
   really do want to have an array with a single value as the result, please use
   `windows = [0.0, ts.sequence_length]`.

`windows = "trees"`
   This says that you want statistics computed separately on the portion of the genome
   spanned by each tree, so is equivalent to passing `windows = ts.breakpoints()`.
   (Beware: there can be a lot of trees!)

`windows = "sites"`
   This says to output one set of values for **each site**,
   and is equivalent to passing `windows = [s.position for s in ts.sites()] + [ts.sequence_length]`.
   This will return one statistic for each site (beware!);
   since the windows are all different sizes you probably want to also pass
   `span_normalise=False` (see below).


(sec_stats_span_normalise)=

#### Span normalise

In addition to windowing there is an option, `span_normalise` (which defaults to `True`),
All the primary statistics defined here are *sums* across locations in the genome:
something is computed for each position, and these values are added up across all positions in each window.
Whether the total span of the window is then taken into account is determined by the option `span_normalise`:
if it is `True` (the default), the sum for each window is converted into an *average*,
by dividing by the window's *span* (i.e. the length of genome that it covers).
Otherwise, the sum itself is returned.
The default is `span_normalise=True`,
because this makes the values comparable across windows of different sizes.
To make this more concrete: {meth}`pairwise sequence divergence <.TreeSequence.divergence>`
between two samples with `mode="site"` is the density of sites that differ between the samples;
this is computed for each window by counting up the number of sites
at which the two differ, and dividing by the total span of the window.
If we wanted the number of sites at which the two differed in each window,
we'd calculate divergence with `span_normalise=False`.

Following on from above, suppose we computed the statistic `S` with
`windows = [0, a, b]` and `span_normalise=True`,
and then computed `T` in just the same way except with `span_normalise=False`.
Then `S[0]` would be equal to `T[0] / a` and `S[1] = T[1] / (b - a)`.
Furthermore, the value obtained with `windows = [0, b]` would be equal to `T[0] + T[1]`.
However, you probably usually want the (default) normalized version:
don't get unnormalised values unless you're sure that's what you want.
The exception is when computing a site statistic with `windows = "sites"`:
this case, computes a statistic with the pattern of genotypes at each site,
and normalising would divide these statistics by the distance to the previous variant site
(probably not what you want to do).

:::{note}
The resulting values are scaled "per unit of sequence length" - for instance, pairwise
sequence divergence is measured in "differences per unit of sequence length". Functions
such as {func}`msprime:msprime.sim_mutations` will by default add mutations in discrete
coordinates, usually interpreted as base pairs, in which
case span normalised statistics are in units of "per base pair".
:::

(sec_stats_output_format)=

#### Output format

Each of the statistics methods returns a `numpy` ndarray.
Suppose that the output is named `out`.
If `windows` has been specified, the number of rows of the output is equal to the
number of windows, so that `out.shape[0]` is equal to `len(windows) - 1`
and `out[i]` is an array of statistics describing the portion of the tree sequence
from `windows[i]` to `windows[i + 1]` (including the left but not the right endpoint).
What is returned within each window depends on the {ref}`mode <sec_stats_mode>`:

`mode="site"` or `mode="branch"`
   The output is a two-dimensional array,
   with columns corresponding to the different statistics computed: `out[i, j]` is the `j`-th statistic
   in the `i`-th window.

`mode="node"`
   The output is a three-dimensional array,
   with the second dimension corresponding to node id.
   In other words, `out.shape[1]` is equal to `ts.num_nodes`,
   and `out[i,j]` is an array of statistics computed for node `j` on the `i`-th window.

The final dimension of the arrays in other cases is specified by the method.

Note, however, that empty dimensions can optionally be dropped,
as described in the {ref}`sec_stats_output_dimensions` section.

A note about **default values** and **division by zero**:
Under the hood, statistics computation fills in zeros everywhere, then updates these
(since statistics are all **additive**, this makes sense).
But now suppose that you've got a statistic that returns `nan`
("not a number") sometimes, like if you're taking the diversity of a sample set with only `n=1` sample,
which involves dividing by `n * (n - 1)`.
Usually, you'll just get `nan` everywhere that the division by zero happens.
But there's a couple of caveats.
For `site` statistics, any windows without any sites in them never get touched,
so they will have a value of 0.
For `branch` statistics, any windows with **no branches** will similarly remain 0.
That said, you should **not** rely on the specific behavior of whether `0` or `nan` is returned
for "empty" cases like these: it is subject to change.


(sec_stats_output_dimensions)=

#### Output dimensions

In the general case, tskit outputs two dimensional (or three dimensional, in the case of node
stats) numpy arrays, as described in the {ref}`sec_stats_output_format` section.
The first dimension corresponds to the window along the genome
such that for some result array `x`, `x[j]` contains information about the jth window.
The last dimension corresponds to the statistics being computed, so that `x[j, k]` is the
value of the kth statistic in the jth window (in the two dimensional case). This is
a powerful and general interface, but in many cases we will not use this full generality
and the extra dimensions in the numpy arrays are inconvenient.

Tskit optionally removes empty dimensions from the output arrays following a few
simple rules.

1. If `windows` is None we are computing over the single window covering the
   full sequence. We therefore drop the first dimension of the array.

2. In one-way stats, if the `sample_sets` argument is a 1D array we interpret
   this as specifying a single sample set (and therefore a single statistic), and
   drop the last dimension of the output array. If `sample_sets` is None
   (the default), we use the sample set `ts.samples()`, invoking
   this rule (we therefore drop the last dimension by default).

3. In k-way stats, if the `indexes` argument is a 1D array of length k
   we intepret this as specifying a single statistic and drop the last
   dimension of the array. If `indexes` is None (the default) and
   there are k sample sets, we compute the statistic from these sample sets
   and drop the last dimension.

4. If, after dropping these dimensions, the dimension is 0, we return a numpy
   scalar (instead of an array of dimension 0).

Rules 2 and 3 can be summarised by "the dimensions of the input determines
the dimensions of the output". Note that dropping these dimensions is
**optional**: it is always possible to keep the full dimensions of the
output arrays.

Please see the {ref}`tutorial <sec_tutorial_stats>` for examples of the
various output dimension options.


(sec_stats_general_api)=

### General API

The methods {meth}`TreeSequence.general_stat` and {meth}`TreeSequence.sample_count_stat`
provide access to the general-purpose algorithm for computing statistics.
Here is a bit more discussion of how to use these.


(sec_stats_polarisation)=

#### Polarisation

Many statistics calculated from genome sequence treat all alleles on equal footing,
as one must without knowledge of the ancestral state and sequence of mutations that produced the data.
Separating out the *ancestral* allele (e.g., as inferred using an outgroup)
is known as *polarisation*.
For instance, in the allele frequency spectrum, a site with alleles at 20% and 80% frequency
is no different than another whose alleles are at 80% and 20%,
unless we know in each case which allele is ancestral,
and so while the unpolarised allele frequency spectrum gives the distribution of frequencies of *all* alleles,
the *polarised* allele frequency spectrum gives the distribution of frequencies of only *derived* alleles.

This concept is extended to more general statistics as follows.
For site statistics, summary functions are applied to the total weight or number of samples
associated with each allele; but if polarised, then the ancestral allele is left out of this sum.
For branch or node statistics, summary functions are applied to the total weight or number of samples
below, and above each branch or node; if polarised, then only the weight below is used.

(sec_stats_strictness)=

### Strictness, and which branches count?

Most statistics are not affected by invariant sites,
and hence do not depend on any part of the tree that is not ancestral to any of the sample sets.
However, some statistics are different: for instance, 
given a pair of samples, {meth}`TreeSequence.genetic_relatedness`
with `centre=False` (and `polarised=True`, the default for that method)
adds up the total number of alleles (or total area of branches) that is
either ancestral to both samples *or ancestral to neither*.
So, it depends on what else is in the tree sequence.
(For this reason, we don't recommend actually *using* this combination of options for genetic
relatedness.)

In terms of the summary function {math}`f(x)`, "not affected by invariant sites" translates to
{math}`f(0) = f(n) = 0`, where {math}`n` is the vector of sample set sizes.
By default, {meth}`TreeSequence.general_stat` checks if the summary function satisfies this condition,
and throws an error if not; this check can be disabled by setting `strict=False`.


(sec_stats_summary_functions)=

#### Summary functions

For convenience, here are the summary functions used for many of the statistics.
Below, {math}`x` denotes the number of samples in a sample set below a node,
`n` denotes the total size of a sample set, {math}`p = x / n`,
and boolean expressions (e.g., {math}`(x > 0)`) are interpreted as 0/1.

`diversity`
: {math}`f(x) = \frac{x (n - x)}{n (n-1)}`

  For an unpolarized statistic with biallelic loci, this calculates
  {math}`2 p (1-p)`.

`segregating_sites`
: {math}`f(x) =  (x > 0) (1 - x / n)`

  (Note: this works because if {math}`\sum_i p_1 = 1` then {math}`\sum_{i=1}^k (1-p_i) = k-1`.)

`Y1`
: {math}`f(x) = \frac{x (n - x) (n - x - 1)}{n (n-1) (n-2)}`

`divergence`
: {math}`f(x_1, x_2) = \frac{x_1 (n_2 - x_2)}{n_1 n_2}`,

  unless the two indices are the same, when the diversity function is used.

  For an unpolarised statistic with biallelic loci, this calculates
  {math}`p_1 (1-p_2) + (1 - p_1) p_2`.

`genetic_relatedness, centre=True`
: {math}`f(x_i, x_j) = (x_i / n_i - m)(x_j / n_j - m)`,

  where {math}`m = \frac{1}{n}\sum_{k=1}^n x_k` with {math}`n` the total number
  of sample sets.
  For a polarised statistic (the default) with biallelic loci, this calculates
  {math}`(p_1 - \bar{p}) (p_2 - \bar{p})`, where {math}`\bar{p}` is the average
  derived allele frequency across sample sets.

`genetic_relatedness, centre=False`
: {math}`f(x_i, x_j) = (x_i / n_i) (x_j / n_j)`.

  For an polarised statistic (the default) with biallelic loci, this calculates
  {math}`p_1 p_2`.

`genetic_relatedness_weighted, centre=True`
: {math}`f(w_i, w_j, x_i, x_j) = (x_i - w_i p) (x_j - w_j p)`,

  where {math}`p` is the proportion of all samples below the focal node,
  and {math}`w_j = \sum_{k=1}^n W_{kj}` is the sum of the weights in the {math}`j`th column of the weight matrix.

`genetic_relatedness_weighted, centre=False`
: {math}`f(w_i, w_j, x_i, x_j) = x_i x_j`.


`Y2`
: {math}`f(x_1, x_2) = \frac{x_1 (n_2 - x_2) (n_2 - x_2 - 1)}{n_1 n_2 (n_2 - 1)}`

`f2`
: {math}`f(x_1, x_2) = \frac{x_1 (x_1 - 1) (n_2 - x_2) (n_2 - x_2 - 1)}{n_1 (n_1 - 1) n_2 (n_2 - 1)} - \frac{x_1 (n_1 - x_1) (n_2 - x_2) x_2}{n_1 (n_1 - 1) n_2 (n_2 - 1)}`

  For an unpolarized statistic with biallelic loci, this calculates
  {math}`((p_1 - p_2)^2 - (p_1 (1-p_2)^2 + (1-p_1) p_2^2)/n_1 - (p_1^2 (1-p_2) + (1-p_1)^2 p_2)/n_2`
  {math}`+ (p_1 p_2 + (1-p_1)(1-p_2))/ n_1 n_2)(1 + \frac{1}{n_1 - 1})(1 + \frac{1}{n_2 - 1})`,
  which is the unbiased estimator for {math}`(p_1 - p_2)^2` from a finite sample.

`Y3`
: {math}`f(x_1, x_2, x_3) = \frac{x_1 (n_2 - x_2) (n_3 - x_3)}{n_1 n_2 n_3}`

`f3`
: {math}`f(x_1, x_2, x_3) = \frac{x_1 (x_1 - 1) (n_2 - x_2) (n_3 - x_3)}{n_1 (n_1 - 1) n_2 n_3} - \frac{x_1 (n_1 - x_1) (n_2 - x_2) x_3}{n_1 (n_1 - 1) n_2 n_3}`

  For an unpolarized statistic with biallelic loci, this calculates
  {math}`((p_1 - p_2)(p_1 - p_3) - p_1 (1-p_2)(1-p_3)/n_1 - (1-p_1) p_2 p_3/n_1)(1 + \frac{1}{n_1 - 1})`,
  which is the unbiased estimator for {math}`(p_1 - p_2)(p_1 - p_3)` from a finite sample.

`f4`
: {math}`f(x_1, x_2, x_3, x_4) = \frac{x_1 x_3 (n_2 - x_2) (n_4 - x_4)}{n_1 n_2 n_3 n_4} - \frac{x_1 x_4 (n_2 - x_2) (n_3 - x_3)}{n_1 n_2 n_3 n_4}`

  For an unpolarized statistic with biallelic loci, this calculates
  {math}`(p_1 - p_2)(p_3 - p_4)`.

`trait_covariance`
: {math}`f(w) = \frac{w^2}{2 (n-1)^2}`,

  where {math}`w` is the sum of all trait values of the samples below the node.

`trait_correlation`
: {math}`f(w, x) = \frac{w^2}{2 x (1 - x/n) (n - 1)}`,

  where as before {math}`x` is the total number of samples below the node,
  and {math}`n` is the total number of samples.

`trait_linear_model`
: {math}`f(w, z, x) = \frac{1}{2}\left( \frac{w - \sum_{j=1}^k z_j v_j}{x - \sum_{j=1}^k z_j^2} \right)^2`,

  where {math}`w` and {math}`x` are as before,
  {math}`z_j` is the sum of the j-th normalised covariate values below the node,
  and {math}`v_j` is the covariance of the trait with the j-th covariate.


## Multi site statistics

:::{todo}
Document statistics that use information about correlation between sites, such as
LdCalculator (and perhaps reference {ref}`sec_identity`). Note that if we have a general
framework which has the same calling conventions as the single site stats,
we can rework the sections above.
:::


(sec_stats_notes)=

## Notes


(sec_stats_notes_afs)=

### Allele frequency spectrum

Most single site statistics are based on the summaries of the allele frequency spectra
(AFS). The `tskit` AFS interface includes windowed and joint spectra,
using the same general pattern as other statistics,
but some of the details about how it is defined,
especially in the presence of multiple alleles per site, need to be explained.
If all sites are biallelic, then the result is just as you'd expect:
see the method documentation at {meth}`~TreeSequence.allele_frequency_spectrum` 
for the description.
Note that with `mode="site"`, we really do tabulate *allele* counts:
if more than one mutation on different parts of the tree produce the same allele,
it is the total number with this allele (i.e., inheriting *either* mutation)
that goes into the AFS.
The AFS with `mode="branch"` is the expected value for the Site AFS
with infinite-sites, biallelic mutation, so there is nothing surprising there,
either.

But, how do we deal with sites at which there are more than two alleles?
At each site, we iterate over the distinct alleles at that site,
and for each allele, count how many samples in each sample set
have inherited that allele.
For a concrete example, suppose that we are computing the AFS of a single
sample set with 10 samples, and are considering a site with three alleles:
*a*, *b*, and *c*,
which have been inherited by 6, 3, and 1 samples, respectively,
and that allele *a* is ancestral.
What we do at this site depends on if the AFS is polarised or not.

If we are computing the *polarised* AFS,
we add 1 to each entry of the output corresponding to each allele count
*except* the ancestral allele.
In our example, we'd add 1 to both `AFS[3]` and `AFS[1]`.
This means that the sum of all entries of a polarised, site AFS
should equal the total number of non-ancestral alleles in the tree sequence
that are ancestral to at least one of the samples in the tree sequence
but not ancestral to all of them.
The reason for this last caveat is that like with most statistics,
mutations that are not ancestral to *any* samples (not just those in the sample sets)
are not counted (and so don't even enter into `AFS[0]`),
and similarly for those alleles inherited by *all* samples.

Now, if we are computing the *unpolarised* AFS,
we add *one half* to each entry of the *folded* output
corresponding to each allele count *including* the ancestral allele.
What does this mean?
Well, `polarised=False` means that we cannot distinguish between an
allele count of 6 and an allele count of 4.
So, *folding* means that we would add our allele that is seen in 6 samples
to `AFS[4]` instead of `AFS[6]`.
So, in total, we will add 0.5 to each of `AFS[4]`, `AFS[3]`, and `AFS[1]`.
This means that the sum of an unpolarised AFS
will be equal to the total number of alleles that are inherited
by any of the samples in the tree sequence, divided by two.
Why one-half? Well, notice that if in fact the mutation that produced the *b*
allele had instead produced an *a* allele,
so that the site had only two alleles, with frequencies 7 and 3.
Then, we would have added 0.5 to `AFS[3]` *twice*.


(sec_stats_notes_trait)=

### Trait correlations

{meth}`~TreeSequence.trait_covariance`, {meth}`~TreeSequence.trait_correlation`, and
{meth}`~TreeSequence.trait_linear_model` compute correlations and covariances of traits
(i.e., an arbitrary vector) with allelic state, possibly in the context of a multivariate
linear model with other covariates (as in GWAS).


(sec_stats_notes_f)=

### Patterson's f statistics

{meth}`~TreeSequence.f4`, {meth}`~TreeSequence.f3`, and {meth}`~TreeSequence.f2`
are the `f` statistics (also called `F` statistics) introduced in
[Reich et al (2009)](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC2842210/).
See the documentation (link below) for the definition,
and [Peter (2016)](https://www.genetics.org/content/202/4/1485) for readable
discussion of their use.


(sec_stats_notes_y)=

### Y statistics

{meth}`~TreeSequence.Y3` and {meth}`~TreeSequence.Y2` are the `Y` statistics introduced
by [Ashander et al (2018)](https://www.biorxiv.org/content/10.1101/354530v1)
as a three-sample intermediate between diversity/divergence (which are
pairwise) and Patterson's f statistics (which are four-way).


(sec_stats_notes_derived)=

### Derived statistics

Most statistics have the property that `mode="branch"` and
`mode="site"` are "dual" in the sense that they are equal, on average, under
a high neutral mutation rate. {meth}`~TreeSequence.Fst` and {meth}`~TreeSequence.Tajimas_D`
do not have this property (since both are ratios of statistics that do have this property).


(sec_stats_notes_gnn)=

### Genealogical nearest neighbours

The {meth}`~TreeSequence.genealogical_nearest_neighbours` statistic is not based on branch
lengths, but on topologies. therefore it currently has a slightly different interface to
the other single site statistics. This may be revised in the future.


--- ../../tskit/docs/introduction.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

```{currentmodule} tskit
```

(sec_introduction)=

# Introduction

This is the documentation for `tskit`, the tree sequence toolkit. Succinct tree sequences
provide a highly efficient way of storing a set of related DNA sequences by encoding
their ancestral history as a set of correlated trees along the genome.  The evolutionary
history of genetic sequences is often technically referred to as an Ancestral
Recombination Graph (ARG); succinct tree sequences are fully compatible with this
formulation, and tskit is a therefore a powerful platform for processing ARGs.

The tree sequence format is output by a number of external software libraries
and programs (such as [msprime](https://github.com/tskit-dev/msprime), 
[SLiM](https://github.com/MesserLab/SLiM), 
[fwdpp](http://molpopgen.github.io/fwdpp/), and 
[tsinfer](https://tsinfer.readthedocs.io/en/latest/)) that either simulate or
infer the evolutionary history of genetic sequences. This library provides the
underlying functionality that such software uses to load, examine, and
manipulate tree sequences, including efficient methods for calculating
{ref}`genetic statistics<sec_stats>`.

For a gentle introduction, you might like to read "{ref}`tutorials:sec_what_is`"
on our {ref}`tutorials site<tutorials:sec_intro>`. There you can also find further
tutorial material to introduce you to the key concepts behind succinct tree sequences.

:::{note}
This documentation is under active development and may be incomplete
in some areas. If you would like to help improve it, please open an issue or
pull request at [GitHub](https://github.com/tskit-dev/tskit).
:::


--- ../../tskit/docs/args.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

```{currentmodule} tskit
```

(sec_args)=

# ARGs as tree sequences

At its heart, a `tskit` {ref}`tree sequence<sec_what_is>` consists of a list of
{ref}`sec_terminology_nodes`, and a list of {ref}`sec_terminology_edges` that connect
parent to child nodes. Therefore a succinct tree sequence is equivalent to a
[directed graph](https://en.wikipedia.org/wiki/Directed_graph),
which is additionally annotated with genomic positions such that at each
position, a path through the edges exists which defines a tree. This graph
interpretation of a tree sequence maps very closely to the concept of
an "ancestral recombination graph" (or ARG). See
[this preprint](https://www.biorxiv.org/content/10.1101/2023.11.03.565466v2) for further details.

## Full ARGs

:::{margin}
An original, narrower definition, which we do not use here, restricts the term ARG to
the neutral coalescent process with simple crossover-recombination, and the
graph structure defined by that process, see e.g.
[Griffiths & Marjoram (1997)](https://research.monash.edu/en/publications/an-ancestral-recombination-graph)
:::

The term "ARG" is [often used](https://doi.org/10.1086%2F508901) to refer to
a structure consisting of nodes and edges that describe the genetic genealogy of a set
of sampled chromosomes which have evolved via a process of inheritance combined
with recombination. We use the term "full ARG" for a commonly-described type of
ARG that contains not just nodes that involve coalescence of ancestral material,
but also additional non-coalescent nodes. These nodes correspond to
recombination events, and common ancestor events that are not associated with
coalescence in any of the local trees. Full ARGs can be stored and analysed in
[tskit](https://tskit.dev) like any other tree sequence. A full ARG can be generated using
{func}`msprime:msprime.sim_ancestry` by specifying `coalescing_segments_only=False` along with
`additional_nodes = msprime.NodeType.COMMON_ANCESTOR | msprime.NodeType.RECOMBINANT`
(or the equivalent `record_full_arg=True`) as described
{ref}`in the msprime docs<msprime:sec_ancestry_full_arg>`:

```{code-cell}
import msprime

parameters = {
    "samples": 3, # Three diploid individuals == six sample genomes
    "sequence_length": 1e4,
    "recombination_rate": 1e-7,
    "population_size": 1e3,
    "random_seed": 333,
}

ts_arg = msprime.sim_ancestry(
    **parameters,
    discrete_genome=False,  # the strict Hudson ARG needs unique crossover positions (i.e. a continuous genome)
    coalescing_segments_only=False,   # setting record_full_arg=True is equivalent to these last 2 parameters
    additional_nodes=msprime.NodeType.COMMON_ANCESTOR | msprime.NodeType.RECOMBINANT,
)

print('Simulated a "full ARG" under the Hudson model:')
print(
    f" ARG stored in a tree sequence with {ts_arg.num_nodes} nodes and"
    f" {ts_arg.num_edges} edges (creating {ts_arg.num_trees} local trees)"
)
```

Like any tree sequence, we can also add mutations to the ARG to generate genetic
variation:


```{code-cell}
import numpy as np
mu = 1e-7
ts_arg = msprime.sim_mutations(ts_arg, rate=mu, random_seed=888)
print("     Sample node:  " + "   ".join(str(u) for u in ts_arg.samples()))
for v in ts_arg.variants():
    print(f"Variable site {v.site.id}:", np.array(v.alleles)[v.genotypes])
```

### Visualization

:::{margin}
In an `msprime` full ARG, recombinations are recorded in a specific way, by storing
the parental genomes of a gamete. This means that
*two* `tskit` nodes are created for each recombination, capturing transmission
to the left vs right of the crossover breakpoint. These two nodes, which exist
at the same timepoint, are identified by the
{data}`~msprime:msprime.msprime.NODE_IS_RE_EVENT` and
are displayed as a single point in the "ortho" viz below, labelled with
two node IDs separated by a slash.
:::

The normal {ref}`sec_tskit_viz` of a tree sequence is as a set of local
trees. However, all tree sequences can also be
{ref}`plotted as graphs<sec_tskit_viz_other_graph>`. In particular, the Hudson "full ARG"
model guarantees that the graph consists of nodes which mark a split into two child lineages
("common ancestor" nodes) or nodes which mark a split into two parent lineages
("recombination" nodes). Such ARGs can be visualized with edges drawn as horizontal and vertical
lines (the "ortho" style in the
[tskit_arg_visualiser](https://github.com/kitchensjn/tskit_arg_visualizer) software):

```{code-cell} ipython3
:"tags": ["hide-input"]
%%javascript
require.config({paths: {d3: 'https://d3js.org/d3.v7.min'}});
require(["d3"], function(d3) {window.d3 = d3;});
```

```{code-cell} ipython3
import tskit_arg_visualizer
d3arg = tskit_arg_visualizer.D3ARG.from_ts(ts=ts_arg)
w, h = 450, 300  # width and height
d3arg.draw(w, h, edge_type="ortho", sample_order=[0, 2, 1, 3, 5, 4])
```

### Local trees and arity

Below is a plot of the equivalent local trees in the ARG above,
colouring recombination nodes in red and common
ancestor nodes (unlabelled) in blue.

```{code-cell}
:"tags": ["hide-input"]
# Plot the recombination nodes in red, with a horizontal line at the time of occurrence,
# and only label nodes that are samples or recombination nodes.
samples = set(ts_arg.samples())
re_nodes = set(nd.id for nd in ts_arg.nodes() if nd.flags & msprime.NODE_IS_RE_EVENT)
ca_nodes = set(np.arange(ts_arg.num_nodes)) - re_nodes - samples
re_times = [int(nd.time) for nd in ts_arg.nodes() if nd.flags & msprime.NODE_IS_RE_EVENT]
style = ".y-axis .grid {stroke: #ff000033} .mut .sym {stroke: goldenrod}"
for u in re_nodes:
    style += f".n{u} > .sym {{fill: red}}"
for u in ca_nodes:
    style += f".n{u} > .sym {{fill: blue}}"
ts_arg.draw_svg(
    size=(600, 300),
    y_axis=True,
    y_ticks=re_times,
    y_gridlines=True,
    style=style,
    mutation_labels={},
    node_labels={u: u for u in samples | re_nodes}
)
```

The number of children descending from a node in a local tree can be termed the
"local arity" of that node. It is clear from the plot above that red nodes always
have a local arity of 1, and blue nodes sometimes do. This may seem an unusual
state of affairs: tree representations often focus on branch-points, and ignore nodes
with a single child. Indeed, it is possible to [simplify](sec_args_simplification) the
ARG above, resulting in a graph whose local trees only contain branch points or tips
(i.e. local arity is never 1). Such a graph is [more compact](sec_args_disadvantages)
than the full ARG, but it omits some information about the timings and
topological operations associated with recombination
events and some common ancestor events. This information, as captured by the local
unary nodes, is useful for

1. Retaining precise information about the time and lineages involved in recombination.
   This is required e.g. to ensure we can always work out the tree editing (or
   {ref}`subtree-prune-and-regraft (SPR)<sec_concepts_sprs>`) moves
   required to change one local tree into another as you move along the genome.

2. Calculating the likelihood of an full ARG under a specific model of evolution
   (most commonly, the neutral coalescent with recombination, or CwR, as modelled e.g. by
   [Hudson (1983)](https://doi.org/10.1016/0040-5809(83)90013-8))

(sec_args_sprs)=
#### SPRs and recombination

The location of the recombination nodes in the trees above imply that the
*recombination events* happened ~588 and ~59 generations ago. The older one,
at 588 generations, involved node
13 (to the left of position 2601.01) and node 14 (to the right). As well as narrowing
down the recombination event to a specific point in time, the position of these two
nodes tells us that the SPR to convert the first into the second tree
involves pruning the branch above samples 1, 3, 4, and 5 and regrafting it onto the
branch above samples 0 and 2, rather than the other way around. Note that this
particular recombination does not change the *topology* of the tree, but simply the
branch lengths.

The recombination event 59 generations ago involved nodes 7 and 8, with the crossover
ocurring at position 6516.94. The SPR operation which converts the middle tree into the
last one involves pruning the branch above sample node 5 and regrafting it onto the
branch above the common ancestor of 1 and 3. In this case, the recombination has led to
a change in topology, such that the closest relative of 5 is node 4 from positions 0
to 6516.94, but 1 and 3 from positions 6516.94 to 10,000.

(sec_args_likelihoods)=

### Calculating likelihoods

Because our ARG above was generated under the standard Hudson model (e.g. neutral
evolution in a large population with unique recombination breakpoints along a continuous
genome), we can calculate its likelihood under that model, for a given recombination
rate and population size, using the {func}`msprime:msprime.log_arg_likelihood` method:

```{code-cell}
print(
    "Log likelihood of the genealogy under the Hudson model:",
    msprime.log_arg_likelihood(
        ts_arg,
        recombination_rate=parameters["recombination_rate"],
        Ne=parameters["population_size"]
    )
)
```

:::{note}
This likelihood calculation is tied to the specific `tskit` representation of the ARG that
is output by the `msprime` simulator. In particular, it expects each recombination event to
correspond to two recombination nodes, which allows so-called "diamond" events to be
represented, in which both parents at a recombination event trace directly back to the
same common ancestor.
:::

(sec_args_simplification)=

## Simplification

If we fully {ref}`simplify<sec_simplification>` the tree above, all remaining nodes
(apart from the samples) will have a local arity greater than one.
This loses information about the timings of recombination and non-coalescent
common ancestry, but it still keeps the local tree structure intact:

```{code-cell}
ts = ts_arg.simplify()
ts.draw_svg(
    size=(600, 300),
    y_axis=True,
    node_labels={u: u for u in ts.samples()},
    mutation_labels={},
    style=".mut .sym {stroke: goldenrod}",
    y_ticks=[t*500 for t in range(4)]
)
```

Note that all recombination nodes have been lost from this graph; the effects of recombination
are instead reflected in more recent coalescent nodes that are recombinant decendants of the
original recombination nodes. This results in graph nodes which simultanousely have
multiple children and multiple parents. Graphs with such nodes are unsuited to
the "ortho" style of
graph plotting. Instead, we can plot the tree sequence graph using diagonal lines:
    
```{code-cell}
d3graph = tskit_arg_visualizer.D3ARG.from_ts(ts=ts)
d3graph.draw(w, h, sample_order=[0, 2, 1, 3, 5, 4])
```

These "simplified" graphs are what are produced as the default `msprime` output. The
exact SPR moves from tree to tree may no longer be obvious, and the ARG likelihood
cannot be calculated from a tree sequence of this form.

Note that we can still calculate the *mutation likelihood*
(i.e. the likelihood of the observed pattern of mutations, given the genealogy) because
the topology and branch lengths of the local trees remain unchanged after simplification:

```{code-cell}
print("Log likelihood of mutations given the genealogy:")
print(' "full" ARG:',  msprime.log_mutation_likelihood(ts_arg, mutation_rate=mu))
print(" simplified:", msprime.log_mutation_likelihood(ts, mutation_rate=mu))
```

(sec_args_disadvantages)=
## Disadvantages of Full ARGs

There are two main reasons why you might *not* want to store a full ARG, but instead use
a simplified version. Firstly if you are inferring ARGs from real data (rather than simulating
them), it may not be possible or even desirable to infer recombination events. Often
there are many possible recombination events which are compatible with a given
set of recombined genome sequences.

Secondly, even if you *are* simulating genetic ancestry, storing full ARG requires many extra nodes.
In fact, as the sequence length increases, the non-coalescent nodes come
to dominate the tree sequence.
We can calculate the percentage of non-coalescent nodes by comparing a full ARG with
its simplified version:

```{code-cell}
large_sim_parameters = parameters.copy()
large_sim_parameters["sequence_length"] *= 1000
large_ts_arg = msprime.sim_ancestry(
    **large_sim_parameters,
    discrete_genome=False,  # not technically needed, as we aren't calculating likelihoods
    coalescing_segments_only=False,
    additional_nodes=msprime.NodeType.COMMON_ANCESTOR | msprime.NodeType.RECOMBINANT,
)
large_ts = large_ts_arg.simplify()

print(
    "Non-coalescent nodes take up "
    f"{(1-large_ts.num_nodes/large_ts_arg.num_nodes) * 100:0.2f}% "
    f"of this {large_ts.sequence_length/1e6:g} megabase {large_ts.num_samples}-tip ARG"
)
```

This is one of the primary reasons that nodes which are never associated with coalescent
events are excluded by default in simulation software such as
[msprime](https://tskit.dev/software/msprime.html) and
[SLiM](https://tskit.dev/software/SLiM.html).

:::{note}
As well as ancestrally relevant nodes, the original (mathematical) ARG formulation by
[Griffiths (1991)](https://www.jstor.org/stable/4355649) includes recombination
nodes that are not ancestral to the samples. This leads to a graph with an
vastly larger number of nodes than even the ARGs simulated here, and using such
structures for simulation or inference is therefore infeasible.
:::

## ARG formats and `tskit`

It is worth noting a subtle and somewhat philosophical
difference between some classical ARG formulations, and the ARG formulation
used in `tskit`. Classically, nodes in an ARG are taken to represent _events_
(specifically, "common ancestor", "recombination", and "sampling" events),
and genomic regions of inheritance are encoded by storing a specific breakpoint location on
each recombination node. In contrast, {ref}`nodes<tskit:sec_data_model_definitions_node>` in a `tskit`
ARG correspond to _genomes_. More crucially, inherited regions are defined by intervals
stored on *edges* (via the {attr}`~Edge.left` and  {attr}`~Edge.right` properties),
rather than on nodes. Here, for example, is the edge table from our ARG:

```{code-cell}
ts_arg.tables.edges
```

Technically therefore, ARGs stored by `tskit` are edge-annotated
"genome ARGs", or [gARGs](https://www.biorxiv.org/content/10.1101/2023.11.03.565466v1).
This flexible format can describe both simulated genetic ancestries (including
those involving {ref}`msprime:sec_ancestry_gene_conversion`), and e.g. real-world
genetic ancestries, such as
[inferrred recombining viral genealogies](https://www.biorxiv.org/content/10.1101/2023.06.08.544212v1).
The focus on genomes rather than events is also what makes
simplification possible, and means `tskit` can encode ancestry without having
to pin down exactly when specific ancestral events took place.


## Working with ARGs in `tskit`

All tree sequences, including, but not limited to full ARGs, can be treated as
directed (acyclic) graphs. Although many tree sequence operations operate from left to
right along the genome, some are more naturally though of as passing from node
to node via the edges, regardless of the genomic position of the edge. This section
describes some of these fundamental graph operations.

### Graph traversal

The standard edge iterator, {meth}`TreeSequence.edge_diffs()`, goes from left to
right along the genome, matching the {meth}`TreeSequence.trees()` iterator. Although
this will visit all the edges in the graph, these will *not* necessarily be grouped
by the node ID either of the edge parent or the edge child. To do this, an alternative
traversal (from top-to-bottom or bottom-to-top of the tree sequence) is required.

To traverse the graph by node ID, the {meth}`TreeSequence.nodes()` iterator can be
used. In particular, because parents are required to be strictly older than their
children, iterating through nodes using `order="timeasc"` will ensure that children
are always visited before their parents (similar to a breadth-first or "level order"
search). However, using {meth}`TreeSequence.nodes()` is inefficient if you also
want to access the *edges* associated with each node. 

The examples below show how to efficiently visit the all the edges in a
tree sequence, grouped by the nodes to which they are connected, while
also ensuring that children are visited before parents (or vice versa).

#### Traversing parent nodes

The most efficient graph traversal method visits all the parent nodes in the
tree sequence, grouping edges for which that node is a parent. This is simple
because edges in a tree sequence are {ref}`ordered<tskit:sec_edge_requirements>`
firstly by the time of the parent node, then by node ID.

```{code-cell}
import itertools
import operator
def edges_by_parent_timeasc(ts):
    return itertools.groupby(ts.edges(), operator.attrgetter("parent"))

for parent_id, edges in edges_by_parent_timeasc(ts):
    t = ts.node(parent_id).time
    children = {e.child for e in edges}
    print(f"Node {parent_id} at time {t} has these child node IDs: {children}")
```

This visits children before parents. To visit parents before children, you can
simply use ``reversed(ts.edges())`` rather than ``ts.edges()`` within the
``groupby`` function. Note that terminal nodes (i.e. which are either isolated
or leaf nodes in all local trees) are not visited by this function: therefore
the method above will omit all the terminal sample nodes.

#### Traversing child nodes

Sometimes you may wish to iterate over all the edges for which a node is a
child (rather than a parent). This can be done by sorting the edges by
child node time (then child node id). This is a little slower, but can be
done relatively efficiently as follows:

```{code-cell}
import itertools
import operator
import numpy as np
def edges_by_child_timeasc(ts):
    # edges sorted by child node time, then child node id using np.lexsort
    it = (ts.edge(u) for u in np.lexsort((ts.edges_child, ts.nodes_time[ts.edges_child])))
    return itertools.groupby(it, operator.attrgetter("child"))

for child_id, edges in edges_by_child_timeasc(ts):
    t = ts.node(child_id).time
    parents = {e.parent for e in edges}
    print(f"Node {child_id} at time {t} has these parent node IDs: {parents}")
```

To visit parents before children, the ``lexsort`` can take the negated
``-ts.nodes_time`` rather than simply using ``ts.nodes_time``. Note that nodes
which are never children of an edge are not visited by this algorithm. Such
nodes are either {ref}`isolated<sec_data_model_tree_isolated_nodes>` or a
{ref}`root<sec_data_model_tree_roots>` in each local tree.


(sec_args_other_analysis)=

### Other graph analysis

For graph-theory based analysis, it can be helpful to convert a tree sequence
to a [networkx](https://networkx.org/) graph. This can be done using the
following code:

```{code-cell} ipython3
import tskit
import networkx as nx
import pandas as pd

def to_networkx_graph(ts, interval_lists=False):
    """
    Make an nx graph from a tree sequence. If `intervals_lists` is True, then
    each graph edge will have an ``intervals`` attribute containing a *list*
    of tskit.Intervals per parent/child combination. Otherwise each graph edge
    will correspond to a tskit edge, with a ``left`` and ``right`` attribute.
    """
    D = dict(source=ts.edges_parent, target=ts.edges_child, left=ts.edges_left, right=ts.edges_right)
    G = nx.from_pandas_edgelist(pd.DataFrame(D), edge_attr=True, create_using=nx.MultiDiGraph)
    if interval_lists:
        GG = nx.DiGraph()  # Mave a new graph with one edge that can contai
        for parent, children in G.adjacency():
            for child, edict in children.items():
                ilist = [tskit.Interval(v['left'], v['right']) for v in edict.values()]
                GG.add_edge(parent, child, intervals=ilist)
        G = GG
    nx.set_node_attributes(G, {n.id: {'flags':n.flags, 'time': n.time} for n in ts.nodes()})
    return G

arg = to_networkx_graph(ts_arg)
```

It is then possible to use the full range of networkx functions to analyse the graph:

```{code-cell} ipython3
assert nx.is_directed_acyclic_graph(arg)  # All ARGs should be DAGs
print("All descendants of node 10 are", nx.descendants(arg, 10))
```

Networkx also has some built-in drawing functions: below is one of the simplest ones
(for other possibilities, see the {ref}`sec_tskit_viz` tutorial).

```{code-cell} ipython3
for layer, nodes in enumerate(nx.topological_generations(arg.reverse())):
    for node in nodes:
        arg.nodes[node]["layer"] = layer
pos = nx.multipartite_layout(arg, subset_key="layer", align='horizontal')

nx.draw_networkx(arg, pos=pos)
```

## Other software

:::{todo}
Show how ARGweaver output can be converted to tskit form.
:::

:::{todo}
Show how KwARG output can be converted to tskit form.
:::

:::{todo}
Implement conversion between the _msprime_ 2 RE node version and the more conventional 1 RE node version. See https://github.com/tskit-dev/msprime/issues/1942 for extensive discussion on the advantages / disadvantages of using 2 nodes vs 1 node-with-metadata.
:::



--- ../../tskit/docs/ibd.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

```{currentmodule} tskit
```


(sec_identity)=

# Identity by descent

The {meth}`.TreeSequence.ibd_segments` method allows us to compute
segments of identity by descent.

:::{note}
This documentation page is preliminary
:::

:::{todo}
Relate the concept of identity by descent to the MRCA spans in the tree sequence.
:::

## Examples

Let's take a simple tree sequence to illustrate the {meth}`.TreeSequence.ibd_segments`
method and associated {ref}`sec_python_api_reference_identity`:

```{code-cell}
:tags: [hide-input]

import tskit
import io
from IPython.display import SVG

nodes = io.StringIO(
    """\
    id      is_sample   time
    0       1           0
    1       1           0
    2       1           0
    3       0           1
    4       0           2
    5       0           3
    """
)
edges = io.StringIO(
    """\
    left    right   parent  child
    2     10     3       0
    2     10     3       2
    0     10     4       1
    0     2      4       2
    2     10     4       3
    0     2      5       0
    0     2      5       4
    """
)
ts = tskit.load_text(nodes=nodes, edges=edges, strict=False)

SVG(ts.draw_svg())
```

### Definition

A pair of nodes ``(u, v)`` has an IBD segment with a left and right
coordinate ``[left, right)`` and ancestral node ``a`` iff the most
recent common ancestor of the segment ``[left, right)`` in nodes ``u``
and ``v`` is ``a``, and the segment has been inherited along the same
genealogical path (ie. it has not been broken by recombination). The
segments returned are the longest possible ones.

Consider the IBD segments that we get from our example tree sequence:

```{code-cell}
segments = ts.ibd_segments(store_segments=True)
for pair, segment_list in segments.items():
    print(pair, list(segment_list))
```

Each of the sample pairs (0, 1), (0, 2) and (1, 2) is associated with
two IBD segments, representing the different paths from these sample
pairs to their common ancestor. Note in particular that (1, 2) has
**two** IBD segments rather than one: even though the MRCA is
4 in both cases, the paths from the samples to the MRCA are different
in the left and right trees.


### Data structures

The result of calling {meth}`.TreeSequence.ibd_segments` is an
{class}`.IdentitySegments` class:

```{code-cell}
segments = ts.ibd_segments()
print(segments)
```

By default this class only stores the high-level summaries of the
IBD segments discovered. As we can see in this example, we have a
total of six segments and
the total span (i.e., the sum lengths of the genomic intervals spanned
by IBD segments) is 30.

If required, we can get more detailed information about particular
segment pairs and the actual segments using the ``store_pairs``
and ``store_segments`` arguments.

:::{warning}
Only use the ``store_pairs`` and ``store_segments`` arguments if you
really need this information! The number of IBD segments can be
very large and storing them all requires a lot of memory. It is
also much faster to just compute the overall summaries, without
needing to store the actual lists.
:::


```{code-cell}
segments = ts.ibd_segments(store_pairs=True)
for pair, value in segments.items():
    print(pair, "::", value)
```

Now we can see the more detailed breakdown of how the identity segments
are distributed among the sample pairs. The {class}`.IdentitySegments`
class behaves like a dictionary, such that ``segments[(a, b)]`` will return
the {class}`.IdentitySegmentList` instance for that pair of samples:

```{code-cell}
seglist = segments[(0, 1)]
print(seglist)
```

If we want to access the detailed information about the actual
identity segments, we must use the ``store_segments`` argument:

```{code-cell}
segments = ts.ibd_segments(store_pairs=True, store_segments=True)
segments[(0, 1)]
```

The {class}`.IdentitySegmentList` behaves like a Python list,
where each element is an instance of {class}`.IdentitySegment`.

:::{warning}
The order of segments in an {class}`.IdentitySegmentList`
is arbitrary, and may change in future versions.
:::


```{eval-rst}
.. todo:: More examples using the other bits of the IdentitySegments
    API here
```

### Controlling the sample sets

By default we get the IBD segments between all pairs of
{ref}`sample<sec_data_model_definitions_sample>` nodes.

#### IBD within a sample set
We can reduce this to pairs within a specific set using the
``within`` argument:


```{eval-rst}
.. todo:: More detail and better examples here.
```

```{code-cell}
segments = ts.ibd_segments(within=[0, 2], store_pairs=True)
print(list(segments.keys()))
```

#### IBD between sample sets

We can also compute IBD **between** sample sets:

```{code-cell}
segments = ts.ibd_segments(between=[[0,1], [2]], store_pairs=True)
print(list(segments.keys()))
```

:::{seealso}
See the {meth}`.TreeSequence.ibd_segments` documentation for
more details.
:::

### Constraints on the segments

The ``max_time`` and ``min_span`` arguments allow us to constrain the
segments that we consider.

```{eval-rst}
.. todo:: Add examples for these arguments.
```


--- ../../tskit/docs/quickstart.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

:::{currentmodule} tskit
:::

```{code-cell} ipython3
:tags: [remove-cell]
import msprime

def basic_sim():
    ts = msprime.sim_ancestry(
        3,
        population_size=1000,
        model="dtwf",
        sequence_length=1e4,
        recombination_rate=1e-7,
        random_seed=665)
    ts = msprime.sim_mutations(ts, rate=2e-7, random_seed=123)
    ts.dump("data/basic_tree_seq.trees")
    
def create_notebook_data():
    basic_sim()

# create_notebook_data()  # uncomment to recreate the tree seqs used in this notebook
```

# Quickstart

Our {ref}`tutorials site<tutorials:sec_intro>` has a more extensive tutorial on
{ref}`sec_tskit_getting_started`. Below we just give a quick flavour of the
{ref}`sec_python_api` (note that
APIs in {ref}`C <sec_c_api>` and Rust exist, and it is also possible to
{ref}`interface to the Python library in R <tutorials:sec_tskit_r>`).

## Basic properties

Any tree sequence, such as one generated by {ref}`msprime <msprime:sec_intro>`, can be
loaded, and a summary table printed. This example uses a small tree sequence, but the
`tskit` library scales effectively to ones encoding millions of genomes and variable
sites. 

```{code-cell}
import tskit

ts = tskit.load("data/basic_tree_seq.trees")  # Or generate using e.g. msprime.sim_ancestry()
ts  # In a Jupyter notebook this displays a summary table. Otherwise use print(ts)
```

## Individual trees

You can get e.g. the first tree in the tree sequence and analyse it.

```{code-cell}
first_tree = ts.first()
print("Total branch length in first tree is", first_tree.total_branch_length, ts.time_units)
print("The first of", ts.num_trees, "trees is plotted below")
first_tree.draw_svg(y_axis=True)  # plot the tree: only useful for small trees
```

## Extracting genetic data

A tree sequence provides an extremely compact way to
{ref}`store genetic variation data <tutorials:sec_what_is_dna_data>`. The trees allow
this data to be {meth}`decoded <Variant.decode>` at each site:

```{code-cell}
for variant in ts.variants():
    print(
        "Variable site", variant.site.id,
        "at genome position", variant.site.position,
        ":", [variant.alleles[g] for g in variant.genotypes],
    )
```

## Analysis

Tree sequences enable {ref}`efficient analysis <tutorials:sec_what_is_analysis>`
of genetic variation using a comprehensive range of built-in {ref}`sec_stats`:

```{code-cell}
genetic_diversity = ts.diversity()
print("Av. genetic diversity across the genome is", genetic_diversity)

branch_diversity = ts.diversity(mode="branch")
print("Av. genealogical dist. between pairs of tips is", branch_diversity,  ts.time_units)
```

## Plotting the whole tree sequence

This can give you a visual feel for small genealogies:

```{code-cell}
ts.draw_svg(
    size=(800, 300),
    y_axis=True,
    mutation_labels={m.id: m.derived_state for m in ts.mutations()},
)
```

## Underlying data structures

The data that defines a tree sequence is stored in a set of tables. These tables
can be viewed, and copies of the tables can be edited to create a new tree sequence.

```{code-cell}
# The sites table is one of several tables that underlie a tree sequence
ts.tables.sites
```

The rest of this documentation gives a comprehensive description of the entire `tskit`
library, including {ref}`descriptions and definitions <sec_table_definitions>` of all
the tables.



--- ../../tskit/docs/analysing_trees.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

```{currentmodule} tskit
```

```{code-cell} ipython3
:tags: [remove-cell]
import io
import pickle

import msprime
import tskit

def tree_traversals():
    nodes = """\
    id      is_sample   time
    0       1           0
    1       1           0
    2       1           0
    3       1           0
    4       1           0
    5       0           1
    6       0           2
    7       0           3
    """
    edges = """\
    left    right   parent  child
    0       1       5       0,1,2
    0       1       6       3,4
    0       1       7       5,6
    """
    # NB same tree as used above, and we're using the same diagram.
    ts = tskit.load_text(
        nodes=io.StringIO(nodes), edges=io.StringIO(edges), strict=False
    )
    ts.dump("data/tree_traversals.trees")

def different_time_samples():
    samples = [
        msprime.Sample(0, 0),
        msprime.Sample(0, 1),
        msprime.Sample(0, 20),
    ]
    ts = msprime.simulate(
        Ne=1e6,
        samples=samples,
        demographic_events=[
            msprime.PopulationParametersChange(time=10, growth_rate=2, population_id=0),
        ],
        random_seed=42,
    )
    ts.dump("data/different_time_samples.trees")

def parsimony_simple():
    ts = msprime.sim_ancestry(3, random_seed=42)
    ts.dump("data/parsimony_simple.trees")

def parsimony_map():
    # pick a seed that gives the tips in the right order
    ts = msprime.sim_ancestry(3, sequence_length=100, random_seed=38)
    ts.dump("data/parsimony_map.trees")
    ts = msprime.sim_mutations(
        ts, rate=0.01, random_seed=123, discrete_genome=False)
    data = [
        {'pos':v.site.position, 'alleles': v.alleles, 'genotypes':v.genotypes}
        for v in ts.variants()]
    with open("data/parsimony_map.pickle", "wb") as f:
        pickle.dump(data, f)


def create_notebook_data():
    tree_traversals()
    different_time_samples()
    parsimony_simple()
    parsimony_map()

# create_notebook_data()  # uncomment to recreate the tree seqs used in this notebook
```

(sec_analysing_trees)=
# Analysing trees

There are a number of different ways we might want to analyse a single  {class}`Tree`.
Most involve some sort of traversal over the nodes, mutations, or branches in the tree.
{program}`tskit` provides various way of traversing through a tree, and also some
built in phylogenetic algorithms such as {meth}`Tree.map_mutations` which efficently
places mutations ("characters" in phylogenetic terminology) on a given tree.


(sec_analysing_trees_traversals)=

## Tree traversals

Given a single {class}`Tree`, traversals in various orders are possible using the
{meth}`~Tree.nodes` iterator. Take the following tree:


```{code-cell} ipython3
import tskit

ts = tskit.load("data/tree_traversals.trees")
tree = ts.first()
tree.draw_svg()
```

We can visit the nodes in different orders by providing an `order` parameter to
the {meth}`Tree.nodes` iterator:

```{code-cell} ipython3
for order in ["preorder", "inorder", "postorder"]:
    print(f"{order}:\t", list(tree.nodes(order=order)))
```

Much of the time, the specific ordering of the nodes is not important
and we can leave it out (defaulting to preorder traversal). For example,
here we compute the total branch length of a tree:

```{code-cell} ipython3
total_branch_length = sum(tree.branch_length(u) for u in tree.nodes())
print(f"Total branch length: {total_branch_length}")
```

Note that this is also available as the {attr}`Tree.total_branch_length` attribute.

### Traversing upwards

For many applications it is useful to be able to traverse upwards from the
leaves. We can do this using the {meth}`Tree.parent` method, which
returns the parent of a node. For example, we can traverse upwards from
each of the samples in the tree:

```{code-cell} ipython3
for u in tree.samples():
    path = []
    v = u
    while v != tskit.NULL:
        path.append(v)
        v = tree.parent(v)
    print(u, "->", path)
```

### Traversals with information

Sometimes we will need to traverse down the tree while maintaining
some information about the nodes that are above it. While this
can be done using recursive algorithms, it is often more convenient
and efficient to use an iterative approach. Here, for example,
we define an iterator that yields all nodes in preorder along with
their path length to root:

```{code-cell} ipython3
def preorder_dist(tree):
    for root in tree.roots:
        stack = [(root, 0)]
        while len(stack) > 0:
            u, distance = stack.pop()
            yield u, distance
            for v in tree.children(u):
                stack.append((v, distance + 1))

print(list(preorder_dist(tree)))
```


(sec_tutorial_networkx)=

## Networkx

Traversals and other network analysis can also be performed using the sizeable
[networkx](https://networkx.github.io/documentation/stable/index.html)
library. This can be achieved by calling {meth}`Tree.as_dict_of_dicts` to
convert a {class}`Tree` instance to a format that can be imported by networkx to
create a graph:

```{code-cell} ipython3
import networkx as nx

g = nx.DiGraph(tree.as_dict_of_dicts())
print(sorted(g.edges))
```

### Traversing upwards in networkx

We can revisit the above examples and traverse upwards with
networkx using a depth-first search algorithm:

```{code-cell} ipython3
g = nx.DiGraph(tree.as_dict_of_dicts())
for u in tree.samples():
    path = [u] + [parent for parent, child, _ in
                  nx.edge_dfs(g, source=u, orientation="reverse")]
    print(u, "->", path)
```


### Calculating distances to the root

Similarly, we can yield the nodes of a tree along with their distance to the
root in pre-order in networkx as well. Running this on the example above gives us
the same result as before:

```{code-cell} ipython3
g = nx.DiGraph(tree.as_dict_of_dicts())
for root in tree.roots:
    print(sorted(list(nx.shortest_path_length(g, source=root).items())))
```


### Finding nearest neighbors

If some samples in a tree are not at time 0, then finding the nearest neighbor
of a sample is a bit more involved. Instead of writing our own traversal code
we can again draw on a networkx algorithm.
Let us start with an example tree with three samples that were sampled at
different time points:

```{code-cell} ipython3
ts = tskit.load("data/different_time_samples.trees")
tree = ts.first()
tree.draw_svg(y_axis=True, time_scale="rank")
```

The generation times for these nodes are as follows:

```{code-cell} ipython3
for u in tree.nodes():
    print(f"Node {u}: time {tree.time(u)}")
```

Note that samples 0 and 1 are about 35 generations apart from each other even though
they were sampled at almost the same time. This is why samples 0 and 1 are
closer to sample 2 than to each other.

For this nearest neighbor search we will be traversing up and down the tree,
so it is easier to treat the tree as an undirected graph:

```{code-cell} ipython3
g = nx.Graph(tree.as_dict_of_dicts())
```

When converting the tree to a networkx graph the edges are annotated with their
branch length:

```{code-cell} ipython3
for e in g.edges(data=True):
    print(e)
```

We can now use the "branch_length" field as a weight for a weighted shortest path
search:

```{code-cell} ipython3
import collections
import itertools

# a dictionary of dictionaries to represent our distance matrix
dist_dod = collections.defaultdict(dict)
for source, target in itertools.combinations(tree.samples(), 2):
    dist_dod[source][target] = nx.shortest_path_length(
        g, source=source, target=target, weight="branch_length"
    )
    dist_dod[target][source] = dist_dod[source][target]

# extract the nearest neighbor of nodes 0, 1, and 2
nearest_neighbor_of = [min(dist_dod[u], key=dist_dod[u].get) for u in range(3)]

print(dict(zip(range(3), nearest_neighbor_of)))
```


(sec_analysing_trees_parsimony)=

## Parsimony

Take a site on the following tree with three allelic states, where the
samples are coloured by the allele they possess, but where we don't know
the position of the mutations that caused this variation:

```{code-cell} ipython3
tree = tskit.load("data/parsimony_simple.trees").first()
alleles = ["red", "blue", "green"]
genotypes = [0, 0, 0, 0, 1, 2]
styles = [f".n{j} > .sym {{fill: {alleles[g]}}}" for j, g in enumerate(genotypes)]
tree.draw_svg(style="".join(styles))
```

The {meth}`Tree.map_mutations` method finds a parsimonious explanation for a
set of discrete character observations on the samples in a tree using classical
phylogenetic algorithms:

```{code-cell} ipython3
ancestral_state, mutations = tree.map_mutations(genotypes, alleles)
print("Ancestral state = ", ancestral_state)
for mut in mutations:
    print(f"Mutation: node = {mut.node} derived_state = {mut.derived_state}")
```

In this case, the algorithm has concluded, quite reasonably, that the most parsimonious
description of this state is that the ancestral state is red and there was
a mutation to blue and green over nodes 4 and 5.

### Building tables

Below we show how a set of tables can be updated using the
{ref}`Tables API<tskit:sec_tables_api>`, taking advantage of the
{meth}`Tree.map_mutations` method to identify parsimonious positions
for mutations on a tree. Here's the tree we'll use:

```{code-cell} ipython3
import pickle
ts = tskit.load("data/parsimony_map.trees")
ts.draw_svg(size=(500, 300), time_scale="rank")
```

Now we can modify the tables by adding mutations. To find the location of mutations,
we infer them from some observed data (some site positions with associated genotypes
and allelic states, in the conventional {class}`tskit encoding <Variant>`):


```{code-cell} ipython3
with open("data/parsimony_map.pickle", "rb") as file:
    data = pickle.load(file)  # Load saved variant data from a file
print("Variant data: each site has a position, allele list, and genotypes array:")
for i, v in enumerate(data):
    print(f"Site {i} (pos {v['pos']:7.4f}): alleles: {v['alleles']}, genotypes: {v['genotypes']}")
print()
tree = ts.first()  # there's only one tree anyway
tables = ts.dump_tables()
# Infer the sites and mutations from the variants.
for variant in data:
    ancestral_state, mutations = tree.map_mutations(variant["genotypes"], variant['alleles'])
    site_id = tables.sites.add_row(variant['pos'], ancestral_state=ancestral_state)
    info = f"Site {site_id}: parsimony sets ancestral state to {ancestral_state}"
    parent_offset = len(tables.mutations)
    for mut in mutations:
        parent = mut.parent
        if parent != tskit.NULL:
            parent += parent_offset
        mut_id = tables.mutations.add_row(
            site_id, node=mut.node, parent=parent, derived_state=mut.derived_state)
        info += f", and places mutation {mut_id} to {mut.derived_state} above node {mut.node}"
    print(info)
new_ts = tables.tree_sequence()
```

And here are the parsimoniously positioned mutations on the tree

```{code-cell} ipython3
mut_labels = {}  # An array of labels for the mutations
for mut in new_ts.mutations():  # Make pretty labels showing the change in state
    site = new_ts.site(mut.site)
    older_mut = mut.parent >= 0  # is there an older mutation at the same position?
    prev = new_ts.mutation(mut.parent).derived_state if older_mut else site.ancestral_state
    mut_labels[site.id] = f"{mut.id}: {prev}→{mut.derived_state}"

new_ts.draw_svg(size=(500, 300), mutation_labels=mut_labels, time_scale="rank")
```


### Parsimony and missing data

We can also take missing data
into account when finding a set of parsimonious state transitions. We do this by
specifying the special value {data}`tskit.MISSING_DATA` (-1) as the state, which is
treated by the algorithm as "could be anything".

For example, here we state that sample 0 is missing, indicated by the colour white:

```{code-cell} ipython3
tree = tskit.load("data/parsimony_simple.trees").first()
alleles = ["red", "blue", "green", "white"]
genotypes = [tskit.MISSING_DATA, 0, 0, 0, 1, 2]
styles = [f".n{j} > .sym {{fill: {alleles[g]}}}" for j, g in enumerate(genotypes)]
tree.draw_svg(style="".join(styles))
```

Now we run the {meth}`Tree.map_mutations` method, which applies the Hartigan parsimony
algorithm:

```{code-cell} ipython3
ancestral_state, mutations = tree.map_mutations(genotypes, alleles)
print("Ancestral state = ", ancestral_state)
for mut in mutations:
    print(f"Mutation: node = {mut.node} derived_state = {mut.derived_state}")
```

The algorithm decided, again, quite reasonably, that the most parsimonious explanation
for the input data is the same as before. Thus, if we used this information to fill
out mutation table as above, we would impute the missing value for 0 as red.

The output of the algorithm can be a little surprising at times. Consider this example::

```{code-cell} ipython3
tree = msprime.simulate(6, random_seed=42).first()
alleles = ["red", "blue", "white"]
genotypes = [1, -1, 0, 0, 0, 0]
node_colours = {j: alleles[g] for j, g in enumerate(genotypes)}
ancestral_state, mutations = tree.map_mutations(genotypes, alleles)
print("Ancestral state = ", ancestral_state)
for mut in mutations:
    print(f"Mutation: node = {mut.node} derived_state = {mut.derived_state}")
tree.draw(node_colours=node_colours)
```


Note that this is putting a mutation to blue over node 6, **not** node 0 as
we might expect. Thus, we impute here that node 1 is blue. It is important
to remember that the algorithm is minimising the number of state transitions;
this may not correspond always to what we might consider the most parsimonious
explanation.


(sec_trees_numba)=
## Fast tree-based algorithms using numba

:::{todo}
Add a few examples here of how to use numba to speed up tree based dynamic programming
algorithms. There are a good number of worked-up examples with timings in
[issue #63](https://github.com/tskit-dev/tutorials/issues/63)
:::



--- ../../tskit/docs/export.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

```{currentmodule} tskit
```

```{code-cell}
:tags: [hide-input]

from IPython.display import display
```

(sec_export)=

# Data export

(sec_export_vcf)=
## Variant Call Format

Tskit supports exporting data to the standard
[Variant Call Format](http://samtools.github.io/hts-specs/VCFv4.3.pdf)
via the `tskit vcf` {ref}`command line interface<sec_cli>` command
and the {meth}`TreeSequence.write_vcf` method in the {ref}`sec_python_api`.
Conversion is quite efficient, with tskit producing VCF data at several
hundred megabytes per second (for large files), which is usually as fast as
it can be written to storage or consumed by programs in a pipeline.

::::{tip}
If we have a tree sequence file the
{ref}`command line interface<sec_cli>` is often the most
convenient way to convert to VCF:

:::{code-block} bash
$ tskit vcf example.trees > example.vcf
:::

See the {ref}`sec_export_vcf_compression` section for information
on how to compress the VCF output.
::::

For tree sequences produced by recent versions of programs such as
``msprime``, ``SLiM``, ``fwdpy11`` or ``tsinfer``, VCF output will
"do the right thing" and no further arguments are needed.
For example, here we simulate 3 diploid individuals
with mutations using ``msprime``, and convert to VCF.

```{code-cell}
import sys
import msprime
ts = msprime.sim_ancestry(
    samples=3, ploidy=2, sequence_length=10, random_seed=2)
ts = msprime.sim_mutations(ts, rate=0.1, random_seed=2)
ts.write_vcf(sys.stdout)
```

In the output VCF we have 3 diploid samples
(see the {ref}`sec_export_vcf_terminology` section)
corresponding to samples specified in the ancestry simulation
with IDs ``tsk_0``, ``tsk_1`` and ``tsk_2``
(see the {ref}`sec_export_vcf_individual_names`
section for how to override these default labels).
We then have a line for every row
in the {ref}`site table<sec_site_table_definition>`, and
the data is derived directly from the {meth}`TreeSequence.variants`
method; e.g.

```{code-cell}
for var in ts.variants():
    print(var.site.position, var.site.id, var.alleles, var.genotypes, sep="\t")
```

We can see the ``POS`` value is equal to the site's position
(see the {ref}`sec_export_vcf_modifying_coordinates` for information
on how we deal with continuous coordinates), the ``ID`` value
is the site's ID, and the ``REF`` and ``ALT`` values
are derived from the variant's ``alleles``.

The ``GT`` values for the three diploid individuals are derived from the
variant's genotypes (see the {ref}`sec_export_vcf_terminology` section).
For this simulation, the diploid individuals correspond to
adjacent sample nodes in order, and we can see there is a direct
correspondence between the phased ``GT`` values and variant's genotypes.
See the {ref}`sec_export_vcf_constructing_gt` section for
more information on how this done in general and for options
to control the VCF sample and ``GT`` values.

::::{important}
In these examples we write the VCF data to ``sys.stdout`` so that we can see
the output. Usually, however, you'd write to a file:

:::{code-block}
with open("output.vcf", "w") as vcf_file:
    ts.write_vcf(vcf_file)
:::

:::{seealso}
See the {ref}`sec_export_vcf_compression` section for information
on how to compress the output or convert to BCF.
:::

::::

(sec_export_vcf_terminology)=

### Terminology

There are some mismatches between the terminology for tskit and VCF.
In VCF a "sample" is a multiploid individual, but in tskit a sample
refers to a single **node** (monoploid genome), and an individual
consists of one or more nodes (e.g., two nodes for a diploid).
Similarly, in VCF a "genotype" refers to the observed allelic state
for a sample **individual** at a particular site,
whereas in tskit a genotype is the observed allelic state
for a **node** (see {attr}`.Variant.genotypes`).

:::{seealso}
See the {ref}`sec_glossary` for more details on tskit's data model
 and terminology.
:::

(sec_export_vcf_compression)=

### Compressed output

The simplest way to compress the VCF output is to use the
`tskit vcf` {ref}`command line interface<sec_cli>`
and pipe the output to `bgzip`:

:::{code-block} bash
$ tskit vcf example.trees | bgzip -c > example.vcf.gz
:::
A general way to convert VCF data to various formats is to pipe the text
produced by ``tskit`` into ``bcftools`` using the command
line interface:

:::{code-block} bash
$ tskit vcf example.trees | bcftools view -O b > example.bcf
:::

If you need more control over the form of the output (or want to work
directly in Python), the following recipe has the same effect:

:::{code-block}

import os
import subprocess

read_fd, write_fd = os.pipe()
write_pipe = os.fdopen(write_fd, "w")
with open("output.bcf", "w") as bcf_file:
    proc = subprocess.Popen(
        ["bcftools", "view", "-O", "b"], stdin=read_fd, stdout=bcf_file
    )
    ts.write_vcf(write_pipe)
    write_pipe.close()
    os.close(read_fd)
    proc.wait()
    if proc.returncode != 0:
        raise RuntimeError("bcftools failed with status:", proc.returncode)
:::


The VCF output can also be compressed using the {mod}`gzip` Python module:

:::{code-block}

import gzip

with gzip.open("output.vcf.gz", "wt") as f:
    ts.write_vcf(f)
:::

However, this gzipped VCF won't be fully compatible with downstream tools
such as tabix, which usually require the VCF to use the specialised bgzip format.

(sec_export_vcf_masking_output)=

### Masking output

The {meth}`TreeSequence.write_vcf` method provides the
``site_mask`` and ``sample_mask`` arguments to
omit or mark parts of the output as missing.

```{code-cell}
ts = msprime.sim_ancestry(
    samples=3, ploidy=2, sequence_length=10, random_seed=2)
ts = msprime.sim_mutations(ts, rate=0.1, random_seed=2)
ts.tables.sites
```

The ``sample_mask`` argument provides a general way to mask out
parts of the output, which can be helpful when simulating missing
data. In this (contrived) example, we create a sample mask function
that marks one genotype missing in each variant in a regular
pattern:

:::{code-block}

def sample_mask(variant):
    sample_mask = np.zeros(ts.num_samples, dtype=bool)
    sample_mask[variant.site.id % ts.num_samples] = 1
    return sample_mask


ts.write_vcf(sys.stdout, sample_mask=sample_mask)
:::

(sec_export_vcf_constructing_gt)=

### Constructing GT values

The core elements of the tskit
{ref}`data model<sec_data_model>`
are {ref}`nodes<sec_node_table_definition>`,
{ref}`edges<sec_node_table_definition>`,
{ref}`sites<sec_site_table_definition>` and
{ref}`mutations<sec_mutation_table_definition>`.
These four tables allow us to completely describe the
genetic ancestry of a set of sampled monoploid
genomes and their genetic variation.
The {ref}`individual table<sec_individual_table_definition>`
defines a set of individual *organisms*, and it can
be used to define the inheritance relationships between
then (the pedigree). An individual may be associated
with one or more nodes, and these nodes may or
may not be samples (see the {ref}`sec_glossary`
for clarification of these terms).
Thus, there is some complexity in how the per-individual GT values
are generated, which we explain in this section.

#### Without individuals

We start with an example in which there are no individuals
defined (which was the default in msprime before version 1.0):

```{code-cell}
import tskit
tables = tskit.Tree.generate_balanced(4, span=10).tree_sequence.dump_tables()
tables.sites.add_row(3, ancestral_state="A")
tables.mutations.add_row(site=0, node=0, derived_state="T")
ts = tables.tree_sequence()
display(ts.draw_svg())
display(ts)
ts.write_vcf(sys.stdout)
```

Here we define a tree sequence consisting of a single tree, which
has a variant site at position 3 and a mutation over node 0.
There is no information about individuals in this tree sequence,
and so we assume that each of the nodes corresponds to a single
haploid individual.

Users of msprime simulations would often be interested in producing
VCFs for diploid organisms. Because of the assumptions made
by these simulations, this means arbitrarily combining the sample
nodes into pairs. This is what the ``ploidy`` option does:

```{code-cell}
ts.write_vcf(sys.stdout, ploidy=2)
```

Thus, the ``GT`` values for the (synthetic) diploid individual ``tsk_0``
are generated by combining nodes 0 and 1, and ``tsk_1``
by combining nodes 2 and 3.

:::{important}
Software packages modelling multiploid individuals are encouraged to
use the individual table to make their assumptions explicit. Recent
versions of simulators and inference methods should all do this,
and so the ``ploidy`` argument is really only intended to support
legacy code. It is therefore an error to supply a value for ``ploidy``
when individual information is present in a tree sequence.
:::

#### With individuals

Extending the example in the previous section, we add some individual data
defining a pair of diploid sibs and their parents.

:::{note}
We set the nodes for (e.g.) individual 2 to [1, 3] here to illustrate
that nodes for a given individual are not necessarily contiguous.
:::

```{code-cell}
tables.individuals.add_row(parents=[-1, -1])
tables.individuals.add_row(parents=[-1, -1])
tables.individuals.add_row(parents=[0, 1])
tables.individuals.add_row(parents=[0, 1])
node_individual = tables.nodes.individual
node_individual[[1, 3]] = 2
node_individual[[0, 2]] = 3
tables.nodes.individual = node_individual
display(tables.individuals)
display(tables.nodes)
ts = tables.tree_sequence()
ts.write_vcf(sys.stdout)
```

In this model we have four individuals defined, but only
individuals 2 and 3 are associated with nodes (more specifically,
**sample** nodes). Thus, we output **two** VCF sample individuals
composed of the linked nodes.

:::{note}
Note that the labels are ``tsk_0`` and ``tsk_1`` even though
the individual IDs are 2 and 3. See the
{ref}`sec_export_vcf_individual_names` section for how to change the
these default labels.
:::

(sec_export_vcf_individual_names)=

### Individual names

By default the VCF samples are given the labels ``tsk_0``, ``tsk_1``,
..., ``tsk_{N - 1}``, where ``N`` is the number of individuals to
be output (see the {ref}`sec_export_vcf_constructing_gt` section).

We can change this default labelling using the ``individual_names``
argument::

```{code-cell}
import sys
import msprime
ts = msprime.sim_ancestry(
    samples=3, ploidy=2, sequence_length=10, random_seed=2)
ts = msprime.sim_mutations(ts, rate=0.1, random_seed=2)
ts.write_vcf(sys.stdout, individual_names=["A", "B", "C"])
```

#### Exporting to plink

The default VCF sample IDs produced by ``tskit`` do not work well
with plink because it parses the individual
IDs based on a particular format, and does not allow ``0`` as a valid
identifier. We get an error like this:

```
Error: Sample ID ends with "_0", which induces an invalid IID of '0`.
```

This can be fixed by using the ``individual_names`` argument
to set the names to anything where the first name doesn't end with ``_0``.
An example implementation for diploid individuals is:

:::{code-block}
n_dip_indv = int(ts.num_samples / 2)
indv_names = [f"tsk_{i}indv" for i in range(n_dip_indv)]
with open("output.vcf", "w") as vcf_file:
    ts.write_vcf(vcf_file, individual_names=indv_names)
:::

Adding a second ``_`` (eg: ``tsk_0_indv``) is not recommended as
``plink`` uses ``_`` as the default separator for separating family
id and individual id, and two underscores will throw an error.

(sec_export_vcf_modifying_coordinates)=

### Modifying coordinates

Tskit supports continuous genome coordinates, but VCF only supports discrete
genome positions. Thus, when we convert a tree sequence that has sites
at continuous positions we must discretise the coordinates in some way.

The ``position_transform`` argument provides a way to flexibly translate
the genomic location of sites in tskit to the appropriate value in VCF.
There are two fundamental differences in the way that tskit and VCF define
genomic coordinates. The first is that tskit uses floating point values
to encode positions, whereas VCF uses integers. Thus, if the tree sequence
contains positions at non-integral locations there is an information loss
incurred by translating to VCF. By default, we round the site positions
to the nearest integer, such that there may be several sites with the
same integer position in the output. The second difference between VCF
and tskit is that VCF is defined to be a 1-based coordinate system, whereas
tskit uses 0-based. However, how coordinates are transformed depends
on the VCF parser, and so we do **not** account for this change in
coordinate system by default.

:::{note}
The msprime 0.x legacy API simulates using continuous coordinates. It may
be simpler to update your code to use the msprime 1.0 API (which uses
discrete coordinates by default) than to work out how to transform
coordinates in a way that is suitable for your application.
:::

:::{todo}
Provide some examples of using position transform.
:::


--- ../../tskit/docs/what_is.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

```{code-cell} ipython3
:tags: [remove-cell]
import msprime
import demes


def whatis_example():
    demes_yml = """\
        description:
          Asymmetric migration between two extant demes.
        time_units: generations
        defaults:
          epoch:
            start_size: 5000
        demes:
          - name: Ancestral_population
            epochs:
              - end_time: 1000
          - name: A
            ancestors: [Ancestral_population]
          - name: B
            ancestors: [Ancestral_population]
            epochs:
              - start_size: 2000
                end_time: 500
              - start_size: 400
                end_size: 10000
        migrations:
          - source: A
            dest: B
            rate: 1e-4
        """
    with open("data/whatis_example.yml", "wt") as f:
        f.write(demes_yml)
    graph = demes.loads(demes_yml)
    demography = msprime.Demography.from_demes(graph)
    # Choose seed so num_trees=3, tips are in same order,
    # first 2 trees are topologically different, and all trees have the same root
    seed = 12581
    ts = msprime.sim_ancestry(
        samples={"A": 2, "B": 3},
        demography=demography,
        recombination_rate=1e-8,
        sequence_length=1000,
        random_seed=seed)
    # Mutate
    # Choose seed to give 12 muts, last one above node 14
    seed = 1476
    ts = msprime.sim_mutations(ts, rate=1e-7, random_seed=seed)
    ts.dump("data/whatis_example.trees")

def create_notebook_data():
    whatis_example()

# create_notebook_data()  # uncomment to recreate the tree seqs used in this notebook
```

(sec_what_is)=

# What is a tree sequence?

A *succinct tree sequence*, or "tree sequence" for short, represents the ancestral
relationships between a set of DNA sequences. Tree sequences are based on fundamental
biological principles of inheritance, DNA duplication, mutation, and recombination;
they can be created by [evolutionary simulation](https://tskit.dev/software/#simulate)
or by [inferring genealogies from empirical DNA data](https://tskit.dev/software/#infer).

:::{margin} Key point
Tree sequences are used to encode and analyse large genetic datasets
:::

Tree sequences provide an efficient way of storing
[genetic variation](https://en.wikipedia.org/wiki/Genetic_variation) data, and can
power analyses of millions of whole [genomes](https://en.wikipedia.org/wiki/Genome).
Plots (a) and (b) below summarize these aspects 
(see additional details on [storage](plot_storing_everyone) and
[compute](plot_incremental_calculation) further down).

```{code-cell} ipython3
:"tags": ["remove-input"]
# This cell deliberately removed (not just hidden via a toggle) as it's not helpful
# for understanding tskit code (it's merely plotting code)
import matplotlib_inline
import matplotlib.pyplot as plt
import numpy as np
%matplotlib inline
matplotlib_inline.backend_inline.set_matplotlib_formats('svg')

data1 = np.genfromtxt("data/storing_everyone.csv", delimiter=",", usecols=np.arange(1,12), names=True)
data2 = np.genfromtxt("data/benchmarks_without_copy_longer_genome.txt", encoding=None, names=True, dtype=None)
fig, (ax1, ax2) = plt.subplots(1,2, figsize=(16, 4.5))
fig.subplots_adjust(wspace=0.5, left=0, right=1)
keep = data1['sample_size'] <= 1e6
x, y = data1['sample_size'][keep], data1['tsk_fit'][keep]/data1['vcf_fit'][keep]
ax1.spines["top"].set_visible(False)
ax1.spines["right"].set_visible(False)
ax1.loglog(x, y, c="C0", linewidth=4)
ax1.set_xlabel('# of 100Mb genomes', fontsize=18)
ax1.set_ylabel('Size of tree sequence\nfile (relative to VCF) ', fontsize=18)
ax1.tick_params(axis="both", labelsize=16)

txt = ax1.text(0.5, 1.3, "(a) Storing a million genomes as a tree sequence takes thousands of times less disk space",
    ha='center', va='top', transform=ax1.transAxes, wrap=True, size=24)
txt._get_wrap_line_width = lambda: 600

ts_time = {n: t for s, n, t in data2[['toolkit','nsam','seconds']] if s == 'tskit'}
libseq_time = {n: t for s, n, t in data2[['toolkit','nsam','seconds']] if s == 'libseq'}
x = np.unique(list(ts_time.keys()) + list(libseq_time.keys()))
y = np.array([libseq_time[time]/ts_time[time] for time in x])
ax2.spines["top"].set_visible(False)
ax2.spines["right"].set_visible(False)
ax2.loglog(x, y, linewidth=4)
ax2.set_xlabel("# of genomes", fontsize=18)
ax2.set_ylabel("Tajima's D calculations per\nunit time (relative to libseq)", fontsize=18)
ax2.tick_params(axis="both", labelsize=16)
txt = ax2.text(0.5, 1.3, "(b) Genetic calculations on millions of genomes can be sped up by many orders of magnitude",
    ha='center', va='top', transform=ax2.transAxes, wrap=True, size=24
)    
txt._get_wrap_line_width = lambda: 600
plt.show()
```

## How to think about tree sequences

(sec_what_is_local_trees)=

### A sequence of trees...

As the name suggests, the simplest way to think about a tree sequence is that it
describes a sequence of correlated "local trees" --- i.e. genetic trees located at
different points along a [chromosome](https://en.wikipedia.org/wiki/Chromosome).
Here's a tiny example based on ten haploid genomes, $\mathrm{a}$ to $\mathrm{j}$,
spanning a short 1000 letter chromosome.

```{code-cell} ipython3
:"tags": ["hide-input"]
import string
import tskit

mutated_ts = tskit.load("data/whatis_example.trees")
ts = mutated_ts.delete_sites(list(range(mutated_ts.num_sites)))
# Extra code to label and order the tips alphabetically rather than numerically
labels = {i: string.ascii_lowercase[i] for i in range(ts.num_nodes)}
genome_order = [n for n in ts.first().nodes(order="minlex_postorder") if ts.node(n).is_sample()]
labels.update({n: labels[i] for i, n in enumerate(genome_order)})
style1 = (
    ".node:not(.sample) > .sym, .node:not(.sample) > .lab {visibility: hidden;}"
    ".mut {font-size: 12px} .y-axis .tick .lab {font-size: 85%}")
sz = (800, 250)  # size of the plot, slightly larger than the default
ticks = [0, 5000, 10000, 15000, 20000]
ts.draw_svg(
    size=sz, node_labels=labels, style=style1, y_label="Time ago",
    y_axis=True, y_ticks=ticks)
```

::::{margin}
:::{note}
For clarity in these examples, we are using letters to label nodes. Normally, however,
the nodes are referred to by {ref}`numerical ID<sec_terminology_nodes>`.
:::
::::

The tickmarks on the X axis and background shading indicate the genomic positions covered
by the trees. The tickmarks indicate recombination events that explain relationships
between the ten genomes. There were two such recombination events, giving us three local trees.
For the first short portion of the chromosome, from the start until position 189,
the relationships between the ten genomes are shown by the first tree.
The second tree shows the relationships between positions 189 and 546.
By inspecting the first and the second local tree we can see that genomes $\mathrm{b}-\mathrm{f}$
changed their "most recent common ancestor" (MRCA) with genome $\mathrm{a}$ to
MRCA with genome $\mathrm{g}$.
The third tree shows the relationships between positions 546 and 1000 (the end).
By inspecting the second and the third local tree we can see that
recombination changed the ancestry of genomes $\mathrm{b}-\mathrm{f}$
back to shared MRCA with genome $\mathrm{g}$.

(sec_what_is_genealogical_network)=

### ... created by a genealogical network

In fact, succinct tree sequences don't store each tree separately, but instead are
based on an interconnected *genetic genealogy*, in which
[genetic recombination](https://en.wikipedia.org/wiki/Genetic_recombination) has led
to different regions of the chromosome having different histories. Another way of
thinking about the tree sequence above is that it describes the full genetic ancestry
of our 10 genomes.

(sec_what_is_dna_data)=

## An efficient encoding of DNA data

A tree sequence can be used to describe patterns of genetic variation by combining the
trees with a knowledge of where *mutations* occur on their branches. Here's how that
might look in our simple example:

```{code-cell} ipython3
:"tags": ["hide-input"]

mut_labels = {}  # An array of labels for the mutations, listing position & allele change
l = "{:g} ({}→{})"
for mut in mutated_ts.mutations():  # This entire loop is just to make pretty labels
    site = mutated_ts.site(mut.site)
    older_mut = mut.parent >= 0  # is there an older mutation at the same position?
    prev = mutated_ts.mutation(mut.parent).derived_state if older_mut else site.ancestral_state
    mut_labels[mut.id] = l.format(site.position, prev, mut.derived_state)

mutated_ts.draw_svg(
    size=sz, style=style1, node_labels=labels, mutation_labels=mut_labels)
```

There are now twelve single nucleotide mutations in the tree sequence. They are shown on the
branches of the trees, and the positions of the twelve variable sites associated with the
mutations are shown along the X axis.

:::{margin} Key point
Mutation on trees are the source of genetic variation
:::

The trees inform us that, for example, the final mutation (at position 987) is inherited
by genomes $\mathrm{h}$ to $\mathrm{j}$. These genomes must have an *T* at that position,
compared to the original value of *G*. In other words, once we know the ancestry, placing
a relatively small number of mutations is enough to explain all the observed genetic
variation. Here's the resulting "variant matrix":

```{code-cell} ipython3
:"tags": ["hide-input"]
haplotypes = ["   ".join(h) for h in mutated_ts.haplotypes()]
print("Position: " + "".join(f"{s.position:^4g}" for s in mutated_ts.sites()))
print("\n".join(sorted(
    [f"Genome {labels[i]}:  {h}" for i, h in zip(mutated_ts.samples(), haplotypes)])))
```

This approach scales effectively to millions of genomes, and to chromosomes of
hundreds of megabases in length. The ability to deal with huge datasets comes down to
one key feature of genomic data: adjacent trees along a chromosome are highly correlated,
that is, they *share structure*. In our example this becomes evident
if we highlight the branches ("edges" in tree sequence terminology) that remain
unchanged between the first and the second tree.

(fig_what_is_edge_diffs)=

```{code-cell} ipython3
:"tags": ["hide-input"]
# Highlight certain edges in certain trees. Other visualization possibilities in tutorials/viz.html
kept_edges = [e for e in ts.edges() if e.left==0 and e.right>ts.breakpoints(True)[1]]
style3 = (
    ",".join(f"#svg1 .tree:not(.t2) .node.a{e.parent}.n{e.child} > .edge" for e in kept_edges)
    + "{stroke:#00DD00; stroke-width: 2px}"
    + style1)
ts.draw_svg(
    size=(500, 250), x_lim=(0, 500), root_svg_attributes={'id':'svg1'},  y_ticks=ticks,
    node_labels=labels, style=style3)
```

:::{margin} Key point
Tree sequences are efficient because they don't store each tree separately
:::

A branch can be shared by many adjacent trees, but is stored as a single edge in the tree
sequence. For large datasets this is a great saving, because typically each tree-change
affects only a few branches at a time, regardless of the tree size.

Below is an extension of the plot at the top of this page, showing predicted
file sizes when storing not just millions, but billions of human-like genomes:
enough to encompass every human on the planet. This demonstrates that the tree sequence
encoding leads to savings of many orders of magnitude, even when compared against
compressed versions of the standard VCF storage format (original published data
[here](https://www.nature.com/articles/s41588-019-0483-y/figures/1)). It's also worth
noting that the efficiency extends to processing time too: tree sequences are often
several orders of magnitude faster to process than other storage formats.

(plot_storing_everyone)=

```{code-cell} ipython3
:"tags": ["remove-input"]
# This cell deliberately removed (not just hidden via a toggle) as it's not helpful
# for understanding tskit code (it's merely plotting code)
x = data1['sample_size']
fig, ax1 = plt.subplots(1, figsize=(10, 4))
ax1.spines["top"].set_visible(False)
ax1.spines["right"].set_visible(False)

plt.loglog(x,  data1['vcf_fit'], c="C1", label="VCF", linewidth=2)
plt.loglog(x,  data1['vcfz_fit'], c="C1", label="compressed VCF", linewidth=2, linestyle=":")

plt.loglog(x, data1['tsk_fit'], c="C0", label="tree sequence", linewidth=2)
plt.loglog(x, data1['tskz_fit'], c="C0", label="compressed tree sequence", linewidth=2, linestyle=":")

plt.xlabel('Number of 100Mb genomes (log scale)', fontsize=12)
plt.ylabel('Space required (GB, log scale)', fontsize=12)
plt.text(16e9, 0.001, 'Size of\nentire\nhuman\npopulation', ha="center", va="bottom", size=14)
plt.annotate('', xy=(16e9, 0.0001), xytext=(16e9, 0.001), 
            arrowprops=dict(facecolor='black', shrink=0))
plt.legend()
plt.show()
```

(sec_what_is_ancestry)=

## A record of genetic ancestry

::::{margin}
:::{note}
The genetic genealogy is sometimes referred to as an ancestral recombination graph,
or ARG, and one way to think of tskit tree sequence is as a way
to store various different sorts of ARGs (see the {ref}`ARG tutorial<sec_args>`)
:::
::::

Often, we're not interested so much in the DNA sequence data as the genetic ancestry
itself (discussed e.g. [here](https://www.nature.com/articles/s41588-019-0492-x)).
In other words, the main consideration is the actual trees in a tree sequence, rather
than the distributions of mutations placed upon them --- indeed in genetic simulations, it
{ref}`may not be necessary<sec_tskit_no_mutations>` to incorporate neutral mutations at all.
The trees reflect, for example, the origin and age of alleles under
selection, the spatial structure of populations, and the effects
of hybridization and admixture in the past.

The tree sequence in this tutorial was actually generated using a model of population
splits and expansions as shown in the following schematic,
{ref}`plotted<sec_tskit_viz_other_demographic>` using the
[DemesDraw](https://pypi.org/project/demesdraw/) package. Our 10 genomes were sampled
from modern day populations A (a constant-size population) and B (a recently expanding
one), where limited migration is occuring from A to B.

```{code-cell} ipython3
:"tags": ["remove-input"]
# This cell deliberately removed (not just hidden via a toggle) as it's not helpful
# for understanding tskit code (it's merely plotting code taken from the demesdraw docs)
import demes
import demesdraw

def size_max(graph):
    return max(
        max(epoch.start_size, epoch.end_size)
        for deme in graph.demes
        for epoch in deme.epochs
    )

graph = demes.load("data/whatis_example.yml")
w = 1.5 * size_max(graph)
positions = dict(Ancestral_population=0, A=-w, B=w)
fig, ax = plt.subplots(1, figsize=(5, 3))
ax = demesdraw.tubes(graph, ax=ax, positions=positions, seed=1)
plt.show(ax.figure)
```

A major benefit of "tree sequence thinking" is the close relationship between the
tree sequence and the underlying biological processes that produced the genetic
sequences in the first place, such as those pictured in the demography above. For
example, each branch point (or "internal node") in one of our trees can be
imagined as a genome which existed at a specific time in the past, and
which is a MRCA of the descendant genomes at that position on the chromosome.
We can mark these extra "ancestral genomes" on our tree diagrams with circular symbols,
distinguishing them from the *sampled* genomes ($\mathrm{a}$ to $\mathrm{j}$)
marked with square symbols. We can even colour the nodes by the population that we know
(or infer) them to belong to at the time:

```{code-cell} ipython3
:"tags": ["hide-input"]
colours = {0: "#1f77b4", 1: "#ff7f0e", 2: "#2ca02c"}
style2 = ".y-axis .tick .lab {font-size: 85%}"
style2 += "#svg2 .node > .sym {visibility: visible;}"  # force-show all nodes: not normally needed
style2 += "".join([f".p{n.population} > .sym {{fill: {colours[n.population]}}}" for n in ts.nodes()])

mutated_ts.draw_svg(
    size=sz, root_svg_attributes={'id':'svg2'}, y_label="Time ago (generations)",
    y_axis=True, y_ticks=ticks, node_labels=labels, mutation_labels={}, style=style2)
```

The diagram shows that most of the ancestral genomes $\mathrm{k}$ to $\mathrm{u}$
lived much longer ago than the population split, 1000 generations back, and
resided in the ancestral (blue) population. The tree sequence also allows us to easily
deduce these MRCA genomes, simply by looking at which mutations they have inherited:

```{code-cell} ipython3
:"tags": ["hide-input"]
import numpy as np
tables = mutated_ts.dump_tables()
# Flip sample and nonsample flags, making the haplotypes() method print out nonsample nodes
s_flags = tables.nodes.flags[ts.samples()[0]]
no_flags = s_flags-s_flags
tables.nodes.flags = np.where(tables.nodes.flags & tskit.NODE_IS_SAMPLE, no_flags, s_flags)
ts_flipped = tables.tree_sequence()
haplotypes = ["   ".join(h) for h in ts_flipped.haplotypes(missing_data_character=" ")]
print(" " * ts_flipped.num_sites, " " * (ts_flipped.num_sites-4), "")
print(
    "||ANCESTRAL GENOMES||     Position:",
    "".join(f"{s.position:^4g}" for s in ts_flipped.sites()))
print(
    "\n".join(reversed(sorted([
        f"Genome {labels[i]} ({ts.node(i).time:7.1f} {ts_flipped.time_units} ago):  {h}"
        for i, h in zip(ts_flipped.samples(), haplotypes)]))))
```

You can see that some ancestors are missing genomic regions, because those parts of
their genome have not been inherited by any of the sampled genomes. In other words, that
ancestral node is not present in the corresponding local tree.


(sec_what_is_analysis)=

## A framework for efficient computation


Using tree structures is a common way to implement efficient computer algorithms, and
many [phylogenetic](https://en.wikipedia.org/wiki/Phylogenetics) methods use the
structure provided by the evolutionary tree to implement efficient 
[dynamic programming](https://en.wikipedia.org/wiki/Dynamic_programming)
algorithms. The tree sequence structure allows these approaches
to be extended to the particular form of
[phylogenetic network](https://en.wikipedia.org/wiki/Phylogenetic_network) defined by
multiple correlated trees along a genome.


```{margin} Key point
Most genetic calculations involve iterating over trees, which is highly efficient in
{program}`tskit`
```

For example, statistical measures of genetic variation can be thought of as a calculation
combining the local trees with the mutations on each branch (or, often preferably, the
length of the branches: see [this summary](https://doi.org/10.1534/genetics.120.303253)).
Because a tree sequence is built on a set of small branch changes along the chromosome,
statistical calculations can often be updated incrementally as we
move along the genome, without having to perform the calculation *de novo* on each tree.
Using tree sequences can result in speed-ups of many
orders of magnitude when perfoming calculations on large datasets, as in this example of
calculating [Tajima's D](https://en.wikipedia.org/wiki/Tajima%27s_D)
(from [here](https://www.genetics.org/content/215/3/779#F9)):

(plot_incremental_calculation)=
```{code-cell} ipython3
:"tags": ["remove-input"]
# This cell deliberately removed (not just hidden via a toggle) as it's not helpful
# for understanding tskit code (it's merely plotting code)
ts_time = np.array([[n,t] for s, n, t in data2[['toolkit','nsam','seconds']] if s == 'tskit'])
ska_time = np.array([[n, t] for s, n, t in data2[['toolkit','nsam','seconds']] if s == 'allel'])
libseq_time = np.array([[n, t] for s, n, t in data2[['toolkit','nsam','seconds']] if s == 'libseq'])
fig, ax1 = plt.subplots(1, figsize=(10, 5))
ax1.spines["top"].set_visible(False)
ax1.spines["right"].set_visible(False)
ax1.loglog(ska_time[:,0], ska_time[:,1], c="C3", linewidth=2, label="scikit-allel library")
ax1.loglog(libseq_time[:,0], libseq_time[:,1], c="C1", linewidth=2, label="libseq library")
ax1.loglog(ts_time[:,0], ts_time[:,1], c="C0", linewidth=2, label="tskit library")
ax1.set_ylabel("Time to calculate Tajima's D (secs/site)", fontsize=12)
ax1.set_xlabel("Number of sampled genomes", fontsize=12)
plt.legend()
plt.show()
```

The {program}`tskit` library has {ref}`extensive support<sec_analysing_tree_sequences>`
for these sorts of population genetic calculations. It provides efficient methods for
traversing through large {ref}`trees<sec_analysing_trees_traversals>` and
{ref}`tree sequences<sec_processing_trees>`, as well as providing other
phylogenetically relevant methods such as
{ref}`parsimonious placement of mutations<sec_analysing_trees_parsimony>`,
and the {ref}`counting of topologies<sec_counting_topologies>` embedded within
larger trees.

If you are new to tree sequences, and want to start finding out about {program}`tskit`,
you might now want to continue to the next tutorial: {ref}`sec_terminology_and_concepts`.


## Further reading

* Jump straight in: if you already have a tree sequence you wish to deal with, the
  {ref}`sec_tskit_getting_started` tutorial show you how to do a number of common tasks.
* How is a tree sequence stored: details in the {ref}`sec_tables` tutorial
* The offical {program}`tskit` [documentation](https://tskit.dev/tskit/docs)


--- ../../tskit/docs/metadata.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

```{currentmodule} tskit
```

(sec_metadata)=

# Metadata

The tree-sequence and all the entities within it (nodes, mutations, edges,  etc.) can
have metadata associated with them. This is intended for storing and passing on
information that tskit itself does not use or interpret, for example information derived
from a VCF INFO field, or administrative information (such as unique identifiers)
relating to samples and populations. Note that provenance information about how a tree
sequence was created should not be stored in metadata, instead the provenance mechanisms
in tskit should be used (see {ref}`sec_provenance`).

The metadata for each entity (e.g. row in a table) is described by a schema for each
entity type (e.g. table). The schemas allow the tskit Python API to encode and decode
metadata automatically and, most importantly, tells downstream users and tools how to
decode and interpret the metadata. For example, the `msprime` schema for populations
requires both a `name` and a `description` for each defined population: these names and
descriptions can assist downstream users in understanding and using `msprime` tree
sequences. It is best practice to populate such metadata fields if your files will be
used by any third party, or if you wish to remember what the rows refer to some time
after making the file!

Technically, schemas describe what information is stored in each metadata record, and
how it is to be encoded, plus some optional rules about the types and ranges of data
that can be stored. Every node's metadata follows the node schema, every mutation's
metadata the mutation schema, and so on. Most users of tree-sequence files will not
need to modify the schemas: typically, as in the example of `msprime` above, schemas are
defined by the software which created the tree-sequence file. The exact metadata stored
depends on the use case; it is also possible for subsequent processes to add or modify
the schemas, if they wish to add to or modify the types (or encoding) of the metadata.

The metadata schemas are in the form of a
[JSON Schema](http://json-schema.org/) (a good guide to JSON Schema is at
[Understanding JSON Schema](https://json-schema.org/understanding-json-schema/)). The
schema must specify an object with properties,
the keys and types of those properties are specified along with optional
long-form names, descriptions and validations such as min/max or regex matching for
strings, see the {ref}`sec_metadata_schema_examples` below.

As a convenience the simplest, permissive JSON schema is available as
{meth}`MetadataSchema.permissive_json()`.

The {ref}`sec_tutorial_metadata` Tutorial shows how to use schemas and access metadata
in the tskit Python API.

Note that the C API simply provides byte-array binary access to the metadata,
leaving the encoding and decoding to the user. The same can be achieved with the Python
API, see {ref}`sec_tutorial_metadata_binary`.


(sec_metadata_examples)=

## Examples

In this section we give some examples of how to define metadata
schemas and how to add metadata to various parts of a tree sequence
using the Python API. For simplicity, these initial examples use the JSON codec
(see {ref}`sec_metadata_codecs`).

(sec_metadata_examples_top_level)=

### Top level

Top level metadata is associated with the tree sequence as a whole, rather than
any specific table. This is used, for example, by programs such as
[SLiM](https://github.com/MesserLab/SLiM) to store information about the sort of
model that was used to generate the tree sequence (but note that detailed information
used to recreate the tree sequence is better stored in {ref}`sec_provenance`).

Here's an example of adding your own top-level metadata to a tree sequence:

```{code-cell}
import tskit
# Define some top-level metadata you might want to add to a tree sequence
top_level_metadata = {
    "taxonomy": {"species": "Arabidopsis lyrata", "subspecies": "petraea"},
    "generation_time": 2,
}

# Generate a simple tree sequence of one random tree.
ts = tskit.Tree.generate_random_binary(8, branch_length=10, random_seed=9).tree_sequence

# To edit a tree sequence, first dump it to tables.
tables = ts.dump_tables()

# Set the metadata schema for the top-level metadata
tables.metadata_schema = tskit.MetadataSchema.permissive_json()  # simplest schema
# Set the metadata itself
tables.metadata = top_level_metadata

ts = tables.tree_sequence()
print(
    "The tree sequence is of",
    ts.metadata["taxonomy"]["species"],
    "subsp.",
    ts.metadata["taxonomy"]["subspecies"],
)
```

In this case, the species and subspecies name are self-explanatory, but
the interpretation of the `generation_time` field is less clear. Setting
a more precise schema will help other users of your tree sequence:

```{code-cell}
schema = {
    "codec": "json",
    "type": "object",
    "properties": {
        "generation_time": {"type": "number", "description": "Generation time in years"},
    },
    "additionalProperties": True,  # optional: True by default anyway
}
tables.metadata_schema = tskit.MetadataSchema(schema)
tables.metadata = top_level_metadata  # put the metadata back in
ts = tables.tree_sequence()
print(ts.metadata)
```

Note that the schema here only describes the `generation_time` field. The
metadata also contains additional fields (such as the species) that are
not in the schema; this is allowed because `additionalProperties` is `True`
(assumed by default in the {ref}`sec_metadata_codecs_json` codec, but shown
above for clarity). 

Explicitly specified fields are *validated* on input, helping to avoid errors.
For example, setting the generation time to a string will now raise an error:

```{code-cell}
:tags: ["raises-exception", "output_scroll"]
tables.metadata = {"generation_time": "two of your earth years"}
```

:::{note}
Although we have stored the generation time in metadata, the
time *units* of a tree sequence should be stored in the 
{attr}`~TreeSequence.time_units` attribute, not in
metadata. For example, we could set `tables.time_units = "generations"`.
:::

(sec_metadata_examples_reference_sequence)=

### Reference sequence

Often a genome will be associated with a
reference sequence for that species. In this case, we might want to
store not just the species name, but also e.g. the build version of
the reference sequence, and possibly the reference sequence itself.
There is built-in support for this in tskit, via the
{attr}`~tskit.ReferenceSequence.metadata` and
{attr}`~tskit.ReferenceSequence.metadata_schema` properties
of the {attr}`TreeSequence.reference_sequence` attribute
(see the {ref}`sec_data_model_reference_sequence` documentation).

:::{todo}
Add examples of reference sequence metadata when the API becomes less
preliminary. This should
include an example where we declare (or better, use on we define
in the library) a standard metadata schema for a species, which
defines and documents accession numbers, genome builds, etc. e.g.

```python
tables.reference_sequence.metadata_schema = standard_schema
tables.reference_sequence.metadata = {...}
ts = tables.tree_sequence()
```
:::

(sec_metadata_examples_tables)=

### Tables

Each table in a tree sequence (apart from the provenance table)
can have its own metadata, and associated metadata schema.

```{code-cell}
tables.individuals.metadata_schema = tskit.MetadataSchema.permissive_json()
tables.individuals.add_row(metadata={"Accession ID": "ABC123"})
ts = tables.tree_sequence()
print(",\n           ".join(str(ts.individual(0)).split(", ")))

```

However, we might want something more descriptive than the default
{meth}`~MetadataSchema.permissive_json()`. schema. We could create
a new schema, or modify the existing one. Modification is useful
if a nontrivial schema has been set already, for example in the
{ref}`case of populations <msprime:sec_demography_populations_metadata>`
when the tree sequence has been generated by
{func}`msprime:msprime.sim_ancestry`.

```{code-cell}
# Modify an existing schema
schema_as_python_dict = tables.individuals.metadata_schema.schema
if "properties" not in schema_as_python_dict:
    schema_as_python_dict["properties"] = {}
schema_as_python_dict["properties"]["Accession ID"] = {
    "type": "string", "description": "An accession ID for this individual"}

# Optional: require an accession id to be specified for all individuals
if "required" not in schema_as_python_dict:
    schema_as_python_dict["required"] = []
schema_as_python_dict["required"].append("Accession ID")

# Set the schema back on the table
tables.individuals.metadata_schema = tskit.MetadataSchema(schema_as_python_dict)

# Put all the metadata back in, using validate_and_encode_row to validate it
tables.individuals.packset_metadata([
    tables.individuals.metadata_schema.validate_and_encode_row(ind.metadata)
    for ind in tables.individuals
])
print("New schema:", tables.individuals.metadata_schema)
```


### Defaults

Since we specified that the `accession_id` property was required in the
example above, the user *always* has to provide it, otherwise it will
fail to validate:

```{code-cell}
:tags: ["raises-exception", "output_scroll"]
tables.individuals.add_row(metadata={"Comment": "This has no accession ID"})
```

However, rather than require a user-specified value, we can provide a
default, which will be returned if the field is absent. In this case the property
should not be marked as `required`.

```{code-cell}
new_schema = {
    "codec": "json",
    "type": "object",
    "properties": {
        "Accession ID": {
            "type": "string",
            "description": "An accession ID for this individual",
            "default": "N/A",  # Default if this property is absent
        },
    },
    "default": {"Accession ID": "N/A"},  # Default if no metadata in this row
}
tables.individuals.metadata_schema = tskit.MetadataSchema(new_schema)
tables.individuals.packset_metadata([
    tables.individuals.metadata_schema.validate_and_encode_row(ind.metadata)
    for ind in tables.individuals
])
tables.individuals.add_row(metadata={"Comment": "This has no accession ID"})
ts = tables.tree_sequence()

print("Newly added individual:")
print(",\n           ".join(str(ts.individual(-1)).split(", ")))
```

:::{note}
In the {ref}`sec_metadata_codecs_json` codec, defaults can only
be set for the shallowest level of the metadata object.
:::

(sec_metadata_codecs)=

## Codecs

The underlying metadata is in raw binary (see
{ref}`data model <sec_metadata_definition>`) and so it
must be encoded and decoded. The C API does not do this, but the Python API will
use the schema to decode the metadata to Python objects.
The encoding for doing this is specified in the top-level schema property `codec`.
Currently the Python API supports the `json` codec which encodes metadata as
[JSON](https://www.json.org/json-en.html), and the `struct` codec which encodes
metadata in an efficient schema-defined binary format using {func}`python:struct.pack` .

(sec_metadata_codecs_json)=

### `json`

When `json` is specified as the `codec` in the schema the metadata is encoded in
the human readable [JSON](https://www.json.org/json-en.html) format. As this format
is human readable and encodes numbers as text it uses more bytes than the `struct`
format. However it is simpler to configure as it doesn't require any format specifier
for each type in the schema. Tskit deviates from standard JSON in that
empty metadata is interpreted as an empty object. This is to allow setting of a schema
to a table with out the need to modify all existing empty rows.

(sec_metadata_codecs_struct)=

### `struct`

When `struct` is specifed as the `codec` in the schema the metadata is encoded
using {func}`python:struct.pack` which results in a compact binary representation which
is much smaller and generally faster to encode/decode than JSON.

This codec places extra restrictions on the schema:

1. Each property must have a `binaryFormat`
    This sets the binary encoding used for the property.

2. All metadata objects must have fixed properties.
    This means that additional properties not listed in the schema are disallowed. Any
    property that does not have a `default` specified in the schema must be present.
    Default values will be encoded.

3. Arrays must be lists of homogeneous objects.
    For example, this is not valid:
    ```
    {"type": "array", "items": [{"type": "number"}, {"type": "string"}]}
    ```

4. Types must be singular and not unions.
    For example, this is not valid:
    ```
    {"type": ["number", "string"]}
    ```
    One exception is that the top-level can be a union of `object` and `null` to
    support the case where some rows do not have metadata.

5. The order that properties are encoded is by default alphabetically by name.
    The order can be overridden by setting an optional numerical `index` on each
    property. This is due to objects being unordered in JSON and Python `dicts`.


#### binaryFormat

To determine the binary encoding of each property in the metadata the `binaryFormat` key is used.
This describes the encoding for each property using `struct`
[format characters](https://docs.python.org/3/library/struct.html#format-characters).
For example an unsigned 8-byte integer can be specified with::

```
{"type": "number", "binaryFormat":"Q"}
```

And a length 10 string with::

```
{"type": "string", "binaryFormat":"10p"}
```

Some of the text below is copied from
[the python docs](https://docs.python.org/3/library/struct.html).


##### Numeric and boolean types

The supported numeric and boolean types are:


```{list-table}
:header-rows: 1
* - Format
  - C Type
  - Python type
  - Numpy type
  - Size in bytes
* - `?`
  - *_Bool*
  - bool
  - bool
  - 1
* - `b`
  - *signed char*
  - integer
  - int8
  - 1
* - `B`
  - *unsigned char*
  - integer
  - uint8
  - 1
* - `h`
  - *short*
  - integer
  - int16
  - 2
* - `H`
  - *unsigned short*
  - integer
  - uint16
  - 2
* - `i`
  - *int*
  - integer
  - int32
  - 4
* - `I`
  - *unsigned int*
  - integer
  - uint32
  - 4
* - `l`
  - `long`
  - integer
  - int32
  - 4
* - `L`
  - *unsigned long*
  - integer
  - uint32
  - 4
* - `q`
  - `long long`
  - integer
  - int64
  - 8
* - `Q`
  - *unsigned long long*
  - integer
  - uint64
  - 8
* - `f`
  - *float*
  - float
  - float32
  - 4
* - `d`
  - *double*
  - float
  - float64
  - 8
```

When attempting to pack a non-integer using any of the integer conversion
codes, if the non-integer has a `__index__` method then that method is
called to convert the argument to an integer before packing.

For the `'f'` and `'d'` conversion codes, the packed
representation uses the IEEE 754 binary32 or binary64 format (for
`'f'` or `'d'` respectively), regardless of the floating-point
format used by the platform.

Note that endian-ness cannot be specified and is fixed at little endian.

When encoding a value using one of the integer formats (`'b'`,
`'B'`, `'h'`, `'H'`, `'i'`, `'I'`, `'l'`, `'L'`,
`'q'`, `'Q'`), if the value is outside the valid range for that format
then {exc}`struct.error` is raised.

For the `'?'` format character, the decoded value will be either `True` or
`False`. When encoding, the truth value of the input is used.


##### Strings

```{list-table}
:header-rows: 1
* - Format
  - C Type
  - Python type
  - Size in bytes
* - `x`
  - pad byte
  - no value
  - as specified
* - `c`
  - *char*
  - string of length 1
  - 1
* - `s`
  - *char[]*
  - string
  - as specified
* - `p`
  - *char[]*
  - string
  - as specified
```

For the `'s'` format character, the number prefixed is interpreted as the length in
bytes, for example,
`'10s'` means a single 10-byte string. For packing, the string is
truncated or padded with null bytes as appropriate to make it fit. For
unpacking, the resulting bytes object always has exactly the specified number
of bytes, unless `nullTerminated` is `true`, in which case it ends at the first
`null`. As a special case, `'0s'` means a single, empty string.

The `'p'` format character encodes a "Pascal string", meaning a short
variable-length string stored in a fixed number of bytes, given by the count.
The first byte stored is the length of the string, or 255, whichever is
smaller.  The bytes of the string follow.  If the string to encode is too long
(longer than the count minus 1), only the leading
`count-1` bytes of the string are stored.  If the string is shorter than
`count-1`, it is padded with null bytes so that exactly count bytes in all
are used.  Note that strings specified with this format cannot be longer than 255.

Strings that are longer than the specified length will be silently truncated,
note that the length is in bytes, not characters.

The string encoding can be set with `stringEncoding` which defaults to `utf-8`.
A list of possible encodings is
[here](https://docs.python.org/3.7/library/codecs.html#standard-encodings).

For most cases, where there are no `null` characters in the metadata
`{"type":"string", "binaryFormat": "1024s", "nullTerminated": True}` is a good option
with the size set to that appropriate for the strings to be encoded.


##### Padding bytes

Unused padding bytes (for compatibility) can be added with a schema entry like:

```
{"type": "null", "binaryFormat":"5x"} # 5 padding bytes
```

##### Arrays

The codec stores the length of the array before the array data. The format used for the
length of the array can be chosen with `arrayLengthFormat` which must be one
of `B`, `H`, `I`, `L` or `Q` which have the same meaning as in the numeric
types above. `L` is the default. As an example:

```
{"type": "array", {"items": {"type":"number", "binaryFormat":"h"}}, "arrayLengthFormat":"B"}
```

Will result in an array of 2 byte integers, prepended by a single-byte array-length.

For dealing with legacy encodings that do not store the
length of the array, setting `noLengthEncodingExhaustBuffer` to `true` will read
elements of the array until the metadata buffer is exhausted. As such an array
with this option must be the last type in the encoded struct.


##### Union typed metadata

As a special case under the `struct` codec, the top-level type of metadata can be a
union of `object` and `null`. Set `"type": ["object", "null"]`. Properties should
be defined as normal, and will be ignored if the metadata is `None`.

(sec_metadata_schema_examples)=

## Schema examples

### Struct codec

As an example here is a schema using the `struct` codec which could apply, for example,
to the individuals in a tree sequence:

```python
complex_struct_schema = {
    "codec": "struct",
    "type": "object",
    "properties": {
        "accession_number": {"type": "integer", "binaryFormat": "i"},
        "collection_date": {
            "description": "Date of sample collection in ISO format",
            "type": "string",
            "binaryFormat": "10p",
            "pattern": "^([1-9][0-9]{3})-(1[0-2]|0[1-9])-(3[01]|0[1-9]|[12][0-9])?$",
        },
        "phenotype": {
            "description": "Phenotypic measurements on this individual",
            "type": "object",
            "properties": {
                "height": {
                    "description": "Height in metres, or NaN if unknown",
                    "type": "number",
                    "binaryFormat": "f",
                    "default": float("NaN"),
                },
                "age": {
                    "description": "Age in years at time of sampling, or -1 if unknown",
                    "type": "number",
                    "binaryFormat": "h",
                    "default": -1,
                },
            },
            "default": {},
        },
    },
    "required": ["accession_number", "collection_date"],
    "additionalProperties": False,
}

# Demonstrate use
tables.individuals.clear()
tables.individuals.metadata_schema = tskit.MetadataSchema(complex_struct_schema)
tables.individuals.add_row(
    metadata={"accession_number": 123, "collection_date": "2011-02-11"}
)
ts = tables.tree_sequence()
print(ts.individual(0).metadata)
```

This schema states that the metadata for each row of the table is an object consisting
of three properties. Property `accession_number` is a number (stored as a 4-byte int).
Property `collection_date` is a string which must satisfy a regex, which checks it is
a valid [ISO8601](https://www.iso.org/iso-8601-date-and-time-format.html) date. Property
`phenotype` is itself an object consisting of the properties `height` (a single precision
floating point number) and age (a 2 byte signed integer).
Because this is a struct codec, and neither of the the first two properties have a
default set, they must be marked as "required" (in the JSON codec if no default is given,
unspecified properties will simply be missing in the returned metadata dictionary).
Also because this is a struct codec, `additionalProperties` must be set to False. This
is assumed by default in the struct codec, but has been shown above for clarity.

(sec_metadata_api_overview)=

## Python Metadata API Overview

Schemas are represented in the Python API by the {class}`tskit.MetadataSchema`
class which can be assigned to, and retrieved from, tables via their `metadata_schema`
attribute (e.g. {attr}`tskit.IndividualTable.metadata_schema`). The schemas
for all tables can be retrieved from a {class}`tskit.TreeSequence` by the
{attr}`tskit.TreeSequence.table_metadata_schemas` attribute.

The top-level tree sequence metadata schema is set via
{attr}`tskit.TableCollection.metadata_schema` and can be accessed via
{attr}`tskit.TreeSequence.metadata_schema`.

Each table's `add_row` method (e.g. {meth}`tskit.IndividualTable.add_row`) will
validate and encode the metadata using the schema. This encoding will also happen when
tree sequence metadata is set (e.g. `table_collection.metadata = {...}`.

Metadata will be lazily decoded if accessed via
`tables.individuals[0].metadata`.  `tree_sequence.individual(0).metadata` or
`tree_sequence.metadata`

In the interests of efficiency the bulk methods of `set_columns`
(e.g. {meth}`tskit.IndividualTable.set_columns`)
and `append_columns` (e.g. {meth}`tskit.IndividualTable.append_columns`) do not
validate or encode metadata. See {ref}`sec_tutorial_metadata_bulk` for how to prepare
metadata for these methods.

Metadata processing can be disabled and raw bytes stored/retrieved. See
{ref}`sec_tutorial_metadata_binary`.

(sec_metadata_schema_schema)=

## Full metaschema

The schema for metadata schemas is formally defined using
[JSON Schema](http://json-schema.org/) and given in full here. Any schema passed to
{class}`tskit.MetadataSchema` is validated against this metaschema.

```{eval-rst}
.. literalinclude:: ../python/tskit/metadata_schema.schema.json
    :language: json
```


--- ../../tskit/docs/work_with_metadata.md ---

---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.12
    jupytext_version: 1.9.1
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---
```{currentmodule} tskit
```


(sec_tutorial_metadata)=

# Working with Metadata

Metadata is information associated with entities that {program}`tskit` doesn't use or
interpret, but which is useful to pass on to downstream analysis such as sample ids,
dates etc. (see {ref}`sec_metadata` for a full discussion). Each
{ref}`table<sec_tables_api_table>` has a {class}`MetadataSchema` which details the
contents and encoding of the metadata for each row. A metadata schema is a JSON document
that conforms to [JSON Schema](https://json-schema.org/understanding-json-schema/)
(The full schema for tskit is at {ref}`sec_metadata_schema_schema`). Here we use an
{ref}`example tree sequence<sec_intro_downloading_datafiles>`
which contains some demonstration metadata:

```{code-cell} ipython3
:tags: [remove-cell]
import msprime
import tskit

def metadata():
    tables = msprime.sim_ancestry(4).dump_tables()
    tables.individuals.metadata_schema = tskit.MetadataSchema(
    {'additionalProperties': False,
     'codec': 'json',
     'properties': {'accession': {'description': 'ENA accession number',
                                  'type': 'string'},
                    'pcr': {'description': 'Was PCR used on this sample',
                            'name': 'PCR Used',
                            'type': 'boolean'}},
     'required': ['accession', 'pcr'],
     'type': 'object'}
    )
    md = [
        {'accession': 'ERS0001', 'pcr': True},
        {'accession': 'ERS0002', 'pcr': True},
        {'accession': 'ERS0003', 'pcr': True},
        {'accession': 'ERS0004', 'pcr': False},
    ]
    table = tables.individuals
    copy = table.copy()
    table.clear()
    for m, row in zip(md, copy):
        table.append(row.replace(metadata=m))
    ts = tables.tree_sequence()
    ts.dump("data/metadata.trees")

def create_notebook_data():
    metadata()

# create_notebook_data()  # uncomment to recreate the tree seqs used in this notebook
```


```{code-cell} ipython3
import tskit
import json

ts = tskit.load("data/metadata.trees")
```
(sec_tutorial_metadata_reading)=

## Reading metadata and schemas

Metadata is automatically decoded using the schema when accessed via a
{class}`TreeSequence` or {class}`TableCollection` Python API. For example:

```{code-cell} ipython3
print("Metadata for individual 0:", ts.individual(0).metadata)  # Tree sequence access
print("Metadata for individual 0:", ts.tables.individuals[0].metadata)  # Table access
```

Viewing the {class}`MetadataSchema` for a table can help with understanding
its metadata, as it can contain descriptions and constraints:

```{code-cell} ipython3
ts.table_metadata_schemas.individual
```

The same schema can be accessed via a {attr}`~IndividualTable.metadata_schema` attribute
on each table (printed prettily here using ``json.dumps``)
 
```{code-cell} ipython3
schema = ts.tables.individuals.metadata_schema
print(json.dumps(schema.asdict(), indent=4))  # Print with indentations
```

The top-level metadata and schemas for the entire tree sequence are similarly
accessed with {attr}`TreeSequence.metadata` and {attr}`TreeSequence.metadata_schema`.

:::{note}
If there is no schema (i.e. it is equal to ``MetadataSchema(None)``) for a table
or top-level metadata, then no decoding is performed and ``bytes`` will be returned.
:::

(sec_tutorial_metadata_modifying)=

## Modifying metadata and schemas

If you are creating or modifying a tree sequence by changing the underlying tables,
you may want to record or add to the metadata. If the change fits into the same schema,
this is relatively simple, you can follow the
{ref}`description of minor table edits<sec_tables_editing_minor>` in the
{ref}`sec_tables` tutorial. However if it requires a change to the schema, this must be
done first, as it is then used to validate and encode the metadata. 

Schemas in tskit are held in a {class}`MetadataSchema`.
A Python dict representation of the schema is passed to its constructor, which
will validate the schema. Here are a few examples: the first one allows arbitrary fields
to be added, the second one (which will construct the schema we printed above) does not:

```{code-cell} ipython3
basic_schema = tskit.MetadataSchema({'codec': 'json'})

complex_schema = tskit.MetadataSchema({
    'codec': 'json',
    'additionalProperties': False,
    'properties': {'accession': {'description': 'ENA accession number',
                                 'type': 'string'},
                   'pcr': {'description': 'Was PCR used on this sample',
                           'name': 'PCR Used',
                           'type': 'boolean'}},
    'required': ['accession', 'pcr'],
    'type': 'object',
})
```

This {class}`MetadataSchema` can then be assigned to a table or the top-level
tree sequence e.g. {attr}`~IndividualTable.metadata_schema`:

```{code-cell} ipython3
tables = tskit.TableCollection(sequence_length=1)  # make a new, empty set of tables
tables.individuals.metadata_schema = complex_schema
```

This will overwrite any existing schema. Note that this will not validate any existing
metadata against the new schema. Now that the table has a schema, calls to
{meth}`~IndividualTable.add_row` will validate and encode the metadata:

```{code-cell} ipython3
row_id = tables.individuals.add_row(0, metadata={"accession": "Bob1234", "pcr": True})
print(f"Row {row_id} added to the individuals table")
```

If we try to add metadata that doesn't fit the schema, such as accidentally using a
string instead of a proper Python boolean, we'll get an error:

```{code-cell} ipython3
:tags: [raises-exception, output_scroll]
tables.individuals.add_row(0, metadata={"accession": "Bob1234", "pcr": "false"})
```

and because we set ``additionalProperties`` to ``False`` in the schema, an error is
also raised if we attempt to add new fields:

```{code-cell} ipython3
:tags: [raises-exception, output_scroll]
tables.individuals.add_row(0, metadata={"accession": "Bob1234", "pcr": True, "newKey": 25})
```


To set the top-level metadata, just assign it. Validation and encoding happen as
specified by the top-level metadata schema

```{code-cell} ipython3
tables.metadata_schema = basic_schema  # Allows new fields to be added that are not validated
tables.metadata = {"mean_coverage": 200.5}
print(tables.metadata)
```

:::{note}
*Provenance* information, detailing the origin of the data, modification timestamps,
and (ideally) how the tree sequence can be reconstructed, should go in
{ref}`sec_provenance`, not metadata.
:::

To modify a schema --- for example to add a key --- first get the dict representation,
modify, then write back:

```{code-cell} ipython3
schema_dict = tables.individuals.metadata_schema.schema
schema_dict["properties"]["newKey"] = {"type": "integer"}
tables.individuals.metadata_schema = tskit.MetadataSchema(schema_dict)
# Now this will work:
new_id = tables.individuals.add_row(metadata={'accession': 'abc123', 'pcr': False, 'newKey': 25})
print(tables.individuals[new_id].metadata)
```

To modify the metadata of rows in tables use the {ref}`sec_tutorial_metadata_bulk`.

(sec_tutorial_metadata_viewing_raw)=

## Viewing raw metadata

If you need to see the raw (i.e. bytes) metadata, you just need to remove the
schema, for instance:

```{code-cell} ipython3
individual_table = tables.individuals.copy()  # don't change the original tables.individual

print("Metadata:\n", individual_table[0].metadata)

individual_table.metadata_schema = tskit.MetadataSchema(None)
print("\nRaw metadata:\n", individual_table[0].metadata)
```

(sec_tutorial_metadata_bulk)=

## Metadata for bulk table methods

In the interests of efficiency each table's {meth}`~NodeTable.packset_metadata` method,
as well as the more general {meth}`~NodeTable.set_columns` and
{meth}`~NodeTable.append_columns` methods, do not attempt to validate or encode metadata.
You can call {meth}`MetadataSchema.validate_and_encode_row` directly to prepare metadata
for these methods:

```{code-cell} ipython3
metadata_column = [
    {"accession": "etho1234", "pcr": True},
    {"accession": "richard1235", "pcr": False},
    {"accession": "albert1236", "pcr": True},
]
encoded_metadata_column = [
    tables.individuals.metadata_schema.validate_and_encode_row(r) for r in metadata_column
]
md, md_offset = tskit.pack_bytes(encoded_metadata_column)
tables.individuals.set_columns(flags=[0, 0, 0], metadata=md, metadata_offset=md_offset)
tables.individuals
```

Or if all columns do not need to be set:

```{code-cell} ipython3
tables.individuals.packset_metadata(
    [tables.individuals.metadata_schema.validate_and_encode_row(r) for r in metadata_column]
)
```

(sec_tutorial_metadata_binary)=

## Binary metadata

To disable the validation and encoding of metadata and store raw bytes pass ``None`` to
{class}`MetadataSchema`

```{code-cell} ipython3
tables.populations.metadata_schema = tskit.MetadataSchema(None)
tables.populations.add_row(metadata=b"SOME CUSTOM BYTES #!@")
print(tables.populations[0].metadata)
```


--- ../../tskit/.github/PULL_REQUEST_TEMPLATE.md ---

## Description

Thanks for contributing to tskit! :heart:
A guide to the PR process is [here](https://tskit.readthedocs.io/en/latest/development.html#development_workflow_git)
Please replace this text with a summary of the change and which issue is fixed, if any. Please also include relevant motivation and context.

Fixes #(issue) <- Putting the issue number here will auto-close the issue when this PR is merged 

# PR Checklist:

- [ ] Tests that fully cover new/changed functionality.
- [ ] Documentation including tutorial content if appropriate.
- [ ] Changelogs, if there are API changes.
