from caffe2.proto import caffe2_pb2
from caffe2.python import core, workspace
import onnx
import onnx.defs
from onnx.helper import make_node, make_graph, make_tensor_value_info, make_model
from onnx.backend.base import namedtupledict
from caffe2.python.models.download import ModelDownloader
import caffe2.python.onnx.backend as c2
from caffe2.python.onnx.workspace import Workspace
from caffe2.python.trt.transform import convert_onnx_model_to_trt_op, transform_caffe2_net
from caffe2.python.onnx.tests.test_utils import TestCase
import numpy as np
import os.path
import time
import unittest
import tarfile
import tempfile
import shutil
from six.moves.urllib.request import urlretrieve

def _print_net(net):
    for i in net.external_input:
        print("Input: {}".format(i))
    for i in net.external_output:
        print("Output: {}".format(i))
    for op in net.op:
        print("Op {}".format(op.type))
        for x in op.input:
            print("  input: {}".format(x))
        for y in op.output:
            print("  output: {}".format(y))


def _base_url(opset_version):
    return 'https://s3.amazonaws.com/download.onnx/models/opset_{}'.format(opset_version)

# TODO: This is copied from https://github.com/onnx/onnx/blob/master/onnx/backend/test/runner/__init__.py. Maybe we should
# expose a model retrival API from ONNX
def _download_onnx_model(model_name, opset_version):
    onnx_home = os.path.expanduser(os.getenv('ONNX_HOME', os.path.join('~', '.onnx')))
    models_dir = os.getenv('ONNX_MODELS',
                           os.path.join(onnx_home, 'models'))
    model_dir = os.path.join(models_dir, model_name)
    if not os.path.exists(os.path.join(model_dir, 'model.onnx')):
        if os.path.exists(model_dir):
            bi = 0
            while True:
                dest = '{}.old.{}'.format(model_dir, bi)
                if os.path.exists(dest):
                    bi += 1
                    continue
                shutil.move(model_dir, dest)
                break
        os.makedirs(model_dir)

        # On Windows, NamedTemporaryFile can not be opened for a
        # second time
        url = '{}/{}.tar.gz'.format(_base_url(opset_version), model_name)
        download_file = tempfile.NamedTemporaryFile(delete=False)
        try:
            download_file.close()
            print('Start downloading model {} from {}'.format(
                model_name, url))
            urlretrieve(url, download_file.name)
            print('Done')
            with tarfile.open(download_file.name) as t:
                t.extractall(models_dir)
        except Exception as e:
            print('Failed to prepare data for model {}: {}'.format(
                model_name, e))
            raise
        finally:
            os.remove(download_file.name)
    return model_dir

