using Random

function mpi_init()
    !MPI.Initialized() && MPI.Init()
    comm = MPI.COMM_WORLD
    rank = MPI.Comm_rank(comm)
    size = MPI.Comm_size(comm)
    return comm, rank, size
end

"""
    get_indexes(manager, action, size, unsolved=[true, false], original=[true, false], merged=[false])


Retrieve the indexes used to process either instances or decompositions.

# Examples (results depends on the database)
```julia-repl
julia> manager = ExperimentManager()
julia> get_indexes(manager, :solve, 4)
4-element Array{Array{Int64,1},1}:
 [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
 [12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22]
 [23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33]
 [34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44]

julia> get_indexes(manager, :solve, 4, [true])
4-element Array{Array{Int64,1},1}:
 [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
 [12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22]
 [23, 24, 25, 26, 27, 28, 29, 30, 31, 32]
 [33, 34, 35, 36, 37, 38, 39, 40, 41, 42]

julia> get_indexes(manager, :gen, 4)
4-element Array{Array{Int64,1},1}:
 [1, 2, 3, 4, 5, 6]
 [7, 8, 9, 10, 11]
 [12, 13, 14, 15, 16]
 [17, 18, 19, 20, 21]
```
"""
function get_indexes(manager, action, size, unsolved=[true, false], original=[true, false], merged=[true, false], combined=[true, false], select=nothing)
    indexes = []
    if action in [:solve, :merge, :combine]
        decompositions_index = db.getdecompositions_index(manager.decompositions;
                                                          unsolved=unsolved,
                                                          original=original,
                                                          merge=merged,
                                                          combined=combined,
                                                          select=select)
        indexes = split_index(length(decompositions_index), size)
        indexes = [decompositions_index[i[begin]:i[end]] for i in indexes]
    elseif action in [:gen, :genoriginal]
        paths_matpower = db.getmatpowerpath_all(manager.instances)
        indexes = split_index(length(paths_matpower), size)
        indexes = [[i[begin]:i[end]...] for i in indexes]
    end
    return indexes
end

function balance_decomposition(decompositions, chunk_size)
    decompositions = vcat([sort(decomposition[n:n + chunk_size], by=d->d["features"]["chordal_graph"]["nv"]) for n in 1:chunk_size:length(decompositions) - chunk_size]...)
    return decompositions
end

function process_decomposition(manager, decompositions, action, rank=nothing;
                               mpi_infos=nothing,
                               merge_type=:molzahn, 
                               merge_treshold_type=:nv,
                               merge_treshold=0.1,
                               combine_type=:intersect,
                               neighbors_g=db.getgraphs_dict())
    for (i, decomposition) in enumerate(decompositions)
        instance_name = decomposition["_id"]["instance_name"]
        oid = decomposition["_id"]["oid"]
        if action == :merge
            @info "[$rank] Merging $(instance_name): $(oid), $(decomposition["nb_added_edges"])"
            genmerge(manager, instance_name, oid;
                     merge_type=merge_type,
                     merge_treshold_type=merge_treshold_type,
                     merge_treshold=merge_treshold)
            @info "[$rank] $i / $(length(decompositions)) $(instance_name): $(oid), $(decomposition["nb_added_edges"]) merged"
        elseif action == :solve
            @info "[$rank] Solving $(instance_name): $(oid), $(decomposition["nb_added_edges"])"
            solve_decomposition(manager, instance_name, oid, true, mpi_infos)
            @info "[$rank] $i / $(length(decompositions)) $(instance_name): $(oid), $(decomposition["nb_added_edges"]) solved"
        elseif action == :combine
            @info "[$rank] Combining $(instance_name): $(oid), $(decomposition["nb_added_edges"])"
            results = gencombine(manager, instance_name, oid; combine_type=combine_type, neighbors_g=neighbors_g)
            if isnothing(results)
                @info "[$rank] $i / $(length(decompositions)) $(instance_name): $(oid), $(decomposition["nb_added_edges"]) could not be combined"
            else
                @info "[$rank] $i / $(length(decompositions)) $(instance_name): $(oid), $(decomposition["nb_added_edges"]) combined"
            end
        end
    end
end

function process_instances(manager, paths_matpower, action, rank=nothing;
                           configfile=nothing,
                           kwargs...)
    for (i, path) in enumerate(paths_matpower)
        @info "[$rank] $i / $(length(paths_matpower)) Generating $(split(basename(path), ".")[1])"
        if action == :genoriginal
            genoriginal(manager, path; kwargs...)
        elseif action == :gen
            # Here we could save some operations by parsing the JSON file only once.
            # Maybe I'll change it later but for the moment I find it clearer to do it this way
            # and I don't think it add too much computation time (I didn't test so it's only guess until I test)
            gendecomposition_json(manager, path, configfile; kwargs...)
        end
        @info "[$rank] $i / $(length(paths_matpower)) $(basename(path)) generated"
    end
end

function mpirun(action;
                port="mongodb://localhost:27017",
                db_name="choleskyexp",
                configfile::String="./config.json",
                mpi_infos::AbstractDict=Dict(),
                unsolved=[true],
                original=[false],
                merged=[true, false],
                combined=[true, false],
                finalize=true,
                merge_treshold_type=:nv,
                merge_treshold=0.1,
                merge_type::Symbol=:molzahn,
                combine_type::Symbol=:intersect,
                seed=1234,
                kwargs...)
    Random.seed!(seed)
    comm, rank, size = mpi_init()
    manager = ExperimentManager(port=port, db=db_name)
    neighbors_g = action == :combine ? db.getgraphs_dict(manager.decompositions) : nothing
    if rank == 0
        @info "Size: $size"
        select = action == :combine ? (x -> x["features"]["solver"]["cholesky_diff"] < 0) : nothing
        indexes = get_indexes(manager, action, size, unsolved, original, merged, combined, select)
        for (i, index) in enumerate(indexes)
            @info "[$rank] Rank $(i - 1): $(index[begin]) -> $(index[end]) | $(length(index))"
            MPI.Isend(index, i - 1, 0, comm)
        end
        if length(indexes) < size
            for i in length(indexes) + 1:size
                MPI.Isend(Vector{Int}(undef,0), i - 1, 0, comm)
            end
        end
    end
    MPI.Barrier(comm)
    @info "[$rank] Retrieving the number of instances/decompositions to process"
    status = MPI.Probe(0, 0, comm)
    count = MPI.Get_count(status, Int)
    if count == 0
        @info "[$rank] Exciting because no instances/decompositions to process"
        MPI.Finalize()
        return
    end
    indexes = Array{Int}(undef, count)
    MPI.Irecv!(indexes, 0, 0, comm)
    @info "[$rank] Processing"
    if action in [:merge, :solve, :combine]
        decompositions = collect(manager.decompositions)[indexes]
        sort!(decompositions, by=d -> d["features"]["chordal_graph"]["nv"])
        process_decomposition(manager, decompositions, action, rank;
                              merge_type=merge_type,
                              merge_treshold_type=merge_treshold_type,
                              merge_treshold=merge_treshold,
                              combine_type=combine_type,
                              neighbors_g=neighbors_g,
                              mpi_infos=mpi_infos)
    elseif action in [:gen, :genoriginal]
        paths_matpower = db.getmatpowerpath_all(manager.instances)[indexes]
        process_instances(manager, paths_matpower, action, rank;
                          configfile,
                          kwargs...)
    end
    @info "[$rank] Process done."
    finalize && MPI.Finalize()
end
