import unittest
import numpy as np
import copy
from hypothesis import given
import hypothesis.strategies as st

from caffe2.python.model_helper import ModelHelper
from caffe2.python.models import resnet
from caffe2.python import workspace, brew
import caffe2.python.hypothesis_test_util as hu
import caffe2.python.mkl.rewrite_graph as rewrite_graph


def deterministic_io(model):
    model = copy.deepcopy(model)
    for i, op in enumerate(model.InitProto().op):
        op.device_option.random_seed = i + 1
    if not model.Proto().external_output:
        model.Proto().external_output.extend([model.Proto().op[-1].output[0]])
    return model

def simple_fc():
    model = ModelHelper(name="r")
    brew.fc(model, "data", "fc", 10, 10)
    return model, [(1, 10)]

def double_matmul():
    model = ModelHelper(name="r")
    fc0 = brew.fc(model, "data", "fc0", 10, 10)
    fc1 = brew.fc(model, fc0, "fc1", 10, 10)
    model.Proto().external_output[:] = [str(fc0), str(fc1)]
    return model, [(1, 10)]

def simple_relu():
    model = ModelHelper(name="r")
    brew.relu(model, "data", "fc")
    return model, [(1, 10)]


def simple_mlp():
    model = ModelHelper(name="r")
    brew.relu(
        model,
        brew.fc(
            model,
            brew.relu(
                model,
                brew.fc(
                    model,
                    "data",
                    "fc1",
                    10,
                    10),
                "rl1"),
            "fc2",
            10,
            10),
        "rl2")
    return model, [(1, 10)]


def simple_cnn():
    model = ModelHelper(name="r", arg_scope={"order": "NCHW", "is_test": True})
    brew.conv(
        model, "data", 'conv1', 3, 16, kernel=3, stride=1
    )
    brew.spatial_bn(
        model, 'conv1', 'conv1_spatbn', 16, epsilon=1e-3
    )
    brew.relu(model, 'conv1_spatbn', 'relu1')
    return model, [(1, 3, 32, 32)]


def alexnet():
    model = ModelHelper(name="r", arg_scope={"order": "NCHW", "is_test": True})
    conv1 = brew.conv(
        model,
        "data",
        "conv1",
        3,
        64,
        11, ('XavierFill', {}), ('ConstantFill', {}),
        stride=4,
        pad=2
    )
    relu1 = brew.relu(model, conv1, "conv1")
    pool1 = brew.max_pool(model, relu1, "pool1", kernel=3, stride=2, pad=0,
                          legacy_pad=3)
    lrn1 = brew.lrn(
        model, pool1, "pool1_lrn", size=5, alpha=1.0e-4, beta=0.75, bias=1.0)
    conv2 = brew.conv(
        model,
        lrn1,
        "conv2",
        64,
        192,
        5,
        ('XavierFill', {}),
        ('ConstantFill', {}),
        pad=2
    )
    relu2 = brew.relu(model, conv2, "conv2")
    pool2 = brew.max_pool(model, relu2, "pool2", kernel=3, stride=2)
    lrn2 = brew.lrn(
        model, pool2, "pool2_lrn", size=5, alpha=1.0e-4, beta=0.75, bias=1.0)
    conv3 = brew.conv(
        model,
        lrn2,
        "conv3",
        192,
        384,
        3,
        ('XavierFill', {}),
        ('ConstantFill', {}),
        pad=1
    )
    relu3 = brew.relu(model, conv3, "conv3")
    conv4 = brew.conv(
        model,
        relu3,
        "conv4",
        384,
        256,
        3,
        ('XavierFill', {}),
        ('ConstantFill', {}),
        pad=1
    )
    relu4 = brew.relu(model, conv4, "conv4")
    conv5 = brew.conv(
        model,
        relu4,
        "conv5",
        256,
        256,
        3,
        ('XavierFill', {}),
        ('ConstantFill', {}),
        pad=1
    )
    relu5 = brew.relu(model, conv5, "conv5")
    pool5 = brew.max_pool(model, relu5, "pool5", kernel=3, stride=2)
    fc6 = brew.fc(
        model,
        pool5, "fc6", 256 * 6 * 6, 4096, ('XavierFill', {}),
        ('ConstantFill', {})
    )
    relu6 = brew.relu(model, fc6, "fc6")
    fc7 = brew.fc(
        model, relu6, "fc7", 4096, 4096, ('XavierFill', {}), ('ConstantFill', {})
    )
    relu7 = brew.relu(model, fc7, "fc7")
    drop7 = brew.dropout(model, relu7, "fc7_dropout", is_test=1, ratio=0.5)
    fc8 = brew.fc(
        model, drop7, "fc8", 4096, 1000, ('XavierFill', {}), ('ConstantFill', {})
    )
    relu8 = brew.relu(model, fc8, "fc8")
    brew.dropout(model, relu8, "fc8_dropout", is_test=1, ratio=0.5)
    return model, [(1, 3, 224, 224)]


