"""
Processes images with the testing model using CSIM, Dynasty or Dongle.
"""
import sys
from typing import List, Optional

import numpy.typing as npt

from python_flow.common import config as conf
from python_flow.common import exceptions
from python_flow.utils import csim
from python_flow.utils import dynasty
from python_flow.utils import dongle
from python_flow.utils import utils

def emulator_csim(config: conf.TestConfig,
                  pre_results: List[npt.ArrayLike],
                  inp_trans: Optional[List[int]] = [0, 3, 1, 2],
                  oup_trans: Optional[List[int]] = [0, 2, 3, 1]) -> List[npt.ArrayLike]:
    """Uses hardware CSIM simulator to run the preprocessed image through the model.

    Arguments:
        config: TestConfig class that holds the user's input config and options.
        pre_results: List of preprocessed NumPy arrays in channel last format.
        inp_trans: List of integers indicating the axes order to transpose the inputs.
        oup_trans: List of integers indicating the axes order to transpose the outputs.

    Returns:
        A list of NumPy arrays representing the outputs of the model.
    """
    data_type = config["emu"]["data_type"]
    dump_folder = config["flow"]["out_folder"]
    first = config["emu"]["channel_first"]
    input_names = config["emu"]["input_names"]
    input_shapes = config["emu"]["input_shapes"]
    output_folder = config["emu"]["csim_output"]
    model_maps = config["emu"]["model_maps"]
    nef = config["emu"]["nef_file"]
    platform = config["emu"]["platform"]
    reordering = config["emu"]["reordering"]

    dump_file_name = "_".join([config["emu"]["model_type"], "inputs"])
    dump_file = ".".join([dump_file_name, "pkl", "xz"])
    input_dump = dump_folder / dump_file

    try:
        inputs = utils.prep_inputs(pre_results, input_names, False,
                                   input_dump, inp_trans, input_shapes)
    except exceptions.InvalidInputError as e:
        sys.exit(f"model {config['emu']['model_type']}: {e}")
    return csim.csim_inference(nef, inputs, reordering, first, model_maps, platform,
                               None, data_type, output_folder, oup_trans)

def emulator_dongle(config: conf.TestConfig,
                    pre_results: List[npt.ArrayLike],
                    inp_trans: Optional[List[int]] = [0, 3, 1, 2],
                    oup_trans: Optional[List[int]] = [0, 2, 3, 1]) -> List[npt.ArrayLike]:
    """Uses Dongle to run the preprocessed image through the model.

    Input pre_results are in format (1, h, w, c).

    Arguments:
        config: TestConfig class that holds the user's input config and options.
        pre_results: List of preprocessed NumPy arrays in channel last format.
        inp_trans: List of integers indicating the axes order to transpose the inputs, unused.
        oup_trans: List of integers indicating the axes order to transpose the outputs, unused.

    Returns:
        A list of NumPy arrays representing the outputs in the model.
    """
    import python_flow.dongle.dongle_config as dongle_config
    import python_flow.utils.dongle as dongle_utils

    input_names = config["emu"]["input_names"]
    first = config["emu"]["channel_first"]
    model_id = config["emu"]["model_id"]
    output_folder = config["emu"]["dongle_output"]
    reordering = config["emu"]["reordering"]

    # convert preprocessed numpy arrays to binary files
    input_files = dongle_utils.prep_bin_input(pre_results, input_names, dongle_config.ioinfo[str(model_id)]['ioinfo_in'])
    # run dongle inference
    response, radices, scales = dongle.run_dongle(output_folder, input_files, model_id)

    # convert dongle response to numpy arrays
    ioinfo_out = dongle_config.ioinfo[str(model_id)]['ioinfo_out']
    fixed_mode = dongle_config.dongle_fixed_mode
    return dongle_utils.dongle_to_np(response, reordering, ioinfo_out, first, radices, scales,
                               fixed_mode, model_id, output_folder)

