text:
module DoubleDicts

import MathOptInterface

const MOI = MathOptInterface

abstract type AbstractDoubleDict{V} <: AbstractDict{MOI.ConstraintIndex,V} end

abstract type AbstractDoubleDictInner{F,S,V} <:
              AbstractDict{MOI.ConstraintIndex{F,S},V} end

"""
    typed_value(dict::AbstractDoubleDictInner{F,S,V}, value) where {F,S,V}

Convert the `value` stored inside `dict` to the equivalent on the outer
`DoubleDict`. This is useful when the value type `V` of the inner dict is
different to the outer dict. (See, e.g., [`IndexDoubleDict`](@ref).)
"""
typed_value(::AbstractDoubleDictInner, value) = value

"""
    DoubleDict{V}

An optimized dictionary to map `MOI.ConstraintIndex` to values of type `V`.

Works as a `AbstractDict{MOI.ConstraintIndex,V}` with minimal differences.

If `V` is also a `MOI.ConstraintIndex`, use [`IndexDoubleDict`](@ref).

Note that `MOI.ConstraintIndex` is not a concrete type, opposed to
`MOI.ConstraintIndex{MOI.VariableIndex, MOI.Integers}`, which is a concrete
type.

When looping through multiple keys of the same Function-in-Set type, use
```julia
inner = dict[F, S]
```
to return a type-stable [`DoubleDictInner`](@ref).
"""
struct DoubleDict{V} <: AbstractDoubleDict{V}
    dict::Dict{Tuple{Type,Type},Dict{Int64,V}}
    function DoubleDict{V}() where {V}
        return new{V}(Dict{Tuple{Type,Type},Dict{Int64,V}}())
    end
end

"""
    DoubleDictInner{F,S,V}

A type stable inner dictionary of [`DoubleDict`](@ref).
"""
mutable struct DoubleDictInner{F,S,V} <: AbstractDoubleDictInner{F,S,V}
    dict::Dict{Int64,V}
    function DoubleDictInner{F,S}(d::DoubleDict{V}) where {F,S,V}
        return new{F,S,V}(get!(d.dict, (F, S), Dict{Int64,V}()))
    end
end

function Base.getindex(d::DoubleDict, ::Type{F}, ::Type{S}) where {F,S}
    return DoubleDictInner{F,S}(d)
end

"""
    IndexDoubleDict

A specialized version of [`DoubleDict`] in which the values are of type
`MOI.ConstraintIndex`

When looping through multiple keys of the same Function-in-Set type, use
```julia
inner = dict[F, S]
```
to return a type-stable [`IndexDoubleDictInner`](@ref).
"""
struct IndexDoubleDict <: AbstractDoubleDict{MOI.ConstraintIndex}
    dict::Dict{Tuple{Type,Type},Dict{Int64,Int64}}
    function IndexDoubleDict()
        return new(Dict{Tuple{Type,Type},Dict{Int64,Int64}}())
    end
end

"""
    IndexDoubleDictInner{F,S}

A type stable inner dictionary of [`IndexDoubleDict`](@ref).
"""
mutable struct IndexDoubleDictInner{F,S} <:
               AbstractDoubleDictInner{F,S,MOI.ConstraintIndex{F,S}}
    dict::Dict{Int64,Int64}
    function IndexDoubleDictInner{F,S}(d::IndexDoubleDict) where {F,S}
        return new{F,S}(get!(d.dict, (F, S), Dict{Int64,Int64}()))
    end
end

function Base.getindex(d::IndexDoubleDict, ::Type{F}, ::Type{S}) where {F,S}
    return IndexDoubleDictInner{F,S}(d)
end

function typed_value(::IndexDoubleDictInner{F,S}, v::Int64) where {F,S}
    return MOI.ConstraintIndex{F,S}(v)
end

# Base.sizehint!

function Base.sizehint!(::AbstractDoubleDict, ::Integer)
    return throw(
        ErrorException(
            "sizehint!(d::DoubleDict, ::Integer) has no proper" *
            " meaning for DoubleDict, use sizehint!(d[F,S], n::Integer) " *
            "instead.",
        ),
    )
end

function Base.sizehint!(d::AbstractDoubleDictInner, n::Integer)
    return sizehint!(d.dict, n)
end

# Base.length

function Base.length(d::AbstractDoubleDict)
    len = 0
    for inner in values(d.dict)
        len += length(inner)
    end
    return len
end

Base.length(d::AbstractDoubleDictInner) = length(d.dict)

# Base.haskey

function Base.haskey(
    d::AbstractDoubleDict,
    key::MOI.ConstraintIndex{F,S},
) where {F,S}
    return haskey(d[F, S], key)
end

function Base.haskey(
    d::AbstractDoubleDictInner{F,S},
    key::MOI.ConstraintIndex{F,S},
) where {F,S}
    return haskey(d.dict, key.value)
end

# Base.get

function Base.get(
    d::AbstractDoubleDict,
    key::MOI.ConstraintIndex{F,S},
    default,
) where {F,S}
    inner = d[F, S]
    return get(inner, key, default)
end

function Base.get(
    d::AbstractDoubleDictInner,
    key::MOI.ConstraintIndex{F,S},
    default,
) where {F,S}
    return typed_value(d, get(d.dict, key.value, default))
