"""
Script for statistically evaluating various aspects of tsinfer performance.
"""
import argparse
import sys
import collections
import random
import concurrent.futures
import time
import warnings
import os.path
import colorsys
import json

import numpy as np
import pandas as pd
import matplotlib as mp
# Force matplotlib to not use any Xwindows backend.
mp.use('Agg')
import matplotlib.pyplot as plt
from matplotlib import collections as mc
import seaborn as sns
import tqdm
import scipy.stats
import colorutils


import tsinfer
import tsinfer.cli as cli
import msprime

def make_errors(v, p):
    """
    For each sample an error occurs with probability p. Errors are generated by
    sampling values from the stationary distribution, that is, if we have an
    allele frequency of f, a 1 is emitted with probability f and a
    0 with probability 1 - f. Thus, there is a possibility that an 'error'
    will in fact result in the same value.
    """
    w = np.copy(v)
    if p > 0:
        m = v.shape[0]
        frequency = np.sum(v) / m
        # Randomly choose samples with probability p
        samples = np.where(np.random.random(m) < p)[0]
        # Generate observations from the stationary distribution.
        errors = (np.random.random(samples.shape[0]) < frequency).astype(int)
        w[samples] = errors
    return w


def generate_samples(ts, error_p):
    """
    Returns samples with a bits flipped with a specified probability.

    Rejects any variants that result in a fixed column.
    """
    G = np.zeros((ts.num_sites, ts.num_samples), dtype=np.int8)
    for variant in ts.variants():
        done = False
        # Reject any columns that have no 1s or no zeros
        while not done:
            G[variant.index] = make_errors(variant.genotypes, error_p)
            s = np.sum(G[variant.index])
            done = 0 < s < ts.sample_size
    return G


def infer_from_simulation(
        ts, recombination_rate=1, input_error=0, sample_error=0, engine="C",
        path_compression=True):
    if input_error == 0:
        genotypes = ts.genotype_matrix()
    else:
        genotypes = generate_samples(ts, input_error)
    positions = [mut.position for mut in ts.mutations()]
    return tsinfer.infer(
        genotypes, positions=positions, sequence_length=ts.sequence_length,
        recombination_rate=recombination_rate, sample_error=sample_error,
        engine=engine, path_compression=path_compression)


def get_mean_rf_distance(ts1, ts2):
    """
    Returns the mean distance between the trees in the specified tree sequences.
    """
    assert ts1.sample_size == ts2.sample_size
    assert ts1.sequence_length == ts2.sequence_length
    trees1 = []
    intervals1 = []
    trees2 = []
    intervals2 = []
    tns = dendropy.TaxonNamespace()
    for t in ts1.trees():
        dt = dendropy.Tree.get(data=t.newick(), schema="newick", taxon_namespace=tns)
        trees1.append(dt)
        intervals1.append(t.interval)
    assert len(trees1) == ts1.num_trees
    for t in ts2.trees():
        dt = dendropy.Tree.get(data=t.newick(), schema="newick", taxon_namespace=tns)
        trees2.append(dt)
        intervals2.append(t.interval)
    assert len(trees2) == ts2.num_trees
    j1 = 0
    j2 = 0
    total_distance = 0
    total_metric = 0
    # I haven't tested this algorithm thoroughly, so there might be corner cases
    # not handled correctly. However, the total_distance assert below should
    # catch the problem if it occurs.
    while j1 < len(trees1) and j2 < len(trees2):
        # Each iteration of this loop considers one overlapping interval and
        # increments the counters.
        l1, r1 = intervals1[j1]
        l2, r2 = intervals2[j2]
        l = max(l1, l2)
        r = min(r1, r2)
        rf_distance = dendropy.calculate.treecompare.symmetric_difference(
                trees1[j1], trees2[j2])
        total_metric += rf_distance * (r - l)
        total_distance += r - l
        if r1 <= r2:
            j1 += 1
        if r1 >= r2:
            j2 += 1
    # assert total_distance, ts1.sequence_length)
    return total_metric / total_distance


def check_basic_performance():
    # Basic check to ensure that we get the same results regardless of the recombination
    # rate

    num_samples = 10
    MB = 10**6
    sites = []
    trees = []
    edges = []
    nodes = []
    rf_distance = []
    recombination_rate = []
    for seed in range(1, 10):
        ts_source = msprime.simulate(
            num_samples, length=1*MB, Ne=10**4, recombination_rate=1e-8, mutation_rate=1e-8,
            random_seed=seed)
        print("sim: n = ",
                ts_source.num_samples, ", m =", ts_source.num_sites,
                "num_trees = ", ts_source.num_trees)
        for exponent in range(0, 6):
            infer_recomb_rate = 10**(-exponent)
            ts_inferred = infer_from_simulation(
                ts_source, recombination_rate=infer_recomb_rate, sample_error=0)
            recombination_rate.append(infer_recomb_rate)
            rf = get_mean_rf_distance(ts_source, ts_inferred)
            rf_distance.append(get_mean_rf_distance(ts_source, ts_inferred))
            sites.append(ts_source.num_sites)
            trees.append(ts_inferred.num_trees / ts_source.num_trees)
            nodes.append(ts_inferred.num_nodes / ts_source.num_nodes)
            edges.append(ts_inferred.num_edges / ts_source.num_edges)

            # print("recombination_rate = ", recombination_rate)
            # print("mean rf = ", rf)
            # print("num_trees = ", ts_inferred.num_trees / ts_source.num_trees)
            # print("num_edges= ", ts_inferred.num_edges/ ts_source.num_edges)
    df = pd.DataFrame(data={
        "num_samples": num_samples,
        "recombination_rate": recombination_rate,
        "sites": sites,
        "trees": trees,
        "edges": edges,
        "nodes": nodes,
        "rf_distance": rf_distance})
    for y_value in ["trees", "nodes", "edges", "rf_distance"]:
        plt.figure()
        plot = sns.boxplot(x=df["recombination_rate"], y=df[y_value])
        plt.savefig("tmp__NOBACKUP__/{}.png".format(y_value))


