text:
export solve

# TODO: better names for the functions

function hamiltonianinterior(a, model::HJBTwoDim, v, t, x, xidx, Δx)
    i,j = xidx
    K = [length(xi) for xi in x]
    invdx = 1.0 ./ Δx
    hdx2 = 0.5 ./ Δx.^2

    @inbounds begin
        idxi = K[2]*(i-1) + j
        idxj1f = idxi + K[2]; idxj1b = idxi - K[2]
        idxj2f = idxi + 1;    idxj2b = idxi - 1

        xij = [x[1][i], x[2][j]]
        bval = model.b(t,xij,a)
        sval2 = model.σ(t,xij,a).^2
        coeff1f = sval2[1]*hdx2[1] + max(bval[1],0.)*invdx[1]
        coeff1b = sval2[1]*hdx2[1] - min(bval[1],0.)*invdx[1]
        coeff2f = sval2[2]*hdx2[2] + max(bval[2],0.)*invdx[2]
        coeff2b = sval2[2]*hdx2[2] - min(bval[2],0.)*invdx[2]
        return (coeff1f*(v[idxi]-v[idxj1f]) + coeff1b*(v[idxi]-v[idxj1b]) +
                coeff2f*(v[idxi]-v[idxj2f]) + coeff2b*(v[idxi]-v[idxj2b]) -
                model.f(t,xij,a))
    end
end

function setIJV!{T<:Real}(I::Vector{Int},J::Vector{Int},V::Vector{T},
                          Ival::Int,Jval::Int,Vval::T,counter::Int)
    counter = counter+1
    I[counter] = Ival
    J[counter] = Jval
    V[counter] = Vval

    return counter
end

function updateboundarysystem!{T<:Real}(I, J, V, rhs, model, v, t::T, x,
                                        a1::Vector{T}, a2::Vector{T},
                                        Δt::T, Δx::Vector{T})
    taux = Δt ./Δx
    htaux2 = 0.5*Δt ./ Δx.^2
    K = [length(xi) for xi in x]

    counter = 0
    # Dirichlet conditions for x_1 = {xmin, xmax}
    for i = [1, K[1]], j = 1:K[2]
        idxi = K[2]*(i-1) + j
        xij = [x[1][i], x[2][j]]
        counter = setIJV!(I,J,V,idxi,idxi,1.0,counter)
        # TODO: we could move rhs assignment outside this function
        # so it doesn't get called on every loop in the Newton solver
        rhs[idxi] = model.Dbound(t, xij)

    end

    # Dirichlet conditions for x_2 = {xmin, xmax}
    for i = 2:K[1]-1, j = [1, K[2]]
        idxi = K[2]*(i-1)+j
        xij = [x[1][i], x[2][j]]
        counter = setIJV!(I,J,V,idxi,idxi,1.0,counter)
        # TODO: we could rhs assignment outside this function
        # so it doesn't get called on every loop in the Newton solver
        rhs[idxi] = model.Dbound(t, xij)
    end

    @assert counter == length(V)
end

function updateinteriorsystem!{T<:Real}(I, J, V, rhs, model, v, t::T, x,
                                        a1::Vector{T}, a2::Vector{T},
                                        Δt::T, Δx::Vector{T})
    # Updates:
    # rhs    = value function at previous timestep + f at current timestep
    # I,J,V  = vectors for creating sparse system matrix (see ?sparse)

    # Input
    # model  = HJBTwoDim object
    # v      = value function at previous timestep
    # t      = value of (forward) time
    # x      = tuple of vectors of x-values
    # a1     = policy-values on interior, element 1
    # a2     = policy-values on interior, element 2
    # Δt     = time-step size
    # Δx     = spacial step length

    taux = Δt ./Δx
    htaux2 = 0.5*Δt ./ Δx.^2
    K = [length(xi) for xi in x]

    counter = 0

    for i = 2:K[1]-1, j = 2:K[2]-1
        @inbounds begin
            idxi = K[2]*(i-1) + j
            idxj1f = idxi + K[2]; idxj1b = idxi - K[2]
            idxj2f = idxi + 1;    idxj2b = idxi - 1
            xij = [x[1][i], x[2][j]]
            aij = [a1[idxi], a2[idxi]]

            bval = model.b(t,xij,aij)
            sval2 = model.σ(t,xij,aij).^2
            coeff1f = -(sval2[1]*htaux2[1] + max(bval[1],0.)*taux[1])
            coeff1b = -(sval2[1]*htaux2[1] - min(bval[1],0.)*taux[1])
            coeff2f = -(sval2[2]*htaux2[2] + max(bval[2],0.)*taux[2])
            coeff2b = -(sval2[2]*htaux2[2] - min(bval[2],0.)*taux[2])
            coeff0 = 1.0-(coeff1f+coeff1b + coeff2f+coeff2b)

            # TODO: does it make a performance difference what order I put these in?
            counter = setIJV!(I,J,V,idxi,idxi,coeff0, counter)
            counter = setIJV!(I,J,V,idxi,idxj1f,coeff1f, counter)
            counter = setIJV!(I,J,V,idxi,idxj1b,coeff1b, counter)
            counter = setIJV!(I,J,V,idxi,idxj2f,coeff2f, counter)
            counter = setIJV!(I,J,V,idxi,idxj2b,coeff2b, counter)

            rhs[idxi] = v[idxi] + Δt*model.f(t,xij,aij)
        end
    end
    @assert counter == length(V)
end

function optimizepol!(pol::Tuple, objective::Function, v, model::HJBTwoDim, t,
                      x, Δx, idxi::Int;
                      tol = 1e-4, maxiter = 1000, optimizer = LBFGS,
                      verbose=false)
    initialguess = [p[idxi] for p in pol]
    g!(a, out) = ForwardDiff.gradient!(out, objective, a)
    diffobj = DifferentiableFunction(objective, g!)
    res = optimize(diffobj, initialguess,
                   model.amin, model.amax, Fminbox(),
                   optimizer = optimizer,
                   show_trace=verbose, extended_trace=verbose,
                   f_tol=tol,x_tol=tol,g_tol=tol,
                   iterations=100,
                   optimizer_o = OptimizationOptions(f_tol=tol,x_tol=tol,
                                                     g_tol=tol, iterations=100,
                                                     show_trace=verbose,
                                                     extended_trace=verbose))
    for (i,val) in enumerate(Optim.minimizer(res))
        pol[i][idxi] = val
    end
end

function updateinteriorpol!(pol::Tuple, v, model::HJBTwoDim, t, x::Tuple, Δx::Vector;
                            tol = 1e-4, maxiter = 1000,
                            optimizer = LBFGS, verbose=false)
    # Loops over each x value and optimises the control
    # TODO: Should we instead optimize the whole control vector
    # by considering the sum of the individual objectives
    # (the gradient should be diagonal?)

    K = [length(xi) for xi in x]
    for i = 2:K[1]-1, j = 2:K[2]-1
        @inbounds begin
            objective(a) = hamiltonianinterior(a, model, v, t, x, [i, j], Δx)
            idxi = K[2]*(i-1) + j
            optimizepol!(pol, objective, v, model, t, x, Δx, idxi;
                         tol=tol, verbose=verbose,
                         optimizer=optimizer, maxiter=maxiter)
        end
    end
end

function policynewtonupdate{T<:Real}(model::HJBTwoDim{T},
                                     v, a1, a2, x::Tuple{Vector{T},Vector{T}},
                                     Δx::Vector{T}, Δt, ti::Int;
                                     tol = 1e-4,
                                     scale = 1.0,
                                     maxpolicyiter::Int = 10)
    # v  = value function at previous time-step
    # an = policy function at previous time-step / initial guess for update
    t = (ti-1)*Δt
    n = length(v)
    K = [length(xi) for xi in x]
    @assert length(a1) == n && length(a2) == n

    # Elements in sparse system matrix (n\times n) size
    interiornnz = 5*prod(K-2)
    boundarynnz = 2*(sum(K)-2)
    #totnnz = interiornnz + boundarynnz
    #I = zeros(Int, totnnz); J = zeros(I); V = zeros(T, totnnz)
    Ii = zeros(Int, interiornnz); Ji = zeros(Ii); Vi = zeros(T, interiornnz)
    Ib = zeros(Int, boundarynnz); Jb = zeros(Ib); Vb = zeros(T, boundarynnz)

    rhs = zeros(v)

    # TODO: copy or pass reference?
    pol1 = copy(a1)
    pol2 = copy(a2)

    vnew = copy(v)

    for k in 0:maxpolicyiter
        updateinteriorpol!((pol1, pol2), vnew, model, t, x, Δx)
        updateboundarysystem!(Ib,Jb,Vb, rhs, model, v, t, x, pol1, pol2, Δt, Δx)
        updateinteriorsystem!(Ii,Ji,Vi, rhs, model, v, t, x, pol1, pol2, Δt, Δx)

        Mat = sparse([Ib;Ii],[Jb;Ji],[Vb;Vi],n,n,(x,y)->error("Each index should be unique"))

        # TODO: Use Krylov solver for high-dimensional PDEs?
        vold = vnew
        vnew = Mat\rhs

        vchange = maximum(abs(vnew-vold)./max(1.,abs(vnew)))
        if vchange < tol && k>0
            break
        end
    end
    # TODO: do we need this one?
    updateinteriorpol!((pol1, pol2), vnew, model, t, x, Δx)

    return vnew, pol1, pol2
end

function policytimestep{T<:Real}(model::HJBTwoDim{T},
                                 v, x::Tuple{Vector{T},Vector{T}},
                                 Δx::Vector{T}, Δt, ti::Int, avals::Tuple)
    # v  = value function at previous time-step
    # an = policy function at previous time-step / initial guess for update
    t = (ti-1)*Δt

    n = length(v)
    K = [length(xi) for xi in x]

    # Elements in sparse system matrix (n\times n) size
    interiornnz = 5*prod(K-2)
    boundarynnz = 2*(sum(K)-2)
    #totnnz = interiornnz + boundarynnz
    #I = zeros(Int, totnnz); J = zeros(I); V = zeros(T, totnnz)
    Ii = zeros(Int, interiornnz); Ji = zeros(Ii); Vi = zeros(T, interiornnz)
    Ib = zeros(Int, boundarynnz); Jb = zeros(Ib); Vb = zeros(T, boundarynnz)

    a1const = zeros(n)
    a2const = zeros(n)
    newind1 = ones(Int, n) # Indices for pol1
    newind2 = ones(Int, n) # Indices for pol2
    @inbounds vnew = -maxintfloat(T)*ones(v)
    indkeep = zeros(Bool, length(vnew)) # Indices to keep vold value

    rhs = zeros(v)

    for i = 1:length(avals[1]), j = 1:length(avals[2])
        a1const[:] = avals[1][i]
        a2const[:] = avals[2][j]

        updateboundarysystem!(Ib,Jb,Vb, rhs, model, v, t, x, a1const, a2const, Δt, Δx)
        updateinteriorsystem!(Ii,Ji,Vi, rhs, model, v, t, x, a1const, a2const, Δt, Δx)

        Mat = sparse([Ib;Ii],[Jb;Ji],[Vb;Vi],n,n,(x,y)->error("Each index should be unique"))

        # TODO: Use Krylov solver for high-dimensional PDEs?
        vold = vnew
        vnew = Mat\rhs

        indkeep[:] = vold .> vnew
        vnew[indkeep] = vold[indkeep]
        newind1[!indkeep] = i
        newind2[!indkeep] = j
    end

    pol1 = avals[1][newind1]
    pol2 = avals[2][newind2]

    return vnew, pol1, pol2
end

function timeloopiteration(model::HJBTwoDim, K::Vector{Int}, N::Int,
                           Δt, vinit, x::Tuple, Δx::Vector)
    # Pass v and pol by reference?
    v = zeros(length(vinit), N+1)
    # No policy at t = T
    pol1 = zeros(length(vinit), N)
    pol2 = zeros(length(vinit), N)
    pol = (pol1, pol2)

    @inbounds v[:,N+1] = vinit # We use forward time t instead of backward time t

    # initial guess for control
    pol1init = fill(0.5*(model.amax[1]+model.amin[1]), prod(K))
    pol2init = fill(0.5*(model.amax[2]+model.amin[2]), prod(K))
    @inbounds v[:,N], pol1[:,N], pol2[:,N] = policynewtonupdate(model, v[:,N+1], pol1init, pol2init,
                                                                x, Δx, Δt, N)

    for j = N-1:-1:1
        # t = (j-1)*Δt
        # TODO: pass v-column, pol-column by reference?
        @inbounds (v[:,j], pol1[:,j],
                   pol2[:,j]) = policynewtonupdate(model, v[:,j+1], pol1[:,j+1], pol2[:,j+1],
                                                   x, Δx, Δt, j)
    end

    return v, pol
end

function timeloopiteration(model::HJBTwoDim, K::Vector{Int}, N::Int,
                           Δt, vinit, x::Tuple, Δx::Vector, avals::Tuple)
    # Pass v and pol by reference?
    v = zeros(length(vinit), N+1)
    # No policy at t = T
    pol1 = zeros(length(vinit), N)
    pol2 = zeros(length(vinit), N)
    pol = (pol1, pol2)

    @inbounds v[:,N+1] = vinit

    @inbounds v[:,N], pol1[:,N], pol2[:,N] = policytimestep(model, v[:,N+1],
                                                            x, Δx, Δt, N, avals)

    for j = N-1:-1:1
        # t = (j-1)*Δt
        # TODO: pass v-column, pol-column by reference?
        @inbounds (v[:,j], pol1[:,j],
                   pol2[:,j]) = policytimestep(model, v[:,j+1],
                                               x, Δx, Δt, j, avals)
    end

    return v, pol
end

function solve{T1<:Real}(model::HJBTwoDim{T1}, K::Vector{Int}, N::Int)
    # K   = number of points in each direction of the space domain
    # N+1 = number of points in time domain
    x1 = linspace(model.xmin[1], model.xmax[1], K[1])
    x2 = linspace(model.xmin[2], model.xmax[2], K[2])
    x = (collect(x1), collect(x2))
    Δx = (model.xmax-model.xmin)./(K-1) # TODO: use diff(x) to accomodate non-uniform grid
    Δt = model.T/N # TODO: introduce non-uniform timesteps?

    vinit = zeros(T1, prod(K))
    for i = 1:K[1], j = 1:K[2]
        @inbounds begin
            idx = K[2]*(i-1)+j
            xij = [x[1][i], x[2][j]]
            vinit[idx] = model.g(xij)
        end
    end

    v, pol = timeloopiteration(model, K, N, Δt, vinit, x, Δx)
    return v, pol
end

solve(model::HJBTwoDim, K::Int, N::Int) = solve(model, [K,K], N)

function solve{T1<:Real}(model::HJBTwoDim{T1}, K::Vector{Int}, N::Int,
                         avals::Tuple)
    # K   = number of points in each direction of the space domain
    # N+1 = number of points in time domain
    x1 = linspace(model.xmin[1], model.xmax[1], K[1])
    x2 = linspace(model.xmin[2], model.xmax[2], K[2])
    x = (collect(x1), collect(x2))
    Δx = (model.xmax-model.xmin)./(K-1) # TODO: use diff(x) to accomodate non-uniform grid
    Δt = model.T/N # TODO: introduce non-uniform timesteps?

    vinit = zeros(T1, prod(K))
    for i = 1:K[1], j = 1:K[2]
        @inbounds begin
            idx = K[2]*(i-1)+j
            xij = [x[1][i], x[2][j]]
            vinit[idx] = model.g(xij)
        end
    end

    v, pol = timeloopiteration(model, K, N, Δt, vinit, x, Δx, avals)
    return v, pol
end

max_stars_repo_name: anriseth/HJBSolver

max_issues_repo_name: anriseth/HJBSolver

max_forks_repo_name: anriseth/HJBSolver


