## @package control_ops_util
# Module caffe2.python.control_ops_util





from caffe2.python import core


def get_external_blob_names(net, lexical_scope):
    """
    Returns a set of blobs a given net depends on and a set of
    output blobs that are written by the net
    Inputs:
        net - net to return input/output blobs for;
        lexical_scope - all external blob names visible to the net
    """
    # Use the blobs that are actually read/written to as external inputs/outputs
    net_proto = net.Proto()
    net_ssa, _ = core.get_ssa(net_proto)
    input_names = core.get_undefined_blobs(net_ssa)
    for input_name in input_names:
        assert str(input_name) in lexical_scope, \
            "Input blob " + input_name + " is undefined"

    output_names = set()
    for op in net_proto.op:
        for output in op.output:
            if output in lexical_scope:
                output_names.add(output)

    return input_names, output_names


def add_if_op(if_net, cond_blob, lexical_scope, then_net, else_net=None):
    """
    A helper function to add an If op to the net.
    Automatically determines whether blobs in the then/else subnets are external
    (from the outer workspace) or local (visible only inside subnet's workspace)
    based on lexical scope - set of all outer blob names visible to the 'If'
    operator. All the blobs in then/else subnets with names matching a name in lexical
    scope and all the blobs that are first used as the operators' inputs are
    considered outer blobs - these blobs must exist in the outer workspace,
    then/else subnets can read their values and new values written into these blobs
    will be visible outside of the 'If' operator. All other blobs are local - exist
    only within inner workspaces for then/else.
    Inputs:
        if_net - net to add an If op to;
        cond_blob - scalar bool blob reference, used as If condition;
        lexical_scope - a set of outer blob names visible to then/else branches;
        then_net/else_net - nets (core.Net) for then/else branches
    """
    then_input_blob_names, then_output_blob_names = get_external_blob_names(
        then_net, lexical_scope)

    else_input_blob_names = set()
    else_output_blob_names = set()
    if else_net:
        else_input_blob_names, else_output_blob_names = get_external_blob_names(
            else_net, lexical_scope)

    input_blob_names = then_input_blob_names | else_input_blob_names
    output_blob_names = then_output_blob_names | else_output_blob_names

    if_inputs = [cond_blob]
    if_inputs += [core.BlobReference(name=b, net=None) for b in input_blob_names]
    if_outputs = [core.BlobReference(name=b, net=None) for b in output_blob_names]

    do_then_net = core.Net('do_then_net')

    then_input_blobs = \
        [core.BlobReference(name=b, net=None) for b in then_input_blob_names]
    then_output_blobs = \
        [core.BlobReference(name=b, net=None) for b in then_output_blob_names]
    then_input_output_names_ordered = [
        str(b) for b in (then_input_blobs + then_output_blobs)]

    then_outer_blob_names = list(then_input_blob_names | then_output_blob_names)
    then_outer_blob_names_idx = [
        then_input_output_names_ordered.index(b) for b in then_outer_blob_names]

    # make sure to use net's name to have unique blob name across multiple subnets
    do_then_workspace_blob = if_net.NextScopedBlob(if_net.Name() + '/workspace_if_then')
    then_input_blobs.append(do_then_workspace_blob)
    then_output_blobs.append(do_then_workspace_blob)
    # make sure that added workspace pointer blobs are in if inputs/outputs
    if_inputs.append(do_then_workspace_blob)
    if_outputs.append(do_then_workspace_blob)

    do_then_net.Do(
        then_input_blobs,
        then_output_blobs,
        net=then_net.Proto(),
        inner_blobs=then_outer_blob_names,
        outer_blobs_idx=then_outer_blob_names_idx)
    do_then_net.AddExternalOutput(*then_output_blobs)

    if_args = {}
    if_args['then_net'] = do_then_net.Proto()

    do_else_workspace_blob = None
    if else_net:
        do_else_net = core.Net('do_else_net')

        else_input_blobs = \
            [core.BlobReference(name=b, net=None) for b in else_input_blob_names]
        else_output_blobs = \
            [core.BlobReference(name=b, net=None) for b in else_output_blob_names]
        else_input_output_names_ordered = [
            str(b) for b in (else_input_blobs + else_output_blobs)]

        else_outer_blob_names = list(else_input_blob_names | else_output_blob_names)
        else_outer_blob_names_idx = [
            else_input_output_names_ordered.index(b) for b in else_outer_blob_names]

        do_else_workspace_blob = \
            if_net.NextScopedBlob(if_net.Name() + '/workspace_if_else')
        else_input_blobs.append(do_else_workspace_blob)
        else_output_blobs.append(do_else_workspace_blob)
        # make sure that added workspace pointer blobs are in if inputs/outputs
        if_inputs.append(do_else_workspace_blob)
        if_outputs.append(do_else_workspace_blob)

        do_else_net.Do(
            else_input_blobs,
            else_output_blobs,
            net=else_net.Proto(),
            inner_blobs=else_outer_blob_names,
            outer_blobs_idx=else_outer_blob_names_idx)
        do_else_net.AddExternalOutput(*else_output_blobs)
        if_args['else_net'] = do_else_net.Proto()

    if_net.CreateScope([], [do_then_workspace_blob])
    if do_else_workspace_blob:
        if_net.CreateScope([], [do_else_workspace_blob])
    if_net.If(if_inputs, if_outputs, **if_args)
    if_net.AddExternalOutput(*if_outputs)