def check_effect_error_param(input_error=0.0):
    num_samples = 10
    MB = 10**6
    sites = []
    trees = []
    edges = []
    nodes = []
    rf_distance = []
    error = []
    for seed in range(1, 10):
        ts_source = msprime.simulate(
            num_samples, length=1*MB, Ne=10**4, recombination_rate=1e-8, mutation_rate=1e-8,
            random_seed=seed)
        print("sim: n = ",
                ts_source.num_samples, ", m =", ts_source.num_sites,
                "num_trees = ", ts_source.num_trees)
        for err in [0] + list(10.0**-np.arange(1, 5)):
            ts_inferred = infer_from_simulation(
                ts_source, recombination_rate=1, input_error=input_error, sample_error=err)
            error.append(err)
            rf = get_mean_rf_distance(ts_source, ts_inferred)
            rf_distance.append(get_mean_rf_distance(ts_source, ts_inferred))
            sites.append(ts_source.num_sites)
            trees.append(ts_inferred.num_trees / ts_source.num_trees)
            nodes.append(ts_inferred.num_nodes / ts_source.num_nodes)
            edges.append(ts_inferred.num_edges / ts_source.num_edges)
        df = pd.DataFrame(data={
            "num_samples": num_samples,
            "error": error,
            "sites": sites,
            "trees": trees,
            "edges": edges,
            "nodes": nodes,
            "rf_distance": rf_distance})
        df.to_csv("tmp__NOBACKUP__/error_{}.csv".format(input_error))
        # df = pd.read_csv("tmp__NOBACKUP__/error_{}.csv".format(input_error))
        plt.figure()
        fig, axs = plt.subplots(nrows=2, ncols=2)
        for j, y_value in enumerate(["trees", "nodes", "edges", "rf_distance"]):
            sns.boxplot(x=df["error"], y=df[y_value], ax=axs[j // 2, j % 2])
        plt.suptitle("input_error = {}".format(input_error))
        plt.savefig("tmp__NOBACKUP__/error_{}.png".format(input_error))
        plt.clf()


def check_variable_recomb():
    rate = 1e-10
    Mb = 10**6
    num_samples = 20
    seed = 10
    recomb_map = msprime.RecombinationMap(
        positions=[0, 1*Mb, 1.1 * Mb, 2 * Mb],
        rates = [rate, 100 * rate, rate, 0], num_loci=10000)
    sites = []
    trees = []
    edges = []
    nodes = []
    rf_distance = []
    flat_rate = []
    for seed in range(1, 100):
        ts_source = msprime.simulate(
            num_samples, recombination_map=recomb_map, Ne=10**4, mutation_rate=1e-8,
            random_seed=seed,)
        print(
            "sim: n = ", ts_source.num_samples, ", m =", ts_source.num_sites,
            "num_trees = ", ts_source.num_trees)
        samples = np.zeros((ts_source.num_samples, ts_source.num_sites), dtype=np.int8)
        for variant in ts_source.variants():
            samples[:, variant.index] = variant.genotypes
        positions = np.array([site.position for site in ts_source.sites()])
        genetic_positions = np.array([
            recomb_map.physical_to_genetic(x) for x in positions])
        recombination_rate = np.zeros(ts_source.num_sites)
        recombination_rate[1:] = genetic_positions[1:] - genetic_positions[:-1]
        recombination_rate /= recomb_map.get_num_loci()

        for recomb_rate in [1, recombination_rate]:
            ts_inferred = tsinfer.infer(
                samples=samples, positions=positions,
                sequence_length=ts_source.sequence_length,
                recombination_rate=recomb_rate, sample_error=0)
            rf = get_mean_rf_distance(ts_source, ts_inferred)
            rf_distance.append(get_mean_rf_distance(ts_source, ts_inferred))
            flat_rate.append(recomb_rate is 1)
            sites.append(ts_source.num_sites)
            trees.append(ts_inferred.num_trees / ts_source.num_trees)
            nodes.append(ts_inferred.num_nodes / ts_source.num_nodes)
            edges.append(ts_inferred.num_edges / ts_source.num_edges)
        df = pd.DataFrame(data={
            "num_samples": num_samples,
            "sites": sites,
            "trees": trees,
            "edges": edges,
            "nodes": nodes,
            "rf_distance": rf_distance,
            "flat_rate": flat_rate})
            # df.to_csv("tmp__NOBACKUP__/error_{}.csv".format(input_error))
    plt.figure()
    fig, axs = plt.subplots(nrows=2, ncols=2)
    for j, y_value in enumerate(["trees", "nodes", "edges", "rf_distance"]):
        sns.boxplot(x=df["flat_rate"], y=df[y_value], ax=axs[j // 2, j % 2])
    plt.suptitle("variable recomb")
    plt.savefig("tmp__NOBACKUP__/variable_recomb.png")
    plt.clf()

    # print(ts_source.num_sites)
    # print(ts_source.num_trees)
    # for t in ts_source.trees():
    #     print(t.interval[1] - t.interval[0])


# TODO these are really unit tests. Move them into the tests directory.
def check_single_tree_one_mutation_per_branch():
    for num_samples in [4]:
        num_samples = 10
        for seed in range(1, 10):
            ts_source = msprime.simulate(num_samples, length=2 * num_samples, random_seed=seed)
            # Put a mutation on every branch.
            t = ts_source.dump_tables()
            for u in range(ts_source.num_nodes - 1):
                t.sites.add_row(position=u, ancestral_state="0")
                t.mutations.add_row(site=u, node=u, derived_state="1")
            ts_source = msprime.load_tables(**t.asdict())
            ts_inferred = infer_from_simulation(ts_source)
            assert ts_inferred.num_trees == 1


def verify_trees_equal(ts_source, ts_inferred):
    """
    Verifies that the inferred trees are topologically equal to the specified
    inferred trees.
    """
    for t in ts_source.trees():
        print(t.interval)
        print(t.draw(format="unicode"))
    print("++")
    for t in ts_inferred.trees():
        print(t.interval)
        print(t.draw(format="unicode"))

    print("==")


# TODO these are really unit tests. Move them into the tests directory.
def check_many_trees_one_mutation_per_branch():
    # for num_samples in [10, 100, 1000, 10000]:

    import daiquiri
    daiquiri.setup(level="DEBUG", outputs=(daiquiri.output.Stream(sys.stdout),))
    num_samples = 4
    print("num_samples = ", num_samples)
    # for seed in range(1, 1000):
    for seed in [19]:
        print(seed)
        recombination_map = msprime.RecombinationMap.uniform_map(
                100, 0.002, num_loci=100)
        ts_source = msprime.simulate(
                num_samples, recombination_map=recombination_map,
                random_seed=seed, model="smc_prime")
        # Put a mutation on every branch.
        tables = ts_source.dump_tables()
        j = 0
        for tree in ts_source.trees():
            left, right = tree.interval
            n = len(list(tree.nodes()))
            delta = (right - left) / n
            x = left
            for u in tree.nodes():
                if u != tree.root:
                    tables.sites.add_row(position=x, ancestral_state="0")
                    tables.mutations.add_row(site=j, node=u, derived_state="1")
                    j += 1
                    x += delta
        print(tables)
        ts_source = msprime.load_tables(**tables.asdict())
        print("samples")
        print(ts_source.genotype_matrix().T)
        ts_inferred = infer_from_simulation(ts_source,
                path_compression=False,
                # engine="Py-matrix")
                engine="P")
        print("num_trees", ts_source.num_trees, ts_inferred.num_trees)
        print("num_edges", ts_source.num_edges, ts_inferred.num_edges)
        verify_trees_equal(ts_source, ts_inferred)
        print(ts_inferred.tables)
        assert ts_source.num_trees >= ts_inferred.num_trees
        assert ts_source.num_edges >= ts_inferred.num_edges


# TODO these are really unit tests. Move them into the tests directory.
def check_single_tree_one_mutation_per_branch():
    for num_samples in [10, 100]:
        print("num_samples = ", num_samples)
        for seed in range(1, 10):
            ts_source = msprime.simulate(num_samples, length=2 * num_samples, random_seed=seed)
            # Put a mutation on every branch.
            t = ts_source.dump_tables()
            for u in range(ts_source.num_nodes - 1):
                t.sites.add_row(position=u, ancestral_state="0")
                t.mutations.add_row(site=u, node=u, derived_state="1")
            ts_source = msprime.load_tables(**t.asdict())
            ts_inferred = infer_from_simulation(ts_source)
            assert ts_inferred.num_trees == 1

def check_single_tree_two_mutations_per_branch():
    for num_samples in [10, 100]:
        print("num_samples = ", num_samples)
        num_mutations = 2 * num_samples * 2
        d = 1 / num_mutations
        for seed in range(1, 10):
            ts_source = msprime.simulate(num_samples, length=num_mutations, random_seed=seed)
            # Put a mutation on every branch.
            t = ts_source.dump_tables()
            j = 0
            for u in range(ts_source.num_nodes - 1):
                for _ in range(2):
                    t.sites.add_row(position=j * d, ancestral_state="0")
                    t.mutations.add_row(site=j, node=u, derived_state="1")
                    j += 1
            ts_source = msprime.load_tables(**t.asdict())
            ts_inferred = infer_from_simulation(ts_source)
            assert ts_inferred.num_trees == 1

def check_single_tree_many_mutations_per_branch():
    num_samples = 10
    for k in range(2, 20):
        print("k = ", k)
        num_mutations = 2 * num_samples * k
        d = 1 / num_mutations
        for seed in range(1, 10):
            ts_source = msprime.simulate(num_samples, length=num_mutations, random_seed=seed)
            # Put a mutation on every branch.
            t = ts_source.dump_tables()
            j = 0
            for u in range(ts_source.num_nodes - 1):
                for _ in range(k):
                    t.sites.add_row(position=j * d, ancestral_state="0")
                    t.mutations.add_row(site=j, node=u, derived_state="1")
                    j += 1
            ts_source = msprime.load_tables(**t.asdict())
            # print(ts_source.tables)
            ts_inferred = infer_from_simulation(ts_source)
            assert ts_inferred.num_trees == 1


def check_single_tree_high_mutation_rate():
    # import daiquiri
    # daiquiri.setup(level="DEBUG")
    # TODO change this test so that we throw a Poisson number of mutations on
    # each branch. It's very hard to have mutations on _every_ branch for larger
    # tree sizes.
    for num_samples in [3, 10, 50]:
        for seed in range(1, 10):
            ts_source = msprime.simulate(num_samples, random_seed=seed, mutation_rate=5000)
            print("sim = ", num_samples, ts_source.num_sites, seed)
            nodes = set()
            for site in ts_source.sites():
                for mutation in site.mutations:
                    nodes.add(mutation.node)
            # if nodes != set(range(ts_source.num_nodes - 1)):
            #     continue
            assert nodes == set(range(ts_source.num_nodes - 1))
            ts_inferred = infer_from_simulation(ts_source)
            # for t in ts_source.trees():
            #     print(t.draw(format="unicode"))
            # print(ts_inferred.num_trees)
            # for t in ts_inferred.trees():
            #     print(t.draw(format="unicode"))
            assert ts_inferred.num_trees == 1


##############################
# Updated code to work with the CLI
##############################


def run_infer(ts, engine="C", path_compression=True, exact_ancestors=False):
    """
    Runs the perfect inference process on the specified tree sequence.
    """

    sample_data = tsinfer.SampleData.from_tree_sequence(ts)

    if exact_ancestors:
        ancestor_data = tsinfer.AncestorData(sample_data)
        tsinfer.build_simulated_ancestors(sample_data, ancestor_data, ts)
        ancestor_data.finalise()
    else:
        ancestor_data = tsinfer.generate_ancestors(sample_data)

    ancestors_ts = tsinfer.match_ancestors(
        sample_data, ancestor_data, path_compression=path_compression,
        engine=engine)
    inferred_ts = tsinfer.match_samples(
        sample_data, ancestors_ts, path_compression=path_compression,
        engine=engine)
    return inferred_ts

def edges_performance_worker(args):
    simulation_args, tree_metrics = args
    before = time.perf_counter()
    smc_ts = msprime.simulate(**simulation_args)
    sim_time = time.perf_counter() - before

    tmp_ts = tsinfer.strip_singletons(smc_ts)
    if tmp_ts.num_sites == 0:
        warnings.warn("Dropping simulation with no variants")
        return {}

    before = time.perf_counter()
    estimated_ancestors_ts = run_infer(smc_ts, exact_ancestors=False)
    estimated_ancestors_time = time.perf_counter() - before
    num_children = []
    for edgeset in estimated_ancestors_ts.edgesets():
        num_children.append(len(edgeset.children))
    estimated_ancestors_num_children = np.array(num_children)

    before = time.perf_counter()
    exact_ancestors_ts = run_infer(smc_ts, exact_ancestors=True)
    exact_ancestors_time = time.perf_counter() - before
    num_children = []
    for edgeset in exact_ancestors_ts.edgesets():
        num_children.append(len(edgeset.children))
    exact_ancestors_num_children = np.array(num_children)

    results = {
        "sim_time": sim_time,
        "estimated_anc_time": estimated_ancestors_time,
        "exact_anc_time": exact_ancestors_time,
        "num_sites": smc_ts.num_sites,
        "source_num_trees": smc_ts.num_trees,
        "estimated_anc_trees": estimated_ancestors_ts.num_trees,
        "exact_anc_trees": exact_ancestors_ts.num_trees,
        "source_edges": smc_ts.num_edges,
        "estimated_anc_edges": estimated_ancestors_ts.num_edges,
        "exact_anc_edges": exact_ancestors_ts.num_edges,
        "estimated_anc_max_children": np.max(estimated_ancestors_num_children),
        "estimated_anc_mean_children":
            np.mean(estimated_ancestors_num_children),
        "exact_anc_max_children": np.max(exact_ancestors_num_children),
        "exact_anc_mean_children":
            np.mean(exact_ancestors_num_children),
    }
    results.update(simulation_args)
    if tree_metrics:
        before = time.perf_counter()
        breakpoints, kc_distance = tsinfer.compare(smc_ts, exact_ancestors_ts)
        d = breakpoints[1:] - breakpoints[:-1]
        d /= breakpoints[-1]
        exact_anc_kc_distance_weighted = np.sum(kc_distance * d)
        exact_anc_perfect_trees = np.sum((kc_distance == 0) * d)
        exact_anc_kc_mean = np.mean(kc_distance)
        breakpoints, kc_distance = tsinfer.compare(smc_ts, estimated_ancestors_ts)
        d = breakpoints[1:] - breakpoints[:-1]
        d /= breakpoints[-1]
        estimated_anc_kc_distance_weighted = np.sum(kc_distance * d)
        estimated_anc_perfect_trees = np.sum((kc_distance == 0) * d)
        estimated_anc_kc_mean = np.mean(kc_distance)
        tree_metrics_time = time.perf_counter() - before
        results.update({
            "tree_metrics_time": tree_metrics_time,
            "exact_anc_kc_distance_weighted": exact_anc_kc_distance_weighted,
            "exact_anc_perfect_trees": exact_anc_perfect_trees,
            "exact_anc_kc_mean": exact_anc_kc_mean,
            "estimated_anc_kc_distance_weighted": estimated_anc_kc_distance_weighted,
            "estimated_anc_perfect_trees": estimated_anc_perfect_trees,
            "estimated_anc_kc_mean": estimated_anc_kc_mean,
        })
    return results


def run_edges_performance(args):
    num_lengths = 10
    MB = 10**6

    work = []
    rng = random.Random()
    if args.random_seed is not None:
        rng.seed(args.random_seed)
    for L in np.linspace(0, args.length, num_lengths + 1)[1:]:
        for _ in range(args.num_replicates):
            sim_args = {
                "sample_size": args.sample_size,
                "length": L * MB,
                "recombination_rate": args.recombination_rate,
                "mutation_rate": args.mutation_rate,
                "Ne": 10**4,
                "model": "smc_prime",
                "random_seed": rng.randint(1, 2**30)}
            work.append((sim_args, args.compute_tree_metrics))

    random.shuffle(work)
    progress = tqdm.tqdm(total=len(work), disable=not args.progress)
    results = []
    try:
        with concurrent.futures.ProcessPoolExecutor(args.num_processes) as executor:
            for result in executor.map(edges_performance_worker, work):
                results.append(result)
                progress.update()

    except KeyboardInterrupt:
        pass
    progress.close()

    df = pd.DataFrame(results)
    df.length /= MB
    dfg = df.groupby(df.length).mean()
    # print(dfg.estimated_anc_edges.describe())
    print(dfg)

    name_format = os.path.join(
        args.destination_dir,
        "ancestors_n={}_L={}_mu={}_rho={}_{{}}.png".format(
            args.sample_size, args.length, args.mutation_rate, args.recombination_rate))

    plt.plot(
        dfg.num_sites, dfg.estimated_anc_edges / dfg.source_edges,
        label="estimated ancestors")
    plt.plot(
        dfg.num_sites, dfg.exact_anc_edges / dfg.source_edges,
        label="exact ancestors")
    plt.title("n = {}, mut_rate={}, rec_rate={}, reps={}".format(
        args.sample_size, args.mutation_rate, args.recombination_rate,
        args.num_replicates))
    plt.ylabel("inferred # edges / source # edges")
    plt.xlabel("Num sites")
    plt.legend()
    plt.savefig(name_format.format("edges"))
    plt.clf()

    plt.plot(
        dfg.num_sites, dfg.estimated_anc_mean_children,
        label="estimated ancestors mean", color="blue")
    plt.plot(
        dfg.num_sites, dfg.estimated_anc_max_children,
        label="estimated ancestors max", color="blue", linestyle=":")
    plt.title("n = {}, mut_rate={}, rec_rate={}, reps={}".format(
        args.sample_size, args.mutation_rate, args.recombination_rate,
        args.num_replicates))
    plt.plot(
        dfg.num_sites, dfg.exact_anc_mean_children,
        label="exact ancestors mean", color="red")
    plt.plot(
        dfg.num_sites, dfg.exact_anc_max_children,
        label="exact ancestors max", color="red", linestyle=":")
    plt.title("n = {}, mut_rate={}, rec_rate={}, reps={}".format(
        args.sample_size, args.mutation_rate, args.recombination_rate,
        args.num_replicates))
    plt.ylabel("num_children")
    plt.xlabel("Num sites")
    plt.legend()
    plt.savefig(name_format.format("num_children"))
    plt.clf()



    if args.compute_tree_metrics:
        plt.plot(
            dfg.num_sites, dfg.estimated_anc_kc_distance_weighted,
            label="estimated ancestors")
        plt.plot(
            dfg.num_sites, dfg.exact_anc_kc_distance_weighted,
            label="exact ancestors")
        plt.title("n = {}, mut_rate={}, rec_rate={}, reps={}".format(
            args.sample_size, args.mutation_rate, args.recombination_rate,
            args.num_replicates))
        plt.ylabel("Distance weighted KC metric")
        plt.xlabel("Num sites")
        plt.legend()
        plt.savefig(name_format.format("kc_distance_weighted"))
        plt.clf()

        plt.plot(
            dfg.num_sites, dfg.estimated_anc_kc_mean,
            label="estimated ancestors")
        plt.plot(
            dfg.num_sites, dfg.exact_anc_kc_mean,
            label="exact ancestors")
        plt.title("n = {}, mut_rate={}, rec_rate={}, reps={}".format(
            args.sample_size, args.mutation_rate, args.recombination_rate,
            args.num_replicates))
        plt.ylabel("Mean KC metric")
        plt.xlabel("Num sites")
        plt.legend()
        plt.savefig(name_format.format("kc_mean"))
        plt.clf()

        plt.plot(
            dfg.num_sites, dfg.estimated_anc_perfect_trees,
            label="estimated ancestors")
        plt.plot(
            dfg.num_sites, dfg.exact_anc_perfect_trees,
            label="exact ancestors")
        plt.title("n = {}, mut_rate={}, rec_rate={}, reps={}".format(
            args.sample_size, args.mutation_rate, args.recombination_rate,
            args.num_replicates))
        plt.ylabel("Mean KC metric")
        plt.xlabel("Num sites")
        plt.legend()
        plt.savefig(name_format.format("perfect_trees"))
        plt.clf()


def get_num_sample_edges(ts):
    """
    Returns the number of edges where the child is a sample.
    """
    count = 0
    for e in ts.edges():
        node = ts.node(e.child)
        count += node.is_sample()
    return count


def effective_recombination_worker(args):
    simulation_args = args
    smc_ts = msprime.simulate(**simulation_args)
    estimated_ancestors_ts = run_infer(smc_ts, exact_ancestors=False)
    exact_ancestors_ts = run_infer(smc_ts, exact_ancestors=True)
    results = {
        "num_sites": smc_ts.num_sites,
        "source_num_trees": smc_ts.num_trees,
        "source_sample_edges": get_num_sample_edges(smc_ts),
        "estimated_anc_sample_edges": get_num_sample_edges(estimated_ancestors_ts),
        "exact_anc_sample_edges": get_num_sample_edges(exact_ancestors_ts),
    }
    results.update(simulation_args)
    return results


def run_effective_recombination(args):
    num_sample_sizes = 10
    MB = 10**6

    work = []
    rng = random.Random()
    if args.random_seed is not None:
        rng.seed(args.random_seed)
    for n in np.linspace(0, args.sample_size, num_sample_sizes + 1)[1:].astype(int):
        for _ in range(args.num_replicates):
            sim_args = {
                "sample_size": n,
                "length": args.length * MB,
                "recombination_rate": args.recombination_rate,
                "mutation_rate": args.mutation_rate,
                "Ne": 10**4,
                "model": "smc_prime",
                "random_seed": rng.randint(1, 2**30)}
            work.append(sim_args)

    random.shuffle(work)
    progress = tqdm.tqdm(total=len(work), disable=not args.progress)
    results = []
    try:
        with concurrent.futures.ProcessPoolExecutor(args.num_processes) as executor:
            for result in executor.map(effective_recombination_worker, work):
                results.append(result)
                progress.update()

    except KeyboardInterrupt:
        pass
    progress.close()

    df = pd.DataFrame(results)
    df.length /= MB
    df.source_sample_edges -= df.sample_size
    df.estimated_anc_sample_edges -= df.sample_size
    df.exact_anc_sample_edges -= df.sample_size

    dfg = df.groupby(df.sample_size).mean()
    print(dfg)

    name_format = os.path.join(
        args.destination_dir,
        "effective_recombination_n={}_L={}_mu={}_rho={}_{{}}.png".format(
            args.sample_size, args.length, args.mutation_rate, args.recombination_rate))

    plt.plot(dfg.source_sample_edges, label="source")
    plt.plot(dfg.estimated_anc_sample_edges, label="estimated ancestors")
    plt.plot(dfg.exact_anc_sample_edges, label="exact ancestors")
    plt.ylabel("Number of sample edges - n")
    plt.xlabel("Sample size")
    plt.legend()
    plt.savefig(name_format.format("sample_edges"))
    plt.clf()


def unrank(samples, n):
    """
    Unranks the specified set of samples from a possible n into its position
    in a lexicographically sorted list of bitstrings.
    """
    bitstring = np.zeros(n, dtype=int)
    for s in samples:
        bitstring[s] = 1
    mult = 2**np.arange(n, dtype=int)
    unranked = np.sum(mult * bitstring)
    return unranked


def edge_plot(ts, filename):
    n = ts.num_samples
    pallete = sns.color_palette("husl", 2**n - 1)
    lines = []
    colours = []
    for tree in ts.trees():
        left, right = tree.interval
        for u in tree.nodes():
            for c in tree.children(u):
                lines.append([(left, c), (right, c)])
                colours.append(pallete[unrank(tree.samples(c), n)])

    lc = mc.LineCollection(lines, linewidths=2, colors=colours)
    fig, ax = plt.subplots()
    ax.add_collection(lc)
    ax.autoscale()
    plt.savefig(filename)
    plt.clf()



def run_hotspot_analysis(args):
    MB = 10**6
    L = args.length * MB

    rng = random.Random()
    if args.random_seed is not None:
        rng.seed(args.random_seed)

    breakpoints = np.linspace(0, L, args.num_hotspots + 2)
    end = breakpoints[1:-1] + L * args.hotspot_width
    breakpoints = np.hstack([breakpoints, end])
    breakpoints.sort()
    rates = np.zeros_like(breakpoints)
    rates[:-1] = args.recombination_rate
    # Set the odd elements of the array to be hotspots.
    rates[1::2] *= args.hotspot_intensity
    recomb_map = msprime.RecombinationMap(list(breakpoints), list(rates))

    sim_args = {
        "sample_size": args.sample_size,
        "recombination_map": recomb_map,
        "mutation_rate": args.mutation_rate,
        "Ne": 10**4,
        "random_seed": rng.randint(1, 2**30)}
    ts = msprime.simulate(**sim_args)
    print("simulated ", ts.num_trees, "trees and", ts.num_sites, "sites")

    inferred_ts = run_infer(ts)

    num_bins = 100
    hotspot_breakpoints = breakpoints

    for density in [True, False]:
        for x in hotspot_breakpoints[1:-1]:
            plt.axvline(x=x, color="k", ls=":")
        breakpoints = np.array(list(inferred_ts.breakpoints()))
        v, bin_edges = np.histogram(breakpoints, num_bins, density=density)
        plt.plot(bin_edges[:-1], v, label="inferred")
        breakpoints = np.array(list(ts.breakpoints()))
        v, bin_edges = np.histogram(breakpoints, num_bins, density=density)
        plt.plot(bin_edges[:-1], v, label="source")
        plt.ylabel("Number of breakpoints")
        plt.legend()

        name_format = os.path.join(
            args.destination_dir,
            "hotspots_n={}_L={}_mu={}_rho={}_N={}_I={}_W={}_{{}}.png".format(
                args.sample_size, args.length, args.mutation_rate, args.recombination_rate,
                args.num_hotspots, args.hotspot_intensity, args.hotspot_width))
        plt.savefig(name_format.format("breakpoints_density={}".format(density)))
        plt.clf()

    print("Generating edge plots")
    # TODO add option for colour mapping.
    edge_plot(ts, name_format.format("source_edges"))
    edge_plot(inferred_ts, name_format.format("dest_edges"))


def error_analysis_worker(args):
    simulation_args, input_error, inference_error = args
    before = time.perf_counter()
    ts = msprime.simulate(**simulation_args)

    # V = generate_samples(ts, input_error)
    ts = tsinfer.insert_errors(ts, inference_error)
    V = ts.genotype_matrix()
    inferred_ts = tsinfer.infer(
        V, [site.position for site in ts.sites()], sample_error=0,
        sequence_length=ts.sequence_length,
        recombination_rate=simulation_args["recombination_rate"])

    results = {
        "input_error": input_error,
        "inference_error": inference_error,
        "num_sites": ts.num_sites,
        "source_num_trees": ts.num_trees,
        "inferred_num_trees": inferred_ts.num_trees,
        "source_edges": ts.num_edges,
        "inferred_edges": inferred_ts.num_edges,
    }
    results.update(simulation_args)

    breakpoints, kc_distance = tsinfer.compare(ts, inferred_ts)
    d = breakpoints[1:] - breakpoints[:-1]
    d /= breakpoints[-1]
    kc_distance_weighted = np.sum(kc_distance * d)
    kc_mean = np.mean(kc_distance)
    results.update({
        "kc_distance_weighted": kc_distance_weighted,
        "kc_mean": kc_mean,
    })
    return results


def run_error_analysis(args):
    MB = 10**6
    L = args.length * MB

    work = []
    rng = random.Random()
    if args.random_seed is not None:
        rng.seed(args.random_seed)
    inference_errors = [
        0, args.error_probability / 10, args.error_probability,
        args.error_probability * 10]
    for e in inference_errors:
        for _ in range(args.num_replicates):
            sim_args = {
                "sample_size": args.sample_size,
                "length": L,
                "recombination_rate": args.recombination_rate,
                "mutation_rate": args.mutation_rate,
                "Ne": 10**4,
                "random_seed": rng.randint(1, 2**30)}
            work.append((sim_args, args.error_probability, e))

    random.shuffle(work)
    progress = tqdm.tqdm(total=len(work), disable=not args.progress)
    results = []
    try:
        with concurrent.futures.ProcessPoolExecutor(args.num_processes) as executor:
            for result in executor.map(error_analysis_worker, work):
                results.append(result)
                progress.update()

    except KeyboardInterrupt:
        pass
    progress.close()

    df = pd.DataFrame(results)

    name_format = os.path.join(
        args.destination_dir,
        "error_n={}_L={}_mu={}_rho={}_e={}_{{}}.png".format(
            args.sample_size, args.length, args.mutation_rate,
            args.recombination_rate, args.error_probability))

    sns.boxplot(x="inference_error", y="kc_mean", data=df)
    plt.ylabel("KC distance")
    plt.xlabel("Error parameter")
    plt.savefig(name_format.format("kc"))
    plt.clf()

    # plt.plot(dfg.inferred_edges / dfg.source_edges)
    # plt.axvline(args.error_probability, ls="--")
    # plt.ylabel("inferred # edges / source # edges")
    # plt.xlabel("Error parameter")
    # plt.savefig(name_format.format("edges"))
    # plt.clf()

    # plt.plot(dfg.kc_mean)
    # plt.axvline(args.error_probability, ls="--")
    # plt.ylabel("KC distance")
    # plt.xlabel("Error parameter")
    # plt.savefig(name_format.format("kc"))
    # plt.clf()



def ancestor_properties_worker(args):
    simulation_args, compute_exact = args
    ts = msprime.simulate(**simulation_args)

    sample_data = tsinfer.SampleData.from_tree_sequence(ts)
    estimated_anc = tsinfer.generate_ancestors(sample_data)
    estimated_anc_length = estimated_anc.end[:] - estimated_anc.start[:]
    focal_sites = estimated_anc.focal_sites[:]
    estimated_anc_focal_distance = np.zeros(estimated_anc.num_ancestors)
    for j in range(estimated_anc.num_ancestors):
        focal = focal_sites[j]
        if len(focal) > 0:
            estimated_anc_focal_distance[j] = focal[-1] - focal[0]

    results = {
        "num_sites": ts.num_sites,
        "num_trees": ts.num_trees,
        "estimated_anc_num": estimated_anc.num_ancestors,
        "estimated_anc_mean_len": np.mean(estimated_anc_length),
        "estimated_anc_mean_focal_distance": np.mean(estimated_anc_focal_distance),
    }

    if compute_exact:
        exact_anc = tsinfer.AncestorData(sample_data)
        tsinfer.build_simulated_ancestors(sample_data, exact_anc, ts)
        exact_anc.finalise()
        exact_anc_length = exact_anc.end[:] - exact_anc.start[:]

        focal_sites = exact_anc.focal_sites[:]
        exact_anc_focal_distance = np.zeros(exact_anc.num_ancestors)
        for j in range(exact_anc.num_ancestors):
            focal = focal_sites[j]
            if len(focal) > 0:
                exact_anc_focal_distance[j] = focal[-1] - focal[0]
        results.update({
            "exact_anc_num": exact_anc.num_ancestors,
            "exact_anc_mean_len": np.mean(exact_anc_length),
            "exact_anc_mean_focal_distance": np.mean(exact_anc_focal_distance),
        })

    results.update(simulation_args)
    return results


def run_ancestor_properties(args):
    num_lengths = 10
    MB = 10**6

    work = []
    rng = random.Random()
    if args.random_seed is not None:
        rng.seed(args.random_seed)
    for L in np.linspace(0, args.length, num_lengths + 1)[1:]:
        for _ in range(args.num_replicates):
            sim_args = {
                "sample_size": args.sample_size,
                "length": L * MB,
                "recombination_rate": args.recombination_rate,
                "mutation_rate": args.mutation_rate,
                "Ne": 10**4,
                "model": "smc_prime",
                "random_seed": rng.randint(1, 2**30)}
            work.append((sim_args, not args.skip_exact))

    random.shuffle(work)
    progress = tqdm.tqdm(total=len(work), disable=not args.progress)
    results = []
    try:
        with concurrent.futures.ProcessPoolExecutor(args.num_processes) as executor:
            for result in executor.map(ancestor_properties_worker, work):
                results.append(result)
                progress.update()

    except KeyboardInterrupt:
        pass
    progress.close()

    df = pd.DataFrame(results)
    df.length /= MB
    dfg = df.groupby(df.length).mean()
    print(dfg)

    name_format = os.path.join(
        args.destination_dir, "anc-prop_n={}_L={}_mu={}_rho={}_{{}}.png".format(
        args.sample_size, args.length, args.mutation_rate, args.recombination_rate))

    # plt.plot(
    #     dfg.num_sites, dfg.exact_anc_num / dfg.estimated_anc_num,
    #     label="")
        # label="exact ancestors")
    plt.plot(dfg.num_sites, dfg.estimated_anc_num, label="estimated ancestors")
    if not args.skip_exact:
        plt.plot(dfg.num_sites, dfg.exact_anc_num, label="exact ancestors")
    plt.title("n = {}, mut_rate={}, rec_rate={}, reps={}".format(
        args.sample_size, args.mutation_rate, args.recombination_rate,
        args.num_replicates))
    # plt.ylabel("inferred # ancestors / exact # ancestors")
    plt.xlabel("Num sites")
    plt.legend()
    plt.savefig(name_format.format("num"))
    plt.clf()

    plt.plot(dfg.num_sites, dfg.estimated_anc_mean_len, label="estimated ancestors")
    if not args.skip_exact:
        plt.plot(dfg.num_sites, dfg.exact_anc_mean_len, label="exact ancestors")
    plt.title("n = {}, mut_rate={}, rec_rate={}, reps={}".format(
        args.sample_size, args.mutation_rate, args.recombination_rate,
        args.num_replicates))
    # plt.ylabel("inferred # ancestors / exact # ancestors")
    plt.xlabel("Num sites")
    plt.legend()
    plt.savefig(name_format.format("mean_len"))
    plt.clf()

    plt.plot(
        dfg.num_sites, dfg.estimated_anc_mean_focal_distance,
        label="estimated ancestors")
    if not args.skip_exact:
        plt.plot(
            dfg.num_sites, dfg.exact_anc_mean_focal_distance,
            label="exact ancestors")
    plt.title("n = {}, mut_rate={}, rec_rate={}, reps={}".format(
        args.sample_size, args.mutation_rate, args.recombination_rate,
        args.num_replicates))
    # plt.ylabel("inferred # ancestors / exact # ancestors")
    plt.xlabel("Num sites")
    plt.legend()
    plt.savefig(name_format.format("mean_focal_distance"))
    plt.clf()

def running_mean(x, N):
    cumsum = np.cumsum(np.insert(x, 0, 0))
    return (cumsum[N:] - cumsum[:-N]) / float(N)

def run_ancestor_comparison(args):
    MB = 10**6
    rng = random.Random(args.random_seed)
    sim_args = {
        "sample_size": args.sample_size,
        "length": args.length * MB,
        "recombination_rate": args.recombination_rate,
        "mutation_rate": args.mutation_rate,
        "Ne": 10**4,
        "model": "smc_prime",
        "random_seed": rng.randint(1, 2**30)}
    ts = msprime.simulate(**sim_args)

    # ts  = tsinfer.insert_errors(ts, args.error_probability, seed=args.random_seed)
    V = generate_samples(ts, args.error_probability)

    sample_data = tsinfer.SampleData(sequence_length=ts.sequence_length)
    for j, v in enumerate(V):
        sample_data.add_site(j, ["0", "1"], v)
    sample_data.finalise()

    estimated_anc = tsinfer.generate_ancestors(sample_data)
    estimated_anc_length = estimated_anc.end[1:] - estimated_anc.start[1:]

    exact_anc = tsinfer.AncestorData(sample_data)
    tsinfer.build_simulated_ancestors(sample_data, exact_anc, ts)
    exact_anc.finalise()
    exact_anc_length = exact_anc.end[1:] - exact_anc.start[1:]


    name_format = os.path.join(
        args.destination_dir, "anc-comp_n={}_L={}_mu={}_rho={}_err={}_{{}}".format(
        args.sample_size, args.length, args.mutation_rate, args.recombination_rate,
        args.error_probability))
    if args.store_data:
        filename = name_format.format("length.json")
        data = {
            "exact_ancestors": exact_anc_length.tolist(),
            "estimated_ancestors": estimated_anc_length.tolist()}
        with open(filename, "w") as f:
            json.dump(data, f)

    plt.hist([exact_anc_length, estimated_anc_length], label=["Exact", "Estimated"])
    plt.legend()
    plt.savefig(name_format.format("length-dist.png"))
    plt.clf()

    frequency = estimated_anc.time[:][1:]
    print(estimated_anc_length[frequency==2])
    plt.hist(estimated_anc_length[frequency==2], bins=50)
    plt.xlabel("doubleton ancestor length")
    plt.savefig(name_format.format("doubleton-length-dist.png"))
    plt.clf()

    # Because we have different numbers of ancestors, we need to rescale time
    # somehow to display them on the same axis. We just map time linearly into
    # 0,1. Possibly this is misleading.
    nbins = 100
    x = running_mean(exact_anc_length, nbins)
    # x = exact_anc_length
    plt.plot(np.linspace(0, 1, x.shape[0]), x, label="Exact")
    x = running_mean(estimated_anc_length, nbins)
    # x = estimated_anc_length
    plt.plot(np.linspace(0, 1, x.shape[0]), x, label="Estimated")
    plt.xlabel("Time (oldest to youngest)")
    plt.ylabel("Length")
    plt.legend()
    plt.savefig(name_format.format("length-time.png"))
    plt.clf()

def get_node_degree_by_depth(ts):
    """
    Returns a tuple (degree, depth) for each node in each tree in the
    specified tree sequence.
    """
    degree = []
    depth = []
    for tree in ts.trees():
        stack = [(tree.root, 0)]
        while len(stack) > 0:
            u, d = stack.pop()
            if len(tree.children(u)) > 0:
                degree.append(len(tree.children(u)))
                depth.append(d)
            for v in tree.children(u):
                stack.append((v, d + 1))
    return np.array(degree), np.array(depth)


def run_node_degree(args):
    MB = 10**6
    rng = random.Random()
    if args.random_seed is not None:
        rng.seed(args.random_seed)
    sim_args = {
        "sample_size": args.sample_size,
        "length": args.length * MB,
        "recombination_rate": args.recombination_rate,
        "mutation_rate": args.mutation_rate,
        "Ne": 10**4,
        "model": "smc_prime",
        "random_seed": rng.randint(1, 2**30)}
    smc_ts = msprime.simulate(**sim_args)

    engine= "C"
    df = pd.DataFrame()
    for path_compression in [True, False]:
        estimated_ancestors_ts = run_infer(
            smc_ts, engine=engine, exact_ancestors=False, path_compression=path_compression)
        degree, depth = get_node_degree_by_depth(estimated_ancestors_ts)
        df = df.append(pd.DataFrame({
            "degree": degree, "depth": depth, "type":"estimated",
            "path_compression": path_compression}))
        exact_ancestors_ts = run_infer(
            smc_ts, engine=engine, exact_ancestors=True, path_compression=path_compression)
        degree, depth = get_node_degree_by_depth(exact_ancestors_ts)
        df = df.append(pd.DataFrame({
            "degree": degree, "depth": depth, "type":"exact",
            "path_compression": path_compression}))

    name_format = os.path.join(
        args.destination_dir, "node-degree_n={}_L={}_mu={}_rho={}_{{}}.png".format(
        args.sample_size, args.length, args.mutation_rate, args.recombination_rate))
    print(df.describe())

    sns.factorplot(x="depth", y="degree",
            hue="path_compression", col="type",
            data=df, kind="bar")
    plt.savefig(name_format.format("path-compression"))
    plt.clf()

    sns.barplot(x="depth", y="degree", hue="type", data=df[df.path_compression])
    plt.savefig(name_format.format("length"))
    plt.clf()


def multiple_recombinations(ts):
    """
    Returns true if the specified tree sequence contains multiple recombinations.
    """
    for _, e_out, _ in ts.edge_diffs():
        if len(e_out) > 4:
            return True
    return False

def run_perfect_inference(args):
    for seed in range(1, args.num_replicates + 1):
        base_ts = msprime.simulate(
            args.sample_size, Ne=10**4, length=args.length * 10**6,
            recombination_rate=1e-8, random_seed=seed, model="smc_prime")
        print("simulated ts with n={} and {} trees; seed={}".format(
            base_ts.num_samples, base_ts.num_trees, seed))
        if multiple_recombinations(base_ts):
            print("Multiple recombinations; skipping")
            continue
        ts, inferred_ts = tsinfer.run_perfect_inference(
            base_ts, num_threads=args.num_threads,
            engine=args.engine, extended_checks=args.extended_checks,
            time_chunking=not args.no_time_chunking,
            path_compression=args.path_compression)
        print("n={} num_trees={} num_sites={}".format(
            ts.num_samples, ts.num_trees, ts.num_sites))
        assert ts.num_samples == inferred_ts.num_samples
        assert ts.num_sites == inferred_ts.num_sites
        if args.path_compression:
            _, distances = tsinfer.compare(ts, inferred_ts)
            assert np.all(distances == 0)
        else:
            assert ts.tables.edges == inferred_ts.tables.edges
            assert np.all(ts.tables.sites.position == inferred_ts.tables.sites.position)
            assert ts.tables.mutations == inferred_ts.tables.mutations
            assert np.array_equal(ts.tables.nodes.flags, inferred_ts.tables.nodes.flags)
            assert np.any(ts.tables.nodes.time != inferred_ts.tables.nodes.time)


def setup_logging(args):
    log_level = "WARN"
    if args.verbosity > 0:
        log_level = "INFO"
    if args.verbosity > 1:
        log_level = "DEBUG"
    if args.log_section is None:
        daiquiri.setup(level=log_level)
    else:
        daiquiri.setup(level="WARN")
        logger = logging.getLogger(args.log_section)
        logger.setLevel(log_level)


if __name__ == "__main__":

    top_parser = argparse.ArgumentParser(
        description="Simple inferface for running various tsinfer evaluations.")
    top_parser.add_argument(
        "-V", "--version", action='version',
        version='%(prog)s {}'.format(tsinfer.__version__))

    subparsers = top_parser.add_subparsers(dest="subcommand")
    subparsers.required = True

    parser = subparsers.add_parser(
        "perfect-inference", aliases=["pi"],
        help="Runs the perfect inference process on simulated tree sequences.")
    cli.add_logging_arguments(parser)
    parser.set_defaults(runner=run_perfect_inference)
    parser.add_argument("--engine", default="C")
    parser.add_argument("--sample-size", "-n", type=int, default=10)
    parser.add_argument(
        "--length", "-l", type=float, default=1, help="Sequence length in MB")
    parser.add_argument("--num-replicates", "-R", type=int, default=1)
    parser.add_argument("--num-threads", "-t", type=int, default=0)
    parser.add_argument(
        "--progress", "-p", action="store_true",
        help="Show a progress monitor.")
    parser.add_argument(
        "--extended-checks", "-X", action="store_true",
        help="Enable extra consistency checking (slow)")
    parser.add_argument(
        "--no-time-chunking", action="store_true",
        help="Disable time-chunking to give each ancestor a distinct time.")
    parser.add_argument(
        "--path-compression", "-c", action="store_true",
        help="Turn on path compression. Makes verification much slower.")

    parser = subparsers.add_parser(
        "edges-performance", aliases=["ep"],
        help="Runs a plot showing performance in terms of the edge ratio.")
    cli.add_logging_arguments(parser)
    parser.set_defaults(runner=run_edges_performance)
    parser.add_argument("--sample-size", "-n", type=int, default=10)
    parser.add_argument(
        "--length", "-l", type=float, default=1, help="Sequence length in MB")
    parser.add_argument(
        "--recombination-rate", "-r", type=float, default=1e-8,
        help="Recombination rate")
    parser.add_argument(
        "--mutation-rate", "-u", type=float, default=1e-8,
        help="Mutation rate")
    parser.add_argument("--num-replicates", "-R", type=int, default=10)
    parser.add_argument("--num-processes", "-P", type=int, default=None)
    parser.add_argument("--random-seed", "-s", type=int, default=None)
    parser.add_argument("--destination-dir", "-d", default="")
    parser.add_argument(
        "--compute-tree-metrics", "-T", action="store_true",
        help="Compute tree metrics")
    parser.add_argument(
        "--progress", "-p", action="store_true",
        help="Show a progress monitor.")

    parser = subparsers.add_parser(
        "effective-recombination", aliases=["er"],
        help="Shows the effective recombination rate against sample size.")
    cli.add_logging_arguments(parser)
    parser.set_defaults(runner=run_effective_recombination)
    parser.add_argument("--sample-size", "-n", type=int, default=1000)
    parser.add_argument(
        "--length", "-l", type=float, default=1, help="Sequence length in MB")
    parser.add_argument(
        "--recombination-rate", "-r", type=float, default=1e-8,
        help="Recombination rate")
    parser.add_argument(
        "--mutation-rate", "-u", type=float, default=1e-8,
        help="Mutation rate")
    parser.add_argument("--num-replicates", "-R", type=int, default=10)
    parser.add_argument("--num-processes", "-P", type=int, default=None)
    parser.add_argument("--random-seed", "-s", type=int, default=None)
    parser.add_argument("--destination-dir", "-d", default="")
    parser.add_argument(
        "--progress", "-p", action="store_true",
        help="Show a progress monitor.")

    parser = subparsers.add_parser(
        "hotspot-analysis", aliases=["ha"],
        help="Runs plots analysing the effects of recombination hotspots.")
    cli.add_logging_arguments(parser)
    parser.set_defaults(runner=run_hotspot_analysis)
    parser.add_argument("--sample-size", "-n", type=int, default=10)
    parser.add_argument(
        "--length", "-l", type=float, default=1, help="Sequence length in MB")
    parser.add_argument(
        "--recombination-rate", "-r", type=float, default=1e-8,
        help="Recombination rate")
    parser.add_argument(
        "--mutation-rate", "-u", type=float, default=1e-8,
        help="Mutation rate")
    parser.add_argument(
        "--num-hotspots", "-N", type=int, default=1,
        help="Number of hotspots")
    parser.add_argument(
        "--hotspot-intensity", "-I", type=float, default=10,
        help="Intensity of hotspots relative to background.")
    parser.add_argument(
        "--hotspot-width", "-W", type=float, default=0.01,
        help="Width of hotspots as a fraction of total genome length.")
    parser.add_argument("--num-replicates", "-R", type=int, default=10)
    parser.add_argument("--num-processes", "-P", type=int, default=None)
    parser.add_argument("--random-seed", "-s", type=int, default=None)
    parser.add_argument("--destination-dir", "-d", default="")
    parser.add_argument(
        "--progress", "-p", action="store_true",
        help="Show a progress monitor.")

    parser = subparsers.add_parser(
        "error-analysis", aliases=["ea"],
        help="Runs plots analysing the effects of inserted errors.")
    cli.add_logging_arguments(parser)
    parser.set_defaults(runner=run_error_analysis)
    parser.add_argument("--sample-size", "-n", type=int, default=10)
    parser.add_argument(
        "--length", "-l", type=float, default=1, help="Sequence length in MB")
    parser.add_argument(
        "--recombination-rate", "-r", type=float, default=1e-8,
        help="Recombination rate")
    parser.add_argument(
        "--mutation-rate", "-u", type=float, default=1e-8,
        help="Mutation rate")
    parser.add_argument(
        "--error-probability", "-e", type=float, default=0.01,
        help="Probability of errors.")
    parser.add_argument("--num-replicates", "-R", type=int, default=10)
    parser.add_argument("--num-processes", "-P", type=int, default=None)
    parser.add_argument("--random-seed", "-s", type=int, default=None)
    parser.add_argument("--destination-dir", "-d", default="")
    parser.add_argument(
        "--progress", "-p", action="store_true",
        help="Show a progress monitor.")

    parser = subparsers.add_parser(
        "ancestor-properties", aliases=["ap"],
        help="Runs plots showing the properties of estimated ancestors.")
    cli.add_logging_arguments(parser)
    parser.set_defaults(runner=run_ancestor_properties)
    parser.add_argument("--sample-size", "-n", type=int, default=10)
    parser.add_argument(
        "--length", "-l", type=float, default=1, help="Sequence length in MB")
    parser.add_argument(
        "--recombination-rate", "-r", type=float, default=1e-8,
        help="Recombination rate")
    parser.add_argument(
        "--mutation-rate", "-u", type=float, default=1e-8,
        help="Mutation rate")
    parser.add_argument("--num-replicates", "-R", type=int, default=10)
    parser.add_argument("--num-processes", "-P", type=int, default=None)
    parser.add_argument("--random-seed", "-s", type=int, default=None)
    parser.add_argument("--destination-dir", "-d", default="")
    parser.add_argument(
        "--progress", "-p", action="store_true",
        help="Show a progress monitor.")
    parser.add_argument(
        "--skip-exact", "-S", action="store_true",
        help="Skip computing the exact ancestors")

    parser = subparsers.add_parser(
        "ancestor-comparison", aliases=["ac"],
        help="Runs plots comparing the real and simulated ancestors for a single instance.")
    cli.add_logging_arguments(parser)
    parser.set_defaults(runner=run_ancestor_comparison)
    parser.add_argument("--sample-size", "-n", type=int, default=10)
    parser.add_argument(
        "--length", "-l", type=float, default=1, help="Sequence length in MB")
    parser.add_argument(
        "--recombination-rate", "-r", type=float, default=1e-8,
        help="Recombination rate")
    parser.add_argument(
        "--mutation-rate", "-u", type=float, default=1e-8,
        help="Mutation rate")
    parser.add_argument(
        "--error-probability", "-e", type=float, default=0,
        help="Error probability")
    parser.add_argument("--random-seed", "-s", type=int, default=None)
    parser.add_argument("--destination-dir", "-d", default="")
    parser.add_argument(
        "--store-data", "-S", action="store_true",
        help="Store the raw data.")

    parser = subparsers.add_parser(
        "node-degree", aliases=["nd"],
        help="Plots node degree vs depth in the tree.")
    cli.add_logging_arguments(parser)
    parser.set_defaults(runner=run_node_degree)
    parser.add_argument("--sample-size", "-n", type=int, default=10)
    parser.add_argument(
        "--length", "-l", type=float, default=1, help="Sequence length in MB")
    parser.add_argument(
        "--recombination-rate", "-r", type=float, default=1e-8,
        help="Recombination rate")
    parser.add_argument(
        "--mutation-rate", "-u", type=float, default=1e-8,
        help="Mutation rate")
    parser.add_argument("--random-seed", "-s", type=int, default=None)
    parser.add_argument("--destination-dir", "-d", default="")

    args = top_parser.parse_args()
    cli.setup_logging(args)
    args.runner(args)



    # check_basic_performance()
    # check_effect_error_param(float(sys.argv[1]))
    # plot_basic_performance()
    # check_variable_recomb()
    # check_single_tree_one_mutation_per_branch()
    # check_single_tree_two_mutations_per_branch()
    # check_single_tree_many_mutations_per_branch()
    # check_single_tree_high_mutation_rate()
    # check_many_trees_one_mutation_per_branch()