def emulator_float(config: conf.TestConfig,
                   pre_results: List[npt.ArrayLike],
                   inp_trans: Optional[List[int]] = [0, 3, 1, 2],
                   oup_trans: Optional[List[int]] = [0, 2, 3, 1]) -> List[npt.ArrayLike]:
    """Uses float Dynasty to run the preprocessed image through the model.

    Arguments:
        config: TestConfig class that holds the user's input config and options.
        pre_results: List of preprocessed NumPy arrays in channel last format.
        inp_trans: List of integers indicating the axes order to transpose the inputs.
        oup_trans: List of integers indicating the axes order to transpose the outputs.

    Returns:
        A list of NumPy arrays representing the outputs of the model.
    """
    dump_folder = config["flow"]["out_folder"]
    first = config["emu"]["channel_first"]
    input_names = config["emu"]["input_names"]
    input_nodes = config["emu"]["input_nodes"]
    input_shapes = config["emu"]["input_shapes"]
    onnx = config["emu"]["onnx_file"]
    out_node_shape = config["emu"]["out_node_shape"]
    output_folder = config["emu"]["onnx_output"]
    reordering = config["emu"]["reordering"]
    use_cuda = config["flow"]["cuda"]
    use_ort = config["flow"]["ort"]

    dump_file_name = "_".join([config["emu"]["model_type"], "inputs"])
    dump_file = ".".join([dump_file_name, "pkl", "xz"])
    input_dump = dump_folder / dump_file

    try:
        inputs = utils.prep_inputs(pre_results, input_names, False,
                                   input_dump, inp_trans, input_shapes)
    except exceptions.InvalidInputError as e:
        sys.exit(f"model {config['emu']['model_type']}: {e}")
    return dynasty.dynasty_inference(
        onnx, inputs, reordering, first, input_nodes, out_node_shape, out_dir=output_folder,
        use_cuda=use_cuda, use_ort=use_ort, oup_trans=oup_trans)

def emulator_fixed(config: conf.TestConfig,
                   pre_results: List[npt.ArrayLike],
                   inp_trans: Optional[List[int]] = [0, 3, 1, 2],
                   oup_trans: Optional[List[int]] = [0, 2, 3, 1]) -> List[npt.ArrayLike]:
    """Uses fixed Dynasty to run the preprocessed image through the model.

    Arguments:
        config: TestConfig class that holds the user's input config and options.
        pre_results: List of preprocessed NumPy arrays in channel last format.
        inp_trans: List of integers indicating the axes order to transpose the inputs.
        oup_trans: List of integers indicating the axes order to transpose the outputs.

    Returns:
        A list of NumPy arrays representing the outputs of the model.
    """
    bie = config["emu"]["bie_file"]
    data_type = config["emu"]["data_type"]
    dump_folder = config["flow"]["out_folder"]
    d_ioinfo = config["emu"]["d_ioinfo"]
    first = config["emu"]["channel_first"]
    input_names = config["emu"]["input_names"]
    input_nodes = config["emu"]["input_nodes"]
    input_shapes = config["emu"]["input_shapes"]
    onnx = config["emu"]["onnx_file"]
    out_node_shape = config["emu"]["out_node_shape"]
    output_folder = config["emu"]["onnx_output"]
    platform = config["emu"]["platform"]
    reordering = config["emu"]["reordering"]
    use_onnx_json = config["emu"]["debug"]

    dump_file_name = "_".join([config["emu"]["model_type"], "inputs"])
    dump_file = ".".join([dump_file_name, "pkl", "xz"])
    input_dump = dump_folder / dump_file

    try:
        inputs = utils.prep_inputs(pre_results, input_names, False,
                                   input_dump, inp_trans, input_shapes)
    except exceptions.InvalidInputError as e:
        sys.exit(f"model {config['emu']['model_type']}: {e}")
    if use_onnx_json:
        return dynasty.dynasty_inference(onnx, inputs, reordering, first, input_nodes,
                                         out_node_shape, True, platform, data_type,
                                         output_folder, oup_trans=oup_trans,
                                         d_ioinfo=d_ioinfo)
    else:
        return dynasty.dynasty_inference(bie, inputs, reordering, first, input_nodes,
                                         out_node_shape, True, platform, data_type,
                                         output_folder, oup_trans=oup_trans, d_ioinfo=d_ioinfo)


EMULATOR_MAPPING = {
    "csim": emulator_csim,
    "fixed": emulator_fixed,
    "float": emulator_float,
    "dongle": emulator_dongle,
    "bypass": lambda x, y: print("Bypassing inference..."),
}

def emulator_main(config: conf.TestConfig,
                  pre_results: List[npt.ArrayLike],
                  inp_trans: Optional[List[int]] = [0, 3, 1, 2],
                  oup_trans: Optional[List[int]] = [0, 2, 3, 1]) -> List[npt.ArrayLike]:
    """Main function that calls the appropriate emulator function.

    Arguments:
        config: TestConfig class that holds the user's input config and options.
        pre_results: List of preprocessed NumPy arrays in channel last format.
        inp_trans: List of integers indicating the axes order to transpose the inputs.
        oup_trans: List of integers indicating the axes order to transpose the outputs.

    Returns:
        A list of NumPy arrays representing the outputs in the model.
    """
    try:
        results = EMULATOR_MAPPING.get(
            config["emu"]["emu_mode"],
            lambda x, y: print("Invalid inference mode. Bypassing..."))(
                config, pre_results, inp_trans, oup_trans)
    except exceptions.LibError as error:
        sys.exit(error)
    except exceptions.ConfigError as error:
        sys.exit(error)

    return results