module Milann

using Flux
import Flux.Tracker: data, @grad, track

using Statistics

# This implemens a MIL version where the bag instances are stored in a continuous tensor and bags are delimited by ranges.

export RangeMIL, segmax, segmean, segmax_naive, segmean_naive, segmaxmean

struct RangeMIL
    premodel
    aggregation
    postmodel
end
@Flux.treelike RangeMIL

function (m::RangeMIL)(X::AbstractArray, B::AbstractArray)
    (nfeatures, ninsts), nbags = size(X), length(B)
    # X: nfeatures x ninsts, B: nbags
    Ypre = m.premodel(X)
    # Ypre: npremodeloutput x ninstances
    Yagg = m.aggregation(Ypre, B)
    # Yagg: npremodeloutput x nbags
    m.postmodel(Yagg)
end

"""
segmax(X, segments)

Compute segmented maximum for instances `X` and bags indexed by `segments`.
"""
function segmax(X, segments)
    Y = similar(X, size(X, 1), length(segments))

    # naive approach, fast on CPU
    #     for (i, seg) in enumerate(segments)
    #         Y[:, i:i] .= mean(view(X, :, seg), dims=2)
    #     end

    # sliced approach much better for GPU
    Y .= 0
    mlen = maximum(length.(segments))
    for i in 1:mlen
        Yindices = [j for (j, e) in enumerate(segments) if length(e) >= i]
        Xindices = map(q -> first(q) + i - 1, filter(e -> length(e) >= i, segments))
        Y[:, Yindices] = max.(view(Y, :, Yindices), view(X, :, Xindices))
    end
    Y
end

"""
segmax_idx(X, segments)

Compute segmented maximum for instances `X` and bags indexed by `segments`. Same as segmax, but returns also the indices of maxima in input (used for grad computation).
"""
function segmax_idx(X, segments)
    Y = similar(X, size(X, 1), length(segments))
    Y2 = similar(X, size(Y)...) # twice the memmory!
    idxs = similar(X, size(Y)...)
    Y .= -Inf
    Y2 .= NaN
    idxs .= 0
    mlen = maximum(length.(segments))
    for i in 1:mlen
        Yindices = [j for (j, e) in enumerate(segments) if length(e) >= i]
        Xindices = map(q -> first(q) + i - 1, filter(e -> length(e) >= i, segments))
        Y2[:, Yindices] = max.(view(Y, :, Yindices), view(X, :, Xindices))
        idxs[:, Yindices] += view(Y, :, Yindices)  .!= view(Y2, :, Yindices) # has been max increased since last iteration?
        tmp = Y; Y = Y2; Y2 = tmp # just copy references
    end
    # idxs now holds indices of maxima relative to bags, make them relative to inputs
    offsets = first.(segments) .- 1 # first indices in X minus 1
    Y, Int.(cpu(idxs) .+ reshape(offsets, 1, length(offsets)))
end

segmax(X::TrackedArray, segments) = track(segmax, X, segments)

function dsegmax(Δ, segments, idxs)
    grads = similar(Δ, size(Δ, 1), last(last(segments)))
    grads .= 0
    for (i, row) in enumerate(eachrow(idxs))
        # @show size(row), minimum(row), size(Δ)
        grads[i, row] .= Δ[i, :]
    end
    grads, nothing
end

@grad function segmax(X, segments)
    Y, idxs = segmax_idx(data(X), segments)
    Y, Δ -> dsegmax(Δ, segments, idxs)
end

"""
segmean(X, segments)

Compute segmented mean for instances `X` and bags indexed by `segments`.
"""
function segmean(X, segments)
    Y = similar(X, size(X, 1), length(segments))

    # naive approach, fast on CPU
    #     for (i, seg) in enumerate(segments)
    #         Y[:, i:i] .= mean(view(X, :, seg), dims=2)
    #     end

    # sliced approach much better for GPU
    Y .= 0
    mlen = maximum(length.(segments))
    for i in 1:mlen
        Yindices = [j for (j, e) in enumerate(segments) if length(e) >= i]
        Xindices = map(q -> first(q) + i - 1, filter(e -> length(e) >= i, segments))
        Y[:, Yindices] += view(X, :, Xindices)
    end
    Y ./= reshape(gpu(length.(segments)), 1, length(segments))
    Y
end

segmean(X::TrackedArray, segments) = track(segmean, X, segments)

function dsegmean(Δ, segments)
    grads = similar(Δ, size(Δ, 1), last(last(segments)))

    # naive approach, fast on CPU
    #     for (i, seg) in enumerate(segments)
    #         grads[:, seg] .= view(Δ, :,i) / length(seg)
    #     end

    # sliced approach much better for GPU
    mlen = maximum(length.(segments))
    for i in 1:mlen
        Yindices = [j for (j, e) in enumerate(segments) if length(e) >= i]
        Xindices = map(q -> first(q) + i - 1, filter(e -> length(e) >= i, segments))
        grads[:, Xindices] .= view(Δ, :, Yindices)
    end
    for i in 2:mlen # normalization
        Xindices = vcat((collect(seg) for seg in segments if length(seg) == i)...)
        if length(Xindices) > 0
            grads[:, Xindices] ./= i
        end
    end

    grads, nothing
end

@grad segmean(X, segments) = segmean(data(X), segments), Δ -> dsegmean(Δ, segments)

"""
segmax_naive(X, segments)

Segmented maximum: naive reference implementation using default AD.
"""
function segmax_naive(X, segments)
    hcat((maximum(X[:,s], dims=ndims(X)) for s in segments)...)
end

"""
    segmean_naive(X, segments)

Segmented mean: naive reference implementation using default AD.
"""
function segmean_naive(X, segments)
    hcat((mean(X[:,s], dims=ndims(X)) for s in segments)...)
end

"""
    segmaxmean(X, segments)

Combine `segmax` and `segmean`.
"""
segmaxmean(X, segments) = vcat(segmax(X, segments), segmean(X, segments))

end # module


## Step2 is iteration to find the stationary distribution

#module Step2

using LinearAlgebra, Random, Distributions, Statistics, Plots
using BenchmarkTools, Interpolations
include("Step1Grid.jl")

"""
    Relies on step1 (either with simple grid search or endogenous) and finds
    stationary distribution.
"""

# Define function for inverse of policy function a^{-1}
function pol_inv(A, pol_func, aprime, eprime)

    """
    Function that finds the inverse of the optimal policy function.
    Inputs:
        - A: set of assets
        - pol_func: optimal policy function
        - aprime: assets in the next period
        - eprime: exogenous state variable in the next period
    Output: inverse of the policy function
    """

    # if a' < minimum value of policy function, then a must be a_lowerbar
    if minimum(pol_func[:, eprime]) > A[aprime]

        return 1

    # when e'=2, distinguish pre vs post cross 45° line (endogenous a_UpperBar)
    elseif eprime == 2

        return min(findlast(pol_func[:, eprime] .<= A[aprime]),
            findfirst((pol_func[:, eprime] - A) .< 0.01))

    # any other case, find element of A s.t. policy function is closest to a'
    else

        return findlast(pol_func[:, eprime] .<= A[aprime])

    end

end

# Define function for Step2
function step2(A, E, Π, pol_func; tol = 0.01, maxT = 600)

    """
    Function that finds the stationary distribution Ψ.
    Relies on pol_inv in the code.
    Inputs:
        - A: set of assets
        - E: set of exogenous shocks
        - Π: transition matrix
        - pol_func: optimal policy function
    Output:
        - F_update: stationary distribution Ψ
    """

    # retrieve dimensions
    a_size = length(A)
    e_size = size(E)[1]

    # Check that everything works
    @assert size(pol_func)[1] == a_size
    @assert size(pol_func)[2] == e_size

    # Find endogenous position where a' = a
    aBar = findlast((pol_func[:,2] - A) .< tol)
    cut = max(aBar, 2)           # if aBar = 1 then range will not work

    aux1 = range(zero(eltype(A)), one(eltype(A))/e_size, length = cut)
    aux2 = ones(eltype(A), a_size - cut)/e_size
    aux = [aux1; aux2]
    F0 = repeat(aux, 1, e_size)

    F_update = copy(F0)

    # Iteration of distribution
    normdiff = Inf
    t = 1

    while normdiff > tol && t <= maxT

        for a = 1:a_size
            for e = 1:e_size

                F_update[a,e] = Π[1, e] * F0[pol_inv(A, pol_func, a, 1), 1] +
                          Π[2, e] * F0[pol_inv(A, pol_func, a, 2), 2]

            end
        end

        t += 1
        normdiff = maximum(abs.(F_update - F0))

        F0 = copy(F_update)

    end

    return F_update

end

#export pol_inv, step2

#end


module RandomWalk

using Distributed

using Interpolations
using HDF5

export dom_random_walk, biased_random_walk

struct Trajectory
    positions::Array{Int, 2}
    times::Array{Float64, 1}
    box::Int
end

function translate(s::Array{Int, 1}, step::Array{Int, 1}, N::Int)
    p  = s + step
    for i = 1:3
        if p[i] > N
            p[i] = p[i] - N
        elseif p[i] < 1
            p[i] = N - p[i]
        end
    end
    return p
end

function next_step(traj::Trajectory, step::Int, rates)
    # choose random number and calc wait times
    z = rand(2 * size(traj.positions, 2))
    tw = - 1 ./ rates(traj.positions[step, :]) .* log.(1 .- z)
    # select direction of shortest tw
    t = minimum(tw)
    i = indexin([t], tw)[1]
    d = ceil.(Int, i/2)
    p1 = traj.positions[step, :]
    p1[d] += 2 * (i % 2) - 1
    # append new position to trajectory
    traj.times[step + 1] = traj.times[step] + t
    traj.positions[step + 1, :] = p1
end

function random_walk(steps::Int, timemax::Float64, rates, N; dim::Int=3)
    traj = Trajectory(Array{Int}(undef, steps, dim), Array{Float64}(undef, steps), N)
    traj.times[1] = 0.0
    traj.positions[1, :] = ceil.(Int, rand(dim) * N)
    for i = 1:steps-1
        if traj.times[i] > timemax
            # when maximum time is reached, fill array and end loop
            traj.times[i + 1:end] .= traj.times[i] + 1.0
            # println("Reached maxtime in $i steps.")
            break
        end
        next_step(traj, i, rates)
    end
    if traj.times[end] < timemax
        println("Maxtime was not reached: times[end] = $(traj.times[end])")
    end
    return traj
end

function to_uniform_time(time::AbstractArray{Float64, 1}, tr::Trajectory)
    pos = Array{Int}(undef, length(time), size(tr.positions, 2))
    j = 1
    for (i, t) in enumerate(time)
        while (j < length(tr.times)) & (tr.times[j+1] <= t)
            j += 1
        end
        # pos[i, :] = tr.positions[tr.times .<= t, :][end, :]
        pos[i, :] = tr.positions[j, :]
    end
    return pos
end

function savedata(fname::String, time, traj, box)
    # traj has shape (times, 3, numtraj)
    h5open(fname, "w") do file
        write(file, "box", box)
        write(file, "time", time)
        g = HDF5.g_create(file, "positions")
        for i = 1:length(time)
            g["$(i-1)"] = traj[i, :, :]
        end
    end
end

function dom_random_walk(numtraj::Int, N::Int, steps::Int, time::AbstractArray{Float64,1}; rate_1=1e-2, rate_2=1.0)#::Array{Float64, 3}
    timemax = maximum(time)

    traj = @distributed (a, b) -> cat(a, b, dims=3) for i = 1:numtraj
        dom_rates(pos::Array{Int, 1}) = Bool(sum(mod1.(pos, N) .> N/2) % 2) ? rate_1 : rate_2
        tr = random_walk(steps, timemax, dom_rates, N)
        to_uniform_time(time, tr)
    end
    return traj
end

function biased_random_walk(numtraj::Int, N::Int, steps::Int, time::AbstractArray{Float64, 1}; rate_1=1e-2, rate_2=1.0, bias=0.25)
    timemax = maximum(time)

    traj = @distributed (a, b) -> cat(a, b, dims=3) for i = 1:numtraj

        function biased_rates(pos::Array{Int, 1})
            x = mod1.(pos, N)[1]
            if (x == 1) | (x == N)
                b = exp(bias)
            elseif (x == (N / 2)) | (x == (N / 2 + 1))
                b = exp(-bias)
            else
                b = 1
            end
            # b = 1 + bias * cos(2 * pi * (pos[1] - 0.5) / N)
            k = Bool(sum(x .> N/2) % 2) ? rate_1 : rate_2
            [1/b, b] .* k
        end

        tr = random_walk(steps, timemax, biased_rates, N, dim=1)
        to_uniform_time(time, tr)
    end
    return traj

end

end


using DataStructures

using PyCall
p_VIEW = p["VIEW"]
p_STRUCT = p["STRUCT"]
p_MKPOL = p["MKPOL"]

using LinearAlgebraicRepresentation
using Plasm

import Base.cat

L = LinearAlgebraicRepresentation

View = Plasm.view

"""
	apply(affineMatrix::Matrix)(larmodel::LAR)::LAR

Apply the `affineMatrix` parameter to the vertices of `larmodel`.

# Example

```
julia> square = LinearAlgebraicRepresentation.cuboid([1,1])
([0.0 0.0 1.0 1.0; 0.0 1.0 0.0 1.0], Array{Int64,1}[[1, 2, 3, 4]])

julia> Plasm.apply(LinearAlgebraicRepresentation.t(1,2))(square)
([1.0 1.0 2.0 2.0; 2.0 3.0 2.0 3.0], Array{Int64,1}[[1, 2, 3, 4]])
```
"""
function apply(affineMatrix)
	function apply0(larmodel)
		return L.struct2lar(L.Struct([ affineMatrix,larmodel ]))
	end
	return apply0
end

"""
	comp(funs::Array)

Standard mathematical composition.

Pipe execution from right to left on application to actual parameter.
"""
function comp(funs)
    function compose(f,g)
	  return x -> f(g(x))
	end
    id = x->x
    return reduce(compose, funs; init=id)
end

"""
	cons(funs::Array)(x::Any)::Array

*Construction* functional of FL and PLaSM languages.

Provides a *vector* functional that returns the array of
applications of component functions to actual parameter.

# Example

```
julia> Plasm.cons([cos,sin])(0)
2-element Array{Float64,1}:
 1.0
 0.0
```
"""
function cons(funs)
	return x -> [f(x) for f in funs]
end

"""
	k(Any)(x)

*Constant* functional of FL and PLaSM languages.

Gives a constant functional that always returns the actual parameter
when applied to another parameter.

#	Examples

```
julia> Plasm.k(10)(100)
10

julia> Plasm.k(sin)(cos)
sin
```
"""
function k(Any)
	x->Any
end

"""
	aa(fun::Function)(args::Array)::Array

AA applies fun to each element of the args sequence

# Example

```
julia> Plasm.aa(sqrt)([1,4,9,16])
4-element Array{Float64,1}:
 1.0
 2.0
 3.0
 4.0
```
"""
function aa(fun)
	function aa1(args::Array)
		map(fun,args)
	end
	return aa1
end

"""
	id(x::Anytype)

Identity function.  Return the argument.

"""
id = x->x

"""
	distr(args::Union{Tuple,Array})::Array

Distribute right. The parameter `args` must contain a `list` and an element `x`.
Return the `pair` array with the elements of `args` coupled with `x`

# Example

```
julia> Plasm.distr(([1,2,3],10))
3-element Array{Array{Int64,1},1}:
 [1, 10]
 [2, 10]
 [3, 10]

julia> Plasm.distr([[1,2,3],10])
3-element Array{Array{Int64,1},1}:
 [1, 10]
 [2, 10]
 [3, 10]
```
"""
function distr(args)
	list,element = args
	return [ [e,element] for e in list ]
end

"""
	distl(args::Union{Tuple,Array})::Array

Distribute right. The parameter `args` must contain an element `x` and a `list`.
Return the `pair` array with `x` coupled with the elements of `args`.

# Example

```
julia> Plasm.distl((10, [1,2,3]))
3-element Array{Array{Int64,1},1}:
 [10, 1]
 [10, 2]
 [10, 3]

julia> Plasm.distl([10, [1,2,3]])
3-element Array{Array{Int64,1},1}:
 [10, 1]
 [10, 2]
 [10, 3]
```
"""
function distl(args)
	element, list = args
	return [ [element, e] for e in list ]
end

# Vector definition of printable ASCII codes as one-dimensional LAR models.
# Font design for *Geometric Programming for Computer-Aided Design*, Wiley, 2003.

