## @package control
# Module caffe2.python.control
"""
Implement functions for controlling execution of nets and steps, including
  Do
  DoParallel
  For-loop
  While-loop
  Do-While-loop
  Switch
  If
"""


from caffe2.python import core
from future.utils import viewitems


# Used to generate names of the steps created by the control functions.
# It is actually the internal index of these steps.
_current_idx = 1
_used_step_names = set()


def _get_next_step_name(control_name, base_name):
    global _current_idx, _used_step_names
    concat_name = '%s/%s' % (base_name, control_name)
    next_name = concat_name
    while next_name in _used_step_names:
        next_name = '%s_%d' % (concat_name, _current_idx)
        _current_idx += 1
    _used_step_names.add(next_name)
    return next_name


def _MakeList(input):
    """ input is a tuple.
    Example:
    (a, b, c)   --> [a, b, c]
    (a)         --> [a]
    ([a, b, c]) --> [a, b, c]
    """
    if len(input) == 0:
        raise ValueError(
            'input cannot be empty.')
    elif len(input) == 1:
        output = input[0]
        if not isinstance(output, list):
            output = [output]
    else:
        output = list(input)
    return output


def _IsNets(nets_or_steps):
    if isinstance(nets_or_steps, list):
        return all(isinstance(n, core.Net) for n in nets_or_steps)
    else:
        return isinstance(nets_or_steps, core.Net)


def _PrependNets(nets_or_steps, *nets):
    nets_or_steps = _MakeList((nets_or_steps,))
    nets = _MakeList(nets)
    if _IsNets(nets_or_steps):
        return nets + nets_or_steps
    else:
        return [Do('prepend', nets)] + nets_or_steps


def _AppendNets(nets_or_steps, *nets):
    nets_or_steps = _MakeList((nets_or_steps,))
    nets = _MakeList(nets)
    if _IsNets(nets_or_steps):
        return nets_or_steps + nets
    else:
        return nets_or_steps + [Do('append', nets)]


def GetConditionBlobFromNet(condition_net):
    """
    The condition blob is the last external_output that must
    be a single bool
    """
    assert len(condition_net.Proto().external_output) > 0, (
        "Condition net %s must has at least one external output" %
        condition_net.Proto.name)
    # we need to use a blob reference here instead of a string
    # otherwise, it will add another name_scope to the input later
    # when we create new ops (such as OR of two inputs)
    return core.BlobReference(condition_net.Proto().external_output[-1])


def BoolNet(*blobs_with_bool_value):
    """A net assigning constant bool values to blobs. It is mainly used for
    initializing condition blobs, for example, in multi-task learning, we
    need to access reader_done blobs before reader_net run. In that case,
    the reader_done blobs must be initialized.

    Args:
    blobs_with_bool_value: one or more (blob, bool_value) pairs. The net will
    assign each bool_value to the corresponding blob.

    returns
    bool_net: A net assigning constant bool values to blobs.

    Examples:
    - BoolNet((blob_1, bool_value_1), ..., (blob_n, bool_value_n))
    - BoolNet([(blob_1, net1), ..., (blob_n, bool_value_n)])
    - BoolNet((cond_1, bool_value_1))
    """
    blobs_with_bool_value = _MakeList(blobs_with_bool_value)
    bool_net = core.Net('bool_net')
    for blob, bool_value in blobs_with_bool_value:
        out_blob = bool_net.ConstantFill(
            [],
            [blob],
            shape=[],
            value=bool_value,
            dtype=core.DataType.BOOL)
        bool_net.AddExternalOutput(out_blob)

    return bool_net


def NotNet(condition_blob_or_net):
    """Not of a condition blob or net

    Args:
    condition_blob_or_net can be either blob or net. If condition_blob_or_net
    is Net, the condition is its last external_output
    that must be a single bool.

    returns
    not_net: the net NOT the input
    out_blob: the output blob of the not_net
    """
    if isinstance(condition_blob_or_net, core.Net):
        condition_blob = GetConditionBlobFromNet(condition_blob_or_net)
    else:
        condition_blob = condition_blob_or_net

    not_net = core.Net('not_net')
    out_blob = not_net.Not(condition_blob)
    not_net.AddExternalOutput(out_blob)

    return not_net, out_blob


def _CopyConditionBlobNet(condition_blob):
    """Make a condition net that copies the condition_blob

    Args:
    condition_blob is a single bool.

    returns
    not_net: the net NOT the input
    out_blob: the output blob of the not_net
    """
    condition_net = core.Net('copy_condition_blob_net')
    out_blob = condition_net.Copy(condition_blob)
    condition_net.AddExternalOutput(out_blob)

    return condition_net, out_blob