end

# Base.getindex

function Base.getindex(
    d::AbstractDoubleDict,
    key::MOI.ConstraintIndex{F,S},
) where {F,S}
    inner = d[F, S]
    return inner[key]
end

function Base.getindex(
    d::AbstractDoubleDictInner{F,S},
    key::MOI.ConstraintIndex{F,S},
) where {F,S}
    x = get(d.dict, key.value) do
        return throw(KeyError(key))
    end
    return typed_value(d, x)
end

# Base.setindex!

function Base.setindex!(
    d::AbstractDoubleDict{V},
    value::V,
    key::MOI.ConstraintIndex{F,S},
) where {F,S,V}
    inner = d[F, S]
    inner[key] = value
    return value
end

function Base.setindex!(
    d::AbstractDoubleDictInner{F,S,V},
    value::V,
    key::MOI.ConstraintIndex{F,S},
) where {F,S,V}
    d.dict[key.value] = value
    return value
end

function Base.setindex!(
    d::IndexDoubleDict,
    value::MOI.ConstraintIndex{F,S},
    key::MOI.ConstraintIndex{F,S},
) where {F,S}
    inner = d[F, S]
    inner[key] = value
    return value
end

function Base.setindex!(
    d::IndexDoubleDictInner{F,S},
    value::MOI.ConstraintIndex{F,S},
    key::MOI.ConstraintIndex{F,S},
) where {F,S}
    d.dict[key.value] = value.value
    return value
end

# Base.empty!

function Base.empty!(d::AbstractDoubleDict)
    Base.empty!(d.dict)
    return d
end

function Base.empty!(d::AbstractDoubleDictInner)
    empty!(d.dict)
    return d
end

# Base.delete!

function Base.delete!(
    d::AbstractDoubleDict,
    key::MOI.ConstraintIndex{F,S},
) where {F,S}
    delete!(d[F, S], key)
    return d
end

function Base.delete!(
    d::AbstractDoubleDictInner{F,S},
    key::MOI.ConstraintIndex{F,S},
) where {F,S}
    delete!(d.dict, key.value)
    return d
end

# Base.isempty

function Base.isempty(d::AbstractDoubleDict)
    return isempty(d.dict) || all(isempty, values(d.dict))
end

Base.isempty(d::AbstractDoubleDictInner) = isempty(d.dict)

# Base.values

function Base.values(d::AbstractDoubleDict{V}) where {V}
    out = V[]
    for (F, S) in keys(d.dict)
        append!(out, values(d[F, S]))
    end
    return out
end

Base.values(d::AbstractDoubleDictInner) = typed_value.(Ref(d), values(d.dict))

# Base.keys

function Base.keys(d::AbstractDoubleDict)
    out = MOI.ConstraintIndex[]
    for (F, S) in keys(d.dict)
        append!(out, keys(d[F, S]))
    end
    return out
end

function Base.keys(d::AbstractDoubleDictInner{F,S}) where {F,S}
    return MOI.ConstraintIndex{F,S}.(keys(d.dict))
end

# Base.iterate

function Base.iterate(d::AbstractDoubleDict)
    outer_next = iterate(d.dict)
    if outer_next === nothing
        return  # There are no keys.
    end
    # The result is a (F,S)=>inner pair.
    ((F, S), inner), outer_state = outer_next
    inner_next = iterate(inner)
    while inner_next === nothing
        # It may be that the inner dictionary is empty! If so, we should go to
        # the next element in the outer dictionary.
        outer_next = iterate(d.dict, outer_state)
        if outer_next === nothing
            return
        end
        ((F, S), inner), outer_state = outer_next
        # Start iterating from scratch on this new `inner` object.
        inner_next = iterate(inner)
    end
    (k, v), inner_state = inner_next
    result = MOI.ConstraintIndex{F,S}(k) => typed_value(d[F, S], v)
    return result, (inner_state, outer_next)
end

function Base.iterate(d::AbstractDoubleDictInner{F,S}) where {F,S}
    next = iterate(d.dict)
    if next === nothing
        return
    end
    (k, v), inner_state = next
    result = MOI.ConstraintIndex{F,S}(k) => typed_value(d, v)
    return result, (d.dict, inner_state)
end

function Base.iterate(d::AbstractDoubleDict, state)
    inner_state, outer_next = state
    ((F, S), inner), outer_state = outer_next
    inner_next = iterate(inner, inner_state)
    while inner_next === nothing
        # We may have reached the end of this inner dictionary. Get the next
        # element of the outer dictionary:
        outer_next = iterate(d.dict, outer_state)
        if outer_next === nothing
            return
        end
        ((F, S), inner), outer_state = outer_next
        # Start iterating from scratch on this new `inner` object.
        inner_next = iterate(inner)
    end
    (k, v), inner_state = inner_next
    result = MOI.ConstraintIndex{F,S}(k) => typed_value(d[F, S], v)
    return result, (inner_state, outer_next)
end

function Base.iterate(d::AbstractDoubleDictInner{F,S}, state) where {F,S}
    inner, inner_state = state
    next = iterate(inner, inner_state)
    if next === nothing
        return
    end
    ((k, v), next_inner_state) = next
    result = MOI.ConstraintIndex{F,S}(k) => typed_value(d, v)
    return result, (inner, next_inner_state)
