# SPDX-License-Identifier: MIT
# Copyright © 2020 Patrick Levin
"""Public API of the tensorflowjs graph model Converter"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import json
import os
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import tensorflow as tf
from tensorflow.compat.v1 import GraphDef
from numpy import ndarray as Tensor

from tensorflowjs.converters.common import ARTIFACT_MODEL_JSON_FILE_NAME
from tensorflowjs.converters.common import ARTIFACT_MODEL_TOPOLOGY_KEY
from tensorflowjs.converters.common import ARTIFACT_WEIGHTS_MANIFEST_KEY
from tensorflowjs.converters.common import SIGNATURE_KEY
from tensorflowjs.converters.common import USER_DEFINED_METADATA_KEY
from tensorflowjs.converters.common import FORMAT_KEY
from tensorflowjs.converters.common import TFJS_GRAPH_MODEL_FORMAT

from tensorflowjs.read_weights import read_weights
from google.protobuf.json_format import ParseDict

import tfjs_graph_converter.common as common
from tfjs_graph_converter.compat import convert_int64_to_int32
from tfjs_graph_converter.convert_prelu import replace_prelu, split_fused_prelu
from tfjs_graph_converter.convert_fused_depthwise import split_fused_depthwise
from tfjs_graph_converter.graph_rewrite_util import validate_supported_ops
from tfjs_graph_converter.graph_rewrite_util import harmonize_dtypes
from tfjs_graph_converter.optimization import optimize_graph
import tfjs_graph_converter.quirks as quirks
import tfjs_graph_converter.util as util

SIGNATURE_OUTPUTS = 'outputs'
SIGNATURE_METHOD = 'method_name'


class ModelFormatError(Exception):
    """Format error while trying to parse model"""
    def __init__(self, msg, fmt):
        """Set error message and parsed format

            Args:
                msg: Error message
                fmt: Actual format as found in the model
        """
        super().__init__(msg)
        self.message = msg
        self.format = fmt
        self.input_path = ''
        self.output_path = ''


# this class exists as a way to ensure valid mappings without having check
# inside consumers of the type
class RenameMap:
    """Mapping from input and output nodes to new names"""
    def __init__(self, mapping):
        """Set input and output remapping

            Args:
                mapping: mapping from old to new names - any type that
                         can be converted to a Dict[str, str].
                         New names must be unique.
        """
        def _is_valid_str(s): return isinstance(s, str) and len(s.strip()) > 0
        def _is_valid(k, v): return _is_valid_str(k) and _is_valid_str(v)
        if not isinstance(mapping, dict):
            try:
                mapping = dict(mapping)
            except TypeError as t:
                raise ValueError(f'{t}')
        if any((k, v) for k, v in mapping.items() if not _is_valid(k, v)):
            raise ValueError('mapping must be between non-empty strings')
        self.mapping = mapping

    def apply(self, signature: util.SignatureDef) -> util.SignatureDef:
        """Apply the rename mapping to a SignatureDef"""
        if not isinstance(signature, util.SignatureDef):
            raise ValueError('signature must be a SignatureDef proto')

        # replace the contents of a mapfield with another mapfield
        def _replace_contents(target, replacement):
            keys = [key for key in target]
            # clear the target
            for key in keys:
                del target[key]
            # copy the contents
            for key in replacement:
                target[key].CopyFrom(replacement[key])

        updated = signature
        if any(self.mapping):
            # operate on a local copy such that the order of replacements
            # doesn't matter (i.e. allow for A=B & B=A name swapping)
            inputs, outputs = signature.inputs, signature.outputs
            temp = util.SignatureDef(inputs=inputs, outputs=outputs)
            for old_name, new_name in self.mapping.items():
                if old_name in inputs:
                    temp.inputs[new_name].CopyFrom(inputs[old_name])
                    del temp.inputs[old_name]
                elif old_name in outputs:
                    temp.outputs[new_name].CopyFrom(outputs[old_name])
                    del temp.outputs[old_name]
            _replace_contents(inputs, temp.inputs)
            _replace_contents(outputs, temp.outputs)
        return updated


def _parse_path_and_model_json(model_dir: str) -> Tuple[str, str]:
    """
    Parse model directory name and return path and file name

    Args:
        model_dir: Model file path - either directory name or path + file name

    Raises:
        ValueError: Model or model directory don't exist

    Returns:
        Tuple of directory name and model file name (without directory)
    """
    if model_dir.endswith('.json'):
        if not os.path.isfile(model_dir):
            raise ValueError(f'Model not found: {model_dir}')
        return os.path.split(model_dir)
    if os.path.isdir(model_dir):
        return model_dir, ARTIFACT_MODEL_JSON_FILE_NAME
    raise ValueError(f'Model path is not a directory: {model_dir}')


def _convert_graph_def(message_dict: Dict[str, Any]) -> GraphDef:
    """
    Convert JSON to TF GraphDef message

    Args:
        message_dict: deserialised JSON message

    Returns:
        TF GraphDef message
    """
    message_dict = quirks.fix_node_attributes(message_dict)
    return ParseDict(message_dict, tf.compat.v1.GraphDef())


def _extract_signature_def(model_json: Dict[str, Any]
                           ) -> Optional[util.SignatureDef]:
    """
    Extract the signature definition from the model's meta data.

    Args:
        model_json: JSON dict from TFJS model file

    Returns:
        TF SignatureDef proto; None if meta data is missing or incomplete
    """
    # three possible scenarios:
    #   1. meta data contains a valid signature w/ inputs and outputs
    #   2. meta data contains incomplete signature (missing in- or outputs)
    #   3. meta data is missing or doesn't contain signature
    # this function works for scenario 1)
    if USER_DEFINED_METADATA_KEY not in model_json:
        return None
    meta_data = model_json[USER_DEFINED_METADATA_KEY]
    if SIGNATURE_KEY not in meta_data:
        return None
    signature = meta_data[SIGNATURE_KEY]
    if tf.saved_model.PREDICT_INPUTS not in signature:
        return None
    if tf.saved_model.PREDICT_OUTPUTS not in signature:
        return None
    signature_def = ParseDict(signature, util.SignatureDef())
    if len(signature_def.method_name) == 0:
        signature_def.method_name = tf.saved_model.PREDICT_METHOD_NAME

    def _remove_channel_from_key(mapfield):
        names_to_fix = [key for key in mapfield if key.endswith(':0')]
        for name in names_to_fix:
            mapfield[name[0:-2]].CopyFrom(mapfield[name])
            del mapfield[name]

    _remove_channel_from_key(signature_def.inputs)
    _remove_channel_from_key(signature_def.outputs)
    return signature_def


def _set_signature_dtypes(graph: tf.Graph, signature_def: util.SignatureDef
                          ) -> Tuple[tf.Graph, util.SignatureDef]:
    """Set the dtype of each input and output to match the graph and return
       both
    """
    for key, value in signature_def.inputs.items():
        node = graph.get_tensor_by_name(value.name)
        value.dtype = node.dtype.as_datatype_enum
    for key, value in signature_def.outputs.items():
        node = graph.get_tensor_by_name(value.name)
        value.dtype = node.dtype.as_datatype_enum
    return graph, signature_def


def _create_graph(graph_def: GraphDef,
                  weight_dict: Dict[str, Tensor],
                  modifiers: Dict[str, Callable]) -> tf.Graph:
    """
    Create a TF Graph from nodes

    Args:
        graph_def: TF GraphDef message containing the node graph
        weight_dict: Dictionary from node names to tensor data
        modifiers: Operations to be performed on weights before the conversion

    Raises:
        ValueError: The given graph def contains unsupported operations

    Returns:
        TF Graph for inference or saving
    """
    graph = tf.Graph()
    validate_supported_ops(graph_def)
    with tf.compat.v1.Session(graph=graph):
        for key, value in weight_dict.items():
            if key in modifiers:
                value = (modifiers[key])(value)
            weight_dict[key] = tf.convert_to_tensor(value)
        weight_dict = harmonize_dtypes(graph_def, weight_dict)
        tf.graph_util.import_graph_def(graph_def, weight_dict, name='')

    graph_def = optimize_graph(graph)
    return graph_def_to_graph_v1(graph_def)


def _replace_unsupported_operations(
        input_graph_def: GraphDef) -> Tuple[GraphDef, Dict[str, Callable]]:
    """Replace known unsupported operations by rewriting the input graph"""
    weight_modifiers = dict()
    # split fused ops that contain unsupported activations
    new_graph, modifiers = split_fused_prelu(input_graph_def)
    weight_modifiers.update(modifiers)
    # split fused depthwise convolutions
    new_graph, modifiers = split_fused_depthwise(new_graph)
    weight_modifiers.update(modifiers)
    # replace unsupported activations
    new_graph, modifiers = replace_prelu(new_graph)
    weight_modifiers.update(modifiers)
    return new_graph, weight_modifiers


def _convert_graph_model_to_graph(model_json: Dict[str, Any],
                                  base_path: str,
                                  compat_mode: bool = False
                                  ) -> Tuple[tf.Graph, util.SignatureDef]:
    """
    Convert TFJS JSON model to TF Graph

    Args:
        model_json: JSON dict from TFJS model file
        base_path:  Path to the model file (where to find the model weights)
        compat_mode: True, if only TFJS datatypes should be used

    Returns:
        Tuple of TF Graph for inference or saving and TF signature definition
    """
    if ARTIFACT_MODEL_TOPOLOGY_KEY not in model_json:
        raise ValueError(
            f"model_json is missing key '{ARTIFACT_MODEL_TOPOLOGY_KEY}'")
    model_format = model_json[FORMAT_KEY] if FORMAT_KEY in model_json else ''
    if model_format != TFJS_GRAPH_MODEL_FORMAT:
        raise ModelFormatError(f"unsupported model format: '{model_format}'",
                               model_format)

    topology = model_json[ARTIFACT_MODEL_TOPOLOGY_KEY]

    if ARTIFACT_WEIGHTS_MANIFEST_KEY not in model_json:
        raise ValueError(f'{ARTIFACT_MODEL_JSON_FILE_NAME} is missing key '
                         f"'{ARTIFACT_WEIGHTS_MANIFEST_KEY}'")

    weights_manifest = model_json[ARTIFACT_WEIGHTS_MANIFEST_KEY]
    weight_list = read_weights(weights_manifest, base_path, flatten=True)

    graph_def = _convert_graph_def(topology)
    name, data = common.TFJS_NAME_KEY, common.TFJS_DATA_KEY
    weight_dict = dict((weight[name], weight[data]) for weight in weight_list)
    graph_def, weight_modifiers = _replace_unsupported_operations(graph_def)
    if compat_mode:
        graph_def = convert_int64_to_int32(graph_def)
    graph = _create_graph(graph_def, weight_dict, weight_modifiers)
    signature_def = _extract_signature_def(model_json) or util.infer_signature(
        graph)
    return _set_signature_dtypes(graph, signature_def)


def enable_cuda():
    """
    Call this function *BEFORE* any tensorflow or tfjs_graph_converter function
    to (re-)enable CUDA devices.

    By default, CUDA devices are disabled when using tfjs_graph_converter to
    ensure successful model conversion regardless of GPU capabilities and
    avaiable GPU memory.

    If you need to use CUDA devices and tfjs_graph_converter simultaneously,
    call this function right after importing tfjs_graph_converter:

    ```
    import tfjs_graph_converter as tfjs
    import tensorflow as tf
    tfjs.api.enable_cuda()


    def main() -> None:
        graph = tfjs.api.load_graph_model('models/tfjs_model')
        model = tfjs.api.graph_to_function_v2(graph)
        model_input = ...
        # inference will run on the GPU if available
        result = model(model_input)
        ...

    if __name__ == '__main__'
        main()
    ```
    """
    os.environ['CUDA_VISIBLE_DEVICES'] = ''


def load_graph_model_and_signature(model_dir: str,
                                   compat_mode: bool = False
                                   ) -> Tuple[tf.Graph, util.SignatureDef]:
    """
    Load a TFJS Graph Model from a directory

    Args:
        model_dir: Directory that contains the tfjs model.json and weights;
                alternatively name and path of the model.json if the name
                differs from the default ("model.json")
        compat_mode: If True, only TFJS-compatible datatypes are used

    Returns:
        Tupel of TF frozen graph for inference or saving and TF signature def
    """
    model_path, model_name = _parse_path_and_model_json(model_dir)
    model_file_path = os.path.join(model_path, model_name)
    with open(model_file_path, "r") as model_file:
        model_json = json.load(model_file)
    return _convert_graph_model_to_graph(model_json, model_path, compat_mode)


def load_graph_model(model_dir: str, compat_mode: bool = False) -> tf.Graph:
    """
    Load a TFJS Graph Model from a directory

    Args:
        model_dir: Directory that contains the tfjs model.json and weights;
                alternatively name and path of the model.json if the name
                differs from the default ("model.json")
        compat_mode: If True, only TFJS-compatible datatypes are used

    Returns:
        TF frozen graph for inference or saving
    """
    graph, _ = load_graph_model_and_signature(model_dir, compat_mode)
    return graph


def graph_def_to_graph_v1(graph_def: GraphDef) -> tf.Graph:
    """Convert a GraphDef protobuf message to a tf.Graph

    Use this function to convert the graph message returned by
    `load_graph_model` to a tf.Graph that can be used for inference.

    Args:
        graph_def: GraphDef protobuf message as returned by `load_graph_model`

    Returns:
        tf.Graph for inference.
    """
    graph = tf.Graph()
    with tf.compat.v1.Session(graph=graph):
        tf.graph_util.import_graph_def(graph_def, name='')
    return graph


def graph_to_function_v2(graph: Union[GraphDef, tf.Graph]) -> Callable:
    """Wrap a GraphDef or TF1 frozen graph in a TF2 function for easy inference

    Use this function to convert a GraphDef returned by `load_graph_model` or
    a TF v1 frozen graph into a callable TF2 function.

    Args:
        graph: GraphDef protocol buffer message or TF1 frozen graph

    Returns:
        The function returns a TF2 wrapped function that is callable with
        input tensors or `numpy` arrays as arguments and returns a list of
        model outputs as tensors.
    """
    graph_def = graph.as_graph_def() if isinstance(graph, tf.Graph) else graph

    def _imports_graph_def():
        tf.graph_util.import_graph_def(graph_def, name='')

    wrapped_import = tf.compat.v1.wrap_function(_imports_graph_def, [])
    import_graph = wrapped_import.graph
    inputs = util.get_input_tensors(graph_def)
    outputs = util.get_output_tensors(graph_def)
    return wrapped_import.prune(
        tf.nest.map_structure(import_graph.as_graph_element, inputs),
        tf.nest.map_structure(import_graph.as_graph_element, outputs))


def graph_model_to_frozen_graph(model_dir: str, export_path: str,
                                compat_mode: bool = False) -> str:
    """
    Convert a TFJS graph model to a frozen TF graph

    Args:
        model_dir: Directory that contains the TFJS JSON model and weights
        export_path: Path to the frozen graph (e.g. './output.pb')
        compat_mode: If True, only TFJS-compatible datatypes are used

    Returns:
        The path to the output proto-file.
    """
    export_dir = os.path.dirname(export_path)
    model_name = os.path.basename(export_path)

    graph = load_graph_model(model_dir, compat_mode)
    return tf.io.write_graph(graph, export_dir, model_name, as_text=False)


def graph_model_to_saved_model(model_dir: str,
                               export_dir: str,
                               tags: List[str] = None,
                               signature_def_map: dict = None,
                               signature_key_map: RenameMap = None,
                               compat_mode: bool = False) -> str:
    """
    Convert a TFJS graph model to a SavedModel

    Args:
        model_dir: Directory that contains the TFJS JSON model and weights
        export_dir: Target directory to save the TF model in
        tags: Tags for the SavedModel, optional
        signature_def_map: Dictionary that maps signature keys to model
            signatures;
            A model signature is a dictionary that maps a signature key
            (the name of the signature) - which can be empty or `None` - to
            a dictionary that contains two keys:
            - `outputs`, which maps to a flat list of output tensor names
            - `method_name`, which is an optional name describing the signature
              (defaults to `tf.saved_model.PREDICT_METHOD_NAME`)

            Example:
            {'model_signature': {'outputs': ['Output1', 'Output2']},
            {'debug_signature': {'outputs': ['Output1', 'Output2', 'Detail']}}}

            The names given in the 'outputs' must refer to output tensors in
            the model; otherwise a `ValueError` is raised.
        signature_key_map: Mapping from input- or output tensors to signature
                           keys. The default signature uses tensor names for
                           signature keys. This argument allows to map tensor
                           names to different keys.
        compat_mode: If True, only TFJS-compatible datatypes are used

    Returns:
        The path to which the model was written.
    """
    graph, signature_def = load_graph_model_and_signature(model_dir,
                                                          compat_mode)
    builder = tf.compat.v1.saved_model.Builder(export_dir)
    signature_map = _get_signature_map(graph, signature_def, signature_def_map)
    tags = _get_tags(tags)
    if signature_key_map is not None:
        for signature in signature_map.values():
            signature_key_map.apply(signature)

    with tf.compat.v1.Session(graph=graph) as sess:
        builder.add_meta_graph_and_variables(sess, tags=tags,
                                             signature_def_map=signature_map)
    return builder.save()


def graph_models_to_saved_model(model_list: List[Tuple[str, List[str]]],
                                export_dir: str,
                                signatures: dict = None,
                                signature_keys: Dict[str, RenameMap] = None,
                                compat_mode: bool = False
                                ) -> str:
    """
    Read multiple TFJS graph models and saves them in a single SavedModel

    Args:
        model_list: List of tuples containing TFJS model dir and tags, e.g.
            [("./models/model1", ["step1"]), ("./models/model2": ["step2"])]
        export_dir: Target directory to save the TF model in
        signatures: Dictionary of model names (i.e. the first item in
            `model_list` tuples) to per-model signatures.
            A model signature is a dictionary that maps a signature key
            (the name of the signature) - which can be empty or `None` - to
            a dictionary that contains two keys:
            - `outputs`, which maps to a flat list of output tensor names
            - `method_name`, which is an optional name describing the signature
              (defaults to `tf.saved_model.PREDICT_METHOD_NAME`)
            Empty keys map to the default signature key.
            Keys that don't map to a model in ``model_list`` are ignored.

            {'./models/model2': {'': {'outputs': ['model_output1']},
                                 'debug': {'outputs': ['model_output1',
                                                       'details']}}

            If ``signatures`` is not provided or doesn't contain a key matching
            a model in ``model_list``, the default signature is used.
        signature_keys: Dictionary of model names (i.e. the first item in
            `model_list` tuples) to per-model signature key mappings. This
            allows a remapping of signature inputs and outputs to different
            keys (the tensor names stay unaffected).
        compat_mode: If True, only TFJS-compatible datatypes are used

    Returns:
        The path to which the model was written.
    """
    signatures = signatures or {}
    signature_keys = signature_keys or {}
    builder = tf.compat.v1.saved_model.Builder(export_dir)

    def _apply_key_map(model, signature_map):
        if model in signature_keys and signature_keys[model] is not None:
            signature_key_map = signature_keys[model]
            for signature in signature_map.values():
                signature_key_map.apply(signature)

    model_dir, tags = model_list[0]
    tags = _get_tags(tags)
    graph, signature_def = load_graph_model_and_signature(model_dir,
                                                          compat_mode)
    signature = signatures[model_dir] if model_dir in signatures else None
    signature_map = _get_signature_map(graph, signature_def, signature)
    _apply_key_map(model_dir, signature_map)

    with tf.compat.v1.Session(graph=graph) as sess:
        builder.add_meta_graph_and_variables(sess, tags=tags,
                                             signature_def_map=signature_map)

    for model_dir, tags in model_list[1:]:
        graph, signature_def = load_graph_model_and_signature(model_dir)
        signature = signatures[model_dir] if model_dir in signatures else None
        signature_map = _get_signature_map(graph, signature_def, signature)
        _apply_key_map(model_dir, signature_map)
        tags = _get_tags(tags)
        with tf.compat.v1.Session(graph=graph):
            builder.add_meta_graph(tags=tags, signature_def_map=signature_map)

    return builder.save()


def _get_signature_map(graph: tf.Graph, default_signature_def: dict,
                       signature_def_map: Optional[dict]) -> dict:
    """Select the signature def map from default & user provided signatures"""
    default_key = tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY
    if signature_def_map is None:
        return {default_key: default_signature_def}
    else:
        return _build_signatures(graph, signature_def_map)


def _build_signatures(graph: tf.Graph, signature_def_map: dict) -> dict:
    """
    Turn a signature map into a proper argument for ``add_meta_graph`` by
    creating tensor info for each output name given in ``signature_def_map``.

    Args:
        graph: TF Graph instance
        signature_def_map: Dictionary that maps signature keys to output names

    Returns:
        Signature definition map for passing to
        ``SavedModelBuilder.add_meta_graph``
    """
    inputs = {info.name: util._build_tensor_info(graph, info)
              for info in util.get_input_nodes(graph)}
    default_key = tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY
    output_tensors = set(name[0:-2] for name in util.get_output_tensors(graph))
    for key, value in signature_def_map.items():
        if SIGNATURE_OUTPUTS not in value:
            raise ValueError(f'Signature "{key or default_key}" is invalid: '
                             f'the key "{SIGNATURE_OUTPUTS}" is missing')
        if len(value[SIGNATURE_OUTPUTS]) == 0:
            raise ValueError(f'Signature key "{SIGNATURE_OUTPUTS}" must not'
                             'be empty')
        for name in value[SIGNATURE_OUTPUTS]:
            if name not in output_tensors:
                valid_outputs = str(output_tensors)[1:-1]
                raise ValueError(f'Signature "{key or default_key}" is '
                                 f'invalid: "{name}" is not an output tensor.'
                                 f' Valid outputs are {valid_outputs}')
    return {key or default_key: _build_signature(graph, inputs, info)
            for key, info in signature_def_map.items()}


def _build_signature(graph: tf.Graph, inputs: dict,
                     info: dict) -> util.SignatureDef:
    """
    Return tensor info for inputs and outputs.

    Args:
        graph: TF Graph instance
        inputs: dict that maps input names to tensor information
        info: dict containing the key 'outputs' that maps to flat list of
            output names and an optional key 'method_name'

    Returns:
        dict that maps inputs and outputs to tensor information
        (name, type, and shape)
    """
    def _info(name):
        return util.NodeInfo(name=name, shape=(), dtype=0, tensor=name+':0')

    output_names = info[SIGNATURE_OUTPUTS]
    outputs = {name: util._build_tensor_info(graph, _info(name))
               for name in output_names}
    name = tf.saved_model.PREDICT_METHOD_NAME
    if SIGNATURE_METHOD in info:
        name = info[SIGNATURE_METHOD] or name
    return util.SignatureDef(inputs=inputs, outputs=outputs, method_name=name)


def _get_tags(tags):
    """Reformat tags to a proper list of trimmed strings"""
    if tags is None:
        tags = tf.saved_model.SERVING
    if isinstance(tags, str):
        tags = [tags]
    # remove all empty and whitespace-only tags and return trimmed strings
    return [tag.strip() for tag in tags if bool(tag) and len(tag.strip()) > 0]