class TensorRTOpTest(TestCase):
    def setUp(self):
        self.opset_version = onnx.defs.onnx_opset_version()

    def _test_relu_graph(self, X, batch_size, trt_max_batch_size):
        node_def = make_node("Relu", ["X"], ["Y"])
        Y_c2 = c2.run_node(node_def, {"X": X})
        graph_def = make_graph(
            [node_def],
            name="test",
            inputs=[make_tensor_value_info("X", onnx.TensorProto.FLOAT, [batch_size, 1, 3, 2])],
            outputs=[make_tensor_value_info("Y", onnx.TensorProto.FLOAT, [batch_size, 1, 3, 2])])
        model_def = make_model(graph_def, producer_name='relu-test')
        op_outputs = [x.name for x in model_def.graph.output]
        op = convert_onnx_model_to_trt_op(model_def, max_batch_size=trt_max_batch_size)
        device_option = core.DeviceOption(caffe2_pb2.CUDA, 0)
        op.device_option.CopyFrom(device_option)
        Y_trt = None
        ws = Workspace()
        with core.DeviceScope(device_option):
            ws.FeedBlob("X", X)
            ws.RunOperatorsOnce([op])
            output_values = [ws.FetchBlob(name) for name in op_outputs]
            Y_trt = namedtupledict('Outputs', op_outputs)(*output_values)
        np.testing.assert_almost_equal(Y_c2, Y_trt)


    @unittest.skipIf(not workspace.C.use_trt, "No TensortRT support")
    def test_relu_graph_simple(self):
        X = np.random.randn(1, 1, 3, 2).astype(np.float32)
        self._test_relu_graph(X, 1, 50)


    @unittest.skipIf(not workspace.C.use_trt, "No TensortRT support")
    def test_relu_graph_big_batch(self):
        X = np.random.randn(52, 1, 3, 2).astype(np.float32)
        self._test_relu_graph(X, 52, 50)

    def _test_onnx_importer(self, model_name, data_input_index, opset_version=onnx.defs.onnx_opset_version()):
        model_dir = _download_onnx_model(model_name, opset_version)
        model_def = onnx.load(os.path.join(model_dir, 'model.onnx'))
        input_blob_dims = [int(x.dim_value) for x in model_def.graph.input[data_input_index].type.tensor_type.shape.dim]
        op_inputs = [x.name for x in model_def.graph.input]
        op_outputs = [x.name for x in model_def.graph.output]
        print("{}".format(op_inputs))
        data = np.random.randn(*input_blob_dims).astype(np.float32)
        Y_c2 = c2.run_model(model_def, {op_inputs[data_input_index]: data})
        op = convert_onnx_model_to_trt_op(model_def, verbosity=3)
        device_option = core.DeviceOption(caffe2_pb2.CUDA, 0)
        op.device_option.CopyFrom(device_option)
        Y_trt = None
        ws = Workspace()
        with core.DeviceScope(device_option):
            ws.FeedBlob(op_inputs[data_input_index], data)
            if opset_version >= 5:
                # Some newer models from ONNX Zoo come with pre-set "data_0" input
                ws.FeedBlob("data_0", data)
            ws.RunOperatorsOnce([op])
            output_values = [ws.FetchBlob(name) for name in op_outputs]
            Y_trt = namedtupledict('Outputs', op_outputs)(*output_values)
        np.testing.assert_allclose(Y_c2, Y_trt, rtol=1e-3)

    @unittest.skipIf(not workspace.C.use_trt, "No TensortRT support")
    def test_resnet50(self):
        self._test_onnx_importer('resnet50', 0, 9)

    @unittest.skipIf(not workspace.C.use_trt, "No TensortRT support")
    def test_bvlc_alexnet(self):
        self._test_onnx_importer('bvlc_alexnet', 0, 9)

    @unittest.skip("Until fixing Unsqueeze op")
    def test_densenet121(self):
        self._test_onnx_importer('densenet121', -1, 3)

    @unittest.skipIf(not workspace.C.use_trt, "No TensortRT support")
    def test_inception_v1(self):
        self._test_onnx_importer('inception_v1', -3, 9)

    @unittest.skip("Until fixing Unsqueeze op")
    def test_inception_v2(self):
        self._test_onnx_importer('inception_v2', 0, 9)

    @unittest.skip('Need to revisit our ChannelShuffle exporter to avoid generating 5D tensor')
    def test_shufflenet(self):
        self._test_onnx_importer('shufflenet', 0)

    @unittest.skipIf(not workspace.C.use_trt, "No TensortRT support")
    def test_squeezenet(self):
        self._test_onnx_importer('squeezenet', -1, 9)

    @unittest.skipIf(not workspace.C.use_trt, "No TensortRT support")
    def test_vgg16(self):
        self._test_onnx_importer('vgg16', 0, 9)

    @unittest.skipIf(not workspace.C.use_trt, "No TensortRT support")
    def test_vgg19(self):
        self._test_onnx_importer('vgg19', -2, 9)


