"""
This module contains classes that store information about the testing flow.
"""
import json
import pathlib
import pprint
import re
import shutil
import sys

import onnxruntime as ort
from sys_flow.compiler_v2 import unpack_nefs as unpack_nefs_v1
from sys_flow.inference import get_model_io as get_model_io_v1
from sys_flow_v2.compiler_v2 import unpack_nefs as unpack_nefs_v2
from sys_flow_v2.inference import get_model_io as get_model_io_v2

import python_flow.common.constants as constants
import python_flow.common.exceptions as exceptions
import python_flow.internal.internal as internal

def clean_string(in_string):
    """Cleans input string to allow API to handle clean directories."""
    return re.sub(r'[^a-zA-Z0-9\/_-]', '', in_string)

def check_key_existence(keys, config, process, outer_json):
    """Check if all the keys in the 'keys' list is in the 'config' dictionary."""
    for key in keys:
        if key not in config:
            raise exceptions.RequiredConfigError(
                f"[emu][{process}] is missing {key} key from input JSON {outer_json}")

def flatten_emu(config):
    """Flattens the needed dictionary in the input JSON."""
    if config["emu_mode"] == "csim":
        for key in config["csim"]:
            config[key] = config["csim"][key]
    elif config["emu_mode"] == "fixed":
        for key in config["fixed"]:
            config[key] = config["fixed"][key]
    elif config["emu_mode"] == "float":
        for key in config["float"]:
            config[key] = config["float"][key]
    elif config["emu_mode"] == "dongle":
        for key in config["csim"]:
            config[key] = config["csim"][key]
        for key in config["dongle"]:
            config[key] = config["dongle"][key]

    del config["csim"]
    del config["fixed"]
    del config["float"]
    if 'dongle' in config.keys():
        del config['dongle']