def add_while_op(
        while_net, cond_blob, lexical_scope, loop_body_net, condition_body_net=None):
    """
    A helper function to add a While op to the net. Same rules for determining
    outer and inner blobs as for the 'If' operator apply for the 'While' operator
    loop and condition subnets. If specified, condition net is executed in a separate
    workspace before the first and after each iteration, the last operator must have
    a single scalar boolean output that is written into the condition blob.
    Inputs:
        while_net - net to add a While op to;
        cond_blob - scalar bool blob reference, used as a stop condition;
        lexical_scope - a set of outer blob names visible to the loop's body;
        loop_body_net - net to execute on each iteration;
        condition_body_net - net to compute condition value
    """
    input_blob_names, output_blob_names = get_external_blob_names(
        loop_body_net, lexical_scope)

    # Since it's possible that loop is not going to run even once
    # we have to add loop's external outputs into inputs
    input_blob_names |= output_blob_names

    loop_inputs = [core.BlobReference(name=b, net=None) for b in input_blob_names]
    loop_outputs = [core.BlobReference(name=b, net=None) for b in output_blob_names]

    while_inputs = [cond_blob] + loop_inputs
    while_outputs = [] + loop_outputs

    do_loop_body_net = core.Net('do_loop_body_net')

    loop_input_output_names_ordered = [
        str(b) for b in (loop_inputs + loop_outputs)]
    loop_body_outer_blob_names = list(input_blob_names | output_blob_names)
    loop_body_outer_blob_names_idx = [
        loop_input_output_names_ordered.index(b) for b in loop_body_outer_blob_names]

    do_loop_body_workspace_blob = \
        while_net.NextScopedBlob(while_net.Name() + '/workspace_loop_body')

    loop_inputs.append(do_loop_body_workspace_blob)
    loop_outputs.append(do_loop_body_workspace_blob)
    # make sure that added workspace pointer blobs are in While inputs/outputs
    while_inputs.append(do_loop_body_workspace_blob)
    while_outputs.append(do_loop_body_workspace_blob)

    do_loop_body_net.Do(
        loop_inputs,
        loop_outputs,
        net=loop_body_net.Proto(),
        inner_blobs=loop_body_outer_blob_names,
        outer_blobs_idx=loop_body_outer_blob_names_idx,
        copy_external_blobs=True)
    do_loop_body_net.AddExternalOutput(*loop_outputs)

    while_args = {}
    while_args['loop_net'] = do_loop_body_net.Proto()

    cond_workspace_blob = None
    if condition_body_net:
        cond_input_blob_names, cond_output_blob_names = get_external_blob_names(
            condition_body_net, lexical_scope)

        # make sure condition blob is written by condition net and is
        # visible outside of it
        found_condition_output = False
        for op in condition_body_net.Proto().op:
            if str(cond_blob) in op.output:
                found_condition_output = True
                break
        assert found_condition_output, \
            "Condition net does not write into condition blob"
        if str(cond_blob) not in cond_output_blob_names:
            cond_output_blob_names.add(str(cond_blob))

        cond_inputs = [core.BlobReference(name=b, net=None)
                        for b in cond_input_blob_names]
        assert str(cond_blob) in cond_output_blob_names, \
            'Condition blob expected in condition net output'
        cond_outputs = [core.BlobReference(name=b, net=None)
                        for b in cond_output_blob_names]

        condition_net = core.Net('do_loop_condition_net')

        cond_input_output_names_ordered = [
            str(b) for b in (cond_inputs + cond_outputs)]
        cond_body_outer_blob_names = \
            list(cond_input_blob_names | cond_output_blob_names)
        cond_body_outer_blob_names_idx = [
            cond_input_output_names_ordered.index(b)
            for b in cond_body_outer_blob_names]

        cond_workspace_blob = \
            while_net.NextScopedBlob(while_net.Name() + '/workspace_loop_cond')
        cond_inputs.append(cond_workspace_blob)
        cond_outputs.append(cond_workspace_blob)

        condition_net.Do(
            cond_inputs,
            cond_outputs,
            net=condition_body_net.Proto(),
            inner_blobs=cond_body_outer_blob_names,
            outer_blobs_idx=cond_body_outer_blob_names_idx)
        condition_net.AddExternalOutput(*cond_outputs)

        while_args['cond_net'] = condition_net.Proto()

        while_inputs += [b for b in cond_inputs
                            if str(b) not in input_blob_names]
        while_outputs += [b for b in cond_outputs
                            if str(b) not in output_blob_names]

        if str(cond_blob) not in lexical_scope:
            while_net.ConstantFill(
                [],
                cond_blob,
                dtype=core.DataType.BOOL,
                value=False)

    while_net.CreateScope([], [do_loop_body_workspace_blob])
    if cond_workspace_blob:
        while_net.CreateScope([], [cond_workspace_blob])
    while_net.While(while_inputs, while_outputs, **while_args)
    while_net.AddExternalOutput(*while_outputs)