end

end


text:
# This file is a part of BAT.jl, licensed under the MIT License (MIT).

const standard_confidence_vals = [0.68, 0.95, 0.997]

function rectangle_path(lo::Vector{<:Real}, hi::Vector{<:Real})
    [
        lo[1] lo[2];
        hi[1] lo[2];
        hi[1] hi[2];
        lo[1] hi[2];
        lo[1] lo[2];
    ]
end

function err_ellipsis_path(μ::Vector{<:Real}, Σ::Matrix{<:Real}, confidence::Real = 0.68, npts = 256)
    σ_sqr, A = eig(Hermitian(Σ))
    σ = sqrt.(σ_sqr)
    ϕ = linspace(0, 2π, 100)
    σ_scaled = σ .* sqrt(invlogcdf(Chisq(2), log(confidence)))
    xy = hcat(σ_scaled[1] * cos.(ϕ), σ_scaled[2] * sin.(ϕ)) * [A[1,1] A[1,2]; A[2,1] A[2,2]]
    xy .+ μ'
end

@recipe function f(samples::DensitySampleVector, parsel::NTuple{2,Integer})
    pi_x, pi_y = parsel

    acc = find(x -> x > 0, samples.weight)
    rej = find(x -> x <= 0, samples.weight)

    base_markersize = get(d, :markersize, 1.5)
    seriestype = get(d, :seriestype, :scatter)

    plot_bounds = get(d, :bounds, true)
    delete!(d, :bounds)

    if seriestype == :scatter
        color = parse(RGBA{Float64}, get(d, :seriescolor, :green))
        label = get(d, :label, isempty(rej) ? "samples" : "accepted")

        @series begin
            seriestype := :scatter
            label := label
            markersize := [w < 1 ? base_markersize : base_markersize * sqrt(w) for w in samples.weight[acc]]
            markerstrokewidth := 0
            color := [w >= 1 ? color : RGBA(convert(RGB, color), color.alpha * w) for w in samples.weight[acc]]
            xlabel --> "\$\\theta_$(pi_x)\$"
            ylabel --> "\$\\theta_$(pi_y)\$"
            (samples.params[pi_x, acc], samples.params[pi_y, acc])
        end

        if !isempty(rej)
            @series begin
                seriestype := :scatter
                label := "rejected"
                markersize := base_markersize
                markerstrokewidth := 0
                color := :red
                (samples.params[pi_x, rej], samples.params[pi_y, rej])
            end
        end
    elseif seriestype == :histogram2d
        @series begin
            seriestype := :histogram2d
            label --> "samples"
            xlabel --> "\$\\theta_$(pi_x)\$"
            ylabel --> "\$\\theta_$(pi_y)\$"
            weights := samples.weight[:]
            (samples.params[pi_x, :], samples.params[pi_y, :])
        end
    else
        error("seriestype $seriestype not supported")
    end

    nothing
end

@recipe function f(bounds::HyperRectBounds, parsel::NTuple{2,Integer})
    pi_x, pi_y = parsel

    vol = spatialvolume(bounds)
    vhi = vol.hi[[pi_x, pi_y]]; vlo = vol.lo[[pi_x, pi_y]]
    rect_xy = rectangle_path(vlo, vhi)
    bext = 0.1 * (vhi - vlo)
    xlims = (vlo[1] - bext[1], vhi[1] + bext[1])
    ylims = (vlo[2] - bext[2], vhi[2] + bext[2])

    @series begin
        seriestype := :path
        label --> "bounds"
        color --> :darkred
        alpha --> 0.75
        linewidth --> 2
        xlims --> xlims
        ylims --> ylims
        (rect_xy[:,1], rect_xy[:,2])
    end

    nothing
end

@recipe function f(stats::MCMCBasicStats, parsel::NTuple{2,Integer})
    pi_x, pi_y = parsel

    Σ_all = stats.param_stats.cov
    Σ = [Σ_all[pi_x, pi_x] Σ_all[pi_x, pi_y]; Σ_all[pi_y, pi_x] Σ_all[pi_y, pi_y]]

    μ = stats.param_stats.mean[[pi_x, pi_y]]
    mode_xy = stats.mode[[pi_x, pi_y]]

    conf = standard_confidence_vals

    color --> :darkviolet
    alpha --> 0.75
    linewidth --> 2

    for i in eachindex(conf)
        xy = err_ellipsis_path(μ, Σ, conf[i])
        @series begin
            seriestype := :path
            label --> "$(100 *conf[i])%"
            (xy[:,1], xy[:,2])
        end
    end

    markercolor --> :white
    markeralpha --> 0
    markersize --> 7
    markerstrokecolor --> :black
    markerstrokealpha --> 1
    markerstrokewidth --> 2

    @series begin
        seriestype := :scatter
        label := "mean"
        markershape := :circle
        ([μ[1]], [μ[2]])
    end

    @series begin
        seriestype := :scatter
        label := "mode"
        markershape := :rect
        ([mode_xy[1]], [mode_xy[2]])
    end

    vlo = stats.param_stats.minimum[[pi_x, pi_y]]
    vhi = stats.param_stats.maximum[[pi_x, pi_y]]
    rect_xy = rectangle_path(vlo, vhi)

    @series begin
        seriestype := :path
        label --> "bbox"
        (rect_xy[:,1], rect_xy[:,2])
    end

    nothing