def simple_resnet():
    model = ModelHelper(name="r", arg_scope={"order": "NCHW", "is_test": True})
    resnet.create_resnet_32x32(
        model, "data", num_input_channels=1, num_groups=1, num_labels=5,
        is_test=True)
    return model, [(1, 1, 32, 32)]


def complex_resnet():
    model = ModelHelper(name="r", arg_scope={"order": "NCHW", "is_test": True})
    resnet.create_resnet50(
        model, "data", num_input_channels=1, num_labels=5, is_test=True,
        no_loss=True)
    return model, [(1, 1, 224, 224)]


@unittest.skipIf(not workspace.C.use_mkldnn, "No MKLDNN support.")
class MKLRewriteTest(hu.HypothesisTestCase):
    @given(gen=st.sampled_from([simple_relu, simple_fc,
                                simple_mlp, simple_cnn]))
    def test_mkl_simple_rewrite(self, gen):
        cpu_model, (shape,) = gen()
        cpu_model = deterministic_io(cpu_model)
        mkl_model = rewrite_graph.rewrite_model_helper_simple(cpu_model)
        X = np.random.randn(*shape).astype(np.float32)

        def run(model):
            self.ws.run(model.InitProto())
            self.ws.create_blob(model.Proto().external_input[0]).feed(X)
            self.ws.run(model.Proto())
            return self.ws.blobs[model.Proto().external_output[0]].fetch()

        np.testing.assert_allclose(run(cpu_model), run(mkl_model),
                                   atol=1e-4, rtol=1e-4)

    def test_mkl_resnet_rewrite(self):
        cpu_model, (shape,) = complex_resnet()
        cpu_model = deterministic_io(cpu_model)
        mkl_model = rewrite_graph.rewrite_model_helper_simple(cpu_model)
        np.random.seed(1701)
        X = np.random.randn(*shape).astype(np.float32)

        def run(model):
            self.ws.run(model.InitProto())
            self.ws.create_blob(model.Proto().external_input[0]).feed(X)
            self.ws.run(model.Proto())
            return self.ws.blobs[model.Proto().external_output[0]].fetch()
        np.testing.assert_allclose(run(cpu_model), run(mkl_model),
                                   atol=1e-4, rtol=1e-4)

    def test_mkl_multi_output_rewrite(self):
        cpu_model, shapes = double_matmul()
        cpu_model = deterministic_io(cpu_model)
        mkl_model = rewrite_graph.rewrite_model_helper_simple(cpu_model)
        np.random.seed(1701)
        Xs = [np.random.randn(*shape).astype(np.float32) for shape in shapes]

        def run(model):
            self.ws.run(model.InitProto())
            for (name, X) in zip(model.Proto().external_input, Xs):
                self.ws.create_blob(name).feed(X)
            print(model.Proto())
            self.ws.run(model.Proto())
            return [self.ws.blobs[name].fetch()
                    for name in model.Proto().external_output]

        run(mkl_model)

        np.testing.assert_allclose(run(cpu_model), run(mkl_model),
                                   atol=1e-4, rtol=1e-4)

    def test_mkl_alexnet_rewrite(self):
        cpu_model, (shape,) = alexnet()
        cpu_model = deterministic_io(cpu_model)
        mkl_model = rewrite_graph.rewrite_model_helper_simple(cpu_model)
        np.random.seed(1701)
        X = np.random.randn(*shape).astype(np.float32)

        def run(model):
            self.ws.run(model.InitProto())
            self.ws.create_blob(model.Proto().external_input[0]).feed(X)
            self.ws.run(model.Proto())
            return self.ws.blobs[model.Proto().external_output[0]].fetch()
        np.testing.assert_allclose(run(cpu_model), run(mkl_model),
                                   atol=1e-4, rtol=1e-4)

if __name__ == "__main__":
    import unittest
    unittest.main()