using StaticArrays, LinearAlgebra, BenchmarkTools

# ３次元配列を行列の配列に変換する

"Just specify that [an element of] the argument is an array of matrices"
struct MatrixArray end

"""
    R = matrixarray(S::AbstractArray{T,3})

`R[k][i, j] = S[i, k, j]`
"""
matrixarray(S) = [SMatrix{size(S[:, i, :])...}(S[:, i, :]) for i in axes(S, 2)]
#matrixarray(S) = [S[:, i, :] for i in axes(S, 2)] # without StaticArrays

# 要素ごとに行列積を行う（元のコードを整理した）

reconst1(G) = reconst1(MatrixArray(), matrixarray.(G))

function reconst1(::MatrixArray, G)
    [tr(prod(k -> G[k][CI[k]], eachindex(G))) 
        for CI in CartesianIndices(Tuple(length.(G)))]
end

# 行列積の途中結果を再利用する

reconst2(G) = reconst2(MatrixArray(), matrixarray.(G))
reconst2(::MatrixArray, G) = tr.(reduce(eachprod, G))

"""
    C = eachprod(A, B)

`C[i, j, k] = A[i, j] * B[k]`

`A, B, C :: Array{Matrix}`
"""
eachprod(A, B) = A .* expanddim(B, A)

"""
    Bx = expanddim(B, A)

`Bx = reshape(B, (1, 1, 1, m, n))` where `ndims(A) == 3`, `size(B) == (m, n)`
"""
expanddim(B, A) = reshape(B, (ntuple(_ -> 1, ndims(A))..., size(B)...))

# 行列積の最終段を省き、トレースを直接求める

reconst3(G) = reconst3(MatrixArray(), matrixarray.(G))
reconst3(::MatrixArray, G) = _reconst3(G...)

# Gs を真ん中で分けないと遅くなる
function _reconst3(Gs...)
    h = length(Gs) ÷ 2
    _reconst3(reduce(eachprod, Gs[begin:h]), reduce(eachprod, Gs[h+1:end]))
end

#_reconst3(G1, G2, Gs...) = _reconst3(eachprod(G1, G2), Gs...) # 遅い
_reconst3(G1, G2) = trprod.(G1, expanddim(G2, G1))

"""
    trprod(A, B)

Returns `tr(A * B)`
"""
trprod(A, B) = dot(vec(A'), vec(B))

function genG(I, r)
    r2 = circshift(r, -1)
    randn.(r, I, r2)
end

function mybench(K::Int = 1; I = [3, 4, 5, 6], r = [2, 2, 3, 4])
    I, r = K*I, K*r
    @show I, r
    G = genG(I, r)
    @assert reconst1(G) ≈ reconst2(G) ≈ reconst3(G)
    
    @btime reconst1($G)
    @btime reconst2($G)
    @btime reconst3($G)
    return
end

mybench(1)

mybench(2)

mybench(5)