"""
    getcollection(db_name::String, collection_name::String, client=Mongoc.Client())

Return a `Mongoc.Collection` object corresponding to a MongoDB collection.
"""
function getcollection(db_name::String, collection_name::String, client=Mongoc.Client())
    db = client[db_name]
    return db[collection_name]
end

"""
    isinstance(collection::Mongoc.Collection, instance_name::String)

Check if an instance with "_id" == `instance_name` exists in `collection`.
See also [`getinstance`].
"""
function isinstance(collection::Mongoc.Collection, instance_name::String)
    document = Mongoc.BSON("_id" => instance_name)
    return Mongoc.count_documents(collection, document) >= 1
end

"""
    isdecomposition(collection::Mongoc.Collection, instance_name::String, added_edges::AbstractVector{N})) where N<:AbstractVector{Int}

Check if a decomposition with "_id" = {"instance_name" : `instance_name`, "added_edges" : `added_edges`} exists in `collection`.
See also [`getdecomposition`].
"""
function isdecomposition(collection::Mongoc.Collection, instance_name::String, oid::Union{UUID,AbstractString})
    document = Mongoc.BSON("_id" => Mongoc.BSON("instance_name" => instance_name, "oid" => oid))
    return Mongoc.count_documents(collection, document) >= 1
end

"""
    getinstance(collection::Mongoc.Collection, instance_name::String)

Return a `Mongoc.BSON` corresponding to the document in `collection` with "_id" == `instance_name`. Return `nothing` if document does not exist.
"""
function getinstance(collection::Mongoc.Collection, instance_name::String)
    selector = Mongoc.BSON("_id" => instance_name)
    return Mongoc.find_one(collection, selector)
end

"""
    getdecomposition(collection::Mongoc.Collection, instance_name::String, added_edges::Vector{Vector{Int}})

Return a `Mongoc.BSON` correspong to the document in `collection` width "_id" == { "instance_name" : `instance_name`, "added_edges" : `added_edges`}. Return `nothing` if document does not exist.
"""
function getdecomposition(collection::Mongoc.Collection, instance_name::String, oid::Union{UUID,AbstractString})
    selector = Mongoc.BSON("_id" => Mongoc.BSON("instance_name" => instance_name, "oid" => oid))
    return Mongoc.find_one(collection, selector)
end


"""
    issolved(collection::Mongoc.Collection, instance_name::String, added_edges::Vector{Vector{Int}})

Check if the field `features.solver.solving_time` exists in the specified decomposition.
"""
function issolved(collection::Mongoc.Collection, instance_name::String, oid::Union{UUID,AbstractString})
    document = Mongoc.BSON("_id" => Mongoc.BSON("instance_name" => instance_name, "oid" => oid),
                           "features.solver.solving_time" => Mongoc.BSON("\$exists" => true))
    return Mongoc.count_documents(collection, document) >= 1
end

"""
    getOPFpath_all(collection::Mongoc.Collection)

Return a dict containing the fields "paths.OPF_ctr" and "paths.OPF_mat" for each instance.

# Examples
```julia-repl
julia> getOPFpath_all(instances)
Dict{Any,Any} with 2 entries:
  "case1354pegase" => Dict("mat"=>"path/to/OPF/mat","ctr"=>"path/to/OPF/ctr")
  "case9241pegase" => Dict("mat"=>"path/to/OPF/mat","ctr"=>"path/to/OPF/ctr")
```
"""
function getOPFpath_all(collection::Mongoc.Collection)
    query = Mongoc.BSON("""
        { "paths.OPF_ctr": { "\$exists": true },
          "paths.OPF_mat": { "\$exists": true }
        }""")
    options = Mongoc.BSON("""
        { "projection": 
            { 
                "features": false,
                "paths.matpower": false,
                "decompositions": false
            }
        }""")
    path_dict = Dict()
    for instance in Mongoc.find(collection, query; options=options)
        path_dict[instance["_id"]] = Dict("ctr" => instance["paths"]["OPF_ctr"],
                                           "mat" => instance["paths"]["OPF_mat"])
    end
    return path_dict
end

"""
    getOPFpath_one(collection::Mongoc.Collection, instance_name::String)

Return a dict containing the path of the specified instance.

```julia-repl
julia> getOPFpath_one(instances, "case118mod")
Dict{String, String} with 2 entries:
  "mat" => "path/to/OPF/ctr"
  "ctr" => "path/to/OPF/ctr"
```
"""
function getOPFpath_one(collection::Mongoc.Collection, instance_name::String)
    query = Mongoc.BSON("_id" => instance_name)
    options = Mongoc.BSON("projection" => Mongoc.BSON("paths" => true))
    path_dict = Mongoc.find_one(collection, query; options=options)["paths"]
    isempty(path_dict) && return
    path_dict = Dict("ctr" => path_dict["OPF_ctr"], "mat" => path_dict["OPF_mat"])
    return path_dict
end


function getpath_one(collection::Mongoc.Collection, instance_name::String)
    query = Mongoc.BSON("_id" => instance_name)
    options = Mongoc.BSON("projection" => Mongoc.BSON("paths" => true))
    path_dict = Mongoc.find_one(collection, query; options=options)["paths"]
    isempty(path_dict) && return
    return path_dict
end

