#!/usr/bin/env python

"""
Benchmark different terminal color libraries.
"""

from timeit import Timer


def benchmark(stmt, n=1000, r=3):
    setup = (
        "from ansimarkup import parse;"
        "from colorama import Style as S, Fore as F;"
        "from termcolor import colored;"
        "from plumbum import colors;"
        "from pastel import colorize;"
        "from rich.markup import render"
    )
    timer = Timer(stmt, setup=setup)
    best = min(timer.repeat(r, n))

    usec = best * 1e6 / n
    return usec


def run_tests(title, tests):
    print(title)
    results = sorted((benchmark(v), k) for k, v in tests.items())
    for usec, name in results:
        print("  {:<12} {:01.4f} μs".format(name, usec))
    print()


tests_simple_1 = {
    "termcolor": 'colored("red bold", "red", attrs=["bold"])',
    "colorama": 'S.BRIGHT + F.RED + "red bold" + S.RESET_ALL',
    "ansimarkup": 'parse("<r><b>red bold</b></r>")',
    "pastel": 'colorize("<fg=red;options=bold>red bold</>")',
    "plumbum": 'colors.red | colors.bold | "red bold"',
    "rich": 'render("[bold red]red bold[/bold red]")',
}

tests_simple_2 = {
    "termcolor": 'colored("red bold", "red", attrs=["bold"]) + colored("red", "red") + colored("bold", attrs=["bold"])',
    "colorama": 'S.BRIGHT + F.RED + "red bold" + S.RESET_ALL + F.RED + "red" + S.RESET_ALL + S.BRIGHT + "bold" + S.RESET_ALL',
    "ansimarkup": 'parse("<r><b>red bold</b>red</r><b>bold</b>")',
    "pastel": 'colorize("<fg=red;options=bold>red bold</><fg=red>red</><options=bold>red</>")',
    "plumbum": '(colors.red + colors.bold | "red bold") + (colors.red | "red") + (colors.bold | "bold")',
    "rich": 'render("[red][bold]red bold[/bold]red[/red][bold]bold[/bold]")',
}

run_tests("Benchmark 1: <r><b>red bold</b></r>", tests_simple_1)
run_tests("Benchmark 2: <r><b>red bold</b>red</r><b>bold</b>", tests_simple_2)