text:
#Versions:
#Plots v0.28.2
#PyPlot v2.8.2
using Plots; pyplot()

#colors extracted from the Julia logo
jred = [RGB(213/255,99/255,92/255), RGB(203/255,60/255,51/255)]
jgreen = [RGB(96/255,173/255,81/255), RGB(56/255,152/255,38/255)]
jblue = [RGB(102/255,130/255, 233/255), RGB(64/255,99/255,216/255)]
jpurple = [RGB(170/255,121/255,193/255),RGB(149/255,88/255,178/255)]
jtext = [[jpurple[2],jpurple[2],jpurple[2]],[jpurple[2],jpurple[2],jpurple[2]]]
jdet = [:white, :gray]

function plot_frame(bc = :transparent)

    #detector geometry
    top=3.5
    top2=2
    bot=6.8
    inner=2
    middle=5.3
    outer=6.3
    ratio=4
    lwidth=5

    plot(xlims = (-6.5,8.5), ylims = (-9.15,5.85), legend = false, size = (429,429),
    showaxis = false, grid = false, background_color = bc)
    if bc == :transparent
        #plot detector background
        rectangle(w, h, x, y) = Shape(x .+ [0,w,w,0], y .+ [0,0,h,h])
        plot!(rectangle(2*outer,top2+bot,-outer,-bot), color = jdet[1], linecolor = :transparent)
        plot!(Shape([-outer,outer,middle,-middle],[top2-0.1,top2-0.1,top,top]), color = jdet[1], linecolor = :transparent)
        for i in inner:0.1:middle; plot!([i*cos(t) for t in 0:pi/100:2pi], [i/ratio*sin(t) + top for t in 0:pi/100:2pi], linewidth = 3, color = jdet[1]); end
        for i in 0:0.1:outer; plot!([i*cos(t) for t in 0:pi/100:2pi], [i/ratio*sin(t) - bot for t in 0:pi/100:2pi], linewidth = 3, color = jdet[1]); end
    end
    plot!()
end

function plot_detector()

    #detector geometry
    top=3.5
    top2=2
    bot=6.8
    inner=2
    middle=5.3
    outer=6.3
    ratio=4
    lwidth=5

    #plot detector lines and contact
    for i in 0:0.05:0.95 plot!([i*inner*cos(t) for t in 0:pi/100:2pi], [i*inner/ratio*sin(t) + top for t in 0:pi/100:2pi], linewidth = 3, color = jblue[1]) end
    plot!([inner*cos(t) for t in 0:pi/100:2pi], [inner/ratio*sin(t)+top for t in 0:pi/100:2pi], linewidth = lwidth, color = jblue[2])
    plot!([middle*cos(t) for t in 0:pi/100:2pi], [middle/ratio*sin(t)+top for t in 0:pi/100:2pi], linewidth = lwidth, color = jdet[2])
    plot!(vcat([-middle,-outer],[outer*cos(t) for t in -pi:pi/100:0],[outer,middle]), vcat([top,top2],[outer/ratio*sin(t)-bot for t in -pi:pi/100:0],[top2,top]), linewidth = lwidth, color = jdet[2])
end

function SSD(k::Integer)
    if k >= 420 return 3 end
    if k >= 300 return 2 end
    if k >= 150 return 1 end
end

function plot_rest(k::Integer, label::Bool)

    #detector geometry
    top=3.5
    top2=2
    bot=6.8
    inner=2
    middle=5.3
    outer=6.3
    ratio=4

    #plot drift paths of electrons and holes
    ex = [0,0]
    ey = [sqrt(25-6.25)-5, top]
    hx = [6.3*cos(t) for t in -2pi/3.2:-pi/100:-pi]
    hy = [6.3/4*sin(t)-3.6 for t in -2pi/3.2:-pi/100:-pi]

    #plot incoming γ ray
    γmax = 9      #length of γ ray
    ϕ0 = -0.8-pi  #phase of γ ray
    α= -pi*1/2.8  #rotation angle for γ ray
    x = collect(γmax:-0.02:2.9)
    y = sin.(x * 0.7pi .+ ϕ0)
    newx = vcat(4 .+ x * cos(α) .+ y * sin(α), [4.5+1.5*cos(t) for t in pi/2-0.1:-pi/100:-pi/1.5])
    newy = vcat(-4 .- x * sin(α) .+ y * cos(α), [-2.6+1.5*sin(t) for t in pi/2-0.1:-pi/100:-pi/1.5])

    lwidth = 5    #width of lines

    if k > 400
        kk = k - 400
        hx = [6.3*cos(t) for t in -2pi/3.2:-pi/300:-pi][1:min(kk,end)]
        hy = [6.3/4*sin(t)-3.6 for t in -2pi/3.2:-pi/300:-pi][1:min(kk,end)]
        ex = [0,0]
        ey = [sqrt(25-6.25)-5, min(sqrt(25-6.25)-5 + (top - sqrt(25-6.25)+5)*kk/length(-2pi/3.2:-pi/300:-pi),top)]
        plot!(ex,ey, width = lwidth, color = jgreen[2])
        plot!(hx,hy, width = lwidth, color = jred[2])
    end
    plot_detector()

    plot!(newx[1:min(k,length(x))], newy[1:min(k,length(x))], width = 2*lwidth, color = :white)
    plot!(newx[min(k,length(x)):min(k,end)], newy[min(k,length(x)):min(k,end)], width = 2.3*lwidth, color = :white)
    plot!(newx[1:min(k,end)], newy[1:min(k,end)], width = lwidth, color = jtext[2][2])

    #plot the three Julia dots
    if k > 350
        msize = min(k - 350,78)
        scatter!([2.5],[-5], color = jpurple[1], markerstrokecolor = jpurple[2], markerstrokewidth = lwidth, markersize = msize)
        scatter!([-2.5],[-5], color = jred[1], markerstrokecolor = jred[2], markerstrokewidth = lwidth, markersize = msize)
        scatter!([0],[sqrt(25-6.25)-5], color = jgreen[1], markerstrokecolor = jgreen[2], markerstrokewidth = lwidth, markersize = msize)
    end

    if label
        if k >= 150
            annotate!([(8.15, 2.1, Plots.text("olid", 50, jtext[1][1], :left)),
               (6.9,-0.45, Plots.text("tate", 50, jtext[1][2], :left)),
               (6.55,-3.1, Plots.text("etectors", 50, jtext[1][3], :left))][1:SSD(k)])
        end

        plot!(xlims = (-6.5,16.6), ylims = (-8.5,5.2), size = (649,391))
    end
    plot!()
end

function get_logo(; animate::Bool=false, label::Bool=true, transparent=true, step::Int = 5)
    if animate
        ks::Vector{Int} = collect(0:step:600)
        if !(600 in ks) push!(ks, 600) end
        return @animate for k in ks
            k = (k == 0 ? k = 600 : k)
            plot_frame(:white)
            plot_rest(k, label)
        end
    else
        plot_frame(transparent ? :transparent : :white)
        plot_rest(600,label)
    end
end

"""
    get_and_save_logo( [; animate::Bool=false, label::Bool=true])

Compute and save the SolidStateDetectors.jl logo.

# Keyword arguments
- `animate::Bool=false`: if set to `true`, the logo will be saved as animated GIF, otherwise as static SVG file.
- `label::Bool=true`: if set to `true`, the words "Solid State Detectors" will appear on the right of the logo.

"""
function get_and_save_logo(; animate::Bool = false, label::Bool = true, transparent::Bool = true, fps = 15, step::Int = 5, loop::Bool = false)
    logo = get_logo(animate = animate, label = label, transparent = transparent, step = step);

    basename = "logo" * (label ? "" : "_no_name")
    animate ? gif(logo, basename*".gif", fps = fps, loop = loop) : savefig(basename*".svg")
end

# get_and_save_logo(animate = true, label = false, fps = 300, step = 5, loop = false)

ext: jl


text:
# # Schelling's segregation model

# ![](schelling.gif)

# In this introductory example we demonstrate Agents.jl's architecture and
# features through building
# the following definition of Schelling's segregation model:

# * Agents belong to one of two groups (0 or 1).
# * The agents live in a two-dimensional Moore grid (8 neighbors per node).
# * If an agent is in the same group with at least three neighbors, then it is happy.
# * If an agent is unhappy, it keeps moving to new locations until it is happy.

# Schelling's model shows that even small preferences of agents to have neighbors
# belonging to the same group (e.g. preferring that at least 30% of neighbors to
# be in the same group) could lead to total segregation of neighborhoods.

# This model is also available as [`Models.schelling`](@ref).

# ## Defining the agent type

using Agents, AgentsPlots
gr() # hide

mutable struct SchellingAgent <: AbstractAgent
    id::Int # The identifier number of the agent
    pos::Tuple{Int,Int} # The x, y location of the agent on a 2D grid
    mood::Bool # whether the agent is happy in its node. (true = happy)
    group::Int # The group of the agent,  determines mood as it interacts with neighbors
end

# Notice that the position of this Agent type is a `Tuple{Int,Int}` because
# we will use a `GridSpace`.

# We added two more fields for this model, namely a `mood` field which will
# store `true` for a happy agent and `false` for an unhappy one, and an `group`
# field which stores `0` or `1` representing two groups.

# ## Creating a space

# For this example, we will be using a Moore 2D grid, e.g.

space = GridSpace((10, 10), moore = true)

# ## Creating an ABM

# To make our model we follow the instructions of [`AgentBasedModel`](@ref).
# We also want to include a property `min_to_be_happy` in our model, and so we have:

properties = Dict(:min_to_be_happy => 3)
schelling = ABM(SchellingAgent, space; properties = properties)

# Here we used the default scheduler (which is also the fastest one) to create
# the model. We could instead try to activate the agents according to their
# property `:group`, so that all agents of group 1 act first.
# We would then use the scheduler [`property_activation`](@ref) like so:

schelling2 = ABM(
    SchellingAgent,
    space;
    properties = properties,
    scheduler = property_activation(:group),
)

# Notice that `property_activation` accepts an argument and returns a function,
# which is why we didn't just give `property_activation` to `scheduler`.

# ## Creating the ABM through a function

# Here we put the model instantiation in a function so that
# it will be easy to recreate the model and change its parameters.

# In addition, inside this function, we populate the model with some agents.
# We also change the scheduler to [`random_activation`](@ref).
# Because the function is defined based on keywords,
# it will be of further use in [`paramscan`](@ref) below.

function initialize(; numagents = 320, griddims = (20, 20), min_to_be_happy = 3)
    space = GridSpace(griddims, moore = true)
    properties = Dict(:min_to_be_happy => min_to_be_happy)
    model = ABM(SchellingAgent, space;
                properties = properties, scheduler = random_activation)
    ## populate the model with agents, adding equal amount of the two types of agents
    ## at random positions in the model
    for n in 1:numagents
        agent = SchellingAgent(n, (1, 1), false, n < numagents / 2 ? 1 : 2)
        add_agent_single!(agent, model)
    end
    return model
end
nothing # hide

# Notice that the position that an agent is initialized does not matter
# in this example.
# This is because we use [`add_agent_single!`](@ref), which places the agent in a random,
# empty location on the grid, thus updating its position.

# ## Defining a step function

# Finally, we define a _step_ function to determine what happens to an
# agent when activated.

function agent_step!(agent, model)
    agent.mood == true && return # do nothing if already happy
    minhappy = model.min_to_be_happy
    neighbor_cells = node_neighbors(agent, model)
    count_neighbors_same_group = 0
    ## For each neighbor, get group and compare to current agent's group
    ## and increment count_neighbors_same_group as appropriately.
    for neighbor_cell in neighbor_cells
        node_contents = get_node_contents(neighbor_cell, model)
        ## Skip iteration if the node is empty.
        length(node_contents) == 0 && continue
        ## Otherwise, get the first agent in the node...
        agent_id = node_contents[1]
        ## ...and increment count_neighbors_same_group if the neighbor's group is
        ## the same.
        neighbor_agent_group = model[agent_id].group
        if neighbor_agent_group == agent.group
            count_neighbors_same_group += 1
        end
    end
    ## After counting the neighbors, decide whether or not to move the agent.
    ## If count_neighbors_same_group is at least the min_to_be_happy, set the
    ## mood to true. Otherwise, move the agent to a random node.
    if count_neighbors_same_group ≥ minhappy
        agent.mood = true
    else
        move_agent_single!(agent, model)
    end
    return
end
nothing # hide

# For the purpose of this implementation of Schelling's segregation model,
# we only need an agent step function.

# When defining `agent_step!`, we used some of the built-in functions of Agents.jl,
# such as [`node_neighbors`](@ref) that returns the neighboring nodes of the
# node on which the agent resides, [`get_node_contents`](@ref) that returns the
# IDs of the agents on a given node, and [`move_agent_single!`](@ref) which moves
# agents to random empty nodes on the grid. A full list of built-in functions
# and their explanations are available in the [API](@ref) page.

# ## Stepping the model

# Let's initialize the model with 370 agents on a 20 by 20 grid.

model = initialize()

# We can advance the model one step
step!(model, agent_step!)

# Or for three steps
step!(model, agent_step!, 3)

# ## Running the model and collecting data

# We can use the [`run!`](@ref) function with keywords to run the model for
# multiple steps and collect values of our desired fields from every agent
# and put these data in a `DataFrame` object.
# We define vector of `Symbols`
# for the agent fields that we want to collect as data
adata = [:pos, :mood, :group]

model = initialize()
data, _ = run!(model, agent_step!, 5; adata = adata)
data[1:10, :] # print only a few rows

# We could also use functions in `adata`, for example we can define
x(agent) = agent.pos[1]
model = initialize()
adata = [x, :mood, :group]
data, _ = run!(model, agent_step!, 5; adata = adata)
data[1:10, :]

# With the above `adata` vector, we collected all agent's data.
# We can instead collect aggregated data for the agents.
# For example, let's only get the number of happy individuals, and the
# maximum of the "x" (not very interesting, but anyway!)

model = initialize();
adata = [(:mood, sum), (x, maximum)]
data, _ = run!(model, agent_step!, 5; adata = adata)
data

# Other examples in the documentation are more realistic, with a much more meaningful
# collected data. Don't forget to use the function [`aggname`](@ref) to access the
# columns of the resulting dataframe by name.

# ## Visualizing the data

# We can use the [`plotabm`](@ref) function to plot the distribution of agents on a
# 2D grid at every generation, via the
# [AgentsPlots](https://github.com/JuliaDynamics/AgentsPlots.jl) package.
# Let's color the two groups orange and blue and make one a square and the other a circle.
groupcolor(a) = a.group == 1 ? :blue : :orange
groupmarker(a) = a.group == 1 ? :circle : :square
plotabm(model; ac = groupcolor, am = groupmarker, as = 4)

# ## Animating the evolution

# The function [`plotabm`](@ref) can be used to make your own animations
cd(@__DIR__) #src
model = initialize();
anim = @animate for i in 0:10
    p1 = plotabm(model; ac = groupcolor, am = groupmarker, as = 4)
    title!(p1, "step $(i)")
    step!(model, agent_step!, 1)
end

gif(anim, "schelling.gif", fps = 2)

# ## Replicates and parallel computing

# We can run replicates of a simulation and collect all of them in a single `DataFrame`.
# To that end, we only need to specify `replicates` in the `run!` function:

model = initialize(numagents = 370, griddims = (20, 20), min_to_be_happy = 3)
data, _ = run!(model, agent_step!, 5; adata = adata, replicates = 3)
data[(end - 10):end, :]

# It is possible to run the replicates in parallel.
# For that, we should start julia with `julia -p n` where is the number
# of processing cores. Alternatively, we can define the number of cores from
# within a Julia session:

# ```julia
# using Distributed
# addprocs(4)
# ```

# For distributed computing to work, all definitions must be preceded with
# `@everywhere`, e.g.

