import tempfile
import time
import pickle
import pandas as pd
import os
import subprocess
import pprint
import threading
from multiprocessing.pool import ThreadPool

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_build_config(benchmark):
    config = BuildConfigM33()
    config.bench_name = benchmark
    bench_repeat_count = config.bench_repeat_count_small[benchmark]
    config.bench_repeat_count = bench_repeat_count

    config.insert_compiler_checkpoints = True
    config.enable_extension = True
    config.use_checkpoint_pass_counter = False
    config.use_checkpoint_voltage_check = True
    config.print_recovery_message = True
    config.split_loop = True
    config.enable_static_loop_pass_count = False
    config.enable_adaptive_loop_pass_count = True
    config.max_loop_ids = 15
    config.print_stats = False
    config.custom_unroll = False
    config.print_latency_overhead = True
    return config

def main():

    pps = PPS_E36311A()

    target_current_limit = 0.015
    total_iterations = 20

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

    # benchmarks = ["vFFT"]

    for benchmark in benchmarks:
        config = get_build_config(benchmark)
        env = TestEnv(WORKSPACE_ROOT, NVM_RESET_BIN, OPENOCD_SCRIPT)

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

        with tempfile.TemporaryDirectory() as build_dir:
            binary = env.build_binary(config, build_dir)
            env.clear_nvm_and_load_binary(binary, resume=False)
        
        pps.set_current(target_current_limit, 1)
        time.sleep(1)

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

        all_records = []
        for record in records:
            stats = record["stats"]
            for line in stats:
                key, val = line.split(":")
                val = int(val.strip())
                key = key.strip()
                record[key] = val
            all_records.append(record)
        
        df = pd.DataFrame(all_records)
        rename_columns = {
            "cycles in init()": "init",
            "cycles in getCSC()": "getCSC",
            "cycles in adjust()": "adjust",
        }
        df = df.rename(columns=rename_columns)
        columns = ["bench_name", "init", "getCSC", "adjust", "total cycles"]
        print(df[columns])
        save_records(benchmark, "adaptive", df)
        
    pps.output_off(1)


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



if __name__ == "__main__":
    main()
