import json
import click
import numpy as np

import matplotlib as mpl
import seaborn as sns
mpl.use("agg")

import matplotlib.pyplot as plt  # noqa: E402

sns.set_style("ticks")

COMPONENT_COLORS = {
    "numpy": "orangered",
    "numpy-mpi": "coral",
    "jax": "dodgerblue",
    "jax-mpi": "steelblue",
    "jax-gpu": "teal",
    "fortran": "0.4",
    "fortran-mpi": "0.2",
}


@click.argument("INFILES", nargs=-1, type=click.Path(dir_okay=False, exists=True))
@click.option("--xaxis", type=click.Choice(["nproc", "size"]), required=True)
@click.option("--norm-component", default=None)
@click.command()
def plot_benchmarks(infiles, xaxis, norm_component):
    benchmarks = set()
    components = set()
    sizes = set()
    nprocs = set()

    for infile in infiles:
        with open(infile) as f:
            data = json.load(f)

        meta = data["settings"]
        benchmarks |= set(meta["only"])
        components |= set(meta["components"])
        sizes |= set(meta["sizes"])
        nprocs.add(meta["nproc"])

    if xaxis == "nproc":
        assert len(sizes) == 1
        xvals = np.array(sorted(nprocs))
    elif xaxis == "size":
        assert len(nprocs) == 1
        xvals = np.array(sorted(sizes))
    else:
        assert False

    if norm_component is not None and norm_component not in components:
        raise ValueError(f"Did not find norm component {norm_component} in data")

    component_data = {
        benchmark: {
            comp: np.full(len(xvals), np.nan)
            for comp in components
        }
        for benchmark in benchmarks
    }

    for infile in infiles:
        with open(infile) as f:
            data = json.load(f)

        for benchmark, bench_res in data["benchmarks"].items():
            for res in bench_res:
                if xaxis == "size":
                    # sizes are approximate, take the closest one
                    x_idx = np.argmin(np.abs(np.array(xvals) - res["size"]))
                else:
                    x_idx = xvals.tolist().index(data["settings"]["nproc"])

                time = float(res["per_iteration"]["mean"])
                component_data[benchmark][res["component"]][x_idx] = time

    for benchmark in benchmarks:
        fig, ax = plt.subplots(1, 1, figsize=(5.5, 4), dpi=150)

        last_coords = {}
        for component in components:
            if norm_component:
                # compute rel. speedup
                yvals = component_data[benchmark][norm_component] / component_data[benchmark][component]
            else:
                yvals = component_data[benchmark][component]

            plt.plot(xvals, yvals, ".--", color=COMPONENT_COLORS[component])

            finite_mask = np.isfinite(yvals)
            if finite_mask.any():
                last_coords[component] = (xvals[finite_mask][-1], yvals[finite_mask][-1])
            else:
                last_coords[component] = (xvals[0], 1)

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

        title_kwargs = dict(
            fontdict=dict(weight="bold", size=11),
            ha="left", x=0, y=1.05
        )
        if xaxis == "nproc":
            plt.xlabel("Number of MPI processes")
            mantissa, exponent = f"{list(sizes)[0]:.1e}".split("e")
            exponent = exponent.lstrip("+0")
            plt.title(f"Benchmark '{benchmark}' for size {mantissa} $\\times$ 10$^{{{exponent}}}$", **title_kwargs)

        elif xaxis == "size":
            nproc = list(nprocs)[0]
            plt.xlabel("Problem size (# elements)")
            plt.title(f"Benchmark '{benchmark}' on {nproc} processes", **title_kwargs)

            if norm_component:
                plt.axhline(nproc, linestyle="dashed", alpha=0.4, lw=1, color="C0")
                plt.annotate("Perfect CPU scaling", (min(xvals), nproc), xytext=(0, 2), textcoords="offset points", alpha=0.4, color="C0")

        if norm_component:
            plt.ylabel("Relative speedup")
        else:
            plt.ylabel("Time per iteration (s)")

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

        fig.canvas.draw()

        # add annotations, make sure they don"t overlap
        last_text_pos = 0
        for component, (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(
                component, (x, y), xytext=(10, 0), textcoords="offset points",
                annotation_clip=False, color=COMPONENT_COLORS[component], va="center",
                weight="bold",
            )

            last_text_pos = tp

        fig.tight_layout()
        fig.savefig(f"{benchmark}.png")
        plt.close(fig)


if __name__ == "__main__":
    plot_benchmarks()

@@ -85,16 +85,17 @@
  as sns%0A
+%0A
 mpl.use(
@@ -1422,25 +1422,16 @@
 data = %7B
-%0A        
 benchmar
@@ -1434,29 +1434,16 @@
 hmark: %7B
-%0A            
 comp: np
@@ -1467,28 +1467,16 @@
  np.nan)
-%0A           
  for com
@@ -1494,26 +1494,9 @@
 ents
-%0A        %7D%0A       
+%7D
  for
@@ -1519,21 +1519,16 @@
 nchmarks
-%0A    
 %7D%0A%0A    f
@@ -2979,29 +2979,16 @@
  = dict(
-%0A            
 fontdict
@@ -3017,28 +3017,16 @@
 ize=11),
-%0A           
  ha=%22lef
@@ -3040,25 +3040,16 @@
 , y=1.05
-%0A        
 )%0A      
@@ -3701,16 +3701,37 @@
 nnotate(
+%0A                    
 %22Perfect
@@ -3744,16 +3744,36 @@
 caling%22,
+%0A                   
  (min(xv
@@ -3789,73 +3789,171 @@
 oc),
- xytext=(0, 2), textcoords=%22offset points%22, alpha=0.4, color=%22C0%22
+%0A                    xytext=(0, 2),%0A                    textcoords=%22offset points%22,%0A                    alpha=0.4,%0A                    color=%22C0%22,%0A                
 )%0A%0A 
@@ -4566,32 +4566,80 @@
 ent,
- (x, y), xytext=(10, 0),
+%0A                (x, y),%0A                xytext=(10, 0),%0A               
  tex
@@ -4701,16 +4701,32 @@
 p=False,
+%0A               
  color=C
@@ -4752,16 +4752,32 @@
 ponent%5D,
+%0A               
  va=%22cen

apply black