ascii32 = ([0.0; 0.0],Array{Int64,1}[[1]])
ascii33 = ([1.75 1.75 2.0 2.0 1.5 1.5; 1.75 5.5 0.5 1.0 1.0 0.5],Array{Int64,1}[[1,2],[3,4],[4,5],[5,6],[6,3]])
ascii34 = ([1.0 2.0 2.0 2.0 1.5 1.5 2.0 3.0 3.0 3.0 2.5 2.5; 4.0 5.0 5.5 6.0 6.0 5.5 4.0 5.0 5.5 6.0 6.0 5.5],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,3],[7,8],[8,9],[9,10],[10,11],[11,12],[12,9]])
ascii35 = ([1.0 3.0 1.0 3.0 1.25 1.75 2.25 2.75; 2.5 2.5 3.5 3.5 1.75 4.0 1.75 4.0],Array{Int64,1}[[1,2],[3,4],[5,6],[7,8]])
ascii36 = ([0.0 1.0 3.0 4.0 4.0 3.0 1.0 0.0 0.0 1.0 3.0 4.0 2.0 2.0; 1.0 0.0 0.0 1.0 2.0 3.0 3.0 4.0 5.0 6.0 6.0 5.0 -0.5 6.5],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,9],[9,10],[10,11],[11,12],[13,14]])
ascii37 = ([2.5 2.5 2.0 2.0 2.5 2.5 2.0 2.0 1.5 3.5; 0.0 0.5 0.5 0.0 5.5 6.0 6.0 5.5 5.5 11.5],Array{Int64,1}[[1,2],[2,3],[3,4],[4,1],[5,6],[6,7],[7,8],[8,5],[9,10]])
ascii38 = ([4.0 3.0 1.0 0.0 0.0 1.0 2.0 3.0 3.0 2.0 1.0 1.0 3.0 4.0; 1.0 0.0 0.0 1.0 2.0 3.0 3.0 4.0 5.0 6.0 5.0 4.0 2.0 2.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,9],[9,10],[10,11],[11,12],[12,13],[13,14]])
ascii39 = ([1.0 2.0 2.0 2.0 1.5 1.5; 4.0 5.0 5.5 6.0 6.0 5.5],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,3]])
ascii40 = ([2.0 1.0 0.5 1.0 2.0; 0.0 1.0 3.0 5.0 6.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5]])
ascii41 = ([2.0 3.0 3.5 3.0 2.0; 0.0 1.0 3.0 5.0 6.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5]])
ascii42 = ([1.0 3.0 2.0 2.0 1.0 3.0 1.0 3.0; 3.0 3.0 2.0 4.0 2.0 4.0 4.0 2.0],Array{Int64,1}[[1,2],[3,4],[5,6],[7,8]])
ascii43 = ([1.0 3.0 2.0 2.0; 3.0 3.0 2.0 4.0],Array{Int64,1}[[1,2],[3,4]])
ascii44 = ([1.0 2.0 2.0 2.0 1.5 1.5; -1.0 0.0 0.5 1.0 1.0 0.5],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,3]])
ascii45 = ([1.0 3.0; 3.0 3.0],Array{Int64,1}[[1,2]])
ascii46 = ([2.0 2.0 1.5 1.5 2.0; 0.0 0.5 0.5 0.0 0.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5]])
ascii47 = ([1.0 3.0; 0.0 6.0],Array{Int64,1}[[1,2]])
ascii48 = ([4.0 3.0 1.0 0.0 0.0 1.0 3.0 4.0; 1.0 0.0 0.0 1.0 5.0 6.0 6.0 5.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,1],[4,8]])
ascii49 = ([0.0 2.0 2.0 0.0 4.0; 4.0 6.0 0.0 0.0 0.0],Array{Int64,1}[[1,2],[2,3],[4,5]])
ascii50 = ([0.0 0.0 1.0 3.0 4.0 4.0 0.0 4.0; 4.0 5.0 6.0 6.0 5.0 4.0 0.0 0.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8]])
ascii51 = ([0.0 4.0 2.0 4.0 4.0 3.0 1.0 0.0 0.0; 6.0 6.0 4.0 2.0 1.0 0.0 0.0 1.0 2.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,9]])
ascii52 = ([4.0 0.0 4.0 4.0; 1.0 1.0 6.0 0.0],Array{Int64,1}[[1,2],[2,3],[3,4]])
ascii53 = ([4.0 0.0 0.0 3.0 4.0 4.0 3.0 1.0 0.0 0.0; 6.0 6.0 4.0 4.0 3.0 1.0 0.0 0.0 1.0 2.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,9],[9,10]])
ascii54 = ([4.0 1.0 0.0 0.0 1.0 3.0 4.0 4.0 3.0 1.0 0.0; 6.0 6.0 5.0 1.0 0.0 0.0 1.0 3.0 4.0 4.0 3.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,9],[9,10],[10,11]])
ascii55 = ([0.0 0.0 4.0 0.0; 5.0 6.0 6.0 0.0],Array{Int64,1}[[1,2],[2,3],[3,4]])
ascii56 = ([1.0 3.0 4.0 4.0 3.0 1.0 0.0 0.0 4.0 4.0 3.0 1.0 0.0 0.0; 0.0 0.0 1.0 2.0 3.0 3.0 2.0 1.0 4.0 5.0 6.0 6.0 5.0 4.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,1],[6,5],[5,9],[9,10],[10,11],[11,12],[12,13],[13,14],[14,6]])
ascii57 = ([0.0 3.0 4.0 4.0 3.0 1.0 0.0 0.0 1.0 3.0 4.0; 0.0 0.0 1.0 5.0 6.0 6.0 5.0 3.0 2.0 2.0 3.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,9],[9,10],[10,11]])
ascii58 = ([2.0 2.0 1.5 1.5 2.0 2.0 1.5 1.5; 1.0 1.5 1.5 1.0 3.0 3.5 3.5 3.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,1],[5,6],[6,7],[7,8],[8,5]])
ascii59 = ([2.0 2.0 1.5 1.5 1.0 2.0 2.0 2.0 1.5 1.5; 3.0 3.5 3.5 3.0 -0.5 0.5 1.0 1.5 1.5 1.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,1],[5,6],[6,7],[7,8],[8,9],[9,10],[10,7]])
ascii60 = ([3.0 0.0 3.0; 6.0 3.0 0.0],Array{Int64,1}[[1,2],[2,3]])
ascii61 = ([1.0 3.0 1.0 3.0; 2.5 2.5 3.5 3.5],Array{Int64,1}[[1,2],[3,4]])
ascii62 = ([1.0 4.0 1.0; 6.0 3.0 0.0],Array{Int64,1}[[1,2],[2,3]])
ascii63 = ([2.0 2.0 1.5 1.5 1.75 1.75 3.0 3.0 2.0 1.0 0.0 0.0; 0.0 0.5 0.5 0.0 1.0 2.75 4.0 5.0 6.0 6.0 5.0 4.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,1],[5,6],[6,7],[7,8],[8,9],[9,10],[10,11],[11,12]])
ascii64 = ([4.0 1.0 0.0 0.0 1.0 3.0 4.0 4.0 2.0 1.0 2.0 3.0 2.0; 0.0 0.0 1.0 3.0 4.0 4.0 3.0 1.0 1.0 2.0 3.0 2.0 1.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,9],[9,10],[10,11],[11,12],[12,13]])
ascii65 = ([0.0 0.0 1.0 3.0 4.0 4.0 0.0 4.0; 0.0 5.0 6.0 6.0 5.0 0.0 2.0 2.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[7,8]])
ascii66 = ([0.0 0.0 3.0 4.0 4.0 3.0 4.0 4.0 3.0 0.0 0.0 3.0; 0.0 6.0 6.0 5.0 4.0 3.0 2.0 1.0 0.0 0.0 3.0 3.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,9],[9,10],[11,12]])
ascii67 = ([4.0 3.0 1.0 0.0 0.0 1.0 3.0 4.0; 1.0 0.0 0.0 1.0 5.0 6.0 6.0 5.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8]])
ascii68 = ([0.0 0.0 3.0 4.0 4.0 3.0 0.0; 0.0 6.0 6.0 5.0 1.0 0.0 0.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7]])
ascii69 = ([4.0 0.0 0.0 4.0 0.0 3.0; 0.0 0.0 6.0 6.0 3.0 3.0],Array{Int64,1}[[1,2],[2,3],[3,4],[5,6]])
ascii70 = ([0.0 0.0 4.0 0.0 3.0; 0.0 6.0 6.0 3.0 3.0],Array{Int64,1}[[1,2],[2,3],[4,5]])
ascii71 = ([2.0 4.0 4.0 3.0 1.0 0.0 0.0 1.0 3.0 4.0; 3.0 3.0 1.0 0.0 0.0 1.0 5.0 6.0 6.0 5.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,9],[9,10]])
ascii72 = ([0.0 0.0 4.0 4.0 0.0 4.0; 0.0 6.0 6.0 0.0 3.0 3.0],Array{Int64,1}[[1,2],[3,4],[5,6]])
ascii73 = ([2.0 2.0 1.0 3.0 1.0 3.0; 0.0 6.0 0.0 0.0 6.0 6.0],Array{Int64,1}[[1,2],[3,4],[5,6]])
ascii74 = ([0.0 1.0 2.0 3.0 3.0 2.0 4.0; 1.0 0.0 0.0 1.0 6.0 6.0 6.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[6,7]])
ascii75 = ([4.0 0.0 4.0 0.0 0.0; 6.0 3.0 0.0 0.0 6.0],Array{Int64,1}[[1,2],[2,3],[4,5]])
ascii76 = ([4.0 0.0 0.0; 0.0 0.0 6.0],Array{Int64,1}[[1,2],[2,3]])
ascii77 = ([0.0 0.0 2.0 4.0 4.0; 0.0 6.0 4.0 6.0 0.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5]])
ascii78 = ([0.0 0.0 4.0 4.0 4.0; 0.0 6.0 2.0 0.0 6.0],Array{Int64,1}[[1,2],[2,3],[4,5]])
ascii79 = ([4.0 3.0 1.0 0.0 0.0 1.0 3.0 4.0 4.0; 1.0 0.0 0.0 1.0 5.0 6.0 6.0 5.0 1.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,9]])
ascii80 = ([0.0 0.0 3.0 4.0 4.0 3.0 0.0; 0.0 6.0 6.0 5.0 3.0 2.0 2.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7]])
ascii81 = ([4.0 3.0 1.0 0.0 0.0 1.0 3.0 4.0 4.0 3.0 4.0; 1.0 0.0 0.0 1.0 5.0 6.0 6.0 5.0 1.0 1.0 0.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,9],[10,11]])
ascii82 = ([0.0 0.0 3.0 4.0 4.0 3.0 0.0 3.0 4.0; 0.0 6.0 6.0 5.0 3.0 2.0 2.0 2.0 0.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[8,9]])
ascii83 = ([0.0 1.0 3.0 4.0 4.0 3.0 1.0 0.0 0.0 1.0 3.0 4.0; 1.0 0.0 0.0 1.0 2.0 3.0 3.0 4.0 5.0 6.0 6.0 5.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,9],[9,10],[10,11],[11,12]])
ascii84 = ([2.0 2.0 0.0 4.0; 0.0 6.0 6.0 6.0],Array{Int64,1}[[1,2],[3,4]])
ascii85 = ([0.0 0.0 1.0 3.0 4.0 4.0; 6.0 1.0 0.0 0.0 1.0 6.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6]])
ascii86 = ([0.0 2.0 4.0; 6.0 0.0 6.0],Array{Int64,1}[[1,2],[2,3]])
ascii87 = ([0.0 0.0 1.0 2.0 3.0 4.0 4.0; 6.0 3.0 0.0 3.0 0.0 3.0 6.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7]])
ascii88 = ([0.0 4.0 0.0 4.0; 0.0 6.0 6.0 0.0],Array{Int64,1}[[1,2],[3,4]])
ascii89 = ([0.0 2.0 4.0 2.0 2.0; 6.0 2.0 6.0 2.0 0.0],Array{Int64,1}[[1,2],[2,3],[4,5]])
ascii90 = ([0.0 4.0 0.0 4.0; 6.0 6.0 0.0 0.0],Array{Int64,1}[[1,2],[2,3],[3,4]])
ascii91 = ([2.0 1.0 1.0 2.0; 0.0 0.0 6.0 6.0],Array{Int64,1}[[1,2],[2,3],[3,4]])
ascii92 = ([1.0 3.0; 6.0 0.0],Array{Int64,1}[[1,2]])
ascii93 = ([2.0 3.0 3.0 2.0; 0.0 0.0 6.0 6.0],Array{Int64,1}[[1,2],[2,3],[3,4]])
ascii94 = ([1.0 2.0 3.0; 5.0 6.0 5.0],Array{Int64,1}[[1,2],[2,3]])
ascii95 = ([1.0 4.0; 0.0 0.0],Array{Int64,1}[[1,2]])
ascii96 = ([2.0 2.0 3.0 2.5 2.5 2.0; 4.5 5.0 6.0 4.0 4.5 4.0],Array{Int64,1}[[1,2],[2,3],[4,5],[5,1],[1,6],[6,4]])
ascii97 = ([4.0 3.0 1.0 0.0 0.0 1.0 3.0 4.0 4.0 4.0; 1.0 0.0 0.0 1.0 2.0 3.0 3.0 2.0 0.0 3.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,1],[9,10]])
ascii98 = ([4.0 3.0 1.0 0.0 0.0 1.0 3.0 4.0 0.0 0.0 1.0; 1.0 0.0 0.0 1.0 2.0 3.0 3.0 2.0 0.0 5.0 5.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,1],[9,10],[10,11]])
ascii99 = ([4.0 3.0 1.0 0.0 0.0 1.0 3.0 4.0; 1.0 0.0 0.0 1.0 2.0 3.0 3.0 2.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8]])
ascii100 = ([4.0 3.0 1.0 0.0 0.0 1.0 3.0 4.0 4.0 4.0 3.0; 1.0 0.0 0.0 1.0 2.0 3.0 3.0 2.0 0.0 5.0 5.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,1],[9,10],[10,11]])
ascii101 = ([4.0 1.0 0.0 0.0 1.0 3.0 4.0 4.0 0.0; 0.0 0.0 1.0 2.0 3.0 3.0 2.0 1.0 1.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,9]])
ascii102 = ([4.0 4.0 3.0 2.0 1.0 1.0 0.0 2.0; 3.0 4.0 5.0 5.0 4.0 0.0 1.0 1.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[7,8]])
ascii103 = ([4.0 3.0 1.0 0.0 0.0 1.0 3.0 4.0 4.0 3.0 1.0 0.0; 1.0 0.0 0.0 1.0 2.0 3.0 3.0 2.0 0.0 -1.0 -1.0 0.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,1],[1,9],[9,10],[10,11],[11,12]])
ascii104 = ([4.0 4.0 3.0 1.0 0.0 0.0 0.0 1.0; 0.0 2.0 3.0 3.0 2.0 0.0 5.0 5.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[6,7],[7,8]])
ascii105 = ([1.0 3.0 1.0 3.0 2.0 2.0 2.25 2.25 1.75 1.75; 0.0 0.0 3.0 3.0 0.0 3.0 3.75 4.25 4.25 3.75],Array{Int64,1}[[1,2],[3,4],[5,6],[7,8],[8,9],[9,10],[10,7]])
ascii106 = ([1.0 3.0 2.0 2.0 1.0 0.0 2.25 2.25 1.75 1.75; 3.0 3.0 3.0 0.0 -1.0 0.0 3.75 4.25 4.25 3.75],Array{Int64,1}[[1,2],[3,4],[4,5],[5,6],[7,8],[8,9],[9,10],[10,7]])
ascii107 = ([0.0 1.0 1.0 0.0 4.0 2.0 1.0 3.0 4.0; 0.0 0.0 3.0 3.0 0.0 0.0 1.0 3.0 3.0],Array{Int64,1}[[1,2],[2,3],[3,4],[5,6],[6,7],[7,8],[8,9]])
ascii108 = ([2.0 2.0 1.0 1.0 3.0; 0.0 5.0 5.0 0.0 0.0],Array{Int64,1}[[1,2],[2,3],[4,5]])
ascii109 = ([4.0 4.0 2.0 2.0 2.0 0.0 0.0 0.0; 0.0 3.0 2.0 0.0 3.0 2.0 0.0 3.0],Array{Int64,1}[[1,2],[2,3],[4,5],[5,6],[7,8]])
ascii110 = ([3.0 3.0 1.0 1.0 1.0; 0.0 3.0 2.0 0.0 3.0],Array{Int64,1}[[1,2],[2,3],[4,5]])
ascii111 = ([4.0 3.0 1.0 0.0 0.0 1.0 3.0 4.0 4.0; 1.0 0.0 0.0 1.0 2.0 3.0 3.0 2.0 1.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,9]])
ascii112 = ([4.0 3.0 1.0 0.0 0.0 1.0 3.0 4.0 0.0 0.0; 1.0 0.0 0.0 1.0 2.0 3.0 3.0 2.0 3.0 -1.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,1],[9,10]])
ascii113 = ([4.0 3.0 1.0 0.0 0.0 1.0 3.0 4.0 4.0 4.0; 1.0 0.0 0.0 1.0 2.0 3.0 3.0 2.0 3.0 -1.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8],[8,1],[9,10]])
ascii114 = ([0.0 2.0 1.0 1.0 1.0 2.0 3.0 4.0; 0.0 0.0 0.0 3.0 2.0 3.0 3.0 2.0],Array{Int64,1}[[1,2],[3,4],[5,6],[6,7],[7,8]])
ascii115 = ([0.0 4.0 3.0 1.0 0.0 1.0 3.0 4.0; 0.0 0.0 1.0 1.0 2.0 3.0 3.0 2.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[7,8]])
ascii116 = ([1.0 3.0 2.0 2.0 2.0 3.0; 0.0 0.0 0.0 5.0 4.0 4.0],Array{Int64,1}[[1,2],[3,4],[5,6]])
ascii117 = ([0.0 1.0 1.0 2.0 3.0 4.0 4.0; 3.0 3.0 1.0 0.0 0.0 1.0 3.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7]])
ascii118 = ([0.0 1.0 3.0 4.0; 3.0 0.0 3.0 3.0],Array{Int64,1}[[1,2],[2,3],[3,4]])
ascii119 = ([0.0 0.0 1.0 2.0 3.0 4.0 4.0; 3.0 2.0 0.0 2.0 0.0 2.0 3.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7]])
ascii120 = ([0.0 1.0 4.0 1.0 4.0; 3.0 3.0 0.0 0.0 3.0],Array{Int64,1}[[1,2],[2,3],[4,5]])
ascii121 = ([0.0 1.0 2.5 0.0 1.0 4.0; 3.0 3.0 1.5 0.0 0.0 3.0],Array{Int64,1}[[1,2],[2,3],[4,5],[5,6]])
ascii122 = ([0.0 0.0 3.0 0.0 3.0 4.0; 2.0 3.0 3.0 0.0 0.0 1.0],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6]])
ascii123 = ([2.5 2.0 2.0 1.5 2.0 2.0 2.5; 6.5 6.0 3.5 3.0 2.5 0.0 -0.5],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7]])
ascii124 = ([2.0 2.0; 0.0 5.0],Array{Int64,1}[[1,2]])
ascii125 = ([1.5 2.0 2.0 2.5 2.0 2.0 1.5; 6.5 6.0 3.5 3.0 2.5 0.0 -0.5],Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7]])
ascii126 = ([1.0 1.75 2.75 3.5; 5.0 5.5 5.0 5.5],Array{Int64,1}[[1,2],[2,3],[3,4]])

