"""
$(SIGNATURES)

Returns a partial LMM model (non-functional due to lacking data) to be used in likelihoodratiotests.
`k` to selcet which of the modelfit's to fake
"""
function fake_lmm(m::UnfoldLinearMixedModel,k::Int)
    (feterm,reterm) = Unfold.designmatrix(m).Xs
    fakeY = zeros(size(feterm,1))
    lmm = Unfold.LinearMixedModel_wrapper(Unfold.formula(m),fakeY,designmatrix(m).Xs)

    fcoll = Unfold.modelfit(m)
    #lmm.objective .= fcoll.fits[1].objective
    #lmm.optsum.feval .= 1
    #lmm.optsum.fmin .= 1
    lmm.optsum.sigma = fcoll.fits[k].σ
    lmm.optsum.optimizer = :unfold
    lmm.optsum.returnvalue = :success
    setθ!(lmm,fcoll.fits[k].θ)
    return lmm
end
    fake_lmm(m::UnfoldLinearMixedModel) = fake_lmm.(m,1:length(modelfit(m).fits))

    """
    $(SIGNATURES)

    Calculate likelihoodratiotest
    """
	function MixedModels.likelihoodratiotest(m::UnfoldLinearMixedModel...)
		#@info lrtest(fake_lmm.(m,1)...)
		n = length(Unfold.modelfit(m[1]).fits)
		lrt = Array{MixedModels.LikelihoodRatioTest}(undef,n)
		sizehint!(lrt,n)
		for k = 1:n
			ms = fake_lmm.(m,k)
			@info objective.(ms)
			lrt[k] = MixedModels.likelihoodratiotest(ms...)
		end
		return lrt
	end

    """
    $(SIGNATURES)
    Unfold-Method: return pvalues of likelihoodratiotests, typically calculated:

    # Examples
    julia> pvalues(likelihoodratiotest(m1,m2))

    where m1/m2 are UnfoldLinearMixedModel's

    Tipp: if you only compare two models you can easily get a vector of p-values:

    julia> vcat(pvalues(likelihoodratiotest(m1,m2))...)

    Multiple channels are returned linearized at the moment, as we do not have access to the amount of channels after the LRT, you can do:

    julia> reshape(vcat(pvalues(likelihoodratiotest(m1,m2))...),ntimes,nchan)'

    """
	function pvalues(lrtvec::Vector{MixedModels.LikelihoodRatioTest})
		[lrt.pvalues for lrt in lrtvec]
	end


@inline getbody(mechanism::Mechanism, id::Int64) = mechanism.bodies[id]
@inline getbody(mechanism::Mechanism, id::Nothing) = mechanism.origin
function getbody(mechanism::Mechanism, name::String)
    if mechanism.origin.name == name
        return mechanism.origin
    else
        for body in mechanism.bodies
            if body.name == name
                return body
            end
        end
    end
    return nothing
end
@inline geteqconstraint(mechanism::Mechanism, id::Int64) = mechanism.eqconstraints[id]
function geteqconstraint(mechanism::Mechanism, name::String)
    for eqc in mechanism.eqconstraints
        if eqc.name == name
            return eqc
        end
    end
    return nothing
end
@inline getineqconstraint(mechanism::Mechanism, id::Int64) = mechanism.ineqconstraints[id]
function getineqconstraint(mechanism::Mechanism, name::String)
    for ineqc in mechanism.eqconstraints
        if ineqc.name == name
            return ineqc
        end
    end
    return nothing
end

getcomponent(mechanism::Mechanism, id::Nothing) = mechanism.origin
function getcomponent(mechanism::Mechanism, id::Int64)
    if haskey(mechanism.bodies, id)
        return getbody(mechanism, id)
    elseif haskey(mechanism.eqconstraints, id)
        return geteqconstraint(mechanism, id)
    elseif haskey(mechanism.ineqconstraints, id)
        return getineqconstraint(mechanism, id)
    else
        return nothing
    end
end

function getcomponent(mechanism::Mechanism, name::String)
    component = getbody(mechanism,name)
    if component == nothing
        component = geteqconstraint(mechanism,name)
    end
    if component == nothing
        component = getineqconstraint(mechanism,name)
    end
    return component
end

function getshape(mechanism::Mechanism, id)
    for shape in mechanism.shapes
        for bodyid in shape.bodyids
            if bodyid == id
                return shape
            end
        end
    end

    return nothing
end

function setentries!(mechanism::Mechanism)
    graph = mechanism.graph
    ldu = mechanism.ldu

    for (id, body) in pairs(mechanism.bodies)
        for cid in directchildren(graph, id)
            setLU!(mechanism, getentry(ldu, (id, cid)), id, geteqconstraint(mechanism, cid))
        end

        diagonal = getentry(ldu, id)
        setDandΔs!(mechanism, diagonal, body)
        for cid in ineqchildren(graph, id)
            extendDandΔs!(mechanism, diagonal, body, getineqconstraint(mechanism, cid))
        end
    end

    for node in mechanism.eqconstraints
        id = node.id

        for cid in directchildren(graph, id)
            setLU!(mechanism, getentry(ldu, (id, cid)), node, cid)
        end

        for cid in loopchildren(graph, id)
            setLU!(getentry(ldu, (id, cid)))
        end

        diagonal = getentry(ldu, id)
        setDandΔs!(mechanism, diagonal, node)
    end
end

@inline function normf(mechanism::Mechanism, body::Body{T}) where T
    f = dynamics(mechanism, body)
    return dot(f, f)
end

@inline function normf(mechanism::Mechanism, c::EqualityConstraint)
    f = g(mechanism, c)
    return dot(f, f)
end

@inline function normf(mechanism::Mechanism, ineqc::InequalityConstraint)
    f = gs(mechanism, ineqc)
    d = h(ineqc)
    return dot(f, f) + dot(d, d)
end

@inline function normfμ(mechanism::Mechanism, ineqc::InequalityConstraint)
    f = gs(mechanism, ineqc)
    d = hμ(ineqc, mechanism.μ)
    return dot(f, f) + dot(d, d)
end

@inline function GtλTof!(mechanism::Mechanism, body::Body, eqc::EqualityConstraint)
    body.f -= ∂g∂pos(mechanism, eqc, body.id)' * eqc.s1
    return
end

@inline function NtγTof!(mechanism::Mechanism, body::Body, ineqc::InequalityConstraint)
    body.f -= ∂g∂pos(mechanism, ineqc, body)' * ineqc.γ1
    return
end

@inline function normf(mechanism::Mechanism)
    mechanism.normf = 0

    # Allocates otherwise
    for body in mechanism.bodies
        mechanism.normf += normf(mechanism, body)
    end
    foreach(addNormf!, mechanism.eqconstraints, mechanism)
    foreach(addNormf!, mechanism.ineqconstraints, mechanism)

    return sqrt(mechanism.normf)
end

@inline function meritf(mechanism::Mechanism)
    mechanism.normf = 0

    # Allocates otherwise
    for body in mechanism.bodies
        mechanism.normf += normf(mechanism, body)
    end
    foreach(addNormf!, mechanism.eqconstraints, mechanism)
    foreach(addNormfμ!, mechanism.ineqconstraints, mechanism)

    return sqrt(mechanism.normf)
end

@inline function normΔs(mechanism::Mechanism)
    mechanism.normΔs = 0

    # Allocates otherwise
    mechanism.normΔs += mapreduce(normΔs, +, mechanism.bodies)
    foreach(addNormΔs!, mechanism.eqconstraints, mechanism)
    foreach(addNormΔs!, mechanism.ineqconstraints, mechanism)

    return sqrt(mechanism.normΔs)
end

@inline function addNormf!(ineqc::InequalityConstraint, mechanism::Mechanism)
    mechanism.normf += normf(mechanism, ineqc)
    return
end

@inline function addNormfμ!(ineqc::InequalityConstraint, mechanism::Mechanism)
    mechanism.normf += normfμ(mechanism, ineqc)
    return
end

@inline function addNormf!(eqc::EqualityConstraint, mechanism::Mechanism)
    mechanism.normf += normf(mechanism, eqc)
    return
end

@inline function addNormΔs!(component::Component, mechanism::Mechanism)
    mechanism.normΔs += normΔs(component)
    return
end

function feasibilityStepLength!(mechanism::Mechanism)
    ldu = mechanism.ldu

    τ = 0.995
    mechanism.α = 1.

    for ineqc in mechanism.ineqconstraints
        feasibilityStepLength!(mechanism, ineqc, getineqentry(ldu, ineqc.id), τ)
    end

    return
end

function feasibilityStepLength!(mechanism, ineqc::InequalityConstraint{T,N}, ineqentry::InequalityEntry, τ) where {T,N}
    s1 = ineqc.s1
    γ1 = ineqc.γ1
    Δs = ineqentry.Δs
    Δγ = ineqentry.Δγ

    for i = 1:N
        αmax = τ * s1[i] / Δs[i]
        (αmax > 0) && (αmax < mechanism.α) && (mechanism.α = αmax)
        αmax = τ * γ1[i] / Δγ[i]
        (αmax > 0) && (αmax < mechanism.α) && (mechanism.α = αmax)
    end
    return
end

function setFrictionForce!(mechanism)
    foreach(setFrictionForce!, mechanism.ineqconstraints, mechanism)
end

function saveToStorage!(mechanism::Mechanism, t)
    No = mechanism.No
    for (ind, body) in enumerate(mechanism.bodies)
        mechanism.storage.x[ind][t] = body.x[No]
        mechanism.storage.q[ind][t] = body.q[No]
    end
    for (ind, eqc) in enumerate(mechanism.eqconstraints)
        mechanism.storage.eqmultipliers[ind][t] = eqc.s1
    end
end

@inline function updatePos!(body::Body, Δt)
    body.x[1] = body.x[2]
    body.x[2] = getx3(body, Δt)
    body.q[1] = body.q[2]
    body.q[2] = getq3(body, Δt)
    return
end

function verifyConstraints!(mechanism::Mechanism)
    for eqc in mechanism.eqconstraints
        if norm(g(mechanism, eqc)) > 1e-3
            @info string("Probably disconnected bodies at constraint: ", eqc.id)
        end
    end
end

function simulate!(mechanism::Mechanism;save::Bool = false,debug::Bool = false)
    debug && verifyConstraints!(mechanism)
    bodies = mechanism.bodies
    eqcs = mechanism.eqconstraints
    ineqcs = mechanism.ineqconstraints
    Δt = mechanism.Δt
    foreach(s0tos1!, bodies)
    foreach(s0tos1!, eqcs)
    foreach(s0tos1!, ineqcs)

    for i = mechanism.steps
        newton!(mechanism, warning = debug)
        save && saveToStorage!(mechanism, i)
        foreach(updatePos!, bodies, Δt)
    end
    return
end

function simulate!(mechanism::Mechanism, control!::Function;save::Bool = false,debug::Bool = false)
    debug && verifyConstraints!(mechanism)
    bodies = mechanism.bodies
    eqcs = mechanism.eqconstraints
    ineqcs = mechanism.ineqconstraints
    Δt = mechanism.Δt
    foreach(s0tos1!, bodies)
    foreach(s0tos1!, eqcs)
    foreach(s0tos1!, ineqcs)

    for i = mechanism.steps
        control!(mechanism, (i - 1) * Δt)
        newton!(mechanism, warning = debug)
        save && saveToStorage!(mechanism, i)
        foreach(updatePos!, bodies, Δt)
    end
    return
end

function simulate!(mechanism::Mechanism, controller::Controller;save::Bool = false,debug::Bool = false)
    debug && verifyConstraints!(mechanism)
    bodies = mechanism.bodies
    eqcs = mechanism.eqconstraints
    ineqcs = mechanism.ineqconstraints
    Δt = mechanism.Δt
    foreach(s0tos1!, bodies)
    foreach(s0tos1!, eqcs)
    foreach(s0tos1!, ineqcs)

    for i = mechanism.steps
        control!(mechanism, controller)
        newton!(mechanism, warning = debug)
        save && saveToStorage!(mechanism, i)
        foreach(updatePos!, bodies, Δt)
    end
    return
end

function plotθ(mechanism::Mechanism{T}, id) where T
    n = length(mechanism.bodies)
    θ = zeros(T, n, length(mechanism.steps))
    for i = 1:n
        qs = mechanism.storage.q[i]
        for (t, q) in enumerate(qs)
            θ[i,t] = angleaxis(q)[1] * sign(angleaxis(q)[2][1])
        end
    end

    p = plot(collect(0:mechanism.Δt:mechanism.tend - mechanism.Δt), θ[id[1],:])
    for ind in Iterators.rest(id, 2)
        plot!(collect(0:mechanism.Δt:mechanism.tend - mechanism.Δt), θ[ind,:])
    end
    return p
end

function plotλ(mechanism::Mechanism{T}, id) where T
    n = sum(length.(mechanism.eqconstraints))
    λ = zeros(T, n, length(mechanism.steps))
    startpos = 1
    endpos = 0
    for i = 1:length(mechanism.eqconstraints)
        endpos = startpos + length(mechanism.eqconstraints[i]) - 1

        λs = mechanism.storage.λ[i]
        for (t, val) in enumerate(λs)
            λ[startpos:endpos,t] = val
        end

        startpos = endpos + 1
    end

    p = plot(collect(0:mechanism.Δt:mechanism.tend - mechanism.Δt), λ[id[1],:])
    for ind in Iterators.rest(id, 2)
        plot!(collect(0:mechanism.Δt:mechanism.tend - mechanism.Δt), λ[ind,:])
    end
    return p
