text:
"""
    Dirichlet

The [Dirichlet distribution](http://en.wikipedia.org/wiki/Dirichlet_distribution) is often
used the conjugate prior for Categorical or Multinomial distributions.
The probability density function of a Dirichlet distribution with parameter
``\\alpha = (\\alpha_1, \\ldots, \\alpha_k)`` is:

```math
f(x; \\alpha) = \\frac{1}{B(\\alpha)} \\prod_{i=1}^k x_i^{\\alpha_i - 1}, \\quad \\text{ with }
B(\\alpha) = \\frac{\\prod_{i=1}^k \\Gamma(\\alpha_i)}{\\Gamma \\left( \\sum_{i=1}^k \\alpha_i \\right)},
\\quad x_1 + \\cdots + x_k = 1
```

```julia
# Let alpha be a vector
Dirichlet(alpha)         # Dirichlet distribution with parameter vector alpha

# Let a be a positive scalar
Dirichlet(k, a)          # Dirichlet distribution with parameter a * ones(k)
```
"""
struct Dirichlet{T<:Real} <: ContinuousMultivariateDistribution
    alpha::Vector{T}
    alpha0::T
    lmnB::T

    function Dirichlet{T}(alpha::Vector{T}) where T
        alpha0::T = zero(T)
        lmnB::T = zero(T)
        for i in 1:length(alpha)
            ai = alpha[i]
            ai > 0 || throw(ArgumentError("Dirichlet: alpha must be a positive vector."))
            alpha0 += ai
            lmnB += lgamma(ai)
        end
        lmnB -= lgamma(alpha0)
        new{T}(alpha, alpha0, lmnB)
    end

    function Dirichlet{T}(d::Integer, alpha::T) where T
        alpha0 = alpha * d
        new{T}(fill(alpha, d), alpha0, lgamma(alpha) * d - lgamma(alpha0))
    end
end

Dirichlet(alpha::Vector{T}) where {T<:Real} = Dirichlet{T}(alpha)
Dirichlet(d::Integer, alpha::T) where {T<:Real} = Dirichlet{T}(d, alpha)
Dirichlet(alpha::Vector{T}) where {T<:Integer} = Dirichlet{Float64}(convert(Vector{Float64},alpha))
Dirichlet(d::Integer, alpha::Integer) = Dirichlet{Float64}(d, Float64(alpha))

struct DirichletCanon
    alpha::Vector{Float64}
end

length(d::DirichletCanon) = length(d.alpha)

#### Conversions
convert(::Type{Dirichlet{Float64}}, cf::DirichletCanon) = Dirichlet(cf.alpha)
convert(::Type{Dirichlet{T}}, alpha::Vector{S}) where {T<:Real, S<:Real} = Dirichlet(convert(Vector{T}, alpha))
convert(::Type{Dirichlet{T}}, d::Dirichlet{S}) where {T<:Real, S<:Real} = Dirichlet(convert(Vector{T}, d.alpha))

Base.show(io::IO, d::Dirichlet) = show(io, d, (:alpha,))

# Properties

length(d::Dirichlet) = length(d.alpha)
mean(d::Dirichlet) = d.alpha .* inv(d.alpha0)
params(d::Dirichlet) = (d.alpha,)
@inline partype(d::Dirichlet{T}) where {T<:Real} = T

function var(d::Dirichlet)
    α = d.alpha
    α0 = d.alpha0
    c = 1.0 / (α0 * α0 * (α0 + 1.0))

    k = length(α)
    v = Vector{Float64}(undef, k)
    for i = 1:k
        @inbounds αi = α[i]
        @inbounds v[i] = αi * (α0 - αi) * c
    end
    return v
end

function cov(d::Dirichlet)
    α = d.alpha
    α0 = d.alpha0
    c = 1.0 / (α0 * α0 * (α0 + 1.0))

    k = length(α)
    C = Matrix{Float64}(undef, k, k)

    for j = 1:k
        αj = α[j]
        αjc = αj * c
        for i = 1:j-1
            @inbounds C[i,j] = - α[i] * αjc
        end
        @inbounds C[j,j] = αj * (α0 - αj) * c
    end

    for j = 1:k-1, i = j+1:k
        @inbounds C[i,j] = C[j,i]
    end
    return C
end

function entropy(d::Dirichlet)
    α = d.alpha
    α0 = d.alpha0
    k = length(α)

    en = d.lmnB + (α0 - k) * digamma(α0)
    for j in 1:k
        @inbounds αj = α[j]
        en -= (αj - 1.0) * digamma(αj)
    end
    return en
end

function dirichlet_mode!(r::Vector{T}, α::Vector{T}, α0::T) where T <: Real
    k = length(α)
    s = α0 - k
    for i = 1:k
        @inbounds αi = α[i]
        if αi <= one(T)
            error("Dirichlet has a mode only when alpha[i] > 1 for all i" )
        end
        @inbounds r[i] = (αi - one(T)) / s
    end
    return r
end

dirichlet_mode(α::Vector{T}, α0::T) where {T <: Real} = dirichlet_mode!(Vector{T}(undef, length(α)), α, α0)

mode(d::Dirichlet) = dirichlet_mode(d.alpha, d.alpha0)
mode(d::DirichletCanon) = dirichlet_mode(d.alpha, sum(d.alpha))

modes(d::Dirichlet) = [mode(d)]

# Evaluation

function insupport(d::Dirichlet, x::AbstractVector{T}) where T<:Real
    n = length(x)
    if length(d.alpha) != n
        return false
    end
    s = 0.0
    for i in 1:n
        xi = x[i]
        if xi < 0.0
            return false
        end
        s += xi
    end
    if abs(s - 1.0) > 1e-8
        return false
    end
    return true
end

function _logpdf(d::Dirichlet, x::AbstractVector{T}) where T<:Real
    a = d.alpha
    s = 0.
    for i in 1:length(a)
        @inbounds s += (a[i] - 1.0) * log(x[i])
    end
    return s - d.lmnB
end

# sampling

function _rand!(d::Union{Dirichlet,DirichletCanon}, x::AbstractVector{T}) where T<:Real
    s = 0.0
    n = length(x)
    α = d.alpha
    for i in 1:n
        @inbounds s += (x[i] = rand(Gamma(α[i])))
    end
    multiply!(x, inv(s)) # this returns x
end

#######################################
#
#  Estimation
#
#######################################

struct DirichletStats <: SufficientStats
    slogp::Vector{Float64}   # (weighted) sum of log(p)
    tw::Float64              # total sample weights

    DirichletStats(slogp::Vector{Float64}, tw::Real) = new(slogp, Float64(tw))
end

length(ss::DirichletStats) = length(s.slogp)

mean_logp(ss::DirichletStats) = ss.slogp * inv(ss.tw)

function suffstats(::Type{Dirichlet}, P::AbstractMatrix{Float64})
    K = size(P, 1)
    n = size(P, 2)
    slogp = zeros(K)
    for i = 1:n
        for k = 1:K
            @inbounds slogp[k] += log(P[k,i])
        end
    end
    DirichletStats(slogp, n)
end

function suffstats(::Type{Dirichlet}, P::AbstractMatrix{Float64}, w::AbstractArray{Float64})
    K = size(P, 1)
    n = size(P, 2)
    if length(w) != n
        throw(DimensionMismatch("Inconsistent argument dimensions."))
    end

    tw = 0.
    slogp = zeros(K)

    for i = 1:n
        @inbounds wi = w[i]
        tw += wi
        for k = 1:K
            @inbounds slogp[k] += log(P[k,i]) * wi
        end
    end
    DirichletStats(slogp, tw)
end

# fit_mle methods

## Initialization

function _dirichlet_mle_init2(μ::Vector{Float64}, γ::Vector{Float64})
    K = length(μ)

    α0 = 0.
    for k = 1:K
        @inbounds μk = μ[k]
        @inbounds γk = γ[k]
        ak = (μk - γk) / (γk - μk * μk)
        α0 += ak
    end
    α0 /= K

    multiply!(μ, α0)
end

function dirichlet_mle_init(P::AbstractMatrix{Float64})
    K = size(P, 1)
    n = size(P, 2)

    μ = Vector{Float64}(undef, K)  # E[p]
    γ = Vector{Float64}(undef, K)  # E[p^2]

    for i = 1:n
        for k = 1:K
            @inbounds pk = P[k, i]
            @inbounds μ[k] += pk
            @inbounds γ[k] += pk * pk
        end
    end

    c = 1.0 / n
    for k = 1:K
        μ[k] *= c
        γ[k] *= c
    end

    _dirichlet_mle_init2(μ, γ)
end

function dirichlet_mle_init(P::AbstractMatrix{Float64}, w::AbstractArray{Float64})
    K = size(P, 1)
    n = size(P, 2)

    μ = Vector{Float64}(undef, K)  # E[p]
    γ = Vector{Float64}(undef, K)  # E[p^2]
    tw = 0.

    for i = 1:n
        @inbounds wi = w[i]
        tw += wi
        for k = 1:K
            pk = P[k, i]
            @inbounds μ[k] += pk * wi
            @inbounds γ[k] += pk * pk * wi
        end
    end

    c = 1.0 / tw
    for k = 1:K
        μ[k] *= c
        γ[k] *= c
    end

    _dirichlet_mle_init2(μ, γ)
end

## Newton-Ralphson algorithm

function fit_dirichlet!(elogp::Vector{Float64}, α::Vector{Float64};
    maxiter::Int=25, tol::Float64=1.0e-12, debug::Bool=false)
    # This function directly overrides α

    K = length(elogp)
    length(α) == K || throw(DimensionMismatch("Inconsistent argument dimensions."))

    g = Vector{Float64}(undef, K)
    iq = Vector{Float64}(undef, K)
    α0 = sum(α)

    if debug
        objv = dot(α - 1.0, elogp) + lgamma(α0) - sum(lgamma(α))
    end

    t = 0
    converged = false
    while !converged && t < maxiter
        t += 1

        # compute gradient & Hessian
        # (b is computed as well)

        digam_α0 = digamma(α0)
        iz = 1.0 / trigamma(α0)
        gnorm = 0.
        b = 0.
        iqs = 0.

        for k = 1:K
            @inbounds ak = α[k]
            @inbounds g[k] = gk = digam_α0 - digamma(ak) + elogp[k]
            @inbounds iq[k] = - 1.0 / trigamma(ak)

            @inbounds b += gk * iq[k]
            @inbounds iqs += iq[k]

            agk = abs(gk)
            if agk > gnorm
                gnorm = agk
            end
        end
        b /= (iz + iqs)

        # update α

        for k = 1:K
            @inbounds α[k] -= (g[k] - b) * iq[k]
            @inbounds if α[k] < 1.0e-12
                α[k] = 1.0e-12
            end
        end
        α0 = sum(α)

        if debug
            prev_objv = objv
            objv = dot(α - 1.0, elogp) + lgamma(α0) - sum(lgamma(α))
            @printf("Iter %4d: objv = %.4e  ch = %.3e  gnorm = %.3e\n",
                t, objv, objv - prev_objv, gnorm)
        end

        # determine convergence

        converged = gnorm < tol
    end

    Dirichlet(α)