end

ext: jl


text:
# getters, e.g. coords(code::TensorNetworkCode)
for data in [:coords, :node_types, :node_indices, :edge_types, :edge_indices]
    eval(quote
        @doc """
            $($data)(code::TensorNetworkCode)

        Return `$($data)` field of the code's [`CodeGraph`](@ref).
        """
        $data(code::TensorNetworkCode) = code.code_graph.$data
        $data(code_graph::CodeGraph) = code_graph.$data
    end)
end

# getters, e.g. coords(code::TensorNetworkCode, label::Int)
for data = (:coords, :node_types, :node_indices)
    eval(quote
        @doc """
            $($data)(code::TensorNetworkCode, label::Int)

        Return `$($data)[label]` for the node from the code's [`CodeGraph`](@ref).
        """
        $data(code::TensorNetworkCode, label::Int) = code.code_graph.$data[label]
        $data(code_graph::CodeGraph, label::Int) = code_graph.$data[label]
    end)
end

# getters, e.g. edge_types(code::TensorNetworkCode, label::Set{Int})
for data = (:edge_types, :edge_indices)
    eval(quote
        @doc """
            $($data)(code::TensorNetworkCode, label::Set{Int})

        Return `$($data)[label]` for the edge from the code's [`CodeGraph`](@ref).
        """
        $data(code::TensorNetworkCode, label::Set{Int}) = code.code_graph.$data[label]
        $data(code_graph::CodeGraph, label::Set{Int}) = code_graph.$data[label]
    end)
end

# NOTE: THESE SETTERS CANNOT WORK BECAUSE CodeGraph IS IMMUTABLE
# # setters, e.g. set_coords!(code::TensorNetworkCode, new_data)
# for data in [:coords, :node_types, :node_indices, :edge_types, :edge_indices]
#     function_name = Symbol(:set_,data,:!)
#     eval(quote
#         function $function_name(code::TensorNetworkCode, new_data)
#             code.code_graph.$data = new_data
#             return nothing
#         end
#         function $function_name(code_graph::CodeGraph, new_data)
#             code_graph.$data = new_data
#             return nothing
#         end
#     end)
# end

# setters, e.g. set_coords!(code::TensorNetworkCode, label:Int, new_data)
for data = (:coords, :node_types, :node_indices)
    function_name = Symbol(:set_,data,:!)
    eval(quote
        @doc """
            $($function_name)(code::TensorNetworkCode, label::Int, new_data)

        Set `$($data)[label]` for the node from the code's [`CodeGraph`](@ref).
        """
        function $function_name(code::TensorNetworkCode, label::Int, new_data)
            code.code_graph.$data[label] = new_data
            return nothing
        end
        function $function_name(code_graph::CodeGraph, label::Int, new_data)
            code_graph.$data[label] = new_data
            return nothing
        end
    end)
end

# setters, e.g. set_edge_types!(code::TensorNetworkCode, label:Set{Int}, new_data)
for data = (:edge_types, :edge_indices)
    function_name = Symbol(:set_,data,:!)
    eval(quote
        @doc """
            $($function_name)(code::TensorNetworkCode, label::Set{Int}, new_data)

        Set `$($data)[label]` for the edge from the code's [`CodeGraph`](@ref).
        """
        function $function_name(code::TensorNetworkCode, label::Set{Int}, new_data)
            code.code_graph.$data[label] = new_data
            return nothing
        end
        function $function_name(code_graph::CodeGraph, label::Set{Int}, new_data)
            code_graph.$data[label] = new_data
            return nothing
        end
    end)
end

function edges(code_graph::CodeGraph)
    output = collect(keys(code_graph.edge_types))
    return output
end
"""
    edges(code) -> Vector{Set{Int}}

Returns a list of edge labels for the code, corresponding to edges between nodes.
"""
edges(code::TensorNetworkCode) = edges(code.code_graph)

function nodes(code_graph::CodeGraph)
    output = collect(keys(coords(code_graph)))
    sort!(output)
    return output
end
"""
    nodes(code::TensorNetworkCode) -> Vector{Int}

Returns an ordered list of node labels for the code, corresponding to the qubits and virtual
nodes.
"""
nodes(code::TensorNetworkCode) = nodes(code.code_graph)

num_nodes(code_graph::CodeGraph) = length(code_graph.coords)
"""
    num_nodes(code::TensorNetworkCode) -> Int

Returns the number of nodes in the code, which is the sum of the number of qubits and the
number of virtual nodes.
"""
num_nodes(code::TensorNetworkCode) = num_nodes(code.code_graph)

"""
    physical_neighbours(code::TensorNetworkCode, node::Int) -> Set{Int}

Returns the node labels of physical qubits connected by an edge to the given node.
"""
function physical_neighbours(code,node)
    all_edges = edges(code)
    neighbour_edges = filter(x -> node in x, all_edges)
    neighbour_nodes = union(neighbour_edges...)
    return filter(x -> x > 0 && x != node, neighbour_nodes)
end