end


#
# This file is a part of MolecularGraph.jl
# Licensed under the MIT License http://opensource.org/licenses/MIT
#

export
    maximalcliques, maximumclique,
    maximalconncliques, maximumconnclique

mutable struct FindCliqueState{T<:UndirectedGraph}
    graph::T
    targetsize::Union{Int,Nothing} # Int

    adjacencies::Dict{Int,Set{Int}}
    expire::Union{UInt64,Nothing} # UInt64, nanoseconds
    Q::Vector{Int}

    cliques::Vector{Set{Int}}
    status::Symbol

    function FindCliqueState{T}(graph; timeout=nothing, targetsize=nothing,
                kwargs...) where {T<:UndirectedGraph}
        if timeout !== nothing
            expire = (time_ns() + timeout * 1_000_000_000)::UInt64
        else
            expire = nothing
        end
        # fast adjacency access
        adj = Dict{Int,Set{Int}}()
        for n in nodeset(graph)
            adj[n] = adjacencies(graph, n)
        end
        new(graph, targetsize, adj, expire, [], [], :ongoing)
    end
end

mutable struct FindConnCliqueState{T<:ModularProduct}
    graph::T
    targetsize::Union{Int,Nothing} # Int

    adjacencies::Dict{Int,Set{Int}}
    connected::Dict{Int,Set{Int}}
    disconn::Dict{Int,Set{Int}}
    expire::Union{UInt64,Nothing} # UInt64, nanoseconds

    cliques::Vector{Set{Int}}
    status::Symbol

    function FindConnCliqueState{T}(graph; timeout=nothing, targetsize=nothing,
                kwargs...) where {T<:ModularProduct}
        if timeout !== nothing
            expire = (time_ns() + timeout * 1_000_000_000)::UInt64
        else
            expire = nothing
        end
        # fast adjacency access
        adj = Dict{Int,Set{Int}}()
        conn = Dict{Int,Set{Int}}()
        disconn = Dict{Int,Set{Int}}()
        for n in nodeset(graph)
            adj[n] = Set{Int}()
            conn[n] = Set{Int}()
            disconn[n] = Set{Int}()
            for (i, a) in neighbors(graph, n)
                push!(adj[n], a)
                if edgeattr(graph, i).hasedge
                    push!(conn[n], a)
                else
                    push!(disconn[n], a)
                end
            end
        end
        new(graph, targetsize, adj, conn, disconn, expire, [], :ongoing)
    end
end

function expand!(state::FindCliqueState, subg, cand)
    (state.status == :timedout || state.status == :targetreached) && return
    if isempty(subg)
        # Report max clique
        push!(state.cliques, Set(state.Q))
        return
    elseif state.expire !== nothing && time_ns() > state.expire
        state.status = :timedout
        return
    elseif state.targetsize !== nothing && length(state.Q) >= state.targetsize
        state.status = :targetreached
        push!(state.cliques, Set(state.Q))
        return
    end
    candnbrcnt(n) = length(intersect(cand, state.adjacencies[n]))
    pivot = sortstablemax(subg, by=candnbrcnt)
    copv = setdiff(cand, state.adjacencies[pivot])

    for q in copv
        push!(state.Q, q)
        qnbrs = state.adjacencies[q]
        subgq = intersect(subg, qnbrs)
        candq = intersect(cand, qnbrs)
        expand!(state, subgq, candq)
        pop!(cand, q)
        pop!(state.Q)
    end
    return
end

function expandconn!(state::FindConnCliqueState, R, P, Q, X, Y)
    (state.status == :timedout || state.status == :targetreached) && return
    if isempty(P) && isempty(X)
        # Report max clique
        push!(state.cliques, copy(R))
        return
    elseif state.expire !== nothing && time_ns() > state.expire
        state.status = :timedout
        return
    elseif state.targetsize !== nothing && length(R) >= state.targetsize
        state.status = :targetreached
        push!(state.cliques, copy(R))
        return
    end
    while !isempty(P)
        n = pop!(P)
        Rnew = union(R, [n])
        Qnew = intersect(Q, state.disconn[n])
        Pnew = union(
            intersect(P, state.adjacencies[n]),
            intersect(Q, state.connected[n]))
        Ynew = intersect(Y, state.disconn[n])
        Xnew = union(
            intersect(X, state.adjacencies[n]),
            intersect(Y, state.connected[n]))
        expandconn!(state, Rnew, Pnew, Qnew, Xnew, Ynew)
        push!(X, n)
    end
    return
end

"""
    maximalcliques(graph::UndirectedGraph; kwargs...
        ) -> Tuple{Vector{Set{Int}}, Symbol}

Return maximal cliques.

# Reference

1. Tomita, E., Tanaka, A., & Takahashi, H. (2006). The worst-case time
   complexity for generating all maximal cliques and computational experiments.
   Theoretical Computer Science, 363(1), 28–42.
   https://doi.org/10.1016/J.TCS.2006.06.015
1. Cazals, F., & Karande, C. (2008). A note on the problem of reporting maximal
   cliques. Theoretical Computer Science, 407(1–3), 564–568.
   https://doi.org/10.1016/j.tcs.2008.05.010
"""
function maximalcliques(graph::T; kwargs...) where {T<:UndirectedGraph}
    state = FindCliqueState{T}(graph; kwargs...)
    expand!(state, nodeset(graph), nodeset(graph))
    if state.status == :ongoing
        state.status = :done
    end
    return (state.cliques, state.status)
end

"""
    maximumclique(graph::UndirectedGraph; kwargs...) -> Tuple{Set{Int}, Symbol}

Return a maximum clique.

"""
function maximumclique(graph::T; kwargs...) where {T<:UndirectedGraph}
    (cliques, status) = maximalcliques(graph; kwargs...)
    return (sortstablemax(cliques, by=length, init=[]), status)
end

"""
    maximalconncliques(graph::ModularProduct; kwargs...
        ) -> Tuple{Vector{Set{Int}}, Symbol}

Return maximal connected cliques.

# Reference

1. Cazals, F., & Karande, C. (2005). An algorithm for reporting maximal
   c-cliques. Theoretical Computer Science, 349(3), 484–490.
   https://doi.org/10.1016/j.tcs.2005.09.038

"""
function maximalconncliques(graph::T; kwargs...) where {T<:UndirectedGraph}
    state = FindConnCliqueState{T}(graph; kwargs...)
    nodes = nodeset(graph)
    done = Set{Int}()
    for n in nodes
        R = Set([n])
        P = intersect(setdiff(nodes, done), state.connected[n])
        Q = intersect(setdiff(nodes, done), state.disconn[n])
        X = intersect(state.connected[n], done)
        Y = intersect(state.disconn[n], done)
        expandconn!(state, R, P, Q, X, Y)
        push!(done, n)
    end
    if state.status == :ongoing
        state.status = :done
    end
    return (state.cliques, state.status)
end

"""
    maximumconnclique(graph::ModularProduct; kwargs...
        ) -> Tuple{Set{Int}, Symbol}

Return a maximum connected clique.

"""
function maximumconnclique(graph::T; kwargs...) where {T<:ModularProduct}
    (cliques, status) = maximalconncliques(graph; kwargs...)
    return (sortstablemax(cliques, by=length, init=[]), status)
end


const ContactImplicitMPC = Main
include(joinpath(@__DIR__, "..", "src/dynamics", "hopper_2D", "visuals.jl"))
T = Float64
vis = Visualizer()
open(vis)

# get hopper model
s_sim = get_simulation("hopper_2D", "sine2_2D_lc", "sinusoidal")
s = get_simulation("hopper_2D", "flat_2D_lc", "flat")
model = s.model
env = s.env

# get trajectory
ref_traj = deepcopy(ContactImplicitMPC.get_trajectory(s.model, s.env,
    joinpath(module_dir(), "src/dynamics/hopper_2D/gaits/gait_forward.jld2"),
    load_type = :joint_traj))
# time
H = ref_traj.H
h = ref_traj.h
N_sample = 5
H_mpc = 10
h_sim = h / N_sample
H_sim = 100*H*N_sample #500*H*N_sample

# Select the initial speed
v0 = 0.2
Tstance = 0.13 # measure using hop-in-place gait
Tflight = 0.62 # measure using hop-in-place gait
p = raibert_policy(s_sim.model, v0=v0, Tstance=Tstance, Tflight=Tflight, h=h)

off0 = SVector{model.dim.q,T}([0.0, 0.5, 0.0, 0.0])
off1 = SVector{model.dim.q,T}([0*v0*h_sim, 0.5, 0.0, 0.0])
q_ref = SVector{model.dim.q,T}([0.0, 0.5, 0.0, 0.5])
q0_sim = copy(q_ref) + off0
q1_sim = copy(q_ref) + off1

sim = ContactImplicitMPC.simulator(s_sim, q0_sim, q1_sim, h_sim, H_sim,
    p = p,
    ip_opts = ContactImplicitMPC.InteriorPointOptions(
		γ_reg = 0.0,
		undercut = Inf,
		r_tol = 1.0e-8,
		κ_tol = 1.0e-8,),
    sim_opts = ContactImplicitMPC.SimulatorOptions(warmstart = true))

status = ContactImplicitMPC.simulate!(sim, verbose = true)

plt = plot(layout=(3,1), legend=false)
plot!(plt[1,1], hcat(Vector.(vcat([fill(ref_traj.q[i], N_sample) for i=1:H]...))...)',
    color=:red, linewidth=3.0)
