import tsinfer
import tskit
from tskit import MISSING_DATA

with tsinfer.SampleData(sequence_length=100) as sd:
    sd.add_site(51, [1, 0, 1, 1, 1], ["A", "T"], ancestral_allele=1)  # Anc allele is "T"
    sd.add_site(52, [0, 0, 0, 1, 1], ["G", "C"], ancestral_allele=0)
    sd.add_site(53, [0, 1, 1, 1, 0], ["C", "A"], ancestral_allele=0)
    sd.add_site(54, [0, 1, 1, 1, 1], ["G", "C"], ancestral_allele=MISSING_DATA)
    sd.add_site(55, [0, 1, 1, 1, 0], ["A", "C"], ancestral_allele=0)
    sd.add_site(56, [0, 1, 2, 0, 0], ["T", "G", "C"], ancestral_allele=0)

ts = tsinfer.infer(sd)
print("Ancestral state at site at position 51 is", ts.site(position=51).ancestral_state)
print("Ancestral state at site at position 54 inferred to be", ts.site(position=54).ancestral_state)

ts.draw_svg(
    size=(700, 300),
    y_axis=True,
    mutation_labels={m.id: f"{s.id}: {s.ancestral_state}→{m.derived_state}" for s in ts.sites() for m in s.mutations},
)

ts.draw_svg(size=(700, 300), x_lim=(0, ts.sequence_length), y_axis=True)

print(ts.num_trees, "trees in the tree sequence")  # 4, not 2, because the empty regions count as trees!

for tree in ts.trees():
    if tree.num_roots != ts.num_samples:
        # this is not an empty tree - mrca calculations will fail on an empty tree
        print(f"MRCA of 3 and 4 in the {tree.interval} is at time {tree.tmrca(3,4)}")

print(f"Time units in the inferred tree sequence are set to '{ts.time_units}'")
print("Sitewise diversity =", ts.diversity())  # This is OK as it isn't calculated form branch lengths
print(ts.diversity(mode="branch"))  # This now fails

import msprime
import tsinfer
import matplotlib.pyplot as plt
import numpy as np

rho = 1e-7
mu = 10 * rho
Ne = 5e3
filename = "simulated_data.samples"
ts = msprime.sim_ancestry(
    10,
    sequence_length=10000,
    population_size=Ne,
    recombination_rate=rho,
    random_seed=333
)
ts = msprime.sim_mutations(ts, rate=mu, random_seed=333)
print("Simulated", ts.num_trees, "trees", ts.num_sites, "sites")
sample_data = tsinfer.SampleData.from_tree_sequence(ts, path=filename)
print(f"Saved genetic variation to '{filename}'")

plot_region = (2e3, 3e3) 
inferred_simplified_ts = tsinfer.infer(sample_data, simplify=True)
print("Default output in tsinfer version 0.2.3")
inferred_simplified_ts.simplify(range(10), keep_unary=True).draw_svg(
    size=(1000, 500),
    x_lim=plot_region,
    node_labels={85: "...", 89: 89, 90: 90, 91: "ultimate anc.", 92: "X"},
    mutation_labels={},
    time_scale="rank",
    style=".n91 > .lab, .n91 > .sym {fill: green}"
)

inferred_postprocessed_ts = tsinfer.infer(sample_data)  # post-process happens by default  
print("New output in tsinfer version 0.3:")
print(" ancestor X removed, and ultimate ancestor split into multiple nodes")
matplotlib_colours = {
    "tab:blue": "#1f77b4", "tab:orange": "#ff7f0e", "tab:green": "#2ca02c", "tab:red": "#d62728"
}
mpl_vals = list(matplotlib_colours.values())
inferred_postprocessed_ts.simplify(range(10), keep_unary=True).draw_svg(
    size=(1000, 500),
    x_lim=(2e3, 3e3),
    time_scale="rank",
    node_labels={85: "...", 89: 89, 90: 90, 91: 91, 92: 92, 93: 93, 94: 94},
    mutation_labels={},
    style="".join([
        f".n{u} > .lab, .n{u} > .sym }"
        for i, u in enumerate(range(91, 95))
    ]),
)

import tsdate
mrcas_orig = [[t.interval.right, t.time(t.root)] for t in ts.trees()]

new_inferred_ts = inferred_postprocessed_ts.simplify()  # remove unary nodes for tsdate
dated_postprocessed_ts = tsdate.date(inferred_ts, mutation_rate=mu, Ne=Ne)
mrcas_new_method = [
    [t.interval.right, t.time(t.root)]
    for t in dated_postprocessed_ts.trees()
    if t.num_roots != dated_postprocessed_ts.num_samples
]

plt.step(*np.array(mrcas_orig).T, c="black", label="true tMRCA")
plt.step(*np.array(mrcas_new_method).T, c="tab:green", label="tsinferred MRCA")
plt.legend()
plt.yscale("log")

old_inferred_simplified_ts = inferred_simplified_ts.simplify()  # remove unary nodes
dated_simplified_ts = tsdate.date(old_inferred_simplified_ts, mutation_rate=mu, Ne=Ne)
mrcas_old_method = [
    [t.interval.right, t.time(t.root)]
    for t in dated_simplified_ts.trees()
    if t.num_roots != dated_simplified_ts.num_samples
]

mrcas_old = [
    [t.interval.right, t.time(t.root)]
    for t in dated_simplified_ts.trees()
    if t.num_roots != dated_simplified_ts.num_samples
]

plt.step(*np.array(mrcas_orig).T, c="black", label="true tMRCA")
plt.step(*np.array(mrcas_old).T, c="tab:orange", label="old tsinferred MRCA")
plt.legend()
plt.yscale("log")

! python3 -m tsinfer generate-ancestors simulated_data.samples

! python3 -m tsinfer infer --recombination_rate=0.4 --mismatch_ratio=0.1 simulated_data.samples -O inferred.ts
# Or python3 -m tsinfer infer --recombination_map=/path/to/hapmap --mismatch_ratio=0.1 simulated_data.samples

import json
inferred_ts = tskit.load("inferred.ts")
print(json.loads(inferred_ts.provenance(-1).record)["parameters"])

! python3 -m tsinfer match-ancestors simulated_data.samples
! python3 -m tsinfer match-samples --no-post-process simulated_data.samples -O simulated_no_pp.trees

no_pp_ts = tskit.load("simulated_no_pp.trees")
pp_ts = tsinfer.post_process(no_pp_ts, split_ultimate=False, erase_flanks=False)

ancestors = tsinfer.generate_ancestors(sd, exclude_positions=[s.position for s in sd.sites()])
anc_ts = tsinfer.match_ancestors(sd, ancestors)

shorter_sample_data = sample_data.subset(
    sites=range(100),
    sequence_length=sample_data.sites_position[:][101])
print("Original data covers ", sample_data.sequence_length, "bp")
print("Subset covers ", shorter_sample_data.sequence_length, "bp")