import tempfile
import time
import pandas as pd
import pickle

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()
    config.bench_name = benchmark

    print("\n\n===== measure original program =====\n\n")
    target_current = 0.1
    currs = measure(config, pps, smu, target_current)
    # with open("output/test_1_original.pickle", "wb") as f:
    #     pickle.dump(currs, f)



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

    target_current = 0.02
    records = measure(config, pps, smu, target_current)
    mask = records["time_taken"] < 1
    n_total = len(records)
    n_passed = len(records[mask])
    print(f"{n_passed}/{n_total} passed SLA (1s)")

    # with open("output/test_1_ckpt.pickle", "wb") as f:
    #     pickle.dump(currs, f)


def measure(config, pps: PPS_E36311A, smu: SMU_B2902A, target_current):
    env = program_with_config(config, pps, smu)

    target_iterations = 10
    benchmark = config.bench_name
    pps.set_voltage(3.3, 1)
    pps.set_current(target_current, 1)
    watcher = SerialWatcher(benchmark, target_iterations)
    print("\n\n===== start measuring =====\n\n")
    # smu.start_measure()
    env.resume_board(terminate=True)
    records = watcher.run()
    # result = smu.stop_measure()
    print("\n\n===== measure finished =====\n\n")
    records = pd.DataFrame(records)
    columns = ["bench_name", "time_taken", "outputs", "is_correct"]
    print(records[columns])
    return records

    # currs = pd.DataFrame(result)[0]
    # return currs


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

    benchmarks = ["vFFT"]
    for benchmark in benchmarks:
        # measure_original(benchmark, pps, smu)
        measure_ckpt(benchmark, pps, smu)
    
    smu.power_off()

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("\n===== build and program binary =====\n")

    with tempfile.TemporaryDirectory() as build_dir:
        binary = env.build_binary(config, build_dir)
        env.clear_nvm_and_load_binary(binary, resume=False)
    
    print("\n===== program finished =====\n")
    
    return env
    

def get_default_build_config():
    config = BuildConfigM33()
    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.bench_repeat_count = 10
    return config


if __name__ == "__main__":
    main()