text:
# returns the list of subexpressions which a given tape depends on directly
function list_subexpressions(nd::Vector{NodeData})

    # TODO: code overlap with compute_gradient_sparsity
    indices = Set{Int}()
    for k in 1:length(nd)
        nod = nd[k]
        if nod.nodetype == SUBEXPRESSION
            push!(indices, nod.index)
        end
    end
    return sort(collect(indices))
end

"""
    _topological_sort(
        starts::Vector{Int},
        subexpressions::Vector{Vector{NodeData}},
        subexpression_dependency_graph::Vector{Vector{Int}} =
            Vector{Vector{Int}}(undef, length(subexpressions)),
    )

Return a topologically sorted list of the integer subexpresssion indices that
need to be computed to evalute `subexpressions[s]` for all `s in starts`.

`starts` should be ordered, and not contain duplicates.

`subexpression_dependency_graph[i]` is a lazily-computed list of "out" edges
from node `i`, in terms of the integer-valued subexpression index (i.e.,
`node.index`). This list should be unique and ordered.

If calling `_topological_sort` a single time, you may omit the
`subexpression_dependency_graph` argument.

However, if calling `_topological_sort` multiple times on the _same_ vector of
subexpresssions, you should create `subexpression_dependency_graph` once (either
as the uninitialized vector, or by explicitly computing the full
`subexpression_dependency_graph`), and pass it in.

## Notes

* It is important to not use recursion here, because expressions may have
  arbitrary levels of nesting!
* This function assumes `subexpressions` is acyclic.
"""
function _topological_sort(
    starts,
    subexpressions::Vector{Vector{NodeData}},
    subexpression_dependency_graph::Vector{Vector{Int}} = Vector{Vector{Int}}(
        undef,
        length(subexpressions),
    ),
)
    ordered = Int[]
    in_order = fill(false, length(subexpressions))
    stack = Tuple{Int,Bool}[]
    for s in starts
        if in_order[s]
            continue  # s is already in `ordered`.
        end
        push!(stack, (s, true))
        while !isempty(stack)
            node, needs_checking = pop!(stack)
            if !needs_checking
                # We must be returning to this node for a second time, and we
                # have already checked all of the children. Therefore, we can
                # add it to the set of ordered nodes.
                push!(ordered, node)
                in_order[node] = true
                continue
            elseif in_order[node]
                continue  # This node has already been added to `ordered`.
            end
            # Re-add the node to the stack, but set the `false` flag this time
            # so next time we visit it, it will go on the `ordered` list
            # instead.
            push!(stack, (node, false))
            if !isassigned(subexpression_dependency_graph, node)
                subexpression_dependency_graph[node] =
                    list_subexpressions(subexpressions[node])
            end
            for child in subexpression_dependency_graph[node]
                if !in_order[child]
                    push!(stack, (child, true))
                end
            end
        end
    end
    return ordered
end

"""
    order_subexpressions(
        main_expressions::Vector{Vector{NodeData}},
        subexpressions::Vector{Vector{NodeData}};
    )

Topologically sort the subexpression needed to evaluate `main_expressions`.

Returns two things:

 * A `Vector{Int}` containing the ordered list of subexpression-indices that
   need to be evaluated to compute all `main_expressions`
 * A `Vector{Vector{Int}}`, containing a list of ordered lists of
   subexpression-indices that need to be evaluated to compute
   `main_expressions[i]`.

**Warning:** This doesn't handle cyclic expressions! But this should be fine
because we can't compute them in JuMP anyway.
"""
function order_subexpressions(
    main_expressions::Vector{Vector{NodeData}},
    subexpressions::Vector{Vector{NodeData}},
)
    # The graph of node dependencies. Constructed lazily.
    subexpression_dependency_graph =
        Vector{Vector{Int}}(undef, length(subexpressions))
    # Node dependencies of the main expressions.
    starts = Set{Int}()
    individual_sorts = Vector{Int}[]
    for expression in main_expressions
        s = list_subexpressions(expression)
        union!(starts, s)
        push!(
            individual_sorts,
            _topological_sort(
                s,
                subexpressions,
                subexpression_dependency_graph,
            ),
        )
    end
    full_sort = _topological_sort(
        starts,
        subexpressions,
        subexpression_dependency_graph,
    )
    return full_sort, individual_sorts
end

"""
    order_subexpressions(
        main_expression::Vector{NodeData},
        subexpressions::Vector{Vector{NodeData}},
    )

Return a `Vector{Int}` containing the topologically ordered list of
subexpression-indices that need to be evaluated to compute `main_expression`.
"""
function order_subexpressions(
    main_expression::Vector{NodeData},
    subexpressions::Vector{Vector{NodeData}},
)
    s = list_subexpressions(main_expression)
    return _topological_sort(s, subexpressions)
end

export list_subexpressions, order_subexpressions

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
"""
Summary of posterior over (f, f′) given current evaluations. The posterior
takes the form of a GP, which can be queried using methods #TODO, and stores
the evaluations so far.
"""

const ζ = 10     # numerical stability in kernel function (any positive const is ok)

mutable struct PLSPosterior{T <: AbstractFloat}
    Ts::Vector{T}
    Y::Vector{T}
    ∇Y::Vector{T}
    σ²_f::T
    σ²_∇::T
    G::Matrix{T}   # posterior kernel matrix
    A::Vector{T}   # posterior mean weights
end

