text:
#=

Author: Shunsuke Hori

=#

using LS

using QuantEcon
using NLopt
using NLsolve
using Dierckx

Para=LS.Para
Planners_Allocation_Sequential = LS.Planners_Allocation_Sequential
"""
Bellman equation for the continuation of the Lucas-Stokey Problem
"""
mutable struct BellmanEquation{TP <: Para, TI <: Integer, TR <: Real}
    para::TP
    S::TI
    xbar::Array{TR}
    time_0::Bool
    z0::Array{Array}
    cFB::Vector{TR}
    nFB::Vector{TR}
    xFB::Vector{TR}
    zFB::Vector{Vector{TR}}
end

"""
Compute the planner's allocation by solving Bellman
equation.
"""
struct Planners_Allocation_Bellman{TP <: Para, TI <: Integer,
                                   TVg <: AbstractVector, TT <: Tuple}
    para::TP
    mc::MarkovChain
    S::TI
    T::BellmanEquation
    mugrid::TVg
    xgrid::TVg
    Vf::Array
    policies::TT
end

"""
Initializes the type from the calibration Para
"""
function Planners_Allocation_Bellman(para::Para, mugrid::AbstractArray)
    G = para.G
    S = size(para.Pi, 1) # number of states
    mc=MarkovChain(para.Pi)
    #now find the first best allocation
    Vf, policies, T, xgrid = solve_time1_bellman(para, mugrid)
    T.time_0 = true #Bellman equation now solves time 0 problem
    return Planners_Allocation_Bellman(para, mc, S, T, mugrid, xgrid, Vf, policies)
end

"""
Solve the time  1 Bellman equation for calibration Para and initial grid mugrid
"""
function solve_time1_bellman{TR <: Real}(para::Para{TR}, mugrid::AbstractArray)
    Pi = para.Pi
    S = size(para.Pi, 1)

    #First get initial fit from lucas stockey solution.
    #Need to change things to be ex_ante
    PP = LS.Planners_Allocation_Sequential(para)

    function incomplete_allocation(PP::LS.Planners_Allocation_Sequential,
                                    mu_::AbstractFloat, s_::Integer)
        c, n, x, V = LS.time1_value(PP, mu_)
        return c,n,dot(Pi[s_, :], x), dot(Pi[s_, :], V)
    end
    cf = Array{Function}(S, S)
    nf = Array{Function}(S, S)
    xprimef = Array{Function}(S, S)
    Vf = Vector{Function}(S)
    # cf = Array{LinInterp}(S, S)
    # nf = Array{LinInterp}(S, S)
    # xprimef = Array{LinInterp}(S, S)
    # Vf = Vector{LinInterp}(S)
    xgrid = Array{TR}(S, length(mugrid))
    for s_ in 1:S
        c = Array{TR}(length(mugrid), S)
        n = Array{TR}(length(mugrid), S)
        x = Array{TR}(length(mugrid))
        V = Array{TR}(length(mugrid))
        for (i_mu, mu) in enumerate(mugrid)
            c[i_mu, :], n[i_mu, :], x[i_mu], V[i_mu] =
                incomplete_allocation(PP, mu, s_)
        end
        xprimes = repmat(x, 1, S)
        xgrid[s_, :] = x
        for sprime = 1:S
            splc = Spline1D(x[end:-1:1], c[:, sprime][end:-1:1], k=3)
            spln = Spline1D(x[end:-1:1], n[:, sprime][end:-1:1], k=3)
            splx = Spline1D(x[end:-1:1], xprimes[:, sprime][end:-1:1], k=3)
            cf[s_, sprime] = y -> splc(y)
            nf[s_, sprime] = y -> spln(y)
            xprimef[s_, sprime] = y -> splx(y)
            # cf[s_, sprime] = LinInterp(x[end:-1:1], c[:, sprime][end:-1:1])
            # nf[s_, sprime] = LinInterp(x[end:-1:1], n[:, sprime][end:-1:1])
            # xprimef[s_, sprime] = LinInterp(x[end:-1:1], xprimes[:, sprime][end:-1:1])
        end
        splV = Spline1D(x[end:-1:1], V[end:-1:1], k=3)
        Vf[s_] = y -> splV(y)
        # Vf[s_] = LinInterp(x[end:-1:1], V[end:-1:1])
    end

    policies = [cf, nf, xprimef]

    #create xgrid
    xbar = [maximum(minimum(xgrid)), minimum(maximum(xgrid))]
    xgrid = linspace(xbar[1], xbar[2], length(mugrid))

    #Now iterate on Bellman equation
    T = BellmanEquation(para, xgrid, policies)
    diff = 1.0
    while diff > 1e-6
        PF = (i_x, x, s) -> get_policies_time1(T, i_x, x, s, Vf, xbar)
        Vfnew, policies = fit_policy_function(T, PF, xgrid)

        diff = 0.0
        for s=1:S
            diff = max(diff,maximum(abs,
                       (Vf[s].(xgrid)-Vfnew[s].(xgrid))./Vf[s].(xgrid)))
        end

        println("diff = $diff")
        Vf = copy(Vfnew)
    end

    return Vf, policies, T, xgrid
end

"""
Fits the policy functions
"""
function fit_policy_function{TF<:AbstractFloat}(
                    T::BellmanEquation,
                    PF::Function, xgrid::AbstractVector{TF})
    S = T.S
    # preallocation
    PFvec = Array{TF}(4S+1, length(xgrid))
    cf = Array{Function}(S, S)
    nf = Array{Function}(S, S)
    xprimef = Array{Function}(S, S)
    TTf = Array{Function}(S, S)
    Vf = Vector{Function}(S)
    # fit policy fuctions
    for s_ in 1:S
        for (i_x, x) in enumerate(xgrid)
            PFvec[:, i_x] = PF(i_x, x, s_)
        end
        splV = Spline1D(xgrid, PFvec[1,:], k=3)
        Vf[s_] = y -> splV(y)
        # Vf[s_] = LinInterp(xgrid, PFvec[1, :])
        for sprime=1:S
            splc = Spline1D(xgrid, PFvec[1+sprime,:], k=3)
            spln = Spline1D(xgrid, PFvec[1+S+sprime,:], k=3)
            splxprime = Spline1D(xgrid, PFvec[1+2S+sprime,:], k=3)
            splTT = Spline1D(xgrid, PFvec[1+3S+sprime,:], k=3)
            cf[s_,sprime] = y -> splc(y)
            nf[s_,sprime] = y -> spln(y)
            xprimef[s_,sprime] = y -> splxprime(y)
            TTf[s_,sprime] = y -> splTT(y)
        end
    end
    policies = (cf, nf, xprimef, TTf)
    return Vf, policies
end

"""
Computes Tau given c,n
"""
function Tau(pab::Planners_Allocation_Bellman,
            c::AbstractArray, n::AbstractArray)
    para = pab.para
    Uc, Un = para.Uc(c, n), para.Un(c, n)
    return 1+Un./(para.Theta .* Uc)
end

Tau(pab::Planners_Allocation_Bellman, c::Real, n::Real) =
    Tau(pab, [c], [n])

"""
Finds the optimal allocation given initial government debt B_ and state s_0
"""
function time0_allocation(pab::Planners_Allocation_Bellman,
                           B_::Real, s0::Integer)
    T, Vf = pab.T, pab.Vf
    xbar = T.xbar
    z0 = get_policies_time0(T, B_, s0, Vf, xbar)

    c0, n0, xprime0, T0 = z0[2], z0[3], z0[4], z0[5]
    return c0, n0, xprime0, T0
end

"""
Simulates planners policies for `T` periods
"""
function simulate{TF <: AbstractFloat}(pab::Planners_Allocation_Bellman,
                   B_::TF, s_0::Integer, T::Integer,
                   sHist::Vector=QuantEcon.simulate(pab.mc, T, init=s_0))
    para, mc, Vf, S = pab.para, pab.mc, pab.Vf, pab.S
    Pi, Uc = para.Pi, para.Uc
    cf, nf, xprimef, TTf = pab.policies

    cHist=Array{TF}(T)
    nHist=Array{TF}(T)
    Bhist=Array{TF}(T)
    xHist=Array{TF}(T)
    TauHist=Array{TF}(T)
    THist=Array{TF}(T)
    muHist=Array{TF}(T)

    #time0
    cHist[1], nHist[1], xHist[1], THist[1]  = time0_allocation(pab, B_, s_0)
    TauHist[1] = Tau(pab, cHist[1], nHist[1])[s_0]
    Bhist[1] = B_
    muHist[1] = Vf[s_0](xHist[1])

    #time 1 onward
    for t in 2:T
        s_, x, s = sHist[t-1], xHist[t-1], sHist[t]
        c = Array{TF}(S)
        n = Array{TF}(S)
        xprime = Array{TF}(S)
        TT = Array{TF}(S)
        for sprime=1:S
            c[sprime], n[sprime], xprime[sprime], TT[sprime] =
                cf[s_, sprime](x), nf[s_, sprime](x),
                xprimef[s_, sprime](x), TTf[s_, sprime](x)
        end

        Tau_val = Tau(pab, c, n)[s]
        u_c = Uc(c, n)
        Eu_c = dot(Pi[s_, :], u_c)

        muHist[t] = Vf[s](xprime[s])

        cHist[t], nHist[t], Bhist[t], TauHist[t] = c[s], n[s], x/Eu_c, Tau_val
        xHist[t], THist[t] = xprime[s], TT[s]
    end
    return cHist, nHist, Bhist, xHist, TauHist, THist, muHist, sHist
end

"""
Initializes the class from the calibration Para
"""
function BellmanEquation{TF <: AbstractFloat}(para::Para{TF},
                                              xgrid::AbstractVector{TF},
                                              policies0::Array)
    S = size(para.Pi, 1) # number of states
    xbar = [minimum(xgrid), maximum(xgrid)]
    time_0 = false
    z0 = Array{Array}(length(xgrid), S)
    cf, nf, xprimef = policies0[1], policies0[2], policies0[3]
    for s in 1:S
        for (i_x, x) in enumerate(xgrid)
            cs=Array{TF}(S)
            ns=Array{TF}(S)
            xprimes=Array{TF}(S)
            for j=1:S
                cs[j], ns[j], xprimes[j] = cf[s, j](x), nf[s, j](x), xprimef[s, j](x)
            end
            z0[i_x, s] = vcat(cs, ns, xprimes, zeros(S))
        end
    end
    cFB, nFB, IFB, xFB, zFB = LS.find_first_best(para, S, 2)
    return BellmanEquation(para, S, xbar, time_0, z0, cFB, nFB, xFB, zFB)
end

"""
Finds the optimal policies
"""
function get_policies_time1(T::BellmanEquation,
        i_x::Integer, x::Real, s_::Integer, Vf::AbstractArray{Function}, xbar::AbstractVector)
    para, S = T.para, T.S
    beta, Theta, G, Pi = para.beta, para.Theta, para.G, para.Pi
    U,Uc,Un = para.U, para.Uc, para.Un

    S_possible = sum(Pi[s_, :].>0)
    sprimei_possible = find(Pi[s_, :].>0)

    function objf(z, grad)
        c, xprime = z[1:S_possible], z[S_possible+1:2S_possible]
        n = (c+G[sprimei_possible])./Theta[sprimei_possible]
        Vprime = [Vf[sprimei_possible[si]](xprime[si]) for si in 1:S_possible]
        return -dot(Pi[s_, sprimei_possible], U.(c, n)+beta*Vprime)
    end

    function cons(out, z, grad)
        c, xprime, TT =
            z[1:S_possible], z[S_possible+1:2S_possible], z[2S_possible+1:3S_possible]
        n = (c+G[sprimei_possible])./Theta[sprimei_possible]
        u_c = Uc.(c, n)
        Eu_c = dot(Pi[s_, sprimei_possible], u_c)
        out .= x*u_c/Eu_c - u_c.*(c-TT)-Un(c, n).*n - beta*xprime
    end
    function cons_no_trans(out, z, grad)
        c, xprime = z[1:S_possible], z[S_possible+1:2S_possible]
        n = (c+G[sprimei_possible])./Theta[sprimei_possible]
        u_c = Uc.(c, n)
        Eu_c = dot(Pi[s_, sprimei_possible], u_c)
        out .= x*u_c/Eu_c - u_c.*c-Un(c, n).*n - beta*xprime
    end

    if para.transfers == true
        lb = vcat(zeros(S_possible), ones(S_possible)*xbar[1], zeros(S_possible))
        if para.n_less_than_one == true
            ub = vcat(100*ones(S_possible)-G[sprimei_possible],
                      ones(S_possible)*xbar[2], 100*ones(S_possible))
        else
            ub = vcat(100*ones(S_possible),
                      ones(S_possible)*xbar[2],
                      100*ones(S_possible))
        end
        init = vcat(T.z0[i_x, s_][sprimei_possible],
                    T.z0[i_x, s_][2S+sprimei_possible],
                    T.z0[i_x, s_][3S+sprimei_possible])
        opt = Opt(:LN_COBYLA, 3S_possible)
        equality_constraint!(opt, cons, zeros(S_possible))
    else
        lb = vcat(zeros(S_possible), ones(S_possible)*xbar[1])
        if para.n_less_than_one == true
            ub = vcat(ones(S_possible)-G[sprimei_possible], ones(S_possible)*xbar[2])
        else
            ub = vcat(ones(S_possible), ones(S_possible)*xbar[2])
        end
        init = vcat(T.z0[i_x, s_][sprimei_possible],
                    T.z0[i_x, s_][2S+sprimei_possible])
        opt = Opt(:LN_COBYLA, 2S_possible)
        equality_constraint!(opt, cons_no_trans, zeros(S_possible))
    end
    init[init .> ub] = ub[init .> ub]
    init[init .< lb] = lb[init .< lb]

    min_objective!(opt, objf)
    lower_bounds!(opt, lb)
    upper_bounds!(opt, ub)
    maxeval!(opt, 10000000)
    maxtime!(opt, 10)
    ftol_rel!(opt, 1e-8)
    ftol_abs!(opt, 1e-8)

    (minf, minx, ret) = optimize(opt, init)

    if ret != :SUCCESS && ret != :ROUNDOFF_LIMITED && ret != :MAXEVAL_REACHED &&
         ret != :FTOL_REACHED && ret != :MAXTIME_REACHED
        error("optimization failed: ret = $ret")
    end

    T.z0[i_x, s_][sprimei_possible] = minx[1:S_possible]
    T.z0[i_x, s_][S+sprimei_possible] = minx[1:S_possible]+G[sprimei_possible]
    T.z0[i_x, s_][2S+sprimei_possible] = minx[S_possible+1:2S_possible]
    if para.transfers == true
        T.z0[i_x, s_][3S+sprimei_possible] = minx[2S_possible+1:3S_possible]
    else
        T.z0[i_x, s_][3S+sprimei_possible] = zeros(S)
    end

    return vcat(-minf, T.z0[i_x, s_])
end