class TestConfig():
    """Configurations setup from an input JSON.

    input_config is given as the data already loaded from the corresponding input JSON file.

    Attributes:
        config: Dictionary of configurations parsed from an input JSON,
    """
    def __init__(self, input_config, options, outer_json, input_image_list):
        app_absolute = pathlib.Path(options["app_folder"]).resolve()
        input_config["pre"] = PreConfig(input_config["pre"], app_absolute, outer_json).config
        input_config["emu"] = EmuConfig(input_config["emu"], app_absolute, outer_json).config
        input_config["post"] = PostConfig(input_config["post"], app_absolute, outer_json).config
        self.orig_emu = input_config["emu"]     # used for updating paths
        self.config = input_config
        self.input_image_list = input_image_list

        self.init_options(options, outer_json)
        self.config = internal.set_reordering(self.config)

    def __repr__(self):
        return ("\n---------------------------Other config--------------------------\n{}"
                "\n------------------------Preprocess config------------------------\n{}"
                "\n------------------------Inference config--------------------------\n{}"
                "\n------------------------Postprocess config------------------------\n{}").format(
                    pprint.pformat(self.config["flow"]), pprint.pformat(self.config["pre"]),
                    pprint.pformat(self.config["emu"]), pprint.pformat(self.config["post"]))

    def clear_results(self, in_beginning):
        """Clear the inference results folder.

        Arguments:
            in_beginning: Flag indicating if this function was called at the beginning of flow.
        """
        emu = self.config["emu"]
        clear_mode = self.config["flow"]["clear"]

        # 3 clear modes currently: none, all_x_result, before
        if in_beginning:
            if clear_mode == "before" or clear_mode == "all_x_result":
                if emu["emu_mode"] == "csim":
                    shutil.rmtree(self.config["flow"]["csim_output"], ignore_errors=True)
                elif emu["emu_mode"] == "dongle":
                    shutil.rmtree(self.config["flow"]["dongle_output"], ignore_errors=True)
                elif emu["emu_mode"] == "float" or emu["emu_mode"] == "fixed":
                    shutil.rmtree(self.config["flow"]["onnx_output"], ignore_errors=True)
                shutil.rmtree(self.config["flow"]["out_folder_base"], ignore_errors=True)
        else:
            if clear_mode == "all_x_result":
                if emu["emu_mode"] == "csim":
                    shutil.rmtree(emu["csim_output"], ignore_errors=True)
                elif emu["emu_mode"] == "dongle":
                    shutil.rmtree(emu["dongle_output"], ignore_errors=True)
                elif emu["emu_mode"] == "float" or emu["emu_mode"] == "fixed":
                    shutil.rmtree(emu["onnx_output"], ignore_errors=True)

                # remake out folder for just result.json for specific image folder
                # this function is called just before dumping result json in simulator
                # out_folder will be set by this point
                shutil.rmtree(self.config["flow"]["out_folder"], ignore_errors=True)
                pathlib.Path(self.config["flow"]["out_folder"]).mkdir(parents=True, exist_ok=True)

    def init_options(self, options, outer_json):
        """Save the command line options into the class dictionary."""
        app = pathlib.Path(options["app_folder"]).resolve()
        flow_config = options.copy()
        flow_config["app_folder"] = app
        flow_config["is_binary"] = options["inputs"] != "image"
        flow_config["input_json"] = outer_json
        flow_config["input_shape"] = options["input_shape"]

        try:
            relative_dir = pathlib.Path(flow_config["image_directory"]).resolve().relative_to(
                pathlib.Path("app").resolve())
        except ValueError: # pretend absolute is relative by stripping front '/', for out folder
            relative_dir = pathlib.Path(
                str(pathlib.Path(flow_config["image_directory"]).resolve()).strip("/"))

        # when using image json, update paths will set image_name as part of folder path instead
        # of image directory
        if self.input_image_list:
            flow_config["csim_output"] = "/".join([self.orig_emu["csim_output"], app.name])
            flow_config["dongle_output"] = "/".join([self.orig_emu["dongle_output"], app.name])
            flow_config["onnx_output"] = "/".join([self.orig_emu["onnx_output"], app.name])
            flow_config["out_folder_base"] = (pathlib.Path("bin/out/") / app.name).resolve()
        else:
            # general directory, specific image directory is updated in update_paths
            flow_config["csim_output"] = self.orig_emu["csim_output"] / (app.name / relative_dir)
            flow_config["dongle_output"] = self.orig_emu["dongle_output"] / (app.name / relative_dir)
            flow_config["onnx_output"] = self.orig_emu["onnx_output"] / (app.name / relative_dir)
            flow_config["out_folder_base"] = ("bin/out/" / (app.name / relative_dir)).resolve()

        self.config["flow"] = flow_config

    def update_paths(self, image_list):
        """Update the class dictionary with the current test image."""
        image_list = [pathlib.Path(image) for image in image_list] # list of absolute path strings
        self.config["flow"]["image_file"] = image_list

        # limit characters in case too long to create folder
        combined_name = "+".join([image.name for image in image_list])[:100]

        # when using image json, update path with input image path rather than image directory
        if self.input_image_list:
            try:
                relative_dir = pathlib.Path(image_list[0].parent).resolve().relative_to(
                    pathlib.Path("app").resolve())
            except ValueError: # pretend absolute is relative by stripping front '/'
                relative_dir = pathlib.Path(str(image_list[0].parent.resolve()).strip("/"))
            combined_name = relative_dir / combined_name

        clean_out = clean_string(str(self.config["flow"]["out_folder_base"] / combined_name))
        self.config["flow"]["out_folder"] = pathlib.Path(clean_out)
        self.config["flow"]["out_folder"].mkdir(parents=True, exist_ok=True)

        # update dump folders to match current test image, mkdir if necessary
        emu = self.orig_emu.copy()
        model = emu["model_type"]
        if emu["emu_mode"] == "csim":
            emu["csim_output"] = clean_string(
                str(self.config["flow"]["csim_output"] / combined_name / model))
            pathlib.Path(emu["csim_output"]).mkdir(parents=True, exist_ok=True)

            if emu["platform"] in constants.INI_PLATFORMS:
                ini_name = model + f"_{emu['platform']}.ini"
                new_ini = str(self.config["flow"]["out_folder"] / ini_name)
                new_ini_path = pathlib.Path(new_ini).resolve()
                emu["ini"] = str(new_ini_path)
        elif emu["emu_mode"] == "dongle":
            emu["dongle_output"] = clean_string(
                str(self.config["flow"]["dongle_output"] / combined_name / model))
            pathlib.Path(emu["dongle_output"]).mkdir(parents=True, exist_ok=True)
        elif emu["emu_mode"] == "float" or emu["emu_mode"] == "fixed":
            emu["onnx_output"] = clean_string(
                str(self.config["flow"]["onnx_output"] / combined_name / model))
            pathlib.Path(emu["onnx_output"]).mkdir(parents=True, exist_ok=True)

        self.config["emu"] = emu

