#!/usr/bin/env python3

import argparse
import re
import pandas as pd

def parse_cycle_deltas(file_path):
    cycle_values = []

    csrr_pattern = re.compile(r"csrr\s+\w+,\s+cycle")
    value_pattern = re.compile(r"0x[0-9a-fA-F]+$")

    with open(file_path, 'r') as f:
        lines = f.readlines()

    i = 0
    while i < len(lines):
        line = lines[i].strip()

        if csrr_pattern.search(line):
            if i + 1 < len(lines):
                next_line = lines[i+1].strip()
                match = value_pattern.search(next_line)

                if match:
                    val = int(match.group(), 16)
                    cycle_values.append(val)
                    i += 1
                else:
                    print(f"Warning: Found csrr at line {i} but couldn't find value on line {i+1}")
        i += 1

    if len(cycle_values) % 2 != 0:
        raise ValueError(
            f"Hanging CSRR detected! Found {len(cycle_values)} cycle reads. "
            "Each 'start' must have a corresponding 'end'."
        )

    data = []
    for j in range(0, len(cycle_values), 2):
        start_val = cycle_values[j]
        end_val = cycle_values[j+1]
        delta = end_val - start_val

        data.append({
            'start_cycle': start_val,
            'end_cycle': end_val,
            'delta': delta
        })

    return pd.DataFrame(data)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Extract cycle deltas from trace files and calculate speedup')
    parser.add_argument('files', nargs='+', help='Path to trace file(s). Last file is baseline if more than one.')
    args = parser.parse_args()

    if len(args.files) == 2:
        files = args.files
        baseline_idx = 1
    elif len(args.files) >= 3:
        files = args.files
        baseline_idx = len(files) - 1
    else:
        print("Error: At least 2 files required")
        exit(1)

    try:
        dfs = [parse_cycle_deltas(f) for f in files]
        baseline_df = dfs[baseline_idx]

        for i, df in enumerate(dfs):
            if len(df) != len(baseline_df):
                raise ValueError(
                    f"Mismatch in number of deltas: {files[i]} has {len(df)} deltas, "
                    f"but {files[baseline_idx]} has {len(baseline_df)} deltas. Cannot pair them."
                )

        result = pd.DataFrame({
            'start_cycle_baseline': baseline_df['start_cycle'],
            'end_cycle_baseline': baseline_df['end_cycle'],
            'delta_baseline': baseline_df['delta']
        })

        for i, df in enumerate(dfs):
            if i == baseline_idx:
                result[f'speedup_{i}'] = 1.0
            else:
                result[f'speedup_{i}'] = baseline_df['delta'] / df['delta']

        print("Cycle Delta Analysis (baseline: " + files[baseline_idx] + "):")
        print(result)

        result.to_parquet("result.parquet")

    except FileNotFoundError as e:
        print(f"Error: {e}")
    except ValueError as e:
        print(f"Error: {e}")