#!/usr/bin/env python3

import argparse
from cmath import nan
import json
import copy
import os
import sys
import subprocess
import shutil
import get_run_val as rval
from enum import Enum, auto

# used for HAPS or real chip
DRAM_BASE_SHIFT = 0
# DRAM_BASE_SHIFT = 0x60000000

class NoValue(Enum):
    def __repr__(self):
        return '<%s.%s>' % (self.__class__.__name__, self.name)

class CType(NoValue):
    WAGNER = "730"
    RAVEL = "1140"

"""
Versions for 730 and 1140.
"""
class VersionX30(NoValue):
    SINGLE = "single"
    MULTI = "multi"
    MODEL = "model"
    MODEL_OPT = "model_opt"
    MODEL_REL = "model_rel"
    MULTI_DBG = "multi_dbg"
    MODEL_DBG = "model_dbg"
    MODEL_DBG_EDA = "model_dbg_eda"
    MODEL_ENC = "model_enc"
    MODEL_CPU = "model_cpu"

def dbg_print_args():
    bin = os.path.basename(sys.argv[0])
    args = {" ".join(sys.argv[1:])}
    print(f"=> [{bin}] arguments: {args}")


def load_json_file(path):
    with open(path) as f:
        return json.load(f)

def to_mebibyte(hex_value):
    if hex_value is None:
        return None
    return int(hex_value, 16) / (2 ** 20)

def to_fw_addr(size):
    return f"{hex(size)}00000"

def get_version_cfg(ctype, version):
    bcfg = get_basic_cfg(ctype)

    if ctype == CType.WAGNER.value:
        config = get_730_version_cfg(bcfg, version)
    if ctype == CType.RAVEL.value:
        config = get_730_version_cfg(bcfg, version)
        config["sram_size"] = 2048
        config["weight_compress"] = True # always enabled

    return config

def get_basic_cfg(ctype):
    BASIC_CFG_730 = {
        "def_data_bitw": 8,
        "dram_base_addr": "0x80000000",
        "dedicated_outbuf": True,
        "sram_size": 1024,
        "weight_compact": False,
        "weight_compress": False,
        "dynamic_fp": True,
        "dynamic_json_path": "XXXXX.json",
        "skip_backend": False,
        "debug_mode": False,
        "optimize": {
            "cmd_size": False,
            "do_loop_for_batch": False,
            "dma": False,
            "getw": False,
            "pooling": False,
            "parallel": False,
            "tiling_for_dma": False,
        },
        "enable_ncore": True,
        "use_ch_compact_fmt": True,
        "fm_cut_not_support_as_cpu_node": True,
    }

    if ctype == CType.WAGNER.value:
        return BASIC_CFG_730
    elif ctype == CType.RAVEL.value:
        return BASIC_CFG_730
    else:
        raise ValueError(f"invalid versionctype [{ctype}]")

def get_730_version_cfg(basic, version):
    basic = update_run_val(basic)

    SINGLE_CFG = merge_dict(basic, {
        "single_layer_test": True,
        "enable_ncore": False,
        "use_ch_compact_fmt": False,
    })
    MULTI_CFG = merge_dict(basic, {
        "optimize": {
            "getw": True,
            "pooling": True,
            "parallel": True,
            "tiling_for_dma": True,
        }})
    MODEL_CFG = merge_dict(basic, {
        "optimize": {
            "dma": True,
            "getw": True,
            "pooling": True,
            "parallel": True,
            "tiling_for_dma": True,
        }
    })
    MULTI_DBG_CFG = merge_dict(MULTI_CFG, {
        "debug_mode": True,
    })
    MODEL_DBG_CFG = merge_dict(MODEL_CFG, {
        "debug_mode": True,
    })
    MODEL_DBG_EDA_CFG = merge_dict(MODEL_CFG, {"optimize": {
        "getw": False,
    }})
    MODEL_OPT_CFG = merge_dict(MODEL_CFG, {"optimize": {
        "cmd_size": True
    }})
    MODEL_REL_CFG = merge_dict(MODEL_OPT_CFG, {
    })
    MODEL_ENC_CFG = merge_dict(MODEL_CFG, {
        "weight_compress": True,
        "encryption_flag": True,
        "encryption_key": "0x1234",
        "encryption_file": "/home/rick/working/kdp720_compiler/resource/keys.bin",
    })

    if version == "single":
        return SINGLE_CFG
    elif version == "multi":
        return MULTI_CFG
    elif version == "model":
        return MODEL_CFG
    elif version == "model_opt" or version == "model_cpu":
        return MODEL_OPT_CFG
    elif version == "model_rel":
        return MODEL_REL_CFG
    elif version == "multi_dbg":
        return MULTI_DBG_CFG
    elif version == "model_dbg":
        return MODEL_DBG_CFG
    elif version == "model_dbg_eda":
        return MODEL_DBG_EDA_CFG
    elif version == "model_enc":
        return MODEL_ENC_CFG
    else:
        raise ValueError(f"invalid version [{version}]")

