"""
    KaHyParBipartite{RT,IT,GM}
    KaHyParBipartite(; sc_target, imbalances=collect(0.0:0.005:0.8),
        max_group_size=40, greedy_config=GreedyMethod())

Optimize the einsum code contraction order using the KaHyPar + Greedy approach.
This program first recursively cuts the tensors into several groups using KaHyPar,
with maximum group size specifed by `max_group_size` and maximum space complexity specified by `sc_target`,
Then finds the contraction order inside each group with the greedy search algorithm. Other arguments are

# Fields
- `sc_target` is the target space complexity, defined as `log2(number of elements in the largest tensor)`,
- `imbalances` is a KaHyPar parameter that controls the group sizes in hierarchical bipartition,
- `max_group_size` is the maximum size that allowed to used greedy search,
- `sub_optimizer` is the sub-optimizer used to find the contraction order when the group size is small enough.

# References
- [Hyper-optimized tensor network contraction](https://arxiv.org/abs/2002.01935)
- [Simulating the Sycamore quantum supremacy circuits](https://arxiv.org/abs/2103.03074)
"""
Base.@kwdef struct KaHyParBipartite{RT,IT,SO} <: CodeOptimizer
    sc_target::RT
    imbalances::IT = 0.0:0.005:0.8
    max_group_size::Int = 40
    sub_optimizer::SO = GreedyMethod()
end

function induced_subhypergraph(s::SparseMatrixCSC, group)
    s0 = s[group,:]
    nvs = vec(sum(s0, dims=1))
    remaining_edges = findall(!iszero, nvs)
    s0[:,remaining_edges], remaining_edges
end

function convert2int(sizes::AbstractVector)
    round.(Int, sizes .* 100)
end

"""
    BipartiteResult{RT}
    BipartiteResult(part1, part2, sc, valid)

Result of the bipartite optimization. `part1` and `part2` are the two parts of the bipartition, `sc` is the space complexity of the bipartition, `valid` is a boolean indicating whether the bipartition is valid.
"""
struct BipartiteResult{RT}
    part1::Vector{Int}
    part2::Vector{Int}
    sc::RT
    valid::Bool
end

function bipartite_sc(bipartiter, adj::SparseMatrixCSC, vertices, log2_sizes)
    error("""Guess you are trying to use the `KaHyParBipartite` optimizer.
Then you need to add `using KaHyPar` first!""")
end

# the space complexity (external degree of freedoms) if we contract this group
function group_sc(adj, group, log2_sizes)
    degree_in = sum(adj[group,:], dims=1)
    degree_all = sum(adj, dims=1)
    sum(i->(degree_in[i]!=0 && degree_in[i]!=degree_all[i] ? Float64(log2_sizes[i]) : 0.0), 1:size(adj,2))
end

function bipartition_recursive(bipartiter, adj::SparseMatrixCSC, vertices::AbstractVector{T}, log2_sizes) where T
    if length(vertices) > bipartiter.max_group_size
        res = bipartite_sc(bipartiter, adj, vertices, log2_sizes)
        !res.valid && @warn "fail to find a valid partition for `sc_target = $(bipartiter.sc_target)`, got minimum value `$(res.sc)` (imbalances = $(bipartiter.imbalances))"
        groups = Vector{T}[]
        for part in (res.part1, res.part2)
            for component in _connected_components(adj, part)
                push!(groups, component)
            end
        end
        newparts = [bipartition_recursive(bipartiter, adj, groups[i], log2_sizes) for i=1:length(groups)]
        if length(groups) > 2
            # number of groups is small (2 or 3), thus I think it is not necessary to use more complex methods.
            tree = coarse_grained_optimize(adj, groups, log2_sizes, GreedyMethod())
            return map_tree_to_parts(tree, newparts)
        else
            return newparts
        end
    else
        return [vertices]
    end
end

# parts are vectors of Ts
function coarse_grained_optimize(adj, parts, log2_sizes, sub_optimizer)
    incidence_list = get_coarse_grained_graph(adj, parts)
    log2_edge_sizes = Dict([i=>log2_sizes[i] for i=1:length(log2_sizes)])
    tree, _, _ = tree_greedy(incidence_list, log2_edge_sizes; α = sub_optimizer.α, temperature = sub_optimizer.temperature)
    return tree
end