def MergeConditionNets(name, condition_nets, relation):
    """
    Merge multi condition nets into a single condition nets.

    Args:
        name: name of the new condition net.
        condition_nets: a list of condition nets. The last external_output
                        of each condition net must be single bool value.
        relation: can be 'And' or 'Or'.

    Returns:
        - A new condition net. Its last external output is relation of all
          condition_nets.
    """
    if not isinstance(condition_nets, list):
        return condition_nets
    if len(condition_nets) <= 1:
        return condition_nets[0] if condition_nets else None

    merged_net = core.Net(name)
    for i in range(len(condition_nets)):
        net_proto = condition_nets[i].Proto()
        assert net_proto.device_option == merged_net.Proto().device_option
        assert net_proto.type == merged_net.Proto().type
        merged_net.Proto().op.extend(net_proto.op)
        merged_net.Proto().external_input.extend(net_proto.external_input)
        # discard external outputs as we're combining them together
        curr_cond = GetConditionBlobFromNet(condition_nets[i])
        if i == 0:
            last_cond = curr_cond
        else:
            last_cond = merged_net.__getattr__(relation)([last_cond, curr_cond])
        # merge attributes
        for k, v in viewitems(condition_nets[i]._attr_dict):
            merged_net._attr_dict[k] += v

    merged_net.AddExternalOutput(last_cond)

    return merged_net


def CombineConditions(name, condition_nets, relation):
    """
    Combine conditions of multi nets into a single condition nets. Unlike
    MergeConditionNets, the actual body of condition_nets is not copied into
    the combine condition net.

    One example is about multi readers. Each reader net has a reader_done
    condition. When we want to check whether all readers are done, we can
    use this function to build a new net.

    Args:
        name: name of the new condition net.
        condition_nets: a list of condition nets. The last external_output
                        of each condition net must be single bool value.
        relation: can be 'And' or 'Or'.

    Returns:
        - A new condition net. Its last external output is relation of all
          condition_nets.
    """
    if not condition_nets:
        return None
    if not isinstance(condition_nets, list):
        raise ValueError('condition_nets must be a list of nets.')

    if len(condition_nets) == 1:
        condition_blob = GetConditionBlobFromNet(condition_nets[0])
        condition_net, _ = _CopyConditionBlobNet(condition_blob)
        return condition_net

    combined_net = core.Net(name)
    for i in range(len(condition_nets)):
        curr_cond = GetConditionBlobFromNet(condition_nets[i])
        if i == 0:
            last_cond = curr_cond
        else:
            last_cond = combined_net.__getattr__(relation)(
                [last_cond, curr_cond])

    combined_net.AddExternalOutput(last_cond)

    return combined_net


def Do(name, *nets_or_steps):
    """
    Execute the sequence of nets or steps once.

    Examples:
    - Do('myDo', net1, net2, ..., net_n)
    - Do('myDo', list_of_nets)
    - Do('myDo', step1, step2, ..., step_n)
    - Do('myDo', list_of_steps)
    """
    nets_or_steps = _MakeList(nets_or_steps)
    if (len(nets_or_steps) == 1 and isinstance(
            nets_or_steps[0], core.ExecutionStep)):
        return nets_or_steps[0]
    else:
        return core.scoped_execution_step(
            _get_next_step_name('Do', name), nets_or_steps)


def DoParallel(name, *nets_or_steps):
    """
    Execute the nets or steps in parallel, waiting for all of them to finish

    Examples:
    - DoParallel('pDo', net1, net2, ..., net_n)
    - DoParallel('pDo', list_of_nets)
    - DoParallel('pDo', step1, step2, ..., step_n)
    - DoParallel('pDo', list_of_steps)
    """
    nets_or_steps = _MakeList(nets_or_steps)
    if (len(nets_or_steps) == 1 and isinstance(
            nets_or_steps[0], core.ExecutionStep)):
        return nets_or_steps[0]
    else:
        return core.scoped_execution_step(
            _get_next_step_name('DoParallel', name),
            nets_or_steps,
            concurrent_substeps=True)


def _RunOnceIf(name, condition_blob_or_net, nets_or_steps):
    """
    Execute nets_or_steps once if condition_blob_or_net evaluates as true.

    If condition_blob_or_net is Net, the condition is its last external_output
    that must be a single bool. And this net will be executed before
    nets_or_steps so as to get the condition.
    """
    condition_not_net, stop_blob = NotNet(condition_blob_or_net)
    if isinstance(condition_blob_or_net, core.Net):
        nets_or_steps = _PrependNets(
            nets_or_steps, condition_blob_or_net, condition_not_net)
    else:
        nets_or_steps = _PrependNets(nets_or_steps, condition_not_net)

    def if_step(control_name):
        return core.scoped_execution_step(
            _get_next_step_name(control_name, name),
            nets_or_steps,
            should_stop_blob=stop_blob,
            only_once=True,
        )

    if _IsNets(nets_or_steps):
        bool_net = BoolNet((stop_blob, False))
        return Do(name + '/_RunOnceIf',
                  bool_net, if_step('_RunOnceIf-inner'))
    else:
        return if_step('_RunOnceIf')


