import tempfile
import time
import pickle
import pandas as pd

from imc_utils.pps_e36311a import PPS_E36311A
from imc_utils.smu_b2902a import SMU_B2902A
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
import imc_utils.serial_device

# from ..common_configs import WORKSPACE_ROOT, NVM_RESET_BIN, OPENOCD_SCRIPT

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 get_config(benchmark, config_name):
    config = get_default_build_config()
    bench_repeat_count = config.bench_repeat_count_small[benchmark]
    config.bench_repeat_count = bench_repeat_count

    if config_name == "original":
        config.insert_compiler_checkpoints = False

    if config_name == "pass_count":
        config.use_checkpoint_pass_counter = True
        config.checkpoint_pass_count = 100000000

    if config_name == "unroll":
        config.custom_unroll = True
        config.use_checkpoint_voltage_check = True
    
    if config_name == "adaptive":
        config.split_loop = True
        config.enable_adaptive_loop_pass_count = True
        config.use_checkpoint_voltage_check = True
        config.max_loop_ids = 20
        config.loop_opt_debug = False

    return config


def main():
    pps = PPS_E36311A()
    smu = SMU_B2902A(voltage=3.3)
    use_smu = True
    config = get_default_build_config()

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

    configs = ["original", "pass_count", "unroll", "adaptive"]
    configs = ["adaptive"]
    total_iterations = 10
    # configs = ["unroll"]
    # total_iterations = 1

    for benchmark in benchmarks:
        for measure_config in configs:
            all_records = []

            config = get_config(benchmark, measure_config)
            config.bench_name = benchmark

            if use_smu:
                smu.set_current_limit(0.1)
                smu.power_on()
            else:
                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=True)

            iterations = total_iterations if measure_config == "adaptive" else 1
            watcher = SerialWatcher(benchmark, iterations)
            records = watcher.run()
            for record in records:
                record.update({
                    "config": measure_config,
                })
                all_records.append(record)

            df = pd.DataFrame(all_records)
            print(df)
            save_records(benchmark, measure_config, df)

    if use_smu:
        smu.power_off()
    else:
        pps.output_off(1)


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 = False
    config.use_checkpoint_voltage_check = False
    config.bench_infinite_loop = True
    config.print_recovery_message = False
    config.split_loop = False
    config.print_stats = True
    return config

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

if __name__ == "__main__":
    main()
