from typing import Set

import onnx
from onnx import ModelProto

from ..onnx_utils import remove_by_name, get_by_name
from ..utils import NNOIR_SUPPORTED_OPS


def fix_postprocess(model: ModelProto):
    """Removes the nnoir-uncompatible nodes at the end of the graph
    Fetch nodes from end until finding a compatible node and then stops.

    Designed to remove unsupported post-process
    (e.g. the part with ArrayFeatureExtractor generated by Azure Custom Vision)

    """
    output_names = set(el.name for el in model.graph.output)
    initializers_names = set(el.name for el in model.graph.initializer)

    idx = len(model.graph.node) - 1
    to_delete_node = []
    to_remove_output = set()
    added_output = set()  # type: Set[str]

    while idx > 0 and not model.graph.node[idx].op_type in NNOIR_SUPPORTED_OPS:
        last_node = model.graph.node[idx]

        to_delete_node.append(last_node)
        if not output_names.union(added_output).isdisjoint(last_node.output):

            to_remove_output.update(output_names.intersection(last_node.output))
            added_output.difference_update(last_node.output)
            added_output.update(set(last_node.input).difference(initializers_names))

        idx -= 1

    node: onnx.NodeProto
    for node in to_delete_node:
        print("deleted: {} ({})".format(node.name, node.op_type))
        model.graph.node.remove(node)

    remove_by_name(model, to_remove_output, 'output')
    inferred_model = onnx.shape_inference.infer_shapes(model)
    value_infos = get_by_name(inferred_model, added_output, "value_info")
    model.graph.output.extend(value_infos.values())
