"""unittest for ModelHelper class"""



import unittest

from caffe2.python import brew, model_helper


class ModelHelperTest(unittest.TestCase):
    def test_get_complete_net_type(self):
        model = model_helper.ModelHelper("test_orig")
        brew.conv(
            model,
            "input",
            "conv",
            dim_in=3,
            dim_out=16,
            weight_init=("MSRAFill", {}),
            kernel=3,
            stride=1,
            pad=0,
        )
        model.net.Proto().type = "async_scheduling"
        net = model.GetCompleteNet()
        model2 = model_helper.ModelHelper("test_new")
        model2.ConstructInitTrainNetfromNet(net)
        self.assertTrue(model2.net.Proto().type, "async_scheduling")
        self.assertTrue(model2.param_init_net.Proto().type, "async_scheduling")

    def test_get_complete_net(self):
        model = model_helper.ModelHelper("test_orig")
        conv = brew.conv(
            model,
            "input",
            "conv",
            dim_in=3,
            dim_out=16,
            weight_init=("MSRAFill", {}),
            kernel=3,
            stride=1,
            pad=0,
        )
        conv = brew.spatial_bn(model, conv, "conv_bn", 16, epsilon=1e-3, is_test=False)
        conv = brew.relu(model, conv, "conv_relu")
        pred = brew.fc(model, conv, "pred", dim_in=16 * 3 * 3, dim_out=10)
        brew.softmax(model, pred, "softmax")
        net = model.GetCompleteNet()
        model2 = model_helper.ModelHelper("test_new")
        model2.ConstructInitTrainNetfromNet(net)

        net = model.param_init_net
        net2 = model2.param_init_net
        for op1, op2 in zip(net.Proto().op, net2.Proto().op):
            op1.debug_info = op1.debug_info + "/param_init_net"
            self.assertEqual(
                op1, op2, "op mismatch between {}\n and {}\n".format(op1, op2)
            )
        net = model.net
        net2 = model2.net
        for op1, op2 in zip(net.Proto().op, net2.Proto().op):
            self.assertEqual(
                op1, op2, "op mismatch between {}\n and {}\n".format(op1, op2)
            )
        # this is not guaranteed in other situations where user define own net
        self.assertEqual(
            sorted(map(str, net.external_inputs)),
            sorted(map(str, net2.external_inputs)),
        )