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_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 = 20
    
    if config_name == "unroll":
        config.use_checkpoint_voltage_check = True
        config.custom_unroll = True

    return config


def parse_energy_trace_nrel(data_file, include_datetime=False):
    df = pd.read_csv(data_file)
    power_col = df.columns[2]
    trace_mA = df[power_col] * 1000

    target_panel_size_m2 = 0.0001 #10cm * 10cm
    target_panel_size_m2 = 0.00013 #1cm * 1cm
    trace_mA = trace_mA * target_panel_size_m2

    target_input_voltage = 3.3
    trace_mA = trace_mA / target_input_voltage

    early = trace_mA[0:200]
    inc1 = trace_mA[200:300]
    high = trace_mA[300:500]
    dec1 = trace_mA[500:600]
    late = trace_mA[600:800]

    if include_datetime:
        df["Global Horizontal [W/m^2]"] = trace_mA
        ret = df
    else:
        ret = trace_mA
    # ret = inc1

    return ret


def simulate_current(smu, wait_s, trace_mA):
    for i, curr in enumerate(trace_mA):
        # pps.set_current(curr / 1000, 1)
        smu.set_current_limit(curr/1000)
        time.sleep(wait_s)
        print(f"Set current to {curr:.3f}mA ({i+1}/{len(trace_mA)})")


def watcher_thread_impl(watcher):
    records = watcher.run()
    return records


def main():
    pps = PPS_E36311A()
    smu = SMU_B2902A(voltage=3.3)
    config = get_default_build_config()
    wait_s = 1

    data_file = "data/2015-04-23-Sacramento.csv"
    data_file = "data/2015-04-22-Sacramento.csv"
    trace_mA = parse_energy_trace_nrel(data_file)

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

    configs = ["pass_count", "adaptive", "unroll"]
    configs = ["adaptive"]

    for benchmark in benchmarks:
        for config_name in configs:
            all_records = []
            config = get_build_config(benchmark, config_name)

            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)
            smu.set_current_limit(0.1)
            smu.power_on()

            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(0, 1)
            smu.set_current_limit(0)
            time.sleep(2)

            watcher_pool = ThreadPool(processes=1)
            watcher = SerialWatcher(benchmark, 9999)
            result = watcher_pool.apply_async(watcher_thread_impl, (watcher,))
            watcher_pool.close()

            start_time = time.time()
            env.resume_board(terminate=True)

            args = (smu, wait_s, trace_mA)
            simulate_thread = threading.Thread(target=simulate_current, args=args)
            simulate_thread.start()

            simulate_thread.join()
            print("simulation thread finished")
            watcher.stop_signal = True
            watcher_pool.join()
            print(result)

            records = result.get()
            print(records)
            
            for record in records:
                record.update({
                    "wait_time": wait_s,
                    # "expr_start_time": start_time,
                })
                all_records.append(record)

            df = pd.DataFrame(all_records)
            df.attrs["start_time"] = start_time
            print(df)
            save_records(benchmark, config_name, df)
    
    # pps.output_off(1)
    smu.power_off()


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, config_name, df):
    with open(f"output/{config_name}/{bench_name}.pickle", "wb") as f:
        pickle.dump(df, f)


if __name__ == "__main__":
    main()