# Constructor
function PLSPosterior(Ts::Vector{T}, Y::Vector{T}, ∇Y::Vector{T}, σ²_f::T, σ²_∇::T) where T <: AbstractFloat

    N = length(Ts)

    # build prior Gram matrix components
    K   = k.(Ts, Ts')
    K∂  = kd.(Ts, Ts')
    ∂K∂ = dkd.(Ts, Ts')

    # build full (noised) Gram matrix
    G = [K K∂; transpose(K∂) ∂K∂];
    G[diagind(G)[1:N]] .+= σ²_f
    G[diagind(G)[N+1:2N]] .+= σ²_∇

    # weights of posterior mean (for linear combo of query points)
    A = G \ [Y; ∇Y];

    return PLSPosterior(Ts, Y, ∇Y, σ²_f, σ²_∇, G, A)
end

#= Integrated Wiener Process Kernel functions (function and derivative)
------------------------------------------------------------------------=#
# kernel:
k(a,b)   = min(a + ζ,b + ζ)^3/3 + abs(a-b) * min(a + ζ,b + ζ)^2/2;
kd(a,b)  = (a<b) * ((a + ζ)^2/2) + (a>=b) * ((a + ζ)*(b + ζ) - (b + ζ)^2/2);
dk(a,b)  = (a>b) * ((b+ζ)^2/2) + (a<=b) .* ((a+ζ)*(b+ζ) - (a+ζ).^2/2);
dkd(a,b) = min(a+ζ,b+ζ);

# further derivatives
ddk(a,b) = (a<=b) * (b-a);
ddkd(a,b) = (a<=b);
dddk(a,b) = -(a<=b);

#= Posterior statistic calculations for PLSPosterior
------------------------------------------------------------------------=#
# posterior mean function and all its derivatives
m(gp::PLSPosterior, t::Vector)   = let Ts = gp.Ts; [k.(t, Ts')    kd.(t,  Ts')  ] * gp.A; end
d1m(gp::PLSPosterior, t::Vector) = let Ts = gp.Ts; [dk.(t, Ts')   dkd.(t,  Ts') ] * gp.A; end
d2m(gp::PLSPosterior, t::Vector) = let Ts = gp.Ts; [ddk.(t, Ts')  ddkd.(t,  Ts')] * gp.A; end
d3m(gp::PLSPosterior{T}, t::Vector) where T <: AbstractFloat = let Ts = gp.Ts; [dddk.(t, Ts')  zeros(T, length(t), length(Ts))] * gp.A; end

# posterior marginal covariance between function and first derivative
Σ(gp::PLSPosterior, t::Vector)   = let Ts = gp.Ts; k.(t,t')   - ([k.(t, Ts')   kd.(t, Ts') ] * (gp.G \ [k.(t, Ts')   kd.(t, Ts') ]')); end
Σd(gp::PLSPosterior, t::Vector)  = let Ts = gp.Ts; kd.(t,t')  - ([k.(t, Ts')   kd.(t, Ts') ] * (gp.G \ [dk.(t, Ts')  dkd.(t, Ts')]')); end
dΣd(gp::PLSPosterior, t::Vector) = let Ts = gp.Ts; dkd.(t,t') - ([dk.(t, Ts')  dkd.(t, Ts')] * (gp.G \ [dk.(t, Ts')  dkd.(t, Ts')]')); end

V(gp::PLSPosterior, t::Vector) = diag(Σ(gp, t))
Vd(gp::PLSPosterior, t::Vector) = diag(Σd(gp, t))
dVd(gp::PLSPosterior, t::Vector) = diag(dΣd(gp, t))

# covariance terms with function (derivative) values at origin
# TODO: vectorised versions (d>1) do not work! (dimensions must match)
Σ0f(gp::PLSPosterior, t::Vector{T}) where T <: Real  = let Ts = gp.Ts; T0 = T(0); k.(T0,t)   - ([k.(T0, Ts')   kd.(T0, Ts') ] * (gp.G \ [k.(t, Ts')   kd.(t, Ts') ]')); end
Σd0f(gp::PLSPosterior, t::Vector{T}) where T <: Real  = let Ts = gp.Ts; T0 = T(0); dk.(T0,t)  - ([dk.(T0, Ts')  dkd.(T0, Ts')] * (gp.G \ [k.(t, Ts')   kd.(t, Ts') ]')); end
Σ0df(gp::PLSPosterior, t::Vector{T}) where T <: Real  = let Ts = gp.Ts; T0 = T(0); kd.(T0,t)  - ([k.(T0, Ts')   kd.(T0, Ts') ] * (gp.G \ [dk.(t, Ts')  dkd.(t, Ts')]')); end
Σd0df(gp::PLSPosterior, t::Vector{T}) where T <: Real = let Ts = gp.Ts; T0 = T(0); dkd.(T0,t) - ([dk.(T0, Ts')  dkd.(T0, Ts')] * (gp.G \ [dk.(t, Ts')  dkd.(t, Ts')]')); end

V0f(gp::PLSPosterior, t::Vector{T}) where T <: Real = diag(Σ0f(gp, t))
Vd0f(gp::PLSPosterior, t::Vector{T}) where T <: Real = diag(Σd0f(gp, t))
V0df(gp::PLSPosterior, t::Vector{T}) where T <: Real = diag(Σ0df(gp, t))
Vd0df(gp::PLSPosterior, t::Vector{T}) where T <: Real = diag(Σd0df(gp, t))

# # ========================= SCALAR VERSIONS =============================
m(gp::PLSPosterior, t::Real) = m(gp, [t]) |> arr2sc
d1m(gp::PLSPosterior, t::Real) = d1m(gp, [t]) |> arr2sc
d2m(gp::PLSPosterior, t::Real) = d2m(gp, [t]) |> arr2sc
d3m(gp::PLSPosterior, t::Real) = d3m(gp, [t]) |> arr2sc
V(gp::PLSPosterior, t::Real) = Σ(gp, [t]) |> arr2sc
Vd(gp::PLSPosterior, t::Real) = Σd(gp, [t]) |> arr2sc
dVd(gp::PLSPosterior, t::Real) = dΣd(gp, [t]) |> arr2sc
V0f(gp::PLSPosterior, t::Real) = Σ0f(gp, [t]) |> arr2sc
Vd0f(gp::PLSPosterior, t::Real) = Σd0f(gp, [t]) |> arr2sc
V0df(gp::PLSPosterior, t::Real) = Σ0df(gp, [t]) |> arr2sc
Vd0df(gp::PLSPosterior, t::Real) = Σd0df(gp, [t]) |> arr2sc

# Update
function update(gp::PLSPosterior{T}, x::PLSEvaluation{T}; update_llh_var::Bool=false) where T <: AbstractFloat
    # Currently this just rebuilds from scratch -- my belief is that
    # this is sufficiently fast, especially the 2-dimensional broadcast on the
    # kernel matrix constructions that there is no need to utilise intermediate
    # results. However, it's an obvious place to look for optimisation if reqd,
    # esp. if you increase the evaluation limit.
    @unpack σ²_f, σ²_∇ = gp
    @unpack t, y, ∇y = x
    N = length(gp) + 1
    Ts, Y, ∇Y = vcat(gp.Ts, t), vcat(gp.Y, y), vcat(gp.∇Y, ∇y)
    update_llh_var && begin; σ²_f, σ²_∇ = (x.σ²_f + (N-1) * σ²_f)/N, (x.σ²_∇ + (N-1) * σ²_∇)/N; end
    return PLSPosterior(Ts, Y, ∇Y, σ²_f, σ²_∇)
end

# Find the minimum mean value over specified Ts (default: current evaluations)
min_mean(x::PLSPosterior) = min_mean(x, x.Ts)
function min_mean(x::PLSPosterior{T}, Ts::Vector{T}) where T <: AbstractFloat
    M  = m(x, Ts)
    minj = argmin(M)
    return M[minj], Ts[minj], minj
end

# Standard methods
Base.length(x::PLSPosterior) = length(x.Ts)
Base.eltype(x::PLSPosterior{T}) where T <: AbstractFloat = T

# allow broadcasting over posterior objects (interpret as scalar)
Broadcast.broadcastable(x::PLSPosterior) = Ref(x)

ext: jl

max_stars_repo_name: ornithos/ProbLineSearch


text:
@testset "Testing MatMul" begin
    # First, test normal matmul.
    a = rand(Float32, 10, 20)
    b = rand(Float32, 20, 10)
    A = OneDNN.Memory(a)
    B = OneDNN.Memory(b)

    Z = OneDNN.materialize(OneDNN.matmul(A, B))
    z = a * b
    @test size(Z) == (size(a, 1), size(b, 2))
    @test isapprox(Z, z)
    @inferred *(a, b)

    Z = OneDNN.materialize(
        OneDNN.matmul(OneDNN.Memory(transpose(b)), OneDNN.Memory(transpose(a)))
    )
    @test isapprox(transpose(Z), z)

    # Do pullbacks work?
    # Note: `rrule` definition lives in ChainRuies.jl under `Base/arraymath.jl`.
    x = randn(Float32, 100, 100)
    y = randn(Float32, 100, 100)
    X = OneDNN.Memory(x)
    Y = OneDNN.Memory(y)

    Z, back_dnnl = Zygote._pullback(*, X, Y)
    z, back_jl = Zygote._pullback(*, x, y)
    @test isapprox(OneDNN.materialize(Z), z)

    dz = Float32(0.125) * randn(Float32, size(z))
    DZ = OneDNN.Memory(dz)

    grads_jl = back_jl(dz)
    grads_dnnl = back_dnnl(DZ)

    @test length(grads_jl) == length(grads_dnnl) == 3
    @test grads_jl[1] === grads_dnnl[1] === nothing
    @test isapprox(grads_jl[2], OneDNN.materialize(grads_dnnl[2]))
    @test isapprox(grads_jl[3], OneDNN.materialize(grads_dnnl[3]))
end

# # Test applying some post ops
# x = rand(Float32, 2, 2)
# y = rand(Float32, 2, 2)
# z = rand(Float32, 2, 2)

# # expected result
# expected = (x * y) + (transpose(z) * x)

# out = OneDNN.matmul(y, x)

# postops = OneDNN.PostOps()
# OneDNN.appendsum!(postops)
# attr = OneDNN.Attributes()
# OneDNN.add!(attr, postops)
# OneDNN.matmul!(out, x, transpose(z); attributes = attr)

# result = OneDNN.materialize(out)
# @test isapprox(result, expected)

# # Test scaling
# x = randn(Float32, 5, 5)
# y = randn(Float32, 5, 5)
# attr = OneDNN.Attributes()
# OneDNN.setscale!(attr, 2)
# z = OneDNN.matmul(y, x; attributes = attr)

# @test isapprox(OneDNN.materialize(z), 2 * x * y)


text:
# %% Calculating Kaplan Yorke dimensions for as many systems as possible
using DrWatson
@quickactivate :FractalDimension # uses DynamicalSystems, PyPlot
include(srcdir("style.jl"))

# Time to evolve for reaching convergence
N = Int(1e6)
Ds = range(4; step = 2, length = 6)

function Δλ(params)
    @unpack N, D = params
    ds = Systems.lorenz96(D, range(0; length = D, step = 0.1); F = 24.0)
    λs = lyapunovspectrum(ds, N; Ttr = 100, FractalDimension.Data.diffeq...)
    Δ = kaplanyorke_dim(λs)
    return @strdict N D λs Δ
end

for D in Ds
    params = @strdict N D
    file, s = produce_or_load(
        datadir("kaplanyorke"), params, Δλ;
        prefix = ("lorenz96_chaotic"), suffix = "jld2", force = true
    )
    println("For Lorenz-96:")
    println("D=$D, Δ=$(file["Δ"])")
end

# %% Show convergence explicitly
ds = Systems.lorenz96(D, range(0; length = D, step = 0.1); F = 24.0)
tinteg = tangent_integrator(ds; FractalDimension.Data.diffeq...)
λs, t = ChaosTools.lyapunovspectrum_convergence(tinteg, N, 0.1, 100.0)
plot(t, λs)

# %% Rest of systems
systems = Symbol[
    :henon_chaotic,
    :towel_chaotic,
    :coupled_logistics, # D = 8
    :roessler_chaotic,
    :henonheiles_chaotic,
    # :lorenz96_chaotic, # D = 8, already calculated above
]
N = Int(1e6)

for s in systems # start producing
    params = @dict N

    ds = if s == :henon_chaotic
        Systems.henon([0.08152, 0.122717]; a = 1.4, b = 0.3)
    elseif s == :towel_chaotic
        Systems.towel([0.085, -0.121, 0.075])
    elseif s == :coupled_logistics
        params[:r] = 4.0
        u0 = range(0.1, 0.9; length = 8) .+ 0.0001
        ds = DiscreteDynamicalSystem(Data.logistics!, u0, [0.1, 3.0])
    elseif s == :roessler_chaotic
        ds = Systems.roessler([0.1, -0.2, 0.1]; c = 5.7, a = 0.2, b = 0.2)
    elseif s == :henonheiles_chaotic
        u0 = [0.0, -0.25, 0.42081, 0.0]
        ds = Systems.henonheiles(u0)
    end

    @show s
    function Δλ2(params)
        @unpack N = params
        λs = lyapunovspectrum(ds, N; Ttr = 100, FractalDimension.Data.diffeq...)
        @show λs
        Δ = kaplanyorke_dim(λs)
        return @strdict N λs Δ
    end

    file, s = produce_or_load(
        datadir("kaplanyorke"), params, Δλ2;
        prefix = string(s), suffix = "jld2", force = false
    )

    println("Δ=$(file["Δ"])")
end

# %% Try coupled tent maps
tr = Data.coupled_logistics(; D = 8, k = 0.1, r = 4.0)

using3D()
x,y,z= columns(tr)
scatter3D(x,y,z)

max_stars_repo_name: Datseris/FractalDimension

max_issues_repo_name: Datseris/FractalDimension

max_forks_repo_name: Datseris/FractalDimension


text:
module Symplectic

import LinearAlgebra
include("DataTypes/SymplecticForm.jl")
include("DataTypes/CoupledResonators.jl")

export SymplecticForm,
       Ω,
       CoupledResonators,
       addActive,
       addPassive,
       addGammaEx,
       addGammaIn,
       scatteringMatrix,
       nonSymplecticity,
       checkSymplectic,
       modes,
       getQ,
       getP,
       ⊗,
       dsum,
       ⊕,
       sct,
       isct,
       symplecticCayleyTransform,
       inverseSymplecticCayleyTransform,
       lstf,
       localSymplecticTransformation,
       ibs,
       interferenceBasedSequence,
       decoupleSequence,
       teleportation,
       genericity,
       generify,
       phaseShifting,
       randomPhaseShifiting,
       localSymplecticMatrix,
       randomLocalSymplecticMatrix,
       beamSplitter,
       circulator,
       toQPQPBasis,
       toQQPPBasis,
       pif,
       preIwasawaFactorization,
       spectrum

# Utilities
using LinearAlgebra:
       checksquare,
       opnorm,
       eigen,
       isposdef,
       issymmetric,
       norm,
       dot

↑(m, i) = 2*m - (i % 2)
↓(m, i) = 2*m -1 + (i  % 2)

function Lkmij(u, v, m, i, j)
    if norm(v[2*m-1:2*m]) ≈ 0
        return i==j ? 1 : 0
    else
        return (-1)^(j+1)*v[↑(m, i)]*u[↓(m, j)]/( v[2*m]^2 + v[2*m-1]^2 ) + (-1)^i*v[↓(m, i)]*u[↑(m, j)]/( u[2*m]^2 + u[2*m-1]^2 )
    end
end

function getSequence(Ss, modeToDecouple::Int; lambda::Real = 1)
       if length(Ss) == 4
           return decoupleSequence(Ss...; m = modeToDecouple, lambda = lambda)
       else
           step = length(Ss) ÷ 4
           Rs = [ getSequence(Ss[i:i+step-1], modeToDecouple-1 ) for i in 1:step:length(Ss) ]

           Ls = decoupleSequence([ *(RLst...) for RLst in Rs]...; m = modeToDecouple, lambda = lambda)
           return vcat([ (i%2==0) ? [ Ls[i] ] : Rs[(i+1) ÷ 2] for i in 1:7 ]...)
       end
end

# calculate the 'nonSymplecticity".
nonSymplecticity(J::Union{SymplecticForm, UniformScaling}) = 0
nonSymplecticity(A::AbstractMatrix) = opnorm(transpose(A)*Ω*A - Ω) / opnorm(Ω)

# check if a matrix is symplectic, returning its non-symplecticity
checkSymplectic(J::Union{SymplecticForm, UniformScaling}) = 0
function checkSymplectic(A::AbstractMatrix)
       eltype(A) <: Real || throw(ArgumentError("matrix is not real-valued"))
       checksquare(A) % 2 == 0 || throw(DimensionMismatch("matrix is not even-dimensional"))
       return nonSymplecticity(A)
end
function checkSymplectic(A...)
       nonSymplecticities = []
       for a in A
              if !( a isa Union{SymplecticForm, UniformScaling} )
                     eltype(a) <: Real || throw(ArgumentError("matrix is not real-valued"))
                     checksquare(a) % 2 == 0 || throw(DimensionMismatch("matrix is not even-dimensional"))
              end
              push!(nonSymplecticities, nonSymplecticity(a))
       end
       return nonSymplecticities
end

genericity(S::AbstractMatrix) = (x -> max(x...) / min(x...))(abs.(S))

phaseShifting(θ::Real) = [cos(θ) sin(θ); -sin(θ) cos(θ)]
phaseShifting(θs...) = dsum( map(phaseShifting, θs)... )
randomPhaseShifiting(n::Int) =
       n == 1 ? phaseShifting(2π * rand()) : phaseShifting( (2π*rand(n))...)

localSymplecticMatrix(θ::Real, l::Real, ϕ::Real) =  phaseShifting(θ) * [l 0 ; 0  1/l ] * phaseShifting(ϕ)
localSymplecticMatrix(θs::Vector, ls::Vector, ϕs::Vector) =
       dsum( map(args -> localSymplecticMatrix(args...), zip(θs, ls, ϕs))... )
randomLocalSymplecticMatrix(n::Int; bound::Real = 0.2) =
       n == 1 ? localSymplecticMatrix(2π * rand(), bound*( rand() -1/2) + 1.0, 2π * rand()) :
              localSymplecticMatrix(2π * rand(n), bound*( rand(n) .-1/2) .+ 1.0, 2π * rand(n))

# Generte a Two-Mode BeamSplitter
beamSplitter(angle::Real) = phaseShifting(angle) ⊗ I(2)
# Generate a BeamSplitter between two modes in a multi-mode System
function beamSplitter(angle::Real, mode1, mode2, modes)
    mode1, mode2 = sort([mode1, mode2])
    lst = [ [[2 * (i + 2) - 1, 2 * (i + 2)] for i in 1:mode1-1]... ; [1, 2] ; [[2 * (i + mode1 + 1) - 1, 2 * (i + mode1 + 1)] for i in 1: mode2 - mode1 - 1]...; [3, 4] ; [[2 * (i + mode2 ) - 1, 2 * (i + mode2 )]  for i in 1:modes-mode2]...  ];
    return dsum( beamSplitter(angle), I(2*modes - 4) )[lst, lst]
end

circulator(perm...) = circulator(perm)
circulator(perm::Vector) =
       hvcat(length(perm), [i == perm[j] ? I(2) : zeros(2, 2)   for i in 1:length(perm), j in 1:length(perm) ]...)

function generify(S::AbstractMatrix; randomize::Function = (n -> I(2 * n)), l::Int = 1)
       n = modes(S)
       seq = vcat( [[randomize(n), S] for i in 1:l ]... , [randomize(n)] )
       g = genericity(*(seq...))
       for i in 1:100
              seqNew = vcat( [[randomize(n), S] for i in 1:l ]... , [randomize(n)] )
              gNew = genericity(*(seqNew...))
              if gNew < g
                     seq = seqNew
                     g = gNew
              end
       end
       return seq
end

modes(A::AbstractMatrix) = ( checkSymplectic(A); size(A, 1) ÷ 2 )

# Get quadtrature (column vector) from a symplectic matrix
getQ(J::Union{SymplecticForm, UniformScaling}, n::Int) = J[:, 2*n-1]
getQ(A::AbstractMatrix, n::Int) =
       n ≤ modes(A) ? A[:, 2*n-1] : throw(ArgumentError("index is out of bound"))
getQ(J::Union{SymplecticForm, UniformScaling}, n::AbstractRange{<:Integer}) = [J[:, 2*i-1] for i in n]
function getQ(A::AbstractMatrix{T}, n::AbstractRange{<:Integer}) where T
       length(n) ≤ modes(A) || throw(ArgumentError("index is out of bound"))
       return [A[:, 2*i-1] for i in n]
end

getP(J::Union{SymplecticForm, UniformScaling}, n::Int) = J[:, 2*n]
getP(A::AbstractMatrix, n::Int) =
       n ≤ modes(A) ? A[:, 2*n] : throw(ArgumentError("index is out of bound"))
getP(J::Union{SymplecticForm, UniformScaling}, n::AbstractRange{<:Integer}) = [J[:, 2*i] for i in n]
function getP(A::AbstractMatrix{T}, n::AbstractRange{<:Integer}) where T
       length(n) ≤ modes(A) || throw(ArgumentError("index is out of bound"))
       return [A[:, 2*i] for i in n]
end

const ⊗ = Base.kron

dsum(M1::AbstractMatrix, M2::AbstractMatrix) = cat(M1, M2; dims=(1,2))
dsum(M1::AbstractMatrix, M2::AbstractMatrix, Ms::Vararg{AbstractMatrix}) = dsum(dsum(M1, M2),  mapreduce(identity, dsum, Ms ))
const ⊕ = dsum

# Transforming a symmetric matrix (member of the symplectic algebra) to a symplectic matrix
symplecticCayleyTransform(M::AbstractMatrix) = I - (Ω * M  + I / 2  )^-1
const sct = symplecticCayleyTransform

# Transforming a symplectic matrix to a symmetric matrix
inverseSymplecticCayleyTransform(S::AbstractMatrix) = -Ω*((I - S)^-1 - I/2)
const isct = inverseSymplecticCayleyTransform

# Construct a local symplectic transformation transforming quadrature u to quadrature λΩv (the conjugate of quadrature v)
function localSymplecticTransformation(u::Vector, v::Vector; lambda::Real = 1)
       ( length(u) == length(v) && length(u) % 2 == 0 && length(v) % 2 == 0 ) || throw(ArgumentError("invalid input"))
       n = length(u) ÷ 2
       if n == 1
              return [ Lkmij(u, lambda*v, 1, i, j) for i in 1:2, j in 1:2 ]
       else
              return ⊕([ [ Lkmij(u, lambda*v, m, i, j) for i in 1:2, j in 1:2 ] for m in 1:n ]...)
       end
end
const lstf = localSymplecticTransformation

# Given symplectic matrices S₄, S₃, S₂, and S₁, return a sequence S₄ L₃ S₃ L₂ S₂ L₁ S₁ decoupling the m-th mode, where the local operations are determined up to the scaling factor λ
function decoupleSequence(S4, S3, S2, S1; m::Int=1, lambda::Real = 1)
       L1 = lstf(S1[:,2*m-1], S2[2*m-1,:]; lambda = lambda)
       L3 = lstf(S3[:,2*m-1], S4[2*m-1,:]; lambda = lambda^-1)
       T1 = S2 * L1 * S1
       T2 = S4 * L3 * S3
       L2 = ⊕([ (i==m ? Ω(2) : lstf(T1[2*i-1:2*i,2*m], T2[2*m, 2*i-1:2*i])) for i in 1:(size(T1, 1)÷2)]...)
       return [S4, L3, S3, L2, S2, L1, S1]
end

# Get local Operations L16, L15, ..., L1 for the interfernce-basded sequence LR S L15 S ... S L1 S = ST, with a given target Symplectic Matrix ST
function interferenceBasedSequence(S::AbstractMatrix; T = I(4), lambda::Real = 1)
       return interferenceBasedSequence([S for i in 1:4^modes(T)]; T = T, lambda = lambda)
end
function interferenceBasedSequence(Ss; T = I(4), lambda::Real = 1)
       n =  modes(Ss[1]) - modes(T)
       nT = modes(T)
       Ss[end] =  Ss[end] * ( inv(T) ⊕ I( 2 * n  ) )
       Ss = getSequence(Ss, nT; lambda = lambda)
       R = ( *(Ss...)[1:2 * nT, 1:2 * nT])^-1  ⊕ I( 2 * n )
       Ss[end] =  Ss[end] * ( T ⊕ I( 2 * n ) )
       return vcat([ R ] , Ss)
end
const ibs = interferenceBasedSequence

function teleportation(S::AbstractMatrix, inModes::Vector, outModes::Vector)
       allModes = [1:(size(S, 1) ÷ 2);]
       In = vcat( [ [2*i - 1, 2 * i] for i in inModes ]... )
       Out = vcat( [ [2*i - 1, 2 * i] for i in outModes ]... )
       ancModes = [i for i in allModes if !(i in inModes)]
       idlModes = [i for i in allModes if !(i in outModes)]
       Usq = 2*ancModes
       Hm = 2*idlModes
       SOutIn = S[ Out, In  ]
       SHmIn = S[ Hm, In ]
       SOutUsq = S[ Out, Usq ]
       SHmUsq = S[ Hm, Usq ]
       return SOutIn - SOutUsq * SHmUsq^-1 * SHmIn
end

function toQQPPBasis(S::AbstractMatrix)
        n = modes(S)
        order = [ [1:2:2*n;]... , [2:2:2*n;]... ]
        return S[order, order]
end

function toQPQPBasis(S::AbstractMatrix)
       n = modes(S)
       order = vcat([ [i, i + n] for i in 1:n]... )
       return S[order, order]
end

function preIwasawaFactorization(S::AbstractMatrix)
       n = modes(S)
       S = toQQPPBasis(S)
       A, B = S[1:n, 1:n], S[1:n, 1+n:2*n]
       C, D = S[1+n:2*n, 1:n], S[1+n:2*n, 1+n:2*n]
       AAtBBt = A*transpose(A) + B*transpose(B)
       P = (C*transpose(A) + D*transpose(B))* AAtBBt^-1
       L = AAtBBt^(1/2)
       X = AAtBBt^(-1/2) * A
       Y = AAtBBt^(-1/2) * B
       return map(toQPQPBasis, [
                     [ I(n) zeros(n, n); P I(n) ],
                     [ L zeros(n, n); zeros(n, n) L^-1 ],
                     [ X Y; -Y X]
              ])
end
const pif = preIwasawaFactorization

# Return Symplectic Spectrum Diagonal([λ..., λ...]) and a symplectic matrix S such that Sᵀ M S = toQPQPBasis( Diagonal([λ..., λ...]) )
function spectrum(M::AbstractMatrix)
       (eltype(M)<:Real && isposdef(M)) || throw(ArgumentError("only implemented for real-valued positive-definite matrix"))
       n = size(M, 1) ÷ 2
       K = toQQPPBasis(Ω * M)
       eig = eigen(K; sortby = (λ ->(abs(λ) , imag(λ))))
       spec = [abs(imag( v )) for v in eig.values]
       es = [ (v=eig.vectors[:, 2*i - 1];real(v)) for i in 1:n]
       fs = [ (v=eig.vectors[:, 2*i - 1];-imag(v)) for i in 1:n]
       for i in 1:n
              λ = es[i] ⋅ (toQQPPBasis(Ω(2*n)) * fs[i])
              es[i] /= sqrt(λ)
              fs[i] /= sqrt(λ)
       end
       S = hcat( es..., fs... )
       return [spec, toQPQPBasis(S)]
end

end


text:
# Copyright (c) 2017: Miles Lubin and contributors
# Copyright (c) 2017: Google Inc.
#
# Use of this source code is governed by an MIT-style license that can be found
# in the LICENSE.md file or at https://opensource.org/licenses/MIT.

"""
    ScalarizeBridge{T,F,S}

`ScalarizeBridge` implements the following reformulations:

  * ``f(x) - a \\in \\mathbb{R}_+`` into ``f_i(x) \\ge a_i`` for all ``i``
  * ``f(x) - a \\in \\mathbb{R}_-`` into ``f_i(x) \\le a_i`` for all ``i``
  * ``f(x) - a \\in \\{0\\}`` into ``f_i(x) == a_i`` for all ``i``

## Source node

`ScalarizeBridge` supports:

  * `G` in [`MOI.Nonnegatives{T}`](@ref)
  * `G` in [`MOI.Nonpositives{T}`](@ref)
  * `G` in [`MOI.Zeros{T}`](@ref)

## Target nodes

`ScalarizeBridge` creates:

  * `F` in `S`, where `S` is one of [`MOI.GreaterThan{T}`](@ref),
    [`MOI.LessThan{T}`](@ref), and [`MOI.EqualTo{T}`](@ref), depending on the
    type of the input set.
"""
mutable struct ScalarizeBridge{T,F,S} <: AbstractBridge
    scalar_constraints::Vector{MOI.ConstraintIndex{F,S}}
    constants::Vector{T}
end

const Scalarize{T,OT<:MOI.ModelLike} =
    SingleBridgeOptimizer{ScalarizeBridge{T},OT}

function bridge_constraint(
    ::Type{ScalarizeBridge{T,F,S}},
    model::MOI.ModelLike,
    f::MOI.AbstractVectorFunction,
    set::MOI.Utilities.VectorLinearSet,
) where {T,F,S}
    dimension = MOI.output_dimension(f)
    constants = MOI.constant(f, T)
    new_f = MOI.Utilities.scalarize(f, true)
    constraints = Vector{MOI.ConstraintIndex{F,S}}(undef, dimension)
    for i in 1:dimension
        constraints[i] = MOI.add_constraint(model, new_f[i], S(-constants[i]))
    end
    return ScalarizeBridge{T,F,S}(constraints, constants)
end

function MOI.supports_constraint(
    ::Type{ScalarizeBridge{T}},
    F::Type{<:MOI.AbstractVectorFunction},
    ::Type{<:MOI.Utilities.VectorLinearSet},
) where {T}
    # If `F` is `MOI.VectorAffineFunction{Complex{Float64}}`, `S` is `MOI.Zeros`
    # and `T` is `Float64`, it would create a set `MOI.EqualTo{Float64}` which
    # is incorrect hence we say we only support it if the coefficient type of
    # `F` is `T`.
    return MOI.Utilities.is_coefficient_type(F, T)
end

function MOI.Bridges.added_constrained_variable_types(::Type{<:ScalarizeBridge})
    return Tuple{Type}[]
end

function MOI.Bridges.added_constraint_types(
    ::Type{ScalarizeBridge{T,F,S}},
) where {T,F,S}
    return Tuple{Type,Type}[(F, S)]
end

function concrete_bridge_type(
    ::Type{<:ScalarizeBridge{T}},
    F::Type{<:MOI.AbstractVectorFunction},
    S::Type{<:MOI.Utilities.VectorLinearSet},
) where {T}
    return ScalarizeBridge{
        T,
        MOI.Utilities.scalar_type(F),
        MOI.Utilities.scalar_set_type(S, T),
    }
end

function MOI.get(
    bridge::ScalarizeBridge{T,F,S},
    ::MOI.NumberOfConstraints{F,S},
)::Int64 where {T,F,S}
    return length(bridge.scalar_constraints)
end

function MOI.get(
    bridge::ScalarizeBridge{T,F,S},
    ::MOI.ListOfConstraintIndices{F,S},
) where {T,F,S}
    return copy(bridge.scalar_constraints)
end

function MOI.delete(model::MOI.ModelLike, bridge::ScalarizeBridge)
    for ci in bridge.scalar_constraints
        MOI.delete(model, ci)
    end
    return
end

function MOI.delete(
    model::MOI.ModelLike,
    bridge::ScalarizeBridge,
    i::MOI.Bridges.IndexInVector,
)
    MOI.delete(model, bridge.scalar_constraints[i.value])
    deleteat!(bridge.scalar_constraints, i.value)
    deleteat!(bridge.constants, i.value)
    return
end

function MOI.get(
    model::MOI.ModelLike,
    attr::MOI.ConstraintFunction,
    bridge::ScalarizeBridge{T},
) where {T}
    func = MOI.Utilities.vectorize(
        MOI.get.(model, attr, bridge.scalar_constraints),
    )
    if func isa MOI.VectorOfVariables
        return func
    end
    # `func` is in terms of bridged variables here while in `bridge_constraint`
    # it was in terms of the solver variables so `MOI.constant(set)` might be
    # different than `bridge.constants[i]`.
    for i in eachindex(bridge.scalar_constraints)
        set = MOI.get(model, MOI.ConstraintSet(), bridge.scalar_constraints[i])
        rhs = MOI.constant(set)
        if !iszero(rhs)
            func = MOI.Utilities.operate_output_index!(-, T, i, func, rhs)
        end
    end
    return func
end

function MOI.get(
    ::MOI.ModelLike,
    ::MOI.ConstraintSet,
    bridge::ScalarizeBridge{T,F,S},
) where {T,F,S}
    return MOI.Utilities.vector_set_type(S)(length(bridge.scalar_constraints))
end

function MOI.supports(
    model::MOI.ModelLike,
    attr::Union{MOI.ConstraintPrimalStart,MOI.ConstraintDualStart},
    ::Type{ScalarizeBridge{T,F,S}},
) where {T,F,S}
    return MOI.supports(model, attr, MOI.ConstraintIndex{F,S})
end

function MOI.get(
    model::MOI.ModelLike,
    attr::MOI.ConstraintPrimal,
    bridge::ScalarizeBridge,
)
    return MOI.get.(model, attr, bridge.scalar_constraints) .+ bridge.constants
end

function MOI.get(
    model::MOI.ModelLike,
    attr::MOI.ConstraintPrimalStart,
    bridge::ScalarizeBridge,
)
    values = MOI.get.(model, attr, bridge.scalar_constraints)
    if any(isnothing, values)
        return
    end
    return values .+ bridge.constants
end

function MOI.set(
    model::MOI.ModelLike,
    attr::MOI.ConstraintPrimalStart,
    bridge::ScalarizeBridge,
    ::Nothing,
)
    for ci in bridge.scalar_constraints
        MOI.set(model, attr, ci, nothing)
    end
    return
end

function MOI.set(
    model::MOI.ModelLike,
    attr::MOI.ConstraintPrimalStart,
    bridge::ScalarizeBridge,
    value,
)
    MOI.set.(model, attr, bridge.scalar_constraints, value .- bridge.constants)
    return
end

function MOI.get(
    model::MOI.ModelLike,
    attr::MOI.ConstraintDual,
    bridge::ScalarizeBridge,
)
    return MOI.get.(model, attr, bridge.scalar_constraints)
end

function MOI.get(
    model::MOI.ModelLike,
    attr::MOI.ConstraintDualStart,
    bridge::ScalarizeBridge,
)
    values = MOI.get.(model, attr, bridge.scalar_constraints)
    if any(isnothing, values)
        return
    end
    return values
end

function MOI.set(
    model::MOI.ModelLike,
    attr::MOI.ConstraintDualStart,
    bridge::ScalarizeBridge,
    value,
)
    MOI.set.(model, attr, bridge.scalar_constraints, value)
    return
end

function MOI.modify(
    model::MOI.ModelLike,
    bridge::ScalarizeBridge{T,F,S},
    change::MOI.VectorConstantChange{T},
) where {T,F,S}
    bridge.constants = change.new_constant
    MOI.set.(
        model,
        MOI.ConstraintSet(),
        bridge.scalar_constraints,
        S.(-change.new_constant),
    )
    return
end

function MOI.modify(
    model::MOI.ModelLike,
    bridge::ScalarizeBridge,
    change::MOI.MultirowChange{T},
) where {T}
    for (index, value) in change.new_coefficients
        MOI.modify(
            model,
            bridge.scalar_constraints[index],
            MOI.ScalarCoefficientChange{T}(change.variable, value),
        )
    end
    return
end

function MOI.set(
    model::MOI.ModelLike,
    ::MOI.ConstraintFunction,
    bridge::ScalarizeBridge{T,F,S},
    func,
) where {T,F,S}
    old_constants = bridge.constants
    bridge.constants = MOI.constant(func, T)
    new_func = MOI.Utilities.scalarize(func, true)
    MOI.set.(
        model,
        MOI.ConstraintFunction(),
        bridge.scalar_constraints,
        new_func,
    )
    for i in eachindex(bridge.constants)
        if bridge.constants[i] != old_constants[i]
            MOI.set(
                model,
                MOI.ConstraintSet(),
                bridge.scalar_constraints[i],
                S(-bridge.constants[i]),
            )
        end
    end
    return
end

# TODO implement transform


text:
# ## RANDOM DEBUG DEV CODE BELOW
#
# hist = getCliqSolveHistory(tree, :x1)
#
# tree, smt, hist = solveTree!(fg, skipcliqids=[:x1;:x6;:x4;:x3], recordcliqs=[:x2;:x0])
#
# plotTreeUpMsgs(fg, tree, :x1, levels=1)
#
#
# cliq = getClique(tree, :x2)
#
# # OLD
# dwinmsgs = IIF.prepCliqInitMsgsDown!(fg, tree, getParent(tree,cliq)[1], cliq, dbgnew=false)
# plotKDE(dwinmsgs[:x1][1], dims=[1;2], levels=2)
#
# # NEW
# dwinmsgs = IIF.prepCliqInitMsgsDown!(fg, tree, getParent(tree,cliq)[1], cliq, dbgnew=true)
# plotKDE(dwinmsgs[:x1][1], dims=[1;2], levels=2)
#
#
#
# ## Check init message for x3
#
# cliq = getClique(tree, :x1)
#
# # OLD
# dwinmsgs = IIF.prepCliqInitMsgsDown!(fg, tree, getParent(tree,cliq)[1], cliq, dbgnew=false)
# plotKDE(dwinmsgs[:x3][1], dims=[1;2], levels=2)
#
# # NEW
# dwinmsgs = IIF.prepCliqInitMsgsDown!(fg, tree, getParent(tree,cliq)[1], cliq, dbgnew=true)
# plotKDE(dwinmsgs[:x3][1], dims=[1;2], levels=2)
#
#
#
# ## Check init message for x4
#
# cliq = getClique(tree, :x4)
#
# # OLD
# dwinmsgs = IIF.prepCliqInitMsgsDown!(fg, tree, getParent(tree,cliq)[1], cliq, dbgnew=false)
# plotKDE(dwinmsgs[:x3][1], dims=[1;2], levels=2)
#
# # NEW
# dwinmsgs = IIF.prepCliqInitMsgsDown!(fg, tree, getParent(tree,cliq)[1], cliq, dbgnew=true)
# plotKDE(dwinmsgs[:x3][1], dims=[1;2], levels=2)
#
#
#
#
# plotKDE(fg, :x1, dims=[1;2])
# plotPose(fg, :x1)
#
# plotKDE(fg, [:x0;:x1;:x2;:x4;:x6], dims=[1;2],levels=1)
#
#
# drawTree(tree, imgs=true)
#
#
#
# ## check contents
#
# cliq = getClique(tree, :x1)
#
# getCliqMsgsUp(cliq)
#
#
# #
#
# # DEBUG after solveCliq :x2
#
#
# # solve third clq
# smt, hist = solveCliq!(fg, tree, :x1, cliqHistories=hist, recordcliq=true)
#
#
# # Plot to see what is going on
# Gadfly.set_default_plot_size(35cm,20cm)
# plotKDE(fg, [:x0;:x1;:x2], dims=[1;2],levels=1)
#
# plotKDE(fg, :x1, dims=[1;2],levels=1)
#
#
#
# ## DEBUG where is down message :x3
#
# prnt = getClique(tree, :x3)
# getCliqInitUpMsgs(prnt)
#
#
# assignTreeHistory!(tree, hist)
# printCliqHistorySummary(tree, :x1)
#
#
# # think the issue is in here
# # 21:01:00.856  8   null         attemptCliqInitUp     false null | upsolved upsolved
#
#
#
# # somethings up with cliq x1
#
# csmc_8_test = getCliqSolveHistory(tree, :x1)[8][4]
# csmc_9_test = getCliqSolveHistory(tree, :x1)[9][4]
#
# # Just before the mistake
# drawTree(csmc_8_test.tree, show=true)
# prnt_8_test = getClique(csmc_8_test.tree, :x3)
# getCliqInitUpMsgs(prnt_8_test)
#
# # just after the mistake
# prnt_9_test = getClique(csmc_9_test.tree, :x3)
# getCliqInitUpMsgs(prnt_9_test)
#
#
# # so develop  based on sandbox step 8
# stuff = sandboxCliqResolveStep(tree,:x1,8)
# getCliqInitUpMsgs(getClique(stuff[4].tree,:x3))
#
#
#
# IIF.getCliqMsgsUp(tree,:x1)
#
#
# ## DEBUG
#
#
#
#
#
#
#
# ## DEBUG after solveCliq :x2 difference in down init cycle when :dontUseParentFactorsInitDown
#
# assignTreeHistory!(tree, hist)
# printCliqHistorySummary(tree, :x2)
#
# # WORKS
#
# delete!(getSolverParams(getCliqSolveHistory(tree, :x2)[15][4].dfg).devParams, :dontUseParentFactorsInitDown)
#
# good = sandboxCliqResolveStep(tree,:x2,15)
#
# # plotPairVariables(csmcX2_ref[16][4].cliqSubFg, stuff[4].cliqSubFg, :x1, title="(X2,) sandbox 15 good,")
# # plotPairVariables(csmcX2_ref[16][4].cliqSubFg, stuff[4].cliqSubFg, :x2, title="(X2,) sandbox 15 good,")
#
#
# # BREAKS
#
# getSolverParams(getCliqSolveHistory(tree, :x2)[15][4].dfg).devParams[:dontUseParentFactorsInitDown] = ""
#
# bad = sandboxCliqResolveStep(tree,:x2,15)
#
# # plotPairVariables(csmcX2_ref[16][4].cliqSubFg, stuff[4].cliqSubFg, :x1, title="(X2,) sandbox 15 bad,")
# # plotPairVariables(csmcX2_ref[16][4].cliqSubFg, stuff[4].cliqSubFg, :x2, title="(X2,) sandbox 15 bad,")
#
# plotPairVariables(good[4].cliqSubFg, bad[4].cliqSubFg, :x2, title="(X2,) sandbox 15,")
# # plotPairPose2(good[4].cliqSubFg, bad[4].cliqSubFg, :x2, title="(X2,) sandbox 15,")
#
# Gadfly.set_default_plot_size(30cm,18cm)
#
# ## DEBUG
#
#
#
#
#
#
#
# #3 First difference, up msg from (:x2) cliq is 'wrong'
#
# Gadfly.set_default_plot_size(30cm,18cm)
# um1 = IIF.getCliqMsgsUp(tree, :x0)
# plotKDE(um1[:x1], dims=[1;2], levels=2)
#
#
#
# um2 = IIF.getCliqMsgsUp(tree, :x2)
# plotKDE(um2[:x1], dims=[1;2], levels=2)
#
#
#
#
# csmcX2_test = getCliqSolveHistory(tree, :x2)
#
# # csmcX2_test[15][4].cliqSubFg
#
# plotKDE(getKDE(csmcX2_test[17][4].cliqSubFg, :x1),dims=[1;2],levels=2,title="csmcX2_test[17]")
#
#
#
#
# ## SOMETHING changes between csmcX2_ref[16/17] and csmcX2_test[16/17]
#
# writeGraphPdf(csmcX2_ref[17][4].cliqSubFg, show=true)
# writeGraphPdf(csmcX2_test[17][4].cliqSubFg, show=true)
#
#
# plotKDE(csmcX2_ref[16][4].cliqSubFg, [:x1;:x2;:x3], dims=[1;2])
# plotKDE(csmcX2_test[16][4].cliqSubFg, [:x1;:x2;:x3], dims=[1;2])
#
#
#
# isInitialized(csmcX2_test[16][4].cliqSubFg, :x3)
#
#
#
#
#
#
# # before init cycle
# plotPairVariables(csmcX2_ref[16][4].cliqSubFg, csmcX2_test[16][4].cliqSubFg, :x1)
# plotPairVariables(csmcX2_ref[16][4].cliqSubFg, csmcX2_test[16][4].cliqSubFg, :x2)
# plotPairVariables(csmcX2_ref[16][4].cliqSubFg, csmcX2_test[16][4].cliqSubFg, :x3)
#
#
# # after solve
# plotPairVariables(csmcX2_ref[17][4].cliqSubFg, csmcX2_test[17][4].cliqSubFg, :x1)
# plotPairVariables(csmcX2_ref[17][4].cliqSubFg, csmcX2_test[17][4].cliqSubFg, :x2)
# plotPairVariables(csmcX2_ref[17][4].cliqSubFg, csmcX2_test[17][4].cliqSubFg, :x3)


text:
using TeX
using POMDPStressTesting
using Distributions
using Parameters

doc = globaldoc("walk1d"; build_dir="output_walk1d", jmlr=true)
doc.title = "POMDPStressTesting.jl Example: Walk1D"
doc.author = "Robert J. Moss"
doc.address = "Computer Science, Stanford University"
doc.email = "mossr@cs.stanford.edu"
doc.title_case_sections = false
doc.use_subsections = true

@tex T"""\begin{abstract}
In this self-contained tutorial, we define a simple problem for adaptive stress testing (AST)
to find failures. This problem, called Walk1D, samples random walking distances from a standard
normal distribution $\mathcal{N}(0,1)$ and defines failures as walking past a certain threshold
(which is set to $\pm 10$ in this example). AST will either select the seed which deterministically
controls the sampled value from the distribution (i.e. from the transition model) or will directly
sample the provided environmental distributions. These action modes are determined by the seed-action or
sample-action options. AST will guide the simulation to failure events using a notion of distance to failure,
while simultaneously trying to find the set of actions that maximizes the log-likelihood of the samples.
\end{abstract}"""

addpackage!(doc, "url")
addkeywords!(["BlackBox", "GrayBox", "Simulation", "Environment", "EnvironmentSample", "ASTSampleAction", "ASTSeedAction", "Walk1DSim", "Walk1DParams", "ASTMDP", "MCTSPWSolver", "CEMSolver", "TRPOSolver", "PPOSolver", "RandomSearchSolver", "Distribution", "Normal", "logpdf"]; num=2)
addkeywords!(["initialize!", "transition!", "evaluate!", "distance", "isevent", "isterminal", "setup_ast", "search!", "playback", "print_metrics", "environment", "solve"]; num=3)

@tex T"""\section{Gray-box Simulator and Environment}
The simulator and environment are treated as gray-box because we need
access to the state-transition distributions and their associated likelihoods.
"""

@tex T"""\paragraph{Parameters.}
First, we define the parameters of our simulation.""" ->
@with_kw mutable struct Walk1DParams
    startx::Float64 = 0 # Starting x-position
    threshx::Float64 = 10 # +- boundary threshold
    endtime::Int64 = 30 # Simulate end time
end

@tex "\\paragraph{Simulation.} Next, we define a \\texttt{GrayBox.Simulation} structure." ->
@with_kw mutable struct Walk1DSim <: GrayBox.Simulation
    params::Walk1DParams = Walk1DParams() # Parameters
    x::Float64 = 0 # Current x-position
    t::Int64 = 0 # Current time ±
    distribution::Distribution = Normal(0, 1) # Transition distribution
end

@tex "Then, we define our \\texttt{GrayBox.Environment} distributions.
When using the \\texttt{ASTSampleAction}, as opposed to \\texttt{ASTSeedAction},
we need to provide access to the sampleable environment." ->
GrayBox.environment(sim::Walk1DSim) = GrayBox.Environment(:x => sim.distribution)

@tex T"""We override the transition function from the \texttt{GrayBox} interface,
which takes an environment sample as input. We apply the sample in our simulator,
and return the log-likelihood.""" ->
function GrayBox.transition!(sim::Walk1DSim, sample::GrayBox.EnvironmentSample)
    sim.t += 1 # Keep track of time
    sim.x += sample[:x].value # Move agent using sampled value from input
    return logpdf(sample)::Real # Summation handled by `logpdf()`
end

@tex T"""\section{Black-box System}
The system under test, in this case a simple single-dimensional moving agent,
is always treated as black-box. The following interface functions are overridden
to minimally interact with the system, and use outputs from the system to
determine failure event indications and distance metrics.
"""

@tex T"""Now we override the \texttt{BlackBox} interface, starting with the
function that initializes the simulation object. Interface functions
ending in \texttt{!} may modify the \texttt{sim} object in place.""" ->
function BlackBox.initialize!(sim::Walk1DSim)
    sim.t = 0
    sim.x = sim.params.startx
end

@tex T"""We define how close we are to a failure event using a non-negative distance metric.""" ->
BlackBox.distance(sim::Walk1DSim) = max(sim.params.threshx - abs(sim.x), 0)

@tex T"""We define an indication that a failure event occurred.""" ->
BlackBox.isevent(sim::Walk1DSim) = abs(sim.x) >= sim.params.threshx

@tex T"""Similarly, we define an indication that the simulation is in a terminal state.""" ->
BlackBox.isterminal(sim::Walk1DSim) =
    BlackBox.isevent(sim) || sim.t >= sim.params.endtime

@tex T"""Lastly, we use our defined interface to evaluate the system under test.
Using the input sample, we return the log-likelihood, distance to an event, and event indication.""" ->
function BlackBox.evaluate!(sim::Walk1DSim, sample::GrayBox.EnvironmentSample)
    logprob::Real = GrayBox.transition!(sim, sample) # Step simulation
    d::Real = BlackBox.distance(sim) # Calculate miss distance
    event::Bool = BlackBox.isevent(sim) # Check event indication
    return (logprob::Real, d::Real, event::Bool)
end

@tex T"""\section{AST Setup and Running}
Setting up our simulation, we instantiate our simulation object and
pass that to the Markov decision proccess (MDP) object of the adaptive stress testing
formulation. We use Monte Carlo tree search (MCTS) with progressive widening on the action
space as our solver. Hyperparameters are passed to \texttt{MCTSPWSolver}, which is
a simple wrapper around the POMDPs.jl implementation of MCTS. Lastly, we solve the MDP
to produce a planner. Note we are using the \texttt{ASTSampleAction}.""" ->
function setup_ast(seed=0)
    # Create gray-box simulation object
    sim::GrayBox.Simulation = Walk1DSim()

    # AST MDP formulation object
    mdp::ASTMDP = ASTMDP{ASTSampleAction}(sim)
    mdp.params.debug = true # record metrics
    mdp.params.top_k = 10 # record top k best trajectories
    mdp.params.seed = seed # set RNG seed for determinism

    # Hyperparameters for MCTS-PW as the solver
    solver = MCTSPWSolver(n_iterations=1000, # number of algorithm iterations
                          exploration_constant=1.0, # UCT exploration
                          k_action=1.0, # action widening
                          alpha_action=0.5, # action widening
                          depth=sim.params.endtime) # tree depth

    # Get online planner (no work done, yet)
    planner = solve(solver, mdp)

    return planner
end

@tex T"""After setup, we \textit{search} for failures using the planner and output the best action trace.""" ->
begin
    planner = setup_ast()
    action_trace = search!(planner)
end

@tex T"""We can also \textit{playback} specific trajectories and print intermediate $x$-values.""" ->
final_state = playback(planner, action_trace, sim->sim.x)

@tex T"""Finally, we can print metrics associated with the AST run for further analysis.""" ->
failure_rate = print_metrics(planner)

@tex T"""\section{Solvers}
The solvers provided by the POMDPStressTesting.jl package include the following.
""" ->
begin
    # Reinforcement learning
        MCTSPWSolver
    # Deep reinforcement learning
        TRPOSolver
        PPOSolver
    # Stochastic optimization
        CEMSolver
    # Baselines
        RandomSearchSolver
end

@tex T"""\section{Reward Function}
The AST reward function gives a reward of $0$ if an event is found,
a reward of negative distance if no event is found at termination,
and the log-likelihood during the simulation.
\blfootnote{File generated using TeX.jl: \url{https://github.com/mossr/TeX.jl}}
"""
@texeq function R(p,e,d,τ)
    if τ && e
        return 0
    elseif τ && !e
        return -d
    else
        return log(p)
    end
end

@attachfile!

texgenerate() # Generate PDF


text:
using SqState
using Plots
plotly()

function main()
    @info "Initialising"

    data_path = joinpath(SqState.PROJECT_PATH, "../data", "dm.hdf5")
    data_name = "sq4"
    ρ = read_ρ(data_path, data_name)

    x_range = -5:0.1:5
    p_range = -5:0.1:5
    start_time = time()
    wf = WignerFunction(x_range, p_range)
    end_time = time()

    @info "Done, took $(end_time - start_time)(s)"

    start_time = time()
    w = wf(ρ)
    lim = maximum(abs.(w))
    p = heatmap(
        x_range,
        p_range,
        w,
        title="Wigner Function",
        xlabel="X",
        ylabel="P",
        # xticks=x_range,
        # yticks=p_range,
        clim=(-lim, lim),
        c=cgrad([
            RGBA(53/255, 157/255, 219/255, 1),
            RGBA(240/255, 240/255, 240/255, 1),
            RGBA(219/255, 64/255, 68/255, 1)]
        ),
        size=(900, 825)
    )
    end_time = time()
    @info "Render time: $(end_time - start_time)(s)"

    return p
end

main()


text:
#Generic Branch Models
abstract type AbstractBranchFormulation <: AbstractDeviceFormulation end
abstract type AbstractBoundedBranchFormulation <: AbstractBranchFormulation end

#Abstract Line Models
struct StaticLine <: AbstractBranchFormulation end
struct StaticLineBounds <: AbstractBoundedBranchFormulation end
struct StaticLineUnbounded <: AbstractBranchFormulation end
struct FlowMonitoredLine <: AbstractBranchFormulation end

#Abstract Transformer Models
struct StaticTransformer <: AbstractBranchFormulation end
struct StaticTransformerBounds <: AbstractBoundedBranchFormulation end
struct StaticTransformerUnbounded <: AbstractBranchFormulation end

# Not implemented yet
# struct TapControl <: AbstractBranchFormulation end
# struct PhaseControl <: AbstractBranchFormulation end

#################################### Branch Variables ##################################################
# Because of the way we integrate with PowerModels, most of the time PowerSimulations will create variables
# for the branch flows either in AC or DC.
flow_variables!(
    psi_container::PSIContainer,
    ::Type{<:PM.AbstractPowerModel},
    ::IS.FlattenIteratorWrapper{<:PSY.ACBranch},
) = nothing

function flow_variables!(
    psi_container::PSIContainer,
    ::Type{<:StandardPTDFModel},
    devices::IS.FlattenIteratorWrapper{B},
) where {B <: PSY.ACBranch}
    add_variable(psi_container, devices, variable_name(FLOW_ACTIVE_POWER, B), false)
    return
end

#################################### Flow Variable Bounds ##################################################
function _get_constraint_data(
    devices::IS.FlattenIteratorWrapper{B},
) where {B <: PSY.ACBranch}
    constraint_data = Vector{DeviceRange}(undef, length(devices))
    for (ix, d) in enumerate(devices)
        limit_values = (min = -1 * PSY.get_rate(d), max = PSY.get_rate(d))
        name = PSY.get_name(d)
        services_ub = Vector{Symbol}()
        for service in PSY.get_services(d)
            SR = typeof(service)
            push!(services_ub, Symbol("R$(PSY.get_name(service))_$SR"))
        end
        constraint_data[ix] = DeviceRange(name, limit_values, services_ub, Vector{Symbol}())
    end
    return constraint_data
end

function branch_rate_bounds!(
    psi_container::PSIContainer,
    devices::IS.FlattenIteratorWrapper{B},
    model::DeviceModel{B, <:AbstractBranchFormulation},
    ::Type{<:PM.AbstractDCPModel},
) where {B <: PSY.ACBranch}
    constraint_data = _get_constraint_data(devices)
    set_variable_bounds!(psi_container, constraint_data, FLOW_ACTIVE_POWER, B)
    return
end

function branch_rate_bounds!(
    psi_container::PSIContainer,
    devices::IS.FlattenIteratorWrapper{B},
    model::DeviceModel{B, <:AbstractBranchFormulation},
    ::Type{<:PM.AbstractPowerModel},
) where {B <: PSY.ACBranch}
    constraint_data = _get_constraint_data(devices)
    set_variable_bounds!(psi_container, constraint_data, FLOW_ACTIVE_POWER_FROM_TO, B)
    set_variable_bounds!(psi_container, constraint_data, FLOW_ACTIVE_POWER_TO_FROM, B)
    return
end

#################################### Rate Limits Constraints ###############################
function branch_rate_constraints!(
    psi_container::PSIContainer,
    devices::IS.FlattenIteratorWrapper{B},
    model::DeviceModel{B, <:AbstractBranchFormulation},
    ::Type{<:PM.AbstractActivePowerModel},
    feedforward::Nothing,
) where {B <: PSY.ACBranch}
    constraint_data = _get_constraint_data(devices)
    device_range(
        psi_container,
        constraint_data,
        constraint_name(RATE_LIMIT, B),
        variable_name(FLOW_ACTIVE_POWER, B),
    )
    return
end

function branch_rate_constraints!(
    psi_container::PSIContainer,
    devices::IS.FlattenIteratorWrapper{B},
    model::DeviceModel{B, <:AbstractBranchFormulation},
    ::Type{<:PM.AbstractPowerModel},
    feedforward::Nothing,
) where {B <: PSY.ACBranch}
    range_data = [(PSY.get_name(h), PSY.get_rate(h)) for h in devices]
    rating_constraint!(
        psi_container,
        range_data,
        constraint_name(RATE_LIMIT_FT, B),
        (
            variable_name(FLOW_ACTIVE_POWER_FROM_TO, B),
            variable_name(FLOW_REACTIVE_POWER_FROM_TO, B),
        ),
    )

    rating_constraint!(
        psi_container,
        range_data,
        constraint_name(RATE_LIMIT_TF, B),
        (
            variable_name(FLOW_ACTIVE_POWER_TO_FROM, B),
            variable_name(FLOW_REACTIVE_POWER_TO_FROM, B),
        ),
    )
    return
end

#################################### Flow Limits Constraints ##################################################
function branch_flow_constraints!(
    psi_container::PSIContainer,
    devices::IS.FlattenIteratorWrapper{PSY.MonitoredLine},
    model::DeviceModel{PSY.MonitoredLine, FlowMonitoredLine},
    ::Type{T},
    feedforward::Union{Nothing, AbstractAffectFeedForward},
) where {T <: PM.DCPPowerModel}
    flow_range_data = Vector{PSI.DeviceRange}(undef, length(devices))
    for (ix, d) in enumerate(devices)
        if PSY.get_flowlimits(d).to_from != PSY.get_flowlimits(d).from_to
            @info("Flow limits in Line $(PSY.get_name(d)) aren't equal. The minimum will be used in formulation $(T)")
        end
        limit = min(
            PSY.get_rate(d),
            PSY.get_flowlimits(d).to_from,
            PSY.get_flowlimits(d).from_to,
        )
        minmax = (min = -1 * limit, max = limit)
        flow_range_data[ix] = DeviceRange(PSY.get_name(d), minmax)
    end
    device_range(
        psi_container,
        flow_range_data,
        constraint_name(FLOW_LIMIT, PSY.MonitoredLine),
        variable_name(FLOW_ACTIVE_POWER, PSY.MonitoredLine),
    )
    return
end

function branch_flow_constraints!(
    psi_container::PSIContainer,
    devices::IS.FlattenIteratorWrapper{PSY.MonitoredLine},
    model::DeviceModel{PSY.MonitoredLine, FlowMonitoredLine},
    ::Type{<:PM.AbstractPowerModel},
    feedforward::Union{Nothing, AbstractAffectFeedForward},
)
    names = Vector{String}(undef, length(devices))
    limit_values_FT = Vector{MinMax}(undef, length(devices))
    limit_values_TF = Vector{MinMax}(undef, length(devices))
    to = Vector{PSI.DeviceRange}(undef, length(devices))
    from = Vector{PSI.DeviceRange}(undef, length(devices))
    for (ix, d) in enumerate(devices)
        limit_values_FT[ix] =
            (min = -1 * PSY.get_flowlimits(d).from_to, max = PSY.get_flowlimits(d).from_to)
        limit_values_TF[ix] =
            (min = -1 * PSY.get_flowlimits(d).to_from, max = PSY.get_flowlimits(d).to_from)
        names[ix] = PSY.get_name(d)
        to[ix] = DeviceRange(names[ix], limit_values_FT[ix])
        from[ix] = DeviceRange(names[ix], limit_values_TF[ix])
    end

    device_range(
        psi_container,
        to,
        constraint_name(FLOW_LIMIT_FROM_TO, PSY.MonitoredLine),
        variable_name(FLOW_ACTIVE_POWER_FROM_TO, PSY.MonitoredLine),
    )
    device_range(
        psi_container,
        from,
        constraint_name(FLOW_LIMIT_TO_FROM, PSY.MonitoredLine),
        variable_name(FLOW_ACTIVE_POWER_TO_FROM, PSY.MonitoredLine),
    )
    return
end


text:
"""
    - free cone
        x ∈ Rⁿ

    A Semismooth Newton Method for Fast, Generic Convex Programming
        https://arxiv.org/abs/1705.00772
"""

# free cone
function κ_free(z)
    z
end

# free cone Jacobian
function Jκ_free(z)
    p = ones(length(z))
    return Diagonal(p)
end


text:
# Based on http://scikit-learn.org/stable/modules/gaussian_process.html
using PyCall
@pyimport numpy as np
@pyimport sklearn.gaussian_process as gausspr

f(x) = x .* sin(x)

xtrain = np.atleast_2d([0.9, 1.0, 3.0, 5.0, 7.0, 9.0, 9.1])'
ytrain = f(xtrain)

gp = gausspr.GaussianProcess(theta0=1e-2, thetaL=1e-4, thetaU=1e-1)
gp["fit"](xtrain, ytrain)

xtest = np.atleast_2d(np.linspace(0, 10, 1000))'
ytest = f(xtest)
y_pred, sigma2_pred = gp["predict"](xtest, eval_MSE=true)

mse(yhat, y) = mean(sum((yhat .- y).^2))
rmse(yhat, y) = sqrt(mse(yhat, y))
mape(yhat, y) = 100.0 * mean(abs( (y .- yhat)./y ))

rmse(y_pred[:], ytest[:])

# Around 1.0 it should be pretty close:
y_pred[95:105]
ytest[95:105]
mape(y_pred[95:105], ytest[95:105])


text:
export @searchreturn

using Base: unwrap_unionall

include("taskmanager.jl")

struct CallableFinder{P <: SearchPolicy}
    modules::AbstractVector{Module}
    sp::P
    should_stop::Base.RefValue{Bool}
end

Base.IteratorSize(::Type{<: CallableFinder}) = Base.SizeUnknown()
Base.IteratorEltype(::Type{<: CallableFinder}) = Base.HasEltype()
Base.eltype(::Type{<: CallableFinder}) = Base.Callable

const SeenKey = Tuple{Module, Symbol}
# As hashing types and functions takes time, convert them into "fully
# qualified name" and has it instead.

seenkey(x::Any) = (Base, Symbol("## dummy ##"))
seenkey(x::Union{Module, DataType, UnionAll, Function}) = _seenkey(x, x)

_seenkey(x::Any, ::Any) = seenkey(x)
_seenkey(x::UnionAll, o) = _seenkey(unwrap_unionall(x), o)
_seenkey(x::Union{Module, DataType, Function}, o) = (parentmodule(x), nameof(o))
# `parentmodule(::UnionAll)` and do `unwrap_unionall`.  However, it may
# be wrapping `Union` and it would throw an error in this case.
# Manually handled here to fallback to `_seenkey(::Any, ::Any)`.

struct CFState
    names::Vector{Symbol}
    modules::Vector{Module}
    i_name::Int
    i_module::Int
    seen::Set{SeenKey}
end

advance_name(state::CFState, i_name) =
    CFState(state.names,
            state.modules,
            i_name,
            state.i_module,
            state.seen)

function advance_module(state::CFState)
    i = state.i_module + 1
    return CFState(
        names(state.modules[i]; all=true),
        state.modules,
        1,
        i,
        state.seen)
end

function Base.iterate(cf::CallableFinder)
    if isempty(cf.modules)
        return nothing
    end
    state = CFState(
        names(cf.modules[1]; all=true),
        collect(cf.modules),
        1,
        1,
        Set([seenkey(nothing)]))
    return Base.iterate(cf, state)
end

function Base.iterate(cf::CallableFinder, state::CFState)
    m = state.modules[state.i_module]
    for i in state.i_name:lastindex(state.names)
        name = state.names[i]
        yield()  # Maybe throttle? Or Julia scheduler would do it?
        if cf.should_stop[]
            return nothing
        end
        if isdefined(m, name)
            x = getproperty(m, name)
            if (cf.sp isa Recursive &&
                    x isa Module &&
                    is_defined_in(x, m) &&
                    !(x in state.modules))
                insert!(state.modules, state.i_module + 1, x)
            elseif is_locatable(x)
                key = seenkey(x)
                if !(key in state.seen)
                    push!(state.seen, key)
                    return (x, advance_name(state, i + 1))
                end
            end
        end
    end
    if state.i_module < length(state.modules)
        return Base.iterate(cf, advance_module(state))
    end
    return nothing
end

mayreturn(inferred::Type, typ::Type) = inferred <: typ

mayreturn(inferred::Union, typ::Type) =
    mayreturn(inferred.a, typ) || mayreturn(inferred.b, typ)

rettype_is(method::Method, typ::Type) =
    rettype_is(method.specializations, typ)

function rettype_is(specializations::Core.TypeMapEntry, typ::Type)
    @goto start
    while true
        specializations = specializations.next
        specializations === nothing && return false
        @label start
        specializations isa Core.TypeMapEntry || continue
        specializations.func isa Core.MethodInstance || continue
        specializations.func.rettype isa Type || continue
        mayreturn(specializations.func.rettype, typ) && return true
    end
end

function rettype_is(specializations::Core.TypeMapLevel, typ::Type)
    specializations.arg1 isa Vector || return false
    for a in specializations.arg1
        rettype_is(a, typ) && return true
    end
    return false
end

rettype_is(::Any, typ::Type) = false

rettype_is(typ::Type) = method -> rettype_is(method, typ)

function search_by_rettype(sp::SearchPolicy,
                           typ::Type,
                           modules::AbstractVector{Module},
                           should_stop::Ref{Bool} = Ref(false))
    if isempty(modules)
        modules = Base.loaded_modules_array()
    end
    callables = CallableFinder(modules, sp, should_stop)
    methiters = Base.Generator(methods, callables)
    meths = Iterators.flatten(methiters)
    return Iterators.filter(rettype_is(typ), meths)
end

struct SearchQuery
    policy::SearchPolicy
    returntype::Type
    modules::Vector{Module}
end

struct BackgroundSearch
    id::Int
    query::SearchQuery
    found::Vector
    task::Task
    done::Base.RefValue{Bool}
    should_stop::Base.RefValue{Bool}
end

function showquery(io::IO, query::SearchQuery)
    print(io, "::")
    printstyled(io, query.returntype, color=:cyan)
    if !isempty(query.modules)
        print(io, " from ")
        printstyled(io, query.modules, color=:cyan)
    end
    if query.policy isa Recursive
        print(io, " recursively")
    end
    return
end

showquery(io, search::BackgroundSearch) = showquery(io, search.query)

function Base.show(io::IO, query::SearchQuery)
    print(io, nameof(typeof(query)), ": ")
    showquery(io, query)
    return
end

function Base.show(io::IO, search::BackgroundSearch)
    print(io, nameof(typeof(search)), " id=", search.id)

    # print " [done]" etc.
    print(io, " [")
    if search.should_stop[]
        printstyled(io, "canceled"; color=:yellow)
    elseif search.done[]
        printstyled(io, "done"; color=:green)
    elseif istaskdone(search.task)
        printstyled(io, "error"; color=:light_red, bold=true)
    elseif !istaskstarted(search.task)
        printstyled(io, "queued"; color=:magenta)
    else
        printstyled(io, "active"; color=:red)
    end
    print(io, "]")
    # Notes: `istaskdone(search.task)` can't be used here since `show`
    # is called at the end of the `search.task`.  (Though this can be
    # worked around by using additional task watching `search.task`.)

    print(io, " ", length(search.found), " found")
    return
end

function Base.show(io::IO, ::MIME"text/plain", search::BackgroundSearch)
    show(io, search)
    println(io)
    if search.done[] || search.should_stop[] || istaskdone(search.task)
        print(io, "Searched ")
    else
        print(io, "Searching ")
    end
    showquery(io, search)
    return
end

function _start_search_return(id, args...)
    query = SearchQuery(args...)
    found = []
    done = Ref(false)
    should_stop = Ref(false)
    task = @task begin
        append!(found, search_by_rettype(args..., should_stop))
        done[] = true

        @info "Finished search id=$id.  Look it up by `@search $id`."
        show(stderr, "text/plain", background_searches[id])
        println(stderr)
    end
    enqueue(task)
    return BackgroundSearch(id, query, found, task, done, should_stop)
end

const background_searches = BackgroundSearch[]

nextid() = length(background_searches) + 1

function schedule_search_return(typ::Type,
                                modules::AbstractVector{Module};
                                policy = Recursive())
    search = _start_search_return(nextid(), policy, typ, modules)
    # TODO: don't start search task immediately
    push!(background_searches, search)
    return search
end

"""
    kill(search::InteractiveCodeSearch.BackgroundSearch)

Stop `search`.
"""
Base.kill(search::BackgroundSearch) = stop(search)

"""
    InteractiveCodeSearch.stop(search::BackgroundSearch)
    InteractiveCodeSearch.stop(id::Int = 0)

Stop `search`.  The first form can be invoked via `kill(search)` as well
(where `kill` is the function exported from `Base`).  See also `bg`.
"""
function stop(search::BackgroundSearch)
    # schedule(...; error=true) is not the right way to stop the task?
    # https://github.com/JuliaLang/julia/issues/25353#issuecomment-354807622
    search.should_stop[] = true
    wait(search.task)
    return
end
stop(id::Int = 0) = kill(bg(id))

"""
    InteractiveCodeSearch.bg(id::Int = 0)

Lookup background search by `id`.  For convenience, `id = 0` is the
last search and `id = -1` is the second last search.  More generally,
non-positive `id` is treated as the offset from `end`.
"""
function bg(id::Int = 0)
    if id < 1
        return background_searches[end + id]
    else
        return background_searches[id]
    end
end

# @search _s1
code_search(::SearchPolicy, search::BackgroundSearch) =
    search_methods(search.found)

# @search 1
code_search(p::SearchPolicy, id::Int) = code_search(p, bg(id))

function searchreturn(typ::Type, modules::AbstractVector{Module};
                      policy = Recursive())
    search_methods(search_by_rettype(policy, typ, modules))
end

"""
    @searchreturn Type [Module...]

Search functions returning type `Type` in `Module`s.  As this search
typically takes some time to finish, interactive matcher will not be
launched by this command.  Instead, a "handle" to the search in
background is returned which can be queried via `@search` later.
Calling `kill` (`Base.kill`) on the handle cancels the search.

# Limitations

* It does not work with Julia >= 1.2.
* Running `@searchreturn` for many modules may be slow for the _first_
  run.  Thus, searching from all modules (i.e., not specifying
  `Module` arguments) may take tens of seconds to minutes, depending
  of what are loaded.  Searching within `Base` takes about 30 seconds.
  After `DifferentialEquations` is loaded, searching for all modules
  takes 1.5 minutes.  Note that searching from the same module for the
  second time is fast (a few seconds), even if different `Type` is
  specified.
* The functions must be executed (JIT'ed) once for `@searchreturn` to
  find their returned by type.
* Any IO operations (like printing in REPL) would be slow while the search
  is active in background.
* Keyboard interruption does not work well while background search is
  active.  You need to hit CTRL-C multiple times to terminate a "foreground"
  code.  Furthermore, it will bring down the background search task as well.

# Examples
```julia-repl
julia> using LinearAlgebra, SparseArrays

julia> spzeros(3, 3)

julia> @searchreturn AbstractMatrix LinearAlgebra SparseArrays
┌ Info: Search result is stored in variable `_s1`.
│ You can interactively narrow down the search result later by
└ `@search _s1` or `@search 1`.

BackgroundSearch id=1 [active] 0 found
Searching ::AbstractArray{T,2} where T from Module[LinearAlgebra SparseArrays] recursively

julia> @search _s1

julia> kill(_s1)  # stop the search
```

If you prefer giving a custom name to the search result, just assign it to
some variable.

```julia-repl
julia> my_search = @searchreturn AbstractMatrix LinearAlgebra SparseArrays
julia> @search my_search
```
"""
macro searchreturn(typ, modules...)
    modules_array = Expr(:ref, Module, esc.(modules)...)
    id = nextid()
    result = Symbol("_s", id)
    quote
        let search = schedule_search_return($(esc(typ)), $modules_array)
            global $(esc(result))
            if isdefined($(__module__), $(QuoteNode(result)))
                @warn $("""
                Variable `$result` exists!
                To refer to the search result, `ans` must be saved to some
                variable **NOW**.
                """)
            else
                @info $("""
                Search result is stored in variable `$result`.
                You can interactively narrow down the search result later by
                `@search $result` or `@search $id`.
                """)
                $(esc(result)) = search
            end
            println(stderr)
            search
        end
    end
end


text:
struct ExternalSample end

function (external_traverse::ExternalSample)(sol::DeepCFRSolver, h, p, t)
    game = sol.game
    current_player = player(game, h)

    if isterminal(game, h)
        return Float32(CFR.utility(game, p, h))

    elseif iszero(current_player)
        a = chance_action(game, h)
        h′ = next_hist(game,h,a)
        return external_traverse(sol, h′, p, t)

    elseif current_player == p
        I = vectorized(game, infokey(game, h)) # info KEY here, not infostate
        A = actions(game, h)
        v_σ_Ia = Vector{Float32}(undef, length(A))
        v_σ = 0.0f0

        σ = regret_match_strategy(sol, I, p)
        for (i,a) in enumerate(A)
            h′ = next_hist(game, h, a)
            v = external_traverse(sol, h′, p, t)
            v_σ_Ia[i] = v
            v_σ += σ[i]*v
        end

        r = v_σ_Ia .-= v_σ
        push!(sol.Mv[p], I,t,r)

        return v_σ

    else
        I = vectorized(game, infokey(game, h)) # info KEY here, not infostate
        A = actions(game, h)
        σ = regret_match_strategy(sol, I, current_player)
        push!(sol.Mπ, I,t,σ)
        a = A[weighted_sample(σ)]
        h′ = next_hist(game, h, a)
        return external_traverse(sol, h′, p, t)
    end
end

"""
Infokey type of game may not be in vectorized form.

`vectorized(game::Game, I::infokeytype(game))` returns the original key type
in vectorized form to be pushed through a neural network.
"""
function vectorized end

vectorized(game::Game, I) = I

"""
Sample index of vector according to weights given by vector (sum of weights assumed to be 1.0)
"""
function weighted_sample(rng::Random.AbstractRNG, σ::AbstractVector)
    t = rand(rng)
    i = 1
    cw = σ[1]
    while cw < t && i < length(σ)
        i += 1
        cw += σ[i]
    end
    return i
end

weighted_sample(σ::AbstractVector) = weighted_sample(Random.GLOBAL_RNG, σ)

function regret_match_strategy(sol::DeepCFRSolver, I::AbstractVector, p)
    return regret_match_strategy(sol.V[p](I))
end

function regret_match_strategy(r::AbstractVector)
    s = 0.0f0
    max_neg_idx = 0
    max_neg = -Inf
    for i in eachindex(r)
        if r[i] > 0.0f0
            s += r[i]
        else
            v = r[i]
            if v > max_neg
                max_neg_idx = i
                max_neg = v
            end
            r[i] = 0.0f0
        end
    end

    if s > 0.0f0
        return r ./= s
    else
        r[max_neg_idx] = 1.0f0
        return r
    end
end


text:
# f(x)=sf*tanh(x)

sf=2.5

Fstanh(x::Array{Float64,2})=(sf*tanh(x),[]);
Fstanh_inplace(value,aux,x::Array{Float64,2})=copy!(value,sf*tanh(x))

Dstanh(derivativeIDX,f_c,faux_c,grad_c,grad_n,x)=axpy!(sf,grad_c.*(1.-(f_c/sf).^2),grad_n)

if PROC=="GPU" ## TODO
    function Fstanh(x::CudaArray)
        tmp=copy(x)
        stanh!(sf,x,tmp)
        return (tmp,[]) # problem here is that I don't know how to free tmp -- this will cause memory leak I think.
    end
    Fstanh_inplace(value::CudaArray,aux,x::CudaArray)=stanh!(sf,x,value)

    Dstanh(derivativeIDX,f_c,faux_c,grad_c::CudaArray,grad_n::CudaArray,x::CudaArray)=Dstanh!(sf,grad_c,f_c,grad_n)
end

Derivative[Fstanh]=Dstanh
Inplace[Fstanh]=Fstanh_inplace

stanh(A::ADnode)=ADnode(Fstanh,A)
export stanh

max_stars_repo_name: davidbarber/AutoDiff

max_issues_repo_name: davidbarber/AutoDiff

max_forks_repo_name: davidbarber/AutoDiff


text:
module AlgorithmTests

using Test
using HDF5
using FFTW

using PicoQuant
include("test_utils.jl")

@testset "Algorithm tests" begin
    @testset "Test state preparation algorithm" begin
        @test begin
            qubits = 1
            circ = create_simple_preparation_circuit(qubits, 4, 42)

            ψ = get_statevector_using_picoquant(circ,
                                                big_endian=false)
            ψ′ = get_statevector_using_qiskit(circ, big_endian=false)

            overlap = (a, b) -> abs(transpose(conj.(a)) * b)

            overlap(ψ, ψ′) ≈ 1.
        end

        @test begin
            qubits = 3
            circ = create_simple_preparation_circuit(qubits, 4, 42)

            ψ = get_statevector_using_picoquant(circ,
                                                big_endian=false)
            ψ′ = get_statevector_using_qiskit(circ, big_endian=false)

            overlap = (a, b) -> abs(transpose(conj.(a)) * b)

            overlap(ψ, ψ′) ≈ 1.
        end
    end

    @testset "Small QFT circuit test" begin
        @test begin
            n = 3

            prep_circ = create_simple_preparation_circuit(n, 3, 43)
            qft_circ = create_qft_circuit(n)

            full_circ = prep_circ.combine(qft_circ)

            qft_input = get_statevector_using_qiskit(prep_circ, big_endian=false)
            ref_output = ifft(qft_input)
            norm = x -> x ./ sqrt(sum(x .* conj.(x)))
            ref_output = norm(ref_output)

            ψ = get_statevector_using_picoquant(full_circ)

            overlap = (a, b) -> abs(transpose(conj.(a)) * b)

            overlap(ψ, ref_output) ≈ 1.
        end
    end

    @testset "Larger QFT circuit test" begin
        @test begin
            n = 8

            prep_circ = create_simple_preparation_circuit(n, 3, 43)
            qft_circ = create_qft_circuit(n)

            full_circ = prep_circ.combine(qft_circ)

            qft_input = get_statevector_using_qiskit(prep_circ, big_endian=false)
            ref_output = ifft(qft_input)
            norm = x -> x ./ sqrt(sum(x .* conj.(x)))
            ref_output = norm(ref_output)

            ψ = get_statevector_using_picoquant(full_circ)

            overlap = (a, b) -> abs(transpose(conj.(a)) * b)

            overlap(ψ, ref_output) ≈ 1.
        end
    end

    @testset "Small RQC circuit test" begin
        @test begin
            n = 3; m = 3; depth = 8

            rqc = create_RQC(n, m, depth)

            qiskit_Ψ = get_statevector_using_qiskit(rqc, big_endian=false)

            picoquant_ψ = get_statevector_using_picoquant(rqc)

            overlap = (a, b) -> abs(transpose(conj.(a)) * b)

            overlap(picoquant_ψ, qiskit_Ψ) ≈ 1.
        end
    end

    @testset "iSWAP RQC circuit test" begin
        @test begin
            n = 3; m = 3; depth = 8

            rqc = create_RQC(n, m, depth;
                            use_iswap=true, final_Hadamard_layer=true)

            qiskit_Ψ = get_statevector_using_qiskit(rqc, big_endian=false)

            picoquant_ψ = get_statevector_using_picoquant(rqc)

            overlap = (a, b) -> abs(transpose(conj.(a)) * b)

            overlap(picoquant_ψ, qiskit_Ψ) ≈ 1.
        end
    end
end
end


text:
function setcompilers(app)

cpustatus0 = Sys.which(app.cpucompiler);
cpustatus1 = Sys.which("g++");
cpustatus2 = Sys.which("/usr/bin/g++");
cpustatus3 = Sys.which("/usr/local/bin/g++");
cpustatus4 = Sys.which("/opt/local/bin/g++");

if cpustatus0 != nothing
    print("Using " * app.cpucompiler * " compiler for CPU source code\n");
elseif cpustatus1 != nothing
    app.cpucompiler = "g++";
    print("Using g++ compiler for CPU source code\n");
elseif cpustatus2 != nothing
    app.cpucompiler = "/usr/bin/g++";
    print("Using /usr/bin/g++ compiler for CPU source code\n");
elseif cpustatus3 != nothing
    app.cpucompiler = "/usr/local/bin/g++";
    print("Using /usr/local/bin/g++ compiler for CPU source code\n");
elseif cpustatus4 != nothing
    app.cpucompiler = "/opt/local/bin/g++";
    print("Using /opt/local/bin/g++ compiler for CPU source code\n");
else
    error("Exasim search in /usr/bin, /usr/local/bin, and /opt/local/bin and could not find C++ compiler. Please see the documentation to install it. After installation, please set its path to app.cpucompiler");
end

if app.mpiprocs>1
    mpistatus0 = Sys.which(app.mpicompiler);
    mpistatus1 = Sys.which("mpicxx");
    mpistatus2 = Sys.which("/usr/bin/mpicxx");
    mpistatus3 = Sys.which("/usr/local/bin/mpicxx");
    mpistatus4 = Sys.which("/opt/local/bin/mpicxx");
    mpistatus5 = Sys.which("mpicxx-openmpi-mp");
    mpistatus6 = Sys.which("/usr/bin/mpicxx-openmpi-mp");
    mpistatus7 = Sys.which("/usr/local/bin/mpicxx-openmpi-mp");
    mpistatus8 = Sys.which("/opt/local/bin/mpicxx-openmpi-mp");
    if mpistatus0 != nothing
        print("Using " * app.mpicompiler * " compiler for MPI source code\n");
    elseif mpistatus5 != nothing
        app.mpicompiler = "mpicxx-openmpi-mp";
        app.mpirun = "mpirun-openmpi-mp";
        print("Using mpicxx-openmpi-mp compiler for MPI source code\n");
    elseif mpistatus6 != nothing
        app.mpicompiler = "/usr/bin/mpicxx-openmpi-mp";
        app.mpirun = "/usr/bin/mpirun-openmpi-mp";
        print("Using /usr/bin/mpicxx-openmpi-mp compiler for MPI source code\n");
    elseif mpistatus7 != nothing
        app.mpicompiler = "/usr/local/bin/mpicxx-openmpi-mp";
        app.mpirun = "/usr/local/bin/mpirun-openmpi-mp";
        print("Using /usr/local/bin/mpicxx-openmpi-mp compiler for MPI source code\n");
    elseif mpistatus8 != nothing
        app.mpicompiler = "/opt/local/bin/mpicxx-openmpi-mp";
        app.mpirun = "/opt/local/bin/mpirun-openmpi-mp";
        print("Using opt/local/bin/mpicxx-openmpi-mp compiler for MPI source code\n");
    elseif mpistatus1 != nothing
        app.mpicompiler = "mpicxx";
        app.mpirun = "mpirun";
        print("Using mpicxx compiler for MPI source code\n");
    elseif mpistatus2 != nothing
        app.mpicompiler = "/usr/bin/mpicxx";
        app.mpirun = "/usr/bin/mpirun";
        print("Using /usr/bin/mpicxx compiler for MPI source code\n");
    elseif mpistatus3 != nothing
        app.mpicompiler = "/usr/local/bin/mpicxx";
        app.mpirun = "/usr/local/bin/mpirun";
        print("Using /usr/local/bin/mpicxx compiler for MPI source code\n");
    elseif mpistatus4 != nothing
        app.mpicompiler = "/opt/local/bin/mpicxx";
        app.mpirun = "/opt/local/bin/mpirun";
        print("Using opt/local/bin/mpicxx compiler for MPI source code\n");
    else
        error("Exasim search in /usr/bin, /usr/local/bin, and /opt/local/bin and could not find MPI compiler. Please see the documentation to install it. After installation, please set its path to app.mpicompiler");
    end
    app.mpirun = replace(app.mpicompiler, "mpicxx" => "mpirun");
end

if app.platform == "gpu"
    gpustatus0 = Sys.which(app.gpucompiler);
    gpustatus1 = Sys.which("nvcc");
    gpustatus2 = Sys.which("/usr/bin/nvcc");
    gpustatus3 = Sys.which("/usr/local/bin/nvcc");
    gpustatus4 = Sys.which("/opt/local/bin/nvcc");

    if gpustatus0 != nothing
        print("Using " * app.gpucompiler * " compiler for GPU source code\n");
    elseif gpustatus1 != nothing
        app.gpucompiler = "nvcc";
        print("Using nvcc compiler for GPU source code\n");
    elseif gpustatus2 != nothing
        app.gpucompiler = "/usr/bin/nvcc";
        print("Using /usr/bin/nvcc compiler for GPU source code\n");
    elseif gpustatus3 != nothing
        app.gpucompiler = "/usr/local/bin/nvcc";
        print("Using /usr/local/bin/nvcc compiler for GPU source code\n");
    elseif gpustatus4 != nothing
        app.gpucompiler = "opt/local/bin/nvcc";
        print("Using opt/local/bin/nvcc compiler for GPU source code\n");
    else
        error("Exasim search in /usr/bin, /usr/local/bin, and /opt/local/bin and could not find NVCC compiler. Please see the documentation to install it. After installation, please set its path to app.gpucompiler");
    end
end

mdir = pwd();
ii = findlast(app.codename, mdir);
coredir = mdir[1:ii[end]] * "/lib";

cpulib = 0;
gpulib = 0;
if Sys.isapple()
    if (isfile(coredir * "/Mac/commonCore.a")) && (isfile(coredir * "/Mac/opuCore.a"))
        cpulib = 1;
    end
    if (isfile(coredir * "/Mac/gpuCore.a"))
        gpulib = 1;
    end
elseif Sys.isunix()
    if (isfile(coredir * "/Linux/commonCore.a")) && (isfile(coredir * "/Linux/opuCore.a"))
        cpulib = 1;
    end
    if (isfile(coredir * "/Linux/gpuCore.a"))
        gpulib = 1;
    end
elseif Sys.iswindows()
    if (isfile(coredir * "/Windows/commonCore.a")) && (isfile(coredir * "/Windows/opuCore.a"))
        cpulib = 1;
    end
    if (isfile(coredir * "/Windows/gpuCore.a"))
        gpulib = 1;
    end
end

if cpulib==0
    print("Generating CPU core libraries.");
    genlib(app.cpucompiler, "", coredir);
end

if gpulib==0
    if (length(app.gpucompiler)>0) && (app.platform == "gpu")
        print("Generating GPU core library.");
        genlib("", app.gpucompiler, coredir);
    end
end

return app

end


text:
# sample_target_mismatch-gaussian_vs_t-kernel_vs_graph
#
# This script compares the values of the univariate graph Stein
# discrepancy and kernel discrepancy obtained when samples are drawn i.i.d.
# from their targets.

using SteinDistributions:
    SteinGaussian,
    SteinScaleLocationStudentT,
    gradlogdensity
using SteinDiscrepancy:
    SteinGaussianKernel,
    ksd,
    gsd

include("experiment_utils.jl")

# random number generator seed
# we allow this to be set on the commandline
@parseintcli seed "s" "seed" 7
# set the target distribution
@parsestringcli distname "t" "target" "gaussian"
# set the discrepancy
@parsestringcli discrepancytype "d" "discrepancy" "gaussiankernel"
# Specify dimension of sampled vectors
d = 1
# Specify a maximum number of samples
maxn = round(Int, 5e5)
# gaussian target
gaussian_target = SteinGaussian(d)
# the graphsolver
graphsolver = "gurobi"
#graphsolver = "clp"
# setup kernel (if necessary)
kernel = nothing
if discrepancytype == "gaussiankernel"
    kernel = SteinGaussianKernel()
end
# setup target
if distname == "studentt"
    df = 10.0
    target = SteinScaleLocationStudentT(df, 0.0, sqrt((df-2.0)/df))
elseif distname == "gaussian"
    target = SteinGaussian(d)
end
# define gradlogp
function gradlogp(x::Array{Float64,1})
    gradlogdensity(target, x)
end
# Draw samples from target distribution
X = @setseed rand(target,maxn)
# Sample sizes at which optimization problem will be solved
ns = vcat(100:100:min(1000,maxn),
          2000:1000:min(10000,maxn),
          20000:10000:min(100000,maxn),
          200000:100000:maxn)

# Solve optimization problem at each sample size
for i in ns
    @printf("[Beginning n=%d]:\n", i)
    # Create program with data subset
    Xi = X[1:i,:]
    discrepancy = solvetime = nothing
    if kernel != nothing
        @printf("[Beginning kernel computations n=%d]\n", i)
        result = ksd(points=Xi,
                     gradlogdensity=gradlogp,
                     kernel=kernel)
        discrepancy = sqrt(result.discrepancy2)
        solvetime = result.solvetime
    elseif discrepancytype == "graph"
        @printf("[Beginning graph solver n=%d]\n", i)
        result = gsd(points=Xi,
                     gradlogdensity=gradlogp,
                     solver=graphsolver)
        discrepancy = sum(result.objectivevalue)
        solvetime = sum(result.solvetime)
    end
    # save data
    instance_data = Dict{Any, Any}(
        "target" => distname,
        "n" => i,
        "d" => d,
        "seed" => seed,
        "discrepancytype" => discrepancytype,
        "discrepancy" => discrepancy,
        "solvetime" => solvetime,
    )
    save_json(
        instance_data;
        dir="sample_target_mismatch-gaussian_vs_t-kernel_vs_graph",
        target=distname,
        discrepancytype=discrepancytype,
        n=i,
        d=d,
        seed=seed
    )
end

@printf("COMPLETE!")


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

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

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

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

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

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


text:
## -------------------------------------------------------------------
function load_raw_model()
    model = iJR.load_model("raw_model")
    model_size = size(model)
    nz_abs = ChU.nzabs_range(model.S)
    @info("MAT MODEL LOADED", model_size, nz_abs)
    ChN.test_fba(model, iJR.BIOMASS_IDER; summary = false)
    return model
end

## -------------------------------------------------------------------
# rescale bounds bounds
# The abs maximum bounds will be set to 100
function rescale_bounds!(model)
    @info("CLAMP BOUNDS", iJR.ABS_MAX_BOUND)
    foreach(model.rxns) do ider
            ChU.isfixxed(model, ider) && return # fixxed reaction are untouched

            old_ub = ChU.ub(model, ider)
            new_ub = old_ub == 0.0 ? 0.0 : iJR.ABS_MAX_BOUND
            ChU.ub!(model, ider, new_ub)

            old_lb = ChU.lb(model, ider)
            new_lb = old_lb == 0.0 ? 0.0 : -iJR.ABS_MAX_BOUND
            ChU.lb!(model, ider, new_lb)
    end
    return model
end

## -------------------------------------------------------------------
# CLOSING EXCHANGES
function close_exchanges!(model, exchs)

    @info("CLOSE EXCANGES", length(exchs))
    # Close, for now, all ChU.exchanges for avoiding it to be in revs
    # The reversible reactions will be splited for modeling cost
    # Exchanges have not associated cost, so, we do not split them
    foreach(exchs) do rxn
        ChU.ub!(model, rxn, 0.0) # Closing all outtakes
        ChU.lb!(model, rxn, 0.0) # Closing all intakes
    end
    return model
end

## -------------------------------------------------------------------
# ENZYMATIC COST
function add_cost(model; cost_ub = 1.0)
    ## -------------------------------------------------------------------
    # ENZYMATIC COST INFO
    # The cost will be introduced as a reaction, we follow the same cost models as
    # Beg et al. (2007): https://doi.org/10.1073/pnas.0609845104.
    # A new balance equations is then added:
    #        Σ(rᵢ*costᵢ) + tot_cost = 0
    #    Because the cost coefficients (costᵢ) < 0 (it resamble a reactant), the system must allocate
    #    the fluxes (rᵢ) so that Σ(rᵢ*costᵢ) = tot_cost, and tot_cost
    #    are usually bounded [0.0, 1.0]
    cost_info = Dict()
    fwd_ider(rxn) = string(rxn, ChU.FWD_SUFFIX);
    bkwd_ider(rxn) = string(rxn, ChU.BKWD_SUFFIX);
    for rxn in model.rxns
        # The ChU.exchanges, the atpm and the biomass are synthetic reactions, so,
        # they have should not have an associated enzimatic cost
        any(startswith.(rxn, ["EX_", "DM_"])) && continue
        rxn == iJR.BIOMASS_IDER && continue
        rxn == iJR.ATPM_IDER && continue

        # Only the internal, non reversible reactions have an associated cost
        # We will split the rev reactions, so we save the cost for both versions (fwd, bkwd)
        if ChU.isrev(model, rxn)
            cost_info[fwd_ider(rxn)] = -iJR.beg_enz_cost(rxn)
            cost_info[bkwd_ider(rxn)] = -iJR.beg_enz_cost(rxn)
        else
            cost_info[rxn] = -iJR.beg_enz_cost(rxn)
        end
    end

    ## -------------------------------------------------------------------
    # SPLITING REVS
    nfwd_suffix = ChU.FWD_SUFFIX
    bkwd_suffix = ChU.BKWD_SUFFIX
    @info("SPLITING REVS", nfwd_suffix, bkwd_suffix)
    model = ChU.split_revs(model;
        get_fwd_ider = fwd_ider,
        get_bkwd_ider = bkwd_ider,
    );

    ## -------------------------------------------------------------------
    # ADDING COST REACCION
    cost_met_id = "cost"
    cost_exch_id = iJR.COST_IDER

    # info
    to_add = cost_info |> length
    min_abs_cost = cost_info |> values .|> abs |> minimum
    max_abs_cost = cost_info |> values .|> abs |> maximum
    @info("ADDING COST",
        to_add, min_abs_cost, max_abs_cost, cost_met_id, cost_exch_id
    )

    M, N = size(model)
    cost_met = ChU.Met(cost_met_id, S = collect(values(cost_info)), rxns = collect(keys(cost_info)), b = 0.0)
    model = ChU.expanded_model(model, M + 1, N + 1)
    ChU.set_met!(model, ChU.findempty(model, :mets), cost_met)
    cost_exch = ChU.Rxn(cost_exch_id, S = [1.0], mets = [cost_met_id], lb = -iJR.ABS_MAX_BOUND, ub = 0.0, c = 0.0)
    ChU.set_rxn!(model, ChU.findempty(model, :rxns), cost_exch);

    # tot_cost is the exchange that controls the bounds of the
    # enzimatic cost contraint, we bound it to [0, 1.0]
    ChU.lb!(model, cost_exch_id, 0.0);
    ChU.ub!(model, cost_exch_id, cost_ub);

    return model
end

## -------------------------------------------------------------------
# SET BASE EXCHANGE
function reset_exchanges(model, exchs)

    model = ChU.fix_dims(model)

    @info("SETTING EXCHANGES")
    # To control the intakes just the metabolites defined in the
    # base_intake_info (The minimum medium) will be opened.
    # The base model will be constraint as in a cultivation with
    # experimental minimum xi
    # see Cossios paper (see README)

    foreach(exchs) do rxn
        (rxn in [iJR.COST_IDER]) && return
        ChU.ub!(model, rxn, iJR.ABS_MAX_BOUND) # Opening all outakes
        ChU.lb!(model, rxn, 0.0) # Closing all intakes
    end

    # see Cossios paper (see README) for details in the Chemostat bound constraint
    xi = minimum(Nd.val(:xi))
    intake_info = iJR.load_base_intake_info()
    ChSS.apply_bound!(model, xi, intake_info; emptyfirst = true)

    ## -------------------------------------------------------------------
    ChN.test_fba(model, iJR.BIOMASS_IDER, iJR.COST_IDER)
    return model
end

## -------------------------------------------------------------------
function scale_model(model, scale_factor)

    base_nzabs_range = ChU.nzabs_range(model.S)
    base_size = size(model)

    # Scale model (reduce S ill-condition)
    scl_model = iszero(scale_factor) ?
        deepcopy(model) :
        ChU.well_scaled_model(model, scale_factor)

    scl_size = size(scl_model)
    scl_nzabs_range = ChU.nzabs_range(scl_model.S)

    @info("Model",
        base_size, base_nzabs_range,
        scl_size, scl_nzabs_range
    ); println()

    return scl_model
end

## -------------------------------------------------------------------
function make_fva_model(model, exp, D; scale_factor = 1000.0)
    @info("DOING FVA", exp, D)

    ## -------------------------------------------------------------------
    # prepare model
    model0 = scale_model(deepcopy(model), scale_factor)

    M, N = size(model0)
    exp_xi = Nd.val(:xi, exp)
    intake_info = iJR.intake_info(exp)
    ChSS.apply_bound!(model0, exp_xi, intake_info;
        emptyfirst = true)

    ChN.test_fba(exp, model0, iJR.BIOMASS_IDER, iJR.COST_IDER)
    fva_model = ChLP.fva_preprocess(model0,
        check_obj = iJR.BIOMASS_IDER,
        verbose = true
    );
    ChN.test_fba(exp, fva_model, iJR.BIOMASS_IDER, iJR.COST_IDER)

    return fva_model

end

## -------------------------------------------------------------------
function make_max_model(model; scale_factor = 1000.0)
    # This model is bounded by the maximum rates found for EColi.
    # Data From:
    # Varma, (1993): 2465–73. https://doi.org/10.1128/AEM.59.8.2465-2473.1993.
    # Extract max exchages from FIG 3 to form the maximum polytope

    @info("DOING MAX MODEL")

    max_model = scale_model(deepcopy(model), scale_factor)

    # Biomass
    # 2.2 1/ h
    ChU.bounds!(max_model, iJR.BIOMASS_IDER, 0.0, 2.2)

    Fd_rxns_map = iJR.load_rxns_map()
    # 40 mmol / gDW h
    # ChU.bounds!(max_model, Fd_rxns_map["GLC"], -40.0, 0.0)
    ChU.bounds!(max_model, Fd_rxns_map["GLC"], -8.0, 0.0)
    # 45 mmol/ gDW
    ChU.bounds!(max_model, Fd_rxns_map["AC"], 0.0, 40.0)
    # 55 mmol/ gDW h
    ChU.bounds!(max_model, Fd_rxns_map["FORM"], 0.0, 55.0)
    # 20 mmol/ gDW h
    ChU.bounds!(max_model, Fd_rxns_map["O2"], -20.0, 0.0)

    # fva
    max_model = ChLP.fva_preprocess(max_model,
        check_obj = iJR.BIOMASS_IDER,
        verbose = true
    );

    ## -------------------------------------------------------------------
    test_model = deepcopy(max_model)
    for (exp, D) in Nd.val(:D) |> enumerate
        cgD_X = Nd.cval(:GLC, exp) * Nd.val(:D, exp) / Nd.val(:X, exp)
        ChU.lb!(test_model, iJR.EX_GLC_IDER, -cgD_X)
        fbaout = ChLP.fba(test_model, iJR.BIOMASS_IDER)
        biom = ChU.av(test_model, fbaout, iJR.BIOMASS_IDER)
        @info("Test", exp, cgD_X, D, biom); println()
    end

    return max_model
end


text:
# run with `julia 4.jl`
using Printf

function toInt(str)::Int64
    return parse(Int64, str)
end

function crossout(boards, call)
    for (b, board) in enumerate(boards)
        for (i, row) in enumerate(board)
            for (j, item) in enumerate(row)
                if item == call
                    boards[b][i][j] = 0
                    @goto endBoard
                end
            end
        end
        @label endBoard
    end
    return
end

function sumBoard(board)
    total = 0
    for row in board
        for item in row
            total += item
        end
    end
    return total
end

function getWinner(boards)
    win = 0
    for board in boards
        # check rows
        for row in board
            for item in row
                if item != 0
                    @goto endRow
                end
            end
            win = sumBoard(board)
            @goto winBoard

            @label endRow
        end
        # check columns
        # transpose does not work :(
        # break
        for (i, row) in enumerate(board)
            for (j, item) in enumerate(row)
                if board[j][i] != 0
                    @goto endRow2
                end
            end
            win = sumBoard(board)
            @goto winBoard

            @label endRow2
        end
    end
    @label winBoard
    return win
end

function main()
    file = "4.in"
    fp = open(file, "r")

    boards = []
    calls = []
    single_board = []
    start = false
    line_count = 0

    for lines in readlines(fp)
        # increment line_count
        line_count += 1

        if (line_count + 4) % 6 == 0
            if start
                push!(boards, single_board)
                single_board = []
            else
                start = true
            end
        end

        if (line_count == 1)
            calls = map(toInt, split(lines, ','))
        else
            if lines == "" continue end
            push!(single_board, map(toInt, split(lines)))
            # println(map(toInt, split(lines)))
        end
    end
    close(fp)

    winner = 0
    for call in calls
        crossout(boards, call)
        # println(boards)
        # break
        winner = getWinner(boards)
        if (winner != 0)

            println("$winner $call")
            winner *= call
            break
        end
    end
    println(winner)
end

main()


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

using Knet

const loop_iters = 500

function loop(x = 2, n = loop_iters)
  r = one(x)/one(x)
  while n > 0
    r *= sin(x)
    n -= 1
  end
  return r
end

@result loop Knet grad(loop)(2)

using Knet

logreg(W, b, x) = sum(sigm.(W * x .+ b))

W = randn(10,28^2)
b = randn(10)
x = rand(28^2)

@result logreg Knet grad(x -> logreg(x...))(($W,$b,$x))