def _RunOnceIfNot(name, condition_blob_or_net, nets_or_steps):
    """
    Similar to _RunOnceIf() but Execute nets_or_steps once if
    condition_blob_or_net evaluates as false.
    """
    if isinstance(condition_blob_or_net, core.Net):
        condition_blob = GetConditionBlobFromNet(condition_blob_or_net)
        nets_or_steps = _PrependNets(nets_or_steps, condition_blob_or_net)
    else:
        copy_net, condition_blob = _CopyConditionBlobNet(condition_blob_or_net)
        nets_or_steps = _PrependNets(nets_or_steps, copy_net)

    return core.scoped_execution_step(
        _get_next_step_name('_RunOnceIfNot', name),
        nets_or_steps,
        should_stop_blob=condition_blob,
        only_once=True,
    )


def For(name, nets_or_steps, iter_num):
    """
    Execute nets_or_steps iter_num times.

    Args:
    nets_or_steps: a ExecutionStep or a Net or a list of ExecutionSteps or
                   a list nets.
    iter_num:    the number times to execute the nets_or_steps.

    Returns:
    A ExecutionStep instance.
    """
    init_net = core.Net('init-net')
    iter_cnt = init_net.CreateCounter([], init_count=iter_num)
    iter_net = core.Net('For-iter')
    iter_done = iter_net.CountDown([iter_cnt])

    for_step = core.scoped_execution_step(
        _get_next_step_name('For-inner', name),
        _PrependNets(nets_or_steps, iter_net),
        should_stop_blob=iter_done)
    return Do(name + '/For',
              Do(name + '/For-init-net', init_net),
              for_step)


def While(name, condition_blob_or_net, nets_or_steps):
    """
    Execute nets_or_steps when condition_blob_or_net returns true.

    Args:
    condition_blob_or_net: If it is an instance of Net, its last
      external_output must be a single bool.
    nets_or_steps: a ExecutionStep or a Net or a list of ExecutionSteps or
                   a list nets.

    Returns:
    A ExecutionStep instance.
    """
    condition_not_net, stop_blob = NotNet(condition_blob_or_net)
    if isinstance(condition_blob_or_net, core.Net):
        nets_or_steps = _PrependNets(
            nets_or_steps, condition_blob_or_net, condition_not_net)
    else:
        nets_or_steps = _PrependNets(nets_or_steps, condition_not_net)

    def while_step(control_name):
        return core.scoped_execution_step(
            _get_next_step_name(control_name, name),
            nets_or_steps,
            should_stop_blob=stop_blob,
        )

    if _IsNets(nets_or_steps):
        # In this case, while_step has sub-nets:
        # [condition_blob_or_net, condition_not_net, nets_or_steps]
        # If stop_blob is pre-set to True (this may happen when While() is
        # called twice), the loop will exit after executing
        # condition_blob_or_net. So we use BootNet to set stop_blob to
        # False.
        bool_net = BoolNet((stop_blob, False))
        return Do(name + '/While', bool_net, while_step('While-inner'))
    else:
        return while_step('While')


def Until(name, condition_blob_or_net, nets_or_steps):
    """
    Similar to While() but execute nets_or_steps when
    condition_blob_or_net returns false
    """
    if isinstance(condition_blob_or_net, core.Net):
        stop_blob = GetConditionBlobFromNet(condition_blob_or_net)
        nets_or_steps = _PrependNets(nets_or_steps, condition_blob_or_net)
    else:
        stop_blob = core.BlobReference(str(condition_blob_or_net))

    return core.scoped_execution_step(
        _get_next_step_name('Until', name),
        nets_or_steps,
        should_stop_blob=stop_blob)


def DoWhile(name, condition_blob_or_net, nets_or_steps):
    """
    Execute nets_or_steps when condition_blob_or_net returns true. It will
    execute nets_or_steps before evaluating condition_blob_or_net.

    Args:
    condition_blob_or_net: if it is an instance of Net, tts last external_output
      must be a single bool.
    nets_or_steps: a ExecutionStep or a Net or a list of ExecutionSteps or
                   a list nets.

    Returns:
    A ExecutionStep instance.
    """
    condition_not_net, stop_blob = NotNet(condition_blob_or_net)
    if isinstance(condition_blob_or_net, core.Net):
        nets_or_steps = _AppendNets(
            nets_or_steps, condition_blob_or_net, condition_not_net)
    else:
        nets_or_steps = _AppendNets(nets_or_steps, condition_not_net)

    # If stop_blob is pre-set to True (this may happen when DoWhile() is
    # called twice), the loop will exit after executing the first net/step
    # in nets_or_steps. This is not what we want. So we use BootNet to
    # set stop_blob to False.
    bool_net = BoolNet((stop_blob, False))
    return Do(name + '/DoWhile', bool_net, core.scoped_execution_step(
        _get_next_step_name('DoWhile-inner', name),
        nets_or_steps,
        should_stop_blob=stop_blob,
    ))


