import os
import sys
import cv2
import time
from datetime import datetime
import ctypes
import numpy as np
import pycuda.autoinit
import pycuda.driver as cuda
import tensorflow as tf
import uff
import tensorrt as trt
import graphsurgeon as gs

ctypes.CDLL("../../PYTHON_SSD_ALPR/lib/libflattenconcat.so")
# Create TRT plugin nodes to replace unsupported ops in Tensorflow graph
 
def add_plugin(graph):
    all_assert_nodes = graph.find_nodes_by_op("Assert")
    graph.remove(all_assert_nodes, remove_exclusive_dependencies=True)

    all_identity_nodes = graph.find_nodes_by_op("Identity")
    graph.forward_inputs(all_identity_nodes)
    
    Input = gs.create_plugin_node(name="Input",
        op="Placeholder",
        dtype=tf.float32,
        shape=[1, 3, 300, 300]    
    )

    PriorBox = gs.create_plugin_node(name="GridAnchor", op="GridAnchor_TRT",
        minSize=0.2,
        maxSize=0.95,
        aspectRatios=[1.0, 2.0, 0.5, 3.0, 0.33],
        variance=[0.1,0.1,0.2,0.2],
        featureMapShapes=[19, 10, 5, 3, 2, 1],
        numLayers=6
    )
    NMS = gs.create_plugin_node(
        name="NMS",
        op="NMS_TRT",
        shareLocation=1,
        varianceEncodedInTarget=0,
        backgroundLabelId=0,
        confidenceThreshold=1e-8,
        nmsThreshold=0.6,
        topK=100,
        keepTopK=100,
        numClasses=91,
        inputOrder=[0, 2, 1],
        confSigmoid=1,
        isNormalized=1
    )
    concat_priorbox = gs.create_node(
        "concat_priorbox",
        op="ConcatV2",
        dtype=tf.float32,
        axis=2
    )
    concat_box_loc = gs.create_plugin_node(
        "concat_box_loc",
        op="FlattenConcat_TRT",
        dtype=tf.float32,
        axis=1,
        ignoreBatch=0
        
    )
    concat_box_conf = gs.create_plugin_node(
        "concat_box_conf",
        op="FlattenConcat_TRT",
        dtype=tf.float32,
        axis=1,
        ignoreBatch=0
    )

    # Create a mapping of namespace names -> plugin nodes.
    namespace_plugin_map = {
        "MultipleGridAnchorGenerator": PriorBox,
        "Postprocessor": NMS,
        "Preprocessor": Input,
        "ToFloat": Input,
        "image_tensor": Input,
        "MultipleGridAnchorGenerator/Concatenate": concat_priorbox,
        #"MultipleGridAnchorGenerator/Identity": concat_priorbox,
        "concat": concat_box_loc,
        "concat_1": concat_box_conf
    }


    graph.collapse_namespaces(namespace_plugin_map)
    graph.remove(graph.graph_outputs, remove_exclusive_dependencies=False)
    
    
    # Create a constant Tensor and set it as input for GridAnchor_TRT
    #data = np.array([1, 1], dtype=np.float32) 
    #anchor_input = gs.create_node("AnchorInput", "Const", value=data)  
    #graph.append(anchor_input)
    #graph.find_nodes_by_op("GridAnchor_TRT")[0].input.insert(0, "AnchorInput")

    return graph


path = '../model/ssd_inception_v2_coco_2017_11_17/ALPR_model/frozen_inference_graph.pb'
output_uff = '../model/ssd_inception_v2_coco_2017_11_17/ALPR_model/sample_ssd_relu6.uff'

#path = '../model/ssd_inception_v2_coco_2017_11_17/standard_coco_91_classes_model/frozen_inference_graph.pb'
#output_uff = '../model/ssd_inception_v2_coco_2017_11_17/standard_coco_91_classes_model/sample_ssd_relu6.uff'

output_name = ['NMS']

dims = [3,300,300]

# compile model into TensorRT
if not os.path.isfile(output_uff):
	dynamic_graph = gs.DynamicGraph(path)
	graph = add_plugin(dynamic_graph)
	uff_model = uff.from_tensorflow(graph.as_graph_def(), output_name, output_filename=output_uff)