end

function fit_mle(::Type{Dirichlet}, P::AbstractMatrix{Float64};
    init::Vector{Float64}=Float64[], maxiter::Int=25, tol::Float64=1.0e-12)

    α = isempty(init) ? dirichlet_mle_init(P) : init
    elogp = mean_logp(suffstats(Dirichlet, P))
    fit_dirichlet!(elogp, α; maxiter=maxiter, tol=tol)
end

function fit_mle(::Type{Dirichlet}, P::AbstractMatrix{Float64}, w::AbstractArray{Float64};
    init::Vector{Float64}=Float64[], maxiter::Int=25, tol::Float64=1.0e-12)

    n = size(P, 2)
    length(w) == n || throw(DimensionMismatch("Inconsistent argument dimensions."))

    α = isempty(init) ? dirichlet_mle_init(P, w) : init
    elogp = mean_logp(suffstats(Dirichlet, P, w))
    fit_dirichlet!(elogp, α; maxiter=maxiter, tol=tol)
end


text:
# Function for running Riemannian manifold Hamiltonian Monte Carlo (RMHMC)
function rmhmc(model::Model, opts::RmhmcOpts)
  mcmc = Array(Float64, opts.mcmc.nPostBurnin, model.nPars)
  z = Array(Float64, opts.mcmc.nPostBurnin, model.nPars)
  proposedGradLogPosterior = Array(Float64, model.nPars)
  invGDerivG = Array(Float64, model.nPars, model.nPars, model.nPars)
  traceInvGDerivG = Array(Float64, model.nPars)
  momentumTerm = Array(Float64, model.nPars)
  proposed, accepted = 0., 0.

  println("Running burn-in iterations...")

  currentPars = model.randPrior()
  currentLogPosterior = model.logPosterior(currentPars)
  currentGradLogPosterior = model.gradLogPosterior(currentPars)
  leapStep = opts.setLeapStep(1, 0., opts.mcmc.n, opts.mcmc.nBurnin, 0.)

  for i = 1:opts.mcmc.n
    proposed += 1

    proposedPars = copy(currentPars)
    G = model.tensor(proposedPars)
    invG = inv(G)
    cholG = chol(G)
    momentum = cholG'*randn(model.nPars)
    currentLogDetG = 0.5*(log(2)+model.nPars*log(pi)+2*sum(log(diag(cholG))))
    currentHamiltonian  =
      -currentLogPosterior+currentLogDetG+(momentum'*(invG*momentum))/2
    derivG = model.derivTensor(proposedPars)
    for j = 1:model.nPars
      invGDerivG[:, :, j] = invG*derivG[:, :, j]
      traceInvGDerivG[j] = trace(invGDerivG[:, :, j])
    end
    timeStep = (randn() > 0.5 ? 1. : -1.)
    randomSteps = ceil(rand()*opts.nLeaps)
    for j = 1:randomSteps
      leapGradLogPosterior  = model.gradLogPosterior(proposedPars)
      leapMomentum = copy(momentum)
      for k = 1:opts.nNewton
        invGMomentum = invG*leapMomentum
        for r = 1:model.nPars
          momentumTerm[r] =
            (0.5*(leapMomentum'*invGDerivG[:, :, r]*invGMomentum))[1]
        end
        leapMomentum = (momentum
          +timeStep*(leapStep/2)*(leapGradLogPosterior-0.5*traceInvGDerivG
          +momentumTerm))
      end
      momentum = copy(leapMomentum)
      leapInvGMomentum = invG*momentum
      leapParameters = copy(proposedPars)
      for k = 1:opts.nNewton
        G = model.tensor(leapParameters)
        invGMomentum = G\momentum
        leapParameters =
          proposedPars+(timeStep*(leapStep/2))*(leapInvGMomentum+invGMomentum)
      end
      proposedPars = copy(leapParameters)
      G = model.tensor(proposedPars)
      invG = inv(G)
      derivG = model.derivTensor(proposedPars)
      for k = 1:model.nPars
        invGDerivG[:, :, k] = invG*derivG[:, :, k]
        traceInvGDerivG[k] = trace(invGDerivG[:, :, k])
      end
      invGMomentum = invG*momentum
      for k = 1:model.nPars
        momentumTerm[k] = (0.5*(momentum'*invGDerivG[:, :, k]*invGMomentum))[1]
      end
      proposedGradLogPosterior = model.gradLogPosterior(proposedPars)
      momentum = (momentum
        +timeStep*(leapStep/2)*(proposedGradLogPosterior-0.5*traceInvGDerivG
        +momentumTerm))
    end
    proposedLogPosterior = model.logPosterior(proposedPars)
    proposedLogDet = 0.5*(log(2)+model.nPars*log(pi)+2*sum(log(diag(chol(G)))))
    proposedHamiltonian =
      -proposedLogPosterior+proposedLogDet+(momentum'*(invG*momentum))/2

    ratio = (currentHamiltonian-proposedHamiltonian)[1]

    if ratio > 0 || (ratio > log(rand()))
      accepted += 1

      currentPars = copy(proposedPars)
      currentLogPosterior = copy(proposedLogPosterior)
      currentGradLogPosterior = copy(proposedGradLogPosterior)
    end

    if i > opts.mcmc.nBurnin
      mcmc[i-opts.mcmc.nBurnin, :] = currentPars
      z[i-opts.mcmc.nBurnin, :] = -currentGradLogPosterior/2
    end

    if mod(i, opts.mcmc.monitorRate) == 0
      acceptanceRatio = accepted/proposed

    leapStep = opts.setLeapStep(i, acceptanceRatio, opts.mcmc.n,
      opts.mcmc.nBurnin, leapStep)

      println("Iteration $i of $(opts.mcmc.n): ", round(100*acceptanceRatio, 2),
        " % acceptance ratio")

      proposed, accepted = 0., 0.
    end

    if i == opts.mcmc.nBurnin
      println("Burn-in completed...\n\nRunning post burn-in MCMC...");
    end
  end

  return mcmc, z
end

ext: jl


text:
"""
    f!_from_f(f, F::Abstractarray)

Return an inplace version of f
"""
function f!_from_f(f, F::AbstractArray, inplace)
    if inplace
        return f
    else
        return function ff!(F, x)
            copyto!(F, f(x))
            F
        end
    end
end
function df!_from_df(g, F::Real, inplace)
    if inplace
        return g
    else
        return function gg!(G, x)
            copyto!(G, g(x))
            G
        end
    end
end
function df!_from_df(j, F::AbstractArray, inplace)
    if inplace
        return j
    else
        return function jj!(J, x)
            copyto!(J, j(x))
            J
        end
    end
end
function fdf!_from_fdf(fg, F::Real, inplace)
    if inplace
        return fg
    else
        return function ffgg!(G, x)
            fx, gx = fg(x)
            copyto!(G, gx)
            fx
        end
    end
end
function fdf!_from_fdf(fj, F::AbstractArray, inplace)
    if inplace
        return fj
    else
        return function ffjj!(F, J, x)
            fx, jx = fj(x)
            copyto!(J, jx)
            copyto!(F, fx)
        end
    end
end
function h!_from_h(h, F::Real, inplace)
    if inplace
        return h
    else
        return function hh!(H, x)
            copyto!(H, h(x))
            H
        end
    end
end


text:
"""
    print_workspace(cur_module)

Lorem ipsum dolor sit amet.
"""
function print_workspace(cur_module)

  module_workspace = OrderedDict()

  skip_var_keys = [ Symbol(cur_module) , :doc ]

  skip_start_chars = [ "#" , "/" ]

  skip_julia_types = [ Function , Type , Module ]

  for cur_var in sort(names(cur_module, true))

    skip_var = false

    for cur_skip_key in skip_var_keys
      skip_var |= ( cur_var == cur_skip_key )
    end

    for cur_skip_char in skip_start_chars
      skip_var |= startswith(string(cur_var), cur_skip_char)
    end

    if skip_var ; continue ; end

    cur_value = getfield(cur_module, cur_var)

    for cur_skip_type in skip_julia_types
      skip_var |= isa(cur_value, cur_skip_type)
    end

    if skip_var ; continue ; end

    module_workspace[cur_var] = cur_value

  end

  longest_name = maximum(
    map(
      x -> length(string(x)),
      keys(module_workspace)
    )
  )

  longest_name += 1

  dash_line_length = length(string(cur_module)) + 2

  truncate_length = ( 80 - 5 - longest_name )

  module_header = "\n\n"
  module_header *= "-" ^ dash_line_length
  module_header *= "\n $(cur_module) \n"
  module_header *= "-" ^ dash_line_length
  module_header *= "\n"

  println(module_header)

  for (cur_var, cur_value) in module_workspace
    padded_var = Base.cpad(string(cur_var), longest_name)

    truncated_value = cur_value

    if isa(cur_value, String)
      truncated_value = SubString(
        cur_value, 1, truncate_length
      )
    end

    if truncated_value != cur_value
      truncated_value *= "..."
    end

    println("$padded_var: $truncated_value")
  end

end


text:
using Junet
using Test

@testset "util.jl" begin
    @testset "DiscreteSampler" begin
        x = rand(1:66, 100)
        dp = Junet.DiscreteSampler(Int[])
        for i = x
            @test_nowarn push!(dp, i)
        end
        for _ = 1:1000
            i, j = Junet.randd(dp, return_residual=true)
            @test 1 <= i <= length(x)
            @test 1 <= j <= x[i]
        end
    end

end

@testset "primitives.jl" begin
    @testset "NodePtr" begin
        @test Junet.NodePtr(1, 2) < Junet.NodePtr(2, 2)
        @test Junet.NodePtr(2, 1) > 1
    end

    @testset "Node" begin
        n = Junet.Node{Int,Int}([], [])
        @test length(n.forward) == length(n.reverse) == 0
        push!(n.forward, Junet.NodePtr(1, 2))
        @test length(n.forward) == 1
    end
end

@testset "graph.jl" begin
    g = Graph()
    ug = Graph(directed=false)

    @testset "directed?" begin
        @test isdirected(g)
        @test !isdirected(ug)

        @test typeof(directed(g)) <: DirectedGraph
        @test typeof(directed(ug)) <: DirectedGraph
        @test typeof(undirected(g)) <: UndirectedGraph
        @test typeof(undirected(ug)) <: UndirectedGraph

        @test typeof(g'') == typeof(g)
        @test typeof(ug') == typeof(ug)
    end

    @testset "simple?" begin
        @test ismultigraph(g)
        @test ismultigraph(ug)

        lg = Graph(TEdge=Nothing)
        @test typeof(lg) <: LightGraph
    end

    @testset "node add/remove/count" begin
        for i = 1:10
            addnode!(g)
            @test nodecount(g) == i
        end
        @test_nowarn addnodes!(g, 10)
        @test_nowarn addnodes!(g, 0)
        @test_throws AssertionError addnodes!(g, -1)
        @test nodecount(g) == 20

        for i = 19:-1:0
            remnode!(g, rand(nodes(g)))
            @test nodecount(g) == i
        end
    end
end

@testset "attributes.jl" begin
    c = Junet.ConstantAttribute(1, ()->69)
    @test length(c) == 69
    @test all(x == 1 for x = c)

    s = Junet.SparseAttribute(c)
    @test length(s) == 69
    @test all(x == 1 for x = s)

    d = Junet.DenseAttribute(s)
    @test length(s) == 69
    @test all(x == 1 for x = s)
end

@testset "generate/" begin
    @testset "classic.jl" begin
        @testset "graph_complete" begin
            for i = 2:10
                @test density(graph_complete(i)) == 1
                @test density(graph_complete(i, directed=false)) == 1
            end
        end

        @testset "graph_tree" begin
            @test_nowarn graph_tree(0, 0)
            counts(g::Graph) = (nodecount(g), edgecount(g))
            @test counts(graph_tree(2, 1)) == (3, 2)
            @test counts(graph_tree(2, 5)) == (63, 62)
            @test counts(graph_tree(4, 4)) == (341, 340)
        end
    end

    @testset "random.jl" begin
        @testset "gilbert_fill!" begin
            g = Graph(directed=true)
            @test_nowarn Junet.gilbert_fill!(g, 1.)
            @test edgecount(g) == 0
            addnode!(g)         # nodecount(g) == 1
            @test_nowarn Junet.gilbert_fill!(g, 1.)
            @test edgecount(g) == 0
            addnode!(g)         # nodecount(g) == 2
            @test_nowarn Junet.gilbert_fill!(g, 1.)
            @test edgecount(g) == 2

            g = Graph(directed=false)
            @test_nowarn Junet.gilbert_fill!(g, 1.)
            @test edgecount(g) == 0
            addnode!(g)         # nodecount(g) == 1
            @test_nowarn Junet.gilbert_fill!(g, 1.)
            @test edgecount(g) == 0
            addnode!(g)         # nodecount(g) == 2
            @test_nowarn Junet.gilbert_fill!(g, 1.)
            @test edgecount(g) == 1
        end

        @testset "graph_gilbert" begin
            @test_nowarn graph_gilbert(0, 0)
            @test_nowarn graph_gilbert(1, 0)
            @test_nowarn graph_gilbert(1, 1)
            @test_throws AssertionError graph_gilbert(10, 2)

            @test density(graph_gilbert(100, 0)) == 0
            @test density(graph_gilbert(100, 0, directed=false)) == 0
            @test density(graph_gilbert(100, 1)) == 1
            @test density(graph_gilbert(100, 1, directed=false)) == 1

            for i = 0:0.1:1.0
                g1 = graph_gilbert(1000, i)
                @test density(g1) ≈ i  atol = .002
                @test selfloopcount(g1) == 0

                g2 = graph_gilbert(1000, i, directed=false)
                @test density(g2) ≈ i  atol = .002
                @test selfloopcount(g2) == 0
            end
        end

        @testset "graph_erdos_renyi" begin
            @test_nowarn graph_erdos_renyi(0, 0)
            @test_nowarn graph_erdos_renyi(1, 0)
            @test_throws ErrorException graph_erdos_renyi(1, 1)
            @test_throws ErrorException graph_erdos_renyi(10, 91)
            @test_throws ErrorException graph_erdos_renyi(10, 46, directed=false)

            for i = 100:100:1000
                j = rand(1:i * (i - 1))
                g1 = graph_erdos_renyi(i, j)
                @test edgecount(g1) == j
                @test selfloopcount(g1) == 0

                j = rand(1:i * (i - 1) >>> 1)
                g2 = graph_erdos_renyi(i, j, directed=false)
                @test edgecount(g2) == j
                @test selfloopcount(g2) == 0
            end
        end

    end

end

@testset "draw/" begin
    @testset "layout.jl" begin
        @testset "rescale" begin
            @test Junet.rescale([], 10, 1) == []
            @test Junet.rescale([0], 10, 1) == [5]
            @test Junet.rescale([0, 0], 10, 1) == [5, 5]
            @test Junet.rescale([0, 1], 10, 1) == [1, 9]
        end

    end

    @testset "plot.jl" begin
        @testset "getrgb" begin
            @test Junet.getrgb("red") == (1, 0, 0)
            @test Junet.getrgb("black") == (0, 0, 0)
            @test Junet.getrgb("#000") == (0, 0, 0)
            @test Junet.getrgb("#ff00ff") == (1, 0, 1)
            @test Junet.getrgb((0, 255, 0)) == (0, 1, 0)

        end

    end
end


text:
# this line because the const objects are not being exported
include(Pkg.dir("TimeSeries/src/.timeseriesrc.jl"))

facts("const values are set the package defaults") do

  context("SHOWINT") do
      @fact SHOWINT => false
  end

  context("DECIMALS") do
      @fact DECIMALS => 2
  end
end


text:
module BarotropicQGQL

export
  Problem,
  set_zeta!,
  updatevars!,

  energy,
  enstrophy,
  dissipation,
  work,
  drag

using
  FFTW,
  CUDA,
  Reexport,
  DocStringExtensions

@reexport using FourierFlows

using LinearAlgebra: mul!, ldiv!
using FourierFlows: parsevalsum, parsevalsum2
import FFTW: rfft

abstract type BarotropicQGQLVars <: AbstractVars end

nothingfunction(args...) = nothing

"""
    Problem(dev::Device=CPU();
                      nx = 256,
                      ny = nx,
                      Lx = 2π,
                      Ly = Lx,
                       β = 0.0,
      deformation_radius = Inf,
                     eta = nothing,
                       ν = 0.0,
                      nν = 1,
                       μ = 0.0,
                      dt = 0.01,
                 stepper = "RK4",
                   calcF = nothingfunction,
              stochastic = false,
        aliased_fraction = 1/3,
                       T = Float64)

Construct a quasi-linear barotropic quasi-geostrophic `problem` on device `dev`.

Keyword arguments
=================
    - `dev`: (required) `CPU()` or `GPU()`; computer architecture used to time-step `problem`.
    - `nx`: Number of grid points in ``x``-domain.
    - `ny`: Number of grid points in ``y``-domain.
    - `Lx`: Extent of the ``x``-domain.
    - `Ly`: Extent of the ``y``-domain.
    - `β`: Planetary vorticity ``y``-gradient.
    - `eta`: Topographic potential vorticity.
    - `ν`: Small-scale (hyper)-viscosity coefficient.
    - `nν`: (Hyper)-viscosity order, `nν```≥ 1``".
    - `μ`: Linear drag coefficient.
    - `dt`: Time-step.
    - `stepper`: The extent of the ``y``-domain.
    - `calcF`: Function that calculates the Fourier transform of the forcing, ``F̂``.
    - `stochastic`: `true` or `false`; boolean denoting whether `calcF` is temporally stochastic.
    - `aliased_fraction`: the fraction of high-wavenubers that are zero-ed out by `dealias!()`.
    - `T`: `Float32` or `Float64`; floating point type used for `problem` data.
"""
function Problem(dev::Device=CPU();
  # Numerical parameters
                nx = 256,
                ny = nx,
                Lx = 2π,
                Ly = Lx,
  # Physical parameters
                 β = 0.0,
               eta = nothing,
  # Drag and/or hyper-/hypo-viscosity
                 ν = 0.0,
                nν = 1,
                 μ = 0.0,
  # Timestepper and equation options
                dt = 0.01,
           stepper = "RK4",
             calcF = nothingfunction,
        stochastic = false,
  # Float type and dealiasing
  aliased_fraction = 1/3,
                 T = Float64)

  # the grid
  grid = TwoDGrid(dev, nx, Lx, ny, Ly; aliased_fraction=aliased_fraction, T=T)
  x, y = gridpoints(grid)

  # topographic PV
  eta === nothing && ( eta = zeros(dev, T, (nx, ny)) )

  params = !(typeof(eta)<:ArrayType(dev)) ?
           Params(grid, β, eta, μ, ν, nν, calcF) :
           Params(β, eta, rfft(eta), μ, ν, nν, calcF)

  vars = calcF == nothingfunction ? DecayingVars(dev, grid) : stochastic ? StochasticForcedVars(dev, grid) : ForcedVars(dev, grid)

  equation = BarotropicQGQL.Equation(params, grid)

  FourierFlows.Problem(equation, stepper, dt, grid, vars, params, dev)
end

# ----------
# Parameters
# ----------

"""
    Params{T, Aphys, Atrans}(β, eta, etah, μ, ν, nν, calcF!)

A struct containing the parameters for a barotropic QL QG problem. Included are:

$(TYPEDFIELDS)
"""
struct Params{T, Aphys, Atrans} <: AbstractParams
    "planetary vorticity y-gradient"
       β :: T
    "topographic potential vorticity"
     eta :: Aphys
    "Fourier transform of topographic potential vorticity"
    etah :: Atrans
    "linear drag coefficient"
       μ :: T
    "small-scale (hyper)-viscosity coefficient"
       ν :: T
    "(hyper)-viscosity order, `nν```≥ 1``"
      nν :: Int
    "function that calculates the Fourier transform of the forcing, ``F̂``"
  calcF! :: Function
end

"""
    Params(grid, β, eta::Function, μ, ν, nν, calcF)

Return the `params` for barotropic QL QG problem on `grid` and with topographic PV prescribed
as function, `eta(x, y)`.
"""
function Params(grid::AbstractGrid{T, A}, β, eta::Function, μ, ν, nν, calcF) where {T, A}
  eta_on_grid = FourierFlows.on_grid(eta, grid)
  etah_on_grid = rfft(eta_on_grid)

  return Params(β, A(eta_on_grid), A(etah_on_grid), μ, ν, nν, calcF)
end

# ---------
# Equations
# ---------

"""
    Equation(params, grid)

Return the equation for two-dimensional barotropic QG QL problem with parameters `params` and
on `grid`. Linear operator ``L`` includes bottom drag ``μ``, (hyper)-viscosity of order ``n_ν``
with coefficient ``ν`` and the ``β`` term:
```math
L = - μ - ν |𝐤|^{2 n_ν} + i β k_x / |𝐤|² .
```
Nonlinear term is computed via `calcN!` function.
"""
function Equation(params::Params, grid::AbstractGrid)
  L = @. - params.μ - params.ν * grid.Krsq^params.nν + im * params.β * grid.kr * grid.invKrsq
  CUDA.@allowscalar L[1, 1] = 0

  return FourierFlows.Equation(L, calcN!, grid)
end

# ----
# Vars
# ----

"""
    Vars{Aphys, Atrans, F, P}(u, v, U, uzeta, vzeta, zeta, Zeta, psi, Psi, N, NZ, uh, vh, Uh, zetah, Zetah, psih, Psih, Fh, prevsol)

The variables for barotropic QL QG:

$(FIELDS)
"""
mutable struct Vars{Aphys, Atrans, F, P} <: BarotropicQGQLVars
    "x-component of small-scale velocity"
        u :: Aphys
    "y-component of small-scale velocity"
        v :: Aphys
    "x-component of large-scale velocity"
        U :: Aphys
    "small-scale u′ζ′"
    uzeta :: Aphys
    "small-scale v′ζ′"
    vzeta :: Aphys
    "small-scale relative vorticity"
     zeta :: Aphys
    "large-scale relative vorticity"
     Zeta :: Aphys
    "small-scale relative vorticity"
      psi :: Aphys
    "large-scale relative vorticity"
      Psi :: Aphys
    "small-scale nonlinear term"
       Nz :: Atrans
    "large-scale nonlinear term"
       NZ :: Atrans
    "Fourier transform of x-component of small-scale velocity"
       uh :: Atrans
    "Fourier transform of y-component of small-scale velocity"
       vh :: Atrans
    "Fourier transform of x-component of large-scale velocity"
       Uh :: Atrans
    "Fourier transform of small-scale relative vorticity"
    zetah :: Atrans
    "Fourier transform of large-scale relative vorticity"
    Zetah :: Atrans
    "Fourier transform of small-scale relative vorticity"
     psih :: Atrans
    "Fourier transform of large-scale relative vorticity"
     Psih :: Atrans
    "Fourier transform of forcing"
       Fh :: F
    "`sol` at previous time-step"
  prevsol :: P
end

const DecayingVars = Vars{<:AbstractArray, <:AbstractArray, Nothing, Nothing}
const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing}
const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray}

"""
    DecayingVars(dev, grid)

Return the vars for unforced two-dimensional quasi-linear barotropic QG problem on device `dev`
and with `grid`.
"""
function DecayingVars(dev::Dev, grid::AbstractGrid) where Dev
  T = eltype(grid)
  @devzeros Dev T (grid.nx, grid.ny) u v U uzeta vzeta zeta Zeta psi Psi
  @devzeros Dev Complex{T} (grid.nkr, grid.nl) N NZ uh vh Uh zetah Zetah psih Psih

  return Vars(u, v, U, uzeta, vzeta, zeta, Zeta, psi, Psi, N, NZ, uh, vh, Uh, zetah, Zetah, psih, Psih, nothing, nothing)
end

"""
    ForcedVars(dev, grid)

Return the `vars` for forced two-dimensional quasi-linear barotropic QG problem on device
`dev` and with `grid`.
"""
function ForcedVars(dev::Dev, grid::AbstractGrid) where Dev
  T = eltype(grid)
  @devzeros Dev T (grid.nx, grid.ny) u v U uzeta vzeta zeta Zeta psi Psi
  @devzeros Dev Complex{T} (grid.nkr, grid.nl) N NZ uh vh Uh zetah Zetah psih Psih Fh

  return Vars(u, v, U, uzeta, vzeta, zeta, Zeta, psi, Psi, N, NZ, uh, vh, Uh, zetah, Zetah, psih, Psih, Fh, nothing)
end

"""
    StochasticForcedVars(dev, grid)

Return the `vars` for stochastically forced two-dimensional quasi-linear barotropic QG problem
on device `dev` and with `grid`.
"""
function StochasticForcedVars(dev::Dev, grid::AbstractGrid) where Dev
  T = eltype(grid)
  @devzeros Dev T (grid.nx, grid.ny) u v U uzeta vzeta zeta Zeta psi Psi
  @devzeros Dev Complex{T} (grid.nkr, grid.nl) N NZ uh vh Uh zetah Zetah psih Psih Fh prevsol

  return Vars(u, v, U, uzeta, vzeta, zeta, Zeta, psi, Psi, N, NZ, uh, vh, Uh, zetah, Zetah, psih, Psih, Fh, prevsol)
end

# -------
# Solvers
# -------

"""
    calcN_advection!(N, sol, t, clock, vars, params, grid)

Calculate the Fourier transform of the advection term for quasi-linear barotropic QG dynamics.
"""
function calcN_advection!(N, sol, t, clock, vars, params, grid)
  @. vars.zetah = sol
  CUDA.@allowscalar @. vars.zetah[1, :] = 0
  @. vars.Zetah = sol
  CUDA.@allowscalar @. vars.Zetah[2:end, :] = 0

  @. vars.uh =  im * grid.l  * grid.invKrsq * vars.zetah
  @. vars.vh = -im * grid.kr * grid.invKrsq * vars.zetah
  @. vars.Uh =  im * grid.l  * grid.invKrsq * vars.Zetah

  ldiv!(vars.zeta, grid.rfftplan, vars.zetah)
  ldiv!(vars.u, grid.rfftplan, vars.uh)
  ldiv!(vars.v, grid.rfftplan, vars.vh)

  ldiv!(vars.Zeta, grid.rfftplan, vars.Zetah)
  ldiv!(vars.U, grid.rfftplan, vars.Uh)

  @. vars.uzeta = vars.u * vars.zeta           # u*ζ
  @. vars.vzeta = vars.v * vars.zeta           # v*ζ

  mul!(vars.uh, grid.rfftplan, vars.uzeta)     # \hat{u*q}
  @. vars.NZ = -im * grid.kr * vars.uh         # -∂[u*q]/∂x
  mul!(vars.vh, grid.rfftplan, vars.vzeta)     # \hat{v*q}
  @. vars.NZ += - im * grid.l * vars.vh        # -∂[v*q]/∂y
  CUDA.@allowscalar @. vars.NZ[2:end, :] = 0

  @. vars.U = vars.U * vars.zeta                # U*ζ
  @. vars.u = vars.u * vars.Zeta                # u*Ζ
  @. vars.v = vars.v * vars.Zeta                # v*Ζ

  mul!(vars.uh, grid.rfftplan, vars.U + vars.u) # \hat{U*ζ + u*Ζ}
  @. vars.Nz = -im * grid.kr*vars.uh            # -∂[U*ζ + u*Ζ]/∂x
  mul!(vars.vh, grid.rfftplan, vars.v)          # \hat{v*Z}
  @. vars.Nz += - im * grid.l*vars.vh           # -∂[v*Z]/∂y
  CUDA.@allowscalar @. vars.Nz[1, :] = 0

  @. N = vars.NZ + vars.Nz

  return nothing
end

"""
    calcN!(N, sol, t, clock, vars, params, grid)

Calculate the nonlinear term, that is the advection term and the forcing,

```math
N = - \\widehat{𝖩(ψ, ζ + η)}^{\\mathrm{QL}} + F̂ .
```
"""
function calcN!(N, sol, t, clock, vars, params, grid)
  dealias!(sol, grid)

  calcN_advection!(N, sol, t, clock, vars, params, grid)

  addforcing!(N, sol, t, clock, vars, params, grid)

  return nothing
end

"""
    addforcing!(N, sol, t, clock, vars, params, grid)

When the problem includes forcing, calculate the forcing term ``F̂`` and add it to the
nonlinear term ``N``.
"""
addforcing!(N, sol, t, cl, vars::Vars, params, grid) = nothing

function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid)
  params.calcF!(vars.Fh, sol, t, clock, vars, params, grid)
  @. N += vars.Fh

  return nothing
end

function addforcing!(N, sol, t, clock, vars::StochasticForcedVars, params, grid)
  if t == clock.t # not a substep
    @. vars.prevsol = sol # sol at previous time-step is needed to compute budgets for stochastic forcing
    params.calcF!(vars.Fh, sol, t, clock, vars, params, grid)
  end
  @. N += vars.Fh

  return nothing
end

# ----------------
# Helper functions
# ----------------

"""
    updatevars!(sol, vars, params, grid)
    updatevars!(prob)

Update the `vars` of a problem `prob` that has `grid` and `params` with the solution in `sol`.
"""
function updatevars!(sol, vars, params, grid)
  dealias!(sol, grid)

  CUDA.@allowscalar sol[1, 1] = 0
  @. vars.zetah = sol
  CUDA.@allowscalar @. vars.zetah[1, :] = 0
  @. vars.Zetah = sol
  CUDA.@allowscalar @. vars.Zetah[2:end, :] = 0

  @. vars.Psih = -vars.Zetah * grid.invKrsq
  @. vars.psih = -vars.zetah * grid.invKrsq
  @. vars.uh = -im * grid.l  * vars.psih
  @. vars.vh =  im * grid.kr * vars.psih
  @. vars.Uh =  im * grid.l  * vars.Zetah * grid.invKrsq

  ldiv!(vars.zeta, grid.rfftplan, deepcopy(vars.zetah))
  ldiv!(vars.Zeta, grid.rfftplan, deepcopy(vars.Zetah))
  ldiv!(vars.psi, grid.rfftplan, vars.psih)
  ldiv!(vars.Psi, grid.rfftplan, vars.Psih)
  ldiv!(vars.u, grid.rfftplan, deepcopy(vars.uh))
  ldiv!(vars.v, grid.rfftplan, deepcopy(vars.vh))
  ldiv!(vars.U, grid.rfftplan, deepcopy(vars.Uh))

  return nothing
end

updatevars!(prob) = updatevars!(prob.sol, prob.vars, prob.params, prob.grid)

"""
    set_zeta!(prob, zeta)
    set_zeta!(sol, vars, grid, zeta)

Set the solution `sol` as the transform of `zeta` and update variables `vars` on the `grid`.
"""
function set_zeta!(sol, vars, params, grid, zeta)
  mul!(vars.zetah, grid.rfftplan, zeta)
  CUDA.@allowscalar vars.zetah[1, 1] = 0.0
  @. sol = vars.zetah

  updatevars!(sol, vars, params, grid)
  return nothing
end

set_zeta!(prob, zeta) = set_zeta!(prob.sol, prob.vars, prob.params, prob.grid, zeta)

"""
    energy(sol, grid)
    energy(prob)

Return the domain-averaged kinetic energy of `sol`.
"""
@inline energy(sol, grid::AbstractGrid) =
  0.5 * (parsevalsum2(grid.kr .* grid.invKrsq .* sol, grid)
        + parsevalsum2(grid.l .* grid.invKrsq .* sol, grid)) / (grid.Lx * grid.Ly)

energy(prob) = energy(prob.sol, prob.grid)

"""
    enstrophy(sol, grid, vars)
    enstrophy(prob)

Return the domain-averaged enstrophy of `sol`.
"""
function enstrophy(sol, grid::AbstractGrid, vars::AbstractVars)
  @. vars.uh = sol
  CUDA.@allowscalar vars.uh[1, 1] = 0

  return 0.5 * parsevalsum2(vars.uh, grid) / (grid.Lx * grid.Ly)
end
enstrophy(prob) = enstrophy(prob.sol, prob.grid, prob.vars)

"""
    dissipation(prob)
    dissipation(sol, vars, params, grid)

Return the domain-averaged energy dissipation rate. `nν` must be >= 1.
"""
@inline function dissipation(sol, vars, params, grid)
  @. vars.uh = grid.Krsq^(params.nν - 1) * abs2(sol)
  CUDA.@allowscalar vars.uh[1, 1] = 0

  return params.ν / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid)
end

@inline dissipation(prob) = dissipation(prob.sol, prob.vars, prob.params, prob.grid)

"""
    work(prob)
    work(sol, vars, params, grid)

Return the domain-averaged rate of work of energy by the forcing, `params.Fh`.
"""
@inline function work(sol, vars::ForcedVars, grid)
  @. vars.uh = grid.invKrsq * sol * conj(vars.Fh)

  return 1 / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid)
end

@inline function work(sol, vars::StochasticForcedVars, grid)
  @. vars.uh = grid.invKrsq * (vars.prevsol + sol) / 2 * conj(vars.Fh)

  return 1 / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid)
end

@inline work(prob) = work(prob.sol, prob.vars, prob.grid)

"""
    drag(prob)
    drag(sol, vars, params, grid)

Return the extraction of domain-averaged energy by drag `μ`.
"""
@inline function drag(prob)
  sol, vars, params, grid = prob.sol, prob.vars, prob.params, prob.grid
  @. vars.uh = grid.invKrsq * abs2(sol)
  CUDA.@allowscalar vars.uh[1, 1] = 0

  return params.μ / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid)