def DoUntil(name, condition_blob_or_net, nets_or_steps):
    """
    Similar to DoWhile() but execute nets_or_steps when
    condition_blob_or_net returns false. It will execute
    nets_or_steps before evaluating condition_blob_or_net.

    Special case: if condition_blob_or_net is a blob and is pre-set to
    true, then only the first net/step of nets_or_steps will be executed and
    loop is exited. So you need to be careful about the initial value the
    condition blob when using DoUntil(), esp when DoUntil() is called twice.
    """
    if not isinstance(condition_blob_or_net, core.Net):
        stop_blob = core.BlobReference(condition_blob_or_net)
        return core.scoped_execution_step(
            _get_next_step_name('DoUntil', name),
            nets_or_steps,
            should_stop_blob=stop_blob)

    nets_or_steps = _AppendNets(nets_or_steps, condition_blob_or_net)
    stop_blob = GetConditionBlobFromNet(condition_blob_or_net)

    # If stop_blob is pre-set to True (this may happen when DoWhile() is
    # called twice), the loop will exit after executing the first net/step
    # in nets_or_steps. This is not what we want. So we use BootNet to
    # set stop_blob to False.
    bool_net = BoolNet((stop_blob, False))
    return Do(name + '/DoUntil', bool_net, core.scoped_execution_step(
        _get_next_step_name('DoUntil-inner', name),
        nets_or_steps,
        should_stop_blob=stop_blob,
    ))


def Switch(name, *conditions):
    """
    Execute the steps for which the condition is true.
    Each condition is a tuple (condition_blob_or_net, nets_or_steps).
    Note:
      1. Multi steps can be executed if their conditions are true.
      2. The conditions_blob_or_net (if it is Net) of all steps will be
         executed once.

    Examples:
    - Switch('name', (cond_1, net_1), (cond_2, net_2), ..., (cond_n, net_n))
    - Switch('name', [(cond_1, net1), (cond_2, net_2), ..., (cond_n, net_n)])
    - Switch('name', (cond_1, net_1))
    """
    conditions = _MakeList(conditions)
    return core.scoped_execution_step(
        _get_next_step_name('Switch', name),
        [_RunOnceIf(name + '/Switch', cond, step) for cond, step in conditions])


def SwitchNot(name, *conditions):
    """
    Similar to Switch() but execute the steps for which the condition is False.
    """
    conditions = _MakeList(conditions)
    return core.scoped_execution_step(
        _get_next_step_name('SwitchNot', name),
        [_RunOnceIfNot(name + '/SwitchNot', cond, step)
         for cond, step in conditions])


def If(name, condition_blob_or_net,
       true_nets_or_steps, false_nets_or_steps=None):
    """
    condition_blob_or_net is first evaluated or executed. If the condition is
    true, true_nets_or_steps is then executed, otherwise, false_nets_or_steps
    is executed.

    If condition_blob_or_net is Net, the condition is its last external_output
    that must be a single bool. And this Net will be executred before both
    true/false_nets_or_steps so as to get the condition.
    """
    if not false_nets_or_steps:
        return _RunOnceIf(name + '/If',
                          condition_blob_or_net, true_nets_or_steps)

    if isinstance(condition_blob_or_net, core.Net):
        condition_blob = GetConditionBlobFromNet(condition_blob_or_net)
    else:
        condition_blob = condition_blob_or_net

    return Do(
        name + '/If',
        _RunOnceIf(name + '/If-true',
                   condition_blob_or_net, true_nets_or_steps),
        _RunOnceIfNot(name + '/If-false', condition_blob, false_nets_or_steps)
    )


def IfNot(name, condition_blob_or_net,
          true_nets_or_steps, false_nets_or_steps=None):
    """
    If condition_blob_or_net returns false, executes true_nets_or_steps,
    otherwise executes false_nets_or_steps
    """
    if not false_nets_or_steps:
        return _RunOnceIfNot(name + '/IfNot',
                             condition_blob_or_net, true_nets_or_steps)

    if isinstance(condition_blob_or_net, core.Net):
        condition_blob = GetConditionBlobFromNet(condition_blob_or_net)
    else:
        condition_blob = condition_blob_or_net

    return Do(
        name + '/IfNot',
        _RunOnceIfNot(name + '/IfNot-true',
                      condition_blob_or_net, true_nets_or_steps),
        _RunOnceIf(name + '/IfNot-false', condition_blob, false_nets_or_steps)
    )