import subprocess
import re
from tabulate import tabulate

import os

env = os.environ.copy()
env["TILELANG_CLEAR_CACHE"] = "1"


def parse_output(output):
    data = {}
    for line in output.split('\n'):
        line = line.strip()
        if line.startswith('Latency:'):
            match = re.search(r'Latency: ([\d.]+)', line)
            data['latency'] = match.group(1) if match else 'N/A'
        elif line.startswith('TFlops:'):
            match = re.search(r'TFlops: ([\d.]+)', line)
            data['best_tflops'] = match.group(1) if match else 'N/A'
        elif line.startswith('Config:'):
            data['config'] = line.split('Config: ')[-1]
        elif line.startswith('Reference TFlops:'):
            match = re.search(r'Reference TFlops: ([\d.]+)', line)
            data['ref_tflops'] = match.group(1) if match else 'N/A'
    return data


output_v1 = subprocess.run(['./tl/bin/python', './maint/scripts/performance.py'],
                           capture_output=True,
                           text=True,
                           env=env).stdout
data_v1 = parse_output(output_v1)

output_v2 = subprocess.run(['./tll/bin/python', './maint/scripts/performance.py'],
                           capture_output=True,
                           text=True,
                           env=env).stdout
data_v2 = parse_output(output_v2)

table = [[
    "original", data_v1['latency'], data_v1['best_tflops'], data_v1['ref_tflops'], data_v1['config']
], [
    "current", data_v2['latency'], data_v2['best_tflops'], data_v2['ref_tflops'], data_v2['config']
]]

headers = ["version", "Best Latency (s)", "Best TFlops", "Reference TFlops", "Best Config"]

print(tabulate(table, headers=headers, tablefmt="github", stralign="left", numalign="decimal"))