from caffe2.python import control, core, test_util, workspace

import logging
logger = logging.getLogger(__name__)


class TestControl(test_util.TestCase):
    def setUp(self):
        super(TestControl, self).setUp()
        self.N_ = 10

        self.init_net_ = core.Net("init-net")
        cnt = self.init_net_.CreateCounter([], init_count=0)
        const_n = self.init_net_.ConstantFill(
            [], shape=[], value=self.N_, dtype=core.DataType.INT64)
        const_0 = self.init_net_.ConstantFill(
            [], shape=[], value=0, dtype=core.DataType.INT64)

        self.cnt_net_ = core.Net("cnt-net")
        self.cnt_net_.CountUp([cnt])
        curr_cnt = self.cnt_net_.RetrieveCount([cnt])
        self.init_net_.ConstantFill(
            [], [curr_cnt], shape=[], value=0, dtype=core.DataType.INT64)
        self.cnt_net_.AddExternalOutput(curr_cnt)

        self.cnt_2_net_ = core.Net("cnt-2-net")
        self.cnt_2_net_.CountUp([cnt])
        self.cnt_2_net_.CountUp([cnt])
        curr_cnt_2 = self.cnt_2_net_.RetrieveCount([cnt])
        self.init_net_.ConstantFill(
            [], [curr_cnt_2], shape=[], value=0, dtype=core.DataType.INT64)
        self.cnt_2_net_.AddExternalOutput(curr_cnt_2)

        self.cond_net_ = core.Net("cond-net")
        cond_blob = self.cond_net_.LT([curr_cnt, const_n])
        self.cond_net_.AddExternalOutput(cond_blob)

        self.not_cond_net_ = core.Net("not-cond-net")
        cond_blob = self.not_cond_net_.GE([curr_cnt, const_n])
        self.not_cond_net_.AddExternalOutput(cond_blob)

        self.true_cond_net_ = core.Net("true-cond-net")
        true_blob = self.true_cond_net_.LT([const_0, const_n])
        self.true_cond_net_.AddExternalOutput(true_blob)

        self.false_cond_net_ = core.Net("false-cond-net")
        false_blob = self.false_cond_net_.GT([const_0, const_n])
        self.false_cond_net_.AddExternalOutput(false_blob)

        self.idle_net_ = core.Net("idle-net")
        self.idle_net_.ConstantFill(
            [], shape=[], value=0, dtype=core.DataType.INT64)

    def CheckNetOutput(self, nets_and_expects):
        """
        Check the net output is expected
        nets_and_expects is a list of tuples (net, expect)
        """
        for net, expect in nets_and_expects:
            output = workspace.FetchBlob(
                net.Proto().external_output[-1])
            self.assertEqual(output, expect)

    def CheckNetAllOutput(self, net, expects):
        """
        Check the net output is expected
        expects is a list of bools.
        """
        self.assertEqual(len(net.Proto().external_output), len(expects))
        for i in range(len(expects)):
            output = workspace.FetchBlob(
                net.Proto().external_output[i])
            self.assertEqual(output, expects[i])

    def BuildAndRunPlan(self, step):
        plan = core.Plan("test")
        plan.AddStep(control.Do('init', self.init_net_))
        plan.AddStep(step)
        self.assertEqual(workspace.RunPlan(plan), True)

    def ForLoopTest(self, nets_or_steps):
        step = control.For('myFor', nets_or_steps, self.N_)
        self.BuildAndRunPlan(step)
        self.CheckNetOutput([(self.cnt_net_, self.N_)])

    def testForLoopWithNets(self):
        self.ForLoopTest(self.cnt_net_)
        self.ForLoopTest([self.cnt_net_, self.idle_net_])

    def testForLoopWithStep(self):
        step = control.Do('count', self.cnt_net_)
        self.ForLoopTest(step)
        self.ForLoopTest([step, self.idle_net_])

    def WhileLoopTest(self, nets_or_steps):
        step = control.While('myWhile', self.cond_net_, nets_or_steps)
        self.BuildAndRunPlan(step)
        self.CheckNetOutput([(self.cnt_net_, self.N_)])

    def testWhileLoopWithNet(self):
        self.WhileLoopTest(self.cnt_net_)
        self.WhileLoopTest([self.cnt_net_, self.idle_net_])

    def testWhileLoopWithStep(self):
        step = control.Do('count', self.cnt_net_)
        self.WhileLoopTest(step)
        self.WhileLoopTest([step, self.idle_net_])

    def UntilLoopTest(self, nets_or_steps):
        step = control.Until('myUntil', self.not_cond_net_, nets_or_steps)
        self.BuildAndRunPlan(step)
        self.CheckNetOutput([(self.cnt_net_, self.N_)])

    def testUntilLoopWithNet(self):
        self.UntilLoopTest(self.cnt_net_)
        self.UntilLoopTest([self.cnt_net_, self.idle_net_])

    def testUntilLoopWithStep(self):
        step = control.Do('count', self.cnt_net_)
        self.UntilLoopTest(step)
        self.UntilLoopTest([step, self.idle_net_])

    def DoWhileLoopTest(self, nets_or_steps):
        step = control.DoWhile('myDoWhile', self.cond_net_, nets_or_steps)
        self.BuildAndRunPlan(step)
        self.CheckNetOutput([(self.cnt_net_, self.N_)])

    def testDoWhileLoopWithNet(self):
        self.DoWhileLoopTest(self.cnt_net_)
        self.DoWhileLoopTest([self.idle_net_, self.cnt_net_])

    def testDoWhileLoopWithStep(self):
        step = control.Do('count', self.cnt_net_)
        self.DoWhileLoopTest(step)
        self.DoWhileLoopTest([self.idle_net_, step])

    def DoUntilLoopTest(self, nets_or_steps):
        step = control.DoUntil('myDoUntil', self.not_cond_net_, nets_or_steps)
        self.BuildAndRunPlan(step)
        self.CheckNetOutput([(self.cnt_net_, self.N_)])

    def testDoUntilLoopWithNet(self):
        self.DoUntilLoopTest(self.cnt_net_)
        self.DoUntilLoopTest([self.cnt_net_, self.idle_net_])

    def testDoUntilLoopWithStep(self):
        step = control.Do('count', self.cnt_net_)
        self.DoUntilLoopTest(step)
        self.DoUntilLoopTest([self.idle_net_, step])

    def IfCondTest(self, cond_net, expect, cond_on_blob):
        if cond_on_blob:
            step = control.Do(
                'if-all',
                control.Do('count', cond_net),
                control.If('myIf', cond_net.Proto().external_output[-1],
                           self.cnt_net_))
        else:
            step = control.If('myIf', cond_net, self.cnt_net_)
        self.BuildAndRunPlan(step)
        self.CheckNetOutput([(self.cnt_net_, expect)])

    def testIfCondTrueOnNet(self):
        self.IfCondTest(self.true_cond_net_, 1, False)

    def testIfCondTrueOnBlob(self):
        self.IfCondTest(self.true_cond_net_, 1, True)

    def testIfCondFalseOnNet(self):
        self.IfCondTest(self.false_cond_net_, 0, False)

    def testIfCondFalseOnBlob(self):
        self.IfCondTest(self.false_cond_net_, 0, True)

    def IfElseCondTest(self, cond_net, cond_value, expect, cond_on_blob):
        if cond_value:
            run_net = self.cnt_net_
        else:
            run_net = self.cnt_2_net_
        if cond_on_blob:
            step = control.Do(
                'if-else-all',
                control.Do('count', cond_net),
                control.If('myIfElse', cond_net.Proto().external_output[-1],
                           self.cnt_net_, self.cnt_2_net_))
        else:
            step = control.If('myIfElse', cond_net,
                              self.cnt_net_, self.cnt_2_net_)
        self.BuildAndRunPlan(step)
        self.CheckNetOutput([(run_net, expect)])

    def testIfElseCondTrueOnNet(self):
        self.IfElseCondTest(self.true_cond_net_, True, 1, False)

    def testIfElseCondTrueOnBlob(self):
        self.IfElseCondTest(self.true_cond_net_, True, 1, True)

    def testIfElseCondFalseOnNet(self):
        self.IfElseCondTest(self.false_cond_net_, False, 2, False)

    def testIfElseCondFalseOnBlob(self):
        self.IfElseCondTest(self.false_cond_net_, False, 2, True)

    def IfNotCondTest(self, cond_net, expect, cond_on_blob):
        if cond_on_blob:
            step = control.Do(
                'if-not',
                control.Do('count', cond_net),
                control.IfNot('myIfNot', cond_net.Proto().external_output[-1],
                              self.cnt_net_))
        else:
            step = control.IfNot('myIfNot', cond_net, self.cnt_net_)
        self.BuildAndRunPlan(step)
        self.CheckNetOutput([(self.cnt_net_, expect)])

    def testIfNotCondTrueOnNet(self):
        self.IfNotCondTest(self.true_cond_net_, 0, False)

    def testIfNotCondTrueOnBlob(self):
        self.IfNotCondTest(self.true_cond_net_, 0, True)

    def testIfNotCondFalseOnNet(self):
        self.IfNotCondTest(self.false_cond_net_, 1, False)

    def testIfNotCondFalseOnBlob(self):
        self.IfNotCondTest(self.false_cond_net_, 1, True)

    def IfNotElseCondTest(self, cond_net, cond_value, expect, cond_on_blob):
        if cond_value:
            run_net = self.cnt_2_net_
        else:
            run_net = self.cnt_net_
        if cond_on_blob:
            step = control.Do(
                'if-not-else',
                control.Do('count', cond_net),
                control.IfNot('myIfNotElse',
                              cond_net.Proto().external_output[-1],
                              self.cnt_net_, self.cnt_2_net_))
        else:
            step = control.IfNot('myIfNotElse', cond_net,
                                 self.cnt_net_, self.cnt_2_net_)
        self.BuildAndRunPlan(step)
        self.CheckNetOutput([(run_net, expect)])

    def testIfNotElseCondTrueOnNet(self):
        self.IfNotElseCondTest(self.true_cond_net_, True, 2, False)

    def testIfNotElseCondTrueOnBlob(self):
        self.IfNotElseCondTest(self.true_cond_net_, True, 2, True)

    def testIfNotElseCondFalseOnNet(self):
        self.IfNotElseCondTest(self.false_cond_net_, False, 1, False)

    def testIfNotElseCondFalseOnBlob(self):
        self.IfNotElseCondTest(self.false_cond_net_, False, 1, True)

    def testSwitch(self):
        step = control.Switch(
            'mySwitch',
            (self.false_cond_net_, self.cnt_net_),
            (self.true_cond_net_, self.cnt_2_net_)
        )
        self.BuildAndRunPlan(step)
        self.CheckNetOutput([(self.cnt_net_, 0), (self.cnt_2_net_, 2)])

    def testSwitchNot(self):
        step = control.SwitchNot(
            'mySwitchNot',
            (self.false_cond_net_, self.cnt_net_),
            (self.true_cond_net_, self.cnt_2_net_)
        )
        self.BuildAndRunPlan(step)
        self.CheckNetOutput([(self.cnt_net_, 1), (self.cnt_2_net_, 0)])

    def testBoolNet(self):
        bool_net = control.BoolNet(('a', True))
        step = control.Do('bool', bool_net)
        self.BuildAndRunPlan(step)
        self.CheckNetAllOutput(bool_net, [True])

        bool_net = control.BoolNet(('a', True), ('b', False))
        step = control.Do('bool', bool_net)
        self.BuildAndRunPlan(step)
        self.CheckNetAllOutput(bool_net, [True, False])

        bool_net = control.BoolNet([('a', True), ('b', False)])
        step = control.Do('bool', bool_net)
        self.BuildAndRunPlan(step)
        self.CheckNetAllOutput(bool_net, [True, False])

    def testCombineConditions(self):
        # combined by 'Or'
        combine_net = control.CombineConditions(
            'test', [self.true_cond_net_, self.false_cond_net_], 'Or')
        step = control.Do('combine',
                          self.true_cond_net_,
                          self.false_cond_net_,
                          combine_net)
        self.BuildAndRunPlan(step)
        self.CheckNetOutput([(combine_net, True)])

        # combined by 'And'
        combine_net = control.CombineConditions(
            'test', [self.true_cond_net_, self.false_cond_net_], 'And')
        step = control.Do('combine',
                          self.true_cond_net_,
                          self.false_cond_net_,
                          combine_net)
        self.BuildAndRunPlan(step)
        self.CheckNetOutput([(combine_net, False)])

    def testMergeConditionNets(self):
        # merged by 'Or'
        merge_net = control.MergeConditionNets(
            'test', [self.true_cond_net_, self.false_cond_net_], 'Or')
        step = control.Do('merge', merge_net)
        self.BuildAndRunPlan(step)
        self.CheckNetOutput([(merge_net, True)])

        # merged by 'And'
        merge_net = control.MergeConditionNets(
            'test', [self.true_cond_net_, self.false_cond_net_], 'And')
        step = control.Do('merge', merge_net)
        self.BuildAndRunPlan(step)
        self.CheckNetOutput([(merge_net, False)])