end

end # module


text:
__precompile__()

module ExplicitQPs

using JuMP
using JuMP: GenericAffExpr, AffExpr

# See Bemporad, "A Survey on Explicit Model Predictive Control",
# section 2.2 (page 353)

getcol(x::Variable) = x.col

struct ExplicitSolution{T, M <: AbstractMatrix{T}}
    model::Model
    solution::Vector{GenericAffExpr{T, Variable}}
    constant::Vector{T}
    jacobian::M
    params::Vector{T}
    variable_map::Vector{Tuple{Bool, Int}}
end

function solution(ex::ExplicitSolution, v::Variable)
    @assert v.m === ex.model
    isparam, idx = ex.variable_map[getcol(v)]
    @assert !isparam "$v is a parameter, so you can access its value with getparameter()"
    ex.solution[idx]
end

function parameter(ex::ExplicitSolution, v::Variable)
    @assert v.m === ex.model
    isparam, idx = ex.variable_map[getcol(v)]
    @assert isparam "$v is not a parameter, so you can access its value with getsolution()"
    ex.params[idx]
end

function gradient(ex::ExplicitSolution, v::Variable)
    @assert v.m === ex.model
    isparam, idx = ex.variable_map[getcol(v)]
    @assert !isparam "cannot extract a jacobian for a parameter"
    ex.jacobian[idx, :]