plot!(plt[1,1], hcat(Vector.(sim.traj.q)...)', color=:blue, linewidth=1.0)
plot!(plt[2,1], hcat(Vector.(vcat([fill(ref_traj.u[i][1:nu], N_sample) for i=1:H]...))...)',
    color=:red, linewidth=3.0)
plot!(plt[2,1], hcat(Vector.([u[1:model.dim.u] for u in sim.traj.u]*N_sample)...)', color=:blue, linewidth=1.0)
plot!(plt[3,1], hcat(Vector.([γ[1:model.dim.c] for γ in sim.traj.γ]*N_sample)...)', color=:blue, linewidth=1.0)

plot_lines!(vis, model, sim.traj.q[1:10:end])
plot_surface!(vis, s_sim.env, n=200, xlims = [-1, 40])
anim = visualize_robot!(vis, model, sim.traj, sample=5)
anim = visualize_force!(vis, model, s_sim.env, sim.traj, anim=anim, h=h_sim, sample = 5)

plot(hcat([u[[1,2]] for u in Vector.(sim.traj.u)[1:end]]...)')

# filename = "hopper_ipopt"
# MeshCat.convert_frames_to_video(
#     "/home/simon/Downloads/$filename.tar",
#     "/home/simon/Documents/$filename.mp4", overwrite=true)
#
# convert_video_to_gif(
#     "/home/simon/Documents/$filename.mp4",
#     "/home/simon/Documents/$filename.gif", overwrite=true)

# const ContactImplicitMPC = Main


using Test
using SparseIR

@testset "basis.jl" begin
    @testset "consistency" begin
        β = 2.0
        wmax = 5.0
        ε = 1e-5
        basis_f, basis_b = SparseIR.finite_temp_bases(β, wmax, ε)
        bs = FiniteTempBasisSet(β, wmax, ε)

        @test length(bs.basis_f) == length(basis_f)
        @test length(bs.basis_b) == length(basis_b)
    end

    @testset "consistency2" begin
        β = 2.0
        wmax = 5.0
        ε = 1e-5

        sve_result = sve_logistic[β * wmax]
        basis_f, basis_b = SparseIR.finite_temp_bases(β, wmax, ε, sve_result)
        smpl_τ_f = TauSampling(basis_f)
        smpl_τ_b = TauSampling(basis_b)
        smpl_wn_f = MatsubaraSampling(basis_f)
        smpl_wn_b = MatsubaraSampling(basis_b)

        bs = FiniteTempBasisSet(β, wmax, ε; sve_result)
        @test smpl_τ_f.sampling_points == smpl_τ_b.sampling_points
        @test bs.smpl_tau_f.matrix == smpl_τ_f.matrix
        @test bs.smpl_tau_b.matrix == smpl_τ_b.matrix

        @test bs.smpl_wn_f.matrix == smpl_wn_f.matrix
        @test bs.smpl_wn_b.matrix == smpl_wn_b.matrix
    end
end


using HDF5

function move_onehot(move::Point)
    board = BitArray(N, N)
    fill!(board, false)
    board[linearindex(move)] = true
    board
end

"Generate training data of form [(features, move), ...] for a single sgf file"
function generate_training_data(filename::AbstractString;
                                features::Vector{Function}=DEFAULT_FEATURES)
    sgf = load_sgf(filename)
    examples = Vector{Tuple{BitArray, BitArray}}()
    if sgf == nothing
        return examples
    end
    board = Board()
    while true
        move = get_next_move(sgf)
        if move == nothing
            break
        end
        if move != PASS_MOVE
            computed_features = get_features(board, features=features)
            answer = move_onehot(move)
            push!(examples, (computed_features, answer))
        end
        play_move(board, move)
    end
    return examples
end

"Generate training data for many sgf files"
function generate_training_data(filenames::Vector{AbstractString};
                                progress_update=100,
                                features::Vector{Function} = DEFAULT_FEATURES)
    examples = Vector{Tuple{BitArray, BitArray}}()
    tm = time()
    for (i, filename) in enumerate(filenames)
        if i % progress_update == 0
            println(STDERR, "$(i)/$(length(filenames)): $(time() - tm)")
        end
	game = generate_training_data(filename, features=features)
        if length(game) > 0
            push!(examples, game...)
        end
    end
    examples
end

"Extract training data directly to an hdf5 file"
function extract_to_hdf5(hf5_filename::AbstractString, filenames::Vector{AbstractString}, features::Vector{Function}; progress_update=1000, nlabels = 81)
    example_data = get_features(Board(), features=features)
    data_dim = size(example_data)

    hf5 = h5open(hf5_filename, "w")
    data = d_create(hf5, "data", datatype(UInt8),
                 ((data_dim..., 1), (data_dim..., -1)),
                 "chunk", (data_dim..., 1))
    label = d_create(hf5, "label", datatype(UInt8),
                     ((nlabels, 1), (nlabels, -1)),
                     "chunk", (nlabels, 1))
    tm = time()
    nexamples = 0
    for (i, filename) in enumerate(filenames)
        if i % progress_update == 0
            println(STDERR, "$(i)/$(length(filenames)): $(time() - tm)")
        end
        game = generate_training_data(filename, features=features)
        if length(game) > 0
            new_N = nexamples + length(game)
            set_dims!(data, (data_dim..., new_N))
            set_dims!(label, (nlabels, new_N))
            data_batch = batch_training_examples([x[1] for x in game])
            label_batch = batch_training_examples([x[2][:] for x in game])
            # Assumes data is 3 dims per example and labels are 1
            data[:, :, :, nexamples+1:new_N] = data_batch
            label[:, nexamples+1:new_N] = label_batch
            nexamples = new_N
        end
    end
    close(hf5)
end

"""
Create a vector by appending the given arrays in memory.
Reshape such that they are concatenated along the last+1 dimension

Wrote this because
X = cat(4,[x[1] for x in examples]...);
is slow for some reason
"""
function batch_training_examples(arrs)
    sizeper = length(arrs[1])
    out = similar(arrs[1], sizeper * length(arrs))
    for (i, arr) in enumerate(arrs)
        start = sizeper * (i - 1) + 1
        finish = sizeper * (i)
        out[start:finish] = arrs[i][:]
    end
    reshape(out, size(arrs[1])..., length(arrs))
end


## ExoplanetsSysSim/examples/test.jl
## (c) 2015 Eric B. Ford

# workspace()
# reload("ExoplanetsSysSim")
using ExoplanetsSysSim

println("# Starting end-to-end test (setting parameters, creating physical catalogs, creating observed catalogs, computing summary statistics, computing the distance between them).  ")
test_eval_model()


"""
    NoOp <: Augmentor.AffineOperation

Identity transformation that does not do anything with the given
image, but instead passes it along unchanged (without copying).

Usually used in combination with [`Either`](@ref) to denote a
"branch" that does not perform any computation.
"""
struct NoOp <: AffineOperation end

@inline supports_eager(::Type{NoOp}) = false
@inline supports_stepview(::Type{NoOp}) = true
@inline supports_view(::Type{NoOp}) = true

# TODO: implement method for n-dim arrays
toaffinemap(::NoOp, img::AbstractMatrix) = AffineMap(@SMatrix([1. 0; 0 1.]), @SVector([0.,0.]))
applyeager(::NoOp, img::AbstractArray, param) = plain_array(img)
applylazy(::NoOp, img::AbstractArray, param) = img

function applyview(::NoOp, img::AbstractArray, param)
    idx = map(i->1:length(i), axes(img))
    indirect_view(img, idx)
end

function applystepview(::NoOp, img::AbstractArray, param)
    idx = map(i->1:1:length(i), axes(img))
    indirect_view(img, idx)
end

function showconstruction(io::IO, op::NoOp)
    print(io, typeof(op).name.name, "()")
end

function Base.show(io::IO, op::NoOp)
    if get(io, :compact, false)
        print(io, "No operation")
    else
        print(io, "Augmentor.")
        showconstruction(io, op)
    end
end


#### PRREQUISITES ####
 include("objectives.jl")

#### for module multiobjectiveGraphs ####
export  multiobj # , getObj, setObj!, summarize
# export descr, defaultValue, mimic, combine # already in objectives.jl
# export lexicoBetter, dominates, dominatesStrictly, objLength

import Base.size, Base.length # Julia 1+ , allow redefinition of theses functions

##############################################################
#################### generic type for multi-objective values     #########
##############################################################

#### note : all objectives do not need to share the same type of values ##################

mutable struct multiobj # <: genericMultiobj # Tobjval is union of all the Tobjvals of the objectives
	nb::UInt8;  ######### number of objectives is (severely :-) limited to 255
	objectives::Vector@{ genericWeightCategory}
	############## constructors
	function multiobj()
			## create an monoobj with initialised array of  1 summable objective
		ar= Vector{ genericWeightCategory }(undef, 1)    # Vector{ weightMinSum{Tobjval} }( 1)  for Julia 0.63
		ar[1]=weightMinSum{Float64}()
		return new(1,ar)
	end													##############

	function multiobj( mo::multiobj )
	 	return new( mo.nb, mo.objectives ) ## clone with same types and same values
	end 												##############
	function multiobj(n::Number; typeofobjectives=weightMinSum{Float64})
			## create an multiobj with initialised array of summable objectives
		nnn::Int8=convert(Int8,n)
		ar= Vector{ genericWeightCategory }(undef,  nnn )  # suppress undef first parameter for Julia 0.63
		for i::Int64 in 1:nnn 	ar[i]=typeofobjectives() 	end
		return new(nnn,ar)
	end													##############
	function multiobj(p::Number,q::Number,r::Number ; typeofvalues=Float64)
			## create an multiobj with initialised array of
			## p summable objectives, q multiplicative and r bottleneck
		ar=Vector{genericWeightCategory}(undef, p+q+r) # suppress undef first parameter for Julia 0.63
		for i::Int64 in 1:p 	ar[i]=weightMinSum{typeofvalues}() 	end
		for i::Int64 in (p+1):(p+q) 	ar[i]=weightMinProduct{typeofvalues}() 	end
		for i::Int64 in (p+q+1):(p+q+r) 	ar[i]=weightMinMax{typeofvalues}() 	end
		return new(p+q+r,ar)
	end													##############
	function multiobj(mobjTypes::Vector{genericWeightCategory } )
			## create an multiobj with initialised types of weights (and matching default values)
		nn = length(mobjTypes)
		ar = Vector{genericWeightCategory}(undef,  nn ) # suppress undef first parameter for Julia 0.63
		   # with the same types in the same indices
		for i::Int64 in 1:nn 	ar[i]=mimic(mobjTypes[i], defaultValue( mobjTypes[i] ) ) 	end
		return new{Tobjval}( nn, ar )
	end
	function multiobj(mobjTypes::Array{genericWeightCategory } )
			## create an multiobj with initialised types of weights (and matching default values)
		nn = length(mobjTypes)
		ar = Vector{genericWeightCategory}(undef,  nn ) # suppress undef first parameter for Julia 0.63
		   # with the same types in the same indices
		for i::Int64 in 1:nn 	ar[i]=mimic(mobjTypes[i], defaultValue( mobjTypes[i] ) ) 	end
		return new{Tobjval}( nn, ar )
	end													##############
	function multiobj(mobjTypes::Vector{ genericWeightCategory },
								mobjVals::Vector{Tvals} )  where  Tvals <: Number
			## create an multiobj with array of objectives types and initialised array of values of weights
		nn = length(mobjVals)
		ar = Vector{genericWeightCategory}(undef, nn )  # suppress undef first parameter for Julia 0.63
		   # with the same types in the same indices
		for i::Int64 in 1:nn 	ar[i]=mimic(mobjTypes[i], mobjVals[i] ) 	end
		return new{Tobjval}( nn, ar )
	end
	function multiobj(a::multiobj,
								mobjVals::Vector{Tvals} )  where Tvals <: Number
			## create an multiobj with model of objectives types and initialised array of values of weights
		nn = length(mobjVals)
		ar = Vector{genericWeightCategory}(undef, nn )  # suppress undef first parameter for Julia 0.63
		   # with the same types in the same indices
		for i::Int64 in 1:nn 	ar[i]=mimic(a.objectives[i], mobjVals[i] ) 	end
		return new( nn, ar )
	end													##############
	function multiobj(a::multiobj,
								mobjVals::Array{Tvals} )  where Tvals <: Number
			## create an multiobj with model of objectives types and initialised array of values of weights
		nn = length(mobjVals)
		ar = Vector{genericWeightCategory}(undef, nn )  # suppress undef first parameter for Julia 0.63
		   # with the same types in the same indices
		for i::Int64 in 1:nn 	ar[i]=mimic(a.objectives[i], mobjVals[i] ) 	end
		return new( nn, ar )
	end													##############
	function multiobj(a::multiobj , b::multiobj )
			## create an multiobj with combinations of weights in a && b
		nn::Int8 = objLength(a)
		ar= Vector{ genericWeightCategory }(undef,  nn )  # a && b MUST be the same size  # suppress undef first parameter for Julia 0.63
		   # with the same types in the same indices
		for i::Int64 in 1:nn 	ar[i]=combine(a.objectives[i],b.objectives[i]) 	end
		return new(nn,ar)
	end
end

######### accessor ##############
 function getObj( mo::multiobj, numobj::Integer )
   return mo.objectives[numobj]
 end
 function setObj!( mo::multiobj, numobj::Integer, objvalue)
   old = mo.objectives[numobj] ; mo.objectives[numobj] = objvalue ; return old
 end
######### displaying a multi-objective ############
 function summarize( mo::multiobj ; log=stdout)
 	print(log,"[ ")
 	 for obj in mo.objectives
 	 	print(log,obj.value," ")
 	 end
 	 print(log,"]");  #### be carefull using it : no line break here
 end

 ######### description of a multi-objective ############
 function descr( mo::multiobj ; withName=true )
  txt::String="[ "
  if withName
 	 txt*=descr(mo.objectives[1])
 	 for i in 2:length(mo.objectives) txt*=","*descr(mo.objectives[i])  end
 	 txt*="]"
  else
 	 txt*=string(mo.objectives[1].value)
 	 for i in 2:length(mo.objectives)  txt*=","*string(mo.objectives[i].value)  end
 	 txt*="]"
  end
  return txt
 end

############## initialisation values of given types
	function defaultValue( mo::multiobj )
		return multiobj( mo.objectives )
	end
############## initialisation values of given types
	function defaultValue( objTypes::Vector{ genericWeightCategory } )
		return multiobj( objTypes )
	end
############## same types with default values ( same than preceding )
	function mimic( objTypes::Vector{ genericWeightCategory} )
	 	return multiobj( objTypes )
	end
############## same types with other values
	function mimic( objTypes::Vector{ genericWeightCategory },
								objVals::Vector{Tvals} ) where   Tvals <: Number
	 	return multiobj( objTypes, objVals )
	end

	function mimic( objTypes::Vector{ genericWeightCategory },
								objVals::Array{Tvals} ) where  Tvals <: Number
	 	return multiobj( objTypes, objVals )
	end
############## same types with other values
	function mimic( objTypes::multiobj,
									objVals::Vector{Tvals} ) where  Tvals <: Number
		 return multiobj( objTypes, objVals )
	end
############## clone
	function mimic( mo::multiobj )
	 	return multiobj( mo )
	end

############# combination ##########
function combine(a::multiobj, b::multiobj )::multiobj
 return multiobj(a , b )
end

############# lexicographic order ##########
function lexicoBetter(a::multiobj, b::multiobj )::Bool
	num::Int8 = 1 ; maxNum::Int8 = objLength(a)
	while num < maxNum && compare(a.objectives[num], b.objectives[num])==(true,true) ## sequential evaluation not needed here
		num += 1
	end
	return compare(a.objectives[num], b.objectives[num])[1]
end
############# lexicographic order cycling, beginning with k-th  ##########
function lexicoBetter(a::multiobj, b::multiobj, k::Int8 )::Bool
	cntr::Int8 = 1 ; maxNum::Int8 = objLength(a) ; num::Int8 =k
	while cntr < maxNum && compare(a.objectives[num], b.objectives[num])==(true,true) ## sequential evaluation not needed here
		cntr += 1 ; num = (num % maxNum) + 1 ;
	end
	return compare(a.objectives[num], b.objectives[num])[1]
end
############# dominance ##########
function compare(a::multiobj, b::multiobj )::Tuple{Bool,Bool}
              ## returns (a<=b, b<=a) ; hence returns  (true,true) if  a == b
 			i::Int8=1; aisbetter::Bool=true  # better and worse must be seen here as large ( _ or equal)
 			while i <= a.nb &&  aisbetter                                # a.nb is assumed == b.nb
 					aisbetter = aisbetter && compare(a.objectives[i],  b.objectives[i])[1]
 				    i += 1
 			end
 			i=1; aisworse::Bool=true
 			while i <= a.nb &&  aisworse                               # a.nb is assumed == b.nb
 					aisworse = aisworse && compare(a.objectives[i], b.objectives[i])[2]
 				    i += 1
 			end
 			return  aisbetter, aisworse
end

function dominates(a::multiobj, b::multiobj )::Bool
  cp = compare(a,b)
  return cp[1]
end

function dominatesStrictly(a::multiobj, b::multiobj )::Bool
  cp = compare(a,b)
  return cp[1] && ! cp[2]
end

######### specific dominance (one objective only -> Dijsktra) ###########
function dominates(a::multiobj, b::multiobj, objnum::Int8 )::Bool
  return compare(a.objectives[objnum],  b.objectives[objnum])[1]
end

############## some well known functions redefined     ##########
function size(mo::multiobj)::Int8  where Tobjval<:Number  return mo.nb end
function length(mo::multiobj)::Int8  where Tobjval<:Number  return mo.nb end
function objLength(mo::multiobj)::Int8  where Tobjval<:Number return mo.nb end

println("multiObjectives loaded")


""" Supertype for all generation technologies"""
abstract type Generator <: StaticInjection end
const Generators = Array{<:Generator, 1}

""" Supertype for all Hydropower generation technologies"""
abstract type HydroGen <: Generator end

"""
Supertype for all renewable generation technologies
Requires the implementation of `get_rating`and `get_power_factor` methods
"""
abstract type RenewableGen <: Generator end

""" Supertype for all Thermal generation technologies"""
abstract type ThermalGen <: Generator end

function IS.get_limits(
    valid_range::Union{NamedTuple{(:min, :max)}, NamedTuple{(:max, :min)}},
    unused::T,
) where {T <: Generator}
    # Gets min and max value defined for a field,
    # e.g. "valid_range": {"min":-1.571, "max":1.571}.
    return (min = valid_range.min, max = valid_range.max, zero = 0.0)
end

"""
Return the max active power for the Renewable Generation calculated as the rating * power_factor
"""
function get_max_active_power(d::T) where {T <: RenewableGen}
    return get_rating(d) * get_power_factor(d)
end

"""
Return the max reactive power for the Renewable Generation calculated as the rating * power_factor
"""
function get_max_reactive_power(d::T) where {T <: RenewableGen}
    return get_rating(d) * sin(acos(get_power_factor(d)))
end


mutable struct Parser
    worddict::Dict
    chardict::Dict
    model
end

function Parser()
    words = h5read(wordembeds_file, "key")
    worddict = Dict(words[i] => i for i=1:length(words))
    chardict = Dict("UNKNOWN" => 1)

    for word in words
        for c in Vector{Char}(word)
            get!(chardict, string(c), length(chardict)+1)
            get!(chardict, string(uppercase(c)), length(chardict)+1)
        end
    end
    Parser(worddict, chardict, nothing)
end

function encode(p::DepParser, words::Vector{String})
    worddict = p.worddict
    chardict = p.chardict
    unkword = worddict["UNKNOWN"]
    unkchar = chardict["UNKNOWN"]
    w = map(w -> get(worddict,lowercase(w),unkword), words)
    #cs = map(words) do w
    #    map(c -> get(chardict,string(c),unkchar), Vector{Char}(w))
    #end
    w
end

function readdata!(path::String, p::DepParser)
    datasets = Dataset[]
    words, tags = String[], String[]
    lines = open(readlines, path)
    for i = 1:length(lines)
        line = lines[i]
        if isempty(line) || i == length(lines)
            isempty(words) && continue
            w, cs = encode(ner, words)
            t = encode(ner.tagset, tags)
            push!(datasets, Dataset(w,cs,t))
            empty!(words)
            empty!(tags)
        else
            items = split(line, "\t")
            push!(words, String(items[1]))
            push!(tags, String(items[2]))
            #word = replace(word, r"[0-9]", '0')
        end
    end
    datasets
end

function setup_data(doc::Vector, worddict, chardict)
    data = []
    unkword = worddict["UNKNOWN"]
    root = worddict["PADDING"]
    for sent in doc
        w = Int[root]
        cs = Var[Var([0])]
        h = Int[0]
        for items in sent
            word = items[2]
            word0 = replace(word, r"[0-9]", '0')
            wordid = get(worddict, lowercase(word0), unkword)

            chars = Vector{Char}(word0)
            charids = map(c -> get!(chardict,c,length(chardict)+1), chars)

            head = parse(Int, items[7]) + 1
            push!(w, wordid)
            push!(cs, Var(charids))
            push!(h, head)
        end
        push!(data, (Var(w),cs,Var(h)))
    end
    data
end


"""
REINFORCE learner (Williams 1992)

Specifically for 2 actions and reward only at the end of episode!

Note
----
The policy parameters θ as well as the feature vector ϕ are
column vectors. Here we use a (states x actions) matrix for simpler access.
also, the feature vector ϕ is a one-hot vector in our case. therefore we
do not need it to keep a vector but access the entries by the
state/action pair [s_t, a_t]
"""
mutable struct REINFORCE
    nr_states::Int
    nr_actions::Int
    α::Float64
    α_eff::Float64
    γ::Float64
    γ_k::Float64
    action_selection_prob_in_episode::Array{Any,1}
    states_in_episode::Array{Any,1}
    actions_in_episode::Array{Any,1}
    θ::Array{Float64, 2}
    target::Array{Float64, 1}
    modulatingfactor_in_episode::Array{Any,1}
    do_debug_print::Bool
end
function REINFORCE(; nr_states = 10, nr_actions = 4, α = 0.1, γ = .9,
    do_debug_print = false)
    θ = zeros(Float64, nr_states, nr_actions)

    REINFORCE(nr_states, nr_actions, α, α, γ, 1., [], [], [], θ, [0.], [], do_debug_print)
end
export REINFORCE
function do_debug_print_final(learner::Union{REINFORCE, REINFORCEModulated, REINFORCEbaseline})
    println("Final:")
    for s in 1:size(learner.θ,1)
        @show learner.θ[s, :]
    end
    println("--------------")
end
""" Update: main function """
function update!(learner::Union{REINFORCE, REINFORCEModulated, REINFORCEbaseline},
                state, action, reward, p_action, modulatingfactor)
    push!(learner.action_selection_prob_in_episode, p_action)
    push!(learner.states_in_episode, state)
    push!(learner.actions_in_episode, action)
    push!(learner.modulatingfactor_in_episode, modulatingfactor)
    if reward>0
         # Backward pass to update the policy
        for k = (length(learner.states_in_episode) : -1 : 1)

            getlearningrate!(learner, learner.modulatingfactor_in_episode[k])

            gettarget!(learner, learner.states_in_episode[k], reward)

            learner.θ[learner.states_in_episode[k], learner.actions_in_episode[k]] +=
                        learner.α_eff * learner.target[1] * (1. - learner.action_selection_prob_in_episode[k])
            learner.θ[learner.states_in_episode[k], 3-learner.actions_in_episode[k]] +=
                        learner.α_eff * learner.target[1] * (0. - (1. - learner.action_selection_prob_in_episode[k]))
            learner.γ_k *= learner.γ
        end
        # Reset
        learner.γ_k = 1.
        learner.action_selection_prob_in_episode = []
        learner.states_in_episode = []
        learner.actions_in_episode = []
        learner.modulatingfactor_in_episode = []
    end
    if learner.do_debug_print
        do_debug_print_final(learner)
    end
end
export update!
function gettarget!(learner::Union{REINFORCE, REINFORCEModulated}, state, reward)
    learner.target[1] = reward * learner.γ_k  #  note: assuming a single reward, set during the forward pass above
end


using .Settings: Config

export plot

mutable struct Plotter
    o::PyObject
end
Plotter(config::Config) = Plotter(qha_plotting.Plotter(config))

@pyinterface Plotter

function plot(cfgfile)
    config = from_yaml(cfgfile)

    plotter = Plotter(config)

    desired_pressures_gpa =
        range(config["P_MIN"]; length = config["NTV"], step = config["DELTA_P"])
    config["DESIRED_PRESSURES_GPa"] = desired_pressures_gpa

    results_folder = expanduser(config["output_directory"])
    calculation_option = Dict(
        "F" => "f_tp",
        "G" => "g_tp",
        "H" => "h_tp",
        "U" => "u_tp",
        "V" => "v_tp",
        "Cv" => "cv_tp_jmolk",
        "Cp" => "cp_tp_jmolk",
        "Bt" => "bt_tp_gpa",
        "Btp" => "btp_tp",
        "Bs" => "bs_tp_gpa",
        "alpha" => "alpha_tp",
        "gamma" => "gamma_tp",
    )

    file_ftv_fitted = joinpath(results_folder, "f_tv_fitted_ev_ang3.txt")
    config["f_tv_fitted"] = file_ftv_fitted

    file_ftv_non_fitted = joinpath(results_folder, "f_tv_nonfitted_ev_ang3.txt")
    config["f_tv_non_fitted"] = file_ftv_non_fitted

    file_ptv_gpa = joinpath(results_folder, "p_tv_gpa.txt")
    config["p_tv_gpa"] = file_ptv_gpa

    plotter.fv_pv()

    for idx in config["thermodynamic_properties"]
        if idx in ("F", "G", "H", "U")
            attr_name = calculation_option[idx] * "_" * config["energy_unit"]
            file_name = attr_name * ".txt"
            file_dir = joinpath(results_folder, file_name)
            config[idx] = file_dir
            plotter.plot2file(idx)
        elseif idx == "V"
            v_ang3 = calculation_option[idx] * "_ang3"
            file_name_ang3 = v_ang3 * ".txt"
            file_dir_ang3 = joinpath(results_folder, file_name_ang3)
            config[idx] = file_dir_ang3
            plotter.plot2file(idx)
        elseif idx in ("Cv", "Cp", "Bt", "Btp", "Bs", "alpha", "gamma")
            attr_name = calculation_option[idx]
            file_name = attr_name * ".txt"
            file_dir = joinpath(results_folder, file_name)
            config[idx] = file_dir
            plotter.plot2file(idx)
        end
    end
end


using LightGraphs

import Base.isequal

# vertex, either SpeciesEMU or EMURxn
abstract type AbstractV
end

struct VSEMU <: AbstractV
    semu::SpeciesEMU
end
function Base.hash(x::VSEMU)
    return hash(string("abc",hash(x.semu)))
end

struct VRxn <: AbstractV
    r::EMURxn
end
function Base.hash(x::VRxn)
    return hash(string("cde",hash(x.r)))
end

function Base.isequal(x::AbstractV,y::AbstractV)
    #println("\n\nEquality Test for AbstractV")
    #println(typeof(x))
    #println(typeof(x))
    if( typeof(x) != typeof(y) )
        return false
    end
    if( typeof(x) == VSEMU )
        return isequal( x.semu , y.semu )
    end
    if( typeof(x) == VRxn )
        return isequal( x.r , y.r )
    end
end

mutable struct AVGraph
    g::SimpleDiGraph
    graph_db::Dict{Int64,AbstractV}
    graph_db_inv::Dict{AbstractV,Int64}
    last_index::Int64
end

function get_all_emurxns(g::AVGraph)
    vrxns = filter( x1 -> typeof(x1) == VRxn , collect( keys( g.graph_db_inv ) ) )
    rxns::Vector{EMURxn} = map( x1 -> x1.r , vrxns )
end

function get_all_semus(g::AVGraph)
    vsemus = filter( x1 -> typeof(x1) == VSEMU , collect( keys( g.graph_db_inv ) ) )
    semus::Vector{SpeciesEMU} = map( x1 -> x1.semu , vsemus )
end

function AVGraph()
    #return AVGraph(graph(Int64[],Edge{Int64}[]),Dict{Int64,AbstractV}(),Dict{AbstractV,Int64}(),0)
    return AVGraph(SimpleDiGraph{Int64}(),Dict{Int64,AbstractV}(),Dict{AbstractV,Int64}(),0)
end
function av_add_vertex(g::AVGraph,v::AbstractV)
    #println("n vertices: $(nv(g.g))")
    if(haskey(g.graph_db_inv,v))
        return
    else
        g.last_index = g.last_index+1
        add_vertex!( g.g )
        g.graph_db[g.last_index] = v
        g.graph_db_inv[v] = g.last_index
    end
end
function av_add_edge(g::AVGraph, v1::AbstractV , v2::AbstractV)
    av1 = g.graph_db_inv[v1]
    av2 = g.graph_db_inv[v2]
    add_edge!(g.g,av1,av2)
end
function split_into_connected_components(g::AVGraph)
    split_components = Vector{AVGraph}()
    wccs = weakly_connected_components( g.g )

    for wcci in wccs
        (g_i,vmap) = induced_subgraph(g.g,wcci)
        new_dict::Dict{Int64,AbstractV} = Dict{Int64,AbstractV}()
        new_dict_rev::Dict{AbstractV,Int64} = Dict{AbstractV,Int64}()
        for zi in 1:length(wcci)
            new_dict[ zi ] = g.graph_db[ vmap[zi] ]
            new_dict_rev[ g.graph_db[ vmap[zi] ] ] = zi
        end
        push!( split_components , AVGraph(g_i,new_dict,new_dict_rev,length(wcci)) )
    end

    return split_components
end

"""
create_emu_graph(rxns::Vector{EMURxn})
returns Dict{Int64,Vector{AVGraph}}

# what it does:
  1. it sorts emu rxns according to size of target EMUs
  2. it computes the connected compoments
  it then returns the connected components, sorted by size
"""
function create_emu_graph(rxns::Vector{EMURxn})

    sizes::Set{Int64} = Set{Int64}( map( x1 -> length(x1.product_emu) , collect(rxns) ) )
    sorted_rxns      = Dict{Int64,Set{EMURxn}}()

    network_components::Dict{Integer,Vector{AVGraph}} = Dict{Integer,Vector{AVGraph}}()

    for si = sizes
        rxi = filter( x1 -> (length(x1.product_emu)==si) , rxns )
        sorted_rxns[si] = Set(rxi)

        # create graphs and find connected compoments..
        #g::GenericGraph{AbstractV,Edge{AbstractV}} = graph(AbstractV[],Edge{AbstractV}[])
        #g::AVGraph(simple_graph(0),Dict{Int64,AbstractV}(),Dict{AbstractV;Int64}(),0)
        g::AVGraph = AVGraph()

        # add all rxns
        for ri in rxi
                semu_p  = pemu2semu( ri.product_emu )
                semus_r = map( xi -> remu2semu(xi) , ri.reactand_emus )

                # add vertices and edges
                vp = VSEMU(semu_p)
                # add_vertex!( g , vp )
                av_add_vertex(g,vp)
                vr = VRxn(ri)
                #add_vertex!( g , vr )
                av_add_vertex(g,vr)
                #add_edge!(g,vr,vp)
                av_add_edge(g,vr,vp)

                v_sr = map( xi -> VSEMU(xi) , semus_r )
                #for vsri in v_sr ; add_vertex!( g , vsri ) ; end
                #for vsri in v_sr ; add_edge!( g , vsri , vr ) ; end
                for vsri in v_sr ; av_add_vertex( g , vsri ) ; end
                for vsri in v_sr ; av_add_edge( g , vsri , vr ) ; end
        end

        #print(g.g)
        #return g

        # now split into connected compoments
        g_split::Vector{AVGraph} = split_into_connected_components(g)
        network_components[si] = g_split
    end
    return network_components
end

"""
get_roots_of_av_graph(g::AVGraph)
returns Tuple{Vector{SpeciesEMU},Vector{SpeciesEMU}}

returns all semu vertices with zero indegree, and the complement
"""
function get_roots_of_av_graph(g::AVGraph)
    roots::Vector{SpeciesEMU} = Vector{AbstractV}()
    non_roots::Vector{SpeciesEMU} = Vector{AbstractV}()
    for vi::AbstractV in keys(g.graph_db_inv)
        if( vi isa VSEMU )
            if( length( inneighbors( g.g , g.graph_db_inv[vi] ) )==0)
                push!(roots,vi.semu)
            else
                push!(non_roots,vi.semu)
            end
        end
    end
    return ( roots , non_roots )
end


#
# Optimizations:
#  - replace leftmap, rightmap with sorted vectors
#

###########################
# SiteOp                  #
###########################

struct SiteOp
  name::String
  site::Int
end

name(s::SiteOp) = s.name
site(s::SiteOp) = s.site

Base.show(io::IO,s::SiteOp) = print(io,"\"$(name(s))\"($(site(s)))")

Base.:(==)(s1::SiteOp,s2::SiteOp) = (s1.site==s2.site && s1.name==s2.name)

function Base.isless(s1::SiteOp,s2::SiteOp)::Bool
  if site(s1) < site(s2)
    return true
  end
  return name(s1) < name(s2)
end

###########################
# OpTerm                  #
###########################

const OpTerm = Vector{SiteOp}
mult(t1::OpTerm,t2::OpTerm) = isempty(t2) ? t1 : vcat(t1,t2)

function isfermionic(t::OpTerm,sites)::Bool
  p = +1
  for op in t
    if has_fermion_string(sites[site(op)],name(op))
      p *= -1
    end
  end
  return (p == -1)
end

###########################
# MPOTerm                 #
###########################

mutable struct MPOTerm
  coef::ComplexF64
  ops::OpTerm
end
coef(op::MPOTerm) = op.coef
ops(op::MPOTerm) = op.ops

function Base.:(==)(t1::MPOTerm,
                    t2::MPOTerm)
  return coef(t1) ≈ coef(t2) && ops(t1) == ops(t2)
end

function Base.isless(t1::MPOTerm, t2::MPOTerm)::Bool
  if !isapprox(coef(t1),coef(t2))
    ct1 = coef(t1)
    ct2 = coef(t2)
    #"lexicographic" ordering on  complex numbers
    return real(ct1) < real(ct2) ||
           (real(ct1) == real(ct2) && imag(ct1) < imag(ct2))
  end
  return ops(t1) < ops(t2)
end

function MPOTerm(c::Number,
                 op1::String,
                 i1::Int)
  return MPOTerm(convert(ComplexF64,c),[SiteOp(op1,i1)])
end

function MPOTerm(c::Number,
                 op1::String,i1::Int,
                 op2::String,i2::Int)
  return MPOTerm(convert(ComplexF64,c),[SiteOp(op1,i1),SiteOp(op2,i2)])
end

function MPOTerm(c::Number,
                 op1::String,i1::Int,
                 op2::String,i2::Int,
                 ops...)
  vop = OpTerm(undef,2+div(length(ops),2))
  vop[1] = SiteOp(op1,i1)
  vop[2] = SiteOp(op2,i2)
  for n = 1:div(length(ops),2)
    vop[2+n] = SiteOp(ops[2*n-1],ops[2*n])
  end
  return MPOTerm(convert(ComplexF64,c),vop)
end

#function MPOTerm(c::Number,
#                 ops::OpTerm)
#  return MPOTerm(convert(ComplexF64,c),ops)
#end

function Base.show(io::IO,
                   op::MPOTerm)
  c = coef(op)
  if iszero(imag(c))
    print(io,"$(real(c)) ")
  elseif iszero(real(c))
    print(io,"$(imag(c))im ")
  else
    print(io,"($c) ")
  end
  for o in ops(op)
    print(io,"\"$(name(o))\"($(site(o))) ")
  end
end

############################
## AutoMPO                 #
############################

"""
An AutoMPO stores a collection of
operator terms, to be later summed
together into an MPO by calling
the function `MPO` on the AutoMPO object.
Each term is a product of local operators
specified by names such as "Sz" or "N",
times an optional coefficient which
can be real or complex.

Which local operator names are available
is determined by the function `op`
associated with the TagType defined by
special Index tags, such as "S=1/2","S=1",
"Fermion", and "Electron".
"""
struct AutoMPO
  data::Vector{MPOTerm}
  AutoMPO(terms::Vector{MPOTerm}) = new(terms)
end

"""
    AutoMPO()

Construct an empty AutoMPO
"""
AutoMPO() = AutoMPO(Vector{MPOTerm}())

data(ampo::AutoMPO) = ampo.data

Base.:(==)(ampo1::AutoMPO,
           ampo2::AutoMPO) = data(ampo1) == data(ampo2)

Base.copy(ampo::AutoMPO) = AutoMPO(copy(data(ampo)))

Base.size(ampo::AutoMPO) = size(data(ampo))

"""
    add!(ampo::AutoMPO,
         op1::String, i1::Int)

    add!(ampo::AutoMPO,
         coef::Number,
         op1::String, i1::Int)

    add!(ampo::AutoMPO,
         op1::String, i1::Int,
         op2::String, i2::Int,
         ops...)

    add!(ampo::AutoMPO,
         coef::Number,
         op1::String, i1::Int,
         op2::String, i2::Int,
         ops...)

    +(ampo:AutoMPO, term::Tuple)

Add a single- or multi-site operator
term to the AutoMPO `ampo`. Each operator
is specified by a name (String) and a
site number (Int). The second version
accepts a real or complex coefficient.

The `+` operator version of this function
accepts a tuple with entries either
(String,Int,String,Int,...) or
(Number,String,Int,String,Int,...)
where these tuple values are the same
as valid inputs to the `add!` function.
For inputting a very large number of
terms (tuples) to an AutoMPO, consider
using the broadcasted operator `.+=`
which avoids reallocating the AutoMPO
after each addition.

# Examples
```julia
ampo = AutoMPO()

add!(ampo,"Sz",2,"Sz",3)

ampo += ("Sz",3,"Sz",4)

ampo += (0.5,"S+",4,"S-",5)

ampo .+= (0.5,"S+",5,"S-",6)
```
"""
function add!(ampo::AutoMPO,
              coef::Number,
              op::String, i::Int)
  push!(data(ampo),MPOTerm(coef,op,i))
  return
end

add!(ampo::AutoMPO,op::String, i::Int) = add!(ampo,1.0,op,i)

function add!(ampo::AutoMPO,
              coef::Number,
              op1::String, i1::Int,
              op2::String, i2::Int)
  push!(data(ampo),MPOTerm(coef,op1,i1,op2,i2))
  return
end

add!(ampo::AutoMPO,
     op1::String, i1::Int,
     op2::String, i2::Int) = add!(ampo,1.0,op1,i1,op2,i2)

function add!(ampo::AutoMPO,
              coef::Number,
              op1::String, i1::Int,
              op2::String, i2::Int,
              ops...)
  push!(data(ampo), MPOTerm(coef, op1, i1, op2, i2, ops...))
  return ampo
end

add!(ampo::AutoMPO,
     op1::String, i1::Int,
     op2::String, i2::Int,
     ops...) = add!(ampo, 1.0, op1, i1, op2, i2, ops...)

"""
    subtract!(ampo::AutoMPO,
              op1::String, i1::Int,
              op2::String, i2::Int,
              ops...)

    subtract!(ampo::AutoMPO,
              coef::Number,
              op1::String, i1::Int,
              op2::String, i2::Int,
              ops...)

Subtract a multi-site operator term
from the AutoMPO `ampo`. Each operator
is specified by a name (String) and a
site number (Int). The second version
accepts a real or complex coefficient.
"""
subtract!(ampo::AutoMPO,
          op1::String, i1::Int,
          op2::String, i2::Int,
          ops...) = add!(ampo, -1.0, op1, i1, op2, i2, ops...)

function subtract!(ampo::AutoMPO,
                   coef::Number,
                   op1::String, i1::Int,
                   op2::String, i2::Int,
                   ops...)
  push!(data(ampo), -MPOTerm(coef, op1, i1, op2, i2, ops...))
  return ampo
end

Base.:-(t::MPOTerm) = MPOTerm(-coef(t), ops(t))

function Base.:+(ampo::AutoMPO,
                 term::Tuple)
  ampo_plus_term = copy(ampo)
  add!(ampo_plus_term, term...)
  return ampo_plus_term
end

function Base.:-(ampo::AutoMPO,
                 term::Tuple)
  ampo_plus_term = copy(ampo)
  subtract!(ampo_plus_term, term...)
  return ampo_plus_term
end

#
# ampo .+= ("Sz",1) syntax using broadcasting
#

struct AutoMPOStyle <: Broadcast.BroadcastStyle end
Base.BroadcastStyle(::Type{<:AutoMPO}) = AutoMPOStyle()

struct AutoMPOAddTermStyle <: Broadcast.BroadcastStyle end

Base.broadcastable(ampo::AutoMPO) = ampo

Base.BroadcastStyle(::AutoMPOStyle,
                    ::Broadcast.Style{Tuple}) = AutoMPOAddTermStyle()

Broadcast.instantiate(bc::Broadcast.Broadcasted{AutoMPOAddTermStyle}) = bc

function Base.copyto!(ampo,
                      bc::Broadcast.Broadcasted{AutoMPOAddTermStyle,
                                                <:Any,
                                                typeof(+)})
  add!(ampo, bc.args[2]...)
  return ampo
end

#
# ampo .-= ("Sz",1) syntax using broadcasting
#

function Base.copyto!(ampo,
                      bc::Broadcast.Broadcasted{AutoMPOAddTermStyle,
                                                <:Any,
                                                typeof(-)})
  subtract!(ampo, bc.args[2]...)
  return ampo
end

function Base.show(io::IO,
                   ampo::AutoMPO)
  println(io,"AutoMPO:")
  for term in data(ampo)
    println(io,"  $term")
  end
end

##################################
# MatElem (simple sparse matrix) #
##################################

struct MatElem{T}
  row::Int
  col::Int
  val::T
end

#function Base.show(io::IO,m::MatElem)
#  print(io,"($(m.row),$(m.col),$(m.val))")
#end

function toMatrix(els::Vector{MatElem{T}})::Matrix{T} where {T}
  nr = 0
  nc = 0
  for el in els
    nr = max(nr,el.row)
    nc = max(nc,el.col)
  end
  M = zeros(T,nr,nc)
  for el in els
    M[el.row,el.col] = el.val
  end
  return M
end

function Base.:(==)(m1::MatElem{T},m2::MatElem{T})::Bool where {T}
  return (m1.row==m2.row && m1.col==m2.col && m1.val==m2.val)
end

function Base.isless(m1::MatElem{T},m2::MatElem{T})::Bool where {T}
  if m1.row != m2.row
    return m1.row < m2.row
  elseif m1.col != m2.col
    return m1.col < m2.col
  end
  return m1.val < m2.val
end

struct QNMatElem{T}
  rowqn::QN
  colqn::QN
  row::Int
  col::Int
  val::T
end

function Base.:(==)(m1::QNMatElem{T},m2::QNMatElem{T})::Bool where {T}
  return (m1.row==m2.row && m1.col==m2.col && m1.val==m2.val && m1.rowqn==m2.rowqn && m1.colqn==m2.colqn)
end

function Base.isless(m1::QNMatElem{T},m2::QNMatElem{T})::Bool where {T}
  if m1.rowqn != m2.rowqn
    return m1.rowqn < m2.rowqn
  elseif m1.colqn != m2.colqn
    return m1.colqn < m2.colqn
  elseif m1.row != m2.row
    return m1.row < m2.row
  elseif m1.col != m2.col
    return m1.col < m2.col
  end
  return m1.val < m2.val
end

function posInLink!(linkmap::Vector{OpTerm},
                    op::OpTerm)::Int
  isempty(op) && return -1
  for n=1:length(linkmap)
    (linkmap[n]==op) && return n
  end
  push!(linkmap,op)
  return length(linkmap)
end

function determineValType(terms::Vector{MPOTerm})
  for t in terms
    (!isreal(coef(t))) && return ComplexF64
  end
  return Float64
end

function computeSiteProd(sites,
                         ops::OpTerm)::ITensor
  i = ops[1].site
  T = op(sites[i],ops[1].name)
  for j=2:length(ops)
    (ops[j].site != i) && error("Mismatch of site number in computeSiteProd")
    opj = op(sites[i],ops[j].name)
    T = product(T, opj)
  end
  return T
end

function remove_dups!(v::Vector{T}) where {T}
  N = length(v)
  (N==0) && return
  sort!(v)
  n = 1
  u = 2
  while u <= N
    while u < N && v[u]==v[n]
      u += 1
    end
    if v[u] != v[n]
      v[n+1] = v[u]
      n += 1
    end
    u += 1
  end
  resize!(v,n)
  return
end #remove_dups!

function svdMPO(ampo::AutoMPO,
                sites;
                kwargs...)::MPO

  mindim::Int = get(kwargs,:mindim,1)
  maxdim::Int = get(kwargs,:maxdim,10000)
  cutoff::Float64 = get(kwargs,:cutoff,1E-13)

  N = length(sites)

  ValType = determineValType(data(ampo))

  Vs = [Matrix{ValType}(undef,1,1) for n=1:N]
  tempMPO = [MatElem{MPOTerm}[] for n=1:N]

  crosses_bond(t::MPOTerm,n::Int) = (ops(t)[1].site <= n <= ops(t)[end].site)

  rightmap = OpTerm[]
  next_rightmap = OpTerm[]

  for n=1:N

    leftbond_coefs = MatElem{ValType}[]

    leftmap = OpTerm[]
    for term in data(ampo)
      crosses_bond(term,n) || continue

      left::OpTerm   = filter(t->(t.site < n),ops(term))
      onsite::OpTerm = filter(t->(t.site == n),ops(term))
      right::OpTerm  = filter(t->(t.site > n),ops(term))

      bond_row = -1
      bond_col = -1
      if !isempty(left)
        bond_row = posInLink!(leftmap,left)
        bond_col = posInLink!(rightmap,mult(onsite,right))
        bond_coef = convert(ValType,coef(term))
        push!(leftbond_coefs,MatElem(bond_row,bond_col,bond_coef))
      end

      A_row = bond_col
      A_col = posInLink!(next_rightmap,right)
      site_coef = 1.0+0.0im
      if A_row == -1
        site_coef = coef(term)
      end
      if isempty(onsite)
        if isfermionic(right,sites)
          #println("Putting F on site $n")
          push!(onsite,SiteOp("F",n))
        else
          #println("Putting Id on site $n")
          push!(onsite,SiteOp("Id",n))
        end
      end
      el = MatElem(A_row,A_col,MPOTerm(site_coef,onsite))
      push!(tempMPO[n],el)
    end
    rightmap = next_rightmap
    next_rightmap = OpTerm[]

    remove_dups!(tempMPO[n])

    if n > 1 && !isempty(leftbond_coefs)
      M = toMatrix(leftbond_coefs)
      U,S,V = svd(M)
      P = S.^2
      truncate!(P;maxdim=maxdim,cutoff=cutoff,mindim=mindim)
      tdim = length(P)
      nc = size(M,2)
      Vs[n-1] = Matrix{ValType}(V[1:nc,1:tdim])
    end

  end

  llinks = Vector{Index{Int}}(undef, N+1)
  llinks[1] = Index(2, "Link,l=0")

  H = MPO(sites)

  # Constants which define MPO start/end scheme
  rowShift = 2
  colShift = 2
  startState = 2
  endState = 1

  for n=1:N
    VL = Matrix{ValType}(undef,1,1)
    if n > 1
      VL = Vs[n-1]
    end
    VR = Vs[n]
    tdim = size(VR,2)

    llinks[n+1] = Index(2+tdim,"Link,l=$n")

    finalMPO = Dict{OpTerm,Matrix{ValType}}()

    ll = llinks[n]
    rl = llinks[n+1]

    idTerm = [SiteOp("Id",n)]
    finalMPO[idTerm] = zeros(ValType,dim(ll),dim(rl))
    idM = finalMPO[idTerm]
    idM[1,1] = 1.0
    idM[2,2] = 1.0

    defaultMat() = zeros(ValType,dim(ll),dim(rl))

    for el in tempMPO[n]
      A_row = el.row
      A_col = el.col
      t = el.val
      (abs(coef(t)) > eps()) || continue

      M = get!(finalMPO,ops(t),defaultMat())

      ct = convert(ValType,coef(t))
      if A_row==-1 && A_col==-1 #onsite term
        M[startState,endState] += ct
      elseif A_row==-1 #term starting on site n
        for c=1:size(VR,2)
          z = ct*VR[A_col,c]
          M[startState,colShift+c] += z
        end
      elseif A_col==-1 #term ending on site n
        for r=1:size(VL,2)
          z = ct*conj(VL[A_row,r])
          M[rowShift+r,endState] += z
        end
      else
        for r=1:size(VL,2),c=1:size(VR,2)
          z = ct*conj(VL[A_row,r])*VR[A_col,c]
          M[rowShift+r,colShift+c] += z
        end
      end
    end

    s = sites[n]
    #H[n] = emptyITensor(dag(s),s',ll,rl)
    H[n] = emptyITensor(Any)
    for (op,M) in finalMPO
      T = itensor(M,ll,rl)
      H[n] += T*computeSiteProd(sites,op)
    end

  end

  L = emptyITensor(llinks[1])
  L[startState] = 1.0

  R = emptyITensor(llinks[N+1])
  R[endState] = 1.0

  H[1] *= L
  H[N] *= R

  return H
end #svdMPO

function qn_svdMPO(ampo::AutoMPO,
                   sites;
                   kwargs...)::MPO

  mindim::Int = get(kwargs,:mindim,1)
  maxdim::Int = get(kwargs,:maxdim,10000)
  cutoff::Float64 = get(kwargs,:cutoff,1E-13)

  N = length(sites)

  ValType = determineValType(data(ampo))

  Vs = [Dict{QN,Matrix{ValType}}() for n=1:N+1]
  tempMPO = [QNMatElem{MPOTerm}[] for n=1:N]

  crosses_bond(t::MPOTerm,n::Int) = (ops(t)[1].site <= n <= ops(t)[end].site)

  rightmap = Dict{QN,Vector{OpTerm}}()
  next_rightmap = Dict{QN,Vector{OpTerm}}()

  for n=1:N

    leftbond_coefs = Dict{QN,Vector{MatElem{ValType}}}()

    leftmap = Dict{QN,Vector{OpTerm}}()
    for term in data(ampo)
      crosses_bond(term,n) || continue

      left::OpTerm   = filter(t->(t.site < n),ops(term))
      onsite::OpTerm = filter(t->(t.site == n),ops(term))
      right::OpTerm  = filter(t->(t.site > n),ops(term))

      function calcQN(term::OpTerm)
        q = QN()
        for st in term
          op_tensor = op(sites[site(st)],name(st))
          q -= flux(op_tensor)
        end
        return q
      end
      lqn = calcQN(left)
      sqn = calcQN(onsite)

      q_leftmap = get!(leftmap,lqn,OpTerm[])
      q_rightmap = get!(rightmap,lqn,OpTerm[])

      bond_row = -1
      bond_col = -1
      if !isempty(left)
        bond_row = posInLink!(q_leftmap,left)
        bond_col = posInLink!(q_rightmap,mult(onsite,right))
        bond_coef = convert(ValType,coef(term))
        q_leftbond_coefs = get!(leftbond_coefs,lqn,MatElem{ValType}[])
        push!(q_leftbond_coefs,MatElem(bond_row,bond_col,bond_coef))
      end

      rqn = sqn+lqn
      q_next_rightmap = get!(next_rightmap,rqn,OpTerm[])
      A_row = bond_col
      A_col = posInLink!(q_next_rightmap,right)
      site_coef = 1.0+0.0im
      if A_row == -1
        site_coef = coef(term)
      end
      if isempty(onsite)
        if isfermionic(right,sites)
          push!(onsite,SiteOp("F",n))
        else
          push!(onsite,SiteOp("Id",n))
        end
      end
      el = QNMatElem(lqn,rqn,A_row,A_col,MPOTerm(site_coef,onsite))
      push!(tempMPO[n],el)
    end
    rightmap = next_rightmap
    next_rightmap = Dict{QN,Vector{OpTerm}}()

    remove_dups!(tempMPO[n])

    if n > 1 && !isempty(leftbond_coefs)
      for (q,mat) in leftbond_coefs
        M = toMatrix(mat)
        U,S,V = svd(M)
        P = S.^2
        truncate!(P;maxdim=maxdim,cutoff=cutoff,mindim=mindim)
        tdim = length(P)
        nc = size(M,2)
        Vs[n][q] = Matrix{ValType}(V[1:nc,1:tdim])
      end
    end
  end

  #
  # Make MPO link indices
  #
  d0 = 2
  llinks = Vector{QNIndex}(undef,N+1)
  llinks[1] = Index(QN()=>d0;tags="Link,l=0")
  for n=1:N
    qi = Vector{Pair{QN,Int}}()
    if !haskey(Vs[n+1],QN())
      # Make sure QN=zero is first in list of sectors
      push!(qi,QN()=>d0)
    end
    for (q,Vq) in Vs[n+1]
      cols = size(Vq,2)
      if q==QN()
        # Make sure QN=zero is first in list of sectors
        insert!(qi,1,q=>d0+cols)
      else
        push!(qi,q=>cols)
      end
    end
    llinks[n+1] = Index(qi...;tags="Link,l=$n")
  end

  H = MPO(N)

  # Constants which define MPO start/end scheme
  startState = 2
  endState = 1

  for n=1:N
    finalMPO = Dict{Tuple{QN,OpTerm},Matrix{ValType}}()

    ll = llinks[n]
    rl = llinks[n+1]

    function defaultMat(ll,rl,lqn,rqn)
      ldim = qnblockdim(ll,lqn)
      rdim = qnblockdim(rl,rqn)
      return zeros(ValType,ldim,rdim)
    end

    idTerm = [SiteOp("Id",n)]
    finalMPO[(QN(),idTerm)] = defaultMat(ll,rl,QN(),QN())
    idM = finalMPO[(QN(),idTerm)]
    idM[1,1] = 1.0
    idM[2,2] = 1.0

    for el in tempMPO[n]
      t = el.val
      (abs(coef(t)) > eps()) || continue
      A_row = el.row
      A_col = el.col

      M = get!(finalMPO,(el.rowqn,ops(t)),defaultMat(ll,rl,el.rowqn,el.colqn))

      # rowShift and colShift account for
      # special entries in the zero-QN sector
      # of the MPO
      rowShift = (el.rowqn == QN()) ? 2 : 0
      colShift = (el.colqn == QN()) ? 2 : 0

      ct = convert(ValType,coef(t))
      if A_row==-1 && A_col==-1 #onsite term
        M[startState,endState] += ct
      elseif A_row==-1 #term starting on site n
        VR = Vs[n+1][el.colqn]
        for c=1:size(VR,2)
          z = ct*VR[A_col,c]
          M[startState,colShift+c] += z
        end
      elseif A_col==-1 #term ending on site n
        VL = Vs[n][el.rowqn]
        for r=1:size(VL,2)
          z = ct*conj(VL[A_row,r])
          M[rowShift+r,endState] += z
        end
      else
        VL = Vs[n][el.rowqn]
        VR = Vs[n+1][el.colqn]
        for r=1:size(VL,2),c=1:size(VR,2)
          z = ct*conj(VL[A_row,r])*VR[A_col,c]
          M[rowShift+r,colShift+c] += z
        end
      end
    end

    s = sites[n]
    #H[n] = emptyITensor(dag(s),s',dag(ll),rl)
    H[n] = emptyITensor(Any)
    for (q_op,M) in finalMPO
      op_prod = q_op[2]
      Op = computeSiteProd(sites,op_prod)

      rq = q_op[1]
      sq = flux(Op)
      cq = rq-sq

      rn = qnblocknum(ll,rq)
      cn = qnblocknum(rl,cq)

      #TODO: wrap following 3 lines into a function
      block = (rn,cn)
      T = BlockSparseTensor(ValType,[block],IndexSet(dag(ll),rl))
      blockview(T,block) .= M

      IT = itensor(T)
      H[n] += IT * Op
    end

  end

  L = emptyITensor(llinks[1])
  L[startState] = 1.0

  R = emptyITensor(dag(llinks[N+1]))
  R[endState] = 1.0

  H[1] *= L
  H[N] *= R

  return H
end #qn_svdMPO

function sorteachterm(ampo::AutoMPO, sites)
  ampo = copy(ampo)
  isless_site(o1::SiteOp, o2::SiteOp) = site(o1) < site(o2)
  for t in data(ampo)
    Nt = length(t.ops)

    # Sort operators in t by site order,
    # and keep the permutation used, perm, for analysis below
    perm = Vector{Int}(undef,Nt)
    sortperm!(perm,t.ops, alg=InsertionSort, lt=isless_site)

    t.ops = t.ops[perm]

    # Identify fermionic operators,
    # zeroing perm for bosonic operators,
    # and inserting string "F" operators
    parity = +1
    for n=Nt:-1:1
      fermionic = has_fermion_string(sites[site(t.ops[n])],name(t.ops[n]))
      if parity == -1
        # Put Jordan-Wigner string emanating
        # from fermionic operators to the right
        # (Remaining F operators will be put in by svdMPO)
        t.ops[n] = SiteOp("$(name(t.ops[n]))*F",site(t.ops[n]))
      end
      if fermionic
        parity = -parity
      else
        # Ignore bosonic operators in perm
        # by zeroing corresponding entries
        perm[n] = 0
      end
    end
    if parity == -1
      error("Parity-odd fermionic terms not yet supported by AutoMPO")
    end
    # Keep only fermionic op positions (non-zero entries)
    filter!(!iszero,perm)
    # Account for anti-commuting, fermionic operators
    # during above sort; put resulting sign into coef

    t.coef *= parity_sign(perm)
  end
  return ampo
end

"""
    MPO(ampo::AutoMPO,sites::Vector{<:Index};kwargs...)

Convert an AutoMPO object `ampo` to an
MPO, with indices given by `sites`. The
resulting MPO will have the indices
`sites[1], sites[1]', sites[2], sites[2]'`
etc. The conversion is done by an algorithm
that compresses the MPO resulting from adding
the AutoMPO terms together, often achieving
the minimum possible bond dimension.

# Examples
```julia
ampo = AutoMPO()
ampo += ("Sz",1,"Sz",2)
ampo += ("Sz",2,"Sz",3)
ampo += ("Sz",3,"Sz",4)

sites = siteinds("S=1/2",4)
H = MPO(ampo,sites)
```
"""
function MPO(ampo::AutoMPO,
             sites::Vector{<:Index};
             kwargs...)::MPO
  ampo = sorteachterm(ampo,sites)
  #for t in data(ampo)
  #  @show t
  #end
  if hasqns(sites[1])
    return qn_svdMPO(ampo,sites;kwargs...)
  end
  return svdMPO(ampo,sites;kwargs...)
end

@deprecate toMPO(args...; kwargs...) MPO(args...; kwargs...)


using MPI
ENV["JULIA_MPI_BINARY"]="system"
using VertexEliminationOrder
using ArgParse

function parse_commandline(ARGS)
    s = ArgParseSettings()
    @add_arg_table! s begin
        "--seed"
            help = "Base seed to use for sampling flow-cutter inputs."
            default = 42
            arg_type = Int64
        "--seed_bound"
            help = "Minimum difference of seed between MPI processes."
            default = 1000
            arg_type = Int64
        "--time"
            help = "The number of seconds to run flow cutter for."
            default = 30
            arg_type = Int64
        "--graph"
            help = "The name of the graph file .gr to open."
            default = "sycamore_53_20.gr"
            arg_type = String
        "--dir_graphfiles"
            help = "Path to your graph files .gr to read them."
            default = "$(@__DIR__)/../test/example_graphs/"
            arg_type = String
        "--nsample"
            help = "The number of sample of flow cutter random inputs per execution."
            default = 20
            arg_type = Int64

        "--sample_augment"
            help = "The number of execution needed before increasing nsample."
            default = 8
            arg_type = Int64
        "--max_imbalances"
            nargs ='*'
            help = "The max imabalances to select best cut from flow-cutter."
            default = [1.0, 0.8, 0.6]
            arg_type = Float64
    end

    return parse_args(ARGS, s)
end

parsed_args = parse_commandline(ARGS)

graph_name = parsed_args["graph"]
dirgraphfiles = parsed_args["dir_graphfiles"]
duration = parsed_args["time"]
max_imbalances = parsed_args["max_imbalances"]
base_seed = parsed_args["seed"]
seed_diff = parsed_args["seed_bound"]
base_nsample = parsed_args["nsample"]
sample_augment = parsed_args["sample_augment"]

println("- - - args - - -")
println("duration = $duration")
println("max_imbalances = $max_imbalances")
println("base_seed = $base_seed")
println("seed_diff = $seed_diff")
println("base_nsample = $base_nsample")
println("sample_augment = $sample_augment")
println("- - - result - - -")

MPI.Init()
comm = MPI.COMM_WORLD
my_rank = MPI.Comm_rank(comm)
root = 0
print("rank $my_rank has $(Threads.nthreads()) threads \n")

g = graph_from_gr(joinpath(dirgraphfiles, graph_name))
order, tw = order_tw_by_dissections_simple(
    g,
    duration;
    max_imbalances=max_imbalances ,
    seed=base_seed + my_rank * seed_diff,
    nsample=base_nsample,
    sample_augment=sample_augment
)
MPI.Barrier(comm)

"""returns the better (rank, treewidth) pair"""
function best(order1::Tuple{Int, Int}, order2::Tuple{Int, Int})
    order1[2] < order2[2] ? order1 : order2
end

rank_with_best, best_tw = MPI.Allreduce((my_rank, tw), best, comm)
MPI.Barrier(comm)

if my_rank == root
    best_order = similar(order)
    rreq = MPI.Irecv!(best_order, rank_with_best, rank_with_best, comm)
end

if my_rank == rank_with_best
    print("rank $my_rank sending elimination order to root \n")
    sreq = MPI.Isend(order, root, my_rank, comm)
end

MPI.Barrier(comm)

if my_rank == root
    @show best_tw
    @show best_order
end


import Base.Cartesian.@ntuple

nparticles(p) = length(p)
nparticles(p::Type{<:AbstractParticles{T,N}}) where {T,N} = N
nparticles(p::AbstractParticles{T,N}) where {T,N} = N
nparticles(p::ParticleArray) = nparticles(eltype(p))
nparticles(p::Type{<:ParticleArray}) = nparticles(eltype(p))

particletype(p::AbstractParticles) = typeof(p)
particletype(::Type{P}) where P <: AbstractParticles = P
particletype(p::AbstractArray{<:AbstractParticles}) = eltype(p)

particleeltype(::AbstractParticles{T,N}) where {T,N} = T
particleeltype(::AbstractArray{<:AbstractParticles{T,N}}) where {T,N} = T

vecindex(p,i) = getindex(p,i)
vecindex(p::AbstractParticles,i) = getindex(p.particles,i)
vecindex(p::ParticleArray,i) = vecindex.(p,i)
vecindex(p::NamedTuple,i) = (; Pair.(keys(p), ntuple(j->arggetter(i,p[j]), fieldcount(typeof(p))))...)

function indexof_particles(args)
    inds = findall(a-> a <: SomeKindOfParticles, args)
    inds === nothing && throw(ArgumentError("At least one argument should be <: AbstractParticles. If particles appear nested as fields inside an argument, see `with_workspace` and `Workspace`"))
    all(nparticles(a) == nparticles(args[inds[1]]) for a in args[inds]) || throw(ArgumentError("All p::Particles must have the same number of particles."))
    (inds...,)
    # TODO: test all same number of particles
end

function arggetter(i,a::Union{SomeKindOfParticles, NamedTuple})
    vecindex(a,i)
end

arggetter(i,a) = a

"""
    @bymap f(p, args...)

Call `f` with particles or vectors of particles by using `map`. This can be utilized if registering `f` using [`register_primitive`](@ref) fails. See also [`Workspace`](@ref) if `bymap` fails.
"""
macro bymap(ex)
    @capture(ex, f_(args__)) || error("expected a function call")
    quote
        bymap($(esc(f)),$(esc.(args)...))
    end
end

"""
    bymap(f, args...)

Uncertainty propagation using the `map` function.

Call `f` with particles or vectors of particles by using `map`. This can be utilized if registering `f` using [`register_primitive`](@ref) fails. See also [`Workspace`](@ref) if `bymap` fails.
"""
function bymap(f::F, args...) where F
    inds = indexof_particles(typeof.(args))
    T,N,PT = particletypetuple(args[first(inds)])
    individuals = map(1:N) do i
        argsi = ntuple(j->arggetter(i,args[j]), length(args))
        f(argsi...)
    end
    PTNT = PT{eltype(eltype(individuals)),N}
    if (eltype(individuals) <: AbstractArray{TT,0} where TT) || eltype(individuals) <: Number
        PTNT(individuals)
    elseif eltype(individuals) <: AbstractArray{TT,1} where TT
        PTNT(copy(reduce(hcat,individuals)'))
    elseif eltype(individuals) <: AbstractArray{TT,2} where TT
        # @show PT{eltype(individuals),N}
        reshape(PTNT(copy(reduce(hcat,vec.(individuals))')), size(individuals[1],1),size(individuals[1],2))::Matrix{PTNT}
    else
        error("Output with dimension >2 is currently not supported by `bymap`. Consider if `ℝⁿ2ℝⁿ_function($(f), $(args...))` works for your use case.")
    end
end

"""
Distributed uncertainty propagation using the `pmap` function. See [`bymap`](@ref) for more details.
"""
function bypmap(f::F, args...) where F
    inds = indexof_particles(typeof.(args))
    T,N,PT = particletypetuple(args[first(inds)])
    individuals = map(1:N) do i
        argsi = ntuple(j->arggetter(i,args[j]), length(args))
        f(argsi...)
    end
    PTNT = PT{eltype(eltype(individuals)),N}
    if (eltype(individuals) <: AbstractArray{TT,0} where TT) || eltype(individuals) <: Number
        PTNT(individuals)
    elseif eltype(individuals) <: AbstractArray{TT,1} where TT
        PTNT(copy(reduce(hcat,individuals)'))
    elseif eltype(individuals) <: AbstractArray{TT,2} where TT
        # @show PT{eltype(individuals),N}
        reshape(PTNT(copy(reduce(hcat,vec.(individuals))')), size(individuals[1],1),size(individuals[1],2))::Matrix{PTNT}
    else
        error("Output with dimension >2 is currently not supported by `bymap`. Consider if `ℝⁿ2ℝⁿ_function($(f), $(args...))` works for your use case.")
    end
end

"""
    @bypmap f(p, args...)

Call `f` with particles or vectors of particles by using parallel `pmap`. This can be utilized if registering `f` using [`register_primitive`](@ref) fails. See also [`Workspace`](@ref) if `bymap` fails.
"""
macro bypmap(ex)
    @capture(ex, f_(args__)) || error("expected a function call")
    quote
        bypmap($(esc(f)),$(esc.(args)...))
    end
end

"""
    @prob a < b

Calculate the probability that an event on any of the forms `a < b, a > b, a <= b, a >= b` occurs, where `a` and/or `b` are of type `AbstractParticles`.
"""
macro prob(ex)
    ex.head == :call && ex.args[1] ∈ (:<,:>,:<=,:>=) || error("Expected an expression on any of the forms `a < b, a > b, a <= b, a >= b`")
    op = ex.args[1]
    a  = ex.args[2]
    b  = ex.args[3]
    quote
        mean($op.(MonteCarloMeasurements.maybe_particles($(esc(a))), MonteCarloMeasurements.maybe_particles($(esc(b)))))
    end
end


#=
Support same input format as asap2 sorting:
    - 2 records per core
    - Sort high to low: pass low keys.
    - Record is 5 words key, 45 word payload
    - 0x2A followed by 0 is the flush signal, which arrives between records.
    -- Treat any non-0 as flush.
    -- Post-flush value is the number of records to pass; 0 for the first core,
       but will generally be >0 within the chain.
    -- The final core will supress the flush and count signals, and just transmits
       output records.
    - 0 is the non-flush signal, indicates a new record to process.
AG layout:
    - AG0 points to low record key's low word, 5 words, stride backwards.
    - AG1 points to low record, 50 words, stride forward.
    - AG2 points to high record, 50 words, stride forward.
Storage:
    - 0: 1 if this is the first core in the chain, read single input.
    - 1: 1 if this is the last core in the chain, send to single output, dont pass count.
Set signed vs. unsigned key to configurable.
=#

function initialize_snakesort(core)
    # Initialize address generators
    AsapSim.set_stride!(core.address_generators[1], 0xFF)
    AsapSim.set_stride!(core.address_generators[2], 1)
    AsapSim.set_stride!(core.address_generators[3], 1)

    # Initialize CXMASK0
    AsapSim.set!(core.condexec[1], 1 << 2)
end

# Helper function for testing. Make a 50 word record with the given key.
function makerecord(key; zerofirst = false)
    if zerofirst
        return vcat(Int16.([0, 0,0,0,0,key]), rand(Int16, 45))
    else
        return vcat(Int16.([0,0,0,0,key]), rand(Int16, 45))
    end
end

#Function(Snakesort)
@asap4asm function snakesort(firstcore, lastcore, lastcore_flush)
    new_key_4 = Loc(AsapSim.DMEM, 1)
    new_key_3 = Loc(AsapSim.DMEM, 2)
    new_key_2 = Loc(AsapSim.DMEM, 3)
    new_key_1 = Loc(AsapSim.DMEM, 4)
    new_key_0 = Loc(AsapSim.DMEM, 5)
    num_records_to_pass = Loc(AsapSim.DMEM, 8)

    # Record 0 stored between 20 and 69
    # Key is 5 word range starting from low word of key and going backward.
    # So 24 down to 20.
    ag_config_full_record_0 = ((69 << 8) + 20)
    ag_config_key_0         = ((20 << 8 + 24))
    local_key_0_4 = Loc(AsapSim.DMEM, 20)
    local_key_0_3 = Loc(AsapSim.DMEM, 21)
    local_key_0_2 = Loc(AsapSim.DMEM, 22)
    local_key_0_1 = Loc(AsapSim.DMEM, 23)
    local_key_0_0 = Loc(AsapSim.DMEM, 24)

    # Record 1 stored between 70 and 119
    ag_config_full_record_1 = ((119 << 8) + 70)
    ag_config_key_1         = ((70 << 8) + 74)
    local_key_1_4 = Loc(AsapSim.DMEM, 70)
    local_key_1_3 = Loc(AsapSim.DMEM, 71)
    local_key_1_2 = Loc(AsapSim.DMEM, 72)
    local_key_1_1 = Loc(AsapSim.DMEM, 73)
    local_key_1_0 = Loc(AsapSim.DMEM, 74)

    low_key_ptr_pi  = :ag0pi
    low_key_ptr     = :ag0
    low_key_ptr_cfg = Loc(AsapSim.AG_START, 0)

    low_record_ptr_pi  = :ag1pi
    low_record_ptr     = :ag1
    low_record_ptr_cfg = Loc(AsapSim.AG_START, 1)

    high_record_ptr_pi  = :ag2pi
    high_record_ptr     = :ag
    high_record_ptr_cfg = Loc(AsapSim.AG_START, 2)

    # Begin main loop.

    # Start a new record batch, after reset.
    @label start
        # Peek at ibuf, check if the end of block code (any non-0) is found,
        #  and go to flush code if so.
        MOVE(null, ibuf0)
        BR(flush_0_local_record_with_count, Z, neg)
        # Otherwise not flushing

        # Reuse the storage to low function.  Cycle low pointer between
        #  memory blocks.
        MOVI(low_record_ptr_cfg, ag_config_full_record_0)
        # Set high to the same location; this is used if a flush signal is
        #  seen before storing another record.
        MOVI(high_record_ptr_cfg, ag_config_full_record_0, nop3) #TODO: nop2 this
        # Load an initial record.
        BRL(store_new_record_at_low, j)

        # Read ibuf, check if the end of block code (any non-0) is found,
        #  and go to flush code if so.
        MOVE(null, ibuf0)
        BR(flush_1_local_record_with_count, Z, neg)

        # Otherwise not flushing, so load a second record.
        # Set low pointer to next block.
        MOVI(low_record_ptr_cfg, ag_config_full_record_1, nop3)
        BRL(store_new_record_at_low, j)

        # Find which record is lower to set proper pointers.
        # This function jumps to get_new_key when done.
        BRL(compare_local_records)

    @label get_new_key
        # Peek at ibuf, check if the end of block code is found.
        # 0 means not end of block; non-zero is end of block.
        # Pass the control word to the next core either way.
        # Last core will not write out if not passing flush signals.
        if(lastcore == 1 && lastcore_flush == 0)
            MOVE(null, ibuf0)
        else
            MOVE(output[0], ibuf0)
        end
        BR(flush_2_local_record, Z, neg)

        # Get the next key.
        MOVE(new_key_4, ibuf0)
        MOVE(new_key_3, ibuf0)
        MOVE(new_key_2, ibuf0)
        MOVE(new_key_1, ibuf0)
        MOVE(new_key_0, ibuf0)

        # Compare to the lowest local key.
        # TODO: fast compare of high words and jump if new key is less, maybe.
        SUBU(null, low_key_ptr_pi, bypass[1])
        SUBCU(null, low_key_ptr_pi, bypass[3])
        SUBCU(null, low_key_ptr_pi, new_key_2)
        SUBCU(null, low_key_ptr_pi, new_key_3)
        #if(Unsigned0_Signed1){
        SUBC(null, low_key_ptr_pi, new_key_4, cxs0) #Set if carry if true.
        #}else{
        #SUBCU(null, low_key_ptr_pi, new_key_4, cxs0) #Set if carry if true.
        #}
        # If carry, new key is higher.

        # If new key lower or equal, pass it with record to output.
        # This function jumps to get_new_key when done.
        BRL(pass_input_to_output, C, neg)

        # If new key higher, send the local low record to output.
        # Store the new key and record at the previous low location.
        # Key.
        MOVE(output[0], low_record_ptr)
        MOVE(low_record_ptr_pi, new_key_4)
        MOVE(output[0], low_record_ptr)
        MOVE(low_record_ptr_pi, new_key_3)
        MOVE(output[0], low_record_ptr)
        MOVE(low_record_ptr_pi, new_key_2)
        MOVE(output[0], low_record_ptr)
        MOVE(low_record_ptr_pi, new_key_1)
        MOVE(output[0], low_record_ptr)
        MOVE(low_record_ptr_pi, new_key_0)

        # Rest of upper half.
        # 2x10=20
        RPT(10)
            MOVE(output[0], low_record_ptr)
            MOVE(low_record_ptr_pi, ibuf0)
            MOVE(output[0], low_record_ptr)
            MOVE(low_record_ptr_pi, ibuf0)
        END_RPT()

        # Lower half.
        # 2x12=24, add one extra at end for 25.
        RPT(12)
            if lastcore
                MOVE(output[0], low_record_ptr)
            else
                MOVE(output[1], low_record_ptr)
            end

            if firstcore
                MOVE(low_record_ptr_pi, ibuf0)
            else
                MOVE(low_record_ptr_pi, ibuf1)
            end

            if lastcore
                MOVE(output[0], low_record_ptr)
            else
                MOVE(output[1], low_record_ptr)
            end

            if firstcore
                MOVE(low_record_ptr_pi, ibuf0)
            else
                MOVE(low_record_ptr_pi, ibuf1)
            end
        END_RPT()
        # Copy of above.
        if lastcore
            MOVE(output[0], low_record_ptr)
        else
            MOVE(output[1], low_record_ptr)
        end

        if firstcore
            MOVE(low_record_ptr_pi, ibuf0)
        else
            MOVE(low_record_ptr_pi, ibuf1)
        end

    # Compare the two stored records to find which is now lower.
    # This block jumps to get_new_key when done.
    # Shared with startup code.
    @label compare_local_records
        # Check the keys against each other.
        SUBU(null, local_key_0_0, local_key_1_0)
        SUBCU(null, local_key_0_1, local_key_1_1)
        SUBCU(null, local_key_0_2, local_key_1_2)
        SUBCU(null, local_key_0_3, local_key_1_3)
        #if(Unsigned0_Signed1){
        SUBC(null, local_key_0_4, local_key_1_4, cxs0) #Set if carry is true.
        #}else{
        #SUBCU(null, local_key_0_4, local_key_1_4, cxs0) #Set if carry is true.
        #}

        # If carry set, key 0 is lower.
        MOVI(low_record_ptr_cfg, ag_config_full_record_0, cxt0)
        MOVI(low_key_ptr_cfg, ag_config_key_0, cxt0)
        MOVI(high_record_ptr_cfg, ag_config_full_record_1, cxt0)

        # Otherwise key 1 is lower.
        MOVI(low_record_ptr_cfg, ag_config_full_record_1, cxf0)
        MOVI(low_key_ptr_cfg, ag_config_key_1, cxf0)
        MOVI(high_record_ptr_cfg, ag_config_full_record_0, cxf0)

        # Ready to get the next record.
        BRL(get_new_key)

    @label pass_input_to_output
        MOVE(output[0], new_key_4)
        MOVE(output[0], new_key_3)
        MOVE(output[0], new_key_2)
        MOVE(output[0], new_key_1)
        MOVE(output[0], new_key_0)
        # 5x4=20
        RPT(5)
            MOVE(output[0], ibuf0)
            MOVE(output[0], ibuf0)
            MOVE(output[0], ibuf0)
            MOVE(output[0], ibuf0)
        END_RPT()
        # 5x5=25
        RPT(5)
            if firstcore && !lastcore
                MOVE(output[1], ibuf0)
                MOVE(output[1], ibuf0)
                MOVE(output[1], ibuf0)
                MOVE(output[1], ibuf0)
                MOVE(output[1], ibuf0)
            elseif (firstcore && lastcore)
                MOVE(output[0], ibuf0)
                MOVE(output[0], ibuf0)
                MOVE(output[0], ibuf0)
                MOVE(output[0], ibuf0)
                MOVE(output[0], ibuf0)
            elseif (!firstcore && lastcore)
                MOVE(output[0], ibuf1)
                MOVE(output[0], ibuf1)
                MOVE(output[0], ibuf1)
                MOVE(output[0], ibuf1)
                MOVE(output[0], ibuf1)
            else
                MOVE(output[1], ibuf1)
                MOVE(output[1], ibuf1)
                MOVE(output[1], ibuf1)
                MOVE(output[1], ibuf1)
                MOVE(output[1], ibuf1)
            end
        END_RPT()
        BRL(get_new_key)

    @label store_new_record_at_low
        # MOVE(low_record_ptr_pi, ibuf0)
        # MOVE(low_record_ptr_pi, ibuf0)
        # MOVE(low_record_ptr_pi, ibuf0)
        # MOVE(low_record_ptr_pi, ibuf0)
        # MOVE(low_record_ptr_pi, ibuf0)
        RPT(25, nop3)
            MOVE(low_record_ptr_pi, ibuf0)
        END_RPT()
        RPT(25, nop3)
            if firstcore
                MOVE(low_record_ptr_pi, ibuf0)
            else
                MOVE(low_record_ptr_pi, ibuf1)
            end
        END_RPT()
        BRL(back)

    @label flush_0_local_record_with_count
        # Get count and add 0 for records being passed.
        # Skip the output for hte last core when not passing flush signals.
        if lastcore == 0 || lastcore_flush == 1
            MOVE(num_records_to_pass, ibuf0_next)
            ADD(output[0], ibuf0, 0)
        else
            MOVE(num_records_to_pass, ibuf0)
        end
        BRL(flush_0_local_record, nop2)

    @label flush_1_local_record_with_count
        # Get count and add 1 for records being passed.
        if lastcore == 0 || lastcore_flush == 1
            MOVE(num_records_to_pass, ibuf0_next)
            ADD(output[0], ibuf0, 1)
        else
            MOVE(num_records_to_pass, ibuf0)
        end
        BRL(flush_1_local_record)

    @label flush_2_local_record
        # Get count and add 2 for records being passed.
        if lastcore == 0 || lastcore_flush == 1
            MOVE(num_records_to_pass, ibuf0_next)
            ADD(output[0], ibuf0, 2)
        else
            MOVE(num_records_to_pass, ibuf0)
        end

        # Send low record.
        # Not expanded due to imem limit.
        RPT(25, nop3)
            MOVE(output[0], low_record_ptr_pi)
        END_RPT()
        RPT(25, nop3)
            if lastcore
                MOVE(output[0], low_record_ptr_pi)
            else
                MOVE(output[1], low_record_ptr_pi)
            end
        END_RPT()

    @label flush_1_local_record
        # Send high record.
        # 5x5=25
        RPT(5)
            MOVE(output[0], high_record_ptr_pi)
            MOVE(output[0], high_record_ptr_pi)
            MOVE(output[0], high_record_ptr_pi)
            MOVE(output[0], high_record_ptr_pi)
            MOVE(output[0], high_record_ptr_pi)
        END_RPT()
        RPT(25, nop3)
            if lastcore
                MOVE(output[0], high_record_ptr_pi)
            else
                MOVE(output[1], high_record_ptr_pi)
            end
            END_RPT()

    @label flush_0_local_record
        # Pass records until count of how many to pass goes negative.
        # Maximum of ~32k records supported.
        # Start with a 0 check.
        MOVE(null, num_records_to_pass)
        BR(start, Z)

    @label pass_record
    # Loop unroll the repeats to remove the need for nops.
    # 5x5=25
    RPT(5)
        MOVE(output[0], ibuf0)
        MOVE(output[0], ibuf0)
        MOVE(output[0], ibuf0)
        MOVE(output[0], ibuf0)
        MOVE(output[0], ibuf0)
    END_RPT()
    # 5x5=25
    RPT(5)
        if firstcore && !lastcore
            MOVE(output[1], ibuf0)
            MOVE(output[1], ibuf0)
            MOVE(output[1], ibuf0)
            MOVE(output[1], ibuf0)
            MOVE(output[1], ibuf0)
        elseif firstcore && lastcore
            MOVE(output[0], ibuf0)
            MOVE(output[0], ibuf0)
            MOVE(output[0], ibuf0)
            MOVE(output[0], ibuf0)
            MOVE(output[0], ibuf0)
        elseif !firstcore && lastcore
            MOVE(output[0], ibuf1)
            MOVE(output[0], ibuf1)
            MOVE(output[0], ibuf1)
            MOVE(output[0], ibuf1)
            MOVE(output[0], ibuf1)
        else
            MOVE(output[1], ibuf1)
            MOVE(output[1], ibuf1)
            MOVE(output[1], ibuf1)
            MOVE(output[1], ibuf1)
            MOVE(output[1], ibuf1)
        end
    END_RPT()

    # Check count and loop.
    SUBU(num_records_to_pass, num_records_to_pass, 1)
    BRL(pass_record, Z, neg)
    BRL(start)

end

# Start_Initialization
#   # Verify storage size.
#   if(m->storage.size() != 3){
#       __debugbreak();
#   }
#
#   # Set up address generator strides.
#   MOVI(AG0STRIDE, 0xFF)
#   MOVI(AG1STRIDE, 1)
#   MOVI(AG2STRIDE, 1)
#
#   # Set condex 0 to look at carry flag.
#   MOVI(CXMASK0, 1<<COND_EXEC_MASK_CARRY_BIT)
#
# End_Initialization


module FiniteHorizonPOMDPs

using POMDPs
using POMDPModelTools
using Random: Random, AbstractRNG

export
    HorizonLength,
    FiniteHorizon,
    InfiniteHorizon,
    horizon,
    stage,
    stage_states,
    stage_stateindex,
    ordered_stage_states,
    stage_observations,
    stage_obsindex,
    ordered_stage_observations

include("interface.jl")

export
    fixhorizon

include("fixhorizon.jl")

end