import pandas as pd
import altair as alt
import re
benchmarks = snakemake.input.benchmarks
replicates = snakemake.params.replicates
color = "#1b9e77"

job_count = lambda replicates: 9 * replicates + 2

benchmark_dfs = []
for benchmark, n in zip(benchmarks, replicates):
    benchmark = pd.read_csv(benchmark, sep="\t")
    benchmark["job_count"] = job_count(n)
    benchmark_dfs.append(benchmark)

data = pd.concat(benchmark_dfs)

base = alt.Chart(data).encode(
    x=alt.X('job_count', axis=alt.Axis(title="jobs", labelAngle=45, values=[20000, 40000, 60000, 80000, 100000])),
)

chart = base.mark_line(color=color).encode(
    y=alt.Y('s', axis=alt.Axis(title="seconds"))
)

chart |= base.mark_line(color=color).encode(
    y=alt.Y('max_rss', axis=alt.Axis(title="memory [MB]"))
)
chart.save(snakemake.output[0])