end

function jacobian(ex::ExplicitSolution, vars::AbstractVector{Variable})
    hcat(gradient.(ex, vars)...)'
end

function variable_map(m::Model, params::AbstractArray{Variable})
    var_index = 1
    param_index = 1
    nvars = length(m.colCat)
    var_map = Tuple{Bool, Int}[]
    for i in 1:nvars
        var = Variable(m, i)
        if var in params
            push!(var_map, (true, param_index))
            param_index += 1
        else
            push!(var_map, (false, var_index))
            var_index += 1
        end
    end
    @assert param_index == (length(params) + 1)
    @assert param_index + var_index == (nvars + 2)
    var_map
end

function active_inequalities(m::Model, params::AbstractArray{Variable}, eps=1e-6)
    nvars = length(m.colCat)
    A_active = SparseVector{Float64, Int64}[]
    b_active = Float64[]

    for i in 1:nvars
        var = Variable(m, i)
        if var in params
            continue
        end
        λ = getdual(var)
        if abs(λ) > eps
            if λ > 0
                ai = sparsevec([i], [-1.0], nvars)
                bi = -m.colLower[i]
            else
                ai = sparsevec([i], [1.0], nvars)
                bi = m.colUpper[i]
            end
            push!(A_active, ai)
            push!(b_active, bi)
        end
    end

    nconstr = length(m.linconstr)
    for (constraint, λ) in zip(m.linconstr, m.linconstrDuals)
        if abs(λ) > eps
            if λ > 0
                ai = -sparsevec([var.col for var in constraint.terms.vars], constraint.terms.coeffs, nvars)
                bi = -constraint.lb
            else
                ai = sparsevec([var.col for var in constraint.terms.vars], constraint.terms.coeffs, nvars)
                bi = constraint.ub
            end
            push!(A_active, ai)
            push!(b_active, bi)
        end
    end

    A::SparseMatrixCSC{Float64,Int} = if isempty(A_active)
        sparse(zeros(0, nvars))
    else
        hcat(A_active...)'
    end
    b::Vector{Float64} = vcat(b_active...)
    A, b