"""
Finds the optimal policies
"""
function get_policies_time0(T::BellmanEquation,
        B_::Real, s0::Integer, Vf::AbstractArray{Function}, xbar::AbstractVector)
    para = T.para
    beta, Theta, G = para.beta, para.Theta, para.G
    U, Uc, Un = para.U, para.Uc, para.Un

    function objf(z, grad)
        c, xprime = z[1], z[2]
        n=(c+G[s0])/Theta[s0]
        return -(U(c, n)+beta*Vf[s0](xprime))
    end

    function cons(z,grad)
        c, xprime, TT = z[1], z[2], z[3]
        n=(c+G[s0])/Theta[s0]
        return -Uc(c, n)*(c-B_-TT)-Un(c, n)*n - beta*xprime
    end
    cons_no_trans(z, grad) = cons(vcat(z, 0), grad)

    if para.transfers == true
        lb = [0.0, xbar[1], 0.0]
        if para.n_less_than_one == true
            ub = [1-G[s0], xbar[2], 100]
        else
            ub = [100.0, xbar[2], 100.0]
        end
        init = vcat(T.zFB[s0][1], T.zFB[s0][3], T.zFB[s0][4])
        init = [0.95124922, -1.15926816,  0.0]
        opt = Opt(:LN_COBYLA, 3)
        equality_constraint!(opt, cons)
    else
        lb = [0.0, xbar[1]]
        if para.n_less_than_one == true
            ub = [1-G[s0], xbar[2]]
        else
            ub = [100, xbar[2]]
        end
        init = vcat(T.zFB[s0][1], T.zFB[s0][3])
        init = [0.95124922, -1.15926816]
        opt = Opt(:LN_COBYLA, 2)
        equality_constraint!(opt, cons_no_trans)
    end
    init[init.> ub] = ub[init.> ub]
    init[init.< lb] = lb[init.< lb]

    min_objective!(opt, objf)
    lower_bounds!(opt, lb)
    upper_bounds!(opt, ub)
    maxeval!(opt, 100000000)
    maxtime!(opt, 30)

    (minf, minx, ret) = optimize(opt, init)

    if ret != :SUCCESS && ret != :ROUNDOFF_LIMITED && ret != :MAXEVAL_REACHED && ret != :FTOL_REACHED
        error("optimization failed: ret = $ret")
    end

    if para.transfers == true
        return -minf, minx[1], minx[1]+G[s0], minx[2], minx[3]
    else
        return -minf, minx[1], minx[1]+G[s0], minx[2], 0
    end
end


text:
using AES256CBC
using Base.Test

include("aes256cbc_tests.jl")

ext: jl


text:
using Test
using RenyiEntropyOfMagic
using OMEinsum

function naive_two_site(A1, A2, O1, O2)
    psi = vec(ein"ji,jk->ik"(A1, A2))
    return psi' * reshape(ein"ik,jl->ijkl"(O1, O2), 4, 4) * psi
end

function naive_three_site(A1, A2, A3, O1, O2, O3)
    psi = vec(ein"ai,abj,bk->ijk"(A1,A2,A3))
    O = reshape(ein"il,jm,kn->ijklmn"(O1,O2,O3), length(psi), length(psi))
    return psi' * O * psi
end

A1, A2 = rand(3, 2), rand(3, 2)
O1, O2 = rand(2, 2), rand(2, 2)

@test naive_two_site(A1, A2, O1, O2) ≈ boundary_vector(A1, O1)' * boundary_vector(A2, O2)

A1, A2, A3 = rand(3, 2), rand(3, 3, 2), rand(3, 2)
O1, O2, O3 = rand(2, 2), rand(2, 2), rand(2, 2)

result = boundary_vector(A1, O1)' * transfer_matrix(A2, O2) * boundary_vector(A3, O3)
target = naive_three_site(A1, A2, A3, O1, O2, O3)
@test result ≈ target

function enumerate_paulistring(L::Int)
    labels_a = Iterators.product([0:1 for _ in 1:L]...)
    labels_b = Iterators.product([0:1 for _ in 1:L]...)

    return map(Iterators.product(labels_a, labels_b)) do (as, bs)
        map(zip(as, bs)) do (a, b)
            if a == 0 && b == 0
                Consts.I2
            elseif a == 0 && b == 1
                Consts.Z
            elseif a == 1 && b == 0
                Consts.X
            else # 11
                Consts.X * Consts.Z
            end
        end
    end |> vec
end

As = rand(3, 2), rand(3, 3, 2), rand(3, 2)
sum(expectation(As, Os)^4 for Os in enumerate_paulistring(3))

using Statistics

mean(1:10^6) do _
    Os = sample_paulistring(3)
    expectation(As, Os)^4
end * 4^3
renyi_entropy(As)

exact_renyi_entropy(As, Os)
A = rand(2, 4, 2)
O = rand(2, 2)
transfer_matrix(A, O)


text:
immutable DiscreteUniform <: DiscreteUnivariateDistribution
    a::Int
    b::Int
    function DiscreteUniform(a::Real, b::Real)
        ia = int(a); ib = int(b)
        ia < ib || error("int(a) must be less than int(b)")
        new(ia, ib)
    end
end

DiscreteUniform(b::Integer) = DiscreteUniform(0, b)
DiscreteUniform() = DiscreteUniform(0, 1)

min(d::DiscreteUniform) = d.a
max(d::DiscreteUniform) = d.b

function cdf(d::DiscreteUniform, k::Real)
    k < d.a ? 0. : (k > d.b ? 1. : (ifloor(k) - d.a + 1.0) / (d.b - d.a + 1.0))
end

entropy(d::DiscreteUniform) = log(d.b - d.a + 1.0)

insupport(d::DiscreteUniform, x::Number) = isinteger(x) && d.a <= x <= d.b

function kurtosis(d::DiscreteUniform)
    n = d.b - d.a + 1.0
    return -(6.0 / 5.0) * (n^2 + 1.0) / (n^2 - 1.0)
end

mean(d::DiscreteUniform) = (d.a + d.b) / 2.0

median(d::DiscreteUniform) = (d.a + d.b) / 2.0

function mgf(d::DiscreteUniform, t::Real)
    a, b = d.a, d.b
    (exp(t * a) - exp(t * (b + 1))) / ((b - a + 1.0) * (1.0 - exp(t)))
end

function cf(d::DiscreteUniform, t::Real)
    a, b = d.a, d.b
    (exp(im * t * a) - exp(im * t * (b + 1))) / ((b - a + 1.0) * (1.0 - exp(im * t)))
end

mode(d::DiscreteUniform) = d.a
modes(d::DiscreteUniform) = [d.a:d.b]

function pdf(d::DiscreteUniform, x::Real)
    insupport(d, x) ? (1.0 / (d.b - d.a + 1)) : 0.0
end

function quantile(d::DiscreteUniform, p::Real)
    d.a + ifloor(p * (d.b - d.a + 1))
end

rand(d::DiscreteUniform) = randi(d.a, d.b)

skewness(d::DiscreteUniform) = 0.0

var(d::DiscreteUniform) = ((d.b - d.a + 1.0)^2 - 1.0) / 12.0

# Fit model

function fit_mle{T <: Real}(::Type{DiscreteUniform}, x::Array{T})
    if isempty(x)
        throw(ArgumentError("x cannot be empty."))
    end

    xmin = xmax = x[1]
    for i = 2:length(x)
        xi = x[i]
        if xi < xmin
            xmin = xi
        elseif xi > xmax
            xmax = xi
        end
    end

    DiscreteUniform(xmin, xmax)
end


text:
# stacked area plot

# Vegalite

using VegaLite, VegaDatasets

dataset("unemployment-across-industries") |>
@vlplot(
    :area,
    width=300, hieght=200,
    x={
        "yearmonth(date)",
        axis={format="%Y"}
    },
    y="sum(count)",
    color={
        :series,
        scale={scheme="category20b"}
    }
)

max_stars_repo_name: Ellon-M/visualizations

max_issues_repo_name: Ellon-M/visualizations

max_forks_repo_name: Ellon-M/visualizations


text:
P = rand(5, 6) |> OTSIN.normalize
Ps = [rand(size(P)...) |> OTSIN.normalize for i in 1:10]
M = randn(size(P)...)