"""
    getcholesky(collection::Mongoc.Collection)

Return an array of BSON document where each BSON document is a decomposition where no edges has been added
(called "cholesky graph" in the context of this experiment).
"""
function getcholesky(collection::Mongoc.Collection)
    query = Mongoc.BSON("nb_added_edges" => 0, "dec_type" => "cholesky")
    return Mongoc.find(collection, query)
end

"""
    get_all_matpower_path(collection::Mongoc.Collection)

Get all the matpower path field in `collection` in the form of an array of `Mongoc.BSON` object.
"""
function getmatpowerpath_all(collection::Mongoc.Collection)
    results = collect(Mongoc.find(collection, Mongoc.BSON(); options=Mongoc.BSON("projection" => Mongoc.BSON("paths.matpower" => true))))
    return [path["paths"]["matpower"] for path in results if haskey(path["paths"], "matpower")]
end

function getdecompositions_index(collection::Mongoc.Collection;
                                 unsolved=[true],
                                 original=[false],
                                 merge=[true, false],
                                 combined=[true, false],
                                 select=nothing)::Vector{Int}
    indexes = []
    for (index, document) in enumerate(collection)
        isoriginal = document["nb_added_edges"] == 0
        isunsolved = !haskey(document["features"], "solver") || !haskey(document["features"]["solver"], "solving_time")
        ismerge = document["nb_added_edges"] == -1
        iscombined = document["nb_added_edges"] == -2
        if (isoriginal in original) && (isunsolved in unsolved) && (ismerge in merge) && (iscombined in combined)
            push!(indexes, index)
        end
    end
    if !isnothing(select) 
        decompositions = collect(collection)
        filter!(i -> select(decompositions[i]), indexes)
    end
    return indexes
end

function getcholeskytime(collection::Mongoc.Collection, instance_name::String)
    query = Mongoc.BSON("_id.instance_name" => instance_name, "nb_added_edges" => 0, "dec_type" => "cholesky")
    options = Mongoc.BSON("projection" => Mongoc.BSON("features.solver.solving_time" => true))
    decomposition = Mongoc.find_one(collection, query; options=options)
    if decomposition != nothing && haskey(decomposition["features"], "solver") && haskey(decomposition["features"]["solver"], "solving_time")
        return decomposition["features"]["solver"]["solving_time"]
    end
end

function getcholeskytime_all(collection::Mongoc.Collection)
    query = Mongoc.BSON("nb_added_edges" => 0, "dec_type" => "cholesky")
    options = Mongoc.BSON("projection" => Mongoc.BSON("features.solver.solving_time" => true))
    decompositions = Mongoc.find(collection, query; options=options)
    if decompositions != nothing
        choleskytimes = Dict()
        for d in decompositions
            if haskey(d["features"], "solver")
                choleskytimes[d["_id"]["instance_name"]] = d["features"]["solver"]["solving_time"]
            end
        end
        return choleskytimes
    end
end

function getinstancevalue(collection::Mongoc.Collection, instance_name::String, key::String)
    query = Mongoc.BSON("_id" => instance_name)
    options = Mongoc.BSON("projection" => Mongoc.BSON(key => true))
    instance = Mongoc.find_one(collection, query; options=options)
    isnothing(instance) && return
    keys = split(key, ".")
    for k in keys
        if haskey(instance, k)
            if typeof(instance[k]) <: AbstractDict
                instance = instance[k]
            else
                return instance[k]
            end
        end
    end
end

function getnv(collection::Mongoc.Collection, instance_name::String)
    query = Mongoc.BSON("_id" => instance_name)
    options = Mongoc.BSON("projection" => Mongoc.BSON("features.graph.nv" => true))
    return Mongoc.find_one(collection, query; options=options)["features"]["graph"]["nv"]
end

function getne(collection::Mongoc.Collection, instance_name::String)
    query = Mongoc.BSON("_id" => instance_name)
    options = Mongoc.BSON("projection" => Mongoc.BSON("features.graph.ne" => true))
    return Mongoc.find_one(collection, query; options=options)["features"]["graph"]["ne"]
end

function getkneighbors_indexes(collection::Mongoc.Collection,
                               instance_name::String,
                               oid::Union{UUID,AbstractString},
                               k=3,
                               neighbors_g=getgraphs_dict(collection, x -> x["_id"]["instance_name"] == instance_name && x["features"]["solver"]["cholesky_diff"] <= 0))
    decomposition = getdecomposition(collection, instance_name, oid)
    g = construct_graph_from_cliques(decomposition["cliques"])
    dists = [edge_distance(g, ng[1]) for ng in neighbors_g[instance_name]]
    indexes = sortperm(dists)[begin:k]
    return [n[2] for n in neighbors_g[instance_name][indexes]]
end

function getgraphs_dict(collection::Mongoc.Collection, select=x -> haskey(x["features"], "solver") && x["features"]["solver"]["cholesky_diff"] <= 0)
    graphs = Dict()
    for (i, d) in enumerate(collection)
        if select(d)
            instance_name = d["_id"]["instance_name"]
            gtuple = (construct_graph_from_cliques(d["cliques"]), i)
            if instance_name in keys(graphs)
                push!(graphs[instance_name], gtuple)
            else
                graphs[instance_name] = [gtuple]
            end
        end
    end
    return graphs
end
