# Generated by pandoc-plot 1.9.1
import matplotlib.pyplot as plt
import numpy as np

np.random.seed(2019)

import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import matplotlib.lines as mlines
import numpy as np
import random

random.seed(2023)

fig, ax = plt.subplots(1, 1, figsize=(8, 4))
ax.xaxis.set_visible(False)
ax.yaxis.set_visible(False)
ax.set_frame_on(False)
ax.set_ylim([-8, 0])
ax.set_xlim([-6, 6])

def discrete_colors(num):
    """Returns a list of discrete colors to plot, for example, various time-traces."""
    cmap = plt.get_cmap("inferno")
    mi, ma = 0.11, 0.75

    if num == 1:
        return [cmap(mi)]
    elif num == 2:
        return [cmap(mi), cmap(ma)]

    step = (ma - mi) / (num - 1)
    return [cmap(mi + i * step) for i in range(num)]


def point(x, y):
    radius = 0.7
    w = radius / (ax.get_xlim()[1] - ax.get_xlim()[0])
    h = radius / (ax.get_ylim()[1] - ax.get_ylim()[0])
    return mpatches.Ellipse(xy=(x, y), width=w, height=h, color="k", clip_on=False, zorder=float('inf'))


def line(x1, x2, y, **kwargs):
    return mpatches.FancyArrow(x=x1, y=y, dx=x2 - x1, dy=1, **kwargs)


def draw_trajectory(ax, tree, **kwargs):
    y0, y1, *ys = sorted(tree.keys())
    x0, *_ = tree[y0]
    x1 = random.choice(tree[y1])
    segments = [((x0, y0), (x1, y1))]
    for y in ys:
        _, prev = segments[-1]
        # For aesthetic purposes, we limit the horizontal range
        nextp = (random.choice([x for x in tree[y] if abs(x - prev[0]) < 2]), y)
        segments.append((prev, nextp))

    for (x1, y1), (x2, y2) in segments:
        ax.add_artist(mlines.Line2D(xdata=[x1, x2], ydata=[y1, y2], **kwargs))


tree = dict()
for y, n in enumerate([1, 3, 2, 6, 5, 4, 3, 2, 1]):
    tree[-y] = (np.arange(start=0, stop=n, step=1) - n / 2).tolist()


patches = list()
for y, xs in tree.items():
    patches.extend([point(x, y) for x in xs])

for ix, c in enumerate(discrete_colors(3), start=1):
    draw_trajectory(ax, tree, color=c)

# This needs to be calculated before plotting, as plotting the patches mutates
# the coordinates
minx = 1.5 * min(p.get_extents().intervalx.min() for p in patches)

for patch in patches:
    ax.add_patch(patch)

for y, label in enumerate(
    [
        "Start", 
        "First driver", 
        "Second driver", 
        "Car 1, Passenger 1/3", 
        "Car 2, Passenger 1/3", 
        "Car 1, Passenger 2/3", 
        "Car 2, Passenger 2/3", 
        "Car 1, Passenger 3/3", 
        "Car 2, Passenger 3/3",
    ]
):
    ax.text(
        x=minx,
        y=-y,
        s=label,
        transform=ax.transData,
        horizontalalignment="right",
        verticalalignment="center",
    )

plt.tight_layout()