"""Original rooted trees with the same radial profile; not experimental data."""
from pathlib import Path
import math

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np


def polar(radius, degrees):
    theta = math.radians(degrees)
    return np.array([radius * math.cos(theta), radius * math.sin(theta)])


def examples():
    common = {"root": np.zeros(2), "v1": polar(1, 0)}
    a = dict(common, v2=polar(2, -50), v3=polar(3, 60),
             a1=polar(4, -105), a2=polar(4, -5),
             b1=polar(4, 20), b2=polar(4, 100))
    ac = {"root": ["v1"], "v1": ["v2", "v3"],
          "v2": ["a1", "a2"], "v3": ["b1", "b2"]}
    b = dict(common, v2=polar(2, -50), v3=polar(3, -90),
             a1=polar(4, -120), a2=polar(4, -60),
             b1=polar(4, 0), b2=polar(4, 70))
    bc = {"root": ["v1"], "v1": ["v2", "b2"],
          "v2": ["v3", "b1"], "v3": ["a1", "a2"]}
    return [(a, ac), (b, bc)]


def edges(children):
    return [(parent, child) for parent, nodes in children.items() for child in nodes]


def intersection_count(points, children, radius):
    hits = set()
    for u, v in edges(children):
        start, delta = points[u], points[v] - points[u]
        aa, bb = float(delta @ delta), float(2 * start @ delta)
        cc = float(start @ start - radius * radius)
        discriminant = bb * bb - 4 * aa * cc
        if discriminant < -1e-9:
            continue
        for t in [(-bb - math.sqrt(max(0, discriminant))) / (2 * aa),
                  (-bb + math.sqrt(max(0, discriminant))) / (2 * aa)]:
            if -1e-9 <= t <= 1 + 1e-9:
                point = start + np.clip(t, 0, 1) * delta
                hits.add(tuple(np.round(point, 8)))
    return len(hits)


def branch_counts(children, node="root", count=0):
    offspring = children.get(node, [])
    if not offspring:
        return [count]
    count += len(offspring) >= 2
    return [value for child in offspring for value in branch_counts(children, child, count)]


def check_geometry(points, children):
    # Nonnegative initial radial derivative implies increasing radius on each edge.
    for u, v in edges(children):
        assert float(points[u] @ (points[v] - points[u])) >= -1e-9
    def orientation(a, b, c):
        p, q = b - a, c - a
        return p[0] * q[1] - p[1] * q[0]
    es = edges(children)
    for i, (u, v) in enumerate(es):
        for x, y in es[i + 1:]:
            if {u, v} & {x, y}:
                continue
            a, b, c, d = points[u], points[v], points[x], points[y]
            crossing = (orientation(a, b, c) * orientation(a, b, d) < -1e-9
                        and orientation(c, d, a) * orientation(c, d, b) < -1e-9)
            assert not crossing


def main():
    trees = examples()
    assert sorted(branch_counts(trees[0][1])) == [2, 2, 2, 2]
    assert sorted(branch_counts(trees[1][1])) == [1, 2, 3, 3]
    radii = np.r_[np.linspace(.05, 4.95, 99), [1., 2., 3., 4.]]
    for points, children in trees:
        check_geometry(points, children)
        for r in radii:
            expected = math.ceil(r - 1e-10) if r <= 4 + 1e-10 else 0
            assert intersection_count(points, children, r) == expected, r
    fig = plt.figure(figsize=(10, 7), dpi=100, facecolor="white")
    axes = [fig.add_axes([.04, .35, .43, .51]), fig.add_axes([.53, .35, .43, .51])]
    for ax, (points, children), color, title in zip(
            axes, trees, ["#276c99", "#b25731"], ["Tree A", "Tree B"]):
        for r in range(1, 5):
            ax.add_patch(plt.Circle((0, 0), r, fill=False, edgecolor="#d4dde3", linewidth=.8))
            ax.text(-r + .09, .1, str(r), color="#71818b", fontsize=9)
        for u, v in edges(children):
            p, q = points[u], points[v]
            ax.plot([p[0], q[0]], [p[1], q[1]], color=color, linewidth=2.5)
        for node in ["v1", "v2", "v3"]:
            p = points[node]
            ax.plot(*p, "o", color=color, markersize=6)
        ax.plot(0, 0, "ko", markersize=5)
        ax.text(-.42, -.38, "root", fontsize=10, color="#243641")
        ax.set(xlim=(-4.35, 4.35), ylim=(-4.35, 4.35), aspect="equal")
        ax.axis("off")
        ax.set_title(title, fontsize=13, weight="bold", color=color)
        counts = ", ".join(map(str, sorted(branch_counts(children))))
        ax.text(.5, -.055, "Branch points along root-to-leaf paths: " + counts,
                transform=ax.transAxes, ha="center", fontsize=10, color="#243641")
    ax = fig.add_axes([.14, .105, .72, .17])
    for left, right, value in [(0, 1, 1), (1, 2, 2), (2, 3, 3), (3, 4, 4), (4, 5, 0)]:
        ax.plot([left, right], [value, value], color="#243641", linewidth=2)
        ax.plot(left, value, "o", color="#243641", markerfacecolor="white", markersize=5)
        if right <= 4:
            ax.plot(right, value, "o", color="#243641", markersize=5)
    ax.set(xlim=(0, 5), ylim=(-.35, 4.6), xticks=range(6), yticks=range(5),
           xlabel="Radius from the root (chosen units)", ylabel=r"$N(r)$")
    ax.grid(color="#e0e7ec", linewidth=.7)
    ax.set_title("Same intersection count at every radius", fontsize=11)
    fig.suptitle("A radial intersection profile does not determine rooted tree topology",
                 y=.975, fontsize=15, weight="bold", color="#172330")
    fig.text(.5, .925, "All edges move outward; branch points lie at radii 1, 2 and 3; every leaf lies at radius 4.",
             ha="center", fontsize=10.5, color="#45535d")
    fig.text(.5, .018, "Original geometric counterexample, not measured anatomy. Counts are of distinct intersections.",
             ha="center", fontsize=10, color="#45535d")
    fig.savefig(Path.cwd() / (Path(__file__).stem + ".png"), dpi=100,
                facecolor="white", transparent=False)
    plt.close(fig)


if __name__ == "__main__":
    main()