# !/usr/bin/env python3
# You may need to install matplotlib:
# pip install matplotlib
# Or use uv: 
# uv run --with=matplotlib scripts/shortinputplots.py

#
# This script reads the output file generated by the shortbench benchmark
# and generates plots for each function showing total time and time per unit
# against input size.
#
# Usage:
# 1. Run the shortbench benchmark to generate output.txt in the benchmarks directory:
#    mybuilddirector/benchmarks/shortbench --all > output.txt
# 2. Execute this script: python scripts/shortinputplots.py output.txt
#

import matplotlib.pyplot as plt
import sys

def parse_output_file(filename):
    functions = {}
    with open(filename, 'r') as f:
        lines = f.readlines()
    
    i = 0
    while i < len(lines):
        line = lines[i].strip()
        if line.startswith('# Benchmarking'):
            # Extract function name
            parts = line.split()
            func_name = parts[2]  # e.g., validate_ascii
            functions[func_name] = {'sizes': [], 'times': [], 'time_per_units': []}
            i += 1
            # Skip the next lines until the table header
            while i < len(lines) and not lines[i].strip().startswith('Size'):
                i += 1
            if i < len(lines):
                i += 1  # Skip the header
                i += 1  # Skip the dashes
                # Now read data lines
                while i < len(lines) and lines[i].strip() and not lines[i].strip().startswith('#'):
                    data_line = lines[i].strip()
                    if data_line:
                        parts = data_line.split()
                        if len(parts) >= 3:
                            size = int(parts[0])
                            time = float(parts[1])
                            time_per_unit = float(parts[2])
                            functions[func_name]['sizes'].append(size)
                            functions[func_name]['times'].append(time)
                            functions[func_name]['time_per_units'].append(time_per_unit)
                    i += 1
        else:
            i += 1
    return functions

def generate_plots(functions):
    for func_name, data in functions.items():
        sizes = data['sizes']
        times = data['times']
        time_per_units = data['time_per_units']
        
        fig, ax1 = plt.subplots(figsize=(10, 6))
        
        # Plot total time
        ax1.plot(sizes, times, marker='o', linestyle='-', color='blue', label='Total Time')
        ax1.set_xlabel('Size (bytes)')
        ax1.set_ylabel('Total Time (ns)', color='blue')
        ax1.tick_params(axis='y', labelcolor='blue')
        ax1.set_ylim(bottom=0)
        ax1.grid(True)
        
        # Create second y-axis for time per unit
        ax2 = ax1.twinx()
        ax2.plot(sizes, time_per_units, marker='x', linestyle='--', color='red', label='Time per Unit')
        ax2.set_ylabel('Time per Byte (ns)', color='red')
        ax2.tick_params(axis='y', labelcolor='red')
        ax2.set_ylim(bottom=0)
        
        plt.title(f'{func_name} - Total Time and Time per Unit vs Size')
        
        # Add legends
        lines1, labels1 = ax1.get_legend_handles_labels()
        lines2, labels2 = ax2.get_legend_handles_labels()
        ax1.legend(lines1 + lines2, labels1 + labels2, loc='upper left')
        
        plt.savefig(f'{func_name}.png')
        plt.close()

if __name__ == '__main__':
    if len(sys.argv) != 2:
        print("Usage: python scripts/shortinputplots.py <output_file>")
        sys.exit(1)
    filename = sys.argv[1]
    functions = parse_output_file(filename)
    generate_plots(functions)