class PreConfig():
    """Preprocess configurations from the input JSON.

    Attributes:
        config: Dictionary of preprocess configurations parsed from an input JSON
    """
    defaults = {"dump_onnx_txt": False, "pre_bypass": False, "pre_mode": "",
                "pre_type": ""}

    def __init__(self, config, app_folder, outer_json):
        self.prep(config, app_folder, outer_json)

    def __repr__(self):
        return ("-------------------------Preprocess config-------------------------\n{}".format(
            pprint.pformat(self.config)))

    def prep(self, config, app_folder, outer_json):
        """Prepares optional input config values that are needed for preprocessing."""
        self.config = {}
        self.config.update(self.defaults)
        self.config.update(config)

        internal.prep_inner(app_folder, self.config, outer_json, "pre")

class EmuConfig():
    """Simulator configurations from the input JSON.

    Attributes:
        config: Dictionary of configurations parsed from an input JSON
    """
    defaults = {"channel_first": True, "debug": False, "data_type": "float", "emu_mode": "csim",
                "ioinfo": "", "ini": constants.TEMPLATE_INI, "model_type": "fd",
                "platform": 520, "radix": 8, "reordering": [], "runner_mode": "algorithm",
                "onnx_output": str(pathlib.Path("bin/dynasty_dump").resolve()), # relative to root
                "csim_output": str(pathlib.Path("bin/csim_dump").resolve()),
                "dongle_output": str(pathlib.Path("bin/dongle_dump").resolve())}

    required_csim = ["nef_file", "input_names"]
    required_bie = ["bie_file", "input_names"]
    required_onnx_json = ["onnx_file", "json_file", "input_names"]
    required_float = ["onnx_file", "input_names"]
    required_dongle = ["nef_file", "input_names"]

    def __init__(self, config, app_folder, outer_json):
        flatten_emu(config)
        self.prep(config, app_folder, outer_json)

    def __repr__(self):
        return ("-------------------------Emulator config-------------------------\n{}".format(
            pprint.pformat(self.config)))

    def verify(self, config, app_folder, outer_json):
        """Verifies that certain emulator configurations exist and are valid."""
        files_to_check = []
        mode = config["emu_mode"]

        if mode == "csim":
            files_to_check.extend(self.required_csim)
            check_key_existence(files_to_check, config, "csim", outer_json)
            files_to_check.remove("input_names")
        elif mode == "fixed":
            if config["debug"]: # onnx/json
                files_to_check.extend(self.required_onnx_json)
            else:               # bie
                files_to_check.extend(self.required_bie)
            check_key_existence(files_to_check, config, "fixed", outer_json)
            files_to_check.remove("input_names")
        elif mode == "float":
            files_to_check.extend(self.required_float)
            check_key_existence(files_to_check, config, "float", outer_json)
            files_to_check.remove("input_names")
        elif mode == "dongle":
            files_to_check.extend(self.required_dongle)
            check_key_existence(files_to_check, config, "dongle", outer_json)
            files_to_check.remove("input_names")

        for input_file in files_to_check:
            if config[input_file] == "":
                raise exceptions.InvalidInputError(f"[emu][{mode}][{input_file}] "
                                                   f"from input JSON {outer_json} cannot be empty")
            file_name = (app_folder / config[input_file]).resolve()
            if not file_name.is_file():
                raise exceptions.InvalidInputError(
                    f"{file_name} does not exist\n\n[emu][{mode}][{input_file}] "
                    f"from input JSON {outer_json}")
            config[input_file] = str(file_name)

        # for shape checking for csim/fixed
        if config["ioinfo"] != "":
            config["ioinfo"] = str((app_folder / config["ioinfo"]).resolve())

    def setup_model(self):
        """Prepares the model using the sys_flow API."""
        input_nodes = None
        out_node_shape = None
        d_ioinfo = None
        input_shapes = None

        if self.config["platform"] in constants.PLATFORMS_MO3:
            get_model_io = get_model_io_v1
            unpack_nefs = unpack_nefs_v1
        else:
            get_model_io = get_model_io_v2
            unpack_nefs = unpack_nefs_v2

        if self.config["emu_mode"] == "csim" or self.config["emu_mode"] == "fixed":
            if self.config["ioinfo"] != "":
                with open(self.config["ioinfo"]) as ioinfo_file:
                    data = json.load(ioinfo_file)
                input_shapes = [node["onnx_shape"] for node in data["input"]]

        if self.config["emu_mode"] == "csim":
            model_maps, p_out = unpack_nefs(self.config["nef_file"], self.config["platform"])
            self.config["model_maps"] = model_maps
            self.config["p_out"] = p_out
        elif self.config["emu_mode"] == "fixed":
            if self.config["debug"]:
                input_nodes, _, out_node_shape, d_ioinfo = get_model_io(
                    pathlib.Path(self.config["onnx_file"]), self.config["platform"])
            else:
                input_nodes, _, out_node_shape, d_ioinfo = get_model_io(
                    pathlib.Path(self.config["bie_file"]), self.config["platform"])
        elif self.config["emu_mode"] == "float":
            input_nodes, _, out_node_shape, _ = get_model_io_v2(pathlib.Path(self.config["onnx_file"]))
            model = ort.InferenceSession(self.config["onnx_file"], providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
            input_shapes = [onnx_input.shape for onnx_input in model.get_inputs()]
        self.config["input_nodes"] = input_nodes
        self.config["out_node_shape"] = out_node_shape
        self.config["d_ioinfo"] = d_ioinfo
        self.config["input_shapes"] = input_shapes

    def prep(self, config, app_folder, outer_json):
        """Prepares optional input config values that are needed for simulation."""
        self.config = {}
        self.config.update(self.defaults)
        self.config.update(config)

        # for backward compatibility
        if "onnx_input" in self.config:
            self.config["input_names"] = self.config["onnx_input"]

        if self.config["emu_mode"] != "bypass": # ignore config if bypass
            try:
                self.verify(self.config, app_folder, outer_json)
            except exceptions.ConfigError as error:
                sys.exit(error)

            self.setup_model()

class PostConfig():
    """Postprocess configurations from the input JSON.

    Attributes:
        config: Dictionary of configurations parsed from an input JSON
    """
    defaults = {"post_bypass": False, "post_mode": "", "post_type": ""}

    def __init__(self, config, app_folder, outer_json):
        self.prep(config, app_folder, outer_json)

    def __repr__(self):
        return ("-------------------------Postprocess config-------------------------\n{}".format(
            pprint.pformat(self.config)))

    def prep(self, config, app_folder, outer_json):
        """Prepares optional input config values that are needed for postprocessing."""
        self.config = {}
        self.config.update(self.defaults)
        self.config.update(config)

        internal.prep_inner(app_folder, self.config, outer_json, "post")