import pickle
import pandas as pd


def get_average_active_secs(df, repeat=1):
    mask = (df > 0.005) & (df < 0.020)
    active = df[mask]
    secs = len(active) / 1000
    return secs / repeat


def analyze_benchmark(benchmark):
    with open(f"output/test_1_{benchmark}_original.pickle", "rb") as f:
        d1 = pickle.load(f)

    with open(f"output/test_1_{benchmark}_ckpt.pickle", "rb") as f:
        d2 = pickle.load(f)
    
    avg_orig = get_average_active_secs(d1, 2)
    avg_ckpt = get_average_active_secs(d2, 5)
    return avg_orig, avg_ckpt


def do_analysis(benchmarks):
    results = []
    for benchmark in benchmarks:
        try:
            avg_orig, avg_ckpt = analyze_benchmark(benchmark)
        except:
            continue
        results.append({
            "benchmark": benchmark,
            "original (s)": avg_orig,
            "ckpt (s)": avg_ckpt,
            "overhead (%)": ((avg_ckpt / avg_orig) - 1) * 100,
        })
    d = pd.DataFrame(results)
    print(d)

    avg_overhead = d["overhead (%)"].mean()

    print(f"average overhead: {avg_overhead:.2f}%")


def main():
    benchmarks = ["vFFT", "vConv2d", "vCrc", "vBasicMath", "vSha", "vStringSearch", "vAes", "vMatMul"]
    benchmarks = ["vFFT", "vCrc", "vBasicMath", "vSha"]
    do_analysis(benchmarks)


if __name__ == "__main__":
    main()