# ```julia
# @everywhere using Agents
# @everywhere mutable struct SchellingAgent ...
# ```

# Then we can tell the `run!` function to run replicates in parallel:

# ```julia
# data, _ = run!(model, agent_step!, 2, adata=adata,
#                replicates=5, parallel=true)
# ```

# ## Scanning parameter ranges

# We often are interested in the effect of different parameters on the behavior of an
# agent-based model. `Agents.jl` provides the function [`paramscan`](@ref) to automatically explore
# the effect of different parameter values.

# We have already defined our model initialization function as `initialize`.
# We now also define a processing function, that returns the percentage of
# happy agents:

happyperc(moods) = count(x -> x == true, moods) / length(moods)

adata = [(:mood, happyperc)]
parameters =
    Dict(:min_to_be_happy => collect(2:5), :numagents => [200, 300], :griddims => (20, 20))

data, _ = paramscan(parameters, initialize; adata = adata, n = 3, agent_step! = agent_step!)
data

# `paramscan` also allows running replicates per parameter setting:

data, _ = paramscan(
    parameters,
    initialize;
    adata = adata,
    n = 3,
    agent_step! = agent_step!,
    replicates = 3,
)

data[(end - 10):end, :]

# We can combine all replicates with an aggregating function, such as mean, using
# the `groupby` and `combine` functions from the `DataFrames` package:

using DataFrames: groupby, combine, Not, select!
using Statistics: mean
gd = groupby(data,[:step, :min_to_be_happy, :numagents])
data_mean = combine(gd,[:happyperc_mood,:replicate] .=> mean)

select!(data_mean, Not(:replicate_mean))

# Note that the second argument takes the column names on which to split the data,
# i.e., it denotes which columns should not be aggregated. It should include
# the `:step` column and any parameter that changes among simulations. But it should
# not include the `:replicate` column.
# So in principle what we are doing here is simply averaging our result across the replicates.

# ## Launching the interactive application
# Given the definitions we have already created for a normal study of the Schelling model,
# it is almost trivial to launch an interactive application for it.
# First, we load `InteractiveChaos` to access `interactive_abm`
# %% #src
# ```julia
# using InteractiveChaos
# using GLMakie # we choose OpenGL as plotting backend
# ```

# Then, we define a dictionary that maps some model-level parameters to a range of potential
# values, so that we can interactively change them.
parange = Dict(:min_to_be_happy => 0:8)

# Due to the different plotting backend (Plots.jl vs Makie.jl) we redefine some of the
# plotting functions (in the near future this won't be necessary, as everything will be Makie.jl based)

groupcolor(a) = a.group == 1 ? :blue : :orange
groupmarker(a) = a.group == 1 ? :circle : :rect

# We define the `alabels` so that we can simple see the plotted timeseries with a
# shorter name (since the defaults can get large)
adata = [(:mood, sum), (x, mean)]
alabels = ["happy", "avg. x"]

model = initialize(; numagents = 300) # fresh model, noone happy

# ```julia
# scene, adf, modeldf =
# interactive_abm(model, agent_step!, dummystep, parange;
#                 ac = groupcolor, am = groupmarker, as = 1,
#                 adata = adata, alabels = alabels)
# ```
#
# ```@raw html
# <video width="100%" height="auto" controls autoplay loop>
# <source src="https://raw.githubusercontent.com/JuliaDynamics/JuliaDynamics/master/videos/agents/schelling_app.mp4?raw=true" type="video/mp4">
# </video>
# ```


text:
@recipe function f(r::PerronFrobenius.RectangularBinningTransferOperator)
    seriestype  :=  :heatmap
    xlabel --> "state #i"
    ylabel --> "state #j"
    r.transfermatrix
end


text:
#!/usr/bin/env julia
#=
Run Graze from the command line.

Usage:
cli.jl URL
=#
import Graze: parse, FeedError, HTTPError, NotImplementedError

function main(arguments=ARGS)
    #=
    Run Graze.

    input:
        arguments: An array of command line arguments.
    =#
    if length(arguments) != 1
        println("USAGE:\n\tcli.jl URL")
        exit(1)
    end

    feed = parse(arguments[1])

    try
        feed = parse(arguments[1])
    catch exception
        if isa(exception, HTTPError)
            println("Unable to access '$(arguments[1])'")
        elseif isa(exception, FeedError)
            println("Error while parsing feed")
        elseif isa(exception, NotImplementedError)
            println("Looks like I'm not done making this yet")
        else
            println("Something went wrong:\n$exception")
        end

        exit(1)
    end

    println(feed)
end

if !isinteractive()
    #  It's not quite __name__ == '__main__', but it will do
    main()
end


text:
include("./imports.jl")

# TODO - ADD TEST BELOW
# @testset ExtendedTestSet "get_affine" begin
#     @testset ExtendedTestSet "get_affine" begin
#     end
# end

@testset ExtendedTestSet "io_orientation" begin
    @testset ExtendedTestSet "io_orientation" begin
        eye = [
            1 0 0 0
            0 1 0 0
            0 0 1 0
            0 0 0 1
        ]
        answer = [
            1.0 1.0
            2.0 1.0
            3.0 1.0
        ]
        @test io_orientation(eye) == answer
    end

    # TODO - FIX TEST BELOW
    # @testset ExtendedTestSet "io_orientation" begin
    #     aff =  [1 0 0 1; 1 0 0 1; 1 0 0 1; 1 0 0 1]
    #     answer = [2.0 1.0; NaN NaN; NaN NaN]
    #     test = io_orientation(aff)
    #     @test test == answer
    # end
end

@testset ExtendedTestSet "axcodes2ornt" begin
    @testset ExtendedTestSet "axcodes2ornt" begin
        answer = [
            1.0 1.0
            2.0 1.0
            3.0 1.0
        ]
        test = axcodes2ornt(("R", "A", "S"))
        @test test == answer
    end

    @testset ExtendedTestSet "axcodes2ornt" begin
        answer = [
            2.0 1.0
            1.0 -1.0
            3.0 1.0
        ]
        test = axcodes2ornt(('F', 'L', 'U'), (('L', 'R'), ('B', 'F'), ('D', 'U')))
        @test test == answer
    end
end

@testset ExtendedTestSet "ornt2axcodes" begin
    @testset ExtendedTestSet "ornt2axcodes" begin
        ornt = [1.0 1.0
                2.0 1.0
                3.0 1.0]
        answer = ["R", "A", "S"]
        test = ornt2axcodes(ornt)
        @test test == answer
    end

    # TODO: DOUBLE CHECK THIS TEST
    # @testset ExtendedTestSet "ornt2axcodes" begin
    #     ornt = [2.0 1.0
    #             1.0 -1.0
    #             3.0 1.0]
    #     answer = ["F", "L", "U"]
    #     test = ornt2axcodes(ornt, (('L','R'),('B','F'),('D','U')))
    #     @test test == answer
    # end
end

@testset ExtendedTestSet "ornt_transform" begin
    @testset ExtendedTestSet "ornt_transform" begin
        src = [
            1.0 1.0
            2.0 1.0
            3.0 1.0
        ]
        dst = [
            2.0 1.0
            1.0 -1.0
            3.0 1.0
        ]
        answer = [
            2.0 -1.0
            1.0 1.0
            3.0 1.0
        ]
        test = ornt_transform(src, dst)
        @test test == answer
    end

    @testset ExtendedTestSet "ornt_transform" begin
        src = [
            1.0 1.0
            2.0 1.0
            3.0 1.0
        ]
        dst = [
            2.0 1.0
            1.0 1.0
            3.0 1.0
        ]
        answer = [
            2.0 1.0
            1.0 1.0
            3.0 1.0
        ]
        test = ornt_transform(src, dst)
        @test test == answer
    end
end

@testset ExtendedTestSet "apply_orientation" begin
    @testset ExtendedTestSet "apply_orientation" begin
        d1 = [
            1 1 1 1
            1 1 1 1
            0 0 0 0
            1 0 0 1
        ]
        d2 = [
            1 1 1 1
            1 1 1 1
            0 0 0 0
            1 0 0 1
        ]
        data = cat(d1, d2; dims=3)
        o = [
            2.0 1.0
            1.0 1.0
            3.0 1.0
        ]
        a1 = [
            1 1 0 1
            1 1 0 0
            1 1 0 0
            1 1 0 1
        ]
        a2 = [
            1 1 0 1
            1 1 0 0
            1 1 0 0
            1 1 0 1
        ]
        answer = cat(a1, a2; dims=3)
        test = apply_orientation(data, o)
        @test test == answer
    end
end

# TODO: FIX THIS TEST AND PROBLEMS WITH `I`
# @testset ExtendedTestSet "inv_ornt_aff" begin
#     @testset ExtendedTestSet "inv_ornt_aff" begin
#         ornt = [1 0 0
# 	            0 1 0
# 	            0 0 1]
#         shape = [4, 4, 2]
#         answer = [0.0 0.0 0.0 1.5
#                   1.0 0.0 0.0 0.0
#                   0.0 0.0 0.0 0.5
#                   0.0 0.0 0.0 1.0]
#         test = inv_ornt_aff(ornt, shape)
#         @test test == answer
#     end

#     @testset ExtendedTestSet "inv_ornt_aff" begin
#         ornt = [1 0 1
# 	            1 0 1
# 	            0 0 1]
#         shape = [4, 4, 2]
#         answer = [0.0 0.0 0.0 1.5
#                   0.0 0.0 0.0 1.5
#                   0.0 0.0 0.0 0.5
#                   0.0 0.0 0.0 1.0]
#         test = inv_ornt_aff(ornt, shape)
#         @test test == answer
#     end
# end


text:
using Merlin

type POSModel
  word_f
  char_f
  sent_f
end

function POSModel(path)
  T = Float32
  #word_f = Lookup(T, 500000, 100)
  word_f = Lookup("$(path)/nyt100.lst", T)
  char_f = [Lookup(T,100,10),
            Window2D(10,5,1,1,0,2),
            Linear(T,50,50),
            Max(2)]
  sent_f = [Window2D(150,5,1,1,0,2),
            Linear(T,750,300),
            ReLU(),
            Linear(T,300,45)]
  POSModel(word_f, char_f, sent_f)
end

function forward(m::POSModel, tokens::Vector{Token})
  wordvec = map(t -> t.wordid, tokens) |> Variable
  wordmat = m.word_f(wordvec)
  charvecs = map(tokens) do t
    v = Variable(t.charids)
    m.char_f(v)
  end
  charmat = charvecs |> Concat(2)
  [wordmat, charmat] |> Concat(1) |> m.sent_f
end

function update!(m::POSModel, opt::Optimizer)
  for f in (m.word_f, m.char_f, m.sent_f)
    Merlin.update!(opt, f)
  end
end

max_stars_repo_name: hshindo/POSTagging-benchmark

max_issues_repo_name: hshindo/POSTagging-benchmark

max_forks_repo_name: hshindo/POSTagging-benchmark


text:
module SmoothedSpectralAbscissa
using LinearAlgebra, Roots

abstract type AbstractIOMatrices end
struct IOIdentity <:AbstractIOMatrices end

abstract type OptimMethod end
struct OptimOrder2 end
struct OptimNewton end

@inline function spectral_abscissa(A::Matrix{R}) where R
   convert(R,maximum(real.(eigvals(A))))
end

# trace of matrix product
@inline function trace_of_product(A::Matrix{R},B::Matrix{R}) where R
    n=size(A,1)
    acc=zero(R)
    for i in 1:n, j in 1:n
        @inbounds acc += A[i,j]*B[j,i]
    end
    return acc
end
# remove value from the diagonal
# equivalent to copy!(A,A-s*I)
@inline function subtract_diagonal!(A::Matrix{R},s::R) where R
    n=size(A,1)
    @simd for i in 1:n
        @inbounds A[i,i] -= s
    end
    return nothing
end

"""
        default_eps_ssa(A::Matrix{Float64}) -> eps_ssa

default value for the ``\\varepsilon`` used to compute the SSA. It scales with the matrix
size, so that it takes the value 0.05 for a ``150 \\times 150`` matrix.
"""
default_eps_ssa(A::Matrix{<:Real}) = 0.01 * 150.0 / size(A,1)

struct SSAAlloc{R,C}
    R::Matrix{R}
    R_alloc::Matrix{R}
    Rt::Matrix{R}
    Z::Matrix{R}
    Zt::Matrix{R}
    D::Matrix{R}
    D_alloc::Matrix{R}
    Dt::Matrix{R}
    At::Matrix{R}
    P::Matrix{R}
    Q::Matrix{R}
    YZt_alloc::Matrix{R}
    A_eigvals::Vector{C}
end
"""
        SSAAlloc(n::Integer) -> SSAAlloc

Memory allocation to avoid re-allocating space for repeated SSA computations.
# Arguments
- `n::Integer`: the size of the matrix (or matrices) on which the SSA will be computed
"""
function SSAAlloc(n::Integer)
  return SSAAlloc(
   map( _ -> Matrix{Float64}(undef,n,n),1:12)... , Vector{ComplexF32}(undef,n))
end

"""
        SSAAlloc(A::Matrix) -> SSAAlloc

Memory allocation to avoid re-allocating space for repeated SSA computations.
This corresponds to `SSAAlloc(size(A,1))`
"""
function SSAAlloc(A::Matrix{<:Real})
  return SSAAlloc(size(A,1))
end

"""
        PQ_init!(PQ::SSAAlloc,A::Matrix) -> nothing
Stores some (costly) pre-computed quantities on SSAAlloc.
Should be called only once before the root-finding procedure.
Schur decompositions are performed here.
"""
function PQ_init!(PQ::SSAAlloc{R,C},A::Matrix{R}) where {R,C}
    At=transpose!(PQ.At,A)
    F = schur(A)
    copyto!(PQ.R,F.T)
    copyto!(PQ.Z,F.Z)
    copyto!(PQ.A_eigvals,F.values)
    F = schur!(At)
    copyto!(PQ.Rt,F.T)
    copyto!(PQ.Zt,F.Z)
    mul!(PQ.D,transpose(PQ.Z),PQ.Z,-1.0,0.0)
    mul!(PQ.Dt,transpose(PQ.Zt),PQ.Zt,-1.0,0.0)
    return nothing
end

@inline function spectral_abscissa(PQ::SSAAlloc{R,C}) where {R,C}
    return convert(R,maximum(real.(PQ.A_eigvals)))
end
function get_P_or_Q!(PorQ::M,s::Real,Z::M,R::M,D::M,
        YZt_alloc::M) where M<:Matrix{<:Real}
    subtract_diagonal!(R,s)
    Y, scale = LAPACK.trsyl!('N','T', R, R, D)
    mul!(YZt_alloc,Y,transpose(Z))
    mul!(PorQ,Z,YZt_alloc,inv(scale),0.0)
    return nothing
end
function get_P!(s::T, PQ::SSAAlloc{T,C}) where {T,C}
    R=copyto!(PQ.R_alloc,PQ.R)
    D=copyto!(PQ.D_alloc,PQ.D)
    get_P_or_Q!(PQ.P,s,PQ.Z,R,D,PQ.YZt_alloc)
    return nothing
end
function get_Q!(s::T, PQ::SSAAlloc{T,C}) where {T,C}
    R=copyto!(PQ.R_alloc,PQ.Rt)
    D=copyto!(PQ.D_alloc,PQ.Dt)
    get_P_or_Q!(PQ.Q,s,PQ.Zt,R,D,PQ.YZt_alloc)
    return nothing
end

function ssa_simple_obj(s::R,PQ::SSAAlloc{R,C},ssa_eps::R,sa::R) where {R,C}
    get_P!( max(sa+eps(10.0),s) ,PQ) # SSA is not defined below SA !
    return inv(tr(PQ.P)) - ssa_eps
end