"""
	hpcs::Array

"""
hpcs = [
	ascii32,ascii33,ascii34,ascii35,ascii36,ascii37,ascii38,ascii39,ascii40,ascii41,
	ascii42,ascii43,ascii44,ascii45,ascii46,ascii47,ascii48,ascii49,ascii50,ascii51,
	ascii52,ascii53,ascii54,ascii55,ascii56,ascii57,ascii58,ascii59,ascii60,ascii61,
	ascii62,ascii63,ascii64,ascii65,ascii66,ascii67,ascii68,ascii69,ascii70,ascii71,
	ascii72,ascii73,ascii74,ascii75,ascii76,ascii77,ascii78,ascii79,ascii80,ascii81,
	ascii82,ascii83,ascii84,ascii85,ascii86,ascii87,ascii88,ascii89,ascii90,ascii91,
	ascii92,ascii93,ascii94,ascii95,ascii96,ascii97,ascii98,ascii99,ascii100,ascii101,
	ascii102,ascii103,ascii104,ascii105,ascii106,ascii107,ascii108,ascii109,ascii110,ascii111,
	ascii112,ascii113,ascii114,ascii115,ascii116,ascii117,ascii118,ascii119,ascii120,ascii121,
	ascii122,ascii123,ascii124,ascii125,ascii126 ]

"""
ascii_LAR = DataStructures.OrderedDict{Int,LinearAlgebraicRepresentation.LAR}()

	ascii_LAR::{Int,LinearAlgebraicRepresentation.LAR}

*Ordered dictionary* of printable ASCII codes as one-dimensional *LAR models* in 2D.

`Key` is the `integer` ASCII code between 32 and 126.
Font design: *Geometric Programming for Computer-Aided Design*, Wiley, 2003.

# Example
```
julia> Plasm.ascii_LAR[46]
([2.0 2.0 … 1.5 2.0; 0.0 0.5 … 0.0 0.0], Array{Int64,1}[[1, 2], [2, 3], [3, 4], [4, 5]])

julia> Plasm.ascii_LAR[126]
([1.0 1.75 2.75 3.5; 5.0 5.5 5.0 5.5], Array{Int64,1}[[1, 2], [2, 3], [3, 4]])
```
"""
ascii_LAR = OrderedDict(zip(32:126,Plasm.hpcs))

# Attributes for `text` 2D graphics primitive.
# Reference to GKS ISO Graphics Standard (ISO/IEC 7942-4:1998)

const textalignment = "centre" #default value
const textangle = pi/4 #default value
const textwidth = 1.0 #default value
const textheight = 1.0 #default value
const textspacing = 0.25 #default value
const fontwidth = 4.0 #default value
const fontheight = 8.0 #default value
const fontspacing = 1.0 #default value

"""
	charpols(mystring::String)::Array{LAR,1}

Return the `array` of `LAR` models, for chars in `mystring`.
"""
function charpols(mystring)
    return [ascii_LAR[code] for code in [Int(char) for char in mystring]]
end

"""
	charseq(mystring::String)::Array{Char,1}

# Example
```
julia> Plasm.charseq("PLaSM")
5-element Array{Char,1}:
 'P'
 'L'
 'a'
 'S'
 'M'
```
"""
function charseq(mystring)
	return [char for char in mystring]
end

"""
	text(mystring::String)::LAR

Compute the one-dim *LAR model* drawing the contents of `mystring`

# Example
```
julia> model = Plasm.text("PLaSM")
# output
([0.0 0.0 3.0 4.0 4.0 3.0 0.0 9.0 5.0 5.0 14.0 13.0 11.0 10.0
10.0 11.0 13.0 14.0 14.0 14.0 15.0 16.0 18.0 19.0 19.0 18.0 16.0 15.0 15.0 16.0 18.0
19.0 20.0 20.0 22.0 24.0 24.0; 0.0 6.0 6.0 5.0 3.0 2.0 2.0 0.0 0.0 6.0 1.0 0.0 0.0 1.0
2.0 3.0 3.0 2.0 0.0 3.0 1.0 0.0 0.0 1.0 2.0 3.0 3.0 4.0 5.0 6.0 6.0 5.0 0.0 6.0 4.0
6.0 0.0],
Array{Int64,1}[[1,2],[2,3],[3,4],[4,5],[5,6],[6,7],[8,9],[9,10],[11,12],[12,13],[13,14],
[14,15],[15,16],[16,17],[17,18],[18,11],[19,20],[21,22],[22,23],[23,24],[24,25],[25,26],
[26,27],[27,28],[28,29],[29,30],[30,31],[31,32],[33,34],[34,35],[35,36],[36,37]])

julia> Plasm.view(model)
```
"""
function text(mystring)
	out = comp([ L.struct2lar, L.Struct, cat, distr,
			cons([ charpols, k(L.t(fontspacing+fontwidth,0)) ]),charseq ])(mystring)
	return out
end

"""
	a2a(mat::Matrix)(models::Array{LAR,1})::Struct.body

Local function, service to `textWithAttributes` implementation.
"""
function a2a(mat)
	function a2a0(models)
		assembly = []
		for model in models
			push!( assembly, LinearAlgebraicRepresentation.Struct([ mat,model ]) )
		end
		assembly
	end
	return a2a0
end

"""
	translate(c::Number)(lar::LAR)::LAR

Local function, service to `textWithAttributes` implementation.
"""
function translate(c)
	function translate0(lar)
		xs = lar[1][1,:]
		width = maximum(xs) - minimum(xs)
		apply(LinearAlgebraicRepresentation.t(width/c,0))(lar)
	end
	return translate0
end

"""
	align(textalignment="centre"::String)(model::LAR)::LAR

Local function, service to `textWithAttributes` implementation.
"""
function align(textalignment="centre")
	function align1(model)
		if ( textalignment == "centre" ) out=translate(-2)(model)
		elseif ( textalignment == "right" ) out=translate(-1)(model)
		elseif ( textalignment == "left" ) out=model
		end
		return out
	end
end

"""
 	cat(args)

Redefined locally, as service to `textWithAttributes` implementation.
"""
function cat(args)
	return reduce( (x,y) -> append!(x,y), args; init=[] )
end

"""
	textWithAttributes(textalignment='centre', textangle=0,
		textwidth=1.0, textheight=2.0, textspacing=0.25)(strand::String)::LAR

Partial implementation of the GKS's graphics primitive `text`.

# Example

```
Plasm.view(Plasm.textWithAttributes("left", pi/4)("PLaSM"))
```
"""
function textWithAttributes(textalignment="centre", textangle=0,
							textwidth=1.0, textheight=2.0, textspacing=0.25)
	function textWithAttributes(strand)
		id = x->x
		mat = LinearAlgebraicRepresentation.s(textwidth/fontwidth,textheight/fontheight)
		comp([
		   apply(LinearAlgebraicRepresentation.r(textangle)),
		   align(textalignment),
		   L.struct2lar,
		   L.Struct,
		   cat,
		   distr,
		   cons([ a2a(mat) ∘ charpols,
				k(LinearAlgebraicRepresentation.t(textwidth+textspacing,0)) ]),
		   charseq ])(strand)
	end
end

"""
	embed(n::Int)(model::LAR)::LAR

Embed a `LAR` `model` of dimension ``d`` in a space ``R^{d+n}``.
The embedding is done by adding ``d`` zero coordinates to each vertex.

# Example

```
julia> square = LinearAlgebraicRepresentation.cuboid([1,1])
([0.0 0.0 1.0 1.0; 0.0 1.0 0.0 1.0], Array{Int64,1}[[1, 2, 3, 4]])

julia> embed(1)(square)
([0.0 0.0 1.0 1.0; 0.0 1.0 0.0 1.0; 0.0 0.0 0.0 0.0], Array{Int64,1}[[1, 2, 3, 4]])
```
"""
function embed(n)
	function embed0( larmodel )
		return [larmodel[1]; zeros(n,size(larmodel[1],2))], larmodel[2]
	end
	return embed0
end

"""
	numbering(model::LARmodel)

Display a *wireframe* of `model` cells, of any dimension, with their *ordinal number*.
Different `colors` and size are used for the various dimensional cells.

# Examples

```
model = LinearAlgebraicRepresentation.cuboidGrid([3,4,2], true)
Plasm.view(Plasm.numbering()(model))

model = LinearAlgebraicRepresentation.cuboidGrid([10,10], true)
Plasm.view(Plasm.numbering(1.5)(model))
```
"""
function numbering(numberSizeScaling=1)
	function numbering0(model)
		V,cells = model
		if size(V,1)==2
			V = Plasm.embed(1)(model)[1]
		end
		wireframe = Plasm.lar2hpc(V,cells[2])
		ns = numberSizeScaling
		gcode = Plasm.textWithAttributes("centre", 0, 0.1ns, 0.2ns, 0.025ns)
		scene = [wireframe]
		for (h,skel) in enumerate(cells)
			colors = [p.GREEN,p.YELLOW,p.CYAN,p.ORANGE]
			nums = []
			for (k,cell) in enumerate(skel)
				center = sum([V[:,v] for v in cell])/length(cell)
				code = Plasm.embed(1)( gcode(string(k)) )
				scaling = (0.6+0.1h,0.6+0.1h,1)
				push!(nums, LinearAlgebraicRepresentation.struct2lar( LinearAlgebraicRepresentation.Struct([
					LinearAlgebraicRepresentation.t(center...), LinearAlgebraicRepresentation.s(scaling...), code ]) ))
			end
			hpc = Plasm.lar2hpc(nums)
			push!( scene, p.COLOR(colors[h])(hpc) )
		end
		p.STRUCT( scene )
	end
	return numbering0
end


# Accumulate statistics (mean, variance etc.) at different times

# Stores statistics at one time

using StreamStats

# Stores statistics at given times:
type StatisticsObject
    data::Dict{Float64, StreamStats.Moments}
end

StatisticsObject() = StatisticsObject(Dict())

# Adds the data x at time t
function add_data!(o::StatisticsObject, t, x)
    if haskey(o.data, t)
        #add_data!(o.data[t], x)
        update!(o.data[t], x)
    else
        o.data[t] = StreamStats.Moments() #StatisticsContainer(x)
        update!(o.data[t], x)
    end
end

function get_statistics(o::StatisticsObject)
    times = sort(collect(keys(o.data)))

    [tuple(t, state(o.data[t])..., o.data[t].n) for t in times]

end


module TestMLBaseWrapper

using Test
using TSML
using DataFrames

@testset "MLBase transformers" begin

  @testset "StandardScaler transforms features" begin
    features = [
      5 10;
      -5 0;
      0 5;
    ] |> DataFrame
    labels = [
      "x";
      "y";
      "z";
    ]
    expected_transformed = [
      1.0 1.0;
      -1.0 -1.0;
      0.0 0.0;
    ] |> DataFrame
    standard_scaler = StandardScaler()
    fit!(standard_scaler, features, labels)
    transformed = transform!(standard_scaler, features)
    @test (transformed .== expected_transformed) |> Matrix |> sum == 6
  end

end

end # module


"""
# References
- Reference model (e.g., xd, vd, ad, ad_dot, ad_ddot)
[1] S. J. Su, Y. Y. Zhu, H. R. Wang, and C. Yun, “A Method to Construct a Reference Model for Model Reference Adaptive Control,” Adv. Mech. Eng., vol. 11, no. 11, pp. 1–9, 2019.
# Notes
d: degree of the reference model
auto_diff: auto_diff = true for tracking time-varying command; otherwise for set-point regulation
x_cmd_func(t): function of time
"""
struct ReferenceModel <: AbstractEnv
    d::Int  # degree
    Ks::AbstractArray
    auto_diff::Bool
    x_cmd_func::Union{Function, Nothing}
end
function ReferenceModel(d::Int; x_cmd_func=nothing)
    @assert d >= 0
    Ks = []
    if d == 4
        # [1]
        push!(Ks, Diagonal(1.0*ones(3)))
        push!(Ks, Diagonal(3.4*ones(3)))
        push!(Ks, Diagonal(5.4*ones(3)))
        push!(Ks, Diagonal(4.9*ones(3)))
        push!(Ks, Diagonal(2.7*ones(3)))
    else
        error("Assign values of matrix `Kx` manually")
    end
    auto_diff = x_cmd_func == nothing ? false : true
    ReferenceModel(d, Ks, auto_diff, x_cmd_func)
end

"""
`x_i` denotes `i`th (time) derivative of `x`
"""
function State(env::ReferenceModel)
    @unpack d = env
    return function (x0)
        xs_dict = Dict(:x_0 => x0)
        for i in 1:d
            xs_dict[Symbol(:x_, i)] = zeros(size(x0))  # e.g., xs_dict[:x1] = zeros(3)
        end
        ComponentArray((; zip(keys(xs_dict), values(xs_dict))...))
    end
end

function Dynamics!(env::ReferenceModel)
    @unpack d, Ks, auto_diff, x_cmd_func = env
    # derivatives for auto_diff
    funcs = nothing
    if auto_diff
        _funcs = Function[x_cmd_func]
        for i in 1:d+1
            push!(_funcs, (t) -> ForwardDiff.derivative(_funcs[i], t))
        end
    end
    funcs(t) = [_func(t) for _func in _funcs]
    @Loggable function dynamics!(dX, X, p, t; x_cmd=nothing)
        @onlylog state = X
        xs = [getproperty(X, Symbol(:x_, 0))]  # x_0, x_1, ...
        for i in 0:d-1
            _x_next = getproperty(X, Symbol(:x_, i+1))
            setproperty!(dX, Symbol(:x_, i), _x_next)  # e.g., dX.x_0 = x_1
            push!(xs, _x_next)
        end
        dx_d = nothing
        if auto_diff
            if x_cmd != nothing
                error("Do not provide a manual command for mode auto_diff=$(auto_diff)")
            end
            # dx_d = -sum(Ks .* xs) + sum([Ks..., I] .* funcs(t))
            desired_cmds = funcs(t)
            dx_d = -sum(Ks .* xs) + sum([Ks..., I] .* desired_cmds)
            @log x_cmd = desired_cmds[1]
        else
            if x_cmd == nothing
                error("Provide a manual command for mode auto_diff=$(auto_diff)")
            end
            dx_d = -sum(Ks .* xs) + Ks[1]*x_cmd
            @log x_cmd
        end
        setproperty!(dX, Symbol(:x_, d), dx_d)
    end
