#!/usr/bin/env python3

import os
import re
from collections import defaultdict

import click
import matplotlib

matplotlib.use("Agg")

# stupid regex matching ahead
RE_RESULT = re.compile(
    r"".join(
        [
            r"\s*",
            r"(?P<size>(?:\d|,)+)\s*",
            r"(?P<backend>\w+)\s*",
            *(
                rf"(?P<{name}>(?:\d|\.|,)+)\s*"
                for name in (
                    "calls",
                    "mean",
                    "stdev",
                    "min",
                    "p25",
                    "median",
                    "p75",
                    "max",
                    "delta",
                )
            ),
        ]
    )
)
RE_BENCHMARK = re.compile(r"benchmarks\.(?P<name>\w+)")
RE_PLATFORM = re.compile(r"Running on (?P<platform>\w+)")

BACKEND_COLORS = {
    "numpy": "C0",
    "aesara": "C1",
    "cupy": "C2",
    "jax": "C3",
    "numba": "C4",
    "pytorch": "C5",
    "tensorflow": "C6",
}


def plot_results(records, benchmark, platform, outfile, plot_delta=False):
    import matplotlib.pyplot as plt

    fig, ax = plt.subplots(1, 1, figsize=(5.5, 4), dpi=75)

    this_record = records[(benchmark, platform)]
    last_coords = {}

    for backend, backend_values in this_record.items():
        x = backend_values["size"]
        if plot_delta:
            y = backend_values["delta"]
            ylabel = "Relative speedup"
        else:
            y = backend_values["mean"]
            ylabel = "Mean runtime (s)"

        x, y = zip(*sorted(zip(x, y), key=lambda ix: ix[0]))

        plt.plot(x, y, "o--", label=backend, color=BACKEND_COLORS[backend])
        last_coords[backend] = (x[-1], y[-1])

        ax.spines["right"].set_visible(False)
        ax.spines["top"].set_visible(False)

        plt.xlabel("Problem size (# elements)")
        plt.ylabel(ylabel)

        plt.xscale("log")
        plt.yscale("log")

    plt.title(f'Benchmark "{benchmark}" on {platform.upper()}')
    fig.canvas.draw()

    # add annotations, make sure they don't overlap
    last_text_pos = 0
    for backend, (x, y) in sorted(last_coords.items(), key=lambda k: k[1][1]):
        trans = ax.transData
        _, tp = trans.transform((0, y))
        tp = max(tp, last_text_pos + 15)
        _, y = trans.inverted().transform((0, tp))

        plt.annotate(
            backend,
            (x, y),
            xytext=(10, 0),
            textcoords="offset points",
            annotation_clip=False,
            color=BACKEND_COLORS[backend],
            va="center",
        )

        last_text_pos = tp

    fig.tight_layout()
    fig.savefig(outfile)
    plt.close(fig)


def _parse_int(string):
    return int(string.replace(",", "_"))


@click.command("plot")
@click.argument("INFILE", type=click.File("r"))
@click.option(
    "-o",
    "--outdir",
    required=True,
    type=click.Path(file_okay=False, writable=True),
    help="Output directory for plots",
)
@click.option(
    "--plot-delta",
    is_flag=True,
    help="Plot relative speedup instead of absolute runtime",
)
def main(infile, outdir, plot_delta):
    """Read a benchmark report from file or stdin and plot the results

    Example:

        $ python run.py benchmarks/equation_of_state > bench.txt

        $ python plot.py bench.txt -o plots

    """
    records = {}

    for line in infile:
        bench_match = RE_BENCHMARK.match(line)
        if bench_match:
            current_benchmark = bench_match.group("name")
            continue

        platform_match = RE_PLATFORM.match(line)
        if platform_match:
            current_platform = platform_match.group("platform")
            continue

        result_match = RE_RESULT.match(line)
        if not result_match:
            continue

        result_line = result_match.groupdict()
        backend = result_line["backend"]

        key = (current_benchmark, current_platform)
        if key not in records:
            records[key] = {}

        if backend not in records[key]:
            records[key][backend] = defaultdict(list)

        record = records[key][backend]

        if _parse_int(result_line["size"]) in record["size"]:
            click.echo(
                f"Warning: duplicate entry for benchmark {current_benchmark} "
                f'on {current_platform}, backend {backend}, size {result_line["size"]} '
                "- skipping"
            )
            continue

        for rkey, rval in result_line.items():
            if rkey in ("calls", "size"):
                rval = _parse_int(rval)
            elif rkey in (
                "mean",
                "stdev",
                "min",
                "p25",
                "median",
                "p75",
                "max",
                "delta",
            ):
                rval = float(rval)

            record[rkey].append(rval)

    os.makedirs(outdir, exist_ok=True)

    for benchmark, platform in records.keys():
        outfile = os.path.join(outdir, f"bench-{benchmark}-{platform}.png")
        plot_results(records, benchmark, platform, outfile, plot_delta)
        click.echo(f"Wrote {outfile}")


if __name__ == "__main__":
    main()