# Find zero with Newton!
function ssa_simple_obj_newton(s::R,PQ::SSAAlloc{R,C},ssa_eps::R,sa::R) where {R,C}
  _s = max(sa+eps(10.0),s)
  get_P!(_s,PQ) # SSA is not defined below SA !
  get_Q!(_s,PQ)
  fs = tr(PQ.P)
  mdfs = 2.0*trace_of_product(PQ.P,PQ.Q)
  obj = inv(fs) - ssa_eps
  dobj = mdfs/(fs*fs)
  return (obj,obj/dobj)
end

"""
    ssa!(A,grad,PQ,ssa_eps=nothing ;
        io_matrices=SmoothedSpectralAbscissa.IOIdentity,
        optim_method=SmoothedSpectralAbscissa.OptimOrder2) -> ssa_value

Computes the smoothed spectral abscissa (SSA) of matrix A and its gradient (optionally).
If `ssa_eps` is not specified, the default value is set by `default_eps_ssa(A)`

# Arguments
- `A::Matrix{<:Real}`: A square matrix
- `grad::Union{Nothing,Matrix{<:Real}}`: `nothing`  skips the gradient computation. To compute the gradient, provide a matrix `grad` of the same size as A. The matrix will be rewritten in-place with entries ``\\frac{\\partial\\; \\tilde{\\alpha}_\\varepsilon(A)}{\\partial\\; A_{i,j}}``
- `alloc::SSAAlloc` : Pre-allocated matrices for iterative computations. Can be generated by `SSAAlloc(A)`
- `ssa_eps::Union{Nothing,R}=nothing` The ``\\varepsilon`` parameter associated with the SSA computation. If not specified, it is set by `default_eps_ssa(A)`
- `io_matrices=SmoothedSpectralAbscissa.IOIdentity` : input and output matrix. **WARNING** for now the only option is identity.
- `optim_method=SmoothedSpectralAbscissa.OptimOrder2` : root finding algorithm. Two methods implemented: `OptimOrder2` and `OptimNewton`.
# Returns
- `ssa_value<:Real`: this is ``\\tilde{\\alpha}_\\varepsilon(A)``
"""
function ssa!(A::Matrix{R},grad::Union{Nothing,Matrix{R}}, alloc::SSAAlloc{R,C},
    ssa_eps::Union{Nothing,R}=nothing ; optim_method::Type=OptimOrder2,
    io_matrices::Type=IOIdentity) where {R,C}
  _ssa_eps = something(ssa_eps, default_eps_ssa(A))
  return ssa!(A,grad,alloc,_ssa_eps,optim_method,io_matrices)
end
# Order2() to find roots
function ssa!(A::Matrix{R},grad::Union{Nothing,Matrix{R}}, alloc::SSAAlloc{R,C},
     ssa_eps::R , optim_method::Type{OptimOrder2},
     io_matrices::Type{IOIdentity}) where {R,C}
  PQ_init!(alloc,A)
  _sa = spectral_abscissa(alloc)
  _start = _sa + 0.5*ssa_eps
  objfun(s)=ssa_simple_obj(s,alloc,ssa_eps,_sa)
  _method = Roots.Order2()
  s_star::R = find_zero(objfun, _start,_method;maxevals=1_000)
  isnothing(grad) && return s_star
  # compute gradient
  get_P!(s_star,alloc)
  get_Q!(s_star,alloc)
  cc=inv(trace_of_product(alloc.Q,alloc.P))
  mul!(grad,alloc.Q,alloc.P,cc,0.0)
  return s_star
end

# newton to find roots
function ssa!(A::Matrix{R},grad::Union{Nothing,Matrix{R}}, alloc::SSAAlloc{R,C},
     ssa_eps::R , optim_method::Type{OptimNewton},
     io_matrices::Type{IOIdentity}) where {R,C}
  PQ_init!(alloc,A)
  _sa = spectral_abscissa(alloc)
  _start = _sa + 0.5*ssa_eps # _start = _sa + 0.1abs(_sa)
  objfun(s)=ssa_simple_obj_newton(s,alloc,ssa_eps,_sa)
  _method = Roots.Newton()
  s_star::R = find_zero(objfun, _start,_method;maxevals=1_000)
  isnothing(grad) && return s_star
  # compute gradient
  get_P!(s_star,alloc)
  get_Q!(s_star,alloc)
  cc=inv(trace_of_product(alloc.Q,alloc.P))
  mul!(grad,alloc.Q,alloc.P,cc,0.0)
  return s_star
end

"""
    ssa_simple!(A,grad,PQ,ssa_eps=nothing) -> ssa

deprecated ... kept for testing
"""
function ssa_simple!(A::Matrix{R},grad::Union{Nothing,Matrix{R}},
        PQ::SSAAlloc{R,C},ssa_eps::Union{Nothing,R}=nothing) where {R,C}
 return ssa!(A,grad,PQ,ssa_eps)
end

# short versions that also allocates the memory
"""
    ssa(A,ssa_eps=nothing) -> ssa_val

Computes the smoothed spectral abscissa (SSA) of matrix A (with identity input-output weighting matrices). If `ssa_eps` is not specified, the default value is set by `default_eps_ssa(A)`

# Arguments
- `A::Matrix{<:Real}`: A square matrix
- `ssa_eps::Union{Nothing,R}=nothing` The ``\\varepsilon`` parameter associated with the SSA computation. If not specified, it is set by `default_eps_ssa(A)`
- `io_matrices=SmoothedSpectralAbscissa.IOIdentity` : input and output matrix. **WARNING** for now the only option is identity.
# Returns
- `ssa_val<:Real`: this is the ``\\tilde{\\alpha}_\\varepsilon(A)``
"""
function ssa(A::Matrix{R},ssa_eps::Union{Nothing,R}=nothing,
    io_matrices::Type=IOIdentity) where R
  alloc=SSAAlloc(size(A,1))
  _epsssa = something(ssa_eps, default_eps_ssa(A))
  return ssa!(copy(A),nothing,alloc, _epsssa;io_matrices=io_matrices)
end

"""
      ssa_withgradient(A,ssa_eps) -> (ssa_val,gradient_matrix)
Computes the Smoothed Spectral Abscissa (SSA) of matrix A
and also its gradient with respect to each element of A.
If `ssa_eps` is not specified, the default value is set by `default_eps_ssa(A)`

In case you need to call the SSA iteratively (e.g. gradient-based optimization),
please consider pre-allocating memory and using `ssa_simple!(...)`.

# Arguments
- `A::Matrix{<:Real}`: A square matrix
- `ssa_eps::Union{Nothing,R}=nothing` The ``\\varepsilon`` parameter associated with the SSA computation. If not specified, it is set by `default_eps_ssa(A)`
- `io_matrices=SmoothedSpectralAbscissa.IOIdentity` : input and output matrix. **WARNING** for now the only option is identity.
# Returns
- `ssa_val<:Real`: this is the ``\\tilde{\\alpha}_\\varepsilon(A)``
- `gradient_matrix::Matrix{<:Real}`: the gradient of the SSA w.r.t. each entry of matrix `A`
"""
function ssa_withgradient(A::Matrix{R},ssa_eps::Union{Nothing,R}=nothing,
        io_matrices::Type=IOIdentity) where R
    alloc=SSAAlloc(size(A,1))
    gradmat=similar(A)
    epsssa = something(ssa_eps, default_eps_ssa(A))
    _ssa = ssa!(copy(A),gradmat,alloc, epsssa;io_matrices=io_matrices)
    return _ssa,gradmat
end

end # module


text:
function mdl_avr_ode!(
    device_states,
    output_ode,
    device::PSY.DynamicGenerator{M, S, PSY.AVRFixed, TG, P},
) where {M <: PSY.Machine, S <: PSY.Shaft, TG <: PSY.TurbineGov, P <: PSY.PSS}

    #TODO Change EMF name for Vf in PowerSystems
    #Update Vf voltage on inner vars
    get_inner_vars(device)[Vf_var] = PSY.get_Vf(PSY.get_avr(device))

    return
end

function mdl_avr_ode!(
    device_states,
    output_ode,
    device::PSY.DynamicGenerator{M, S, PSY.AVRSimple, TG, P},
) where {M <: PSY.Machine, S <: PSY.Shaft, TG <: PSY.TurbineGov, P <: PSY.PSS}

    #Obtain references
    V_ref = PSY.get_ext(device)[CONTROL_REFS][V_ref_index]

    #Obtain indices for component w/r to device
    local_ix = get_local_state_ix(device, PSY.AVRSimple)

    #Define inner states for component
    internal_states = @view device_states[local_ix]
    Vf = internal_states[1]

    #Define external states for device
    V_th = sqrt(get_inner_vars(device)[VR_gen_var]^2 + get_inner_vars(device)[VI_gen_var]^2)

    #Get Parameters
    Kv = PSY.get_Kv(PSY.get_avr(device))

    #Compute ODEs
    output_ode[local_ix[1]] = Kv * (V_ref - V_th)

    #Update inner_vars
    get_inner_vars(device)[Vf_var] = Vf

    return
end

"""
Model of AVR Type I in Julia.
Refer to Power System Modelling and Scripting by F. Milano for the equations
"""

function mdl_avr_ode!(
    device_states,
    output_ode,
    device::PSY.DynamicGenerator{M, S, PSY.AVRTypeI, TG, P},
) where {M <: PSY.Machine, S <: PSY.Shaft, TG <: PSY.TurbineGov, P <: PSY.PSS}

    #Obtain references
    V0_ref = PSY.get_ext(device)[CONTROL_REFS][V_ref_index]

    #Obtain indices for component w/r to device
    local_ix = get_local_state_ix(device, PSY.AVRTypeI)

    #Define inner states for component
    internal_states = @view device_states[local_ix]
    Vf = internal_states[1]
    Vr1 = internal_states[2]
    Vr2 = internal_states[3]
    Vm = internal_states[4]

    #Define external states for device
    V_th = sqrt(get_inner_vars(device)[VR_gen_var]^2 + get_inner_vars(device)[VI_gen_var]^2)
    Vs = get_inner_vars(device)[V_pss_var]

    #Get parameters
    avr = PSY.get_avr(device)
    Ka = PSY.get_Ka(avr)
    Ke = PSY.get_Ke(avr)
    Kf = PSY.get_Kf(avr)
    Ta = PSY.get_Ta(avr)
    Te = PSY.get_Te(avr)
    Tf = PSY.get_Tf(avr)
    Tr = PSY.get_Tr(avr)
    Vr_max = PSY.get_Vr_max(avr)
    Vr_min = PSY.get_Vr_min(avr)
    Ae = PSY.get_Ae(avr)
    Be = PSY.get_Be(avr)

    #Compute auxiliary parameters
    Se_Vf = Ae * exp(Be * abs(Vf)) #16.13
    V_ref = V0_ref + Vs

    #Set anti-windup for Vr1. #TODO in callbacks
    #if Vr1 > Vr_max
    #    Vr1 = Vr_max
    #elseif Vr1 < Vr_min
    #    Vr1 = Vr_min
    #end

    #Compute 4 States AVR ODE:
    output_ode[local_ix[1]] = -(1.0 / Te) * (Vf * (Ke + Se_Vf) - Vr1) #16.12c
    output_ode[local_ix[2]] = (1.0 / Ta) * (Ka * (V_ref - Vm - Vr2 - (Kf / Tf) * Vf) - Vr1) #16.12a
    output_ode[local_ix[3]] = -(1.0 / Tf) * ((Kf / Tf) * Vf + Vr2) #16.12b
    output_ode[local_ix[4]] = (1.0 / Tr) * (V_th - Vm) #16.11

    #Update inner_vars
    get_inner_vars(device)[Vf_var] = Vf

    return
end

"""
Model of AVR Type II in Julia.
Refer to Power System Modelling and Scripting by F. Milano for the equations
"""

function mdl_avr_ode!(
    device_states,
    output_ode,
    device::PSY.DynamicGenerator{M, S, PSY.AVRTypeII, TG, P},
) where {M <: PSY.Machine, S <: PSY.Shaft, TG <: PSY.TurbineGov, P <: PSY.PSS}

    #Obtain references
    V0_ref = PSY.get_ext(device)[CONTROL_REFS][V_ref_index]

    #Obtain indices for component w/r to device
    local_ix = get_local_state_ix(device, PSY.AVRTypeII)

    #Define inner states for component
    internal_states = @view device_states[local_ix]
    Vf = internal_states[1]
    Vr1 = internal_states[2]
    Vr2 = internal_states[3]
    Vm = internal_states[4]

    #Define external states for device
    V_th = sqrt(get_inner_vars(device)[VR_gen_var]^2 + get_inner_vars(device)[VI_gen_var]^2)
    Vs = get_inner_vars(device)[V_pss_var]

    #Get parameters
    avr = PSY.get_avr(device)
    K0 = PSY.get_K0(avr)
    T1 = PSY.get_T1(avr)
    T2 = PSY.get_T2(avr)
    T3 = PSY.get_T3(avr)
    T4 = PSY.get_T4(avr)
    Te = PSY.get_Te(avr)
    Tr = PSY.get_Tr(avr)
    Vr_max = PSY.get_Vr_max(avr)
    Vr_min = PSY.get_Vr_min(avr)
    Ae = PSY.get_Ae(avr)
    Be = PSY.get_Be(avr)

    #Compute auxiliary parameters
    Se_Vf = Ae * exp(Be * abs(Vf)) #16.13
    V_ref = V0_ref + Vs
    Vr = K0 * Vr2 + (T4 / T3) * (Vr1 + K0 * (T2 / T1) * (V_ref - Vm)) #16.21

    #Set anti-windup for Vr1. #TODO in callbacks
    #if Vr > Vr_max
    #    Vr = Vr_max
    #elseif Vr < Vr_min
    #    Vr = Vr_min
    #end

    #Compute 4 States AVR ODE:
    output_ode[local_ix[1]] = -(1.0 / Te) * (Vf * (1.0 + Se_Vf) - Vr) #16.18
    output_ode[local_ix[2]] = (1.0 / T1) * (K0 * (1.0 - (T2 / T1)) * (V_ref - Vm) - Vr1) #16.14
    output_ode[local_ix[3]] =
        (1.0 / (K0 * T3)) *
        ((1.0 - (T4 / T3)) * (Vr1 + K0 * (T2 / T1) * (V_ref - Vm)) - K0 * Vr2)  #16.20
    output_ode[local_ix[4]] = (1.0 / Tr) * (V_th - Vm) #16.11

    #Update inner_vars
    get_inner_vars(device)[Vf_var] = Vf

    return
end


text:
using DiffEqSensitivity, OrdinaryDiffEq, DiffEqCallbacks, Flux
using Random, Test
using Zygote

function test_hybridNODE(sensealg)
    Random.seed!(12345)
    datalength = 100
    tspan = (0.0,100.0)
    t = range(tspan[1],tspan[2],length=datalength)
    target = 3.0*(1:datalength)./datalength  # some dummy data to fit to
    cbinput = rand(1, datalength) #some external ODE contribution
    pmodel = Chain(
        Dense(2, 10, init=zeros),
        Dense(10, 2, init=zeros))
    p, re = Flux.destructure(pmodel)
    dudt(u,p,t) = re(p)(u)

    # callback changes the first component of the solution every time
    # t is an integer
    function affect!(integrator, cbinput)
        event_index = round(Int,integrator.t)
        integrator.u[1] += 0.2*cbinput[event_index]
    end
    callback = PresetTimeCallback(collect(1:datalength),(int)->affect!(int, cbinput))

    # ODE with Callback
    prob = ODEProblem(dudt,[0.0, 1.0],tspan,p)

    function predict_n_ode(p)
        arr = Array(solve(prob, Tsit5(),
            p=p, sensealg=sensealg, saveat=2.0, callback=callback))[1,2:2:end]
        return arr[1:datalength]
    end

    function loss_n_ode()
        pred = predict_n_ode(p)
        loss = sum(abs2,target .- pred)./datalength
    end

    cb = function () #callback function to observe training
        pred = predict_n_ode(p)
        display(loss_n_ode())
    end
    @show sensealg
    Flux.train!(loss_n_ode, Flux.params(p), Iterators.repeated((), 20), ADAM(0.005), cb = cb)
    @test loss_n_ode() < 0.5
    println("  ")