def update_run_val(config):
    run_val = rval.get_run_val()

    return merge_dict(
        config, {
            "run_val": run_val,
        })


def get_add_cfg(add_json):
    cfg = {}
    if add_json:
        cfg = json.loads(add_json)
    return cfg


def get_hack_cfg(hack_cfg):
    cfg = {}
    if hack_cfg:
        cfg = load_json_file(hack_cfg)
    return cfg


def get_radix_cfg(radix_cfg):
    cfg = {}

    if radix_cfg is None:
        cfg["dynamic_fp"] = False
    else:
        if not os.path.isfile(radix_cfg):
            raise FileNotFoundError(radix_cfg)
        cfg = {
            "dynamic_fp": True,
            "dynamic_json_path": os.path.abspath(radix_cfg)
        }

    return cfg

def get_shape_cfg(shape_cfg):
    cfg = {}

    if shape_cfg is None:
        return cfg

    if not os.path.isfile(shape_cfg):
        raise FileNotFoundError(shape_cfg)
    cfg = {
        "shape_info_cfg": os.path.abspath(shape_cfg)
    }

    return cfg

def delete_elements(data, elements_to_delete=[]):
    for element in elements_to_delete:
        if isinstance(element, list):
            delete_elements(data, element)
            continue
        elif isinstance(element, dict):
            for key, val in element.items():
                delete_elements(data[key], [val])
            continue

        if element not in data:
            continue
        
        if isinstance(data, list):
            data.remove(element)
        elif isinstance(data, dict):
            del data[element]

def merge_dict(base_dictionary, additional_dictionary, keys_to_delete=[]):
    if additional_dictionary is None:
        return base_dictionary

    result = copy.deepcopy(base_dictionary)

    delete_elements(result, keys_to_delete)

    # merges additional_dictionary into result
    for key in additional_dictionary:
        if key not in result:
            result[key] = additional_dictionary[key]
            continue

        if isinstance(result[key], dict) and isinstance(additional_dictionary[key], dict):
            result[key] = merge_dict(result[key], additional_dictionary[key])
            continue

        result[key] = additional_dictionary[key]
            
    return result

def env_set_cfg(cfg, ctype):
    if os.getenv("RUN_IP_EVAL") == "1":
        add_ip_eval_cfg(cfg, ctype)

    if os.getenv("WT_COMPRESS") == "1":
        cfg["weight_compress"] = True

    return cfg


def add_ip_eval_cfg(cfg, ctype):
    key = "ip_evaluator_cfg"
    if key not in cfg:
        cfg[key] = get_ip_eval_cfg(ctype)
        return True

    return False


def remove_ip_eval_cfg(cfg):
    key = "ip_evaluator_cfg"
    if key in cfg:
        del cfg[key]


def get_ip_eval_cfg(ctype):
    fpath = f"ip_eval/ip_eval_{ctype}.json"
    return get_resource_fpath(fpath)


def get_resource_fpath(rel_path):
    # /home/kai/compiler/test/gen_rtl_test/utils/gen_config.py
    script_path = os.path.realpath(__file__)
    proj_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(script_path))))
    return os.path.join(proj_dir, "resource", rel_path)


