# Themelis, Stella, Patrinos, "Forward-backward envelope for the sum of two
# nonconvex functions: Further properties and nonmonotone line-search
# algorithms", SIAM Journal on Optimization, vol. 28, no. 3, pp. 2274–2303
# (2018).

using Base.Iterators
using ProximalAlgorithms.IterationTools
using ProximalOperators: Zero
using LinearAlgebra
using Printf

"""
    ZeroFPRIteration(; <keyword-arguments>)

Instantiate the ZeroFPR algorithm (see [1]) for solving optimization problems
of the form

    minimize f(Ax) + g(x),

where `f` is smooth and `A` is a linear mapping (for example, a matrix).

# Arguments
- `x0`: initial point.
- `f=Zero()`: smooth objective term.
- `A=I`: linear operator (e.g. a matrix).
- `g=Zero()`: proximable objective term.
- `Lf=nothing`: Lipschitz constant of the gradient of x ↦ f(Ax).
- `gamma=nothing`: stepsize to use, defaults to `1/Lf` if not set (but `Lf` is).
- `adaptive=false`: forces the method stepsize to be adaptively adjusted.
- `minimum_gamma=1e-7`: lower bound to `gamma` in case `adaptive == true`.
- `max_backtracks=20`: maximum number of line-search backtracks.
- `H=LBFGS(x0, 5)`: variable metric to use to compute line-search directions.

# References
- [1] Themelis, Stella, Patrinos, "Forward-backward envelope for the sum of two
nonconvex functions: Further properties and nonmonotone line-search algorithms",
SIAM Journal on Optimization, vol. 28, no. 3, pp. 2274–2303 (2018).
"""

Base.@kwdef struct ZeroFPRIteration{R,C<:Union{R,Complex{R}},Tx<:AbstractArray{C},Tf,TA,Tg,TH}
    f::Tf = Zero()
    A::TA = I
    g::Tg = Zero()
    x0::Tx
    alpha::R = real(eltype(x0))(0.95)
    beta::R = real(eltype(x0))(0.5)
    Lf::Maybe{R} = nothing
    gamma::Maybe{R} = Lf === nothing ? nothing : (alpha / Lf)
    adaptive::Bool = false
    minimum_gamma::R = real(eltype(x0))(1e-7)
    max_backtracks::Int = 20
    H::TH = LBFGS(x0, 5)
end

Base.IteratorSize(::Type{<:ZeroFPRIteration}) = Base.IsInfinite()

Base.@kwdef mutable struct ZeroFPRState{R,Tx,TAx,TH}
    x::Tx             # iterate
    Ax::TAx           # A times x
    f_Ax::R           # value of smooth term
    grad_f_Ax::TAx    # gradient of f at Ax
    At_grad_f_Ax::Tx  # gradient of smooth term
    gamma::R          # stepsize parameter of forward and backward steps
    y::Tx             # forward point
    xbar::Tx          # forward-backward point
    g_xbar::R         # value of nonsmooth term (at xbar)
    res::Tx           # fixed-point residual at iterate (= x - xbar)
    H::TH             # variable metric
    tau::Maybe{R} = nothing
    Axbar::TAx = zero(Ax)
    grad_f_Axbar::TAx = zero(Ax)
    At_grad_f_Axbar::Tx = zero(x)
    xbarbar::Tx = zero(x)
    res_xbar::Tx = zero(x)
    xbar_prev::Maybe{Tx} = nothing
    res_xbar_prev::Maybe{Tx} = nothing
    d::Tx = zero(x)
    Ad::TAx = zero(Ax)
end

f_model(state::ZeroFPRState) =
    f_model(state.f_Ax, state.At_grad_f_Ax, state.res, state.gamma)