end

function test_hybridNODE2(sensealg)
    Random.seed!(12345)
    u0 = Float32[2.; 0.; 0.; 0.]
    tspan = (0f0,1f0)

    ## Get goal data
    function trueaffect!(integrator)
        integrator.u[3:4] = -3*integrator.u[1:2]
    end
    function trueODEfunc(dx,x,p,t)
        @views dx[1:2] .= x[1:2] + x[3:4]
        dx[1] += x[2]
        dx[2] += x[1]
        dx[3:4] .= 0f0
    end
    cb_ = PeriodicCallback(trueaffect!,0.1f0,save_positions=(true,true),initial_affect=true)
    prob = ODEProblem(trueODEfunc,u0,tspan)
    sol = solve(prob,Tsit5(),callback=cb_,save_everystep=false,save_start=true)
    ode_data = Array(sol)[1:2,1:end]'

    ## Make model
    dudt2 = Chain(Dense(4,50,tanh),
                    Dense(50,2))
    p,re = Flux.destructure(dudt2) # use this p as the initial condition!
    function affect!(integrator)
        integrator.u[3:4] = -3*integrator.u[1:2]
    end
    function ODEfunc(dx,x,p,t)
        dx[1:2] .= re(p)(x)
        dx[3:4] .= 0f0
    end
    z0 = u0
    prob = ODEProblem(ODEfunc,z0,tspan)
    cb = PeriodicCallback(affect!,0.1f0,save_positions=(true,true),initial_affect=true)

    ## Initialize learning functions
    function predict_n_ode()
        _prob = remake(prob,p=p)
        Array(solve(_prob,Tsit5(),u0=z0,p=p,callback=cb,save_everystep=false,save_start=true,sensealg=sensealg))[1:2,:]
    end
    function loss_n_ode()
        pred = predict_n_ode()[1:2,1:end]'
        loss = sum(abs2,ode_data .- pred)
        loss
    end
    loss_n_ode() # n_ode.p stores the initial parameters of the neural ODE
    cba = function ()  #callback function to observe training
        pred = predict_n_ode()[1:2,1:end]'
        display(sum(abs2,ode_data .- pred))
        return false
    end
    cba()

    ## Learn
    ps = Flux.params(p)
    data = Iterators.repeated((), 25)

    @show sensealg

    Flux.train!(loss_n_ode, ps, data, ADAM(0.0025), cb = cba)

    @test loss_n_ode() < 0.5

    println("  ")
end

mutable struct Affect{T}
      callback_data::T
end
compute_index(t) = round(Int,t)+1
function (cb::Affect)(integrator)
    indx = compute_index(integrator.t)
    integrator.u .= integrator.u .+ @view(cb.callback_data[:, indx, 1]) * (integrator.t - integrator.tprev)
end
function test_hybridNODE3(sensealg)
    u0 = Float32[2.; 0.]
    datasize = 101
    tspan = (0.0f0,10.0f0)

    function trueODEfunc(du,u,p,t)
        du .= -u
    end
    t = range(tspan[1],tspan[2],length=datasize)
    prob = ODEProblem(trueODEfunc,u0,tspan)
    ode_data = Array(solve(prob,Tsit5(),saveat=t))

    true_data = reshape(ode_data,(2,length(t),1))
    true_data = convert.(Float32,true_data)
    callback_data = true_data * 1f-3
    train_dataloader = Flux.Data.DataLoader((true_data = true_data,callback_data = callback_data),batchsize=1)
    dudt2 = Chain(Dense(2,50,tanh),
                 Dense(50,2))
    p,re = Flux.destructure(dudt2)
    function dudt(du,u,p,t)
        du .= re(p)(u)
    end
    z0 = Float32[2.; 0.]
    prob = ODEProblem(dudt,z0,tspan)

    function callback_(callback_data)
        affect! = Affect(callback_data)
        condition(u,t,integrator) = integrator.t > 0
        DiscreteCallback(condition,affect!,save_positions=(false,false))
    end

    function predict_n_ode(true_data_0,callback_data, sense)
        _prob = remake(prob,p=p,u0=true_data_0)
        solve(_prob,Tsit5(),callback=callback_(callback_data),saveat=t,sensealg=sense)
    end

    function loss_n_ode(true_data,callback_data)
        sol = predict_n_ode((vec(true_data[:,1,:])),callback_data,sensealg)
        pred = Array(sol)
        loss = Flux.mse((true_data[:,:,1]),pred)
        loss
    end

    ps = Flux.params(p)
    opt = ADAM(0.1)
    epochs = 10
    function cb1(true_data,callback_data)
        display(loss_n_ode(true_data,callback_data))
        return false
    end

    function train!(loss, ps, data, opt, cb)
        ps = Params(ps)
        for (true_data,callback_data) in data
            gs = gradient(ps) do
                loss(true_data,callback_data)
            end
            Flux.update!(opt, ps, gs)
            cb(true_data,callback_data)
        end
        return nothing
    end

    @Flux.epochs epochs train!(loss_n_ode, Params(ps),train_dataloader, opt, cb1)
    loss =  loss_n_ode(true_data[:,:,1],callback_data)
    @info loss
    @test loss < 0.5
end

@testset "PresetTimeCallback" begin
    test_hybridNODE(ForwardDiffSensitivity())
    test_hybridNODE(BacksolveAdjoint())
    test_hybridNODE(InterpolatingAdjoint())
    test_hybridNODE(QuadratureAdjoint())
end

@testset "PeriodicCallback" begin
    test_hybridNODE2(ReverseDiffAdjoint())
    test_hybridNODE2(BacksolveAdjoint())
    test_hybridNODE2(InterpolatingAdjoint())
    test_hybridNODE2(QuadratureAdjoint())
end

@testset "tprevCallback" begin
    test_hybridNODE3(ReverseDiffAdjoint())
    test_hybridNODE3(BacksolveAdjoint())
    test_hybridNODE3(InterpolatingAdjoint())
    test_hybridNODE3(QuadratureAdjoint())
end


text:
struct BindingCurve{T,B,F}
    bond_lengths::B
    potential::Vector{T}
    equilibrium_bond_length::T
    potential_minimum::T
    fit::F
end

function calculate_binding_curve(
    bond_lengths::AbstractVector, model::AdiabaticModel, environment::EvaluationEnvironment
)
    potential = calculate_diatomic_energy.(bond_lengths, model, environment) # Calculate binding curve
    potential_minimum, index = findmin(potential)
    equilibrium_bond_length = bond_lengths[index]
    fit = fit_binding_curve(bond_lengths, potential)
    return BindingCurve(bond_lengths, potential, equilibrium_bond_length, potential_minimum, fit)
end

function fit_binding_curve(bond_lengths, binding_curve)
    itp = interpolate(binding_curve, BSpline(Cubic(Line(OnGrid()))))
    sitp = scale(itp, bond_lengths)
    return sitp
end

function calculate_force_constant(binding_curve)
    return hessian(binding_curve.fit, binding_curve.equilibrium_bond_length)[1]
end

function guess_initial_energy(k, μ, ν, potential_minimum)
    ω = sqrt(k / μ)
    E_guess = (ν + 1/2) * ω
    return E_guess + potential_minimum
end

function plot_binding_curve(bond_lengths, binding_curve, fit)
    plt = lineplot(bond_lengths, binding_curve;
        title="Binding curve", xlabel="Bond length / bohr", ylabel="Energy / Hartree",
        name="Actual values", canvas=DotCanvas, border=:ascii
    )
    lineplot!(plt, bond_lengths, fit.(bond_lengths); name="Fitted curve")
    show(plt)
    println()
    @info "The two lines shown above should closely match. \
        This indicates the evaluation of the potential is working correctly."
end


text:
module SimpleCollisions

import LinearAlgebra as LA
import StaticArrays as SA

const Vector2D{T} = SA.SVector{2, T}

include("utils.jl")
include("position_and_orientation.jl")
include("shapes.jl")
include("collision_detection.jl")
include("collision_manifold.jl")
include("physics.jl")

end


text:
for x in ["OnlineStats",
"StatsBase",
"Distributions",
"DataFrames",
"ProgressMeter",
"Iterators",
"ArrayViews",
"HDF5",
"JLD",
"SQLite",
"ProfileView",
"Coverage",
"PyCall"]
    @show x
    Pkg.add(x)
end
Pkg.update();
using WinRPM
WinRPM.update();


text:
module IUP

export
	Ihandle,
	IupAlarm,
	IupBackgroundBox,
	IupButton,
	IupCanvas, IupCbox, IupClose, IupControlsOpen,
	IupDial, IupDestroy, IupDialog,
	IupFill, IupFrame,
	IupGetAttribute, IupGetClassName, IupGetDialog, IupGetFloat, IupGetFile, IupGetGlobal, IupGetHandle, IupGetInt,
	IupGetDouble, IupGetDoubleId, IupGetInt2, IupGridBox, IupGetParam,
	IupHbox,
	IupHide,
	IupImage, IupImageRGBA, IupItem,
	IupLabel, IupList,
	IupLoad,
	IupMainLoop, IupMap, IupMatrix,
	IupMenu,
	IupMessage,
	IupMultiLine,
	IupOpen,
	IupPopup, IupProgressBar,
	IupRadio,
	IupRefresh,
	IupSeparator,
	IupSetAttribute, IupSetAttributes, IupSetAttributeHandle, IupSetCallback, IupSetCallbacks, IupSetFunction,
	IupSetFloat, IupSetDouble, IupSetDoubleId, IupSetGlobal, IupSetHandle, IupSetLanguage,
	IupShow,
	IupShowXY,
	IupStoreAttribute,
	IupSubmenu,
	IupTabs, IupTabsv, IupText, IupToggle, IupTree,
	IupVal, IupVbox,
	#
	# From iup_plot
	#
	IupPlotOpen,
	IupPlot,
	IupPlotBegin,
	IupPlotAdd,
	IupPlotAddStr,
	IupPlotEnd,
	IupPlotLoadData,
	IupPlotInsertStr,
	IupPlotInsert,
	IupPlotInsertStrPoints,
	IupPlotInsertPoints,
	IupPlotAddPoints,
	IupPlotAddStrPoints,
	IupPlotGetSample,
	IupPlotGetSampleStr,
	IupPlotGetSampleSelection,
	IupPlotSetSample,
	IupPlotSetSampleStr,
	IupPlotSetSampleSelection,
	IupPlotTransform,
	IupPlotTransformTo,
	IupPlotPaintTo,
	#
	# Need to add many more
	#
	OBJC_NEW_PROPERTIES, IUP_CLOSE,
	IUP_ERROR, IUP_NOERROR, IUP_OPENED, IUP_INVALID, IUP_INVALID_ID, IUP_IGNORE,
	IUP_DEFAULT, IUP_CENTER, IUP_LEFT, IUP_TOP, IUP_BOTTOM, IUP_RIGHT, IUP_MOUSEPOS, IUP_CURRENT, IUP_CENTERPARENT,
	IUP_BUTTON1, IUP_BUTTON2, IUP_BUTTON3, IUP_BUTTON4, IUP_BUTTON5

include("libiup_h.jl")
include("libiup.jl")
include("iupcontrols.jl")
include("iup_plot.jl")

end  # module


text:
# DataLoaders - Concrete subtypes of AbstractDataLoader
#
# Lucas Ondel 2020

function Base.show(io::IO, dl::AbstractDataLoader)
    println(io, "$(typeof(dl))")
    println(io, "  data: $(typeof(dl.data))")
    print(io, "  batchsize: $(dl.batchsize)")
end

_index(i, batchsize) = (i-1) * batchsize + 1

#######################################################################
# DataLoader

"""
    struct DataLoader
        data
        batchsize
    end

# Constructor

    DataLoader(data[, batchsize = 1, preprocess = x -> x,
               preprocess_element = x -> x])

where `data` is a sequence of elements to iterate over, `batchsize` is
the size of each batch, `preprocess` is a user-defined function to
apply on each batch and `preprocess_element` is a user-defined function
to apply on each batch's element. By default, `preprocess` and
`preprocess_element` are simply the identity function.

!!! warning

    When iterating, the final batch may have a size smaller
    than `batchsize`.

"""
struct DataLoader{T<:AbstractVector}  <: AbstractDataLoader{T}
    data::T
    batchsize::UInt
    fbatch::Function
    felement::Function

    function DataLoader(data::AbstractVector; batchsize = 1, preprocess = x -> x,
                        preprocess_element = x -> x)
        length(data) > 0 || throw(ArgumentError("cannot create a DataLoader from an empty collection"))
        batchsize >= 1 || throw(ArgumentError("`batchsize = $batchsize` should greater or equal to 1"))
        new{typeof(data)}(data, batchsize, preprocess, preprocess_element)
    end
end

function Base.iterate(dl::DataLoader, state = 1)
    if state > size(dl.data, 1)
        return nothing
    end
    offset = min(state+dl.batchsize-1, size(dl.data,1))
    dl.fbatch(dl.felement.(dl.data[state:offset])), offset+1
end
Base.length(dl::DataLoader) = UInt(ceil(size(dl.data, 1)/dl.batchsize))
Base.eltype(dl::DataLoader) = eltype(dl.data)

function Base.getindex(dl::DataLoader, i)
    1 <= i <= length(dl) || throw(BoundsError(dl, i))
    start = _index(i, dl.batchsize)
    offset = min(start + dl.batchsize - 1, size(dl.data,1))
    dl.fbatch(dl.felement.(dl.data[start:offset]))
end

function Base.getindex(dl::DataLoader, ur::UnitRange)
    1 <= ur.start <= length(dl) || throw(BoundsError(dl, ur.start))
    1 <= ur.stop <= length(dl) || throw(BoundsError(dl, ur.stop))

    N = size(dl.data, 1)
    start = _index(ur.start, dl.batchsize)
    offset = min(_index(ur.stop, dl.batchsize) + dl.batchsize - 1, N)
    DataLoader(dl.data[start:offset], batchsize = dl.batchsize,
               preprocess = dl.fbatch, preprocess_element = dl.felement)
end

Base.firstindex(dl::DataLoader) = 1
Base.lastindex(dl::DataLoader) = length(dl)

max_stars_repo_name: lucasondel/DataLoaders

max_issues_repo_name: lucasondel/DataLoaders

max_forks_repo_name: lucasondel/DataLoaders


text:
module MultiAgentSysAdmin

# Write your package code here.
using StaticArrays
using LightGraphs
using POMDPs
using Random

using MultiAgentPOMDPs
using POMDPModelTools: Deterministic

const MAPOMDPs = MultiAgentPOMDPs

"""
- `status`: {good, faulty, dead}
- `load`: {idle, loaded, success}
"""
const MachineState = SVector{2, Int}
const MachineAction = @SArray [0, 1] # noop, reboot

"""
- `p_fail_base`:
- `p_fail_bonus`:
- `p_dead_base`:
- `p_dead_bonus`:
- `p_load`:  probability of getting a job when idle.
- `p_doneG`: probability of completing a job when good.
- `p_doneF`: probability of completing a job when faulty.
`p_fail_bonus` and `p_dead_bonus` are additional bonuses counted
when all neighbors are faulty. Counted per agent.
If a machine with 2 neighbors has a single faulty neighbor, it will get
an additional failing probability of `p_fail_bonus/2`. If the same machine
has one faulty neighbor and one dead neighbor, it will get a penalty of
`p_fail_bonus/2 + p_dead_bonus/2`.
"""
abstract type AbstractSysAdmin{Bool} <: JointMDP{Vector{MachineState}, Vector{Int}} end