def hack_dram_addr(cfg):
    if DRAM_BASE_SHIFT <= 0:
        return cfg

    print(f"Hack dram addr config by adding \"0x{DRAM_BASE_SHIFT:x}\":", file=sys.stderr)
    for key, val in cfg.items():
        if key[-5:] == "_addr":
            addr = DRAM_BASE_SHIFT + int(val, 16)
            addr = str(hex(addr))
            print(f"  {key}: {val} => {addr}", file=sys.stderr)
    print("", file=sys.stderr)

    return cfg


def save_config(path, cfg):
    with open(path, "w") as fp:
        json.dump(cfg, fp, indent=4, sort_keys=True)
        fp.write('\n')

# use opt_compile.py to find best compiler config
def gen_best_config(ctype, opt_model, image_cut_search_args_str, cfg, cfg_path):
    def get_script_path():
        opt_compile_dir = os.getenv("OPT_COMPILE_DIR")
        if opt_compile_dir and os.path.exists(opt_compile_dir):
            opt_compile = os.path.join(opt_compile_dir, 'opt_compile.py')
        else:
            # get opt_compile.py in default path. ex: compiler/test/gen_rtl_test/utils/gen_config.py
            script_path = os.path.realpath(__file__)
            proj_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(script_path))))
            opt_compile = os.path.join(proj_dir, 'bin/opt_compile/opt_compile.py')

        assert os.path.exists(
            opt_compile
        ), f"environment variable OPT_COMPILE_DIR [{opt_compile_dir}] doesn't exist ! please export OPT_COMPILE_DIR=XXXX"
        return opt_compile

    if not opt_model:
        return 0

    print('=> [opt_compile]')
    # env [COMPILER_BIN_DIR] is used and checked in opt_compile.py
    assert os.path.exists(opt_model), f"opt_model [{opt_model}] doesn't exist"

    if ctype not in ['730', '1140']:
        print('[Not support] only support for 730|1140 !')
        return -1

    # create tmp_cfg, and add ip_evaluator_cfg key
    out_dir = os.path.dirname(os.path.realpath(cfg_path))
    tmp_cfg = copy.deepcopy(cfg)
    is_add_ip_eval = add_ip_eval_cfg(tmp_cfg, ctype)
    tmp_cfg_path = os.path.join(out_dir, 'tmp_cfg.json')
    save_config(tmp_cfg_path, tmp_cfg)

    # run opt_compile.py, and will do:
    #  - create opt_compile. directory to store debug files. (will be removed after run)
    #  - create opt_output. directory to store exported files.  (wll keep after run)
    #  - export opt.log
    #  - export best_config.json (overwrite to cfg_path)
    #  - sub-module depandency:
    #    - image_cut_search: export best_image_cut_config.json
    script = get_script_path()
    opt_log_dir = os.path.join(out_dir, 'opt_compile')
    opt_out_dir = os.path.join(out_dir, 'opt_output')

    basic_args_str = f"{script} {ctype} {opt_model} {tmp_cfg_path} {opt_log_dir} {opt_out_dir}"
    search_args_str = f'--image_cut_search_args "{image_cut_search_args_str}"'
    basic_args = basic_args_str.split()
    search_args =['--image_cut_search_args', image_cut_search_args_str]

    print(f'{basic_args_str} {search_args_str}');
    cmd_args = basic_args + search_args
    o = subprocess.run(cmd_args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    # export opt.log
    opt_log = f"{opt_log_dir}/opt.log"
    if os.path.exists(opt_log):
        shutil.copy(opt_log, out_dir, follow_symlinks=True)

    if o.returncode != 0:
        print('\n[Failed] failed to run opt_compile:')
        if o.stdout:
            print(o.stdout.decode())

        if o.stderr:
            print(o.stderr.decode())

        shutil.rmtree(opt_log_dir)
        sys.exit(o.returncode)

    best_cfg_path = os.path.join(opt_log_dir, 'best_config.json')
    assert os.path.exists(best_cfg_path), f"best_config.json doesn't exist"
    best_cfg = load_json_file(best_cfg_path)
    #export best_image_cut_config.json for image_cut_search
    if 'fmap_cut' in best_cfg and best_cfg['fmap_cut']['mode'] == 'manual':
        org_fmap_cut_cfg = best_cfg['fmap_cut']['cfg']
        new_fmap_cut_cfg = os.path.join(out_dir, 'best_image_cut_config.json')
        shutil.copyfile(org_fmap_cut_cfg, new_fmap_cut_cfg)
        best_cfg['fmap_cut']['cfg'] = new_fmap_cut_cfg
        if is_add_ip_eval:
            remove_ip_eval_cfg(best_cfg)
    # export best_config.json (overwrite to cfg_path)
    save_config(cfg_path, best_cfg)

    shutil.rmtree(opt_log_dir)
    return 0

def gen_all():
    versions = {
        CType.WAGNER : VersionX30,
    }
    for ctype in versions:
        output_directory = f"./{ctype.value}"
        if not os.path.exists(output_directory):
            os.makedirs(output_directory, exist_ok=True)
        for version in versions[ctype]:
            config = get_version_cfg(ctype.value, version.value)
            file_name = f"config_{ctype.value}_{version.value}.json"
            config_path = f"{output_directory}/{file_name}"
            save_config(config_path, config)

################
##    MAIN    ##
################

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="generate config.json for compiler")
    ctypes = "|".join([ctype.value for ctype in CType])
    parser.add_argument("-t", "--ctype", help=f"compiler type [{ctypes}]")
    parser.add_argument(
        "-v",
        "--version",
        help="config version: [single|multi|model|model_rel|model_opt|model_dbg|model_dbg_eda]")
    parser.add_argument("-f", "--tmpl_cfg", help="path to template config (if exists, version is skipped)")
    parser.add_argument("-o", "--output", default="config.json", help="path to output config")
    parser.add_argument("-r", "--radix_cfg", help="path to radix config")
    parser.add_argument("-s", "--shape_cfg", help="path to shape info config")
    parser.add_argument("-k", "--hack_cfg", help="path to hack config")
    parser.add_argument("-a",
                        "--add_json",
                        action='append',
                        help="additional config in json str format, e.g. '{\"single_layer_test\":true,\"hstride\":true}'")
    parser.add_argument(
        "-m",
        "--opt_model",
        help=
        "model path. run opt_compile to find best compiler config. please export COMPILER_BIN_DIR=XXXX to specify compiler bin dir, and export OPT_COMPILE_DIR=XXXX to specify opt compile bin dir."
    )
    parser.add_argument("--image_cut_search_args", type=str, default=" -r -u -t -s",
                        help="arguments for image cut search. only valid when enable [-m|--opt_model]. e.g: this value can be [-r|-u|-t|-s|-pgs]. defult is ' -r -u -t -s'. hint: must add space before each argument.\n"
                        "see bin/opt_compile/modules/image_cut_search/all_rule_auto_release.py -h for more detail.")
    parser.add_argument("-g", "--gen-all-version", action="store_true", help="Generate basic configs for all version with all ctype to verify or to debug.")

    # dbg_print_args()

    args = parser.parse_args()
    if args.gen_all_version:
        gen_all()
        print("Generation done.")
        sys.exit()

    ctype = args.ctype

    tmpl_cfg = args.tmpl_cfg
    if tmpl_cfg:
        cfg = load_json_file(tmpl_cfg)
    else:
        cfg = get_version_cfg(ctype, args.version)

    hack_cfg = get_hack_cfg(args.hack_cfg)
    cfg = merge_dict(cfg, hack_cfg)

    radix_cfg = get_radix_cfg(args.radix_cfg)
    cfg = merge_dict(cfg, radix_cfg)

    shape_cfg = get_shape_cfg(args.shape_cfg)
    cfg = merge_dict(cfg, shape_cfg)

    if args.add_json:
        for js in args.add_json:
            add_cfg = get_add_cfg(js)
            cfg = merge_dict(cfg, add_cfg)

    cfg = env_set_cfg(cfg, ctype)
    cfg = hack_dram_addr(cfg)

    # save to file
    cfg_path = args.output
    save_config(cfg_path, cfg)

    # find best compiler config and replace the old one
    if 0 != gen_best_config(ctype, args.opt_model, args.image_cut_search_args, cfg, cfg_path):
        sys.exit(-1)