function Base.iterate(iter::ZeroFPRIteration{R}) where {R}
    x = copy(iter.x0)
    Ax = iter.A * x
    grad_f_Ax, f_Ax = gradient(iter.f, Ax)

    gamma = iter.gamma

    if gamma === nothing
        # compute lower bound to Lipschitz constant of the gradient of x ↦ f(Ax)
        xeps = x .+ R(1)
        grad_f_Axeps, f_Axeps = gradient(iter.f, iter.A * xeps)
        L = norm(iter.A' * (grad_f_Axeps - grad_f_Ax)) / R(sqrt(length(x)))
        gamma = iter.alpha / L
    end

    # compute initial forward-backward step
    At_grad_f_Ax = iter.A' * grad_f_Ax
    y = x - gamma .* At_grad_f_Ax
    xbar, g_xbar = prox(iter.g, y, gamma)

    # compute initial fixed-point residual
    res = x - xbar

    state = ZeroFPRState(
        x=x, Ax=Ax, f_Ax=f_Ax, grad_f_Ax=grad_f_Ax, At_grad_f_Ax=At_grad_f_Ax,
        gamma=gamma, y=y, xbar=xbar, g_xbar=g_xbar, res=res, H=iter.H,
    )

    return state, state
end

function Base.iterate(
    iter::ZeroFPRIteration{R},
    state::ZeroFPRState{R,Tx,TAx},
) where {R,Tx,TAx}
    f_Axbar_upp = f_model(state)
    # These need to be performed anyway (to compute xbarbar later on)
    mul!(state.Axbar, iter.A, state.xbar)
    f_Axbar = gradient!(state.grad_f_Axbar, iter.f, state.Axbar)

    # backtrack gamma (warn and halt if gamma gets too small)
    while iter.gamma === nothing || iter.adaptive == true
        if state.gamma < iter.minimum_gamma
            @warn "parameter `gamma` became too small ($(state.gamma)), stopping the iterations"
            return nothing
        end
        tol = 10 * eps(R) * (1 + abs(f_Axbar))
        if f_Axbar <= f_Axbar_upp + tol
            break
        end
        state.gamma *= 0.5
        state.y .= state.x .- state.gamma .* state.At_grad_f_Ax
        state.g_xbar = prox!(state.xbar, iter.g, state.y, state.gamma)
        state.res .= state.x .- state.xbar
        reset!(state.H)
        f_Axbar_upp = f_model(state)
        mul!(state.Axbar, iter.A, state.xbar)
        f_Axbar = gradient!(state.grad_f_Axbar, iter.f, state.Axbar)
    end

    if state.xbar_prev === nothing
        state.xbar_prev = zero(state.x)
        state.res_xbar_prev = zero(state.x)
    end

    # compute FBE
    FBE_x = f_Axbar_upp + state.g_xbar

    # compute residual at xbar
    mul!(state.At_grad_f_Axbar, iter.A', state.grad_f_Axbar)
    state.y .= state.xbar .- state.gamma .* state.At_grad_f_Axbar
    g_xbarbar = prox!(state.xbarbar, iter.g, state.y, state.gamma)
    state.res_xbar .= state.xbar .- state.xbarbar

    if state.xbar_prev !== nothing
        # update metric
        update!(state.H, state.xbar - state.xbar_prev, state.res_xbar - state.res_xbar_prev)
        # store vectors for next update
        copyto!(state.xbar_prev, state.xbar)
        copyto!(state.res_xbar_prev, state.res_xbar)
    end

    # compute direction
    mul!(state.d, state.H, -state.res_xbar)

    # Perform line-search over the FBE
    tau = R(1)
    mul!(state.Ad, iter.A, state.d)

    sigma = iter.beta * (0.5 / state.gamma) * (1 - iter.alpha)
    tol = 10 * eps(R) * (1 + abs(FBE_x))
    threshold = FBE_x - sigma * norm(state.res)^2 + tol

    for _ in 1:iter.max_backtracks
        state.x .= state.xbar_prev .+ tau .* state.d
        state.Ax .= state.Axbar .+ tau .* state.Ad
        # TODO: can precompute most of next line in case f is quadratic
        state.f_Ax = gradient!(state.grad_f_Ax, iter.f, state.Ax)
        mul!(state.At_grad_f_Ax, iter.A', state.grad_f_Ax)
        state.y .= state.x .- state.gamma .* state.At_grad_f_Ax
        state.g_xbar = prox!(state.xbar, iter.g, state.y, state.gamma)
        state.res .= state.x .- state.xbar
        FBE_x = f_model(state) + state.g_xbar

        if FBE_x <= threshold
            state.tau = tau
            return state, state
        end

        tau *= 0.5
    end

    @warn "stepsize `tau` became too small ($(tau)), stopping the iterations"
    return nothing
end

# Solver

struct ZeroFPR{R, K}
    maxit::Int
    tol::R
    verbose::Bool
    freq::Int
    kwargs::K
end

function (solver::ZeroFPR)(x0; kwargs...)
    stop(state::ZeroFPRState) = norm(state.res, Inf) / state.gamma <= solver.tol
    disp((it, state)) = @printf(
        "%5d | %.3e | %.3e | %.3e\n",
        it,
        state.gamma,
        norm(state.res, Inf) / state.gamma,
        (state.tau === nothing ? 0.0 : state.tau)
    )
    iter = ZeroFPRIteration(; x0=x0, solver.kwargs..., kwargs...)
    iter = take(halt(iter, stop), solver.maxit)
    iter = enumerate(iter)
    if solver.verbose
        iter = tee(sample(iter, solver.freq), disp)
    end
    num_iters, state_final = loop(iter)
    return state_final.xbar, num_iters
end

ZeroFPR(; maxit=1_000, tol=1e-8, verbose=false, freq=10, kwargs...) =
    ZeroFPR(maxit, tol, verbose, freq, kwargs)


nodetype(a::Type{<:AbstractCell2Node}) = Vec{Float64}
@inline nodetype(a::A) where {A<:AbstractCell2Node} = nodetype(A)
nnodes(a::Type{<:AbstractCell2Node{NN,NF}}) where {NN,NF} = NN
@inline nnodes(a::A) where {A<:AbstractCell2Node} = nnodes(A)
nfaces(a::Type{<:AbstractCell2Node{NN,NF}}) where {NN,NF} = NF
@inline nfaces(a::A) where {A<:AbstractCell2Node} = nfaces(A)

nodetype(a::Type{<:AbstractCell2Node2D}) = Vec2D{Float64}
@inline nodetype(a::A) where {A<:AbstractCell2Node2D} = nodetype(A)

struct TriangleCell <: AbstractCell2Node2D{3,3}
    ind::Tuple{UInt,UInt,UInt}
end

function Base.show(io::IO,t::TriangleCell)
    compact = get(io, :compact, false)

    if !compact
        print(io,"Triangle")
    end
    print(io,"(",t[1],",",t[2],",",t[3],")")
end

vtkcelltype(a::TriangleCell) = 5

function aspect_ratio(a::TriangleCell,m::AbstractMesh)
    @inbounds begin
        l1 = distance(m.nodes[a[1]], m.nodes[a[2]])
        l2 = distance(m.nodes[a[2]], m.nodes[a[3]])
        l3 = distance(m.nodes[a[3]], m.nodes[a[1]])
    end
    return min(l1,l2,l3)/max(l1,l2,l3)
end

cell_faces(a::TriangleCell) = Face2Node{2}((a[1], a[2])), Face2Node{2}((a[2], a[3])), Face2Node{2}((a[3], a[1]))

@inline volume(a::TriangleCell,n1::AbstractVec,n2::AbstractVec,n3::AbstractVec) = triangle_area(n1,n2,n3)
@inline function volume(a::TriangleCell,nodes::AbstractArray{<:Vec2D}) ; @inbounds r = volume(a, nodes[a[1]], nodes[a[2]], nodes[a[3]]); return r;end

@inline perimeter(a::TriangleCell,n1::AbstractVec,n2::AbstractVec,n3::AbstractVec) = triangle_perimeter(n1,n2,n3)
@inline function perimeter(a::TriangleCell,nodes::AbstractArray{<:Vec2D}) ; @inbounds r = perimeter(a, nodes[a[1]], nodes[a[2]], nodes[a[3]]); return r;end

@inline mean_dx(a::TriangleCell,n1::AbstractVec,n2::AbstractVec,n3::AbstractVec) = 4*volume(a,n1,n2,n3)/perimeter(a,n1,n2,n3)
@inline function mean_dx(a::TriangleCell,nodes::AbstractArray{<:Vec2D}) ; @inbounds r = mean_dx(a, nodes[a[1]], nodes[a[2]], nodes[a[3]]); return r; end

@inline center(a::TriangleCell,n1::AbstractVec,n2::AbstractVec,n3::AbstractVec) = triangle_center(n1,n2,n3)
@inline function center(a::TriangleCell,nodes::AbstractArray{<:Vec2D}) ; @inbounds r = center(a, nodes[a[1]], nodes[a[2]], nodes[a[3]]); return r;end

abstract type AbstractCell2Face{T,N} <: NTupleWrap{N,Int} end

#=
Idea: make lists of (volumec1,volume,c2), (centerc1,centerc2), (cellpropertyc1,cellpropertyc2)
=#


using FuzzySets, Test

x = [1, 4.5]
y = [3.5, 4]
@test FuzzySets.distance(x, y) == 2.5495097567963922


__precompile__()

"""
    FunManifolds

Main module for `FunManifolds.jl` -- a Julia package for functional differential geometry.
"""
module FunManifolds

using LinearAlgebra
using Interpolations
using Manifolds
using ManifoldsBase
using Markdown: @doc_str
using QuadGK
using StaticArrays

import Base: +, -, *, isapprox

import ManifoldsBase:
    allocate,
    allocate_result,
    allocate_result_type,
    allocation_promotion_function,
    array_value,
    base_manifold,
    check_point,
    check_point__transparent,
    check_vector,
    decorated_manifold,
    decorator_transparent_dispatch,
    default_decorator_dispatch,
    distance,
    embed,
    embed!,
    exp,
    exp!,
    exp!__intransparent,
    geodesic,
    # TODO: uncomment the import if `flat!` goes to ManifoldsBase
    # flat!__intransparent,
    get_basis,
    get_coordinates,
    get_coordinates!,
    get_embedding,
    get_vector,
    get_vector!,
    get_vectors,
    injectivity_radius,
    inner,
    inner__intransparent,
    is_point,
    is_vector,
    inverse_retract,
    inverse_retract!,
    log,
    log!,
    manifold_dimension,
    mid_point,
    mid_point!,
    norm,
    number_eltype,
    number_of_coordinates,
    project,
    project!,
    representation_size,
    retract,
    retract!,
    shortest_geodesic,
    vector_transport_direction,
    vector_transport_direction!,
    vector_transport_to,
    vector_transport_to!,
    zero_vector,
    zero_vector!

import Manifolds:
    apply,
    apply!,
    compose,
    compose!,
    get_iterator,
    identity_element,
    identity_element!,
    inv,
    inv!,
    optimal_alignment,
    optimal_alignment!,
    zero_vector

using Manifolds: _read, _write, AbstractRiemannianDiffBackend

mutable struct GeneralParams
    quad_rel_tol::Union{Real,Nothing}
    quad_abs_tol::Union{Real,Nothing}
end

const global PARAMS = GeneralParams(nothing, nothing)

function rtoldefault(M::AbstractManifold, x1, x2)
    return 1.e-8
end

function atoldefault(M::AbstractManifold, x1, x2)
    return 0.0
end

function concretize_tols(M::AbstractManifold, x1, x2; reltol=nothing, abstol=nothing)
    rtol = if reltol === nothing
        rtoldefault(M, x1, x2)
    else
        reltol
    end

    atol = if abstol === nothing
        atoldefault(M, x1, x2)
    else
        abstol
    end

    return (rtol, atol)
end

struct ProjectedDifferenceBackend{TDT<:Union{Number,Nothing}} <:
       AbstractRiemannianDiffBackend
    dt::TDT
end

ProjectedDifferenceBackend() = ProjectedDifferenceBackend{Float64}(1e-7)

include("interpolation.jl")

include("DiscretizedCurves.jl")
include("FunctionCurve.jl")
include("CurveWarping.jl")
include("CurveWarpingSRSF.jl")
include("curve_warping_alignment.jl")
include("functional_transformations.jl")

export CurveWarpingAction,
    CurveWarpingGroup,
    CurveWarpingSpace,
    CurveWarpingSRSFAction,
    CurveWarpingSRSFGroup,
    CurveWarpingSRSFSpace,
    DiscretizedCurves,
    FunctionCurveSpace,
    reverse_srsf,
    reverse_srvf,
    srsf,
    srvf,
    tsrvf,
    transport_srvf,
    transport_srvf!,
    UniformDiscretizedCurves

end #module


using Random
using LinearAlgebra

function uniformOnSphere(rng = Random.GLOBAL_RNG)
    phi = 2.0 * pi * rand(rng)
    z = 2.0 * rand(rng) - 1.0;
    r = sqrt(1.0 - z * z)
    return (r * cos(phi), r * sin(phi), z)
end

function exchangeEnergy(s1, M::InteractionMatrix, s2)::Float64
    return s1[1] * (M.m11 * s2[1] + M.m12 * s2[2] + M.m13 * s2[3]) + s1[2] * (M.m21 * s2[1] + M.m22 * s2[2] + M.m23 * s2[3]) + s1[3] * (M.m31 * s2[1] + M.m32 * s2[2] + M.m33 * s2[3])
end

function getEnergy(lattice::Lattice{D,N})::Float64 where {D,N}
    energy = 0.0

    for site in 1:length(lattice)
        s0 = getSpin(lattice, site)

        #two-spin interactions
        interactionSites = getInteractionSites(lattice, site)
        interactionMatrices = getInteractionMatrices(lattice, site)
        for i in 1:length(interactionSites)
            if site > interactionSites[i]
                energy += exchangeEnergy(s0, interactionMatrices[i], getSpin(lattice, interactionSites[i]))
            end
        end

        #onsite interaction
        energy += exchangeEnergy(s0, getInteractionOnsite(lattice, site), s0)

        #field interaction
        energy += dot(s0, getInteractionField(lattice, site))
    end

    return energy
end

function getEnergyDifference(lattice::Lattice{D,N}, site::Int, newState::Tuple{Float64,Float64,Float64})::Float64 where {D,N}
    dE = 0.0
    oldState = getSpin(lattice, site)
    ds = newState .- oldState

    #two-spin interactions
    interactionSites = getInteractionSites(lattice, site)
    interactionMatrices = getInteractionMatrices(lattice, site)
    for i in 1:length(interactionSites)
        dE += exchangeEnergy(ds, interactionMatrices[i], getSpin(lattice, interactionSites[i]))
    end

    #onsite interaction
    interactionOnsite = getInteractionOnsite(lattice, site)
    dE += exchangeEnergy(newState, interactionOnsite, newState) - exchangeEnergy(oldState, interactionOnsite, oldState)

    #field interaction
    dE += dot(ds, getInteractionField(lattice, site))

    return dE
end

function getMagnetization(lattice::Lattice{D,N}) where {D,N}
    mx, my, mz = 0.0, 0.0, 0.0
    for i in 1:length(lattice)
        spin = getSpin(lattice, i)
        mx += spin[1]
        my += spin[2]
        mz += spin[3]
    end
    return [mx, my, mz] / length(lattice)
end

function getCorrelation(lattice::Lattice{D,N}, spin::Int = 1) where {D,N}
    corr = zeros(length(lattice))
    s0 = getSpin(lattice, spin)
    for i in 1:length(lattice)
        corr[i] = dot(s0, getSpin(lattice, i))
    end
    return corr
end


using Parameters
using DifferentialEquations
using Blink
using Interact
using Plots
gr()

# Bring in the model
include("lv_model.jl")

# Build a Graphical User Interface (GUI) to explore the parameters
let
    u0 = [0.5, 0.5]
    t_span = (0.0, 100.0)

    ui = @manipulate for α12 in slider(0.1:0.01:2.0; label = "α12", value = 0.8),
                         α21 in slider(0.1:0.01:2.0; label = "α21", value = 0.3),
                         K1 in slider(0.1:0.01:3.0; label = "K1", value = 1.5),
                         K2 in slider(0.1:0.01:3.0; label = "K2", value = 1.0),
                         r1 in slider(0.1:0.01:10.0; label = "r1", value = 1.0),
                         r2 in slider(0.1:0.01:10.0; label = "r2", value = 2.0)
        p = LVPar(α12 = α12, α21 = α21, K1 = K1, K2 = K2, r1 = r1, r2 = r2)
        prob = ODEProblem(lv_comp, u0, t_span, p)
        sol = solve(prob, reltol = 1e-8)
        plot(sol)
        xlabel!("time")
        ylabel!("Density")
    end

    w = Blink.AtomShell.Window()
    body!(w, ui)
end


# add more julia processes
using Distributed
nprocs() < 4 ? addprocs(5-nprocs()) : nothing

# tell Julia that you want to use these modules/namespaces
using RoME
## Inter-operating visualization packages for Caesar/RoME/IncrementalInference exist
using RoMEPlotting

# start with an empty factor graph object
fg = initfg()

# Add the first pose :x0
addVariable!(fg, :x0, Pose2)

# Add at a fixed location PriorPose2 to pin :x0 to a starting location (10,10, pi/4)
addFactor!(fg, [:x0], PriorPose2( MvNormal([10; 10; pi/6.0], Matrix(Diagonal([0.1;0.1;0.05].^2))) ) )

# Drive around in a hexagon
for i in 0:5
    psym = Symbol("x$i")
    nsym = Symbol("x$(i+1)")
    addVariable!(fg, nsym, Pose2)
    pp = Pose2Pose2(MvNormal([10.0;0;pi/3], Matrix(Diagonal([0.1;0.1;0.1].^2))))
    addFactor!(fg, [psym;nsym], pp )
end

# perform inference, and remember first runs are slower owing to Julia's just-in-time compiling
getSolverParams(fg).drawtree=true
getSolverParams(fg).showtree=true
tree, smt, hist = solveTree!(fg)

pl = drawPoses(fg) #, meanmax=:mean)
# For scripting use-cases you can export the image
pl |> Gadfly.PDF("/tmp/test1.pdf", 20cm, 10cm)  # or PNG(...)

# another type of plot
plotPose(fg, :x6)

# Add landmarks with Bearing range measurements
addVariable!(fg, :l1, Point2, tags=[:LANDMARK])
p2br = Pose2Point2BearingRange(Normal(0,0.1),Normal(20.0,1.0))
addFactor!(fg, [:x0; :l1], p2br )

# Initialize :l1 numerical values but do not rerun solver
ensureAllInitialized!(fg)
pl = drawPosesLandms(fg)
Gadfly.draw(Gadfly.PDF("/tmp/test2.pdf", 20cm, 10cm),pl)  # or PNG(...)

# Add landmarks with Bearing range measurements
p2br2 = Pose2Point2BearingRange(Normal(0,0.1),Normal(20.0,1.0))
addFactor!(fg, [:x6; :l1], p2br2 )

# solve
batchSolve!(fg, drawpdf=true)

# redraw
pl = drawPosesLandms(fg, meanmax=:mean)
Gadfly.draw(Gadfly.PDF("/tmp/test3.pdf", 20cm, 10cm),pl)  # or PNG(...)

#


using InteractionNetworkModels, StructuredDistances, Distributions
using Plots, Random

E = [
    [1,2,1,2],
    [1,2,1],
    [3,4,3]
]
d = MatchingDistance(FastLCS(101))
K_inner, K_outer = (DimensionRange(1,10), DimensionRange(1,25))
model = SIM(
    E, 3.5,
    d,
    1:10,
    K_inner, K_outer
)

# Auxiliary sampler
β = 0.7
mcmc_move = InvMcmcMixtureMove(
    (
        EditAllocationMove(ν=3),
        InsertDeleteMove(ν=3, len_dist=TrGeometric(0.8, 1, model.K_inner.u))
    ),
    (β, 1-β)
)
mcmc_sampler = InvMcmcSampler(
    mcmc_move,
    lag=30, burn_in=3000
)

@time x = mcmc_sampler(
    model, desired_samples=5000, burn_in=0, lag=1
)
plot(x)

@time x = mcmc_sampler(
    model,
    desired_samples=50,
    lag=500,
    burn_in=10000
)
acceptance_prob(mcmc_sampler)

plot(x)
summaryplot(x)

E_prior = SIM(E, 0.1, model.dist, model.V, model.K_inner, model.K_outer)
γ_prior = Uniform(0.5,7.0)

data = x.sample
posterior = SimPosterior(data, E_prior, γ_prior)

# Construct posterior sampler
β = 0.8
mode_move = InvMcmcMixtureMove(
    (
        EditAllocationMove(ν=1),
        PathPermutationMove(ν=1),
        InsertDeleteMove(ν=1, len_dist=TrGeometric(0.8, 1, model.K_inner.u)),
        SplitMergeMove(ν=1)
    ),
    (β/2, β/2, (1-β)/2, (1-β)/2)
    # (β, 1-β)
)

posterior_sampler = IexMcmcSampler(
    mode_move, mcmc_sampler,
    ε=0.3,
    aux_init_at_prev=true
)

# E_init = shuffle.(model.mode)
E_init = [vcat(model.mode...)]
@time x = posterior_sampler(
    posterior,
    desired_samples=1_000,
    S_init=E_init,
    γ_init=3.5
)
acceptance_prob(posterior_sampler)

# With fast matching distance
d_fast = FastMatchingDistance(FastLCS(101), 1000)
model_fast = SIM(
    E, 3.5,
    d_fast,
    1:10,
    K_inner, K_outer
)

E_prior_fast = SIM(E, 0.1, model.dist, model.V, model.K_inner, model.K_outer)
posterior_fast = SimPosterior(data, E_prior_fast, γ_prior)

E_init = [vcat(model.mode...)]
@time x = posterior_sampler(
    posterior_fast,
    desired_samples=1_000,
    S_init=E_init,
    γ_init=3.5
)
acceptance_prob(posterior_sampler)

plot(x.suff_stat_trace)
plot(x,E)
x.S_sample


### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# ╔═╡ 25e36fc6-7b76-11eb-02df-bd87dde5cdfb
begin
	using Flux: Chain, Dense, params, logitcrossentropy, onehotbatch, ADAM, train!, softmax
	using Test
end

# ╔═╡ 8fe7bc4c-7b85-11eb-0287-0f179e325305
md"### Data Generation"

# ╔═╡ fde731a6-7b84-11eb-3758-23db3b0d2057
function fizzbuzz(x::Int)
    is_divisible_by_three = x % 3 == 0
    is_divisible_by_five = x % 5 == 0

    if is_divisible_by_three & is_divisible_by_five
        return "fizzbuzz"
    elseif is_divisible_by_three
        return "fizz"
    elseif is_divisible_by_five
        return "buzz"
    else
        return "else"
    end
end

# ╔═╡ 3a290946-7b85-11eb-39d9-87db6c6beef1
const LABELS = ["fizz", "buzz", "fizzbuzz", "else"];

# ╔═╡ 7f695632-7b85-11eb-1804-a9876b2c027f
md"### Feature Engineering"

# ╔═╡ 735140d0-7b85-11eb-0aaf-e3b29e66fa00
features(x) = float.([x % 3, x % 5, x % 15])

# ╔═╡ 7351a052-7b85-11eb-3a5d-e323871eb7ad
features(x::AbstractArray) = hcat(features.(x)...)

# ╔═╡ a0765410-7b85-11eb-076e-6f877fff0c95
features(1:15)

# ╔═╡ efff94ba-7b85-11eb-0d9e-d333585c6ed3
md"### Data Preparation"

# ╔═╡ f9de33f6-7b85-11eb-3eb0-6fcf7a369345
function getdata()

    @test fizzbuzz.([3, 5, 15, 98]) == LABELS

    raw_x = 1:100;
    raw_y = fizzbuzz.(raw_x);

    X = features(raw_x);
    y = onehotbatch(raw_y, LABELS);
    return X, y
end

# ╔═╡ 05d0994c-7b86-11eb-09ff-a33769bc2fc2
raw_x = 1:100

# ╔═╡ 73acb7d4-7b86-11eb-2666-b91efed75c66
raw_y = fizzbuzz.(raw_x)

# ╔═╡ 73ad3c68-7b86-11eb-27a3-815a9fa17317
X = features(raw_x)

# ╔═╡ 73b0a77c-7b86-11eb-15ae-fbb1afdb09c3
y = onehotbatch(raw_y, LABELS)

# ╔═╡ a45bc3b6-7b86-11eb-2b77-c5ce1e85d93c
m = Chain(Dense(3, 10), Dense(10, 4))

# ╔═╡ 3515921a-7b87-11eb-3639-f1369417430e
l1, l2 = m.layers

# ╔═╡ 4ae19f9c-7b87-11eb-326d-f36d3f129e86
m(features(3))

# ╔═╡ aadac77c-7b87-11eb-1720-07dad016249d
loss(x, y) = logitcrossentropy(m(x), y)

# ╔═╡ ba3420a6-7b87-11eb-33d9-4bf0d6ac359e
loss(X, y)

# ╔═╡ e2a834f0-7b87-11eb-10db-7151b60248ae
opt = ADAM()

# ╔═╡ e3e37032-7b87-11eb-3eff-c7592e3878ce
train!(loss, params(m), [(X, y)], opt)

# ╔═╡ 627c3cf8-7b88-11eb-0992-0fffc92c84ca
for e in 1:500
	train!(loss, params(m), [(X, y)], opt)
end

# ╔═╡ 103c356a-7b88-11eb-0c1b-a37efe27b9ef
deepbuzz(x) = (a = argmax(m(features(x))); a == 4 ? x : LABELS[a])

# ╔═╡ 346d4b2c-7b88-11eb-3713-ef621dc50ef5
deepbuzz.(1:15)

# ╔═╡ 5e38196e-7b88-11eb-3997-b52bcd52d76d

# ╔═╡ Cell order:
# ╠═25e36fc6-7b76-11eb-02df-bd87dde5cdfb
# ╟─8fe7bc4c-7b85-11eb-0287-0f179e325305
# ╠═fde731a6-7b84-11eb-3758-23db3b0d2057
# ╠═3a290946-7b85-11eb-39d9-87db6c6beef1
# ╟─7f695632-7b85-11eb-1804-a9876b2c027f
# ╠═735140d0-7b85-11eb-0aaf-e3b29e66fa00
# ╠═7351a052-7b85-11eb-3a5d-e323871eb7ad
# ╠═a0765410-7b85-11eb-076e-6f877fff0c95
# ╟─efff94ba-7b85-11eb-0d9e-d333585c6ed3
# ╠═f9de33f6-7b85-11eb-3eb0-6fcf7a369345
# ╠═05d0994c-7b86-11eb-09ff-a33769bc2fc2
# ╠═73acb7d4-7b86-11eb-2666-b91efed75c66
# ╠═73ad3c68-7b86-11eb-27a3-815a9fa17317
# ╠═73b0a77c-7b86-11eb-15ae-fbb1afdb09c3
# ╠═a45bc3b6-7b86-11eb-2b77-c5ce1e85d93c
# ╠═3515921a-7b87-11eb-3639-f1369417430e
# ╠═4ae19f9c-7b87-11eb-326d-f36d3f129e86
# ╠═aadac77c-7b87-11eb-1720-07dad016249d
# ╠═ba3420a6-7b87-11eb-33d9-4bf0d6ac359e
# ╠═e2a834f0-7b87-11eb-10db-7151b60248ae
# ╠═e3e37032-7b87-11eb-3eff-c7592e3878ce
# ╠═627c3cf8-7b88-11eb-0992-0fffc92c84ca
# ╠═103c356a-7b88-11eb-0c1b-a37efe27b9ef
# ╠═346d4b2c-7b88-11eb-3713-ef621dc50ef5
# ╠═5e38196e-7b88-11eb-3997-b52bcd52d76d


using Random

## Binary string fitnesses

function onemax(instance:: Dict{String, }, point::BitArray{1})
    sum(point)
end

function zeromax(instance:: Dict{String, }, point::BitArray{1})
    length(point) - sum(point)
end

function fractionmax(instance:: Dict{String, }, fraction::Float64)
    return function _F(instance, point)
        ## length(candidate)*fraction
        lastZero = Int(instance["size"]*fraction)
        ZeroMaxFitness(instance, point[1:lastZero]) + OneMaxFitness(instance, point[lastZero+1:end])
    end
end


module ConstraintProgrammingExtensions
  using MathOptInterface

  """
      AllDifferent(dimension::Int)
  The set corresponding to an all-different constraint.
  All expressions of a vector-valued function are enforced to take distinct
  values in the solution: for all pairs of expressions, their values must
  differ.
  This constraint is sometimes called `distinct`.
  ## Example
      [x, y, z] in AllDifferent(3)
      # enforces `x != y` AND `x != z` AND `y != z`.
  """
  struct AllDifferent <: AbstractVectorSet
      dimension::Int
  end

  """
      Domain{T <: Number}(values::Set{T})
  The set corresponding to an enumeration of constant values.
  The value of a scalar function is enforced to take a value from this set of
  values.
  This constraint is sometimes called `in`.
  ## Example
      x in Domain(Set(1, 2, 3))
      # enforces `x == 1` OR `x == 2` OR `x == 3`.
  """
  struct Domain{T <: Number} <: AbstractScalarSet
      values::Set{T}
  end

  function Base.copy(set::Domain{T}) where T
      return Domain(copy(set.values))
  end

  """
      Membership(dimension)
  The first element of a function of dimension `dimension` must equal at least
  one of the following `dimension - 1` elements of the function.
  This constraint is sometimes called `in_set`.
  ## Example
      [x, y, z] in Membership(dimension)
      # enforces `x == y` OR `x == z`.
  """
  struct Membership <: AbstractScalarSet
      values::Set{AbstractScalarFunction}
  end

  function Base.copy(set::Membership)
      return Domain(copy(set.values))
  end

  """
      DifferentFrom(dimension)
  The first element of a function of dimension `dimension` cannot equal any
  of the following `dimension - 1` elements of the function.
  ## Example
      [x, y, z] in DifferentFrom(dimension)
      # enforces `x != y` AND `x != z`.
  """
  struct DifferentFrom <: AbstractScalarSet
      dimension::Int
  end

  """
      Count{T <: Real}(value::T, dimension::Int)
  ``\\{(y, x) \\in \\mathbb{N} \\times \\mathbb{R}^n : y = |\\{i | x_i = value\\}|\\}``
  """
  struct Count{T <: Real} <: AbstractVectorSet
      value::T
      dimension::Int
  end

  dimension(set::Count{T}) where T = set.dimension + 1

  function Base.copy(set::Count)
      return Count(copy(set.value), value)
  end

  """
      CountDistinct(dimension::Int)
  The first variable in the set is forced to be the number of distinct values in the rest of the expressions.
  Also called `nvalues`. This is a relaxed version of `AllDifferent`; it encodes an `AllDifferent` constraint
  when the first variable is the number of variables in the set.
  """
  struct CountDistinct <: AbstractVectorSet
      dimension::Int
  end

  dimension(set::CountDistinct) where T = set.dimension + 1

  """
      Strictly{S <: Union{LessThan, GreaterThan}}
  Converts an inequality set to a set with the same inequality made strict. For instance, while LessThan(1)
  corresponds to the inequality `<= 1`, Strictly(LessThan(1)) corresponds to the inequality `< 1`, i.e.
  the value 1 is no more allowed.
  """
  struct Strictly{S <: Union{LessThan, GreaterThan}} <: AbstractScalarSet
      set::S
  end

  function Base.copy(set::Strictly{S}) where S
      return Count(copy(set.set))
  end

  """
      Element{T <: Real}(values::Vector{T})
  ``\\{(x, i) \\in \\mathbb{R}^d \\times \\mathbb{N}^d | x_j = values[i_j]\\}``
  Less formally, the j-th element constrained in this set will take the value of `values` at the index given by
  the (j + dimension)-th element.
  ## Examples
      [x, 3] in Element([4, 5, 6])
      # Enforces that x = 6, because 6 is the 3rd element from the array.
      [x, 3, y, j] in Element([4, 5, 6])
      # Enforces that x = 6, because 6 is the 3rd element from the array.
      # Enforces that y = array[j], depending on the value of j (an integer between 1 and 3).
  """
  struct Element{T <: Real} <: AbstractVectorSet
      values::Vector{T}
      dimension::Int
  end

  dimension(set::Element{T}) where T = 2 * set.dimension

  function Base.copy(set::Element{T}) where T
      return Element(copy(set.values), set.dimension)
  end

  """
      Sort(dimension::Int)
  Ensures that the first `dimension` elements is a sorted copy of the next `dimension` elements.
  ## Example
      [a, b, c, d] in Sort(2)
      # Enforces that:
      # - the first part is sorted: a <= b
      # - the first part corresponds to the second one:
      #     - either a = c and b = d
      #     - or a = d and b = c
  """
  struct Sort <: AbstractVectorSet
      dimension::Int
  end

  dimension(set::Sort) where T = 2 * set.dimension

  """
      SortPermutation(dimension::Int)
  Ensures that the first `dimension` elements is a sorted copy of the next `dimension` elements.
  The last `dimension` elemenst give a permutation to get from the original array to its sorted version.
  ## Example
      [a, b, c, d, i, j] in SortPermutation(2)
      # Enforces that:
      # - the first part is sorted: a <= b
      # - the first part corresponds to the second one:
      #     - either a = c and b = d: in this case, i = 1 and j = 2
      #     - or a = d and b = c: in this case, i = 2 and j = 1
  """
  struct SortPermutation <: AbstractVectorSet
      dimension::Int
  end

  dimension(set::SortPermutation) where T = 3 * set.dimension

  """
      BinPacking(n_bins::Int, n_items::Int)
  Implements an uncapacitated version of the bin packing. The first `n_bins` variables give the load in each bin,
  the next `n_items` give the number of the bin to which the item is assigned to, and the last `n_items` ones
  give the size of each item. The load of a bin is defined as the sum of the sizes of the items put in that bin.
  Also called `pack`.
  ## Example
      [a, b, c, d, e] in BinPacking(1, 2)
      # As there is only one bin, the only solution is to put all the items in that bin.
      # Enforces that:
      # - the bin load is the sum of the objects in that bin: a = d + e
      # - the bin number of the two items is 1: b = c = 1
  """
  struct BinPacking <: AbstractVectorSet
      n_bins::Int
      n_items::Int
  end

  dimension(set::BinPacking) where T = set.n_bins + 2 * set.n_items

  """
      CapacitatedBinPacking(n_bins::Int, n_items::Int)
  Implements an uncapacitated version of the bin packing. The first `n_bins` variables give the load in each bin,
  the next `n_bins` are the capacity of each bin, the followin `n_items` give the number of the bin to which the item
  is assigned to, and the last `n_items` ones give the size of each item. The load of a bin is defined as the sum
  of the sizes of the items put in that bin.
  This constraint is equivalent to `BinPacking` with inequality constraints on the loads of the bins. However, there
  are more efficient propagators for the combined constraint (bin packing with maximum load).
  Also called `bin_packing_capa`.
  ## Example
      [a, 2, b, c, d, e] in CapacitatedBinPacking(1, 2)
      # As there is only one bin, the only solution is to put all the items in that bin if its capacity is large enough.
      # Enforces that:
      # - the bin load is the sum of the objects in that bin: a = d + e
      # - the bin load is at most its capacity: a <= 2
      # - the bin number of the two items is 1: b = c = 1
  """
  struct CapacitatedBinPacking <: AbstractVectorSet
      n_bins::Int
      n_items::Int
  end

  dimension(set::CapacitatedBinPacking) where T = 2 * set.n_bins + 2 * set.n_items

  """
      ReificationSet{S <: AbstractScalarSet}(set::S)
  ``\\{(y, x) \\in \\{0, 1\\} \\times \\mathbb{R}^n | y = 1 \\iff x \\in set, y = 0 otherwise\\}``.
  `S` has to be a sub-type of `AbstractScalarSet`.
  This set serves to find out whether a given constraint is satisfied. The only possible values are
  0 and 1.
  """
  struct ReificationSet{S <: AbstractScalarSet} <: AbstractVectorSet
      set::S
      dimension::Int
  end

  dimension(set::ReificationSet{T}) where T = set.dimension + 1

  function Base.copy(set::ReificationSet{S}) where S
      return ReificationSet(copy(set.set), set.dimension)
  end

  # isbits types, nothing to copy
  function Base.copy(
      set::Union{AllDifferent, DifferentFrom, CountDistinct, Element,
          BinPacking, CapacitatedBinPacking, DifferentFrom}
  )
      return set
  end
end


using Statistics: mean

using FFTW
using GPUifyLoops

using Oceananigans.Operators

function ∇²_ppn!(grid::RegularCartesianGrid, f, ∇²f)
    @loop for k in (1:grid.Nz; blockIdx().z)
        @loop for j in (1:grid.Ny; (blockIdx().y - 1) * blockDim().y + threadIdx().y)
            @loop for i in (1:grid.Nx; (blockIdx().x - 1) * blockDim().x + threadIdx().x)
                @inbounds ∇²f[i, j, k] = ∇²_ppn(grid, f, i, j, k)
            end
        end
    end

    @synchronize
end

function mixed_fft_commutes(N)
    A = rand(N, N, N)
    Ã1 = FFTW.dct(FFTW.rfft(A, [1, 2]), 3)
    Ã2 = FFTW.rfft(FFTW.dct(A, 3), [1, 2])
    Ã1 ≈ Ã2
end

function mixed_ifft_commutes(N)
    A = rand(N, N, N)

    Ã1 = FFTW.dct(FFTW.rfft(A, [1, 2]), 3)
    Ã2 = FFTW.rfft(FFTW.dct(A, 3), [1, 2])

    A11 = FFTW.irfft(FFTW.idct(Ã1, 3), N, [1, 2])
    A12 = FFTW.idct(FFTW.irfft(Ã1, N, [1, 2]), 3)
    A21 = FFTW.irfft(FFTW.idct(Ã2, 3), N, [1, 2])
    A22 = FFTW.idct(FFTW.irfft(Ã2, N, [1, 2]), 3)
    A ≈ A11 && A ≈ A12 && A ≈ A21 && A ≈ A22
end

function fftw_planner_works(ft, Nx, Ny, Nz, planner_flag)
    g = RegularCartesianGrid(ft, (Nx, Ny, Nz), (100, 100, 100))
    RHS = CellField(Complex{ft}, CPU(), g)
    solver = PoissonSolver(g, RHS, FFTW.ESTIMATE)
    true  # Just making sure our PoissonSolver does not error/crash.
end

function poisson_ppn_planned_div_free_cpu(ft, Nx, Ny, Nz, planner_flag)
    g = RegularCartesianGrid(ft, (Nx, Ny, Nz), (100, 100, 100))

    RHS = CellField(Complex{ft}, CPU(), g)
    RHS_orig = CellField(Complex{ft}, CPU(), g)
    ϕ = CellField(Complex{ft}, CPU(), g)
    ∇²ϕ = CellField(Complex{ft}, CPU(), g)

    RHS.data .= rand(Nx, Ny, Nz)
    RHS.data .= RHS.data .- mean(RHS.data)

    RHS_orig.data .= copy(RHS.data)

    solver = PoissonSolver(g, RHS, planner_flag)

    solve_poisson_3d_ppn_planned!(solver, g, RHS, ϕ)
    ∇²_ppn!(g, ϕ, ∇²ϕ)

    ∇²ϕ.data ≈ RHS_orig.data
end


abstract type Data end

function clip(
        A::Array{Union{Missing, T}, 2} where T <: Number;
        x::Int64,
        y::Int64,
        distance::Int64
    )

    sizes = size(A)

    xlower = Int64(max(x - distance, 1))
    xupper = Int64(min(x + distance, sizes[2]))
    ylower = Int64(max(y - distance, 1))
    yupper = Int64(min(y + distance, sizes[1]))

    A_sub = A[ylower:yupper, xlower:xupper]

    dim1 = size(A_sub)[1]
    dim2 = size(A_sub)[2]

    new_x = min(distance + 1, x)
    new_y = min(distance + 1, y)

    dist_array = [sqrt((j - new_x)^2 + (i - new_y)^2) for i = 1:dim1, j = 1:dim2]

    A_sub[dist_array .> distance] .= missing

    A_sub
end

function get_targets(
        source_array::Array{Union{T, Missing}, 2} where T <: Number,
        arguments::Dict{String, Int64},
        precision::DataType
    )

    block_size = arguments["block_size"]
    block_radius = arguments["block_radius"]
    nrows = arguments["nrows"]
    ncols = arguments["ncols"]

    start = block_radius + 1

    xs = [start:block_size:(ncols - block_radius);]
    ys = [start:block_size:(nrows - block_radius);]

    ground_points = zeros(precision, (length(xs)*length(ys), 2))

    let
        c = 1
        for i = 1:length(xs)
            for j = 1:length(ys)
                ground_points[c, 1] = xs[i]
                ground_points[c, 2] = ys[j]
                c += 1
            end
        end
    end

    # create column ground_points[, 3] to hold source strengths for each target
    ground_points = cat(ground_points,
                        zeros(precision, size(ground_points)[1], 1);
                        dims = 2
                    )

    # populate ground_points[, 3] with sum of sources in block of size
    # block_size centered on each target
    for i = 1:size(ground_points)[1]
        xlower = Int64(ground_points[i, 1] - block_radius)
        xupper = min(Int64(ground_points[i, 1] + block_radius), ncols)
        ylower = Int64(ground_points[i, 2] - block_radius)
        yupper = min(Int64(ground_points[i, 2] + block_radius), nrows)

        # source_array is Array{Union{T, Missing}, 2} but should have no
        # missings, so not having skipmissings be okay
        ground_points[i, 3] = sum(source_array[ylower:yupper, xlower:xupper])
    end

    # get rid of ground_points with strength equal to 0
    targets = ground_points[ground_points[:, 3] .> 0, 1:3]
    targets
end

# x and y defined by targets object. Ultimately the for loop will be done by
# iterating through rows of targets object
function get_source(
        source_array::Array{Union{Missing, T}, 2} where T <: Number,
        arguments::Dict{String, Int64},
        conditional::Bool,
        condition1_present::Array{Union{Missing, T}, 2} where T <: Number,
        condition1_future::Array{Union{Missing, T}, 2} where T <: Number,
        condition2_present::Array{Union{Missing, T}, 2} where T <: Number,
        condition2_future::Array{Union{Missing, T}, 2} where T <: Number,
        comparison1::String,
        comparison2::String,
        condition1_lower::Number,
        condition1_upper::Number,
        condition2_lower::Number,
        condition2_upper::Number,
        precision::DataType;
        x::Int64,
        y::Int64,
        strength::Number
    )

    block_radius = arguments["block_radius"]
    radius = arguments["radius"]
    buffer = arguments["buffer"]
    nrows = arguments["nrows"]
    ncols = arguments["ncols"]

    source_subset = clip(source_array,
                         x = x,
                         y = y,
                         distance = radius)

    # Append missing if buffer > 0
    if buffer > 0
        ### Columns
        nrow_sub = size(source_subset)[1]
        left_col_num = max(0, min(buffer, x - radius - 1))
        right_col_num = max(0, min(buffer, ncols - (x + radius)))

        # Add left columns
        if left_col_num > 0
            source_subset = hcat(fill(missing, (nrow_sub, left_col_num)),
                                 source_subset)
        end
        # Add right columns
        if right_col_num > 0
            source_subset = hcat(source_subset,
                                 fill(missing, (nrow_sub, right_col_num)))
        end

        ### Rows
        ncol_sub = size(source_subset)[2]
        top_row_num = max(0, min(buffer, y - radius - 1))
        bottom_row_num = max(0, min(buffer, nrows - (y + radius)))

        # Add top rows
        if top_row_num > 0
            source_subset = vcat(fill(missing, (top_row_num, ncol_sub)),
                                 source_subset)
        end
        #Add bottom rows
        if bottom_row_num > 0
            source_subset = vcat(source_subset,
                                 fill(missing, (bottom_row_num, ncol_sub)))
        end
    end

    # Replace nodata vals with 0s
    source_subset[ismissing.(source_subset)] .= 0.0

    # Set any sources inside target block to 0
    xlower_sub = (radius + buffer  + 1) - block_radius
    xupper_sub = min((radius + buffer  + 1)  + block_radius, ncols)
    ylower_sub = (radius + buffer + 1)  - block_radius
    yupper_sub = min((radius + buffer  + 1)  + block_radius, nrows)

    source_subset[ylower_sub:yupper_sub, xlower_sub:xupper_sub] .= 0

    # allocate total current equal to target "strength", divide among sources
    # according to their source strengths
    source_sum = sum(source_subset[coalesce.(source_subset .> 0, false)])
    source_subset[source_subset .> 0] .=
        (source_subset[coalesce.(source_subset .> 0, false)] * strength) / source_sum

    if conditional
        xlower_buffered = Int64(max(x - radius - buffer, 1))
        xupper_buffered = Int64(min(x + radius + buffer, ncols))
        ylower_buffered = Int64(max(y - radius - buffer, 1))
        yupper_buffered = Int64(min(y + radius + buffer, nrows))
        xlower = x - block_radius
        xupper = min(x + block_radius, ncols)
        ylower = y - block_radius
        yupper = min(y + block_radius, nrows)

        source_target_match!(source_subset,
                             arguments["n_conditions"],
                             condition1_present,
                             condition1_future,
                             condition2_present,
                             condition2_future,
                             comparison1,
                             comparison2,
                             condition1_lower,
                             condition1_upper,
                             condition2_lower,
                             condition2_upper,
                             ylower,
                             yupper,
                             xlower,
                             xupper,
                             ylower_buffered,
                             yupper_buffered,
                             xlower_buffered,
                             xupper_buffered
                             )
    end
    source_subset
end

function source_target_match!(source_subset::Array{Union{T, Missing}, 2} where T <: Number,
                              n_conditions::Int64,
                              condition1_present::Array{Union{T, Missing}, 2} where T <: Number,
                              condition1_future::Array{Union{T, Missing}, 2} where T <: Number,
                              condition2_present::Array{Union{T, Missing}, 2} where T <: Number,
                              condition2_future::Array{Union{T, Missing}, 2} where T <: Number,
                              comparison1::String,
                              comparison2::String,
                              condition1_lower::Number,
                              condition1_upper::Number,
                              condition2_lower::Number,
                              condition2_upper::Number,
                              ylower::Int64,
                              yupper::Int64,
                              xlower::Int64,
                              xupper::Int64,
                              ylower_buffered::Int64,
                              yupper_buffered::Int64,
                              xlower_buffered::Int64,
                              xupper_buffered::Int64
                              )
    con1_present_subset = condition1_present[ylower_buffered:yupper_buffered,
                                             xlower_buffered:xupper_buffered]

    if comparison1 == "within"
      value1 = median(skipmissing(condition1_future[ylower:yupper, xlower:xupper]))
      source_subset[coalesce.(((con1_present_subset .- value1) .> condition1_upper) .|
          ((con1_present_subset .- value1) .< condition1_lower), false)] .= 0
    elseif comparison1 == "equals"
      value1 = mode(skipmissing(condition1_future[ylower:yupper, xlower:xupper]))
      source_subset[coalesce.(con1_present_subset .!= value1, false)] .= 0
    end

    if n_conditions == 2
      con2_present_subset = condition2_present[ylower_buffered:yupper_buffered,
                                               xlower_buffered:xupper_buffered]
      if comparison2 == "within"
          value2 = median(skipmissing(condition2_future[ylower:yupper, xlower:xupper]))
          source_subset[coalesce.(((con2_present_subset .- value2) .> condition2_upper) .|
              ((con2_present_subset .- value2) .< condition2_lower), false)] .= 0
      elseif comparison2 == "equals"
          value2 = mode(skipmissing(condition2_future[ylower:yupper, xlower:xupper]))
          source_subset[coalesce.(con2_present_subset .!= value2, false)] .= 0
      end
    end
end

function get_ground(
        arguments::Dict{String, Int64},
        precision::DataType;
        x::Int64,
        y::Int64,

    )
    radius = arguments["radius"]
    buffer = arguments["buffer"]
    distance = radius + buffer

    nrows = arguments["nrows"]
    ncols = arguments["ncols"]

    xlower = Int64(max(x - radius - buffer, 1))
    xupper = Int64(min(x + radius + buffer, ncols))
    ylower = Int64(max(y - radius - buffer, 1))
    yupper = Int64(min(y + radius + buffer, nrows))

    size_x = xupper - xlower + 1
    size_y = yupper - ylower + 1

    ground = fill(convert(precision, 0.0),
                  size_y,
                  size_x)

    new_x = min(distance + 1, x)
    new_y = min(distance + 1, y)

    ground[new_y, new_x] = Inf

    ground
end

function get_conductance(
        resistance::Array{Union{T, Missing}, 2} where T <: Number,
        arguments::Dict{String, Int64},
        x::Int64,
        y::Int64,
        resistance_file_is_conductance::Bool
    )

    radius = arguments["radius"]
    buffer = arguments["buffer"]

    resistance_clipped = clip(resistance,
                              x = x,
                              y = y,
                              distance = radius + buffer)

    if resistance_file_is_conductance
        conductance = resistance_clipped
    else
        conductance = 1 ./ resistance_clipped
    end

    conductance
end

function calculate_current(
        conductance::Array{Union{T, Missing}, 2} where T <: Number,
        source::Array{Union{T, Missing}, 2} where T <: Number,
        ground::Array{T, 2} where T <: Number,
        flags::Circuitscape.RasterFlags,
        cs_cfg::Dict{String, String},
        T::DataType
    )
    V = Int64

    # Replace missings with -9999, then convert to Array{T, 2}
    # prior to circuitscape
    conductance[ismissing.(conductance)] .= -9999
    source[ismissing.(source)] .= -9999
    conductance = convert(Array{T, 2}, conductance)
    source = convert(Array{T, 2}, source)

    # get raster data
    cellmap = conductance
    polymap = Matrix{V}(undef, 0, 0)
    source_map = source
    ground_map = ground
    points_rc = (V[], V[], V[])
    strengths = Matrix{T}(undef, 0, 0)

    included_pairs = Circuitscape.IncludeExcludePairs(:undef,
                                                      V[],
                                                      Matrix{V}(undef,0,0))

    # This is just to satisfy type requirements, most of it not used
    hbmeta = Circuitscape.RasterMeta(size(cellmap)[2],
                                     size(cellmap)[1],
                                     0.,
                                     0.,
                                     1.,
                                     -9999.,
                                     Array{Float64, 1}(undef, 1),
                                     "")

    rasterdata = Circuitscape.RasData(cellmap,
                                      polymap,
                                      source_map,
                                      ground_map,
                                      points_rc,
                                      strengths,
                                      included_pairs,
                                      hbmeta)

    # Generate advanced data
    data = Circuitscape.compute_advanced_data(rasterdata, flags, cs_cfg)

    G = data.G
    nodemap = data.nodemap
    polymap = data.polymap
    hbmeta = data.hbmeta
    sources = data.sources
    grounds = data.grounds
    finitegrounds = data.finite_grounds
    cc = data.cc
    src = data.src
    check_node = data.check_node
    source_map = data.source_map # Need it for one to all mode
    cellmap = data.cellmap

    # Flags
    is_raster = flags.is_raster
    is_alltoone = flags.is_alltoone
    is_onetoall = flags.is_onetoall
    write_v_maps = flags.outputflags.write_volt_maps
    write_c_maps = flags.outputflags.write_cur_maps
    write_cum_cur_map_only = flags.outputflags.write_cum_cur_map_only

    volt = zeros(eltype(G), size(nodemap))
    ind = findall(x -> x != 0, nodemap)
    f_local = Vector{eltype(G)}()
    solver_called = false
    voltages = Vector{eltype(G)}()
    outvolt = Circuitscape.alloc_map(hbmeta)
    outcurr = Circuitscape.alloc_map(hbmeta)

    for c in cc
        if check_node != -1 && !(check_node in c)
            continue
        end

        # a_local = laplacian(G[c, c])
        a_local = G[c,c]
        s_local = sources[c]
        g_local = grounds[c]

        if sum(s_local) == 0 || sum(g_local) == 0
            continue
        end

        if finitegrounds != [-9999.]
            f_local = finitegrounds[c]
        else
            f_local = finitegrounds
        end

        voltages = Circuitscape.multiple_solver(cs_cfg,
                                                a_local,
                                                s_local,
                                                g_local,
                                                f_local)

        local_nodemap = Circuitscape.construct_local_node_map(nodemap,
                                                              c,
                                                              polymap)

        solver_called = true

        Circuitscape.accum_currents!(outcurr,
                                     voltages,
                                     cs_cfg,
                                     a_local,
                                     voltages,
                                     f_local,
                                     local_nodemap,
                                     hbmeta)
    end

    outcurr
end

function solve_target!(
        i::Int64,
        n_targets::Int64,
        int_arguments::Dict{String, Int64},
        targets::Array{T, 2} where T <: Number,
        source_strength::Array{Union{Missing, T}, 2} where T <: Number,
        resistance::Array{Union{Missing, T}, 2} where T <: Number,
        cs_cfg::Dict{String, String},
        flags::Circuitscape.RasterFlags,
        o::Circuitscape.OutputFlags,
        calc_flow_potential::Bool,
        correct_artifacts::Bool,
        conditional::Bool,
        condition1_present::Array{Union{Missing, T}, 2} where T <: Number,
        condition1_future::Array{Union{Missing, T}, 2} where T <: Number,
        condition2_present::Array{Union{Missing, T}, 2} where T <: Number,
        condition2_future::Array{Union{Missing, T}, 2} where T <: Number,
        comparison1::String,
        comparison2::String,
        condition1_lower::Number,
        condition1_upper::Number,
        condition2_lower::Number,
        condition2_upper::Number,
        correction_array::Array{T, 2} where T <: Number,
        cum_currmap::Array{T, 3} where T <: Number,
        fp_cum_currmap::Array{T, 3}  where T <: Number,
        precision::DataType,
        resistance_file_is_conductance::Bool
    )
    ## get source
    x_coord = Int64(targets[i, 1])
    y_coord = Int64(targets[i, 2])
    source = get_source(source_strength,
                        int_arguments,
                        conditional,
                        condition1_present,
                        condition1_future,
                        condition2_present,
                        condition2_future,
                        comparison1,
                        comparison2,
                        condition1_lower,
                        condition1_upper,
                        condition2_lower,
                        condition2_upper,
                        precision,
                        x = x_coord,
                        y = y_coord,
                        strength = float(targets[i, 3]))

    ## get ground
    ground = get_ground(int_arguments,
                        precision,
                        x = x_coord,
                        y = y_coord)

    ## get conductances for Omniscape
    conductance = get_conductance(resistance,
                                  int_arguments,
                                  x_coord,
                                  y_coord,
                                  resistance_file_is_conductance)

    grid_size = size(source)

    ## Run circuitscape
    curr = calculate_current(conductance,
                             source,
                             ground,
                             flags,
                             cs_cfg,
                             precision)

    ## If normalize = True, calculate null map and normalize
    if calc_flow_potential == true
        null_conductance = convert(Array{Union{precision, Missing}, 2}, fill(1, grid_size))

        flow_potential = calculate_current(null_conductance,
                                           source,
                                           ground,
                                           flags,
                                           cs_cfg,
                                           precision)
    end

    if correct_artifacts
        correction_array2 = deepcopy(correction_array)
        lowerxcut = 1
        upperxcut = size(correction_array, 2)
        lowerycut = 1
        upperycut = size(correction_array, 1)

        if x_coord > int_arguments["ncols"] - (int_arguments["radius"] + int_arguments["buffer"])
            upperxcut = upperxcut - (upperxcut - grid_size[2])
        end

        if x_coord < int_arguments["radius"] + int_arguments["buffer"] + 1
            lowerxcut = upperxcut - grid_size[2] + 1
        end

        if y_coord > int_arguments["nrows"] - (int_arguments["radius"] + int_arguments["buffer"])
            upperycut = upperycut - (upperycut - grid_size[1])
        end

        if y_coord < int_arguments["radius"] + int_arguments["buffer"] + 1
            lowerycut = upperycut - grid_size[1] + 1
        end

        correction_array2 = correction_array[lowerycut:upperycut,
                                             lowerxcut:upperxcut]

        curr = curr .* correction_array2

        if calc_flow_potential
            flow_potential = flow_potential .* correction_array2
        end
    end

    ## Accumulate values
    xlower = max(x_coord - int_arguments["radius"] - int_arguments["buffer"], 1)
    xupper = min(x_coord + int_arguments["radius"] + int_arguments["buffer"],
                 int_arguments["ncols"])
    ylower = max(y_coord - int_arguments["radius"] - int_arguments["buffer"], 1)
    yupper = min(y_coord + int_arguments["radius"] + int_arguments["buffer"],
                 int_arguments["nrows"])

    cum_currmap[ylower:yupper, xlower:xupper, threadid()] .=
        cum_currmap[ylower:yupper, xlower:xupper, threadid()] .+ curr

    if calc_flow_potential == true
        fp_cum_currmap[ylower:yupper, xlower:xupper, threadid()] .=
            fp_cum_currmap[ylower:yupper, xlower:xupper, threadid()] .+ flow_potential
    end

end

function calc_correction(
        arguments::Dict{String, Int64},
        cs_cfg::Dict{String, String},
        flags::Circuitscape.RasterFlags,
        o,
        conditional::Bool,
        condition1_present::Array{Union{T, Missing}, 2} where T <: Number,
        condition1_future::Array{Union{T, Missing}, 2} where T <: Number,
        condition2_present::Array{Union{T, Missing}, 2}where T <: Number,
        condition2_future::Array{Union{T, Missing}, 2} where T <: Number,
        comparison1::String,
        comparison2::String,
        condition1_lower::Number,
        condition1_upper::Number,
        condition2_lower::Number,
        condition2_upper::Number,
        precision::DataType
    )
    buffer = arguments["buffer"]
    # This may not apply seamlessly in the case (if I add the option) that source strengths
    # are not adjusted by target weight, but stay the same according to their
    # original values. Something to keep in mind...

    temp_source = convert(Array{Union{precision, Missing}, 2},
                          fill(1.0,
                               arguments["radius"] * 2 + buffer * 2 + 1,
                               arguments["radius"] * 2 + buffer * 2 + 1))

    source_null = clip(temp_source,
                            x = arguments["radius"] + buffer + 1,
                            y = arguments["radius"] + buffer + 1,
                            distance = arguments["radius"])

    # Append NoData (-9999) if buffer > 0
    if buffer > 0
        column_dims = (size(source_null)[1], buffer)
        # Add columns
        source_null = hcat(fill(missing, column_dims),
                                source_null,
                                fill(missing, column_dims))

        row_dims = (buffer, size(source_null)[2])
        # Add rows
        source_null = vcat(fill(missing, row_dims),
                                source_null,
                                fill(missing, row_dims))
    end
    n_sources = sum(source_null[(!).(ismissing.(source_null))])

    source_null[ismissing.(source_null)] .= 0.0
    source_null[arguments["radius"] + arguments["buffer"] + 1,
                arguments["radius"] + arguments["buffer"] + 1] = 0.0
    source_null[source_null .!= 0.0] .= 1 / (n_sources - 1)

    source_blocked = get_source(temp_source,
                              arguments,
                              conditional,
                              condition1_present,
                              condition1_future,
                              condition2_present,
                              condition2_future,
                              comparison1,
                              comparison2,
                              condition1_lower,
                              condition1_upper,
                              condition2_lower,
                              condition2_upper,
                              precision,
                              x = (arguments["radius"] + arguments["buffer"] + 1),
                              y = (arguments["radius"] + arguments["buffer"] + 1),
                              strength = float(arguments["block_size"] ^ 2))

    conductance = clip(temp_source,
                       x = arguments["radius"] + arguments["buffer"] + 1,
                       y = arguments["radius"] + arguments["buffer"] + 1,
                       distance = arguments["radius"] + arguments["buffer"])

    ground = fill(convert(precision, 0.0),
                  size(source_null))

    ground[arguments["radius"] + arguments["buffer"] + 1,
           arguments["radius"] + arguments["buffer"] + 1] = Inf

    block_null_current = calculate_current(conductance,
                                           source_blocked,
                                           ground,
                                           flags,
                                           cs_cfg,
                                           precision)

    null_current =  calculate_current(conductance,
                                      source_null,
                                      ground,
                                      flags,
                                      cs_cfg,
                                      precision)
    null_current_total = fill(convert(precision, 0.),
                              arguments["radius"] * 2 + arguments["buffer"] * 2 + arguments["block_size"],
                              arguments["radius"] * 2 + arguments["buffer"] * 2 + arguments["block_size"])

    for i in 1:arguments["block_size"]
        for j in 1:arguments["block_size"]
            null_current_total[i:(i + arguments["radius"] * 2 + arguments["buffer"] * 2),
                               j:(j + arguments["radius"] * 2 + arguments["buffer"] * 2)] += null_current
        end
    end

    null_current_total = null_current_total[(arguments["block_radius"] + 1):(size(null_current, 1) + arguments["block_radius"]),
                                            (arguments["block_radius"] + 1):(size(null_current, 2) + arguments["block_radius"])]

    null_current_total[block_null_current .== 0.] .= 0

    null_current_total[null_current_total .== 0.0] .= 1.0
    block_null_current[block_null_current .== 0.0] .= 1.0

    correction = null_current_total ./ block_null_current

    correction
end


# Input data
data = readlines("Day04.input");
steps = map(v -> parse(Int64, v), split(data[1], ","));
boards = map(
    rows -> hcat(
        map(
            row -> map(v -> parse(Int64, v), filter(!isempty, split(row, " "))),
            rows[2:end],
        )...,
    ),
    Iterators.partition(data[2:end], 6),
);

# Part 1
stepboard = collect(Iterators.flatten(map(step -> map(v -> (step, v), boards), steps)));
applystep = (A, num) -> map(x -> x == num ? -1 : x, A);
isbingo =
    (A) -> any(
        vcat(map(col -> all(col .< 0), eachcol(A)), map(row -> all(row .< 0), eachrow(A))),
    );

candidates = filter(
    x -> (isbingo(x[2]) != isbingo(applystep(x[2], x[1]))),
    map(
        ((i, _),) -> (
            len = size(boards, 1);
            stepboard[i] = (
                stepboard[i][1],
                i > len ? applystep(stepboard[i-len][2], stepboard[i-len][1]) :
                stepboard[i][2],
            )
        ),
        enumerate(stepboard),
    ),
);

getans =
    (F, candidates) ->
        F(candidates)[1] * sum(filter(>(0), applystep(F(candidates)[2], F(candidates)[1])));

ans = getans(first, candidates);
println("Part 1: $ans")

# Part 2
ans = getans(last, candidates);
println("Part 2: $ans")


# series

# Makie

using CairoMakie, Random, LaTeXStrings
CairoMakie.activate!()
let
    Random.seed!(123)
    m = 200
    function spiral(; a = 1, n = 100, h = 0, k = 0)
        φ = LinRange(rand() + 1, 6π, n)
        h .+ rand(-1:2:1) * a * cos.(φ) ./ φ, k .+ rand(-1:2:1) * a * sin.(φ) ./ φ
    end
    x = raw"x=a\,\cos(\varphi)/\varphi,\quad "
    y = raw"y=a\,\sin(\varphi)/\varphi,\quad "
    curves = [spiral(; a = rand(), h = rand(-1:1)) for i in 1:m]

    with_theme(theme_minimal()) do
        fig = Figure(resolution = (600, 400))
        ax = Axis(fig[1, 1])
        CairoMakie.series!(curves; color = to_colormap(:Hokusai1, m), linewidth = 1.5)
        text!(latexstring(x * y * "\\varphi>0"), position = (-1, -1))
        hidedecorations!(ax; grid = false)
        display(fig)
    end
end;


module KroneckerMatrices

using Iterators

import Base:eltype,convert,size,length,
            rank,det,trace,transpose,ctranspose,inv,*,
            svdfact,svd,svdvals,eigfact,eigvals,chol,qrfact,lufact

export AbstractKronMat,AbstractKronVec,AbstractKronMatOrVec,
       KronProdMat,KronProdVec,KronSumMat,terms,sizes,full,
       kronprod,kronprod,kronsum

include("utils.jl")
include("fullkronfuns.jl")
include("abstract.jl")
include("products.jl")
include("sums.jl")
include("factorization.jl")

end


using Query
using DataFrames
using Statistics
using Test

@testset "a.b Syntax (dplyr API)" begin

    df = DataFrame(name=repeat(["John", "Sally", "Kirk"],inner=[1],outer=[2]),
                   age=vcat([10., 20., 30.],[10., 20., 30.].+3),
                   children=repeat([3,2,2],inner=[1],outer=[2]),state=[:a,:a,:a,:b,:b,:b])

    x = @from i in df begin
        @group i by i.state into g
        @select {group=key(g),mage=mean(g.age), oldest=maximum(g.age), youngest=minimum(g.age)}
        @collect DataFrame
    end

    @test x isa DataFrame
    @test size(x) == (2,4)
    @test x[1,:mage] == 20
    @test x[2,:mage] == 23
    @test x[1,:oldest] == 30
    @test x[2,:oldest] == 33
    @test x[1,:youngest] == 10
    @test x[2,:youngest] == 13
end


using DynamicalBilliards, Plots
include("../billiard_utils.jl")
include("../plot_utils.jl")
include("../randomfinitewall.jl")
include("../simulation.jl")
include("../vgate.jl")

include("agent.jl")
include("setup.jl")

bt = construct_billiard(poly, sp, ep, clr)
u = construct_controlmap(length(clr))

agt = create_agents(num_agents, length(clr)+1, init_width, init_center, u)

# Loop manually from here if you want to see step by step plots
# for j = 1:1000
    simulate_once!(agt, bt, 0.1)
# end

plot_map(agt, bt)


push!(LOAD_PATH,"../src/") # load module from local directory

using PolaronMobility
using Printf # used for some stdout

using Test

@testset "PolaronMobility" begin

include("FrohlichAlpha.jl") # Simple test of Frohlich alpha vs. literature values
include("FeynmanAthermal.jl") # Athermal Feynman tests
include("HellwarthEffectiveFrequency.jl") # Test Hellwarth et al. 1999 PRB 'B' multiple branch reduction scheme
include("FrostPolaronMobility2017.jl") # Reproduce values published in Frost 2017 PRB
include("MultipleBranches.jl") # Test explicit Oct 2019:-> explicit phonon branches
include("MultiplePhonons.jl") # Test 2021 work on multiple phonon modes
end

println("\nThat's me! If I finished without interupting, all tests have passed.")


abstract type AbstractVariable <: JuMP.AbstractVariable end

abstract type HintPoint end

# It will not really be the center and the center for z = 0 is the the same as for <h, x> = 0
struct CenterPoint{T} <: HintPoint
    h::Vector{T}
end

struct InteriorPoint{T} <: HintPoint
    h::Vector{T}
end
_β(model, h::InteriorPoint) = @variable(model)
_b(model, h::InteriorPoint) = @variable(model, [1:length(h.h)])

function polar_perspective_ellipsoid(ell, point::HintPoint, z::SpaceVariable,
                                     x::Vector{SpaceVariable})
    y = [z; x]
    H = Sets._householder(point.h)
    p = y' * H * Sets._HPH(ell) * H * y
    return Sets.perspective_dual(Sets.Householder(ell, p, point.h, z, x))
end
function polar_perspective_ellipsoid(model, Q::Symmetric{JuMP.VariableRef},
                                     point::CenterPoint, z, x)
    psd_constraint(model, Q)
    ell = Sets.Ellipsoid(Q)
    return polar_perspective_ellipsoid(ell, point, z, x)
end
function polar_perspective_ellipsoid(model, Q::Symmetric{JuMP.VariableRef},
                                     point::InteriorPoint, z, x)
    n = LinearAlgebra.checksquare(Q)
    @assert n == length(point.h)
    β = @variable(model, base_name="β")
    b = @variable(model, [1:length(point.h)], base_name="b")
    psd_constraint(model, Symmetric([β+1 b'; b Q]))
    ell = Sets.ShiftedEllipsoid(Q, b, β)
    return polar_perspective_ellipsoid(ell, point, z, x)
end

function perspective_dual_polyset(set, point::HintPoint, z::SpaceVariable,
                                  x::Vector{SpaceVariable})
    y = [z; x]
    H = Sets._householder(point.h)
    p = Sets.perspective_gauge0(set)(y => H * y)
    return Sets.perspective_dual(Sets.Householder(set, p, point.h, z, x))
end
# For `CenterPoint`, `q` should be non-perspective, need refactoring
function perspective_dual_polyset(degree, q, point::InteriorPoint, z, x)
    set = Sets.ConvexPolynomialSet(degree, q, z, x)
    perspective_dual_polyset(set, point, z, x)
end

### Polytope ###
struct Polytope <: AbstractVariable
    symmetric::Bool
    dimension::Union{Nothing, Int}
    piecewise::Union{Polyhedra.Rep, Nothing}
end
function Polytope(;
    symmetric::Bool=false,
    dimension::Union{Int, Nothing}=nothing,
    piecewise::Union{Polyhedra.Rep, Nothing}=nothing,
)
    function update_dim(object, dim_fun)
        if object !== nothing
            d = dim_fun(object)
            if dimension === nothing
                dimension = d
            elseif dimension != d
                throw(DimensionMismatch())
            end
        end
    end
    update_dim(piecewise, Polyhedra.fulldim)
    return Polytope(symmetric, dimension, piecewise)
end
Sets.space_variables(::Polytope) = nothing

function variable_set(model::JuMP.AbstractModel, ell::Polytope, space::Space,
                      space_dimension, space_polyvars)
    n = space_dimension
    if ell.symmetric
        if ell.piecewise === nothing
            set = Sets.PolarPoint(@variable(model, [1:n], base_name = "a"))
        else
            hashyperplanes(ell.piecewise) && error("hyperplanes not supported for piecewise")
            a = @variable(model, [1:nhalfspaces(ell.piecewise), 1:n], base_name="a")
            sets = [Sets.PolarPoint(a[i, :]) for i in 1:nhalfspaces(ell.piecewise)]
            set = Sets.Piecewise(sets, ell.piecewise)
            for i in eachindex(set.graph)
                for (j, v) in set.graph[i]
                    if i < j # The constraints are the same for (i, j) and (j, i)
                        λ = @variable(model, base_name="λc[$i,$j]")
                        # Ensures continuity
                        @constraint(model, sets[i].a - sets[j].a .== λ * v)
                        # Ensures convexity
                        # We need to ensure that `dot(sets[i].a, x)` is maximal
                        # over when `piece[i]`; see Proposition 5 of [R21].
                        #
                        # [R21] Raković, S. V.
                        # *Control Minkowski–Lyapunov functions*
                        # Automatica, Elsevier BV, 2021, 128, 109598
                        #
                        # In fact, we only need this to hold locally to prove
                        # convexity of each point so we can just check for every
                        # neighbor.
                        # Because <ai, x> = <aj, x> by the previous constraint,
                        # we can check the gradient.
                        @constraint(model, dot(sets[i].a, v) <= dot(sets[j].a, v), base_name="conv[$i,$j]")
                    end
                end
            end
        end
        if space == PrimalSpace
            return set
        else
            @assert space == DualSpace
            return Polyhedra.polar(set)
        end
    else
        error("Non-symmetric polytope not supported yet")
    end
end
function JuMP.value(h::Sets.PolarPoint)
    return Sets.PolarPoint(JuMP.value.(h.a))
end

### Ellipsoid ###
struct Ellipsoid <: AbstractVariable
    point::Union{Nothing, HintPoint}
    symmetric::Bool
    dimension::Union{Nothing, Int}
    guaranteed_psd::Bool # Is it already guaranteed that it is PSD ? e.g. by nth_root
    superset::Union{Sets.Ellipsoid, Nothing}
    piecewise::Union{Polyhedra.Rep, Nothing}
end
function Ellipsoid(; point::Union{Nothing, HintPoint}=nothing,
                   symmetric::Bool=false,
                   dimension::Union{Int, Nothing}=nothing,
                   superset::Union{Sets.Ellipsoid, Nothing}=nothing,
                   piecewise::Union{Polyhedra.Rep, Nothing}=nothing)
    function update_dim(object, dim_fun)
        if object !== nothing
            d = dim_fun(object)
            if dimension === nothing
                dimension = d
            elseif dimension != d
                throw(DimensionMismatch())
            end
        end
    end
    update_dim(point, point -> length(point.h))
    update_dim(superset, Sets.dimension)
    update_dim(piecewise, Polyhedra.fulldim)
    return Ellipsoid(point, symmetric, dimension, false, superset, piecewise)
end
Sets.space_variables(::Ellipsoid) = nothing

function variable_set(model::JuMP.AbstractModel, ell::Ellipsoid, space::Space,
                      space_dimension, space_polyvars)
    n = space_dimension
    function new_Q()
        # TODO, we should use constraiend variable instead in case direct mode is used.
        Q = @variable(model, [1:n, 1:n], Symmetric, base_name="Q")
        if !ell.guaranteed_psd
            psd_constraint(model, Q)
        end
        return Q
    end
    if ell.symmetric
        function new_piece()
            if space == PrimalSpace
                if ell.superset !== nothing
                    Q = Symmetric(new_Q() + ell.superset.Q)
                else
                    Q = new_Q()
                end
                return Sets.Ellipsoid(Q)
            else
                ell.superset === nothing || error("superset not supported in dual space")
                @assert space == DualSpace
                return Sets.Ellipsoid(new_Q())
            end
        end
        if ell.piecewise === nothing
            set = new_piece()
        else
            hashyperplanes(ell.piecewise) && error("hyperplanes not supported for piecewise")
            sets = [new_piece() for i in 1:nhalfspaces(ell.piecewise)]
            set = Sets.Piecewise(sets, ell.piecewise)
            DynamicPolynomials.@polyvar x[1:n]
            q = [quad_form(set.Q, x) for set in sets]
            for i in eachindex(set.graph)
                for (j, v) in set.graph[i]
                    if i < j # The constraints are the same for (i, j) and (j, i)
                        @constraint(model, q[i] == q[j], domain = @set x'v == 0)
                        # v corresponds to `-n_ij` in LCSS paper
                        Δ = sets[i].Q * v - sets[j].Q * v
                        inter = set.pieces[i] ∩ set.pieces[j]
                        h = HalfSpace(Δ, zero(eltype(Δ)))
                        @constraint(model, inter ⊆ h)
                    end
                end
            end
        end
        if space == PrimalSpace
            return set
        else
            return Polyhedra.polar(set)
        end
    else
        ell.superset === nothing || error("superset not supported for non-symmetric Ellipsoid")
        ell.piecewise === nothing || error("piecewise not supported for non-symmetric Ellipsoid")
        if space == PrimalSpace
            error("Non-symmetric ellipsoid non implemented yet, use `Ellipsoid(symmetric=true)`.")
        else
            @assert space == DualSpace
            if ell.point === nothing
                throw(ArgumentError("Specify a point for nonsymmetric ellipsoid, e.g. `Ellipsoid(point=InteriorPoint([1.0, 0.0]))"))
            end
            return polar_perspective_ellipsoid(model, new_Q(), ell.point,
                                               data(model).perspective_polyvar,
                                               space_polyvars)
        end
    end
end
function JuMP.value(ell::Sets.Ellipsoid)
    return Sets.Ellipsoid(Symmetric(JuMP.value.(ell.Q)))
end

### PolySet ###
struct PolySet <: AbstractVariable
    point::Union{Nothing, HintPoint}
    symmetric::Bool
    degree::Int
    dimension::Union{Nothing, Int}
    convex::Bool
    variables::Union{Nothing, Vector{SpaceVariable}}
    superset::Union{Sets.PolySet, Nothing}
    basis::Type
end
function PolySet(; point::Union{Nothing, HintPoint}=nothing,
                 symmetric::Bool=false,
                 degree::Union{Int, Nothing}=nothing,
                 dimension::Union{Int, Nothing}=nothing,
                 convex::Bool=false,
                 variables::Union{Vector{SpaceVariable}, Nothing}=nothing,
                 superset::Union{Sets.PolySet, Nothing}=nothing,
                 basis::Type=MultivariateBases.MonomialBasis)
    if degree === nothing
        error("Degree of PolySet not specified, use PolySet(degree=..., ...)")
    end
    if isodd(degree)
        throw(ArgumentError("Degree of PolySet not even"))
    end
    if dimension === nothing
        if point !== nothing
            dimension = length(point.h)
        end
    end
    if superset !== nothing
        if dimension === nothing
            dimension = Sets.dimension(superset)
        elseif dimension != Sets.dimension(superset)
            throw(DimensionMismatch())
        end
        if variables === nothing
            variables = Sets.space_variables(superset)
        elseif variables != Sets.space_variables(superset)
            error("Space variables set does not correspond to superset space variables.")
        end
    end
    return PolySet(point, symmetric, degree, dimension, convex, variables, superset, basis)
end
Sets.space_variables(p::PolySet) = p.variables

function constrain_convex(model, p, vars)
    hessian = MP.differentiate(p, vars, 2)
    # We do not just do `@constraint(model, p in SOSConvex())` as we would
    # like to have access to the PSD matrix of variables for the det volume heuristic
    y = [MP.similarvariable(eltype(hessian), gensym()) for i in 1:LinearAlgebra.checksquare(hessian)]
    q = dot(y, hessian * y)
    X = SumOfSquares.Certificate.monomials_half_newton_polytope(MP.monomials(q), (y,))
    # If `X` is empty, we will need the following bridge
    JuMP.add_bridge(model, SumOfSquares.Bridges.Constraint.EmptyBridge)
    # If `length(X)` is 2, we will need the following bridge
    JuMP.add_bridge(model, SumOfSquares.Bridges.Constraint.PositiveSemidefinite2x2Bridge)
    set = SumOfSquares.matrix_cone(MOI.PositiveSemidefiniteConeTriangle,
                                   length(X))
    Q = @variable(model, [1:MOI.dimension(set)])
    @constraint(model, Q in set)
    s = SumOfSquares.build_gram_matrix(
        Q, MonomialBasis(X), MOI.PositiveSemidefiniteConeTriangle, Float64)
    @constraint(model, q == s)
    return MultivariateMoments.getmat(s)
end

function variable_set(model::JuMP.AbstractModel, set::PolySet, space::Space,
                      space_dimension, space_polyvars)
    n = space_dimension
    d = data(model)
    # General all monomials of degree `degree`, we don't want monomials of
    # lower degree as the polynomial is homogeneous
    @assert iseven(set.degree)
    if set.symmetric
        monos = MP.monomials(space_polyvars, div(set.degree, 2))
    else
        monos = MP.monomials(lift_space_variables(d, space_polyvars),
                             div(set.degree, 2))
    end
    basis = MultivariateBases.basis_covering_monomials(set.basis, monos)
    # TODO If `set.convex` and `set.symmetric`, no need for the poly to be SOS, see Lemma 6.33 of [BPT12]
    p = @variable(model, variable_type=SOSPoly(basis))
    if set.convex
        set.superset === nothing || error("superset not supported for convex PolySet")
        if set.symmetric
            convexity_proof = constrain_convex(model, p, space_polyvars)
            if space == PrimalSpace
                return Sets.ConvexPolySet(set.degree, p, convexity_proof)
            else
                @assert space == DualSpace
                return Polyhedra.polar(Sets.ConvexPolySet(set.degree, p, convexity_proof))
            end
        else
            constrain_convex(model, MP.subs(p, d.perspective_polyvar => 1),
                             space_polyvars)
            if space == PrimalSpace
                error("Non-symmetric PolySet in PrimalSpace not implemented yet")
            else
                @assert space == DualSpace
                if set.point === nothing
                    throw(ArgumentError("Specify a point for nonsymmetric polyset, e.g. `PolySet(point=InteriorPoint([1.0, 0.0]))"))
                end
                return perspective_dual_polyset(set.degree, p, set.point, d.perspective_polyvar, space_polyvars)
            end
        end
    else
        if set.symmetric
            if space == PrimalSpace
                if set.superset !== nothing
                    p = SetProg.SumOfSquares.gram_operate(+, set.superset.p, p)
                end
                return Sets.PolySet(set.degree, p)
            else
                error("Non-convex PolySet not supported in $space")
            end
        else
            error("Non-convex nonsymmetric PolySet not implemented yet.")
        end
    end
end
_value(convexity_proof::Nothing) = nothing
function _value(convexity_proof::MultivariateMoments.SymMatrix)
    return MultivariateMoments.SymMatrix(JuMP.value.(convexity_proof.Q),
                                         convexity_proof.n)
end
function JuMP.value(set::Sets.PolySet)
    return Sets.PolySet(set.degree, JuMP.value(set.p))
end
function JuMP.value(set::Sets.ConvexPolySet)
    return Sets.ConvexPolySet(set.degree, JuMP.value(set.p), _value(set.convexity_proof))
end
function JuMP.value(set::Sets.Polar)
    return Polyhedra.polar(JuMP.value(Polyhedra.polar(set)))
end
function JuMP.value(set::Sets.PerspectiveDual)
    return Sets.perspective_dual(JuMP.value(Sets.perspective_dual(set)))
end
function JuMP.value(set::Sets.Householder)
    return Sets.Householder(JuMP.value(set.set), JuMP.value(set.p), set.h,
                            set.z, set.x)
end
function JuMP.value(set::Sets.ShiftedEllipsoid)
    return Sets.ShiftedEllipsoid(Symmetric(JuMP.value.(set.Q)),
                                 JuMP.value.(set.b), JuMP.value(set.β))
end
function JuMP.value(set::Sets.ConvexPolynomialSet)
    return Sets.ConvexPolynomialSet(set.degree, JuMP.value(set.q), set.z, set.x)
end
function JuMP.value(set::Sets.Piecewise)
    return Sets.Piecewise(JuMP.value.(set.sets), set.polytope, set.pieces, set.graph)
end

### SetVariableRef ###
mutable struct SetVariableRef{M <: JuMP.AbstractModel,
                              S <: AbstractVariable} <: JuMP.AbstractVariableRef
    model::M
    set::S
    name::String
    # `variable` is typically `Sets.AbstractSet{JuMP.VariableRef}` but it can
    # also be `Sets.AbstractSet{JuMP.AffExpr}` with `PolySet(superset = ...)`.
    variable::Union{Nothing, Sets.AbstractSet}
    space_index::Union{Nothing, SpaceIndex}
end
JuMP.name(vref::SetVariableRef) = vref.name
function JuMP.build_variable(_error::Function, info::JuMP.VariableInfo, set::AbstractVariable)
    @assert !info.has_lb && !info.has_ub && !info.has_fix && !info.binary && !info.integer && !info.has_start
    return set
end
function JuMP.add_variable(model::JuMP.AbstractModel, set::AbstractVariable, name::String)
    vref = SetVariableRef(model, set, name, nothing, nothing)
    d = data(model)
    @assert d.state == Modeling
    push!(d.variables, vref)
    return vref
end
JuMP.value(vref::SetVariableRef) = JuMP.value(vref.variable)

function clear_spaces(vref::SetVariableRef)
    vref.space_index = nothing
end
Sets.space_variables(::SetVariableRef) = nothing
Sets.perspective_variable(::SetVariableRef) = nothing
function create_spaces(vref::SetVariableRef, spaces::Spaces)
    if vref.space_index === nothing
        if Sets.space_variables(vref.set) === nothing
            if vref.set.dimension === nothing
                vref.space_index = new_space(spaces)
            else
                vref.space_index = new_space(spaces, vref.set.dimension)
            end
        else
            vref.space_index = new_space(spaces, Sets.space_variables(vref.set))
        end
    end
    return vref.space_index
end
space_index(vref::SetVariableRef) = vref.space_index

function load(model::JuMP.AbstractModel, vref::SetVariableRef)
    d = data(model)
    vref.variable = variable_set(model, vref.set, d.space,
                                 space_dimension(d.spaces, vref.space_index),
                                 space_polyvars(d.spaces, vref.space_index))
end

variablify(p::Sets.Projection) = Polyhedra.project(variablify(p.set), p.indices)


#=

In this file we simulate a simple rocket landing

Units in kg, m, s

=#
include("..\\src\\rocket\\rocket-setup.jl")
include("..\\src\\dynamics\\trajectory-setup.jl")
include("..\\src\\objective\\LQR_objective.jl")
include("..\\src\\constraints\\constraintManager.jl")
include("..\\src\\auglag\\auglag-core.jl")
include("..\\src\\solver\\AL-Primal-Main-Solver.jl")
include("..\\src\\solver\\QP-AffineEquality-Solver.jl")
include("..\\src\\other_utils\\parsePrimalDual.jl")

include("..\\src\\results\\trajectoryParsing.jl")
include("..\\src\\results\\plotTrajectory.jl")
include("..\\src\\results\\plotConstraintViolation.jl")
include("..\\src\\results\\plotObjective.jl")

# Based on the Falcon 9
# 549,054 kg (Mass)
# 282 s (Specific Impulse)
# select "y" as the vertical
mass = 1
isp = 1
grav = [0; -9.81]
deltaTime = 1
rocket = rocket_simple(mass, isp, grav, deltaTime)

# in m
# The Karman Line (100 km)
rocketStart = [2.0; 20.0; 0.0; 0.0]
rocketEnd = [0.0; 0.0; 0.0; 0.0]#[-5.0; 0.0; 0.0; 0.0]

uHover = mass * grav

# Number of time steps to discretize the trajectory
NSteps = 40
# Initialize the trajectory with a line
initTraj = initializeTraj(rocketStart, rocketEnd, uHover, uHover, NSteps)

# Use a Linear Quadratic Regulator as the cost function
lqrQMat = 10 * Diagonal(I, size(rocketStart, 1))
lqrRMat = 2 * Diagonal(I, Int64(size(rocketStart, 1) / 2))
costFun = makeLQR_TrajReferenced(lqrQMat, lqrRMat, NSteps, initTraj)

ADyn, BDyn = rocketDynamicsFull(rocket, rocketStart, rocketEnd, NSteps)
dynConstraint = AL_AffineEquality(ADyn, BDyn)
lambdaInit = zeros(size(BDyn))

cMRocket = constraintManager([dynConstraint], [lambdaInit])

penaltyStart = 1.0

# Test that the evaluations work
println("\n--------------------------------------------")
println(" Testing constraints are inputted correctly ")
println("--------------------------------------------")
println("Starting constraint violation: ")
println([evalConstraints(cMRocket, initTraj, penaltyStart)])
println("Starting gradient of constraints: ")
println(evalGradConstraints(cMRocket, initTraj, penaltyStart))
println("Starting hessian of constraints: ")
println(evalHessConstraints(cMRocket, initTraj))

# Equiped with the constraint term and the objective term, I now build the
# Augmented Lagrangian
alRocket = augLag(costFun, cMRocket, penaltyStart)

# Test that the evaluations work
println("--------------------------------------------")
println("           Testing AL is Functional         ")
println("--------------------------------------------")
println("Evaluating augmented lagrangian: ")
println([evalAL(alRocket, initTraj)])
println("Evaluating gradient of augmented lagrangian: ")
println(evalGradAL(alRocket, initTraj))
println("Evaluating hessian of augmented lagrangian: ")
println(size(evalHessAl(alRocket, initTraj)))

# Next we select resonable solver parameters
currSolveParams = solverParams(0.1, 0.5,
                                6, 2,
                                10^-4,
                                10, 10^6,
                                2.5, 2, 0.2, 0.2, 0.4)
solParamPrint(currSolveParams)
println()
println()

println("--------------------------------------------")
println("             Beginning Solve                ")
println("--------------------------------------------")

# # Solve the Trajectory Optimization problem
# trajStates, resArr = ALPrimalNewtonMain(initTraj, alRocket, currSolveParams)
#
#
# if true
#     # Get the parsed list of trajectories
#     nDim = size(grav, 1)
#     pltTraj, pltCV, pltObj, plts, pltv, pltu = batchPlot(trajStates, nDim)
# end

costQ = costFun.lqr_qp.QR_Full
costP = - (initTraj' * costQ)'
trajLambdaSolved = solveQP_AffineEq(costQ, costP, ADyn, BDyn)
trajSolved = parsePrimalDualVec(trajLambdaSolved, size(initTraj, 1))

if true
    # Get the parsed list of trajectories
    nDim = size(grav, 1)
    pltTraj, pltCV, pltObj, plts, pltv, pltu = batchPlot(
                                    [initTraj, trajSolved.primals], nDim)
end

if true
    header = "DirectLanding"
    saveBulk(pltTraj, pltCV, pltObj, plts, pltv, pltu, header)
end


# LPA_wb_plus.jl
# Label propagation algorithm for weighted bipartite networks that finds modularity.
# Contains the LPAwb+ and the DIRTLPAwb+ algorithms
# Author :  Stephen Beckett ( https://github.com/sjbeckett/weighted-modularity-LPAwbPLUS )
# MIT License

###############################################

function BarbersMatrix(MATRIX)
	return MATRIX - (sum(MATRIX,2)*sum(MATRIX,1))/sum(MATRIX);
end

###############################################

function WEIGHTEDMODULARITY(BMatrix,Matsum,redlabels::Array{Int64,1},bluelabels::Array{Float64,2})
	#see equation 8
	holdsum = 0;

	for rr = 1:length(redlabels)
	    for cc = 1:length(bluelabels)
	        kroneckerdelta = redlabels[rr] == bluelabels[cc];
	        holdsum += BMatrix[rr,cc] * kroneckerdelta;
	    end
	end
	return holdsum/Matsum;
end

###############################################
function TRACE(MATRIX)
	return sum(diag(MATRIX))
end

###############################################

function WEIGHTEDMODULARITY2(BMatrix,Matsum,redlabels,bluelabels)
	#see equation 9
	UNIred = unique(redlabels)
	Lred = length(redlabels)
	UNIblu = unique(bluelabels)
	Lblu = length(bluelabels)
	LABELMAT1 = zeros(length(UNIred),Lred)
	LABELMAT2 = zeros(Lblu,length(UNIblu))

	for aa = 1:Lred
		LABELMAT1[findin(UNIred,redlabels[aa]),aa] = 1
	end

	for aa = 1:Lblu
		LABELMAT2[aa,findin(UNIblu,bluelabels[aa])] = 1
	end

	return TRACE(LABELMAT1 * BMatrix * LABELMAT2)/Matsum

end

###############################################
function DIVISION(redlabels::Array{Int64,1},bluelabels::Array{Float64,2})
	divisionsFound= Int64[];

	for aa = 1:min(maximum(redlabels),maximum(bluelabels))
		if( (isempty(findin(redlabels,aa))==false) && (isempty(findin(bluelabels,aa))==false) )
			push!(divisionsFound,aa)
		end
	end

	return divisionsFound;
end

###############################################

function StageTwo_LPAwbdash(row_marginals,col_marginals,MATRIX,BMatrix,Matsum,redlabels,bluelabels, Qb_now)

	divisionsFound = DIVISION(redlabels,bluelabels);
	NUMdiv = length(divisionsFound);
	IterateFlag = 1;
	while IterateFlag == 1
		CombinedDivisionsThisTime = 0;
		if NUMdiv > 1
			for div1check = 1:NUMdiv-1
				Mod1 = divisionsFound[div1check]
				for div2check = div1check+1:NUMdiv
					CHECK_RED = copy(redlabels);
					CHECK_RED[findin(redlabels,Mod1)] = divisionsFound[div2check];
					CHECK_BLUE = copy(bluelabels);
					CHECK_BLUE[findin(bluelabels,Mod1)] = divisionsFound[div2check];
					QQ = WEIGHTEDMODULARITY2(BMatrix,Matsum,CHECKRED,CHECKBLUE);
					if QQ > Qb_now #If a division will improve modularity - find the best way to do this
                    				FoundBetter = 0;
                    				for aa = 1:NUMdiv
                        				CHECK_RED2 = copy(redlabels);
                        				CHECK_RED2[findin(redlabels,divisionsFound[aa])] = Mod1;
                        				CHECK_BLUE2 = copy(bluelabels);
                        				CHECK_BLUE2[findin(bluelabels,divisionsFound[aa])] = Mod1;
                        				if WEIGHTEDMODULARITY2(BMatrix,Matsum,CHECK_RED2,CHECK_BLUE2) > QQ
                            					FoundBetter = 1;
                        				end
                        				CHECK_RED2 = copy(redlabels);
                        				CHECK_RED2[findin(redlabels,divisionsFound[aa])] = divisionsFound[div2check];
                        				CHECK_BLUE2 = copy(bluelabels);
                        				CHECK_BLUE2[findin(bluelabels,divisionsFound[aa])] = divisionsFound[div2check];
                        				if WEIGHTEDMODULARITY2(BMatrix,Matsum,CHECK_RED2,CHECK_BLUE2) > QQ
                            					FoundBetter = 1;
                        				end
                    				end
                    				if FoundBetter == 0 #If no better configuration found - JOIN.
                        				redlabels = copy(CHECK_RED);
                        				bluelabels = copy(CHECK_BLUE);
                        				CombinedDivisionsThisTime += 1;
                    				end
                			end
				end
			end
			if CombinedDivisionsThisTime == 0 #If no divisions were joined move on
		            IterateFlag = 0;
		        end
		else
		IterateFlag = 0;
		end

		(redlabels,bluelabels,Qb_now) = StageOne_LPAwbdash(row_marginals,col_marginals,MATRIX,BMatrix,Matsum,redlabels,bluelabels);
		divisionsFound = DIVISION(redlabels,bluelabels);
		NUMdiv = length(divisionsFound);
	end

return redlabels, bluelabels, Qb_now;

end

###############################################

function StageOne_LPAwbdash(row_marginals,col_marginals,MATRIX,BMatrix,Matsum,redlabels,bluelabels)
	#Create storage containers for total marginals attached to each red(row)
	#label and blue(column) label

	BLUELABELLENGTH = length(bluelabels);
	REDLABELLENGTH = length(redlabels);
	TotalRedDegrees = NaN*ones(maximum(redlabels),1);

	TotalBlueDegrees = NaN*ones(max(BLUELABELLENGTH,REDLABELLENGTH),1);

	#Fill up these containers according to current labels
	#Red
	for aa = 1:REDLABELLENGTH
	    if isnan(TotalRedDegrees[redlabels[aa]])
        	TotalRedDegrees[redlabels[aa]] = row_marginals[aa];
	    else
        	TotalRedDegrees[redlabels[aa]] += row_marginals[aa];
	    end
	end

	#Blue
	if sum(isnan(bluelabels)) != BLUELABELLENGTH #occurs first time through as blue nodes unlabelled
	    for bb = 1:BLUELABELLENGTH
        	if isnan(TotalBlueDegrees[bluelabels[bb]])
	            TotalBlueDegrees[bluelabels[bb]] = col_marginals[bb];
        	else
        	    TotalBlueDegrees[bluelabels[bb]] += col_marginals[bb];
        	end
	    end
	else
	    TotalBlueDegrees = zeros(max(BLUELABELLENGTH,REDLABELLENGTH),1);
	end

	#locally maximise modularity!
	(redlabels,bluelabels,Qb_now) = LOCALMAXIMISATION(row_marginals,col_marginals,MATRIX,BMatrix,Matsum,redlabels,bluelabels,TotalRedDegrees,TotalBlueDegrees);

	return redlabels, bluelabels, Qb_now;

end

###############################################

function LOCALMAXIMISATION(row_marginals,col_marginals,MATRIX,BMatrix,Matsum,redlabels,bluelabels,TotalRedDegrees,TotalBlueDegrees)

	#Find score for current partition
	QbAfter = WEIGHTEDMODULARITY2(BMatrix,Matsum,redlabels,bluelabels);

	IterateFlag = 1;
	while IterateFlag == 1
		#Save old information
		QbBefore = QbAfter;
		old_redlabels = copy(redlabels);
		old_bluelabels = copy(bluelabels);
		old_TRD = copy(TotalRedDegrees);
		old_TBD = copy(TotalBlueDegrees); #13

		#Update Blue Nodes using red node information (see equation 10)
		bluelabelchoices = 1:maximum(redlabels);

		for bb = 1:length(bluelabels)
			if isnan(bluelabels[bb]) == false
				TotalBlueDegrees[bluelabels[bb]] -= col_marginals[bb];
			end
			changebluelabeltest = NaN*ones(1,length(bluelabelchoices));

			for ww = 1:length(bluelabelchoices)
				changebluelabeltest[ww] = sum( (redlabels .== bluelabelchoices[ww]) .* MATRIX[:,bb])  -  col_marginals[bb]*TotalRedDegrees[bluelabelchoices[ww]]/Matsum ;
			end

			#assign new label based on maximisation of above condition

			labels = findin(changebluelabeltest, maximum(changebluelabeltest));
			newlabelindex = labels[rand(1:length(labels),1)];
			bluelabels[bb] = bluelabelchoices[newlabelindex[1]];

			if bluelabels[bb] > length(TotalBlueDegrees)
	                	TotalBlueDegrees[bluelabels[bb]] = 0;
	        	end

        		#Update total marginals on new labelling
            		TotalBlueDegrees[bluelabels[bb]] += col_marginals[bb];
        	end#40

		#Now update red node labels based on blue node information (see equation 10)
		redlabelchoices = 1:maximum(bluelabels);

		for aa = 1:length(redlabels)
			TotalRedDegrees[redlabels[aa]] -= row_marginals[aa];
			changeredlabeltest = NaN*ones(1,length(redlabelchoices));

			for ww = 1:length(redlabelchoices)
				changeredlabeltest[ww] = sum( (bluelabels .== redlabelchoices[ww]) .* MATRIX[aa,:])  -  row_marginals[aa]*TotalBlueDegrees[redlabelchoices[ww]]/Matsum ;
			end

			#assign new label based on maximisation of above condition
			labels = findin(changeredlabeltest,maximum(changeredlabeltest));
			newlabelindex = labels[rand(1:length(labels),1)];
 			redlabels[aa] = redlabelchoices[newlabelindex[1]];

			if redlabels[aa] > length(TotalRedDegrees)
				TotalRedDegrees[redlabels[aa]] = 0;
			end
			TotalRedDegrees[redlabels[aa]] +=  row_marginals[aa];
		end

		#Find the new modularity score based on node label updates.
        	QbAfter = WEIGHTEDMODULARITY(BMatrix,Matsum,redlabels,bluelabels);

        	#If this modularity is not as good as previous stop iterating and
        	#use that previous best information
        	if QbAfter <= QbBefore
            		redlabels = copy(old_redlabels);
	            	bluelabels = copy(old_bluelabels);
        	    	TotalRedDegrees = copy(old_TRD);
        	    	TotalBlueDegrees = copy(old_TBD);
        	    	IterateFlag = 0;
        	end

	end

	Qb_now = QbAfter;

	return redlabels, bluelabels, Qb_now;
end

###############################################

function LPA_wb_plus(MATRIX, initialmoduleguess = NaN)

	#Make sure the smallest matrix dimension represent the red labels by making them the rows (If matrix is transposed here, will be transposed back at the end)
	flipped = 0;
	if size(MATRIX,1) > size(MATRIX,2)
		MATRIX = MATRIX';
		flipped = 1;
	end

	Matsum = sum(MATRIX);
	col_marginals = sum(MATRIX,1);
	row_marginals = sum(MATRIX,2);
	BMatrix = BarbersMatrix(MATRIX);

	#initiliase labels
	bluelabels = [NaN*ones(1,size(MATRIX,2))];
	if isnan(initialmoduleguess)
		redlabels = [1:size(MATRIX,1)];
	else
		redlabels = rand(1:(initialmoduleguess+1),size(MATRIX,1));
	end

	#Run Phase 1: Locally update lables to maximise Qb
	(redlabels,bluelabels,Qb_now) = StageOne_LPAwbdash(row_marginals,col_marginals,MATRIX,BMatrix,Matsum,redlabels,bluelabels);

	#Run Phase 2: Connect divisions from top-down if it improves Qb, then run
	#phase 1 again. Repeat till Qb cannot be improved.
	(redlabels,bluelabels,Qb_now) = StageTwo_LPAwbdash(row_marginals,col_marginals,MATRIX,BMatrix,Matsum,redlabels,bluelabels,Qb_now);

	if flipped==1 #If matrix was flipped, swap row and column labels
    		holder = copy(redlabels);
		redlabels = bluelabels;
		bluelabels = holder;
	end

	return redlabels, bluelabels, Qb_now;
end

###############################################

function DIRT_LPA_wb_plus(MATRIX, mini = 4, reps = 10)

	A = LPA_wb_plus(MATRIX);

	mods = length(unique(A[1]));

	if (mods-mini) > 0
		for aa = mini:mods
			for bb = 1:reps
				B = LPA_wb_plus(MATRIX,aa)
				if B[3] > A[3]
					A = B
				end
			end
		end
	end

	return A[1], A[2], A[3];
end