import subprocess


class Bench:
    def __init__(self, line):
        self.name = line.split(' ')[1]
        self.speed = float(line.split('... bench:')[1].split('/iter')[0].strip().split(' ')[0].replace(',', ''))
        self.margin = float(line.split('(+/-')[1].split(')')[0].strip().replace(',', ''))


def run_bench():
    res = subprocess.run(["cargo", "+nightly", "bench"], check=True, stdout=subprocess.PIPE)
    return res.stdout.decode('utf-8')


def get_benches():
    print("> Retrieving benches...")
    out = run_bench()
    print("> Done")

    lines = out.split("\n")
    i = 0
    found_bench = False
    while i < len(lines):
        line = lines[i].strip()
        i += 1
        if line.startswith("test result:"):
            found_bench = True
            break
    if not found_bench:
        raise "Not found benchmarks in:\n{}".format(out)
    benches = {}
    while i < len(lines):
        line = lines[i].strip()
        i += 1
        if line.startswith("test result:"):
            break
        if not line.startswith("test bench_"):
            continue
        bench = Bench(line)
        benches[bench.name] = bench
    return benches


def get_good_results():
    benches = get_benches()

    for i in range(10):
        new_benches = get_benches()
        worst_rel_margin_1000 = 0
        updates = 0
        for key in benches.keys():
            if new_benches[key].speed < benches[key].speed and new_benches[key].margin < 2 * benches[key].margin:
                benches[key] = new_benches[key]
                updates += 1
            rel_margin_1000 = 1000 * benches[key].margin / benches[key].speed
            if rel_margin_1000 > worst_rel_margin_1000:
                worst_rel_margin_1000 = rel_margin_1000
        if worst_rel_margin_1000 > 5:
            print(">> Worst relative margin (iteration {}, updates: {}): {}".format(
                i, updates, worst_rel_margin_1000))
        else:
            break
    return benches


def get_str_for_bench(bench):
    return "{} ns/iter (+/- {})".format(bench.speed, bench.margin)


def show_results(benches):
    print("==== BENCHES RESULTS =====")
    for value in benches.values():
        print("{}: {}".format(value.name, get_str_for_bench(value)))


def show_comparison(current, new_ones):
    print("=== COMPARISON ===")
    print("| bench name | current | new | diff |")
    print("|-|-|-|-|")
    for key in current.keys():
        print("| {} | {} | {} | {:.1f}% |".format(
            key,
            get_str_for_bench(current[key]),
            get_str_for_bench(new_ones[key]),
            (new_ones[key].speed * 100. / current[key].speed) - 100.,
        ))


def main():
    subprocess.run(["git", "checkout", "main"], check=True)
#    subprocess.run(["git", "checkout", "inline-main"], check=True)
#    subprocess.run(["git", "checkout", "nonzero-for-cstr"], check=True)
    current = get_good_results()
    show_results(current)
    
#    subprocess.run(["git", "checkout", "inline"], check=True)
    subprocess.run(["git", "checkout", "inline-main"], check=True)
#    subprocess.run(["git", "checkout", "nonzero-for-cstr"], check=True)
    new_ones = get_good_results()
    show_results(new_ones)
    
    show_comparison(current, new_ones)


if __name__ == "__main__":
    main()