Base.@kwdef struct UniSysAdmin{T} <: AbstractSysAdmin{T}
    nagents::Int = 4
    # status
    p_fail_base::Float64 = 0.4
    p_fail_bonus::Float64 = 0.2
    p_dead_base::Float64 = 0.1
    p_dead_bonus::Float64 = 0.5
    # load
    p_load::Float64 = 0.6
    p_doneG::Float64 = 0.9
    p_doneF::Float64 = 0.6

    discount::Float64 = 0.9
    reboot_penalty = -0.7
end

Base.@kwdef struct BiSysAdmin{T} <: AbstractSysAdmin{T}
    nagents::Int = 4
    # status
    p_fail_base::Float64 = 0.4
    p_fail_bonus::Float64 = 0.2
    p_dead_base::Float64 = 0.1
    p_dead_bonus::Float64 = 0.5
    # load
    p_load::Float64 = 0.6
    p_doneG::Float64 = 0.9
    p_doneF::Float64 = 0.6

    discount::Float64 = 0.9
    reboot_penalty = -0.0
end

Base.@kwdef struct RingofRingSysAdmin{T} <: AbstractSysAdmin{T}
    nrings::Int = 3
    nagents_per_ring::Int = 3
    # status
    p_fail_base::Float64 = 0.4
    p_fail_bonus::Float64 = 0.2
    p_dead_base::Float64 = 0.1
    p_dead_bonus::Float64 = 0.5
    # load
    p_load::Float64 = 0.6
    p_doneG::Float64 = 0.9
    p_doneF::Float64 = 0.6

    discount::Float64 = 0.9
    reboot_penalty = -0.0
end

Base.@kwdef struct StarSysAdmin{T} <: AbstractSysAdmin{T}
    nagents::Int = 4
    # status
    p_fail_base::Float64 = 0.4
    p_fail_bonus::Float64 = 0.2
    p_dead_base::Float64 = 0.1
    p_dead_bonus::Float64 = 0.5
    # load
    p_load::Float64 = 0.6
    p_doneG::Float64 = 0.9
    p_doneF::Float64 = 0.6

    discount::Float64 = 0.9
    reboot_penalty = -0.0
end

Base.@kwdef struct RandomSysAdmin{T} <: AbstractSysAdmin{T}
    nagents::Int = 4
    nedges::Int = 5
    seed::Int = 1
    # status
    p_fail_base::Float64 = 0.4
    p_fail_bonus::Float64 = 0.2
    p_dead_base::Float64 = 0.1
    p_dead_bonus::Float64 = 0.5
    # load
    p_load::Float64 = 0.6
    p_doneG::Float64 = 0.9
    p_doneF::Float64 = 0.6

    discount::Float64 = 0.9
    reboot_penalty = -0.0
end

POMDPs.discount(p::AbstractSysAdmin) = p.discount
#POMDPs.isterminal(p::AbstractSysAdmin, s) = all(x->x[2] == 3) # XXX
POMDPs.isterminal(p::AbstractSysAdmin, s) = false

POMDPs.actionindex(p::AbstractSysAdmin, a, c) = findfirst(isequal(a), MachineAction)

MAPOMDPs.n_agents(p::AbstractSysAdmin) = p.nagents
MAPOMDPs.n_agents(p::RingofRingSysAdmin) = p.nagents_per_ring * p.nrings

MAPOMDPs.agent_actions(p::AbstractSysAdmin, idx::Int64, s::MachineState) = MachineAction
MAPOMDPs.agent_actions(p::AbstractSysAdmin, idx::Int64) = MachineAction
MAPOMDPs.agent_actionindex(p::AbstractSysAdmin, idx::Int64, a) = findfirst(isequal(a), MAPOMDPs.agent_actions(p, idx))
POMDPs.actions(p::AbstractSysAdmin) = vec(map(collect, Iterators.product((MAPOMDPs.agent_actions(p, i) for i in 1:MAPOMDPs.n_agents(p))...)))
POMDPs.actionindex(p::AbstractSysAdmin, a) = findfirst(isequal(a), POMDPs.actions(p))

function coord_graph_adj_mat(p::UniSysAdmin)
    mat = zeros(Int64, p.nagents, p.nagents)
    for i in 1:p.nagents-1
        mat[i, i+1] = 1
    end
    mat[p.nagents, 1] = 1
    return mat
end

function coord_graph_adj_mat(p::BiSysAdmin)
    mat = zeros(Int64, p.nagents, p.nagents)
    for i in 1:p.nagents-1
        mat[i, i+1] = 1
    end
    mat[p.nagents, 1] = 1
    mat = mat + mat'
    return mat
end

function coord_graph_adj_mat(p::RingofRingSysAdmin)
    na = MAPOMDPs.n_agents(p)
    mat = zeros(Int64, na, na)
    # Inner ring
    for idx in Iterators.product(1:p.nrings:na, 1:p.nrings:na)
        if idx[1] == idx[2]
            continue
        end
        mat[idx...] = 1
    end
    # rings on inner ring
    for i in 1:p.nrings:na-1
        for j in i:(i+p.nagents_per_ring-2)
            mat[j, j+1] = 1
            mat[j+1, j] = 1
        end
        mat[i, i+p.nagents_per_ring-1] = 1
        mat[i+p.nagents_per_ring-1, i] = 1
    end
    return mat
end

function coord_graph_adj_mat(p::RandomSysAdmin)
    return Matrix(adjacency_matrix(coordination_graph(p)))
end

function MAPOMDPs.coordination_graph(p::UniSysAdmin)
    DiGraph(coord_graph_adj_mat(p))
end

function MAPOMDPs.coordination_graph(p::BiSysAdmin)
    SimpleGraph(coord_graph_adj_mat(p))
end

function MAPOMDPs.coordination_graph(p::RandomSysAdmin)
    SimpleGraph{Int}(p.nagents, p.nedges, seed=p.seed)
end

function MAPOMDPs.coordination_graph(p::RingofRingSysAdmin)
    SimpleGraph(coord_graph_adj_mat(p))
end

function coord_graph_adj_mat(p::StarSysAdmin)
    mat = zeros(Int64, p.nagents, p.nagents)
    for i in 2:p.nagents
        mat[1, i] = 1
        mat[i, 1] = 1
    end
    return mat
end

function MAPOMDPs.coordination_graph(p::StarSysAdmin)
    SimpleGraph(coord_graph_adj_mat(p))
end

# status: good, fail, dead
# load: idle, work, done
MAPOMDPs.agent_states(p::AbstractSysAdmin, idx::Int64) = vec(MachineState[MachineState(status,load) for status in 1:3, load in 1:3])
MAPOMDPs.agent_stateindex(p::AbstractSysAdmin, idx::Int64, s) = findfirst(isequal(s), MAPODMPs.agent_states(p, idx))
POMDPs.states(p::AbstractSysAdmin) = vec(map(collect, Iterators.product((MAPOMDPs.agent_states(p, i) for i in 1:MAPOMDPs.n_agents(p))...)))
POMDPs.stateindex(p::AbstractSysAdmin, s) = findfirst(isequal(s), POMDPs.states(p))

function POMDPs.initialstate(p::AbstractSysAdmin)
    return Deterministic(MachineState[MachineState(1, 1) for _ in 1:MAPOMDPs.n_agents(p)])
end

function sysadmin_loop(p, s, a, rng)
    coordgraph = MAPOMDPs.coordination_graph(p) #SimpleGraph(coord_graph_adj_mat(p))
    sp_vec = Vector{MachineState}(undef, MAPOMDPs.n_agents(p))
    r_vec = Vector{Float64}(undef, MAPOMDPs.n_agents(p))
    for aidx in 1:MAPOMDPs.n_agents(p)
        rew = 0.0
        bonus = 0.0
        neighs = neighbors(coordgraph, aidx)
        for neigh in neighs
            status = s[neigh][1]
            if status == 2  # neighbor Fail
                bonus += p.p_fail_bonus
            elseif status == 3 # neighbor dead
                bonus += p.p_dead_bonus
            end
        end
        bonus /= length(neighs)
        p_fail = p.p_fail_base + bonus
        p_dead = p.p_dead_base + bonus

        # Rewards only if noop
        if a[aidx] == 0         # noop
            status = s[aidx][1]
            if status == 1      # Good
                if rand(rng) < p_fail
                    newstatus = 2
                else
                    newstatus = 1
                end
            elseif status == 2
                if rand(rng) < p_dead
                    newstatus = 3
                else
                    newstatus = 2
                end
            elseif status == 3
                newstatus = 3
            end

            load = s[aidx][1]
            if load == 1        # idle
                if newstatus == 1
                    if rand(rng) < p.p_load
                        newload = 2
                    else
                        newload = 1
                    end
                elseif newstatus == 2
                    if rand(rng) < p.p_load
                        newload = 2
                    else
                        newload = 1
                    end
                elseif newstatus == 3
                    newload = 1
                end
            elseif load == 2    # work
                if newstatus == 1
                    if rand(rng) < p.p_doneG
                        newload = 3
                        rew = 1.0 # finish reward
                    else
                        newload = 2
                    end
                elseif newstatus == 2
                    if rand(rng) < p.p_doneF
                        newload = 3
                        rew = 1.0 # finish reward
                    else
                        newload = 2
                    end
                elseif newstatus == 3 # dead
                    newload = 1
                end
            elseif load == 3    # done
                newload = 3
            end
        else                    # reboot
            newstatus = 1       # Good
            newload = 1
            rew += p.reboot_penalty
        end
        sp_vec[aidx] = MachineState(newstatus, newload)
        r_vec[aidx] = rew
    end
    return (sp=sp_vec, r=r_vec)
end

"""
Basically, the only way we can get reward is by:
- Starting from the Load state (since it's the only one that can complete)
- Doing action 0;
- And ending up in the Done state.
dead machine increases the probability that its neighbors become faulty and die
system receives a reward of 1 if a process terminates successfully
status is faulty, processes take longer to terminate
If the machine dies, the process is lost.
"""
function POMDPs.gen(p::AbstractSysAdmin{false}, s, a, rng)
    return sysadmin_loop(p, s, a, rng)
end

function POMDPs.gen(p::AbstractSysAdmin{true}, s, a, rng)
    sp_vec, r_vec = sysadmin_loop(p, s, a, rng)
    return (sp=sp_vec, r=sum(r_vec))
end

export UniSysAdmin, BiSysAdmin, RingofRingSysAdmin, RandomSysAdmin

end

max_stars_repo_name: JuliaPOMDP/MultiAgentSysAdmin

max_issues_repo_name: JuliaPOMDP/MultiAgentSysAdmin

max_forks_repo_name: JuliaPOMDP/MultiAgentSysAdmin


text:
using Pkg
Pkg.activate(pwd())

# # Neural networks
# ## Prepare data

using RDatasets

iris = dataset("datasets", "iris")
X = Matrix(iris[:, 1:4])
y = iris.Species

#### Exercise:
#Write the `split` function, which randomly splits the dataset and the labels into training
#and testing sets. Its input should be the dataset `X` and the labels `y`. It should have
#four outputs. Include 80% of data in the training set and 20% of data in the testing set
# by default.
#
#**Hints:**
#- Use the `randperm` function from the `Random` package.
#- While `y` can be assumed to a vector, `X` is a matrix or a more-dimensional array. Then
#it is beneficial to use the `selectdim` function to select subindices along the correct
#dimension.
#
#---
#### Solution:

using Random

function split(X, y::AbstractVector; dims=1, ratio_train=0.8, kwargs...)
    n = length(y)
    size(X, dims) == n || throw(DimensionMismatch("..."))

    n_train = round(Int, ratio_train*n)
    i_rand = randperm(n)
    i_train = i_rand[1:n_train]
    i_test = i_rand[n_train+1:end]

    return selectdim(X, dims, i_train), y[i_train], selectdim(X, dims, i_test), y[i_test]
end

#+

X_train, y_train, X_test, y_test = split(X, y)

# ---
#
# ### Exercise:
# Write the `normalize` functions as described above. It should have two inputs and two
# outputs. The keyword argument `dims` should also be included.
#
# **Hint**: check the help for the `mean` function.
#
# ---
# ### Solution:

using Statistics

function normalize(X_train, X_test; dims=1, kwargs...)
    col_mean = mean(X_train; dims)
    col_std = std(X_train; dims)

    return (X_train .- col_mean) ./ col_std, (X_test .- col_mean) ./ col_std
end

#+

X_train, X_test = normalize(X_train, X_test)

# ---
#
# ### Exercise:
# Write the `onehot` function that converts the labels `y` into their one-hot
# representation. The samples should be along the second dimension. Write the `onecold`
# function that converts the one-hot representation into the one-cold (original)
# representation. Both these functions need to have two arguments; the second one is
# `classes`, which equals `unique(y)`.
#
# Write a check that both functions work correctly.
#
# **Hints:**
# - The one-hot representation for a label has the size equalling to the number of classes.
# All entries besides one are zeros.
# - Since the one-hot representation represents probabilities, the prediction is the class
# with the highest probability.
#
# ---
# ### Solution:

function onehot(y, classes)
    y_onehot = falses(length(classes), length(y))
    for (i, class) in enumerate(classes)
        y_onehot[i, y .== class] .= 1
    end
    return y_onehot
end

onecold(y, classes) = [classes[argmax(y_col)] for y_col in eachcol(y)]

#+

classes = unique(y)
isequal(onecold(onehot(y, classes), classes), y)

# ---

using LinearAlgebra

function prepare_data(X, y; do_normal=true, do_onehot=true, kwargs...)
    X_train, y_train, X_test, y_test = split(X, y; kwargs...)

    if do_normal
        X_train, X_test = normalize(X_train, X_test; kwargs...)
    end

    classes = unique(y)

    if do_onehot
        y_train = onehot(y_train, classes)
        y_test = onehot(y_test, classes)
    end

    return X_train, y_train, X_test, y_test, classes
end

#+

Random.seed!(666)

iris = dataset("datasets", "iris")
X = Matrix(iris[:, 1:4])
y = iris.Species

