#!/usr/bin/env python3
import os

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = ""
import numpy
import cv2
import tensorflow as tf
from tensorflow.tools.graph_transforms import TransformGraph


def run2d():
    batch_size = 4
    input_shape = [batch_size, 32, 32, 2]
    features = tf.placeholder(tf.float32, input_shape, name='input')

    features0 = tf.layers.conv2d(inputs=features, filters=8, kernel_size=[3, 3], padding='same')
    features1 = tf.layers.conv2d(inputs=features, filters=16, kernel_size=[3, 3], padding='same')
    features2 = tf.concat([features0, features1], axis=-1)
    features3 = tf.layers.conv2d(inputs=features2, filters=8, kernel_size=[3, 3], padding='same')

    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        constant_graph = tf.graph_util.convert_variables_to_constants(sess, sess.graph.as_graph_def(),
                                                                      ['conv2d_2/BiasAdd'])
        tf.train.write_graph(constant_graph, "", "graph_final.pb", as_text=False)

    # export
    with tf.gfile.FastGFile("graph_final.pb", 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        graph_def = TransformGraph(graph_def, ['input'], ['conv2d_2/BiasAdd'], ['strip_unused_nodes'])
        with tf.gfile.FastGFile('saved_model.pb', 'wb') as f:
            f.write(graph_def.SerializeToString())

    # read model
    cvNet = cv2.dnn.readNetFromTensorflow('./saved_model.pb')
    img = numpy.zeros((batch_size, 2, 32, 32), dtype='float32')
    cvNet.setInput(img)
    cvOut = cvNet.forward()
    print(cvOut.shape)


def run3d():
    batch_size = 4
    input_shape = [batch_size, 32, 32, 32, 2]
    features = tf.placeholder(tf.float32, input_shape, name='input')

    features0 = tf.layers.conv3d(inputs=features, filters=8, kernel_size=[3, 3, 3], padding='same')
    features1 = tf.layers.conv3d(inputs=features, filters=16, kernel_size=[3, 3, 3], padding='same')
    features2 = tf.concat([features0, features1], axis=-1)
    features3 = tf.layers.conv3d(inputs=features2, filters=8, kernel_size=[3, 3, 3], padding='same')

    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        constant_graph = tf.graph_util.convert_variables_to_constants(sess, sess.graph.as_graph_def(),
                                                                      ['conv3d_2/BiasAdd'])
        tf.train.write_graph(constant_graph, "", "graph_final.pb", as_text=False)

    # export
    with tf.gfile.FastGFile("graph_final.pb", 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        graph_def = TransformGraph(graph_def, ['input'], ['conv3d_2/BiasAdd'], ['strip_unused_nodes'])
        with tf.gfile.FastGFile('saved_model.pb', 'wb') as f:
            f.write(graph_def.SerializeToString())

    # read model
    cvNet = cv2.dnn.readNetFromTensorflow('./saved_model.pb')
    img = numpy.zeros((batch_size, 2, 32, 32, 32), dtype='float32')
    cvNet.setInput(img)
    cvOut = cvNet.forward()
    print(cvOut.shape)


if __name__ == '__main__':
    print(cv2.__version__)
    print(tf.__version__)
    run2d()
    run3d()