import tempfile
import time
import pandas as pd
import pickle
import subprocess

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

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 measure_original(benchmark, pps, smu):
    config = get_default_build_config(benchmark)

    print("\n===== A. measure original program =====\n")
    target_current = 0.1
    currs = measure(config, pps, smu, target_current, repeat=2)
    with open(f"output/test_1_{benchmark}_original.pickle", "wb") as f:
        pickle.dump(currs, f)


def measure_ckpt(benchmark, pps, smu: SMU_B2902A):
    config = get_default_build_config(benchmark)
    config.insert_compiler_checkpoints = False
    config.use_checkpoint_pass_counter = False
    config.split_loop = False
    config.enable_adaptive_loop_pass_count = False
    config.use_checkpoint_voltage_check = False
    config.max_loop_ids = 15
    config.enable_dynamic_checkpoint = True

    # print("\n===== B. measure checkpoint-instrumented program =====\n")
    target_current = 0.015
    wait_s = 3
    repeat = 200
    records = measure(config, pps, smu, target_current, repeat=repeat, wait_before_start_s=wait_s)
    records = pd.DataFrame(records)
    columns = ["bench_name", "time_taken", "outputs", "is_correct"]
    print(records[columns])
    with open(f"output/test_2_{benchmark}_ckpt.pickle", "wb") as f:
        pickle.dump(records, f)


def measure(config, pps: PPS_E36311A, smu: SMU_B2902A, target_current, repeat=1, wait_before_start_s=0):
    env = program_with_config(config, pps, smu)

    benchmark = config.bench_name
    pps.set_voltage(3.3, 1)
    pps.set_current(target_current, 1)
    watcher = SerialWatcher(benchmark, repeat)

    if wait_before_start_s > 0:
        print(f" > wait for {wait_before_start_s} secs")
        time.sleep(wait_before_start_s)
        env.resume_board(terminate=True, verbose=False)

    print(" > start measuring\n")
    if wait_before_start_s == 0:
        env.resume_board(terminate=True, verbose=False)
    records = watcher.run()
    print("\n > measure finished")
    return records


def main():
    pps = PPS_E36311A()
    smu = SMU_B2902A(voltage=0, use_external_power=True)
    smu.setup_sense()
    smu.set_current_limit(0.1)

    benchmarks = ["vFFT", "vCrc", "vAes", "vSha"]
    benchmarks = ["vSha"]

    total = len(benchmarks)
    for i, benchmark in enumerate(benchmarks):
        print(f"\n+++ ({i+1}/{total}) {benchmark} +++")
        # measure_original(benchmark, pps, smu)
        measure_ckpt(benchmark, pps, smu)


def program_with_config(config, pps, smu):
    env = TestEnv(WORKSPACE_ROOT, NVM_RESET_BIN, OPENOCD_SCRIPT)

    smu.power_on()
    pps.set_voltage(3.3, 1)
    pps.set_current(0.1, 1)
    pps.output_on(1)
    time.sleep(1)

    print(" > build and program binary")

    with tempfile.TemporaryDirectory() as build_dir:
        binary = env.build_binary(config, build_dir, redirect_file=subprocess.DEVNULL)
        env.clear_nvm_and_load_binary(binary, resume=False, verbose=False)
    
    print(" > program finished")
    
    return env


def get_bench_repeat_count(benchmark):
    d = {
        "vFFT": 10,
        "vConv2d": 15000,
        "vCrc": 1000,
        "vBasicMath": 1,
        "vSha": 40,
        "vStringSearch": 3000,
        "vAes": 1300,
        "vMatMul": 20000,
    }
    return d[benchmark]
    

def get_default_build_config(benchmark):
    config = BuildConfigM33()
    config.bench_name = benchmark
    config.insert_compiler_checkpoints = False
    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 = True
    config.split_loop = False
    config.enable_static_loop_pass_count = False
    config.enable_adaptive_loop_pass_count = False
    config.print_stats = False
    config.custom_unroll = False
    config.enable_dynamic_checkpoint = False
    config.bench_repeat_count = get_bench_repeat_count(benchmark)
    return config


if __name__ == "__main__":
    main()