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

from imc_utils.pps_e36311a import PPS_E36311A
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_name):
    config = get_default_build_config()
    config.bench_name = benchmark
    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 = config.pass_count_10ms[benchmark]

    if config_name == "adaptive":
        config.use_checkpoint_voltage_check = True
        config.split_loop = True
        config.enable_adaptive_loop_pass_count = True
        config.max_loop_ids = 30
    
    if config_name == "unroll":
        config.custom_unroll = True

    return config


def get_flash_size(binary):
    flash_sections = [
        ".isr_vector",
        ".text",
        ".rodata",
        ".ARM.extab",
        ".ARM",
        ".preinit_array",
        ".init_array",
        ".fini_array",
        ".data"
    ]
    total = 0
    text = 0
    inst = f"objdump -h {binary}".split()
    output = subprocess.run(inst, capture_output=True)
    lines = output.stdout.decode().split("\n")
    for line in lines:
        tokens = line.strip().split()
        if len(tokens) < 1:
            continue
        if tokens[0].isdigit():
            section, size_in_hex = tokens[1], tokens[2]
            size = int(size_in_hex, base=16)
            if section in flash_sections:
                print(section, size_in_hex, size)
                total += size
                if section == ".text":
                    text = size
    return total, text


def get_func_size(binary, bench_name):
    function_names = {
        "vBasicMath": ["vBasicMath"],
        "vCrc": ["updateCRC32", "crc32file", "crc32buf", "vCrc"],
        "vFFT": ["vFFT"],
        "vSha": [
            "sha_init",
            "sha_update",
            "byte_reverse",
            "sha_transform",
            "sha_final",
            "sha_stream",
            "sha_print",
            "vSha",
        ],
        "vStringSearch": ["vStringSearch"],
        "vMatMul": ["vMatMul"],
        "vConv2d": ["vConv2d"],
        "vAes": [
            "AES_init_ctx_iv",
            "KeyExpansion",
            "AES_CBC_encrypt_buffer",
            "XorWithIv",
            "Cipher",
            "vAes",
            "AddRoundKey",
            "SubBytes",
            "ShiftRows",
            "MixColumns",
            "xtime",
        ],
    }
    funcs = function_names[bench_name]
    inst = f"nm -S -t d {binary}".split()
    output = subprocess.run(inst, capture_output=True)
    lines = output.stdout.decode().split("\n")
    total_size = 0
    for line in lines:
        if len(line) == 0:
            continue
        tokens = line.strip().split()
        if tokens[-1] in funcs:
            size = int(tokens[1])
            total_size += size
    return total_size


def main():
    config = get_default_build_config()

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

    duplication_ratios = list(range(10, 101, 10))

    for benchmark in benchmarks:
        for duplication_ratio in duplication_ratios:
            all_records = []
            config_name = "adaptive"
            config = get_build_config(benchmark, config_name)
            config.duplication_ratio = duplication_ratio

            env = TestEnv(WORKSPACE_ROOT, NVM_RESET_BIN, OPENOCD_SCRIPT)

            with tempfile.TemporaryDirectory() as build_dir:
                binary = env.build_binary(config, build_dir)
                size, text_size = get_flash_size(binary)
                func_size = get_func_size(binary, benchmark)
            
            record = {
                "benchmark": benchmark,
                # "config": config_name,
                "duplication_ratio": duplication_ratio,
                "flash": size,
                "text": text_size,
                "func": func_size,
            }

            all_records.append(record)
            df = pd.DataFrame(all_records)
            print(df)
            save_records(benchmark, duplication_ratio, df)


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


def save_records(bench_name, duplication_ratio, df):
    if not os.path.exists("output"):
        os.makedirs("output")
    if not os.path.exists(f"output/{bench_name}"):
        os.makedirs(f"output/{bench_name}")
    with open(f"output/{bench_name}/duplication_ratio_{duplication_ratio}.pickle", "wb") as f:
        pickle.dump(df, f)


if __name__ == "__main__":
    main()

