# 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

const PARTICLE_VALUE_SIZE::Int = 48
const FLOAT_SIZE::Int = 8

"""
    parse_nodes(input::AbstractString)

Parse the given string into a vector of strings containing each node.
"""
function parse_nodes(input::AbstractString)
    regex = r"'([^']*)'"
    matches = eachmatch(regex, input)
    output = [match.captures[1] for match in matches]
    return output
end

"""
    parse_edges(input::AbstractString)

Parse the given string into a vector of strings containing each edge. Currently unused since the entire graph can be read from just the node names.
"""
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

"""
    parse_dag(filename::String, model::ABCModel; verbose::Bool = false)

Read an abc-model process from the given file. If `verbose` is set to true, print some progress information to stdout.

Returns a valid [`DAG`](@ref).
"""
function parse_dag(filename::AbstractString, model::ABCModel, verbose::Bool = false)
    file = open(filename, "r")

    if (verbose)
        println("Opened file")
    end
    nodes_string = readline(file)
    nodes = parse_nodes(nodes_string)

    close(file)
    if (verbose)
        println("Read file")
    end

    graph = DAG()

    # estimate total number of nodes
    # try to slightly overestimate so no resizing is necessary
    # data nodes are not included in length(nodes) and there are a few more than compute nodes
    estimate_no_nodes = round(Int, length(nodes) * 4)
    if (verbose)
        println("Estimating ", estimate_no_nodes, " Nodes")
    end
    sizehint!(graph.nodes, estimate_no_nodes)

    sum_node = insert_node!(graph, make_node(ComputeTaskABC_Sum(0)), track = false, invalidate_cache = false)
    global_data_out = insert_node!(graph, make_node(DataTask(FLOAT_SIZE)), track = false, invalidate_cache = false)
    insert_edge!(graph, sum_node, global_data_out, track = false, invalidate_cache = false)

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

    if (verbose)
        println("Building graph")
    end
    noNodes = 0
    nodesToRead = length(nodes)
    while !isempty(nodes)
        node = popfirst!(nodes)
        noNodes += 1
        if (noNodes % 100 == 0)
            if (verbose)
                percent = string(round(100.0 * noNodes / nodesToRead, digits = 2), "%")
                print("\rReading Nodes... $percent")
            end
        end
        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(PARTICLE_VALUE_SIZE), string(node)),
                track = false,
                invalidate_cache = false,
            ) # read particle data node
            compute_P = insert_node!(graph, make_node(ComputeTaskABC_P()), track = false, invalidate_cache = false) # compute P node
            data_Pu =
                insert_node!(graph, make_node(DataTask(PARTICLE_VALUE_SIZE)), track = false, invalidate_cache = false) # transfer data from P to u (one ABCParticleValue object)
            compute_u = insert_node!(graph, make_node(ComputeTaskABC_U()), track = false, invalidate_cache = false) # compute U node
            data_out =
                insert_node!(graph, make_node(DataTask(PARTICLE_VALUE_SIZE)), track = false, invalidate_cache = false) # transfer data out from u (one ABCParticleValue object)

            insert_edge!(graph, data_in, compute_P, track = false, invalidate_cache = false)
            insert_edge!(graph, compute_P, data_Pu, track = false, invalidate_cache = false)
            insert_edge!(graph, data_Pu, compute_u, track = false, invalidate_cache = false)
            insert_edge!(graph, compute_u, data_out, track = false, invalidate_cache = false)

            # 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(ComputeTaskABC_V()), track = false, invalidate_cache = false)
            data_out =
                insert_node!(graph, make_node(DataTask(PARTICLE_VALUE_SIZE)), track = false, invalidate_cache = false)

            if (occursin(regex_c, in1))
                # put an S node after this input
                compute_S = insert_node!(graph, make_node(ComputeTaskABC_S1()), track = false, invalidate_cache = false)
                data_S_v = insert_node!(
                    graph,
                    make_node(DataTask(PARTICLE_VALUE_SIZE)),
                    track = false,
                    invalidate_cache = false,
                )

                insert_edge!(graph, dataOutNodes[in1], compute_S, track = false, invalidate_cache = false)
                insert_edge!(graph, compute_S, data_S_v, track = false, invalidate_cache = false)

                insert_edge!(graph, data_S_v, compute_v, track = false, invalidate_cache = false)
            else
                insert_edge!(graph, dataOutNodes[in1], compute_v, track = false, invalidate_cache = false)
            end

            if (occursin(regex_c, in2))
                # 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(ComputeTaskABC_S1()), track = false, invalidate_cache = false)
                data_S_v = insert_node!(
                    graph,
                    make_node(DataTask(PARTICLE_VALUE_SIZE)),
                    track = false,
                    invalidate_cache = false,
                )

                insert_edge!(graph, dataOutNodes[in2], compute_S, track = false, invalidate_cache = false)
                insert_edge!(graph, compute_S, data_S_v, track = false, invalidate_cache = false)

                insert_edge!(graph, data_S_v, compute_v, track = false, invalidate_cache = false)
            else
                insert_edge!(graph, dataOutNodes[in2], compute_v, track = false, invalidate_cache = false)
            end

            insert_edge!(graph, compute_v, data_out, track = false, invalidate_cache = false)
            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(ComputeTaskABC_V()), track = false, invalidate_cache = false)
            data_v =
                insert_node!(graph, make_node(DataTask(PARTICLE_VALUE_SIZE)), track = false, invalidate_cache = false)

            insert_edge!(graph, dataOutNodes[in2], compute_v, track = false, invalidate_cache = false)
            insert_edge!(graph, dataOutNodes[in3], compute_v, track = false, invalidate_cache = false)
            insert_edge!(graph, compute_v, data_v, track = false, invalidate_cache = false)

            # combine with the v of the combined other input
            compute_S2 = insert_node!(graph, make_node(ComputeTaskABC_S2()), track = false, invalidate_cache = false)
            data_out = insert_node!(graph, make_node(DataTask(FLOAT_SIZE)), track = false, invalidate_cache = false) # output of a S2 task is only a float

            insert_edge!(graph, data_v, compute_S2, track = false, invalidate_cache = false)
            insert_edge!(graph, dataOutNodes[in1], compute_S2, track = false, invalidate_cache = false)
            insert_edge!(graph, compute_S2, data_out, track = false, invalidate_cache = false)

            insert_edge!(graph, data_out, sum_node, track = false, invalidate_cache = false)
            add_child!(task(sum_node))
        elseif occursin(regex_plus, node)
            if (verbose)
                println("\rReading Nodes Complete    ")
                println("Added ", length(graph.nodes), " nodes")
            end
        else
            @assert false ("Unknown node '$node' while reading from file $filename")
        end
    end

    #put all nodes into dirty nodes set
    graph.dirtyNodes = copy(graph.nodes)

    if (verbose)
        println("Generating the graph's properties")
    end
    graph.properties = GraphProperties(graph)

    if (verbose)
        println("Done")
    end

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

"""
    parse_process(string::AbstractString, model::ABCModel)

Parse a string representation of a process, such as "AB->ABBB" into the corresponding [`ABCProcessDescription`](@ref).
"""
function parse_process(str::AbstractString, model::ABCModel)
    inParticles = Dict{Type, Int}()
    outParticles = Dict{Type, Int}()

    if !(contains(str, "->"))
        throw("Did not find -> while parsing process \"$str\"")
    end

    (inStr, outStr) = split(str, "->")

    if (isempty(inStr) || isempty(outStr))
        throw("Process (\"$str\") input or output part is empty!")
    end

    for t in types(model)
        inCount = count(x -> x == String(t)[1], inStr)
        outCount = count(x -> x == String(t)[1], outStr)
        if inCount != 0
            inParticles[t] = inCount
        end
        if outCount != 0
            outParticles[t] = outCount
        end
    end

    if length(inStr) != sum(values(inParticles))
        throw("Encountered unknown characters in the input part of process \"$str\"")
    elseif length(outStr) != sum(values(outParticles))
        throw("Encountered unknown characters in the output part of process \"$str\"")
    end

    return ABCProcessDescription(inParticles, outParticles)
end