end


function trim_arr(arr; trim_at_begin, trim_at_end)
    if (trim_at_begin + trim_at_end) > 1
        error("you are trying to remove more than you have")
    elseif (trim_at_begin < 0) | (trim_at_end < 0)
        error("you are trying to remove less than zero")
    end
    l = length(arr)
    trim_b = round(Int, float(l*trim_at_begin), RoundDown)
    trim_e = round(Int, float(l*trim_at_end), RoundDown)
    return arr[begin+trim_b:end-trim_e]
end

rms(ar) = sqrt(mean(ar.^2))

######

function arr_amplitude(;testarr, trim_at_begin=0, trim_at_end=0)
    ar = trim_arr(testarr; trim_at_begin, trim_at_end)
    return (;y=(maximum(ar)-minimum(ar)))
end

function arr_length(;testarr, trim_at_begin=0, trim_at_end=0)
    ar = trim_arr(testarr; trim_at_begin, trim_at_end)
    return (;y=length(ar))
end

function arr_mean(;testarr, trim_at_begin=0, trim_at_end=0)
    ar = trim_arr(testarr; trim_at_begin, trim_at_end)
    return (;y=mean(ar))
end

function arr_median(;testarr, trim_at_begin=0, trim_at_end=0)
    ar = trim_arr(testarr; trim_at_begin, trim_at_end)
    return (;y=median(ar))
end

function arr_std(;testarr, trim_at_begin=0, trim_at_end=0)
    ar = trim_arr(testarr; trim_at_begin, trim_at_end)
    return (;y=std(ar))
end

function arr_rms(;testarr, trim_at_begin=0, trim_at_end=0)
    ar = trim_arr(testarr; trim_at_begin, trim_at_end)
    return (;y=rms(ar))
end

function vec_amplify(;testarr, trim_at_begin=0, trim_at_end=0, p=1)
    vec = trim_arr(testarr; trim_at_begin, trim_at_end).*p
    return (;y=rms(vec), bigarrs=(;vec))
end

function vec_attenuate(;testarr, trim_at_begin=0, trim_at_end=0, p=1)
    vec = trim_arr(testarr; trim_at_begin, trim_at_end)./p
    return (;y=rms(vec), bigarrs=(;vec))
end

function vec_power(;testarr, trim_at_begin=0, trim_at_end=0, p=1)
    vec = abs.(trim_arr(testarr; trim_at_begin, trim_at_end)).^p
    return (;y=rms(vec), bigarrs=(;vec))
end


using QuantumCumulants
using Test
using OrdinaryDiffEq
using ModelingToolkit

@testset "scaling" begin

order = 2
N_c = 1 #number of clusters
@cnumbers Δc κ Γ2 Γ3 Γ23 η ν3 ν2
Δ2 = [Parameter(Symbol(:Δ2_, i)) for i=1:N_c]
Δ3 = [Parameter(Symbol(:Δ3_, i)) for i=1:N_c]
Ω3 = [Parameter(Symbol(:Ω3_, i)) for i=1:N_c]
g = [Parameter(Symbol(:g_, i)) for i=1:N_c]
N = [Parameter(Symbol(:N_, i)) for i=1:N_c]

# Define hilbert space
hf = FockSpace(:cavity)
ha = [NLevelSpace(Symbol(:atoms, j),3) for j=1:N_c]
ha_c = [ClusterSpace(ha[j],N[j],order) for j=1:N_c]
h = ⊗(hf, ha_c...)
# Define the fundamental operators
@qnumbers a::Destroy(h,1)
S(i,j,c) = Transition(h,Symbol(:σ, c),i, j, 1+c) #c=cluster