function get_coarse_grained_graph(adj, parts)
    ADJ = vcat([sum(adj[part,:], dims=1) for part in parts]...)
    degree_in = sum(ADJ, dims=1)
    degree_all = sum(adj, dims=1)
    openedges = filter(i->degree_in[i]!=0 && degree_in[i]!=degree_all[i], 1:size(adj, 2))
    v2e = Dict{Int,Vector{Int}}()
    for v=1:size(ADJ, 1)
        v2e[v] = findall(!iszero, view(ADJ,v,:))
    end
    incidence_list = IncidenceList(v2e; openedges=openedges)
    return incidence_list
end

function map_tree_to_parts(tree, parts)
    if tree isa ContractionTree
        [map_tree_to_parts(tree.left, parts), map_tree_to_parts(tree.right, parts)]
    else
        parts[tree]
    end
end

# KaHyPar
function adjacency_matrix(ixs::AbstractVector)
    rows = Int[]
    cols = Int[]
    edges = unique!([Iterators.flatten(ixs)...])
    for (i,ix) in enumerate(ixs)
        push!(rows, map(x->i, ix)...)
        push!(cols, map(x->findfirst(==(x), edges), ix)...)
    end
    return sparse(rows, cols, ones(Int, length(rows)), length(ixs), length(edges)), edges
end

function recursive_bipartite_optimize(bipartiter, code::AbstractEinsum, size_dict)
    ixs, iy = getixsv(code), getiyv(code)
    ixv = [ixs..., iy]
    
    adj, edges = adjacency_matrix(ixv)
    vertices = collect(1:length(ixv))
    parts = bipartition_recursive(bipartiter, adj, vertices, [log2(size_dict[e]) for e in edges])
    optcode = recursive_construct_nestedeinsum(ixv, empty(iy), parts, size_dict, 0, bipartiter.sub_optimizer)
    return fix_binary_tree(pivot_tree(optcode, length(ixs) + 1))
end

maplocs(ne::NestedEinsum{ET}, parts) where ET = isleaf(ne) ? NestedEinsum{ET}(parts[ne.tensorindex]) : NestedEinsum(maplocs.(ne.args, Ref(parts)), ne.eins)

function kahypar_recursive(ne::NestedEinsum; log2_size_dict, sc_target, min_size, imbalances=0.0:0.04:0.8)
    if length(ne.args >= min_size) && all(isleaf, ne.args)
        bipartite_eincode(adj, ne.args, ne.eins; log2_size_dict=log2_size_dict, sc_target=sc_target, min_size=min_size, imbalances=imbalances)
    end
    kahypar_recursive(ne.args; log2_size_dict, sc_target=sc_target, min_size=min_size, imbalances=imbalances)
end

recursive_flatten(obj::Tuple) = vcat(recursive_flatten.(obj)...)
recursive_flatten(obj::AbstractVector) = vcat(recursive_flatten.(obj)...)
recursive_flatten(obj) = obj

function recursive_construct_nestedeinsum(ixs::AbstractVector{<:AbstractVector}, iy::AbstractVector{L}, parts::AbstractVector, size_dict, level, sub_optimizer) where L
    if length(parts) == 2
        # code is a nested einsum
        code1 = recursive_construct_nestedeinsum(ixs, iy, parts[1], size_dict, level+1, sub_optimizer)
        code2 = recursive_construct_nestedeinsum(ixs, iy, parts[2], size_dict, level+1, sub_optimizer)
        AB = recursive_flatten(parts[2]) ∪ recursive_flatten(parts[1])
        inset12, outset12 = ixs[AB], ixs[setdiff(1:length(ixs), AB)]
        iy12 = Iterators.flatten(inset12) ∩  (Iterators.flatten(outset12) ∪ iy)
        iy1, iy2 = getiyv(code1.eins), getiyv(code2.eins)
        return NestedEinsum([code1, code2], EinCode([iy1, iy2], L[(level==0 ? iy : iy12)...]))
    elseif length(parts) == 1
        return recursive_construct_nestedeinsum(ixs, iy, parts[1], size_dict, level, sub_optimizer)
    else
        error("not a bipartition, got size $(length(parts))")
    end
end

function recursive_construct_nestedeinsum(ixs::AbstractVector{<:AbstractVector}, iy::AbstractVector{L}, parts::AbstractVector{<:Integer}, size_dict, level, sub_optimizer) where L
    if isempty(parts)
        error("got empty group!")
    end
    inset, outset = ixs[parts], ixs[setdiff(1:length(ixs), parts)]
    iy1 = level == 0 ? iy : Iterators.flatten(inset) ∩  (Iterators.flatten(outset) ∪ iy)
    res = optimize_code(EinCode(inset, iy1), size_dict, sub_optimizer)
    if res isa SlicedEinsum
        @assert length(res.slicing) == 0
        res = res.eins
    end
    return maplocs(res, parts)
end
