################### The data types in OMEinsum ###################
"""
    AbstractEinsum

Abstract type for einsum notations.

# Required Interfaces
- `getixsv`: a vector of vectors, each vector represents the labels associated with a input tensor.
- `getiyv`: a vector of labels associated with the output tensor.
- `uniquelabels`: a vector of labels that are unique in the einsum notation.

# Derived interfaces
- `labeltype`: the data type to represent the labels in the einsum notation.
"""
abstract type AbstractEinsum end

"""
    getixsv(code::AbstractEinsum) -> Vector{Vector{LT}}

Returns the input indices of the einsum notation.
Each vector represents the labels associated with a input tensor.
"""
function getixsv end

"""
    getiyv(code::AbstractEinsum) -> Vector{LT}

Returns the output index of the einsum notation.
"""
function getiyv end

"""
    uniquelabels(code::AbstractEinsum) -> Vector{LT}

Returns the unique labels in the einsum notation. The labels are the indices of the tensors.
"""
uniquelabels(code::AbstractEinsum) = unique!(vcat(getixsv(code)..., getiyv(code)))

"""
    labeltype(code::AbstractEinsum) -> Type

Returns the data type to represent the labels in the einsum notation.
"""
labeltype(code::AbstractEinsum) = eltype(getiyv(code))

"""
    EinCode{LT} <: AbstractEinsum
    EinCode(ixs::Vector{Vector{LT}}, iy::Vector{LT})

Einsum code with input indices `ixs` and output index `iy`.

# Examples
The einsum notation for matrix multiplication is:
```jldoctest
julia> code = OMEinsumContractionOrders.EinCode([[1,2], [2, 3]], [1, 3])
1∘2, 2∘3 -> 1∘3

julia> OMEinsumContractionOrders.getixsv(code)
2-element Vector{Vector{Int64}}:
 [1, 2]
 [2, 3]

julia> OMEinsumContractionOrders.getiyv(code)
2-element Vector{Int64}:
 1
 3
```
"""
struct EinCode{LT} <: AbstractEinsum
    ixs::Vector{Vector{LT}}
    iy::Vector{LT}
end
getixsv(rc::EinCode) = rc.ixs
getiyv(rc::EinCode) = rc.iy
Base.:(==)(a::EinCode, b::EinCode) = a.ixs == b.ixs && a.iy == b.iy

"""
    CodeOptimizer

Abstract type for code optimizers.
"""
abstract type CodeOptimizer end

"""
    CodeSlicer

Abstract type for code slicers.
"""
abstract type CodeSlicer end

"""
    AbstractDecompositionType

Abstract type for decomposition types, which includes [`TreeDecomp`](@ref) and [`PathDecomp`](@ref).
"""
abstract type AbstractDecompositionType end

"""
    TreeDecomp <: AbstractDecompositionType

Tree decomposition type.
"""
struct TreeDecomp <: AbstractDecompositionType end

"""
    PathDecomp <: AbstractDecompositionType

Path decomposition type.
"""
struct PathDecomp <: AbstractDecompositionType end

"""
    ScoreFunction

A function to compute the score of a contraction code:

```
score = tc_weight * 2^tc + rw_weight * 2^rw + sc_weight * max(0, 2^sc - 2^sc_target)
```

# Fields
- `tc_weight`: the weight of the time complexity, default is 1.0.
- `sc_weight`: the weight of the space complexity (the size of the largest tensor), default is 1.0.
- `rw_weight`: the weight of the read-write complexity, default is 0.0.
- `sc_target`: the target space complexity, below which the `sc_weight` will be set to 0 automatically, default is 0.0.
"""
struct ScoreFunction
    tc_weight::Float64
    sc_weight::Float64
    rw_weight::Float64
    sc_target::Float64
    function ScoreFunction(; tc_weight=1.0, sc_weight=1.0, rw_weight=0.0, sc_target=20.0)
        @assert tc_weight >= 0
        @assert sc_weight >= 0
        @assert rw_weight >= 0
        @assert sc_target >= 0
        new(Float64(tc_weight), Float64(sc_weight), Float64(rw_weight), Float64(sc_target))
    end
end

function (score::ScoreFunction)(tc, sc, rw)
    return score.tc_weight * exp2(tc) + score.rw_weight * exp2(rw) + score.sc_weight * max(0.0, exp2(sc) - exp2(score.sc_target))
end