# Hamiltonian
H = Δc*a'*a + sum(Δ2[c]*sum(S(2,2,c)) for c=1:N_c) + sum(Δ3[c]*sum(S(3,3,c)) for c=1:N_c) +
    sum(Ω3[c]*(sum(S(3,1,c)) + sum(S(1,3,c))) for c=1:N_c) + sum(g[c]*(a'*sum(S(1,2,c)) + a*sum(S(2,1,c))) for c=1:N_c)
# Collapse operators
J = [a,[S(1,2,c) for c=1:N_c]..., [S(1,3,c) for c=1:N_c]..., [S(2,3,c) for c=1:N_c]..., [S(3,3,c) for c=1:N_c]..., [S(2,2,c) for c=1:N_c]..., a'a]
rates = [κ,[Γ2 for i=1:N_c]...,[Γ3 for i=1:N_c]...,[Γ23 for i=1:N_c]...,[ν3 for i=1:N_c]...,[ν2 for i=1:N_c]...,η]

# Derive equation for average photon number
ops = [a'a, S(2,2,1)[1], a'*S(1,2,1)[1]]
he = meanfield(ops,H,J;rates=rates,order=2)

# Custom filter function -- include only phase-invariant terms
ϕ(x) = 0
ϕ(x::Destroy) = -1
ϕ(x::Create) = 1
function ϕ(t::Transition)
    if (t.i==1 && t.j==2) || (t.i==3 && t.j==2) || (t.i==:g && t.j==:e)
        -1
    elseif (t.i==2 && t.j==1) || (t.i==2 && t.j==3) || (t.i==:e && t.j==:g)
        1
    else
        0
    end
end
ϕ(avg::Average) = ϕ(avg.arguments[1])
function ϕ(t::QuantumCumulants.QMul)
    p = 0
    for arg in t.args_nc
        p += ϕ(arg)
    end
    return p
end
phase_invariant(x) = iszero(ϕ(x))

he_avg = cumulant_expansion(he,2)
he_scale = complete(he_avg;filter_func=phase_invariant, order=order, multithread=false)
@test length(he_scale) == 9

ps = [Δc; κ; Γ2; Γ3; Γ23; η; ν3; ν2; Δ2; Δ3; Ω3; g; N]
@named sys = ODESystem(he_scale)
u0 = zeros(ComplexF64, length(he_scale))

N0 = 1000/N_c
N_ = N0*[1.0 for c=1:N_c]

p0 = ps .=> [ones(length(ps)-1); N_]
prob1 = ODEProblem(sys,u0,(0.0,1.0),p0)
sol1 = solve(prob1, Tsit5(), reltol=1e-12, abstol=1e-12)

@test sol1.u[end][1] ≈ 0.0758608728203

avg = average(a*S(2,1,1)[2])
# TODO: @test sol1[avg], sol1, he_scale) == map(conj, get_solution(average(a'*S(1,2,1)[1]), sol1, he_scale)) == conj.(getindex.(sol1.u, 3))

## Two-level laser
M = 2
@cnumbers N
hf = FockSpace(:cavity)
ha = NLevelSpace(:atom, (:g,:e))
hc = ClusterSpace(ha, N, M)
h = tensor(hf, hc)

@qnumbers a::Destroy(h)
σ(i,j) = Transition(h,:σ,i,j)

@cnumbers Δ g κ γ ν

H = Δ*a'*a + g*sum(a'*σ(:g,:e)[i] + a*σ(:e,:g)[i] for i=1:M)
J = [a;[σ(:g,:e)[i] for i=1:M];[σ(:e,:g)[i] for i=1:M]]
rates = [κ; [γ for i=1:M]; [ν for i=1:M]]

he = meanfield(a'*a, H, J; rates=rates, order=2)

# Complete
he_scaled = complete(he;filter_func=phase_invariant)

names = he_scaled.names
avg = average(σ(:e,:g)[1]*σ(:e,:e)[2])
@test isequal(average(σ(:e,:e)[1]*σ(:e,:g)[2]), QuantumCumulants.substitute_redundants(avg,[QuantumCumulants.ClusterAon(2,1),QuantumCumulants.ClusterAon(2,2)],names))

@test QuantumCumulants.lt_reference_order(σ(:e,:g)[1],σ(:g,:e)[2])
@test !QuantumCumulants.lt_reference_order(σ(:g,:e)[1],σ(:e,:g)[2])

he_avg = cumulant_expansion(he_scaled,2)
@test isempty(find_missing(he_avg))

ps = (Δ, g, γ, κ, ν, N)
@named sys = ODESystem(he_avg)
p0 = ps .=> (0, 1.5, 0.25, 1, 4, 7)
u0 = zeros(ComplexF64, length(he_scaled))
prob = ODEProblem(sys, u0, (0.0, 50.0), p0)
sol = solve(prob, RK4(), abstol=1e-10, reltol=1e-10)

@test sol.u[end][1] ≈ 12.601868534

# Spectrum
corr = CorrelationFunction(a',a,he_avg;steady_state=true,filter_func=phase_invariant)
Spec = Spectrum(corr,ps)
s = Spec(range(-π, π, length=301), sol.u[end], getindex.(p0,2))
@test all(s .>= 0.0)

## Some abstract tests
M = 4
@cnumbers N
hc = FockSpace(:cavity)
hvib = FockSpace(:mode)
hcluster = ClusterSpace(hvib,N,M)
h = ⊗(hc, hcluster)
a = Destroy(h,:a,1)
b = Destroy(h,:b,2)

names = [:a,[Symbol(:b_,i) for i=1:M]]
scale_aons = [QuantumCumulants.ClusterAon(2,i) for i=1:M]

avg = average(a*b[1]*b[2]*b[2])
avg_sub = QuantumCumulants.substitute_redundants(avg, scale_aons, names)
@test isequal(average(a*b[1]*b[1]*b[2]), avg_sub)

avg = average(b[1]'*b[2]'*b[2])
avg_sub = QuantumCumulants.substitute_redundants(avg, scale_aons, names)
@test isequal(avg_sub, average(b[1]'*b[1]*b[2]'))

avg = average(b[1]*b[1]*b[2]'*b[2])
avg_sub = QuantumCumulants.substitute_redundants(avg, scale_aons, names)
@test isequal(avg_sub, average(b[1]'*b[1]*b[2]*b[2]))

avg = average(b[1]*b[2]'*b[3]*b[4]')
avg_sub = QuantumCumulants.substitute_redundants(avg, scale_aons, names)
@test isequal(avg_sub, average(b[1]'*b[2]'*b[3]*b[4]))

# Test Holstein
M = 2
hc = FockSpace(:cavity)
hvib = FockSpace(:mode)
hcluster = ClusterSpace(hvib,N,M)
h = ⊗(hc, hcluster)
@cnumbers G Δ κ γ Ω
a = Destroy(h,:a,1)
b = Destroy(h,:b,2)

H = Δ*a'*a + G*sum(b[i] + b[i]' for i=1:M)*a'*a + Ω*(a+a')
J = [a,b]
rates = [κ,γ]
ops = [a,a'*a,a*a,b[1],a*b[1],a'*b[1],b[1]'*b[1],b[1]*b[1],b[1]'*b[2],b[1]*b[2]]
he = meanfield(ops,H,J;rates=rates)

he_avg = cumulant_expansion(he,2)
@test isempty(find_missing(he_avg))

ps = (G,Δ,κ,γ,Ω,N)
@named sys = ODESystem(he_avg)

u0 = zeros(ComplexF64, length(he_avg))
p0 = ps .=> ones(length(ps))
prob = ODEProblem(sys, u0, (0.0,1.0), p0)
sol = solve(prob, Tsit5())

# Test molecule
M = 2 # Order
# Prameters
@cnumbers λ ν Γ η Δ γ N
# Hilbert space
h_in = NLevelSpace(:internal, 2)
hv = FockSpace(:vib)
hc = ClusterSpace(hv, N, M)
h = ⊗(h_in, hc)
# Operators
σ(i,j) = Transition(h, :σ, i, j)
b = Destroy(h, :b)
# Hamiltonian
H0 = Δ*σ(2,2) + ν*sum(b_'*b_ for b_ in b)
H_holstein = -1*λ*sum((b_' + b_ for b_ in b))*σ(2,2)
Hl = η*(σ(1,2) + σ(2,1))
H = H0 + H_holstein + Hl
# Jumps
J = [σ(1,2), b]
rates = [γ,Γ]
# Equations
ops = [σ(2,2),σ(1,2),b[1],σ(1,2)*b[1],σ(2,1)*b[1],σ(2,2)*b[1],b[1]'*b[1],b[1]*b[1],b[1]'*b[2],b[1]*b[2]]
he = meanfield(ops, H, J; rates=rates)
he_avg = cumulant_expansion(he,2)
@test isempty(find_missing(he_avg))
ps = (Δ,η,γ,λ,ν,Γ,N)
# Generate function
@named sys = ODESystem(he_avg)
p0 = ps .=> [ones(length(ps)-1); 4]
u0 = zeros(ComplexF64,length(he_avg))
prob1 = ODEProblem(sys,u0,(0.0,1.0),p0)
sol1 = solve(prob1,Tsit5(),abstol=1e-12,reltol=1e-12)
bdb1 = sol1[b[1]'*b[1]][end]
σ22_1 = sol1[σ(2,2)][end]
σ12_1 = sol1[σ(1,2)][end]

## Two clusters
N_c = 2
N = cnumbers([Symbol(:N_, i) for i=1:N_c]...)
M = 2
hf = FockSpace(:cavity)
ha = NLevelSpace(:atom, (:g,:e))
hc = [ClusterSpace(ha, N[i], M) for i=1:N_c]
h = tensor(hf, hc...)

@qnumbers a::Destroy(h)
σ(i,j,c) = Transition(h,Symbol(:σ_, c),i,j,c+1)

@cnumbers κ
ν = cnumbers([Symbol(:ν_, c) for c=1:N_c]...)
γ = cnumbers([Symbol(:γ_, c) for c=1:N_c]...)
Δ = cnumbers([Symbol(:Δ_, c) for c=1:N_c]...)
g = cnumbers([Symbol(:g_, c) for c=1:N_c]...)

H = sum(Δ[c]*σ(:e,:e,c)[k] for c=1:N_c, k=1:M) + sum(g[c]*(a'*σ(:g,:e,c)[i] + a*σ(:e,:g,c)[i]) for i=1:M, c=1:N_c)
J = [a;[σ(:g,:e,c) for c=1:N_c];[σ(:e,:g,c) for c=1:N_c]]
rates = [κ,γ...,ν...]

ops = [a'*a]
he = meanfield(ops, H, J; rates=rates, order=2)

# Scale
he_scaled = complete(he; filter_func=phase_invariant)

@test isempty(find_missing(he_scaled))

ps = (κ, Δ..., g..., γ..., ν..., N...)
@named sys = ODESystem(he_scaled)
if N_c==2
    p0 = ps .=> (1, [0 for i=1:N_c]..., [1.5 for i=1:N_c]..., [0.25 for i=1:N_c]..., [4 for i=1:N_c]..., 4, 3)
elseif N_c==3
    p0 = ps .=> (1, [0 for i=1:N_c]..., [1.5 for i=1:N_c]..., [0.25 for i=1:N_c]..., [4 for i=1:N_c]..., 2, 3, 2)
end
u0 = zeros(ComplexF64, length(he_scaled))
prob = ODEProblem(sys, u0, (0.0, 50.0), p0)
sol = solve(prob, RK4(), abstol=1e-10, reltol=1e-10)

@test sol.u[end][1] ≈ 12.601868534

### 6-level laser

order = 2
# Define parameters
@cnumbers Δc κ g Γ12 Γ13 Γ24 Γ34 Γ54 Γ16 Γ64 Δ3 Δ4 Δ5 Δ6 Ω13 Ω34 Ω54 Ω64 η ν12 ν13 ν34 ν54 ν64 N

# Define hilbert space
hf = FockSpace(:cavity)
ha = NLevelSpace(:atom,6)
ha_c = ClusterSpace(ha,N,order)
h = ⊗(hf, ha_c)
# Define the fundamental operators
a = Destroy(h,:a,1)
σ(i,j) = Transition(h,:σ,i,j)
@test length(σ(6,6)) == order

# Hamiltonian
H = Δc*a'a + Δ3*sum(σ(3,3)) + Δ4*sum(σ(4,4)) + Δ5*sum(σ(5,5)) + Δ6*sum(σ(6,6)) +
    Ω13*(sum(σ(3,1)) + sum(σ(1,3))) + Ω34*(sum(σ(3,4)) + sum(σ(4,3))) + Ω54*(sum(σ(5,4)) + sum(σ(4,5))) + Ω64*(sum(σ(6,4)) + sum(σ(4,6))) +
    g*(a'*sum(σ(1,2)) + a*sum(σ(2,1)))
# Collapse operators
J = [a, σ(1,2), σ(1,3), σ(2,4), σ(3,4), σ(5,4), σ(6,4), σ(1,6), σ(2,2), σ(3,3)+σ(4,4)+σ(5,5)+σ(6,6), σ(4,4)+σ(5,5)+σ(6,6), σ(5,5), σ(6,6), a'a]
rates = [κ, Γ12, Γ13, Γ24, Γ34, Γ54, Γ64, Γ16, ν12, ν13, ν34, ν54, ν64, η]

# Derive equation for average photon number
ops = [a'a, σ(2,2)[1], σ(3,3)[1], σ(4,4)[1], σ(5,5)[1], σ(6,6)[1]]
eqs_ops = meanfield(ops,H,J;rates=rates, order=order, multithread=true)

@test length(eqs_ops) == length(ops)

### 4th order 2-level laser
order = 4
# Define parameters
@cnumbers δA ΩA wA νA γ κ δc NA
# Define hilbert space
hf = FockSpace(:cavity)
haA = NLevelSpace(:atomA,2)
heA = ClusterSpace(haA, NA, order) #atom ensemble 1
h = ⊗(hf, heA)
a = Destroy(h,:a,1)
σA(i,j) = Transition(h,:σA,i,j,2)
# Hamiltonian
H = δc*a'a + δA*sum(σA(2,2)) + ΩA/2*(a'sum(σA(1,2)) + a*sum(σA(2,1)))
# dissipative processes
J = [a, σA(1,2), σA(2,1), σA(2,2)]
rates = [κ, γ, wA, γ, νA]
# Derive equation for average photon number
ops = [a'a, σA(2,2)[1]]
he_ops = meanfield(ops,H,J;rates=rates, multithread=true, order=order)

he_scale = complete(he_ops; filter_func=phase_invariant, order=order, multithread=true)
@test length(he_scale) == 18
@test isempty(find_missing(he_scale))

### 4th order synchronization

order = 4
# Define parameters
@cnumbers δA δB ΩA ΩB wA wB νA νB γ κ δc NA NB
# Define hilbert space
hf = FockSpace(:cavity)
haA = NLevelSpace(:atomA,2)
haB = NLevelSpace(:atomB,2)
heA = ClusterSpace(haA, NA, order) #atom ensemble 1
heB = ClusterSpace(haB, NB, order) #atom ensemble 2
h = ⊗(hf, heA, heB)
# Define the fundamental operators
a = Destroy(h,:a,1)
σA(i,j) = Transition(h,:σA,i,j,2)
σB(i,j) = Transition(h,:σB,i,j,3)
# Hamiltonian
H = δc*a'a + δA*sum(σA(2,2)) + δB*sum(σB(2,2)) +
    ΩA/2*(a'sum(σA(1,2)) + a*sum(σA(2,1))) + ΩB/2*(a'sum(σB(1,2)) + a*sum(σB(2,1)))
# dissipative processes
J = [a, σA(1,2), σA(2,1), σB(1,2), σB(2,1), σA(2,2), σB(2,2)]
rates = [κ, γ, wA, γ, wB, νA, νB]
# Derive equation for average photon number
ops = [a'a, σA(2,2)[1], σB(2,2)[1]]
he_ops = meanfield(ops,H,J;rates=rates, multithread=true, order=order)

he_scale = complete(he_ops; filter_func=phase_invariant, order=order, multithread=true)
@test length(he_scale) == 66
@test isempty(find_missing(he_scale))

@named sys = ODESystem(he_scale)
u0 = zeros(ComplexF64, length(he_scale))
ps = [δA, δB, ΩA, ΩB, wA, wB, νA, νB, γ, κ, δc, NA, NB]
p0 = ps.=>[1.0 + i/20 for i=1:length(ps)]
prob = ODEProblem(sys,u0,(0.0,1.0),p0)
sol = solve(prob,RK4())

uend = copy(sol.u[end])
@test length(unique(uend)) == 66
uend_f = filter(x->imag(x) != 0, uend)
@test 2*length(uend_f) == length(unique([uend_f; adjoint.(uend_f)]))

end # testset


using YaoExtensions, YaoZX, Yao
using Test

@testset "qft" begin
    c = qft_circuit(4)
    zxd = ZXDiagram(4)
    push_gate!(zxd, c)
    circ = toqbir(zxd)
    @test operator_fidelity(circ, c) ≈ 1
end


# ## [Aeroelastic Analysis of a Typical Section](@id section-stability)
#
# In this example, we perform a two-dimensional aeroelastic analysis for a two degree of
# freedom typical section model.
#
# ![](../assets/section-drawing.svg)
#
#-
#md # !!! tip
#md #     This example is also available as a Jupyter notebook:
#md #     [`section-stability.ipynb`](@__NBVIEWER_ROOT_URL__/examples/section-stability.ipynb).
#-
#
# The equations of motion for this model are
# ```math
# m \left(\ddot{h}+b x_\theta \ddot{\theta} \right) + k_h h = -L \\
# I_\theta \ddot{\theta} + m b x_\theta \ddot{h} + k_\theta = M
# ```
# where ``b`` is the semichord length, ``k_h`` is the linear spring constant, ``k_\theta``
# is the torsional spring constant, ``m`` is the mass per unit span, ``x_\theta`` is the
# distance to the center of mass from the reference point, ``I_θ`` is the moment of
# inertia about the reference point, ``L`` is the lift per unit span, and ``M`` is the
# moment per unit span about the reference point.  To learn more about how this model is
# implemented in this package, see [`Section`](@ref).
#
# This package has a number of pre-implemented aerodynamic models which may be used to
# model the aerodynamics of the typical section model.  These models include a steady-state
# thin airfoil theory model (see [`Steady`](@ref)), a quasi-steady thin airfoil theory
# model (see [`QuasiSteady`](@ref)), an unsteady aerodynamic model based on Wagner's
# function (see [`Wagner`](@ref)), and Peters' finite state aerodynamic model
# (see [`Peters`](@ref)).  We will perform aeroelastic analyses using each of these models
# and compare the results.
#
# The non-dimensional parameters we will use are
# ```math
# a = -1/5 \quad e = -1/10  \\
# r^2 = \frac{I_P}{m b^2} \quad \sigma = \frac{\omega_h}{\omega_\theta} \\
# \mu = \frac{m}{\rho_\infty \pi b^2} \quad V = \frac{U}{b \omega_\theta}
# ```
# where ``a`` is the normalized distance from the semichord to the reference point, ``e``
# is the normalized distance from the semichord to the center of mass, and ``\omega_h`` and
# ``\omega_\theta`` are the uncoupled natural frequencies, defined as
# ```math
# \omega_h = \sqrt{\frac{k_h}{m}} \quad \omega_\theta = \sqrt{\frac{k_\theta}{I_P}}
# ```

using Aeroelasticity, DifferentialEquations, LinearAlgebra

## define non-dimensional parameters
V = range(1e-6, 3.1, length=1000) # = U/(b*ωθ) (reduced velocity)
a = -1/5 # reference point normalized location
e = -1/10 # center of mass normalized location
μ = 20 # = m/(ρ*pi*b^2) (mass ratio)
r2 = 6/25 # = Iθ/(m*b^2) (radius of gyration about P)
σ = 2/5 # = ωh/ωθ (natural frequency ratio)
xθ = e - a # distance from center of mass to reference point
a0 = 2*pi # lift curve slope
α0 = 0 # zero lift angle
cd0 = 0 # drag coefficient
cm0 = 0 # moment coefficient

## choose dimensional parameters
b = 1 # semichord
ρ = 1 # air density
ωθ = 1 # pitch natural frequency
c = 343 # air speed of sound

## calculate dimensionalized parameters
U = V*b*ωθ # freestrean velocity
m = μ*ρ*pi*b^2 # mass
Sθ = m*xθ*b # mass imbalance
Iθ = r2*m*b^2 # inertia
ωh = σ*ωθ # plunge natural frequency
kh = m*ωh^2 # plunge spring constant
kθ = Iθ*ωθ^2 # pitch spring constant

## define aerodynamic models
aerodynamic_models = (Steady(), QuasiSteady(), Wagner(), Peters{6}())

## initialize eigenvalue/eigenvector storage
λ = Vector{Matrix{ComplexF64}}(undef, length(aerodynamic_models))
Uλ = Vector{Array{ComplexF64,3}}(undef, length(aerodynamic_models))
Vλ = Vector{Array{ComplexF64,3}}(undef, length(aerodynamic_models))

## perform an analysis for each aerodynamic model
for imodel = 1:length(aerodynamic_models)

    ## define coupled model
    model = assemble_model(;
        aerodynamic_model = aerodynamic_models[imodel],
        structural_model = Section())

    ## define ODE function
    f = ODEFunction(model)

    ## eigenvalue/eigenvector storage
    nλ = number_of_states(model)
    λ[imodel] = zeros(ComplexF64, nλ, length(V))
    Uλ[imodel] = zeros(ComplexF64, nλ, nλ, length(V))
    Vλ[imodel] = zeros(ComplexF64, nλ, nλ, length(V))

    ## loop through each reduced frequency
    for i = 1:length(V)

        ## define aerodynamic parameters
        aerodynamic_parameters = (; a = a, b = b, a0 = a0, alpha0 = α0, cd0 = cd0, cm0 = cm0)

        ## define structural parameters
        structural_parameters = (; kh = kh, ktheta = kθ, m = m, Stheta = Sθ, Itheta = Iθ)

        ## define additional parameters
        additional_parameters = (; U = U[i], rho = ρ, c = c)

        ## define parameter vector
        p = assemble_parameters(model;
            aerodynamic_parameters = aerodynamic_parameters,
            structural_parameters = structural_parameters,
            additional_parameters = additional_parameters)

        ## define initial guess for equilibrium states
        x0 = assemble_states(model)

        ## find equilibrium point
        x = solve(SteadyStateProblem(f, x0, p))

        ## linearize about equilibrium point
        K, M = linearize(model, x, p)

        ## perform linear stability analysis
        λi, Uλi, Vλi = get_eigen(model, K, M)

        ## correlate eigenvalues
        if i > 1
            ## previous left eigenvector matrix
            Uλpi = Uλ[imodel][:,:,i-1]

            ## use correlation matrix to correlate eigenmodes
            perm, corruption = correlate_eigenmodes(Uλpi, M, Vλi)

            ## re-arrange eigenmodes
            λi = λi[perm]
            Uλi = Uλi[perm,:]
            Vλi = Vλi[:,perm]
        end

        ## save eigenvalues/eigenvectors
        λ[imodel][:,i] = λi
        Uλ[imodel][:,:,i] = Uλi
        Vλ[imodel][:,:,i] = Vλi
    end
end

#!jl nothing #hide

# We now plot the results for each aerodynamic model.

using Plots
pyplot()

sp1 = plot(
    title = "Non-Dimensional Frequency",
    xlim = (0,3.1),
    xtick = 0.0:0.5:3.0,
    xlabel = "\$ \\frac{U}{b \\omega_\\theta} \$",
    ylim = (0, 1.05),
    ytick = 0.0:0.2:1.0,
    ylabel = "\$ \\frac{\\Omega}{\\omega_\\theta} \$",
    titlefontsize = 10,
    guidefontsize = 10,
    legendfontsize = 8,
    tickfontsize = 9,
    foreground_color_legend = nothing,
    background_color_legend = nothing,
    minorgrid=false
    )

sp2 = plot(
    title = "Non-Dimensional Damping",
    xlim = (0,3.1),
    xtick = 0.0:0.5:3.0,
    xlabel = "\$ \\frac{U}{b \\omega_\\theta} \$",
    ylim = (-0.7, 0.605),
    ytick = -0.6:0.2:0.6,
    ylabel = "\$ \\frac{Γ}{\\omega_\\theta} \$",
    framestyle = :zerolines,
    titlefontsize = 10,
    guidefontsize = 10,
    legendfontsize = 8,
    tickfontsize = 9,
    legend = :topleft,
    foreground_color_legend = nothing,
    background_color_legend = nothing,
    minorgrid = false
    )

labels = ["Steady", "Quasi-Steady", "Wagner", "Peters (N=6)"]

for ia = 1:length(aerodynamic_models)

    plot!(sp1, V, imag.(λ[ia][1,:])/ωθ,
        label = labels[ia],
        color = ia,
        markersize = 1,
        markerstrokewidth = 0,
        )

    for i = 2:size(λ[ia], 1)
        plot!(sp1, V, imag.(λ[ia][i,:])/ωθ,
            label = "",
            color = ia,
            markersize = 1,
            markerstrokewidth = 0,
            )
    end

    plot!(sp2, V, real.(λ[ia][1,:])/ωθ,
        label = labels[ia],
        color = ia,
        markersize = 1,
        markerstrokewidth = 0,
        )

    for i = 2:size(λ[ia], 1)
        plot!(sp2, V, real.(λ[ia][i,:])/ωθ,
            label = "",
            color = ia,
            markersize = 1,
            markerstrokewidth = 0,
            )
    end
end

p1 = plot(sp1, sp2, layout = (1, 2), size = (800, 300))

#md savefig(p1, "../assets/section-stability.svg") #hide
#md nothing #hide

#md # ![](../assets/section-stability.svg)

# Using the `Wagner` or `Peters` aerodynamic models yields a flutter reduced velocity
# around 2.2, while the `Steady` and `QuasiSteady` aerodynamic models predict
# significantly lower flutter velocities.  The aerodynamic state variables of the
# `Wagner` and `Peters` models allows these models to capture the impact of vortex
# shedding on the lift and drag of the profile, therefore we can expect these models to
# yield more accurate results than the `Steady` and `QuasiSteady` models.
#
# The non-dimensional parameters we use for this example match those used by Hodges and
# Pierce in "Introduction to Structural Dynamics and Aeroelasticity".  Hodges and Pierce
# performed the analysis using a steady-state model and Peter's finite state model with
# six state variables.   The results presented here for the steady-state and Peters'
# finite state models match the results presented by Hodges and Pierce in "Introduction
# to Structural Dynamics and Aeroelasticity", which validates our implementation of
# these models.  Additionally, since the flutter speed predicted by the `Wagner` and
# `Peters` models match, we can be reasonably confident that the Wagner unsteady
# aerodynamic model is also implemented correctly.


# Script to run a test case for understanding how M scales with β #

# Add paths
loc_gm_package = "../src/"
if !(loc_gm_package in LOAD_PATH)
    push!(LOAD_PATH, "../src/")
end

# Required packages
using LinearAlgebra

using GraphicalModelLearning
using LightGraphs

# Function definitions
using Printf, SparseArrays
include("beta_scaling.jl")

using DelimitedFiles

# parameters
FLAG_create_struct_gm = true
FLAG_regular_random_gm = true
FLAG_lattice_gm = false
sampling_regime = M_regime()
learning_algo = NLP()

@printf("Creating a simple model!")
file_M_opt_gm = "M_opt_small_system_M_RPLE_regularization.txt"

# Simple test
N = 3
d = 2
α = 0.4
β_array = [1.0 + 0.1*i for i=0:20]

# Create and plot the initial graphical graphical model
β = copy(β_array[4])
adj_matrix = [0.0 β β; β 0.0 α; β α 0.0]
struct_adj_matrix = sign.(abs.(adj_matrix))

## Start the complexity studies
τ = α/2
L = 45
M_factor = 0.05
M_opt = Array{Int64,1}(undef,length(β_array))

# To prevent rewriting
adj_matrix_orig = deepcopy(adj_matrix)

let
    M_guess = 1000
    for i = 1:length(β_array)
        # define β
        β = copy(β_array[i])
        println(β)
        println(M_guess)

        # Create the adjacency matrix
        adj_matrix = replace_beta_gm(adj_matrix_orig, α, β)

        # Get the optimum number of samples
        M_opt[i] = get_M_opt_glauber_dynamics_regularization(adj_matrix, RPLE(), learning_algo, sampling_regime, τ, L, M_guess, M_factor)

        # Update the guess of M_opt
        @printf("beta=%f, M_opt=%d \n", β, M_opt[i])
        M_guess = Int(floor(0.8*copy(M_opt[i])))
    end
end

writedlm(file_M_opt_gm,M_opt)


module Test_clyap

using LinearAlgebra
using MatrixEquations
using Test

@testset "Testing continuous Lyapunov equation solvers" begin

n = 10
Ty = Float64

@testset "Continuous Lyapunov equations" begin

reltol = eps(float(100))
a = -1; b = 2im; @time x = lyapc(a,b)
@test abs(a*x+x*a'+b) < reltol

reltol = eps(float(100f0))
a = 1f0-2f0im; b = 2f0; @time x = lyapc(a,b)
@test abs(a*x+x*a'+b) < reltol

try
  x = lyapc(zeros(1,1),ones(1,1))
  @test false
catch
  @test true
end

try
  x = lyapc([0 1;  -1 0 ],ones(2,2))
  @test false
catch
  @test true
end

try
  x = lyapc([1 1;  0 -1],ones(2,2))
  @test false
catch
  @test true
end

try
  x = lyapc([1 1 1 1;  -1 1 0 1; 0 0 -1 1; 0 0 -1 -1],ones(4,4))
  @test false
catch
  @test true
end

for Ty in (Float64, Float32)

ar = rand(Ty,n,n);
ac = rand(Ty,n,n)+im*rand(Ty,n,n);
c = rand(Ty,n,n)+im*rand(Ty,n,n);
qc = c'*c;
Qr = real(qc);
Ty == Float64 ? reltol = eps(float(100)) : reltol = eps(100*n*one(Ty))

@time x = lyapc(ac,qc);
@test norm(ac*x+x*ac'+qc)/norm(x)/norm(ac) < reltol

α = 3+im; # α = im  # SingularException
@time x = lyapc(α*I,qc);
@test norm(α*x+x*α'+qc)/norm(x) < reltol

α = 3+im; # α = im  # SingularException
@time x = lyapc(α,Qr);
@test norm(α*x+x*α'+Qr)/norm(x) < reltol

@time x = lyapc(ac',qc);
@test norm(ac'*x+x*ac+qc)/norm(x)/norm(ac) < reltol

@time x = lyapc(ar,Qr)
@test norm(ar*x+x*ar'+Qr)/norm(x)/norm(ar)  < reltol

@time x = lyapc(ar',Qr)
@test norm(ar'*x+x*ar+Qr)/norm(x)/norm(ar) < reltol

@time x = lyapc(ac,Qr);
@test norm(ac*x+x*ac'+Qr)/norm(x)/norm(ac) < reltol

@time x = lyapc(ac',Qr);
@test norm(ac'*x+x*ac+Qr)/norm(x)/norm(ac) < reltol

@time x = lyapc(ar,qc);
@test norm(ar*x+x*ar'+qc)/norm(x)/norm(ar) < reltol

@time x = lyapc(ar',qc);
@test norm(ar'*x+x*ar+qc)/norm(x)/norm(ar)  < reltol
end
end

@testset "Continuous generalized Lyapunov equations" begin

reltol = eps(float(100))
a = -1+im; ee = 3im; b = 2; @time x = lyapc(a,ee,b)
@test abs(a*x*ee'+ee*x*a'+b) < reltol

reltol = eps(float(100f0))
a = 1f0-2f0im; ee = 3f0im; b = 2f0; @time x = lyapc(a,ee,b)
@test abs(a*x*ee'+ee*x*a'+b) < reltol

try
  x = lyapc(zeros(1,1),ones(1,1),ones(1,1))
  @test false
catch
  @test true
end

try
  x = lyapc([0 2;  -2 0 ],[2 0; 0 2],ones(2,2))
  @test false
catch
  @test true
end

try
  x = lyapc([2 2;  0 -2],[2 0; 0 2],ones(2,2))
  @test false
catch
  @test true
end

try
  x = lyapc(2*[1 1 1 1;  -1 1 0 1; 0 0 -1 1; 0 0 -1 -1],2*Matrix{Float64}(I,4,4), ones(4,4))
  @test false
catch
  @test true
end

for Ty in (Float64, Float32)

ar = rand(Ty,n,n)
ac = rand(Ty,n,n)+im*rand(Ty,n,n)
er = rand(Ty,n,n)
ec = er+im*rand(Ty,n,n)

c = rand(Ty,n,n)+im*rand(Ty,n,n)
qc = c'*c
Qr = real(qc)
Ty == Float64 ? reltol = eps(float(100)) : reltol = eps(100*n*one(Ty))

@time x = lyapc(ac,ec,qc);
@test norm(ac*x*ec'+ec*x*ac'+qc)/norm(x)/norm(ac)/norm(ec) < reltol

β = 3
@time x = lyapc(ac,β*I,qc);
@test norm(ac*x*β'+β*x*ac'+qc)/norm(x)/norm(ac)/norm(ec) < reltol

β = (1+im);
@time x = lyapc(ac,β,qc);
@test norm(ac*x*β'+β*x*ac'+qc)/norm(x)/norm(ac)/norm(ec) < reltol

α = 2+3im
@time x = lyapc(α*I,ec,qc);
@test norm(α*x*ec'+ec*x*α'+qc)/norm(x)/norm(ac)/norm(ec) < reltol

α = 2+3im; β = (1+im); # α = im; β = 1  # SingularException
@time x = lyapc(α,β,qc);
@test norm(α*x*β'+β*x*α'+qc)/norm(x)/norm(ac)/norm(ec) < reltol

α = 2+3im; β = (1+im); # α = im; β = 1  # SingularException
@time x = lyapc(α*I,β*I,qc);
@test norm(α*x*β'+β*x*α'+qc)/norm(x)/norm(ac)/norm(ec) < reltol

@time x = lyapc(ac',ec',qc);
@test norm(ac'*x*ec+ec'*x*ac+qc)/norm(x)/norm(ac)/norm(ec)  < reltol

β = (1+im);
@time x = lyapc(ac',β*I,qc);
@test norm(ac'*x*β+β'*x*ac+qc)/norm(x)/norm(ac)/norm(ec)  < reltol

α = 2+3im
@time x = lyapc(α*I,ec',qc);
@test norm(α'*x*ec+ec'*x*α+qc)/norm(x)/norm(ac)/norm(ec) < reltol

@time x = lyapc(ac',ec,qc);
@test norm(ac'*x*ec'+ec*x*ac+qc)/norm(x)/norm(ac)/norm(ec)  < reltol

@time x = lyapc(ac,ec',qc);
@test norm(ac*x*ec+ec'*x*ac'+qc)/norm(x)/norm(ac)/norm(ec)  < reltol

@time x = lyapc(ar,er,Qr);
@test norm(ar*x*er'+er*x*ar'+Qr)/norm(x)/norm(ar)/norm(er)  < reltol

@time x = lyapc(ar,er,qc);
@test norm(ar*x*er'+er*x*ar'+qc)/norm(x)/norm(ar)/norm(er)  < reltol

@time x = lyapc(ar',er',Qr);
@test norm(ar'*x*er+er'*x*ar+Qr)/norm(x)/norm(ar)/norm(er) < reltol

@time x = lyapc(ar',er,Qr);
@test norm(ar'*x*er'+er*x*ar+Qr)/norm(x)/norm(ar)/norm(er) < reltol

@time x = lyapc(ar,er',Qr);
@test norm(ar*x*er+er'*x*ar'+Qr)/norm(x)/norm(ar)/norm(er) < reltol

@time x = lyapc(ac,ec,Qr);
@test norm(ac*x*ec'+ec*x*ac'+Qr)/norm(x)/norm(ac)/norm(ec)   < reltol

@time x = lyapc(ac',ec',Qr);
@test norm(ac'*x*ec+ec'*x*ac+Qr)/norm(x)/norm(ac)/norm(ec)   < reltol

@time x = lyapc(ac',ec,Qr);
@test norm(ac'*x*ec'+ec*x*ac+Qr)/norm(x)/norm(ac)/norm(ec)   < reltol

@time x = lyapc(ac,ec',Qr);
@test norm(ac*x*ec+ec'*x*ac'+Qr)/norm(x)/norm(ac)/norm(ec)   < reltol

@time x = lyapc(ar,er,qc);
@test norm(ar*x*er'+er*x*ar'+qc)/norm(x)/norm(ar)/norm(er) < reltol

@time x = lyapc(ar',er',qc);
@test norm(ar'*x*er+er'*x*ar+qc)/norm(x)/norm(ar)/norm(er) < reltol

@time x = lyapc(ar',er,qc);
@test norm(ar'*x*er'+er*x*ar+qc)/norm(x)/norm(ar)/norm(er) < reltol

@time x = lyapc(ar,er',qc);
@test norm(ar*x*er+er'*x*ar'+qc)/norm(x)/norm(ar)/norm(er) < reltol
end
end

@testset "Continuous Lyapunov equations - Schur form" begin

for Ty in (Float64, Float32)

ar = rand(Ty,n,n);
ac = rand(Ty,n,n)+im*rand(Ty,n,n);
er = rand(Ty,n,n);
ec = er+im*rand(Ty,n,n);
as, es = schur(ar,er);
acs, ecs = schur(ac,ec);

c = rand(Ty,n,n)+im*rand(Ty,n,n);
qc = c'*c;
Qr = real(qc);
Ty == Float64 ? reltol = eps(float(100)) : reltol = eps(100*n*one(Ty))

x = copy(qc)
@time lyapcs!(acs,ecs,x);
@test norm(acs*x*ecs'+ecs*x*acs'+qc)/norm(x)/norm(acs)/norm(ecs) < reltol

x = copy(qc)
@time lyapcs!(acs,I,x);
@test norm(acs*x+x*acs'+qc)/norm(x)/norm(acs)/norm(ecs) < reltol

x = copy(qc)
@time lyapcs!(acs,ecs,x,adj=true);
@test norm(acs'*x*ecs+ecs'*x*acs+qc)/norm(x)/norm(acs)/norm(ecs) < reltol

# test
x = copy(Qr);
@time lyapcs!(as,es,x);
@test norm(as*x*es'+es*x*as'+Qr)/norm(x)/norm(as)/norm(es) < reltol

x = copy(Qr);
@time lyapcs!(as,es,x,adj=true);
@test norm(as'*x*es+es'*x*as+Qr)/norm(x)/norm(as)/norm(es) < reltol

x = copy(Qr)
@time lyapcs!(as,I,x);
@test norm(as*x+x*as'+Qr)/norm(x)/norm(as) < reltol

x = copy(Qr)
@time lyapcs!(as,x);
@test norm(as*x+x*as'+Qr)/norm(x)/norm(as) < reltol

if Ty == Float64
try
  x = convert(Matrix{Complex{Float32}},copy(qc))
  @time lyapcs!(as,x);
  @test false
catch
  @test true
end
end

x = copy(Qr)
@time lyapcs!(as,es,x,adj=true);
@test norm(as'*x*es+es'*x*as+Qr)/norm(x)/norm(as)/norm(es) < reltol

x = copy(qc)
@time lyapcs!(acs,x);
@test norm(acs*x+x*acs'+qc)/norm(x)/norm(acs) < reltol

x = copy(qc)
@time lyapcs!(acs,x,adj=true);
@test norm(acs'*x+x*acs+qc)/norm(x)/norm(acs) < reltol

x = copy(Qr)
@time lyapcs!(as,x);
@test norm(as*x+x*as'+Qr)/norm(x)/norm(as) < reltol

x = copy(Qr)
@time lyapcs!(as,x,adj=true);
@test norm(as'*x+x*as+Qr)/norm(x)/norm(as) < reltol
end
end

end

end


include("hyperelasticity.jl")
#
ts = ExplicitSolverTime(Δt, NT)
ubd, abd = compute_boundary_info(domain, globaldata, ts)
Fext = compute_external_force(domain, globaldata, ts)

d0 = zeros(2domain.nnodes)
v0 = zeros(2domain.nnodes)
a0 = zeros(2domain.nnodes)

mode = "consistent_tangent"
if length(ARGS)>=1
  global mode = ARGS[1]
end

if mode=="linear"
  global Dc, nn_law
  Dc = Variable(rand(3,3))
  Dc = spd(Dc)
  function nn_law(strain)
    strain_tensor = strain_voigt_to_tensor(strain)
    stress = strain*Dc
    stress
  end

elseif mode=="consistent_tangent"
  global Dc, nn_law
  Dc = Variable(rand(3,3))
  Dc = spd(Dc)
  function nn_law(strain)
    coef = ae(strain, [20,20,20,6])
    coef = [coef constant(ones(size(coef,1), 1))]
    H = consistent_tangent_matrix(coef, Dc)
    stress = batch_matmul(H, strain)
    stress
  end

elseif mode=="nn"
  global Dc, nn_law
  # general neural network
  function nn_law(ε)
    ae(ε, [20,20,20,3])
  end
elseif mode=="free_energy"
  global Dc, nn_law
  # free energy
  function nn_law(ε)
    φ = squeeze(ae(ε, [20,20,20,1]))
    tf.gradients(φ, ε)[1]
  end
else
  error("$mode not valid")
end

d, v, a= ExplicitSolver(globaldata, domain, d0, v0, a0, Δt, NT, nn_law, Fext, ubd, abd; strain_type="finite")

# test  = open("test.txt", "w")
sess = Session(); init(sess)
# run(sess, d)
idx = @. (0:n)*m + 1
idx = [idx; (@. idx + domain.nnodes)]
loss = sum((d[:, idx] - d_[:, idx])^2)
sess = Session(); init(sess)
@info run(sess, loss)
# error()

!isdir(mode) && mkdir(mode)

for i = 1:100
  loss_ = BFGS!(sess, loss, 1000)
  d0 = run(sess, d)

  # visualize
  close("all")
  p = visualize_displacement(d0, domain)
  saveanim(p, "$mode/$i.gif")
  close("all")
  plot(d_[:,1], "-", color="C1")
  plot(d0[:,1], "--", color="C1")
  plot(d_[:,1+domain.nnodes], color="C2")
  plot(d0[:,1+domain.nnodes], "--", color="C2")
  savefig("$mode/$i.png")

  open("$mode/loss.txt", "a") do io
    writedlm(io, loss_)
  end
  ADCME.save(sess, "$mode/$i.mat")
  if length(loss_)<1000
    break
  end
end


module lattice

export Lattice,LatticeD2Q9,LatticeD2Q4,LatticeD2Q5

abstract type Lattice end

# D2Q9 lattice
struct LatticeD2Q9 <: Lattice
  Δx::AbstractFloat
  Δt::AbstractFloat
  cc::AbstractFloat
  cs::AbstractFloat
  cssq::AbstractFloat
  f::Array{Float64, 3}
  c::Matrix{Int64}
  w::Vector{Float64}
  nw::Int
  nI::Int               # This should not belong to lattice?
  nJ::Int

  function LatticeD2Q9(dx::AbstractFloat, dt::AbstractFloat, nI::Int, nJ::Int,
                       rho::AbstractFloat=0.0)
    # Default lattice speed vectors and associated weights
    cdef = [1 0; 0 1; -1 0; 0 -1; 1 1; -1 1; -1 -1; 1 -1; 0 0]'
    wdef = [1.0/9.0; 1.0/9.0; 1.0/9.0; 1.0/9.0; 1.0/36.0; 1.0/36.0;
            1.0/36.0; 1.0/36.0; 4.0/9.0]
    nw = length(wdef)
    # Lattice functions for computing `c`, `cs`, and `cssq`
    cf(dx, dt)    = dx/dt
    csf(dx, dt)   = cf(dx, dt) / sqrt(3)
    cssqf(dx, dt) = csf(dx, dt) * csf(dx, dt)
    f = zeros(Float64, (nw, nI, nJ))
    for k=1:nw
      f[k,:,:] = fill(rho * wdef[k], (nI, nJ))
    end
    new(dx, dt, cf(dx, dt), csf(dx, dt), cssqf(dx, dt), f, cdef, wdef,
      nw, nI, nJ)
  end

end

# D2Q4 lattice
struct LatticeD2Q4 <: Lattice
  Δx::AbstractFloat     # Spatial interval
  Δt::AbstractFloat     # Time interval
  cc::AbstractFloat     # Lattice speed
  cs::AbstractFloat     # Sound speed
  cssq::AbstractFloat   # Sound speed squared
  f::Array{Float64, 3}  # Distribution
  c::Matrix{Int64}      # Speed vectors
  w::Vector{Float64}    # Weights
  nw::Int               # Number of velocity directions
  nI::Int               # This should not belong to lattice?
  nJ::Int

  function LatticeD2Q4(dx::AbstractFloat, dt::AbstractFloat, nI::Int, nJ::Int,
                       rho::Float64=0.0)
    # Default lattice speed vectors and associated weights
    cdef = [1 0; -1 0; 0 1; 0 -1]'
    wdef = [1.0/4.0; 1.0/4.0; 1.0/4.0; 1.0/4.0]
    nw = size(cdef, 2)

    # Lattice functions for computing `c`, `cs`, and `cssq`
    cf(dx, dt)    = dx/dt
    csf(dx, dt)   = cf(dx, dt) / sqrt(2)
    cssqf(dx, dt) = csf(dx, dt) * csf(dx, dt)

    f = zeros(Float64, (nw, nI, nJ))
    for k=1:nw
      f[k,:,:] = fill(rho * wdef[k], (nI, nJ))
    end
    new(dx, dt, cf(dx, dt), csf(dx, dt), cssqf(dx, dt), f, cdef, wdef,
      nw, nI, nJ)
  end

#  LatticeD2Q4(dx::AbstractFloat, dt::AbstractFloat, f::Array{Float64, 3}) =
#    new(dx, dt, cf(dx, dt), csf(dx, dt), cssqf(dx, dt), f, cdef, wdef, n)
end

# D2Q4 lattice
struct LatticeD2Q5 <: Lattice
  Δx::AbstractFloat     # Spatial interval
  Δt::AbstractFloat     # Time interval
  cc::AbstractFloat     # Lattice speed
  cs::AbstractFloat     # Sound speed
  cssq::AbstractFloat   # Sound speed squared
  f::Array{Float64, 3}  # Distribution
  c::Matrix{Int64}      # Speed vectors
  w::Vector{Float64}    # Weights
  nw::Int               # Number of velocity directions
  nI::Int               # This should not belong to lattice?
  nJ::Int

  function LatticeD2Q5(dx::AbstractFloat, dt::AbstractFloat, nI::Int, nJ::Int,
                       rho::Float64=0.0)
    # Default lattice speed vectors and associated weights
    cdef = [1 0; -1 0; 0 1; 0 -1; 0 0]'
    wdef = [2.0/12.0; 2.0/12.0; 2.0/12.0; 2.0/12.0; 4.0/12.0]
    nw = size(cdef, 2)

    # Lattice functions for computing `c`, `cs`, and `cssq`
    cf(dx, dt)    = dx/dt
    csf(dx, dt)   = cf(dx, dt) / sqrt(3)
    cssqf(dx, dt) = csf(dx, dt) * csf(dx, dt)

    f = zeros(Float64, (nw, nI, nJ))
    for k=1:nw
      f[k,:,:] = fill(rho * wdef[k], (nI, nJ))
    end
    new(dx, dt, cf(dx, dt), csf(dx, dt), cssqf(dx, dt), f, cdef, wdef,
      nw, nI, nJ)
  end

#  LatticeD2Q4(dx::AbstractFloat, dt::AbstractFloat, f::Array{Float64, 3}) =
#    new(dx, dt, cf(dx, dt), csf(dx, dt), cssqf(dx, dt), f, cdef, wdef, n)
end

# Opposite direction lattice vector index
#
# \param lat Lattice
# \param k Index of lattice vector
# \return Index of opposite direction
function opp_lat_vec(lat::LatticeD2Q9, k::Int)
  if k == 1
    return 3
  elseif k == 2
    return 4
  elseif k == 3
    return 1
  elseif k == 4
    return 2
  elseif k == 5
    return 7
  elseif k == 6
    return 8
  elseif k == 7
    return 5
  elseif k == 8
    return 6
  elseif k == 9
    return 9
  else
    error("$k > 9 for a D2Q9 lattice. Only nine vectors (1-9) possible");
  end
end

# Opposite direction lattice vector index
#
# \param lat Lattice
# \param k Index of lattice vector
# \return Index of opposite direction
function opp_lat_vec(lat::LatticeD2Q4, k::Int)
if k == 1
    return 2
elseif k == 2
    return 1
elseif k == 3
    return 4
elseif k == 4
    return 3
else
    error("$k > 4 for a D2Q4 lattice. Only four vectors (1-4) possible")
end

end

#=
# Helper function for initializing lattice
function _fill_lat(lat::Lattice, i_range::UnitRange{Int}, j_range::UnitRange{Int},
                   rho::Real)
  for k=1:lat.n, j=j_range, i=i_range
    lat.f[k, i, j] = rho * lat.w[k];
  end
end
=#

end


using CalibrationTests
using CalibrationErrors, Distributions
using Random, Test, Statistics

Random.seed!(1234)

# sample data
function generate_binary_data(nsamples)
    # generate predictions
    predictions = rand(Dirichlet(2, 1), nsamples)

    # generate labels
    labels_consistent = [rand() < predictions[1, i] ? 1 : 2 for i in 1:nsamples]
    labels_only_two = fill(2, nsamples)

    (predictions, labels_consistent), (predictions, labels_only_two)
end
const data_consistent, data_only_two = generate_binary_data(500)

# define matrix-valued kernel (using the mean total variation distance as bandwidth)
const kernel = UniformScalingKernel(ExponentialKernel(3))

@testset "Consistency test" begin
    println("consistency tests")

    # define estimators
    estimators = (ECE(UniformBinning(10)), BiasedSKCE(kernel),
                  LinearUnbiasedSKCE(kernel), QuadraticUnbiasedSKCE(kernel))

    for estimator in estimators
        test_consistent = ConsistencyTest(estimator, data_consistent)
        @test pvalue(test_consistent) > 0.7

        test_only_two = ConsistencyTest(estimator, data_only_two)
        @test pvalue(test_only_two) < 1e-6
    end
end

@testset "Distribution-free tests" begin
    println("distribution-free tests")

    # define estimators
    estimators = (BiasedSKCE(kernel), LinearUnbiasedSKCE(kernel),
                  QuadraticUnbiasedSKCE(kernel))

    for estimator in estimators
        test_consistent = DistributionFreeTest(estimator, data_consistent)
        @test pvalue(test_consistent) > 0.7

        test_only_two = ConsistencyTest(estimator, data_only_two)
        @test pvalue(test_only_two) < 1e-6
    end
end

@testset "Asymptotic linear test" begin
    println("asymptotic linear test")

    test_consistent = AsymptoticLinearTest(kernel, data_consistent)
    @test pvalue(test_consistent) > 0.7

    test_only_two = AsymptoticLinearTest(kernel, data_only_two)
    @test pvalue(test_only_two) < 1e-6
end

@testset "Asymptotic quadratic test" begin
    println("asymptotic quadratic test")

    test_consistent = AsymptoticQuadraticTest(kernel, data_consistent)
    @test pvalue(test_consistent) > 0.7

    test_only_two = AsymptoticQuadraticTest(kernel, data_only_two)
    @test pvalue(test_only_two) < 1e-6
end


import Pkg

if haskey(Pkg.installed(),"MosekTools")
    using MosekTools
    const dpkf_ok = true
else
    dpkf_ok = false
    println("WARNING: MosekTools.jl not installed, the functions
    for differentially private Kalman filtering cannot be used!")
end

#=
if haskey(Pkg.installed(),"SCS")
    using SCS  # does not work with SCS currently for some reason
    const dpkf_ok = true
else
    dpkf_ok = false
    println("WARNING: SCS.jl not installed, the functions for differentially
    private Kalman filtering cannot be used!")
end
=#

using JuMP
using ControlSystems
using LinearAlgebra

"""    (D, P_val, X_val, Ω_val, M) =
            staticInputBlock_DPKF_ss(Ls, As, Cs, Winvs, Vs, Vinvs, k_priv, ρ,
                                    m=size(As,1), p=size(Cs,1), r=size(Ls,1))

Compute an input matrix D for the two-block differentially private
steady-state Kalman filter mechanism. D takes linear combinations of
individual signals before privacy-preserving Gaussian noise injection.

Note that we do not guarantee observability of the pair (A,D*C). If
observability is lost and needed, call the dfactor function with the
returned matrix M and a truncation threshold lower than 1e-6 in order
to compute a new matrix D with more rows.

Inputs: one matrix per individual

- Ls: size(r,nusers*m), or size (r,m,nusers) if want estimator of sum_i L_i x_i
- As: size (m,m,nusers)
- Cs: size (p,m,nusers)
- Winvs: size (m,m,nusers) -- inverses of process noise cov. matrices
- Vs: size (p,p,nusers) -- measurement noise cov. matrices
- Vinvs: inverses of Vs (computed externally for modularity/efficiency)
- ρ: vector of size nusers, for the adjacency definition: ||yᵢ-yᵢ'||₂ <= ρᵢ
for user i.
- k_priv: proportionality constant for Gaussian privacy-preserving noise
injection. Compute it with k_priv = gaussianMechConstant(ϵ,δ) for your
choice of ϵ, δ.
"""

function staticInputBlock_DPKF_ss(Ls, As, Cs, Winvs, Vs, Vinvs, ρ, k_priv,
                                  nusers=size(As,3), m=size(As,1),
                                  p=size(Cs,1), r=size(Ls,1))

    modl = Model(Mosek.Optimizer)
    #modl = Model(with_optimizer(Mosek.Optimizer))
    #modl = Model(with_optimizer(SCS.Optimizer))

    Nm = nusers * m
    Np = nusers * p

    A = zeros(Nm, Nm)
    C = zeros(Np, Nm)
    V = zeros(Np, Np)
    Winv = zeros(Nm, Nm)
    Vinv = zeros(Np, Np)

    for i=1:nusers
        A[(i-1)*m+1:i*m, (i-1)*m+1:i*m] = As[:,:,i]
        C[(i-1)*p+1:i*p, (i-1)*m+1:i*m] = Cs[:,:,i]
        V[(i-1)*p+1:i*p, (i-1)*p+1:i*p] = Vs[:,:,i]
        Winv[(i-1)*m+1:i*m, (i-1)*m+1:i*m] = Winvs[:,:,i]
        Vinv[(i-1)*p+1:i*p, (i-1)*p+1:i*p] = Vinvs[:,:,i]
    end

    if size(Ls,3) > 1
        L = zeros(r, Nm)
        for i=1:nusers
            L[:, (i-1)*m+1:i*m] = Ls[:,:,i]
        end
    else
        L = copy(Ls)
    end

    if r == 1
        @variable(modl, X >= 0)
        @objective(modl, Min, X)
    else
        @variable(modl, X[1:r, 1:r], PSD)
        @objective(modl, Min, tr(X))
    end
    if Nm == 1
        @variable(modl, Ω >= 0)
    else
        @variable(modl, Ω[1:Nm, 1:Nm], PSD)
    end
    if Np == 1
        @variable(modl, P >= 0)
    else
        @variable(modl, P[1:Np, 1:Np], PSD)
    end

    @SDconstraint(modl, hvcat((2,2), X, L, L', Ω) >= zeros(r+Nm,r+Nm))

    @SDconstraint(modl, hvcat((2,2),
        C'*P*C-Ω+Winv, Winv*A, A'*Winv, Ω+A'*Winv*A) >= zeros(2*Nm,2*Nm))

    for i=1:nusers
        e = zeros(p,Np); e[:,(i-1)*p+1:i*p]=Matrix{Float64}(I,p,p)
        @SDconstraint(modl,
          hvcat((2,2), Matrix{Float64}(I,p,p)/(k_priv^2*ρ[i]^2)+Vinvs[:,:,i], e, e', V-V*P*V)
          >= zeros(Np+p,Np+p))
    end

    optimize!(modl)
    status = termination_status(modl)

    println("============================================")
    println("Solution status: ", status)
    println("Objective value: ", objective_value(modl))
    println("============================================")

    P_val = value.(P)

    X_val = value.(X)
    Ω_val = value.(Ω)

    #= # Debug
    Σ_val = inv(Ω_val)
    println("Norm of the difference between Σ^{-1} and (A Σ A' + W)^{-1}+C' P C")
    E1 = C'*P_val*C+inv(A*Σ_val*A'+inv(Winv))-inv(Σ_val)
    println(norm(E1))
    =#

    # Matrix to factorize
    M = k_priv^2 * (inv(V-V*P_val*V)-Vinv)
    # Compute D matrix
    D = dfactor(M; svaltol=1e-6)

    # We return M also, so that the factorization can be redone if needed
    return (D, P_val, X_val, Ω_val, M)
end

"""    (D, P_val, X_val, Ω_val, M) =
            staticInputBlock_DPLQG_ss(Ls, As, Bs, Cs, Vs, Vinvs, Winvs, k_priv, ρ,
                                    m=size(As,1), p=size(Cs,1), r=size(Ls,1))

Compute an input matrix D for the two-block differentially private
steady-state LQG mechanism. D takes linear combinations of
individual signals before privacy-preserving Gaussian noise injection.

Note that we do not guarantee observability of the pair (A,D*C). If
observability is lost and needed, call the dfactor function with the
returned matrix M and a truncation threshold lower than 1e-6 in order
to compute a new matrix D with more rows.

Inputs: one matrix per individual

- Q: size(nusers*m,nusers*m) (m = dim of individual state-space)
- R: size(du,du) (du = dim of control input signal)
- As: size (m,m,nusers)
- Bs: size (m,du,nsusers) - All B matrices pre-multiply the same input u
- Cs: size (p,m,nusers)
- Vs: size (p,p,nusers)
- Vinvs: inverses of Vs (computed externally for modularity/efficiency)
- Winvs: size (m,m,nusers)
- ρ: vector of size nusers, for the adjacency definition: ||yᵢ-yᵢ'||₂ <= ρᵢ
for user i.
- k_priv: proportionality constant for Gaussian privacy-preserving noise
injection. Compute it with k_priv = gaussianMechConstant(ϵ,δ) for your
choice of ϵ, δ.
"""
function staticInputBlock_DPLQG_ss(Q, R, As, Bs, Cs, Ws, Winvs, Vs, Vinvs,
                                ρ, k_priv, nusers=size(As,3), m=size(As,1),
                                p=size(Cs,1), du=size(Bs,2))

    Nm = nusers * m

    A = zeros(Nm, Nm)
    B = zeros(Nm, du)  # the same control input is shared by all agents
    W = zeros(Nm, Nm)

    for i=1:nusers
        A[(i-1)*m+1:i*m, (i-1)*m+1:i*m] = As[:,:,i]
        B[(i-1)*m+1:i*m, :] = Bs[:,:,i]  # stack the Bs in one column
        W[(i-1)*m+1:i*m, (i-1)*m+1:i*m] = Ws[:,:,i]
    end

    P = dare(A, B, Q, R)  #  computes ss LQR cost
    N = A'*P*A+Q-P
    L = dfactor(N, svaltol=1e-5)  # factorize L' L = P
    r = size(L,1)

    (D, perf_val, X_val, Ω_val, M) =
        staticInputBlock_DPKF_ss(L, As, Cs, Winvs, Vs, Vinvs, ρ, k_priv, nusers, m, p, r)
    return (D, tr(P*W)+tr(X_val), P, X_val, M)
end

"""    dfactor(M; svaltol=1e-4)

Compute D such that D'D = M for M positive semidefinite, via SVD. Tries to
output a wide D matrix by neglecting the singular values of M smaller than
svaltol*(max svals).
"""
function dfactor(M; svaltol=1e-4)
    (U,S,V) = svd(M)
    svalmax = maximum(S)
    threshold = svaltol*svalmax
    cols = []
    for i=1:length(S)
        if S[i] >= threshold
            push!(cols, i)
        end
    end
    D = zeros(size(M,1), length(cols))
    for j=1:length(cols)
        D[:,j] = U[:,cols[j]] * sqrt(S[cols[j]])
    end
    return sign(D[1,1])*D'
end

"""    evaluateKFperf(D, Ls, As, Cs, Vs, Ws, ρ, k_priv)

Compute the steady-state MSE of the two-block differentially private
Kalman filter, with a static matrix D to combine input signals.
The performance is computed by solving an algebraic Riccati equation.
It corresponds to the MSE at the end of a measurement update state in
the Kalman filter.

Inputs:
- D: size (q,nusers*p)  (q can be arbitrary, p = individual output signal dim.)
one matrix per individual
- Ls: size (r,m,nusers)  (if want estimator of sum_i L_i x_i)
- As: size (m,m,nusers)
- Cs: size (p,m,nusers)
- Ws: size (m,m,nusers) -- process noise cov. matrix
- Vs: size (p,p,nusers) -- measurement noise cov. matrix
- ρ: vector of size nusers, for the adjacency definition: ||yᵢ-yᵢ'||₂ <= ρᵢ
for user i.
- k_priv: proportionality constant for Gaussian privacy-preserving noise
injection. Compute it with k_priv = gaussianMechConstant(ϵ,δ) for your
choice of ϵ, δ.

Outputs: (performance, Σ, Σupdate)
where Σ and Σupdate are the error covariance matrices after the
prediction and measurement update step respectively
"""
function evaluateKFperf(D, Ls, As, Cs, Ws, Vs, ρ, k_priv)
    nusers=size(As,3)
    m=size(As,1)
    p=size(Cs,1)
    r=size(Ls,1)

    Nm = nusers * m
    Np = nusers * p

    L = zeros(r, Nm)
    A = zeros(Nm, Nm)
    C = zeros(Np, Nm)
    W = zeros(Nm, Nm)
    V = zeros(Np, Np)

    for i=1:nusers
        L[:, (i-1)*m+1:i*m] = Ls[:,:,i]
        A[(i-1)*m+1:i*m, (i-1)*m+1:i*m] = As[:,:,i]
        C[(i-1)*p+1:i*p, (i-1)*m+1:i*m] = Cs[:,:,i]
        W[(i-1)*m+1:i*m, (i-1)*m+1:i*m] = Ws[:,:,i]
        V[(i-1)*p+1:i*p, (i-1)*p+1:i*p] = Vs[:,:,i]
        #Winv[(i-1)*m+1:i*m, (i-1)*m+1:i*m] = Winvs[:,:,i]
        #Vinv[(i-1)*p+1:i*p, (i-1)*p+1:i*p] = Vinvs[:,:,i]
    end

    # compute sensitivity
    Δ₂ = 0
    for i=1:nusers
        Δ₂ = max(Δ₂, ρ[i] * maximum(svdvals(D[:,(i-1)*p+1:i*p])))
    end

    V₁ = D*V*D' + k_priv^2  * Δ₂^2 * Matrix{Float64}(I,size(D,1),size(D,1))
    C₁ = D*C
    V₁ = 0.5*(V₁+V₁')  # might have lost perfect symmetry due to numerical issues
    Σ = dare(A', C₁', W, V₁)  #  computes ss cov. after time update step
    Σupdate = Σ - Σ*C₁'*inv(C₁*Σ*C₁'+V₁)*C₁*Σ  # cov. afer meas. update step

    #return (trace(L*Σupdate*L'), Δ₂, Σupdate)
    return tr(L*Σupdate*L'), Σ, Σupdate
end

"""    evaluateLQGperf(D, Q, R, As, Bs, Cs, Ws, Vs, ρ, k_priv)

Compute the steady-state MSE of the two-block differentially private
LQG controller, with a static matrix D to combine input signals.
The performance is computed by solving two algebraic Riccati equations.

Inputs:
- D: size (q,nusers*p)  (q can be arbitrary, p = individual output signal dim.)
one matrix per individual
- Q: size(nusers*m,nusers*m) (m = dim of individual state-space)
- R: size(du,du) (du = dim of control input signal)
- As: size (m,m,nusers)
- Bs: size (m,du,nsusers) - All B matrices pre-multiply the same input u
- Cs: size (p,m,nusers)
- Ws: size (m,m,nusers) -- Process noise
- Vs: size (p,p,nusers) -- Measurement noise
- ρ: vector of size nusers, for the adjacency definition: ||yᵢ-yᵢ'||₂ <= ρᵢ
for user i.
- k_priv: proportionality constant for Gaussian privacy-preserving noise
injection. Compute it with k_priv = gaussianMechConstant(ϵ,δ) for your
choice of ϵ, δ.

Outputs: (performance, Σ, Σupdate, P)
where Σ and Σupdate are the error covariance matrices after the
prediction and measurement update step respectively, P is the solution
of the control algebraic Riccati equation
"""
function evaluateLQGperf(D, Q, R, As, Bs, Cs, Ws, Vs, ρ, k_priv)
    nusers=size(As,3)
    m=size(As,1)
    p=size(Cs,1)
    du=size(R,1)

    Nm = nusers * m
    Np = nusers * p

    A = zeros(Nm, Nm)
    B = zeros(Nm, du)  # the same control input is shared by all agents
    C = zeros(Np, Nm)
    W = zeros(Nm, Nm)
    V = zeros(Np, Np)

    for i=1:nusers
        A[(i-1)*m+1:i*m, (i-1)*m+1:i*m] = As[:,:,i]
        B[(i-1)*m+1:i*m, :] = Bs[:,:,i]  # stack the Bs in one column
        C[(i-1)*p+1:i*p, (i-1)*m+1:i*m] = Cs[:,:,i]
        W[(i-1)*m+1:i*m, (i-1)*m+1:i*m] = Ws[:,:,i]
        V[(i-1)*p+1:i*p, (i-1)*p+1:i*p] = Vs[:,:,i]
    end

    P = dare(A, B, Q, R)  #  computes ss LQR cost matrix

    # compute sensitivity
    Δ₂ = 0
    for i=1:nusers
        Δ₂ = max(Δ₂, ρ[i] * maximum(svdvals(D[:,(i-1)*p+1:i*p])))
    end

    V₁ = D*V*D' + k_priv^2  * Δ₂^2 * Matrix{Float64}(I,size(D,1),size(D,1))
    C₁ = D*C
    V₁ = 0.5*(V₁+V₁')  # might have lost perfect symmetry due to numerical issues
    Σ = dare(A', C₁', W, V₁)  #  computes ss cov. after time update step
    Σupdate = Σ - Σ*C₁'*inv(C₁*Σ*C₁'+V₁)*C₁*Σ  # cov. afer meas. update step

    N = Q+A'*P*A-P
    return tr(P*W + N*Σupdate), P, Σ, Σupdate
end


using SpectralKit, PGFPlotsX

basis = smolyak_basis(Chebyshev, EndpointGrid(), Val(3),
                      let b = BoundedLinear(-1,1); (b, b); end)
g = grid(basis)
p = @pgf Axis({scale_only_axis = true,
               width = "6cm",
               height = "6cm" },
              Plot({ only_marks, red }, Coordinates(Tuple.(g))))
pgfsave("smolyak_grid.png", p; dpi = 600)


using Manopt, Manifolds, ManifoldsBase, Random, LinearAlgebra, BenchmarkTools
Random.seed!(42)

function run_rayleigh_experiment(n::Int)
    A = randn(n, n)
    A = (A + A') / 2
    F(X::Array{Float64,1}) = X' * A * X
    ∇F(X::Array{Float64,1}) = 2 * (A * X - X * X' * A * X)
    M = Sphere(n - 1)
    x = random_point(M)
    return quasi_Newton(
        M,
        F,
        ∇F,
        x;
        #memory_size=-1,
        stopping_criterion=StopWhenAny(
            StopAfterIteration(max(1000)), StopWhenGradientNormLess(10^(-6))
        ),
        debug=[:Iteration, " ", :Cost, "\n", 1, :Stop],
    )
end
io = IOBuffer()

for n in [100]
    b = @benchmark run_rayleigh_experiment($n) samples = 30
    show(io, "text/plain", b)
    s = String(take!(io))
    println("Benchmarking $(n):\n", s, "\n\n")
end


@inline nrows(v::V) where {T, V<:AbstractVector{T}} = length(v)
@inline ncols(v::V) where {T, V<:AbstractVector{T}} = 1

@inline nrows(m::M) where {T, M<:AbstractMatrix{T}} = size(m)[1]
@inline ncols(m::M) where {T, M<:AbstractMatrix{T}} = size(m)[2]

@inline nrows(a::A) where {T, N, A<:AbstractArray{T,N}} = size(m)[1]
@inline ncols(a::A) where {T, N, A<:AbstractArray{T,N}} = N > 1 ? size(m)[2] : throw(BoundsError("$(size(a))[2]"))
@inline nstrata(a::A) where {T, N, A<:AbstractArray{T,N}} = N > 2 ? size(m)[3] : throw(BoundsError("$(size(a))[3]"))

#=
     these benchmark favorably relative to these canoncial formulations

@inline linear_to_cartesian(m::Array{T,N}, linearidx::Int) where {N, T} =
    Tuple(CartesianIndices(m)[linearidx])

@inline cartesian_to_linear(m::Array{T,N}, cartesianidx::C) where {T,N, C<:Tuple} =
     (LinearIndices(m))[cartesianidx...,]

@inline cartesian_to_linear(m::Array{T,N}, cartesianidx::C) where {N, T, C<:Vector{T}} =
     (LinearIndices(m))[cartesianidx...,]
=#

@inline cartesian_to_linear(nrows::Int, row::Int, col::Int) = row + (col-1)*nrows

@inline cartesian_to_linear(m::M, row::Int, col::Int) where {T,M<:AbstractMatrix{T}} =
    row + (col-1)*size(m)[1]

@inline linear_to_cartesian(m::M, linearidx::Int) where {T,M<:AbstractMatrix{T}} =
    linear_to_cartesian(size(m)[1], linearidx)

@inline function linear_to_cartesian(nrows::Int, linearidx::Int)
    row = rem(linearidx, nrows)
    row = ifelse(row === 0, nrows, row)
    col = div(linearidx-row, nrows) + 1
    return row,col
end


################################################
# Unconstrained
################################################
function fitnessToMass(fitness::Vector{Float64})
    m = minimum(fitness)

    if m < 0
        fitness = 2 * abs(m) .+ fitness
    end

    fitness = 2 * maximum(fitness) .- fitness

    return fitness
end

"""
    getMass(U, max_fitness, epsilon)
    get the mass in a vector `m[i] = f[i] + 2*max_fitness*sum_vio[i] `, tolerance
    in equality constraints are given by `epsilon`
"""
function getMass(U::Array{xf_indiv,1})
    n, d = length(U), length(U[1].x)

    fitness = zeros(Float64, n)

    for i = 1:n
        fitness[i] = U[i].f
    end

    return fitnessToMass(fitness)
end


isperfect(n::Integer) = n == sum([n % i == 0 ? i : 0 for i = 1:(n - 1)])
perfects(n::Integer) = filter(isperfect, 1:n)

@show perfects(10000)


"""
    SplitQuaternion{T <: Real} = Hyperbolic{Binion{T}}

A split-quaternion is a hyperbolic Cayley-Dickson construct with complex numbers.
"""
const SplitQuaternion{T <: Real} = Hyperbolic{Binion{T}}

function SplitQuaternion(z::Binion{T}) where T <: Real
    SplitQuaternion{T}(z, zero(Binion{T}))
end

function SplitQuaternion(a::T, b::T, c::T, d::T) where T <: Real
    SplitQuaternion{T}(Binion{T}(a, b), Binion{T}(c, d))
end

function SplitQuaternion(a::T, b::T, c::T) where T <: Real
    SplitQuaternion{T}(Binion{T}(a, b), Binion{T}(c, zero(T)))
end

function SplitQuaternion(a::T, b::T) where T <: Real
    SplitQuaternion{T}(Binion{T}(a, b), zero(Binion{T}))
end

function SplitQuaternion(a::T) where T <: Real
    SplitQuaternion{T}(Binion{T}(a, zero(T)), zero(Binion{T}))
end

function SplitQuaternion(a::Real, b::Real)
    SplitQuaternion(promote(a, b)...)
end

function SplitQuaternion(a::Real, b::Real, c::Real)
    SplitQuaternion(promote(a, b, c)...)
end

function SplitQuaternion(a::Real, b::Real, c::Real, d::Real)
    SplitQuaternion(promote(a, b, c, d)...)
end

"""
    unreal(z::SplitQuaternion)

The unreal part of a split-quaternion is a 3-vector.
"""
function unreal(z::SplitQuaternion)
    vcat(unreal(z.l), array(z.r))
end

function show(io::IO, z::SplitQuaternion)
    print(io, "[1: ")
    print(io, z.l.l)
    print(io, ", i: ")
    print(io, z.l.r)
    print(io, ", s: ")
    print(io, z.r.l)
    print(io, ", is: ")
    print(io, z.r.r)
    print(io, "]")
end

function random(::Type{SplitQuaternion{T}}) where T <: Real
    SplitQuaternion{T}(random(Binion{T}), random(Binion{T}))
end


using KernelAbstractions, CUDA, SimpleDiffEq
using StaticArrays

function loop(u, p, t)
    σ = p[1]; ρ = p[2]; β = p[3]
    du1 = σ*(u[2]-u[1])
    du2 = u[1]*(ρ-u[3]) - u[2]
    du3 = u[1]*u[2] - β*u[3]
    return SVector{3}(du1, du2, du3)
end
const func = ODEFunction(loop)
u0 = 10ones(Float32,3)
const su0= SVector{3}(u0)
const dt = 1f-1
const tspan = (0.0f0, 10.0f0)

const odeoop = ODEProblem{false}(loop, SVector{3}(u0), (0.0f0, 10.0f0),  Float32[10, 28, 8/3])
sol2 = solve(odeoop,GPUSimpleTsit5(),dt=dt)
cps = Array([@SVector [10f0,28f0,8/3f0] for i in 1:32])
ps = CuArray([@SVector [10f0,28f0,8/3f0] for i in 1:32])
CUDA.allowscalar(false)

function f(p)
    solve(odeoop,GPUSimpleTsit5(),dt=dt)[end]
end

map(f,cps)

@kernel function f_kernel(ps)
    solve(odeoop,GPUSimpleTsit5(),dt=dt)[end]
end

version = CUDADevice()
wait(version, f_kernel(version)(ps;ndrange=length(ps)))

function DiffEqBase.solve(prob::ODEProblem,
                          alg::GPUSimpleTsit5;
                          dt = 0.1f0)
  u0 = prob.u0
  tspan = prob.tspan
  f = prob.f
  p = prob.p
  t = tspan[1]
  tf = prob.tspan[2]
  #ts = tspan[1]:dt:tspan[2]
  us = MVector{101,typeof(u0)}(undef)
  us[1] = u0
  u = u0
  k7 = f(u, p, t)

  cs, as, btildes, rs = SimpleDiffEq._build_atsit5_caches(eltype(u0))
  c1, c2, c3, c4, c5, c6 = cs
  a21, a31, a32, a41, a42, a43, a51, a52, a53, a54,
  a61, a62, a63, a64, a65, a71, a72, a73, a74, a75, a76 = as
  btilde1, btilde2, btilde3, btilde4, btilde5, btilde6, btilde7 = btildes

  # FSAL
  for i in 2:101
      uprev = u; k1 = k7; t = tspan[1] + dt * i
      tmp = uprev+dt*a21*k1
      k2 = f(tmp, p, t+c1*dt)
      tmp = uprev+dt*(a31*k1+a32*k2)
      k3 = f(tmp, p, t+c2*dt)
      tmp = uprev+dt*(a41*k1+a42*k2+a43*k3)
      k4 = f(tmp, p, t+c3*dt)
      tmp = uprev+dt*(a51*k1+a52*k2+a53*k3+a54*k4)
      k5 = f(tmp, p, t+c4*dt)
      tmp = uprev+dt*(a61*k1+a62*k2+a63*k3+a64*k4+a65*k5)
      k6 = f(tmp, p, t+dt)
      u = uprev+dt*(a71*k1+a72*k2+a73*k3+a74*k4+a75*k5+a76*k6)
      k7 = f(u, p, t+dt)
      us[i] = u
  end

  SArray(us)
end

# Works now!
function f(p)
    solve(odeoop,GPUSimpleTsit5(),dt=dt)[end]
end

map(f,cps)

@kernel function f_kernel(ps)
    solve(odeoop,GPUSimpleTsit5(),dt=dt)[end]
end

version = CUDADevice()
wait(version, f_kernel(version)(ps;ndrange=length(ps)))

# Fails to change parameters
function f(p)
    remake(odeoop,p=p)
    solve(odeoop,GPUSimpleTsit5(),dt=dt)[end]
end

map(f,cps)

@kernel function f_kernel(ps)
    i = @index(Global)
    remake(odeoop,p=ps[i])
    solve(odeoop,GPUSimpleTsit5(),dt=dt)[end]
end

version = CUDADevice()
wait(version, f_kernel(version)(ps;ndrange=length(ps)))


include("grad.jl")

using Flux, Flux.Tracker, Distributions
import Flux.Tracker: lgamma
import Random

Random.seed!(0)

# Normal
println("Test Normal")
m = param(1)
log_s = param(log(.5))
# x = 2.0
z = rand()
x = Tracker.data(m + z * exp(log_s))

@time back!(logpdf(Normal(m, exp(log_s)), x))
@time res = grad_logq_normal(x, m.data, log_s.data)
@assert (m.tracker.grad, log_s.tracker.grad) == res

B = Int(1e5)
@time for i in 1:B
  back!(logpdf(Normal(m, exp(log_s)), x))
end
@time for i in 1:B
  res = grad_logq_normal(x, m.data, log_s.data)
end
# @assert (m.tracker.grad, log_s.tracker.grad) == res

# Gamma
println("Test Gamma")
log_a = param(log(1))
log_b = param(log(2))
a = exp(log_a)
b = exp(log_b)
x = rand(Gamma(a.data, b.data))

back!(logpdf(Gamma(a, b), x))
res_ad = (log_a.tracker.grad, log_b.tracker.grad)
res = grad_logq_gamma(x, log_a.data, log_b.data)
@assert res_ad == res

# Beta
println("Test Beta")
log_a = param(log(1))
log_b = param(log(2))
a = exp(log_a)
b = exp(log_b)
x = rand(Beta(a.data, b.data))

back!(logpdf(Beta(a, b), x))
res_ad = (log_a.tracker.grad, log_b.tracker.grad)
res = grad_logq_beta(x, log_a.data, log_b.data)
@assert res_ad == res

# Dirichlet
println("Test Dirichlet")
lpdf_dirichlet(p, a) = lgamma(sum(a)) - sum(lgamma.(a)) + sum((a .- 1) .* log.(p))
log_alpha = param(rand(10))
alpha = exp.(log_alpha)
x = rand(Dirichlet(Tracker.data(alpha)))

back!(lpdf_dirichlet(x, alpha))
res_ad = log_alpha.grad
res = grad_logq_dirichlet(x, log_alpha.data)
@assert all(abs.(res_ad - res) .< 1e-6)