#
# Copyright (C) 2015-2021 University of Oxford
#
# This file is part of msprime.
#
# msprime is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# msprime is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with msprime.  If not, see <http://www.gnu.org/licenses/>.
#
"""
Test cases for basic ancestry simulation operations.
"""

import collections
import datetime
import json
import logging
import math
import random
import warnings

import numpy as np
import pytest
import tskit

import msprime
from msprime import _msprime, ancestry


def tree_sequences_equal(ts1, ts2):
    """
    Returns True if the specified tree sequences are equal, ignoring
    their provenances.
    """
    t1 = ts1.dump_tables()
    t2 = ts2.dump_tables()
    t1.provenances.clear()
    t2.provenances.clear()
    return t1 == t2


def has_discrete_genome(ts):
    """
    Returns True if the specified tree sequence has discrete genome coordinates.
    """
    tables = ts.tables
    edges_left = np.all(tables.edges.left == np.floor(tables.edges.left))
    edges_right = np.all(tables.edges.right == np.floor(tables.edges.right))
    migrations_left = np.all(tables.migrations.left == np.floor(tables.migrations.left))
    migrations_right = np.all(
        tables.migrations.right == np.floor(tables.migrations.right)
    )
    sites = np.all(tables.sites.position == np.floor(tables.sites.position))
    return edges_left and edges_right and migrations_left and migrations_right and sites