"""
    NestedEinsum{LT} <: AbstractEinsum
    NestedEinsum(args::Vector{NestedEinsum}, eins::EinCode)

The einsum notation with a contraction order specified as a tree data structure.
It is automatically generated by the contraction code optimizer with the [`optimize_code`](@ref) function.

# Fields
- `args`: the children of the current node
- `tensorindex`: the index of the input tensor, required only for leaf nodes. For non-leaf nodes, it is `-1`.
- `eins`: the einsum notation for the operation at the current node.
"""
struct NestedEinsum{LT} <: AbstractEinsum
    args::Vector{NestedEinsum{LT}}
    tensorindex::Int  # -1 if not leaf
    eins::EinCode{LT}
    NestedEinsum(args::Vector{NestedEinsum{LT}}, eins::EinCode) where LT = new{LT}(args, -1, eins)
    NestedEinsum{LT}(arg::Int) where LT = new{LT}(NestedEinsum{LT}[], arg)
end
function Base.:(==)(a::NestedEinsum, b::NestedEinsum)
    return a.args == b.args && a.tensorindex == b.tensorindex && if isdefined(a, :eins)
        isdefined(b, :eins) && a.eins == b.eins
    else
        !isdefined(b, :eins)
    end
end
isleaf(ne::NestedEinsum) = ne.tensorindex != -1
function getixsv(ne::NestedEinsum{LT}) where LT
    d = collect_ixs!(ne, Dict{Int,Vector{LT}}())
    ks = sort!(collect(keys(d)))
    return @inbounds [d[i] for i in ks]
end
function collect_ixs!(ne::NestedEinsum, d::Dict{Int,Vector{LT}}) where LT
    @inbounds for i=1:length(ne.args)
        arg = ne.args[i]
        if isleaf(arg)
            d[arg.tensorindex] = getixsv(ne.eins)[i]
        else
            collect_ixs!(arg, d)
        end
    end
    return d
end
getiyv(ne::NestedEinsum) = getiyv(ne.eins)
function is_binary(code::NestedEinsum)
    isleaf(code) && return true
    length(code.args) == 2 && all(is_binary, code.args)
end
# a path decomposition is represented as a contraction tree like
# a,b,c,d,e -> (((a,b),c),d),e
is_path_decomposition(ne::NestedEinsum) = isleaf(ne) || length(ne.args) <= 2 && is_path_decomposition(ne.args[1]) && (length(ne.args) == 1 || isleaf(ne.args[2]))

"""
    SlicedEinsum{LT,ET<:Union{EinCode{LT},NestedEinsum{LT}}} <: AbstractEinsum
    SlicedEinsum(slicing::Vector{LT}, eins::ET)

The einsum notation with sliced indices. The sliced indices are the indices enumerated manually at the top level.
By slicing the indices, the space complexity of the einsum notation can be reduced.

# Fields
- `slicing`: the sliced indices.
- `eins`: the einsum notation of the current node, which is a [`NestedEinsum`](@ref) object.
"""
struct SlicedEinsum{LT,ET<:Union{EinCode{LT},NestedEinsum{LT}}} <: AbstractEinsum
    slicing::Vector{LT}
    eins::ET
end
Base.:(==)(a::SlicedEinsum, b::SlicedEinsum) = a.slicing == b.slicing && a.eins == b.eins
getixsv(se::SlicedEinsum) = getixsv(se.eins)
getiyv(se::SlicedEinsum) = getiyv(se.eins)
is_binary(code::SlicedEinsum) = is_binary(code.eins)
is_path_decomposition(ne::SlicedEinsum) = is_path_decomposition(ne.eins)

# Better printing
struct LeafString
    str::String
end
function AbstractTrees.children(ne::NestedEinsum)
    [isleaf(item) ? LeafString(_join(getixsv(ne.eins)[k])) : item for (k,item) in enumerate(ne.args)]
end
function AbstractTrees.printnode(io::IO, x::NestedEinsum)
    isleaf(x) ? print(io, x.tensorindex) : print(io, x.eins)
end
AbstractTrees.printnode(io::IO, e::LeafString) = print(io, e.str)
function Base.show(io::IO, e::EinCode)
    s = join([_join(ix) for ix in getixsv(e)], ", ") * " -> " * _join(getiyv(e))
    print(io, s)
end
function Base.show(io::IO, e::NestedEinsum)
    print_tree(io, e)
