#!/usr/bin/env python3

from typing import Any
import matplotlib.pyplot as plot
import numpy as np

dtype = np.dtype([
    ("implementation", "U20"),
    ("task", "U20"),
    ("num_threads", np.int32),
    ("batch_size", np.int32),
    ("throughput", np.float32),
])

def split(x: np.ndarray, c: str) -> list[tuple[Any, np.ndarray]]:
    x = np.sort(x, order=[c])
    vs, vis = np.unique(x[c], return_index=True)
    xs = np.split(x, vis[1:])
    return list(zip(vs, xs))

def only(x: np.ndarray, c: str, v: Any) -> np.ndarray:
    return x[np.nonzero(x[c] == v)]

def plot_compare(data: np.ndarray):
    data = only(data, "batch_size", 4096)
    for (task, data) in split(data, "task"):
        fig, ax = plot.subplots()
        ax.set_title(f"takeaway vs. crossbeam ({task} tasks, batch size 4096)")
        ax.set_xlabel("# threads")
        ax.set_ylabel("throughput (tasks/s/thread)")
        for (impl, data) in split(data, "implementation"):
            ax.plot(data["num_threads"], data["throughput"] / data["num_threads"], label=f"{impl}")
        ax.legend()
        fig.savefig(f"assets/compare-{task}-4096b.svg")

def plot_batch_sizes(data: np.ndarray):
    data = only(data, "task", "noop")
    for (impl, data) in split(data, "implementation"):
        fig, ax = plot.subplots()
        ax.set_title(f"{impl} across batch sizes (noop tasks)")
        ax.set_xlabel("# threads")
        ax.set_ylabel("throughput (tasks/s/thread)")
        for (batch_size, data) in split(data, "batch_size"):
            ax.plot(data["num_threads"], data["throughput"] / data["num_threads"], label=f"{batch_size}")
        ax.legend()
        fig.savefig(f"assets/batch-sizes-noop-{impl}.svg")

data = np.loadtxt("daemon.bench", dtype)

plot_compare(data)
plot_batch_sizes(data)

plot.show()
