import tempfile
import time
import pickle
import pandas as pd
import pprint

from imc_utils.pps_e36311a import PPS_E36311A
from imc_utils.build_config.cortex_m33 import BuildConfigM33
from imc_utils.build_config.test_env import TestEnv
from imc_utils.serial_watch import SerialWatcher

WORKSPACE_ROOT = "/home/ybkim/workspace/imc/imc_freertos_app_m33"
NVM_RESET_BIN = f"{WORKSPACE_ROOT}/imc/utils/nvm_reset.elf"
OPENOCD_SCRIPT = f"{WORKSPACE_ROOT}/imc_freertos_app_m33.cfg"

def main():
    pps = PPS_E36311A()
    config = get_default_build_config()

    benchmarks = [
        "vBasicMath",
        "vCrc",
        "vFFT",
        "vSha",
        "vStringSearch",
        "vMatMul",
        "vConv2d",
        "vAes",
    ]
    # benchmarks = ["vCrc"]

    target_period_ms = 10
    target_period_ms = 400 # for 10mF
    target_period_ms = 334/2 # for 4.7mF
    target_period_ms = 309/2

    output_dict = {}

    for benchmark in benchmarks:
        bench_repeat_count = config.bench_repeat_count_small[benchmark]

        config.bench_name = benchmark
        config.bench_repeat_count = bench_repeat_count

        pps.set_voltage(3.3, 1)
        pps.set_current(0.1, 1)
        pps.output_on(1)

        env = TestEnv(WORKSPACE_ROOT, NVM_RESET_BIN, OPENOCD_SCRIPT)

        with tempfile.TemporaryDirectory() as build_dir:
            binary = env.build_binary(config, build_dir)
            env.clear_nvm_and_load_binary(binary, resume=False)

        total_iterations = 1
        watcher = SerialWatcher(benchmark, total_iterations)
        env.resume_board(terminate=True)
        records = watcher.run()

        record = records[0]
        time_taken = record["time_taken"]
        stats = record["stats"]

        for stat in stats:
            if stat.startswith("checkpoint triggered"):
                checkpoint_triggered = int(stat.split(":")[1].strip())
                break

        trigger_period = time_taken / checkpoint_triggered * 1000
        pass_count = target_period_ms / trigger_period

        print(f"checkpoint triggered: {checkpoint_triggered}")
        print(f"checkpoint trigger frequency: {checkpoint_triggered/time_taken}/s")
        print(f"trigger period: {trigger_period}ms")
        print(f"pass count: {pass_count}")

        output_dict[benchmark] = int(pass_count)

    pps.output_off(1)

    pprint.pprint(output_dict)


def get_default_build_config():
    config = BuildConfigM33()
    config.bench_name = "vBasicMath"
    config.insert_compiler_checkpoints = True
    config.enable_extension = True
    config.use_checkpoint_pass_counter = True
    config.checkpoint_pass_count = 10000000
    config.use_checkpoint_voltage_check = False
    config.bench_infinite_loop = True
    config.split_loop = False
    config.enable_static_loop_pass_count = False
    config.loop_pass_count = 0
    config.enable_adaptive_loop_pass_count = False
    config.max_loop_ids = 30
    config.print_recovery_message = False
    config.print_stats = True
    return config


def save_records(bench_name, df):
    with open(f"output/{bench_name}.pickle", "wb") as f:
        pickle.dump(df, f)


if __name__ == "__main__":
    main()
