"""
Utility functions to help perform Dongle inference.
"""
import os
import re
import sys
import time
import json

import numpy as np

from python_flow.common import directory_manager as dm
from python_flow.common import exceptions
from python_flow.utils import csim, utils

from sys_flow.compiler_v2 import combine_nef as sys_combine_nef
from sys_flow.inference import data_converter

def run_dongle(output_folder, inputs, model_id):
    """Performs Dongle inference.

    Arguments:
        output_folder: String path to directory where outputs will be stored.
        model_id: Integer model to run.
        inputs: List of string paths to each of the input image RGBA binaries.

    Returns:
        A tuple of 3 lists, where the first list holds the Dongle response data, the
        second list holds the radices for each output node, and the third list holds the
        scales for each output node.
    """
    import python_flow.dongle.dongle_client as dongle
    from python_flow.dongle import dongle_config

    internal_id = dongle_config.pId_internalId[os.getpid()]

    # Extract the values and make them into a list
    inputs = [value for values in inputs.values() for value in values]

    with dm.DirectoryManager(output_folder):
        try:
            infer_start = time.time()
            # dongle inference
            response, radices, scales = dongle.dongle_clients[internal_id].kdp_inference(
                input_files=inputs, model_id=model_id,
                fixed_output=dongle_config.dongle_fixed_mode)
            infer_end = time.time()
            print(f"Dongle inference time: {infer_end - infer_start}")
        except:
            raise exceptions.LibError(f"Unknown Dongle Error!\n")

    return response, radices, scales

def dongle_to_np(response, reordering, ioinfo_out, first, radices, scales,
                 fixed_output=False, model_id=None, output_dir=None):
    """Reorder dongle output into a list of NumPy arrays according to ioinfo_file.

    The resulting NumPy arrays are values in channel first format if first is true (bchw).
    Otherwise, it will be in channel last format (bhwc).

    Arguments:
        response: response data from dongle
        reordering: List of indices to set output node order
        ioinfo_out: Ioinfo of output nodes.
        first: Boolean to keep result in channel first format
    """
    from python_flow.dongle import dongle_config

    # dump the inference results when dongle debug mode is Enabled
    if dongle_config.dongle_debug_mode:
        for idx, data in enumerate(response):
            if fixed_output:
                file_name = 'dongle_output_fixed_model_{}_node_{}_shape_{}.txt'.format(
                    model_id, idx, data.shape)
                file_path = os.path.join(output_dir, file_name)
                np.savetxt(file_path, data.flatten(), fmt='%d')
            else:
                file_name = 'dongle_output_float_model_{}_node_{}_shape_{}.txt'.format(
                    model_id, idx, data.shape)
                file_path = os.path.join(output_dir, file_name)
                np.savetxt(file_path, data.flatten(), fmt='%f')

    # convert fixed result to float result using e2e method
    if fixed_output:
        response = csim.convert_results(response, radices, scales, "fx2fl")

    reorder_output = [None] * len(response)
    try:
        ioinfo_out_order = {}
        for idx, out_node in enumerate(ioinfo_out):
            ioinfo_out_order[out_node['name']] = idx

        for i, order_name in enumerate(reordering):
            response_idx = int(ioinfo_out_order[order_name])
            reorder_output[i] = response[response_idx]
            if first:
                reorder_output[i] = reorder_output[i].transpose(0, 3, 1, 2)
    except exceptions.ConfigError as error:
        sys.exit(error)

    return reorder_output

def prep_bin_input(pre_results, input_names, ioinfo_in):
    """
    Prepare binary input files as the inputs of dongle inference service.

    Arguments:
        pre_results: List of preprocessed NumPy arrays in channel last format.
        input_names: Node names of the ordered inputs. 
        ioinfo_in: Input ioinfo.
    Returns:
        Paths to the binary files.
    """
    inputs = utils.prep_inputs(pre_results, input_names, False, None)
    csim_bin_list, _, _ = data_converter(inputs, ioinfo_in, p_working=None)

    return csim_bin_list


def combine_nef(user_config, platform):
    """
    Combine one or more nef models to a single combined nef model.
    Extract the ioinfo data stored in each model.

    Arguments:
        user_config: Config params of current settings.
        platform: Interger Field(520, 630, 720, 730, etc...)

    Returns:
        combined_nef_path: Path to the combined nef file.
        combined_ioinfo: ioinfo of all models({model_id: ioinfo, ...})
    """
    nef_paths = []
    for key in user_config.keys():
        # skip the model with emu mode != dongle
        try:
            emu_mode = user_config[key].config["emu"]["emu_mode"]
        except:
            continue
        # add model info into list waiting for compiling if emu mode == dongle
        if emu_mode == "dongle":
            try:
                nef_file = user_config[key].orig_emu['nef_file']
                assert os.path.exists(nef_file) == True
                nef_paths.append(nef_file)
                app_name = re.findall(r'app/(.*?)/models', nef_file)[0]
                dongle_out_dir = os.path.join(user_config[key].orig_emu['dongle_output'], app_name)
            except KeyError as error:
                raise KeyError(r'Failed to collect nef model from inputs. Please check if you include nef_file field in the dongle related input json files.', error)
            except AssertionError:
                raise FileNotFoundError(r'Failed to find the model files. Please check if the nef_file of the dongle related models exist.')

    if len(nef_paths) == 0:
        raise ValueError("Failed to collect nef model from inputs. Please check the input json files.")

    # combine several nef files to a single nef file
    _, p_nef, p_ioinfo, _ = sys_combine_nef(nef_paths, platform, dongle_out_dir)

    # read ioinfo and combined nef path
    combined_nef_path = p_nef
    with open(p_ioinfo, 'r') as f_ioinfo:
        combined_ioinfo = json.load(f_ioinfo)

    return combined_nef_path, combined_ioinfo