import tempfile
import time
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

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_exec_time(curr):
    capacitance_F = 0.0047
    # capacitance_F = 0.01
    volt_high = 4.8
    volt_low = 3.7
    energy_j = 0.5 * capacitance_F * (volt_high**2 - volt_low**2)
    volt_mcu = 3.3
    avg_time = energy_j / volt_mcu / curr
    print(f"estimated running time (curr): {avg_time}")


def main():

    max_curr = 0.0214817
    get_exec_time(max_curr)
    return

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

    volt = 0
    curr = 0.1
    smu = SMU_B2902A(voltage=volt, curr_limit=curr, use_external_power=True)
    smu.power_on()

    pps = PPS_E36311A()
    pps.set_voltage(3.3, 1)
    pps.set_current(curr, 1)
    pps.output_on(1)

    env = TestEnv(WORKSPACE_ROOT, NVM_RESET_BIN, OPENOCD_SCRIPT)
    config = get_default_build_config()

    for benchmark in benchmarks:
        config.bench_name = benchmark
        smu.power_on()
        with tempfile.TemporaryDirectory() as build_dir:
            binary = env.build_binary(config, build_dir)
            env.clear_nvm_and_load_binary(binary, resume=True)

        smu.start_measure()
        time.sleep(measure_time)
        currs = smu.stop_measure()
        df = pd.DataFrame(currs)
        # print(df)
        avg_curr = float(df.mean().iloc[0])
        print(avg_curr)

        max_curr = float(df.max().iloc[0])
        print(max_curr)

        capacitance_F = 0.047
        volt_high = 4.8
        volt_low = 3.6
        energy_j = 0.5 * capacitance_F * (volt_high**2 - volt_low**2)
        volt_mcu = 3.3
        avg_time = energy_j / volt_mcu / avg_curr
        print(f"estimated running time (avg curr): {avg_time}")
        max_time = energy_j / volt_mcu / max_curr
        print(f"estimated running time (max curr): {max_time}")

    smu.power_off()
    pps.output_off(1)


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 = False
    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
    return config

if __name__ == "__main__":
    main()
