from os import makedirs
from os.path import join as pjoin
import json
import logging
from tqdm import tqdm
import numpy as np

from remend.tools.disassemble import DisassemblerX64, DisassemblerARM32, DisassemblerAArch64

# Dont show warnings
logging.getLogger("cle").setLevel(logging.ERROR)

BUILD="build"
DATASET="dataset"
#     # "linux_color_cvt": [
#         # {"name": "transfer_rgb_to_oprgb", "eqn": "pow(x, 1/2.19921875)"},
#         # {"name": "transfer_rgb_to_dcip3", "eqn": "pow(x, 1/2.6)"},
#         # {"name": "transfer_rgb_to_smpte2084", "eqn": "pow(x, 1/2.6)"},
#     # ]

with open("dataset_info.json", "r") as f:
    BINS = json.load(f)

DTYPES = ["f", "d"]
OPTS = [0, 1, 2] # 3
ARCH = [
    ("arm32", DisassemblerARM32),
    ("aarch64", DisassemblerAArch64),
    ("x64", DisassemblerX64),
]

def match_constants(asm_c, eqn_c):
    from math import pi, tan, sqrt
    cmap = {}
    for ac, acv in asm_c.items():
        acv = float(acv)
        for ec, ecv in eqn_c.items():
            ecv = eval(ecv)
            if abs(acv - ecv) < 1e-3:
                cmap[ec] = str(ac)
    return cmap

makedirs(DATASET, exist_ok=True)
for arch, Dclass in ARCH:
    asmf = open(pjoin(DATASET, arch + ".asm"), "w")
    eqnf = open(pjoin(DATASET, arch + ".eqn"), "w")
    infof = open(pjoin(DATASET, arch + ".info"), "w")
    print("Running:", arch)
    samples = [(binary, func, data, opt, dtype) 
                    for binary in BINS
                    for func, data in BINS[binary].items()
                    for opt in OPTS
                    for dtype in DTYPES
              ]
    for binary, func, data, opt, dtype in tqdm(samples):
        binpath = pjoin(BUILD, f"{binary}_{arch}_O{opt}.elf")
        fname = f"{func}_{dtype}"
        if binary == "nn_funcs":
            fname += "_"
        D = Dclass(binpath)
        diss = D.disassemble(fname)
        info = {"constants": D.constants, "bin": binpath, "func": fname,
                "eqn": data["eqn"], "eqn_constants": data["constants"]}
        prefix = data["prefix"][:]

        cmap = match_constants(D.constants, data["constants"])
        if len(cmap) != len(data["constants"]):
            info["cmap"] = cmap
        else:
            for ec, ac in cmap.items():
                prefix = prefix.replace(ec, "<PH>"+ac)
            prefix = prefix.replace("<PH>", "k")

        asmf.write(diss + "\n")
        eqnf.write(prefix + "\n")
        json.dump(info, infof)
        infof.write("\n")
    asmf.close()
    infof.close()