def get_bottleneck_examples():
    """
    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),
    ]
    for n in [3, 10, 100]:
        ts = msprime.simulate(
            n,
            length=100,
            recombination_rate=1,
            demographic_events=bottlenecks,
            random_seed=n,
        )
        yield ts


class TestFullArg:
    """
    Tests for recording the full ARG.
    """

    def verify(self, sim, multiple_mergers=False):
        sim.random_generator.seed = 1234
        sim.run()
        tree_sequence = next(sim.run_replicates(1))
        # Check if we have multiple merger somewhere.
        found = False
        for edgeset in tree_sequence.edgesets():
            if len(edgeset.children) > 2:
                found = True
                break
        assert multiple_mergers == found

        flags = tree_sequence.tables.nodes.flags
        time = tree_sequence.tables.nodes.time
        # TODO add checks for migrations.
        re_nodes = np.where(flags == msprime.NODE_IS_RE_EVENT)[0]
        ca_nodes = np.where(flags == msprime.NODE_IS_CA_EVENT)[0]
        gc_nodes = np.where(flags == msprime.NODE_IS_GC_EVENT)[0]
        coal_nodes = np.where(flags == 0)[0]
        # There should be two recombination nodes for every event
        assert np.array_equal(time[re_nodes[::2]], time[re_nodes[1::2]])  # Odd indexes
        assert re_nodes.shape[0] / 2 == sim.num_recombination_events
        # There should be two gene conversion nodes for every effective event
        assert np.array_equal(time[gc_nodes[::2]], time[gc_nodes[1::2]])  # Odd indexes
        assert (
            gc_nodes.shape[0] / 2
            == sim.num_internal_gene_conversion_events
            - sim.num_noneffective_gene_conversion_events
        )
        if not multiple_mergers:
            assert (
                ca_nodes.shape[0] + coal_nodes.shape[0] == sim.num_common_ancestor_events
            )
        else:
            assert ca_nodes.shape[0] > 0
        # After simplification, all the RE, GC, and CA nodes should be gone.
        ts_simplified = tree_sequence.simplify()
        new_flags = ts_simplified.tables.nodes.flags
        new_time = ts_simplified.tables.nodes.time
        assert np.sum(new_flags == msprime.NODE_IS_RE_EVENT) == 0
        assert np.sum(new_flags == msprime.NODE_IS_CA_EVENT) == 0
        assert np.sum(new_flags == msprime.NODE_IS_GC_EVENT) == 0
        # All coal nodes from the original should be identical to the originals
        assert np.array_equal(time[coal_nodes], new_time[new_flags == 0])
        assert ts_simplified.num_nodes <= tree_sequence.num_nodes
        assert ts_simplified.num_edges <= tree_sequence.num_edges
        return tree_sequence

    def test_no_recombination(self):
        sim = ancestry._parse_simulate(10, record_full_arg=True)
        ts = self.verify(sim)
        ts_simplified = ts.simplify()
        t1 = ts.tables
        t2 = ts_simplified.tables
        assert t1.nodes == t2.nodes
        assert t1.edges == t2.edges

    def test_recombination_n25(self):
        sim = ancestry._parse_simulate(
            25,
            recombination_rate=1,
            record_full_arg=True,
        )
        self.verify(sim)

    def test_recombination_n5(self):
        sim = ancestry._parse_simulate(
            5,
            recombination_rate=10,
            record_full_arg=True,
        )
        self.verify(sim)

    def test_recombination_n50(self):
        sim = ancestry._parse_simulate(
            50,
            recombination_rate=2,
            record_full_arg=True,
        )
        self.verify(sim)

    def test_recombination_n100(self):
        sim = ancestry._parse_simulate(
            100,
            recombination_rate=0.2,
            record_full_arg=True,
        )
        self.verify(sim)

    def test_multimerger(self):
        sim = ancestry._parse_simulate(
            100,
            recombination_rate=10,
            record_full_arg=True,
            demographic_events=[
                msprime.InstantaneousBottleneck(time=0.1, population=0, strength=5)
            ],
        )
        self.verify(sim, multiple_mergers=True)

    def test_multimerger_beta(self):
        sim = ancestry._parse_sim_ancestry(
            100,
            recombination_rate=0.1,
            record_full_arg=True,
            sequence_length=10,
            model=msprime.BetaCoalescent(alpha=1.1),
        )
        self.verify(sim, multiple_mergers=True)

    def test_multimerger_dirac(self):
        sim = ancestry._parse_sim_ancestry(
            100,
            recombination_rate=0.1,
            record_full_arg=True,
            sequence_length=10,
            model=msprime.DiracCoalescent(psi=0.5, c=100),
        )
        self.verify(sim, multiple_mergers=True)

    def test_gene_conversion_n25(self):
        sim = ancestry._parse_sim_ancestry(
            25,
            gene_conversion_rate=1,
            gene_conversion_tract_length=1,
            sequence_length=10,
            record_full_arg=True,
        )
        self.verify(sim)

    def test_gene_conversion_n5(self):
        sim = ancestry._parse_sim_ancestry(
            5,
            gene_conversion_rate=10,
            gene_conversion_tract_length=1.5,
            sequence_length=10,
            record_full_arg=True,
        )
        self.verify(sim)

    def test_gene_conversion_n50(self):
        sim = ancestry._parse_sim_ancestry(
            50,
            gene_conversion_rate=2,
            gene_conversion_tract_length=1,
            sequence_length=5,
            record_full_arg=True,
        )
        self.verify(sim)

    def test_gene_conversion_n100(self):
        sim = ancestry._parse_sim_ancestry(
            100,
            gene_conversion_rate=0.2,
            gene_conversion_tract_length=2,
            sequence_length=10,
            record_full_arg=True,
        )
        self.verify(sim)


class TestStoreUnary:
    """
    Tests for recording unary nodes.
    """

    def verify_store_unary(self, ts):
        self.verify_min_max_children(ts)
        tss = ts.simplify()
        self.verify_simple_stats(ts, tss)

    def verify_min_max_children(self, ts):
        min_children = np.zeros(ts.num_nodes, dtype=int)
        max_children = np.zeros_like(min_children)

        for tree in ts.trees():
            for i in range(ts.num_nodes):
                n = tree.num_children_array[i]
                if n > 0:
                    if min_children[i] > 0:
                        min_children[i] = min(min_children[i], n)
                    else:
                        min_children[i] = n
                    max_children[i] = max(max_children[i], n)

        assert np.any(min_children == 1)
        for minc, maxc in zip(min_children, max_children):
            if minc == 1:
                assert maxc >= 2

    def verify_simple_stats(self, ts, ts_simplified):
        assert ts.num_samples == ts_simplified.num_samples
        assert ts.num_nodes == ts_simplified.num_nodes

    def test_unary(self):
        for n in [25, 50, 100]:
            ts = msprime.sim_ancestry(
                samples=n,
                sequence_length=100,
                recombination_rate=0.1,
                coalescing_segments_only=False,
            )
            self.verify_store_unary(ts)

    def test_unary_smc(self):
        for n in [25, 50, 100]:
            ts = msprime.sim_ancestry(
                samples=n,
                sequence_length=100,
                recombination_rate=0.1,
                coalescing_segments_only=False,
                model="smc",
            )
            self.verify_store_unary(ts)

    def test_unary_smc_prime(self):
        for n in [25, 50, 100]:
            ts = msprime.sim_ancestry(
                samples=n,
                sequence_length=100,
                recombination_rate=0.1,
                coalescing_segments_only=False,
                model="smc_prime",
            )
            self.verify_store_unary(ts)

    def test_unary_sweep_coalescence(self):
        N = 1e6
        model = msprime.SweepGenicSelection(
            position=0.5,
            start_frequency=1.0 / (2 * N),
            end_frequency=1.0 - (1.0 / (2 * N)),
            s=100,
            dt=1e-6,
        )
        ts = msprime.sim_ancestry(
            10,
            recombination_rate=1e-4,
            model=[model, "hudson"],
            population_size=N,
            sequence_length=100,
            coalescing_segments_only=False,
            random_seed=6974,
        )
        assert ts.num_trees > 1
        assert all(tree.num_roots == 1 for tree in ts.trees())
        self.verify_store_unary(ts)

    def test_multimerger(self):
        demography = msprime.Demography()
        demography.add_population(initial_size=10_000)
        demography.add_instantaneous_bottleneck(time=0.1, population=0, strength=5)
        ts = msprime.sim_ancestry(
            samples=20,
            sequence_length=100,
            recombination_rate=0.1,
            coalescing_segments_only=False,
            demography=demography,
        )
        self.verify_store_unary(ts)

    def test_unary_gene_conversion(self):
        rate = 100
        sim = ancestry._parse_sim_ancestry(
            10,
            sequence_length=100,
            gene_conversion_rate=rate,
            gene_conversion_tract_length=5,
            coalescing_segments_only=False,
        )
        sim.run()
        assert sim.num_gene_conversion_events > 1
        ts = tskit.TableCollection.fromdict(sim.tables.asdict()).tree_sequence()
        self.verify_store_unary(ts)

    def test_no_recombination(self):
        # no recombination -> no unary nodes
        ts = msprime.sim_ancestry(
            samples=25, sequence_length=100, coalescing_segments_only=False
        )
        tss = ts.simplify()
        assert ts.equals(tss, ignore_provenance=True)

    def test_predefined_scenario(self):
        # initial state
        # 1.50┊ 5  6  ┊ 5  7  ┊
        #     ┊ ┃  ┃  ┊ ┃  ┃  ┊
        # 0.70┊ ┃  ┃  ┊ ┃  4  ┊
        #     ┊ ┃  ┃  ┊ ┃ ┏┻┓ ┊
        # 0.50┊ ┃  3  ┊ ┃ ┃ ┃ ┊
        #     ┊ ┃ ┏┻┓ ┊ ┃ ┃ ┃ ┊
        # 0.00┊ 0 1 2 ┊ 0 1 2 ┊
        #     0       5      10
        # completed simulation
        #     ┊       ┊   9   ┊
        #     ┊       ┊ ┏━┻┓  ┊
        #     ┊   8   ┊ 8  ┃  ┊
        #     ┊ ┏━┻┓  ┊ ┃  ┃  ┊
        #     ┊ 5  6  ┊ 5  7  ┊
        #     ┊ ┃  ┃  ┊ ┃  ┃  ┊
        #     ┊ ┃  ┃  ┊ ┃  4  ┊
        #     ┊ ┃  ┃  ┊ ┃ ┏┻┓ ┊
        #     ┊ ┃  3  ┊ ┃ ┃ ┃ ┊
        #     ┊ ┃ ┏┻┓ ┊ ┃ ┃ ┃ ┊
        #     ┊ 0 1 2 ┊ 0 1 2 ┊
        #     0       5      10

        tables = tskit.TableCollection(sequence_length=10)
        tables.time_units = "generations"
        tables.nodes.set_columns(
            time=np.array([0, 0, 0, 0.5, 0.7, 1.5, 1.5, 1.5]),
            flags=np.array([1, 1, 1, 0, 0, 0, 0, 0], dtype="uint32"),
            population=np.array([0, 0, 0, 0, 0, 1, 1, 0], dtype="int32"),
        )
        tables.edges.set_columns(
            left=np.array([0, 0, 5, 5, 0, 0, 5]),
            right=np.array([5, 5, 10, 10, 10, 5, 10]),
            parent=np.array([3, 3, 4, 4, 5, 6, 7], dtype="int32"),
            child=np.array([1, 2, 1, 2, 0, 3, 4], dtype="int32"),
        )
        tables.populations.add_row()
        tables.populations.add_row()
        demography = msprime.Demography()
        demography.add_population(initial_size=100)
        demography.add_population(initial_size=100)
        demography.add_mass_migration(1e6, source=1, dest=0, proportion=1.0)
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            simulator = msprime.ancestry._parse_sim_ancestry(
                initial_state=tables,
                demography=demography,
                coalescing_segments_only=False,
            )
        simulator._run_until(event_chunk=1, end_time=math.inf)
        new_tables = tskit.TableCollection.fromdict(simulator.tables.asdict())
        ts = new_tables.tree_sequence()
        tree = ts.first()
        left_root = tree.root
        tree.next()
        assert tree.parent_array[left_root] == tree.root


class TestAdditionalNodes:
    """
    Tests for recording the full ARG.
    """

    def verify(self, sim, multiple_mergers=False, additional_nodes=None):
        sim.random_generator.seed = 21
        sim.run()
        tree_sequence = next(sim.run_replicates(1))
        # Check if we have multiple merger somewhere.
        found = False
        for edgeset in tree_sequence.edgesets():
            if len(edgeset.children) > 2:
                found = True
                break
        assert multiple_mergers == found

        flags = tree_sequence.tables.nodes.flags
        time = tree_sequence.tables.nodes.time
        re_nodes = np.where(np.bitwise_and(flags, msprime.NODE_IS_RE_EVENT))[0]
        ca_nodes = np.where(np.bitwise_and(flags, msprime.NODE_IS_CA_EVENT))[0]
        gc_nodes = np.where(np.bitwise_and(flags, msprime.NODE_IS_GC_EVENT))[0]
        mig_nodes = np.where(np.bitwise_and(flags, msprime.NODE_IS_MIG_EVENT))[0]
        coal_nodes = np.where(flags == 0)[0]

        if additional_nodes is None:
            additional_nodes = msprime.NodeType(0)
        for name, member in msprime.NodeType.__members__.items():
            value = (member & additional_nodes).value
            if value > 0:
                if name == "RECOMBINANT":
                    # There should be two recombination nodes for every event
                    assert np.array_equal(
                        time[re_nodes[::2]], time[re_nodes[1::2]]
                    )  # Odd indexes
                    assert re_nodes.shape[0] / 2 == sim.num_recombination_events
                elif name == "GENE_CONVERSION":
                    # There should be two gene conversion nodes for every effective event
                    assert np.array_equal(
                        time[gc_nodes[::2]], time[gc_nodes[1::2]]
                    )  # Odd indexes
                    assert (
                        gc_nodes.shape[0] / 2
                        == sim.num_internal_gene_conversion_events
                        - sim.num_noneffective_gene_conversion_events
                    )
                elif name == "COMMON_ANCESTOR":
                    if not multiple_mergers:
                        assert (
                            ca_nodes.shape[0] + coal_nodes.shape[0]
                            == sim.num_common_ancestor_events
                        )
                    else:
                        assert ca_nodes.shape[0] > 0
                elif name == "MIGRANT":
                    assert np.sum(mig_nodes) > 0
                elif name == "PASS_THROUGH":
                    pass
                else:
                    raise (ValueError)

        # After simplification, all the RE, GC, and CA nodes should be gone.
        ts_simplified = tree_sequence.simplify()
        new_flags = ts_simplified.tables.nodes.flags
        new_time = ts_simplified.tables.nodes.time
        assert np.sum(new_flags == msprime.NODE_IS_RE_EVENT) == 0
        assert np.sum(new_flags == msprime.NODE_IS_CA_EVENT) == 0
        assert np.sum(new_flags == msprime.NODE_IS_GC_EVENT) == 0
        # All coal nodes from the original should be identical to the originals
        assert np.array_equal(time[coal_nodes], new_time[new_flags == 0])
        assert ts_simplified.num_nodes <= tree_sequence.num_nodes
        assert ts_simplified.num_edges <= tree_sequence.num_edges
        return tree_sequence

    def verify_dtwf(self, sim, additional_nodes):
        # additional checks needed here
        sim.random_generator.seed = 21
        sim.run()
        ts = next(sim.run_replicates(1))

        flags = ts.tables.nodes.flags
        assert np.all(flags[: ts.num_samples] == 1)
        time = ts.tables.nodes.time
        # nodes can be multiple things at the same time
        ca_nodes = np.where(np.bitwise_and(flags, msprime.NODE_IS_CA_EVENT))[0]
        re_nodes = np.where(np.bitwise_and(flags, msprime.NODE_IS_RE_EVENT))[0]
        pt_nodes = np.where(np.bitwise_and(flags, msprime.NODE_IS_PASS_THROUGH))[0]
        for name, member in msprime.NodeType.__members__.items():
            value = (member & additional_nodes).value
            if value > 0:
                if name == "RECOMBINANT":
                    assert np.array_equal(time[re_nodes[::2]], time[re_nodes[1::2]])
                    self.verify_dtwf_recombination(ts.tables)
                elif name == "COMMON_ANCESTOR":
                    for tree in ts.trees():
                        assert np.all(tree.num_children_array[ca_nodes] <= 1)
                elif name == "PASS_THROUGH":
                    for tree in ts.trees():
                        assert np.all(tree.num_children_array[pt_nodes] <= 1)
                else:
                    raise (ValueError)

        return ts

    def verify_dtwf_pass_through(self, ts):
        for tree in ts.trees():
            queue = [tree.root]
            while queue:
                parent = queue.pop(0)
                time = tree.time(parent) - 1
                num_children = 0
                for child in tree.children(parent):
                    assert tree.time(child) == time
                    queue.append(child)
                    num_children += 1
                if num_children == 0:
                    assert tree.time(parent) == 0

    def verify_dtwf_recombination(self, tables):
        children = collections.defaultdict(set)
        for edge in tables.edges:
            children[edge.child].add(edge.parent)
        for parents in children.values():
            num_parents = len(parents)
            if num_parents == 2:
                for parent in parents:
                    assert tables.nodes[parent].flags & msprime.NODE_IS_RE_EVENT
            else:
                assert num_parents == 1

    def test_no_recombination(self):
        node_type = msprime.NodeType.COMMON_ANCESTOR
        sim = ancestry._parse_sim_ancestry(
            10, additional_nodes=node_type, coalescing_segments_only=False
        )
        ts = self.verify(sim)
        ts_simplified = ts.simplify()
        t1 = ts.tables
        t2 = ts_simplified.tables
        assert t1.nodes == t2.nodes
        assert t1.edges == t2.edges

    def test_recombination_re(self):
        node_type = msprime.NodeType.RECOMBINANT
        sim = ancestry._parse_sim_ancestry(
            25,
            recombination_rate=1,
            sequence_length=1,
            discrete_genome=False,
            additional_nodes=node_type,
            coalescing_segments_only=False,
        )
        self.verify(sim, additional_nodes=node_type)

    def test_recombination_ca(self):
        node_type = msprime.NodeType.COMMON_ANCESTOR
        sim = ancestry._parse_sim_ancestry(
            5,
            discrete_genome=False,
            sequence_length=1,
            recombination_rate=10,
            additional_nodes=node_type,
            coalescing_segments_only=False,
        )
        self.verify(sim, additional_nodes=node_type)

    def test_recombination_mig(self):
        node_type = msprime.NodeType.MIGRANT
        demography = msprime.Demography()
        demography.add_population(name="A", initial_size=5)
        demography.add_population(name="B", initial_size=5)
        demography.set_symmetric_migration_rate(populations=("A", "B"), rate=0.1)
        sim = ancestry._parse_sim_ancestry(
            samples={0: 2, 1: 2},
            discrete_genome=False,
            sequence_length=1,
            recombination_rate=10,
            additional_nodes=node_type,
            coalescing_segments_only=False,
            demography=demography,
        )
        self.verify(sim, additional_nodes=node_type)

    def test_recombination_re_ca(self):
        node_type = msprime.NodeType.COMMON_ANCESTOR | msprime.NodeType.RECOMBINANT
        sim = ancestry._parse_sim_ancestry(
            5,
            discrete_genome=False,
            sequence_length=1,
            recombination_rate=10,
            additional_nodes=node_type,
            coalescing_segments_only=False,
        )
        self.verify(sim, additional_nodes=node_type)

    def test_recombination_ca_mig(self):
        node_type = msprime.NodeType.COMMON_ANCESTOR | msprime.NodeType.MIGRANT
        demography = msprime.Demography()
        demography.add_population(name="A", initial_size=5)
        demography.add_population(name="B", initial_size=5)
        demography.set_symmetric_migration_rate(populations=("A", "B"), rate=0.1)
        sim = ancestry._parse_sim_ancestry(
            samples={0: 2, 1: 2},
            recombination_rate=10,
            sequence_length=1,
            discrete_genome=False,
            additional_nodes=node_type,
            coalescing_segments_only=False,
            demography=demography,
        )
        self.verify(sim, additional_nodes=node_type)

    def test_recombination_re_mig(self):
        node_type = msprime.NodeType.RECOMBINANT | msprime.NodeType.MIGRANT
        demography = msprime.Demography()
        demography.add_population(name="A", initial_size=1)
        demography.add_population(name="B", initial_size=1)
        demography.set_symmetric_migration_rate(populations=("A", "B"), rate=0.1)
        sim = ancestry._parse_sim_ancestry(
            samples=[
                msprime.SampleSet(2, population="A"),
                msprime.SampleSet(2, population="B"),
            ],
            sequence_length=1,
            discrete_genome=False,
            recombination_rate=10,
            additional_nodes=node_type,
            coalescing_segments_only=False,
            demography=demography,
        )
        self.verify(sim, additional_nodes=node_type)

    def test_multimerger(self):
        node_type = msprime.NodeType.RECOMBINANT
        demography = msprime.Demography()
        demography.add_population(initial_size=1000)
        demography.add_instantaneous_bottleneck(time=0.1, population=0, strength=5)
        sim = ancestry._parse_sim_ancestry(
            100,
            sequence_length=1,
            discrete_genome=False,
            recombination_rate=0.010,
            additional_nodes=node_type,
            coalescing_segments_only=False,
            demography=demography,
        )
        self.verify(sim, multiple_mergers=True, additional_nodes=node_type)

    def test_multimerger_beta(self):
        node_type = msprime.NodeType.RECOMBINANT
        sim = ancestry._parse_sim_ancestry(
            100,
            recombination_rate=0.1,
            additional_nodes=node_type,
            coalescing_segments_only=False,
            sequence_length=10,
            model=msprime.BetaCoalescent(alpha=1.1),
        )
        self.verify(sim, multiple_mergers=True, additional_nodes=node_type)

    def test_multimerger_dirac(self):
        node_type = msprime.NodeType.RECOMBINANT
        sim = ancestry._parse_sim_ancestry(
            100,
            recombination_rate=0.1,
            additional_nodes=node_type,
            coalescing_segments_only=False,
            sequence_length=10,
            model=msprime.DiracCoalescent(psi=0.5, c=5),
        )
        self.verify(sim, multiple_mergers=True, additional_nodes=node_type)

    def test_gene_conversion(self):
        node_type = msprime.NodeType.GENE_CONVERSION
        sim = ancestry._parse_sim_ancestry(
            25,
            gene_conversion_rate=1,
            gene_conversion_tract_length=1,
            sequence_length=10,
            additional_nodes=node_type,
            coalescing_segments_only=False,
        )
        self.verify(sim, additional_nodes=node_type)

    def test_dtwf_mig(self):
        node_type = msprime.NodeType.MIGRANT
        demography = msprime.Demography()
        demography.add_population(name="A", initial_size=5)
        demography.add_population(name="B", initial_size=5)
        demography.set_symmetric_migration_rate(populations=("A", "B"), rate=0.1)
        with pytest.raises(
            ValueError,
            match="Recording MIGRANT nodes is currently"
            " not supported in DTWF simulation.",
        ):
            _ = ancestry._parse_sim_ancestry(
                samples={0: 2, 1: 2},
                discrete_genome=False,
                sequence_length=1,
                recombination_rate=10,
                additional_nodes=node_type,
                coalescing_segments_only=False,
                demography=demography,
                model="dtwf",
            )

    def test_dtwf_full_arg(self):
        with pytest.raises(
            ValueError,
            match="Full ARG recording not supported in FixedPedigree and "
            "DTWF simulation",
        ):
            _ = msprime.sim_ancestry(5, model="dtwf", record_full_arg=True)
        with pytest.raises(
            ValueError,
            match="Full ARG recording not supported in FixedPedigree and "
            "DTWF simulation",
        ):
            _ = msprime.simulate(
                5, model=msprime.DiscreteTimeWrightFisher(), record_full_arg=True
            )

    def test_dtwf_store_additional_nodes(self):
        node_values = [
            msprime.NODE_IS_RE_EVENT | msprime.NODE_IS_PASS_THROUGH,
            msprime.NODE_IS_CA_EVENT | msprime.NODE_IS_PASS_THROUGH,
            None,
        ]
        node_values[-1] = node_values[0] | node_values[1]
        for node_value in node_values:
            additional_nodes = msprime.NodeType(node_value)
            sim = ancestry._parse_sim_ancestry(
                10,
                population_size=100,
                model="dtwf",
                ploidy=2,
                recombination_rate=1e-3,
                random_seed=1234,
                sequence_length=100,
                additional_nodes=additional_nodes,
                coalescing_segments_only=False,
            )
            ts = self.verify_dtwf(sim, additional_nodes)
            if (additional_nodes.value & msprime.NODE_IS_PASS_THROUGH) and (
                additional_nodes.value & msprime.NODE_IS_CA_EVENT
            ):
                self.verify_dtwf_pass_through(ts)


class TestSimulator:
    """
    Runs tests on the underlying Simulator object.
    """

    def verify_simulation(self, n, m, r):
        """
        Verifies a simulation for the specified parameters.
        """
        sim = ancestry._parse_sim_ancestry(
            n,
            sequence_length=m,
            recombination_rate=r,
            random_seed=32,
        )
        sim.run()
        assert sim.num_breakpoints == len(sim.breakpoints)
        assert sim.time > 0
        assert sim.num_avl_node_blocks > 0
        assert sim.num_segment_blocks > 0
        assert sim.num_node_mapping_blocks > 0
        tree_sequence = next(sim.run_replicates(1))
        t = 0.0
        for record in tree_sequence.nodes():
            if record.time > t:
                t = record.time
        assert sim.time == t
        assert sim.num_common_ancestor_events > 0
        assert sim.num_recombination_events >= 0
        assert np.sum(sim.num_migration_events) >= 0
        assert sim.num_multiple_recombination_events >= 0

    def test_random_parameters(self):
        num_random_sims = 10
        for _ in range(num_random_sims):
            n = random.randint(2, 100)
            m = random.randint(10, 100)
            r = random.random()
            self.verify_simulation(n, m, r)

    def test_perf_parameters(self):
        sim = ancestry._parse_simulate(10)
        sim.run()
        assert sim.avl_node_block_size > 0
        assert sim.segment_block_size > 0
        assert sim.node_mapping_block_size > 0

    def test_event_chunk(self):
        sim = ancestry._parse_simulate(10)
        for bad_chunk in [-(2**32), -1, 0]:
            with pytest.raises(ValueError):
                sim.run(event_chunk=bad_chunk)
        sim.reset()
        sim.run(event_chunk=2**32 + 1)
        sim.reset()
        sim.run(event_chunk=2**64 + 1)

    def test_debug_func(self):
        sim = ancestry._parse_simulate(10)
        count = 0

        def f(sim):
            nonlocal count
            count += 1

        sim.run(event_chunk=1, debug_func=f)
        assert count > 0

    def test_info_logging(self, caplog):
        sim = ancestry._parse_simulate(10)
        sim.random_generator.seed = 42
        with caplog.at_level(logging.INFO):
            sim.run()
        assert len(caplog.records) == 3
        assert caplog.messages[0].startswith("model[0] {'name': 'hudson'}")
        assert (
            caplog.messages[1] == "Running model {'name': 'hudson'} until max time: inf"
        )
        assert caplog.messages[-1].startswith("Completed at time")

    def test_debug_logging(self, caplog):
        sim = ancestry._parse_simulate(3)
        sim.random_generator.seed = 42
        with caplog.at_level(logging.DEBUG):
            sim.run(event_chunk=1)
        assert len(caplog.records) == 6
        assert caplog.messages[0].startswith("model[0] {'name': 'hudson'}")
        assert caplog.messages[-1].startswith("Completed at time")
        assert caplog.messages[2].startswith("time=0.312845 ancestors=2 ret=")

    def test_debug_logging_dtwf(self, caplog):
        sim = ancestry._parse_simulate(3, Ne=10, model="dtwf")
        sim.random_generator.seed = 42
        with caplog.at_level(logging.DEBUG):
            sim.run(event_chunk=1)
            assert len(caplog.records) == 8
            assert (
                caplog.messages[1] == "Running model {'name': 'dtwf'}"
                " until max time: inf"
            )
            assert caplog.messages[-1].startswith("Completed at time")
            assert caplog.messages[2].startswith("time=1 ancestors=3 ret=")

    def test_str(self):
        sim = ancestry._parse_simulate(3)
        s = str(sim)
        assert len(s) > 0


class TestSimulateAfterLocalMRCA:
    def check_roots(self, ts, end_time=None, allow_multiple_roots=False):
        if allow_multiple_roots:
            root_times = [tree.time(tree.roots[0]) for tree in ts.trees()]
        else:
            root_times = [tree.time(tree.root) for tree in ts.trees()]
        assert len(set(root_times)) == 1
        if end_time is not None:
            assert root_times[0] == end_time

    @pytest.mark.parametrize(
        "model",
        [
            "dtwf",
            "hudson",
            msprime.BetaCoalescent(alpha=1.5),
            msprime.DiracCoalescent(psi=0.1, c=2),
        ],
    )
    def test_single_models(self, model):
        tss = msprime.sim_ancestry(
            10,
            population_size=10,
            stop_at_local_mrca=False,
            random_seed=1,
            sequence_length=10,
            recombination_rate=0.1,
            model=model,
            num_replicates=10,
        )
        for ts in tss:
            self.check_roots(ts)

    def test_fixed_pedigree_full_trace(self):
        # The pedigree used has this shape:
        #         O
        #         |
        #         O
        #        / \
        #       O   O
        #       \\ //
        #         O
        #        / \
        #       O   O
        #       \\ //
        #         O
        eldest_time = 5
        pb = msprime.PedigreeBuilder()
        gggpa = pb.add_individual(time=eldest_time)
        ggpa = pb.add_individual(time=4, parents=[gggpa, gggpa])
        # first diamound
        gmom = pb.add_individual(time=3, parents=[ggpa, ggpa])
        gdad = pb.add_individual(time=3, parents=[ggpa, ggpa])
        grandpa_id = pb.add_individual(time=2, parents=[gmom, gdad])
        # second diamond
        mom_id = pb.add_individual(time=1, parents=[grandpa_id, grandpa_id])
        dad_id = pb.add_individual(time=1, parents=[grandpa_id, grandpa_id])
        pb.add_individual(time=0, parents=[mom_id, dad_id], is_sample=True)
        pedigree = pb.finalise(100)

        ts = msprime.sim_ancestry(
            initial_state=pedigree,
            stop_at_local_mrca=False,
            coalescing_segments_only=False,
            additional_nodes=(
                msprime.NodeType.RECOMBINANT
                | msprime.NodeType.PASS_THROUGH
                | msprime.NodeType.COMMON_ANCESTOR
            ),
            recombination_rate=0.01,
            model="fixed_pedigree",
        )
        for tree in ts.trees():
            for sample in ts.samples():
                path = [sample] + list(tree.ancestors(sample))
                time = ts.nodes_time[path]
                # Each sample should pass through each level of the pedigree.
                np.testing.assert_array_equal(time, np.arange(eldest_time + 1))

    def test_multiple_models(self):
        """
        To test multiple models, we run simulations with end time < first models
        duration and test (that would be using the first model). Then we run again
        until end_time too large to cover the second model.
        """
        end_time = 499
        ts = msprime.sim_ancestry(
            2,
            population_size=100,
            model=[
                msprime.DiscreteTimeWrightFisher(duration=500),
                msprime.StandardCoalescent(),
            ],
            random_seed=2,
            recombination_rate=0.1,
            sequence_length=10,
            stop_at_local_mrca=False,
            end_time=end_time,
        )

        self.check_roots(ts, end_time, allow_multiple_roots=True)

        ts = msprime.sim_ancestry(
            2,
            population_size=100,
            model=[
                msprime.DiscreteTimeWrightFisher(duration=500),
                msprime.StandardCoalescent(),
            ],
            random_seed=2,
            recombination_rate=0.1,
            sequence_length=10,
            stop_at_local_mrca=False,
        )
        self.check_roots(ts)

    @pytest.mark.parametrize(
        "model",
        [
            "dtwf",
            "hudson",
        ],
    )
    def test_migration(self, model):
        demography = msprime.Demography.stepping_stone_model(
            [100, 100], migration_rate=0.01
        )
        ts = msprime.sim_ancestry(
            {0: 2, 1: 1},
            demography=demography,
            stop_at_local_mrca=False,
            recombination_rate=0.01,
            sequence_length=10,
            model=model,
        )
        self.check_roots(ts)

    @pytest.mark.parametrize(
        "model",
        [
            "dtwf",
            "hudson",
        ],
    )
    def test_additional_nodes(self, model):

        ts = msprime.sim_ancestry(
            2,
            population_size=10,
            random_seed=2,
            additional_nodes=(
                msprime.NodeType.RECOMBINANT
                | msprime.NodeType.PASS_THROUGH
                | msprime.NodeType.COMMON_ANCESTOR
            ),
            coalescing_segments_only=False,
            stop_at_local_mrca=False,
            sequence_length=5,
            model=model,
        )

        self.check_roots(ts)


class TestParseRandomSeed:
    """
    Tests for parsing the random seed values.
    """

    def test_default(self):
        # Make sure we get different random seeds when calling sequentially.
        seeds = [ancestry._parse_random_seed(None) for _ in range(100)]
        assert len(set(seeds)) == len(seeds)
        assert all(isinstance(seed, int) for seed in seeds)

    def test_numpy(self):
        seed = 12345
        ret_seed = ancestry._parse_random_seed(np.array([seed], dtype=int)[0])
        assert ret_seed == seed
        ret_seed = ancestry._parse_random_seed(np.array([seed], dtype=int))
        assert ret_seed == seed
        assert isinstance(ret_seed, int)

    def test_ints(self):
        # Anything that can be cast to an int is fine.
        for seed in [1234, 12.0, "12"]:
            ret_seed = ancestry._parse_random_seed(seed)
            assert ret_seed == int(seed)


class TestParseReplicateIndex:
    def test_none(self):
        assert (
            ancestry._parse_replicate_index(
                replicate_index=None, random_seed=None, num_replicates=None
            )
            is None
        )
        assert (
            ancestry._parse_replicate_index(
                replicate_index=None, random_seed=1, num_replicates=None
            )
            is None
        )
        assert (
            ancestry._parse_replicate_index(
                replicate_index=None, random_seed=1, num_replicates=1
            )
            is None
        )

    def test_random_seed(self):
        with pytest.raises(ValueError):
            ancestry._parse_replicate_index(
                replicate_index=0, random_seed=None, num_replicates=None
            )
        assert (
            ancestry._parse_replicate_index(
                replicate_index=1234, random_seed=1, num_replicates=None
            )
            == 1234
        )

    def test_num_replicates(self):
        with pytest.raises(ValueError):
            ancestry._parse_replicate_index(
                replicate_index=0, random_seed=1, num_replicates=1
            )

    def test_values(self):
        assert (
            ancestry._parse_replicate_index(
                replicate_index=0, random_seed=1, num_replicates=None
            )
            == 0
        )
        assert (
            ancestry._parse_replicate_index(
                replicate_index="1234", random_seed=1, num_replicates=None
            )
            == 1234
        )
        with pytest.raises(ValueError):
            ancestry._parse_replicate_index(
                replicate_index=-1, random_seed=1, num_replicates=None
            )


class TestParseSimAncestry:
    """
    Tests that the front-end for the sim_ancestry function correctly
    creates simulators with the required parameter values.
    """

    def test_sequence_length(self):
        # a single locus simulation will have sequence_length = 1
        sim = ancestry._parse_sim_ancestry(10)
        assert sim.sequence_length == 1
        assert sim.copy_tables().sequence_length == 1
        assert sim.recombination_map.total_mass == 0
        assert sim.gene_conversion_map.total_mass == 0

        # if we specify a rate_map for either GC or recomb this defines length.
        rate_map = msprime.RateMap.uniform(101, 0)
        sim = ancestry._parse_sim_ancestry(10, recombination_rate=rate_map)
        assert sim.sequence_length == rate_map.sequence_length
        assert sim.copy_tables().sequence_length == sim.sequence_length

        sim = ancestry._parse_sim_ancestry(
            10, gene_conversion_rate=rate_map, gene_conversion_tract_length=1
        )
        assert sim.sequence_length == rate_map.sequence_length
        assert sim.copy_tables().sequence_length == sim.sequence_length

        # If we have an initial_state this defines sequence_length
        initial_state = tskit.TableCollection(1234)
        initial_state.populations.add_row()
        with pytest.warns(msprime.IncompletePopulationMetadataWarning):
            sim = ancestry._parse_sim_ancestry(
                initial_state=initial_state, population_size=1
            )
        assert sim.sequence_length == 1234
        assert sim.copy_tables().sequence_length == sim.sequence_length

    def test_sequence_length_errors(self):
        # scaler rate values with no squence length is an error
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(10, recombination_rate=1)
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(
                10, gene_conversion_rate=1, gene_conversion_tract_length=1
            )

        # A rate map with a value that disagrees with sequence length
        rate_map = msprime.RateMap.uniform(101, 0)
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(
                10, recombination_rate=rate_map, sequence_length=1
            )
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(
                10,
                gene_conversion_rate=rate_map,
                gene_conversion_tract_length=1,
                sequence_length=1,
            )

        # A different rate map with a sequence_length that disagrees
        other_rate_map = msprime.RateMap.uniform(1, 0)
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(
                10,
                recombination_rate=other_rate_map,
                gene_conversion_rate=rate_map,
                gene_conversion_tract_length=1,
            )
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(
                10,
                recombination_rate=other_rate_map,
                gene_conversion_rate=rate_map,
                gene_conversion_tract_length=1,
                sequence_length=other_rate_map.sequence_length,
            )
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(
                10,
                recombination_rate=other_rate_map,
                gene_conversion_rate=rate_map,
                gene_conversion_tract_length=1,
                sequence_length=rate_map.sequence_length,
            )
        # Both maps disagree with sequence_length
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(
                10,
                recombination_rate=other_rate_map,
                gene_conversion_rate=rate_map,
                gene_conversion_tract_length=1,
                sequence_length=56789,
            )

        # An initial state with a sequence_length that disagrees.
        initial_state = tskit.TableCollection(1234)
        initial_state.time_units = "generations"
        initial_state = initial_state.tree_sequence()
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(initial_state=initial_state, sequence_length=1)
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(
                initial_state=initial_state, recombination_rate=rate_map
            )
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(
                initial_state=initial_state,
                gene_conversion_rate=rate_map,
                gene_conversion_tract_length=1,
            )
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(
                initial_state=initial_state,
                recombination_rate=other_rate_map,
                gene_conversion_rate=rate_map,
                gene_conversion_tract_length=1,
            )

    def test_sequence_length_discrete_genome(self):
        # Can't have floating point sequence_length with discrete_genome
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(10, sequence_length=1.1, discrete_genome=True)
        # Anything goes if we have a continuous genome, though.
        sim = ancestry._parse_sim_ancestry(
            10, sequence_length=0.1, discrete_genome=False
        )
        assert sim.sequence_length == 0.1

    def test_sequence_length_bad_arguments(self):
        for bad_value in ["x", b"sdf"]:
            with pytest.raises(ValueError):
                ancestry._parse_sim_ancestry(10, sequence_length=bad_value)

        for bad_type in [[], {}]:
            with pytest.raises(TypeError):
                ancestry._parse_sim_ancestry(10, sequence_length=bad_type)

    def test_gene_conversion_simple(self):
        for rate in ["1234", 1234, 1234.0]:
            sim = ancestry._parse_sim_ancestry(
                10,
                sequence_length=10,
                gene_conversion_rate=rate,
                gene_conversion_tract_length=5,
            )
            assert sim.sequence_length == 10
            gc_map = sim.gene_conversion_map
            assert gc_map.sequence_length == 10
            assert len(gc_map.rate) == 1
            assert gc_map.rate[0] == 1234
            assert sim.gene_conversion_tract_length == 5

    def test_gene_conversion_errors(self):
        # No tract length is an error
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(
                10, sequence_length=10, gene_conversion_rate=1234
            )
        # Specifying a tract_length and no map is also an error.
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(
                10, sequence_length=10, gene_conversion_tract_length=5.5
            )

        for bad_type in [[], {}]:
            with pytest.raises(TypeError):
                ancestry._parse_sim_ancestry(
                    10,
                    sequence_length=10,
                    gene_conversion_rate=bad_type,
                    gene_conversion_tract_length=1,
                )

            with pytest.raises(TypeError):
                ancestry._parse_sim_ancestry(
                    10,
                    sequence_length=10,
                    gene_conversion_rate=1,
                    gene_conversion_tract_length=bad_type,
                )

    def test_discrete_genome(self):
        # default is True
        sim = ancestry._parse_sim_ancestry(10, sequence_length=10)
        assert sim.discrete_genome
        for discrete_genome in [True, False]:
            sim = ancestry._parse_sim_ancestry(
                10, sequence_length=10, discrete_genome=discrete_genome
            )
            assert sim.discrete_genome == discrete_genome
        # Falsey values are not OK
        for discrete_genome in ["", []]:
            with pytest.raises(TypeError):
                ancestry._parse_sim_ancestry(
                    10, sequence_length=10, discrete_genome=discrete_genome
                )

    def test_start_time(self):
        # default is 0
        sim = ancestry._parse_sim_ancestry(10)
        assert sim.start_time == 0
        for start_time in [1234, 1234.34, "1", "1.234"]:
            sim = ancestry._parse_sim_ancestry(10, start_time=start_time)
            assert sim.start_time == float(start_time)
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(10, start_time="bad value")
        with pytest.raises(TypeError):
            ancestry._parse_sim_ancestry(10, start_time=[])

    def test_end_time(self):
        # default is inf
        sim = ancestry._parse_sim_ancestry(10)
        assert sim.end_time == math.inf
        for end_time in [1234, 1234.34, "1", "1.234"]:
            sim = ancestry._parse_sim_ancestry(10, end_time=end_time)
            assert sim.end_time == float(end_time)
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(10, end_time="bad value")
        with pytest.raises(TypeError):
            ancestry._parse_sim_ancestry(10, end_time=[])

    def test_record_migrations(self):
        # default is False
        sim = ancestry._parse_sim_ancestry(10)
        assert not sim.record_migrations
        for record_migrations in [True, False]:
            sim = ancestry._parse_sim_ancestry(10, record_migrations=record_migrations)
            assert sim.record_migrations == bool(record_migrations)
        for truthy in [0, []]:
            with pytest.raises(TypeError):
                ancestry._parse_sim_ancestry(10, record_migrations=truthy)

    def test_record_full_arg(self):
        # default is False
        sim = ancestry._parse_sim_ancestry(10)
        assert sim.additional_nodes == 0
        node_values = [sum(2**i for i in (17, 18, 19, 21)), 0]
        for record_full_arg, node_value in zip([True, False], node_values):
            sim = ancestry._parse_sim_ancestry(10, record_full_arg=record_full_arg)
            assert sim.additional_nodes == node_value

        for truthy in [0, []]:
            with pytest.raises(TypeError):
                ancestry._parse_sim_ancestry(10, record_full_arg=truthy)

    def test_ploidy(self):
        # default is 2
        sim = ancestry._parse_sim_ancestry(10)
        assert sim.ploidy == 2
        for ploidy in [1, 2, np.array([33])[0]]:
            sim = ancestry._parse_sim_ancestry(10, ploidy=ploidy)
            assert sim.ploidy == int(ploidy)

        for bad_ploidy in [0, -1]:
            with pytest.raises(ValueError):
                ancestry._parse_sim_ancestry(10, ploidy=bad_ploidy)

        for bad_ploidy in ["1", "0.1", 0.1, np.array([0.1])[0]]:
            with pytest.raises(TypeError):
                ancestry._parse_sim_ancestry(10, ploidy=bad_ploidy)

    def test_population_size(self):
        # default is 1
        sim = ancestry._parse_sim_ancestry(10)
        assert sim.demography.num_populations == 1
        assert sim.demography.populations[0].initial_size == 1
        assert sim.demography.populations[0].growth_rate == 0
        for pop_size in [2, 0.1, 100, 1e6, "100"]:
            sim = ancestry._parse_sim_ancestry(10, population_size=pop_size)
            assert sim.demography.num_populations == 1
            assert sim.demography.populations[0].initial_size == float(pop_size)
            assert sim.demography.populations[0].growth_rate == 0
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(10, population_size=-1)
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(10, population_size="bad value")
        with pytest.raises(TypeError):
            ancestry._parse_sim_ancestry(10, population_size=[])

        # Cannot specify a population_size and demography args.
        demography = msprime.Demography.stepping_stone_model([1], 0)
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(10, demography=demography, population_size=1)

        # Cannot specify a population_size and a initial state with > 1 pop.
        demography = msprime.Demography.isolated_model([10, 10])
        ts1 = msprime.sim_ancestry(
            {0: 2}, demography=demography, end_time=0, random_seed=1
        )
        with pytest.raises(ValueError, match="Must specify demography"):
            msprime.sim_ancestry(initial_state=ts1, population_size=100, random_seed=2)

    def test_demography(self):
        demography = msprime.Demography.stepping_stone_model([1] * 5, 0.1)
        samples = {0: 2}
        sim = ancestry._parse_sim_ancestry(samples, demography=demography)
        assert sim.demography is not demography
        assert sim.num_populations == demography.num_populations
        assert np.array_equal(sim.migration_matrix, demography.migration_matrix)
        # Numeric samples fail here as we have more than 1 population
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(5, demography=demography)
        with pytest.raises(TypeError):
            ancestry._parse_sim_ancestry(samples, demography="not a demography")

        demography.populations[0].initial_size = -1
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(samples, demography=demography)

    def test_model(self):
        # Extensive testing of the model parsing is done elsewhere.
        sim = ancestry._parse_sim_ancestry(10, model="smc")
        assert len(sim.models) == 1
        assert isinstance(sim.models[0], msprime.SmcApproxCoalescent)
        assert sim.models[0].duration is None

        sim = ancestry._parse_sim_ancestry(
            10, model=(msprime.SmcPrimeApproxCoalescent(duration=10), "hudson")
        )
        assert len(sim.models) == 2
        assert isinstance(sim.models[0], msprime.SmcPrimeApproxCoalescent)
        assert sim.models[0].duration == 10
        assert isinstance(sim.models[1], msprime.StandardCoalescent)
        assert sim.models[1].duration is None

    def test_dtwf_population_size(self):
        # It's an error to not specify a pop size for dtwf.
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(10, model="dtwf", ploidy=2)

    def test_initial_state_errors(self):
        tables = tskit.TableCollection(1)
        tables.time_units = "generations"
        tables.populations.add_row()
        # sequence_length doesn't match.
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(initial_state=tables, sequence_length=100)
        # Must have at least one population
        tables = tskit.TableCollection(1)
        tables.time_units = "generations"
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(initial_state=tables)
        for bad_type in [[], "sdf", {}]:
            with pytest.raises(TypeError):
                ancestry._parse_sim_ancestry(initial_state=bad_type)

        # Must have a demography if we have more than 1 pop.
        demography = msprime.Demography.isolated_model([10, 10])
        ts1 = msprime.sim_ancestry(
            {0: 2}, demography=demography, end_time=0, random_seed=1
        )
        with pytest.raises(ValueError, match="Must specify either a demography"):
            msprime.sim_ancestry(initial_state=ts1, random_seed=2)

        demography = msprime.Demography.isolated_model([10])
        ts1 = msprime.sim_ancestry(
            {0: 2}, demography=demography, end_time=0, random_seed=1
        )
        with pytest.raises(ValueError, match="Must specify either a demography"):
            msprime.sim_ancestry(initial_state=ts1, random_seed=2)

    def test_initial_state(self):
        ts = msprime.sim_ancestry(10, end_time=0.01, random_seed=2)
        # Same if we use either the tables or tree sequence object.
        sim = ancestry._parse_sim_ancestry(initial_state=ts, population_size=1)
        assert sim.copy_tables() == ts.tables
        sim = ancestry._parse_sim_ancestry(initial_state=ts.tables, population_size=1)
        assert sim.copy_tables() == ts.tables

    def test_num_labels(self):
        for num_labels in [1, 2, 10]:
            sim = ancestry._parse_sim_ancestry(10, num_labels=num_labels)
            assert sim.num_labels == num_labels

    def test_samples_and_initial_state(self):
        # If we specify neither samples or initial_state we get an error
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(samples=None)

        # Specifying both is also an error.
        tables = tskit.TableCollection(1)
        tables.time_units = "generations"
        tables.populations.add_row()
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(2, initial_state=tables)


class TestParseSimAncestryPedigree:
    def test_pedigree_requires_initial_state(self):
        with pytest.raises(ValueError, match="Must specify an input pedigree"):
            ancestry._parse_sim_ancestry(model="fixed_pedigree")

    def test_pedigree_requires_no_demography(self):
        demography = msprime.Demography.isolated_model([1])
        with pytest.raises(ValueError, match="Cannot specify demography"):
            ancestry._parse_sim_ancestry(model="fixed_pedigree", demography=demography)

    def test_pedigree_requires_no_population_size(self):
        with pytest.raises(ValueError, match="Cannot specify population_size"):
            ancestry._parse_sim_ancestry(model="fixed_pedigree", population_size=1)

    @pytest.mark.parametrize("ploidy", [1, 3])
    def test_pedigree_requires_ploidy2(self, ploidy):
        tables = tskit.TableCollection(1)
        tables.time_units = "generations"
        with pytest.raises(ValueError, match="must have ploidy=2"):
            ancestry._parse_sim_ancestry(
                model="fixed_pedigree",
                initial_state=tables,
                ploidy=ploidy,
            )

    def test_pedigree_requires_no_samples(self):
        tables = tskit.TableCollection(1)
        tables.time_units = "generations"
        with pytest.raises(ValueError, match="Cannot specify both samples and"):
            ancestry._parse_sim_ancestry(
                model="fixed_pedigree",
                initial_state=tables,
                samples=10,
            )

    def test_pedigree_requires_no_gc(self):
        pb = msprime.PedigreeBuilder()
        with pytest.raises(ValueError, match="Gene conversion not supported"):
            ancestry._parse_sim_ancestry(
                model="fixed_pedigree",
                initial_state=pb.finalise(10),
                gene_conversion_rate=1,
                gene_conversion_tract_length=2,
            )

    def test_pedigree_requires_no_full_arg(self):
        pb = msprime.PedigreeBuilder()
        with pytest.raises(ValueError, match="Full ARG recording not supported"):
            ancestry._parse_sim_ancestry(
                model="fixed_pedigree",
                initial_state=pb.finalise(10),
                record_full_arg=True,
            )

    def test_pedigree_requires_no_record_migrations(self):
        pb = msprime.PedigreeBuilder()
        with pytest.raises(ValueError, match="Migration recording not supported"):
            ancestry._parse_sim_ancestry(
                model="fixed_pedigree",
                initial_state=pb.finalise(10),
                record_migrations=True,
            )

    def test_pedigree_requires_no_start_time(self):
        pb = msprime.PedigreeBuilder()
        with pytest.raises(ValueError, match="Cannot specify start_time"):
            ancestry._parse_sim_ancestry(
                model="fixed_pedigree", initial_state=pb.finalise(10), start_time=1
            )

    def test_pedigree_with_following_model(self):
        pb = msprime.PedigreeBuilder()
        with pytest.raises(
            ValueError, match="Cannot use FixedPedigree simulation in conjunction"
        ):
            ancestry._parse_sim_ancestry(
                model=["fixed_pedigree", "hudson"], initial_state=pb.finalise(10)
            )

    def test_pedigree_with_previous_model(self):
        pb = msprime.PedigreeBuilder()
        with pytest.raises(
            ValueError, match="Cannot use FixedPedigree simulation in conjunction"
        ):
            ancestry._parse_sim_ancestry(
                model=["hudson", "fixed_pedigree"], initial_state=pb.finalise(10)
            )


class TestSimAncestrySamples:
    def test_negative_samples(self):
        for ploidy in [1, 2, 5]:
            with pytest.raises(ValueError):
                ancestry._parse_sim_ancestry(-1, ploidy=ploidy)

    def test_numeric_samples(self):
        for n in [1, 3, 10]:
            for ploidy in [1, 2, 3]:
                sim = ancestry._parse_sim_ancestry(n, ploidy=ploidy)
                assert sim.ploidy == ploidy
                tables = sim.copy_tables()
                assert len(tables.individuals) == n
                for individual in tables.individuals:
                    assert individual.flags == 0
                    assert len(individual.location) == 0
                assert len(tables.nodes) == n * ploidy
                assert len(tables.populations) == 1
                for node_id, node in enumerate(tables.nodes):
                    assert node.individual == node_id // ploidy
                    assert node.time == 0
                    assert node.flags == tskit.NODE_IS_SAMPLE
                    assert node.population == 0

    def test_sample_sets_override_ploidy(self):
        samples = [
            msprime.SampleSet(1, ploidy=2),
            msprime.SampleSet(1),  # No ploidy, should use default
            msprime.SampleSet(1, ploidy=1),
        ]
        sim = ancestry._parse_sim_ancestry(samples, ploidy=3)
        tables = sim.copy_tables()
        assert len(tables.individuals) == 3
        assert len(tables.nodes) == 6
        assert np.all(tables.nodes.individual[0:2] == 0)
        assert np.all(tables.nodes.individual[2:5] == 1)
        assert tables.nodes.individual[5] == 2

    def test_sample_sets_override_time(self):
        demography = msprime.Demography.isolated_model([1])
        demography.populations[0].default_sampling_time = 10
        samples = [
            msprime.SampleSet(1, ploidy=2, time=2),
            msprime.SampleSet(1, ploidy=3),  # No time, should use default
            msprime.SampleSet(1, ploidy=1, time=0),
        ]
        sim = ancestry._parse_sim_ancestry(samples, demography=demography)
        tables = sim.copy_tables()
        assert len(tables.individuals) == 3
        assert len(tables.nodes) == 6
        assert np.all(tables.nodes.time[0:2] == 2)
        assert np.all(tables.nodes.time[2:5] == 10)
        assert tables.nodes.time[5] == 0

    def test_numeric_samples_types(self):
        # Make sure the various different ways we can specify a numeric
        # set of samples all give the same answer
        values = [10, 10.0, np.array([10], dtype=int)[0]]
        for value in values:
            sim = ancestry._parse_sim_ancestry(value)
            assert len(sim.copy_tables().individuals) == 10

    def test_numeric_samples_only_simple_demography(self):
        # A simple 1-population model is fine.
        demography = msprime.Demography.isolated_model([1])
        sim = ancestry._parse_sim_ancestry(10, demography=demography)
        assert len(sim.copy_tables().individuals) == 10

        demography = msprime.Demography.stepping_stone_model([1, 1], 0)
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry(10, demography=demography)

    def test_population_name_sampling_errors(self):
        model = msprime.Demography.island_model([1] * 2, migration_rate=1)
        model.populations[0].name = "A"
        model.populations[1].name = "B"
        with pytest.raises(ValueError):
            ancestry._parse_sim_ancestry({}, demography=model)

        for bad_sample in [{"A": 1, "B": -1}, {"A": -1}, {"A": 0, "B": -10}]:
            with pytest.raises(ValueError):
                ancestry._parse_sim_ancestry(bad_sample, demography=model)
        for bad_type in [6.6, "a"]:
            with pytest.raises(TypeError):
                ancestry._parse_sim_ancestry({"A": bad_type}, demography=model)
        for bad_name in ["C", "AC", "", "X" * 100]:
            with pytest.raises(KeyError):
                ancestry._parse_sim_ancestry({bad_name: 1}, demography=model)

    def test_population_name_samples_two_populations(self):
        model = msprime.Demography.island_model([1] * 2, migration_rate=1)
        model.populations[0].name = "A"
        model.populations[1].name = "B"
        sim = ancestry._parse_sim_ancestry({"A": 1}, demography=model)
        tables = sim.copy_tables()
        assert len(tables.nodes) == 2
        assert np.all(tables.nodes.population == 0)

        sim = ancestry._parse_sim_ancestry({"B": 1}, demography=model)
        tables = sim.copy_tables()
        assert len(tables.nodes) == 2
        assert np.all(tables.nodes.population == 1)

        sim = ancestry._parse_sim_ancestry({"A": 1, "B": 1}, demography=model)
        tables = sim.copy_tables()
        assert len(tables.nodes) == 4
        assert np.all(tables.nodes.population[:2] == 0)
        assert np.all(tables.nodes.population[2:] == 1)

        sim = ancestry._parse_sim_ancestry({"B": 1, "A": 1}, demography=model)
        tables = sim.copy_tables()
        assert len(tables.nodes) == 4
        assert np.all(tables.nodes.population[:2] == 1)
        assert np.all(tables.nodes.population[2:] == 0)

        sim = ancestry._parse_sim_ancestry({0: 1, "B": 1}, demography=model)
        tables = sim.copy_tables()
        assert len(tables.nodes) == 4
        assert np.all(tables.nodes.population[:2] == 0)
        assert np.all(tables.nodes.population[2:] == 1)

        sim = ancestry._parse_sim_ancestry({1: 1, 0: 1}, demography=model)
        tables = sim.copy_tables()
        assert len(tables.nodes) == 4
        assert np.all(tables.nodes.population[:2] == 1)
        assert np.all(tables.nodes.population[2:] == 0)

        # Referring to the same population by name and ID means we do the
        # sampling twice.
        sim = ancestry._parse_sim_ancestry({0: 1, "A": 1}, demography=model)
        tables = sim.copy_tables()
        assert len(tables.nodes) == 4
        assert np.all(tables.nodes.population[:4] == 0)

    def verify_samples_map(self, samples, demography, ploidy):
        demography = demography.validate()
        sim = ancestry._parse_sim_ancestry(
            samples=samples, demography=demography, ploidy=ploidy
        )
        n = sum(samples.values())
        assert sim.ploidy == ploidy
        tables = sim.copy_tables()
        assert len(tables.individuals) == n
        assert len(tables.nodes) == n * ploidy
        assert len(tables.populations) == demography.num_populations
        node_id = 0
        ind_id = 0
        for pop_id, num_samples in samples.items():
            for _ in range(num_samples):
                individual = tables.individuals[ind_id]
                for _ in range(ploidy):
                    node = tables.nodes[node_id]
                    assert node.individual == ind_id
                    assert (
                        node.time == demography.populations[pop_id].default_sampling_time
                    )
                    assert node.population == pop_id
                    assert node.flags == tskit.NODE_IS_SAMPLE
                    node_id += 1
                assert individual.flags == 0
                assert len(individual.location) == 0
                ind_id += 1

    def test_sample_demography(self):
        demography = msprime.Demography.isolated_model([1])
        self.verify_samples_map({0: 10}, demography, ploidy=1)
        self.verify_samples_map({0: 10}, demography, ploidy=2)

        demography = msprime.Demography.stepping_stone_model([1] * 5, 0)
        samples = {j: j for j in range(5)}
        self.verify_samples_map(samples, demography, ploidy=1)
        self.verify_samples_map(samples, demography, ploidy=2)

        samples = {4: 15}
        self.verify_samples_map(samples, demography, ploidy=1)
        self.verify_samples_map(samples, demography, ploidy=2)

    def test_sample_time_with_map(self):
        demography = msprime.Demography.stepping_stone_model([1, 1], 0)
        demography.populations[0].default_sampling_time = 1
        demography.populations[1].default_sampling_time = 2
        samples = {0: 5, 1: 5}
        self.verify_samples_map(samples, demography, ploidy=1)
        self.verify_samples_map(samples, demography, ploidy=2)

    def test_bad_list_samples_population(self):
        demography = msprime.Demography.stepping_stone_model([1, 1], 0)
        for bad_pop in [-1, 2]:
            samples = [msprime.SampleSet(2, time=0, population=bad_pop)]
            with pytest.raises(KeyError):
                ancestry._parse_sim_ancestry(samples=samples, demography=demography)
        for bad_pop in [1.1, ValueError]:
            samples = [msprime.SampleSet(2, time=0, population=bad_pop)]
            with pytest.raises(KeyError):
                ancestry._parse_sim_ancestry(samples=samples, demography=demography)

    def test_list_samples_no_population(self):
        # No population value is fine when we have one population
        samples = [msprime.SampleSet(2)]
        ts = ancestry.sim_ancestry(samples)
        assert ts.num_populations == 1

        # But not if we have 2 or more pops.
        demography = msprime.Demography.stepping_stone_model([1, 1], 0)
        with pytest.raises(ValueError, match="SampleSet population"):
            ancestry._parse_sim_ancestry(samples=samples, demography=demography)

    def test_sample_list_types(self):
        bad_sample_types = [ValueError, msprime.Sample(0, 0), "asdrf"]
        for bad_sample_type in bad_sample_types:
            with pytest.raises(TypeError, match="SampleSet object"):
                ancestry._parse_sim_ancestry([msprime.SampleSet(1), bad_sample_type])

    def test_bad_top_level_sample_types(self):
        bad_sample_types = [ValueError]
        for bad_sample_type in bad_sample_types:
            with pytest.raises(TypeError, match="different forms"):
                ancestry._parse_sim_ancestry(bad_sample_type)


class TestStartTime:
    """
    Test the detailed semantics of start_time in terms of how
    it interacts with demographic events, sampling, etc.
    """

    def test_ancient_sample_before_start_time(self):
        ts = msprime.sim_ancestry(
            [
                msprime.SampleSet(1, time=0, ploidy=1),
                msprime.SampleSet(1, time=1, ploidy=1),
            ],
            start_time=2,
        )
        assert ts.num_nodes == 3
        time = ts.tables.nodes.time
        assert time[0] == 0
        assert time[1] == 1
        assert time[2] > 2

    def test_all_ancient_sample_before_start_time(self):
        ts = msprime.sim_ancestry(
            [
                msprime.SampleSet(1, time=1, ploidy=1),
                msprime.SampleSet(1, time=2, ploidy=1),
            ],
            start_time=3,
        )
        assert ts.num_nodes == 3
        time = ts.tables.nodes.time
        assert time[0] == 1
        assert time[1] == 2
        assert time[2] > 3

    def test_population_size_change_before_start_time(self):
        demography = msprime.Demography()
        # Effectively zero size.
        demography.add_population(initial_size=1e-200)
        with pytest.raises(_msprime.LibraryError, match="<= to its child"):
            msprime.sim_ancestry(2, demography=demography, start_time=5, random_seed=1)
        demography.add_population_parameters_change(time=4.999, initial_size=1)
        msprime.sim_ancestry(2, demography=demography, start_time=5, random_seed=1)

    def test_population_split_before_start_time(self):
        demography = msprime.Demography()
        demography.add_population(name="A", initial_size=1)
        demography.add_population(name="B", initial_size=1)
        demography.add_population(name="C", initial_size=1)
        demography.add_population_split(time=3, ancestral="C", derived=["A", "B"])
        # Here we sample lineages at time 0 and the these get moved to "C" at
        # time 3. Only when we reach "start_time" do we start generating
        # events.
        ts = msprime.sim_ancestry(
            {"A": 1, "B": 1},
            ploidy=1,
            demography=demography,
            start_time=15,
            random_seed=1,
        )
        nodes = ts.tables.nodes
        assert nodes.time[-1] > 15
        assert nodes.population[-1] == 2

    def test_population_sample_inactive_before_start_time(self):
        demography = msprime.Demography()
        demography.add_population(name="A", initial_size=1)
        demography.add_population(name="B", initial_size=1)
        demography.add_population(name="C", initial_size=1)
        demography.add_population_split(time=3, ancestral="C", derived=["A", "B"])
        with pytest.raises(_msprime.InputError, match="from an inactive population"):
            msprime.sim_ancestry(
                [
                    msprime.SampleSet(1, population="B", time=0),
                    msprime.SampleSet(1, population="A", time=4),
                ],
                demography=demography,
                start_time=15,
                random_seed=1,
            )


class TestParseSimulate:
    """
    Tests that the front-end for the simulate function correctly
    creates simulators with the required parameter values.
    """

    def test_length(self):
        for bad_length in [-1, 0, -1e-6]:
            with pytest.raises(ValueError):
                ancestry._parse_simulate(10, length=bad_length)

    def test_num_labels(self):
        for bad_value in [-1, 0, 0.1]:
            with pytest.raises(ValueError):
                ancestry._parse_simulate(10, num_labels=bad_value)

    def test_sample_size(self):
        with pytest.raises(ValueError):
            ancestry._parse_simulate()
        with pytest.raises(ValueError):
            ancestry._parse_simulate(1)
        with pytest.raises(ValueError):
            ancestry._parse_simulate(sample_size=1)
        for n in [2, 100, 1000]:
            sim = ancestry._parse_simulate(n)
            tables = sim.copy_tables()
            assert len(tables.populations) == 1
            assert len(tables.individuals) == 0
            assert len(tables.edges) == 0
            assert len(tables.nodes) == n
            assert np.all(tables.nodes.flags == tskit.NODE_IS_SAMPLE)
            assert np.all(tables.nodes.time == 0)
            assert np.all(tables.nodes.individual == tskit.NULL)

    def test_effective_population_size(self):
        def f(Ne):
            return ancestry._parse_simulate(10, Ne=Ne)

        for bad_value in [-1, -1e16, 0]:
            with pytest.raises(ValueError):
                f(bad_value)
        for Ne in [1, 10, 1e5]:
            sim = f(Ne)
            assert sim.demography.populations[0].initial_size == Ne
        # Test the default.
        sim = ancestry._parse_simulate(10)
        assert sim.demography.populations[0].initial_size == 1

    def test_population_configurations(self):
        def f(configs):
            return ancestry._parse_simulate(population_configurations=configs)

        for bad_type in [10, ["sdf"], "sdfsd"]:
            with pytest.raises(TypeError):
                f(bad_type)
        # Just test the basic equalities here. The actual
        # configuration options are tested elewhere.
        for N in range(1, 10):
            pop_configs = [
                msprime.PopulationConfiguration(5, initial_size=5) for _ in range(N)
            ]
            sample_size = 5 * N
            sim = ancestry._parse_simulate(population_configurations=pop_configs)
            assert len(sim.demography.populations) == len(pop_configs)
            for pop, pop_config in zip(sim.demography.populations, pop_configs):
                assert pop.initial_size == pop_config.initial_size
                assert pop.growth_rate == pop_config.growth_rate
            tables = sim.copy_tables()
            assert len(tables.nodes) == sample_size
            assert len(sim.population_configuration) == N
        # The default is a single population
        sim = ancestry._parse_simulate(10)
        assert len(sim.population_configuration) == 1

    def test_sample_size_population_configuration(self):
        for d in range(1, 5):
            # Zero sample size is always an error
            configs = [msprime.PopulationConfiguration(0) for _ in range(d)]
            with pytest.raises(ValueError):
                ancestry._parse_simulate(population_configurations=configs)
            configs = [msprime.PopulationConfiguration(2) for _ in range(d)]
            sim = ancestry._parse_simulate(population_configurations=configs)
            tables = sim.copy_tables()
            assert len(tables.nodes) == 2 * d
            i = 0
            for j in range(d):
                for _ in range(2):
                    node = tables.nodes[i]
                    assert node.population == j
                    assert node.time == 0
                    assert node.flags == tskit.NODE_IS_SAMPLE
                    i += 1

    def test_migration_matrix(self):
        # Cannot specify a migration matrix without population
        # configurations
        with pytest.raises(ValueError):
            ancestry._parse_simulate(10, migration_matrix=[])
        for N in range(1, 10):
            pop_configs = [msprime.PopulationConfiguration(5) for _ in range(N)]
            sim = ancestry._parse_simulate(population_configurations=pop_configs)
            # If we don't specify a matrix, it's 0 everywhere.
            matrix = np.zeros((N, N))
            np.testing.assert_array_equal(sim.migration_matrix, matrix)

            def f(matrix):
                return ancestry._parse_simulate(
                    population_configurations=pop_configs,  # noqa: B023
                    migration_matrix=matrix,
                )

            matrix = [[(j + k) * int(j != k) for j in range(N)] for k in range(N)]
            sim = f(matrix)
            np.testing.assert_array_equal(sim.demography.migration_matrix, matrix)
            # Try with equivalent numpy array.
            sim = f(np.array(matrix))
            np.testing.assert_array_equal(sim.demography.migration_matrix, matrix)
            np.testing.assert_array_equal(sim.migration_matrix, matrix)
            for bad_type in [{}, "", 234, 1.2]:
                with pytest.raises(ValueError):
                    f(bad_type)
            # Now check for the structure of the matrix.
            matrix[0][0] = "bad value"
            with pytest.raises(TypeError):
                f(matrix)
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                matrix[0] = None
                with pytest.raises(ValueError):
                    f(matrix)
                matrix[0] = []
                with pytest.raises(ValueError):
                    f(matrix)
            # Simple numpy array.
            matrix = np.ones((N, N))
            np.fill_diagonal(matrix, 0)
            sim = f(matrix)
            np.testing.assert_array_equal(
                np.array(sim.demography.migration_matrix), matrix
            )
            sim.run()
            events = np.array(sim.num_migration_events)
            assert events.shape == (N, N)
            assert np.all(events >= 0)

    def test_default_migration_matrix(self):
        sim = ancestry._parse_simulate(10)
        assert sim.migration_matrix == [0.0]

    def test_demographic_events(self):
        for bad_type in ["sdf", 234, [12], [None]]:
            with pytest.raises(TypeError):
                ancestry._parse_simulate(2, demographic_events=bad_type)
        # TODO test for bad values.

    def test_recombination_rate(self):
        def f(recomb_rate):
            return ancestry._parse_simulate(10, recombination_rate=recomb_rate)

        for bad_type in ["", {}, []]:
            with pytest.raises(TypeError):
                f(bad_type)
        for bad_value in [-1, -1e15]:
            with pytest.raises(ValueError):
                f(bad_value)
        for rate in [0, 1e-3, 10]:
            sim = f(rate)
            recomb_map = sim.recombination_map
            assert list(recomb_map.position) == [0, 1]
            assert list(recomb_map.rate) == [rate]
            assert sim.sequence_length == recomb_map.sequence_length

    def test_recombination_map(self):
        def f(recomb_map):
            return ancestry._parse_simulate(10, recombination_map=recomb_map)

        with pytest.raises(TypeError):
            f("wrong type")
        for n in range(2, 10):
            positions = list(range(n))
            rates = [0.1 * j for j in range(n - 1)]
            # Use the old-form RecombinationMap
            recomb_map = msprime.RecombinationMap(positions, rates + [0.0])
            sim = ancestry._parse_simulate(10, recombination_map=recomb_map)
            other_map = sim.recombination_map
            assert list(other_map.position) == positions
            assert list(other_map.rate) == rates
            assert sim.sequence_length == other_map.sequence_length
            # Use the new-form RateMap
            rate_map = msprime.RateMap(position=positions, rate=rates)
            sim = ancestry._parse_simulate(10, recombination_map=rate_map)
            other_map = sim.recombination_map
            assert list(other_map.position) == positions
            assert list(other_map.rate) == rates
            assert sim.sequence_length == other_map.sequence_length

    def test_combining_recomb_map_and_rate_length(self):
        recomb_map = msprime.RecombinationMap([0, 1], [1, 0])
        with pytest.raises(ValueError):
            ancestry._parse_simulate(
                10,
                recombination_map=recomb_map,
                length=1,
            )
        with pytest.raises(ValueError):
            ancestry._parse_simulate(
                10,
                recombination_map=recomb_map,
                recombination_rate=100,
            )
        with pytest.raises(ValueError):
            ancestry._parse_simulate(
                10,
                recombination_map=recomb_map,
                length=1,
                recombination_rate=1,
            )

    def test_sample_combination_errors(self):
        # Make sure that the various ways we can specify the samples
        # operate correctly.
        s = msprime.Sample(time=0.0, population=0)
        with pytest.raises(ValueError):
            ancestry._parse_simulate()
        # Cannot provide sample_size with either population configurations
        # or samples
        with pytest.raises(ValueError):
            ancestry._parse_simulate(sample_size=2, samples=[s, s])
        pop_configs = [msprime.PopulationConfiguration(sample_size=2)]
        with pytest.raises(ValueError):
            ancestry._parse_simulate(
                sample_size=2,
                population_configurations=pop_configs,
            )
        # If we provide samples and population_configurations we cannot
        # have a sample size for the config.
        pop_configs = [msprime.PopulationConfiguration(sample_size=2)]
        with pytest.raises(ValueError):
            ancestry._parse_simulate(
                samples=[s, s],
                population_configurations=pop_configs,
            )
        pop_configs = [
            msprime.PopulationConfiguration(sample_size=None),
            msprime.PopulationConfiguration(sample_size=2),
        ]
        with pytest.raises(ValueError):
            ancestry._parse_simulate(
                samples=[s, s],
                population_configurations=pop_configs,
            )

    def test_samples(self):
        pop_configs = [
            msprime.PopulationConfiguration(),
            msprime.PopulationConfiguration(),
            msprime.PopulationConfiguration(),
        ]
        samples = [
            msprime.Sample(population=0, time=0),
            msprime.Sample(population=1, time=1),
            msprime.Sample(population=2, time=2),
        ]
        sim = ancestry._parse_simulate(
            samples=samples, population_configurations=pop_configs
        )
        tables = sim.copy_tables()
        assert len(tables.nodes) == len(samples)
        for node, sample in zip(tables.nodes, samples):
            assert node.population == sample.population
            assert node.time == sample.time
            assert node.flags == tskit.NODE_IS_SAMPLE

    def test_sample_is_tuple(self):
        s = msprime.Sample(population=1, time=2)
        assert s.population == 1
        assert s.time == 2
        assert s == (1, 2)

    def test_model_change_events(self):
        main_model = msprime.SmcApproxCoalescent()
        sim = ancestry._parse_simulate(
            Ne=100,
            sample_size=2,
            model=main_model,
            demographic_events=[
                msprime.SimulationModelChange(1, msprime.DiscreteTimeWrightFisher()),
                msprime.SimulationModelChange(2, None),
            ],
        )
        assert len(sim.models) == 3
        assert sim.models[0].duration == 1
        assert isinstance(sim.models[0], msprime.SmcApproxCoalescent)
        assert sim.models[1].duration == 1
        assert isinstance(sim.models[1], msprime.DiscreteTimeWrightFisher)
        # When model=None we change to the standard coalescent
        assert sim.models[2].duration is None
        assert isinstance(sim.models[2], msprime.StandardCoalescent)

    def test_model_change_events_error(self):
        with pytest.raises(ValueError, match="durations must be >= 0"):
            msprime.simulate(
                Ne=100,
                sample_size=2,
                random_seed=2,
                model=None,
                demographic_events=[
                    msprime.SimulationModelChange(1, msprime.DiscreteTimeWrightFisher()),
                    msprime.SimulationModelChange(0.5, None),
                ],
            )

    def test_bad_sample_population_reference(self):
        # What happens when we reference a population that doesn't exist?
        with pytest.raises(
            ValueError, match="Invalid population reference '1' in sample at index 1"
        ):
            msprime.simulate(
                samples=[
                    msprime.Sample(population=0, time=0),
                    msprime.Sample(population=1, time=0),
                ]
            )

        with pytest.raises(
            ValueError, match="Negative population ID in sample at index 2"
        ):
            msprime.simulate(
                samples=[
                    msprime.Sample(population=0, time=0),
                    msprime.Sample(population=0, time=0),
                    msprime.Sample(population=-1, time=0),
                ]
            )

    def test_bad_additional_node_flag(self):
        with pytest.raises(
            ValueError,
            match="Set coalescing_segments_only to False when recording "
            "additional_nodes.",
        ):
            msprime.sim_ancestry(10, additional_nodes=msprime.NodeType.COMMON_ANCESTOR)
        with pytest.raises(
            ValueError,
            match="additional_nodes should be an instance of msprime.NodeType.",
        ):
            msprime.sim_ancestry(10, additional_nodes=1 << 17)

    def test_bad_full_arg_flag(self):
        with pytest.raises(
            ValueError, match="Cannot set both record_full_arg and additional_nodes."
        ):
            msprime.sim_ancestry(
                10,
                additional_nodes=msprime.NodeType.COMMON_ANCESTOR,
                record_full_arg=True,
                coalescing_segments_only=None,
            )


class TestSimAncestryInterface:
    """
    Some simple tests cases for the sim_ancestry interface.
    """

    def test_defaults(self):
        n = 10
        # Diploid sim by default.
        ts = msprime.sim_ancestry(n)
        assert ts.num_samples == 2 * n
        assert ts.num_individuals == n
        assert ts.num_trees == 1
        assert ts.num_sites == 0
        assert ts.sequence_length == 1

    def test_ploidy(self):
        n = 10
        for k in [1, 2, 3, 4]:
            ts = msprime.sim_ancestry(n, ploidy=k)
            assert ts.num_samples == k * n
            assert ts.num_trees == 1
            assert ts.num_sites == 0
            assert ts.sequence_length == 1
            # TODO check for individuals

    def test_hudson_time_scale(self):
        n = 10
        seed = 1234
        for ploidy in [1, 2, 3, 7]:
            ts1 = msprime.sim_ancestry(n * ploidy, ploidy=1, random_seed=seed)
            ts2 = msprime.sim_ancestry(n, ploidy=ploidy, random_seed=seed)
            t1 = ts1.tables
            t2 = ts2.tables
            assert np.allclose(t1.nodes.time * ploidy, t2.nodes.time)
            assert t1.edges == t2.edges

    def test_ploidy_demography(self):
        n = 2
        demography = msprime.Demography.island_model([100] * 2, 0.1)
        for k in [1, 2, 3, 4]:
            ts = msprime.sim_ancestry(
                samples={0: n, 1: n}, ploidy=k, demography=demography
            )
            assert ts.num_samples == 2 * n * k
            assert ts.num_trees == 1
            assert ts.num_sites == 0
            assert ts.sequence_length == 1
            assert ts.num_populations == 2
            assert ts.num_individuals == 2 * n

    def test_ploidy_bottleneck(self):
        def run_sim(ploidy):
            demography = msprime.Demography.isolated_model([100])
            demography.add_instantaneous_bottleneck(time=10, strength=20, population=0)
            return msprime.sim_ancestry(
                [msprime.SampleSet(10, ploidy=1)],
                demography=demography,
                ploidy=ploidy,
                random_seed=1234,
            )

        ts_p1 = run_sim(1)
        ts_p2 = run_sim(2)
        # These should be different simulations and not just simple rescalings
        # of the times, because the bottleneck changes depending on the ploidy.
        # NOTE: this may be brittle: we're depending on the bottleneck actually
        # creating a ternary node in one but not the other. This is just a
        # random chance and may need to be changed in future.
        assert ts_p1.num_nodes != ts_p2.num_nodes

    def test_random_seed(self):
        ts1 = msprime.sim_ancestry(10, random_seed=1)
        ts2 = msprime.sim_ancestry(10, random_seed=1)
        assert tree_sequences_equal(ts1, ts2)

        ts2 = msprime.sim_ancestry(10, random_seed=2)
        assert not tree_sequences_equal(ts1, ts2)

    def test_population_size(self):
        ts1 = msprime.sim_ancestry(10, population_size=1, random_seed=2)
        # Defaults to 1
        ts2 = msprime.sim_ancestry(10, random_seed=2)
        assert tree_sequences_equal(ts1, ts2)
        ts2 = msprime.sim_ancestry(10, population_size=100, random_seed=2)
        # Acts as a simple scaling factor on times.
        assert ts1.tables.edges == ts2.tables.edges
        assert np.allclose(100 * ts1.tables.nodes.time, ts2.tables.nodes.time)

    def test_replicates(self):
        ts = msprime.sim_ancestry(10)
        assert isinstance(ts, tskit.TreeSequence)
        for n in [0, 1, 2, 5]:
            ts_list = list(msprime.sim_ancestry(10, num_replicates=n))
            assert len(ts_list) == n
            for ts in ts_list:
                assert isinstance(ts, tskit.TreeSequence)
                assert ts.num_individuals == 10
                assert ts.num_samples == 20
                assert ts.num_trees == 1

    def test_recombination_rate(self):
        ts = msprime.sim_ancestry(10, recombination_rate=1, sequence_length=10)
        assert ts.num_samples == 20
        assert ts.sequence_length == 10
        assert ts.num_trees > 1
        assert has_discrete_genome(ts)
        # A non-zero recombination_rate and no sequence length is an error
        with pytest.raises(ValueError):
            msprime.sim_ancestry(10, recombination_rate=1)
        # But if we specify a rate map, that's OK.
        rate_map = msprime.RateMap.uniform(sequence_length=10, rate=1)
        ts = msprime.sim_ancestry(10, recombination_rate=rate_map)
        assert ts.num_samples == 20
        assert ts.sequence_length == 10
        assert ts.num_trees > 1
        assert has_discrete_genome(ts)

        # We should get precisely the same ts if we have the same seed
        ts1 = msprime.sim_ancestry(
            10, recombination_rate=1, sequence_length=10, random_seed=1
        )
        ts2 = msprime.sim_ancestry(10, recombination_rate=rate_map, random_seed=1)
        assert tree_sequences_equal(ts1, ts2)

    def test_gene_conversion_rate(self):
        ts = msprime.sim_ancestry(
            10,
            gene_conversion_rate=1,
            gene_conversion_tract_length=2,
            sequence_length=10,
            random_seed=14,
        )
        assert ts.num_trees > 1
        for tree in ts.trees():
            assert tree.num_roots == 1

    def test_gene_conversion_simple_map(self):
        n = 10
        ts = msprime.sim_ancestry(
            n,
            gene_conversion_rate=1,
            gene_conversion_tract_length=1,
            sequence_length=10,
            recombination_rate=1,
            discrete_genome=True,
        )
        assert isinstance(ts, tskit.TreeSequence)
        assert ts.num_samples == 2 * n
        assert ts.num_trees > 1

    def test_gene_conversion_continuous(self):
        with pytest.raises(ValueError):
            msprime.sim_ancestry(
                10,
                gene_conversion_rate=1,
                gene_conversion_tract_length=1,
                discrete_genome=False,
            )

    def test_gene_conversion_default_map(self):
        n = 10
        ts = msprime.sim_ancestry(
            n,
            sequence_length=10,
            gene_conversion_rate=1,
            gene_conversion_tract_length=1,
            discrete_genome=True,
        )
        assert isinstance(ts, tskit.TreeSequence)
        assert ts.num_samples == 2 * n
        assert ts.num_trees > 1

    def test_model(self):
        ts1 = msprime.sim_ancestry(10, population_size=100, random_seed=2)
        ts2 = msprime.sim_ancestry(
            10, population_size=100, model="hudson", random_seed=2
        )
        assert tree_sequences_equal(ts1, ts2)
        ts2 = msprime.sim_ancestry(
            10, population_size=100, model="dtwf", random_seed=2, ploidy=2
        )
        assert not tree_sequences_equal(ts1, ts2)

    def test_continuous_genome(self):
        ts = msprime.sim_ancestry(
            10, recombination_rate=10, sequence_length=1, discrete_genome=False
        )
        assert ts.num_samples == 20
        assert ts.sequence_length == 1
        assert ts.num_trees > 1
        assert not has_discrete_genome(ts)

    def test_discrete_genome(self):
        # Default to discrete_genome=True
        ts = msprime.sim_ancestry(
            10, recombination_rate=10, sequence_length=10, random_seed=2
        )
        assert ts.num_trees > 1
        assert has_discrete_genome(ts)

        ts = msprime.sim_ancestry(
            10,
            recombination_rate=10,
            sequence_length=10,
            random_seed=2,
            discrete_genome=True,
        )
        assert ts.num_trees > 1
        assert has_discrete_genome(ts)

        ts = msprime.sim_ancestry(
            10,
            recombination_rate=1,
            sequence_length=10,
            random_seed=2,
            discrete_genome=False,
        )
        assert ts.num_trees > 1
        assert not has_discrete_genome(ts)

    def test_discrete_genome_migrations(self):
        def run_sim(discrete_genome=None):
            demography = msprime.Demography.stepping_stone_model([1, 1], 0.1)
            return msprime.sim_ancestry(
                samples={0: 5, 1: 5},
                demography=demography,
                sequence_length=5,
                recombination_rate=1,
                discrete_genome=discrete_genome,
                record_migrations=True,
                random_seed=2134,
            )

        ts_discrete = run_sim(True)
        assert ts_discrete.num_trees > 1
        assert ts_discrete.num_migrations > 1
        assert has_discrete_genome(ts_discrete)

        ts_continuous = run_sim(False)
        assert ts_continuous.num_trees > 1
        assert ts_continuous.num_migrations > 1
        assert not has_discrete_genome(ts_continuous)

    def test_numpy_random_seed(self):
        seed = np.array([12345], dtype=np.int64)[0]
        assert seed.dtype == np.int64
        ts1 = msprime.sim_ancestry(10, random_seed=seed)
        ts2 = msprime.sim_ancestry(10, random_seed=seed)
        assert ts1.tables.nodes == ts2.tables.nodes

    def test_record_provenance(self):
        # The content of the provenances is tested elsewhere.
        ts = msprime.sim_ancestry(10, random_seed=2)
        assert ts.num_provenances == 1
        ts = msprime.sim_ancestry(10, random_seed=2, record_provenance=False)
        assert ts.num_provenances == 0

    def test_dtwf(self):
        ts = msprime.sim_ancestry(
            10, population_size=100, model="dtwf", ploidy=2, random_seed=1234
        )
        assert ts.num_trees == 1
        assert ts.first().num_roots == 1
        # All node times should be integers
        time = ts.tables.nodes.time
        assert np.all(time == np.floor(time))

    def test_dtwf_non_diploid(self):
        for ploidy in [1, 3, 7]:
            with pytest.raises(_msprime.LibraryError):
                msprime.sim_ancestry(
                    10, population_size=100, model="dtwf", ploidy=ploidy
                )

    def test_sweep_coalescence(self):
        N = 1e6
        model = msprime.SweepGenicSelection(
            position=0.5,
            start_frequency=1.0 / (2 * N),
            end_frequency=1.0 - (1.0 / (2 * N)),
            s=1000,
            dt=1e-6,
        )
        ts = msprime.sim_ancestry(10, model=model)
        assert ts.num_trees == 1
        assert ts.first().num_roots == 1

    def test_start_time(self):
        ts = msprime.sim_ancestry(10, ploidy=1, random_seed=42, start_time=100)
        assert np.all(ts.tables.nodes.time[10:] > 100)

    def test_end_time(self):
        ts = msprime.sim_ancestry(10, random_seed=42, end_time=0.01)
        assert np.all(ts.tables.nodes.time <= 0.01)
        assert ts.first().num_roots > 1

    def test_start_time_model_duration(self):
        ts = msprime.sim_ancestry(
            10,
            random_seed=42,
            population_size=100,
            start_time=5,
            model=msprime.DiscreteTimeWrightFisher(duration=10),
        )
        assert ts.first().num_roots > 1
        assert np.all(ts.tables.nodes.time <= 15)
        assert ts.tables.nodes.time[-1] == 15

    def test_end_time_model_duration(self):
        # Should take the min of the model duration and end_time
        ts = msprime.sim_ancestry(
            10,
            random_seed=42,
            end_time=1,
            model=msprime.StandardCoalescent(duration=0.01),
        )
        assert np.all(ts.tables.nodes.time <= 0.01)
        assert ts.tables.nodes.time[-1] == 0.01
        assert ts.first().num_roots > 1

        ts = msprime.sim_ancestry(
            10,
            random_seed=42,
            end_time=0.01,
            model=msprime.StandardCoalescent(duration=1),
        )
        assert np.all(ts.tables.nodes.time <= 0.01)
        assert ts.first().num_roots > 1

    def test_model_durations_accumulate(self):
        dt = 0.125
        n = 8
        t = dt * n
        ts = msprime.sim_ancestry(
            10, random_seed=42, model=[msprime.StandardCoalescent(duration=dt)] * n
        )
        assert np.all(ts.tables.nodes.time <= t)
        assert ts.first().num_roots > 1
        assert ts.tables.nodes.time[-1] == t

    def test_record_migrations(self):
        demography = msprime.Demography.stepping_stone_model([1, 1], 0.1)
        samples = {0: 2, 1: 2}
        ts = msprime.sim_ancestry(samples, demography=demography, random_seed=42)
        assert ts.first().num_roots == 1
        # Migrations are off by default
        assert ts.num_migrations == 0

        ts = msprime.sim_ancestry(
            samples, demography=demography, random_seed=42, record_migrations=True
        )
        assert ts.first().num_roots == 1
        # Migrations are off by default
        assert ts.num_migrations > 0

        ts = msprime.sim_ancestry(
            samples, demography=demography, random_seed=42, record_migrations=False
        )
        assert ts.first().num_roots == 1
        assert ts.num_migrations == 0

    def test_record_full_arg(self):
        ts = msprime.sim_ancestry(
            4,
            recombination_rate=1,
            random_seed=2,
            sequence_length=10,
            record_full_arg=True,
        )
        flags = ts.tables.nodes.flags
        assert np.sum(flags == msprime.NODE_IS_RE_EVENT) > 0
        for record_full_arg in [None, False]:
            ts = msprime.sim_ancestry(
                4,
                recombination_rate=1,
                random_seed=2,
                sequence_length=10,
                record_full_arg=record_full_arg,
            )
            flags = ts.tables.nodes.flags
            assert np.sum(flags == msprime.NODE_IS_RE_EVENT) == 0

    def test_initial_tables_recapitate(self):
        # Simple recapitate scenario
        ts = msprime.sim_ancestry(5, end_time=0.5, random_seed=53)
        assert ts.first().num_roots > 1
        recapitated1 = msprime.sim_ancestry(
            initial_state=ts, population_size=1, random_seed=234
        )
        assert recapitated1.num_trees == 1
        assert recapitated1.first().num_roots == 1

        # We should get the same answer from the providing the tables argument.
        recapitated2 = msprime.sim_ancestry(
            initial_state=ts.tables, population_size=1, random_seed=234
        )
        assert tree_sequences_equal(recapitated1, recapitated2)

    def test_replicate_index(self):
        reps = msprime.sim_ancestry(12, random_seed=52, num_replicates=10)
        for index, ts1 in enumerate(reps):
            ts2 = msprime.sim_ancestry(12, random_seed=52, replicate_index=index)
            t1 = ts1.dump_tables()
            t2 = ts2.dump_tables()
            t1.provenances.clear()
            t2.provenances.clear()
            assert t1 == t2

    def test_bad_additional_node_flag(self):
        with pytest.raises(ValueError):
            msprime.sim_ancestry(
                10,
                additional_nodes=msprime.NodeType.COMMON_ANCESTOR,
            )


class TestSimulateInterface:
    """
    Some simple test cases for the simulate() interface.
    """

    def test_defaults(self):
        n = 10
        ts = msprime.simulate(n)
        assert isinstance(ts, tskit.TreeSequence)
        assert ts.get_sample_size() == n
        assert ts.get_num_trees() == 1
        assert ts.get_num_mutations() == 0
        assert ts.get_sequence_length() == 1
        assert len(list(ts.provenances())) == 1

    def test_positional_args_not_allowed(self):
        with pytest.raises(TypeError):
            msprime.simulate(2, 100)

    def verify_provenance(self, provenance):
        """
        Checks that the specified provenance object has the right sort of
        properties.
        """
        # Generate the ISO 8601 time for now, without the high precision suffix,
        # and compare the prefixes.
        today = datetime.date.today().isoformat()
        k = len(today)
        assert provenance.timestamp[:k] == today
        assert provenance.timestamp[k] == "T"
        d = json.loads(provenance.record)
        assert len(d) > 0

    def test_provenance(self):
        ts = msprime.simulate(10)
        assert ts.num_provenances == 1
        self.verify_provenance(ts.provenance(0))
        for ts in msprime.simulate(10, num_replicates=10):
            assert ts.num_provenances == 1
            self.verify_provenance(ts.provenance(0))

    def test_end_time(self):
        ts = msprime.simulate(15, recombination_rate=2, random_seed=8, end_time=0.1)
        for tree in ts.trees():
            for root in tree.roots:
                assert tree.time(root) == 0.1

    def test_replicates(self):
        n = 20
        num_replicates = 10
        count = 0
        for ts in msprime.simulate(n, num_replicates=num_replicates):
            count += 1
            assert isinstance(ts, tskit.TreeSequence)
            assert ts.get_sample_size() == n
            assert ts.get_num_trees() == 1
        assert num_replicates == count

    def test_mutations(self):
        n = 10
        ts = msprime.simulate(n, mutation_rate=10)
        assert isinstance(ts, tskit.TreeSequence)
        assert ts.get_sample_size() == n
        assert ts.get_num_trees() == 1
        assert ts.get_num_mutations() > 0

    def test_no_mutations_with_start_time(self):
        with pytest.raises(ValueError):
            msprime.simulate(10, mutation_rate=10, start_time=3)
        # But fine if we set start_time = None
        ts = msprime.simulate(10, mutation_rate=10, start_time=None, random_seed=1)
        assert ts.num_sites > 0

    def test_mutation_generator_unsupported(self):
        with pytest.raises(ValueError):
            msprime.simulate(10, mutation_generator="some non-None value")

    def test_mutation_interface(self):
        for bad_type in [{}, [], self]:
            with pytest.raises(TypeError):
                msprime.simulate(10, mutation_rate=bad_type)
        for bad_value in ["x", "234x"]:
            with pytest.raises(ValueError):
                msprime.simulate(10, mutation_rate=bad_value)

    def test_recombination(self):
        n = 10
        ts = msprime.simulate(n, recombination_rate=10)
        assert isinstance(ts, tskit.TreeSequence)
        assert ts.sample_size == n
        assert ts.num_trees > 1
        assert ts.num_mutations == 0

    def test_num_labels(self):
        # Running simulations with different numbers of labels in the default
        # setting should have no effect.
        tables = [
            msprime.simulate(10, num_labels=num_labels, random_seed=1).dump_tables()
            for num_labels in range(1, 5)
        ]
        for t in tables:
            t.provenances.clear()
        for t in tables:
            assert t == tables[0]

    def test_replicate_index(self):
        reps = msprime.simulate(10, random_seed=52, num_replicates=10)
        for index, ts1 in enumerate(reps):
            ts2 = msprime.simulate(10, random_seed=52, replicate_index=index)
            t1 = ts1.dump_tables()
            t2 = ts2.dump_tables()
            t1.provenances.clear()
            t2.provenances.clear()
            assert t1 == t2


class TestDiscreteGenomeSimulate:
    """
    Test that the num_loci==sequence_length trick for doing a discrete genome
    simulation in simulate() still works.
    """

    def test_simple_example(self):
        rmap = msprime.RecombinationMap.uniform_map(100, 0.1, num_loci=100)
        ts = msprime.simulate(10, recombination_map=rmap, random_seed=1)
        assert has_discrete_genome(ts)
        assert ts.sequence_length == 100
        assert ts.num_trees > 1

    def test_single_locus(self):
        rmap = msprime.RecombinationMap.uniform_map(1, 100, num_loci=1)
        ts = msprime.simulate(10, recombination_map=rmap, random_seed=1)
        assert has_discrete_genome(ts)
        assert ts.sequence_length == 1
        assert ts.num_trees == 1

    @pytest.mark.parametrize("m", [2, 5, 7])
    def test_small_examples(self, m):
        rmap = msprime.RecombinationMap.uniform_map(m, 100, num_loci=m)
        ts = msprime.simulate(10, recombination_map=rmap, random_seed=1)
        assert has_discrete_genome(ts)
        assert ts.sequence_length == m
        assert ts.num_trees == m

    def test_simulate_from_example(self):
        # Example from the docs on a WF simulation
        N = 10
        num_loci = 2
        recomb_map = msprime.RecombinationMap.uniform_map(
            num_loci, 1 / num_loci, num_loci
        )
        ts = msprime.simulate(5, Ne=N / 2, recombination_map=recomb_map, random_seed=5)
        assert has_discrete_genome(ts)
        assert ts.sequence_length == num_loci
        assert ts.num_trees > 1


class TestReprRoundTrip:
    """
    Tests that we can eval the repr of objects to round trip them.
    """

    def assert_repr_round_trip(self, obj_list):
        for obj in obj_list:
            obj_copy = eval(repr(obj), globals(), msprime.__dict__)
            assert obj_copy == obj
            assert obj_copy is not obj

    def test_population(self):
        examples = [
            msprime.Population(initial_size=2),
            msprime.Population(2, growth_rate=5),
            msprime.Population(initial_size=234, growth_rate=10),
        ]
        self.assert_repr_round_trip(examples)

    def test_population_parameters_change(self):
        examples = [
            msprime.PopulationParametersChange(time=1, initial_size=1),
            msprime.PopulationParametersChange(time=1, growth_rate=2),
            msprime.PopulationParametersChange(time=1, growth_rate=1, population=2),
            msprime.PopulationParametersChange(
                time=3, initial_size=3, growth_rate=1, population=2
            ),
        ]
        self.assert_repr_round_trip(examples)

    def test_migration_rate_change(self):
        examples = [
            msprime.MigrationRateChange(time=1, rate=1),
            msprime.MigrationRateChange(time=1, rate=1, source=1, dest=2),
        ]
        self.assert_repr_round_trip(examples)

    def test_mass_migration(self):
        examples = [
            msprime.MassMigration(time=1, source=1, dest=2),
            msprime.MassMigration(time=1, source=1, dest=2, proportion=0.2),
        ]
        self.assert_repr_round_trip(examples)

    def test_simulation_model_change(self):
        examples = [
            msprime.SimulationModelChange(),
            msprime.SimulationModelChange(model="hudson"),
            msprime.SimulationModelChange(model=msprime.DiscreteTimeWrightFisher()),
            msprime.SimulationModelChange(
                model=msprime.BetaCoalescent(alpha=1, truncation_point=2)
            ),
        ]
        self.assert_repr_round_trip(examples)

    def test_simple_bottleneck(self):
        examples = [
            msprime.SimpleBottleneck(time=10, population=2),
            msprime.SimpleBottleneck(time=10, population=2, proportion=0.5),
        ]
        self.assert_repr_round_trip(examples)

    def test_instantaneous_bottleneck(self):
        examples = [
            msprime.InstantaneousBottleneck(time=10, population=1),
            msprime.InstantaneousBottleneck(time=10, population=1, strength=10),
        ]
        self.assert_repr_round_trip(examples)

    def test_census_event(self):
        examples = [
            msprime.CensusEvent(time=10),
        ]
        self.assert_repr_round_trip(examples)

    def test_simulation_models(self):
        examples = [
            msprime.StandardCoalescent(),
            msprime.SmcApproxCoalescent(),
            msprime.SmcPrimeApproxCoalescent(),
            msprime.DiscreteTimeWrightFisher(),
            msprime.FixedPedigree(),
            msprime.BetaCoalescent(),
            msprime.BetaCoalescent(alpha=1, truncation_point=10),
            msprime.DiracCoalescent(),
            msprime.DiracCoalescent(psi=1234, c=56),
            msprime.SweepGenicSelection(
                position=1, start_frequency=0.5, end_frequency=0.9, s=0.1, dt=1e-4
            ),
        ]
        self.assert_repr_round_trip(examples)


class TestUnknownGenomeRegions:
    def test_sim_ancestry_unknown_flanks(self):
        rate_map = msprime.RateMap(position=[0, 1, 9, 10], rate=[np.nan, 0, np.nan])
        ts = msprime.sim_ancestry(2, recombination_rate=rate_map, random_seed=1)
        assert ts.num_trees == 3
        tree = ts.first()
        assert tree.interval == (0, 1)
        assert tree.num_roots == ts.num_samples
        tree.next()
        assert tree.interval == (1, 9)
        assert tree.num_roots == 1
        tree.next()
        assert tree.interval == (9, 10)
        assert tree.num_roots == ts.num_samples

    def test_sim_ancestry_unknown_left(self):
        rate_map = msprime.RateMap(position=[0, 1, 10], rate=[np.nan, 0])
        ts = msprime.sim_ancestry(2, recombination_rate=rate_map, random_seed=1)
        assert ts.num_trees == 2
        tree = ts.first()
        assert tree.interval == (0, 1)
        assert tree.num_roots == ts.num_samples
        tree.next()
        assert tree.interval == (1, 10)
        assert tree.num_roots == 1

    def test_sim_ancestry_unknown_right(self):
        rate_map = msprime.RateMap(position=[0, 9, 10], rate=[0, np.nan])
        ts = msprime.sim_ancestry(2, recombination_rate=rate_map, random_seed=1)
        assert ts.num_trees == 2
        tree = ts.first()
        assert tree.interval == (0, 9)
        assert tree.num_roots == 1
        tree.next()
        assert tree.interval == (9, 10)
        assert tree.num_roots == ts.num_samples

    def test_sim_ancestry_unknown_left_full_arg(self):
        rate_map = msprime.RateMap(position=[0, 1, 10], rate=[np.nan, 1])
        ts = msprime.sim_ancestry(
            4, recombination_rate=rate_map, record_full_arg=True, random_seed=1
        )
        assert ts.num_trees > 2
        tree = ts.first()
        assert tree.interval == (0, 1)
        assert tree.num_roots == ts.num_samples
        flags = ts.tables.nodes.flags
        assert np.sum(flags == msprime.NODE_IS_RE_EVENT) > 0

    def test_sim_ancestry_unknown_mid(self):
        rate_map = msprime.RateMap(position=[0, 1, 9, 10], rate=[0, np.nan, 0])
        with pytest.raises(ValueError, match="Missing regions of the genome"):
            msprime.sim_ancestry(2, recombination_rate=rate_map, random_seed=1)

    def test_sim_ancestry_unknown_mid_plus_flanks(self):
        rate_map = msprime.RateMap(
            position=[0, 1, 5, 6, 7, 10], rate=[np.nan, 0, np.nan, 0, np.nan]
        )
        with pytest.raises(ValueError, match="Missing regions of the genome"):
            msprime.sim_ancestry(2, recombination_rate=rate_map, random_seed=1)


class TestTimeUnits:
    def test_simulate_gives_generations(self):
        ts = msprime.simulate(2, random_seed=1)
        assert ts.time_units == "generations"

    def test_sim_ancestry_gives_generations(self):
        ts = msprime.sim_ancestry(2, random_seed=1)
        assert ts.time_units == "generations"

    @pytest.mark.parametrize("time_units", ["", "unknown", "mya", "ticks"])
    def test_initial_state_warns_not_generations(self, time_units):
        tables = msprime.sim_ancestry(2, end_time=0, random_seed=1).dump_tables()
        tables.time_units = time_units
        with pytest.warns(msprime.TimeUnitsMismatchWarning, match="time_units"):
            ts = msprime.sim_ancestry(
                initial_state=tables, population_size=10, random_seed=1
            )
        assert ts.time_units == time_units

    def test_initial_state_suppress_message(self):
        tables = msprime.sim_ancestry(2, end_time=0, random_seed=1).dump_tables()
        tables.time_units = "ticks"
        with warnings.catch_warnings(record=True) as w:
            with warnings.catch_warnings():
                warnings.simplefilter("ignore", msprime.TimeUnitsMismatchWarning)
                msprime.sim_ancestry(
                    initial_state=tables, population_size=10, random_seed=1
                )
        assert len(w) == 0

    def test_initial_state_errors_uncalibrated(self):
        tables = msprime.sim_ancestry(2, end_time=0, random_seed=1).dump_tables()
        tables.time_units = tskit.TIME_UNITS_UNCALIBRATED
        with pytest.raises(ValueError, match="time_units"):
            msprime.sim_ancestry(initial_state=tables, population_size=10, random_seed=1)


class TestSMCK:
    @pytest.mark.parametrize("seed", [4512, 873561])
    def test_discrete(self, seed):
        tss = msprime.sim_ancestry(
            samples=10,
            model=msprime.SMCK(1),
            recombination_rate=0.005,
            sequence_length=1000,
            random_seed=seed,
            num_replicates=10,
        )
        for ts in tss:
            assert max(tree.num_roots for tree in ts.trees()) == 1

    def test_continuous(self):
        tss = msprime.sim_ancestry(
            samples=10,
            model=msprime.SMCK(1),
            recombination_rate=0.005,
            sequence_length=1000,
            num_replicates=10,
            random_seed=7598782,
            discrete_genome=False,
        )
        for ts in tss:
            assert max(tree.num_roots for tree in ts.trees()) == 1

    def test_ancient_samples(self):
        tables = tskit.TableCollection(100)
        tables.time_units = "generations"
        tables.populations.add_row()
        tables.nodes.add_row(flags=1, time=0.0, population=0)
        tables.nodes.add_row(flags=1, time=1.0, population=0)
        tables.nodes.add_row(flags=1, time=10.0, population=0)
        tables.nodes.add_row(flags=1, time=30.0, population=0)
        tables.populations.metadata_schema = tskit.MetadataSchema.permissive_json()
        ts = msprime.sim_ancestry(
            initial_state=tables,
            population_size=10_000,
            model=msprime.SMCK(1),
            recombination_rate=1e-6,
        )
        for tree in ts.trees():
            assert tree.num_roots == 1

    def test_model_switch(self):
        ts = msprime.sim_ancestry(
            samples=10,
            population_size=10_000,
            model=[
                msprime.StandardCoalescent(duration=10),
                msprime.SMCK(1, duration=10),
                msprime.StandardCoalescent(),
            ],
            random_seed=10,
            recombination_rate=1e-5,
            sequence_length=100,
        )
        for tree in ts.trees():
            assert tree.num_roots == 1

    def test_model_switch_high_rec(self):
        ts = msprime.sim_ancestry(
            samples=10,
            population_size=10_000,
            model=[
                msprime.StandardCoalescent(duration=100),
                msprime.SMCK(1),
            ],
            random_seed=10,
            recombination_rate=1e-4,
            sequence_length=100,
        )
        for tree in ts.trees():
            assert tree.num_roots == 1

    @pytest.mark.parametrize("hull_offset", [2, 0.5, 1e-6, 2.133])
    @pytest.mark.parametrize("discrete_genome", [True, False])
    def test_smc_k_plus(self, hull_offset, discrete_genome):
        tss = msprime.sim_ancestry(
            samples=10,
            population_size=10_000,
            model=msprime.SMCK(hull_offset),
            random_seed=10,
            recombination_rate=1e-5,
            sequence_length=100,
            num_replicates=10,
            discrete_genome=discrete_genome,
        )
        for ts in tss:
            assert ts.num_trees > 1
            for tree in ts.trees():
                assert tree.num_roots == 1

    def test_two_pops(self):
        demography = msprime.Demography()
        demography.add_population(initial_size=10_000)
        demography.add_population(initial_size=10_000)
        demography.add_mass_migration(1e6, source=1, dest=0, proportion=1.0)
        ts = msprime.sim_ancestry(
            samples={0: 2, 1: 2},
            demography=demography,
            model=msprime.SMCK(1),
            random_seed=74024,
            recombination_rate=1e-5,
            sequence_length=100,
        )
        for tree in ts.trees():
            assert tree.num_roots == 1

    def test_gc(self):

        with pytest.raises(
            ValueError,
            match="Gene conversion is not supported for the SMCK model. "
            "Please refer to issue #2399 on GitHub for details.",
        ):
            msprime.sim_ancestry(
                samples=10,
                model=msprime.SMCK(1),
                sequence_length=100,
                gene_conversion_rate=1.0,
                gene_conversion_tract_length=5,
                num_replicates=5,
                recombination_rate=0.01,
            )