@testset "fitting" begin

    @testset "regularization" begin
        @test r(M, L2()) ≈ r(M, ElasticNet(0.0)) < r(10M, L2())
        @test r(M, L1()) ≈ r(M, ElasticNet(1.0)) < r(10M, L1())
        @test r(M, ElasticNet(0.5)) ≈ 0.5r(M, L1()) + 0.5r(M, L2())
    end

    @testset "fitting" begin
        Mfit = fitM(P, fix_a=false, fix_b=true, γ=0.01)
        @test sum(Mfit) < 1e-5  # approximately 0
        Mfit = fitM(Ps, fix_a=true, fix_b=true)
        #Mfit, losses = fitMsgd(Ps, fix_a=true, fix_b=true)
    end

    @testset "uneven marginals" begin
        # a is fixed, b not
        a = rand(10) |> OTSIN.normalize
        Ps = [a * (OTSIN.normalize(rand(1, 5))) for i in 1:10]
        #Mfree, losses = fitMsgd(Ps, fix_a=false, fix_b=false)
    end

    @testset "parallel cross-entropy" begin

        @testset "fix both" begin
            ce = 0.0
            for Pᵢ in Ps
                a, b = marginals(Pᵢ)
                Qᵢ = optimaltransport(M, a, b)
                ce += relative_entropy(Pᵢ, Qᵢ)
            end
            A, B = marginals(Ps)
            @test ce ≈ parallel_cross_entropy(M, Ps, A=A, B=B)
        end

        @testset "fix a" begin
            ce = 0.0
            for Pᵢ in Ps
                a, b = marginals(Pᵢ)
                Qᵢ = optimaltransport(M, a)
                ce += relative_entropy(Pᵢ, Qᵢ)
            end
            A, B = marginals(Ps)
            @test ce ≈ parallel_cross_entropy_A(M, Ps, A=A)
        end

        @testset "fix b" begin
            ce = 0.0
            for Pᵢ in Ps
                a, b = marginals(Pᵢ)
                Qᵢ = optimaltransport(M', b)'
                ce += relative_entropy(Pᵢ, Qᵢ)
            end
            A, B = marginals(Ps)
            @test ce ≈ parallel_cross_entropy_B(M, Ps, B=B)
        end

        @testset "fix none" begin
            ce = 0.0
            for Pᵢ in Ps
                a, b = marginals(Pᵢ)
                Qᵢ = optimaltransport(M)
                ce += relative_entropy(Pᵢ, Qᵢ)
            end
            A, B = marginals(Ps)
            @test ce ≈ parallel_cross_entropy_free(M, Ps)
        end

    end

end

max_stars_repo_name: PoisotLab/OTSIN

max_issues_repo_name: PoisotLab/OTSIN

max_forks_repo_name: PoisotLab/OTSIN


text:
using ITensors

"""
    trg(T::ITensor; χmax::Int, nsteps::Int) -> κ, T

Perform the TRG algorithm on the partition function composed of the ITensor T.

The indices of T must obey come in pairs `(sₕ => sₕ')` and  `(sᵥ => sᵥ').

χmax is the maximum renormalized bond dimension.

nsteps are the number of renormalization steps performed.

The outputs are κ, the partition function per site, and the final renormalized
ITensor T.
"""
function trg(T::ITensor;
             χmax::Int,
             nsteps::Int,
             cutoff = 0.0,
             svd_alg = "divide_and_conquer")
  sₕ, sᵥ = filterinds(T; plev = 0)
  @assert hassameinds((sₕ, sₕ', sᵥ, sᵥ'), T)

  # Keep track of the partition function per site
  κ = 1.0
  for n in 1:nsteps
    Fₕ, Fₕ′ = factorize(T, (sₕ', sᵥ');
                       ortho = "none",
                       maxdim = χmax,
                       cutoff = cutoff,
                       tags = tags(sₕ),
                       svd_alg = svd_alg)

    s̃ₕ = commonind(Fₕ, Fₕ′)
    Fₕ′ *= δ(dag(s̃ₕ), s̃ₕ')

    Fᵥ, Fᵥ′ = factorize(T, (sₕ, sᵥ');
                       ortho = "none",
                       maxdim = χmax,
                       cutoff = cutoff,
                       tags = tags(sᵥ),
                       svd_alg = svd_alg)

    s̃ᵥ = commonind(Fᵥ, Fᵥ′)
    Fᵥ′ *=  δ(dag(s̃ᵥ), s̃ᵥ')

    T = (Fₕ * δ(dag(sₕ'), sₕ)) *
        (Fᵥ * δ(dag(sᵥ'), sᵥ)) *
        (Fₕ′ * δ(dag(sₕ), sₕ')) *
        (Fᵥ′ * δ(dag(sᵥ), sᵥ'))

    sₕ, sᵥ = s̃ₕ, s̃ᵥ

    trT = abs((T * δ(sₕ, sₕ') * δ(sᵥ, sᵥ'))[])
    T = T / trT
    κ *= trT ^ (1 / 2 ^ n)
  end
  return κ, T
end


text:
module Zoo

using Photon
using Knet
using Photon.Layers
import Photon: hasgpu, KorA

include("alexnet.jl")
include("densenet.jl")
include("vgg.jl")

end


text:
# TESTSET FOR ALL DEFAULT LABEL FUNCTIONS
@testset "Default labels" begin

    #-----------------
    #   ULTIMATE DEFAULT (no specification)
    #-----------------
    @testset "Ultimate Default" begin

        # only one function to test
        @test_nowarn getDefaultLabel(Int64)
        @test_nowarn getDefaultLabel(Float64)
        @test_nowarn getDefaultLabel(Complex)
        @test_nowarn getDefaultLabel(String)
        @test_nowarn getDefaultLabel(Symbol)

    end

    #-----------------
    #   AB DEFAULT (sublattice)
    #-----------------
    @testset "AB Default" begin

        # A part
        @test_nowarn getDefaultLabelA(Int64)
        @test_nowarn getDefaultLabelA(Float64)
        @test_nowarn getDefaultLabelA(Complex)
        @test_nowarn getDefaultLabelA(String)
        @test_nowarn getDefaultLabelA(Symbol)

        # B part
        @test_nowarn getDefaultLabelB(Int64)
        @test_nowarn getDefaultLabelB(Float64)
        @test_nowarn getDefaultLabelB(Complex)
        @test_nowarn getDefaultLabelB(String)
        @test_nowarn getDefaultLabelB(Symbol)

    end

    #-----------------
    #   XYZ DEFAULT (Kitaev)
    #-----------------
    @testset "XYZ Default" begin

        # X part
        @test_nowarn getDefaultLabelX(Int64)
        @test_nowarn getDefaultLabelX(Float64)
        @test_nowarn getDefaultLabelX(Complex)
        @test_nowarn getDefaultLabelX(String)
        @test_nowarn getDefaultLabelX(Symbol)

        # Y part
        @test_nowarn getDefaultLabelY(Int64)
        @test_nowarn getDefaultLabelY(Float64)
        @test_nowarn getDefaultLabelY(Complex)
        @test_nowarn getDefaultLabelY(String)
        @test_nowarn getDefaultLabelY(Symbol)

        # Z part
        @test_nowarn getDefaultLabelZ(Int64)
        @test_nowarn getDefaultLabelZ(Float64)
        @test_nowarn getDefaultLabelZ(Complex)
        @test_nowarn getDefaultLabelZ(String)
        @test_nowarn getDefaultLabelZ(Symbol)

        # Symbol part with extra prefix
        @test_nowarn getDefaultLabelX(Symbol, "J")
        @test_nowarn getDefaultLabelY(Symbol, "J")
        @test_nowarn getDefaultLabelZ(Symbol, "J")

        # String part with extra prefix
        @test_nowarn getDefaultLabelX(String, "t")
        @test_nowarn getDefaultLabelY(String, "t")
        @test_nowarn getDefaultLabelZ(String, "t")

    end

    #-----------------
    #   N DEFAULT (labels 1...N)
    #-----------------
    @testset "N Default" begin

        # iterate over certain range of numbers
        for n in Int64[1,2,3,4,5,10,20,50,100]
            # N part for this number
            @test_nowarn getDefaultLabelN(Int64,   n)
            @test_nowarn getDefaultLabelN(Float64, n)
            @test_nowarn getDefaultLabelN(Complex, n)
            @test_nowarn getDefaultLabelN(String,  n)
            @test_nowarn getDefaultLabelN(Symbol,  n)
        end

    end

# end the testset
end


text:
function parse(obj::EphemerisFileParser, arg0::BufferedReader, arg1::JString)
    return jcall(obj, "parse", EphemerisFile, (BufferedReader, JString), arg0, arg1)
end

function parse(obj::EphemerisFileParser, arg0::JString)
    return jcall(obj, "parse", EphemerisFile, (JString,), arg0)
end


text:
# Lattices from B. Militzer, High Energy Density Phys. 21 (2016) 8–15. doi:10.1016/j.hedp.2016.09.003. Tables 1-4

fcc1 = [1 0 0; 1 -1 0; 0 0 -1];
fcc4 = [1 1 -1; 1 -1 1; 1 -1 -1];
fcc7 = [1 1 0; -1 2 -1; -1 0 2];
fcc13 = [2 1 -1; 2 -1 -2; -1 1 -2];
fcc23 = [3 -1 0; -1 3 -1; -1 0 3];
fcc26 = [2 1 0; -2 3 -1; -2 0 3];
fcc32 = [2 2 -2; 2 -2 2; 2 -2 -2];
fcc38 = [2 2 -1; -2 3 -2; -2 1 3];
fcc50 = [2 1 1; -2 4 -1; -2 -1 4];
fcc70 = [3 2 -2; -2 4 -3; -2 2 3];
fcc92 = [4 1 -2; -3 4 2; -1 3 -4];
fcc96 = [4 -4 0; -3 -3 3; 0 0 -4];
fcc100 = [4 1 -1; -2 5 -1; -1 -1 5];
fcc108 = [3 3 -3; 3 -3 3; 3 -3 -3];
fcc_cells = [fcc1, fcc4, fcc7, fcc13, fcc23, fcc26, fcc32, fcc38, fcc50, fcc70, fcc92, fcc96, fcc100, fcc108];

cubic1 = [1 0 0; 0 1 0; 0 0 -1];
cubic3 = [1 0 1; 1 1 0; 1 -1 -1];
cubic7 = [2 0 1; 1 2 0; 1 -1 -1];
cubic8 = [2 0 0; 0 2 0; 0 0 -2];
cubic13 = [2 1 1; 1 -2 -1; 0 -1 2];
cubic14 = [2 1 1; 1 -2 -1; -1 -1 2];
cubic24 = [3 0 0; 0 2 2; 0 2 -2];
cubic27 = [3 0 0; 0 3 0; 0 0 -3];
cubic33 = [3 1 0; 1 -3 -1; 0 -1 3];
cubic38 = [3 1 1; 2 -3 0; -1 -1 3];
cubic52 = [4 0 0; 0 3 2; 0 2 -3];
cubic56 = [3 2 1; 2 -3 -2; -1 2 -3];
cubic64 = [4 0 0; 0 4 0; 0 0 -4];
cubic73 = [4 1 1; 1 -4 -1; 0 -1 4];
cubic76 = [4 1 1; 1 -4 -1; -1 -1 4];
cubic93 = [3 3 1; 3 -3 -2; -1 2 -4];
cubic100 = [5 0 0; 0 4 2; 0 2 -4];
cubic_cells = [cubic1, cubic3, cubic7, cubic8, cubic13, cubic14, cubic24, cubic27, cubic33, cubic38, cubic52, cubic56, cubic64, cubic73, cubic76, cubic93, cubic100];

bcc1 = [1 0 0; 1 -1 -1; 0 -1 0];
bcc2 = [1 0 -1; 1 -1 0; 0 -1 -1];
bcc6 = [2 -1 -2; -1 2 1; -1 0 -1];
bcc9 = [2 0 -1; 1 -2 -2; -1 1 -1];
bcc16 = [2 0 -2; 2 -2 0; 0 -2 -2];
bcc21 = [3 -1 -2; -1 3 1; -1 0 -2];
bcc25 = [3 -1 -3; -2 3 1; -1 -1 -2];
bcc28 = [3 0 -1; 2 -3 -3; -1 1 -2];
bcc35 = [3 0 -2; -2 3 0; 0 -2 -3];
bcc48 = [3 0 -3; 2 -4 -2; -2 0 -2];
bcc53 = [3 1 -1; 3 -3 -4; -2 3 -1];
bcc54 = [3 0 -3; 3 -3 0; 0 -3 -3];
bcc65 = [4 0 -1; 3 -4 -4; -1 1 -3];
bcc84 = [4 0 -2; 2 -4 -4; -2 3 -2];
bcc91 = [4 0 -3; -3 4 0; 0 -3 -4];
bcc103 = [5 -1 -2; -2 5 3; 0 -1 4];
bcc107 = [5 -1 -2; 2 -5 -4; -1 2 -3];
bcc112 = [5 -1 -1; -1 5 1; -1 1 5];
bcc128 = [4 0 -4; 4 -4 0; 0 -4 -4];
bcc_cells = [bcc1, bcc2, bcc6, bcc9, bcc16, bcc21, bcc25, bcc28, bcc35, bcc48, bcc53, bcc54, bcc65, bcc84, bcc91, bcc103, bcc107, bcc112, bcc128];

hcp1 = [1 0 0; 1 1 0; 0 0 1];
hcp3 = [2 1 0; 1 -1 0; 0 0 -1];
hcp6 = [2 2 1; 1 1 -1; -1 1 0];
hcp8 = [2 0 0; 1 2 1; 1 2 -1];
hcp11 = [3 1 0; 1 2 1; 0 -2 1];
hcp14 = [2 -1 1; -2 -2 1; 0 -2 -1];
hcp16 = [3 1 0; -1 -3 0; 0 0 -2];
hcp21 = [3 1 1; 2 3 -1; -1 2 1];
hcp27 = [3 0 1; 3 3 -1; 0 -3 -1];
hcp37 = [4 1 0; -1 -4 -1; -1 -2 2];
hcp39 = [4 2 1; 1 -3 -1; -1 -2 2];
hcp42 = [3 -1 1; -3 -4 1; 0 -2 -2];
hcp53 = [4 0 1; 3 4 -1; 0 -3 -2];
hcp62 = [5 2 0; -1 -4 2; -1 -3 -2];
hcp65 = [4 0 1; 3 4 -2; 0 -3 -2];
hcp69 = [5 1 0; -2 -5 0; 0 0 -3];
hcp76 = [4 2 2; 3 4 -2; -2 3 1];
hcp80 = [5 0 0; 2 4 2; 2 4 -2];
hcp88 = [5 0 1; 3 4 -2; -1 -4 -2];
hcp90 = [6 3 0; 0 -5 0; 0 0 -3];
hcp100 = [5 0 0; 3 6 1; 1 2 -3];
hcp_cells = [hcp1, hcp3, hcp6, hcp8, hcp11, hcp14, hcp16, hcp21, hcp27, hcp37, hcp39, hcp42, hcp53, hcp62, hcp65, hcp69, hcp76, hcp80, hcp88, hcp90];


text:
module Graphs
include("graphs/digraph.jl")
include("graphs/search.jl")
include("graphs/decomp.jl")
end


text:
"""
    ESS

Elliptical slice sampling algorithm.

# Examples
```jldoctest; setup = :(Random.seed!(1))
julia> @model gdemo(x) = begin
           m ~ Normal()
           x ~ Normal(m, 0.5)
       end
gdemo (generic function with 2 methods)

julia> sample(gdemo(1.0), ESS(), 1_000) |> mean
Mean

│ Row │ parameters │ mean     │
│     │ Symbol     │ Float64  │
├─────┼────────────┼──────────┤
│ 1   │ m          │ 0.824853 │
```
"""
struct ESS{space} <: InferenceAlgorithm end

ESS() = ESS{()}()
ESS(space::Symbol) = ESS{(space,)}()

mutable struct ESSState{V<:VarInfo} <: AbstractSamplerState
    vi::V
end

function Sampler(alg::ESS, model::Model, s::Selector)
    # sanity check
    vi = VarInfo(model)
    space = getspace(alg)
    vns = _getvns(vi, s, Val(space))
    length(vns) == 1 ||
        error("[ESS] does only support one variable ($(length(vns)) variables specified)")
    for vn in vns[1]
        dist = getdist(vi, vn)
        isgaussian(dist) ||
            error("[ESS] only supports Gaussian prior distributions")
    end

    state = ESSState(vi)
    info = Dict{Symbol, Any}()

    return Sampler(alg, info, s, state)
end

isgaussian(dist) = false
isgaussian(::Normal) = true
isgaussian(::NormalCanon) = true
isgaussian(::AbstractMvNormal) = true

# always accept in the first step
function AbstractMCMC.step!(
    ::AbstractRNG,
    model::Model,
    spl::Sampler{<:ESS},
    ::Integer,
    ::Nothing;
    kwargs...
)
    return Transition(spl)
end

function AbstractMCMC.step!(
    rng::AbstractRNG,
    model::Model,
    spl::Sampler{<:ESS},
    ::Integer,
    transition;
    kwargs...
)
    # obtain mean of distribution
    vi = spl.state.vi
    vns = _getvns(vi, spl)
    μ = mapreduce(vcat, vns[1]) do vn
        dist = getdist(vi, vn)
        vectorize(dist, mean(dist))
    end

    # obtain previous sample
    f = vi[spl]

    # recompute log-likelihood in logp
    if spl.selector.tag !== :default
        runmodel!(model, vi, spl)
    end

    # sample log-likelihood threshold for the next sample
    threshold = getlogp(vi) - randexp(rng)

    # sample from the prior
    set_flag!(vi, vns[1][1], "del")
    runmodel!(model, vi, spl)
    ν = vi[spl]

    # sample initial angle
    θ = 2 * π * rand(rng)
    θmin = θ - 2 * π
    θmax = θ

    while true
        # compute proposal and apply correction for distributions with nonzero mean
        sinθ, cosθ = sincos(θ)
        a = 1 - (sinθ + cosθ)
        vi[spl] = @. f * cosθ + ν * sinθ + μ * a

        # recompute log-likelihood and check if threshold is reached
        runmodel!(model, vi, spl)
        if getlogp(vi) > threshold
            break
        end

        # shrink the bracket
        if θ < 0
            θmin = θ
        else
            θmax = θ
        end

        # sample new angle
        θ = θmin + rand(rng) * (θmax - θmin)
    end

    return Transition(spl)
end

function tilde(ctx::DefaultContext, sampler::Sampler{<:ESS}, right, vn::VarName, inds, vi)
    if vn in getspace(sampler)
        return tilde(LikelihoodContext(), SampleFromPrior(), right, vn, inds, vi)
    else
        return tilde(ctx, SampleFromPrior(), right, vn, inds, vi)
    end
end

function tilde(ctx::DefaultContext, sampler::Sampler{<:ESS}, right, left, vi)
    return tilde(ctx, SampleFromPrior(), right, left, vi)
end

function dot_tilde(ctx::DefaultContext, sampler::Sampler{<:ESS}, right, left, vn::VarName, inds, vi)
    if vn in getspace(sampler)
        return dot_tilde(LikelihoodContext(), SampleFromPrior(), right, left, vn, inds, vi)
    else
        return dot_tilde(ctx, SampleFromPrior(), right, left, vn, inds, vi)
    end
end

function dot_tilde(ctx::DefaultContext, sampler::Sampler{<:ESS}, right, left, vi)
    return dot_tilde(ctx, SampleFromPrior(), right, left, vi)
end


text:
using StanSamples, Test

import StanSamples: # test internals
    iscommentline, fields, ColVar, combined_size, StanVar, StanScalar, StanArray,
    combine_colvars, empty_values, ncols, empty_vars_values, _read_values, read_values

@testset "raw reading" begin
    @test iscommentline("# this is a comment")
    @test iscommentline(" # this is a comment")
    @test iscommentline("#")
    @test !iscommentline("fly in the ointment #")
    @test !iscommentline("99,12")
    @test !iscommentline("99,12#comment")
    @test fields("a,b,c") == ["a","b","c"]
    @test fields("") == [""] # corner cases, should not appear in a CSV produced by CmdStan
    @test fields(",") == ["",""]
end

@testset "parsing variable names" begin
    @test ColVar("a") == ColVar(:a)
    @test ColVar("b99") == ColVar(:b99)
    @test ColVar("accept_stat__") == ColVar(:accept_stat__)
    @test ColVar("a.1.2.3") == ColVar(:a, 1, 2, 3)
    @test_throws ArgumentError ColVar("a.foo")
    @test_throws ArgumentError ColVar("a.0")
    @test_throws ArgumentError ColVar("a.0.")
end

@testset "combined size" begin
    @test combined_size([CartesianIndex((i,)) for i in 1:3]) == (3,)
    @test combined_size(vec([CartesianIndex((i,j))
                             for i in 1:3, j in 1:4])) == (3,4)
    @test_throws ArgumentError combined_size([CartesianIndex((i,)) for i in [1,3]])
    @test_throws ArgumentError combined_size([CartesianIndex((i,)) for i in [2,1]])
end

@testset "parsing header" begin
    let h = ColVar.([:a, :b, :c])
        @test combine_colvars(h) == StanScalar.((:a, :b, :c))
    end
    @test combine_colvars(ColVar.(["a", "b.1", "b.2", "c"])) ==
        (StanScalar(:a), StanArray(:b, (2,)), StanScalar(:c))
    @test combine_colvars(ColVar.(["a", "b.1.1", "b.2.1", "b.1.2", "b.2.2", "c"])) ==
        (StanScalar(:a), StanArray(:b, (2, 2)), StanScalar(:c))
    @test_throws ArgumentError combine_colvars(ColVar.(["a", "b.1", "b.2.1", "c"]))
end

@testset "variable types" begin
    a = StanScalar(:A)
    b, c, d = [StanArray(arg...) for arg in [(:B,2), (:C,2,3), (:D,2,3,5)]]
    @test empty_values(a) == Vector{Float64}()
    @test ncols(a) == 1
    @test empty_values(b) == Matrix{Float64}(undef, 2, 0)
    @test ncols(b) == 2
    @test empty_values(c) == Array{Float64,3}(undef, 2, 3, 0)
    @test ncols(c) == 6
    @test empty_values(d) == Array{Float64,4}(undef, 2, 3, 5, 0)
    @test ncols(d) == 30
end

@testset "empty var named tuple" begin
    vars = (StanScalar(:A), StanArray(:B,1), StanArray(:C,1,2), StanArray(:D,1,2,3))
    @test empty_vars_values(vars) == (A = Vector{Float64}(),
                                      B = Matrix{Float64}(undef, 1, 0),
                                      C = Array{Float64,3}(undef, 1, 2, 0),
                                      D = Array{Float64,4}(undef, 1, 2, 3, 0))
end

@testset "_read_values" begin
    @test _read_values(StanScalar(:a), [1.0]) == 1.0
    @test _read_values(StanArray(:a,2), [1.0, 2.0]) == [1.0, 2.0]
    @test _read_values(StanArray(:a,2,2), 1.0:4.0) == [1.0 3.0; 2.0 4.0]
    @test_throws BoundsError _read_values(StanArray(:a, 9), 1.0:4.0)
end

@testset "read values" begin
    # line with values
    io = IOBuffer("""
1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0
  # next line is too long, one after that is incomplete, then unparsable
1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0
1.0,2.0,3.0,4.0,5.0,6.0,7.0
1.0,foo,3.0,4.0,5.0,6.0,7.0,8.0
""")
    vars = (StanScalar(:a), StanArray(:b, 3), StanArray(:c, 2, 2))
    @test sum(ncols, vars) == 8
    vars_values = empty_vars_values(vars)
    buffer = Vector{Float64}(undef, sum(ncols, vars))
    @test read_values(io, vars, vars_values, buffer)
    @test vars_values == (a = [1.0], b = reshape(Float64.(2:4), 3, 1),
                          c = reshape(Float64.(5:8), 2, 2, 1))
    # comment line
    @test !read_values(io, vars, vars_values, buffer)
    # line too long
    @test_throws DimensionMismatch read_values(io, vars, vars_values, buffer)
    # incomplete line
    @test_throws DimensionMismatch read_values(io, vars, vars_values, buffer)
    # parser error
    @test_throws ArgumentError read_values(io, vars, vars_values, buffer)
end

"Simple samples we use for testing."
const SAMPLES_SIMPLE = """
# a comment
a,b.1,b.2,c.1.1,c.2.1,c.1.2,c.2.2
1.0,2.0,3.0,4.0,5.0,6.0,7.0
8.0,9.0,10.0,11.0,12.0,13.0,14.0
# a terminating comment
"""

@testset "read samples simple" begin
    io = IOBuffer(SAMPLES_SIMPLE)
    samples = read_samples(io)
    @test samples.a == [1.0, 8.0]
    @test samples.b == permutedims([2.0 3.0;
                                    9.0 10.0])
    @test size(samples.c) == (2, 2, 2)
    @test vec(samples.c) == Float64.(vcat(4:7, 11:14))
end

@testset "read samples large" begin
    samples = read_samples(joinpath(@__DIR__, "testmodel", "test-samples-1.csv"))
    scalar_vars = [:lp__,:accept_stat__,:stepsize__,:treedepth__,:n_leapfrog__,
                   :divergent__,:energy__,:mu,:sigma, :nu]
    @test Set(keys(samples)) == Set(vcat(scalar_vars, :alpha))
    N = 1000                    # hardcoded
    for v in scalar_vars
        @test isa(samples[v], Vector{Float64})
        @test length(samples[v]) == N
    end
    α = samples[:alpha]
    @test isa(α, AbstractArray{Float64,3})
    @test size(α) == (3, 5, N)
end

@testset "empty file (no header)" begin
    mktemp() do path, io
        write(io, """
# empty file
# just comments
""")
        @test_throws ErrorException read_samples(path)
    end
end

@testset "sample matrix" begin
    io = IOBuffer(SAMPLES_SIMPLE)
    header, matrix = read_sample_matrix(io)
    @test header == ["a", "b.1", "b.2", "c.1.1", "c.2.1", "c.1.2", "c.2.2"]
    @test matrix == permutedims(reshape(1:14, :, 2))
end


text:
using EasyHamiltonian,Test

@testset "Basic Usage" begin
    kitaev_hamiltonian = HamiltonianExpression()
    kitaev_hamiltonian += -1.0 ,"Cdagup",1,"Cup",2
    kitaev_hamiltonian += -1.0 ,"Cdagup",2,"Cup",1
    kitaev_hamiltonian += -1.0 ,"Cdagdn",1,"Cdn",2
    kitaev_hamiltonian += -1.0 ,"Cdagdn",2,"Cdn",1

    N = 2
    site_type = "Electron"
    hamiltonian = Hamiltonian(kitaev_hamiltonian,site_type,N)

    @test isa(Matrix(hamiltonian),Matrix)
    @test isa(MPO(hamiltonian),MPO)
end


text:
# the minimum Julia's sleep time is 1 us # sleep(1e-3)
# Average minimal RasPi's time (500 MHz) is 2 ns

export DisplayP4, DisplayP8
export write, clear_display, return_home, function_set, display_on_off_control, entry_mode_set, set_cursor, cursor_or_display_shift

# UInt8
# RS RW (E)
# D7 D6 D5 D4 D4 D3 D2 D1 D0

abstract type AbstractDisplay end
abstract type DisplayP end

"""
    function DisplayP8(
        RS::Int,
        RW::Union{Int,Nothing},
        E::Int,
        D7::Int,
        D6::Int,
        D5::Int,
        D4::Int
        D3::Int,
        D2::Int,
        D1::Int,
        D0::Int
    )

Initialize object representing character display with parallel bus in 8 bit mode
and prepare it for writing characters. Cursor will be set to the first line and first symbol: (0,0) pisition.

## Arguments

- `RS`-`D0` : GPIO numer connected to corresponding pins.
    You have not to connect `RW` pin, `RW` is allowed to be `nothing`.
"""
struct DisplayP8 <: DisplayP
    RS::Int  # 0 : manage, 1 : transfer
    RW::Union{Int,Nothing}  # read/write pin, 0 means writing, 1 : reading
    E::Int   # clock pin
    D7::Int
    D6::Int
    D5::Int
    D4::Int
    D3::Int
    D2::Int
    D1::Int
    D0::Int
    max_rate::Int
end

"""
    function DisplayP4(
        RS::Int,
        RW::Union{Int,Nothing},
        E::Int,
        D7::Int,
        D6::Int,
        D5::Int,
        D4::Int
    )

Initialize object representing character display with parallel bus in 4 bit mode
and prepare it for writing characters. Cursor will be set to the first line and first symbol: (0,0) pisition.

## Arguments

- `RS`-`D4` : GPIO numer connected to corresponding pins.
    You have not to connect `RW` pin, `RW` is allowed to be `nothing`.
"""
struct DisplayP4 <: DisplayP
    RS::Int  # 0 : manage, 1 : transfer
    RW::Union{Int,Nothing}  # read/write pin, 0 means writing, 1 : reading
    E::Int   # clock pin
    D7::Int
    D6::Int
    D5::Int
    D4::Int
    max_rate::Int # Byte/s

    function DisplayP4(
        RS::Int,
        RW::Union{Int,Nothing},
        E::Int,
        D7::Int,
        D6::Int,
        D5::Int,
        D4::Int;
        max_rate::Int = 25_000 # Byte/s theoretical: ~ 1Byte/40us = 25000 Byte/s, actual ~ 1Byte/3ms = 333 Byte/s
    )
        #@assert max_rate <= 25_000 "max_rate must be less than 25000 Byte/s, got $max_rate"

        init_gpio()

        pin_vector = [RS, E, D7, D6, D5, D4,]
        gpio_set_mode(pin_vector, :out)
        gpio_set_mode(RW, :out)

        gpio_clear(pin_vector)
        gpio_clear(RW)

        d = new(RS, RW, E, D7, D6, D5, D4, max_rate)

        # init commands
        sleep(20e-3) # wait > 20 ms after power on
        reset(d)

        return d
    end
end

function reset(display::DisplayP)
    # set 8bit mode 3 times (from datasheet)
    gpio_clear([display.RS, display.RW])
    sleep(0.001) # T_AS > 60 ns
    _send_semibyte(display, 0b0011_0000)
    sleep(4.1e-3) # > 4.1 ms
    _send_semibyte(display, 0b0011_0000)
    sleep(0.001) # > 100 us
    _send_semibyte(display, 0b0011_0000)
    sleep(0.001) # t_wait > 37 us

    is_8bit = typeof(display) == DisplayP8

    if !is_8bit
        _send_semibyte(display, 0b0010_0000) # switch to 4bit mode
        sleep(0.001) # t_wait > 37 us
    end

    function_set(display; _8bits = is_8bit) # function set
    display_on_off_control(display) # Display on
    clear_display(display) # clear display
    entry_mode_set(display) # Entry Mode Set

    return nothing
end

"""
    function write(display::DisplayP, symbol::Char)

Write character to the current position and shift cursor to the next position.

## Arguments

- `display` : DisplayP object.
- `symbol` : character to print. The symbol code must be convertable to `UInt8`.
    Full list of available characters can be found in datasheet for the display.
    If some of symbol cannot be converted, "?" sign will be printed.
"""
function Base.write(display::DisplayP, symbol::Char)
    try
        byte = UInt8(symbol)
        write(display, byte)
    catch e
        if isa(e, InexactError)
            @warn "Symbol '$symbol' cannot be converted to UInt8"
            write(display, '?')
        else
            throw(e)
        end
    end
end

"""
    function write(display::DisplayP, string::String)

Write character string starting from the current position.

## Arguments

- `display` : DisplayP object.
- `string` : string to print. Each symbol code must be convertable to `UInt8`.
    Full list of available characters can be found in datasheet for the display.
"""
function Base.write(display::DisplayP, string::String)
    for letter in string
        write(display, letter)
    end
end

"""
    set_cursor(display::DisplayP, col::Int = 0, row::Int = 0)

Set cursor position.
__

## Arguments

- `col` : number of position: integer from `0` to `15`
- `row` : number of string: integer `0` or `1`
"""
function set_cursor(display::DisplayP, col::Int = 0, row::Int = 0)
    if row == 0
        byte = UInt8(col)
    elseif row == 1
        byte = UInt8(4 * 16 + col)
    else
        throw("row must be 0 or 1, got $row")
    end
    set_DDRAM_address(display, byte)
end

### auxilary ###

function _command_byte(display::DisplayP, byte::UInt8)
    gpio_clear([display.RS, display.RW])
    sleep(0.001) # T_AS # > 60 ns

    _send_byte(display, byte)
    sleep(0.001) # t_wait > 37 us

    return nothing
end

function _send_semibyte(display::DisplayP, byte::UInt8)
    pin_array = [
        display.D4,
        display.D5,
        display.D6,
        display.D7,
    ]

    # high order bits
    gpio_set(display.E)
    for i in 0:3
        if (byte >> (0x4 + i)) % 0x2 == 0x1
            gpio_set(pin_array[i + 1])
        else
            gpio_clear(pin_array[i + 1])
        end
    end
    sleep(0.001) # T_DSW > 195 ns
    gpio_clear(display.E)
    sleep(0.001) # T_H > 10 ns

    return nothing
end

# to send splitted by two steps
# width between two cycles t_cyc_E must be > 1000 ns, actual > 1 ms
# PW_EH (E up width) must be > 450 ns, actual > 1 ms
function _send_byte(display::DisplayP4, byte::UInt8)
    pin_array = [
        display.D4,
        display.D5,
        display.D6,
        display.D7,
    ]

    # high order bits
    gpio_set(display.E)
    for i in 0:3
        if (byte >> (0x4 + i)) % 0x2 == 0x1
            gpio_set(pin_array[i + 1])
        else
            gpio_clear(pin_array[i + 1])
        end
    end
    sleep(0.001) # T_DSW > 195 ns
    gpio_clear(display.E)
    sleep(0.001) # T_H > 10 ns

    # low order bits
    gpio_set(display.E)
    for i in 0:3
        if (byte >> i) % 0x2 == 0x1
            gpio_set(pin_array[i + 1])
        else
            gpio_clear(pin_array[i + 1])
        end
    end
    sleep(0.001) # T_DSW > 195 ns
    gpio_clear(display.E)
    sleep(0.001) # T_H > 10 ns

    return nothing
end

function _send_byte(display::DisplayP8, byte::UInt8)
    pin_array = [
        display.D0,
        display.D1,
        display.D2,
        display.D3,
        display.D4,
        display.D5,
        display.D6,
        display.D7,
    ]

    # high order bits
    gpio_set(display.E)
    for i in 0:7
        if (byte >> i) % 0x2 == 0x1
            gpio_set(pin_array[i + 1])
        else
            gpio_clear(pin_array[i + 1])
        end
    end
    sleep(0.001) # T_DSW > 195 ns
    gpio_clear(display.E)
    sleep(0.001) # T_H > 10 ns

    return nothing
end

### low level write ###

"""
    write(display::DisplayP, byte::UInt8)

Write data to CG or DDRAM

_DB7 DB6 DB5 DB4 DB4 DB3 DB2 DB1 DB0_

## Arguments

- `byte` : Byte representing symbol.
"""
function Base.write(display::DisplayP, byte::UInt8)
    gpio_set(display.RS); gpio_clear(display.RW)
    sleep(0.001) # T_AS > 60 ns

    _send_byte(display, byte)
    sleep(0.001) # t_wait # > 37 us

    return nothing
end

### low level commands ###

"""
    clear_display(display::DisplayP)

Clears entire display and sets DDRAM address 0 in address counter.

_0 0 0 0 0 0 0 1_
"""
function clear_display(display::DisplayP)
    _command_byte(display, 0b0000_0001)
    sleep(1.5e-3) # > t_wait_clear > 1500 us
end

"""
    return_home(display::DisplayP)

Sets DDRAM address 0 in address counter.
Also returns display from being shifted to original position.
DDRAM contents remain unchanged.

_0 0 0 0 0 0 1 -_
"""
function return_home(display::DisplayP)
    _command_byte(display, 0b0000_0010)
    sleep(1.5e-3) # t_wait_clear > 1500 us
end

"""
    entry_mode_set(display::DisplayP; increment::Bool = true, shift::Bool = false)

Sets cursor move direction and specifies display shift.
These operations are performed during data write and read.

_0 0 0 0 0 1 I/D S_

## Arguments

- `increment` : `true` means increment, `false` decrement
- `shift` : `true` accompanies display shift
"""
function entry_mode_set(display::DisplayP; increment::Bool = true, shift::Bool = false)
    byte = UInt8(0b0000_0100 | (increment << 1) | shift)
    _command_byte(display, byte)
end

"""
    display_on_off_control(
        display::DisplayP;
        display_on::Bool = true,
        cursor_on::Bool = true,
        blinking::Bool = false
    )

Sets entire display (D) on/off, cursor on/off (C), and blinking of cursor position character (B).

_0 0 0 0 1 D C B_

## Arguments

- `display_on` : `true` means display on, `false` is off.
- `cursor_on` : `true` means cursor on, `false` is off.
- `blinking` : `true` means blinking cursor character.
"""
function display_on_off_control(display::DisplayP; display_on::Bool = true, cursor_on::Bool = true, blinking::Bool = false)
    byte = UInt8(0b0000_1000 | (display_on << 2) | (cursor_on << 1) | blinking)
    _command_byte(display, byte)
end

"""
    cursor_shift(display::DisplayP; right::Bool = false)

Moves cursor without changing DDRAM contents.

_0 0 0 1 0 R/L - -_
"""
function cursor_shift(display::DisplayP; right::Bool = false)
    byte = UInt8(0b0001_0000 | (right << 2))
    _command_byte(display, byte)
end

"""
    display_shift(display::DisplayP; right::Bool = false)

Shifts display without changing DDRAM contents.

_# 0 0 0 1 1 R/L - -_
"""
function display_shift(display::DisplayP; right::Bool = false)
    byte = UInt8(0b0001_1000 | (right << 2))
    _command_byte(display, byte)
end

"""
    function_set(
        display::DisplayP;
        _8bits::Bool = false,
        _2lines::Bool = true,
        _10pixels::Bool = false,
        second_set::Bool = false
    )

Sets interface data length
(DL), number of display lines
(N), and character font (F)

0 0 1 DL N F P —

## Arguments

- `_8bits` : set 8 bit mode (`true`) or 4 bit model (`false`)
- `_2lines` : `true` means 2 lines, `false` means one line.
- `_10pixels` : `true` for 5x10 pixel symbols modules. Default (`false`) is 5x8.
- `second_set` : Some modules supports 2 symbols fonts. See the datasheet.
"""
function function_set(display::DisplayP; _8bits::Bool = false, _2lines::Bool = true, _10pixels::Bool = false, second_set::Bool = false)
    byte = UInt8(0b0010_0000 | (_8bits << 4) | (_2lines << 3) | (_10pixels << 2) | (second_set << 1))
    _command_byte(display, byte)
end

"""
    set_CGRAM_address(display::DisplayP, byte::UInt8)

Sets CGRAM address.
CGRAM data is sent and
received after this setting

_0 1 ACG ACG ACG ACG ACG ACG_
"""
function set_CGRAM_address(display::DisplayP, byte::UInt8)
    @assert byte <= 0b0011_1111 "Address must be between 0x00 and 0x3f, got $byte"
    _command_byte(display, 0b0100_0000 | byte)
end

"""
    set_DDRAM_address(display::DisplayP, byte::UInt8)

Sets DDRAM address.
DDRAM data is sent and
received after this setting.

_1 ADD ADD ADD ADD ADD ADD ADD_
"""
function set_DDRAM_address(display::DisplayP, byte::UInt8)
    @assert byte <= 0b0111_1111 "Address must be between 0x00 and 0x7f, got $byte"
    _command_byte(display, 0b1000_0000 | byte)
end

### low level read ###

# 0 1 BF AC AC AC AC AC AC AC
# read_busy_flag_and_address

# ? ? ? ? ? ? ? ?
# Read data from CG or DDRAM
# read


text:
###############################################################################
# Type for efficient storage of GPC realizations
type GaussianProcessSample
    knot_locs::Array
    knot_values::Array
    dim::Integer
    nknot::Integer
    iter::Integer
    warmup::Integer

    GaussianProcessSample(knot_locs::Array, iters::Integer, warmup::Integer) =
      new(knot_locs,
          Array(Float64, size(knot_locs, 1), iters),
          size(knot_locs, 2),
          size(knot_locs, 1),
          iters,
          warmup)
end

# Extract a specific realization as a GPC
GaussianProcessConvolution(GPS::GaussianProcessSample, realization::Integer) =
  new(GPS.knot_locs,
      GPS.knot_values[:, realization],
      GPS.dim,
      GPS.nknot)


text:
using RobotZoo
using StaticArrays, LinearAlgebra

model = RobotZoo.Cartpole()
n,m = size(model)
N = 11
x,u = rand(model)
dt = 0.1
z = KnotPoint(x,u,dt)

#--- Generate some constraints
# Circle Constraint
xc = SA[1,1,1]
yc = SA[1,2,3]
r  = SA[1,1,1]
cir = CircleConstraint(n, xc, yc, r)

# Goal Constraint
xf = @SVector rand(n)
goal = GoalConstraint(xf)

# Linear Constraint
p = 5
A = @SMatrix rand(p,n+m)
b = @SVector rand(p)
lin = LinearConstraint(n,m,A,b, Inequality())

# Bound Constraint
xmin = -@SVector rand(n)
xmax = +@SVector rand(n)
umin = -@SVector rand(m)
umax = +@SVector rand(m)
bnd = BoundConstraint(n,m, x_min=xmin, x_max=xmax, u_min=umin, u_max=umax)

# Dynamics Constraint
dyn = TO.DynamicsConstraint(model, N)

##--- Create the Constraint List
cons = ConstraintList(n,m,N)
add_constraint!(cons, cir, 1:N)
add_constraint!(cons, goal, N)
add_constraint!(cons, lin, 2:N-1)
add_constraint!(cons, bnd, 1:N-1)
add_constraint!(cons, dyn, 1:N-1)

##--- Create an ALConstraintSet
conset = Altro.ALConstraintSet(cons, model)
Altro.reset!(conset, SolverOptions())
@test all(conset.convals .=== conset.errvals)

cval = conset.convals[1]
@test cval isa Altro.ALConVal{typeof(cir)}
@test cval.inds == cons.inds[1]
@test size(cval.jac[1]) == (length(cir), n)
@test size(cval.jac) == (N,1)
cval.jac[1] .= 1
@test cval.jac[1] == ones(length(cir), n)
@test cval.iserr == false

cval = conset.convals[3]
@test cval isa Altro.ALConVal{typeof(lin)}
@test cval.inds == cons.inds[3]
@test size(cval.jac[1]) == (length(lin), n+m)
@test size(cval.jac) == (N-2,1)
cval.jac[2] .= 1
cval.jac[2][:,n+1:end] .= 2
@test cval.jac[2] ≈ [ones(length(lin), n) 2*ones(length(lin),m)]
@test cval.iserr == false

cval = conset.convals[end]
@test cval isa Altro.ALConVal{typeof(dyn)}
@test cval.inds == cons.inds[end]
@test size(cval.jac[1]) == (n, n+m)
@test size(cval.jac) == (N-1,2)
@test cval.iserr == false
@test cval.jac[3,2] == zeros(n,n)

# Test iteration
@test all([con for con in conset] .=== [con for con in cons])
@test length(conset) == length(cons)

##--- Test evaluation
Z = Traj([rand(n) for k = 1:N], [rand(m) for k = 1:N-1], fill(dt,N))
TO.evaluate!(conset, Z)
@test conset.errvals[end].vals[1] ≈ discrete_dynamics(RK3, model, Z[1]) - state(Z[2])
@test conset.errvals[3].vals[2] ≈ TO.evaluate(lin, Z[3])
@test conset.errvals[2].vals[1] ≈ TO.evaluate(goal, Z[end])

TO.jacobian!(conset, Z)
∇c = TO.gen_jacobian(dyn)
discrete_jacobian!(RK3, ∇c, model, Z[1])
@test conset.errvals[end].jac[1] ≈ ∇c
@test conset.errvals[end].jac[1,2] ≈ -I(n)
@test conset.errvals[2].jac[1] ≈ I(n)
∇c = TO.gen_jacobian(lin)
TO.jacobian!(∇c, lin, Z[5])
@test conset.errvals[3].jac[5] ≈ ∇c ≈ A

##--- Stats functions
vals = [conval.vals for conval in conset.convals]
viols = map(enumerate(vals)) do (i,val)
    pos(x) = x > 0 ? x : zero(x)
    if TO.sense(cons[i]) == Inequality()
        [pos.(v) for v in val]
    else
        [abs.(v) for v in val]
    end
end
c_max = map(enumerate(vals)) do (i,val)
    if TO.sense(cons[i]) == Inequality()
        maximum(maximum.(val))
    else
        maximum(norm.(val,Inf))
    end
end
@test maximum(c_max) ≈ max_violation(conset)
maximum(c_max)
@test maximum(maximum.([maximum.(viol) for viol in viols])) ≈ max_violation(conset)

p = 2
@test norm([norm(norm.(val,p),p) for val in viols],p) ≈ TO.norm_violation(conset, p)
p = 1
@test norm([norm(norm.(val,p),p) for val in viols],p) ≈ TO.norm_violation(conset, p)
p = Inf
@test norm([norm(norm.(val,p),p) for val in viols],p) ≈ TO.norm_violation(conset, p)
@test !(TO.norm_violation(conset, Inf) ≈ TO.norm_violation(conset, 2))
@test_throws ArgumentError TO.norm_violation(conset, 3)

# TODO: test penalty reset
@test Altro.max_penalty(conset) ≈ 1.0
conset.convals[3].μ[6] = @SVector fill(30, length(cons[3]))
@test Altro.max_penalty(conset) == 30
Altro.reset_penalties!(conset)
@test Altro.max_penalty(conset) ≈ 1.0
conset.convals[1].params.μ0 = 10
Altro.reset_penalties!(conset)
@test Altro.max_penalty(conset) ≈ 10

TO.evaluate!(conset, Z)
conset.convals[2].vals[1][2] = 2*max_violation(conset)
@test TO.findmax_violation(conset) == "GoalConstraint at time step 11 at index 2"
conset.convals[4].vals[2][3] = 2*max_violation(conset)
@test TO.findmax_violation(conset) == "BoundConstraint at time step 2 at x max 3"

#--- AL Updates
Altro.reset_duals!(conset)
Altro.dual_update!(conset)
@test conset.convals[2].λ[1] ≈ conset.convals[2].vals[1]
@test conset.convals[1].λ[3] ≈ clamp.(10*conset.convals[1].vals[3], 0, Inf)
λ0 = copy(conset.convals[1].λ[3])

Altro.dual_update!(conset)
@test conset.convals[1].λ[3] ≈ clamp.(λ0 .+ 10*conset.convals[1].vals[3], 0, Inf)
Altro.reset_duals!(conset)
@test conset.convals[1].λ[3] == zeros(length(cons[1]))

Altro.reset_penalties!(conset)
Altro.penalty_update!(conset)
@test Altro.max_penalty(conset) ≈ 100
@test conset.convals[2].μ[1] ≈ fill(10, length(cons[2]))

#--- Active Set
Altro.dual_update!(conset)
Altro.update_active_set!(conset)
@test conset.convals[2].active[1] == ones(n)
@test conset.convals[end].active == [ones(n) for k = 1:N-1]
for i = 1:length(conset.convals[1].active)
    for j = 1:length(cons[1])
        if conset.convals[1].active[i][j]
            @test (conset.convals[1].vals[i][j] > 0) | (conset.convals[1].λ[i][j] > 0)
        else
            @test (conset.convals[1].vals[i][j] <= 0) & (conset.convals[1].λ[i][j] <= 0)
        end
    end
end

#--- Cost
J = zeros(N)
TO.cost!(J, conset.convals[2]) #, conset.λ[2], conset.μ[2], conset.active[2])
@test J[1:N-1] == zeros(N-1)
@test J[N] ≈ (conset.convals[2].λ[1] + 0.5* conset.convals[2].μ[1] .* conset.convals[2].vals[1])' *
    conset.convals[2].vals[1]
TO.cost!(J, conset.convals[3])#, conset.λ[3], conset.μ[3], conset.active[3])
@test J[2] ≈ (conset.convals[3].λ[1] + 0.5* conset.convals[3].μ[1] .* conset.convals[3].active[1] .* conset.convals[3].vals[1])' *
    conset.convals[3].vals[1]
@test J[1] ≈ 0

E = TO.CostExpansion(n,m,N)
TO.cost_expansion!(E, conset.convals[1])
cx = conset.convals[1].jac[1]
Iμ = Diagonal(conset.convals[1].active[1] .* conset.convals[1].μ[1])
g = Iμ * conset.convals[1].vals[1] .+ conset.convals[1].λ[1]
@test E[1].Q ≈ cx'Iμ*cx + I(n)
@test E[1].q ≈ cx'g

E = TO.CostExpansion(n,m,N)
i = 3
k = 4
TO.cost_expansion!(E, conset.convals[i])#, conset.λ[i], conset.μ[i], conset.active[i])
cx = conset.convals[i].jac[k]
Iμ = Diagonal(conset.convals[i].active[k] .* conset.convals[i].μ[k])
H = cx'Iμ*cx + I
g = cx'*(Iμ * conset.convals[i].vals[k] .+ conset.convals[i].λ[k])
@test E[k+1].Q ≈ H[1:n,1:n]
@test E[k+1].H ≈ H[1:n, n+1:end]'
@test E[k+1].R ≈ H[n+1:end, n+1:end]
@test E[k+1].q ≈ g[1:n]
@test E[k+1].r ≈ g[n+1:end]

#--- Link constraints
cons2 = ConstraintList(n,m,N)
add_constraint!(cons2, cir, 1:N)
add_constraint!(cons2, goal, N)
add_constraint!(cons2, lin, 2:N-1)
add_constraint!(cons2, dyn, 1:N-1)
conset2 = Altro.ALConstraintSet(cons2, model)

max_violation(conset2) ≈ 0
@test TO.findmax_violation(conset2) == "No constraints violated"

@test conset.convals[1].con === conset2.convals[1].con
@test conset.convals[1].vals !== conset2.convals[1].vals
@test conset.convals[1].jac !== conset2.convals[1].jac

@test conset.convals[3].con === conset2.convals[3].con
@test conset.convals[3].vals !== conset2.convals[3].vals
@test conset.convals[3].jac !== conset2.convals[3].jac

@test conset.convals[end].con === conset2.convals[end].con
@test conset.convals[end].vals !== conset2.convals[end].vals
@test conset.convals[end].jac !== conset2.convals[end].jac

Altro.link_constraints!(conset2, conset)
@test conset.convals[1].vals === conset2.convals[1].vals
@test conset.convals[1].jac === conset2.convals[1].jac
@test conset.convals[3].vals === conset2.convals[3].vals
@test conset.convals[3].jac === conset2.convals[3].jac
@test conset.convals[end].vals === conset2.convals[end].vals
@test conset.convals[end].jac === conset2.convals[end].jac
@test length(conset) == 5
@test length(conset2) == 4


text:
using Dates: length
using Base: func_for_method_checked
# This file is a part of TypeDBClient.  License is MIT: https://github.com/Humans-of-Julia/TypeDBClient.jl/blob/main/LICENSE

using Behavior
using TypeDBClient
using Dates

g = TypeDBClient

function _attribute(transaction, label::String)
    res = g.match(transaction, "match \$x type $label;")
    erg = isempty(res) ? [] : collect(Iterators.flatten([values(rm.data) for rm in res]))
    return erg
end

# Scenario: Attribute with value type boolean can be created
@given("put attribute type: is-alive, with value type: boolean") do context
    res = _put_attribute_to_db(context, "is-alive", g.Proto.AttributeType_ValueType.BOOLEAN)
    @expect typeof(res) == g.Proto.Transaction_Res
    @expect typeof(res.concept_manager_res) == g.Proto.ConceptManager_Res
    @expect typeof(res.concept_manager_res.put_attribute_type_res) == g.Proto.ConceptManager_PutAttributeType_Res
end

@given("put attribute type: age, with value type: long") do context
    _put_attribute_to_db(context, "age", g.Proto.AttributeType_ValueType.LONG)
end

@given("put attribute type: score, with value type: double") do context
    _put_attribute_to_db(context, "score", g.Proto.AttributeType_ValueType.DOUBLE)
end

@given("put attribute type: birth-date, with value type: datetime") do context
    _put_attribute_to_db(context, "birth-date", g.Proto.AttributeType_ValueType.DATETIME)
end

@given("put attribute type: name, with value type: string") do context
    _put_attribute_to_db(context, "name", g.Proto.AttributeType_ValueType.STRING)
end

@given("put attribute type: email, with value type: string") do context
    _put_attribute_to_db(context, "email", g.Proto.AttributeType_ValueType.STRING)
end

@given("attribute(email) as(string) set regex: \\S+@\\S+\\.\\S+") do context
    cm = g.ConceptManager(context[:transaction])
    res = get(cm, g.AttributeType, "email")
    res_rem = g.as_remote(res, context[:transaction])
    res_regex = g.set_regex(res_rem, raw"\S+@\S+\.\S+")
end

@when("\$x = attribute(is-alive) as(boolean) put: true") do context
    attr_type = g.get(context[:cm], AttributeType, "is-alive")
    res = g.put(g.as_remote(attr_type, context[:transaction]), true)
    context[:x] = res
end

@then("attribute \$x is null: false") do context
    @expect context[:x] !== nothing
end

@then("attribute \$x has type: is-alive") do context
    @expect context[:x].type.label.name == "is-alive"
end

@then("attribute \$x has value type: boolean") do context
    @expect typeof(context[:x].value) == Bool
end

@then("attribute \$x has boolean value: true") do context
    @expect typeof(context[:x].value) == Bool
    @expect context[:x].value
end

@when("\$x = attribute(is-alive) as(boolean) get: true") do context
    @expect context[:x].type.label.name == "is-alive"
    @expect context[:x].value
end

# Scenario: Attribute with value type long can be created

@when("\$x = attribute(age) as(long) put: 21") do context
    attr_type = g.get(context[:cm], AttributeType, "age")
    res = g.put(g.as_remote(attr_type, context[:transaction]), 21)
    context[:x] = res
end

@then("attribute \$x has type: age") do context
    erg = context[:x]
    @expect erg.type.label.name == "age"
end

@then("attribute \$x has value type: long") do context
    erg = context[:x]
    @expect erg.type == g.AttributeType(g.Label("","age"), false, VALUE_TYPE.LONG)
    @expect typeof(erg.value) == Int64
end

@then("attribute \$x has long value: 21") do context
    erg = context[:x]
    @expect typeof(erg.value) == Int64
    @expect erg.value == 21
end

@when("\$x = attribute(age) as(long) get: 21") do context
    @expect context[:x].value == 21
end

#  Scenario: Attribute with value type double can be created
@when("\$x = attribute(score) as(double) put: 123.456") do context
    attr_type = g.get(context[:cm], AttributeType, "score")
    res = g.put(g.as_remote(attr_type, context[:transaction]), 123.456)
    context[:x] = res
end

@then("attribute \$x has type: score") do context
    erg = context[:x]
    @expect erg.type.label.name == "score"
end

@then("attribute \$x has value type: double") do context
    erg = context[:x]
    @expect erg.type == g.AttributeType(g.Label("","score"), false, VALUE_TYPE.DOUBLE)
    @expect typeof(erg.value) == Float64
end

@then("attribute \$x has double value: 123.456") do context
    erg = context[:x]
    @expect typeof(erg.value) == Float64
    @expect erg.value == 123.456
end

# Scenario: Attribute with value type string can be created
@when("\$x = attribute(name) as(string) put: alice") do context
    attr_type = g.get(context[:cm], AttributeType, "name")
    res = g.put(g.as_remote(attr_type, context[:transaction]), "alice")
    context[:x] = res
end

@then("attribute \$x has type: name") do context
    erg = context[:x]

    @expect erg.type.label.name == "name"
end

@then("attribute \$x has value type: string") do context
    erg = context[:x]
    @expect typeof(erg.value) == String
end

@then("attribute \$x has string value: alice") do context
    erg = context[:x]
    @expect typeof(erg.value) == String
    @expect erg.value == "alice"
end

@when("\$x = attribute(name) as(string) get: alice") do context
    @expect context[:x].value == "alice"
end

# Scenario: Attribute with value type string and satisfies a regular expression can be created

@when("\$x = attribute(email) as(string) put: alice@email.com") do context
    attr_type = g.get(context[:cm], AttributeType, "email")
    res = g.put(g.as_remote(attr_type, context[:transaction]), "alice@email.com")
    context[:x] = res
end

@then("attribute \$x has type: email") do context
    @expect context[:x].type.label.name == "email"
end

@then("attribute \$x has string value: alice@email.com") do context
    @expect context[:x].value == "alice@email.com"
end

@when("\$x = attribute(email) as(string) get: alice@email.com") do context
    @expect typeof(context[:x].value) == String
    @expect context[:x].value == "alice@email.com"
end

# Scenario: Attribute with value type string but does not satisfy a regular expression cannot be created
@when("attribute(email) as(string) put: alice-email-com; throws exception") do context
    try
        attr_type = g.get(context[:cm], AttributeType, "email")
        res = g.put(g.as_remote(attr_type, context[:transaction]), alice-email-com)
    catch ex
        @expect ex !== nothing
    end
end

#  Scenario: Attribute with value type datetime can be created
@when("\$x = attribute(birth-date) as(datetime) put: 1990-01-01 11:22:33") do context
    attr_type = g.get(context[:cm], AttributeType, "birth-date")
    res = g.put(g.as_remote(attr_type, context[:transaction]), parse(DateTime,"1990-01-01T11:22:33"))
    context[:x] = res
end

@then("attribute \$x has type: birth-date") do context
    erg = context[:x]

    @expect erg.type.label.name == "birth-date"
end

@then("attribute \$x has value type: datetime") do context
    erg = context[:x]
    @expect typeof(erg.value) == Int64
end

@then("attribute \$x has datetime value: 1990-01-01 11:22:33") do context
    erg = context[:x]
    @expect typeof(erg.value) == Int64
    @expect unix2datetime(erg.value / 1000) == parse(DateTime,"1990-01-01T11:22:33")
end

@when("\$x = attribute(birth-date) as(datetime) get: 1990-01-01 11:22:33") do context
    erg = context[:x]
    @expect erg.type.label.name == "birth-date"
    @expect typeof(erg.value) == Int64
    @expect unix2datetime(erg.value / 1000) == parse(DateTime,"1990-01-01T11:22:33")
end

# Scenario: Attribute with value type boolean can be retrieved from its type
@then("attribute(is-alive) get instances contain: \$x") do context
    type = g.get(context[:cm], AttributeType, "is-alive")
    instances_type = g.get_instances(g.as_remote(type, context[:transaction]))
    @expect in(context[:x], instances_type)
end

# Scenario: Attribute with value type long can be retrieved from its type
@then("attribute(age) get instances contain: \$x") do context
    attr = g.get(context[:cm], AttributeType, "age")
    instances_type = g.get_instances(g.as_remote(attr, context[:transaction]))
    @expect in(context[:x], instances_type)
end

# Scenario: Attribute with value type double can be retrieved from its type
@then("attribute(score) get instances contain: \$x") do context
    attr = g.get(context[:cm], AttributeType, "score")
    instances_type = g.get_instances(g.as_remote(attr, context[:transaction]))
    @expect in(context[:x], instances_type)
end

# Scenario: Attribute with value type string can be retrieved from its type
@then("attribute(name) get instances contain: \$x") do context
    attr = g.get(context[:cm], AttributeType, "name")
    instances_type = g.get_instances(g.as_remote(attr, context[:transaction]))
    @expect in(context[:x], instances_type)
end

# Scenario: Attribute with value type datetime can be retrieved from its type
@then("attribute(birth-date) get instances contain: \$x") do context
    attr = g.get(context[:cm], AttributeType, "birth-date")
    instances_type = g.get_instances(g.as_remote(attr, context[:transaction]))
    @expect in(context[:x], instances_type)
end

# Scenario: Attribute with value type boolean can be deleted
@when("delete attribute: \$x") do context
    g.delete(g.as_remote(context[:x], context[:transaction]))
end

@then("attribute \$x is deleted: true") do context
    res = g.is_deleted(g.as_remote(context[:x], context[:transaction]))
    @expect res
end

@then("attribute \$x is null: true") do context
    res = g.get(context[:cm], context[:x])
    @expect res === nothing
end

# Scenario: Attribute with value type double can be deleted
@when("\$x = attribute(score) as(double) get: 123.456") do context
    @expect context[:x].value == 123.456
end

# Scenario: Attribute with value type double is assignable and retrievable from a 'long' value
@when("\$x = attribute(score) as(double) put: 123") do context
    attr_type = g.get(context[:cm], AttributeType, "score")
    res = g.put(g.as_remote(attr_type, context[:transaction]), 123.0)
    context[:x] = res
end

@then("attribute \$x has double value: 123") do context
    erg = context[:x]
    @expect erg.value == 123
end

@when("\$x = attribute(score) as(double) get: 123") do context
    attr = g.get(context[:cm], AttributeType, "score")
    context[:x] = g.get(g.as_remote(attr,context[:transaction]), 123.0)
end


text:
__precompile__()

module AndersonMoore
# http://www.stochasticlifestyle.com/finalizing-julia-package-documentation-testing-coverage-publishing/

using Compat
using Compat.LinearAlgebra

# Walk through all the directories in src and add the files
# If you add/remove/rename a directory, mirror the change here
dirs = ["AndersonMoore", "gensys", "sparseAim", "util"]
for dir in dirs
    for (root, _, files) in walkdir(joinpath(dirname(@__FILE__), dir))
        for file in files
            include(joinpath(root, file))
        end
    end
end

# Export all functions
export exactShift!, numericShift!, shiftRight!, buildA!, augmentQ!, eigenSys!, reducedForm,
AndersonMooreAlg, sameSpan, deleteCols, deleteRows, loadSparseAim, callSparseAim, checkAM, err, gensysToAMA

end # module


text:
abstract type DeconvolutionMethod end

#######################################
# Interface for Fourier Deconvolution #
#######################################
struct FourierDeconv <: DeconvolutionMethod
    bandwidth::Float64
    Xgrid::Vector{Float64}
    num_t::Int
end

function piecewise_uniform(bin_edges::AbstractVector, weights::Weights)
    @assert length(bin_edges) == length(weights)+1
    @assert all(diff(bin_edges) .> 0)
    bin_uniforms = Uniform.(bin_edges[1:end-1], prevfloat.(bin_edges[2:end]))
    weights_vector = convert(Vector, weights)
    mixture = MixtureModel(bin_uniforms, weights_vector / sum(weights)) # mixture of uniforms with equal weights
end

# function decon(X::Vector, σ_X::Vector, bw::Float64, grid::Vector; num_t=50, fixup=true)
function decon(dm::FourierDeconv, X, ϵdistr)
    Fhat_xx = Fhat(dm.Xgrid, X, dm.num_t, dm.bandwidth, ϵdistr)
    fix_CDF!(Fhat_xx)
    return piecewise_uniform(dm.Xgrid, Weights(diff(Fhat_xx)))
    # Phat_xx = diff(Fhat_xx)
    # return Distributions.Generic(dm.Xgrid[2:end], Phat_xx)
end

##################################################
# Interface for Exponential Spline Deconvolution #
##################################################
struct EfronDeconv <: DeconvolutionMethod
    dimension::Int
    Xgrid::Vector{Float64}
    c0::Float64
end

function decon(dm::EfronDeconv, X, ϵdistr)
    d = ExpoSpline(dm.Xgrid, zeros(dm.dimension))
    d_optim = DeconvolveDistribution.decon(d, ϵdistr, X, dm.c0)
    return d_optim
end


text:
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

module TestExceptions

using MXNet
using Test

struct MXError′ <: AbstractMXError
  msg::String
end

function test_show()
  @info "AbstractMXError::Base.show"

  io = IOBuffer()
  e = MXError′("magic")
  print(io, e)
  str = String(take!(io))
  @test str == "magic"
end

@testset "Exception Test" begin
  test_show()
end

end  # module TestExceptions


text:
function DE(x, p, t, derivatives, scalings, constants, BCs)
    ν₀, ν₋, ΔRi, Riᶜ, Pr = p

    Nz, H, τ, f, g, α  = constants.Nz, constants.H, constants.τ, constants.f, constants.g, constants.α
    σ_uw, σ_vw, σ_wT = scalings.uw.σ, scalings.vw.σ, scalings.wT.σ
    μ_u, μ_v, σ_u, σ_v, σ_T = scalings.u.μ, scalings.v.μ, scalings.u.σ, scalings.v.σ, scalings.T.σ

    D_cell = derivatives.cell
    D_face = derivatives.face

    u = @view x[1:Nz]
    v = @view x[Nz + 1:2Nz]
    T = @view x[2Nz + 1:3Nz]

    ϵ = 1f-7
    ∂u∂z = D_face * u
    ∂v∂z = D_face * v
    ∂T∂z = D_face * T
    Ri = local_richardson.(∂u∂z .+ ϵ, ∂v∂z .+ ϵ, ∂T∂z .+ ϵ, H, g, α, σ_u, σ_v, σ_T)

    ν = ν₀ .+ ν₋ .* tanh_step.((Ri .- Riᶜ) ./ ΔRi)
    # ν = zeros(Float32, 33)

    ν∂u∂z = [-(BCs.uw.bottom - scalings.uw(0f0)); σ_u / σ_uw / H .* ν[2:end-1] .* ∂u∂z[2:end-1]; -(BCs.uw.top - scalings.uw(0f0))]
    ν∂v∂z = [-(BCs.vw.bottom - scalings.vw(0f0)); σ_v / σ_vw / H .* ν[2:end-1] .* ∂v∂z[2:end-1]; -(BCs.vw.top - scalings.vw(0f0))]
    ν∂T∂z = [-(BCs.wT.bottom - scalings.wT(0f0)); σ_T / σ_wT / H .* ν[2:end-1] ./ Pr .* ∂T∂z[2:end-1]; -(BCs.wT.top - scalings.wT(0f0))]

    ∂u∂t = τ / H * σ_uw / σ_u .* derivatives.cell * ν∂u∂z .+ f * τ / σ_u .* (σ_v .* v .+ μ_v)
    ∂v∂t = τ / H * σ_vw / σ_v .* derivatives.cell * ν∂v∂z .- f * τ / σ_v .* (σ_u .* u .+ μ_u)
    ∂T∂t = τ / H * σ_wT / σ_T .* derivatives.cell * ν∂T∂z

    return [∂u∂t; ∂v∂t; ∂T∂t]
end

function optimise_modified_pacanowski_philander(train_files, tsteps, timestepper, optimizers, maxiters, FILE_PATH;
                                                n_simulations, ν₀ = 1f-4, ν₋ = 1f-1, ΔRi=0.1f0, Riᶜ=0.25f0, Pr=1f0,
                                                train_gradient=true, gradient_scaling=5f-3,
                                                training_fractions=nothing)
    @info "Loading training data"
    𝒟 = WindMixing.data(train_files, scale_type=ZeroMeanUnitVarianceScaling, enforce_surface_fluxes=false)

    @info "Preparing constants and parameters"

    ν₀_scaling = 1 / ν₀
    ν₋_scaling = 1 / ν₋
    ΔRi_scaling = 1 / ΔRi
    Riᶜ_scaling = 1 / Riᶜ
    Pr_scaling = 1 / Pr

    function scale_parameter(parameter, scaling)
        return parameter * scaling
    end

    function unscale_parameter(parameter, scaling)
        return parameter / scaling
    end

    function prepare_parameters()
        Nz = length(𝒟.u.z)
        H = abs(𝒟.uw.z[end] - 𝒟.uw.z[1])
        τ = abs(𝒟.t[:,1][end] - 𝒟.t[:,1][1])
        u_scaling = 𝒟.scalings["u"]
        v_scaling = 𝒟.scalings["v"]
        T_scaling = 𝒟.scalings["T"]
        uw_scaling = 𝒟.scalings["uw"]
        vw_scaling = 𝒟.scalings["vw"]
        wT_scaling = 𝒟.scalings["wT"]

        constants = (H=H, τ=τ, Nz=Nz, f=1f-4, α=2f-4, g=9.80665f0)
        scalings = (u=u_scaling, v=v_scaling, T=T_scaling, uw=uw_scaling, vw=vw_scaling, wT=wT_scaling, parameters=[ν₀_scaling, ν₋_scaling, ΔRi_scaling, Riᶜ_scaling, Pr_scaling])
        derivatives = (cell=Float32.(Dᶜ(Nz, 1 / Nz)), face=Float32.(Dᶠ(Nz, 1 / Nz)))
        parameters = [ν₀, ν₋, ΔRi, Riᶜ, Pr]
        scaled_parameters = scale_parameter.(parameters, scalings.parameters)

        return constants, scalings, derivatives, scaled_parameters
    end

    constants, scalings, derivatives, scaled_parameters = prepare_parameters()

    Nz = constants.Nz

    D_face = derivatives.face

    n_steps = Int(length(@view(𝒟.t[:,1])) / n_simulations)

    uvT₀s = [𝒟.uvT_scaled[:,n_steps * i + tsteps[1]] for i in 0:n_simulations - 1]
    t_train = 𝒟.t[:,1][tsteps] ./ constants.τ
    tspan_train = (t_train[1], t_train[end])
    uvT_trains = [𝒟.uvT_scaled[:,n_steps * i + 1:n_steps * (i + 1)][:, tsteps] for i in 0:n_simulations - 1]

    u_trains, v_trains, T_trains = split_u.(uvT_trains, Nz), split_v.(uvT_trains, Nz), split_T.(uvT_trains, Nz)

    if train_gradient
        u_trains_gradients = [∂_∂z(sol, D_face) for sol in u_trains]
        v_trains_gradients = [∂_∂z(sol, D_face) for sol in v_trains]
        T_trains_gradients = [∂_∂z(sol, D_face) for sol in T_trains]
    end

    @info "Setting up BCs"

    @inline function BC(i)
        index = n_steps * i + tsteps[1]
        return (uw=(bottom=𝒟.uw.scaled[1,index], top=𝒟.uw.scaled[end, index]),
                vw=(bottom=𝒟.vw.scaled[1,index], top=𝒟.vw.scaled[end, index]),
                wT=(bottom=𝒟.wT.scaled[1,index], top=𝒟.wT.scaled[end, index]))
    end

    BCs = [BC(i) for i in 0:n_simulations - 1]

    @info "Setting up differential equations"

    prob_NDEs = [ODEProblem((x, p, t) -> DE(x, p, t, derivatives, scalings, constants, BCs[i]), uvT₀s[i], tspan_train) for i in 1:n_simulations]

    function determine_loss_scalings()
        if training_fractions === nothing
            loss_scalings = (u=1, v=1, T=1, ∂u∂z=gradient_scaling, ∂v∂z=gradient_scaling, ∂T∂z=gradient_scaling)
        else
            unscaled_parameters = unscale_parameter.(scaled_parameters, scalings.parameters)
            sols = [Array(solve(prob_NDEs[i], timestepper, p=unscaled_parameters, reltol=1f-3, sensealg=InterpolatingAdjoint(autojacvec=ZygoteVJP()), saveat=t_train)) for i in 1:n_simulations]
            u_sols, v_sols, T_sols = split_u.(sols, Nz), split_v.(sols, Nz), split_T.(sols, Nz)

            u_loss = mean(loss.(u_trains, u_sols))
            v_loss = mean(loss.(v_trains, v_sols))
            T_loss = mean(loss.(T_trains, T_sols))

            if train_gradient
                u_sols_gradients = [∂_∂z(sol, D_face) for sol in u_sols]
                v_sols_gradients = [∂_∂z(sol, D_face) for sol in v_sols]
                T_sols_gradients = [∂_∂z(sol, D_face) for sol in T_sols]

                ∂u∂z_loss = mean(loss.(u_trains_gradients, u_sols_gradients))
                ∂v∂z_loss = mean(loss.(v_trains_gradients, v_sols_gradients))
                ∂T∂z_loss = mean(loss.(T_trains_gradients, T_sols_gradients))
            else
                ∂u∂z_loss = 0
                ∂v∂z_loss = 0
                ∂T∂z_loss = 0
            end

            losses = (u=u_loss, v=v_loss, T=T_loss, ∂u∂z=∂u∂z_loss, ∂v∂z=∂v∂z_loss, ∂T∂z=∂T∂z_loss)
            loss_scalings = calculate_loss_scalings(losses, training_fractions, train_gradient)
        end
        return loss_scalings
    end

    @info "Determining training scalings"

    loss_scalings = determine_loss_scalings()

    function loss_mpp(parameters, p)
        unscaled_parameters = unscale_parameter.(parameters, scalings.parameters)
        sols = [Array(solve(prob_NDEs[i], timestepper, p=unscaled_parameters, reltol=1f-3, sensealg=InterpolatingAdjoint(autojacvec=ZygoteVJP()), saveat=t_train)) for i in 1:n_simulations]

        u_sols, v_sols, T_sols = split_u.(sols, Nz), split_v.(sols, Nz), split_T.(sols, Nz)
        u_loss = mean(loss.(u_trains, u_sols))
        v_loss = mean(loss.(v_trains, v_sols))
        T_loss = mean(loss.(T_trains, T_sols))

        losses = (u=u_loss, v=v_loss, T=T_loss, ∂u∂z=0, ∂v∂z=0, ∂T∂z=0)
        scaled_losses = apply_loss_scalings(losses, loss_scalings)

        return sum(scaled_losses), scaled_losses
    end

    function loss_gradient_mpp(parameters, p)
        unscaled_parameters = unscale_parameter.(parameters, scalings.parameters)
        sols = [Array(solve(prob_NDEs[i], timestepper, p=unscaled_parameters, reltol=1f-3, sensealg=InterpolatingAdjoint(autojacvec=ZygoteVJP()), saveat=t_train)) for i in 1:n_simulations]

        u_sols, v_sols, T_sols = split_u.(sols, Nz), split_v.(sols, Nz), split_T.(sols, Nz)

        u_sols_gradients = [∂_∂z(sol, D_face) for sol in u_sols]
        v_sols_gradients = [∂_∂z(sol, D_face) for sol in v_sols]
        T_sols_gradients = [∂_∂z(sol, D_face) for sol in T_sols]

        u_loss = mean(loss.(u_trains, u_sols))
        v_loss = mean(loss.(v_trains, v_sols))
        T_loss = mean(loss.(T_trains, T_sols))

        ∂u∂z_loss = mean(loss.(u_trains_gradients, u_sols_gradients))
        ∂v∂z_loss = mean(loss.(v_trains_gradients, v_sols_gradients))
        ∂T∂z_loss = mean(loss.(T_trains_gradients, T_sols_gradients))

        losses = (u=u_loss, v=v_loss, T=T_loss, ∂u∂z=∂u∂z_loss, ∂v∂z=∂v∂z_loss, ∂T∂z=∂T∂z_loss)
        scaled_losses = apply_loss_scalings(losses, loss_scalings)

        return sum(scaled_losses), scaled_losses
    end

    @info "Setting up optimization objective"

    if train_gradient
        f_loss = OptimizationFunction(loss_gradient_mpp, GalacticOptim.AutoZygote())
    else
        f_loss = OptimizationFunction(loss_mpp, GalacticOptim.AutoZygote())
    end

    prob_loss = OptimizationProblem(f_loss, scaled_parameters, lb=[0f0, 0f0, 0f0, 0f0, 0f0], ub=[10f0, 10f0, 10f0, 10f0, 10f0])

    @info "Writing metadata"

    write_metadata_modified_pacanowski_philander_optimisation(FILE_PATH, train_files, maxiters, tsteps, unscale_parameter.(scaled_parameters, scalings.parameters), loss_scalings, optimizers)

    @info "Beginning Training"

    res = 0
    for i in 1:length(optimizers)
        iter = 1
        opt = optimizers[i]
        function cb(args...)
            if iter <= maxiters
                parameters = args[1]
                unscaled_parameters = unscale_parameter.(parameters, scalings.parameters)
                ν₀, ν₋, ΔRi, Riᶜ, Pr = unscaled_parameters
                losses = args[3]
                profile_loss = losses.u + losses.v + losses.T
                gradient_loss = losses.∂u∂z + losses.∂v∂z + losses.∂T∂z
                total_loss = profile_loss + gradient_loss
                @info "ν₀ = $ν₀, ν₋ = $ν₋, ΔRi = $ΔRi, Riᶜ = $Riᶜ, Pr = $Pr, loss = $(total_loss): profile = $profile_loss, gradient = $gradient_loss (u = $(losses.u), v = $(losses.v), T = $(losses.T)) , ∂u∂z = $(losses.∂u∂z), ∂v∂z = $(losses.∂v∂z), ∂T∂z = $(losses.∂T∂z)), optimizer $i/$(length(optimizers)), iteration = $iter/$maxiters"
                write_data_modified_pacanowski_philander_optimisation(FILE_PATH, losses, unscaled_parameters)
            end
            iter += 1
            false
        end

        res = solve(prob_loss, opt, cb=cb, maxiters=maxiters)
        scaled_parameters .= res.minimizer
    end

    ν₀, ν₋, ΔRi, Riᶜ = unscale_parameter.(scaled_parameters, scalings.parameters)
    @info "ν₀ = $ν₀, ν₋ = $ν₋, ΔRi = $ΔRi, Riᶜ = $Riᶜ, Pr = $Pr, loss = $loss"
end


text:
## Example of connected components in an algorithm
using SparseArrays
using LinearAlgebra
A = [0 1 1 1 1 0 1 0 0 0
     1 0 0 0 0 0 0 0 0 0
     0 0 0 0 1 1 1 0 0 0
     1 1 0 0 1 1 0 1 0 0
     1 1 1 1 0 0 0 0 1 0
     0 0 1 0 0 0 1 0 1 1
     0 0 0 0 0 1 0 0 0 1
     0 0 0 0 0 0 0 0 1 0
     0 0 0 0 0 0 0 1 0 0
     0 0 0 0 0 0 0 0 0 0]
n = size(A,1)
x = ones(n)
for i=1:20
  global x
  x = A*x
  x = x/norm(x)
end
x
##
x = zeros(n)
x[9] = 1
for i=1:20
  global x
  x = A'*x
  x = x/norm(x)
end
x


text:
include("load.jl")
using Plots

include("flow.jl")
include("coefficient.jl")
include("prop.jl")

deltamT2_success=-5.759e7
deltamT2_failed= -5.8e7
const deltamT2 = (deltamT2_success+deltamT2_failed)/2

test_zc=1.0

# -5.677822*10^5

sol.u

function flow!(du, u, p, t)
    za = @view u[:, 1] # za
    zc = @view u[:, 2] #zc
    λcA = @view u[:, 3]
    λ3A = @view u[:, 4]
    dzadt = @view du[:, 1]
    dzcdt = @view du[:, 2]
    dλcAdt = @view du[:, 3]
    dλ3Adt = @view du[:, 4]
    zafun = extrapolate(interpolate((pgrid,), za, Gridded(Linear())), Flat())
    # zafun = Spline1D(pgrid,za)
    zcfun = extrapolate(interpolate((pgrid,), zc, Gridded(Linear())), Flat())
    λcAfun = extrapolate(interpolate((pgrid,), λcA, Gridded(Linear())), Flat())
    λ3Afun = extrapolate(interpolate((pgrid,), λ3A, Gridded(Linear())), Flat())
    # zcfun = Spline1D(pgrid,zc)
    m2 = za[1]
    # hatza = @. pgrid^2 * (pgrid^2 + m2)^-1 * za
    # hatzafun =
    #     extrapolate(interpolate((pgrid,), hatza, Gridded(Linear())), Flat())
    # hatzafun(x) = x^2 * (x^2 + m2) * ((x^2 + m2)^2 + 0.00001 * m2)^-1 * zafun(x)
    # hatzafun(x) = x^2 * (x^2 + m2) ^-1 * zafun(x)
    dzadtfun =
        extrapolate(interpolate((pgrid,), dzadt, Gridded(Linear())), Flat())
    # dzadtfun=Spline1D(pgrid,dzadt)
    dzcdtfun =
        extrapolate(interpolate((pgrid,), dzcdt, Gridded(Linear())), Flat())
    # dzcdtfun=Spline1D(pgrid,dzcdt)
    k = Λ * exp(t)
    println(m2," ","k=",k," ","zc=",zc[1])
    # display(plot(pgrid, [zc], xaxis = :log, seriestype = :scatter))
    # display( Plots.plot(x->GAqmp(x, 1000.0, 0.3, m2, hatzafun.(@. max(abs(x - 1000.0), 1)), t),1,3k,xaxis=:log))
    # println(dzadt)
    # println(dzcdt)
    tildeZA = zafun((kbar^6 + k^6)^(1 / 6))
    DtildeZA =
        dzadtfun((kbar^6 + k^6)^(1 / 6)) +
        k^6 *
        (kbar^6 + k^6)^(-5 / 6) *
        Interpolations.gradient(zafun, (kbar^6 + k^6)^(1 / 6))[1]
    # DtildeZA =
    #         dzadtfun((kbar^6 + k^6)^(1 / 6)) +
    #         k^6 *
    #         (kbar^6 + k^6)^(-5 / 6) *
    #         derivative(zafun, (kbar^6 + k^6)^(1 / 6))
    Dtildezc = dzcdtfun(k) + k * Interpolations.gradient(zcfun, k)[1]

    dzadt .= Cubature.hcubature_v(N_pgrid, (x,v) -> flowA_v!(v,
            x,
            m2,
            zafun,
            tildeZA,
            DtildeZA,
            zcfun(k),
            Dtildezc,
            zcfun,
            λcAfun,
            λ3Afun,
            t,
        ),[1.0, -1.0],[3 * k, 1.0],maxevals = 2400)[1]

    dzcdt .= Cubature.hcubature_v(N_pgrid, (x,v) -> flowG_v!(v,
            x,
            m2,
            zafun,
            tildeZA,
            DtildeZA,
            zcfun(k),
            Dtildezc,
            zcfun,
            λcAfun,
            λ3Afun,
            t,
        ),[1.0, -1.0],[3 * k, 1.0],maxevals = 2400)[1]

    dλcAdt .=Cubature.hcubature_v(N_pgrid, (x,v) -> flowλcA_v!(v,
            x,
            m2,
            zafun,
            tildeZA,
            DtildeZA,
            zcfun(k),
            Dtildezc,
            zcfun,
            λcAfun,
            λ3Afun,
            t,
        ),[1.0, -1.0, -1.0],[3 * k, 1.0, 1.0],maxevals = 9600)[1]

    dλ3Adt .= Cubature.hcubature_v(N_pgrid, (x,v) -> flowλ3A_v!(v,
            x,
            m2,
            zafun,
            tildeZA,
            DtildeZA,
            zcfun(k),
            Dtildezc,
            zcfun,
            λcAfun,
            λ3Afun,
            t,
        ),[1.0, -1.0, -1.0],[3 * k, 1.0, 1.0],maxevals = 9600)[1]
end

# const deltamT2 = -5.73125 * 10^7# initial condition for mk^2

writedlm(
        "/home/wjfu1/tyy/YM-Zero-Vertex/YM-Zero-Vertex-2/log.txt","")
argmax(sol.u[1][:,1].^-1)

for i in 1:2
    include("solve.jl")
end

for i = 1:20
    include("solve.jl")
end

rm("/home/wjfu1/tyy/YM-Zero-Vertex/YM-Zero-Vertex-2/fig/", recursive=true)
rm("/home/wjfu1/tyy/YM-Zero-Vertex/YM-Zero-Vertex-2/data2/", recursive=true)

mkdir("/home/wjfu1/tyy/YM-Zero-Vertex/YM-Zero-Vertex-2/fig/")

mkdir("/home/wjfu1/tyy/YM-Zero-Vertex/YM-Zero-Vertex-2/data2/")

writedlm(
    "/home/wjfu1/tyy/YM-Zero-Vertex/YM-Zero-Vertex-2/data2/pgrid.dat",pgrid)

savefig(plot(
    pgrid,
    pgrid .^ -2 .* sol(-11.51292546497)[:, 1] .^ -1,
    xaxis = :log,
    seriestype = :scatter,
),"/home/wjfu1/tyy/YM-Zero-Vertex/YM-Zero-Vertex-2/fig/1.pdf")

using Plots
plot(
    pgrid,
    pgrid .^ -2 .* sol(-11.51292546497)[:, 1] .^ -1,
    xaxis = :log,
    seriestype = :scatter,
)

argmax(sol(-11.51292546497)[:, 1].^-1)

plot(
    pgrid,
    sol(-11.51292546497)[:, 2],
    xaxis = :log,
    yaxis =:log,
    seriestype = :scatter,
)

zc0004=readdlm("/home/wjfu1/tyy/YM-Zero-Vertex/YM-Zero-Vertex-2/data2/-5.7755567878484726e7.dat")

zc0019=readdlm("/home/wjfu1/tyy/YM-Zero-Vertex/YM-Zero-Vertex-2/data2/-5.775556778907776e7.dat")
zc001=readdlm("/home/wjfu1/tyy/YM-Zero-Vertex/YM-Zero-Vertex-2/data2/-5.77555678486824e7.dat")
zc015=readdlm("/home/wjfu1/tyy/YM-Zero-Vertex/YM-Zero-Vertex-2/data2/-5.7755568981170654e7.dat")

zc0019
zc0004[:,2][1]
plot(
    pgrid,
    [zc0004[:, 2] zc001[:, 2] zc0019[:, 2] zc015[:, 2]],
    xaxis = :log,
    yaxis = :log,
    seriestype = :scatter,
)

plot(
    pgrid,
    [zc0004[:, 1] zc001[:, 1] zc0019[:, 1] zc015[:, 1]],
    xaxis = :log,
    yaxis = :log,
    seriestype = :scatter,
)

sol(-11.51292546497)[:,3]
plot(
    pgrid,
    zc0019[:,3],
    xaxis = :log,
)

plot(
    pgrid,
    zc0019[:,3] .^2 .* zc0019[:,1] .^-1 .* zc0019[:,2] .^-2,
    xaxis = :log,
    yaxis = :log,
)

plot(
    pgrid,
    [zc0019[:,4] .^2 .* zc0019[:,1] .^-3,zc0019[:,3] .^2 .* zc0019[:,1] .^-1 .* zc0019[:,2] .^-2],
    xaxis = :log,
    yaxis = :log,
)

zc0019[:,1]

max_stars_repo_name: Yangyang-Tan/FRG-Code-Collection

max_issues_repo_name: Yangyang-Tan/FRG-Code-Collection

max_forks_repo_name: Yangyang-Tan/FRG-Code-Collection


text:
# NOTE: Modify these things!
simname = "test-sim"

function outdir_suffix(scale, kmcmc, seed_data)
  return "seed_$(seed_data)/scale_$(scale)/Kmcmc_$(kmcmc)"
end

settings = [Dict(:simname => simname,
                 :repfam_dist_scale => scale,
                 :Kmcmc => kmcmc,
                 :seed_data => seed_data,
                 :seed_mcmc => 0,
                 :outdir_suffix => outdir_suffix(scale, kmcmc, seed_data))
            for scale in [0, 10],
            kmcmc in [2, 3, 5],
            seed_data in 2:3]

settings = vec(settings)


text:
using Test

using ForwardDiff
using Equilibrium
using EFIT

mu0 = 4pi*10^-7

@testset "Equilibrium" begin

include("solovev.jl")

include("efit.jl")

end


text:
using RecipesBase

custom_colorant = [
    colorant"#44AA99",
    colorant"#332288",
    colorant"#DDCC77",
    colorant"#CC6677",
    colorant"#AA4499",
    colorant"#DDDDDD",
    colorant"#117733",
    colorant"#882255",
    colorant"#999933",
    colorant"#88CCEE",
]

# create various plot types:
# Sensitivity
# Performance Curves
# HeatMap Sensitivity
stats_plot_types = [:violin, :boxplot, :dotplot, :density, :histogram, :histogram2d, :groupedhist]

@recipe function f(dc::DataCollection, params::Dict)
    idx = findall(dc.data) do ld
        all([params[i] == ld.line_params[i] for i in keys(params)])
    end
    if length(idx) == 1
        dc[idx[1]]
    else
        dc[idx]
    end
end

@recipe function f(dc::DataCollection, params::AbstractVector)
    for arg in params
        @series begin
            dc, arg
        end
    end
end

@recipe function f(func::Function, dc::DataCollection)
    idx = findall(func, dc.data)
    dc[idx]
end

@recipe function f(dc::DataCollection; color_func=nothing)

    legendtitle --> "$(join(sort(collect(keys(dc[1].line_params))), ' '))"
    for d ∈ dc.data
        @series begin
            if !(color_func isa Nothing)
                color --> color_func(d)
            end
            d
        end
    end
end

@recipe function f(ld_vec::DataCollection{LD}; sort_idx=nothing, z=1.0, color_func=nothing) where {F<:Number, LD<:LineData{Vector{F}}}

    ld_sorted = if sort_idx isa Nothing
        ld_vec.data
    else
        sf = (ld1::LineData, ld2::LineData) -> isless(ld1.line_params[sort_idx], ld2.line_params[sort_idx])
        sort(ld_vec.data, lt=sf)
    end

    if !(color_func isa Nothing)
        color --> color_func(d)
    end

    st = get(plotattributes, :seriestype, :path)
    if st ∈ stats_plot_types
        for ld in ld_sorted
            @series begin
                label_idx --> sort_idx
                ld
            end
        end
    else
        x = [ld.line_params[sort_idx] for ld in ld_sorted]
        y = [mean(ld.data) for ld in ld_sorted]
        yerror := [z*sqrt(var(ld.data)/length(ld.data)) for ld in ld_sorted]

        (x, y)
    end
end

@recipe function f(ld::LineData{Vector{F}}; label_idx=nothing, make_label_string=false, z=1.0) where F<:Number
    st = get(plotattributes, :seriestype, :boxplot)
    if st ∈ stats_plot_types
        y = ld.data
        if !(label_idx isa Nothing)
            x = if make_label_string
                [string(ld.line_params[label_idx])]
            else
                [ld.line_params[label_idx]]
            end
            (x, y)
        else
            x = if :x ∈ keys(plotattributes)
                plotattributes[:x]
            elseif :label ∈ keys(plotattributes)
                [plotattributes[:label]]
            else
                throw("Need :x or :label when not using label_idx.")
            end
            (x, y)
        end
    else
        error("$(typeof(ld)) not supported for series type: $(st).")
    end
end

@recipe function f(ld::LineData{Vector{Vector{F}}}; z=1.0, color_dict=nothing) where F<:Number

    if :label ∈ keys(plotattributes)
	label := "$(plotattributes[:label]), $(ld.swept_params))"
    else
        ks = sort(collect(keys(ld.line_params)))
        lbl = prod(string(ld.line_params[k]) * " " for k in ks)
	label := "$(lbl), $(ld.swept_params)"
    end

    if !(color_dict isa Nothing)
        color --> color_dict[ld.line_params]
    end

    st = get(plotattributes, :seriestype, :path)
    if st ∈ stats_plot_types
        y = vcat(transpose.(ld.data)...)
        if :x ∈ keys(plotattributes)
            (plotattributes[:x], y)
        else
            (1:length(ld.data[1]), y)
        end
    else
        ribbon := std_uneven(ld.data; z=z)
        μ = mean_uneven(ld.data)
        y = μ
        x --> 1:length(μ)

        (plotattributes[:x], y)
    end

end


text:
export minreal

"""
### function minreal(G::TransferFunction)

Compute the minimal realization of a dynamic system represented by a transfer
function.  In other words, compute a transfer function without common poles and
zeros.

##### Args

* G: Transfer function.

##### Returns

The transfer function of the minimal realization of the dynamic system.

"""

function minreal(G::TransferFunction)
    # Get the list of poles and zeros.
    p = pole(G)
    z = izero(G)

    # Get the system gain.
    k = G.num[end]/G.den[end]

    # Get the poles and zeros of the minimal system.
    mp = p
    mz = Array{Complex128}(0)

    # Auxiliary variable to compute tolarance.
    #
    # This algorithm was adapted from file `__minreal__.m` of Octave Forge.
    sqrt_eps = sqrt(eps())

    # Remove the poles that can be cancelled with zeros.
    for i = 1:length(z)
        # Compute the tolerance.
        tol = (abs(z[i]) < sqrt_eps) ? 10000*sqrt_eps : 10000*abs(z[i])*sqrt_eps

        # Compute the distance from the poles to the zero.
        Δ = abs.(p-z[i])

        # Find if there are any poles within a distance of `tol` from the
        # analyzed zero.
        ind = find(Δ .< tol)

        if !isempty(ind)
            # Delete the closest pole to the current zero.
            deleteat!(mp, indmin(Δ))
        else
            push!(mz,round(z[i],6))
        end
    end

    # Return the minimal transfer function.
    zpk(mz, mp, k)
end


text:
module Ops20

import Base: show # this will be unnecessary if we don't use structs -- only in rnn.jl
using AutoGrad: AutoGrad, @primitive, @primitive1, value, Param # we should not need Param in ops -- only in rnn.jl
using Base: has_offset_axes, unsafe_convert
using LinearAlgebra.BLAS: libblas, @blasfunc
using LinearAlgebra: chkstride1, BlasInt, lmul!
using NNlib: conv, DenseConvDims, maxpool, meanpool, PoolDims, ∇conv_data, ∇conv_filter, ∇maxpool, ∇meanpool
using Random: rand!, seed!
using Statistics: mean, var

include("activation.jl"); export elu, relu, selu, sigm, invx
include("dropout.jl");    export dropout
include("bmm.jl");        export bmm
include("conv.jl");       export conv4, deconv4, mat, pool, unpool     # rename conv/deconv, deprecate mat?
include("softmax.jl");    export logsoftmax, logsumexp, softmax, logp  # deprecate one of logp/logsoftmax?
include("loss.jl");       export accuracy, bce, logistic, nll, zeroone # deprecate one of bce/logistic, old signatures of nll/acc
include("batchnorm.jl");  export batchnorm, bnmoments, bnparams        # rethink the interface: single struct? no struct?
include("rnn.jl");        export RNN, rnninit, rnnforw, rnnparam, rnnparams # rethink the interface. no struct?

end


text:
"""
Compute the `n-th` degree Chebyshev polynomial of first and second kind and its derivatives

 * `chebyshev(x, n)` Chebyshev polynomial of the first kind
 * `chebyshev2(x, n)` Chebyshev polynomial of the second kind
 * `dchebyshev(x, n)` Derivative of Chebyshev polynomial of the first kind
 * `dchebyshev2(x, n)` Derivative of Chebyshev polynomial of the second kind

There are also functions that compute the zeros of Chebyshev polynomials:

 * `chebyshev_zeros`
 * `chebyshev2_zeros`
 * `chebyshev_zeros!`
 * `chebyshev2_zeros!`

"""
function chebyshev(x, n)
    if n==0
        return one(x)
    elseif n==1
        return x
    end

    T0 = one(x)
    T1 = x

    for i = 2:n
        T2 = 2*x*T1 - T0
        T0 = T1
        T1 = T2
    end

    return T1

end

function chebyshev2(x, n)

    if n==0
        return one(x)
    elseif n==1
        return 2x
    end

    U0 = one(x)
    U1 = 2*x

    for i = 2:n
        U2 = 2*x*U1 - U0
        U0 = U1
        U1 = U2
    end

    return U1

end

dchebyshev(x, n) = n * chebyshev2(x, n-1)

function dchebyshev2(x, n)

    if n == 0
        return zero(x)
    elseif n==1
        return 2*one(x)
    end

    dU0 = zero(x)
    dU1 = 2
    U0 = one(x)
    U1 = 2*x

    for i = 2:n
        U2 = 2*x*U1 - U0
        dU2 = 2*U1 + 2*x*dU1 - dU0

        U0 = U1
        U1 = U2
        dU0 = dU1
        dU1 = dU2
    end

    return dU1
end

function chebyshev_zeros!(n, x::AbstractArray{T}) where {T<:Number}

    for k in 1:n
        num::T =  (2*k - 1)*pi

        x[k] = -cos( num/(2*n) )
    end

    return x

end

chebyshev_zeros(n, ::Type{T}=Float64) where {T<:Number} = chebyshev_zeros!(n, zeros(T, n))

function chebyshev2_zeros!(n, x::AbstractArray{T}) where {T<:Number}

    for k in 1:n
        num::T = k*pi

        x[k] = -cos( num/(n+1) )
    end

    return x

end
chebyshev2_zeros(n, ::Type{T}=Float64) where {T<:Number} = chebyshev2_zeros!(n, zeros(T, n))

@doc (@doc chebyshev) chebyshev2

@doc (@doc chebyshev) dchebyshev
@doc (@doc chebyshev) dchebyshev2

@doc (@doc chebyshev) chebyshev_zeros!
@doc (@doc chebyshev) chebyshev2_zeros!
@doc (@doc chebyshev) chebyshev_zeros
@doc (@doc chebyshev) chebyshev2_zeros