"""
    set_coords!(
        code::TensorNetworkCode,
        new_coords::AbstractVector{<:AbstractVector{<:Real}}
    )

Assign new coordinates to all (physical and virtual) nodes of the code, where the order of
the coordinates matches that of the nodes returned by [`nodes`](@ref), i.e. in ascending
order.
"""
function set_coords!(
        code::TensorNetworkCode,
        new_coords::AbstractVector{<:AbstractVector{<:Real}}
)
    num_nodes(code) == length(new_coords) || error("different number of nodes and coords!")
    for (i, node) in enumerate(nodes(code))
        set_coords!(code, node, new_coords[i])
    end
end

"""
    shift_coords!(code::TensorNetworkCode, shift::AbstractVector{<:Real})

Shift the coordinates of all (physical and virtual) nodes of the code by the given shift.
"""
function shift_coords!(code::TensorNetworkCode, shift::AbstractVector{<:Real})
    for node in nodes(code)
        new_coords = coords(code,node) + shift
        set_coords!(code, node, new_coords)
    end
end

"""
    new_indices(code::TensorNetworkCode) -> TensorNetworkCode

Deepcopies the code but changes all `ITensor` indices.
"""
function new_indices(code::TensorNetworkCode)
    new_code = deepcopy(code)

    for edge in edges(new_code)
        old_indices = edge_indices(new_code,edge)
        # Create new indices with same dims and tags
        new_indices = [Index(dim(old_index),tags(old_index)) for old_index in old_indices]
        for node in nodes(new_code)
            if node < 0
                new_node_indices = node_indices(new_code,node)
                for α in 1:length(old_indices)
                    β = findfirst(x->x==old_indices[α],new_node_indices)
                    if β == nothing
                        continue
                    end
                    new_node_indices[β] = new_indices[α]
                end
                set_node_indices!(new_code,node,new_node_indices)
            end
        end
        set_edge_indices!(new_code,edge,new_indices)
    end

    return new_code
end


text:
using Random

export randsubset

"""
    randsubset([rng=GLOBAL_RNG,] A, m::Integer)

Return random subset of `A` with `m` elements.
If `A` is an `AbstractSet` return a set of the same type.
Otherwise return a `Vector` with the element type of `A`.
If `A` is not unique, result is not unique.
"""
randsubset(A, m::Integer) = randsubset(Random.GLOBAL_RNG, A, m)
function randsubset(rng::AbstractRNG, A, m::Integer)
    randsubset(rng, collect(A), m)
end
function randsubset(rng::AbstractRNG, A::AbstractSet, m::Integer)
    typeof(A)(randsubset(rng, collect(A), m))
end
function randsubset(rng::AbstractRNG, A::AbstractVector, m::Integer)
    A[randsubset(rng, axes(A, 1), m)]
end
function randsubset(rng::AbstractRNG, A::AbstractUnitRange, m::Integer)
    n = length(A)
    0 <= m <= n || throw(ArgumentError("required 0 <= m <= $n but m = $m"))
    m == 0 && return eltype(A)[]
    m == n && return collect(A)
    if m<<1 > n
        deleteat!(collect(A), _randsubset(rng, 1:n, n-m))
    else
        _randsubset(rng, A, m)
    end
end

function _randsubset(rng::AbstractRNG, A::AbstractUnitRange, m::Integer)
    x = randsubseq(rng, A, m / length(A))
    k = length(x)
    while k < m
        x = unique!(sort!(vcat(x, randsubset(rng, A, m - k))))
        k = length(x)
    end
    while k > m
        y = unique!(sort!(rand(rng, 1:k, k-m)))
        deleteat!(x, y)
        k = length(x)
    end
    x
end


text:
using Einsum
using LinearAlgebra