class TensorRTTransformTest(TestCase):
    def setUp(self):
        self.model_downloader = ModelDownloader()

    def _add_head_tail(self, pred_net, new_head, new_tail):
        orig_head = pred_net.external_input[0]
        orig_tail = pred_net.external_output[0]

        # Add head
        head = caffe2_pb2.OperatorDef()
        head.type = "Copy"
        head.input.append(new_head)
        head.output.append(orig_head)
        dummy = caffe2_pb2.NetDef()
        dummy.op.extend(pred_net.op)
        del pred_net.op[:]
        pred_net.op.extend([head])
        pred_net.op.extend(dummy.op)
        pred_net.external_input[0] = new_head

        # Add tail
        tail = caffe2_pb2.OperatorDef()
        tail.type = "Copy"
        tail.input.append(orig_tail)
        tail.output.append(new_tail)
        pred_net.op.extend([tail])
        pred_net.external_output[0] = new_tail

    @unittest.skipIf(not workspace.C.use_trt, "No TensortRT support")
    def test_resnet50_core(self):
        N = 2
        warmup = 20
        repeat = 100
        print("Batch size: {}, repeat inference {} times, warmup {} times".format(N, repeat, warmup))
        init_net, pred_net, _ = self.model_downloader.get_c2_model('resnet50')
        self._add_head_tail(pred_net, 'real_data', 'real_softmax')
        input_blob_dims = (N, 3, 224, 224)
        input_name = "real_data"

        device_option = core.DeviceOption(caffe2_pb2.CUDA, 0)
        init_net.device_option.CopyFrom(device_option)
        pred_net.device_option.CopyFrom(device_option)
        for op in pred_net.op:
            op.device_option.CopyFrom(device_option)
            op.engine = 'CUDNN'
        net_outputs = pred_net.external_output
        Y_c2 = None
        data =  np.random.randn(*input_blob_dims).astype(np.float32)
        c2_time = 1
        workspace.SwitchWorkspace("gpu_test", True)
        with core.DeviceScope(device_option):
            workspace.FeedBlob(input_name, data)
            workspace.RunNetOnce(init_net)
            workspace.CreateNet(pred_net)
            for _ in range(warmup):
                workspace.RunNet(pred_net.name)
            start = time.time()
            for _ in range(repeat):
                workspace.RunNet(pred_net.name)
            end = time.time()
            c2_time = end - start
            output_values = [workspace.FetchBlob(name) for name in net_outputs]
            Y_c2 = namedtupledict('Outputs', net_outputs)(*output_values)
        workspace.ResetWorkspace()

        # Fill the workspace with the weights
        with core.DeviceScope(device_option):
            workspace.RunNetOnce(init_net)

        # Cut the graph
        start = time.time()
        pred_net_cut = transform_caffe2_net(pred_net,
                                            {input_name: input_blob_dims},
                                            build_serializable_op=False)
        del init_net, pred_net
        pred_net_cut.device_option.CopyFrom(device_option)
        for op in pred_net_cut.op:
            op.device_option.CopyFrom(device_option)
        #_print_net(pred_net_cut)

        Y_trt = None
        input_name = pred_net_cut.external_input[0]
        print("C2 runtime: {}s".format(c2_time))
        with core.DeviceScope(device_option):
            workspace.FeedBlob(input_name, data)
            workspace.CreateNet(pred_net_cut)
            end = time.time()
            print("Conversion time: {:.2f}s".format(end -start))

            for _ in range(warmup):
                workspace.RunNet(pred_net_cut.name)
            start = time.time()
            for _ in range(repeat):
                workspace.RunNet(pred_net_cut.name)
            end = time.time()
            trt_time = end - start
            print("TRT runtime: {}s, improvement: {}%".format(trt_time, (c2_time-trt_time)/c2_time*100))
            output_values = [workspace.FetchBlob(name) for name in net_outputs]
            Y_trt = namedtupledict('Outputs', net_outputs)(*output_values)
        np.testing.assert_allclose(Y_c2, Y_trt, rtol=1e-3)