metagraphoptimization.jl/src/import.jl

# functions for importing DAGs from a file

regex_a = r"^[A-C]\d+$"                     # Regex for the initial particles
regex_c = r"^[A-C]\(([^']*),([^']*)\)$"     # Regex for the combinations of 2 particles
regex_m = r"^M\(([^']*),([^']*),([^']*)\)$" # Regex for the combinations of 3 particles
regex_plus = r"^\+$"                        # Regex for the sum

function parse_nodes(input::AbstractString)
    regex = r"'([^']*)'"
    matches = eachmatch(regex, input)
    output = [match.captures[1] for match in matches]
    return output
end

function parse_edges(input::AbstractString)
    regex = r"\('([^']*)', '([^']*)'\)"
    matches = eachmatch(regex, input)
    output = [(match.captures[1], match.captures[2]) for match in matches]
    return output
end

function importTxt(filename::String)
    file = open(filename, "r")

    nodes_string = readline(file)
    nodes = parse_nodes(nodes_string)

    close(file)

    graph = DAG()

    sum_node = insert_node(graph, make_node(ComputeTaskSum()))
    global_data_out = insert_node(graph, make_node(DataTask(10)))
    insert_edge(graph, make_edge(sum_node, global_data_out))

    # remember the data out nodes for connection
    dataOutNodes = Dict()

    for node in nodes
        if occursin(regex_a, node)
            # add nodes and edges for the state reading to u(P(Particle))
            data_in = insert_node(graph, make_node(DataTask(4))) # read particle data node
            compute_P = insert_node(graph, make_node(ComputeTaskP())) # compute P node
            data_Pu = insert_node(graph, make_node(DataTask(6))) # transfer data from P to u
            compute_u = insert_node(graph, make_node(ComputeTaskU())) # compute U node
            data_out = insert_node(graph, make_node(DataTask(3))) # transfer data out from u

            insert_edge(graph, make_edge(data_in, compute_P))
            insert_edge(graph, make_edge(compute_P, data_Pu))
            insert_edge(graph, make_edge(data_Pu, compute_u))
            insert_edge(graph, make_edge(compute_u, data_out))
            
            # remember the data_out node for future edges
            dataOutNodes[node] = data_out
        elseif occursin(regex_c, node)
            capt = match(regex_c, node)
            
            in1 = capt.captures[1]
            in2 = capt.captures[2]

            compute_v = insert_node(graph, make_node(ComputeTaskV()))
            data_out = insert_node(graph, make_node(DataTask(5)))

            if (occursin(regex_c, capt.captures[1]))
                # put an S node after this input
                compute_S = insert_node(graph, make_node(ComputeTaskS1()))
                data_S_v = insert_node(graph, make_node(DataTask(5)))

                insert_edge(graph, make_edge(dataOutNodes[capt.captures[1]], compute_S))
                insert_edge(graph, make_edge(compute_S, data_S_v))

                insert_edge(graph, make_edge(data_S_v, compute_v))
            else
                insert_edge(graph, make_edge(dataOutNodes[capt.captures[1]], compute_v))
            end

            if (occursin(regex_c, capt.captures[2]))
                # i think the current generator only puts the combined particles in the first space, so this case might never be entered
                # put an S node after this input
                compute_S = insert_node(graph, make_node(ComputeTaskS1()))
                data_S_v = insert_node(graph, make_node(DataTask(5)))

                insert_edge(graph, make_edge(dataOutNodes[capt.captures[2]], compute_S))
                insert_edge(graph, make_edge(compute_S, data_S_v))

                insert_edge(graph, make_edge(data_S_v, compute_v))
            else
                insert_edge(graph, make_edge(dataOutNodes[capt.captures[2]], compute_v))
            end
 
            insert_edge(graph, make_edge(compute_v, data_out))
            dataOutNodes[node] = data_out

        elseif occursin(regex_m, node)
            # assume for now that only the first particle of the three is combined and the other two are "original" ones
            capt = match(regex_m, node)
            in1 = capt.captures[1]
            in2 = capt.captures[2]
            in3 = capt.captures[3]

            # in2 + in3 with a v
            compute_v = insert_node(graph, make_node(ComputeTaskV()))
            data_v = insert_node(graph, make_node(DataTask(5)))

            insert_edge(graph, make_edge(dataOutNodes[in2], compute_v))
            insert_edge(graph, make_edge(dataOutNodes[in3], compute_v))
            insert_edge(graph, make_edge(compute_v, data_v))

            # combine with the v of the combined other input
            compute_S2 = insert_node(graph, make_node(ComputeTaskS2()))
            data_out = insert_node(graph, make_node(DataTask(10)))

            insert_edge(graph, make_edge(data_v, compute_S2))
            insert_edge(graph, make_edge(dataOutNodes[in1], compute_S2))
            insert_edge(graph, make_edge(compute_S2, data_out))

            insert_edge(graph, make_edge(data_out, sum_node))
        elseif occursin(regex_plus, node)
            println("Found sum node, end")
        else
            error("Unknown node '", node, "' while reading from file ", filename)
        end
    end

    # don't actually need to read the edges
    return graph
end