@testset "simple" begin

    ## the very basics
    bc = rand(2,3)
    cd = rand(3,4)

    @mul A[b,d] := bc[b,c] * cd[c,d]
    @test A == bc * cd

    B = similar(A)
    @mul B[b,d] = bc[b,c] * cd[c,d] # in-place
    @test B == A

    ## and with vectors
    c = randn(3)
    @mul b[b] := bc[b,c] * c[c]
    @test b == bc * c

    @mul c[c] := b[b] * bc[b,c]
    @test c == vec(b' * bc)

    @mul z[] := b[b] * b[b]
    @test z[] == dot(b,b)

    b′ = similar(b)
    @mul b′[b] = bc[b,c] * c[c] # in-place
    @test b′ == bc * c

    c′ = similar(c)
    @mul c′[c] = b[b] * bc[b,c]

end
@testset "permutedims" begin

    bbb = rand(2,2,2);
    bbb2 = randn(2,2,2);

    @mul    A[i,j] := bbb[i,k,k2] * bbb2[j,k,k2]
    @reduce B[i,j] := sum(k,k2) bbb[i,k,k2] * bbb2[j,k,k2]
    @einsum C[i,j] := bbb[i,k,k2] * bbb2[j,k,k2]
    @test A ≈ B ≈ C

    D = similar(A);
    @mul    D[i,j] = bbb[i,k,k2] * bbb2[j,k,k2] # in-place
    @test A == D

    @mul    A[i,j] := bbb[k2,i,k] * bbb2[j,k,k2]
    @reduce B[i,j] := sum(k,k2) bbb[k2,i,k] * bbb2[j,k,k2]
    @einsum C[i,j] := bbb[k2,i,k] * bbb2[j,k,k2]
    @test A ≈ B ≈ C

    D = similar(A);
    @mul D[i,j] = bbb[k2,i,k] * bbb2[j,k,k2] # in-place
    @test A == D

    cccc = rand(3,3,3,3);
    cccc2 = randn(3,3,3,3);

    @mul    A[i,j] := cccc[x,y,j,z] * cccc2[z,i,y,x]
    @reduce B[i,j] := sum(x,y,z) cccc[x,y,j,z] * cccc2[z,i,y,x]
    @einsum C[i,j] := cccc[x,y,j,z] * cccc2[z,i,y,x]
    @test A ≈ B ≈ C

    D = similar(A);
    @mul    D[i,j] = cccc[x,y,j,z] * cccc2[z,i,y,x]  # in-place
    @test A ≈ D

    @mul    A[i,j,k,l] := cccc[i,x,y,k] * cccc2[l,j,y,x]
    @reduce B[i,j,k,l] := sum(x,y) cccc[i,x,y,k] * cccc2[l,j,y,x]
    @einsum C[i,j,k,l] := cccc[i,x,y,k] * cccc2[l,j,y,x]
    @test A ≈ B ≈ C

    D = similar(A);
    @mul    D[i,j,k,l] = cccc[i,x,y,k] * cccc2[l,j,y,x]   # in-place
    @test A ≈ D

    @mul    A[i,j,k,l,m,n] := cccc[x,k,i,m] * cccc2[l,x,j,n]
    @reduce B[i,j,k,l,m,n] := sum(x) cccc[x,k,i,m] * cccc2[l,x,j,n]
    @einsum C[i,j,k,l,m,n] := cccc[x,k,i,m] * cccc2[l,x,j,n]
    @test A ≈ B  ≈ C

    D = similar(A);
    @mul    D[i,j,k,l,m,n] = cccc[x,k,i,m] * cccc2[l,x,j,n]  # in-place
    @test A ≈ D

    ## repeat with non-alphabeitcal LHS!
    @mul    A[j,i,l,k,n,m] := cccc[x,k,i,m] * cccc2[l,x,j,n]
    @reduce B[j,i,l,k,n,m] := sum(x) cccc[x,k,i,m] * cccc2[l,x,j,n]
    @einsum C[j,i,l,k,n,m] := cccc[x,k,i,m] * cccc2[l,x,j,n]
    @test A ≈ B  ≈ C

    D = similar(A);
    @mul    D[j,i,l,k,n,m] = cccc[x,k,i,m] * cccc2[l,x,j,n]  # in-place
    @test A ≈ D

    ## Mason Potter's example:
    # W = rand(2,2,2,2); M = rand(2,2,2);
    W = rand(6,7,4,5); M = rand(7,2,3);

    @reduce N[σ, b\a, b′\a′] := sum(σ′) W[σ,σ′,b,b′] * M[σ′,a,a′]
    @mul N2[σ, b\a, b′\a′] := W[σ,σ′,b,b′] * M[σ′,a,a′]
    @test N ≈ N2

    @reduce R[σ, b,a, b′\a′] := sum(σ′) W[σ,σ′,b,b′] * M[σ′,a,a′]  lazy
    @mul R2[σ, b,a, b′\a′] := W[σ,σ′,b,b′] * M[σ′,a,a′]
    @test R ≈ R2
    # invperm((4,2,1,3)) == (3, 2, 4, 1)

    @reduce S[σ, b,a, b′,a′] := sum(σ′) W[σ,σ′,b,b′] * M[σ′,a,a′]
    @mul S2[σ, b,a, b′,a′] := W[σ,σ′,b,b′] * M[σ′,a,a′]
    @test S ≈ S2
    # invperm((4,2,1,3,5)) == (3, 2, 4, 1, 5)

    ## in-place version
    N3 = similar(N); N4 = similar(N);
    @reduce N3[σ, b\a, b′\a′] = sum(σ′) W[σ,σ′,b,b′] * M[σ′,a,a′]  lazy
    @mul N4[σ, b\a, b′\a′] = W[σ,σ′,b,b′] * M[σ′,a,a′]
    @test N ≈ N3 ≈ N4

    ## these were OK:
    M = rand(3,3,3);

    @reduce Q[σ, b\a, a′] := sum(σ′) M[σ,σ′,b] * M[σ′,a,a′];
    @mul Q2[σ, b\a, a′] := M[σ,σ′,b] * M[σ′,a,a′];
    @test Q ≈ Q2

    @reduce Q3[σ, b,a, a′] := sum(σ′) M[σ,σ′,b] * M[σ′,a,a′];
    @mul    Q4[σ, b,a, a′] := sum(σ′) M[σ,σ′,b] * M[σ′,a,a′];
    @test Q3 ≈ Q4

end
@testset "in & out shapes" begin

    bc = rand(2,3)
    cd = rand(3,4)

    @mul A[b,_,d] := bc[b,c] * cd[c,d]
    @mul A1[b,1,d] := bc[b,c] * cd[c,d]
    @test size(A) == size(A1) == (2,1,4)

    B = similar(A)
    @mul B[b,_,d] = bc[b,c] * cd[c,d] # in-place
    @test B == A

    @mul C[b] := bc[b,c] * cd[c,3]
    @test C == bc * cd[:,3]

    D = similar(C)
    @mul D[b] = bc[b,c] * cd[c,3] # in-place
    @test D == C

end
@testset "batch" begin

    bbb = rand(2,2,2)
    bbb2 = randn(2,2,2)

    @mul A[i,k,n] := bbb[i,j,n] * bbb2[j,k,n]
    @reduce B[i,k,n] := sum(j) bbb[i,j,n] * bbb2[j,k,n]
    @einsum C[i,k,n] := bbb[i,j,n] * bbb2[j,k,n]
    @test A == B == C

    D = similar(A);
    E = similar(B);
    F = similar(C);
    @mul D[i,k,n] = bbb[i,j,n] * bbb2[j,k,n] # in-place
    @reduce E[i,k,n] = sum(j) bbb[i,j,n] * bbb2[j,k,n]
    @einsum F[i,k,n] = bbb[i,j,n] * bbb2[j,k,n]
    @test A == D == E == F

    ## permute
    @mul A[n,k,i] := bbb[i,j,n] * bbb2[n,k,j]
    @reduce B[n,k,i] := sum(j) bbb[i,j,n] * bbb2[n,k,j]
    @einsum C[n,k,i] := bbb[i,j,n] * bbb2[n,k,j]
    @test A == B == C

    D = similar(A);
    E = similar(B);
    F = similar(C);
    @mul D[n,k,i] := bbb[i,j,n] * bbb2[n,k,j] # in-place
    @reduce E[n,k,i] := sum(j) bbb[i,j,n] * bbb2[n,k,j]
    @einsum F[n,k,i] := bbb[i,j,n] * bbb2[n,k,j]
    @test A == D == E == F

    ## readme & help
    X = rand(2,3,4)
    Y = rand(3,2,4)
    @mul W[β][i,j] := X[i,k,β] * Y[k,j,β]
    @test W[2] ≈ X[:,:,2] * Y[:,:,2]

    B = rand(8)
    C = rand(3,2,4)
    @mul A[_,i][j] := B[i\k] * C[j,2,k]
    rB = reshape(B,2,4)
    @einsum D[i,j] := rB[i,k] * C[j,2,k]
    @test A[1,2] ≈ D[2,:]

end
@testset "recursion" begin

    A = rand(2,3); B = rand(3,4); C = rand(4,5);

    @einsum V[i,l] := A[i,j] * B[j,k] * C[k,l]
    @reduce V2[i,l] := sum(j,k) A[i,j] * B[j,k] * C[k,l]

    @reduce W[i,l] := sum(j) A[i,j] * @mul [j,l] := B[j,k] * C[k,l]

    # @mul W2[i,l] := sum(j) A[i,j] * @mul [j,l] := B[j,k] * C[k,l]
    # maybe I decided not to allow this for now.

    @test V ≈ V2 ≈ W #≈ W2

end


text:
@testset "Sanity Row/String" begin
    @test distance_row("", "")    == 0

    @testset "Sanity Row/String/Left" begin
        @test distance_row("a", "")   == 1
        @test distance_row("ab", "")  == 2
        @test distance_row("abc", "") == 3
    end

    @testset "Sanity Row/String/Right" begin
        @test distance_row("", "a")   == 1
        @test distance_row("", "ab")  == 2
        @test distance_row("", "abc") == 3
    end

    @testset "Sanity Row/String/Both" begin
        @test distance_row("a", "a")     == 0
        @test distance_row("a", "b")     == 1
        @test distance_row("ab", "bb")   == 1
        @test distance_row("abc", "bbc") == 1
        @test distance_row("abc", "abb") == 1
    end
end

max_stars_repo_name: jlegare/LevenshteinToolkit

max_issues_repo_name: jlegare/LevenshteinToolkit

max_forks_repo_name: jlegare/LevenshteinToolkit


text:
module Sandcorn
using Statistics: mean, var
using Distributions
using CorrelationFunctions
using Optim
using Base.Iterators

include("generate.jl")
include("likelihood.jl")
include("optimize.jl")

export generate_image, sandcorn_parameters

end # module


text:
"""
    AbstractSolver

Supertype for solvers that may be used to solve implicit systems of equations.
"""
abstract type AbstractSolver end

"""
    create_mutating_solver_cache(component::AbstractImplicitComponent, x, solver::AbstractSolver)

Construct a cache for a nonlinear solver which uses mutating functions.  `x` is
a pointer to an array that will be modified to contain the correct input values.
"""
create_mutating_solver_cache(component::AbstractImplicitComponent, x, ::AbstractSolver) = component

"""
    create_nonmutating_solver_cache(component::AbstractImplicitComponent, x, solver::AbstractSolver)

Construct a cache for a nonlinear solver which uses nonmutating functions.  `x` is
a pointer to an array that will be modified to contain the correct input values.
"""
create_nonmutating_solver_cache(component::AbstractImplicitComponent, x, ::AbstractSolver) = component

# --- Newton Solver --- #

"""
    Newton{T} <: AbstractSolver

Newton solver, as provided by the NLsolve package.  Stores the keyword arguments
which are passed to NLsolve in the field `kwargs`
"""
struct Newton{T} <: AbstractSolver
    kwargs::T
end

"""
    Newton(; kwargs...)

Newton solver, as provided by the NLsolve package.  Uses the same default
arguments as NLsolve, but defaults to using the BackTracking linesearch from
the LineSearches package rather than no linesearch.

The relevant keyword arguments and their default values are repeated here for
convenience.

# Keyword Arguments
 - `xtol = 0.0`: Step size tolerance for ending iterations
 - `ftol = 1e-8`: Function tolerance for ending iterations
 - `iterations = 1000`: Maximum number of iterations
 - `linesearch = LineSearches.BackTracking(maxstep=1e6)`: Linesearch algorithm
 - `linsolve = (x, A, b) -> copyto!(x, A\b)`: Linear solver function
"""
function Newton(; linesearch=BackTracking(maxstep=1e6), kwargs...)
    kwargs = (linesearch=linesearch, kwargs...)
    return Newton(kwargs)
end

function (solver::Newton)(cache, x, y0)
    # solve nonlinear system
    results = nlsolve(cache, y0, method=:newton; solver.kwargs...)
    # return NaNs if not converged
    if !results.f_converged
        copyto!(results.zero, NaN)
    end
    # return results
    return results.zero
end

# --- Trust Region Solver --- #

"""
    TrustRegion{T} <: AbstractSolver

Trust region solver, as provided by the NLsolve package.  Stores the keyword
arguments which are passed to NLsolve in the field `kwargs`
"""
struct TrustRegion{T} <: AbstractSolver
    kwargs::T
end

"""
    TrustRegion(; kwargs...)

Trust region solver, as provided by the NLsolve package.  Uses the same default
arguments as NLsolve.

The relevant keyword arguments and their default values are repeated here for
convenience.

# Keyword Arguments
 - `xtol = 0.0`: Step size tolerance for ending iterations
 - `ftol = 1e-8`: Function tolerance for ending iterations
 - `iterations = 1000`: Maximum number of iterations
 - `factor = 1.0`: Factor for determining initial size of trust region
 - `autoscale = true`: Autoscale variables?
 - `linsolve = (x, A, b) -> copyto!(x, A\b)`: Linear solver function
"""
function TrustRegion(; kwargs...)
    return TrustRegion(kwargs)
end

function (solver::TrustRegion)(cache, x, y0)
    # solve nonlinear system
    results = nlsolve(cache, y0, method=:trust_region; solver.kwargs...)
    # return NaNs if not converged
    if !results.f_converged
        copyto!(results.zero, NaN)
    end
    # return results
    return results.zero
end

# NLsolve uses the same cache for TrustRegion and Newton
function create_mutating_solver_cache(component::AbstractImplicitComponent, x, ::Union{TrustRegion, Newton})
    f! = function(r, y)
        residuals!(component, x, y)
        copyto!(r, residuals(component))
        return r
    end
    j! = function(drdy, y)
        residual_output_jacobian!(component, x, y)
        copyto!(drdy, residual_output_jacobian(component))
        return drdy
    end
    fj! = function(r, drdy, y)
        residuals_and_output_jacobian!(component, x, y)
        copyto!(r, residuals(component))
        copyto!(drdy, residual_output_jacobian(component))
        return r, drdy
    end
    r = residuals(component)
    drdy = residual_output_jacobian(component)
    y_f = component.y_f
    y_df = component.y_dfdy
    f_calls = [0,]
    df_calls = [0,]
    return OnceDifferentiable(f!, j!, fj!, r, drdy, y_f, y_df, f_calls, df_calls)
end

# NLsolve uses the same cache for TrustRegion and Newton
function create_nonmutating_solver_cache(component::AbstractImplicitComponent, x, ::Union{TrustRegion, Newton})
    f! = (r, y) -> residuals!(component, r, x, y)
    j! = (drdy, y) -> residual_output_jacobian!(component, drdy, x, y)
    fj! = (r, drdy, y) -> residuals_and_output_jacobian!(component, r, drdy, x, y)
    r = residuals(component)*one(eltype(x))
    drdy = residual_output_jacobian(component)*one(eltype(x))
    y_f = component.y_f*one(eltype(x))
    y_df = component.y_dfdy*one(eltype(x))
    f_calls = [0,]
    df_calls = [0,]
    return OnceDifferentiable(f!, j!, fj!, r, drdy, y_f, y_df, f_calls, df_calls)
end


text:
# Finite-difference, single-active electron
struct FDSAEWfn{T} <: AbstractSAEWfn{T}
    data::Matrix{Complex{T}}
    dr::T
end

norm(wfn::FDSAEWfn) = vecnorm(wfn.data)*sqrt(wfn.dr)

*(A::Union{AbstractMatrix,AbstractLinearMap}, wfn::FDSAEWfn) =
    FDSAEWfn(reshape(A*view(wfn.data, :), size(wfn.data)), wfn.dr)

@recipe function plot(wfn::FDSAEWfn)
    wfn.dr*(1:size(wfn.data,1)), abs2.(wfn.data)
end

export FDSAEWfn