X_train, y_train, X_test, y_test, classes = prepare_data(X', y; dims=2)

#+

Random.seed!(666)
aux1 = prepare_data(X, y; dims=1)

#+

Random.seed!(666)
aux2 = prepare_data(X', y; dims=2)

#+

norm(aux1[1] - aux2[1]')

# ## Create the network

struct SimpleNet{T<:Real}
    W1::Matrix{T}
    b1::Vector{T}
    W2::Matrix{T}
    b2::Vector{T}
end

# ### Exercise:
# Write an outer constructor for `SimpleNet`. Its inputs should be three integers
# representing the input size of the three layers. All matrices should be initialized based
# on the normal distribution.
#
# **Hint**: think about the representation of the dense layer.
#
# ---
# ### Solution:

SimpleNet(n1, n2, n3) = SimpleNet(randn(n2, n1), randn(n2), randn(n3, n2), randn(n3))

# ---

Random.seed!(666)
m = SimpleNet(size(X_train,1), 5, size(y_train,1))

# ### Exercise:
# Write a functor `function (m::SimpleNet)(x)` which computes the prediction (forward pass)
# of the neural network `SimpleNet`.
#
# **Bonus**: try to make the functor work for both vectors (one sample) and matrices
# (multiple samples) `x`.
#
# ---
# ### Solution:

function (m::SimpleNet)(x)
    z1 = m.W1*x .+ m.b1
    a1 = max.(z1, 0)
    z2 = m.W2*a1 .+ m.b2
    return exp.(z2) ./ sum(exp.(z2), dims=1)
end

# ---

m(X_train[:,1:2])

# ## Train the network

function grad(m::SimpleNet, x::AbstractVector, y; ϵ=1e-10)
    z1 = m.W1*x .+ m.b1
    a1 = max.(z1, 0)
    z2 = m.W2*a1 .+ m.b2
    a2 = exp.(z2) ./ sum(exp.(z2), dims=1)
    l = -sum(y .* log.(a2 .+ ϵ))

    e_z2 = exp.(z2)
    l_part = (- e_z2 * e_z2' + Diagonal(e_z2 .* sum(e_z2))) / sum(e_z2)^2

    l_a2 = - y ./ (a2 .+ ϵ)
    l_z2 = l_part * l_a2
    l_a1 = m.W2' * l_z2
    l_z1 = l_a1 .* (a1 .> 0)
    l_x = m.W1' * l_z1

    l_W2 = l_z2 * a1'
    l_b2 = l_z2
    l_W1 = l_z1 * x'
    l_b1 = l_z1

    return l, l_W1, l_b1, l_W2, l_b2
end

#+

g_all = [grad(m, X_train[:,k], y_train[:,k]) for k in 1:size(X_train,2)]
typeof(g_all)

#+

function mean_tuple(d::AbstractArray{<:Tuple})
    Tuple([mean([d[k][i] for k in 1:length(d)]) for i in 1:length(d[1])])
end

#+

g_mean = mean_tuple(g_all)
typeof(g_mean)

# ### Exercise:
# Train the network with a gradient descent with stepsize ``\alpha=0.1`` for ``200``
# iterations. Save the objective value at each iteration and plot the results.
#
# ---
# ### Solution:

α = 1e-1
max_iter = 200
L = zeros(max_iter)

for iter in 1:max_iter
    grad_all = [grad(m, X_train[:,k], y_train[:,k]) for k in 1:size(X_train,2)]
    grad_mean = mean_tuple(grad_all)

    L[iter] = grad_mean[1]

    m.W1 .-= α*grad_mean[2]
    m.b1 .-= α*grad_mean[3]
    m.W2 .-= α*grad_mean[4]
    m.b2 .-= α*grad_mean[5]
end

# ---

using Plots

plot(L;
    xlabel="Iteration",
    ylabel="Loss function",
    label="",
    title="Loss function on the training set"
)

# ## Prediction
# ### Exercise:
#
# Write a function which predict the labels for samples. Show the accuracy on both training # and testing sets.
#
# ---
# ### Solution:

predict(X) = m(X)
accuracy(X, y) = mean(onecold(predict(X), classes) .== onecold(y, classes))

println("Train accuracy = ", accuracy(X_train, y_train))
println("Test accuracy = ", accuracy(X_test, y_test))


text:
import GeoStatsBase: solve, solvesingle
import MLJModelInterface
import TableOperations

const MI = MLJModelInterface
const TO = TableOperations

########################
# DUMMY LEARNING MODEL
########################

struct DummyModel <: MI.Supervised end
MI.fit(m::DummyModel, v, X, y) = rand(unique(y),length(y)), 0, 0
MI.predict(m::DummyModel, θ, X) = rand(θ, Tables.rowcount(X))
MI.target_scitype(m::DummyModel) = AbstractVector{<:MI.Finite}

#################
# DUMMY SOLVERS
#################

@estimsolver DummyEstimSolver begin end
function solve(problem::EstimationProblem, solver::DummyEstimSolver)
  sdat = data(problem)
  sdom = domain(problem)
  npts = nelements(sdom)
  vars = name.(variables(problem))
  μs = [v => fill(mean(sdat[v]), npts) for v in vars]
  σs = [Symbol(v,:Var) => fill(var(sdat[v]), npts) for v in vars]
  georef((; μs..., σs...), sdom)
end

@estimsolver ESolver begin
  @param A=1.0
  @param B=2
  @jparam J="foo"
  @global C=true
end

@simsolver SSolver begin
  @param A=1.0
  @param B=2
  @jparam J="foo"
  @global C=true
end

function solvesingle(problem::SimulationProblem, covars, ::SSolver, preproc)
  Dict(var => rand(100) for var in covars.names)
end

struct DummyLearnSolver{M} <: LearningSolver
  model::M
end

function solve(problem::LearningProblem, solver::DummyLearnSolver)
  sdata = sourcedata(problem)
  tdata = targetdata(problem)
  ptask = task(problem)
  model = solver.model

  # learn task with source data
  stable = values(sdata)
  if issupervised(ptask)
    X = TO.select(stable, features(ptask)...)
    y = Tables.getcolumn(stable, label(ptask))
    θ, _, __ = MI.fit(model, 0, X, y)
  else
    X = TO.select(stable, features(ptask)...)
    θ, _, __ = MI.fit(model, 0, X)
  end

  # perform task with target data
  ttable = values(tdata)
  X = TO.select(ttable, features(ptask)...)
  ŷ = MI.predict(model, θ, X)

  # post-process result
  var = outputvars(ptask)[1]
  val = if issupervised(ptask)
    isprobabilistic(model) ? mode.(ŷ) : ŷ
  else
    ŷ
  end

  # georeference on target domain
  ctor = constructor(typeof(tdata))
  dom  = domain(tdata)
  dat  = Dict(paramdim(dom) => (; var=>val))

  ctor(dom, dat)
end


text:
# canonical data version: 1.3.0

using Test

include("complex-numbers.jl")

@test ComplexNumber <: Number

@test ComplexNumber(0, 1)^2 == ComplexNumber(-1, 0)

@testset "Arithmetic" begin
    @testset "Addition" begin
        @test ComplexNumber(1, 0) + ComplexNumber(2, 0) == ComplexNumber(3, 0)
        @test ComplexNumber(0, 1) + ComplexNumber(0, 2) == ComplexNumber(0, 3)
        @test ComplexNumber(1, 2) + ComplexNumber(3, 4) == ComplexNumber(4, 6)
    end

    @testset "Subtraction" begin
        @test ComplexNumber(1, 0) - ComplexNumber(2, 0) == ComplexNumber(-1, 0)
        @test ComplexNumber(0, 1) - ComplexNumber(0, 2) == ComplexNumber(0, -1)
        @test ComplexNumber(1, 2) - ComplexNumber(3, 4) == ComplexNumber(-2, -2)
    end

    @testset "Multiplication" begin
        @test ComplexNumber(1, 0) * ComplexNumber(2, 0) == ComplexNumber(2, 0)
        @test ComplexNumber(0, 1) * ComplexNumber(0, 2) == ComplexNumber(-2, 0)
        @test ComplexNumber(1, 2) * ComplexNumber(3, 4) == ComplexNumber(-5, 10)
    end

    @testset "Division" begin
        @test ComplexNumber(1, 0) / ComplexNumber(2, 0) == ComplexNumber(0.5, 0)
        @test ComplexNumber(0, 1) / ComplexNumber(0, 2) == ComplexNumber(0.5, 0)
        @test ComplexNumber(1, 2) / ComplexNumber(3, 4) ≈ ComplexNumber(0.44, 0.08)
    end
end

@testset "Absolute value" begin
    @test abs(ComplexNumber(5, 0))  == 5
    @test abs(ComplexNumber(-5, 0)) == 5
    @test abs(ComplexNumber(0, 5))  == 5
    @test abs(ComplexNumber(0, -5)) == 5
    @test abs(ComplexNumber(3, 4))  == 5
end

@testset "Complex conjugate" begin
    @test conj(ComplexNumber(5, 0))  == ComplexNumber(5, 0)
    @test conj(ComplexNumber(0, 5))  == ComplexNumber(0, -5)
    @test conj(ComplexNumber(1, 1))  == ComplexNumber(1, -1)
end

@testset "Real part" begin
    @test real(ComplexNumber(1, 0)) == 1
    @test real(ComplexNumber(0, 1)) == 0
    @test real(ComplexNumber(1, 2)) == 1
end

@testset "Imaginary part" begin
    @test imag(ComplexNumber(1, 0)) == 0
    @test imag(ComplexNumber(0, 1)) == 1
    @test imag(ComplexNumber(1, 2)) == 2
end

# Bonus A
@testset "Complex exponential" begin
    @test_skip exp(ComplexNumber(0, π)) ≈ ComplexNumber(-1, 0)
    @test_skip exp(ComplexNumber(0, 0)) == ComplexNumber(1, 0)
    @test_skip exp(ComplexNumber(1, 0)) ≈ ComplexNumber(ℯ, 0)
    @test_skip exp(ComplexNumber(log(2), π)) ≈ ComplexNumber(-2, 0)
end

# Bonus B
@testset "Syntax sugar jm" begin
    @test_skip ComplexNumber(0, 1)  == jm
    @test_skip ComplexNumber(1, 0)  == 1 + 0jm
    @test_skip ComplexNumber(1, 1)  == 1 + 1jm
    @test_skip ComplexNumber(-1, 0) == jm^2
end


text:
function batch_plot(sim::EPOCHSimulation, f, folder_name, file_name, args...; overwrite=true, kwargs...)
    if !ispath(joinpath(sim.dir, folder_name))
        mkdir(joinpath(sim.dir, folder_name))
    end

    @time @progress name = "Creating $f plots for $sim.dir" for file in sim
        fn = joinpath(sim.dir, folder_name, file_name(file))
        if isfile(fn)
            @info "File $fn already exists"
            if overwrite
                @debug "Overwriting $fn..."
            else
                @debug "Skipping"
                continue
            end
        end
        plt = f(file, args...; kwargs...)
        isnothing(plt) && continue
        Plots.savefig(plt, fn)
    end
end

function batch_plot(sims, f, file_name::Function, args...; overwrite=true, kwargs...)
    @progress name="Batch plot" for sim in sims
        fn = joinpath(sim.dir, file_name(sim))
        if isfile(fn)
            @info "File $fn already exists"
            if overwrite
                @debug "Overwriting $fn..."
            else
                @debug "Skipping"
                continue
            end
        end
        plt = f(sim, args...; kwargs...)
        isnothing(plt) && continue
        Plots.savefig(plt, fn)
    end
end

function batch_plot(sims, f, folder_name, file_name::Function, args...; kwargs...)
    @progress name="Batch plot" for sim in sims
        batch_plot(sim, f, folder_name, file_name, args...; kwargs...)
    end
end


text:
import Base.iterate
export NotFactSampler, trace

# pdmp_inner!(rng, ∇ϕ!, ∇ϕx, t, x, θ, c::Bound, abc, (t′, renew), τref, v, (acc, num),
# Flow::Union{BouncyParticle, Boomerang}, args...; subsample=false, factor=1.5, adapt=false)

struct NotFactSampler{TF,T∇ϕ!,Tc,Tu0,Trng,Targs,Tkargs} <: PDMPSampler
    F::TF
    ∇ϕ!::T∇ϕ!
    c::Tc
    u0::Tu0 # t0 => (x0, θ0)
    rng::Trng

    args::Targs

    kargs::Tkargs
end

function NotFactSampler(∇ϕ!, u0, c, F::Union{BouncyParticle,Boomerang}, args...;
    factor=1.8, subsample=subsample, adapt=false, seed=Seed())
    kargs = (factor=factor, adapt=adapt, subsample=subsample)
    return NotFactSampler(F, ∇ϕ!, c, u0, Rng(seed), args, kargs)
end

function iterate(FS::NotFactSampler)
    t0, (x0, θ0) = FS.u0
    Flow = FS.F
    n = length(x0)
    t, x, θ, ∇ϕx = t0, copy(x0), copy(θ0), copy(θ0)
    c = FS.c
    rng = FS.rng
    Ξ = Trace(t0, x0, θ0, Flow)
    τref = waiting_time_ref(rng, Flow)
    ∇ϕx, v = FS.∇ϕ!(∇ϕx, t, x, θ, FS.args...)
    ∇ϕx = grad_correct!(∇ϕx, x, Flow)
    num = acc = 0
    abc = ab(x, θ, c, ∇ϕx, v, Flow)

    t′, renew = next_time(t, abc, rand(rng))

    iterate(FS, ((t => (x, θ)), ∇ϕx, (acc, num), c, abc, (t′, renew), τref, v))
end

function iterate(FS::NotFactSampler,  (u, ∇ϕx, (acc, num), c, abc, (t′, renew), τref, v))
    t, (x, θ) = u
    t, x, θ, (acc, num), c, abc, (t′, renew), τref, v = pdmp_inner!(FS.rng, FS.∇ϕ!, ∇ϕx, t, x, θ, c, abc, (t′, renew), τref, v, (acc, num), FS.F, FS.args...; FS.kargs...)
    ev = event(t, x, θ, FS.F)
    u = t => (x, θ)
    return ev, (u, ∇ϕx, (acc, num), c, abc, (t′, renew), τref, v)
end


text:
elements = [SA[1 0 0 0; 0 1 0 0; 0 0 1 0; 0 0 0 1],
SA[0 -1 0 0; 1 0 0 0; 0 0 1 0; 0 0 0 1],
SA[0 0 -1 0; 0 1 0 0; 1 0 0 0; 0 0 0 1],
SA[0 0 -1 0; 1 0 0 0; 0 -1 0 0; 0 0 0 1],
SA[0 0 0 -1; 0 1 0 0; 0 0 1 0; 1 0 0 0],
SA[0 0 0 -1; 1 0 0 0; 0 0 1 0; 0 -1 0 0],
SA[0 0 0 -1; 0 1 0 0; 1 0 0 0; 0 0 -1 0],
SA[0 0 0 -1; 1 0 0 0; 0 -1 0 0; 0 0 -1 0],
SA[-1 0 0 0; 0 -1 0 0; 0 0 1 0; 0 0 0 1],
SA[0 -1 0 0; 0 0 -1 0; 1 0 0 0; 0 0 0 1],
SA[-1 0 0 0; 0 0 -1 0; 0 -1 0 0; 0 0 0 1],
SA[0 -1 0 0; 0 0 0 -1; 0 0 1 0; 1 0 0 0],
SA[-1 0 0 0; 0 0 0 -1; 0 0 1 0; 0 -1 0 0],
SA[0 -1 0 0; 0 0 0 -1; 1 0 0 0; 0 0 -1 0],
SA[-1 0 0 0; 0 0 0 -1; 0 -1 0 0; 0 0 -1 0],
SA[-1 0 0 0; 0 1 0 0; 0 0 -1 0; 0 0 0 1],
SA[0 1 0 0; 1 0 0 0; 0 0 -1 0; 0 0 0 1],
SA[0 0 -1 0; 0 1 0 0; 0 0 0 -1; 1 0 0 0],
SA[0 0 -1 0; 1 0 0 0; 0 0 0 -1; 0 -1 0 0],
SA[-1 0 0 0; 0 1 0 0; 0 0 0 -1; 0 0 -1 0],
SA[0 1 0 0; 1 0 0 0; 0 0 0 -1; 0 0 -1 0],
SA[0 0 -1 0; 0 -1 0 0; -1 0 0 0; 0 0 0 1],
SA[0 1 0 0; 0 0 -1 0; -1 0 0 0; 0 0 0 1],
SA[0 0 -1 0; 0 0 0 -1; 0 -1 0 0; 1 0 0 0],
SA[0 0 -1 0; 0 0 0 -1; -1 0 0 0; 0 -1 0 0],
SA[0 1 0 0; 0 0 0 -1; -1 0 0 0; 0 0 -1 0],
SA[-1 0 0 0; 0 1 0 0; 0 0 1 0; 0 0 0 -1],
SA[0 1 0 0; 1 0 0 0; 0 0 1 0; 0 0 0 -1],
SA[0 0 1 0; 0 1 0 0; 1 0 0 0; 0 0 0 -1],
SA[0 0 1 0; 1 0 0 0; 0 -1 0 0; 0 0 0 -1],
SA[0 0 0 -1; 0 -1 0 0; 0 0 1 0; -1 0 0 0],
SA[0 0 0 -1; 0 0 -1 0; 1 0 0 0; 0 -1 0 0],
SA[0 0 0 -1; 0 0 -1 0; 0 -1 0 0; -1 0 0 0],
SA[0 1 0 0; 0 0 0 -1; 0 0 1 0; -1 0 0 0],
SA[0 0 1 0; 0 0 0 -1; 1 0 0 0; 0 -1 0 0],
SA[0 0 1 0; 0 0 0 -1; 0 -1 0 0; -1 0 0 0],
SA[0 0 0 -1; 0 1 0 0; 0 0 -1 0; -1 0 0 0],
SA[0 0 0 -1; 1 0 0 0; 0 0 -1 0; 0 1 0 0],
SA[0 0 1 0; 0 1 0 0; 0 0 0 -1; -1 0 0 0],
SA[0 0 1 0; 1 0 0 0; 0 0 0 -1; 0 1 0 0],
SA[0 0 0 -1; 0 -1 0 0; -1 0 0 0; 0 0 -1 0],
SA[0 0 0 -1; 0 0 -1 0; -1 0 0 0; 0 1 0 0],
SA[0 0 1 0; 0 0 0 -1; -1 0 0 0; 0 1 0 0],
SA[0 1 0 0; -1 0 0 0; 0 0 1 0; 0 0 0 1],
SA[0 0 1 0; 0 -1 0 0; 1 0 0 0; 0 0 0 1],
SA[0 0 1 0; -1 0 0 0; 0 -1 0 0; 0 0 0 1],
SA[0 0 0 1; 0 -1 0 0; 0 0 1 0; 1 0 0 0],
SA[0 0 0 1; -1 0 0 0; 0 0 1 0; 0 -1 0 0],
SA[0 0 0 1; 0 -1 0 0; 1 0 0 0; 0 0 -1 0],
SA[0 0 0 1; -1 0 0 0; 0 -1 0 0; 0 0 -1 0],
SA[0 -1 0 0; -1 0 0 0; 0 0 -1 0; 0 0 0 1],
SA[0 -1 0 0; 0 0 -1 0; 0 0 0 -1; 1 0 0 0],
SA[-1 0 0 0; 0 0 -1 0; 0 0 0 -1; 0 -1 0 0],
SA[0 -1 0 0; -1 0 0 0; 0 0 0 -1; 0 0 -1 0],
SA[0 0 1 0; 0 1 0 0; -1 0 0 0; 0 0 0 1],
SA[0 0 0 1; 0 0 -1 0; 0 -1 0 0; 1 0 0 0],
SA[0 0 0 1; 0 0 -1 0; -1 0 0 0; 0 -1 0 0],
SA[0 0 0 1; 0 1 0 0; -1 0 0 0; 0 0 -1 0],
SA[0 -1 0 0; -1 0 0 0; 0 0 1 0; 0 0 0 -1],
SA[0 -1 0 0; 0 0 1 0; 1 0 0 0; 0 0 0 -1],
SA[-1 0 0 0; 0 0 1 0; 0 -1 0 0; 0 0 0 -1],
SA[0 0 0 1; 0 1 0 0; 0 0 1 0; -1 0 0 0],
SA[0 0 0 1; 0 0 1 0; 1 0 0 0; 0 -1 0 0],
SA[0 0 0 1; 0 0 1 0; 0 -1 0 0; -1 0 0 0],
SA[0 -1 0 0; 0 0 0 -1; 0 0 -1 0; -1 0 0 0],
SA[-1 0 0 0; 0 0 0 -1; 0 0 -1 0; 0 1 0 0],
SA[0 -1 0 0; 0 0 1 0; 0 0 0 -1; -1 0 0 0],
SA[-1 0 0 0; 0 0 1 0; 0 0 0 -1; 0 1 0 0],
SA[0 0 0 1; 0 0 1 0; -1 0 0 0; 0 1 0 0],
SA[0 0 1 0; 1 0 0 0; 0 1 0 0; 0 0 0 1],
SA[0 0 0 1; 0 1 0 0; 0 0 -1 0; 1 0 0 0],
SA[0 0 0 1; 1 0 0 0; 0 0 -1 0; 0 -1 0 0],
SA[0 0 0 1; 1 0 0 0; 0 1 0 0; 0 0 -1 0],
SA[1 0 0 0; 0 -1 0 0; 0 0 -1 0; 0 0 0 1],
SA[1 0 0 0; 0 0 -1 0; 0 1 0 0; 0 0 0 1],
SA[0 1 0 0; 0 0 0 -1; 0 0 -1 0; 1 0 0 0],
SA[1 0 0 0; 0 0 0 -1; 0 0 -1 0; 0 -1 0 0],
SA[1 0 0 0; 0 0 0 -1; 0 1 0 0; 0 0 -1 0],
SA[0 0 -1 0; 0 1 0 0; -1 0 0 0; 0 0 0 -1],
SA[0 0 -1 0; 1 0 0 0; 0 1 0 0; 0 0 0 -1],
SA[0 0 -1 0; 0 -1 0 0; 0 0 0 -1; -1 0 0 0],
SA[0 1 0 0; 0 0 -1 0; 0 0 0 -1; -1 0 0 0],
SA[0 0 -1 0; 0 0 0 -1; 0 1 0 0; -1 0 0 0],
SA[0 0 0 1; 1 0 0 0; 0 0 1 0; 0 1 0 0],
SA[1 0 0 0; 0 -1 0 0; 0 0 0 -1; 0 0 -1 0],
SA[1 0 0 0; 0 0 -1 0; 0 0 0 -1; 0 1 0 0],
SA[1 0 0 0; 0 0 0 -1; 0 0 1 0; 0 1 0 0],
SA[0 0 -1 0; -1 0 0 0; 0 1 0 0; 0 0 0 1],
SA[0 0 -1 0; 0 -1 0 0; 0 0 0 1; 1 0 0 0],
SA[0 0 -1 0; -1 0 0 0; 0 0 0 1; 0 -1 0 0],
SA[-1 0 0 0; 0 -1 0 0; 0 0 0 1; 0 0 -1 0],
SA[0 1 0 0; -1 0 0 0; 0 0 0 1; 0 0 -1 0],
SA[0 1 0 0; 0 0 -1 0; 0 0 0 1; 1 0 0 0],
SA[1 0 0 0; 0 0 -1 0; 0 0 0 1; 0 -1 0 0],
SA[1 0 0 0; 0 1 0 0; 0 0 0 1; 0 0 -1 0],
SA[0 0 -1 0; -1 0 0 0; 0 -1 0 0; 0 0 0 -1],
SA[0 1 0 0; 0 0 1 0; -1 0 0 0; 0 0 0 -1],
SA[0 0 -1 0; 0 1 0 0; 0 0 0 1; -1 0 0 0],
SA[-1 0 0 0; 0 0 1 0; 0 0 0 1; 0 -1 0 0],
SA[0 1 0 0; 0 0 1 0; 0 0 0 1; -1 0 0 0],
SA[1 0 0 0; 0 0 1 0; 0 0 0 1; 0 1 0 0],
SA[0 0 0 1; 0 1 0 0; 1 0 0 0; 0 0 1 0],
SA[0 0 0 1; 1 0 0 0; 0 -1 0 0; 0 0 1 0],
SA[1 0 0 0; 0 -1 0 0; 0 0 1 0; 0 0 0 -1],
SA[0 1 0 0; 0 0 -1 0; 1 0 0 0; 0 0 0 -1],
SA[1 0 0 0; 0 0 -1 0; 0 -1 0 0; 0 0 0 -1],
SA[0 1 0 0; 0 0 0 -1; 1 0 0 0; 0 0 1 0],
SA[1 0 0 0; 0 0 0 -1; 0 -1 0 0; 0 0 1 0],
SA[1 0 0 0; 0 1 0 0; 0 0 -1 0; 0 0 0 -1],
SA[0 -1 0 0; 1 0 0 0; 0 0 -1 0; 0 0 0 -1],
SA[1 0 0 0; 0 1 0 0; 0 0 0 -1; 0 0 1 0],
SA[0 -1 0 0; 1 0 0 0; 0 0 0 -1; 0 0 1 0],
SA[0 0 1 0; 0 -1 0 0; -1 0 0 0; 0 0 0 -1],
SA[0 -1 0 0; 0 0 -1 0; -1 0 0 0; 0 0 0 -1],
SA[0 -1 0 0; 0 0 0 -1; -1 0 0 0; 0 0 1 0],
SA[0 0 0 -1; -1 0 0 0; 0 0 1 0; 0 1 0 0],
SA[0 0 0 -1; 0 -1 0 0; 1 0 0 0; 0 0 1 0],
SA[0 0 0 -1; -1 0 0 0; 0 -1 0 0; 0 0 1 0],
SA[0 0 0 -1; -1 0 0 0; 0 0 -1 0; 0 -1 0 0],
SA[0 0 1 0; -1 0 0 0; 0 0 0 -1; 0 -1 0 0],
SA[0 0 0 -1; 0 1 0 0; -1 0 0 0; 0 0 1 0],
SA[0 1 0 0; 0 0 1 0; 1 0 0 0; 0 0 0 1],
SA[1 0 0 0; 0 0 1 0; 0 -1 0 0; 0 0 0 1],
SA[1 0 0 0; 0 0 1 0; 0 0 0 -1; 0 -1 0 0],
SA[0 -1 0 0; 0 0 1 0; -1 0 0 0; 0 0 0 1],
SA[0 0 0 -1; 1 0 0 0; 0 1 0 0; 0 0 1 0],
SA[0 0 0 -1; 0 -1 0 0; 0 0 -1 0; 1 0 0 0],
SA[0 0 0 -1; 0 0 -1 0; 0 1 0 0; 1 0 0 0],
SA[0 0 1 0; 0 0 0 -1; 0 1 0 0; 1 0 0 0],
SA[0 0 1 0; 0 -1 0 0; 0 0 0 -1; 1 0 0 0],
SA[0 0 0 -1; -1 0 0 0; 0 1 0 0; 0 0 -1 0],
SA[0 0 1 0; 0 -1 0 0; 0 0 0 1; -1 0 0 0],
SA[0 0 1 0; -1 0 0 0; 0 0 0 1; 0 1 0 0],
SA[-1 0 0 0; 0 0 -1 0; 0 0 0 1; 0 1 0 0],
SA[0 0 1 0; 0 1 0 0; 0 0 0 1; 1 0 0 0],
SA[0 -1 0 0; 0 0 1 0; 0 0 0 1; 1 0 0 0],
SA[0 1 0 0; 0 0 0 1; 0 0 1 0; 1 0 0 0],
SA[1 0 0 0; 0 0 0 1; 0 0 1 0; 0 -1 0 0],
SA[0 1 0 0; 0 0 0 1; 1 0 0 0; 0 0 -1 0],
SA[1 0 0 0; 0 0 0 1; 0 -1 0 0; 0 0 -1 0],
SA[0 0 1 0; 0 0 0 1; 0 -1 0 0; 1 0 0 0],
SA[0 0 1 0; 0 0 0 1; -1 0 0 0; 0 -1 0 0],
SA[0 -1 0 0; 0 0 0 1; -1 0 0 0; 0 0 -1 0],
SA[0 0 -1 0; 0 -1 0 0; 1 0 0 0; 0 0 0 -1],
SA[0 -1 0 0; 0 0 0 1; 0 0 1 0; -1 0 0 0],
SA[0 0 -1 0; 0 0 0 1; 1 0 0 0; 0 -1 0 0],
SA[0 0 -1 0; 0 0 0 1; 0 -1 0 0; -1 0 0 0],
SA[0 0 0 1; 0 -1 0 0; 0 0 -1 0; -1 0 0 0],
SA[0 0 0 1; -1 0 0 0; 0 0 -1 0; 0 1 0 0],
SA[0 0 -1 0; -1 0 0 0; 0 0 0 -1; 0 1 0 0],
SA[0 0 -1 0; 0 0 0 1; -1 0 0 0; 0 1 0 0],
SA[-1 0 0 0; 0 0 1 0; 0 1 0 0; 0 0 0 1],
SA[0 -1 0 0; 0 0 0 1; 0 0 -1 0; 1 0 0 0],
SA[-1 0 0 0; 0 0 0 1; 0 0 -1 0; 0 -1 0 0],
SA[-1 0 0 0; 0 0 0 1; 0 1 0 0; 0 0 -1 0],
SA[-1 0 0 0; 0 0 -1 0; 0 1 0 0; 0 0 0 -1],
SA[0 0 0 1; 0 0 -1 0; 0 1 0 0; -1 0 0 0],
SA[-1 0 0 0; 0 0 0 1; 0 0 1 0; 0 1 0 0],
SA[0 0 1 0; 1 0 0 0; 0 0 0 1; 0 -1 0 0],
SA[0 -1 0 0; 1 0 0 0; 0 0 0 1; 0 0 -1 0],
SA[0 -1 0 0; 0 0 -1 0; 0 0 0 1; -1 0 0 0],
SA[0 0 -1 0; 1 0 0 0; 0 0 0 1; 0 1 0 0],
SA[0 -1 0 0; 0 0 0 1; 1 0 0 0; 0 0 1 0],
SA[-1 0 0 0; 0 0 0 1; 0 -1 0 0; 0 0 1 0],
SA[-1 0 0 0; 0 -1 0 0; 0 0 -1 0; 0 0 0 -1],
SA[-1 0 0 0; 0 -1 0 0; 0 0 0 -1; 0 0 1 0],
SA[0 0 0 1; 0 -1 0 0; -1 0 0 0; 0 0 1 0],
SA[-1 0 0 0; 0 0 0 -1; 0 1 0 0; 0 0 1 0],
SA[0 0 0 1; 0 0 1 0; 0 1 0 0; 1 0 0 0],
SA[0 1 0 0; 0 0 1 0; 0 0 0 -1; 1 0 0 0],
SA[0 0 0 1; 0 0 -1 0; 1 0 0 0; 0 1 0 0],
SA[0 0 -1 0; 0 0 0 -1; 1 0 0 0; 0 1 0 0],
SA[0 1 0 0; -1 0 0 0; 0 0 -1 0; 0 0 0 -1],
SA[1 0 0 0; 0 0 1 0; 0 1 0 0; 0 0 0 -1],
SA[0 0 0 -1; 0 0 1 0; -1 0 0 0; 0 -1 0 0],
SA[0 0 0 -1; 0 0 1 0; 0 1 0 0; -1 0 0 0],
SA[0 0 0 -1; 0 0 1 0; 1 0 0 0; 0 1 0 0],
SA[-1 0 0 0; 0 1 0 0; 0 0 0 1; 0 0 1 0],
SA[0 1 0 0; 1 0 0 0; 0 0 0 1; 0 0 1 0],
SA[0 1 0 0; -1 0 0 0; 0 0 0 -1; 0 0 1 0],
SA[0 0 0 -1; 0 0 1 0; 0 -1 0 0; 1 0 0 0],
SA[0 0 -1 0; 0 0 0 1; 0 1 0 0; 1 0 0 0],
SA[0 1 0 0; 0 0 0 1; 0 0 -1 0; -1 0 0 0],
SA[1 0 0 0; 0 0 0 1; 0 0 -1 0; 0 1 0 0],
SA[0 1 0 0; 0 0 0 1; -1 0 0 0; 0 0 1 0],
SA[1 0 0 0; 0 -1 0 0; 0 0 0 1; 0 0 1 0],
SA[0 0 1 0; -1 0 0 0; 0 1 0 0; 0 0 0 -1],
SA[0 -1 0 0; -1 0 0 0; 0 0 0 1; 0 0 1 0],
SA[0 0 1 0; 0 0 0 1; 1 0 0 0; 0 1 0 0],
SA[0 0 1 0; 0 0 0 1; 0 1 0 0; -1 0 0 0],
SA[0 0 0 1; -1 0 0 0; 0 1 0 0; 0 0 1 0],
SA[1 0 0 0; 0 0 0 1; 0 1 0 0; 0 0 1 0]]