end
Base.show(io::IO, ::MIME"text/plain", e::NestedEinsum) = show(io, e)
Base.show(io::IO, ::MIME"text/plain", e::EinCode) = show(io, e)
_join(ix) = isempty(ix) ? "" : join(ix, connector(eltype(ix)))
connector(::Type{Char}) = ""
connector(::Type{Int}) = "∘"
connector(::Type) = "-"

function is_unary_or_binary(code::NestedEinsum)
    if isleaf(code) return true end
    if length(code.args) > 2 return false end
    return all(is_unary_or_binary, code.args)
end


# reformulate the nested einsum, removing a given tensor without change the space complexity
# consider only binary contraction tree with no openedges
function pivot_tree(code::NestedEinsum{LT}, removed_tensor_id::Int) where LT
    @assert is_unary_or_binary(code) "The contraction tree is not binary"
    @assert isempty(getiyv(code)) "The contraction tree has open edges"

    path = path_to_tensor(code, removed_tensor_id)
    isempty(path) && return code   # the tensor is at the root?

    right = popfirst!(path)
    left = right == 1 ? 2 : 1

    if isleaf(code.args[left]) && isleaf(code.args[right])
        ixsv = getixsv(code.eins)
        return NestedEinsum([code.args[left]], EinCode([ixsv[left]], ixsv[right]))
    elseif isleaf(code.args[right])
        return NestedEinsum([code.args[left].args...], EinCode(getixsv(code.args[left].eins), getixsv(code.eins)[right]))
    else
        # update the ein code to make sure the root of the left part and the right part are the same
        left_code = code.args[left]
        right_code = NestedEinsum([code.args[right].args...], EinCode(getixsv(code.args[right].eins), getixsv(code.eins)[left]))
    end
    tree = _pivot_tree!(left_code, right_code, path)

    return tree
end


function _pivot_tree!(left_code::NestedEinsum{LT}, right_code::NestedEinsum{LT}, path::Vector{Int}) where{LT}
    if !isleaf(right_code)
        right = popfirst!(path)
        left = right == 1 ? 2 : 1
        if length(right_code.args) == 1
            # orign: left: a, right: b -> a
            # reformulated: left: a -> b, right: b
            new_eins = EinCode([getiyv(right_code.eins)], getixsv(right_code.eins)[1])
            left_code = NestedEinsum([left_code], new_eins)
            left_code = _pivot_tree!(left_code, right_code.args[1], path)
        elseif length(right_code.args) == 2
            # origin: left: a, right: b, c -> a
            # reformulated: left: a, b -> c, right: c
            new_eins = EinCode([getiyv(right_code.eins), getixsv(right_code.eins)[left]], getixsv(right_code.eins)[right])
            left_code = NestedEinsum([left_code, right_code.args[left]], new_eins)
            left_code = _pivot_tree!(left_code, right_code.args[right], path)
        else
            error("The contraction tree is not binary")
        end
    end
    return left_code
end

# find the path to a given tensor in a nested einsum
function path_to_tensor(code::NestedEinsum, index::Int)
    path = Vector{Int}()
    _find_root!(code, index, path)
    return path
end

function _find_root!(code::NestedEinsum, index::Int, path::Vector{Int})
    if isleaf(code) return code.tensorindex == index end

    for (i, arg) in enumerate(code.args)
        if _find_root!(arg, index, path)
            pushfirst!(path, i)
            return true
        end
    end
    return false
end

# convert a tree with 1-2 children to a binary tree
function fix_binary_tree(code::NestedEinsum)
    isleaf(code) && return code
    args, ixs, iy = code.args, code.eins.ixs, code.eins.iy
    if length(args) == 1  # fix unary tree
        isleaf(args[1]) && return code
        subtree = fix_binary_tree(args[1])
        # Note: the resulting tree might still be unary, so the output is not guaranteed to be binary
        return NestedEinsum(subtree.args, EinCode(subtree.eins.ixs, iy))
    end
    @assert length(args) == 2 "The contraction tree can not be fixed, it has more than 2 children, got code = $code"
    newargs = empty(args)
    newixs = empty(ixs)
    for (ix, arg) in zip(ixs, args)
        arg = fix_binary_tree(arg)
        if length(arg.args) == 1  # remove the unary node
            push!(newargs, arg.args[1])
            push!(newixs, arg.eins.ixs[1])
        else
            push!(newargs, arg)
            push!(newixs, ix)
        end
    end
    return NestedEinsum(newargs, EinCode(newixs, iy))
end