end

# Handling affine cost terms
# Let's say we have a problem of the form:
# min. 1/2 u' H u + x' F u + q' u
#  u
# s.t. G u <= W + S x
#
# We can perform a change of variables:
# let u = y + v
# where v is a constant. Then the problem becomes:
# min 1/2 (y + v)' H (y + v) + x' F (y + v) + q' (y + v)
#  y
# s.t. G (y + v) <= W + S x
#
# We can rewrite the cost as:
# min 1/2 y' H y + x' F y + q' y + v' H y + 1/2 v' H v + x' F v + q' v
#  y
# s.t. G y <= W - G v + S x
#
# and we can drop constants from the cost to get:
# min 1/2 y' H y + x' F y + q' y + v' H y
#
# If we choose v such that
# q' y + v' H y = 0
#
# then we end up with a problem in the standard explicit QP form:
# min 1/2 y' H y + x' F y
#  y
# s.t. G y <= (W - G v) + S x
#
# To do this, we need:
# q' = -v' H
# or
# H' v = -q
#
# So the procedure is:
# Solve the explicit QP with H, F, G and S unchanged, but with
# W <- W - G v
# where v = -H' \ q
# then substitute u = y + v when returning the answer

function explicit_solution(m::Model, params::AbstractArray{Variable}, eps=1e-3)
    nvars = length(m.colCat)
    Ã, W̃ = active_inequalities(m, params)

    param_cols = Set([v.col for v in params])
    isparam = collect(1:nvars) .∈ param_cols

    G̃ = Ã[:, .!isparam]
    S̃ = .-Ã[:, isparam]

    Q = sparse(getcol.(m.obj.qvars1), getcol.(m.obj.qvars2), m.obj.qcoeffs, nvars, nvars)
    Q = (Q .+ Q')

    # todo: cholesky instead?
    H = lufact(Q[.!isparam, .!isparam])

    # Handle affine cost terms by performing change of variables
    # to solve for (y + v) instead of u
    if !isempty(m.obj.aff)
        q = zeros(nvars)
        for i in eachindex(m.obj.aff.vars)
            var, coeff = m.obj.aff.vars[i], m.obj.aff.coeffs[i]
            q[var.col] += coeff
        end
        q̃ = q[.!isparam]
        ṽ = -(H \ q̃)
    else
        ṽ = zeros(size(G̃, 2))
    end
    W̃ .-= G̃ * ṽ

    F = Q[.!isparam, isparam]

    # v_test = rand(nvars)
    # x_test = v_test[isparam]
    # z_test = v_test[.!isparam]
    # Y = Q[isparam, isparam]
    # @assert 0.5 * v_test' * Q * v_test ≈ (0.5 * z_test' * H * z_test + x_test' * F' * z_test + 0.5 * x_test' * Y * x_test)

    x = getvalue(params)
    HiG = H \ full(G̃')
    HiF = H \ full(F)
    λ_active = -(G̃ * HiG) * (W̃ + (S̃ + G̃ * HiF) * x)

    GHiG = lufact(G̃ * HiG)

    jacobian = HiG * (GHiG \ (S̃ + G̃ * HiF)) - HiF
    constant = HiG * (GHiG \ full(W̃))

    # Undo change of variables from u to y
    constant .+= ṽ

    solution = [AffExpr(params, jacobian[i, :], constant[i]) for i in 1:length(constant)]
    @assert isapprox(JuMP.getvalue.(solution), constant + jacobian * x)

    ExplicitSolution(m, solution, constant, jacobian, x, variable_map(m, params))
end

end

max_stars_repo_name: rdeits/ExplicitQPs

max_issues_repo_name: rdeits/ExplicitQPs

max_forks_repo_name: rdeits/ExplicitQPs


text:
using Test
using JustSayIt

functions = [:click_double, :click_left, :click_middle, :click_right, :click_triple, :press_left, :release_left]

@testset "$(basename(@__FILE__))" begin
    @testset "Module Mouse available" begin
        @test @isdefined Mouse
    end
    @testset "function $f available" for f in functions
        @test f in names(Mouse; all=true)
    end;
end;


text:
module NeuralNetwork

include("network.jl")
include("SGD.jl")

export Chain, Layer, DSGD!, SGD!, GD!, test

end # module


text:
import IndexedTables: stack, unstack
export stack, unstack

function stack(t::DDataset, by = pkeynames(t); select = isa(t, DNDSparse) ? valuenames(t) : excludecols(t, by),
    variable = :variable, value = :value)

    function stackchunk(x)
        stack(x, by; select=select, variable=variable, value=value)
    end

    return fromchunks(delayedmap(stackchunk, t.chunks))
end

function unstack(::Type{D}, ::Type{T}, key, val, cols::AbstractVector{S}) where {D <:DDataset, T, S}
    D1 = D isa DNDSparse ? NDSparse : NextTable
    function unstackchunk(x, y)
        unstack(D1, T, x, y, cols)
    end
    fromchunks(delayedmap(unstackchunk, key.chunks, val.chunks))
end

function unstack(t::D, by = pkeynames(t); variable = :variable, value = :value) where {D<:DDataset}
    tgrp = groupby((value => identity,), t, by, select = (variable, value))
    S = eltype(colnames(t))
    col = column(t, variable)
    cols = S.(collect(Dagger.treereduce(delayed(union), delayedmap(unique, col.chunks))))
    T = eltype(columns(t, value))
    unstack(D, T isa Type{<:DataValue} ? eltype(T) : T, pkeys(tgrp), columns(tgrp, value), cols)
end


text:
"""
Unsupervised Clustetaring of Depth Images using Watson Mixture Mode
Tesis Doctoral: https://tel.archives-ouvertes.fr/tel-01160770v2/document
[Unsupervised Clustering of Depth Images Using Watson Mixture Model](http://ieeexplore.ieee.org/document/6976757/)
"""
module DepthWMM
export WatsonMixtureModel,
      clusterize
import Distributions: concentration

include("DiametricalClustering.jl")
include("../JCSA/ExponentialDistributions.jl")
#using Distributions
using Clustering
mutable struct WatsonMixtureModel
    dim::Integer
    weights::Vector{Float64}
    watsons::Vector{MvWatsonExponential}
end

function WatsonMixtureModel(dim::Integer, k::Integer)
    weights = zeros(k)
    watsons = MvWatsonExponential[]
    for j=1:k
        push!(watsons, MvWatsonExponential(dim))
    end
    return WatsonMixtureModel(dim, weights, watsons)
end
function dimension(w::WatsonMixtureModel)
    return w.dim
end
function calculateV( wd::MvWatsonExponential, dim::Integer, data::Vector{<:Vector{<:AbstractFloat}})
  n  = length(data)
  x  = sufficientStatistic(wd,dim)
  for i=1:n
      x = x + sufficientStatistic(data[i])
  end
  return x/n
end
function expectation(wmm::WatsonMixtureModel, x::Matrix{<:AbstractFloat})
    local k = length(wmm.weights)
    (dim, npatterns) = size(x)
    local pij = zeros(Float64, k,npatterns)
    local d = dimension(wmm)
    local a=0.5
    for i=1:npatterns
        xx = x[:,i]
        for j=1:k
            pij[j,i] = wmm.weights[j]*exp(bregmanDivergence(wmm.watsons[j],xx))
        end
        pij[:,i] = pij[:,i]/sum(pij[:,i])
    end
    return pij
end
function maximization(wmm::WatsonMixtureModel,pij::Matrix{<:AbstractFloat}, x::Matrix{<:AbstractFloat})
    k = length(wmm.weights)
    (dim,nvectors) = size(x)
    for j=1:k
        fill!(wmm.watsons[j].n,0)
        sumpij = 0
        for  i =1:nvectors
            wmm.watsons[j].n = wmm.watsons[j].n + pij[j,i]* x[:,i]
            sumpij = sumpij + pij[j,i]

        end
        wmm.watsons[j].n = wmm.watsons[j].n / (sumpij+eps())
        wmm.weights[j] = (sumpij) /nvectors
    end
end
function _clusterize(wmm::WatsonMixtureModel, x::Matrix{<:AbstractFloat})
    local k          = length(wmm.weights)
    local clustering = zeros(Integer, size(x,2))
    local d = dimension(wmm)
    for j=1:size(x,2)
        vals =[ -bregmanDivergence(wmm.watsons[i], x[:,j]) for i=1:k]
        clustering[j] = indmin(vals)
    end

    return clustering
end
function _clusterize( pij::Matrix{<:AbstractFloat}, k::Integer)
    clustering = ind2sub(size(pij),vec(findmax(pij,1)[2]))[1]
    return clustering
end
function loss(wmm::WatsonMixtureModel, indexes::Vector, x::Matrix{<:AbstractFloat})
    local d = dimension(wmm)
    local k = length(wmm.weights)
    (p,nvectors) = size(x)
    lossSum = 0
    a  = 0.5
    for i=1:length(indexes)
        xx = x[:,i]
        for j=1:k
            lossSum = lossSum + bregmanDivergence(wmm, xx, j)
        end
    end
    return log(-lossSum)
end
function computeParameters(wmm::WatsonMixtureModel)
	local k = length(wmm.weights)
	local d = dimension(wmm)
	for i=1:k
        updateParameters!(wmm.watsons[i])
    end
end
function  getExtensionMatrices(wmm::WatsonMixtureModel, X, mu)
	(d,n) = size(X)
	X_ex = sufficientStatistic(wmm.watsons[1], d, n)
	M_ex = sufficientStatistic(wmm.watsons[1], d, size(mu,2))

	for i=1:n
    	X_ex[:,i] = sufficientStatistic(wmm.watsons[1],X[:,i])
	end
	for i=1:size(mu,2)
		M_ex[:,i] = sufficientStatistic(wmm.watsons[1], mu[:,i])
	end
	return (X_ex, M_ex)
end
function clusterize(x::Matrix{<:AbstractFloat}, k::Integer)
    k1          = k
    k           = k
    wmm         = WatsonMixtureModel(size(x,1),k)
    k           = length(wmm.weights)
    d           = dimension(wmm)

    (assignments, representatives) = DiametricalClustering.clusterize(x,k)
	(xss, mu) = getExtensionMatrices(wmm, x,representatives)
    for i=1:k
        cluster = find(assignments.==i)
        wmm.weights[i] = length(cluster)/length(assignments)
		wmm.watsons[i].n = vec(mean(xss[:,cluster],2) )
	end
	computeParameters(wmm)
    prevLoss = 999999999
    currLoss = 999999998
    #while (currLoss < prevLoss)
    for i=1:300
        pij  = expectation(wmm, xss)
        maximization(wmm,pij,xss )
		computeParameters(wmm)
        println(wmm.weights)
        b = _clusterize(pij,k)

        #currLoss = loss(wmm,b,xss)
        #println(currLoss)
       imshow(reshape(b,480,640))
    end
    MM = rand(k,k)
    M = zeros(k)
    q = zeros(k)
    for i=1:k
        q[i] = kummerRatio(0.5, d/5, wmm.watsons[i].concentration)
        M[i] = log(WatsonDistribution.M(0.5,p/5, wmm.watsons[i].concentration))
    end
    for i=1:k
        ni = wmm.watsons[i].n
        ki = wmm.watsons[i].concentration

        for j=1:k
            kj = wmm.watsons[j].concentration
            nj = wmm.watsons[j].n
            Dij = dot(ni, (ni*ki ) /q[i]) - M[i] +M[j] - dot(ni, (nj*kj)/q[j])
            MM[i,j] = wmm.weights[i] * wmm.weights[j] * Dij
        end
    end
    h = hclust(Symmetric(MM), :average)
    indexes = cutree(h, k=k1)
    println(indexes)

    n = zeros(k1, p)
    pimerged = zeros(k1)
    for i = 1:maximum(indexes)
        pimerged[i] = sum(wmm.weights[indexes.==i])
        for j in find(indexes.==i)
           n[i,:]      =  n[i,:] + wmm.weights[j]*wmm.n[j,:]
        end
        n[i,:] = n[i,:] / pimerged[i]

    end
    wmm.weights = pimerged
    wmm.n      = n
    return _clusterize(wmm,xss)
end

#=facts("DepethWMM") do
    context("Concentrarion") do
	    v  = [0.01;0.05;0.2]
	    k  = norm(v)
        kn = estimateConcentration(length(v),v, tol=0.00000001)
		@fact kummerRatio(0.5, 3/2, kn, tol =0.00000001) --> roughly(k,atol=0.00001)
	end
end=#
end


text:
using MomentDistances
using Statistics: mean
using Test

@testset "distance checks" begin
    s1 = rand() * 2
    s2 = rand() * 2

    A1 = rand(4, 4)
    A2 = rand(4, 4)

    for (ms, d) in ((AbsoluteRelative(), abs(s1 - s2)),
                    (AbsoluteRelative(; relative_adjustment = Inf),
                     abs(s1 - s2) / max(abs(s1), abs(s2))),
                    (AbsoluteRelative(; relative_adjustment = 0.3),
                     abs(s1 - s2) / max(1, 0.3 * max(abs(s1), abs(s2)))))
        @test iszero(@inferred(distance(ms, s1, s1)))
        ds = @inferred(distance(ms, s1, s2))
        @test ds ≈ d

        mA = ElementwiseMean(ms)
        @test iszero(@inferred(distance(mA, A1, A1)))
        dA = @inferred distance(mA, A1, A2)
        @test dA ≈ mean(distance.(Ref(ms), A1, A2))
        @test_throws DimensionMismatch distance(mA, A1, ones(3, 7))

        w = rand()
        mN = NamedSum((s = ms, A = Weighted(mA, w)))
        nt = (s = s1, A = A1)
        @test iszero(@inferred(distance(mN, nt, nt)))
        @test @inferred(distance(mN, nt, (s = s2, A = A2))) ≈
            distance(ms, s1, s2) + distance(mA, A1, A2) * w
    end
end

@testset "inference checks" begin
    @test @inferred(distance(AbsoluteRelative(), 1.0, 1.0)) == 0
    @test @inferred(distance(Weighted(AbsoluteRelative(), 1.0), 1.0, 1.0)) == 0
    @test @inferred(distance(ElementwiseMean(AbsoluteRelative()), ones(3), ones(3))) == 0
    metric = NamedSum((a = AbsoluteRelative(), b = Weighted(AbsoluteRelative(), 2.0),
                       c = AbsoluteRelative(), d = AbsoluteRelative()))
    nt = (a = 1, b = 2, c = 3.0, d = 4.0)
    @test @inferred(distance(metric, nt, nt)) == 0
end

@testset "constructor checks" begin
    @test_throws ArgumentError Weighted(AbsoluteRelative(), -9)
end

@testset "text summaries" begin
    s1 = 1
    s2 = 2
    ms = AbsoluteRelative()
    A1 = ones(2, 2)
    A2 = reshape(0:3, 2, 2)
    mA = ElementwiseMean(AbsoluteRelative())

    @test MomentDistances.summary(ms, s1, s2) == "‹1.0 ↔ 2.0: 1.0›"
    @test MomentDistances.summary(mA, A1, A2) ==
        """
        elementwise mean distance: 1.0
          [1,1]  ‹1.0 ↔ 0.0: 1.0›
          [2,1]  ‹1.0 ↔ 1.0: 0.0›
          [1,2]  ‹1.0 ↔ 2.0: 1.0›
          [2,2]  ‹1.0 ↔ 3.0: 2.0›"""
    @test MomentDistances.summary(Weighted(mA, 0.7), A1, A2) ==
        """
        weighted: 0.7
          elementwise mean distance: 1.0
            [1,1]  ‹1.0 ↔ 0.0: 1.0›
            [2,1]  ‹1.0 ↔ 1.0: 0.0›
            [1,2]  ‹1.0 ↔ 2.0: 1.0›
            [2,2]  ‹1.0 ↔ 3.0: 2.0›"""
    @test MomentDistances.summary(NamedSum((s = ms, A = Weighted(mA, 0.7))),
                                  (s = s1, A = A1), (s = s2, A = A2)) ==
        """
        total: 1.7
          from s:
            ‹1.0 ↔ 2.0: 1.0›
          from A:
            weighted: 0.7
              elementwise mean distance: 1.0
                [1,1]  ‹1.0 ↔ 0.0: 1.0›
                [2,1]  ‹1.0 ↔ 1.0: 0.0›
                [1,2]  ‹1.0 ↔ 2.0: 1.0›
                [2,2]  ‹1.0 ↔ 3.0: 2.0›"""

    @test sprint(summarize, ms, s1, s2) == MomentDistances.summary(ms, s1, s2)
end


text:
using Parameters

@with_kw struct Lap_Para
    # mesh
    Nintervals::Int = 32        # number of intervals for your mesh
    basisorder::Int = 2         # element basis function order (1 or 2)
    xmin::Float64 = -2.0        # minimum x-value for mesh
    xmax::Float64 = 2.0         # minimum x-value for mesh
    ymin::Float64 = -1.0        # minimum x-value for mesh
    ymax::Float64 = 1.0         # minimum x-value for mesh

    # boundary conditions
    lBCval::Float64 = 1.0       # (dirichlet) value at left boundary
    rBCval::Float64 = 0.0       # (dirichlet) value at right boundary

    # visualization
    solFolder::String = "solutions"         # folder for solution to go in
    solFile::String = "laplace_example"     # name of vtk solutions output
    solName::String = "temperature"         # name of unknown
end

max_stars_repo_name: pseastham/eFEMpart

max_issues_repo_name: pseastham/eFEMpart


text:
# ---
# title: 91. Decode Ways
# id: problem91
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: String, Dynamic Programming
# link: <https://leetcode.com/problems/decode-ways/description/>
# hidden: true
# ---
#
# A message containing letters from `A-Z` is being encoded to numbers using the
# following mapping:
#
#
#
#     'A' -> 1
#     'B' -> 2
#     ...
#     'Z' -> 26
#
#
# Given a **non-empty** string containing only digits, determine the total
# number of ways to decode it.
#
# The answer is guaranteed to fit in a **32-bit** integer.
#
#
#
# **Example 1:**
#
#
#
#     Input: s = "12"
#     Output: 2
#     Explanation: It could be decoded as "AB" (1 2) or "L" (12).
#
#
# **Example 2:**
#
#
#
#     Input: s = "226"
#     Output: 3
#     Explanation: It could be decoded as "BZ" (2 26), "VF" (22 6), or "BBF" (2 2 6).
#
#
# **Example 3:**
#
#
#
#     Input: s = "0"
#     Output: 0
#     Explanation: There is no character that is mapped to a number starting with '0'. We cannot ignore a zero when we face it while decoding. So, each '0' should be part of "10" --> 'J' or "20" --> 'T'.
#
#
# **Example 4:**
#
#
#
#     Input: s = "1"
#     Output: 1
#
#
#
#
# **Constraints:**
#
#   * `1 <= s.length <= 100`
#   * `s` contains only digits and may contain leading zero(s).
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
# old versions of the code

"""
`Metropolis` - Sample Ising on a N×N square lattice with spin flip proposals
and Metropolis rates

### Fields
* `x0` - Initial of the lattice
* `β` - Inverse temperature
* `niters` - Number of iterations to run
### Optional Fields
* `J = 1.0` - Coupling constant>0
* `nsave_iters = 1` - Number of iterations between saves
"""
function Metropolis(x0, β, niters; J=1.0, nsave_iters=1)

    N = size(x0)[1];
    x = copy(x0);
    xp = copy(x0);

    x_trajectory = typeof(x0)[];

    E = Energy(x, J = J);

    𝒩 = [zeros(MVector{2,Int8}) for _ in 1:4];

    for n in 1:niters
        i = rand(1:N);
        j = rand(1:N);
        xp[i,j] *= -1;

        # compute local energy contribution
        e = 0.0;
        GetNeighbors!(𝒩, i, j, N);
        for (k,l) in 𝒩
            e += - 0.5 * J * xp[i,j] * xp[k,l];
        end

        # double once for double counting of bonds and again for the sign change
        ΔE = 2 * 2 * e;

        a = min(1, exp(-β * ΔE))
        ζ = rand();
        if ζ < a
            x[i,j] *= -1;
            E += ΔE;
        else
            # restore state if rejected
            xp[i,j] *= -1;
        end
        if(mod(n,nsave_iters) == 0)
            push!(x_trajectory,copy(x));
        end
    end

    return x_trajectory
end

"""
`Metropolis!` - Sample Ising on a N×N square lattice with spin flip proposals
and Metropolis rates.  This is the in place version.

### Fields
* `x` - State of the lattice
* `β` - Inverse temperature
* `niters` - Number of iterations to run
### Optional Fields
* 'J=1.0' - Coupling constant>0
"""
function Metropolis!(x, β, niters; J=1.0)

    N = size(x)[1];
    xp = copy(x)

    E = Energy(x, J = J);

    𝒩 = [zeros(MVector{2,Int8}) for _ in 1:4];

    for n in 1:niters
        i = rand(1:N);
        j = rand(1:N);
        xp[i,j] *= -1;

        # compute local energy contribution
        e = 0.0;
        GetNeighbors!(𝒩, i, j, N);
        for (k,l) in 𝒩
            e += - 0.5 * J * xp[i,j] * xp[k,l];
        end

        # double once for double counting of bonds and again for the sign change
        ΔE = 2 * 2 * e;

        a = min(1, exp(-β * ΔE))
        ζ = rand();
        if ζ < a
            x[i,j] *= -1;
            E += ΔE;
        else
            # restore state if rejected
            xp[i,j] *= -1;
        end
    end
    x
end

"""
`Wolff` - Sample Ising on a N×N square lattice with the Wolff cluster algorithm

### Fields
* `x0` - Initial state of the lattice
* `β` - Inverse temperature
* `niters` - Number of iterations to run
### Optional Fields
* 'J = 1.0' - Coupling constant>0
* `nsave_iters = 1` - Number of iterations between saves
"""
function Wolff(x0, β, niters; J = 1.0, nsave_iters=1)

    p = 1 - exp(-2 * β * J);

    N = size(x0)[1];
    x = copy(x0);

    x_trajectory = typeof(x0)[];

    for n in 1:niters
        i = rand(1:N);
        j = rand(1:N);

        D = x[i,j];
        x[i,j] *=-1;

        k,l = mod1(i+1,N), j;
        if (x[k,l]==D)
            if (rand() < p)
                x[k,l] *=-1;
                FlipNeighbors!(x, k, l, p, D, N)
            end
        end

        k,l = mod1(i-1,N), j;
        if (x[k,l]==D)
            if (rand() < p)
                x[k,l] *=-1;
                FlipNeighbors!(x, k, l, p, D, N)
            end
        end

        k,l = i, mod1(j+1,N);
        if (x[k,l]==D)
            if (rand() < p)
                x[k,l] *=-1;
                FlipNeighbors!(x, k, l, p, D, N)
            end
        end

        k,l = i, mod1(j-1,N);
        if (x[k,l]==D)
            if (rand() < p)
                x[k,l] *=-1;
                FlipNeighbors!(x, k, l, p, D, N)
            end
        end
        if(mod(n,nsave_iters) == 0)
            push!(x_trajectory,copy(x));
        end
    end

    return x_trajectory
end

"""
`Wolff!` - Sample Ising on a N×N square lattice with the Wolff cluster
algorithm.  In place version.

### Fields
* `x` - State of the lattice
* `β` - Inverse temperature
* `niters` - Number of iterations to run
### Optional Fields
* 'J=1.0' - Coupling constant>0
"""
function Wolff!(x, β, niters; J = 1.0)

    p = 1 - exp(-2 * β * J);

    N = size(x)[1];

    for n in 1:niters
        i = rand(1:N);
        j = rand(1:N);

        D = x[i,j];
        x[i,j] *=-1;

        k,l = mod1(i+1,N), j;
        if (x[k,l]==D)
            if (rand() < p)
                x[k,l] *=-1;
                FlipNeighbors!(x, k, l, p, D, N)
            end
        end

        k,l = mod1(i-1,N), j;
        if (x[k,l]==D)
            if (rand() < p)
                x[k,l] *=-1;
                FlipNeighbors!(x, k, l, p, D, N)
            end
        end

        k,l = i, mod1(j+1,N);
        if (x[k,l]==D)
            if (rand() < p)
                x[k,l] *=-1;
                FlipNeighbors!(x, k, l, p, D, N)
            end
        end

        k,l = i, mod1(j-1,N);
        if (x[k,l]==D)
            if (rand() < p)
                x[k,l] *=-1;
                FlipNeighbors!(x, k, l, p, D, N)
            end
        end

    end

    x
end


text:
"""
Actor-critic
------------

This implementation follows Sutton & Barto, Chapter 7.7
"Eligibility Traces for Actor-Critic Methods"
The policy update follows the scheme given in Chapter 6.6 where
the policy changes inversely proportional to the action selection probability.

See learner/actorcriticlearner.jl for the core implementation.

Note
----
get_all_puzzles(), get_all_episodes_in_puzzle(), get_data_by_puzzle_episode()
are in rl_data_toolsforalgos.jl
"""
function getLL_ActorCritic_TDLambda(data::Array{Int,2}, alpha::Float64, eta::Float64,
                                    gamma::Float64, lambda_actor::Float64,
                                    lambda_critic::Float64,
                                    temp::Float64;
                                    init_policy=nothing,
                                    init_Vs = nothing,
                                    use_replacing_trace=true,
                                    do_debug_print = false)
    data_stats, nr_states, nr_actions, nr_samples = get_data_stats_and_dimensions(data)
    if data_stats.nr_subjects != 1
        error("Invalid argument: getLL_ActorCritic_TDLambda supports only one subject")
    end

    LL_unbiased = 0.
    action_selection_probs = zeros(Float64, nr_samples)
    # signals:
    TD_record = zeros(Float64, nr_samples)
    V_state = zeros(Float64, nr_samples + data_stats.nr_rewards) # V_state = zeros(Float64, nr_samples + data_stats.nr_rewards)
    finalV_record = Dict{Int,Any}()
    final_policy_record = Dict{Int,Any}()
    Δpolicy_record = zeros(Float64, nr_samples)

    PolicyParametersImbalance_record = zeros(Float64, nr_samples)
    PolicyImbalance_record = zeros(Float64, nr_samples)
    PolicyParametersImbalanceAllstates_record = Array{Array{Float64,1}, 1}(undef, nr_samples)
    [PolicyParametersImbalanceAllstates_record[i] = zeros(Float64, nr_states) for i in 1:nr_samples]
    PolicyParametersDiffChosenVsUnchosen_record = zeros(Float64, nr_samples)
    PolicyDiffChosenVsUnchosen_record = zeros(Float64, nr_samples)

    current_puzzle_id = -1
    modulatingfactor = 1.
    learner = []
    j = 1
    for i = 1:nr_samples
        if current_puzzle_id != data[i,Int(c_puzzle)]
            # Puzzle = independent block of the experiment.
            # The code was made for general purposes, in this experiment we
            # have only ONE puzzle.
            current_puzzle_id = data[i,Int(c_puzzle)]
            if do_debug_print
              println("new puzzle: $current_puzzle_id")
            end
            if init_policy == nothing # providing init_policy w/o init_Vs is considered a bug
                learner = ActorCritic(nr_states = nr_states,
                                    nr_actions = nr_actions,
                                    α_critic = eta,
                                    α_actor = alpha,
                                    γ = gamma,
                                    λ_critic = lambda_critic,
                                    λ_actor = lambda_actor,
                                    use_replacing_trace = use_replacing_trace,
                                    do_debug_print = do_debug_print)
            else
                learner = ActorCritic(nr_states = nr_states,
                                    nr_actions = nr_actions,
                                    α_critic = eta,
                                    α_actor = alpha,
                                    γ = gamma,
                                    λ_critic = lambda_critic,
                                    λ_actor = lambda_actor,
                                    policyparameters = copy(init_policy[current_puzzle_id]),
                                    V = copy(init_Vs[current_puzzle_id]),
                                    use_replacing_trace = use_replacing_trace,
                                    do_debug_print = do_debug_print)
            end
        end
        state = data[i, Int(c_state)]
        action = data[i, Int(c_action)]
        reward::Float64= data[i, Int(c_reward)]
        next_state = data[i, Int(c_nextState)]
        done = reward>0 ? true : false

        p_actions = getSoftmaxProbs(learner.policyparameters[state,:], temp, [0., 0.])

        unbiased_p_sa = p_actions[action]
        action_selection_probs[i] = copy(unbiased_p_sa)
        LL_unbiased += log(unbiased_p_sa)

        # Track V values (used as fmri regressor)
        V_state[j] = learner.V[state]
        if done
            j+=1
            if j < (nr_samples + data_stats.nr_rewards)
                V_state[j] = learner.V[next_state] # Track V value of landing state in case it is the Goal
            end
        end
        j+=1

        PolicyParametersImbalance_record[i] = abs(learner.policyparameters[state, 1] - learner.policyparameters[state, 2])
        PolicyImbalance_record[i] = abs(p_actions[1] - p_actions[2])
        PolicyParametersImbalanceAllstates_record[i] = abs.(learner.policyparameters[:, 1] - learner.policyparameters[:, 2])
        PolicyParametersDiffChosenVsUnchosen_record[i] = learner.policyparameters[state, action] - learner.policyparameters[state, mod(action,2)+1]
        PolicyDiffChosenVsUnchosen_record[i] = p_actions[action] - p_actions[mod(action,2)+1]

        policyparameters_before = deepcopy(learner.policyparameters)

        # Actor-Critic
        update!(learner, state, action, next_state, reward, unbiased_p_sa, modulatingfactor)
        TD_record[i] = learner.td_error[1]

        finalV_record[current_puzzle_id] = learner.V
        final_policy_record[current_puzzle_id] = learner.policyparameters

        Δpolicy_record[i] = sum(abs.(learner.policyparameters .- policyparameters_before))

    end

    maxLL = LL_unbiased
    trial_by_trial_p_sa = (action_selection_probs)

    trace_type = use_replacing_trace ? "ReplTrc" : "EligTrc"
    algo = "Actor-Critic (TD-Lambda, $trace_type)"
    param_v = [alpha, eta, gamma, lambda_actor, lambda_critic, temp]
    param_str = "alpha(actor)=$(round(alpha,digits=4)), eta(critic)=$(round(eta,digits=4)),
                gamma=$(round(gamma,digits=4)),
                lambda_actor=$(round(lambda_actor,digits=4)),
                lambda_critic=$(round(lambda_critic,digits=4)),
                temp=$(round(temp,digits=4)),
                init_policy=$(repr(init_policy)),
                init_Vs=$(repr(init_Vs)), use_replacing_trace=$use_replacing_trace"

    results_dict = Dict{String,Any}(
        "trialByTrial_Vs" => V_state,
        "final_V" => finalV_record,
        "final_Policy" => final_policy_record,
        "DeltaPolicy" => Δpolicy_record,
        "PolicyParametersImbalance" => PolicyParametersImbalance_record,
        "PolicyImbalance" => PolicyImbalance_record,
        "PolicyParametersImbalanceAllstates" => PolicyParametersImbalanceAllstates_record,
        "PolicyParametersDiffChosenVsUnchosen" => PolicyParametersDiffChosenVsUnchosen_record,
        "PolicyDiffChosenVsUnchosen" => PolicyDiffChosenVsUnchosen_record,
        )

    signals = Step_by_step_signals(
    data_stats.all_subject_ids[1],
    TD_record,  # RPE
    [],  # no SPE
    trial_by_trial_p_sa,
    maxLL,
    data_stats,
    algo,
    param_v,
    param_str,
    99, # dummy value
    results_dict)
    return maxLL, signals
end


text:
# ==============================================================================
# Max-Plus Productive (flowshops).
# ==============================================================================

using LightGraphs

A = MP([1.0 -Inf; 0.0 2.0])
G1 = mpgraph(A)
G1.weights
@test nonzeros(G1.weights) == MP([1.0; 0.0; 2.0])
@test nv(G1) == 2
@test ne(G1) == 3

G2 = mpgraph(mpsparse(A))
@test G2.weights == G1.weights
@test nv(G2) == nv(G1)
@test ne(G2) == ne(G1)


text:
@testset "Add Route" begin
    @testset "Handles" begin
        app = App()

        add_get!(app, "/get") do req
            return "Test GET"
        end

        add_post!(app, "/post") do req
            return "Test POST"
        end

        add_put!(app, "/put") do req
            return "Test PUT"
        end

        add_patch!(app, "/patch") do req
            return "Test PATCH"
        end

        add_delete!(app, "/delete") do req
            return "Test DELETE"
        end

        add_options!(app, "/options") do req
            return "Test OPTIONS"
        end

        add_head!(app, "/head") do req
            return "Test HEAD"
        end

        get_req = HTTP.Request("GET", "/get")
        post_req = HTTP.Request("POST", "/post")
        put_req = HTTP.Request("PUT", "/put")
        patch_req = HTTP.Request("PATCH", "/patch")
        delete_req = HTTP.Request("DELETE", "/delete")
        options_req = HTTP.Request("OPTIONS", "/options")
        head_req = HTTP.Request("HEAD", "/head")

        get_handler = HTTP.Handlers.gethandler(app.router, get_req)
        post_handler = HTTP.Handlers.gethandler(app.router, post_req)
        put_handler = HTTP.Handlers.gethandler(app.router, put_req)
        patch_handler = HTTP.Handlers.gethandler(app.router, patch_req)
        delete_handler = HTTP.Handlers.gethandler(app.router, delete_req)
        options_handler = HTTP.Handlers.gethandler(app.router, options_req)
        head_handler = HTTP.Handlers.gethandler(app.router, head_req)

        @test get_handler.func(get_req, "") == "Test GET"
        @test post_handler.func(post_req, "") == "Test POST"
        @test put_handler.func(put_req, "") == "Test PUT"
        @test patch_handler.func(patch_req, "") == "Test PATCH"
        @test delete_handler.func(delete_req, "") == "Test DELETE"
        @test options_handler.func(options_req, "") == "Test OPTIONS"
        @test head_handler.func(head_req, "") == "Test HEAD"
    end

    @testset "Params and Query" begin
        app = App()

        add_get!(app, "/get/:x/:y") do req
            return Dict(req.params)
        end

        add_get!(app, "/get") do req
            return Dict(req.query)
        end

        add_post!(app, "/post/:x/:y") do req
            return Dict(req.params)
        end

        add_post!(app, "/post") do req
            return Dict(req.query)
        end

        add_put!(app, "/put/:x/:y") do req
            return Dict(req.params)
        end

        add_put!(app, "/put") do req
            return Dict(req.query)
        end

        add_patch!(app, "/patch/:x/:y") do req
            return Dict(req.params)
        end

        add_patch!(app, "/patch") do req
            return Dict(req.query)
        end

        add_delete!(app, "/delete/:x/:y") do req
            return Dict(req.params)
        end

        add_delete!(app, "/delete") do req
            return Dict(req.query)
        end

        add_options!(app, "/options/:x/:y") do req
            return Dict(req.params)
        end

        add_options!(app, "/options") do req
            return Dict(req.query)
        end

        add_head!(app, "/head/:x/:y") do req
            return Dict(req.params)
        end

        add_head!(app, "/head") do req
            return Dict(req.query)
        end

        get_params_req = HTTP.Request("GET", "/get/10/20")
        get_query_req = HTTP.Request("GET", "/get?x=10&y=20")

        get_params_handler = HTTP.Handlers.gethandler(app.router, get_params_req)
        get_query_handler = HTTP.Handlers.gethandler(app.router, get_query_req)

        @test get_params_handler.func(get_params_req, "") == Dict(:x => "10", :y => "20")
        @test get_query_handler.func(get_query_req, "") == Dict(:x => "10", :y => "20")

        post_params_req = HTTP.Request("POST", "/post/10/20")
        post_query_req = HTTP.Request("POST", "/post?x=10&y=20")

        post_params_handler = HTTP.Handlers.gethandler(app.router, post_params_req)
        post_query_handler = HTTP.Handlers.gethandler(app.router, post_query_req)

        @test post_params_handler.func(post_params_req, "") == Dict(:x => "10", :y => "20")
        @test post_query_handler.func(post_query_req, "") == Dict(:x => "10", :y => "20")

        put_params_req = HTTP.Request("PUT", "/put/10/20")
        put_query_req = HTTP.Request("PUT", "/put?x=10&y=20")

        put_params_handler = HTTP.Handlers.gethandler(app.router, put_params_req)
        put_query_handler = HTTP.Handlers.gethandler(app.router, put_query_req)

        @test put_params_handler.func(put_params_req, "") == Dict(:x => "10", :y => "20")
        @test put_query_handler.func(put_query_req, "") == Dict(:x => "10", :y => "20")

        patch_params_req = HTTP.Request("PATCH", "/patch/10/20")
        patch_query_req = HTTP.Request("PATCH", "/patch?x=10&y=20")

        patch_params_handler = HTTP.Handlers.gethandler(app.router, patch_params_req)
        patch_query_handler = HTTP.Handlers.gethandler(app.router, patch_query_req)

        @test patch_params_handler.func(patch_params_req, "") == Dict(:x => "10", :y => "20")
        @test patch_query_handler.func(patch_query_req, "") == Dict(:x => "10", :y => "20")

        delete_params_req = HTTP.Request("DELETE", "/delete/10/20")
        delete_query_req = HTTP.Request("DELETE", "/delete?x=10&y=20")

        delete_params_handler = HTTP.Handlers.gethandler(app.router, delete_params_req)
        delete_query_handler = HTTP.Handlers.gethandler(app.router, delete_query_req)

        @test delete_params_handler.func(delete_params_req, "") == Dict(:x => "10", :y => "20")
        @test delete_query_handler.func(delete_query_req, "") == Dict(:x => "10", :y => "20")

        options_params_req = HTTP.Request("OPTIONS", "/options/10/20")
        options_query_req = HTTP.Request("OPTIONS", "/options?x=10&y=20")

        options_params_handler = HTTP.Handlers.gethandler(app.router, options_params_req)
        options_query_handler = HTTP.Handlers.gethandler(app.router, options_query_req)

        @test options_params_handler.func(options_params_req, "") == Dict(:x => "10", :y => "20")
        @test options_query_handler.func(options_query_req, "") == Dict(:x => "10", :y => "20")

        head_params_req = HTTP.Request("HEAD", "/head/10/20")
        head_query_req = HTTP.Request("HEAD", "/head?x=10&y=20")

        head_params_handler = HTTP.Handlers.gethandler(app.router, head_params_req)
        head_query_handler = HTTP.Handlers.gethandler(app.router, head_query_req)

        @test head_params_handler.func(head_params_req, "") == Dict(:x => "10", :y => "20")
        @test head_query_handler.func(head_query_req, "") == Dict(:x => "10", :y => "20")
    end
end